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)