Brute force bifurcation diagram - need help with stiff solver

[sdwfrost]

Hi All,

I’m trying to reproduce a bifurcation diagram (brute force) computed using XPP-AUT. My code is here and the original XPP code from Krylova and Earn (2013) is here.

The XPP code seems to give a nice bifurcation plot:

However, I’ve tried multiple solvers, and many of them fail with this very stiff problem. Here’s the best I’ve done so far, using a PositiveDomain callback:

You can see gaps in the paths which aren’t present in the XPP code (which doesn’t throw any errors). Any idea of how to potentially improve the Julia code to be more robust?

To save you a trip to GitHub, here’s the Julia code:

using ModelingToolkit
using OrdinaryDiffEq
using DiffEqCallbacks
using DataFrames
using CSV

@parameters t R₀ γ μ a
@variables S(t) I(t) β(t)
D = Differential(t)

eqs = [D(S) ~ μ - β*S*I - μ*S,
       D(I) ~ β*S*I - (γ+μ)*I,
       β ~ R₀*(γ+μ)*(1+a*cos(2*π*t))]

@named sys = ODESystem(eqs)
simpsys = structural_simplify(sys)

p = [μ => 0.02, γ => 28.08, R₀ => 17.0, a => 0.08]
u₀ = [S => 0.9, I => 0.001]
tmin = 0.0
tmax = 650
transient = 600
strobe = 1.0
tspan = (tmin, tmax)
R0vec = collect(1.0:0.01:30.0)
solver = Tsit5()

prob = ODEProblem(simpsys, u₀, tspan, p; jac=true)

function prob_func(prob, i, repeat)
    remake(prob, p=[μ => 0.02, γ => 28.08, R₀ => R0vec[i], a => 0.08])
end

function output_func(sol, i)
    strobetimes = collect(transient:strobe:tmax)
    df = DataFrame(sol(strobetimes))
    rename!(df,[:t, :S, :I])
    df[!,"R0"] = fill(R0vec[i],length(strobetimes))
    df[!,"LogS"] = log10.(abs.(df.S))
    df[!,"LogI"] = log10.(abs.(df.I))
    return (df, false)
end

ensemble_prob = EnsembleProblem(prob,
                                prob_func = prob_func,
                                output_func = output_func)

@time sim = solve(ensemble_prob,
                  solver,
                  EnsembleThreads(),
                  trajectories = length(R0vec);
                  callback = PositiveDomain())

results = vcat(sim...)

CSV.write("bruteforceSIR_R0.dat", results; delim=" ", header=false)

using StatsPlots
using LaTeXStrings

p = @df results scatter(:R0,
                    :LogI,
                    xlabel=L"R_0",
                    ylabel=L"log_{10} I",
                    markersize=1.0,
                    color=:gray,
                    legend=false,
                    ylims=(-6,-2))
savefig(p, "bruteforceSIR_R0_statsplots.png")

[Oscar_Smith]

if this is a stiff problem, why are you using tsit5? does Rodas5P do better?

[sdwfrost]

I get lots of

Warning: Instability detected. Aborting

with Rodas5P (without changing anything else). Hence, I get a plot like this:

[Oscar_Smith]

That makes very little sense. What about Rosenbrock23 or TRBDF2?

[sdwfrost]

They’re worse! Here’s TRBDF2 (again, using default settings as above). I agree it doesn’t make sense…

Just for fun, here’s RK4

Agreed that this doesn’t make a great deal of sense on first impression.

[rveltz]

Did you try lsoda or cvode_bdf? Although i would be surprised if juliabdoes not beat them; i would expext xpp to use one of them

[rveltz]

I would also say that you need multiple trials per R0

[Sevi]

I think the problem is the behavior of PositiveDomain. Its documentations says

It reduces the next step by a certain scale factor until the extrapolated value at the next
time point is non-negative with a certain tolerance. Extrapolations are cheap to compute but might be inaccurate, so if a time step is changed it is additionally reduced by a safety factor of 0.9. Since extrapolated values are only non-negative up to a certain tolerance and in addition actual calculations might lead to negative values, also any negative values at the current time point are set to 0.

The last point is crucial. We want to see very small concentrations up to 1e-6 or so, but looking at the output data, they are simply set to zero by the callback.

Using isoutofdomain=(u,p,t) -> any(x -> x < 0, u) instead of callback = PositiveDomain() looks a lot better, filling out the “gaps” in the diagram. I also lowered the tolerances to reltol=1e-11, abstol=1e-11 just to be sure, but that’s probably too low (and isoutofdomain should have a similar effect).

UPDATE: A combination of low abstol and the isoutofdomain check seems to be enough – no need for such a low reltol which just slows everything down.

PS: Lowering the tolerance inside PositiveDomain might also have the same effect, didn’t try that yet.

[sdwfrost]

@Sevi

eltol=1e-11, abstol=1e-11

Thanks! This looks a lot better (though not as good as the XPP code). I’ll dig in to the settings in XPP (I think it uses a Dormand-Prince scheme as default).

Any idea why ‘non-stiff’ solvers do better here? The system has an extreme transient at the beginning, which I can probably overcome with some fiddly initial condition wrangling (but this doesn’t seem necessary in XPP).

Here’s the code so far with a non-stiff solver (RK4)

using ModelingToolkit
using OrdinaryDiffEq
using DataFrames
using CSV

@parameters t R₀ γ μ a
@variables S(t) I(t) β(t)
D = Differential(t)

eqs = [D(S) ~ μ - β*S*I - μ*S,
       D(I) ~ β*S*I - (γ+μ)*I,
       β ~ R₀*(γ+μ)*(1+a*cos(2*π*t))]

@named sys = ODESystem(eqs)
simpsys = structural_simplify(sys)

p = [μ => 0.02, γ => 28.08, R₀ => 17.0, a => 0.08]
u₀ = [S => 0.9, I => 0.001]
tmin = 0.0
tmax = 650
transient = 600
strobe = 1.0
tspan = (tmin, tmax)
R0vec = collect(1.0:0.01:30.0)
solver = RK4()
tol = 1e-11

prob = ODEProblem(simpsys, u₀, tspan, p; jac=true) 

function prob_func(prob, i, repeat)
    remake(prob, p=[μ => 0.02, γ => 28.08, R₀ => R0vec[i], a => 0.08])
end

function output_func(sol, i)
    strobetimes = collect(transient:strobe:tmax)
    df = DataFrame(sol(strobetimes))
    rename!(df,[:t, :S, :I])
    df[!,"R0"] = fill(R0vec[i],length(strobetimes))
    df[!,"LogS"] = log10.(abs.(df.S))
    df[!,"LogI"] = log10.(abs.(df.I))
    return (df, false)
end

ensemble_prob = EnsembleProblem(prob,
                                prob_func = prob_func,
                                output_func = output_func)

@time sim = solve(ensemble_prob,
                  solver,
                  EnsembleThreads(),
                  trajectories = length(R0vec);
                  isoutofdomain=(u,p,t) -> any(x -> x < 0, u),
                  abstol=tol)

results = vcat(sim...)

CSV.write("bruteforceSIR_R0.dat", results; delim=" ", header=false)

using StatsPlots
using LaTeXStrings

p = @df results scatter(:R0,
                    :LogI,
                    xlabel=L"R_0",
                    ylabel=L"log_{10} I",
                    markersize=1.0,
                    color=:gray,
                    legend=false,
                    ylims=(-6,-2))
savefig(p, "bruteforceSIR_R0_statsplots.png")

[sdwfrost]

I did some more investigation, and I see what’s going on with the stiff solvers; if you set a=0, the system has two stable states, one with S=1 and I=0, and one with S<1 and I>0. Some of the stiff solvers seem to get stuck in the first, while others fail.

Here’s some example code with different solvers.

using ModelingToolkit
using OrdinaryDiffEq
using LSODA
using Sundials

@parameters t R₀ γ μ a
@variables S(t) I(t) β(t)
D = Differential(t)

eqs = [D(S) ~ μ - β*S*I - μ*S,
       D(I) ~ β*S*I - (γ+μ)*I,
       β ~ R₀*(γ+μ)*(1+a*cos(2*π*t))]

@named sys = ODESystem(eqs)
simpsys = structural_simplify(sys)

u₀ = [S => 0.9, I => 0.001]
tmin = 0.0
tmax = 650
tol = 1e-11
maxiters = 1e7

p₀ = [μ => 0.02, γ => 28.08, R₀ => 17.0, a => 0.0]

prob = ODEProblem(simpsys, u₀, (tmin, tmax), p₀; jac=true)
sol_rk4 = solve(prob, RK4();
                maxiters = maxiters,
                isoutofdomain=(u,p,t) -> any(x -> (x <= 0 || x >= 1), u),
                abstol=tol)

sol_lsoda = solve(prob, lsoda();
                maxiters = maxiters,
                abstol=tol)

sol_rosenbrock23 = solve(prob, Rosenbrock23();
                maxiters = maxiters,
                isoutofdomain=(u,p,t) -> any(x -> (x <= 0 || x >= 1), u),
                abstol = tol)

sol_cvodebdf = solve(prob, CVODE_BDF())

Here’s the output of the final state.

julia> print(sol_rk4.u[end])
[0.05882344413440804, 0.0006699054322332516]
julia> print(sol_lsoda.u[end])
[0.05882352941176473, 0.0006698764915218742]
julia> print(sol_rosenbrock23.u[end])
[0.9999994236820176, 1.604027727172643e-47]

CVODE_BDF is interesting. It works if I don’t set the tolerance to a low value.

julia> sol_cvodebdf = solve(prob, CVODE_BDF())
retcode: Success
Interpolation: 3rd order Hermite
t: 2679-element Vector{Float64}:
   0.0
   9.365055056001597e-5
   0.00018730110112003193
   0.00045236993728053627
   ⋮
 649.6776241741102
 649.912024827687
 650.0
u: 2679-element Vector{Vector{Float64}}:
 [0.9, 0.001]
 [0.8999583517992265, 0.001039101101715838]
 [0.8999150700341084, 0.0010797289334348362]
 [0.8997831768370831, 0.0012034963123768742]
 ⋮
 [0.05857054135397698, 0.0006870603487333439]
 [0.05853834794611831, 0.0006657353926255108]
 [0.05856604433672487, 0.0006581134035398938]

but fails if I drop the tolerance.

julia> sol_cvodebdf = solve(prob, CVODE_BDF(); abstol=1e-11)

[CVODES WARNING]  CVode
  Internal t = 5.43326 and h = 2.9051e-16 are such that t + h = t on the next step. The solver will continue anyway.


[CVODES WARNING]  CVode
  Internal t = 5.43326 and h = 2.9051e-16 are such that t + h = t on the next step. The solver will continue anyway.


[CVODES WARNING]  CVode
  Internal t = 5.43326 and h = 2.9051e-16 are such that t + h = t on the next step. The solver will continue anyway.


[CVODES WARNING]  CVode
  Internal t = 5.43326 and h = 2.9051e-16 are such that t + h = t on the next step. The solver will continue anyway.


[CVODES WARNING]  CVode
  Internal t = 5.43326 and h = 1.79233e-16 are such that t + h = t on the next step. The solver will continue anyway.


[CVODES WARNING]  CVode
  Internal t = 5.43326 and h = 1.79233e-16 are such that t + h = t on the next step. The solver will continue anyway.


[CVODES WARNING]  CVode
  Internal t = 5.43326 and h = 1.79233e-16 are such that t + h = t on the next step. The solver will continue anyway.


[CVODES WARNING]  CVode
  Internal t = 5.43326 and h = 1.79233e-16 are such that t + h = t on the next step. The solver will continue anyway.


[CVODES WARNING]  CVode
  Internal t = 5.43326 and h = 1.11224e-16 are such that t + h = t on the next step. The solver will continue anyway.


[CVODES WARNING]  CVode
  Internal t = 5.43326 and h = 1.11224e-16 are such that t + h = t on the next step. The solver will continue anyway.


[CVODES WARNING]  CVode
  The above warning has been issued mxhnil times and will not be issued again for this problem.

retcode: MaxIters
Interpolation: 3rd order Hermite
t: 100001-element Vector{Float64}:
 0.0
 6.622041520638907e-5
 0.00013244083041277815
 0.0003449609238189219
 ⋮
 5.433256962605888
 5.433256962605888
 5.433256962605888
u: 100001-element Vector{Vector{Float64}}:
 [0.9, 0.001]
 [0.8999708850285875, 0.001027335790869192]
 [0.8999409716210338, 0.001055417802667344]
 [0.899839159128681, 0.0011509726910748413]
 ⋮
 [6.134513963058041e152, -6.13451396305804e152]
 [6.134513963058041e152, -6.13451396305804e152]

If I change R0 to 2.0 (which should make the system less stiff), then lsoda and CVODE_BDF fails too:

julia> print(sol_rk4.u[end])
[0.5000006296727381, 0.00035587073771302234]
julia> print(sol_lsoda.u[end])
[NaN, NaN]
julia> print(sol_rosenbrock23.u[end])
[0.9999995646381588, 5.974907074363282e-99]
julia> print(sol_cvodebdf.u[end])
[1.7885027908424116e153, -1.788502790842412e153]

[sdwfrost]

I still don’t know why the stiff solvers can’t handle the transient properly, but with initial conditions closer to the equilibrium cycle, stiff solvers like Rosenbrock23 now give better results (though why it’s not working still for low R0 is odd to me). Code below.

using ModelingToolkit
using OrdinaryDiffEq
using DataFrames
using CSV

@parameters t R₀ γ μ a
@variables S(t) I(t) β(t)
D = Differential(t)

eqs = [D(S) ~ μ - β*S*I - μ*S,
       D(I) ~ β*S*I - (γ+μ)*I,
       β ~ R₀*(γ+μ)*(1+a*cos(2*π*t))]

@named sys = ODESystem(eqs)
simpsys = structural_simplify(sys)

S₀ = 1/R₀
I₀ = (μ/(μ+γ))*(1-S₀)
u₀ = [S => S₀, I => I₀]
p = [μ => 0.02, γ => 28.08, R₀ => 17.0, a => 0.08]

tmin = 0.0
tmax = 650
transient = 600
strobe = 1.0
tspan = (tmin, tmax)
R0vec = collect(1.0:0.01:30.0)
solver = Rosenbrock23()
tol = 1e-11
maxiters = 1e7

prob = ODEProblem(simpsys, u₀, tspan, p; jac=true) 

# Utility functions

function prob_func(prob, i, repeat)
    remake(prob, p=[μ => 0.02, γ => 28.08, R₀ => R0vec[i], a => 0.08])
end

function output_func(sol, i)
    strobetimes = collect(transient:strobe:tmax)
    df = DataFrame(sol(strobetimes))
    rename!(df,[:t, :S, :I])
    df[!,"R0"] = fill(R0vec[i],length(strobetimes))
    df[!,"LogS"] = log10.(abs.(df.S))
    df[!,"LogI"] = log10.(abs.(df.I))
    return (df, false)
end

# Run ensemble

ensemble_prob = EnsembleProblem(prob,
                                prob_func = prob_func,
                                output_func = output_func)

@time sim = solve(ensemble_prob,
                  solver,
                  EnsembleThreads(),
                  trajectories = length(R0vec);
                  maxiters = maxiters,
                  isoutofdomain=(u,p,t) -> any(x -> x < 0, u),
                  abstol=tol)

# Output

results = vcat(sim...)
CSV.write("bruteforceSIR_R0.dat", results; delim=" ", header=false)

# Plotting

using StatsPlots
using LaTeXStrings

p = @df results scatter(:R0,
                    :LogI,
                    xlabel=L"R_0",
                    ylabel=L"log_{10} I",
                    markersize=1.0,
                    color=:gray,
                    legend=false,
                    ylims=(-6,-2))
savefig(p, "bruteforceSIR_R0_statsplots.png")

[ChrisRackauckas]

@sdwfrost

I still don’t know why the stiff solvers can’t handle the transient properly

They are not stability bound but tolerance bound, so they might be taking large steps where the difference between two trajectories is within the tolerances. RK methods are forced to take really small steps and thus solve below the required error tolerance.

[rveltz]

I found a trick for a different context. Use I=exp(Ilog) and

function sir!(du, X, p, t = 0)
    @unpack R₀, γ, μ, a = p
    S, Ilog, u, v = X
    I = exp(Ilog)
    β = R₀ * (γ + μ) * (1 + a * cos(2pi*t)
    du[1] = μ - β * S * I - μ * S
    du[2] = β * S - (γ + μ)
    du
end