sbiax 0.0.57

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.

---

Contributing

Want to add something? See CONTRIBUTING.md.

---

Citation

If you found this library to be useful in academic work, please cite:

@misc{homer2024simulationbasedinferencedodelsonschneidereffect,
      title={Simulation-based inference has its own Dodelson-Schneider effect (but it knows that it does)}, 
      author={Jed Homer and Oliver Friedrich and Daniel Gruen},
      year={2024},
      eprint={2412.02311},
      archivePrefix={arXiv},
      primaryClass={astro-ph.CO},
      url={https://arxiv.org/abs/2412.02311}, 
}

@article{
    Homer2025, 
    doi = {10.21105/joss.07606}, 
    url = {https://doi.org/10.21105/joss.07606}, 
    year = {2025}, 
    publisher = {The Open Journal}, 
    volume = {10}, 
    number = {105}, 
    pages = {7606}, 
    author = {Jed Homer and Oliver Friedrich}, 
    title = {SBIAX: Density-estimation simulation-based inference in JAX}, 
    journal = {Journal of Open Source Software} 
}