diff --git a/changelog.d/2175.changed.md b/changelog.d/2175.changed.md new file mode 100644 index 0000000000..041e0d1f84 --- /dev/null +++ b/changelog.d/2175.changed.md @@ -0,0 +1 @@ +Use public Core operations for UK parameter initialization and dated-cache clearing while preserving parameter and reform behavior. diff --git a/docs/book/policy/model-baseline.md b/docs/book/policy/model-baseline.md index 29f7dafb3c..138f52dd4b 100644 --- a/docs/book/policy/model-baseline.md +++ b/docs/book/policy/model-baseline.md @@ -122,3 +122,9 @@ Notable exclusions: - **Some devolved tax policies**: Beyond property transaction taxes, other devolved policies may have limited coverage All parameter values include references to primary legislation and can be found in the [PolicyEngine UK parameters directory](https://github.com/PolicyEngine/policyengine-uk/tree/master/policyengine_uk/parameters). + +## Parameter cache ownership + +The UK tax-benefit system initializes Core state and installs parameter trees through Core's public operations. The existing `reset_parameter_caches()` method clears dated views on the system, the root and descendant parameter nodes without replacing the live tree or its leaves. Parameter processing retains its statutory timelines, baseline, structural reform nodes and rollback behavior; this change does not alter policy values or reform timing. + +The selected operations are available at the existing Core minimum version, 3.32.27, and in the locked 3.33.2 release, so this parameter-only migration does not change the dependency or lockfile. It implements part of [issue #2175](https://github.com/PolicyEngine/policyengine-uk/issues/2175). Supplied-input tracking, model-value provenance and clone, deletion and rebuild compatibility remain separate work. Those compatibility paths are retained, including the distinction between an amount supplied by a user and one stored by the model. diff --git a/policyengine_uk/tax_benefit_system.py b/policyengine_uk/tax_benefit_system.py index cfd1f67801..9d49758d48 100644 --- a/policyengine_uk/tax_benefit_system.py +++ b/policyengine_uk/tax_benefit_system.py @@ -4,7 +4,7 @@ from collections import OrderedDict from dataclasses import replace from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Tuple # PolicyEngine core imports from policyengine_core.parameters import Parameter, ParameterNode @@ -81,6 +81,14 @@ logger = logging.getLogger(__name__) + +def _install_parameters(system, parameters: ParameterNode) -> None: + # A Core Reform delegates unknown instance methods to its baseline. Keep + # this orchestration local and call the public base operation explicitly: + # Reform.modify_parameters would also start from the baseline's tree. + TaxBenefitSystem.modify_parameters(system, lambda _: parameters) + + # Module constants COUNTRY_DIR = Path(__file__).parent DEFAULT_DATASET_ENV_VAR = "POLICYENGINE_UK_DEFAULT_DATASET" @@ -239,11 +247,13 @@ def neutralize_variable(self, variable_name: str) -> None: super().neutralize_variable(part) def reset_parameter_caches(self): - """Reset all caches in the tax-benefit system.""" - self._parameters_at_instant_cache = {} - for parameter in self.parameters.get_descendants(): - parameter._at_instant_cache = {} - self.parameters._at_instant_cache = {} + """Clear dated views through their Core owners, retaining the tree.""" + for node in [self.parameters, *self.parameters.get_descendants()]: + if isinstance(node, ParameterNode): + # The public operation clears this node and its ancestors; + # root-only clearing would leave cached descendants behind. + node.clear_parent_cache() + _install_parameters(self, self.parameters) def reset_parameters(self) -> None: """Reset parameters by reloading from the parameters directory. @@ -253,7 +263,6 @@ def reset_parameters(self) -> None: 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 @@ -262,7 +271,7 @@ def reset_parameters(self) -> None: ): self._parameters_before_reset = current self.load_parameters(self.parameters_dir) - self.parameters = add_removed_parameter_aliases(self.parameters) + _install_parameters(self, add_removed_parameter_aliases(self.parameters)) def process_parameters(self) -> None: """Process and transform parameters with UK-specific adjustments. @@ -284,7 +293,7 @@ def process_parameters(self) -> None: """ previous = (self.parameters, self._projected_timelines, self._raw_parameters) before_reset = getattr(self, "_parameters_before_reset", None) - self._parameters_at_instant_cache = {} + 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. @@ -343,12 +352,12 @@ def apply_parameter_assignments( sorted(assignments, key=lambda assignment: assignment.order) ) ] - self._parameters_at_instant_cache = {} + self.reset_parameter_caches() statutory = self._raw_parameters is None key = _assignments_key(self.dated_assignments) if statutory else None cached = _rebuilt_parameters_cache.get(key) if key is not None else None if cached is not None: - self.parameters = cached[0].clone() + _install_parameters(self, cached[0].clone()) # Shared with the memo: read only through the copying # ``projected_timelines`` property. self._projected_timelines = cached[1] @@ -392,7 +401,7 @@ def parameters_with(self, assignments: List[ParameterAssignment]) -> ParameterNo returned is never this system's own. """ probe = copy.copy(self) - probe.parameters = self.parameters.clone() + _install_parameters(probe, self.parameters.clone()) probe.apply_parameter_assignments(assignments) return probe.parameters @@ -452,7 +461,8 @@ def _rebuild(self, raw: ParameterNode, statutory: bool) -> None: # reloaded parameters are part of it. self._raw_parameters = raw.clone() if not assignments: - self.parameters, self._projected_timelines = process_parameter_tree(raw) + parameters, self._projected_timelines = process_parameter_tree(raw) + _install_parameters(self, parameters) return # Changes to parameters processing generates are written as each is # generated; the rest go on the statutory timeline first. @@ -468,9 +478,10 @@ def _rebuild(self, raw: ParameterNode, statutory: bool) -> None: projected = prior_law_projection(on_timeline, unreformed, statutory, generated) for assignment in on_timeline: write_assignment(raw.get_child(assignment.path), assignment, projected) - self.parameters, self._projected_timelines = process_parameter_tree( + parameters, self._projected_timelines = process_parameter_tree( raw, unreformed, generated=generated ) + _install_parameters(self, parameters) def _apply_legacy_overrides(self) -> None: if self.legacy_overrides: @@ -478,29 +489,9 @@ def _apply_legacy_overrides(self) -> None: def __init__(self, capital_gains_projection: CapitalGainsProjection | None = None): """Initialize the UK tax-benefit system with entities and parameters.""" - self._parameters_at_instant_cache: Dict[str, Any] = {} - self.variables = {} + super().__init__(entities=[Person, BenUnit, Household]) self.capital_gains_projection = capital_gains_projection - # Create copies of entity classes to avoid modifying originals - person, benunit, household = ( - copy.copy(Person), - copy.copy(BenUnit), - copy.copy(Household), - ) - - # Set up entities - self.entities = [person, benunit, household] - self.person_entity = person - self.group_entities = [benunit, household] - self.group_entity_keys = [entity.key for entity in self.group_entities] - - # Link entities to this tax-benefit system - for entity in self.entities: - entity.set_tax_benefit_system(self) - - self.variable_module_metadata = {} - # Load all variables from the variables directory self.add_variables_from_directory(COUNTRY_DIR / "variables") @@ -528,8 +519,7 @@ def __init__(self, capital_gains_projection: CapitalGainsProjection | None = Non # the full pipeline (saves ~0.5s from convert_to_fiscal_year_parameters). # Reforms rebuild the parameters from the statutory tree through # apply_parameter_assignments, so they still get the full pipeline. - self._parameters_at_instant_cache = {} - self.parameters = _processed_parameters_cache.clone() + _install_parameters(self, _processed_parameters_cache.clone()) else: self.reset_parameters() _statutory_parameters_cache = self.parameters.clone() diff --git a/policyengine_uk/tests/test_parameter_cache_apis.py b/policyengine_uk/tests/test_parameter_cache_apis.py new file mode 100644 index 0000000000..08b6e5a822 --- /dev/null +++ b/policyengine_uk/tests/test_parameter_cache_apis.py @@ -0,0 +1,207 @@ +"""Dated parameter views survive the UK public cache-reset lifecycle. + +These are behavior-preservation controls for an ownership refactor. The +separate architecture gate checks that country code leaves Core cache fields +to Core; no tax-policy correction is asserted by these tests. +""" + +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor + +import pytest +from policyengine_core.parameters import ParameterNode +from policyengine_core.reforms import Reform +from policyengine_core.taxbenefitsystems import TaxBenefitSystem + +from policyengine_uk import CountryTaxBenefitSystem, Simulation +from policyengine_uk import tax_benefit_system as tbs_module +from policyengine_uk.entities import BenUnit, Household, Person +from policyengine_uk.utils.dated_reforms import read_parameter_changes + + +def parameter_tree(): + def branch(value): + return {"nested": {"amount": {"values": {"2025-01-01": value}}}} + + # Keep the real UK subtrees, including NI scales and their compatibility + # views. Mutate only this clone, never the constructor's cached roots. + root = CountryTaxBenefitSystem().parameters.clone() + root.add_child("branch", ParameterNode("branch", data=branch(11))) + root.baseline.add_child("branch", ParameterNode("branch", data=branch(23))) + return root + + +def small_system(): + system = CountryTaxBenefitSystem.__new__(CountryTaxBenefitSystem) + TaxBenefitSystem.__init__(system, entities=[Person, BenUnit, Household]) + system.modify_parameters(lambda _: parameter_tree()) + system.dated_assignments = [] + system.legacy_overrides = [] + system._projected_timelines = {} + system._raw_parameters = None + system._parameters_before_reset = None + system.capital_gains_projection = None + return system + + +def readings(system): + root = system.parameters + return ( + system.get_parameters_at_instant("2026-01-01").branch.nested.amount, + root("2026-01-01").branch.nested.amount, + root.branch("2026-01-01").nested.amount, + root.branch.nested("2026-01-01").amount, + root.baseline("2026-01-01").branch.nested.amount, + root.baseline.branch.nested("2026-01-01").amount, + ) + + +def test_reset_refreshes_system_and_every_descendant_without_replacing_tree(): + system = small_system() + root = system.parameters + leaf = root.branch.nested.amount + baseline = root.baseline + assert readings(system) == (11, 11, 11, 11, 23, 23) + old_view = system.get_parameters_at_instant("2026-01-01") + # A structural modifier can edit a timeline in place before asking for a + # reset. This deliberately does not invoke Parameter.update's own clearing. + leaf.values_list[0].value = 31 + root.baseline.branch.nested.amount.values_list[0].value = 47 + system.reset_parameter_caches() + assert readings(system) == (31, 31, 31, 31, 47, 47) + assert system.parameters is root + assert root.branch.nested.amount is leaf + assert root.baseline is baseline + assert system.get_parameters_at_instant("2026-01-01") is not old_view + + +def test_reset_keeps_an_actual_reforms_current_edits_and_baseline_separate(): + original = small_system() + + class Changed(Reform): + def apply(self): + self.parameters.branch.nested.amount.update(period="2026", value=31) + + reformed = Changed(original) + root = reformed.parameters + assert readings(original) == (11, 11, 11, 11, 23, 23) + assert readings(reformed) == (31, 31, 31, 31, 23, 23) + CountryTaxBenefitSystem.reset_parameter_caches(reformed) + assert reformed.parameters is root + assert readings(reformed) == (31, 31, 31, 31, 23, 23) + assert readings(original) == (11, 11, 11, 11, 23, 23) + + +def test_resets_can_repeat_inside_a_public_parameter_modifier(): + system = small_system() + root = system.parameters + assert readings(system)[0] == 11 + + def modifier(parameters): + parameters.branch.nested.amount.values_list[0].value = 37 + system.reset_parameter_caches() + system.reset_parameter_caches() + return parameters + + TaxBenefitSystem.modify_parameters(system, modifier) + assert system.parameters is root + assert readings(system) == (37, 37, 37, 37, 23, 23) + + +def test_cloned_systems_clear_only_their_own_dated_views_in_threads(): + original = small_system() + assert readings(original)[0] == 11 + + def change(amount): + clone = original.clone() + assert readings(clone)[0] == 11 + clone.parameters.branch.nested.amount.values_list[0].value = amount + clone.reset_parameter_caches() + return readings(clone), clone.parameters + + with ThreadPoolExecutor(max_workers=3) as pool: + results = list(pool.map(change, [29, 31, 37])) + for (values, root), amount in zip(results, [29, 31, 37]): + assert values == (amount, amount, amount, amount, 23, 23) + assert root is not original.parameters + assert readings(original) == (11, 11, 11, 11, 23, 23) + + +@pytest.mark.parametrize("route", ["process", "assignments"]) +def test_a_failed_rebuild_restores_tree_timelines_and_readable_views( + monkeypatch, route +): + system = small_system() + root = system.parameters + timelines = system._projected_timelines + original_values = readings(system) + + def fail(*args, **kwargs): + raise RuntimeError("deliberate processor failure") + + monkeypatch.setattr(tbs_module, "_rebuilt_parameters_cache", OrderedDict()) + monkeypatch.setattr(tbs_module, "process_parameter_tree", fail) + with pytest.raises(RuntimeError, match="deliberate processor failure"): + if route == "process": + system.process_parameters() + else: + # Use a real statutory path so assignment validation succeeds; + # the injected processor failure must exercise the rollback. + system.apply_parameter_assignments( + read_parameter_changes( + {"gov.hmrc.income_tax.rates.uk[0].rate": {"2026": 0.31}} + ) + ) + assert system.parameters is root + assert system._projected_timelines is timelines + assert system._raw_parameters is None + assert system.dated_assignments == [] + assert readings(system) == original_values + + +def test_real_cold_and_warm_uk_construction_keeps_entities_and_default_tax( + monkeypatch, +): + for name in ( + "_processed_parameters_cache", + "_statutory_parameters_cache", + "_projected_timelines_cache", + ): + monkeypatch.setattr(tbs_module, name, None) + cold = CountryTaxBenefitSystem() + warm = CountryTaxBenefitSystem() + assert cold.parameters is not warm.parameters + path = "gov.hmrc.income_tax.allowances.personal_allowance.amount" + for system in (cold, warm): + assert system.parameters.get_child(path)("2026-01-01") == 12_570 + for entity in system.entities: + assert entity.get_variable("age") is system.get_variable("age") + situation = { + "people": {"person": {"age": {2026: 40}, "employment_income": {2026: 30_000}}}, + "benunits": {"benunit": {"members": ["person"]}}, + "households": {"household": {"members": ["person"]}}, + } + for system in (cold, warm): + monkeypatch.setattr( + "policyengine_uk.simulation.CountryTaxBenefitSystem", system.clone + ) + simulation = Simulation(situation=situation) + assert simulation.calculate("income_tax", 2026).tolist() == pytest.approx( + [3486] + ) + + +def test_real_reload_and_processing_keep_structural_nodes_and_baseline(): + system = CountryTaxBenefitSystem() + added = ParameterNode( + "extra_structure", data={"amount": {"values": {"2025-01-01": 17}}} + ) + system.parameters.add_child("extra_structure", added) + path = "gov.hmrc.income_tax.allowances.personal_allowance.amount" + assert system.parameters.get_child(path)("2026-01-01") == 12_570 + system.get_parameters_at_instant("2026-01-01") + system.reset_parameters() + system.process_parameters() + assert system.parameters.extra_structure("2026-01-01").amount == 17 + assert system.parameters.get_child(path)("2026-01-01") == 12_570 + assert system.parameters.baseline.get_child(path)("2026-01-01") == 12_570