Allow different numbers of draws in loo_model_weights()

NEWS.md

@@ -1,5 +1,6 @@
 # loo (development version)
 
+* Allow different numbers of draws in `loo_model_weights()` by @avehtari in #409
 * Improve numerical stability in `loo()`, `psis()`, model weighting, subsampling, 
 and moment matching in #395 
 * Fix `loo_compare()` when used with subsampling: compute model comparison by 

R/loo_model_weights.R

@@ -11,12 +11,14 @@
 #'   elements are named the names will be used to label the models in the
 #'   results. Each matrix/object should have dimensions \eqn{S} by \eqn{N},
 #'   where \eqn{S} is the size of the posterior sample (with all chains merged)
-#'   and \eqn{N} is the number of data points. If `x` is a list of
+#'   and \eqn{N} is the number of data points. All models must have the same
+#'   \eqn{N}, but \eqn{S} can differ between models. If `x` is a list of
 #'   log-likelihood matrices then [loo()] is called internally on each matrix.
 #'   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.
+#'   function directly. Objects from [loo_subsample()] are not supported,
+#'   because they have pointwise values 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
@@ -420,7 +422,8 @@ validate_r_eff_list <- function(r_eff_list, K, N) {
 #' Validate log-likelihood list argument
 #'
 #' Checks that log-likelihood list has at least 2 elements and that each element
-#' has the same dimensions.
+#' has the same number of observations (columns). The number of posterior
+#' draws (rows) can differ, because the weights use only the pointwise elpd.
 #'
 #' @noRd
 #' @param log_lik_list User's list of log-likelihood matrices (the `x` argument
@@ -432,9 +435,9 @@ validate_log_lik_list <- function(log_lik_list) {
   if (length(log_lik_list) < 2) {
     stop("At least two models are required.", call. = FALSE)
   }
-  if (length(unique(sapply(log_lik_list, ncol))) != 1 |
-     length(unique(sapply(log_lik_list, nrow))) != 1) {
-    stop("Each log-likelihood matrix must have the same dimensions.", call. = FALSE)
+  if (length(unique(sapply(log_lik_list, ncol))) != 1) {
+    stop("Each log-likelihood matrix must have the same number of ",
+         "observations (columns).", call. = FALSE)
   }
   invisible(TRUE)
 }
@@ -447,11 +450,20 @@ validate_psis_loo_list <- function(psis_loo_list) {
   if (!all(sapply(psis_loo_list, is.psis_loo))) {
     stop("List elements must all be 'psis_loo' objects or log-likelihood matrices.")
   }
+  # 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)
+  }
 
-  dims <- sapply(psis_loo_list, dim)
-  if (length(unique(dims[1, ])) != 1 |
-      length(unique(dims[2, ])) != 1) {
-    stop("Each object in the list must have the same dimensions.", 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))
+  if (length(unique(n_obs)) != 1) {
+    stop("Each object in the list must have the same number of observations.",
+         call. = FALSE)
   }
   invisible(TRUE)
 }

tests/testthat/test_model_weighting.R

@@ -55,6 +55,56 @@ test_that("stacking handles equal infinite log predictive densities", {
   expect_equal(as.numeric(stacking_weights(lpd)), c(0, 1, 0), tolerance = 1e-6)
 })
 
+test_that("loo_model_weights accepts models with different numbers of draws", {
+  ll_short <- list(log_lik1, log_lik2[1:250, ], log_lik3)
+  loo_short <- loo_list
+  loo_short[[2]] <- loo(log_lik2[1:250, ], r_eff = r_eff_list[[2]])
+
+  for (method in c("stacking", "pseudobma")) {
+    w_mat <- loo_model_weights(ll_short, method = method, BB = FALSE,
+                               r_eff_list = r_eff_list)
+    w_loo <- loo_model_weights(loo_short, method = method, BB = FALSE)
+    expect_equal(as.numeric(w_mat), as.numeric(w_loo))
+    expect_true(all(w_loo >= 0))
+    expect_equal(sum(w_loo), 1)
+  }
+})
+
+test_that("loo_model_weights rejects objects from loo_subsample()", {
+  # 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)
+  set.seed(409)
+  N <- 200
+  dat <- data.frame(y = rbinom(N, size = 10, prob = 0.7), K = 10)
+  llfun <- function(data_i, draws) {
+    dbinom(data_i$y, size = data_i$K, prob = draws, log = TRUE)
+  }
+  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)))
+  }
+  loo_ss1 <- ss(post1)
+  loo_ss2 <- ss(post2)
+  # the two subsamples contain different observations
+  expect_false(setequal(obs_idx(loo_ss1), obs_idx(loo_ss2)))
+  expect_error(
+    loo_model_weights(list(loo_ss1, loo_ss2)),
+    "does not support objects from loo_subsample()",
+    fixed = TRUE
+  )
+  loo_full <- suppressWarnings(loo(llfun, draws = post1, data = dat, r_eff = rep(1, N)))
+  expect_error(
+    loo_model_weights(list(loo_full, loo_ss2)),
+    "does not support objects from loo_subsample()",
+    fixed = TRUE
+  )
+})
+
 test_that("loo_model_weights throws correct errors and warnings", {
   expect_error(
     loo_model_weights(log_lik1),
@@ -67,18 +117,14 @@ test_that("loo_model_weights throws correct errors and warnings", {
     "At least two models"
   )
 
-  expect_error(
-    loo_model_weights(list(log_lik1, log_lik2[-1, ])),
-    "same dimensions"
-  )
   expect_error(
     loo_model_weights(list(log_lik1, log_lik2, log_lik3[, -1])),
-    "same dimensions"
+    "same number of observations"
   )
 
   loo_list2 <- loo_list
-  attr(loo_list2[[3]], "dims") <- c(10, 10)
-  expect_error(loo_model_weights(loo_list2), "same dimensions")
+  loo_list2[[3]] <- loo(log_lik3[, -1], r_eff = rep(0.9, 49))
+  expect_error(loo_model_weights(loo_list2), "same number of observations")
 
   expect_error(
     loo_model_weights(ll_list, r_eff_list = r_eff_list[-1]),