Stan model index

All Stan models used in the course

This page shows the source of every Stan model included in the NFIDD 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.

Probability distributions and parameter estimation

coin

data {
  int<lower = 1> N;  // integer, minimum 1
  int<lower = 0> x; // integer, minimum 0
}

parameters {
  real<lower = 0, upper = 1> theta; // real, between 0 and 1
}

model {
  // Uniform prior
  theta ~ uniform(0, 1);
  // Binomial likelihood
  x ~ binomial(N, theta);
}

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);
}

Delay distributions

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]);
    }
  }
}

Forecasting concepts

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;
  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]);
    }
  }
}

Forecasting models

mechanistic-r

functions {
  #include "functions/convolve_with_delay.stan"
  #include "functions/pop_bounded_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;
  int h;                // number of days to forecast
  array[2] real N_prior;      // prior for total population
}

transformed data {
   int m = n + h;
}

parameters {
  real<lower = 0> R;    // initial reproduction number
  real<lower = 0> N;   // total population
}

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

model {
  // priors
  R ~ normal(1, 0.5) T[0,];
  N ~ normal(N_prior[1], N_prior[2]) T[0,];
  obs ~ poisson(onsets[1:n]);
}

generated quantities {
  array[h] real forecast;
  if (h > 0) {
    for (i in 1:h) {
      forecast[i] = poisson_rng(onsets[n + i]);
    }
  }
}

statistical-r

functions {
  #include "functions/convolve_with_delay.stan"
  #include "functions/renewal.stan"
  #include "functions/geometric_diff_ar.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
}

transformed data {
   int m = n + h;
}

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

transformed parameters {
  array[m] real R = geometric_diff_ar(init_R, rw_noise, rw_sd, damp);
  array[m] real <upper = 1e5> infections = renewal(I0, R, gen_time_pmf);
  array[m] real onsets = convolve_with_delay(infections, ip_pmf);
}

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

generated quantities {
  array[h] real forecast;
  if (h > 0) {
    for (i in 1:h) {
      forecast[i] = poisson_rng(onsets[n + i]);
    }
  }
}