Оптимизация команды с участием суммы и экспоненты

Question

Я кодирую в python, используя numpy. Я хочу оптимизировать формулу, которая выглядит следующим образом, я использую картинку для удобства чтения.

В этом примере времена t взяты из разных списков, указанных в верхнем индексе. Соответствующим вектором здесь является T_t, который является списком списка. Вот мой оригинальный код:

def first_version(m, n, k, T_t, BETA):
    if k == 1:
        return 0
    ans = 0
    for i in range(len(T_t[n])):
        if T_t[n][i] < T_t[m][k - 1]:
            ans += (T_t[m][k - 1] - T_t[n][i]) * np.exp(-BETA[m, n] * (T_t[m][k - 1] - T_t[n][i]))
        else:
            break
    return ans

Перерыв в конце позволяет мне сэкономить некоторое время. У меня была блестящая идея использовать библиотеку numpy для улучшения производительности:

def second_version(m, n, k, T_t, BETA):
    if k == 1:
        return 0
    the_times = np.maximum( T_t[m][k - 1] - np.array(T_t[n]) , 0  )
    ans = sum(the_times * np.exp( -BETA[m, n] * the_times  ))
    return ans

Для сравнения, второй алгоритм работает в 100 раз быстрее. Можно ли сделать лучше? В частности, я сожалею о том, что numpy вычисляет максимум по всему вектору, когда, вероятно, половина его будет равна 0 в конце.

У вас есть идеи, как улучшить эти биты кода?

---

Я забыл сумму в коде номер 2. Это замедляет код и делает его только в 20 раз быстрее .

Answer 1

У меня есть 2 основных предложения:

• Использование np.sum() вместо sum() примерно в три раза увеличивает скорость second_version
• Использование numba.jit снова увеличивает скорость примерно на 8й. (На самом деле вы можете jit-скомпилировать любую версию и получить примерно одинаковую скорость)

Пример полного кода:

import numpy as np
import numba
import timeit


def first_version(m, n, k, T_t, BETA):
    if k == 1:
        return 0
    ans = 0
    for i in range(len(T_t[n])):
        if T_t[n][i] < T_t[m][k - 1]:
            ans += (T_t[m][k - 1] - T_t[n][i]) * np.exp(-BETA[m, n] * (T_t[m][k - 1] - T_t[n][i]))
        else:
            break
    return ans


def second_version(m, n, k, T_t, BETA):
    if k == 1:
        return 0
    the_times = np.maximum( T_t[m][k - 1] - np.array(T_t[n]) , 0  )
    ans = np.sum(the_times * np.exp( -BETA[m, n] * the_times  ))
    return ans


def jit_version(m, n, k, T_t, BETA):
    # wrapper makes it to that numba doesn't have to deal with 
    # the list-of-arrays data type
    return jit_version_core(k, T_t[m], T_t[n], BETA[m, n])


@numba.jit(nopython=True)
def jit_version_core(k, t1, t2, b):
    if k == 1:
        return 0
    ans = 0
    for i in range(len(t2)):
        if t2[i] < t1[k - 1]:
            ans += (t1[k - 1] - t2[i]) * np.exp(-b * (t1[k - 1] - t2[i]))
        else:
            break
    return ans


N = 10000
t1 = np.cumsum(np.random.random(size=N))
t2 = np.cumsum(np.random.random(size=N))
beta = np.random.random(size=(2, 2))

for fn in ['first_version', 'second_version', 'jit_version']:
    print("------", fn)
    v = globals()[fn](0, 1, len(t1), [t1, t2], beta)
    t = timeit.timeit('%s(0, 1, len(t1), [t1, t2], beta)' % fn, number=100, globals=globals())
    print("output:", v, "time:", t)

И вывод:

------ first_version
output: 3.302938986817431 time: 2.900316455983557
------ second_version
output: 3.3029389868174306 time: 0.12064526398899034
------ jit_version
output: 3.302938986817431 time: 0.013476221996825188