Python 快速(er)numpy花式索引和缩减?

Python 快速(er)numpy花式索引和缩减?,python,optimization,numpy,scipy,cython,Python,Optimization,Numpy,Scipy,Cython,我正在尝试使用和加速奇特的索引来“连接”两个数组,并在其中一个结果轴上求和 大概是这样的: $ 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

我正在尝试使用和加速奇特的索引来“连接”两个数组,并在其中一个结果轴上求和

大概是这样的:

$ 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

有没有一种简单的方法来加速[5]中的
语句?我应该使用OpenMP和类似于
scipy.weave
Cython
prange

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

我认为,除此之外,你将看到速度提升的唯一其他方法是使用类似的东西为GPU编写自定义内核。

另一个相关问题是,我将如何使用
pandas
来做同样的事情?Numpy是以C速度做的,因此你可能无法使用weave来加快速度。它只是默认情况下,将数组视为平面,您仍然可以使用
关键字。例如,
np.take(np.arange(10).重塑(-1,2)),[0],axis=0)
将选择第一行。@jorgeca:对,但我认为,除非对平面数组进行索引,否则不能通过指定行和列来提取单个元素。