sbiax 0.0.12

Fast, parallel and lightweight simulation-based inference in JAX.

sbiax

Fast, lightweight and parallel simulation-based inference.

sbiax is a lightweight library for simulation-based inference (SBI) with a fixed grid of simulations.

[!WARNING] :building_construction: Note this repository is under construction, expect changes. :building_construction:

---

Design

In a typical inference problem the data likelihood is unknown. Using density-estimation SBI, we can proceed by

• simulating a set of data and model parameters ${(\boldsymbol{\xi}, \boldsymbol{\pi})_0, ..., (\boldsymbol{\xi}, \boldsymbol{\pi})_N}$,
• obtaining a measurement $\hat{\boldsymbol{\xi}}$,
• compressing the simulations and the measurements - usually with a neural network or linear compression - to a set of summaries ${(\boldsymbol{x}, \boldsymbol{\pi})_0, ..., (\boldsymbol{x}, \boldsymbol{\pi})_N}$ and $\hat{\boldsymbol{x}}$,
• fitting an ensemble of normalising flow or similar density estimation algorithms (e.g. a Gaussian mixture model),
• the optional optimisation of the parameters for the architecture and fitting hyperparameters of the algorithms,
• sampling the ensemble posterior (using an MCMC sampler if the likelihood is fit directly) conditioned on the datavector to obtain parameter constraints on the parameters of a physical model, $\boldsymbol{\pi}$.

sbiax is a code for implementing each of these steps.

---

Usage

Install via

pip install sbiax

and have a look at examples.