9

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?

chimeracoder
  • 20,648
  • 21
  • 60
  • 60
  • 1
    Possible duplicate of http://stackoverflow.com/questions/1589706/iterating-over-arbitrary-dimension-of-numpy-array – Katriel Aug 18 '10 at 15:29
  • Possible duplicate of [Iterating over arbitrary dimension of numpy.array](https://stackoverflow.com/questions/1589706/iterating-over-arbitrary-dimension-of-numpy-array) – Mr Tsjolder from codidact Aug 02 '18 at 12:15

2 Answers2

6

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)
unutbu
  • 842,883
  • 184
  • 1,785
  • 1,677
5

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]
tom10
  • 67,082
  • 10
  • 127
  • 137