[avehtari]

Summary

This PR runs kmeans() in get_p_clust() on the coordinates of the reference model draws in their low-dimensional span instead of on the S × n matrix t(eta). The clusters stay the same. varsel() with 50 000 observations, 4 000 draws and 20 clusters gets about 4.5 times faster. For small numbers of observations, the run time does not change.

This PR was assisted by Claude.

Overhead

In current master: get_p_clust() runs kmeans() on t(eta), an S × n matrix, so each k-means iteration costs O(S K n) for K clusters. The draws of a GLM-type reference model span at most as many dimensions as the model has coefficients, often far fewer than n. In a profile of varsel() for a Gaussian model with n = 10 000, 4 000 draws and 20 clusters, kmeans() took 65 % of the time.

Changes

• The new internal eta_for_kmeans() centers the draws and finds an orthonormal basis Q of their span with a pivoted qr() (rank tolerance 1e-10) of a subset of the draws. It starts with 64 evenly spaced draws and doubles the subset until the subset is rank deficient. It accepts Q only if Q reproduces all draws (relative tolerance 1e-8), and then returns the S × r matrix of coordinates.
• get_p_clust() calls kmeans() on the result of eta_for_kmeans().

The Euclidean distances between the draws, and from a draw to any mean of draws, are the same in the coordinates as in t(eta). kmeans() takes its initial centers with sample.int(S, K), which gives the same rows for both inputs. So k-means finds the same clustering up to rounding. The subset is deterministic, so the function does not change the state of the random number generator. If the span has more than min(n, S) / 4 dimensions, if the check fails, or if eta has missing values, kmeans() gets t(eta) as before.

Results

Gaussian reference model with P = 50 and 4 000 draws, 20 search clusters; all results identical to master (difference 0, identical rankings):

case	master	this PR
varsel(), n = 50 000, 20 evaluation clusters	792 s	174 s
varsel(), n = 10 000, 20 evaluation clusters	130 s	46 s
varsel(), n = 10 000, 400 evaluation draws	119 s	81 s
varsel(), n = 1 000, 20 evaluation clusters	19 s	16 s

For small n (N = 50, P = 200), the run time does not change. In a local check with rank 21, n = 5 000 and S = 1 000, the clusters were identical and kmeans() took 0.18 s instead of 2.8 s. The peak memory is about the same as on master: eta_for_kmeans() holds about two n × S matrices for a short time, and on master t(eta) and the copy that kmeans() passes to Fortran also take two.

Tests

The new test_eta_for_kmeans.R checks that:

• for low-rank draws (rank 6, n = 500, S = 300), eta_for_kmeans() returns at most 6 columns, keeps the pairwise distances (tolerance 1e-10), and gives the same kmeans() clusters as t(eta);
• for full-rank draws, it returns t(eta);
• it does not change the state of the random number generator.

All three tests fail on master (function not found). The full suite passes with NOT_CRAN=true: 256 tests, 55 812 expectations, 0 failures, the same 4 warnings as master. Tested with this change applied to master 73cf0ee (projpred 2.11.0.9000), R 4.6.1, brms 2.23.0, rstanarm 2.32.2, cmdstanr 0.9.0.9001, CmdStan 2.39.0, on CSC Roihu (Linux, AMD EPYC).