Add RSSSM: regime-switching sequential sampling models - #1162
Conversation
…d make it idempotent
…ood wiring in sample
…, plot/vi/graph surface
…, and visualizations
…vector validation, LOO caveat and API docs
…nment, lapse idempotency, degeneracy warning, FFBS brute-force test
…idation, fixed-value bounds checks, registry-first choices and emission resolution, inherited-method guards, pyrefly and ruff format fixes
…ite fixed values rejected, per-regime p_outlier and LAN backend default on the config path, extra_fields rejection, notebook lint
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
|
Warning Review limit reached
Next review available in: 5 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThe PR adds the ChangesRSSSM regime-switching model
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant RSSSM
participant PyMC
participant FFBS
participant DataTree
Client->>RSSSM: construct model with data and configuration
RSSSM->>PyMC: build emissions and scalar forward likelihood
Client->>RSSSM: sample posterior
RSSSM->>PyMC: run sampling
Client->>RSSSM: infer regimes or compute log likelihood
RSSSM->>FFBS: process posterior draws
FFBS->>DataTree: write regime and likelihood outputs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
📝 Generate docstrings
🧪 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.
Actionable comments posted: 6
🧹 Nitpick comments (4)
tests/hmm/test_rsssm_ffbs.py (2)
378-387: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDerive the
vparticipant dimension fromn.Lines 385-387 size
a,z, andtwithn = model.n_participants. Line 383 hardcodes the participant dimension ofvto 2 through the literal[[-1.0, 1.0], [-1.0, 1.0]]. The two agree only becausemake_panel(2, 40)creates two participants. If that fixture changes,vand the other parameters disagree in shape and the failure is opaque.♻️ Proposed fix
"v": ( ("chain", "draw", "p", "k"), - np.tile([[-1.0, 1.0], [-1.0, 1.0]], (1, 4, 1, 1)), + np.tile([-1.0, 1.0], (1, 4, n, 1)), ),🤖 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 `@tests/hmm/test_rsssm_ffbs.py` around lines 378 - 387, Update the `"v"` dataset construction in the test to derive its participant dimension from `n = model.n_participants` instead of hardcoding two participant rows. Preserve the existing values and `(chain, draw, p, k)` dimensions while ensuring `v` remains shape-compatible with `a`, `z`, and `t` if the fixture participant count changes.
140-151: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExercise the production per-trial decomposition.
test_per_trial_delta_sums_to_marginalconstructsdeltafromlogZin the test, so the assertion is a telescoping identity. It does not testcompute_log_likelihood. The production split is inline, and no helper exists to call. Extract the split into a helper and test it, or comparemodel.compute_log_likelihoodoutput with an independently calculated per-trial vector. The existing marginal test checks only the total.🤖 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 `@tests/hmm/test_rsssm_ffbs.py` around lines 140 - 151, Replace the self-derived telescoping assertion in test_per_trial_delta_sums_to_marginal with coverage of the production per-trial decomposition used by compute_log_likelihood. Extract that inline split into a callable helper and test its output, or independently compute the expected per-trial vector and compare it with model.compute_log_likelihood, while retaining the existing total-marginal validation.src/hssm/hmm/specs.py (1)
44-79: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winValidate that Dirichlet concentrations are positive.
StickyDirichletandDirichletConcentrationaccept non-positivealpha. A Dirichlet requires strictly positive concentrations.config.pycallsconcentration(K)eagerly so a mismatched matrix fails early; a zero or negative entry currently survives that check and fails later inside PyMC, or yields a degenerate prior.FixedInitialDistribution(lines 115-122) already validates its values eagerly, so the same treatment here is consistent.Apply the same check in `DirichletInitialDistribution.concentration` (lines 137-148).♻️ Proposed fix
def concentration(self, K: int) -> np.ndarray: """Return the ``(K, K)`` Dirichlet concentration matrix.""" + if self.diag <= 0 or self.offdiag <= 0: + raise ValueError( + f"transition_prior concentrations must be > 0, got " + f"diag={self.diag}, offdiag={self.offdiag}." + ) alpha = np.full((K, K), float(self.offdiag))def concentration(self, K: int) -> np.ndarray: """Return the ``(K, K)`` Dirichlet concentration matrix.""" alpha = np.asarray(self.alpha, dtype=float) + if np.any(alpha <= 0): + raise ValueError( + f"transition_prior alpha must be > 0 everywhere, got {alpha.tolist()}." + ) if alpha.ndim == 0:src/hssm/hmm/likelihoods/builder.py (1)
130-148: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConstrain
poolingto a literal type.
poolingis typed asstr, and any value other than"full"selects the per-participant branch. A typo such as"None"therefore silently builds the no-pooling expansion instead of raising. UseLiteral["full", "none"]in bothmake_hmm_logp_opandbuild_log_emission, or validate the value once at the top ofbuild_log_emission.♻️ Proposed validation
N, T = n_participants, n_trials n_rows = N * T + if pooling not in ("full", "none"): + raise ValueError(f"`pooling` must be 'full' or 'none', got {pooling!r}.")🤖 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/hmm/likelihoods/builder.py` around lines 130 - 148, Constrain the pooling parameter in both make_hmm_logp_op and build_log_emission to Literal["full", "none"], or validate it at the start of build_log_emission and reject any other value. Ensure the existing branching only handles these two recognized pooling modes instead of treating arbitrary strings as "none".
🤖 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.
Inline comments:
In `@design/hssm_hmm.md`:
- Line 900: Escape the union-type pipe characters in the inline-code expressions
in the table rows for switching_params and the corresponding row at line 903,
using the table-safe Markdown escape so each expression remains in a single
column and passes MD056.
- Line 270: Update the fenced module-tree code block in hssm_hmm.md by adding
the text language tag to its opening fence, changing it to a text-designated
fence while leaving the block contents unchanged.
- Line 703: Update the §5.5 and §5.6 documentation to describe infer_regimes and
compute_log_likelihood, including their FFBS helpers, as accepting and returning
xarray.DataTree rather than az.InferenceData. Revise the related post-fit
container and return descriptions to match the DataTree structure while
preserving the existing likelihood contents and dimensions.
In `@src/hssm/hmm/config.py`:
- Around line 130-142: Update the per-regime detection in the p_outlier
validation block to use the same one-dimensional fixed-vector predicate as the
checks around the fixed-vector handling, requiring ndarray values to have ndim
== 1. Reuse a shared predicate across all three p_outlier checks so 0-D arrays
are treated as shared scalars and rejected by the decision-10.1.9 guard.
In `@src/hssm/hmm/specs.py`:
- Around line 295-311: Update resolve_ordering to validate dictionary keys
against the supported OrderByParam fields before calling OrderByParam(**spec),
and raise the established clear configuration error for any unexpected key
instead of allowing the constructor’s raw TypeError. Preserve handling of valid
ordering dictionaries and all other input forms.
In `@tests/hmm/test_rsssm_ffbs.py`:
- Line 211: Update the layout comment next to the `ll` assignment to describe
the actual dimensions as `(chain, draw, __obs__)`, matching the flattened
observation group validated by the later assertion.
---
Nitpick comments:
In `@src/hssm/hmm/likelihoods/builder.py`:
- Around line 130-148: Constrain the pooling parameter in both make_hmm_logp_op
and build_log_emission to Literal["full", "none"], or validate it at the start
of build_log_emission and reject any other value. Ensure the existing branching
only handles these two recognized pooling modes instead of treating arbitrary
strings as "none".
In `@tests/hmm/test_rsssm_ffbs.py`:
- Around line 378-387: Update the `"v"` dataset construction in the test to
derive its participant dimension from `n = model.n_participants` instead of
hardcoding two participant rows. Preserve the existing values and `(chain, draw,
p, k)` dimensions while ensuring `v` remains shape-compatible with `a`, `z`, and
`t` if the fixture participant count changes.
- Around line 140-151: Replace the self-derived telescoping assertion in
test_per_trial_delta_sums_to_marginal with coverage of the production per-trial
decomposition used by compute_log_likelihood. Extract that inline split into a
callable helper and test its output, or independently compute the expected
per-trial vector and compare it with model.compute_log_likelihood, while
retaining the existing total-marginal validation.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 549c3a6a-1847-4201-8d51-b51a2e27b332
📒 Files selected for processing (28)
.github/workflows/check_notebooks.ymldesign/hssm_hmm.mddesign/hssm_hmm_overview.mddocs/api/rsssm.mddocs/changelog.mddocs/tutorials/rsssm_k3.ipynbdocs/tutorials/rsssm_regime_switching.ipynbmkdocs.ymlsrc/hssm/__init__.pysrc/hssm/base.pysrc/hssm/hmm/__init__.pysrc/hssm/hmm/config.pysrc/hssm/hmm/ffbs.pysrc/hssm/hmm/likelihoods/__init__.pysrc/hssm/hmm/likelihoods/builder.pysrc/hssm/hmm/likelihoods/emissions.pysrc/hssm/hmm/likelihoods/forward.pysrc/hssm/hmm/ordering.pysrc/hssm/hmm/rsssm.pysrc/hssm/hmm/specs.pysrc/hssm/hmm/utils.pytests/hmm/__init__.pytests/hmm/conftest.pytests/hmm/test_rsssm.pytests/hmm/test_rsssm_ffbs.pytests/hmm/test_rsssm_lan.pytests/hmm/test_rsssm_sampling.pytests/test_hssm.py
…y and Dirichlet-positivity validation, pooling literal, per-trial delta helper, design-doc DataTree alignment
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@src/hssm/hmm/config.py`:
- Around line 174-175: Update the per-regime detection logic around
_is_fixed_vector and has_per_regime so fixed vectors count as
regime-distinguishing only when their values are not all identical. Preserve
switching_params behavior, and add a regression test covering identical fixed
vectors such as {"v": [0.5, 0.5]} that verifies the documented identifiability
warning is emitted.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 42fd1bd5-ab33-4045-8f8c-d4a0a41c81a8
📒 Files selected for processing (8)
design/hssm_hmm.mdsrc/hssm/hmm/config.pysrc/hssm/hmm/ffbs.pysrc/hssm/hmm/likelihoods/builder.pysrc/hssm/hmm/rsssm.pysrc/hssm/hmm/specs.pytests/hmm/test_rsssm.pytests/hmm/test_rsssm_ffbs.py
🚧 Files skipped from review as they are similar to previous changes (6)
- src/hssm/hmm/specs.py
- src/hssm/hmm/ffbs.py
- src/hssm/hmm/likelihoods/builder.py
- src/hssm/hmm/rsssm.py
- tests/hmm/test_rsssm_ffbs.py
- design/hssm_hmm.md
| has_per_regime = bool(self.switching_params) or any( | ||
| _is_fixed_vector(spec) for spec in self.param_specs.values() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Warn for identical fixed per-regime values.
Line 175 treats every length-K fixed vector as regime-distinguishing. If switching_params is empty and param_specs={"v": [0.5, 0.5]}, no emission parameter differs across regimes. The model is still unidentifiable, but it skips the documented warning.
Treat a fixed vector as distinguishing only when at least two regime values differ. Add a warning regression test for identical fixed vectors.
🤖 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/hmm/config.py` around lines 174 - 175, Update the per-regime
detection logic around _is_fixed_vector and has_per_regime so fixed vectors
count as regime-distinguishing only when their values are not all identical.
Preserve switching_params behavior, and add a regression test covering identical
fixed vectors such as {"v": [0.5, 0.5]} that verifies the documented
identifiability warning is emitted.
…istics result to Dataset in model_cartoon
fmuia
left a comment
fmuia
left a comment
There was a problem hiding this comment.
Addressed all six actionable comments and the four nitpicks in f613688:
- Per-trial split extracted into
ffbs._per_trial_delta(now used bycompute_log_likelihood) and tested against brute-force prefix marginals instead of the self-derived telescoping identity. vparticipant dimension in the no-pooling FFBS test is derived frommodel.n_participants.StickyDirichlet,DirichletConcentration, andDirichletInitialDistributionreject non-positive concentrations eagerly (tests added).poolingis constrained toLiteral["full", "none"]inmake_hmm_logp_op/build_log_emission, with a runtime guard.
1 participant
RSSSMclass (src/hssm/hmm/): regime-switching sequential sampling models with HMM dynamics over latent regimes, analytical and LAN emission backends, per-regimep_outlierlapse mixture, and post-hoc FFBS regime recovery.tests/hmm/), tutorials (K=2 and K=3), API reference, and changelog entry.Closes #957
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation
Bug Fixes