Retirement surfaces are owned by the base builder: no release-time refit

changelog.d/keogh-puf-refit.fixed.md

@@ -0,0 +1 @@
+Preserve the base builder's completed retirement-distribution surface after frozen-support selection, preventing a selected rare Keogh donor from being refit and broadcast across PUF-support rows.

packages/populace-build/src/populace/build/us_runtime/retirement_distributions.py

@@ -595,16 +595,27 @@ def with_us_retirement_distribution_inputs(
     *,
     seed: int,
     time_period: int,
+    force_puf_imputation: bool = False,
 ) -> Frame:
-    """Materialize measured retirement-distribution leaves on a US frame."""
+    """Materialize measured retirement-distribution leaves on a US frame.
+
+    ``force_puf_imputation`` belongs only at the base builder's post-clone
+    boundary.  Every later support-frame call is consume-only, including when
+    a frozen selection is missing or has flattened a rare leaf.  Refitting
+    there would make support selection redefine the donor universe and can
+    broadcast a rare leaf such as ``keogh_distributions`` across the retained
+    PUF rows.  The downstream signal gate, rather than a refit, owns support
+    surface completeness and signal.
+    """
 
     if frame.schema != US_SCHEMA:
         raise ValueError("US retirement distributions require the US schema.")
     person = frame.table("person")
     has_support_channels = _PERSON_SUPPORT_CHANNEL_COLUMN in person.columns
-    if (
-        _retirement_distribution_surface_carries_signal(frame)
-        and not has_support_channels
+    if has_support_channels and not force_puf_imputation:
+        return frame
+    if not has_support_channels and _retirement_distribution_surface_carries_signal(
+        frame
     ):
         return frame
 

packages/populace-build/tests/test_us_puf_support_base_builder.py

@@ -1043,7 +1043,7 @@ def test_main_runs_cps_only_inputs_before_clone_and_after_puf_then_fails_gate(
     salt_refund_gate_frames: list[object] = []
     adult_care_gate_frames: list[object] = []
     energy_subsidy_gate_frames: list[object] = []
-    retirement_distribution_calls: list[tuple[object, int, int]] = []
+    retirement_distribution_calls: list[tuple[object, int, int, bool]] = []
     retirement_distribution_gate_frames: list[object] = []
     prior_year_income_calls: list[tuple[object, int, int]] = []
     prior_year_income_gate_frames: list[object] = []
@@ -1255,8 +1255,16 @@ def fake_with_weeks_unemployed(
         lambda frame, *, seed, time_period: frame,
     )
 
-    def fake_retirement_distributions(frame, *, seed, time_period):
-        retirement_distribution_calls.append((frame, seed, time_period))
+    def fake_retirement_distributions(
+        frame,
+        *,
+        seed,
+        time_period,
+        force_puf_imputation=False,
+    ):
+        retirement_distribution_calls.append(
+            (frame, seed, time_period, force_puf_imputation)
+        )
         if frame == "disability-benefits-direct":
             return "retirement-distributions-direct"
         return "retirement-distributions-puf"
@@ -1652,10 +1660,12 @@ def fake_retirement_distributions_signal_gate(frame):
         if failing_gate in {"energy_subsidy", "retirement_distributions"}
         else []
     )
-    expected_retirement_distribution_calls = [("disability-benefits-direct", 7, 2024)]
+    expected_retirement_distribution_calls = [
+        ("disability-benefits-direct", 7, 2024, False)
+    ]
     if failing_gate == "retirement_distributions":
         expected_retirement_distribution_calls.append(
-            ("disability-benefits-puf", 7, 2024)
+            ("disability-benefits-puf", 7, 2024, True)
         )
     assert retirement_distribution_calls == expected_retirement_distribution_calls
     assert retirement_distribution_gate_frames == (
@@ -1665,6 +1675,22 @@ def fake_retirement_distributions_signal_gate(frame):
     )
 
 
+def test_resume_retirement_stage_forces_puf_imputation() -> None:
+    """The resume-side ownership boundary is pinned (PR #557 round 2, low).
+
+    The live post-clone boundary is behaviorally asserted above; this pins
+    the named-stage resume branch so deleting its force flag fails a test
+    (source-pin precedent: the main-summary gate tests below).
+    """
+    builder = _load_support_builder_module()
+    source = Path(builder.__file__).read_text(encoding="utf-8")
+    marker = 'elif stage == "retirement_distributions_post_clone":'
+    assert marker in source
+    window = source.split(marker, 1)[1].split("elif ", 1)[0]
+    assert "with_us_retirement_distribution_inputs(" in window
+    assert "force_puf_imputation=True" in window
+
+
 def test_main_summary_records_retirement_distribution_gate() -> None:
     builder = _load_support_builder_module()
     source = Path(builder.__file__).read_text(encoding="utf-8")

packages/populace-build/tests/test_us_retirement_distributions.py

@@ -316,7 +316,12 @@ def fit(
             return FakeFitted()
 
     monkeypatch.setattr(module, "QRF", FakeQRF)
-    result = with_us_retirement_distribution_inputs(expanded, seed=7, time_period=2024)
+    result = with_us_retirement_distribution_inputs(
+        expanded,
+        seed=7,
+        time_period=2024,
+        force_puf_imputation=True,
+    )
 
     assert calls["init"] == {"n_estimators": 100, "seed": 7}
     assert len(calls["training"]) == len(direct.table("person"))
@@ -342,6 +347,126 @@ def fit(
     assert gate.passed, gate.failures
 
 
+def test_completed_puf_surface_survives_narrowed_support_without_refit(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    direct = with_us_retirement_distribution_inputs(_frame(), seed=0, time_period=2024)
+    expanded = clone_us_frame_for_puf_support(direct)
+
+    class ZeroFitted:
+        def predict(self, test: pd.DataFrame, **kwargs) -> pd.DataFrame:
+            return pd.DataFrame(
+                0.0,
+                index=test.index,
+                columns=list(_PUF_QRF_OUTPUTS),
+            )
+
+    class ZeroQRF:
+        def __init__(self, **kwargs: object) -> None:
+            pass
+
+        def fit(self, *args: object, **kwargs: object) -> ZeroFitted:
+            return ZeroFitted()
+
+    monkeypatch.setattr(module, "QRF", ZeroQRF)
+    completed = with_us_retirement_distribution_inputs(
+        expanded,
+        seed=7,
+        time_period=2024,
+        force_puf_imputation=True,
+    )
+
+    # Model frozen-support recovery by removing one all-zero PUF household.
+    # Frame.select deliberately preserves the surviving person index, so the
+    # selected frame also covers the non-RangeIndex path that changed the
+    # historical 5,000-row donor sample.
+    person = completed.table("person")
+    puf_zero = person["person_support_channel"].eq("puf_tax_detail") & ~person[
+        list(_PUF_QRF_OUTPUTS)
+    ].any(axis=1)
+    drop_index = person.index[puf_zero][0]
+    selected = completed.select(person.index != drop_index)
+    before = us_retirement_distributions_signal_gate(selected)
+    assert before.passed, before.failures
+    before_share = before.details["nonzero_shares"]["keogh_distributions"]
+    assert 0.0000001 <= before_share <= 0.005
+    before_values = selected.table("person")[list(_OUTPUTS)].copy()
+    before_keogh_carriers = int((before_values["keogh_distributions"] > 0).sum())
+
+    class UnexpectedQRF:
+        def __init__(self, **kwargs: object) -> None:
+            raise AssertionError("a completed retirement surface must not be refit")
+
+    monkeypatch.setattr(module, "QRF", UnexpectedQRF)
+    result = with_us_retirement_distribution_inputs(
+        selected,
+        seed=7,
+        time_period=2024,
+    )
+
+    assert result is selected
+    after = us_retirement_distributions_signal_gate(result)
+    assert after.passed, after.failures
+    assert after.details["nonzero_shares"]["keogh_distributions"] == before_share
+    pd.testing.assert_frame_equal(
+        result.table("person")[list(_OUTPUTS)],
+        before_values,
+    )
+    assert (
+        result.table("person")["keogh_distributions"].gt(0).sum()
+        == before_keogh_carriers
+    )
+
+    # If support selection removes every rare carrier from one leaf, preserve
+    # the completed surface and fail closed at the gate instead of refitting.
+    # Removing the measured carrier rows keeps source reconciliation valid and
+    # isolates the degeneration/prevalence checks this regression owns.
+    keogh_carriers = result.table("person")["keogh_distributions"].gt(0)
+    assert keogh_carriers.any()
+    selected_away = result.select(~keogh_carriers)
+    degenerate = with_us_retirement_distribution_inputs(
+        selected_away,
+        seed=7,
+        time_period=2024,
+    )
+    assert degenerate is selected_away
+    degenerate_gate = us_retirement_distributions_signal_gate(degenerate)
+    assert not degenerate_gate.passed
+    assert degenerate_gate.failures == (
+        "keogh_distributions: degenerate with 1 distinct value(s).",
+        "keogh_distributions: weighted nonzero share 0.00000000 outside "
+        "[0.00000010, 0.00500000].",
+    )
+    assert degenerate_gate.details["source_mismatches"]["keogh_distributions"] == 0
+
+
+@pytest.mark.parametrize("missing", _OUTPUTS)
+def test_incomplete_puf_surface_is_consume_only_and_fails_at_the_signal_gate(
+    monkeypatch: pytest.MonkeyPatch,
+    missing: str,
+) -> None:
+    direct = with_us_retirement_distribution_inputs(_frame(), seed=0, time_period=2024)
+    incomplete = clone_us_frame_for_puf_support(direct)
+    incomplete.table("person").drop(columns=[missing], inplace=True)
+
+    class UnexpectedQRF:
+        def __init__(self, **kwargs: object) -> None:
+            raise AssertionError("an incomplete support surface must not be refit")
+
+    monkeypatch.setattr(module, "QRF", UnexpectedQRF)
+    result = with_us_retirement_distribution_inputs(
+        incomplete,
+        seed=7,
+        time_period=2024,
+    )
+
+    assert result is incomplete
+    gate = us_retirement_distributions_signal_gate(result)
+    assert not gate.passed
+    assert gate.failures == (f"person columns missing: {[missing]}.",)
+    assert gate.details == {"missing": [missing]}
+
+
 def test_gate_rejects_a_default_or_source_divergent_leaf() -> None:
     result = with_us_retirement_distribution_inputs(_frame(), seed=0, time_period=2024)
     result.table("person")["keogh_distributions"] = 0.0