Create an array where each element stores its indices

Question

I want to create a 2d numpy array where every element is a tuple of its indices.

Example (4x5):

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


        

Answer 1

You can use numpy.mgrid and swap it's axes:

>>> # assuming a 3x3 array
>>> np.mgrid[:3, :3].swapaxes(-1, 0)
array([[[0, 0],
        [1, 0],
        [2, 0]],

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

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

That still differs a bit from your desired array so you can roll your axes:

>>> np.mgrid[:3, :3].swapaxes(2, 0).swapaxes(0, 1)
array([[[0, 0],
        [0, 1],
        [0, 2]],

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

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

---

Given that someone timed the results I also want to present a manual numba based version that "beats 'em all":

import numba as nb
import numpy as np

@nb.njit
def _indexarr(a, b, out):
    for i in range(a):
        for j in range(b):
            out[i, j, 0] = i
            out[i, j, 1] = j
    return out

def indexarr(a, b):
    arr = np.empty([a, b, 2], dtype=int)
    return _indexarr(a, b, arr)

Timed:

a, b = 400, 500
indexarr(a, b)  # numba needs a warmup run
%timeit indexarr(a, b)                                  # 1000 loops, best of 3: 1.5 ms per loop
%timeit np.mgrid[:a, :b].swapaxes(2, 0).swapaxes(0, 1)  # 100 loops, best of 3: 7.17 ms per loop
%timeit np.mgrid[:a, :b].transpose(1,2,0)               # 100 loops, best of 3: 7.47 ms per loop
%timeit create_grid(a, b)                               # 100 loops, best of 3: 2.26 ms per loop

and on a smaller array:

a, b = 4, 5
indexarr(a, b)
%timeit indexarr(a, b)                                 # 100000 loops, best of 3: 13 µs per loop
%timeit np.mgrid[:a, :b].swapaxes(2, 0).swapaxes(0, 1) # 10000 loops, best of 3: 181 µs per loop
%timeit np.mgrid[:a, :b].transpose(1,2,0)              # 10000 loops, best of 3: 182 µs per loop
%timeit create_grid(a, b)                              # 10000 loops, best of 3: 32.3 µs per loop

As promised it "beats 'em all" in terms of performance :-)