From f91065f4c6972a923c19d5f28de66c90fb80d7c2 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 12:12:39 -0400 Subject: [PATCH 1/3] Export identifiers for every year an input frame gives values for Simulation.to_input_dataframe() now gives the person, benefit unit and household identifiers for each whole year the frame gives other values for. The dataset a frame loads as holds them in each of those years, so a simulation built from the frame records them there and exported them, while the simulation that made the frame held them in a year after its data only by carrying them forward. Since #2251 every simulation records uc_receives_new_claimant_health_element for 2026 to 2029, so every frame export failed the round trip that #2270's property test checks: 18 columns exported, 38 re-exported. Adds the minimal counterexample as an explicit example, regression tests for the rebalancing's years and for inputs set by hand before, among and after them, a unit test of the helper, and a property that a frame over several years, given inputs for other years once built, exports the same frame again. Co-Authored-By: Claude Opus 5.5 --- .../fix-roundtrip-recorded-inputs.fixed.md | 1 + docs/book/usage/simulations.md | 6 + policyengine_uk/data/input_dataframe.py | 92 ++++++++- policyengine_uk/simulation.py | 28 ++- .../tests/test_dataframe_dataset.py | 193 +++++++++++++++++- 5 files changed, 313 insertions(+), 7 deletions(-) create mode 100644 changelog.d/fix-roundtrip-recorded-inputs.fixed.md diff --git a/changelog.d/fix-roundtrip-recorded-inputs.fixed.md b/changelog.d/fix-roundtrip-recorded-inputs.fixed.md new file mode 100644 index 0000000000..9d7e60fdcf --- /dev/null +++ b/changelog.d/fix-roundtrip-recorded-inputs.fixed.md @@ -0,0 +1 @@ +- Fixed `Simulation.to_input_dataframe()` leaving out the person, benefit unit and household identifiers for years in which a simulation was given inputs after it was built, such as the Universal Credit health element status for 2026 to 2029, so that an exported frame now loads back to a simulation that exports the same frame. diff --git a/docs/book/usage/simulations.md b/docs/book/usage/simulations.md index 4128daf55b..aa00d945df 100644 --- a/docs/book/usage/simulations.md +++ b/docs/book/usage/simulations.md @@ -200,6 +200,12 @@ print(sim.calculate("council_tax", 2025)) # One value per household period that is not one year are set on the simulation directly. A variable given for both months and the year is set column by column in the frame's order, since a year's value is spread over the months not already set. +- `to_input_dataframe()` gives the identifier columns for every year it gives + other values for, which is where the dataset the frame loads as holds them. + So an exported frame loads back to a simulation that exports the same frame, + even when the simulation was given inputs for years after its data, such as + the Universal Credit health element status the rebalancing records for 2026 + to 2029. ```{important} The model stores `person_household_id` and `person_benunit_id` as 32-bit diff --git a/policyengine_uk/data/input_dataframe.py b/policyengine_uk/data/input_dataframe.py index fcc6b133d1..60da596084 100644 --- a/policyengine_uk/data/input_dataframe.py +++ b/policyengine_uk/data/input_dataframe.py @@ -15,6 +15,11 @@ Identifiers are read as whole numbers, value by value, and a float that may stand for more than one whole number is refused rather than allowed to merge two people or groups. + +The identifiers hold for every period, and the dataset a frame loads as holds +them in each year the frame gives data for. ``Simulation.to_input_dataframe`` +writes them for each of those years too (``with_identifiers_in_each_year``), +so that an exported frame loads back to a simulation that exports it again. """ import warnings @@ -267,6 +272,17 @@ def _holds_floats(identifiers: pd.Series) -> bool: return values.dtype.kind == "f" +def _is_whole_year(period) -> bool: + """Whether ``period`` is one calendar year, the period of a dataset + year.""" + start = period.start + return ( + period.unit == "year" + and period.size == 1 + and (start.month, start.day) == (1, 1) + ) + + def _same(left: np.ndarray, right: np.ndarray) -> np.ndarray: """Whether each pair of values is equal, counting two missing values as equal.""" @@ -362,11 +378,7 @@ def no_tables() -> Dict[str, Dict[str, np.ndarray]]: continue period = period_(period_text) start = period.start - whole_year = ( - period.unit == "year" - and period.size == 1 - and (start.month, start.day) == (1, 1) - ) + whole_year = _is_whole_year(period) if variable in IDENTIFIER_VARIABLES: if whole_year: identifier_years.add(start.year) @@ -425,3 +437,73 @@ def no_tables() -> Dict[str, Dict[str, np.ndarray]]: for fiscal_year, tables in sorted(years.items()) ] return UKMultiYearDataset(datasets=datasets), replayed, identifier_inputs + + +def with_identifiers_in_each_year( + df: pd.DataFrame, tax_benefit_system: "TaxBenefitSystem" +) -> pd.DataFrame: + """``df``, an exported frame, with each identifier it gives for some year + given for every year it gives other values for. + + The dataset a frame loads as holds the people and groups in each year the + frame gives values for (``dataset_from_input_dataframe``), so a simulation + built from it holds the identifiers as inputs in those years, and exports + them there. A simulation that records an input for a later year once it + is built, as the Universal Credit rebalancing records each benefit unit's + health element status for 2026 to 2029, holds its identifiers in that + year only by carrying them forward, and policyengine-core exports only + the periods an input was set for. Giving the identifiers for that year + makes the frame load back to a simulation that exports the same frame. + + A simulation holds one set of people and groups, so every column of an + identifier names the same ones (the loader refuses a frame whose columns + differ). An identifier added for a year copies its column for the latest + year before it the frame gives, or else the earliest year it gives. Each + goes beside the identifier's other columns, in order of period, where + policyengine-core exports an input set for that year. + + Years are those ``dataset_from_input_dataframe`` makes dataset years: the + whole calendar years that columns of other variables of + ``tax_benefit_system`` name. ``df`` is not changed. + """ + named: Dict[str, Dict[int, object]] = {} + years = set() + for column in df.columns: + variable, period_text = split_column(column) + if variable not in tax_benefit_system.variables: + continue + period = period_(period_text) + if not _is_whole_year(period): + continue + if variable in IDENTIFIER_VARIABLES: + named.setdefault(variable, {})[period.start.year] = column + else: + years.add(period.start.year) + added = { + variable: sorted(years - set(columns)) + for variable, columns in named.items() + if years - set(columns) + } + if not added: + return df + result = df.copy(deep=False) + for variable, missing in added.items(): + columns = named[variable] + for year in missing: + earlier = [named_year for named_year in columns if named_year < year] + source = columns[max(earlier) if earlier else min(columns)] + name = f"{variable}{COLUMN_SEPARATOR}{year}" + # The identifier's columns, by period text, as the exporter sorts + # them; the new one goes before the first that sorts after it. + own = [ + position + for position, column in enumerate(result.columns) + if split_column(column)[0] == variable + ] + later = [ + position + for position in own + if split_column(result.columns[position])[1] > str(year) + ] + result.insert(later[0] if later else own[-1] + 1, name, df[source].copy()) + return result diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py index 2acef5aefe..a709ffd65d 100644 --- a/policyengine_uk/simulation.py +++ b/policyengine_uk/simulation.py @@ -32,7 +32,10 @@ reset_growthfactor_uprating, ) from policyengine_uk.data.dataset_sources import materialize_gcs_dataset_url -from policyengine_uk.data.input_dataframe import dataset_from_input_dataframe +from policyengine_uk.data.input_dataframe import ( + dataset_from_input_dataframe, + with_identifiers_in_each_year, +) from policyengine_uk.utils.dependencies import get_variable_dependencies from policyengine_uk.utils.supplied_inputs import ( SUPPLIED_INPUT_VARIABLES, @@ -426,6 +429,29 @@ def build_from_dataframe(self, df: pd.DataFrame) -> None: for variable, time_period, values in replayed: self.set_input(variable, time_period, values) + def to_input_dataframe( + self, include_computed_variables: bool = False + ) -> pd.DataFrame: + """Export the simulation's inputs as a DataFrame that + ``Simulation(dataset=...)`` loads back. + + policyengine-core exports each input for the periods it was set for. + The identifiers are also given for every year the frame gives other + values for, as the dataset the frame loads as holds them, so that the + frame loads back to a simulation that exports the same frame. See + ``policyengine_uk.data.input_dataframe.with_identifiers_in_each_year``. + + Args: + include_computed_variables: If True, export every stored value, + calculated ones included, as policyengine-core does. + """ + return with_identifiers_in_each_year( + super().to_input_dataframe( + include_computed_variables=include_computed_variables + ), + self.tax_benefit_system, + ) + def build_from_dataset(self, dataset: Dataset) -> None: """Build simulation from a Dataset object. diff --git a/policyengine_uk/tests/test_dataframe_dataset.py b/policyengine_uk/tests/test_dataframe_dataset.py index c3793a4f51..30d5e9b595 100644 --- a/policyengine_uk/tests/test_dataframe_dataset.py +++ b/policyengine_uk/tests/test_dataframe_dataset.py @@ -30,6 +30,12 @@ 5. Round trip: the frame a simulation exports loads back to the same people, groups and inputs, and exports the same frame again; so does that frame saved as CSV, which holds the exported 32-bit identifiers as 64-bit floats. + This holds for frames that name the identifiers and give values for + different years, and for simulations given inputs for other years once + built, as the Universal Credit rebalancing gives each benefit unit's + health element status for 2026 to 2029. The export gives the identifiers, + exactly, for every year it gives other values for, which is where the + dataset the frame loads as holds them. 6. Differential: a simulation built from a dataset and one built from the frame it exports give the same value for every variable of the model. 7. Exact names: a column is read only if the part of its name before ``__`` @@ -67,7 +73,7 @@ import numpy as np import pandas as pd import pytest -from hypothesis import HealthCheck, given, settings +from hypothesis import HealthCheck, example, given, settings from hypothesis import strategies as st from policyengine_uk import CountryTaxBenefitSystem, Microsimulation, Simulation @@ -77,6 +83,7 @@ IDENTIFIER_VARIABLES, dataset_from_input_dataframe, split_column, + with_identifiers_in_each_year, ) from policyengine_uk.scenarios import no_economic_assumptions @@ -601,6 +608,57 @@ def test_a_frame_that_names_no_whole_year_is_held_in_the_default_input_year(): ) == {1: True} +def test_an_export_gives_its_identifiers_for_each_year_it_gives_values_for( + tax_benefit_system, +): + # Values for 2024, 2025, 2027 and 2029 make those years of the dataset, + # which holds the identifiers in each. A month, a year only an identifier + # names, and a column that names no variable add none. Each added column + # copies the identifier's latest earlier column, or else its earliest: + # here person_id__2028 is given as floats, so person_id__2029 is too. + frame_ = pd.DataFrame( + { + "benunit_id__2025": [7, 7], + "person_id__2025": [1, 2], + "person_id__2028": [1.0, 2.0], + "age__2024": [39, 37], + "age__2025": [40, 38], + "age__2027": [42, 40], + "ebr_energy_bills_credit_eligible__2026-01": [True, True], + "council_tax__2029": [1_500.0, 1_500.0], + "previous_household_id__2030": [5, 5], + } + ) + before = frame_.copy() + + exported = with_identifiers_in_each_year(frame_, tax_benefit_system) + + assert exported.columns.tolist() == [ + "benunit_id__2024", + "benunit_id__2025", + "benunit_id__2027", + "benunit_id__2029", + "person_id__2024", + "person_id__2025", + "person_id__2027", + "person_id__2028", + "person_id__2029", + *frame_.columns[3:], + ] + for year, source in [(2024, 2025), (2027, 2025), (2029, 2028)]: + pd.testing.assert_series_equal( + exported[f"person_id__{year}"], + frame_[f"person_id__{source}"], + check_names=False, + ) + pd.testing.assert_frame_equal(exported[frame_.columns], frame_) + pd.testing.assert_frame_equal(frame_, before) + # A frame that gives its identifiers for each such year gains nothing. + pd.testing.assert_frame_equal( + with_identifiers_in_each_year(exported, tax_benefit_system), exported + ) + + # The dataset a frame-built simulation holds. @@ -1140,6 +1198,29 @@ def test_members_may_all_leave_a_group_value_missing(tax_benefit_system): suppress_health_check=[HealthCheck.too_slow, HealthCheck.data_too_large], ) ROUND_TRIP_SETTINGS = settings(SIMULATION_SETTINGS, max_examples=15) +# The smallest example of the round trip failing once the Universal Credit +# rebalancing recorded each benefit unit's health element status for 2026 to +# 2029 (#2251). The export gave that status for those years but the person +# and their groups for 2025 alone. Loaded back, the frame was a dataset with a +# year for each of 2025 to 2029, which holds the identifiers in each, so the +# second export gave them for 2026 to 2029 as well: 38 columns, not 18. +ONE_PERSON = Population( + person_ids=np.array([1]), + groups={"benunit": np.array([1]), "household": np.array([1])}, + values={ + "age": np.array([0]), + "hours_worked": np.array([0]), + "gender": np.array(["MALE"]), + "is_disabled_for_benefits": np.array([False]), + "benunit_reported_capital": np.array([0]), + "would_claim_uc": np.array([False]), + "council_tax": np.array([0]), + "household_weight": np.array([0]), + "region": np.array(["LONDON"]), + }, +) +# Years around the data year, for frames and inputs that span several. +YEARS = list(range(YEAR - 1, YEAR + 6)) TABLE_SETTINGS = settings( max_examples=200, deadline=None, @@ -1165,6 +1246,7 @@ def test_a_frame_builds_its_people_groups_and_values(population, id_columns): @ROUND_TRIP_SETTINGS @given(population=populations(), id_columns=id_columns) +@example(population=ONE_PERSON, id_columns="both") def test_an_exported_frame_loads_back_and_exports_itself(population, id_columns): simulation = Simulation(dataset=frame(population, id_columns)) exported = simulation.to_input_dataframe() @@ -1187,6 +1269,115 @@ def test_an_exported_frame_loads_back_and_exports_itself(population, id_columns) assert_holds(Simulation(dataset=from_csv), population) +def columns_of(frame_: pd.DataFrame, variable: str) -> list: + return [column for column in frame_ if split_column(column)[0] == variable] + + +def test_an_export_names_the_people_in_each_year_the_rebalancing_records(): + # The rebalancing records the status of a simulation built from data for + # 2025 for 2026 to 2029, where it holds its people and groups by carrying + # them forward from 2025. + exported = Simulation(dataset=frame(ONE_PERSON)).to_input_dataframe() + + status = "uc_receives_new_claimant_health_element" + years = range(YEAR, 2030) + assert columns_of(exported, status) == [f"{status}__{year}" for year in years[1:]] + for identifier in IDENTIFIER_VARIABLES: + columns = columns_of(exported, identifier) + assert columns == [f"{identifier}__{year}" for year in years] + # Together, where the exporter puts an input set for each year. + start = exported.columns.get_loc(columns[0]) + assert exported.columns[start : start + len(columns)].tolist() == columns + for column in columns: + assert exported[column].tolist() == [1] + reloaded = Simulation(dataset=exported) + pd.testing.assert_frame_equal(reloaded.to_input_dataframe(), exported) + + +@pytest.mark.parametrize("year", [YEAR - 1, YEAR + 2, YEAR + 7]) +def test_an_input_set_once_built_for_another_year_round_trips(year): + # Not only the rebalancing: any input set for a year the data does not + # give, before it, among the rebalancing's years or after them. + simulation = Simulation(dataset=frame(ONE_PERSON)) + simulation.set_input("hours_worked", year, np.array([1_000])) + exported = simulation.to_input_dataframe() + + for identifier in IDENTIFIER_VARIABLES: + assert exported[f"{identifier}__{year}"].tolist() == [1] + reloaded = Simulation(dataset=exported) + assert reloaded.calculate("hours_worked", year).tolist() == [1_000] + pd.testing.assert_frame_equal(reloaded.to_input_dataframe(), exported) + + +def multi_year_frame( + population: Population, id_columns: str, identifier_years, data_years +) -> pd.DataFrame: + """The frame of ``population`` with its identifiers given for + ``identifier_years`` and its values for ``data_years``.""" + columns = {} + for year in sorted(set(identifier_years) | set(data_years)): + for column, values in frame(population, id_columns, year).items(): + identifier = split_column(column)[0] in IDENTIFIER_VARIABLES + if year in (identifier_years if identifier else data_years): + columns[column] = values + return pd.DataFrame(columns) + + +@ROUND_TRIP_SETTINGS +@given( + population=populations(), + id_columns=id_columns, + identifier_years=st.sets(st.sampled_from(YEARS), min_size=1, max_size=2), + data_years=st.sets(st.sampled_from(YEARS), min_size=1, max_size=3), + later_years=st.sets(st.sampled_from(YEARS), max_size=2), +) +@example( + population=ONE_PERSON, + id_columns="both", + identifier_years={YEAR}, + data_years={YEAR}, + later_years=set(), +) +@example( + population=ONE_PERSON, + id_columns="own", + identifier_years={YEAR + 2}, + data_years={YEAR - 1, YEAR + 3}, + later_years={YEAR + 5}, +) +def test_a_frame_over_several_years_exports_the_same_frame_again( + population, id_columns, identifier_years, data_years, later_years +): + # 5. Round trip, for frames that name the identifiers and give values + # for different years, and simulations given an input for other years + # once built. + simulation = Simulation( + dataset=multi_year_frame(population, id_columns, identifier_years, data_years) + ) + for year in sorted(later_years): + simulation.set_input( + "hours_worked", year, population.values["hours_worked"] + year + ) + exported = simulation.to_input_dataframe() + + reloaded = Simulation(dataset=exported) + + pd.testing.assert_frame_equal(reloaded.to_input_dataframe(), exported) + # The export gives the exact identifiers in every year it gives values + # for, and in each year the frame named them for. + years = { + int(period) + for variable, period in map(split_column, exported.columns) + if variable not in IDENTIFIER_VARIABLES and period.isdigit() + } + assert years >= set(data_years) | set(later_years) + for year in years | set(identifier_years): + assert exported[f"person_id__{year}"].tolist() == population.person_ids.tolist() + for group in GROUPS: + exact = population.groups[group].tolist() + assert exported[f"{group}_id__{year}"].tolist() == exact + + @TABLE_SETTINGS @given(population=populations(), id_columns=id_columns, data=st.data()) def test_tables_hold_the_frame_in_any_row_order( From 49e7ef3e01ea4bc20b27c4db19b1930bab9f3d10 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 13:29:48 -0400 Subject: [PATCH 2/3] Give every identifier for each year any column names The loader sets all five identifiers for each whole year an identifier column names, as well as for each year of the dataset, so an export that gives one identifier for a year alone (person_id set by hand for 2032) still re-exported four more columns. The helper now adds every identifier the frame gives for each whole year a model-variable column names, identifier columns included. Also corrects the account of a year before the data, where inputs do not carry back and the exporting simulation's identifiers are zero, and states that the ordering matches the default (sorted) export. Adds an identifier-only case to the regression and the property. Co-Authored-By: Claude Opus 5.5 --- .../fix-roundtrip-recorded-inputs.fixed.md | 2 +- docs/book/usage/simulations.md | 13 ++-- policyengine_uk/data/input_dataframe.py | 52 ++++++++------ policyengine_uk/simulation.py | 6 +- .../tests/test_dataframe_dataset.py | 72 +++++++++++++------ 5 files changed, 90 insertions(+), 55 deletions(-) diff --git a/changelog.d/fix-roundtrip-recorded-inputs.fixed.md b/changelog.d/fix-roundtrip-recorded-inputs.fixed.md index 9d7e60fdcf..69bb07a587 100644 --- a/changelog.d/fix-roundtrip-recorded-inputs.fixed.md +++ b/changelog.d/fix-roundtrip-recorded-inputs.fixed.md @@ -1 +1 @@ -- Fixed `Simulation.to_input_dataframe()` leaving out the person, benefit unit and household identifiers for years in which a simulation was given inputs after it was built, such as the Universal Credit health element status for 2026 to 2029, so that an exported frame now loads back to a simulation that exports the same frame. +- Fixed `Simulation.to_input_dataframe()` leaving out the person, benefit unit and household identifiers for years in which a simulation was given inputs after it was built, such as the Universal Credit health element status for 2026 to 2029, so that an exported frame loads back to a simulation that exports the same frame. diff --git a/docs/book/usage/simulations.md b/docs/book/usage/simulations.md index aa00d945df..7b54a694da 100644 --- a/docs/book/usage/simulations.md +++ b/docs/book/usage/simulations.md @@ -200,12 +200,13 @@ print(sim.calculate("council_tax", 2025)) # One value per household period that is not one year are set on the simulation directly. A variable given for both months and the year is set column by column in the frame's order, since a year's value is spread over the months not already set. -- `to_input_dataframe()` gives the identifier columns for every year it gives - other values for, which is where the dataset the frame loads as holds them. - So an exported frame loads back to a simulation that exports the same frame, - even when the simulation was given inputs for years after its data, such as - the Universal Credit health element status the rebalancing records for 2026 - to 2029. +- `to_input_dataframe()` gives the identifier columns for every year that one + of its columns gives a whole year's value for, since a simulation built from + the frame holds the identifiers in each of those years. So an exported frame + loads back to a simulation that exports the same frame, even when the + simulation was given inputs for years outside its data, such as the + Universal Credit health element status the rebalancing records for 2026 to + 2029. ```{important} The model stores `person_household_id` and `person_benunit_id` as 32-bit diff --git a/policyengine_uk/data/input_dataframe.py b/policyengine_uk/data/input_dataframe.py index 60da596084..220aad79dc 100644 --- a/policyengine_uk/data/input_dataframe.py +++ b/policyengine_uk/data/input_dataframe.py @@ -16,10 +16,11 @@ stand for more than one whole number is refused rather than allowed to merge two people or groups. -The identifiers hold for every period, and the dataset a frame loads as holds -them in each year the frame gives data for. ``Simulation.to_input_dataframe`` -writes them for each of those years too (``with_identifiers_in_each_year``), -so that an exported frame loads back to a simulation that exports it again. +The identifiers hold for every period, and a simulation built from a frame +holds them all as inputs in each whole year a column names. +``Simulation.to_input_dataframe`` gives them for each such year too +(``with_identifiers_in_each_year``), so that an exported frame loads back to +a simulation that exports it again. """ import warnings @@ -443,28 +444,36 @@ def with_identifiers_in_each_year( df: pd.DataFrame, tax_benefit_system: "TaxBenefitSystem" ) -> pd.DataFrame: """``df``, an exported frame, with each identifier it gives for some year - given for every year it gives other values for. - - The dataset a frame loads as holds the people and groups in each year the - frame gives values for (``dataset_from_input_dataframe``), so a simulation - built from it holds the identifiers as inputs in those years, and exports - them there. A simulation that records an input for a later year once it - is built, as the Universal Credit rebalancing records each benefit unit's - health element status for 2026 to 2029, holds its identifiers in that - year only by carrying them forward, and policyengine-core exports only - the periods an input was set for. Giving the identifiers for that year - makes the frame load back to a simulation that exports the same frame. + given for every whole calendar year a column of the frame names. + + ``dataset_from_input_dataframe`` holds every identifier in each such year: + a year a column of another variable names is a year of the dataset, which + holds the people and groups, and each other year an identifier column + names gets them all as inputs. So a simulation built from the frame holds + the identifiers as inputs in those years, and exports them there. + + The simulation that exported the frame may not have. One given an input + for a year after its data once it was built, as the Universal Credit + rebalancing records each benefit unit's health element status for 2026 + to 2029, holds its identifiers in that year only by carrying them forward + from an earlier one; and in a year before its data, where inputs do not + carry back, its identifiers are the default, zero. policyengine-core + exports an + input only for the periods it was set for. Giving every identifier for + each year makes the frame load back to a simulation that exports the + same frame. A simulation holds one set of people and groups, so every column of an identifier names the same ones (the loader refuses a frame whose columns differ). An identifier added for a year copies its column for the latest year before it the frame gives, or else the earliest year it gives. Each - goes beside the identifier's other columns, in order of period, where - policyengine-core exports an input set for that year. + goes beside the identifier's other columns, in order of period, which is + where the default export (sorted by period) puts an input set for that + year. - Years are those ``dataset_from_input_dataframe`` makes dataset years: the - whole calendar years that columns of other variables of - ``tax_benefit_system`` name. ``df`` is not changed. + Only columns that name a variable of ``tax_benefit_system`` count, and an + identifier the frame does not give for any year is not added. ``df`` is + not changed. """ named: Dict[str, Dict[int, object]] = {} years = set() @@ -477,8 +486,7 @@ def with_identifiers_in_each_year( continue if variable in IDENTIFIER_VARIABLES: named.setdefault(variable, {})[period.start.year] = column - else: - years.add(period.start.year) + years.add(period.start.year) added = { variable: sorted(years - set(columns)) for variable, columns in named.items() diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py index a709ffd65d..26bbbe18f7 100644 --- a/policyengine_uk/simulation.py +++ b/policyengine_uk/simulation.py @@ -436,9 +436,9 @@ def to_input_dataframe( ``Simulation(dataset=...)`` loads back. policyengine-core exports each input for the periods it was set for. - The identifiers are also given for every year the frame gives other - values for, as the dataset the frame loads as holds them, so that the - frame loads back to a simulation that exports the same frame. See + The identifiers are also given for every whole year a column names, + where a simulation built from the frame holds them, so that the frame + loads back to a simulation that exports the same frame. See ``policyengine_uk.data.input_dataframe.with_identifiers_in_each_year``. Args: diff --git a/policyengine_uk/tests/test_dataframe_dataset.py b/policyengine_uk/tests/test_dataframe_dataset.py index 30d5e9b595..6ba665b8fa 100644 --- a/policyengine_uk/tests/test_dataframe_dataset.py +++ b/policyengine_uk/tests/test_dataframe_dataset.py @@ -33,9 +33,10 @@ This holds for frames that name the identifiers and give values for different years, and for simulations given inputs for other years once built, as the Universal Credit rebalancing gives each benefit unit's - health element status for 2026 to 2029. The export gives the identifiers, - exactly, for every year it gives other values for, which is where the - dataset the frame loads as holds them. + health element status for 2026 to 2029. The export gives every + identifier for each whole year a column names, which is where a + simulation built from it holds them, and person_id and the benefit unit + and household identifiers exactly. 6. Differential: a simulation built from a dataset and one built from the frame it exports give the same value for every variable of the model. 7. Exact names: a column is read only if the part of its name before ``__`` @@ -612,10 +613,11 @@ def test_an_export_gives_its_identifiers_for_each_year_it_gives_values_for( tax_benefit_system, ): # Values for 2024, 2025, 2027 and 2029 make those years of the dataset, - # which holds the identifiers in each. A month, a year only an identifier - # names, and a column that names no variable add none. Each added column - # copies the identifier's latest earlier column, or else its earliest: - # here person_id__2028 is given as floats, so person_id__2029 is too. + # which holds the identifiers in each, and the loader sets them all for + # 2028, which an identifier names. A month and a column that names no + # variable add no year. Each added column copies the identifier's latest + # earlier column, or else its earliest: here person_id__2028 is given as + # floats, so person_id__2029 is too. frame_ = pd.DataFrame( { "benunit_id__2025": [7, 7], @@ -637,6 +639,7 @@ def test_an_export_gives_its_identifiers_for_each_year_it_gives_values_for( "benunit_id__2024", "benunit_id__2025", "benunit_id__2027", + "benunit_id__2028", "benunit_id__2029", "person_id__2024", "person_id__2025", @@ -645,10 +648,15 @@ def test_an_export_gives_its_identifiers_for_each_year_it_gives_values_for( "person_id__2029", *frame_.columns[3:], ] - for year, source in [(2024, 2025), (2027, 2025), (2029, 2028)]: + for identifier, year, source in [ + ("person_id", 2024, 2025), + ("person_id", 2027, 2025), + ("person_id", 2029, 2028), + ("benunit_id", 2028, 2025), + ]: pd.testing.assert_series_equal( - exported[f"person_id__{year}"], - frame_[f"person_id__{source}"], + exported[f"{identifier}__{year}"], + frame_[f"{identifier}__{source}"], check_names=False, ) pd.testing.assert_frame_equal(exported[frame_.columns], frame_) @@ -1294,18 +1302,27 @@ def test_an_export_names_the_people_in_each_year_the_rebalancing_records(): pd.testing.assert_frame_equal(reloaded.to_input_dataframe(), exported) -@pytest.mark.parametrize("year", [YEAR - 1, YEAR + 2, YEAR + 7]) -def test_an_input_set_once_built_for_another_year_round_trips(year): +@pytest.mark.parametrize( + "variable, year, value", + [ + ("hours_worked", YEAR - 1, 1_000), + ("hours_worked", YEAR + 2, 1_000), + ("hours_worked", YEAR + 7, 1_000), + # One identifier for a year: the loader sets them all for it. + ("person_id", YEAR + 7, 1), + ], +) +def test_an_input_set_once_built_for_another_year_round_trips(variable, year, value): # Not only the rebalancing: any input set for a year the data does not # give, before it, among the rebalancing's years or after them. simulation = Simulation(dataset=frame(ONE_PERSON)) - simulation.set_input("hours_worked", year, np.array([1_000])) + simulation.set_input(variable, year, np.array([value])) exported = simulation.to_input_dataframe() for identifier in IDENTIFIER_VARIABLES: assert exported[f"{identifier}__{year}"].tolist() == [1] reloaded = Simulation(dataset=exported) - assert reloaded.calculate("hours_worked", year).tolist() == [1_000] + assert reloaded.calculate(variable, year).tolist() == [value] pd.testing.assert_frame_equal(reloaded.to_input_dataframe(), exported) @@ -1330,6 +1347,7 @@ def multi_year_frame( identifier_years=st.sets(st.sampled_from(YEARS), min_size=1, max_size=2), data_years=st.sets(st.sampled_from(YEARS), min_size=1, max_size=3), later_years=st.sets(st.sampled_from(YEARS), max_size=2), + later_identifier_years=st.sets(st.sampled_from(YEARS), max_size=1), ) @example( population=ONE_PERSON, @@ -1337,6 +1355,7 @@ def multi_year_frame( identifier_years={YEAR}, data_years={YEAR}, later_years=set(), + later_identifier_years=set(), ) @example( population=ONE_PERSON, @@ -1344,13 +1363,19 @@ def multi_year_frame( identifier_years={YEAR + 2}, data_years={YEAR - 1, YEAR + 3}, later_years={YEAR + 5}, + later_identifier_years={YEAR + 4}, ) def test_a_frame_over_several_years_exports_the_same_frame_again( - population, id_columns, identifier_years, data_years, later_years + population, + id_columns, + identifier_years, + data_years, + later_years, + later_identifier_years, ): # 5. Round trip, for frames that name the identifiers and give values - # for different years, and simulations given an input for other years - # once built. + # for different years, and simulations given an input, or a person's + # identifier, for other years once built. simulation = Simulation( dataset=multi_year_frame(population, id_columns, identifier_years, data_years) ) @@ -1358,20 +1383,21 @@ def test_a_frame_over_several_years_exports_the_same_frame_again( simulation.set_input( "hours_worked", year, population.values["hours_worked"] + year ) + for year in sorted(later_identifier_years): + simulation.set_input("person_id", year, population.person_ids) exported = simulation.to_input_dataframe() reloaded = Simulation(dataset=exported) pd.testing.assert_frame_equal(reloaded.to_input_dataframe(), exported) - # The export gives the exact identifiers in every year it gives values - # for, and in each year the frame named them for. + # The export gives the exact identifiers in every whole year it names. years = { int(period) - for variable, period in map(split_column, exported.columns) - if variable not in IDENTIFIER_VARIABLES and period.isdigit() + for _, period in map(split_column, exported.columns) + if period.isdigit() } - assert years >= set(data_years) | set(later_years) - for year in years | set(identifier_years): + assert years >= set(data_years) | set(later_years) | set(identifier_years) + for year in years: assert exported[f"person_id__{year}"].tolist() == population.person_ids.tolist() for group in GROUPS: exact = population.groups[group].tolist() From a9664455e12e415ab96dd5cce7f0fe03f14a75a1 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 14:02:52 -0400 Subject: [PATCH 3/3] Reflow the with_identifiers_in_each_year docstring Co-Authored-By: Claude Opus 5.5 --- policyengine_uk/data/input_dataframe.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/policyengine_uk/data/input_dataframe.py b/policyengine_uk/data/input_dataframe.py index 220aad79dc..6ff5c077a4 100644 --- a/policyengine_uk/data/input_dataframe.py +++ b/policyengine_uk/data/input_dataframe.py @@ -458,10 +458,9 @@ def with_identifiers_in_each_year( to 2029, holds its identifiers in that year only by carrying them forward from an earlier one; and in a year before its data, where inputs do not carry back, its identifiers are the default, zero. policyengine-core - exports an - input only for the periods it was set for. Giving every identifier for - each year makes the frame load back to a simulation that exports the - same frame. + exports an input only for the periods it was set for. Giving every + identifier for each year makes the frame load back to a simulation that + exports the same frame. A simulation holds one set of people and groups, so every column of an identifier names the same ones (the loader refuses a frame whose columns