update extract function to work with simulations

epiforecasts · Feb 19, 2024 · 9581066 · 9581066
1 parent b70d79d
commit 9581066
Showing 1 changed file with 22 additions and 19 deletions.
diff --git a/R/extract.R b/R/extract.R
@@ -181,28 +181,31 @@ extract_parameter_samples <- function(stan_fit, data, reported_dates,
     samples,
     reported_dates
   )
-  if (data$estimate_r == 1) {
-    out$R <- extract_parameter(
-      "R",
-      samples,
-      reported_dates
-    )
-    if (data$bp_n > 0) {
-      out$breakpoints <- extract_parameter(
-        "bp_effects",
+  if ("estimate_r" %in% names(data)) {
+    if (data$estimate_r == 1) {
+      out$R <- extract_parameter(
+        "R",
         samples,
-        1:data$bp_n
+        reported_dates
+      )
+      if (data$bp_n > 0) {
+        out$breakpoints <- extract_parameter(
+          "bp_effects",
+          samples,
+          1:data$bp_n
+        )
+        out$breakpoints <- out$breakpoints[
+          ,
+          strat := date
+        ][, c("time", "date") := NULL]
+      }
+    } else {
+      out$R <- extract_parameter(
+        "gen_R",
+        samples,
+        reported_dates
       )
-      out$breakpoints <- out$breakpoints[,
-        strat := date][, c("time", "date") := NULL
-      ]
     }
-  } else {
-    out$R <- extract_parameter(
-      "gen_R",
-      samples,
-      reported_dates
-    )
   }
   out$growth_rate <- extract_parameter(
     "r",