Cache ADD and DIVIDE results only where a plain read returns them

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.

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,59 @@ 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 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
+        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
+        if self._reads_input_at(variable.name, period):
+            return
+        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:
         """
         Calculate the value of a variable using the ``calculate_output`` attribute of the variable.

tests/core/test_option_result_cache.py

@@ -0,0 +1,176 @@
+"""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)]
+
+
+@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]

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)