https://github.com/homerjed/sbiax

Fast, lightweight and parallelised simulation-based inference in JAX.
https://github.com/homerjed/sbiax

generative-models jax normalizing-flows simulation-based-inference

Last synced: about 22 hours ago
JSON representation

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

• Host: GitHub
• URL: https://github.com/homerjed/sbiax
• Owner: homerjed
• License: mit
• Created: 2024-02-08T15:00:43.000Z (12 months ago)
• Default Branch: main
• Last Pushed: 2024-10-24T11:07:01.000Z (3 months ago)
• Last Synced: 2024-10-25T10:05:22.046Z (3 months ago)
• Topics: generative-models, jax, normalizing-flows, simulation-based-inference
• Language: Python
• Homepage:
• Size: 4.07 MB
• Stars: 0
• Watchers: 1
• Forks: 0
• Open Issues: 0
• Metadata Files:
  • Readme: README.md
  • License: LICENSE

Awesome Lists containing this project

• awesome-neural-sbi - [Code - estimation Simulation-based inference in JAX. (Code Packages and Benchmarks)

README

        
sbiax

Fast, lightweight and parallel simulation-based inference.


[![DOI](https://joss.theoj.org/papers/10.21105/joss.07606/status.svg)](https://doi.org/10.21105/joss.07606) [![Project Status: Active – The project has reached a stable, usable state and is being actively developed.](https://www.repostatus.org/badges/latest/active.svg)](https://www.repostatus.org/#active) [![arXiv](https://img.shields.io/badge/arXiv-2412.02311-b31b1b.svg)](https://arxiv.org/abs/2412.02311)



  

    

    

    

  



`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](https://github.com/homerjed/sbiax/tree/main/examples).

-----

### Contributing

Want to add something? See `CONTRIBUTING.md`.

-----

### Citation

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

```bibtex

@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}, 

}

```

```bibtex

@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} 

} 

```