Add RSSSM: regime-switching sequential sampling models#1000
Draft
fmuia wants to merge 16 commits into
Draft
Conversation
…d make it idempotent
…ood wiring in sample
…, plot/vi/graph surface
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
…, and visualizations
…vector validation, LOO caveat and API docs
AlexanderFengler
added a commit
that referenced
this pull request
Jul 13, 2026
Full-suite enumeration (sequential, CI ordering, no --exitfirst) surfaced 7
failures in two leak directions plus one integration gap:
- Inbound (fast-CI blockers): pytest imports every module at collection, so a
later module's import-time `hssm.set_floatX("float32")` flipped JAX x64 off
after tests/addm's import-time `set_jax_precision(True)`, desyncing the
vendored kernel's cached dtype — builder/equiv/subclass bodies then computed
float32 against a float64 contract. New tests/addm/conftest.py autouse
fixture pins float64 per-test (pytensor floatX + JAX x64 + kernel cache).
- Outbound: test_addm_waic_loo's body set_floatX("float64") leaked into later
float32-dependent LAN/ONNX tests (tests/slow/test_mcmc approx_differentiable
failed). The fixture restores the previous precision after each aDDM test.
- tests/rl choice-only smoke: ssm-simulators 0.13 ships the ssms.rl presets
(2AB_RW_InvTempSoftmax et al.) so the old guard un-skipped them, but the
presets are not registered in ssms' model_config, which HSSM's RV path
requires — skip until the RLSSM wiring lands (#1000).
Validated: inbound repro (builder/equiv/subclass after a float32 import) 40
passed; outbound pair (waic_loo then the mcmc param) 2 passed; rl smoke 4
skipped; full tests/addm 68 passed.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
AlexanderFengler
added a commit
that referenced
this pull request
Jul 13, 2026
…on + tutorial (#967) * reformulation of addm-integration plan but as a model under the base HSSM class * Recent modifications to class-based addm integration * v3 edits of ADDM-HSSM integration plan * Commit 1: copied over relevant scripts from the efficient-fpt repository and refactored them to fit within HSSM ecosystem. * Commit 2: add builder and separate attention process file. * Commit 3: created ADDM config * Commit 4 v0: Initial attempt at ADDM subclass. Caveat: Post-sampling calculation of log-likelihood and regression on core parameters not yet working. * Current development summary file * Commit 4 test sampling notebook. Additional changes: switched default truncation number to 6 and added an ADDM simulated dataset for sampling testing. * CB: move aDDM tests into the pytest suite; drop shipped data blobs Move the four commit-numbered dev tests out of the shipped package (src/hssm/addm/commit_tests/) into tests/addm/ with feature names, so pytest (testpaths=["tests"]) collects them and they stop shipping in the wheel. Repoint test_addm_vendored_kernel's vendored-jax path to the repo src/ layout (keeps it jax+numpy-only, no full-hssm import). Remove the 1.5MB copied addm_data_*.npz and the 543KB recovery notebook (unused fixtures that were shipping in src/; a proper marimo recovery tutorial comes later). NOTICE_commit_*.md kept for provenance consolidation in the ownership commit. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * CB: add efpt NumPy reference oracle + FPT equivalence tests Vendor efpt's pure-NumPy backend (single_stage/multi_stage + shared _defaults/quadrature/utils/validation) into tests/addm/oracle/ as a self-contained, independent numerical oracle kept at efpt's original DEFAULT_TRUNC_NUM=100. Add test_addm_equiv_efpt.py asserting HSSM's vendored JAX single-stage FPT kernel matches the NumPy oracle to ~1e-6 under x64 (scalar + vectorized cases), mirroring efpt's own cross-backend equivalence tests. This is the correctness anchor the trial-wise (CC) and TRUNC_NUM (CA) commits build on. Source: efficient-fpt @ d97a451 (MIT, (c) 2025 Sicheng Liu). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * CA: keep JAX TRUNC_NUM=6 (documented, RAM-bounded); add efpt MIT LICENSE + provenance - Keep _defaults.DEFAULT_TRUNC_NUM=6 (efpt upstream is 100) as a deliberate, documented HSSM modification: the JAX kernel unrolls the FPT truncation series into the compiled graph and its reverse-mode gradient (VJP) materializes per-term intermediates, so trunc=100 can exhaust RAM at compile/grad time. The cost is ~6e-4 relative FPT-density error vs the efpt numpy oracle (acceptable for inference). Guarded by a test (DEFAULT_TRUNC_NUM <= 6) and documented in NOTICE.md -- this was the audit's actual concern (the fork was undocumented). - Add likelihoods/jax/LICENSE (efpt MIT, (c) 2025 Sicheng Liu), reproduced verbatim as MIT requires since HSSM's own top-level license differs. - Consolidate provenance into likelihoods/jax/NOTICE.md beside the vendored code (matches the jax/__init__.py reference), with attribution + a relicense TODO. - Remove src/hssm/addm/commit_tests/ (dev NOTICE_commit_*.md; provenance now in NOTICE.md). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * CC: support trial-wise / regression core aDDM params Regression or hierarchical priors on a core param (eta, kappa, a, b, x0) make it vary per trial, so it reaches the log-likelihood Op as an (n_obs,) array. The prior aDDM hard-coded params_is_trialwise=[False, ...], bypassing HSSM's standard trial-wise machinery: a regressed param arrived as (n_obs,) and the scalar batched kernel mis-shaped it against the stage dimension and crashed. Such models built but could not be sampled. Fix, reusing established HSSM patterns (no new machinery): - addm.py: replace the [False, ...] freeze with the canonical [param.is_trialwise for ... if name != "p_outlier"] list (identical to hssm.py / RLSSM), so dist.py broadcasts regressed params to (n_obs,). - builder.py: the logp detects per-trial core params by runtime shape (shape[0] == n_obs) at trace time and vmaps the vendored single-trial kernel compute_addm_logfptd with in_axes = 0-if-trial-wise-else-None (the same trial-wise pattern as the ONNX path). All-scalar models keep the optimized batchscan kernel. The vendored jax/ kernel is untouched (stays re-vendorable). sigma stays a scalar diffusion constant. A custom attention process combined with per-trial core params raises NotImplementedError (documented gap). TDD (x64, RED->GREEN against trusted oracles): - tests/addm/test_addm_trialwise.py: constant regression == scalar path bit-for-bit; genuinely per-trial params == the single-trial kernel evaluated trial by trial (aDDM trials are conditionally independent); mixed scalar/trial-wise in_axes; grad finite; custom-process guard. - tests/addm/test_addm_subclass.py: a hierarchical regression on eta now SAMPLES cleanly (upgrades Andrew's build-only check) and the posterior carries the per-participant structure; the regressed param is flagged is_trialwise. Full parameter recovery from known ground truth is deferred to CF (it needs the ssm-simulators aDDM RV from CE to generate data). Divergence D1 from Andrew's plan (he deferred trial-wise core params as out of scope). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * CD: add non-decision time t (decision-time shift) Adds `t` as a sampled non-decision time (Andrew's Commit 5), applied entirely in the builder so the vendored likelihoods/jax/ kernel stays untouched. t is added LAST to the default list_params [eta, kappa, a, b, x0, t] with default 0.0, so the stock model sits on the exact t=0 identity path. Non-decision time is not merely rt-t: fixation onsets (sacc_array) are recorded from trial start but the likelihood reasons in decision time, so removing the first t seconds shifts rt (rt_eff = rt - t) AND slides every fixation onset back by t with the first anchored at 0. Verified against the kernel internals: it reads sacc only via diff-durations (shift-invariant except stage 0) and a final-stage `rt - sacc[safe_d_idx]` (safe_d_idx = min(d-1, max_d-1) >= 1, invariant), so the only live effect is to shorten the first stage by t; d==1 trials get a plain rt-t. Constraint (t inside the first fixation, rt-t>0) is enforced per-trial with -inf. For gradient safety the kernel is evaluated on a clamped in-support t_safe = clip(t, 0, first_fix_end - eps) while `reject` is computed from the ORIGINAL t and applied via jnp.where(reject, -inf, ll) — so a rejected trial never produces a NaN that could leak through the mask onto valid trials. t composes with CC: scalar or per-trial (regressed), classified/broadcast in the shift; no kernel/in_axes change (t folds into rt_eff/sacc_eff before the vmap). dist.py's ensure_positive_ndt is left untouched (it fires by name for "t" and benignly softens the rt<=t offenders to LOGP_LB at the PyTensor level; it has no sacc access so it cannot be the gate). TDD (x64, RED->GREEN vs independent oracles) in tests/addm/test_addm_ndt.py: - t=0 reproduces Commit 4 bit-for-bit (scalar-batch AND per-trial vmap paths); - t>0 == a per-trial-loop manual shift through both the batched kernel and the vendored single-trial kernel, and differs from a naive rt-only shift; - invalid t (outside first fixation, t>=rt, incl. d==1 and the boundary) -> -inf per trial while valid trials keep their manual-shift value; - grad w.r.t. scalar and per-trial t finite, incl. a mixed valid/rejected batch; - per-trial t composes with a trial-wise eta and with a custom process (from_mu). Existing C2-C4 tests updated for the new 6-param arity (t after x0). Config gains t in list_params/bounds/params_default; addm.py needs no change (Op is built from config.list_params; params_is_trialwise auto-flags a regressed t). The vendored likelihoods/jax/ kernel and DEFAULT_TRUNC_NUM=6 are unchanged. Divergence from Andrew: none — this keeps his C5 design; recovery of a regressed t from known data lands with CF (needs the CE simulator). Full addm suite: 59 pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * CE-2 (HSSM): condition aDDM posterior-predictive on observed fixations The aDDM likelihood already receives the observed fixations (extra_fields); the generative rng_fn did not, so posterior/prior-predictive draws self-sampled new fixations instead of conditioning on the real ones. Close the gap reusing the standard RV machinery (no new RV type): - dist.py make_hssm_rv: give HSSMRV a mutable class attr `_extra_fields = None`; rng_fn forwards it to ssms_rng_fn when set. Default None keeps every existing model's generative path byte-for-byte unchanged (verified: a ddm RV keeps _extra_fields None; 17 shared PPC/simulator tests still pass). - addm.py: `_push_rv_extra_fields` stashes {name: array} of the observed covariates on the RV class (rv_op is an instance, so its class carries the attr the rng_fn classmethod reads). Wired into _make_model_distribution (initial data) and _update_extra_fields (PPC / new-data refresh). Each aDDM builds its own RV class, so this never leaks across models. Requires the CE-1 ssm-simulators build (cssm.addm) installed; the ssm-sim side (simulator() extra_fields + rng_fn covariate broadcast) lands in a sibling commit. TDD (tests/addm/test_addm_ppc.py, skips if cssm.addm absent): - make_hssm_rv('addm') builds a real RV (no 'not supported' warning); - the RV carries the observed fixations; a non-aDDM RV carries none (backward compat); - sim<->likelihood recovery: data simulated by the ported engine (CE-1) is best-fit by its true params under the vendored JAX likelihood — ties the two efpt halves. (Surfaced a real aDDM fact: observed fixations are truncated at rt, since the likelihood needs rt >= sacc[d-1]; real data satisfies this. The "keep saccading past the last observed fixation" refinement for PPC is Andrew's C6 / a follow-up.) - [slow] a fitted aDDM's PPC conditions on the observed fixations and its predictive distribution differs from the self-sampled fallback. Full addm suite: 64 pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * CF+CG: aDDM recovery script, tutorial, docs, cleanup CF (recovery + tutorial + docs): - scripts/addm_parameter_recovery.py: off-CI end-to-end check — simulate an aDDM dataset with the ported ssm-simulators engine, fit hssm.aDDM (vendored JAX likelihood), and verify the posterior recovers ground truth within ~2 sd. Confirmed: all six params recovered (e.g. 250 trials / 60 draws -> every |z| < 2). - docs/tutorials/attentional_ddm.py: marimo tutorial (ecosystem convention) demonstrating a trial-wise regression on eta (eta ~ 1 + x), recovery, and posterior-predictive checks conditioned on the observed fixations. marimo-check clean. Not yet wired into mkdocs nav / executed in CI — it needs the aDDM ssm-simulators build; regen + nav instructions are in the file header. - docs/changelog.md: aDDM entry. CG (cleanup + cross-repo pin): - Delete the addm-andrew-dev/ scratch design notes. - pyproject.toml: document that aDDM inference needs no ssm-simulators (vendored likelihood) while PPC/simulation needs cssm.addm, with a release-time TODO to bump the ssm-simulators floor once the aDDM model is published. hssm.aDDM stays exported — every capability (likelihood, trial-wise, ndt, simulator, PPC) landed. A real aDDM fact surfaced during recovery and captured in both the script and the tutorial: observed fixations are truncated at rt (the likelihood requires rt >= sacc[d-1]); real data satisfies this. The "keep saccading past the last observed fixation" refinement for unbiased PPC is a follow-up (Andrew's C6). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * CH: pointwise log-likelihood for aDDM (WAIC/LOO) Model comparison via az.waic / az.loo works out of the box: idata_kwargs={"log_likelihood": True} yields a finite (chain, draw, n_obs) pointwise log-likelihood. No dedicated draw-dimension kernel was needed — the CC trial-wise builder picks its vmap in_axes from the runtime parameter shapes, so it handles PyMC's post-hoc re-evaluation (which adds a draw dimension to the params) for free. This is the capability Andrew flagged as an out-of-scope "bigger kernel lift"; CC's shape-driven design made it emergent. - tests/addm/test_addm_waic_loo.py: sample with log_likelihood=True, assert the log-likelihood group is present, correctly shaped, and finite, and that az.loo / az.waic return finite estimates. Needs only the vendored likelihood (no ssm-simulators), so it runs in the normal suite; marked slow (it samples). - test_addm_subclass.py: correct the now-stale comment claiming the pointwise log-likelihood is unsupported (it is disabled there only for speed). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * CF: make the aDDM tutorial interactive + show the handshake Rework the marimo tutorial so it is a hands-on test bench, not just a linear demo: - Section 1: interactive sliders (eta, kappa, a, b) that re-simulate and re-plot the RT distribution + choice proportion live — poke the simulator. - Section 2: the covariate handshake, shown at the simulator level — the same params simulated Mode 2 (conditioned on observed fixations) vs Mode 1 (self-sampled) give different predictive distributions, plus an ASCII diagram of the full extra_fields -> RV._extra_fields -> rng_fn -> simulator(extra_fields=) -> cssm.addm Mode-2 flow. - Section 4: prints the built model's rv_op._extra_fields (the observed fixations actually wired into the generative path). - Sampling / recovery / PPC are gated behind a mo.ui.run_button, so the simulator and handshake sections are instant on load. Verified headlessly (marimo export html): every simulator/handshake/model-build cell runs; only the button-gated sample/recover/PPC cells wait (as intended). Header documents the local-build + `uv run --no-sync marimo edit` launch recipe. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * CF: drive the aDDM tutorial through the Simulator class Use ssm-simulators' high-level `Simulator("addm")` (configured once, then `sim.simulate(theta=..., extra_fields=...)`) instead of calling the low-level `cssm.addm` Cython function directly. This is the idiomatic public API AND the exact path HSSM's posterior-predictive checks go through, so the covariate handshake the notebook demonstrates is the real one, not a shortcut: - §1 explore: theta = (n_trials, 6) tiled + extra_fields(fixations). - §2 handshake: the same Simulator run with vs without extra_fields — Mode 2 (conditioned) vs Mode 1 (self-sampled) — confirmed visibly different (KS 0.078, p 2e-8), and the wiring diagram now ends at simulator()/Simulator. - §3 dataset: per-trial eta column in theta + extra_fields. Depends on the sibling ssm-simulators commit that adds extra_fields to Simulator.simulate(). Verified headlessly (marimo export html): all simulator/handshake/model cells run; only the button-gated sample/PPC cells wait. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * CF: document aDDM chain-parallelism reality in the tutorial The tutorial's sample() uses cores=1, which runs the 2 chains sequentially — this is deliberate, not a bug, but it looked suspicious. Add a comment with the actual reason, established empirically: - aDDM auto-selects the numpyro (JAX) NUTS sampler (approx_differentiable + jax), so sampling is already in JAX. - On CPU there is no chain parallelism to exploit: the FPT likelihood is heavy and one JAX chain already saturates all cores. Measured: two independent spawn processes -> 1.97x (contend, no win); numpyro chain_method="vectorized" -> 2.86x (batches ~2x compute, slower); numpyro default "parallel" and PyMC cores>1 fork both deadlock (JAX threads + process spawn). cores=1 is the safe, correct choice. - Genuine parallel chains need a GPU: numpyro chain_method="vectorized" vmaps both chains onto one device. Documented in the cell. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * CF: add diagnostic plots to the aDDM tutorial + fix sampler comment Make the fit section a proper diagnostic workflow (all button-gated, so load stays instant) and lighten the run: - Diagnostics section with four plots (verified to render): recovery — az.plot_posterior with ref_val = the ground-truth params, so each posterior visibly covers (or not) its truth line; trace — chain mixing/convergence; pair — posterior geometry with divergences; posterior predictive — model.plot_predictive(x_range=(-4, 4)), observed vs predicted RT densities. arviz imported once in the setup cell. - Drop draws 500 -> 300 and the dataset 800 -> 300 trials so the gated run is ~2-3 min instead of much longer (the per-iter FPT cost scales with trials). Also correct the (wrong) sampler comment. Measured on an 18-core CPU: aDDM already uses the numpyro JAX sampler (auto-selected for approx_differentiable + jax); a single chain uses only ~6/18 cores, but the FPT gradient is memory-bandwidth-bound (materializes per-term intermediates — the reason for TRUNC_NUM<=6), not core-bound, so parallel chains don't help: spawn 1.97x, chain_method="vectorized" 2.86x, forced-device "parallel" 2.24x, PyMC fork deadlocks. cores=1 is correct on CPU; real chain parallelism wants a GPU (numpyro chain_method="vectorized"). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * Gitignore marimo __marimo__/ session cache Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * refactor: ponytail trims (reuse sacc parse, pytest.raises) Minor over-engineering cleanup from a ponytail review (no behaviour change): - addm.py _validate_addm_columns: compute max_d by reusing _stack_sacc_array (single sacc-parsing path) instead of re-deriving rows + max_d, and drop the dead `len(rows) != n` check (a DataFrame column always has n rows) and the now-unused `n`. - tests: replace the hand-rolled _assert_raises helper (duplicated across two files) with the stdlib pytest.raises it reinvented. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(addm): type loglik_kind and prior_settings as Literals (mypy) Clears the 2 mypy errors failing #967's 3.13 CI job: - aDDMConfig.loglik_kind: str -> LoglikKind | None (matches BaseModelConfig's field type) - aDDM.__init__ prior_settings: str | None -> Literal["safe"] | None (matches HSSMBase) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(plotting): support aDDM (and RLSSM) model cartoons `plot_model_cartoon` is generic — it renders from *simulator* metadata — but it resolved `list_params`/`choices` via `default_model_config`, which is keyed only by built-in `SupportedModels`. HSSMBase subclasses (aDDM, RLSSM) are absent, so the lookup raised `KeyError("addm")`. - Add `_cartoon_model_meta(model)`: resolve `(list_params, choices)` from the registry for built-ins, else from `model.model_config` (aDDMConfig/RLSSMConfig). - Thread the resolved `(list_params, choices)` from `plot_model_cartoon` through `_plot_model_cartoon_1D`/`_2D`; the >2-choice renderer (`plot_func_model_n`) takes an explicit `choices` (registry fallback). The 2-choice path (aDDM) derives choices from sim metadata, so it needs no lookup. - Test: slow aDDM cartoon test (collapsing boundary + drift path + sample trajectories); skips on an ssm-sim build without the cartoon metadata. Registry generalization for subclass models is tracked in #1031. Requires the companion ssm-simulators change (boundary/trajectory/z metadata). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * style(addm): fix ruff lint + format across the aDDM likelihood, tests, and tutorial Clears the pre-existing ruff `check src/hssm` + `format --check .` failures on this branch: - likelihoods/jax/*: import sorting (I001), E501 (ruff-format wrapping + `# noqa` on the long vmap lambda headers), a D205/D202 docstring blank-line, and an unused `global _QUAD_CACHE` (it is only mutated via `.clear()`, not rebound — PLW0602). - tests/addm/* and docs/tutorials/attentional_ddm.py: `ruff format` normalization. Behaviour-neutral (whitespace / comments / import order only): full tests/addm/ green (65). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(addm): per-call fixation-continuation policy for posterior predictive Exposes the ssm-simulators fixation-continuation strategies on the aDDM: an aDDMConfig default plus a per-PPC-call override, so one fitted model can be swept across continuation policies (prolong_last_fixation / sample_continuation / resample_all_fixations) with no rebuild/re-fit. - aDDMConfig: `continuation_mode` / `continuation_params` fields (default = prolong, i.e. today's behaviour) + validation against the ssm-simulators registries (defensive import so an older ssms doesn't break aDDM import). - `aDDM._push_rv_extra_fields` seeds the generative `_extra_fields` conduit with the policy; the per-call override lives on `self._continuation_override`, which `_push_rv_extra_fields` reads — so the base PPC path's `_update_extra_fields` refresh preserves an active override instead of resetting it to the config default. - `aDDM.sample_posterior_predictive` override: `continuation_mode`/`continuation_params` applied for the call only (leak-free, restored in `finally`), then delegates to the base. - Tests: config validation + a per-call PPC override test (threads to the RV conduit, runs all three modes end-to-end, reverts with no leakage). Updated `test_addm_ppc`'s exact-keys assertion (the conduit now also carries the continuation policy). Full tests/addm/ green (67). Requires the companion ssm-simulators fixation-continuation change; test skip-guarded. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs(addm): thorough PPC + model-cartoon tutorial sections Demonstrates the full aDDM diagnostic surface on one fitted model: - §6 Posterior predictive: the default observed-fixation-conditioned check, plus the per-call fixation-continuation policies (prolong_last_fixation / sample_continuation / resample_all_fixations) — selectable on the SAME fitted model with no re-fit. - §7 Model cartoon: hssm.plotting.plot_model_cartoon for the aDDM (collapsing bounds, drift path, stochastic sample trajectories, non-decision time, choice-split RT histograms). Cells verified against the combined ssm-simulators build (cartoon metadata + continuation); `marimo check` clean; `ruff format --check .` clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs(addm): caveat that the model cartoon is not the fixation-conditioned PPC The aDDM cartoon re-simulates at the posterior-mean parameters with the simulator self-sampling its own fixations (Mode 1) and regenerates its predictive with the DEFAULT continuation policy, so it neither conditions on the observed fixations nor honors a per-call `continuation_mode`. Document this in the tutorial §7 and the `plot_model_cartoon` docstring so the cartoon is read as a drift/boundary schematic, not as the §6 fixation-conditioned check. The real fix (thread extra_fields + continuation policy through the cartoon path) is tracked in #1039. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(addm): adapt to pymc6 / arviz 1.0 after merging main Main landed the pymc6 migration (#1019) before this branch, so the merge needed arviz-1.0/DataTree adaptation on the aDDM code: - addm.py: type the PPC override with `xarray.DataTree` (was `az.InferenceData`); the delegation is positional so it already matched the base's renamed `dt`. - tests: `model.sample()` returns a DataTree — swap `isinstance(..., InferenceData)` for `DataTree` and `.groups()` membership for `in`. arviz 1.x dropped `az.waic` and renamed the LOO estimate to `ELPDData.elpd`; the WAIC/LOO test now samples enough draws for a valid PSIS tail and asserts on `az.loo(...).elpd`. - tutorial: `az.plot_posterior` was removed in arviz 1.x -> `az.plot_forest`. - config.py: drop the defensive `try/except ImportError` guard now that the `ssm-simulators>=0.13.1` pin (set in the merge) guarantees the aDDM engine + fixation_continuation module. Housekeeping so the merged config's stricter gates pass: - docstrings on the aDDM test functions (repo selects `D`, no test exemption); wrap two long lines; `# ruff: noqa` the generated marimo tutorial. - annotate `_QUAD_CACHE` (mypy) and widen `plot_function_kwargs` to `dict[str, Any]` so pyrefly accepts the dynamic `choices` kwarg. Full tests/addm suite green against released ssm-simulators 0.13.1 (68 passed). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * docs(addm): fix arviz-1.x tutorial breakage + stale docs found in review Adversarially-verified review findings on the merge/adaptation: - tutorial: az.plot_pair no longer accepts `kind`/`divergences` (hard ValueError on arviz 1.2) — drop them; update the diagnostics prose. - tutorial: plot_model_cartoon's data arg was renamed idata->dt on main; the old keyword was silently swallowed by **kwargs (falling back to model.traces), teaching a dead parameter — pass dt=. - tutorial: restore the ground-truth values lost in the plot_posterior->plot_forest swap (truth now shown in the figure title). - tutorial: module docstring still gave pre-release "install a local ssm-simulators build" instructions; ssm-simulators >= 0.13.1 ships the aDDM engine, so a plain `uv sync` suffices. - changelog: v0.4.0 is already tagged — move the aDDM (and this week's noncentered) entries out of released sections into "Unreleased", note the bambi/ssm-simulators floor bumps there, and restore the released 0.4.0 section to what actually shipped (Bambi 0.18+). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * test(addm): fix global-precision leaks that broke the full suite Full-suite enumeration (sequential, CI ordering, no --exitfirst) surfaced 7 failures in two leak directions plus one integration gap: - Inbound (fast-CI blockers): pytest imports every module at collection, so a later module's import-time `hssm.set_floatX("float32")` flipped JAX x64 off after tests/addm's import-time `set_jax_precision(True)`, desyncing the vendored kernel's cached dtype — builder/equiv/subclass bodies then computed float32 against a float64 contract. New tests/addm/conftest.py autouse fixture pins float64 per-test (pytensor floatX + JAX x64 + kernel cache). - Outbound: test_addm_waic_loo's body set_floatX("float64") leaked into later float32-dependent LAN/ONNX tests (tests/slow/test_mcmc approx_differentiable failed). The fixture restores the previous precision after each aDDM test. - tests/rl choice-only smoke: ssm-simulators 0.13 ships the ssms.rl presets (2AB_RW_InvTempSoftmax et al.) so the old guard un-skipped them, but the presets are not registered in ssms' model_config, which HSSM's RV path requires — skip until the RLSSM wiring lands (#1000). Validated: inbound repro (builder/equiv/subclass after a float32 import) 40 passed; outbound pair (waic_loo then the mcmc param) 2 passed; rl smoke 4 skipped; full tests/addm 68 passed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * test(rl): point the choice-only preset skip at tracking issue #1052 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Alexander Fengler <alexanderfengler@gmx.de> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
hssm.RSSSM, a new top-level class for fitting regime-switching sequential sampling models (issue #957). Each trial belongs to one ofKhidden 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 batchedpytensor.scan, contributed as a scalarpm.Potential), leaving only continuous parameters for NUTS — the same approach as the hand-builtdocs/tutorials/hmm_ddm_regime_switching.ipynb, but as a high-levelHSSM(...)/RLSSM(...)-style class.Implements the design on the
issue_957/plan_hssm_hmm_classbranch (docs/design/hssm_hmm.md), Phases 2–5.What's included
src/hssm/hmm/): subpackage withrsssm.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).Kand inN(participants); unbalanced panels via end-padding + emission mask.approx_differentiable, jax & pytensor).ordered-transform anchor —AutoOrdering(default),OrderByParam,NoOrdering.switching_params= inferred-per-regime.pooling="none");P/pi0stay global. Estimable globalpi0.p_outlierlapse mixture ((1 - p_k)·SSM_k + p_k·lapse).infer_regimes(FFBS),compute_log_likelihood(per-trial logp forarviz.loo/waic),plot_regime_recovery.Deliberate deviations from the design doc (please review)
Two change documented behavior and warrant explicit sign-off:
log_likelihoodgroup 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 toarviz.looaslogp = 0("perfectly predicted") observations, biasingloo/waicon unbalanced panels.Minor/cosmetic:
vi()raisesNotImplementedError(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
p_outlier, end-to-end FFBS +loo). 88% line coverage onhmm/.sum(delta) == marginal == model Potential; unbalancedloocounts real trials only.p_outlierrecovery; LAN backends agree to 1e-15.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-drivenP, per-regime regressions, cross-emission/semi-Markov, the posterior-predictive family.Closes #957.
🤖 Generated with Claude Code