[Draft] Nested sampling implementation

[yallup]

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

• If I add a new sampler, there is an issue discussing it already; Nested Sampling implementation #753
• 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;
• If I add a new sampler* I added/updated related examples

Consider opening a Draft PR if your work is still in progress but you would like some feedback from other contributors.

High level description of changes

The following files are included in the folder ns:

• base: The base nested sampler. Detailed more below, should be somewhat familiar to the SMC structure. Nested sampling as an outer kernel with a delete function (resampling) to remove the lowest likelihood points, and then maps a vectorized update over those “deleted” particles to replace them with new particles subject to the hard likelihood constraint.

• adaptive: Similar to the SMC inner kernel tuning, wraps the base sampler with a parameter update function to tune the inner kernel parameters.

• utils: Useful calculations, particularly for extracting log_weights, weighted (at a specified inverse temperature) samples of the target

• vectorized_slice: A compatible inner kernel for the nested sampling kernels, this is non-standard for the rest of the library so opinions on how best to do this are most welcome, we tried to follow the SMC design of flexible choice of inner kernel, but currently only this one works... Currently this explicitly loads both the prior logdensity and loglikelihood as functions, as we would think about them in nested sampling. But I suspect there is a clever way to lift this to be in the mcmc folder, and overload the extra loglikelihood condition for use in nested sampling. For now we have a practical implementation here that works for our purpose. Currently this doesn’t use a proposal distribution as in the mh kernels, allowing a more flexible definition of a random slice direction, and instead hardcodes a derived from a covariance.

Out of these there are currently 3 top level APIs defined (which is somewhat overkill as things stand but hopefully it translates). Base and adaptive both have top level apis, named generically as per the usual design. Inside adaptive we have put a top level api for nss or "nested slice sampling", that loads explicitly the vectorized slice inner kernel and corresponding tuning.

Example usage

Lastly there is an example usage script (not to be included in the final PR but to help demonstrate how we intend these components to be used on a simple gaussian-gaussian integration problem. Under docs/examples/nested_sampling.py (this has an external dependency of distrax). I have added a number of inline comments in this to explain some choices, this would all be extended at some point and folded into the sampling book rather than here but I’ve included it as a tracked file for convenience.

As there are quite a few non standard parts here I will submit this as a draft PR for now, hoping for some higher level feedback before getting too far into the weeds. Hopefully there is enough here for you to go on @AdrienCorenflos as an initial look and the example works out of the box for you.

[AdrienCorenflos]

Apologies, I was travelling, I will check this ASAP