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

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]],
...
```
User contributions licensed under: CC BY-SA
10 People found this is helpful