[Python] Parallelisation/Optimisation of integrals in Python

I need to compute a large set of INDEPENDENT integrals and I have currently written the following code for it in Python:

# We define the integrand
def integrand(tau,k,t,z,config, epsilon=1e-7):
u = np.sqrt(np.maximum(epsilon,tau**2 - z**2))
return np.sin(config.omega * (t - tau))*k/2, np.sin(config.omega * (t - tau)) * j1(k * u) / u


NON-VECTORISED
partial_integral = np.zeros((len(n_values),len(t_values),len(z_values)))

for n in range(1, len(n_values)): # We skip the n=0 case as it is trivially = 0
for j in range(1, len(z_values)): # We skip the z=0 case as it is trivially = 0
for i in range(1, len(t_values)): # We skip the t=0 case as it is trivially = 0
partial_integral[n,i,j],_ = quad(integrand, x_min[n,i,j], x_max[n,i,j], args=(k_n_values[n],t_values,z_values[j],config), limit=np.inf) # We use quad


Now, all of len(n_values),len(t_values),len(z_values) are big numbers, so I would like to accelerate the code as much as possible. Any suggestions?


Apart from playing with different integration libraries (among which Scipy.quad seems the best), I have thought about vectorising the code:

# VECTORISED
def compute_integral(n,i,j):
quad(integrand, x_min[n,i,j], x_max[n,i,j], args=(k_n_values[n],t_values,z_values[j],config), limit=np.inf) # We use quad

# Use np.meshgrid to create the index grids for n, i, j (starting from 1 to avoid 0-index)
n_grid, i_grid, j_grid = np.meshgrid(np.arange(0, len(n_values)), np.arange(0, len(t_values)), np.arange(0, len(z_values)), indexing='ij')

# Flatten the grids to vectorize the loop over n, i, j
indices = np.vstack([n_grid.ravel(), i_grid.ravel(), j_grid.ravel()]).T

# Vectorize the integral computation using np.vectorize
vectorized_integral = np.vectorize(lambda n, i, j: compute_integral(n, i, j))

# Apply the vectorized function to all combinations of (n, i, j)
partial_integral = np.empty((len(n_values),len(t_values),len(z_values)))
partial_integral[tuple(indices.T)] = vectorized_integral(*indices.T)


But it is not clear that this wins me much...

I have also tried to use Numba (with numba-scipy for the j1) to JIT the integrand in the hope of gaining a bit of performance, and indeed this has served me to get a x5 performance!

PS: As an extra, I am currently running this script in my PC, but I would probably be able to run it on a cluster.

Continue reading...