Iterate over last axis of a numpy array

Question

Let's say we have a (20, 5) array. We can iterate over each row very pythonically:

import numpy as np
xs = np.array(range(100)).reshape(20, 5)
for x in xs:
    print(x)

If we want to iterate over another axis (here in the example, iterate over columns, but I'm looking for a solution for each possible axis in a ndarray), it's less direct, we can use the method from Iterating over arbitrary dimension of numpy.array:

for i in range(xs.shape[-1]):
    x = xs[..., i]
    print(x)

Is there a more direct way to iterate over another axis, like (pseudo-code):

for x in xs.iterator(axis=-1):
    print(x)

?

you could could transpose that axis to the front. but I think the index approach is best — hpaulj, Jun 25 '21 at 10:06
@hpaulj I'm looking for a solution for a n-dimensional array, and let's say we want to iterate over axis 3 among 4. — Basj, Jun 25 '21 at 10:55
@basj transposition works for n-dimentional arrays. Note that (lazy default) transposition as well as direct indexing will be innefficient on big arrays. An eager transposition could be a good idea if you plan to perform many operation like that on the target arrays. — Jérôme Richard, Jun 25 '21 at 11:30
Why not [the second answer in the link you shared](https://stackoverflow.com/a/5923332/913098)? `np.rollaxis` — Gulzar, Jun 25 '21 at 11:38
@Gulzar Isn't this taking new memory to perform the operation? — Basj, Jun 25 '21 at 11:46
Use `np.take` to identify the indexing axis. Another approach is to construct a indexing tuple that includes the necessay `slice` objects. — hpaulj, Jun 25 '21 at 12:32
@Basj according to the answer it creates a view, meaning no extra memory. At least that's what I understand. However, the doc does not state it, so I am not sure. — Gulzar, Jun 25 '21 at 12:43
`transpose` takes an argument. Swap and roll can do the same thing. The link discusses most of the options. — hpaulj, Jun 25 '21 at 12:44
`np.rollaxis`, `np.moveaxis` both use `np.transpose. Look at their code. — hpaulj, Jun 25 '21 at 14:01

score 1 · Answer 1 · answered Jun 27 '21 at 14:27

I think that as_strided from the stride tricks module should do the work here.

It creates a view into the array and not a copy (as stated by the docs).

Here is a simple demonstration of as_stided capabilities:

from numpy.lib.stride_tricks import as_strided
import numpy as np
xs = np.array(range(3 *3 * 4)).reshape(3,3, 4)
for x in xs:
    print(x)

output:

[[ 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]]

function to iterate over array specific axis:

def iterate_over_axis(arr, axis=0):
    strides = arr.strides
    strides_ = [strides[axis], *strides[0:axis], *strides[(axis+1):]]
    shape = arr.shape
    shape_ = [shape[axis], *shape[0:axis], *shape[(axis+1):]]
    return as_strided(arr,  strides=strides_, shape=shape_)

for x in iterate_over_axis(xs, axis=1):
    print(x)

output:

[[ 0  1  2  3]
 [12 13 14 15]
 [24 25 26 27]]
[[ 4  5  6  7]
 [16 17 18 19]
 [28 29 30 31]]
[[ 8  9 10 11]
 [20 21 22 23]
 [32 33 34 35]]

Iterate over last axis of a numpy array

1 Answers1