Самый быстрый Python log-sum-exp в 'Reduat'

Question

Как часть пакета статистического программирования, мне нужно добавить преобразованные в лог значения вместе с LogSumExp Function . Это значительно менее эффективно, чем сложение незарегистрированных значений.

Кроме того, мне нужно добавить значения вместе, используя функциональность numpy .ufun c .reduecat .

Существуют различные варианты, которые я рассмотрел, с код ниже:

1. (для сравнения в не лог-пространстве) используйте numpy .add.reduceat
2. Numpy ufun c для добавления зарегистрированных значений вместе: np.logaddexp.reduceat
3. Рукописная функция сокращения со следующими функциями logsumexp:
   • реализация scipy logsumexp
   • функция logsumexp в Python (с numba )
   • Функция потокового logsumexp in Python (с numba )

def logsumexp_reduceat(arr, indices, logsum_exp_func):
    res = list()
    i_start = indices[0]
    for cur_index, i in enumerate(indices[1:]):
        res.append(logsum_exp_func(arr[i_start:i]))
        i_start = i

    res.append(logsum_exp_func(arr[i:]))
    return res

@numba.jit(nopython=True)
def logsumexp(X):
    r = 0.0
    for x in X:
        r += np.exp(x)  
    return np.log(r)

@numba.jit(nopython=True)
def logsumexp_stream(X):
    alpha = -np.Inf
    r = 0.0
    for x in X:
        if x != -np.Inf:
            if x <= alpha:
                r += np.exp(x - alpha)
            else:
                r *= np.exp(alpha - x)
                r += 1.0
                alpha = x
    return np.log(r) + alpha

arr = np.random.uniform(0,0.1, 10000)
log_arr = np.log(arr)
indices = sorted(np.random.randint(0, 10000, 100))

# approach 1
%timeit np.add.reduceat(arr, indices)
12.7 µs ± 503 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

# approach 2
%timeit np.logaddexp.reduceat(log_arr, indices)
462 µs ± 17.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

# approach 3, scipy function
%timeit logsum_exp_reduceat(arr, indices, scipy.special.logsumexp)
3.69 ms ± 273 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# approach 3 handwritten logsumexp
%timeit logsumexp_reduceat(log_arr, indices, logsumexp)
139 µs ± 7.1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

# approach 3 streaming logsumexp
%timeit logsumexp_reduceat(log_arr, indices, logsumexp_stream)
164 µs ± 10.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Результаты timeit показывают, что рукописные функции logsumexp с numba являются самыми быстрыми вариантами, но все же в 10 раз медленнее, чем numpy .add.reduceat.

Несколько вопросов:

1. Существуют ли другие подходы (или настройки для вариантов, которые я представил? d) что быстрее? Например, есть ли способ использовать таблицу поиска для вычисления функции logsumexp?
2. Почему функция Себастьяна Новозина «streaming logsumexp» не быстрее, чем наивный подход?

Answer 1

Есть некоторые возможности для улучшения

Но никогда не ожидайте, что logsumexp будет столь же быстрым, как стандартное суммирование, потому что exp довольно дорогая операция.

Пример

import numpy as np

#from version 0.43 until 0.47 this has to be set before importing numba
#Bug: https://github.com/numba/numba/issues/4689
from llvmlite import binding
binding.set_option('SVML', '-vector-library=SVML')
import numba as nb

@nb.njit(fastmath=True,parallel=False)
def logsum_exp_reduceat(arr, indices):
    res = np.empty(indices.shape[0],dtype=arr.dtype)

    for i in nb.prange(indices.shape[0]-1):
        r = 0.
        for j in range(indices[i],indices[i+1]):
            r += np.exp(arr[j])  
        res[i]=np.log(r)

    r = 0.
    for j in range(indices[-1],arr.shape[0]):
        r += np.exp(arr[j])  
    res[-1]=np.log(r)
    return res

Время

#small example where parallelization doesn't make sense
arr = np.random.uniform(0,0.1, 10_000)
log_arr = np.log(arr)
#use arrays if possible
indices = np.sort(np.random.randint(0, 10_000, 100))

%timeit logsum_exp_reduceat(arr, indices)
#without parallelzation 22 µs ± 173 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
#with parallelization   84.7 µs ± 32.2 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit np.add.reduceat(arr, indices)
#4.46 µs ± 61.1 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

#large example where parallelization makes sense
arr = np.random.uniform(0,0.1, 1000_000)
log_arr = np.log(arr)
indices = np.sort(np.random.randint(0, 1000_000, 100))

%timeit logsum_exp_reduceat(arr, indices)
#without parallelzation 1.57 ms ± 14.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
#with parallelization   409 µs ± 14.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit np.add.reduceat(arr, indices)
#340 µs ± 11.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)