gensbi-examples 0.4.0

Examples for the GenSBI library

GenSBI Examples

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

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

Prerequisites

You need to have GenSBI and the examples package installed.

With CUDA 12 support (Recommended):

pip install gensbi[cuda12, examples]

CPU-only:

pip install gensbi

Install Examples Package

To download these examples, clone the github repository:

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 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 models 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: