diff --git a/changelog.d/2417.fixed.md b/changelog.d/2417.fixed.md new file mode 100644 index 0000000000..beff819bf6 --- /dev/null +++ b/changelog.d/2417.fixed.md @@ -0,0 +1 @@ +- Fixed `Simulation.dataset` retaining projected amounts under `no_economic_assumptions` after the simulation loaded frozen inputs. diff --git a/docs/book/usage/simulations.md b/docs/book/usage/simulations.md index 7b54a694da..71ffe5b26f 100644 --- a/docs/book/usage/simulations.md +++ b/docs/book/usage/simulations.md @@ -294,6 +294,10 @@ sim = Simulation(dataset=dataset) print(sim.calculate("household_net_income", 2026)) ``` +With `scenario=no_economic_assumptions`, `sim.dataset[year]` holds the reset +monetary inputs used by the simulation. Helpers such as `filter_dataset` read +those frozen amounts while retaining each year's demographic inputs. + `Simulation` and `Microsimulation` can also load H5 files from local paths, Hugging Face URLs, or Google Cloud Storage URLs: diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py index 3f769e8950..9a495e84e1 100644 --- a/policyengine_uk/simulation.py +++ b/policyengine_uk/simulation.py @@ -792,7 +792,6 @@ def build_from_dataset_source( multi_year_dataset = _url_dataset_cache[cache_key] self.build_from_multi_year_dataset(multi_year_dataset) self._restore_cgt_extension_provenance(multi_year_dataset) - self.dataset = multi_year_dataset return dataset_file = materialize_gcs_dataset_url(dataset_source) self.build_from_file(dataset_file, cache_key=dataset_source) @@ -833,7 +832,6 @@ def build_from_file( self.build_from_multi_year_dataset(multi_year_dataset) self._restore_cgt_extension_provenance(multi_year_dataset) - self.dataset = multi_year_dataset def build_from_url(self, url: str) -> None: """Build simulation from a HuggingFace dataset URL. @@ -854,7 +852,6 @@ def build_from_url(self, url: str) -> None: multi_year_dataset = _url_dataset_cache[cache_key] self.build_from_multi_year_dataset(multi_year_dataset) self._restore_cgt_extension_provenance(multi_year_dataset) - self.dataset = multi_year_dataset return # Parse HuggingFace URL components @@ -977,7 +974,6 @@ def get_first_array(variable_name: str) -> np.ndarray: self.build_from_multi_year_dataset(multi_year_dataset) self._restore_cgt_extension_provenance(multi_year_dataset) - self.dataset = multi_year_dataset def build_from_single_year_dataset(self, dataset: UKSingleYearDataset) -> None: """Build simulation from a single-year UK dataset. @@ -990,7 +986,6 @@ def build_from_single_year_dataset(self, dataset: UKSingleYearDataset) -> None: _pre_encode_enum_columns(dataset, self.tax_benefit_system) self.build_from_multi_year_dataset(dataset) self._restore_cgt_extension_provenance(dataset) - self.dataset = dataset def build_from_multi_year_dataset(self, dataset: UKMultiYearDataset) -> None: """Build simulation from a multi-year UK dataset. diff --git a/policyengine_uk/tests/test_simulation_dataset_frozen.py b/policyengine_uk/tests/test_simulation_dataset_frozen.py new file mode 100644 index 0000000000..ef9f2bd28b --- /dev/null +++ b/policyengine_uk/tests/test_simulation_dataset_frozen.py @@ -0,0 +1,311 @@ +"""Stored MOCK inputs must describe the same frozen data as calculations.""" + +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +import yaml +from hypothesis import HealthCheck, example, given, settings +from hypothesis import strategies as st +from policyengine_core.data import Dataset + +import policyengine_uk.simulation as simulation_module +from policyengine_uk import Simulation +from policyengine_uk.data import UKSingleYearDataset +from policyengine_uk.data import economic_assumptions +from policyengine_uk.data.economic_assumptions import ( + DATASET_UPRATING_EXCLUSIONS, + extend_single_year_dataset, +) +from policyengine_uk.data.filter_dataset import filter_dataset +from policyengine_uk.scenarios import no_economic_assumptions +from policyengine_uk.system import system + + +pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system") + +BASE_YEAR = 2025 +END_YEAR = 2028 +UPRATING_INDICES = yaml.safe_load( + (Path(economic_assumptions.__file__).parent / "uprating_indices.yaml").read_text() +) +INDEXED_INPUTS = { + name + for names in UPRATING_INDICES.values() + for name in names + if name in system.variables +} +FALLBACK_INPUTS = { + name + for name, variable in system.variables.items() + if variable.uprating is not None + and variable.value_type in (int, float) + and name not in INDEXED_INPUTS + and name not in DATASET_UPRATING_EXCLUSIONS +} +UPRATED_INPUTS = INDEXED_INPUTS | FALLBACK_INPUTS | {"rent", "council_tax"} +# Simulation transfers these survey observations to their reported or +# pre-response input names after the loader has installed the arrays. +LOADED_INPUT_NAMES = { + "employment_income": "employment_income_before_lsr", + "capital_gains": "capital_gains_before_response", + "employee_pension_contributions": "employee_pension_contributions_reported", + "bus_fare_spending": "bus_fare_spending_reported", +} +GROWTH_PATHS = ( + "gov.economic_assumptions.yoy_growth.obr.average_earnings", + "gov.economic_assumptions.yoy_growth.obr.consumer_price_index", + "gov.economic_assumptions.yoy_growth.obr.per_capita.gdp", + "gov.economic_assumptions.yoy_growth.obr.per_capita.mixed_income", + "gov.economic_assumptions.yoy_growth.ons.population", + "gov.economic_assumptions.yoy_growth.obr.council_tax.england", + "gov.economic_assumptions.yoy_growth.ons.private_rental_prices.LONDON", +) + + +@pytest.fixture +def single_year(): + """One synthetic household with every model-defined uprated input.""" + single = UKSingleYearDataset( + person=pd.DataFrame( + { + "person_id": [1], + "person_benunit_id": [1], + "person_household_id": [1], + "age": [40], + "employment_income": [30_000.0], + "self_employment_income": [2_000.0], + "savings_interest_income": [100.0], + "capital_gains_before_response": [1_000.0], + "childcare_expenses": [500.0], + # filter_dataset can carry these without calculating taxes + # and benefits just to impute population-wide values. + "months_since_last_birthday": [3], + "attends_private_school": [False], + } + ), + benunit=pd.DataFrame( + {"benunit_id": [1], "benunit_reported_capital": [5_000.0]} + ), + household=pd.DataFrame( + { + "household_id": [1], + "household_weight": [1_000.0], + "region": ["LONDON"], + "tenure_type": ["RENT_PRIVATELY"], + "rent": [12_000.0], + "council_tax": [1_500.0], + "mortgage_interest_repayment": [200.0], + "food_and_non_alcoholic_beverages_consumption": [2_000.0], + } + ), + fiscal_year=BASE_YEAR, + ) + # Derive coverage from the public index mapping and variable metadata, + # including inputs whose own index only applies after the legacy horizon. + tables = { + name: getattr(single, name).to_dict(orient="list") + for name in single.table_names + } + for name in sorted(UPRATED_INPUTS): + variable = system.variables[name] + amount = 1_000 if variable.value_type is int else 1_000.0 + tables[variable.entity.key].setdefault(name, [amount]) + for donor, target in LOADED_INPUT_NAMES.items(): + entity = system.variables[donor].entity.key + tables[entity][target] = tables[entity][donor].copy() + return UKSingleYearDataset( + **{name: pd.DataFrame(values) for name, values in tables.items()}, + fiscal_year=BASE_YEAR, + ) + + +def _project(single_year, growth_changes=None): + parameters = simulation_module.CountryTaxBenefitSystem().parameters.clone() + for path, growth in (growth_changes or {}).items(): + parameters.get_child(path).update(period="2026", value=growth) + projected = extend_single_year_dataset(single_year, parameters, end_year=END_YEAR) + # Annual demographic inputs are authoritative independently of economic + # assumptions. Give the fixture changing ages so the reset must retain them. + for year in projected.years: + projected[year].person = projected[year].person.assign( + age=single_year.person.age + year - BASE_YEAR + ) + projected._cgt_extension_provenance = { + "capital_gains_before_response": { + str(year): { + "status": "projected", + "method": "dataset_extension", + "source_period": str(year - 1), + "projection_fingerprint": None, + } + for year in projected.years[1:] + } + } + return projected + + +def _build_frozen(monkeypatch, tmp_path, single_year, projected, route): + restored_from = [] + restore = Simulation._restore_cgt_extension_provenance + + def capture_restore(self, dataset): + restored_from.append(dataset) + restore(self, dataset) + + monkeypatch.setattr( + Simulation, "_restore_cgt_extension_provenance", capture_restore + ) + if route in ("hf-cache", "gs-cache"): + scheme = "hf" if route == "hf-cache" else "gs" + source = f"{scheme}://mock/frozen-inputs/dataset.h5" + monkeypatch.setattr( + simulation_module, "_url_dataset_cache", {source: projected} + ) + # #2405 prevents ordinary frozen simulations from reusing this cache. + # Force the historical route to test the caller's ownership of the + # real multi-year loader's reset copy independently of cache policy. + monkeypatch.setattr(Simulation, "_dataset_cache_key", lambda self, url: url) + + def unexpected_download(*args, **kwargs): + pytest.fail("The MOCK URL must load from the seeded in-memory cache") + + monkeypatch.setattr( + simulation_module, "download_huggingface_dataset", unexpected_download + ) + monkeypatch.setattr( + simulation_module, "materialize_gcs_dataset_url", unexpected_download + ) + else: + # The scenario normally freezes growth before extension. Supplying a + # genuinely projected extension exercises reset in these callers too; + # the multi-year builder and reset_growthfactor_uprating are real. + monkeypatch.setattr( + Simulation, + "_extend_dataset_with_cgt_provenance", + lambda self, dataset: projected, + ) + if route == "single-year": + source = single_year + elif route == "file": + source = tmp_path / "single-year.h5" + single_year.save(str(source)) + else: + assert route == "core-dataset" + dataset_path = tmp_path / "core-dataset.h5" + + class MockCoreDataset(Dataset): + name = "mock-frozen-inputs" + label = "MOCK frozen input regression" + data_format = Dataset.TIME_PERIOD_ARRAYS + file_path = dataset_path + + source = MockCoreDataset() + source.save_dataset( + { + column: {str(BASE_YEAR): table[column].to_numpy()} + for table in single_year.tables + for column in table.columns + if table[column].dtype.kind in "biuf" + } + ) + + sim = Simulation(dataset=source, scenario=no_economic_assumptions) + # The scenario also builds a current-law baseline through this route. + assert restored_from and all(original is projected for original in restored_from) + return sim + + +def _assert_inputs_match(sim, base): + supplied = {name for table in base.tables for name in table.columns} + assert UPRATED_INPUTS <= supplied + assert sim.dataset.years == list(range(BASE_YEAR, END_YEAR + 1)) + for year in sim.dataset.years: + for table_name in base.table_names: + table = getattr(sim.dataset[year], table_name) + base_table = getattr(base, table_name) + for variable in sorted(set(base_table.columns) & UPRATED_INPUTS): + calculated_variable = LOADED_INPUT_NAMES.get(variable, variable) + calculated = np.asarray(sim.calculate(calculated_variable, year)) + np.testing.assert_allclose( + table[variable].to_numpy(), + calculated, + rtol=1e-6, + err_msg=f"{variable} in {year}: stored data differs from loaded inputs", + ) + np.testing.assert_allclose( + table[variable].to_numpy(), + base_table[variable].to_numpy(), + rtol=1e-6, + err_msg=f"{variable} in {year}: economic growth was not reset", + ) + + +@pytest.mark.parametrize( + "route", ["single-year", "file", "hf-cache", "gs-cache", "core-dataset"] +) +def test_frozen_dataset_matches_loaded_inputs( + monkeypatch, tmp_path, single_year, route +): + projected = _project(single_year) + projected_income = projected[END_YEAR].person.employment_income.copy() + assert projected_income.iloc[0] != single_year.person.employment_income.iloc[0] + + sim = _build_frozen(monkeypatch, tmp_path, single_year, projected, route) + + _assert_inputs_match(sim, single_year) + assert sim.dataset is not projected + pd.testing.assert_series_equal( + projected[END_YEAR].person.employment_income, projected_income + ) + # Provenance is restored from the original even though the reset copy + # deliberately does not carry its private extension annotation. + record = sim.capital_gains_projection_metadata["application"]["by_variable"][ + "capital_gains_before_response" + ][str(END_YEAR)] + assert record["status"] == "unknown" + assert record["projection_fingerprint"] is None + + +@settings( + max_examples=8, + deadline=None, + derandomize=True, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +@example(growth_changes={GROWTH_PATHS[0]: 0.2}) +@given( + growth_changes=st.dictionaries( + st.sampled_from(GROWTH_PATHS), + st.floats(min_value=-0.4, max_value=0.6, allow_nan=False, allow_infinity=False), + min_size=1, + max_size=len(GROWTH_PATHS), + ) +) +def test_frozen_dataset_matches_loaded_inputs_for_random_growth( + monkeypatch, tmp_path, single_year, growth_changes +): + projected = _project(single_year, growth_changes) + with monkeypatch.context() as patch: + sim = _build_frozen(patch, tmp_path, single_year, projected, "hf-cache") + _assert_inputs_match(sim, single_year) + + +def test_filter_dataset_uses_frozen_inputs(monkeypatch, tmp_path, single_year): + projected = _project(single_year) + sim = _build_frozen(monkeypatch, tmp_path, single_year, projected, "hf-cache") + + extracted = filter_dataset(sim, household_id=1, year=END_YEAR) + + assert extracted.time_period == str(END_YEAR) + np.testing.assert_array_equal( + extracted.person.age, single_year.person.age + END_YEAR - BASE_YEAR + ) + for table_name in single_year.table_names: + table = getattr(extracted, table_name) + base_table = getattr(single_year, table_name) + for variable in base_table.select_dtypes(include="number"): + if variable == "age": + continue + np.testing.assert_allclose(table[variable], base_table[variable], rtol=1e-6) diff --git a/policyengine_uk/tests/test_simulation_dataset_sources.py b/policyengine_uk/tests/test_simulation_dataset_sources.py index 278d0bbc4d..4569b4f127 100644 --- a/policyengine_uk/tests/test_simulation_dataset_sources.py +++ b/policyengine_uk/tests/test_simulation_dataset_sources.py @@ -110,6 +110,7 @@ def fake_build_from_file(dataset_file, *, cache_key=None): def fake_build_from_multi_year_dataset(dataset): captured["dataset"] = dataset + simulation.dataset = dataset simulation_module._url_dataset_cache.pop(url, None) simulation_module._url_dataset_cache[url] = cached_dataset @@ -161,6 +162,7 @@ def __init__(self, file_path): def fake_build_from_multi_year_dataset(dataset): built["dataset"] = dataset + simulation.dataset = dataset simulation_module._url_dataset_cache.pop(cache_key, None) monkeypatch.setattr(simulation_module, "UKMultiYearDataset", FakeMultiYearDataset)