Package: jaxns
Main
Status:Develop
Status:Mission:To make nested samplingfaster, easier, and more powerfulWhat is it?JAXNS is:a probabilistic programming framework using nested sampling as the engine;coded in JAX in a manner that allows lowering the entire inference algorithm to XLA primitives, which are
JIT-compiled for high performance;continuously improving on its mission of making nested sampling faster, easier, and more powerful; andcitable, use the(old) pre-print here.JAXNS Probabilistic Programming FrameworkJAXNS provides a powerful JAX-based probabilistic programming framework, which allows you to define probabilistic
models easily, and use them for advanced purposes. Probabilistic models can have both Bayesian and parameterised
variables.
Bayesian variables are random variables, and are sampled from a prior distribution.
Parameterised variables are point-wise representations of a prior distribution, and are thus not random.
Associated with them is the log-probability of the prior distribution at that point.Let's break apart an example of a simple probabilistic model. Note, this example can also be followed
indocs/examples/intro_example.ipynb.Defining a probabilistic modelPrior models are functions that produce generators ofPriorobjects.
The function must eventually return the inputs to the likelihood function.
The returned values of a yieldedPrioris a simple JAX array, i.e. you can do anything you want to it with JAX ops.
The rules of static programming apply, i.e. you cannot dynamically allocate arrays.JAXNS makes use of the Tensorflow Probability library for defining prior distributions, thus you can usealmostany of the TFP distributions. You can also use any of the TFP bijectors to define transformed distributions.Distributions do have some requirements to be valid for use in JAXNS.They must have a quantile function, i.e.dist.quantile(dist.cdf(x)) == x.They must have alog_probmethod that returns the log-probability of the distribution at a given value.Most of the TFP distributions satisfy these requirements.JAXNS has some special priors defined that can't be defined from TFP, seejaxns.framework.special_priors. You can
always request more if you need them.Prior variablesmaybe named but don't have to be. If they are named then they can be collected later via a
transformation, otherwise they are deemed hidden variables.The output values of prior models are the inputs to the likelihood function. They can be PyTree's,
e.g.typing.NamedTuple's.Finally, priors can become point-wise estimates of the prior distribution, by callingparametrised(). This turns a
Bayesian variable into a parameterised variable, e.g. one which can be used in optimisation.importjaximporttensorflow_probability.substrates.jaxastfptfpd=tfp.distributionsfromjaxns.framework.modelimportModelfromjaxns.framework.priorimportPriordefprior_model():mu=yieldPrior(tfpd.Normal(loc=0.,scale=1.))# Let's make sigma a parameterised variablesigma=yieldPrior(tfpd.Exponential(rate=1.),name='sigma').parametrised()x=yieldPrior(tfpd.Cauchy(loc=mu,scale=sigma),name='x')uncert=yieldPrior(tfpd.Exponential(rate=1.),name='uncert')returnx,uncertdeflog_likelihood(x,uncert):returntfpd.Normal(loc=0.,scale=uncert).log_prob(x)model=Model(prior_model=prior_model,log_likelihood=log_likelihood)# You can sanity check the model (always a good idea when exploring)model.sanity_check(key=jax.random.PRNGKey(0),S=100)# The size of the Bayesian part of the prior space is `model.U_ndims`.Sampling and transforming variablesThere are two spaces of samples:U-space: samples in base measure space, and is dimensionless, or rather has units of probability.X-space: samples in the space of the model, and has units of the prior variable.# Sample the prior in U-space (base measure)U=model.sample_U(key=jax.random.PRNGKey(0))# Transform to X-spaceX=model.transform(U=U)# Only named Bayesian prior variables are returned, the rest are treated as hidden variables.assertset(X.keys())=={'x','uncert'}# Get the return value of the prior model, i.e. the input to the likelihoodx_sample,uncert_sample=model.prepare_input(U=U)Computing log-probabilitiesAll computations are based on the U-space variables.# Evaluate different parts of the modellog_prob_prior=model.log_prob_prior(U)log_prob_likelihood=model.log_prob_likelihood(U,allow_nan=False)log_prob_joint=model.log_prob_joint(U,allow_nan=False)Computing gradients of the joint probability w.r.t. parametersinit_params=model.paramsdeflog_prob_joint_fn(params,U):# Calling model with params returns a new model with the params setreturnmodel(params).log_prob_joint(U,allow_nan=False)value,grad=jax.value_and_grad(log_prob_joint_fn)(init_params,U)Nested Sampling EngineGiven a probabilistic model, JAXNS can perform nested sampling on it. This allows computing the Bayesian evidence and
posterior samples.fromjaxnsimportDefaultNestedSamplerns=DefaultNestedSampler(model=model,max_samples=1e5)# Run the samplertermination_reason,state=ns(jax.random.PRNGKey(42))# Get the resultsresults=ns.to_results(termination_reason=termination_reason,state=state)To AOT or JIT-compile the sampler# Ahead of time compilation (sometimes useful)ns_aot=jax.jit(ns).lower(jax.random.PRNGKey(42)).compile()# Just-in-time compilation (usually useful)ns_jit=jax.jit(ns)You can inspect the results, and plot them.fromjaxnsimportsummary,plot_diagnostics,plot_cornerplotsummary(results)plot_diagnostics(results)plot_cornerplot(results)Output:--------
Termination Conditions:
Small remaining evidence
--------
likelihood evals: 149918
samples: 3780
phantom samples: 1710
likelihood evals / sample: 39.7
phantom fraction (%): 45.2%
--------
logZ=-1.65 +- 0.15
H=-1.13
ESS=132
--------
uncert: mean +- std.dev. | 10%ile / 50%ile / 90%ile | MAP est. | max(L) est.
uncert: 0.68 +- 0.58 | 0.13 / 0.48 / 1.37 | 0.0 | 0.0
--------
x: mean +- std.dev. | 10%ile / 50%ile / 90%ile | MAP est. | max(L) est.
x: 0.07 +- 0.62 | -0.57 / 0.06 / 0.73 | 0.0 | 0.0
--------Using the posterior samplesNested sampling produces weighted posterior samples. To use for most use cases, you can simply resample (with
replacement).fromjaxnsimportresamplesamples=resample(key=jax.random.PRNGKey(0),samples=results.samples,log_weights=results.log_dp_mean,S=1000,replace=True)DocumentationYou can read the documentationhere. In addition, JAXNS is partially
described in theoriginal paper, as well as the paper onPhantom-Powered Nested
Sampling paper.InstallNotes:JAXNS requires >= Python 3.9. It is always highly recommended to use the latest version of Python.It is always highly recommended to use a unique virtual environment for each project.
To useminiconda, ensure it is installed on your system, then run the following commands:# To create a new env, if necessarycondacreate-njaxns_pypython=3.11
condaactivatejaxns_pyFor end usersInstall directly from PyPi,pipinstalljaxnsFor developmentClone repogit clone https://www.github.com/JoshuaAlbert/jaxns.git, and install:cdjaxns
pipinstall-rrequirements.txt
pipinstall-rrequirements-tests.txt
pipinstall-rrequirements-examples.txt
pipinstall.Getting help and contributing examplesDo you have a neat Bayesian problem, and want to solve it with JAXNS?
I'm really encourage anyone in either the scientific community or industry to get involved and join the discussion
forum.
Please use thegithub discussion forumfor getting help, or
contributing examples/neat use cases.Quick startCheckout the exampleshere.CaveatsThe 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).Speed test comparison with other nested sampling packagesJAXNS is really fast because it uses JAX.
JAXNS is much faster than PolyChord, MultiNEST, and dynesty, typically achieving two to three orders of magnitude
improvement in run time, for models with cheap likelihood evaluations.
This is shown in (https://arxiv.org/abs/2012.15286).Recently JAXNS has implemented Phantom-Powered Nested Sampling, which significantly reduces the number of required
likelihood evaluations for inferring the posterior. This is shown in (https://arxiv.org/abs/2312.11330).Note on performance with parallelisationNote, that this is an experimental feature.If you setnum_parallel_workers > 1you will usejax.pmapunder the hood for parallelisation.
This is a very powerful feature, but it is important to understand how it works.
It runs identical copies of the nested sampling algorithm on multiple devices.
There is a two-part stopping condition.
First, each copy goes until the user defined stopping condition is metper device.
Then, it performs an all-gather and finds at the highest likelihood contour among all copies, and continues all copies
hit this likelihood contour.
This ensures consistency of depth across all copies.
We then merge the copies and compute the final results.The algorithm is fairly memory bound, so running parallelisation over multiple CPUs on the same machine may not yield
the expected speed up, and depends on how expensive the likelihood evaluations are. Running over separate physical
devices
is the best way to achieve speed up.Change Log23 Feb, 2024 -- JAXNS 2.4.10 released. Hotfix for import error.21 Feb, 2024 -- JAXNS 2.4.9 released. Minor improvements to some priors, and bug fixes.31 Jan, 2024 -- JAXNS 2.4.8 released. Improved global optimisation performance using gradient slicing.
Improved evidence maximisation.25 Jan, 2024 -- JAXNS 2.4.6/7 released. Added logging. Use L-BFGS for Evidence Maximisation M-step. Fix bug in finetune.24 Jan, 2024 -- JAXNS 2.4.5 released. Gradient based finetuning global optimisation using L-BFGS. Added ability to
simulate prior models without bulding model (for data generation.)15 Jan, 2024 -- JAXNS 2.4.4 released. Fix performance issue for largermax_samples. Fixed bug in termination
conditions. Improved parallel performance.10 Jan, 2024 -- JAXNS 2.4.2/3 released. Another performance boost, and experimental global optimiser.9 Jan, 2024 -- JAXNS 2.4.1 released. Improve performance slightly for largermax_samples, still a performance issue.8 Jan, 2024 -- JAXNS 2.4.0 released. Python 3.9+ becomes supported. Migrate parametrised models to stable.
All models are now default able to be parametrised, so you can use hk.Parameter anywhere in the model.21 Dec, 2023 -- JAXNS 2.3.4 released. Correction for ESS and logZ uncert.parameter_estimationmode.20 Dec, 2023 -- JAXNS 2.3.2/3 released. Improved default parameters.difficult_modelmode. Improve plotting.18 Dec, 2023 -- JAXNS 2.3.1 released. Paper open science release. Default parameters from paper.11 Dec, 2023 -- JAXNS 2.3.0 released. Release of Phantom-Powered Nested Sampling algorithm.5 Oct, 2023 -- JAXNS 2.2.6 released. Minor update to evidence maximisation.3 Oct, 2023 -- JAXNS 2.2.5 released. Parametrised priors, and evidence maximisation added.24 Sept, 2023 -- JAXNS 2.2.4 released. Add marginalising from saved U samples.28 July, 2023 -- JAXNS 2.2.3 released. Bug fix for singular priors.26 June, 2023 -- JAXNS 2.2.1 released. Multi-ellipsoidal sampler added back in. Adaptive refinement disabled, as a bias
has been detected in it.15 June, 2023 -- JAXNS 2.2.0 released. Added support to allow TFP bijectors to defined transformed distributions. Other
minor improvements.15 April, 2023 -- JAXNS 2.1.0 released. pmap used on outer-most loops allowing efficient device-device communication
during parallel runs.8 March, 2023 -- JAXNS 2.0.1 released. Changed how we're doing annotations to support python 3.8 again.3 January, 2023 -- JAXNS 2.0 released. Complete overhaul of components. New way to build models.5 August, 2022 -- JAXNS 1.1.1 released. Pytree shaped priors.2 June, 2022 -- JAXNS 1.1.0 released. Dynamic sampling takes advantage of adaptive refinement. Parallelisation. Bayesian
opt and global opt modules.30 May, 2022 -- JAXNS 1.0.1 released. Improvements to speed, parallelisation, and structure of code.9 April, 2022 -- JAXNS 1.0.0 released. Parallel sampling, dynamic search, and adaptive refinement. Global optimiser
released.2 Jun, 2021 -- JAXNS 0.0.7 released.13 May, 2021 -- JAXNS 0.0.6 released.8 Mar, 2021 -- JAXNS 0.0.5 released.8 Mar, 2021 -- JAXNS 0.0.4 released.7 Mar, 2021 -- JAXNS 0.0.3 released.28 Feb, 2021 -- JAXNS 0.0.2 released.28 Feb, 2021 -- JAXNS 0.0.1 released.1 January, 2021 -- Paper submittedStar History