From 3bc26d4c3efc83863d7bb8d6a9a87d8a1897b77e Mon Sep 17 00:00:00 2001
From: Sam <s.e.abbott12@gmail.com>
Date: Fri, 20 Dec 2024 17:25:17 +0000
Subject: [PATCH] make it possible to either use pop out of sample or fit with
 it in sample

---
 R/create.R                          |  2 +-
 R/opts.R                            | 27 +++++++++++++++++++++------
 inst/stan/functions/infections.stan |  2 +-
 man/rt_opts.Rd                      | 10 +++++++++-
 4 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/R/create.R b/R/create.R
index ef7c02237..97ae12fda 100644
--- a/R/create.R
+++ b/R/create.R
@@ -263,7 +263,7 @@ create_rt_data <- function(rt = rt_opts(), breakpoints = NULL,
     breakpoints = breakpoints,
     future_fixed =  as.numeric(future_rt$fixed),
     fixed_from = future_rt$from,
-    use_pop = as.integer(rt$pop != Fixed(0)),
+    use_pop = as.integer(rt$pop != Fixed(0)) + as.integer(rt$estimate_pop),
     stationary =  as.numeric(rt$gp_on == "R0"),
     future_time = horizon - future_rt$from
   )
diff --git a/R/opts.R b/R/opts.R
index e0a478d97..6dc1f545f 100644
--- a/R/opts.R
+++ b/R/opts.R
@@ -324,6 +324,13 @@ trunc_opts <- function(dist = Fixed(0), default_cdf_cutoff = 0.001,
 #' Used to adjust Rt estimates based on the proportion of the population that
 #' is susceptible. Defaults to `Fixed(0)` which means no population adjustment
 #' is done.
+#' 
+#' @param pop_within_horizon Logical, defaults to `FALSE`. Should the
+#' susceptible population adjustment be applied within the time horizon of data
+#' or just beyond it. Note that if `pop_within_horizon = TRUE` the Rt estimate
+#' will be unadjusted for susceptible depletion but the resulting posterior
+#' predictions for infections and cases will be adjusted for susceptible
+#' depletion.
 #'
 #' @param gp_on Character string, defaulting to  "R_t-1". Indicates how the
 #' Gaussian process, if in use, should be applied to Rt. Currently supported
@@ -354,7 +361,8 @@ rt_opts <- function(prior = LogNormal(mean = 1, sd = 1),
                     use_breakpoints = TRUE,
                     future = "latest",
                     gp_on = c("R_t-1", "R0"),
-                    pop = Fixed(0)) {
+                    pop = Fixed(0),
+                    pop_within_horizon = FALSE) {
   rt <- list(
     use_rt = use_rt,
     rw = rw,
@@ -388,15 +396,22 @@ rt_opts <- function(prior = LogNormal(mean = 1, sd = 1),
   }
 
   if (is.numeric(pop)) {
-    cli_warn(
-      c(
-        "!" = "Specifying {.var pop} as a numeric value is deprecated.",
-        "i" = "Use a {.cls dist_spec} instead, e.g. Fixed({pop})."
-      )
+    lifecycle::deprecate_warn(  
+      "1.7.0",  
+      "rt_opts(pop = 'must be a `<dist_spec>`')",  
+      details = "For specifying a fixed population size, use `Fixed(pop)`"
     )
     pop <- Fixed(pop)
   }
   rt$pop <- pop
+  if (isTRUE(pop_within_horizon) && pop == Fixed(0)) {
+    cli_abort(
+      c(
+        "!" = "pop_within_horizon = TRUE but pop is fixed at 0."
+      )
+    )
+  }
+  rt$estimate_pop <- TRUE
 
   if (rt$use_rt) {
     rt$prior <- prior
diff --git a/inst/stan/functions/infections.stan b/inst/stan/functions/infections.stan
index f5e9f78c4..2b626f4ad 100644
--- a/inst/stan/functions/infections.stan
+++ b/inst/stan/functions/infections.stan
@@ -48,7 +48,7 @@ vector generate_infections(vector oR, int uot, vector gt_rev_pmf,
   // iteratively update infections
   for (s in 1:ot) {
     infectiousness[s] = update_infectiousness(infections, gt_rev_pmf, uot, s);
-    if (use_pop && s > nht) {
+    if (use_pop == 2 && s > nht) {
       exp_adj_Rt = exp(-R[s] * infectiousness[s] / (pop - cum_infections[nht]));
       exp_adj_Rt = exp_adj_Rt > 1 ? 1 : exp_adj_Rt;
       infections[s + uot] = (pop - cum_infections[s]) * (1 - exp_adj_Rt);
diff --git a/man/rt_opts.Rd b/man/rt_opts.Rd
index 3155f9101..358117438 100644
--- a/man/rt_opts.Rd
+++ b/man/rt_opts.Rd
@@ -11,7 +11,8 @@ rt_opts(
   use_breakpoints = TRUE,
   future = "latest",
   gp_on = c("R_t-1", "R0"),
-  pop = Fixed(0)
+  pop = Fixed(0),
+  pop_within_horizon = FALSE
 )
 }
 \arguments{
@@ -54,6 +55,13 @@ estimates, which may not be desirable.}
 Used to adjust Rt estimates based on the proportion of the population that
 is susceptible. Defaults to \code{Fixed(0)} which means no population adjustment
 is done.}
+
+\item{pop_within_horizon}{Logical, defaults to \code{FALSE}. Should the
+susceptible population adjustment be applied within the time horizon of data
+or just beyond it. Note that if \code{pop_within_horizon = TRUE} the Rt estimate
+will be unadjusted for susceptible depletion but the resulting posterior
+predictions for infections and cases will be adjusted for susceptible
+depletion.}
 }
 \value{
 An \verb{<rt_opts>} object with settings defining the time-varying