[digicosmos86]

This PR brings HSSM up to the Bambi rewrite (bambinos/bambi#1002). It is large because Bambi changed a lot underneath us, so the notes below explain the most important changes and why they were made.

Tracking issue: #1305.

Summary

Table of contents

• Summary
• Breaking changes
  • Custom links take a single inverse_link
  • Dimension names follow Bambi 0.20
  • Traces must match the model they predict from
• Changes to existing functionality
  • Centered intercepts and initial values
  • Likelihood parameters and out-of-sample log-likelihood
  • Out-of-sample posterior predictive
  • Predictions for new or unidentified participants
  • Trace contents
  • Errors for unknown formula variables
• Fixes
• Dependency, packaging, and CI changes
• Significant internal changes
• Tests
• Documentation
• Merge from main
• Compatibility and release plan

Bambi 0.20 separates its frontend, which describes the model, from its backend, which builds and evaluates the PyMC graph. Prediction, log-likelihood computation, and parameter transforms now run on the PyMC/PyTensor graph instead of being rebuilt with NumPy/xarray. Bambi also renamed its model vocabulary (components → parameters) and reorganized its modules.

HSSM relied on several internals that were removed or moved: Model._compute_likelihood_params, Family._make_dist_kwargs_and_coords, response_evaluate_new_data, sample_new_groups, and the split linkinv/linkinv_backend links, among others. This PR replaces each of them with the graph-based paths Bambi now provides. Where Bambi's new behaviour would have changed what HSSM users see, the PR keeps HSSM's public results stable. Where Bambi's new behaviour is more correct, the PR adopts it and documents the change.

The work began by marking every test failure under Bambi dev as xfail and grouping them by root cause (design/bambi-dev-upgrade-test-failures.md, design/bambi-migration-fix-plan.md). Each group was then fixed under its own sub-issue and its xfail marks removed. No migration xfail marks remain.

Breaking changes

Custom links take a single inverse_link

Bambi no longer keeps one inverse link for NumPy computations and another for the PyTensor graph (#1311). hssm.Link follows suit:

# before
hssm.Link("custom_log", link=np.log, linkinv=np.exp, linkinv_backend=pt.exp)
# after
hssm.Link("custom_log", link=np.log, inverse_link=np.exp)

inverse_link must work on PyTensor tensors. NumPy ufuncs such as np.exp already do, because they dispatch to symbolic operations. Reading link.linkinv or link.linkinv_backend now raises AttributeError. The forward link is optional. The built-in bounded gen_logit exposes the same transformation through the new attribute, so its callers are unaffected.

Dimension names follow Bambi 0.20

• Group-specific terms use {factor}_dim instead of {factor}__factor_dim.
• A categorical predictor in a formula with an intercept uses {var}_dim_reduced.
• A nested factor such as (1|g:h) now gets two dims, g_dim and h_dim, instead of one.
• The 2-D response dim is rt,response_dim in every group ([Migration] Regenerate the recorded trace fixtures that predate bambi 0.20 (*_Intercept_centered) #1336). Bambi names it rt_dim, and observed_data previously used rt,response_extra_dim_0. SSMFamily declares RESPONSE_NDIM = 2 so Bambi gives the response its correct dims.

Code that selects coordinates by dim name must be updated.

Traces must match the model they predict from

sample_posterior_predictive now raises a ValueError when the posterior is missing a free random variable of the model (#1336). Before, an out-of-sample prediction silently drew the missing variable from its prior. The common case is a trace saved before Bambi 0.20, which has no v_Intercept_centered. Missing *_offset variables are allowed when their base term is present. Old traces need to be resampled.

Changes to existing functionality

Centered intercepts and initial values

When a parameter has an intercept, at least one common predictor, and center_predictors=True (the default), Bambi now samples {param}_Intercept_centered and records the original-scale {param}_Intercept as a deterministic (#1375). az.summary and plot_trace therefore show two rows for such an intercept. Both are kept, because Bambi's predict and compute_log_likelihood need the centered one. The intercept prior means the same thing as before. To report only the usual intercept, use:

az.summary(idata, var_names=["~_centered"], filter_vars="like")

Initial values follow the free variable:

• HSSM's INITVAL_SETTINGS defaults and any prior initval keyed on {param}_Intercept are applied to {param}_Intercept_centered ([Migration] Let truncated-prior callables accept dims #1312).
• sample(initvals={"v_Intercept": ...}) is moved to v_Intercept_centered, because PyMC silently ignores an initial value on a deterministic ([Migration] Document the *_Intercept_centered / *_Intercept pair and predictor centering #1375). An explicit v_Intercept_centered key takes precedence.

The new explanation page Predictor centering and the intercept covers this in detail.

Likelihood parameters and out-of-sample log-likelihood

HSSM now evaluates the trial-wise likelihood parameters with Model.predict(kind="response_params") instead of the removed Bambi internals (#1313). In sample, they are still added to posterior, and keep_likelihood_params / add_likelihood_parameters_to_datatree behave as before. Out of sample, Bambi writes them to a predictions group. log_likelihood(data=...) uses that group for its computation and then removes it, so the only change to the trace is the new log_likelihood group. It no longer overwrites posterior with out-of-sample parameters.

Rebuilding the response term on new data now goes through formulae.design_matrices, which replaces Bambi's removed response_evaluate_new_data. For that reason formulae>=0.7.0 is now a direct dependency.

Out-of-sample posterior predictive

Bambi now stores out-of-sample predictions in predictions and predictions_constant_data instead of posterior_predictive (#1315). HSSM moves the response draws back into posterior_predictive, so in-sample and out-of-sample results keep one layout. This also keeps the plotting helpers working, since they always predict on the data frame they are given. Stale prediction groups are cleared before a new response prediction. kind="response_params" still returns Bambi's result unchanged.

Predictions for new or unidentified participants

Bambi retired sample_new_groups (HSSM always passed False) and now decides how to handle each observation from its grouping value (#1316):

• A missing value (None, np.nan, pd.NA) means the participant's identity is unknown. The observation borrows the coefficients of a randomly chosen fitted group, which matches the old behaviour.
• A non-missing value not seen during fitting is a new group. Its coefficients are drawn from the population-level model using the fitted hyperparameters, and all of its observations share them. Previously an unseen level was handled like a missing one.

This affects sample_posterior_predictive(data=...) and log_likelihood(data=...). Because those coefficients are random, log_likelihood now takes a random_seed to make the result reproducible (#1376). The new how-to guide Predict for new or unidentified participants explains the semantics and shows how to average draws into a trial-level or participant-level predictive score.

Trace contents

• Traces gain a constant_data group with the design-matrix data (for example v_x_data) and the group index arrays (for example g__idx) Bambi used to build the model.
• Bambi 0.20 records fixed parameters (for example the default p_outlier, or a fixed z) as constant deterministics. These are now dropped from posterior ([Migration] Drop the constant parameters and *_Intercept_centered RVs bambi 0.20 keeps in posterior #1330) and from prior ([Migration] Stop pymc's forward sampler rejecting the aDDM p_outlier TensorConstant #1329). In prior they broke forward sampling, such as the prior-predictive mean in plot_model_cartoon.
• Posterior cleanup now drops only the parent parameter's {parent}_mean alias (e.g. v_mean). Before, it dropped every variable whose name contained _mean, so a covariate such as rt_mean lost its coefficient v_rt_mean ([Migration] Fix the _mean substring filter in posterior cleanup and drop dead helpers #1378).

Errors for unknown formula variables

A formula that names a column missing from data now raises a ValueError that names the parameter and the formula (#1317). Before, users saw a bare KeyError/IndexError from formulae/pandas, and its type had changed across releases.

Fixes

• vi(backend="jax") works again on every model, including missing-data models ([Migration] Make VI on the JAX compile backend trace bambi 0.20's symbolic __obs__ alloc #1328). Bambi 0.20 keeps the response and the __obs__ observation count as shared variables and derives shapes from them, which the JAX linker rejects as dynamic. HSSM now passes pm.fit a more_replacements that pins these to constants, the same fix PyMC applies in its own JAX samplers. A user-supplied more_replacements still takes precedence entry by entry, and an explicit None is accepted.
• Truncated-prior callables accept dims and shape ([Migration] Let truncated-prior callables accept dims #1312). Bambi now instantiates callable priors as dist(name, dims=...) when it builds the centered intercept.
• plot_model_cartoon casts traces to the model's dtypes before computing deterministics ([Migration] Regenerate the recorded trace fixtures that predate bambi 0.20 (*_Intercept_centered) #1336).

Dependency, packaging, and CI changes

• bambi tracks the dev branch through [tool.uv.sources] until 0.20 is released. The floor must be pinned to the published bambi>=0.20 before this ships (see the release plan below).
• formulae>=0.7.0 is now a direct dependency (see above).
• The pytest filterwarnings setting turns Bambi's "X is deprecated. Use Y instead." FutureWarnings into errors, so any deprecated accessor still in use fails CI ([Migration] Migrate deprecated bambi accessors and refresh test doubles #1314). All of HSSM's uses have been migrated: components → parameters, distributional_components → conditional_parameters, response_component.term → response_term.
• pyrefly's floor and pre-commit hook are bumped to 1.3.0, and the new errors it reported are fixed.
• The version is bumped to 0.6.0.

Significant internal changes

• base.py no longer rebuilds likelihood parameters itself. It calls Model.predict and reorganizes the output groups (predictions → posterior_predictive, removing temporary groups).
• utils.log_likelihood rebuilds the response with formulae.design_matrices and evaluates the parameters returned by _compute_likelihood_params.
• _vi_compat.py holds the shared-variable freeze for JAX VI.
• link.py keeps a single inverse_link path.
• _clean_posterior_group drops every parameter Bambi built as a pm.Deterministic except *_Intercept_centered.
• _remap_centered_intercept_initvals handles user initvals passed to sample().

Tests

• The trace fixtures (cavanagh_idata.nc, idata_cavanagh_cartoon.nc, test_idata_race.nc) are regenerated for Bambi 0.20 by the new tests/trace_fixtures.py. The stale idata_race_cartoon.nc is removed.
• New tests/test_new_group_predictions.py covers the unknown-identity and new-group prediction paths, including random_seed.
• New tests cover centered-intercept initvals, the remap of sample(initvals=...), constant parameters dropped from prior/posterior, the _mean cleanup, the trace guard in sample_posterior_predictive, JAX VI on missing-data models, and the ValueError for unknown formula variables.
• A test asserts that HSSM's out-of-sample path raises no FutureWarning.
• Every migration xfail mark has been removed.

Documentation

• New: Predictor centering and the intercept and Predict for new or unidentified participants.
• The link_functions tutorial is migrated to inverse_link, and the initial-values tutorial now uses a_Intercept_centered.
• Notebook outputs, archived notebooks, and the cached scientific_workflow_hssm traces are regenerated with Bambi 0.20 dim names ([Migration] Refresh notebooks that show pre-0.20 dim and variable names #1377). Machine-local paths are stripped from outputs.
• random_slope_safe_priors now shows HSSM rejecting the legacy free-mean priors under noncentered=True. tutorial_p_outlier_regression models v as population terms plus participant deviations, which satisfies the safe-prior guard on bounded v ([Migration] Rework random_slope_safe_priors and tutorial_p_outlier_regression for the new prior validation #1382).
• The changelog has a 0.6.0 (unreleased) section covering all of the above.

Merge from main

main is merged into this branch. The conflicts were resolved as follows:

• base.py initvals: both sides are kept. The branch's _centered-aware lookup (settings_key) is combined with main's clamping of default initvals into declared bounds (Clamp default initial values into the model's declared bounds #1293). A default for *_Intercept_centered is therefore looked up under *_Intercept and clamped into the parameter's bounds.
• utils.py: the branch is kept. Main's pyrefly fix (Lift the pyrefly <1.3 cap: 1.3.0 reports 5 new type errors #1334) applied to code the rewrite replaced.
• pyproject.toml: the branch's 0.6.0 and pyrefly>=1.3.0 are kept. Main capped pyrefly at <1.3, but this branch already fixes those errors. The pyrefly and mypy hooks pass on the merged tree.
• Tests: both sides' new tests are kept. The invalid-formula case expects the branch's ValueError rather than main's KeyError (see [Migration] Reconcile two drifted test assertions (formula error type, noncentered message) #1317). Two of main's new tests were adapted to Bambi 0.20:
  • test_identity_link_override_clamps_into_bounds and test_explicit_log_link_keeps_default_start read the default from {param}_Intercept_centered, since t ~ 1 + stim now centers the intercept. The clamped value is unchanged (0.25 + 0.05 * 1.75).
  • test_jax_callable_non_trialwise_params_logp[jax_linker] freezes the model with freeze_dims_and_data before compiling under the JAX linker, the same way PyMC's JAX samplers do and for the same reason as [Migration] Make VI on the JAX compile backend trace bambi 0.20's symbolic __obs__ alloc #1328. The test still guards Fixing a parameter with backend='jax' fails with a bare AssertionError in specifyshape #1092.
• Tutorials (random_slope_safe_priors, tutorial_p_outlier_regression): the branch is kept. Both sides fixed the same failures (Two tutorials fail at construction since the 0.5.0 safe-prior guards (p_outlier regression, random-slope safe priors) #1363 on main, [Migration] Rework random_slope_safe_priors and tutorial_p_outlier_regression for the new prior validation #1382 here), and only the branch's versions were executed against Bambi 0.20.
• Changelog: 0.6.0 (unreleased) is placed above main's 0.5.1.

Compatibility and release plan

• Do not release from this branch until Bambi 0.20 is on PyPI. It currently installs Bambi from the dev branch via [tool.uv.sources], which pip install hssm ignores. Before 0.6.0 is tagged, remove the git source and pin bambi>=0.20.
• Users should expect to update code that relies on custom Link objects (linkinv → inverse_link), selects coordinates by dim name, or reads *_Intercept as a free variable. Traces saved before 0.6.0 must be resampled before they can be used for out-of-sample prediction.
• Rolling back means reverting this merge and staying on Bambi 0.19. No user data or on-disk formats are migrated, but traces fit with 0.6.0 carry the new dim names and *_Intercept_centered variables, which 0.5.x will not interpret identically.

Summary by CodeRabbit

• New Features
  • Updated HSSM for Bambi 0.20, including improved prediction and likelihood handling for fitted, missing, and unseen groups.
  • Posterior predictive sampling now defaults to response predictions and supports reproducible predictions for new groups.
  • Added guidance on predictions for new participants and predictor centering.
• Bug Fixes
  • Improved error messages for missing formula columns and validation of posterior traces.
  • Improved compatibility of JAX variational inference with models whose data or observation dimensions change.
• Documentation
  • Updated custom-link examples to use the current inverse_link interface and added 0.6.0 migration notes.