求与Python(scipy.integrate)速度相当的三重积分Julia代码

问题描述 #

需要计算积分范围包含其他变量的三重积分，当前编写的Julia代码比Python版本慢1000倍（单行形式慢10000倍）。

Python参考代码 #

import scipy.integrate as spi

integrand = lambda x,y,z: x**2 + 2*y + z**2

result, err = spi.quad(lambda z: spi.quad(lambda y: spi.quad(lambda x: integrand(x, y, z), 0, y**2)[0], 0, 10)[0], 2, 4)

原Julia代码（性能不佳版本） #

using QuadGK

integrand0 = (x,y,z) -> x^2 + 2*y + z^2

integrand1 = (y,z) -> quadgk((x) -> integrand0(x,y,z), 0,y^2)[1]
    
integrand2 = (z) -> quadgk((y) -> integrand1(y,z), 0,10)[1]

result, err = quadgk((z) -> integrand2(z), 2, 4)

单行版Julia代码 #

result, err = quadgk((z)->quadgk((y)-> quadgk((x)->integrand0(x,y,z), 0,y^2)[1], 0,10)[1], 2, 4)

性能差异原因 #

1. JIT编译延迟：Julia是即时编译语言，第一次执行代码时会编译所有涉及的函数，这个编译时间会被计入首次运行耗时，而Python是解释型语言无此开销。仅测试单次运行会严重低估Julia的实际运行速度。
2. 嵌套匿名函数的类型不稳定：原代码大量使用嵌套匿名函数，这类函数捕获外部变量时容易导致类型不稳定，JIT编译器无法充分优化，进而拖慢运行效率。

优化后的Julia代码 #

方案1：优化QuadGK调用（消除类型不稳定） #

将匿名函数改为显式命名函数，帮助编译器进行类型推断和优化，同时规避首次编译时间的干扰（运行多次取平均）：

using QuadGK

# 定义为普通函数，而非匿名函数，提升类型稳定性
integrand(x, y, z) = x^2 + 2y + z^2

# 显式定义内层积分函数，便于编译器优化
function inner_integral(y, z)
    return quadgk(x -> integrand(x, y, z), 0, y^2)[1]
end

function middle_integral(z)
    return quadgk(y -> inner_integral(y, z), 0, 10)[1]
end

# 第一次运行包含编译时间，第二次及之后为实际运行速度
@time result, err = quadgk(z -> middle_integral(z), 2, 4)
@time result, err = quadgk(z -> middle_integral(z), 2, 4)

方案2：使用HCubature.jl（专为多重积分设计） #

对于多重积分，HCubature.jl比嵌套调用单变量积分工具更高效，语法也更简洁：

using HCubature

# 将积分变量打包为数组，适配HCubature的接口
integrand(x) = x[1]^2 + 2x[2] + x[3]^2

# 定义积分区域：x∈[0, y²], y∈[0,10], z∈[2,4]
# 对依赖于其他变量的维度，传入函数作为上下限
result, err = hcubature(integrand, (0, 0, 2), (y -> y^2, 10, 4))

# 忽略首次运行的编译时间
@time result, err = hcubature(integrand, (0, 0, 2), (y -> y^2, 10, 4))

测试说明 #

运行Julia代码时，务必忽略第一次运行的耗时，第二次及之后的运行时间才是实际执行速度。优化后的代码在多次运行后，性能会显著超过Python版本，符合Julia的速度优势特性。

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