diff --git a/changelog.d/fix-carry-over-order.fixed.md b/changelog.d/fix-carry-over-order.fixed.md new file mode 100644 index 000000000..a25488556 --- /dev/null +++ b/changelog.d/fix-carry-over-order.fixed.md @@ -0,0 +1 @@ +Auto-carry-over now carries only inputs, taking the latest one stored for a period that starts no later than the requested period, so a carried value no longer depends on which periods were calculated first or on a later input. diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 9b70e2e99..f1c3304f4 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -38,6 +38,15 @@ def _read_only_view(array: numpy.ndarray) -> numpy.ndarray: # share nothing, and an empty set is 216 bytes. _NOTHING_SHARED: FrozenSet[str] = frozenset() +# What ``InMemoryStorage._inputs`` is while a storage holds no input, for the +# same reason: most storages in a simulation hold no input. +_NO_INPUTS: FrozenSet[str] = frozenset() + +# Key a pickled storage's state carries once it records its inputs; a state +# without it comes from a version that recorded neither inputs nor derived +# values (see ``__setstate__``). +_INPUTS_RECORDED = "_inputs_recorded" + class InMemoryStorage: """ @@ -52,12 +61,47 @@ class InMemoryStorage: # this class attribute. A key left in the set after code outside this # class empties ``_arrays`` costs one extra copy at most. _shared: Union[Set[str], FrozenSet[str]] = _NOTHING_SHARED + # Keys whose value was stored as an input: with ``put(..., derived=False)``, + # the default. Every other stored value was calculated by the simulation + # (``put(..., derived=True)``), which ``is_derived`` reports. A key counts + # only while it is stored, and every ``put`` sets or clears its mark. + # + # The storage records inputs rather than derived values because a + # simulation calculates far more values than it is given (in a + # policyengine-us household, 5,201 of 5,204 stored values are derived). + # As with ``_shared``, a storage has a set of its own only while it holds + # an input; otherwise it reads this class attribute. + # + # A storage's own inputs are a set it alone changes, or a frozenset it + # shares with storages cloned from it (or it from them): ``clone`` hands + # both the same frozenset, and whichever changes its inputs first + # replaces it with a set of its own (``_own_inputs``). A branch clones + # every holder, so copying at each clone would cost a set per input + # holder per branch for inputs most branches never change. + _inputs: Union[Set[str], FrozenSet[str]] = _NO_INPUTS is_eternal: bool def __init__(self, is_eternal: bool): self._arrays = {} self.is_eternal = is_eternal + def __getstate__(self) -> dict: + state = self.__dict__.copy() + state[_INPUTS_RECORDED] = True + return state + + def __setstate__(self, state: dict) -> None: + if not state.pop(_INPUTS_RECORDED, False): + # Pickled before storages recorded their inputs, when nothing + # marked a calculated value: every stored value counts as an + # input, as it did then, except any a development version of this + # change marked derived. + derived = state.pop("_derived", ()) + inputs = set(state.get("_arrays", {})).difference(derived) + if inputs: + state["_inputs"] = inputs + self.__dict__.update(state) + def clone(self, share_arrays: bool = False) -> "InMemoryStorage": """Copy this storage. @@ -95,6 +139,19 @@ def clone(self, share_arrays: bool = False) -> "InMemoryStorage": clone._shared = shared else: clone._arrays = {key: array.copy() for key, array in self._arrays.items()} + # Share the inputs (see ``_inputs``). A mark left behind by code + # outside this class emptying ``_arrays`` never counts (see + # ``is_derived``), so the clone gets only those of keys it holds. + inputs = self._inputs + if inputs: + if inputs <= clone._arrays.keys(): + if not isinstance(inputs, frozenset): + inputs = self.__dict__["_inputs"] = frozenset(inputs) + clone._inputs = inputs + else: + held = inputs.intersection(clone._arrays) + if held: + clone._inputs = frozenset(held) return clone def get(self, period: Period, branch_name: str = "default") -> ArrayLike: @@ -113,6 +170,23 @@ def get(self, period: Period, branch_name: str = "default") -> ArrayLike: self._stop_sharing(key) return values + def has(self, period: Period, branch_name: str = "default") -> bool: + """Whether a value is stored for ``period`` under ``branch_name``. + + Unlike ``get``, this never copies an array shared by ``clone``. + """ + if self.is_eternal: + period = periods.period(periods.ETERNITY) + return f"{branch_name}:{periods.period(period)}" in self._arrays + + def is_derived(self, period: Period, branch_name: str = "default") -> bool: + """Whether the value stored for ``period`` under ``branch_name`` was + stored with ``derived=True``; ``False`` if none is stored.""" + if self.is_eternal: + period = periods.period(periods.ETERNITY) + key = f"{branch_name}:{periods.period(period)}" + return key in self._arrays and key not in self._inputs + # Each method below reads ``self._shared`` once and works on that set, and # releases it with one ``dict.pop`` that never raises. A storage is still # not safe to read from several threads at once (see ``clone``), but @@ -144,8 +218,48 @@ def _release_if_empty(self, shared: Set[str]) -> None: if not shared: self.__dict__.pop("_shared", None) + def _own_inputs(self) -> Set[str]: + """This storage's inputs as a set that only it changes, made from the + inputs it shares, or from nothing, if need be.""" + inputs = self.__dict__.get("_inputs") + if inputs is None: + inputs = self.__dict__["_inputs"] = set() + elif isinstance(inputs, frozenset): + inputs = self.__dict__["_inputs"] = set(inputs) + return inputs + + def _mark_input(self, key: str) -> None: + """Record that the value stored for ``key`` is an input.""" + if key not in self._inputs: + self._own_inputs().add(key) + + def _unmark_input(self, key: str) -> None: + """Record that the value stored for ``key`` was calculated.""" + if key in self._inputs: + inputs = self._own_inputs() + inputs.discard(key) + self._release_inputs_if_empty(inputs) + + def _unmark_dropped_keys(self) -> None: + """Forget the input marks of keys ``_arrays`` no longer has.""" + inputs = self._inputs + if inputs and not inputs <= self._arrays.keys(): + inputs = self._own_inputs() + inputs.intersection_update(self._arrays) + self._release_inputs_if_empty(inputs) + + def _release_inputs_if_empty(self, inputs: Set[str]) -> None: + """Go back to the class's no-inputs object once no stored value is an + input, as ``_release_if_empty`` does for ``_shared``.""" + if not inputs: + self.__dict__.pop("_inputs", None) + def put( - self, value: ArrayLike, period: Period, branch_name: str = "default" + self, + value: ArrayLike, + period: Period, + branch_name: str = "default", + derived: bool = False, ) -> None: if self.is_eternal: period = periods.period(periods.ETERNITY) @@ -169,6 +283,10 @@ def put( key = f"{branch_name}:{period}" self._arrays[key] = value self._stop_sharing(key) + if derived: + self._unmark_input(key) + else: + self._mark_input(key) def delete(self, period: Period = None, branch_name: str = "default") -> None: if period is None: @@ -181,6 +299,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: if not period_item.startswith(branch_prefix) } self._stop_sharing_dropped_keys() + self._unmark_dropped_keys() return if self.is_eternal: @@ -199,6 +318,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: ) } self._stop_sharing_dropped_keys() + self._unmark_dropped_keys() def get_known_periods(self) -> list: # Split on the first colon only: an anchored period's string form diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 3563c805d..db3e2cad0 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -22,10 +22,19 @@ def __init__( ): self._files = {} self._enums = {} + # File keys stored with ``put(..., derived=True)``; see + # ``InMemoryStorage``. + self._derived = set() self.is_eternal = is_eternal self.preserve_storage_dir = preserve_storage_dir self.storage_dir = storage_dir + def __setstate__(self, state: dict) -> None: + # A storage pickled before derived marks existed has none: its values + # count as inputs. + state.setdefault("_derived", set()) + self.__dict__.update(state) + def clone(self) -> "OnDiskStorage": """Create a private metadata view over this storage directory. @@ -43,6 +52,7 @@ def clone(self) -> "OnDiskStorage": ) clone._files = self._files.copy() clone._enums = self._enums.copy() + clone._derived = set(self._derived) clone._storage_dir_owner = getattr(self, "_storage_dir_owner", self) return clone @@ -63,8 +73,29 @@ def get(self, period: Period, branch_name: str = "default") -> ArrayLike: return None return self._decode_file(values) + def has(self, period: Period, branch_name: str = "default") -> bool: + """Whether a value is stored for ``period`` under ``branch_name``. + + Unlike ``get``, this reads no file. + """ + if self.is_eternal: + period = periods.period(periods.ETERNITY) + return f"{branch_name}_{periods.period(period)}" in self._files + + def is_derived(self, period: Period, branch_name: str = "default") -> bool: + """Whether the value stored for ``period`` under ``branch_name`` was + stored with ``derived=True``; ``False`` if none is stored.""" + if self.is_eternal: + period = periods.period(periods.ETERNITY) + key = f"{branch_name}_{periods.period(period)}" + return key in self._derived and key in self._files + def put( - self, value: ArrayLike, period: Period, branch_name: str = "default" + self, + value: ArrayLike, + period: Period, + branch_name: str = "default", + derived: bool = False, ) -> None: if self.is_eternal: period = periods.period(periods.ETERNITY) @@ -77,6 +108,10 @@ def put( value = value.view(numpy.ndarray) numpy.save(path, value) self._files[filename] = path + if derived: + self._derived.add(filename) + else: + self._derived.discard(filename) def delete(self, period: Period = None, branch_name: str = "default") -> None: if period is None: @@ -89,6 +124,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: for period_item, value in self._files.items() if not period_item.startswith(branch_prefix) } + self._derived.intersection_update(self._files) return if self.is_eternal: @@ -101,6 +137,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: for period_item, value in self._files.items() if not period_item == f"{branch_name}_{period}" } + self._derived.intersection_update(self._files) def get_known_periods(self) -> list: return list([periods.period(x.split("_")[1]) for x in self._files.keys()]) @@ -113,6 +150,8 @@ def get_known_branch_periods(self) -> list: def restore(self) -> None: self._files = files = {} + # Files read back from a directory carry no derived marks. + self._derived = set() # Restore self._files from content of storage_dir. for filename in os.listdir(self.storage_dir): if not filename.endswith(".npy"): diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 11a8ae50b..a02208894 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -347,9 +347,15 @@ def _set( value: ArrayLike, branch_name: str = "default", validate_nan: bool = False, + derived: bool = False, ) -> None: simulation = getattr(self, "simulation", None) - user_input_contexts = getattr(simulation, "_user_input_contexts", None) + # A value calculated while an input is being set (say, by a + # ``set_input`` helper that calculates) is not part of that input: it + # belongs to the branch it was calculated on. + user_input_contexts = ( + None if derived else getattr(simulation, "_user_input_contexts", None) + ) if user_input_contexts and branch_name == "default": branch_name = user_input_contexts[-1] value = self._to_array(value, validate_nan=validate_nan) @@ -367,17 +373,34 @@ def _set( ) if should_store_on_disk: - self._disk_storage.put(value, period, branch_name) + self._disk_storage.put(value, period, branch_name, derived=derived) else: - self._memory_storage.put(value, period, branch_name) + self._memory_storage.put(value, period, branch_name, derived=derived) if user_input_contexts: if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() simulation._user_input_keys.add((self.variable.name, branch_name, period)) def put_in_cache( - self, value: ArrayLike, period: Period, branch_name: str = "default" + self, + value: ArrayLike, + period: Period, + branch_name: str = "default", + derived: bool = False, ) -> None: + """Cache ``value`` for ``period``. + + ``derived`` marks a value the simulation calculated rather than took + as input: a formula result, a carried, uprated or default value, a + twelfth of a yearly flow cached at a month by ``calculate_divide``, + or a sum over several sub-periods cached by ``calculate_add``. + Auto-carry-over never carries such a value into another period (see + ``is_derived``). The mark is stored with the value, for the (branch, + period) key written, and any later write to that key replaces it. + + A derived value never replaces an input that ``get_array(period, + branch_name)`` reads: the input is kept and nothing is stored. + """ if self._do_not_store: return @@ -388,7 +411,14 @@ def put_in_cache( ): return - self._set(period, value, branch_name) + if ( + derived + and self._branch_storing(period, branch_name) is not None + and not self.is_derived(period, branch_name) + ): + return + + self._set(period, value, branch_name, derived=derived) def default_array(self) -> ArrayLike: """ @@ -396,3 +426,80 @@ def default_array(self) -> ArrayLike: """ return self.variable.default_array(self.population.count) + + def _stores(self, period: Period, branch_name: str) -> bool: + """Whether a value is stored for ``period`` under ``branch_name``, + without reading or copying it.""" + return self._memory_storage.has(period, branch_name) or ( + self._disk_storage is not None + and self._disk_storage.has(period, branch_name) + ) + + def _readable_branches(self, branch_name: str = "default") -> List[str]: + """``get_array``'s lookup order: the branch, its ``parent_branch`` + ancestors, then ``default``.""" + names = [branch_name] + if branch_name != "default": + parent = ( + getattr(self.simulation, "parent_branch", None) + if self.simulation + else None + ) + while parent is not None: + names.append(parent.branch_name) + parent = getattr(parent, "parent_branch", None) + names.append("default") + return list(dict.fromkeys(names)) + + def _branch_storing(self, period: Period, branch_name: str = "default") -> str: + """The branch whose stored value ``get_array(period, branch_name)`` + reads, or ``None`` if none stores one.""" + for name in self._readable_branches(branch_name): + if self._stores(period, name): + return name + return None + + def is_derived(self, period: Period, branch_name: str = "default") -> bool: + """Whether the value ``get_array(period, branch_name)`` reads was + calculated by the simulation rather than set as an input. + + The answer comes from the branch that stores the value read: the + branch itself, else its ``parent_branch`` ancestors, else + ``default``. ``False`` if no value is stored for ``period``. + """ + storing = self._branch_storing(period, branch_name) + if storing is None: + return False + if self._memory_storage.has(period, storing): + return self._memory_storage.is_derived(period, storing) + return self._disk_storage.is_derived(period, storing) + + def get_input_periods(self, branch_name: str = "default") -> List[Period]: + """The periods for which the value ``get_array(period, branch_name)`` + reads is an input rather than a value the simulation calculated (see + ``put_in_cache``). Periods stored only under branches this one cannot + read are left out. + + One pass over the stored keys: for each period, the key ``get_array`` + reads first (the branch before its ancestors, memory before disk). + """ + rank = { + name: index + for index, name in enumerate(self._readable_branches(branch_name)) + } + storages = [self._memory_storage] + if self._disk_storage is not None: + storages.append(self._disk_storage) + read = {} + for order, storage in enumerate(storages): + for stored_branch, period in storage.get_known_branch_periods(): + if stored_branch not in rank: + continue + key = (rank[stored_branch], order) + if period not in read or key < read[period][0]: + read[period] = (key, storage, stored_branch) + return [ + period + for period, (_, storage, stored_branch) in read.items() + if not storage.is_derived(period, stored_branch) + ] diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 2652d9723..fa3acfa2e 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -103,6 +103,38 @@ def _uprating_index_value(parameter, instant) -> Optional[float]: return defined[-1].value +_DAYS_BEFORE_MONTH = (0, 31, 59, 90, 120, 151, 181, 212, 243, 273, 304, 334) + + +def _end_order(period: Period) -> float: + """Sorts periods by when they end: the number of the day after the + period's last day, counting days as ``date.toordinal`` does. + + ``period.stop`` gives the same order wherever it has a value, but it + raises for a period that ends after 9999-12-31 (``day:9999-12-30:3``), + the last date ``datetime`` can represent. This is integer arithmetic on + the same (proleptic Gregorian) calendar, so every period has a value and + one that ends later always sorts later. + """ + unit, (year, month, day), size = period + if unit == ETERNITY: + return float("inf") + if unit == periods.DAY: + day += size + elif unit == MONTH: + year, month = divmod(year * 12 + month - 1 + size, 12) + month += 1 + else: + year += size + # The first day of that month, then ``day - 1`` days on: a day past the + # month's end runs into the next month, as it does in ``period.stop``. + leap = year % 4 == 0 and (year % 100 != 0 or year % 400 == 0) + days_before_year = ( + (year - 1) * 365 + (year - 1) // 4 - (year - 1) // 100 + (year - 1) // 400 + ) + return days_before_year + _DAYS_BEFORE_MONTH[month - 1] + (leap and month > 2) + day + + if TYPE_CHECKING: from policyengine_core.taxbenefitsystems import TaxBenefitSystem @@ -394,8 +426,10 @@ def _invalidate_all_caches(self) -> None: for population in self.populations.values(): for holder in population._holders.values(): holder._memory_storage._arrays = {} + holder._memory_storage._unmark_dropped_keys() if holder._disk_storage is not None: holder._disk_storage._files = {} + holder._disk_storage._derived = set() # Replay preserved user inputs so ``calculate`` still sees them. for user_input in preserved: holder = self.get_holder(user_input.variable_name) @@ -881,7 +915,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if np.all(~mask): array = holder.default_array() array = self._cast_formula_result(array, variable) - holder.put_in_cache(array, period, self.branch_name) + holder.put_in_cache(array, period, self.branch_name, derived=True) return array array = None @@ -942,18 +976,63 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: and variable.calculate_output is None and len(known_periods) > 0 ): - # Variables with a calculate-output property specify - # Sort by period.start (temporal order). Sorting Period - # tuples lexicographically puts "year" before "month" - # alphabetically, so a known "2023" annual value would - # win over a later "2024-06" monthly value (bug H1). - last_known_period = max(known_periods, key=lambda p: p.start) - if last_known_period.start > period.start: + # Carry over the latest input: of the stored periods that + # start no later than ``period``, the one that starts last + # (on a tie, the one that ends last, then the larger unit, + # so the choice never depends on the order inputs were + # stored), preferring periods at + # the variable's own definition-period unit and using + # another unit only when there is none, as for an input to + # a variable with no ``set_input`` helper. Compare + # period.start (temporal order): sorting Period tuples + # lexicographically puts "year" before "month" + # alphabetically, so a known "2023" annual value would win + # over a later "2024-06" monthly value (bug H1). + # + # Only inputs carry. Every value this simulation + # calculated is marked derived when cached (formula + # results, carried, uprated and default values, a twelfth + # cached by ``calculate_divide``, a sum cached by + # ``calculate_add``), and carrying one would make the + # result depend on what was calculated first: a later + # period's carried value would hide an earlier input, and + # a value already masked by ``defined_for``, or given by a + # formula that has since ended, would carry forward. + # A later input does not carry backwards. + last_known_period = max( + ( + input_period + for input_period in holder.get_input_periods( + self.branch_name + ) + if input_period.start <= period.start + ), + key=lambda p: ( + p.unit == variable.definition_period, + p.start, + _end_order(p), + periods.unit_weight(p.unit), + str(p), + ), + default=None, + ) + if last_known_period is not None: + # Pass branch_name through so auto-carry-over respects + # the active branch instead of reaching for the + # "default" branch's cache (bug H2). + array = holder.get_array(last_known_period, self.branch_name) + elif any( + known_period.start > period.start + for known_period in known_periods + ): + # No input to carry, but a later period is stored: as + # before, return the default without caching it. A + # cached default would change what a formula testing + # whether a value is stored sees, and the uprating path + # above would uprate later periods from it. return holder.default_array() - # Pass branch_name through so auto-carry-over respects the - # active branch instead of reaching for the "default" - # branch's cache (bug H2). - array = holder.get_array(last_known_period, self.branch_name) + else: + array = holder.default_array() else: array = holder.default_array() @@ -969,7 +1048,8 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: array = EnumArray(array, variable.possible_values) array = self._cast_formula_result(array, variable) - holder.put_in_cache(array, period, self.branch_name) + # Calculated, not input: auto-carry-over never carries it. + holder.put_in_cache(array, period, self.branch_name, derived=True) except SpiralError: array = holder.default_array() @@ -1044,12 +1124,16 @@ def calculate_add( ) ) + sub_periods = list(period.get_subperiods(variable.definition_period)) result = sum( - self.calculate(variable_name, sub_period) - for sub_period in period.get_subperiods(variable.definition_period) + self.calculate(variable_name, sub_period) for sub_period in sub_periods ) - holder = self.get_holder(variable.name) - holder.put_in_cache(result, period, self.branch_name) + # Cache only a sum over several sub-periods, as derived from them. A + # single sub-period's value is already stored, as an input or + # derived, by ``calculate``. + if len(sub_periods) > 1: + holder = self.get_holder(variable.name) + holder.put_in_cache(result, period, self.branch_name, derived=True) return result def calculate_divide( @@ -1082,7 +1166,7 @@ def calculate_divide( computation_period = period.this_year result = self.calculate(variable_name, period=computation_period) / 12.0 holder = self.get_holder(variable.name) - holder.put_in_cache(result, period, self.branch_name) + holder.put_in_cache(result, period, self.branch_name, derived=True) return result elif period.unit == periods.YEAR: return self.calculate(variable_name, period) diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index c381b11c4..b7073081d 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -87,7 +87,15 @@ class TaxBenefitSystem: parameters_dir: str = None """Directory containing the YAML parameter tree.""" auto_carry_over_input_variables: bool = False - """Whether to automatically carry over input variables when calculating a variable for a period different from the period of the input variables.""" + """Whether a variable with no formula result for a period (and no + ``uprating`` path from an earlier period) takes an input from another + period. If so, the period takes the input stored for the latest-starting + period that starts no later than it (on a tie, the one that ends last), + preferring inputs at the variable's own definition-period unit, masked by + the variable's ``defined_for``; with no such input, the default. An input + stored for the period itself is read back as stored. An input never carries backwards, and values the simulation + calculated never carry (see ``Holder.is_derived``), so the result does + not depend on which periods were calculated first.""" basic_inputs: List[str] = None """Short list of basic inputs to get medium accuracy.""" modelled_policies: str = None diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index c3db0c4fa..65b3c6644 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -9,6 +9,10 @@ from policyengine_core.periods import ETERNITY from policyengine_core.simulations import Simulation +# Periods, one per line, whose dumped value the simulation calculated (see +# ``Holder.is_derived``), so a restored simulation does not carry them over. +DERIVED_PERIODS_FILE = "derived_periods.txt" + def dump_simulation(simulation, directory): """ @@ -66,9 +70,18 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): def _dump_holder(holder, directory): disk_storage = holder.create_disk_storage(directory, preserve=True) + derived_periods = set() for period in holder.get_known_periods(): value = holder.get_array(period) disk_storage.put(value, period) + # Read the mark of exactly the value dumped: the same period on the + # same branch as ``get_array``. + if holder.is_derived(period): + derived_periods.add(str(period)) + if derived_periods: + path = os.path.join(disk_storage.storage_dir, DERIVED_PERIODS_FILE) + with open(path, "w") as file: + file.write("\n".join(sorted(derived_periods)) + "\n") def _dump_entity(population, directory): @@ -135,6 +148,12 @@ def _restore_holder(simulation, variable, directory): holder = simulation.get_holder(variable) + derived_periods_path = os.path.join(storage_dir, DERIVED_PERIODS_FILE) + derived_periods = set() + if os.path.exists(derived_periods_path): + with open(derived_periods_path) as file: + derived_periods = set(file.read().split()) + for period in disk_storage.get_known_periods(): value = disk_storage.get(period) - holder.put_in_cache(value, period) + holder.put_in_cache(value, period, derived=str(period) in derived_periods) diff --git a/tests/core/test_carry_over_order.py b/tests/core/test_carry_over_order.py new file mode 100644 index 000000000..0fc88939f --- /dev/null +++ b/tests/core/test_carry_over_order.py @@ -0,0 +1,731 @@ +"""Auto-carry-over depends only on the inputs, not on what was calculated first. + +With ``auto_carry_over_input_variables``, a variable with no formula result +for a period takes its value from another period. Before this fix, +``Simulation._calculate`` took the latest-starting stored period other than +a ``calculate_add``/``calculate_divide`` cache (#557), and returned the +default if that period started after the requested one. Two things made the +result depend on calculation order: + +* Any later stored period blocked carry-over. A later input did (inputs for + 2012 and 2014 gave the default for 2013), and so did a value the simulation + had itself calculated: once 2014 had been calculated, carrying the 2012 + input into it, 2013 came out as the default instead of the 2012 input. +* Values the simulation calculated were carried like inputs. A value masked by + ``defined_for`` carried its mask into later periods, a default cached for a + period where ``defined_for`` was false everywhere replaced the input, and a + formula result carried past the formula's ``end``. + +The rule now is the reference rule in ``tests/fixtures/carry_over.py``. Every +value the simulation caches is marked derived for the (branch, period) key it +is stored under, ``Holder.is_derived`` reports the mark of the value a branch +reads, and only values that are not derived carry. A derived value never +replaces an input a branch can read. + +Each regression below checks a simulation that calculated other periods first +against a fresh simulation that calculates only the period in question; the +properties are in ``test_carry_over_order_property.py``. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.data_storage import InMemoryStorage, OnDiskStorage +from policyengine_core.experimental import MemoryConfig +from policyengine_core.reforms import Reform +from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, +) +from tests.fixtures.carry_over import alone, build_system, simulation + + +@pytest.fixture(scope="module") +def system(): + return build_system() + + +def _after(system, inputs, requests, variable, period): + built = simulation(system, inputs) + for requested_variable, requested_period in requests: + built.calculate(requested_variable, requested_period) + return built.calculate(variable, period) + + +@pytest.mark.parametrize( + "variable,earlier_input,later_period,between", + [ + ("carried", "2012", "2014", "2013"), + ("carried_count", "2012", "2014", "2013"), + ("carried_monthly", "2012-01", "2012-06", "2012-03"), + ], +) +def test_calculated_later_period_does_not_hide_earlier_input( + system, variable, earlier_input, later_period, between +): + inputs = {variable: {earlier_input: [7, 8]}} + result = _after(system, inputs, [(variable, later_period)], variable, between) + np.testing.assert_array_equal(result, [7, 8]) + np.testing.assert_array_equal(result, alone(system, inputs, variable, between)) + + +@pytest.mark.parametrize( + "variable,earlier,later,between", + [ + ("carried", "2012", "2014", "2013"), + ("carried_monthly", "2012-01", "2012-06", "2012-03"), + ], +) +def test_later_input_does_not_hide_earlier_input( + system, variable, earlier, later, between +): + inputs = {variable: {earlier: [7, 8], later: [9, 10]}} + np.testing.assert_array_equal(alone(system, inputs, variable, between), [7, 8]) + + +def test_inputs_are_stored_in_any_order(system): + inputs = {"carried": {"2014": [9, 10], "2012": [7, 8]}} + for year, expected in [("2013", [7, 8]), ("2014", [9, 10]), ("2016", [9, 10])]: + np.testing.assert_array_equal(alone(system, inputs, "carried", year), expected) + + +@pytest.mark.parametrize( + "first,second", [("2013", "2013-01"), ("2013-01", "2013")], ids=["year", "month"] +) +def test_inputs_starting_together_carry_the_one_ending_last(system, first, second): + values = {"2013": [1, 1], "2013-01": [2, 2]} + inputs = { + "year_input_without_helper": {first: values[first], second: values[second]} + } + np.testing.assert_array_equal( + alone(system, inputs, "year_input_without_helper", "2014"), [1, 1] + ) + + +@pytest.mark.parametrize( + "first,second", [("2013", "year:2013:2"), ("year:2013:2", "2013")] +) +def test_inputs_at_one_unit_starting_together_carry_the_one_ending_last( + system, first, second +): + values = {"2013": [1, 1], "year:2013:2": [2, 2]} + inputs = { + "year_input_without_helper": {first: values[first], second: values[second]} + } + np.testing.assert_array_equal( + alone(system, inputs, "year_input_without_helper", "2015"), [2, 2] + ) + + +@pytest.mark.parametrize("reverse", [False, True], ids=["month-first", "day-first"]) +def test_inputs_with_the_same_extent_carry_the_larger_unit(system, reverse): + """Two inputs covering the same days in other units than the variable's + resolve the same way whichever was stored first.""" + values = {"month:2013-01:2": [10, 10], "day:2013-01-01:59": [20, 20]} + order = list(values)[::-1] if reverse else list(values) + inputs = {"year_input_without_helper": {key: values[key] for key in order}} + np.testing.assert_array_equal( + alone(system, inputs, "year_input_without_helper", "2014"), [10, 10] + ) + + +@pytest.mark.parametrize( + "stored,asked", + [("day:9999-12-30:3", "9999-12-31"), ("day:2012-01-01:4000000", "2013-01-01")], +) +def test_an_input_ending_after_the_last_date_carries(system, stored, asked): + """A period that ends after 9999-12-31 has no ``stop`` (it raises + ``OverflowError``); its input carries all the same.""" + inputs = {"day_input_without_helper": {stored: [10, 20]}} + np.testing.assert_array_equal( + alone(system, inputs, "day_input_without_helper", asked), [10, 20] + ) + + +@pytest.mark.parametrize("longer_first", [True, False], ids=["longer", "shorter"]) +def test_an_input_ending_after_the_last_date_ends_last(system, longer_first): + values = {"day:2012-01-01:4000000": [10, 20], "day:2012-01-01:2": [1, 2]} + order = list(values) if longer_first else list(values)[::-1] + inputs = {"day_input_without_helper": {key: values[key] for key in order}} + np.testing.assert_array_equal( + alone(system, inputs, "day_input_without_helper", "2013-01-01"), [10, 20] + ) + + +@pytest.mark.parametrize("longer_first", [True, False], ids=["longer", "shorter"]) +def test_of_two_inputs_ending_after_the_last_date_the_later_carries( + system, longer_first +): + """Both end after 9999-12-31. As strings, ``:10`` sorts before ``:9``.""" + values = {"day:9999-12-30:10": [10, 10], "day:9999-12-30:9": [9, 9]} + order = list(values) if longer_first else list(values)[::-1] + inputs = {"day_input_without_helper": {key: values[key] for key in order}} + np.testing.assert_array_equal( + alone(system, inputs, "day_input_without_helper", "9999-12-31"), [10, 10] + ) + + +@pytest.mark.parametrize("months_first", [True, False], ids=["months", "days"]) +def test_inputs_in_two_units_ending_after_the_last_date_carry_the_later( + system, months_first +): + """800 days from 1 February 9999 end after 24 months from the same day. + The larger unit and the string order both favour the months, so only + the end picks the days.""" + values = {"day:9999-02-01:800": [8, 8], "month:9999-02:24": [24, 24]} + order = list(values)[::-1] if months_first else list(values) + inputs = {"year_input_without_helper": {key: values[key] for key in order}} + np.testing.assert_array_equal( + alone(system, inputs, "year_input_without_helper", "year:9999-02"), [8, 8] + ) + + +def test_later_input_does_not_carry_backwards(system): + built = simulation(system, {"carried": {"2012": [7, 8], "2014": [9, 10]}}) + np.testing.assert_array_equal(built.calculate("carried", "2011"), [0, 0]) + np.testing.assert_array_equal(built.calculate("carried", "2015"), [9, 10]) + np.testing.assert_array_equal(built.calculate("carried", "2011"), [0, 0]) + + +def test_value_masked_by_defined_for_does_not_carry_its_mask(system): + inputs = { + "carried_if_eligible": {"2012": [7, 8]}, + "eligible": {"2013": [False, True], "2014": [True, True]}, + } + built = simulation(system, inputs) + np.testing.assert_array_equal( + built.calculate("carried_if_eligible", "2013"), [0, 8] + ) + np.testing.assert_array_equal( + built.calculate("carried_if_eligible", "2014"), [7, 8] + ) + + +def test_default_cached_where_defined_for_is_false_does_not_carry(system): + inputs = { + "carried_if_eligible": {"2012": [7, 8]}, + "eligible": {"2013": [False, False], "2014": [True, True]}, + } + built = simulation(system, inputs) + np.testing.assert_array_equal( + built.calculate("carried_if_eligible", "2013"), [0, 0] + ) + np.testing.assert_array_equal( + built.calculate("carried_if_eligible", "2014"), [7, 8] + ) + + +@pytest.mark.parametrize("own_input", [None, "2012"]) +def test_formula_result_does_not_carry_past_the_formula_end(system, own_input): + inputs = {"carried": {"2012": [1, 2]}} + if own_input is not None: + inputs["formula_until_2013"] = {own_input: [40, 50]} + expected = alone(system, inputs, "formula_until_2013", "2014") + np.testing.assert_array_equal(expected, [0, 0] if own_input is None else [40, 50]) + result = _after( + system, inputs, [("formula_until_2013", "2013")], "formula_until_2013", "2014" + ) + np.testing.assert_array_equal(result, expected) + + +def test_input_set_after_a_period_defaulted_still_carries(system): + built = simulation(system, {}) + np.testing.assert_array_equal(built.calculate("carried", "2013"), [0, 0]) + built.set_input("carried", "2012", np.array([7.0, 8.0])) + 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).""" + built = simulation(system, {}) + 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 + built = simulation(system, {"carried_monthly": {"2013-06": [1, 2]}}) + 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 + + +def test_branch_input_is_not_hidden_by_a_value_the_parent_carried(system): + built = simulation(system, {"carried": {"2012": [7, 8]}}) + built.calculate("carried", "2013") + branch = built.get_branch("reform") + branch.set_input("carried", "2012", np.array([70.0, 80.0])) + np.testing.assert_array_equal(branch.calculate("carried", "2014"), [70, 80]) + np.testing.assert_array_equal(built.calculate("carried", "2014"), [7, 8]) + + +def test_derived_mark_belongs_to_the_branch_that_stored_the_value(system): + """A value calculated on one branch does not make another branch's input + for the same period derived.""" + + def build(): + built = simulation(system, {"carried": {"2011": [7, 7]}}) + built.get_holder("formula_until_2013").set_input( + periods.period("2012"), np.array([30.0, 30.0]), "other" + ) + return built + + expected = build().get_branch("other").calculate("formula_until_2013", "2014") + np.testing.assert_array_equal(expected, [30, 30]) + built = build() + np.testing.assert_array_equal( + built.calculate("formula_until_2013", "2012"), [12, 12] + ) + other = built.get_branch("other") + holder = other.get_holder("formula_until_2013") + assert holder.is_derived(periods.period("2012"), "default") + assert not holder.is_derived(periods.period("2012"), "other") + np.testing.assert_array_equal( + other.calculate("formula_until_2013", "2014"), expected + ) + + +def test_add_over_a_calculated_period_keeps_it_derived(system): + inputs = { + "carried_if_eligible": {"2012": [7, 8]}, + "eligible": {"2013": [False, True], "2014": [True, True]}, + } + built = simulation(system, inputs) + np.testing.assert_array_equal( + built.calculate_add("carried_if_eligible", "2013"), [0, 8] + ) + np.testing.assert_array_equal( + built.calculate("carried_if_eligible", "2014"), [7, 8] + ) + + +def test_add_over_an_input_period_keeps_it_an_input(system): + built = simulation(system, {"carried": {"2012": [7, 8]}}) + np.testing.assert_array_equal(built.calculate_add("carried", "2012"), [7, 8]) + np.testing.assert_array_equal(built.calculate("carried", "2013"), [7, 8]) + + +@pytest.mark.parametrize( + "variable,input_period,kind,request_period,later_period", + [ + # The sum over the two years would replace the two-year input. + ("year_input_without_helper", "year:2012:2", "add", "year:2012:2", "2014"), + # The twelfth of 2012 would replace the January input. + ("year_input_without_helper", "2012-01", "divide", "2012-01", "2014"), + # The sum of the months would replace the year input. + ("month_input_without_helper", "2013", "add", "2013", "2014-01"), + ], +) +@pytest.mark.parametrize("on_branch", [False, True], ids=["simulation", "branch"]) +def test_add_or_divide_never_replaces_an_input( + system, variable, input_period, kind, request_period, later_period, on_branch +): + inputs = {variable: {input_period: [120, 24]}} + expected = alone(system, inputs, variable, later_period) + np.testing.assert_array_equal(expected, [120, 24]) + built = simulation(system, inputs) + if on_branch: + built = built.get_branch("reform") + if kind == "add": + built.calculate_add(variable, request_period) + else: + built.calculate_divide(variable, request_period) + holder = built.get_holder(variable) + np.testing.assert_array_equal( + holder.get_array(periods.period(input_period), built.branch_name), [120, 24] + ) + np.testing.assert_array_equal(built.calculate(variable, later_period), expected) + + +def test_add_over_an_anchored_year_keeps_the_anchored_input(system): + """``set_input`` stores a month of a yearly flow at the year starting + then; ``ADD`` over that year is a single calendar-year sub-period.""" + inputs = {"carried": {"2012": [1, 2], "2013-06": [5, 6]}} + expected = alone(system, inputs, "carried", "2014") + np.testing.assert_array_equal(expected, [5, 6]) + built = simulation(system, inputs) + built.calculate_add("carried", "year:2013-06") + np.testing.assert_array_equal(built.calculate("carried", "2014"), expected) + + +def test_month_cache_of_a_year_input_does_not_carry_into_the_next_year(system): + """policyengine-us computes ``monthly_age`` from ``age`` one month at a + time; each month of a yearly flow caches a twelfth of the year (#557).""" + inputs = {"carried": {"2024": [40, 6]}} + built = simulation(system, inputs) + for month in range(1, 13): + built.calculate("carried", f"2024-{month:02d}") + np.testing.assert_allclose(built.calculate("carried", "2024-12"), [40 / 12, 6 / 12]) + np.testing.assert_array_equal(built.calculate("carried", "2025"), [40, 6]) + + +def test_own_unit_input_wins_over_an_input_at_another_unit(system): + """As before 3.24.0, and as in #557: the year input, not the later month + input of a variable with no ``set_input`` helper.""" + inputs = {"year_input_without_helper": {"2024": [120, 12], "2024-12": [9, 9]}} + np.testing.assert_array_equal( + alone(system, inputs, "year_input_without_helper", "2025"), [120, 12] + ) + + +# Lifecycle paths that write or drop stored values outside ``Holder._set`` +# (the cases #557's review found to leave per-period marks stale). + +FLOW = "year_input_without_helper" + + +def test_reform_replay_restores_an_input_at_another_unit(system): + built = simulation(system, {FLOW: {"2024-06": [600, 60]}}) + branch = built.get_branch("child") + branch.calculate_divide(FLOW, "2024-06") + branch.apply_reform(_noop) + np.testing.assert_array_equal(branch.calculate(FLOW, "2025"), [600, 60]) + + +def test_dump_and_restore_keep_carrying_the_input(system, tmp_path): + built = simulation(system, {FLOW: {"2024": [120, 12]}}) + built.calculate_divide(FLOW, "2024-06") + dump_simulation(built, str(tmp_path / "dump")) + restored = restore_simulation(str(tmp_path / "dump"), system) + np.testing.assert_array_equal(built.calculate(FLOW, "2025"), [120, 12]) + np.testing.assert_array_equal(restored.calculate(FLOW, "2025"), [120, 12]) + + +def test_deleting_a_branch_cache_exposes_the_ancestor_input(system): + built = simulation(system, {FLOW: {"2024-06": [600, 60]}}) + branch = built.get_branch("child") + branch.calculate_divide(FLOW, "2024-06") + branch.get_holder(FLOW).delete_arrays(periods.period("2024-06"), "child") + np.testing.assert_array_equal(branch.calculate(FLOW, "2025"), [600, 60]) + np.testing.assert_array_equal(built.calculate(FLOW, "2025"), [600, 60]) + + +def test_deleting_a_branch_input_does_not_expose_the_ancestor_twelfth(system): + built = simulation(system, {FLOW: {"2024": [120, 12]}}) + built.calculate_divide(FLOW, "2024-06") + branch = built.get_branch("child") + branch.set_input(FLOW, "2024-06", np.array([600.0, 60.0])) + branch.get_holder(FLOW).delete_arrays(periods.period("2024-06"), "child") + np.testing.assert_array_equal(branch.calculate(FLOW, "2025"), [120, 12]) + np.testing.assert_array_equal(built.calculate(FLOW, "2025"), [120, 12]) + + +@pytest.mark.parametrize("delete_first", [True, False], ids=["deleted", "replaced"]) +def test_a_value_written_straight_to_storage_is_not_derived(system, delete_first): + """The mark is stored with the value: a value written later at the same + key, even straight to the storage as ``apply_reform``'s input replay + does, is not taken for a derived one.""" + built = simulation(system, {"carried": {"2012": [7, 8]}}) + built.calculate("carried", "2013") + holder = built.get_holder("carried") + assert holder.is_derived(periods.period("2013")) + if delete_first: + holder.delete_arrays(periods.period("2013")) + assert not holder.is_derived(periods.period("2013")) + holder._memory_storage.put(np.array([5.0, 5.0]), periods.period("2013")) + assert not holder.is_derived(periods.period("2013")) + np.testing.assert_array_equal(built.calculate("carried", "2014"), [5, 5]) + + +def test_carry_over_in_a_branch_copies_only_the_array_it_reads(system): + """Deciding which periods are inputs reads no shared array (#556): the + branch copies only the input it carries.""" + built = simulation( + system, {"carried": {"2012": [1, 1], "2014": [2, 2], "2016": [3, 3]}} + ) + built.calculate("carried", "2017") + branch = built.get_branch("reform") + storage = branch.get_holder("carried")._memory_storage + shared = len(storage._shared) + np.testing.assert_array_equal(branch.calculate("carried", "2020"), [3, 3]) + assert len(storage._shared) == shared - 1 + + +def _marks(storage): + """A storage's own record of which values are which: an in-memory + storage records its inputs, an on-disk one its derived values.""" + return storage._derived if isinstance(storage, OnDiskStorage) else storage._inputs + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_storages_keep_the_derived_mark_with_the_value(on_disk, tmp_path): + storage = ( + OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + if on_disk + else InMemoryStorage(is_eternal=False) + ) + year = periods.period("2013") + value = np.array([1.0, 2.0]) + assert not storage.has(year) and not storage.is_derived(year) + storage.put(value, year, derived=True) + assert storage.has(year) and storage.is_derived(year) + assert not storage.has(year, "reform") and not storage.is_derived(year, "reform") + clone = storage.clone() + assert clone.is_derived(year) + storage.put(value, year) + assert storage.has(year) and not storage.is_derived(year) + assert clone.is_derived(year) + storage.put(value, year, derived=True) + storage.delete(year) + assert not storage.has(year) and not storage.is_derived(year) + assert not _marks(storage) + storage.put(value, year, derived=True) + storage.put(value, year, "reform", derived=True) + storage.delete(branch_name="reform") + assert storage.is_derived(year) and not storage.has(year, "reform") + storage.delete() + assert not _marks(storage) + storage.put(value, year) + assert not storage.is_derived(year) + + +def test_carry_over_reads_derived_marks_from_disk_storage(system): + built = simulation(system, {}) + 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 + built.set_input("carried", "2012", np.array([7.0, 8.0])) + np.testing.assert_array_equal(built.calculate("carried", "2014"), [7, 8]) + assert holder._disk_storage.is_derived(periods.period("2014")) + np.testing.assert_array_equal(built.calculate("carried", "2013"), [7, 8]) + holder.delete_arrays(periods.period("2014")) + built.set_input("carried", "2014", np.array([9.0, 10.0])) + assert not holder.is_derived(periods.period("2014")) + np.testing.assert_array_equal(built.calculate("carried", "2015"), [9, 10]) + + +def test_a_default_cached_before_an_input_does_not_become_an_uprating_base(system): + """With no input at or before 2012 the default is not cached for an + uprated variable, so 2014 still carries the 2013-01 input.""" + inputs = {"uprated_input_without_helper": {"2013-01": [10, 20]}} + expected = alone(system, inputs, "uprated_input_without_helper", "2014") + np.testing.assert_array_equal(expected, [10, 20]) + built = simulation(system, inputs) + np.testing.assert_array_equal( + built.calculate("uprated_input_without_helper", "2012"), [0, 0] + ) + np.testing.assert_array_equal( + built.calculate("uprated_input_without_helper", "2014"), expected + ) + + +def test_a_value_calculated_while_an_input_is_set_stays_on_its_own_branch(system): + """A ``set_input`` helper that calculates writes the calculated value to + the simulation's own branch, as a derived value, not to the input's.""" + built = simulation(system, {"carried": {"2012": [1, 2]}}) + holder = built.get_holder("formula_until_2013") + holder.set_input(periods.period("2012"), np.array([20.0, 20.0]), "other") + built.get_holder("input_with_calculating_helper").set_input( + periods.period("2013-06"), np.array([1.0, 1.0]), "other" + ) + np.testing.assert_array_equal( + holder.get_array(periods.period("2012"), "other"), [20, 20] + ) + np.testing.assert_array_equal( + holder.get_array(periods.period("2012"), "default"), [6, 7] + ) + assert holder.is_derived(periods.period("2012"), "default") + assert not holder.is_derived(periods.period("2012"), "other") + assert ("formula_until_2013", "default", periods.period("2012")) not in ( + built._user_input_keys + ) + other = built.get_branch("other") + np.testing.assert_array_equal( + other.calculate("formula_until_2013", "2014"), [20, 20] + ) + + +def test_carry_over_reads_provenance_in_one_pass(system, monkeypatch): + """With many calculated periods and deep branches, a decision makes one + pass over the stored keys instead of walking the branches per period.""" + from policyengine_core.holders import Holder + + built = simulation(system, {"carried": {"1900": [1, 1]}}) + for year in range(1901, 2010): + built.calculate("carried", str(year)) + branch = built + for depth in range(20): + branch = branch.get_branch(f"b{depth}") + calls = {"stores": 0, "passes": 0} + stores, passes = Holder._stores, Holder.get_input_periods + + def counting_stores(self, *args, **kwargs): + calls["stores"] += 1 + return stores(self, *args, **kwargs) + + def counting_passes(self, *args, **kwargs): + calls["passes"] += 1 + return passes(self, *args, **kwargs) + + monkeypatch.setattr(Holder, "_stores", counting_stores) + monkeypatch.setattr(Holder, "get_input_periods", counting_passes) + np.testing.assert_array_equal(branch.calculate("carried", "2020"), [1, 1]) + # One pass for the decision; caching the result checks the period on each + # branch in the chain once (21), not every stored period on each (2,310). + assert calls["passes"] == 1 + assert calls["stores"] <= 21 + + +def test_a_period_stored_only_on_another_branch_does_not_carry(system): + """A branch never takes a period it cannot read for an input.""" + built = simulation(system, {"year_input_without_helper": {"2013-01": [20, 20]}}) + built.get_holder("year_input_without_helper").set_input( + periods.period("2012"), np.array([10.0, 10.0]), "other" + ) + other = built.get_branch("other") + np.testing.assert_array_equal( + built.calculate("year_input_without_helper", "2014"), [20, 20] + ) + np.testing.assert_array_equal( + other.calculate("year_input_without_helper", "2014"), [10, 10] + ) + + +def test_marks_are_cleared_with_the_values_apply_reform_wipes(system): + built = simulation(system, {}) + 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)) + built.apply_reform(_noop) + assert not holder._memory_storage._inputs + 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.memory_config = MemoryConfig(max_memory_occupation=0) + holder = built.get_holder("carried") + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + writer = OnDiskStorage(holder._disk_storage.storage_dir, preserve_storage_dir=True) + built.calculate("carried", "2012") + assert holder.is_derived(periods.period("2012")) + built.apply_reform(_noop) + writer.put(np.array([10.0, 10.0]), periods.period("2012")) + holder._disk_storage.restore() + assert not holder.is_derived(periods.period("2012")) + np.testing.assert_array_equal(built.calculate("carried", "2013"), [10, 10]) + + +def test_marks_are_cleared_with_in_memory_values_apply_reform_wipes(system): + built = simulation(system, {}) + holder = built.get_holder("carried") + for year in range(2010, 2020): + built.calculate("carried", str(year)) + built.apply_reform(_noop) + assert not holder._memory_storage._inputs + + +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.memory_config = MemoryConfig(max_memory_occupation=0) + holder = built.get_holder("carried") + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + built.calculate("carried", "2012") + assert holder.is_derived(periods.period("2012")) + writer = OnDiskStorage(holder._disk_storage.storage_dir, preserve_storage_dir=True) + writer.put(np.array([10.0, 10.0]), periods.period("2012")) + holder._disk_storage.restore() + assert not holder.is_derived(periods.period("2012")) + np.testing.assert_array_equal(built.calculate("carried", "2013"), [10, 10]) + + +def test_a_value_in_memory_takes_precedence_over_one_on_disk(system): + """``get_array`` reads memory before disk for the same key, and so does + the input test.""" + built = simulation(system, {}) + holder = built.get_holder("carried") + holder._disk_storage = holder.create_disk_storage() + year = periods.period("2012") + holder._disk_storage.put(np.array([0.0, 0.0]), year, derived=True) + holder._memory_storage.put(np.array([7.0, 8.0]), year) + assert year in holder.get_input_periods() + np.testing.assert_array_equal(built.calculate("carried", "2013"), [7, 8]) + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +def test_storages_pickled_without_marks_still_work(on_disk, tmp_path): + import pickle + + storage = ( + OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + if on_disk + else InMemoryStorage(is_eternal=False) + ) + year = periods.period("2012") + storage.put(np.array([10.0]), year) + # The state an older version pickled: its arrays or files, and no marks. + state = dict(storage.__dict__) + for name in ("_derived", "_inputs", "_shared"): + state.pop(name, None) + old = type(storage).__new__(type(storage)) + old.__setstate__(state) + restored = pickle.loads(pickle.dumps(old)) + assert restored.has(year) and not restored.is_derived(year) + restored.clone() + restored.put(np.array([20.0]), periods.period("2013"), derived=True) + assert restored.is_derived(periods.period("2013")) + restored.delete(periods.period("2013")) + np.testing.assert_array_equal(restored.get(year), [10.0]) + + +def test_input_starting_with_the_requested_period_carries(system): + """An input in another unit that starts when the period does carries.""" + built = simulation(system, {"year_input_without_helper": {"2013-01": [3, 4]}}) + np.testing.assert_array_equal( + built.calculate("year_input_without_helper", "2013"), [3, 4] + ) + + +def test_carried_value_is_derived_and_the_input_is_not(system): + built = simulation(system, {"carried": {"2012": [7, 8]}}) + built.calculate("carried", "2013") + holder = built.get_holder("carried") + assert holder.is_derived(periods.period("2013")) + assert not holder.is_derived(periods.period("2012")) + assert not holder.is_derived(periods.period("2015")) + branch = built.get_branch("reform") + assert branch.get_holder("carried").is_derived(periods.period("2013"), "reform") + + +def test_restored_simulation_keeps_derived_marks(system, tmp_path): + inputs = {"carried": {"2012": [1, 2]}} + built = simulation(system, inputs) + np.testing.assert_array_equal(built.calculate("formula_until_2013", "2012"), [6, 7]) + dump_simulation(built, str(tmp_path / "dump")) + restored = restore_simulation(str(tmp_path / "dump"), system) + expected = alone(system, inputs, "formula_until_2013", "2014") + np.testing.assert_array_equal(expected, [0, 0]) + np.testing.assert_array_equal( + built.calculate("formula_until_2013", "2014"), expected + ) + np.testing.assert_array_equal( + restored.calculate("formula_until_2013", "2014"), expected + ) + np.testing.assert_array_equal(restored.calculate("carried", "2014"), [1, 2]) + + +class _noop(Reform): + def apply(self): + pass + + +def test_inputs_still_carry_after_apply_reform(system): + inputs = {"carried": {"2012": [7, 8]}} + built = simulation(system, inputs) + built.calculate("carried", "2014") + built.apply_reform(_noop) + holder = built.get_holder("carried") + assert not any(holder.is_derived(p) for p in holder.get_known_periods()) + np.testing.assert_array_equal(built.calculate("carried", "2013"), [7, 8]) + np.testing.assert_array_equal(built.calculate("carried", "2014"), [7, 8]) diff --git a/tests/core/test_carry_over_order_property.py b/tests/core/test_carry_over_order_property.py new file mode 100644 index 000000000..23111e262 --- /dev/null +++ b/tests/core/test_carry_over_order_property.py @@ -0,0 +1,329 @@ +"""Properties of auto-carry-over: order independence and the reference rule. + +For any inputs and any earlier requests (plain calculations in either unit, +``ADD`` and ``DIVIDE``), what a simulation then returns for a period, read +from the simulation or from a branch or nested branch forked after those +requests, equals what a fresh simulation given the same inputs returns for +that period alone, and equals the reference rule in +``tests/fixtures/carry_over.py``. Regressions are in +``test_carry_over_order.py``. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +# The smoke job installs Core without the dev extra but collects every module. +pytest.importorskip("hypothesis") + +from hypothesis import HealthCheck, example, given, settings # noqa: E402 +from hypothesis import strategies as st # noqa: E402 + +from policyengine_core import periods # noqa: E402 +from tests.fixtures.carry_over import ( # noqa: E402 + COUNT, + alone, + build_system, + reference, + request, + simulation, +) + +SYSTEM = build_system() + +YEARS = [str(year) for year in range(2010, 2017)] +MONTHS = [f"{year}-{month:02d}" for year in (2012, 2013) for month in (1, 2, 6, 12)] +MULTI_YEAR = ["year:2012:2"] + +YEAR_VARIABLES = [ + "carried", + "carried_count", + "eligible", + "carried_if_eligible", + "formula_until_2013", + "year_input_without_helper", +] +MONTH_VARIABLES = ["carried_monthly", "month_input_without_helper"] +# What each variable's result reads: itself and its formula and defined_for. +READS = { + "carried_if_eligible": ["carried_if_eligible", "eligible"], + "formula_until_2013": ["formula_until_2013", "carried"], +} + +FLOAT = st.floats(min_value=-100, max_value=100, allow_nan=False, width=32) +VALUES = { + "carried_count": st.integers(min_value=-100, max_value=100), + "eligible": st.booleans(), +} +INPUT_PERIODS = {variable: YEARS for variable in YEAR_VARIABLES} +INPUT_PERIODS["carried_monthly"] = MONTHS +# Variables with no set_input helper store inputs at any period. +INPUT_PERIODS["year_input_without_helper"] = YEARS + MONTHS + MULTI_YEAR +INPUT_PERIODS["month_input_without_helper"] = MONTHS + YEARS + + +def _shift(period, step): + period = periods.period(period) + return str(period.offset(step, period.unit)) + + +@st.composite +def scenarios(draw): + inputs = {} + for variable, input_periods in INPUT_PERIODS.items(): + values = draw( + st.dictionaries( + st.sampled_from(input_periods), + st.lists(VALUES.get(variable, FLOAT), min_size=COUNT, max_size=COUNT), + max_size=3, + ) + ) + if values: + inputs[variable] = values + variable = draw(st.sampled_from(YEAR_VARIABLES + MONTH_VARIABLES)) + unit_periods = YEARS if variable in YEAR_VARIABLES else MONTHS + period = draw(st.sampled_from(unit_periods)) + # Bias requests toward what the target reads and toward periods next to + # the inputs and the target: each bug needs a short chain of those. + nearby = {period} + for name in READS.get(variable, [variable]): + nearby.update(inputs.get(name, {})) + nearby.update( + _shift(known, step) + for known in list(nearby) + if not known.startswith("year:") + for step in (-1, 1) + ) + nearby = sorted(nearby) + request_variables = st.one_of( + st.sampled_from(READS.get(variable, [variable])), + st.sampled_from(YEAR_VARIABLES + MONTH_VARIABLES), + ) + + @st.composite + def one_request(draw): + name = draw(request_variables) + kind = draw(st.sampled_from(["calculate", "calculate", "add", "divide"])) + candidates = nearby + YEARS + MONTHS + if name in YEAR_VARIABLES: + candidates += MULTI_YEAR + requested = draw(st.sampled_from(candidates)) + unit = periods.period(requested).unit + if kind == "divide" and ( + name not in YEAR_VARIABLES or periods.period(requested).size != 1 + ): + kind = "calculate" + if name in MONTH_VARIABLES and requested.startswith("year:"): + kind = "add" + if kind == "add" and name in YEAR_VARIABLES and unit == periods.MONTH: + kind = "calculate" + if kind == "calculate" and requested.startswith("year:"): + kind = "add" + return kind, name, requested + + requests = draw(st.lists(one_request(), max_size=8)) + branch = draw(st.sampled_from([None, "reform", "nested"])) + branch_requests = draw(st.lists(one_request(), max_size=4)) if branch else [] + return inputs, requests, branch, branch_requests, (variable, period) + + +@settings( + max_examples=400, + deadline=None, + derandomize=True, + suppress_health_check=[HealthCheck.too_slow], +) +@given(scenario=scenarios()) +# One example per regression shape in test_carry_over_order.py. +@example( + scenario=( + {"carried": {"2012": [7, 8]}}, + [("calculate", "carried", "2014")], + None, + [], + ("carried", "2013"), + ) +) +@example( + scenario=( + { + "carried_if_eligible": {"2012": [7, 8]}, + "eligible": {"2013": [False, True], "2014": [True, True]}, + }, + [("calculate", "carried_if_eligible", "2013")], + None, + [], + ("carried_if_eligible", "2014"), + ) +) +@example( + scenario=( + { + "carried_if_eligible": {"2012": [7, 8]}, + "eligible": {"2013": [False, False], "2014": [True, True]}, + }, + [("add", "carried_if_eligible", "2013")], + "reform", + [], + ("carried_if_eligible", "2014"), + ) +) +@example( + scenario=( + {"carried": {"2012": [1, 2]}}, + [("calculate", "formula_until_2013", "2013")], + None, + [], + ("formula_until_2013", "2014"), + ) +) +@example( + scenario=( + {"year_input_without_helper": {"year:2012:2": [120, 24]}}, + [("add", "year_input_without_helper", "year:2012:2")], + None, + [], + ("year_input_without_helper", "2014"), + ) +) +@example( + scenario=( + {"year_input_without_helper": {"2012-01": [120, 24]}}, + [("divide", "year_input_without_helper", "2012-01")], + "nested", + [], + ("year_input_without_helper", "2014"), + ) +) +@example( + scenario=( + {"month_input_without_helper": {"2013": [120, 24]}}, + [("add", "month_input_without_helper", "2013")], + None, + [], + ("month_input_without_helper", "2013-06"), + ) +) +@example( + scenario=( + {"year_input_without_helper": {"2013-01": [2, 2], "2013": [1, 1]}}, + [], + None, + [], + ("year_input_without_helper", "2014"), + ) +) +@example( + scenario=( + {"year_input_without_helper": {"2013-01": [3, 4]}}, + [], + None, + [], + ("year_input_without_helper", "2013"), + ) +) +@example( + scenario=( + {"year_input_without_helper": {"year:2012:2": [120, 24]}}, + [], + "reform", + [("add", "year_input_without_helper", "year:2012:2")], + ("year_input_without_helper", "2014"), + ) +) +def test_carry_over_depends_only_on_the_inputs(scenario): + inputs, requests, branch, branch_requests, (variable, period) = scenario + built = simulation(SYSTEM, inputs) + for kind, name, requested in requests: + request(built, kind, name, requested) + if branch == "reform": + built = built.get_branch("reform") + elif branch == "nested": + built = built.get_branch("reform").get_branch("nested") + for kind, name, requested in branch_requests: + request(built, kind, name, requested) + result = built.calculate(variable, period) + + message = ( + f"{variable} {period} after {requests}, then {branch_requests} on {branch}" + ) + np.testing.assert_array_equal( + result, alone(SYSTEM, inputs, variable, period), err_msg=message + ) + np.testing.assert_array_equal( + result, reference(SYSTEM, inputs, variable, period), err_msg=message + ) + + +# ----- The order of period ends ----- # + +from datetime import date # noqa: E402 + +from policyengine_core.simulations.simulation import _end_order # noqa: E402 + +_starts = st.dates(min_value=date(1, 1, 1), max_value=date(9999, 12, 31)) +_units = st.sampled_from([periods.DAY, periods.MONTH, periods.YEAR]) + + +def _period(unit, start, size): + return periods.Period( + (unit, periods.Instant((start.year, start.month, start.day)), size) + ) + + +def _numpy_day_after(period): + """The day after the period's last day, on numpy's calendar, which runs + past year 9999: an implementation that shares nothing with + ``_end_order`` or ``Period.stop``.""" + unit, (year, month, day), size = period + if unit == periods.DAY: + return np.datetime64(f"{year:04d}-{month:02d}-{day:02d}") + np.timedelta64( + size, "D" + ) + months = size if unit == periods.MONTH else 12 * size + first = np.datetime64(f"{year:04d}-{month:02d}") + np.timedelta64(months, "M") + return first.astype("datetime64[D]") + np.timedelta64(day - 1, "D") + + +@settings(max_examples=1000, deadline=None, derandomize=True) +@given(unit=_units, start=_starts, size=st.integers(min_value=1, max_value=5_000_000)) +@example(unit=periods.DAY, start=date(9999, 12, 30), size=3) +@example(unit=periods.MONTH, start=date(2012, 1, 31), size=1) +@example(unit=periods.YEAR, start=date(2012, 2, 29), size=1) +@example(unit=periods.MONTH, start=date(9999, 2, 1), size=24) +def test_end_order_is_the_day_after_the_last_day(unit, start, size): + """Against two other implementations: ``Period.stop`` where it has a + value, and numpy's calendar everywhere.""" + period = _period(unit, start, size) + day_one = np.datetime64("0001-01-01") + expected = int((_numpy_day_after(period) - day_one) / np.timedelta64(1, "D")) + 1 + assert _end_order(period) == expected + try: + stop = date(*period.stop) + except (OverflowError, ValueError): + # Ends after 9999-12-31, the last date ``datetime`` has: ``stop`` + # raises, or gives a year no ``date`` can hold. + assert expected > date.max.toordinal() + 1 + else: + assert _end_order(period) == stop.toordinal() + 1 + + +@settings(max_examples=500, deadline=None, derandomize=True) +@given( + unit=_units, + start=_starts, + size=st.integers(min_value=1, max_value=5_000_000), + more=st.integers(min_value=1, max_value=1000), +) +def test_a_longer_period_from_the_same_day_ends_later(unit, start, size, more): + assert _end_order(_period(unit, start, size + more)) > _end_order( + _period(unit, start, size) + ) + + +def test_eternity_ends_after_every_other_period(): + eternity = periods.period(periods.ETERNITY) + assert _end_order(eternity) > _end_order( + _period(periods.YEAR, date(9999, 1, 1), 5_000_000) + ) diff --git a/tests/core/test_storage_input_index.py b/tests/core/test_storage_input_index.py new file mode 100644 index 000000000..ffbd75bcd --- /dev/null +++ b/tests/core/test_storage_input_index.py @@ -0,0 +1,285 @@ +"""A storage records which of its values are inputs, at no cost when none is. + +``InMemoryStorage.is_derived`` tells a value the simulation calculated +(``put(..., derived=True)``) from an input, so that auto-carry-over and +uprating start only from inputs. The storage records the inputs, in +``_inputs``, because a simulation calculates far more values than it is +given. Like ``_shared`` (``test_storage_shared_index.py``), ``_inputs`` is +one empty object for every storage until a storage holds an input: a +simulation has a storage for every variable of the tax-benefit system, and +an empty set in each is 216 bytes per variable per simulation (#577). A +storage and its clones share one frozenset of inputs until one of them +changes its own. ``test_storage_shared_index_differential.py`` checks random +sequences of operations against a storage that always had its own sets. +""" + +from __future__ import annotations + +import copy +import pickle + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.data_storage import InMemoryStorage +from policyengine_core.data_storage.in_memory_storage import ( + _INPUTS_RECORDED, + _NO_INPUTS, + _NOTHING_SHARED, +) +from policyengine_core.reforms import Reform +from tests.fixtures.branch_shared_arrays import build_simulation + +JANUARY = periods.period("2017-01") + + +def _storages(simulation): + return { + name: holder._memory_storage + for population in simulation.populations.values() + for name, holder in population._holders.items() + } + + +def test_the_no_inputs_index_is_empty_and_cannot_change(): + assert isinstance(_NO_INPUTS, frozenset) + assert not _NO_INPUTS + + +def test_storing_reading_and_deleting_derived_values_never_needs_a_set(): + storage = InMemoryStorage(is_eternal=False) + assert storage._inputs is _NO_INPUTS + storage.put(np.array([1.0]), "2017-01", derived=True) + assert storage.is_derived("2017-01") + storage.get("2017-01") + storage.delete("2017-01") + storage.put(np.array([1.0]), "2017-02", derived=True) + storage.delete() + assert storage._inputs is _NO_INPUTS + assert "_inputs" not in vars(storage) + + +@pytest.mark.parametrize( + "drop", + [ + lambda storage: storage.put(np.array([2.0]), "2017-01", derived=True), + lambda storage: storage.delete("2017-01"), + lambda storage: storage.delete("2017"), + lambda storage: storage.delete(), + ], + ids=[ + "replaced_by_a_derived_value", + "deleted", + "deleted_by_a_containing_period", + "wiped", + ], +) +def test_an_input_gets_a_set_that_is_released_with_it(drop): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0]), "2017-01") + assert storage._inputs == {"default:2017-01"} + assert "_inputs" in vars(storage) + assert not storage.is_derived("2017-01") + + drop(storage) + # Released, not replaced by another empty object. + assert "_inputs" not in vars(storage) + assert storage._inputs is _NO_INPUTS + + +def test_the_set_is_kept_while_another_value_is_an_input(): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0]), "2017-01") + storage.put(np.array([1.0]), "2017-02") + storage.delete("2017-01") + assert storage._inputs == {"default:2017-02"} + storage.put(np.array([3.0]), "2017-02", derived=True) + assert storage._inputs is _NO_INPUTS + assert storage.is_derived("2017-02") + + +def test_only_stored_values_are_derived(): + storage = InMemoryStorage(is_eternal=False) + assert not storage.is_derived("2017-01") + storage.put(np.array([1.0]), "2017-01", "reform", derived=True) + assert not storage.is_derived("2017-01") + assert storage.is_derived("2017-01", "reform") + + +@pytest.mark.parametrize("share_arrays", [False, True], ids=["copy", "share"]) +def test_clones_share_the_inputs_until_one_of_them_changes_its_own(share_arrays): + derived_only = InMemoryStorage(is_eternal=False) + derived_only.put(np.array([1.0]), "2017-01", derived=True) + assert derived_only.clone(share_arrays)._inputs is _NO_INPUTS + + source = InMemoryStorage(is_eternal=False) + source.put(np.array([1.0]), "2017-01") + source.put(np.array([1.0]), "2017-02") + first = source.clone(share_arrays) + second = source.clone(share_arrays) + grandchild = first.clone(share_arrays) + # One immutable set for all four: a clone costs no set of its own. + inputs = source._inputs + assert isinstance(inputs, frozenset) + assert inputs == {"default:2017-01", "default:2017-02"} + assert first._inputs is inputs and second._inputs is inputs + assert grandchild._inputs is inputs + + # Storing an input over an input changes nothing. + first.put(np.array([5.0]), "2017-01") + assert first._inputs is inputs + + # The first storage to change its inputs gets a set of its own. + first.put(np.array([2.0]), "2017-01", derived=True) + assert first.is_derived("2017-01") and not first.is_derived("2017-02") + assert isinstance(first._inputs, set) + for other in (source, second, grandchild): + assert other._inputs is inputs and not other.is_derived("2017-01") + + source.put(np.array([3.0]), "2017-03") + assert source._inputs == inputs | {"default:2017-03"} + second.put(np.array([3.0]), "2017-03", derived=True) + assert second.is_derived("2017-03") and not source.is_derived("2017-03") + second.delete("2017-02") + assert second._inputs == {"default:2017-01"} + assert grandchild._inputs is inputs + grandchild.delete() + assert grandchild._inputs is _NO_INPUTS + assert not source.is_derived("2017-01") and not source.is_derived("2017-02") + + +def test_clone_of_an_emptied_storage_gets_only_the_inputs_it_holds(): + source = InMemoryStorage(is_eternal=False) + source.put(np.array([1.0]), "2017-01") + source._arrays.clear() # as policyengine-uk clears a dropped branch + assert source.clone()._inputs is _NO_INPUTS + source.put(np.array([1.0]), "2017-02") + clone = source.clone() + assert clone._inputs == {"default:2017-02"} + assert isinstance(clone._inputs, frozenset) + + +@pytest.mark.parametrize( + "duplicate", + [copy.deepcopy, lambda storage: pickle.loads(pickle.dumps(storage))], + ids=["deepcopy", "pickle"], +) +def test_copied_or_unpickled_storage_keeps_the_marks(duplicate): + derived = InMemoryStorage(is_eternal=False) + derived.put(np.array([1.0]), "2017-01", derived=True) + twin = duplicate(derived) + assert twin._inputs is _NO_INPUTS + assert twin._shared is _NOTHING_SHARED + assert twin.is_derived("2017-01") + + for cloned_first in (False, True): + inputs = InMemoryStorage(is_eternal=False) + inputs.put(np.array([1.0]), "2017-01") + if cloned_first: + inputs.clone() # its inputs are now a frozenset it shares + twin = duplicate(inputs) + assert not twin.is_derived("2017-01") + # Changing either one's inputs leaves the other's alone. + twin.put(np.array([2.0]), "2017-01", derived=True) + assert twin.is_derived("2017-01") + assert not inputs.is_derived("2017-01") + inputs.put(np.array([2.0]), "2017-02") + assert not twin.has("2017-02") and not twin.is_derived("2017-02") + + +def _older_state(storage, **extra): + """The state a version that recorded no inputs pickled.""" + state = {"_arrays": dict(storage._arrays), "is_eternal": storage.is_eternal} + state.update(extra) + return state + + +def test_values_from_an_older_pickle_count_as_inputs(): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0]), "2017-01", derived=True) + storage.put(np.array([1.0]), "2017-02", derived=True) + restored = InMemoryStorage.__new__(InMemoryStorage) + # 3.32.12 pickled an empty set of shared keys in every storage. + restored.__setstate__(_older_state(storage, _shared=set())) + assert not restored.is_derived("2017-01") + assert not restored.is_derived("2017-02") + assert restored._inputs == {"default:2017-01", "default:2017-02"} + assert _INPUTS_RECORDED not in vars(restored) + + # A development version of this change recorded derived values instead. + restored = InMemoryStorage.__new__(InMemoryStorage) + restored.__setstate__(_older_state(storage, _derived={"default:2017-02"})) + assert not restored.is_derived("2017-01") + assert restored.is_derived("2017-02") + assert "_derived" not in vars(restored) + # A state with no set of shared keys gets none. + assert "_shared" not in vars(restored) + + # An older storage with nothing stored needs no set. + restored = InMemoryStorage.__new__(InMemoryStorage) + restored.__setstate__(_older_state(InMemoryStorage(is_eternal=False))) + assert restored._inputs is _NO_INPUTS + assert restored._shared is _NOTHING_SHARED + restored.put(np.array([2.0]), "2017-03", derived=True) + assert restored.is_derived("2017-03") + + +def test_a_pickle_records_that_the_inputs_were_recorded(): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0]), "2017-01", derived=True) + state = storage.__getstate__() + assert state[_INPUTS_RECORDED] is True + assert _INPUTS_RECORDED not in vars(storage) + assert pickle.loads(pickle.dumps(storage)).is_derived("2017-01") + + +def test_only_storages_holding_an_input_have_a_set(tax_benefit_system): + simulation = build_simulation(tax_benefit_system) + simulation.calculate("disposable_income", JANUARY) + + storages = _storages(simulation) + assert len(storages) == len(tax_benefit_system.variables) + with_inputs = {name for name, storage in storages.items() if storage._inputs} + with_derived = { + name + for name, storage in storages.items() + if any( + storage.is_derived(key.split(":", 1)[1], key.split(":", 1)[0]) + for key in storage._arrays + ) + } + # The situation's inputs, and the values calculated from them. + assert with_inputs and with_derived + assert with_inputs | with_derived < set(storages) + for name, storage in storages.items(): + assert ("_inputs" in vars(storage)) == (name in with_inputs), name + if name not in with_inputs: + assert storage._inputs is _NO_INPUTS, name + assert storage._shared is _NOTHING_SHARED, name + + +def test_a_branch_shares_its_parents_inputs(tax_benefit_system): + simulation = build_simulation(tax_benefit_system) + simulation.calculate("disposable_income", JANUARY) + branch = simulation.get_branch("branch") + parents = _storages(simulation) + assert any(storage._inputs for storage in parents.values()) + for name, storage in _storages(branch).items(): + assert storage._inputs is parents[name]._inputs, name + + +def test_apply_reform_keeps_only_the_replayed_inputs(tax_benefit_system): + class _noop(Reform): + def apply(self): + pass + + simulation = build_simulation(tax_benefit_system) + simulation.calculate("disposable_income", JANUARY) + before = { + name: set(storage._inputs) for name, storage in _storages(simulation).items() + } + simulation.apply_reform(_noop) + for name, storage in _storages(simulation).items(): + assert set(storage._inputs) == before.get(name, set()), name + assert set(storage._arrays) == set(storage._inputs), name diff --git a/tests/core/test_storage_shared_index_differential.py b/tests/core/test_storage_shared_index_differential.py index 792ec7a34..343229466 100644 --- a/tests/core/test_storage_shared_index_differential.py +++ b/tests/core/test_storage_shared_index_differential.py @@ -1,12 +1,16 @@ """A storage without a set of its own behaves like one that always had one. ``_EagerIndexStorage`` below is ``InMemoryStorage`` as it was when every -storage owned a set of shared keys. Random sequences of operations run on -both, and after every step the two must hold the same arrays and the same -shared keys; at the end every read must agree. The storage under test must -also keep its own invariants: it refers to the one shared-nothing object -exactly when it shares nothing, and no two storages hold the same set. -``test_storage_shared_index.py`` pins the same behaviour with examples. +storage owned a set of shared keys and a set of derived keys. The storage +under test records its inputs instead, and has each set only while it is +not empty. Random sequences of operations, storing inputs and derived +values, run on both, and after every step the two must hold the same arrays, +the same shared keys, and the same stored values marked derived; at the end +every read and every ``is_derived`` must agree. The storage under test must +also keep its own invariants: it refers to the one shared-nothing (no-inputs) +object exactly when it shares nothing (holds no input), and no two storages +hold the same mutable set. ``test_storage_shared_index.py`` and +``test_storage_input_index.py`` pin the same behaviour with examples. """ from __future__ import annotations @@ -21,6 +25,7 @@ from policyengine_core import periods from policyengine_core.data_storage import InMemoryStorage from policyengine_core.data_storage.in_memory_storage import ( + _NO_INPUTS, _NOTHING_SHARED, _can_share, _read_only_view, @@ -28,11 +33,13 @@ class _EagerIndexStorage: - """The storage with an empty set of shared keys in every instance.""" + """The storage with an empty set of shared keys and an empty set of + derived keys in every instance.""" def __init__(self, is_eternal): self._arrays = {} self._shared = set() + self._derived = set() self.is_eternal = is_eternal def clone(self, share_arrays=False): @@ -46,8 +53,15 @@ def clone(self, share_arrays=False): clone._arrays[key] = array.copy() else: clone._arrays = {key: array.copy() for key, array in self._arrays.items()} + clone._derived = set(self._derived) return clone + def is_derived(self, period, branch_name="default"): + if self.is_eternal: + period = periods.period(periods.ETERNITY) + key = f"{branch_name}:{periods.period(period)}" + return key in self._derived and key in self._arrays + def get(self, period, branch_name="default"): if self.is_eternal: period = periods.period(periods.ETERNITY) @@ -62,13 +76,17 @@ def get(self, period, branch_name="default"): self._shared.discard(key) return values - def put(self, value, period, branch_name="default"): + def put(self, value, period, branch_name="default", derived=False): if self.is_eternal: period = periods.period(periods.ETERNITY) period = periods.period(period) key = f"{branch_name}:{period}" self._arrays[key] = value self._shared.discard(key) + if derived: + self._derived.add(key) + else: + self._derived.discard(key) def delete(self, period=None, branch_name="default"): if period is None: @@ -79,6 +97,7 @@ def delete(self, period=None, branch_name="default"): if not period_item.startswith(branch_prefix) } self._shared.intersection_update(self._arrays) + self._derived.intersection_update(self._arrays) return if self.is_eternal: period = periods.period(periods.ETERNITY) @@ -92,6 +111,7 @@ def delete(self, period=None, branch_name="default"): ) } self._shared.intersection_update(self._arrays) + self._derived.intersection_update(self._arrays) PERIODS = ["2017-01", "2017-02", "2017"] @@ -121,6 +141,7 @@ def _make_array(kind, seed): _branch, st.sampled_from(ARRAY_KINDS), _seed, + st.booleans(), ), st.tuples(st.just("get"), _storage, _period, _branch), st.tuples(st.just("write"), _storage, _period, _branch, _seed), @@ -147,8 +168,8 @@ def _apply(storages, operation): kind, index = operation[0], operation[1] % len(storages) storage = storages[index] if kind == "put": - _, _, period, branch, array_kind, seed = operation - storage.put(_make_array(array_kind, seed), period, branch) + _, _, period, branch, array_kind, seed, derived = operation + storage.put(_make_array(array_kind, seed), period, branch, derived=derived) return None if kind == "get": _, _, period, branch = operation @@ -181,6 +202,12 @@ def _assert_same_state(tested, reference): for key, array in storage._arrays.items(): assert _same_array(array, expected._arrays[key]), (index, key) assert storage._shared == expected._shared, index + # The stored values marked derived. A storage emptied from outside + # keeps the marks of keys it no longer holds, which the reference + # copies into clones; they never count. + assert storage._arrays.keys() - storage._inputs == ( + expected._derived & expected._arrays.keys() + ), index def _assert_index_invariants(storages, cleared): @@ -198,9 +225,19 @@ def _assert_index_invariants(storages, cleared): for key in shared & storage._arrays.keys(): # A shared array cannot be written through. assert not storage._arrays[key].flags.writeable, (index, key) - # A set belongs to one storage. - assert len({id(shared) for shared in own_sets}) == len(own_sets) + inputs = storage._inputs + assert (inputs is _NO_INPUTS) == (not inputs), index + assert ("_inputs" in vars(storage)) == bool(inputs), index + # Inputs are a set of the storage's own, or a frozenset it may share + # with storages cloned from it or it from them. + if isinstance(inputs, set): + own_sets.append(inputs) + if index not in cleared: + assert inputs <= storage._arrays.keys(), index + # A set belongs to one storage, and to one of its two indexes. + assert len({id(own) for own in own_sets}) == len(own_sets) assert not _NOTHING_SHARED + assert not _NO_INPUTS @hypothesis.settings( @@ -230,6 +267,9 @@ def test_storage_without_its_own_set_matches_storage_with_one(is_eternal, operat for index, (storage, expected) in enumerate(zip(tested, reference)): for branch in BRANCHES: for period in PERIODS: + assert storage.is_derived(period, branch) == expected.is_derived( + period, branch + ), (index, branch, period) assert _same_array( storage.get(period, branch), expected.get(period, branch) ), (index, branch, period) diff --git a/tests/fixtures/carry_over.py b/tests/fixtures/carry_over.py new file mode 100644 index 000000000..c68c08047 --- /dev/null +++ b/tests/fixtures/carry_over.py @@ -0,0 +1,219 @@ +"""Variables, simulations and a reference rule for the auto-carry-over tests. + +Shared by ``tests/core/test_carry_over_order.py`` (regressions) and +``tests/core/test_carry_over_order_property.py`` (properties). + +The reference rule, for a variable with no formula result for a period: an +input stored for the period itself is read back as stored; otherwise the +period takes the input stored for the latest-starting period that starts no +later than it (on a tie, the one that ends last, then the larger unit), +preferring inputs at the +variable's own definition-period unit, masked by the variable's +``defined_for`` at the period; with no such input, the default. +""" + +from __future__ import annotations + +import numpy as np + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem, entities +from policyengine_core.parameters import ParameterNode +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables import Variable + +COUNT = 2 +FORMULA_END_YEAR = 2013 + + +class carried(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "Yearly flow input with no formula" + + +class carried_count(Variable): + value_type = int + entity = entities.Person + definition_period = periods.YEAR + label = "Yearly integer input with no formula" + + +class eligible(Variable): + value_type = bool + entity = entities.Person + definition_period = periods.YEAR + label = "Yearly boolean input that carried_if_eligible is defined for" + + +class carried_if_eligible(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + defined_for = "eligible" + label = "Yearly flow input with no formula, defined for eligible people" + + +class formula_until_2013(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + end = f"{FORMULA_END_YEAR}-12-31" + label = "Yearly variable whose formula ends after 2013" + + def formula(person, period, parameters): + return person("carried", period) + 5 + + +class carried_monthly(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "Monthly flow input with no formula" + + +class year_input_without_helper(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + set_input = None + label = "Yearly input with no set_input helper, stored at any period" + + +class month_input_without_helper(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + set_input = None + label = "Monthly input with no set_input helper, stored at any period" + + +class day_input_without_helper(Variable): + value_type = float + entity = entities.Person + definition_period = periods.DAY + set_input = None + label = "Daily input with no set_input helper, stored at any period" + + +class uprated_input_without_helper(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + set_input = None + uprating = "carry_over_test.index" + label = "Yearly uprated input with no set_input helper" + + +def _calculate_while_setting(holder, period, array): + """A ``set_input`` helper that calculates another variable first.""" + holder.simulation.calculate("formula_until_2013", "2012") + holder._set(period.this_year, array) + + +class input_with_calculating_helper(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + set_input = _calculate_while_setting + label = "Yearly input whose set_input helper calculates" + + +VARIABLES = ( + carried, + carried_count, + eligible, + carried_if_eligible, + formula_until_2013, + carried_monthly, + year_input_without_helper, + month_input_without_helper, + day_input_without_helper, + uprated_input_without_helper, + input_with_calculating_helper, +) + + +def build_system() -> CountryTaxBenefitSystem: + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = True + system.parameters.add_child( + "carry_over_test", + ParameterNode( + "carry_over_test", data={"index": {"values": {"2010-01-01": 1.0}}} + ), + ) + system.add_variables(*VARIABLES) + return system + + +def simulation(system, inputs): + """A simulation of ``COUNT`` people given ``{variable: {period: values}}``.""" + built = SimulationBuilder().build_default_simulation(system, count=COUNT) + for variable, values in inputs.items(): + for period, array in values.items(): + built.set_input(variable, period, np.array(array)) + return built + + +def alone(system, inputs, variable, period): + """What a fresh simulation given ``inputs`` returns for one request.""" + return simulation(system, inputs).calculate(variable, period) + + +def request(built, kind, variable, period): + if kind == "add": + built.calculate_add(variable, period) + elif kind == "divide": + built.calculate_divide(variable, period) + else: + built.calculate(variable, period) + + +def _stored_inputs(system, inputs, variable): + """The inputs ``set_input`` stores: it ignores those after the ``end``.""" + end = system.get_variable(variable).end + return { + periods.period(stored): values + for stored, values in inputs.get(variable, {}).items() + if end is None or periods.period(stored).start.date <= end + } + + +def reference(system, inputs, variable, period): + """The carry-over rule, computed from the inputs alone. + + ``period`` is in the variable's definition period. + """ + period = periods.period(period) + default = system.get_variable(variable).default_array(COUNT) + stored_inputs = _stored_inputs(system, inputs, variable) + if period in stored_inputs: + # An input for the period itself is read back as stored: neither a + # formula nor ``defined_for`` applies to it. + return np.array(stored_inputs[period], dtype=default.dtype) + if variable == "formula_until_2013" and period.start.year <= FORMULA_END_YEAR: + return (reference(system, inputs, "carried", period) + 5).astype(default.dtype) + earlier = [stored for stored in stored_inputs if stored.start <= period.start] + own_unit = [ + stored + for stored in earlier + if stored.unit == system.get_variable(variable).definition_period + ] + if earlier: + latest = max( + own_unit or earlier, + key=lambda stored: ( + stored.start, + stored.stop, + periods.unit_weight(stored.unit), + str(stored), + ), + ) + value = np.array(stored_inputs[latest]).astype(default.dtype) + else: + value = default + if variable == "carried_if_eligible": + value = np.where(reference(system, inputs, "eligible", period), value, default) + return value