Parallelize some functions

.Rbuildignore

@@ -19,7 +19,7 @@ vignettes/loo2-non-factorized_cache/*
 .vscode/*
 ^\.github$
 ^vignettes/online-only$
-
+^\.git-blame-ignore-revs$
 ^CRAN-SUBMISSION$
 ^release-prep\.R$
 ^_pkgdown\.yml$

R/loo_subsample.R

@@ -536,7 +536,7 @@ elpd_loo_approximation <- function(.llfun, data, draws, cores, loo_approximation
 
   if (loo_approximation %in% c("tis", "sis")) {
     draws <- .thin_draws(draws, loo_approximation_draws)
-    is_values <- suppressWarnings(loo.function(.llfun, data = data, draws = draws, is_method = loo_approximation))
+    is_values <- suppressWarnings(loo.function(.llfun, data = data, draws = draws, is_method = loo_approximation, cores = cores))
     return(is_values$pointwise[, "elpd_loo"])
   }
 

R/psis_approximate_posterior.R

@@ -97,11 +97,10 @@ ap_psis.array <-
     log_ratios <- validate_ll(log_ratios)
     log_ratios <- llarray_to_matrix(log_ratios)
 
-    r_eff <- prepare_psis_r_eff(r_eff, len = ncol(log_ratios))
     ap_psis.matrix(log_ratios = log_ratios,
                    log_p = log_p,
                    log_g = log_g,
-                   cores = 1)
+                   cores = cores)
   }
 
 #' @export
@@ -132,4 +131,3 @@ ap_psis.default <- function(log_ratios, log_p, log_g, ...) {
     ap_psis.matrix(as.matrix(log_ratios), log_p, log_g, cores = 1)
   }
 
-

R/waic.R

@@ -100,27 +100,68 @@ waic.function <-
   function(x,
            ...,
            data = NULL,
-           draws = NULL) {
+           draws = NULL,
+           cores = getOption("mc.cores", 1)) {
+    cores <- loo_cores(cores)
     stopifnot(is.data.frame(data) || is.matrix(data), !is.null(draws))
 
     .llfun <- validate_llfun(x)
     N <- dim(data)[1]
     S <- length(as.vector(.llfun(data_i = data[1,, drop=FALSE], draws = draws, ...)))
-    waic_list <- lapply(seq_len(N), FUN = function(i) {
-      ll_i <- .llfun(data_i = data[i,, drop=FALSE], draws = draws, ...)
-      ll_i <- as.vector(ll_i)
-      lpd_i <- logMeanExp(ll_i)
-      p_waic_i <- var(ll_i)
-      elpd_waic_i <- lpd_i - p_waic_i
-      c(elpd_waic = elpd_waic_i, p_waic = p_waic_i)
-    })
+    if (cores == 1) {
+      waic_list <-
+        lapply(
+          X = seq_len(N),
+          FUN = .waic_i,
+          llfun = .llfun,
+          data = data,
+          draws = draws,
+          ...
+        )
+    } else {
+      if (!os_is_windows()) {
+        # On Mac or Linux use mclapply() for multiple cores
+        waic_list <-
+          parallel::mclapply(
+            mc.cores = cores,
+            X = seq_len(N),
+            FUN = .waic_i,
+            llfun = .llfun,
+            data = data,
+            draws = draws,
+            ...
+          )
+      } else {
+        # On Windows use makePSOCKcluster() and parLapply() for multiple cores
+        cl <- parallel::makePSOCKcluster(cores)
+        on.exit(parallel::stopCluster(cl))
+        waic_list <-
+          parallel::parLapply(
+            cl = cl,
+            X = seq_len(N),
+            fun = .waic_i,
+            llfun = .llfun,
+            data = data,
+            draws = draws,
+            ...
+          )
+      }
+    }
     pointwise <- do.call(rbind, waic_list)
     pointwise <- cbind(pointwise, waic = -2 * pointwise[, "elpd_waic"])
 
     throw_pwaic_warnings(pointwise[, "p_waic"], digits = 1)
     waic_object(pointwise, dims = c(S, N))
   }
 
+.waic_i <- function(i, llfun, data, draws, ...) {
+  ll_i <- llfun(data_i = data[i,, drop=FALSE], draws = draws, ...)
+  ll_i <- as.vector(ll_i)
+  lpd_i <- logMeanExp(ll_i)
+  p_waic_i <- var(ll_i)
+  elpd_waic_i <- lpd_i - p_waic_i
+  c(elpd_waic = elpd_waic_i, p_waic = p_waic_i)
+}
 
 #' @export
 dim.waic <- function(x) {
@@ -164,4 +205,3 @@ throw_pwaic_warnings <- function(p, digits = 1, warn = TRUE) {
   }
   invisible(NULL)
 }
-

tests/testthat/test_loo_and_waic.R

@@ -223,6 +223,22 @@ test_that("loo.function runs with multiple cores", {
   expect_identical(loo_with_fn2$estimates, loo_with_fn1$estimates)
 })
 
+test_that("waic.function runs with multiple cores", {
+  waic_with_fn1 <- waic(
+    llfun,
+    data = data,
+    draws = draws,
+    cores = 1
+  )
+  waic_with_fn2 <- waic(
+    llfun,
+    data = data,
+    draws = draws,
+    cores = 2
+  )
+  expect_identical(waic_with_fn2$estimates, waic_with_fn1$estimates)
+})
+
 test_that("save_psis option to loo.function makes correct psis object", {
   loo_with_fn2 <- loo.function(
     llfun,

tests/testthat/test_psis_approximate_posterior.R

@@ -325,3 +325,18 @@ test_that("Deprecation of log_q argument", {
   expect_s3_class(psis_lap, "psis")
   expect_lt(pareto_k_values(psis_lap), 0.7)
 })
+
+test_that("ap_psis.array works as ap_psis.matrix", {
+  log_p <- test_data_psis_approximate_posterior$laplace_independent$log_p
+  log_g <- test_data_psis_approximate_posterior$laplace_independent$log_q
+  ll <- test_data_psis_approximate_posterior$laplace_independent$log_liks
+  ll_array <- array(ll, dim = c(nrow(ll) / 2, 2, ncol(ll)))
+  ll_matrix <- loo:::llarray_to_matrix(ll_array)
+
+  expect_silent(
+    psis_array <- ap_psis(log_ratios = -ll_array, log_p = log_p, log_g = log_g)
+  )
+  psis_matrix <- ap_psis(log_ratios = -ll_matrix, log_p = log_p, log_g = log_g)
+
+  expect_equal(psis_array, psis_matrix)
+})