Comparison of sampling algorithms for calibration of a line with unknown model error#

from collections import OrderedDict

import corner

import numpy as np

from matplotlib import pyplot as plt
from scipy import stats

import rxmc

Using database version X4-2025-12-31 located in: /mnt/home/beyerkyl/x4db/unpack_exfor-2025/X4-2025-12-31

true_params = OrderedDict(
    [
        ("m", 2),
        ("b", 4),
    ]
)

rng = np.random.default_rng(43)

noise = 0.05
N = 12
# x_data = rng.random(N)
x_data = np.linspace(0, 1, N)
y_true = true_params["m"] * x_data + true_params["b"]
y_err = np.array([rng.normal(0, noise * y) for y in y_true])
y_data = y_true + y_err

# Define priors: (mean, variance) for normal, (alpha, beta) for inverse gamma
reported_stat_err = noise * np.mean(y_data)
m_prior = (2, 10)
b_prior = (1, 10)
sigma2_prior = (3, reported_stat_err / (3 - 1))

class LinearModel(rxmc.physical_model.PhysicalModel):
    def __init__(self):
        params = [
            rxmc.params.Parameter("m", float, "no-units"),
            rxmc.params.Parameter("b", float, "y-units"),
        ]
        super().__init__(params)

    def evaluate(self, observation, m, b):
        return self.y(observation.x, m, b)

    def y(self, x, m, b):
        # useful to have a function hat takes in an array-like x
        # rather than an Observation, e.g. for plotting
        return m * x + b

prior_mean = OrderedDict(
    [
        ("m", 1),
        ("b", 5),
    ]
)
prior_std_dev = OrderedDict(
    [
        ("m", 0.5),
        ("b", 0.5),
    ]
)

covariance = np.diag(list(prior_std_dev.values())) ** 2
mean = np.array(list(prior_mean.values()))
prior_distribution = stats.multivariate_normal(mean, covariance)

my_model = LinearModel()

observation = rxmc.observation.Observation(
    x=x_data,
    y=y_data,
    y_stat_err=y_true * noise,
)

plt.errorbar(
    x_data,
    y_data,
    noise * y_data,
    color="k",
    marker="o",
    linestyle="none",
    label="experiment with bias",
)
plt.plot(x_data, y_true, "k", label="truth")

[<matplotlib.lines.Line2D at 0x14f21fc9ed50>]

../_images/7b9e1411b972a5090f3c13dcc2adba325b79ae852b6fa2e3e9baf7fe369b9e30.png

likelihood = rxmc.likelihood_model.UnknownNoiseFractionErrorModel()

evidence = rxmc.evidence.Evidence(
    parametric_constraints=[
        rxmc.constraint.Constraint(
            [observation],
            my_model,
            likelihood,
        )
    ]
)

Sampling configurations for the model parameters#

def proposal_distribution(x, rng):
    return stats.multivariate_normal.rvs(
        mean=x, cov=prior_distribution.cov / 100, random_state=rng
    )


metropolis_model = rxmc.param_sampling.MetropolisHastingsSampler(
    params=my_model.params,
    starting_location=prior_distribution.mean,
    proposal=proposal_distribution,
    prior=prior_distribution,
)

adaptive_model = rxmc.param_sampling.BatchedAdaptiveMetropolisSampler(
    params=my_model.params,
    starting_location=prior_distribution.mean,
    prior=prior_distribution,
    initial_proposal_cov=prior_distribution.cov / 1000,
)

Sampling configurations for the likelihood parameters#

noise_prior = stats.norm(loc=np.log(0.01), scale=1)

def proposal_distribution_log_noise(x, rng):
    return np.atleast_1d(stats.norm.rvs(loc=x, scale=0.1, random_state=rng))


metropolis_likelihood = rxmc.param_sampling.MetropolisHastingsSampler(
    params=likelihood.params,
    starting_location=np.array(noise_prior.mean()),
    proposal=proposal_distribution_log_noise,
    prior=noise_prior,
)

adaptive_likelihood = rxmc.param_sampling.BatchedAdaptiveMetropolisSampler(
    params=likelihood.params,
    starting_location=np.array(noise_prior.mean()),
    prior=noise_prior,
    initial_proposal_cov=np.array([[1]]),
)

Walkers#

walker_metropolis = rxmc.walker.Walker(
    metropolis_model,
    evidence,
    likelihood_samplers=[metropolis_likelihood],
)
walker_adaptive = rxmc.walker.Walker(
    adaptive_model,
    evidence,
    likelihood_samplers=[adaptive_likelihood],
)

%%time
walker_adaptive.walk(n_steps=5000, burnin=1000, batch_size=100)

Burn-in batch 1/10 completed, 100 steps.
Burn-in batch 2/10 completed, 100 steps.
Burn-in batch 3/10 completed, 100 steps.
Burn-in batch 4/10 completed, 100 steps.
Burn-in batch 5/10 completed, 100 steps.
Burn-in batch 6/10 completed, 100 steps.
Burn-in batch 7/10 completed, 100 steps.
Burn-in batch 8/10 completed, 100 steps.
Burn-in batch 9/10 completed, 100 steps.
Burn-in batch 10/10 completed, 100 steps.
Batch: 1/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.310
  Likelihood parameter acceptance fractions: [0.46]
Batch: 2/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.310
  Likelihood parameter acceptance fractions: [0.53]
Batch: 3/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.650
  Likelihood parameter acceptance fractions: [0.49]
Batch: 4/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.390
  Likelihood parameter acceptance fractions: [0.39]
Batch: 5/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.320
  Likelihood parameter acceptance fractions: [0.44]
Batch: 6/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.330
  Likelihood parameter acceptance fractions: [0.43]
Batch: 7/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.290
  Likelihood parameter acceptance fractions: [0.44]
Batch: 8/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.600
  Likelihood parameter acceptance fractions: [0.42]
Batch: 9/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.380
  Likelihood parameter acceptance fractions: [0.51]
Batch: 10/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.290
  Likelihood parameter acceptance fractions: [0.52]
Batch: 11/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.440
  Likelihood parameter acceptance fractions: [0.39]
Batch: 12/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.500
  Likelihood parameter acceptance fractions: [0.26]
Batch: 13/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.330
  Likelihood parameter acceptance fractions: [0.47]
Batch: 14/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.410
  Likelihood parameter acceptance fractions: [0.48]
Batch: 15/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.440
  Likelihood parameter acceptance fractions: [0.41]
Batch: 16/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.380
  Likelihood parameter acceptance fractions: [0.44]
Batch: 17/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.190
  Likelihood parameter acceptance fractions: [0.46]
Batch: 18/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.450
  Likelihood parameter acceptance fractions: [0.34]
Batch: 19/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.350
  Likelihood parameter acceptance fractions: [0.49]
Batch: 20/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.370
  Likelihood parameter acceptance fractions: [0.5]
Batch: 21/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.450
  Likelihood parameter acceptance fractions: [0.52]
Batch: 22/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.380
  Likelihood parameter acceptance fractions: [0.44]
Batch: 23/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.240
  Likelihood parameter acceptance fractions: [0.48]
Batch: 24/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.330
  Likelihood parameter acceptance fractions: [0.5]
Batch: 25/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.470
  Likelihood parameter acceptance fractions: [0.34]
Batch: 26/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.380
  Likelihood parameter acceptance fractions: [0.51]
Batch: 27/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.310
  Likelihood parameter acceptance fractions: [0.43]
Batch: 28/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.540
  Likelihood parameter acceptance fractions: [0.45]
Batch: 29/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.290
  Likelihood parameter acceptance fractions: [0.45]
Batch: 30/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.350
  Likelihood parameter acceptance fractions: [0.42]
Batch: 31/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.390
  Likelihood parameter acceptance fractions: [0.37]
Batch: 32/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.410
  Likelihood parameter acceptance fractions: [0.52]
Batch: 33/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.450
  Likelihood parameter acceptance fractions: [0.46]
Batch: 34/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.440
  Likelihood parameter acceptance fractions: [0.49]
Batch: 35/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.310
  Likelihood parameter acceptance fractions: [0.42]
Batch: 36/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.430
  Likelihood parameter acceptance fractions: [0.51]
Batch: 37/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.310
  Likelihood parameter acceptance fractions: [0.58]
Batch: 38/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.420
  Likelihood parameter acceptance fractions: [0.46]
Batch: 39/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.260
  Likelihood parameter acceptance fractions: [0.59]
Batch: 40/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.360
  Likelihood parameter acceptance fractions: [0.38]
Batch: 41/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.550
  Likelihood parameter acceptance fractions: [0.38]
Batch: 42/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.200
  Likelihood parameter acceptance fractions: [0.44]
Batch: 43/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.610
  Likelihood parameter acceptance fractions: [0.43]
Batch: 44/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.430
  Likelihood parameter acceptance fractions: [0.38]
Batch: 45/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.160
  Likelihood parameter acceptance fractions: [0.32]
Batch: 46/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.330
  Likelihood parameter acceptance fractions: [0.56]
Batch: 47/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.410
  Likelihood parameter acceptance fractions: [0.37]
Batch: 48/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.320
  Likelihood parameter acceptance fractions: [0.36]
Batch: 49/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.230
  Likelihood parameter acceptance fractions: [0.4]
Batch: 50/50 completed, 100 steps. 
  Model parameter acceptance fraction: 0.390
  Likelihood parameter acceptance fractions: [0.43]
CPU times: user 4.8 s, sys: 224 ms, total: 5.02 s
Wall time: 4.78 s

%%time
walker_metropolis.walk(n_steps=20000, burnin=1000, batch_size=1000)

Burn-in batch 1/1 completed, 1000 steps.
Batch: 1/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.670
  Likelihood parameter acceptance fractions: [0.825]
Batch: 2/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.750
  Likelihood parameter acceptance fractions: [0.825]
Batch: 3/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.799
  Likelihood parameter acceptance fractions: [0.839]
Batch: 4/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.870
  Likelihood parameter acceptance fractions: [0.83]
Batch: 5/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.725
  Likelihood parameter acceptance fractions: [0.827]
Batch: 6/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.797
  Likelihood parameter acceptance fractions: [0.839]
Batch: 7/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.690
  Likelihood parameter acceptance fractions: [0.839]
Batch: 8/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.614
  Likelihood parameter acceptance fractions: [0.847]
Batch: 9/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.651
  Likelihood parameter acceptance fractions: [0.835]
Batch: 10/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.784
  Likelihood parameter acceptance fractions: [0.834]
Batch: 11/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.722
  Likelihood parameter acceptance fractions: [0.82]
Batch: 12/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.821
  Likelihood parameter acceptance fractions: [0.816]
Batch: 13/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.766
  Likelihood parameter acceptance fractions: [0.819]
Batch: 14/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.637
  Likelihood parameter acceptance fractions: [0.815]
Batch: 15/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.768
  Likelihood parameter acceptance fractions: [0.839]
Batch: 16/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.642
  Likelihood parameter acceptance fractions: [0.851]
Batch: 17/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.621
  Likelihood parameter acceptance fractions: [0.833]
Batch: 18/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.646
  Likelihood parameter acceptance fractions: [0.838]
Batch: 19/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.753
  Likelihood parameter acceptance fractions: [0.849]
Batch: 20/20 completed, 1000 steps. 
  Model parameter acceptance fraction: 0.718
  Likelihood parameter acceptance fractions: [0.834]
CPU times: user 13.2 s, sys: 130 ms, total: 13.4 s
Wall time: 13.4 s

def plot_chains(walker, model, true_params):
    fig, axes = plt.subplots(
        walker.model_sampler.chain.shape[1] + 3, 1, figsize=(8, 8), sharex=True
    )
    for i in range(walker.model_sampler.chain.shape[1]):
        axes[i].plot(walker.model_sampler.chain[:, i])
        axes[i].set_ylabel(f"${model.params[i].latex_name}$ [{model.params[i].unit}]")
        true_value = true_params[model.params[i].name]
        axes[i].hlines(
            true_value, 0, len(walker.model_sampler.chain), "r", linestyle="--"
        )

    axes[-3].plot(walker.model_sampler.logp_chain)
    axes[-3].set_ylabel(r"$\log{\mathcal{L}(\alpha_i | \mathcal{O})}$")

    lmp = walker.likelihood_samplers[0].params[0]
    axes[-2].plot(walker.likelihood_samplers[0].chain)
    axes[-2].set_ylabel(f"${lmp.latex_name}$ [{lmp.unit}]")
    axes[-2].hlines(
        np.log(noise), 0, len(walker.likelihood_samplers[0].chain), "r", linestyle="--"
    )

    axes[-1].plot(walker.likelihood_samplers[0].logp_chain)
    axes[-1].set_ylabel(r"$\log{\mathcal{L}(\alpha_i | \mathcal{O})}$")

    axes[-1].set_xlabel(r"$i$")
    # plt.legend(title="chains", ncol=3,)

plot_chains(walker_metropolis, my_model, true_params)

../_images/4ddd2646529084607f7dc0aaf3b893e4f19302d3eca62a0b2b21ae8af93c099c.png

plot_chains(walker_adaptive, my_model, true_params)

../_images/bae8ea6f2b314fd16199e59bd2a31cdb7e0feeee2f72cc798b51a5d5767ec620.png

fig = corner.corner(
    np.hstack(
        [
            walker_metropolis.model_sampler.chain,
            walker_metropolis.likelihood_samplers[0].chain,
        ]
    ),
    labels=[p.name for p in my_model.params]
    + [walker_metropolis.likelihood_samplers[0].params[0].name],
    label="posterior",
    truths=[true_params["m"], true_params["b"], np.log(noise)],
    color="tab:blue",
)
corner.corner(
    np.hstack(
        [
            walker_adaptive.model_sampler.chain,
            walker_adaptive.likelihood_samplers[0].chain,
        ]
    ),
    fig=fig,
    color="tab:orange",
)

fig.suptitle("posterior")

Text(0.5, 0.98, 'posterior')

../_images/fcb3c7af93da3fc7ea3b2d5add56e233ca3586607fc427a2cbb82d40db6a40dd.png

x = np.linspace(-0.5, 1.5, 10)
n_posterior_samples = walker_metropolis.model_sampler.chain.shape[0]
y1 = np.zeros((n_posterior_samples, len(x)))
for i in range(n_posterior_samples):
    sample = walker_metropolis.model_sampler.chain[i, :]
    y1[i, :] = my_model.y(x, *sample)

n_posterior_samples = walker_adaptive.model_sampler.chain.shape[0]
y2 = np.zeros((n_posterior_samples, len(x)))
for i in range(n_posterior_samples):
    sample = walker_adaptive.model_sampler.chain[i, :]
    y2[i, :] = my_model.y(x, *sample)

# Metropolis Hastings
upper, med, lower = np.percentile(y1, [5, 50, 95], axis=0)
p = plt.fill_between(
    x, lower, upper, alpha=0.3, zorder=2, color="tab:blue", label="Metropolis-Hastings"
)
plt.plot(x, med, ":", color=p.get_facecolor())

# adaptive
upper, med, lower = np.percentile(y2, [5, 50, 95], axis=0)
p = plt.fill_between(
    x,
    lower,
    upper,
    alpha=0.3,
    zorder=2,
    color="tab:orange",
    label="Adaptive Metropolis",
)
plt.plot(x, med, ":", color=p.get_facecolor())

# truth
plt.plot(x, my_model.y(x, *true_params.values()), "k--", label="truth")
plt.errorbar(
    x_data,
    y_data,
    noise * y_data,
    color="k",
    marker="o",
    linestyle="none",
    label="experiment with bias",
)

plt.xlabel(r"$x$")
plt.ylabel(r"$y$")
plt.legend()

<matplotlib.legend.Legend at 0x14f21fca6120>

../_images/6af579a5d4718c3825e33d8d17c3eea428b2ffcc391dede0bee29b102c236c62.png