blackjax-devs / blackjax

BlackJAX is a Bayesian Inference library designed for ease of use, speed and modularity.
https://blackjax-devs.github.io/blackjax/
Apache License 2.0
806 stars 105 forks source link

Add MCLMC sampler #586

Closed reubenharry closed 10 months ago

reubenharry commented 10 months ago

Todos:

As per #530, this PR adds MCLMC (Microcanonical Langevin Monte Carlo) as a sampling method to BlackJax. It is an adaptation of the original repo: https://github.com/JakobRobnik/MicroCanonicalHMC/tree/master.

In particular, this is an implementation of the single chain partially momentum updating algorithm (as opposed to full updates every n steps). It also will implement (once this PR is ready) the MCLMC tuning algorithm.

codecov[bot] commented 10 months ago

Codecov Report

Attention: 1 lines in your changes are missing coverage. Please review.

Comparison is base (f49945d) 99.16% compared to head (70f1dd5) 99.16%.

Files Patch % Lines
blackjax/mcmc/mclmc.py 97.67% 1 Missing :warning:
Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #586 +/- ## ======================================== Coverage 99.16% 99.16% ======================================== Files 52 54 +2 Lines 2391 2515 +124 ======================================== + Hits 2371 2494 +123 - Misses 20 21 +1 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

reubenharry commented 10 months ago

After a bit of a mammoth effort, the PR is ready for review.

One possible point of objection is that there are still some poorly named variables, now confined to certain functions in the adaptation code. These I have not renamed because I don't fully understand what all of them are doing, so some reading/talking to Jakob is in order.

But there's now a test of both the tuning and the kernel, that uses pytrees, showing that the code is pytree generic. Moreover, the implementation reproduces the original MCLMC implementation perfectly (with appropriate choices of random seeds), so other than a few details (I removed the preconditioning for now, for simplicity), it's hopefully looking good.