Search code examples
pythonnumpyiterable-unpacking

Unpack `np.unravel_index()` in for loop


I'm trying to directly unpack the tuple returned from np.unravel_index() in a for loop definition, but I'm running into the following issue:

for i,j in np.unravel_index(range(len(10)), (2, 5)):
    print(i,j)

returns:

ValueError: too many values to unpack (expected 2)

I can solve the issue by doing:

idx = np.unravel_index(range(len(10)), (2, 5))
for i, j in zip(idx[0], idx[1]):
    print(i, j)

but I really feel I should be able to do it all in the for loop assignment.

I have looked through StackOverflow and found nothing that could help me with my specific question.

Solution: As the accepted solution I think that this does exactly what I want i.e., unpack directly in the for loop assignment and without previous knowledge of the dimensions of idx:

for i, j in zip(*np.unravel_index(range(len(10)), (2, 5)):
    print(i, j)

Solution

  • Your idx is a tuple of arrays:

    In [559]: idx = np.unravel_index(np.arange(5),(2,5))
    In [560]: idx
    Out[560]: (array([0, 0, 0, 0, 0]), array([0, 1, 2, 3, 4]))
    

    The tuple works great for indexing, e.g. data[idx]. In fact that's its intended purpose.

    Your zip turns that into a list/iteration on 2 element tuples:

    In [561]: list(zip(*idx))
    Out[561]: [(0, 0), (0, 1), (0, 2), (0, 3), (0, 4)]
    

    np.transpose can also turn it into a (n,2) array:

    In [562]: np.transpose(idx)
    Out[562]: 
    array([[0, 0],
           [0, 1],
           [0, 2],
           [0, 3],
           [0, 4]])
    

    Iteration on [562] will be just as slow as the iteration on the zip, possibly slower. But if you don't need to iterate, [562] may be better.

    Notice I used zip(*idx) above, so your expression could written as:

    for i, j in zip(*np.unravel_index(range(len(neuron_sample)), (2, 5))):
        print(i, j)