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 stackoverflow.com/questions/1280667/… ? — tom10
– tom10, Commented Aug 22, 2009 at 14:29

Otto Allmendinger · Accepted Answer · 2009-08-22 16:14:33Z

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, 2009 at 16:14

answered Aug 22, 2009 at 14:38

Otto Allmendinger

28.5k7 gold badges71 silver badges79 bronze badges

Sign up to request clarification or add additional context in comments.

7 Comments

Nathan Fellman Over a year ago

This results in the error: ValueError: shape mismatch: objects cannot be broadcast to a single shape

Nathan Fellman Over a year ago

...However, using (x,y,z) instead of offset in your example fixes that

Nathan Fellman Over a year ago

What does the * in *map(...) do?

Otto Allmendinger Over a year ago

@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)

Nathan Fellman Over a year ago

as opposed to foo(s) which would pass foo only one argument, the list s?

|

nosklo · Accepted Answer · 2009-08-22 14:29:59Z

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, 2009 at 14:29

nosklo

224k58 gold badges300 silver badges299 bronze badges

Collectives™ on Stack Overflow

Pythonic way of iterating over 3D array

2 Answers 2

7 Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

7 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related