Add RSSSM: regime-switching sequential sampling models

[fmuia]

Summary

Adds hssm.RSSSM, a new top-level class for fitting regime-switching sequential sampling models (issue #957). Each trial belongs to one of K hidden cognitive regimes that evolve as a Markov chain; within a regime the (rt, response) emission is a standard SSM (e.g. DDM) with regime-specific switching parameters. The discrete regimes are marginalised out by the forward algorithm (one batched pytensor.scan, contributed as a scalar pm.Potential), leaving only continuous parameters for NUTS — the same approach as the hand-built docs/tutorials/hmm_ddm_regime_switching.ipynb, but as a high-level HSSM(...)/RLSSM(...)-style class.

Implements the design on the issue_957/plan_hssm_hmm_class branch (docs/design/hssm_hmm.md), Phases 2–5.

What's included

• Core class (src/hssm/hmm/): subpackage with rsssm.py, config.py, specs.py, ordering.py, utils.py, and the layered likelihood (likelihoods/{forward,emissions,builder}.py). Builds the PyMC model directly (bambi deferred, decision 10.1.8).
• Generic in K and in N (participants); unbalanced panels via end-padding + emission mask.
• Both emission backends: analytical and LAN (approx_differentiable, jax & pytensor).
• Label-switching: ordered-transform anchor — AutoOrdering (default), OrderByParam, NoOrdering.
• Three-mode parameter rule: scalar = shared, length-K list = fixed-per-regime, prior/switching_params = inferred-per-regime.
• Pooling: full (default) and none (pooling="none"); P/pi0 stay global. Estimable global pi0.
• Per-regime p_outlier lapse mixture ((1 - p_k)·SSM_k + p_k·lapse).
• Post-hoc helpers: infer_regimes (FFBS), compute_log_likelihood (per-trial logp for arviz.loo/waic), plot_regime_recovery.
• Docs: two executed tutorials (single- and multi-participant), API reference page, changelog entry.

Deliberate deviations from the design doc (please review)

Two change documented behavior and warrant explicit sign-off:

1. Forward algorithm is the scaled (normalised) recursion, not the textbook un-normalised one in §5.2. The un-normalised version produces NaN gradients past ~400 trials per participant (both PyTensor and JAX), while the value stays finite — masked entirely by numpyro's init jitter. The scaled recursion is value-identical (matched to machine precision) and gradient-stable.
2. log_likelihood group is a flat (chain, draw, __obs__) over real trials only, not (chain, draw, n_participants, n_trials) (§5.6). The rectangular shape fed padded trials to arviz.loo as logp = 0 ("perfectly predicted") observations, biasing loo/waic on unbalanced panels.

Minor/cosmetic: vi() raises NotImplementedError (deferred, not implemented as §6.3 suggested); bambi-specific inherited kwargs (link_settings, process_initvals, …) are not accepted; the tutorial regression is asserted as deterministic likelihood-level equivalence rather than draw-for-draw posterior bit-for-bit (the two parameterisations can't match draw-for-draw).

Testing

• 76 fast tests + a slow recovery suite (K=2/K=3/angle-LAN, per-regime p_outlier, end-to-end FFBS + loo). 88% line coverage on hmm/.
• Forward correctness vs brute-force enumeration, an independent numpy forward, and long-sequence gradient finiteness (T=400/500/1000, both backends).
• FFBS vs exact forward-backward smoothing; sum(delta) == marginal == model Potential; unbalanced loo counts real trials only.
• A 26-config build/gradient sweep; per-regime p_outlier recovery; LAN backends agree to 1e-15.
• Edge cases: unoccupied/imbalanced regime (no NaN), descending anchor, LAN-backend FFBS, no-pooling fixed params, plot_regime_recovery, vi/log_likelihood/graph.

Known untested (intentional): the full (K, n_trials, n_participants) recovery grid (representative points tested instead); end-to-end sampling on an unbalanced panel (the exact masked-marginal property is tested deterministically).

Scope

v1 only. Deferred to later PRs (architecture leaves hooks): partial/hierarchical pooling, per-participant pi0, covariate-driven P, per-regime regressions, cross-emission/semi-Markov, the posterior-predictive family.

Closes #957.

🤖 Generated with Claude Code

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 80dadd6f-28ed-4293-8ea5-3ec5e2a0b1ba

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch implement_hmm_class

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[fmuia]

Superseded by #1162 — identical content on a clean history (this branch's history could not be rewritten due to repository rules blocking force-pushes).