I'm trying to reshape my input arrays in a way that's described below to be able to fit a curve on the data points, plot it, etc. It works fine, but I'm afraid it's not the most efficient way of doing it, and also it's hard to understand and read. It's part of a bigger function, but I only post here the critical part and not the whole function. I also added little comments to help you understand what's going on. I know it's ugly.
import numpy as np from itertools import chain #example inputs delay = np.array([10,20,30,40,50,60,70,None]) omega = [[1,2,3,4],[1.1,2.2,3.3,None],[1.4,2.4,3.4,None],[1.5,2.8,None,None], [1.8,2.9,None,None],[1.9,None,None,None],[2.0,None,None,None],[None,None,None,None]] """ desired output explained: with delay's first component 10 I want to assign 1, 2, 3 and 4 so the outputs starts like this: [10,10,10,10] , [1,2,3,4] then with delay's second component 20 I want to assign 1.1, 2.2, 3.3 and drop the None value. at the moment the outputs should look like this: [10,10,10,10,20,20,20] , [1,2,3,4,1.1,2.2,3.3] and so on. """ def spp_method(omegas, delays): #dropping all None values from delays delays = delays[delays != np.array(None)] #dropping all the None values from omegas omegas_ext =  for element in omegas: item = [x for x in element if x is not None] omegas_ext.append(item) #expand delays with the number of values in omegas_ext appropriate component delays_exp =  for idx in range(len(omegas_ext)): if len(omegas_ext[idx])>1: value = delays[idx] item_to_add = [value]*(len(omegas_ext[idx])) delays_exp.append(item_to_add) elif len(omegas_ext[idx]) == 1: value = delays[idx] delays_exp.append([value]) #put the values into simple array to plot and to fit curve. delays_unpacked =  omegas_unpacked =  for element in omegas_ext: for item in element: omegas_unpacked.append(item) delays_unpacked = list(chain.from_iterable(delays_exp)) return np.array(delays_unpacked), np.array(omegas_unpacked) y, x = spp_method(omega, delay) print(x) #outputs to: [1. 2. 3. 4. 1.1 2.2 3.3 1.4 2.4 3.4 1.5 2.8 1.8 2.9 1.9 2. ] print(y) #outputs to: [10 10 10 10 20 20 20 30 30 30 40 40 50 50 60 70]
which are correct.
Any improvements in the code are well-appreciated.