From 14785411902510c18f27e8dde3712f6ce5c61fae Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 01:15:36 -0400 Subject: [PATCH 1/5] Uprate only from inputs, never from calculated values Simulation._calculate uprated a variable from its latest earlier stored period in its own unit, including periods the simulation had calculated itself, so the result depended on calculation order. Executed on fix-carry-over-order (index +3.7%/yr, input at 2012, 2015 asked alone vs after 2013 and 2014): - int: [1116, 85] alone, [1115, 83] after (truncation compounds) - float: float32 rounding compounds - defined_for false for one person in 2013 only: [1115.16, 1115.16] alone, [0.0, 1115.16] after (the 2013 mask carries into 2015) Also: a cached default, a monthly input carried into a calculated year, and a reform formula's result before its end each became an uprating source. For variables with uprating, keep only the earlier same-unit periods in Holder.get_input_periods(branch_name) (the carry-over PR's branch-aware provenance): periods whose value this branch reads is an input. That also drops periods stored only under a branch this one cannot read, which read back as None and raised TypeError. With no earlier input, the carry-over path or the default applies, as when the period is asked alone. Tests: tests/core/test_uprating_order.py (regressions, no hypothesis), tests/core/test_uprating_order_property.py (derandomized Hypothesis: any requests, branches and branch inputs, carry-over on/off == alone, byte for byte, and == reference rule), shared fixtures in tests/fixtures/uprating_order.py. #551's path-independence test is now exact instead of rel=1e-5. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-uprating-order.fixed.md | 1 + policyengine_core/simulations/simulation.py | 27 +- tests/core/test_uprating_order.py | 285 +++++++++++++++++ tests/core/test_uprating_order_property.py | 299 ++++++++++++++++++ .../core/variables/test_variable_uprating.py | 11 +- tests/fixtures/uprating_order.py | 215 +++++++++++++ 6 files changed, 828 insertions(+), 10 deletions(-) create mode 100644 changelog.d/fix-uprating-order.fixed.md create mode 100644 tests/core/test_uprating_order.py create mode 100644 tests/core/test_uprating_order_property.py create mode 100644 tests/fixtures/uprating_order.py diff --git a/changelog.d/fix-uprating-order.fixed.md b/changelog.d/fix-uprating-order.fixed.md new file mode 100644 index 000000000..ae873047d --- /dev/null +++ b/changelog.d/fix-uprating-order.fixed.md @@ -0,0 +1 @@ +Uprating now starts from the latest earlier input in the variable's own unit and skips periods the simulation calculated, so calculating intermediate periods first no longer compounds rounding or truncation, or carries an eligibility mask or a default into later uprated values. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 710165141..c5f977f9b 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -935,13 +935,31 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if known_period.unit == variable.definition_period and known_period.start < period.start ] - if variable.uprating is not None and len(earlier_known_periods) > 0: + # Uprate only from an input this branch reads (see + # ``Holder.get_input_periods``). Every value this simulation + # calculated is marked derived when cached, and uprating from + # one would make the result depend on which periods were + # calculated first: an integer truncated, or a float32 + # rounded, at an intermediate period would compound, and a + # value masked by ``defined_for``, a default, or a value + # carried from another unit, cached there, would replace the + # input. A period stored only under a branch this one cannot + # read would read back as ``None``. + earlier_input_periods = [] + if variable.uprating is not None: + input_periods = set(holder.get_input_periods(self.branch_name)) + earlier_input_periods = [ + known_period + for known_period in earlier_known_periods + if known_period in input_periods + ] + if earlier_input_periods: # Take the latest period from the filtered list itself. # Indexing ``known_periods`` with a position in the # filtered list picked the wrong period whenever a later # one was stored first. latest_known_period = max( - earlier_known_periods, key=lambda p: p.start + earlier_input_periods, key=lambda p: p.start ) try: uprating_parameter = get_parameter( @@ -1028,8 +1046,9 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # No input to carry, but a later period is stored: as # before, return the default without caching it. A # cached default would change what a formula testing - # whether a value is stored sees, and the uprating path - # above would uprate later periods from it. + # whether a value is stored sees. (The uprating path + # above skips derived periods, so it would not uprate + # from one.) return holder.default_array() else: array = holder.default_array() diff --git a/tests/core/test_uprating_order.py b/tests/core/test_uprating_order.py new file mode 100644 index 000000000..d67ffebee --- /dev/null +++ b/tests/core/test_uprating_order.py @@ -0,0 +1,285 @@ +"""Uprating depends only on the inputs, not on what was calculated first. + +A variable with ``uprating`` and no formula result for a period takes the +latest earlier value stored in its own unit and multiplies it by the ratio of +the uprating index at the two period starts. Until this fix, +``Simulation._calculate`` took that value from any stored period, including +one the simulation had itself calculated, so calculating an intermediate +period first changed the result: + +* An integer value truncated at each intermediate year compounded: an input of + 1001 for 2012 gave 1116 for 2015 asked alone, and 1115 after 2013 and 2014. +* A float32 value rounded at each intermediate year compounded the rounding. +* A value masked by ``defined_for``, or a default cached at an intermediate + year, replaced the input from then on: someone ineligible in 2013 only got + 0 for 2015 after 2013 was calculated, and the input uprated otherwise. +* A default, an input stored in another unit and carried into an + intermediate year, or the result of a formula a reform gave the variable up + to an ``end`` became a value to uprate, where the period asked alone gave + the default, the carried input or the uprated input. + +Uprating now starts only from an input this branch reads +(``Holder.get_input_periods``): values the simulation calculated are marked +derived when cached and skipped, as auto-carry-over skips them, and so are +periods stored only under branches this one cannot read. The rule is in +``tests/fixtures/uprating_order.py`` +and the order-independence property in ``test_uprating_order_property.py``. +Each regression here compares with a fresh simulation that calculates only +the period in question, byte for byte. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import entities +from policyengine_core.reforms import Reform +from policyengine_core.variables import Variable +from tests.fixtures.uprating_order import ( + UPRATING, + alone, + assert_bitwise_equal, + build_system, + index, + simulation, +) + +SYSTEMS = {True: build_system(True), False: build_system(False)} + + +@pytest.fixture(params=[True, False], ids=["carry_over", "no_carry_over"]) +def system(request): + return SYSTEMS[request.param] + + +def _after(system, inputs, requests, variable, period): + built = simulation(system, inputs) + for requested_variable, requested_period in requests: + built.calculate(requested_variable, requested_period) + return built.calculate(variable, period) + + +def test_integer_truncation_does_not_compound(system): + inputs = {"uprated_count": {"2012": [1001, 77]}} + requests = [("uprated_count", "2013"), ("uprated_count", "2014")] + result = _after(system, inputs, requests, "uprated_count", "2015") + # Before the fix: [1115, 83], truncated at 2013 and again at 2014. + np.testing.assert_array_equal(result, [1116, 85]) + assert_bitwise_equal(result, alone(system, inputs, "uprated_count", "2015")) + + +def test_float32_rounding_does_not_compound(system): + inputs = {"uprated": {"2012": [1001.3, 77.7]}} + requests = [("uprated", "2013"), ("uprated", "2014")] + result = _after(system, inputs, requests, "uprated", "2015") + assert_bitwise_equal(result, alone(system, inputs, "uprated", "2015")) + np.testing.assert_allclose( + result, np.array([1001.3, 77.7]) * index(2015) / index(2012), rtol=1e-6 + ) + + +def test_value_masked_by_defined_for_is_not_uprated(system): + inputs = { + "uprated_if_eligible": {"2012": [1000, 1000]}, + "eligible": { + "2012": [True, True], + "2013": [False, True], + "2014": [True, True], + "2015": [True, True], + }, + } + requests = [("uprated_if_eligible", "2013"), ("uprated_if_eligible", "2014")] + result = _after(system, inputs, requests, "uprated_if_eligible", "2015") + # Before the fix: [0, 1115.16]; the 2013 mask carried into 2015. + np.testing.assert_allclose( + result, [1000 * index(2015) / index(2012)] * 2, rtol=1e-6 + ) + assert_bitwise_equal(result, alone(system, inputs, "uprated_if_eligible", "2015")) + + +def test_default_cached_where_defined_for_is_false_is_not_uprated(system): + inputs = { + "uprated_if_eligible": {"2012": [1000, 1000]}, + "eligible": { + "2012": [True, True], + "2013": [False, False], + "2014": [True, True], + }, + } + requests = [("uprated_if_eligible", "2013")] + result = _after(system, inputs, requests, "uprated_if_eligible", "2014") + # Before the fix: [0, 0], the default cached for 2013. + np.testing.assert_allclose( + result, [1000 * index(2014) / index(2012)] * 2, rtol=1e-6 + ) + assert_bitwise_equal(result, alone(system, inputs, "uprated_if_eligible", "2014")) + + +def test_cached_default_value_is_not_uprated(system): + requests = [("uprated_with_default", "2013")] + result = _after(system, {}, requests, "uprated_with_default", "2015") + # Before the fix: 10 * index(2015) / index(2013). + np.testing.assert_array_equal(result, [10, 10]) + assert_bitwise_equal(result, alone(system, {}, "uprated_with_default", "2015")) + + +def test_input_in_another_unit_is_not_uprated_through_a_calculated_year(system): + inputs = {"uprated_any_unit": {"2012-03": [10, 20]}} + requests = [("uprated_any_unit", "2013")] + result = _after(system, inputs, requests, "uprated_any_unit", "2015") + # Asked alone, 2015 has no earlier yearly input to uprate: auto-carry-over + # carries the monthly input unchanged, or it gets the default. Before the + # fix, calculating 2013 first cached that value at 2013, and 2015 uprated + # it. + expected = [10, 20] if system.auto_carry_over_input_variables else [0, 0] + np.testing.assert_array_equal(result, expected) + assert_bitwise_equal(result, alone(system, inputs, "uprated_any_unit", "2015")) + + +def test_reform_formula_result_is_not_uprated_past_the_formula_end(system): + """Core rejects a variable declaring both a formula and ``uprating``, but + a reform can give an uprated variable a formula, which keeps the + inherited ``uprating``. Past the formula's ``end`` the variable is uprated + again, from its input, not from a formula result.""" + + class formula_until_2013(Reform): + def apply(self): + class uprated(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + end = "2013-12-31" + label = "Uprated yearly input given a formula up to 2013" + + def formula(person, period, parameters): + return person.filled_array(42) + + self.update_variable(uprated) + + reformed = formula_until_2013(system) + assert reformed.get_variable("uprated").uprating == UPRATING + inputs = {"uprated": {"2012": [1, 2]}} + result = _after(reformed, inputs, [("uprated", "2013")], "uprated", "2015") + # Before the fix: the 2013 formula result, 42, uprated to 2015. + np.testing.assert_allclose( + result, np.array([1, 2]) * index(2015) / index(2012), rtol=1e-6 + ) + assert_bitwise_equal(result, alone(reformed, inputs, "uprated", "2015")) + + +def test_monthly_rounding_does_not_compound(system): + inputs = {"uprated_monthly": {"2011-06": [1001.3, 77.7]}} + requests = [("uprated_monthly", f"2012-{month:02d}") for month in range(1, 13)] + requests += [("uprated_monthly", "2013-01")] + result = _after(system, inputs, requests, "uprated_monthly", "2014-03") + assert_bitwise_equal(result, alone(system, inputs, "uprated_monthly", "2014-03")) + np.testing.assert_allclose( + result, np.array([1001.3, 77.7]) * index(2014) / index(2011), rtol=1e-6 + ) + + +@pytest.mark.parametrize("nested", [False, True]) +def test_branch_forked_after_intermediate_years_uprates_the_input(system, nested): + inputs = {"uprated_count": {"2012": [1001, 77]}} + built = simulation(system, inputs) + built.calculate("uprated_count", "2013") + built.calculate("uprated_count", "2014") + branch = built.get_branch("reform") + if nested: + branch = branch.get_branch("nested") + result = branch.calculate("uprated_count", "2015") + np.testing.assert_array_equal(result, [1116, 85]) + assert_bitwise_equal(result, alone(system, inputs, "uprated_count", "2015")) + + +def test_parent_calculation_does_not_hide_a_branch_input(): + """An input stored only under a branch keeps its provenance when the + parent calculates a value for the same period.""" + system = SYSTEMS[True] + built = simulation(system, {"uprated": {"2011": [1, 2]}}) + built.get_holder("uprated").set_input( + periods.period("2012"), np.array([300.0, 400.0]), "other" + ) + built.calculate("uprated", "2012") + branch = built.get_branch("other") + np.testing.assert_allclose( + branch.calculate("uprated", "2014"), + np.array([300, 400]) * index(2014) / index(2012), + rtol=1e-6, + ) + + +def test_period_stored_only_under_an_unreadable_branch_is_not_uprated(system): + """A holder can store a period under a branch this simulation cannot + read (here through ``Holder.set_input`` with another branch name). It is + not an input this simulation can uprate from.""" + inputs = {"uprated_count": {"2012": [1001, 77]}} + built = simulation(system, inputs) + built.get_holder("uprated_count").set_input( + periods.period("2013"), np.array([2001, 177]), "sibling" + ) + # Before the fix: TypeError, the sibling's 2013 read back as None here. + result = built.calculate("uprated_count", "2015") + np.testing.assert_array_equal(result, [1116, 85]) + assert_bitwise_equal(result, alone(system, inputs, "uprated_count", "2015")) + + +# Guards against over-correction: inputs at intermediate periods still count. + + +def test_latest_earlier_input_is_uprated(system): + inputs = {"uprated": {"2012": [1, 2], "2014": [300, 400]}} + for requests in ( + [], + [("uprated", "2013")], + [("uprated", "2013"), ("uprated", "2016")], + ): + result = _after(system, inputs, requests, "uprated", "2015") + np.testing.assert_allclose( + result, np.array([300, 400]) * index(2015) / index(2014), rtol=1e-6 + ) + + +def test_input_set_over_a_calculated_period_is_uprated(system): + built = simulation(system, {"uprated": {"2012": [1, 2]}}) + built.calculate("uprated", "2013") + built.set_input("uprated", "2013", np.array([300.0, 400.0])) + np.testing.assert_allclose( + built.calculate("uprated", "2015"), + np.array([300, 400]) * index(2015) / index(2013), + rtol=1e-6, + ) + + +def test_branch_input_over_a_period_the_parent_calculated_is_uprated(system): + built = simulation(system, {"uprated": {"2012": [1, 2]}}) + built.calculate("uprated", "2013") + branch = built.get_branch("reform") + branch.set_input("uprated", "2013", np.array([300.0, 400.0])) + np.testing.assert_allclose( + branch.calculate("uprated", "2015"), + np.array([300, 400]) * index(2015) / index(2013), + rtol=1e-6, + ) + np.testing.assert_allclose( + built.calculate("uprated", "2015"), + np.array([1, 2]) * index(2015) / index(2012), + rtol=1e-6, + ) + + +def test_add_over_an_input_period_keeps_it_an_input(system): + built = simulation(system, {"uprated_count": {"2012": [1001, 77]}}) + built.calculate_add("uprated_count", "2012") + built.calculate("uprated_count", "2013") + np.testing.assert_array_equal(built.calculate("uprated_count", "2015"), [1116, 85]) + + +def test_uprated_value_is_cached_as_derived(system): + built = simulation(system, {"uprated": {"2012": [1, 2]}}) + built.calculate("uprated", "2013") + holder = built.get_holder("uprated") + assert holder.is_derived(periods.period("2013")) + assert not holder.is_derived(periods.period("2012")) diff --git a/tests/core/test_uprating_order_property.py b/tests/core/test_uprating_order_property.py new file mode 100644 index 000000000..5f015cc3c --- /dev/null +++ b/tests/core/test_uprating_order_property.py @@ -0,0 +1,299 @@ +"""Properties of uprating: order independence and the reference rule. + +Over the domain generated here, what a simulation returns for an uprated +variable, after earlier requests, equals byte for byte (dtype, shape and +bytes, so ``-0.0`` differs from ``0.0``) what a fresh simulation given the +same inputs returns for that period alone. It also equals the reference rule +in ``tests/fixtures/uprating_order.py``. Checked with and without +auto-carry-over. Regressions are in ``test_uprating_order.py``. + +The domain: + +* Inputs are stored in each variable's own unit. ``uprated_any_unit`` has no + ``set_input`` helper, so its inputs are years or months. +* Earlier requests are plain calculations in either unit, ``ADD`` and + ``DIVIDE``. +* The value may be read from the simulation, or from a branch or nested + branch forked after those requests. The branch may set its own input for + the target variable. The earlier requests then avoid the target period of + the target variable: a value the parent calculated before the fork stays + cached for the branch whatever its inputs are, which is not what this + module tests. +* Targets are in the variable's definition unit. +""" + +from __future__ import annotations + +import pytest + +# The smoke job installs Core without the dev extra but collects every module. +pytest.importorskip("hypothesis") + +from hypothesis import HealthCheck, example, given, settings # noqa: E402 +from hypothesis import strategies as st # noqa: E402 + +from policyengine_core import periods # noqa: E402 +from tests.fixtures.uprating_order import ( # noqa: E402 + COUNT, + assert_bitwise_equal, + build_system, + reference, + request, + simulation, +) + +SYSTEMS = {True: build_system(True), False: build_system(False)} + +YEARS = [str(year) for year in range(2009, 2019)] +INPUT_YEARS = [str(year) for year in range(2010, 2017)] +MONTHS = [f"{year}-{month:02d}" for year in range(2011, 2016) for month in range(1, 13)] +INPUT_MONTHS = [month for month in MONTHS if "2011-06" <= month <= "2013-12"] +REQUEST_MONTHS = [month for month in MONTHS if "2011-11" <= month <= "2015-06"] + +YEAR_VARIABLES = [ + "uprated", + "uprated_count", + "eligible", + "uprated_if_eligible", + "uprated_with_default", + "uprated_any_unit", +] +ALL_VARIABLES = YEAR_VARIABLES + ["uprated_monthly"] +FLOATS = st.floats(min_value=-2000, max_value=2000, allow_nan=False, width=32) +VALUES = {variable: FLOATS for variable in ALL_VARIABLES} +VALUES["uprated_count"] = st.integers(min_value=-2000, max_value=2000) +VALUES["eligible"] = st.booleans() +INPUT_PERIODS = {variable: INPUT_YEARS for variable in YEAR_VARIABLES} +INPUT_PERIODS["uprated_monthly"] = INPUT_MONTHS +INPUT_PERIODS["uprated_any_unit"] = INPUT_YEARS + INPUT_MONTHS +# Other variables read ``eligible`` (``uprated_if_eligible`` is defined for +# it), so a parent request for them caches it at their period: no branch +# input for it. +BRANCH_INPUT_VARIABLES = [name for name in ALL_VARIABLES if name != "eligible"] + + +def _variable_inputs(variable, min_size=0): + return st.dictionaries( + st.sampled_from(INPUT_PERIODS[variable]), + st.lists(VALUES[variable], min_size=COUNT, max_size=COUNT), + min_size=min_size, + max_size=3, + ) + + +def _requests_on(variable): + """Requests for ``variable``: plain calculations in either unit (a yearly + flow asked for a month caches a twelfth there, a monthly flow asked for a + year caches the sum of its months), and the ADD and DIVIDE options.""" + years = st.sampled_from(YEARS) + any_periods = st.sampled_from(YEARS + REQUEST_MONTHS) + if variable == "uprated_monthly": + return st.tuples( + st.sampled_from(["calculate", "add"]), st.just(variable), any_periods + ) + return st.one_of( + st.tuples(st.just("calculate"), st.just(variable), any_periods), + st.tuples(st.just("add"), st.just(variable), years), + st.tuples(st.just("divide"), st.just(variable), any_periods), + ) + + +def _overlaps(first, second): + first, second = periods.period(first), periods.period(second) + return first.start <= second.stop and second.start <= first.stop + + +target_strategy = st.one_of( + st.tuples(st.sampled_from(YEAR_VARIABLES), st.sampled_from(YEARS)), + st.tuples(st.just("uprated_monthly"), st.sampled_from(REQUEST_MONTHS)), +) + + +@st.composite +def scenarios(draw): + """A target, the inputs, the requests made before the target, the branch + it is read from and that branch's own input, if any. The target variable + always has an input, and about half the requests are for it: those are + the requests that cache values it could be uprated from.""" + variable, period = draw(target_strategy) + inputs = { + name: draw(_variable_inputs(name, min_size=int(name == variable))) + for name in ALL_VARIABLES + } + inputs = {name: values for name, values in inputs.items() if values} + any_request = st.one_of(*(_requests_on(name) for name in ALL_VARIABLES)) + requests = draw( + st.lists(st.one_of(_requests_on(variable), any_request), max_size=8) + ) + branch = draw(st.sampled_from([None, "reform", "nested"])) + branch_input = None + if branch is not None and variable in BRANCH_INPUT_VARIABLES: + branch_input = draw( + st.none() + | st.tuples( + st.sampled_from(INPUT_PERIODS[variable]), + st.lists(VALUES[variable], min_size=COUNT, max_size=COUNT), + ) + ) + if branch_input is not None: + requests = [ + (kind, name, requested) + for kind, name, requested in requests + if not (name == variable and _overlaps(requested, period)) + ] + return inputs, requests, variable, period, branch, branch_input + + +def _read_from(built, branch, variable, branch_input): + if branch == "reform": + built = built.get_branch("reform") + elif branch == "nested": + built = built.get_branch("reform").get_branch("nested") + if branch_input is not None: + built.set_input(variable, branch_input[0], branch_input[1]) + return built + + +@settings( + max_examples=500, + deadline=None, + derandomize=True, + suppress_health_check=[HealthCheck.too_slow], +) +@given(scenario=scenarios(), auto_carry_over=st.booleans()) +# The regressions in test_uprating_order.py, as examples of the property. +@example( + scenario=( + {"uprated_count": {"2012": [1001, 77]}}, + [ + ("calculate", "uprated_count", "2013"), + ("calculate", "uprated_count", "2014"), + ], + "uprated_count", + "2015", + None, + None, + ), + auto_carry_over=True, +) +@example( + scenario=( + {"uprated": {"2012": [1001.3, 77.7]}}, + [("calculate", "uprated", "2013"), ("divide", "uprated", "2014")], + "uprated", + "2015", + "nested", + None, + ), + auto_carry_over=False, +) +@example( + scenario=( + { + "uprated_if_eligible": {"2012": [1000, 1000]}, + "eligible": {"2012": [True, True], "2013": [False, True]}, + }, + [("calculate", "uprated_if_eligible", "2013")], + "uprated_if_eligible", + "2015", + "reform", + None, + ), + auto_carry_over=True, +) +@example( + scenario=( + { + "uprated_if_eligible": {"2012": [1000, 1000]}, + "eligible": {"2012": [True, True], "2013": [False, False]}, + }, + [("add", "uprated_if_eligible", "2013")], + "uprated_if_eligible", + "2014", + None, + None, + ), + auto_carry_over=True, +) +@example( + scenario=( + {"uprated_with_default": {"2016": [1, 2]}}, + [("calculate", "uprated_with_default", "2013")], + "uprated_with_default", + "2015", + None, + None, + ), + auto_carry_over=False, +) +@example( + scenario=( + {"uprated_any_unit": {"2012-03": [10, 20]}}, + [ + ("divide", "uprated_any_unit", "2012-03"), + ("calculate", "uprated_any_unit", "2013"), + ], + "uprated_any_unit", + "2015", + None, + None, + ), + auto_carry_over=True, +) +@example( + scenario=( + {"uprated_monthly": {"2011-06": [1001.3, 77.7]}}, + [ + ("calculate", "uprated_monthly", "2012-02"), + ("add", "uprated_monthly", "2013"), + ], + "uprated_monthly", + "2014-03", + None, + None, + ), + auto_carry_over=True, +) +# A branch input over a period the parent calculated before the fork. +@example( + scenario=( + {"uprated": {"2012": [1, 2]}}, + [("calculate", "uprated", "2013")], + "uprated", + "2015", + "reform", + ("2013", [300.0, 400.0]), + ), + auto_carry_over=True, +) +# Signed zero survives uprating from an input. +@example( + scenario=( + {"uprated": {"2012": [-0.0, 0.0]}}, + [("calculate", "uprated", "2013"), ("calculate", "uprated", "2014")], + "uprated", + "2015", + None, + None, + ), + auto_carry_over=False, +) +def test_uprating_depends_only_on_the_inputs(scenario, auto_carry_over): + inputs, requests, variable, period, branch, branch_input = scenario + system = SYSTEMS[auto_carry_over] + + built = simulation(system, inputs) + for kind, requested_variable, requested_period in requests: + request(built, kind, requested_variable, requested_period) + result = _read_from(built, branch, variable, branch_input).calculate( + variable, period + ) + + fresh = _read_from(simulation(system, inputs), branch, variable, branch_input) + message = f"{variable} {period} after {requests}, branch {branch} {branch_input}" + assert_bitwise_equal(result, fresh.calculate(variable, period), message) + + seen = {name: dict(values) for name, values in inputs.items()} + if branch_input is not None: + seen.setdefault(variable, {})[branch_input[0]] = branch_input[1] + assert_bitwise_equal(result, reference(system, seen, variable, period), message) diff --git a/tests/core/variables/test_variable_uprating.py b/tests/core/variables/test_variable_uprating.py index eb634a92b..7cf6cc056 100644 --- a/tests/core/variables/test_variable_uprating.py +++ b/tests/core/variables/test_variable_uprating.py @@ -211,16 +211,15 @@ def test_matches_explicitly_backdated_index(system, backdated_system, known, req @pytest.mark.parametrize("known", range(2008, 2016)) def test_result_does_not_depend_on_intermediate_years_computed(system, known): - """Path independence: computing every year in turn (each uprating from the - one before, which is then known) gives the same final value as uprating - straight from the known year.""" + """Path independence: computing every year in turn first gives exactly + the value uprated straight from the known year. Each year uprates from + the input, not from the year before (see test_uprating_order.py); chained + from the year before, float32 rounding compounded.""" stepwise = simulate(system, {known: VALUE}) for year in range(known + 1, 2022): stepwise.calculate("uprated_income", year) stepwise_result = float(stepwise.calculate("uprated_income", 2022)[0]) - assert stepwise_result == pytest.approx( - calculate(system, {known: VALUE}, 2022), rel=1e-5 - ) + assert stepwise_result == calculate(system, {known: VALUE}, 2022) @pytest.mark.parametrize( diff --git a/tests/fixtures/uprating_order.py b/tests/fixtures/uprating_order.py new file mode 100644 index 000000000..826fcc8b3 --- /dev/null +++ b/tests/fixtures/uprating_order.py @@ -0,0 +1,215 @@ +"""Variables, systems and a reference rule for the uprating order tests. + +Shared by ``tests/core/test_uprating_order.py`` (regressions) and +``tests/core/test_uprating_order_property.py`` (properties). The simulation +helpers are the auto-carry-over tests' (``tests/fixtures/carry_over.py``). + +The reference rule, for a variable with ``uprating`` and no formula result +for a period: an input stored for the period itself is read back as stored; +otherwise the latest earlier input stored in the variable's own unit, times +the ratio of the uprating index at the two period starts, cast to the +variable's type. With no such input, the period gets what auto-carry-over +gives it (the input stored for the latest-starting period that starts no +later than it, on a tie the one that ends last, preferring the variable's own +unit) or, without auto-carry-over or such an input, the default. Either way the value is masked +by the variable's ``defined_for`` at the period. +""" + +from __future__ import annotations + +import numpy as np + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem, entities +from policyengine_core.parameters import ParameterNode, get_parameter +from policyengine_core.simulations.simulation import _uprating_index_value +from policyengine_core.variables import Variable +from tests.fixtures.carry_over import COUNT, alone, request, simulation + +__all__ = [ + "COUNT", + "INDEX", + "UPRATING", + "alone", + "assert_bitwise_equal", + "build_system", + "index", + "reference", + "request", + "simulation", +] + +# 3.7% a year: ratios that are not exact in float32, so rounding at an +# intermediate year shows. +INDEX = {f"{year}-01-01": 100 * 1.037 ** (year - 2010) for year in range(2010, 2021)} +UPRATING = "uprating_order.index" + + +class uprated(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + uprating = UPRATING + label = "Uprated yearly input" + + +class uprated_count(Variable): + value_type = int + entity = entities.Person + definition_period = periods.YEAR + uprating = UPRATING + label = "Uprated yearly integer input" + + +class eligible(Variable): + value_type = bool + entity = entities.Person + definition_period = periods.YEAR + label = "Yearly boolean input that uprated_if_eligible is defined for" + + +class uprated_if_eligible(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + uprating = UPRATING + defined_for = "eligible" + label = "Uprated yearly input, defined for eligible people" + + +class uprated_with_default(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + uprating = UPRATING + default_value = 10 + label = "Uprated yearly input with a non-zero default" + + +class uprated_any_unit(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + uprating = UPRATING + set_input = None + label = "Uprated yearly input with no set_input helper, stored at any period" + + +class uprated_monthly(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + uprating = UPRATING + label = "Uprated monthly input" + + +VARIABLES = ( + uprated, + uprated_count, + eligible, + uprated_if_eligible, + uprated_with_default, + uprated_any_unit, + uprated_monthly, +) + + +def build_system(auto_carry_over: bool = True) -> CountryTaxBenefitSystem: + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = auto_carry_over + system.parameters.add_child( + "uprating_order", + ParameterNode("uprating_order", data={"index": {"values": INDEX}}), + ) + system.add_variables(*VARIABLES) + return system + + +def index(year: int) -> float: + return INDEX[f"{year}-01-01"] + + +def assert_bitwise_equal(actual, expected, message=""): + """Same dtype, shape and bytes: ``-0.0`` and ``0.0`` differ.""" + actual, expected = np.asarray(actual), np.asarray(expected) + assert actual.dtype == expected.dtype, ( + f"{message}: {actual.dtype} != {expected.dtype}" + ) + assert actual.shape == expected.shape, ( + f"{message}: {actual.shape} != {expected.shape}" + ) + assert actual.tobytes() == expected.tobytes(), f"{message}: {actual} != {expected}" + + +def _stored_inputs(system, inputs, variable): + """The inputs as ``set_input`` stores them, cast to the variable's type. + + The reference rule covers inputs in the variable's own unit, and inputs + in any unit to a variable with no ``set_input`` helper (stored as given). + A helper would split or copy an input in another unit across the + variable's own periods first, which the rule does not model. + """ + definition = system.get_variable(variable) + stored = {} + for period, values in inputs.get(variable, {}).items(): + period = periods.period(period) + if period.unit != definition.definition_period and definition.set_input: + raise ValueError( + f"{variable}: an input for {period} goes through set_input, " + "which the reference rule does not model" + ) + stored[period] = np.asarray(values, dtype=definition.dtype) + return stored + + +def _uprating_factor(system, definition, earlier, period): + parameter = get_parameter(system.parameters, definition.uprating) + value_then = _uprating_index_value(parameter, earlier.start) + value_now = _uprating_index_value(parameter, period.start) + if value_then is None or value_now is None or value_then == 0: + return 1 + return value_now / value_then + + +def reference(system, inputs, variable, period): + """The uprating and carry-over rules, computed from the inputs alone. + + ``period`` is in the variable's definition period. + """ + definition = system.get_variable(variable) + period = periods.period(period) + default = definition.default_array(COUNT) + stored = _stored_inputs(system, inputs, variable) + if period in stored: + # An input for the period itself is read back as stored: + # ``defined_for`` does not apply to it. + return stored[period] + earlier = [ + stored_period + for stored_period in stored + if stored_period.unit == definition.definition_period + and stored_period.start < period.start + ] + carried = [ + stored_period for stored_period in stored if stored_period.start <= period.start + ] + if definition.uprating is not None and earlier: + latest = max(earlier, key=lambda stored_period: stored_period.start) + value = stored[latest] * _uprating_factor(system, definition, latest, period) + elif system.auto_carry_over_input_variables and carried: + own_unit = [ + stored_period + for stored_period in carried + if stored_period.unit == definition.definition_period + ] + latest = max( + own_unit or carried, + key=lambda stored_period: (stored_period.start, stored_period.stop), + ) + value = stored[latest] + else: + value = default + if definition.defined_for is not None: + mask = reference(system, inputs, definition.defined_for, period) + value = np.where(mask, value, default) + return np.asarray(value).astype(definition.dtype) From 168a710a195a208ba2986fb79914da8001e86790 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 4 Oct 2026 12:03:47 -0400 Subject: [PATCH 2/5] Let a set_input helper's yearly input replace months already calculated set_input_dispatch_by_period and set_input_divide_by_period skipped any month that held a stored value, including one the simulation had calculated and cached as derived. Given a yearly input after a month of that year was calculated: - dispatch kept the calculated month and copied it into the later months as an input; - divide took it out of the year's total and divided the rest among the other months. The month read back as calculated, and since uprating now starts only from inputs, the next month was uprated from an earlier month (review finding P1 on #563, review/review_final_rebase.md). The helpers now treat only a stored input as present (get_stored_input): a derived month is replaced, as an input set for that month alone would replace it. Simulation.set_input also drops the fast-cache entries for the months a helper writes, which otherwise returned the calculated value. Executed probe (p1_helper_probe.py: stock and flow, carry-over on/off, root/branch/nested, three calculate-first patterns): 108 cases, mismatches vs a fresh simulation given the same input were 84 on master, 96 on #563 before this commit, 0 after. Tests: regressions for both helpers in test_uprating_order.py, and a derandomized Hypothesis property (300 examples) in test_uprating_order_property.py. A fresh simulation given the same inputs in the same order returns the same bytes for every month of the input's year and for another month, and every month of the year is an input. The fixtures add uprated_monthly_stock, with a month-by-month index. Ablations: the regressions fail 56/56 without either change, and 52/56 with the helper change alone. The property's generated examples alone fail without the change and with either half alone. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-uprating-order.fixed.md | 2 +- policyengine_core/holders/helpers.py | 33 ++++- policyengine_core/simulations/simulation.py | 15 +++ tests/core/test_uprating_order.py | 106 ++++++++++++++++ tests/core/test_uprating_order_property.py | 131 ++++++++++++++++++++ tests/fixtures/uprating_order.py | 30 ++++- 6 files changed, 309 insertions(+), 8 deletions(-) diff --git a/changelog.d/fix-uprating-order.fixed.md b/changelog.d/fix-uprating-order.fixed.md index ae873047d..ddb4df046 100644 --- a/changelog.d/fix-uprating-order.fixed.md +++ b/changelog.d/fix-uprating-order.fixed.md @@ -1 +1 @@ -Uprating now starts from the latest earlier input in the variable's own unit and skips periods the simulation calculated, so calculating intermediate periods first no longer compounds rounding or truncation, or carries an eligibility mask or a default into later uprated values. +Uprating now starts from the latest earlier input in the variable's own unit and skips periods the simulation calculated, and a yearly input given through a `set_input` helper now replaces months already calculated instead of keeping them, so calculating intermediate periods first no longer compounds rounding or truncation, or carries an eligibility mask, a default or a calculated month into later uprated values. diff --git a/policyengine_core/holders/helpers.py b/policyengine_core/holders/helpers.py index 790f136d1..664cd3850 100644 --- a/policyengine_core/holders/helpers.py +++ b/policyengine_core/holders/helpers.py @@ -22,6 +22,31 @@ def get_stored_array(holder: Holder, period: Period, branch_name: str) -> ArrayL return holder._get_array_from_storage(period, branch_name) +def get_stored_input(holder: Holder, period: Period, branch_name: str) -> ArrayLike: + """The input stored for ``period`` under ``branch_name`` itself, or + ``None``. + + A value the simulation calculated and cached there (see + ``Holder.put_in_cache``) is not an input, so the helpers below treat the + sub-period as empty and the new input replaces it, as an input set for + that sub-period alone would. Treating it as an input kept the calculated + value under the input, copied it into the later sub-periods (dispatch), + or took it out of the total (divide), so the result depended on what had + been calculated before the input was set. + """ + if holder._memory_storage.has(period, branch_name): + storage = holder._memory_storage + elif holder._disk_storage is not None and holder._disk_storage.has( + period, branch_name + ): + storage = holder._disk_storage + else: + return None + if storage.is_derived(period, branch_name): + return None + return storage.get(period, branch_name) + + def set_input_dispatch_by_period(holder: Holder, period: Period, array: ArrayLike): """ This function can be declared as a ``set_input`` attribute of a variable. @@ -46,11 +71,11 @@ def set_input_dispatch_by_period(holder: Holder, period: Period, array: ArrayLik after_instant = period.start.offset(period_size, period_unit) - # Cache the input data, skipping the existing cached months + # Cache the input data, skipping the months that already have an input branch_name = get_input_branch(holder) sub_period = period.start.period(cached_period_unit) while sub_period.start < after_instant: - existing_array = get_stored_array(holder, sub_period, branch_name) + existing_array = get_stored_input(holder, sub_period, branch_name) if existing_array is None: holder._set(sub_period, array, branch_name) else: @@ -91,7 +116,7 @@ def set_input_divide_by_period(holder: Holder, period: Period, array: ArrayLike) sub_period = period.start.period(cached_period_unit) sub_periods_count = 0 while sub_period.start < after_instant: - existing_array = get_stored_array(holder, sub_period, branch_name) + existing_array = get_stored_input(holder, sub_period, branch_name) if existing_array is not None: remaining_array -= existing_array else: @@ -103,7 +128,7 @@ def set_input_divide_by_period(holder: Holder, period: Period, array: ArrayLike) divided_array = remaining_array / sub_periods_count sub_period = period.start.period(cached_period_unit) while sub_period.start < after_instant: - if get_stored_array(holder, sub_period, branch_name) is None: + if get_stored_input(holder, sub_period, branch_name) is None: holder._set(sub_period, divided_array, branch_name) sub_period = sub_period.offset(1) elif not (remaining_array == 0).all(): diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index c5f977f9b..9aac75e2c 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1544,6 +1544,21 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non _fast_cache = getattr(self, "_fast_cache", None) if _fast_cache is not None: _fast_cache.pop((variable_name, period), None) + if variable.set_input and period.unit != variable.definition_period: + # The helper wrote the input's sub-periods, replacing any + # value calculated there, so what ``calculate`` returned for + # them is stale too. (``_end_order``, not ``stop``: ``stop`` + # raises for a period that ends after 9999-12-31.) + stale = [ + key + for key in _fast_cache + if key[0] == variable_name + and isinstance(key[1], Period) + and period.start <= key[1].start + and _end_order(key[1]) <= _end_order(period) + ] + for key in stale: + del _fast_cache[key] def get_variable_population(self, variable_name: str) -> Population: variable = self.tax_benefit_system.get_variable( diff --git a/tests/core/test_uprating_order.py b/tests/core/test_uprating_order.py index d67ffebee..de660c43f 100644 --- a/tests/core/test_uprating_order.py +++ b/tests/core/test_uprating_order.py @@ -26,6 +26,12 @@ and the order-independence property in ``test_uprating_order_property.py``. Each regression here compares with a fresh simulation that calculates only the period in question, byte for byte. + +A yearly input for a monthly variable is stored month by month by its +``set_input`` helper. The helpers took a month the simulation had calculated +for an input and kept it, so the uprating rule above skipped it. They now +replace it, and ``Simulation.set_input`` drops what ``calculate`` returned +for those months (the regressions at the end of this module). """ from __future__ import annotations @@ -38,6 +44,7 @@ from policyengine_core.reforms import Reform from policyengine_core.variables import Variable from tests.fixtures.uprating_order import ( + MONTHLY_INDEX, UPRATING, alone, assert_bitwise_equal, @@ -283,3 +290,102 @@ def test_uprated_value_is_cached_as_derived(system): holder = built.get_holder("uprated") assert holder.is_derived(periods.period("2013")) assert not holder.is_derived(periods.period("2012")) + + +# A yearly input for a monthly variable goes through its ``set_input`` helper, +# which stores it in each month of the year: copied to each month for a stock +# (``set_input_dispatch_by_period``), divided between them for a flow +# (``set_input_divide_by_period``). A month the simulation calculated before +# the input was set is not an input, so the helper replaces it as it fills a +# month nothing is stored for. Before the fix, the helpers treated it as one: +# dispatch kept it and copied it into the later months as an input, and divide +# took it out of the year's total. The month also read back as calculated, and +# a later month was uprated from an earlier input month, or from the copies. + +HELPER_INPUTS = { + "uprated_monthly_stock": ([7.0, 9.0], [7.0, 9.0]), + "uprated_monthly": ([84.0, 108.0], [7.0, 9.0]), +} +CALCULATED_FIRST = [["2012-12"], ["2012-06"], ["2012-06", "2012-12"], ["2012-01"]] + + +def _on(built, branch): + if branch in ("reform", "nested"): + built = built.get_branch("reform") + if branch == "nested": + built = built.get_branch("nested") + return built + + +@pytest.mark.parametrize("branch", [None, "reform", "nested"]) +@pytest.mark.parametrize("calculated", CALCULATED_FIRST, ids="+".join) +@pytest.mark.parametrize("variable", list(HELPER_INPUTS), ids=["dispatch", "divide"]) +def test_helper_input_replaces_months_calculated_before_it( + system, variable, calculated, branch +): + yearly, monthly = HELPER_INPUTS[variable] + built = _on(simulation(system, {}), branch) + for month in calculated: + built.calculate(variable, month) + built.set_input(variable, "2012", np.array(yearly)) + + fresh = _on(simulation(system, {}), branch) + fresh.set_input(variable, "2012", np.array(yearly)) + holder = built.get_holder(variable) + for month in ("2012-01", "2012-06", "2012-12"): + result = built.calculate(variable, month) + np.testing.assert_array_equal(result, monthly) + assert_bitwise_equal(result, fresh.calculate(variable, month), month) + assert not holder.is_derived(periods.period(month), built.branch_name) + for month in ("2013-01", "2014-03"): + assert_bitwise_equal( + built.calculate(variable, month), fresh.calculate(variable, month), month + ) + + +@pytest.mark.parametrize("variable", list(HELPER_INPUTS), ids=["dispatch", "divide"]) +def test_helper_input_is_uprated_from_its_last_month(system, variable): + """The month after the year is uprated from December's input, not from an + earlier month the helper filled while December held a calculated value.""" + if variable == "uprated_monthly_stock": + ratio = MONTHLY_INDEX["2013-01-01"] / MONTHLY_INDEX["2012-12-01"] + else: + ratio = index(2013) / index(2012) + yearly, monthly = HELPER_INPUTS[variable] + built = simulation(system, {}) + built.calculate(variable, "2012-12") + built.set_input(variable, "2012", np.array(yearly)) + np.testing.assert_allclose( + built.calculate(variable, "2013-01"), np.array(monthly) * ratio, rtol=1e-6 + ) + + +@pytest.mark.parametrize( + "variable, expected", + [ + # Dispatch copies an input it meets into the months after it. + ("uprated_monthly_stock", {"2012-02": [7, 9], "2012-06": [1, 2]}), + # Divide shares what the inputs it meets leave of the total. + ( + "uprated_monthly", + {"2012-02": [83 / 11, 106 / 11], "2012-06": [83 / 11, 106 / 11]}, + ), + ], + ids=["dispatch", "divide"], +) +def test_helper_input_keeps_an_input_already_set_in_its_year( + system, variable, expected +): + """Guard against over-correction: an input for a month of the year is + still kept, whether or not other months were calculated.""" + yearly, _ = HELPER_INPUTS[variable] + for calculated in ([], ["2012-06"], ["2012-03", "2012-06"]): + built = simulation(system, {variable: {"2012-03": [1.0, 2.0]}}) + for month in calculated: + built.calculate(variable, month) + built.set_input(variable, "2012", np.array(yearly)) + np.testing.assert_array_equal(built.calculate(variable, "2012-03"), [1, 2]) + for month, values in expected.items(): + np.testing.assert_allclose( + built.calculate(variable, month), values, rtol=1e-6 + ) diff --git a/tests/core/test_uprating_order_property.py b/tests/core/test_uprating_order_property.py index 5f015cc3c..1aa0349da 100644 --- a/tests/core/test_uprating_order_property.py +++ b/tests/core/test_uprating_order_property.py @@ -20,6 +20,10 @@ cached for the branch whatever its inputs are, which is not what this module tests. * Targets are in the variable's definition unit. + +A second property, ``test_helper_input_replaces_values_calculated_in_its_year``, +covers yearly inputs given through a ``set_input`` helper after calculations +in that year; its domain is stated above it. """ from __future__ import annotations @@ -297,3 +301,130 @@ def test_uprating_depends_only_on_the_inputs(scenario, auto_carry_over): if branch_input is not None: seen.setdefault(variable, {})[branch_input[0]] = branch_input[1] assert_bitwise_equal(result, reference(system, seen, variable, period), message) + + +# Inputs given through a ``set_input`` helper after calculations. +# +# A yearly input for a monthly variable is stored in each month of the year by +# its helper: copied for a stock, divided for a flow. The property: whatever +# was calculated in that year before the input was set, the simulation then +# returns, byte for byte, what a fresh simulation given the same inputs in the +# same order returns, for every month of the year and for a later or earlier +# month, from the simulation, a branch or a nested branch. Every month of the +# year then holds an input. +# +# The domain: earlier inputs are months, set on the simulation first; the +# yearly input and the requests before it are on the simulation or branch +# read from, and those requests are inside the input's year (plain +# calculations of its months, or ADD over it). A value calculated outside the +# year before the input stays cached whatever the input is, which is not what +# this module tests. + +HELPER_VARIABLES = ["uprated_monthly_stock", "uprated_monthly"] +HELPER_YEARS = ["2011", "2012", "2013"] + + +@st.composite +def helper_scenarios(draw): + variable = draw(st.sampled_from(HELPER_VARIABLES)) + earlier_inputs = draw( + st.dictionaries( + st.sampled_from(INPUT_MONTHS), + st.lists(FLOATS, min_size=COUNT, max_size=COUNT), + max_size=3, + ) + ) + year = draw(st.sampled_from(HELPER_YEARS)) + yearly = draw(st.lists(FLOATS, min_size=COUNT, max_size=COUNT)) + months_of_year = [f"{year}-{month:02d}" for month in range(1, 13)] + requests = draw( + st.lists( + st.one_of( + st.tuples(st.just("calculate"), st.sampled_from(months_of_year)), + st.tuples(st.just("add"), st.just(year)), + ), + max_size=5, + ) + ) + branch = draw(st.sampled_from([None, "reform", "nested"])) + target = draw(st.sampled_from(REQUEST_MONTHS)) + return variable, earlier_inputs, year, yearly, requests, branch, target + + +def _helper_simulation(system, variable, earlier_inputs, branch): + built = simulation(system, {variable: earlier_inputs} if earlier_inputs else {}) + return _read_from(built, branch, variable, None) + + +@settings( + max_examples=300, + deadline=None, + derandomize=True, + suppress_health_check=[HealthCheck.too_slow], +) +@given(scenario=helper_scenarios(), auto_carry_over=st.booleans()) +# December calculated, then the year's input (review finding P1). +@example( + scenario=( + "uprated_monthly_stock", + {}, + "2012", + [7.0, 9.0], + [("calculate", "2012-12")], + None, + "2013-01", + ), + auto_carry_over=True, +) +# A month calculated mid-year: divide took it out of the year's total. +@example( + scenario=( + "uprated_monthly", + {"2011-06": [5.0, 6.0]}, + "2012", + [84.0, 108.0], + [("calculate", "2012-06")], + "nested", + "2012-07", + ), + auto_carry_over=False, +) +# ADD over the year, and an input already set in it. +@example( + scenario=( + "uprated_monthly_stock", + {"2012-03": [1.0, 2.0]}, + "2012", + [7.0, 9.0], + [("add", "2012")], + "reform", + "2013-02", + ), + auto_carry_over=True, +) +def test_helper_input_replaces_values_calculated_in_its_year(scenario, auto_carry_over): + variable, earlier_inputs, year, yearly, requests, branch, target = scenario + system = SYSTEMS[auto_carry_over] + + built = _helper_simulation(system, variable, earlier_inputs, branch) + for kind, requested in requests: + request(built, kind, variable, requested) + built.set_input(variable, year, yearly) + + fresh = _helper_simulation(system, variable, earlier_inputs, branch) + fresh.set_input(variable, year, yearly) + + message = f"{variable} after {requests}, then {year} = {yearly}, branch {branch}" + holder = built.get_holder(variable) + for month in [f"{year}-{month:02d}" for month in range(1, 13)]: + assert not holder.is_derived(periods.period(month), built.branch_name), month + assert_bitwise_equal( + built.calculate(variable, month), + fresh.calculate(variable, month), + f"{message}: {month}", + ) + assert_bitwise_equal( + built.calculate(variable, target), + fresh.calculate(variable, target), + f"{message}: {target}", + ) diff --git a/tests/fixtures/uprating_order.py b/tests/fixtures/uprating_order.py index 826fcc8b3..0f2a66fc9 100644 --- a/tests/fixtures/uprating_order.py +++ b/tests/fixtures/uprating_order.py @@ -23,12 +23,13 @@ from policyengine_core.country_template import CountryTaxBenefitSystem, entities from policyengine_core.parameters import ParameterNode, get_parameter from policyengine_core.simulations.simulation import _uprating_index_value -from policyengine_core.variables import Variable +from policyengine_core.variables import QuantityType, Variable from tests.fixtures.carry_over import COUNT, alone, request, simulation __all__ = [ "COUNT", "INDEX", + "MONTHLY_INDEX", "UPRATING", "alone", "assert_bitwise_equal", @@ -43,6 +44,13 @@ # intermediate year shows. INDEX = {f"{year}-01-01": 100 * 1.037 ** (year - 2010) for year in range(2010, 2021)} UPRATING = "uprating_order.index" +# 0.3% a month, so that which month of a year a value is uprated from shows. +MONTHLY_INDEX = { + f"{year}-{month:02d}-01": 100 * 1.003 ** ((year - 2010) * 12 + month - 1) + for year in range(2010, 2021) + for month in range(1, 13) +} +MONTHLY_UPRATING = "uprating_order.monthly_index" class uprated(Variable): @@ -100,7 +108,16 @@ class uprated_monthly(Variable): entity = entities.Person definition_period = periods.MONTH uprating = UPRATING - label = "Uprated monthly input" + label = "Uprated monthly input (a flow: set_input_divide_by_period)" + + +class uprated_monthly_stock(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + quantity_type = QuantityType.STOCK + uprating = MONTHLY_UPRATING + label = "Uprated monthly input (a stock: set_input_dispatch_by_period)" VARIABLES = ( @@ -111,6 +128,7 @@ class uprated_monthly(Variable): uprated_with_default, uprated_any_unit, uprated_monthly, + uprated_monthly_stock, ) @@ -119,7 +137,13 @@ def build_system(auto_carry_over: bool = True) -> CountryTaxBenefitSystem: system.auto_carry_over_input_variables = auto_carry_over system.parameters.add_child( "uprating_order", - ParameterNode("uprating_order", data={"index": {"values": INDEX}}), + ParameterNode( + "uprating_order", + data={ + "index": {"values": INDEX}, + "monthly_index": {"values": MONTHLY_INDEX}, + }, + ), ) system.add_variables(*VARIABLES) return system From 1ee63a80aafdd8525e455f3de5f18fbaeac0a3e4 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 4 Oct 2026 13:23:12 -0400 Subject: [PATCH 3/5] Never write over an on-disk file a clone also reads OnDiskStorage.clone shares its directory and every file stored so far, and put wrote each value to its key's path. So a write through one storage changed what the other read. The clone docstring noted it. - Simulation.clone() keeps the branch name. A set_input on the clone wrote over the source's file, and with the previous commit so did a yearly helper input replacing a calculated month (review of 168a710a, review/review_final_rebase2.md): the source's December became the clone's input. - The source writing a period again after cloning (delete, then calculate) changed the clone's value (round-1 review, finding P3). A storage now writes over a file only if it wrote it and has not shared it since (_own_files). Otherwise, if any storage in its clone family wrote or stored that file (_family_files, one set per family), the value goes to a new file of its own under replaced/, which restore() never reads. Storages not cloned from one another, such as restore() and a separate writer, still write over each other's files as before. Pickling counts as cloning. A state pickled before this change treats every file it stores as shared. Tests: tests/core/test_on_disk_storage_clones.py covers the review's sequence for both helpers and both carry-over modes, own-unit set_input on a clone, the source rewriting a period, and storage-level cases: - clones, a clone of a clone and in-place writes of own files; - independent storages over one directory, and restore() reading only usual names; - a restored storage's clone, deepcopy/pickle, and old pickles. On 168a710a, 14 of the 15 fail; the one that passes is the independent-storage guard. The helper property in test_uprating_order_property.py now also runs on disk, and with a clone taken before the input whose stored months must not change. With explicit rows removed and no shrinking, it fails on 168a710a. Co-Authored-By: Claude Opus 5.5 --- .../fix-disk-storage-clone-writes.fixed.md | 1 + .../data_storage/on_disk_storage.py | 69 ++++- tests/core/test_on_disk_storage_clones.py | 236 ++++++++++++++++++ tests/core/test_uprating_order_property.py | 95 +++++-- 4 files changed, 381 insertions(+), 20 deletions(-) create mode 100644 changelog.d/fix-disk-storage-clone-writes.fixed.md create mode 100644 tests/core/test_on_disk_storage_clones.py diff --git a/changelog.d/fix-disk-storage-clone-writes.fixed.md b/changelog.d/fix-disk-storage-clone-writes.fixed.md new file mode 100644 index 000000000..44d2e3eb4 --- /dev/null +++ b/changelog.d/fix-disk-storage-clone-writes.fixed.md @@ -0,0 +1 @@ +Simulations and storages cloned from one another no longer write over each other's values in on-disk storage: a value stored for a key whose file a clone also reads goes to a new file of its own. diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index db3e2cad0..24501bf33 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -1,5 +1,6 @@ import os import shutil +import tempfile import numpy from numpy.typing import ArrayLike @@ -8,6 +9,12 @@ from policyengine_core.enums import EnumArray from policyengine_core.periods import Period +# Subdirectory of a storage directory for the files ``put`` writes when a +# key's usual file is shared with a clone (see ``OnDiskStorage._path_to_write``). +# ``restore`` reads only the files directly in the directory, so it never +# reads these. +REPLACEMENTS_DIR = "replaced" + class OnDiskStorage: """ @@ -25,14 +32,32 @@ def __init__( # File keys stored with ``put(..., derived=True)``; see # ``InMemoryStorage``. self._derived = set() + # Paths of the files this storage wrote and has not shared with a + # clone since: the only shared-directory files ``put`` writes over + # (see ``_path_to_write``). + self._own_files = set() + # Paths of every file this storage, the storage it was cloned from or + # any other clone of either wrote or stored when cloned: one set, + # shared by all of them. + self._family_files = set() self.is_eternal = is_eternal self.preserve_storage_dir = preserve_storage_dir self.storage_dir = storage_dir + def __getstate__(self) -> dict: + # Whatever this state is read into reads the same files, so from now + # on this storage writes over none of them, as after ``clone``. + self._own_files = set() + return self.__dict__.copy() + def __setstate__(self, state: dict) -> None: # A storage pickled before derived marks existed has none: its values # count as inputs. state.setdefault("_derived", set()) + # Nor did it record the files it wrote: it writes over none it stores. + state.setdefault("_own_files", set()) + if "_family_files" not in state: + state["_family_files"] = set(state.get("_files", {}).values()) self.__dict__.update(state) def clone(self) -> "OnDiskStorage": @@ -40,10 +65,10 @@ def clone(self) -> "OnDiskStorage": The file and enum mappings are copied so deleting or rewiring entries through the clone does not mutate the source storage. The underlying - ``.npy`` files remain shared: writing the same ``{branch}_{period}`` - key from two views targets the same path and can overwrite the file. - Clones retain the original cleanup owner so the shared directory stays - alive, but never own cleanup themselves. + ``.npy`` files are shared, so neither storage writes over them: a later + ``put`` of one of their keys, through either storage, writes a new + file (see ``_path_to_write``). Clones retain the original cleanup owner + so the shared directory stays alive, but never own cleanup themselves. """ clone = OnDiskStorage( self.storage_dir, @@ -54,6 +79,11 @@ def clone(self) -> "OnDiskStorage": clone._enums = self._enums.copy() clone._derived = set(self._derived) clone._storage_dir_owner = getattr(self, "_storage_dir_owner", self) + # Both storages now read every file stored so far, including any this + # family did not write (read back by ``restore``, say). + self._family_files.update(self._files.values()) + clone._family_files = self._family_files + self._own_files = set() return clone def _decode_file(self, file: str) -> ArrayLike: @@ -102,17 +132,45 @@ def put( period = periods.period(period) filename = f"{branch_name}_{period}" - path = os.path.join(self.storage_dir, filename) + ".npy" + path = self._path_to_write(filename) if isinstance(value, EnumArray): self._enums[path] = value.possible_values value = value.view(numpy.ndarray) numpy.save(path, value) self._files[filename] = path + self._own_files.add(path) + self._family_files.add(path) if derived: self._derived.add(filename) else: self._derived.discard(filename) + def _path_to_write(self, filename: str) -> str: + """The path ``put`` writes the value for key ``filename`` to. + + Storages cloned from one another share a directory and the files + stored before cloning, and each can store the same key afterwards, + which names the same path. Writing over a file another of them reads + would change that storage's value. So ``put`` writes over one of + their files only if this storage wrote it and has not shared it since; + otherwise the value goes to a new file of its own, in + ``REPLACEMENTS_DIR``. A file only storages outside this family wrote + is written over, as before. + """ + current = self._files.get(filename) + if current is not None and current in self._own_files: + return current + path = os.path.join(self.storage_dir, filename) + ".npy" + if path in self._own_files or path not in self._family_files: + return path + directory = os.path.join(self.storage_dir, REPLACEMENTS_DIR) + os.makedirs(directory, exist_ok=True) + descriptor, path = tempfile.mkstemp( + prefix=f"{filename}.", suffix=".npy", dir=directory + ) + os.close(descriptor) + return path + def delete(self, period: Period = None, branch_name: str = "default") -> None: if period is None: # Only wipe files belonging to the requested branch (previously @@ -152,6 +210,7 @@ def restore(self) -> None: self._files = files = {} # Files read back from a directory carry no derived marks. self._derived = set() + self._own_files = set() # Restore self._files from content of storage_dir. for filename in os.listdir(self.storage_dir): if not filename.endswith(".npy"): diff --git a/tests/core/test_on_disk_storage_clones.py b/tests/core/test_on_disk_storage_clones.py new file mode 100644 index 000000000..c4c425dea --- /dev/null +++ b/tests/core/test_on_disk_storage_clones.py @@ -0,0 +1,236 @@ +"""Storages cloned from one another never write over each other's files. + +A cloned ``OnDiskStorage`` shares its source's directory and the files stored +before cloning, and the two can each store the same key later, which names +the same path. ``put`` wrote every value to its key's path, so a write +through one storage changed what the other read: + +* ``Simulation.clone()`` keeps the branch name, so a ``set_input`` on the + clone wrote over the source's file for that period. A yearly input given + through a ``set_input`` helper, which now replaces a month the simulation + calculated, did the same. +* The source writing a period again after cloning (say, calculating it after + deleting it) changed the clone's value too. + +``put`` now writes over one of the family's files only if the storage wrote +it and has not shared it since; otherwise it writes a new file of its own, +in a subdirectory ``restore`` does not read. Storages that were not cloned +from one another write over each other's files, as before. +""" + +from __future__ import annotations + +import copy +import glob +import os +import pickle + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.data_storage import OnDiskStorage +from policyengine_core.experimental import MemoryConfig +from tests.fixtures.uprating_order import build_system, simulation + +SYSTEMS = {True: build_system(True), False: build_system(False)} +HELPER_INPUTS = { + "uprated_monthly_stock": [7.0, 9.0], + "uprated_monthly": [84.0, 108.0], +} + + +@pytest.fixture(params=[True, False], ids=["carry_over", "no_carry_over"]) +def system(request): + return SYSTEMS[request.param] + + +def _on_disk(built, variable): + built.memory_config = MemoryConfig(max_memory_occupation=0) + holder = built.get_holder(variable) + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + return holder + + +def _stored(holder, period): + """The stored value, read past the simulation's fast cache.""" + return holder.get_array(periods.period(period)) + + +@pytest.mark.parametrize("variable", list(HELPER_INPUTS), ids=["dispatch", "divide"]) +def test_helper_input_on_a_clone_leaves_the_source_unchanged(system, variable): + parent = simulation(system, {}) + holder = _on_disk(parent, variable) + calculated = parent.calculate(variable, "2012-12").copy() + assert holder._disk_storage.is_derived(periods.period("2012-12")) + + child = parent.clone() + child.set_input(variable, "2012", np.array(HELPER_INPUTS[variable])) + + np.testing.assert_array_equal(child.calculate(variable, "2012-12"), [7, 9]) + np.testing.assert_array_equal(_stored(holder, "2012-12"), calculated) + assert holder.is_derived(periods.period("2012-12")) + + +def test_set_input_on_a_clone_leaves_the_source_unchanged(): + parent = simulation(SYSTEMS[True], {}) + holder = _on_disk(parent, "uprated") + parent.set_input("uprated", "2012", np.array([1.0, 2.0])) + + child = parent.clone() + child.set_input("uprated", "2012", np.array([3.0, 4.0])) + + np.testing.assert_array_equal(_stored(holder, "2012"), [1, 2]) + np.testing.assert_array_equal(_stored(child.get_holder("uprated"), "2012"), [3, 4]) + + +def test_source_writing_a_period_again_leaves_the_clone_unchanged(): + parent = simulation(SYSTEMS[True], {}) + _on_disk(parent, "uprated") + parent.set_input("uprated", "2011", np.array([1001.0, 77.0])) + parent.set_input("uprated", "2012", np.array([2222.0, 222.0])) + branch = parent.get_branch("branch") + expected = branch.calculate("uprated", "2015").copy() + + parent.delete_arrays("uprated", "2012") + parent.calculate("uprated", "2012") + + np.testing.assert_array_equal( + _stored(branch.get_holder("uprated"), "2012"), [2222, 222] + ) + branch.delete_arrays("uprated", "2015") + np.testing.assert_array_equal(branch.calculate("uprated", "2015"), expected) + + +# ----- Storage level ----- # + + +@pytest.fixture +def storage(tmp_path): + return OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + + +def _npy_files(storage): + return sorted( + glob.glob(os.path.join(storage.storage_dir, "**", "*.npy"), recursive=True) + ) + + +def test_storages_cloned_from_one_another_keep_their_own_values(storage): + storage.put(np.array([1.0]), "2012") + clone = storage.clone() + + storage.put(np.array([2.0]), "2012") + np.testing.assert_array_equal(clone.get("2012"), [1]) + clone.put(np.array([3.0]), "2012") + np.testing.assert_array_equal(storage.get("2012"), [2]) + np.testing.assert_array_equal(clone.get("2012"), [3]) + + # A key both store after cloning. + storage.put(np.array([4.0]), "2013") + clone.put(np.array([5.0]), "2013") + np.testing.assert_array_equal(storage.get("2013"), [4]) + np.testing.assert_array_equal(clone.get("2013"), [5]) + + +def test_a_clone_of_a_clone_keeps_its_own_values(storage): + storage.put(np.array([1.0]), "2012") + clone = storage.clone() + clone.put(np.array([2.0]), "2012") + grandchild = clone.clone() + + clone.put(np.array([3.0]), "2012") + grandchild.put(np.array([4.0]), "2012") + storage.put(np.array([5.0]), "2012") + + np.testing.assert_array_equal(storage.get("2012"), [5]) + np.testing.assert_array_equal(clone.get("2012"), [3]) + np.testing.assert_array_equal(grandchild.get("2012"), [4]) + + +def test_a_storage_writes_over_its_own_files(storage): + """No file per write: a value the storage alone reads is replaced in + place, including one it moved to a new file of its own.""" + storage.put(np.array([1.0]), "2012") + storage.put(np.array([2.0]), "2012", derived=True) + assert len(_npy_files(storage)) == 1 + np.testing.assert_array_equal(storage.get("2012"), [2]) + assert storage.is_derived(periods.period("2012")) + + storage.clone() + storage.put(np.array([3.0]), "2012") + storage.put(np.array([4.0]), "2012") + assert len(_npy_files(storage)) == 2 + np.testing.assert_array_equal(storage.get("2012"), [4]) + assert not storage.is_derived(periods.period("2012")) + + +def test_storages_not_cloned_from_one_another_write_over_each_others_files( + storage, +): + """Two storages made for one directory, as ``restore`` and a separate + writer are, still name the same file for a key.""" + storage.put(np.array([1.0]), "2012") + writer = OnDiskStorage(storage.storage_dir, preserve_storage_dir=True) + writer.put(np.array([2.0]), "2012") + np.testing.assert_array_equal(storage.get("2012"), [2]) + assert len(_npy_files(storage)) == 1 + + +def test_restore_reads_only_the_files_named_for_their_keys(storage): + storage.put(np.array([1.0]), "2012") + storage.put(np.array([2.0]), "month:2012-03:3") + clone = storage.clone() + clone.put(np.array([3.0]), "2012") + assert len(_npy_files(storage)) == 3 + + restored = OnDiskStorage(storage.storage_dir, preserve_storage_dir=True) + restored.restore() + + assert sorted(map(str, restored.get_known_periods())) == [ + "2012", + "month:2012-03:3", + ] + np.testing.assert_array_equal(restored.get("2012"), [1]) + + +def test_a_restored_storage_and_its_clone_keep_their_own_values(storage): + storage.put(np.array([1.0]), "2012") + restored = OnDiskStorage(storage.storage_dir, preserve_storage_dir=True) + restored.restore() + clone = restored.clone() + + clone.put(np.array([2.0]), "2012") + + np.testing.assert_array_equal(restored.get("2012"), [1]) + np.testing.assert_array_equal(clone.get("2012"), [2]) + + +@pytest.mark.parametrize( + "copier", [copy.deepcopy, lambda s: pickle.loads(pickle.dumps(s))] +) +def test_a_copied_storage_and_its_source_keep_their_own_values(storage, copier): + storage.put(np.array([1.0]), "2012") + copied = copier(storage) + + copied.put(np.array([2.0]), "2012") + np.testing.assert_array_equal(storage.get("2012"), [1]) + storage.put(np.array([3.0]), "2012") + np.testing.assert_array_equal(copied.get("2012"), [2]) + np.testing.assert_array_equal(storage.get("2012"), [3]) + + +def test_a_storage_unpickled_from_before_writes_over_none_of_its_files(storage): + """A storage pickled before storages recorded the files they wrote may + share them with another, so it writes a new file for each.""" + storage.put(np.array([1.0]), "2012") + state = pickle.loads(pickle.dumps(storage.__dict__)) + del state["_own_files"], state["_family_files"] + old = OnDiskStorage.__new__(OnDiskStorage) + old.__setstate__(state) + + old.put(np.array([2.0]), "2012") + + np.testing.assert_array_equal(storage.get("2012"), [1]) + np.testing.assert_array_equal(old.get("2012"), [2]) diff --git a/tests/core/test_uprating_order_property.py b/tests/core/test_uprating_order_property.py index 1aa0349da..d829f6768 100644 --- a/tests/core/test_uprating_order_property.py +++ b/tests/core/test_uprating_order_property.py @@ -35,8 +35,10 @@ from hypothesis import HealthCheck, example, given, settings # noqa: E402 from hypothesis import strategies as st # noqa: E402 +import numpy as np # noqa: E402 from policyengine_core import periods # noqa: E402 +from policyengine_core.experimental import MemoryConfig # noqa: E402 from tests.fixtures.uprating_order import ( # noqa: E402 COUNT, assert_bitwise_equal, @@ -311,14 +313,16 @@ def test_uprating_depends_only_on_the_inputs(scenario, auto_carry_over): # returns, byte for byte, what a fresh simulation given the same inputs in the # same order returns, for every month of the year and for a later or earlier # month, from the simulation, a branch or a nested branch. Every month of the -# year then holds an input. +# year then holds an input. And the input changes nothing a clone of the +# simulation, taken just before it, has stored: a clone keeps the branch +# name, and on disk it shares the files stored before cloning. # -# The domain: earlier inputs are months, set on the simulation first; the -# yearly input and the requests before it are on the simulation or branch -# read from, and those requests are inside the input's year (plain -# calculations of its months, or ADD over it). A value calculated outside the -# year before the input stays cached whatever the input is, which is not what -# this module tests. +# The domain: the variable is stored in memory or on disk. Earlier inputs are +# months, set on the simulation first; the yearly input and the requests +# before it are on the simulation or branch read from, and those requests are +# inside the input's year (plain calculations of its months, or ADD over it). +# A value calculated outside the year before the input stays cached whatever +# the input is, which is not what this module tests. HELPER_VARIABLES = ["uprated_monthly_stock", "uprated_monthly"] HELPER_YEARS = ["2011", "2012", "2013"] @@ -351,18 +355,40 @@ def helper_scenarios(draw): return variable, earlier_inputs, year, yearly, requests, branch, target -def _helper_simulation(system, variable, earlier_inputs, branch): - built = simulation(system, {variable: earlier_inputs} if earlier_inputs else {}) +def _helper_simulation(system, variable, earlier_inputs, branch, on_disk): + built = simulation(system, {}) + if on_disk: + built.memory_config = MemoryConfig(max_memory_occupation=0) + holder = built.get_holder(variable) + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + for period, values in earlier_inputs.items(): + built.set_input(variable, period, np.array(values)) return _read_from(built, branch, variable, None) +def _stored_months(built, variable, months): + holder = built.get_holder(variable) + stored = {} + for month in months: + value = holder.get_array(periods.period(month), built.branch_name) + stored[month] = None if value is None else value.copy() + return stored + + +@pytest.mark.filterwarnings("ignore:Memory configuration is a feature") @settings( max_examples=300, deadline=None, derandomize=True, suppress_health_check=[HealthCheck.too_slow], ) -@given(scenario=helper_scenarios(), auto_carry_over=st.booleans()) +@given( + scenario=helper_scenarios(), + auto_carry_over=st.booleans(), + on_disk=st.booleans(), + with_clone=st.booleans(), +) # December calculated, then the year's input (review finding P1). @example( scenario=( @@ -375,6 +401,24 @@ def _helper_simulation(system, variable, earlier_inputs, branch): "2013-01", ), auto_carry_over=True, + on_disk=False, + with_clone=False, +) +# The same on disk, with a clone: the input replaced the clone's December too +# (review finding on 168a710a). +@example( + scenario=( + "uprated_monthly_stock", + {}, + "2012", + [7.0, 9.0], + [("calculate", "2012-12")], + None, + "2013-01", + ), + auto_carry_over=True, + on_disk=True, + with_clone=True, ) # A month calculated mid-year: divide took it out of the year's total. @example( @@ -388,6 +432,8 @@ def _helper_simulation(system, variable, earlier_inputs, branch): "2012-07", ), auto_carry_over=False, + on_disk=True, + with_clone=True, ) # ADD over the year, and an input already set in it. @example( @@ -401,22 +447,41 @@ def _helper_simulation(system, variable, earlier_inputs, branch): "2013-02", ), auto_carry_over=True, + on_disk=False, + with_clone=True, ) -def test_helper_input_replaces_values_calculated_in_its_year(scenario, auto_carry_over): +def test_helper_input_replaces_values_calculated_in_its_year( + scenario, auto_carry_over, on_disk, with_clone +): variable, earlier_inputs, year, yearly, requests, branch, target = scenario system = SYSTEMS[auto_carry_over] + months = [f"{year}-{month:02d}" for month in range(1, 13)] - built = _helper_simulation(system, variable, earlier_inputs, branch) + built = _helper_simulation(system, variable, earlier_inputs, branch, on_disk) for kind, requested in requests: request(built, kind, variable, requested) + if with_clone: + clone = built.clone() + stored_by_clone = _stored_months(clone, variable, months) built.set_input(variable, year, yearly) - fresh = _helper_simulation(system, variable, earlier_inputs, branch) + fresh = _helper_simulation(system, variable, earlier_inputs, branch, on_disk) fresh.set_input(variable, year, yearly) - message = f"{variable} after {requests}, then {year} = {yearly}, branch {branch}" + message = ( + f"{variable} after {requests}, then {year} = {yearly}, branch {branch}, " + f"on disk {on_disk}" + ) + if with_clone: + for month, value in _stored_months(clone, variable, months).items(): + if stored_by_clone[month] is None: + assert value is None, f"{message}: clone's {month}" + else: + assert_bitwise_equal( + value, stored_by_clone[month], f"{message}: clone's {month}" + ) holder = built.get_holder(variable) - for month in [f"{year}-{month:02d}" for month in range(1, 13)]: + for month in months: assert not holder.is_derived(periods.period(month), built.branch_name), month assert_bitwise_equal( built.calculate(variable, month), From c5c7bf4c40bbc753b8839e93e41a60602065de64 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 4 Oct 2026 15:25:12 -0400 Subject: [PATCH 4/5] Give each clone its own input record; keep a storage's own files across restore Round-3 review of #563 (review/review_final_rebase3.md), two blocking findings: 1. Simulation.clone() shared _user_input_keys with its source. A yearly helper input on the clone now replaces a month the source had calculated, and it recorded that month's key in the shared set. The source's apply_reform then kept its own calculated month as an input and uprated from it. The shared record predates this PR, but the helper change made it reachable this way. Each clone now gets a copy of the record. Nothing read the record across simulations: apply_reform and the input export look up keys in the simulation's own holders. 2. restore() cleared the storage's own files, so a storage that was never cloned moved its next write to replaced/, where a separate reader of the directory no longer saw it. A storage now remembers, for each key, the file it last wrote and has not shared since (_own_paths, replacing _own_files). restore() keeps it, and so do delete and apply_reform's wipe. Writing a deleted key again therefore reuses that file instead of making a new one each time (the review's P3: 100 delete/write cycles after cloning had left 101 files). Tests: - test_uprating_order.py: test_an_input_set_on_a_clone_is_not_replayed_by_its_source (both helpers, carry-over on/off, memory and disk; 8/8 fail on 1ee63a80). - test_on_disk_storage_clones.py: restore-then-write and delete-then-write reuse (both fail on 1ee63a80). The restore test now uses a month period, because Windows file names cannot hold the ':' of month:2012-03:3; that is what failed Windows CI on 1ee63a80. - The helper property: the clone also applies a no-op reform, and its input periods must not change. Its generated examples alone fail on 1ee63a80. Co-Authored-By: Claude Opus 5.5 --- .../data_storage/on_disk_storage.py | 26 +++++------ policyengine_core/simulations/simulation.py | 5 +++ tests/core/test_on_disk_storage_clones.py | 29 +++++++++++-- tests/core/test_uprating_order.py | 43 +++++++++++++++++++ tests/core/test_uprating_order_property.py | 15 ++++++- 5 files changed, 101 insertions(+), 17 deletions(-) diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 24501bf33..26750e389 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -32,10 +32,11 @@ def __init__( # File keys stored with ``put(..., derived=True)``; see # ``InMemoryStorage``. self._derived = set() - # Paths of the files this storage wrote and has not shared with a - # clone since: the only shared-directory files ``put`` writes over - # (see ``_path_to_write``). - self._own_files = set() + # For each key, the file this storage last wrote for it, while it has + # not shared that file since: the only files of its family ``put`` + # writes over (see ``_path_to_write``). Kept when the key is deleted, + # so that writing it again reuses the file. + self._own_paths = {} # Paths of every file this storage, the storage it was cloned from or # any other clone of either wrote or stored when cloned: one set, # shared by all of them. @@ -47,7 +48,7 @@ def __init__( def __getstate__(self) -> dict: # Whatever this state is read into reads the same files, so from now # on this storage writes over none of them, as after ``clone``. - self._own_files = set() + self._own_paths = {} return self.__dict__.copy() def __setstate__(self, state: dict) -> None: @@ -55,7 +56,7 @@ def __setstate__(self, state: dict) -> None: # count as inputs. state.setdefault("_derived", set()) # Nor did it record the files it wrote: it writes over none it stores. - state.setdefault("_own_files", set()) + state.setdefault("_own_paths", {}) if "_family_files" not in state: state["_family_files"] = set(state.get("_files", {}).values()) self.__dict__.update(state) @@ -83,7 +84,7 @@ def clone(self) -> "OnDiskStorage": # family did not write (read back by ``restore``, say). self._family_files.update(self._files.values()) clone._family_files = self._family_files - self._own_files = set() + self._own_paths = {} return clone def _decode_file(self, file: str) -> ArrayLike: @@ -138,7 +139,7 @@ def put( value = value.view(numpy.ndarray) numpy.save(path, value) self._files[filename] = path - self._own_files.add(path) + self._own_paths[filename] = path self._family_files.add(path) if derived: self._derived.add(filename) @@ -157,11 +158,11 @@ def _path_to_write(self, filename: str) -> str: ``REPLACEMENTS_DIR``. A file only storages outside this family wrote is written over, as before. """ - current = self._files.get(filename) - if current is not None and current in self._own_files: - return current + own = self._own_paths.get(filename) + if own is not None: + return own path = os.path.join(self.storage_dir, filename) + ".npy" - if path in self._own_files or path not in self._family_files: + if path not in self._family_files: return path directory = os.path.join(self.storage_dir, REPLACEMENTS_DIR) os.makedirs(directory, exist_ok=True) @@ -210,7 +211,6 @@ def restore(self) -> None: self._files = files = {} # Files read back from a directory carry no derived marks. self._derived = set() - self._own_files = set() # Restore self._files from content of storage_dir. for filename in os.listdir(self.storage_dir): if not filename.endswith(".npy"): diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 9aac75e2c..1a21bc8fc 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1622,6 +1622,11 @@ def clone( ): new_dict[key] = value new._fast_cache = {} + # Each simulation records its own inputs. A shared record let an + # input set on one replay, in the other's ``apply_reform``, whatever + # the other had calculated for that period, as an input. + if hasattr(self, "_user_input_keys"): + new._user_input_keys = set(self._user_input_keys) # Only pass ``share_arrays`` when sharing, so a population or holder # ``clone`` override with the earlier signature still deep-copies. diff --git a/tests/core/test_on_disk_storage_clones.py b/tests/core/test_on_disk_storage_clones.py index c4c425dea..f8de89a45 100644 --- a/tests/core/test_on_disk_storage_clones.py +++ b/tests/core/test_on_disk_storage_clones.py @@ -180,7 +180,7 @@ def test_storages_not_cloned_from_one_another_write_over_each_others_files( def test_restore_reads_only_the_files_named_for_their_keys(storage): storage.put(np.array([1.0]), "2012") - storage.put(np.array([2.0]), "month:2012-03:3") + storage.put(np.array([2.0]), "2012-03") clone = storage.clone() clone.put(np.array([3.0]), "2012") assert len(_npy_files(storage)) == 3 @@ -190,7 +190,7 @@ def test_restore_reads_only_the_files_named_for_their_keys(storage): assert sorted(map(str, restored.get_known_periods())) == [ "2012", - "month:2012-03:3", + "2012-03", ] np.testing.assert_array_equal(restored.get("2012"), [1]) @@ -226,7 +226,7 @@ def test_a_storage_unpickled_from_before_writes_over_none_of_its_files(storage): share them with another, so it writes a new file for each.""" storage.put(np.array([1.0]), "2012") state = pickle.loads(pickle.dumps(storage.__dict__)) - del state["_own_files"], state["_family_files"] + del state["_own_paths"], state["_family_files"] old = OnDiskStorage.__new__(OnDiskStorage) old.__setstate__(state) @@ -234,3 +234,26 @@ def test_a_storage_unpickled_from_before_writes_over_none_of_its_files(storage): np.testing.assert_array_equal(storage.get("2012"), [1]) np.testing.assert_array_equal(old.get("2012"), [2]) + + +def test_a_storage_restored_then_written_still_writes_its_own_file(storage): + """``restore`` shares nothing, so a storage that was never cloned keeps + writing over the files it wrote, which another storage reads.""" + storage.put(np.array([5.0]), "2013") + storage.restore() + storage.put(np.array([6.0]), "2013") + reader = OnDiskStorage(storage.storage_dir, preserve_storage_dir=True) + reader.restore() + np.testing.assert_array_equal(reader.get("2013"), [6]) + assert len(_npy_files(storage)) == 1 + + +def test_deleting_and_writing_a_key_again_reuses_its_file(storage): + storage.put(np.array([0.0]), "2012", derived=True) + clone = storage.clone() + for value in range(1, 21): + storage.delete("2012") + storage.put(np.array([float(value)]), "2012", derived=True) + np.testing.assert_array_equal(storage.get("2012"), [20]) + np.testing.assert_array_equal(clone.get("2012"), [0]) + assert len(_npy_files(storage)) == 2 diff --git a/tests/core/test_uprating_order.py b/tests/core/test_uprating_order.py index de660c43f..05285239e 100644 --- a/tests/core/test_uprating_order.py +++ b/tests/core/test_uprating_order.py @@ -41,6 +41,7 @@ from policyengine_core import periods from policyengine_core.country_template import entities +from policyengine_core.experimental import MemoryConfig from policyengine_core.reforms import Reform from policyengine_core.variables import Variable from tests.fixtures.uprating_order import ( @@ -389,3 +390,45 @@ def test_helper_input_keeps_an_input_already_set_in_its_year( np.testing.assert_allclose( built.calculate(variable, month), values, rtol=1e-6 ) + + +class _no_change(Reform): + def apply(self): + pass + + +@pytest.mark.filterwarnings("ignore:Memory configuration is a feature") +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +@pytest.mark.parametrize("variable", list(HELPER_INPUTS), ids=["dispatch", "divide"]) +def test_an_input_set_on_a_clone_is_not_replayed_by_its_source( + system, variable, on_disk +): + """``apply_reform`` keeps a simulation's inputs and drops what it + calculated. A clone recorded its inputs in its source's record, so once + a helper input on the clone replaced a month the source had calculated, + the source's reform kept its calculated month as an input, and uprated + from it.""" + inputs = {variable: {"2011-06": [5.0, 6.0]}} + source = simulation(system, {}) + if on_disk: + source.memory_config = MemoryConfig(max_memory_occupation=0) + holder = source.get_holder(variable) + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + source.set_input(variable, "2011-06", np.array(inputs[variable]["2011-06"])) + source.calculate(variable, "2012-12") + clone = source.clone() + yearly, monthly = HELPER_INPUTS[variable] + clone.set_input(variable, "2012", np.array(yearly)) + + source.apply_reform(_no_change) + clone.apply_reform(_no_change) + + december = periods.period("2012-12") + assert (variable, "default", december) not in source._user_input_keys + assert source.get_holder(variable).get_array(december) is None + assert_bitwise_equal( + source.calculate(variable, "2013-01"), + alone(system, inputs, variable, "2013-01"), + ) + np.testing.assert_array_equal(clone.calculate(variable, "2012-12"), monthly) diff --git a/tests/core/test_uprating_order_property.py b/tests/core/test_uprating_order_property.py index d829f6768..1b39b23c2 100644 --- a/tests/core/test_uprating_order_property.py +++ b/tests/core/test_uprating_order_property.py @@ -39,6 +39,7 @@ from policyengine_core import periods # noqa: E402 from policyengine_core.experimental import MemoryConfig # noqa: E402 +from policyengine_core.reforms import Reform # noqa: E402 from tests.fixtures.uprating_order import ( # noqa: E402 COUNT, assert_bitwise_equal, @@ -315,7 +316,8 @@ def test_uprating_depends_only_on_the_inputs(scenario, auto_carry_over): # month, from the simulation, a branch or a nested branch. Every month of the # year then holds an input. And the input changes nothing a clone of the # simulation, taken just before it, has stored: a clone keeps the branch -# name, and on disk it shares the files stored before cloning. +# name, and on disk it shares the files stored before cloning. Nor does it +# make the clone's ``apply_reform`` keep any value the clone calculated. # # The domain: the variable is stored in memory or on disk. Earlier inputs are # months, set on the simulation first; the yearly input and the requests @@ -367,6 +369,11 @@ def _helper_simulation(system, variable, earlier_inputs, branch, on_disk): return _read_from(built, branch, variable, None) +class _no_change(Reform): + def apply(self): + pass + + def _stored_months(built, variable, months): holder = built.get_holder(variable) stored = {} @@ -463,6 +470,8 @@ def test_helper_input_replaces_values_calculated_in_its_year( if with_clone: clone = built.clone() stored_by_clone = _stored_months(clone, variable, months) + clone_holder = clone.get_holder(variable) + inputs_of_clone = set(clone_holder.get_input_periods(clone.branch_name)) built.set_input(variable, year, yearly) fresh = _helper_simulation(system, variable, earlier_inputs, branch, on_disk) @@ -480,6 +489,10 @@ def test_helper_input_replaces_values_calculated_in_its_year( assert_bitwise_equal( value, stored_by_clone[month], f"{message}: clone's {month}" ) + clone.apply_reform(_no_change) + assert ( + set(clone_holder.get_input_periods(clone.branch_name)) == inputs_of_clone + ), f"{message}: the clone's inputs after a reform" holder = built.get_holder(variable) for month in months: assert not holder.is_derived(periods.period(month), built.branch_name), month From 45f2bfc355119ca79de67b2ba50f53dd9ef657c2 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 4 Oct 2026 23:18:24 -0400 Subject: [PATCH 5/5] Say in the changelog that each clone keeps its own input record The fourth commit (c5c7bf4c) gave each Simulation.clone() a copy of _user_input_keys instead of sharing the source's set, which changes what apply_reform keeps as an input after an input on the other simulation. The fragment only described the on-disk file change. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-disk-storage-clone-writes.fixed.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changelog.d/fix-disk-storage-clone-writes.fixed.md b/changelog.d/fix-disk-storage-clone-writes.fixed.md index 44d2e3eb4..1977f85df 100644 --- a/changelog.d/fix-disk-storage-clone-writes.fixed.md +++ b/changelog.d/fix-disk-storage-clone-writes.fixed.md @@ -1 +1 @@ -Simulations and storages cloned from one another no longer write over each other's values in on-disk storage: a value stored for a key whose file a clone also reads goes to a new file of its own. +Simulations and storages cloned from one another no longer write over each other's values in on-disk storage: a value stored for a key whose file a clone also reads goes to a new file of its own. Each cloned simulation also keeps its own record of user inputs, so `apply_reform` on one no longer keeps, as an input, a value it calculated for a period the other was given as input.