在DifferentialEquations.jsl中减少内存分配

Question

我正在使用DifferentialEquations.jsl来求解一个ODE系统，如下所示。结果实际上并不相关，因为p只包含用于生成MWE的测试参数，但关键是，尽管使用了就地ODE函数，我还是看到了大量的内存分配。

using DifferentialEquations
function ode_fun!(du,u,p,t)
a,b,c,d,e = p
X = @. u[1] * a * ((b-c)/b)
Y = @. u[2] * d * ((b-e)/b)
du[1] = -sum(X) + sum(Y) - u[1]*u[2]
du[2] = sum(X) - sum(Y) - u[1]*u[2]
end
#exemplary parameters 
a = collect(10:-0.1:0.1)
b = a.^2
c = b*0.7
d = collect(0.01:0.01:1)
e = b*0.3
u0 = [1.0, 0.5]
p = [a,b,c,d,e]
tspan = [0.0, 100.0]
t = collect(0:0.01:100) 
prob = ODEProblem(ode_fun!,u0,tspan,p,saveat=t) 
@time sol = solve(prob)
1.837609 seconds (5.17 M allocations: 240.331 MiB, 2.31% gc time) #Julia 1.5.2

由于我需要反复解决这个ODE系统，我想尽可能减少分配，并想知道是否可以对此做些什么。我一直想知道问题是否出在X和Y身上，并试图在ODE函数之外预先分配它们，但不幸的是，未能成功减少分配。

Answer 1

我很确定这应该更快，并且是分配的一半

function ode_fun!(du,u,p,t)
a,b,c,d,e = p
XmY = @. u[1] * a * (1-c/b) - u[2] * d * (1-e/b)
sXmY = sum(XmY)
du[1] = -sXmY - u[1]*u[2]
du[2] = sXmY - u[1]*u[2]
end

可能有一种方法可以消除所有这些，但我不是DifferentialEquations专家。

Answer 2

序言：我意识到这是一个老问题，询问者可能已经学会了如何避免大量分配。这个答案是针对一个仍然在帖子中犯错误的读者。

我相信你看到的分配(以及放缓(主要来自汇编。当我在Julia 1.9.1上运行您的脚本两次时，我得到¹

julia> include("your_script.jl")
2.287537 seconds (3.33 M allocations: 216.865 MiB, 3.31% gc time, 99.91% compilation time)
julia> include("your_script.jl")
0.069226 seconds (46.83 k allocations: 4.283 MiB, 97.70% compilation time: 100% of which was recompilation)

在第一次调用时，Julia必须编译运行代码所需的所有函数，包括DifferentialEquations.jl中的函数(请参阅性能提示中的本节(。这主导了代码的实际运行时间，而实际运行时间只需要几分之一秒。在随后的调用中，Julia只需要执行明显更快的重新编译(在这种情况下，它仍然主导着运行时(。

第二个重大错误是性能关键代码应该在函数内部。简单地将代码封装在一个函数中(下面代码中的main(可以显著节省时间和分配：

# new_script.jl
using DifferentialEquations
function ode_fun!(du,u,p,t)
# ...
end
function main()
# rest of the code
end

julia> include("new_script.jl")
main (generic function with 1 method)
julia> main()
2.334540 seconds (3.33 M allocations: 216.774 MiB, 2.50% gc time, 99.92% compilation time)
julia> main()
0.001639 seconds (10.94 k allocations: 2.260 MiB)    # 40x speedup

和以前一样，我们必须让Julia编译代码，然后才能获得实际的运行时间。

重复使用阵列

最后，如果您需要多次运行解算器，那么预分配数组并重复使用会有所帮助，如下所示。与上述变化相比，这导致相对较小的节省。

using DifferentialEquations
function ode_fun!(du,u,p,t)
a,b,c,d,e,X,Y = p              # X, Y from p
@. X = u[1] * a * ((b-c)/b)    # avoid allocating new X, Y
@. Y = u[2] * d * ((b-e)/b)
du[1] = -sum(X) + sum(Y) - u[1]*u[2]
du[2] = sum(X) - sum(Y) - u[1]*u[2]
end
function main()
#exemplary parameters
a = collect(10:-0.1:0.1)
b = a.^2
c = b*0.7
d = collect(0.01:0.01:1)
e = b*0.3
X = similar(a)             # preallocate X, Y
Y = similar(d)
u0 = [1.0, 0.5]
p = [a,b,c,d,e,X,Y]        # pass X, Y to solver through p
tspan = [0.0, 100.0]
t = collect(0:0.01:100)
prob = ODEProblem(ode_fun!,u0,tspan,p,saveat=t)
function solve_many_times()
for i in 1:10000
sol = solve(prob)
end
end
@time solve_many_times()
return
end

---

1. 注意，在Julia的当前版本中，@time还报告了编译所花费的时间百分比