Fill 2-D numpy array with index location

Question

I've been trying to figure out a clean, pythonic way to fill each element of an empty numpy array with the index value(s) of that element, without using for loops. For 1-D, it's easy, you can just use something like np.arange or just a basic range. But at 2-D and higher dimensions, I'm stumped on how to easily do this.

(Edit: Or just build a regular list like this, then np.array(lst) it. I think I just answered my question - use a list comprehension?)

Example:

rows = 4
cols = 4
arr = np.empty((rows, cols, 2))  # 4x4 matrix with [x,y] location

for y in range(rows):
    for x in range(cols):
        arr[y, x] = [y, x]

'''
Expected output:
[[[0,0], [0,1], [0,2], [0,3]],
 [[1,0], [1,1], [1,2], [1,3]],
 [[2,0], [2,1], [2,2], [2,3]],
 [[3,0], [3,1], [3,2], [3,3]]]
'''

Answer 1

What you are showing is a meshgrid of a 4X4 matrix; You can either use np.mgrid, then transpose the result:

np.moveaxis(np.mgrid[:rows,:cols], 0, -1)
#array([[[0, 0],
#        [0, 1],
#        [0, 2],
#        [0, 3]],

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

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

#       [[3, 0],
#        [3, 1],
#        [3, 2],
#        [3, 3]]])

Or use np.meshgrid with matrix indexing ij:

np.dstack(np.meshgrid(np.arange(rows), np.arange(cols), indexing='ij'))
#array([[[0, 0],
#        [0, 1],
#        [0, 2],
#        [0, 3]],

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

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

#       [[3, 0],
#        [3, 1],
#        [3, 2],
#        [3, 3]]])

Answer 2

another way using np.indices and concatenate

 np.concatenate([x.reshape(4,4,1) for x in np.indices((4,4))],2)

or with np.dstack

np.dstack(np.indices((4,4)))

---

Some bench marking since you have a ton of possibilities

def Psidom_mrgid(rows,cols):
    np.mgrid[:rows, :cols].transpose((1, 2, 0))

def Psidom_mesh(rows,cols):
    np.dstack(np.meshgrid(np.arange(rows), np.arange(cols), indexing='ij'))

def Mad_tile(rows,cols):
    r = np.tile(np.arange(rows).reshape(rows, 1), (1, cols))
    c = np.tile(np.arange(cols), (rows, 1))
    result = np.stack((r, c), axis=-1)

def bora_comp(rows,cols):
    x = [[[i, j] for j in range(rows)] for i in range(cols)]

def djk_ind(rows,cols):
    np.concatenate([x.reshape(rows, cols, 1) for x in np.indices((rows, cols))], 2)

def devdev_mgrid(rows,cols):
    index_tuple = np.mgrid[0:rows, 0:cols]
    np.dstack(index_tuple).reshape((rows, cols, 2)


In[8]: %timeit Psidom_mrgid(1000,1000)
100 loops, best of 3: 15 ms per loop

In[9]: %timeit Psidom_mesh(1000,1000)
100 loops, best of 3: 9.98 ms per loop

In[10]: %timeit Mad_tile(1000,1000)
100 loops, best of 3: 15.3 ms per loop

In[11]: %timeit bora_comp(1000,1000)
1 loop, best of 3: 221 ms per loop

In[12]: %timeit djk_ind(1000,1000)
100 loops, best of 3: 9.72 ms per loop

In[13]: %timeit devdev_mgrid(1000,1000)
10 loops, best of 3: 20.6 ms per loop

Answer 3

I guess that's pretty pythonic:

[[[i,j] for j in range(5)] for i in range(5)]

Output:

[[[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]],

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

Answer 4

Check out numpy.mgrid, which will return two arrays with the i and j indices. To combine them you can stack the arrays and reshape them. Something like this:

import numpy as np

def index_pair_array(rows, cols):
    index_tuple = np.mgrid[0:rows, 0:cols]
    return np.dstack(index_tuple).reshape((rows, cols, 2))

Answer 5

There are a few ways of doing this numpythonically.

One way is using np.tile and np.stack:

r = np.tile(np.arange(rows).reshape(rows, 1), (1, cols))
c = np.tile(np.arange(cols), (rows, 1))
result = np.stack((r, c), axis=-1)

A better way of getting the coordinates might be np.meshgrid:

rc = np.meshgrid(np.arange(rows), np.arange(cols), indexing='ij')
result = np.stack(rc, axis=-1)