From f227bfc50e8477b03848a962778bf68531144e10 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 12:34:40 -0400 Subject: [PATCH 1/6] Split private renters across BRMAs and tie the split in calibration One BRMA draw per household makes published aggregates depend on the draw: a few heavily weighted private renters carry most of its variance, and calibration passes it on. Each private-renting household now becomes k records (PE_UK_DATA_BRMA_SPLIT, default 4) with an equal share of its weight, at evenly spaced quantiles of its BRMA distribution walked in LHA rate order. Calibration ties each household's records (brma_split_group): group-mean target rows, one shared prior, equal saved weights. Co-Authored-By: Claude Opus 5.5 --- .../brma-split-private-renters.changed.md | 1 + docs/imputations.md | 7 + policyengine_uk_data/datasets/brma.py | 252 ++++++++++++- .../datasets/create_datasets.py | 30 ++ policyengine_uk_data/tests/test_brma_split.py | 343 ++++++++++++++++++ policyengine_uk_data/utils/calibrate.py | 53 ++- 6 files changed, 681 insertions(+), 5 deletions(-) create mode 100644 changelog.d/brma-split-private-renters.changed.md create mode 100644 policyengine_uk_data/tests/test_brma_split.py diff --git a/changelog.d/brma-split-private-renters.changed.md b/changelog.d/brma-split-private-renters.changed.md new file mode 100644 index 00000000..cc5f9233 --- /dev/null +++ b/changelog.d/brma-split-private-renters.changed.md @@ -0,0 +1 @@ +- Split each private-renting household in the enhanced FRS into 4 records (`PE_UK_DATA_BRMA_SPLIT`) at evenly spaced, LHA-rate-ordered quantiles of its BRMA distribution, with calibration tying the records' weights (`brma_split_group`). This cuts the noise one BRMA draw puts into Universal Credit, Housing Benefit and AHC poverty. diff --git a/docs/imputations.md b/docs/imputations.md index 75fd1a56..17794821 100644 --- a/docs/imputations.md +++ b/docs/imputations.md @@ -222,6 +222,13 @@ The FRS identifies only the region, but Local Housing Allowance rates vary by Br A household takes one of its benefit units' BRMAs, chosen at random. Sources, method and validation are in `storage/BRMA_DATA_SOURCES.md`. +One draw per household is noisy: a few heavily weighted private renters carry most of the variation, and calibration passes it on. So the enhanced dataset splits each private-renting household into `PE_UK_DATA_BRMA_SPLIT` records (default 4) just before the output-area clone step: +- each record has an equal share of the household's weight; +- each record takes the BRMA at one of evenly spaced quantiles of the household's BRMA distribution, so each record has exactly the distribution of a single draw; +- the quantiles walk the region's BRMAs from the lowest LHA rate to the highest, so the records' average outcome stays within a fraction of the household's range of outcomes of its expectation. + +The records share `brma_split_group` (the original household id). Calibration treats each group as one household: its records take the group's mean target rows, share one prior, and keep equal weights. + --- ## Calibration Targets diff --git a/policyengine_uk_data/datasets/brma.py b/policyengine_uk_data/datasets/brma.py index 9aec9bc6..4764bf3a 100644 --- a/policyengine_uk_data/datasets/brma.py +++ b/policyengine_uk_data/datasets/brma.py @@ -1,4 +1,4 @@ -"""Broad Rental Market Area (BRMA) assignment for FRS benefit units. +"""Broad Rental Market Area (BRMA) assignment for FRS households. The FRS identifies only a household's region, but Local Housing Allowance rates vary by BRMA. Each benefit unit is given a BRMA drawn within its @@ -14,6 +14,18 @@ These weights replace the row counts of the 2019-20 LHA list of rents. Its Scottish, Welsh and Northern Ireland rows were copies of English BRMAs' lists, so their counts said nothing about those nations' rental markets. + +The base FRS dataset gets one draw per household (``assign_brmas`` and +``pick_household_brmas``). One draw is noisy: a few heavily weighted private +renters carry most of the variation, so a different seed moved UK Universal +Credit by about £0.3bn. The enhanced dataset therefore splits each +private-renting household into ``k`` records (``split_private_renters_across_brmas``), +each with ``1/k`` of its weight and one of ``k`` evenly spaced quantiles of its +BRMA distribution, walking its region's BRMAs from the lowest LHA rate to the +highest. Each record keeps exactly the distribution of a single draw; the +records' average follows the expectation far more closely. Calibration ties +the ``k`` records of a household (``BRMA_SPLIT_GROUP_COLUMN``) so that they +keep equal weights. """ from pathlib import Path @@ -28,6 +40,9 @@ # Census bedroom band whose private-rented households weight each LHA category. LHA_CATEGORY_BEDROOMS = {"A": "1", "B": "1", "C": "2", "D": "3", "E": "4+"} +# Household column holding the id of the household a split record came from. +BRMA_SPLIT_GROUP_COLUMN = "brma_split_group" + def load_brma_weights(path: Path = BRMA_HOUSEHOLDS_PATH) -> pd.DataFrame: """Return BRMA sampling weights by region and LHA category. @@ -104,3 +119,238 @@ def pick_household_brmas( return units.groupby("household_id").brma.aggregate( lambda x: x.sample(n=1, random_state=rng).iloc[0] ) + + +def household_brma_probabilities( + region: np.ndarray, + benunit_household: np.ndarray, + benunit_category: np.ndarray, + weights: pd.DataFrame | None = None, +) -> tuple[dict[str, list[str]], np.ndarray]: + """Each household's BRMA distribution under ``assign_brmas`` and ``pick_household_brmas``. + + Each benefit unit draws from its region × LHA category cell and the + household takes one unit's BRMA at random, so the household's distribution + is the mean of its units' cell distributions. + + Args: + region: Region name of each household. + benunit_household: Index into ``region`` of each benefit unit's household. + benunit_category: LHA category (A-E) of each benefit unit. + weights: Output of ``load_brma_weights``; loaded from storage if omitted. + + Returns: + ``(brmas, probabilities)``: ``brmas`` maps each region to its BRMA names + in alphabetical order; ``probabilities[h, j]`` is household h's + probability of the j-th BRMA of its region (zero past the region's count). + + Raises: + ValueError: If a household has no benefit unit, or a unit's cell has no + BRMA with a positive weight. + """ + if weights is None: + weights = load_brma_weights() + region = np.asarray(region).astype(str) + benunit_household = np.asarray(benunit_household) + benunit_category = np.asarray(benunit_category).astype(str) + brmas = {r: sorted(g.brma.unique()) for r, g in weights.groupby("region")} + width = max(len(names) for names in brmas.values()) + probabilities = np.zeros((len(region), width)) + units = np.bincount(benunit_household, minlength=len(region)) + unit_region = region[benunit_household] + covered = np.zeros(len(benunit_household), bool) + for (cell_region, category), cell in weights.groupby(["region", "lha_category"]): + mask = (unit_region == cell_region) & (benunit_category == category) + if not mask.any(): + continue + p = cell.set_index("brma").weight.reindex(brmas[cell_region]).fillna(0) + p = np.pad(p.to_numpy(float) / p.sum(), (0, width - len(p))) + np.add.at(probabilities, benunit_household[mask], p) + covered |= mask + if (units == 0).any(): + raise ValueError("Every household needs at least one benefit unit.") + if not covered.all(): + cells = sorted(set(zip(unit_region[~covered], benunit_category[~covered]))) + raise ValueError(f"No BRMA weights for region × LHA category cells {cells}.") + return brmas, probabilities / units[:, None] + + +def lha_rate_keys( + simulation, + year: int, + brmas: dict[str, list[str]], + region: np.ndarray, + benunit_household: np.ndarray, +) -> np.ndarray: + """Each household's LHA rate in each BRMA of its region. + + The mean, over the household's benefit units, of PolicyEngine UK's + ``uncapped_BRMA_LHA_rate`` (``BRMA_LHA_rate`` in versions without it) with + the household moved to that BRMA; ``inf`` past the region's count. It only + orders BRMAs for the spaced draws: any order keeps every household's + distribution exact. + """ + region = np.asarray(region).astype(str) + count = np.array([len(brmas[r]) for r in region]) + width = max(len(names) for names in brmas.values()) + units = np.maximum(np.bincount(benunit_household, minlength=len(region)), 1) + keys = np.full((len(region), width), np.inf) + variable = "uncapped_BRMA_LHA_rate" + if variable not in simulation.tax_benefit_system.variables: + variable = "BRMA_LHA_rate" + for j in range(width): + slot = np.minimum(j, count - 1) + names = np.array([brmas[r][i] for r, i in zip(region, slot)]) + simulation.set_input("brma", year, names) + simulation.delete_arrays(variable) + rate = np.asarray(simulation.calculate(variable, year)) + mean = np.bincount(benunit_household, weights=rate, minlength=len(region)) + keys[:, j] = np.where(j < count, mean / units, np.inf) + simulation.delete_arrays(variable) + return keys + + +def brmas_at_quantiles( + probabilities: np.ndarray, keys: np.ndarray | None, quantiles: np.ndarray +) -> np.ndarray: + """Column of each row's distribution at its quantile in [0, 1). + + Columns are walked in ascending ``keys`` order (ties, and ``keys=None``, by + column). A column is returned only where its probability is positive: the + cumulative probability is set to exactly 1 from the last positive column + on, so rounding cannot carry a quantile past it. + """ + quantiles = np.asarray(quantiles, float) + if np.any((quantiles < 0) | (quantiles >= 1)): + raise ValueError("Quantiles must lie in [0, 1).") + rows, width = probabilities.shape + if keys is None: + order = np.broadcast_to(np.arange(width), probabilities.shape) + else: + order = np.argsort( + np.where(probabilities > 0, keys, np.inf), axis=1, kind="stable" + ) + p = np.take_along_axis(probabilities, order, axis=1) + cumulative = np.cumsum(p, axis=1) + last = width - 1 - np.argmax(p[:, ::-1] > 0, axis=1) + cumulative[np.arange(width)[None, :] >= last[:, None]] = 1.0 + position = (cumulative <= quantiles[:, None]).sum(axis=1) + return order[np.arange(rows), position] + + +def spaced_quantiles(rows: int, k: int, rng: np.random.Generator) -> np.ndarray: + """``(rows, k)`` quantiles ``(u + i / k) mod 1`` with one uniform ``u`` per row. + + Each quantile is uniform on [0, 1), and a row's ``k`` quantiles fall one + in each interval of width ``1 / k``. + """ + if k < 1: + raise ValueError("k must be at least 1.") + u = rng.random(rows) + return (u[:, None] + np.arange(k)[None, :] / k) % 1.0 + + +def split_private_renters_across_brmas( + dataset, k: int, seed: int = 0, weights: pd.DataFrame | None = None +): + """Split each private-renting household into ``k`` records across its BRMAs. + + Each record has ``1/k`` of the household's weight and the BRMA at one of + ``k`` evenly spaced quantiles of the household's distribution + (``household_brma_probabilities``), walking its region's BRMAs from the + lowest LHA rate to the highest. Other households keep their record and + BRMA. Every household gets ``BRMA_SPLIT_GROUP_COLUMN``: the id of the + household its record came from. + + Copies offset household ids by ``i * step``, benefit unit ids by + ``i * step * 100`` and person ids by ``i * step * 1000`` (``step`` a power + of ten above every household id, raised if needed so no id collides), so + ``benunit_id // 100 == household_id`` still holds wherever it held. + + Returns: + The new dataset. With ``k == 1`` it is the input plus the group column. + """ + from policyengine_uk import Microsimulation + from policyengine_uk.data import UKSingleYearDataset + + household = dataset.household.copy() + household[BRMA_SPLIT_GROUP_COLUMN] = household.household_id.to_numpy() + if k == 1: + return UKSingleYearDataset( + person=dataset.person, + benunit=dataset.benunit, + household=household, + fiscal_year=dataset.time_period, + ) + year = int(dataset.time_period) + simulation = Microsimulation(dataset=dataset) + region = np.asarray(simulation.calculate("region", year)).astype(str) + renter = ( + np.asarray(simulation.calculate("tenure_type", year)).astype(str) + == "RENT_PRIVATELY" + ) + benunit_household = pd.Series( + np.arange(len(household)), index=household.household_id.to_numpy() + ).loc[np.asarray(simulation.populations["benunit"].household("household_id", year))] + benunit_household = benunit_household.to_numpy() + category = np.asarray(simulation.calculate("LHA_category", year)).astype(str) + brmas, probabilities = household_brma_probabilities( + region, benunit_household, category, weights + ) + keys = lha_rate_keys(simulation, year, brmas, region, benunit_household) + columns = brmas_at_quantiles( + np.repeat(probabilities[renter], k, axis=0), + np.repeat(keys[renter], k, axis=0), + spaced_quantiles(int(renter.sum()), k, np.random.default_rng(seed)).ravel(), + ).reshape(-1, k) + names = np.array( + [[brmas[r][j] for j in row] for r, row in zip(region[renter], columns)], + dtype=object, + ).reshape(-1, k) + + person, benunit = dataset.person, dataset.benunit + renter_ids = household.household_id.to_numpy()[renter] + benunit_of = ( + person.drop_duplicates("person_benunit_id") + .set_index("person_benunit_id") + .person_household_id + ) + renter_benunits = benunit[ + np.isin(benunit_of.loc[benunit.benunit_id].to_numpy(), renter_ids) + ] + renter_people = person[person.person_household_id.isin(renter_ids)] + step = 10 ** len(str(int(household.household_id.max()))) + benunit_step = max(step * 100, 10 ** len(str(int(benunit.benunit_id.max())))) + person_step = max(step * 1000, 10 ** len(str(int(person.person_id.max())))) + + base = household.copy() + base.loc[renter, "brma"] = names[:, 0] + base.loc[renter, "household_weight"] = household.household_weight[renter] / k + households, benunits, people = [base], [benunit], [person] + for i in range(1, k): + copy = base[renter].copy() + copy["brma"] = names[:, i] + copy["household_id"] += i * step + households.append(copy) + units = renter_benunits.copy() + units["benunit_id"] += i * benunit_step + benunits.append(units) + members = renter_people.copy() + members["person_household_id"] += i * step + members["person_benunit_id"] += i * benunit_step + members["person_id"] += i * person_step + people.append(members) + result = UKSingleYearDataset( + person=pd.concat(people, ignore_index=True), + benunit=pd.concat(benunits, ignore_index=True), + household=pd.concat(households, ignore_index=True), + fiscal_year=dataset.time_period, + ) + for table, column in ( + (result.household, "household_id"), + (result.benunit, "benunit_id"), + (result.person, "person_id"), + ): + if not table[column].is_unique: + raise ValueError(f"Splitting private renters duplicated {column}.") + return result diff --git a/policyengine_uk_data/datasets/create_datasets.py b/policyengine_uk_data/datasets/create_datasets.py index 31008183..f3a12bb4 100644 --- a/policyengine_uk_data/datasets/create_datasets.py +++ b/policyengine_uk_data/datasets/create_datasets.py @@ -1,6 +1,8 @@ import logging import os +import pandas as pd + from policyengine_uk_data.utils.build_environment import ( assert_local_build_environment, ) @@ -17,6 +19,10 @@ VALIDATION_ONLY_LOCAL_TARGETS = ["hmrc/employment_income/count"] +# Records per private-renting household across its BRMAs (datasets/brma.py). +BRMA_SPLIT = 4 + + def _get_positive_int_env(name: str, default: int) -> int: raw_value = os.environ.get(name) if raw_value is None: @@ -82,6 +88,8 @@ def main(): "PE_UK_DATA_OA_CLONES", 2 if is_testing else 10, ) + # Records per private-renting household across its BRMAs (1 = one draw). + brma_split = _get_positive_int_env("PE_UK_DATA_BRMA_SPLIT", BRMA_SPLIT) frs_release = CURRENT_FRS_RELEASE align_to_base_year = frs_release.base_year != frs_release.survey_year align_step = f"Align to {frs_release.base_year} base year" @@ -110,6 +118,7 @@ def main(): "Impute salary sacrifice", "Impute student loan plan", "Assign Pension Credit take-up", + "Split private renters across BRMAs", "Clone and assign OA geography", "Calibrate constituency weights", "Calibrate local authority weights", @@ -224,6 +233,15 @@ def main(): ) update_dataset("Assign Pension Credit take-up", "completed") + update_dataset("Split private renters across BRMAs", "processing") + from policyengine_uk_data.datasets.brma import ( + BRMA_SPLIT_GROUP_COLUMN, + split_private_renters_across_brmas, + ) + + frs = split_private_renters_across_brmas(frs, k=brma_split) + update_dataset("Split private renters across BRMAs", "completed") + # Clone households and assign OA geography update_dataset("Clone and assign OA geography", "processing") from policyengine_uk_data.calibration.clone_and_assign import ( @@ -265,6 +283,16 @@ def main(): get_performance, ) + # A split household's records (within one OA clone) calibrate as one. + brma_groups = pd.factorize( + pd.MultiIndex.from_arrays( + [ + frs.household[BRMA_SPLIT_GROUP_COLUMN].to_numpy(), + frs.household["clone_index"].to_numpy(), + ] + ) + )[0] + # Run calibration with verbose progress frs_calibrated_constituencies = calibrate_local_areas( dataset=frs, @@ -281,6 +309,7 @@ def main(): get_performance=get_performance, nested_progress=nested_progress, # Pass the nested progress manager time_period=frs_release.calibration_year, + groups=brma_groups, ) update_dataset("Calibrate constituency weights", "completed") @@ -308,6 +337,7 @@ def main(): get_performance=get_la_performance, nested_progress=nested_progress, # Pass the nested progress manager time_period=frs_release.calibration_year, + groups=brma_groups, ) update_dataset("Calibrate local authority weights", "completed") diff --git a/policyengine_uk_data/tests/test_brma_split.py b/policyengine_uk_data/tests/test_brma_split.py new file mode 100644 index 00000000..29d4b66a --- /dev/null +++ b/policyengine_uk_data/tests/test_brma_split.py @@ -0,0 +1,343 @@ +"""Splitting private renters across BRMAs, and the calibration that ties them. + +Invariants, for any distributions, keys and generator state: + +1. Support: a quantile never selects a BRMA with zero probability, and + always a column of the household's region. +2. Exact marginals: each record's BRMA has exactly its household's + distribution (``household_brma_probabilities``), whatever the walk order. +3. Spacing: a household's ``k`` quantiles fall one in each interval of width + ``1/k``, so walking BRMAs in increasing order of an outcome, the records' + mean of that outcome is within ``range / k`` of its expectation. +4. Conservation: a split household's records carry exactly its weight, people + and benefit units; other households are untouched; every id stays unique; + ``benunit_id // 100 == household_id`` holds wherever it held. +5. Determinism: the same seed gives the same records. +6. Tying: a group's calibrated weights are equal, its target rows identical, + and its zero-weight rows share one household's prior. +""" + +import importlib.util + +import numpy as np +import pandas as pd +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st +from hypothesis.extra import numpy as hnp + +from policyengine_uk_data.datasets.brma import ( + BRMA_SPLIT_GROUP_COLUMN, + brmas_at_quantiles, + household_brma_probabilities, + load_brma_weights, + spaced_quantiles, +) +from policyengine_uk_data.utils.calibrate import ( + initialize_weight_priors, + tie_rows, + tie_weights, +) + + +@st.composite +def distributions(draw, max_rows=6, max_width=7): + """Rows of probabilities with some zero columns, and keys to order them.""" + rows = draw(st.integers(1, max_rows)) + width = draw(st.integers(1, max_width)) + raw = draw( + hnp.arrays( + float, + (rows, width), + elements=st.sampled_from([0.0, 0.0, 1e-9, 0.1, 0.5, 1.0, 3.0, 1e6]), + ) + ) + raw[raw.sum(axis=1) == 0, 0] = 1.0 # every row needs some probability + keys = draw( + hnp.arrays(float, (rows, width), elements=st.floats(-5, 5, allow_nan=False)) + ) + return raw / raw.sum(axis=1, keepdims=True), keys + + +def exact_frequencies(probabilities, keys, grid=20_000): + """Share of a fine quantile grid landing on each column, per row.""" + tau = (np.arange(grid) + 0.5) / grid + out = np.zeros_like(probabilities) + for row in range(len(probabilities)): + columns = brmas_at_quantiles( + np.repeat(probabilities[row : row + 1], grid, axis=0), + None if keys is None else np.repeat(keys[row : row + 1], grid, axis=0), + tau, + ) + out[row] = np.bincount(columns, minlength=probabilities.shape[1]) / grid + return out + + +@settings(max_examples=60, deadline=None, derandomize=True) +@given(distributions(), st.booleans()) +def test_quantiles_have_exact_marginals_and_support(case, ordered): + probabilities, keys = case + keys = keys if ordered else None + frequencies = exact_frequencies(probabilities, keys) + # A grid of n points matches each column's probability to within 1/n per + # boundary; zero-probability columns are never chosen. + assert np.abs(frequencies - probabilities).max() <= 2 / 20_000 + 1e-12 + assert (frequencies[probabilities == 0] == 0).all() + + +@settings(max_examples=200, deadline=None, derandomize=True) +@given(distributions(), st.floats(0, 1, exclude_max=True)) +def test_quantiles_never_leave_the_support(case, tau): + probabilities, keys = case + for t in (tau, np.nextafter(1.0, 0.0), 0.0): + columns = brmas_at_quantiles(probabilities, keys, np.full(len(keys), t)) + assert (probabilities[np.arange(len(columns)), columns] > 0).all() + + +@settings(max_examples=100, deadline=None, derandomize=True) +@given(distributions()) +def test_walk_follows_the_keys(case): + probabilities, keys = case + tau = np.linspace(0, 1, 50, endpoint=False) + for row in range(len(probabilities)): + columns = brmas_at_quantiles( + np.repeat(probabilities[row : row + 1], len(tau), axis=0), + np.repeat(keys[row : row + 1], len(tau), axis=0), + tau, + ) + assert (np.diff(keys[row][columns]) >= 0).all() + + +@settings(max_examples=100, deadline=None, derandomize=True) +@given(distributions(), st.integers(1, 9), st.integers(0, 2**32 - 1)) +def test_spaced_records_track_the_expectation(case, k, seed): + probabilities, keys = case + quantiles = spaced_quantiles(len(probabilities), k, np.random.default_rng(seed)) + assert ((quantiles >= 0) & (quantiles < 1)).all() + # One quantile in each interval [i/k, (i+1)/k). + assert (np.sort(np.floor(quantiles * k), axis=1) == np.arange(k)).all() + columns = brmas_at_quantiles( + np.repeat(probabilities, k, axis=0), np.repeat(keys, k, axis=0), quantiles.ravel() + ).reshape(-1, k) + # Outcome = the key itself (any outcome monotone in the walk order). + outcome = np.where(probabilities > 0, keys, 0) + mean = np.take_along_axis(outcome, columns, axis=1).mean(axis=1) + expected = (probabilities * outcome).sum(axis=1) + support = np.where(probabilities > 0, keys, np.nan) + spread = np.nanmax(support, axis=1) - np.nanmin(support, axis=1) + assert (np.abs(mean - expected) <= spread / k + 1e-9).all() + + +def test_spaced_quantiles_reject_k_below_one(): + with pytest.raises(ValueError): + spaced_quantiles(3, 0, np.random.default_rng(0)) + + +@pytest.fixture(scope="module") +def weights(): + return load_brma_weights() + + +def test_household_distribution_is_the_mean_of_its_benefit_units(weights): + region = np.array(["SCOTLAND", "SCOTLAND", "NORTHERN_IRELAND"]) + benunit_household = np.array([0, 1, 1, 2]) + category = np.array(["C", "A", "E", "B"]) + brmas, p = household_brma_probabilities(region, benunit_household, category, weights) + assert np.allclose(p.sum(axis=1), 1) + + def cell(r, c): + g = weights[(weights.region == r) & (weights.lha_category == c)] + v = g.set_index("brma").weight.reindex(brmas[r]).fillna(0).to_numpy() + return np.pad(v / v.sum(), (0, p.shape[1] - len(v))) + + assert np.allclose(p[0], cell("SCOTLAND", "C")) + assert np.allclose(p[1], (cell("SCOTLAND", "A") + cell("SCOTLAND", "E")) / 2) + assert np.allclose(p[2], cell("NORTHERN_IRELAND", "B")) + assert brmas["SCOTLAND"] == sorted(brmas["SCOTLAND"]) + assert (p[:, len(brmas["NORTHERN_IRELAND"]) :][2] == 0).all() + + +def test_household_distribution_fails_closed(weights): + table = weights[~((weights.region == "WALES") & (weights.lha_category == "D"))] + with pytest.raises(ValueError, match="No BRMA weights"): + household_brma_probabilities( + np.array(["WALES"]), np.array([0]), np.array(["D"]), table + ) + with pytest.raises(ValueError, match="benefit unit"): + household_brma_probabilities( + np.array(["WALES", "WALES"]), np.array([0]), np.array(["C"]), weights + ) + + +# --- splitting a dataset (runs PolicyEngine UK on a few synthetic households) --- + +needs_policyengine = pytest.mark.skipif( + importlib.util.find_spec("policyengine_uk") is None, + reason="policyengine_uk not available", +) + + +def toy_dataset(): + from policyengine_uk.data import UKSingleYearDataset + + regions = ["LONDON", "LONDON", "SCOTLAND", "WALES", "NORTH_WEST", "LONDON"] * 2 + tenure = ["RENT_PRIVATELY", "OWNED_OUTRIGHT", "RENT_PRIVATELY"] * 4 + n = len(regions) + household = pd.DataFrame( + { + "household_id": np.arange(1, n + 1), + "household_weight": np.r_[np.linspace(100, 2_000, n - 2), 0.0, 0.0], + "region": regions, + "tenure_type": tenure, + "rent": np.linspace(4_000, 30_000, n), + "brma": ["INNER_NORTH_LONDON"] * n, + "council_tax": np.full(n, 1_500.0), + } + ) + people, units = [], [] + for h in household.household_id: + # Households 3 and 6 hold two benefit units (a lodger family). + for b in range(1, 3 if h % 3 == 0 else 2): + units.append({"benunit_id": h * 100 + b}) + for p in range(1, 3): + people.append( + { + "person_id": h * 1000 + b * 10 + p, + "person_household_id": h, + "person_benunit_id": h * 100 + b, + "age": 25 + 10 * b + p, + } + ) + return UKSingleYearDataset( + person=pd.DataFrame(people), + benunit=pd.DataFrame(units), + household=household, + fiscal_year=2024, + ) + + +@needs_policyengine +@pytest.mark.parametrize("k", [1, 3]) +def test_split_conserves_households_and_keeps_ids_unique(k, weights): + from policyengine_uk_data.datasets.brma import split_private_renters_across_brmas + + dataset = toy_dataset() + result = split_private_renters_across_brmas(dataset, k=k, seed=7, weights=weights) + before, after = dataset.household, result.household + renter = before.tenure_type.astype(str) == "RENT_PRIVATELY" + + group = after[BRMA_SPLIT_GROUP_COLUMN] + assert set(group) == set(before.household_id) + sizes = group.value_counts() + assert (sizes.loc[before.household_id[renter]] == k).all() + assert (sizes.loc[before.household_id[~renter]] == 1).all() + # Each household's weight is conserved across its records. + weight = after.groupby(group).household_weight.sum() + assert np.allclose(weight.loc[before.household_id], before.household_weight) + # Non-renters keep their record exactly. + kept = after.set_index("household_id").loc[before.household_id[~renter]] + assert (kept.brma.astype(str) == "INNER_NORTH_LONDON").all() + # Every record's BRMA is a positive-weight BRMA of its own region. + supported = set(zip(weights.region, weights.brma)) + renters_after = after[after.tenure_type.astype(str) == "RENT_PRIVATELY"] + if k > 1: + assert all( + (r, b) in supported + for r, b in zip(renters_after.region.astype(str), renters_after.brma.astype(str)) + ) + for table, column in ( + (result.household, "household_id"), + (result.benunit, "benunit_id"), + (result.person, "person_id"), + ): + assert table[column].is_unique + assert set(result.person.person_household_id) == set(after.household_id) + assert set(result.person.person_benunit_id) == set(result.benunit.benunit_id) + # The id scheme benunit_id // 100 == household_id survives the copies. + owner = result.person.drop_duplicates("person_benunit_id").set_index( + "person_benunit_id" + ).person_household_id + assert (result.benunit.benunit_id // 100 == owner.loc[result.benunit.benunit_id].to_numpy()).all() + assert len(result.person) == len(dataset.person) + (k - 1) * len( + dataset.person[dataset.person.person_household_id.isin(before.household_id[renter])] + ) + + +@needs_policyengine +def test_split_is_deterministic(weights): + from policyengine_uk_data.datasets.brma import split_private_renters_across_brmas + + a = split_private_renters_across_brmas(toy_dataset(), k=4, seed=3, weights=weights) + b = split_private_renters_across_brmas(toy_dataset(), k=4, seed=3, weights=weights) + pd.testing.assert_frame_equal(a.household, b.household) + pd.testing.assert_frame_equal(a.person, b.person) + + +# --- calibration ties --- + + +@settings(max_examples=100, deadline=None, derandomize=True) +@given( + st.lists(st.integers(0, 4), min_size=1, max_size=12), + st.integers(1, 3), + st.integers(0, 2**32 - 1), +) +def test_tied_rows_and_weights_are_group_means_and_keep_totals(labels, areas, seed): + rng = np.random.default_rng(seed) + groups = np.array(labels) + rows = rng.random((len(groups), 3)) + tied = tie_rows(pd.DataFrame(rows), groups) + for g in set(labels): + m = groups == g + assert np.allclose(tied.to_numpy()[m], rows[m].mean(axis=0)) + w = rng.random((areas, len(groups))) + tw = tie_weights(w, groups) + for g in set(labels): + m = groups == g + assert np.allclose(tw[:, m], tw[:, m][:, :1]) + assert np.allclose(tw[:, m].sum(axis=1), w[:, m].sum(axis=1)) + # Predictions w @ rows are unchanged once rows are tied. + assert np.allclose(tw @ tied.to_numpy(), w @ tied.to_numpy()) + + +def test_split_zero_weight_rows_share_one_households_prior(): + weights = np.array([1_500.0, 0.0, 0.0, 0.0, 625.0, 0.0]) + groups = np.array([1, 2, 2, 2, 3, 4]) # rows 1-3 split one SPI household + priors = initialize_weight_priors(weights, groups=groups) + ungrouped = initialize_weight_priors(np.array([1_500.0, 0.0, 625.0, 0.0])) + assert priors[[1, 2, 3]].sum() == pytest.approx(ungrouped[1]) + assert priors[5] == pytest.approx(ungrouped[3]) + assert np.allclose(priors[[0, 4]], ungrouped[[0, 2]]) + assert priors.sum() == pytest.approx(weights.sum()) + + +@pytest.mark.skipif(importlib.util.find_spec("torch") is None, reason="needs torch") +def test_calibration_keeps_a_groups_records_equal(tmp_path, monkeypatch): + from policyengine_uk_data.utils import calibrate as calibrate_module + from policyengine_uk_data.utils.calibrate import calibrate_local_areas + + monkeypatch.setattr(calibrate_module, "STORAGE_FOLDER", tmp_path) + # Records 0-2 split one household; their rows differ before tying. + matrix = pd.DataFrame([[1.0, 0.0], [3.0, 0.0], [2.0, 0.0], [0.0, 1.0], [0.0, 2.0]]) + local = pd.DataFrame([[600.0, 300.0], [500.0, 400.0]]) + mask = np.ones((2, 5)) + + class Data: + def __init__(self, w): + self.household = pd.DataFrame({"household_weight": w}) + + def copy(self): + return Data(self.household.household_weight.to_numpy().copy()) + + result = calibrate_local_areas( + dataset=Data(np.array([100.0, 100.0, 100.0, 300.0, 300.0])), + matrix_fn=lambda d: (matrix.copy(), local.copy(), mask.copy()), + national_matrix_fn=lambda d: (pd.DataFrame(np.ones((5, 1))), pd.Series([1_500.0])), + area_count=2, + weight_file="tied.h5", + dataset_key="2024", + epochs=21, + groups=np.array([0, 0, 0, 1, 2]), + ) + w = result.household.household_weight.to_numpy() + assert w[0] == pytest.approx(w[1]) and w[1] == pytest.approx(w[2]) diff --git a/policyengine_uk_data/utils/calibrate.py b/policyengine_uk_data/utils/calibrate.py index f97c4e1e..c6d2bdfb 100644 --- a/policyengine_uk_data/utils/calibrate.py +++ b/policyengine_uk_data/utils/calibrate.py @@ -110,6 +110,7 @@ def load_weights( def initialize_weight_priors( original_weights: np.ndarray, zero_weight_total_share: float = DEFAULT_ZERO_WEIGHT_PRIOR_TOTAL_SHARE, + groups: np.ndarray | None = None, ) -> np.ndarray: """Build deterministic positive household priors for calibration. @@ -119,6 +120,10 @@ def initialize_weight_priors( zero-weight rows are present, preserve the relative distribution of positive survey weights while reserving a fixed share of total prior mass for the zero-weight rows. + + ``groups`` (one label per row) marks rows that split one household, such + as a private renter's BRMA records: zero-weight rows of a group then share + one household's prior between them. """ weights = np.asarray(original_weights, dtype=np.float64) if weights.ndim != 1: @@ -141,10 +146,34 @@ def initialize_weight_priors( priors = np.empty_like(weights, dtype=np.float64) priors[positive_mask] = weights[positive_mask] * (1 - zero_weight_total_share) - priors[zero_mask] = positive_total * zero_weight_total_share / zero_mask.sum() + if groups is None: + priors[zero_mask] = positive_total * zero_weight_total_share / zero_mask.sum() + else: + zero_groups = pd.factorize(np.asarray(groups)[zero_mask])[0] + size = np.bincount(zero_groups) + priors[zero_mask] = ( + positive_total * zero_weight_total_share / len(size) / size[zero_groups] + ) return priors +def tie_rows(values, groups: np.ndarray): + """Replace each row (a household) by the mean of its group's rows.""" + codes = pd.factorize(np.asarray(groups))[0] + array = np.asarray(values, dtype=float) + sums = np.zeros((codes.max() + 1,) + array.shape[1:]) + np.add.at(sums, codes, array) + tied = sums[codes] / np.bincount(codes).reshape((-1,) + (1,) * (array.ndim - 1))[codes] + if isinstance(values, pd.DataFrame): + return pd.DataFrame(tied, index=values.index, columns=values.columns) + return tied + + +def tie_weights(weights: np.ndarray, groups: np.ndarray) -> np.ndarray: + """Average area × household weights over each group's households.""" + return tie_rows(np.asarray(weights).T, groups).T + + def _as_bool_mask(series: pd.Series) -> np.ndarray: if pd.api.types.is_bool_dtype(series): return series.fillna(False).to_numpy(dtype=bool) @@ -205,6 +234,11 @@ def _household_weight_diagnostics( return diagnostics +def _saved_weights(weights, r, groups): + final_weights = (torch.exp(weights) * r).detach().numpy() + return final_weights if groups is None else tie_weights(final_weights, groups) + + def calibrate_local_areas( dataset: UKSingleYearDataset, matrix_fn, @@ -221,6 +255,7 @@ def calibrate_local_areas( nested_progress=None, time_period: int | str | None = None, zero_weight_prior_total_share: float = DEFAULT_ZERO_WEIGHT_PRIOR_TOTAL_SHARE, + groups: np.ndarray | None = None, ): """ Generic calibration function for local areas (constituencies, local authorities, etc.) @@ -239,6 +274,11 @@ def calibrate_local_areas( area_name: Name of the area type for logging zero_weight_prior_total_share: Share of prior household mass to reserve for rows whose incoming household_weight is zero. + groups: Optional label per household marking records that split one + household (a private renter's BRMA records). The records of a group + are calibrated as one: each takes the group's mean target rows, they + share one prior, and their saved weights are averaged, so they keep + equal weights. """ if dataset_key is None: dataset_key = default_weight_dataset_key() @@ -257,12 +297,16 @@ def track_stage(stage_name: str): with track_stage(f"{area_name}: build local target matrix"): matrix, y, r = _call_matrix_fn(matrix_fn, dataset, time_period) + if groups is not None: + matrix = tie_rows(matrix, groups) m_c, y_c = matrix.copy(), y.copy() with track_stage(f"{area_name}: build national target matrix"): m_national, y_national = _call_matrix_fn( national_matrix_fn, dataset, time_period ) + if groups is not None: + m_national = tie_rows(m_national, groups) m_n, y_n = m_national.copy(), y_national.copy() with track_stage(f"{area_name}: prepare tensors and optimizer"): @@ -280,6 +324,7 @@ def track_stage(stage_name: str): household_prior_weights = initialize_weight_priors( dataset.household.household_weight.values, zero_weight_total_share=zero_weight_prior_total_share, + groups=groups, ) area_prior_weights = household_prior_weights / areas_per_household original_weights = np.log(np.clip(area_prior_weights, 1e-12, None)) @@ -395,7 +440,7 @@ def dropout_weights(weights, p): return masked_weights optimizer = torch.optim.Adam([weights], lr=1e-1) - final_weights = (torch.exp(weights) * r).detach().numpy() + final_weights = _saved_weights(weights, r, groups) performance = pd.DataFrame() def log_performance( @@ -477,7 +522,7 @@ def log_performance( ) if epoch % 10 == 0: - final_weights = (torch.exp(weights) * r).detach().numpy() + final_weights = _saved_weights(weights, r, groups) log_performance( epoch=epoch, @@ -517,7 +562,7 @@ def log_performance( ) if epoch % 10 == 0: - final_weights = (torch.exp(weights) * r).detach().numpy() + final_weights = _saved_weights(weights, r, groups) log_performance( epoch=epoch, From 7ea4a2efd17a2078f9022b5c7806a7001a4be661 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 12:59:40 -0400 Subject: [PATCH 2/6] Calibrate without groups when no household is split PE_UK_DATA_BRMA_SPLIT=1 now leaves calibration exactly as before the split existed (singleton groups reordered the floating-point prior arithmetic). Co-Authored-By: Claude Opus 5.5 --- .../datasets/create_datasets.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/policyengine_uk_data/datasets/create_datasets.py b/policyengine_uk_data/datasets/create_datasets.py index f3a12bb4..0af72911 100644 --- a/policyengine_uk_data/datasets/create_datasets.py +++ b/policyengine_uk_data/datasets/create_datasets.py @@ -284,14 +284,16 @@ def main(): ) # A split household's records (within one OA clone) calibrate as one. - brma_groups = pd.factorize( - pd.MultiIndex.from_arrays( - [ - frs.household[BRMA_SPLIT_GROUP_COLUMN].to_numpy(), - frs.household["clone_index"].to_numpy(), - ] - ) - )[0] + brma_groups = None + if brma_split > 1: + brma_groups = pd.factorize( + pd.MultiIndex.from_arrays( + [ + frs.household[BRMA_SPLIT_GROUP_COLUMN].to_numpy(), + frs.household["clone_index"].to_numpy(), + ] + ) + )[0] # Run calibration with verbose progress frs_calibrated_constituencies = calibrate_local_areas( From 778e3d65aca05542a5ae91bc3969b2bf223805f5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 13:32:16 -0400 Subject: [PATCH 3/6] Calibrate each split household as one household MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Averaging tied records' rows and saved weights still optimised one log-weight per record, and the calibration's dropout (a dropped log-weight becomes the mean of all log-weights) lifts records that start at 1/k of their household's weight. That pushed private renters' weight down by about 5% and moved UK Universal Credit by about +£0.6bn against one draw. RecordGroups now collapses each group to one row (its records' mean target rows, its total weight as the prior), the optimiser sees one weight per household, and each record gets an equal share of its group's weight. A test checks the grouped calibration equals calibrating the unsplit households on those mean rows, with the same seed. Co-Authored-By: Claude Opus 5.5 --- policyengine_uk_data/tests/test_brma_split.py | 152 +++++++++++------- policyengine_uk_data/utils/calibrate.py | 114 +++++++------ 2 files changed, 161 insertions(+), 105 deletions(-) diff --git a/policyengine_uk_data/tests/test_brma_split.py b/policyengine_uk_data/tests/test_brma_split.py index 29d4b66a..862a2ee5 100644 --- a/policyengine_uk_data/tests/test_brma_split.py +++ b/policyengine_uk_data/tests/test_brma_split.py @@ -13,8 +13,9 @@ and benefit units; other households are untouched; every id stays unique; ``benunit_id // 100 == household_id`` holds wherever it held. 5. Determinism: the same seed gives the same records. -6. Tying: a group's calibrated weights are equal, its target rows identical, - and its zero-weight rows share one household's prior. +6. Tying: calibrating a split household's records with ``groups`` is + calibrating one household on their mean target rows: the same weights as + an unsplit dataset holding those mean rows, split equally across the records. """ import importlib.util @@ -33,11 +34,7 @@ load_brma_weights, spaced_quantiles, ) -from policyengine_uk_data.utils.calibrate import ( - initialize_weight_priors, - tie_rows, - tie_weights, -) +from policyengine_uk_data.utils.calibrate import RecordGroups @st.composite @@ -117,7 +114,9 @@ def test_spaced_records_track_the_expectation(case, k, seed): # One quantile in each interval [i/k, (i+1)/k). assert (np.sort(np.floor(quantiles * k), axis=1) == np.arange(k)).all() columns = brmas_at_quantiles( - np.repeat(probabilities, k, axis=0), np.repeat(keys, k, axis=0), quantiles.ravel() + np.repeat(probabilities, k, axis=0), + np.repeat(keys, k, axis=0), + quantiles.ravel(), ).reshape(-1, k) # Outcome = the key itself (any outcome monotone in the walk order). outcome = np.where(probabilities > 0, keys, 0) @@ -142,7 +141,9 @@ def test_household_distribution_is_the_mean_of_its_benefit_units(weights): region = np.array(["SCOTLAND", "SCOTLAND", "NORTHERN_IRELAND"]) benunit_household = np.array([0, 1, 1, 2]) category = np.array(["C", "A", "E", "B"]) - brmas, p = household_brma_probabilities(region, benunit_household, category, weights) + brmas, p = household_brma_probabilities( + region, benunit_household, category, weights + ) assert np.allclose(p.sum(axis=1), 1) def cell(r, c): @@ -243,7 +244,9 @@ def test_split_conserves_households_and_keeps_ids_unique(k, weights): if k > 1: assert all( (r, b) in supported - for r, b in zip(renters_after.region.astype(str), renters_after.brma.astype(str)) + for r, b in zip( + renters_after.region.astype(str), renters_after.brma.astype(str) + ) ) for table, column in ( (result.household, "household_id"), @@ -254,12 +257,19 @@ def test_split_conserves_households_and_keeps_ids_unique(k, weights): assert set(result.person.person_household_id) == set(after.household_id) assert set(result.person.person_benunit_id) == set(result.benunit.benunit_id) # The id scheme benunit_id // 100 == household_id survives the copies. - owner = result.person.drop_duplicates("person_benunit_id").set_index( - "person_benunit_id" - ).person_household_id - assert (result.benunit.benunit_id // 100 == owner.loc[result.benunit.benunit_id].to_numpy()).all() + owner = ( + result.person.drop_duplicates("person_benunit_id") + .set_index("person_benunit_id") + .person_household_id + ) + assert ( + result.benunit.benunit_id // 100 + == owner.loc[result.benunit.benunit_id].to_numpy() + ).all() assert len(result.person) == len(dataset.person) + (k - 1) * len( - dataset.person[dataset.person.person_household_id.isin(before.household_id[renter])] + dataset.person[ + dataset.person.person_household_id.isin(before.household_id[renter]) + ] ) @@ -282,62 +292,88 @@ def test_split_is_deterministic(weights): st.integers(1, 3), st.integers(0, 2**32 - 1), ) -def test_tied_rows_and_weights_are_group_means_and_keep_totals(labels, areas, seed): +def test_groups_collapse_and_expand_without_changing_predictions(labels, areas, seed): rng = np.random.default_rng(seed) - groups = np.array(labels) - rows = rng.random((len(groups), 3)) - tied = tie_rows(pd.DataFrame(rows), groups) - for g in set(labels): - m = groups == g - assert np.allclose(tied.to_numpy()[m], rows[m].mean(axis=0)) - w = rng.random((areas, len(groups))) - tw = tie_weights(w, groups) - for g in set(labels): - m = groups == g - assert np.allclose(tw[:, m], tw[:, m][:, :1]) - assert np.allclose(tw[:, m].sum(axis=1), w[:, m].sum(axis=1)) - # Predictions w @ rows are unchanged once rows are tied. - assert np.allclose(tw @ tied.to_numpy(), w @ tied.to_numpy()) - - -def test_split_zero_weight_rows_share_one_households_prior(): - weights = np.array([1_500.0, 0.0, 0.0, 0.0, 625.0, 0.0]) - groups = np.array([1, 2, 2, 2, 3, 4]) # rows 1-3 split one SPI household - priors = initialize_weight_priors(weights, groups=groups) - ungrouped = initialize_weight_priors(np.array([1_500.0, 0.0, 625.0, 0.0])) - assert priors[[1, 2, 3]].sum() == pytest.approx(ungrouped[1]) - assert priors[5] == pytest.approx(ungrouped[3]) - assert np.allclose(priors[[0, 4]], ungrouped[[0, 2]]) - assert priors.sum() == pytest.approx(weights.sum()) + groups = RecordGroups(np.array(labels), len(labels)) + rows = rng.random((len(labels), 3)) + means = groups.means(pd.DataFrame(rows)).to_numpy() + for g, label in enumerate(pd.unique(np.array(labels))): + assert np.allclose(means[g], rows[np.array(labels) == label].mean(axis=0)) + w = rng.random((areas, len(groups.sizes))) + records = groups.expand(w) + # Equal shares within a group, conserving the group's weight. + for g in range(len(groups.sizes)): + m = groups.codes == g + assert np.allclose(records[:, m], w[:, [g]] / m.sum()) + assert np.allclose(records.sum(axis=1), w.sum(axis=1)) + # Record weights on record rows predict what group weights do on mean rows. + assert np.allclose(records @ rows, w @ means) + totals = groups.totals(np.arange(len(labels), dtype=float)) + assert np.allclose(totals.sum(), np.arange(len(labels)).sum()) + + +def test_identity_groups_do_nothing(): + groups = RecordGroups(None, 4) + w = np.arange(8.0).reshape(2, 4) + first = w[0] + assert groups.expand(w) is w + assert groups.totals(first) is first @pytest.mark.skipif(importlib.util.find_spec("torch") is None, reason="needs torch") -def test_calibration_keeps_a_groups_records_equal(tmp_path, monkeypatch): +def test_grouped_records_calibrate_as_one_household(tmp_path, monkeypatch): + import torch + from policyengine_uk_data.utils import calibrate as calibrate_module from policyengine_uk_data.utils.calibrate import calibrate_local_areas monkeypatch.setattr(calibrate_module, "STORAGE_FOLDER", tmp_path) - # Records 0-2 split one household; their rows differ before tying. - matrix = pd.DataFrame([[1.0, 0.0], [3.0, 0.0], [2.0, 0.0], [0.0, 1.0], [0.0, 2.0]]) - local = pd.DataFrame([[600.0, 300.0], [500.0, 400.0]]) - mask = np.ones((2, 5)) class Data: def __init__(self, w): - self.household = pd.DataFrame({"household_weight": w}) + self.household = pd.DataFrame({"household_weight": np.asarray(w, float)}) def copy(self): return Data(self.household.household_weight.to_numpy().copy()) - result = calibrate_local_areas( - dataset=Data(np.array([100.0, 100.0, 100.0, 300.0, 300.0])), - matrix_fn=lambda d: (matrix.copy(), local.copy(), mask.copy()), - national_matrix_fn=lambda d: (pd.DataFrame(np.ones((5, 1))), pd.Series([1_500.0])), - area_count=2, - weight_file="tied.h5", - dataset_key="2024", - epochs=21, - groups=np.array([0, 0, 0, 1, 2]), + local = pd.DataFrame([[600.0, 300.0], [500.0, 400.0]]) + + def run(matrix, weights, groups): + torch.manual_seed(0) + n = len(weights) + return calibrate_local_areas( + dataset=Data(weights), + matrix_fn=lambda d: (pd.DataFrame(matrix), local.copy(), np.ones((2, n))), + national_matrix_fn=lambda d: ( + pd.DataFrame(np.ones((n, 1))), + pd.Series([1_500.0]), + ), + area_count=2, + weight_file="w.h5", + dataset_key="2024", + epochs=31, + groups=groups, + ).household.household_weight.to_numpy() + + # Household 0 split into three records (0-2) with different rows; a + # zero-weight household split into two (5-6). + split_rows = np.array( + [ + [1.0, 0.0], + [3.0, 0.0], + [2.0, 0.0], + [0.0, 1.0], + [0.0, 2.0], + [1.0, 1.0], + [0.0, 3.0], + ] ) - w = result.household.household_weight.to_numpy() - assert w[0] == pytest.approx(w[1]) and w[1] == pytest.approx(w[2]) + split_weights = np.array([100.0, 100.0, 100.0, 300.0, 300.0, 0.0, 0.0]) + groups = np.array([0, 0, 0, 1, 2, 3, 3]) + grouped = run(split_rows, split_weights, groups) + # The same households unsplit, each holding its records' mean row. + one_rows = np.array([[2.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.5, 2.0]]) + unsplit = run(one_rows, np.array([300.0, 300.0, 300.0, 0.0]), None) + assert np.allclose(grouped[:3], unsplit[0] / 3) + assert np.allclose(grouped[3:5], unsplit[1:3]) + assert np.allclose(grouped[5:], unsplit[3] / 2) diff --git a/policyengine_uk_data/utils/calibrate.py b/policyengine_uk_data/utils/calibrate.py index c6d2bdfb..8a104bad 100644 --- a/policyengine_uk_data/utils/calibrate.py +++ b/policyengine_uk_data/utils/calibrate.py @@ -110,7 +110,6 @@ def load_weights( def initialize_weight_priors( original_weights: np.ndarray, zero_weight_total_share: float = DEFAULT_ZERO_WEIGHT_PRIOR_TOTAL_SHARE, - groups: np.ndarray | None = None, ) -> np.ndarray: """Build deterministic positive household priors for calibration. @@ -120,10 +119,6 @@ def initialize_weight_priors( zero-weight rows are present, preserve the relative distribution of positive survey weights while reserving a fixed share of total prior mass for the zero-weight rows. - - ``groups`` (one label per row) marks rows that split one household, such - as a private renter's BRMA records: zero-weight rows of a group then share - one household's prior between them. """ weights = np.asarray(original_weights, dtype=np.float64) if weights.ndim != 1: @@ -146,32 +141,49 @@ def initialize_weight_priors( priors = np.empty_like(weights, dtype=np.float64) priors[positive_mask] = weights[positive_mask] * (1 - zero_weight_total_share) - if groups is None: - priors[zero_mask] = positive_total * zero_weight_total_share / zero_mask.sum() - else: - zero_groups = pd.factorize(np.asarray(groups)[zero_mask])[0] - size = np.bincount(zero_groups) - priors[zero_mask] = ( - positive_total * zero_weight_total_share / len(size) / size[zero_groups] - ) + priors[zero_mask] = positive_total * zero_weight_total_share / zero_mask.sum() return priors -def tie_rows(values, groups: np.ndarray): - """Replace each row (a household) by the mean of its group's rows.""" - codes = pd.factorize(np.asarray(groups))[0] - array = np.asarray(values, dtype=float) - sums = np.zeros((codes.max() + 1,) + array.shape[1:]) - np.add.at(sums, codes, array) - tied = sums[codes] / np.bincount(codes).reshape((-1,) + (1,) * (array.ndim - 1))[codes] - if isinstance(values, pd.DataFrame): - return pd.DataFrame(tied, index=values.index, columns=values.columns) - return tied +class RecordGroups: + """Records that split one household, calibrated as that household. + ``labels`` gives each record's group. The calibration optimises one + weight per group, on the mean of its records' target rows (the household's + expected contribution per unit of weight) with the group's total weight as + its prior; each record then gets an equal share of its group's weight. + """ -def tie_weights(weights: np.ndarray, groups: np.ndarray) -> np.ndarray: - """Average area × household weights over each group's households.""" - return tie_rows(np.asarray(weights).T, groups).T + def __init__(self, labels: np.ndarray | None, records: int): + """``labels=None``: every record is its own group (no work is done).""" + self.identity = labels is None + if self.identity: + labels = np.arange(records) + self.codes = pd.factorize(np.asarray(labels))[0] + self.sizes = np.bincount(self.codes) + self.first = np.zeros(len(self.sizes), dtype=int) + self.first[self.codes[::-1]] = np.arange(len(self.codes))[::-1] + + def means(self, values): + """Per-group means of record rows (a DataFrame keeps its columns).""" + array = np.asarray(values, dtype=float) + sums = np.zeros((len(self.sizes),) + array.shape[1:]) + np.add.at(sums, self.codes, array) + means = sums / self.sizes.reshape((-1,) + (1,) * (array.ndim - 1)) + if isinstance(values, pd.DataFrame): + return pd.DataFrame(means, columns=values.columns) + return means + + def totals(self, values: np.ndarray) -> np.ndarray: + if self.identity: + return values + return np.bincount(self.codes, weights=np.asarray(values, dtype=float)) + + def expand(self, weights: np.ndarray) -> np.ndarray: + """Group weights (last axis) to records, split equally.""" + if self.identity: + return weights + return np.asarray(weights)[..., self.codes] / self.sizes[self.codes] def _as_bool_mask(series: pd.Series) -> np.ndarray: @@ -234,11 +246,6 @@ def _household_weight_diagnostics( return diagnostics -def _saved_weights(weights, r, groups): - final_weights = (torch.exp(weights) * r).detach().numpy() - return final_weights if groups is None else tie_weights(final_weights, groups) - - def calibrate_local_areas( dataset: UKSingleYearDataset, matrix_fn, @@ -275,9 +282,8 @@ def calibrate_local_areas( zero_weight_prior_total_share: Share of prior household mass to reserve for rows whose incoming household_weight is zero. groups: Optional label per household marking records that split one - household (a private renter's BRMA records). The records of a group - are calibrated as one: each takes the group's mean target rows, they - share one prior, and their saved weights are averaged, so they keep + household (a private renter's BRMA records). Each group is + calibrated as one household (``RecordGroups``), so its records keep equal weights. """ if dataset_key is None: @@ -297,8 +303,15 @@ def track_stage(stage_name: str): with track_stage(f"{area_name}: build local target matrix"): matrix, y, r = _call_matrix_fn(matrix_fn, dataset, time_period) + record_groups = RecordGroups( + groups, len(dataset.household.household_weight.values) + ) if groups is not None: - matrix = tie_rows(matrix, groups) + r = np.asarray(r) + if not (r == r[:, record_groups.first[record_groups.codes]]).all(): + raise ValueError("A record group spans areas of different countries.") + matrix = record_groups.means(matrix) + r = r[:, record_groups.first] m_c, y_c = matrix.copy(), y.copy() with track_stage(f"{area_name}: build national target matrix"): @@ -306,7 +319,7 @@ def track_stage(stage_name: str): national_matrix_fn, dataset, time_period ) if groups is not None: - m_national = tie_rows(m_national, groups) + m_national = record_groups.means(m_national) m_n, y_n = m_national.copy(), y_national.copy() with track_stage(f"{area_name}: prepare tensors and optimizer"): @@ -322,9 +335,8 @@ def track_stage(stage_name: str): areas_per_household, 1 ) # avoid division by zero household_prior_weights = initialize_weight_priors( - dataset.household.household_weight.values, + record_groups.totals(dataset.household.household_weight.values), zero_weight_total_share=zero_weight_prior_total_share, - groups=groups, ) area_prior_weights = household_prior_weights / areas_per_household original_weights = np.log(np.clip(area_prior_weights, 1e-12, None)) @@ -440,7 +452,7 @@ def dropout_weights(weights, p): return masked_weights optimizer = torch.optim.Adam([weights], lr=1e-1) - final_weights = _saved_weights(weights, r, groups) + final_weights = (torch.exp(weights) * r).detach().numpy() performance = pd.DataFrame() def log_performance( @@ -480,8 +492,8 @@ def log_performance( diagnostics = _household_weight_diagnostics( dataset, - final_weights.sum(axis=0), - household_prior_weights, + record_groups.expand(final_weights).sum(axis=0), + record_groups.expand(household_prior_weights), ) for key, value in diagnostics.items(): performance_step[key] = value @@ -522,7 +534,7 @@ def log_performance( ) if epoch % 10 == 0: - final_weights = _saved_weights(weights, r, groups) + final_weights = (torch.exp(weights) * r).detach().numpy() log_performance( epoch=epoch, @@ -535,9 +547,13 @@ def log_performance( # Save weights with h5py.File(STORAGE_FOLDER / weight_file, "w") as f: - f.create_dataset(dataset_key, data=final_weights) + f.create_dataset( + dataset_key, data=record_groups.expand(final_weights) + ) - dataset.household.household_weight = final_weights.sum(axis=0) + dataset.household.household_weight = record_groups.expand( + final_weights + ).sum(axis=0) else: for epoch in range(epochs): optimizer.zero_grad() @@ -562,7 +578,7 @@ def log_performance( ) if epoch % 10 == 0: - final_weights = _saved_weights(weights, r, groups) + final_weights = (torch.exp(weights) * r).detach().numpy() log_performance( epoch=epoch, @@ -575,8 +591,12 @@ def log_performance( # Save weights with h5py.File(STORAGE_FOLDER / weight_file, "w") as f: - f.create_dataset(dataset_key, data=final_weights) + f.create_dataset( + dataset_key, data=record_groups.expand(final_weights) + ) - dataset.household.household_weight = final_weights.sum(axis=0) + dataset.household.household_weight = record_groups.expand( + final_weights + ).sum(axis=0) return dataset From 3a097486d99e8b3cb3e85df03a21eca05c21d57b Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 18:23:34 -0400 Subject: [PATCH 4/6] Address review: source identity, float32 weights, sharper invariants - Split records keep the household they came from as source_household_id, and the OA clone step keeps it, so geography support counts one source household per split household instead of k. - Expanding group weights keeps the optimiser's float32 dtype, household totals are summed before expansion, and the cross-country check uses one fingerprint per record rather than a second full mask. - A benefit unit without people, or a supplied cell with no positive weight, now raises; the split skips rate keys when nobody rents privately. - Tests: each split record's marginal over a uniform phase; quantile boundaries with zero, tied and infinite-key columns; rate keys against fresh simulations per BRMA; BRMA-dependent national rows and the saved h5 in the grouped calibration; mixed-country groups; source identity. The four mutants that survived review (< for <=, a fixed phase, first national row, a stale rate cache) now fail. Co-Authored-By: Claude Opus 5.5 --- .../calibration/clone_and_assign.py | 8 +- policyengine_uk_data/datasets/brma.py | 64 +++-- policyengine_uk_data/tests/test_brma_split.py | 223 +++++++++++++++++- policyengine_uk_data/utils/calibrate.py | 20 +- 4 files changed, 280 insertions(+), 35 deletions(-) diff --git a/policyengine_uk_data/calibration/clone_and_assign.py b/policyengine_uk_data/calibration/clone_and_assign.py index f1907b9d..f5fb132c 100644 --- a/policyengine_uk_data/calibration/clone_and_assign.py +++ b/policyengine_uk_data/calibration/clone_and_assign.py @@ -189,7 +189,13 @@ def clone_and_assign( # Clone household table hh_clone = hh.copy() - hh_clone["source_household_id"] = hh_id_col + # A record split from another household (datasets/brma.py) keeps that + # household as its source, so survey support is not overcounted. + hh_clone["source_household_id"] = ( + hh["source_household_id"].to_numpy() + if "source_household_id" in hh.columns + else hh_id_col + ) if "source_year" not in hh_clone.columns: hh_clone["source_year"] = dataset.time_period hh_clone["household_id"] = new_hh_ids diff --git a/policyengine_uk_data/datasets/brma.py b/policyengine_uk_data/datasets/brma.py index 4764bf3a..9746a49e 100644 --- a/policyengine_uk_data/datasets/brma.py +++ b/policyengine_uk_data/datasets/brma.py @@ -164,6 +164,11 @@ def household_brma_probabilities( if not mask.any(): continue p = cell.set_index("brma").weight.reindex(brmas[cell_region]).fillna(0) + if not (p >= 0).all() or p.sum() <= 0: + raise ValueError( + f"No BRMA weights for region × LHA category cells " + f"{[(cell_region, category)]}." + ) p = np.pad(p.to_numpy(float) / p.sum(), (0, width - len(p))) np.add.at(probabilities, benunit_household[mask], p) covered |= mask @@ -262,10 +267,14 @@ def split_private_renters_across_brmas( BRMA. Every household gets ``BRMA_SPLIT_GROUP_COLUMN``: the id of the household its record came from. - Copies offset household ids by ``i * step``, benefit unit ids by - ``i * step * 100`` and person ids by ``i * step * 1000`` (``step`` a power - of ten above every household id, raised if needed so no id collides), so - ``benunit_id // 100 == household_id`` still holds wherever it held. + Copies offset household ids by ``i * step`` (``step`` a power of ten above + every household id), benefit unit ids by ``i * step * 100`` and person ids + by ``i * step * 1000``, each raised to a power of ten above that table's + largest id if needed so no id collides. When no raise is needed, + ``benunit_id // 100 == household_id`` still holds wherever it held. Every + benefit unit must have at least one person. Each record also gets + ``source_household_id``, the household it came from, which the OA clone + step keeps, so split records count as one source household. Returns: The new dataset. With ``k == 1`` it is the input plus the group column. @@ -275,6 +284,7 @@ def split_private_renters_across_brmas( household = dataset.household.copy() household[BRMA_SPLIT_GROUP_COLUMN] = household.household_id.to_numpy() + household["source_household_id"] = household.household_id.to_numpy() if k == 1: return UKSingleYearDataset( person=dataset.person, @@ -282,17 +292,39 @@ def split_private_renters_across_brmas( household=household, fiscal_year=dataset.time_period, ) + person, benunit = dataset.person, dataset.benunit + owner = ( + person.drop_duplicates("person_benunit_id") + .set_index("person_benunit_id") + .person_household_id + ) + empty = ~benunit.benunit_id.isin(owner.index) + if empty.any(): + raise ValueError( + f"{int(empty.sum())} benefit units have no people; every benefit " + "unit needs one to place it in a household." + ) + benunit_owner = owner.loc[benunit.benunit_id].to_numpy() + renter = household.tenure_type.astype(str).to_numpy() == "RENT_PRIVATELY" + if not renter.any(): + return UKSingleYearDataset( + person=person, + benunit=benunit, + household=household, + fiscal_year=dataset.time_period, + ) year = int(dataset.time_period) simulation = Microsimulation(dataset=dataset) region = np.asarray(simulation.calculate("region", year)).astype(str) - renter = ( - np.asarray(simulation.calculate("tenure_type", year)).astype(str) - == "RENT_PRIVATELY" - ) - benunit_household = pd.Series( + position = pd.Series( np.arange(len(household)), index=household.household_id.to_numpy() - ).loc[np.asarray(simulation.populations["benunit"].household("household_id", year))] - benunit_household = benunit_household.to_numpy() + ) + benunit_household = position.loc[benunit_owner].to_numpy() + sim_order = np.asarray(simulation.calculate("benunit_id", year)) + if not np.array_equal(sim_order, benunit.benunit_id.to_numpy()): + raise ValueError( + "The simulation's benefit units are not in the dataset's order." + ) category = np.asarray(simulation.calculate("LHA_category", year)).astype(str) brmas, probabilities = household_brma_probabilities( region, benunit_household, category, weights @@ -308,16 +340,8 @@ def split_private_renters_across_brmas( dtype=object, ).reshape(-1, k) - person, benunit = dataset.person, dataset.benunit renter_ids = household.household_id.to_numpy()[renter] - benunit_of = ( - person.drop_duplicates("person_benunit_id") - .set_index("person_benunit_id") - .person_household_id - ) - renter_benunits = benunit[ - np.isin(benunit_of.loc[benunit.benunit_id].to_numpy(), renter_ids) - ] + renter_benunits = benunit[np.isin(benunit_owner, renter_ids)] renter_people = person[person.person_household_id.isin(renter_ids)] step = 10 ** len(str(int(household.household_id.max()))) benunit_step = max(step * 100, 10 ** len(str(int(benunit.benunit_id.max())))) diff --git a/policyengine_uk_data/tests/test_brma_split.py b/policyengine_uk_data/tests/test_brma_split.py index 862a2ee5..135458d5 100644 --- a/policyengine_uk_data/tests/test_brma_split.py +++ b/policyengine_uk_data/tests/test_brma_split.py @@ -127,6 +127,63 @@ def test_spaced_records_track_the_expectation(case, k, seed): assert (np.abs(mean - expected) <= spread / k + 1e-9).all() +class GridGenerator: + """Stands in for a Generator: ``random(n)`` returns the midpoints of n equal cells.""" + + def random(self, n): + return (np.arange(n) + 0.5) / n + + +@settings(max_examples=60, deadline=None, derandomize=True) +@given(distributions(max_rows=3), st.integers(1, 6)) +def test_every_split_record_has_its_households_distribution(case, k): + # Over a uniform phase, record i's quantile (u + i/k) mod 1 is uniform, so + # its BRMA has exactly the household's distribution. + probabilities, keys = case + n = 6_000 + for row in range(len(probabilities)): + quantiles = spaced_quantiles(n, k, GridGenerator()) + for i in range(k): + columns = brmas_at_quantiles( + np.repeat(probabilities[row : row + 1], n, axis=0), + np.repeat(keys[row : row + 1], n, axis=0), + quantiles[:, i], + ) + frequency = np.bincount(columns, minlength=probabilities.shape[1]) / n + assert np.abs(frequency - probabilities[row]).max() <= 2 / n + 1e-12 + + +@pytest.mark.parametrize( + "probabilities, keys, quantile, expected", + [ + ([0, 0.5, 0, 0.5], None, 0.0, 1), + ([0, 0.5, 0, 0.5], None, np.nextafter(0.5, 0), 1), + ([0, 0.5, 0, 0.5], None, 0.5, 3), + ([0, 0.5, 0, 0.5], None, np.nextafter(1.0, 0), 3), + ([0.2, 0.3, 0.5], [1.0, 1.0, 1.0], 0.0, 0), # tied keys keep column order + ([0.2, 0.3, 0.5], [1.0, 1.0, 1.0], 0.2, 1), + ([0, 0.4, 0, 0.6], [np.inf, 2.0, np.inf, 1.0], 0.0, 3), # lowest key first + ([0, 0.4, 0, 0.6], [np.inf, 2.0, np.inf, 1.0], 0.6, 1), + ([0, 0.4, 0, 0.6], [np.inf, 2.0, np.inf, 1.0], np.nextafter(1.0, 0), 1), + ( + [0.5, 0.5, 0, 0], + [1.0, 2.0, np.inf, np.inf], + np.nextafter(1.0, 0), + 1, + ), # padding + ], +) +def test_quantile_boundaries(probabilities, keys, quantile, expected): + p = np.array([probabilities], dtype=float) + k = None if keys is None else np.array([keys], dtype=float) + assert brmas_at_quantiles(p, k, np.array([quantile]))[0] == expected + + +def test_quantiles_outside_the_unit_interval_are_rejected(): + with pytest.raises(ValueError): + brmas_at_quantiles(np.array([[1.0]]), None, np.array([1.0])) + + def test_spaced_quantiles_reject_k_below_one(): with pytest.raises(ValueError): spaced_quantiles(3, 0, np.random.default_rng(0)) @@ -170,6 +227,16 @@ def test_household_distribution_fails_closed(weights): ) +def test_household_distribution_rejects_zero_mass_cells(weights): + table = weights.copy() + cell = (table.region == "WALES") & (table.lha_category == "C") + table.loc[cell, "weight"] = 0.0 + with pytest.raises(ValueError, match="No BRMA weights"): + household_brma_probabilities( + np.array(["WALES"]), np.array([0]), np.array(["C"]), table + ) + + # --- splitting a dataset (runs PolicyEngine UK on a few synthetic households) --- needs_policyengine = pytest.mark.skipif( @@ -322,6 +389,7 @@ def test_identity_groups_do_nothing(): @pytest.mark.skipif(importlib.util.find_spec("torch") is None, reason="needs torch") def test_grouped_records_calibrate_as_one_household(tmp_path, monkeypatch): + import h5py import torch from policyengine_uk_data.utils import calibrate as calibrate_module @@ -338,15 +406,15 @@ def copy(self): local = pd.DataFrame([[600.0, 300.0], [500.0, 400.0]]) - def run(matrix, weights, groups): + def run(matrix, national, weights, groups): torch.manual_seed(0) n = len(weights) - return calibrate_local_areas( + result = calibrate_local_areas( dataset=Data(weights), matrix_fn=lambda d: (pd.DataFrame(matrix), local.copy(), np.ones((2, n))), national_matrix_fn=lambda d: ( - pd.DataFrame(np.ones((n, 1))), - pd.Series([1_500.0]), + pd.DataFrame(national), + pd.Series([1_500.0, 900.0]), ), area_count=2, weight_file="w.h5", @@ -354,6 +422,12 @@ def run(matrix, weights, groups): epochs=31, groups=groups, ).household.household_weight.to_numpy() + with h5py.File(tmp_path / "w.h5") as f: + saved = f["2024"][...] + # The saved area weights are the records' weights, in float32 like the optimiser's. + assert saved.shape == (2, n) and saved.dtype == np.float32 + assert np.allclose(saved.sum(axis=0), result, rtol=1e-5) + return result # Household 0 split into three records (0-2) with different rows; a # zero-weight household split into two (5-6). @@ -368,12 +442,147 @@ def run(matrix, weights, groups): [0.0, 3.0], ] ) + # National rows differ within the split groups too (e.g. a BRMA-dependent benefit). + split_national = np.array( + [ + [1.0, 0.0], + [1.0, 4.0], + [1.0, 2.0], + [1.0, 1.0], + [1.0, 0.0], + [1.0, 3.0], + [1.0, 1.0], + ] + ) split_weights = np.array([100.0, 100.0, 100.0, 300.0, 300.0, 0.0, 0.0]) groups = np.array([0, 0, 0, 1, 2, 3, 3]) - grouped = run(split_rows, split_weights, groups) - # The same households unsplit, each holding its records' mean row. + grouped = run(split_rows, split_national, split_weights, groups) + # The same households unsplit, each holding its records' mean rows. one_rows = np.array([[2.0, 0.0], [0.0, 1.0], [0.0, 2.0], [0.5, 2.0]]) - unsplit = run(one_rows, np.array([300.0, 300.0, 300.0, 0.0]), None) + one_national = np.array([[1.0, 2.0], [1.0, 1.0], [1.0, 0.0], [1.0, 2.0]]) + unsplit = run(one_rows, one_national, np.array([300.0, 300.0, 300.0, 0.0]), None) assert np.allclose(grouped[:3], unsplit[0] / 3) assert np.allclose(grouped[3:5], unsplit[1:3]) assert np.allclose(grouped[5:], unsplit[3] / 2) + + +@pytest.mark.skipif(importlib.util.find_spec("torch") is None, reason="needs torch") +def test_a_group_spanning_countries_is_rejected(tmp_path, monkeypatch): + from policyengine_uk_data.utils import calibrate as calibrate_module + from policyengine_uk_data.utils.calibrate import calibrate_local_areas + + monkeypatch.setattr(calibrate_module, "STORAGE_FOLDER", tmp_path) + + class Data: + def __init__(self, w): + self.household = pd.DataFrame({"household_weight": np.asarray(w, float)}) + + def copy(self): + return Data(self.household.household_weight.to_numpy().copy()) + + mask = np.array( + [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]] + ) # records 0 and 1 in different countries + with pytest.raises(ValueError, match="different countries"): + calibrate_local_areas( + dataset=Data([1.0, 1.0, 1.0]), + matrix_fn=lambda d: ( + pd.DataFrame(np.eye(3, 2)), + pd.DataFrame(np.ones((2, 2))), + mask, + ), + national_matrix_fn=lambda d: ( + pd.DataFrame(np.ones((3, 1))), + pd.Series([3.0]), + ), + area_count=2, + weight_file="w.h5", + dataset_key="2024", + epochs=1, + groups=np.array([0, 0, 1]), + ) + + +@needs_policyengine +def test_split_records_keep_their_source_household(weights): + from policyengine_uk_data.datasets.brma import split_private_renters_across_brmas + + result = split_private_renters_across_brmas( + toy_dataset(), k=3, seed=1, weights=weights + ) + household = result.household + assert (household.source_household_id == household[BRMA_SPLIT_GROUP_COLUMN]).all() + + +@needs_policyengine +def test_a_benefit_unit_without_people_is_rejected(weights): + from policyengine_uk.data import UKSingleYearDataset + + from policyengine_uk_data.datasets.brma import split_private_renters_across_brmas + + dataset = toy_dataset() + benunit = pd.concat( + [dataset.benunit, pd.DataFrame({"benunit_id": [99_999]})], ignore_index=True + ) + broken = UKSingleYearDataset( + person=dataset.person, + benunit=benunit, + household=dataset.household, + fiscal_year=2024, + ) + with pytest.raises(ValueError, match="no people"): + split_private_renters_across_brmas(broken, k=2, seed=0, weights=weights) + + +@needs_policyengine +def test_rate_keys_match_fresh_simulations(weights): + # The keys come from one simulation with the BRMA reset and the rate's cache + # cleared per slot; each slot must equal a simulation built with that BRMA. + from policyengine_uk import Microsimulation + from policyengine_uk.data import UKSingleYearDataset + + from policyengine_uk_data.datasets.brma import lha_rate_keys + + dataset = toy_dataset() + year = int(dataset.time_period) + simulation = Microsimulation(dataset=dataset) + region = dataset.household.region.astype(str).to_numpy() + position = pd.Series(np.arange(len(region)), index=dataset.household.household_id) + owner = ( + dataset.person.drop_duplicates("person_benunit_id") + .set_index("person_benunit_id") + .person_household_id + ) + benunit_household = position.loc[owner.loc[dataset.benunit.benunit_id]].to_numpy() + brmas, _ = household_brma_probabilities( + region, + benunit_household, + np.asarray(simulation.calculate("LHA_category", year)).astype(str), + weights, + ) + keys = lha_rate_keys(simulation, year, brmas, region, benunit_household) + variable = "uncapped_BRMA_LHA_rate" + if variable not in simulation.tax_benefit_system.variables: + variable = "BRMA_LHA_rate" + distinct = set() + for j in (0, 3, 7): + household = dataset.household.copy() + household["brma"] = [brmas[r][min(j, len(brmas[r]) - 1)] for r in region] + fresh = Microsimulation( + dataset=UKSingleYearDataset( + person=dataset.person, + benunit=dataset.benunit, + household=household, + fiscal_year=year, + ) + ) + rate = np.asarray(fresh.calculate(variable, year)) + units = np.bincount(benunit_household, minlength=len(region)) + mean = ( + np.bincount(benunit_household, weights=rate, minlength=len(region)) / units + ) + has_slot = np.array([j < len(brmas[r]) for r in region]) + assert np.allclose(keys[has_slot, j], mean[has_slot]) + assert np.isinf(keys[~has_slot, j]).all() + distinct.add(tuple(np.round(mean[has_slot], 2))) + assert len(distinct) == 3 # the slots really differ, so a stale cache would show diff --git a/policyengine_uk_data/utils/calibrate.py b/policyengine_uk_data/utils/calibrate.py index 8a104bad..fbd8f78d 100644 --- a/policyengine_uk_data/utils/calibrate.py +++ b/policyengine_uk_data/utils/calibrate.py @@ -183,7 +183,9 @@ def expand(self, weights: np.ndarray) -> np.ndarray: """Group weights (last axis) to records, split equally.""" if self.identity: return weights - return np.asarray(weights)[..., self.codes] / self.sizes[self.codes] + weights = np.asarray(weights) + sizes = self.sizes[self.codes].astype(weights.dtype, copy=False) + return weights[..., self.codes] / sizes def _as_bool_mask(series: pd.Series) -> np.ndarray: @@ -308,7 +310,11 @@ def track_stage(stage_name: str): ) if groups is not None: r = np.asarray(r) - if not (r == r[:, record_groups.first[record_groups.codes]]).all(): + # One fingerprint per record's area mask; a group's must all agree. + fingerprint = r.T @ np.random.default_rng(0).random(r.shape[0]) + if not np.allclose( + fingerprint, fingerprint[record_groups.first[record_groups.codes]] + ): raise ValueError("A record group spans areas of different countries.") matrix = record_groups.means(matrix) r = r[:, record_groups.first] @@ -492,7 +498,7 @@ def log_performance( diagnostics = _household_weight_diagnostics( dataset, - record_groups.expand(final_weights).sum(axis=0), + record_groups.expand(final_weights.sum(axis=0)), record_groups.expand(household_prior_weights), ) for key, value in diagnostics.items(): @@ -552,8 +558,8 @@ def log_performance( ) dataset.household.household_weight = record_groups.expand( - final_weights - ).sum(axis=0) + final_weights.sum(axis=0) + ) else: for epoch in range(epochs): optimizer.zero_grad() @@ -596,7 +602,7 @@ def log_performance( ) dataset.household.household_weight = record_groups.expand( - final_weights - ).sum(axis=0) + final_weights.sum(axis=0) + ) return dataset From 03c066613d19d9999b0901357d2bc1ea7f2a6ef2 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 07:36:13 -0400 Subject: [PATCH 5/6] Check entity order and tenure through the simulation; compare area masks exactly - The split checks that the household and benefit unit tables are in PolicyEngine UK's population order (sorted by id), comparing the populations' ids rather than the benunit_id input, which a swapped row passed and int64 ids above 2**31 failed. - Private renters are read from the simulation's tenure_type, so a tenure column holding enum codes splits the same households as one holding names. - RecordGroups.columns_agree compares each grouped record's area mask with its group's first exactly, in chunks, replacing a seeded projection that near-equal one-hot masks could pass. Co-Authored-By: Claude Opus 5.5 --- policyengine_uk_data/datasets/brma.py | 26 ++-- policyengine_uk_data/tests/test_brma_split.py | 112 +++++++++++++++++- policyengine_uk_data/utils/calibrate.py | 19 ++- 3 files changed, 140 insertions(+), 17 deletions(-) diff --git a/policyengine_uk_data/datasets/brma.py b/policyengine_uk_data/datasets/brma.py index 9746a49e..010c66a4 100644 --- a/policyengine_uk_data/datasets/brma.py +++ b/policyengine_uk_data/datasets/brma.py @@ -272,7 +272,8 @@ def split_private_renters_across_brmas( by ``i * step * 1000``, each raised to a power of ten above that table's largest id if needed so no id collides. When no raise is needed, ``benunit_id // 100 == household_id`` still holds wherever it held. Every - benefit unit must have at least one person. Each record also gets + benefit unit must have at least one person, and the household and benefit + unit tables must be sorted by id (as PolicyEngine UK needs). Each record also gets ``source_household_id``, the household it came from, which the OA clone step keeps, so split records count as one source household. @@ -305,7 +306,21 @@ def split_private_renters_across_brmas( "unit needs one to place it in a household." ) benunit_owner = owner.loc[benunit.benunit_id].to_numpy() - renter = household.tenure_type.astype(str).to_numpy() == "RENT_PRIVATELY" + year = int(dataset.time_period) + simulation = Microsimulation(dataset=dataset) + for entity, ids in ( + ("household", household.household_id), + ("benunit", benunit.benunit_id), + ): + # PolicyEngine UK orders each entity by id but reads input columns in + # table order, so the two orders must agree. + if not np.array_equal(simulation.populations[entity].ids, ids.to_numpy()): + raise ValueError( + f"The {entity} table is not sorted by {entity}_id, so the " + "simulation would misalign its rows." + ) + tenure = np.asarray(simulation.calculate("tenure_type", year)).astype(str) + renter = tenure == "RENT_PRIVATELY" if not renter.any(): return UKSingleYearDataset( person=person, @@ -313,18 +328,11 @@ def split_private_renters_across_brmas( household=household, fiscal_year=dataset.time_period, ) - year = int(dataset.time_period) - simulation = Microsimulation(dataset=dataset) region = np.asarray(simulation.calculate("region", year)).astype(str) position = pd.Series( np.arange(len(household)), index=household.household_id.to_numpy() ) benunit_household = position.loc[benunit_owner].to_numpy() - sim_order = np.asarray(simulation.calculate("benunit_id", year)) - if not np.array_equal(sim_order, benunit.benunit_id.to_numpy()): - raise ValueError( - "The simulation's benefit units are not in the dataset's order." - ) category = np.asarray(simulation.calculate("LHA_category", year)).astype(str) brmas, probabilities = household_brma_probabilities( region, benunit_household, category, weights diff --git a/policyengine_uk_data/tests/test_brma_split.py b/policyengine_uk_data/tests/test_brma_split.py index 135458d5..8089db5d 100644 --- a/policyengine_uk_data/tests/test_brma_split.py +++ b/policyengine_uk_data/tests/test_brma_split.py @@ -245,7 +245,7 @@ def test_household_distribution_rejects_zero_mass_cells(weights): ) -def toy_dataset(): +def toy_dataset(id_offset: int = 0): from policyengine_uk.data import UKSingleYearDataset regions = ["LONDON", "LONDON", "SCOTLAND", "WALES", "NORTH_WEST", "LONDON"] * 2 @@ -253,7 +253,7 @@ def toy_dataset(): n = len(regions) household = pd.DataFrame( { - "household_id": np.arange(1, n + 1), + "household_id": np.arange(1, n + 1) + id_offset, "household_weight": np.r_[np.linspace(100, 2_000, n - 2), 0.0, 0.0], "region": regions, "tenure_type": tenure, @@ -265,7 +265,7 @@ def toy_dataset(): people, units = [], [] for h in household.household_id: # Households 3 and 6 hold two benefit units (a lodger family). - for b in range(1, 3 if h % 3 == 0 else 2): + for b in range(1, 3 if (h - id_offset) % 3 == 0 else 2): units.append({"benunit_id": h * 100 + b}) for p in range(1, 3): people.append( @@ -586,3 +586,109 @@ def test_rate_keys_match_fresh_simulations(weights): assert np.isinf(keys[~has_slot, j]).all() distinct.add(tuple(np.round(mean[has_slot], 2))) assert len(distinct) == 3 # the slots really differ, so a stale cache would show + + +def _with_tables(dataset, **tables): + from policyengine_uk.data import UKSingleYearDataset + + parts = { + "person": dataset.person, + "benunit": dataset.benunit, + "household": dataset.household, + } + parts.update(tables) + return UKSingleYearDataset(**parts, fiscal_year=dataset.time_period) + + +@needs_policyengine +@pytest.mark.parametrize("table", ["benunit", "household"]) +def test_tables_out_of_id_order_are_rejected(table, weights): + # PolicyEngine UK sorts entities by id but reads input columns in table + # order, so swapping two rows would misalign every input. + from policyengine_uk_data.datasets.brma import split_private_renters_across_brmas + + dataset = toy_dataset() + frame = getattr(dataset, table) + swapped = frame.iloc[[1, 0] + list(range(2, len(frame)))].reset_index(drop=True) + with pytest.raises(ValueError, match=f"{table} table is not sorted"): + split_private_renters_across_brmas( + _with_tables(dataset, **{table: swapped}), k=2, seed=0, weights=weights + ) + + +@needs_policyengine +def test_ids_beyond_int32_split_like_small_ones(weights): + # Benefit unit ids above 2**31 must not trip the order check. + from policyengine_uk_data.datasets.brma import split_private_renters_across_brmas + + offset = 30_000_000 + small = split_private_renters_across_brmas( + toy_dataset(), k=3, seed=5, weights=weights + ) + large = split_private_renters_across_brmas( + toy_dataset(offset), k=3, seed=5, weights=weights + ) + assert large.benunit.benunit_id.max() > 2**31 + assert len(large.household) == len(small.household) + assert ( + large.household.brma.astype(str).to_numpy() + == small.household.brma.astype(str).to_numpy() + ).all() + + +@needs_policyengine +def test_encoded_tenure_splits_the_same_renters(weights): + # A tenure column holding enum codes must select the same private renters + # as one holding names. + from policyengine_uk.variables.household.demographic.tenure_type import ( + TenureType, + ) + + from policyengine_uk_data.datasets.brma import split_private_renters_across_brmas + + dataset = toy_dataset() + household = dataset.household.copy() + members = list(TenureType) + household["tenure_type"] = [ + members.index(TenureType[name]) for name in household.tenure_type.astype(str) + ] + assert pd.api.types.is_integer_dtype(household.tenure_type) + named = split_private_renters_across_brmas(dataset, k=2, seed=4, weights=weights) + coded = split_private_renters_across_brmas( + _with_tables(dataset, household=household), k=2, seed=4, weights=weights + ) + assert len(named.household) > len(dataset.household) + assert len(coded.household) == len(named.household) + assert ( + coded.household.brma.astype(str).to_numpy() + == named.household.brma.astype(str).to_numpy() + ).all() + + +@settings(max_examples=200, deadline=None, derandomize=True) +@given( + st.lists(st.integers(0, 3), min_size=1, max_size=10), + st.integers(1, 4), + st.integers(1, 3), + st.integers(0, 2**32 - 1), +) +def test_column_agreement_is_exact(labels, rows, chunk, seed): + rng = np.random.default_rng(seed) + groups = RecordGroups(np.array(labels), len(labels)) + matrix = (rng.random((rows, len(labels))) < 0.5).astype(float) + if rng.random() < 0.5: # often make the groups agree + matrix = matrix[:, groups.first[groups.codes]] + first = groups.first[groups.codes] + expected = all( + np.array_equal(matrix[:, j], matrix[:, first[j]]) for j in range(len(labels)) + ) + assert groups.columns_agree(matrix, chunk=chunk) == expected + + +def test_one_hot_area_masks_that_differ_never_agree(): + # Two records of a group placed in different single areas of 650 (masks + # whose seeded random projections nearly coincide) must be told apart. + mask = np.zeros((650, 2)) + mask[506, 0] = mask[173, 1] = 1.0 + assert not RecordGroups(np.array([0, 0]), 2).columns_agree(mask) + assert RecordGroups(np.array([0, 1]), 2).columns_agree(mask) diff --git a/policyengine_uk_data/utils/calibrate.py b/policyengine_uk_data/utils/calibrate.py index fbd8f78d..512e5702 100644 --- a/policyengine_uk_data/utils/calibrate.py +++ b/policyengine_uk_data/utils/calibrate.py @@ -179,6 +179,19 @@ def totals(self, values: np.ndarray) -> np.ndarray: return values return np.bincount(self.codes, weights=np.asarray(values, dtype=float)) + def columns_agree(self, matrix: np.ndarray, chunk: int = 4096) -> bool: + """Whether every record's column (last axis) equals its group's first. + + Compared exactly, a chunk of records at a time. + """ + records = np.flatnonzero(self.first[self.codes] != np.arange(len(self.codes))) + for start in range(0, len(records), chunk): + columns = records[start : start + chunk] + firsts = self.first[self.codes[columns]] + if not np.array_equal(matrix[..., columns], matrix[..., firsts]): + return False + return True + def expand(self, weights: np.ndarray) -> np.ndarray: """Group weights (last axis) to records, split equally.""" if self.identity: @@ -310,11 +323,7 @@ def track_stage(stage_name: str): ) if groups is not None: r = np.asarray(r) - # One fingerprint per record's area mask; a group's must all agree. - fingerprint = r.T @ np.random.default_rng(0).random(r.shape[0]) - if not np.allclose( - fingerprint, fingerprint[record_groups.first[record_groups.codes]] - ): + if not record_groups.columns_agree(r): raise ValueError("A record group spans areas of different countries.") matrix = record_groups.means(matrix) r = r[:, record_groups.first] From 7270d61557888221bc2556b25b5f76323b1b3124 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 07:58:38 -0400 Subject: [PATCH 6/6] Refuse a split without tenure or with duplicate ids The model defaults tenure_type to renting privately, so a table without the column would have split every household; it now raises. Duplicate household or benefit unit ids get their own error rather than the unsorted-table one. Tests cover both, the no-renter path, and check column agreement against every same-group pair rather than the helper's own first-record index. Co-Authored-By: Claude Opus 5.5 --- policyengine_uk_data/datasets/brma.py | 9 +++++ policyengine_uk_data/tests/test_brma_split.py | 39 ++++++++++++++++++- 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/policyengine_uk_data/datasets/brma.py b/policyengine_uk_data/datasets/brma.py index 010c66a4..09887a30 100644 --- a/policyengine_uk_data/datasets/brma.py +++ b/policyengine_uk_data/datasets/brma.py @@ -306,6 +306,15 @@ def split_private_renters_across_brmas( "unit needs one to place it in a household." ) benunit_owner = owner.loc[benunit.benunit_id].to_numpy() + if "tenure_type" not in household.columns: + # The model would default every household to renting privately. + raise ValueError("The household table needs tenure_type to find renters.") + for entity, ids in ( + ("household", household.household_id), + ("benunit", benunit.benunit_id), + ): + if not ids.is_unique: + raise ValueError(f"The {entity} table has duplicate {entity}_id values.") year = int(dataset.time_period) simulation = Microsimulation(dataset=dataset) for entity, ids in ( diff --git a/policyengine_uk_data/tests/test_brma_split.py b/policyengine_uk_data/tests/test_brma_split.py index 8089db5d..f4f52316 100644 --- a/policyengine_uk_data/tests/test_brma_split.py +++ b/policyengine_uk_data/tests/test_brma_split.py @@ -678,9 +678,11 @@ def test_column_agreement_is_exact(labels, rows, chunk, seed): matrix = (rng.random((rows, len(labels))) < 0.5).astype(float) if rng.random() < 0.5: # often make the groups agree matrix = matrix[:, groups.first[groups.codes]] - first = groups.first[groups.codes] expected = all( - np.array_equal(matrix[:, j], matrix[:, first[j]]) for j in range(len(labels)) + np.array_equal(matrix[:, i], matrix[:, j]) + for i in range(len(labels)) + for j in range(len(labels)) + if labels[i] == labels[j] ) assert groups.columns_agree(matrix, chunk=chunk) == expected @@ -692,3 +694,36 @@ def test_one_hot_area_masks_that_differ_never_agree(): mask[506, 0] = mask[173, 1] = 1.0 assert not RecordGroups(np.array([0, 0]), 2).columns_agree(mask) assert RecordGroups(np.array([0, 1]), 2).columns_agree(mask) + + +@needs_policyengine +def test_a_dataset_without_private_renters_is_returned_unsplit(weights): + from policyengine_uk_data.datasets.brma import split_private_renters_across_brmas + + dataset = toy_dataset() + household = dataset.household.copy() + household["tenure_type"] = "OWNED_OUTRIGHT" + owners = _with_tables(dataset, household=household) + result = split_private_renters_across_brmas(owners, k=3, seed=0, weights=weights) + assert len(result.household) == len(household) + assert (result.household[BRMA_SPLIT_GROUP_COLUMN] == household.household_id).all() + assert (result.household.brma.astype(str) == "INNER_NORTH_LONDON").all() + + +@needs_policyengine +@pytest.mark.parametrize("problem", ["no_tenure", "duplicate_household"]) +def test_split_refuses_ambiguous_tables(problem, weights): + from policyengine_uk_data.datasets.brma import split_private_renters_across_brmas + + dataset = toy_dataset() + household = dataset.household.copy() + if problem == "no_tenure": + household = household.drop(columns="tenure_type") + match = "needs tenure_type" + else: + household.loc[1, "household_id"] = household.household_id[0] + match = "duplicate household_id" + with pytest.raises(ValueError, match=match): + split_private_renters_across_brmas( + _with_tables(dataset, household=household), k=2, seed=0, weights=weights + )