Делаем четыре вложенных цикла for еще быстрее с помощью Numba

Используя numpy и некоторую математику, можно ускорить ваш код, поэтому он становится быстрее, чем текущая numba-версия на порядок. Мы также увидим, что использование numba в улучшенной функции делает ее еще быстрее.

Довольно часто numba используется слишком часто - часто можно написать код только для numpy, который весьма эффективен - это также имеет место здесь.

Проблема с имеющимся кодом numpy: не следует обращаться к отдельным элементам, а использовать встроенные функции numpy — они работают так же быстро, как и большую часть времени. Только если невозможно использовать эти функции numpy, можно использовать numba или cython.

Однако самая большая проблема здесь заключается в формулировке проблемы. Для фиксированных i и j у нас есть следующая формула для расчета (я немного упростил ее):

 g[i,j]=sum_ii sum_jj exp(value_ii+value_jj)
       =sum_ii sum_jj exp(value_ii)*exp(value_jj)
       =sum_ii exp(value_ii) * sum_jj exp(value_jj)

Для вычисления последней формулы нужно O(n+m) операций, а для первой, наивной формулы O(n*m) - совсем другое дело!

Первая версия, использующая функциональность numpy, может быть похожа на:

def calc_ead(b,s0,e,dx):
    n,m=b.shape
    norm = 1.0/(2*np.pi*s0**2)
    gb = np.zeros((n,m))
    vI=np.arange(n)
    vJ=np.arange(m)
    for i in range(n):
        for j in range(m):
            II=(i-vI)*dx
            JJ=(j-vJ)*dx
            denom=2.0*(s0*(1.0+e*b[i,j]))**2
            expII=np.exp(-II*II/denom)
            expJJ=np.exp(-JJ*JJ/denom)
            gb[i,j]=norm*(expII.sum()*expJJ.sum())
    return gb

А теперь по сравнению с оригинальной numba-реализацией:

>>> a=np.random.random((256,256))

>>> print(calc_gb_gauss_2d_nb(a,0.1,1.0,0.5)[1,1])
15.9160709993
>>> %timeit -n1 -r1 calc_gb_gauss_2d_nb(a,0.1,1.0,0.5)
1min 6s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

и теперь numpy-функция:

>>> print(calc_ead(a,0.1,1.0,0.5)[1,1])
15.9160709993
>>> %timeit -n1 -r1 calc_ead(a,0.1,1.0,0.5)
1.8 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

Есть два наблюдения:

1. результаты такие же.
2. версия numpy в 37 раз быстрее, для больших задач эта разница станет еще больше.

Очевидно, вы можете использовать numba для еще большего ускорения. Тем не менее, по-прежнему рекомендуется использовать numpy-функциональность, когда это возможно - довольно удивительно, насколько тонкими могут быть самые простые вещи - например, даже подсчет суммы:

>>> nb_calc_ead = njit(double[:, :](double[:, :],double,double,double))(calc_ead)
>>>print(nb_calc_ead(a,0.1,1.0,0.5)[1,1])
15.9160709993
>>>%timeit -n1 -r1 nb_calc_ead(a,0.1,1.0,0.5)
587 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

Есть еще фактор 3!

Эту задачу можно было бы распараллелить, но сделать это правильно не так просто. Моя дешевая попытка использовать явное распараллеливание циклов:

from numba import njit, prange
import math

@njit(parallel=True)                 #needed, so it is parallelized
def parallel_nb_calc_ead(b,s0,e,dx):
    n,m=b.shape
    norm = 1.0/(2*np.pi*s0**2)
    gb = np.zeros((n,m))
    vI=np.arange(n)
    vJ=np.arange(m)
    for i in prange(n):             #outer loop = explicit prange-loop
        for j in range(m):
            denom=2.0*(s0*(1.0+e*b[i,j]))**2
            expII=np.zeros((n,))
            expJJ=np.zeros((m,))
            for k in range(n):
                II=(i-vI[k])*dx
                expII[k]=math.exp(-II*II/denom)

            for k in range(m):
                JJ=(j-vJ[k])*dx
                expJJ[k]=math.exp(-JJ*JJ/denom)
            gb[i,j]=norm*(expII.sum()*expJJ.sum())
    return gb

И сейчас:

>>> print(parallel_nb_calc_ead(a,0.1,1.0,0.5)[1,1])
15.9160709993
>>> %timeit -n1 -r1 parallel_nb_calc_ead(a,0.1,1.0,0.5)
349 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

означает почти еще один коэффициент 2 (у моей машины только два процессора, в зависимости от аппаратного обеспечения ускорение может быть больше). Кстати, мы почти в 200 раз быстрее оригинальной версии.

Бьюсь об заклад, можно было бы улучшить приведенный выше код, но я не собираюсь этого делать.

Список текущей версии, с которой сравнивается calc_ead:

import numpy as np
from numba import njit, double

def calc_gb_gauss_2d(b,s0,e,dx):
    n,m=b.shape
    norm = 1.0/(2*np.pi*s0**2)
    gb = np.zeros((n,m))
    for i in range(n):
        for j in range(m):
            for ii in range(n):
                for jj in range(m):
                    gb[i,j]+=np.exp(-(((i-ii)*dx)**2+((j-jj)*dx)**2)/(2.0*(s0*(1.0+e*b[i,j]))**2))
            gb[i,j]*=norm
    return gb

calc_gb_gauss_2d_nb = njit(double[:, :](double[:, :],double,double,double))(calc_gb_gauss_2d)