From 9348ed782429afc98e4d1aee5c93753931a6bfff Mon Sep 17 00:00:00 2001
From: Joshua Lambert <joshua.lambert@lshtm.ac.uk>
Date: Mon, 9 Sep 2024 14:43:10 +0100
Subject: [PATCH] calculate and set stat_threshold and add progress bar to
 simulate_scenarios

---
 R/simulate.R | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/R/simulate.R b/R/simulate.R
index 40c712fd..c610e797 100644
--- a/R/simulate.R
+++ b/R/simulate.R
@@ -535,6 +535,12 @@ simulate_scenarios <- function(statistic,
                                breaks,
                                ...,
                                include_index_case = TRUE) {
+  # this is to ensure the trick to set the stat_threshold works
+  stopifnot(
+    "The last element in `breaks` needs to be at least 1 greater than
+      the second to last element" =
+      (breaks[length(breaks) - 1] + 1) < breaks[length(breaks)]
+  )
   checkmate::assert_logical(include_index_case, any.missing = FALSE, len = 1)
   scenarios <- expand.grid(
     offspring_dist = offspring_dist,
@@ -555,12 +561,15 @@ simulate_scenarios <- function(statistic,
     scenarios$k[scenarios$offspring_dist == "rpois"] <- NA_real_
   }
 
+  pb <- utils::txtProgressBar(max = nrow(scenarios), style = 3)
   df_list <- vector(mode = "list", length = nrow(scenarios))
   for (i in seq_len(nrow(scenarios))) {
+    utils::setTxtProgressBar(pb, i)
     args <- list(
-      n_chains = 1000,
+      n_chains = 10, # 1000
       offspring_dist = match.fun(scenarios[i, "offspring_dist"]),
-      statistic = scenarios[i, "statistic"]
+      statistic = scenarios[i, "statistic"],
+      stat_threshold = breaks[length(breaks) - 1] + 1
     )
     if (scenarios[i, "offspring_dist"] == "rpois") {
       args$lambda <- scenarios[i, "R"]
@@ -596,6 +605,7 @@ simulate_scenarios <- function(statistic,
     df_$statistic <- scenarios[i, "statistic"]
     df_list[[i]] <- df_
   }
+  close(pb)
   df <- do.call(rbind, df_list)
   return(df)
 }