使用numpy或scipy将3D数据数组拟合到1D函数

Question

我现在正在尝试把很多数据拟合到一个正弦函数上。如果我只有一组数据（1D数组），那么使用scipy.optimize.curve_fit()就能很好地工作。但是，似乎如果函数本身是单维的，它不允许输入更高维度的数据。我不想用for循环来遍历数组，因为在Python中这样做会非常慢。

到目前为止，我的代码应该类似于这个：

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

这里的数据是一个三维数组（shape=(x,y,z)），我想把每一行data[:,a,b]拟合到这个函数上，param的输出应该是一个(4,y,z)形状的数组。

当然，对于多维数据，这会导致一个错误：

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

也许有简单的解决办法，但我不太确定该怎么做。有什么建议吗？

我在寻找答案时发现这很困难，因为大多数相关主题都与更高维度函数的拟合有关。

Answer 1

使用 np.apply_along_axis() 可以解决你的问题。只需要这样做：

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 )

下面是一个例子：

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 )