Iterating over arbitrary dimension of numpy.array

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

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.

Solution 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:

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

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

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

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

Solution 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.

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

Iterating over arbitrary dimension of numpy.array

Solution 1:

Solution 2:

Solution 3:

Related

Recent Posts