From 65b1656f93f89f2c3992f491b01bc2e1b63b6943 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 08:06:04 -0400 Subject: [PATCH 1/4] Cache ADD and DIVIDE results only where a plain read returns them calculate_add and calculate_divide stored their result at the requested period, where every later plain calculate of that variable and period found it. For a STOCK variable that is not its plain value (last month of the year, or the year's value for a month), so a monthly STOCK read 120 instead of 10 after an ADD and a yearly one 1 instead of 12 after a DIVIDE. The same happened over several periods of a variable's own unit (a plain read raises), for day variables, for integer or boolean FLOW results the cache truncates, and over an input stored at that period. Results are now cached only for the case _calculate itself routes to these options (a FLOW variable over a period of another unit), when the result has the variable's dtype and nothing is stored there yet. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-stock-option-caches.fixed.md | 1 + policyengine_core/simulations/simulation.py | 46 +++++- tests/core/test_option_result_cache.py | 148 ++++++++++++++++++ .../core/test_option_result_cache_property.py | 80 ++++++++++ tests/fixtures/option_caches.py | 98 ++++++++++++ 5 files changed, 369 insertions(+), 4 deletions(-) create mode 100644 changelog.d/fix-stock-option-caches.fixed.md create mode 100644 tests/core/test_option_result_cache.py create mode 100644 tests/core/test_option_result_cache_property.py create mode 100644 tests/fixtures/option_caches.py diff --git a/changelog.d/fix-stock-option-caches.fixed.md b/changelog.d/fix-stock-option-caches.fixed.md new file mode 100644 index 000000000..6de3bef3f --- /dev/null +++ b/changelog.d/fix-stock-option-caches.fixed.md @@ -0,0 +1 @@ +Stop `calculate_add` and `calculate_divide` from caching results a plain `calculate` of the same period would not return (STOCK variables, several periods of a variable's own unit, day variables, values whose type the cache would truncate) or over an input stored there, so a later read no longer depends on whether the option ran first. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 2652d9723..874da9856 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1048,8 +1048,7 @@ def calculate_add( self.calculate(variable_name, sub_period) for sub_period in period.get_subperiods(variable.definition_period) ) - holder = self.get_holder(variable.name) - holder.put_in_cache(result, period, self.branch_name) + self._cache_option_result(variable, period, result) return result def calculate_divide( @@ -1081,8 +1080,7 @@ def calculate_divide( if period.unit == periods.MONTH: computation_period = period.this_year result = self.calculate(variable_name, period=computation_period) / 12.0 - holder = self.get_holder(variable.name) - holder.put_in_cache(result, period, self.branch_name) + self._cache_option_result(variable, period, result) return result elif period.unit == periods.YEAR: return self.calculate(variable_name, period) @@ -1093,6 +1091,46 @@ def calculate_divide( ) ) + def _cache_option_result( + self, variable: Variable, period: Period, result: ArrayLike + ) -> None: + """Cache an ADD or DIVIDE result at ``period`` if a plain read would return it. + + A value cached at ``period`` is what every later ``calculate`` of the + variable at ``period`` returns. ``_calculate`` computes a FLOW + variable over a period of another unit with these same options (a + monthly variable over a year with ``calculate_add``, a yearly one over + a month with ``calculate_divide``), so their result is the plain value + there and is cached. Anywhere else it is not: + + - A STOCK variable's plain value over a year is its last month's, and + over a month the year's, not the sum or the twelfth. + - Over several periods of the variable's own unit, a plain read + raises instead. + - A day variable's plain read over a month or a year does not sum. + - Over a single period of its own unit, the sum is the value + ``calculate`` has already stored. + + Caching there would make a later plain read depend on whether the + option ran first. Nor does an option result replace a value that a + plain read already finds at ``period``, such as an input there, or + get cached when storing it would change it: the twelfth of an integer + or a count of true months is stored as the variable's own type, so a + later read would return the truncated value where the first returned + the exact one. + """ + if variable.quantity_type == QuantityType.STOCK: + return + routed = (variable.definition_period == MONTH and period.unit == YEAR) or ( + variable.definition_period == YEAR and period.unit == MONTH + ) + if not routed or np.asarray(result).dtype != variable.dtype: + return + holder = self.get_holder(variable.name) + if holder.get_array(period, self.branch_name) is not None: + return + holder.put_in_cache(result, period, self.branch_name) + def calculate_output(self, variable_name: str, period: Period = None) -> ArrayLike: """ Calculate the value of a variable using the ``calculate_output`` attribute of the variable. diff --git a/tests/core/test_option_result_cache.py b/tests/core/test_option_result_cache.py new file mode 100644 index 000000000..c13859452 --- /dev/null +++ b/tests/core/test_option_result_cache.py @@ -0,0 +1,148 @@ +"""ADD and DIVIDE results are cached only where a plain read returns them. + +``calculate_add`` and ``calculate_divide`` used to store their result at the +requested period, where every later plain ``calculate`` of the variable at +that period found it. That is right for a FLOW variable, whose plain value +over a period of another unit is exactly that sum or twelfth, and wrong +elsewhere: a STOCK variable's plain value over a year is its last month's, +and over a month the year's. A plain read then returned 120 or 1 instead of +10 or 12, depending only on whether the option had run first. +""" + +import pytest + +from tests.fixtures.option_caches import ( + DAY, + FLOW, + MONTH, + PROBE, + STOCK, + YEAR, + build, + make_probe, + run, +) + + +def fresh_and_after(probe, inputs, target, prior, carry_over=False): + """``target`` in a new simulation, and after running ``prior`` first.""" + fresh = build(probe, inputs, carry_over) + used = build(probe, inputs, carry_over) + prior_result = run(used, *prior) + return run(fresh, *target), run(used, *target), prior_result + + +def test_add_over_a_monthly_stock_leaves_its_annual_value_alone(): + probe = make_probe(MONTH, STOCK) + fresh, after, prior = fresh_and_after( + probe, + {"2012-01": 10}, + ("calculate", "2012"), + ("add", "2012"), + carry_over=True, + ) + assert prior == [120.0] + assert fresh == after == [10.0] + + +def test_divide_over_a_yearly_stock_leaves_its_monthly_value_alone(): + probe = make_probe(YEAR, STOCK, value_type=int) + fresh, after, prior = fresh_and_after( + probe, + {"2012": 12}, + ("calculate", "2012-06"), + ("divide", "2012-06"), + ) + assert prior == [1.0] + assert fresh == after == [12.0] + + +def test_stock_with_a_formula_keeps_its_last_month_value(): + probe = make_probe(MONTH, STOCK, with_formula=True) + fresh, after, prior = fresh_and_after( + probe, {}, ("calculate", "2012"), ("add", "2012") + ) + assert prior == [sum(1200.0 + month for month in range(1, 13))] + assert fresh == after == [1212.0] + + +def test_flow_add_is_still_cached_as_the_plain_annual_value(): + probe = make_probe(MONTH, FLOW, with_formula=True) + simulation = build(probe) + total = simulation.calculate_add(PROBE, "2012") + holder = simulation.get_holder(PROBE) + assert holder.get_array("2012").tolist() == total.tolist() + assert simulation.calculate(PROBE, "2012").tolist() == total.tolist() + + +def test_flow_divide_is_still_cached_as_the_plain_monthly_value(): + probe = make_probe(YEAR, FLOW, with_formula=True) + simulation = build(probe) + twelfth = simulation.calculate_divide(PROBE, "2012-03") + holder = simulation.get_holder(PROBE) + assert holder.get_array("2012-03").tolist() == twelfth.tolist() + assert twelfth.tolist() == [pytest.approx(1012.0 / 12)] + + +def test_add_does_not_replace_an_input_stored_at_the_year(): + # A monthly FLOW variable whose input is stored at the year itself. + probe = make_probe(MONTH, FLOW, with_formula=True, set_input=None) + fresh, after, prior = fresh_and_after( + probe, {"2012": 7}, ("calculate", "2012"), ("add", "2012") + ) + assert prior == [sum(1200.0 + month for month in range(1, 13))] + assert fresh == after == [7.0] + + +def test_divide_does_not_replace_an_input_stored_at_the_month(): + probe = make_probe(YEAR, FLOW, with_formula=True, set_input=None) + fresh, after, prior = fresh_and_after( + probe, {"2012-02": 24}, ("calculate", "2012-02"), ("divide", "2012-02") + ) + assert prior == [pytest.approx(1012.0 / 12)] + assert fresh == after == [24.0] + + +@pytest.mark.parametrize( + "definition_period, period", + [(YEAR, "year:2012:2"), (MONTH, "month:2012-01:3")], +) +def test_add_over_several_own_periods_does_not_make_a_plain_read_succeed( + definition_period, period +): + probe = make_probe(definition_period, FLOW, with_formula=True) + fresh, after, prior = fresh_and_after( + probe, {}, ("calculate", period), ("add", period) + ) + assert isinstance(prior, list) + assert fresh == after == "ValueError" + + +def test_add_over_a_day_variable_leaves_its_monthly_value_alone(): + # A plain read of a day variable over a month runs its formula for the + # month; it does not sum the days. + probe = make_probe(DAY, FLOW, with_formula=True) + fresh, after, prior = fresh_and_after( + probe, {}, ("calculate", "2012-01"), ("add", "2012-01") + ) + assert prior == [31 * 1201.0] + assert fresh == after == [1201.0] + + +def test_stock_add_in_a_branch_leaves_the_branch_annual_value_alone(): + probe = make_probe(MONTH, STOCK, with_formula=True) + fresh = build(probe).get_branch("other") + used = build(probe).get_branch("other") + used.calculate_add(PROBE, "2012") + assert run(used, "calculate", "2012") == run(fresh, "calculate", "2012") + assert run(used, "calculate", "2012") == [1212.0] + + +def test_a_twelfth_of_an_integer_flow_is_not_cached_truncated(): + # Storing 1012 / 12 in an integer variable would truncate it to 84, so a + # second monthly read would differ from the first. + probe = make_probe(YEAR, FLOW, value_type=int, with_formula=True) + fresh, after, prior = fresh_and_after( + probe, {}, ("calculate", "2012-01"), ("calculate", "2012-01") + ) + assert fresh == after == prior == [pytest.approx(1012 / 12)] diff --git a/tests/core/test_option_result_cache_property.py b/tests/core/test_option_result_cache_property.py new file mode 100644 index 000000000..11f84ba1a --- /dev/null +++ b/tests/core/test_option_result_cache_property.py @@ -0,0 +1,80 @@ +"""Property: running ADD or DIVIDE first never changes a later result. + +For a monthly or yearly variable of any quantity type and value type, with +or without a formula, and with inputs stored at any periods, every request +(plain ``calculate``, ``calculate_add`` or ``calculate_divide`` at a single +or multi-period period) returns the same values, or raises the same kind of +error, whatever requests ran before it. The reference is the same request in +a new simulation with the same inputs. + +Auto-carry-over is off here: on its own, carrying over calculated values is +a separate order dependence (policyengine-core#562). +""" + +import pytest + +from tests.fixtures.option_caches import ( + FLOW, + MONTH, + STOCK, + YEAR, + build, + make_probe, + run, +) + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +OPERATIONS = ["calculate", "add", "divide"] +PERIODS = [ + "2012", + "2013", + "2012-01", + "2012-06", + "2012-12", + "2013-01", + "year:2012:2", + "month:2012-01:3", +] +INPUT_PERIODS = ["2012", "2013", "2012-01", "2012-06", "2012-12"] +REQUESTS = st.tuples(st.sampled_from(OPERATIONS), st.sampled_from(PERIODS)) + + +@st.composite +def scenarios(draw): + probe = dict( + definition_period=draw(st.sampled_from([MONTH, YEAR])), + quantity_type=draw(st.sampled_from([STOCK, FLOW])), + value_type=draw(st.sampled_from([float, int, bool])), + with_formula=draw(st.booleans()), + set_input=draw(st.sampled_from(["default", None])), + ) + inputs = draw( + st.dictionaries(st.sampled_from(INPUT_PERIODS), st.integers(0, 40), max_size=3) + ) + prior = draw(st.lists(REQUESTS, max_size=5)) + target = draw(REQUESTS) + return probe, inputs, prior, target + + +def build_or_none(probe, inputs): + try: + return build(probe, inputs) + except (TypeError, ValueError): + # Inputs this variable cannot take at those periods (spreading a + # boolean over months, for instance, subtracts booleans). + return None + + +@hypothesis.settings(max_examples=400, deadline=None) +@hypothesis.given(scenarios()) +def test_earlier_requests_do_not_change_a_result(scenario): + attributes, inputs, prior, target = scenario + probe = make_probe(**attributes) + fresh = build_or_none(probe, inputs) + hypothesis.assume(fresh is not None) + used = build(probe, inputs) + for request in prior: + run(used, *request) + assert run(used, *target) == run(fresh, *target) diff --git a/tests/fixtures/option_caches.py b/tests/fixtures/option_caches.py new file mode 100644 index 000000000..b28de4004 --- /dev/null +++ b/tests/fixtures/option_caches.py @@ -0,0 +1,98 @@ +"""One-variable simulations for testing what ADD and DIVIDE results cache. + +Shared by ``tests/core/test_option_result_cache.py`` and its Hypothesis +property module, which pytest must be able to skip on its own when +Hypothesis is not installed. +""" + +from typing import Any, Dict, Optional + +import numpy as np + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.entities import Person +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables import QuantityType, Variable + +PROBE = "probe" + + +def period_code(period) -> float: + """A value that differs for every month and year, for probe formulas.""" + start = period.start + if period.unit == periods.YEAR: + return 1000.0 + start.year % 100 + return float(100 * (start.year % 100) + start.month) + + +def make_probe( + definition_period: str, + quantity_type: str, + value_type: type = float, + with_formula: bool = False, + set_input: Optional[Any] = "default", +) -> type: + """A Person variable named ``probe`` with the given time behaviour. + + With ``with_formula``, its formula returns ``period_code(period)``, so + sums over months and twelfths of years are easy to tell apart from + single-period values. ``set_input=None`` stores inputs at the period + given instead of spreading them over the variable's own periods. + """ + attributes = dict( + value_type=value_type, + entity=Person, + definition_period=definition_period, + quantity_type=quantity_type, + label="Option cache probe", + ) + if set_input is None: + attributes["set_input"] = None + if with_formula: + + def formula(person, period): + return person.filled_array(period_code(period)) + + attributes["formula"] = formula + return type(PROBE, (Variable,), attributes) + + +def build( + probe: type, + inputs: Optional[Dict[str, Any]] = None, + carry_over: bool = False, +): + """One-person simulation of ``probe`` with ``inputs`` ({period: value}).""" + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = carry_over + system.add_variable(probe) + simulation = SimulationBuilder().build_default_simulation(system, count=1) + for period, value in (inputs or {}).items(): + simulation.set_input(PROBE, period, np.array([value])) + return simulation + + +def run(simulation, operation: str, period: str): + """Run ``calculate``, ``add`` or ``divide`` on ``probe``. + + Returns the value as a list, or the error's type name, so that results of + fresh and reused simulations compare equal exactly when both return the + same values or both raise the same kind of error. + """ + method = { + "calculate": simulation.calculate, + "add": simulation.calculate_add, + "divide": simulation.calculate_divide, + }[operation] + try: + return np.asarray(method(PROBE, period)).astype(float).tolist() + except ValueError as error: + return type(error).__name__ + + +STOCK = QuantityType.STOCK +FLOW = QuantityType.FLOW +MONTH = periods.MONTH +YEAR = periods.YEAR +DAY = periods.DAY From 268bee26a50a96823f8ec1cc788f6d2d032bc324 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 07:27:04 -0400 Subject: [PATCH 2/4] Let ADD and DIVIDE replace an aggregate they calculated before Review of 65b1656f: refusing to replace any stored value also stopped an explicit ADD or DIVIDE from refreshing an aggregate cached before an input changed (set January to 12, ADD 2012, set January to 24, ADD 2012: a plain read of 2012 gave 12, where master gives 24). Only an input that a plain read finds there, on this branch, an ancestor or default, is now protected; a calculated value is replaced. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/simulations/simulation.py | 31 +++++++++++++++------ tests/core/test_option_result_cache.py | 28 +++++++++++++++++++ 2 files changed, 50 insertions(+), 9 deletions(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 874da9856..4cb1aff84 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1112,12 +1112,12 @@ def _cache_option_result( ``calculate`` has already stored. Caching there would make a later plain read depend on whether the - option ran first. Nor does an option result replace a value that a - plain read already finds at ``period``, such as an input there, or - get cached when storing it would change it: the twelfth of an integer - or a count of true months is stored as the variable's own type, so a - later read would return the truncated value where the first returned - the exact one. + option ran first. Nor is an option result cached when storing it would + change it (the twelfth of an integer or a count of true months is + stored as the variable's own type, so a later read would return the + truncated value where the first returned the exact one), or over an + input a plain read finds at ``period``. It does replace a value + calculated there before, which may predate a change to the inputs. """ if variable.quantity_type == QuantityType.STOCK: return @@ -1126,10 +1126,23 @@ def _cache_option_result( ) if not routed or np.asarray(result).dtype != variable.dtype: return - holder = self.get_holder(variable.name) - if holder.get_array(period, self.branch_name) is not None: + if self._reads_input_at(variable.name, period): return - holder.put_in_cache(result, period, self.branch_name) + self.get_holder(variable.name).put_in_cache(result, period, self.branch_name) + + def _reads_input_at(self, variable_name: str, period: Period) -> bool: + """Whether a plain read of the variable at ``period`` finds an input. + + The read takes the value stored under the first of this branch, its + ancestors and ``default`` that has one; it is an input if that branch + stored it with ``set_input`` (or loaded it from the dataset). + """ + holder = self.get_holder(variable_name) + input_keys = getattr(self, "_user_input_keys", None) or () + for branch_name in self._get_visible_branch_names(): + if holder._get_array_from_storage(period, branch_name) is not None: + return (variable_name, branch_name, period) in input_keys + return False def calculate_output(self, variable_name: str, period: Period = None) -> ArrayLike: """ diff --git a/tests/core/test_option_result_cache.py b/tests/core/test_option_result_cache.py index c13859452..212e6f10d 100644 --- a/tests/core/test_option_result_cache.py +++ b/tests/core/test_option_result_cache.py @@ -146,3 +146,31 @@ def test_a_twelfth_of_an_integer_flow_is_not_cached_truncated(): probe, {}, ("calculate", "2012-01"), ("calculate", "2012-01") ) assert fresh == after == prior == [pytest.approx(1012 / 12)] + + +@pytest.mark.parametrize( + "definition_period, native, target, option, first, second", + [ + (MONTH, "2012-01", "2012", "add", [12.0], [24.0]), + (YEAR, "2012", "2012-01", "divide", [1.0], [2.0]), + ], +) +def test_an_option_refreshes_the_aggregate_it_cached_before( + definition_period, native, target, option, first, second +): + # The cached aggregate predates the new input; running the option again + # replaces it, as on master, so the plain read follows the input. + simulation = build(make_probe(definition_period, FLOW)) + simulation.set_input(PROBE, native, [12]) + assert run(simulation, option, target) == first + simulation.set_input(PROBE, native, [24]) + assert run(simulation, option, target) == second + assert run(simulation, "calculate", target) == second + + +def test_an_option_in_a_branch_does_not_replace_a_parent_input(): + probe = make_probe(MONTH, FLOW, with_formula=True, set_input=None) + simulation = build(probe, {"2012": 7}) + branch = simulation.get_branch("other") + assert run(branch, "add", "2012") == [sum(1200.0 + month for month in range(1, 13))] + assert run(branch, "calculate", "2012") == [7.0] From b46440d86f5419508f163f14e3e14dd9ad31c1ef Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:03:18 -0400 Subject: [PATCH 3/4] Use storage provenance for ADD and DIVIDE input protection Remove the append-only input-record guard so explicit FLOW options can refresh derived aggregates after a target input is deleted. Keep visible inputs via Holder.put_in_cache(derived=True), including restored inputs and nearest-ancestor inputs. Preserve existing ordering properties and add focused refresh differential and provenance regressions. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-stock-option-caches.fixed.md | 2 +- policyengine_core/simulations/simulation.py | 28 ++--- tests/core/test_option_result_cache.py | 103 +++++++++++++++++- .../core/test_option_result_cache_property.py | 56 ++++++++++ 4 files changed, 165 insertions(+), 24 deletions(-) diff --git a/changelog.d/fix-stock-option-caches.fixed.md b/changelog.d/fix-stock-option-caches.fixed.md index 6de3bef3f..3c6d80e77 100644 --- a/changelog.d/fix-stock-option-caches.fixed.md +++ b/changelog.d/fix-stock-option-caches.fixed.md @@ -1 +1 @@ -Stop `calculate_add` and `calculate_divide` from caching results a plain `calculate` of the same period would not return (STOCK variables, several periods of a variable's own unit, day variables, values whose type the cache would truncate) or over an input stored there, so a later read no longer depends on whether the option ran first. +Cache `calculate_add` and `calculate_divide` results as derived only where a plain read returns them without changing dtype, preserving visible inputs and allowing explicit aggregate refresh after an input changes or is deleted. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 60ffb8325..8ed2fe858 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1364,9 +1364,17 @@ def _cache_option_result( option ran first. Nor is an option result cached when storing it would change it (the twelfth of an integer or a count of true months is stored as the variable's own type, so a later read would return the - truncated value where the first returned the exact one), or over an - input a plain read finds at ``period``. It does replace a value + truncated value where the first returned the exact one). + + The result is cached as derived, so auto-carry-over and uprating + never take it for an input, and ``put_in_cache`` keeps an input a + plain read finds at ``period`` (in this branch, an ancestor or + ``default``) instead of storing it. It does replace a value calculated there before, which may predate a change to the inputs. + Whether the value read is an input comes from the storage's mark for + that key, which every write sets, not from the simulation's record of + ``set_input`` calls: that record keeps a key after its value is + deleted, and has no key for an input restored from a dump. """ if variable.quantity_type == QuantityType.STOCK: return @@ -1375,26 +1383,10 @@ def _cache_option_result( ) if not routed or np.asarray(result).dtype != variable.dtype: return - if self._reads_input_at(variable.name, period): - return self.get_holder(variable.name).put_in_cache( result, period, self.branch_name, derived=True ) - def _reads_input_at(self, variable_name: str, period: Period) -> bool: - """Whether a plain read of the variable at ``period`` finds an input. - - The read takes the value stored under the first of this branch, its - ancestors and ``default`` that has one; it is an input if that branch - stored it with ``set_input`` (or loaded it from the dataset). - """ - holder = self.get_holder(variable_name) - input_keys = getattr(self, "_user_input_keys", None) or () - for branch_name in self._get_visible_branch_names(): - if holder._get_array_from_storage(period, branch_name) is not None: - return (variable_name, branch_name, period) in input_keys - return False - def calculate_output(self, variable_name: str, period: Period = None) -> ArrayLike: """ Calculate the value of a variable using the ``calculate_output`` attribute of the variable. diff --git a/tests/core/test_option_result_cache.py b/tests/core/test_option_result_cache.py index 212e6f10d..5bed17d93 100644 --- a/tests/core/test_option_result_cache.py +++ b/tests/core/test_option_result_cache.py @@ -11,6 +11,11 @@ import pytest +from policyengine_core import periods +from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, +) from tests.fixtures.option_caches import ( DAY, FLOW, @@ -72,6 +77,8 @@ def test_flow_add_is_still_cached_as_the_plain_annual_value(): total = simulation.calculate_add(PROBE, "2012") holder = simulation.get_holder(PROBE) assert holder.get_array("2012").tolist() == total.tolist() + assert holder.is_derived(periods.period("2012")) + assert periods.period("2012") not in holder.get_input_periods() assert simulation.calculate(PROBE, "2012").tolist() == total.tolist() @@ -81,6 +88,8 @@ def test_flow_divide_is_still_cached_as_the_plain_monthly_value(): twelfth = simulation.calculate_divide(PROBE, "2012-03") holder = simulation.get_holder(PROBE) assert holder.get_array("2012-03").tolist() == twelfth.tolist() + assert holder.is_derived(periods.period("2012-03")) + assert periods.period("2012-03") not in holder.get_input_periods() assert twelfth.tolist() == [pytest.approx(1012.0 / 12)] @@ -168,9 +177,93 @@ def test_an_option_refreshes_the_aggregate_it_cached_before( assert run(simulation, "calculate", target) == second -def test_an_option_in_a_branch_does_not_replace_a_parent_input(): - probe = make_probe(MONTH, FLOW, with_formula=True, set_input=None) - simulation = build(probe, {"2012": 7}) +@pytest.mark.parametrize( + "definition_period, target, option, result", + [ + (MONTH, "2012", "add", [sum(1200.0 + month for month in range(1, 13))]), + (YEAR, "2012-02", "divide", [pytest.approx(1012.0 / 12)]), + ], +) +@pytest.mark.parametrize("nested", [False, True]) +def test_an_option_in_a_branch_does_not_replace_a_parent_input( + definition_period, target, option, result, nested +): + probe = make_probe(definition_period, FLOW, with_formula=True, set_input=None) + simulation = build(probe, {target: 7}) + if nested: + simulation = simulation.get_branch("parent") + # A raw holder write has an input mark, without a set_input record. + # The nearest ancestor must take precedence over default's input. + simulation.get_holder(PROBE).put_in_cache( + [9], periods.period(target), simulation.branch_name + ) branch = simulation.get_branch("other") - assert run(branch, "add", "2012") == [sum(1200.0 + month for month in range(1, 13))] - assert run(branch, "calculate", "2012") == [7.0] + assert run(branch, option, target) == result + assert run(branch, "calculate", target) == [9.0 if nested else 7.0] + assert not branch.get_holder(PROBE).is_derived( + periods.period(target), branch.branch_name + ) + + +@pytest.mark.parametrize( + "definition_period, deleted, native, target, option, first, second", + [ + (MONTH, "2012", "2012-01", "2012", "add", [12.0], [24.0]), + (YEAR, "2012-01", "2012", "2012-01", "divide", [1.0], [2.0]), + ], +) +def test_an_option_refreshes_its_aggregate_where_a_deleted_input_was( + definition_period, deleted, native, target, option, first, second +): + # An input stored at the target itself, then deleted. The simulation's + # record of what ``set_input`` stored keeps it, but the value a plain + # read finds at the target afterwards is the aggregate the option + # cached, which is calculated: running the option again replaces it. + probe = make_probe(definition_period, FLOW, set_input=None) + simulation = build(probe, {deleted: 7}) + simulation.delete_arrays(PROBE, deleted) + simulation.set_input(PROBE, native, [12]) + assert run(simulation, option, target) == first + simulation.set_input(PROBE, native, [24]) + assert run(simulation, option, target) == second + assert run(simulation, "calculate", target) == second + + +@pytest.mark.parametrize( + "definition_period, stored, option, result", + [ + (MONTH, "2012", "add", [sum(1200.0 + month for month in range(1, 13))]), + (YEAR, "2012-02", "divide", [pytest.approx(1012.0 / 12)]), + ], +) +def test_an_option_does_not_replace_an_input_restored_from_a_dump( + definition_period, stored, option, result, tmp_path +): + # ``restore_simulation`` stores the dumped input with ``put_in_cache``, + # not ``set_input``: only the storage's input mark says it is one. + probe = make_probe(definition_period, FLOW, with_formula=True, set_input=None) + built = build(probe, {stored: 7}) + dump_simulation(built, str(tmp_path / "dump")) + restored = restore_simulation(str(tmp_path / "dump"), built.tax_benefit_system) + assert run(restored, option, stored) == result + assert run(restored, "calculate", stored) == [7.0] + assert not restored.get_holder(PROBE).is_derived(periods.period(stored)) + + +@pytest.mark.parametrize( + "definition_period, option, target, later", + [ + (MONTH, "add", "2012", "2013-01"), + (YEAR, "divide", "2012-06", "2013"), + ], +) +def test_an_option_result_does_not_carry_past_a_formula_end( + definition_period, option, target, later +): + probe = make_probe(definition_period, FLOW, with_formula=True) + probe.end = "2012-12-31" + fresh = build(probe, carry_over=True) + used = build(probe, carry_over=True) + run(used, option, target) + assert used.get_holder(PROBE).is_derived(periods.period(target)) + assert run(used, "calculate", later) == run(fresh, "calculate", later) == [0.0] diff --git a/tests/core/test_option_result_cache_property.py b/tests/core/test_option_result_cache_property.py index 11f84ba1a..d8a71fda8 100644 --- a/tests/core/test_option_result_cache_property.py +++ b/tests/core/test_option_result_cache_property.py @@ -13,9 +13,11 @@ import pytest +from policyengine_core import periods from tests.fixtures.option_caches import ( FLOW, MONTH, + PROBE, STOCK, YEAR, build, @@ -78,3 +80,57 @@ def test_earlier_requests_do_not_change_a_result(scenario): for request in prior: run(used, *request) assert run(used, *target) == run(fresh, *target) + + +@hypothesis.settings(max_examples=100, deadline=None) +@hypothesis.example( + definition_period=MONTH, + values=[12, 24], + deleted=True, + depth=2, + trace=False, +) +@hypothesis.example( + definition_period=YEAR, + values=[12, 24], + deleted=True, + depth=2, + trace=True, +) +@hypothesis.given( + definition_period=st.sampled_from([MONTH, YEAR]), + values=st.lists( + st.integers(-100, 100).map(lambda value: value * 12), min_size=2, max_size=5 + ), + deleted=st.booleans(), + depth=st.integers(0, 2), + trace=st.booleans(), +) +def test_explicit_refresh_matches_a_fresh_plain_read( + definition_period, values, deleted, depth, trace +): + """Refreshing a derived aggregate follows current inputs, even after deletion.""" + native, target, option = ( + ("2012-01", "2012", "add") + if definition_period == MONTH + else ("2012", "2012-01", "divide") + ) + probe = make_probe(definition_period, FLOW, set_input=None) + used = build(probe, {target: 7} if deleted else {}) + if deleted: + used.delete_arrays(PROBE, target) + used.trace = trace + for index in range(depth): + used = used.get_branch(f"branch_{index}") + for value in values: + used.set_input(PROBE, native, [value]) + fresh = build(probe, {native: value}) + expected = run(fresh, "calculate", target) + assert run(used, option, target) == expected + assert run(used, "calculate", target) == expected + assert used.get_holder(PROBE).is_derived( + periods.period(target), used.branch_name + ) + assert periods.period(target) not in used.get_holder(PROBE).get_input_periods( + used.branch_name + ) From 77df88302596d0552a9b065a1ab112701b0fa1a8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 01:54:30 -0400 Subject: [PATCH 4/4] Test option cache provenance after storage tier merge --- policyengine_core/simulations/simulation.py | 8 +- tests/core/test_option_result_cache.py | 11 +- .../test_option_result_input_provenance.py | 120 ++++++++++++++++++ 3 files changed, 130 insertions(+), 9 deletions(-) create mode 100644 tests/core/test_option_result_input_provenance.py diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 02df5701d..96d05624e 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1417,9 +1417,11 @@ def _cache_option_result( ``default``) instead of storing it. It does replace a value calculated there before, which may predate a change to the inputs. Whether the value read is an input comes from the storage's mark for - that key, which every write sets, not from the simulation's record of - ``set_input`` calls: that record keeps a key after its value is - deleted, and has no key for an input restored from a dump. + the first visible value, which every write sets. The simulation's + supplied-input record follows inputs across storage tiers for replay + and export; it can name a disk input hidden by a derived memory value, + and raw holder inputs need not be registered there. Accepted cache + writes never register the aggregate as a supplied input. """ if variable.quantity_type == QuantityType.STOCK: return diff --git a/tests/core/test_option_result_cache.py b/tests/core/test_option_result_cache.py index 5bed17d93..921e7ba44 100644 --- a/tests/core/test_option_result_cache.py +++ b/tests/core/test_option_result_cache.py @@ -215,10 +215,9 @@ def test_an_option_in_a_branch_does_not_replace_a_parent_input( def test_an_option_refreshes_its_aggregate_where_a_deleted_input_was( definition_period, deleted, native, target, option, first, second ): - # An input stored at the target itself, then deleted. The simulation's - # record of what ``set_input`` stored keeps it, but the value a plain - # read finds at the target afterwards is the aggregate the option - # cached, which is calculated: running the option again replaces it. + # Deleting the target input also removes its supplied-input record. The + # aggregate subsequently cached there is derived, so running the option + # again replaces it with the result from the current native inputs. probe = make_probe(definition_period, FLOW, set_input=None) simulation = build(probe, {deleted: 7}) simulation.delete_arrays(PROBE, deleted) @@ -239,8 +238,8 @@ def test_an_option_refreshes_its_aggregate_where_a_deleted_input_was( def test_an_option_does_not_replace_an_input_restored_from_a_dump( definition_period, stored, option, result, tmp_path ): - # ``restore_simulation`` stores the dumped input with ``put_in_cache``, - # not ``set_input``: only the storage's input mark says it is one. + # ``restore_simulation`` restores storage's input mark and registers the + # supplied input for replay. The option must preserve the visible input. probe = make_probe(definition_period, FLOW, with_formula=True, set_input=None) built = build(probe, {stored: 7}) dump_simulation(built, str(tmp_path / "dump")) diff --git a/tests/core/test_option_result_input_provenance.py b/tests/core/test_option_result_input_provenance.py new file mode 100644 index 000000000..7464c6c8a --- /dev/null +++ b/tests/core/test_option_result_input_provenance.py @@ -0,0 +1,120 @@ +"""Option caches remain derived alongside supplied-input tier records.""" + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.experimental import MemoryConfig +from tests.fixtures.option_caches import ( + FLOW, + MONTH, + PROBE, + YEAR, + build, + make_probe, + run, +) + + +def _store_on_disk(simulation): + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + holder = simulation.get_holder(PROBE) + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + + +@pytest.mark.parametrize( + "definition_period, native, target, option, expected", + [ + (MONTH, "2012-01", "2012", "add", [12.0]), + (YEAR, "2012", "2012-01", "divide", [1.0]), + ], +) +@pytest.mark.parametrize("on_disk", [False, True]) +@pytest.mark.parametrize("deleted_target", [False, True]) +def test_option_cache_is_never_registered_exported_or_replayed_as_input( + definition_period, native, target, option, expected, on_disk, deleted_target +): + simulation = build(make_probe(definition_period, FLOW, set_input=None)) + if on_disk: + _store_on_disk(simulation) + holder = simulation.get_holder(PROBE) + target_period = periods.period(target) + native_period = periods.period(native) + target_key = (PROBE, "default", target_period) + if deleted_target: + simulation.set_input(PROBE, target, [7]) + simulation.delete_arrays(PROBE, target) + assert target_key not in simulation._user_input_keys + simulation.set_input(PROBE, native, [12]) + + assert run(simulation, option, target) == expected + + storage = holder._disk_storage if on_disk else holder._memory_storage + assert storage.get(target_period).tolist() == expected + assert storage.is_derived(target_period) + assert not holder._stores_user_input(target_period, "default") + assert target_key not in simulation._user_input_keys + assert (PROBE, "default", native_period) in simulation._user_input_keys + frame = simulation.to_input_dataframe() + assert { + column: frame[column].tolist() + for column in frame.columns + if column.startswith(f"{PROBE}__") + } == {f"{PROBE}__{native}": [12.0]} + assert simulation.to_input_dict()[PROBE] == {native: [12.0]} + + simulation._invalidate_all_caches() + + assert storage.get(target_period) is None + assert storage.get(native_period).tolist() == [12.0] + assert target_key not in simulation._user_input_keys + assert holder._stores_user_input(native_period, "default") + assert simulation.to_input_dataframe()[f"{PROBE}__{native}"].tolist() == [12.0] + assert run(simulation, "calculate", target) == expected + assert storage.is_derived(target_period) + + +@pytest.mark.parametrize( + "definition_period, native, target, option, expected", + [ + (MONTH, "2012-01", "2012", "add", [12.0]), + (YEAR, "2012", "2012-01", "divide", [1.0]), + ], +) +def test_option_refresh_preserves_a_supplied_disk_input_hidden_by_memory_cache( + definition_period, native, target, option, expected +): + simulation = build(make_probe(definition_period, FLOW, set_input=None)) + _store_on_disk(simulation) + simulation.set_input(PROBE, native, [12]) + simulation.set_input(PROBE, target, [7]) + holder = simulation.get_holder(PROBE) + target_period = periods.period(target) + input_path = holder._disk_storage._files[f"default_{target}"] + # Model an existing derived value in the higher-priority memory tier. + # The supplied disk value remains independently recorded for replay. + holder._memory_storage.put( + np.array([0.0], dtype=holder.variable.dtype), target_period, derived=True + ) + + assert run(simulation, option, target) == expected + assert run(simulation, "calculate", target) == expected + assert holder._memory_storage.get(target_period).tolist() == expected + assert holder._memory_storage.is_derived(target_period) + assert holder._disk_storage.get(target_period).tolist() == [7.0] + assert holder._user_input_storage[("default", target)] == ( + target_period, + frozenset({"disk"}), + ) + assert holder._stores_user_input(target_period, "default") + assert (PROBE, "default", target_period) in simulation._user_input_keys + + simulation._invalidate_all_caches() + + assert holder._memory_storage.get(target_period) is None + assert holder._disk_storage._files[f"default_{target}"] == input_path + assert not holder._disk_storage.is_derived(target_period) + assert run(simulation, "calculate", target) == [7.0] + assert holder._stores_user_input(target_period, "default") + assert (PROBE, "default", target_period) in simulation._user_input_keys