Field-level inference: Hamiltonian Monte Carlo for non-linear models¶

import numpy as np
import os
import jax
import jax.numpy as jnp
from cycler import cycler
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import LogNorm, SymLogNorm
from mpl_toolkits.axes_grid1 import make_axes_locatable
np.random.seed(123456)

plt.rcParams.update({'lines.linewidth': 2})
plt.rcParams.update({'text.usetex': True})
plt.rcParams.update({'text.latex.preamble': r"\usepackage{amsmath}\usepackage{upgreek}"})
plt.rcParams.update({'font.family': 'serif'})
plt.rcParams.update({'font.size': 15})
dir="./plots/HMC_nonlinear_model/"
os.makedirs(dir, exist_ok=True)

# Download and define the Planck color map
from matplotlib.colors import ListedColormap
import os
if not os.path.isfile("data/Planck_Parchment_RGB.txt"):
    !wget https://raw.githubusercontent.com/zonca/paperplots/master/data/Planck_Parchment_RGB.txt --directory-prefix=data/
planck = ListedColormap(np.loadtxt("data/Planck_Parchment_RGB.txt")/255.)
planck.set_bad("C7") # color of missing pixels

N = 32
L = 1.0 # box size
A_s  = 6e-9 # power spectrum normalisation, arbitrary units
n_s  = 0.96 # scalar spectral index
f_NL = 2000.0 # non-linear coupling parameter
D1 = 1.732e7 # growth factor of fluctuations at z=0, arbitrary units

def build_power_spectrum(N, A_s=1.0, n_s=0.96):
    # real‐space grid spacing
    dx = L / N

    # build k‐space grid for rfft2
    kx = jnp.fft.fftfreq(N, d=dx) * 2*np.pi         # length N
    ky = jnp.fft.rfftfreq(N, d=dx)  * 2*np.pi       # length N//2+1
    kx, kx = jnp.meshgrid(kx, ky, indexing='ij')    # shape (N, N//2+1)
    k = jnp.sqrt(kx**2 + ky**2)
    k = jnp.where(k == 0, 1.0, k)  # prevent division by zero

    # build Pk safely (zero at k=0)
    Pk = A_s * k**(n_s - 1.0)
    Pk = jnp.where(k == 0, 0, Pk)  # Set mean mode to zero power
    return Pk
Pkgrid = build_power_spectrum(N, A_s, n_s)

@jax.jit
def phiL_from_real_noise(white_noise):
    """
    Given real-space white noise, construct a Gaussian random field with power spectrum
    P(k) = A_s * k^(n_s - 1).
    Args:
        white_noise: real array of shape [N, N] (from jax.random.normal(key, (N, N)))
        A_s, n_s: parameters of power spectrum
    Returns:
        field: real array [N, N]
    """
    N = white_noise.shape[0]
    noise_k = jnp.fft.rfft2(white_noise)
    # Multiply by sqrt of power spectrum
    filtered_k = noise_k * jnp.sqrt(Pkgrid)
    # IFFT back to real space
    field = jnp.fft.irfft2(filtered_k).real
    return field

def build_transfer_function(N):
    """Transfer function from primordial potential to density contrast."""
    # real‐space grid spacing
    dx = L / N

    # build k‐space grid for rfft2
    kx = jnp.fft.fftfreq(N, d=dx) * 2*np.pi         # length N
    ky = jnp.fft.rfftfreq(N, d=dx)  * 2*np.pi       # length N//2+1
    kx, ky = jnp.meshgrid(kx, ky, indexing='ij')    # shape (N, N//2+1)
    k = jnp.sqrt(kx**2 + ky**2)
    k = jnp.where(k == 0, 1.0, k)  # prevent division by zero

    # build Tk safely (zero at k=0)
    T = np.zeros_like(k)

    # shape parameter, characterising the wavenumber at equality
    Omega_b = 0.049
    Omega_m = 0.315
    h = 0.674
    fb = Omega_b / Omega_m
    shape = Omega_m * h * np.exp(-Omega_b - np.sqrt(2.0*h) * fb)

    # power spectrum, BBKS style
    q = k / shape
    alpha = 2.34; beta = 3.89; gamma = 16.1; delta = 5.46; epsilon = 6.71
    aux = 1.0 + beta * q + (gamma * q)**2 + (delta * q)**3 + (epsilon * q)**4
    T = (jnp.log(1.0 + alpha * q) / (alpha * q) * aux**-0.25);
    T *= q**(1./2.) # correct scaling for delta(k) = D1 * T(k) * Phi_NL(k) since we want P_delta(k) \propto k^n_s T(k)**2 \propto  k * T(k)**2 * P_phi(k)
    T = jnp.where(k == 0, 0, T)  # Set mean mode to zero transfer function

    return T
Tgrid = build_transfer_function(N)

@jax.jit
def delta_from_phi(phi, f_NL=f_NL, D1=D1):
    """
    Given real-space Phi, compute Phi_NL and then delta such that
    delta(k) = D1 * T(k)**2 * Phi_NL(k), and inverse FFT to real-space delta.

    Args:
        phi: [N, N] real array, the field Phi in real space
        f_NL: scalar
    Returns:
        delta_real: [N, N] real array, delta in real space
    """
    N = phi.shape[0]
    # 1. Construct Phi_NL
    Phi_NL = phi + f_NL * phi**2

    # 2. FFT to Fourier space
    PhiNL_k = jnp.fft.rfft2(Phi_NL)

    # 3. Multiply by transfer function T(k)
    delta_k = D1 * Tgrid * PhiNL_k
    delta_k = delta_k.at[0, 0].set(0.0)  # Explicitly zero the mean mode

    # 4. IFFT back to real space
    delta_real = jnp.fft.irfft2(delta_k).real

    return delta_real

@jax.jit
def data_model(white_noise, f_NL=f_NL, D1=D1):
    """
    Given real-space white noise, construct a Gaussian random field with power spectrum
    P(k) = A_s * k^(n_s - 1), and compute the corresponding delta field.

    Args:
        white_noise: real array of shape [N, N] (from jax.random.normal(key, (N, N)))
        f_NL: scalar
        D1: scalar
    Returns:
        delta_real: [N, N] real array, delta in real space
    """
    phi = phiL_from_real_noise(white_noise)
    delta = delta_from_phi(phi, f_NL=f_NL, D1=D1)
    return delta

key = jax.random.PRNGKey(12)
white_noise = jax.random.normal(key, (N, N))
phi = phiL_from_real_noise(white_noise)
delta = delta_from_phi(phi)

def make_noise_field():
    """
    Build a 32×32 array “field” of noise variances:

     - Border of width 3 pixels everywhere: value = 1e0
     - Central 26×26 block: two regions of different noise variances
        - A low‐noise region (1e-5)
        - A medium‐noise region (5e-5) in the lower left corner
     - A 2×2 high‐noise patch (1e0) centered in the array

    Returns
    -------
    field : ndarray, shape (32,32)
        The resulting noise‐variance field.
    """
    N = 32
    field = np.zeros((N, N), dtype=float)
    mask = np.zeros((N, N), dtype=bool)

    # 1) High‐noise border
    bw = 3
    high = 1
    field[:bw, :]   = high
    field[-bw:, :]  = high
    field[:, :bw]   = high
    field[:, -bw:]  = high
    mask[:bw, :]   = True
    mask[-bw:, :]  = True
    mask[:, :bw]   = True
    mask[:, -bw:]  = True

    # 2) Central block (size 26×26)
    low  = 1e-5
    med  = 5e-5
    for i in range(bw,N-bw):
        for j in range(bw,N-bw):
            if i-j<-3:
                field[i,j] = med
            else:
                field[i,j] = low

    # 3) Sprinkle a 2×2 high‐noise patch
    ci = 24
    cj = 24
    field[ci-1:ci+1, cj-1:cj+1] = high
    mask[ci-1:ci+1, cj-1:cj+1] = True

    return field, mask

noise_variance_field, mask = make_noise_field()
invN = np.diag(1.0 / noise_variance_field.flatten())
invN[np.where(invN <= 1)] = 0.0

# noise: one realisation
noise = np.random.normal(size=(N, N)) * np.sqrt(noise_variance_field)
noise_v = np.ma.masked_where(mask, noise) # for visualization purposes only

# data = delta + noise
data = delta + noise
data_v = np.ma.masked_where(mask, data)  # for visualization purposes only

def log_prior(signal):
    """
    Compute the log-prior of the signal (white noise field).
    Args:
        signal: [N, N] array, the signal (white noise field)
    Returns:
        log_prior: scalar, the log-prior value
    """
    # Assuming a Gaussian prior with zero mean and unit variance
    log_prior = -0.5 * jnp.sum(signal**2)
    return log_prior

def log_likelihood(signal, data, noise_variance_field):
    """
    Compute the (unnormalised) log-likelihood of the data given the signal and noise variance field.

    Args:
        data: [N, N] array, the observed data
        signal: [N, N] array, the signal (white noise field)
        noise_variance_field: [N, N] array, the noise variance field

    Returns:
        log_likelihood: scalar, the log-likelihood value
    """
    delta = data_model(signal)
    residual = data - delta
    log_likelihood = -0.5 * jnp.sum(residual**2 / noise_variance_field)
    return log_likelihood

def log_posterior(signal, data, noise_variance_field):
    """
    Compute the (unnormalised) log-posterior of the signal given the data and noise variance field.

    Args:
        data: [N, N] array, the observed data
        signal: [N, N] array, the signal (white noise field)
        noise_variance_field: [N, N] array, the noise variance field

    Returns:
        log_posterior: scalar, the log-posterior value
    """
    log_likelihood_value = log_likelihood(signal, data, noise_variance_field)
    log_prior_value = log_prior(signal)
    return log_likelihood_value + log_prior_value

# Compute the gradients of log_prior, log_likelihood, and log_posterior w.r.t. white noise using JAX autodiff
def d_log_prior_d_white_noise(white_noise):
    """
    Compute the gradient of `log_prior` w.r.t. all elements of `white_noise`.

    Args:
        white_noise: [N, N] array

    Returns:
        grad: [N, N] array, the gradient of log_prior w.r.t. white_noise
    """
    grad = -white_noise  # Gradient of Gaussian prior with zero mean
    return grad

def d_log_likelihood_d_white_noise_autodiff(white_noise, data, noise_variance_field):
    """
    Compute the gradient of `log_likelihood` w.r.t. all elements of `white_noise` using JAX autodiff.
    This function uses JAX's automatic differentiation to compute the gradient efficiently.

    Args:
        white_noise: [N, N] array
        data: [N, N] array
        noise_variance_field: [N, N] array

    Returns:
        grad: [N, N] array, the gradient of log_likelihood w.r.t. white_noise
    """
    grad = jax.grad(log_likelihood)(white_noise, data, noise_variance_field)  # shape (N,N)
    return grad

def d_log_posterior_d_white_noise_autodiff(white_noise, data, noise_variance_field):
    """
    Compute the gradient of `log_posterior` w.r.t. all elements of `white_noise` using JAX autodiff.
    This function uses JAX's automatic differentiation to compute the gradient efficiently.

    Args:
        white_noise: [N, N] array
        data: [N, N] array
        noise_variance_field: [N, N] array

    Returns:
        grad: [N, N] array, the gradient of log_posterior w.r.t. white_noise
    """
    grad = jax.grad(log_posterior)(white_noise, data, noise_variance_field)  # shape (N,N)
    return grad

# Compute the gradients of log_prior, log_likelihood, and log_posterior w.r.t. white noise using finite differences
def d_log_prior_d_white_noise_fd(white_noise, epsilon=1e-3):
    """
    Compute the gradient of `log_prior` w.r.t. all elements of `white_noise` using finite differences.

    Args:
        white_noise: [N, N] array
        epsilon: finite difference step size

    Returns:
        grad_fd: [N, N] array, finite difference gradient
    """
    N = white_noise.shape[0]
    grad_fd = jnp.zeros_like(white_noise)

    for i in range(N):
        for j in range(N):
            white_noise_up   = white_noise.at[i, j].add(+epsilon)
            white_noise_down = white_noise.at[i, j].add(-epsilon)
            f_up   = log_prior(white_noise_up)
            f_down = log_prior(white_noise_down)
            grad_fd = grad_fd.at[i, j].set((f_up - f_down) / (2 * epsilon))
    return grad_fd

def d_log_likelihood_d_white_noise_fd(white_noise, data, noise_variance_field, epsilon=1e-3):
    """
    Compute the gradient of `log_likelihood` w.r.t. all elements of `white_noise` using finite differences.

    Args:
        white_noise: [N, N] array
        data: [N, N] array
        noise_variance_field: [N, N] array
        epsilon: finite difference step size

    Returns:
        grad_fd: [N, N] array, finite difference gradient
    """
    N = white_noise.shape[0]
    grad_fd = jnp.zeros_like(white_noise)

    for i in range(N):
        for j in range(N):
            white_noise_up   = white_noise.at[i, j].add(+epsilon)
            white_noise_down = white_noise.at[i, j].add(-epsilon)
            f_up   = log_likelihood(white_noise_up,   data, noise_variance_field)
            f_down = log_likelihood(white_noise_down, data, noise_variance_field)
            grad_fd = grad_fd.at[i, j].set((f_up - f_down) / (2 * epsilon))
    return grad_fd

def d_log_posterior_d_white_noise_fd(white_noise, data, noise_variance_field, epsilon=1e-3):
    """
    Compute the gradient of `log_posterior` w.r.t. all elements of `white_noise` using finite differences.

    Args:
        white_noise: [N, N] array
        data: [N, N] array
        noise_variance_field: [N, N] array
        epsilon: finite difference step size

    Returns:
        grad_fd: [N, N] array, finite difference gradient
    """
    N = white_noise.shape[0]
    grad_fd = jnp.zeros_like(white_noise)

    for i in range(N):
        for j in range(N):
            white_noise_up   = white_noise.at[i, j].add(+epsilon)
            white_noise_down = white_noise.at[i, j].add(-epsilon)
            f_up   = log_posterior(white_noise_up,   data, noise_variance_field)
            f_down = log_posterior(white_noise_down, data, noise_variance_field)
            grad_fd = grad_fd.at[i, j].set((f_up - f_down) / (2 * epsilon))
    return grad_fd

grad_prior = d_log_prior_d_white_noise(white_noise)
grad_prior_fd = d_log_prior_d_white_noise_fd(white_noise)

grad_lh = d_log_likelihood_d_white_noise_autodiff(white_noise, data, noise_variance_field)
grad_lh_fd = d_log_likelihood_d_white_noise_fd(white_noise, data, noise_variance_field)

grad_post = d_log_posterior_d_white_noise_autodiff(white_noise, data, noise_variance_field)
grad_post_fd = d_log_posterior_d_white_noise_fd(white_noise, data, noise_variance_field)

import blackjax
from typing import Any, Tuple

def sample_white_noise_hmc(
    white_noise_init: jnp.ndarray,
    data: jnp.ndarray,
    noise_variance_field: jnp.ndarray,
    n_samples: int,
    n_adapt: int = 1000,  # Number of adaptation steps
    rng_key: jax.random.PRNGKey = None,
) -> Tuple[jnp.ndarray, Any]:
    if rng_key is None:
        rng_key = jax.random.PRNGKey(0)

    N = white_noise_init.shape[0]

    def logprob(flat_white_noise):
        white_noise = flat_white_noise.reshape((N, N))
        return log_posterior(white_noise, data, noise_variance_field)

    initial_position = jnp.ravel(white_noise_init)

    # Set up the adaptation routine for HMC
    adapt = blackjax.window_adaptation(
        blackjax.hmc,
        logprob,
        num_steps=n_adapt,  # Number of adaptation steps here
        num_integration_steps=10,  # Number of leapfrog steps
        target_acceptance_rate=0.8,  # Target acceptance rate
    )

    # Run the adaptation phase
    state, kernel, adaptation_state = adapt.run(rng_key, initial_position)

    # Sampling loop
    def inference_loop(rng_key, kernel, initial_state, num_samples):
        @jax.jit
        def one_step(state, rng_key):
            state, _ = kernel(rng_key, state)
            return state, state

        keys = jax.random.split(rng_key, num_samples)
        _, states = jax.lax.scan(one_step, initial_state, keys)

        return states

    # Run the sampling
    states = inference_loop(rng_key, kernel, state, n_samples)

    # Reshape the samples to the original field shape
    samples = states.position.reshape(n_samples, N, N)

    return samples, states

n_samples = 10000  # Number of samples to draw
N_chains = 5  # Number of chains to run
initial_scaling = 0.001  # Initial scaling for the white noise field, start from an overdispersed state

try:
    samples_chain = np.load('data/HMC_nonlinear_model/samples_chain.npy', allow_pickle=True).item()
except FileNotFoundError:
    samples_chain = {}
    for c in range(N_chains):
        rng_key = jax.random.PRNGKey(42 + c)
        field_init = jax.random.normal(rng_key, (N, N)) * initial_scaling
        samples_chain[c], infos = sample_white_noise_hmc(
            field_init,
            data,
            noise_variance_field,
            n_samples=n_samples,
            n_adapt=10,
            rng_key=rng_key
        )
    os.makedirs('data/HMC_nonlinear_model/', exist_ok=True)
    np.save('data/HMC_nonlinear_model/samples_chain.npy', samples_chain)

log_likelihoods_chain = {}
for c in range(N_chains):
    log_likelihoods_chain[c] = jax.vmap(log_likelihood, in_axes=(0, None, None))(samples_chain[c], data, noise_variance_field)

# Define a function to compute the power spectrum of the sampled delta fields

@jax.jit
def power_spectrum_2d_jitted(field, nbins=200, L=1.0):
    # real‐space grid spacing
    N = field.shape[0]
    dx = L / N

    # build k‐space grid for rfft2
    kx = jnp.fft.fftfreq(N, d=dx) * 2 * np.pi  # length N
    ky = jnp.fft.rfftfreq(N, d=dx) * 2 * np.pi  # length N//2+1
    kx, ky = jnp.meshgrid(kx, ky, indexing='ij')  # shape (N, N//2+1)
    k = jnp.sqrt(kx**2 + ky**2)
    kmag = k.flatten()

    # Compute rfft2
    ft = jnp.fft.rfft2(field)
    power2d = (jnp.abs(ft) ** 2) * (dx ** 2) / (N ** 2)
    power_flat = power2d.flatten()

    # Bin edges (exclude k=0 when setting min)
    k_nonzero = kmag[1:] if kmag.size > 1 else kmag  # Remove k=0 for k_min
    k_min = k_nonzero.min()
    k_max = kmag.max()
    bins = jnp.logspace(jnp.log10(k_min), jnp.log10(k_max), nbins + 1)
    bin_idx = jnp.digitize(kmag, bins) - 1  # bin assignments

    # Bin means
    Pk_sum = jnp.bincount(bin_idx, weights=power_flat, length=nbins)
    counts = jnp.bincount(bin_idx, length=nbins)
    Pk = jnp.where(counts > 0, Pk_sum / counts, 0.0)
    k_sum = jnp.bincount(bin_idx, weights=kmag, length=nbins)
    k_bin = jnp.where(counts > 0, k_sum / counts, 0.0)

    mask = counts > 0
    return k_bin, Pk, counts

def power_spectrum_2d(field, nbins=200, L=1.0):
    """
    Compute the 2D power spectrum of a 2D field using FFT and binning.

    Args:
        field: [N, N] array, the 2D field
        nbins: int, number of bins for the power spectrum
        L: float, box size

    Returns:
        k_bin: [nbins] array, the binned wavenumbers
        Pk: [nbins] array, the binned power spectrum values
        counts: [nbins] array, number of modes in each bin
    """
    k_bin, Pk, counts = power_spectrum_2d_jitted(field)
    mask = counts > 0
    return k_bin[mask], Pk[mask]

k_vals, Pk_signal_gt = power_spectrum_2d(white_noise, nbins=200)
k_vals, Pk_delta_gt = power_spectrum_2d(delta, nbins=200)

# Compute the power spectrum of the signal (white noise)
Pk = np.zeros((Nburnin, len(k_vals)))
for i in range(0, Nburnin, 50):
    k_vals, Pk[i] = power_spectrum_2d(samples_chain[0][i])

# Reconstruct the delta field from the sampled white noise fields
delta_samples_chain1 = jax.vmap(data_model)(samples_chain[0])

# Compute the power spectrum of the delta field samples
Pk_delta = np.zeros((Nburnin, len(k_vals)))
for i in range(0, Nburnin, 50):
    k_vals, Pk_delta[i] = power_spectrum_2d(delta_samples_chain1[i])

# Based on python code from the emcee tutorials, https://emcee.readthedocs.io/en/stable/tutorials/autocorr/

def next_power_of_2(n: int) -> int:
    """Smallest power of two ≥ n."""
    return 1 << (n - 1).bit_length() if n > 0 else 1

# 1D autocorrelation function
def autocorr_func_1d(x, norm=True):
    """
    Compute the 1D autocorrelation via FFT in O(N log N).
    If norm=True, normalize so acf[0] = 1.
    """
    x = np.asarray(x, dtype=float)
    n = x.size
    nfft = 2 * next_power_of_2(n)

    # real FFT
    f = np.fft.rfft(x - np.mean(x), n=nfft)
    ps = (f * f.conjugate()).real      # power spectrum
    acf = np.fft.irfft(ps, n=nfft)[:n]
    acf /= 2 * nfft

    # normalise
    if norm:
        if acf[0] <= 0:
            return 0 # or raise ValueError("Autocorrelation function is zero or negative at lag 0.")
        else:
            acf /= acf[0]
    return acf

# Automated windowing procedure following Sokal (1989)
def auto_window(taus, c):
    """
    Return the first lag k for which k < c * tau_k fails.
    If none fail, return len(taus)-1.
    """
    k = np.arange(len(taus))
    mask = k < c * taus
    # find first index where mask is False
    idx = np.argmax(~mask)
    return idx if mask[idx] == False else len(taus) - 1

# Following the suggestion from Goodman & Weare (2010)
def autocorr_gw2010(x, c: float = 5.0) -> float:
    """
    Estimate the integrated autocorrelation time following
    Goodman & Weare (2010), with window parameter c.
    """
    acf = autocorr_func_1d(x, norm=True)
    taus = 2.0 * np.cumsum(acf) - 1.0
    window = auto_window(taus, c)
    return taus[window]

def N_eff(x) -> float:
    """
    Effective number of independent samples in x.
    Accepts input as either a list or a 1D numpy array.
    """
    x = np.asarray(x)
    tau = autocorr_gw2010(x)
    if tau <= 0:
        return 0 # or raise ValueError("Autocorrelation time is zero or negative.")
    return x.size / tau

n_thin_corr = 200
N_eff_10_20 = np.zeros((N_chains, n_samples // n_thin_corr), dtype=float)
N_eff_25_10 = np.zeros((N_chains, n_samples // n_thin_corr), dtype=float)
for c in range(N_chains):
    for i, samp in enumerate(np.arange(n_thin_corr, n_samples + 1, n_thin_corr)):
        # Use up-to samp (not samp:) and start with at least n_thin_corr samples
        N_eff_10_20[c, i] = N_eff(samples_chain[c][:samp, 10, 20])
        N_eff_25_10[c, i] = N_eff(samples_chain[c][:samp, 25, 10])

def gelman_rubin(chain):
    # between chains variance
    Psi_dotj = np.mean(chain, axis=1)
    Psi_dotdot = np.mean(Psi_dotj, axis=0)
    m = chain.shape[0]
    n = chain.shape[1]
    B = n / (m - 1.) * np.sum((Psi_dotj - Psi_dotdot)**2, axis=0)

    # within chains variance
    sj2 = np.var(chain, axis=1, ddof=1)
    W = np.mean(sj2, axis=0)

    # estimators
    var_minus = W
    var_plus = (n - 1.) / n * W + 1. / n * B
    R_hat = np.sqrt(var_plus / var_minus)
    return R_hat

# The input array must have dimensions (nchains, nsamp, npars) = (m, n, npars).

chain = jnp.array([samples_chain[c] for c in range(N_chains)])
chain = chain.reshape((N_chains, n_samples, N*N))  # Reshape to (nchains, nsamp, npars)
Rstat = gelman_rubin(chain).reshape((N, N))

# Join the different chains into one array for the posterior samples
samples = jnp.concatenate([samples_chain[c][Nburnin:] for c in range(N_chains)], axis=0)
PhiL = phiL_from_real_noise(white_noise)
PhiL_samples = jax.vmap(phiL_from_real_noise)(samples)
PhiNL = PhiL + f_NL * PhiL**2
PhiNL_samples = PhiL_samples + f_NL * PhiL_samples**2
delta_samples = jax.vmap(data_model)(samples)

@jax.jit
def compute_summaries(samples):
    empirical_mean = jnp.mean(samples, axis=0)
    empirical_var = jnp.var(samples, axis=0)
    return empirical_mean, empirical_var
empirical_mean, empirical_var = compute_summaries(samples)

PhiL_mean, PhiL_var = compute_summaries(PhiL_samples)
PhiNL_mean, PhiNL_var = compute_summaries(PhiNL_samples)
delta_mean, delta_var = compute_summaries(delta_samples)