Stan model index

All Stan models used in the course

This page shows the source of every Stan model included in the nfidd.nowcasting package, grouped by the session in which it is first introduced. Reading them in order makes it easier to see how the models build up across the course.

Delay distributions

gamma

// gamma_model.stan
data {
  int<lower=0> N;
  array[N] real y;
}

parameters {
  real<lower=0> alpha;
  real<lower=0> beta;
}

model {
  alpha ~ normal(0, 10) T[0,];
  beta ~ normal(0, 10) T[0,];
  y ~ gamma(alpha, beta);
}

lognormal

// lognormal_model.stan
data {
  int<lower=0> N; // number of data points
  array[N] real y; // data
}

parameters {
  real meanlog;
  real<lower=0> sdlog;
}

model {
  meanlog ~ normal(0, 10);  // prior distribution
  sdlog ~ normal(0, 10) T[0, ]; // prior distribution

  y ~ lognormal(meanlog, sdlog);
}

Biases in delay distributions

censored-delay-model

data {
  int<lower = 0> n;
  array[n] int<lower = 0> onset_to_hosp;
}

parameters {
  real meanlog;
  real<lower = 0> sdlog;
  array[n] real<lower = 0, upper = 1> onset_day_time;
  array[n] real<lower = 0, upper = 1> hosp_day_time;
}

transformed parameters {
  array[n] real<lower = 0> true_onset_to_hosp;
  for (i in 1:n) {
    true_onset_to_hosp[i] =
      onset_to_hosp[i] + hosp_day_time[i] - onset_day_time[i];
  }
}

model {
  meanlog ~ normal(0, 10);
  sdlog ~ normal(0, 10) T[0, ];
  onset_day_time ~ uniform(0, 1);
  hosp_day_time ~ uniform(0, 1);

  true_onset_to_hosp ~ lognormal(meanlog, sdlog);
}

truncated-delay-model

data {
  int<lower = 0> n;
  array[n] real<lower = 0> onset_to_hosp;
  array[n] real<lower = 0> time_since_onset;
}

parameters {
  real meanlog;
  real<lower = 0> sdlog;
}

model {
  meanlog ~ normal(0, 10);
  sdlog ~ normal(0, 10) T[0, ];

  for (i in 1:n) {
    onset_to_hosp[i] ~ lognormal(meanlog, sdlog) T[0, time_since_onset[i]];
  }
}

Using delay distributions to model the data generating process

estimate-infections

functions {
  #include "functions/convolve_with_delay.stan"
}

data {
  int n;            // number of time days
  array[n] int obs; // observed onsets
  int<lower = 1> ip_max; // max incubation period
  // probability mass function of incubation period distribution (first index zero)
  array[ip_max + 1] real ip_pmf;
}

parameters {
  array[n] real<lower = 0> infections;
}

transformed parameters {
  array[n] real onsets = convolve_with_delay(infections, ip_pmf);
}

model {
  // priors
  infections ~ normal(0, 10) T[0, ];
  obs ~ poisson(onsets);
}

\(R_t\) estimation and the renewal equation

estimate-r

functions {
  #include "functions/renewal.stan"
}

data {
  int n;                // number of days
  int I0;              // number initially infected
  array[n] int obs;     // observed infections
  int gen_time_max;     // maximum generation time
  array[gen_time_max] real gen_time_pmf;  // pmf of generation time distribution
}

parameters {
  array[n] real<lower = 0> R;
}

transformed parameters {
  array[n] real infections = renewal(I0, R, gen_time_pmf);
}

model {
  // priors
  R ~ normal(1, 1) T[0, ];
  obs ~ poisson(infections);
}

estimate-inf-and-r

functions {
  #include "functions/convolve_with_delay.stan"
  #include "functions/renewal.stan"
}

data {
  int n;                // number of days
  int I0;              // number initially infected
  array[n] int obs;     // observed symptom onsets
  int gen_time_max;     // maximum generation time
  array[gen_time_max] real gen_time_pmf;  // pmf of generation time distribution
  int<lower = 1> ip_max; // max incubation period
  array[ip_max + 1] real ip_pmf;
}

parameters {
  array[n] real<lower = 0> R;
}

transformed parameters {
  array[n] real infections = renewal(I0, R, gen_time_pmf);
  array[n] real onsets = convolve_with_delay(infections, ip_pmf);
}

model {
  // priors
  R ~ normal(1, 1) T[0, ];
  obs ~ poisson(onsets);
}

estimate-inf-and-r-rw

functions {
  #include "functions/convolve_with_delay.stan"
  #include "functions/renewal.stan"
  #include "functions/geometric_random_walk.stan"
}

data {
  int n;                // number of days
  int I0;              // number initially infected
  array[n] int obs;     // observed symptom onsets
  int gen_time_max;     // maximum generation time
  array[gen_time_max] real gen_time_pmf;  // pmf of generation time distribution
  int<lower = 1> ip_max; // max incubation period
  array[ip_max + 1] real ip_pmf;
}

parameters {
  real<lower = 0> init_R;         // initial reproduction number
  array[n-1] real rw_noise; // random walk noise
  real<lower = 0> rw_sd; // random walk standard deviation
}

transformed parameters {
  array[n] real R = geometric_random_walk(init_R, rw_noise, rw_sd);
  array[n] real infections = renewal(I0, R, gen_time_pmf);
  array[n] real onsets = convolve_with_delay(infections, ip_pmf);
}

model {
  // priors
  init_R ~ normal(1, 0.5) T[0, ];
  rw_noise ~ std_normal();
  rw_sd ~ normal(0, 0.05) T[0, ];
  obs ~ poisson(onsets);
}

Nowcasting

simple-nowcast

functions {
  #include "functions/condition_onsets_by_report.stan"
}

data {
  int n;                // number of days
  array[n] int obs;     // observed symptom onsets
  int report_max;       // max reporting delay
  array[report_max + 1] real report_cdf;
}

parameters {
  array[n] real<lower = 0> onsets;
}

transformed parameters {
  array[n] real reported_onsets = condition_onsets_by_report(onsets, report_cdf);
}

model {
  onsets ~ normal(5, 20) T[0,];
  // Likelihood
  obs ~ poisson(reported_onsets);
}

simple-nowcast-rw

functions {
  #include "functions/geometric_random_walk.stan"
  #include "functions/condition_onsets_by_report.stan"
}

data {
  int n;                // number of days
  array[n] int obs;     // observed symptom onsets
  int report_max;       // max reporting delay
  array[report_max + 1] real report_cdf;
}

parameters {
  real<lower=0> init_onsets;
  array[n-1] real rw_noise;
  real<lower=0> rw_sd;
}

transformed parameters {
  array[n] real onsets = geometric_random_walk(init_onsets, rw_noise, rw_sd);
  array[n] real reported_onsets = condition_onsets_by_report(onsets, report_cdf);
}

model {
  init_onsets ~ lognormal(0, 1) T[0,];
  rw_noise ~ std_normal();
  rw_sd ~ normal(0, 5) T[0,];
  //Likelihood
  obs ~ poisson(reported_onsets);
}

Joint nowcasting with an unknown reporting delay

joint-nowcast

functions {
  #include "functions/geometric_random_walk.stan"
  #include "functions/observe_onsets_with_delay.stan"
  #include "functions/combine_obs_with_predicted_obs_rng.stan"
}

data {
  int n;                // number of days
  int m;                // number of reports
  array[n] int p;       // number of observations per day
  array[m] int obs;     // observed symptom onsets
  int d;                // number of reporting delays
}

transformed data{
  array[n] int P = to_int(cumulative_sum(p));
  array[n] int D = to_int(cumulative_sum(rep_array(d, n)));
}

parameters {
  real<lower=0> init_onsets;
  array[n-1] real rw_noise;
  real<lower=0> rw_sd;
  simplex[d] reporting_delay; // reporting delay distribution
}

transformed parameters {
  array[n] real onsets = geometric_random_walk(init_onsets, rw_noise, rw_sd);
  array[m] real onsets_by_report = observe_onsets_with_delay(onsets, reporting_delay, P, p);
}

model {
  // Prior
  init_onsets ~ normal(1, 5) T[0,];
  rw_noise ~ std_normal();
  rw_sd ~ normal(0, 0.05) T[0,];
  reporting_delay ~ dirichlet(rep_vector(1, d));
  // Likelihood
  obs ~ poisson(onsets_by_report);
}

generated quantities {
  array[d*n] real complete_onsets_by_report = observe_onsets_with_delay(onsets, reporting_delay, D, rep_array(d, n));
  array[n] int nowcast = combine_obs_with_predicted_obs_rng(obs, complete_onsets_by_report, P, p, d, D);
}

joint-nowcast-with-r

functions {
  #include "functions/geometric_random_walk.stan"
  #include "functions/renewal.stan"
  #include "functions/convolve_with_delay.stan"
  #include "functions/observe_onsets_with_delay.stan"
  #include "functions/combine_obs_with_predicted_obs_rng.stan"
}

data {
  int n;                // number of days
  int m;                // number of reports
  array[n] int p;       // number of observations per day
  array[m] int obs;     // observed symptom onsets
  int d;                // number of reporting delays
  int gen_time_max;     // maximum generation time
  array[gen_time_max] real gen_time_pmf;  // pmf of generation time distribution
  int<lower = 1> ip_max; // max incubation period
  array[ip_max + 1] real ip_pmf;
  int h;                // number of days to forecast
}

transformed data{
  array[n] int P = to_int(cumulative_sum(p));
  array[n] int D = to_int(cumulative_sum(rep_array(d, n)));
}

parameters {
  real<lower = 0> init_I;         // initial number of infected
  real<lower = 0> init_R;         // initial reproduction number
  array[n-1] real rw_noise;       // random walk noise
  real<lower = 0> rw_sd; // random walk standard deviation
  simplex[d] reporting_delay; // reporting delay distribution
}

transformed parameters {
  array[n] real R = geometric_random_walk(init_R, rw_noise, rw_sd);
  array[n] real infections = renewal(init_I, R, gen_time_pmf);
  array[n] real onsets = convolve_with_delay(infections, ip_pmf);
  array[m] real onsets_by_report = observe_onsets_with_delay(onsets, reporting_delay, P, p);
}

model {
  // Prior
  init_I ~ lognormal(0, 1);
  init_R ~ normal(1, 0.5) T[0, ];
  rw_noise ~ std_normal();
  rw_sd ~ normal(0, 0.05) T[0,];
  reporting_delay ~ dirichlet(rep_vector(1, d));
  // Likelihood
  obs ~ poisson(onsets_by_report);
}

generated quantities {
  array[d*n] real complete_onsets_by_report = observe_onsets_with_delay(onsets, reporting_delay, D, rep_array(d, n));
  array[n] int nowcast = combine_obs_with_predicted_obs_rng(obs, complete_onsets_by_report, P, p, d, D);

  // Forecast the underlying onsets
  array[h] real forecast;
  if (h > 0) {
    array[h + n - 1] real f_rw_noise;
    for (i in 1:n-1) {
      f_rw_noise[i] = rw_noise[i];
    }
    for (i in n:(h + n - 1)) {
      f_rw_noise[i] = normal_rng(0, 1);
    }
    array[h + n] real f_R = geometric_random_walk(init_R, f_rw_noise, rw_sd);
    array[h + n] real f_infections = renewal(init_I, f_R, gen_time_pmf);
    array[h + n] real f_onsets = convolve_with_delay(f_infections, ip_pmf);
    for (i in 1:h) {
      forecast[i] = poisson_rng(f_onsets[n + i]);
    }
  }
}

Jointly fitting multiple data sources

estimate-inf-and-r-multi-stream

functions {
  #include "functions/convolve_with_delay.stan"
  #include "functions/renewal_seeded.stan"
  #include "functions/geometric_random_walk.stan"
}

data {
  int n;                // number of days
  int I0;               // number initially infected

  // shared latent process
  int gen_time_max;     // maximum generation time
  array[gen_time_max] real gen_time_pmf; // pmf of generation time distribution

  // per-stream switches: set to 1 to fit a stream, 0 to leave it out
  int<lower = 0, upper = 1> use_cases;
  int<lower = 0, upper = 1> use_deaths;
  int<lower = 0, upper = 1> use_ww;

  // switch on a random walk on the IFR rather than a constant IFR
  int<lower = 0, upper = 1> tv_death_scale;

  // switch cases to a negative binomial observation model
  int<lower = 0, upper = 1> use_nb_cases;

  // priors for each stream's scaling, as (mean, sd) pairs
  array[2] real ascertainment_p; // ascertainment prior, truncated to [0, 1]
  array[2] real ifr_p;           // IFR prior, truncated to [0, 1]
  array[2] real ww_scale_p;      // wastewater scaling prior, truncated to [0, ]

  // stream 1: cases (infection-to-report delay, ascertainment)
  array[n] int cases;
  int<lower = 1> case_delay_max;
  array[case_delay_max + 1] real case_delay_pmf;

  // stream 2: deaths (infection-to-death delay, IFR)
  array[n] int deaths;
  int<lower = 1> death_delay_max;
  array[death_delay_max + 1] real death_delay_pmf;

  // stream 3: wastewater (infection-to-shedding delay, scaling)
  array[n] real ww;     // log-scale wastewater concentration
  int<lower = 1> ww_delay_max;
  array[ww_delay_max + 1] real ww_delay_pmf;
}

parameters {
  real<lower = 0> seed_base;       // infections at the start of the window
  real initial_growth;             // daily (log) growth rate over the seed
  real<lower = 0> init_R;          // initial reproduction number
  array[n-1] real rw_noise;        // random walk noise
  real<lower = 0> rw_sd;           // random walk standard deviation
  real<lower = 0, upper = 1> ascertainment; // proportion of infections reported
  real<lower = 0, upper = 1> ifr;  // infection fatality ratio (initial level)
  real<lower = 0> ww_scale;        // wastewater scaling (signal per infection)
  real<lower = 0> ww_sigma;        // wastewater obs sd (log scale)
  // reciprocal overdispersion for the cases, 1 / sqrt(phi)
  real<lower = 0> cases_overdispersion;
  // random walk on the death scaling, sized to zero when not used
  array[tv_death_scale ? n - 1 : 0] real ifr_rw_noise;
  array[tv_death_scale ? 1 : 0] real<lower = 0> ifr_rw_sd;
}

transformed parameters {
  // one shared infection / Rt process feeds every stream
  array[n] real R = geometric_random_walk(init_R, rw_noise, rw_sd);
  // initial history over gen_time_max days, grown from seed_base, so the
  // renewal has a complete generation time window on the first observed day
  array[gen_time_max] real seed_infections;
  for (t in 1:gen_time_max) {
    seed_infections[t] =
      seed_base * exp(initial_growth * (t - gen_time_max));
  }
  array[n] real infections =
    renewal_seeded(seed_infections, R, gen_time_pmf);

  // death scaling: constant ifr, or a random walk starting from it
  array[n] real death_scale;
  if (tv_death_scale) {
    death_scale = geometric_random_walk(ifr, ifr_rw_noise, ifr_rw_sd[1]);
  } else {
    for (i in 1:n) death_scale[i] = ifr;
  }

  // each stream convolves the same infections with its own delay
  array[n] real exp_cases;
  array[n] real exp_deaths;
  array[n] real exp_ww;
  {
    array[n] real case_conv = convolve_with_delay(infections, case_delay_pmf);
    array[n] real death_conv = convolve_with_delay(infections, death_delay_pmf);
    array[n] real ww_conv = convolve_with_delay(infections, ww_delay_pmf);
    for (i in 1:n) {
      exp_cases[i] = ascertainment * case_conv[i];
      exp_deaths[i] = death_scale[i] * death_conv[i];
      exp_ww[i] = ww_scale * ww_conv[i];
    }
  }
}

model {
  // priors
  seed_base ~ normal(I0, I0) T[0, ];
  initial_growth ~ normal(0, 0.2);
  init_R ~ normal(1, 0.5) T[0, ];
  rw_noise ~ std_normal();
  rw_sd ~ normal(0, 0.05) T[0, ];
  ascertainment ~ normal(ascertainment_p[1], ascertainment_p[2]) T[0, 1];
  ifr ~ normal(ifr_p[1], ifr_p[2]) T[0, 1];
  ww_scale ~ normal(ww_scale_p[1], ww_scale_p[2]) T[0, ];
  ww_sigma ~ normal(0, 0.5) T[0, ];
  cases_overdispersion ~ normal(0, 1) T[0, ];
  if (tv_death_scale) {
    ifr_rw_noise ~ std_normal();
    ifr_rw_sd ~ normal(0, 0.1);   // half-normal via <lower = 0> on parameter
  }

  // the streams are conditionally independent given the infections, so the
  // joint log-likelihood is the sum of the terms we switch on
  if (use_cases) {
    if (use_nb_cases) {
      // variance = mu + mu^2 * cases_overdispersion^2
      cases ~ neg_binomial_2(exp_cases, inv_square(cases_overdispersion));
    } else {
      cases ~ poisson(exp_cases);
    }
  }
  if (use_deaths) {
    deaths ~ poisson(exp_deaths);
  }
  if (use_ww) {
    ww ~ normal(log(exp_ww), ww_sigma);
  }
}

generated quantities {
  // posterior predictive draws, with observation error, for each stream
  array[n] int pp_cases;
  array[n] int pp_deaths;
  array[n] real pp_ww;
  for (i in 1:n) {
    if (use_nb_cases) {
      pp_cases[i] = neg_binomial_2_rng(
        exp_cases[i], inv_square(cases_overdispersion)
      );
    } else {
      pp_cases[i] = poisson_rng(exp_cases[i]);
    }
    pp_deaths[i] = poisson_rng(exp_deaths[i]);
    pp_ww[i] = normal_rng(log(exp_ww[i]), ww_sigma);
  }
}

From nowcasting to forecasting

estimate-inf-and-r-rw-forecast

functions {
  #include "functions/convolve_with_delay.stan"
  #include "functions/renewal.stan"
  #include "functions/geometric_random_walk.stan"
}

data {
  int n;                // number of days
  int I0;              // number initially infected
  array[n] int obs;     // observed symptom onsets
  int gen_time_max;     // maximum generation time
  array[gen_time_max] real gen_time_pmf;  // pmf of generation time distribution
  int<lower = 1> ip_max; // max incubation period
  array[ip_max + 1] real ip_pmf;
  int h;                // number of days to forecast
}

parameters {
  real<lower = 0> init_R;         // initial reproduction number
  array[n-1] real rw_noise; // random walk noise
  real<lower = 0, upper = 1> rw_sd; // random walk standard deviation
}

transformed parameters {
  array[n] real R = geometric_random_walk(init_R, rw_noise, rw_sd);
  array[n] real infections = renewal(I0, R, gen_time_pmf);
  array[n] real onsets = convolve_with_delay(infections, ip_pmf);
}

model {
  // priors
  init_R ~ normal(1, 0.5) T[0, ];
  rw_noise ~ std_normal();
  rw_sd ~ normal(0, 0.05) T[0,];
  obs ~ poisson(onsets[1:n]);
}
generated quantities {
  array[h] real forecast;
  array[h] real forecast_R;
  if (h > 0) {
    array[n + h - 1] real f_rw_noise;
    for (i in 1:(n-1)) {
      f_rw_noise[i] = rw_noise[i];
    }
    for (i in n:(n + h - 1)) {
      f_rw_noise[i] = normal_rng(0, 1);
    }
    array[h + n] real f_R = geometric_random_walk(init_R, f_rw_noise, rw_sd);
    array[h + n] real f_infections = renewal(I0, f_R, gen_time_pmf);
    array[h + n] real f_onsets = convolve_with_delay(f_infections, ip_pmf);
    for (i in 1:h) {
      forecast[i] = poisson_rng(f_onsets[n + i]);
      forecast_R[i] = f_R[n + i];
    }
  }
}