Replace each element of multidimensional array with tuple of respective indices

Question

In numpy, let's say I have an array k that has shape (2, 3, 2, 2).

k = np.array([[[[-0.08759809, -0.10987781],
                [-0.18387192, -0.2109216 ]],

               [[ 0.21027089,  0.21661097],
                [ 0.22847626,  0.23004637]],

               [[ 0.50813986,  0.54309974],
                [ 0.64082444,  0.67101435]]],


              [[[-0.98053589, -1.03143541],
                [-1.19128892, -1.24695841]],

               [[ 0.69108355,  0.66880383],
                [ 0.59480972,  0.56776003]],

               [[ 2.36270298,  2.36904306],
                [ 2.38090835,  2.38247847]]]])

How would I create a new array j of shape (2, 3, 2, 2) such that each element of j is the index of the corresponding value in k ?

Example in the first element of the first dimension and first element of the second dimension. (Corresponding to

[[[[-0.08759809, -0.10987781],
   [-0.18387192, -0.2109216 ]],

)

[[[[(0, 0, 0, 0), (0, 0, 0, 1)],
   [(0, 0, 1, 0), (0, 0, 1, 1)]],

.... and so on.

Answer 1

A simple solution would be to iterate over j and explicitly fill each element with its index j[idx, :] = idx :

k = np.round(np.random.random((4, 5)), 2)
j = np.empty(k.shape+(k.ndim,))
for idx in np.ndindex(k.shape):
    j[idx, :] = idx

# array([[[0., 0.],
#         [0., 1.],
#         [0., 2.],
#         [0., 3.],
#         [0., 4.]],
#        [[1., 0.],
#         [1., 1.],
#         [1., 2.],
#         [1., 3.],
#         [1., 4.]],
#        [[2., 0.],
#         [2., 1.],
#         [2., 2.],
#         [2., 3.],
#         [2., 4.]],
#        [[3., 0.],
#         [3., 1.],
#         [3., 2.],
#         [3., 3.],
#         [3., 4.]]])

Answer 2

The idea is to use np.ndindex as follows:

j = np.fromiter(np.ndindex(k.shape), dtype='i4,'*k.ndim).reshape(k.shape)

Result:

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

        [[(0, 1, 0, 0), (0, 1, 0, 1)],
          [(0, 1, 1, 0), (0, 1, 1, 1)]],

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


        [[[(1, 0, 0, 0), (1, 0, 0, 1)],
          [(1, 0, 1, 0), (1, 0, 1, 1)]],

        [[(1, 1, 0, 0), (1, 1, 0, 1)],
          [(1, 1, 1, 0), (1, 1, 1, 1)]],

        [[(1, 2, 0, 0), (1, 2, 0, 1)],
          [(1, 2, 1, 0), (1, 2, 1, 1)]]]],
      dtype=[('f0', '<i4'), ('f1', '<i4'), ('f2', '<i4'), ('f3', '<i4')])