<jupyter_start><jupyter_text>Quickstart<jupyter_code>%load_ext autoreload
%autoreload 2
from os import environ
GOOGLE_COLAB = True if "COLAB_GPU" in environ else False
if GOOGLE_COLAB:
    !pip install git+https://github.com/undark-lab/swyft.git
import numpy as np
import pylab as plt
import swyft<jupyter_output><empty_output><jupyter_text>Setup the forward model<jupyter_code>device = "cpu"
n_training_samples = 10_000
n_parameters = 2
marginal_indices_1d, marginal_indices_2d = swyft.utils.get_corner_marginal_indices(n_parameters)
observation_key = "x"

n_weighted_samples = 10_000
def model(v, sigma = 0.2):
    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[0]
observation_shapes = {key: value.shape for key, value in observation_o.items()}
simulator = swyft.Simulator(
    model, 
    n_parameters, 
    sim_shapes=observation_shapes
)<jupyter_output><empty_output><jupyter_text>Setup the prior and storage<jupyter_code>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)
store.add(n_training_samples, prior)
store.simulate()
dataset = swyft.Dataset(n_training_samples, prior, store)<jupyter_output><empty_output><jupyter_text>Train a 1d marginal estimator<jupyter_code>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)<jupyter_output>training: lr=5e-06, epoch=25, validation loss=1.6534<jupyter_text>Create a simple violin plot and a row of histograms to view the 1d marginals<jupyter_code>n_rejection_samples = 5_000

posterior_1d = swyft.MarginalPosterior(mre_1d, prior)
samples_1d = posterior_1d.sample(n_rejection_samples, observation_o)
_ = swyft.violin(samples_1d)
_, _ = swyft.hist1d(samples_1d, kde=True)<jupyter_output><empty_output><jupyter_text>Train a 2d marginal estimator<jupyter_code>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)<jupyter_output>training: lr=5e-07, epoch=25, validation loss=0.4701<jupyter_text>Combine the two to create a corner plot<jupyter_code>posterior_1d = swyft.MarginalPosterior(mre_1d, prior)
weighted_samples_1d = posterior_1d.weighted_sample(n_weighted_samples, observation_o)

posterior_2d = swyft.MarginalPosterior(mre_2d, prior)
weighted_samples_2d = posterior_2d.weighted_sample(n_weighted_samples, observation_o)
_, _ = swyft.corner(
    weighted_samples_1d, 
    weighted_samples_2d,
    kde=True,
    truth=v_o,
    labels=["x1", "x2"]
)<jupyter_output>1it [00:04,  4.01s/it]<jupyter_text>Test the expected coverage probability<jupyter_code>n_observations = 200
n_posterior_samples = 5_000

empirical_mass_1d, _ = posterior_1d.empirical_mass(n_observations, n_posterior_samples, dataset)
empirical_mass_2d, _ = posterior_2d.empirical_mass(n_observations, n_posterior_samples, dataset)
_, _ = swyft.empirical_z_score_corner(empirical_mass_1d, empirical_mass_2d, figsize=(10,10))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.0918394828979262, 'fraction_numerical': 0.029356315647724212, 'mean_word_length': 7.290178571428571, 'pattern_counts': {'":': 0, '<': 30, '<?xml version=': 0, '>': 30, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '281662', 'license': 'MIT', 'n_tokens_mistral': 1501, 'n_tokens_neox': 1430, 'n_words': 212, 'path': 'notebooks/Quickstart.ipynb', 'repo_name': 'undark-lab/swyft', 'size': '3713'}