Skip to main content

Nested Sampling in JAX

Project description

JAXNS

What is it?

Enables probabilistic programming using nested sampling. It's coded in JAX in a manner that allows lowering the entire inference algorithm to XLA primitives, which are JIT-compiled for high performance. You can read about it here: (https://arxiv.org/abs/2012.15286)

JAXNS provides a constrained likelihood sampler which combines and modifies ideas from MultiNest (F. Feroz et al. 2008; https://arxiv.org/pdf/0809.3437.pdf) and PolyChord (W.J. Handley et al. 2015; https://arxiv.org/abs/1506.00171). There are two samplers available provided by setting sampler_name=slice for slice sampling, and sampler_name=multi_ellipsoid for rejection sampling.

Install

Make sure you have JAX and the usual suspects with pip install jax jaxlib numpy matplotlib scipy. Install with pip install jaxns or pip install git+http://github.com/Joshuaalbert/jaxns.git.

Quick start

JAXNS is really fast because it uses JAX. The caveat is that you need to be able to define your likelihood function with JAX. This is usually no big deal because JAX is just a replacement for numpy and many likelihoods can be expressed such. If you're unfamiliar, take a quick tour of JAX (https://jax.readthedocs.io/en/latest/notebooks/quickstart.html). For a simple example see the simple multivariate normal likelihood with conjugate prior example

Speed test comparison with other nested sampling packages

JAXNS is much faster than PolyChord, MultiNEST, and dynesty, typically achieving two to three orders of magnitude improvement in speed. I show this in (https://arxiv.org/abs/2012.15286).

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

jaxns-0.0.3.tar.gz (70.4 kB view hashes)

Uploaded Source

Built Distribution

jaxns-0.0.3-py3-none-any.whl (101.4 kB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page