GenSBI Examples

This repository contains a collection of examples, tutorials, and recipes for GenSBI, a JAX-based library for Simulation-Based Inference using generative methods.

These examples demonstrate how to use GenSBI for various tasks, including:

• Defining and running inference pipelines.
• Using different embedding networks (MLP, ResNet, etc.).
• Handling various data types (1D signals, 2D images).

Installation

To run these examples you need both GenSBI and this examples helper package (gensbi-examples). The helper package is installed automatically through the examples extra of GenSBI.

Using uv (recommended)

uv add gensbi[examples]
# or, for a standalone install:
uv pip install gensbi[examples]

For GPU support (CUDA 12):

uv add gensbi[cuda12,examples]
# or
uv pip install gensbi[cuda12,examples]

Using pip

pip install gensbi[examples]
# with GPU support (CUDA 12):
pip install gensbi[cuda12,examples]

For more installation options, including how to install uv, see the Installation Guide.

Download the example notebooks and scripts

The gensbi-examples package provides the helper utilities, but the notebooks and training scripts live in this repository. To get them, clone the repo:

git clone https://github.com/aurelio-amerio/GenSBI-examples.git

Structure

• examples/: Contains standalone example scripts and notebooks.
• src/gensbi_examples: Helper utilities for the examples.

Getting Started

• My First Model:

Training Models

Models can be trained from the command line by passing a YAML configuration file to one of the provided training scripts. Each configuration file specifies the task, model architecture, training methodology (flow matching, diffusion, or score matching), optimizer settings, and training hyperparameters. Example configuration files can be found alongside each benchmark task under the config/ subdirectory (e.g., examples/sbi-benchmarks/two_moons/flow_flux/config/config_flow_flux.yaml).

General SBI tasks. For the standard SBI benchmark tasks (Two Moons, Gaussian Linear, SLCP, etc.), use the general-purpose training script:

python scripts/train_sbi_model.py --config <path_to_config.yaml>

This script trains the model, runs sampling, and computes diagnostic metrics (C2ST, TARP, SBC, LC2ST) automatically.

SBIBM benchmarks (budget/model/methodology scans). To reproduce the results reported in the main paper—including scans over training budget, model type, and methodology—use the dedicated SBIBM training script:

python scripts/train_sbi_model_sbibm.py --config <path_to_config.yaml>

This script additionally accepts a --dsize flag to control the training dataset size (default: 100,000), which is used for the budget scan experiments:

python scripts/train_sbi_model_sbibm.py --config <path_to_config.yaml> --dsize 10000

Advanced tasks. For more complex tasks such as gravitational waves and strong lensing, dedicated training scripts are provided in the corresponding example directories (e.g., examples/sbi-benchmarks/gravitational_waves/train-gw.py and examples/sbi-benchmarks/lensing/train-lensing.py).

Neural Density Estimators (NDE)

These examples demonstrate the usage of Neural Density Estimators for unconditional density estimation tasks.

• Diffusion EDM 2D Unconditional:
• Diffusion SM 2D Unconditional:
• Flow Matching 2D Unconditional:
• Flow Matching 2D Unconditional (Flux1Joint):

SBI Benchmark Examples

This repository includes a comprehensive suite of Simulation-Based Inference (SBI) benchmarks. These examples cover a range of standard tasks used to evaluate SBI methods, including simple distributions, physical systems, and complex toy problems. For each task, we provide implementations using various generative methods available in GenSBI, such as Flow Matching and Diffusion models with different architectures (Flux, SimFormer).

Two Moons

• Diffusion Flux:
• Diffusion Flux1Joint:
• Diffusion SimFormer:
• Flow Flux:
• Flow Flux1Joint:
• Flow SimFormer:

Bernoulli GLM

• Flow Flux:
• Flow Flux1Joint:

Gaussian Linear

• Flow Flux:
• Flow Flux1Joint:

Gaussian Mixture

• Flow Flux:
• Flow Flux1Joint:

Gravitational Waves

• GW Example:

Lensing

• Lensing Example:

SLCP

• Flow Flux:
• Flow Flux1Joint:
• Flow SimFormer: