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

• stormont: This is the most concise example of np.mgrid that I've yet seen. The docs are gibberish to me and I couldn't make heads or tails out of other SO posts.
• stormont: This seems easy to extend to higher dimensions, too, which is nice. I haven't yet figured out the right args to transpose the way I'd like for 3-D, but at least I can tell it's supported.
• akuiper: To generalize, you can use np.moveaxis (which is a special transpose) instead of transpose here. It should apply to nd as well without modification of args.
• lesolorzanov: How about if one needs the value as a string? like arr[y, x] = "y_x"

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

• Devin Cairns: @Psidom 's answer is better, you don't need to stack the np.mgrid output, as it is already an ndarray object

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)