Numpy.ndarray: iteration by specific axes of an array

Question

I have a multidimensional array. And I have to iterate via some of its axes. If I needed all axes, I could use nditer , but if I need only the specific ones, I have to do it manually:

my_array = np.arange(3 * 4 * 5).reshape((3, 4, 5))
for i in range(my_array.shape[0]):
    for j in range(my_array.shape[1]):
        print(i, j)
        # Here should be some processing of the 3rd dimension items of the (i,j)

Cannot you advice me a simpler way to do it?

Answer 1

Consider passing to a single loop, and using ndindex ( docs ):

my_array = np.arange(3 * 4 * 5).reshape((3, 4, 5))
for ij in np.ndindex(my_array.shape[:2]):
    i,j=ij
    print(i,j)