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/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 9aec9bc6..09887a30 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,279 @@ 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) + 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 + 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`` (``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, 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. + + 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() + household["source_household_id"] = household.household_id.to_numpy() + if k == 1: + return UKSingleYearDataset( + person=dataset.person, + benunit=dataset.benunit, + 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() + 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 ( + ("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, + benunit=benunit, + household=household, + fiscal_year=dataset.time_period, + ) + 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() + 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) + + renter_ids = household.household_id.to_numpy()[renter] + 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())))) + 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..0af72911 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,18 @@ def main(): get_performance, ) + # A split household's records (within one OA clone) calibrate as one. + 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( dataset=frs, @@ -281,6 +311,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 +339,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..f4f52316 --- /dev/null +++ b/policyengine_uk_data/tests/test_brma_split.py @@ -0,0 +1,729 @@ +"""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: 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 + +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 RecordGroups + + +@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() + + +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)) + + +@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 + ) + + +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( + importlib.util.find_spec("policyengine_uk") is None, + reason="policyengine_uk not available", +) + + +def toy_dataset(id_offset: int = 0): + 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) + id_offset, + "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 - id_offset) % 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_groups_collapse_and_expand_without_changing_predictions(labels, areas, seed): + rng = np.random.default_rng(seed) + 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_grouped_records_calibrate_as_one_household(tmp_path, monkeypatch): + import h5py + 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) + + 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()) + + local = pd.DataFrame([[600.0, 300.0], [500.0, 400.0]]) + + def run(matrix, national, weights, groups): + torch.manual_seed(0) + n = len(weights) + 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(national), + pd.Series([1_500.0, 900.0]), + ), + area_count=2, + weight_file="w.h5", + dataset_key="2024", + 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). + 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], + ] + ) + # 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_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]]) + 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 + + +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]] + expected = all( + 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 + + +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) + + +@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 + ) diff --git a/policyengine_uk_data/utils/calibrate.py b/policyengine_uk_data/utils/calibrate.py index f97c4e1e..512e5702 100644 --- a/policyengine_uk_data/utils/calibrate.py +++ b/policyengine_uk_data/utils/calibrate.py @@ -145,6 +145,62 @@ def initialize_weight_priors( return priors +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 __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 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: + return weights + 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: if pd.api.types.is_bool_dtype(series): return series.fillna(False).to_numpy(dtype=bool) @@ -221,6 +277,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 +296,10 @@ 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). Each group is + calibrated as one household (``RecordGroups``), so its records keep + equal weights. """ if dataset_key is None: dataset_key = default_weight_dataset_key() @@ -257,12 +318,23 @@ 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: + r = np.asarray(r) + 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] 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 = 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"): @@ -278,7 +350,7 @@ 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, ) area_prior_weights = household_prior_weights / areas_per_household @@ -435,8 +507,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 @@ -490,9 +562,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() @@ -530,8 +606,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