Architecture redesign

.gitignore

@@ -1,3 +1,6 @@
+# git worktrees
+.worktrees/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

README.md

@@ -1,64 +1,59 @@
 # sbijax <img src="https://raw.githubusercontent.com/dirmeier/sbijax/main/docs/_static/sticker.png" align="right" width="160px"/>
 
-[![active](https://www.repostatus.org/badges/latest/active.svg)](https://www.repostatus.org/#active)
 [![ci](https://github.com/dirmeier/sbijax/actions/workflows/ci.yaml/badge.svg)](https://github.com/dirmeier/sbijax/actions/workflows/ci.yaml)
 [![codecov](https://codecov.io/gh/dirmeier/sbijax/branch/main/graph/badge.svg?token=dn1xNBSalZ)](https://codecov.io/gh/dirmeier/sbijax)
 [![documentation](https://readthedocs.org/projects/sbijax/badge/?version=latest)](https://sbijax.readthedocs.io/en/latest/?badge=latest)
 [![version](https://img.shields.io/pypi/v/sbijax.svg?colorB=black&style=flat)](https://pypi.org/project/sbijax/)
 
 > Simulation-based inference in JAX
 
-## About
-
 ``Sbijax`` is a Python library for neural simulation-based inference and
 approximate Bayesian computation using [JAX](https://github.com/google/jax).
-It implements recent methods, such as *Simulated-annealing ABC*,
+It implements recent methods, such as *Simulated Annealing ABC*,
 *Surjective Neural Likelihood Estimation*, *Neural Approximate Sufficient Statistics*
-or *Consistency model posterior estimation*, as well as methods to compute model
-diagnostics and for visualizing posterior distributions.
+or *Neural Posterior Score Estimation*.
 
 > [!CAUTION]
 > ⚠️ As per the LICENSE file, there is no warranty whatsoever for this free software tool. If you discover bugs, please report them.
 
-## Examples
+## Quick start
 
-`Sbijax` implements a slim object-oriented API with functional elements stemming from
-JAX. All a user needs to define is a prior model, a simulator function and an inferential algorithm.
-For example, you can define a neural likelihood estimation method and generate posterior samples like this:
+`Sbijax` implements a fully functional API in the idiom of [Haiku](https://github.com/google-deepmind/dm-haiku):
+every method is a factory returning a record of pure functions, with parameters
+threaded explicitly. All a user needs to define is a prior, a simulator function
+and an inferential algorithm. For example, you can define a neural likelihood
+estimation method and generate posterior samples like this:
 
 ```python
 from jax import numpy as jnp, random as jr
-from sbijax import NLE
-from sbijax.nn import make_maf
 from tensorflow_probability.substrates.jax import distributions as tfd
 
-def prior_fn():
-    prior = tfd.JointDistributionNamed(dict(
-        theta=tfd.Normal(jnp.zeros(2), jnp.ones(2))
-    ), batch_ndims=0)
-    return prior
+from sbijax import nle, train, sample, simulate
+from sbijax.mcmc import make_sampler, nuts
+from sbijax.nn import make_maf
+
+prior = tfd.JointDistributionNamed(dict(
+    theta=tfd.Normal(jnp.zeros(2), jnp.ones(2))
+), batch_ndims=0)
 
 def simulator_fn(seed, theta):
     p = tfd.Normal(jnp.zeros_like(theta["theta"]), 0.1)
     y = theta["theta"] + p.sample(seed=seed)
     return y
 
-
-fns = prior_fn, simulator_fn
-model = NLE(fns, make_maf(2))
+estimator = nle(make_maf(2))
 
 y_observed = jnp.array([-1.0, 1.0])
-data, _ = model.simulate_data(jr.PRNGKey(1))
-params, _ = model.fit(jr.PRNGKey(2), data=data)
-posterior, _ = model.sample_posterior(jr.PRNGKey(3), params, y_observed)
+data = simulate(jr.key(1), prior, simulator_fn, n=10_000)
+params, info = train(jr.key(2), estimator, data)
+samples, _ = sample(
+    jr.key(3), estimator, params, y_observed,
+    sampler=make_sampler(nuts, prior=prior),
+)
 ```
 
 More self-contained examples can be found in [examples](https://github.com/dirmeier/sbijax/tree/main/examples).
 
-## Documentation
-
-Documentation can be found [here](https://sbijax.readthedocs.io/en/latest/).
-
 ## Installation
 
 Make sure to have a working `JAX` installation. Depending whether you want to use CPU/GPU/TPU,
@@ -76,36 +71,9 @@ To install the latest GitHub <RELEASE>, use:
 pip install git+https://github.com/dirmeier/sbijax@<RELEASE>
 ```
 
-## Contributing
-
-Contributions in the form of pull requests are more than welcome. A good way to start is to check out issues labelled
-[good first issue](https://github.com/dirmeier/sbijax/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22).
-
-In order to contribute:
-
-1) Clone `sbijax` and install `uv` from [here](https://docs.astral.sh/uv/getting-started/installation/).
-2) Install all dependencies using `uv sync --all-groups`.
-3) Install `pre-commit` and `gitlint` via:
-
-   ```shell
-   pre-commit install
-   gitlint install-hook
-   ```
-4) Create a new branch locally `git checkout -b feature/my-new-feature` or `git checkout -b issue/fixes-bug`.
-5) Implement your contribution and ideally a test case.
-6) Test, lint and format your contribution by running:
-
-   ```shell
-   uv run pytest                         # run the test suite
-   uv run ruff check sbijax examples     # lint
-   uv run ruff format sbijax examples    # format
-   uv run mypy sbijax                    # type-check
-   ```
+## Documentation
 
-   The `pre-commit` hook installed in step 3 runs `ruff` and `mypy` on every
-   commit, so these checks also run automatically. To build the docs locally,
-   run `make html` from within the `docs` directory.
-7) Submit a PR 🙂.
+Documentation can be found [here](https://sbijax.readthedocs.io/en/latest/).
 
 ## Citing sbijax
 
@@ -123,4 +91,4 @@ If you find our work relevant to your research, please consider citing:
 ## Acknowledgements
 
 > [!NOTE]
-> 📝 The API of the package is heavily inspired by the excellent Pytorch-based [`sbi`](https://github.com/sbi-dev/sbi) package.
+> 📝 The API of the package is heavily inspired by [`Haiku`](https://github.com/google-deepmind/dm-haiku).

docs/_static/theme.css

@@ -1,26 +1,15 @@
 html[data-theme="light"] {
-    --pst-color-primary: rgb(121, 40, 161);
+    /* --pst-color-primary: rgb(121, 40, 161); */
+    /* --pst-color-primary:#b26679;
     --pst-color-primary-bg: #ffe9dd;
     --pst-color-secondary: #b26679;
-    --pst-color-inline-code-links: #b26679;
+    --pst-color-inline-code-links: #b26679; */
 }
 
 pre > span {
     line-height: 20px;
 }
 
-span.kn {
-     color: rgb(0, 120, 161) !important;
-}
-
-span.ml, span.mi, span.nb {
-     color: lightcoral !important;
-}
-
-span.k, span.nn {
-     color: rgb(168, 70, 185) !important;
-}
-
 h1 > code > span {
     font-weight: 300 !important;
 }
@@ -34,9 +23,6 @@ pre {
 h1 {
     margin-bottom: 50px;
 }
-h3, h2, h1 {
-
-}
 
 nav > li > a > code.literal {
     padding-top: 0;
@@ -49,18 +35,18 @@ nav.bd-links p.caption {
     text-transform: uppercase;
 }
 
-code.literal {
+/* code.literal {
   background-color: white;
   border: 0;
   border-radius: 0;
-}
+} */
 
 a > code {
     font-weight: 575;
 }
 
 a:hover {
-      text-decoration-thickness: 1px !important;
+    text-decoration-thickness: 1px !important;
 }
 
 ul.bd-breadcrumbs li.breadcrumb-item a:hover {

docs/api/index.rst

@@ -0,0 +1,12 @@
+API Reference
+=============
+
+..  toctree::
+    :maxdepth: 4
+
+    sbijax
+    sbijax.experimental
+    sbijax.mcmc
+    sbijax.nn
+    sbijax.simulators
+    sbijax.util

docs/api/sbijax.experimental.rst

@@ -0,0 +1,38 @@
+sbijax.experimental
+===================
+
+.. currentmodule:: sbijax.experimental
+
+``sbijax.experimental`` contains experimental code that might get ported to the
+main code base or possibly deleted again.
+
+``cmpe`` (consistency-model posterior estimation) and ``aio`` are functional
+factories; ``aio`` delegates to the ``fmpe`` core, and
+``make_truncated_proposal`` builds the truncated-prior proposal used with
+:func:`sbijax.run_sequential`. The score networks below are consumed by
+:func:`sbijax.npse`, which now lives in the main package.
+
+.. autosummary::
+    cmpe
+    aio
+    make_truncated_proposal
+
+.. autofunction:: cmpe
+
+.. autofunction:: aio
+
+.. autofunction:: make_truncated_proposal
+
+.. currentmodule:: sbijax.experimental.nn
+
+.. autosummary::
+    make_score_model
+    make_simformer_based_score_model
+    ScoreModel
+
+.. autofunction:: make_simformer_based_score_model
+
+.. autofunction:: make_score_model
+
+..  autoclass:: ScoreModel
+    :members: __call__

docs/api/sbijax.mcmc.rst

@@ -0,0 +1,53 @@
+sbijax.mcmc
+===========
+
+.. currentmodule:: sbijax.mcmc
+
+``sbijax.mcmc`` builds the posterior samplers
+and exposes the low-level MCMC routines they are built on.
+
+:func:`make_sampler` bundles a :class:`Kernel` -- a handle identifying a
+BlackJAX MCMC algorithm -- with the prior and ``N(0, I)`` chain initialisation
+into a sampler that is passed to :func:`sbijax.sample`. The available algorithms
+are ``nuts``, ``mala``, ``rmh`` and ``imh``::
+
+    from sbijax.mcmc import make_sampler, nuts
+
+    sampler = make_sampler(nuts, prior=prior)
+    samples, info = sample(key, estimator, params, y_obs, sampler=sampler)
+
+.. autosummary::
+    make_sampler
+    imh
+    mala
+    nuts
+    rmh
+    sample_with_imh
+    sample_with_mala
+    sample_with_nuts
+    sample_with_rmh
+    sample_with_slice
+
+.. autofunction:: make_sampler
+
+.. autofunction:: sample_with_imh
+
+.. autofunction:: sample_with_mala
+
+.. autofunction:: sample_with_nuts
+
+.. autofunction:: sample_with_rmh
+
+.. autofunction:: sample_with_slice
+
+.. autodata:: imh
+    :no-value:
+
+.. autodata:: mala
+    :no-value:
+
+.. autodata:: nuts
+    :no-value:
+
+.. autodata:: rmh
+    :no-value:

docs/sbijax.nn.rst → docs/api/sbijax.nn.rst

@@ -1,5 +1,5 @@
-``sbijax.nn``
-=============
+sbijax.nn
+=========
 
 .. currentmodule:: sbijax.nn