EMAUS (Ensemble Microcanonical Adjusted-Unadjusted Sampler)

[reubenharry]

This is in essence a tuning scheme for running multiple chains in parallel (ensemble). The key trick is to use MCLMC without MH (i.e. unadjusted) first in order to very quickly find the typical set, and then to use MH for the steps you actually count. The details are in how to decide when to switch between these two phases, and how to tune the hyperparameters using estimates from the ensemble of chains.

The upshot is that the method is fast, even compared to other ensemble methods: sampling book: TODO

Changes:

• a new method, emaus: this consists of two files, both in the adaptation folder
• util.py: some general purpose helper functions, run_eca, eca_step and ensemble_execute_fn. These are helper functions for running batches (ensembles) of chains, optionally across multiple nodes

A few important guidelines and requirements before we can merge your PR:

• We should be able to understand what the PR does from its title only;
• There is a high-level description of the changes;
• There are links to all the relevant issues, discussions and PRs;
• The branch is rebased on the latest main commit;
• Commit messages follow these guidelines;
• The code respects the current naming conventions;
• Docstrings follow the numpy style guide
• pre-commit is installed and configured on your machine, and you ran it before opening the PR;
• There are tests covering the changes;
• The doc is up-to-date (pending EMAUS (Ensemble Microcanonical Adjusted-Unadjusted Sampler) sampling-book#68)
• If I add a new sampler* I added/updated related examples

[reubenharry]

I unmarked this as a draft, even though there's still some rough edges, just in case @junpenglao or anyone else wanted to give high level feedback over the main changes (e.g. should I rearrange the files and put everything in one file, etc).

I see I have also somehow added spaces to a bunch of files, not sure how this happened or how to fix...

[reubenharry]

This should be ready to review @junpenglao . I realize it's a large PR, so let me know if anything can be done to make it easier.

Confusingly, the new test passes locally, but fails on CI. I don't understand what has changed (the random seed seems to be the same).

[junpenglao]

Great work, let me take a look this week.

[junpenglao]

I don't understand what has changed (the random seed seems to be the same).

The latest release of JAX seems to change some of the seeding output, try either a new random key or in crease the number of samples.