create a numpy array with its index

There are several numpy tools for generating a mesh, such as np.meshgrid. I’ll use mgrid since it produces one array. To make the action a bit clearer I’m using different dimensions.

In [410]: np.mgrid[0:2,0:3,0:4]
Out[410]: 
array([[[[0, 0, 0, 0],
         [0, 0, 0, 0],
         [0, 0, 0, 0]],

        [[1, 1, 1, 1],
         [1, 1, 1, 1],
         [1, 1, 1, 1]]],


       [[[0, 0, 0, 0],
         [1, 1, 1, 1],
         [2, 2, 2, 2]],

        [[0, 0, 0, 0],
         [1, 1, 1, 1],
         [2, 2, 2, 2]]],


       [[[0, 1, 2, 3],
         [0, 1, 2, 3],
         [0, 1, 2, 3]],

        [[0, 1, 2, 3],
         [0, 1, 2, 3],
         [0, 1, 2, 3]]]])
In [411]: _.shape
Out[411]: (3, 2, 3, 4)

This has put three ‘dimensions’ first; you want it to be last, so we need to do a transpose.

In [412]: np.mgrid[0:2,0:3,0:4].transpose(1,2,3,0)
Out[412]: 
array([[[[0, 0, 0],
         [0, 0, 1],
         [0, 0, 2],
         [0, 0, 3]],

        [[0, 1, 0],
         [0, 1, 1],
         [0, 1, 2],
         [0, 1, 3]],

        [[0, 2, 0],
         [0, 2, 1],
         [0, 2, 2],
         [0, 2, 3]]],


       [[[1, 0, 0],
         [1, 0, 1],
         [1, 0, 2],
         [1, 0, 3]],

        [[1, 1, 0],
         [1, 1, 1],
         [1, 1, 2],
         [1, 1, 3]],

        [[1, 2, 0],
         [1, 2, 1],
         [1, 2, 2],
         [1, 2, 3]]]])
In [413]: _.shape
Out[413]: (2, 3, 4, 3)

BUT, are you clear as to why you need such a large array? Maybe a sparse mesh would be just as useful.

In [416]: np.ogrid[0:2,0:3,0:4]
Out[416]: 
[array([[[0]],
 
        [[1]]]),
 array([[[0],
         [1],
         [2]]]),
 array([[[0, 1, 2, 3]]])]

With broadcasting these 3 arrays will work just as well as the 4d array.

Don’t skimp on the basic numpy reading.