diff --git a/changelog.d/core-3-32-12-shared-branch-arrays.changed.md b/changelog.d/core-3-32-12-shared-branch-arrays.changed.md new file mode 100644 index 000000000..e3b9cc7bf --- /dev/null +++ b/changelog.d/core-3-32-12-shared-branch-arrays.changed.md @@ -0,0 +1 @@ +Require policyengine-core 3.32.12 or later, whose simulation branches copy cached arrays only when they first read them, so marginal tax rates and labour supply responses use less memory with unchanged results. diff --git a/policyengine_uk/tests/code_health/test_cached_arrays_not_written_in_place.py b/policyengine_uk/tests/code_health/test_cached_arrays_not_written_in_place.py index d862762f8..321edd943 100644 --- a/policyengine_uk/tests/code_health/test_cached_arrays_not_written_in_place.py +++ b/policyengine_uk/tests/code_health/test_cached_arrays_not_written_in_place.py @@ -19,8 +19,15 @@ functions. Writes into projections (``person.benunit("x", period)``) are allowed: core builds a new array for those. - A run-time check that makes every cached array read-only as it is stored and - runs a simulation through ``Simulation.__init__`` (which applies the UC - rebalancing modifier) and the PIP phase-in scenario. + as it is read, and runs a simulation through ``Simulation.__init__`` (which applies the UC + rebalancing modifier), the PIP phase-in scenario, and the code that branches + a simulation (marginal tax rates, labour supply responses and the capital + gains realisation response). + +Branches raise the stakes: since policyengine-core 3.32.12 a branch shares the +simulation's cached arrays and copies each one only when it first reads it, so +a write in place into one of them after branching also reaches any branch that +has not read it yet. """ import ast @@ -451,15 +458,30 @@ def _uc_and_pip_claimant() -> dict: @pytest.fixture def read_only_cache(monkeypatch): - """Make every array read-only as the simulation stores it.""" + """Make every cached array read-only as it is stored and as it is read. + + Freezing on ``put`` alone misses arrays that reach a storage without it: + the copy a branch makes the first time it reads an array it shares, and + the copies ``Simulation.clone()`` makes. A branch's own branches (the + labour supply measurement under the ``baseline`` branch) read those, so + a write into them after branching would otherwise go unnoticed. + """ original_put = InMemoryStorage.put + original_get = InMemoryStorage.get def put(self, value, period, branch_name="default"): if isinstance(value, np.ndarray): value.flags.writeable = False return original_put(self, value, period, branch_name) + def get(self, period, branch_name="default"): + value = original_get(self, period, branch_name) + if isinstance(value, np.ndarray): + value.flags.writeable = False + return value + monkeypatch.setattr(InMemoryStorage, "put", put) + monkeypatch.setattr(InMemoryStorage, "get", get) @pytest.mark.parametrize("scenario", [None, "reform_pip_phase_in"]) @@ -475,3 +497,71 @@ def test_simulation_runs_with_read_only_cache(read_only_cache, scenario): # The modifiers ran: the claimant has a health element and PIP. assert sim.calculate("uc_LCWRA_element", 2026)[0] > 0 assert sim.calculate("pip", 2025)[0] > 0 + + +def _earning_couple_with_gains() -> dict: + members = ["adult_1", "adult_2", "child"] + return { + "people": { + "adult_1": { + "age": {year: 45 for year in YEARS}, + "employment_income": {year: 60_000 for year in YEARS}, + "capital_gains": {year: 50_000 for year in YEARS}, + }, + "adult_2": { + "age": {year: 43 for year in YEARS}, + "employment_income": {year: 18_000 for year in YEARS}, + }, + "child": {"age": {year: 7 for year in YEARS}}, + }, + "benunits": {"benunit": {"members": members}}, + "households": {"household": {"members": members}}, + } + + +BRANCHING_SCENARIOS = { + "labour_supply_responses": { + "gov.simulation.labour_supply_responses.substitution_elasticity": 0.25, + "gov.simulation.labour_supply_responses.income_elasticity": -0.05, + "gov.hmrc.income_tax.allowances.personal_allowance.amount": 13_070, + }, + "capital_gains_responses": { + "gov.simulation.capital_gains_responses.elasticity": 1.0, + "gov.hmrc.cgt.basic_rate": 0.20, + "gov.hmrc.cgt.higher_rate": 0.40, + }, +} + + +@pytest.mark.parametrize( + "case", ["marginal_rates", "labour_supply_responses", "capital_gains_responses"] +) +def test_branching_runs_with_read_only_cache(read_only_cache, case): + from policyengine_uk import Simulation + from policyengine_uk.model_api import Scenario + + if case == "marginal_rates": + sim = Simulation(situation=_earning_couple_with_gains()) + for year in YEARS: + sim.calculate("marginal_tax_rate", year) + sim.calculate("marginal_tax_rate_on_capital_gains", year) + # The branches ran: both earners face a rate, and so do the gains. + assert {"adult_1_pay_rise", "adult_2_pay_rise"} <= set(sim.branches) + assert (sim.calculate("marginal_tax_rate", 2026)[:2] > 0).all() + assert sim.calculate("marginal_tax_rate_on_capital_gains", 2026)[0] > 0 + return + changes = { + name: {str(year): value for year in YEARS} + for name, value in BRANCHING_SCENARIOS[case].items() + } + sim = Simulation( + situation=_earning_couple_with_gains(), + scenario=Scenario(parameter_changes=changes), + ) + for year in YEARS: + sim.calculate("household_net_income", year) + if case == "labour_supply_responses": + assert "lsr_measurement" in sim.branches + else: + # The measurement branches are deleted; the response shows they ran. + assert sim.calculate("capital_gains_behavioural_response", 2026)[0] < 0 diff --git a/pyproject.toml b/pyproject.toml index 8df817cbb..8127f8c3f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ classifiers = [ ] requires-python = ">=3.11" dependencies = [ - "policyengine-core>=3.32.9", + "policyengine-core>=3.32.12", "microdf-python>=1.2.1", "pydantic>=2.11.7", "tables>=3.10.2", diff --git a/uv.lock b/uv.lock index 320085977..e3fd3ac60 100644 --- a/uv.lock +++ b/uv.lock @@ -1453,7 +1453,7 @@ wheels = [ [[package]] name = "policyengine-core" -version = "3.32.9" +version = "3.32.13" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "dpath" }, @@ -1473,14 +1473,14 @@ dependencies = [ { name = "standard-imghdr" }, { name = "wheel" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/92/dd/b739b08d9d0964e22f328370735e8b472a749a7414539d8523a234b20c43/policyengine_core-3.32.9.tar.gz", hash = "sha256:1d7bba75d93abbea13e50b561459efb61b8fef8172007b7ab6d063e6f1d45131", size = 427371, upload-time = "2026-09-29T16:48:09.056Z" } +sdist = { url = "https://files.pythonhosted.org/packages/40/0a/b6c27953e1083d3ab2c40330c49580130fff48a2fe9cbfc8d731e92882e6/policyengine_core-3.32.13.tar.gz", hash = "sha256:a870f7e212fdfa1b5fdad7df92aa2b325fd0280ede16a40aeeceea74d44fb864", size = 457899, upload-time = "2026-10-03T11:30:20.572Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a0/c3/631b6c4ada9c7c4b5e9af89ae00ed2d0a27b450be437d8eace8d37d26a79/policyengine_core-3.32.9-py3-none-any.whl", hash = "sha256:07a7da1989c73437a14f654d81bc74c12a720919c5b087d2448b2cecb3cf8873", size = 247395, upload-time = "2026-09-29T16:48:07.453Z" }, + { url = "https://files.pythonhosted.org/packages/d8/01/00ebf4c13152347bc5396b40e80445fa92ea96712b9ec0bfdca7c5d11413/policyengine_core-3.32.13-py3-none-any.whl", hash = "sha256:262543c3ac696207f304f55664711b72e894fbad7b566407afa9d30b16086a41", size = 251251, upload-time = "2026-10-03T11:30:19.171Z" }, ] [[package]] name = "policyengine-uk" -version = "2.107.0" +version = "2.109.5" source = { editable = "." } dependencies = [ { name = "microdf-python" }, @@ -1514,7 +1514,7 @@ requires-dist = [ { name = "hypothesis", marker = "extra == 'dev'" }, { name = "jupyter-book", marker = "extra == 'dev'", specifier = ">=2.0.0a0" }, { name = "microdf-python", specifier = ">=1.2.1" }, - { name = "policyengine-core", specifier = ">=3.32.9" }, + { name = "policyengine-core", specifier = ">=3.32.12" }, { name = "pydantic", specifier = ">=2.11.7" }, { name = "pytest-cov", marker = "extra == 'dev'" }, { name = "pytest-xdist", marker = "extra == 'dev'" },