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

changelog.d/fix-stock-option-caches.fixed.md

@@ -0,0 +1 @@
+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.

policyengine_core/simulations/simulation.py

@@ -1342,12 +1342,7 @@ def calculate_add(
         result = sum(
             self.calculate(variable_name, sub_period) for sub_period in sub_periods
         )
-        # Cache only a sum over several sub-periods, as derived from them. A
-        # single sub-period's value is already stored, as an input or
-        # derived, by ``calculate``.
-        if len(sub_periods) > 1:
-            holder = self.get_holder(variable.name)
-            holder.put_in_cache(result, period, self.branch_name, derived=True)
+        self._cache_option_result(variable, period, result)
         return result
 
     def calculate_divide(
@@ -1379,8 +1374,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, derived=True)
+            self._cache_option_result(variable, period, result)
             return result
         elif period.unit == periods.YEAR:
             return self.calculate(variable_name, period)
@@ -1391,6 +1385,55 @@ 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).
+
+        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
+        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
+        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
+        self.get_holder(variable.name).put_in_cache(
+            result, period, self.branch_name, derived=True
+        )
+
     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,268 @@
+"""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 policyengine_core import periods
+from policyengine_core.tools.simulation_dumper import (
+    dump_simulation,
+    restore_simulation,
+)
+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 holder.is_derived(periods.period("2012"))
+    assert periods.period("2012") not in holder.get_input_periods()
+    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 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)]
+
+
+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
+
+
+@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, 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
+):
+    # 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)
+    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`` 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"))
+    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]