From 5042d2f8a5a1b913ad3e725eee5a1d1bdcf80c17 Mon Sep 17 00:00:00 2001
From: Sebastian Funk <sebastian.funk@lshtm.ac.uk>
Date: Tue, 26 Jul 2022 15:44:33 +0100
Subject: [PATCH] implement suggestions by @hsbadr

See
https://github.com/epiforecasts/EpiNow2/pull/213#discussion_r550701239
https://github.com/epiforecasts/EpiNow2/pull/213#discussion_r550701313
https://github.com/epiforecasts/EpiNow2/pull/213#discussion_r550701732

For now not implementing comparison to the approximate growth rate as
this seems quite a specific use case that could also be done outside the
stan model.

Also not implementing any approximate growth rate from seeding time -
instead minimum seeding time is now set to 1, so the last seeding time
is used to calculate the first growth rate.
---
 R/extract.R                        | 2 +-
 inst/stan/estimate_infections.stan | 5 ++---
 inst/stan/simulate_infections.stan | 4 ++--
 3 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/R/extract.R b/R/extract.R
index 4201ef393..9e21e6a8a 100644
--- a/R/extract.R
+++ b/R/extract.R
@@ -119,7 +119,7 @@ extract_parameter_samples <- function(stan_fit, data, reported_dates, reported_i
   out$growth_rate <- extract_parameter(
     "r",
     samples,
-    reported_dates[-1]
+    reported_dates
   )
   if (data$week_effect > 1) {
     out$day_of_week <- extract_parameter(
diff --git a/inst/stan/estimate_infections.stan b/inst/stan/estimate_infections.stan
index c36ea58b0..24b3b56c9 100644
--- a/inst/stan/estimate_infections.stan
+++ b/inst/stan/estimate_infections.stan
@@ -138,7 +138,7 @@ model {
 generated quantities {
   int imputed_reports[ot_h];
   vector[estimate_r > 0 ? 0: ot_h] R;
-  real r[ot_h - 1];
+  real r[ot_h];
   vector[return_likelihood > 1 ? ot : 0] log_lik;
   if (estimate_r == 0){
     // sample generation time
@@ -148,8 +148,7 @@ generated quantities {
     R = calculate_Rt(infections, seeding_time, gt_mean_sample, gt_sd_sample,
                          max_gt, rt_half_window);
   }
-  // estimate growth from infections
-  r = calculate_growth(infections, seeding_time + 1);
+  r = calculate_growth(infections, seeding_time);
   // simulate reported cases
   imputed_reports = report_rng(reports, rep_phi, obs_dist);
   // log likelihood of model
diff --git a/inst/stan/simulate_infections.stan b/inst/stan/simulate_infections.stan
index 8275990a5..f1c11906b 100644
--- a/inst/stan/simulate_infections.stan
+++ b/inst/stan/simulate_infections.stan
@@ -26,7 +26,7 @@ generated quantities {
   matrix[n, t] infections; //latent infections
   matrix[n, t - seeding_time] reports; // observed cases
   int imputed_reports[n, t - seeding_time];
-  real r[n, t - seeding_time - 1];
+  real r[n, t - seeding_time];
   vector[seeding_time] uobs_inf;
   for (i in 1:n) {
     uobs_inf = generate_seed(initial_infections[i], initial_growth[i], seeding_time);
@@ -49,6 +49,6 @@ generated quantities {
     }
    // simulate reported cases
    imputed_reports[i] = report_rng(to_vector(reports[i]), rep_phi[i], obs_dist);
-   r[i] = calculate_growth(to_vector(infections[i]), seeding_time + 1);
+   r[i] = calculate_growth(to_vector(infections[i]), seeding_time);
   }
 }