Issue 347

R/index.R

@@ -5,8 +5,9 @@
 #' @param bias_correct Should bias correction be implemented [TMB::sdreport()]?
 #' @param level The confidence level.
 #' @param area Grid cell area. A vector of length `newdata` from
-#'   [predict.sdmTMB()] or a value of length 1, which will be repeated
-#'   internally to match.
+#'   [predict.sdmTMB()] *or* a value of length 1 which will be repeated
+#'   internally to match *or* a character value representing the column
+#'   name of the offset.
 #' @param silent Silent?
 #' @param ... Passed to [TMB::sdreport()].
 #'
@@ -107,6 +108,11 @@
 #' }
 #' @export
 get_index <- function(obj, bias_correct = FALSE, level = 0.95, area = 1, silent = TRUE, ...)  {
+  # if offset is a character vector, use the value in the dataframe
+  if (is.character(area)) {
+    area <- obj$data[[area]]
+  }
+
   d <- get_generic(obj, value_name = "link_total",
     bias_correct = bias_correct, level = level, trans = exp, area = area, ...)
   names(d)[names(d) == "trans_est"] <- "log_est"
@@ -119,6 +125,12 @@ get_index <- function(obj, bias_correct = FALSE, level = 0.95, area = 1, silent
 #' @export
 get_cog <- function(obj, bias_correct = FALSE, level = 0.95, format = c("long", "wide"), area = 1, silent = TRUE, ...)  {
   if (bias_correct) cli_abort("Bias correction with get_cog() is currently disabled.")
+
+  # if offset is a character vector, use the value in the dataframe
+  if (is.character(area)) {
+    area <- obj$data[[area]]
+  }
+
   d <- get_generic(obj, value_name = c("cog_x", "cog_y"),
     bias_correct = bias_correct, level = level, trans = I, area = area, ...)
   d <- d[, names(d) != "trans_est", drop = FALSE]
@@ -145,6 +157,12 @@ get_eao <- function(obj,
   ...
 )  {
   if (bias_correct) cli_abort("Bias correction with get_eao() is currently disabled.")
+
+  # if offset is a character vector, use the value in the dataframe
+  if (is.character(area)) {
+    area <- obj$data[[area]]
+  }
+
   d <- get_generic(obj, value_name = c("log_eao"),
     bias_correct = bias_correct, level = level, trans = exp, area = area, ...)
   names(d)[names(d) == "trans_est"] <- "log_est"
@@ -155,6 +173,11 @@ get_eao <- function(obj,
 get_generic <- function(obj, value_name, bias_correct = FALSE, level = 0.95,
   trans = I, area = 1, silent = TRUE, ...) {
 
+  # if offset is a character vector, use the value in the dataframe
+  if (is.character(area)) {
+    area <- obj$data[[area]]
+  }
+
   reinitialize(obj$fit_obj)
 
   if ((!isTRUE(obj$do_index) && value_name[1] == "link_total") || value_name[1] == "cog_x" || value_name[[1]] == "log_eao") {

tests/testthat/test-index.R

@@ -114,6 +114,43 @@ test_that("Index integration with area vector works with extra time and possibly
   expect_equal(ind$se - ind0$se[ind0$year %in% seq(2011, 2017, 2)], c(0, 0, 0, 0))
 })
 
+test_that("get_index works", {
+  skip_on_cran()
+
+  pcod_spde <- make_mesh(pcod, c("X", "Y"), n_knots = 50, type = "kmeans")
+  m <- sdmTMB(
+    data = pcod,
+    formula = density ~ 0 + as.factor(year),
+    spatiotemporal = "off", # speed
+    time = "year", mesh = pcod_spde,
+    family = tweedie(link = "log")
+  )
+
+  # add some jittered area data to qcs_grid for testing
+  qcs_grid$area <- runif(nrow(qcs_grid), 0.9, 1.1)
+  nd <- replicate_df(qcs_grid, "year", unique(pcod$year))
+
+  predictions <- predict(m, newdata = nd, return_tmb_object = TRUE)
+
+  # get predictions with area passed as vector
+  ind <- get_index(predictions, area = nd$area)
+  # get predictions with area as a named column
+  ind2 <- get_index(predictions, area = "area")
+  expect_equal(ind, ind2)
+
+  # get predictions with area passed as vector
+  eao <- get_eao(predictions, area = nd$area)
+  # get predictions with area as a named column
+  eao2 <- get_eao(predictions, area = "area")
+  expect_equal(eao, eao2)
+
+  # get predictions with area passed as vector
+  cog <- get_cog(predictions, area = nd$area)
+  # get predictions with area as a named column
+  cog2 <- get_cog(predictions, area = "area")
+  expect_equal(cog, cog2)
+})
+
 # test_that("get_index faster epsilon bias correction", {
 #   skip_on_cran()
 #