Fast (er) numpy fancy Indexing und Reduktion?

Question

Ich versuche zu verwenden und zu beschleunigen Phantasie Indexierung zu "Join" zwei Arrays und Summe über eine der Achsen der Ergebnisse.

Etwas wie folgt aus:

$ 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 

Gibt es eine einfache Möglichkeit, die Aussage in In [5] zu beschleunigen? Sollte ich mit OpenMP gehen und so etwas wie scipy.weave oder Cythonprange?

Answer 1

numpy.take ist aus irgendeinem Grund viel schneller als Phantasieindexierung. Der einzige Trick ist, dass das Array als flach behandelt wird.

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 

Ich denke, die einzige andere Weg Sie viel von einer Verbesserung der Geschwindigkeit sehen werden darüber hinaus mit so etwas wie PyCUDA einen angepassten Kernel für eine GPU zu schreiben wäre.