From 22a9b3134c3dbce7477e9bf590170e5cbdb9d199 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 18:08:17 -0400 Subject: [PATCH 1/9] Count only inputs as already set when a set_input helper splits a longer period set_input_divide_by_period and set_input_dispatch_by_period treated any value stored for a sub-period as already set, including values the simulation had calculated (a cached default or formula result). The same annual input then gave different months depending on what was calculated first: a divided input skipped calculated months and shared the rest among the others, or failed as inconsistent after the year had been read; a dispatched input reused a calculated month for every later month. The helpers now read the simulation's record of inputs (_user_input_keys): a recorded sub-period keeps its input (and, for dispatch, passes it on to the later sub-periods, as before), and a calculated one is replaced. After storing, they drop the variable's calculated values over overlapping periods (the sums calculate_add caches, the twelfths calculate_divide caches) under the input's branch and the branches it reads through, record what they store as inputs, and evict calculate's fast cache for those periods. Co-Authored-By: Claude Opus 5.5 --- changelog.d/set-input-helper-order.fixed.md | 1 + policyengine_core/holders/helpers.py | 195 +++++++- tests/core/test_set_input_helper_order.py | 469 ++++++++++++++++++ .../test_set_input_helper_order_property.py | 254 ++++++++++ tests/fixtures/set_input_helper_order.py | 337 +++++++++++++ 5 files changed, 1235 insertions(+), 21 deletions(-) create mode 100644 changelog.d/set-input-helper-order.fixed.md create mode 100644 tests/core/test_set_input_helper_order.py create mode 100644 tests/core/test_set_input_helper_order_property.py create mode 100644 tests/fixtures/set_input_helper_order.py diff --git a/changelog.d/set-input-helper-order.fixed.md b/changelog.d/set-input-helper-order.fixed.md new file mode 100644 index 000000000..37e80fa43 --- /dev/null +++ b/changelog.d/set-input-helper-order.fixed.md @@ -0,0 +1 @@ +`set_input_divide_by_period` and `set_input_dispatch_by_period` now leave only inputs in place when they spread a longer period's input over a variable's months or years: a value the simulation calculated for a sub-period (a default, a formula result) is replaced, and a sum or twelfth of the variable it calculated over an overlapping period is dropped, so the same input gives the same values whatever was calculated before it was set (previously a calculated month kept its value and changed the others' share, or made the input fail as inconsistent). diff --git a/policyengine_core/holders/helpers.py b/policyengine_core/holders/helpers.py index 790f136d1..d8f7f8bb2 100644 --- a/policyengine_core/holders/helpers.py +++ b/policyengine_core/holders/helpers.py @@ -1,4 +1,5 @@ import logging +from typing import List import numpy from numpy.typing import ArrayLike @@ -22,12 +23,155 @@ def get_stored_array(holder: Holder, period: Period, branch_name: str) -> ArrayL return holder._get_array_from_storage(period, branch_name) +def _is_input(holder: Holder, period: Period, branch_name: str) -> bool: + """Whether the value stored for ``period`` under ``branch_name`` was set + as an input, rather than calculated by the simulation (a formula result, + or a default, carried-over or uprated value). + + The simulation's record of inputs is ``_user_input_keys``. A holder with + no simulation, or a simulation with no record, cannot tell the two apart, + so every stored value counts as an input there. + """ + simulation = getattr(holder, "simulation", None) + input_keys = getattr(simulation, "_user_input_keys", None) + if input_keys is None: + return True + name = holder.variable.name + if (name, branch_name, period) in input_keys: + return True + # Storage keys twelve months from the first of a month as that year, so + # an input set for ``month:2013-01:12`` is stored under ``2013``. + unit, start, size = period + if unit == periods.YEAR and size == 1: + return (name, branch_name, Period((periods.MONTH, start, 12))) in input_keys + return False + + +def _get_input_array(holder: Holder, period: Period, branch_name: str) -> ArrayLike: + """The input stored for ``period`` under ``branch_name``, or ``None`` if + nothing is stored there or the value was calculated.""" + # Checked before reading, so that a calculated value a branch still + # shares with its parent is not copied only to be replaced. + if not _is_input(holder, period, branch_name): + return None + return get_stored_array(holder, period, branch_name) + + +def _store_input( + holder: Holder, period: Period, array: ArrayLike, branch_name: str +) -> None: + """Store ``array`` as the input for ``period``, replacing any value the + simulation calculated for it. + + The value is added to the simulation's record of inputs even when the + helper is called directly instead of through ``Holder.set_input``, and + ``calculate``'s fast cache drops the value it held for the period. + """ + holder._set(period, array, branch_name) + simulation = getattr(holder, "simulation", None) + input_keys = getattr(simulation, "_user_input_keys", None) + if input_keys is not None: + input_keys.add((holder.variable.name, branch_name, period)) + fast_cache = getattr(simulation, "_fast_cache", None) + if fast_cache: + fast_cache.pop((holder.variable.name, period), None) + + +def _branches_read_with(holder: Holder, branch_name: str) -> List[str]: + """``branch_name`` and the branches whose values the holder's simulation + reads when ``branch_name`` stores none: its ancestors, then ``default``.""" + simulation = getattr(holder, "simulation", None) + get_visible_branch_names = getattr(simulation, "_get_visible_branch_names", None) + if get_visible_branch_names is not None: + visible_branch_names = get_visible_branch_names() + if branch_name in visible_branch_names: + return visible_branch_names + return [branch_name] + + +def _drop_calculated_overlapping( + holder: Holder, period: Period, branch_name: str +) -> None: + """Drop what the simulation calculated for the variable over ``period``, + once a helper has set inputs for its sub-periods. + + Besides the sub-periods the helper replaced, the simulation may hold + values it calculated from them: ``calculate_add`` caches the sum of a + variable over a longer period (a monthly variable's ``2013``), and + ``calculate_divide`` a twelfth of a yearly variable at a month. Any such + value for a period that overlaps ``period`` is out of date, and + ``calculate`` would return it instead of adding up or dividing the + input. Inputs are kept. + + This holder's storage is searched under ``branch_name`` and under the + branches read with it (see ``_branches_read_with``): a branch starts with + its own index of what its parent calculated, and reads those values + until they are dropped. The simulation the branch was created from keeps + its values. + """ + branch_names = set(_branches_read_with(holder, branch_name)) + # Memory keys are "{branch}:{period}"; branch names cannot contain ":". + # Disk keys are "{branch}_{period}"; branch names can contain "_", but + # period strings cannot, so the period follows the last "_". + memory = holder._memory_storage + for key in list(memory._arrays): + stored_branch_name, period_string = key.split(":", 1) + if _is_calculated_overlapping( + holder, stored_branch_name, period_string, branch_names, period + ): + del memory._arrays[key] + memory._shared.discard(key) + disk = holder._disk_storage + if disk is not None: + for key in list(disk._files): + stored_branch_name, period_string = key.rsplit("_", 1) + if _is_calculated_overlapping( + holder, stored_branch_name, period_string, branch_names, period + ): + del disk._files[key] + + +def _is_calculated_overlapping( + holder: Holder, + stored_branch_name: str, + period_string: str, + branch_names: set, + period: Period, +) -> bool: + if stored_branch_name not in branch_names: + return False + try: + stored_period = periods.period(period_string) + except ValueError: + # Not a key ``put`` wrote (say, a file ``restore`` found). + return False + if stored_period.unit == periods.ETERNITY: + return False + if stored_period.start > period.stop or stored_period.stop < period.start: + return False + return not _is_input(holder, stored_period, stored_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. In this case, the variable will accept inputs on larger periods that its definition period, and the value for the larger period will be applied to all its subperiods. + Only inputs count as already known. A sub-period whose value the + simulation calculated (a formula result, or a default, carried-over or + uprated value) takes the new input like a sub-period with no value, so + the result does not depend on what was calculated before the input was + set. + + A sub-period that already has an input keeps it, and that input, not the + value given for the larger period, is applied to the sub-periods after + it that have none: with ``3`` set for March, setting ``7`` for the year + gives ``7`` in January and February and ``3`` from March to December. + The value of a stock known at some month holds for the rest of the + period. (Setting the two inputs in the other order gives ``7`` in every + month but March.) + To read more about ``set_input`` attributes, check the `documentation `_. """ array = holder._to_array(array, validate_nan=True) @@ -46,18 +190,22 @@ 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 + # Store the input data, skipping the sub-periods that already have an input branch_name = get_input_branch(holder) sub_period = period.start.period(cached_period_unit) + stored = False while sub_period.start < after_instant: - existing_array = get_stored_array(holder, sub_period, branch_name) - if existing_array is None: - holder._set(sub_period, array, branch_name) + existing_input = _get_input_array(holder, sub_period, branch_name) + if existing_input is None: + _store_input(holder, sub_period, array, branch_name) + stored = True else: - # The array of the current sub-period is reused for the next ones. - # TODO: refactor or document this behavior - array = existing_array + # The input of the current sub-period is applied to the next + # ones (see the docstring). + array = existing_input sub_period = sub_period.offset(1) + if stored: + _drop_calculated_overlapping(holder, period, branch_name) def set_input_divide_by_period(holder: Holder, period: Period, array: ArrayLike): @@ -66,6 +214,13 @@ def set_input_divide_by_period(holder: Holder, period: Period, array: ArrayLike) In this case, the variable will accept inputs on larger periods that its definition period, and the value for the larger period will be divided between its subperiods. + Sub-periods that already have an input keep it, and what is left of the + value is divided between the others. Only inputs count as already known. + A sub-period whose value the simulation calculated (a formula result, or + a default, carried-over or uprated value) takes its share of the new + input like a sub-period with no value, so the result does not depend on + what was calculated before the input was set. + To read more about ``set_input`` attributes, check the `documentation `_. """ if not isinstance(array, numpy.ndarray): @@ -85,27 +240,25 @@ def set_input_divide_by_period(holder: Holder, period: Period, array: ArrayLike) after_instant = period.start.offset(period_size, period_unit) - # Count the number of elementary periods to change, and the difference with what is already known. + # Find the elementary periods to change, and the difference with the inputs already known. branch_name = get_input_branch(holder) remaining_array = array.copy() sub_period = period.start.period(cached_period_unit) - sub_periods_count = 0 + sub_periods_to_set = [] while sub_period.start < after_instant: - existing_array = get_stored_array(holder, sub_period, branch_name) - if existing_array is not None: - remaining_array -= existing_array + existing_input = _get_input_array(holder, sub_period, branch_name) + if existing_input is not None: + remaining_array -= existing_input else: - sub_periods_count += 1 + sub_periods_to_set.append(sub_period) sub_period = sub_period.offset(1) - # Cache the input data - if sub_periods_count > 0: - 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: - holder._set(sub_period, divided_array, branch_name) - sub_period = sub_period.offset(1) + # Store the input data + if sub_periods_to_set: + divided_array = remaining_array / len(sub_periods_to_set) + for sub_period in sub_periods_to_set: + _store_input(holder, sub_period, divided_array, branch_name) + _drop_calculated_overlapping(holder, period, branch_name) elif not (remaining_array == 0).all(): raise ValueError( "Inconsistent input: variable {0} has already been set for all months contained in period {1}, and value {2} provided for {1} doesn't match the total ({3}). This error may also be thrown if you try to call set_input twice for the same variable and period.".format( diff --git a/tests/core/test_set_input_helper_order.py b/tests/core/test_set_input_helper_order.py new file mode 100644 index 000000000..636cf5768 --- /dev/null +++ b/tests/core/test_set_input_helper_order.py @@ -0,0 +1,469 @@ +"""An input set over a longer period does not depend on what was calculated first. + +``set_input_divide_by_period`` and ``set_input_dispatch_by_period`` spread an +input given for a longer period over a variable's own periods, leaving the +sub-periods that already have an input as they are. They used to treat any +stored value as such an input, including one the simulation had calculated +(a default, a formula result), so the same input gave different values +depending on what had been calculated before it was set. + +``test_set_input_helper_order_property.py`` checks the same rule over random +sequences. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.holders import ( + set_input_dispatch_by_period, + set_input_divide_by_period, +) +from policyengine_core.simulations import SimulationBuilder +from tests.fixtures.set_input_helper_order import ( + BRANCH_NAME, + MONTHS, + Status, + build_simulation, + build_system, + read, +) + +MONTHS_2013 = MONTHS[:12] +YEARLY_INPUT = np.array([1200.0, 2400.0], dtype=np.float32) +MONTHLY_SHARE = np.array([100.0, 200.0], dtype=np.float32) +STATUS_INDEX = {status.name: status.index for status in Status} + + +def assert_reads(simulation, name, expected_by_period): + for period, expected in expected_by_period.items(): + np.testing.assert_array_equal( + read(simulation, name, period), expected, err_msg=f"{name} {period}" + ) + + +# Divided inputs + + +def test_divided_input_replaces_a_month_calculated_before_it(): + simulation = build_simulation() + assert_reads(simulation, "flow_m", {"2013-01": [0, 0]}) + + simulation.set_input("flow_m", "2013", YEARLY_INPUT) + + assert_reads( + simulation, + "flow_m", + {**{month: MONTHLY_SHARE for month in MONTHS_2013}, "2013": YEARLY_INPUT}, + ) + + +def test_divided_input_is_accepted_after_the_year_was_calculated(): + simulation = build_simulation() + assert_reads(simulation, "flow_m", {"2013": [0, 0]}) + + # Every month now holds a calculated default, and the year their sum. + simulation.set_input("flow_m", "2013", YEARLY_INPUT) + + assert_reads( + simulation, + "flow_m", + {"2013-05": MONTHLY_SHARE, "2013-12": MONTHLY_SHARE, "2013": YEARLY_INPUT}, + ) + + +def test_divided_input_keeps_month_inputs(): + simulation = build_simulation() + simulation.set_input("flow_m", "2013-01", np.array([300.0, 600.0])) + + simulation.set_input("flow_m", "2013", YEARLY_INPUT) + + rest = (np.array([900.0, 1800.0], dtype=np.float32) / 11).astype(np.float32) + assert_reads( + simulation, + "flow_m", + {"2013-01": [300, 600], **{month: rest for month in MONTHS_2013[1:]}}, + ) + + +def test_divided_input_keeps_month_inputs_and_replaces_calculated_months(): + simulation = build_simulation() + simulation.set_input("flow_m", "2013-01", np.array([300.0, 600.0])) + assert_reads(simulation, "flow_m", {"2013-02": [0, 0], "2013-01": [300, 600]}) + + simulation.set_input("flow_m", "2013", YEARLY_INPUT) + + rest = (np.array([900.0, 1800.0], dtype=np.float32) / 11).astype(np.float32) + assert_reads( + simulation, + "flow_m", + {"2013-01": [300, 600], **{month: rest for month in MONTHS_2013[1:]}}, + ) + + +def test_divided_input_is_set_from_a_situation_with_a_month_and_its_year(): + simulation = SimulationBuilder().build_from_entities( + build_system(), + {"persons": {"a": {"flow_m": {"2013": 1200, "2013-01": 300}}}}, + ) + + rest = (np.array([900.0], dtype=np.float32) / 11).astype(np.float32) + assert_reads(simulation, "flow_m", {"2013-01": [300], "2013-07": rest}) + + +def test_divided_input_that_contradicts_the_inputs_is_refused(): + simulation = build_simulation() + simulation.set_input("flow_m", "2013", YEARLY_INPUT) + # Reading the months does not turn them into calculated values. + assert_reads(simulation, "flow_m", {month: MONTHLY_SHARE for month in MONTHS_2013}) + + with pytest.raises(ValueError, match="Inconsistent input"): + simulation.set_input("flow_m", "2013", YEARLY_INPUT + 100) + + simulation.set_input("flow_m", "2013", YEARLY_INPUT) + assert_reads(simulation, "flow_m", {"2013-06": MONTHLY_SHARE, "2013": YEARLY_INPUT}) + + +def test_divided_input_replaces_formula_results(): + simulation = build_simulation() + assert_reads( + simulation, "formula_flow_m", {"2013-01": [11, 11], "2013": [198, 198]} + ) + + simulation.set_input("formula_flow_m", "2013", YEARLY_INPUT) + + assert_reads( + simulation, + "formula_flow_m", + {"2013-01": MONTHLY_SHARE, "2013": YEARLY_INPUT, "2014-01": [11, 11]}, + ) + + +def test_divided_input_to_a_yearly_variable_replaces_a_calculated_year(): + simulation = build_simulation() + assert_reads(simulation, "flow_y", {"2013": [0, 0]}) + + simulation.set_input("flow_y", "month:2013-01:24", YEARLY_INPUT * 2) + + assert_reads(simulation, "flow_y", {"2013": YEARLY_INPUT, "2014": YEARLY_INPUT}) + + +def test_divided_input_to_a_yearly_variable_drops_the_twelfth_calculated_before(): + simulation = build_simulation() + # A month of a yearly flow is a twelfth of the year, cached at the month. + assert_reads(simulation, "flow_y", {"2013-05": [0, 0]}) + + simulation.set_input("flow_y", "month:2013-01:12", YEARLY_INPUT) + + assert_reads(simulation, "flow_y", {"2013": YEARLY_INPUT, "2013-05": MONTHLY_SHARE}) + + +# Dispatched inputs + + +@pytest.mark.parametrize("calculated_first", ["2013-01", "2013-05", "2013-12", "2013"]) +def test_dispatched_input_replaces_months_calculated_before_it(calculated_first): + simulation = build_simulation() + assert_reads(simulation, "count_m", {calculated_first: [0, 0]}) + + simulation.set_input("count_m", "2013", np.array([7, 9])) + + assert_reads( + simulation, + "count_m", + {**{month: [7, 9] for month in MONTHS_2013}, "2013": [7, 9]}, + ) + + +def test_dispatched_input_applies_an_earlier_month_input_to_the_months_after_it(): + simulation = build_simulation() + simulation.set_input("count_m", "2013-03", np.array([3, 4])) + + simulation.set_input("count_m", "2013", np.array([7, 9])) + + assert_reads( + simulation, + "count_m", + { + "2013-01": [7, 9], + "2013-02": [7, 9], + **{month: [3, 4] for month in MONTHS_2013[2:]}, + }, + ) + + +def test_dispatched_input_keeps_month_inputs_and_replaces_calculated_months(): + simulation = build_simulation() + simulation.set_input("count_m", "2013-03", np.array([3, 4])) + assert_reads(simulation, "count_m", {"2013-02": [0, 0], "2013-06": [0, 0]}) + + simulation.set_input("count_m", "2013", np.array([7, 9])) + + assert_reads( + simulation, + "count_m", + { + "2013-01": [7, 9], + "2013-02": [7, 9], + **{month: [3, 4] for month in MONTHS_2013[2:]}, + }, + ) + + +def test_dispatched_input_changes_nothing_when_every_month_has_an_input(): + simulation = build_simulation() + simulation.set_input("count_m", "2013", np.array([7, 9])) + assert_reads(simulation, "count_m", {month: [7, 9] for month in MONTHS_2013}) + + simulation.set_input("count_m", "2013", np.array([1, 2])) + + assert_reads(simulation, "count_m", {month: [7, 9] for month in MONTHS_2013}) + + +def test_dispatched_flow_input_drops_the_sum_calculated_before_it(): + simulation = build_simulation() + assert_reads(simulation, "dispatched_flow_m", {"2013": [0, 0]}) + + simulation.set_input("dispatched_flow_m", "2013", np.array([5.0, 6.0])) + + assert_reads(simulation, "dispatched_flow_m", {"2013-04": [5, 6], "2013": [60, 72]}) + + +def test_dispatched_input_replaces_calculated_booleans_and_enums(): + simulation = build_simulation() + assert_reads(simulation, "flag_m", {"2013-01": [False, False]}) + assert_reads(simulation, "status_m", {"2013-01": [STATUS_INDEX["none"]] * 2}) + + simulation.set_input("flag_m", "2013", np.array([True, False])) + simulation.set_input("status_m", "2013", np.array(["all", "some"])) + + assert_reads( + simulation, "flag_m", {"2013-01": [True, False], "2013-08": [True, False]} + ) + assert_reads( + simulation, + "status_m", + { + month: [STATUS_INDEX["all"], STATUS_INDEX["some"]] + for month in ("2013-01", "2013-08") + }, + ) + + +def test_dispatched_input_to_a_yearly_variable_replaces_a_calculated_year(): + simulation = build_simulation() + assert_reads(simulation, "count_y", {"2013": [0, 0], "2014-06": [0, 0]}) + + simulation.set_input("count_y", "month:2013-01:24", np.array([7, 9])) + + assert_reads( + simulation, "count_y", {"2013": [7, 9], "2014": [7, 9], "2014-06": [7, 9]} + ) + + +# What the helpers leave alone + + +def test_values_calculated_outside_the_input_period_are_kept(): + simulation = build_simulation() + assert_reads(simulation, "flow_m", {"2014-03": [0, 0], "2014": [0, 0]}) + assert_reads(simulation, "flow_m", {"2012-12": [0, 0], "2012": [0, 0]}) + holder = simulation.get_holder("flow_m") + + simulation.set_input("flow_m", "2013", YEARLY_INPUT) + + for period in ("2012-12", "2012", "2014-03", "2014"): + assert holder.get_array(period) is not None, period + + +def test_a_sum_calculated_over_a_period_that_overlaps_the_input_is_dropped(): + simulation = build_simulation() + assert_reads(simulation, "flow_m", {"year:2013:2": [0, 0]}) + + simulation.set_input("flow_m", "2014", YEARLY_INPUT) + + assert_reads( + simulation, "flow_m", {"year:2013:2": YEARLY_INPUT, "2014": YEARLY_INPUT} + ) + + +def test_an_input_stored_at_a_longer_period_is_kept(): + simulation = build_simulation() + # Twelve months from January have the unit of a monthly variable, so the + # input is stored as given, under the year. + simulation.set_input("flow_m", "month:2013-01:12", np.array([50.0, 60.0])) + + simulation.set_input("flow_m", "year:2013:2", YEARLY_INPUT * 2) + + holder = simulation.get_holder("flow_m") + np.testing.assert_array_equal(holder.get_array("2013"), [50, 60]) + assert_reads( + simulation, "flow_m", {"2013-01": MONTHLY_SHARE, "2014-12": MONTHLY_SHARE} + ) + + +# The record of inputs, and calculate's fast cache + + +def test_a_replaced_value_is_recorded_as_an_input(): + simulation = build_simulation() + assert_reads(simulation, "flow_m", {"2013-01": [0, 0]}) + + simulation.set_input("flow_m", "2013", YEARLY_INPUT) + + january = periods.period("2013-01") + assert ("flow_m", "default", january) in simulation._user_input_keys + # Inputs survive the cache wipe that follows a reform; calculated values do not. + simulation._invalidate_all_caches() + np.testing.assert_array_equal( + simulation.get_holder("flow_m").get_array(january), MONTHLY_SHARE + ) + + +def test_calculate_returns_the_input_after_a_repeated_read(): + simulation = build_simulation() + # The second read is answered from calculate's fast cache. + assert_reads(simulation, "flow_m", {"2013-01": [0, 0]}) + assert_reads(simulation, "flow_m", {"2013-01": [0, 0]}) + assert_reads(simulation, "count_m", {"2013-01": [0, 0]}) + assert_reads(simulation, "count_m", {"2013-01": [0, 0]}) + + simulation.set_input("flow_m", "2013", YEARLY_INPUT) + simulation.set_input("count_m", "2013", np.array([7, 9])) + + assert_reads(simulation, "flow_m", {"2013-01": MONTHLY_SHARE}) + assert_reads(simulation, "count_m", {"2013-01": [7, 9]}) + + +def test_a_helper_called_directly_stores_inputs(): + simulation = build_simulation() + assert_reads(simulation, "flow_m", {"2013-01": [0, 0]}) + assert_reads(simulation, "count_m", {"2013-01": [0, 0]}) + year = periods.period("2013") + + set_input_divide_by_period(simulation.get_holder("flow_m"), year, YEARLY_INPUT) + set_input_dispatch_by_period( + simulation.get_holder("count_m"), year, np.array([7, 9]) + ) + + january = periods.period("2013-01") + assert ("flow_m", "default", january) in simulation._user_input_keys + assert ("count_m", "default", january) in simulation._user_input_keys + assert_reads(simulation, "flow_m", {"2013-01": MONTHLY_SHARE}) + assert_reads(simulation, "count_m", {"2013-01": [7, 9]}) + # A second call finds those inputs, as it would after ``set_input``. + with pytest.raises(ValueError, match="Inconsistent input"): + set_input_divide_by_period( + simulation.get_holder("flow_m"), year, YEARLY_INPUT + 100 + ) + set_input_dispatch_by_period( + simulation.get_holder("count_m"), year, np.array([1, 2]) + ) + assert_reads(simulation, "count_m", {"2013-01": [7, 9]}) + + +def test_every_stored_value_counts_as_an_input_without_a_record(): + simulation = build_simulation() + holder = simulation.get_holder("flow_m") + holder.put_in_cache(np.array([300.0, 600.0]), periods.period("2013-01")) + # A simulation with no record of its inputs cannot tell them apart from + # what it calculated. + del simulation._user_input_keys + + set_input_divide_by_period(holder, periods.period("2013"), YEARLY_INPUT) + + rest = (np.array([900.0, 1800.0], dtype=np.float32) / 11).astype(np.float32) + np.testing.assert_array_equal(holder.get_array("2013-01"), [300, 600]) + np.testing.assert_array_equal(holder.get_array("2013-02"), rest) + + +# Branches + + +def test_input_on_a_branch_drops_the_sum_its_parent_calculated(): + simulation = build_simulation() + assert_reads(simulation, "flow_m", {"2013": [0, 0]}) + branch = simulation.get_branch(BRANCH_NAME) + + branch.set_input("flow_m", "2013", YEARLY_INPUT) + + assert_reads(branch, "flow_m", {"2013-01": MONTHLY_SHARE, "2013": YEARLY_INPUT}) + # The simulation the branch came from keeps what it calculated. + assert simulation.get_holder("flow_m").get_array("2013") is not None + assert_reads(simulation, "flow_m", {"2013-01": [0, 0], "2013": [0, 0]}) + + +def test_input_on_a_branch_replaces_what_the_branch_calculated(): + simulation = build_simulation() + branch = simulation.get_branch(BRANCH_NAME) + assert_reads(branch, "flow_m", {"2013-01": [0, 0], "2013": [0, 0]}) + assert_reads(branch, "count_m", {"2013-06": [0, 0]}) + + branch.set_input("flow_m", "2013", YEARLY_INPUT) + branch.set_input("count_m", "2013", np.array([7, 9])) + + assert_reads(branch, "flow_m", {"2013-01": MONTHLY_SHARE, "2013": YEARLY_INPUT}) + assert_reads(branch, "count_m", {"2013-06": [7, 9], "2013-12": [7, 9]}) + assert_reads(simulation, "flow_m", {"2013-01": [0, 0]}) + assert_reads(simulation, "count_m", {"2013-06": [0, 0]}) + + +def test_input_stored_under_default_replaces_what_the_branch_calculated(): + simulation = build_simulation() + branch = simulation.get_branch(BRANCH_NAME) + # The branch stores what it calculates under its own name, which it + # reads before ``default``. + assert_reads(branch, "flow_m", {"2013-01": [0, 0], "2013": [0, 0]}) + + # ``Holder.set_input`` stores under ``default`` unless given a branch. + branch.get_holder("flow_m").set_input(periods.period("2013"), YEARLY_INPUT) + + assert_reads(branch, "flow_m", {"2013-01": MONTHLY_SHARE, "2013": YEARLY_INPUT}) + + +# Values stored on disk + + +def test_divided_input_replaces_values_calculated_and_stored_on_disk(): + simulation = build_simulation(on_disk=True) + assert_reads(simulation, "flow_m", {"2013-01": [0, 0], "2013": [0, 0]}) + holder = simulation.get_holder("flow_m") + assert holder._disk_storage.get(periods.period("2013")) is not None + + simulation.set_input("flow_m", "2013", YEARLY_INPUT) + + assert_reads(simulation, "flow_m", {"2013-01": MONTHLY_SHARE, "2013": YEARLY_INPUT}) + + +def test_dispatched_input_replaces_values_calculated_and_stored_on_disk(): + simulation = build_simulation(on_disk=True) + assert_reads(simulation, "count_m", {"2013-05": [0, 0]}) + assert_reads(simulation, "dispatched_flow_m", {"2013": [0, 0]}) + + simulation.set_input("count_m", "2013", np.array([7, 9])) + simulation.set_input("dispatched_flow_m", "2013", np.array([5.0, 6.0])) + + assert_reads(simulation, "count_m", {"2013-05": [7, 9], "2013-09": [7, 9]}) + assert_reads(simulation, "dispatched_flow_m", {"2013-04": [5, 6], "2013": [60, 72]}) + + +def test_input_on_a_branch_drops_a_sum_stored_on_disk_under_its_name(): + simulation = build_simulation(on_disk=True) + # The branch name contains the separator of disk storage keys. (Disk + # storage cannot list the periods stored under such a name, so the sum + # is stored directly instead of being calculated on the branch.) + assert "_" in BRANCH_NAME + branch = simulation.get_branch(BRANCH_NAME) + holder = branch.get_holder("flow_m") + year = periods.period("2013") + holder._disk_storage.put(np.array([5.0, 5.0], dtype=np.float32), year, BRANCH_NAME) + + branch.set_input("flow_m", "2013", YEARLY_INPUT) + + assert holder._disk_storage.get(year, BRANCH_NAME) is None + np.testing.assert_array_equal( + holder._disk_storage.get(periods.period("2013-01"), BRANCH_NAME), + MONTHLY_SHARE, + ) diff --git a/tests/core/test_set_input_helper_order_property.py b/tests/core/test_set_input_helper_order_property.py new file mode 100644 index 000000000..7294fa34e --- /dev/null +++ b/tests/core/test_set_input_helper_order_property.py @@ -0,0 +1,254 @@ +"""What a simulation calculated before its inputs were set does not change them. + +Each example sets some inputs, has one simulation calculate a random list of +requests, then sets more inputs on that simulation and on one that +calculated nothing, and compares the two. It also compares both with a model +that holds inputs only (``tests/fixtures/set_input_helper_order.py``). +``test_set_input_helper_order.py`` pins the same behaviour with examples. + +Outside the properties, because they are other order dependences: + +- Other variables. A value calculated earlier from the variable an input is + set for stays cached, as after any ``set_input``, so the properties read + only the variables the inputs are set for, whose one formula reads nothing. +- Auto-carry-over and uprating, which read whatever periods are stored + (policyengine-core#562, #563): the test system uses neither. +- ``calculate_add`` and ``calculate_divide`` called directly, whose results + are cached where a plain read does not return them (#571): the requests + are plain ``calculate`` calls. +- An input set for one of the variable's own periods (one month, or one + year) after a sum or a twelfth of the variable was calculated: the helpers + are not involved and that sum stays cached (policyengine-core#579). The + periods of another size are compared only for variables given no such + input after the calculations. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +# The country-package smoke job installs no dev dependencies. +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from tests.fixtures.set_input_helper_order import ( + BRANCH_NAME, + COUNT, + DIVIDED, + FLOWS, + HELPER_PERIODS, + NAMES, + OTHER_READS, + OWN_PERIODS, + OWN_READS, + Reference, + Status, + apply_input, + build_simulation, + build_system, + input_record, + read, + read_all, + same_arrays, + sub_periods, +) + +SYSTEM = build_system() +STATUSES = [status.name for status in Status] + + +def _values(name): + if name == "flag_m": + element = st.booleans() + elif name == "status_m": + element = st.sampled_from(STATUSES) + elif name in ("count_m", "count_y"): + element = st.integers(-50, 50) + else: + # Whole numbers, so that every stored float32 is the same on any + # platform; the shares of a divided input are not whole. + element = st.integers(-1200, 1200).map(float) + return st.lists(element, min_size=COUNT, max_size=COUNT) + + +def _input(periods_of): + return st.sampled_from(NAMES).flatmap( + lambda name: st.tuples( + st.just(name), st.sampled_from(periods_of(name)), _values(name) + ) + ) + + +own_input = _input(lambda name: OWN_PERIODS[name]) +helper_input = _input(lambda name: HELPER_PERIODS[name]) +any_input = st.one_of(own_input, helper_input, helper_input) +request = st.sampled_from(NAMES).flatmap( + lambda name: st.tuples( + st.just(name), st.sampled_from(OWN_READS[name] + OTHER_READS[name]) + ) +) + +SETTINGS = dict( + deadline=None, + suppress_health_check=[ + hypothesis.HealthCheck.too_slow, + hypothesis.HealthCheck.data_too_large, + ], +) + + +def _run(first_inputs, requests_before_branch, requests, later_inputs, on_disk, branch): + """Run the example; return the simulation the later inputs were set on, + each later input's outcome, and the model's.""" + simulation = build_simulation(SYSTEM, on_disk=on_disk) + reference = Reference() + for name, period, values in first_inputs: + assert apply_input(simulation, name, period, values) == reference.set_input( + name, period, values + ) + for name, period in requests_before_branch: + read(simulation, name, period) + branch_name = "default" + if branch: + simulation = simulation.get_branch(BRANCH_NAME) + branch_name = BRANCH_NAME + for name, period in requests: + read(simulation, name, period) + outcomes = [ + apply_input(simulation, name, period, values) + for name, period, values in later_inputs + ] + expected = [ + reference.set_input(name, period, values, branch_name) + for name, period, values in later_inputs + ] + return simulation, outcomes, expected, reference, branch_name + + +@hypothesis.settings(max_examples=500, **SETTINGS) +@hypothesis.given( + first_inputs=st.lists(any_input, max_size=4), + requests_before_branch=st.lists(request, max_size=6), + requests=st.lists(request, max_size=8), + later_inputs=st.lists(any_input, min_size=1, max_size=6), + on_disk=st.booleans(), + branch=st.booleans(), +) +def test_inputs_do_not_depend_on_what_was_calculated_before( + first_inputs, requests_before_branch, requests, later_inputs, on_disk, branch +): + # Branches of a simulation that stores on disk share its files + # (policyengine-core#558), so the two are not combined. + on_disk = on_disk and not branch + calculated, outcomes, expected, reference, branch_name = _run( + first_inputs, requests_before_branch, requests, later_inputs, on_disk, branch + ) + fresh, fresh_outcomes, _, _, _ = _run( + first_inputs, [], [], later_inputs, on_disk, branch + ) + + # The same inputs are accepted or refused, and recorded. + assert outcomes == fresh_outcomes == expected + assert input_record(calculated) == input_record(fresh) + + # A sum or a twelfth calculated before an input for one of the + # variable's own periods stays cached (policyengine-core#579). + given_own_period_input = { + name for name, period, _ in later_inputs if period in OWN_PERIODS[name] + } + other_periods_of = [name for name in NAMES if name not in given_own_period_input] + values = read_all(calculated, other_periods_of=other_periods_of) + fresh_values = read_all(fresh, other_periods_of=other_periods_of) + assert values.keys() == fresh_values.keys() + for key in values: + assert same_arrays(values[key], fresh_values[key]), ( + key, + values[key], + fresh_values[key], + ) + + # Both hold the inputs the inputs-only model holds. + for name in NAMES: + for period in OWN_READS[name]: + expected_value = reference.value(name, period, branch_name) + assert same_arrays(values[(name, period)], expected_value), ( + name, + period, + values[(name, period)], + expected_value, + ) + + +@hypothesis.settings(max_examples=300, **SETTINGS) +@hypothesis.given( + first_inputs=st.lists(own_input, max_size=4), + requests=st.lists(request, max_size=8), + name=st.sampled_from(DIVIDED), + data=st.data(), +) +def test_a_divided_input_adds_up_to_the_value_given(first_inputs, requests, name, data): + """Conservation: a divided input is accepted while a sub-period is left + without an input, the variable's inputs over the period then add up to + it, and a plain read of the period returns that sum, whatever was + calculated before.""" + period = data.draw(st.sampled_from(HELPER_PERIODS[name])) + values = data.draw(_values(name)) + simulation = build_simulation(SYSTEM) + for input_name, input_period, input_values in first_inputs: + assert apply_input(simulation, input_name, input_period, input_values) is None + for request_name, request_period in requests: + read(simulation, request_name, request_period) + + # The input is refused only when every sub-period already has one. + with_input = { + input_period + for input_name, input_period, _ in first_inputs + if input_name == name + } + hypothesis.assume(set(sub_periods(name, period)) - with_input) + assert apply_input(simulation, name, period, values) is None + + given = np.array(values, dtype=np.float32) + total = sum( + read(simulation, name, sub_period).astype(np.float64) + for sub_period in sub_periods(name, period) + ) + np.testing.assert_allclose(total, given, rtol=1e-5, atol=1e-2) + if name in FLOWS and period in OTHER_READS[name]: + # The same period, read as ``calculate`` adds it up. + np.testing.assert_allclose( + read(simulation, name, period), given, rtol=1e-5, atol=1e-2 + ) + + +@hypothesis.settings(max_examples=300, **SETTINGS) +@hypothesis.given( + first_inputs=st.lists(own_input, max_size=6), + requests=st.lists(request, max_size=8), + later_input=helper_input, +) +def test_a_helper_leaves_the_inputs_it_finds_unchanged( + first_inputs, requests, later_input +): + """An input already set for one of the variable's own periods is never + changed by an input set over a longer period, and every sub-period of + that longer period holds an input afterwards.""" + simulation = build_simulation(SYSTEM) + for name, period, values in first_inputs: + assert apply_input(simulation, name, period, values) is None + for name, period in requests: + read(simulation, name, period) + before = { + (name, period): read(simulation, name, period).copy() + for name, period, _ in first_inputs + } + + name, period, values = later_input + apply_input(simulation, name, period, values) + + for key, value in before.items(): + assert same_arrays(read(simulation, *key), value), key + recorded = input_record(simulation) + for sub_period in sub_periods(name, period): + assert (name, "default", sub_period) in recorded diff --git a/tests/fixtures/set_input_helper_order.py b/tests/fixtures/set_input_helper_order.py new file mode 100644 index 000000000..d61dbf366 --- /dev/null +++ b/tests/fixtures/set_input_helper_order.py @@ -0,0 +1,337 @@ +"""Variables, simulations and a reference model for the ``set_input`` helper tests. + +Shared by ``tests/core/test_set_input_helper_order.py`` (regressions) and +``tests/core/test_set_input_helper_order_property.py`` (properties). + +The reference model (``Reference``) holds inputs only: it has no notion of a +value the simulation calculated. For each variable and branch it keeps the +input of each of the variable's own periods (months or years), and applies +the two helpers' rules to those: + +- divide: the sub-periods that already have an input keep it, and what is + left of the value given for the longer period is divided between the + others (an error if there are none and something is left); +- dispatch: each sub-period with no input takes the value given for the + longer period, or the input of the latest earlier sub-period that had one. + +A simulation agrees with the model only if what it calculated before an +input was set has no effect on the inputs it holds afterwards. +""" + +from __future__ import annotations + +import warnings +from typing import Dict, List, Optional, Tuple + +import numpy as np + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem, entities +from policyengine_core.enums import Enum +from policyengine_core.experimental import MemoryConfig +from policyengine_core.holders import ( + set_input_dispatch_by_period, + set_input_divide_by_period, +) +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables import Variable + +COUNT = 2 +BRANCH_NAME = "what_if" + + +class Status(Enum): + none = "None" + some = "Some" + all = "All" + + +class flow_m(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "Monthly flow input with no formula (divided by default)" + + +class count_m(Variable): + value_type = int + entity = entities.Person + definition_period = periods.MONTH + label = "Monthly integer input with no formula (dispatched by default)" + + +class flag_m(Variable): + value_type = bool + entity = entities.Person + definition_period = periods.MONTH + label = "Monthly boolean input with no formula (dispatched by default)" + + +class status_m(Variable): + value_type = Enum + possible_values = Status + default_value = Status.none + entity = entities.Person + definition_period = periods.MONTH + label = "Monthly enum input with no formula (dispatched by default)" + + +class dispatched_flow_m(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + set_input = set_input_dispatch_by_period + label = "Monthly flow input that repeats a longer period's value in each month" + + +class formula_flow_m(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "Monthly flow with a formula that reads no other variable" + + def formula(person, period, parameters): + return person.filled_array(10.0 + period.start.month) + + +class flow_y(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "Yearly flow input with no formula (divided by default)" + + +class count_y(Variable): + value_type = int + entity = entities.Person + definition_period = periods.YEAR + label = "Yearly integer input with no formula (dispatched by default)" + + +VARIABLES = [ + flow_m, + count_m, + flag_m, + status_m, + dispatched_flow_m, + formula_flow_m, + flow_y, + count_y, +] +NAMES = [variable.__name__ for variable in VARIABLES] +MONTHLY = [ + name + for name in NAMES + if getattr(globals()[name], "definition_period") == periods.MONTH +] +YEARLY = [name for name in NAMES if name not in MONTHLY] +DIVIDED = ["flow_m", "formula_flow_m", "flow_y"] +#: Variables whose plain read over a longer (or shorter) period is a sum (or +#: a twelfth) that ``calculate`` caches at that period. +FLOWS = ["flow_m", "dispatched_flow_m", "formula_flow_m", "flow_y"] + +MONTHS = [f"{year}-{month:02d}" for year in (2013, 2014) for month in range(1, 13)] +YEARS = ["2013", "2014"] + +#: The variable's own periods an input can be set for directly. +OWN_PERIODS = {name: (MONTHS if name in MONTHLY else YEARS) for name in NAMES} +#: Longer periods, which go through the variable's ``set_input`` helper. +HELPER_PERIODS = { + name: ( + ["2013", "2014", "year:2013:2"] + if name in MONTHLY + else ["month:2013-01:12", "month:2014-01:12", "month:2013-01:24"] + ) + for name in NAMES +} +#: Periods read with a plain ``calculate``: the variable's own, with one on +#: each side of the two years inputs are set in. +OWN_READS = { + name: ( + ["2012-12"] + MONTHS + ["2015-01"] + if name in MONTHLY + else ["2012"] + YEARS + ["2015"] + ) + for name in NAMES +} +#: Periods of another size a plain ``calculate`` accepts: a monthly +#: variable over one or two years (its sum if a flow, else its last month), +#: a yearly one at a month (a twelfth if a flow, else the year's value). +OTHER_READS = { + name: ( + ["2013", "2014", "year:2013:2"] + if name in MONTHLY + else ["2013-05", "2014-01", "2014-12"] + ) + for name in NAMES +} + + +def build_system() -> CountryTaxBenefitSystem: + system = CountryTaxBenefitSystem() + system.add_variables(*VARIABLES) + return system + + +def build_simulation(system=None, on_disk: bool = False): + """A simulation of ``COUNT`` people with no inputs. + + With ``on_disk``, every value is stored on disk rather than in memory. + """ + simulation = SimulationBuilder().build_default_simulation( + system or build_system(), count=COUNT + ) + if on_disk: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + # A holder takes its storage from the configuration it is created + # under, and building the simulation creates them all: drop the + # (still empty) ones of the fixture's variables. + for name in NAMES: + del simulation.get_variable_population(name)._holders[name] + return simulation + + +def sub_periods(name: str, period: str) -> List[str]: + """The variable's own periods within ``period``, as a helper walks them.""" + unit = periods.MONTH if name in MONTHLY else periods.YEAR + period = periods.period(period) + after = period.start.offset(period.size, period.unit) + result = [] + sub_period = period.start.period(unit) + while sub_period.start < after: + result.append(str(sub_period)) + sub_period = sub_period.offset(1) + return result + + +def _dtype(name: str): + return { + "flow_m": np.float32, + "dispatched_flow_m": np.float32, + "formula_flow_m": np.float32, + "flow_y": np.float32, + "count_m": np.int32, + "count_y": np.int32, + "flag_m": np.bool_, + "status_m": np.int16, + }[name] + + +def as_input(name: str, values) -> np.ndarray: + """``values`` as the array ``set_input`` takes for the variable.""" + if name == "status_m": + return np.array(values, dtype=str) + return np.array(values, dtype=_dtype(name)) + + +def _stored(name: str, values) -> np.ndarray: + """``values`` as the simulation stores them.""" + if name == "status_m": + return np.array([Status[value].index for value in values], dtype=np.int16) + return np.array(values, dtype=_dtype(name)) + + +class Reference: + """Inputs-only model of ``set_input`` (see the module docstring).""" + + def __init__(self): + # (name, branch, own period) -> array + self.inputs: Dict[Tuple[str, str, str], np.ndarray] = {} + + def copy(self) -> "Reference": + other = Reference() + other.inputs = dict(self.inputs) + return other + + def set_input( + self, name: str, period: str, values, branch: str = "default" + ) -> Optional[str]: + """Apply an input; return ``"ValueError"`` if the helper refuses it.""" + array = _stored(name, values) + if period in OWN_PERIODS[name]: + self.inputs[(name, branch, period)] = array + return None + own = sub_periods(name, period) + if name in DIVIDED: + remaining = array.copy() + free = [] + for sub_period in own: + known = self.inputs.get((name, branch, sub_period)) + if known is not None: + remaining -= known + else: + free.append(sub_period) + if free: + share = (remaining / len(free)).astype(array.dtype) + for sub_period in free: + self.inputs[(name, branch, sub_period)] = share + elif not (remaining == 0).all(): + return "ValueError" + return None + for sub_period in own: + known = self.inputs.get((name, branch, sub_period)) + if known is None: + self.inputs[(name, branch, sub_period)] = array + else: + array = known + return None + + def value(self, name: str, period: str, branch: str = "default") -> np.ndarray: + """What a plain ``calculate`` of one of the variable's own periods + returns: the branch's input, else the default branch's, else the + formula result or the default value.""" + for branch_name in dict.fromkeys([branch, "default"]): + known = self.inputs.get((name, branch_name, period)) + if known is not None: + return known + if name == "formula_flow_m": + month = periods.period(period).start.month + return np.full(COUNT, 10.0 + month, dtype=np.float32) + return np.zeros(COUNT, dtype=_dtype(name)) + + +def apply_input(simulation, name: str, period: str, values) -> Optional[str]: + """``simulation.set_input``; return the error's type name if it raises.""" + try: + simulation.set_input(name, period, as_input(name, values)) + except ValueError as error: + return type(error).__name__ + return None + + +def read(simulation, name: str, period: str) -> np.ndarray: + """A plain ``calculate``, as a plain array (enums as their indices).""" + return np.asarray(simulation.calculate(name, period)) + + +def read_all(simulation, names=NAMES, other_periods_of=NAMES) -> dict: + """Every read of ``names`` in a fixed order: each variable's own periods, + then, for those in ``other_periods_of``, the periods of another size.""" + result = {} + for name in names: + for period in OWN_READS[name]: + result[(name, period)] = read(simulation, name, period) + for name in names: + if name in other_periods_of: + for period in OTHER_READS[name]: + result[(name, period)] = read(simulation, name, period) + return result + + +def input_record(simulation) -> set: + """The simulation's record of inputs for the fixture's variables.""" + return { + (name, branch, str(period)) + for name, branch, period in simulation._user_input_keys + if name in NAMES + } + + +def same_arrays(left: np.ndarray, right: np.ndarray) -> bool: + return ( + left.dtype == right.dtype + and left.shape == right.shape + and np.array_equal(left, right) + ) From af877e804b6c316eabc9f7882be7e766f8cbf23a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 18:13:07 -0400 Subject: [PATCH 2/9] Evict calculate's fast cache for the periods a set_input helper drops calculate does not put sums or twelfths in its fast cache today, so this is defensive: a value held there for a period whose stored value the helper drops goes with it. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/holders/helpers.py | 7 +++++++ tests/core/test_set_input_helper_order.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/policyengine_core/holders/helpers.py b/policyengine_core/holders/helpers.py index d8f7f8bb2..3bd0a7015 100644 --- a/policyengine_core/holders/helpers.py +++ b/policyengine_core/holders/helpers.py @@ -110,6 +110,7 @@ def _drop_calculated_overlapping( its values. """ branch_names = set(_branches_read_with(holder, branch_name)) + dropped_periods = set() # Memory keys are "{branch}:{period}"; branch names cannot contain ":". # Disk keys are "{branch}_{period}"; branch names can contain "_", but # period strings cannot, so the period follows the last "_". @@ -121,6 +122,7 @@ def _drop_calculated_overlapping( ): del memory._arrays[key] memory._shared.discard(key) + dropped_periods.add(period_string) disk = holder._disk_storage if disk is not None: for key in list(disk._files): @@ -129,6 +131,11 @@ def _drop_calculated_overlapping( holder, stored_branch_name, period_string, branch_names, period ): del disk._files[key] + dropped_periods.add(period_string) + fast_cache = getattr(getattr(holder, "simulation", None), "_fast_cache", None) + if fast_cache and dropped_periods: + for dropped_period in dropped_periods: + fast_cache.pop((holder.variable.name, periods.period(dropped_period)), None) def _is_calculated_overlapping( diff --git a/tests/core/test_set_input_helper_order.py b/tests/core/test_set_input_helper_order.py index 636cf5768..2f58c6700 100644 --- a/tests/core/test_set_input_helper_order.py +++ b/tests/core/test_set_input_helper_order.py @@ -337,6 +337,20 @@ def test_calculate_returns_the_input_after_a_repeated_read(): assert_reads(simulation, "count_m", {"2013-01": [7, 9]}) +def test_a_dropped_sum_leaves_calculates_fast_cache(): + simulation = build_simulation() + assert_reads(simulation, "flow_m", {"year:2013:2": [0, 0]}) + # ``calculate`` does not put a sum in its fast cache today; a value held + # there for a period whose stored value is dropped goes with it. + two_years = periods.period("year:2013:2") + simulation._fast_cache[("flow_m", two_years)] = np.zeros(2, dtype=np.float32) + + simulation.set_input("flow_m", "2014", YEARLY_INPUT) + + assert ("flow_m", two_years) not in simulation._fast_cache + assert_reads(simulation, "flow_m", {"year:2013:2": YEARLY_INPUT}) + + def test_a_helper_called_directly_stores_inputs(): simulation = build_simulation() assert_reads(simulation, "flow_m", {"2013-01": [0, 0]}) From 1f5e3ecb375ba006dff998a793ace481d0d9e409 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 22:50:26 -0400 Subject: [PATCH 3/9] Restore dumped values as inputs; test nested branches and rolling years A dump does not record which values were inputs, and restore_simulation stored every value with put_in_cache, so a restored simulation had an empty input record and the helpers replaced restored inputs (review finding 1). Restore now stores each value as a recorded input, the rule #576 applies to dumps without an inputs.txt. Tests: a restored month keeps its value under a yearly input (divide and dispatch); an input on a nested branch drops the sum its parent branch calculated; an input stored for twelve months from March is kept. The order-independence property now creates up to two nested branches, with calculations at each level. Co-Authored-By: Claude Opus 5.5 --- changelog.d/set-input-helper-order.fixed.md | 2 +- policyengine_core/tools/simulation_dumper.py | 23 ++++++- tests/core/test_set_input_helper_order.py | 64 +++++++++++++++++++ .../test_set_input_helper_order_property.py | 44 +++++++------ tests/fixtures/set_input_helper_order.py | 2 + 5 files changed, 113 insertions(+), 22 deletions(-) diff --git a/changelog.d/set-input-helper-order.fixed.md b/changelog.d/set-input-helper-order.fixed.md index 37e80fa43..50ca56498 100644 --- a/changelog.d/set-input-helper-order.fixed.md +++ b/changelog.d/set-input-helper-order.fixed.md @@ -1 +1 @@ -`set_input_divide_by_period` and `set_input_dispatch_by_period` now leave only inputs in place when they spread a longer period's input over a variable's months or years: a value the simulation calculated for a sub-period (a default, a formula result) is replaced, and a sum or twelfth of the variable it calculated over an overlapping period is dropped, so the same input gives the same values whatever was calculated before it was set (previously a calculated month kept its value and changed the others' share, or made the input fail as inconsistent). +`set_input_divide_by_period` and `set_input_dispatch_by_period` now leave only inputs in place when they spread a longer period's input over a variable's months or years: a value the simulation calculated for a sub-period (a default, a formula result) is replaced, and the variable's values it calculated over an overlapping period (a sum, a twelfth) are dropped, so the same input gives the same values whatever was calculated before it was set (previously a calculated month kept its value and changed the others' share, or made the input fail as inconsistent). `restore_simulation` restores every dumped value as an input, since a dump does not say which values were inputs, so a later input over a longer period keeps them as it did before. diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index c3db0c4fa..b48dc047d 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -38,6 +38,12 @@ def dump_simulation(simulation, directory): def restore_simulation(directory, tax_benefit_system, **kwargs): """ Restore simulation from directory + + A dump does not say which of its values were inputs, so every value is + restored as an input (recorded in ``_user_input_keys``, as ``set_input`` + records one). An input set later over a longer period then keeps the + restored values of its sub-periods, as it would have kept them in the + dumped simulation if they were inputs there. """ simulation = Simulation( tax_benefit_system, tax_benefit_system.instantiate_entities() @@ -137,4 +143,19 @@ def _restore_holder(simulation, variable, directory): for period in disk_storage.get_known_periods(): value = disk_storage.get(period) - holder.put_in_cache(value, period) + _restore_input(simulation, holder, period, value) + + +def _restore_input(simulation, holder, period, value): + """Store ``value`` as an input, recorded as ``set_input`` records one. + + ``Holder.set_input`` would also run the variable's ``set_input`` helper, + but a dump holds values already split into the variable's own periods. + """ + if not hasattr(simulation, "_user_input_contexts"): + simulation._user_input_contexts = [] + simulation._user_input_contexts.append(simulation.branch_name) + try: + holder._set(period, value, simulation.branch_name) + finally: + simulation._user_input_contexts.pop() diff --git a/tests/core/test_set_input_helper_order.py b/tests/core/test_set_input_helper_order.py index 2f58c6700..688633b06 100644 --- a/tests/core/test_set_input_helper_order.py +++ b/tests/core/test_set_input_helper_order.py @@ -13,6 +13,8 @@ from __future__ import annotations +import tempfile + import numpy as np import pytest @@ -22,6 +24,10 @@ set_input_divide_by_period, ) from policyengine_core.simulations import SimulationBuilder +from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, +) from tests.fixtures.set_input_helper_order import ( BRANCH_NAME, MONTHS, @@ -304,6 +310,18 @@ def test_an_input_stored_at_a_longer_period_is_kept(): ) +def test_an_input_stored_for_twelve_months_from_another_month_is_kept(): + simulation = build_simulation() + # Storage keys twelve months from March as the year starting in March. + simulation.set_input("flow_m", "month:2013-03:12", np.array([50.0, 60.0])) + + simulation.set_input("flow_m", "year:2013:2", YEARLY_INPUT * 2) + + holder = simulation.get_holder("flow_m") + np.testing.assert_array_equal(holder.get_array("year:2013-03"), [50, 60]) + assert_reads(simulation, "flow_m", {"2013-05": MONTHLY_SHARE}) + + # The record of inputs, and calculate's fast cache @@ -437,6 +455,52 @@ def test_input_stored_under_default_replaces_what_the_branch_calculated(): assert_reads(branch, "flow_m", {"2013-01": MONTHLY_SHARE, "2013": YEARLY_INPUT}) +def test_input_on_a_nested_branch_drops_the_sum_its_parent_branch_calculated(): + simulation = build_simulation() + parent = simulation.get_branch(BRANCH_NAME) + # The parent branch stores the sum under its own name, which the nested + # branch reads before ``default``. + assert_reads(parent, "flow_m", {"2013": [0, 0]}) + nested = parent.get_branch("nested") + + nested.set_input("flow_m", "2013", YEARLY_INPUT) + + assert_reads(nested, "flow_m", {"2013-01": MONTHLY_SHARE, "2013": YEARLY_INPUT}) + assert_reads(parent, "flow_m", {"2013": [0, 0]}) + + +# Restored simulations + + +@pytest.mark.parametrize( + "name, month, month_value, year_value, later, later_value", + [ + ("flow_m", "2013-01", [300.0, 600.0], YEARLY_INPUT, "2013-02", None), + ("count_m", "2013-03", [3, 4], [7, 9], "2013-04", [3, 4]), + ], +) +def test_a_restored_month_keeps_its_value_under_a_yearly_input( + name, month, month_value, year_value, later, later_value +): + simulation = build_simulation() + simulation.set_input(name, month, np.array(month_value)) + with tempfile.TemporaryDirectory(prefix="core-set-input-helpers-") as directory: + dump_simulation(simulation, directory) + restored = restore_simulation(directory, simulation.tax_benefit_system) + + # A dump does not say which values were inputs: all of them are restored + # as inputs, as an input set before the dump would have been kept. + restored.set_input(name, "2013", np.array(year_value)) + + if later_value is None: + rest = np.array(year_value, dtype=np.float32) - np.array( + month_value, dtype=np.float32 + ) + later_value = (rest / 11).astype(np.float32) + assert_reads(restored, name, {month: month_value, later: later_value}) + + +# Values stored on disk # Values stored on disk diff --git a/tests/core/test_set_input_helper_order_property.py b/tests/core/test_set_input_helper_order_property.py index 7294fa34e..c0d068ae8 100644 --- a/tests/core/test_set_input_helper_order_property.py +++ b/tests/core/test_set_input_helper_order_property.py @@ -1,8 +1,9 @@ """What a simulation calculated before its inputs were set does not change them. -Each example sets some inputs, has one simulation calculate a random list of -requests, then sets more inputs on that simulation and on one that -calculated nothing, and compares the two. It also compares both with a model +Each example sets some inputs, has one simulation calculate random lists of +requests (on it, then on up to two branches created one inside the other), +then sets more inputs on the last of these and on the same branch of a +simulation that calculated nothing, and compares the two. It also compares both with a model that holds inputs only (``tests/fixtures/set_input_helper_order.py``). ``test_set_input_helper_order.py`` pins the same behaviour with examples. @@ -33,7 +34,7 @@ st = hypothesis.strategies from tests.fixtures.set_input_helper_order import ( - BRANCH_NAME, + BRANCH_NAMES, COUNT, DIVIDED, FLOWS, @@ -98,23 +99,27 @@ def _input(periods_of): ) -def _run(first_inputs, requests_before_branch, requests, later_inputs, on_disk, branch): +def _run(first_inputs, requests_by_level, later_inputs, on_disk, depth): """Run the example; return the simulation the later inputs were set on, - each later input's outcome, and the model's.""" + each later input's outcome, and the model's. + + ``depth`` branches are created one inside the other, and each list of + ``requests_by_level`` is calculated before the next branch is created + (the last one on the simulation the later inputs are set on). + """ simulation = build_simulation(SYSTEM, on_disk=on_disk) reference = Reference() for name, period, values in first_inputs: assert apply_input(simulation, name, period, values) == reference.set_input( name, period, values ) - for name, period in requests_before_branch: - read(simulation, name, period) branch_name = "default" - if branch: - simulation = simulation.get_branch(BRANCH_NAME) - branch_name = BRANCH_NAME - for name, period in requests: - read(simulation, name, period) + for level in range(depth + 1): + if level: + branch_name = BRANCH_NAMES[level - 1] + simulation = simulation.get_branch(branch_name) + for name, period in requests_by_level[level] if requests_by_level else []: + read(simulation, name, period) outcomes = [ apply_input(simulation, name, period, values) for name, period, values in later_inputs @@ -129,23 +134,22 @@ def _run(first_inputs, requests_before_branch, requests, later_inputs, on_disk, @hypothesis.settings(max_examples=500, **SETTINGS) @hypothesis.given( first_inputs=st.lists(any_input, max_size=4), - requests_before_branch=st.lists(request, max_size=6), - requests=st.lists(request, max_size=8), + requests_by_level=st.lists(st.lists(request, max_size=6), min_size=3, max_size=3), later_inputs=st.lists(any_input, min_size=1, max_size=6), on_disk=st.booleans(), - branch=st.booleans(), + depth=st.integers(0, 2), ) def test_inputs_do_not_depend_on_what_was_calculated_before( - first_inputs, requests_before_branch, requests, later_inputs, on_disk, branch + first_inputs, requests_by_level, later_inputs, on_disk, depth ): # Branches of a simulation that stores on disk share its files # (policyengine-core#558), so the two are not combined. - on_disk = on_disk and not branch + on_disk = on_disk and not depth calculated, outcomes, expected, reference, branch_name = _run( - first_inputs, requests_before_branch, requests, later_inputs, on_disk, branch + first_inputs, requests_by_level, later_inputs, on_disk, depth ) fresh, fresh_outcomes, _, _, _ = _run( - first_inputs, [], [], later_inputs, on_disk, branch + first_inputs, [], later_inputs, on_disk, depth ) # The same inputs are accepted or refused, and recorded. diff --git a/tests/fixtures/set_input_helper_order.py b/tests/fixtures/set_input_helper_order.py index d61dbf366..1cecf0bdd 100644 --- a/tests/fixtures/set_input_helper_order.py +++ b/tests/fixtures/set_input_helper_order.py @@ -38,6 +38,8 @@ COUNT = 2 BRANCH_NAME = "what_if" +#: A branch, and a branch of that branch. +BRANCH_NAMES = [BRANCH_NAME, "nested"] class Status(Enum): From f7ac58b056bd667045f3fcb666e0dd34a08c9870 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 22:58:17 -0400 Subject: [PATCH 4/9] Keep the property test off disk file names Windows rejects On disk, a value calculated for a period like year:2013:2 is stored in a file named after the period, and Windows rejects ":" in file names (policyengine-core#526). The Windows CI jobs failed on such a read in the on-disk examples. In disk mode the property now skips those periods. Co-Authored-By: Claude Opus 5.5 --- .../test_set_input_helper_order_property.py | 8 +++++--- tests/fixtures/set_input_helper_order.py | 17 ++++++++++++++--- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/tests/core/test_set_input_helper_order_property.py b/tests/core/test_set_input_helper_order_property.py index c0d068ae8..b5af64df1 100644 --- a/tests/core/test_set_input_helper_order_property.py +++ b/tests/core/test_set_input_helper_order_property.py @@ -52,6 +52,7 @@ read, read_all, same_arrays, + storable_on_disk, sub_periods, ) @@ -119,7 +120,8 @@ def _run(first_inputs, requests_by_level, later_inputs, on_disk, depth): branch_name = BRANCH_NAMES[level - 1] simulation = simulation.get_branch(branch_name) for name, period in requests_by_level[level] if requests_by_level else []: - read(simulation, name, period) + if not on_disk or storable_on_disk(period): + read(simulation, name, period) outcomes = [ apply_input(simulation, name, period, values) for name, period, values in later_inputs @@ -162,8 +164,8 @@ def test_inputs_do_not_depend_on_what_was_calculated_before( name for name, period, _ in later_inputs if period in OWN_PERIODS[name] } other_periods_of = [name for name in NAMES if name not in given_own_period_input] - values = read_all(calculated, other_periods_of=other_periods_of) - fresh_values = read_all(fresh, other_periods_of=other_periods_of) + values = read_all(calculated, other_periods_of=other_periods_of, on_disk=on_disk) + fresh_values = read_all(fresh, other_periods_of=other_periods_of, on_disk=on_disk) assert values.keys() == fresh_values.keys() for key in values: assert same_arrays(values[key], fresh_values[key]), ( diff --git a/tests/fixtures/set_input_helper_order.py b/tests/fixtures/set_input_helper_order.py index 1cecf0bdd..8d114cdac 100644 --- a/tests/fixtures/set_input_helper_order.py +++ b/tests/fixtures/set_input_helper_order.py @@ -308,9 +308,19 @@ def read(simulation, name: str, period: str) -> np.ndarray: return np.asarray(simulation.calculate(name, period)) -def read_all(simulation, names=NAMES, other_periods_of=NAMES) -> dict: +def storable_on_disk(period: str) -> bool: + """Whether disk storage can hold a value calculated for ``period`` on any + platform: it names the file after the period, and Windows rejects ":" in + file names (policyengine-core#526).""" + return ":" not in period + + +def read_all( + simulation, names=NAMES, other_periods_of=NAMES, on_disk: bool = False +) -> dict: """Every read of ``names`` in a fixed order: each variable's own periods, - then, for those in ``other_periods_of``, the periods of another size.""" + then, for those in ``other_periods_of``, the periods of another size + (with ``on_disk``, only those disk storage can hold).""" result = {} for name in names: for period in OWN_READS[name]: @@ -318,7 +328,8 @@ def read_all(simulation, names=NAMES, other_periods_of=NAMES) -> dict: for name in names: if name in other_periods_of: for period in OTHER_READS[name]: - result[(name, period)] = read(simulation, name, period) + if not on_disk or storable_on_disk(period): + result[(name, period)] = read(simulation, name, period) return result From 29410c3d17986103256c52ad974a7964bb0ad077 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 07:27:51 -0400 Subject: [PATCH 5/9] Record which dumped values were inputs; forget an input a calculation replaces Review round 2: 1. Restoring every dumped value as an input froze calculated values through a later reform. The dumper now writes, next to each variable's arrays, the periods whose value was an input (inputs.txt), and restore records exactly those; a dump without the file restores every value as an input. This is policyengine-core#576's dumper change, taken byte for byte so the two PRs merge without conflict. 2. A value stored by put_in_cache over an input (say a calculate_add sum over an input set for twelve months) left the input's record entry in place, so the helpers and apply_reform kept treating the calculated value as an input. Holder._set now drops the entry, in both of its forms for twelve months from the first of a month, when it stores a value outside set_input. Co-Authored-By: Claude Opus 5.5 --- changelog.d/set-input-helper-order.fixed.md | 2 +- policyengine_core/holders/holder.py | 25 ++++++ policyengine_core/tools/simulation_dumper.py | 87 +++++++++++++++++--- tests/core/test_set_input_helper_order.py | 43 ++++++++++ tests/fixtures/set_input_helper_order.py | 17 ++++ 5 files changed, 163 insertions(+), 11 deletions(-) diff --git a/changelog.d/set-input-helper-order.fixed.md b/changelog.d/set-input-helper-order.fixed.md index 50ca56498..9552c7ff7 100644 --- a/changelog.d/set-input-helper-order.fixed.md +++ b/changelog.d/set-input-helper-order.fixed.md @@ -1 +1 @@ -`set_input_divide_by_period` and `set_input_dispatch_by_period` now leave only inputs in place when they spread a longer period's input over a variable's months or years: a value the simulation calculated for a sub-period (a default, a formula result) is replaced, and the variable's values it calculated over an overlapping period (a sum, a twelfth) are dropped, so the same input gives the same values whatever was calculated before it was set (previously a calculated month kept its value and changed the others' share, or made the input fail as inconsistent). `restore_simulation` restores every dumped value as an input, since a dump does not say which values were inputs, so a later input over a longer period keeps them as it did before. +`set_input_divide_by_period` and `set_input_dispatch_by_period` now leave only inputs in place when they spread a longer period's input over a variable's months or years: a value the simulation calculated for a sub-period (a default, a formula result) is replaced, and the variable's values it calculated over an overlapping period (a sum, a twelfth) are dropped, so the same input gives the same values whatever was calculated before it was set (previously a calculated month kept its value and changed the others' share, or made the input fail as inconsistent). A value calculated over an input now also removes that input from the simulation's record of inputs, so `apply_reform` and the helpers stop treating it as one. diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 11a8ae50b..35c4a8285 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -374,6 +374,31 @@ def _set( if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() simulation._user_input_keys.add((self.variable.name, branch_name, period)) + else: + self._forget_replaced_input(period, branch_name) + + def _forget_replaced_input(self, period: Period, branch_name: str) -> None: + """Drop the record of an input that a calculated value has replaced. + + ``_user_input_keys`` says which stored values are inputs. When a value + stored outside ``set_input`` (by ``put_in_cache``) replaces one, the + stored value is no longer that input, so its entry goes: + ``apply_reform`` would otherwise keep the calculated value, and the + ``set_input`` helpers would treat it as already set. Storage keys + twelve months from the first of a month as that year, so both forms + of the entry are dropped. + """ + input_keys = getattr(self.simulation, "_user_input_keys", None) + if not input_keys or period is None: + return + name = self.variable.name + period = periods.period(period) + input_keys.discard((name, branch_name, period)) + unit, start, size = period + if unit == periods.YEAR and size == 1: + input_keys.discard((name, branch_name, Period((periods.MONTH, start, 12)))) + elif unit == periods.MONTH and size == 12: + input_keys.discard((name, branch_name, Period((periods.YEAR, start, 1)))) def put_in_cache( self, value: ArrayLike, period: Period, branch_name: str = "default" diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index b48dc047d..1320bd073 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -2,13 +2,22 @@ import os +import warnings import numpy as np from policyengine_core.data_storage import OnDiskStorage +from policyengine_core import periods from policyengine_core.periods import ETERNITY from policyengine_core.simulations import Simulation +# Next to each variable's arrays: the periods, one per line, whose dumped +# value was an input (stored through ``set_input``). ``restore_simulation`` +# registers exactly these as inputs, so ``apply_reform``, which keeps inputs +# and drops calculated values, keeps the same values in the restored +# simulation as in the dumped one. +INPUT_PERIODS_FILE = "inputs.txt" + def dump_simulation(simulation, directory): """ @@ -26,24 +35,27 @@ def dump_simulation(simulation, directory): entities_dump_dir = os.path.join(directory, "__entities__") os.mkdir(entities_dump_dir) + input_keys = _input_storage_keys(simulation) for entity in simulation.populations.values(): # Dump entity structure _dump_entity(entity, entities_dump_dir) # Dump variable values for holder in entity._holders.values(): - _dump_holder(holder, directory) + _dump_holder(holder, directory, input_keys) def restore_simulation(directory, tax_benefit_system, **kwargs): """ Restore simulation from directory - A dump does not say which of its values were inputs, so every value is - restored as an input (recorded in ``_user_input_keys``, as ``set_input`` - records one). An input set later over a longer period then keeps the - restored values of its sub-periods, as it would have kept them in the - dumped simulation if they were inputs there. + Values the dumped simulation stored as inputs are restored as inputs + (recorded in ``_user_input_keys``, as ``set_input`` records them), and + every other value as a calculated one, so ``apply_reform`` keeps and drops + the same values it would have in the dumped simulation. A dump written + before inputs were recorded (no ``inputs.txt``) does not say which values + were inputs, so every value in it is restored as an input, with a + warning: ``apply_reform`` then keeps its calculated values as dumped. """ simulation = Simulation( tax_benefit_system, tax_benefit_system.instantiate_entities() @@ -64,17 +76,58 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): variables_to_restore = ( variable for variable in os.listdir(directory) if variable != "__entities__" ) - for variable in variables_to_restore: - _restore_holder(simulation, variable, directory) + without_input_record = [ + variable + for variable in variables_to_restore + if not _restore_holder(simulation, variable, directory) + ] + if without_input_record: + warnings.warn( + f"The simulation dump in {directory} does not record which values " + f"were inputs ({len(without_input_record)} variables have no " + f"{INPUT_PERIODS_FILE}; it was written by an earlier version of " + "policyengine-core). Every value in it is restored as an input, " + "so apply_reform keeps the calculated values as dumped instead of " + "recalculating them. Dump the simulation again to record its " + "inputs.", + stacklevel=2, + ) return simulation -def _dump_holder(holder, directory): +def _dump_holder(holder, directory, input_keys=frozenset()): disk_storage = holder.create_disk_storage(directory, preserve=True) + input_periods = [] for period in holder.get_known_periods(): value = holder.get_array(period) disk_storage.put(value, period) + # The input record of exactly the value dumped: ``get_array`` above + # reads the default branch. + if (holder.variable.name, "default", str(period)) in input_keys: + input_periods.append(str(period)) + path = os.path.join(disk_storage.storage_dir, INPUT_PERIODS_FILE) + with open(path, "w") as file: + file.write("".join(f"{period}\n" for period in dict.fromkeys(input_periods))) + + +def _input_storage_keys(simulation): + """The storage keys ``_user_input_keys`` records as inputs. + + Each record entry becomes ``(variable, branch, period)`` with the period + as storage writes it, as ``Simulation._invalidate_all_caches`` reads the + record back through the storage: an ETERNITY variable's one value is an + input whatever period its entry names. + """ + input_keys = set() + for name, branch_name, period in getattr(simulation, "_user_input_keys", ()): + variable = simulation.tax_benefit_system.get_variable(name) + if variable is not None and variable.definition_period == ETERNITY: + period = ETERNITY + elif period is None: + continue + input_keys.add((name, branch_name, str(periods.period(period)))) + return input_keys def _dump_entity(population, directory): @@ -129,6 +182,7 @@ def _restore_entity(population, directory): def _restore_holder(simulation, variable, directory): + """Restore one variable's values; return whether its inputs were recorded.""" storage_dir = os.path.join(directory, variable) is_variable_eternal = ( simulation.tax_benefit_system.get_variable(variable).definition_period @@ -141,9 +195,22 @@ def _restore_holder(simulation, variable, directory): holder = simulation.get_holder(variable) + input_periods_path = os.path.join(storage_dir, INPUT_PERIODS_FILE) + if os.path.exists(input_periods_path): + with open(input_periods_path) as file: + input_periods = set(file.read().split()) + else: + # Dumped before inputs were recorded: nothing says which values were + # calculated, so keep every value as an input. + input_periods = None + for period in disk_storage.get_known_periods(): value = disk_storage.get(period) - _restore_input(simulation, holder, period, value) + if input_periods is None or str(period) in input_periods: + _restore_input(simulation, holder, period, value) + else: + holder.put_in_cache(value, period) + return input_periods is not None def _restore_input(simulation, holder, period, value): diff --git a/tests/core/test_set_input_helper_order.py b/tests/core/test_set_input_helper_order.py index 688633b06..802cb61f5 100644 --- a/tests/core/test_set_input_helper_order.py +++ b/tests/core/test_set_input_helper_order.py @@ -30,6 +30,7 @@ ) from tests.fixtures.set_input_helper_order import ( BRANCH_NAME, + FormulaReturns999, MONTHS, Status, build_simulation, @@ -500,6 +501,48 @@ def test_a_restored_month_keeps_its_value_under_a_yearly_input( assert_reads(restored, name, {month: month_value, later: later_value}) +def test_a_restored_calculated_value_is_recalculated_after_a_reform(): + simulation = build_simulation() + assert_reads(simulation, "formula_flow_m", {"2013-01": [11, 11]}) + with tempfile.TemporaryDirectory(prefix="core-set-input-helpers-") as directory: + dump_simulation(simulation, directory) + restored = restore_simulation(directory, simulation.tax_benefit_system) + + # The dump records which values were inputs, so a calculated one is + # restored as calculated and a reform recalculates it. + restored.apply_reform(FormulaReturns999) + + assert_reads(restored, "formula_flow_m", {"2013-01": [999, 999]}) + + +# A calculated value replacing an input + + +def test_a_sum_cached_over_an_input_stops_counting_as_an_input(): + simulation = build_simulation() + simulation.set_input("flow_m", "month:2013-01:12", np.array([50.0, 60.0])) + # ``calculate_add`` caches the sum of the months over the input stored + # for the year. + simulation.calculate_add("flow_m", "2013") + + simulation.set_input("flow_m", "2013", YEARLY_INPUT) + + assert_reads(simulation, "flow_m", {"2013-01": MONTHLY_SHARE, "2013": YEARLY_INPUT}) + simulation._invalidate_all_caches() + assert_reads(simulation, "flow_m", {"2013-01": MONTHLY_SHARE, "2013": YEARLY_INPUT}) + + +def test_a_calculated_value_replacing_an_input_is_not_kept_by_a_reform(): + simulation = build_simulation() + simulation.set_input("formula_flow_m", "2013-01", np.array([5.0, 5.0])) + holder = simulation.get_holder("formula_flow_m") + holder.put_in_cache(np.array([7.0, 7.0]), periods.period("2013-01")) + + simulation._invalidate_all_caches() + + assert_reads(simulation, "formula_flow_m", {"2013-01": [11, 11]}) + + # Values stored on disk # Values stored on disk diff --git a/tests/fixtures/set_input_helper_order.py b/tests/fixtures/set_input_helper_order.py index 8d114cdac..79aace4a2 100644 --- a/tests/fixtures/set_input_helper_order.py +++ b/tests/fixtures/set_input_helper_order.py @@ -29,6 +29,7 @@ from policyengine_core.country_template import CountryTaxBenefitSystem, entities from policyengine_core.enums import Enum from policyengine_core.experimental import MemoryConfig +from policyengine_core.reforms import Reform from policyengine_core.holders import ( set_input_dispatch_by_period, set_input_divide_by_period, @@ -96,6 +97,22 @@ def formula(person, period, parameters): return person.filled_array(10.0 + period.start.month) +class FormulaReturns999(Reform): + """Replaces ``formula_flow_m``'s formula with one that returns 999.""" + + def apply(self): + class formula_flow_m(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "Monthly flow with a formula that returns 999" + + def formula(person, period, parameters): + return person.filled_array(999.0) + + self.update_variable(formula_flow_m) + + class flow_y(Variable): value_type = float entity = entities.Person From ac4576bbd23908c7d424dcc46b3187e0e1bd9799 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 07:53:07 -0400 Subject: [PATCH 6/9] Test a calculated value replacing an input with put_in_cache, not calculate_add With policyengine-core#571, calculate_add no longer caches a sum over an input a plain read finds, so the input is kept and the test's premise (master's calculate_add overwriting it) did not hold. Storing the calculated value with put_in_cache replaces the input with or without #571. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_set_input_helper_order.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/core/test_set_input_helper_order.py b/tests/core/test_set_input_helper_order.py index 802cb61f5..e53255594 100644 --- a/tests/core/test_set_input_helper_order.py +++ b/tests/core/test_set_input_helper_order.py @@ -521,9 +521,11 @@ def test_a_restored_calculated_value_is_recalculated_after_a_reform(): def test_a_sum_cached_over_an_input_stops_counting_as_an_input(): simulation = build_simulation() simulation.set_input("flow_m", "month:2013-01:12", np.array([50.0, 60.0])) - # ``calculate_add`` caches the sum of the months over the input stored - # for the year. - simulation.calculate_add("flow_m", "2013") + # A calculated value stored over the input for the year, as master's + # ``calculate_add`` stores the sum of the months there. + simulation.get_holder("flow_m").put_in_cache( + np.zeros(2, dtype=np.float32), periods.period("2013") + ) simulation.set_input("flow_m", "2013", YEARLY_INPUT) From 917f0c23e71b8a818b64b5fcafca7fe867bd2333 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 10:12:28 -0400 Subject: [PATCH 7/9] Drop the record cleanup on calculated writes; evict the fast cache only for read branches Review round 3: - Holder._set no longer drops a record entry when a calculated value is stored over an input. A calculate_add at a variable's own period stores the unchanged input and lost its record; a clone, which shares its source's record until #561, deleted the source's entries; and ETERNITY variables keep entries under several periods. The case the cleanup was for, calculate_add caching a sum over an input, is what #571 stops. - Master now gives a storage that shares nothing a frozenset for _shared (#578). The drop step deleted from it directly; it now leaves that to InMemoryStorage._stop_sharing_dropped_keys. - The helpers evicted calculate's fast cache for a sub-period stored under any branch name, including one the simulation does not read. They now evict only for branches it reads, as #576's holder-write property requires. Co-Authored-By: Claude Opus 5.5 --- changelog.d/set-input-helper-order.fixed.md | 2 +- policyengine_core/holders/helpers.py | 34 ++++++++++++---- policyengine_core/holders/holder.py | 25 ------------ tests/core/test_set_input_helper_order.py | 45 +++++++-------------- 4 files changed, 42 insertions(+), 64 deletions(-) diff --git a/changelog.d/set-input-helper-order.fixed.md b/changelog.d/set-input-helper-order.fixed.md index 9552c7ff7..d9cbfe756 100644 --- a/changelog.d/set-input-helper-order.fixed.md +++ b/changelog.d/set-input-helper-order.fixed.md @@ -1 +1 @@ -`set_input_divide_by_period` and `set_input_dispatch_by_period` now leave only inputs in place when they spread a longer period's input over a variable's months or years: a value the simulation calculated for a sub-period (a default, a formula result) is replaced, and the variable's values it calculated over an overlapping period (a sum, a twelfth) are dropped, so the same input gives the same values whatever was calculated before it was set (previously a calculated month kept its value and changed the others' share, or made the input fail as inconsistent). A value calculated over an input now also removes that input from the simulation's record of inputs, so `apply_reform` and the helpers stop treating it as one. +`set_input_divide_by_period` and `set_input_dispatch_by_period` now leave only inputs in place when they spread a longer period's input over a variable's months or years: a value the simulation calculated for a sub-period (a default, a formula result) is replaced, and the variable's values it calculated over an overlapping period (a sum, a twelfth) are dropped, so the same input gives the same values whatever was calculated before it was set (previously a calculated month kept its value and changed the others' share, or made the input fail as inconsistent). `dump_simulation` records which values were inputs and `restore_simulation` restores exactly those as inputs, so a restored simulation's inputs keep counting as inputs. diff --git a/policyengine_core/holders/helpers.py b/policyengine_core/holders/helpers.py index 3bd0a7015..914bd57c7 100644 --- a/policyengine_core/holders/helpers.py +++ b/policyengine_core/holders/helpers.py @@ -72,9 +72,28 @@ def _store_input( input_keys = getattr(simulation, "_user_input_keys", None) if input_keys is not None: input_keys.add((holder.variable.name, branch_name, period)) + _evict_fast_cache(holder, [period], branch_name) + + +def _evict_fast_cache(holder: Holder, changed_periods, branch_name: str) -> None: + """Drop ``calculate``'s fast-cache entries for the variable's + ``changed_periods``, if the simulation reads ``branch_name``. + + The fast cache holds what this simulation's ``calculate`` returned; a + value stored under a branch it does not read changes none of that. + """ + simulation = getattr(holder, "simulation", None) fast_cache = getattr(simulation, "_fast_cache", None) - if fast_cache: - fast_cache.pop((holder.variable.name, period), None) + if not fast_cache: + return + get_visible_branch_names = getattr(simulation, "_get_visible_branch_names", None) + if ( + get_visible_branch_names is not None + and branch_name not in get_visible_branch_names() + ): + return + for changed_period in changed_periods: + fast_cache.pop((holder.variable.name, changed_period), None) def _branches_read_with(holder: Holder, branch_name: str) -> List[str]: @@ -121,8 +140,8 @@ def _drop_calculated_overlapping( holder, stored_branch_name, period_string, branch_names, period ): del memory._arrays[key] - memory._shared.discard(key) dropped_periods.add(period_string) + memory._stop_sharing_dropped_keys() disk = holder._disk_storage if disk is not None: for key in list(disk._files): @@ -132,10 +151,11 @@ def _drop_calculated_overlapping( ): del disk._files[key] dropped_periods.add(period_string) - fast_cache = getattr(getattr(holder, "simulation", None), "_fast_cache", None) - if fast_cache and dropped_periods: - for dropped_period in dropped_periods: - fast_cache.pop((holder.variable.name, periods.period(dropped_period)), None) + _evict_fast_cache( + holder, + [periods.period(dropped_period) for dropped_period in dropped_periods], + branch_name, + ) def _is_calculated_overlapping( diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 35c4a8285..11a8ae50b 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -374,31 +374,6 @@ def _set( if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() simulation._user_input_keys.add((self.variable.name, branch_name, period)) - else: - self._forget_replaced_input(period, branch_name) - - def _forget_replaced_input(self, period: Period, branch_name: str) -> None: - """Drop the record of an input that a calculated value has replaced. - - ``_user_input_keys`` says which stored values are inputs. When a value - stored outside ``set_input`` (by ``put_in_cache``) replaces one, the - stored value is no longer that input, so its entry goes: - ``apply_reform`` would otherwise keep the calculated value, and the - ``set_input`` helpers would treat it as already set. Storage keys - twelve months from the first of a month as that year, so both forms - of the entry are dropped. - """ - input_keys = getattr(self.simulation, "_user_input_keys", None) - if not input_keys or period is None: - return - name = self.variable.name - period = periods.period(period) - input_keys.discard((name, branch_name, period)) - unit, start, size = period - if unit == periods.YEAR and size == 1: - input_keys.discard((name, branch_name, Period((periods.MONTH, start, 12)))) - elif unit == periods.MONTH and size == 12: - input_keys.discard((name, branch_name, Period((periods.YEAR, start, 1)))) def put_in_cache( self, value: ArrayLike, period: Period, branch_name: str = "default" diff --git a/tests/core/test_set_input_helper_order.py b/tests/core/test_set_input_helper_order.py index e53255594..1884de3b2 100644 --- a/tests/core/test_set_input_helper_order.py +++ b/tests/core/test_set_input_helper_order.py @@ -370,6 +370,20 @@ def test_a_dropped_sum_leaves_calculates_fast_cache(): assert_reads(simulation, "flow_m", {"year:2013:2": YEARLY_INPUT}) +def test_an_input_under_a_branch_the_simulation_does_not_read_keeps_its_fast_cache(): + simulation = build_simulation() + assert_reads(simulation, "flow_m", {"2013-01": [0, 0]}) + cached = simulation._fast_cache[("flow_m", periods.period("2013-01"))] + + simulation.get_holder("flow_m").set_input( + periods.period("2013"), YEARLY_INPUT, "elsewhere" + ) + + # The simulation reads only ``default``, so what it returned is unchanged. + assert simulation._fast_cache[("flow_m", periods.period("2013-01"))] is cached + assert_reads(simulation, "flow_m", {"2013-01": [0, 0]}) + + def test_a_helper_called_directly_stores_inputs(): simulation = build_simulation() assert_reads(simulation, "flow_m", {"2013-01": [0, 0]}) @@ -515,37 +529,6 @@ def test_a_restored_calculated_value_is_recalculated_after_a_reform(): assert_reads(restored, "formula_flow_m", {"2013-01": [999, 999]}) -# A calculated value replacing an input - - -def test_a_sum_cached_over_an_input_stops_counting_as_an_input(): - simulation = build_simulation() - simulation.set_input("flow_m", "month:2013-01:12", np.array([50.0, 60.0])) - # A calculated value stored over the input for the year, as master's - # ``calculate_add`` stores the sum of the months there. - simulation.get_holder("flow_m").put_in_cache( - np.zeros(2, dtype=np.float32), periods.period("2013") - ) - - simulation.set_input("flow_m", "2013", YEARLY_INPUT) - - assert_reads(simulation, "flow_m", {"2013-01": MONTHLY_SHARE, "2013": YEARLY_INPUT}) - simulation._invalidate_all_caches() - assert_reads(simulation, "flow_m", {"2013-01": MONTHLY_SHARE, "2013": YEARLY_INPUT}) - - -def test_a_calculated_value_replacing_an_input_is_not_kept_by_a_reform(): - simulation = build_simulation() - simulation.set_input("formula_flow_m", "2013-01", np.array([5.0, 5.0])) - holder = simulation.get_holder("formula_flow_m") - holder.put_in_cache(np.array([7.0, 7.0]), periods.period("2013-01")) - - simulation._invalidate_all_caches() - - assert_reads(simulation, "formula_flow_m", {"2013-01": [11, 11]}) - - -# Values stored on disk # Values stored on disk From 1fa8b034b20aa4a410edfb5914a47eaca2c73660 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 13:30:13 -0400 Subject: [PATCH 8/9] Drop sums and twelfths an own-period input makes stale set_input for one of a variable's own periods (a month of a monthly variable, a year of a yearly one) replaced the stored value but left what calculate had derived from it at periods of another size: the sum of a monthly flow over a year, a yearly flow's twelfth at a month. A later calculate returned the old result, so it depended on whether that period was read before the input was set (#579). Warm annual flow 0, January set to 1: the annual read stayed 0. Holder.set_input now applies the rule the set_input helpers follow (_drop_calculated_overlapping): after storing an input for one of the variable's own periods, it drops the values the simulation calculated for the variable at overlapping periods, on the input's branch and the branches it reads, and evicts calculate's fast-cache entries for overlapping periods. Inputs are kept. ETERNITY variables are left alone. Stacked on #581, whose helper it reuses. Removes #581's property exclusion for #579, so test_inputs_do_not_depend_on_what_was_calculated_before now compares reads of other-size periods for every variable. On #581's head that property fails, shrunk to an annual flow_m read on a branch of depth 2; with this change it passes. Adds seven regression examples, all of which fail on #581's head. Co-Authored-By: Claude Opus 5.5 --- .../fix-own-period-input-stale-sums.fixed.md | 1 + policyengine_core/holders/helpers.py | 40 +++++++ policyengine_core/holders/holder.py | 9 +- ...test_own_period_input_overlapping_cache.py | 104 ++++++++++++++++++ .../test_set_input_helper_order_property.py | 20 ++-- 5 files changed, 160 insertions(+), 14 deletions(-) create mode 100644 changelog.d/fix-own-period-input-stale-sums.fixed.md create mode 100644 tests/core/test_own_period_input_overlapping_cache.py diff --git a/changelog.d/fix-own-period-input-stale-sums.fixed.md b/changelog.d/fix-own-period-input-stale-sums.fixed.md new file mode 100644 index 000000000..97ee376eb --- /dev/null +++ b/changelog.d/fix-own-period-input-stale-sums.fixed.md @@ -0,0 +1 @@ +`set_input` for one of a variable's own periods (one month, or one year) now drops the sums and twelfths the simulation calculated for that variable over overlapping periods, so a later `calculate` adds up or divides the new input instead of returning the old result (#579). diff --git a/policyengine_core/holders/helpers.py b/policyengine_core/holders/helpers.py index 914bd57c7..382d84b4e 100644 --- a/policyengine_core/holders/helpers.py +++ b/policyengine_core/holders/helpers.py @@ -158,6 +158,46 @@ def _drop_calculated_overlapping( ) +def _drop_calculated_over_own_period( + holder: Holder, period: Period, branch_name: str +) -> None: + """Drop what the simulation calculated for the variable over periods + that overlap ``period``, once an input is set for ``period``, one of the + variable's own periods (policyengine-core#579). + + ``calculate`` adds up a monthly flow over a year, or divides a yearly + flow into a month, and caches the result at that period. The value the + input replaced went into it, so it is out of date, and ``calculate`` + would return it instead of adding up or dividing the input. The rule is + the one the ``set_input`` helpers follow (``_drop_calculated_overlapping``), + applied to an input that needs no helper. Inputs are kept. + + ``calculate``'s fast cache also drops what it holds for the variable at + periods that overlap ``period``, ``period`` included, so an input set + through the holder rather than ``Simulation.set_input`` is read back too. + """ + _drop_calculated_overlapping(holder, period, branch_name) + simulation = getattr(holder, "simulation", None) + fast_cache = getattr(simulation, "_fast_cache", None) + if not fast_cache: + return + name = holder.variable.name + _evict_fast_cache( + holder, + [ + cached_period + for cached_name, cached_period in list(fast_cache) + if cached_name == name + and isinstance(cached_period, Period) + and cached_period.unit != periods.ETERNITY + and not ( + cached_period.start > period.stop or cached_period.stop < period.start + ) + ], + branch_name, + ) + + def _is_calculated_overlapping( holder: Holder, stored_branch_name: str, diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 11a8ae50b..166be6541 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -280,7 +280,14 @@ def set_input( and period.unit != self.variable.definition_period ): return self.variable.set_input(self, period, array) - return self._set(period, array, branch_name, validate_nan=True) + self._set(period, array, branch_name, validate_nan=True) + if self.variable.definition_period != periods.ETERNITY: + # Imported here because the helpers module imports this one. + from policyengine_core.holders.helpers import ( + _drop_calculated_over_own_period, + ) + + _drop_calculated_over_own_period(self, period, branch_name) finally: if simulation is not None: simulation._user_input_contexts.pop() diff --git a/tests/core/test_own_period_input_overlapping_cache.py b/tests/core/test_own_period_input_overlapping_cache.py new file mode 100644 index 000000000..0fea5632f --- /dev/null +++ b/tests/core/test_own_period_input_overlapping_cache.py @@ -0,0 +1,104 @@ +"""An input for one of a variable's own periods updates its sums and twelfths. + +``calculate`` adds up a monthly flow over a year and divides a yearly flow +into a month, and caches the result at that period. An input set for one +month (or one year) of the variable replaced the value stored there but left +those results in place, so a later ``calculate`` returned the old sum or +twelfth (policyengine-core#579). The result depended on whether the longer +(or shorter) period had been read before the input was set. + +``set_input`` now drops what the simulation calculated for the variable over +periods that overlap the input's, by the rule the ``set_input`` helpers +follow (``test_set_input_helper_order.py``). Inputs are kept. +``test_set_input_helper_order_property.py`` compares such reads with a +simulation that calculated nothing first. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from tests.fixtures.set_input_helper_order import ( + BRANCH_NAME, + build_simulation, + read, +) + + +def _sequence(simulation, warm): + """Issue #579's table: read the other period or not, set, read again.""" + if warm: + read(simulation, "flow_m", "2013") + read(simulation, "flow_y", "2013-05") + simulation.set_input("flow_m", "2013-05", np.array([50.0, 0.0])) + simulation.set_input("flow_y", "2013", np.array([1200.0, 0.0])) + return read(simulation, "flow_m", "2013"), read(simulation, "flow_y", "2013-05") + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_sum_and_twelfth_follow_an_input_set_after_them(on_disk): + cold = _sequence(build_simulation(on_disk=on_disk), warm=False) + warm = _sequence(build_simulation(on_disk=on_disk), warm=True) + + for result in (cold, warm): + np.testing.assert_array_equal(result[0], [50.0, 0.0]) + np.testing.assert_array_equal(result[1], [100.0, 0.0]) + + +def test_annual_sum_follows_a_january_input(): + # The audit witness: a warm annual flow of 0, then January set to 1. + simulation = build_simulation() + np.testing.assert_array_equal(read(simulation, "flow_m", "2013"), [0.0, 0.0]) + + simulation.set_input("flow_m", "2013-01", np.array([1.0, 1.0])) + + np.testing.assert_array_equal(read(simulation, "flow_m", "2013"), [1.0, 1.0]) + + +def test_inputs_for_other_months_are_kept(): + simulation = build_simulation() + simulation.set_input("flow_m", "2013-04", np.array([7.0, 0.0])) + np.testing.assert_array_equal(read(simulation, "flow_m", "2013"), [7.0, 0.0]) + + simulation.set_input("flow_m", "2013-05", np.array([50.0, 0.0])) + + np.testing.assert_array_equal(read(simulation, "flow_m", "2013-04"), [7.0, 0.0]) + np.testing.assert_array_equal(read(simulation, "flow_m", "2013"), [57.0, 0.0]) + assert ("flow_m", "default", "2013-04") in { + (name, branch, str(period)) + for name, branch, period in simulation._user_input_keys + } + + +def test_sum_over_two_years_follows_an_input(): + simulation = build_simulation() + np.testing.assert_array_equal(read(simulation, "flow_m", "year:2013:2"), [0.0, 0.0]) + + simulation.set_input("flow_m", "2014-12", np.array([3.0, 0.0])) + + np.testing.assert_array_equal(read(simulation, "flow_m", "year:2013:2"), [3.0, 0.0]) + + +def test_branch_input_updates_the_sum_the_branch_reads(): + # The branch reads the sum its parent calculated until it is dropped; + # the parent keeps its own sum, which its inputs still give. + simulation = build_simulation() + np.testing.assert_array_equal(read(simulation, "flow_m", "2013"), [0.0, 0.0]) + branch = simulation.get_branch(BRANCH_NAME) + + branch.set_input("flow_m", "2013-05", np.array([50.0, 0.0])) + + np.testing.assert_array_equal(read(branch, "flow_m", "2013"), [50.0, 0.0]) + np.testing.assert_array_equal(read(simulation, "flow_m", "2013"), [0.0, 0.0]) + + +def test_input_through_the_holder_updates_the_sum(): + # ``Simulation.set_input`` drops the month from ``calculate``'s fast + # cache; an input set on the holder has to as well. + simulation = build_simulation() + np.testing.assert_array_equal(read(simulation, "flow_m", "2013"), [0.0, 0.0]) + + simulation.get_holder("flow_m").set_input("2013-05", np.array([50.0, 0.0])) + + np.testing.assert_array_equal(read(simulation, "flow_m", "2013"), [50.0, 0.0]) diff --git a/tests/core/test_set_input_helper_order_property.py b/tests/core/test_set_input_helper_order_property.py index b5af64df1..ecf58da5e 100644 --- a/tests/core/test_set_input_helper_order_property.py +++ b/tests/core/test_set_input_helper_order_property.py @@ -17,11 +17,11 @@ - ``calculate_add`` and ``calculate_divide`` called directly, whose results are cached where a plain read does not return them (#571): the requests are plain ``calculate`` calls. -- An input set for one of the variable's own periods (one month, or one - year) after a sum or a twelfth of the variable was calculated: the helpers - are not involved and that sum stays cached (policyengine-core#579). The - periods of another size are compared only for variables given no such - input after the calculations. + +An input set for one of the variable's own periods (one month, or one year) +after a sum or a twelfth of the variable was calculated involves no helper, +but follows the same rule (policyengine-core#579), so the periods of another +size are compared for every variable. """ from __future__ import annotations @@ -158,14 +158,8 @@ def test_inputs_do_not_depend_on_what_was_calculated_before( assert outcomes == fresh_outcomes == expected assert input_record(calculated) == input_record(fresh) - # A sum or a twelfth calculated before an input for one of the - # variable's own periods stays cached (policyengine-core#579). - given_own_period_input = { - name for name, period, _ in later_inputs if period in OWN_PERIODS[name] - } - other_periods_of = [name for name in NAMES if name not in given_own_period_input] - values = read_all(calculated, other_periods_of=other_periods_of, on_disk=on_disk) - fresh_values = read_all(fresh, other_periods_of=other_periods_of, on_disk=on_disk) + values = read_all(calculated, on_disk=on_disk) + fresh_values = read_all(fresh, on_disk=on_disk) assert values.keys() == fresh_values.keys() for key in values: assert same_arrays(values[key], fresh_values[key]), ( From adc18425ee153799911678a379fc31c3fd65c9ce Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 19 Oct 2026 04:06:25 -0400 Subject: [PATCH 9/9] Preserve master period bounds and provenance coverage after merge Keep ordinary overlap comparisons unchanged and use master's Gregorian end ordering only when a period stop overflows. Cover warmed stored and fast-cache-only defaults beyond datetime's last date. Recreate shadow disk caches after own-period cleanup so the provenance regressions and property retain their original witnesses. Clarify dependency invalidation versus same-variable overlap cleanup in the merged branch docs. --- docs/usage/branches.md | 25 ++++++++++---------- policyengine_core/holders/helpers.py | 18 ++++++++++---- tests/core/test_carry_over_order.py | 21 ++++++++++++++++ tests/core/test_user_input_keys_soundness.py | 12 ++++++++++ 4 files changed, 60 insertions(+), 16 deletions(-) diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 1465c01e7..8a103f6fc 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -149,18 +149,19 @@ therefore settle without exhausting the retry budget. An input set again with the value the branch already reads for that very period, as an input (not through a `set_input` helper that spreads it over -other periods, and not `-0.0` for `0.0`), changes nothing, so it drops nothing -and is no input change. An input stored for the very period being calculated -after the calculation began (by its own formula, say; under the branch's name -or any it reads, such as `default` through `Holder.set_input`) is the result, +other periods, and not `-0.0` for `0.0`), does not trigger branch dependency +invalidation or count as an input change. An input stored for the very period +being calculated after the calculation began (by its own formula, say; under the +branch's name or any it reads, such as `default` through `Holder.set_input`) is the result, as it would be had it been set first. Input helpers that divide or dispatch a larger-period input replace calculated sub-periods while retaining supplied inputs in their storage tiers. Their shared -overlap cleanup also removes calculated sums or twelfths for the variable, -including inherited values in the branch's storage, without changing the parent. +overlap cleanup also runs for an input at a non-eternal variable's own period and +removes calculated sums or twelfths for the variable, including inherited values +in the branch's storage, without changing the parent. Branch dependency invalidation additionally drops later calculated values that -may depend on the replaced input. On a simulation without a parent, helper +may depend on the replaced input. On a simulation without a parent, overlap cleanup still runs while branch dependency invalidation does not. A custom `set_input` handler that calculates values between its own stores @@ -169,9 +170,9 @@ anything (through `calculate`, `calculate_add` or `_calculate`), the branch drops again, by the same rule, once the handler returns or raises (the inputs it stored before raising stay). -If there is no such record, the branch drops nothing. That is the case when a -formula creates the branch while it is still calculating the variable the -branch overrides, unless another branch the formula created for the same +If there is no such record, branch dependency invalidation drops nothing. That is +the case when a formula creates the branch while it is still calculating the +variable the branch overrides, unless another branch the formula created for the same comparison already handed back a value calculated from that variable: then the branch drops what came back and what was calculated after it. @@ -200,8 +201,8 @@ or carried-over value also counts. A result whose calculation was running when an input changed is not kept unless it was calculated again from the new inputs. An input set again with the value the branch already reads for that period, as an input, changes no value read anywhere, so -dropping nothing then is sound; one set where the value was calculated still -drops, since it changes which periods hold inputs. +skipping dependency invalidation then is sound; one set where the value was +calculated still drops, since it changes which periods hold inputs. The rule can drop more than it needs to (a value stored later that does not depend on the input is calculated again) but not less, within these limits: diff --git a/policyengine_core/holders/helpers.py b/policyengine_core/holders/helpers.py index adfe0258f..1eba51fff 100644 --- a/policyengine_core/holders/helpers.py +++ b/policyengine_core/holders/helpers.py @@ -199,14 +199,24 @@ def _drop_calculated_over_own_period( if cached_name == name and isinstance(cached_period, Period) and cached_period.unit != periods.ETERNITY - and not ( - cached_period.start > period.stop or cached_period.stop < period.start - ) + and _periods_overlap(cached_period, period) ], branch_name, ) +def _periods_overlap(first: Period, second: Period) -> bool: + """Compare extents even when a period ends beyond datetime's last date.""" + try: + return not (first.start > second.stop or first.stop < second.start) + except OverflowError: + from policyengine_core.simulations.simulation import _end_order + + return first.start.date.toordinal() < _end_order( + second + ) and second.start.date.toordinal() < _end_order(first) + + def _is_calculated_overlapping( holder: Holder, stored_branch_name: str, @@ -224,7 +234,7 @@ def _is_calculated_overlapping( return False if stored_period.unit == periods.ETERNITY: return False - if stored_period.start > period.stop or stored_period.stop < period.start: + if not _periods_overlap(stored_period, period): return False return not _is_input(holder, stored_period, stored_branch_name, storage_name) diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py index 0fc88939f..a144fd78f 100644 --- a/tests/core/test_carry_over_order.py +++ b/tests/core/test_carry_over_order.py @@ -145,6 +145,27 @@ def test_an_input_ending_after_the_last_date_carries(system, stored, asked): ) +@pytest.mark.parametrize("cache_only", [False, True]) +def test_an_input_ending_after_the_last_date_drops_overlapping_defaults( + system, cache_only +): + built = simulation(system, {}) + name = "day_input_without_helper" + holder = built.get_holder(name) + requested = periods.period("9999-12-31") + holder._do_not_store = cache_only + np.testing.assert_array_equal(built.calculate(name, requested), [0, 0]) + assert (name, requested) in built._fast_cache + assert holder._memory_storage.has(requested) == (not cache_only) + holder._do_not_store = False + + built.set_input(name, "day:9999-12-30:3", [10, 20]) + + assert holder.get_array(requested) is None + assert (name, requested) not in built._fast_cache + np.testing.assert_array_equal(built.calculate(name, requested), [10, 20]) + + @pytest.mark.parametrize("longer_first", [True, False], ids=["longer", "shorter"]) def test_an_input_ending_after_the_last_date_ends_last(system, longer_first): values = {"day:2012-01-01:4000000": [10, 20], "day:2012-01-01:2": [1, 2]} diff --git a/tests/core/test_user_input_keys_soundness.py b/tests/core/test_user_input_keys_soundness.py index 9fad83d1e..898001b75 100644 --- a/tests/core/test_user_input_keys_soundness.py +++ b/tests/core/test_user_input_keys_soundness.py @@ -157,6 +157,10 @@ def test_memory_input_wins_over_a_conflicting_disk_cache(invalidate): holder.put_in_cache(np.array([99.0]), JANUARY) simulation.memory_config.max_memory_occupation_pc = 101 simulation.set_input("rent", JANUARY, [42.0]) + assert holder._disk_storage.get(JANUARY) is None + # Own-period cleanup removed the old cache; model a later conflicting + # disk cache so export and replay still exercise both tiers. + holder._disk_storage.put(np.array([99.0]), JANUARY) np.testing.assert_array_equal(holder._memory_storage.get(JANUARY), [42.0]) np.testing.assert_array_equal(holder._disk_storage.get(JANUARY), [99.0]) @@ -234,6 +238,10 @@ def test_deleted_memory_input_does_not_turn_a_disk_calculation_into_an_input(): assert _key("rent", JANUARY) not in simulation._user_input_keys simulation.memory_config.max_memory_occupation_pc = 101 simulation.set_input("rent", JANUARY, [500.0]) + assert holder._disk_storage.get(JANUARY) is None + # Keep a calculated disk value behind the input to test its provenance + # after the memory input is deleted. + holder._disk_storage.put(np.array([0.0]), JANUARY, derived=True) np.testing.assert_array_equal(holder._memory_storage.get(JANUARY), [500.0]) np.testing.assert_array_equal(holder._disk_storage.get(JANUARY), [0.0]) @@ -373,6 +381,10 @@ def test_tier_provenance_replays_the_same_inputs_as_a_fresh_simulation( holder.put_in_cache(np.array([float(disk_value)]), JANUARY) simulation.memory_config.max_memory_occupation_pc = 101 simulation.set_input("rent", JANUARY, [memory_value]) + if not disk_is_input: + assert holder._disk_storage.get(JANUARY) is None + # Preserve the generated conflicting-tier witness after cleanup. + holder._disk_storage.put(np.array([float(disk_value)]), JANUARY) if delete_memory: simulation.delete_arrays("rent", "2025")