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 :-)

Answer 2

Here's an initialization based method -

def create_grid(m,n):
    out = np.empty((m,n,2),dtype=int) #Improvement suggested by @AndrasDeak
    out[...,0] = np.arange(m)[:,None]
    out[...,1] = np.arange(n)
    return out

Sample run -

In [47]: create_grid(4,5)
Out[47]: 
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]]])

Runtime test for all approaches posted thus far on (4,5) grided and bigger sizes -

In [111]: %timeit np.moveaxis(np.indices((4,5)), 0, -1)
     ...: %timeit np.mgrid[:4, :5].swapaxes(2, 0).swapaxes(0, 1)
     ...: %timeit np.mgrid[:4,:5].transpose(1,2,0)
     ...: %timeit create_grid(4,5)
     ...: 
100000 loops, best of 3: 11.1 µs per loop
100000 loops, best of 3: 17.1 µs per loop
100000 loops, best of 3: 17 µs per loop
100000 loops, best of 3: 2.51 µs per loop

In [113]: %timeit np.moveaxis(np.indices((400,500)), 0, -1)
     ...: %timeit np.mgrid[:400, :500].swapaxes(2, 0).swapaxes(0, 1)
     ...: %timeit np.mgrid[:400,:500].transpose(1,2,0)
     ...: %timeit create_grid(400,500)
     ...: 
1000 loops, best of 3: 351 µs per loop
1000 loops, best of 3: 1.01 ms per loop
1000 loops, best of 3: 1.03 ms per loop
10000 loops, best of 3: 190 µs per loop

Answer 3

You can abuse numpy.mgrid or meshgrid for this purpose:

>>> import numpy as np
>>> np.mgrid[:4,:5].transpose(1,2,0)
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 4

Do you do this because you need it or just for sport? In the former case:

np.moveaxis(np.indices((4,5)), 0, -1)

np.indices does precisely what its name suggests. Only it arranges axes differently to you. So we move them with moveaxis

As @Eric points out one attractive feature of this method is that it works unmodified at arbitrary number of dimensions:

dims = tuple(np.multiply.reduceat(np.zeros(16,int)+2, np.r_[0, np.sort(np.random.choice(16, np.random.randint(10)))]))
# len(dims) == ?
np.moveaxis(np.indices(dims), 0, -1) # works