From a7988ca62e9a0c161444dfb3dc080eb45cc04102 Mon Sep 17 00:00:00 2001 From: Aki Vehtari Date: Thu, 22 Oct 2026 23:17:13 +0300 Subject: [PATCH 1/2] Vectorize the draw-wise predictive variance of the Gaussian and Student-t families --- NEWS.md | 2 ++ R/extend_family.R | 17 +++++++----- tests/testthat/test_predvar.R | 49 +++++++++++++++++++++++++++++++++++ 3 files changed, 62 insertions(+), 6 deletions(-) create mode 100644 tests/testthat/test_predvar.R diff --git a/NEWS.md b/NEWS.md index 94ffe21cd..5006039ff 100644 --- a/NEWS.md +++ b/NEWS.md @@ -7,6 +7,8 @@ If you read this from a place other than ) + ## Minor changes ## Bug fixes and maintenance diff --git a/R/extend_family.R b/R/extend_family.R index e3f9c3999..de9726733 100644 --- a/R/extend_family.R +++ b/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) { diff --git a/tests/testthat/test_predvar.R b/tests/testthat/test_predvar.R new file mode 100644 index 000000000..c2f82c7e0 --- /dev/null +++ b/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) +}) From f28c0ba7dbeb81415fa4bb8d7ccbc45c48a7f025 Mon Sep 17 00:00:00 2001 From: Aki Vehtari Date: Sat, 24 Oct 2026 15:29:41 +0300 Subject: [PATCH 2/2] Add the PR number to the NEWS entry --- NEWS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/NEWS.md b/NEWS.md index 5006039ff..6290978d1 100644 --- a/NEWS.md +++ b/NEWS.md @@ -7,7 +7,7 @@ If you read this from a place other than ) +* 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