Python itertools.product reorder the generation

Question

I have this:

shape = (2, 4) # arbitrary, could be 3 dimensions such as (3, 5, 7), etc...

for i in itertools.product(*(range(x) for x in shape)):
    print(i)

# output: (0, 0) (0, 1) (0, 2) (0, 3) (1, 0) (1, 1) (1, 2) (1, 3)

So far, so good, itertools.product advances the rightmost element on every iteration. But now I want to be able to specify the iteration order according to the following:

axes = (0, 1) # normal order
# output: (0, 0) (0, 1) (0, 2) (0, 3) (1, 0) (1, 1) (1, 2) (1, 3)

axes = (1, 0) # reversed order
# output: (0, 0) (1, 0) (2, 0) (3, 0) (0, 1) (1, 1) (2, 1) (3, 1)

If shapes had three dimensions, axes could have been for instance (0, 1, 2) or (2, 0, 1) etc, so it's not a matter of simply using reversed(). So I wrote some code that does that but seems very inefficient:

axes = (1, 0)

# transposed axes
tpaxes = [0]*len(axes)
for i in range(len(axes)):
    tpaxes[axes[i]] = i

for i in itertools.product(*(range(x) for x in shape)):
    # reorder the output of itertools.product
    x = (i[y] for y in tpaxes)
    print(tuple(x))

Any ideas on how to properly do this?

for your examples `tpaxes` is `[1, 0]` and `axes` is `(1, 0)`. You might want to change your example data for clarity so they are different :) — hochl, Mar 28 '12 at 08:25
True, axes=tpaxes because that's the only possible way to reorder the axes of a 2d matrix. For a 3d matrix, that's not the case. If axes was `(2, 0, 1)` then tpaxes would be `(1, 2, 0)` for example. — Giovanni Funchal, Mar 28 '12 at 08:28
I know -- just wanted to point out that a more complicated example would be better in this case; No offence. — hochl, Mar 28 '12 at 08:37
The only way to do this without an extra step is to write your own implementation of `product`. I link to a couple you could start with in [this post about `itertools.product`](http://stackoverflow.com/questions/9857608/how-can-i-find-all-the-combinations-of-elements-in-a-set-of-arrays/9857635#9857635) from the other day. My question is __why__. If you really needed to do this in some specific situation, you just reorder the arguments you fed to product into the correct order to begin with, and you wouldn't need to change the order of the generation. — agf, Mar 30 '12 at 10:10
Sorry if I'm not completely following your logic, but why don't you first reorder the ranges you're passing to `product`? Do you want to change the axes order (as it is in the example) or rather the iteration order (then the 2nd output would be `(0, 0), (1, 0), (0, 1), (1, 1), (0, 2), (1, 2), (0, 3), (1, 3)`) — bereal, Apr 01 '12 at 09:31
Ah, now I see, you want to change both, coords order and iteration order... — bereal, Apr 01 '12 at 09:39

bereal · Accepted Answer · 2012-04-06T16:21:28.933

Well, this is in fact a manual specialised product. It should be faster since axes are reordered only once:

def gen_chain(dest, size, idx, parent):
    # iterate over the axis once
    # then trigger the previous dimension to update
    # until everything is exhausted
    while True:
        if parent: next(parent) # StopIterator is propagated upwards

        for i in xrange(size):
            dest[idx] = i
            yield 

        if not parent: break

def prod(shape, axes):
    buf = [0] * len(shape)
    gen = None

    # EDIT: fixed the axes order to be compliant with the example in OP 
    for s, a in zip(shape, axes):
        # iterate over the axis and put to transposed
        gen = gen_chain(buf, s, a, gen)

    for _ in gen:
        yield tuple(buf)


print list(prod((2,4), (0,1)))
# [(0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2), (1, 3)]
print list(prod((2,4), (1,0)))
# [(0, 0), (1, 0), (2, 0), (3, 0), (0, 1), (1, 1), (2, 1), (3, 1)]
print list(prod((4,3,2),(1,2,0)))
# [(0, 0, 0), (1, 0, 0), (0, 0, 1), (1, 0, 1), (0, 0, 2), (1, 0, 2), ...

Reinstate Monica · Answer 2 · 2012-04-05T22:40:47.757

If you can afford it memory-wise: Let itertools.product do the hard work, and use zip to switch the axes around.

import itertools
def product(shape, axes):
    prod_trans = tuple(zip(*itertools.product(*(range(shape[axis]) for axis in axes))))

    prod_trans_ordered = [None] * len(axes)
    for i, axis in enumerate(axes):
        prod_trans_ordered[axis] = prod_trans[i]
    return zip(*prod_trans_ordered)

Small test:

>>> print(*product((2, 2, 4), (1, 2, 0)))
(0, 0, 0) (1, 0, 0) (0, 0, 1) (1, 0, 1) (0, 0, 2) (1, 0, 2) (0, 0, 3) (1, 0, 3) (0, 1, 0) (1, 1, 0) (0, 1, 1) (1, 1, 1) (0, 1, 2) (1, 1, 2) (0, 1, 3) (1, 1, 3)

The above version is fast if there are not too may products. For large result sets, the following is faster, but... uses eval (although in a rather safe way):

def product(shape, axes):
    d = dict(("r%i" % axis, range(shape[axis])) for axis in axes)
    text_tuple = "".join("x%i, " % i for i in range(len(axes)))
    text_for = " ".join("for x%i in r%i" % (axis, axis) for axis in axes)
    return eval("((%s) %s)" % (text_tuple, text_for), d)

Edit: If you want to not only change the order of iteration, but also the shape (as in the OP's example), small changes are needed:

import itertools
def product(shape, axes):
    prod_trans = tuple(zip(*itertools.product(*(range(s) for s in shape))))

    prod_trans_ordered = [None] * len(axes)
    for i, axis in enumerate(axes):
        prod_trans_ordered[axis] = prod_trans[i]
    return zip(*prod_trans_ordered)

And the eval version:

def product(shape, axes):
    d = dict(("r%i" % axis, range(s)) for axis, s in zip(axes, shape))
    text_tuple = "".join("x%i, " % i for i in range(len(axes)))
    text_for = " ".join("for x%i in r%i" % (axis, axis) for axis in axes)
    return eval("((%s) %s)" % (text_tuple, text_for), d)

Test:

>>> print(*product((2, 2, 4), (1, 2, 0)))
(0, 0, 0) (1, 0, 0) (2, 0, 0) (3, 0, 0) (0, 0, 1) (1, 0, 1) (2, 0, 1) (3, 0, 1) (0, 1, 0) (1, 1, 0) (2, 1, 0) (3, 1, 0) (0, 1, 1) (1, 1, 1) (2, 1, 1) (3, 1, 1)

+1, I especially like the `eval` solution. This is a creative piece of coding that one doesn't see every day! It might be ugly, but it's rather clear what it does and if performance really is critical, I think it's the only way :) — Niklas B., Apr 02 '12 at 02:52
Hm, where is (0, 0, 3) coming from, the last axis should be moved into middle, shouldn't it? — bereal, Apr 05 '12 at 04:51
@bereal: The values of axis 1 are increased first, then of axis 2, then of axis 0. Which values appear on which axis is not affected by the `axes` parameter; it only determines the order in which the values are increased. — Reinstate Monica, Apr 05 '12 at 16:45
@WolframH in the OP's example with pairs, the order of axes is changed. — bereal, Apr 05 '12 at 16:50
As far as I understand, the 4-axis (the last one, having index 2) should be moved to the middle in the result, since `(1, 2, 0)` places it there. So there should be `(0, 3, 0)`, not `(3, 0, 0)`. — bereal, Apr 06 '12 at 12:36
@bereal: I checked my output against the OP's code and it was OK. Do you get `(0,3,0)` or `(3,0,0)` when executing the OP's code? — Reinstate Monica, Apr 06 '12 at 15:49
@WolframH, ah, finally I'm getting the OP's logic. 0 in the last position should mean that last axis goes to the first. Ok, thanks. — bereal, Apr 06 '12 at 16:03

score 1 · Answer 3 · answered Mar 30 '12 at 20:55

I don't know how efficient this is, but you should be able to do something like this...

shape = (2, 4, 3)
axes = (2, 0, 1)

# Needed to get the original ordering back
axes_undo = tuple(reversed(axes))

# Reorder the shape in a configuration so that .product will give you
# the order you want.
reordered = tuple(reversed(map(lambda x: shape[x], list(axes))))

# When printing out the results from .product, put the results back
# into the original order.
for i in itertools.product(*(range(x) for x in reordered)):
    print(tuple(map(lambda x: i[x], list(axes_undo))))

I tried is up to 4 dimensions and it seems to work. ;)

I'm just swapping the dimensions around and then swapping them back.

score 1 · Answer 4 · answered Mar 31 '12 at 01:59

Have you tried timing to see how much longer it takes? What you have shouldn't be much slower than without reordering.

You could try modify what you have to use in-place splice assignment.

tpaxes = tuple(tpaxes)
for i in itertools.product(*(range(x) for x in shape)):
    # reorder the output of itertools.product
    i[:] = (i[y] for y in tpaxes)
    print(tuple(x))

Also you could get a speedup by making tpaxes a local variable of a function rather than a global variable (which has slower lookup times)

Otherwise my suggestion is somehow write your own product function..

score 0 · Answer 5 · answered Mar 28 '12 at 08:11

0

for i in itertools.product(*(range(x) for x in reversed(shape))):
    print tuple(reversed(i))

answered Mar 28 '12 at 08:11

Janne Karila

24,266
6
53
94

That works in the 2D case, but I'd like to be able to specify the how the axes are reversed. I have updated the question to reflect that more clearly. – Giovanni Funchal Mar 28 '12 at 08:15
@GiovanniFunchal Then I think your approach is fine: `tuple(i[y] for y in tpaxes)` – Janne Karila Mar 28 '12 at 08:20

score 0 · Answer 6 · answered Apr 02 '12 at 00:06

import itertools

normal = (0, 1)
reverse = (1, 0)

def axes_ordering(x):
    a, b = x
    return b - a

shape = (2, 4)

for each in itertools.product(*(range(x) for x in shape)):
    print(each[::axes_ordering(normal)], each[::axes_ordering(reverse)])

result:

(0, 0) (0, 0)
(0, 1) (1, 0)
(0, 2) (2, 0)
(0, 3) (3, 0)
(1, 0) (0, 1)
(1, 1) (1, 1)
(1, 2) (2, 1)
(1, 3) (3, 1)

Python itertools.product reorder the generation

6 Answers6

Linked