Add NPSE + AiO and add more unit tests (#55)

Author: dirmeier

.github/workflows/ci.yaml

@@ -73,7 +73,7 @@ jobs:
           pip install hatch
       - name: Build package
         run: |
-          pip install jaxlib==0.4.24 jax==0.4.24
+          pip install jaxlib jax
       - name: Run tests
         run: |
           hatch run test:test

.github/workflows/examples.yaml

@@ -35,12 +35,14 @@ jobs:
           pip install .
       - name: Run tests
         run: |
-            python examples/bivariate_gaussian-smcabc.py --n-rounds 1
+            python examples/gaussian_linear-aio.py --n-iter 10
+            python examples/gaussian_linear-smcabc.py --n-rounds 1
             python examples/mixture_model-cmpe.py --n-iter 10
             python examples/mixture_model-nle.py --n-iter 10
             python examples/mixture_model-nle.py --n-iter 10 --use-spf
             python examples/mixture_model-npe.py --n-iter 10
             python examples/mixture_model-nre.py --n-iter 10
+            python examples/mixture_model-npse.py --n-iter 10
             python examples/slcp-fmpe.py --n-iter 10
             python examples/slcp-nass_nle.py --n-iter 10 --n-rounds 1
             python examples/slcp-nass_smcabc.py --n-iter 10 --n-rounds 1

.pre-commit-config.yaml

@@ -15,7 +15,7 @@ repos:
     - id: trailing-whitespace
 
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v0.910-1
+  rev: v1.14.1
   hooks:
   - id: mypy
     args: ["--ignore-missing-imports"]

README.md

@@ -89,6 +89,19 @@ In order to contribute:
 4) test it by calling `make tests`, `make lints` and `make format` on the (Unix) command line,
 5) submit a PR 🙂
 
+## Citing sbijax
+
+If you find our work relevant to your research, please consider citing:
+
+```
+@article{dirmeier2024simulation,
+  title={Simulation-based inference with the Python Package sbijax},
+  author={Dirmeier, Simon and Ulzega, Simone and Mira, Antonietta and Albert, Carlo},
+  journal={arXiv preprint arXiv:2409.19435},
+  year={2024}
+}
+```
+
 ## Acknowledgements
 
 > [!NOTE]

docs/index.rst

@@ -130,6 +130,7 @@ License
     :hidden:
 
     sbijax
+    sbijax.experimental
     sbijax.mcmc
     sbijax.nn
     sbijax.util

docs/references.bib

@@ -36,23 +36,20 @@ @article{dirmeier2023simulation
 @inproceedings{papama2019neural,
   title         = {Sequential Neural Likelihood: Fast Likelihood-free Inference with Autoregressive Flows},
   author        = {Papamakarios, George and Sterratt, David and Murray, Iain},
-  booktitle     = {Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics},
-  year          = {2019},
-  doi           = {10.48550/arXiv.1805.07226}
+  booktitle     = {International Conference on Artificial Intelligence and Statistics},
+  year          = {2019}
 }
 @inproceedings{greenberg2019automatic,
   title         = {Automatic posterior transformation for likelihood-free inference},
   author        = {Greenberg, David and Nonnenmacher, Marcel and Macke, Jakob},
   booktitle     = {International Conference on Machine Learning},
-  year          = {2019},
-  doi           = {10.48550/arXiv.1905.0748}
+  year          = {2019}
 }
 @inproceedings{miller2022contrast,
   author        = {Miller, Benjamin K and Weniger, Christoph and Forr\'{e}, Patrick},
   booktitle     = {Advances in Neural Information Processing Systems},
   title         = {Contrastive Neural Ratio Estimation},
-  year          = {2022},
-  doi           = {10.48550/arXiv.2210.06170}
+  year          = {2022}
 }
 @article{beaumont2009adaptive,
   title={Adaptive approximate {B}ayesian computation},
@@ -62,8 +59,7 @@ @article{beaumont2009adaptive
   number={4},
   pages={983--990},
   year={2009},
-  publisher={Oxford University Press},
-  doi={10.1093/biomet/asp052}
+  publisher={Oxford University Press}
 }
 @article{albert2015simulated,
   title={A simulated annealing approach to approximate {B}ayes computations},
@@ -74,3 +70,15 @@ @article{albert2015simulated
   year={2015},
   publisher={Springer}
 }
+@inproceedings{sharrock2024sequential,
+  title={Sequential Neural Score Estimation: Likelihood-Free Inference with Conditional Score Based Diffusion Models},
+  author={Louis Sharrock and Jack Simons and Song Liu and Mark Beaumont},
+  booktitle={International Conference on Machine Learning},
+  year={2024}
+}
+@inproceedings{gloeckler2024allinone,
+  title={All-in-one simulation-based inference},
+  author={Manuel Gloeckler and Michael Deistler and Christian Dietrich Weilbach and Frank Wood and Jakob H. Macke},
+  booktitle={International Conference on Machine Learning},
+  year={2024},
+}

docs/sbijax.experimental.rst

@@ -0,0 +1,32 @@
+``sbijax.experimental``
+=======================
+
+.. currentmodule:: sbijax.experimental
+
+``sbijax.experimental`` contains experimental code that might get ported to the
+main code base or possibly deleted again.
+
+.. autosummary::
+    AiO
+    NPSE
+
+..  autoclass:: AiO
+    :members: fit, simulate_data, simulate_data_and_possibly_append, sample_posterior
+
+..  autoclass:: NPSE
+    :members: fit, simulate_data, simulate_data_and_possibly_append, sample_posterior
+
+
+.. 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__

examples/gaussian_linear-aio.py

@@ -0,0 +1,51 @@
+"""All-in-one simulation-based inference.
+
+Demonstrates AiO on a linear Gaussian model.
+"""
+import matplotlib.pyplot as plt
+import numpy as np
+from jax import numpy as jnp, random as jr
+from tensorflow_probability.substrates.jax import distributions as tfd
+
+from sbijax import plot_posterior
+from sbijax.experimental import AiO
+from sbijax.experimental.nn import make_simformer_based_score_model
+
+
+def prior_fn():
+    prior = tfd.JointDistributionNamed(dict(
+        theta=tfd.Normal(jnp.zeros(5), 1)
+    ), batch_ndims=0)
+    return prior
+
+
+def simulator_fn(seed, theta):
+    mean = theta["theta"].reshape(-1, 5)
+    y = tfd.Normal(mean, 0.1).sample(seed=seed)
+    return y
+
+
+def run(n_iter):
+    y_observed = jnp.linspace(-2.0, 2.0, 5)
+    fns = prior_fn, simulator_fn
+    mask = jnp.zeros((10, 10))
+    mask = mask.at[np.arange(5, 10), np.arange(5)].set(1)
+    mask = mask + mask.T + jnp.eye(10)
+
+    neural_network = make_simformer_based_score_model(5, mask, 1, 1)
+    model = AiO(fns, neural_network)
+
+    data, _ = model.simulate_data(jr.PRNGKey(1), n_simulations=10_000)
+    params, info = model.fit(jr.PRNGKey(2), data=data, n_early_stopping_patience=25, n_iter=n_iter)
+    inference_result, _ = model.sample_posterior(jr.PRNGKey(3), params, y_observed)
+
+    plot_posterior(inference_result, point_estimate="mean")
+    plt.show()
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n-iter", type=int, default=1_000)
+    args = parser.parse_args()
+    run(args.n_iter)

examples/bivariate_gaussian-smcabc.py examples/gaussian_linear-smcabc.pyexamples/bivariate_gaussian-smcabc.py renamed to examples/gaussian_linear-smcabc.py

@@ -0,0 +1,50 @@
+"""Flow matching posterior estimation example.
+
+Demonstrates FPME on a simple mixture model.
+"""
+import matplotlib.pyplot as plt
+from jax import numpy as jnp, random as jr
+from tensorflow_probability.substrates.jax import distributions as tfd
+
+from sbijax import plot_posterior
+from sbijax import FMPE
+from sbijax.nn import make_cnf
+
+
+def prior_fn():
+    prior = tfd.JointDistributionNamed(dict(
+        theta=tfd.Normal(jnp.zeros(2), 1)
+    ), batch_ndims=0)
+    return prior
+
+
+def simulator_fn(seed, theta):
+    mean = theta["theta"].reshape(-1, 2)
+    n = mean.shape[0]
+    data_key, cat_key = jr.split(seed)
+    categories = tfd.Categorical(logits=jnp.zeros(2)).sample(seed=cat_key, sample_shape=(n,))
+    scales = jnp.array([1.0, 0.1])[categories].reshape(-1, 1)
+    y = tfd.Normal(mean, scales).sample(seed=data_key)
+    return y
+
+
+def run(n_iter):
+    y_observed = jnp.array([-2.0, 2.0])
+    fns = prior_fn, simulator_fn
+    neural_network = make_cnf(2)
+    model = FMPE(fns, neural_network)
+
+    data, _ = model.simulate_data(jr.PRNGKey(1), n_simulations=20_000)
+    params, info = model.fit(jr.PRNGKey(2), data=data, n_early_stopping_patience=25, n_iter=n_iter)
+    inference_result, _ = model.sample_posterior(jr.PRNGKey(3), params, y_observed)
+
+    plot_posterior(inference_result)
+    plt.show()
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--n-iter", type=int, default=1_000)
+    args = parser.parse_args()
+    run(args.n_iter)