Improve documentation (#26)

README.md

@@ -12,16 +12,20 @@
 [JAX](https://github.com/google/jax) using [Haiku](https://github.com/deepmind/dm-haiku),
 [Distrax](https://github.com/deepmind/distrax) and [BlackJAX](https://github.com/blackjax-devs/blackjax). Specifically, `sbijax` implements
 
-- [Sequential Monte Carlo ABC](https://www.routledge.com/Handbook-of-Approximate-Bayesian-Computation/Sisson-Fan-Beaumont/p/book/9780367733728) (`SMCABC`),
+- [Sequential Monte Carlo ABC](https://www.routledge.com/Handbook-of-Approximate-Bayesian-Computation/Sisson-Fan-Beaumont/p/book/9780367733728) (`SMCABC`)
 - [Neural Likelihood Estimation](https://arxiv.org/abs/1805.07226) (`SNL`)
 - [Surjective Neural Likelihood Estimation](https://arxiv.org/abs/2308.01054) (`SSNL`)
 - [Neural Posterior Estimation C](https://arxiv.org/abs/1905.07488) (short `SNP`)
 - [Contrastive Neural Ratio Estimation](https://arxiv.org/abs/2210.06170) (short `SNR`)
 - [Neural Approximate Sufficient Statistics](https://arxiv.org/abs/2010.10079) (`SNASS`)
 - [Neural Approximate Slice Sufficient Statistics](https://openreview.net/forum?id=jjzJ768iV1) (`SNASSS`)
+- [Flow matching posterior estimation](https://openreview.net/forum?id=jjzJ768iV1) (`SFMPE`)
 
 where the acronyms in parentheses denote the names of the methods in `sbijax`.
 
+> [!CAUTION]
+> ⚠️ As per the LICENSE file, there is no warranty whatsoever for this free software tool. If you discover bugs, please report them.
+
 ## Examples
 
 You can find several self-contained examples on how to use the algorithms in [examples](https://github.com/dirmeier/sbijax/tree/main/examples).

docs/index.rst

@@ -13,13 +13,18 @@
 `JAX <https://github.com/google/jax>`_ using `Haiku <https://github.com/deepmind/dm-haiku>`_,
 `Distrax <https://github.com/deepmind/distrax>`_ and `BlackJAX <https://github.com/blackjax-devs/blackjax>`_. Specifically, :code:`sbijax` implements
 
-- `Sequential Monte Carlo ABC <https://www.routledge.com/Handbook-of-Approximate-Bayesian-Computation/Sisson-Fan-Beaumont/p/book/9780367733728>`_ (:code:`SMCABC`),
+- `Sequential Monte Carlo ABC <https://www.routledge.com/Handbook-of-Approximate-Bayesian-Computation/Sisson-Fan-Beaumont/p/book/9780367733728>`_ (:code:`SMCABC`)
 - `Neural Likelihood Estimation <https://arxiv.org/abs/1805.07226>`_ (:code:`SNL`)
 - `Surjective Neural Likelihood Estimation <https://arxiv.org/abs/2308.01054>`_ (:code:`SSNL`)
 - `Neural Posterior Estimation C <https://arxiv.org/abs/1905.07488>`_ (short :code:`SNP`)
 - `Contrastive Neural Ratio Estimation <https://arxiv.org/abs/2210.06170>`_ (short :code:`SNR`)
 - `Neural Approximate Sufficient Statistics <https://arxiv.org/abs/2010.10079>`_ (:code:`SNASS`)
 - `Neural Approximate Slice Sufficient Statistics <https://openreview.net/forum?id=jjzJ768iV1>`_ (:code:`SNASSS`)
+- `Flow matching posterior estimation <https://arxiv.org/abs/2305.17161>`_ (:code:`SFMPE`)
+
+.. caution::
+
+ ⚠️ As per the LICENSE file, there is no warranty whatsoever for this free software tool. If you discover bugs, please report them.
 
 Installation
 ------------

docs/sbijax.nn.rst

@@ -12,6 +12,7 @@ networks and normalizing flows.
  make_affine_maf
  make_surjective_affine_maf
  make_resnet
+ make_ccnf
  make_snass_net
  make_snasss_net
 
@@ -22,6 +23,8 @@ networks and normalizing flows.
 
 .. autofunction:: make_resnet
 
+.. autofunction:: make_ccnf
+
 .. autofunction:: make_snass_net
 
 .. autofunction:: make_snasss_net

docs/sbijax.rst

@@ -15,6 +15,7 @@ Methods
  SNL
  SNP
  SNR
+ SFMPE
  SNASS
  SNASSS
 
@@ -42,6 +43,12 @@ SNR
 .. autoclass:: SNR
  :members: fit, simulate_data_and_possibly_append, sample_posterior
 
+SFMPE
+~~~~~
+
+.. autoclass:: SFMPE
+ :members: fit, simulate_data_and_possibly_append, sample_posterior
+
 SNASS
 ~~~~~
 

examples/bivariate_gaussian_fmpe.py → examples/bivariate_gaussian_sfmpe.py

@@ -10,7 +10,7 @@
 from jax import numpy as jnp
 from jax import random as jr
 
-from sbijax import FMPE
+from sbijax import SFMPE
 from sbijax.nn import CCNF
 
 
@@ -45,7 +45,7 @@ def run():
  prior_simulator_fn, prior_logdensity_fn = prior_model_fns()
  fns = (prior_simulator_fn, prior_logdensity_fn), simulator_fn
 
- estim = FMPE(fns, make_model(2))
+ estim = SFMPE(fns, make_model(2))
  optimizer = optax.adam(1e-3)
 
  data, params = None, {}

sbijax/__init__.py

@@ -2,11 +2,11 @@
 sbijax: Simulation-based inference in JAX
 """
 
-__version__ = "0.1.8"
+__version__ = "0.1.9"
 
 
 from sbijax._src.abc.smc_abc import SMCABC
-from sbijax._src.fmpe import FMPE
+from sbijax._src.fmpe import SFMPE
 from sbijax._src.snass import SNASS
 from sbijax._src.snasss import SNASSS
 from sbijax._src.snl import SNL

sbijax/_src/fmpe.py

@@ -9,7 +9,6 @@
 from tqdm import tqdm
 
 from sbijax._src._sne_base import SNE
-from sbijax._src.nn.continuous_normalizing_flow import CCNF
 from sbijax._src.util.early_stopping import EarlyStopping
 
 
@@ -59,8 +58,14 @@ def _cfm_loss(
 
 
 # pylint: disable=too-many-arguments,unused-argument,useless-parent-delegation
-class FMPE(SNE):
- """Flow matching posterior estimation.
+class SFMPE(SNE):
+ r"""Sequential flow matching posterior estimation.
+
+ Implements a sequential version of the FMPE algorithm introduced in [1]_.
+ For all rounds $r > 1$ parameter samples
+ :math:`\theta \sim \hat{p}^r(\theta)` are drawn from
+ the approximate posterior instead of the prior when computing the flow
+ matching loss.
 
  Args:
  model_fns: a tuple of tuples. The first element is a tuple that
@@ -71,23 +76,23 @@ class FMPE(SNE):
 
  Examples:
  >>> import distrax
- >>> from sbijax import SNP
- >>> from sbijax.nn import make_affine_maf
+ >>> from sbijax import SFMPE
+ >>> from sbijax.nn import make_ccnf
  >>>
  >>> prior = distrax.Normal(0.0, 1.0)
  >>> s = lambda seed, theta: distrax.Normal(theta, 1.0).sample(seed=seed)
  >>> fns = (prior.sample, prior.log_prob), s
- >>> flow = make_affine_maf()
+ >>> flow = make_ccnf(1)
  >>>
- >>> snr = SNP(fns, flow)
+ >>> estim = SFMPE(fns, flow)
 
  References:
  .. [1] Wildberger, Jonas, et al. "Flow Matching for Scalable
  Simulation-Based Inference." Advances in Neural Information
  Processing Systems, 2024.
  """
 
- def __init__(self, model_fns, density_estimator: CCNF):
+ def __init__(self, model_fns, density_estimator):
  """Construct a FMPE object.
 
  Args:

sbijax/_src/snp.py

@@ -34,9 +34,9 @@ class SNP(SNE):
  >>> prior = distrax.Normal(0.0, 1.0)
  >>> s = lambda seed, theta: distrax.Normal(theta, 1.0).sample(seed=seed)
  >>> fns = (prior.sample, prior.log_prob), s
- >>> flow = make_affine_maf()
+ >>> flow = make_affine_maf(1)
  >>>
- >>> snr = SNP(fns, flow)
+ >>> estim = SNP(fns, flow)
 
  References:
  .. [1] Greenberg, David, et al. "Automatic posterior transformation for