seid.stan

functions {
  real mm_binom_w_beta_lpdf(real y, real n, real p){
    real shape = p*n;
    real rate = (1-p)*n;
    
    return beta_lpdf(y | shape, rate);
  }
}
data {
  int<lower=0> numobs; 
  array[numobs] int obs; 
  real<lower = 0> N;
}
parameters {
  real <lower = 0> dispersion; 
  real <lower = 0, upper = 1> sus_frac; 
  real <lower = 0, upper = 1> inf_frac;

  
  real <lower=0.0001,upper= 3> beta; // Exposure rate
  real <lower=0.0001,upper= 100> inc_p; // Inverse of infection rate
  real <lower=0.0001,upper= 100> inf_p; // Inverse of death rate
    


  array[numobs] real<lower=0, upper = 1> exposures; 
  array[numobs] real<lower=0, upper = 1> infections; 
  array[numobs] real<lower=0, upper = 1> deaths; 
}
transformed parameters{
  real sigma; 
  real gamma; // can't use name_p because of reserved keywords
  
  vector[numobs + 1] S_full; 
  vector[numobs + 1] E_full;
  vector[numobs + 1] I_full; 
  vector[numobs + 1] D_full;
  
  real init_sus;
  real init_exposures;
  real init_infections;
  real init_deaths;
  
  sigma =  1/inc_p;
  gamma =  1/inf_p;
  
  
  init_sus = N*sus_frac; // Initial values. 
  init_exposures = inc_p/inf_p * N * sus_frac * inf_frac;
  init_infections = N * sus_frac * inf_frac;
  init_deaths = 0;

  
  S_full[1] = init_sus;
  E_full[1] = init_exposures;
  I_full[1] = init_infections;
  D_full[1] = init_deaths;
  
  
  for (t in 1:numobs) { 
    S_full[t+1] = S_full[t] - S_full[t]*exposures[t];
    E_full[t+1] = E_full[t] - E_full[t]*infections[t] + S_full[t]*exposures[t];
    I_full[t+1] = I_full[t] + E_full[t]*infections[t] - I_full[t]*deaths[t];
    D_full[t+1] = D_full[t] + I_full[t]*deaths[t];
  }
}
model {
  dispersion ~ exponential(6);

  // Priors on parameters (commented out => uniform)
  beta ~ gamma(2, 0.25);
  inc_p ~ gamma(6, 0.5);
  inf_p ~ gamma(20, 0.11);
  
  sus_frac ~ beta(1, 1);
  inf_frac ~ beta(0.2, 199.8);
  

  // Now all observations in the loop
  for (t in 1:numobs) { 
    
      // keep in mind, this array shifted by one. 
      exposures[t] ~ mm_binom_w_beta(S_full[t], -expm1(-beta*I_full[t] / (S_full[t] + E_full[t] + I_full[t])));
      infections[t] ~ mm_binom_w_beta(E_full[t], -expm1(-sigma));
      deaths[t] ~ mm_binom_w_beta(I_full[t], -expm1(-gamma));

      obs[t] ~ neg_binomial_2(I_full[t]*deaths[t], dispersion);
  }
}
generated quantities {
  array[numobs] real u;
  
  u = deaths;
}