Vectorize the draw-wise predictive variance of the Gaussian and Student-t

NEWS.md

@@ -7,6 +7,8 @@ If you read this from a place other than <https://mc-stan.org/projpred/news/inde
 
 ## Major changes
 
+* Faster computation of the predictive variance of the reference model for the Gaussian and Student-t families. This matters for a large number of observations and many projected draws. For example, `project()` with 10 000 observations and 4 000 projected draws is about 7 times faster. (GitHub: #575)
+
 ## Minor changes
 
 ## Bug fixes and maintenance

R/extend_family.R

@@ -600,6 +600,15 @@ extend_family_poisson <- function(family) {
   return(family)
 }
 
+# For each observation (row of `mu`), the weighted variance of `mu` across the
+# draws (columns of `mu`), with weights `wdraws` (the caller normalizes them to
+# sum to 1). Vectorized form of
+# `sum(wdraws * (mu[i, ] - weighted.mean(mu[i, ], wdraws))^2)`.
+rowvar_wdraws <- function(mu, wdraws) {
+  mu_mean <- drop(mu %*% wdraws) / sum(wdraws)
+  return(drop(((mu - mu_mean)^2) %*% wdraws))
+}
+
 extend_family_gaussian <- function(family) {
   # ce_gauss() does not give the actual cross-entropy (not even the one which
   # would result from dropping terms which would cancel out when calculating the
@@ -615,9 +624,7 @@ extend_family_gaussian <- function(family) {
   }
   predvar_gauss <- function(mu, dis, wdraws = 1) {
     wdraws <- wdraws / sum(wdraws)
-    mu_var <- do.call(rbind, lapply(seq_len(nrow(mu)), function(i) {
-      sum(wdraws * (mu[i, ] - weighted.mean(mu[i, ], wdraws))^2)
-    }))
+    mu_var <- rowvar_wdraws(mu, wdraws)
     as.vector(sum(wdraws * dis^2) + mu_var)
   }
   ll_gauss <- function(mu, dis, y, weights = 1) {
@@ -718,9 +725,7 @@ extend_family_student_t <- function(family) {
   }
   predvar_student_t <- function(mu, dis, wdraws = 1) {
     wdraws <- wdraws / sum(wdraws)
-    mu_var <- do.call(rbind, lapply(seq_len(nrow(mu)), function(i) {
-      sum(wdraws * (mu[i, ] - weighted.mean(mu[i, ], wdraws))^2)
-    }))
+    mu_var <- rowvar_wdraws(mu, wdraws)
     as.vector(family$nu / (family$nu - 2) * sum(wdraws * dis^2) + mu_var)
   }
   ll_student_t <- function(mu, dis, y, weights = 1) {

tests/testthat/test_predvar.R

@@ -0,0 +1,49 @@
+context("predictive variance of the Gaussian and Student-t families")
+
+# Per-observation weighted variance across the draws, computed row by row (the
+# reference for rowvar_wdraws()):
+rowvar_loop <- function(mu, wdraws) {
+  vapply(seq_len(nrow(mu)), function(i) {
+    sum(wdraws * (mu[i, ] - weighted.mean(mu[i, ], wdraws))^2)
+  }, numeric(1))
+}
+
+predvar_inputs <- function(n = 30, S = 7, seed = 8152) {
+  set.seed(seed)
+  mu <- matrix(rnorm(n * S, mean = 2), n, S)
+  wdraws <- runif(S, 0.2, 1)
+  list(mu = mu, dis = runif(S, 0.5, 1.5), wdraws = wdraws / sum(wdraws))
+}
+
+test_that("rowvar_wdraws() equals the row-wise weighted variance", {
+  inp <- predvar_inputs()
+  expect_equal(rowvar_wdraws(inp$mu, inp$wdraws),
+               rowvar_loop(inp$mu, inp$wdraws), tolerance = 1e-12)
+  # Equal weights:
+  w_eq <- rep(1 / ncol(inp$mu), ncol(inp$mu))
+  expect_equal(rowvar_wdraws(inp$mu, w_eq), rowvar_loop(inp$mu, w_eq),
+               tolerance = 1e-12)
+  # One draw: the variance is exactly zero.
+  expect_identical(rowvar_wdraws(inp$mu[, 1, drop = FALSE], 1),
+                   rep(0, nrow(inp$mu)))
+  # A missing value gives a missing variance for its observation only.
+  mu_na <- inp$mu
+  mu_na[4, 2] <- NA
+  v_na <- rowvar_wdraws(mu_na, inp$wdraws)
+  expect_identical(which(is.na(v_na)), 4L)
+  expect_equal(v_na[-4], rowvar_loop(mu_na, inp$wdraws)[-4], tolerance = 1e-12)
+})
+
+test_that(paste(
+  "predvar() of the Gaussian and Student-t families adds the weighted",
+  "variance of mu to the mean of the squared dispersion"
+), {
+  inp <- predvar_inputs()
+  v_mu <- rowvar_loop(inp$mu, inp$wdraws)
+  fam_gauss <- extend_family(gaussian())
+  expect_equal(fam_gauss$predvar(inp$mu, inp$dis, inp$wdraws),
+               sum(inp$wdraws * inp$dis^2) + v_mu, tolerance = 1e-12)
+  fam_t <- extend_family(Student_t(nu = 5))
+  expect_equal(fam_t$predvar(inp$mu, inp$dis, inp$wdraws),
+               5 / 3 * sum(inp$wdraws * inp$dis^2) + v_mu, tolerance = 1e-12)
+})