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