Examples#
In the following, we demonstrate several sbijax methods using the complex “Simple Liklelihood Complex Posterior” model.
[1]:
import jax
import optax
import os
import sbijax
import seaborn as sns
%matplotlib inline
import matplotlib.pyplot as plt
from matplotlib.ticker import AutoLocator, MaxNLocator
from jax import numpy as jnp, random as jr
from tensorflow_probability.substrates.jax import distributions as tfd
We remove some warnings that TFP is emitting, when using 64-bit arithmetic instead of 32-bit.
[2]:
import warnings
warnings.filterwarnings("ignore")
We implement a custom function to visualize posterior pairs.
[3]:
def plot_posteriors(obj):
cmap = sns.color_palette("rocket", as_cmap=False, desat=0.6, n_colors=10)
cmap = sns.blend_palette(cmap, as_cmap=True)
_, axes = plt.subplots(figsize=(12, 10), nrows=5, ncols=5)
for i in range(0, 5):
for j in range(0, 5):
ax = axes[i, j]
if i < j:
ax.axis('off')
else:
ax.hexbin(obj[..., j], obj[..., i], gridsize=50, bins='log', cmap=cmap)
ax.spines.left.set_linewidth(.5)
ax.spines.bottom.set_linewidth(.5)
ax.spines.right.set_linewidth(.5)
ax.spines.top.set_linewidth(.5)
ax.xaxis.set_major_locator(MaxNLocator(2))
ax.yaxis.set_major_locator(MaxNLocator(2))
ax.xaxis.set_tick_params(width=1, length=2, labelsize=25)
ax.yaxis.set_tick_params(width=1, length=2, labelsize=25)
if i != j:
ax.set_yticks([-3, 0, 3])
ax.set_xticks([-3, 0, 3])
else:
ax.set_yticklabels([])
if i < 4:
ax.set_xticklabels([])
ax.xaxis.set_tick_params(width=0., length=0)
if j != 0:
ax.set_yticklabels([])
ax.yaxis.set_tick_params(width=0., length=0)
ax.grid(which='major', axis='both', alpha=0.5)
for i in range(5):
axes[i, i].hist(obj[..., i], color="black")
return axes
[4]:
def plot_ess_and_trace(samples_arr):
"""Plot trace lines for each parameter across chains."""
from matplotlib.ticker import AutoLocator
colors = sns.color_palette("rocket_r", as_cmap=False, desat=0.6, n_colors=10)
n_chains, n_draws, n_params = samples_arr.shape
_, axes = plt.subplots(figsize=(6, 2 * n_params), nrows=n_params, ncols=1)
axes = list(axes)
for i, ax in enumerate(axes):
for j in range(n_chains):
ax.plot(samples_arr[j, :, i], color=colors[j % len(colors)], alpha=0.4, lw=0.8)
ax.set_ylabel(rf"$\theta_{i}$", fontsize=13)
ax.spines[['right', 'top']].set_visible(False)
ax.yaxis.set_major_locator(AutoLocator())
axes[-1].set_xlabel("draw", fontsize=13)
plt.tight_layout()
return axes
We then define the generative model.
[5]:
def prior_fn():
prior = tfd.JointDistributionNamed(dict(
theta=tfd.Uniform(jnp.full(5, -3.0), jnp.full(5, 3.0))
), batch_ndims=0)
return prior
def simulator_fn(seed, theta):
theta = theta["theta"]
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 = tfd.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
)
y = y.reshape((*theta.shape[:1], 8))
return y
[6]:
y_obs = jnp.array([[
-0.9707123,
-2.9461224,
-0.4494722,
-3.4231849,
-0.13285634,
-3.364017,
-0.85367596,
-2.4271638,
]])
MCMC#
We first sample from the “true” posterior using MCMC, specifically a slice sampler.
[7]:
from functools import partial
from jax import scipy as jsp
from sbijax.mcmc import sample_with_nuts, sample_with_slice
[8]:
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 = tfd.MultivariateNormalFullCovariance(mu, cov)
return p.log_prob(y)
def log_density_fn(theta, y):
prior_lp = tfd.JointDistributionNamed(dict(
theta=tfd.Uniform(jnp.full(5, -3.0), jnp.full(5, 3.0))
)).log_prob(theta)
likelihood_lp = likelihood_fn(theta, y)
lp = jnp.sum(prior_lp) + jnp.sum(likelihood_lp)
return lp
[9]:
log_density = partial(log_density_fn, y=y_obs)
def lp(theta):
return log_density(theta["theta"])
slice_samples_raw, _ = sample_with_slice(
jr.PRNGKey(0),
lp,
prior_fn(),
n_chains=10,
n_samples=10_000,
n_warmup=5_000
)
slice_samples = slice_samples_raw["theta"]
We then compute model diagnostics.
[10]:
slice_samples_dict = {"theta": slice_samples.reshape(10, 5000, 5)}
print("ESS:", sbijax.ess(slice_samples_dict))
print("R-hat:", sbijax.rhat(slice_samples_dict))
ESS: {'theta': Array([49247.61 , 45751.457, 34526.8 , 37308.582, 52743.375], dtype=float32)}
R-hat: {'theta': Array([1.0031444, 1.0033463, 1.0376533, 1.0050069, 1.0019763], dtype=float32)}
[11]:
plot_posteriors(slice_samples.reshape(-1, 5))
plt.tight_layout()
plt.show()
SNLE#
Next, we use surjective neural likelihood estimation to compute a posterior distribution.
[12]:
from sbijax import snle, run_sequential
from sbijax.nn import make_maf
[13]:
n_dim_data = 8
n_layer_dimensions, hidden_sizes = (8, 8, 5, 5, 5), (64, 64)
neural_network = make_maf(
n_dim_data,
n_layer_dimensions=n_layer_dimensions,
hidden_sizes=hidden_sizes
)
prior = prior_fn()
model_snle = snle(prior, neural_network)
[14]:
snle_params, info = run_sequential(
jr.PRNGKey(1),
model_snle,
prior,
simulator_fn,
y_obs,
n_rounds=15,
n_simulations_per_round=1_000,
)
14%|███▊ | 141/1000 [00:44<04:33, 3.14it/s]
59%|███████████████▉ | 588/1000 [03:33<02:29, 2.76it/s]
35%|█████████▌ | 354/1000 [02:29<04:33, 2.36it/s]
32%|████████▋ | 320/1000 [02:31<05:22, 2.11it/s]
25%|██████▋ | 248/1000 [02:11<06:39, 1.88it/s]
20%|█████▎ | 199/1000 [01:59<08:00, 1.67it/s]
26%|███████ | 261/1000 [02:46<07:51, 1.57it/s]
50%|█████████████▍ | 497/1000 [05:43<05:47, 1.45it/s]
31%|████████▎ | 308/1000 [03:49<08:35, 1.34it/s]
29%|███████▊ | 291/1000 [03:52<09:27, 1.25it/s]
22%|█████▉ | 222/1000 [03:12<11:12, 1.16it/s]
20%|█████▎ | 195/1000 [03:03<12:36, 1.06it/s]
21%|█████▋ | 209/1000 [03:28<13:09, 1.00it/s]
15%|████▏ | 154/1000 [03:18<18:11, 1.29s/it]
34%|█████████ | 335/1000 [07:52<15:38, 1.41s/it]
[15]:
snle_samples, _ = model_snle.sample(
jr.PRNGKey(5), snle_params, y_obs, n_samples=5_000, n_warmup=2_500, n_chains=10
)
[16]:
plot_posteriors(
snle_samples["theta"].reshape(-1, 5),
)
plt.tight_layout()
plt.show()
FMPE#
As a comparison, we use flow matching posterior estimation.
[17]:
from sbijax import fmpe, simulate
from sbijax.nn import make_cnf
[18]:
n_dim_theta = 5
n_layers, hidden_size = 5, 128
neural_network = make_cnf(n_dim_theta, n_layers, hidden_size)
model_fmpe = fmpe(prior, neural_network)
[19]:
data = simulate(
jr.PRNGKey(1),
prior,
simulator_fn,
n=20_000,
)
fmpe_params, info = model_fmpe.fit(
jr.PRNGKey(2),
data,
optimizer=optax.adam(0.001),
n_early_stopping_delta=0.00001,
n_early_stopping_patience=30
)
6%|█▋ | 62/1000 [01:20<20:12, 1.29s/it]
[20]:
fmpe_samples, _ = model_fmpe.sample(
jr.PRNGKey(5), fmpe_params, y_obs, n_samples=25_000
)
[21]:
plot_posteriors(
fmpe_samples["theta"].reshape(-1, 5),
)
plt.tight_layout()
plt.show()
SMC-ABC#
Finally, we evaluate SMC-ABC using neural sufficient statistics.
[22]:
from sbijax import nass, smcabc, simulate
from sbijax.nn import make_nass_net
[23]:
n_embedding_dim, hidden_sizes = 5, (64, 64)
neural_network = make_nass_net(n_embedding_dim, hidden_sizes)
model_nass = nass(neural_network)
data = simulate(jr.PRNGKey(1), prior, simulator_fn, n=20_000)
params_nass, _ = model_nass.fit(jr.PRNGKey(2), data, n_early_stopping_patience=25)
22%|██████ | 225/1000 [03:22<11:38, 1.11it/s]
[24]:
def summary_fn(y):
s = model_nass.summarize(params_nass, y)
return s
def distance_fn(y_simulated, y_observed):
diff = y_simulated - y_observed
dist = jax.vmap(lambda el: jnp.linalg.norm(el))(diff)
return dist
[27]:
model_smc = smcabc(prior, simulator_fn, summary_fn, distance_fn)
smc_samples, _ = model_smc.sample(
jr.PRNGKey(5),
y_obs,
n_rounds=10,
n_particles=5_000,
eps_step=0.9,
ess_min=2_000
)
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [12:44<00:00, 76.48s/it]
[28]:
plot_posteriors(
smc_samples["theta"].reshape(-1, 5),
)
plt.tight_layout()
plt.show()
Session info#
[31]:
import session_info
session_info.show(html=False)
-----
haiku 0.0.16
jax 0.10.2
jaxlib 0.10.2
matplotlib 3.11.0
optax 0.2.8
sbijax 0.4.0
seaborn 0.13.2
session_info v1.0.1
tensorflow_probability 0.26.0-dev20260318
-----
IPython 9.11.0
jupyter_client 8.8.0
jupyter_core 5.9.1
jupyterlab 4.5.6
notebook 7.5.5
-----
Python 3.12.10 (main, May 30 2025, 05:53:56) [Clang 20.1.4 ]
macOS-26.2-arm64-arm-64bit
-----
Session information updated at 2026-07-03 19:06