2functions {
3 vector sir(real t, vector y, real beta, real gamma, int N) {
4 real S = y[1];
5 real I = y[2];
6 real R = y[3];
8 real dS_dt = -beta * I * S / N;
9 real dI_dt = beta * I * S / N - gamma * I;
10 real dR_dt = gamma * I;
12 return [dS_dt, dI_dt, dR_dt]';
13 }
14}
15data {
16 int<lower=1> n_days;
17 vector[3] y0;
18 real t0;
19 array[n_days] real ts;
20 int N;
21 array[n_days] int cases; // how many people are sick on day n
22}
23parameters {
24 real<lower=0> gamma;
25 real<lower=0> beta;
26 real<lower=0> phi_inv;
27}
28transformed parameters {
29 real phi = 1. / phi_inv;
30 array[n_days] vector[3] y = ode_rk45(sir, y0, t0, ts, beta, gamma, N);
31}
32model {
33 //priors
34 beta ~ normal(2, 1);
35 gamma ~ normal(0.4, 0.5);
36 phi_inv ~ exponential(5);
38 cases ~ poisson(y[ : , 2]);
39 // try it: a overdispersed likelihood instead
40 // cases ~ neg_binomial_2(y[,2], phi);
41}
42generated quantities {
43 // R0 is the expected number of new infections caused by a single infected individual
44 real R0 = beta / gamma;
45 real recovery_time = 1 / gamma;
46 array[n_days] real pred_cases = poisson_rng(y[ : , 2]);
47 // array[n_days] real pred_cases = neg_binomial_2_rng(y[,2], phi);
48}