Question

我正在使用python脚本进行数据分析，并从分析中了解到，执行以下操作np.sum(C[np.isin(A, b)])的行占用了95％以上的计算时间，其中A，{{ 1}}是等尺寸C的2D NumPy数组，而m x n是可变长度的1D数组。我想知道是否不是专用的NumPy函数，是否可以加速这种计算？

b，A (int64)：C (float64)

的典型大小

10M x 100的典型大小：b (int64)

Answer 1

由于标签的整数范围较小，因此使用下面的np.bincount（pp）可以大大提高速度。另外，您可以通过创建遮罩（p2）来加快查找速度。这-就像您的原始代码一样-允许将np.sum替换为math.fsum，以保证在机器精度结果（p3之内保持精确。另外，我们可以将其pythranize进行另一次40%加速（p4）。

在我的装备上，numba soln（mx）的运行速度与pp差不多，但是也许我做得不好。

import numpy as np
import math
from subsum import pflat

MAXIND = 120_000

def OP():
    return sum(C[np.isin(A, b)])

def pp():
    return np.bincount(A.reshape(-1), C.reshape(-1), MAXIND)[np.unique(b)].sum()
def p2():
    grid = np.zeros(MAXIND, bool)
    grid[b] = True
    return C[grid[A]].sum()
def p3():
    grid = np.zeros(MAXIND, bool)
    grid[b] = True
    return math.fsum(C[grid[A]])
def p4():
    return pflat(A.ravel(), C.ravel(), b, MAXIND)

import numba as nb

@nb.njit(parallel=True,fastmath=True)
def nb_ss(A,C,b):
    s=set(b)
    sum=0.
    for i in nb.prange(A.shape[0]):
        for j in range(A.shape[1]):
            if A[i,j] in s:
                sum+=C[i,j]
    return sum

def mx():
    return nb_ss(A,C,b)

sh = 100_000, 100

A = np.random.randint(0, MAXIND, sh)
C = np.random.random(sh)
b = np.random.randint(0, MAXIND, 1000)

print(OP(), pp(), p2(), p3(), p4(), mx())

from timeit import timeit

print("OP", timeit(OP, number=4)*250)
print("pp", timeit(pp, number=10)*100)
print("p2", timeit(p2, number=10)*100)
print("p3", timeit(p3, number=10)*100)
print("p4", timeit(p4, number=10)*100)
print("mx", timeit(mx, number=10)*100)

pythran模块的代码：

[subsum.py]

import numpy as np

#pythran export pflat(int[:], float[:], int[:], int)

def pflat(A, C, b, MAXIND):
    grid = np.zeros(MAXIND, bool)
    grid[b] = True
    return C[grid[A]].sum()

编译就像pythran subsum.py

一样简单

样品运行：

41330.15849965791 41330.15849965748 41330.15849965747 41330.158499657475 41330.15849965791 41330.158499657446
OP 1963.3807722493657
pp 53.23419079941232
p2 21.8758742994396
p3 26.829131800332107
p4 12.988955597393215
mx 52.37018179905135

Answer 2

我假设您已根据需要将int64更改为int8。

您可以使用Numba的parallel和It功能进行更快的Numpy计算并利用内核。

@numba.jit(nopython=True, parallel=True)
def (A,B,c):
    return np.sum(C[np.isin(A, b)])

Documentation for Numba Parallel

Answer 3

我不知道为什么np.isin这么慢，但是您可以更快地实现您的功能。以下Numba解决方案使用一组可快速查找值并进行并行化。内存占用量也小于Numpy实现中的内存占用量。

代码

import numpy as np
import numba as nb


@nb.njit(parallel=True,fastmath=True)
def nb_pp(A,C,b):
    s=set(b)
    sum=0.
    for i in nb.prange(A.shape[0]):
        for j in range(A.shape[1]):
            if A[i,j] in s:
                sum+=C[i,j]
    return sum

时间

pp实现和第一个数据样本是Paul Panzers上面回答的形式。

MAXIND = 120_000
sh = 100_000, 100
A = np.random.randint(0, MAXIND, sh)
C = np.random.random(sh)
b = np.random.randint(0, MAXIND, 1000)

MAXIND = 120_000
%timeit res_1=np.sum(C[np.isin(A, b)])
1.5 s ± 10.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit res_2=pp(A,C,b)
62.5 ms ± 624 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit res_3=nb_pp(A,C,b)
17.1 ms ± 141 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)


MAXIND = 10_000_000
%timeit res_1=np.sum(C[np.isin(A, b)])
2.06 s ± 27.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit res_2=pp(A,C,b)
206 ms ± 3.67 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit res_3=nb_pp(A,C,b)
17.6 ms ± 332 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

MAXIND = 100
%timeit res_1=np.sum(C[np.isin(A, b)])
1.01 s ± 20.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit res_2=pp(A,C,b)
46.8 ms ± 538 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit res_3=nb_pp(A,C,b)
3.88 ms ± 84.8 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

实现numpy.isin后跟sum的更快方法

3 个答案: