fix(loo): derive the refit path's observed-variable name instead of hard-coding y_post

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,
     )

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