Iterating over arbitrary dimension of numpy.array

Question

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


        

Answer 1

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