Numerical Double Integration using numba and scipy

Question

Please see the following code. I am using it to calculate the double integration. Please help me to improve the code.

import numpy as np
from scipy import integrate
from scipy.special import j1
from scipy.special import j0
from numpy import sin
from numpy import cos
from numpy import pi
import time
import numba as nb
from numba import jit
from numba import cfunc
from numba.types import intc, CPointer, float64
from scipy import LowLevelCallable
from matplotlib import pyplot as plt
q = np.linspace(0.03, 1.0, 1000)
start = time.time()
length = 20000
pd = 0.1
radius = 40
start = time.time()
def jit_integrand_function(integrand_function):
 jitted_function = jit(integrand_function, nopython=True)
 @cfunc(float64(intc, CPointer(float64)),error_model="numpy",fastmath=True)
 def wrapped(n, xx):
 ar = nb.carray(xx, n)
 return jitted_function(ar[0], ar[1], ar[2])
 return LowLevelCallable(wrapped.ctypes)
@jit_integrand_function
def f(theta,z,q):
 def ff(theta, z, q):
 l = length / 2
 qz = q * sin(theta)
 qr = q * cos(theta)
 return (4 * pi * sin(qz * l) * z * j1(qr * z) / (qr * qz)) ** 2
 def g_distribution(z, radius, pd):
 c = radius
 s = pd * c
 return (1 / s)) * np.exp(-0.5 * (((z - radius) / s) ** 2))
 return ff(theta,z,q)*g_distribution(z,radius,pd)* sin(theta)
def lower_inner(z):
 return 0.
def upper_inner(z):
 return 120.
y = np.empty(len(q))
for n in range(len(q)):
 y[n] = integrate.dblquad(f, 0, 1*pi/180, lower_inner, upper_inner,args=(q[n],))[0]
end = time.time()
print(end - start)
plt.loglog(q,y)
plt.show()

Also, I was wondering if I can use GPU for the above code? How much gain I would have for the trouble(I have zero knowledge of programming for GPU use)?

• mle: Why do u use the cfunc decorator instead of njit/jit decorators? Maybe writing the integration scheme entirely in numba would be best I think in terms of speed.

Answer 1

Prior to vectorisation, the simple version of this can look like

import time
import numpy as np
import scipy
from matplotlib import pyplot as plt
def f(theta: float, z: float, q: float, l: float, radius: float, sinv: float) -> float:
 qz = q * np.sin(theta)
 qr = q * np.cos(theta)
 ff = (
 4*np.pi*z/(qr * qz)*np.sin(qz * l)*scipy.special.j1(qr * z)
 )**2
 g_distribution = sinv*np.exp(-0.5 * ((z - radius)*sinv)**2)
 return ff*g_distribution*np.sin(theta)
def lower_inner(z: float) -> float:
 return 0.
def upper_inner(z: float) -> float:
 return 120.
def solve(
 length: float = 20_000,
 pd: float = 0.1,
 radius: float = 40,
 n_q: int = 1_000,
) -> tuple[np.ndarray, np.ndarray]:
 q = np.linspace(start=0.03, stop=1.0, num=n_q)
 y = np.empty(n_q)
 l = 0.5*length
 c = radius
 s = pd * c
 sinv = 1/s
 for n in range(n_q):
 y[n], err = scipy.integrate.dblquad(
 func=f,
 a=0, b=np.pi/180,
 gfun=lower_inner, hfun=upper_inner,
 args=(q[n], l, radius, sinv),
 )
 return q, y
def regression_test() -> None:
 q, y = solve(n_q=50)
 assert q.shape == (50,)
 assert y.shape == (50,)
 assert np.isclose(y.min(), -1.0677255659807734e-26, atol=0, rtol=1e-14)
 assert np.isclose(y.max(), 5.1716382391107805e-27, atol=0, rtol=1e-14)
 assert np.isclose(y.sum(), 3.2115029907631903e-27, atol=0, rtol=1e-14)
def demo() -> None:
 start = time.monotonic()
 q, y = solve(n_q=200)
 end = time.monotonic()
 print('Duration:', end - start)
 plt.plot(q, y)
 plt.show()
if __name__ == '__main__':
 regression_test()
 demo()

but are you sure that your expressions are correct? These quantities seem awfully low:

graph

More important than numba, I think, is actual vectorisation. Whereas a vector-valued double integral isn't supported in scipy yet, there are draft implementations of an nquad_vec; see for instance scikit-fda.