numpy 数组的任意列之间的（内存）高效操作

Question

我有一个大型 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)

我知道的方法是：

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)

4. 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

Answer 2

如果您想击败 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

Answer 3

您可以使用 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

Answer 4

使用 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 装饰器也非常高效。