Fastest way to create numpy 2d array of indices
I want to create a numpy 2d array containing cell indices, for example such a 2x2 math can be created using:
np.array([[[0,0],[0,1]],[[1,0],[1,1]]])
In other words, the cell with the index i,j
must contain a list [i,j]
.
I could do a nested loop to do this, but I'm wondering if you have a quick pythonic way to do this?
source to share
For performance with NumPy, I would suggest an array initialization approach -
def indices_array(n):
r = np.arange(n)
out = np.empty((n,n,2),dtype=int)
out[:,:,0] = r[:,None]
out[:,:,1] = r
return out
For general (m,n,2)
-shaped output, we need some modifications:
def indices_array_generic(m,n):
r0 = np.arange(m) # Or r0,r1 = np.ogrid[:m,:n], out[:,:,0] = r0
r1 = np.arange(n)
out = np.empty((m,n,2),dtype=int)
out[:,:,0] = r0[:,None]
out[:,:,1] = r1
return out
Note. Also, read the 2019 addendum in this post. increase with large m
, n
.
Trial run -
In [145]: n = 3
In [146]: indices_array(n)
Out[146]:
array([[[0, 0],
[0, 1],
[0, 2]],
[[1, 0],
[1, 1],
[1, 2]],
[[2, 0],
[2, 1],
[2, 2]]])
If you want an 2D
array of 2
columns, just change the shape -
In [147]: indices_array(n).reshape(-1,2)
Out[147]:
array([[0, 0],
[0, 1],
[0, 2],
[1, 0],
[1, 1],
[1, 2],
[2, 0],
[2, 1],
[2, 2]])
Timing and verification -
In [141]: n = 100
...: out1 = np.array(list(product(range(n), repeat=2))).reshape(n,n,2)
...: out2 = indices_array(n)
...: print np.allclose(out1, out2)
...:
True
# @Ofek Ron solution
In [26]: %timeit np.array(list(product(range(n), repeat=2))).reshape(n,n,2)
100 loops, best of 3: 2.69 ms per loop
In [27]: # @Brad Solomon soln
...: def ndindex_app(n):
...: row, col = n,n
...: return np.array(list(np.ndindex((row, col)))).reshape(row, col, 2)
...:
# @Brad Solomon soln
In [28]: %timeit ndindex_app(n)
100 loops, best of 3: 5.72 ms per loop
# Proposed earlier in this post
In [29]: %timeit indices_array(n)
100000 loops, best of 3: 12.1 µs per loop
In [30]: 2690/12.1
Out[30]: 222.31404958677686
200x+
acceleration for n=100
with initialization!
Application for 2019
We can also use - np.indices
def indices_array_generic_builtin(m,n):
return np.indices((m,n)).transpose(1,2,0)
Timing -
In [115]: %timeit indices_array_generic(1000,1000)
...: %timeit indices_array_generic_builtin(1000,1000)
100 loops, best of 3: 2.92 ms per loop
1000 loops, best of 3: 1.37 ms per loop
source to share
You want np.ndindex
.
def coords(row, col):
return np.array(list(np.ndindex((row, col)))).reshape(row, col, 2)
coords(3, 2)
Out[32]:
array([[[0, 0],
[0, 1]],
[[1, 0],
[1, 1]],
[[2, 0],
[2, 1]]])
source to share