Более быстрый способ реализации numpy.isin с последующим суммированием

Question

Я выполняю анализ данных с использованием скрипта Python и узнал из профилирования, что более 95% времени вычислений занимает строка, которая выполняет следующую операцию np.sum(C[np.isin(A, b)]), где A, C - это 2D NumPy. массивы одинаковой размерности m x n, а b - это одномерный массив переменной длины. Мне интересно, если не выделенная функция NumPy, есть ли способ ускорить такие вычисления?

Типичные размеры A (int64), C (float64): 10M x 100

Типичный размер b (int64): 1000

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

Я не знаю, почему 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 и первый пример данных приведены в ответе Пола Панцерса выше.

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)

Answer 3

Я предполагаю, что вы изменили int64 на int8, где это необходимо.

Вы можете использовать параллельные функции Numba и функцию It для более быстрых вычислений Numpy и использования ядер.

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

Документация для Numba Parallel