Return to Revisions

I'm a bit new to working with Numba, but I got the gist of it. I wonder if there any more advanced tricks to make four nested for loops even faster that what I have now. In particular, I need to calculate the following integral:

$$ G_B(\mathbf X, T) = \Lambda \int_\Omega G(\mathbf X, \mathbf X', T) W(\mathbf X', T) \ d\mathbf X' \\ G(\mathbf X, \mathbf X', T) = \frac{1}{2\pi S_0^2} \exp\left[-\frac{\left| \mathbf X -\mathbf X'\right|^2}{2[S_0(1+EB(\mathbf X, T))]^2}\right] $$

Where \$B\$ is a 2D array, and \$S_0\$ and \$E\$ are certain parameters. My code is the following:

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):
 sigma = 2.0*(s0*(1.0+e*b[i,j]))**2
 for ii in range(n):
 for jj in range(m):
 gb[i,j]+=np.exp(-(((i-ii)*dx)**2+((j-jj)*dx)**2)/sigma)
 gb[i,j]*=norm
 return gb
calc_gb_gauss_2d_nb = njit(double[:, :](double[:, :],double,double,double))(calc_gb_gauss_2d)

For and input array of size ×ばつ256 the calculation speed is:

In [4]: a=random.random((256,256))
In [5]: %timeit calc_gb_gauss_2d_nb(a,0.1,1.0,0.5)
The slowest run took 8.46 times longer than the fastest. This could mean that an intermediate result is being cached.
1 loop, best of 3: 1min 1s per loop

Comparison between pure Python and Numba calculation speed give me this picture: Comparative performance plot

Is there any way to optimize my code for better performance?