[pymc6 migration] Updated main documentation

[digicosmos86]

Mainly updated .md file and one tutorial (getting_started.ipynb)

Summary by CodeRabbit

Release Notes: HSSM v0.4.0

• New Features

  • Added support for reinforcement learning sequential sampling models (RLSSM).

• Breaking Changes

  • sample() now returns xarray.DataTree instead of arviz.InferenceData.
  • vi() return type updated; returns Approximation and DataTree.
  • Removed plot_trace() and summary() convenience methods; use ArviZ directly.
  • Installation now requires pip/uv only; Conda support removed.
  • Python 3.11 no longer supported; requires Python 3.12+.
  • Default sampling backend changed to numba.

• Documentation

  • Updated installation guides, changelog, and API documentation for v0.4.0.

• Chores

  • Streamlined core dependencies to PyMC 6.0+, ArviZ 1.0+, and Bambi 0.18+.
  • Updated Python version matrix in CI workflows to 3.12–3.13.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[coderabbitai]

Caution

Review failed

Failed to post review comments.

We encountered an issue with GitHub. Use @coderabbitai full review to retry the review.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 4d168dd0-e3fa-4ac6-b27b-051cd16e964d

📥 Commits

Reviewing files that changed from the base of the PR and between 1956e91 and fc7dd26.

⛔ Files ignored due to path filters (2)

• test_ellipse_output.png is excluded by !**/*.png
• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (37)

• .github/workflows/check_notebooks.yml
• .github/workflows/run_tests.yml
• .gitignore
• README.md
• docs/api/hssm.md
• docs/changelog.md
• docs/getting_started/getting_started.ipynb
• docs/getting_started/installation.md
• docs/index.md
• docs/local_development.md
• docs/overrides/main.html
• pyproject.toml
• src/hssm/base.py
• src/hssm/distribution_utils/dist.py
• src/hssm/likelihoods/analytical.py
• src/hssm/plotting/model_cartoon.py
• src/hssm/plotting/predictive.py
• src/hssm/plotting/quantile_probability.py
• src/hssm/plotting/utils.py
• src/hssm/utils.py
• tests/conftest.py
• tests/rl/test_rl_likelihood_builder.py
• tests/slow/test_choice_only.py
• tests/slow/test_mcmc.py
• tests/slow/test_missing_data_and_deadline_mcmc.py
• tests/slow/test_missing_data_and_deadline_vi.py
• tests/slow/test_missing_data_mcmc.py
• tests/slow/test_missing_data_vi.py
• tests/slow/test_vi.py
• tests/test_hssm.py
• tests/test_initvals.py
• tests/test_likelihoods.py
• tests/test_plotting.py
• tests/test_rlssm.py
• tests/test_sample_posterior_predictive.py
• tests/test_save_load.py
• tests/test_utils.py

💤 Files with no reviewable changes (1)

• docs/api/hssm.md

🧰 Additional context used

📓 Path-based instructions (5)

**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

**/*.py: Use ruff for linting and formatting (via pre-commit hooks)
Run mypy for type checking on all Python files

Files:

• tests/test_likelihoods.py
• tests/test_save_load.py
• src/hssm/likelihoods/analytical.py
• tests/test_utils.py
• tests/test_initvals.py
• tests/slow/test_missing_data_and_deadline_vi.py
• tests/test_rlssm.py
• tests/rl/test_rl_likelihood_builder.py
• tests/conftest.py
• tests/test_hssm.py
• src/hssm/distribution_utils/dist.py
• tests/slow/test_mcmc.py
• tests/slow/test_choice_only.py
• tests/slow/test_vi.py
• tests/test_sample_posterior_predictive.py
• src/hssm/plotting/predictive.py
• tests/slow/test_missing_data_and_deadline_mcmc.py
• src/hssm/plotting/utils.py
• tests/slow/test_missing_data_mcmc.py
• tests/slow/test_missing_data_vi.py
• src/hssm/utils.py
• src/hssm/plotting/quantile_probability.py
• tests/test_plotting.py
• src/hssm/plotting/model_cartoon.py
• src/hssm/base.py

tests/**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

Mark slow tests with @pytest.mark.slow decorator

Files:

• tests/test_likelihoods.py
• tests/test_save_load.py
• tests/test_utils.py
• tests/test_initvals.py
• tests/slow/test_missing_data_and_deadline_vi.py
• tests/test_rlssm.py
• tests/rl/test_rl_likelihood_builder.py
• tests/conftest.py
• tests/test_hssm.py
• tests/slow/test_mcmc.py
• tests/slow/test_choice_only.py
• tests/slow/test_vi.py
• tests/test_sample_posterior_predictive.py
• tests/slow/test_missing_data_and_deadline_mcmc.py
• tests/slow/test_missing_data_mcmc.py
• tests/slow/test_missing_data_vi.py
• tests/test_plotting.py

.github/workflows/check_notebooks.yml

📄 CodeRabbit inference engine (CLAUDE.md)

Use SKIP_NOTEBOOKS environment variable in CI workflows to skip notebook execution during notebook checks

Files:

• .github/workflows/check_notebooks.yml

docs/changelog.md

📄 CodeRabbit inference engine (CLAUDE.md)

Changelog headings should have no prefix (e.g., ### 0.3.0)

Files:

• docs/changelog.md

pyproject.toml

📄 CodeRabbit inference engine (CLAUDE.md)

pyproject.toml: Version in pyproject.toml should have no prefix (e.g., 0.3.0)
bayesflow should be installed from dev branch (git+https://github.com/bayesflow-org/bayesflow@dev) in the notebook dependency group only

Files:

• pyproject.toml

🧠 Learnings (2)

📚 Learning: 2026-06-17T23:10:29.060Z

Learnt from: digicosmos86
Repo: lnccbrown/HSSM PR: 1001
File: tests/conftest.py:185-185
Timestamp: 2026-06-17T23:10:29.060Z
Learning: When reviewing HSSM code that uses ArviZ, note that for ArviZ 1.0+ (arviz>=1.0) `arviz.from_netcdf()` returns an `xarray.DataTree` (not `arviz.InferenceData`). Therefore, assignments like `model._inference_obj = az.from_netcdf(...)` should be treated as type-compatible when `_inference_obj` is annotated as `xr.DataTree | None`. Do not flag a type-mismatch for this pattern under ArviZ>=1.0; only flag if the target type is incompatible or if the project is pinned to an older ArviZ version.

Applied to files:

• tests/test_likelihoods.py
• tests/test_save_load.py
• src/hssm/likelihoods/analytical.py
• tests/test_utils.py
• tests/test_initvals.py
• tests/slow/test_missing_data_and_deadline_vi.py
• tests/test_rlssm.py
• tests/rl/test_rl_likelihood_builder.py
• tests/conftest.py
• tests/test_hssm.py
• src/hssm/distribution_utils/dist.py
• tests/slow/test_mcmc.py
• tests/slow/test_choice_only.py
• tests/slow/test_vi.py
• tests/test_sample_posterior_predictive.py
• src/hssm/plotting/predictive.py
• tests/slow/test_missing_data_and_deadline_mcmc.py
• src/hssm/plotting/utils.py
• tests/slow/test_missing_data_mcmc.py
• tests/slow/test_missing_data_vi.py
• src/hssm/utils.py
• src/hssm/plotting/quantile_probability.py
• tests/test_plotting.py
• src/hssm/plotting/model_cartoon.py
• src/hssm/base.py

📚 Learning: 2026-06-16T20:52:09.884Z

Learnt from: digicosmos86
Repo: lnccbrown/HSSM PR: 984
File: tests/slow/test_mcmc.py:139-139
Timestamp: 2026-06-16T20:52:09.884Z
Learning: In this repository, do not suggest adding `arviz` as an explicit top-level dependency in `pyproject.toml`. `arviz` is intentionally pulled transitively via `bambi` (which declares `arviz` as a direct dependency) to keep `arviz` version compatibility aligned with `bambi`. Only consider making `arviz` explicit if there’s a deliberate, repo-wide reason to override this transitive/version-coupling behavior.

Applied to files:

• pyproject.toml

🪛 ast-grep (0.44.0)

src/hssm/base.py

[warning] 1479-1479: File path is request-/variable-derived; validate and normalize to prevent path traversal.
Context: open(model_path.joinpath("model.pkl"), "wb")
Note: [CWE-22] Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal').

(open-filename-from-request)

🪛 markdownlint-cli2 (0.22.1)

docs/changelog.md

[warning] 3-3: Heading levels should only increment by one level at a time
Expected: h2; Actual: h3

(MD001, heading-increment)

docs/index.md

[warning] 154-154: Code block style
Expected: fenced; Actual: indented

(MD046, code-block-style)

📝 Walkthrough

Walkthrough

HSSM v0.4.0 replaces arviz.InferenceData with xarray.DataTree as the return type for all sampling, VI, log-likelihood, and predictive methods across HSSMBase and all plotting modules. The release also drops Python 3.11, switches from conda to pip/uv installation, bumps core dependencies (PyMC 6, ArviZ 1.0+, Bambi 0.18+), deprecates model.summary()/model.plot_trace() to NotImplementedError, and adds an optional hssm[io] extra for HDF5 backends.

Changes

HSSM v0.4.0: xarray DataTree migration, PyPI packaging, Python 3.12+

Layer / File(s)	Summary
Version bump, dependency updates, CI, and tooling pyproject.toml, .github/workflows/run_tests.yml, .github/workflows/check_notebooks.yml, .gitignore	Bumps version to 0.4.0, raises Python floor to 3.12, replaces hatchling with uv_build, adds io optional extra (h5netcdf/h5py), drops Python 3.11 from CI matrix, adds uv.lock to .gitignore.
DataTree/Approximation type contracts in HSSMBase and utils src/hssm/base.py, src/hssm/utils.py	Introduces DataTree and Approximation imports, retypes _inference_obj/_inference_obj_vi, updates sample()/vi() return types, adds _requires_io_backends IO guard, renames predictive_idata_to_dataframe to predictive_dt_to_dataframe, updates traces/vi_idata/vi_approx properties.
Sampling, VI, log-likelihood, and posterior cleanup src/hssm/base.py, src/hssm/utils.py	Rewrites vi(), _clean_posterior_group, log_likelihood(), add_likelihood_parameters_to_datatree(), and _compute_log_likelihood() to operate on DataTree groups; removes numpyro jitter auto-edit and tqdm progress bar; updates DataTree-specific cleanup helpers.
Posterior/prior predictive sampling, save/load IO src/hssm/base.py	Rewrites sample_posterior_predictive() with DataTree draw-slicing and xr.concat in safe mode; updates predict(), sample_do(), sample_prior_predictive() to return DataTree; renames save_idata_only to save_traces_only; implements load_model_traces; deprecates plot_trace/summary to raise NotImplementedError.
Plotting utilities: DataTree replacing InferenceData src/hssm/plotting/utils.py, src/hssm/plotting/predictive.py	Migrates _get_plotting_df, _use_traces_or_sample, and _to_dt_group to use xr.DataTree; changes predictive sample extraction to dt[predictive_group][response_str]; changes _histogram return to pd.Series.
Public plotting functions: predictive, quantile_probability, model_cartoon src/hssm/plotting/predictive.py, src/hssm/plotting/quantile_probability.py, src/hssm/plotting/model_cartoon.py	Updates all three public plot functions to accept `dt: xr.DataTree
PyTensor floatX path fix in likelihoods src/hssm/distribution_utils/dist.py, src/hssm/likelihoods/analytical.py	Replaces pm.floatX with pm.pytensorf.floatX for LOGP_LB, parameter bounds, fixed-vector substitutions, and local constants in logp_poisson_race.
Test fixtures and core HSSM tests tests/conftest.py, tests/test_hssm.py, tests/test_save_load.py, tests/test_sample_posterior_predictive.py, tests/test_utils.py, tests/test_likelihoods.py, tests/rl/test_rl_likelihood_builder.py, tests/test_rlssm.py	Renames cav_idata fixture to cav_dt; wraps mock posteriors in xr.DataTree; updates core tests to use add_likelihood_parameters_to_datatree, dt= keyword, and predictive_dt_to_dataframe; asserts xr.DataTree load types; marks RL/RLSSM tests xfail.
Slow test suite and plotting tests tests/slow/test_*.py, tests/test_plotting.py	Updates all slow MCMC/VI/choice-only/missing-data tests to assert xr.DataTree, compare groups with xr.testing.assert_equal, call az.plot_trace_dist instead of model.plot_trace, and pass progressbar=False; updates all plotting tests to use cav_dt fixture and dict-style trace access.
Documentation and changelog for v0.4.0 README.md, docs/changelog.md, docs/getting_started/installation.md, docs/index.md, docs/local_development.md, docs/overrides/main.html, docs/api/hssm.md	Adds v0.4.0 changelog; rewrites installation docs to pip/uv (removing conda); adds RLSSM/ArviZ/ONNX feature bullets; updates troubleshooting to Python 3.12–3.14; bumps site announcement banner; removes summary/plot_trace from API docs.

Sequence Diagrams

sequenceDiagram
  participant User
  participant HSSMBase
  participant PyMC
  participant DataTree
  participant PlottingModule

  User->>HSSMBase: model.sample(draws, tune, ...)
  HSSMBase->>PyMC: pm.sample(...)
  PyMC-->>HSSMBase: raw InferenceData (internal)
  HSSMBase->>DataTree: _clean_posterior_group(dt=...)
  DataTree-->>HSSMBase: cleaned DataTree
  HSSMBase-->>User: DataTree

  User->>HSSMBase: model.log_likelihood(dt)
  HSSMBase->>DataTree: read dt["posterior"] chain/draw
  HSSMBase->>DataTree: assign dt["log_likelihood"]
  DataTree-->>HSSMBase: updated DataTree
  HSSMBase-->>User: DataTree

  User->>HSSMBase: model.plot_predictive(dt=dt)
  HSSMBase->>PlottingModule: plot_predictive(dt=dt)
  PlottingModule->>DataTree: dt[predictive_group][response_str]
  PlottingModule-->>User: Axes | FacetGrid

Loading

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120 minutes

Possibly related PRs

• lnccbrown/HSSM#984: Updates slow MCMC test expectations and replaces model.plot_trace(...) with az.plot_trace_dist(...), directly overlapping with this PR's test changes in tests/slow/test_*_mcmc.py.
• lnccbrown/HSSM#991: Refactors sample_posterior_predictive, sample_prior_predictive, and sample_do in src/hssm/base.py to operate on xarray.DataTree groups, which is the same core change in this PR.
• lnccbrown/HSSM#1001: Modifies src/hssm/plotting/model_cartoon.py to switch from ArviZ InferenceData to xarray DataTree, directly overlapping with the model cartoon plotting changes in this PR.

Suggested reviewers

• AlexanderFengler
• krishnbera

Poem

🐇 A bunny hops through DataTree boughs,
No InferenceData slowing the cause!
pip install hssm, Python 3.12,
Traces now branch like leaves on the shelves.
v0.4.0 — the forest grows wise! 🌳

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Title check	⚠️ Warning	The title '[pymc6 migration] Updated main documentation' is partially related to the changeset but does not accurately reflect the scope of changes.	The PR includes comprehensive code changes far beyond documentation updates—including ArviZ-to-DataTree migration, API updates, type system changes, and test modifications. A more accurate title would be '[pymc6 migration] Migrate InferenceData to DataTree and update documentation' or similar.
Docstring Coverage	⚠️ Warning	Docstring coverage is 74.07% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch pymc6-migration-doc-updates

⚔️ Resolve merge conflicts

• Resolve merge conflict in branch pymc6-migration-doc-updates

---

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