Combining Models

A common pattern in HEP is combining independent statistical models that share a parameter of interest but use different naming conventions. This page walks through the full workflow: aligning parameter names, merging states, and running a joint fit.

Setup: two independent measurements

Two experiments each measure a particle mass with their own calibration uncertainties:

import jax
import jax.numpy as jnp
import everwillow as ew
import everwillow.statelib as sl
from everwillow.uncertainty import uncertainties

jax.config.update("jax_enable_x64", True)


# Experiment A: measures "mass" with a scale uncertainty
def nll_a(params, obs):
    mass = params["mass"]
    scale = params["scale_a"]
    pred = mass * scale
    return (
        0.5 * ((pred - obs["m_a"]) / obs["err_a"]) ** 2
        + 0.5 * (scale - 1.0) ** 2 / 0.02**2
    )


# Experiment B: measures the same quantity but calls it "m"
def nll_b(params, obs):
    mass = params["m"]
    scale = params["scale_b"]
    pred = mass * scale
    return (
        0.5 * ((pred - obs["m_b"]) / obs["err_b"]) ** 2
        + 0.5 * (scale - 1.0) ** 2 / 0.03**2
    )

Each model was developed independently, so the shared parameter has a different name:

state_a = sl.State.from_pytree({"mass": 125.0, "scale_a": 1.0})
state_b = sl.State.from_pytree({"m": 125.0, "scale_b": 1.0})

Aligning parameter names

Use apply_transformations() to rename "m" → "mass" in model B:

state_b_aligned = sl.apply_transformations(
    state_b,
    {"m": sl.Transform(new_key="mass")},
)

The renamed state remembers the original key. When you call to_pytree(), model B still receives {"m": ...} - the rename is transparent to the NLL function.

Combined fit

Before fitting, call prepare() to merge the states and build a combined NLL that dispatches each sub-pytree to the correct model. This is necessary because fit() takes a single NLL and a single state - prepare() produces both from the individual pieces:

obs = {"m_a": 125.5, "err_a": 1.5, "m_b": 124.8, "err_b": 2.0}

combined_nll, combined = ew.prepare([nll_a, nll_b], [state_a, state_b_aligned])

result = ew.fit(combined_nll, combined, obs)
print(result.params.to_pytree())
# ({'mass': 125.27, 'scale_a': 1.0}, {'m': 125.27, 'scale_b': 1.0})
#           ^ same value in both models

To recover the individual states after fitting, use split():

state_a_fitted, state_b_fitted = sl.split(result.params)
print(state_a_fitted.to_pytree())  # {'mass': 125.27, 'scale_a': 1.0}
print(state_b_fitted.to_pytree())  # {'m': 125.27, 'scale_b': 1.0}

Comparing uncertainties

The combined measurement has a smaller uncertainty on "mass" than either experiment alone:

# Individual fits
result_a = ew.fit(nll_a, state_a, obs)
result_b = ew.fit(nll_b, state_b, obs)

unc_a = uncertainties(nll_a, result_a.params, obs)
unc_b = uncertainties(nll_b, result_b.params, obs)
unc_combined = uncertainties(combined_nll, result.params, obs)

print(f"σ(mass) exp A only:  {float(unc_a['mass']):.3f}")
print(f"σ(mass) exp B only:  {float(unc_b['m']):.3f}")
print(f"σ(mass) combined:    {float(unc_combined['mass']):.3f}")
# σ(mass) exp A only:  2.924
# σ(mass) exp B only:  4.245
# σ(mass) combined:    2.408

The combined constraint is tighter than either individual measurement - the merged state ensures both models are optimized with a single shared mass parameter.