Estimate R0 Stan

gistfile1.txt

# in the R script: 

init_fun <- function(){list(R = array(rep(rgamma(n = data$t, 
                                                    shape = (rt_prior$mean / rt_prior$sd)^2, 
                                                    scale = (rt_prior$sd^2) / rt_prior$mean),
                                             data$k),dim = c(data$k, data$t)),
                              phi = array(rexp(data$k, 1/2)))}
                              
                              



# in estimate_R.stan

// Calculate infectiousness at each timestep for each sample in turn
  for (j in 1:k){
      // Initialise infectiousness as zero initially
      infectiousness[j] = rep_vector(0, t);
      for (s in 2:t){
          for (i in 1:(min((s - 1), n - 1))){
            infectiousness[j][s] += (obs_imported[s - i, j] + obs_local[s - i, j]) * intervals[i + 1, j];
          
          if (infectiousness[j][s] == 0)
            infectiousness[j][s] = 0.0000001;
         
      }
    }
  }
  
  
  
  ## estimate with reduced sum - not working, needs updating
  
  functions {
  real partial_sum(int[] x_slice,
                   int start, int end,
                   int window,
                   int[] obs_local,
                   vector R,
                   vector infectiousness, 
                   real phi) {
    
    real out = 0;
    
    for (s in x_slice)
      for (i in (s - window + 1):s){
        out += neg_binomial_2_lpmf(obs_local[i] | R[s] * infectiousness[i], phi);
      }
    
    return out;
    // return neg_binomial_2_lpmf(obs_local[i, j] | R[j][s] * infectiousness[j][i], phi[j]);
    // return neg_binomial_2_lpmf(y_slice | R[start:end] * infectiousness[j][i], phi[j]);
    // return bernoulli_logit_lpmf(y_slice | beta[1] + beta[2] * x[start:end]);
  }
}

data {
  int t; // number of time steps
  int k; // number of interval and case samples
  int n; // length of interval distributions
  int <lower = 0> obs_imported[t, k]; // imported cases
  int <lower = 0> obs_local[t, k]; // local cases
  int window; // length of window
  real intervals[n, k]; // matrix of different interval distributions
  real <lower = 0> r_mean;
  real <lower = 0> r_sd;
}

transformed data{
  // Set up transformed data objects
  vector[t] infectiousness[k];
  real r_alpha; //alpha parameter of the R gamma prior
  real r_beta;  //beta parameter of the R gamma prior
  int indices[(t - window)];

  // calculate alpha and beta for gamma distribution
  r_alpha = (r_mean / r_sd)^2;
  r_beta = (r_sd^2) / r_mean;
  
  for (s in 1:(t - window))
    indices[s] = s + window;
    
  
  // Calculate infectiousness at each timestep for each sample in turn
  for (j in 1:k){
      // Initialise infectiousness as zero initially
      infectiousness[j] = rep_vector(0, t);
      for (s in 2:t){
          for (i in 1:(min((s - 1), n - 1))){
            infectiousness[j][s] += (obs_imported[s - i, j] + obs_local[s - i, j]) * intervals[i + 1, j];
          
          if (infectiousness[j][s] == 0)
            infectiousness[j][s] = 0.0000001;
         
      }
    }
  }

}

parameters{
  vector<lower=0>[t] R[k]; // Effective reproduction number over time
  real <lower = 0> phi[k]; // Dispersion of negative binomial distribution
}

model {
  int grainsize = 1;  
  for (j in 1:k) {
     R[j] ~ gamma(r_alpha, r_beta); // Prior  on Rt
    phi[j] ~ exponential(0.5); //Prior on Phi
  }

  // print(Rt)
  // print(phi)

  // print(infectiousness[1]);
  // print(". . .");
  // print(infectiousness[2]);

 //Build likelihood across all samples
 for(j in 1:k) {
   target += reduce_sum(partial_sum, 
                       indices,
                       grainsize,
                       window, 
                       obs_local[j],
                       R[j],
                       infectiousness[j],
                       phi[j]);
 }
}