Bayesian calibration of polynomials, including inference of overall normalization

import corner

import numpy as np
from matplotlib import pyplot as plt
from scipy import stats

import rxmc

Using database version X4-2024-12-31 located in: /home/kyle/db/exfor/unpack_exfor-2024/X4-2024-12-31

rng = np.random.default_rng(42)

poly4 = rxmc.physical_model.Polynomial(4)

true_params = [1, 0.5, -0.1, -0.4, 0.1]

settings = [
    {
        "domain": [-0.3, 0.4],
        "N": 50,
        "noise": 0.1,
        "systematic_err": 0.1,
    },
    {
        "domain": [0.3, 0.5],
        "N": 30,
        "noise": 0.1,
        "systematic_err": 0.5,
    },
    {
        "domain": [0.1, 0.6],
        "N": 25,
        "noise": 0.1,
        "systematic_err": 0.2,
    },
    {
        "domain": [-0.5, 0.1],
        "N": 15,
        "noise": 0.2,
        "systematic_err": 0.6,
    },
]

def generate_observations(settings, true_model, rng, true_params, scale_err=True):
    obs = []
    for setting in settings:
        x0, x1 = setting["domain"]
        synthetic_obs = rxmc.observation.Observation(
            x=rng.random(setting["N"]) * (x1 - x0) + x0,
            y=np.zeros(setting["N"]),
            y_stat_err=np.ones(setting["N"]) * setting["noise"],
        )
        renormalization = rng.normal(1, setting["systematic_err"])
        y_true = true_model(synthetic_obs, *true_params)
        if scale_err:
            synthetic_obs.y = rng.normal(y_true, setting["noise"]) * renormalization
        else:
            synthetic_obs.y = rng.normal(y_true * renormalization, setting["noise"])

        synthetic_obs.renormalization = renormalization
        synthetic_obs.sys_norm_err = setting["systematic_err"]

        obs.append(synthetic_obs)
    return obs

observations = generate_observations(settings, poly4, rng, true_params, scale_err=False)

observations_unreported_sys_err = [
    rxmc.observation.Observation(
        x=obs.x,
        y=obs.y,
        y_stat_err=obs.y_stat_err,
    )
    for obs in observations
]

N_fine = 100
domain_fine = (-1, 1)
truth = rxmc.observation.Observation(
    x=np.linspace(*domain_fine, N_fine), y=np.zeros(N_fine)
)
truth.y = poly4(truth, *true_params)

plt.plot(truth.x, truth.y, "k:", label="truth")
for synthetic_observation in observations:
    plt.errorbar(
        synthetic_observation.x,
        synthetic_observation.y,
        synthetic_observation.y_stat_err,
        linestyle="none",
        marker=".",
        label=f"renormalization = {synthetic_observation.renormalization:1.3f}",
    )
plt.xlabel("x")
plt.ylabel("y")
plt.legend()

<matplotlib.legend.Legend at 0x780de1416090>

correct_model = poly4

# block supports come from the library helper
_block_supports = rxmc.covariance.stacked_supports


evidence_models = {}

# Unknown per-dataset normalization: a latent scale rho_i for each dataset, routed
# to it by identity and sampled as ordinary *model* parameters.  This is the v2
# form of the old per-dataset UnknownNormalizationModel -- rho moves onto the
# PhysicalModel (via rxmc.transforms.per_observation_scaling) and is sampled in
# the model block (jointly with the physics) rather than as a per-constraint
# covariance nuisance.  All constraints share the one model instance, so they
# still share a single model-parameter vector.
norm_model = rxmc.physical_model.Polynomial(
    order=correct_model.order,
    transform=rxmc.transforms.per_observation_scaling(observations),
)
evidence_models["unknown_norm"] = rxmc.evidence.Evidence(
    [rxmc.constraint.Constraint([obs], norm_model) for obs in observations]
)

# reported systematic errors folded in as fixed per-dataset normalization modes
evidence_models["marginalized_sys_err"] = rxmc.evidence.Evidence(
    [
        rxmc.constraint.Constraint(
            observations,
            correct_model,
            extra_terms=[
                rxmc.covariance.normalization_term(
                    magnitude=o.sys_norm_err, support=sup
                )
                for o, sup in zip(observations, _block_supports(observations))
            ],
        )
    ]
)

# systematic errors omitted entirely (statistical only)
evidence_models["unreported_sys_err"] = rxmc.evidence.Evidence(
    [rxmc.constraint.Constraint(observations_unreported_sys_err, correct_model)]
)

# omitted systematics absorbed by a single unknown (averaging) model-error term
gamma = rxmc.params.Parameter(
    "log fractional err", float, latex_name=r"\gamma", unit="dimensionless"
)
_N_unreported = sum(o.n_data_pts for o in observations_unreported_sys_err)
evidence_models["unreported_sys_err_with_unknown_model_err"] = rxmc.evidence.Evidence(
    [
        rxmc.constraint.Constraint(
            observations_unreported_sys_err,
            correct_model,
            extra_terms=[
                rxmc.covariance.model_error_term(
                    gamma,
                    averaging=True,
                    support=np.arange(_N_unreported),
                )
            ],
        )
    ]
)

evidence_models.keys()

dict_keys(['unknown_norm', 'marginalized_sys_err', 'unreported_sys_err', 'unreported_sys_err_with_unknown_model_err'])

print("hyperparameters for each model:")
hyperparams = {}
for key, evidence_model in evidence_models.items():
    hyperparams[key] = []
    for constraint in evidence_model.parametric_constraints:
        for p in constraint.params:
            hyperparams[key].append(p)
    print(f"{key}: {[p.name for p in hyperparams[key]]}")

hyperparameters for each model:
unknown_norm: []
marginalized_sys_err: []
unreported_sys_err: []
unreported_sys_err_with_unknown_model_err: ['log fractional err']

Priors

We will put a tight prior on $a_0$, to avoid identifiability issues, as $a_0$ corresponds to an additive offset of the entire model, which is pretty confounding with a multiplicative bias.

cov = np.diag(np.ones(len(correct_model.params)))
cov[0, 0] = 0.001
model_prior = stats.multivariate_normal(mean=np.array([1, 0, 0, 0, 0]), cov=cov)

unknown_log_norm_priors = [
    stats.multivariate_normal(
        mean=[0], cov=[[np.log(1 + (setting["systematic_err"]) ** 2)]]
    )
    for setting in settings
]

unknown_model_err_prior = stats.multivariate_normal(mean=[np.log(0.1)], cov=[[0.01]])

Walkers

walkers = {}

# combined prior over [physics params, per-dataset log_rho] (block-diagonal):
# the per-dataset scales are now model parameters, so they are sampled in the
# model block -- this walker has no separate likelihood samplers.
n_physics = len(correct_model.params)
rho_prior_vars = np.array(
    [np.log(1 + setting["systematic_err"] ** 2) for setting in settings]
)
norm_prior_mean = np.concatenate([model_prior.mean, np.zeros(len(observations))])
norm_prior_cov = np.zeros((n_physics + len(observations),) * 2)
norm_prior_cov[:n_physics, :n_physics] = model_prior.cov
norm_prior_cov[n_physics:, n_physics:] = np.diag(rho_prior_vars)
norm_model_prior = stats.multivariate_normal(mean=norm_prior_mean, cov=norm_prior_cov)

walkers["unknown_norm"] = rxmc.walker.Walker(
    model_sampler=rxmc.param_sampling.BatchedAdaptiveMetropolisSampler(
        norm_model.params,
        starting_location=norm_model_prior.mean,
        prior=norm_model_prior,
        initial_proposal_cov=norm_model_prior.cov,
    ),
    evidence=evidence_models["unknown_norm"],
)

walkers["unreported_sys_err_with_unknown_model_err"] = rxmc.walker.Walker(
    model_sampler=rxmc.param_sampling.BatchedAdaptiveMetropolisSampler(
        correct_model.params,
        starting_location=model_prior.mean,
        prior=model_prior,
        initial_proposal_cov=model_prior.cov,
    ),
    evidence=evidence_models["unreported_sys_err_with_unknown_model_err"],
    likelihood_samplers=[
        rxmc.param_sampling.BatchedAdaptiveMetropolisSampler(
            params=hyperparams["unreported_sys_err_with_unknown_model_err"],
            starting_location=unknown_model_err_prior.mean,
            prior=unknown_model_err_prior,
            initial_proposal_cov=unknown_model_err_prior.cov,
        )
    ],
)

walkers["marginalized_sys_err"] = rxmc.walker.Walker(
    model_sampler=rxmc.param_sampling.BatchedAdaptiveMetropolisSampler(
        correct_model.params,
        starting_location=model_prior.mean,
        prior=model_prior,
        initial_proposal_cov=model_prior.cov,
    ),
    evidence=evidence_models["marginalized_sys_err"],
)

walkers["unreported_sys_err"] = rxmc.walker.Walker(
    model_sampler=rxmc.param_sampling.BatchedAdaptiveMetropolisSampler(
        correct_model.params,
        starting_location=model_prior.mean,
        prior=model_prior,
        initial_proposal_cov=model_prior.cov,
    ),
    evidence=evidence_models["unreported_sys_err"],
)

for key, walker in walkers.items():
    print(f"\nRunning {key}")
    walker.walk(n_steps=20000, burnin=8000, batch_size=2000)

Running unknown_norm

Burn-in batch 1/4 completed, 2000 steps.

Burn-in batch 2/4 completed, 2000 steps.

Burn-in batch 3/4 completed, 2000 steps.

Burn-in batch 4/4 completed, 2000 steps.

Batch: 1/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.247

Batch: 2/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.306

Batch: 3/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.266

Batch: 4/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.277

Batch: 5/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.294

Batch: 6/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.286

Batch: 7/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.276

Batch: 8/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.248

Batch: 9/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.292

Batch: 10/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.288

Running unreported_sys_err_with_unknown_model_err

Burn-in batch 1/4 completed, 2000 steps.

Burn-in batch 2/4 completed, 2000 steps.

Burn-in batch 3/4 completed, 2000 steps.

Burn-in batch 4/4 completed, 2000 steps.

Batch: 1/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.283
  Likelihood parameter acceptance fractions: [0.423]

Batch: 2/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.281
  Likelihood parameter acceptance fractions: [0.451]

Batch: 3/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.291
  Likelihood parameter acceptance fractions: [0.4215]

Batch: 4/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.302
  Likelihood parameter acceptance fractions: [0.455]

Batch: 5/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.311
  Likelihood parameter acceptance fractions: [0.436]

Batch: 6/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.290
  Likelihood parameter acceptance fractions: [0.4425]

Batch: 7/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.278
  Likelihood parameter acceptance fractions: [0.4445]

Batch: 8/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.267
  Likelihood parameter acceptance fractions: [0.474]

Batch: 9/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.305
  Likelihood parameter acceptance fractions: [0.4405]

Batch: 10/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.292
  Likelihood parameter acceptance fractions: [0.4275]

Running marginalized_sys_err

Burn-in batch 1/4 completed, 2000 steps.

Burn-in batch 2/4 completed, 2000 steps.

Burn-in batch 3/4 completed, 2000 steps.

Burn-in batch 4/4 completed, 2000 steps.

Batch: 1/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.294

Batch: 2/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.288

Batch: 3/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.288

Batch: 4/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.264

Batch: 5/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.290

Batch: 6/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.311

Batch: 7/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.287

Batch: 8/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.292

Batch: 9/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.310

Batch: 10/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.301

Running unreported_sys_err

Burn-in batch 1/4 completed, 2000 steps.

Burn-in batch 2/4 completed, 2000 steps.

Burn-in batch 3/4 completed, 2000 steps.

Burn-in batch 4/4 completed, 2000 steps.

Batch: 1/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.301

Batch: 2/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.284

Batch: 3/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.289

Batch: 4/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.295

Batch: 5/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.282

Batch: 6/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.296

Batch: 7/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.310

Batch: 8/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.315

Batch: 9/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.290

Batch: 10/10 completed, 2000 steps. 
  Model parameter acceptance fraction: 0.291

colors = dict(
    zip(walkers.keys(), ["tab:blue", "tab:orange", "tab:purple", "tab:green"])
)
colors

{'unknown_norm': 'tab:blue',
 'unreported_sys_err_with_unknown_model_err': 'tab:orange',
 'marginalized_sys_err': 'tab:purple',
 'unreported_sys_err': 'tab:green'}

model = correct_model
fig, axes = plt.subplots(
    walkers["marginalized_sys_err"].model_sampler.chain.shape[1] + 1,
    1,
    figsize=(8, 8),
    sharex=True,
)
for i in range(walker.model_sampler.chain.shape[1]):
    # plot walkers
    for key, walker in walkers.items():
        axes[i].plot(walker.model_sampler.chain[:, i], alpha=0.4, color=colors[key])

    axes[i].set_ylabel(f"${model.params[i].latex_name}$")
    true_value = true_params[i]
    axes[i].hlines(true_value, 0, len(walker.model_sampler.chain), "r", linestyle="--")

# plot likelihoods
for key, walker in walkers.items():
    axes[-1].plot(walker.model_sampler.logp_chain, alpha=0.4, color=colors[key])

prior_samples = model_prior.rvs(walker.model_sampler.chain.shape[0])
prior_samples.shape

(20000, 5)

domain = np.zeros((len(walkers), correct_model.n_params, 2))
for i, (key, walker) in enumerate(walkers.items()):
    chain = walker.model_sampler.chain[:, : correct_model.n_params]
    domain[i, ...] = np.array(
        [
            np.min(chain, axis=0),
            np.max(chain, axis=0),
        ]
    ).T
domain = np.array([np.min(domain[:, :, 0], axis=0), np.max(domain[..., 1], axis=0)])

domain.shape

(2, 5)

fig = plt.figure()
for key, walker in walkers.items():
    corner.corner(
        walker.model_sampler.chain[:, : correct_model.n_params],
        fig=fig,
        color=colors[key],
        range=domain.T,
        labels=[f"${p.latex_name}$" for p in correct_model.params],
        truths=true_params,
        labelpad=0.1,
        max_n_ticks=2,
        truth_color="k",
    )

    plt.plot([], [], color=colors[key], label=key)
fig.legend()
# plt.tight_layout()

<matplotlib.legend.Legend at 0x780ddc0c65a0>

As expected, unknown_norm and marginalized_sys_err are the same. This is because they originate from the same exact statistical model, in which the normalization of each data set is a random variable. The latter explicitly marginalizes over the renormalization according to $\mathcal{N}(1,\sigma_{sys})$, whereas the former takes that distribution as a prior but attempts to learn the real distribution by conditioning the data.

In this case the underlying process for generating $N$ is exactly what the experimentalists report (that is, it is $\mathcal{N}(1,\sigma_{sys})$). If the experimentally reported $\sigma_{sys}$ was incorrect, then marginalized_sys_err would not be marginalizing over the correct distribution, and would converge to something wrong (try it!). On the other hand, with unknown_norm, incorrect experimentally reported $\sigma_{sys}$ just means a bad prior: $\sigma_{sys}$ doesn’t enter into the likelihood at all. Eventually, sampling should converge to the appropriate distribution.

Of course, the other two methods, which do not account for systematic error at all, fail to converge to the region of the truth.

Let’s explore the performance of the inference of the normalizations from the “unknown norm” model

We will look at the posteriors of $\rho_i$, the renormalization of the model predictions with respect to each data set. We will look at the maxima a posteriori (MAP) and compare them to the values we actually renormalized by when generating the synthetic data. Then we will plot the experimental values renormalized by the MAP $\rho$s, and we should see that they re-align themselves with the truth.

# per-dataset rho posteriors live in the model-block chain, after the physics
n_physics = len(correct_model.params)
norm_chain = np.exp(walkers["unknown_norm"].model_sampler.chain[:, n_physics:])
norm_chain.shape

(20000, 4)

norm_log_posterior_vals = walkers["unknown_norm"].model_sampler.logp_chain
norm_log_posterior_vals.shape

(20000,)

map_idx = np.argmax(walkers["unknown_norm"].model_sampler.logp_chain)
N_data_sets = len(settings)
maps = norm_chain[map_idx, :]
maps

array([1.0519558 , 0.53718276, 1.2928494 , 1.00432021])

np.array([obs.renormalization for obs in observations])

array([1.06789136, 0.53671203, 1.3071512 , 1.05429387])

np.mean(norm_chain, axis=0)

array([1.04462468, 0.5287628 , 1.28785673, 0.99731649])

fig = corner.corner(
    norm_chain,
    labels=range(len(settings)),
    truths=[obs.renormalization for obs in observations],
    labelpad=0.1,
    max_n_ticks=2,
    color="tab:blue",
    truth_color="red",
)

plt.plot(truth.x, truth.y, "k:", label="truth")
for i, synthetic_observation in enumerate(observations):
    rho_hat = maps[i]
    plt.errorbar(
        synthetic_observation.x,
        synthetic_observation.y / rho_hat,
        synthetic_observation.y_stat_err,
        linestyle="none",
        marker=".",
        label=r"$\rho_{\text{truth}}$ = "
        + f"{synthetic_observation.renormalization:1.3f},"
        + r" MAP = "
        + f"{rho_hat:1.2f}",
    )
plt.xlabel("$x$")
plt.ylabel("$y$")
plt.legend()
plt.title(r"Renormalizing experimental data sets based on MAP of $p(\rho)$")

Text(0.5, 1.0, 'Renormalizing experimental data sets based on MAP of $p(\\rho)$')

predictive posteriors

walker.model_sampler.chain.shape

(20000, 5)

def predictive_post(chain, model, obs, n_samples, intervals):
    draw_idxs = rng.choice(np.arange(chain.shape[0]), n_samples)
    draws = chain[draw_idxs, :]
    ym = np.array([model(obs, *p) for p in draws])
    return np.percentile(ym, intervals, axis=0)

fig = plt.figure(figsize=(8, 4))
plt.plot(truth.x, truth.y, "k:", label="truth")

for synthetic_observation in observations:
    plt.errorbar(
        synthetic_observation.x,
        synthetic_observation.y,
        synthetic_observation.y_stat_err,
        linestyle="none",
        marker=".",
        alpha=0.4,
        # label=f"renormalization = {synthetic_observation.renormalization:1.3f}",
    )

hatches = ["//\\//\\", "|-", "", ""]
alphas = [0.1, 0.25, 0.25, 0.25]
for i, (key, walker) in enumerate(walkers.items()):
    intervals = predictive_post(
        walker.model_sampler.chain[:, : correct_model.n_params],
        correct_model,
        truth,
        1000,
        [16, 84],
    )
    plt.fill_between(
        truth.x,
        intervals[0],
        intervals[1],
        alpha=alphas[i],
        color=colors[key],
        label=key,
        hatch=hatches[i],
        zorder=99,
    )
plt.xlabel("x")
plt.ylabel("y")
plt.ylim([0, 2])
fig.legend(loc="upper left", framealpha=1)

<matplotlib.legend.Legend at 0x780dd75b3470>

Covariance-structure gallery

The sections above inferred unknown normalisations — a single rank-one systematic per dataset. But a Constraint’s covariance is assembled from arbitrary Terms, so the same machinery expresses far richer uncertainty structure. This gallery surveys four cases the covariance API handles, each just a different Term (or support) on the stacked residual:

1. a systematic correlated across $x$ (smoothly, not a flat normalisation);

2. correlated statistical errors — and why ignoring them is overconfident;

3. unknown / misreported magnitudes (the free-$\gamma$ model error above);

4. a systematic shared across datasets (cross-block coupling).

A small helper normalises a covariance to a correlation matrix for plotting.

from sklearn.gaussian_process.kernels import RBF, ConstantKernel


def correlation(Sigma):
    d = np.sqrt(np.diag(Sigma))
    return Sigma / np.outer(d, d)

1. A systematic correlated across $x$

A flat normalization_term is a rank-one mode: every point is 100 % correlated (correlation matrix all ones off-diagonal). A systematic that varies smoothly with $x$ — e.g. an energy-dependent efficiency — instead has correlation that decays with separation. That is a kernel_term (a GP prior on the systematic). Same API, richer structure.

o0 = observations[0]
xspan = o0.x.max() - o0.x.min()

# (a) flat normalization: rank-one, uniform correlation
c_flat = rxmc.constraint.Constraint(
    [o0],
    correct_model,
    extra_terms=[rxmc.covariance.normalization_term(magnitude=0.05)],
)
# (b) smooth correlated systematic: a GP (kernel_term) over x
sys_kernel = ConstantKernel(0.05**2) * RBF(length_scale=xspan / 4)
c_smooth = rxmc.constraint.Constraint(
    [o0],
    correct_model,
    extra_terms=[rxmc.covariance.kernel_term(sys_kernel)],
)

S_flat = c_flat.covariance_matrix(true_params)
S_smooth = c_smooth.covariance_matrix(true_params, tuple(sys_kernel.theta))

fig, ax = plt.subplots(1, 2, figsize=(9, 4))
for a, S, t in [
    (ax[0], S_flat, "flat normalization (rank-one)"),
    (ax[1], S_smooth, "smooth systematic (kernel_term)"),
]:
    im = a.imshow(correlation(S), vmin=-1, vmax=1, cmap="RdBu_r")
    a.set_title(t)
    fig.colorbar(im, ax=a, fraction=0.046)
fig.suptitle("correlation of one dataset under two systematic structures");

2. Correlated statistical errors

Statistical errors are usually taken as independent (a diagonal covariance). When they are not — shared backgrounds, unfolding, detector resolution — ignoring the correlation makes the posterior overconfident. We draw line data with exponentially-correlated noise and fit it two ways: a naive diagonal, and the correct correlated covariance supplied as a fixed Term.

line = rxmc.physical_model.Polynomial(1)  # a0 + a1 x
a_true = [1.0, 0.8]
xs = np.linspace(0.0, 4.0, 20)
ell, sig = 1.2, 0.15
R = np.exp(-np.abs(xs[:, None] - xs[None, :]) / ell)  # exponential correlation
C = sig**2 * R
y_corr = line(rxmc.observation.Observation(x=xs, y=np.zeros_like(xs)), *a_true)
y_corr = y_corr + rng.multivariate_normal(np.zeros(xs.size), C)
obs_corr = rxmc.observation.Observation(x=xs, y=y_corr)

c_naive = rxmc.constraint.Constraint(
    [obs_corr],
    line,
    extra_terms=[rxmc.covariance.Term(sig * np.ones(xs.size), kind="diag")],
)
c_correlated = rxmc.constraint.Constraint(
    [obs_corr],
    line,
    extra_terms=[rxmc.covariance.Term(C)],
)


def fit_line(constraint, seed=7):
    evidence = rxmc.evidence.Evidence([constraint])
    prior = stats.multivariate_normal(mean=a_true, cov=np.diag([0.5, 0.5]) ** 2)
    sampler = rxmc.param_sampling.BatchedAdaptiveMetropolisSampler(
        params=line.params,
        starting_location=prior.mean,
        prior=prior,
        initial_proposal_cov=prior.cov / 100,
    )
    walker = rxmc.walker.Walker(sampler, evidence, rng=np.random.default_rng(seed))
    walker.walk(n_steps=6000, burnin=2000, batch_size=1000, verbose=False)
    return walker.model_sampler.chain


chain_naive = fit_line(c_naive)
chain_correlated = fit_line(c_correlated)
for name, ch in [("naive diagonal", chain_naive), ("correlated", chain_correlated)]:
    print(
        f"{name:14s} a0 = {ch[:,0].mean():.3f} ± {ch[:,0].std():.3f}   "
        f"a1 = {ch[:,1].mean():.3f} ± {ch[:,1].std():.3f}"
    )

naive diagonal a0 = 1.101 ± 0.062   a1 = 0.758 ± 0.027
correlated     a0 = 1.048 ± 0.129   a1 = 0.779 ± 0.048

fig = corner.corner(
    chain_naive, labels=["a0", "a1"], color="tab:red", truths=a_true, truth_color="k"
)
corner.corner(chain_correlated, fig=fig, color="tab:blue")
plt.plot([], [], color="tab:red", label="naive diagonal (overconfident)")
plt.plot([], [], color="tab:blue", label="correlated fixed Term")
fig.legend(loc="upper right");

The naive-diagonal posterior is visibly tighter than the correct one: treating correlated noise as independent over-counts the information. The a fixed Term with the true correlation restores honest uncertainty.

3. Unknown / misreported magnitudes

When the size of an uncertainty is itself unknown, it becomes a sampled Parameter. The unreported_sys_err_with_unknown_model_err scenario above does exactly this with a free $\gamma$ (model_error_term, a diagonal inflation). Here we just visualise how the covariance diagonal grows as $\gamma$ increases — inference slides along this family to whatever the data support.

o = observations[0]
gamma = rxmc.params.Parameter("log gamma", float)
c_me = rxmc.constraint.Constraint(
    [o],
    correct_model,
    extra_terms=[rxmc.covariance.model_error_term(gamma, averaging=True)],
)
plt.figure(figsize=(7, 4))
for g in [0.02, 0.05, 0.10]:
    S = c_me.covariance_matrix(true_params, (np.log(g),))
    plt.plot(o.x, np.sqrt(np.diag(S)), marker=".", label=f"$\gamma$ = {g:.2f}")
plt.plot(o.x, o.y_stat_err, "k:", label="reported stat. err")
plt.xlabel("x")
plt.ylabel("total std. dev.")
plt.legend()
plt.title("unknown model error inflates the diagonal by a sampled $\gamma$");

<>:12: SyntaxWarning: invalid escape sequence '\g'
<>:17: SyntaxWarning: invalid escape sequence '\g'
<>:12: SyntaxWarning: invalid escape sequence '\g'
<>:17: SyntaxWarning: invalid escape sequence '\g'
/tmp/ipykernel_402344/1563835932.py:12: SyntaxWarning: invalid escape sequence '\g'
  plt.plot(o.x, np.sqrt(np.diag(S)), marker=".", label=f"$\gamma$ = {g:.2f}")
/tmp/ipykernel_402344/1563835932.py:17: SyntaxWarning: invalid escape sequence '\g'
  plt.title("unknown model error inflates the diagonal by a sampled $\gamma$");

4. A systematic shared across datasets

If two datasets share a systematic (a common calibration), the coupling is a single Term whose support spans both blocks — off-diagonal correlation between datasets. Treating them independently throws that coupling away. (The dedicated correlated_observations.ipynb notebook explores the inference consequences; here we just show the structure.)

pair = observations[:2]
supports = _block_supports(pair)
full = np.concatenate(supports)

c_shared = rxmc.constraint.Constraint(
    pair,
    correct_model,
    extra_terms=[rxmc.covariance.normalization_term(magnitude=0.06, support=full)],
)
c_separate = rxmc.constraint.Constraint(
    pair,
    correct_model,
    extra_terms=[
        rxmc.covariance.normalization_term(magnitude=0.06, support=s) for s in supports
    ],
)
n1 = pair[0].n_data_pts
fig, ax = plt.subplots(1, 2, figsize=(9, 4))
for a, c, t in [
    (ax[0], c_shared, "shared across datasets (case A)"),
    (ax[1], c_separate, "independent per dataset"),
]:
    im = a.imshow(
        correlation(c.covariance_matrix(true_params)), vmin=-1, vmax=1, cmap="RdBu_r"
    )
    a.axhline(n1 - 0.5, color="k", lw=0.6)
    a.axvline(n1 - 0.5, color="k", lw=0.6)
    a.set_title(t)
    fig.colorbar(im, ax=a, fraction=0.046)
fig.suptitle("a shared systematic couples the two datasets' blocks");

Summary

Every case above is the same Constraint machinery with a different Term:

case	Term	structure
flat normalisation	normalization_term	rank-one, uniform
correlated-across-$x$ systematic	kernel_term	smooth/banded
correlated statistical errors	fixed Term	arbitrary off-diagonal
unknown magnitude	model_error_term (free $\gamma$)	sampled diagonal
shared across datasets	cross-block normalization_term	off-diagonal blocks

Declaring the uncertainty is the modelling choice — the inference machinery is unchanged.