python - Fitting a 3D array of data to a 1D function with numpy or scipy

Question

4I am currently trying to fit a lot of data to a sine function. In the case where I only have one set of data (1D array), scipy.optimize.curve_fit() works fine. However it does not permit a higher dimensional data input if the function itself is only one dimensional as far as i can see. I don't want to iterate over the array using for loops as that works incredibly slow in python.

My code so far should look similar to this:

from scipy import optimize
import numpy as np    
def f(x,p1,p2,p3,p4): return p1 + p2*np.sin(2*np.pi*p3*x + p4)      #fit function

def fit(data,guess):
   n = data.shape[0] 
   leng = np.arange(n)
   param, pcov = optimize.curve_fit(f,leng,data,guess)
   return param, pcov

where data is a threedimensional array (shape=(x,y,z)) and I would like to fit each line data[:,a,b] to the function with param being a (4,y,z) shaped array as output. Of course, for multidimensional data this results in a

ValueError: operands could not be broadcast together with shapes (2100,2100) (5)

Maybe there is an easy solution to this but I am not sure how to do it. Any suggestions?

Searching for an answer to my question proofed quite difficult since most topics with those keywords relate to the fitting of higher dimensional functions.

Answer 1

Using np.apply_along_axis() solves your problem. Just do this:

func1d = lambda y, *args: optimize.curve_fit(f, xdata=x, ydata=y, *args)[0] #<-- [0] to get only popt
param = np.apply_along_axis( func1d, axis=2, arr=data )

See the example below:

from scipy import optimize
import numpy as np
def f(x,p1,p2,p3,p4):
    return p1 + p2*np.sin(2*np.pi*p3*x + p4)
sx = 50  # size x
sy = 200 # size y
sz = 100 # size z
# creating the reference parameters
tmp = np.empty((4,sy,sz))
tmp[0,:,:] = (1.2-0.8) * np.random.random_sample((sy,sz)) + 0.8
tmp[1,:,:] = (1.2-0.8) * np.random.random_sample((sy,sz)) + 0.8
tmp[2,:,:] = np.ones((sy,sz))
tmp[3,:,:] = np.ones((sy,sz))*np.pi/4
param_ref = np.empty((4,sy,sz,sx))     # param_ref in this shape will allow an
for i in range(sx):                    # one-shot evaluation of f() to create 
    param_ref[:,:,:,i] = tmp           # the data sample
# creating the data sample
x = np.linspace(0,2*np.pi)
factor = (1.1-0.9)*np.random.random_sample((sy,sz,sx))+0.9
data = f(x, *param_ref) * factor       # the one-shot evalution is here
# finding the adjusted parameters
func1d = lambda y, *args: optimize.curve_fit(f, xdata=x, ydata=y, *args)[0] #<-- [0] to get only popt
param = np.apply_along_axis( func1d, axis=2, arr=data )