Add SRNE and remove rejection ABC (#21)

* Add SRNE
* Remove rejection ABC

README.md

@@ -1,29 +1,34 @@
 # sbijax
 
-[![status](http://www.repostatus.org/badges/latest/concept.svg)](http://www.repostatus.org/#concept)
+[![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)
 [![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 implements several algorithms for simulation-based inference using
-[JAX](https://github.com/google/jax), [Haiku](https://github.com/deepmind/dm-haiku) and [BlackJAX](https://github.com/blackjax-devs/blackjax).
+`sbijax` implements several algorithms for simulation-based inference in
+[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
 
-SbiJAX so far implements
+- [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`)
 
-- Rejection ABC (`RejectionABC`),
-- Sequential Monte Carlo ABC (`SMCABC`),
-- Sequential Neural Likelihood Estimation (`SNL`)
-- Surjective Sequential Neural Likelihood Estimation (`SSNL`)
-- Sequential Neural Posterior Estimation C (short `SNP`)
+where the acronyms in parentheses denote the names of the methods in `sbijax`.
 
 ## Examples
 
-You can find several self-contained examples on how to use the algorithms in `examples`.
+You can find several self-contained examples on how to use the algorithms in [examples](https://github.com/dirmeier/sbijax/tree/main/examples).
 
-## Usage
+## Documentation
+
+Documentation can be found [here](https://sbijax.readthedocs.io/en/latest/).
 
 ## Installation
 
@@ -42,6 +47,11 @@ To install the latest GitHub <RELEASE>, use:
 pip install git+https://github.com/dirmeier/sbijax@<RELEASE>
 ```
 
+## Acknowledgements
+
+> 📝 The package draws significant inspiration from the excellent Pytorch-based [`sbi`](https://github.com/sbi-dev/sbi) package which is substantially more
+feature-complete and user-friendly, and better documented.
+
 ## Author
 
 Simon Dirmeier <a href="mailto:sfyrbnd @ pm me">sfyrbnd @ pm me</a>

examples/bivariate_gaussian_snl.py

@@ -13,12 +13,12 @@
 from jax import numpy as jnp
 from jax import random as jr
 from surjectors import (
- MADE,
  Chain,
  MaskedAutoregressive,
  Permutation,
  TransformedDistribution,
 )
+from surjectors.nn import MADE
 from surjectors.util import unstack
 
 from sbijax import SNL

examples/bivariate_gaussian_snp.py

@@ -1,5 +1,5 @@
 """
-Example using sequential posterior estimation on a bivariate Gaussian
+Example using sequential neural posterior estimation on a bivariate Gaussian
 """
 
 import distrax

examples/bivariate_gaussian_snr.py

@@ -0,0 +1,72 @@
+"""
+Example using sequential neural ratio estimation on a bivariate Gaussian
+"""
+
+import distrax
+import haiku as hk
+import matplotlib.pyplot as plt
+import optax
+import seaborn as sns
+from jax import numpy as jnp
+from jax import random as jr
+
+from sbijax import SNR
+
+
+def prior_model_fns():
+ p = distrax.Independent(distrax.Normal(jnp.zeros(2), jnp.ones(2)), 1)
+ return p.sample, p.log_prob
+
+
+def simulator_fn(seed, theta):
+ p = distrax.Normal(jnp.zeros_like(theta), 1.0)
+ y = theta + p.sample(seed=seed)
+ return y
+
+
+def make_model():
+ @hk.without_apply_rng
+ @hk.transform
+ def _mlp(inputs, **kwargs):
+ return hk.nets.MLP([64, 64, 1])(inputs)
+
+ return _mlp
+
+
+def run():
+ y_observed = jnp.array([2.0, -2.0])
+
+ prior_simulator_fn, prior_logdensity_fn = prior_model_fns()
+ fns = (prior_simulator_fn, prior_logdensity_fn), simulator_fn
+
+ snr = SNR(fns, make_model())
+ optimizer = optax.adam(1e-3)
+
+ data, params = None, {}
+ for i in range(5):
+ data, _ = snr.simulate_data_and_possibly_append(
+ jr.fold_in(jr.PRNGKey(1), i),
+ params=params,
+ observable=y_observed,
+ data=data,
+ )
+ params, info = snr.fit(
+ jr.fold_in(jr.PRNGKey(2), i),
+ data=data,
+ optimizer=optimizer,
+ batch_size=100,
+ )
+
+ rng_key = jr.PRNGKey(23)
+ snr_samples, _ = snr.sample_posterior(rng_key, params, y_observed)
+ fig, axes = plt.subplots(2)
+ for i, ax in enumerate(axes):
+ sns.histplot(snr_samples[:, i], color="darkblue", ax=ax)
+ ax.set_xlim([-3.0, 3.0])
+ sns.despine()
+ plt.tight_layout()
+ plt.show()
+
+
+if __name__ == "__main__":
+ run()

examples/bivariate_gaussian_snasss.py → examples/slcp_snass.py

@@ -1,6 +1,5 @@
 """
-Example using sequential neural approximate (slice) summary statistics on a
-bivariate Gaussian with repeated dimensions
+Example SNASS on the SLCP experiment
 """
 
 import distrax
@@ -11,6 +10,7 @@
 import seaborn as sns
 from jax import numpy as jnp
 from jax import random as jr
+from jax import scipy as jsp
 from surjectors import (
  Chain,
  MaskedAutoregressive,
@@ -23,20 +23,64 @@
 from sbijax import SNASSS
 from sbijax.nn import make_snasss_net
 
-W = jr.normal(jr.PRNGKey(0), (2, 10))
-
 
 def prior_model_fns():
- p = distrax.Independent(distrax.Normal(jnp.zeros(2), jnp.ones(2)), 1)
+ p = distrax.Independent(
+ distrax.Uniform(jnp.full(5, -3.0), jnp.full(5, 3.0)), 1
+ )
  return p.sample, p.log_prob
 
 
 def simulator_fn(seed, theta):
- y = theta @ W
- y = y + distrax.Normal(jnp.zeros_like(y), 0.1).sample(seed=seed)
+ orig_shape = theta.shape
+ if theta.ndim == 2:
+ theta = theta[:, None, :]
+ us_key, noise_key = jr.split(seed)
+
+ def _unpack_params(ps):
+ m0 = ps[..., [0]]
+ m1 = ps[..., [1]]
+ s0 = ps[..., [2]] ** 2
+ s1 = ps[..., [3]] ** 2
+ r = jnp.tanh(ps[..., [4]])
+ return m0, m1, s0, s1, r
+
+ m0, m1, s0, s1, r = _unpack_params(theta)
+ us = distrax.Normal(0.0, 1.0).sample(
+ seed=us_key, sample_shape=(theta.shape[0], theta.shape[1], 4, 2)
+ )
+ xs = jnp.empty_like(us)
+ xs = xs.at[:, :, :, 0].set(s0 * us[:, :, :, 0] + m0)
+ y = xs.at[:, :, :, 1].set(
+ s1 * (r * us[:, :, :, 0] + jnp.sqrt(1.0 - r**2) * us[:, :, :, 1]) + m1
+ )
+ if len(orig_shape) == 2:
+ y = y.reshape((*theta.shape[:1], 8))
+ else:
+ y = y.reshape((*theta.shape[:2], 8))
  return y
 
 
+def likelihood_fn(theta, y):
+ mu = jnp.tile(theta[:2], 4)
+ s1, s2 = theta[2] ** 2, theta[3] ** 2
+ corr = s1 * s2 * jnp.tanh(theta[4])
+ cov = jnp.array([[s1**2, corr], [corr, s2**2]])
+ cov = jsp.linalg.block_diag(*[cov for _ in range(4)])
+ p = distrax.MultivariateNormalFullCovariance(mu, cov)
+ return p.log_prob(y)
+
+
+def log_density_fn(theta, y):
+ prior_lp = distrax.Independent(
+ distrax.Uniform(jnp.full(5, -3.0), jnp.full(5, 3.0)), 1
+ ).log_prob(theta)
+ likelihood_lp = likelihood_fn(theta, y)
+
+ lp = jnp.sum(prior_lp) + jnp.sum(likelihood_lp)
+ return lp
+
+
 def make_model(dim):
  def _bijector_fn(params):
  means, log_scales = unstack(params, -1)
@@ -50,7 +94,7 @@ def _flow(method, **kwargs):
  bijector_fn=_bijector_fn,
  conditioner=MADE(
  dim,
- [50, 50, dim * 2],
+ [64, 64, dim * 2],
  2,
  w_init=hk.initializers.TruncatedNormal(0.001),
  b_init=jnp.zeros,
@@ -75,38 +119,39 @@ def _flow(method, **kwargs):
 
 
 def run():
- y_observed = jnp.array([[2.0, -2.0]]) @ W
+ thetas = jnp.linspace(-2.0, 2.0, 5)
+ y_0 = simulator_fn(jr.PRNGKey(0), thetas.reshape(-1, 5)).reshape(-1, 8)
 
  prior_simulator_fn, prior_logdensity_fn = prior_model_fns()
  fns = (prior_simulator_fn, prior_logdensity_fn), simulator_fn
 
  estim = SNASSS(
  fns,
- make_model(2),
- make_snasss_net([64, 64, 2], [64, 64, 1], [64, 64, 1]),
+ make_model(5),
+ make_snasss_net([64, 64, 5], [64, 64, 1], [64, 64, 1]),
  )
  optimizer = optax.adam(1e-3)
 
  data, params = None, {}
- for i in range(2):
+ for i in range(5):
  data, _ = estim.simulate_data_and_possibly_append(
- jr.fold_in(jr.PRNGKey(1), i),
+ jr.fold_in(jr.PRNGKey(12), i),
  params=params,
- observable=y_observed,
+ observable=y_0,
  data=data,
  )
  params, _ = estim.fit(
- jr.fold_in(jr.PRNGKey(2), i),
+ jr.fold_in(jr.PRNGKey(23), i),
  data=data,
  optimizer=optimizer,
  batch_size=100,
  )
 
  rng_key = jr.PRNGKey(23)
- snp_samples, _ = estim.sample_posterior(rng_key, params, y_observed)
- fig, axes = plt.subplots(2)
+ snasss_samples, _ = estim.sample_posterior(rng_key, params, y_0)
+ fig, axes = plt.subplots(5)
  for i, ax in enumerate(axes):
- sns.histplot(snp_samples[:, i], color="darkblue", ax=ax)
+ sns.histplot(snasss_samples[:, i], color="darkblue", ax=ax)
  ax.set_xlim([-3.0, 3.0])
  sns.despine()
  plt.tight_layout()

examples/slcp_ssnl.py

@@ -4,13 +4,11 @@
 
 import argparse
 from functools import partial
-from timeit import default_timer as timer
 
 import distrax
 import haiku as hk
 import jax
 import matplotlib.pyplot as plt
-import numpy as np
 import optax
 import pandas as pd
 import seaborn as sns
@@ -25,11 +23,11 @@
  Permutation,
  TransformedDistribution,
 )
-from surjectors.conditioners import MADE, mlp_conditioner
+from surjectors.nn import MADE, make_mlp
 from surjectors.util import unstack
 
 from sbijax import SNL
-from sbijax.mcmc.slice import sample_with_slice
+from sbijax.mcmc import sample_with_slice
 
 
 def prior_model_fns():
@@ -50,7 +48,7 @@ def _unpack_params(ps):
  m1 = ps[..., [1]]
  s0 = ps[..., [2]] ** 2
  s1 = ps[..., [3]] ** 2
- r = np.tanh(ps[..., [4]])
+ r = jnp.tanh(ps[..., [4]])
  return m0, m1, s0, s1, r
 
  m0, m1, s0, s1, r = _unpack_params(theta)
@@ -60,7 +58,7 @@ def _unpack_params(ps):
  xs = jnp.empty_like(us)
  xs = xs.at[:, :, :, 0].set(s0 * us[:, :, :, 0] + m0)
  y = xs.at[:, :, :, 1].set(
- s1 * (r * us[:, :, :, 0] + np.sqrt(1.0 - r**2) * us[:, :, :, 1]) + m1
+ s1 * (r * us[:, :, :, 0] + jnp.sqrt(1.0 - r**2) * us[:, :, :, 1]) + m1
  )
  if len(orig_shape) == 2:
  y = y.reshape((*theta.shape[:1], 8))
@@ -95,7 +93,7 @@ def _bijector_fn(params):
  return distrax.ScalarAffine(means, jnp.exp(log_scales))
 
  def _decoder_fn(n_dim):
- decoder_net = mlp_conditioner(
+ decoder_net = make_mlp(
  [50, n_dim * 2],
  w_init=hk.initializers.TruncatedNormal(stddev=0.001),
  )
@@ -189,7 +187,6 @@ def run(use_surjectors):
  optimizer = optax.adam(1e-3)
 
  data, params = None, {}
- start = timer()
  for i in range(5):
  data, _ = snl.simulate_data_and_possibly_append(
  jr.fold_in(jr.PRNGKey(12), i),
@@ -201,8 +198,6 @@ def run(use_surjectors):
  params, info = snl.fit(
  jr.fold_in(jr.PRNGKey(23), i), data=data, optimizer=optimizer
  )
- end = timer()
- print(end - start)
 
  sample_key, rng_key = jr.split(jr.PRNGKey(123))
  snl_samples, _ = snl.sample_posterior(sample_key, params, y_observed)

pyproject.toml

@@ -26,7 +26,8 @@ dependencies = [
  "dm-haiku>=0.0.9",
  "optax>=0.1.3",
  "surjectors>=0.3.0",
- "tfp-nightly>=0.20.0.dev20230404"
+ "tfp-nightly>=0.20.0.dev20230404",
+ "tqdm>=4.64.1"
 ]
 dynamic = ["version"]
 
@@ -46,6 +47,12 @@ exclude = [
  "/.pre-commit-config.yaml"
 ]
 
+[tool.hatch.envs.examples]
+dependencies = [
+ "matplotlib>=3.6.1",
+ "seaborn>=0.12.2"
+]
+
 [tool.hatch.envs.test]
 dependencies = [
  "pylint>=2.15.10",
@@ -57,7 +64,6 @@ dependencies = [
 lint = 'pylint sbijax'
 test = 'pytest -v --doctest-modules --cov=./sbijax --cov-report=xml sbijax'
 
-
 [tool.black]
 line-length = 80
 target-version = ['py311']