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