Распараллеливание/оптимизация интегралов в Python ⇐ Python

[Anonymous]

Мне нужно вычислить большой набор НЕЗАВИСИМЫХ интегралов, и сейчас я написал для этого следующий код на 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

Теперь все len(n_values),len(t_values),len(z_values) являются большими числами, поэтому я хотел бы максимально ускорить код. Есть предложения?

Помимо экспериментов с различными библиотеками интеграции (среди которых Scipy.quad кажется лучшей), я думал о векторизации кода:
# 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, epsilon), limit=10000) # 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)

Но не факт, что это меня сильно выиграет...
Я также пытался использовать Numba (с numba-scipy для j1 ) для JIT подынтегральной функции в надежде немного повысить производительность, и это действительно помогло мне получить производительность x5! Для этого я просто изменил свою функцию на следующую:
from numba import njit # Remember to also have numba-scipy installed!!!!!

# We define the integrand and JIT it with Numba (and numba-scipy) for a faster performance
@njit
def integrand(tau,k,t,z,omega, epsilon):
u = np.sqrt(np.maximum(0,tau**2 - z**2))
if u < epsilon:
return np.sin(omega * (t - tau))*k/2
else:
return np.sin(omega * (t - tau)) * j1(k * u) / u

PS: Кроме того, в настоящее время я запускаю этот сценарий на своем ПК, но, вероятно, смогу запустить его и в кластере.

Полный код:
import os
import numpy as np
from scipy.special import j1
from scipy.integrate import quad
from numba import njit # Remember to also have numba-scipy installed!!!!!

from tqdm import tqdm

def perform_integrals(config):
'''
We perform the integrals using quad
'''

# We store the range of n, kn and gn in arrays of length N_max
n_values = np.linspace(0, config.N_max-1, config.N_max, dtype=int)
k_n_values = 2 * np.pi * n_values

# We store the range of t, z in the arrays of dimension (N_t) and (N_z)
t_values = np.linspace(0, config.N_t*config.delta_t, config.N_t)
z_values = np.linspace(0, config.N_z*config.delta_z, config.N_z)

# Preallocate the result arrays (shape: len(n_values) x len(z_values))
x_min = np.zeros((len(t_values), len(z_values)))
x_max = np.empty((len(t_values), len(z_values)))

# Compute the values
t1_values = np.roll(t_values, 1)
t1_values[0] = 0.
x_min = np.maximum(z_values[None, :], t1_values[:, None]) # Max between z_values[j] and t_values[i-1]
x_max = np.maximum(z_values[None, :], t_values[:, None]) # Max between z_values[j] and t_values

# We define the integrand and JIT it with Numba (and numba-scipy) for a faster performance
@njit
def integrand(tau,k,t,z,omega, epsilon=1e-7):
u = np.sqrt(np.maximum(0,tau**2 - z**2))
if u < epsilon:
return np.sin(omega * (t - tau))*k/2
else:
return np.sin(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 tqdm(range(1, len(n_values))): # We skip the n=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
for j in range(1, len(z_values)): # We skip the z=0 case as it is trivially = 0
partial_integral[n,i,j],_ = quad(integrand, x_min[i,j], x_max[i,j], args=(k_n_values[n],t_values,z_values[j],config.omega, 1e-7), limit=10000, epsabs=1e-7, epsrel=1e-4) # We use quad

return partial_integral

class TalbotConfig:
def __init__(self):
self.A = 1. # Amplitude of signal
self.c = 1. # Speed of light
self.d = 1. # Distance between gratings we fix it = 1
self._lambda = self.d / 10. # Wavelength
self.w = 2 * self._lambda # Width of the gratings

# Other relevant magnitudes
self.omega = 2 * np.pi * self.c / self._lambda # Frequency of the signal
self.z_T = self._lambda/(1. - np.sqrt(1.-(self._lambda/self.d) ** 2)) # Talbot distance = 2 d^2/λ

# Simulation parameters
self.N_x = 27*2 -1 # Number of samples in x direction
self.N_z = 192*2 -1 # Number of samples in z direction
self.N_t = 100 -1 # Number of samples in time
self.N_max = int(self.d / self._lambda * 4) # Number of terms in the series

# Other relevant magnitudes
self.last_t_zT = 1. # Final time / Z_t
self.delta_t = self.z_T/self.c/self.N_t * self.last_t_zT # Time between photos
self.delta_x = self.d/2/self.N_x # X-Distance between points
self.delta_z = self.z_T/self.N_z # Z-Distance between points

def __str__(self):
params = {
"Amplitude of signal (A)": self.A,
"Speed of light (c)": self.c,
"Distance between gratings (d)": self.d,
"Wavelength (lambda)": self._lambda,
"Width of the gratings (w)": self.w,
"Frequency of the signal (omega)": self.omega,
"Talbot distance (z_T)": self.z_T,
"Number of samples in x direction (N_x)": self.N_x,
"Number of samples in z direction (N_z)": self.N_z,
"Number of samples in time (N_t)": self.N_t,
"Number of terms in the series (N_max)": self.N_max,
"Time between photos (delta_t)": self.delta_t,
"X-Distance between points (delta_x)": self.delta_x,
"Z-Distance between points (delta_z)": self.delta_z,
"Final time / z_T": self.last_t_zT
}

print("{:

Подробнее здесь: https://stackoverflow.com/questions/793 ... -in-python