Python中大规模独立积分的并行化与优化方案咨询

我现在需要计算一大批独立的积分，目前已经用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[i],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[i],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编译，结果确实获得了5倍的性能提升！修改后的函数如下：

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：补充一下，我现在是在个人PC上运行这个脚本，但之后应该可以在集群上运行。

完整代码 #

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[i]


    # 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[i],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("{:<45} {:<40}".format('\nParameter', 'Value'))
        print("-" * 65)
        for key, value in params.items():
            print("{:<45} {:<40}".format(key, value))
        return ""
    
    def debugging(self):
        if self.Debugging:
            # Simulation parameters
            self.N_x = 5 # Number of samples in x direction
            self.N_z = 5 # Number of samples in z direction
            self.N_t = 5 # Number of samples in time
            self.N_max = int(self.d / self._lambda)*2 # Number of terms in the series

        return

if __name__ == "__main__":
    config = TalbotConfig()

    # Are we debugging?
    config.Debugging = False
    config.debugging()

    print(config)

 

    integral = perform_integrals(config)

备注：内容来源于stack exchange，提问作者Gabriel Ybarra Marcaida