遍历 ndarray 的切片

Question

Say I have a 3D numpy.array, eg with dimensions xyz, 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, ie 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!

Answer 1

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

Answer 2

This could probably be solved more elegantly than this, but one way of doing it if you know dim beforehand(eg 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.

Answer 3

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]

Answer 4

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.

Answer 5

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 eg 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 second case above 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]