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 🙂