Multivariate birth process parameterisation of an SIR model

sir_mbp.jl

# Libraries
using ModelingToolkit
using OrdinaryDiffEq
using LinearAlgebra: Diagonal
import Base: sqrt

# Utility function

function sqrt(x::Diagonal)
    Base.sqrt.(x)
end

# First system

@parameters t β c γ S₀ I₀ N
@variables lnSI(t) lnIR(t)
@derivatives D'~t

nSI = exp(lnSI)-1.0
nIR = exp(lnIR)-1.0
S = S₀-nSI
I = I₀+nSI-nIR;
R = N-S-I

eqs = [D(lnSI) ~ (exp(-lnSI)-0.5*exp(-2.0*lnSI))*(β*c*I/N*S),
       D(lnIR) ~ (exp(-lnIR)-0.5*exp(-2.0*lnIR))*(γ*I)]

tspan = (0.0,40.0)
u0 = [lnSI => 0.0, lnIR => 0.0]
p = [β => 0.05,
     c => 10.0,
     γ => 0.25,
     S₀ => 990.0,
     I₀ => 10.0,
     N => 1000.0]

sys = ODESystem(eqs)

# Second system
F = calculate_jacobian(sys)
@variables Σ[1:2,1:2](t)
X = [lnSI,lnIR]
Λ = sqrt(Diagonal([β*c*I/N*S,γ*I]))
eqs2 = [D.(Σ) .~ F*Σ .+ Σ*F' .+ Diagonal(-exp.(X))*Λ]