我有一个大的二维 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
关于python - numpy数组任意列之间的(内存)高效操作,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/57199248/