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

akuiper · Accepted Answer · 2017-09-28 02:30:33Z

7

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

edited Sep 28, 2017 at 2:30

answered Sep 28, 2017 at 1:53

akuiper

216k33 gold badges362 silver badges379 bronze badges

Sign up to request clarification or add additional context in comments.

4 Comments

stormont Over a year ago

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 Over a year ago

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 Over a year ago

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 Over a year ago

How about if one needs the value as a string? like arr[y, x] = "y_x"

DJK · Accepted Answer · 2017-09-28 02:31:01Z

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

Bora · Accepted Answer · 2017-09-28 01:54:42Z

1

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

answered Sep 28, 2017 at 1:54

Bora

1,53216 silver badges20 bronze badges

Comments

Devin Cairns · Accepted Answer · 2017-09-28 01:59:38Z

1

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

answered Sep 28, 2017 at 1:59

Devin Cairns

7006 silver badges10 bronze badges

1 Comment

Devin Cairns Over a year ago

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

Mad Physicist · Accepted Answer · 2017-09-28 02:06:43Z

1

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)

edited Sep 28, 2017 at 2:06

answered Sep 28, 2017 at 1:55

Mad Physicist

116k29 gold badges202 silver badges292 bronze badges

Collectives™ on Stack Overflow

Fill 2-D numpy array with index location

5 Answers 5

4 Comments

Comments

Comments

1 Comment

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

5 Answers 5

4 Comments

Comments

Comments

1 Comment

Comments

Your Answer

Sign up or log in

Post as a guest

Related