diff --git a/changelog.d/2374.fixed.md b/changelog.d/2374.fixed.md new file mode 100644 index 0000000000..1d71ba5221 --- /dev/null +++ b/changelog.d/2374.fixed.md @@ -0,0 +1 @@ +- A failed batch of legacy parameter overrides is now rolled back whole, nodes a structural reform adds to the parameters survive later rebuilds and reprocessing, and the scenario docs say that `Scenario.apply` on a built simulation edits the annualised values. diff --git a/docs/book/usage/scenarios.md b/docs/book/usage/scenarios.md index ffc01ede7b..d50098a3f0 100644 --- a/docs/book/usage/scenarios.md +++ b/docs/book/usage/scenarios.md @@ -396,7 +396,9 @@ 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. They 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. 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. + +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. ### Creating complex scenarios with simulation modifiers diff --git a/policyengine_uk/tax_benefit_system.py b/policyengine_uk/tax_benefit_system.py index 9e4175a690..23ef292ec4 100644 --- a/policyengine_uk/tax_benefit_system.py +++ b/policyengine_uk/tax_benefit_system.py @@ -53,6 +53,7 @@ Timeline, apply_legacy_overrides, copy_timelines, + graft_missing_nodes, requires_projection, snapshot_timelines, tree_defines, @@ -173,8 +174,21 @@ def reset_parameter_caches(self): self.parameters._at_instant_cache = {} def reset_parameters(self) -> None: - """Reset parameters by reloading from the parameters directory.""" + """Reset parameters by reloading from the parameters directory. + + The processed parameters it replaces are kept until + ``process_parameters`` runs, so nodes a structural reform added to + them survive the reload. Only a processed tree is kept, and a second + reset before processing keeps the first one. + """ self._parameters_at_instant_cache = {} + current = getattr(self, "parameters", None) + if ( + getattr(self, "_parameters_before_reset", None) is None + and current is not None + and "baseline" in current.children + ): + self._parameters_before_reset = current self.load_parameters(self.parameters_dir) self.parameters = add_removed_parameter_aliases(self.parameters) @@ -188,14 +202,34 @@ def process_parameters(self) -> None: The system's dated reform changes (``dated_assignments``) are written onto the parameters first, so a modifier that reloads and edits the parameters before processing them again (as ``no_economic_assumptions`` - does) keeps the reform. Its legacy overrides of the annualised - parameters are then applied again. The parameters as they stood - before the changes are kept as the base for later rebuilds. + does) keeps the reform. Nodes a structural reform added to the + processed parameters before ``reset_parameters`` are added back (so + a node a modifier deletes after the reset comes back too), and the + legacy overrides of the annualised parameters are applied again. + The parameters as they stood before the changes are kept as the base + for later rebuilds. If processing fails, the system is left as it + was. """ + previous = (self.parameters, self._projected_timelines, self._raw_parameters) + before_reset = getattr(self, "_parameters_before_reset", None) self._parameters_at_instant_cache = {} - self._rebuild(self.parameters, statutory=False) - self._apply_legacy_overrides() - self.reset_parameter_caches() + try: + # Processing changes the tree it is given, so it gets a copy: if + # it fails, the system keeps the tree it had. + self._rebuild(self.parameters.clone(), statutory=False) + if before_reset is not None: + graft_missing_nodes(before_reset, self.parameters) + self._apply_legacy_overrides() + except Exception: + ( + self.parameters, + self._projected_timelines, + self._raw_parameters, + ) = previous + raise + finally: + self.reset_parameter_caches() + self._parameters_before_reset = None def apply_parameter_assignments( self, assignments: List[ParameterAssignment] @@ -213,10 +247,10 @@ def apply_parameter_assignments( parameter takes under the reform's other changes (see ``policyengine_uk.utils.dated_reforms``). - A rebuild keeps the system's dated changes and legacy overrides. It - does not keep other edits made to the processed parameters, such as - a structural reform's, as ``reset_parameters`` never has. If the - rebuild fails, the system is left as it was. + A rebuild keeps the system's dated changes and legacy overrides, and + the nodes a structural reform added to the processed parameters. It + does not keep other edits made directly to existing processed + parameters. If the rebuild fails, the system is left as it was. """ if not assignments: return @@ -240,6 +274,8 @@ def apply_parameter_assignments( cached = _rebuilt_parameters_cache.get(key) if key is not None else None if cached is not None: self.parameters = cached[0].clone() + # Shared with the memo: read only through the copying + # ``projected_timelines`` property. self._projected_timelines = cached[1] else: raw = ( @@ -254,6 +290,7 @@ def apply_parameter_assignments( self.parameters.clone(), copy_timelines(self._projected_timelines), ) + graft_missing_nodes(previous[1], self.parameters) self._apply_legacy_overrides() except Exception: ( @@ -270,12 +307,31 @@ def add_legacy_overrides(self, overrides: List[LegacyOverride]) -> None: """Apply changes to the annualised parameters and keep them. Kept so a later rebuild of the parameters applies them again, after - the dated changes. Overrides that fail to apply are not kept. + the dated changes. If any override fails, none is applied or kept. """ if not overrides: return + # Every target is found before anything changes; ``update`` replaces + # a parameter's ``values_list`` and marks it and its ancestors + # modified, so keeping each list and those flags is enough to undo + # the batch. + targets = [self.parameters.get_child(override.path) for override in overrides] + saved = [(target, target.values_list) for target in targets] + flags = {} + for target in targets: + node = target + while node is not None and id(node) not in flags: + flags[id(node)] = (node, getattr(node, "modified", False)) + node = getattr(node, "parent", None) try: apply_legacy_overrides(self.parameters, overrides) + except Exception: + for target, values_list in saved: + target.values_list = values_list + target.parent.clear_parent_cache() + for node, modified in flags.values(): + node.modified = modified + raise finally: self.reset_parameter_caches() self.legacy_overrides = self.legacy_overrides + list(overrides) @@ -352,6 +408,9 @@ def __init__(self): # when it has been processed itself (None: current law's). self._raw_parameters: Optional[ParameterNode] = None self._projected_timelines: Optional[Dict[str, Timeline]] = None + # The processed parameters ``reset_parameters`` replaced, until + # ``process_parameters`` adds back the nodes only they had. + self._parameters_before_reset: Optional[ParameterNode] = None # Set up and process parameters self.parameters_dir = COUNTRY_DIR / "parameters" @@ -400,7 +459,11 @@ def statutory_parameter_defines() -> Callable[[str], bool]: def projected_timelines() -> Dict[str, Timeline]: - """Current law's projected timeline for each gov parameter, by name.""" + """Current law's projected timeline for each gov parameter, by name. + + Shared: do not change them. ``CountryTaxBenefitSystem.projected_timelines`` + returns a copy. + """ _ensure_caches() return _projected_timelines_cache diff --git a/policyengine_uk/tests/test_dated_reform_routes.py b/policyengine_uk/tests/test_dated_reform_routes.py index 383231dfe6..da92893998 100644 --- a/policyengine_uk/tests/test_dated_reform_routes.py +++ b/policyengine_uk/tests/test_dated_reform_routes.py @@ -1024,3 +1024,210 @@ def test_preserve_calendar_dates_changes_keep_their_dates(start, length, value): while day < datetime.date(2035, 1, 1): assert parameter(day.isoformat()) == timeline_value(expected, day), day day += datetime.timedelta(days=37) + + +# Review round 3 (GPT-6.1 Sol): failed override batches roll back, and +# nodes a structural reform adds survive rebuilds and reprocessing. + + +class _add_bonus(Reform): + def apply(self): + from policyengine_core.parameters import ParameterNode + + self.parameters.add_child( + "test_budget", + ParameterNode("test_budget", data={"bonus": {"values": {"2000-01-01": 1}}}), + ) + + +def test_a_failed_override_batch_applies_none_of_its_overrides(): + simulation = Simulation(situation=SITUATION) + system = simulation.tax_benefit_system + parameters = system.parameters + with pytest.raises(ValueError): + Scenario.from_reform( + {BASIC_RATE: {"2027": 0.3}, "gov.not_a_real_parameter": {"2027": 1}} + ).apply(simulation) + assert annual(simulation, BASIC_RATE, 2027) == 0.20 + assert system.legacy_overrides == [] + # An unknown path is refused before anything changes, and the tree is + # the one the system had. + assert system.parameters is parameters + Scenario.from_reform({BASIC_RATE: {"2027": 0.3}}).apply(simulation) + assert system.parameters is parameters + assert annual(simulation, BASIC_RATE, 2027) == 0.3 + simulation.apply_parameter_changes({UC_SINGLE_OLD: {"2027": 500}}) + assert annual(simulation, BASIC_RATE, 2027) == 0.3 + assert annual(simulation, UC_SINGLE_OLD, 2027) == 500 + + +def test_a_batch_failing_after_an_update_restores_values_and_flags(monkeypatch): + simulation = Simulation(situation=SITUATION) + system = simulation.tax_benefit_system + rate = system.parameters.get_child(BASIC_RATE) + values_list = rate.values_list + # Processing marks parameters modified, so clear the flags on the path + # to the rate to see whether a failed batch leaves them set. + path, node = [], rate + while node is not None: + node.modified = False + path.append(node) + node = getattr(node, "parent", None) + # Reform.from_dict keys reach the overrides unparsed, so the second + # override fails after the first has updated the rate. + reform = Reform.from_dict( + {BASIC_RATE: {"2027-01-01": 0.3}, UC_SINGLE_OLD: {"not-a-date": 1}}, + country_id="uk", + ) + # Guard against the test going empty: the batch must fail inside the + # overrides, after the rate was updated, so the rollback runs. + original = tbs_module.apply_legacy_overrides + seen = {} + + def spy(parameters, overrides): + try: + return original(parameters, overrides) + except Exception: + seen["rate_updated"] = rate.values_list is not values_list + raise + + monkeypatch.setattr(tbs_module, "apply_legacy_overrides", spy) + with pytest.raises(ValueError): + Scenario.from_reform(reform).apply(simulation) + assert seen == {"rate_updated": True} + assert rate.values_list is values_list + assert annual(simulation, BASIC_RATE, 2027) == 0.20 + assert system.legacy_overrides == [] + assert not any(node.modified for node in path) + + +def test_a_rebuild_keeps_a_structural_reforms_parameter(): + simulation = Simulation( + situation=SITUATION, reform=(_add_bonus, {"test_budget.bonus": 7}) + ) + simulation.apply_parameter_changes({BASIC_RATE: {"2027": 0.21}}) + parameters = simulation.tax_benefit_system.parameters + assert parameters.test_budget.bonus("2027-01-01") == 7 + assert annual(simulation, BASIC_RATE, 2027) == 0.21 + # Repeated rebuilds keep it too. + simulation.apply_parameter_changes({UC_SINGLE_OLD: {"2027": 500}}) + assert simulation.tax_benefit_system.parameters.test_budget.bonus("2027-01-01") == 7 + + +def test_reprocessing_keeps_a_structural_reforms_parameter(): + from policyengine_uk.scenarios.no_economic_assumptions import ( + no_economic_assumptions, + ) + + simulation = Simulation( + situation=SITUATION, reform=(_add_bonus, {"test_budget.bonus": 7}) + ) + no_economic_assumptions.apply(simulation) + parameters = simulation.tax_benefit_system.parameters + assert parameters.test_budget.bonus("2027-01-01") == 7 + assert annual(simulation, PERSONAL_ALLOWANCE, 2035) == 12_570 + assert simulation.tax_benefit_system._parameters_before_reset is None + + +def test_failed_reprocessing_leaves_the_system_as_it_was(monkeypatch): + simulation = Simulation(situation=SITUATION) + Scenario.from_reform({BASIC_RATE: {"2027": 0.3}}).apply(simulation) + system = simulation.tax_benefit_system + system.reset_parameters() + reloaded = (system.parameters, system._projected_timelines, system._raw_parameters) + + # Fail after the rebuild has replaced the parameters, while replaying + # the legacy override. + def fail(*args, **kwargs): + raise RuntimeError("replay failed") + + monkeypatch.setattr(tbs_module, "apply_legacy_overrides", fail) + with pytest.raises(RuntimeError, match="replay failed"): + system.process_parameters() + assert system.parameters is reloaded[0] + assert system._projected_timelines is reloaded[1] + assert system._raw_parameters is reloaded[2] + # The tree is as reloaded, not processed: no baseline child, and the + # basic rate keeps its legal dates rather than fiscal-year values. + assert "baseline" not in system.parameters.children + # Processing again then works. + monkeypatch.undo() + system.process_parameters() + assert "baseline" in system.parameters.children + assert annual(simulation, BASIC_RATE, 2027) == 0.3 + + +class _add_top_band(Reform): + def apply(self): + scale = self.parameters.gov.hmrc.income_tax.rates.uk + band = scale.brackets[-1].clone() + band.parent = scale + scale.brackets.append(band) + band.threshold.update(period="year:2000:200", value=200_000) + band.rate.update(period="year:2000:200", value=0.5) + + +def test_a_rebuild_keeps_a_band_a_structural_reform_appended(): + simulation = Simulation(situation=SITUATION, reform=_add_top_band) + # The scale has four bands, uk[0] to uk[3]; the reform adds uk[4]. + top = "gov.hmrc.income_tax.rates.uk[4].rate" + assert annual(simulation, top, 2027) == 0.5 + simulation.apply_parameter_changes({BASIC_RATE: {"2027": 0.21}}) + assert annual(simulation, top, 2027) == 0.5 + assert annual(simulation, BASIC_RATE, 2027) == 0.21 + + +def test_a_second_reset_keeps_the_processed_tree_for_grafting(): + simulation = Simulation( + situation=SITUATION, reform=(_add_bonus, {"test_budget.bonus": 7}) + ) + system = simulation.tax_benefit_system + system.reset_parameters() + system.reset_parameters() + system.process_parameters() + assert system.parameters.test_budget.bonus("2027-01-01") == 7 + + +def test_grafting_adds_only_missing_nodes(): + from policyengine_core.parameters import ParameterNode + + from policyengine_uk.utils.dated_reforms import graft_missing_nodes + + source = ParameterNode( + "", + data={ + "a": { + "x": {"values": {"2000-01-01": 1}}, + "y": {"values": {"2000-01-01": 2}}, + }, + "baseline": {"z": {"values": {"2000-01-01": 3}}}, + }, + ) + target = ParameterNode("", data={"a": {"x": {"values": {"2000-01-01": 10}}}}) + added = graft_missing_nodes(source, target) + assert added == ["a.y"] + assert target.a.x("2001-01-01") == 10 + assert target.a.y("2001-01-01") == 2 + assert "baseline" not in target.children + + +def test_grafting_appends_missing_scale_brackets(): + from policyengine_core.parameters import ParameterNode + + from policyengine_uk.utils.dated_reforms import graft_missing_nodes + + def band(threshold, rate): + return { + "threshold": {"values": {"2000-01-01": threshold}}, + "rate": {"values": {"2000-01-01": rate}}, + } + + source = ParameterNode( + "", data={"s": {"brackets": [band(0, 0.1), band(10, 0.2), band(20, 0.3)]}} + ) + target = ParameterNode("", data={"s": {"brackets": [band(0, 0.15), band(10, 0.2)]}}) + assert graft_missing_nodes(source, target) == ["s[2]"] + assert len(target.s.brackets) == 3 + assert target.s.brackets[0].rate("2001-01-01") == 0.15 + assert target.s.brackets[2].rate("2001-01-01") == 0.3 + assert target.s.brackets[2].parent is target.s diff --git a/policyengine_uk/utils/dated_reforms.py b/policyengine_uk/utils/dated_reforms.py index ae243793f6..edc9ed837d 100644 --- a/policyengine_uk/utils/dated_reforms.py +++ b/policyengine_uk/utils/dated_reforms.py @@ -79,9 +79,10 @@ been given. Any later rebuild of its parameters (another call to ``Simulation.apply_parameter_changes``, or a modifier that reloads and reprocesses the parameters, as ``no_economic_assumptions`` does) writes them -all again, onto the parameters as that modifier left them. A rebuild does -not keep other edits to the processed parameters, such as a structural -reform's. +all again, onto the parameters as that modifier left them. Nodes a +structural reform added to the processed parameters (a new parameter, or a +band appended to a scale) are added back (``graft_missing_nodes``); other +edits made directly to existing processed parameters are not kept. ``Scenario.apply`` on an already-built simulation never rebuilds: it edits the annualised parameters, as it always has, so every change it makes is a @@ -96,7 +97,7 @@ from dataclasses import dataclass from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple -from policyengine_core.parameters import Parameter, ParameterNode +from policyengine_core.parameters import Parameter, ParameterNode, ParameterScale from policyengine_core.parameters.parameter_at_instant import ParameterAtInstant from policyengine_core.periods import Instant, Period, instant, period @@ -502,6 +503,43 @@ def copy_timelines(timelines: Dict[str, Timeline]) -> Dict[str, Timeline]: } +def graft_missing_nodes( + source: ParameterNode, target: ParameterNode, root: bool = True +) -> List[str]: + """Add to ``target`` copies of the nodes ``source`` has and it lacks. + + A rebuild or a reload starts from the YAML, so nodes a structural + reform added to the processed parameters (``parameters.add_child``, or + a bracket appended to a scale) would otherwise be lost. Existing + parameters are left as ``target`` has them, and the root's ``baseline`` + child, which every processing builds again, is skipped. Added nodes are + copies of ``source``'s, not processed again. Returns the names of the + nodes added. + """ + added = [] + for key, child in source.children.items(): + if root and key == "baseline": + continue + existing = target.children.get(key) + if existing is None: + target.add_child(key, child.clone()) + added.append(child.name) + elif isinstance(child, ParameterNode) and isinstance(existing, ParameterNode): + added += graft_missing_nodes(child, existing, root=False) + elif isinstance(child, ParameterScale) and isinstance(existing, ParameterScale): + for bracket, existing_bracket in zip(child.brackets, existing.brackets): + added += graft_missing_nodes(bracket, existing_bracket, root=False) + extra = child.brackets[len(existing.brackets) :] + for bracket in extra: + copy = bracket.clone() + copy.parent = existing + existing.brackets.append(copy) + added.append(bracket.name) + if extra: + existing.clear_parent_cache() + return added + + def tree_defines(parameters: ParameterNode) -> Callable[[str], bool]: """``defines(path)``: whether a parameter tree has a leaf at ``path``."""