Iterate across arbitrary dimension in numpy

Question

I have a multidimensional numpy array, and I need to iterate across a given dimension. Problem is, I won't know which dimension until runtime. In other words, given an array m, I could want

m[:,:,:,i] for i in xrange(n)

or I could want

m[:,:,i,:] for i in xrange(n)

etc.

I imagine that there must be a straightforward feature in numpy to write this, but I can't figure out what it is/what it might be called. Any thoughts?

Answer 1

There are many ways to do this. You could build the right index with a list of slices, or perhaps alter m's strides. However, the simplest way may be to use np.swapaxes:

import numpy as np
m=np.arange(24).reshape(2,3,4)
print(m.shape)
# (2, 3, 4)

Let axis be the axis you wish to loop over. m_swapped is the same as m except the axis=1 axis is swapped with the last (axis=-1) axis.

axis=1
m_swapped=m.swapaxes(axis,-1)
print(m_swapped.shape)
# (2, 4, 3)

Now you can just loop over the last axis:

for i in xrange(m_swapped.shape[-1]):
    assert np.all(m[:,i,:] == m_swapped[...,i])

Note that m_swapped is a view, not a copy, of m. Altering m_swapped will alter m.

m_swapped[1,2,0]=100
print(m)
assert(m[1,0,2]==100)

Answer 2

You can use slice(None) in place of the :. For example,

from numpy import *

d = 2  # the dimension to iterate

x = arange(5*5*5).reshape((5,5,5))
s = slice(None)  # :

for i in range(5):
    slicer = [s]*3  # [:, :, :]
    slicer[d] = i   # [:, :, i]
    print x[slicer] # x[:, :, i]