Implement group level mixtures

[GidonFrischkorn]

Closes #1659.

The changes in this PR implement the feature requested in #1659: mixture() can compute the mixture over the levels of a grouping variable instead of over individual observations, so a whole group (e.g. a participant) is assigned to one component. The motivation, the likelihood distinction, and the Stan sketch are available in the issue and its linked note.

In two places I departed from the proposals from the issue:

• The grouping is specified in the call to mixture(), not in a mix() aterm in the brmsformula. The issue proposed y | mix(gr = "ID") ~ .... That works, but a mix addition term collides with the existing mi term through R's $ partial matching (adforms$mi resolves to a mix element when no mi is present), which would force $mi → [["mi"]] edits across ~8 unrelated files. Putting gr on the family avoids that and keeps the change scoped. The only thing the family route doesn't get for free is the grouping column reaching the model frame (built from formula variables), so it is added to allvars explicitly in R/brmsterms.R. Feel free to say if you would prefer the grouping specification as an aterm in the brmsformula instead.
• I used log_sum_exp over a ps[] array, instead of the log_mix proposed in the issue. This follows other brms code that already builds every mixture as log_sum_exp over per-component terms rather than log_mix, so the group-level path reuses that and generalizes to N components for free.

Background on some coding choices

During implementation I noted some things that I want to shortly elaborate on. If you notice errors in my thinking, please say so, and I am happy to adapt the PR accordingly.

• The normalized _lpdf in the accumulator is required. Normalization constants cancel out of an observation-level log_sum_exp but are summed over the group before mixing here, so dropping them (_lupdf) would change the relative component weights. The grouped path therefore never uses the unnormalized form even when normalize = FALSE.
• Predicting theta is restricted to group-constant predictors. A per-observation-varying mixing proportion is not meaningful when the component is shared across a group. This is validated in standata; when constant, the proportion is read from a representative observation per group (Jmixrep).
• log_lik/loo need to change unit for grouped mixtures. The likelihood for grouped mixtures doesn't factorize over observations, so log_lik returns one column per group and loo/waic become leave-one-group-out — the correct CV unit here. Base loo/waic/loo_compare/r_eff are column-count agnostic and handle it. The changes in the PR follow the same unit: kfold forms folds over whole groups, reloo and loo_moment_match refit/match per group, and loo_subsample works because the pointwise log_lik is already per group. kfold_predict works on saved grouped kfolds; predictions stay per observation while the elpd unit is the group.
• Bundled upstream bugfix (if you prefer I can add this in a separate PR). The Pointwise log_lik (pointwise = TRUE) was broken for all multivariate models without rescor, independent of this feature: log_lik_pointwise passed each response's prepared draws as data_i instead of draws (lapply(draws$resps, log_lik_pointwise, i = i)). Only loo_subsample exercises this path, which is presumably why it went unnoticed until now.
• With its multivariate machinery, brms did already support independent per-response grouped mixtures. So, in the supported multivariate case, each response carries its own mixing proportion and its own per-group component assignment; responses can be coupled only through shared/correlated group-level effects (e.g. (1 | p | ID) across formulas). Thus, a group can be in component 1 for one response and component 2 for another: the joint per-group likelihood is the product of per-response mixtures, not a mixture of products. A single latent class driving all responses (one shared assignment) would require cross-response accumulation into one Lmix before a single log_sum_exp, cutting against the one-likelihood-block-per-response architecture; I felt that this is out of scope here. I added information to the mixture() docs, so they state the independent-assignment semantics explicitly.
• To achieve well-behaved posterior predictives, multivariate grouped mixtures are allowed if all responses are
  grouped mixtures over the same grouping variable with rescor = FALSE (rescor is independently excluded for mixtures by allow_rescor). Then the group remains the pointwise unit jointly: its log-likelihood is the sum of the per-response per-group terms, and every CV method above carries over unchanged. Asymmetric models (grouped mixture + ordinary response) and mismatched gr variables are refused with errors that say why: they leave log_lik/loo without a common pointwise unit.

Evaluation of the implementation

To test the behavior of the implementation, I reproduced Musfeld, Souza & Oberauer (2023, PNAS 120:e2218042120), a per-participant Hebb-repetition mixture (learner vs. non-learner):

• Univariate accuracy model: the generated Stan matches their per-subject log_mix accumulation, including a shared baseline across components (shared non-linear parameters) and a non-linear hinged-onset learner curve. Fitting all three visual conditions on the authors' OSF data recovers the learner proportion at 0.58–0.59 vs. the published ~0.60, near-identical between groups (their central between-group null), with learning curves matching their Fig 3D.
• Multivariate joint model: their combined learning + awareness Stan model is two independent per-subject mixtures coupled by five 2×2 participant-effect correlations. A prototype fit on their data reproduces the key published result, the learning↔awareness onset correlation: 0.83 [0.63, 0.95] vs. published 0.82 [0.63, 0.94], with awareness onset preceding learning onset as in the paper. (caveats: their ordered-beta awareness family was approximated with Beta plus endpoint compression; zero_one_inflated_beta is currently not allowed inside mixture().)

[paul-buerkner]

Thank you very much for this PR! I will review it when I find the time, hopefully during August.

Just a quick note: please disclose the use of AI within contributions to brms. Based on the way comments and details are written, I can see AI at play here and I would like to keep a bit track of it where it enters brms. Nothing bad about it obviously (perhaps apart from the very wordy style of commenting), more like a matter of principle here.

[GidonFrischkorn]

Hi Paul,

sure, sorry for not directly disclosing the AI use. I sketched out the math myself based on the issue report and explored some additional considerations for the posterior predictive and then used Claude Code for implementation and then checked the code against the math for verification. I am also happy to go over the comments to focus them on the relevant points.

If you need more details, feel free to say so.