Pythonic way of iterating over 3D array

Question

I have a 3D array in Python and I need to iterate over all the cubes in the array. That is, for all (x,y,z) in the array's dimensions I need to access the cube:

array[(x + 0, y + 0, z + 0)]
array[(x + 1, y + 0, z + 0)]
array[(x + 0, y + 1, z + 0)]
array[(x + 1, y + 1, z + 0)]
array[(x + 0, y + 0, z + 1)]
array[(x + 1, y + 0, z + 1)]
array[(x + 0, y + 1, z + 1)]
array[(x + 1, y + 1, z + 1)]

The array is a Numpy array, though that's not really necessary. I just found it very easy to read the data in with a one-liner using numpy.fromfile().

Is there any more Pythonic way to iterate over these than the following? That simply looks like C using Python syntax.

for x in range(x_dimension):
    for y in range(y_dimension):
        for z in range(z_dimension):
            work_with_cube(array[(x + 0, y + 0, z + 0)],
                           array[(x + 1, y + 0, z + 0)],
                           array[(x + 0, y + 1, z + 0)],
                           array[(x + 1, y + 1, z + 0)],
                           array[(x + 0, y + 0, z + 1)],
                           array[(x + 1, y + 0, z + 1)],
                           array[(x + 0, y + 1, z + 1)],
                           array[(x + 1, y + 1, z + 1)])

Will this be answered by http://stackoverflow.com/questions/1280667/in-python-is-there-an-easier-way-to-write-6-nested-for-loops ? — tom10, Aug 22 '09 at 14:29

Otto Allmendinger · Accepted Answer · 2009-08-22T16:14:33.053

19

Have a look at itertools, especially itertools.product. You can compress the three loops into one with

import itertools

for x, y, z in itertools.product(*map(xrange, (x_dim, y_dim, z_dim)):
    ...

You can also create the cube this way:

cube = numpy.array(list(itertools.product((0,1), (0,1), (0,1))))
print cube
array([[0, 0, 0],
       [0, 0, 1],
       [0, 1, 0],
       [0, 1, 1],
       [1, 0, 0],
       [1, 0, 1],
       [1, 1, 0],
       [1, 1, 1]])

and add the offsets by a simple addition

print cube + (10,100,1000)
array([[  10,  100, 1000],
       [  10,  100, 1001],
       [  10,  101, 1000],
       [  10,  101, 1001],
       [  11,  100, 1000],
       [  11,  100, 1001],
       [  11,  101, 1000],
       [  11,  101, 1001]])

which would to translate to cube + (x,y,z) in your case. The very compact version of your code would be

import itertools, numpy

cube = numpy.array(list(itertools.product((0,1), (0,1), (0,1))))

x_dim = y_dim = z_dim = 10

for offset in itertools.product(*map(xrange, (x_dim, y_dim, z_dim))):
    work_with_cube(cube+offset)

Edit: itertools.product makes the product over the different arguments, i.e. itertools.product(a,b,c), so I have to pass map(xrange, ...) with as *map(...)

edited Aug 22 '09 at 16:14

answered Aug 22 '09 at 14:38

Otto Allmendinger

27,448
7
68
79

This results in the error: ValueError: shape mismatch: objects cannot be broadcast to a single shape – Nathan Fellman Aug 22 '09 at 15:41
...However, using `(x,y,z)` instead of `offset` in your example fixes that – Nathan Fellman Aug 22 '09 at 15:52
1

What does the `*` in `*map(...)` do? – Nathan Fellman Aug 23 '09 at 17:44
@Nathan: it calls the function with a list as the arguments. `s=[1,2,3]; foo(*s)` is the same as `foo(1,2,3)` – Otto Allmendinger Aug 23 '09 at 20:55
as opposed to `foo(s)` which would pass `foo` only one argument, the list `s`? – Nathan Fellman Oct 04 '09 at 21:20
I see a problem with your suggested solution. X is my LSB. In the original question you can see that it flips at each line, while in your suggested example it's the MSB, since it flips only once, in the middle of the list. How can I generate the list `(0,0,0), (1,0,0), (0,1,0),(1,1,0) etc...`? – Nathan Fellman Oct 04 '09 at 21:23
reverse every vector with `[v[::-1] for v in itertools.product((0,1), (0,1), (0,1))]` – Otto Allmendinger Oct 05 '09 at 09:19

score 8 · Answer 2 · answered Aug 22 '09 at 14:29

8

import itertools
for x, y, z in itertools.product(xrange(x_size), 
                                 xrange(y_size), 
                                 xrange(z_size)):
    work_with_cube(array[x, y, z])

answered Aug 22 '09 at 14:29

nosklo

217,122
57
293
297

Pythonic way of iterating over 3D array

2 Answers2

Linked