I wanted to create this kind of array using numpy:
[[[0,0,0], [1,0,0], ..., [1919,0,0]],
[[0,1,0], [1,1,0], ..., [1919,1,0]],
...,
[[0,1019,0], [1,1019,0], ..., [1919,1019,0]]]
To which I can access via:
>>> data[25][37]
array([25, 37, 0])
I've tried to create an array this way, but it's not complete:
>>> data = np.mgrid[0:1920:1, 0:1080:1].swapaxes(0,2).swapaxes(0,1)
>>> data[25][37]
array([25, 37])
Do you have any idea how to solve this using numpy?
Approach #1 : Here's one way with np.ogrid and array-initialization -
def indices_zero_grid(m,n):
I,J = np.ogrid[:m,:n]
out = np.zeros((m,n,3), dtype=int)
out[...,0] = I
out[...,1] = J
return out
Sample run -
In [550]: out = indices_zero_grid(1920,1080)
In [551]: out[25,37]
Out[551]: array([25, 37, 0])
Approach #2 : A modification of @senderle's cartesian_product and also inspired by @unutbu's modification to it -
import functools
def indices_zero_grid_v2(m,n):
"""
Based on cartesian_product
http://stackoverflow.com/a/11146645 (@senderle)
Inspired by : https://stackoverflow.com/a/46135435 (@unutbu)
"""
shape = m,n
arrays = [np.arange(s, dtype='int') for s in shape]
broadcastable = np.ix_(*arrays)
broadcasted = np.broadcast_arrays(*broadcastable)
rows, cols = functools.reduce(np.multiply, broadcasted[0].shape), \
len(broadcasted)+1
out = np.zeros(rows * cols, dtype=int)
start, end = 0, rows
for a in broadcasted:
out[start:end] = a.reshape(-1)
start, end = end, end + rows
return out.reshape(-1,m,n).transpose(1,2,0)
Runtime test -
In [2]: %timeit indices_zero_grid(1920,1080)
100 loops, best of 3: 8.4 ms per loop
In [3]: %timeit indices_zero_grid_v2(1920,1080)
100 loops, best of 3: 8.14 ms per loop
In [50]: data = np.mgrid[:1920, :1080, :1].transpose(1,2,3,0)[..., 0, :]
In [51]: data[25][37]
Out[51]: array([25, 37, 0])
Note that data[25][37] two calls to __getitem__. With NumPy arrays, you can access the same value more efficiently (with one call to __getitem__) using data[25, 37]:
In [54]: data[25, 37]
Out[54]: array([25, 37, 0])
[...,0,:] at the end/squeeze to remove the singleton dim.
