setu-sbi 0.1.0

JAX-native simulation-based inference with PyMC and NumPyro integration

A bridge from simulation-based inference to PyMC and NumPyro.

Many scientific models are simulators: they generate data from parameters, but do not provide the likelihood $p(x \mid \theta)$, the probability of an observation for a given set of parameters. That is exactly what Bayesian tools like PyMC and NumPyro need to recover the parameters behind observed data, so a simulator cannot be plugged into them directly.

setu (Sanskrit for bridge) closes that gap. It learns the likelihood from simulated data and exports it to a probabilistic programming language (PPL) in a single call, so the model, including priors, hierarchy, and posterior-predictive checks, is written in PyMC or NumPyro as usual. setu provides NLE, NRE, and MixedNLE estimators; training is backend-agnostic.

sbi implements neural likelihood estimation too; setu is a JAX-native complement to it, not a replacement. Because the estimator, its gradient, and the sampler compile into one program, setu targets JAX-native samplers such as BlackJAX and NumPyro and tends to be faster on the benchmarked task (see the benchmarks).

Installation

pip install setu-sbi                     # core (JAX + FlowJAX)
pip install "setu-sbi[pymc]"             # + PyMC / BlackJAX
pip install "setu-sbi[numpyro]"          # + NumPyro
pip install "setu-sbi[tuning]"           # + Optuna hyperparameter tuning
pip install "setu-sbi[all]"              # everything above

The distribution is named setu-sbi on PyPI; the import name is setu.

Documentation

Full documentation lives at https://aai-institute.github.io/setu/:

• Tutorials: hand-held walkthroughs, from a simulator to a posterior in one sitting.
• How-to guides: task-oriented recipes, starting with installation.
• Reference: the full API, the (G, S, T, *E) shape conventions, and a glossary.
• Explanation: what SBI is, why validation matters, and why hierarchical inference is powerful.

New here? Start with the tutorials.

Contributing

Contributions and questions are welcome, from bug reports to documentation fixes to code. See CONTRIBUTING.md for setup, workflow, and how to file an issue. New contributors: CONTEXT.md is the project glossary, so a quick skim is the fastest way to get oriented.

uv sync --dev
uv run pytest tests/ -m "not slow"  # fast tests (~2 min)
uv run ruff check setu/ tests/