From 42b1af334d6d0c5eb1fdc3b65ca81b756d3951f1 Mon Sep 17 00:00:00 2001
From: Sean Anderson <sean@seananderson.ca>
Date: Wed, 20 Mar 2024 20:38:06 -0700
Subject: [PATCH] Add mle-mvn approximate posterior residuals

---
 DESCRIPTION             |  2 +-
 NEWS.md                 |  5 +++++
 R/residuals.R           | 32 +++++++++++++++++++++++++++-----
 man/residuals.sdmTMB.Rd |  8 +++++++-
 4 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/DESCRIPTION b/DESCRIPTION
index dc6ea9590..c6ca78b0b 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,7 +1,7 @@
 Type: Package
 Package: sdmTMB
 Title: Spatial and Spatiotemporal SPDE-Based GLMMs with 'TMB'
-Version: 0.4.2.9003
+Version: 0.4.2.9004
 Authors@R: c(
     person(c("Sean", "C."), "Anderson", , "sean@seananderson.ca",
            role = c("aut", "cre"),
diff --git a/NEWS.md b/NEWS.md
index 2f2cc5852..5bb7512d6 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,5 +1,10 @@
 # sdmTMB (development version)
 
+* Add experimental residuals option "mle-mvn" where a single approximate 
+  posterior sample of the random effects is drawn and these are combined
+  with the MLE fixed effects to produce residuals. This may become the
+  default option.
+
 * Add the generalized gamma distribution (thanks to J.T. Thorson with additional
   work by J.C. Dunic.) See `gengamma()`. This distribution is still in a testing
   phase and is not recommended for applied use yet. #286
diff --git a/R/residuals.R b/R/residuals.R
index 4c0a13757..a0a80cb88 100644
--- a/R/residuals.R
+++ b/R/residuals.R
@@ -240,10 +240,16 @@ qres_beta <- function(object, y, mu, ...) {
 #'   qqnorm(r1)
 #'   qqline(r1)
 #'
+#'   # "mle-mvn" residuals with the fixed effects at their MLE and the random
+#'   # effects sampled from their approximate posterior
+#'   r2 <- residuals(fit, type = "mle-mvn")
+#'   qqnorm(r2)
+#'   qqline(r2)
+#'
 #'   # see also "mle-mcmc" residuals with the help of the sdmTMBextra package
 
 residuals.sdmTMB <- function(object,
-                             type = c("mle-laplace", "mle-mcmc", "mvn-laplace", "response", "pearson"),
+                             type = c("mle-laplace", "mle-mcmc", "mle-mvn", "mvn-laplace", "response", "pearson"),
                              model = c(1, 2),
                              mcmc_samples = NULL,
                              qres_func = NULL,
@@ -301,9 +307,6 @@ residuals.sdmTMB <- function(object,
   }
 
   if (type %in% c("mle-laplace", "response", "pearson")) {
-    # mu <- tryCatch({linkinv(predict(object, newdata = NULL)[[est_column]])}, # newdata = NULL; fast
-    #   error = function(e) NA)
-    # if (is.na(mu[[1]])) {
     mu <- linkinv(predict(object, newdata = object$data, offset = object$tmb_data$offset_i)[[est_column]]) # not newdata = NULL
     # }
   } else if (type == "mvn-laplace") {
@@ -318,6 +321,25 @@ residuals.sdmTMB <- function(object,
     mcmc_samples <- as.numeric(mcmc_samples)
     assert_that(length(mcmc_samples) == nrow(object$data))
     mu <- linkinv(mcmc_samples)
+  } else if (type == "mle-mvn") {
+    ## see TMB:::oneSamplePosterior()
+    tmp <- object$tmb_obj$env$MC(n = 1L, keep = TRUE, antithetic = FALSE)
+    re_samp <- as.vector(attr(tmp, "samples"))
+    lp <- object$tmb_obj$env$last.par.best
+    p <- numeric(length(lp))
+    fe <- object$tmb_obj$env$lfixed()
+    re <- object$tmb_obj$env$lrandom()
+    p[re] <- re_samp
+    p[fe] <- lp[fe]
+    pred <- predict(
+      object,
+      newdata = object$data,
+      mcmc_samples = matrix(p, ncol = 1L),
+      model = model[[1L]],
+      nsim = 1L,
+      offset = object$tmb_data$offset_i
+    )
+    mu <- linkinv(pred[, 1L, drop = TRUE])
   } else {
     cli_abort("residual type not implemented")
   }
@@ -335,7 +357,7 @@ residuals.sdmTMB <- function(object,
 
   if (type == "response") {
     if (!prop_binomial) r <- y - mu else r <- y / size - mu
-  } else if (type == "mle-laplace" || type == "mvn-laplace") {
+  } else if (type == "mle-laplace" || type == "mvn-laplace" || type == "mle-mvn") {
     r <- res_func(object, y, mu, .n = size, ...)
   } else if (type == "mle-mcmc") {
     r <- res_func(object, y, mu, .n = size, ...)
diff --git a/man/residuals.sdmTMB.Rd b/man/residuals.sdmTMB.Rd
index 731d4fb1b..9fc3dcf44 100644
--- a/man/residuals.sdmTMB.Rd
+++ b/man/residuals.sdmTMB.Rd
@@ -6,7 +6,7 @@
 \usage{
 \method{residuals}{sdmTMB}(
   object,
-  type = c("mle-laplace", "mle-mcmc", "mvn-laplace", "response", "pearson"),
+  type = c("mle-laplace", "mle-mcmc", "mle-mvn", "mvn-laplace", "response", "pearson"),
   model = c(1, 2),
   mcmc_samples = NULL,
   qres_func = NULL,
@@ -92,6 +92,12 @@ based on simulations drawn from the assumed multivariate normal distribution
   qqnorm(r1)
   qqline(r1)
 
+  # "mle-mvn" residuals with the fixed effects at their MLE and the random
+  # effects sampled from their approximate posterior
+  r2 <- residuals(fit, type = "mle-mvn")
+  qqnorm(r2)
+  qqline(r2)
+
   # see also "mle-mcmc" residuals with the help of the sdmTMBextra package
 }
 \references{