Fix date-dimension mismatch in forecast_infections

NEWS.md

@@ -1,5 +1,9 @@
 # EpiNow2 (development version)
 
+## Bug fixes
+
+- Fixed a bug in `forecast_infections()` where the summary call to extract dates was using modified args instead of the original fit dimensions, causing a date-dimension mismatch when extending the R trajectory beyond the original observation period.
+
 ## Breaking changes
 
 - Removed deprecated arguments that have been erroring since v1.7.0/v1.8.0:

R/simulate_infections.R

@@ -398,6 +398,9 @@ forecast_infections <- function(estimates,
     draws$R <- R_draws
   }
 
+  # Extract R dates from original fit before modifying args
+  summarised <- summary(estimates, type = "parameters")
+
   # redefine time if Rt != data$t
   est_time <- estimates$args$t
   horizon <- estimates$args$horizon
@@ -422,7 +425,6 @@ forecast_infections <- function(estimates,
   }
 
   # define dates of interest
-  summarised <- summary(estimates, type = "parameters")
   dates <- seq(
     min(na.omit(unique(summarised[variable == "R"]$date))) - days(shift),
     by = "day", length.out = dim(draws$R)[2] + shift

tests/testthat/test-forecast-infections.R

@@ -173,3 +173,23 @@ test_that("get_predictions produces expected output with format = 'summary'", {
   expect_true("mean" %in% names(preds))
   expect_false("confirm" %in% names(preds))
 })
+
+test_that("forecast_infections works with R extended beyond original fit", {
+  skip_integration()
+  fixtures <- get_test_fixtures()
+  # Get the original R length from the fit
+ original_R_length <- nrow(
+    summary(fixtures$estimate_infections, type = "parameters", param = "R")
+  )
+  # Create an R vector longer than the original fit (extend by 10 days)
+  extended_R <- c(rep(NA_real_, original_R_length), rep(0.8, 10))
+  # This would fail before the fix due to date-dimension mismatch
+  sims <- forecast_infections(
+    fixtures$estimate_infections, R = extended_R, samples = 10
+  )
+  expect_equal(names(sims), c("samples", "summarised", "observations"))
+  expect_true(nrow(sims$samples) > 0)
+  # Verify get_samples works (this was part of the failing call chain)
+  samples <- get_samples(sims)
+  expect_s3_class(samples, "data.table")
+})