Simulation Based Inference: Exercises

Credits
AstroStat Academy logo This notebook contains original work by the authors unless stated otherwise. Any external material is properly credited to its sources.
References to papers, datasets, and software are acknowledged. Original content is licensed under the GNU General Public License v3.0 (GNU GPLv3).

Simulation Based Inference: Exercises#

# --- Compatibility patch for GetDist + NumPy 2.0 ---
# NumPy 2.0 removed np.infty, but older versions of GetDist still use it.
# This patch restores it at runtime before importing GetDist.

import numpy as np
if not hasattr(np, 'infty'):
    np.infty = np.inf  # Restore deprecated alias for compatibility

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader, random_split
import numpy as np
import emcee
from scipy.stats import poisson
from scipy.stats import chi2
from scipy.optimize import basinhopping
from tqdm import tqdm
import matplotlib.pyplot as plt
import pytorch_lightning as pl
from scipy.stats import norm
from getdist import plots, MCSamples
from scipy.integrate import quad
import multiprocessing
from multiprocessing import Pool
from nflows.flows.base import Flow
from nflows.distributions.normal import StandardNormal
from nflows.transforms.base import CompositeTransform
from nflows.transforms.autoregressive import MaskedAffineAutoregressiveTransform
from nflows.transforms.permutations import ReversePermutation
from sbi.inference import SNPE  # old name used in v0.22.0
from sbi.utils.torchutils import BoxUniform

%matplotlib inline

Introduction#

In the main notebook, our mock observations were generated by simply adding some Gaussian noise to the theoretical data. However, in a real experiment, our observations may be subject to various other effects, which may be caused by the instruments, the environment, physical processes, etc. These effects are often very hard to model with a simple analytical function, and may introduce additional biases in the data.

In this exercise, we will assume that our observations are contaminated by two new effects: random instrumental drift and intermittent calibration errors.

  1. Random instrumental drift: The instrument that we use to measure the data is not perfect, and its calibration may drift over time due to environmental factors or aging hardware. This drift will accumulate over time and so the measurements that we take at different times may acquire a random offset. We can model this effect by adding a cumulative sum of normally distributed random values to the true magnitudes.

    Why? [Click to expand]

    This simulates the effect of a drift that changes the measurements in a continuous but random manner. The random values are drawn from a normal distribution with a mean of zero, ensuring that the drift can go in either direction (positive or negative) and has a variance that reflects the magnitude of the drift per observation.

  2. Intermittent calibration errors: The instrument may also suffer from occasional large calibration errors that occur at random intervals. This may be caused by a variety of factors, such as power surges, cosmic rays, mechanical disturbances, miscalibrations during the data collection process etc. We model these errors by introducing occasional large spikes in the observed magnitudes.

#random seed
np.random.seed(1235)
torch.manual_seed(1234)

# Constants
c = 299792.458  # Speed of light in km/s

# True cosmological parameters
true_Omega_m = 0.3
true_H0 = 70.0  # Hubble constant in km/s/Mpc

# Redshifts of the supernovae
z = np.linspace(0.01, 1.0, 50)

# Noise level in the observed magnitudes
sigma = 0.1

# Luminosity distance function
def luminosity_distance(z, Omega_m, H0):
    """ Calculate the luminosity distance for a given redshift, Omega_m and H0 """
    integrand = lambda z_prime: 1.0 / np.sqrt(Omega_m * (1 + z_prime)**3 + (1 - Omega_m))
    d_L = np.array([quad(integrand, 0, z_i)[0] for z_i in z])
    return (c * (1 + z) * d_L) / H0

# Generate synthetic data
# Generate synthetic data with non-linear transformation
d_L_true = luminosity_distance(z, true_Omega_m, true_H0)
m_true = 5 * np.log10(d_L_true) + 19.3

# Introduce random walk drift
def random_walk_drift(m, z):
    drift = np.cumsum(np.random.normal(0, 0.2, len(m)))
    return m + drift

# Introduce intermittent calibration errors
def intermittent_calibration_errors(m, z):
    errors = np.zeros_like(m)
    error_indices = np.random.choice(len(m), size=int(len(m) * 0.1), replace=False)
    errors[error_indices] = np.random.normal(0, 1, len(error_indices))
    return m + errors

# Apply both effects to the observed magnitudes
m_obs_drift = random_walk_drift(m_true, z)
m_obs = intermittent_calibration_errors(m_obs_drift, z) + np.random.normal(0, sigma, len(m_true))


# Plot synthetic data
plt.errorbar(z, m_obs, yerr=0.5, fmt='.', label='Observed magnitudes')
plt.plot(z, m_true, label='True magnitudes')
plt.xlabel('Redshift $z$')
plt.ylabel('Apparent magnitude $m$')
plt.legend()
plt.show()

Part 1: Solving with Gaussian likelihood#

Problem: Given the observed magnitudes \(m_{\text{obs}}\) and redshifts \(z\), infer the cosmological parameters \(\Omega_m\) and \(H_0\) using the traditional Bayesian approach assuming a Gaussian likelihood:

\[ \Large \mathcal{L}(\Omega_m, H_0) = \prod_{i=1}^N \frac{1}{\sqrt{2\pi\sigma_i^2}} \exp\left(-\frac{(m_{\text{obs},i} - m_{\text{th},i})^2}{2\sigma_i^2}\right), \tag{3} \]

Tasks:

  1. Define the prior, a Gaussian likelihood with some \(\sigma\) and the posterior.

  2. Run MCMC to sample from the posterior.

  3. Plot the corner plot of the posterior samples and see if the results are consistent with the true values.

(Worked solutions are in the Solutions section at the end of this notebook.)

Click for tip

Do exactly what was done in the main notebook:

  • For the prior, use a flat prior over \(\Omega_m\) and \(H_0\).

  • For the Gaussian likelihood, you can assume a \(\sigma\) looking at the data (I would assume \(\sigma=0.5\), as in the main notebook).

  • For MCMC, you can use the emcee package.

Task 1:

# Log-prior function
def log_prior(params):
    Omega_m, H0 = params
    if ...
        return ...
    else:
        return ...
    
# Define the log-likelihood function
def log_like(params, z, m_obs):
    Omega_m, H0 = params
    d_L = ...
    m_th = ...
    sigma_m = ...  # our assumed standard deviation of magnitudes
    return ... 

# Log-posterior function
def log_post(params, z, m_obs):
    lp = log_prior(...)
    if not np.isfinite(lp):
        return -np.inf # If prior is -inf, return directly -inf
    else:
        return lp + log_like(params, z, m_obs)

Task 2:

# Sampling with `emcee` 
ndim, nwalkers = ... , ...
ncpus = multiprocessing.cpu_count()

# Initialize walkers 
pos = ...

with Pool(ncpus) as pool:
    sampler = emcee.EnsembleSampler(nwalkers, ndim, log_post, args=(z, m_obs), pool=pool)
    sampler_output = sampler.run_mcmc(pos, nsteps=..., progress=True)

Task 3:

# Get samples and convert to GetDist MCSamples
flat_samples = sampler.get_chain(discard=..., flat=True)

# GetDist plotting
samples = MCSamples(samples=flat_samples, names=['\Omega_m', 'H_0'], labels=['\Omega_m', 'H_0'])
g = plots.get_subplot_plotter(width_inch=4)
g.triangle_plot(samples, filled=True, markers={'\Omega_m': true_Omega_m, 'H_0': true_H0})
plt.show()

Part 2: Solving with Neural Posterior Estimation#

Problem: Given the observed magnitudes \(m_{\text{obs}}\) and redshifts \(z\), infer the cosmological parameters \(\Omega_m\) and \(H_0\) using the Neural Posterior Estimation approach.

Tasks:

  1. Construct a simulator that generates the observed magnitudes \(m_{\text{obs}}\) given the cosmological parameters \(\Omega_m\) and \(H_0\). Encode the random instrumental drifts and intermittent calibration errors in the simulator.

  2. Generate a training dataset of observed magnitudes \(m_{\text{obs}}\) and corresponding cosmological parameters \(\Omega_m\) and \(H_0\), by running the simulator for a range of randomly drawn parameters \(\Omega_m\) and \(H_0\).

  3. Define the Neural Posterior Estimator class using the sbi package that will estimate the posterior with a normalizing flow.

  4. Sample from the NPE, conditioning on the observed magnitudes \(m_{\text{obs}}\) that we had in the beginning.

  5. Plot the corner plot of the posterior samples and see if the results are consistent with the true values.

(Worked solutions are in the Solutions section at the end of this notebook.)

Click for tip

Again, do exactly what was done in the main notebook. The actual changes that you need to implement are:

  • In the simulator, introduce the random instrumental drifts and intermittent calibration errors.

  • For the test data, use instead the observed magnitudes \(m_{\text{obs}}\) that we showed in the beginning.

Task 1:

# Forward model to simulate observed magnitudes
def luminosity_distance_simulator(params, z):
    Omega_m, H0 = params
    d_L = ...
    m_th = ...
    m_drift = ... # add the random walk drift here
    m_drift_calibration = ... # add the intermittent calibration errors here
    m_sim = ... # add the standard noise here, to get the final simulated magnitudes
    
    return ...

Task 2:

np.random.seed(1234) # use random seed for reproducibility

# Generate training data
n_train = 20_000
param_samples = np.random.uniform(low=[0, 50], high=[1, 100], size=(n_train, 2)) # Parameter proposal
x_samples = np.array([luminosity_distance_simulator(theta, z) for theta in tqdm(param_samples)])

# Convert to tensors
param_samples = ...
x_samples = ...

Task 3:

prior = ... 
inference = ...
inference.append_simulations(...)
density_estimator = ...
posterior = ...

Tasks 4 & 5:

samples_npe = ...

# Convert to numpy (ensure on CPU and detached from graph)
if isinstance(samples_npe, torch.Tensor):
    samples_npe = samples_npe.detach().cpu().numpy()
else:
    samples_npe = np.asarray(samples_npe)


# Create a getdist MCSamples object
gd_samples_npe = MCSamples(samples=samples_npe, labels=['\Omega_m', 'H_0'], names=['\Omega_m', 'H_0'])

# Make a triangle/contour plot with getdist
g = plots.get_subplot_plotter()
g.triangle_plot([gd_samples_npe, samples], filled=True, markers={'\Omega_m': true_Omega_m, 'H_0': true_H0}, legend_labels=['NPE', 'MCMC'])

🔑 Solutions#

End of the exercise — solutions below. Try the exercises yourself first, then check your work here. (These cells only need the imports and the data-generation cell from the top of the notebook.)

Part 1, Task 1 — Prior, likelihood, and posterior#

# Log-prior function
def log_prior(params):
    Omega_m, H0 = params
    if 0 < Omega_m < 1 and 50 < H0 < 100:
        return 0
    else:
        return -np.inf
    
# Define the log-likelihood function
def log_like(params, z, m_obs):
    Omega_m, H0 = params
    d_L = luminosity_distance(z, Omega_m, H0)
    m_th = 5 * np.log10(d_L) + 19.3
    sigma_m = 0.5  # Assumed standard deviation of magnitudes
    return -0.5 * np.sum(((m_obs - m_th) / sigma_m)**2)    

# Log-posterior function
def log_post(params, z, m_obs):
    lp = log_prior(params)
    if not np.isfinite(lp):
        return -np.inf
    else:
        return lp + log_like(params, z, m_obs)

Part 1, Task 2 — MCMC sampling#

# Sampling with `emcee` 
ndim, nwalkers = 2, 16
ncpus = multiprocessing.cpu_count()

# Initialize walkers around a random position withing the prior volume
pos = np.random.uniform([0, 50], [1, 100], (nwalkers, ndim))

sampler = emcee.EnsembleSampler(nwalkers, ndim, log_post, args=(z, m_obs))
sampler.run_mcmc(pos, 2000, progress=True)

Part 1, Task 3 — Corner plot#

# Get samples and convert to GetDist MCSamples
flat_samples = sampler.get_chain(discard=1000, flat=True)

# GetDist plotting
samples = MCSamples(samples=flat_samples, names=['\Omega_m', 'H_0'], labels=['\Omega_m', 'H_0'])
g = plots.get_subplot_plotter(width_inch=4)
g.triangle_plot(samples, filled=True, markers={'\Omega_m': true_Omega_m, 'H_0': true_H0})
plt.show()

Part 2, Task 1 — The simulator#

# Forward model to simulate observed magnitudes
def luminosity_distance_simulator(params, z):
    Omega_m, H0 = params
    d_L = luminosity_distance(z, Omega_m, H0)
    m_th = 5 * np.log10(d_L) + 19.3
    m_sim_drift = random_walk_drift(m_th, z)
    m_sim = intermittent_calibration_errors(m_sim_drift, z) + np.random.normal(0, sigma, len(m_th))

    return m_sim

Part 2, Task 2 — Training data#

np.random.seed(1234) # use random seed for reproducibility

# Generate training data
n_train = 20_000
param_samples = np.random.uniform(low=[0, 50], high=[1, 100], size=(n_train, 2)) # Parameter proposal
x_samples = np.array([luminosity_distance_simulator(theta, z) for theta in tqdm(param_samples)])

# Convert to tensors
param_samples = torch.tensor(param_samples, dtype=torch.float32)
x_samples = torch.tensor(x_samples, dtype=torch.float32)

Part 2, Task 3 — NPE class#

# Define the prior, and the SBI model and give it parameters theta and corresponding simulations x
prior = BoxUniform(low=torch.tensor([0.0, 50.0]), high=torch.tensor([1.0, 100.0]))

inference = SNPE(prior=prior)

inference.append_simulations(param_samples, x_samples)

density_estimator = inference.train()

posterior = inference.build_posterior(density_estimator)

Part 2, Tasks 4 & 5 — Sample and plot#

samples_npe = posterior.sample((10000,), x=m_obs)

# Convert to numpy (ensure on CPU and detached from graph)
if isinstance(samples_npe, torch.Tensor):
    samples_npe = samples_npe.detach().cpu().numpy()
else:
    samples_npe = np.asarray(samples_npe)


# Create a getdist MCSamples object
gd_samples_npe = MCSamples(samples=samples_npe, labels=['\Omega_m', 'H_0'], names=['\Omega_m', 'H_0'])

# Make a triangle/contour plot with getdist
g = plots.get_subplot_plotter()
g.triangle_plot([gd_samples_npe, samples], filled=True, markers={'\Omega_m': true_Omega_m, 'H_0': true_H0}, legend_labels=['NPE', 'MCMC'])