Hi everyone,
I’m trying to speed up the following ODE system, which is a backward-in-space discretization of a set of PDEs, which I’m generating using ModelingToolkit.jl. The bottleneck is in the construction of the system, which I’d like to keep as intact as possible due to readability. I’d like to increase the size of the system (to K=80, da=52.0) in the below, but the existing code is already quite slow. I’m using Julia 1.9, so I was hoping that the time of first execution wasn’t a huge hit, as it’s a bit clumsy to create a small system just to warm up the functions, but it still seems to be necessary to get the performance.
The timings on my M1 Mac (on first execution) are:
Model construction: 237.319118 seconds (1.73 G allocations: 99.757 GiB, 13.78% gc time, 11.59% compilation time)
ODEProblem: 18.548050 seconds (84.13 M allocations: 4.451 GiB, 11.38% gc time, 58.80% compilation time)
ODESolution: 602.915088 seconds (277.28 M allocations: 12.561 GiB, 1.28% gc time, 100.00% compilation time)
On second execution, I get the following:
Model construction: 197.906369 seconds (1.68 G allocations: 96.424 GiB, 8.34% gc time)
ODEProblem: 7.157878 seconds (54.68 M allocations: 2.833 GiB, 7.54% gc time)
ODESolution: 0.000588 seconds (579 allocations: 750.375 KiB)
Here’s the complete code:
using ModelingToolkit
using OrdinaryDiffEq
using Plots
function cohort_age_structured_sir(;K, S₀, I₀, R₀, B, C, G, DA)
@parameters t β c[1:K] γ Δa
@variables (S(t))[1:K] (I(t))[1:K] (R(t))[1:K] (N(t))[1:K] (λ(t))[1:K] (p(t))[1:K,1:K]
Dt = Differential(t)
eqs = [ Dt(S[1]) ~ -λ[1]*S[1],
[(S[i] - S[i-1])/Δa + Dt(S[i]) ~ -λ[i]*S[i] for i in 2:K]...,
Dt(I[1]) ~ λ[1]*S[1]-γ*I[1],
[(I[i] - I[i-1])/Δa + Dt(I[i]) ~ λ[i]*S[i]-γ*I[i] for i in 2:K]...,
Dt(R[1]) ~ γ*I[1],
[(R[i] - R[i-1])/Δa + Dt(R[i]) ~ γ*I[i] for i in 2:K]...,
[N[i] ~ S[i]+I[i]+R[i] for i in 1:K]...,
[λ[i] ~ sum([β*c[i]*p[i,j]*I[j]/N[j] for j in 1:K]) for i in 1:K]...,
[p[i,j] ~ c[j]*N[j]/sum([c[k]*N[k] for k in 1:K]) for j in 1:K for i in 1:K]...
]
@named sys = ODESystem(eqs);
sys = structural_simplify(sys)
u0 = [[S[i] => S₀[i] for i in 1:K]...,
[I[i] => I₀[i] for i in 1:K]...,
[R[i] => R₀[i] for i in 1:K]...];
p = [β=>B, [c[i]=>C[i] for i in 1:K]..., γ=>G, Δa => DA];
return sys, u0, p
end
K = 40 # Change K to change the size of the system
δa = 80*52/K
δt = 0.1
tmax = 52.0
tspan = (0.0,tmax)
S₀ = [990.0/K for i in 1:K]
I₀ = [10.0/K for i in 1:K]
R₀ = [0.0 for i in 1:K]
β = 0.05
C = [10.0 for i in 1:K]
γ = 0.25
@time sys, u0, p = cohort_age_structured_sir(K=K,
S₀=S₀,
I₀=I₀,
R₀=R₀,
B=β,
C=C,
G=γ,
DA=δa);
@time prob = ODEProblem(sys, u0, tspan, p, sparse=true, jac=false);
@time sol = solve(prob, Tsit5(), saveat=δt);
all_states = states(sys);
indexof(sym,syms) = findfirst(isequal(sym), syms);
@parameters t
@variables (S(t))[1:K] (I(t))[1:K] (R(t))[1:K]
S_indexes = [indexof(S[k],all_states) for k in 1:K]
I_indexes = [indexof(I[k],all_states) for k in 1:K]
R_indexes = [indexof(R[k],all_states) for k in 1:K]
Smat = sol[S_indexes,:]
Imat = sol[I_indexes,:]
Rmat = sol[R_indexes,:]
Stotal = sum(Smat,dims=1)'
Itotal = sum(Imat,dims=1)'
Rtotal = sum(Rmat,dims=1)'
t = sol.t
plot(t, Stotal, label="S", xlabel="Time", ylabel="Number")
plot!(t, Itotal, label="I")
plot!(t, Rtotal, label="R")