Quiero crear una matriz nud de 2d donde cada elemento es una tupla de sus índices.

Ejemplo (4x5):

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

Crearía una pitón list con la siguiente lista de comprensión:

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

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

5
ovs 17 feb. 2017 a las 19:05

4 respuestas

La mejor respuesta

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 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) rejilla y tamaños 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
5
Divakar 17 feb. 2017 a las 17:46

¿Haces esto 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 sugiere su nombre. Solo que organiza los ejes de manera diferente a ti. Entonces los movemos con moveaxis

Como @Eric señala, 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
8
Paul Panzer 17 feb. 2017 a las 19:33

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]]])
3
Andras Deak 17 feb. 2017 a las 16:11

Puede 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 todavía difiere un poco de su matriz deseada para que pueda rodar 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 un manual versión basada 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)

Cronometrada:

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 prometió, "los supera a todos" en términos de rendimiento :-)

3
MSeifert 17 feb. 2017 a las 18:30