如何让整合更快？

Question

我正在使用以下代码，但需要大约一个小时。函数 chi 和被积函数工作正常，但 chi_dop 花费了太多时间。

如何让它更快？ 除了 scipy.integrate.quad 之外，还有更好的集成方法吗？ 积分虚部和先积分复数方程然后取虚部哪个更快？

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, d-6834te6*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, t,d 2*np.pi*384.2292416894837e12)

Answer 1

我建议进行三项修改：删除矢量化、向

@njit

添加

integrand()

注释，以及对矩阵问题使用封闭形式解决方案。

让我们从删除这些行开始：

chi=np.vectorize(chi)
chi_db = np.vectorize(chi_dop)

您没有以矢量化方式使用这些函数，即使您是这样，最好使用

numba.vectorize

而不是

numpy.vectorize

，因为后者是使用 Python 循环实现的。

我对运行

integrand()

函数（有或没有此更改）进行了基准测试。

using np.vectorize
35.4 µs ± 1.91 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
after removing np.vectorize
6.62 µs ± 30.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

接下来，让我们告诉 Numba 也编译您的被积函数。将

@njit

注释添加到

integrand()

。

@njit
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)

基准：

Previous result
6.62 µs ± 30.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
After @njit
2.18 µs ± 11.6 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

接下来，让我们找到一个封闭形式来解决您的矩阵问题。我将使用 SymPy，它可以找到我们可以重复用于许多值的代数解。最终的解决方案不会使用 SymPy。

import sympy
slt = sympy.Symbol('slt')
r = sympy.Symbol('r')
dt = sympy.Symbol('dt')
td = sympy.Symbol('td')
A = sympy.Matrix([[-1/slt, 1j*r, -1j*r], [-1j*r/2, 0, (1j*dt - 1/td)], [1j*r/2, (-1j*dt - 1/td), 0]])
B = sympy.Matrix([0, -1j*r/2, 1j*r/2])
x = A.solve(B)
print(x.tolist()[2])

输出：

[(-0.5*dt*r*td**2 + 0.5*I*r*td)/(1.0*dt**2*td**2 + 1.0*r**2*slt*td + 1.0)]

这是

x[2]

的代数值。直接计算这个会比每次求解矩阵更快。要在 Python 中使用它，我们必须将表示虚数的

I

替换为

1j

。

chi()

的新定义，使用 SymPy 找到的解决方案：

@njit
def chi(r, dt, dm, slt, td, tf):
    x = (-0.5*dt*r*td**2 + 0.5*1j*r*td)/(1.0*dt**2*td**2 + 1.0*r**2*slt*td + 1.0)
    return (2*tf/c)* 2*dm**2*np.imag(x)/(hbar*r*ep)*1e14

基准：

Matrix Solve
2.18 µs ± 11.6 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
Closed Form
671 ns ± 15.7 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

总体加速了 50 倍。