Use subsampling estimates for pseudo-BMA weights in loo_model_weights()

R/loo_model_weights.R

@@ -17,8 +17,16 @@
 #'   Currently the `loo_model_weights()` function is not implemented to be used
 #'   with results from K-fold CV, but you can still obtain weights using K-fold
 #'   CV results by calling the `stacking_weights()` or `pseudobma_weights()`
-#'   function directly. Objects from [loo_subsample()] are not supported,
-#'   because they have pointwise values only for the subsampled observations.
+#'   function directly. Objects from [loo_subsample()] are supported only with
+#'   `method = "pseudobma"` and `BB = FALSE`, which use the subsampling
+#'   estimates of `elpd_loo`. The weights do not account for the subsampling
+#'   uncertainty of these estimates. The subsamples can differ between the
+#'   models, but the differences between the estimates are more accurate when
+#'   all models use the same observations (pass the first object as the
+#'   `observations` argument of [loo_subsample()] for the other models).
+#'   Stacking and the Bayesian bootstrap need pointwise values for all
+#'   observations, and objects from [loo_subsample()] have them only for the
+#'   subsampled observations.
 #' @param method Either `"stacking"` (the default) or `"pseudobma"`, indicating which method
 #'   to use for obtaining the weights. `"stacking"` refers to stacking of
 #'   predictive distributions and  `"pseudobma"` refers to pseudo-BMA+ weighting
@@ -185,6 +193,7 @@ loo_model_weights.default <-
     cores <- loo_cores(cores)
     method <- match.arg(method)
     K <- length(x) # number of models
+    subsampled <- FALSE # set below for a list of loo objects
 
     if (is.matrix(x[[1]])) {
       N <- ncol(x[[1]]) # number of data points
@@ -200,9 +209,14 @@ loo_model_weights.default <-
         elpd_loo[k] <- loo_object$estimates["elpd_loo", "Estimate"]
       }
     } else if (is.psis_loo(x[[1]])) {
-      validate_psis_loo_list(x)
-      lpd_point <- do.call(cbind, lapply(x, function(obj) obj$pointwise[, "elpd_loo"]))
+      # pseudo-BMA without the Bayesian bootstrap needs only the elpd of each
+      # model, which loo_subsample() estimates for all observations
+      validate_psis_loo_list(x, allow_subsampling = method == "pseudobma" && !BB)
       elpd_loo <- sapply(x, function(obj) obj$estimates["elpd_loo", "Estimate"])
+      subsampled <- any(sapply(x, inherits, "psis_loo_ss"))
+      if (!subsampled) {
+        lpd_point <- do.call(cbind, lapply(x, function(obj) obj$pointwise[, "elpd_loo"]))
+      }
     } else {
       stop("'x' must be a list of matrices or a list of 'psis_loo' objects.")
     }
@@ -215,6 +229,10 @@ loo_model_weights.default <-
         optim_control = optim_control
       )
 
+    } else if (subsampled) {
+      # method == "pseudobma" with BB = FALSE (checked above): the subsampling
+      # estimates of elpd_loo
+      wts <- pseudobma_weights_from_elpd(elpd_loo)
     } else {
       # method =="pseudobma"
       wts <- pseudobma_weights(
@@ -328,13 +346,7 @@ pseudobma_weights <-
     }
 
     if (!BB) {
-      uwts <- exp(elpd - max(elpd))
-      wts <- structure(
-        uwts / sum(uwts),
-        names = paste0("model", 1:K),
-        class = "pseudobma_weights"
-      )
-      return(wts)
+      return(pseudobma_weights_from_elpd(elpd))
     }
 
     BB_weighting <- dirichlet_rng(BB_n, rep(alpha, N))
@@ -348,6 +360,29 @@ pseudobma_weights <-
   }
 
 
+#' Pseudo-BMA weights from the elpd of each model
+#'
+#' @noRd
+#' @param elpd Vector of elpd estimates, one for each model. Each must be finite
+#'   or `-Inf`, and at least one must be finite.
+#' @return Weights proportional to `exp(elpd)`, with class
+#'   `"pseudobma_weights"`.
+pseudobma_weights_from_elpd <- function(elpd) {
+  if (anyNA(elpd) || any(elpd == Inf)) {
+    stop("All elpd_loo estimates must be finite or -Inf.", call. = FALSE)
+  }
+  if (!any(is.finite(elpd))) {
+    stop("At least one model must have a finite total predictive density.",
+         call. = FALSE)
+  }
+  uwts <- exp(elpd - max(elpd))
+  structure(
+    uwts / sum(uwts),
+    names = paste0("model", seq_along(elpd)),
+    class = "pseudobma_weights"
+  )
+}
+
 #' @export
 print.stacking_weights <- function(x, digits = 3, ...) {
   cat("Method: stacking\n------\n")
@@ -442,7 +477,7 @@ validate_log_lik_list <- function(log_lik_list) {
   invisible(TRUE)
 }
 
-validate_psis_loo_list <- function(psis_loo_list) {
+validate_psis_loo_list <- function(psis_loo_list, allow_subsampling = FALSE) {
   stopifnot(is.list(psis_loo_list))
   if (length(psis_loo_list) < 2) {
     stop("At least two models are required.", call. = FALSE)
@@ -452,15 +487,25 @@ validate_psis_loo_list <- function(psis_loo_list) {
   }
   # the pointwise values of a subsampled object cover only its subsample, so
   # the objects would not have values for the same observations
-  if (any(sapply(psis_loo_list, inherits, "psis_loo_ss"))) {
-    stop("loo_model_weights() does not support objects from loo_subsample(), ",
-         "because they have pointwise values only for the subsampled ",
-         "observations. Use loo() for all observations.", call. = FALSE)
+  subsampled <- sapply(psis_loo_list, inherits, "psis_loo_ss")
+  if (any(subsampled) && !allow_subsampling) {
+    stop("loo_model_weights() supports objects from loo_subsample() only with ",
+         "method = \"pseudobma\" and BB = FALSE, because they have pointwise ",
+         "values only for the subsampled observations. Use loo() for all ",
+         "observations.", call. = FALSE)
   }
 
-  # the weights use only the pointwise elpd, so the number of posterior draws
-  # can differ between the objects
-  n_obs <- sapply(psis_loo_list, function(x) nrow(x$pointwise))
+  # the weights use only the pointwise elpd or its total, so the number of
+  # posterior draws can differ between the objects
+  # a subsampled object has an approximation of elpd_loo for every observation
+  # (nobs() is the size of the subsample)
+  n_obs <- sapply(psis_loo_list, function(x) {
+    if (inherits(x, "psis_loo_ss")) {
+      length(x$loo_subsampling$elpd_loo_approx)
+    } else {
+      nrow(x$pointwise)
+    }
+  })
   if (length(unique(n_obs)) != 1) {
     stop("Each object in the list must have the same number of observations.",
          call. = FALSE)

tests/testthat/test_model_weighting.R

@@ -70,7 +70,7 @@ test_that("loo_model_weights accepts models with different numbers of draws", {
   }
 })
 
-test_that("loo_model_weights rejects objects from loo_subsample()", {
+test_that("loo_model_weights accepts loo_subsample() objects only for pseudo-BMA", {
   # restore the random number stream, which later tests in this file use
   old_seed <- .Random.seed
   on.exit(assign(".Random.seed", old_seed, envir = globalenv()), add = TRUE)
@@ -83,26 +83,62 @@ test_that("loo_model_weights rejects objects from loo_subsample()", {
   post1 <- as.matrix(rbeta(500, 140, 60))
   post2 <- as.matrix(rbeta(500, 30, 20))
   # the worse model gives high Pareto k values and a warning
-  ss <- function(post) {
-    suppressWarnings(loo_subsample(llfun, draws = post, data = dat,
-                                   observations = 50, loo_approximation = "plpd",
-                                   r_eff = rep(1, N)))
+  ss <- function(post, observations, data = dat) {
+    suppressWarnings(loo_subsample(llfun, draws = post, data = data,
+                                   observations = observations,
+                                   loo_approximation = "plpd",
+                                   r_eff = rep(1, nrow(data))))
   }
-  loo_ss1 <- ss(post1)
-  loo_ss2 <- ss(post2)
-  # the two subsamples contain different observations
+  # different subsamples of different sizes
+  loo_ss1 <- ss(post1, 50)
+  loo_ss2 <- ss(post2, 80)
   expect_false(setequal(obs_idx(loo_ss1), obs_idx(loo_ss2)))
+  loo_full <- suppressWarnings(loo(llfun, draws = post1, data = dat, r_eff = rep(1, N)))
+  error <- "supports objects from loo_subsample() only with method = \"pseudobma\" and BB = FALSE"
+
+  # stacking and pseudo-BMA+ need pointwise values for all observations
+  expect_error(loo_model_weights(list(loo_ss1, loo_ss2)), error, fixed = TRUE)
+  expect_error(loo_model_weights(list(loo_ss1, loo_ss2), BB = FALSE), error,
+               fixed = TRUE)
+  expect_error(loo_model_weights(list(loo_full, loo_ss2)), error, fixed = TRUE)
   expect_error(
-    loo_model_weights(list(loo_ss1, loo_ss2)),
-    "does not support objects from loo_subsample()",
+    loo_model_weights(list(loo_ss1, loo_ss2), method = "pseudobma"),
+    error,
     fixed = TRUE
   )
-  loo_full <- suppressWarnings(loo(llfun, draws = post1, data = dat, r_eff = rep(1, N)))
+
+  # pseudo-BMA uses the subsampling estimates of elpd_loo
+  elpd <- c(loo_ss1$estimates["elpd_loo", "Estimate"],
+            loo_ss2$estimates["elpd_loo", "Estimate"])
+  w <- loo_model_weights(list(a = loo_ss1, b = loo_ss2), method = "pseudobma",
+                         BB = FALSE)
+  expect_s3_class(w, "pseudobma_weights")
+  expect_named(w, c("a", "b"))
+  expect_equal(as.numeric(w), exp(elpd - max(elpd)) / sum(exp(elpd - max(elpd))))
+
+  # a full loo object and a subsampled object can be combined
+  w2 <- loo_model_weights(list(loo_full, loo_ss2), method = "pseudobma", BB = FALSE)
+  elpd2 <- c(loo_full$estimates["elpd_loo", "Estimate"], elpd[2])
+  expect_equal(as.numeric(w2), exp(elpd2 - max(elpd2)) / sum(exp(elpd2 - max(elpd2))))
+
+  # the models must still have the same number of observations
+  loo_ss_short <- ss(post2, 50, data = dat[1:150, ])
   expect_error(
-    loo_model_weights(list(loo_full, loo_ss2)),
-    "does not support objects from loo_subsample()",
-    fixed = TRUE
+    loo_model_weights(list(loo_ss1, loo_ss_short), method = "pseudobma", BB = FALSE),
+    "same number of observations"
+  )
+})
+
+test_that("pseudo-BMA weights from elpd estimates reject NA and Inf", {
+  error <- "All elpd_loo estimates must be finite or -Inf."
+  expect_error(pseudobma_weights_from_elpd(c(-10, NA)), error, fixed = TRUE)
+  expect_error(pseudobma_weights_from_elpd(c(-10, NaN)), error, fixed = TRUE)
+  expect_error(pseudobma_weights_from_elpd(c(-10, Inf)), error, fixed = TRUE)
+  expect_error(
+    pseudobma_weights_from_elpd(c(-Inf, -Inf)),
+    "At least one model must have a finite total predictive density."
   )
+  expect_equal(as.numeric(pseudobma_weights_from_elpd(c(-10, -Inf))), c(1, 0))
 })
 
 test_that("loo_model_weights throws correct errors and warnings", {