StochasticDiffEq Tau-Leaping maxiters

I’m relatively new to Julia and JumpProcesses.jl, so apologies if this is somewhat basic, but I’m having an issue trying to convert the example SIR model with tau-leaping from density-dependent to frequency dependent i.e., divide infection generation by N.

using JumpProcesses, StochasticDiffEq, DifferentialEquations, LinearAlgebra
using Catalyst, Random

function rate(out, u, p, t)
    out[1] = p[1] * u[1] * u[2] / (u[1] + u[2] + u[3])           # β * X * Y / N
    out[2] = p[2] * u[2]                                         # ν * Y
    nothing
end

c = zeros(3, 2)
# X + Y --> Y
c[1, 1] = -1        # X -> X - 1
c[2, 1] = 1         # Y -> Y + 1

# Y --> Z
c[2, 2] = -1        # Y -> Y - 1
c[3, 2] = 1         # Z -> Z + 1

function change!(du, u, p, t, counts, mark)
    mul!(du, c, counts)
    nothing
end

rj = RegularJump(rate, change!, 2)

tmax = 100.0
δt = 0.001
tspan = (0.0, tmax)
tlength = length(tspan[1]:δt:tspan[2])
Random.seed!(1234);

μ = 0.001
u₀ = [1000.0, 50.0, 0.0]

β = 0.1/1000.0
γ = 0.01
p = (β, γ)

prob = DiscreteProblem(u₀, tspan, p)
jump_prob = JumpProblem(prob, Direct(), rj)
sol = solve(jump_prob, TauLeaping(); dt = δt)
sol_array = zeros(4, length(sol))
sol_array[1:3, :] = Array(sol)
sol_array[4, :] = sum(sol_array[1:3, :], dims = 1)

which produces the following error:

┌ Warning: Interrupted. Larger maxiters is needed. If you are using an integrator for non-stiff ODEs or an automatic switching algorithm (the default), you may want to consider using a method for stiff equations. See the solver pages for more details (e.g. https://docs.sciml.ai/DiffEqDocs/stable/solvers/ode_solve/#Stiff-Problems).
└ @ SciMLBase ~/.julia/packages/SciMLBase/KsLES/src/integrator_interface.jl:515
1×1 Matrix{Float64}:
 1050.0

If I change out[1] to remove u[3] it runs. It also runs if I add an N state to u₀ and add the corresponding jumps to rate() and RegularHump(). It also runs if I replace (u[1] + u[2] + u[3]) by 10000, so I’m not sure it’s an issue with how small the rate becomes.

Any help would be appreciated!

Shouldn’t c be 3x2 not 3x3?

Yes, good catch. That was a typo on my part in the question, but my code has c = zeros(3, 2) and still has the same error.

I’ve done a little more digging and it seems like it’s running into a problem when any of the states equal 0. Do you know why this would be the case, as the states all sum to > 0 at all time points.

This is obviously OK and can be worked around in a simple model as you can initialize with all states > 0, but I think was causing an issue when I added births and deaths into the model, as some states will get close to, if not equal to, 0. Do I need use callback functions to try and ensure the jumps never push a state equal or are less than 0 (therefore could divide by 0 in the next step)?

Could be adaptive step size selection is causing the issue, but @ChrisRackauckas could better comment on that than me. You could try passing adaptive=false to the call to solve and see if that gets rid of the error.

Thanks @isaacsas - adaptive = false does fix the issue.

This may be a separate issue, but is there are way to ensure the jumps never result in a state < 0? I have updated the code to include births and deaths, and whilst the code below works correctly, if you change out[1] = p[1] * u[1] * u[2] / (u[1] + u[2] + u[3]) , then it returns the X state as -2.0 at one point.

function rate(out, u, p, t)
    out[1] = p[1] * u[1] * u[2] / (u[1] + u[2] + u[3] + u[4])         # β * X * Y / N
    out[2] = p[2] * u[2]                        # σ * E
    out[3] = p[3] * u[3]                        # γ * Y
    out[4] = p[4] * (u[1] + u[2] + u[3] + u[4])        # μ * N
    out[5] = p[4] * u[1]                        # μ * X
    out[6] = p[4] * u[2]                        # μ * E
    out[7] = p[4] * u[3]                        # μ * Y
    out[8] = p[4] * u[4]                        # μ * Z
    nothing
end

c = zeros(4, 8)
# X + Y --> E
c[1, 1] = -1        # X -> X - 1
c[2, 1] = 1         # E -> E + 1

# E --> Y
c[2, 2] = -1        # E -> E - 1
c[3, 2] = 1         # Y -> Y + 1

# Y --> Z
c[3, 3] = -1        # Y -> Y - 1
c[4, 3] = 1         # Z -> Z + 1

# X --> 2X
c[1, 4] = 1         # X -> X + 1

# X --> 0X
c[1, 5] = -1        # X -> X - 1

# E --> 0E
c[2, 6] = -1        # X -> X - 1

# Y --> 0Y
c[3, 7] = -1        # Y -> Y - 1

# Z --> 0Z
c[4, 8] = -1        # Z -> Z - 1

function change!(du, u, p, t, counts, mark)
    mul!(du, c, counts)
    nothing
end

rj = RegularJump(rate, change!, 8)

tmax = 2.0
δt = 0.1/365
tspan = (0.0, tmax)
tlength = length(tspan[1]:δt:tspan[2])
Random.seed!(1234);


u₀ = [999.0, 10.0, 0.0, 0.0]

μ = 1
β = 1000
σ = 365/8
γ = 365/5
p = (β, σ, γ, μ)

prob = DiscreteProblem(u₀, tspan, p)
jump_prob = JumpProblem(prob, Direct(), rj)
sol = solve(jump_prob, TauLeaping(); dt = δt, maxiters = 1e5, adaptive = false)
sol_array = zeros(5, length(sol))
sol_array[1:4, :] = Array(sol)
sol_array[5, :] = sum(sol_array[1:4, :], dims = 1)

min(sol_array[1:5, :]...)

Going negative is kind of the big problem with tau-leaping, and why one shouldn’t use it for small populations usually.

I guess you could try passing an isoutofdomain function to solve that would reject such steps, see

https://docs.sciml.ai/DiffEqDocs/stable/basics/common_solver_opts/#CommonSolve.solve-Tuple{SciMLBase.AbstractDEProblem,%20Vararg{Any}}

Just be aware this is introducing a bias by throwing away the Poisson random samples that would have been used. There are bridging approaches to get around that, but they’re not implemented in StochasticDiffEq yet.

Just to add, I think the real way to handle this issue is by having a multiscale algorithm that will resolve low population species with exact jumps and high population species with tau-leaping, but we don’t have that implemented yet.

Someone had a similar issue recently and posted a callback approach to having both methods, so maybe you could give that a try:

I’m surprised it’s actually working with DiscreteProblem: it should be acting on SDEProblem?

Thanks @isaacsas, that looks pretty applicable. I’ll give that a go.

Thanks for the response @ChrisRackauckas. I’ll look into SDEProblem, but the tau-leaping example in JumpProcesses.jl uses a DiscreteProblem, so I just based mine on that. I’ll look into SDE implementations more and can report back if there are differences

I made the following additions to the code above to convert it to use SDEProblem.

forcing!(du,u,p,t) = (du .= 0)
noise!(du,u,p,t) = (du .= 0)
W = WienerProcess(0.0,0.0,0.0)

prob = SDEProblem(forcing!, noise!, u₀, tspan, p, noise = W)
jump_prob = JumpProblem(prob, Direct(), rj)
sol = solve(jump_prob, TauLeaping(), dt = δt, maxiters = 1e5)
sol_array = zeros(5, length(sol))
sol_array[1:4, :] = Array(sol)
sol_array[5, :] = sum(sol_array[1:4, :], dims = 1)

This works, but only with EM() solver, and not with TauLeaping(). It seems like only EM() and ImplicitEM() are compatible with RegularJump problems, according to the JumpProccesses.jl documentation.

TauLeaping() produces the following error.

julia> sol = solve(jump_prob, TauLeaping(), dt = δt, maxiters = 1e5)
ERROR: The algorithm is not compatible with the chosen noise type. Please see the documentation on the solver methods
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] __init(_prob::JumpProblem{true, SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Tuple{Int64, Float64, Float64, Int64}, NoiseProcess{Float64, 1, Float64, Float64, Float64, Vector{Float64}, typeof(DiffEqNoiseProcess.WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.WHITE_NOISE_BRIDGE), false, ResettableStacks.ResettableStack{Tuple{Float64, Float64, Float64}, false}, ResettableStacks.ResettableStack{Tuple{Float64, Float64, Float64}, false}, RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, SDEFunction{true,SciMLBase.FullSpecialize,…}, typeof(noise!), Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, Nothing}, Direct, CallbackSet{Tuple{}, Tuple{}}, Nothing, Tuple{}, RegularJump{true, typeof(rate), typeof(change!), Nothing}, Nothing, TaskLocalRNG, Base.Pairs{Symbol, Nothing, Tuple{Symbol}, NamedTuple{(:callback,), Tuple{Nothing}}}}, alg::TauLeaping, timeseries_init::Vector{Any}, ts_init::Vector{Any}, ks_init::Vector{Any}, recompile::Type{Val{true}}; saveat::Tuple{}, tstops::Tuple{}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_noise::Bool, save_on::Bool, save_start::Bool, save_end::Nothing, callback::Nothing, dense::Bool, calck::Bool, dt::Float64, adaptive::Bool, gamma::Rational{Int64}, abstol::Nothing, reltol::Nothing, qmin::Rational{Int64}, qmax::Rational{Int64}, qsteady_min::Int64, qsteady_max::Int64, beta2::Nothing, beta1::Nothing, qoldinit::Rational{Int64}, controller::Nothing, fullnormalize::Bool, failfactor::Int64, delta::Rational{Int64}, maxiters::Float64, dtmax::Float64, dtmin::Float64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), isoutofdomain::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), unstable_check::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), verbose::Bool, force_dtmin::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), userdata::Nothing, initialize_integrator::Bool, seed::UInt64, alias_u0::Bool, alias_jumps::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
   @ StochasticDiffEq ~/.julia/packages/StochasticDiffEq/0PuS7/src/solve.jl:110
 [3] init_call(::JumpProblem{true, SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Tuple{Int64, Float64, Float64, Int64}, NoiseProcess{Float64, 1, Float64, Float64, Float64, Vector{Float64}, typeof(DiffEqNoiseProcess.WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.WHITE_NOISE_BRIDGE), false, ResettableStacks.ResettableStack{Tuple{Float64, Float64, Float64}, false}, ResettableStacks.ResettableStack{Tuple{Float64, Float64, Float64}, false}, RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, SDEFunction{true,SciMLBase.FullSpecialize,…}, typeof(noise!), Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, Nothing}, Direct, CallbackSet{Tuple{}, Tuple{}}, Nothing, Tuple{}, RegularJump{true, typeof(rate), typeof(change!), Nothing}, Nothing, TaskLocalRNG, Base.Pairs{Symbol, Nothing, Tuple{Symbol}, NamedTuple{(:callback,), Tuple{Nothing}}}}, ::TauLeaping, ::Vararg{Any}; merge_callbacks::Bool, kwargshandle::KeywordArgError, kwargs::Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol}, NamedTuple{(:dt, :maxiters), Tuple{Float64, Float64}}})
   @ DiffEqBase ~/.julia/packages/DiffEqBase/JH4gt/src/solve.jl:425
 [4] init(::JumpProblem{true, SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Tuple{Int64, Float64, Float64, Int64}, NoiseProcess{Float64, 1, Float64, Float64, Float64, Vector{Float64}, typeof(DiffEqNoiseProcess.WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.WHITE_NOISE_BRIDGE), false, ResettableStacks.ResettableStack{Tuple{Float64, Float64, Float64}, false}, ResettableStacks.ResettableStack{Tuple{Float64, Float64, Float64}, false}, RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, SDEFunction{true,SciMLBase.FullSpecialize,…}, typeof(noise!), Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, Nothing}, Direct, CallbackSet{Tuple{}, Tuple{}}, Nothing, Tuple{}, RegularJump{true, typeof(rate), typeof(change!), Nothing}, Nothing, TaskLocalRNG, Base.Pairs{Symbol, Nothing, Tuple{Symbol}, NamedTuple{(:callback,), Tuple{Nothing}}}}, ::TauLeaping, ::Vararg{Any}; kwargs::Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol}, NamedTuple{(:dt, :maxiters), Tuple{Float64, Float64}}})
   @ DiffEqBase ~/.julia/packages/DiffEqBase/JH4gt/src/solve.jl:442
 [5] __solve(jump_prob::JumpProblem{true, SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Tuple{Int64, Float64, Float64, Int64}, NoiseProcess{Float64, 1, Float64, Float64, Float64, Vector{Float64}, typeof(DiffEqNoiseProcess.WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.WHITE_NOISE_BRIDGE), false, ResettableStacks.ResettableStack{Tuple{Float64, Float64, Float64}, false}, ResettableStacks.ResettableStack{Tuple{Float64, Float64, Float64}, false}, RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, SDEFunction{true,SciMLBase.FullSpecialize,…}, typeof(noise!), Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, Nothing}, Direct, CallbackSet{Tuple{}, Tuple{}}, Nothing, Tuple{}, RegularJump{true, typeof(rate), typeof(change!), Nothing}, Nothing, TaskLocalRNG, Base.Pairs{Symbol, Nothing, Tuple{Symbol}, NamedTuple{(:callback,), Tuple{Nothing}}}}, alg::TauLeaping, timeseries::Vector{Any}, ts::Vector{Any}, ks::Vector{Any}, recompile::Type{Val{true}}; kwargs::Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol}, NamedTuple{(:dt, :maxiters), Tuple{Float64, Float64}}})
   @ JumpProcesses ~/.julia/packages/JumpProcesses/8AMkl/src/solve.jl:5
 [6] solve(prob::JumpProblem{true, SDEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, Tuple{Int64, Float64, Float64, Int64}, NoiseProcess{Float64, 1, Float64, Float64, Float64, Vector{Float64}, typeof(DiffEqNoiseProcess.WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.WHITE_NOISE_BRIDGE), false, ResettableStacks.ResettableStack{Tuple{Float64, Float64, Float64}, false}, ResettableStacks.ResettableStack{Tuple{Float64, Float64, Float64}, false}, RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, SDEFunction{true,SciMLBase.FullSpecialize,…}, typeof(noise!), Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, Nothing}, Direct, CallbackSet{Tuple{}, Tuple{}}, Nothing, Tuple{}, RegularJump{true, typeof(rate), typeof(change!), Nothing}, Nothing, TaskLocalRNG, Base.Pairs{Symbol, Nothing, Tuple{Symbol}, NamedTuple{(:callback,), Tuple{Nothing}}}}, args::TauLeaping; kwargs::Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol}, NamedTuple{(:dt, :maxiters), Tuple{Float64, Float64}}})
   @ DiffEqBase ~/.julia/packages/DiffEqBase/JH4gt/src/solve.jl:965
 [7] top-level scope
   @ code:1