From d084bfad9e074d9c6eede65e43a7f4ba2264b7b0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 06:02:18 -0400 Subject: [PATCH 1/6] Keep each part's timing in a sum of scenarios Scenario.__add__ dropped applied_before_data_load, so every sum was applied after the data load: no_economic_assumptions + reform froze the parameters after the data had been uprated with growth. A sum now keeps its parts when they are applied at different times, and Simulation applies the parts applied before data load before it builds the data and the rest after (Scenario.parts_by_data_load). The sum's fields still hold the whole sum. A sum of scenarios that are all applied before data load is applied before it, and a scenario that changes nothing has no timing of its own. Fixes #2406. Co-Authored-By: Claude Opus 5.5 --- policyengine_uk/simulation.py | 28 +++--- policyengine_uk/utils/scenario.py | 145 +++++++++++++++++++++++++++--- 2 files changed, 151 insertions(+), 22 deletions(-) diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py index 494463d9c..1f278a909 100644 --- a/policyengine_uk/simulation.py +++ b/policyengine_uk/simulation.py @@ -210,13 +210,20 @@ def __init__( self.branches: Dict[str, Simulation] = {} + # The parts of the scenario applied before the data is built and + # after it. A scenario is one or the other; a sum of scenarios + # applied at different times keeps each part's timing. + before_data_load, after_data_load = ( + scenario.parts_by_data_load() if scenario is not None else (None, None) + ) + # A scenario's parameter changes rebuild the parameters from the # statutory tree, so they come before its modifier, whose edits a # rebuild would discard. - if scenario is not None and scenario.applied_before_data_load: - self.apply_scenario_parameters(scenario) - if scenario.simulation_modifier is not None: - scenario.simulation_modifier(self) + if before_data_load is not None: + self.apply_scenario_parameters(before_data_load) + if before_data_load.simulation_modifier is not None: + before_data_load.simulation_modifier(self) self._refresh_capital_gains_projection() # Build simulation from appropriate source if situation is not None: @@ -240,9 +247,10 @@ def __init__( # (a changed economic assumption does not re-uprate it), but come # before the steps below that read the parameters: the Universal # Credit rebalancing inputs and the structural reforms switched on by - # parameters. - if scenario is not None and not scenario.applied_before_data_load: - self.apply_scenario_parameters(scenario) + # parameters. They rebuild the parameters from those the part applied + # before data load left, so they keep that part's changes. + if after_data_load is not None: + self.apply_scenario_parameters(after_data_load) # Universal Credit reform (July 2025). Needs closer integration in the baseline, # but adding here for ease of toggling on/off via the 'active' parameter. @@ -295,9 +303,9 @@ def __init__( self.calculated_periods = [] - if scenario is not None and not scenario.applied_before_data_load: - if scenario.simulation_modifier is not None: - scenario.simulation_modifier(self) + if after_data_load is not None: + if after_data_load.simulation_modifier is not None: + after_data_load.simulation_modifier(self) self._refresh_capital_gains_projection() self._capture_cgt_supplied_inputs() self._cgt_warning_enabled = True diff --git a/policyengine_uk/utils/scenario.py b/policyengine_uk/utils/scenario.py index d51ab2d71..720fe96f4 100644 --- a/policyengine_uk/utils/scenario.py +++ b/policyengine_uk/utils/scenario.py @@ -1,7 +1,10 @@ +import copy +import functools import logging +import operator -from pydantic import BaseModel, ConfigDict -from typing import Any, Literal, List, Optional, Callable, Dict, Type, Union +from pydantic import BaseModel, ConfigDict, PrivateAttr +from typing import Any, Literal, List, Optional, Callable, Dict, Tuple, Type, Union from policyengine_core.simulations import Simulation from policyengine_core.reforms import Reform from policyengine_uk.utils.dated_reforms import ( @@ -106,15 +109,26 @@ def _apply_reform_class(reform: Type[Reform], simulation: Simulation) -> None: simulation.tax_benefit_system.reset_parameter_caches() +def _equal(value: Any, kept: Any) -> bool: + """Whether a field still equals the copy kept of it.""" + try: + return bool(value == kept) + except Exception: + # A value whose == gives no single truth value, such as an array. + return repr(value) == repr(kept) + + class Scenario(BaseModel): """Represents a scenario configuration for policy simulations. A scenario can include parameter changes and/or simulation modifications that are applied before running a simulation. Scenarios can be combined - using the + operator. + using the + operator, and a sum keeps each part's timing (see + ``parts_by_data_load``). """ applied_before_data_load: bool = False + """Whether ``Simulation`` applies the scenario before it builds its data, so the data's projected years follow its changes. A sum is applied before data load when every part is; a sum of parts applied at different times reads False here and applies each part at its own time.""" parameter_changes: Optional[ Dict[ @@ -137,9 +151,25 @@ class Scenario(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) + # A sum of parts applied at different times keeps the parts, each applied + # at one time, in the order added, for ``parts_by_data_load``; and its own + # fields as they were when it was made, so a sum whose fields are changed + # afterwards is applied as one scenario, as those fields say. + _parts: Optional[Tuple["Scenario", ...]] = PrivateAttr(default=None) + _fields_when_made: Optional[tuple] = PrivateAttr(default=None) + def __add__(self, other: "Scenario") -> "Scenario": """Combine two scenarios by merging parameter changes and chaining modifiers. + The fields of the sum hold both scenarios, as they always have: + ``Scenario.apply`` and code that reads the fields see the whole sum. + ``Simulation`` applies each part at its own time, before or after it + builds its data (see ``parts_by_data_load``), so adding a scenario + applied before data load, such as ``no_economic_assumptions``, to + another keeps it before data load. A scenario that changes nothing + has no timing of its own: ``Scenario() + scenario`` is applied when + ``scenario`` is. + Args: other: Another Scenario to combine with this one @@ -165,28 +195,111 @@ def __add__(self, other: "Scenario") -> "Scenario": # Simple override merged_params[key] = value - # Chain simulation modifiers + # Chain simulation modifiers, as they are when the scenarios are + # added, as the parts below keep them. combined_modifier = None + first, second = self.simulation_modifier, other.simulation_modifier - if self.simulation_modifier and other.simulation_modifier: + if first and second: # Both have modifiers - chain them def combined_modifier(simulation: Simulation) -> None: - self.simulation_modifier(simulation) - other.simulation_modifier(simulation) + first(simulation) + second(simulation) - elif self.simulation_modifier: - combined_modifier = self.simulation_modifier - elif other.simulation_modifier: - combined_modifier = other.simulation_modifier + elif first: + combined_modifier = first + elif second: + combined_modifier = second dated_reforms = (self.dated_reforms or []) + (other.dated_reforms or []) + parts = self._timed_parts() + other._timed_parts() + timings = {part.applied_before_data_load for part in parts} + # Return new scenario with merged configuration - return Scenario( + combined = Scenario( + applied_before_data_load=timings == {True}, parameter_changes=merged_params if merged_params else None, dated_reforms=dated_reforms or None, simulation_modifier=combined_modifier, ) + if len(timings) > 1: + combined._parts = parts + combined._fields_when_made = ( + combined.applied_before_data_load, + copy.deepcopy(combined.parameter_changes), + copy.deepcopy(combined.dated_reforms), + combined.simulation_modifier, + ) + return combined + + def _current_parts(self) -> Optional[Tuple["Scenario", ...]]: + """A sum's parts, unless its fields have changed since it was made.""" + if self._parts is None: + return None + flag, parameter_changes, dated_reforms, modifier = self._fields_when_made + if ( + self.applied_before_data_load == flag + and self.simulation_modifier is modifier + and _equal(self.parameter_changes, parameter_changes) + and _equal(self.dated_reforms, dated_reforms) + ): + return self._parts + return None + + def _timed_parts(self) -> Tuple["Scenario", ...]: + """The scenario as parts each applied at one time, in order: a sum's + parts, or a copy of the scenario (so later edits to it do not reach + a sum it is in), or none if it changes nothing.""" + parts = self._current_parts() + if parts is not None: + return parts + if not ( + self.parameter_changes or self.dated_reforms or self.simulation_modifier + ): + return () + part = self.model_copy() + part._parts = None + part._fields_when_made = None + return (part,) + + def parts_by_data_load( + self, + ) -> Tuple[Optional["Scenario"], Optional["Scenario"]]: + """The parts of the scenario applied before and after data load. + + ``Simulation`` applies the first part's parameter changes and then + its modifier before it builds its data, and the second part's + parameter changes just after the data is built and its modifier at + the end, as for a scenario on its own. A scenario is one or the + other, as ``applied_before_data_load`` says. A sum of parts applied + at different times is split into the sum of the parts applied + before data load and the sum of the rest, each in the order added. + A sum whose fields have been changed since it was made is applied + as one scenario, as its fields say. + + Where parts applied at different times change the same parameter, + the later change holds: the data is built with the value of the part + applied before data load, and the part applied after it then sets + the parameter. A part applied after data load rebuilds the + parameters from those the parts before it left, so it keeps their + parameter changes, and the edits of a modifier that reloads and + reprocesses the parameters (as ``no_economic_assumptions`` does), + but not a modifier's direct edits to the processed parameters. + + Returns: + The part applied before data load and the part applied after it, + each a Scenario, or None where there is no such part. + """ + parts = self._current_parts() + if parts is None: + return (self, None) if self.applied_before_data_load else (None, self) + before = [part for part in parts if part.applied_before_data_load] + after = [part for part in parts if not part.applied_before_data_load] + return ( + functools.reduce(operator.add, before), + functools.reduce(operator.add, after), + ) def dated_assignments(self, order_start: int = 0) -> List[ParameterAssignment]: """The statutory-timeline changes of ``dated_reforms``, in order.""" @@ -377,4 +490,12 @@ def __repr__(self) -> str: parts.append(f"dated_reforms={len(self.dated_reforms)} items") if self.simulation_modifier: parts.append("simulation_modifier=") + timed = self._current_parts() + if timed is not None: + before = sum(part.applied_before_data_load for part in timed) + parts.append( + f"parts={before} before data load, {len(timed) - before} after" + ) + elif self.applied_before_data_load: + parts.append("applied_before_data_load=True") return f"Scenario({', '.join(parts)})" From b3d0539d4fe197b2f94456863523ab0b389e3281 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 06:19:58 -0400 Subject: [PATCH 2/6] Split a sum by its bracketing, and test the timing of each part The merge of parameter changes is not associative (a value, then a dict of periods, then another dict), so a mixed sum keeps its two operands and splits through them, and each part's changes are those `+` gives with the same bracketing. Operands are kept as they were when added: their changes are copied, and their modifiers are the objects the sum's chained modifier calls. Tests (MOCK datasets): no_economic_assumptions + reform, in either order, freezes the data and the parameters as no_economic_assumptions alone does and applies the reform, on single-year and multi-year datasets and on a situation; a macro scenario applied before data load still reaches the data when a reform is added; a data-editing modifier applied after data load still sees the data; a simulation runs each part once at its own time (property); and the split partitions the parts in order, follows the bracketing, keeps Scenario() as an identity and leaves the sum's fields as the previous + made them (properties, with a copy of the previous + as reference). Docs and changelog. Co-Authored-By: Claude Opus 5.5 --- changelog.d/2406.fixed.md | 1 + docs/book/usage/scenarios.md | 13 + .../tests/test_scenario_sum_timing.py | 677 ++++++++++++++++++ policyengine_uk/utils/scenario.py | 115 +-- 4 files changed, 759 insertions(+), 47 deletions(-) create mode 100644 changelog.d/2406.fixed.md create mode 100644 policyengine_uk/tests/test_scenario_sum_timing.py diff --git a/changelog.d/2406.fixed.md b/changelog.d/2406.fixed.md new file mode 100644 index 000000000..1d818e5d2 --- /dev/null +++ b/changelog.d/2406.fixed.md @@ -0,0 +1 @@ +- A sum of scenarios keeps each part's timing. A part applied before data load, such as `no_economic_assumptions`, is still applied before it when added to a reform, so `no_economic_assumptions + reform` freezes the data as well as the parameters, and the sum of a macro scenario applied before data load and a reform projects the data with the macro growth. Parts applied after data load, including modifiers that edit the data, still run after it. The sum's fields still hold the whole sum. `Scenario.parts_by_data_load()` returns the two parts. diff --git a/docs/book/usage/scenarios.md b/docs/book/usage/scenarios.md index d50098a3f..7bef155a8 100644 --- a/docs/book/usage/scenarios.md +++ b/docs/book/usage/scenarios.md @@ -153,6 +153,8 @@ for name, scenario in [ print(f"{name}: £{net_income / 12:.2f} per month") ``` +A sum keeps each part's timing. A simulation applies the parts set with `applied_before_data_load=True` before it builds its data and the rest after, each in the order added. So `no_economic_assumptions + reform` freezes the data and the parameters and applies the reform, and a modifier that edits the data still runs once the data exists. A sum reads `applied_before_data_load=True` only when every part is applied before data load. `Scenario.parts_by_data_load()` returns the two parts. Where both change the same parameter, the part applied after data load sets it last. If you set `applied_before_data_load=True` on a sum yourself, or change its fields after making it, the whole sum is applied as one scenario at that time. + ## Defining structural scenarios Beyond simple parameter changes, you sometimes need scenarios that make structural modifications to how the tax-benefit system works. Let's walk through how to create these, using examples similar to those in PolicyEngine UK's scenarios module. @@ -314,6 +316,17 @@ print(float(baseline.calculate("employment_income", 2026)[0])) print(float(static.calculate("employment_income", 2026)[0])) ``` +Add it to a reform to measure the reform without growth. The sum keeps the freeze before the data load and applies the reform as usual: + +```python +from policyengine_uk.model_api import Scenario + +reform = Scenario.from_reform( + {"gov.hmrc.income_tax.allowances.personal_allowance.amount": {"2025.2099": 20_000}} +) +static_reform = Simulation(situation=situation, scenario=no_economic_assumptions + reform) +``` + ## Advanced scenario techniques ### Time-varying parameters diff --git a/policyengine_uk/tests/test_scenario_sum_timing.py b/policyengine_uk/tests/test_scenario_sum_timing.py new file mode 100644 index 000000000..b4582fe67 --- /dev/null +++ b/policyengine_uk/tests/test_scenario_sum_timing.py @@ -0,0 +1,677 @@ +"""A sum of scenarios applies each part at its own time. + +``Simulation`` applies a scenario before it builds its data +(``applied_before_data_load``) or after. A sum of scenarios applied at +different times, such as ``no_economic_assumptions + reform``, keeps its parts +(``Scenario.parts_by_data_load``), and the simulation applies the parts +applied before data load before it builds the data and the rest after. + +Invariants: + +1. Partition and order: every part's parameter changes and modifier are + applied exactly once, at the part's own time. A modifier applied before + data load sees no populations; one applied after sees them. Within each + time, parts are applied in the order added. +2. Associativity and identity: how a sum is bracketed does not change its + split, and ``Scenario()`` (which changes nothing) adds no part. +3. Backward compatibility: a sum's fields (parameter changes, dated reforms + and the chained modifier) equal those the previous ``Scenario.__add__`` + made, so ``Scenario.apply`` and code reading the fields see the whole sum. + A sum reads ``applied_before_data_load=True`` exactly when every part that + changes something is applied before data load. +4. Freeze: ``no_economic_assumptions + reform``, in either order, gives the + data and the economic assumption parameters of ``no_economic_assumptions`` + alone, and the reform's own changes. +5. A sum is fixed when made: later edits to its operands do not reach it, + and a sum whose own fields are changed is applied as one scenario, as + those fields say. + +All datasets here are MOCK data: synthetic households, not survey records. +""" + +import functools +import operator + +import numpy as np +import pandas as pd +import pytest +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st +from policyengine_core.parameters import Parameter + +from policyengine_uk import Simulation +from policyengine_uk.data.dataset_schema import ( + UKMultiYearDataset, + UKSingleYearDataset, +) +from policyengine_uk.scenarios import no_economic_assumptions +from policyengine_uk.utils.scenario import Scenario + +AVERAGE_EARNINGS = "gov.economic_assumptions.yoy_growth.obr.average_earnings" +PERSONAL_ALLOWANCE = "gov.hmrc.income_tax.allowances.personal_allowance.amount" +BASIC_RATE = "gov.hmrc.income_tax.rates.uk[0].rate" +REFORM = {PERSONAL_ALLOWANCE: {"2025.2099": 20_000}} +CURRENT_LAW_ALLOWANCE = 12_570 +CURRENT_LAW_GROWTH_2027 = 0.024 # yoy_growth.obr.average_earnings for 2027 +MACRO_GROWTH_2027 = 0.2 +# Calendar-year growth series are keyed to 1 January. +MACRO = Scenario( + parameter_changes={AVERAGE_EARNINGS: {"year:2027-01-01:1": MACRO_GROWTH_2027}}, + applied_before_data_load=True, +) +UPRATED_COLUMNS = {"person": ["employment_income"], "household": ["rent"]} +YEARS = range(2025, 2031) + +SITUATION = { + "people": {"person": {"age": {2025: 40}, "employment_income": {2025: 30_000}}}, + "benunits": {"benunit": {"members": ["person"]}}, + "households": {"household": {"members": ["person"]}}, +} + +PROPERTY_SETTINGS = settings( + max_examples=60, + deadline=None, + derandomize=True, + suppress_health_check=[HealthCheck.too_slow], +) +SIMULATION_PROPERTY_SETTINGS = settings( + max_examples=6, + deadline=None, + derandomize=True, + suppress_health_check=[HealthCheck.too_slow], +) + + +def mock_dataset() -> UKSingleYearDataset: + """MOCK data: two synthetic households with columns the uprating reads.""" + person = pd.DataFrame( + { + "person_id": [1, 2], + "person_benunit_id": [1, 2], + "person_household_id": [1, 2], + "age": [40, 70], + "employment_income": [30_000.0, 0.0], + } + ) + benunit = pd.DataFrame({"benunit_id": [1, 2]}) + household = pd.DataFrame( + { + "household_id": [1, 2], + "household_weight": [1.0, 1.0], + "region": ["LONDON", "WALES"], + "tenure_type": ["RENT_PRIVATELY", "OWNED_OUTRIGHT"], + "rent": [12_000.0, 0.0], + "council_tax": [1_500.0, 1_800.0], + } + ) + return UKSingleYearDataset( + person=person, benunit=benunit, household=household, fiscal_year=2025 + ) + + +def mock_multi_year_dataset() -> UKMultiYearDataset: + """MOCK data: one synthetic person, with 2026 already uprated.""" + person_2025 = pd.DataFrame( + { + "person_id": [1], + "person_benunit_id": [1], + "person_household_id": [1], + "employment_income": [30_000.0], + "age": [40], + } + ) + person_2026 = person_2025.copy() + person_2026["employment_income"] = [31_019.96875] + person_2026["age"] = [41] + benunit = pd.DataFrame({"benunit_id": [1]}) + household = pd.DataFrame({"household_id": [1]}) + return UKMultiYearDataset( + datasets=[ + UKSingleYearDataset( + person=person, + benunit=benunit.copy(), + household=household.copy(), + fiscal_year=year, + ) + for year, person in ((2025, person_2025), (2026, person_2026)) + ] + ) + + +def parameter(simulation, path, instant): + return simulation.tax_benefit_system.parameters.get_child(path)(instant) + + +def growth(simulation, year=2027): + """Employment income growth into ``year`` in the simulation's data.""" + data = simulation.dataset + return ( + data[year].person.employment_income.iloc[0] + / data[year - 1].person.employment_income.iloc[0] + - 1 + ) + + +def economic_assumptions(simulation) -> dict: + """Every growth series and statutory uprating input, over the years.""" + node = simulation.tax_benefit_system.parameters.gov.economic_assumptions + return { + p.name: [p(f"{year}-01-01") for year in YEARS] + for subtree in (node.yoy_growth, node.statutory_uprating_inputs) + for p in subtree.get_descendants() + if isinstance(p, Parameter) + } + + +def assert_same_data(simulation, expected): + assert list(simulation.dataset.years) == list(expected.dataset.years) + for year in expected.dataset.years: + for table in ("person", "benunit", "household"): + pd.testing.assert_frame_equal( + getattr(simulation.dataset[year], table), + getattr(expected.dataset[year], table), + obj=f"{table} {year}", + ) + + +# --- the freeze, on MOCK data --------------------------------------------- + + +@pytest.fixture(scope="module") +def frozen(): + return Simulation(dataset=mock_dataset(), scenario=no_economic_assumptions) + + +@pytest.fixture(scope="module", params=["freeze first", "reform first"]) +def frozen_with_reform(request): + reform = Scenario.from_reform(REFORM) + scenario = ( + no_economic_assumptions + reform + if request.param == "freeze first" + else reform + no_economic_assumptions + ) + return Simulation(dataset=mock_dataset(), scenario=scenario) + + +def test_no_economic_assumptions_plus_reform_freezes_the_data( + frozen_with_reform, frozen +): + simulation = frozen_with_reform + assert simulation.disable_economic_assumptions + # Every projected year holds the base year's uprated amounts. + base = simulation.dataset[2025] + for year in simulation.dataset.years: + for table, columns in UPRATED_COLUMNS.items(): + for column in columns: + np.testing.assert_array_equal( + getattr(simulation.dataset[year], table)[column], + getattr(base, table)[column], + err_msg=f"{table}.{column} {year}", + ) + assert_same_data(simulation, frozen) + + +def test_no_economic_assumptions_plus_reform_freezes_the_parameters( + frozen_with_reform, frozen +): + simulation = frozen_with_reform + assert parameter(simulation, AVERAGE_EARNINGS, "2027-01-01") == 0 + # The reform rebuilt the parameters after the freeze and kept it. + assert economic_assumptions(simulation) == economic_assumptions(frozen) + + +def test_no_economic_assumptions_plus_reform_still_applies_the_reform( + frozen_with_reform, frozen +): + simulation = frozen_with_reform + for year in YEARS: + assert parameter(simulation, PERSONAL_ALLOWANCE, str(year)) == 20_000 + assert parameter(frozen, PERSONAL_ALLOWANCE, str(year)) == CURRENT_LAW_ALLOWANCE + # Frozen income above the reformed allowance, at the basic rate: on main + # the income had grown by 2027, so this was £2,352.90. + income_tax = simulation.calculate("income_tax", 2027, map_to="person") + assert income_tax[0] == pytest.approx((30_000 - 20_000) * 0.2) + # The baseline is current law, with growth. + baseline = simulation.baseline + assert parameter(baseline, PERSONAL_ALLOWANCE, "2027") == CURRENT_LAW_ALLOWANCE + assert growth(baseline) == pytest.approx(CURRENT_LAW_GROWTH_2027) + + +@pytest.mark.parametrize("freeze_first", [True, False]) +def test_no_economic_assumptions_plus_reform_resets_a_multi_year_dataset( + freeze_first, +): + dataset = mock_multi_year_dataset() + reform = Scenario.from_reform(REFORM) + scenario = ( + no_economic_assumptions + reform + if freeze_first + else reform + no_economic_assumptions + ) + simulation = Simulation(dataset=dataset, scenario=scenario) + assert float(simulation.calculate("employment_income", 2026)[0]) == 30_000 + assert int(simulation.calculate("age", 2026)[0]) == 41 + assert parameter(simulation, PERSONAL_ALLOWANCE, "2026") == 20_000 + # The input dataset is left as it was. + assert dataset[2026].person["employment_income"].iloc[0] == 31_019.96875 + + +def test_no_economic_assumptions_plus_reform_freezes_a_situation(): + simulation = Simulation( + situation=SITUATION, + scenario=no_economic_assumptions + Scenario.from_reform(REFORM), + ) + # Before the labour supply response, which compares with the baseline. + income = [ + float(simulation.calculate("employment_income_before_lsr", year)[0]) + for year in YEARS + ] + assert income == [30_000.0] * len(income) + assert parameter(simulation, PERSONAL_ALLOWANCE, "2027") == 20_000 + + +# --- other parts applied before data load --------------------------------- + + +def test_before_load_parameter_changes_in_a_sum_reach_the_data(): + """The documented macro route keeps its timing when a reform is added.""" + macro = Simulation(dataset=mock_dataset(), scenario=MACRO) + for scenario in ( + MACRO + Scenario.from_reform(REFORM), + Scenario.from_reform(REFORM) + MACRO, + ): + simulation = Simulation(dataset=mock_dataset(), scenario=scenario) + # On main the sum was applied after data load: 2.4%. + assert growth(simulation) == pytest.approx(MACRO_GROWTH_2027) + assert_same_data(simulation, macro) + assert parameter(simulation, PERSONAL_ALLOWANCE, "2027") == 20_000 + + +def test_a_sum_of_scenarios_applied_before_data_load_is_applied_before_it(): + seen = [] + + def record(simulation): + seen.append(getattr(simulation, "populations", None) is None) + + scenario = MACRO + Scenario( + simulation_modifier=record, applied_before_data_load=True + ) + assert scenario.applied_before_data_load + assert scenario.parts_by_data_load() == (scenario, None) + simulation = Simulation(dataset=mock_dataset(), scenario=scenario) + assert seen == [True] + assert growth(simulation) == pytest.approx(MACRO_GROWTH_2027) + + +# --- a modifier applied after data load still sees the data --------------- + + +def give_pensions(simulation): + """A data-editing modifier: it reads the populations, so it can only run + once the data is built.""" + income = simulation.calculate("employment_income_before_lsr", 2025) + simulation.set_input("private_pension_income", 2027, income / 10) + + +@pytest.mark.parametrize("freeze_first", [True, False]) +def test_an_after_load_data_editing_modifier_in_a_sum_sees_the_data(freeze_first): + edit = Scenario(simulation_modifier=give_pensions) + scenario = ( + no_economic_assumptions + edit + if freeze_first + else edit + no_economic_assumptions + ) + before, after = scenario.parts_by_data_load() + assert before.simulation_modifier is not None + assert after.simulation_modifier is give_pensions + simulation = Simulation(dataset=mock_dataset(), scenario=scenario) + pensions = simulation.calculate("private_pension_income", 2027, map_to="person") + np.testing.assert_array_equal(pensions, [3_000, 0]) + assert growth(simulation) == 0 + assert simulation.disable_economic_assumptions + + +def test_flagging_the_whole_sum_would_run_the_data_edit_without_data(): + """Why a sum keeps each part's timing rather than one flag: the edit + needs populations that do not exist before data load.""" + scenario = no_economic_assumptions + Scenario(simulation_modifier=give_pensions) + scenario.applied_before_data_load = True + with pytest.raises(AttributeError, match="populations"): + Simulation(dataset=mock_dataset(), scenario=scenario) + + +# --- the split, in a real simulation -------------------------------------- + + +@st.composite +def timed_modifier_parts(draw): + """Parts that each record, when their modifier runs, whether the data had + been built.""" + timings = draw(st.lists(st.booleans(), min_size=1, max_size=5)) + events = [] + parts = [] + for index, before in enumerate(timings): + + def record(simulation, index=index): + events.append((index, getattr(simulation, "populations", None) is not None)) + + parts.append( + Scenario(simulation_modifier=record, applied_before_data_load=before) + ) + return timings, parts, events + + +@SIMULATION_PROPERTY_SETTINGS +@given(timed_modifier_parts()) +def test_a_simulation_runs_each_part_once_at_its_own_time(drawn): + timings, parts, events = drawn + Simulation(situation=SITUATION, scenario=functools.reduce(operator.add, parts)) + before = [i for i, b in enumerate(timings) if b] + after = [i for i, b in enumerate(timings) if not b] + assert events == [(i, False) for i in before] + [(i, True) for i in after] + + +def test_parts_applied_later_set_a_parameter_both_change(): + before = Scenario( + parameter_changes={PERSONAL_ALLOWANCE: {"2025": 15_000}}, + applied_before_data_load=True, + ) + after = Scenario.from_reform({PERSONAL_ALLOWANCE: {"2025.2099": 20_000}}) + for scenario in (before + after, after + before): + simulation = Simulation(situation=SITUATION, scenario=scenario) + assert parameter(simulation, PERSONAL_ALLOWANCE, "2025") == 20_000 + + +def test_scenario_apply_runs_every_part_of_a_sum_on_a_built_simulation(): + calls = [] + scenario = ( + Scenario( + simulation_modifier=lambda s: calls.append("before"), + applied_before_data_load=True, + ) + + Scenario.from_reform({BASIC_RATE: 0.21}) + + Scenario(simulation_modifier=lambda s: calls.append("after")) + ) + simulation = Simulation(situation=SITUATION) + scenario.apply(simulation) + assert calls == ["before", "after"] + assert parameter(simulation, BASIC_RATE, "2027") == 0.21 + + +# --- a sum is fixed when made --------------------------------------------- + + +def recorder(log, name): + return lambda simulation: log.append(name) + + +def test_edits_to_an_operand_after_adding_do_not_reach_the_sum(): + log = [] + first = Scenario(simulation_modifier=recorder(log, "first")) + second = Scenario( + simulation_modifier=recorder(log, "second"), applied_before_data_load=True + ) + scenario = first + second + first.applied_before_data_load = True + first.simulation_modifier = recorder(log, "replaced") + before, after = scenario.parts_by_data_load() + before.simulation_modifier(None) + after.simulation_modifier(None) + scenario.simulation_modifier(None) + assert log == ["second", "first", "first", "second"] + + +def test_a_sum_flagged_before_data_load_is_applied_before_it_whole(): + scenario = Scenario.from_reform(REFORM) + no_economic_assumptions + assert scenario.parts_by_data_load()[0] is not scenario + scenario.applied_before_data_load = True + assert scenario.parts_by_data_load() == (scenario, None) + copied = (Scenario.from_reform(REFORM) + no_economic_assumptions).model_copy( + update={"applied_before_data_load": True} + ) + assert copied.parts_by_data_load() == (copied, None) + + +def test_a_sum_whose_fields_change_in_place_is_applied_as_they_say(): + scenario = no_economic_assumptions + Scenario( + parameter_changes={PERSONAL_ALLOWANCE: 20_000} + ) + before, after = scenario.parts_by_data_load() + assert before is not None and after is not None + scenario.parameter_changes[BASIC_RATE] = 0.21 + assert scenario.parts_by_data_load() == (None, scenario) + + +def test_a_mixed_sum_survives_copying(): + scenario = no_economic_assumptions + Scenario.from_reform(REFORM) + for copied in (scenario.model_copy(), scenario.model_copy(deep=True)): + before, after = copied.parts_by_data_load() + assert before.applied_before_data_load + assert not after.applied_before_data_load + + +# --- properties of the split, without a simulation ------------------------ + +PATHS = ["a.x", "a.y", "b.z"] + + +def reference_add(self: Scenario, other: Scenario) -> Scenario: + """The fields ``Scenario.__add__`` made before parts kept their timing.""" + merged_params = {} + if self.parameter_changes: + merged_params.update(self.parameter_changes) + if other.parameter_changes: + for key, value in other.parameter_changes.items(): + if ( + key in merged_params + and isinstance(merged_params[key], dict) + and isinstance(value, dict) + ): + merged_params[key] = {**merged_params[key], **value} + else: + merged_params[key] = value + combined_modifier = None + if self.simulation_modifier and other.simulation_modifier: + + def combined_modifier(simulation): + self.simulation_modifier(simulation) + other.simulation_modifier(simulation) + + elif self.simulation_modifier: + combined_modifier = self.simulation_modifier + elif other.simulation_modifier: + combined_modifier = other.simulation_modifier + dated_reforms = (self.dated_reforms or []) + (other.dated_reforms or []) + return Scenario( + parameter_changes=merged_params if merged_params else None, + dated_reforms=dated_reforms or None, + simulation_modifier=combined_modifier, + ) + + +changes = st.one_of( + st.none(), + st.dictionaries( + st.sampled_from(PATHS), + st.one_of( + st.integers(0, 9), + st.dictionaries(st.sampled_from(["2025", "2026"]), st.integers(0, 9)), + ), + max_size=2, + ), +) + + +@st.composite +def leaves(draw): + """Scenarios that each record their label when their modifier runs, with + some that change nothing.""" + specs = draw( + st.lists( + st.tuples(st.booleans(), st.booleans(), changes, st.booleans()), + min_size=1, + max_size=6, + ) + ) + log = [] + scenarios = [] + for index, (before, has_modifier, parameter_changes, dated) in enumerate(specs): + scenarios.append( + Scenario( + applied_before_data_load=before, + parameter_changes=parameter_changes, + dated_reforms=( + [{"reading": "reform_dict", "changes": {f"c.{index}": index}}] + if dated + else None + ), + simulation_modifier=recorder(log, index) if has_modifier else None, + ) + ) + return scenarios, log + + +def changes_something(scenario): + return bool( + scenario.parameter_changes + or scenario.dated_reforms + or scenario.simulation_modifier + ) + + +def draw_tree(draw, low, high): + """A random bracketing of the leaves numbered low to high - 1.""" + if high - low == 1: + return low + cut = draw(st.integers(low + 1, high - 1)) + return (draw_tree(draw, low, cut), draw_tree(draw, cut, high)) + + +def evaluate(tree, leaf, add): + """Add up a bracketing, skipping leaves that ``leaf`` gives as None.""" + if isinstance(tree, int): + return leaf(tree) + first, second = (evaluate(branch, leaf, add) for branch in tree) + if first is None or second is None: + return second if first is None else first + return add(first, second) + + +def run(log, scenario): + """The labels a scenario's modifier records.""" + del log[:] + if scenario is not None and scenario.simulation_modifier is not None: + scenario.simulation_modifier(None) + return list(log) + + +def fields(scenario): + """Parameter changes and dated reforms; empty and None alike mean none.""" + if scenario is None: + return None, None + return scenario.parameter_changes or None, scenario.dated_reforms or None + + +def merges_associatively(scenarios): + """Whether no parameter is changed by a value in one scenario and by + a dict of periods in another: the merge of parameter changes is then + associative.""" + kinds = {} + for scenario in scenarios: + for path, value in (scenario.parameter_changes or {}).items(): + kinds.setdefault(path, set()).add(isinstance(value, dict)) + return all(len(kind) == 1 for kind in kinds.values()) + + +@PROPERTY_SETTINGS +@given(leaves(), st.data()) +def test_a_sum_splits_into_its_parts_by_timing(drawn, data): + scenarios, log = drawn + count = len(scenarios) + trees = [ + functools.reduce(lambda tree, i: (tree, i), range(1, count), 0), + functools.reduce( + lambda tree, i: (i, tree), range(count - 2, -1, -1), count - 1 + ), + draw_tree(data.draw, 0, count), + ] + active = [changes_something(s) for s in scenarios] + has_before = any( + a and s.applied_before_data_load for a, s in zip(active, scenarios) + ) + has_after = any( + a and not s.applied_before_data_load for a, s in zip(active, scenarios) + ) + if count == 1: + # Not a sum: its own flag says when it is applied. + has_before = scenarios[0].applied_before_data_load + has_after = not has_before + + def phase(tree, before): + """The previous ``+`` over one phase's parts, bracketed as the sum.""" + return evaluate( + tree, + lambda i: ( + scenarios[i] + if active[i] and scenarios[i].applied_before_data_load == before + else None + ), + reference_add, + ) + + splits = [] + for tree in trees: + total = evaluate(tree, lambda i: scenarios[i], operator.add) + before, after = total.parts_by_data_load() + splits.append((before, after)) + # 1. Partition and order: each phase holds exactly its parts' + # changes, bracketed as the sum, and runs their modifiers in the + # order added. + assert fields(before) == fields(phase(tree, True)) + assert fields(after) == fields(phase(tree, False)) + assert run(log, before) == run(log, phase(tree, True)) + assert run(log, after) == run(log, phase(tree, False)) + assert (before is None) == (not has_before) + assert (after is None) == (has_before and not has_after) + assert before is None or before.applied_before_data_load + assert after is None or not after.applied_before_data_load + # 3. Backward compatibility: the fields hold the whole sum, as the + # previous ``+`` made them. + reference = evaluate(tree, lambda i: scenarios[i], reference_add) + assert fields(total) == fields(reference) + assert run(log, total) == run(log, reference) + assert total.applied_before_data_load == (has_before and not has_after) + # 2. Identity: Scenario() adds no part. + for padded in (Scenario() + total, total + Scenario()): + pad_before, pad_after = padded.parts_by_data_load() + assert fields(pad_before) == fields(before) + assert fields(pad_after) == fields(after) + assert run(log, pad_before) == run(log, before) + assert run(log, pad_after) == run(log, after) + # 2. Associativity: modifiers run in the order added however the sum is + # bracketed, and the parts agree wherever the merge itself is + # associative. + for before, after in splits[1:]: + assert run(log, before) == run(log, splits[0][0]) + assert run(log, after) == run(log, splits[0][1]) + if merges_associatively(scenarios): + assert fields(before) == fields(splits[0][0]) + assert fields(after) == fields(splits[0][1]) + + +@PROPERTY_SETTINGS +@given(leaves()) +def test_a_sum_of_one_timing_is_unchanged_but_for_its_flag(drawn): + scenarios, log = drawn + for before in (True, False): + same = [ + s.model_copy(update={"applied_before_data_load": before}) for s in scenarios + ] + total = functools.reduce(operator.add, same, Scenario()) + reference = functools.reduce(reference_add, same, Scenario()) + assert fields(total) == fields(reference) + assert run(log, total) == run(log, reference) + if any(changes_something(s) for s in same): + assert total.applied_before_data_load == before + expected = (total, None) if before else (None, total) + assert total.parts_by_data_load() == expected diff --git a/policyengine_uk/utils/scenario.py b/policyengine_uk/utils/scenario.py index 720fe96f4..67909d6ff 100644 --- a/policyengine_uk/utils/scenario.py +++ b/policyengine_uk/utils/scenario.py @@ -1,7 +1,5 @@ import copy -import functools import logging -import operator from pydantic import BaseModel, ConfigDict, PrivateAttr from typing import Any, Literal, List, Optional, Callable, Dict, Tuple, Type, Union @@ -118,6 +116,15 @@ def _equal(value: Any, kept: Any) -> bool: return repr(value) == repr(kept) +def _plus( + first: Optional["Scenario"], second: Optional["Scenario"] +) -> Optional["Scenario"]: + """The sum of two parts applied at the same time, either of them None.""" + if first is None or second is None: + return second if first is None else first + return first + second + + class Scenario(BaseModel): """Represents a scenario configuration for policy simulations. @@ -151,11 +158,11 @@ class Scenario(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - # A sum of parts applied at different times keeps the parts, each applied - # at one time, in the order added, for ``parts_by_data_load``; and its own - # fields as they were when it was made, so a sum whose fields are changed - # afterwards is applied as one scenario, as those fields say. - _parts: Optional[Tuple["Scenario", ...]] = PrivateAttr(default=None) + # A sum of parts applied at different times keeps its two operands, as + # they were when added, for ``parts_by_data_load``; and its own fields as + # they were when it was made, so a sum whose fields are changed afterwards + # is applied as one scenario, as those fields say. + _operands: Optional[Tuple["Scenario", "Scenario"]] = PrivateAttr(default=None) _fields_when_made: Optional[tuple] = PrivateAttr(default=None) def __add__(self, other: "Scenario") -> "Scenario": @@ -196,7 +203,7 @@ def __add__(self, other: "Scenario") -> "Scenario": merged_params[key] = value # Chain simulation modifiers, as they are when the scenarios are - # added, as the parts below keep them. + # added, as the operands kept below are. combined_modifier = None first, second = self.simulation_modifier, other.simulation_modifier @@ -213,8 +220,7 @@ def combined_modifier(simulation: Simulation) -> None: dated_reforms = (self.dated_reforms or []) + (other.dated_reforms or []) - parts = self._timed_parts() + other._timed_parts() - timings = {part.applied_before_data_load for part in parts} + timings = self._timings() | other._timings() # Return new scenario with merged configuration combined = Scenario( @@ -224,7 +230,7 @@ def combined_modifier(simulation: Simulation) -> None: simulation_modifier=combined_modifier, ) if len(timings) > 1: - combined._parts = parts + combined._operands = (self._as_added(), other._as_added()) combined._fields_when_made = ( combined.applied_before_data_load, copy.deepcopy(combined.parameter_changes), @@ -233,9 +239,14 @@ def combined_modifier(simulation: Simulation) -> None: ) return combined - def _current_parts(self) -> Optional[Tuple["Scenario", ...]]: - """A sum's parts, unless its fields have changed since it was made.""" - if self._parts is None: + def _changes_something(self) -> bool: + return bool( + self.parameter_changes or self.dated_reforms or self.simulation_modifier + ) + + def _current_operands(self) -> Optional[Tuple["Scenario", "Scenario"]]: + """A sum's operands, unless its fields have changed since it was made.""" + if self._operands is None: return None flag, parameter_changes, dated_reforms, modifier = self._fields_when_made if ( @@ -244,24 +255,44 @@ def _current_parts(self) -> Optional[Tuple["Scenario", ...]]: and _equal(self.parameter_changes, parameter_changes) and _equal(self.dated_reforms, dated_reforms) ): - return self._parts + return self._operands return None - def _timed_parts(self) -> Tuple["Scenario", ...]: - """The scenario as parts each applied at one time, in order: a sum's - parts, or a copy of the scenario (so later edits to it do not reach - a sum it is in), or none if it changes nothing.""" - parts = self._current_parts() - if parts is not None: - return parts - if not ( - self.parameter_changes or self.dated_reforms or self.simulation_modifier - ): - return () - part = self.model_copy() - part._parts = None - part._fields_when_made = None - return (part,) + def _timings(self) -> set: + """When the scenario's parts are applied: before data load (True), + after it (False), both, or neither for a scenario that changes + nothing.""" + if self._current_operands() is not None: + return {True, False} + if not self._changes_something(): + return set() + return {self.applied_before_data_load} + + def _as_added(self) -> "Scenario": + """A copy that later edits to this scenario's fields do not reach. + + Parameter changes and dated reforms are copied; the modifier is the + same object, as the sum's chained modifier calls it. + """ + return self.model_copy( + update={ + "parameter_changes": copy.deepcopy(self.parameter_changes), + "dated_reforms": copy.deepcopy(self.dated_reforms), + } + ) + + def _split(self) -> Tuple[Optional["Scenario"], Optional["Scenario"]]: + """``parts_by_data_load``, with no part for a scenario that changes + nothing.""" + operands = self._current_operands() + if operands is None: + if not self._changes_something(): + return None, None + return (self, None) if self.applied_before_data_load else (None, self) + (first_before, first_after), (second_before, second_after) = ( + operand._split() for operand in operands + ) + return _plus(first_before, second_before), _plus(first_after, second_after) def parts_by_data_load( self, @@ -274,9 +305,9 @@ def parts_by_data_load( the end, as for a scenario on its own. A scenario is one or the other, as ``applied_before_data_load`` says. A sum of parts applied at different times is split into the sum of the parts applied - before data load and the sum of the rest, each in the order added. - A sum whose fields have been changed since it was made is applied - as one scenario, as its fields say. + before data load and the sum of the rest, each added in the order + and bracketing of the whole sum. A sum whose fields have been changed + since it was made is applied as one scenario, as its fields say. Where parts applied at different times change the same parameter, the later change holds: the data is built with the value of the part @@ -291,15 +322,9 @@ def parts_by_data_load( The part applied before data load and the part applied after it, each a Scenario, or None where there is no such part. """ - parts = self._current_parts() - if parts is None: + if self._current_operands() is None: return (self, None) if self.applied_before_data_load else (None, self) - before = [part for part in parts if part.applied_before_data_load] - after = [part for part in parts if not part.applied_before_data_load] - return ( - functools.reduce(operator.add, before), - functools.reduce(operator.add, after), - ) + return self._split() def dated_assignments(self, order_start: int = 0) -> List[ParameterAssignment]: """The statutory-timeline changes of ``dated_reforms``, in order.""" @@ -490,12 +515,8 @@ def __repr__(self) -> str: parts.append(f"dated_reforms={len(self.dated_reforms)} items") if self.simulation_modifier: parts.append("simulation_modifier=") - timed = self._current_parts() - if timed is not None: - before = sum(part.applied_before_data_load for part in timed) - parts.append( - f"parts={before} before data load, {len(timed) - before} after" - ) + if self._current_operands() is not None: + parts.append("parts applied before and after data load") elif self.applied_before_data_load: parts.append("applied_before_data_load=True") return f"Scenario({', '.join(parts)})" From cbd909dba915c98d4f46d82896acb76b0f4fe3a8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 06:20:14 -0400 Subject: [PATCH 3/6] Format the docs example Co-Authored-By: Claude Opus 5.5 --- docs/book/usage/scenarios.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/book/usage/scenarios.md b/docs/book/usage/scenarios.md index 7bef155a8..63d903562 100644 --- a/docs/book/usage/scenarios.md +++ b/docs/book/usage/scenarios.md @@ -324,7 +324,9 @@ from policyengine_uk.model_api import Scenario reform = Scenario.from_reform( {"gov.hmrc.income_tax.allowances.personal_allowance.amount": {"2025.2099": 20_000}} ) -static_reform = Simulation(situation=situation, scenario=no_economic_assumptions + reform) +static_reform = Simulation( + situation=situation, scenario=no_economic_assumptions + reform +) ``` ## Advanced scenario techniques From 196ae9ba6f19eada65bb447f7b3adad23bbc7ac8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 06:45:45 -0400 Subject: [PATCH 4/6] Say when a changed sum is applied Co-Authored-By: Claude Opus 5.5 --- docs/book/usage/scenarios.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/book/usage/scenarios.md b/docs/book/usage/scenarios.md index 63d903562..1ee4421a2 100644 --- a/docs/book/usage/scenarios.md +++ b/docs/book/usage/scenarios.md @@ -153,7 +153,7 @@ for name, scenario in [ print(f"{name}: £{net_income / 12:.2f} per month") ``` -A sum keeps each part's timing. A simulation applies the parts set with `applied_before_data_load=True` before it builds its data and the rest after, each in the order added. So `no_economic_assumptions + reform` freezes the data and the parameters and applies the reform, and a modifier that edits the data still runs once the data exists. A sum reads `applied_before_data_load=True` only when every part is applied before data load. `Scenario.parts_by_data_load()` returns the two parts. Where both change the same parameter, the part applied after data load sets it last. If you set `applied_before_data_load=True` on a sum yourself, or change its fields after making it, the whole sum is applied as one scenario at that time. +A sum keeps each part's timing. A simulation applies the parts set with `applied_before_data_load=True` before it builds its data and the rest after, each in the order added. So `no_economic_assumptions + reform` freezes the data and the parameters and applies the reform, and a modifier that edits the data still runs once the data exists. A sum reads `applied_before_data_load=True` only when every part is applied before data load. `Scenario.parts_by_data_load()` returns the two parts. Where both change the same parameter, the part applied after data load sets it last. If you set `applied_before_data_load=True` on a sum yourself, or change its fields after making it, the whole sum is applied as one scenario, at the time its own `applied_before_data_load` says. ## Defining structural scenarios From b20ac13a066a7ba01929c6757c2623ef2c2b5b18 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 09:50:29 -0400 Subject: [PATCH 5/6] Apply a mixed sum's parameter changes before every modifier, and keep it out of the URL dataset cache Review r1 (Opus) found two regressions in the split by timing. The after-load part's parameter changes rebuilt the parameters after the before-load part's modifier had run, which discarded that modifier's direct edits (a structural Reform class applied before data load, or a Parameter.update modifier). And a simulation with a part applied before data load cached its projected dataset under the plain hf:// or gs:// key, so its baseline and later simulations got frozen data (as no_economic_assumptions alone already did on main). A sum with parts applied at different times now applies all of its parameter changes, as + merged them, before the data is built and before any modifier; only its modifiers are split by timing. That keeps modifier edits, gives the precedence Scenario.apply gives, and removes the copies of operands the bracketing-aware split needed. A simulation with a part applied before data load neither reads nor writes the URL dataset cache. Scenario.__deepcopy__ shares one memo between fields and private attributes, so a deep copy keeps the split with callable-object modifiers, and parts_by_data_load returns copies. Tests (MOCK data): the URL cache for hf:// and gs:// in both orders and for a macro sum; before-load direct edits and Reform classes against an after-load reform, against the all-after-load sum; the precedence of +; every modifier sees every parameter change (property over a real Simulation); deep copies with functools.partial; returned parts are copies; the split's modifiers by timing, bracketing, identity and the previous + as reference (properties). Co-Authored-By: Claude Opus 5.5 --- changelog.d/2406.fixed.md | 2 +- docs/book/usage/scenarios.md | 4 +- policyengine_uk/simulation.py | 32 +- .../tests/test_scenario_sum_timing.py | 357 +++++++++++++----- policyengine_uk/utils/scenario.py | 197 +++++----- 5 files changed, 377 insertions(+), 215 deletions(-) diff --git a/changelog.d/2406.fixed.md b/changelog.d/2406.fixed.md index 1d818e5d2..d8bd8eeea 100644 --- a/changelog.d/2406.fixed.md +++ b/changelog.d/2406.fixed.md @@ -1 +1 @@ -- A sum of scenarios keeps each part's timing. A part applied before data load, such as `no_economic_assumptions`, is still applied before it when added to a reform, so `no_economic_assumptions + reform` freezes the data as well as the parameters, and the sum of a macro scenario applied before data load and a reform projects the data with the macro growth. Parts applied after data load, including modifiers that edit the data, still run after it. The sum's fields still hold the whole sum. `Scenario.parts_by_data_load()` returns the two parts. +- A sum of scenarios keeps the timing of each part's modifier. A part applied before data load, such as `no_economic_assumptions`, is still applied before it when added to a reform, so `no_economic_assumptions + reform` freezes the data as well as the parameters. Likewise a macro scenario applied before data load, plus a reform, projects the data with the macro growth. Modifiers of parts applied after data load, including ones that edit the data, still run after it. A sum with a part applied before data load applies all of its parameter changes first, as `+` merges them, so the data follows them all and a modifier's edits survive. A simulation with a scenario applied before data load no longer reads or writes the `hf://` and `gs://` dataset cache, so its projection no longer reaches its baseline or later simulations. The sum of scenarios that are all applied before data load is now applied before it. A sum now calls the modifiers its operands had when it was made. Its fields still hold the whole sum, and `Scenario.parts_by_data_load()` returns what is applied before and after data load. diff --git a/docs/book/usage/scenarios.md b/docs/book/usage/scenarios.md index 1ee4421a2..70421363d 100644 --- a/docs/book/usage/scenarios.md +++ b/docs/book/usage/scenarios.md @@ -153,7 +153,7 @@ for name, scenario in [ print(f"{name}: £{net_income / 12:.2f} per month") ``` -A sum keeps each part's timing. A simulation applies the parts set with `applied_before_data_load=True` before it builds its data and the rest after, each in the order added. So `no_economic_assumptions + reform` freezes the data and the parameters and applies the reform, and a modifier that edits the data still runs once the data exists. A sum reads `applied_before_data_load=True` only when every part is applied before data load. `Scenario.parts_by_data_load()` returns the two parts. Where both change the same parameter, the part applied after data load sets it last. If you set `applied_before_data_load=True` on a sum yourself, or change its fields after making it, the whole sum is applied as one scenario, at the time its own `applied_before_data_load` says. +A sum keeps the timing of each part's modifier. A simulation runs the modifiers of the parts set with `applied_before_data_load=True` before it builds its data, and the rest once the data is built, each in the order added. A sum with a part applied before data load applies all of its parameter changes first, before the data is built, with the precedence `+` gives them: changing the parameters rebuilds them, which would discard the edits of a modifier that has already run. The data's projected years then follow every parameter change of the sum. So `no_economic_assumptions + reform` freezes the data and the parameters and applies the reform, and a modifier that edits the data still runs once the data exists. Such a simulation builds its own copy of a `hf://` or `gs://` dataset rather than sharing the cached one, so its baseline and later simulations keep current law's projection. A sum reads `applied_before_data_load=True` only when every part that changes something is applied before data load. `Scenario.parts_by_data_load()` returns what is applied before data load and what is applied after. If you set a sum's `applied_before_data_load` or `simulation_modifier` yourself after making it, the whole sum is applied as one scenario, at the time its own `applied_before_data_load` says. ## Defining structural scenarios @@ -316,7 +316,7 @@ print(float(baseline.calculate("employment_income", 2026)[0])) print(float(static.calculate("employment_income", 2026)[0])) ``` -Add it to a reform to measure the reform without growth. The sum keeps the freeze before the data load and applies the reform as usual: +Add it to a reform to measure the reform without growth. The sum still applies the freeze before the data is built, so a dataset's projected years keep its first year's amounts, and it applies the reform as usual: ```python from policyengine_uk.model_api import Scenario diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py index 1f278a909..7647c4a70 100644 --- a/policyengine_uk/simulation.py +++ b/policyengine_uk/simulation.py @@ -147,6 +147,10 @@ class Simulation(CoreSimulation): # true for a region or constituency filtered from the data. The builders # set it, so rebuilding a simulation in place keeps it right. built_from_dataset: bool = False + # True when part of the scenario is applied before data load, so the + # parameters the data is built with may not be current law's. Such a + # simulation neither reads nor writes the URL dataset cache. + _scenario_before_data_load: bool = False def __init__( self, @@ -210,12 +214,13 @@ def __init__( self.branches: Dict[str, Simulation] = {} - # The parts of the scenario applied before the data is built and - # after it. A scenario is one or the other; a sum of scenarios - # applied at different times keeps each part's timing. + # What the scenario applies before the data is built and after it. A + # scenario is one or the other; a sum of scenarios applied at + # different times runs each part's modifier at the part's own time. before_data_load, after_data_load = ( scenario.parts_by_data_load() if scenario is not None else (None, None) ) + self._scenario_before_data_load = before_data_load is not None # A scenario's parameter changes rebuild the parameters from the # statutory tree, so they come before its modifier, whose edits a @@ -247,8 +252,7 @@ def __init__( # (a changed economic assumption does not re-uprate it), but come # before the steps below that read the parameters: the Universal # Credit rebalancing inputs and the structural reforms switched on by - # parameters. They rebuild the parameters from those the part applied - # before data load left, so they keep that part's changes. + # parameters. if after_data_load is not None: self.apply_scenario_parameters(after_data_load) @@ -598,6 +602,15 @@ def capital_gains_projection_metadata(self) -> dict: } def _dataset_cache_key(self, source: str): + """The URL dataset cache key for a source, or None not to cache. + + A cached dataset is already projected to later years, with current + law's economic assumptions unless a scenario applied before data load + changed them, so a simulation with such a scenario neither reads nor + writes the cache. + """ + if self._scenario_before_data_load: + return None if self.tax_benefit_system.capital_gains_projection is None: return source metadata = self.capital_gains_projection_metadata @@ -690,7 +703,7 @@ def build_from_dataset_source( return if dataset_source.startswith("gs://"): cache_key = self._dataset_cache_key(dataset_source) - if cache_key in _url_dataset_cache: + if cache_key is not None and cache_key in _url_dataset_cache: 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) @@ -729,8 +742,9 @@ def build_from_file( # Pre-encode string enum columns to int16 once before caching so # subsequent loads skip the expensive astype(str) + searchsorted path. _pre_encode_enum_columns(multi_year_dataset, self.tax_benefit_system) - if cache_key is not None: - _url_dataset_cache[self._dataset_cache_key(cache_key)] = multi_year_dataset + key = self._dataset_cache_key(cache_key) if cache_key is not None else None + if key is not None: + _url_dataset_cache[key] = multi_year_dataset self.build_from_multi_year_dataset(multi_year_dataset) self._restore_cgt_extension_provenance(multi_year_dataset) @@ -751,7 +765,7 @@ def build_from_url(self, url: str) -> None: # Return early from in-memory cache if available: skips HDF5 reading, # uprating and enum encoding (~2.2s on the first load). cache_key = self._dataset_cache_key(url) - if cache_key in _url_dataset_cache: + if cache_key is not None and cache_key in _url_dataset_cache: 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) diff --git a/policyengine_uk/tests/test_scenario_sum_timing.py b/policyengine_uk/tests/test_scenario_sum_timing.py index b4582fe67..63e8e9b2b 100644 --- a/policyengine_uk/tests/test_scenario_sum_timing.py +++ b/policyengine_uk/tests/test_scenario_sum_timing.py @@ -1,34 +1,42 @@ -"""A sum of scenarios applies each part at its own time. +"""A sum of scenarios runs each part's modifier at the part's own time. ``Simulation`` applies a scenario before it builds its data (``applied_before_data_load``) or after. A sum of scenarios applied at -different times, such as ``no_economic_assumptions + reform``, keeps its parts -(``Scenario.parts_by_data_load``), and the simulation applies the parts -applied before data load before it builds the data and the rest after. +different times, such as ``no_economic_assumptions + reform``, keeps which +modifiers belong to which time (``Scenario.parts_by_data_load``). The +simulation applies all of the sum's parameter changes, then the modifiers of +the parts applied before data load, then builds the data, and runs the other +modifiers at the end. Parameter changes rebuild the parameters, which would +discard a modifier's edits, so they come before every modifier. Invariants: -1. Partition and order: every part's parameter changes and modifier are - applied exactly once, at the part's own time. A modifier applied before - data load sees no populations; one applied after sees them. Within each - time, parts are applied in the order added. -2. Associativity and identity: how a sum is bracketed does not change its - split, and ``Scenario()`` (which changes nothing) adds no part. +1. Partition and order: every modifier runs exactly once, at its part's + time. One applied before data load sees no populations; one applied after + sees them. Within each time, modifiers run in the order added, however + the sum is bracketed. +2. Parameter changes: a sum's parameter changes are applied once, as ``+`` + merged them (the precedence ``Scenario.apply`` gives them), before any + modifier, so the timing of the parts changes when parameters are set but + not their values, and a modifier applied before data load keeps its edits. 3. Backward compatibility: a sum's fields (parameter changes, dated reforms and the chained modifier) equal those the previous ``Scenario.__add__`` - made, so ``Scenario.apply`` and code reading the fields see the whole sum. - A sum reads ``applied_before_data_load=True`` exactly when every part that - changes something is applied before data load. + made. A sum reads ``applied_before_data_load=True`` exactly when every part + that changes something is applied before data load, and ``Scenario()`` + (which changes nothing) adds no part. 4. Freeze: ``no_economic_assumptions + reform``, in either order, gives the data and the economic assumption parameters of ``no_economic_assumptions`` - alone, and the reform's own changes. -5. A sum is fixed when made: later edits to its operands do not reach it, - and a sum whose own fields are changed is applied as one scenario, as - those fields say. + alone, and the reform's own changes, also for a URL dataset, whose cache + it neither reads nor writes, so the baseline and later simulations keep + current law's projection. +5. A sum is fixed when made: later edits to its operands do not reach it; a + sum whose flag or modifier is changed is applied as one scenario, as its + fields say; and its parts are copies. All datasets here are MOCK data: synthetic households, not survey records. """ +import copy import functools import operator @@ -38,7 +46,10 @@ from hypothesis import HealthCheck, given, settings from hypothesis import strategies as st from policyengine_core.parameters import Parameter +from policyengine_core.periods import period as period_ +from policyengine_core.reforms import set_parameter +import policyengine_uk.simulation as simulation_module from policyengine_uk import Simulation from policyengine_uk.data.dataset_schema import ( UKMultiYearDataset, @@ -216,7 +227,8 @@ def test_no_economic_assumptions_plus_reform_freezes_the_parameters( ): simulation = frozen_with_reform assert parameter(simulation, AVERAGE_EARNINGS, "2027-01-01") == 0 - # The reform rebuilt the parameters after the freeze and kept it. + # The freeze reloads the parameters after the reform changes them, and + # reprocesses them with the reform's changes. assert economic_assumptions(simulation) == economic_assumptions(frozen) @@ -340,23 +352,161 @@ def test_flagging_the_whole_sum_would_run_the_data_edit_without_data(): Simulation(dataset=mock_dataset(), scenario=scenario) +# --- a sum with a URL dataset --------------------------------------------- + +HF_URL = "hf://policyengine/mock-repo/MOCK_frs_2025_26.h5" +GCS_URL = "gs://mock-bucket/MOCK_frs_2025_26.h5" + + +@pytest.fixture +def url_dataset(monkeypatch, tmp_path): + """Serve both URL schemes from a MOCK file, with an empty cache.""" + path = str(tmp_path / "MOCK_frs_2025_26.h5") + mock_dataset().save(path) + monkeypatch.setattr(simulation_module, "_url_dataset_cache", {}) + monkeypatch.setattr( + simulation_module, "download_huggingface_dataset", lambda **_: path + ) + monkeypatch.setattr( + simulation_module, "materialize_gcs_dataset_url", lambda url, **_: path + ) + + +@pytest.mark.parametrize("url", [HF_URL, GCS_URL]) +@pytest.mark.parametrize("scenario_first", [True, False]) +def test_a_sum_with_a_before_load_part_keeps_out_of_the_url_dataset_cache( + url_dataset, url, scenario_first +): + if not scenario_first: + assert growth(Simulation(dataset=url)) == pytest.approx(CURRENT_LAW_GROWTH_2027) + frozen = Simulation( + dataset=url, scenario=no_economic_assumptions + Scenario.from_reform(REFORM) + ) + assert growth(frozen) == 0 + assert parameter(frozen, PERSONAL_ALLOWANCE, "2027") == 20_000 + # On main the frozen data was cached under the URL: the baseline and + # every later simulation got it. + for simulation in (frozen.baseline, Simulation(dataset=url)): + assert growth(simulation) == pytest.approx(CURRENT_LAW_GROWTH_2027) + + +def test_a_macro_sum_with_a_url_dataset_projects_its_own_data(url_dataset): + Simulation(dataset=GCS_URL) + macro = Simulation(dataset=GCS_URL, scenario=MACRO + Scenario.from_reform(REFORM)) + assert growth(macro) == pytest.approx(MACRO_GROWTH_2027) + assert growth(Simulation(dataset=GCS_URL)) == pytest.approx(CURRENT_LAW_GROWTH_2027) + + +# --- parameter changes come before every modifier ------------------------- + + +def zero_allowance_directly(simulation): + """A modifier that edits the processed parameters directly, which a + rebuild of the parameters would discard.""" + simulation.tax_benefit_system.parameters.get_child(PERSONAL_ALLOWANCE).update( + period=period_("year:2026:1"), value=0 + ) + simulation.tax_benefit_system.reset_parameter_caches() + + +def zero_allowance(before): + """MOCK reform parts that set the allowance to 0 in 2026 by editing the + processed parameters: a modifier, and a structural Reform class.""" + return { + "modifier": Scenario( + simulation_modifier=zero_allowance_directly, + applied_before_data_load=before, + ), + "reform class": Scenario.from_reform( + set_parameter(PERSONAL_ALLOWANCE, 0, period="year:2026:1") + ).model_copy(update={"applied_before_data_load": before}), + } + + +@pytest.mark.parametrize("kind", ["modifier", "reform class"]) +@pytest.mark.parametrize("before_first", [True, False]) +def test_a_before_load_parts_edits_survive_an_after_load_reform(kind, before_first): + rate = Scenario.from_reform({BASIC_RATE: {"2025.2099": 0.25}}) + part = zero_allowance(before=True)[kind] + scenario = part + rate if before_first else rate + part + simulation = Simulation(situation=SITUATION, scenario=scenario) + assert parameter(simulation, PERSONAL_ALLOWANCE, "2026-06-01") == 0 + assert parameter(simulation, BASIC_RATE, "2026-06-01") == 0.25 + # The same parts all applied after data load, as main applied any sum. + after = zero_allowance(before=False)[kind] + reference = Simulation( + situation=SITUATION, scenario=after + rate if before_first else rate + after + ) + assert simulation.calculate("income_tax", 2026)[0] == pytest.approx( + reference.calculate("income_tax", 2026)[0] + ) + + +@pytest.mark.parametrize( + "parts, expected", + [ + ((True, {"2025": 15_000}), (False, {"2025": 20_000})), + ((False, {"2025": 20_000}), (True, {"2025": 15_000})), + ], +) +def test_a_sums_parameter_changes_keep_the_precedence_of_plus(parts, expected): + """Timing changes when a sum's parameter changes are applied, not which + of them holds: the later operand's, as ``+`` merges them.""" + (first_before, first), (second_before, second) = parts, expected + timed = Scenario( + parameter_changes={PERSONAL_ALLOWANCE: first}, + applied_before_data_load=first_before, + ) + Scenario( + parameter_changes={PERSONAL_ALLOWANCE: second}, + applied_before_data_load=second_before, + ) + untimed = Scenario(parameter_changes={PERSONAL_ALLOWANCE: first}) + Scenario( + parameter_changes={PERSONAL_ALLOWANCE: second} + ) + for scenario in (timed, untimed): + simulation = Simulation(situation=SITUATION, scenario=scenario) + assert parameter(simulation, PERSONAL_ALLOWANCE, "2025") == second["2025"] + + +def test_a_mixed_sums_parameter_changes_all_reach_the_data(): + """Intended: with a part applied before data load, every parameter change + of the sum comes before the data is built, an after-load part's growth + change included.""" + record = Scenario(simulation_modifier=lambda s: None, applied_before_data_load=True) + growth_change = Scenario( + parameter_changes={AVERAGE_EARNINGS: {"year:2027-01-01:1": MACRO_GROWTH_2027}} + ) + simulation = Simulation(dataset=mock_dataset(), scenario=record + growth_change) + assert growth(simulation) == pytest.approx(MACRO_GROWTH_2027) + + # --- the split, in a real simulation -------------------------------------- @st.composite def timed_modifier_parts(draw): """Parts that each record, when their modifier runs, whether the data had - been built.""" + been built and the allowance a parameter change of every part set.""" timings = draw(st.lists(st.booleans(), min_size=1, max_size=5)) events = [] parts = [] for index, before in enumerate(timings): def record(simulation, index=index): - events.append((index, getattr(simulation, "populations", None) is not None)) + events.append( + ( + index, + getattr(simulation, "populations", None) is not None, + parameter(simulation, PERSONAL_ALLOWANCE, "2025"), + ) + ) parts.append( - Scenario(simulation_modifier=record, applied_before_data_load=before) + Scenario( + parameter_changes={PERSONAL_ALLOWANCE: {"2025": 10_000 + index}}, + simulation_modifier=record, + applied_before_data_load=before, + ) ) return timings, parts, events @@ -368,18 +518,12 @@ def test_a_simulation_runs_each_part_once_at_its_own_time(drawn): Simulation(situation=SITUATION, scenario=functools.reduce(operator.add, parts)) before = [i for i, b in enumerate(timings) if b] after = [i for i, b in enumerate(timings) if not b] - assert events == [(i, False) for i in before] + [(i, True) for i in after] - - -def test_parts_applied_later_set_a_parameter_both_change(): - before = Scenario( - parameter_changes={PERSONAL_ALLOWANCE: {"2025": 15_000}}, - applied_before_data_load=True, - ) - after = Scenario.from_reform({PERSONAL_ALLOWANCE: {"2025.2099": 20_000}}) - for scenario in (before + after, after + before): - simulation = Simulation(situation=SITUATION, scenario=scenario) - assert parameter(simulation, PERSONAL_ALLOWANCE, "2025") == 20_000 + # Every modifier sees every part's parameter change, the last one added + # holding. + allowance = 10_000 + len(timings) - 1 + assert events == [(i, False, allowance) for i in before] + [ + (i, True, allowance) for i in after + ] def test_scenario_apply_runs_every_part_of_a_sum_on_a_built_simulation(): @@ -432,22 +576,52 @@ def test_a_sum_flagged_before_data_load_is_applied_before_it_whole(): assert copied.parts_by_data_load() == (copied, None) -def test_a_sum_whose_fields_change_in_place_is_applied_as_they_say(): +def test_a_sum_applies_its_current_parameter_changes(): scenario = no_economic_assumptions + Scenario( parameter_changes={PERSONAL_ALLOWANCE: 20_000} ) - before, after = scenario.parts_by_data_load() - assert before is not None and after is not None scenario.parameter_changes[BASIC_RATE] = 0.21 + before, after = scenario.parts_by_data_load() + assert before.parameter_changes == {PERSONAL_ALLOWANCE: 20_000, BASIC_RATE: 0.21} + assert before.simulation_modifier is not None + assert after is None # No part applied after data load has a modifier. + # A sum whose modifier is replaced is applied as one scenario. + scenario.simulation_modifier = recorder([], "replaced") assert scenario.parts_by_data_load() == (None, scenario) +def test_parts_by_data_load_returns_copies(): + scenario = no_economic_assumptions + Scenario.from_reform( + {PERSONAL_ALLOWANCE: {"2026": 1}} + ) + before, _ = scenario.parts_by_data_load() + before.dated_reforms[0].changes[PERSONAL_ALLOWANCE]["2026"] = 999 + assert scenario.dated_reforms[0].changes[PERSONAL_ALLOWANCE]["2026"] == 1 + before, _ = scenario.parts_by_data_load() + assert before.dated_reforms[0].changes[PERSONAL_ALLOWANCE]["2026"] == 1 + + def test_a_mixed_sum_survives_copying(): - scenario = no_economic_assumptions + Scenario.from_reform(REFORM) - for copied in (scenario.model_copy(), scenario.model_copy(deep=True)): - before, after = copied.parts_by_data_load() - assert before.applied_before_data_load - assert not after.applied_before_data_load + log = [] + callable_object = functools.partial(lambda name, simulation: log.append(name), "x") + sums = [ + no_economic_assumptions + Scenario.from_reform(REFORM), + # The sum's own modifier is then a callable object, which a deep copy + # copies. + MACRO + Scenario(simulation_modifier=callable_object), + Scenario(simulation_modifier=callable_object, applied_before_data_load=True) + + Scenario(parameter_changes={PERSONAL_ALLOWANCE: 20_000}), + ] + for scenario in sums: + expected = [part is not None for part in scenario.parts_by_data_load()] + for copied in ( + scenario.model_copy(), + scenario.model_copy(deep=True), + copy.deepcopy(scenario), + ): + before, after = copied.parts_by_data_load() + assert before is not None and before.applied_before_data_load + assert [part is not None for part in (before, after)] == expected # --- properties of the split, without a simulation ------------------------ @@ -548,12 +722,10 @@ def draw_tree(draw, low, high): def evaluate(tree, leaf, add): - """Add up a bracketing, skipping leaves that ``leaf`` gives as None.""" + """Add up a bracketing.""" if isinstance(tree, int): return leaf(tree) first, second = (evaluate(branch, leaf, add) for branch in tree) - if first is None or second is None: - return second if first is None else first return add(first, second) @@ -572,20 +744,9 @@ def fields(scenario): return scenario.parameter_changes or None, scenario.dated_reforms or None -def merges_associatively(scenarios): - """Whether no parameter is changed by a value in one scenario and by - a dict of periods in another: the merge of parameter changes is then - associative.""" - kinds = {} - for scenario in scenarios: - for path, value in (scenario.parameter_changes or {}).items(): - kinds.setdefault(path, set()).add(isinstance(value, dict)) - return all(len(kind) == 1 for kind in kinds.values()) - - @PROPERTY_SETTINGS @given(leaves(), st.data()) -def test_a_sum_splits_into_its_parts_by_timing(drawn, data): +def test_a_sum_splits_its_modifiers_by_timing(drawn, data): scenarios, log = drawn count = len(scenarios) trees = [ @@ -596,67 +757,63 @@ def test_a_sum_splits_into_its_parts_by_timing(drawn, data): draw_tree(data.draw, 0, count), ] active = [changes_something(s) for s in scenarios] + before_labels = [ + i + for i, s in enumerate(scenarios) + if active[i] and s.applied_before_data_load and s.simulation_modifier + ] + after_labels = [ + i + for i, s in enumerate(scenarios) + if active[i] and not s.applied_before_data_load and s.simulation_modifier + ] has_before = any( a and s.applied_before_data_load for a, s in zip(active, scenarios) ) has_after = any( a and not s.applied_before_data_load for a, s in zip(active, scenarios) ) - if count == 1: - # Not a sum: its own flag says when it is applied. - has_before = scenarios[0].applied_before_data_load - has_after = not has_before - - def phase(tree, before): - """The previous ``+`` over one phase's parts, bracketed as the sum.""" - return evaluate( - tree, - lambda i: ( - scenarios[i] - if active[i] and scenarios[i].applied_before_data_load == before - else None - ), - reference_add, - ) + mixed = count > 1 and has_before and has_after - splits = [] for tree in trees: total = evaluate(tree, lambda i: scenarios[i], operator.add) before, after = total.parts_by_data_load() - splits.append((before, after)) - # 1. Partition and order: each phase holds exactly its parts' - # changes, bracketed as the sum, and runs their modifiers in the - # order added. - assert fields(before) == fields(phase(tree, True)) - assert fields(after) == fields(phase(tree, False)) - assert run(log, before) == run(log, phase(tree, True)) - assert run(log, after) == run(log, phase(tree, False)) - assert (before is None) == (not has_before) - assert (after is None) == (has_before and not has_after) - assert before is None or before.applied_before_data_load - assert after is None or not after.applied_before_data_load # 3. Backward compatibility: the fields hold the whole sum, as the - # previous ``+`` made them. + # previous ``+`` made them, bracketed the same way. reference = evaluate(tree, lambda i: scenarios[i], reference_add) assert fields(total) == fields(reference) assert run(log, total) == run(log, reference) - assert total.applied_before_data_load == (has_before and not has_after) - # 2. Identity: Scenario() adds no part. + assert total.applied_before_data_load == ( + count > 1 + and has_before + and not has_after + or count == 1 + and scenarios[0].applied_before_data_load + ) + if not mixed: + # One timing: the scenario is applied whole, at that time. + expected = ( + (total, None) if total.applied_before_data_load else (None, total) + ) + assert (before, after) == expected + continue + # 1. Partition and order: each time runs exactly its parts' + # modifiers, in the order added, however the sum is bracketed. + assert run(log, before) == before_labels + assert run(log, after) == after_labels + assert before.applied_before_data_load + assert after is None or not after.applied_before_data_load + assert (after is None) == (not after_labels) + # 2. Parameter changes: all of them, as ``+`` merged them, come with + # the part applied before data load. + assert fields(before) == fields(total) + assert fields(after) == (None, None) + # 3. Identity: Scenario() adds no part. for padded in (Scenario() + total, total + Scenario()): pad_before, pad_after = padded.parts_by_data_load() assert fields(pad_before) == fields(before) - assert fields(pad_after) == fields(after) - assert run(log, pad_before) == run(log, before) - assert run(log, pad_after) == run(log, after) - # 2. Associativity: modifiers run in the order added however the sum is - # bracketed, and the parts agree wherever the merge itself is - # associative. - for before, after in splits[1:]: - assert run(log, before) == run(log, splits[0][0]) - assert run(log, after) == run(log, splits[0][1]) - if merges_associatively(scenarios): - assert fields(before) == fields(splits[0][0]) - assert fields(after) == fields(splits[0][1]) + assert run(log, pad_before) == before_labels + assert run(log, pad_after) == after_labels @PROPERTY_SETTINGS diff --git a/policyengine_uk/utils/scenario.py b/policyengine_uk/utils/scenario.py index 67909d6ff..a6dc530a9 100644 --- a/policyengine_uk/utils/scenario.py +++ b/policyengine_uk/utils/scenario.py @@ -107,22 +107,19 @@ def _apply_reform_class(reform: Type[Reform], simulation: Simulation) -> None: simulation.tax_benefit_system.reset_parameter_caches() -def _equal(value: Any, kept: Any) -> bool: - """Whether a field still equals the copy kept of it.""" - try: - return bool(value == kept) - except Exception: - # A value whose == gives no single truth value, such as an array. - return repr(value) == repr(kept) +def _chain( + first: Optional[Callable[[Simulation], None]], + second: Optional[Callable[[Simulation], None]], +) -> Optional[Callable[[Simulation], None]]: + """A modifier that runs ``first`` and then ``second``, either of them None.""" + if first and second: + def chained(simulation: Simulation) -> None: + first(simulation) + second(simulation) -def _plus( - first: Optional["Scenario"], second: Optional["Scenario"] -) -> Optional["Scenario"]: - """The sum of two parts applied at the same time, either of them None.""" - if first is None or second is None: - return second if first is None else first - return first + second + return chained + return first or second class Scenario(BaseModel): @@ -130,12 +127,12 @@ class Scenario(BaseModel): A scenario can include parameter changes and/or simulation modifications that are applied before running a simulation. Scenarios can be combined - using the + operator, and a sum keeps each part's timing (see - ``parts_by_data_load``). + using the + operator, and a sum runs each part's modifier at the part's + own time (see ``parts_by_data_load``). """ applied_before_data_load: bool = False - """Whether ``Simulation`` applies the scenario before it builds its data, so the data's projected years follow its changes. A sum is applied before data load when every part is; a sum of parts applied at different times reads False here and applies each part at its own time.""" + """Whether ``Simulation`` applies the scenario before it builds its data, so the data's projected years follow its changes. A sum is applied before data load when every part that changes something is; a sum of parts applied at different times reads False here, runs each part's modifier at the part's own time and applies all its parameter changes before data load.""" parameter_changes: Optional[ Dict[ @@ -158,24 +155,25 @@ class Scenario(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - # A sum of parts applied at different times keeps its two operands, as - # they were when added, for ``parts_by_data_load``; and its own fields as - # they were when it was made, so a sum whose fields are changed afterwards - # is applied as one scenario, as those fields say. - _operands: Optional[Tuple["Scenario", "Scenario"]] = PrivateAttr(default=None) - _fields_when_made: Optional[tuple] = PrivateAttr(default=None) + # A sum of parts applied at different times keeps the modifiers of the + # parts applied before data load and of the rest, each chained in the + # order added, for ``parts_by_data_load``; and its own flag and modifier + # as they were when it was made, so a sum whose flag or modifier is + # changed afterwards is applied as one scenario, as its fields say. + _timed_modifiers: Optional[tuple] = PrivateAttr(default=None) + _made_with: Optional[tuple] = PrivateAttr(default=None) def __add__(self, other: "Scenario") -> "Scenario": """Combine two scenarios by merging parameter changes and chaining modifiers. The fields of the sum hold both scenarios, as they always have: ``Scenario.apply`` and code that reads the fields see the whole sum. - ``Simulation`` applies each part at its own time, before or after it - builds its data (see ``parts_by_data_load``), so adding a scenario - applied before data load, such as ``no_economic_assumptions``, to - another keeps it before data load. A scenario that changes nothing - has no timing of its own: ``Scenario() + scenario`` is applied when - ``scenario`` is. + ``Simulation`` runs each part's modifier at the part's own time, + before or after it builds its data (see ``parts_by_data_load``), so + adding a scenario applied before data load, such as + ``no_economic_assumptions``, to another keeps it before data load. A + scenario that changes nothing has no timing of its own: + ``Scenario() + scenario`` is applied when ``scenario`` is. Args: other: Another Scenario to combine with this one @@ -203,20 +201,8 @@ def __add__(self, other: "Scenario") -> "Scenario": merged_params[key] = value # Chain simulation modifiers, as they are when the scenarios are - # added, as the operands kept below are. - combined_modifier = None - first, second = self.simulation_modifier, other.simulation_modifier - - if first and second: - # Both have modifiers - chain them - def combined_modifier(simulation: Simulation) -> None: - first(simulation) - second(simulation) - - elif first: - combined_modifier = first - elif second: - combined_modifier = second + # added, as the modifiers kept by timing below are. + combined_modifier = _chain(self.simulation_modifier, other.simulation_modifier) dated_reforms = (self.dated_reforms or []) + (other.dated_reforms or []) @@ -230,101 +216,106 @@ def combined_modifier(simulation: Simulation) -> None: simulation_modifier=combined_modifier, ) if len(timings) > 1: - combined._operands = (self._as_added(), other._as_added()) - combined._fields_when_made = ( - combined.applied_before_data_load, - copy.deepcopy(combined.parameter_changes), - copy.deepcopy(combined.dated_reforms), - combined.simulation_modifier, + (first_before, first_after), (second_before, second_after) = ( + self._modifiers_by_timing(), + other._modifiers_by_timing(), ) + combined._timed_modifiers = ( + _chain(first_before, second_before), + _chain(first_after, second_after), + ) + combined._made_with = (False, combined_modifier) return combined + def __deepcopy__(self, memo: Optional[dict] = None) -> "Scenario": + # One memo for the fields and the private attributes, so a copy of a + # sum keeps the modifiers it keeps by timing those its fields chain. + return super().__deepcopy__({} if memo is None else memo) + def _changes_something(self) -> bool: return bool( self.parameter_changes or self.dated_reforms or self.simulation_modifier ) - def _current_operands(self) -> Optional[Tuple["Scenario", "Scenario"]]: - """A sum's operands, unless its fields have changed since it was made.""" - if self._operands is None: + def _current_timed_modifiers(self) -> Optional[tuple]: + """A sum's modifiers by timing, unless its flag or modifier has + changed since it was made.""" + if self._timed_modifiers is None: return None - flag, parameter_changes, dated_reforms, modifier = self._fields_when_made + flag, modifier = self._made_with if ( self.applied_before_data_load == flag and self.simulation_modifier is modifier - and _equal(self.parameter_changes, parameter_changes) - and _equal(self.dated_reforms, dated_reforms) ): - return self._operands + return self._timed_modifiers return None def _timings(self) -> set: """When the scenario's parts are applied: before data load (True), after it (False), both, or neither for a scenario that changes nothing.""" - if self._current_operands() is not None: + if self._current_timed_modifiers() is not None: return {True, False} if not self._changes_something(): return set() return {self.applied_before_data_load} - def _as_added(self) -> "Scenario": - """A copy that later edits to this scenario's fields do not reach. - - Parameter changes and dated reforms are copied; the modifier is the - same object, as the sum's chained modifier calls it. - """ - return self.model_copy( - update={ - "parameter_changes": copy.deepcopy(self.parameter_changes), - "dated_reforms": copy.deepcopy(self.dated_reforms), - } - ) - - def _split(self) -> Tuple[Optional["Scenario"], Optional["Scenario"]]: - """``parts_by_data_load``, with no part for a scenario that changes - nothing.""" - operands = self._current_operands() - if operands is None: - if not self._changes_something(): - return None, None - return (self, None) if self.applied_before_data_load else (None, self) - (first_before, first_after), (second_before, second_after) = ( - operand._split() for operand in operands - ) - return _plus(first_before, second_before), _plus(first_after, second_after) + def _modifiers_by_timing(self) -> tuple: + """The modifiers run before data load and after it.""" + timed = self._current_timed_modifiers() + if timed is not None: + return timed + if self.applied_before_data_load: + return self.simulation_modifier, None + return None, self.simulation_modifier def parts_by_data_load( self, ) -> Tuple[Optional["Scenario"], Optional["Scenario"]]: - """The parts of the scenario applied before and after data load. + """What ``Simulation`` applies before it builds its data, and after. ``Simulation`` applies the first part's parameter changes and then - its modifier before it builds its data, and the second part's + its modifier before it builds its data; and the second part's parameter changes just after the data is built and its modifier at - the end, as for a scenario on its own. A scenario is one or the - other, as ``applied_before_data_load`` says. A sum of parts applied - at different times is split into the sum of the parts applied - before data load and the sum of the rest, each added in the order - and bracketing of the whole sum. A sum whose fields have been changed - since it was made is applied as one scenario, as its fields say. - - Where parts applied at different times change the same parameter, - the later change holds: the data is built with the value of the part - applied before data load, and the part applied after it then sets - the parameter. A part applied after data load rebuilds the - parameters from those the parts before it left, so it keeps their - parameter changes, and the edits of a modifier that reloads and - reprocesses the parameters (as ``no_economic_assumptions`` does), - but not a modifier's direct edits to the processed parameters. + the end. A scenario is one or the other, as + ``applied_before_data_load`` says. + + A sum of parts applied at different times runs each part's modifier + at the part's own time: those of the parts applied before data load + before it, in the order added, and the rest at the end. Its parameter + changes all come first, before the data is built, as ``+`` merged + them (the precedence ``Scenario.apply`` gives them): changing the + parameters rebuilds them, which would discard the edits of a + modifier that has already run. So the first part has every parameter + change of the sum, and the second, if any part applied after data + load has a modifier, has only that modifier. The data's projected + years therefore follow the parameter changes of every part of such + a sum. + + A sum whose ``applied_before_data_load`` or ``simulation_modifier`` + has been changed since it was made is applied as one scenario, as + its fields say. Returns: The part applied before data load and the part applied after it, each a Scenario, or None where there is no such part. """ - if self._current_operands() is None: + timed = self._current_timed_modifiers() + if timed is None: return (self, None) if self.applied_before_data_load else (None, self) - return self._split() + before_modifier, after_modifier = timed + before = Scenario( + applied_before_data_load=True, + parameter_changes=copy.deepcopy(self.parameter_changes), + dated_reforms=copy.deepcopy(self.dated_reforms), + simulation_modifier=before_modifier, + ) + after = ( + Scenario(simulation_modifier=after_modifier) + if after_modifier is not None + else None + ) + return before, after def dated_assignments(self, order_start: int = 0) -> List[ParameterAssignment]: """The statutory-timeline changes of ``dated_reforms``, in order.""" @@ -515,8 +506,8 @@ def __repr__(self) -> str: parts.append(f"dated_reforms={len(self.dated_reforms)} items") if self.simulation_modifier: parts.append("simulation_modifier=") - if self._current_operands() is not None: - parts.append("parts applied before and after data load") + if self._current_timed_modifiers() is not None: + parts.append("modifiers applied before and after data load") elif self.applied_before_data_load: parts.append("applied_before_data_load=True") return f"Scenario({', '.join(parts)})" From ec25db7be8ae952c9aa6bc2b8e23e239d40de382 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 14 Oct 2026 14:26:00 -0400 Subject: [PATCH 6/6] Pin the cache skip for structural reforms, and assert calculated inputs under the freeze Review r2 (Opus, APPROVE at 0d2111325) follow-ups: - A before-load structural reform that changes no parameter (a MOCK re-ordered tenure_type Enum) keeps out of the URL dataset cache, whose enum columns are encoded with the system that built it; a parameter fingerprint alone would not catch it. - The freeze tests also assert the inputs the model calculates, not only the stored dataset. - Docs: the after-load paragraph points to the mixed-sum exception; "change to another value"; the dataset wording before the situation example. Comment on what a deep copy copies. Co-Authored-By: Claude Opus 5.5 --- docs/book/usage/scenarios.md | 6 +-- .../tests/test_scenario_sum_timing.py | 48 ++++++++++++++++++- policyengine_uk/utils/scenario.py | 2 + 3 files changed, 52 insertions(+), 4 deletions(-) diff --git a/docs/book/usage/scenarios.md b/docs/book/usage/scenarios.md index 70421363d..7333467d0 100644 --- a/docs/book/usage/scenarios.md +++ b/docs/book/usage/scenarios.md @@ -153,7 +153,7 @@ for name, scenario in [ print(f"{name}: £{net_income / 12:.2f} per month") ``` -A sum keeps the timing of each part's modifier. A simulation runs the modifiers of the parts set with `applied_before_data_load=True` before it builds its data, and the rest once the data is built, each in the order added. A sum with a part applied before data load applies all of its parameter changes first, before the data is built, with the precedence `+` gives them: changing the parameters rebuilds them, which would discard the edits of a modifier that has already run. The data's projected years then follow every parameter change of the sum. So `no_economic_assumptions + reform` freezes the data and the parameters and applies the reform, and a modifier that edits the data still runs once the data exists. Such a simulation builds its own copy of a `hf://` or `gs://` dataset rather than sharing the cached one, so its baseline and later simulations keep current law's projection. A sum reads `applied_before_data_load=True` only when every part that changes something is applied before data load. `Scenario.parts_by_data_load()` returns what is applied before data load and what is applied after. If you set a sum's `applied_before_data_load` or `simulation_modifier` yourself after making it, the whole sum is applied as one scenario, at the time its own `applied_before_data_load` says. +A sum keeps the timing of each part's modifier. A simulation runs the modifiers of the parts set with `applied_before_data_load=True` before it builds its data, and the rest once the data is built, each in the order added. A sum with a part applied before data load applies all of its parameter changes first, before the data is built, with the precedence `+` gives them: changing the parameters rebuilds them, which would discard the edits of a modifier that has already run. The data's projected years then follow every parameter change of the sum. So `no_economic_assumptions + reform` freezes the data and the parameters and applies the reform, and a modifier that edits the data still runs once the data exists. Such a simulation builds its own copy of a `hf://` or `gs://` dataset rather than sharing the cached one, so its baseline and later simulations keep current law's projection. A sum reads `applied_before_data_load=True` only when every part that changes something is applied before data load. `Scenario.parts_by_data_load()` returns what is applied before data load and what is applied after. If you change a sum's `applied_before_data_load` or `simulation_modifier` to another value after making it, the whole sum is applied as one scenario, at the time its own `applied_before_data_load` says. ## Defining structural scenarios @@ -316,7 +316,7 @@ print(float(baseline.calculate("employment_income", 2026)[0])) print(float(static.calculate("employment_income", 2026)[0])) ``` -Add it to a reform to measure the reform without growth. The sum still applies the freeze before the data is built, so a dataset's projected years keep its first year's amounts, and it applies the reform as usual: +Add it to a reform to measure the reform without growth. The sum still applies the freeze before the data is built, so with a dataset (`Simulation(dataset=...)`) the projected years keep the first year's amounts too, and it applies the reform as usual: ```python from policyengine_uk.model_api import Scenario @@ -411,7 +411,7 @@ Each way of passing a reform reads period keys as it always has: A single-date key in a `reform=` dict overrides the annualised value on that day only. Formulas read annual values on 1 January, so `"2027-01-01"` sets fiscal year 2027 and `"2027-04-06"` changes nothing. The model logs a warning when a key can change nothing. Use `"start.stop"` keys for dated measures. -Changes applied after the data is loaded (the default) leave the data as built, so a changed economic assumption does not re-uprate incomes. Dated changes (`parameter_changes` and `"start.stop"` or bare-year keys) still come before the steps that read parameters while the simulation is built: the Universal Credit rebalancing awards and the structural reforms that parameters switch on. Single-day keys and a scenario's own simulation modifier run after those steps, as they always have. Set `applied_before_data_load=True` on a scenario to project the data with changed economic assumptions too. +Changes applied after the data is loaded (the default) leave the data as built, so a changed economic assumption does not re-uprate incomes. Dated changes (`parameter_changes` and `"start.stop"` or bare-year keys) still come before the steps that read parameters while the simulation is built: the Universal Credit rebalancing awards and the structural reforms that parameters switch on. Single-day keys and a scenario's own simulation modifier run after those steps, as they always have. Set `applied_before_data_load=True` on a scenario to project the data with changed economic assumptions too. A sum with a part applied before data load applies all its parameter changes before the data is built (see "Combining scenarios" above). All of this applies when the scenario or reform is passed to `Simulation(...)`. `Scenario.apply(simulation)` on a simulation that is already built edits its annualised values instead, as it always has: a `"2027-04-06.2028-04-05"` key there overrides those dates of the fiscal-year values, so it does not set fiscal year 2027, and the model logs a warning saying so. Pass the scenario when constructing the simulation for dated measures. diff --git a/policyengine_uk/tests/test_scenario_sum_timing.py b/policyengine_uk/tests/test_scenario_sum_timing.py index 63e8e9b2b..78386ded3 100644 --- a/policyengine_uk/tests/test_scenario_sum_timing.py +++ b/policyengine_uk/tests/test_scenario_sum_timing.py @@ -47,7 +47,7 @@ from hypothesis import strategies as st from policyengine_core.parameters import Parameter from policyengine_core.periods import period as period_ -from policyengine_core.reforms import set_parameter +from policyengine_core.reforms import Reform, set_parameter import policyengine_uk.simulation as simulation_module from policyengine_uk import Simulation @@ -55,6 +55,7 @@ UKMultiYearDataset, UKSingleYearDataset, ) +from policyengine_uk.model_api import YEAR, Enum, Household, Variable from policyengine_uk.scenarios import no_economic_assumptions from policyengine_uk.utils.scenario import Scenario @@ -220,6 +221,10 @@ def test_no_economic_assumptions_plus_reform_freezes_the_data( err_msg=f"{table}.{column} {year}", ) assert_same_data(simulation, frozen) + # What the model calculates from the data, not only the stored dataset. + for year in YEARS: + income = simulation.calculate("employment_income_before_lsr", year) + np.testing.assert_array_equal(income, [30_000, 0], err_msg=str(year)) def test_no_economic_assumptions_plus_reform_freezes_the_parameters( @@ -397,6 +402,47 @@ def test_a_macro_sum_with_a_url_dataset_projects_its_own_data(url_dataset): assert growth(Simulation(dataset=GCS_URL)) == pytest.approx(CURRENT_LAW_GROWTH_2027) +class ReorderedTenure(Enum): + """MOCK: tenure_type's members in another order.""" + + OWNED_OUTRIGHT = "Owned outright" + OWNED_WITH_MORTGAGE = "Owned with a mortgage" + RENT_FROM_COUNCIL = "Rented from Council" + RENT_FROM_HA = "Rented from a Housing Association" + RENT_PRIVATELY = "Rented privately" + + +class tenure_type(Variable): + value_type = Enum + possible_values = ReorderedTenure + default_value = ReorderedTenure.RENT_PRIVATELY + entity = Household + label = "Tenure type of the household" + definition_period = YEAR + + +class reorder_tenure(Reform): + """A MOCK structural reform that changes no parameter.""" + + def apply(self): + self.update_variable(tenure_type) + + +def test_a_before_load_structural_reform_keeps_out_of_the_url_dataset_cache( + url_dataset, +): + """The cached dataset's enum columns are encoded with the system that + built it, so a part applied before data load keeps out of the cache even + when it changes no parameter the projection reads.""" + part = Scenario.from_reform(reorder_tenure).model_copy( + update={"applied_before_data_load": True} + ) + reformed = Simulation(dataset=GCS_URL, scenario=part) + assert reformed.calculate("tenure_type", 2025)[0] == "RENT_PRIVATELY" + later = Simulation(dataset=GCS_URL) + assert later.calculate("tenure_type", 2025)[0] == "RENT_PRIVATELY" + + # --- parameter changes come before every modifier ------------------------- diff --git a/policyengine_uk/utils/scenario.py b/policyengine_uk/utils/scenario.py index a6dc530a9..86923a252 100644 --- a/policyengine_uk/utils/scenario.py +++ b/policyengine_uk/utils/scenario.py @@ -230,6 +230,8 @@ def __add__(self, other: "Scenario") -> "Scenario": def __deepcopy__(self, memo: Optional[dict] = None) -> "Scenario": # One memo for the fields and the private attributes, so a copy of a # sum keeps the modifiers it keeps by timing those its fields chain. + # A callable object among them is copied; a function, such as the + # chained modifier, is not, so it still calls the original. return super().__deepcopy__({} if memo is None else memo) def _changes_something(self) -> bool: