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.

## Advertisement

## 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:

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)}

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]], ...