I spent some time yesterday getting as SIR epidemiology model working in Soss. I see this a lot as a differential equation, but that doesn’t make sense for real data, because what we really observe are counts. So I used a discrete Markov chain. The transitions are “global” relative to the Markov chain, so this was an interesting design challenge.
EDIT: I should say, it doesn’t make sense to me for real data with small counts.
I ended up with this for the Markov kernel:
using Soss
mstep = @model pars,state begin
# Parameters
α = pars.α # Transmission rate
β = pars.β # Recovery rate
γ = pars.γ # Case fatality rate
# Starting counts
s0 = state.s # Susceptible
i0 = state.i # Infected
r0 = state.r # Recovered
d0 = state.d # Deceased
n0 = state.n # Population size
# Transitions between states
si ~ Binomial(s0, α * i0 / n)
ir ~ Binomial(i0, β)
id ~ Binomial(i0-ir, γ)
# Updated counts
s = s0 - si
i = i0 + si - ir - id
r = r0 + ir
d = d0 + id
n = n0 - id
next = (pars=pars, state=(s=s,i=i,r=r,d=d, n=n))
end;
This is wrapped in another model that puts priors on the rate parameters:
m = @model s0 begin
α ~ Uniform()
β ~ Uniform()
γ ~ Uniform()
pars = (α=α, β=β, γ=γ)
x ~ MarkovChain(pars, mstep(pars=pars, state=s0))
end
EDIT: Benjamin Vincent caught some errors in the model and its interpretation, see here.
And here’s inference:
julia> s0 = namedtuple(df[1,:])
(i = 1, r = 0, d = 0, n = 331000000, s = 330999999, si = 0, ir = 0, id = 0)
julia> post = dynamicHMC(m(s0=s0),(x=namedtuple.(eachrow(df)),));
julia> ppost=particles(post)
(γ = 0.00423 ± 0.00011, β = 0.00234 ± 7.9e-5, α = 0.279 ± 0.00084)
The next obvious thing is to use this for a forecast. I don’t want to share this yet, as the result is really bleak. I haven’t yet done any model checking, but I expect the model is much too simple for realistic forecasting.
To draw from the prior, you can do
s0 = (i = 1, r = 0, d = 0, n = 331000000, s = 330999999);
r = rand(m(s0=s0));
for (n,s) in enumerate(r.x)
n>10 && break
println(s)
end
Note that I had to add a few things to get this to work. Full code is available here
EDIT: I’ve updated the model to work with the current master
branch of Soss; MarkovChain
is now included.