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)

我知道的方法是：

带有列列表的花式索引

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。

Answer 1

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

在这个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 3

您可以使用Numba。为了获得最佳性能，通常需要像在C中那样编写简单的循环。（Numba基本上是Python到LLVM-IR的代码转换器，就像C语言的C语言一样）

代码

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。

请注意，使用transonic.jit装饰器，Pythran的效率也很高。

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

4 个答案: