From 55a516a0a527121b183c3de3a2c635f8c0a79192 Mon Sep 17 00:00:00 2001 From: Aki Vehtari Date: Sat, 24 Oct 2026 15:35:35 +0300 Subject: [PATCH 1/2] Cluster the reference model draws in the coordinates of their low-dimensional span --- NEWS.md | 2 ++ R/misc.R | 45 +++++++++++++++++++++++++++- tests/testthat/test_eta_for_kmeans.R | 38 +++++++++++++++++++++++ 3 files changed, 84 insertions(+), 1 deletion(-) create mode 100644 tests/testthat/test_eta_for_kmeans.R diff --git a/NEWS.md b/NEWS.md index 94ffe21cd..442546848 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/misc.R b/R/misc.R index 28ee8b9c9..bab4e4ef9 100644 --- a/R/misc.R +++ b/R/misc.R @@ -379,6 +379,49 @@ get_refdist <- function(refmodel, ndraws = NULL, nclusters = NULL, return(p_ref) } +# Returns the draws (columns) of `eta` as an S x r matrix of coordinates in an +# orthonormal basis of their centered span. These coordinates give the same +# Euclidean distances as `t(eta)`, so kmeans() finds the same clustering (up to +# rounding). The cost of kmeans() then does not grow with the number of +# observations. The draws of a GLM-type reference model span at most as many +# dimensions as the model has coefficients. The basis comes from a +# deterministic subset of the draws, so the state of the random number +# generator does not change. The function uses the basis only if it reproduces +# all draws (relative tolerance `100 * tol`; `tol` is the rank tolerance of +# qr()). +# Otherwise, the function returns `t(eta)`. +eta_for_kmeans <- function(eta, rank_max = floor(min(dim(eta)) / 4), + tol = 1e-10) { + S <- ncol(eta) + if (rank_max < 2 || anyNA(eta)) { + return(t(eta)) + } + eta_c <- eta - rowMeans(eta) + scale_c <- max(abs(eta_c)) + if (scale_c == 0) { + return(t(eta)) + } + k <- min(64L, rank_max) + repeat { + idx <- unique(round(seq(1, S, length.out = k))) + qr_sub <- qr(eta_c[, idx, drop = FALSE], tol = tol) + if (qr_sub$rank < length(idx)) { + # The subset has a lower rank than its number of draws, so its span can + # be the span of all draws (checked next): + q <- qr.Q(qr_sub)[, seq_len(qr_sub$rank), drop = FALSE] + coords <- crossprod(q, eta_c) + if (max(abs(eta_c - q %*% coords)) <= tol * scale_c * 1e2) { + return(t(coords)) + } + return(t(eta)) + } + if (k >= rank_max) { + return(t(eta)) + } + k <- min(2L * k, rank_max) + } +} + # Function for clustering the parameter draws: get_p_clust <- function(family, eta, mu, mu_offs, dis, nclusters = 10, wobs = rep(1, dim(mu)[1]), @@ -389,7 +432,7 @@ get_p_clust <- function(family, eta, mu, mu_offs, dis, nclusters = 10, # Note: A seed is not set here because this function is not exported and has # a calling stack at the beginning of which a seed is set. - out <- kmeans(t(eta), nclusters, iter.max = 50) + out <- kmeans(eta_for_kmeans(eta), nclusters, iter.max = 50) cl <- out$cluster # cluster indices for each draw } else if (typeof(cl) == "list") { # old clustering solution provided, so fetch the cluster indices diff --git a/tests/testthat/test_eta_for_kmeans.R b/tests/testthat/test_eta_for_kmeans.R new file mode 100644 index 000000000..e189db394 --- /dev/null +++ b/tests/testthat/test_eta_for_kmeans.R @@ -0,0 +1,38 @@ +context("eta_for_kmeans()") + +test_that(paste( + "eta_for_kmeans() reduces low-rank draws and keeps their distances and the", + "k-means clustering" +), { + set.seed(8571) + x <- cbind(1, matrix(rnorm(500 * 5), 500, 5)) + beta <- matrix(rnorm(6 * 300), 6, 300) + eta <- x %*% beta + coords <- eta_for_kmeans(eta) + expect_identical(nrow(coords), 300L) + expect_lte(ncol(coords), 6L) + expect_equal(as.matrix(dist(coords)), as.matrix(dist(t(eta))), + tolerance = 1e-10) + set.seed(2318) + cl_full <- kmeans(t(eta), 5, iter.max = 50)$cluster + set.seed(2318) + cl_coords <- kmeans(coords, 5, iter.max = 50)$cluster + expect_identical(cl_coords, cl_full) +}) + +test_that("eta_for_kmeans() returns `t(eta)` for full-rank draws", { + set.seed(8571) + eta <- matrix(rnorm(300 * 200), 300, 200) + expect_identical(eta_for_kmeans(eta), t(eta)) +}) + +test_that("eta_for_kmeans() does not change the state of the RNG", { + set.seed(8571) + x <- cbind(1, matrix(rnorm(500 * 5), 500, 5)) + eta <- x %*% matrix(rnorm(6 * 300), 6, 300) + set.seed(5544) + invisible(eta_for_kmeans(eta)) + u_after <- runif(1) + set.seed(5544) + expect_identical(runif(1), u_after) +}) From 9c12d54e23a3dc931edb2c049c8ede66a30f3932 Mon Sep 17 00:00:00 2001 From: Aki Vehtari Date: Sat, 24 Oct 2026 15:56:06 +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 442546848..1ed9e6a6a 100644 --- a/NEWS.md +++ b/NEWS.md @@ -7,7 +7,7 @@ If you read this from a place other than ) +* Faster clustering of the reference model draws for a large number of observations: `kmeans()` runs on the coordinates of the draws in their low-dimensional span, which gives the same clusters. For example, `varsel()` with 50 000 observations, 4 000 draws and 20 clusters is about 4.5 times faster. (GitHub: #576) ## Minor changes