Simulation-based inference (SBI) acceleration

In this tutorial you will learn:

How to specify a mock data generator in addition to the likelihood
How to train a simulation-based inference model to predict a posterior distribution (neural posterior estimation)
How to verify that the simulation-based inference trained well
How to use SBI-NPE to accelerate the UltraNest nested sampling run

In this example, we fit a Voigt line in a spectrum.

[1]:

import numpy as np
import matplotlib.pyplot as plt
import corner
import ultranest
import ultranest.simbase.minisbi

rng = np.random.default_rng(123)

Defining the model

Gaussian measurement noise

First, we need our physical model function of the intrinsic spectrum, and also predict the noise level. We assume a Voigt line shape and constant noise.

[2]:

noise_std = 0.1

[3]:

N_DATA = 1000
x = np.geomspace(0.55, 0.75, N_DATA)

[4]:

parameter_names = ["amplitude", "center_x", "width", "background", "slope"]
N_PARAMS = len(parameter_names)

Define the model spectrum function, given model parameters

[5]:

from scipy.special import voigt_profile

[6]:

sigma_instrumental = 0.001 # nominal resolution of the spectrograph

[7]:

def generate_mean_and_noise(theta, idx=0, seed=0):
    amplitude, center_x, width, background, slope = theta
    signal = amplitude * voigt_profile(x - center_x, sigma_instrumental, width) + background + (x / center_x) ** slope

    return {
        'mean':      signal,
        'noise_std': signal * 0 + noise_std,  # same noise for each data point
    }

The received idx and seed are not used here. They are integers, in case randomness in the physics model is needed.

Our true parameters for this example:

[8]:

theta_true = [0.0042, 0.656279, 0.001, 0.2, -1.]

[9]:

plt.errorbar(
    x=x,
    y=generate_mean_and_noise(theta_true)['mean'],
    capsize=0,
)
plt.fill_between(
    x,
    generate_mean_and_noise(theta_true)['mean'] - generate_mean_and_noise(theta_true)['noise_std'],
    generate_mean_and_noise(theta_true)['mean'] + generate_mean_and_noise(theta_true)['noise_std'],
    alpha=0.1,
)
plt.yscale('log');

_images/example-simulation-based-inference_15_0.svg

A beautiful spectrum of a line, on top of a slightly sloped continuum.

Prior distribution

Define our prior distribution as usual, based on unit cube transforms.

[10]:

def prior_transform(u):
    result = u.copy()
    # amplitude (index 0): log-uniform [1e-2, 1e2]
    result[0] = 10**(-4 + u[0] * 4)
    # center_x (index 1): uniform [0.65, 0.66]
    result[1] = 0.65 + u[1] * 0.01
    # width (index 2): uniform [1e-4, 0.01]
    result[2] = 10**(-4 + 2 * u[2])
    # background level (index 3): uniform [-1, 1]
    result[3] = -1 + 2 * u[3]
    # slope (index 4): uniform [-2, +2]
    result[4] = -2 + 4 * u[4]
    return result

Generate noisy data

The observed data we will fit

Going from the model above to data is trivial, with a Gaussian random number generator:

[11]:

def inject_noise(mean_props, rng):
    return rng.normal(mean_props['mean'], mean_props['noise_std'])

[12]:

obs_spectrum = inject_noise(generate_mean_and_noise(theta_true), rng)

plt.plot(x, obs_spectrum)
plt.yscale('log')

_images/example-simulation-based-inference_24_0.svg

Generate mock data - Prior predictive checks

Here are some example data sets when we sampled from the prior. These show the diversity of data we consider reasonable a priori:

[13]:

for i in range(10):
    u = rng.uniform(size=N_PARAMS)
    params = prior_transform(u)
    plt.plot(x, inject_noise(generate_mean_and_noise(params), rng))
    plt.yscale('log')
    plt.ylim(1e-2, 40)

_images/example-simulation-based-inference_27_0.svg

Data likelihood

We adopt a corresponding Gaussian likelihood function, that we will use at the end with ultranest to improve upon the SBI approach.

[14]:

def loglikelihood(mean_props, noisy_data):
    # the usual Gaussian likelihood function, comparing data to the model, under noise
    return np.sum(
        -0.5 * np.log(2.0 * np.pi)
        - np.log(mean_props['noise_std'])
        - 0.5 * ((noisy_data - mean_props['mean']) / mean_props['noise_std']) ** 2
    )

Simulation-based inference

The simulation-based inference trains a neural network to predict a simple factorized posterior distribution on the unit hypercube.

The training pipeline is highly efficient, reusing generated model with different noise.

Sampled datasets will be cached internally to a folder.

[15]:

folder = 'mygaussline'

[16]:

# optional: clear cache whenever you change the prior or model function
!rm -rf mygaussline/

This happens efficiently in batches (using joblib.Memory). Some helper functions that call our functions in a loop:

[17]:

from ultranest.simbase.minisbi.utils import make_memory, make_cached_generate

[18]:

memory = make_memory(folder)
generate_noiseless_batch = make_cached_generate(
    memory, prior_transform, generate_mean_and_noise
)

Training the model

Next train a simple neural network to predict a posterior distribution.

[19]:

from ultranest.simbase.minisbi.npe import train_npe

[20]:

model = train_npe(
    # our folder
    folder=folder,
    # our (costly) model function
    generate_noiseless_batch=generate_noiseless_batch,
    # our cheap noisy generator:
    inject_noise=inject_noise,
    # our model parameters
    n_params=N_PARAMS,

    # Parameters related to training strategy:

    # seed for randomness
    base_seed=12345,
    # how long to train: Choose how long you want to wait - the higher the better.
    max_model_evals=100000,
    # optional: if > 1, evaluate models in parallel (with joblib.Parallel)
    num_processes=1,
    # the full history of noise-free models generated with generate_mean_and_noise is kept.
    # they are reused and additional datasets are generated with fresh noise
    # fresh_example_fraction sets the balance between freshly generated models and reuse of previously generated models
    # if the model is cheap, can be 1.0.
    # Probably should not be much lower than 0.001. -- if the loss becomes extremely small (-300), the network memorized all samples.
    fresh_example_fraction=0.01,
    # choose how many new samples are generated in each epoch
    # this sets the number of model evaluations per epoch
    # Note: This sets the batch size seen by the network in each evaluation, by batch_size=fresh_sim_batch_size / fresh_example_fraction - you probably want batch_size of O(1000)
    fresh_sim_batch_size=32,
    # neural network training is somewhat stochastic and may not give an improvement in each epoch
    # set the patience to how many epochs to wait for a significant improvement
    # in any case, the best neural network is kept.
    patience=10,
    # what loss improvement is designated a substantial improvement
    patience_min_delta=0.01,

    # parameters for the neural network architecture and learning

    # learning rate - need to chose this depending on the batch size.
    # Increase if loss is changing slowly, decrease if the loss becomes erratic.
    # Try 1e-3, 1e-4, etc.
    npe_lr=3e-4,
    # maximum width of the layers, 1024 is a decent default, but worth trying other values
    npe_width=512,
    # maximum number of layers. 'cascade' is not very sensitive to this parameter.
    npe_depth=6,
    # shape:
    # - 'rectangular' - a classic MLP where each of the npe_depth layer has npe_width neurons
    # - 'triangular' - a funnel-shaped MLP, with each next layer decreasing linearly in width
    # - 'cascade' - subsequent layers are half the width, half the neurons are forwarded to the next layer, the other half skip to the final layer.
    # 'cascade' is recommended, because it avoids having to play much with npe_width and npe_depth - it contains both a wide and a deep network.
    # the neural network layers can be of equal layer width
    layer_shape='cascade',
    # activation function
    npe_activation_cls='ReLU',
)

fresh_example_fraction=0.010 =>  n_fresh=32, replay_example_count=3168 per sub-batch
Starting parallel simulation generator (3220 batches total, 1 workers) ...
Detected n_data=1000 from probe simulation.
hidden layer widths: [512, 256, 128, 64, 32, 16]
Training NPE (product of Kumaraswamy-Logistic distributions) with on-the-fly batch generation ...
Computing normalisation statistics ...
[ZScoreNorm] fitted 1000 features from 2000 samples.
  mean : min=1.023231  max=10.345800
  std  : min=0.592311   max=27.783062
Input normalisation fitted.
Generating fixed NPE validation set (val_size=1024) ...
Training NPE with on-the-fly batch generation ...
  Epoch   1/390  train_loss=-1.0859  val_loss=-1.4058  fresh_simulator_evals=256  *
  Epoch   2/390  train_loss=-2.1006  val_loss=-2.1500  fresh_simulator_evals=512  *
  Epoch   3/390  train_loss=-2.7577  val_loss=-2.6609  fresh_simulator_evals=768  *
  Epoch   4/390  train_loss=-2.9499  val_loss=-2.8278  fresh_simulator_evals=1024  *
  Epoch   5/390  train_loss=-3.1321  val_loss=-2.9819  fresh_simulator_evals=1280  *
  Epoch   6/390  train_loss=-3.2450  val_loss=-3.1825  fresh_simulator_evals=1536  *
  Epoch   7/390  train_loss=-3.4213  val_loss=-3.3585  fresh_simulator_evals=1792  *
  Epoch   8/390  train_loss=-3.6505  val_loss=-3.4168  fresh_simulator_evals=2048  *
  Epoch   9/390  train_loss=-3.6947  val_loss=-3.5661  fresh_simulator_evals=2304  *
  Epoch  10/390  train_loss=-3.8109  val_loss=-3.8310  fresh_simulator_evals=2560  *
  Epoch  11/390  train_loss=-4.0804  val_loss=-3.7576  fresh_simulator_evals=2816  [1/10]
  Epoch  12/390  train_loss=-4.2388  val_loss=-4.3213  fresh_simulator_evals=3072  *
  Epoch  13/390  train_loss=-4.5041  val_loss=-4.5648  fresh_simulator_evals=3328  *
  Epoch  14/390  train_loss=-4.7074  val_loss=-4.6511  fresh_simulator_evals=3584  *
  Epoch  15/390  train_loss=-4.9803  val_loss=-4.9996  fresh_simulator_evals=3840  *
  Epoch  16/390  train_loss=-5.1568  val_loss=-5.2287  fresh_simulator_evals=4096  *
  Epoch  17/390  train_loss=-5.2509  val_loss=-5.2191  fresh_simulator_evals=4352  [1/10]
  Epoch  18/390  train_loss=-5.1856  val_loss=-5.2102  fresh_simulator_evals=4608  [2/10]
  Epoch  19/390  train_loss=-5.4123  val_loss=-4.8198  fresh_simulator_evals=4864  [3/10]
  Epoch  20/390  train_loss=-4.5319  val_loss=-4.4832  fresh_simulator_evals=5120  [4/10]
  Epoch  21/390  train_loss=-4.6432  val_loss=-4.7081  fresh_simulator_evals=5376  [5/10]
  Epoch  22/390  train_loss=-5.0012  val_loss=-5.1695  fresh_simulator_evals=5632  [6/10]
  Epoch  23/390  train_loss=-5.3541  val_loss=-5.4500  fresh_simulator_evals=5888  *
  Epoch  24/390  train_loss=-5.6689  val_loss=-5.6907  fresh_simulator_evals=6144  *
  Epoch  25/390  train_loss=-5.7531  val_loss=-5.9757  fresh_simulator_evals=6400  *
  Epoch  26/390  train_loss=-6.1245  val_loss=-6.1615  fresh_simulator_evals=6656  *
  Epoch  27/390  train_loss=-5.9483  val_loss=-5.7712  fresh_simulator_evals=6912  [1/10]
  Epoch  28/390  train_loss=-6.0647  val_loss=-6.0812  fresh_simulator_evals=7168  [2/10]
  Epoch  29/390  train_loss=-6.2917  val_loss=-5.9434  fresh_simulator_evals=7424  [3/10]
  Epoch  30/390  train_loss=-6.1702  val_loss=-6.2135  fresh_simulator_evals=7680  *
  Epoch  31/390  train_loss=-6.4975  val_loss=-6.3496  fresh_simulator_evals=7936  *
  Epoch  32/390  train_loss=-6.2768  val_loss=-6.4711  fresh_simulator_evals=8192  *
  Epoch  33/390  train_loss=-6.6392  val_loss=-6.4348  fresh_simulator_evals=8448  [1/10]
  Epoch  34/390  train_loss=-5.9726  val_loss=-5.9031  fresh_simulator_evals=8704  [2/10]
  Epoch  35/390  train_loss=-6.2425  val_loss=-6.4218  fresh_simulator_evals=8960  [3/10]
  Epoch  36/390  train_loss=-6.5963  val_loss=-6.3069  fresh_simulator_evals=9216  [4/10]
  Epoch  37/390  train_loss=-6.1365  val_loss=-6.3211  fresh_simulator_evals=9472  [5/10]
  Epoch  38/390  train_loss=-6.4193  val_loss=-6.6708  fresh_simulator_evals=9728  *
  Epoch  39/390  train_loss=-6.7756  val_loss=-6.5556  fresh_simulator_evals=9984  [1/10]
  Epoch  40/390  train_loss=-6.5203  val_loss=-6.6194  fresh_simulator_evals=10240  [2/10]
  Epoch  41/390  train_loss=-6.8407  val_loss=-7.0332  fresh_simulator_evals=10496  *
  Epoch  42/390  train_loss=-6.9086  val_loss=-7.0496  fresh_simulator_evals=10752  *
  Epoch  43/390  train_loss=-7.1285  val_loss=-7.0607  fresh_simulator_evals=11008  *
  Epoch  44/390  train_loss=-7.0242  val_loss=-6.8919  fresh_simulator_evals=11264  [1/10]
  Epoch  45/390  train_loss=-6.9475  val_loss=-6.9319  fresh_simulator_evals=11520  [2/10]
  Epoch  46/390  train_loss=-7.2324  val_loss=-7.2551  fresh_simulator_evals=11776  *
  Epoch  47/390  train_loss=-7.3234  val_loss=-7.3166  fresh_simulator_evals=12032  *
  Epoch  48/390  train_loss=-7.3908  val_loss=-7.4270  fresh_simulator_evals=12288  *
  Epoch  49/390  train_loss=-7.5515  val_loss=-7.4421  fresh_simulator_evals=12544  *
  Epoch  50/390  train_loss=-7.5906  val_loss=-7.3770  fresh_simulator_evals=12800  [1/10]
  Epoch  51/390  train_loss=-7.5328  val_loss=-7.5716  fresh_simulator_evals=13056  *
  Epoch  52/390  train_loss=-7.6398  val_loss=-7.5584  fresh_simulator_evals=13312  [1/10]
  Epoch  53/390  train_loss=-7.6474  val_loss=-7.5450  fresh_simulator_evals=13568  [2/10]
  Epoch  54/390  train_loss=-7.4859  val_loss=-7.6796  fresh_simulator_evals=13824  *
  Epoch  55/390  train_loss=-7.7760  val_loss=-7.6803  fresh_simulator_evals=14080  [1/10]
  Epoch  56/390  train_loss=-7.5714  val_loss=-7.6185  fresh_simulator_evals=14336  [2/10]
  Epoch  57/390  train_loss=-7.8554  val_loss=-7.8747  fresh_simulator_evals=14592  *
  Epoch  58/390  train_loss=-7.7418  val_loss=-7.9448  fresh_simulator_evals=14848  *
  Epoch  59/390  train_loss=-7.8398  val_loss=-7.7532  fresh_simulator_evals=15104  [1/10]
  Epoch  60/390  train_loss=-7.9498  val_loss=-7.9569  fresh_simulator_evals=15360  *
  Epoch  61/390  train_loss=-7.7907  val_loss=-7.8651  fresh_simulator_evals=15616  [1/10]
  Epoch  62/390  train_loss=-8.0232  val_loss=-8.0780  fresh_simulator_evals=15872  *
  Epoch  63/390  train_loss=-7.9672  val_loss=-8.0524  fresh_simulator_evals=16128  [1/10]
  Epoch  64/390  train_loss=-8.0996  val_loss=-8.1067  fresh_simulator_evals=16384  *
  Epoch  65/390  train_loss=-8.0187  val_loss=-8.0637  fresh_simulator_evals=16640  [1/10]
  Epoch  66/390  train_loss=-8.2671  val_loss=-8.3265  fresh_simulator_evals=16896  *
  Epoch  67/390  train_loss=-8.3237  val_loss=-7.7841  fresh_simulator_evals=17152  [1/10]
  Epoch  68/390  train_loss=-8.2265  val_loss=-8.2941  fresh_simulator_evals=17408  [2/10]
  Epoch  69/390  train_loss=-8.3394  val_loss=-8.1816  fresh_simulator_evals=17664  [3/10]
  Epoch  70/390  train_loss=-8.3642  val_loss=-8.3410  fresh_simulator_evals=17920  *
  Epoch  71/390  train_loss=-8.3273  val_loss=-8.3490  fresh_simulator_evals=18176  [1/10]
  Epoch  72/390  train_loss=-8.3108  val_loss=-8.2912  fresh_simulator_evals=18432  [2/10]
  Epoch  73/390  train_loss=-8.4326  val_loss=-8.2621  fresh_simulator_evals=18688  [3/10]
  Epoch  74/390  train_loss=-8.2535  val_loss=-8.4724  fresh_simulator_evals=18944  *
  Epoch  75/390  train_loss=-8.5584  val_loss=-8.4537  fresh_simulator_evals=19200  [1/10]
  Epoch  76/390  train_loss=-8.5312  val_loss=-8.6644  fresh_simulator_evals=19456  *
  Epoch  77/390  train_loss=-8.6219  val_loss=-8.7943  fresh_simulator_evals=19712  *
  Epoch  78/390  train_loss=-8.6888  val_loss=-8.7142  fresh_simulator_evals=19968  [1/10]
  Epoch  79/390  train_loss=-8.6953  val_loss=-8.8388  fresh_simulator_evals=20224  *
  Epoch  80/390  train_loss=-8.8266  val_loss=-8.5743  fresh_simulator_evals=20480  [1/10]
  Epoch  81/390  train_loss=-8.7541  val_loss=-8.6009  fresh_simulator_evals=20736  [2/10]
  Epoch  82/390  train_loss=-8.7485  val_loss=-8.4907  fresh_simulator_evals=20992  [3/10]
  Epoch  83/390  train_loss=-8.6595  val_loss=-8.9014  fresh_simulator_evals=21248  *
  Epoch  84/390  train_loss=-8.9132  val_loss=-8.2830  fresh_simulator_evals=21504  [1/10]
  Epoch  85/390  train_loss=-8.7887  val_loss=-8.8570  fresh_simulator_evals=21760  [2/10]
  Epoch  86/390  train_loss=-9.0390  val_loss=-8.6846  fresh_simulator_evals=22016  [3/10]
  Epoch  87/390  train_loss=-8.9010  val_loss=-8.3657  fresh_simulator_evals=22272  [4/10]
  Epoch  88/390  train_loss=-8.8941  val_loss=-8.8016  fresh_simulator_evals=22528  [5/10]
  Epoch  89/390  train_loss=-9.0972  val_loss=-8.6109  fresh_simulator_evals=22784  [6/10]
  Epoch  90/390  train_loss=-8.8621  val_loss=-9.0763  fresh_simulator_evals=23040  *
  Epoch  91/390  train_loss=-9.0900  val_loss=-8.9305  fresh_simulator_evals=23296  [1/10]
  Epoch  92/390  train_loss=-8.9290  val_loss=-8.9283  fresh_simulator_evals=23552  [2/10]
  Epoch  93/390  train_loss=-8.9878  val_loss=-9.1271  fresh_simulator_evals=23808  *
  Epoch  94/390  train_loss=-9.1127  val_loss=-9.3154  fresh_simulator_evals=24064  *
  Epoch  95/390  train_loss=-9.1178  val_loss=-9.2105  fresh_simulator_evals=24320  [1/10]
  Epoch  96/390  train_loss=-9.3148  val_loss=-9.3123  fresh_simulator_evals=24576  [2/10]
  Epoch  97/390  train_loss=-9.1447  val_loss=-9.0263  fresh_simulator_evals=24832  [3/10]
  Epoch  98/390  train_loss=-9.1825  val_loss=-9.3789  fresh_simulator_evals=25088  *
  Epoch  99/390  train_loss=-9.3067  val_loss=-9.1505  fresh_simulator_evals=25344  [1/10]
  Epoch 100/390  train_loss=-9.0960  val_loss=-9.2701  fresh_simulator_evals=25600  [2/10]
  Epoch 101/390  train_loss=-9.2554  val_loss=-9.2724  fresh_simulator_evals=25856  [3/10]
  Epoch 102/390  train_loss=-9.3337  val_loss=-9.4309  fresh_simulator_evals=26112  *
  Epoch 103/390  train_loss=-9.2619  val_loss=-9.2754  fresh_simulator_evals=26368  [1/10]
  Epoch 104/390  train_loss=-9.3361  val_loss=-9.4550  fresh_simulator_evals=26624  *
  Epoch 105/390  train_loss=-9.0943  val_loss=-8.9833  fresh_simulator_evals=26880  [1/10]
  Epoch 106/390  train_loss=-9.3134  val_loss=-9.3988  fresh_simulator_evals=27136  [2/10]
  Epoch 107/390  train_loss=-9.4243  val_loss=-9.0658  fresh_simulator_evals=27392  [3/10]
  Epoch 108/390  train_loss=-9.1585  val_loss=-9.0934  fresh_simulator_evals=27648  [4/10]
  Epoch 109/390  train_loss=-9.3143  val_loss=-9.3066  fresh_simulator_evals=27904  [5/10]
  Epoch 110/390  train_loss=-9.5174  val_loss=-9.5110  fresh_simulator_evals=28160  *
  Epoch 111/390  train_loss=-9.5782  val_loss=-9.6021  fresh_simulator_evals=28416  *
  Epoch 112/390  train_loss=-9.7335  val_loss=-9.7840  fresh_simulator_evals=28672  *
  Epoch 113/390  train_loss=-9.4793  val_loss=-9.5026  fresh_simulator_evals=28928  [1/10]
  Epoch 114/390  train_loss=-9.5387  val_loss=-9.6813  fresh_simulator_evals=29184  [2/10]
  Epoch 115/390  train_loss=-9.6575  val_loss=-9.6273  fresh_simulator_evals=29440  [3/10]
  Epoch 116/390  train_loss=-9.7741  val_loss=-9.6144  fresh_simulator_evals=29696  [4/10]
  Epoch 117/390  train_loss=-9.7656  val_loss=-9.6001  fresh_simulator_evals=29952  [5/10]
  Epoch 118/390  train_loss=-9.5755  val_loss=-9.7174  fresh_simulator_evals=30208  [6/10]
  Epoch 119/390  train_loss=-9.6204  val_loss=-9.7690  fresh_simulator_evals=30464  [7/10]
  Epoch 120/390  train_loss=-9.8515  val_loss=-10.0108  fresh_simulator_evals=30720  *
  Epoch 121/390  train_loss=-9.9032  val_loss=-9.6525  fresh_simulator_evals=30976  [1/10]
  Epoch 122/390  train_loss=-9.7036  val_loss=-9.4521  fresh_simulator_evals=31232  [2/10]
  Epoch 123/390  train_loss=-9.6374  val_loss=-9.8079  fresh_simulator_evals=31488  [3/10]
  Epoch 124/390  train_loss=-9.8491  val_loss=-9.8141  fresh_simulator_evals=31744  [4/10]
  Epoch 125/390  train_loss=-9.6394  val_loss=-9.5636  fresh_simulator_evals=32000  [5/10]
  Epoch 126/390  train_loss=-9.7156  val_loss=-9.9512  fresh_simulator_evals=32256  [6/10]
  Epoch 127/390  train_loss=-9.8119  val_loss=-9.8539  fresh_simulator_evals=32512  [7/10]
  Epoch 128/390  train_loss=-10.0046  val_loss=-9.9384  fresh_simulator_evals=32768  [8/10]
  Epoch 129/390  train_loss=-9.9534  val_loss=-9.7611  fresh_simulator_evals=33024  [9/10]
  Epoch 130/390  train_loss=-9.8177  val_loss=-9.9452  fresh_simulator_evals=33280  [10/10]
Early stopping at epoch 130 (patience=10, fresh_simulator_evals=33280)
Restored best model (val_loss=-10.0108)
Model saved to mygaussline/minisbi_npe_1000d_512_512_512_512_512_512_ReLU_KLP_cascade.pt

The meaning of the training output is as follows:

The loss is the log-density of the predicted posterior distribution at the true parameter values.

The training loss is what the optimizer sees. Ideally, this should not stall.

The validation loss is the posterior density of independent samples generated at the beginning and kept separate. These mock data and their true parameters are not seen by the training and kept fixed. The validation loss should keep improving (decreasing).

The fresh_simulator_evals gives the number of model calls performed so far.

The last column gives a star when the model is improving, and otherwise counting up until patience epochs give no improvements.

At the end, the best model is kept.

SBI-NPE posterior distribution

The neural posterior may loose a lot of information, may be biased, etc. We will check this below.

But let’s apply the neural model to our data and see what we get:

[21]:

from ultranest.simbase.minisbi.npe import sample_posterior

posterior_samples_u, posterior_samples_theta = sample_posterior(
    model=model,
    observed_data=obs_spectrum,
    prior_transform=prior_transform,
    n_params=N_PARAMS,
    n_posterior_samples=100000,
)


Predicted posterior (Kumaraswamy-Logistic product, unit-cube space):
  Param 0: loc=0.9768  scale=0.2157  a=1.0192  b=0.2207
  Param 1: loc=0.2924  scale=0.0882  a=1.5984  b=4.2158
  Param 2: loc=0.1643  scale=0.0922  a=1.5926  b=6.1267
  Param 3: loc=0.5780  scale=0.0049  a=0.1956  b=0.4017
  Param 4: loc=0.3793  scale=0.0131  a=0.3000  b=0.2114

/home/user/Downloads/UltraNest/ultranest/simbase/minisbi/logistic.py:39: RuntimeWarning: overflow encountered in exp
  return 1.0 / (1.0 + np.exp(-np.clip(x, -500.0, 500.0)))

[22]:

print("Posterior summary:")
for i, name in enumerate(parameter_names):
    samp = posterior_samples_theta[:, i]
    print(f"  {name:12s}: mean={samp.mean():.3f}  std={samp.std():.3f}"
          f"  true={theta_true[i]:.3f}")

Posterior summary:
  amplitude   : mean=0.005  std=0.014  true=0.004
  center_x    : mean=0.655  std=0.002  true=0.656
  width       : mean=0.001  std=0.002  true=0.001
  background  : mean=0.202  std=0.008  true=0.200
  slope       : mean=-0.990  std=0.035  true=-1.000

[23]:

corner.corner(
    posterior_samples_theta,
    labels=parameter_names,
    truths=theta_true,
    truth_color="red",
    show_titles=True,
    quantiles=[0.16, 0.5, 0.84],
    color="steelblue",
    plot_datapoints=False,
    fill_contours=True,
    smooth=1.0,
);

_images/example-simulation-based-inference_48_0.svg

Getting the NPE distribution parameters

The above is a Kumaraswamy-Logistic chained distribution defined on the unit hypercube. It can be summarized with a few distribution parameters for each model parameters.

[24]:

from ultranest.simbase.minisbi.nested import get_distribution_parameters, KLPTransform

[25]:

distribution_parameters = get_distribution_parameters(model, obs_spectrum)

[26]:

distribution_parameters

[26]:

{'loc': [0.9768165946006775,
  0.2923585772514343,
  0.16428230702877045,
  0.5779978036880493,
  0.3793465793132782],
 'scale': [0.2157188355922699,
  0.08823172003030777,
  0.0922355204820633,
  0.004875042010098696,
  0.013096657581627369],
 'a': [1.0191576480865479,
  1.5984463691711426,
  1.5925710201263428,
  0.19558104872703552,
  0.3000223636627197],
 'b': [0.22065173089504242,
  4.215761661529541,
  6.126714706420898,
  0.40171727538108826,
  0.21142494678497314]}

It might be useful to store for reuse:

[27]:

import json
with open(folder + '/distribution.json', 'w') as f:
    json.dump(distribution_parameters, f)

Checking the SBI model: rank histogram

First, we check that the model gives reliable posteriors, with a rank histogram.

We generate from the prior, and sample posterior samples. Ideally, the rank of the prior samples and the rank of the posterior samples are on a line (see Simulation-based calibration paper).

[28]:

from ultranest.simbase.minisbi.plot import rank_histogram

[29]:

ranks, fig, axes = rank_histogram(
    model=model,
    generate_noiseless_batch=generate_noiseless_batch,
    inject_noise=inject_noise,
    prior_transform=prior_transform,
    n_params=N_PARAMS,
    folder=folder,
    n_test=5000,
    seed=98765,
    param_names=parameter_names,
    n_bins=20,
)
fig;

Rank histogram:   0%|                                                                | 0/5000 [00:00<?, ?sim/s]/home/user/Downloads/UltraNest/ultranest/simbase/minisbi/logistic.py:39: RuntimeWarning: overflow encountered in exp
  return 1.0 / (1.0 + np.exp(-np.clip(x, -500.0, 500.0)))
Rank histogram: 100%|███████████████████████████████████████████████████| 5000/5000 [00:00<00:00, 8673.05sim/s]

_images/example-simulation-based-inference_58_1.svg

These histograms should be flat.

If they a U-shaped, the posterior is over-confident. If they are inverse U-shaped, the posterior is under-confident.

You might be able to improve this by choosing a larger neural network.

[30]:

print("  Per-parameter rank statistics:")
for i, name in enumerate(parameter_names):
    print(f"    {name:12s}: mean={ranks[:, i].mean():.1f}  "
          f"std={ranks[:, i].std():.1f}  expected_mean=100.0")

  Per-parameter rank statistics:
    amplitude   : mean=97.8  std=51.0  expected_mean=100.0
    center_x    : mean=94.5  std=57.5  expected_mean=100.0
    width       : mean=97.2  std=57.8  expected_mean=100.0
    background  : mean=91.3  std=51.5  expected_mean=100.0
    slope       : mean=101.0  std=57.7  expected_mean=100.0

Checking the SBI model: coverage test

Next, we check whether the x% credible intervals contain the true value in x% of cases:

[31]:

from ultranest.simbase.minisbi.plot import parameter_coverage_test

[32]:

coverage, fig, axes = parameter_coverage_test(
    model=model,
    generate_noiseless_batch=generate_noiseless_batch,
    inject_noise=inject_noise,
    prior_transform=prior_transform,
    n_params=N_PARAMS,
    folder=folder,
    n_test=500,
    credible_levels=np.linspace(0.05, 0.99, 20).tolist(),
    seed=11223,
    param_names=parameter_names,
)
fig;

Coverage test:   0%|                                                                  | 0/500 [00:00<?, ?sim/s]/home/user/Downloads/UltraNest/ultranest/simbase/minisbi/logistic.py:39: RuntimeWarning: overflow encountered in exp
  return 1.0 / (1.0 + np.exp(-np.clip(x, -500.0, 500.0)))
Coverage test: 100%|██████████████████████████████████████████████████████| 500/500 [00:00<00:00, 8032.36sim/s]

_images/example-simulation-based-inference_64_1.svg

These should follow the diagonal (red dashed).

Most of them look good. The amplitude parameter is slightly off, but not too wild.

If not, the posterior for that parameter is not well-calibrated (too narrow or too wide error bars). You might be able to improve this by choosing a larger neural network.

[33]:

print("  Empirical coverage at nominal 90% level:")
levels = np.linspace(0.05, 0.99, 20)
idx_90 = np.argmin(np.abs(levels - 0.90))
for i, name in enumerate(parameter_names):
    print(f"    {name:12s}: {coverage[idx_90, i]:.3f}  (nominal=0.900)")

  Empirical coverage at nominal 90% level:
    amplitude   : 0.930  (nominal=0.900)
    center_x    : 0.914  (nominal=0.900)
    width       : 0.888  (nominal=0.900)
    background  : 0.940  (nominal=0.900)
    slope       : 0.886  (nominal=0.900)

ultranest-sbi: SBI-accelerated nested sampling

Setting up the accelerated inference

Here we run nested sampling using the NPE posterior as a reparameterisation.

This allows nested sampling to start at a shrunk-down prior. To ensure identical inference as if we had started with the full prior, we need to correct the likelihood function by a penalty.

The nested sampler samples t ~ Uniform(0,1)^d.
The prior transform maps: t -> u -> theta
The log-likelihood includes a correction: loglike(theta) - log q(u | x_obs)

The correction ensures that the evidence integral and posterior samples are correct with respect to the original flat prior over u (and thus the physical prior over theta).

[34]:

# a helper object defining the distribution, its transform and log-posterior
klp_transform = KLPTransform(**distribution_parameters)

[35]:

class _State:
    last_t = None
state = _State()

# warp the prior by the NPE distribution
def warped_prior_transform(t):
    # keep the last state for likelihood
    state.last_t = t.copy()
    # [0,1]^d -> [0,1]^d with the SBI posterior warp
    u     = klp_transform.transform(t).astype(float)
    # [0,1]^d -> physical parameters
    theta = prior_transform(u)
    return theta

# correct the likelihood for the warp
def corrected_log_likelihood(theta):
    mp   = generate_mean_and_noise(theta=theta, idx=0, seed=0)
    # compute our model likelihood
    logL = loglikelihood(mp, obs_spectrum)
    # compute neural posterior density
    log_jac = klp_transform.log_jacobian(state.last_t)
    return logL + log_jac

Setting up the accelerated sampler

Build our accelerated ultranest-sbi sampler:

[36]:

sampler = ultranest.ReactiveNestedSampler(
    parameter_names,
    corrected_log_likelihood,
    warped_prior_transform,
    log_dir=folder,
    resume='overwrite',
)

Torch modified the logging to show all messages, let’s suppress that

[37]:

import logging

logger = logging.getLogger('ultranest')
logger.setLevel(logging.ERROR)
#logging.basicConfig(level=logging.INFO, stream=sys.stderr)

Run the accelerated ultranest-sbi sampler

[38]:

results = sampler.run(frac_remain=0.5, max_num_improvement_loops=0)

Z=876.7(49.95%) | Like=887.92..889.45 [887.9212..887.9236]*| it/evals=4720/12885 eff=37.8054% N=400

Posterior distribution:

Looking at the posterior distribution of the ultranest-sbi sampler, it is much narrower than the SBI one:

[39]:

corner.corner(
    results['samples'],
    labels=parameter_names,
    color="darkorange",
    plot_datapoints=False,
    fill_contours=False,
    smooth=1.0,
    contour_kwargs={"linestyles": "dashed"},
);

_images/example-simulation-based-inference_81_0.svg

Combined posterior distribution

Comparing them on the same plot, the SBI posterior is extremely wide, while ultranest manages to zoom in and get the full information.

[40]:

figure = corner.corner(
    posterior_samples_theta,
    labels=parameter_names,
    truths=theta_true,
    truth_color="red",
    show_titles=True,
    quantiles=[0.16, 0.5, 0.84],
    color="steelblue",
    plot_datapoints=False,
    fill_contours=True,
    smooth=1.0,
)
# overplot nested sampling posterior on the same figure
corner.corner(
    results['samples'],
    labels=parameter_names,
    fig=figure,
    color="darkorange",
    plot_datapoints=False,
    fill_contours=False,
    smooth=1.0,
    contour_kwargs={"linestyles": "dashed"},
);

_images/example-simulation-based-inference_84_0.svg

Notice that the ultranest-sbi posterior (orange) is much narrower than the SBI posterior (blue), because SBI lost a lot of information. Nevertheless, the blue distribution provides a useful starting point for nested sampling.

Takeaways

In addition to the usual model and its likelihood, we defined a data simulator that can generate data.
We have learned to train a simulation-based inference (SBI) neural posterior estimator (NPE) model for obtaining a fast posterior distribution.
We checked that SBI posterior is amortized to behave well across the entire prior.
We used the SBI posterior on our data to get a starting point for a posterior distribution. This distribution has lost some information.
We used ultranest on top of the SBI posterior to get the posterior samples from the true posterior.
This ultranest-sbi approach effectively “skips” many initial nested sampling iterations, zooming into a data-driven guess for the posterior location. Because the SBI posterior is virtually always too wide, nested sampling is then fast to integrate to the actual posterior bulk.

Limitations: The posterior distribution used by the neural network is factorized and not optimal for challenging distributions. With a different, more sophisticated sbi package, better neural posteriors could be learned, but that would mean more complexity as well.

When using this in publications (such as legend entries), please refer to it as ultranest-sbi, to distinguish it from ultranest.

https://arxiv.org/abs/1804.06788, “Validating Bayesian Inference Algorithms with Simulation-Based Calibration”, Talts e al., 2018
https://doi.org/10.1073/pnas.1912789117 “The frontier of simulation-based inference”, Cranmer, Brehmer, Louppe, 2020