views:

378

answers:

2

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)])
+5  A: 
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])
nosklo
+4  A: 

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(...)

Otto Allmendinger
This results in the error:ValueError: shape mismatch: objects cannot be broadcast to a single shape
Nathan Fellman
...However, using `(x,y,z)` instead of `offset` in your example fixes that
Nathan Fellman
\*sigh* always test your code before putting it up
Otto Allmendinger
The method 'product' of 'itertools' is not available in python 2.5.2...
mshsayem
What does the `*` in `*map(...)` do?
Nathan Fellman
@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
as opposed to `foo(s)` which would pass `foo` only one argument, the list `s`?
Nathan Fellman
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
reverse every vector with `[v[::-1] for v in itertools.product((0,1), (0,1), (0,1))]`
Otto Allmendinger