From f383f6e512b8d7572ec2a85b6cb4948bfe4625c3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 07:46:49 -0400 Subject: [PATCH 1/3] Require policyengine-core 3.32.12 for branches that share cached arrays policyengine-core 3.32.12 (PolicyEngine/policyengine-core#556) makes Simulation.get_branch share the simulation's cached arrays with the new branch and copy each one only on first read, instead of deep-copying every array up front. Marginal tax rates, labour supply responses and the capital gains marginal tax rate all branch, so they stop paying for copies their branches never read. Raise the floor from 3.32.9 and relock (uv.lock moves core 3.32.9 -> 3.32.12; its entry for this package catches up from 2.102.3 to 2.104.7). The one behavioural difference core documents is a write in place into a cached array after branching, which a branch that has not read that array yet would now see. Extend the run-time half of the cached-array guard test to run the branching code paths (marginal tax rates, labour supply responses, capital gains realisation response) with every stored array read-only, so such a write fails the suite. Co-Authored-By: Claude Opus 5.5 --- ...re-3-32-12-shared-branch-arrays.changed.md | 1 + ...test_cached_arrays_not_written_in_place.py | 77 ++++++++++++++++++- pyproject.toml | 2 +- uv.lock | 10 +-- 4 files changed, 83 insertions(+), 7 deletions(-) create mode 100644 changelog.d/core-3-32-12-shared-branch-arrays.changed.md 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 0000000000..656b931547 --- /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.get_branch` shares the cached arrays with a new branch and copies each one only when the branch first reads it, instead of copying every cached array up front. Code that branches a simulation (marginal tax rates, labour supply responses, the capital gains marginal tax rate) no longer copies arrays its branches never read, and results are unchanged. 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 d862762f82..2a4e02634d 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 @@ -20,7 +20,14 @@ 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. + 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 @@ -475,3 +482,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 c7ea644519..7fde24a78f 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 47b04e0136..f6845ad64d 100644 --- a/uv.lock +++ b/uv.lock @@ -1444,7 +1444,7 @@ wheels = [ [[package]] name = "policyengine-core" -version = "3.32.9" +version = "3.32.12" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "dpath" }, @@ -1464,14 +1464,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/e3/cd/2e9c7e4a0c4395a0d556cebde1fa92b94c4c2c594815938e5bfd052b3b47/policyengine_core-3.32.12.tar.gz", hash = "sha256:1d163936b3e4b2704cf0fe33eee3b93832a84aa4bf027b552a745d6f6d35bad8", size = 451537, upload-time = "2026-10-02T04:58:58.42Z" } 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/36/9e/f8f41e2f33454be51e9116c1b96db77ec9982d5d97d8691a0d8bd7b35b99/policyengine_core-3.32.12-py3-none-any.whl", hash = "sha256:44f0cdc100ff171f107671965bdbb41aa0c9965f632e400734e9eb2bf6ed80ee", size = 250582, upload-time = "2026-10-02T04:58:56.986Z" }, ] [[package]] name = "policyengine-uk" -version = "2.102.3" +version = "2.104.7" source = { editable = "." } dependencies = [ { name = "microdf-python" }, @@ -1504,7 +1504,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 = "rich", marker = "extra == 'dev'" }, From b5159dff8be7c23e009b7df782cf928e751d87be Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 09:36:30 -0400 Subject: [PATCH 2/3] Catch the lock's own package entry up with pyproject (2.106.1) Co-Authored-By: Claude Opus 5.5 --- uv.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/uv.lock b/uv.lock index f6845ad64d..6c7f17d06d 100644 --- a/uv.lock +++ b/uv.lock @@ -1471,7 +1471,7 @@ wheels = [ [[package]] name = "policyengine-uk" -version = "2.104.7" +version = "2.106.1" source = { editable = "." } dependencies = [ { name = "microdf-python" }, From 2808bca32b8d9495857bd6fa0b068646f713f5fe Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 07:32:54 -0400 Subject: [PATCH 3/3] Freeze cached arrays as storage returns them; shorten the fragment Review of #2054 found that freezing arrays only in InMemoryStorage.put misses arrays that reach a storage without it: the copy a branch makes on its first read of a shared array (core 3.32.12) and the copies Simulation.clone() makes. A write into one of those after a nested branch exists (the labour supply measurement under the "baseline" branch) passed the run-time guard on both core versions; with the fixture also freezing what InMemoryStorage.get returns it fails, and the unmutated file still passes on both. Shorten the changelog fragment to one sentence, as docs/engineering/skills/github-prs.md asks. Fixes #2109 Co-Authored-By: Claude Opus 5.5 --- ...re-3-32-12-shared-branch-arrays.changed.md | 2 +- ...test_cached_arrays_not_written_in_place.py | 19 +++++++++++++++++-- 2 files changed, 18 insertions(+), 3 deletions(-) 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 index 656b931547..e3b9cc7bf1 100644 --- a/changelog.d/core-3-32-12-shared-branch-arrays.changed.md +++ b/changelog.d/core-3-32-12-shared-branch-arrays.changed.md @@ -1 +1 @@ -Require policyengine-core 3.32.12 or later, whose `Simulation.get_branch` shares the cached arrays with a new branch and copies each one only when the branch first reads it, instead of copying every cached array up front. Code that branches a simulation (marginal tax rates, labour supply responses, the capital gains marginal tax rate) no longer copies arrays its branches never read, and results are unchanged. +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 2a4e02634d..321edd943d 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,7 +19,7 @@ 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 + 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). @@ -458,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"])