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 streams

estimate-inf-and-r-multi-stream

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

  // 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.
  // This lets the SAME model fit one stream on its own, two streams linked
  // together, or all three jointly, so we can build the model up in parts.
  int<lower = 0, upper = 1> use_cases;
  int<lower = 0, upper = 1> use_deaths;
  int<lower = 0, upper = 1> use_ww;

  // switch on a time-varying death scaling (a random walk on the IFR) to let
  // the model explain a drifting severity rather than forcing a compromise.
  int<lower = 0, upper = 1> tv_death_scale;

  // 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> 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)
  // optional random walk on the (log) death scaling, only used when
  // tv_death_scale = 1. Sized to zero otherwise so it costs nothing.
  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);
  array[n] real infections = renewal(I0, R, gen_time_pmf);

  // death scaling: either constant (ifr) or a geometric random walk starting
  // from ifr, which lets a drifting severity be absorbed rather than forced
  // into the shared infections.
  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 is a convolution of 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
  init_R ~ normal(1, 0.5) T[0, ];
  rw_noise ~ std_normal();
  rw_sd ~ normal(0, 0.05) T[0, ];
  ascertainment ~ beta(2, 2);
  ifr ~ beta(1, 50);
  ww_scale ~ normal(1, 1) T[0, ];
  ww_sigma ~ normal(0, 0.5) T[0, ];
  if (tv_death_scale) {
    ifr_rw_noise ~ std_normal();
    ifr_rw_sd ~ normal(0, 0.1);   // half-normal via <lower = 0> on the parameter
  }

  // joint likelihood: each stream contributes its own term off infections,
  // and only the streams we switch on are added. Because the streams are
  // conditionally independent given the infections, the joint log-likelihood
  // is simply the sum of the per-stream terms.
  if (use_cases) {
    cases ~ poisson(exp_cases);
  }
  if (use_deaths) {
    deaths ~ poisson(exp_deaths);
  }
  if (use_ww) {
    ww ~ normal(log(exp_ww), 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];
    }
  }
}