迭代numpy.array的任意维度

Question

有没有什么函数可以让我们在numpy数组的任意维度上获取一个迭代器呢？

在第一个维度上迭代是很简单的……

In [63]: c = numpy.arange(24).reshape(2,3,4)

In [64]: for r in c :
   ....:     print r
   ....: 
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
[[12 13 14 15]
 [16 17 18 19]
 [20 21 22 23]]

但是在其他维度上迭代就比较难了。例如，最后一个维度：

In [73]: for r in c.swapaxes(2,0).swapaxes(1,2) :
   ....:     print r
   ....: 
[[ 0  4  8]
 [12 16 20]]
[[ 1  5  9]
 [13 17 21]]
[[ 2  6 10]
 [14 18 22]]
[[ 3  7 11]
 [15 19 23]]

我正在自己制作一个生成器来实现这个功能，但我很惊讶为什么没有一个叫做numpy.ndarray.iterdim(axis=0)这样的函数可以自动完成这个任务。

Answer 1

所以，你可以很容易地遍历第一个维度，就像你展示的那样。对于任意维度，另一种方法是使用numpy.rollaxis()，把你想要的维度移动到第一个位置（这是默认的行为），然后用返回的数组（这个数组是一个视图，所以速度很快）作为迭代器。

In [1]: array = numpy.arange(24).reshape(2,3,4)

In [2]: for array_slice in np.rollaxis(array, 1):
   ....:     print array_slice.shape
   ....:
(2, 4)
(2, 4)
(2, 4)

编辑：我想说的是，我提交了一个请求给numpy来解决这个问题，链接在这里：https://github.com/numpy/numpy/pull/3262。大家的共识是，这个改动不足以加入到numpy的代码库中。我认为使用np.rollaxis是最好的方法，如果你想要一个迭代器，可以用iter()来包装它。

Answer 2

我会使用以下代码：

c = numpy.arange(2 * 3 * 4)
c.shape = (2, 3, 4)

for r in numpy.rollaxis(c, 2):
    print(r)

这个rollaxis函数是用来给数组创建一个新的视图的。简单来说，它把数组的第二个轴（也就是第三个维度）移动到了最前面，这个操作和 c.transpose(2, 0, 1) 是一样的。

Answer 3

你提到的方法速度很快，但可读性可以通过更清晰的形式来提高：

for i in range(c.shape[-1]):
    print c[:,:,i]

或者，更好一点（更快、更通用、更明确）：

for i in range(c.shape[-1]):
    print c[...,i]

不过，上面提到的第一种方法似乎比 swapaxes() 方法慢了大约一倍：

python -m timeit -s 'import numpy; c = numpy.arange(24).reshape(2,3,4)' \
    'for r in c.swapaxes(2,0).swapaxes(1,2): u = r'
100000 loops, best of 3: 3.69 usec per loop

python -m timeit -s 'import numpy; c = numpy.arange(24).reshape(2,3,4)' \
    'for i in range(c.shape[-1]): u = c[:,:,i]'
100000 loops, best of 3: 6.08 usec per loop

python -m timeit -s 'import numpy; c = numpy.arange(24).reshape(2,3,4)' \
    'for r in numpy.rollaxis(c, 2): u = r'
100000 loops, best of 3: 6.46 usec per loop

我猜这是因为 swapaxes() 不会复制任何数据，而且处理 c[:,:,i] 的方式可能是通过一些通用的代码来完成的（这些代码可以处理更复杂的切片情况，而不仅仅是 :）。

不过要注意，更明确的第二种解决方案 c[...,i] 既可读性强又相当快：

python -m timeit -s 'import numpy; c = numpy.arange(24).reshape(2,3,4)' \
    'for i in range(c.shape[-1]): u = c[...,i]'
100000 loops, best of 3: 4.74 usec per loop