This page was generated from notebooks/Examples - 2. Truncation.ipynb.

# Truncation¶

When the prior is very wide and simulation expense is high, it makes sense to focus our simulations on a certain observation $$b_o$$. We are effectively estimating the likelihood-to-evidence ratio on a small region around the $$\theta_o$$ which produced $$x_o$$. We do this marginally, therefore we take the product of marginal estimates and let that be our truncated region on which to estimate the likelihood-to-evidence ratio. This notebook demonstrates that technique.

:

%load_ext autoreload

:

# DON'T FORGET TO ACTIVATE THE GPU when on google colab (Edit > Notebook settings)
from os import environ
GOOGLE_COLAB = True if "COLAB_GPU" in environ else False
!pip install git+https://github.com/undark-lab/swyft.git

:

import numpy as np
import pylab as plt
import torch

import swyft

:

device = 'cuda' if swyft.utils.is_cuda_available() else "cpu"
n_training_samples = 3000
extra_training_samples_factor = 1.10  # number of drawn samples is stochastic, we draw more than we need by this factor.
n_parameters = 2
marginal_indices_1d, marginal_indices_2d = swyft.utils.get_corner_marginal_indices(n_parameters)
observation_key = "x"

n_posterior_samples_for_truncation = 10_000
n_weighted_samples = 10_000

:

def model(v, sigma = 0.01):
x = v + np.random.randn(n_parameters)*sigma
return {observation_key: x}

v_o = np.zeros(n_parameters)
observation_o = model(v_o, sigma = 0.)

n_observation_features = observation_o[observation_key].shape
observation_shapes = {key: value.shape for key, value in observation_o.items()}

:

simulator = swyft.Simulator(
model,
n_parameters,
sim_shapes=observation_shapes,
)

low = -1 * np.ones(n_parameters)
high = 1 * np.ones(n_parameters)
prior = swyft.get_uniform_prior(low, high)

store = swyft.Store.memory_store(simulator)
# drawing samples from the store is Poisson distributed. Simulating slightly more than we need avoids attempting to draw more than we have.
store.add(extra_training_samples_factor * n_training_samples, prior)
store.simulate()

Creating new store.
Store: Adding 3296 new samples to simulator store.


## creating a do_round function¶

We call the process of training the marginal likelihood-to-evidence ratio estimator and estimating the support of the truncated prior a round. The output of a round is a bound object which is then used in the next round. It makes sense to encapsulate your round in a function which can be called repeatedly.

We will start by truncating in one dimension (with hyperrectangles as bounds)

:

def do_round_1d(bound, observation_focus):
store.add(extra_training_samples_factor * n_training_samples, prior, bound=bound)
store.simulate()

dataset = swyft.Dataset(n_training_samples, prior, store, bound = bound)

network_1d = swyft.get_marginal_classifier(
observation_key=observation_key,
marginal_indices=marginal_indices_1d,
observation_shapes=observation_shapes,
n_parameters=n_parameters,
hidden_features=32,
num_blocks=2,
)
mre_1d = swyft.MarginalRatioEstimator(
marginal_indices=marginal_indices_1d,
network=network_1d,
device=device,
)
mre_1d.train(dataset)
posterior_1d = swyft.MarginalPosterior(mre_1d, prior, bound)
new_bound = posterior_1d.truncate(n_posterior_samples_for_truncation, observation_focus)
return posterior_1d, new_bound


## Truncating and estimating the marginal posterior over the truncated region¶

First we train the one-dimensional likelihood-to-evidence ratios, then we train the two-dimensional estimator on the truncated region. This region is defined by the bound object.

:

bound = None
for i in range(3):
posterior_1d, bound = do_round_1d(bound, observation_o)

training: lr=5e-05, epoch=25, validation loss=0.29923
Store: Adding 1872 new samples to simulator store.
training: lr=0.0005, epoch=25, validation loss=0.5188
Store: Adding 2356 new samples to simulator store.
training: lr=0.0005, epoch=25, validation loss=0.6468

:

network_2d = swyft.get_marginal_classifier(
observation_key=observation_key,
marginal_indices=marginal_indices_2d,
observation_shapes=observation_shapes,
n_parameters=n_parameters,
hidden_features=32,
num_blocks=2,
)
mre_2d = swyft.MarginalRatioEstimator(
marginal_indices=marginal_indices_2d,
network=network_2d,
device=device,
)

:

store.add(extra_training_samples_factor * n_training_samples, prior, bound=bound)
store.simulate()
dataset = swyft.Dataset(n_training_samples, prior, store, bound = bound)

Store: Adding 68 new samples to simulator store.

:

_ = mre_2d.train(dataset)

training: lr=5e-05, epoch=25, validation loss=0.096541

:

weighted_samples_1d = posterior_1d.weighted_sample(n_weighted_samples, observation_o)

posterior_2d = swyft.MarginalPosterior(mre_2d, prior, bound)
weighted_samples_2d = posterior_2d.weighted_sample(n_weighted_samples, observation_o)

:

_, _ = swyft.plot.corner(
weighted_samples_1d,
weighted_samples_2d,
kde=True,
truth=v_o,
xlim=[-0.15, 0.15],
ylim_lower=[-0.15, 0.15],
bins=200,
)

1it [00:04,  4.59s/it] ## Repeat but truncate in two dimensions¶

First we train the two-dimensional likelihood-to-evidence ratio. We use that bound to estimate the one-dimensional likelihood-to-evidence ratios.

:

def do_round_2d(bound, observation_focus):
store.add(extra_training_samples_factor * n_training_samples, prior, bound=bound)
store.simulate()

dataset = swyft.Dataset(n_training_samples, prior, store, bound = bound)

network_2d = swyft.get_marginal_classifier(
observation_key=observation_key,
marginal_indices=marginal_indices_2d,
observation_shapes=observation_shapes,
n_parameters=n_parameters,
hidden_features=32,
num_blocks=2,
)
mre_2d = swyft.MarginalRatioEstimator(
marginal_indices=marginal_indices_2d,
network=network_2d,
device=device,
)
mre_2d.train(dataset)

posterior_2d = swyft.MarginalPosterior(mre_2d, prior, bound)
new_bound = posterior_1d.truncate(n_posterior_samples_for_truncation, observation_focus)

return posterior_2d, new_bound

:

bound = None
for i in range(3):
posterior_2d, bound = do_round_2d(bound, observation_o)

training: lr=5e-05, epoch=25, validation loss=0.0023252
Store: Adding 16 new samples to simulator store.
training: lr=5e-05, epoch=25, validation loss=0.044781
Store: Adding 10 new samples to simulator store.
training: lr=5e-05, epoch=25, validation loss=0.09438

:

network_1d = swyft.get_marginal_classifier(
observation_key=observation_key,
marginal_indices=marginal_indices_1d,
observation_shapes=observation_shapes,
n_parameters=n_parameters,
hidden_features=32,
num_blocks=2,
)
mre_1d = swyft.MarginalRatioEstimator(
marginal_indices=marginal_indices_1d,
network=network_1d,
device=device,
)
_ = mre_1d.train(dataset)

training: lr=0.0005, epoch=25, validation loss=0.7137

:

store.add(extra_training_samples_factor * n_training_samples, prior, bound=bound)
store.simulate()
dataset = swyft.Dataset(n_training_samples, prior, store, bound = bound)

Store: Adding 3 new samples to simulator store.

:

_ = mre_1d.train(dataset)

training: lr=5e-08, epoch=25, validation loss=0.7239

:

posterior_1d = swyft.MarginalPosterior(mre_1d, prior, bound)
weighted_samples_1d = posterior_1d.weighted_sample(n_weighted_samples, observation_o)

weighted_samples_2d = posterior_2d.weighted_sample(n_weighted_samples, observation_o)

:

_, _ = swyft.plot.corner(
weighted_samples_1d,
weighted_samples_2d,
kde=True,
truth=v_o,
xlim=[-0.15, 0.15],
ylim_lower=[-0.15, 0.15],
bins=200,
)

1it [00:04,  4.41s/it] [ ]:




This page was generated from notebooks/Examples - 2. Truncation.ipynb.