¿Rápido (e) nimio lujo de indexación y reducción?

Question

Estoy tratando de usar y acelerar la indexación elegante para "unir" dos matrices y sumar en uno de los ejes de los resultados.

Algo como esto:

$ ipython 
In [1]: import numpy as np 
In [2]: ne, ds = 12, 6 
In [3]: i = np.random.randn(ne, ds).astype('float32') 
In [4]: t = np.random.randint(0, ds, size=(1e5, ne)).astype('uint8') 

In [5]: %timeit i[np.arange(ne), t].sum(-1) 
10 loops, best of 3: 44 ms per loop 

¿Existe una manera sencilla de acelerar la declaración en In [5]? ¿Debería ir con OpenMP y algo así como scipy.weave o ? prange?

Answer 1

numpy.take es mucho más rápido que la indexación elegante por alguna razón. El único truco es que trata la matriz como plana.

In [1]: a = np.random.randn(12,6).astype(np.float32) 

In [2]: c = np.random.randint(0,6,size=(1e5,12)).astype(np.uint8) 

In [3]: r = np.arange(12) 

In [4]: %timeit a[r,c].sum(-1) 
10 loops, best of 3: 46.7 ms per loop 

In [5]: rr, cc = np.broadcast_arrays(r,c) 

In [6]: flat_index = rr*a.shape[1] + cc 

In [7]: %timeit a.take(flat_index).sum(-1) 
100 loops, best of 3: 5.5 ms per loop 

In [8]: (a.take(flat_index).sum(-1) == a[r,c].sum(-1)).all() 
Out[8]: True 

creo que la única otra forma en que vamos a ver mucho de una mejora en la velocidad más allá de esto sería escribir un kernel personalizado para una GPU usando algo como PyCUDA.