Compare BRMA distributions at the sampler's native cell margin

The distribution diagnostic compared household BRMA shares per region
against the rents table's own regional margin — mixing our benunit LHA
category composition with the table's lettings-weighted composition, which
read as double-digit z-scores on a correct sampler. The comparison now
re-derives the benunit-level assignment and checks each (region, LHA
category) cell against its conditional distribution, reporting per-row z
and the max |z|, with the same SDC masking.

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

Author: juaristi22

packages/microcosm-build/tests/test_uk_stochastic_tools.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import pandas as pd
+import pytest
 
 from microcosm.build.uk_runtime.national_frame import uk_national_frame
 
@@ -155,27 +156,48 @@ def continuous_entry(self, key: str):
     }
 
 
-def test_brma_distribution_masks_small_counts() -> None:
+def test_brma_cell_distribution_masks_small_counts_and_reports_z() -> None:
     tool = _load_tool("emit_uk_brma_distribution")
-    household = pd.DataFrame(
+    benunit = pd.DataFrame(
         {
-            "household_id": [1, 2, 3, 4],
-            "region": ["LONDON", "LONDON", "LONDON", "LONDON"],
-            "brma": ["A", "A", "A", "B"],
+            "benunit_id": [10, 20, 30, 40],
+            "region": ["LONDON"] * 4,
+            "LHA_category": ["A"] * 4,
+            "brma": ["CENTRAL", "CENTRAL", "CENTRAL", "OUTER"],
         }
     )
-    resource = {"cells": {"LONDON": {"A": {"A": 3, "B": 1}}}}
+    resource = {"cells": {"LONDON": {"A": {"CENTRAL": 1, "OUTER": 1}}}}
 
-    payload = tool.brma_distribution(
-        household,
+    payload = tool.brma_cell_distribution(
+        benunit,
         count_resource=resource,
         minimum_count=3,
     )
 
     rows = {row["brma"]: row for row in payload["rows"]}
-    assert rows["A"]["built_count"] == 3
-    assert rows["B"]["built_count"] == "<3"
-    assert rows["B"]["built_share"] is None
+    assert rows["CENTRAL"]["built_count"] == 3
+    assert rows["CENTRAL"]["cell_n"] == 4
+    assert rows["CENTRAL"]["z"] is not None
+    assert rows["OUTER"]["built_count"] == "<3"
+    assert rows["OUTER"]["built_share"] is None
+    assert rows["OUTER"]["z"] is None
+    assert payload["max_abs_z"] > 0
+
+
+def test_brma_cell_distribution_fails_closed_on_missing_cell() -> None:
+    tool = _load_tool("emit_uk_brma_distribution")
+    benunit = pd.DataFrame(
+        {
+            "benunit_id": [10],
+            "region": ["LONDON"],
+            "LHA_category": ["B"],
+            "brma": ["CENTRAL"],
+        }
+    )
+    resource = {"cells": {"LONDON": {"A": {"CENTRAL": 1}}}}
+
+    with pytest.raises(KeyError, match="LHA_category"):
+        tool.brma_cell_distribution(benunit, count_resource=resource)
 
 
 def _frame():

tools/emit_uk_brma_distribution.py

@@ -1,61 +1,94 @@
-"""Emit disclosure-safe BRMA distribution diagnostics for a UK frame."""
+"""Emit disclosure-safe BRMA distribution diagnostics for a UK frame.
+
+The comparison runs at the sampler's native margin — (region, LHA category)
+cells — because that is where the assignment is defined: within a cell the
+built benunit-level BRMA shares are multinomial around the count table's
+conditional distribution. A region-level margin would mix our benunit
+category composition with the rents table's lettings-weighted composition
+and mislead.
+"""
 
 from __future__ import annotations
 
 import argparse
 import json
+import math
 from collections.abc import Mapping
 from pathlib import Path
 
 import pandas as pd
 
-from microcosm.build.uk_runtime.frs_brma import load_brma_count_resource
+from microcosm.build.uk_runtime.frs_brma import (
+    UK_BRMA_DECLARED_SEEDS,
+    _benunit_regions,
+    _enum_name,
+    assign_brma_by_cell,
+    load_brma_count_resource,
+)
 from microcosm.build.uk_runtime.national_build import load_uk_national_frame
+from microcosm.build.uk_runtime.national_frame import uk_time_period
 
 
-def brma_distribution(
-    household: pd.DataFrame,
+def brma_cell_distribution(
+    benunit: pd.DataFrame,
     *,
     count_resource: Mapping[str, object],
     minimum_count: int = 3,
 ) -> dict[str, object]:
-    """Compare built household BRMA shares with count-table region priors."""
+    """Compare built benunit BRMA shares per (region, LHA category) cell."""
 
-    required = {"region", "brma"}
-    missing = required - set(household.columns)
+    required = {"region", "LHA_category", "brma"}
+    missing = required - set(benunit.columns)
     if missing:
-        raise ValueError(f"household table is missing column(s): {sorted(missing)}.")
-    rows: list[dict[str, object]] = []
+        raise ValueError(f"benunit table is missing column(s): {sorted(missing)}.")
     cells = count_resource["cells"]
-    for region, region_table in sorted(cells.items()):
-        expected_counts: dict[str, int] = {}
-        for category_counts in region_table.values():
-            for brma, count in category_counts.items():
-                expected_counts[brma] = expected_counts.get(brma, 0) + int(count)
-        built = household.loc[household["region"].astype(str) == str(region), "brma"]
-        built_counts = built.astype(str).value_counts().to_dict()
-        built_total = int(sum(built_counts.values()))
-        expected_total = int(sum(expected_counts.values()))
+    rows: list[dict[str, object]] = []
+    max_abs_z = 0.0
+    for (region, category), group in benunit.groupby(
+        ["region", "LHA_category"], sort=True
+    ):
+        expected_counts = cells.get(str(region), {}).get(str(category))
+        if expected_counts is None:
+            raise KeyError(
+                f"missing count-table cell for region={region!r}, "
+                f"LHA_category={category!r}."
+            )
+        expected_total = float(sum(expected_counts.values()))
+        built_counts = group["brma"].astype(str).value_counts().to_dict()
+        cell_n = int(len(group))
         for brma, expected_count in sorted(expected_counts.items()):
+            expected_share = float(expected_count) / expected_total
             built_count = int(built_counts.get(brma, 0))
+            row: dict[str, object] = {
+                "region": str(region),
+                "lha_category": str(category),
+                "brma": str(brma),
+                "cell_n": cell_n,
+                "expected_share": expected_share,
+            }
             if 0 < built_count < minimum_count:
-                built_count_out: int | str = f"<{minimum_count}"
-                built_share: float | None = None
+                row["built_count"] = f"<{minimum_count}"
+                row["built_share"] = None
+                row["z"] = None
             else:
-                built_count_out = built_count
-                built_share = built_count / built_total if built_total else 0.0
-            rows.append(
-                {
-                    "region": region,
-                    "brma": brma,
-                    "built_count": built_count_out,
-                    "built_share": built_share,
-                    "count_table_share": expected_count / expected_total,
-                }
-            )
+                built_share = built_count / cell_n if cell_n else 0.0
+                sigma = (
+                    math.sqrt(expected_share * (1.0 - expected_share) / cell_n)
+                    if cell_n and 0.0 < expected_share < 1.0
+                    else None
+                )
+                z = (built_share - expected_share) / sigma if sigma else None
+                row["built_count"] = built_count
+                row["built_share"] = built_share
+                row["z"] = z
+                if z is not None:
+                    max_abs_z = max(max_abs_z, abs(z))
+            rows.append(row)
     return {
-        "check": "uk_brma_distribution",
+        "check": "uk_brma_cell_distribution",
+        "margin": "benunit BRMA within (region, LHA category) cells",
         "minimum_count": minimum_count,
+        "max_abs_z": max_abs_z,
         "rows": rows,
     }
 
@@ -66,13 +99,34 @@ def main() -> None:
     parser.add_argument("--output", type=Path, required=True)
     parser.add_argument("--minimum-count", type=int, default=3)
     args = parser.parse_args()
+
+    from microcosm.frame.adapters.policyengine_uk import PolicyEngineUKEngine
+
     frame, _provenance = load_uk_national_frame(args.input_h5)
-    payload = brma_distribution(
-        frame.table("household"),
-        count_resource=load_brma_count_resource(),
+    engine = PolicyEngineUKEngine()
+    lha_category = engine.materialize(frame, ("LHA_category",), uk_time_period(frame))[
+        "LHA_category"
+    ]
+    benunit = frame.table("benunit").copy()
+    benunit["LHA_category"] = [_enum_name(value) for value in lha_category]
+    benunit["region"] = _benunit_regions(
+        frame.table("person"), frame.table("household"), benunit
+    )
+    count_resource = load_brma_count_resource()
+    # Deterministic re-derivation of the benunit-level assignment that fed the
+    # stored household collapse (proven identical by the identity receipt).
+    benunit["brma"] = assign_brma_by_cell(
+        benunit,
+        count_resource=count_resource,
+        seed=UK_BRMA_DECLARED_SEEDS["brma"],
+    )
+    payload = brma_cell_distribution(
+        benunit,
+        count_resource=count_resource,
         minimum_count=args.minimum_count,
     )
     args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+    print("brma cell distribution: max |z| =", round(payload["max_abs_z"], 2))
 
 
 if __name__ == "__main__":