如何向量化scipy.integrate.nquad被积函数以适配qmc_quad？

现有使用scipy.integrate.nquad计算积分的代码，为提升运算速度想改用scipy.integrate.qmc_quad，但qmc_quad要求被积函数支持向量化，当前函数仅能单次计算一个参数组合，需要对被积函数进行向量化改造以适配qmc_quad。

原代码如下：

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import nquad, qmc_quad


X = [11.3, 14.8, 7.6, 10.5, 12.7, 3.9, 11.2, 5.4, 8.5, 5.0, 4.4, 7.3, 2.9, 5.7, 6.2, 7.3, 3.3, 4.2, 5.5, 4.2]
N = len(X)
I = np.arange(1, N + 1)


def pdf_normal(m1, mn, s1, sn, i, n, x):
    m = (m1 * (n - i) / (n - 1)) + (mn * (i - 1) / (n - 1))
    s = (s1 * (n - i) / (n - 1)) + (sn * (i - 1) / (n - 1))
    y = (x - m) / s
    f = np.exp(-0.5 * y * y) / np.sqrt(2 * np.pi)
    pdf = f / s
    return pdf


def function_g_vec(args):
    [m1, mn, s1, sn] = args
    xi = X
    i = I
    n = N
    values = [pdf_normal(m1, mn, s1, sn, i_val, n, x_val) for i_val, x_val in zip(i, xi)]
    result = np.prod(values)
    return result


limits = [(-10, 30), (-10, 20), (0, 15), (0, 10)]
K = nquad(lambda *x: function_g_vec(x), ranges=limits)
print(K)
# K = 2.9335981345167932e-18
# error_K = 2.9282172140557324e-18

核心思路是让被积函数能够处理批量输入的参数组合（即输入是二维数组，每行是一组[m1, mn, s1, sn]），利用numpy的广播机制实现向量化运算，避免循环。

步骤1：向量化改造pdf_normal函数 #

原函数只能处理单个i和x，改造后让它能同时处理所有i和x的组合，并且支持批量参数输入：

def pdf_normal_vec(m1, mn, s1, sn, i, n, x):
    # m1, mn, s1, sn: 形状为(M,)的数组（M是参数组合数量）
    # i, x: 形状为(N,)的数组（N是样本数量）
    # 利用广播扩展维度，让参数与样本维度匹配
    m = (m1[:, None] * (n - i) / (n - 1)) + (mn[:, None] * (i - 1) / (n - 1))
    s = (s1[:, None] * (n - i) / (n - 1)) + (sn[:, None] * (i - 1) / (n - 1))
    y = (x - m) / s
    f = np.exp(-0.5 * y ** 2) / np.sqrt(2 * np.pi)
    pdf = f / s
    return pdf  # 返回形状为(M, N)的数组，每行对应一组参数的所有样本pdf值

步骤2：改造被积函数适配qmc_quad #

让它接受二维数组输入（每行一组参数），并对每组参数计算乘积：

def function_g_qmc(params):
    # params: 形状为(M, 4)的数组，M是批量处理的参数组合数
    m1, mn, s1, sn = params.T  # 拆分出每个参数的数组，形状均为(M,)
    xi = np.array(X)
    i_arr = np.array(I)
    n = N
    # 获取所有参数组合对应的样本pdf值，形状(M, N)
    pdf_vals = pdf_normal_vec(m1, mn, s1, sn, i_arr, n, xi)
    # 对每行（每组参数）计算乘积，得到形状(M,)的结果数组
    result = np.prod(pdf_vals, axis=1)
    return result

步骤3：调用qmc_quad计算积分 #

qmc_quad要求积分限是一个列表，每个元素是对应维度的(low, high)，直接使用原有的limits即可：

# 调用qmc_quad，n_points控制采样点数量（值越大精度越高，速度越慢）
K_qmc, error_qmc = qmc_quad(function_g_qmc, limits=limits, n_points=10000)
print(f"QMC积分结果: {K_qmc}")
print(f"积分误差估计: {error_qmc}")

完整改造后代码 #

import numpy as np
from scipy.integrate import nquad, qmc_quad


X = [11.3, 14.8, 7.6, 10.5, 12.7, 3.9, 11.2, 5.4, 8.5, 5.0, 4.4, 7.3, 2.9, 5.7, 6.2, 7.3, 3.3, 4.2, 5.5, 4.2]
N = len(X)
I = np.arange(1, N + 1)


def pdf_normal_vec(m1, mn, s1, sn, i, n, x):
    m = (m1[:, None] * (n - i) / (n - 1)) + (mn[:, None] * (i - 1) / (n - 1))
    s = (s1[:, None] * (n - i) / (n - 1)) + (sn[:, None] * (i - 1) / (n - 1))
    y = (x - m) / s
    f = np.exp(-0.5 * y ** 2) / np.sqrt(2 * np.pi)
    pdf = f / s
    return pdf


def function_g_qmc(params):
    m1, mn, s1, sn = params.T
    xi = np.array(X)
    i_arr = np.array(I)
    n = N
    pdf_vals = pdf_normal_vec(m1, mn, s1, sn, i_arr, n, xi)
    return np.prod(pdf_vals, axis=1)


# 原nquad计算（用于对比）
def function_g_vec(args):
    [m1, mn, s1, sn] = args
    xi = X
    i = I
    n = N
    values = [pdf_normal_vec(np.array([m1]), np.array([mn]), np.array([s1]), np.array([sn]), i_val, n, x_val)[0,0] for i_val, x_val in zip(i, xi)]
    result = np.prod(values)
    return result

limits = [(-10, 30), (-10, 20), (0, 15), (0, 10)]
K_nquad, error_nquad = nquad(lambda *x: function_g_vec(x), ranges=limits)
print(f"Nquad积分结果: {K_nquad}")
print(f"Nquad积分误差估计: {error_nquad}")

# QMC计算
K_qmc, error_qmc = qmc_quad(function_g_qmc, limits=limits, n_points=10000)
print(f"\nQMC积分结果: {K_qmc}")
print(f"QMC积分误差估计: {error_qmc}")

关键说明 #

• 利用numpy的广播机制：通过[:, None]给参数数组增加一个维度，让它能与样本数组（N维）进行逐元素运算，避免循环。
• qmc_quad的输入要求：被积函数必须接受形状为(M, d)的二维数组（d是积分维度，这里是4），返回形状为(M,)的数组，每组参数对应一个函数值。
• 精度与速度：n_points参数控制采样点数量，数值越大精度越高但速度越慢，可根据需求调整。

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