Skip to content
aai-institutePublic

Repository files navigation

setu

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

PyPI Python versions Documentation License

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.

When to use setu

setu is built for one situation: an intractable likelihood combined with repeated or hierarchical structure in the data (trials, subjects, groups). It is a good fit when:

  • a hierarchical or trial-based model has a simulator but no tractable likelihood, so it cannot enter a PPL directly;
  • an existing PyMC or NumPyro model needs a more realistic simulator component that makes its likelihood intractable;
  • an SBI analysis is held back by hierarchical structure or slow sampling, where gradient-based NUTS inside a PPL would help.

Other tools fit better in the neighboring cases:

  • with a tractable likelihood, model it directly in a PPL such as PyMC, NumPyro, or Stan;
  • with high-dimensional observations and no repeated or hierarchical structure, use a general SBI toolkit such as sbi or BayesFlow.

setu complements both rather than replacing them. It learns a likelihood as an SBI method does, then hands it to a PPL so priors, hierarchy, and fast JAX samplers (BlackJAX, NumPyro) all apply. Because the estimator, its gradient, and the sampler compile into one program, it tends to be faster on the benchmarked task.

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/

Releases

Packages

Contributors

Languages