From 390e9e80a58176296ec05003f2e4633d64e933f5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 01:14:51 -0400 Subject: [PATCH 01/12] Keep the record of set_input values in step with storage Simulation._user_input_keys records each (variable, branch, period) stored through set_input. _invalidate_all_caches (run by apply_reform) keeps the values it names, to_input_dataframe exports them, and country packages read it to tell an entered value from a calculated one. It drifted from storage in two ways (#559): - delete_arrays deleted the values but kept their entries, so a formula result calculated later for the same period counted as an input: it survived apply_reform and was exported. Holder.delete_arrays, which Simulation.delete_arrays calls for each branch it deletes from, now drops the entries for the variable, that branch and the periods in-memory storage deletes (all of them for an eternal variable). Code that deletes through the holder, as country packages do when they move an input to another variable, is covered too. Disk storage deletes only the period asked for, so the entry for a value it still holds is kept. - clone (so also get_branch) shared the record between simulations that store their values separately, so an input set on a clone, on a branch's parent after the branch was made, or on the original after cloning was recorded for both. The copy now gets its own record, and its own empty list of running set_input calls. Tests: example regressions (11 of 12 fail before the fix; the twelfth guards against dropping too much), and a Hypothesis property that runs random set_input / calculate / delete_arrays / clone / get_branch / _invalidate_all_caches sequences against a reference model of each simulation's inputs. Fixes #559 Co-Authored-By: Claude Opus 5.5 --- .../user-input-keys-follow-storage.fixed.md | 1 + policyengine_core/holders/holder.py | 45 +++ policyengine_core/simulations/simulation.py | 14 +- tests/core/test_user_input_keys.py | 254 +++++++++++++++++ tests/core/test_user_input_keys_property.py | 266 ++++++++++++++++++ 5 files changed, 579 insertions(+), 1 deletion(-) create mode 100644 changelog.d/user-input-keys-follow-storage.fixed.md create mode 100644 tests/core/test_user_input_keys.py create mode 100644 tests/core/test_user_input_keys_property.py 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 000000000..585cf3932 --- /dev/null +++ b/changelog.d/user-input-keys-follow-storage.fixed.md @@ -0,0 +1 @@ +`delete_arrays` now drops the record of the inputs it deletes, and `clone` (so also `get_branch`) gives the copy its own record, so `apply_reform` and `to_input_dataframe` no longer treat a later formula result, or another simulation's input, as an input. diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 11a8ae50b..9dfed4945 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -100,11 +100,56 @@ 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. """ self._memory_storage.delete(period, branch_name) if self._disk_storage: self._disk_storage.delete(period, branch_name) + self._forget_deleted_inputs(period, branch_name) + + def _forget_deleted_inputs(self, period: Period, branch_name: str) -> None: + """Drop the simulation's record of the inputs ``delete_arrays`` deleted. + + ``_user_input_keys`` records each (variable, branch, period) stored + through ``set_input``. ``_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 entries dropped are this variable's, for ``branch_name``, in the + periods in-memory storage deletes: every period if ``period`` is + ``None`` or the variable is eternal, otherwise those within + ``period``. Disk storage deletes only ``period`` itself, so the entry + for a value it still holds is kept. + """ + simulation = getattr(self, "simulation", None) + user_input_keys = getattr(simulation, "_user_input_keys", None) + if not user_input_keys: + return + if period is not None: + period = periods.period(period) + name = self.variable.name + deleted = [ + key + for key in user_input_keys + if key[0] == name + and key[1] == branch_name + and ( + period is None + or self._memory_storage.is_eternal + or period.contains(periods.period(key[2])) + ) + ] + if self._disk_storage is not None: + deleted = [ + key + for key in deleted + if self._disk_storage.get(key[2], branch_name) is None + ] + user_input_keys.difference_update(deleted) def _get_array_from_storage( self, period: Period, branch_name: str = "default" diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 2652d9723..903b2888d 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -208,7 +208,10 @@ 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: ``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 +1357,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 +1529,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 diff --git a/tests/core/test_user_input_keys.py b/tests/core/test_user_input_keys.py new file mode 100644 index 000000000..393c8e181 --- /dev/null +++ b/tests/core/test_user_input_keys.py @@ -0,0 +1,254 @@ +"""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 + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.entities import Person +from policyengine_core.experimental import MemoryConfig +from policyengine_core.model_api import MONTH, 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 _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_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, so the entry for a monthly value it still holds stays.""" + simulation = _simulation() + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + holder = simulation.get_holder("rent") + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + 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") == { + ("default", JANUARY), + ("default", FEBRUARY), + } + + simulation.delete_arrays("rent", JANUARY) + assert _input_periods(simulation, "rent") == {("default", FEBRUARY)} + assert holder.get_array(FEBRUARY)[0] == 600.0 + + +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_does_not_make_the_original_record_inputs(): + """While ``set_input`` runs on a clone, values the original calculates + are formula results, not inputs.""" + tax_benefit_system = CountryTaxBenefitSystem() + simulation = _simulation(tax_benefit_system) + clone = simulation.clone() + + def set_rent_and_calculate_on_the_original(holder, period, array): + simulation.calculate("income_tax", JANUARY) + for month in period.get_subperiods(periods.MONTH): + holder._set(month, array) + + tax_benefit_system.variables[ + "rent" + ].set_input = set_rent_and_calculate_on_the_original + clone.set_input("rent", "2025", [500.0]) + + assert _key("income_tax", 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 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 000000000..9076e8561 --- /dev/null +++ b/tests/core/test_user_input_keys_property.py @@ -0,0 +1,266 @@ +"""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 +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 the periods the model says are inputs. + +After ``_invalidate_all_caches`` the simulation and its branches store their +inputs and nothing else. ``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.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"] +EXPORTABLE = ["salary", "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(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 == "rent": + return [float(number)] + return [float(number), float(number) / 2] + + +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 = {} + for key in root._user_input_keys: + variable, branch, period = key + stored = root.get_holder(variable)._memory_storage._arrays + inputs[key] = stored[_slot(key)[1]].copy() + 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) + key = (variable, model.branch_name, periods.period(period)) + array = simulation.get_holder(variable)._to_array(np.asarray(value)) + # An eternal variable stores one value whatever the period, so a + # value set for another period replaces it, and every period it + # was set for keeps its entry. + for other in model.inputs: + if _slot(other) == _slot(key): + model.inputs[other] = array + 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 + visible = model.visible_branch_names() + for variable in EXPORTABLE: + exported = simulation._get_exportable_input_periods(variable, False) + expected_periods = { + periods.period(period) + for name, branch, period in model.inputs + if name == variable and branch in visible + } + if variable == "birth": + assert bool(exported) == bool(expected_periods) + else: + assert set(exported) == expected_periods + + +_months_or_all = st.sampled_from(MONTHS + ["2025", 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("calculate"), + _index, + st.sampled_from(CALCULATED), + st.sampled_from(MONTHS), + ), + st.tuples( + st.just("delete"), + _index, + st.sampled_from(["salary", "income_tax", "rent", "birth"]), + _months_or_all, + ), + st.tuples( + st.just("holder_delete"), + _index, + st.sampled_from(["salary", "income_tax", "rent", "birth"]), + _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", "b2"])), + 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(tax_benefit_system, 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 From a9bdeff03dbcd553a123b3be03ec5500adf8b1f3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 02:28:18 -0400 Subject: [PATCH 02/12] Prune the record from what storage deleted; record each value once Follow-up to the review of the first commit: - Holder.delete_arrays no longer looks through the whole record. It compares the periods memory stores for the branch before and after the deletion and discards exactly those entries, so its cost does not grow with the record. Country marginal-rate code deletes every variable on a branch: with 9,000 entries and 3,024 variables the loop took 0.73 s with the scan and takes 0.018 s now (0.004 s without any pruning). A holder with disk storage, which cannot list its periods for every branch name, still looks through the record for the variable's entries in the deleted periods and drops those whose value neither storage holds. - Holder._set records the period as storage keys the value: eternity for an eternal variable whatever period it was set for, and a Period for a handler that passes a string. Each entry names one stored value, so a string-period input is exported and deleting an eternal input drops its entry. - put_in_cache stores with is_input=False (the same keyword as #560), so values a custom set_input handler calculates are formula results, which apply_reform recalculates, not inputs. - subsample starts the record again before it rebuilds the simulation, so it records only what the rebuild stores. Tests: a guard that a memory-only deletion never iterates the record; the eternal, string-period, calculating-handler, disk and subsample cases; the property's model now includes a handler that calculates and stores months under string periods. Co-Authored-By: Claude Opus 5.5 --- .../user-input-keys-follow-storage.fixed.md | 2 +- policyengine_core/holders/holder.py | 113 ++++++++++----- policyengine_core/simulations/simulation.py | 12 +- tests/core/test_user_input_keys.py | 136 +++++++++++++++++- tests/core/test_user_input_keys_property.py | 62 ++++++-- 5 files changed, 261 insertions(+), 64 deletions(-) diff --git a/changelog.d/user-input-keys-follow-storage.fixed.md b/changelog.d/user-input-keys-follow-storage.fixed.md index 585cf3932..7c4704611 100644 --- a/changelog.d/user-input-keys-follow-storage.fixed.md +++ b/changelog.d/user-input-keys-follow-storage.fixed.md @@ -1 +1 @@ -`delete_arrays` now drops the record of the inputs it deletes, and `clone` (so also `get_branch`) gives the copy its own record, so `apply_reform` and `to_input_dataframe` no longer treat a later formula result, or another simulation's input, as an input. +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). diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 9dfed4945..99eac83d1 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 @@ -104,53 +104,79 @@ def delete_arrays( 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) + in_memory = self._periods_in_memory(branch_name) if user_input_keys else None self._memory_storage.delete(period, branch_name) if self._disk_storage: self._disk_storage.delete(period, branch_name) - self._forget_deleted_inputs(period, branch_name) - def _forget_deleted_inputs(self, period: Period, branch_name: str) -> None: + if user_input_keys: + self._forget_deleted_inputs(user_input_keys, period, branch_name, in_memory) + + def _periods_in_memory(self, branch_name: str) -> set: + return { + stored_period + for stored_branch, stored_period in ( + self._memory_storage.get_known_branch_periods() + ) + if stored_branch == branch_name + } + + def _forget_deleted_inputs( + self, + user_input_keys: set, + period: Period, + branch_name: str, + in_memory_before: 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``. ``_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 entries dropped are this variable's, for ``branch_name``, in the - periods in-memory storage deletes: every period if ``period`` is - ``None`` or the variable is eternal, otherwise those within - ``period``. Disk storage deletes only ``period`` itself, so the entry - for a value it still holds is kept. + 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 entries dropped are those for the values the deletion removed. + In memory, those are found by comparing what memory stores for + ``branch_name`` before and after, so the cost does not grow with the + record. Disk storage cannot list its periods for every branch name, + so for a holder with disk storage this also looks through the + record for the variable's entries in the deleted periods, and drops + those whose value neither storage still holds (disk storage deletes + only ``period`` itself, not the periods within it). """ - simulation = getattr(self, "simulation", None) - user_input_keys = getattr(simulation, "_user_input_keys", None) - if not user_input_keys: - return - if period is not None: - period = periods.period(period) name = self.variable.name - deleted = [ - key - for key in user_input_keys - if key[0] == name - and key[1] == branch_name - and ( - period is None - or self._memory_storage.is_eternal - or period.contains(periods.period(key[2])) - ) - ] + in_memory = self._periods_in_memory(branch_name) + deleted = { + (name, branch_name, removed) for removed in in_memory_before - in_memory + } if self._disk_storage is not None: - deleted = [ + within = None if period is None else periods.period(period) + deleted.update( key - for key in deleted - if self._disk_storage.get(key[2], branch_name) is None - ] + for key in user_input_keys + if key[0] == name + and key[1] == branch_name + and ( + within is None + or self._memory_storage.is_eternal + or within.contains(key[2]) + ) + and key[2] not in in_memory + and self._disk_storage.get(key[2], branch_name) is None + ) user_input_keys.difference_update(deleted) + def _storage_period(self, period: Period) -> Period: + """The period storage keys a value for ``period`` under.""" + if self._memory_storage.is_eternal: + return periods.period(periods.ETERNITY) + return periods.period(period) + def _get_array_from_storage( self, period: Period, branch_name: str = "default" ) -> ArrayLike: @@ -392,10 +418,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: @@ -415,10 +447,15 @@ 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), 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" @@ -433,7 +470,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 903b2888d..1092eca3f 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -209,9 +209,10 @@ def __init__( # 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. The record follows this simulation's - # storage: ``delete_arrays`` drops the entries for the values it - # deletes, and ``clone`` (so also ``get_branch``) gives the copy its - # own record. + # 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 @@ -1987,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 index 393c8e181..2bbd5f671 100644 --- a/tests/core/test_user_input_keys.py +++ b/tests/core/test_user_input_keys.py @@ -16,10 +16,13 @@ from __future__ import annotations import numpy as np +import pandas as pd 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.experimental import MemoryConfig from policyengine_core.model_api import MONTH, Reform, Variable from policyengine_core.simulations import SimulationBuilder @@ -117,6 +120,16 @@ def test_delete_for_a_period_drops_only_entries_within_it(): 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.""" @@ -187,6 +200,50 @@ def test_entry_is_kept_for_a_value_disk_storage_did_not_delete(): 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() @@ -232,23 +289,88 @@ def test_parent_does_not_see_inputs_set_on_its_branch(): assert _input_periods(simulation, "salary") == {("default", JANUARY)} -def test_set_input_on_a_clone_does_not_make_the_original_record_inputs(): - """While ``set_input`` runs on a clone, values the original calculates - are formula results, not inputs.""" +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_calculate_on_the_original(holder, period, array): + 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_calculate_on_the_original + 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], + ) diff --git a/tests/core/test_user_input_keys_property.py b/tests/core/test_user_input_keys_property.py index 9076e8561..5c5c3123c 100644 --- a/tests/core/test_user_input_keys_property.py +++ b/tests/core/test_user_input_keys_property.py @@ -12,6 +12,9 @@ - each entry names a value the simulation stores, equal to the input; - ``to_input_dataframe`` exports the periods the model says are inputs. +One input variable is set through a custom ``set_input`` handler that +calculates another variable first and stores months under string periods. + After ``_invalidate_all_caches`` the simulation and its branches store their inputs and nothing else. ``test_user_input_keys.py`` pins the same behaviour with examples. @@ -27,6 +30,9 @@ 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.model_api import MONTH, Variable from policyengine_core.simulations import SimulationBuilder SITUATION = { @@ -37,7 +43,29 @@ "households": {"h": {"parents": ["a", "b"], "rent": {"2025-01": 800}}}, } MONTHS = ["2025-01", "2025-02", "2025-03"] -EXPORTABLE = ["salary", "rent", "birth"] + + +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", @@ -64,7 +92,7 @@ def _slot(key): def _input_value(variable, number): if variable == "birth": return [f"19{50 + number % 50}-01-01"] * 2 - if variable == "rent": + if variable in ("rent", "quarterly_rent"): return [float(number)] return [float(number), float(number) / 2] @@ -133,15 +161,14 @@ def apply(self, operation): variable, period, number = arguments value = _input_value(variable, number) simulation.set_input(variable, period, value) - key = (variable, model.branch_name, periods.period(period)) array = simulation.get_holder(variable)._to_array(np.asarray(value)) - # An eternal variable stores one value whatever the period, so a - # value set for another period replaces it, and every period it - # was set for keeps its entry. - for other in model.inputs: - if _slot(other) == _slot(key): - model.inputs[other] = array - model.inputs[key] = array + # 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) @@ -224,6 +251,13 @@ def assert_records_match(self): st.sampled_from(["ETERNITY", "2025", "2025-02"]), st.integers(min_value=0, max_value=49), ), + 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, @@ -233,13 +267,13 @@ def assert_records_match(self): st.tuples( st.just("delete"), _index, - st.sampled_from(["salary", "income_tax", "rent", "birth"]), + st.sampled_from(INPUTS), _months_or_all, ), st.tuples( st.just("holder_delete"), _index, - st.sampled_from(["salary", "income_tax", "rent", "birth"]), + st.sampled_from(INPUTS), _months_or_all, st.integers(min_value=0, max_value=3), ), @@ -255,8 +289,8 @@ def assert_records_match(self): suppress_health_check=[hypothesis.HealthCheck.too_slow], ) @hypothesis.given(operations=st.lists(_operation, max_size=30)) -def test_user_input_keys_match_reference_model(tax_benefit_system, operations): - family = _Family(tax_benefit_system) +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) From da05eeb43f3ed0515633754e8a9690917232166a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 08:04:40 -0400 Subject: [PATCH 03/12] Keep the entry of an input a storage still holds; record twelve months as stored Round-2 review fixes: - An entry is dropped only when neither storage still holds its value, so an input deleted from memory that survives on disk stays an input. - Twelve months starting on the first of a month are recorded as the year storage keys them under, so deleting that year drops the entry. - Deleting compares the keys each storage holds before and after: it no longer goes through the record or loads files for a disk-backed holder. - The property model is seeded from the situation, checks to_input_dataframe itself, and a second property covers memory and disk storage. - Regression for the carry-over case found in the review of #562. Co-Authored-By: Claude Opus 5.5 --- .../user-input-keys-follow-storage.fixed.md | 2 +- policyengine_core/holders/holder.py | 103 +++---- tests/core/test_user_input_keys.py | 257 +++++++++++++++++- tests/core/test_user_input_keys_property.py | 231 ++++++++++++++-- 4 files changed, 506 insertions(+), 87 deletions(-) diff --git a/changelog.d/user-input-keys-follow-storage.fixed.md b/changelog.d/user-input-keys-follow-storage.fixed.md index 7c4704611..1efbd3633 100644 --- a/changelog.d/user-input-keys-follow-storage.fixed.md +++ b/changelog.d/user-input-keys-follow-storage.fixed.md @@ -1 +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). +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 99eac83d1..f85bb80db 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -106,30 +106,24 @@ def delete_arrays( """ simulation = getattr(self, "simulation", None) user_input_keys = getattr(simulation, "_user_input_keys", None) - in_memory = self._periods_in_memory(branch_name) if user_input_keys else 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, period, branch_name, in_memory) + self._forget_deleted_inputs(user_input_keys, stored_before) - def _periods_in_memory(self, branch_name: str) -> set: - return { - stored_period - for stored_branch, stored_period in ( - self._memory_storage.get_known_branch_periods() - ) - if stored_branch == branch_name - } + 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, - period: Period, - branch_name: str, - in_memory_before: set, + self, user_input_keys: set, stored_before: Tuple[set, set] ) -> None: """Drop the simulation's record of the inputs ``delete_arrays`` deleted. @@ -140,42 +134,56 @@ def _forget_deleted_inputs( deleted value would make a formula result stored later for that period count as an input. - The entries dropped are those for the values the deletion removed. - In memory, those are found by comparing what memory stores for - ``branch_name`` before and after, so the cost does not grow with the - record. Disk storage cannot list its periods for every branch name, - so for a holder with disk storage this also looks through the - record for the variable's entries in the deleted periods, and drops - those whose value neither storage still holds (disk storage deletes - only ``period`` itself, not the periods within it). + 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 - in_memory = self._periods_in_memory(branch_name) - deleted = { - (name, branch_name, removed) for removed in in_memory_before - in_memory - } - if self._disk_storage is not None: - within = None if period is None else periods.period(period) - deleted.update( - key - for key in user_input_keys - if key[0] == name - and key[1] == branch_name - and ( - within is None - or self._memory_storage.is_eternal - or within.contains(key[2]) - ) - and key[2] not in in_memory - and self._disk_storage.get(key[2], branch_name) is None - ) - user_input_keys.difference_update(deleted) + memory_before, disk_before = stored_before + memory_after, disk_after = self._stored_keys() + # Memory keys are "{branch}:{period}"; branch names cannot contain ":". + removed = [key.split(":", 1) for key in memory_before - memory_after] + # 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 - disk_after] + 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.""" + """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(period) + return periods.period(str(periods.period(period))) def _get_array_from_storage( self, period: Period, branch_name: str = "default" @@ -451,8 +459,9 @@ def _set( if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() # Record the period as storage keys the value (eternity for an - # eternal variable, whatever period it was set for), so each - # entry names one stored value. + # 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)) ) diff --git a/tests/core/test_user_input_keys.py b/tests/core/test_user_input_keys.py index 2bbd5f671..b656a3b84 100644 --- a/tests/core/test_user_input_keys.py +++ b/tests/core/test_user_input_keys.py @@ -17,14 +17,16 @@ 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, Reform, Variable +from policyengine_core.model_api import MONTH, YEAR, Reform, Variable from policyengine_core.simulations import SimulationBuilder JANUARY = "2025-01" @@ -45,6 +47,32 @@ 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)) @@ -180,22 +208,21 @@ def test_branch_delete_drops_entries_for_the_branches_it_deletes_from(): 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, so the entry for a monthly value it still holds stays.""" + within it (policyengine-core#564), so the entry for a monthly value it + still holds stays.""" simulation = _simulation() - simulation.memory_config = MemoryConfig(max_memory_occupation=0) + _store_on_disk(simulation, "rent") holder = simulation.get_holder("rent") - holder._disk_storage = holder.create_disk_storage() - holder._on_disk_storable = True 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") == { - ("default", JANUARY), - ("default", FEBRUARY), - } + 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 @@ -374,3 +401,215 @@ def test_subsample_records_only_the_inputs_it_stores(): [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 index 5c5c3123c..a7db3b57e 100644 --- a/tests/core/test_user_input_keys_property.py +++ b/tests/core/test_user_input_keys_property.py @@ -2,22 +2,29 @@ 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 -tracks, for each simulation, which (variable, branch, period) values were -stored through ``set_input`` and not deleted since. After every step, for -every simulation: +``_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 the periods the model says are inputs. +- ``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. ``test_user_input_keys.py`` pins the same behaviour -with examples. +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 @@ -32,6 +39,7 @@ 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 @@ -43,6 +51,17 @@ "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): @@ -81,7 +100,7 @@ class quarterly_rent(Variable): def _canonical(variable, period): if variable == "birth": return periods.period(periods.ETERNITY) - return periods.period(period) + return periods.period(TWELVE_MONTHS.get(period, period)) def _slot(key): @@ -97,6 +116,12 @@ def _input_value(variable, 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.""" @@ -137,11 +162,10 @@ def delete(self, variable, period, branch_name): class _Family: def __init__(self, tax_benefit_system): root = SimulationBuilder().build_from_entities(tax_benefit_system, SITUATION) - inputs = {} - for key in root._user_input_keys: - variable, branch, period = key - stored = root.get_holder(variable)._memory_storage._arrays - inputs[key] = stored[_slot(key)[1]].copy() + 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)] @@ -161,7 +185,7 @@ def apply(self, operation): variable, period, number = arguments value = _input_value(variable, number) simulation.set_input(variable, period, value) - array = simulation.get_holder(variable)._to_array(np.asarray(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. @@ -220,21 +244,54 @@ def assert_records_match(self): stored = holder._memory_storage._arrays.get(storage_key) assert stored is not None, key assert np.array_equal(stored, expected), key - visible = model.visible_branch_names() - for variable in EXPORTABLE: - exported = simulation._get_exportable_input_periods(variable, False) - expected_periods = { - periods.period(period) - for name, branch, period in model.inputs - if name == variable and branch in visible - } - if variable == "birth": - assert bool(exported) == bool(expected_periods) - else: - assert set(exported) == expected_periods - - -_months_or_all = st.sampled_from(MONTHS + ["2025", None]) + 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( @@ -251,6 +308,13 @@ def assert_records_match(self): 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, @@ -278,7 +342,7 @@ def assert_records_match(self): st.integers(min_value=0, max_value=3), ), st.tuples(st.just("clone"), _index), - st.tuples(st.just("branch"), _index, st.sampled_from(["b1", "b2"])), + st.tuples(st.just("branch"), _index, st.sampled_from(["b1", "no_salt"])), st.tuples(st.just("invalidate"), _index), ) @@ -298,3 +362,110 @@ def test_user_input_keys_match_reference_model(operations): 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 From 579c375144b7aee4427bcd76325c46f96f1c6643 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 08:11:55 -0400 Subject: [PATCH 04/12] Find the keys a delete removed without a second snapshot Co-Authored-By: Claude Opus 5.5 --- policyengine_core/holders/holder.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index f85bb80db..0c8473b54 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -144,12 +144,13 @@ def _forget_deleted_inputs( """ name = self.variable.name memory_before, disk_before = stored_before - memory_after, disk_after = self._stored_keys() + 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 - memory_after] + 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 - disk_after] + removed += [key.rsplit("_", 1) for key in disk_before if key not in files] forgotten = [] for branch, period_string in removed: try: From 64b515768d004fa4d89e853ca06d00e41be8f95e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:13:31 -0400 Subject: [PATCH 05/12] Track actual input provenance across storage tiers and record early periods safely Preserve replayed-input tuples and master carry-over marks; add r3b soundness witnesses and value-aware disk properties. Co-Authored-By: Claude Opus 5.5 --- .../user-input-keys-follow-storage.fixed.md | 2 +- policyengine_core/holders/holder.py | 122 +++++- policyengine_core/simulations/simulation.py | 11 +- tests/core/test_user_input_keys_property.py | 138 ++++++- tests/core/test_user_input_keys_soundness.py | 370 ++++++++++++++++++ 5 files changed, 602 insertions(+), 41 deletions(-) create mode 100644 tests/core/test_user_input_keys_soundness.py diff --git a/changelog.d/user-input-keys-follow-storage.fixed.md b/changelog.d/user-input-keys-follow-storage.fixed.md index 1efbd3633..0e8893a34 100644 --- a/changelog.d/user-input-keys-follow-storage.fixed.md +++ b/changelog.d/user-input-keys-follow-storage.fixed.md @@ -1 +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). +Keep the record of `set_input` values in step with each simulation's storage, including mixed memory/disk provenance after deletion, isolated copies, calculating handlers, subsampling, storage-period aliases, and valid early-year inputs. diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 36c998952..d6062c802 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -22,6 +22,8 @@ class Holder: A holder keeps tracks of a variable values after they have been calculated, or set as an input. """ + _user_input_storage: dict[tuple[str, str], tuple[Period, frozenset[str]]] + def __init__(self, variable: "Variable", population: "Population"): self.population = population self.variable = variable @@ -62,6 +64,7 @@ def clone(self, population: "Population", share_arrays: bool = False) -> "Holder "simulation", "_memory_storage", "_disk_storage", + "_user_input_storage", ): new_dict[key] = value @@ -73,6 +76,9 @@ def clone(self, population: "Population", share_arrays: bool = False) -> "Holder new._disk_storage = ( self._disk_storage.clone() if self._disk_storage is not None else None ) + if hasattr(self, "_user_input_storage"): + # Locations are immutable; only the holder's index needs copying. + new._user_input_storage = self._user_input_storage.copy() new_dict["population"] = population new_dict["simulation"] = population.simulation @@ -141,10 +147,9 @@ def _forget_deleted_inputs( 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. + read. An entry for a removed key stays only while a storage still + holds an input for it. A computed disk value does not become an input + when the memory input for the same key is deleted. """ name = self.variable.name memory_before, disk_before = stored_before @@ -156,28 +161,98 @@ def _forget_deleted_inputs( # period strings cannot, so the period follows the last "_". removed += [key.rsplit("_", 1) for key in disk_before if key not in files] forgotten = [] + recorded = getattr(self, "_user_input_storage", {}) 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): + slot = (branch, period_string) + input_storage = recorded.get(slot) + if input_storage is not None: + input_period, locations = input_storage + locations = frozenset( + location + for location in locations + if ( + f"{branch}:{period_string}" in arrays + if location == "memory" + else f"{branch}_{period_string}" in files + ) + ) + recorded[slot] = (input_period, locations) + else: + try: + input_period = periods.period(period_string) + except ValueError: + # Not a key ``put`` wrote (say, a file ``restore`` found). + continue + key = (name, branch, input_period) + if key in user_input_keys and not self._stores_user_input( + input_period, branch + ): forgotten.append(key) + recorded.pop(slot, None) user_input_keys.difference_update(forgotten) + def _record_input_storage( + self, period: Period, branch_name: str, storage: str, is_input: bool = True + ) -> None: + """Remember which tier received the actual ``set_input`` value. + + Storage's carry-over marks are broader: ``put_in_cache`` defaults to + ``derived=False`` without making its value a replay/export input. + Keep this metadata only for slots that have received an input, with + immutable locations so clones can copy the index independently. + """ + recorded = getattr(self, "_user_input_storage", None) + slot = (branch_name, str(period)) + if not is_input and (recorded is None or slot not in recorded): + return + if recorded is None: + recorded = self._user_input_storage = {} + _, locations = recorded.get(slot, (period, frozenset())) + locations = locations | {storage} if is_input else locations - {storage} + recorded[slot] = (period, locations) + + def _stores_user_input(self, period: Period, branch_name: str) -> 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. + """ + recorded = getattr(self, "_user_input_storage", {}).get( + (branch_name, str(period)) + ) + locations = recorded[1] if recorded is not None else {"memory", "disk"} + for location, storage in ( + ("memory", self._memory_storage), + ("disk", self._disk_storage), + ): + if ( + location in locations + and storage is not None + and storage.has(period, branch_name) + and not storage.is_derived(period, branch_name) + ): + return True + return False + 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``). + Match the string key's aliases using the period fields: eternity for + an eternal variable, and a year for twelve months. Do not reparse the + string: valid early-year periods stringify without leading zeros. """ if self._memory_storage.is_eternal: return periods.period(periods.ETERNITY) - return periods.period(str(periods.period(period))) + period = periods.period(period) + if period.unit in (periods.MONTH, periods.YEAR): + unit, size = period.unit, period.size + if unit == periods.MONTH and size == 12: + unit, size = periods.YEAR, 1 + start = periods.instant((period.start.year, period.start.month, 1)) + return Period((unit, start, size)) + return period def _get_array_from_storage( self, period: Period, branch_name: str = "default" @@ -448,10 +523,23 @@ def _set( >= self.simulation.memory_config.max_memory_occupation_pc ) + stored_period = ( + self._storage_period(period) + if is_input or hasattr(self, "_user_input_storage") + else None + ) + if should_store_on_disk: self._disk_storage.put(value, period, branch_name, derived=derived) else: self._memory_storage.put(value, period, branch_name, derived=derived) + if stored_period is not None: + self._record_input_storage( + stored_period, + branch_name, + "disk" if should_store_on_disk else "memory", + is_input=is_input, + ) if is_input and simulation is not None: if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() @@ -460,7 +548,7 @@ def _set( # 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)) + (self.variable.name, branch_name, stored_period) ) def put_in_cache( diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 6b79c5409..2ebd8c378 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -501,6 +501,8 @@ def _invalidate_all_caches(self) -> None: if holder._disk_storage is not None: holder._disk_storage._files = {} holder._disk_storage._derived = set() + if hasattr(holder, "_user_input_storage"): + holder._user_input_storage = {} # Replay preserved user inputs so ``calculate`` still sees them. for user_input in preserved: holder = self.get_holder(user_input.variable_name) @@ -516,6 +518,9 @@ def _invalidate_all_caches(self) -> None: user_input.period, user_input.branch_name, ) + holder._record_input_storage( + user_input.period, user_input.branch_name, user_input.storage + ) for branch in self.branches.values(): branch._invalidate_all_caches() @@ -1872,11 +1877,7 @@ 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", ())) + # A ``set_input`` running on this simulation is not running on the copy. new._user_input_contexts = [] # A branch shares its parent's baseline: formulas that run in a diff --git a/tests/core/test_user_input_keys_property.py b/tests/core/test_user_input_keys_property.py index a7db3b57e..4ca22c743 100644 --- a/tests/core/test_user_input_keys_property.py +++ b/tests/core/test_user_input_keys_property.py @@ -22,9 +22,10 @@ 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. +memory or on disk, switching between the two. Its model tracks input values +separately in each tier and checks the record, reads and exported values; +a calculation surviving in one tier cannot stand in for a deleted input in +the other. ``test_user_input_keys.py`` pins storage behavior with examples. """ from __future__ import annotations @@ -369,7 +370,12 @@ def test_user_input_keys_match_reference_model(operations): class _Storages: """One simulation whose holders store in memory or on disk, and the - inputs ``set_input`` stored that some storage still holds.""" + input values independently tracked in each storage tier. + + A surviving calculation for the same key cannot replace a deleted + tier's input in the model. The model retains both copies when both + were supplied through ``set_input``, with memory first for reads. + """ def __init__(self, tax_benefit_system): self.simulation = SimulationBuilder().build_from_entities( @@ -380,18 +386,33 @@ def __init__(self, tax_benefit_system): 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 + self.memory_inputs = { + (variable, branch, _canonical(variable, period)): value.copy() + for (variable, branch, period), value in SITUATION_INPUTS.items() } + self.disk_inputs = {} + self.on_disk = True - def holds(self, key): + def holds(self, key, storage): 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 - ) + tier = holder._memory_storage if storage == "memory" else holder._disk_storage + return tier.has(period, branch) + + def _forget_deleted_inputs(self): + self.memory_inputs = { + key: value + for key, value in self.memory_inputs.items() + if self.holds(key, "memory") + } + self.disk_inputs = { + key: value + for key, value in self.disk_inputs.items() + if self.holds(key, "disk") + } + + def readable_inputs(self): + return {**self.disk_inputs, **self.memory_inputs} def stored(self): return { @@ -409,26 +430,98 @@ def apply(self, 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))) + key = (variable, "default", _canonical(variable, period)) + holder = simulation.get_holder(variable) + # Existing memory values are replaced in memory, regardless of + # the threshold controlling subsequent writes to empty slots. + writes_to_disk = self.on_disk and not holder._memory_storage.has( + key[2], key[1] + ) + value = _input_value(variable, number) + simulation.set_input(variable, period, value) + inputs = self.disk_inputs if writes_to_disk else self.memory_inputs + inputs[key] = _expected_array(variable, value) elif kind == "calculate": simulation.calculate(*arguments) elif kind == "delete": simulation.delete_arrays(*arguments) + self._forget_deleted_inputs() elif kind == "holder_delete": variable, period = arguments simulation.get_holder(variable).delete_arrays(period) + self._forget_deleted_inputs() elif kind == "on_disk": (on_disk,) = arguments + self.on_disk = on_disk simulation.memory_config.max_memory_occupation_pc = 0 if on_disk else 101 elif kind == "invalidate": + expected = self.readable_inputs() 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)} + # Replay preserves memory before disk and keeps one copy. It + # must discard every stored value that the model did not input. + assert self.stored() == set(expected) + self.disk_inputs = { + key: value + for key, value in self.disk_inputs.items() + if key not in self.memory_inputs + } def assert_record_matches(self): - assert set(self.simulation._user_input_keys) == self.inputs + simulation = self.simulation + inputs = self.readable_inputs() + assert set(simulation._user_input_keys) == set(inputs) + for tier, expected_inputs in ( + ("memory", self.memory_inputs), + ("disk", self.disk_inputs), + ): + for (variable, branch, period), expected in expected_inputs.items(): + holder = simulation.get_holder(variable) + storage = ( + holder._memory_storage if tier == "memory" else holder._disk_storage + ) + actual = storage.get(period, branch) + assert actual is not None, (tier, variable, branch, period) + assert np.array_equal(actual, expected), ( + tier, + variable, + branch, + period, + ) + for (variable, branch, period), expected in inputs.items(): + actual = simulation.get_holder(variable).get_array(period, branch) + assert actual is not None, (variable, branch, period) + assert np.array_equal(actual, expected), (variable, branch, period) + self._assert_export_matches(inputs) + + def _assert_export_matches(self, inputs): + simulation = self.simulation + expected = {} + for (variable, branch, period), value in inputs.items(): + definition_period = simulation.tax_benefit_system.get_variable( + variable + ).definition_period + if ( + branch == "default" + and variable in EXPORTABLE + and period.unit == definition_period + ): + if variable in ("rent", "quarterly_rent"): + value = np.repeat(value, 2) + expected[f"{variable}__{period}"] = value + stored_before = self.stored() + fast_cache = dict(simulation._fast_cache) + + exported = simulation.to_input_dataframe() + + # Preserve the generated sequence's cache state while checking + # export's actual columns and values against the independent model. + simulation._fast_cache = fast_cache + assert self.stored() == 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 _disk_operation = st.one_of( @@ -459,6 +552,15 @@ def assert_record_matches(self): deadline=None, suppress_health_check=[hypothesis.HealthCheck.too_slow], ) +@hypothesis.example( + operations=[ + ("calculate", "rent", "2025-02"), + ("on_disk", False), + ("set_input", "rent", "2025-02", 500), + ("delete", "rent", "2025"), + ("invalidate",), + ] +) @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) diff --git a/tests/core/test_user_input_keys_soundness.py b/tests/core/test_user_input_keys_soundness.py new file mode 100644 index 000000000..d8375e655 --- /dev/null +++ b/tests/core/test_user_input_keys_soundness.py @@ -0,0 +1,370 @@ +"""Value and provenance regressions for the input-record soundness review. + +The replay/export record is narrower than storage's carry-over input marks: +only values written through ``set_input`` belong to that record. Two tiers +holding the same key can hold values with different input provenance. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.entities import Person +from policyengine_core.model_api import Reform, Variable +from tests.core.test_user_input_keys import ( + JANUARY, + _input_periods, + _key, + _simulation, + _store_on_disk, +) + + +class NoopReform(Reform): + def apply(self): + pass + + +class raw_year(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Raw yearly input" + + +class raw_month(Variable): + value_type = float + entity = Person + definition_period = periods.MONTH + label = "Raw monthly input" + + +class raw_day(Variable): + value_type = float + entity = Person + definition_period = periods.DAY + label = "Raw daily input" + + +@pytest.mark.parametrize( + "period_string", + ["day:2025-01-03:2", "year:2025:2", "month:2025-01:2"], +) +def test_raw_period_keeps_its_unit_and_size_through_invalidation(period_string): + simulation = _simulation() + holder = simulation.get_holder("income_tax") + # The fixture's divide handler otherwise expands a year into months and + # converts a day period into an unsupported mid-month monthly period. + holder.variable.set_input = None + period = periods.period(period_string) + + holder.set_input(period, [42.0]) + + assert _key("income_tax", period) in simulation._user_input_keys + assert _input_periods(simulation, "income_tax") == {("default", period_string)} + np.testing.assert_array_equal(holder.get_array(period), [42.0]) + + simulation._invalidate_all_caches() + + assert _key("income_tax", period) in simulation._user_input_keys + assert _input_periods(simulation, "income_tax") == {("default", period_string)} + np.testing.assert_array_equal(holder.get_array(period), [42.0]) + + +def test_invisible_memory_branch_deletion_removes_its_input_record(): + simulation = _simulation() + holder = simulation.get_holder("salary") + holder.set_input(JANUARY, [42.0], "hidden") + assert _key("salary", JANUARY, "hidden") in simulation._user_input_keys + + holder.delete_arrays(None, "hidden") + + assert _key("salary", JANUARY, "hidden") not in simulation._user_input_keys + assert not holder._memory_storage.has(JANUARY, "hidden") + assert _input_periods(simulation, "salary") == {("default", JANUARY)} + np.testing.assert_array_equal(holder.get_array(JANUARY), [1_000.0]) + + +def test_custom_handler_preserves_its_explicit_destination_branch(): + simulation = _simulation() + + def foreign_branch_input(holder, period, array): + holder._set(period.start.period(periods.MONTH), array, "foreign") + + simulation.tax_benefit_system.variables["rent"].set_input = foreign_branch_input + simulation.set_input("rent", "2025", [42.0]) + holder = simulation.get_holder("rent") + + assert _input_periods(simulation, "rent") == {("foreign", JANUARY)} + np.testing.assert_array_equal( + holder._get_array_from_storage(JANUARY, "foreign"), [42.0] + ) + assert holder._get_array_from_storage(JANUARY, "default") is None + assert simulation._user_input_contexts == [] + + +def test_disk_deletion_preserves_digits_in_the_branch_name(): + simulation = _simulation() + _store_on_disk(simulation, "rent") + holder = simulation.get_holder("rent") + holder.set_input(JANUARY, [42.0], "b1") + assert _key("rent", JANUARY, "b1") in simulation._user_input_keys + + holder.delete_arrays(None, "b1") + + assert _input_periods(simulation, "rent") == set() + assert not holder._disk_storage.has(JANUARY, "b1") + + +def test_branch_read_returns_a_writable_copy_without_changing_the_parent(): + simulation = _simulation() + branch = simulation.get_branch("b1") + array = branch.get_holder("salary").get_array(JANUARY) + + assert array.flags.writeable + array[0] = 42.0 + + np.testing.assert_array_equal( + branch.get_holder("salary").get_array(JANUARY), [42.0] + ) + np.testing.assert_array_equal( + simulation.get_holder("salary").get_array(JANUARY), [1_000.0] + ) + + +def test_disk_input_keeps_its_value_in_storage_and_export(): + simulation = _simulation() + _store_on_disk(simulation, "salary") + simulation.delete_arrays("salary", JANUARY) + + simulation.set_input("salary", JANUARY, [42.0]) + + assert _key("salary", JANUARY) in simulation._user_input_keys + np.testing.assert_array_equal( + simulation.get_holder("salary")._disk_storage.get(JANUARY), [42.0] + ) + assert simulation.to_input_dataframe()["salary__2025-01"].tolist() == [42.0] + + +@pytest.mark.parametrize("invalidate", [False, True]) +def test_memory_input_wins_over_a_conflicting_disk_cache(invalidate): + simulation = _simulation() + _store_on_disk(simulation, "rent") + holder = simulation.get_holder("rent") + holder.put_in_cache(np.array([99.0]), JANUARY) + simulation.memory_config.max_memory_occupation_pc = 101 + simulation.set_input("rent", JANUARY, [42.0]) + np.testing.assert_array_equal(holder._memory_storage.get(JANUARY), [42.0]) + np.testing.assert_array_equal(holder._disk_storage.get(JANUARY), [99.0]) + + if invalidate: + simulation._invalidate_all_caches() + + np.testing.assert_array_equal(holder.get_array(JANUARY), [42.0]) + assert _key("rent", JANUARY) in simulation._user_input_keys + assert simulation.to_input_dataframe()["rent__2025-01"].tolist() == [42.0] + if invalidate: + assert holder._disk_storage.get(JANUARY) is None + + +@pytest.mark.parametrize("year", [1, 99, 999, 1000]) +@pytest.mark.parametrize("variable_type", [raw_year, raw_month, raw_day]) +@pytest.mark.parametrize("entrypoint", ["simulation", "holder"]) +@pytest.mark.parametrize("on_disk", [False, True]) +def test_early_year_input_is_stored_and_recorded( + year, variable_type, entrypoint, on_disk +): + system = CountryTaxBenefitSystem() + system.add_variable(variable_type) + simulation = _simulation(system) + variable = variable_type.__name__ + holder = simulation.get_holder(variable) + if on_disk: + _store_on_disk(simulation, variable) + if variable_type is raw_year: + period = periods.period(year) + elif variable_type is raw_month: + period = periods.Period((periods.MONTH, periods.instant((year, 3, 1)), 1)) + else: + period = periods.Period((periods.DAY, periods.instant((year, 3, 17)), 2)) + + if entrypoint == "simulation": + simulation.set_input(variable, period, [120.0]) + else: + holder.set_input(period, [120.0]) + + assert (variable, "default", period) in simulation._user_input_keys + np.testing.assert_array_equal(holder.get_array(period), [120.0]) + storage = holder._disk_storage if on_disk else holder._memory_storage + np.testing.assert_array_equal(storage.get(period), [120.0]) + assert simulation._user_input_contexts == [] + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_zero_padded_early_year_string_records_the_stored_period(on_disk): + system = CountryTaxBenefitSystem() + system.add_variable(raw_year) + simulation = _simulation(system) + if on_disk: + _store_on_disk(simulation, "raw_year") + + simulation.set_input("raw_year", "0999", [123.0]) + + assert _key("raw_year", periods.period(999)) in simulation._user_input_keys + np.testing.assert_array_equal( + simulation.get_holder("raw_year").get_array(periods.period(999)), [123.0] + ) + assert simulation._user_input_contexts == [] + + +def test_deleted_memory_input_does_not_turn_a_disk_calculation_into_an_input(): + simulation = _simulation() + _store_on_disk(simulation, "rent") + holder = simulation.get_holder("rent") + np.testing.assert_array_equal(simulation.calculate("rent", JANUARY), [0.0]) + assert _key("rent", JANUARY) not in simulation._user_input_keys + simulation.memory_config.max_memory_occupation_pc = 101 + simulation.set_input("rent", JANUARY, [500.0]) + np.testing.assert_array_equal(holder._memory_storage.get(JANUARY), [500.0]) + np.testing.assert_array_equal(holder._disk_storage.get(JANUARY), [0.0]) + + simulation.delete_arrays("rent", "2025") + + assert holder._memory_storage.get(JANUARY) is None + assert _key("rent", JANUARY) not in simulation._user_input_keys + assert "rent__2025-01" not in simulation.to_input_dataframe().columns + + simulation.apply_reform(NoopReform) + + assert holder.get_array(JANUARY) is None + assert holder._disk_storage.get(JANUARY) is None + assert _key("rent", JANUARY) not in simulation._user_input_keys + + +def test_deleted_memory_input_keeps_a_surviving_real_disk_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, [700.0]) + np.testing.assert_array_equal(holder._memory_storage.get(JANUARY), [700.0]) + np.testing.assert_array_equal(holder._disk_storage.get(JANUARY), [500.0]) + + simulation.delete_arrays("rent", "2025") + + assert holder._memory_storage.get(JANUARY) is None + assert _key("rent", JANUARY) in simulation._user_input_keys + assert simulation.to_input_dataframe()["rent__2025-01"].tolist() == [500.0] + + simulation.apply_reform(NoopReform) + + assert holder._memory_storage.get(JANUARY) is None + np.testing.assert_array_equal(holder._disk_storage.get(JANUARY), [500.0]) + assert _key("rent", JANUARY) in simulation._user_input_keys + assert simulation.to_input_dataframe()["rent__2025-01"].tolist() == [500.0] + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_reform_does_not_replay_an_unrecorded_carry_over_input(on_disk): + simulation = _simulation() + if on_disk: + _store_on_disk(simulation, "rent") + holder = simulation.get_holder("rent") + holder.put_in_cache(np.array([38.0]), JANUARY) + # Carry-over marks intentionally include ordinary cache writes; the + # replay/export record includes only inputs explicitly recorded by load. + assert not holder.is_derived(JANUARY) + assert _key("rent", JANUARY) not in simulation._user_input_keys + assert "rent__2025-01" not in simulation.to_input_dataframe().columns + + simulation.apply_reform(NoopReform) + + assert holder.get_array(JANUARY) is None + assert _key("rent", JANUARY) not in simulation._user_input_keys + np.testing.assert_array_equal( + simulation.get_holder("salary").get_array(JANUARY), [1_000.0] + ) + + +# Smoke environments may omit dev dependencies. Only the property is skipped; +# ordinary soundness regressions above must still run in those environments. +try: + import hypothesis +except ImportError: + hypothesis = None + + +if hypothesis is not None: + st = hypothesis.strategies + + @hypothesis.settings(max_examples=30, deadline=None) + @hypothesis.example( + disk_is_input=False, delete_memory=True, disk_value=0, memory_value=500 + ) + @hypothesis.example( + disk_is_input=True, delete_memory=True, disk_value=500, memory_value=700 + ) + @hypothesis.example( + disk_is_input=False, delete_memory=False, disk_value=99, memory_value=42 + ) + @hypothesis.given( + disk_is_input=st.booleans(), + delete_memory=st.booleans(), + disk_value=st.integers(min_value=0, max_value=1_000), + memory_value=st.integers(min_value=0, max_value=1_000), + ) + def test_tier_provenance_replays_the_same_inputs_as_a_fresh_simulation( + disk_is_input, delete_memory, disk_value, memory_value + ): + simulation = _simulation() + _store_on_disk(simulation, "rent") + holder = simulation.get_holder("rent") + if disk_is_input: + simulation.set_input("rent", JANUARY, [disk_value]) + else: + holder.put_in_cache(np.array([float(disk_value)]), JANUARY) + simulation.memory_config.max_memory_occupation_pc = 101 + simulation.set_input("rent", JANUARY, [memory_value]) + if delete_memory: + simulation.delete_arrays("rent", "2025") + + expected = ( + (disk_value if disk_is_input else None) if delete_memory else memory_value + ) + reference = _simulation() + if expected is not None: + reference.set_input("rent", JANUARY, [expected]) + + def rent_export(sim): + frame = sim.to_input_dataframe() + return { + column: frame[column].tolist() + for column in frame.columns + if column.startswith("rent__") + } + + assert _input_periods(simulation, "rent") == _input_periods(reference, "rent") + assert rent_export(simulation) == rent_export(reference) + + simulation.apply_reform(NoopReform) + reference.apply_reform(NoopReform) + + assert _input_periods(simulation, "rent") == _input_periods(reference, "rent") + assert rent_export(simulation) == rent_export(reference) + if expected is None: + assert holder.get_array(JANUARY) is None + assert holder._disk_storage.get(JANUARY) is None + else: + np.testing.assert_array_equal( + holder.get_array(JANUARY), + reference.get_holder("rent").get_array(JANUARY), + ) +else: + + @pytest.mark.skip(reason="Hypothesis is not installed") + def test_tier_provenance_replays_the_same_inputs_as_a_fresh_simulation(): + pass From 94d6f82629041fe95097809442e58eca91de82e3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 10:44:37 -0400 Subject: [PATCH 06/12] Preserve registered input tiers and master holder call semantics Seed existing recorded input locations before writes, retain positional derived semantics, and validate period aliases and early-year replay and deletion. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/holders/holder.py | 36 +++++++++++++++- tests/core/test_user_input_keys_property.py | 32 ++++++++++++++ tests/core/test_user_input_keys_soundness.py | 44 ++++++++++++++++++++ 3 files changed, 111 insertions(+), 1 deletion(-) diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index d6062c802..b6dda37f1 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -191,6 +191,34 @@ def _forget_deleted_inputs( recorded.pop(slot, None) user_input_keys.difference_update(forgotten) + def _seed_input_storage(self, period: Period, branch_name: str) -> None: + """Capture registered input locations before their first new write. + + Restoration can register inputs without this holder's tier metadata. + Only an already-recorded slot can inherit storage's input marks: an + unrecorded cache value must never become a replayable input. + """ + recorded = getattr(self, "_user_input_storage", None) + slot = (branch_name, str(period)) + if recorded is not None and slot in recorded: + return + keys = getattr(self.simulation, "_user_input_keys", ()) + if (self.variable.name, branch_name, period) not in keys: + return + locations = frozenset( + location + for location, storage in ( + ("memory", self._memory_storage), + ("disk", self._disk_storage), + ) + if storage is not None + and storage.has(period, branch_name) + and not storage.is_derived(period, branch_name) + ) + if recorded is None: + recorded = self._user_input_storage = {} + recorded[slot] = (period, locations) + def _record_input_storage( self, period: Period, branch_name: str, storage: str, is_input: bool = True ) -> None: @@ -388,6 +416,10 @@ def set_input( If a ``set_input`` property has been set for the variable, this method may accept inputs for periods not matching the ``definition_period`` of the variable. To read more about this, check the `documentation `_. + + Custom handlers store replayable inputs with ``holder._set`` or + ``holder.set_input``. Values written through ``put_in_cache`` are + excluded from input export and reform replay. """ period = periods.period(period) @@ -495,8 +527,8 @@ def _set( value: ArrayLike, branch_name: str = "default", validate_nan: bool = False, - is_input: Optional[bool] = None, derived: bool = False, + is_input: Optional[bool] = None, ) -> None: simulation = getattr(self, "simulation", None) user_input_contexts = ( @@ -528,6 +560,8 @@ def _set( if is_input or hasattr(self, "_user_input_storage") else None ) + if stored_period is not None: + self._seed_input_storage(stored_period, branch_name) if should_store_on_disk: self._disk_storage.put(value, period, branch_name, derived=derived) diff --git a/tests/core/test_user_input_keys_property.py b/tests/core/test_user_input_keys_property.py index 4ca22c743..995e181e3 100644 --- a/tests/core/test_user_input_keys_property.py +++ b/tests/core/test_user_input_keys_property.py @@ -365,6 +365,38 @@ def test_user_input_keys_match_reference_model(operations): raise AssertionError((step, operation)) from error +@hypothesis.settings(max_examples=300, deadline=None) +@hypothesis.given( + unit=st.sampled_from([periods.DAY, periods.MONTH, periods.YEAR]), + year=st.integers(min_value=1, max_value=20_000), + month=st.integers(min_value=1, max_value=12), + day=st.integers(min_value=1, max_value=28), + size=st.integers(min_value=1, max_value=36), + eternal=st.booleans(), +) +def test_period_normalization_matches_storage_strings( + unit, year, month, day, size, eternal +): + """Normalization preserves storage keys and agrees with the old parser + where it accepts the year, while accepting early and late years too.""" + from policyengine_core.data_storage import InMemoryStorage + from policyengine_core.holders import Holder + + holder = object.__new__(Holder) + holder._memory_storage = InMemoryStorage(is_eternal=eternal) + period = periods.Period((unit, periods.instant((year, month, day)), size)) + + canonical = holder._storage_period(period) + + assert holder._storage_period(canonical) == canonical + if eternal: + assert canonical == periods.period(periods.ETERNITY) + else: + assert str(canonical) == str(period) + if 1_000 <= year <= 9_999: + assert canonical == periods.period(str(period)) + + DISK_INPUTS = ["salary", "income_tax", "rent", "birth"] diff --git a/tests/core/test_user_input_keys_soundness.py b/tests/core/test_user_input_keys_soundness.py index d8375e655..9fad83d1e 100644 --- a/tests/core/test_user_input_keys_soundness.py +++ b/tests/core/test_user_input_keys_soundness.py @@ -202,6 +202,12 @@ def test_early_year_input_is_stored_and_recorded( np.testing.assert_array_equal(storage.get(period), [120.0]) assert simulation._user_input_contexts == [] + simulation.apply_reform(NoopReform) + np.testing.assert_array_equal(storage.get(period), [120.0]) + simulation.delete_arrays(variable) + assert not storage.has(period) + assert (variable, "default", period) not in simulation._user_input_keys + @pytest.mark.parametrize("on_disk", [False, True]) def test_zero_padded_early_year_string_records_the_stored_period(on_disk): @@ -290,6 +296,44 @@ def test_reform_does_not_replay_an_unrecorded_carry_over_input(on_disk): ) +def test_registered_disk_input_survives_a_new_memory_input_and_clone_deletion(): + """Restoration can register an input without holder tier metadata (#576).""" + simulation = _simulation() + _store_on_disk(simulation, "rent") + holder = simulation.get_holder("rent") + holder._disk_storage.put(np.array([500.0]), JANUARY) + simulation._user_input_keys.add(_key("rent", JANUARY)) + assert not hasattr(holder, "_user_input_storage") + simulation.memory_config.max_memory_occupation_pc = 101 + + simulation.set_input("rent", JANUARY, [700.0]) + clone = simulation.clone() + clone.delete_arrays("rent", "2025") + + assert _key("rent", JANUARY) in clone._user_input_keys + assert clone.to_input_dataframe()["rent__2025-01"].tolist() == [500.0] + assert simulation.to_input_dataframe()["rent__2025-01"].tolist() == [700.0] + simulation.delete_arrays("rent", "2025") + for sim in (simulation, clone): + assert _key("rent", JANUARY) in sim._user_input_keys + assert sim.to_input_dataframe()["rent__2025-01"].tolist() == [500.0] + sim.apply_reform(NoopReform) + assert sim.get_holder("rent")._memory_storage.get(JANUARY) is None + np.testing.assert_array_equal( + sim.get_holder("rent")._disk_storage.get(JANUARY), [500.0] + ) + + +def test_positional_derived_argument_keeps_masters_meaning(): + simulation = _simulation() + holder = simulation.get_holder("rent") + + holder._set(JANUARY, np.array([42.0]), "default", False, True) + + assert holder.is_derived(JANUARY) + assert _key("rent", JANUARY) not in simulation._user_input_keys + + # Smoke environments may omit dev dependencies. Only the property is skipped; # ordinary soundness regressions above must still run in those environments. try: From 062224ddc7eda6e086ce0a95cc6e79d44bd13e11 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 11:45:56 -0400 Subject: [PATCH 07/12] Exercise twelve-month normalization before input handlers Co-Authored-By: Claude Opus 5.5 --- tests/core/test_user_input_keys_property.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/core/test_user_input_keys_property.py b/tests/core/test_user_input_keys_property.py index 995e181e3..18561468c 100644 --- a/tests/core/test_user_input_keys_property.py +++ b/tests/core/test_user_input_keys_property.py @@ -366,6 +366,12 @@ def test_user_input_keys_match_reference_model(operations): @hypothesis.settings(max_examples=300, deadline=None) +@hypothesis.example( + unit=periods.MONTH, year=2025, month=1, day=1, size=12, eternal=False +) +@hypothesis.example( + unit=periods.MONTH, year=2025, month=3, day=1, size=12, eternal=False +) @hypothesis.given( unit=st.sampled_from([periods.DAY, periods.MONTH, periods.YEAR]), year=st.integers(min_value=1, max_value=20_000), From 7fa7d284d9bdea1e6938b8c957b6c096868a1f4c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 19:34:39 -0400 Subject: [PATCH 08/12] Clarify custom input handler cache contract --- changelog.d/user-input-keys-follow-storage.fixed.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changelog.d/user-input-keys-follow-storage.fixed.md b/changelog.d/user-input-keys-follow-storage.fixed.md index 0e8893a34..c38a8e51c 100644 --- a/changelog.d/user-input-keys-follow-storage.fixed.md +++ b/changelog.d/user-input-keys-follow-storage.fixed.md @@ -1 +1 @@ -Keep the record of `set_input` values in step with each simulation's storage, including mixed memory/disk provenance after deletion, isolated copies, calculating handlers, subsampling, storage-period aliases, and valid early-year inputs. +Keep input records aligned with memory/disk storage across deletion, copies, handlers, subsampling, period aliases, and early years; custom handlers must use `holder._set` or `holder.set_input` for inputs, since `put_in_cache` values are excluded from input export and reform replay and no longer redirected to the input's branch. From 1e079e4dd2a532bd831496169b2e8ce74f314cff Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 20:43:57 -0400 Subject: [PATCH 09/12] Forget restored inputs deleted in early years --- policyengine_core/holders/holder.py | 9 +- tests/core/test_registered_input_deletion.py | 93 ++++++++++++++++++++ 2 files changed, 101 insertions(+), 1 deletion(-) create mode 100644 tests/core/test_registered_input_deletion.py diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index b6dda37f1..4500a118c 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -179,7 +179,14 @@ def _forget_deleted_inputs( recorded[slot] = (input_period, locations) else: try: - input_period = periods.period(period_string) + # Storage strings omit early years' leading zeros; + # the period parser requires a four-digit year. + components = period_string.split(":") + date_index = 1 if len(components) > 1 else 0 + date_components = components[date_index].split("-") + date_components[0] = date_components[0].zfill(4) + components[date_index] = "-".join(date_components) + input_period = periods.period(":".join(components)) except ValueError: # Not a key ``put`` wrote (say, a file ``restore`` found). continue diff --git a/tests/core/test_registered_input_deletion.py b/tests/core/test_registered_input_deletion.py new file mode 100644 index 000000000..0d1f02ab1 --- /dev/null +++ b/tests/core/test_registered_input_deletion.py @@ -0,0 +1,93 @@ +"""Restored input records follow deletion even for early-year storage keys.""" + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem +from tests.core.test_user_input_keys import _simulation, _store_on_disk +from tests.core.test_user_input_keys_soundness import ( + NoopReform, + raw_day, + raw_month, + raw_year, +) + + +def _raw_simulation(variable_type, on_disk): + system = CountryTaxBenefitSystem() + system.add_variable(variable_type) + simulation = _simulation(system) + if on_disk: + _store_on_disk(simulation, variable_type.__name__) + return simulation + + +@pytest.mark.parametrize("year", [1, 99, 999, 1000]) +@pytest.mark.parametrize("variable_type", [raw_year, raw_month, raw_day]) +@pytest.mark.parametrize("on_disk", [False, True]) +def test_deleted_registered_early_year_input_is_forgotten(year, variable_type, on_disk): + """#576 restore can register an input without Holder tier metadata.""" + simulation = _raw_simulation(variable_type, on_disk) + holder = simulation.get_holder(variable_type.__name__) + period = periods.Period( + (variable_type.definition_period, periods.instant((year, 1, 1)), 1) + ) + storage = holder._disk_storage if on_disk else holder._memory_storage + storage.put(np.array([42.0]), period) + key = (variable_type.__name__, "default", period) + simulation._user_input_keys.add(key) + assert not hasattr(holder, "_user_input_storage") + + holder.delete_arrays() + + assert not storage.has(period) + assert key not in simulation._user_input_keys + holder.put_in_cache(np.array([99.0]), period, derived=True) + simulation.apply_reform(NoopReform) + assert holder.get_array(period) is None + + +@pytest.mark.parametrize("year", [1, 999, 1000]) +@pytest.mark.parametrize( + "variable_type, unit, month, day, size", + [ + (raw_year, periods.YEAR, 1, 1, 2), + (raw_year, periods.YEAR, 3, 1, 1), + (raw_month, periods.MONTH, 3, 1, 2), + (raw_month, periods.MONTH, 3, 1, 12), + (raw_day, periods.DAY, 3, 17, 2), + ], +) +@pytest.mark.parametrize("on_disk", [False, True]) +def test_registered_deletion_matches_set_input_for_period_aliases( + year, variable_type, unit, month, day, size, on_disk +): + """Register-only and set_input paths forget the same normalized slot. + + Deletion removes the input from the record regardless of the stored + period's unit, size, early year, alias, tier, or underscored branch. + A write afterward remains a cache value and is cleared by a reform. + """ + restored = _raw_simulation(variable_type, on_disk) + tracked = _raw_simulation(variable_type, on_disk) + period = periods.Period((unit, periods.instant((year, month, day)), size)) + branch = "restore_branch" + holder = restored.get_holder(variable_type.__name__) + storage = holder._disk_storage if on_disk else holder._memory_storage + storage.put(np.array([42.0]), period, branch) + stored_period = holder._storage_period(period) + key = (variable_type.__name__, branch, stored_period) + restored._user_input_keys.add(key) + tracked.get_holder(variable_type.__name__).set_input(period, [42.0], branch) + assert restored._user_input_keys == tracked._user_input_keys + + for simulation in (restored, tracked): + current = simulation.get_holder(variable_type.__name__) + current.delete_arrays(None, branch) + current.put_in_cache(np.array([99.0]), stored_period, branch, derived=True) + assert key not in simulation._user_input_keys + simulation.apply_reform(NoopReform) + assert current._get_array_from_storage(stored_period, branch) is None + + assert restored._user_input_keys == tracked._user_input_keys From 42c6541310689b247274968a77dc154dd21f5cc5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 20:46:46 -0400 Subject: [PATCH 10/12] Use portable compound-period disk filenames --- .../portable-disk-period-filenames.fixed.md | 1 + .../data_storage/on_disk_storage.py | 35 +++- tests/core/test_disk_period_filenames.py | 167 ++++++++++++++++++ 3 files changed, 199 insertions(+), 4 deletions(-) create mode 100644 changelog.d/portable-disk-period-filenames.fixed.md create mode 100644 tests/core/test_disk_period_filenames.py diff --git a/changelog.d/portable-disk-period-filenames.fixed.md b/changelog.d/portable-disk-period-filenames.fixed.md new file mode 100644 index 000000000..94eb07b65 --- /dev/null +++ b/changelog.d/portable-disk-period-filenames.fixed.md @@ -0,0 +1 @@ +Fix compound period disk storage on Windows while preserving legacy filename restores and clone isolation. diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 3c5dbbb3f..d4123a7f5 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -51,6 +51,11 @@ def _epoch() -> tuple: class OnDiskStorage: """ Low-level class responsible for storing and retrieving calculated vectors on disk + + Compound periods contain colons, which Windows forbids in filenames. + Their physical filenames use semicolons in the period suffix instead; + logical storage keys keep the original period strings. ``restore`` reads + both these portable filenames and legacy filenames containing colons. """ # Defaults for ``__del__``, which also runs on an instance whose @@ -313,13 +318,19 @@ def _path_to_write(self, filename: str) -> str: own = self._own_paths.get(filename) if own is not None: return own - path = os.path.join(self.storage_dir, filename) + ".npy" + # A period string can contain colons, which Windows forbids in + # filenames. Period strings never contain semicolons or underscores, + # so translating only the final suffix is reversible and leaves + # branch names untouched. + branch, separator, period_string = filename.rpartition("_") + portable_filename = f"{branch}{separator}{period_string.replace(':', ';')}" + path = os.path.join(self.storage_dir, portable_filename) + ".npy" if path not in self._family_files and not self._writes_only_new_files(): return path directory = os.path.join(self.storage_dir, REPLACEMENTS_DIR) os.makedirs(directory, exist_ok=True) descriptor, path = tempfile.mkstemp( - prefix=f"{filename}.", suffix=".npy", dir=directory + prefix=f"{portable_filename}.", suffix=".npy", dir=directory ) os.close(descriptor) return path @@ -363,13 +374,21 @@ def restore(self) -> None: """Read back the values stored in this storage's directory: for each key, the file named for it (not those in ``REPLACEMENTS_DIR``). - This storage writes over those files, as over those it wrote, until + This storage writes over portable files, as over those it wrote, until it is cloned or copied, or this process forks: from then on it writes a new file instead, since the clone, copy or forked process reads them. A storage made separately for the directory, the one that wrote them say, still writes over them, as before. So may another storage in this one's family, when ``restore`` runs after it was cloned and this process has not forked since. + + The period suffix of new filenames uses semicolons for the colons + in compound period strings. Legacy filenames containing colons are + still accepted. If both versions of a key exist after updating a + legacy value, the portable filename holds the current value and wins + regardless of directory listing order. Readers still indexed to the + legacy file keep its old value until they restore again. Branch names + are unchanged. """ self._files = files = {} # Files read back from a directory carry no derived marks. @@ -380,7 +399,15 @@ def restore(self) -> None: continue path = os.path.join(self.storage_dir, filename) filename_core = filename.rsplit(".", 1)[0] - files[filename_core] = path + branch, separator, period_string = filename_core.rpartition("_") + key = filename_core + if separator: + key = f"{branch}{separator}{period_string.replace(';', ':')}" + # Updating a legacy value may leave its old colon filename + # beside the portable file. Prefer the latter, independent of + # directory listing order, without parsing arbitrary basenames. + if key not in files or key != filename_core: + files[key] = path self._restored_paths.update(files.values()) def __del__(self, _rmtree=shutil.rmtree, _getpid=os.getpid) -> None: diff --git a/tests/core/test_disk_period_filenames.py b/tests/core/test_disk_period_filenames.py new file mode 100644 index 000000000..5c68eeb5a --- /dev/null +++ b/tests/core/test_disk_period_filenames.py @@ -0,0 +1,167 @@ +"""Disk filenames remain portable without changing logical period keys.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.data_storage import InMemoryStorage, OnDiskStorage + + +COMPOUND_PERIODS = [ + periods.period("day:2025-03-17:2"), + periods.period("month:2025-03:2"), + periods.period("year:2025:2"), + periods.period("year:2025-03"), +] +EARLY_DAYS = [ + periods.Period((periods.DAY, periods.instant((year, 3, 17)), 2)) + for year in [1, 99, 999, 1000] +] +WINDOWS_RESERVED = set('<>:"/\\|?*') + + +@pytest.fixture +def portable_writes(monkeypatch): + """Emulate Windows filename rejection on every CI operating system.""" + save = np.save + + def checked_save(file, *args, **kwargs): + name = Path(file).name + assert not WINDOWS_RESERVED.intersection(name), name + return save(file, *args, **kwargs) + + monkeypatch.setattr(np, "save", checked_save) + + +@pytest.mark.parametrize("period", COMPOUND_PERIODS + EARLY_DAYS, ids=str) +def test_compound_period_has_a_portable_disk_filename( + tmp_path, portable_writes, period +): + storage = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + storage.put(np.array([42.0]), period) + + np.testing.assert_array_equal(storage.get(period), [42.0]) + assert storage.has(period) + assert set(storage._files) == {f"default_{period}"} + restored = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + restored.restore() + np.testing.assert_array_equal(restored.get(period), [42.0]) + assert set(restored._files) == {f"default_{period}"} + + +@pytest.mark.parametrize("period", COMPOUND_PERIODS, ids=str) +def test_compound_period_clones_use_portable_replacement_paths( + tmp_path, portable_writes, period +): + storage = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + storage.put(np.array([1.0]), period) + clone = storage.clone() + + storage.put(np.array([2.0]), period) + clone.put(np.array([3.0]), period) + clone.delete(period) + clone.put(np.array([4.0]), period) + + np.testing.assert_array_equal(storage.get(period), [2.0]) + np.testing.assert_array_equal(clone.get(period), [4.0]) + files = list(tmp_path.rglob("*.npy")) + assert len(files) == 3 + assert all(not WINDOWS_RESERVED.intersection(file.name) for file in files) + restored = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + restored.restore() + np.testing.assert_array_equal(restored.get(period), [1.0]) + + +@pytest.mark.parametrize("period", COMPOUND_PERIODS, ids=str) +@pytest.mark.parametrize("branch", ["default", "policy;percent%3A"]) +def test_disk_and_memory_keep_the_same_values_and_logical_periods( + tmp_path, portable_writes, period, branch +): + """Differential: physical encoding preserves values, keys and isolation.""" + disk = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + memory = InMemoryStorage(is_eternal=False) + simple = periods.period("2025-01") + for storage in [disk, memory]: + storage.put(np.array([11.0, 12.0]), simple, branch) + storage.put(np.array([21.0, 22.0]), period, branch, derived=True) + + assert set(disk.get_known_branch_periods()) == set( + memory.get_known_branch_periods() + ) + np.testing.assert_array_equal(disk.get(period, branch), memory.get(period, branch)) + assert disk.is_derived(period, branch) == memory.is_derived(period, branch) + clones = [storage.clone() for storage in [disk, memory]] + for clone in clones: + clone.put(np.array([31.0, 32.0]), period, branch) + np.testing.assert_array_equal( + clones[0].get(period, branch), clones[1].get(period, branch) + ) + np.testing.assert_array_equal(disk.get(period, branch), memory.get(period, branch)) + + restored = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + restored.restore() + assert set(restored.get_known_branch_periods()) == set( + memory.get_known_branch_periods() + ) + for known in [simple, period]: + np.testing.assert_array_equal( + restored.get(known, branch), memory.get(known, branch) + ) + + +@pytest.mark.skipif( + os.name == "nt", reason="Windows cannot create legacy colon filenames" +) +def test_restore_still_reads_legacy_compound_period_filename(tmp_path): + period = periods.period("day:2025-03-17:2") + branch = "legacy;branch_with_underscore" + np.save(tmp_path / f"{branch}_{period}.npy", np.array([42.0])) + restored = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + + restored.restore() + + np.testing.assert_array_equal(restored.get(period, branch), [42.0]) + assert set(restored._files) == {f"{branch}_{period}"} + + +@pytest.mark.parametrize("name", ["notes", "notes;percent%3A"]) +def test_restore_tolerates_basenames_without_a_key_separator(tmp_path, name): + np.save(tmp_path / f"{name}.npy", np.array([42.0])) + restored = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + + restored.restore() + + assert set(restored._files) == {name} + np.testing.assert_array_equal(restored._decode_file(restored._files[name]), [42.0]) + + +@pytest.mark.skipif( + os.name == "nt", reason="Windows cannot create legacy colon filenames" +) +@pytest.mark.parametrize("reverse_listing", [False, True]) +def test_portable_write_after_legacy_restore_wins_on_restore( + tmp_path, monkeypatch, reverse_listing +): + """An old file left beside its portable replacement cannot mask updates.""" + period = periods.period("month:2025-03:2") + legacy_name = f"default_{period}.npy" + np.save(tmp_path / legacy_name, np.array([1.0])) + storage = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + storage.restore() + storage.put(np.array([2.0]), period) + names = os.listdir(tmp_path) + assert len(names) == 2 + monkeypatch.setattr( + os, "listdir", lambda directory: sorted(names, reverse=reverse_listing) + ) + restored = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + + restored.restore() + + np.testing.assert_array_equal(restored.get(period), [2.0]) + assert set(restored._files) == {f"default_{period}"} From a4f384d77ed0607b6e74e5d269dcbf09c979c50d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 20:52:39 -0400 Subject: [PATCH 11/12] Replay only surviving supplied input tiers after reforms --- policyengine_core/holders/holder.py | 4 +- policyengine_core/simulations/simulation.py | 31 +++-- tests/core/test_input_tier_invalidation.py | 129 ++++++++++++++++++++ 3 files changed, 156 insertions(+), 8 deletions(-) create mode 100644 tests/core/test_input_tier_invalidation.py diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 4500a118c..a46cf4f61 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -564,7 +564,9 @@ def _set( stored_period = ( self._storage_period(period) - if is_input or hasattr(self, "_user_input_storage") + if is_input + or hasattr(self, "_user_input_storage") + or (not derived and getattr(simulation, "_user_input_keys", None)) else None ) if stored_period is not None: diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 2ebd8c378..55db03bf4 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -437,11 +437,11 @@ def _invalidate_all_caches(self) -> None: Called after ``apply_reform`` and any other operation that changes the tax-benefit system underneath an already-calculated simulation. - Every (variable, branch, period) that was populated via - ``set_input`` is preserved — those are source data, not stale - formula output — so a structural reform applied after dataset - load doesn't silently discard the dataset. Everything else - (formula outputs, cached short-path results, on-disk caches) is + Still-present supplied inputs are preserved from the tiers that + hold them, so a structural reform applied after dataset load + doesn't silently discard the dataset. A cache write replacing an + input does not make that cached value a supplied input. Everything + else (formula outputs, cached short-path results, on-disk caches) is wiped so the next ``calculate`` recomputes under the new tax-benefit system. """ @@ -455,7 +455,16 @@ def _invalidate_all_caches(self) -> None: user_input_keys = getattr(self, "_user_input_keys", None) or set() for variable_name, branch_name, period in user_input_keys: holder = self.get_holder(variable_name) - stored_value = holder._memory_storage.get(period, branch_name) + holder._seed_input_storage(period, branch_name) + locations = holder._user_input_storage[(branch_name, str(period))][1] + # Carry-over marks also include external cache writes. Replay + # only tiers that still hold the recorded, supplied input. + stored_value = ( + holder._memory_storage.get(period, branch_name) + if "memory" in locations + and not holder._memory_storage.is_derived(period, branch_name) + else None + ) if stored_value is not None: preserved.append( PreservedUserInput( @@ -467,7 +476,11 @@ def _invalidate_all_caches(self) -> None: ) ) continue - if holder._disk_storage is not None: + if ( + "disk" in locations + and holder._disk_storage is not None + and not holder._disk_storage.is_derived(period, branch_name) + ): disk_period = ( periods.period(periods.ETERNITY) if holder._disk_storage.is_eternal @@ -488,6 +501,10 @@ def _invalidate_all_caches(self) -> None: disk_enum=holder._disk_storage._enums.get(disk_file), ) ) + self._user_input_keys = { + (user_input.variable_name, user_input.branch_name, user_input.period) + for user_input in preserved + } # Iterate only over holders that already exist on each population — # lazy-creating a holder for every variable in the tax-benefit # system (thousands in policyengine-us) inflated the cost of diff --git a/tests/core/test_input_tier_invalidation.py b/tests/core/test_input_tier_invalidation.py new file mode 100644 index 000000000..f241ec6de --- /dev/null +++ b/tests/core/test_input_tier_invalidation.py @@ -0,0 +1,129 @@ +"""Reform replay selects surviving input tiers, rather than cached values.""" + +import numpy as np +import pytest + +from tests.core.test_user_input_keys import ( + JANUARY, + _key, + _simulation, + _store_on_disk, +) +from tests.core.test_user_input_keys_soundness import NoopReform + + +@pytest.mark.parametrize("registered", [False, True]) +def test_reform_keeps_disk_input_instead_of_unrecorded_memory_cache(registered): + simulation = _simulation() + _store_on_disk(simulation, "rent") + if registered: + simulation.get_holder("rent")._disk_storage.put(np.array([500.0]), JANUARY) + simulation._user_input_keys.add(_key("rent", JANUARY)) + else: + simulation.set_input("rent", JANUARY, [500.0]) + simulation.memory_config.max_memory_occupation_pc = 101 + holder = simulation.get_holder("rent") + holder.put_in_cache(np.array([700.0]), JANUARY, derived=False) + + simulation.apply_reform(NoopReform) + + assert holder._memory_storage.get(JANUARY) is None + np.testing.assert_array_equal(holder._disk_storage.get(JANUARY), [500.0]) + assert simulation.to_input_dataframe()["rent__2025-01"].tolist() == [500.0] + + +@pytest.mark.parametrize("registered", [False, True]) +def test_reform_drops_record_when_cache_replaced_the_only_input(registered): + simulation = _simulation() + holder = simulation.get_holder("rent") + if registered: + holder._memory_storage.put(np.array([500.0]), JANUARY) + simulation._user_input_keys.add(_key("rent", JANUARY)) + else: + simulation.set_input("rent", JANUARY, [500.0]) + holder.put_in_cache(np.array([700.0]), JANUARY, derived=False) + + simulation.apply_reform(NoopReform) + + assert holder.get_array(JANUARY) is None + assert _key("rent", JANUARY) not in simulation._user_input_keys + assert "rent__2025-01" not in simulation.to_input_dataframe().columns + + +try: + from hypothesis import example, given, settings, strategies as st +except ImportError: + + @pytest.mark.skip(reason="Hypothesis is only installed with the dev extra") + def test_reform_tier_selection_matches_fresh_inputs(): + pass + +else: + + @settings(max_examples=60, deadline=None) + @example( + disk_input=True, + memory_input=False, + memory_cache=True, + clone=False, + registered=True, + disk_value=500, + memory_value=700, + ) + @given( + disk_input=st.booleans(), + memory_input=st.booleans(), + memory_cache=st.booleans(), + clone=st.booleans(), + registered=st.booleans(), + disk_value=st.integers(0, 1_000), + memory_value=st.integers(0, 1_000), + ) + def test_reform_tier_selection_matches_fresh_inputs( + disk_input, + memory_input, + memory_cache, + clone, + registered, + disk_value, + memory_value, + ): + simulation = _simulation() + _store_on_disk(simulation, "rent") + if disk_input: + if registered: + simulation.get_holder("rent")._disk_storage.put( + np.array([disk_value], dtype=np.float32), JANUARY + ) + simulation._user_input_keys.add(_key("rent", JANUARY)) + else: + simulation.set_input("rent", JANUARY, [disk_value]) + simulation.memory_config.max_memory_occupation_pc = 101 + if memory_input: + simulation.set_input("rent", JANUARY, [memory_value]) + if memory_cache: + simulation.get_holder("rent").put_in_cache( + np.array([float(memory_value)]), JANUARY, derived=False + ) + if clone: + simulation = simulation.clone() + + expected = ( + memory_value + if memory_input and not memory_cache + else disk_value + if disk_input + else None + ) + fresh = _simulation() + if expected is not None: + fresh.set_input("rent", JANUARY, [expected]) + + for _ in range(2): + simulation.apply_reform(NoopReform) + assert simulation._user_input_keys == fresh._user_input_keys + assert simulation.to_input_dataframe().equals(fresh.to_input_dataframe()) + np.testing.assert_array_equal( + simulation.calculate("rent", JANUARY), + fresh.calculate("rent", JANUARY), + ) From 5a69ca8d835484fc582a6c638c231da7c9af2001 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 21:23:37 -0400 Subject: [PATCH 12/12] Check live and restored input provenance against independent references --- tests/core/test_restore_input_provenance.py | 50 ++++++ .../test_restore_input_registry_property.py | 162 ++++++++++++------ 2 files changed, 159 insertions(+), 53 deletions(-) create mode 100644 tests/core/test_restore_input_provenance.py diff --git a/tests/core/test_restore_input_provenance.py b/tests/core/test_restore_input_provenance.py new file mode 100644 index 000000000..3d59366a5 --- /dev/null +++ b/tests/core/test_restore_input_provenance.py @@ -0,0 +1,50 @@ +"""Dump restoration promotes non-derived caches without changing live provenance.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, +) +from tests.fixtures.uprated_inputs import NoOp, build_simulation, build_system + + +@pytest.mark.parametrize("cached_value", [False, True], ids=["same", "changed"]) +def test_restored_cache_is_an_input_while_live_reform_drops_it(tmp_path, cached_value): + """Identical keys and storage marks do not imply identical input provenance.""" + period = periods.period("2012") + key = ("eligible", "default", period) + simulation = build_simulation( + build_system(), [("eligible", period, [False, False])] + ) + holder = simulation.get_holder("eligible") + cached = np.array([cached_value, cached_value], dtype=bool) + holder.put_in_cache(cached, period, derived=False) + + directory = str(tmp_path / "dump") + dump_simulation(simulation, directory) + restored = restore_simulation(directory, simulation.tax_benefit_system) + restored_holder = restored.get_holder("eligible") + + for stored_holder in (holder, restored_holder): + assert not stored_holder.is_derived(period) + np.testing.assert_array_equal(stored_holder.get_array(period), cached) + + for _ in range(2): + simulation.apply_reform(NoOp) + restored.apply_reform(NoOp) + + assert simulation._user_input_keys == set() + assert holder.get_array(period) is None + assert restored._user_input_keys == {key} + np.testing.assert_array_equal(restored_holder.get_array(period), cached) + assert not restored_holder.is_derived(period) + + assert "eligible__2012" not in simulation.to_input_dataframe().columns + assert ( + restored.to_input_dataframe()["eligible__2012"].tolist() == cached.tolist() + ) diff --git a/tests/core/test_restore_input_registry_property.py b/tests/core/test_restore_input_registry_property.py index b9fb60eeb..5905c7453 100644 --- a/tests/core/test_restore_input_registry_property.py +++ b/tests/core/test_restore_input_registry_property.py @@ -1,4 +1,4 @@ -"""Property: a restored simulation keeps and drops what the dumped one would. +"""Property: restores preserve values and reforms keep their own leaf inputs. For random inputs (including inputs replaced after calculating, inputs split by a ``set_input`` helper, ETERNITY inputs set for a year, and inputs set on @@ -12,20 +12,22 @@ derived. 2. **Same answers.** Calculating any sequence of requests on the restored simulation gives what the dumped simulation gives. -3. **Reforms keep the same values.** When the dumped simulation's input - record agrees with its derived marks (it records exactly the stored - values not marked derived), the restored simulation records the same - inputs, and after ``apply_reform`` with a reform that changes nothing, - the restored simulation and the dumped simulation store the same values - and give the same answers to the same requests, byte for byte. So does a - new simulation given only the explicit inputs, when the sequence contains - only input writes, calculations and branch input writes. - -The record and the marks disagree only after a value is cached without -``derived=True`` (an input to carry-over that ``set_input`` did not record) -or after an input is deleted (``delete_arrays`` leaves its record, which -policyengine-core#561 drops). The restored simulation then follows the marks, -by property 1. +3. **Reforms keep the expected inputs.** Successful leaf writes made by + explicit input operations, and not subsequently overwritten by cache or + calculation writes or deleted, are exactly what the dumped simulation + keeps after a reform that changes nothing. The restored simulation keeps + exactly the dumped values not marked derived: restoring promotes these + values to new explicit inputs. These two independent references are + checked even when they differ. When they agree, the simulations also + store the same values and give the same answers after reform, byte for + byte. So does a new simulation given only the explicit inputs, when the + sequence contains only input writes, calculations and branch input writes. + +An external cache write is not a supplied input, even when it replaces a +supplied value with identical bytes. Its non-derived mark makes it available +to carry-over and uprating, and a dump restores it as a new recorded input. +The leaf-write reference observes storage writes by operation kind; it does +not consult the input registry or the supplied-input tier metadata. The new simulation is left out after caching, deletion or a reform, since it is given only the explicit inputs. Period-splitting helpers ignore @@ -37,6 +39,7 @@ from __future__ import annotations import tempfile +from unittest.mock import patch import numpy as np import pytest @@ -46,6 +49,7 @@ st = hypothesis.strategies from policyengine_core import periods +from policyengine_core.data_storage import InMemoryStorage from policyengine_core.enums import EnumArray from policyengine_core.tools.simulation_dumper import ( dump_simulation, @@ -175,43 +179,79 @@ def _input_keys(simulation): steps=st.lists(_step, max_size=12), requests=st.lists(_request, max_size=10), ) +@hypothesis.example( + steps=[ + ("input", ("eligible", "2012", [False, False])), + ("cache", ("eligible", "2012", [False, False])), + ], + requests=[], +) def test_restored_simulation_keeps_and_drops_what_the_dumped_one_would( system, steps, requests ): dumped = build_simulation(system) inputs = [] - for kind, step in steps: - if kind == "input": - # Some inputs are refused (an annual input twice for a variable - # split into months): the new simulation gets the same refusal. - inputs.append(step) - _result(lambda: dumped.set_input(*step)) - elif kind == "branch_input": - # Stays on the branch: the dumped simulation's own values and - # inputs are what a new simulation is given. - _result(lambda: dumped.get_branch("reform").set_input(*step)) - elif kind == "cache": - name, period, values = step - holder = dumped.get_holder(name) - _result( - lambda: holder.put_in_cache( - holder._to_array(values), periods.period(period) + supplied = {} + kind = None + original_put = InMemoryStorage.put + + def observe_put(storage, value, period, branch_name="default", derived=False): + result = original_put(storage, value, period, branch_name, derived=derived) + if branch_name != "default" or kind not in ("input", "cache", "calculate"): + return result + # Observe successful writes to this simulation's leaf storage, not + # the registry/provenance helpers that cache invalidation uses. + for population in dumped.populations.values(): + for name, holder in population._holders.items(): + if holder._memory_storage is storage: + stored_period = periods.ETERNITY if storage.is_eternal else period + key = (name, str(periods.period(stored_period))) + if kind == "input" and not derived: + supplied[key] = (value.dtype.str, value.tobytes()) + else: + supplied.pop(key, None) + return result + return result + + with patch.object(InMemoryStorage, "put", observe_put): + for kind, step in steps: + if kind == "input": + # Observe individual helper writes, including partial writes + # before a refused input, without claiming skipped slots. + inputs.append(step) + _result(lambda: dumped.set_input(*step)) + elif kind == "branch_input": + # Stays on the branch: the dumped simulation's own values and + # inputs are what a new simulation is given. + _result(lambda: dumped.get_branch("reform").set_input(*step)) + elif kind == "cache": + name, period, values = step + holder = dumped.get_holder(name) + _result( + lambda: holder.put_in_cache( + holder._to_array(values), periods.period(period) + ) ) - ) - elif kind == "delete": - _result(lambda: dumped.delete_arrays(*step)) - elif kind == "reform": - dumped.apply_reform(NoOp) - else: - _result(lambda: dumped.calculate(*step)) - - stored_keys = {(name, "default", period) for name, period in _stored(dumped)} - marked_inputs = { - (name, "default", period) - for (name, period), derived in _derived_marks(dumped).items() - if not derived + elif kind == "delete": + _result(lambda: dumped.delete_arrays(*step)) + remaining = _stored(dumped) + supplied = { + key: value for key, value in supplied.items() if key in remaining + } + elif kind == "reform": + dumped.apply_reform(NoOp) + assert _stored(dumped) == supplied + else: + _result(lambda: dumped.calculate(*step)) + + dumped_values = _stored(dumped) + stored_keys = {(name, "default", period) for name, period in dumped_values} + derived_marks = _derived_marks(dumped) + restored_inputs = { + key: value for key, value in dumped_values.items() if not derived_marks[key] } - records_agree = _input_keys(dumped) & stored_keys == marked_inputs + marked_inputs = {(name, "default", period) for name, period in restored_inputs} + references_agree = supplied == restored_inputs with tempfile.TemporaryDirectory(prefix="core-restore-") as directory: dump_simulation(dumped, directory) @@ -221,7 +261,7 @@ def test_restored_simulation_keeps_and_drops_what_the_dumped_one_would( assert _stored(restored) == _stored(dumped) assert _derived_marks(restored) == _derived_marks(dumped) assert _input_keys(restored) == marked_inputs - if records_agree: + if references_agree: assert _input_keys(restored) == _input_keys(dumped) & stored_keys # 2. Same answers. @@ -230,19 +270,32 @@ def test_restored_simulation_keeps_and_drops_what_the_dumped_one_would( lambda: dumped.calculate(*request) ), request - if not records_agree: - return - - # 3. Reforms keep the same values. + # 3. Every case checks both reform contracts, including cache overwrites. fresh = build_simulation(system) for step in inputs: _result(lambda: fresh.set_input(*step)) compare_with_fresh = all( kind in ("input", "calculate", "branch_input") for kind, _ in steps ) - for simulation in (dumped, restored, fresh): + # A separate fresh simulation receives the serialized leaf snapshot; + # no restore/provenance helper defines this reference's input set. + snapshot_reference = build_simulation(system) + for (name, period), (dtype, data) in restored_inputs.items(): + values = np.frombuffer(data, dtype=np.dtype(dtype)).copy() + snapshot_reference.get_holder(name)._set( + periods.period(period), values, is_input=True + ) + for simulation in (dumped, restored, fresh, snapshot_reference): simulation.apply_reform(NoOp) - assert _stored(restored) == _stored(dumped) + assert _stored(dumped) == supplied + assert _stored(restored) == restored_inputs + assert _input_keys(dumped) == { + (name, "default", period) for name, period in supplied + } + assert _input_keys(restored) == marked_inputs + assert _stored(restored) == _stored(snapshot_reference) + if references_agree: + assert _stored(restored) == _stored(dumped) if compare_with_fresh: assert _stored(restored) == _stored(fresh) for request in requests: @@ -252,8 +305,11 @@ def test_restored_simulation_keeps_and_drops_what_the_dumped_one_would( ("dumped", dumped), ("restored", restored), ("fresh", fresh), + ("snapshot", snapshot_reference), ) } - assert answers["restored"] == answers["dumped"], (request, answers) + assert answers["restored"] == answers["snapshot"], (request, answers) + if references_agree: + assert answers["restored"] == answers["dumped"], (request, answers) if compare_with_fresh: assert answers["restored"] == answers["fresh"], (request, answers)