Iterate over inner axes of an array

I want to iterate over some inner dimensions of an array without knowing in advance how many dimensions to iterate over. Furthermore I only know that the last two dimensions should not be iterated over.

For example assume the array has dimension 5 and shape (i,j,k,l,m) and I want to iterate over the second and third dimension. In each step of the iteration I expect an array of shape (i,l,m).

Example 1: Iterate over second dimension of x with x.ndim=4, x.shape=(I,J,K,L). Expected: x[:,j,:,:] for j=0,...,J-1

Example 2: Iterate over second and third dimension of x with x.ndim=5, x.shape=(I,J,K,L,M) Expected: x[:,j,k,:,:] for j=0,...,J-1, k=0,...,K-1

Example 3: Iterate over second, third and fourth dimension of x with x.ndim=6, x.shape=(I,J,K,L,M,N) Expected: x[:,j,k,l,:,:] for j=0,...,J-1, k=0,...,K-1 and l=0,...,L-1

Assume the array has dimension 5 and shape (i,j,k,l,m).

If I know which dimension to iterate over, for example the second and third axis, this is possible with a nested for-loop:

for j in range(x.shape[1]):
    for k in range(x.shape[2]):
        x[...,j,k,:,:] 

However since I do not know in advance how many dimensions I want to iterate over for-loops are not an option. I found a way to generate the indices based on the shapes of the dimensions I want to iterate over.

for b in product(*map(range, x.shape[2:4])):
    print(b)   
>>> (0, 0)
>>> (0, 1)
>>> ...
>>> (0, k)
>>> ...
>>> (j, k)

This yields the indices for arbitrary inner dimension which is what I want. However I’m not aware of a way to use this tuple directly to slice into an an array. Therefore I first need to assign these entries to variables and then use these variables for slicing.

for b in product(*map(range,x.shape[2:4])):
    j,k=b
    x[...,j,k,:,:]

But this approach again only works if I know in advance how many dimensions to iterate over.

Answer

With the caveat that I don’t fully understand how you want to use this, I reckon the following should do what you are asking for:

from itertools import product

def iterover(x, axes):
    x = np.moveaxis(x, axes, np.arange(len(axes)))
    subshape = x.shape[:len(axes)]
    ixi = product(*[range(k) for k in subshape])
    for ix in ixi:
        yield ix, x[ix]

Example:

def genrange(shape):
    return np.arange(np.product(shape)).reshape(shape)

x = genrange((2,3,4,5))
>>> x.shape
(2, 3, 4, 5)

>>> {ix: xx.shape for ix, xx in iterover(x, (1,2))}
{(0, 0): (2, 5),
 (0, 1): (2, 5),
 (0, 2): (2, 5),
 (0, 3): (2, 5),
 (1, 0): (2, 5),
 (1, 1): (2, 5),
 (1, 2): (2, 5),
 (1, 3): (2, 5),
 (2, 0): (2, 5),
 (2, 1): (2, 5),
 (2, 2): (2, 5),
 (2, 3): (2, 5)}

Notes:

1. The axes to iterate over may be in any order, e.g.:

   >>> {ix: xx.shape for ix, xx in iterover(x, (2,0))}
   {(0, 0): (3, 5),
    (0, 1): (3, 5),
    (1, 0): (3, 5),
    (1, 1): (3, 5),
    (2, 0): (3, 5),
    (2, 1): (3, 5),
    (3, 0): (3, 5),
    (3, 1): (3, 5)}
   

2. The sub arrays can be used as L-values (modified in place), with the original array modified accordingly. This is due to np.moveaxis() returning a view of the original array (numpy never ceases to amaze me):

   for ix, xx in iterover(x, (2,0)):
       if ix == (0,0):
           xx *= -1
   
   >>> x
   array([[[[  0,  -1,  -2,  -3,  -4],
            [  5,   6,   7,   8,   9],
            [ 10,  11,  12,  13,  14],
            [ 15,  16,  17,  18,  19]],
   
           [[-20, -21, -22, -23, -24],
            [ 25,  26,  27,  28,  29],
            [ 30,  31,  32,  33,  34],
            [ 35,  36,  37,  38,  39]],
   ...