Created
March 12, 2021 15:54
-
-
Save sdwfrost/792d953e23e03a63297d7bc6d369146c to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# Moment closure of an SIR reaction network model using MomentClosure | |
Simon Frost (@sdwfrost), 2021-03-10 | |
## Libraries | |
```julia | |
using DifferentialEquations | |
using OrdinaryDiffEq | |
using MomentClosure | |
using ModelingToolkit | |
using DataFrames | |
using StatsPlots | |
``` | |
## Transitions | |
```julia | |
@parameters t β c γ | |
@variables S(t) I(t) R(t) | |
``` | |
```julia | |
rxs1 = [Reaction(β*c, [S,I], [I], [1,1], [2]) | |
Reaction(γ, [I], [R])] | |
rs1 = ReactionSystem(rxs1, t, [S,I,R], [β,c,γ]); | |
``` | |
## Time domain | |
We set the timespan for simulations, `tspan`, initial conditions, `u0`, and parameter values, `p` (which are unpacked above as `[β,γ]`). | |
```julia | |
δt = 0.1 | |
tmax = 40.0 | |
tspan = (0.0,tmax) | |
ts = 0:δt:tmax | |
``` | |
## Initial conditions | |
In `ModelingToolkit`, the initial values are defined by an vector of `Pair`s. | |
```julia | |
u0 = [S => 990.0, | |
I => 10.0, | |
R => 0.0] | |
``` | |
We will also need this as a vector of type `Real` for `MomentClosure.jl`. | |
```julia | |
u0v = [x[2] for x in u0] | |
``` | |
## Parameter values | |
Similarly, the parameter values are defined by a dictionary. | |
```julia | |
p1 = [β=>0.00005, | |
c=>10.0, | |
γ=>0.25]; | |
``` | |
## Generating central moment equations | |
We often deal with central moments (mean, variances, etc.) in epidemiological models. For polynomial rates (e.g. λ=βSI), we only need to specify the order of the moments we want. For demonstration purposes, we'll set the order, `m` to be 2 i.e. to just output means and (co)variances. | |
```julia | |
central_eqs1 = generate_central_moment_eqs(rs1, 2, combinatoric_ratelaw=false) | |
``` | |
## Moment closure | |
`MomentClosure.jl` provides many ways to close the system. For each system, we also need to generate a set of corresponding initial conditions. | |
```julia | |
closure_methods = ["zero","normal","poisson","log-normal","gamma","derivative matching"]; | |
``` | |
```julia | |
closed_central_eqs1 = Dict(cm=>moment_closure(central_eqs1,cm) for cm in closure_methods); | |
``` | |
```julia | |
u0map1 = Dict(cm=> deterministic_IC(u0v,closed_central_eqs1[cm]) for cm in closure_methods); | |
``` | |
## Defining and solving the closed equations | |
The problem can now be defined and solved. Here, I cycle through the closure methods. | |
```julia | |
closed_central_eqs1_df = Dict{String,DataFrame}() | |
for cm in closure_methods | |
println(cm) | |
prob = ODEProblem(closed_central_eqs1[cm], u0map1[cm], tspan, p1) | |
sol = solve(prob) | |
df = DataFrame(sol(ts)') | |
rename!(df,[replace(string(x[1]),"(t)" => "") for x in u0map1[cm]]) | |
df[!,:t] = ts | |
closed_central_eqs1_df[cm] = df | |
end; | |
``` |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment