0

I'd like to do something like this:

    if dim==2:
        a,b=grid_shape
        for i in range(a):
            for j in range(b):
                A[i,j] = ...things...

where dim is simply the number of elements in my tuple grid_shape. A is a numpy array of dimension dim. Is there a way to do it without being dimension specific? Without having to write ugly code like 

    if dim==2:
        a,b=grid_shape
        for i in range(a):
            for j in range(b):
                A[i,j] = ...things...
    if dim==3:
        a,b,c=grid_shape
        for i in range(a):
            for j in range(b):
                for k in range(c):
                    A[i,j,k] = ...things...
LukeMathWalker
  • 3
  • 1
  • IIUC, you're looking for [ndenumerate](http://stackoverflow.com/questions/11697274/iterate-over-numpy-matrix-of-unknown-dimension). Can you confirm? – DSM Oct 16 '16 at 13:43
  • Not exactly, even though that function will turn out to be useful in other part of the routine, so thanks anyway :D – LukeMathWalker Oct 16 '16 at 13:46

2 Answers2

0

Using itertools, you can do it like this:

for index in itertools.product(*(range(x) for x in grid_shape)):
    A[index] = ...things...

This relies on a couple of tricks. First, itertools.product() is a function which generates tuples from iterables.

for i in range(a):
    for j in range(b):
        index = i,j
        do_something_with(index)

can be reduced to

for index in itertools.product(range(a),range(b)):
    do_something_with(index)

This works for any number of arguments to itertools.product(), so you can effectively create nested loops of arbitrary depth.

The other trick is to convert your grid shape into the arguments for itertools.product:

(range(x) for x in grid_shape)

is equivalent to

(range(grid_shape[0]),range(grid_shape[1]),...)

That is, it is a tuple of ranges for each grid_shape dimension. Using * then expands this into the arguments.

itertools.product(*(range(x1),range(x2),...))

is equivalent to

itertools.product(range(x1),range(x2),...)

Also, since A[i,j,k] is equivalent to A[(i,j,k)], we can just use A[index] directly.

As DSM points out, since you are using numpy, you can reduce

itertools.product(*(for range(x) for x in grid_shape))

to

numpy.ndindex(grid_shape)

So the final loop becomes

for index in numpy.ndindex(grid_shape):
    A[index] = ...things...
Community
  • 1
  • 1
Vaughn Cato
  • 63,448
  • 5
  • 82
  • 132
0

You can catch the rest of the tuple by putting a star in front of the last variable and make a an array by putting parentheses around it.

>>> tupl = ((1, 2), 3, 4, 5, 6)
>>> a, *b = tupl
>>> a
(1, 2)
>>> b
[3, 4, 5, 6]
>>>

And then you can loop through b. So it would look something like

a,*b=grid_shape
for i in a:
    for j in range(i):
        for k in b:
            for l in range(k):
                A[j, l] = ...things...
rassar
  • 5,412
  • 3
  • 25
  • 41