Rápida indexación de lujo numpy

Mi código para cortar una matriz numpy (a través de una indexación elegante) es muy lento. Actualmente es un cuello de botella en el progtwig.

a.shape (3218, 6) ts = time.time(); a[rows][:, cols]; te = time.time(); print('%.8f' % (te-ts)); 0.00200009 

¿Cuál es la llamada numpy correcta para obtener una matriz que consiste en el subconjunto de filas ‘filas’ y columnas ‘columna’ de la matriz a? (De hecho, necesito la transposición de este resultado)

Puede obtener un poco de velocidad si corta utilizando una indexación y transmisión de lujo:

 from __future__ import division import numpy as np def slice_1(a, rs, cs) : return a[rs][:, cs] def slice_2(a, rs, cs) : return a[rs[:, None], cs] >>> rows, cols = 3218, 6 >>> rs = np.unique(np.random.randint(0, rows, size=(rows//2,))) >>> cs = np.unique(np.random.randint(0, cols, size=(cols//2,))) >>> a = np.random.rand(rows, cols) >>> import timeit >>> print timeit.timeit('slice_1(a, rs, cs)', 'from __main__ import slice_1, a, rs, cs', number=1000) 0.24083110865 >>> print timeit.timeit('slice_2(a, rs, cs)', 'from __main__ import slice_2, a, rs, cs', number=1000) 0.206566124519 

Si piensa en términos de porcentajes, hacer algo un 15% más rápido siempre es bueno, pero en mi sistema, por el tamaño de su matriz, esto nos está costando 40 menos para hacer el corte, y es difícil creer que una operación que tome 240 nosotros seremos tu cuello de botella.

Intenté resumir las excelentes respuestas de Jaime y TheodrosZelleke y mezclar algunos comentarios.

  1. La indexación avanzada (sofisticada) siempre devuelve una copia, nunca una vista.
  2. a[rows][:,cols] implica dos operaciones de indexación sofisticadas, por lo que se crea y descarta una copia intermedia a[rows] . Práctico y legible, pero no muy eficiente. Además, tenga en cuenta que [:,cols] generalmente genera una copia contigua de Fortran desde un C-cont. fuente.
  3. a[rows.reshape(-1,1),cols] es una expresión de indexación avanzada única rows.reshape(-1,1) en el hecho de que rows.reshape(-1,1) y cols se transmiten a la forma del resultado deseado.
  4. Una experiencia común es que la indexación en una matriz aplanada puede ser más eficiente que la indexación elegante, por lo que otro enfoque es

     indx = rows.reshape(-1,1)*a.shape[1] + cols a.take(indx) 

    o

     a.take(indx.flat).reshape(rows.size,cols.size) 
  5. La eficiencia dependerá de los patrones de acceso a la memoria y de si la matriz de inicio es C-continua o Fortran continua, por lo que se necesita experimentación.

  6. Use la indexación sofisticada solo si es realmente necesario: corte básico a[rstart:rstop:rstep, cstart:cstop:cstep] devuelve una vista (aunque no es continua) y debería ser más rápido.

Para mi sorpresa, este tipo de expresión, que calcula los primeros índices 1D lineales, es más de un 50% más rápido que la indexación de matriz consecutiva presentada en la pregunta:

 (a.ravel()[( cols + (rows * a.shape[1]).reshape((-1,1)) ).ravel()]).reshape(rows.size, cols.size) 

ACTUALIZACIÓN: OP actualizó la descripción de la forma de la matriz inicial. Con el tamaño actualizado, el aumento de velocidad ahora está por encima del 99% :

 In [93]: a = np.random.randn(3218, 1415) In [94]: rows = np.random.randint(a.shape[0], size=2000) In [95]: cols = np.random.randint(a.shape[1], size=6) In [96]: timeit a[rows][:, cols] 10 loops, best of 3: 186 ms per loop In [97]: timeit (a.ravel()[(cols + (rows * a.shape[1]).reshape((-1,1))).ravel()]).reshape(rows.size, cols.size) 1000 loops, best of 3: 1.56 ms per loop 

RESPUESTA INITAL: Aquí está la transcripción:

 In [79]: a = np.random.randn(3218, 6) In [80]: a.shape Out[80]: (3218, 6) In [81]: rows = np.random.randint(a.shape[0], size=2000) In [82]: cols = np.array([1,3,4,5]) 

Método de tiempo 1:

 In [83]: timeit a[rows][:, cols] 1000 loops, best of 3: 1.26 ms per loop 

Método de tiempo 2:

 In [84]: timeit (a.ravel()[(cols + (rows * a.shape[1]).reshape((-1,1))).ravel()]).reshape(rows.size, cols.size) 1000 loops, best of 3: 568 us per loop 

Compruebe que los resultados son en realidad los mismos:

 In [85]: result1 = a[rows][:, cols] In [86]: result2 = (a.ravel()[(cols + (rows * a.shape[1]).reshape((-1,1))).ravel()]).reshape(rows.size, cols.size) In [87]: np.sum(result1 - result2) Out[87]: 0.0 

Usando np.ix_ puedes una velocidad similar para desplazarte / remodelar, pero con un código que es más claro:

 a = np.random.randn(3218, 1415) rows = np.random.randint(a.shape[0], size=2000) cols = np.random.randint(a.shape[1], size=6) a = np.random.randn(3218, 1415) rows = np.random.randint(a.shape[0], size=2000) cols = np.random.randint(a.shape[1], size=6) %timeit (a.ravel()[(cols + (rows * a.shape[1]).reshape((-1,1))).ravel()]).reshape(rows.size, cols.size) #101 µs ± 2.36 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each) %timeit ix_ = np.ix_(rows, cols); a[ix_] #135 µs ± 7.47 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each) ix_ = np.ix_(rows, cols) result1 = a[ix_] result2 = (a.ravel()[(cols + (rows * a.shape[1]).reshape((-1,1))).ravel()]).reshape(rows.size, cols.size)​ np.sum(result1 - result2) 0.0