Make the identity-stability instrument actually verify the claim

The tool's main() wrote a stub receipt ("requires caller-supplied
transform") without checking anything. It now recomputes every E4 column
from the pure derivations in original and permuted row order, un-permutes
by entity id, compares both against each other and against the stored
artifact columns, reports missing stored columns explicitly, and exits
nonzero on any mismatch.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: juaristi22

packages/microcosm-build/tests/test_uk_stochastic_tools.py

@@ -95,6 +95,66 @@ def transform(candidate):
     assert receipt["mismatches"] == {}
 
 
+def test_e4_identity_receipt_survives_permutation_on_synthetic_frame() -> None:
+    tool = _load_tool("verify_uk_identity_stability")
+
+    class Contract:
+        def rate(self, key: str) -> float:
+            return 0.5 if not key.startswith("scp") else 0.9
+
+        def continuous_entry(self, key: str):
+            return {"mean": 15.019, "sd": 4.972, "lower": 0, "upper": 30}
+
+    person = pd.DataFrame(
+        {
+            "person_id": [101, 102, 201, 301],
+            "person_benunit_id": [10, 10, 20, 30],
+            "person_household_id": [1, 1, 1, 2],
+            "age": [5, 6, 40, 70],
+            "child_benefit_reported": [0.0, 10.0, 0.0, 0.0],
+            "pension_credit_reported": [0.0, 0.0, 0.0, 5.0],
+            "universal_credit_reported": [0.0, 0.0, 20.0, 0.0],
+        }
+    )
+    frame = uk_national_frame(
+        person=person,
+        benunit=pd.DataFrame({"benunit_id": [10, 20, 30]}),
+        household=pd.DataFrame(
+            {
+                "household_id": [1, 2],
+                "region": ["LONDON", "SCOTLAND"],
+                "household_weight": [2.0, 3.0],
+            }
+        ),
+        time_period="2023",
+    )
+    count_resource = {
+        "cells": {
+            "LONDON": {"A": {"CENTRAL_LONDON": 3, "OUTER_LONDON": 1}},
+            "SCOTLAND": {"A": {"LOTHIAN": 2}},
+        }
+    }
+
+    receipt = tool.e4_identity_receipt(
+        frame,
+        contract=Contract(),
+        count_resource=count_resource,
+        lha_category=["A", "A", "A"],
+        permutation_seed=7,
+    )
+
+    assert receipt["identical_under_permutation"] is True
+    assert receipt["permutation_mismatches"] == {}
+    # The synthetic frame carries no stored E4 columns; the receipt says so
+    # explicitly instead of silently passing the stored comparison.
+    assert receipt["matches_stored_columns"] is False
+    assert set(receipt["stored_columns_missing"]) == {
+        "person",
+        "benunit",
+        "household",
+    }
+
+
 def test_brma_distribution_masks_small_counts() -> None:
     tool = _load_tool("emit_uk_brma_distribution")
     household = pd.DataFrame(

tools/verify_uk_identity_stability.py

@@ -1,13 +1,38 @@
-"""Verify E4 stochastic stage identity stability on an existing UK frame."""
+"""Verify E4 stochastic stage identity stability on an existing UK frame.
+
+Recomputes every E4 column twice from the pure derivations — once in the
+frame's row order, once on row-permuted tables — un-permutes by entity id,
+and also compares the original-order recomputation against the columns
+stored in the artifact. Exit status is nonzero on any mismatch.
+"""
 
 from __future__ import annotations
 
 import argparse
 import json
-from collections.abc import Callable, Sequence
+import sys
+from collections.abc import Callable, Mapping, Sequence
 from pathlib import Path
 
+import numpy as np
+import pandas as pd
+
+from microcosm.build.uk_runtime.frs_brma import (
+    UK_BRMA_DECLARED_SEEDS,
+    _benunit_regions,
+    _enum_name,
+    assign_brma_by_cell,
+    collapse_benunit_brma_to_household,
+    load_brma_count_resource,
+)
+from microcosm.build.uk_runtime.frs_household_draws import derive_frs_household_draws
+from microcosm.build.uk_runtime.frs_person_draws import derive_frs_person_draws
+from microcosm.build.uk_runtime.frs_take_up import (
+    aggregate_person_reported_to_benunit,
+    derive_frs_take_up,
+)
 from microcosm.build.uk_runtime.national_build import load_uk_national_frame
+from microcosm.build.uk_runtime.national_frame import uk_time_period
 
 
 def identity_stability_receipt(
@@ -45,28 +70,140 @@ def identity_stability_receipt(
     }
 
 
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--input-h5", type=Path, required=True)
-    parser.add_argument("--output", type=Path, required=True)
-    args = parser.parse_args()
-    frame, _provenance = load_uk_national_frame(args.input_h5)
-    receipt = {
+def e4_identity_receipt(
+    frame,
+    *,
+    contract,
+    count_resource: Mapping[str, object],
+    lha_category: Sequence[object],
+    permutation_seed: int,
+) -> dict[str, object]:
+    """Recompute every E4 column in original and permuted row order.
+
+    Two claims are receipted: a row permutation of the input tables changes
+    no assignment per entity id, and the original-order recomputation equals
+    the columns stored in the artifact (re-derivation identity).
+    """
+
+    person = frame.table("person")
+    benunit = frame.table("benunit").copy()
+    household = frame.table("household")
+    if len(lha_category) != len(benunit):
+        raise ValueError("LHA_category materialization must align to benunit rows.")
+    benunit["LHA_category"] = [_enum_name(value) for value in lha_category]
+    benunit["region"] = _benunit_regions(person, household, benunit)
+
+    def recompute(person_t, benunit_t, household_t) -> dict[str, pd.DataFrame]:
+        anchors = aggregate_person_reported_to_benunit(person_t, benunit_t)
+        take_up = derive_frs_take_up(benunit_t, anchors=anchors, contract=contract)
+        take_up.index = benunit_t["benunit_id"].to_numpy()
+        person_draws = derive_frs_person_draws(person_t, contract=contract)
+        person_draws.index = person_t["person_id"].to_numpy()
+        household_draws = derive_frs_household_draws(household_t, contract=contract)
+        household_draws.index = household_t["household_id"].to_numpy()
+        seed = UK_BRMA_DECLARED_SEEDS["brma"]
+        benunit_brma = pd.DataFrame(
+            {
+                "benunit_id": benunit_t["benunit_id"].to_numpy(),
+                "brma": assign_brma_by_cell(
+                    benunit_t, count_resource=count_resource, seed=seed
+                ),
+            }
+        )
+        household_draws["brma"] = collapse_benunit_brma_to_household(
+            person_t, benunit_brma, household_t, seed=seed
+        )
+        return {
+            "benunit": take_up,
+            "person": person_draws,
+            "household": household_draws,
+        }
+
+    original = recompute(person, benunit, household)
+    rng = np.random.default_rng(permutation_seed)
+    permuted = recompute(
+        person.iloc[rng.permutation(len(person))].reset_index(drop=True),
+        benunit.iloc[rng.permutation(len(benunit))].reset_index(drop=True),
+        household.iloc[rng.permutation(len(household))].reset_index(drop=True),
+    )
+
+    stored = {
+        "person": frame.table("person").set_index("person_id"),
+        "benunit": frame.table("benunit").set_index("benunit_id"),
+        "household": frame.table("household").set_index("household_id"),
+    }
+    permutation_mismatches: dict[str, list[str]] = {}
+    stored_mismatches: dict[str, list[str]] = {}
+    stored_columns_missing: dict[str, list[str]] = {}
+    for entity, values in original.items():
+        for column in values.columns:
+            left = values[column]
+            right = permuted[entity][column].reindex(left.index)
+            if not np.array_equal(left.to_numpy(), right.to_numpy()):
+                permutation_mismatches.setdefault(entity, []).append(column)
+            if column not in stored[entity].columns:
+                stored_columns_missing.setdefault(entity, []).append(column)
+                continue
+            kept = stored[entity][column].reindex(left.index)
+            if not np.array_equal(
+                left.to_numpy(), kept.to_numpy().astype(left.to_numpy().dtype)
+            ):
+                stored_mismatches.setdefault(entity, []).append(column)
+    return {
         "check": "uk_e4_identity_stability",
-        "input_h5": str(args.input_h5),
-        "status": "requires caller-supplied E4 transform in acceptance harness",
+        "permutation_seed": permutation_seed,
+        "identical_under_permutation": not permutation_mismatches,
+        "matches_stored_columns": not stored_mismatches and not stored_columns_missing,
+        "permutation_mismatches": permutation_mismatches,
+        "stored_mismatches": stored_mismatches,
+        "stored_columns_missing": stored_columns_missing,
+        "columns_by_entity": {
+            entity: list(values.columns) for entity, values in original.items()
+        },
         "entity_row_counts": {
             entity: int(len(frame.table(entity))) for entity in frame.entities
         },
     }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--input-h5", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--permutation-seed", type=int, default=123)
+    args = parser.parse_args()
+
+    from microcosm.build.uk_runtime.take_up_contract import load_uk_take_up_contract
+    from microcosm.frame.adapters.policyengine_uk import PolicyEngineUKEngine
+
+    frame, _provenance = load_uk_national_frame(args.input_h5)
+    engine = PolicyEngineUKEngine()
+    lha_category = engine.materialize(frame, ("LHA_category",), uk_time_period(frame))[
+        "LHA_category"
+    ]
+    receipt = e4_identity_receipt(
+        frame,
+        contract=load_uk_take_up_contract(),
+        count_resource=load_brma_count_resource(),
+        lha_category=lha_category,
+        permutation_seed=args.permutation_seed,
+    )
+    receipt["input_h5"] = str(args.input_h5)
     args.output.write_text(json.dumps(receipt, indent=2, sort_keys=True) + "\n")
+    ok = bool(
+        receipt["identical_under_permutation"] and receipt["matches_stored_columns"]
+    )
+    print(
+        "identity stability:",
+        "PASS" if ok else f"FAIL ({args.output})",
+    )
+    return 0 if ok else 1
 
 
 def _reverse_rows(frame):
     from microcosm.build.uk_runtime.national_frame import (
         uk_household_weight_kind,
         uk_national_frame,
-        uk_time_period,
     )
 
     return uk_national_frame(
@@ -81,4 +218,4 @@ def _reverse_rows(frame):
 
 
 if __name__ == "__main__":
-    main()
+    sys.exit(main())