diff --git a/changelog.d/user-input-keys-follow-storage.fixed.md b/changelog.d/user-input-keys-follow-storage.fixed.md new file mode 100644 index 00000000..1efbd363 --- /dev/null +++ b/changelog.d/user-input-keys-follow-storage.fixed.md @@ -0,0 +1 @@ +Keep the record of `set_input` values (`Simulation._user_input_keys`) in step with storage, so `apply_reform` and `to_input_dataframe` no longer treat a formula result, or another simulation's input, as an input: `delete_arrays` drops the entries for the values it deletes, `clone` (so also `get_branch`) gives the copy its own record, values a custom `set_input` handler calculates are not recorded, `subsample` records only what it stores, and each entry names the period storage keys the value under (eternity for an eternal variable, the year for twelve months starting on the first of a month). diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 11a8ae50..0c8473b5 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -1,6 +1,6 @@ import os import warnings -from typing import TYPE_CHECKING, Any, List, Tuple +from typing import TYPE_CHECKING, Any, List, Optional, Tuple import numpy import psutil @@ -100,12 +100,92 @@ def delete_arrays( If ``period`` is ``None``, remove all known values of the variable. If ``period`` is not ``None``, only remove all values for any period included in period (e.g. if period is "2017", values for "2017-01", "2017-07", etc. would be removed) + + A deleted value set with ``set_input`` stops counting as an input, so + a value calculated later for the same period is not taken for one. """ + simulation = getattr(self, "simulation", None) + user_input_keys = getattr(simulation, "_user_input_keys", None) + stored_before = self._stored_keys() if user_input_keys else None self._memory_storage.delete(period, branch_name) if self._disk_storage: self._disk_storage.delete(period, branch_name) + if user_input_keys: + self._forget_deleted_inputs(user_input_keys, stored_before) + + def _stored_keys(self) -> Tuple[set, set]: + """The keys memory and disk storage hold values under.""" + return ( + set(self._memory_storage._arrays), + set(self._disk_storage._files) if self._disk_storage is not None else set(), + ) + + def _forget_deleted_inputs( + self, user_input_keys: set, stored_before: Tuple[set, set] + ) -> None: + """Drop the simulation's record of the inputs ``delete_arrays`` deleted. + + ``_user_input_keys`` records each (variable, branch, period) stored + through ``set_input``, with the period as storage keys it (see + ``_set``). ``_invalidate_all_caches`` keeps the values it names and + ``to_input_dataframe`` exports them, so an entry left behind for a + deleted value would make a formula result stored later for that + period count as an input. + + The keys the deletion removed are found by comparing the keys each + storage holds before and after, so the cost depends on what this + holder stores, not on the size of the record, and no stored file is + read. An entry for a removed key is dropped only if neither storage + still holds a value for it: disk storage deletes only the period it + is given, not the periods within it (policyengine-core#564), so a + value deleted from memory can survive on disk, and its entry stays. + """ + name = self.variable.name + memory_before, disk_before = stored_before + arrays = self._memory_storage._arrays + files = self._disk_storage._files if self._disk_storage is not None else {} + # Memory keys are "{branch}:{period}"; branch names cannot contain ":". + removed = [key.split(":", 1) for key in memory_before if key not in arrays] + # Disk keys are "{branch}_{period}"; branch names can contain "_" but + # period strings cannot, so the period follows the last "_". + removed += [key.rsplit("_", 1) for key in disk_before if key not in files] + forgotten = [] + for branch, period_string in removed: + try: + key = (name, branch, periods.period(period_string)) + except ValueError: + # Not a key ``put`` wrote (say, a file ``restore`` found), so + # no entry names it. + continue + if key in user_input_keys and not self._stores(key[2], branch): + forgotten.append(key) + user_input_keys.difference_update(forgotten) + + def _stores(self, period: Period, branch_name: str) -> bool: + """Whether either storage holds a value for exactly this branch and + period, checked from the storages' keys without reading the value.""" + period = self._storage_period(period) + if f"{branch_name}:{period}" in self._memory_storage._arrays: + return True + return ( + self._disk_storage is not None + and f"{branch_name}_{period}" in self._disk_storage._files + ) + + def _storage_period(self, period: Period) -> Period: + """The period storage keys a value for ``period`` under. + + Storage keys a value by the period's string form, read back as a + period: eternity for an eternal variable, and for twelve months + starting on the first of a month, the year starting then + (``month:2025-01:12`` is stored as ``2025``). + """ + if self._memory_storage.is_eternal: + return periods.period(periods.ETERNITY) + return periods.period(str(periods.period(period))) + def _get_array_from_storage( self, period: Period, branch_name: str = "default" ) -> ArrayLike: @@ -347,10 +427,16 @@ def _set( value: ArrayLike, branch_name: str = "default", validate_nan: bool = False, + is_input: Optional[bool] = None, ) -> None: simulation = getattr(self, "simulation", None) user_input_contexts = getattr(simulation, "_user_input_contexts", None) - if user_input_contexts and branch_name == "default": + # A value is an input when stored while ``set_input`` runs, unless the + # caller says otherwise: ``put_in_cache`` stores calculated values, + # including those a custom ``set_input`` handler calculates. + if is_input is None: + is_input = bool(user_input_contexts) + if is_input and user_input_contexts and branch_name == "default": branch_name = user_input_contexts[-1] value = self._to_array(value, validate_nan=validate_nan) if self.variable.definition_period != periods.ETERNITY: @@ -370,10 +456,16 @@ def _set( self._disk_storage.put(value, period, branch_name) else: self._memory_storage.put(value, period, branch_name) - if user_input_contexts: + if is_input and simulation is not None: if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() - simulation._user_input_keys.add((self.variable.name, branch_name, period)) + # Record the period as storage keys the value (eternity for an + # eternal variable, whatever period it was set for; the year for + # twelve months starting on the first of a month), so each entry + # names one stored value. + simulation._user_input_keys.add( + (self.variable.name, branch_name, self._storage_period(period)) + ) def put_in_cache( self, value: ArrayLike, period: Period, branch_name: str = "default" @@ -388,7 +480,7 @@ def put_in_cache( ): return - self._set(period, value, branch_name) + self._set(period, value, branch_name, is_input=False) def default_array(self) -> ArrayLike: """ diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 2652d972..1092eca3 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -208,7 +208,11 @@ def __init__( # populates so ``_invalidate_all_caches`` can tell user-provided # source data apart from formula-computed caches. Without this the # post-``apply_reform`` cache wipe would also wipe the dataset the - # simulation was loaded from. + # simulation was loaded from. The record follows this simulation's + # storage: each entry names one stored value (the period is the one + # storage keys it under), ``delete_arrays`` drops the entries for the + # values it deletes, and ``clone`` (so also ``get_branch``) gives the + # copy its own record. self._user_input_keys: set[tuple[str, str, Period]] = set() self.debug: bool = False self.trace: bool = trace @@ -1354,6 +1358,9 @@ def delete_arrays(self, variable: str, period: Period = None) -> None: The calling branch, each ancestor branch, and the default branch are purged from this simulation's private holder storage. Other branch names and the parent simulation's holder storage remain unchanged. + Deleted inputs stop counting as inputs: a value calculated later for + the same period is a formula result, which ``apply_reform`` discards + and ``to_input_dataframe`` does not export. :param variable: the variable whose cached values should be deleted :param period: the period to delete, or all periods when omitted @@ -1523,6 +1530,12 @@ def clone( new.tax_benefit_system = self.tax_benefit_system new.debug = debug new.trace = trace + # The copy stores its own values, so it keeps its own record of + # which are inputs: one set on either simulation afterwards reaches + # only that simulation's storage and record. A ``set_input`` running + # on this simulation is not running on the copy. + new._user_input_keys = set(getattr(self, "_user_input_keys", ())) + new._user_input_contexts = [] return new @@ -1975,8 +1988,11 @@ def subsample( df = subset_df - # Update the dataset and rebuild the simulation + # Update the dataset and rebuild the simulation. Rebuilding replaces + # every stored value, so the record of inputs starts again with the + # ones the rebuild sets. self.dataset = Dataset.from_dataframe(df, self.dataset.time_period) + self._user_input_keys = set() self.build_from_dataset() # Purge ``_fast_cache`` entries populated by ``to_input_dataframe`` diff --git a/tests/core/test_user_input_keys.py b/tests/core/test_user_input_keys.py new file mode 100644 index 00000000..b656a3b8 --- /dev/null +++ b/tests/core/test_user_input_keys.py @@ -0,0 +1,615 @@ +"""Regression tests: the record of ``set_input`` values follows storage. + +``Simulation._user_input_keys`` records each (variable, branch, period) a +simulation stored through ``set_input``. ``_invalidate_all_caches`` (run by +``apply_reform``) keeps the values it names and ``to_input_dataframe`` +exports them, and country packages read it to tell a value entered directly +from one calculated by a formula. + +The record used to drift from storage (policyengine-core#559): +``delete_arrays`` deleted values but kept their entries, so a value +calculated later for the same period counted as an input, and ``clone`` +(so also ``get_branch``) shared one record between simulations that store +their values separately. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template import Simulation as CountryTemplateSimulation +from policyengine_core.country_template.entities import Person +from policyengine_core.data import Dataset +from policyengine_core.data_storage import OnDiskStorage +from policyengine_core.experimental import MemoryConfig +from policyengine_core.model_api import MONTH, YEAR, Reform, Variable +from policyengine_core.simulations import SimulationBuilder + +JANUARY = "2025-01" +FEBRUARY = "2025-02" + + +def _simulation(tax_benefit_system=None): + return SimulationBuilder().build_from_entities( + tax_benefit_system or CountryTaxBenefitSystem(), + { + "persons": {"bill": {"salary": {JANUARY: 1_000}}}, + "households": {"household": {"parents": ["bill"]}}, + }, + ) + + +def _key(variable, period, branch="default"): + return (variable, branch, periods.period(period)) + + +def _store_on_disk(simulation, *variables): + """Send every value stored from now on for ``variables`` to disk.""" + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + for variable in variables: + holder = simulation.get_holder(variable) + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + + +def _held(simulation, variable, entries): + """The (branch, period) entries whose value some storage still holds. + + Disk storage deletes only the period it is given (policyengine-core#564) + and, for a branch, the files of every branch whose name starts with that + branch's name and "_" (#552), so after a delete it can hold values memory + no longer does. Tests of disk deletes compare the record with this, so + they hold however disk storage deletes. + """ + holder = simulation.get_holder(variable) + return { + (branch, period) + for branch, period in entries + if holder._get_array_from_storage(period, branch) is not None + } + + +def _input_periods(simulation, variable): + return { + (branch, str(period)) + for name, branch, period in simulation._user_input_keys + if name == variable + } + + +class income_tax(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "Income tax at twice the rate" + + def formula(person, period, parameters): + return person("salary", period) * 0.3 + + +class DoubleIncomeTaxRate(Reform): + def apply(self): + self.update_variable(income_tax) + + +def test_formula_result_after_delete_is_not_kept_by_apply_reform(): + """A value calculated after its input was deleted is recalculated.""" + simulation = _simulation() + simulation.set_input("income_tax", JANUARY, [5.0]) + simulation.delete_arrays("income_tax") + + assert simulation.calculate("income_tax", JANUARY)[0] == 150.0 + assert _key("income_tax", JANUARY) not in simulation._user_input_keys + + simulation.apply_reform(DoubleIncomeTaxRate) + + assert simulation.get_holder("income_tax").get_array(JANUARY) is None + assert simulation.calculate("income_tax", JANUARY)[0] == 300.0 + + +def test_inputs_not_deleted_are_still_kept_by_apply_reform(): + simulation = _simulation() + simulation.set_input("income_tax", FEBRUARY, [5.0]) + simulation.delete_arrays("income_tax", JANUARY) + + simulation.apply_reform(DoubleIncomeTaxRate) + + assert simulation.calculate("income_tax", FEBRUARY)[0] == 5.0 + assert simulation.calculate("salary", JANUARY)[0] == 1_000.0 + + +def test_default_value_after_delete_is_not_exported(): + """An input variable read after its value was deleted gets its default, + which ``to_input_dataframe`` must not export as an input.""" + simulation = _simulation() + simulation.delete_arrays("salary") + + assert simulation.calculate("salary", JANUARY)[0] == 0.0 + assert "salary__2025-01" not in simulation.to_input_dataframe().columns + + +def test_delete_for_a_period_drops_only_entries_within_it(): + simulation = _simulation() + for month in ("2024-12", FEBRUARY, "2025-03"): + simulation.set_input("salary", month, [2_000.0]) + + simulation.delete_arrays("salary", FEBRUARY) + assert _input_periods(simulation, "salary") == { + ("default", "2024-12"), + ("default", JANUARY), + ("default", "2025-03"), + } + + simulation.delete_arrays("salary", "2025") + assert _input_periods(simulation, "salary") == {("default", "2024-12")} + + +def test_eternal_input_is_recorded_for_the_one_value_it_stores(): + """An eternal variable stores one value whatever period it was set for, + so its entry names that value's period, eternity.""" + simulation = _simulation() + simulation.set_input("birth", "2020", ["1980-01-01"]) + simulation.set_input("birth", "2025-03", ["1981-01-01"]) + + assert _input_periods(simulation, "birth") == {("default", "ETERNITY")} + + +def test_delete_of_an_eternal_variable_drops_its_entry_for_any_period(): + """Eternal values are stored once, whatever period they were set for, + so deleting any period deletes the value and its entry.""" + simulation = _simulation() + simulation.set_input("birth", "2020", ["1980-01-01"]) + + simulation.delete_arrays("birth", "2031-07") + + assert simulation.get_holder("birth").get_known_periods() == [] + assert _input_periods(simulation, "birth") == set() + + +def test_holder_delete_drops_the_entry(): + """Code that moves an input to another variable through the holder (as + country packages do before modelling a behavioural response) leaves no + entry for the variable it moved the value from.""" + simulation = _simulation() + holder = simulation.get_holder("salary") + simulation.set_input("pension", JANUARY, holder.get_array(JANUARY)) + holder.delete_arrays(JANUARY) + + assert _input_periods(simulation, "salary") == set() + assert _key("pension", JANUARY) in simulation._user_input_keys + assert simulation.calculate("salary", JANUARY)[0] == 0.0 + + simulation.apply_reform(DoubleIncomeTaxRate) + + assert holder.get_array(JANUARY) is None + assert "salary__2025-01" not in simulation.to_input_dataframe().columns + + +def test_branch_delete_drops_entries_for_the_branches_it_deletes_from(): + simulation = _simulation() + branch = simulation.get_branch("branch") + nested = branch.get_branch("nested") + nested.set_input("salary", FEBRUARY, [3_000.0]) + holder = nested.get_holder("salary") + holder._memory_storage.put(np.array([7.0]), JANUARY, "unrelated") + nested._user_input_keys.add(_key("salary", JANUARY, "unrelated")) + + nested.delete_arrays("salary") + + assert _input_periods(nested, "salary") == {("unrelated", JANUARY)} + assert _input_periods(branch, "salary") == {("default", JANUARY)} + assert _input_periods(simulation, "salary") == {("default", JANUARY)} + assert branch.calculate("salary", JANUARY)[0] == 1_000.0 + + +def test_entry_is_kept_for_a_value_disk_storage_did_not_delete(): + """Disk storage deletes only the period asked for, not the periods + within it (policyengine-core#564), so the entry for a monthly value it + still holds stays.""" + simulation = _simulation() + _store_on_disk(simulation, "rent") + holder = simulation.get_holder("rent") + simulation.set_input("rent", JANUARY, [500.0]) + simulation.set_input("rent", FEBRUARY, [600.0]) + assert holder._memory_storage.get(JANUARY) is None + + simulation.delete_arrays("rent", "2025") + assert _input_periods(simulation, "rent") == _held( + simulation, "rent", {("default", JANUARY), ("default", FEBRUARY)} + ) + + simulation.set_input("rent", FEBRUARY, [600.0]) + simulation.delete_arrays("rent", JANUARY) + assert _input_periods(simulation, "rent") == {("default", FEBRUARY)} + assert holder.get_array(FEBRUARY)[0] == 600.0 + + simulation.delete_arrays("rent") + assert _input_periods(simulation, "rent") == set() + + +def test_entry_is_dropped_for_an_eternal_value_deleted_from_disk(): + simulation = _simulation() + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + holder = simulation.get_holder("birth") + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + simulation.set_input("birth", "2020", ["1980-01-01"]) + assert holder._memory_storage.get("2020") is None + + simulation.delete_arrays("birth", "2031-07") + + assert holder.get_array("2020") is None + assert _input_periods(simulation, "birth") == set() + + +class _RecordThatMustNotBeScanned(set): + def __iter__(self): + raise AssertionError("delete_arrays looked through the whole record") + + +def test_delete_does_not_look_through_the_whole_record(): + """Deleting reads what the holder stores, not every entry of the record, + so its cost does not grow with the number of inputs (country packages + delete every variable on a branch for each marginal rate).""" + simulation = _simulation() + simulation.set_input("salary", FEBRUARY, [2_000.0]) + simulation.calculate("income_tax", JANUARY) + simulation._user_input_keys = _RecordThatMustNotBeScanned( + simulation._user_input_keys + | {_key(f"other_{index}", JANUARY) for index in range(1_000)} + ) + + for variable in ("income_tax", "salary", "rent", "birth"): + simulation.delete_arrays(variable, JANUARY) + simulation.delete_arrays("salary") + + assert _key("salary", JANUARY) not in simulation._user_input_keys + assert _key("salary", FEBRUARY) not in simulation._user_input_keys + assert len(simulation._user_input_keys) == 1_000 + + +def test_clone_keeps_its_own_record(): + simulation = _simulation() + simulation.calculate("income_tax", JANUARY) + clone = simulation.clone() + + assert clone._user_input_keys == simulation._user_input_keys + assert clone._user_input_keys is not simulation._user_input_keys + + clone.set_input("income_tax", JANUARY, [7.0]) + simulation.set_input("salary", FEBRUARY, [2_000.0]) + + assert _key("income_tax", JANUARY) not in simulation._user_input_keys + assert _key("salary", FEBRUARY) not in clone._user_input_keys + assert simulation.get_holder("income_tax").get_array(JANUARY)[0] == 150.0 + + simulation.apply_reform(DoubleIncomeTaxRate) + clone.apply_reform(DoubleIncomeTaxRate) + + assert simulation.calculate("income_tax", JANUARY)[0] == 300.0 + assert clone.calculate("income_tax", JANUARY)[0] == 7.0 + assert clone.calculate("salary", JANUARY)[0] == 1_000.0 + + +def test_branch_does_not_see_inputs_set_on_its_parent_after_it_was_created(): + simulation = _simulation() + branch = simulation.get_branch("branch") + simulation.set_input("salary", FEBRUARY, [2_000.0]) + + assert _key("salary", FEBRUARY) not in branch._user_input_keys + assert branch.calculate("salary", FEBRUARY)[0] == 0.0 + exported = branch.to_input_dataframe() + assert "salary__2025-02" not in exported.columns + assert exported["salary__2025-01"].tolist() == [1_000.0] + + +def test_parent_does_not_see_inputs_set_on_its_branch(): + simulation = _simulation() + branch = simulation.get_branch("branch") + branch.set_input("salary", FEBRUARY, [2_000.0]) + + assert _key("salary", FEBRUARY, "branch") in branch._user_input_keys + assert _input_periods(simulation, "salary") == {("default", JANUARY)} + + +def test_set_input_on_a_clone_is_not_running_on_the_original(): + """While ``set_input`` runs on a clone, values the original stores, by + calculating or through its holders, are not inputs of that call.""" + tax_benefit_system = CountryTaxBenefitSystem() + simulation = _simulation(tax_benefit_system) + clone = simulation.clone() + + def set_rent_and_store_on_the_original(holder, period, array): + simulation.calculate("income_tax", JANUARY) + simulation.get_holder("accommodation_size")._set(JANUARY, [80.0]) + for month in period.get_subperiods(periods.MONTH): + holder._set(month, array) + + tax_benefit_system.variables["rent"].set_input = set_rent_and_store_on_the_original + clone.set_input("rent", "2025", [500.0]) + + assert _key("income_tax", JANUARY) not in simulation._user_input_keys + assert _key("accommodation_size", JANUARY) not in simulation._user_input_keys + assert _key("rent", JANUARY) in clone._user_input_keys + assert _key("rent", JANUARY) not in simulation._user_input_keys + + +def test_values_a_set_input_handler_calculates_are_not_inputs(): + """A custom ``set_input`` handler may calculate other variables before it + stores the input; those are formula results, which ``apply_reform`` + recalculates.""" + tax_benefit_system = CountryTaxBenefitSystem() + simulation = _simulation(tax_benefit_system) + + def calculate_then_set_rent(holder, period, array): + holder.simulation.calculate("income_tax", JANUARY) + holder._set(period.start.period(periods.MONTH), array) + + tax_benefit_system.variables["rent"].set_input = calculate_then_set_rent + simulation.set_input("rent", "2025", [500.0]) + + assert _input_periods(simulation, "income_tax") == set() + assert _input_periods(simulation, "rent") == {("default", JANUARY)} + + simulation.set_input("salary", JANUARY, [2_000.0]) + simulation._invalidate_all_caches() + + assert simulation.calculate("income_tax", JANUARY)[0] == 300.0 + + +def test_input_a_handler_sets_for_a_period_string_is_recorded_as_that_period(): + tax_benefit_system = CountryTaxBenefitSystem() + simulation = _simulation(tax_benefit_system) + + def set_january_rent(holder, period, array): + holder._set(JANUARY, array) + + tax_benefit_system.variables["rent"].set_input = set_january_rent + simulation.set_input("rent", "2025", [123.0]) + + assert _key("rent", JANUARY) in simulation._user_input_keys + assert simulation.to_input_dataframe()["rent__2025-01"].tolist() == [123.0] + + +def test_subsample_records_only_the_inputs_it_stores(): + data = pd.DataFrame( + { + "person_id__2022": [1, 2], + "household_id__2022": [1, 2], + "person_household_id__2022": [1, 2], + "household_weight__2022": [1.0, 1.0], + "salary__2022-01": [1_000.0, 2_000.0], + } + ) + simulation = CountryTemplateSimulation(dataset=Dataset.from_dataframe(data, "2022")) + branch = simulation.get_branch("sample") + branch.set_input("salary", "2022-01", [3_000.0, 4_000.0]) + + branch.subsample(n=1, seed="user-input-keys", time_period="2022") + + holder = branch.get_holder("salary") + for name, branch_name, period in branch._user_input_keys: + assert ( + branch.get_holder(name)._memory_storage.get(period, branch_name) is not None + ) + assert _input_periods(branch, "salary") == {("sample", "2022-01")} + assert holder._memory_storage.get("2022-01", "sample").tolist() in ( + [3_000.0], + [4_000.0], + ) + + +def _ended(value): + class ended(Variable): + value_type = float + entity = Person + definition_period = YEAR + end = "2012-12-31" + label = "A formula that applies until 2012" + + def formula(person, period, parameters): + return person.filled_array(value) + + return ended + + +class ReformEnded(Reform): + def apply(self): + self.update_variable(_ended(9.0)) + + +def _carrying_over_simulation(on_disk): + tax_benefit_system = CountryTaxBenefitSystem() + tax_benefit_system.auto_carry_over_input_variables = True + tax_benefit_system.add_variable(_ended(7.0)) + simulation = _simulation(tax_benefit_system) + if on_disk: + _store_on_disk(simulation, "ended") + return simulation + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_formula_result_for_a_deleted_input_is_not_carried_over(on_disk): + """A value calculated for the period of a deleted input is a formula + result. ``apply_reform`` discards it, so it is not carried over to later + periods: after the reform the simulation gives what a simulation that + never had the input gives (found in the review of policyengine-core#562). + """ + simulation = _carrying_over_simulation(on_disk) + holder = simulation.get_holder("ended") + simulation.set_input("ended", "2012", [20.0]) + simulation.delete_arrays("ended", "2012") + assert simulation.calculate("ended", "2012")[0] == 7.0 + assert (holder._memory_storage.get("2012") is None) == on_disk + + simulation.apply_reform(ReformEnded) + + never_had_the_input = _carrying_over_simulation(on_disk) + never_had_the_input.apply_reform(ReformEnded) + assert simulation.calculate("ended", "2013")[0] == 0.0 + assert never_had_the_input.calculate("ended", "2013")[0] == 0.0 + assert simulation.calculate("ended", "2012")[0] == 9.0 + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_input_that_was_not_deleted_is_carried_over_after_apply_reform(on_disk): + simulation = _carrying_over_simulation(on_disk) + simulation.set_input("ended", "2012", [20.0]) + + simulation.apply_reform(ReformEnded) + + assert simulation.calculate("ended", "2013")[0] == 20.0 + assert simulation.calculate("ended", "2012")[0] == 20.0 + + +def test_entry_is_kept_while_disk_still_holds_a_value_deleted_from_memory(): + """A value can be stored both in memory and on disk. Deleting a year + deletes the months within it from memory, but disk storage deletes only + the year itself (policyengine-core#564), so the input survives on disk + and stays an input.""" + simulation = _simulation() + _store_on_disk(simulation, "rent") + holder = simulation.get_holder("rent") + simulation.set_input("rent", JANUARY, [500.0]) + simulation.memory_config.max_memory_occupation_pc = 101 + simulation.set_input("rent", JANUARY, [500.0]) + assert holder._memory_storage.get(JANUARY) is not None + assert holder._disk_storage.get(JANUARY) is not None + + simulation.delete_arrays("rent", "2025") + + assert holder._memory_storage.get(JANUARY) is None + assert _input_periods(simulation, "rent") == _held( + simulation, "rent", {("default", JANUARY)} + ) + if _input_periods(simulation, "rent"): + exported = simulation.to_input_dataframe() + assert exported["rent__2025-01"].tolist() == [500.0] + simulation._invalidate_all_caches() + assert holder.get_array(JANUARY)[0] == 500.0 + + +@pytest.mark.parametrize( + "twelve_months, stored_as, salary_month", + [ + ("month:2025-01:12", "2025", JANUARY), + ("month:2025-03:12", "year:2025-03", "2025-03"), + ], +) +def test_twelve_month_input_is_recorded_as_the_year_it_is_stored_as( + twelve_months, stored_as, salary_month +): + """Storage keys twelve months starting on the first of a month as the + year starting then, so deleting that year deletes the input's entry.""" + simulation = _simulation() + simulation.set_input("salary", salary_month, [1_000.0]) + simulation.set_input("income_tax", twelve_months, [5.0]) + assert _input_periods(simulation, "income_tax") == {("default", stored_as)} + + simulation.get_holder("income_tax").delete_arrays(twelve_months) + + assert _input_periods(simulation, "income_tax") == set() + assert simulation.calculate("income_tax", stored_as)[0] == 150.0 + simulation.apply_reform(DoubleIncomeTaxRate) + assert simulation.calculate("income_tax", stored_as)[0] == 300.0 + + +def test_disk_delete_reads_no_file_and_does_not_look_through_the_record( + monkeypatch, +): + """With disk storage configured, deleting checks which keys each storage + still holds without loading any stored file, and without going through + every entry of the record.""" + simulation = _simulation() + _store_on_disk(simulation, "rent", "income_tax", "birth", "salary") + simulation.set_input("rent", JANUARY, [500.0]) + simulation.set_input("rent", FEBRUARY, [600.0]) + simulation.set_input("rent", "2025-03", [700.0]) + simulation.set_input("birth", "2020", ["1980-01-01"]) + simulation.calculate("income_tax", JANUARY) + # March's rent is also stored in memory, so deleting the year deletes + # memory's copy and leaves the one on disk. + simulation.memory_config.max_memory_occupation_pc = 101 + simulation.set_input("rent", "2025-03", [700.0]) + simulation._user_input_keys = _RecordThatMustNotBeScanned( + simulation._user_input_keys + | {_key(f"other_{index}", JANUARY) for index in range(1_000)} + ) + + def no_file_reads(self, file): + raise AssertionError(f"delete_arrays read {file}") + + monkeypatch.setattr(OnDiskStorage, "_decode_file", no_file_reads) + + simulation.delete_arrays("rent", "2025") + simulation.delete_arrays("rent", JANUARY) + for variable in ("income_tax", "salary", "birth"): + simulation.delete_arrays(variable, JANUARY) + simulation.delete_arrays("salary") + + # Any period of an eternal variable deletes its one value. + monkeypatch.undo() + assert set(simulation._user_input_keys) == { + _key("rent", period, branch) + for branch, period in _held( + simulation, "rent", {("default", FEBRUARY), ("default", "2025-03")} + ) + } | {_key(f"other_{index}", JANUARY) for index in range(1_000)} + + +def test_delete_skips_disk_files_storage_did_not_name(): + """A file ``restore`` finds in the storage directory need not be named + after a period; deleting it removes no entry and raises nothing.""" + simulation = _simulation() + _store_on_disk(simulation, "rent") + holder = simulation.get_holder("rent") + simulation.set_input("rent", JANUARY, [500.0]) + holder._disk_storage._files["default_notes"] = "notes.npy" + + simulation.delete_arrays("rent") + + assert holder._disk_storage._files == {} + assert _input_periods(simulation, "rent") == set() + + +def test_disk_delete_on_a_branch_whose_name_contains_an_underscore(): + """Disk keys join the branch name and the period with "_", and branch + names can contain "_" (country packages use names like ``no_salt``).""" + simulation = _simulation() + _store_on_disk(simulation, "rent") + branch = simulation.get_branch("no_salt") + branch.set_input("rent", JANUARY, [500.0]) + branch.set_input("rent", FEBRUARY, [600.0]) + assert branch.get_holder("rent")._memory_storage.get(JANUARY, "no_salt") is None + + branch.delete_arrays("rent", JANUARY) + assert _input_periods(branch, "rent") == {("no_salt", FEBRUARY)} + + branch.get_holder("rent").delete_arrays(None, "no_salt") + assert _input_periods(branch, "rent") == set() + + +def test_entry_stays_while_memory_holds_a_value_disk_deleted_for_another_branch(): + """Deleting a branch's values from disk also deletes the files of every + branch whose name starts with that branch's name and "_" (fixed in + policyengine-core#552). The entry of an input memory still holds stays; + the entry of one only disk held goes.""" + simulation = _simulation() + _store_on_disk(simulation, "rent") + holder = simulation.get_holder("rent") + holder.set_input(JANUARY, [500.0], "x_y") + holder.set_input(FEBRUARY, [600.0], "x_y") + simulation.memory_config.max_memory_occupation_pc = 101 + holder.set_input(JANUARY, [500.0], "x_y") + assert holder._memory_storage.get(FEBRUARY, "x_y") is None + + holder.delete_arrays(None, "x") + + assert holder._memory_storage.get(JANUARY, "x_y")[0] == 500.0 + assert _input_periods(simulation, "rent") == _held( + simulation, "rent", {("x_y", JANUARY), ("x_y", FEBRUARY)} + ) diff --git a/tests/core/test_user_input_keys_property.py b/tests/core/test_user_input_keys_property.py new file mode 100644 index 00000000..a7db3b57 --- /dev/null +++ b/tests/core/test_user_input_keys_property.py @@ -0,0 +1,471 @@ +"""The record of ``set_input`` values matches a reference model. + +Random sequences of ``set_input``, ``calculate``, ``delete_arrays`` (through +the simulation and through a holder), ``clone``, ``get_branch`` and +``_invalidate_all_caches`` run on a family of simulations. A reference model, +seeded from the situation the simulations are built from, tracks for each +simulation which (variable, branch, period) values were stored through +``set_input`` and not deleted since. After every step, for every simulation: + +- ``_user_input_keys`` is exactly the model's set, so a step on one + simulation never changes another's record; +- each entry names a value the simulation stores, equal to the input; +- ``to_input_dataframe`` exports exactly the model's inputs, with their + values, and stores nothing. + +One input variable is set through a custom ``set_input`` handler that +calculates another variable first and stores months under string periods. +Another is set for twelve months, which storage keys as a year. + +After ``_invalidate_all_caches`` the simulation and its branches store their +inputs and nothing else. + +A second property runs ``set_input``, ``calculate``, deletes and +``_invalidate_all_caches`` on one simulation whose holders store values in +memory or on disk, switching between the two: the record names exactly the +inputs some storage still holds. ``test_user_input_keys.py`` pins the same +behaviour with examples. +""" + +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 policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.entities import Household +from policyengine_core.experimental import MemoryConfig +from policyengine_core.model_api import MONTH, Variable +from policyengine_core.simulations import SimulationBuilder + +SITUATION = { + "persons": { + "a": {"birth": {"ETERNITY": "1960-05-01"}, "salary": {"2025-01": 3_000}}, + "b": {"birth": {"ETERNITY": "1990-01-01"}}, + }, + "households": {"h": {"parents": ["a", "b"], "rent": {"2025-01": 800}}}, +} +MONTHS = ["2025-01", "2025-02", "2025-03"] +# The inputs SITUATION sets, as the simulation should store them. +SITUATION_INPUTS = { + ("salary", "default", "2025-01"): np.array([3_000.0, 0.0]), + ("birth", "default", "ETERNITY"): np.array( + ["1960-05-01", "1990-01-01"], dtype="datetime64[D]" + ), + ("rent", "default", "2025-01"): np.array([800.0]), +} +# Storage keys a value by its period's string form: twelve months starting on +# the first of a month are the year starting then. +TWELVE_MONTHS = {"month:2025-01:12": "2025", "month:2025-03:12": "year:2025-03"} + + +def _calculate_then_set_months(holder, period, array): + holder.simulation.calculate("income_tax", MONTHS[0]) + for month in MONTHS: + holder._set(month, array) + + +class quarterly_rent(Variable): + """Set for a year through a custom handler that calculates income tax + first, then stores the first quarter's months under string periods.""" + + value_type = float + entity = Household + definition_period = MONTH + label = "Rent for the first quarter" + set_input = _calculate_then_set_months + + +TAX_BENEFIT_SYSTEM = CountryTaxBenefitSystem() +TAX_BENEFIT_SYSTEM.add_variable(quarterly_rent) +INPUTS = ["salary", "income_tax", "rent", "quarterly_rent", "birth"] +EXPORTABLE = ["salary", "rent", "quarterly_rent", "birth"] +CALCULATED = [ + "salary", + "income_tax", + "social_security_contribution", + "disposable_income", + "rent", + "housing_allowance", + "birth", +] +MAX_SIMULATIONS = 6 + + +def _canonical(variable, period): + if variable == "birth": + return periods.period(periods.ETERNITY) + return periods.period(TWELVE_MONTHS.get(period, period)) + + +def _slot(key): + variable, branch, period = key + return variable, f"{branch}:{_canonical(variable, period)}" + + +def _input_value(variable, number): + if variable == "birth": + return [f"19{50 + number % 50}-01-01"] * 2 + if variable in ("rent", "quarterly_rent"): + return [float(number)] + return [float(number), float(number) / 2] + + +def _expected_array(variable, value): + if variable == "birth": + return np.array(value, dtype="datetime64[D]") + return np.array(value, dtype=float) + + +class _Model: + """What one simulation should record, and the inputs it should store.""" + + def __init__(self, branch_name, parent, inputs): + self.branch_name = branch_name + self.parent = parent + self.inputs = dict(inputs) + self.branches = {} + + def copy(self, branch_name, parent): + return _Model(branch_name, parent, self.inputs) + + def visible_branch_names(self): + names = [self.branch_name] + parent = self.parent + while parent is not None: + names.append(parent.branch_name) + parent = parent.parent + names.append("default") + return list(dict.fromkeys(names)) + + def delete(self, variable, period, branch_name): + if period is not None: + period = periods.period(period) + for key in list(self.inputs): + if ( + key[0] == variable + and key[1] == branch_name + and ( + period is None + or variable == "birth" + or period.contains(periods.period(key[2])) + ) + ): + del self.inputs[key] + + +class _Family: + def __init__(self, tax_benefit_system): + root = SimulationBuilder().build_from_entities(tax_benefit_system, SITUATION) + inputs = { + (variable, branch, _canonical(variable, period)): value + for (variable, branch, period), value in SITUATION_INPUTS.items() + } + self.simulations = [root] + self.models = [_Model("default", None, inputs)] + + def pick(self, index): + index %= len(self.simulations) + return self.simulations[index], self.models[index] + + def add(self, simulation, model): + if simulation not in self.simulations: + self.simulations.append(simulation) + self.models.append(model) + + def apply(self, operation): + kind, index, *arguments = operation + simulation, model = self.pick(index) + if kind == "set_input": + variable, period, number = arguments + value = _input_value(variable, number) + simulation.set_input(variable, period, value) + array = _expected_array(variable, value) + # The handler stores the quarter's months. An eternal variable + # stores one value whatever the period it is set for, and its + # entry names that value. + stored = MONTHS if variable == "quarterly_rent" else [period] + for stored_period in stored: + key = (variable, model.branch_name, _canonical(variable, stored_period)) + model.inputs[key] = array + elif kind == "calculate": + variable, period = arguments + simulation.calculate(variable, period) + elif kind == "delete": + variable, period = arguments + simulation.delete_arrays(variable, period) + for branch_name in model.visible_branch_names(): + model.delete(variable, period, branch_name) + elif kind == "holder_delete": + variable, period, branch_choice = arguments + names = model.visible_branch_names() + branch_name = names[branch_choice % len(names)] + simulation.get_holder(variable).delete_arrays(period, branch_name) + model.delete(variable, period, branch_name) + elif kind == "clone": + if len(self.simulations) < MAX_SIMULATIONS: + self.add( + simulation.clone(), model.copy(model.branch_name, model.parent) + ) + elif kind == "branch": + (name,) = arguments + if name == model.branch_name or name in model.branches: + return + if len(self.simulations) < MAX_SIMULATIONS: + branch_model = model.copy(name, model) + model.branches[name] = branch_model + self.add(simulation.get_branch(name), branch_model) + elif kind == "invalidate": + simulation._invalidate_all_caches() + self._assert_only_inputs_stored(simulation, model) + + def _assert_only_inputs_stored(self, simulation, model): + stored = { + (variable, storage_key) + for population in simulation.populations.values() + for variable, holder in population._holders.items() + for storage_key in holder._memory_storage._arrays + } + assert stored == {_slot(key) for key in model.inputs} + for name, branch in simulation.branches.items(): + self._assert_only_inputs_stored(branch, model.branches[name]) + + def assert_records_match(self): + for simulation, model in zip(self.simulations, self.models): + assert set(simulation._user_input_keys) == set(model.inputs) + for key, expected in model.inputs.items(): + variable, storage_key = _slot(key) + holder = simulation.get_holder(variable) + stored = holder._memory_storage._arrays.get(storage_key) + assert stored is not None, key + assert np.array_equal(stored, expected), key + self._assert_export_matches(simulation, model) + + def _assert_export_matches(self, simulation, model): + """``to_input_dataframe`` exports the inputs of the branches the + simulation reads (for each period, the first such branch's), for + periods of the variable's own unit, and stores nothing.""" + visible = model.visible_branch_names() + expected = {} + for branch in reversed(visible): + for (variable, input_branch, period), value in model.inputs.items(): + definition_period = TAX_BENEFIT_SYSTEM.get_variable( + variable + ).definition_period + if ( + input_branch == branch + and variable in EXPORTABLE + and period.unit == definition_period + ): + if variable in ("rent", "quarterly_rent"): + # One household of both people. + value = np.repeat(value, 2) + expected[f"{variable}__{period}"] = value + stored_before = _stored_keys(simulation) + fast_cache = dict(simulation._fast_cache) + + exported = simulation.to_input_dataframe() + + # Exporting reads stored inputs only. Restore the reads it cached, + # so checking does not change what later steps calculate. + simulation._fast_cache = fast_cache + assert _stored_keys(simulation) == stored_before + assert set(exported.columns) == set(expected) + for column, value in expected.items(): + assert np.array_equal( + exported[column].to_numpy().astype(value.dtype), value + ), column + + +def _stored_keys(simulation): + return { + (variable, key) + for population in simulation.populations.values() + for variable, holder in population._holders.items() + for key in holder._memory_storage._arrays + } + + +_months_or_all = st.sampled_from(MONTHS + ["2025", "year:2025-03", None]) +_index = st.integers(min_value=0, max_value=MAX_SIMULATIONS - 1) +_operation = st.one_of( + st.tuples( + st.just("set_input"), + _index, + st.sampled_from(["salary", "income_tax", "rent"]), + st.sampled_from(MONTHS), + st.integers(min_value=0, max_value=5_000), + ), + st.tuples( + st.just("set_input"), + _index, + st.just("birth"), + st.sampled_from(["ETERNITY", "2025", "2025-02"]), + st.integers(min_value=0, max_value=49), + ), + st.tuples( + st.just("set_input"), + _index, + st.just("salary"), + st.sampled_from(sorted(TWELVE_MONTHS)), + st.integers(min_value=0, max_value=5_000), + ), + st.tuples( + st.just("set_input"), + _index, + st.just("quarterly_rent"), + st.just("2025"), + st.integers(min_value=0, max_value=5_000), + ), + st.tuples( + st.just("calculate"), + _index, + st.sampled_from(CALCULATED), + st.sampled_from(MONTHS), + ), + st.tuples( + st.just("delete"), + _index, + st.sampled_from(INPUTS), + _months_or_all, + ), + st.tuples( + st.just("holder_delete"), + _index, + st.sampled_from(INPUTS), + _months_or_all, + st.integers(min_value=0, max_value=3), + ), + st.tuples(st.just("clone"), _index), + st.tuples(st.just("branch"), _index, st.sampled_from(["b1", "no_salt"])), + st.tuples(st.just("invalidate"), _index), +) + + +@hypothesis.settings( + max_examples=300, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given(operations=st.lists(_operation, max_size=30)) +def test_user_input_keys_match_reference_model(operations): + family = _Family(TAX_BENEFIT_SYSTEM) + family.assert_records_match() + for step, operation in enumerate(operations): + family.apply(operation) + try: + family.assert_records_match() + except AssertionError as error: + raise AssertionError((step, operation)) from error + + +DISK_INPUTS = ["salary", "income_tax", "rent", "birth"] + + +class _Storages: + """One simulation whose holders store in memory or on disk, and the + inputs ``set_input`` stored that some storage still holds.""" + + def __init__(self, tax_benefit_system): + self.simulation = SimulationBuilder().build_from_entities( + tax_benefit_system, SITUATION + ) + self.simulation.memory_config = MemoryConfig(max_memory_occupation=0) + for variable in tax_benefit_system.variables: + holder = self.simulation.get_holder(variable) + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + self.inputs = { + (variable, branch, _canonical(variable, period)) + for variable, branch, period in SITUATION_INPUTS + } + + def holds(self, key): + variable, branch, period = key + holder = self.simulation.get_holder(variable) + return ( + f"{branch}:{period}" in holder._memory_storage._arrays + or f"{branch}_{period}" in holder._disk_storage._files + ) + + def stored(self): + return { + (variable, branch, periods.period(period)) + for population in self.simulation.populations.values() + for variable, holder in population._holders.items() + for branch, period in [ + key.split(":", 1) for key in holder._memory_storage._arrays + ] + + [key.rsplit("_", 1) for key in holder._disk_storage._files] + } + + def apply(self, operation): + kind, *arguments = operation + simulation = self.simulation + if kind == "set_input": + variable, period, number = arguments + simulation.set_input(variable, period, _input_value(variable, number)) + self.inputs.add((variable, "default", _canonical(variable, period))) + elif kind == "calculate": + simulation.calculate(*arguments) + elif kind == "delete": + simulation.delete_arrays(*arguments) + elif kind == "holder_delete": + variable, period = arguments + simulation.get_holder(variable).delete_arrays(period) + elif kind == "on_disk": + (on_disk,) = arguments + simulation.memory_config.max_memory_occupation_pc = 0 if on_disk else 101 + elif kind == "invalidate": + simulation._invalidate_all_caches() + assert self.stored() == set(simulation._user_input_keys) + # An input stops being one when no storage holds its value any more. + self.inputs = {key for key in self.inputs if self.holds(key)} + + def assert_record_matches(self): + assert set(self.simulation._user_input_keys) == self.inputs + + +_disk_operation = st.one_of( + st.tuples( + st.just("set_input"), + st.sampled_from(["salary", "income_tax", "rent"]), + st.sampled_from(MONTHS + sorted(TWELVE_MONTHS)), + st.integers(min_value=0, max_value=5_000), + ), + st.tuples( + st.just("set_input"), + st.just("birth"), + st.sampled_from(["ETERNITY", "2025", "2025-02"]), + st.integers(min_value=0, max_value=49), + ), + st.tuples( + st.just("calculate"), st.sampled_from(CALCULATED), st.sampled_from(MONTHS) + ), + st.tuples(st.just("delete"), st.sampled_from(DISK_INPUTS), _months_or_all), + st.tuples(st.just("holder_delete"), st.sampled_from(DISK_INPUTS), _months_or_all), + st.tuples(st.just("on_disk"), st.booleans()), + st.tuples(st.just("invalidate")), +) + + +@hypothesis.settings( + max_examples=200, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given(operations=st.lists(_disk_operation, max_size=25)) +def test_user_input_keys_follow_memory_and_disk_storage(operations): + storages = _Storages(TAX_BENEFIT_SYSTEM) + storages.assert_record_matches() + for step, operation in enumerate(operations): + storages.apply(operation) + try: + storages.assert_record_matches() + except AssertionError as error: + raise AssertionError((step, operation)) from error