>我有一个大的2D numpy数组。我希望能够在不复制数据的情况下有效地对列的子集运行逐行操作。
在下文中，a = np.arange(1000000).reshape(1000, 10000)和columns = np.arange(1, 1000, 2).供参考，
In [4]: %timeit a.sum(axis=1)
7.26 ms ± 431 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
我知道的方法有：
- 使用列列表进行花式索引
In [5]: %timeit a[:, columns].sum(axis=1)
42.5 ms ± 197 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
- 带有列掩码的花哨索引
In [6]: cols_mask = np.zeros(10000, dtype=bool)
...: cols_mask[columns] = True
In [7]: %timeit a[:, cols_mask].sum(axis=1)
42.1 ms ± 302 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
- 屏蔽数组
In [8]: cells_mask = np.ones((1000, 10000), dtype=bool)
In [9]: cells_mask[:, columns] = False
In [10]: am = np.ma.masked_array(a, mask=cells_mask)
In [11]: %timeit am.sum(axis=1)
80 ms ± 2.71 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
- python循环
In [12]: %timeit sum([a[:, i] for i in columns])
31.2 ms ± 531 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
令我惊讶的是，最后一种方法是最有效的：此外，它避免了复制完整数据，这对我来说是先决条件。但是，它仍然比简单和(数据大小的两倍)慢得多，最重要的是，推广到其他操作(例如，cumsum)并非易事。
我缺少什么方法吗？我可以编写一些cython代码，但我希望该方法适用于任何numpy函数，而不仅仅是sum。
在这个pythran上似乎比numba快一点，至少在我的装备上是这样：
import numpy as np
#pythran export col_sum(float[:,:], int[:])
#pythran export col_sum(int[:,:], int[:])
def col_sum(data, idx):
return data.T[idx].sum(0)
使用pythran <filename.py>编译
计时：
timeit(lambda:cs_pythran.col_sum(a, columns),number=1000)
# 1.644187423051335
timeit(lambda:cs_numba.col_sum(a, columns),number=1000)
# 2.635075871949084
如果你想击败c编译的块求和，你可能最好使用numba。 任何保留在python中的索引(numba用jit创建c编译的函数)都会有python开销。
from numba import jit
@jit
def col_sum(block, idx):
return block[:, idx].sum(1)
%timeit a.sum(axis=1)
100 loops, best of 3: 5.25 ms per loop
%timeit a[:, columns].sum(axis=1)
100 loops, best of 3: 7.24 ms per loop
%timeit col_sum(a, columns)
100 loops, best of 3: 2.46 ms per loop
你可以使用 Numba。为了获得最佳性能，通常需要像在 C 中那样编写简单的循环。 (Numba 基本上是一个 Python 到 LLVM-IR 的代码转换器，很像 C 的 Clang)
法典
import numpy as np
import numba as nb
@nb.njit(fastmath=True,parallel=True)
def row_sum(arr,columns):
res=np.empty(arr.shape[0],dtype=arr.dtype)
for i in nb.prange(arr.shape[0]):
sum=0.
for j in range(columns.shape[0]):
sum+=arr[i,columns[j]]
res[i]=sum
return res
计时
a = np.arange(1_000_000).reshape(1_000, 1_000)
columns = np.arange(1, 1000, 2)
%timeit res_1=a[:, columns].sum(axis=1)
1.29 ms ± 8.05 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit res_2=row_sum(a,columns)
59.3 µs ± 4.35 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
np.allclose(res_1,res_2)
True
使用 Transonic(https://transonic.readthedocs.io)，编写可以通过不同的 Python 加速器(实际上为 Cython、Pythran 和 Numba)加速的代码很容易。
例如，使用boost装饰器，可以编写
import numpy as np
from transonic import boost
T0 = "int[:, :]"
T1 = "int[:]"
@boost
def row_sum_loops(arr: T0, columns: T1):
# locals type annotations are used only by Cython
i: int
j: int
sum_: int
res: "int[]" = np.empty(arr.shape[0], dtype=arr.dtype)
for i in range(arr.shape[0]):
sum_ = 0
for j in range(columns.shape[0]):
sum_ += arr[i, columns[j]]
res[i] = sum_
return res
@boost
def row_sum_transpose(arr: T0, columns: T1):
return arr.T[columns].sum(0)
在我的计算机上，我获得：
TRANSONIC_BACKEND="python" python row_sum_boost.py
Checks passed: results are consistent
Python
row_sum_loops 108.57 s
row_sum_transpose 1.38
TRANSONIC_BACKEND="cython" python row_sum_boost.py
Checks passed: results are consistent
Cython
row_sum_loops 0.45 s
row_sum_transpose 1.32 s
TRANSONIC_BACKEND="numba" python row_sum_boost.py
Checks passed: results are consistent
Numba
row_sum_loops 0.27 s
row_sum_transpose 1.16 s
TRANSONIC_BACKEND="pythran" python row_sum_boost.py
Checks passed: results are consistent
Pythran
row_sum_loops 0.27 s
row_sum_transpose 0.76 s
有关完整代码和此问题示例的更完整比较，请参阅 https://transonic.readthedocs.io/en/stable/examples/row_sum/txt.html。
请注意，Pythran 对transonic.jit装饰器也非常有效。