Update API (#14)

* Refactor
* Increment API
* Fix lints

README.md

@@ -16,11 +16,15 @@ SbiJAX so far implements
 - 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`)
 
 ## Examples
 
 You can find several self-contained examples on how to use the algorithms in `examples`.
 
+## Usage
+
 ## Installation
 
 Make sure to have a working `JAX` installation. Depending whether you want to use CPU/GPU/TPU,

examples/bivariate_gaussian_smcabc.py

@@ -7,6 +7,7 @@
 import matplotlib.pyplot as plt
 import seaborn as sns
 from jax import numpy as jnp
+from jax import random as jr
 
 from sbijax import SMCABC
 
@@ -42,8 +43,9 @@ def run():
  fns = (prior_simulator_fn, prior_logdensity_fn), simulator_fn
 
  smc = SMCABC(fns, summary_fn, distance_fn)
- smc.fit(23, y_observed)
- smc_samples, _ = smc.sample_posterior(10, 1000, 1000, 0.8, 500)
+ smc_samples, _ = smc.sample_posterior(
+ jr.PRNGKey(22), y_observed, 10, 1000, 1000, 0.6, 500
+ )
 
  fig, axes = plt.subplots(2)
  for i in range(2):

examples/bivariate_gaussian_snl.py

@@ -11,7 +11,7 @@
 import optax
 import seaborn as sns
 from jax import numpy as jnp
-from jax import random
+from jax import random as jr
 from surjectors import (
  Chain,
  MaskedAutoregressive,
@@ -31,16 +31,14 @@ def prior_model_fns():
 
 
 def simulator_fn(seed, theta):
- p = distrax.Normal(jnp.zeros_like(theta), 0.1)
+ p = distrax.Normal(jnp.zeros_like(theta), 1.0)
  y = theta + p.sample(seed=seed)
  return y
 
 
 def log_density_fn(theta, y):
  prior = distrax.Independent(distrax.Normal(jnp.zeros(2), jnp.ones(2)), 1)
- likelihood = distrax.MultivariateNormalDiag(
- theta, 0.1 * jnp.ones_like(theta)
- )
+ likelihood = distrax.MultivariateNormalDiag(theta, jnp.ones_like(theta))
 
  lp = jnp.sum(prior.log_prob(theta)) + jnp.sum(likelihood.log_prob(y))
  return lp
@@ -94,26 +92,28 @@ def run():
 
  snl = SNL(fns, make_model(2))
  optimizer = optax.adam(1e-3)
- params, info = snl.fit(
- random.PRNGKey(23),
- y_observed,
- optimizer=optimizer,
- n_rounds=3,
- max_n_iter=100,
- batch_size=64,
- n_early_stopping_patience=5,
- sampler="slice",
- )
 
+ data, params = None, {}
+ for i in range(2):
+ data, _ = snl.simulate_data_and_possibly_append(
+ jr.fold_in(jr.PRNGKey(12), i),
+ params=params,
+ observable=y_observed,
+ data=data,
+ )
+ params, info = snl.fit(
+ jr.fold_in(jr.PRNGKey(23), i), data=data, optimizer=optimizer
+ )
+
+ sample_key, rng_key = jr.split(jr.PRNGKey(123))
  slice_samples = sample_with_slice(
- hk.PRNGSequence(0), log_density, 4, 2000, 1000, prior_simulator_fn
+ sample_key, log_density, prior_simulator_fn
  )
  slice_samples = slice_samples.reshape(-1, 2)
- snl_samples, _ = snl.sample_posterior(
- params, 4, 2000, 1000, sampler="slice"
- )
 
- print(f"Took n={snl.n_total_simulations} simulations in total")
+ sample_key, rng_key = jr.split(rng_key)
+ snl_samples, _ = snl.sample_posterior(sample_key, params, y_observed)
+
  fig, axes = plt.subplots(2, 2)
  for i in range(2):
  sns.histplot(

examples/bivariate_gaussian_snp.py

@@ -9,7 +9,7 @@
 import optax
 import seaborn as sns
 from jax import numpy as jnp
-from jax import random
+from jax import random as jr
 from surjectors import Chain, TransformedDistribution
 from surjectors.bijectors.masked_autoregressive import MaskedAutoregressive
 from surjectors.bijectors.permutation import Permutation
@@ -25,7 +25,7 @@ def prior_model_fns():
 
 
 def simulator_fn(seed, theta):
- p = distrax.Normal(jnp.zeros_like(theta), 0.1)
+ p = distrax.Normal(jnp.zeros_like(theta), 1.0)
  y = theta + p.sample(seed=seed)
  return y
 
@@ -72,21 +72,25 @@ def run():
  prior_simulator_fn, prior_logdensity_fn = prior_model_fns()
  fns = (prior_simulator_fn, prior_logdensity_fn), simulator_fn
 
- optimizer = optax.adamw(1e-04)
  snp = SNP(fns, make_model(2))
- params, info = snp.fit(
- random.PRNGKey(2),
- y_observed,
- n_rounds=3,
- optimizer=optimizer,
- n_early_stopping_patience=10,
- batch_size=64,
- n_atoms=10,
- max_n_iter=100,
- )
-
- print(f"Took n={snp.n_total_simulations} simulations in total")
- snp_samples, _ = snp.sample_posterior(params, 10000)
+ optimizer = optax.adam(1e-3)
+
+ data, params = None, {}
+ for i in range(2):
+ data, _ = snp.simulate_data_and_possibly_append(
+ jr.fold_in(jr.PRNGKey(1), i),
+ params=params,
+ observable=y_observed,
+ data=data,
+ )
+ params, info = snp.fit(
+ jr.fold_in(jr.PRNGKey(2), i),
+ data=data,
+ optimizer=optimizer,
+ )
+
+ rng_key = jr.PRNGKey(23)
+ snp_samples, _ = snp.sample_posterior(rng_key, params, y_observed)
  fig, axes = plt.subplots(2)
  for i, ax in enumerate(axes):
  sns.histplot(snp_samples[:, i], color="darkblue", ax=ax)