From bac6bdf290e7941f80129cf54e962655be027cea Mon Sep 17 00:00:00 2001
From: Sean Anderson <sean@seananderson.ca>
Date: Tue, 26 Mar 2024 10:55:55 -0700
Subject: [PATCH] Finish fixes so that #323 works

---
 R/index.R                   |  5 +++++
 tests/testthat/test-index.R | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/R/index.R b/R/index.R
index 8197f77e7..a2075b0c3 100644
--- a/R/index.R
+++ b/R/index.R
@@ -137,6 +137,11 @@ get_generic <- function(obj, value_name, bias_correct = FALSE, level = 0.95,
       }
     }
 
+    assert_that(!is.null(area))
+    if (length(area) > 1L) {
+      n_fakend <- if (!is.null(obj$fake_nd)) nrow(obj$fake_nd) else 0L
+      area <- c(area, rep(1, n_fakend)) # pad area with any extra time
+    }
     if (length(area) != nrow(obj$pred_tmb_data$proj_X_ij[[1]]) && length(area) != 1L) {
       cli_abort("`area` should be of the same length as `nrow(newdata)` or of length 1.")
     }
diff --git a/tests/testthat/test-index.R b/tests/testthat/test-index.R
index 4ed7cd30d..f6d6f92de 100644
--- a/tests/testthat/test-index.R
+++ b/tests/testthat/test-index.R
@@ -72,6 +72,39 @@ test_that("get_index works with subsets of years", {
   expect_equal(index_apply, index_full)
 })
 
+test_that("Index integration with area vector works with extra time and possibly not all time elements in prediction data #323", {
+  skip_on_cran()
+  fit <- sdmTMB(
+    density ~ s(depth),
+    time_varying_type = 'ar1',
+    time_varying = ~ 1,
+    time = 'year',
+    spatial = 'off',
+    spatiotemporal = 'off',
+    extra_time = c(2012, 2014, 2016),
+    data = pcod_2011,
+    family = tweedie(link = "log")
+  )
+  # with all years:
+  nd <- replicate_df(qcs_grid, "year", seq(2011, 2017))
+  nd$area <- 4
+  p <- predict(fit, newdata = nd, return_tmb_object = TRUE)
+  ind0 <- get_index(p, area = nd$area)
+
+  # newdata doesn't have all fitted years:
+  nd <- replicate_df(qcs_grid, "year", unique(pcod_2011$year))
+  nd$area <- 4
+  p <- predict(fit, newdata = nd, return_tmb_object = TRUE)
+  ind <- get_index(p, area = nd$area)
+  if (FALSE) {
+    library(ggplot2)
+    ggplot(ind, aes(year, est, ymin = lwr, ymax = upr)) + geom_pointrange() +
+      geom_pointrange(data = ind0, colour = "red", mapping = aes(x = year + 0.05))
+  }
+  expect_equal(ind$est - ind0$est[ind0$year %in% seq(2011, 2017, 2)], c(0, 0, 0, 0))
+  expect_equal(ind$se - ind0$se[ind0$year %in% seq(2011, 2017, 2)], c(0, 0, 0, 0))
+})
+
 # test_that("get_index faster epsilon bias correction", {
 #   skip_on_cran()
 #