Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changelog.d/2175.changed.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Use public Core operations for UK parameter initialization and dated-cache clearing while preserving parameter and reform behavior.
6 changes: 6 additions & 0 deletions docs/book/policy/model-baseline.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
64 changes: 27 additions & 37 deletions policyengine_uk/tax_benefit_system.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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.
Expand All @@ -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.
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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.
Expand All @@ -468,39 +478,20 @@ 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:
apply_legacy_overrides(self.parameters, self.legacy_overrides)

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")

Expand Down Expand Up @@ -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()
Expand Down
207 changes: 207 additions & 0 deletions policyengine_uk/tests/test_parameter_cache_apis.py
Original file line number Diff line number Diff line change
@@ -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
Loading