Fix indexes for truncated NB1/NB2 #350

pbs-assess · Jun 18, 2024 · 75b4ecc · 75b4ecc
1 parent f5b4bcd
commit 75b4ecc
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 1 deletion.
diff --git a/src/sdmTMB.cpp b/src/sdmTMB.cpp
@@ -1297,7 +1297,13 @@ Type objective_function<Type>::operator()()
           }
           total(proj_year(i)) += mu_combined(i) * area_i(i);
         } else { // non-delta model
-          mu_combined(i) = InverseLink(proj_eta(i,0), link(0));
+          if (truncated_dist) {
+            // convert from mean of *un-truncated* to mean of *truncated* distribution
+            Type log_nzprob = calc_log_nzprob(exp(proj_eta(i,0)), phi(0), family(0));
+            mu_combined(i) = exp(proj_eta(i,0)) / exp(log_nzprob);
+          } else  {
+            mu_combined(i) = InverseLink(proj_eta(i,0), link(0));
+          }
           total(proj_year(i)) += mu_combined(i) * area_i(i);
         }
       }

diff --git a/tests/testthat/test-truncated-dists.R b/tests/testthat/test-truncated-dists.R
@@ -125,3 +125,67 @@ test_that("Truncated NB1 works with various types of predictions", {
   p2 <- predict(fit2, type = "response")
   expect_equal(p2, p$est, tolerance = 1e-4)
 })
+
+test_that("Truncated NB1/2 indexes are right", {
+  skip_on_cran()
+  set.seed(1)
+  d0 <- data.frame(y = rnbinom(n = 1e3, mu = 3, size = 5))
+  d0$year <- 1
+
+  # D-TNB2
+  fit <- sdmTMB(
+    y ~ 1,
+    data = d0,
+    time = "year",
+    spatiotemporal = "off",
+    spatial = FALSE,
+    family = delta_truncated_nbinom2()
+  )
+  p <- predict(fit, newdata = d0, return_tmb_object = TRUE)
+  x <- get_index(p, area = 1/nrow(d0))
+  pp <- predict(fit, type = "response")
+  expect_equal(x$est, pp$est[1], tolerance = 1e-3)
+
+  # D-TNB1
+  fit <- sdmTMB(
+    y ~ 1,
+    data = d0,
+    time = "year",
+    spatiotemporal = "off",
+    spatial = FALSE,
+    family = delta_truncated_nbinom1()
+  )
+  p <- predict(fit, newdata = d0, return_tmb_object = TRUE)
+  x <- get_index(p, area = 1/nrow(d0))
+  pp <- predict(fit, type = "response")
+  expect_equal(x$est, pp$est[1], tolerance = 1e-3)
+
+  # TNB2
+  dd <- subset(d0, y > 0)
+  fit <- sdmTMB(
+    y ~ 1,
+    data = dd,
+    time = "year",
+    spatiotemporal = "off",
+    spatial = FALSE,
+    family = truncated_nbinom2()
+  )
+  p <- predict(fit, newdata = d0, return_tmb_object = TRUE)
+  x <- get_index(p, area = 1/nrow(d0))
+  pp <- predict(fit, type = "response")
+  expect_equal(x$est, pp$est[1], tolerance = 1e-3)
+
+  # TNB1
+  fit <- sdmTMB(
+    y ~ 1,
+    data = dd,
+    time = "year",
+    spatiotemporal = "off",
+    spatial = FALSE,
+    family = truncated_nbinom2()
+  )
+  p <- predict(fit, newdata = d0, return_tmb_object = TRUE)
+  x <- get_index(p, area = 1/nrow(d0))
+  pp <- predict(fit, type = "response")
+  expect_equal(x$est, pp$est[1], tolerance = 1e-3)
+})