From 65b1656f93f89f2c3992f491b01bc2e1b63b6943 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 08:06:04 -0400 Subject: [PATCH 1/2] 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/2] 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]