iterate over some (but not all) dimensions of a ndarray

Question

I have a three-dimensional ndarray in python and would like to iterate over it in an element-wise manner along two of the three margins.

Put more literally, I would for example like to iterate over all (x,y) pairs but keep the z data together as an array.

As pseudocode, the expression I ultimately am after would be something along these lines

[ f(z) for z in all_xy_pairs(the_ndarray) if g(z) == True ]

I considered using 'reshape' as follows

import numpy as np
# silly example
ii=np.arange(0,3*9,1).reshape(3,3,3)
[ z for z in ii.reshape(9,-1) if z[1]>10 ]

but I would prefer an iterator to which I could pass the array margins over which to iterate (in the example above margins=[0,1]. In pseudocode, the above example would then become

[ z for z in iterate_over_margins(ii, margins=[0,1]) if z[1]>10 ]

Before I start programming this myself, isn't there such an iterator in numpy or a related package? I checked nditer but it doesn't do what I am after.

What's wrong with `z[i,:,j] for i in range(z.shape[0]) for j in range(z.shape[2])`? — user545424, Dec 05 '18 at 16:33
Thanks, this works of course. I would accept this as an answer. — user52366, Dec 05 '18 at 16:40

score 1 · Answer 1 · answered Dec 05 '18 at 16:46

You can select certain rows/columns of a numpy array by indexing along those columns, i.e. z[i,j,k]. In order to select all elements from a particular dimension you can use :. For example, to iterate over the first and last dimensions of a 3d array:

for i in range(z.shape[0]):
    for j in range(z.shape[2]):
        print(z[i,:,j])

score 0 · Answer 2 · answered Dec 05 '18 at 16:57

This answers a slightly different question, but, as you surely know, NumPy generally benefits greatly from using vectorized operations, so if your f and g can be vectorized you could also consider operating over the array containing all the iterated elements in sequence. You could do that with just some reshaping:

import numpy as np

# "Unrolls" an array along the given axes
def unroll_axis(a, axis):
    a = np.asarray(a)
    # This so it works with a single dimension or a sequence of them
    axis = np.atleast_1d(axis)
    # Put unrolled axes at the beginning
    a = np.moveaxis(a, axis, range(len(axis)))
    # Unroll
    return a.reshape((-1,) + a.shape[len(axis):])

# Example
a = np.arange(27).reshape((3, 3, 3))
print(unroll_axis(a, (0, 2)))
# [[ 0  3  6]
#  [ 1  4  7]
#  [ 2  5  8]
#  [ 9 12 15]
#  [10 13 16]
#  [11 14 17]
#  [18 21 24]
#  [19 22 25]
#  [20 23 26]]

So, if g and f are vectorized, you could just do

the_array_unrolled = unroll_axis(the_array, (0, 2))
result = f(the_array_unrolled[g(the_array_unrolled)])

This however does take some more memory, since you are making a new array with the whole sequence of elements instead of processing one at a time.

iterate over some (but not all) dimensions of a ndarray

2 Answers2

Linked