如何加速Python积分计算？优化chi_dop函数及积分方案咨询

我目前用以下Python代码做计算，整体耗时约1小时，其中chi和integrand函数运行正常，但chi_dop函数耗时过长。想咨询三个问题：

1. 如何优化以提升运行速度？
2. 除scipy.integrate.quad外，是否有更高效的积分方法？
3. 直接积分虚部与先积分复方程再取虚部，哪种方式效率更高？

import numpy as np
import scipy as sp
from numba import njit

hbar = 1.05e-34
c = 3e8
ep = 8.854187817e-12
td = 23.93e-9
slt = 24.24e-9
dt=np.linspace(-6, 16, 12027)*1e9
dw = 200 #m/s

@njit
def chi(r, dt, dm, slt, td, tf):
    A = np.array([[-1/slt, 1j*r, -1j*r], [-1j*r/2, 0, (1j*dt - 1/td)], [1j*r/2, (-1j*dt - 1/td), 0]])
    B = np.array([0, -1j*r/2, 1j*r/2])
    x = np.linalg.solve(A, B)
    return (2*tf/c)* 2*dm**2*np.imag(x[2])/(hbar*r*ep)*1e14

def integrand(v, r, dt, dw, dm, slt, td, tf):
    int = chi(r,dt-v/c*tf, dm, slt, td, tf) 
    return int* np.exp(-v**2 / (2 * dw**2)) / (np.sqrt(2 * np.pi) * dw)

def chi_dop(r, dt, dw, dm, slt, td,tf):
    chi_int = sp.integrate.quad(integrand, -10*dw, 10*dw, args=(r, dt, dw, dm, slt, td,tf), limit=100, complex_func=True)[0]
    return  1e-14*chi_int
    
chi=np.vectorize(chi)
chi_db = np.vectorize(chi_dop)

F1= chi_db(1/9, dt, dw, 1/18*3.3029259999999997e-29, slt, td, 2*np.pi*384.2276916107209e12) + chi_db(5/9, dt-157e6*2*np.pi, dw, 5/18*3.3029259999999997e-29, slt, td, 2*np.pi*384.2278485512209e12) + chi_db(14/9, dt-423e6*2*np.pi, dw, 7/9*3.3029259999999997e-29, slt, td, 2*np.pi*384.2281881125209e12)

F2 = chi_db(2/9, dt-6765e6*2*np.pi, dw, 1/9*3.3029259999999997e-29,slt, td, 2*np.pi*384.2344540713318e12) + chi_db(5/9, dt-6834e6*2*np.pi, dw, 5/18*3.3029259999999997e-29, slt, td, 2*np.pi*384.23452629333184e12) + chi_db(5/9, dt-6992e6*2*np.pi, dw, 5/18*3.3029259999999997e-29, slt, td, 2*np.pi*384.2349498859318e12)

F3 = chi_db(20/81, dt-1443e6*2*np.pi, dw, 10/81*3.3029259999999997e-29, slt, td, 2*np.pi*384.2320640089228e12) + chi_db(70/81, dt-1510e6*2*np.pi, dw, 35/81*3.3029259999999997e-29, slt, td, 2*np.pi*384.2320933819228e12) + chi_db(2, dt-1630e6*2*np.pi, dw, 1*3.3029259999999997e-29, slt, td, 2*np.pi*384.2321567819228e12)

F4 = chi_db(2/3, dt-4436e6*2*np.pi, dw, 1/3*3.3029259999999997e-29,slt, td, 2*np.pi*384.2290576494837e12) + chi_db(70/81, dt-4478e6*2*np.pi, dw, 35/81*3.3029259999999997e-29, slt, td, 2*np.pi*384.2291210494837e12) + chi_db(56/81, dt-4545e6*2*np.pi, dw, 28/81*3.3029259999999997e-29, slt, td, 2*np.pi*384.2292416894837e12)

问题解答 #

1. 运行速度优化方案 #

• 替换np.vectorize：vectorize是纯Python循环的包装，效率极低。改为向量化实现：将dt作为数组传入，利用numpy广播批量计算所有dt对应的积分，避免逐个元素处理。
• 预计算常数项：把高斯因子里的1/(np.sqrt(2*np.pi)*dw)提前计算好存入变量，不用每次调用integrand都重复计算。
• 手动求解线性方程组：chi函数里的3x3矩阵可以用克莱姆法则推导解析解，替代np.linalg.solve，减少数值求解的开销，配合njit能进一步提速。
• 缩小积分区间：高斯分布在±3σ（即±3*dw）之外的概率小于0.3%，可以把积分区间从-10*dw到10*dw改成-3*dw到3*dw，大幅减少积分迭代次数。
• 并行化处理：每个dt的积分任务相互独立，用joblib或multiprocessing将dt数组拆分后并行计算，充分利用多核CPU资源。

2. 替代scipy.integrate.quad的高效积分方法 #

• 高斯-埃尔米特积分：你的被积函数是高斯权重乘以目标函数，刚好匹配高斯-埃尔米特积分的形式（积分区间(-∞,∞)，权重为高斯函数）。用scipy.special.hermgauss获取节点和权重，直接计算加权和，速度比quad快数倍甚至一个数量级。
• 向量化数值积分：用trapz或simpson等梯形/辛普森积分法，先生成足够密集的v采样点，一次性计算所有dt对应的积分，利用numpy的向量化运算优势，避免quad的单元素循环开销。
• GPU加速：如果有NVIDIA GPU，用numba.cuda或CuPy将chi和积分逻辑移植到GPU上，大规模数组计算的速度会有质的提升。

3. 积分虚部vs积分复方程取虚部的效率对比 #

直接积分虚部效率更高。原因如下：

• scipy.integrate.quad处理复数函数时，会分别对实部和虚部执行两次实积分，相当于双倍计算量。
• 你的chi函数已经直接返回虚部，integrand中的被积函数是实数，直接积分实数不仅减少了一半的计算量，还避免了复数运算的额外开销。

内容的提问来源于stack exchange，提问作者Purple Monkey