SDAE for Stochastic SIR

I am fairly new at Julia and still learning the ropes. DifferentialEquations.jl package is my primary interest at preset. I want to model/plot a stochastic SIR model using SDAE. The reason my model is in terms of percentages of the population being represented by the S: susceptible, I: infective and R: removed compartments. I define my drift and diffusion functions as something such as:

function f(du,u,p,t)
du[1] = -ß*u[1]*u[2]
du[2] = ß*u[1]*u[2] -µ*u[2]
du[3] = µ*u[2]
end

function g(du,u,p,t)
du[1] = -1/2*u[1]*u[2]*u[3]
du[2] = u[1]*u[2]*u[3]
du[3] = -1/2*u[1]*u[2]*u[3]
end

with say u0 = (0.8, 0.1,0.1)
Now, if I sum

du[1]+du[2]+du[3]=0

I want to be able to input the constraint:

u[1]+u[2]+u[3] = 1

From the documentation at DifferentialEquations.jl: Scientific Machine Learning (SciML) Enabled Simulation and Estimation · DifferentialEquations.jl

It seems I should perhaps add the following

du[4] = u[1] + u[2] + u[3] -1

Do I add this to both drift and diffusion? And what do I give for the mass matrix?

If you add the extra equation then you have to eliminate one of the others. But, are you sure you don’t want to just use two Brownian motions here instead of 3 ? You’d get conservative solutions if you do the chemical Langevin like form

Sorry, not clear to me when you say eliminate one of the others. Which others to be specific?

To answer your question, I wanted 3 BM out of curiosity.

If I understand it correctly, then it seems to me that your problem is something like:

dx[1] = f_1(x[1],x[2])
dx[2] = f_2(x[1],x[2])

with

x[1] + x[2] = c = constant

Here you know that

x[2] = c - x[1]

for all time steps and you can put x[2] into f_1 like

dx[1] = f_1(x[1], c - x[1])

So you can eliminate the ODE dx[2]=f_2(x[1],x[2]).

Somewhat. Let me explain the non-DAE form first, since that’s probably the right way to go here, but I’ll also do the other options. Notice that if you do two Brownian motions, you can associate the randomness with reactions instead of states.


function f(du,u,p,t)

  du[1] = -ß*u[1]*u[2]

  du[2] = ß*u[1]*u[2] -µ*u[2]

  du[3] = µ*u[2]

end

function g(du,u,p,t)

  du[1,1] = sqrt(abs(ß*u[1]*u[2]))

  du[2,1] = -sqrt(abs(ß*u[1]*u[2]))

  du[3,1] = 0

  du[1,2] = 0

  du[2,2] = sqrt(abs(µ*u[2]))

  du[3,2] = -sqrt(abs(µ*u[2]))

end

prob = SDEProblem(f,g,[0.8, 0.1,0.1],(0.0,1.0),noise_rate_prototype=zeros(3,2))

This model is naturally conservative because you have randomness in the reaction, i.e. 1->2 is a stochastic reaction, and 2->3 is a stochastic reaction, and so the randomness is either increasing or decreasing those rates at any instantaneous time, but it’s still mass action kinetics. It’s just:


u1' = -b*u1*u2*dt + sqrt(b*u1*u2)*dW1

u2' = b*u1*u2*dt - m*u2*dt - sqrt(b*u1*u2)*dW1 + sqrt(m*u2)*dW2

u3' = m*u2*dt - sqrt(m*u2)*dW2

and so from reading the equations you can very easily see the conservative-ness property. This is actually known as the Chemical Langevin Equations and is and intermediate result where the noise is the approximate of the Poisson process of a Gillespie process, and with higher copy numbers it converges to the mass action kinetic ODE models. And BTW, Catalyst is a help library which generates this stuff directly from reaction descriptions:

https://catalyst.sciml.ai/dev/tutorials/using_catalyst/

But if you’re willing to drop mass action kinetics, you can use independent Brownian processes and keep the constraint. You just need to choose an equation to drop so that it’s not overdetermined, i.e.:


function f(du,u,p,t)

  du[1] = -ß*u[1]*u[2]

  du[2] = ß*u[1]*u[2] -µ*u[2]

  du[3] = u[1] + u[2] + u[3] -1

end

mass_matrix = [1 0 0; 0 1 0; 0 0 0]

or


function f(du,u,p,t)

  du[1] = -ß*u[1]*u[2]

  du[2] = u[1] + u[2] + u[3] -1

  du[3] = µ*u[2]

end

mass_matrix = [1 0 0; 0 0 0; 0 0 1]

or


function f(du,u,p,t)

  du[1] = u[1] + u[2] + u[3] -1

  du[2] = ß*u[1]*u[2] -µ*u[2]

  du[3] = µ*u[2]

end

mass_matrix = [0 0 0; 0 1 0; 0 0 1]

All 3 of those are fine models, not necessarily MAK but an extension of sorts which could be what you want. But notice that in each case you’re eliminating one of the equations, because the dynamics of the other two fully determine the dynamics of the third.

However, as noted by @stephans3, this is a case where the constraint is linear, so you can just do the elimination by hand and then end up with a non-DAE SDE in any of those three cases. This is actually preferred since then it would be numerically easier to solve.

ModelingToolkit digression

But if you’re lazy, you can use ModelingToolkit.jl to eliminate the equations for you. for example:


using ModelingToolkit

@variables t u[1:3](t)

@parameters ß µ

D = Differential(t)

eqs = [D(u[1]) ~ -ß*u[1]*u[2],

       D(u[2]) ~ ß*u[1]*u[2] - µ*u[2],

       0 ~ u[1] + u[2] + u[3] - 1]

@named sys = ODESystem(eqs,t)

simpsys = structural_simplify(sys)

neqs = [

    D(u[1]) ~ -1/2*u[1]*u[2]*u[3]

    D(u[2]) ~ u[1]*u[2]*u[3]

]

@named stochsys = SDESystem(simpsys,neqs)

julia> equations(stochsys)

2-element Vector{Equation}:

 Differential(t)(u[1](t)) ~ -ß*u[1](t)*u[2](t)

 Differential(t)(u[2](t)) ~ ß*u[1](t)*u[2](t) - μ*u[2](t)

 

julia> ModelingToolkit.get_noiseeqs(stochsys)

2-element Vector{Equation}:

 Differential(t)(u[1](t)) ~ -0.5u[1](t)*u[2](t)*u[3](t)

 Differential(t)(u[2](t)) ~ u[1](t)*u[2](t)*u[3](t)

which will almost work. @yingboma how do I force the substitution of the u[3] out of the noise equations? The constructor that’s required for this is:


function SDESystem(sys::ODESystem, neqs; name=nothing, kwargs...)

   

    name === nothing && throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))

    SDESystem(

        equations(sys), neqs, get_iv(sys), states(sys), parameters(sys),

        get_var_to_name(sys), get_ctrls(sys), get_observed(sys), get_tgrad(sys),

        get_jac(sys), get_ctrl_jac(sys), get_Wfact(sys), get_Wfact_t(sys),

        name, get_systems(sys), get_defaults(sys),

        get_connector_type(sys); kwargs...)

end

If we fix that up, this is automated. Of course, we should finish up structural_simplify on SDESystem which would then make it much simpler.

2 Likes
using ModelingToolkit

@variables t u[1:3](t)

@parameters ß µ

D = Differential(t)

eqs = [D(u[1]) ~ -ß*u[1]*u[2],
       D(u[2]) ~ ß*u[1]*u[2] - µ*u[2],
       0 ~ u[1] + u[2] + u[3] - 1]

@named sys = ODESystem(eqs,t)

simpsys = structural_simplify(sys)

neqs = [
    D(u[1]) ~ -1/2*u[1]*u[2]*u[3]
    D(u[2]) ~ u[1]*u[2]*u[3]
]

neqs = substitute.(neqs, (Dict(eq.lhs => eq.rhs for eq in observed(simpsys)),))

@named stochsys = SDESystem(simpsys,neqs)

You can do something like this.