10

Say I have a 3D numpy.array, e.g. with dimensions x y z, is there a way to iterate over slices along a particular axis? Something like:

for layer in data.slices(dim=2):
    # do something with layer

Edit: To clarify, the example is a dim=3 array, i.e. shape=(len_x, len_y, len_z). Elazar and equivalently kamjagin's solutions work, but aren't that general - you have to construct the [:, :, i] by hand, which means you need to know the dimensions, and the code isn't general enough to handle arrays of arbitrary dimensions. You can fill missing dimension by using something like [..., :], but again you still have to construct this yourself.

Sorry, should have been clearer, the example was a bit too simple!

lost
  • 2,210
  • 2
  • 20
  • 34
  • 1
    Check out http://stackoverflow.com/questions/1589706/iterating-over-arbitrary-dimension-of-numpy-array . The swapaxes method is the fastest, but least clear. The approach by kamjagin can be generalized by constructing a tuple to pass into the brackets (i.e. `data[tuple(slice(None), slice(None), i)]` is the same as `data[:,:,i]`). – AFoglia Jun 28 '13 at 02:01
  • @AFoglia I don't think that the timings on that question are very relevant. Using `swapaxes` or `rollaxis` will take a little more time in setting up the loop, but the actual iterations are faster, see the timings added to my answer. In your very small example, setting up was dominating over the actual iteration. I disagree on the readability, but I may be too used to numpy to notice it. – Jaime Jun 28 '13 at 16:10
  • 1
    @AFoglia numpy has a neat trick for building slice (tuples) separately from the object they're to index: `numpy.s_[:,:,i]` is equivalent to `tuple(slice(None), slice(None), i)`. – JAB Jun 28 '13 at 16:19
  • (Note that `numpy.s_` has an attribute, `maketuple`, that determines whether or not usage of a single slice will result in a tuple or not [i.e. `slice(start, end, step)` vs. `(slice(start, end, step),)`]. The default is `False`.) – JAB Jun 28 '13 at 16:34
  • @JAB Yes, but if he wants a general solution for an arbitrary dimension of an arbitrary rank array, it's easier to programmatically make a tuple with `(slice(None),)*(idim)+(idx,)` than with the `:` character. You can't do `data[(:,)*(idim)+(idx,)]`. – AFoglia Jun 28 '13 at 19:58
  • @AFoglia but you can do `data[np.index_exp[:]*idim + np.index_exp[idx]]`, or `data[np.index_exp[:]*idim + (np.s_[idx],)]`, or even `s=np.index_exp` followed by `data[s[:]*idim + s[idx]]`. (I like that last one the best, personally. Also, I'd forgotten that `numpy.index_exp` was the `IndexExpression` instance with `maketuple` set to `True` by default.) Though in your specific case `data[..., idx]` would probably be the smartest choice. – JAB Jun 28 '13 at 20:31

5 Answers5

13

Iterating over the first dimension is very easy, see below. To iterate over the others, roll that dimension to the front and do the same:

>>> data = np.arange(24).reshape(2, 3, 4)
>>> for dim_0_slice in data: # the first dimension is easy
...     print dim_0_slice
... 
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
[[12 13 14 15]
 [16 17 18 19]
 [20 21 22 23]]
>>> for dim_1_slice in np.rollaxis(data, 1): # for the others, roll it to the front
...     print dim_1_slice
... 
[[ 0  1  2  3]
 [12 13 14 15]]
[[ 4  5  6  7]
 [16 17 18 19]]
[[ 8  9 10 11]
 [20 21 22 23]]
>>> for dim_2_slice in np.rollaxis(data, 2):
...     print dim_2_slice
... 
[[ 0  4  8]
 [12 16 20]]
[[ 1  5  9]
 [13 17 21]]
[[ 2  6 10]
 [14 18 22]]
[[ 3  7 11]
 [15 19 23]]

EDIT Some timings, to compare different methods for largish arrays:

In [7]: a = np.arange(200*100*300).reshape(200, 100, 300)

In [8]: %timeit for j in xrange(100): a[:, j]
10000 loops, best of 3: 60.2 us per loop

In [9]: %timeit for j in xrange(100): a[:, j, :]
10000 loops, best of 3: 82.8 us per loop

In [10]: %timeit for j in np.rollaxis(a, 1): j
10000 loops, best of 3: 28.2 us per loop

In [11]: %timeit for j in np.swapaxes(a, 0, 1): j
10000 loops, best of 3: 26.7 us per loop
Jaime
  • 65,696
  • 17
  • 124
  • 159
  • 1
    Since Numpy 1.11 use [`np.moveaxis( data, k, 0 )`](https://numpy.org/doc/stable/reference/generated/numpy.moveaxis.html) instead. – Jonathan H Jan 21 '22 at 17:52
3

This could probably be solved more elegantly than this, but one way of doing it if you know dim beforehand(e.g. 2) is:

for i in range(data.shape[dim]):
    layer = data[:,:,i]

or if dim=0

for i in range(data.shape[dim]):
    layer = data[i,:,:]

etc.

kamjagin
  • 3,614
  • 1
  • 22
  • 24
  • +1 That's the idea, but you would need to build your indexing tuple dynamically, e.g. `idx = (slice(None),)*dim + (i,) + (slice(None),)*(2-dim)`, and then `layer = data[idx]`. It's less of a hassle to roll the dimension to a set position (the first is the most convenient) and then have always the same iteration machinery, see my answer. – Jaime Jun 27 '13 at 23:14
  • Ahh great. Definitely more elegant with the roll-solution. – kamjagin Jun 28 '13 at 09:45
1

I think the original question was ambiguous, and so is the title:

  • Iterating over x[k,:,:,...] for all k yields as many items as x.shape[0], so I would call this iterating over a dimension;
  • In contrast, iterating over a slice in my mind means iterating e.g. over x[:,i,j,k,...] for all i,j,k. For example, iterating over all columns in an ndarray.

Although that's not what the OP asked (after clarification), people looking for a solution to the latter might find the following code useful:

from itertools import product

def iterslice(x,axis=0):
    sub = [ range(s) for s in x.shape ]
    sub[axis] = (slice(None),)
    for p in product(*sub):
        yield x[p]
Jonathan H
  • 7,591
  • 5
  • 47
  • 80
0

something like that?

>>> data = np.array([[1,2,3],[4,5,6],[7,8,9]])
>>> for layer in [data[:,i] for i in range(3)]:
...     print layer
... 
[1 4 7]
[2 5 8]
[3 6 9]
Elazar
  • 20,415
  • 4
  • 46
  • 67
0

Correct me if I am wrong but it seems to me that your 3D array is gonna look like :

>>> my_array.shape
    (3,N)

Where N is the size of your array. So if you want to iterate over one dimension, you can just do :

>>> for item in my_array[1,:]:

This will iterate on the second dimension.

C.J
  • 55
  • 7
  • ah sorry, should have been clearer - this is a dim=3 array, i.e. shape=(nx, ny, nz) – lost Jun 27 '13 at 22:58