Generate NumPy array containing the indices of another NumPy array

Question

I'd like to generate a np.ndarray NumPy array for a given shape of another NumPy array. The former array should contain the corresponding indices for each cell of the latter array.

Example 1

Let's say we have a = np.ones((3,)) which has a shape of (3,). I'd expect

[[0]
 [1]
 [2]]

since there is a[0], a[1] and a[2] in a which can be accessed by their indices 0, 1 and 2.

Example 2

For a shape of (3, 2) like b = np.ones((3, 2)) there is already very much to write. I'd expect

[[[0 0]
  [0 1]]

 [[1 0]
  [1 1]]

 [[2 0]
  [2 1]]]

since there are 6 cells in b which can be accessed by the corresponding indices b[0][0], b[0][1] for the first row, b[1][0], b[1][1] for the second row and b[2][0], b[2][1] for the third row. Therefore we get [0 0], [0 1], [1 0], [1 1], [2 0] and [2 1] at the matching positions in the generated array.

Thank you very much for taking the time. Let me know if I can clarify the question in any way.

Answer 1

One way to do it with np.indices and np.stack:

np.stack(np.indices((3,)), -1)

#array([[0],
#       [1],
#       [2]])

np.stack(np.indices((3,2)), -1)

#array([[[0, 0],
#        [0, 1]],
#       [[1, 0],
#        [1, 1]],
#       [[2, 0],
#        [2, 1]]])

np.indices returns an array of index grid where each subarray represents an axis:

np.indices((3, 2))

#array([[[0, 0],
#        [1, 1],
#        [2, 2]],        
#       [[0, 1],
#        [0, 1],
#        [0, 1]]])

Then transpose the array with np.stack, stacking index for each element from different axis:

np.stack(np.indices((3,2)), -1)

#array([[[0, 0],
#        [0, 1]],
#       [[1, 0],
#        [1, 1]],
#       [[2, 0],
#        [2, 1]]])