"numpy沿着n维空间应用"

Question

我有一个4d数组，我想将一个函数应用到通过迭代最后两个维度获得的每个2d切片。也就是说，将f（2d_array）应用于（x，y，0，0），f（2d_array）应用于（x，y，0，1）等。我的函数对数组进行了适当的操作，因此维数是相同的，但一般的解决方案将返回一个形状为（x'，y'，w，z）的数组，其中w和z是原始数组的最后两个维度。在

这显然可以推广到nD阵列上的mD切片。在

有没有什么内置的功能可以做到这一点？在

Answer 1

我想您正在寻找类似于^{}与for循环相结合来迭代其他可变轴。在

Answer 2

我自己滚的。我很想知道这个方法和@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

Answer 3

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

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问题可以通过按摩来解决。在

^{pr2}$

相当于：

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需要一个接受1d数组的函数，但可以这样应用：

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

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

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

进行测试计算以确定回路的形状。vectorize和{}做这种测试。在

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.]]]])