如何提升Python数值积分速度？多方案对比及优化方法咨询

待积分函数 #

我有如下形式的待积分函数：

def f(z, t, q):
    return t * 0.5 * (erf((t - z) / 3) - 1) * j0(q * t) * np.exp(-0.5 * ((z - 40) / 2) ** 2)

我将该函数作为示例，用于理解和展示不同积分方法的差异。结合现有技术资料和官方文档说明，我尝试了多种方法提升该函数的积分速度，以下是简要说明和对比：

方案1：仅使用Python和scipy.quad实现 #

耗时：202.76 s

import numpy as np
from scipy import integrate
from scipy.special import erf
from scipy.special import j0
import time

q = np.linspace(0.03, 1.0, 1000)

start = time.time()

def f(q, z, t):
    return t * 0.5 * (erf((t - z) / 3) - 1) * j0(q * t)  * np.exp(-0.5 * ((z - 40) / 2) ** 2)


y = np.empty([len(q)])
for n in range(len(q)):
    y[n] = integrate.dblquad(lambda t, z: f(q[n], z, t), 0, 50, lambda z: 10, lambda z: 60)[0]

end = time.time()
print(end - start)

# 输出：202.76 s

如预期所示，该方案速度较慢。

方案2：使用低级可调用对象配合Numba编译函数实现 #

耗时：6.46 s

import numpy as np
from scipy import integrate
from scipy.special import erf
from scipy.special import j0
import time
import numba as nb
from numba import cfunc
from numba.types import intc, CPointer, float64
from scipy import LowLevelCallable

q = np.linspace(0.03, 1.0, 1000)

start = time.time()

def jit_integrand_function(integrand_function):
    jitted_function = nb.njit(integrand_function, nopython=True)

    # error_model="numpy" -> 不检查除零错误
    @cfunc(float64(intc, CPointer(float64)),error_model="numpy",fastmath=True)
    def wrapped(n, xx):
        ar = nb.carray(xx, n)
        return jitted_function(ar[0], ar[1], ar[2])
    return LowLevelCallable(wrapped.ctypes)

@jit_integrand_function
def f(t, z, q):
    return t * 0.5 * (erf((t - z) / 3) - 1) * j0(q * t) * (1 / (np.sqrt(2 * np.pi) * 2)) * np.exp(
        -0.5 * ((z - 40) / 2) ** 2)

def lower_inner(z):
    return 10.

def upper_inner(z):
    return 60.


y = np.empty(len(q))
for n in range(len(q)):
    y[n] = integrate.dblquad(f, 0, 50, lower_inner, upper_inner,args=(q[n],))[0]

end = time.time()
print(end - start)

# 输出：6.46 s

该方案比第一种方法速度更快。

方案3：使用scipy.quad_vec实现 #

耗时：2.87 s

import numpy as np
from scipy import integrate
from scipy.special import erf
from scipy.special import j0
import time
import numba as nb

q = np.linspace(0.03, 1.0, 1000)

start = time.time()

def f(z, t, q):
    return t * 0.5 * (erf((t - z) / 3) - 1) * j0(q * t) * np.exp(-0.5 * ((z - 40) / 2) ** 2)

ans2 = integrate.quad_vec(lambda t: integrate.quad_vec(lambda z: f(z,t,q),10, 60)[0],0,50)[0]

end = time.time()
print(end - start)
# 输出：2.87s

scipy.quad_vec方案是三者中速度最快的。

待解决问题 #

我想知道如何能让scipy.quad_vec的速度进一步提升？遗憾的是scipy.quad_vec不支持低级可调用函数，请问有没有办法能让scipy.quad或scipy.dblquad实现和scipy.quad_vec同等效率的向量化？或者有没有其他可用的优化方案？

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