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..ae9d1f560 --- /dev/null +++ b/changelog.d/set-input-helper-order.fixed.md @@ -0,0 +1 @@ +Period-splitting input helpers now preserve supplied inputs in memory or on disk, replace calculated sub-period values, and invalidate overlapping caches consistently with the supplied-input tier record. diff --git a/policyengine_core/holders/helpers.py b/policyengine_core/holders/helpers.py index 664cd3850..33414a0c7 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,29 +23,178 @@ 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, storage_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). + + Use the same supplied-input tier record as reform replay. A cache write + can leave a registry entry behind or hide an input in another tier, and + its carry-over mark (``derived=False``) does not make it a supplied input. + A holder with no simulation, or a simulation with no record, keeps the + legacy behavior: 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 + period = holder._storage_period(period) + if (holder.variable.name, branch_name, period) not in input_keys: + return False + holder._seed_input_storage(period, branch_name) + return holder._stores_user_input(period, branch_name, storage_name) + + def get_stored_input(holder: Holder, period: Period, branch_name: str) -> ArrayLike: - """The input stored for ``period`` under ``branch_name`` itself, or - ``None``. - - A value the simulation calculated and cached there (see - ``Holder.put_in_cache``) is not an input, so the helpers below treat the - sub-period as empty and the new input replaces it, as an input set for - that sub-period alone would. Treating it as an input kept the calculated - value under the input, copied it into the later sub-periods (dispatch), - or took it out of the total (divide), so the result depended on what had - been calculated before the input was set. + """The supplied input under this exact branch, preferring memory to disk. + + A surviving disk input is still known when a non-input memory cache + hides it; overlap cleanup removes that cache after the helper succeeds. """ - if holder._memory_storage.has(period, branch_name): - storage = holder._memory_storage - elif holder._disk_storage is not None and holder._disk_storage.has( - period, branch_name + # Checked before reading, so that a calculated value a branch still + # shares with its parent is not copied only to be replaced. + for storage_name, storage in ( + ("memory", holder._memory_storage), + ("disk", holder._disk_storage), ): - storage = holder._disk_storage - else: - return None - if storage.is_derived(period, branch_name): - return None - return storage.get(period, branch_name) + if storage is not None and _is_input(holder, period, branch_name, storage_name): + value = storage.get(period, branch_name) + if value is not None: + return value + return None + + +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. + """ + simulation = getattr(holder, "simulation", None) + input_keys = getattr(simulation, "_user_input_keys", None) + # Let the holder normalize and record the supplied-input tier, including + # direct helper calls. Keep a legacy holder without an input record in + # that mode for the whole call, rather than creating a partial record. + holder._set(period, array, branch_name, is_input=input_keys is not None) + + +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 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]: + """``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)) + 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 "_". + 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, "memory" + ): + del memory._arrays[key] + dropped_periods.add(_period_from_storage_key(period_string)) + memory._stop_sharing_dropped_keys() + memory._unmark_dropped_keys() + 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, "disk" + ): + del disk._files[key] + dropped_periods.add(_period_from_storage_key(period_string)) + disk._derived.intersection_update(disk._files) + _evict_fast_cache(holder, dropped_periods, branch_name) + + +def _period_from_storage_key(period_string: str) -> Period: + """Parse a stored period whose early year may have lost leading zeros.""" + components = period_string.split(":") + date_index = 1 if len(components) > 1 else 0 + date = components[date_index].split("-") + date[0] = date[0].zfill(4) + components[date_index] = "-".join(date) + return periods.period(":".join(components)) + + +def _is_calculated_overlapping( + holder: Holder, + stored_branch_name: str, + period_string: str, + branch_names: set, + period: Period, + storage_name: str, +) -> bool: + if stored_branch_name not in branch_names: + return False + try: + stored_period = _period_from_storage_key(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, storage_name) def set_input_dispatch_by_period(holder: Holder, period: Period, array: ArrayLike): @@ -53,6 +203,20 @@ def set_input_dispatch_by_period(holder: Holder, period: Period, array: ArrayLik 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) @@ -71,18 +235,19 @@ 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 months that already have an input + # 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) while sub_period.start < after_instant: - existing_array = get_stored_input(holder, sub_period, branch_name) - if existing_array is None: - holder._set(sub_period, array, branch_name) + existing_input = get_stored_input(holder, sub_period, branch_name) + if existing_input is None: + _store_input(holder, sub_period, array, branch_name) 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) + _drop_calculated_overlapping(holder, period, branch_name) def set_input_divide_by_period(holder: Holder, period: Period, array: ArrayLike): @@ -91,6 +256,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): @@ -110,30 +282,28 @@ 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_input(holder, sub_period, branch_name) - if existing_array is not None: - remaining_array -= existing_array + existing_input = get_stored_input(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_input(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) 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( holder.variable.name, period, array, array - remaining_array ) ) + _drop_calculated_overlapping(holder, period, branch_name) diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index bde1e6df3..1f393fba5 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -259,13 +259,18 @@ def _record_input_storage( locations = locations | {storage} if is_input else locations - {storage} recorded[slot] = (period, locations) - def _stores_user_input(self, period: Period, branch_name: str) -> bool: + def _stores_user_input( + self, period: Period, branch_name: str, storage_name: Optional[str] = None + ) -> bool: """Whether an actual input survives in either tier, without reading it. Called only for a slot already in the simulation's replay record. Register-only inputs (including restored inputs) may have no location metadata yet; their existing storage marks identify the surviving tier. Marks never add an unrecorded slot to the replay record. + + With ``storage_name``, check only that tier so input helpers can + discard a cache in one tier while keeping an input in the other. """ recorded = getattr(self, "_user_input_storage", {}).get( (branch_name, str(period)) @@ -276,7 +281,8 @@ def _stores_user_input(self, period: Period, branch_name: str) -> bool: ("disk", self._disk_storage), ): if ( - location in locations + (storage_name is None or location == storage_name) + and location in locations and storage is not None and storage.has(period, branch_name) and not storage.is_derived(period, branch_name) 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..b4fb8fc3b --- /dev/null +++ b/tests/core/test_set_input_helper_order.py @@ -0,0 +1,771 @@ +"""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 tempfile + +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 policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, +) +from tests.fixtures.set_input_helper_order import ( + BRANCH_NAME, + FormulaReturns999, + 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}" + ) + + +def assert_matches_fresh_inputs_before_and_after_replay(simulation, fresh, name): + """Compare helper outputs and exported inputs before and after a cache wipe.""" + for replay in (False, True): + if replay: + simulation._invalidate_all_caches() + assert simulation._user_input_keys == fresh._user_input_keys + assert_reads( + simulation, + name, + {period: read(fresh, name, period) for period in [*MONTHS_2013, "2013"]}, + ) + assert simulation.to_input_dataframe().equals(fresh.to_input_dataframe()) + + +# 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} + ) + + +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 + + +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_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_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]}) + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_a_helper_called_directly_stores_inputs(on_disk): + simulation = build_simulation(on_disk=on_disk) + 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 + location = "disk" if on_disk else "memory" + for name in ("flow_m", "count_m"): + holder = simulation.get_holder(name) + for month in MONTHS_2013: + month_period = periods.period(month) + assert holder._user_input_storage[("default", month)] == ( + month_period, + frozenset({location}), + ) + assert holder._stores_user_input(month_period, "default") + 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 + ) + if on_disk: + simulation.memory_config.max_memory_occupation_pc = 101 + for name in ("flow_m", "count_m"): + simulation.get_holder(name).put_in_cache( + np.array([700, 900]), january, derived=False + ) + 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([1, 2]) + ) + assert_reads(simulation, "count_m", {"2013-01": [7, 9]}) + simulation._invalidate_all_caches() + assert_reads(simulation, "flow_m", {"2013-01": MONTHLY_SHARE, "2013": YEARLY_INPUT}) + assert_reads(simulation, "count_m", {"2013-01": [7, 9], "2013": [7, 9]}) + assert simulation.to_input_dataframe()["flow_m__2013-01"].tolist() == [100, 200] + assert simulation.to_input_dataframe()["count_m__2013-01"].tolist() == [7, 9] + + +@pytest.mark.parametrize( + "name, month_value, year_value", + [ + ("flow_m", [300.0, 600.0], YEARLY_INPUT), + ("count_m", [3, 4], [7, 9]), + ], +) +@pytest.mark.parametrize("on_disk", [False, True]) +def test_helper_ignores_unregistered_carry_over_cache( + name, month_value, year_value, on_disk +): + simulation = build_simulation(on_disk=on_disk) + holder = simulation.get_holder(name) + march = periods.period("2013-03") + holder.put_in_cache(np.array(month_value), march, derived=False) + # Carry-over provenance does not make a cache write a supplied input. + assert not holder.is_derived(march) + assert (name, "default", march) not in simulation._user_input_keys + read(simulation, name, "2013") + fresh = build_simulation() + + simulation.set_input(name, "2013", np.array(year_value)) + fresh.set_input(name, "2013", np.array(year_value)) + + assert_matches_fresh_inputs_before_and_after_replay(simulation, fresh, name) + + +@pytest.mark.parametrize( + "name, month_value, year_value", + [ + ("flow_m", [300.0, 600.0], YEARLY_INPUT), + ("count_m", [3, 4], [7, 9]), + ], +) +@pytest.mark.parametrize("registered", [False, True]) +def test_helper_ignores_a_record_after_cache_replaced_the_only_input( + name, month_value, year_value, registered +): + simulation = build_simulation() + holder = simulation.get_holder(name) + march = periods.period("2013-03") + if registered: + holder._memory_storage.put( + np.array(month_value, dtype=holder.variable.dtype), march + ) + simulation._user_input_keys.add((name, "default", march)) + assert not hasattr(holder, "_user_input_storage") + else: + simulation.set_input(name, march, np.array(month_value)) + holder.put_in_cache(np.array([700, 900]), march, derived=False) + assert (name, "default", march) in simulation._user_input_keys + assert not holder._stores_user_input(march, "default") + read(simulation, name, "2013") + # The replaced input no longer survives in either tier. + fresh = build_simulation() + + simulation.set_input(name, "2013", np.array(year_value)) + fresh.set_input(name, "2013", np.array(year_value)) + + assert_matches_fresh_inputs_before_and_after_replay(simulation, fresh, name) + + +@pytest.mark.parametrize( + "name, month_value, year_value", + [ + ("flow_m", [300.0, 600.0], YEARLY_INPUT), + ("count_m", [3, 4], [7, 9]), + ], +) +@pytest.mark.parametrize("registered", [False, True]) +@pytest.mark.parametrize("derived", [False, True]) +def test_helper_keeps_disk_input_hidden_by_memory_cache( + name, month_value, year_value, registered, derived +): + simulation = build_simulation(on_disk=True) + holder = simulation.get_holder(name) + march = periods.period("2013-03") + if registered: + holder._disk_storage.put( + np.array(month_value, dtype=holder.variable.dtype), march + ) + simulation._user_input_keys.add((name, "default", march)) + assert not hasattr(holder, "_user_input_storage") + else: + simulation.set_input(name, march, np.array(month_value)) + simulation.memory_config.max_memory_occupation_pc = 101 + if derived: + # Model a derived shadow already stored above the supplied disk input. + holder._memory_storage.put( + np.array([700, 900], dtype=holder.variable.dtype), march, derived=True + ) + else: + holder.put_in_cache(np.array([700, 900]), march, derived=False) + assert holder._stores_user_input(march, "default") + np.testing.assert_array_equal(holder.get_array(march), [700, 900]) + read(simulation, name, "2013") + fresh = build_simulation() + fresh.set_input(name, march, np.array(month_value)) + + simulation.set_input(name, "2013", np.array(year_value)) + fresh.set_input(name, "2013", np.array(year_value)) + + np.testing.assert_array_equal(holder._disk_storage.get(march), month_value) + assert holder._memory_storage.get(march) is None + assert_matches_fresh_inputs_before_and_after_replay(simulation, fresh, name) + + +@pytest.mark.parametrize( + "name, year_value", + [("flow_m", YEARLY_INPUT), ("count_m", [7, 9])], +) +def test_repeated_helper_clears_cache_shadow_when_every_month_has_a_disk_input( + name, year_value +): + simulation = build_simulation(on_disk=True) + fresh = build_simulation() + for candidate in (simulation, fresh): + candidate.set_input(name, "2013", np.array(year_value)) + holder = simulation.get_holder(name) + january = periods.period("2013-01") + simulation.memory_config.max_memory_occupation_pc = 101 + holder.put_in_cache(np.array([700, 900]), january, derived=False) + read(simulation, name, "2013") + + # The repeated helper writes no new months, but the cache shadow still + # must be removed so reads agree with the supplied disk inputs. + simulation.set_input(name, "2013", np.array(year_value)) + + assert holder._memory_storage.get(january) is None + assert_matches_fresh_inputs_before_and_after_replay(simulation, fresh, name) + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_helper_drops_overlapping_early_year_cache(on_disk): + simulation = build_simulation(on_disk=on_disk) + holder = simulation.get_holder("flow_m") + year = periods.period("0025") + holder.put_in_cache(np.array([5.0, 6.0]), year, derived=True) + # Plain reads do not fast-cache an aggregate today; model an existing + # fast-cache entry to prove cleanup uses the early-year Period itself. + simulation._fast_cache[("flow_m", year)] = np.array([5.0, 6.0]) + fresh = build_simulation() + + set_input_divide_by_period(holder, year, YEARLY_INPUT) + fresh.set_input("flow_m", year, YEARLY_INPUT) + + assert holder.get_array(year) is None + assert ("flow_m", year) not in simulation._fast_cache + january = year.start.period(periods.MONTH) + for month in [january.offset(index) for index in range(12)]: + np.testing.assert_array_equal( + simulation.calculate("flow_m", month), fresh.calculate("flow_m", month) + ) + np.testing.assert_array_equal( + simulation.calculate("flow_m", year), fresh.calculate("flow_m", year) + ) + + +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}) + + +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) + + # The dump restores the month as an input, so the yearly helper keeps it. + 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}) + + +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]}) + + +# 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, derived=True + ) + + 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..17afd435b --- /dev/null +++ b/tests/core/test_set_input_helper_order_property.py @@ -0,0 +1,260 @@ +"""What a simulation calculated before its inputs were set does not change them. + +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. + +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_NAMES, + 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, + storable_on_disk, + 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_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. + + ``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 + ) + branch_name = "default" + 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 []: + 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 + ] + 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_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(), + depth=st.integers(0, 2), +) +def test_inputs_do_not_depend_on_what_was_calculated_before( + first_inputs, requests_by_level, later_inputs, on_disk, depth +): + # Keep disk and branch examples separate: the storage period-listing + # parser does not support branch names containing underscores. + on_disk = on_disk and not depth + calculated, outcomes, expected, reference, branch_name = _run( + first_inputs, requests_by_level, later_inputs, on_disk, depth + ) + fresh, fresh_outcomes, _, _, _ = _run( + first_inputs, [], later_inputs, on_disk, depth + ) + + # 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, 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]), ( + 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..e0a184647 --- /dev/null +++ b/tests/fixtures/set_input_helper_order.py @@ -0,0 +1,369 @@ +"""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.reforms import Reform +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" +#: A branch, and a branch of that branch. +BRANCH_NAMES = [BRANCH_NAME, "nested"] + + +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 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 + 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 storable_on_disk(period: str) -> bool: + """The properties' conservative disk period domain. + + Compound-period filename portability has separate coverage in + ``test_disk_period_filenames.py``; these properties retain simple periods. + """ + 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 + (with ``on_disk``, only those disk storage can hold).""" + 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]: + if not on_disk or storable_on_disk(period): + 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) + )