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..69bb07a587 --- /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 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..7b54a694da 100644 --- a/docs/book/usage/simulations.md +++ b/docs/book/usage/simulations.md @@ -200,6 +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 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 fcc6b133d1..6ff5c077a4 100644 --- a/policyengine_uk/data/input_dataframe.py +++ b/policyengine_uk/data/input_dataframe.py @@ -15,6 +15,12 @@ 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 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 @@ -267,6 +273,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 +379,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 +438,79 @@ 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 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, which is + where the default export (sorted by period) puts an input set for that + year. + + 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() + 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 + 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..26bbbe18f7 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 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: + 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..6ba665b8fa 100644 --- a/policyengine_uk/tests/test_dataframe_dataset.py +++ b/policyengine_uk/tests/test_dataframe_dataset.py @@ -30,6 +30,13 @@ 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 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 ``__`` @@ -67,7 +74,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 +84,7 @@ IDENTIFIER_VARIABLES, dataset_from_input_dataframe, split_column, + with_identifiers_in_each_year, ) from policyengine_uk.scenarios import no_economic_assumptions @@ -601,6 +609,64 @@ 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, 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], + "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__2028", + "benunit_id__2029", + "person_id__2024", + "person_id__2025", + "person_id__2027", + "person_id__2028", + "person_id__2029", + *frame_.columns[3:], + ] + 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"{identifier}__{year}"], + frame_[f"{identifier}__{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 +1206,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 +1254,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 +1277,133 @@ 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( + "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(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(variable, year).tolist() == [value] + 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), + later_identifier_years=st.sets(st.sampled_from(YEARS), max_size=1), +) +@example( + population=ONE_PERSON, + id_columns="both", + identifier_years={YEAR}, + data_years={YEAR}, + later_years=set(), + later_identifier_years=set(), +) +@example( + population=ONE_PERSON, + id_columns="own", + 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, + later_identifier_years, +): + # 5. Round trip, for frames that name the identifiers and give values + # 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) + ) + for year in sorted(later_years): + 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 whole year it names. + years = { + int(period) + for _, period in map(split_column, exported.columns) + if period.isdigit() + } + 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() + 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(