👋 Welcome to sbijax!

👋 Welcome to sbijax!#

Simulation-based inference in JAX

Sbijax is a Python library for neural simulation-based inference and approximate Bayesian computation using JAX. It implements recent methods, such as Sequential Monte Carlo ABC, Surjective Neural Likelihood Estimation, Neural Approximate Sufficient Statistics or Consistency model posterior estimation, as well as methods to compute model diagnostics and for visualizing posterior distributions.

Caution

⚠️ As per the LICENSE file, there is no warranty whatsoever for this free software tool. If you discover bugs, please report them.

Example#

Sbijax implements a slim object-oriented API with functional elements stemming from JAX. All a user needs to define is a prior model, a simulator function and an inferential algorithm. For example, you can define a neural likelihood estimation method and generate posterior samples like this:

from jax import numpy as jnp, random as jr
from sbijax import NLE
from sbijax.nn import make_maf
from tensorflow_probability.substrates.jax import distributions as tfd

def prior_fn():
    prior = tfd.JointDistributionNamed(dict(
        theta=tfd.Normal(jnp.zeros(2), jnp.ones(2))
    ), batch_ndims=0)
    return prior

def simulator_fn(seed, theta):
    p = tfd.Normal(jnp.zeros_like(theta["theta"]), 0.1)
    y = theta["theta"] + p.sample(seed=seed)
    return y


fns = prior_fn, simulator_fn
model = NLE(fns, make_maf(2))

y_observed = jnp.array([-1.0, 1.0])
data, _ = model.simulate_data(jr.PRNGKey(1))
params, _ = model.fit(jr.PRNGKey(2), data=data)
posterior, _ = model.sample_posterior(jr.PRNGKey(3), params, y_observed)

Installation#

You can install sbijax from PyPI using:

pip install sbijax

To install the latest GitHub <RELEASE>, just call the following on the command line:

pip install git+https://github.com/dirmeier/sbijax@<RELEASE>

See also the installation instructions for JAX, if you plan to use sbijax on GPU/TPU.

Contributing#

Contributions in the form of pull requests are more than welcome. A good way to start is to check out issues labelled “good first issue”.

In order to contribute:

  1. Clone sbijax and install uv from here,

  2. install all dependencies using `uv sync,

  3. create a new branch locally git checkout -b feature/my-new-feature or git checkout -b issue/fixes-bug,

  4. implement your contribution and ideally a test case,

  5. test it by calling make tests, make lints and make format on the (Unix) command line,

  6. submit a PR 🙂

Acknowledgements#

Note

📝 The API of the package is heavily inspired by the excellent Pytorch-based sbi package.

License#

sbijax is licensed under the Apache 2.0 License.