inc2prev/stan/functions/rt.stan at e1596adb2e2f4d4826c4ea9de0a30480cd32dfe2 · epiforecasts/inc2prev · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
// Code from:
// https://github.com/epiforecasts/EpiNow2/tree/master/inst/stan/functions
// discretised truncated gamma pmf
vector discretised_gamma_pmf(int[] y, real mu, real sigma, int max_val, int off) {
  int n = num_elements(y);
  vector[n] pmf;
  real trunc_pmf;
  // calculate alpha and beta for gamma distribution
  real small = 1e-5;
  real large = 1e8;
  real c_sigma = sigma < small ? small : sigma;
  real c_mu = mu < small ? small : mu;
  real alpha = ((c_mu) / c_sigma)^2;
  real beta = (c_mu) / (c_sigma^2);
  // account for numerical issues
  alpha = alpha < small ? small : alpha;
  alpha = alpha > large ? large : alpha;
  beta = beta < small ? small : beta;
  beta = beta > large ? large : beta;
  // calculate pmf
  trunc_pmf = gamma_cdf(max_val + 1 | alpha, beta) - gamma_cdf(off | alpha, beta);
  for (i in 1:n){
    pmf[i] = (gamma_cdf(y[i] | alpha, beta) - gamma_cdf(y[i] - 1 | alpha, beta)) /
    trunc_pmf;
  }
  return(pmf);
}
// calculate infectiousness (weighted sum of the generation time and infections)
// for a single time point
real update_infectiousness(vector infections, vector gt_pmf,
                           int seeding_time, int max_gt, int index){
  int inf_start = max(1, (index + seeding_time - max_gt));
  int inf_end = (index + seeding_time - 1);
  int pmf_accessed = min(max_gt, index + seeding_time - 1);
  real new_inf = dot_product(infections[inf_start:inf_end], tail(gt_pmf, pmf_accessed));
  return(new_inf);
}
// calculate Rt directly from inferred infections
vector calculate_Rt(vector infections, int seeding_time,
                    real gt_mean, real gt_sd, int max_gt,
                    int smooth) {
  vector[max_gt] gt_pmf;
  int gt_indexes[max_gt];
  int t = num_elements(infections);
  int ot = t - seeding_time;
  vector[ot] R;
  vector[ot] sR;
  vector[ot] infectiousness = rep_vector(1e-5, ot);
  // calculate PMF of the generation time
  for (i in 1:(max_gt)) {
    gt_indexes[i] = max_gt - i + 2;
  }
  gt_pmf = discretised_gamma_pmf(gt_indexes, gt_mean, gt_sd, max_gt, 1);
  // calculate Rt using Cori et al. approach
  for (s in 1:ot) {
    infectiousness[s] += update_infectiousness(infections, gt_pmf, seeding_time,
                                               max_gt, s);
    R[s] = infections[s + seeding_time] / infectiousness[s];
  }
  if (smooth) {
    for (s in 1:ot) {
      real window = 0;
      sR[s] = 0;
      for (i in max(1, s - smooth):min(ot, s + smooth)) {
        sR[s] += R[i];
        window += 1;
      }
      sR[s] = sR[s] / window;
    }
  }else{
    sR = R;
  }
  return(sR);
}