numpy沿n-空间应用

“基本”沿轴应用模型是在一个轴上迭代,并将另一个轴传递给函数:

In [197]: def foo(x):         # return same size
     ...:     return x*2
     ...: np.array([foo(x) for x in np.arange(12).reshape(3,4)])
     ...: 
Out[197]: 
array([[ 0,  2,  4,  6],
       [ 8, 10, 12, 14],
       [16, 18, 20, 22]])
In [198]: def foo(x):
     ...:     return x.sum()   # return one less dim
     ...: np.array([foo(x) for x in np.arange(12).reshape(3,4)])
     ...: 
Out[198]: array([ 6, 22, 38])
In [199]: def foo(x):
     ...:     return x.sum(keepdims=True)   # condense the dim
     ...: np.array([foo(x) for x in np.arange(12).reshape(3,4)])
     ...: 
Out[199]: 
array([[ 6],
       [22],
       [38]])

你的4d问题可以通过按摩来解决。

In [200]: arr_4d = np.arange(24).reshape(2,3,2,2)
In [201]: arr_2d = arr_4d.reshape(6,4).T
In [202]: res = np.array([foo(x) for x in arr_2d])
In [203]: res
Out[203]: 
array([[60],
       [66],
       [72],
       [78]])
In [204]: res.reshape(2,2)
Out[204]: 
array([[60, 66],
       [72, 78]])

这相当于:

In [205]: arr_4d[:,:,0,0].sum()
Out[205]: 60
In [206]: foo(arr_4d[:,:,0,0].ravel())
Out[206]: array([60])

apply_along_axis

In [209]: np.apply_along_axis(foo,0,arr_4d.reshape(6,2,2))
Out[209]: 
array([[[60, 66],
        [72, 78]]])

foo 可以将其输入重塑为2d,并将其传递给采用2d的函数。 apply_along_index 使用 np.ndindex

In [212]: list(np.ndindex(2,2))
Out[212]: [(0, 0), (0, 1), (1, 0), (1, 1)]

np.vectorize signature 参数,我相信它可以用于处理您的案例。它可能需要转换输入,以便在前两个轴上迭代,将后两个轴传递给函数。请参阅我的答案 https://stackoverflow.com/a/46004266/901925 .

无需重塑或交换,我可以借助 ndindex .

def foo2(x):
    return x.sum(axis=1, keepdims=True) # 2d

的最后2维的索引迭代器 arr_4d :

In [260]: idx = np.ndindex(arr_4d.shape[-2:])

vectorize 和 apply...

In [261]: r1 = foo2(arr_4d[:,:,0,0]).shape
In [262]: r1
Out[262]: (2, 1)

结果数组:

In [263]: res = np.zeros(r1+arr_4d.shape[-2:])
In [264]: res.shape
Out[264]: (2, 1, 2, 2)

现在迭代:

In [265]: for i,j in idx:
     ...:     res[...,i,j] = foo2(arr_4d[...,i,j])
     ...:     
In [266]: res
Out[266]: 
array([[[[ 12.,  15.],
         [ 18.,  21.]]],


       [[[ 48.,  51.],
         [ 54.,  57.]]]])

numpy.apply_over_axes 与一个 for 循环以迭代其他变化的轴。

我自己滚。我很想知道这个方法和@hpaulj的方法之间是否有任何性能差异,以及是否有理由相信编写自定义c模块会带来显著的改进。当然@hpaulj的方法更通用,因为这是我只需要在适当的阵列上执行操作所特有的。

def apply_along_space(f, np_array, axes):
    # apply the function f on each subspace given by iterating over the axes listed in axes, e.g. axes=(0,2)
    for slic in itertools.product(*map(lambda ax: range(np_array.shape[ax]) if ax in axes else [slice(None,None,None)], range(len(np_array.shape)))):
        f(np_array[slic])
    return np_array