Help with speeding up large ODE system from ModelingToolkit

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")
1 Like

This is something that’s being improved by the structure retention. There will be a new module added to the system hopefully quite soon specifically for PDE lowering (it’s already written and being tested)

Dear Chris,

Thanks! Can you ping here once the module has been committed? A system with K=80 has the following timings (most of which spent in structural_simplify:

2655.486365 seconds (20.54 G allocations: 1.318 TiB, 18.86% gc time, 1.03% compilation time)

That’s a lot of allocations! Anything to improve this would be much appreciated; the above is quite a simple/relatively small system of its type.

Best wishes
Simon

PS. How can I use GPU to speed up the solver (only for larger systems) when the model is generated using ModelingToolkit (which has Dicts for initial conditions and parameters)?