Metadata-Version: 2.4
Name: sbiax
Version: 0.0.41
Summary: Fast, parallel and lightweight simulation-based inference in JAX.
Project-URL: repository, https://github.com/homerjed/sbiax
Author-email: Jed Homer <jedhmr@gmail.com>
License: MIT License
        
        Copyright (c) [2024] [Jed Homer]
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
License-File: LICENSE
Keywords: bayesian-inference,deep-learning,generative-models,jax,likelihood-free-inference,simulation-based-inference
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Natural Language :: English
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: beartype==0.19.0
Requires-Dist: blackjax==1.2.4
Requires-Dist: chainconsumer==1.1.2
Requires-Dist: cloudpickle==3.1.0
Requires-Dist: diffrax==0.6.0
Requires-Dist: einops==0.8.0
Requires-Dist: equinox==0.11.8
Requires-Dist: flowjax==16.0.0
Requires-Dist: jax==0.4.33
Requires-Dist: matplotlib==3.9.2
Requires-Dist: ml-collections==0.1.1
Requires-Dist: numpy==1.26.4
Requires-Dist: optax==0.2.3
Requires-Dist: optuna==4.0.0
Requires-Dist: plotly==5.24.1
Requires-Dist: tensorflow-probability==0.24.0
Requires-Dist: tqdm==4.66.5
Description-Content-Type: text/markdown

<h1 align='center'>sbiax</h1>
<h2 align='center'>Fast, lightweight and parallel simulation-based inference.</h2>

<!-- 
<picture>
  <source srcset="https://github.com/homerjed/sbipdf/blob/main/assets/cover_dark.png" media="(prefers-color-scheme: dark)">
  <source srcset="https://github.com/homerjed/sbipdf/blob/main/assets/cover.png" media="(prefers-color-scheme: light)">
  <img src="https://github.com/homerjed/sbipdf/blob/main/assets/cover.png" alt="Your image description">
</picture> -->
<p align="center">
  <picture>
    <source srcset="https://github.com/homerjed/sbiax/blob/main/assets/cover_dark.png" media="(prefers-color-scheme: dark)">
    <source srcset="https://github.com/homerjed/sbiax/blob/main/assets/cover.png" media="(prefers-color-scheme: light)">
    <img src="https://github.com/homerjed/sbiax/blob/main/assets/cover.png" alt="Your image description">
  </picture>
</p>

<!-- <p align="center">
    <picture>
        <source srcset="assets/cover_dark.png" media="(prefers-color-scheme: dark)">
        <source srcset="assets/cover.png" media="(prefers-color-scheme: light)">
        <img src="assets/cover.png" alt="Your image description">
    </picture>
</p> -->

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

<!-- The design puts the neural density estimator (NDE) models at the centre of the code, allowing for flexible combinations of different models.  -->

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

-----

### Design

<!-- A typical inference with SBI occurs with  

* fitting a density estimator to a set of simulations and parameters $(\xi, \pi)$ that may be compressed to summary statistics,
* the measurement of a datavector $\hat{\xi}$,
* the sampling of a posterior $p(\pi|\hat{\xi})$ conditioned on the measurement $\hat{\xi}$.

`sbiax` is designed to perform such an inference.  -->

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

<!-- Bayesian analyses where the likelihood function is unknown can proceed with density-estimation simulation-based inference methods, which typically involve -->

* 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.


<!-- #### a) Configuration

An inference is defined by a `config` file. This is a dictionary that includes

* the architecture(s) of the NDEs,
* how to train these models,
* how to sample these models (e.g. MCMC, ...),
* where to save models, posteriors and figures,
* and generally any other information for your experiments.

NDEs are grouped in an ensemble that defines its own ensemble-likelihood function given an observation.

#### b) Density estimation

A posterior or likelihood is derived from a set of simulations and parameters by fitting a generative model with some loss - this may be a diffusion model or a normalising flow. 

`sbiax` is designed to be centred around these algorithms and to adopt the latest innovations from the machine learning literature.

#### c) Compression 

Density estimation is one of the oldest problems in machine learning. To avoid the difficulties of fitting high-dimensional models to data it is common to compress the data. 

`sbiax` gives you common compression methods that use linear methods or neural networks.  -->

-----

### Usage

Install via

```pip install sbiax```

and have a look at [examples](https://github.com/homerjed/sbiax/tree/main/examples).

### Citation

If you found this library to be useful in academic work, please cite:  <!--([arXiv link](https://arxiv.org/abs/2111.00254)) -->

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

```