有没有办法在整组值上并行化 scipy.integrate 函数？

Question

我正在使用计算速度较慢的似然函数计算某些参数 $ heta$ 的后验概率，我想知道是否有办法加快速度。 我的代码的缓慢部分是卡方的计算，它提供了需要为 z_sn 中的每个样本计算的积分

import scipy.constants as cte
from scipy.integrate import quad

def luminosity_integrand(z, omgM):
    Ez = np.sqrt((1 - omgM) + omgM * np.power(1 + z, 3))
    return 1. / Ez

def luminosity_distance(z, h, omgM):
    integral, _ = quad(luminosity_integrand, 0, z, epsrel=1e-8, args=(omgM))
    return (cte.c / 10. ** 5) / h * (1 + z) * integral

def distance_modulus(z, h, omgM):
    return 5. * np.log10(luminosity_distance(z, h, omgM)) + 25.




def chisq_sn(h, omgM):
    m_model = np.array([distance_modulus(z, h, omgM) for z in z_sn])
    diffs = m_obs-m_model
 
    
    maha_distances = np.dot(np.dot(diffs,inv_cov_plus),diffs) #mahalanobis distance
    return maha_distances

我使用了列表理解，但我不知道这是否是最快的路线。我还读过一些关于使用 cython 的稀疏评论，但我从未使用过它。如果重要的话，我正在使用 emcee 来计算后验

Answer 1

如果您不介意使用私有函数（该函数不一定在 SciPy 的未来版本中可用并且不受官方支持），SciPy 1.12.0 中有一个矢量化积分器，可以提供 10 倍的加速。

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import quad
from scipy.integrate._tanhsinh import _tanhsinh

def luminosity_integrand(z, omgM):
    Ez = np.sqrt((1 - omgM) + omgM * np.power(1 + z, 3))
    return 1. / Ez

z_sn = np.linspace(0, 10, 1000)
omgM = 0.5

%timeit [quad(luminosity_integrand, 0, z, args=(omgM))[0] for z in z_sn]
# 146 ms ± 23.3 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit _tanhsinh(luminosity_integrand, 0, z_sn, args=(omgM,)).integral
# 10.4 ms ± 1.35 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

luminosity1 = [quad(luminosity_integrand, 0, z, args=(omgM))[0] for z in z_sn]
luminosity2 = _tanhsinh(luminosity_integrand, 0, z_sn, args=(omgM,)).integral

np.testing.assert_allclose(luminosity1, luminosity2)

plt.plot(z_sn, luminosity1, '-', label='quad')
plt.plot(z_sn, luminosity2, '--', label='tanh-sinh')
plt.xlabel('z_sn')
plt.ylabel('luminosity')
plt.legend()