PyMC implementation of WALNUTS sampler

[fonnesbeck]

Description

First pass at implementing WALNUTS, partly as a learning opportunity. It runs with pm.sample but does not yet produce a good sample. Still lots of work to do.

Related Issue

• Closes #
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pymc--7836.org.readthedocs.build/en/7836/

[twiecki]

Why not hook up to the walnuts implrementation Bob et al provide?

[fonnesbeck]

This borrows directly and heavily from Bob's repo. Its really not a lot of code.

[codecov]

Codecov Report

❌ Patch coverage is 6.87773% with 853 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.81%. Comparing base (3ae5095) to head (4b474b5).
⚠️ Report is 42 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc/step_methods/hmc/adaptive_integrators.py	3.41%	424 Missing ⚠️
pymc/step_methods/hmc/walnuts.py	8.41%	381 Missing ⚠️
pymc/step_methods/hmc/p2_quantile.py	14.28%	48 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7836      +/-   ##
==========================================
- Coverage   92.91%   88.81%   -4.11%     
==========================================
  Files         107      111       +4     
  Lines       18286    19201     +915     
==========================================
+ Hits        16991    17053      +62     
- Misses       1295     2148     +853     

Files with missing lines	Coverage Δ
pymc/step_methods/__init__.py	100.00% <100.00%> (ø)
pymc/step_methods/hmc/__init__.py	100.00% <100.00%> (ø)
pymc/step_methods/hmc/walnuts_constants.py	100.00% <100.00%> (ø)
pymc/step_methods/hmc/p2_quantile.py	14.28% <14.28%> (ø)
pymc/step_methods/hmc/walnuts.py	8.41% <8.41%> (ø)
pymc/step_methods/hmc/adaptive_integrators.py	3.41% <3.41%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[twiecki]

This borrows directly and heavily from Bob's repo. Its really not a lot of code.

Sure, but they'll keep iterating on it so we'd always need to keep up with changes. Anyway, I'm not against it, we can always revert back to their implementation if we find it's a great sampler. Or pymc-extras might also be a place.

[fonnesbeck]

Yeah, this should probably be in pymc-extras. I will move it.

I think that if it turns out to be a performant algorithm we will want our own implementation of it (doesn't have to be this one, of course).

[fonnesbeck]

Moved to pymc-devs/pymc-extras#530