How to set iteration depth when interating numpy array with e.g. nditer?

Question

supose I have the following array, representing the structure of an rgb image:

[[[ 0,  1,  2], [ 3,  4,  5]]
 [[ 6,  7,  8], [ 9, 10, 11]]
 [[12, 13, 14], [15, 16, 17]]]

How can I iterate over the pixels, eg [0, 1, 2] then [3, 4, 5], and receive the coresponding index? With numpys nditer function I can't define a depth/axis wher it should stop, so it would iterate over each value, eg 0 then 1 then 2 and so on.

Is ther a numpy methode where I can define the iteration depth?

Answer 1

If I understood your question correctly you could just use a simple nested loop

A = np.array([[[ 0,  1,  2], [ 3,  4,  5]],
              [[ 6,  7,  8], [ 9, 10, 11]],
              [[12, 13, 14], [15, 16, 17]]])

for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        print(i, j, A[i,j,...])

0 0 [0 1 2]
0 1 [3 4 5]
1 0 [6 7 8]
1 1 [ 9 10 11]
2 0 [12 13 14]
2 1 [15 16 17]

Or, in a more numpythonic way with np.ndindex

for i in np.ndindex(A.shape[:2]):
    print(i, A[i])

(0, 0) [0 1 2]
(0, 1) [3 4 5]
(1, 0) [6 7 8]
(1, 1) [ 9 10 11]
(2, 0) [12 13 14]
(2, 1) [15 16 17]