快速（呃）numpy花式索引和減少？

Question

我正在嘗試使用並加速花式索引來「連接」兩個數組並總結其中一個結果軸。

事情是這樣的：

$ 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 

有一個簡單的方法來加速In [5]的聲明？我應該去OpenMP和類似scipy.weave或Cython的prange？

Answer 1

numpy.take由於某種原因比幻想索引快得多。唯一的竅門是它將數組視爲平坦的。

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 

我認爲唯一的其他方式，你會看到很多的速度提高之外，這將是寫一個使用類似PyCUDA一個GPU定製的內核。