[avehtari]

Summary

In cv_varsel() with validate_search = TRUE and LOO CV, get_p_clust() reweights the clusters of the reference draws twice per left-out observation (for the search and for the evaluation). It loops over the clusters and calls which(cl == j) in each iteration. Each iteration computes the center, dispersion and predictive variance of one cluster.

With 400 clusters, this PR reduces the time of cv_varsel() with validate_search = TRUE and full LOO by 54% (table below).

This PR was assisted by Claude.

Change

This PR replaces the loop in two cases:

• One draw per cluster (thinned draws, or as many clusters as draws): matrix operations apply the same weight factor 1 - 1e-10 as the loop. The results are identical.
• Several draws per cluster: the new C++ function clust_moments_c() makes two passes over the draws. It computes the weighted means of mu, mu_offs, the dispersion and the squared dispersion, and the weighted variance of mu_offs. The new family element predvar_mix turns these moments into the predictive variances of all clusters. The Gaussian, Student-t, binomial and Poisson families have it, and so does the latent projection (through its Gaussian family).

A user can replace family$predvar. So get_p_clust() uses predvar_mix only if family$predvar is identical to the "predvar" attribute of predvar_mix. In all other cases, get_p_clust() computes the predictive variances in the old loop over the clusters. Examples are the Gamma family, augmented-data families and a replaced predvar. get_p_clust() also uses the old loop for clusters with a weight sum of 0 or a non-finite weight sum.

Speed benchmark

The table gives the time of one run of cv_varsel() with validate_search = TRUE and full LOO per version. Both versions ran in the same SLURM job.

model	master (s)	this PR (s)
Gaussian, N 382, 26 predictor terms, 4000 draws, nclusters = 400, ndraws_pred = 4000	194.7	90.0
Gaussian, N 50, 200 predictor terms, nclusters = 20, ndraws_pred = 400	53.1	53.7

In row 1, the predictor ranking is the same, and all 24 compared elements of the vsel object are identical. In row 2, the predictor ranking is the same. With 20 clusters, the time does not change measurably.

Tests

test_get_p_clust.R (535 expectations) keeps a copy of the old loop. It compares get_p_clust() with the loop for several families and weights, clusters with weight 0, non-finite values and a replaced predvar. It checks that the Gaussian, Student-t, binomial, Poisson and latent families take the predvar_mix path, and it tests the input checks of clust_moments_c(). The full test suite passes on Triton (R 4.5.3, testthat 3.3.2) with 0 failures and the same 19 warnings as on master 7da1b5a.

[avehtari]

Thanks. Fewer cases need the second loop than use_mix covers. It computes the variance of mu_offs, and the first loop also computes the mean of dis^2. Only the Gaussian and Student-t predvar_mix() use these two moments, and the latent projection uses the Gaussian predvar_mix(). For binomial and Poisson, predvar_mix() returns an NA matrix and uses only the length of dis2.

The new commit adds the attribute uses_moments to predvar_mix() (FALSE for binomial and Poisson). It also adds the argument need_var to clust_moments_c(). get_p_clust() sets need_var = TRUE only if it uses predvar_mix() and uses_moments is not FALSE. With need_var = FALSE, clust_moments_c() skips the second loop and the dis^2 sum, and returns NA for the two moments.

The table gives the time of one get_p_clust() call with N = 407, 4000 draws and 400 clusters. Each value is the median of 30 calls; a range shows the values of 3 rounds in one Triton job. The timer has 1 ms resolution.

family	master 4e01313 (ms)	this PR without the skip (ms)	this PR with the skip (ms)
Gaussian	53	8	8-9
binomial	36	8	5-6
Poisson	46	8	5-5.5

For binomial and Poisson, the skip saves 2-3 ms of 8 ms per call. cv_varsel() with validate_search = TRUE and refit_prj = TRUE calls get_p_clust() twice per left-out observation. I did not measure the gain in a full cv_varsel() run. With the skip, the results of get_p_clust() are identical to those without the skip for all three families.

New tests check that clust_moments_c() with need_var = FALSE gives the same means as with need_var = TRUE. They also check that get_p_clust() sets need_var = TRUE only for Gaussian, Student-t, the latent projection and a predvar_mix() without uses_moments.