Используя 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)
Есть два замечания:
- результаты совпадают.
- NumPy версия в 37 раз быстрее, для больших проблем эта разница станет еще больше.
Очевидно, что вы можете использовать нумбу для еще большего ускорения. Тем не менее, это по-прежнему хорошая идея использовать 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)