Last active
March 12, 2021 19:45
-
-
Save sdwfrost/8394921f50fd6b4da0e77a72ac3f96b5 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 OrdinaryDiffEq | |
using MomentClosure | |
using ModelingToolkit | |
using Tables | |
using DataFrames | |
using StatsPlots | |
``` | |
## Transitions | |
```julia | |
@parameters t β c γ | |
@variables S(t) I(t) R(t) | |
``` | |
The first system has polynomial rates. | |
```julia | |
rxs1 = [Reaction(β*c, [S,I], [I], [1,1], [2]) | |
Reaction(γ, [I], [R])] | |
rs1 = ReactionSystem(rxs1, t, [S,I,R], [β,c,γ]); | |
``` | |
The second system has non-polynomial rates. | |
```julia | |
rxs2 = [Reaction((β*c)/(S+I+R), [S,I], [I], [1,1], [2]) | |
Reaction(γ, [I], [R])] | |
rs2 = ReactionSystem(rxs2, 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) | |
tvec = 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]; | |
``` | |
```julia | |
p2 = [β=>0.05, | |
c=>10.0, | |
γ=>0.25]; | |
``` | |
## Simulating the ODE | |
```julia | |
odesys1 = convert(ODESystem, rs1) | |
oprob1 = ODEProblem(odesys1, u0, tspan, p1) | |
osol1 = solve(oprob1, Tsit5()) | |
plot(osol1) | |
``` | |
```julia | |
odesys2 = convert(ODESystem, rs2) | |
oprob2 = ODEProblem(odesys2, u0, tspan, p2) | |
osol2 = solve(oprob2, Tsit5()) | |
plot(osol2) | |
``` | |
## Moment closure | |
`q_order` does not need to be specified for polynomial rate terms. | |
```julia | |
central_eqs1 = generate_central_moment_eqs(rs1, 2, combinatoric_ratelaw=false) | |
closed_central_eqs1 = moment_closure(central_eqs1, "normal") | |
u0map1 = deterministic_IC(u0v, closed_central_eqs1) | |
``` | |
```julia | |
central_eqs2 = generate_central_moment_eqs(rs2, 2, 3, combinatoric_ratelaw=false) | |
closed_central_eqs2 = moment_closure(central_eqs2, "zero") | |
u0map2 = deterministic_IC(u0v, closed_central_eqs2) | |
``` | |
The problem can now be defined and solved. | |
```julia | |
ocprob1 = ODEProblem(closed_central_eqs1, u0map1, tspan, p1) | |
ocsol1 = solve(ocprob1,Tsit5()) | |
ocdf1 = DataFrame(ocsol1(tvec)') | |
rename!(ocdf1,[replace(string(x[1]),"(t)" => "") for x in u0map1]) | |
ocdf1[!,:t] = tvec; | |
``` | |
```julia | |
@df ocdf1 plot(:t,[:μ₁₀₀,:μ₀₁₀,:μ₀₀₁], | |
ribbon=[2 * sqrt.(:M₂₀₀), | |
2 * sqrt.(:M₀₂₀), | |
2 * sqrt.(:M₀₀₂)], | |
label=["S" "I" "R"], | |
xlabel="Time", | |
ylabel="Number") | |
``` | |
```julia | |
ocprob2 = ODEProblem(closed_central_eqs2, u0map2, tspan, p2) | |
ocsol2 = solve(ocprob2,Tsit5()) | |
ocdf2 = DataFrame(ocsol2(tvec)') | |
rename!(ocdf2,[replace(string(x[1]),"(t)" => "") for x in u0map2]) | |
ocdf2[!,:t] = tvec; | |
``` | |
```julia | |
@df ocdf2 plot(:t,[:μ₁₀₀,:μ₀₁₀,:μ₀₀₁], | |
ribbon=[2*sqrt.(:M₂₀₀), | |
2*sqrt.(:M₀₂₀), | |
2*sqrt.(:M₀₀₂)], | |
label=["S" "I" "R"], | |
xlabel="Time", | |
ylabel="Number") | |
``` |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment