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.