快速（呃）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定制的内核。