Sampler plugin for using numpyro samplers with Bilby
bilby-numpyro is a sampler plugin for using numpyro samplers with Bilby likelihoods and priors.
Any likelihood that uses only JAX JIT-compilable functions/classes works automatically.
The sampler should always be given as numpyro regardless of which sampler (e.g., HMC, NUTS) is actually being used.
The sampler is specified with sampler_name and any sampler keyword arguments should be passed in the usual way to
bilby.core.sampler.run_sampler for example
run_sampler(
...,
sampler="numpyro",
sampler_name="NUTS",
num_warmup=500,
num_samples=1000,
)This has not been tested with numpyro parallelisation.
The likelihood class you use should be JIT compilable with JAX.
Generally, this will be the case if you use jax.numpy and jax.scipy instead of numpy/scipy.
Some Bilby priors have direct equivalents in numpyro and so we use those automatically.
You will see logged warnings if a prior is not available and a Uniform prior over the same range will be used instead.