diff --git a/changelog.d/fix-uprating-branch-visibility.fixed.md b/changelog.d/fix-uprating-branch-visibility.fixed.md new file mode 100644 index 000000000..a2c5c9e58 --- /dev/null +++ b/changelog.d/fix-uprating-branch-visibility.fixed.md @@ -0,0 +1 @@ +Scope known periods and dumps to the branch's readable values while preserving derived marks, leave root uprating and carry-over fallback defaults uncached through period conversions and even after branch mutations converge so later inputs take effect, return inputs supplied during conversion on the first call, parse underscore-containing disk branch names safely, delete contained disk periods consistently with memory, reject ambiguous or unparseable storage keys before writing, and warn and skip malformed restored filenames. diff --git a/docs/python_api/data_storage.md b/docs/python_api/data_storage.md index cfaf5ba4d..096ad8dfc 100644 --- a/docs/python_api/data_storage.md +++ b/docs/python_api/data_storage.md @@ -31,6 +31,20 @@ writing when its existing values must remain isolated. :show-inheritance: ``` +## Storage keys + +Both storage backends reject branch names containing `:` and month or year +periods anchored after the first day of a month. They also reject periods whose +string form cannot be parsed back, such as the year `0999`, before storing a +value. Existing aliases, including twelve-month periods serialized as one year, +remain supported. Disk restoration warns and skips filenames without a branch +separator and a canonical period string. + +```{eval-rst} +.. automodule:: policyengine_core.data_storage.storage_keys + :members: +``` + ## TemporaryStorageDirectory A caller-supplied `simulation._data_storage_dir` is created on first use if it diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 8a103f6fc..eec1ae1df 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -237,7 +237,15 @@ depend on the input is calculated again) but not less, within these limits: input. Input helpers still remove overlapping calculated values of the variable they write. A formula's input for its own period is still returned as described above, and results calculated while an input is being set do not - enter the fast cache. + enter the fast cache. When uprating or automatic carry-over finds no source + input, the root returns the fallback default without caching it, including + when repeated branch mutations converge to a fixed point. Period conversion + preserves this rule: a divided, summed or stock-selected value containing + such a fallback also stays uncached. If a formula supplies an input for the + requested period during conversion, that input is returned on the first + call, before deciding whether to cache the converted result. Setting an + earlier input can therefore change the next calculation of the requested period. + Branches can cache these defaults because later inputs invalidate them. - **Existing child branches keep their values.** An input set on a branch does not reach branches already created from it. - **Values kept elsewhere stay.** A value one simulation calculated from @@ -276,9 +284,9 @@ say what each value was calculated from, nor what was read without being kept, so the restored simulation records that values from that number on may depend on anything: an input set for any variable on a branch of it drops all of them. A dump written before calculated values were recorded is restored with every -value as an input, as before, so such values never drop. With disk storage, a branch whose -name contains `_` cannot be dumped yet: `OnDiskStorage.get_known_periods` -splits its keys on every `_` (as on master). +value as an input, as before, so such values never drop. Disk storage splits +keys at the last underscore, so branch names containing `_` also round-trip +through dumps and restoration. Two related behaviours: a branch stops reading macro-cache files once an input is set on it, even one it already reads, since they are keyed by branch name and diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 423d19b9a..bc60c8fe3 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -4,6 +4,7 @@ from numpy.typing import ArrayLike from policyengine_core import periods +from policyengine_core.data_storage.storage_keys import check_storable from policyengine_core.data_storage.store_history import ( _period_from_storage_key, advance_sequence_past, @@ -289,20 +290,8 @@ def put( period = periods.period(period) # Keys embed f"{branch_name}:{period}", so reject inputs whose key - # would be ambiguous to parse back: branch names containing the - # separator, and month/year periods anchored mid-month (their string - # form drops the day). See policyengine-core#526 for the structured- - # key refactor that will lift these restrictions. - if ":" in branch_name: - raise ValueError( - f"Branch name {branch_name!r} may not contain ':' " - "(see policyengine-core#526)." - ) - if period.unit in (periods.YEAR, periods.MONTH) and period.start.day != 1: - raise ValueError( - f"Cannot cache period {period} anchored mid-month: its " - "string form is lossy (see policyengine-core#526)." - ) + # would be ambiguous to parse back, as ``OnDiskStorage.put`` does. + check_storable(branch_name, period) key = f"{branch_name}:{period}" self._arrays[key] = value self._stop_sharing(key) diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 9cac6e2c7..13c899f91 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -1,6 +1,7 @@ import os import shutil import tempfile +import warnings import weakref from typing import Dict, List, Optional, Set, Tuple @@ -17,6 +18,7 @@ directory_containing, path_key, ) +from policyengine_core.data_storage.storage_keys import check_storable from policyengine_core.enums import EnumArray from policyengine_core.periods import Period @@ -54,6 +56,41 @@ def _epoch() -> tuple: return os.getpid(), _forks +def _split_key(key: str) -> tuple: + """Split a ``f"{branch_name}_{period}"`` file key into its two parts. + + Branch names often contain ``_`` (policyengine-us uses ``no_salt`` and + ``mtr_for_adult_1``) but a period's string form never does, so the key + splits on its last ``_``. + """ + branch_name, period = key.rsplit("_", 1) + return branch_name, period + + +def _is_key(key: str) -> bool: + """Whether ``key`` ends with a canonical period after the branch separator. + + Legacy branch names and eternal-storage files retain their restore + behavior; this check only validates the period's spelling. + """ + _, separator, period = key.rpartition("_") + if not separator: + return False + try: + return str(periods.period(period)) == period + except ValueError: + return False + + +def _is_within(key: str, branch_name: str, period: Period) -> bool: + """Whether file key ``key`` stores a value of ``branch_name`` for a + period that ``period`` contains.""" + key_branch_name, key_period = _split_key(key) + return key_branch_name == branch_name and period.contains( + periods.period(key_period) + ) + + class OnDiskStorage: """ Low-level class responsible for storing and retrieving calculated vectors on disk @@ -303,6 +340,11 @@ def put( if self.is_eternal: period = periods.period(periods.ETERNITY) period = periods.period(period) + # The key must split back into this branch name and period (see + # ``get_known_branch_periods`` and ``delete``): reject what + # ``InMemoryStorage.put`` rejects. A month anchored mid-month, say, + # would be stored under the key of the calendar month. + check_storable(branch_name, period) filename = f"{branch_name}_{period}" path = self._path_to_write(filename) @@ -417,12 +459,13 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: if period is None: # Only wipe files belonging to the requested branch (previously # this wiped every branch regardless of ``branch_name`` — same - # class of bug as C2 in InMemoryStorage). - branch_prefix = f"{branch_name}_" + # class of bug as C2 in InMemoryStorage). Compare the parsed + # branch name, not a prefix: deleting ``pre_tcja`` must not + # also wipe ``pre_tcja_ctc``. self._files = { period_item: value for period_item, value in self._files.items() - if not period_item.startswith(branch_prefix) + if _split_key(period_item)[0] != branch_name } self._forget_dropped_keys() return @@ -431,27 +474,34 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: period = periods.period(periods.ETERNITY) period = periods.period(period) + # Delete every period of the branch that ``period`` contains, as + # ``InMemoryStorage.delete`` does and ``Holder.delete_arrays`` + # documents: deleting ``2025`` also deletes ``2025-01``. Previously + # only the file keyed exactly ``period`` was deleted. if period is not None: self._files = { period_item: value for period_item, value in self._files.items() - if not period_item == f"{branch_name}_{period}" + if not _is_within(period_item, branch_name, period) } self._forget_dropped_keys() def get_known_periods(self) -> list: - return list([periods.period(x.split("_")[1]) for x in self._files.keys()]) + return [period for _, period in self.get_known_branch_periods()] def get_known_branch_periods(self) -> list: return [ (branch_name, periods.period(period)) - for branch_name, period in map(lambda x: x.split("_"), self._files.keys()) + for branch_name, period in map(_split_key, self._files.keys()) ] 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``). + A ``.npy`` file whose name lacks a branch separator and a canonical + period string is left out, with a warning. + 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 @@ -471,6 +521,7 @@ def restore(self) -> None: self._files = files = {} # Files read back from a directory carry no derived marks or numbers. self._derived = set() + not_keys = [] self._sequence_numbers = {} # Restore self._files from content of storage_dir. for filename in os.listdir(self.storage_dir): @@ -482,12 +533,25 @@ def restore(self) -> None: key = filename_core if separator: key = f"{branch}{separator}{period_string.replace(';', ':')}" + if not _is_key(key): + # Reading it back would put a key in ``_files`` that names no + # branch and period, which ``delete`` and + # ``get_known_periods`` could not parse. + not_keys.append(filename) + continue # 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. + # directory listing order. if key not in files or key != filename_core: files[key] = path self._restored_paths.update(files.values()) + if not_keys: + warnings.warn( + f"Not restoring {', '.join(sorted(not_keys))} in " + f"{self.storage_dir}: a stored value's file is named for a " + "branch name, '_' and a period.", + stacklevel=2, + ) def __del__(self, _rmtree=shutil.rmtree, _getpid=os.getpid) -> None: # (The defaults keep the two functions reachable while the diff --git a/policyengine_core/data_storage/storage_keys.py b/policyengine_core/data_storage/storage_keys.py new file mode 100644 index 000000000..d17401b88 --- /dev/null +++ b/policyengine_core/data_storage/storage_keys.py @@ -0,0 +1,43 @@ +"""What a storage key can hold. + +Both storages key a value by its branch name and its period's string form: +``f"{branch_name}:{period}"`` in memory, ``f"{branch_name}_{period}"`` on disk +(the name of the value's file). ``check_storable`` enforces the same branch +and period serialization restrictions in both storages (``Holder._set`` +chooses between them by how much memory is in use). +""" + +from policyengine_core import periods +from policyengine_core.periods import Period + + +def check_storable(branch_name: str, period: Period) -> None: + """Raise ``ValueError`` for an ambiguous branch or period serialization. + + Reject a branch name containing ``:``, which separates the parts of an + in-memory key, and a year or month anchored mid-month, whose string form + drops the day (``month:2025-03-15`` is ``2025-03``), so that its key would + be the key of another period. Reject unparseable period strings too, + including years below 1000 whose string form loses leading zeroes. + Supported aliases, such as twelve months serialized as one year, retain + their existing meaning. See policyengine-core#526 for the structured-key + refactor. + """ + if ":" in branch_name: + raise ValueError( + f"Branch name {branch_name!r} may not contain ':' " + "(see policyengine-core#526)." + ) + if period.unit in (periods.YEAR, periods.MONTH) and period.start.day != 1: + raise ValueError( + f"Cannot cache period {period} anchored mid-month (it starts on " + f"{period.start}): its string form is lossy (see " + "policyengine-core#526)." + ) + try: + periods.period(str(period)) + except ValueError as error: + raise ValueError( + f"Cannot cache period {period!r}: its string form cannot be parsed " + "(see policyengine-core#526)." + ) from error diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 4b3c117fb..2d2ee195b 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -538,11 +538,25 @@ def get_memory_usage(self) -> dict: return usage - def get_known_periods(self) -> List[Period]: + def get_known_periods(self, branch_name: str = None) -> List[Period]: """ Get the list of periods the variable value is known for. - """ + With ``branch_name``, list each period ``get_array(period, + branch_name)`` reads a stored value for, once: those stored under + that branch, its ``parent_branch`` ancestors or ``default`` (see + ``_readable_branches``). Without it, list the period of every value + stored, under any branch; that branch may not read some of them. + """ + if branch_name is not None: + readable = set(self._readable_branches(branch_name)) + return list( + dict.fromkeys( + period + for stored_branch, period in self.get_known_branch_periods() + if stored_branch in readable + ) + ) return list(self._memory_storage.get_known_periods()) + list( (self._disk_storage.get_known_periods() if self._disk_storage else []) ) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index de5824c59..d9bb9e8a7 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -283,6 +283,7 @@ class _Frame: "stale_parent", "unobserved_activity", "attempt_cache", + "uncached_fallbacks", ) def __init__(self, simulation: "Simulation", stale_parent: bool = False): @@ -307,6 +308,9 @@ def restart(self) -> None: # Restricted children share this dictionary only until this attempt # restarts. Foreign results never enter persistent caches this way. self.attempt_cache = {} + # Cache eligibility belongs to the requested value, not to every + # formula sharing this frame. A retry may find a real source input. + self.uncached_fallbacks = set() # Retain only creations whose original registration still names # the same object. Saved or replaced snapshots keep their identities. self.carried_branches = _registered_creations( @@ -1487,6 +1491,8 @@ def calculate( transition = frame.transition_signature(result) if transition is not None and transition == previous_transition: frame.settle_reads() + if (variable_name, period) in frame.uncached_fallbacks: + break # This run repeated the complete input transition and its # result. Preserve holder input/blacklist rules and the # existing fast-cache guard when keeping the fixed point. @@ -1657,6 +1663,8 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if variable_name not in self.tax_benefit_system.variables: raise ValueError(f"Variable {variable_name} does not exist.") input_state = self._calculation_start() + frames = _calculation_frames.get() + self._set_fallback_cache_status(variable_name, period, uncached=False) # Lets a custom ``set_input`` handler tell whether it calculated. self._calculations_started = self._calculations_started + 1 population = self.get_variable_population(variable_name) @@ -1682,7 +1690,6 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if cached_array is not None: return cached_array - frames = _calculation_frames.get() if frames and frames[-1].simulation is self: cached_array = frames[-1].recall(variable_name, period) if cached_array is not None: @@ -1741,7 +1748,13 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if variable.definition_period == MONTH and period.unit == YEAR: if variable.quantity_type == QuantityType.STOCK: contained_months = period.get_subperiods(MONTH) - values = self._calculate(variable_name, contained_months[-1]) + source = contained_months[-1] + values = self._calculate(variable_name, source) + self._set_fallback_cache_status( + variable_name, + period, + uncached=self._fallback_is_uncached(variable_name, source), + ) else: values = self.calculate_add(variable_name, period) alternate_period_handling = True @@ -1749,11 +1762,29 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: alternate_period_handling = True if variable.quantity_type == QuantityType.STOCK: values = self._calculate(variable_name, period.this_year) + self._set_fallback_cache_status( + variable_name, + period, + uncached=self._fallback_is_uncached( + variable_name, period.this_year + ), + ) else: values = self.calculate_divide(variable_name, period) if alternate_period_handling: - if is_cache_available and self._may_keep(input_state): + # A formula for the source period can supply the requested + # period itself, including during STOCK selection. That input + # wins even when the converted fallback must stay uncached. + stored_input = self._input_set_meanwhile(holder, period, input_state[2]) + if stored_input is not None: + values = stored_input + self._set_fallback_cache_status(variable_name, period, uncached=False) + if ( + is_cache_available + and not self._fallback_is_uncached(variable_name, period) + and self._may_keep(input_state) + ): smc.set_cache_value(cache_path, values) return values @@ -1791,20 +1822,25 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: finally: _formula_simulation.reset(token) - # If no result, use the default value and cache it + # If no result, carry/uprate an input or use the default value. if array is None: - if variable.uprating is not None or ( + cross_period_fallback = variable.uprating is not None or ( self.tax_benefit_system.auto_carry_over_input_variables and variable.calculate_output is None - ): + ) + if cross_period_fallback: # The value is uprated or carried over from another # period, or defaults for lack of one: an input set later # for another period of the variable can change it. self._get_store_history().record_derived( variable_name, next_sequence_number() ) - # Check if the variable has a previously defined value - known_periods = holder.get_known_periods() + # Check if the variable has a previously defined value. + # Only periods this branch can read count: a value stored + # only under a branch this one cannot read must not change + # what it calculates or caches (below, a later stored period + # leaves the default uncached). + known_periods = holder.get_known_periods(self.branch_name) earlier_known_periods = [ known_period for known_period in known_periods @@ -1923,10 +1959,17 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # whether a value is stored sees. (The uprating path # above skips derived periods, so it would not uprate # from one.) - return holder.default_array() - else: - array = holder.default_array() - else: + return self._uncached_fallback(holder, period, input_state) + if array is None: + if ( + cross_period_fallback + and getattr(self, "parent_branch", None) is None + ): + # Root inputs do not invalidate dependent values. + # Leave this fallback out of every cache so a later + # input for an earlier period takes effect. An input + # set by the formula for this period still wins. + return self._uncached_fallback(holder, period, input_state) array = holder.default_array() if variable.defined_for is not None: @@ -1984,6 +2027,37 @@ def _calculation_start(self) -> Tuple[int, int, int]: and inputs were set here (``_inputs_set``), and a sequence number.""" return self._input_epoch, self._inputs_set, next_sequence_number() + def _uncached_fallback( + self, holder: Holder, period: Period, input_state: Tuple[int, int, int] + ) -> ArrayLike: + """Return a default that must stay uncached even after convergence.""" + stored_input = self._input_set_meanwhile(holder, period, input_state[2]) + if stored_input is not None: + return stored_input + self._set_fallback_cache_status(holder.variable.name, period, uncached=True) + return holder.default_array() + + def _set_fallback_cache_status( + self, variable_name: str, period: Period, *, uncached: bool + ) -> None: + # A conversion helper or a stale caller can have its own frame. Keep + # the returned value's eligibility visible after a nested frame exits, + # including when a retry replaces a fallback with a real result. + key = (variable_name, period) + for frame in _calculation_frames.get(): + if frame.simulation is self: + if uncached: + frame.uncached_fallbacks.add(key) + else: + frame.uncached_fallbacks.discard(key) + + def _fallback_is_uncached(self, variable_name: str, period: Period) -> bool: + key = (variable_name, period) + return any( + frame.simulation is self and key in frame.uncached_fallbacks + for frame in _calculation_frames.get() + ) + def _input_set_meanwhile( self, holder: Holder, period: Period, started_at: int ) -> Optional[ArrayLike]: @@ -2091,9 +2165,15 @@ def calculate_add( sub_periods = list(period.get_subperiods(variable.definition_period)) def total(): - return sum( - self.calculate(variable_name, sub_period) for sub_period in sub_periods - ) + result = 0 + uncached = False + for sub_period in sub_periods: + result = result + self.calculate(variable_name, sub_period) + # Preserve each returned term's status: a later term can + # recalculate an earlier period without changing this sum. + uncached |= self._fallback_is_uncached(variable_name, sub_period) + self._set_fallback_cache_status(variable_name, period, uncached=uncached) + return result # As in ``calculate``: only an outermost sum runs again, and while it # sums it is in flight, so its terms do not run again on their own @@ -2104,6 +2184,7 @@ def total(): try: with frame if frame is not None else nullcontext(): input_state = self._calculation_start() + started_at = input_state[2] result = total() previous_transition = None for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): @@ -2119,7 +2200,9 @@ def total(): frame.restart() input_state = self._calculation_start() result = total() - return self._cache_option_result(variable, period, result, input_state) + return self._cache_option_result( + variable, period, result, input_state, started_at=started_at + ) finally: self._calculations_in_flight -= 1 @@ -2155,6 +2238,13 @@ def calculate_divide( input_state = self._calculation_start() computation_period = period.this_year result = self.calculate(variable_name, period=computation_period) / 12.0 + self._set_fallback_cache_status( + variable_name, + period, + uncached=self._fallback_is_uncached( + variable_name, computation_period + ), + ) return self._cache_option_result(variable, period, result, input_state) elif period.unit == periods.YEAR: return self.calculate(variable_name, period) @@ -2171,6 +2261,8 @@ def _cache_option_result( period: Period, result: ArrayLike, input_state: Tuple[int, int, int], + *, + started_at: Optional[int] = None, ) -> ArrayLike: """Cache an ADD or DIVIDE result at ``period`` if a plain read would return it. @@ -2206,17 +2298,30 @@ def _cache_option_result( and export; it can name a disk input hidden by a derived memory value, and raw holder inputs need not be registered there. Accepted cache writes never register the aggregate as a supplied input. + + ``started_at`` preserves the full sum's input window across retries; + ``input_state`` controls cache eligibility for the final attempt. """ - if variable.quantity_type == QuantityType.STOCK: + holder = self.get_holder(variable.name) + # Resolve formula-time inputs before any cache rejection: returning + # an uncached fallback (or a result of another dtype) must not hide + # an input supplied for the requested period during conversion. + stored_input = self._input_set_meanwhile( + holder, period, input_state[2] if started_at is None else started_at + ) + if stored_input is not None: + self._set_fallback_cache_status(variable.name, period, uncached=False) + return stored_input + if variable.quantity_type == QuantityType.STOCK or self._fallback_is_uncached( + variable.name, period + ): return result routed = (variable.definition_period == MONTH and period.unit == YEAR) or ( variable.definition_period == YEAR and period.unit == MONTH ) if not routed or np.asarray(result).dtype != variable.dtype: return result - return self._cache_result( - self.get_holder(variable.name), result, period, input_state - ) + return self._cache_result(holder, result, period, input_state) def calculate_output(self, variable_name: str, period: Period = None) -> ArrayLike: """ diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index d351bb5d0..e89089830 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -137,9 +137,13 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): def _dump_holder(holder, directory): disk_storage = holder.create_disk_storage(directory, preserve=True) + # Dump what the simulation reads: on a branch, its own value, else its + # nearest ancestor's or the default one. A value stored only under a + # branch the simulation cannot read is left out: read here, it would be + # ``None``, saved as an array ``restore_simulation`` cannot load. branch_name = holder.simulation.branch_name derived_periods = set() - for period in dict.fromkeys(holder.get_known_periods()): + for period in holder.get_known_periods(branch_name): # What the simulation itself reads: on a branch, its own value, else # its nearest ancestor's or the default one. stored_on = holder._branch_storing(period, branch_name) diff --git a/tests/core/test_branch_helper_invalidation.py b/tests/core/test_branch_helper_invalidation.py index ba66b7847..deca3f205 100644 --- a/tests/core/test_branch_helper_invalidation.py +++ b/tests/core/test_branch_helper_invalidation.py @@ -78,15 +78,14 @@ def build(): january = year.first_month months = [january.offset(index) for index in range(12)] parent = build() + if year_string == "0025": + with pytest.raises(ValueError, match="string form cannot be parsed"): + parent.calculate("helper_independent", year) + assert parent.get_holder("helper_independent").get_known_periods() == [] + return parent.calculate("helper_independent", year) parent.set_input("helper_flow", january, [12.0, 24.0]) parent_months = [january] - if year_string == "0025": - # Early-year fallback getters have a separate legacy parser issue. - # Supply every month so this case exercises history and input replay. - for month in months[1:]: - parent.set_input("helper_flow", month, [0.0, 0.0]) - parent_months = months parent_inputs = {("helper_flow", "default", month) for month in parent_months} np.testing.assert_array_equal(parent.calculate("helper_flow", year), [12, 24]) np.testing.assert_array_equal(parent.calculate("helper_result", year), [24, 48]) diff --git a/tests/core/test_branch_input_fixed_point.py b/tests/core/test_branch_input_fixed_point.py index b51f4964e..eb851991e 100644 --- a/tests/core/test_branch_input_fixed_point.py +++ b/tests/core/test_branch_input_fixed_point.py @@ -10,6 +10,8 @@ import contextvars import threading +from types import SimpleNamespace +from unittest.mock import Mock import numpy as np import pytest @@ -19,7 +21,7 @@ from policyengine_core.experimental import MemoryConfig from policyengine_core.simulations import SimulationBuilder import policyengine_core.simulations.simulation as simulation_module -from policyengine_core.variables import Variable +from policyengine_core.variables import QuantityType, Variable def _variable(name, definition_period, formula=None): @@ -164,6 +166,401 @@ def _input_first_result(values): return branch.calculate("fp_capped", "2020") +def _fallback_simulation(cache_mode="memory", definition_period=periods.YEAR): + """The review's branch reset, whose formula selects cross-period fallback.""" + calls = [] + + def cap(person, period): + return np.full(person.count, 1000.0) + + def work(person, period): + return np.full(person.count, 5.0) + + def carried(person, period): + calls.append(str(period)) + person("cap", period) + branch = person.simulation.get_branch("worker") + branch.delete_arrays("cap") + maximum = branch.calculate("work", period) + branch.set_input("cap", period, maximum) + # Returning None selects the default or a real input to carry. + + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = True + system.add_variables( + _variable("cap", definition_period, cap), + _variable("work", definition_period, work), + _variable("carried", definition_period, carried), + ) + simulation = SimulationBuilder().build_default_simulation(system) + if cache_mode == "disk": + simulation.memory_config = MemoryConfig(0) + for population in simulation.populations.values(): + population._holders = {} + elif cache_mode == "drop": + simulation.memory_config = MemoryConfig(1, variables_to_drop=["carried"]) + for population in simulation.populations.values(): + population._holders = {} + elif cache_mode == "blacklist": + system.cache_blacklist = ["carried"] + simulation.opt_out_cache = True + return simulation, calls + + +@pytest.mark.parametrize("cache_mode", ["memory", "disk", "drop", "blacklist"]) +@pytest.mark.parametrize("future_input", [False, True]) +def test_fixed_point_fallback_does_not_hide_later_input(cache_mode, future_input): + simulation, calls = _fallback_simulation(cache_mode) + if future_input: + simulation.set_input("carried", "2022", [999]) + + assert simulation.calculate("carried", "2020").tolist() == [0.0] + assert len(calls) == 2 + + simulation.set_input("carried", "2019", [123]) + assert simulation.calculate("carried", "2020").tolist() == [123.0] + + +def test_fixed_point_fallback_is_absent_from_every_cache(monkeypatch, tmp_path): + simulation, calls = _fallback_simulation() + simulation.dataset = SimpleNamespace(file_path=tmp_path / "dataset.h5", name="test") + cache = Mock() + cache.get_cache_path.return_value = tmp_path / "missing-cache.h5" + monkeypatch.setattr(simulation_module, "SimulationMacroCache", lambda system: cache) + monkeypatch.setattr( + simulation, "check_macro_cache", lambda variable, period: variable == "carried" + ) + + assert simulation.calculate("carried", "2020").tolist() == [0.0] + assert len(calls) == 2 + assert simulation.get_array("carried", "2020") is None + assert ("carried", periods.period("2020")) not in simulation._fast_cache + cache.set_cache_value.assert_not_called() + + simulation.set_input("carried", "2019", [123]) + assert simulation.calculate("carried", "2020").tolist() == [123.0] + assert simulation.get_array("carried", "2020").tolist() == [123.0] + cache.set_cache_value.assert_called_once() + + +@pytest.mark.parametrize("cache_mode", ["memory", "disk", "drop", "blacklist"]) +@pytest.mark.parametrize( + "definition_period,quantity_type,method,requested,source,expected", + [ + (periods.YEAR, QuantityType.FLOW, "calculate", "2020-01", "2019", 10.25), + (periods.YEAR, QuantityType.FLOW, "calculate_divide", "2020-01", "2019", 10.25), + (periods.YEAR, QuantityType.STOCK, "calculate", "2020-01", "2019", 123), + (periods.MONTH, QuantityType.FLOW, "calculate", "2020", "2019-12", 1476), + ( + periods.MONTH, + QuantityType.FLOW, + "calculate_add", + "year:2020:2", + "2019-12", + 2952, + ), + (periods.MONTH, QuantityType.STOCK, "calculate", "2020", "2019-12", 123), + ], +) +def test_converted_fixed_point_fallback_stays_out_of_every_cache( + cache_mode, + definition_period, + quantity_type, + method, + requested, + source, + expected, + monkeypatch, + tmp_path, +): + simulation, _ = _fallback_simulation(cache_mode, definition_period) + simulation.tax_benefit_system.get_variable("carried").quantity_type = quantity_type + simulation.dataset = SimpleNamespace(file_path=tmp_path / "dataset.h5", name="test") + cache = Mock() + cache.get_cache_path.return_value = tmp_path / "missing-cache.h5" + monkeypatch.setattr(simulation_module, "SimulationMacroCache", lambda system: cache) + monkeypatch.setattr( + simulation, "check_macro_cache", lambda variable, period: variable == "carried" + ) + + assert getattr(simulation, method)("carried", requested).tolist() == [0.0] + holder = simulation.get_holder("carried") + assert holder.get_known_periods(simulation.branch_name) == [] + assert not any(name == "carried" for name, _ in simulation._fast_cache) + cache.set_cache_value.assert_not_called() + assert not simulation._open_frames + + simulation.set_input("carried", source, [123]) + fresh, _ = _fallback_simulation(cache_mode, definition_period) + fresh.tax_benefit_system.get_variable("carried").quantity_type = quantity_type + fresh.set_input("carried", source, [123]) + assert getattr(fresh, method)("carried", requested).tolist() == [expected] + assert getattr(simulation, method)("carried", requested).tolist() == [expected] + + +@pytest.mark.parametrize("method", ["calculate", "calculate_add"]) +def test_sum_with_some_real_inputs_does_not_cache_its_fallback_terms(method): + simulation, _ = _fallback_simulation(definition_period=periods.MONTH) + simulation.set_input("carried", "2020-06", [12]) + assert getattr(simulation, method)("carried", "2020").tolist() == [84.0] + assert simulation.get_array("carried", "2020") is None + assert ("carried", periods.period("2020")) not in simulation._fast_cache + + simulation.set_input("carried", "2019-12", [6]) + assert getattr(simulation, method)("carried", "2020").tolist() == [114.0] + + +@pytest.mark.parametrize("method", ["calculate", "calculate_divide"]) +def test_divided_fallback_yields_to_input_supplied_during_calculation(method): + class supplied_during_conversion(Variable): + value_type = float + entity = entities.Person + label = "Input supplied while dividing a fallback" + definition_period = periods.YEAR + quantity_type = QuantityType.FLOW + set_input = None + + def formula(person, period): + person.simulation.set_input("supplied_during_conversion", "2020-02", [123]) + return None + + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = True + system.add_variable(supplied_during_conversion) + simulation = SimulationBuilder().build_default_simulation(system) + + assert getattr(simulation, method)( + "supplied_during_conversion", "2020-02" + ).tolist() == [123.0] + holder = simulation.get_holder("supplied_during_conversion") + assert holder._is_input(periods.period("2020-02"), simulation.branch_name) + assert simulation.get_array("supplied_during_conversion", "2020-02").tolist() == [ + 123.0 + ] + assert simulation.get_array("supplied_during_conversion", "2020") is None + assert ("supplied_during_conversion", periods.period("2020")) not in ( + simulation._fast_cache + ) + assert not simulation._open_frames + + +@pytest.mark.parametrize("method", ["calculate", "calculate_add"]) +def test_summed_fallback_yields_to_input_supplied_during_calculation(method): + class supplied_during_conversion(Variable): + value_type = float + entity = entities.Person + label = "Input supplied while summing fallbacks" + definition_period = periods.MONTH + quantity_type = QuantityType.FLOW + set_input = None + + def formula(person, period): + if period.start.month == 6: + person.simulation.set_input("supplied_during_conversion", "2020", [120]) + return None + + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = True + system.add_variable(supplied_during_conversion) + simulation = SimulationBuilder().build_default_simulation(system) + + assert getattr(simulation, method)( + "supplied_during_conversion", "2020" + ).tolist() == [120.0] + holder = simulation.get_holder("supplied_during_conversion") + assert holder._is_input(periods.period("2020"), simulation.branch_name) + assert simulation.get_array("supplied_during_conversion", "2020").tolist() == [ + 120.0 + ] + for month in range(1, 6): + fallback_period = periods.period(f"2020-{month:02d}") + assert ( + simulation.get_array("supplied_during_conversion", fallback_period) is None + ) + assert ("supplied_during_conversion", fallback_period) not in ( + simulation._fast_cache + ) + assert not simulation._open_frames + + +@pytest.mark.parametrize( + "definition_period,quantity_type,value_type,method,formula_result", + [ + pytest.param( + periods.YEAR, QuantityType.STOCK, float, "calculate", None, id="stock-month" + ), + pytest.param( + periods.YEAR, + QuantityType.STOCK, + float, + "calculate_divide", + None, + id="stock-divide", + ), + pytest.param( + periods.MONTH, QuantityType.STOCK, float, "calculate", 7, id="stock-year" + ), + pytest.param( + periods.YEAR, + QuantityType.FLOW, + int, + "calculate", + None, + id="int-fallback-month", + ), + pytest.param( + periods.YEAR, + QuantityType.FLOW, + int, + "calculate_divide", + None, + id="int-fallback-divide", + ), + pytest.param( + periods.YEAR, + QuantityType.FLOW, + int, + "calculate", + 84, + id="int-formula-month", + ), + pytest.param( + periods.YEAR, + QuantityType.FLOW, + int, + "calculate_divide", + 84, + id="int-formula-divide", + ), + ], +) +def test_conversion_boundaries_preserve_input_supplied_during_calculation( + definition_period, quantity_type, value_type, method, formula_result +): + requested = "2020" if definition_period == periods.MONTH else "2020-02" + supplied = 120 if definition_period == periods.MONTH else 123 + + def formula(person, period): + if definition_period == periods.YEAR or period.start.month == 12: + person.simulation.set_input( + "supplied_at_conversion_boundary", requested, [supplied] + ) + # A real result distinguishes STOCK's selected December value and + # integer DIVIDE's floating result from the newly supplied input. + if formula_result is not None: + return np.full(person.count, formula_result) + return None + + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = True + system.add_variable( + type( + "supplied_at_conversion_boundary", + (Variable,), + { + "value_type": value_type, + "entity": entities.Person, + "label": "Input supplied at a period conversion boundary", + "definition_period": definition_period, + "quantity_type": quantity_type, + "set_input": None, + "formula": formula, + }, + ) + ) + simulation = SimulationBuilder().build_default_simulation(system) + + assert getattr(simulation, method)( + "supplied_at_conversion_boundary", requested + ).tolist() == [supplied] + holder = simulation.get_holder("supplied_at_conversion_boundary") + assert holder._is_input(periods.period(requested), simulation.branch_name) + assert simulation.get_array( + "supplied_at_conversion_boundary", requested + ).tolist() == [supplied] + if formula_result is None: + assert simulation.get_array("supplied_at_conversion_boundary", "2020") is None + assert ("supplied_at_conversion_boundary", periods.period("2020")) not in ( + simulation._fast_cache + ) + assert not simulation._open_frames + + +def test_formula_using_uncached_fallback_still_keeps_its_fixed_point(): + simulation, calls = _fallback_simulation() + + def outer(person, period): + return person("carried", period) + 1 + + simulation.tax_benefit_system.add_variable( + _variable("fp_fallback_outer", periods.YEAR, outer) + ) + assert simulation.calculate("fp_fallback_outer", "2020").tolist() == [1.0] + assert len(calls) == 2 + assert simulation.get_array("carried", "2020") is None + assert simulation.get_array("fp_fallback_outer", "2020").tolist() == [1.0] + + +def test_retry_can_cache_a_formula_result_after_an_uncached_fallback(): + simulation, calls = _fallback_simulation() + original = simulation.tax_benefit_system.get_variable("carried").get_formula("2020") + + def carried(person, period): + original(person, period) + if len(calls) > 1: + return np.full(person.count, 7.0) + + simulation.tax_benefit_system.update_variable( + _variable("carried", periods.YEAR, carried) + ) + assert simulation.calculate("carried", "2020").tolist() == [7.0] + assert len(calls) == 3 + assert simulation.get_array("carried", "2020").tolist() == [7.0] + assert simulation.calculate("carried", "2020").tolist() == [7.0] + assert len(calls) == 3 + + +@pytest.mark.parametrize("method", ["calculate", "calculate_divide"]) +def test_converted_retry_can_cache_a_real_result_after_fallback(method): + simulation, calls = _fallback_simulation() + original = simulation.tax_benefit_system.get_variable("carried").get_formula("2020") + + def carried(person, period): + original(person, period) + if len(calls) > 1: + return np.full(person.count, 84.0) + + simulation.tax_benefit_system.update_variable( + _variable("carried", periods.YEAR, carried) + ) + assert getattr(simulation, method)("carried", "2020-01").tolist() == [7.0] + assert len(calls) == 3 + assert simulation.get_array("carried", "2020-01").tolist() == [7.0] + assert simulation.calculate("carried", "2020-01").tolist() == [7.0] + assert len(calls) == 3 + + +@pytest.mark.parametrize("method", ["calculate", "calculate_add"]) +def test_sum_retains_fallback_status_if_a_later_term_recalculates_its_source(method): + simulation, _ = _fallback_simulation(definition_period=periods.MONTH) + + def carried(person, period): + if period.start.month == 2: + # January's already-returned fallback remains a term of the sum, + # even though its current calculation now reads a real input. + person.simulation.set_input("carried", "2020-01", [12]) + person.simulation.calculate("carried", "2020-01") + + simulation.tax_benefit_system.update_variable( + _variable("carried", periods.MONTH, carried) + ) + assert getattr(simulation, method)("carried", "2020").tolist() == [132.0] + assert simulation.get_array("carried", "2020") is None + assert ("carried", periods.period("2020")) not in simulation._fast_cache + assert getattr(simulation, method)("carried", "2020").tolist() == [144.0] + assert simulation.get_array("carried", "2020").tolist() == [144.0] + + def _assert_settled(simulation, calls, expected, kept_in_holder=True): result = simulation.calculate("fp_outer", "2020") np.testing.assert_array_equal(result, expected) diff --git a/tests/core/test_branch_input_fixed_point_property.py b/tests/core/test_branch_input_fixed_point_property.py index 2d5358132..90e236ebc 100644 --- a/tests/core/test_branch_input_fixed_point_property.py +++ b/tests/core/test_branch_input_fixed_point_property.py @@ -3,15 +3,216 @@ import pytest pytest.importorskip("hypothesis") -from hypothesis import given, settings, strategies as st +from hypothesis import example, given, settings, strategies as st + +from policyengine_core import periods from tests.core.test_branch_input_fixed_point import ( _assert_settled, + _fallback_simulation, _input_first_result, _simulation, ) +@example( + operations=[ + ("calculate", 2020, None), + ("set", 2019, 123), + ("calculate", 2020, None), + ], + cache_mode="memory", +) +@example( + operations=[ + ("calculate", "2020-01", None), + ("set", 2019, 123), + ("calculate", "2020-01", None), + ], + cache_mode="memory", +) +@example( + operations=[ + ("set", 2018, 0), + ("calculate", "2018-01", None), + ("set", 2018, 1), + ("calculate", "2018-01", None), + ], + cache_mode="memory", +) +@example( + operations=[ + ("set", 2019, 1), + ("calculate", 2018, None), + ("calculate", "2020-01", None), + ("set", 2019, 0), + ("calculate", 2020, None), + ], + cache_mode="memory", +) +@given( + operations=st.lists( + st.one_of( + st.tuples(st.just("set"), st.integers(2018, 2022), st.integers(-500, 500)), + st.tuples( + st.just("calculate"), + st.one_of( + st.integers(2018, 2022), + st.tuples(st.integers(2018, 2022), st.integers(1, 12)).map( + lambda value: f"{value[0]}-{value[1]:02d}" + ), + ), + st.none(), + ), + ), + min_size=1, + max_size=25, + ), + cache_mode=st.sampled_from(["memory", "disk", "drop", "blacklist"]), +) +@settings(max_examples=100, deadline=None) +def test_fixed_point_fallback_never_shadows_later_input(operations, cache_mode): + simulation, _ = _fallback_simulation(cache_mode) + inputs = {} + kept = {} + for action, requested, value in operations: + period = periods.period(requested) + year = period.start.year + if action == "set": + simulation.set_input("carried", period, [value]) + inputs[year] = value + # Writing an annual input removes overlapping calculated months + # of this variable, even on the root (independent of branch + # dependency invalidation). + kept = { + known: result + for known, result in kept.items() + if known.start.year != year + } + kept[period] = value + continue + + source = max((known for known in inputs if known <= year), default=None) + # Root inputs keep previous input-derived calculations, including a + # converted month or its source year. Derive each expected value from + # explicit inputs and previously modeled values, never cache contents. + annual = kept.get(period.this_year, inputs[source] if source is not None else 0) + converted = annual / 12 if period.unit == periods.MONTH else annual + expected = kept.get(period, converted) + assert simulation.calculate("carried", period).tolist() == pytest.approx( + [expected] + ) + + # A reused worker can leave the source year fresh enough to cache; + # drop/blacklist modes can omit the converted month. Observe only + # whether caches retain the independently modeled values, validating + # every retained value before remembering its expected counterpart. + candidates = {period.this_year: annual, period: expected} + for candidate, modeled in candidates.items(): + stored = simulation.get_array("carried", candidate) + fast = simulation._fast_cache.get(("carried", candidate)) + for cached in (stored, fast): + if cached is not None: + assert cached.tolist() == pytest.approx([modeled]) + if stored is not None or fast is not None: + kept[candidate] = modeled + else: + kept.pop(candidate, None) + if source is None: + assert simulation.get_array("carried", period) is None + assert simulation.get_array("carried", period.this_year) is None + assert ("carried", period) not in simulation._fast_cache + assert ("carried", period.this_year) not in simulation._fast_cache + + +@example( + definition_period=periods.YEAR, + use_option=False, + month=2, + value=123, + supplied_during_calculation=True, + cache_mode="memory", +) +@example( + definition_period=periods.MONTH, + use_option=False, + month=6, + value=120, + supplied_during_calculation=True, + cache_mode="memory", +) +@example( + definition_period=periods.MONTH, + use_option=True, + month=6, + value=120, + supplied_during_calculation=True, + cache_mode="memory", +) +@given( + definition_period=st.sampled_from([periods.YEAR, periods.MONTH]), + use_option=st.booleans(), + month=st.integers(2, 12), + value=st.integers(-500, 500), + supplied_during_calculation=st.booleans(), + cache_mode=st.sampled_from(["memory", "disk", "drop", "blacklist"]), +) +@settings(max_examples=100, deadline=None) +def test_converted_fallback_matches_input_supplied_during_calculation( + definition_period, + use_option, + month, + value, + supplied_during_calculation, + cache_mode, +): + simulation, _ = _fallback_simulation(cache_mode, definition_period) + variable = simulation.tax_benefit_system.get_variable("carried") + original = variable.get_formula("2020") + requested = f"2020-{month:02d}" if definition_period == periods.YEAR else "2020" + supplied = False + + def carried(person, period): + nonlocal supplied + original(person, period) + if ( + supplied_during_calculation + and not supplied + and (definition_period == periods.YEAR or period.start.month == month) + ): + person.simulation.set_input("carried", requested, [value]) + # Retrying the sum must still recognize this first-attempt input + # without relying on the formula to write it again. + supplied = True + # Returning None still selects fallback; a newly supplied input for + # the requested period takes precedence over the converted result. + + # Preserve the variable shared with any holders already constructed for + # disk/drop storage, including the input handler those holders consult. + variable.formulas["2020-01-01"] = carried + variable.set_input = None + option = ( + "calculate_divide" if definition_period == periods.YEAR else "calculate_add" + ) + method = option if use_option else "calculate" + + # The model depends only on whether the formula supplies the requested + # input. In particular, it never divides that input or sums carried terms. + expected = value if supplied_during_calculation else 0 + assert getattr(simulation, method)("carried", requested).tolist() == pytest.approx( + [expected] + ) + if supplied_during_calculation: + assert simulation.get_array("carried", requested).tolist() == [value] + else: + assert ( + simulation.get_holder("carried").get_known_periods(simulation.branch_name) + == [] + ) + assert not any(name == "carried" for name, _ in simulation._fast_cache) + assert not simulation._open_frames + + @given( values=st.lists( st.floats(min_value=-500, max_value=500, allow_nan=False, width=32), diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py index a144fd78f..df1231027 100644 --- a/tests/core/test_carry_over_order.py +++ b/tests/core/test_carry_over_order.py @@ -154,6 +154,10 @@ def test_an_input_ending_after_the_last_date_drops_overlapping_defaults( holder = built.get_holder(name) requested = periods.period("9999-12-31") holder._do_not_store = cache_only + # Seed existing caches: root carry-over fallbacks no longer store defaults. + default = holder.default_array() + holder.put_in_cache(default, requested, derived=True) + built._fast_cache[(name, requested)] = default np.testing.assert_array_equal(built.calculate(name, requested), [0, 0]) assert (name, requested) in built._fast_cache assert holder._memory_storage.has(requested) == (not cache_only) @@ -259,21 +263,27 @@ def test_input_set_after_a_period_defaulted_still_carries(system): np.testing.assert_array_equal(built.calculate("carried", "2014"), [7, 8]) -def test_defaults_are_cached_as_before(system): - """With nothing stored, or only earlier calculated values, the default is - cached (as master cached the value it carried); before a later stored - period it is not (as master returned it).""" +@pytest.mark.parametrize("on_branch", [False, True], ids=["root", "branch"]) +def test_fallback_defaults_are_cached_only_on_branches(system, on_branch): + """Only branches invalidate these defaults when earlier inputs arrive. + Before a later stored period, neither root nor branch caches them.""" built = simulation(system, {}) + if on_branch: + built = built.get_branch("fallback_defaults") holder = built.get_holder("carried_monthly") for month in range(1, 13): for _ in range(2): built.calculate("carried_monthly", f"2012-{month:02d}") for month in range(1, 13): - assert holder.get_array(periods.period(f"2012-{month:02d}")) is not None + period = periods.period(f"2012-{month:02d}") + assert (holder.get_array(period, built.branch_name) is not None) == on_branch + assert (("carried_monthly", period) in built._fast_cache) == on_branch built = simulation(system, {"carried_monthly": {"2013-06": [1, 2]}}) + if on_branch: + built = built.get_branch("fallback_defaults") holder = built.get_holder("carried_monthly") np.testing.assert_array_equal(built.calculate("carried_monthly", "2012-03"), [0, 0]) - assert holder.get_array(periods.period("2012-03")) is None + assert holder.get_array(periods.period("2012-03"), built.branch_name) is None def test_branch_input_is_not_hidden_by_a_value_the_parent_carried(system): @@ -609,19 +619,21 @@ def test_a_period_stored_only_on_another_branch_does_not_carry(system): def test_marks_are_cleared_with_the_values_apply_reform_wipes(system): - built = simulation(system, {}) + built = simulation(system, {"carried": {"2009": [7, 8]}}) built.memory_config = MemoryConfig(max_memory_occupation=0) holder = built.get_holder("carried") holder._disk_storage = holder.create_disk_storage() holder._on_disk_storable = True for year in range(2010, 2020): built.calculate("carried", str(year)) + assert holder.is_derived(periods.period(year)) built.apply_reform(_noop) - assert not holder._memory_storage._inputs + assert holder.get_known_periods() == [periods.period("2009")] + assert holder.get_input_periods() == [periods.period("2009")] assert not holder._disk_storage._derived # A file written for a period whose value was derived, then read back # by rebuilding the index, is an input. - built = simulation(system, {}) + built = simulation(system, {"carried": {"2011": [7, 8]}}) built.memory_config = MemoryConfig(max_memory_occupation=0) holder = built.get_holder("carried") holder._disk_storage = holder.create_disk_storage() @@ -637,17 +649,19 @@ def test_marks_are_cleared_with_the_values_apply_reform_wipes(system): def test_marks_are_cleared_with_in_memory_values_apply_reform_wipes(system): - built = simulation(system, {}) + built = simulation(system, {"carried": {"2009": [7, 8]}}) holder = built.get_holder("carried") for year in range(2010, 2020): built.calculate("carried", str(year)) + assert holder.is_derived(periods.period(year)) built.apply_reform(_noop) - assert not holder._memory_storage._inputs + assert holder.get_known_periods() == [periods.period("2009")] + assert holder.get_input_periods() == [periods.period("2009")] def test_rebuilding_a_disk_index_reads_files_as_inputs(system): """A derived file replaced by an input, then re-indexed, is an input.""" - built = simulation(system, {}) + built = simulation(system, {"carried": {"2011": [7, 8]}}) built.memory_config = MemoryConfig(max_memory_occupation=0) holder = built.get_holder("carried") holder._disk_storage = holder.create_disk_storage() diff --git a/tests/core/test_disk_period_filenames.py b/tests/core/test_disk_period_filenames.py index 5c68eeb5a..5f9407a77 100644 --- a/tests/core/test_disk_period_filenames.py +++ b/tests/core/test_disk_period_filenames.py @@ -38,7 +38,7 @@ def checked_save(file, *args, **kwargs): monkeypatch.setattr(np, "save", checked_save) -@pytest.mark.parametrize("period", COMPOUND_PERIODS + EARLY_DAYS, ids=str) +@pytest.mark.parametrize("period", COMPOUND_PERIODS + EARLY_DAYS[-1:], ids=str) def test_compound_period_has_a_portable_disk_filename( tmp_path, portable_writes, period ): @@ -54,6 +54,19 @@ def test_compound_period_has_a_portable_disk_filename( assert set(restored._files) == {f"default_{period}"} +@pytest.mark.parametrize("period", EARLY_DAYS[:-1], ids=str) +def test_unparseable_early_period_is_rejected_before_writing( + tmp_path, portable_writes, period +): + storage = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + + with pytest.raises(ValueError, match="string form cannot be parsed"): + storage.put(np.array([42.0]), period) + + assert not storage._files + assert not list(tmp_path.iterdir()) + + @pytest.mark.parametrize("period", COMPOUND_PERIODS, ids=str) def test_compound_period_clones_use_portable_replacement_paths( tmp_path, portable_writes, period @@ -130,14 +143,15 @@ def test_restore_still_reads_legacy_compound_period_filename(tmp_path): @pytest.mark.parametrize("name", ["notes", "notes;percent%3A"]) -def test_restore_tolerates_basenames_without_a_key_separator(tmp_path, name): +def test_restore_skips_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() + with pytest.warns(UserWarning, match="Not restoring"): + restored.restore() - assert set(restored._files) == {name} - np.testing.assert_array_equal(restored._decode_file(restored._files[name]), [42.0]) + assert not restored._files + np.testing.assert_array_equal(np.load(tmp_path / f"{name}.npy"), [42.0]) @pytest.mark.skipif( diff --git a/tests/core/test_disk_storage_delete_contained_periods.py b/tests/core/test_disk_storage_delete_contained_periods.py new file mode 100644 index 000000000..8cdfc01ce --- /dev/null +++ b/tests/core/test_disk_storage_delete_contained_periods.py @@ -0,0 +1,475 @@ +"""``OnDiskStorage.delete(period, branch_name)`` deletes the periods within +``period`` (policyengine-core#564). + +``Holder.delete_arrays`` documents that deleting a period removes "all values +for any period included in period", and ``InMemoryStorage.delete`` does that +(deleting ``2025`` removes ``2025-01`` ... ``2025-12``). ``OnDiskStorage.delete`` +used to remove only the file keyed exactly ``f"{branch_name}_{period}"``, so a +disk-backed simulation kept values the caller had deleted. + +Each example runs on both backends. ``test_disk_and_memory_delete_the_same_keys`` +checks every pair of a stored and a deleted period on a small grid, and +``test_disk_storage_delete_property.py`` checks random sequences of puts and +deletes. +""" + +import os +import warnings + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.data_storage import InMemoryStorage, OnDiskStorage +from policyengine_core.experimental import MemoryConfig +from policyengine_core.simulations import SimulationBuilder + +BACKENDS = ("memory", "disk") + + +@pytest.fixture +def make_storage(tmp_path): + def make(backend, is_eternal=False): + if backend == "memory": + return InMemoryStorage(is_eternal=is_eternal) + directory = tmp_path / f"storage_{len(os.listdir(tmp_path))}" + directory.mkdir() + return OnDiskStorage( + str(directory), is_eternal=is_eternal, preserve_storage_dir=True + ) + + return make + + +def put_all(storage, keys): + for index, (branch_name, period) in enumerate(keys): + storage.put(np.array([float(index)]), periods.period(period), branch_name) + + +def known(storage): + """The stored (branch name, period string) pairs, sorted.""" + return sorted( + (branch_name, str(period)) + for branch_name, period in storage.get_known_branch_periods() + ) + + +# ----- Examples on both backends ---------------------------------------------- + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_delete_exact_period(make_storage, backend): + storage = make_storage(backend) + put_all(storage, [("default", "2025-01"), ("default", "2025-02")]) + + storage.delete(periods.period("2025-01"), "default") + + assert known(storage) == [("default", "2025-02")] + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_delete_year_deletes_its_months_and_days(make_storage, backend): + storage = make_storage(backend) + put_all( + storage, + [ + ("default", "2024-12"), + ("default", "2025"), + ("default", "2025-01"), + ("default", "2025-12"), + ("default", "2025-03-15"), + ("default", "2026-01"), + ], + ) + + storage.delete(periods.period(2025), "default") + + assert known(storage) == [("default", "2024-12"), ("default", "2026-01")] + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_delete_month_keeps_the_year_that_contains_it(make_storage, backend): + storage = make_storage(backend) + put_all( + storage, + [ + ("default", "2025"), + ("default", "2025-01"), + ("default", "2025-01-31"), + ("default", "2025-02"), + ("default", "2025-02-01"), + ], + ) + + storage.delete(periods.period("2025-01"), "default") + + assert known(storage) == [ + ("default", "2025"), + ("default", "2025-02"), + ("default", "2025-02-01"), + ] + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_delete_keeps_other_branches(make_storage, backend): + storage = make_storage(backend) + put_all( + storage, + [ + ("default", "2025-01"), + ("reform", "2025-01"), + ("reform", "2025-02"), + ], + ) + + storage.delete(periods.period(2025), "reform") + + assert known(storage) == [("default", "2025-01")] + np.testing.assert_array_equal(storage.get("2025-01", "default"), [0.0]) + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_delete_matches_whole_branch_names_containing_underscores( + make_storage, backend +): + # ``pre`` and ``pre_tcja`` are prefixes of ``pre_tcja_ctc``; the ``y`` + # branch's prefix ``y_`` starts the ``y_2019`` branch's keys. + branch_names = ("pre", "pre_tcja", "pre_tcja_ctc", "y", "y_2019", "trailing_") + storage = make_storage(backend) + put_all( + storage, + [ + (branch_name, month) + for branch_name in branch_names + for month in ("2025-01", "2025-02") + ], + ) + + storage.delete(periods.period(2025), "pre_tcja") + storage.delete(periods.period(2025), "y") + storage.delete(periods.period("2025-01"), "trailing_") + + assert known(storage) == [ + ("pre", "2025-01"), + ("pre", "2025-02"), + ("pre_tcja_ctc", "2025-01"), + ("pre_tcja_ctc", "2025-02"), + ("trailing_", "2025-02"), + ("y_2019", "2025-01"), + ("y_2019", "2025-02"), + ] + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_eternal_storage_deletes_only_the_branch_eternity_key(make_storage, backend): + storage = make_storage(backend, is_eternal=True) + put_all(storage, [("default", "2025"), ("no_salt", "2025"), ("no", "2025")]) + + # Any period deletes the branch's single ETERNITY value. + storage.delete(periods.period("2025-01"), "no_salt") + + assert known(storage) == [("default", "ETERNITY"), ("no", "ETERNITY")] + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_deleting_eternity_deletes_every_period_of_the_branch(make_storage, backend): + storage = make_storage(backend) + put_all( + storage, + [("default", "2024"), ("default", "2025-01"), ("reform", "2025-01")], + ) + + storage.delete(periods.period(periods.ETERNITY), "default") + + assert known(storage) == [("reform", "2025-01")] + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_delete_multi_unit_periods(make_storage, backend): + storage = make_storage(backend) + put_all( + storage, + [ + ("default", "month:2025-01:3"), + ("default", "year:2025-03"), + ("default", "2025-02"), + ("default", "2025-04"), + ], + ) + + # A quarter contains its months, but a month does not contain the + # quarter, and 2025 does not contain the year starting in March 2025. + storage.delete(periods.period("month:2025-01:3"), "default") + assert known(storage) == [("default", "2025-04"), ("default", "year:2025-03")] + + storage.delete(periods.period("2025-04"), "default") + assert known(storage) == [("default", "year:2025-03")] + + storage.delete(periods.period(2025), "default") + assert known(storage) == [("default", "year:2025-03")] + + storage.delete(periods.period("year:2025:2"), "default") + assert known(storage) == [] + + +# ----- Disk and memory agree on every pair of periods ------------------------- + +PERIODS = [ + "2024", + "2025", + "2025-01", + "2025-02", + "2025-12", + "2025-01-01", + "2025-01-31", + "2025-02-01", + "2026-01", + "ETERNITY", + "month:2025-01:2", + "year:2025-03", + "year:2024:2", +] +BRANCH_NAMES = ["default", "pre_tcja", "pre_tcja_ctc"] + + +@pytest.mark.parametrize("deleted_branch_name", BRANCH_NAMES) +@pytest.mark.parametrize("deleted_period", PERIODS) +def test_disk_and_memory_delete_the_same_keys( + make_storage, deleted_period, deleted_branch_name +): + keys = [(branch_name, period) for branch_name in BRANCH_NAMES for period in PERIODS] + memory, disk = make_storage("memory"), make_storage("disk") + put_all(memory, keys) + put_all(disk, keys) + deleted_period = periods.period(deleted_period) + + memory.delete(deleted_period, deleted_branch_name) + disk.delete(deleted_period, deleted_branch_name) + + assert known(disk) == known(memory) + # And both keep exactly the keys the period does not contain. + assert known(memory) == sorted( + (branch_name, str(periods.period(period))) + for branch_name, period in keys + if branch_name != deleted_branch_name + or not deleted_period.contains(periods.period(period)) + ) + for branch_name, period in known(memory): + np.testing.assert_array_equal( + disk.get(period, branch_name), memory.get(period, branch_name) + ) + + +# ----- Through the simulation (the #564 reproduction) ------------------------- + + +def simulation_with_rent(on_disk): + simulation = SimulationBuilder().build_from_entities( + CountryTaxBenefitSystem(), + {"persons": {"a": {}}, "households": {"h": {"parents": ["a"]}}}, + ) + if on_disk: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + holder = simulation.get_holder("rent") + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + simulation.set_input("rent", "2025-01", [500.0]) + simulation.set_input("rent", "2025-02", [600.0]) + simulation.set_input("rent", "2026-01", [700.0]) + return simulation + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_simulation_delete_arrays_deletes_the_months_of_a_year(on_disk): + simulation = simulation_with_rent(on_disk) + holder = simulation.get_holder("rent") + if on_disk: + assert holder._memory_storage.get_known_periods() == [] + assert len(holder._disk_storage.get_known_periods()) == 3 + + simulation.delete_arrays("rent", "2025") + + assert holder.get_known_periods() == [periods.period("2026-01")] + assert holder.get_array("2025-01") is None + assert holder.get_array("2025-02") is None + np.testing.assert_array_equal(holder.get_array("2026-01"), [700.0]) + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_holder_delete_arrays_deletes_the_months_of_a_year(on_disk): + simulation = simulation_with_rent(on_disk) + holder = simulation.get_holder("rent") + + holder.delete_arrays("2025", "default") + + assert holder.get_known_periods() == [periods.period("2026-01")] + + +# ----- Keys no storage can hold (#565 review, finding 1) ---------------------- + +# A year or month anchored mid-month has a string form that drops the day +# (``month:2025-03-15`` is ``2025-03``), and ``:`` separates the parts of an +# in-memory key. In-memory storage rejected both (policyengine-core#526); disk +# storage stored a mid-month month under the calendar month's key, so reading +# the calendar month returned it, and deleting the period it was put for, +# which does not contain that calendar month, left it. +UNSTORABLE_KEYS = [ + ("default", periods.Period((periods.MONTH, periods.Instant((2025, 3, 15)), 1))), + ("default", periods.Period((periods.MONTH, periods.Instant((2025, 3, 2)), 2))), + ("default", periods.Period((periods.YEAR, periods.Instant((2025, 3, 15)), 1))), + ("no_salt", periods.Period((periods.YEAR, periods.Instant((2025, 1, 2)), 2))), + ("my:reform", periods.period(2025)), +] + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize( + "branch_name, period", + UNSTORABLE_KEYS, + ids=[ + f"{branch}-{period.unit}-{period.start}" for branch, period in UNSTORABLE_KEYS + ], +) +def test_put_rejects_keys_no_storage_can_hold( + make_storage, backend, branch_name, period +): + storage = make_storage(backend) + put_all(storage, [("default", "2025-03")]) + + with pytest.raises(ValueError, match="policyengine-core#526"): + storage.put(np.array([9.0]), period, branch_name) + + assert known(storage) == [("default", "2025-03")] + np.testing.assert_array_equal(storage.get("2025-03", "default"), [0.0]) + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_mid_month_input_is_rejected_whichever_storage_it_goes_to(on_disk): + # The review's reproduction: on disk, this input was accepted, then + # ``delete_arrays`` of its own period left it ([500.] instead of None). + simulation = simulation_with_rent(on_disk) + holder = simulation.get_holder("rent") + + with pytest.raises(ValueError, match="anchored mid-month"): + simulation.set_input("rent", "month:2025-03-15", [500.0]) + simulation.delete_arrays("rent", "month:2025-03-15") + + assert holder.get_array("month:2025-03-15") is None + assert holder.get_array("2025-03") is None + assert sorted(map(str, holder.get_known_periods())) == [ + "2025-01", + "2025-02", + "2026-01", + ] + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("period", [999, "0999", "0999-03", "0999-03-15"]) +def test_put_rejects_unparseable_period_without_changing_values( + make_storage, backend, period +): + # The full #565 review also found put(periods.period(999)): its "999" + # key cannot be parsed by containment deletion or restored from disk. + storage = make_storage(backend) + storage.put(np.array([500.0]), "2025-03", "no_salt", derived=True) + period = periods.period(period) + + with pytest.raises(ValueError, match="string form cannot be parsed"): + storage.put(np.array([9.0]), period, "no_salt") + + assert storage.get(period, "no_salt") is None + assert known(storage) == [("no_salt", "2025-03")] + assert storage.is_derived("2025-03", "no_salt") + np.testing.assert_array_equal(storage.get("2025-03", "no_salt"), [500.0]) + storage.delete(periods.period(2025), "no_salt") + assert known(storage) == [] + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("month", [1, 3]) +def test_twelve_month_alias_remains_storable(make_storage, backend, month): + period = periods.Period((periods.MONTH, periods.Instant((2025, month, 1)), 12)) + storage = make_storage(backend) + + storage.put(np.array([500.0]), period, "no_salt", derived=True) + + alias = periods.period(str(period)) + assert alias.unit == periods.YEAR + assert alias.start == period.start + assert alias.stop == period.stop + assert storage.get_known_branch_periods() == [("no_salt", alias)] + assert storage.is_derived(alias, "no_salt") + np.testing.assert_array_equal(storage.get(alias, "no_salt"), [500.0]) + storage.delete(period, "no_salt") + assert known(storage) == [] + + +# ----- Files restored from a directory (#565 review, finding 2) --------------- + +# ``restore`` read back every ``.npy`` file in the directory. One whose name +# is not a key ``put`` writes (no ``_``, no period after the last ``_``, or a +# period not in its string form) names no value, and deleting any period of +# any branch then raised parsing it, leaving the value it was to delete. +# The raw filenames below bypass storage's semicolon encoder. Only literal +# colon filenames remain excluded on Windows, which cannot create them. +NOT_KEYS = [ + "stray.npy", + "default_garbage.npy", + "default_.npy", + "_.npy", + "default_2025-3.npy", + "default_2025-03-1.npy", + "default_999.npy", +] +if os.name != "nt": + NOT_KEYS += ["default_month:2025-01:12.npy", "no_salt_year:2025:1.npy"] + + +def test_restore_leaves_out_files_that_are_not_keys(tmp_path): + storage = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + put_all(storage, [("default", "2025-01"), ("no_salt", "2025-02")]) + for name in NOT_KEYS: + np.save(str(tmp_path / name), np.array([9.0])) + + restored = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + with pytest.warns(UserWarning, match="Not restoring") as caught: + restored.restore() + + assert known(restored) == [("default", "2025-01"), ("no_salt", "2025-02")] + (warning,) = caught + for name in NOT_KEYS: + assert name in str(warning.message) + + # Deleting works as before the files were added. + restored.delete(periods.period(2025), "default") + assert restored.get("2025-01", "default") is None + assert known(restored) == [("no_salt", "2025-02")] + restored.delete(None, "no_salt") + assert known(restored) == [] + + +def test_restore_reads_back_every_key_without_warning(tmp_path): + keys = [ + ("default", "2025"), + ("no_salt", "2025-03"), + ("y_2019", "2025-03-05"), + ("trailing_", "ETERNITY"), + ("pre_tcja", "year:2025-03"), + ("default", "month:2025-01:3"), + ] + storage = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + put_all(storage, keys) + + restored = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + with warnings.catch_warnings(): + warnings.simplefilter("error") + restored.restore() + + assert known(restored) == known(storage) + for branch_name, period in keys: + np.testing.assert_array_equal( + restored.get(period, branch_name), storage.get(period, branch_name) + ) diff --git a/tests/core/test_disk_storage_delete_property.py b/tests/core/test_disk_storage_delete_property.py new file mode 100644 index 000000000..06a884f58 --- /dev/null +++ b/tests/core/test_disk_storage_delete_property.py @@ -0,0 +1,282 @@ +"""Disk and memory storage keep the same values through any puts and deletes. + +Random sequences of ``put``, ``delete(period, branch_name)`` and +``delete(None, branch_name)`` run on an ``InMemoryStorage`` and an +``OnDiskStorage``, and on a reference model of the documented behaviour: +``put`` stores a value unless no key can hold its branch name and period (a +branch name containing ``:``, or a period whose string form names a period +starting on another day: a year or month anchored mid-month), and then both +storages reject it, with the same error; ``delete(period, branch_name)`` +removes the branch's values for every period that ``period`` contains, +``delete(None, branch_name)`` removes all of the branch's values, and eternal +storage keeps one ``ETERNITY`` value per branch. After every step both +storages hold exactly the model's (branch, period) keys, with the values, +input/derived provenance and sequence numbers last put there. Deleting one +period must preserve all three for a derived value surviving in another. + +A second property restores a disk storage's directory, holding the files of +random keys and of names that are not keys, and checks that ``restore`` reads +back exactly the keys, warns about the other names, and that deletes then +agree with memory. + +Branch names include ones that contain ``_`` or start with another branch's +name followed by ``_``, since disk keys join branch and period with ``_``. +``test_disk_storage_delete_contained_periods.py`` pins the same behaviour +with examples. +""" + +import os +import tempfile +import warnings + +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.data_storage import InMemoryStorage, OnDiskStorage + +BRANCH_NAMES = ( + "default", + "reform", + "pre", + "pre_tcja", + "pre_tcja_ctc", + "y", + "y_2019", + "trailing_", + # No key can hold a branch name containing ``:``. + "my:reform", +) + + +@st.composite +def period_strategy(draw): + if draw(st.integers(0, 19)) == 0: + return periods.period(periods.ETERNITY) + unit = draw(st.sampled_from([periods.YEAR, periods.MONTH, periods.DAY])) + year = draw(st.one_of(st.sampled_from([999, 1000]), st.integers(2023, 2026))) + # Disk filenames encode compound periods' colons as semicolons, so these + # periods are supported on Windows too. + month = draw(st.integers(1, 12)) + size = draw(st.sampled_from([1, 2, 3, 12])) + # Both storages reject years and months anchored mid-month, whose string + # form drops the day (before writing a file, so on Windows too). + if unit == periods.DAY: + day = draw(st.integers(1, 28)) + else: + day = draw(st.one_of(st.just(1), st.integers(2, 28))) + return periods.Period((unit, periods.Instant((year, month, day)), size)) + + +def storable(branch_name, period): + """The model of which keys a storage holds: those whose branch name has no + ``:`` and whose period's string form names a period starting the same + day.""" + try: + parsed = periods.period(str(period)) + except ValueError: + return False + return ":" not in branch_name and parsed.start == periods.period(period).start + + +def put_both( + memory, disk, value, period, branch_name, *, derived=False, sequence_number=None +): + """Put ``value`` in both storages; whether both stored it. Both store it, + or both raise the same ``ValueError``, as the model says.""" + errors = [] + for storage in (memory, disk): + try: + storage.put( + np.array([float(value)]), + period, + branch_name, + derived=derived, + sequence_number=sequence_number, + ) + errors.append(None) + except ValueError as error: + errors.append(str(error)) + assert errors[0] == errors[1] + if memory.is_eternal: + period = periods.period(periods.ETERNITY) + assert (errors[0] is None) == storable(branch_name, period) + return errors[0] is None + + +put_step = st.tuples( + st.just("put"), + period_strategy(), + st.sampled_from(BRANCH_NAMES), + st.integers(0, 1_000), + st.booleans(), +) +delete_step = st.tuples( + st.just("delete"), + st.one_of(st.none(), period_strategy()), + st.sampled_from(BRANCH_NAMES), + st.none(), + st.none(), +) + + +def stored(storage): + """Map each stored (branch name, period string) to its value.""" + return { + (branch_name, str(period)): float(storage.get(period, branch_name)[0]) + for branch_name, period in storage.get_known_branch_periods() + } + + +def provenance(storage): + """Map each stored key to its derived mark and original sequence number.""" + separator = ":" if isinstance(storage, InMemoryStorage) else "_" + return { + (branch_name, str(period)): ( + storage.is_derived(period, branch_name), + storage._sequence_numbers[f"{branch_name}{separator}{period}"], + ) + for branch_name, period in storage.get_known_branch_periods() + } + + +@hypothesis.settings( + max_examples=200, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.example( + is_eternal=False, + steps=[ + ("put", periods.period(2025), "default", 7, False), + ("put", periods.period(2026), "default", 9, True), + ("delete", periods.period(2025), "default", None, None), + ], +) +@hypothesis.given( + is_eternal=st.booleans(), + steps=st.lists(st.one_of(put_step, put_step, delete_step), max_size=25), +) +def test_disk_and_memory_storage_delete_the_same_keys(is_eternal, steps): + model = {} + memory = InMemoryStorage(is_eternal=is_eternal) + with tempfile.TemporaryDirectory() as directory: + disk = OnDiskStorage( + directory, is_eternal=is_eternal, preserve_storage_dir=True + ) + for sequence_number, (action, period, branch_name, value, derived) in enumerate( + steps, start=1 + ): + if period is not None and is_eternal: + period = periods.period(periods.ETERNITY) + if action == "put": + if put_both( + memory, + disk, + value, + period, + branch_name, + derived=derived, + sequence_number=sequence_number, + ): + model[(branch_name, str(period))] = ( + float(value), + derived, + sequence_number, + ) + else: + memory.delete(period, branch_name) + disk.delete(period, branch_name) + model = { + (key_branch_name, key_period): key_value + for (key_branch_name, key_period), key_value in model.items() + if key_branch_name != branch_name + or ( + period is not None + and not period.contains(periods.period(key_period)) + ) + } + + assert ( + stored(disk) + == stored(memory) + == {key: value for key, (value, _, _) in model.items()} + ) + assert ( + provenance(disk) + == provenance(memory) + == { + key: (derived, sequence_number) + for key, (_, derived, sequence_number) in model.items() + } + ) + + +# Names of ``.npy`` files that are not keys: no ``_``, no period after the last +# ``_``, or a period not in its string form. (Not ``eternity``: on a +# case-insensitive file system that is the ``ETERNITY`` key's file.) +NON_CANONICAL_PERIODS = ["2025-3", "2025-03-1", "999", "x", ""] +if os.name != "nt": + # These are deliberately malformed raw filenames, written without the + # storage encoder. Windows cannot create literal-colon filenames. + NON_CANONICAL_PERIODS += ["year:2025", "month:2025-01:12", "day:2025-03-01"] +not_key_names = st.one_of( + st.text(alphabet="abcxyz0123456789-", min_size=1, max_size=8), + st.builds( + "{}_{}".format, + st.sampled_from([name for name in BRANCH_NAMES if ":" not in name]), + st.one_of( + st.sampled_from(NON_CANONICAL_PERIODS), + st.text(alphabet="abcxyz", min_size=1, max_size=6), + ), + ), +).map(lambda name: name + ".npy") + + +@hypothesis.settings( + max_examples=100, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given( + is_eternal=st.booleans(), + puts=st.lists(put_step, max_size=12), + not_keys=st.lists(not_key_names, max_size=4, unique=True), + deletes=st.lists(delete_step, max_size=4), +) +def test_restore_reads_back_the_keys_and_only_the_keys( + is_eternal, puts, not_keys, deletes +): + memory = InMemoryStorage(is_eternal=is_eternal) + with tempfile.TemporaryDirectory() as directory: + disk = OnDiskStorage( + directory, is_eternal=is_eternal, preserve_storage_dir=True + ) + for _, period, branch_name, value, derived in puts: + put_both(memory, disk, value, period, branch_name, derived=derived) + for name in not_keys: + np.save(os.path.join(directory, name), np.array([-1.0])) + + restored = OnDiskStorage( + directory, is_eternal=is_eternal, preserve_storage_dir=True + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + restored.restore() + + assert stored(restored) == stored(memory) + if not_keys: + (warning,) = caught + for name in not_keys: + assert name in str(warning.message) + else: + assert caught == [] + + for _, period, branch_name, _, _ in deletes: + memory.delete(period, branch_name) + restored.delete(period, branch_name) + assert stored(restored) == stored(memory) diff --git a/tests/core/test_holder_write_fast_cache.py b/tests/core/test_holder_write_fast_cache.py index f928f04bb..3c53f2fc1 100644 --- a/tests/core/test_holder_write_fast_cache.py +++ b/tests/core/test_holder_write_fast_cache.py @@ -20,6 +20,7 @@ from policyengine_core import periods from policyengine_core.experimental import MemoryConfig +from policyengine_core.variables import Variable from tests.fixtures.uprated_inputs import build_simulation, build_system @@ -208,14 +209,25 @@ def test_a_write_to_disk_storage_replaces_a_calculated_value(system): def test_a_write_replaces_an_eternity_value_requested_without_a_period(system): - # ``calculate`` with no period caches the result under ``None``. - simulation = build_simulation(system) - assert simulation.calculate("eternal_code").tolist() == [0, 0] - assert ("eternal_code", None) in simulation._fast_cache + class eternal_code_plus_one(Variable): + def formula(person, period): + # The top-level request intentionally has no period. Population + # lookups inside formulas still require an explicit one. + return person("eternal_code", periods.ETERNITY) + 1 + + local_system = system.clone() + local_system.update_variable(eternal_code_plus_one) + # Calculate a formula over a real input: supplied inputs bypass the fast + # cache, while a missing root's default must stay uncached. + simulation = build_simulation(local_system, [("eternal_code", "2012", [0, 0])]) + assert simulation.calculate("eternal_code_plus_one").tolist() == [1, 1] + assert ("eternal_code_plus_one", None) in simulation._fast_cache - simulation.get_holder("eternal_code").set_input(periods.period("2015"), [8, 9]) + simulation.get_holder("eternal_code_plus_one").set_input( + periods.period("2015"), [8, 9] + ) - assert simulation.calculate("eternal_code").tolist() == [8, 9] + assert simulation.calculate("eternal_code_plus_one").tolist() == [8, 9] def _disk_simulation(system): diff --git a/tests/core/test_known_periods_branch_visibility.py b/tests/core/test_known_periods_branch_visibility.py new file mode 100644 index 000000000..2d51e4082 --- /dev/null +++ b/tests/core/test_known_periods_branch_visibility.py @@ -0,0 +1,573 @@ +"""Known periods a branch cannot read. + +Holder storage keys embed the branch name (``"no_salt:2018"`` in memory, +``"no_salt_2018"`` on disk), and a holder can hold keys for branches other +than the simulation's own. ``Holder.get_known_periods()`` lists the periods of +every key, while ``Holder.get_array(period, branch_name)`` reads only that +branch, its ``parent_branch`` ancestors and ``default``. +``Holder.get_known_periods(branch_name)`` lists exactly the periods +``get_array`` reads for that branch, each once. Three defects came from the +gap between the two: + +* ``Simulation._calculate`` checked the unscoped list for stored periods. When + it took the latest earlier period from it, a period stored only under a + branch the simulation cannot read gave ``None``: uprating raised + ``TypeError`` (``None * factor``) and auto-carry-over cached ``NaN``. Since + both read only inputs the branch reads (``Holder.get_input_periods``), such + a period still made the simulation leave a default uncached that it would + otherwise cache. +* ``OnDiskStorage`` parsed ``f"{branch}_{period}"`` keys with + ``split("_")[1]``. A branch name containing ``_`` raised ``ValueError`` + (``no_salt`` -> period ``"salt"``) or listed the wrong period + (``y_2019`` -> ``2019``). ``delete(None, "pre_tcja")`` also wiped + ``pre_tcja_ctc``, whose key shares the prefix. +* ``dump_simulation`` read every listed period under ``default``, saving + ``None`` for a period stored only on the dumped branch, which + ``restore_simulation`` could not load, and the derived marks of the + ``default`` branch's values rather than those of the values dumped. + +``test_known_periods_branch_visibility_property.py`` checks the listing +against ``get_array`` and ``get_input_periods`` for random stored values. +""" + +import itertools +import math +import os +import tempfile +import warnings +from types import SimpleNamespace +from unittest.mock import Mock + +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.data_storage import OnDiskStorage +from policyengine_core.experimental import MemoryConfig +from policyengine_core.model_api import YEAR, Variable +from policyengine_core.parameters import ParameterNode +from policyengine_core.reforms import Reform +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.tools import simulation_dumper + +INDEX_START = 2015 +# The index grows 10% a year, so ratios between years are powers of 1.1. +INDEX = { + f"{year}-01-01": 100 * 1.1 ** (year - INDEX_START) for year in range(2015, 2023) +} + + +def growth(from_year: int, to_year: int) -> float: + return 1.1 ** (to_year - from_year) + + +def build_system(auto_carry_over: bool = False) -> CountryTaxBenefitSystem: + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = auto_carry_over + system.parameters.add_child( + "test_uprating", + ParameterNode("test_uprating", data={"index": {"values": INDEX}}), + ) + + class uprated_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Uprated yearly income" + uprating = "test_uprating.index" + + class carried_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Yearly income carried over without uprating" + + system.add_variable(uprated_income) + system.add_variable(carried_income) + return system + + +@pytest.fixture(scope="module") +def system(): + return build_system() + + +@pytest.fixture(scope="module") +def carry_over_system(): + return build_system(auto_carry_over=True) + + +def new_simulation(system): + return SimulationBuilder().build_default_simulation(system, count=1) + + +def store(simulation, variable, year, value, branch_name): + """Store ``value`` for ``year`` under ``branch_name`` in the holder of + ``simulation``, the way ``Holder.set_input`` does when a caller names a + branch other than the simulation's own.""" + simulation.get_holder(variable).set_input( + periods.period(year), np.array([value]), branch_name + ) + + +def only(array) -> float: + (value,) = array + return float(value) + + +# ----- Simulation._calculate -------------------------------------------------- + + +def test_uprating_skips_period_stored_only_on_unrelated_branch(system): + simulation = new_simulation(system) + store(simulation, "uprated_income", 2016, 1_000.0, "default") + store(simulation, "uprated_income", 2018, 5_000.0, "other") + + result = simulation.calculate("uprated_income", 2020) + + # Uprated from 2016, the latest year the default branch can read. + assert only(result) == pytest.approx(1_000 * growth(2016, 2020)) + + +def test_uprating_with_no_readable_earlier_period_uses_default(system): + simulation = new_simulation(system) + store(simulation, "uprated_income", 2018, 5_000.0, "other") + + assert only(simulation.calculate("uprated_income", 2020)) == 0 + + +def test_uprating_in_branch_skips_sibling_branch_period(system): + simulation = new_simulation(system) + simulation.set_input("uprated_income", 2016, [1_000.0]) + branch = simulation.get_branch("reform") + store(branch, "uprated_income", 2018, 5_000.0, "baseline") + + result = branch.calculate("uprated_income", 2020) + + assert only(result) == pytest.approx(1_000 * growth(2016, 2020)) + + +def test_uprating_in_nested_branch_reads_ancestor_periods(system): + simulation = new_simulation(system) + simulation.set_input("uprated_income", 2016, [1_000.0]) + itemizing = simulation.get_branch("itemizing") + itemizing.set_input("uprated_income", 2018, [5_000.0]) + no_salt = itemizing.get_branch("no_salt") + + result = no_salt.calculate("uprated_income", 2020) + + assert only(result) == pytest.approx(5_000 * growth(2018, 2020)) + + +def test_carry_over_skips_period_stored_only_on_unrelated_branch( + carry_over_system, +): + simulation = new_simulation(carry_over_system) + store(simulation, "carried_income", 2016, 1_000.0, "default") + store(simulation, "carried_income", 2018, 5_000.0, "other") + + result = simulation.calculate("carried_income", 2020) + + assert not math.isnan(only(result)) + assert only(result) == 1_000 + + +def test_carry_over_with_no_readable_period_uses_default(carry_over_system): + simulation = new_simulation(carry_over_system) + store(simulation, "carried_income", 2018, 5_000.0, "other") + + result = simulation.calculate("carried_income", 2020) + + assert not math.isnan(only(result)) + assert only(result) == 0 + + +def test_carry_over_ignores_later_period_on_unrelated_branch(carry_over_system): + simulation = new_simulation(carry_over_system) + store(simulation, "carried_income", 2016, 1_000.0, "default") + store(simulation, "carried_income", 2022, 9_000.0, "other") + + assert only(simulation.calculate("carried_income", 2020)) == 1_000 + + +def test_root_default_before_unreadable_future_input_does_not_hide_later_input( + carry_over_system, +): + simulation = new_simulation(carry_over_system) + store(simulation, "carried_income", 2022, 9_000.0, "other") + + assert only(simulation.calculate("carried_income", 2020)) == 0 + + simulation.set_input("carried_income", 2019, [123]) + + assert only(simulation.calculate("carried_income", 2020)) == 123 + + +@pytest.mark.parametrize( + "variable, auto_carry_over", + [("carried_income", True), ("uprated_income", False)], +) +def test_root_fallback_default_is_not_kept_in_any_cache( + variable, auto_carry_over, monkeypatch, tmp_path +): + import policyengine_core.simulations.simulation as simulation_module + + simulation = new_simulation(build_system(auto_carry_over)) + simulation.dataset = SimpleNamespace(file_path=tmp_path / "dataset.h5", name="test") + cache = Mock() + cache.get_cache_path.return_value = tmp_path / "missing-cache.h5" + monkeypatch.setattr(simulation_module, "SimulationMacroCache", lambda system: cache) + monkeypatch.setattr(simulation, "check_macro_cache", lambda *args: True) + + assert only(simulation.calculate(variable, 2020)) == 0 + + period = periods.period(2020) + assert simulation.get_holder(variable).get_array(period) is None + assert (variable, period) not in simulation._fast_cache + cache.set_cache_value.assert_not_called() + + simulation.set_input(variable, 2019, [123]) + expected = 123 * growth(2019, 2020) if variable == "uprated_income" else 123 + assert only(simulation.calculate(variable, 2020)) == pytest.approx(expected) + assert simulation.get_holder(variable).get_array(period) is not None + cache.set_cache_value.assert_called_once() + + +def test_root_uprating_fallback_uses_input_set_during_formula(): + class self_supplied_income(Reform): + def apply(self): + class uprated_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + + def formula(person, period): + person.simulation.set_input("uprated_income", period, [123]) + + self.update_variable(uprated_income) + + simulation = new_simulation(self_supplied_income(build_system())) + + assert only(simulation.calculate("uprated_income", 2020)) == 123 + holder = simulation.get_holder("uprated_income") + assert not holder.is_derived(periods.period(2020)) + + +# Invariant: values stored only under branches a simulation cannot read do not +# change what it calculates, nor what it stores where it reads. Checked +# exhaustively against a simulation that never had those values, for every +# combination of readable years, unreadable years and requested year below. +READABLE_YEAR_SETS = [(), (2016,), (2018,), (2016, 2018)] +UNREADABLE_YEAR_SETS = [ + years + for size in range(1, 4) + for years in itertools.combinations((2015, 2017, 2019, 2021), size) +] +REQUESTED_YEARS = (2017, 2018, 2019, 2020) + + +def _calculate_with(system, variable, readable, unreadable, requested): + """The branch's result, and the (period, derived) pairs it then reads.""" + simulation = new_simulation(system) + branch = simulation.get_branch("reform") + for year in readable: + store(branch, variable, year, 1_000.0 * (year - 2000), "reform") + for year in unreadable: + store(branch, variable, year, 7_777.0, "baseline") + result = only(branch.calculate(variable, requested)) + holder = branch.get_holder(variable) + read = sorted( + (str(period), holder.is_derived(period, "reform")) + for period in holder.get_known_periods("reform") + ) + return result, read + + +@pytest.mark.parametrize("requested", REQUESTED_YEARS) +@pytest.mark.parametrize("unreadable", UNREADABLE_YEAR_SETS) +@pytest.mark.parametrize("readable", READABLE_YEAR_SETS) +@pytest.mark.parametrize( + "variable, auto_carry_over", + [("uprated_income", False), ("carried_income", True)], +) +def test_unreadable_branch_periods_do_not_change_results( + system, + carry_over_system, + variable, + auto_carry_over, + readable, + unreadable, + requested, +): + tax_benefit_system = carry_over_system if auto_carry_over else system + + with_unreadable = _calculate_with( + tax_benefit_system, variable, readable, unreadable, requested + ) + without_unreadable = _calculate_with( + tax_benefit_system, variable, readable, (), requested + ) + + assert math.isfinite(with_unreadable[0]) + assert with_unreadable == without_unreadable + + +# ----- Holder.get_known_periods(branch_name) agrees with get_array ------------ + +# default -> a -> a_b is one lineage; c is a sibling of a, and a +# branch named "default" nested under a reads only "default" in get_array. +STORED_BRANCHES = ("default", "a", "a_b", "c") + + +def _lineage(system): + simulation = new_simulation(system) + a = simulation.get_branch("a") + return { + "default": simulation, + "a": a, + "a_b": a.get_branch("a_b"), + "c": simulation.get_branch("c"), + "default under a": a.get_branch("default"), + } + + +@pytest.mark.parametrize("on_disk", [False, True]) +@pytest.mark.parametrize( + "stored", + [ + branches + for size in range(1, len(STORED_BRANCHES) + 1) + for branches in itertools.combinations(STORED_BRANCHES, size) + ], +) +def test_known_periods_for_branch_are_exactly_the_readable_ones( + system, tmp_path, stored, on_disk +): + """For every reader, ``get_known_periods(branch_name)`` lists a period if + and only if ``get_array(period, branch_name)`` returns a value.""" + for reader_name, reader in _lineage(system).items(): + holder = reader.get_holder("uprated_income") + if on_disk: + directory = tmp_path / reader_name.replace(" ", "-") + directory.mkdir() + holder._disk_storage = holder.create_disk_storage( + str(directory), preserve=True + ) + storage = holder._disk_storage if on_disk else holder._memory_storage + # One distinct year per stored branch, so each listed period traces + # back to exactly one branch. + for offset, branch_name in enumerate(stored): + storage.put(np.array([1.0]), periods.period(2016 + offset), branch_name) + + stored_periods = {period for _, period in holder.get_known_branch_periods()} + readable = { + period + for period in stored_periods + if holder.get_array(period, reader.branch_name) is not None + } + + listed = holder.get_known_periods(reader.branch_name) + + assert set(listed) == readable, reader_name + assert len(listed) == len(set(listed)), reader_name + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_known_periods_for_branch_list_each_period_once(system, tmp_path, on_disk): + """A period stored under several branches a reader reads, or both in + memory and on disk, is listed once.""" + period = periods.period(2016) + for reader_name, reader in _lineage(system).items(): + holder = reader.get_holder("uprated_income") + directory = tmp_path / reader_name.replace(" ", "-") + directory.mkdir() + holder._disk_storage = holder.create_disk_storage(str(directory), preserve=True) + storages = [holder._memory_storage, holder._disk_storage] + for branch_name in STORED_BRANCHES: + for storage in storages if on_disk else storages[:1]: + storage.put(np.array([1.0]), period, branch_name) + + assert holder.get_known_periods(reader.branch_name) == [period], reader_name + + +def test_get_known_periods_without_branch_lists_every_branch(system): + simulation = new_simulation(system) + store(simulation, "uprated_income", 2016, 1.0, "default") + store(simulation, "uprated_income", 2018, 1.0, "other") + + holder = simulation.get_holder("uprated_income") + + assert sorted(holder.get_known_periods()) == [ + periods.period(2016), + periods.period(2018), + ] + assert holder.get_known_periods("default") == [periods.period(2016)] + + +# ----- OnDiskStorage key parsing ---------------------------------------------- + +BRANCH_NAMES = ( + "default", + "no_salt", + "mtr_for_adult_1", + "pre_tcja_ctc", + "y_2019", + "trailing_", +) +# A multi-unit or rolling-year period's string form contains ``:``, which a +# Windows file name cannot, so disk storage cannot keep those periods there. +skip_on_windows = pytest.mark.skipif( + os.name == "nt", reason="Windows file names cannot contain ':'." +) +PERIODS = ( + periods.period(2025), + periods.period("2025-03"), + periods.period("2025-03-05"), + *( + pytest.param(periods.period(period), id=period, marks=skip_on_windows) + for period in ("month:2025-01:3", "year:2025-03", "year:2024:2") + ), +) + + +@pytest.fixture +def disk_storage(tmp_path): + return OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + + +@pytest.mark.parametrize("period", PERIODS, ids=str) +@pytest.mark.parametrize("branch_name", BRANCH_NAMES) +def test_on_disk_keys_round_trip(disk_storage, branch_name, period): + disk_storage.put(np.array([3.0]), period, branch_name) + + assert disk_storage.get_known_branch_periods() == [(branch_name, period)] + assert disk_storage.get_known_periods() == [period] + np.testing.assert_array_equal(disk_storage.get(period, branch_name), [3.0]) + + # The same holds after rebuilding the index from the files on disk. + disk_storage.restore() + assert disk_storage.get_known_branch_periods() == [(branch_name, period)] + + +@pytest.mark.parametrize("branch_name", BRANCH_NAMES) +def test_on_disk_eternal_keys_round_trip(tmp_path, branch_name): + storage = OnDiskStorage(str(tmp_path), is_eternal=True, preserve_storage_dir=True) + storage.put(np.array([3.0]), periods.period(2025), branch_name) + + assert storage.get_known_branch_periods() == [ + (branch_name, periods.period(periods.ETERNITY)) + ] + + +def test_on_disk_delete_branch_leaves_branches_sharing_its_prefix(disk_storage): + period = periods.period(2025) + for branch_name in ("pre_tcja", "pre_tcja_ctc", "pre"): + disk_storage.put(np.array([1.0]), period, branch_name) + + disk_storage.delete(None, "pre_tcja") + + assert sorted(disk_storage.get_known_branch_periods()) == [ + ("pre", period), + ("pre_tcja_ctc", period), + ] + + +@pytest.mark.parametrize("branch_name", ["no_salt", "y_2019"]) +def test_disk_backed_branch_with_underscore_uprates(branch_name): + """Only public API: a variable added after ``memory_config`` is set gets + disk storage, and the branch's input is uprated from its disk key.""" + tax_benefit_system = build_system() + simulation = new_simulation(tax_benefit_system) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + + class disk_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Uprated yearly income stored on disk" + uprating = "test_uprating.index" + + tax_benefit_system.add_variable(disk_income) + branch = simulation.get_branch(branch_name) + branch.set_input("disk_income", 2018, [5_000.0]) + disk_keys = list(branch.get_holder("disk_income")._disk_storage._files) + assert disk_keys == [f"{branch_name}_2018"] + + result = branch.calculate("disk_income", 2020) + + assert only(result) == pytest.approx(5_000 * growth(2018, 2020)) + + +# ----- dump_simulation -------------------------------------------------------- + + +def test_dump_branch_saves_the_values_the_branch_reads(system): + simulation = new_simulation(system) + simulation.set_input("uprated_income", 2016, [1_000.0]) + simulation.set_input("uprated_income", 2017, [2_000.0]) + branch = simulation.get_branch("reform") + branch.set_input("uprated_income", 2017, [3_000.0]) + branch.set_input("uprated_income", 2018, [5_000.0]) + + directory = os.path.join(tempfile.mkdtemp(), "dump") + simulation_dumper.dump_simulation(branch, directory) + restored = simulation_dumper.restore_simulation(directory, system) + + holder = restored.get_holder("uprated_income") + assert sorted(holder.get_known_periods()) == [ + periods.period(2016), + periods.period(2017), + periods.period(2018), + ] + assert only(holder.get_array(2016)) == 1_000 + assert only(holder.get_array(2017)) == 3_000 + assert only(holder.get_array(2018)) == 5_000 + + +def test_restored_branch_dump_calculates_what_the_branch_calculates(system): + """The dump keeps the derived mark of each value the branch reads, so the + restored simulation uprates from the same inputs as the branch.""" + simulation = new_simulation(system) + simulation.set_input("uprated_income", 2016, [1_000.0]) + # Calculated on ``default``, then set as an input on the branch. + simulation.calculate("uprated_income", 2019) + branch = simulation.get_branch("reform") + branch.set_input("uprated_income", 2019, [3_000.0]) + # Calculated on the branch only. + branch.calculate("uprated_income", 2017) + + directory = os.path.join(tempfile.mkdtemp(), "dump") + simulation_dumper.dump_simulation(branch, directory) + restored = simulation_dumper.restore_simulation(directory, system) + + holder = restored.get_holder("uprated_income") + assert sorted( + (str(period), holder.is_derived(period)) + for period in holder.get_known_periods() + ) == [("2016", False), ("2017", True), ("2019", False)] + for year in range(2016, 2023): + assert only(restored.calculate("uprated_income", year)) == pytest.approx( + only(branch.calculate("uprated_income", year)) + ), year + assert only(restored.calculate("uprated_income", 2021)) == pytest.approx( + 3_000 * growth(2019, 2021) + ) + + +def test_restore_simulation_leaves_out_files_that_are_not_values(system, tmp_path): + simulation = new_simulation(system) + simulation.set_input("uprated_income", 2016, [1_000.0]) + directory = str(tmp_path / "dump") + simulation_dumper.dump_simulation(simulation, directory) + np.save(os.path.join(directory, "uprated_income", "stray.npy"), np.array([9.0])) + + with pytest.warns(UserWarning, match="stray.npy"): + restored = simulation_dumper.restore_simulation(directory, system) + + holder = restored.get_holder("uprated_income") + assert holder.get_known_periods() == [periods.period(2016)] + assert only(holder.get_array(2016)) == 1_000 diff --git a/tests/core/test_known_periods_branch_visibility_property.py b/tests/core/test_known_periods_branch_visibility_property.py new file mode 100644 index 000000000..e2f828ee3 --- /dev/null +++ b/tests/core/test_known_periods_branch_visibility_property.py @@ -0,0 +1,111 @@ +"""``Holder.get_known_periods(branch_name)`` agrees with ``get_array`` for any +stored values. + +Random values are stored, in memory or on disk, as inputs or derived, under +the branches of one lineage: ``default``, its branch ``a``, ``a``'s branch +``a_b``, ``default``'s other branch ``c``, and a branch named ``default`` +nested under ``a``. Each branch of the lineage, reading its own holder, sees: + +* ``get_known_periods(branch_name)`` lists each period ``get_array(period, + branch_name)`` reads a stored value for, once, and no other period: the + listing against the per-period reference, ``get_array``; +* ``get_input_periods(branch_name)`` lists the periods it lists whose value + is an input (``not is_derived(period, branch_name)``): the one pass over + the stored keys against the per-period reference. + +``test_known_periods_branch_visibility.py`` pins the same with examples. +""" + +import os +import tempfile + +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 Person +from policyengine_core.model_api import YEAR, Variable +from policyengine_core.simulations import SimulationBuilder + +STORED_BRANCHES = ("default", "a", "a_b", "c", "unrelated") +YEARS = (2015, 2016, 2017) + + +class yearly_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Yearly income" + + +@pytest.fixture(scope="module") +def system(): + system = CountryTaxBenefitSystem() + system.add_variable(yearly_income) + return system + + +def lineage(system): + simulation = SimulationBuilder().build_default_simulation(system, count=1) + a = simulation.get_branch("a") + return [ + simulation, + a, + a.get_branch("a_b"), + simulation.get_branch("c"), + a.get_branch("default"), + ] + + +stored_value = st.tuples( + st.sampled_from(STORED_BRANCHES), + st.sampled_from(YEARS), + st.booleans(), # derived + st.booleans(), # on disk +) + + +@hypothesis.settings( + max_examples=100, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given(stored=st.lists(stored_value, max_size=12)) +def test_known_periods_for_a_branch_are_the_periods_it_reads(system, stored): + with tempfile.TemporaryDirectory() as directory: + for index, reader in enumerate(lineage(system)): + holder = reader.get_holder("yearly_income") + reader_directory = os.path.join(directory, str(index)) + os.mkdir(reader_directory) + holder._disk_storage = holder.create_disk_storage( + reader_directory, preserve=True + ) + for value, (branch_name, year, derived, on_disk) in enumerate(stored): + storage = holder._disk_storage if on_disk else holder._memory_storage + storage.put( + np.array([float(value)]), + periods.period(year), + branch_name, + derived=derived, + ) + + branch_name = reader.branch_name + listed = holder.get_known_periods(branch_name) + read = [ + period + for period in map(periods.period, YEARS) + if holder.get_array(period, branch_name) is not None + ] + + assert len(listed) == len(set(listed)) + assert set(listed) == set(read) + assert set(holder.get_input_periods(branch_name)) == { + period + for period in listed + if not holder.is_derived(period, branch_name) + } diff --git a/tests/core/test_on_disk_storage_clones.py b/tests/core/test_on_disk_storage_clones.py index 911a7d3e2..cf348530c 100644 --- a/tests/core/test_on_disk_storage_clones.py +++ b/tests/core/test_on_disk_storage_clones.py @@ -62,7 +62,9 @@ def _stored(holder, period): @pytest.mark.parametrize("variable", list(HELPER_INPUTS), ids=["dispatch", "divide"]) def test_helper_input_on_a_clone_leaves_the_source_unchanged(system, variable): - parent = simulation(system, {}) + # A real uprating source makes December a stored derived value under + # both carry-over settings; a missing root input stays uncached. + parent = simulation(system, {variable: {"2011-12": [1.0, 2.0]}}) holder = _on_disk(parent, variable) calculated = parent.calculate(variable, "2012-12").copy() assert holder._disk_storage.is_derived(periods.period("2012-12")) diff --git a/tests/core/test_registered_input_deletion.py b/tests/core/test_registered_input_deletion.py index 0d1f02ab1..1fb9db1a9 100644 --- a/tests/core/test_registered_input_deletion.py +++ b/tests/core/test_registered_input_deletion.py @@ -1,4 +1,4 @@ -"""Restored input records follow deletion even for early-year storage keys.""" +"""Restored input records respect storage period validation and deletion.""" import numpy as np import pytest @@ -34,6 +34,15 @@ def test_deleted_registered_early_year_input_is_forgotten(year, variable_type, o (variable_type.definition_period, periods.instant((year, 1, 1)), 1) ) storage = holder._disk_storage if on_disk else holder._memory_storage + if year < 1000: + input_keys = simulation._user_input_keys.copy() + with pytest.raises(ValueError, match="string form cannot be parsed"): + storage.put(np.array([42.0]), period) + assert not storage.has(period) + assert not storage._sequence_numbers + assert simulation._user_input_keys == input_keys + assert not hasattr(holder, "_user_input_storage") + return storage.put(np.array([42.0]), period) key = (variable_type.__name__, "default", period) simulation._user_input_keys.add(key) @@ -65,9 +74,10 @@ def test_registered_deletion_matches_set_input_for_period_aliases( ): """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. + Unsupported early years are rejected before registration. For supported + years, deletion removes the input regardless of the stored period's unit, + size, 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) @@ -75,6 +85,22 @@ def test_registered_deletion_matches_set_input_for_period_aliases( branch = "restore_branch" holder = restored.get_holder(variable_type.__name__) storage = holder._disk_storage if on_disk else holder._memory_storage + if year < 1000: + simulations = (restored, tracked) + input_keys = [simulation._user_input_keys.copy() for simulation in simulations] + with pytest.raises(ValueError, match="string form cannot be parsed"): + storage.put(np.array([42.0]), period, branch) + with pytest.raises(ValueError, match="string form cannot be parsed"): + tracked.get_holder(variable_type.__name__).set_input(period, [42.0], branch) + for simulation, before in zip(simulations, input_keys): + current = simulation.get_holder(variable_type.__name__) + tier = current._disk_storage if on_disk else current._memory_storage + assert not tier.has(period, branch) + assert not tier._sequence_numbers + assert simulation._user_input_keys == before + assert simulation._user_input_contexts == [] + assert not hasattr(current, "_user_input_storage") + return storage.put(np.array([42.0]), period, branch) stored_period = holder._storage_period(period) key = (variable_type.__name__, branch, stored_period) diff --git a/tests/core/test_restore_input_registry.py b/tests/core/test_restore_input_registry.py index dbaa825ff..ecf429c71 100644 --- a/tests/core/test_restore_input_registry.py +++ b/tests/core/test_restore_input_registry.py @@ -96,6 +96,7 @@ def test_dump_lists_the_calculated_periods_and_marks_its_format(system, tmp_path [ ("uprated_count", "2012", [1001, 77]), ("uprated_count", "2014", [2001, 177]), + ("uprated_amount", "2012", [100.0, 50.0]), ], ) simulation.calculate("uprated_count", "2013") @@ -107,7 +108,7 @@ def calculated(variable): return sorted(path.read_text().split()) if path.exists() else [] assert calculated("uprated_count") == ["2013"] - # Calculated from the defaulted input, which is calculated too. + # Both the uprated value and the formula reading it are derived. assert calculated("doubled_amount") == ["2013"] assert calculated("uprated_amount") == ["2013"] assert (tmp_path / DUMP_FORMAT_FILE).exists() diff --git a/tests/core/test_set_input_helper_order.py b/tests/core/test_set_input_helper_order.py index b4fb8fc3b..4a1620de9 100644 --- a/tests/core/test_set_input_helper_order.py +++ b/tests/core/test_set_input_helper_order.py @@ -581,29 +581,24 @@ def test_repeated_helper_clears_cache_shadow_when_every_month_has_a_disk_input( @pytest.mark.parametrize("on_disk", [False, True]) -def test_helper_drops_overlapping_early_year_cache(on_disk): +def test_helper_rejects_early_year_input_and_cache_without_recording(on_disk): simulation = build_simulation(on_disk=on_disk) holder = simulation.get_holder("flow_m") year = periods.period("0025") - holder.put_in_cache(np.array([5.0, 6.0]), year, derived=True) - # Plain reads do not fast-cache an aggregate today; model an existing - # fast-cache entry to prove cleanup uses the early-year Period itself. - simulation._fast_cache[("flow_m", year)] = np.array([5.0, 6.0]) - fresh = build_simulation() + input_keys = simulation._user_input_keys.copy() - set_input_divide_by_period(holder, year, YEARLY_INPUT) - fresh.set_input("flow_m", year, YEARLY_INPUT) + with pytest.raises(ValueError, match="string form cannot be parsed"): + holder.put_in_cache(np.array([5.0, 6.0]), year, derived=True) + with pytest.raises(ValueError, match="string form cannot be parsed"): + set_input_divide_by_period(holder, year, YEARLY_INPUT) - assert holder.get_array(year) is None + assert holder.get_known_periods() == [] + storage = holder._disk_storage if on_disk else holder._memory_storage + assert not storage._sequence_numbers assert ("flow_m", year) not in simulation._fast_cache - january = year.start.period(periods.MONTH) - for month in [january.offset(index) for index in range(12)]: - np.testing.assert_array_equal( - simulation.calculate("flow_m", month), fresh.calculate("flow_m", month) - ) - np.testing.assert_array_equal( - simulation.calculate("flow_m", year), fresh.calculate("flow_m", year) - ) + assert simulation._user_input_keys == input_keys + assert not hasattr(holder, "_user_input_storage") + assert not getattr(simulation, "_user_input_contexts", []) def test_every_stored_value_counts_as_an_input_without_a_record(): diff --git a/tests/core/test_subsample_deleted_inputs.py b/tests/core/test_subsample_deleted_inputs.py index e6b276673..5dbbe5d31 100644 --- a/tests/core/test_subsample_deleted_inputs.py +++ b/tests/core/test_subsample_deleted_inputs.py @@ -70,16 +70,22 @@ def delete_override(simulation, variable, through_holder): @pytest.mark.parametrize("variable", ["doubled", "deleted_leaf", "deleted_eternal"]) @pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) @pytest.mark.parametrize("through_holder", [False, True], ids=["simulation", "holder"]) -def test_deleted_input_recalculated_as_derived_is_not_exportable( +def test_deleted_input_recalculation_is_not_exportable( variable, on_disk, through_holder ): simulation = prepare(variable, on_disk) delete_override(simulation, variable, through_holder) - simulation.calculate(variable, DATASET_YEAR) + result = simulation.calculate(variable, DATASET_YEAR) - assert simulation.get_holder(variable).is_derived( - period(DATASET_YEAR), simulation.branch_name - ) + holder = simulation.get_holder(variable) + if variable == "deleted_leaf": + # A missing input-only root returns its default without storing it. + np.testing.assert_array_equal(result, np.full(simulation.persons.count, 17)) + assert holder.get_array(period(DATASET_YEAR), simulation.branch_name) is None + assert (variable, period(DATASET_YEAR)) not in simulation._fast_cache + else: + # Formula results remain stored as derived values, never new inputs. + assert holder.is_derived(period(DATASET_YEAR), simulation.branch_name) assert simulation._get_set_input_periods(variable) == [] diff --git a/tests/core/test_user_input_keys_soundness.py b/tests/core/test_user_input_keys_soundness.py index 898001b75..e8bf43b32 100644 --- a/tests/core/test_user_input_keys_soundness.py +++ b/tests/core/test_user_input_keys_soundness.py @@ -178,7 +178,7 @@ def test_memory_input_wins_over_a_conflicting_disk_cache(invalidate): @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( +def test_early_year_input_respects_storage_period_validation( year, variable_type, entrypoint, on_disk ): system = CountryTaxBenefitSystem() @@ -195,6 +195,21 @@ def test_early_year_input_is_stored_and_recorded( else: period = periods.Period((periods.DAY, periods.instant((year, 3, 17)), 2)) + if year < 1000: + input_keys = simulation._user_input_keys.copy() + with pytest.raises(ValueError, match="string form cannot be parsed"): + if entrypoint == "simulation": + simulation.set_input(variable, period, [120.0]) + else: + holder.set_input(period, [120.0]) + storage = holder._disk_storage if on_disk else holder._memory_storage + assert not storage.has(period) + assert not storage._sequence_numbers + assert simulation._user_input_keys == input_keys + assert not hasattr(holder, "_user_input_storage") + assert simulation._user_input_contexts == [] + return + if entrypoint == "simulation": simulation.set_input(variable, period, [120.0]) else: @@ -214,19 +229,23 @@ def test_early_year_input_is_stored_and_recorded( @pytest.mark.parametrize("on_disk", [False, True]) -def test_zero_padded_early_year_string_records_the_stored_period(on_disk): +def test_zero_padded_early_year_string_is_rejected_before_recording(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]) + input_keys = simulation._user_input_keys.copy() + holder = simulation.get_holder("raw_year") + with pytest.raises(ValueError, match="string form cannot be parsed"): + 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] - ) + storage = holder._disk_storage if on_disk else holder._memory_storage + assert not storage.has(periods.period(999)) + assert not storage._sequence_numbers + assert simulation._user_input_keys == input_keys + assert not hasattr(holder, "_user_input_storage") assert simulation._user_input_contexts == [] @@ -258,7 +277,7 @@ def test_deleted_memory_input_does_not_turn_a_disk_calculation_into_an_input(): assert _key("rent", JANUARY) not in simulation._user_input_keys -def test_deleted_memory_input_keeps_a_surviving_real_disk_input(): +def test_contained_deletion_forgets_both_real_input_tiers(): simulation = _simulation() _store_on_disk(simulation, "rent") holder = simulation.get_holder("rent") @@ -271,15 +290,16 @@ def test_deleted_memory_input_keeps_a_surviving_real_disk_input(): 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] + assert holder._disk_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._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] + assert holder._disk_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 @pytest.mark.parametrize("on_disk", [False, True]) @@ -304,7 +324,7 @@ def test_reform_does_not_replay_an_unrecorded_carry_over_input(on_disk): ) -def test_registered_disk_input_survives_a_new_memory_input_and_clone_deletion(): +def test_registered_disk_input_is_forgotten_on_contained_clone_deletion(): """Restoration can register an input without holder tier metadata (#576).""" simulation = _simulation() _store_on_disk(simulation, "rent") @@ -318,18 +338,19 @@ def test_registered_disk_input_survives_a_new_memory_input_and_clone_deletion(): 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 _key("rent", JANUARY) not in clone._user_input_keys + assert "rent__2025-01" not in clone.to_input_dataframe().columns + assert clone.get_holder("rent")._disk_storage.get(JANUARY) is None + np.testing.assert_array_equal(holder._memory_storage.get(JANUARY), [700.0]) + np.testing.assert_array_equal(holder._disk_storage.get(JANUARY), [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] + assert _key("rent", JANUARY) not in sim._user_input_keys + assert "rent__2025-01" not in sim.to_input_dataframe().columns 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] - ) + assert sim.get_holder("rent")._disk_storage.get(JANUARY) is None def test_positional_derived_argument_keeps_masters_meaning(): @@ -355,22 +376,22 @@ def test_positional_derived_argument_keeps_masters_meaning(): @hypothesis.settings(max_examples=30, deadline=None) @hypothesis.example( - disk_is_input=False, delete_memory=True, disk_value=0, memory_value=500 + disk_is_input=False, delete_period=True, disk_value=0, memory_value=500 ) @hypothesis.example( - disk_is_input=True, delete_memory=True, disk_value=500, memory_value=700 + disk_is_input=True, delete_period=True, disk_value=500, memory_value=700 ) @hypothesis.example( - disk_is_input=False, delete_memory=False, disk_value=99, memory_value=42 + disk_is_input=False, delete_period=False, disk_value=99, memory_value=42 ) @hypothesis.given( disk_is_input=st.booleans(), - delete_memory=st.booleans(), + delete_period=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 + disk_is_input, delete_period, disk_value, memory_value ): simulation = _simulation() _store_on_disk(simulation, "rent") @@ -385,12 +406,10 @@ def test_tier_provenance_replays_the_same_inputs_as_a_fresh_simulation( assert holder._disk_storage.get(JANUARY) is None # Preserve the generated conflicting-tier witness after cleanup. holder._disk_storage.put(np.array([float(disk_value)]), JANUARY) - if delete_memory: + if delete_period: simulation.delete_arrays("rent", "2025") - expected = ( - (disk_value if disk_is_input else None) if delete_memory else memory_value - ) + expected = None if delete_period else memory_value reference = _simulation() if expected is not None: reference.set_input("rent", JANUARY, [expected])