Python 2.7：在多维 Numpy 数组中循环一维纤维

Question

我正在寻找一种在 3+ 维数组中沿任意维度循环 1D 纤维（行、列和多维等效项）的方法。

在2D 数组中，这是相当简单的，因为纤维是行和列，所以只要说for row in A 就可以完成工作。但对于 3D 数组，此表达式迭代 2D 切片，而不是 1D 纤维。

以下是可行的解决方案：

import numpy as np
A = np.arange(27).reshape((3,3,3))
func = np.sum
for fiber_index in np.ndindex(A.shape[:-1]):
    print func(A[fiber_index])

但是，我想知道是否有这样的东西：

1. 更惯用的
2. 更快

希望你能帮忙！

Answer 1

我想你可能正在寻找numpy.apply_along_axis

In [10]: def my_func(x):
   ...:     return x**2 + x

In [11]: np.apply_along_axis(my_func, 2, A)
Out[11]: 
array([[[  0,   2,   6],
        [ 12,  20,  30],
        [ 42,  56,  72]],

       [[ 90, 110, 132],
        [156, 182, 210],
        [240, 272, 306]],

       [[342, 380, 420],
        [462, 506, 552],
        [600, 650, 702]]])

尽管许多 NumPy 函数（包括 sum）都有自己的 axis 参数来指定使用哪个轴：

In [12]: np.sum(A, axis=2)
Out[12]: 
array([[ 3, 12, 21],
       [30, 39, 48],
       [57, 66, 75]])

Answer 2

numpy 提供了多种不同的方式来循环一维或多维。

你的例子：

func = np.sum
for fiber_index in np.ndindex(A.shape[:-1]):
    print func(fiber_index)
    print A[fiber_index]

产生类似的东西：

(0, 0)
[0 1 2]
(0, 1)
[3 4 5]
(0, 2)
[6 7 8]
...

在 1st 2 dim 上生成所有索引组合，为您的函数提供最后一个 1D 光纤。

查看ndindex 的代码。这是有启发性的。我试图在https://stackoverflow.com/a/25097271/901925 中提取它的本质。

它使用as_strided 生成一个虚拟矩阵，nditer 在该矩阵上进行迭代。它使用“multi_index”模式来生成索引集，而不是该虚拟对象的元素。迭代本身是使用__next__ 方法完成的。这与numpy 编译代码中当前使用的索引样式相同。

http://docs.scipy.org/doc/numpy-dev/reference/arrays.nditer.html Iterating Over Arrays 有很好的解释，包括在 cython 中这样做的示例。

许多函数，其中sum、max、product，让您指定要迭代的轴（轴）。你的例子，sum，可以写成：

np.sum(A, axis=-1)
np.sum(A, axis=(1,2))   # sum over 2 axes

等价物是

np.add.reduce(A, axis=-1)

np.add 是ufunc，reduce 指定迭代模式。还有很多其他ufunc，以及其他迭代模式——accumulate、reduceat。你也可以定义自己的ufunc。

xnx 建议

np.apply_along_axis(np.sum, 2, A)

值得深入研究apply_along_axis，看看它是如何遍历A 的维度的。在您的示例中，它在while 循环中遍历所有可能的i,j，计算：

outarr[(i,j)] = np.sum(A[(i, j, slice(None))])

在索引元组中包含slice 对象是一个不错的技巧。请注意，它编辑列表，然后将其转换为元组以进行索引。那是因为元组是不可变的。

---

您的迭代可以通过将轴滚动到末端来沿任何轴应用。这是一个“便宜”的操作，因为它只是改变了步幅。

def with_ndindex(A, func, ax=-1):
    # apply func along axis ax
    A = np.rollaxis(A, ax, A.ndim) # roll ax to end (changes strides)
    shape = A.shape[:-1]
    B = np.empty(shape,dtype=A.dtype)
    for ii in np.ndindex(shape):
        B[ii] = func(A[ii])
    return B

我对 3x3x3、10x10x10 和 100x100x100 A 数组进行了一些计时。这种np.ndindex 方法始终比apply_along_axis 方法快三分之一。直接使用np.sum(A, -1) 会快很多。

因此，如果func 仅限于在一维光纤上运行（与sum 不同），那么ndindex 方法是一个不错的选择。