多个数组的Python apply_along_axis

Question

如果我有一个函数,f(x)将一个1d数组作为参数并产生一个1d数组作为输出,我可以使用numpy.apply_along_axis将该函数应用于一行的二维数组X的每一行是f的有效参数.

现在我想用一个带两个参数的函数做类似的事情.例如,我有一个函数f(x,y),它将两个1d数组作为参数,我还有两个2d数组X,Y都有n行.我想对每对行应用f,产生一个又有n行的数组.

如何以有效的方式实现这一目标？

我也对变体感兴趣,其中f需要更多参数或涉及更高维数组:

例如,f可以取3个x,y,z形状的阵列(2,2); (3); (5,)并产生形状(4,4)的结果.

我有X,Y,Z形状(50,100,2,2); (50,100,3); (50,100,5)并想要一个形状的结果(50,100,4,4)

Answer 1

查看代码, numpy.apply_along_axis我看到它只是迭代其他维度,将您的函数应用于每个'行'.有额外的代码允许维度大约为2.但对于2d,X它归结为:

result = np.empty_like(X)
for i, x in enumerate(X):
    result[i] = func1d(x)

还有代码来推断result应该具有的形状.例如,如果func1d是np.sum,那么result将是1d,而不是输入的2d.

所以这个功能没有特殊的"效率".多个输入的扩展可以是普通的Python zip:

 result = np.empty_like(X)
 for i,(x,y) in enumerate(zip(X,Y)):
     result[i] = func1d(x,y)

np.ndindex是生成指数的便利工具.值得一看的是它的代码.它使用通用的numpy迭代器,np.nditer见:http://docs.scipy.org/doc/numpy/reference/arrays.nditer.html

例如,f可以取3个x,y,z形状的阵列(2,2); (3); (5,)并产生形状(4,4)的结果.

我有X,Y,Z形状(50,100,2,2); (50,100,3); (50,100,5)并想要一个形状的结果(50,100,4,4)

for i,j in np.ndindex(50,100):
    result[i,j,:,:] = f(X[i,j,:,:], Y[i,j,:,:], Z[i,j,:,:])

':'不是必需的,但要明确我们正在对2个维度进行索引,并对其余维度进行切片.如果你想迭代1和3维,并切片第2维,则需要它们.