diff --git a/changelog.d/2278.fixed.md b/changelog.d/2278.fixed.md new file mode 100644 index 0000000000..f46e05da0b --- /dev/null +++ b/changelog.d/2278.fixed.md @@ -0,0 +1 @@ +- Fixed group inputs being assigned to the wrong benefit unit or household when dataset tables were not sorted by identifier. diff --git a/docs/book/usage/simulations.md b/docs/book/usage/simulations.md index e41b453994..583a5b887a 100644 --- a/docs/book/usage/simulations.md +++ b/docs/book/usage/simulations.md @@ -271,6 +271,11 @@ Each table should contain the ID columns needed to link entities together: - `benunit`: `benunit_id`, plus any benefit-unit variables - `household`: `household_id`, plus any household-level variables +Benefit-unit and household rows can be in any order. The loader aligns each +year's group inputs by its ID column to the simulation's sorted group IDs; +person rows keep their supplied order. This alignment does not reorder the +caller's tables. Multi-year tables must still describe the same population. + Use `UKSingleYearDataset` when you have one cross-section and want PolicyEngine UK to extend it forward using its standard uprating assumptions. diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py index 3f769e8950..3eb5c63d4c 100644 --- a/policyengine_uk/simulation.py +++ b/policyengine_uk/simulation.py @@ -959,12 +959,30 @@ def get_first_array(variable_name: str) -> np.ndarray: ) ) + # Populations use sorted group IDs; arrays in a Core dataset follow + # the order of that period's supplied identifier array instead. + group_orders = { + entity: { + period: np.argsort(ids) for period, ids in data[f"{entity}_id"].items() + } + for entity in ("benunit", "household") + } + # Load all variable values for variable in data: for time_period in data[variable]: if variable not in self.tax_benefit_system.variables: continue - self.set_input(variable, time_period, data[variable][time_period]) + values = data[variable][time_period] + entity = self.tax_benefit_system.variables[variable].entity.key + # Leave scalars for Core to normalize and validate against + # the population size; only arrays have a group row order. + if entity in group_orders and np.ndim(values) > 0: + order = group_orders[entity].get( + time_period, group_orders[entity][first_time_period] + ) + values = np.asarray(values)[order] + self.set_input(variable, time_period, values) # Now convert to the new UKSingleYearDataset self.input_variables = self.get_known_variables() @@ -1000,6 +1018,13 @@ def build_from_multi_year_dataset(self, dataset: UKMultiYearDataset) -> None: """ if self.disable_economic_assumptions: dataset = dataset.copy() + # The reset copies base-year values positionally. Align only the + # copied group tables first when their row order differs by year. + for year in dataset.years: + for entity in ("benunit", "household"): + table = getattr(dataset[year], entity) + order = np.argsort(table[f"{entity}_id"].to_numpy()) + setattr(dataset[year], entity, table.iloc[order]) reset_growthfactor_uprating(dataset) # Ensure enum columns are encoded and _enum_columns is populated so @@ -1018,7 +1043,11 @@ def build_from_multi_year_dataset(self, dataset: UKMultiYearDataset) -> None: # Load variable values for all years for year in dataset.years: - for table in dataset[year].tables: + for entity, table in zip(dataset[year].table_names, dataset[year].tables): + if entity in ("benunit", "household"): + # Align each year's group inputs to the sorted population + # without reordering the caller's table or the people. + table = table.iloc[np.argsort(table[f"{entity}_id"].to_numpy())] for variable in table.columns: if variable not in self.tax_benefit_system.variables: continue diff --git a/policyengine_uk/tests/test_dataset_group_order.py b/policyengine_uk/tests/test_dataset_group_order.py new file mode 100644 index 0000000000..219b3f0501 --- /dev/null +++ b/policyengine_uk/tests/test_dataset_group_order.py @@ -0,0 +1,268 @@ +"""Group table rows name groups by identifier, independently in each year.""" + +import h5py +import numpy as np +import pandas as pd +import pytest +from policyengine_core.data import Dataset + +from policyengine_uk import Microsimulation, Simulation +from policyengine_uk.data import UKMultiYearDataset, UKSingleYearDataset +from policyengine_uk.scenarios import no_economic_assumptions +from policyengine_uk.system import system + +pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system") + +# A synthetic £100 monthly allowance for every claimant type makes the final +# award £1,200 when capital permits it, with no earnings, rent or deductions. +REFORM = { + f"gov.dwp.universal_credit.standard_allowance.amount.{kind}": { + 2025: 100, + 2026: 100, + } + for kind in ["SINGLE_YOUNG", "SINGLE_OLD", "COUPLE_YOUNG", "COUPLE_OLD"] +} +REFORM.update( + { + f"gov.dwp.universal_credit.deductions.abolish.{kind}": {2025: True, 2026: True} + for kind in ["advance", "third_party", "government"] + } +) + + +def group_dataset(year=2025, benunits=(70, 10, 40), households=(900, 100)): + person = pd.DataFrame( + { + "person_id": [303, 101, 202, 404], + "person_benunit_id": [70, 10, 10, 40], + "person_household_id": [900, 100, 100, 900], + "age": [40, 40, 39, 50], + "employment_income": 0.0, + } + ) + benunit = ( + pd.DataFrame( + { + "benunit_id": [10, 40, 70], + "benunit_reported_capital": ( + [0.0, 6_000.0, 20_000.0] if year == 2025 else [20_000.0, 0.0, 0.0] + ), + "would_claim_uc": True, + } + ) + .set_index("benunit_id", drop=False) + .loc[list(benunits)] + ) + household = ( + pd.DataFrame( + { + "household_id": [100, 900], + "household_weight": [10.0, 20.0] if year == 2025 else [11.0, 21.0], + "region": ["WALES", "LONDON"], + "rent": 0.0, + "tenure_type": "OWNED_OUTRIGHT", + "council_tax": [1_100.0, 2_200.0], + } + ) + .set_index("household_id", drop=False) + .loc[list(households)] + ) + return UKSingleYearDataset( + person=person, benunit=benunit, household=household, fiscal_year=year + ) + + +def core_dataset_file(dataset, path): + """The Core time-period-array layout, with group arrays in table row order.""" + with h5py.File(path, "w") as file: + for table in dataset.tables: + for variable in table.columns: + values = table[variable].to_numpy() + possible_values = getattr( + system.variables[variable], "possible_values", None + ) + if possible_values is not None: + values = np.asarray(possible_values.encode(values)) + file.create_dataset(f"{variable}/{dataset.time_period}", data=values) + return path + + +def array(sim, variable, year, **kwargs): + result = sim.calculate(variable, year, **kwargs) + if hasattr(result, "decode_to_str"): + return result.decode_to_str() + return np.asarray(result) + + +def assert_group_values(sim, year=2025): + np.testing.assert_array_equal(sim.populations["person"].ids, [303, 101, 202, 404]) + for entity, expected in [("benunit", [10, 40, 70]), ("household", [100, 900])]: + np.testing.assert_array_equal(sim.populations[entity].ids, expected) + np.testing.assert_array_equal(array(sim, f"{entity}_id", year), expected) + np.testing.assert_array_equal(array(sim, "person_id", year), [303, 101, 202, 404]) + np.testing.assert_array_equal( + array(sim, "person_benunit_id", year), [70, 10, 10, 40] + ) + np.testing.assert_array_equal( + array(sim, "region", year, map_to="person"), + ["LONDON", "WALES", "WALES", "LONDON"], + ) + weights = [10, 20] if year == 2025 else [11, 21] + np.testing.assert_array_equal(array(sim, "household_weight", year), weights) + np.testing.assert_array_equal( + array(sim, "household_weight", year, map_to="person"), + [weights[1], weights[0], weights[0], weights[1]], + ) + capital = [0, 6_000, 20_000] if year == 2025 else [20_000, 0, 0] + np.testing.assert_array_equal(array(sim, "benunit_reported_capital", year), capital) + np.testing.assert_array_equal( + array(sim, "benunit_reported_capital", year, map_to="person"), + [capital[2], capital[0], capital[0], capital[1]], + ) + expected_uc = [1_200, 1_200, 0] if year == 2025 else [0, 1_200, 1_200] + np.testing.assert_allclose(array(sim, "universal_credit", year), expected_uc) + + +@pytest.mark.parametrize( + "source", ["single", "multi", "single_h5", "multi_h5", "core", "core_h5"] +) +def test_unsorted_group_tables_keep_their_inputs(source, tmp_path): + single = group_dataset() + before = {name: getattr(single, name).copy() for name in single.table_names} + multi = UKMultiYearDataset(datasets=[single]) + if source in ["single", "multi"]: + dataset = single if source == "single" else multi + elif source in ["single_h5", "multi_h5"]: + dataset = tmp_path / "uk.h5" + (single if source == "single_h5" else multi).save(dataset) + else: + dataset = core_dataset_file(single, tmp_path / "core.h5") + if source == "core": + dataset = Dataset.from_file(dataset) + + sim = Simulation(dataset=dataset, reform=REFORM) + + assert_group_values(sim) + # Existing enum encoding may change storage; alignment must not reorder + # caller tables or alter their public values and identifier dtypes. + for name, frame in before.items(): + actual = getattr(single, name) + pd.testing.assert_frame_equal(actual, frame, check_dtype=False) + for column in frame.columns: + if column.endswith("_id"): + assert actual[column].dtype == frame[column].dtype + + +def test_each_year_aligns_its_own_group_row_order(): + multi = UKMultiYearDataset( + datasets=[group_dataset(), group_dataset(2026, (40, 70, 10), (100, 900))] + ) + sim = Microsimulation(dataset=multi, reform=REFORM) + + assert_group_values(sim, 2025) + assert_group_values(sim, 2026) + + +def test_sorted_group_tables_keep_the_same_results(): + sim = Simulation( + dataset=group_dataset(benunits=(10, 40, 70), households=(100, 900)), + reform=REFORM, + ) + + assert_group_values(sim) + + +def test_person_dataframe_rows_keep_their_own_group_values(): + frame = pd.DataFrame( + { + "person_id__2025": [303, 101, 202, 404], + "benunit_id__2025": [70, 10, 10, 40], + "household_id__2025": [900, 100, 100, 900], + "age__2025": [40, 40, 39, 50], + "employment_income__2025": 0.0, + "benunit_reported_capital__2025": [20_000.0, 0.0, 0.0, 6_000.0], + "household_weight__2025": [20.0, 10.0, 10.0, 20.0], + "region__2025": ["LONDON", "WALES", "WALES", "LONDON"], + "tenure_type__2025": "OWNED_OUTRIGHT", + "council_tax__2025": 0.0, + "rent__2025": 0.0, + } + ) + before = frame.copy() + sim = Simulation(dataset=frame, reform=REFORM) + + assert_group_values(sim) + pd.testing.assert_frame_equal(frame, before) + + +def test_frozen_assumptions_align_each_year_before_reset(): + source = UKMultiYearDataset( + datasets=[group_dataset(), group_dataset(2026, (40, 70, 10), (100, 900))] + ) + before = source[2026].copy() + sim = Simulation(dataset=source, scenario=no_economic_assumptions) + + for year in (2025, 2026): + np.testing.assert_array_equal(array(sim, "household_id", year), [100, 900]) + np.testing.assert_array_equal(array(sim, "benunit_id", year), [10, 40, 70]) + np.testing.assert_array_equal(array(sim, "household_weight", year), [10, 20]) + np.testing.assert_array_equal(array(sim, "council_tax", year), [1_100, 2_200]) + np.testing.assert_array_equal( + array(sim, "benunit_reported_capital", year), [0, 6_000, 20_000] + ) + for name in before.table_names: + pd.testing.assert_frame_equal( + getattr(source[2026], name), getattr(before, name) + ) + + +def scalar_core_dataset_file(path, group_count): + """A real Core file with a zero-dimensional council-tax observation.""" + ids = np.arange(group_count, 0, -1) + data = { + "person_id": ids, + "person_benunit_id": ids * 10, + "person_household_id": ids * 100, + "benunit_id": ids * 10, + "household_id": ids * 100, + "age": np.full(group_count, 40), + "household_weight": np.ones(group_count), + "rent": np.zeros(group_count), + "council_tax": 1_000.0, + } + for variable, value in [("region", "WALES"), ("tenure_type", "OWNED_OUTRIGHT")]: + data[variable] = np.asarray( + system.variables[variable].possible_values.encode( + np.full(group_count, value) + ) + ) + with h5py.File(path, "w") as file: + for variable, values in data.items(): + file.create_dataset(f"{variable}/2025", data=values) + return path + + +@pytest.mark.parametrize("source", ["core", "core_h5"]) +def test_scalar_core_input_for_one_household(source, tmp_path): + path = scalar_core_dataset_file(tmp_path / "scalar.h5", group_count=1) + dataset = Dataset.from_file(path) if source == "core" else path + + sim = Simulation(dataset=dataset) + + np.testing.assert_array_equal(sim.populations["household"].ids, [100]) + np.testing.assert_array_equal(array(sim, "council_tax", 2025), [1_000]) + np.testing.assert_array_equal( + array(sim, "council_tax", 2025, map_to="person"), [1_000] + ) + + +@pytest.mark.parametrize("source", ["core", "core_h5"]) +def test_scalar_core_input_does_not_broadcast_to_multiple_households(source, tmp_path): + path = scalar_core_dataset_file(tmp_path / "scalar.h5", group_count=2) + dataset = Dataset.from_file(path) if source == "core" else path + + # Core normalizes a scalar to length one and rejects a population mismatch. + with pytest.raises( + ValueError, match=r'variable "council_tax".*length is 1.*2 households' + ): + Simulation(dataset=dataset)