python - numpy数组任意列之间的(内存)高效操作

Question

我有一个大的二维 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)

我知道的方法是:

1. 花哨的列列表索引

In [5]: %timeit a[:, columns].sum(axis=1)
42.5 ms ± 197 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

1. 使用列掩码进行漂亮的索引

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)

1. 掩码数组

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)

1. 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。

Answer 1

关于这个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