From de570b81e2493368152586446761d61f622d3d20 Mon Sep 17 00:00:00 2001 From: Frank Buckley Date: Sun, 26 Jul 2026 15:55:41 +0100 Subject: [PATCH] fix(loo): derive the refit path's observed-variable name instead of hard-coding y_post MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The exact-refit repair read `idata.log_likelihood["y_post"]` in its alignment guards and in the held-out density evaluation. That name is a property of the likelihood, not a constant: the Beta-Binomial mechanism models register `y_post`, but the floor-rule likelihood registers `y_offfloor`. Registering an off-floor mechanism model — the primary estimand proposed in #433 — would have turned every guard into a KeyError raised at comparison time, after the fits had already run. Derive the name from the built model's single observed RV and verify it against the stored trace, so the failure mode is "this model is not supported", stated before any sampling, rather than a crash an hour in. Multi-outcome models are now refused explicitly: they carry one obs_id axis per outcome, so there is no well-defined single row for reloo to hold out. No behaviour change for the models that exist today — both live traces still resolve to `y_post` and pass every guard unchanged. Refs #433, #438 Co-Authored-By: Claude Opus 5 --- .../statistical_models/loo_refit.py | 53 ++++++++- tests/statistical_models/test_loo_refit.py | 102 ++++++++++++++++++ 2 files changed, 152 insertions(+), 3 deletions(-) diff --git a/src/language_reading_predictors/statistical_models/loo_refit.py b/src/language_reading_predictors/statistical_models/loo_refit.py index 0ea8f8bd..9cb46e54 100644 --- a/src/language_reading_predictors/statistical_models/loo_refit.py +++ b/src/language_reading_predictors/statistical_models/loo_refit.py @@ -72,6 +72,46 @@ def _as_dataset(group): return inner if hasattr(inner, "data_vars") else group +def _observed_variable_name(model: pm.Model, idata_orig, model_id: str) -> str: + """Name of the model's single observed node, verified against the stored trace. + + Hard-coding ``"y_post"`` is correct only for the Beta-Binomial mechanism models + that exist today. The floor-rule likelihood registers its observed node as + ``y_offfloor`` (see ``factories``), so a hard-coded name would turn the alignment + guards below into a ``KeyError`` raised at *comparison* time — after the fits have + already run — the moment an off-floor mechanism model is registered (#433). Deriving + the name from the built model instead keeps the failure mode "this model is not + supported", stated up front, rather than "this crashed after an hour of sampling". + + Refuses anything the wrapper's single-variable pointwise indexing cannot represent: + a multi-outcome model has one ``obs_id`` axis per outcome, and there is no + well-defined single row for ``reloo`` to hold out. + """ + observed = [rv.name for rv in model.observed_RVs] + if len(observed) != 1: + raise ValueError( + f"{model_id}: rebuilt model has {len(observed)} observed variables " + f"({sorted(observed) if observed else 'none'}); the exact-refit repair " + "splices a single pointwise log-likelihood and cannot hold out a row of a " + "multi-outcome model — refusing to refit" + ) + name = observed[0] + stored = set(_as_dataset(idata_orig.log_likelihood).data_vars) + if len(stored) != 1: + raise ValueError( + f"{model_id}: stored log_likelihood group holds {sorted(stored)}; the " + "exact-refit repair cannot tell which variable the Pareto-k indices refer " + "to — refusing to refit" + ) + if name not in stored: + raise ValueError( + f"{model_id}: rebuilt model observes {name!r} but the stored trace holds " + f"{sorted(stored)}; the construction path has drifted from the one that " + "produced this fit — refusing to refit" + ) + return name + + @dataclass(frozen=True) class RefitPlan: """Sampler settings for a refit, taken from the original fit's ``config.json``. @@ -124,6 +164,7 @@ def __init__( refit: RefitPlan, full_model: pm.Model, design, + obs_var: str, *, progressbar: bool = False, ) -> None: @@ -133,6 +174,10 @@ def __init__( self.design = design self.refit = refit self.full_model = full_model + # Derived from the built model rather than assumed, so a likelihood change + # renames the node without silently breaking the splice (see + # :func:`_observed_variable_name`). + self.obs_var = obs_var self.progressbar = progressbar self.n_refits = 0 @@ -244,7 +289,7 @@ def log_likelihood__i(self, excluded_observed_data, idata__i): log_lik = pm.compute_log_likelihood( idata__i, extend_inferencedata=False, progressbar=False ) - return log_lik["y_post"].isel(obs_id=int(excluded_observed_data)) + return log_lik[self.obs_var].isel(obs_id=int(excluded_observed_data)) def build_mechanism_wrapper( @@ -262,10 +307,11 @@ def build_mechanism_wrapper( """ plan = _mechanism.resolve_mechanism_plan(spec) built = _mechanism.build_mechanism_for_plan(plan) + obs_var = _observed_variable_name(built.model, idata_orig, spec.model_id) # Take the observation count from the log-likelihood group, which *is* the # authoritative observation index for LOO, rather than from a posterior dim that # only exists because some deterministic happens to carry ``obs_id``. - stored_ll = idata_orig.log_likelihood["y_post"] + stored_ll = _as_dataset(idata_orig.log_likelihood)[obs_var] n_trace = int(stored_ll.sizes["obs_id"]) if built.prepared.n_obs != n_trace: raise ValueError( @@ -278,7 +324,7 @@ def build_mechanism_wrapper( recomputed = pm.compute_log_likelihood( idata_orig, model=built.model, extend_inferencedata=False, progressbar=False ) - delta = float(np.abs(stored_ll.values - recomputed["y_post"].values).max()) + delta = float(np.abs(stored_ll.values - recomputed[obs_var].values).max()) if delta > 1e-6: raise ValueError( f"{spec.model_id}: rebuilt model does not reproduce the stored " @@ -338,5 +384,6 @@ def build_mechanism_wrapper( refit=RefitPlan.from_config(config), full_model=built.model, design=design, + obs_var=obs_var, progressbar=progressbar, ) diff --git a/tests/statistical_models/test_loo_refit.py b/tests/statistical_models/test_loo_refit.py index 537f1387..c782f261 100644 --- a/tests/statistical_models/test_loo_refit.py +++ b/tests/statistical_models/test_loo_refit.py @@ -179,6 +179,108 @@ def test_frozen_design_refuses_when_it_cannot_reproduce(): ).require_moderator_scaler() +class _FakeRV: + def __init__(self, name: str) -> None: + self.name = name + + +class _FakeModel: + """Stand-in exposing only ``observed_RVs``, which is all the guard reads.""" + + def __init__(self, *names: str) -> None: + self.observed_RVs = [_FakeRV(n) for n in names] + + +def _fake_idata(*names: str): + import xarray as xr + + from types import SimpleNamespace + + ds = xr.Dataset( + {n: (("chain", "draw", "obs_id"), np.zeros((2, 3, 4))) for n in names} + ) + return SimpleNamespace(log_likelihood=ds) + + +def test_observed_variable_name_follows_the_likelihood_rather_than_assuming_y_post(): + """The observed node's name is a property of the likelihood, not a constant. + + Beta-Binomial mechanism models register ``y_post``; the floor-rule likelihood + registers ``y_offfloor``. Hard-coding the former made the alignment guards raise + ``KeyError`` at *comparison* time — after the fits — the moment an off-floor + mechanism model was registered (#433).""" + from language_reading_predictors.statistical_models.loo_refit import ( + _observed_variable_name, + ) + + assert ( + _observed_variable_name(_FakeModel("y_post"), _fake_idata("y_post"), "m") + == "y_post" + ) + assert ( + _observed_variable_name( + _FakeModel("y_offfloor"), _fake_idata("y_offfloor"), "m" + ) + == "y_offfloor" + ) + + +def test_observed_variable_name_refuses_a_multi_outcome_model(): + """A joint model has one ``obs_id`` axis per outcome, so there is no single row + for ``reloo`` to hold out. Refuse rather than splice against the first one.""" + from language_reading_predictors.statistical_models.loo_refit import ( + _observed_variable_name, + ) + + with pytest.raises(ValueError, match="multi-outcome"): + _observed_variable_name( + _FakeModel("y_W", "y_N"), _fake_idata("y_W", "y_N"), "m" + ) + + +def test_observed_variable_name_refuses_when_the_trace_names_something_else(): + from language_reading_predictors.statistical_models.loo_refit import ( + _observed_variable_name, + ) + + with pytest.raises(ValueError, match="drifted"): + _observed_variable_name(_FakeModel("y_offfloor"), _fake_idata("y_post"), "m") + + with pytest.raises(ValueError, match="which variable"): + _observed_variable_name( + _FakeModel("y_post"), _fake_idata("y_post", "y_aux"), "m" + ) + + +def test_held_out_density_is_read_from_the_derived_observed_name(monkeypatch): + """Closes the loop: the wrapper must *use* the derived name, not just record it.""" + import contextlib + + import pymc as pm_mod + import xarray as xr + + from language_reading_predictors.statistical_models.loo_refit import ( + MechanismSamplingWrapper, + ) + + log_lik = xr.Dataset( + { + "y_offfloor": (("chain", "draw", "obs_id"), np.arange(24.0).reshape(2, 3, 4)), + } + ) + monkeypatch.setattr( + pm_mod, "compute_log_likelihood", lambda *a, **k: log_lik, raising=True + ) + + wrapper = object.__new__(MechanismSamplingWrapper) + wrapper.obs_var = "y_offfloor" + wrapper.full_model = contextlib.nullcontext() + + got = wrapper.log_likelihood__i(2, idata__i=None) + assert got.name == "y_offfloor" + np.testing.assert_allclose(got.values, log_lik["y_offfloor"].isel(obs_id=2).values) + + def test_as_dataset_unwraps_a_datatree_group(): """ArviZ 1.x groups are ``DataTree``s whose ``.items()`` yields child *nodes*, not variables — iterating one directly and reading ``.values`` raises instead of