Fix HSGP prior handling: hsgp() terms are usable through HSSM - #1121
Conversation
make_safe_priors blanketed every common design-matrix term with a scalar default bmb.Prior — hsgp() terms included, which bambi rejects at build time since an HSGP prior must be None or a dict of covariance-function priors. This made HSGP terms unusable through HSSM even with no user prior at all. Skip the injection for terms bambi itself identifies as HSGP (bambi.utils.is_hsgp_term); bambi's automatic HSGP priors take over. The term stays in self.terms for parameterization bookkeeping, and an HSGP term can never be an intercept, so nothing else in the loop is affected. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
_make_prior_dict collapsed every dict-valued prior into a single bmb.Prior(**spec). For hsgp() terms that either shipped bambi an object its HSGPTerm prior setter rejects, or crashed earlier on the bambi-native covariance dict for lack of a "name" key. Keys addressing hsgp() formula terms now pass their dict through: inner HSSM-style specs (with a "name") still convert to bmb.Prior, a top-level "name": "HSGP" marker is dropped as convention residue, and everything else reaches bambi untouched. Detection is by key prefix because this runs before design matrices exist; "hsgp" is the reserved stateful-transform name in bambi's formula namespace, so the prefix is the semantics. Tests cover the gh-624 reproduction, the bambi-documented dict shape, no-prior (automatic priors), JSON-able inner specs, the by= variant, two HSGP terms in one formula, mixed formulas keeping safe defaults on linear terms, and both prior_settings modes — the None path skips make_safe_priors entirely, so the pass-through carries the fix alone there. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- tests/slow: end-to-end fit of an hsgp() drift regression (25 draws), asserting the full HSGP posterior is present — covariance parameters, basis weights, and the per-observation GP contribution. - tests/test_hsgp.py: a dict-valued HSGP prior must fall through the parameterization checks' isinstance guards — the group term still gets its *_offset RV and no spurious warning references the HSGP entry. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Simulates a DDM whose drift varies smoothly with a stimulus variable (v = A*sin(stim), trial-wise theta), specifies the hsgp() covariance prior dict, fits, and recovers the curve against ground truth. Wired into the docs nav; outputs committed and execution skipped in both the docs build (mkdocs execute_ignore) and notebook CI (SKIP_NOTEBOOKS), matching the other sampling tutorials — notebook execution is headed out of GitHub CI into a dedicated job. End-to-end regression coverage for gh-624 lives in tests/slow/test_hsgp_mcmc.py instead. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughHSSM now preserves HSGP covariance-prior dictionaries for Bambi, skips scalar defaults for HSGP terms, adds regression tests, and documents the workflow in a new tutorial. ChangesHSGP regression support
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant User
participant HSSM
participant Bambi
participant Sampler
participant Notebook
User->>Notebook: define HSGP regression and priors
Notebook->>HSSM: construct model
HSSM->>Bambi: pass normalized HSGP prior dictionary
Bambi->>Sampler: build and sample regression model
Sampler-->>Notebook: return HSGP posterior variables
Notebook->>Notebook: extract and plot posterior smooth effect
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/hssm/param/regression_param.py (1)
413-444: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFix the malformed nested type hint in
_make_prior_dict.The parameter annotation at Line 414 declares
dict[dict, Any]for the nested-dict branch of the union. The key type should bestr, matching the actual key type of a covariance-parameter prior dict (e.g."sigma","ell").dictis not a valid semantic key type here, and using it obscures the true shape of the accepted input.🔧 Proposed fix
def _make_prior_dict( - prior: dict[str, float | dict[dict, Any] | bmb.Prior], + prior: dict[str, float | dict[str, Any] | bmb.Prior], ) -> dict[str, float | bmb.Prior | dict[str, Any]]:The dispatch logic below (HSGP dict pass-through vs. recursive conversion vs. pass-through scalar) is otherwise correct and matches the new tests.
As per coding guidelines, "Use pyrefly as the primary type checker and maintain mypy compatibility."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/hssm/param/regression_param.py` around lines 413 - 444, Update the parameter annotation of _make_prior_dict so its nested dictionary branch uses string keys, changing dict[dict, Any] to the appropriate dict[str, Any] shape while preserving the existing HSGP dispatch and recursive conversion logic.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@src/hssm/param/regression_param.py`:
- Around line 413-444: Update the parameter annotation of _make_prior_dict so
its nested dictionary branch uses string keys, changing dict[dict, Any] to the
appropriate dict[str, Any] shape while preserving the existing HSGP dispatch and
recursive conversion logic.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: ad9731c3-8280-4fa8-8d4e-57cff8abd7e3
📒 Files selected for processing (8)
.github/workflows/check_notebooks.ymldocs/tutorials/hsgp_regression.ipynbmkdocs.ymlsrc/hssm/hssm.pysrc/hssm/param/regression_param.pytests/slow/test_hsgp_mcmc.pytests/test_hsgp.pytests/unit/param/test_regression_param.py
Pre-existing code, unrelated to the HSGP change: pyrefly 1.2.0 (with pandas 3.0.5 stubs) types `.loc` column-list selection as DataFrame | Series and rejects the declared DataFrame return in _stack_and_rename. CI resolves dependencies fresh on every run because uv.lock is gitignored, so the floating toolchain surfaced this on the first PR after the pyrefly release. A two-column selection is always a DataFrame at runtime; cast accordingly. Verified against the CI-equivalent environment (no lockfile, fresh resolve, Python 3.13, pyrefly 1.2.0): 0 errors. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
CI note: the initial lint failure was not from this PR's changes. Fixed in eddc896 with a one-line Worth a separate discussion: re-tracking 🤖 Generated with Claude Code |
# Conflicts: # src/hssm/plotting/utils.py
1 participant
Closes #624.
Severity was higher than reported: not only did user-specified HSGP priors fail —
hsgp()terms could not be used through HSSM at all, because the safe-prior machinery injected a scalar default that bambi rejects even when no prior was given.What changed
make_safe_priorsno longer injects scalar defaults onto HSGP terms (detected with bambi's ownis_hsgp_term); bambi's automatic HSGP priors apply when none are given_make_prior_dictpasses anhsgp(...)key's dict through as the covariance-prior dictionary bambi requires, instead of collapsing it into a singlebmb.Prior; inner{"name": ...}specs still convert, and a top-level"name": "HSGP"(the shape from HSSM does not properly send priors to Bambi for an HSGP #624's report) is accepted and droppedsrc/hssm/param/regression_param.py(+54 lines); docstring for the prior shapes inhssm.pydocs/tutorials/hsgp_regression.ipynb(nav-wired, outputs committed, execution skipped in docs build and notebook CI like the other sampling tutorials): recovers a smooth drift effectv = A*sin(stim)end to endTests
by=variant, two HSGP terms, mixed formulas keeping safe defaults on linear terms, bothprior_settingsmodes, non-centering inertness, and a slow end-to-end fit asserting the full HSGP posterior (sigma,ell,weights, GP contribution)Verification
uv run pytest -m "not slow"→ 549 passeduv run pytest tests/slow/test_hsgp_mcmc.py -m slow→ passed (19s)uv run prek run --all-files(ruff, ruff-format, pyrefly, mypy) → green per commitbambi @ address-815(Raise an error for prior names that match no term bambinos/bambi#1001) withUNUSED_PRIORS="error"→ 549 passed — HSSM has no reliance on bambi's silent prior leniency, and correct HSGP prior keys validate under [pymc6 migration] Fixed model cartoon plots #1001Notes
bambi.utils.is_hsgp_termexists at the current floor (0.19.0)hsgp(...)term name still falls to bambi's silent-drop behavior; Raise an error for prior names that match no term bambinos/bambi#1001 (open) turns that into a warning/error🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation
Bug Fixes