https://github.com/aurelio-amerio/gensbi
A JAX machine learning library to connect generative models with scientific discovery
https://github.com/aurelio-amerio/gensbi
density-estimation diffusion-models flax flow-matching jax sbi
Last synced: 4 months ago
JSON representation
A JAX machine learning library to connect generative models with scientific discovery
- Host: GitHub
- URL: https://github.com/aurelio-amerio/gensbi
- Owner: aurelio-amerio
- License: other
- Created: 2025-04-01T16:14:51.000Z (about 1 year ago)
- Default Branch: main
- Last Pushed: 2026-02-24T10:38:11.000Z (4 months ago)
- Last Synced: 2026-02-24T15:08:53.221Z (4 months ago)
- Topics: density-estimation, diffusion-models, flax, flow-matching, jax, sbi
- Language: Python
- Homepage: https://gensbi.com
- Size: 142 MB
- Stars: 8
- Watchers: 1
- Forks: 0
- Open Issues: 13
-
Metadata Files:
- Readme: README.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
Awesome Lists containing this project
README
# GenSBI
[](https://github.com/aurelio-amerio/GenSBI/actions/workflows/python-app.yml)
[](https://pypi.org/project/gensbi/)
[](https://pepy.tech/project/gensbi)
> [!IMPORTANT]
> This library is at an early stage of development. The API is currently stable, but it is subject to change.
## Overview
**GenSBI** is a powerful JAX-based library for Simulation-Based Inference (SBI) using state-of-the-art generative models, currently revolving around Optimal Transport Flow Matching and Diffusion Models.
It is designed for researchers and practitioners who need a flexible, high-performance toolkit to solve complex inference problems where the likelihood function is intractable.
## Key Features
- **Modern SBI Algorithms**: Implements cutting-edge techniques like **Optimal Transport Conditional Flow Matching** and **Diffusion Models** for robust and flexible posterior inference.
- **Built on JAX and Flax NNX**: Leverages the power of JAX for automatic differentiation, vectorization, and seamless execution on CPUs, GPUs, and TPUs.
- **High-Level Recipes API**: A simplified interface for common workflows, allowing you to train models and run inference with just a few lines of code.
- **Powerful Transformer Models**: Includes implementations of recent, high-performing models like **Flux1**, **Flux1Join**, and **Simformer** for handling complex, high-dimensional data.
- **Modular and Extensible**: A clean, well-structured codebase that is easy to understand, modify, and extend for your own research.
## Installation
Using [uv](https://docs.astral.sh/uv/) (recommended):
```bash
uv add gensbi
# or, for a standalone install:
uv pip install gensbi
```
For GPU support:
```bash
uv add gensbi[cuda12]
# or
uv pip install gensbi[cuda12]
```
Or using pip:
```bash
pip install gensbi
```
For GPU support and other options, including how to install `uv`, see the [Installation Guide](https://aurelio-amerio.github.io/GenSBI/getting_started/installation.html).
## Quick Start
To get started immediately, you can use the high-level API to train a model.
> [!TIP]
> Check out the **[my_first_model.ipynb](https://github.com/aurelio-amerio/GenSBI-examples/blob/main/examples/getting_started/my_first_model.ipynb)** notebook for a complete, step-by-step introductory tutorial.
```python
from flax import nnx
from gensbi.recipes import Flux1FlowPipeline
from gensbi.models import Flux1Params
train_dataset = ... # define a training dataset (infinite iterator)
val_dataset = ... # define a validation dataset (infinite iterator)
dim_obs = ... # dimension of the parameters (theta)
dim_cond = ... # dimension of the simulator observations (x)
params = Flux1Params(...) # the parameters for your model
# Instantiate the pipeline
pipeline = Flux1FlowPipeline(
train_dataset,
val_dataset,
dim_obs,
dim_cond,
params=params,
)
# Train the model
# Note: GenSBI uses Flax NNX, so we pass a random key generator
pipeline.train(rngs=nnx.Rngs(0))
# After training, get a sampler for posterior sampling
key = jax.random.PRNGKey(42)
samples = pipeline.sample(key, x_observed, num_samples=10_000)
```
## Examples
Examples for this library are available separately in the [GenSBI-examples](https://github.com/aurelio-amerio/GenSBI-examples) repository.
Some key examples include:
**Getting Started:**
- `my_first_model.ipynb` [](https://colab.research.google.com/github/aurelio-amerio/GenSBI-examples/blob/main/examples/getting_started/my_first_model.ipynb)
A beginner-friendly notebook introducing the core concepts of GenSBI on a simple problem.
**Unconditional Density Estimation:**
- `flow_matching_2d_unconditional.ipynb` [](https://colab.research.google.com/github/aurelio-amerio/GenSBI-examples/blob/main/examples/NDE/flow_matching_2d_unconditional.ipynb)
Demonstrates how to use flow matching in 2D for unconditional density estimation.
- `diffusion_2d_unconditional.ipynb` [](https://colab.research.google.com/github/aurelio-amerio/GenSBI-examples/blob/main/examples/NDE/diffusion_2d_unconditional.ipynb)
Demonstrates how to use diffusion models in 2D for unconditional density estimation.
**Conditional Density Estimation:**
- `two_moons_flow_flux.ipynb` [](https://colab.research.google.com/github/aurelio-amerio/GenSBI-examples/blob/main/examples/sbi-benchmarks/two_moons/flow_flux/two_moons_flow_flux.ipynb)
Uses the Flux1 model for posterior density estimation on the two-moons benchmark.
- `two_moons_diffusion_flux.ipynb` [](https://colab.research.google.com/github/aurelio-amerio/GenSBI-examples/blob/main/examples/sbi-benchmarks/two_moons/diffusion_flux/two_moons_diffusion_flux.ipynb)
Uses the Diffusion model for posterior density estimation on the two-moons benchmark.
> [!NOTE]
> A complete list of the currently available examples can be found at the [examples](https://aurelio-amerio.github.io/GenSBI/examples.html) documentation page.
## Citing GenSBI
If you use this library, please consider citing this work and the original methodology papers, see [references](https://aurelio-amerio.github.io/GenSBI/references.html).
```bibtex
@misc{GenSBI,
author = {Amerio, Aurelio},
title = "{GenSBI: Generative models for Simulation-Based Inference}",
year = {2026},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/aurelio-amerio/GenSBI}}
}
```
### Reference implementations
- **Facebook Flow Matching library**: [https://github.com/facebookresearch/flow_matching](https://github.com/facebookresearch/flow_matching)
- **Elucidating the Design Space of Diffusion-Based Generative Models**: [https://github.com/NVlabs/edm](https://github.com/NVlabs/edm)
- **Simformer model**: [https://github.com/mackelab/simformer](https://github.com/mackelab/simformer)
- **Flux1 model from BlackForest Lab**: [https://github.com/black-forest-labs/flux](https://github.com/black-forest-labs/flux)
- **Simulation-Based Inference Benchmark**: [https://github.com/sbi-benchmark/sbibm](https://github.com/sbi-benchmark/sbibm)
> [!NOTE]
> **AI Usage Disclosure**
> This project utilized large language models, specifically Google Gemini and GitHub Copilot, to assist with code suggestions, documentation drafting, and grammar corrections. All AI-generated content has been manually reviewed and verified by human authors to ensure accuracy and adherence to scientific standards.