迭代numpy.array的任意维度

Is there function to get an iterator over an arbitrary dimension of a numpy array?

是否有函数在numpy数组的任意维度上获取迭代器？

Iterating over the first dimension is easy...

迭代第一维很容易......

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

But iterating over other dimensions is harder. For example, the last dimension:

但迭代其他维度更难。例如，最后一个维度：

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

I'm making a generator to do this myself, but I'm surprised there isn't a function named something like numpy.ndarray.iterdim(axis=0) to do this automatically.

我自己做了一个生成器，但我很惊讶没有像numpy.ndarray.iterdim（axis = 0）这样的函数自动执行此操作。

4 个解决方案

#1

What you propose is quite fast, but the legibility can be improved with the clearer forms:

你的建议很快，但更清晰的形式可以提高易读性：

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

or, better (faster, more general and more explicit):

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

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

However, the first approach above appears to be about twice as slow as the swapaxes() approach:

但是，上面的第一种方法似乎是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

I would guess that this is because swapaxes() does not copy any data, and because the handling of c[:,:,i] might be done through general code (that handles the case where : is replaced by a more complicated slice).

我猜这是因为swapaxes（）不复制任何数据，因为c [：，：，i]的处理可能是通过通用代码完成的（处理以下情况：由更复杂的切片替换）。

Note however that the more explicit second solution c[...,i] is both quite legible and quite fast:

但请注意，更明确的第二个解决方案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

#2

I'd use the following:

我使用以下内容：

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

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

The function rollaxis creates a new view on the array. In this case it's moving axis 2 to the front, equivalent to the operation c.transpose(2, 0, 1).

功能rollaxis在阵列上创建一个新视图。在这种情况下，它将轴2移动到前面，相当于操作c.transpose（2,0,1）。

#3

So, one can iterate over the first dimension easily, as you've shown. Another way to do this for arbitrary dimension is to use numpy.rollaxis() to bring the given dimension to the first (the default behavior), and then use the returned array (which is a view, so this is fast) as an iterator.

因此，正如您所示，可以轻松地迭代第一维。对任意维度执行此操作的另一种方法是使用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)

EDIT: I'll comment that I submitted a PR to numpy to address this here: https://github.com/numpy/numpy/pull/3262. The concensus was that this wasn't enough to add to the numpy codebase. I think using np.rollaxis is the best way to do this, and if you want an interator, wrap it in iter().

编辑：我会评论我提交了一个PR到numpy来解决这个问题：https：//github.com/numpy/numpy/pull/3262。一致认为，这还不足以增加numpy代码库。我认为使用np.rollaxis是最好的方法，如果你想要一个interator，请将它包装在iter（）中。

#4

I guess there is no function. When I wrote my function, I ended up taking the iteration EOL also suggested. For future readers, here it is:

我猜没有功能。当我编写我的函数时，我最终还是采用了EOL的迭代建议。对于未来的读者，这里是：

def iterdim(a, axis=0) :
  a = numpy.asarray(a);
  leading_indices = (slice(None),)*axis
  for i in xrange(a.shape[axis]) :
    yield a[leading_indices+(i,)]

#1