Crea una matriz donde cada elemento almacene sus índices.

Quiero crear una matriz numpy 2d donde cada elemento sea una tupla de sus índices.

Ejemplo (4×5):

array([[[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]]]) 

Me gustaría crear una list python con la siguiente lista de comprensión:

 [[(y,x) for x in range(width)] for y in range(height)] 

¿Hay una forma más rápida de lograr lo mismo, tal vez con métodos numpy?

Aquí hay un método basado en la inicialización:

 def create_grid(m,n): out = np.empty((m,n,2),dtype=int) #Improvement suggested by @AndrasDeak out[...,0] = np.arange(m)[:,None] out[...,1] = np.arange(n) return out 

Ejecución de la muestra

 In [47]: create_grid(4,5) Out[47]: array([[[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]]]) 

Prueba de tiempo de ejecución para todos los enfoques publicados hasta ahora en (4,5) tamaños de grids y más grandes –

 In [111]: %timeit np.moveaxis(np.indices((4,5)), 0, -1) ...: %timeit np.mgrid[:4, :5].swapaxes(2, 0).swapaxes(0, 1) ...: %timeit np.mgrid[:4,:5].transpose(1,2,0) ...: %timeit create_grid(4,5) ...: 100000 loops, best of 3: 11.1 µs per loop 100000 loops, best of 3: 17.1 µs per loop 100000 loops, best of 3: 17 µs per loop 100000 loops, best of 3: 2.51 µs per loop In [113]: %timeit np.moveaxis(np.indices((400,500)), 0, -1) ...: %timeit np.mgrid[:400, :500].swapaxes(2, 0).swapaxes(0, 1) ...: %timeit np.mgrid[:400,:500].transpose(1,2,0) ...: %timeit create_grid(400,500) ...: 1000 loops, best of 3: 351 µs per loop 1000 loops, best of 3: 1.01 ms per loop 1000 loops, best of 3: 1.03 ms per loop 10000 loops, best of 3: 190 µs per loop 

¿Lo haces porque lo necesitas o solo por deporte? En el primer caso:

 np.moveaxis(np.indices((4,5)), 0, -1) 

np.indices hace precisamente lo que su nombre sugiere. Solo que dispone de ejes diferente a ti. Así que los movemos con moveaxis

Como @Eric señala que una característica atractiva de este método es que funciona sin modificaciones en un número arbitrario de dimensiones:

 dims = tuple(np.multiply.reduceat(np.zeros(16,int)+2, np.r_[0, np.sort(np.random.choice(16, np.random.randint(10)))])) # len(dims) == ? np.moveaxis(np.indices(dims), 0, -1) # works 

Puede abusar de numpy.mgrid o meshgrid para este propósito:

 >>> import numpy as np >>> np.mgrid[:4,:5].transpose(1,2,0) array([[[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]]]) 

Puedes usar numpy.mgrid e intercambiar sus ejes:

 >>> # assuming a 3x3 array >>> np.mgrid[:3, :3].swapaxes(-1, 0) array([[[0, 0], [1, 0], [2, 0]], [[0, 1], [1, 1], [2, 1]], [[0, 2], [1, 2], [2, 2]]]) 

Eso aún difiere un poco de la matriz deseada para que pueda girar sus ejes:

 >>> np.mgrid[:3, :3].swapaxes(2, 0).swapaxes(0, 1) array([[[0, 0], [0, 1], [0, 2]], [[1, 0], [1, 1], [1, 2]], [[2, 0], [2, 1], [2, 2]]]) 

Dado que alguien cronometró los resultados, también quiero presentar una versión basada en numba manual que “los supera a todos”:

 import numba as nb import numpy as np @nb.njit def _indexarr(a, b, out): for i in range(a): for j in range(b): out[i, j, 0] = i out[i, j, 1] = j return out def indexarr(a, b): arr = np.empty([a, b, 2], dtype=int) return _indexarr(a, b, arr) 

Cronometrado

 a, b = 400, 500 indexarr(a, b) # numba needs a warmup run %timeit indexarr(a, b) # 1000 loops, best of 3: 1.5 ms per loop %timeit np.mgrid[:a, :b].swapaxes(2, 0).swapaxes(0, 1) # 100 loops, best of 3: 7.17 ms per loop %timeit np.mgrid[:a, :b].transpose(1,2,0) # 100 loops, best of 3: 7.47 ms per loop %timeit create_grid(a, b) # 100 loops, best of 3: 2.26 ms per loop 

y en una matriz más pequeña:

 a, b = 4, 5 indexarr(a, b) %timeit indexarr(a, b) # 100000 loops, best of 3: 13 µs per loop %timeit np.mgrid[:a, :b].swapaxes(2, 0).swapaxes(0, 1) # 10000 loops, best of 3: 181 µs per loop %timeit np.mgrid[:a, :b].transpose(1,2,0) # 10000 loops, best of 3: 182 µs per loop %timeit create_grid(a, b) # 10000 loops, best of 3: 32.3 µs per loop 

Como se prometió, “supera a todos” en términos de rendimiento 🙂