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.
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]]])