diff --git a/changelog.d/branch-set-input-invalidation.added.md b/changelog.d/branch-set-input-invalidation.added.md new file mode 100644 index 000000000..b173c9e98 --- /dev/null +++ b/changelog.d/branch-set-input-invalidation.added.md @@ -0,0 +1 @@ +Simulation.drop_computed_arrays() deletes every value a simulation holds except inputs, for branches whose policy changes after they are created. diff --git a/changelog.d/branch-set-input-invalidation.changed.md b/changelog.d/branch-set-input-invalidation.changed.md new file mode 100644 index 000000000..2cc17f5d5 --- /dev/null +++ b/changelog.d/branch-set-input-invalidation.changed.md @@ -0,0 +1 @@ +Branch input changes prevent caching stale calculations and rerun outermost `calculate` and `calculate_add` calls up to ten times, with exact repeated input transitions and results in other branches able to settle under the usual cache restrictions; unchanged direct inputs skip invalidation, branch input stores disable macro-cache reads, `apply_reform` preserves surviving supplied input tiers, branch input helpers share overlapping-cache cleanup with ordinary input helpers, and restored calculated values are invalidated conservatively. diff --git a/changelog.d/branch-set-input-invalidation.fixed.md b/changelog.d/branch-set-input-invalidation.fixed.md new file mode 100644 index 000000000..390794265 --- /dev/null +++ b/changelog.d/branch-set-input-invalidation.fixed.md @@ -0,0 +1 @@ +Branch input changes invalidate derived values while preserving supplied inputs, independent caches, guarded retry effects, fixed-point settlement, branch dumps and disk invalidation order, with read-only foreign calculations reused within each attempt only while their captured input and branch-registration dependencies remain current. diff --git a/docs/_toc.yml b/docs/_toc.yml index c909f933e..7835e4b66 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -7,6 +7,7 @@ parts: - caption: Using PolicyEngine Core chapters: - file: usage/simulation + - file: usage/branches - file: usage/country - file: usage/cli - file: usage/parameters diff --git a/docs/usage/branches.md b/docs/usage/branches.md new file mode 100644 index 000000000..1465c01e7 --- /dev/null +++ b/docs/usage/branches.md @@ -0,0 +1,310 @@ +# Branches + +A branch is a copy of a simulation that calculates under different inputs or +policy without changing the simulation it came from. Formulas use branches to +compare alternatives (for example, tax liability if itemizing), and marginal +tax rates use them to recalculate with slightly higher earnings. + +```python +branch = simulation.get_branch("itemizing") +branch.set_input("tax_unit_itemizes", 2026, itemizes) +tax_if_itemizing = branch.calculate("income_tax", 2026) +``` + +## What a branch starts with + +The first `get_branch(name)` call returns a new branch holding every value its +parent holds at that moment: inputs and calculated values alike. Later calls +with the same name return that branch as it is. What the parent stores after +the branch is created does not reach the branch, and what the branch stores +does not reach the parent. A branch of a branch reads its own values first, +then its parent's, then those of the simulation at the root. + +## Inputs set on a branch + +`set_input` on a branch stores the new value and drops every value the branch +holds that may have been calculated from the value it replaces. What the branch +calculates next therefore uses the input, as a simulation given the input +before calculating anything would, whatever its parent had calculated before +the branch was created. + +To decide what to drop, every stored value carries a sequence number from one +counter shared by the whole process, so a larger number means a later store. +Each simulation also keeps a store history: for each variable and period, the +number of the first store its values may have been calculated from, and for +each variable, the first time one of its values was uprated or carried over +from another period. The history records: + +- the simulation's own stores, including values a holder calculates but does + not keep (`variables_to_drop`, the cache blacklist), values read from the + macro cache, and the default a spiral returns; +- the number from which values may have been calculated from anything: the + first value read from the macro cache (what it was calculated from was + never calculated here) and the values restored from a dump (see below), so + an input for any variable drops them; +- a copy of its parent's history, taken when the branch is created; +- the history of any simulation its formulas calculate in, taken in each time + `calculate` there returns or raises (a formula that branches, sets an input + and calculates in the branch hands back values calculated there); a branch + calculated from a thread with no formula context hands its history to every + ancestor with a calculation running. + +When `set_input(variable, period, value)` is called on a branch, the branch +drops each value it holds, other than an input, whose number is at least the +earliest recorded store of the variable for a period that shares a day with +`period`, or the earliest recorded uprated or carried-over value of the +variable. It then forgets the records numbered from there on (its remaining +values were all stored earlier) and records again the inputs it keeps, unless a +formula is still running in the branch: such a formula may hold, in its own +variables, a value it read before the drop, so the records stay. Each drop +counts as an input change of the branch. A calculation running in the branch +then (one whose formula sets an input, say) may have read the replaced value, +so its result is kept neither in storage nor in the macro cache. Neither is any +result calculated from a value the branch gave before the change, in any +simulation: each calculation notes, for every other simulation it got a value +from (directly or through the calculations it called), that simulation's count +of input changes when the value's calculation began, and keeps its own result +only if none has changed since. Nested calculations in the same simulation +share this record while it is current. After it becomes stale, later nested +calculations can still keep results that use only that simulation's own values. +During these stale attempts, results that read another simulation, change inputs, +write raw cache values, create branches, or observe untracked activity remain +absent from holder storage, the fast cache and the macro cache so the caller's +retry observes the same reads and repeats those operations. Read-only results +that read another simulation can be reused temporarily within the same attempt. +Reuse checks the input-change and input-store counts in the calculating +simulation and every simulation read, the registrations of named branches, and +whether activity occurred without the formula's context. Named branch +registrations, including registered ancestor links, are captured when the +value's calculation begins or a cached value is read. Existing `get_branch` +lookups also capture the name resolved, including aliases whose names differ +from the branch object's own name. Replacing a registration while the producing +formula runs prevents reuse of its earlier read. Nested calculations and reused +results pass these captured registrations and the original dependency reads +back to the caller, so a later input change still +makes that caller stale. Observed mutations, including deletes and raw cache +writes, and branch creations clear this temporary reuse; mutations, creations +and untracked activity within a calculation prevent its result from entering it. +The temporary results are discarded when the outer attempt restarts or ends. + +Within a stable phase of one stale attempt, each distinct variable and period in +an eligible read-only foreign-read dependency graph needs at most one formula +evaluation per simulation. Formula calls therefore grow at most linearly with +the graph's distinct calculations, even when many calculations share the same +descendants. Dependency validation and replay still visit their recorded reads. +This reuse retains the result arrays and dependency records in memory until +invalidated entries are removed, the attempt restarts or the outer calculation +ends; it does not make those values part of the simulation's stored values. + +Raw writes through `Holder.put_in_cache` without `derived=True` count as mutation +activity too, including writes from threads without the formula's context, while +remaining excluded from the supplied-input record used for reform replay and +input export. So a parent formula +calculating in a branch whose formula calls back into the parent and then changes +the branch's input does not keep what it got, nor does a +formula that read a branch and then calculated there something that changed the +branch's input, or set an input on it directly. A value calculated in a thread +started without the formula's context counts too: with no calculation above it +in that thread, every calculation running in the simulation's ancestors notes +it (`asyncio.to_thread` and `contextvars.copy_context` keep the context, so the +formula notes it directly). A calculation whose call into another simulation +settled before returning keeps its result, and unrelated simulations keep +caching. The outermost calculation running in each simulation that does not +keep its result (a `calculate`, or a direct `calculate_add`, whose terms run +within it) then runs again, inner calculations included, until a run reads +nothing that has changed since, and keeps that result, as a simulation given +the new inputs first would. After ten reruns it stops: a formula that keeps +changing inputs returns its last result without keeping it. The budget is per +simulation, so calculations nested across several such simulations can rerun +more. + +An outermost `calculate` or `calculate_add` can also settle a repeated input +transition in branches other than the simulation doing the calculation. Two +consecutive stale attempts must make exactly the same ordered input stores and +deletions and return exactly the same numeric result. Each store is compared +after conversion to the variable's stored dtype; comparisons include dtype, +shape, enum type and floating-point sign bits. Mutations identify the variable, +period and branch. The attempts must also read the same set of simulation +paths. Branches created by `get_branch` during each attempt compare by their +path of names from the same existing ancestor, so a branch recreated under the +same name on each attempt can settle too. A branch created during an earlier +attempt of this calculation also keeps its path while its original parent and +name still register that same object; this is checked both before restarting +and when comparing transitions. Other existing branches, saved snapshots and +direct clones retain separate identities even when their names match. Using +multiple distinct branches at the same path within one attempt prevents this +shortcut. Exact comparisons copy input and result array bytes and retain the +current and previous transition signatures while the calculation is running. + +Every intervening drop in a simulation whose values were read must have been +observed in the calculation's context. A read from a thread without that +context, an unexplained drop, an explicit cache reset for a policy change, an +object, string or masked stored input or result, or an input change in the simulation +doing the calculation prevents this shortcut. After a fixed point is proved, +the calculation retains its dependencies at their settled input-change counts +and keeps the result using the usual rules that preserve inputs and respect +cache restrictions. A persistent branch that deletes an input, reads another +value and sets the deleted input to the same value on each attempt can +therefore settle without exhausting the retry budget. + +An input set again with the value the branch already reads for that very +period, as an input (not through a `set_input` helper that spreads it over +other periods, and not `-0.0` for `0.0`), changes nothing, so it drops nothing +and is no input change. An input stored for the very period being calculated +after the calculation began (by its own formula, say; under the branch's name +or any it reads, such as `default` through `Holder.set_input`) is the result, +as it would be had it been set first. + +Input helpers that divide or dispatch a larger-period input replace calculated +sub-periods while retaining supplied inputs in their storage tiers. Their shared +overlap cleanup also removes calculated sums or twelfths for the variable, +including inherited values in the branch's storage, without changing the parent. +Branch dependency invalidation additionally drops later calculated values that +may depend on the replaced input. On a simulation without a parent, helper +overlap cleanup still runs while branch dependency invalidation does not. + +A custom `set_input` handler that calculates values between its own stores +calculates them from inputs it has not yet replaced, so if it calculated +anything (through `calculate`, `calculate_add` or `_calculate`), the branch +drops again, by the same rule, once the handler returns or raises (the inputs +it stored before raising stay). + +If there is no such record, the branch drops nothing. That is the case when a +formula creates the branch while it is still calculating the variable the +branch overrides, unless another branch the formula created for the same +comparison already handed back a value calculated from that variable: then the +branch drops what came back and what was calculated after it. + +A value written into a holder's storage directly, not through the storage's +`put`, has no number; an input for a variable holding such a value drops every +calculated value. + +### Why this is sound + +A formula's result is stored after everything the formula read was stored or +recorded, so a calculated value carries a larger number than each value it was +calculated from, directly or through other calculated values. Every value a +simulation holds was calculated there, inherited from its parent, handed back +by another simulation's `calculate`, or restored from a dump. In the first +three cases, for every variable it was calculated from, the simulation's +history holds a record of that variable for an overlapping period numbered no +later than the value read (for a value summed or divided from other periods, +the record may be of those periods), unless it was calculated from a +macro-cache read; restored values and values calculated from a macro-cache +read are covered by the record that values from a number on may depend on +anything. So any value that depends on the overridden variable at an +overlapping period was stored after the earliest record that applies. Uprating +and carry-over read which periods hold inputs, so a value they give can change +with an input for any period of the variable, which is why the first uprated +or carried-over value also counts. A result whose calculation was +running when an input changed is not kept unless it was calculated again from +the new inputs. An input set again with the value the branch +already reads for that period, as an input, changes no value read anywhere, so +dropping nothing then is sound; one set where the value was calculated still +drops, since it changes which periods hold inputs. + +The rule can drop more than it needs to (a value stored later that does not +depend on the input is calculated again) but not less, within these limits: + +- **Policy changes are not tracked.** A branch whose tax-benefit system or + parameters differ from its parent's still holds the parent's values + calculated under the parent's policy. Call `branch.drop_computed_arrays()` + after changing them. +- **Formulas must not write into arrays they read.** A formula that changes a + cached array in place (`array += x`) changes a value after its sequence + number was assigned, so values calculated from it may be dropped too late or + not at all. +- **Formulas should calculate, not inspect storage.** A formula that reads + `holder.get_known_periods()`, `simulation.get_array()` or another + simulation's storage directly, and acts on what it finds, depends on values + that are not recorded. +- **Unrelated simulations in other threads.** A formula that calculates in a + simulation other than its own branches, from a thread it starts without + copying its context (`contextvars.copy_context`, which `asyncio.to_thread` + does), does not take in that simulation's history, nor note what it read + there. Its own branches are covered: their history, and what was read in + them, go to every ancestor with a calculation running. +- **Concurrent use of one simulation is not supported.** A formula can wait for + a calculation in a branch from a worker thread, but simultaneous calculations + or mutations in the same simulation or branch are not covered by these + guarantees. +- **A branch a formula keeps between calls is a snapshot.** It holds what its + parent held when it was created, so inputs set on the parent afterwards do + not reach what the formula reads from it. +- **Inputs on the root simulation skip dependency invalidation.** `set_input` + on a simulation that is not a branch keeps values calculated from the replaced + input. Input helpers still remove overlapping calculated values of the + variable they write. A formula's input for its own period is still returned + as described above, and results calculated while an input is being set do not + enter the fast cache. +- **Existing child branches keep their values.** An input set on a branch does + not reach branches already created from it. +- **Values kept elsewhere stay.** A value one simulation calculated from + another's (its branch's, its parent's, or another family's) and kept stays + when that other simulation's input changes after the calculation ended: only + calculations still running then are not kept, and run again. A rerun reads + such a kept value as it is. Variables in `cache_blacklist` or + `variables_to_drop` are still kept in the simulation's fast cache, as + before, so this applies to them too. +- **Branch resets that do not settle.** A formula that reads a branch and then + sets an input the branch does not hold can trigger a rerun, whether the branch + is new or the formula deleted that input. If its transition and result do not + settle under the rules above, it uses the full retry budget and returns its + last result without keeping it. Each later `calculate` or `calculate_add` + pays those reruns again; the result is also absent from dumps and the macro + cache. + +With disk storage (`MemoryConfig`), dropping a value removes it from the +storage's index and leaves its file on disk, since a clone may still read it. +Cloned and copied storages in the same process keep the storage that owns +cleanup, and the containing temporary directory, alive. Unpreserved storage +directories are removed when their cleanup owners are collected; a simulation's +temporary directory is removed when nothing in that process keeps it alive, or +at interpreter exit, except explicitly preserved storage subdirectories and +the folders leading to them. Readers in another process do not extend that +lifetime. +Storing a dropped key again reuses a file only if that storage wrote it and has +not shared it since; cloning, copying or forking can require a new file instead +(see `OnDiskStorage._path_to_write`). + +A simulation dump (`dump_simulation`) holds the values the simulation reads +(on a branch, its own and those it inherited) and records which were +calculated (`derived_periods.txt`). `restore_simulation` restores the others +as inputs and every calculated value under one later number. The dump does not +say what each value was calculated from, nor what was read without being kept, +so the restored simulation records that values from that number on may depend +on anything: an input set for any variable on a branch of it drops all of them. +A dump written before calculated values were recorded is restored with every +value as an input, as before, so such values never drop. With disk storage, a branch whose +name contains `_` cannot be dumped yet: `OnDiskStorage.get_known_periods` +splits its keys on every `_` (as on master). + +Two related behaviours: a branch stops reading macro-cache files once an input +is set on it, even one it already reads, since they are keyed by branch name and +period but not by inputs (a file for its name may have been written by this +branch before the input, or by another simulation's branch of the same name), +and a branch that read one drops, on any input that changes what it reads, +everything calculated from the first read on; and +`requires_computation_after` is satisfied by a prerequisite successfully +requested earlier, even if its values were deleted, dropped, or not kept. + +A formula that reads a branch and then changes an input may return the value +from a rerun with the new input, as if the input had been set first. Use two +branches for comparisons that need both the earlier and the later values. + +## Dropping calculated values + +`simulation.drop_computed_arrays()` deletes every value the simulation holds +except inputs (the dataset or situation it was built from, values set with +`set_input` on it, and, for a branch, values set on the simulations it was +created from before it was created), and returns how many arrays it deleted. +Values cached with `Holder.put_in_cache`, including those without `derived=True`, +and values a custom `set_input` handler calculates are excluded from supplied +inputs. Use it on a +branch after changing its policy: + +```python +branch = simulation.get_branch("pre_reform_rules", clone_system=True) +branch.tax_benefit_system.parameters.gov.some_rate.update(period="2026", value=0.2) +branch.drop_computed_arrays() +``` diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index f1c3304f4..423d19b9a 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -1,9 +1,14 @@ -from typing import Dict, FrozenSet, Set, Union +from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Union import numpy from numpy.typing import ArrayLike from policyengine_core import periods +from policyengine_core.data_storage.store_history import ( + _period_from_storage_key, + advance_sequence_past, + next_sequence_number, +) from policyengine_core.periods import Period @@ -83,6 +88,10 @@ class InMemoryStorage: def __init__(self, is_eternal: bool): self._arrays = {} + # When each array was stored (see ``store_history``). A number + # describes the stored value, as its input mark does, so ``clone`` + # copies the numbers with the arrays. + self._sequence_numbers: Dict[str, int] = {} self.is_eternal = is_eternal def __getstate__(self) -> dict: @@ -100,7 +109,13 @@ def __setstate__(self, state: dict) -> None: inputs = set(state.get("_arrays", {})).difference(derived) if inputs: state["_inputs"] = inputs + # Pickled before stores were numbered: no value has a number. + state.setdefault("_sequence_numbers", {}) self.__dict__.update(state) + # Numbers from the process that pickled this storage must stay below + # those of stores made after unpickling it. + if self._sequence_numbers: + advance_sequence_past(max(self._sequence_numbers.values())) def clone(self, share_arrays: bool = False) -> "InMemoryStorage": """Copy this storage. @@ -152,6 +167,7 @@ def clone(self, share_arrays: bool = False) -> "InMemoryStorage": held = inputs.intersection(clone._arrays) if held: clone._inputs = frozenset(held) + clone._sequence_numbers = dict(self._sequence_numbers) return clone def get(self, period: Period, branch_name: str = "default") -> ArrayLike: @@ -260,7 +276,14 @@ def put( period: Period, branch_name: str = "default", derived: bool = False, + sequence_number: Optional[int] = None, ) -> None: + """Store ``value`` for ``period`` on ``branch_name``. + + ``derived`` marks a value the simulation calculated rather than took + as input (see ``_inputs``), and ``sequence_number`` records when the + value was stored (a new number by default); see :meth:`drop_computed`. + """ if self.is_eternal: period = periods.period(periods.ETERNITY) period = periods.period(period) @@ -287,6 +310,58 @@ def put( self._unmark_input(key) else: self._mark_input(key) + self._sequence_numbers[key] = ( + next_sequence_number() if sequence_number is None else sequence_number + ) + + def drop_computed(self, *, since: Optional[int] = None) -> int: + """Delete stored values that are not inputs, and return how many. + + With ``since``, only values stored with that sequence number or a + later one are deleted (a value with no recorded number counts as + later). Inputs, which ``put`` received without ``derived``, are kept + whatever their number. + """ + inputs = self._inputs + numbers = self._sequence_numbers + dropped = [ + key + for key in self._arrays + if key not in inputs and (since is None or numbers.get(key, since) >= since) + ] + for key in dropped: + del self._arrays[key] + numbers.pop(key, None) + self._stop_sharing(key) + return len(dropped) + + def inputs_since(self, since: Optional[int] = None) -> List[Tuple[Period, int]]: + """The period and number of each input stored at ``since`` or later (or ever).""" + numbers = self._sequence_numbers + return [ + (_period_from_storage_key(key.split(":", 1)[1]), numbers[key]) + for key in self._inputs + if key in self._arrays + and key in numbers + and (since is None or numbers[key] >= since) + ] + + def mark_derived_except(self, keys: Set[str]) -> None: + """Mark every stored value derived except those stored for ``keys``.""" + for key in [key for key in self._inputs if key not in keys]: + self._unmark_input(key) + + def has_unnumbered_values(self) -> bool: + """Whether a value was stored without ``put`` (so without a number).""" + return any(key not in self._sequence_numbers for key in self._arrays) + + def _forget_dropped_numbers(self) -> None: + """Forget the numbers of keys ``_arrays`` no longer has.""" + numbers = self._sequence_numbers + if not numbers.keys() <= self._arrays.keys(): + self._sequence_numbers = { + key: number for key, number in numbers.items() if key in self._arrays + } def delete(self, period: Period = None, branch_name: str = "default") -> None: if period is None: @@ -300,6 +375,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: } self._stop_sharing_dropped_keys() self._unmark_dropped_keys() + self._forget_dropped_numbers() return if self.is_eternal: @@ -319,6 +395,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: } self._stop_sharing_dropped_keys() self._unmark_dropped_keys() + self._forget_dropped_numbers() 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 1c7f8bf63..9cac6e2c7 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -2,11 +2,17 @@ import shutil import tempfile import weakref +from typing import Dict, List, Optional, Set, Tuple import numpy from numpy.typing import ArrayLike from policyengine_core import periods +from policyengine_core.data_storage.store_history import ( + _period_from_storage_key, + advance_sequence_past, + next_sequence_number, +) from policyengine_core.data_storage.storage_directory import ( directory_containing, path_key, @@ -91,6 +97,8 @@ def _start(self, storage_dir: str, is_eternal: bool, parent_directory) -> None: # File keys stored with ``put(..., derived=True)``; see # ``InMemoryStorage``. self._derived = set() + # When each file key was stored; see ``InMemoryStorage``. + self._sequence_numbers: Dict[str, int] = {} # For each key, the file this storage last wrote for it, while it has # not shared that file since: the only files of its family ``put`` # writes over (see ``_path_to_write``). Kept when the key is deleted, @@ -204,7 +212,15 @@ def __setstate__(self, state: dict) -> None: state["_storage_dir_owner"] = owner if state.get("_parent_directory") is None: state["_parent_directory"] = directory_containing(state["storage_dir"]) + # Nor numbers: pickled before stores were numbered. Copies have their + # own numbers, like their file index: replacing or dropping a key + # through one must not change when the other's value was stored. + state["_sequence_numbers"] = dict(state.get("_sequence_numbers", {})) self.__dict__.update(state) + # Numbers from the process that pickled this storage must stay below + # those of stores made after unpickling it. + if self._sequence_numbers: + advance_sequence_past(max(self._sequence_numbers.values())) def clone(self) -> "OnDiskStorage": """Create a private metadata view over this storage directory. @@ -225,6 +241,7 @@ def clone(self) -> "OnDiskStorage": clone._files = self._files.copy() clone._enums = self._enums.copy() clone._derived = set(self._derived) + clone._sequence_numbers = dict(self._sequence_numbers) # One set of read-back paths for the family, so a clone made before # ``restore`` still writes over none of them after a fork. clone._restored_paths = self._restored_paths @@ -276,7 +293,13 @@ def put( period: Period, branch_name: str = "default", derived: bool = False, + sequence_number: Optional[int] = None, ) -> None: + """Store ``value`` for ``period`` on ``branch_name``. + + ``derived`` and ``sequence_number`` are as in + :meth:`InMemoryStorage.put`. + """ if self.is_eternal: period = periods.period(periods.ETERNITY) period = periods.period(period) @@ -296,6 +319,59 @@ def put( self._derived.add(filename) else: self._derived.discard(filename) + self._sequence_numbers[filename] = ( + next_sequence_number() if sequence_number is None else sequence_number + ) + + def drop_computed(self, *, since: Optional[int] = None) -> int: + """Forget stored values that are not inputs, and return how many. + + As :meth:`InMemoryStorage.drop_computed`. The files stay on disk + (clones may still read them) until the storage directory is removed; + writing a dropped key again reuses this storage's own file for it, + as after ``delete``. + """ + numbers = self._sequence_numbers + dropped = [ + key + for key in self._files + if key in self._derived + and (since is None or numbers.get(key, since) >= since) + ] + for key in dropped: + del self._files[key] + numbers.pop(key, None) + self._derived.difference_update(dropped) + return len(dropped) + + def inputs_since(self, since: Optional[int] = None) -> List[Tuple[Period, int]]: + """As :meth:`InMemoryStorage.inputs_since`.""" + numbers = self._sequence_numbers + # Period strings contain no "_"; branch names may. + return [ + (_period_from_storage_key(key.rsplit("_", 1)[1]), numbers[key]) + for key in self._files + if key not in self._derived + and key in numbers + and (since is None or numbers[key] >= since) + ] + + def mark_derived_except(self, keys: Set[str]) -> None: + """As :meth:`InMemoryStorage.mark_derived_except`.""" + self._derived.update(key for key in self._files if key not in keys) + + def has_unnumbered_values(self) -> bool: + """As :meth:`InMemoryStorage.has_unnumbered_values`.""" + return any(key not in self._sequence_numbers for key in self._files) + + def _forget_dropped_keys(self) -> None: + """Forget the derived marks and numbers of keys ``_files`` no longer has.""" + self._derived.intersection_update(self._files) + numbers = self._sequence_numbers + if not numbers.keys() <= self._files.keys(): + self._sequence_numbers = { + key: number for key, number in numbers.items() if key in self._files + } def _path_to_write(self, filename: str) -> str: """The path ``put`` writes the value for key ``filename`` to. @@ -348,7 +424,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) + self._forget_dropped_keys() return if self.is_eternal: @@ -361,7 +437,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) + self._forget_dropped_keys() def get_known_periods(self) -> list: return list([periods.period(x.split("_")[1]) for x in self._files.keys()]) @@ -393,8 +469,9 @@ def restore(self) -> None: are unchanged. """ self._files = files = {} - # Files read back from a directory carry no derived marks. + # Files read back from a directory carry no derived marks or numbers. self._derived = set() + self._sequence_numbers = {} # 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/data_storage/store_history.py b/policyengine_core/data_storage/store_history.py new file mode 100644 index 000000000..2c30c84c8 --- /dev/null +++ b/policyengine_core/data_storage/store_history.py @@ -0,0 +1,247 @@ +"""The order in which a simulation's values were stored. + +Every array a holder stores gets a sequence number from one process-wide +counter, so a larger number means a later store. A value calculated by a +formula is stored after everything the formula read, so it carries a larger +number than each stored value it was calculated from, directly or through +other calculated values. + +``Simulation.set_input`` on a branch uses this to drop only the values that +can depend on the input it replaces: see +:meth:`policyengine_core.simulations.Simulation.set_input`. +""" + +import itertools +import weakref +from typing import Dict, List, Optional, Tuple + +from policyengine_core import periods +from policyengine_core.periods import Period + +_sequence = itertools.count(1) + + +def next_sequence_number() -> int: + """Return a sequence number larger than every one returned before.""" + return next(_sequence) + + +def advance_sequence_past(number: int) -> None: + """Make every later sequence number larger than ``number``. + + Numbers come from a counter in each process, so values unpickled from + another process may carry larger numbers than this process has handed + out; stores made after them must still be numbered later. + """ + global _sequence + _sequence = itertools.count(max(next(_sequence), number + 1)) + + +def _period_from_storage_key(period_string: str) -> Period: + """Parse a stored period whose early year may have lost leading zeros.""" + components = period_string.split(":") + date_index = 1 if len(components) > 1 else 0 + date = components[date_index].split("-") + date[0] = date[0].zfill(4) + components[date_index] = "-".join(date) + return periods.period(":".join(components)) + + +def periods_overlap(first: Period, second: Period) -> bool: + """Whether two periods share at least one day.""" + if periods.ETERNITY in (first.unit, second.unit): + return True + return first.start <= second.stop and second.start <= first.stop + + +class StoreHistory: + """What the values a simulation holds may have been calculated from. + + Each simulation has its own history. For each variable it records: + + - for each period, the sequence number of the first value stored or + calculated for that period that a value the simulation holds may have + been calculated from, including values a holder calculates but does + not keep (``variables_to_drop``, the cache blacklist), whose readers + carry later numbers all the same; + - the sequence number of the first time such a value of the variable was + derived from its values for other periods: uprated or carried over (or + given the default because no period it could be uprated or carried over + from holds an input). Such a value depends on which other periods hold + inputs, so an input set for any period can change it. + + It also records the first sequence number from which held values may + have been calculated from anything (:meth:`record_unknown_sources`): + values restored from a dump, which does not say what each was + calculated from, and values calculated from a macro-cache read, whose + sources were never calculated here. An input for any variable can + change those. + + The simulation's own stores are recorded as they happen. A branch starts + with a copy of its parent's history, as it starts with a copy of its + parent's values. When a formula running in one simulation calculates a + value in another (a branch it created, for instance), the caller merges + the other's history into its own, since what the caller stores next may + be calculated from what it got back. A drop removes the records at or + after its sequence number (see :meth:`prune`). + """ + + def __init__(self): + self._first_stored: Dict[str, Dict[Period, int]] = {} + self._first_derived: Dict[str, int] = {} + # Values numbered from here on may have been calculated from anything. + self._unknown_sources_since: Optional[int] = None + # Every change to the records since this history was created or last + # pruned, in order, so a merge takes in only what changed since it + # last read this history (``period`` is None for a derived record, + # and ``variable_name`` too for an unknown-sources record). + self._journal: List[Tuple[Optional[str], Optional[Period], int]] = [] + self._generation = 0 # changes when a prune empties the journal + # For each history merged into this one, the generation and journal + # length read last. Keyed weakly: an entry goes with the history it + # describes (a formula's temporary branch, say). + self._merged: "weakref.WeakKeyDictionary[StoreHistory, Tuple[int, int]]" = ( + weakref.WeakKeyDictionary() + ) + + def __getstate__(self) -> dict: + state = dict(self.__dict__) + # Weak references do not pickle; merging again only repeats work. + state["_merged"] = {} + return state + + def __setstate__(self, state: dict) -> None: + state.setdefault("_journal", []) + state.setdefault("_generation", 0) + state.setdefault("_unknown_sources_since", None) + self.__dict__.update(state) + self._merged = weakref.WeakKeyDictionary() + numbers = [ + number + for stored in self._first_stored.values() + for number in stored.values() + ] + list(self._first_derived.values()) + if self._unknown_sources_since is not None: + numbers.append(self._unknown_sources_since) + if numbers: + advance_sequence_past(max(numbers)) + + def copy(self) -> "StoreHistory": + new = StoreHistory() + new._first_stored = { + variable_name: dict(stored) + for variable_name, stored in self._first_stored.items() + } + new._first_derived = dict(self._first_derived) + new._unknown_sources_since = self._unknown_sources_since + new._merged = weakref.WeakKeyDictionary(self._merged) + return new + + def record_store(self, variable_name: str, period: Period, sequence_number: int): + stored = self._first_stored.setdefault(variable_name, {}) + recorded = stored.get(period) + if recorded is None or sequence_number < recorded: + stored[period] = sequence_number + self._journal.append((variable_name, period, sequence_number)) + + def record_derived(self, variable_name: str, sequence_number: int): + recorded = self._first_derived.get(variable_name) + if recorded is None or sequence_number < recorded: + self._first_derived[variable_name] = sequence_number + self._journal.append((variable_name, None, sequence_number)) + + def record_unknown_sources(self, sequence_number: int): + """Values numbered ``sequence_number`` or later may depend on any input.""" + recorded = self._unknown_sources_since + if recorded is None or sequence_number < recorded: + self._unknown_sources_since = sequence_number + self._journal.append((None, None, sequence_number)) + + def merge(self, other: "StoreHistory") -> None: + """Take in ``other``'s records, keeping the earlier number of each.""" + if other is self: + return + # Note how far ``other`` had got before reading it: anything it + # records meanwhile (another thread) is read next time. + generation, end = other._generation, len(other._journal) + read = self._merged.get(other) + if read is not None and read[0] == generation: + changes = other._journal[read[1] : end] + else: + changes = [ + (variable_name, period, sequence_number) + for variable_name, stored in list(other._first_stored.items()) + for period, sequence_number in list(stored.items()) + ] + [ + (variable_name, None, sequence_number) + for variable_name, sequence_number in list(other._first_derived.items()) + ] + unknown_sources_since = other._unknown_sources_since + if unknown_sources_since is not None: + changes.append((None, None, unknown_sources_since)) + for variable_name, period, sequence_number in changes: + if variable_name is None: + self.record_unknown_sources(sequence_number) + elif period is None: + self.record_derived(variable_name, sequence_number) + else: + self.record_store(variable_name, period, sequence_number) + self._merged[other] = (generation, end) + + def prune(self, since: Optional[int] = None) -> None: + """Forget the records numbered ``since`` or later (all, without ``since``). + + Called after the simulation drops every value it holds, other than + inputs, numbered ``since`` or later: nothing it still holds was + calculated from what those records describe. The caller records again + the inputs it keeps. + """ + self._first_stored = { + variable_name: kept + for variable_name, stored in self._first_stored.items() + if ( + kept := { + period: sequence_number + for period, sequence_number in stored.items() + if since is not None and sequence_number < since + } + ) + } + self._first_derived = { + variable_name: sequence_number + for variable_name, sequence_number in self._first_derived.items() + if since is not None and sequence_number < since + } + if since is None or ( + self._unknown_sources_since is not None + and self._unknown_sources_since >= since + ): + self._unknown_sources_since = None + self._journal = [] + self._generation += 1 + # Merging a history again must restore what was forgotten here. + self._merged = weakref.WeakKeyDictionary() + + def earliest_dependency(self, variable_name: str, period: Period) -> Optional[int]: + """The first sequence number from which a value may depend on ``variable_name`` at ``period``. + + This is the earliest recorded store of the variable for any period + that shares a day with ``period``, the earliest recorded derivation + of one of its values from other periods, or the first number from + which values may depend on anything, whichever came first. + ``None`` means no value the simulation holds can depend on the + variable's value at ``period``. + """ + candidates = [ + sequence_number + for stored_period, sequence_number in self._first_stored.get( + variable_name, {} + ).items() + if periods_overlap(stored_period, period) + ] + derived = self._first_derived.get(variable_name) + if derived is not None: + candidates.append(derived) + if self._unknown_sources_since is not None: + candidates.append(self._unknown_sources_since) + return min(candidates) if candidates else None diff --git a/policyengine_core/holders/helpers.py b/policyengine_core/holders/helpers.py index 33414a0c7..b6e4b9828 100644 --- a/policyengine_core/holders/helpers.py +++ b/policyengine_core/holders/helpers.py @@ -5,6 +5,7 @@ from numpy.typing import ArrayLike from policyengine_core import periods +from policyengine_core.data_storage.store_history import _period_from_storage_key from policyengine_core.holders.holder import Holder from policyengine_core.periods import Period @@ -152,6 +153,7 @@ def _drop_calculated_overlapping( dropped_periods.add(_period_from_storage_key(period_string)) memory._stop_sharing_dropped_keys() memory._unmark_dropped_keys() + memory._forget_dropped_numbers() disk = holder._disk_storage if disk is not None: for key in list(disk._files): @@ -161,20 +163,10 @@ def _drop_calculated_overlapping( ): del disk._files[key] dropped_periods.add(_period_from_storage_key(period_string)) - disk._derived.intersection_update(disk._files) + disk._forget_dropped_keys() _evict_fast_cache(holder, dropped_periods, branch_name) -def _period_from_storage_key(period_string: str) -> Period: - """Parse a stored period whose early year may have lost leading zeros.""" - components = period_string.split(":") - date_index = 1 if len(components) > 1 else 0 - date = components[date_index].split("-") - date[0] = date[0].zfill(4) - components[date_index] = "-".join(date) - return periods.period(":".join(components)) - - def _is_calculated_overlapping( holder: Holder, stored_branch_name: str, diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 1f393fba5..91c52b8a8 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -9,6 +9,7 @@ from policyengine_core import commons, periods, tools from policyengine_core.data_storage import InMemoryStorage, OnDiskStorage +from policyengine_core.data_storage.store_history import next_sequence_number from policyengine_core.enums import Enum from policyengine_core.errors import PeriodMismatchError from policyengine_core.periods import Period @@ -130,6 +131,12 @@ def delete_arrays( user_input_keys = getattr(simulation, "_user_input_keys", None) stored_before = self._stored_keys() if user_input_keys else None + if self.simulation is not None: + from policyengine_core.simulations.simulation import _note_input_change + + _note_input_change( + self.simulation, "delete", self.variable.name, period, branch_name + ) self._memory_storage.delete(period, branch_name) if self._disk_storage: self._disk_storage.delete(period, branch_name) @@ -308,6 +315,134 @@ def _storage_period(self, period: Period) -> Period: return Period((unit, start, size)) return period + def _drop_computed(self, since: Optional[int] = None) -> int: + """Delete every stored value that is not an input, and return how many. + + With ``since``, only values stored with that sequence number or a + later one are deleted. Inputs are values stored through + :meth:`set_input`, on any branch. + """ + if (self._memory_storage._inputs or self._disk_storage is not None) and getattr( + self.simulation, "_user_input_keys", None + ) is not None: + from .helpers import _is_input, _period_from_storage_key + + # Carry-over marks also include low-level cache writes. Only an + # actual supplied input in this tier survives a branch's drop. + for storage_name, storage in ( + ("memory", self._memory_storage), + ("disk", self._disk_storage), + ): + if storage is None: + continue + candidate_keys = ( + storage._inputs + if storage_name == "memory" + else storage._files.keys() - storage._derived + ) + kept = set() + recorded = getattr(self, "_user_input_storage", {}) + for key in candidate_keys: + branch_name, period_string = ( + key.split(":", 1) + if storage_name == "memory" + else key.rsplit("_", 1) + ) + input_storage = recorded.get((branch_name, period_string)) + if input_storage is not None: + stored_period = input_storage[0] + else: + try: + stored_period = _period_from_storage_key(period_string) + except ValueError: + continue + if _is_input(self, stored_period, branch_name, storage_name): + kept.add(key) + storage.mark_derived_except(kept) + dropped = self._memory_storage.drop_computed(since=since) + if self._disk_storage is not None: + dropped += self._disk_storage.drop_computed(since=since) + return dropped + + def _record_inputs(self, since: Optional[int] = None) -> None: + """Record in the simulation's history each input stored at ``since`` or later.""" + for storage in (self._memory_storage, self._disk_storage): + if storage is not None: + for period, sequence_number in storage.inputs_since(since): + self._record_store(period, sequence_number) + + def _is_input(self, period: Period, branch_name: str = "default") -> bool: + """Whether a value is stored for ``period`` under ``branch_name`` + itself, as the supplied input that ``get_array`` reads.""" + from .helpers import _is_input + + for storage_name, storage in ( + ("memory", self._memory_storage), + ("disk", self._disk_storage), + ): + if storage is not None and storage.has(period, branch_name): + return _is_input(self, period, branch_name, storage_name) + return False + + def _input_unchanged( + self, period: Period, array: ArrayLike, branch_name: str = "default" + ) -> bool: + """Whether :meth:`set_input` would store, for ``period`` itself, the + input :meth:`get_array` already reads there for ``branch_name``. + + Then nothing calculated from that input changes. An input a + ``set_input`` helper spreads over other periods never counts, nor + does a value equal only after conversion to the variable's type (a + float rounded to float32, say), nor ``-0.0`` for ``0.0`` (``1 / x`` + tells them apart). + """ + if self.variable.set_input and period.unit != self.variable.definition_period: + return False + stored_on = self._branch_storing(period, branch_name) + if stored_on is None or not self._is_input(period, stored_on): + return False + try: + stored = self._get_array_from_storage(period, stored_on) + given = numpy.asarray(array) + return bool( + numpy.array_equal(stored, given) + and numpy.array_equal(numpy.signbit(stored), numpy.signbit(given)) + ) + except Exception: # values that cannot be compared are not equal + return False + + def _stored_sequence_number( + self, period: Period, branch_name: str = "default" + ) -> Optional[int]: + """The sequence number of the value stored for ``period`` on ``branch_name``.""" + if self.variable.definition_period == periods.ETERNITY: + period = periods.period(periods.ETERNITY) + period = periods.period(period) + number = self._memory_storage._sequence_numbers.get(f"{branch_name}:{period}") + if number is None and self._disk_storage is not None: + number = self._disk_storage._sequence_numbers.get(f"{branch_name}_{period}") + return number + + def _mark_derived_except(self, kept) -> None: + """Mark every stored value derived except those stored for the + ``(branch_name, period)`` pairs in ``kept``.""" + memory_keys, disk_keys = set(), set() + for branch_name, period in kept: + if self.variable.definition_period == periods.ETERNITY: + period = periods.ETERNITY + period = periods.period(period) + memory_keys.add(f"{branch_name}:{period}") + disk_keys.add(f"{branch_name}_{period}") + self._memory_storage.mark_derived_except(memory_keys) + if self._disk_storage is not None: + self._disk_storage.mark_derived_except(disk_keys) + + def _has_unnumbered_values(self) -> bool: + return self._memory_storage.has_unnumbered_values() or ( + self._disk_storage is not None + and self._disk_storage.has_unnumbered_values() + ) + def _get_array_from_storage( self, period: Period, branch_name: str = "default" ) -> ArrayLike: @@ -475,6 +610,24 @@ def set_input( self._raise_if_input_contains_nan(numpy.asarray(array)) simulation = getattr(self, "simulation", None) if simulation is not None: + # On a branch, drop what may have been calculated from the value + # this input replaces (see ``Simulation.set_input``), unless it + # replaces the same input. Either way the branch stops reading + # macro-cache files, which are keyed by branch name and period, + # not by inputs: a file for its name (written by this branch, or + # by another simulation's branch of the same name) may hold values + # calculated without this input, even before the branch read any. + if getattr(simulation, "parent_branch", None) is None: + simulation._drop_values_that_may_depend_on(self.variable.name, period) + else: + simulation.macro_cache_read = False + if not self._input_unchanged(period, array, branch_name): + simulation._drop_values_that_may_depend_on( + self.variable.name, period + ) + # A calculation running meanwhile looks for inputs set for its own + # period (``Simulation._cache_result``). + simulation._inputs_set += 1 if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() if not hasattr(simulation, "_user_input_contexts"): @@ -485,7 +638,20 @@ def set_input( self.variable.set_input and period.unit != self.variable.definition_period ): - return self.variable.set_input(self, period, array) + started = getattr(simulation, "_calculations_started", 0) + try: + return self.variable.set_input(self, period, array) + finally: + if ( + simulation is not None + and simulation._calculations_started != started + ): + # The handler calculated, maybe between its stores, + # from inputs it had not yet replaced (also if it + # then failed: the inputs it stored stay). + simulation._drop_values_that_may_depend_on( + self.variable.name, period + ) return self._set(period, array, branch_name, validate_nan=True) finally: if simulation is not None: @@ -555,6 +721,7 @@ def _set( validate_nan: bool = False, derived: bool = False, is_input: Optional[bool] = None, + sequence_number: Optional[int] = None, ) -> None: simulation = getattr(self, "simulation", None) user_input_contexts = ( @@ -591,10 +758,19 @@ def _set( if stored_period is not None: self._seed_input_storage(stored_period, branch_name) - if should_store_on_disk: - self._disk_storage.put(value, period, branch_name, derived=derived) - else: - self._memory_storage.put(value, period, branch_name, derived=derived) + if sequence_number is None or not derived: + # Inputs are always numbered when stored: a calculation running + # meanwhile looks for inputs stored after it began. + sequence_number = next_sequence_number() + storage = self._disk_storage if should_store_on_disk else self._memory_storage + storage.put( + value, + period, + branch_name, + derived=derived, + sequence_number=sequence_number, + ) + self._record_store(period, sequence_number) self._evict_fast_cache(period, branch_name) if stored_period is not None: self._record_input_storage( @@ -603,6 +779,19 @@ def _set( "disk" if should_store_on_disk else "memory", is_input=is_input, ) + if not derived and simulation is not None: + from policyengine_core.simulations.simulation import _note_input_change + + # Raw cache writes are effects a stale caller must repeat, even + # though they are not supplied inputs for replay or export. + _note_input_change( + simulation, + "set", + self.variable.name, + stored_period or self._storage_period(period), + branch_name, + value, + ) if is_input and simulation is not None: if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() @@ -634,14 +823,13 @@ def put_in_cache( 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 - - if ( + if self._do_not_store or ( self.simulation.opt_out_cache and self.simulation.tax_benefit_system.cache_blacklist and self.variable.name in self.simulation.tax_benefit_system.cache_blacklist ): + # Not kept, but whatever reads it is stored after it all the same. + self._record_store(period, next_sequence_number()) return if ( @@ -653,6 +841,19 @@ def put_in_cache( self._set(period, value, branch_name, is_input=False, derived=derived) + def _record_store(self, period: Period, sequence_number: int) -> None: + simulation = getattr(self, "simulation", None) + if simulation is None: + return + variable = self.variable + if variable.definition_period == periods.ETERNITY: + period = periods.period(periods.ETERNITY) + elif not isinstance(period, Period): + period = periods.period(period) + simulation._get_store_history().record_store( + variable.name, period, sequence_number + ) + def default_array(self) -> ArrayLike: """ Return a new array of the appropriate length for the entity, filled with the variable default values. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 96d05624e..e4628e244 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1,9 +1,11 @@ import hashlib import os +from contextlib import nullcontext import types from contextvars import ContextVar from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union +from threading import Lock +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union import numpy as np import pandas as pd @@ -13,6 +15,10 @@ from policyengine_core import commons, periods from policyengine_core.data.dataset import Dataset +from policyengine_core.data_storage.store_history import ( + StoreHistory, + next_sequence_number, +) from policyengine_core.data_storage.storage_directory import ( TemporaryStorageDirectory, directory_containing, @@ -176,6 +182,494 @@ def _latest_input_key(period: Period, definition_period: str) -> tuple: ) +# The simulation whose formula is running. When that formula calculates a +# value in another simulation (a branch it created, say), what the formula +# returns, and so what its simulation stores, may be calculated from the other +# simulation's values: ``calculate`` then merges the other's store history into +# this one's (see ``StoreHistory``). +_formula_simulation: ContextVar[Optional["Simulation"]] = ContextVar( + "_formula_simulation", default=None +) + +# How many times ``calculate`` runs a calculation again when an input set on +# the simulation while it ran (by a formula, say) dropped values. A formula +# that changes an input on every run is kept from looping; its last result is +# returned but not kept. +_RERUNS_AFTER_INPUT_CHANGE = 10 + +# The calculations running in this context, innermost last (``_Frame``). +_calculation_frames: ContextVar[tuple] = ContextVar("_calculation_frames", default=()) + +# Only these frames try to keep work that the stale shared frame refused. +# Context-free activity anywhere prevents that optional cache write, without +# changing ordinary frames' scoped read tracking or convergence decisions. +_restricted_frames = set() +_restricted_frames_lock = Lock() +# Activity between restricted calculations must also invalidate their shared +# attempt cache, even when no restricted child is currently running. +_unobserved_activity_generation = 0 + + +def _note_unobserved_activity() -> None: + global _unobserved_activity_generation + with _restricted_frames_lock: + _unobserved_activity_generation += 1 + for frame in _restricted_frames: + frame.unobserved_activity = True + + +@dataclass +class _AttemptResult: + """Read-only work reusable within one stale attempt, with its dependencies.""" + + array: ArrayLike + epoch: int + inputs_set: int + reads: tuple + registrations: tuple + activity_generation: int + + def is_current(self, simulation: "Simulation") -> bool: + return ( + self.activity_generation == _unobserved_activity_generation + and self.epoch == simulation._input_epoch + and self.inputs_set == simulation._inputs_set + and all( + other._input_epoch == epoch and other._inputs_set == inputs_set + for other, epoch, inputs_set in self.reads + ) + and all( + parent.branches.get(name) is branch + for parent, name, branch in self.registrations + ) + ) + + +class _Frame: + """A calculation running in ``simulation``. + + ``start`` is the simulation's ``_input_epoch`` when the calculation (last) + began, and ``reads`` holds, for each other simulation it got a value + from, directly or through the calculations it called, that simulation's + ``_input_epoch`` when the calculation that produced the value began. If + any of those has changed since (a drop there), the value may come from a + replaced one, and so may what the frame calculates from it, which is then + not kept (:meth:`is_stale`). + + Used as a context manager: entering pushes the frame on + ``_calculation_frames`` and lists it in the simulation's + ``_open_frames``, so that a value read in a thread with no frame reaches + the frames waiting for it (``_hand_to_waiting_ancestors``). Leaving, + whatever way the calculation ends (a value, an early default, an error), + tells the caller what it read (``_hand_to_caller``). A class with slots, + not a generator or a dict: one runs for every calculation. + """ + + __slots__ = ( + "simulation", + "start", + "reads", + "registrations", + "outer", + "token", + "registry", + "changes", + "drop_epochs", + "untracked_changes", + "inputs_set", + "branch_identities", + "created_branches", + "carried_branches", + "stale_parent", + "unobserved_activity", + "attempt_cache", + ) + + def __init__(self, simulation: "Simulation", stale_parent: bool = False): + self.simulation = simulation + self.created_branches = {} + self.carried_branches = {} + self.stale_parent = stale_parent + + def restart(self) -> None: + """Begin again, as when the calculation runs again.""" + self.start = self.simulation._input_epoch + self.inputs_set = self.simulation._inputs_set + self.reads = {} + # Capture this calculation's named path before its formula can + # replace a registration. Foreign callers inherit these exact links. + self.registrations = _branch_registrations(self.simulation) + self.changes = [] + self.drop_epochs = {} + self.untracked_changes = False + self.unobserved_activity = False + self.branch_identities = {} + # Restricted children share this dictionary only until this attempt + # restarts. Foreign results never enter persistent caches this way. + self.attempt_cache = {} + # Retain only creations whose original registration still names + # the same object. Saved or replaced snapshots keep their identities. + self.carried_branches = _registered_creations( + self.carried_branches | self.created_branches + ) + self.created_branches = {} + + def note_branch(self, simulation: "Simulation") -> None: + """A name path must refer to just one simulation within an attempt. + + A direct clone retains the branch's name and ancestors. Distinct + clones at that path can take turns making identical stores without + having reached the same state. Recreating one branch between attempts + is still comparable, since this map resets on each restart. + """ + path = _branch_path(simulation) + previous = self.branch_identities.setdefault(path, simulation) + if previous is not simulation: + self.untracked_changes = True + + def transition_signature(self, result: ArrayLike) -> Optional[tuple]: + """A complete, comparable branch transition, or no convergence proof. + + Two identical transitions returning identical values reach a fixed + point for deterministic formulas. Only foreign branches qualify: + own-simulation changes still require the existing input-first rerun. + Every stale read's intervening drops must have been observed in this + context; a thread's unobserved writes cannot prove convergence. + """ + # Include read-only clones too: the same logged mutation can target + # a different clone on each attempt after inspecting several clones. + for other in tuple(self.reads): + self.note_branch(other) + if ( + self.untracked_changes + or not self.changes + or self.start != self.simulation._input_epoch + or self.inputs_set != self.simulation._inputs_set + ): + return None + for other, epoch in tuple(self.reads.items()): + changed = other._input_epoch - epoch + observed = self.drop_epochs.get(other, ()) + if changed < 0 or sum(number > epoch for number in observed) != changed: + return None + value = _value_signature(result) + if value is None: + return None + # A saved branch snapshot can share a name with a later recreation. + # Creations observed in this frame may keep their name path across + # retries while still registered. Recheck here: a formula can replace + # a registration after reading it, before the attempt ends. + comparable_branches = ( + self.created_branches.keys() + | _registered_creations(self.carried_branches).keys() + ) + changes = tuple( + (operation, _branch_path(other, comparable_branches), *details) + for operation, other, *details in self.changes + ) + reads = frozenset( + _branch_path(other, comparable_branches) for other in self.reads + ) + return changes, reads, value + + def settle_reads(self) -> None: + """Keep dependencies, at their settled epochs, after a proven fixed point.""" + self.reads = {other: other._input_epoch for other in self.reads} + + def is_stale(self) -> bool: + """Whether what the frame calculates may come from a value since replaced.""" + if self.simulation._input_epoch != self.start: + return True + reads = self.reads + # A snapshot: a thread handing a read to this frame may add one meanwhile. + return bool(reads) and any( + other._input_epoch != epoch for other, epoch in tuple(reads.items()) + ) + + def may_keep(self) -> bool: + """Cache independent work after a stale caller without hiding foreign reads. + + A fresh nested frame can keep work using only its own simulation. + Reads, mutations and branch creations must still reach the stale + caller on its next attempt: caching them here could hide a later + input change or skip a mutation when that caller retries. + """ + return ( + not ( + self.stale_parent + and ( + self.reads + or self.changes + or self.untracked_changes + or self.created_branches + or self.unobserved_activity + ) + ) + and not self.is_stale() + ) + + def remember(self, variable: str, period: Period, array: ArrayLike) -> None: + """Reuse stable foreign reads without hiding reads or effects on retry.""" + if ( + not self.stale_parent + or not self.reads + or self.changes + or self.untracked_changes + or self.created_branches + or self.unobserved_activity + or self.is_stale() + or any( + parent.branches.get(name) is not branch + for parent, name, branch in self.registrations + ) + ): + return + self.attempt_cache[(self.simulation, variable, period)] = _AttemptResult( + array, + self.start, + self.inputs_set, + tuple( + (other, epoch, other._inputs_set) for other, epoch in self.reads.items() + ), + tuple(self.registrations), + _unobserved_activity_generation, + ) + + def recall(self, variable: str, period: Period) -> Optional[ArrayLike]: + """Validate a memoized result and pass its foreign reads to this frame.""" + key = (self.simulation, variable, period) + cached = self.attempt_cache.get(key) + if cached is None: + return None + if not cached.is_current(self.simulation): + del self.attempt_cache[key] + return None + for other, epoch, _ in cached.reads: + self.reads.setdefault(other, epoch) + self.registrations.update(cached.registrations) + # This simulation's history cannot have been pruned without changing + # its epoch, so it still records every source of the reused value. + return cached.array + + def __enter__(self) -> "_Frame": + simulation = self.simulation + self.restart() + self.outer = outer = _calculation_frames.get() + if self.stale_parent: + self.attempt_cache = outer[-1].attempt_cache + self.token = _calculation_frames.set(outer + (self,)) + registry = simulation._open_frames + if registry is None: + # ``setdefault``: two threads starting here at once share one set. + registry = simulation.__dict__.setdefault("_open_frames", set()) + registry.add(self) + self.registry = registry + if self.stale_parent: + with _restricted_frames_lock: + _restricted_frames.add(self) + return self + + def __exit__(self, *exc_info) -> bool: + if self.stale_parent: + with _restricted_frames_lock: + _restricted_frames.discard(self) + self.registry.discard(self) + _calculation_frames.reset(self.token) + _hand_to_caller(self.outer, self) + return False + + +def _frame_for(simulation: "Simulation", outermost: bool) -> Optional[_Frame]: + """A new frame for a calculation in ``simulation``, or ``None`` to share + the frame it runs in. + + A calculation nested in another in the same simulation shares that one's + frame while it can cache: the shared frame began no later and notes at + least the same reads, each at the first epoch noted, so it is stale + whenever the nested calculation's own would be, and the nested one keeps + a result only when it would have anyway. The outermost calculation has its + own, to run again from (``Simulation.calculate``). Once a shared frame + cannot cache, independent nested work gets a fresh frame so it can still + cache. Reads and mutations remain guarded (``_Frame.may_keep``). + """ + if outermost: + return _Frame(simulation) + frames = _calculation_frames.get() + if not frames or frames[-1].simulation is not simulation: + return _Frame(simulation) + if not frames[-1].may_keep(): + return _Frame(simulation, stale_parent=True) + return None + + +def _hand_to_caller(outer: tuple, frame: _Frame) -> None: + """Pass what a calculation read on to the frame that called it. + + The calculation's own simulation counts at its ``start``: if an input + there changed while it ran, its caller's result is stale too. With no + calling frame in this context (``outer`` empty), the reads go to the + frames open in the simulation's ancestors (``_hand_to_waiting_ancestors``). + """ + reads = frame.reads + reads = tuple(reads.items()) if reads else () + simulation = frame.simulation + if not outer: + _hand_to_waiting_ancestors(simulation, frame.start, reads) + return + caller = outer[-1] + caller.untracked_changes |= frame.untracked_changes + caller.registrations.update(frame.registrations) + caller_simulation = caller.simulation + caller_reads = caller.reads + for other, epoch in reads: + if other is not caller_simulation and other not in caller_reads: + caller_reads[other] = epoch + if simulation is not caller_simulation and simulation not in caller_reads: + caller_reads[simulation] = frame.start + + +def _hand_to_waiting_ancestors( + simulation: "Simulation", start: int, reads: tuple +) -> None: + """Note a value read in ``simulation`` with no calling frame in this context. + + That is a call from code outside any formula, or from a thread a formula + started without copying its context: then a formula in one of the + simulation's ancestors is waiting for the value, so every frame open in + those ancestors notes the read, as it would had the formula called + directly (as for the store history; see + ``Simulation._share_store_history_with_caller``). A frame keeps the first + epoch it notes for a simulation. + """ + _note_unobserved_activity() + ancestor = getattr(simulation, "parent_branch", None) + while ancestor is not None: + for open_frame in tuple(getattr(ancestor, "_open_frames", None) or ()): + open_frame.untracked_changes = True + open_reads = open_frame.reads + for other, epoch in reads: + if other is not ancestor: + open_reads.setdefault(other, epoch) + open_reads.setdefault(simulation, start) + ancestor = getattr(ancestor, "parent_branch", None) + + +def _value_signature(value: ArrayLike) -> Optional[tuple]: + """An immutable exact value comparison, including float signs and enum type. + + Object and string arrays have no supported convergence comparison. Bytes + avoid both hash collisions and later mutation of a retained array view. + """ + if np.ma.isMaskedArray(value): + return None + enum = value.possible_values if isinstance(value, EnumArray) else None + array = np.asarray(value) + if array.dtype.kind not in "biufcmM": + return None + return array.dtype.str, array.shape, enum, array.tobytes() + + +def _branch_registrations(simulation: "Simulation") -> set: + """Named branch links present when a calculation starts or a value is read. + + A standalone clone need not occupy its source's registration, but its + registered ancestors still constrain reuse through that named path. + """ + registrations = set() + while getattr(simulation, "parent_branch", None) is not None: + parent = simulation.parent_branch + name = simulation.branch_name + if parent.branches.get(name) is simulation: + registrations.add((parent, name, simulation)) + simulation = parent + return registrations + + +def _registered_creations(branches: dict) -> dict: + """Observed creations still registered under their original parent and name.""" + return { + branch: (parent, name) + for branch, (parent, name) in branches.items() + if branch.parent_branch is parent + and branch.branch_name == name + and parent.branches.get(name) is branch + } + + +def _branch_path( + simulation: "Simulation", created_branches: Optional[set] = None +) -> tuple: + """Identify recreated ``get_branch`` branches, keeping direct clones distinct. + + A direct clone retains its source's name and parent but has independent + inputs. It anchors a new path by identity, even if later registered under + the source's name; only branches made by ``get_branch`` compare by name. + """ + names = [] + while ( + getattr(simulation, "parent_branch", None) is not None + and getattr(simulation, "_fixed_point_branch", False) + and (created_branches is None or simulation in created_branches) + ): + names.append(simulation.branch_name) + simulation = simulation.parent_branch + return simulation, tuple(reversed(names)) + + +def _note_input_change( + simulation: "Simulation", + operation: str, + variable: str, + period: Optional[Period], + branch_name: str, + value: Optional[ArrayLike] = None, +) -> None: + """Record actual stores and deletes in each frame observing this context.""" + frames = _calculation_frames.get() + if not frames: + _taint_waiting_frames(simulation) + return + # Deletes and raw holder writes need not change an input epoch. They can + # still change what a foreign calculation reads during this attempt. + for frame in frames: + frame.attempt_cache.clear() + signature = _value_signature(value) if operation == "set" else None + if operation == "set" and signature is None: + for frame in frames: + frame.untracked_changes = True + return + change = ( + operation, + simulation, + branch_name, + variable, + period, + signature, + ) + for frame in frames: + frame.note_branch(simulation) + if ( + simulation is frame.simulation + or getattr(simulation, "parent_branch", None) is None + ): + frame.untracked_changes = True + frame.changes.append(change) + + +def _taint_waiting_frames(simulation: "Simulation") -> None: + """An unobserved mutation cannot certify a waiting calculation's fixed point. + + A thread without a formula context can write before the first read, so + counting only drops since that read would miss its mutation entirely. + The simulation and its ancestors may have frames waiting on that thread. + """ + _note_unobserved_activity() + while simulation is not None: + for frame in tuple(getattr(simulation, "_open_frames", None) or ()): + frame.untracked_changes = True + simulation = getattr(simulation, "parent_branch", None) + + @dataclass(frozen=True) class PreservedUserInput: variable_name: str @@ -196,6 +690,17 @@ class Simulation: default_tax_benefit_system: Type["TaxBenefitSystem"] = None """The default tax-benefit system class to use if none is provided.""" + # Counters read on every calculation (see ``set_input``). Class defaults, + # so a subclass that skips ``__init__`` reads them too; incrementing one + # gives the simulation its own. + _input_epoch: int = 0 # drops on this simulation + _inputs_set: int = 0 # inputs stored through ``Holder.set_input`` + _calculations_in_flight: int = 0 # ``calculate``/``calculate_add`` running + _calculations_started: int = 0 # ``_calculate`` calls begun + _requested_variables: Optional[set] = None + _store_history: Optional[StoreHistory] = None # see ``_get_store_history`` + _open_frames: Optional[set] = None # the ``_Frame``s open here, any thread + default_tax_benefit_system_instance: "TaxBenefitSystem" = None """The default tax-benefit system instance to use if none is provided. This requires that the tax-benefit system is initialised when importing a country package. This will slow down the import, but may speed up individual simulations.""" @@ -307,6 +812,9 @@ def __init__( # values it deletes, and ``clone`` (so also ``get_branch``) gives the # copy its own record. self._user_input_keys: set[tuple[str, str, Period]] = set() + # What this simulation's values may have been calculated from; each + # branch starts with a copy (see ``set_input``). + self._store_history = StoreHistory() self.debug: bool = False self.trace: bool = trace self.tracer: SimpleTracer = SimpleTracer() if not trace else FullTracer() @@ -445,7 +953,6 @@ def _invalidate_all_caches(self) -> None: wiped so the next ``calculate`` recomputes under the new tax-benefit system. """ - self._fast_cache = {} self.invalidated_caches = set() # Snapshot user-provided inputs before wiping so they can be # replayed into the fresh storage. Use the storage API instead of @@ -515,16 +1022,25 @@ def _invalidate_all_caches(self) -> None: for holder in population._holders.values(): holder._memory_storage._arrays = {} holder._memory_storage._unmark_dropped_keys() + holder._memory_storage._stop_sharing_dropped_keys() + holder._memory_storage._forget_dropped_numbers() if holder._disk_storage is not None: holder._disk_storage._files = {} holder._disk_storage._derived = set() + holder._disk_storage._forget_dropped_keys() if hasattr(holder, "_user_input_storage"): holder._user_input_storage = {} + for frame in _calculation_frames.get(): + frame.untracked_changes = True + self._drop_computed() # Replay preserved user inputs so ``calculate`` still sees them. for user_input in preserved: holder = self.get_holder(user_input.variable_name) if user_input.storage == "disk" and holder._disk_storage is not None: holder._disk_storage._files[user_input.disk_key] = user_input.disk_file + holder._disk_storage._sequence_numbers[user_input.disk_key] = ( + next_sequence_number() + ) if user_input.disk_enum is not None: holder._disk_storage._enums[user_input.disk_file] = ( user_input.disk_enum @@ -538,6 +1054,9 @@ def _invalidate_all_caches(self) -> None: holder._record_input_storage( user_input.period, user_input.branch_name, user_input.storage ) + for population in self.populations.values(): + for holder in population._holders.values(): + holder._record_inputs(None) for branch in self.branches.values(): branch._invalidate_all_caches() @@ -881,6 +1400,13 @@ def calculate( if _fast_cache is not None: _cached = _fast_cache.get(_fast_key) if _cached is not None: + self._share_store_history_with_caller() + frames = _calculation_frames.get() + if not frames: + _hand_to_waiting_ancestors(self, self._input_epoch, ()) + elif frames[-1].simulation is not self: + frames[-1].reads.setdefault(self, self._input_epoch) + frames[-1].registrations.update(_branch_registrations(self)) return _cached self.tracer.record_calculation_start(variable_name, period, self.branch_name) @@ -890,8 +1416,63 @@ def calculate( # check_formula_determinism), so there is nothing to make reproducible # here. + # Formulas running in this simulation may hold values they read in + # local variables; a drop meanwhile must not forget what they came + # from (see ``_drop_computed``). + # Only the outermost calculation in this simulation runs again after + # an input change here: running it again runs the inner ones too. + outermost = not self._calculations_in_flight + self._calculations_in_flight += 1 + frame = _frame_for(self, outermost) + if frame is not None: + frame.__enter__() try: result = self._calculate(variable_name, period) + previous_transition = None + for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): + if not frame.is_stale(): + break + transition = frame.transition_signature(result) + if transition is not None and transition == previous_transition: + frame.settle_reads() + # This run repeated the complete input transition and its + # result. Preserve holder input/blacklist rules and the + # existing fast-cache guard when keeping the fixed point. + keep_state = self._calculation_start() + result = self._cache_result( + self.get_holder(variable_name), + result, + period, + keep_state, + ) + if self._may_keep(keep_state): + if hasattr(self, "_fast_cache"): + self._fast_cache[(variable_name, period)] = result + if self.check_macro_cache(variable_name, str(period)): + macro = SimulationMacroCache(self.tax_benefit_system) + macro.set_cache_path( + self.dataset.file_path.parent, + self.dataset.name, + variable_name, + str(period), + self.branch_name, + ) + macro.set_cache_value(macro.get_cache_path(), result) + break + previous_transition = transition + # An input set while it ran (by a formula, say), here or in a + # simulation it read from, dropped values, so the result was + # not kept (``_cache_result``). Calculate it again from the new + # inputs, until a run changes none, and keep that result, as a + # simulation given those inputs first would. + frame.restart() + result = self._calculate(variable_name, period) + # Satisfies ``requires_computation_after`` from now on, even if a + # branch input later drops the values. + requested = self._requested_variables + if requested is None: + requested = self._get_requested_variables() + requested.add(variable_name) if isinstance(result, EnumArray) and decode_enums: result = result.decode_to_str() self.tracer.record_calculation_result(result) @@ -902,8 +1483,16 @@ def calculate( result = self.map_result(result, source_entity, map_to) return result finally: - self.tracer.record_calculation_end() - self.purge_cache_of_invalid_values() + try: + # Also when the calculation fails: whether it fails can depend + # on what it read, and a calling formula may catch the error. + self._share_store_history_with_caller() + finally: + self._calculations_in_flight -= 1 + if frame is not None: + frame.__exit__(None, None, None) + self.tracer.record_calculation_end() + self.purge_cache_of_invalid_values() def map_result( self, @@ -1015,6 +1604,9 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: """ if variable_name not in self.tax_benefit_system.variables: raise ValueError(f"Variable {variable_name} does not exist.") + input_state = self._calculation_start() + # Lets a custom ``set_input`` handler tell whether it calculated. + self._calculations_started = self._calculations_started + 1 population = self.get_variable_population(variable_name) holder = population.get_holder(variable_name) variable = self.tax_benefit_system.get_variable( @@ -1038,6 +1630,12 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if cached_array is not None: return cached_array + frames = _calculation_frames.get() + if frames and frames[-1].simulation is self: + cached_array = frames[-1].recall(variable_name, period) + if cached_array is not None: + return cached_array + # Check if cache can be used, if available, check if path exists is_cache_available = self.check_macro_cache(variable_name, str(period)) if is_cache_available: @@ -1057,6 +1655,13 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: value = smc.get_cache_value(cache_path) if value is not None: + # Served without being stored: record it, so values + # calculated from it count as later (see ``set_input``). + # What it was calculated from was never calculated here, + # so values from here on may depend on any input. + sequence_number = next_sequence_number() + holder._record_store(period, sequence_number) + self._get_store_history().record_unknown_sources(sequence_number) return value if variable.requires_computation_after is not None: @@ -1067,7 +1672,16 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: required_is_known_periods = self.get_holder( variable.requires_computation_after ).get_known_periods() - if (not variable_in_stack) and (not len(required_is_known_periods) > 0): + # A branch input may have dropped the prerequisite's values; it + # was still requested. + required_was_requested = ( + variable.requires_computation_after in self._get_requested_variables() + ) + if ( + (not variable_in_stack) + and (not len(required_is_known_periods) > 0) + and not required_was_requested + ): raise ValueError( f"Variable {variable_name} requires {variable.requires_computation_after} to be requested first. That variable is known in: {required_is_known_periods}. The full stack is: {variables_in_stack}. {variable_in_stack, len(required_is_known_periods) > 0}" ) @@ -1087,7 +1701,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: values = self.calculate_divide(variable_name, period) if alternate_period_handling: - if is_cache_available: + if is_cache_available and self._may_keep(input_state): smc.set_cache_value(cache_path, values) return values @@ -1112,18 +1726,31 @@ 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, derived=True) - return array + return self._cache_result(holder, array, period, input_state) array = None # First, try to run a formula try: self._check_for_cycle(variable.name, period) - array = self._run_formula(variable, population, period) + token = _formula_simulation.set(self) + try: + array = self._run_formula(variable, population, period) + finally: + _formula_simulation.reset(token) # If no result, use the default value and cache it if array is None: + if variable.uprating is not None or ( + self.tax_benefit_system.auto_carry_over_input_variables + and variable.calculate_output is None + ): + # The value is uprated or carried over from another + # period, or defaults for lack of one: an input set later + # for another period of the variable can change it. + self._get_store_history().record_derived( + variable_name, next_sequence_number() + ) # Check if the variable has a previously defined value known_periods = holder.get_known_periods() earlier_known_periods = [ @@ -1262,11 +1889,12 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: array = EnumArray(array, variable.possible_values) array = self._cast_formula_result(array, variable) - # Calculated, not input: auto-carry-over never carries it. - holder.put_in_cache(array, period, self.branch_name, derived=True) + array = self._cache_result(holder, array, period, input_state) except SpiralError: array = holder.default_array() + # Not stored, but what reads it is (see ``set_input``). + holder._record_store(period, next_sequence_number()) except RecursionError as e: if isinstance(self.tracer, FullTracer): self.tracer.print_computation_log() @@ -1281,14 +1909,84 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: f"RecursionError while calculating {variable_name} for period {period}. The full computation stack is:\n{stack_formatted}" ) - if is_cache_available: + # Neither cache keeps a result an input change may have made obsolete + # (see ``_cache_result``). + unchanged = self._may_keep(input_state) + if is_cache_available and unchanged: smc.set_cache_value(cache_path, array) - if hasattr(self, "_fast_cache"): + if hasattr(self, "_fast_cache") and unchanged: self._fast_cache[(variable_name, period)] = array return array + def _may_keep(self, input_state: Tuple[int, int, int]) -> bool: + """Whether a result calculated since ``input_state`` may be kept (see ``_cache_result``).""" + if input_state[0] != self._input_epoch or input_state[1] != self._inputs_set: + return False + frames = _calculation_frames.get() + return not frames or frames[-1].may_keep() + + def _calculation_start(self) -> Tuple[int, int, int]: + """When a calculation begins: how many drops ran here (``_input_epoch``) + and inputs were set here (``_inputs_set``), and a sequence number.""" + return self._input_epoch, self._inputs_set, next_sequence_number() + + def _input_set_meanwhile( + self, holder: Holder, period: Period, started_at: int + ) -> Optional[ArrayLike]: + """The input this simulation reads for ``period``, if it was stored after ``started_at``. + + That is an input set while the calculation that began at + ``started_at`` ran (by its own formula, say), under any branch name + the simulation reads. + """ + stored_on = holder._branch_storing(period, self.branch_name) + if ( + stored_on is not None + and holder._is_input(period, stored_on) + and (holder._stored_sequence_number(period, stored_on) or 0) > started_at + ): + return holder._get_array_from_storage(period, stored_on) + return None + + def _cache_result( + self, + holder: Holder, + array: ArrayLike, + period: Period, + input_state: Tuple[int, int, int], + ) -> ArrayLike: + """Cache a calculated value, and return the value to use for it. + + ``input_state`` is :meth:`_calculation_start` when the calculation began. + If an input for the same period was set meanwhile (by the formula + itself, say), that input is the value, as it would be had it been set + first. A calculation that was running when an input set on this + simulation dropped values may have read the replaced value, so its + result is returned but not kept (see ``_drop_computed``). + """ + epoch, inputs_set, started_at = input_state + if inputs_set != self._inputs_set: + stored_input = self._input_set_meanwhile(holder, period, started_at) + if stored_input is not None: + return stored_input + frames = _calculation_frames.get() + may_keep = not frames or frames[-1].may_keep() + if epoch == self._input_epoch and may_keep: + holder.put_in_cache(array, period, self.branch_name, derived=True) + else: + # Not kept, but whatever reads it is stored after it all the same. + holder._record_store(period, next_sequence_number()) + if ( + frames + and frames[-1].simulation is self + and epoch == self._input_epoch + and inputs_set == self._inputs_set + ): + frames[-1].remember(holder.variable.name, period, array) + return array + def purge_cache_of_invalid_values(self) -> None: # We wait for the end of calculate(), signalled by an empty stack, before purging the cache if self.tracer.stack: @@ -1339,11 +2037,39 @@ def calculate_add( ) sub_periods = list(period.get_subperiods(variable.definition_period)) - result = sum( - self.calculate(variable_name, sub_period) for sub_period in sub_periods - ) - self._cache_option_result(variable, period, result) - return result + + def total(): + return sum( + self.calculate(variable_name, sub_period) for sub_period in sub_periods + ) + + # As in ``calculate``: only an outermost sum runs again, and while it + # sums it is in flight, so its terms do not run again on their own + # (and a drop meanwhile keeps the records of what they read). + outermost = not self._calculations_in_flight + self._calculations_in_flight += 1 + frame = _frame_for(self, outermost) + try: + with frame if frame is not None else nullcontext(): + input_state = self._calculation_start() + result = total() + previous_transition = None + for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): + if not frame.is_stale(): + break + transition = frame.transition_signature(result) + if transition is not None and transition == previous_transition: + frame.settle_reads() + break + previous_transition = transition + # An input changed while summing: earlier terms may be + # obsolete (see ``calculate``). Sum again. + frame.restart() + input_state = self._calculation_start() + result = total() + return self._cache_option_result(variable, period, result, input_state) + finally: + self._calculations_in_flight -= 1 def calculate_divide( self, @@ -1372,10 +2098,12 @@ def calculate_divide( ) if period.unit == periods.MONTH: - computation_period = period.this_year - result = self.calculate(variable_name, period=computation_period) / 12.0 - self._cache_option_result(variable, period, result) - return result + frame = _frame_for(self, outermost=False) + with frame if frame is not None else nullcontext(): + input_state = self._calculation_start() + computation_period = period.this_year + result = self.calculate(variable_name, period=computation_period) / 12.0 + return self._cache_option_result(variable, period, result, input_state) elif period.unit == periods.YEAR: return self.calculate(variable_name, period) @@ -1386,8 +2114,12 @@ def calculate_divide( ) def _cache_option_result( - self, variable: Variable, period: Period, result: ArrayLike - ) -> None: + self, + variable: Variable, + period: Period, + result: ArrayLike, + input_state: Tuple[int, int, int], + ) -> ArrayLike: """Cache an ADD or DIVIDE result at ``period`` if a plain read would return it. A value cached at ``period`` is what every later ``calculate`` of the @@ -1424,14 +2156,14 @@ def _cache_option_result( writes never register the aggregate as a supplied input. """ if variable.quantity_type == QuantityType.STOCK: - return + return result routed = (variable.definition_period == MONTH and period.unit == YEAR) or ( variable.definition_period == YEAR and period.unit == MONTH ) if not routed or np.asarray(result).dtype != variable.dtype: - return - self.get_holder(variable.name).put_in_cache( - result, period, self.branch_name, derived=True + return result + return self._cache_result( + self.get_holder(variable.name), result, period, input_state ) def calculate_output(self, variable_name: str, period: Period = None) -> ArrayLike: @@ -1739,6 +2471,30 @@ def delete_arrays(self, variable: str, period: Period = None) -> None: if _fast_cache is not None: _fast_cache.pop((variable, period), None) + def drop_computed_arrays(self) -> int: + """Delete every value this simulation holds except inputs. + + Inputs are the values stored through ``set_input``: the dataset or + situation the simulation was built from, inputs set on it and, for a + branch, inputs set on the simulations it was created from before it + was created. Values a custom ``set_input`` handler calculates are not + inputs. Every other value is calculated again when next requested, + and the simulation stops reading macro-cache files. + + Use this on a branch whose tax-benefit system or parameters differ + from its parent's, whose values the branch would otherwise inherit; + ``set_input`` on a branch drops what depends on the input by itself. + Branches already created from this simulation keep their values. + + Returns: + int: The number of arrays deleted. + """ + # Recalculate rather than read macro-cache files written before. + self.macro_cache_read = False + for frame in _calculation_frames.get(): + frame.untracked_changes = True + return self._drop_computed() + def get_known_periods(self, variable: str) -> List[Period]: """ Get a list variable's known period, i.e. the periods where a value has been initialized and @@ -1772,6 +2528,55 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non array([12, 14], dtype=int32) If a ``set_input`` property has been set for the variable, this method may accept inputs for periods not matching the ``definition_period`` of the variable. To read more about this, check the `documentation `_. + + On a branch (see :meth:`get_branch`), the input also drops what the + branch holds that may have been calculated from the value it + replaces, so what the branch calculates next uses the input, as a + simulation given the input before calculating anything would. Every + value is stored after everything it was calculated from, and each + simulation records the first stores its values may have been + calculated from (see :mod:`policyengine_core.data_storage.store_history`): + its own, its parent's when it was created, and those of simulations + its formulas calculated in. The branch drops each value it holds, + other than an input, stored at or after the earliest recorded store of + ``variable_name`` for a period that shares a day with ``period``, or + the earliest recorded value of ``variable_name`` uprated or carried + over from another period (or given the default for want of one), or + the first value that may depend on anything (restored from a dump, + or calculated from a macro-cache read). If there is none of these, + nothing it holds depends on the value and nothing is dropped. That is the case when a formula creates the branch while + still calculating the variable it overrides, unless another branch + it created for the same comparison already returned a value + calculated from the variable; then only what came back from there, + and what was calculated after, is dropped. + + Once an input is set on it, the branch stops reading macro-cache + files, which are keyed by branch name and period but not by inputs. + Each drop counts as an input change of the branch. A calculation + that may have read a value from before the change, running in the + branch or in any simulation that got a value from it, is not kept; + the outermost such calculation in each simulation runs again until it + reads nothing changed since (at most ten times; after that its result + is returned but not kept), and an input stored for the very period it + calculates after it began is its result. Setting an input to the + value the branch already reads for that period, as an input, drops + nothing. + + What this does not track: a branch given a different tax-benefit + system or parameters (call :meth:`drop_computed_arrays` on it); + formulas that write into an array they read instead of returning a + new one; formulas that test whether a value is stored + (``get_known_periods``, ``get_array``) or read another simulation's + storage directly, rather than calculating; a simulation other than + the formula's own branches calculated from a thread the formula + starts without copying its context; and branches a formula keeps + between calls, which hold what their parent held when they were + created. On any simulation, input helpers remove overlapping + calculated values of the variable they write. Dependency invalidation + runs only on branches; branches already created from the branch keep + their values; and a value another simulation calculated from the + branch and kept stays when the branch's input changes after that + calculation ended. """ period = periods.period(period) if self.start_instant is None or self.start_instant > period.start: @@ -1785,21 +2590,89 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non _fast_cache = getattr(self, "_fast_cache", None) if _fast_cache is not None: _fast_cache.pop((variable_name, period), None) - if variable.set_input and period.unit != variable.definition_period: - # The helper wrote the input's sub-periods, replacing any - # value calculated there, so what ``calculate`` returned for - # them is stale too. (``_end_order``, not ``stop``: ``stop`` - # raises for a period that ends after 9999-12-31.) - stale = [ - key - for key in _fast_cache - if key[0] == variable_name - and isinstance(key[1], Period) - and period.start <= key[1].start - and _end_order(key[1]) <= _end_order(period) - ] - for key in stale: - del _fast_cache[key] + + def _get_store_history(self) -> StoreHistory: + history = self._store_history + if history is None: + history = self._store_history = StoreHistory() + return history + + def _get_requested_variables(self) -> set: + requested = self._requested_variables + if requested is None: + requested = self._requested_variables = set() + return requested + + def _share_store_history_with_caller(self) -> None: + """Merge this simulation's store history into that of a formula calling it.""" + caller = _formula_simulation.get() + if caller is not None: + if caller is not self: + caller._get_store_history().merge(self._get_store_history()) + return + # No formula is visible here: the call comes from code outside any + # formula, or from a thread a formula started without copying its + # context. In the second case a formula in one of this simulation's + # ancestors is waiting for the result, so give it to every ancestor + # with a calculation running (more records only make drops broader). + ancestor = getattr(self, "parent_branch", None) + while ancestor is not None: + if ancestor._calculations_in_flight: + ancestor._get_store_history().merge(self._get_store_history()) + ancestor = getattr(ancestor, "parent_branch", None) + + def _drop_values_that_may_depend_on( + self, variable_name: str, period: Period + ) -> int: + """On a branch, drop what may depend on ``variable_name`` at ``period``. + + Called before an input for ``variable_name`` at ``period`` is stored; + see :meth:`set_input`. Returns the number of arrays dropped. + """ + if getattr(self, "parent_branch", None) is None: + return 0 + # (``Holder.set_input`` has already stopped macro-cache reads here.) + since = self._get_store_history().earliest_dependency( + variable_name, periods.period(period) + ) + if self.get_holder(variable_name)._has_unnumbered_values(): + # Written into storage directly, so neither numbered nor recorded: + # anything may have been calculated from it. + since = 0 + if since is None: + return 0 + return self._drop_computed(since) + + def _drop_computed(self, since: Optional[int] = None) -> int: + """Drop every non-input value numbered ``since`` or later (all, without it).""" + dropped = 0 + for population in self.populations.values(): + for holder in population._holders.values(): + dropped += holder._drop_computed(since) + # The fast cache can also hold values a holder does not keep. + self._fast_cache = {} + # A calculation that got a value from here before the drop, running + # here or in another simulation, may hold it or what it calculated + # from it: none of them keeps its result (``_Frame.is_stale``), and + # the outermost one in each simulation runs again (``calculate``). + self._input_epoch = self._input_epoch + 1 + frames = _calculation_frames.get() + if not frames: + _taint_waiting_frames(self) + for frame in frames: + frame.drop_epochs.setdefault(self, set()).add(self._input_epoch) + if self._calculations_in_flight: + # A formula running here may still hold values calculated from + # what the records describe: keep the records. + return dropped + # Nothing the simulation still holds was calculated from what the + # records numbered ``since`` or later describe, except the inputs it + # keeps, which are recorded again. + self._get_store_history().prune(since) + for population in self.populations.values(): + for holder in population._holders.values(): + holder._record_inputs(since) + return dropped def get_variable_population(self, variable_name: str) -> Population: variable = self.tax_benefit_system.get_variable( @@ -1956,6 +2829,17 @@ def clone( new.tax_benefit_system = self.tax_benefit_system new.debug = debug new.trace = trace + # The copy holds what this simulation holds, so it starts from what + # those values may have been calculated from, and diverges from there. + new._store_history = self._get_store_history().copy() + new._requested_variables = set(self._get_requested_variables()) + # Calculations running in this simulation are not running in the copy. + new._calculations_in_flight = 0 + new._open_frames = set() + # A direct clone can retain its source's branch name and parent while + # taking different inputs. ``get_branch`` marks the returned branch + # after assigning its ancestry, for comparable recreated branches. + new._fixed_point_branch = False # A ``set_input`` running on this simulation is not running on the copy. new._user_input_contexts = [] @@ -2023,6 +2907,13 @@ def get_branch( several threads at once: two first reads of the same array can each make a copy. + A new branch starts with the values this simulation holds; asked for + a name it already has, this returns that branch as it is. An input + set on the branch drops those that may have been calculated from the + value it replaces (see :meth:`set_input`). A branch whose + tax-benefit system or parameters are changed should call + :meth:`drop_computed_arrays` before calculating. + Args: name (str, optional): Name of the branch. Defaults to "branch". clone_system (bool, optional): Whether to clone the tax-benefit system. Use this if you're changing policy parameters. Defaults to False. @@ -2033,7 +2924,13 @@ def get_branch( if name == self.branch_name: return self if name in self.branches: - return self.branches[name] + branch = self.branches[name] + frames = _calculation_frames.get() + if frames: + # The lookup name can differ from the branch object's own + # name, so retain the registration actually resolved. + frames[-1].registrations.add((self, name, branch)) + return branch request = _BranchClone(self) token = _branch_clone.set(request) try: @@ -2044,6 +2941,13 @@ def get_branch( self.branches[name] = branch branch.branch_name = name branch.parent_branch = self + branch._fixed_point_branch = True + frames = _calculation_frames.get() + if not frames: + _note_unobserved_activity() + for frame in frames: + frame.attempt_cache.clear() + frame.created_branches[branch] = (self, name) if self.trace: branch.trace = True branch.tracer = self.tracer @@ -2071,9 +2975,7 @@ def derivative( period = periods.period(self.default_calculation_period) alt_sim = self.clone() - for computed_variable in alt_sim.tax_benefit_system.variables: - if computed_variable not in self.input_variables: - alt_sim.delete_arrays(computed_variable) + alt_sim.drop_computed_arrays() alt_sim.set_input(wrt, period, self.calculate(wrt, period) + delta) original_value = self.calculate(variable, period) new_value = alt_sim.calculate(variable, period) @@ -2445,10 +3347,9 @@ def subsample( df = subset_df - # Update the dataset and rebuild the simulation. Rebuilding replaces - # every stored value, so the record of inputs starts again with the - # ones the rebuild sets. + # Rebuilding replaces storage, its input record, and store history. self.dataset = Dataset.from_dataframe(df, self.dataset.time_period) + self._store_history = StoreHistory() self._user_input_keys = set() self.build_from_dataset() diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index d17c42905..daa5045e8 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -7,6 +7,7 @@ import numpy as np from policyengine_core.data_storage import OnDiskStorage +from policyengine_core.data_storage.store_history import next_sequence_number from policyengine_core.periods import ETERNITY from policyengine_core.simulations import Simulation @@ -31,8 +32,8 @@ def dump_simulation(simulation, directory): """ Write simulation data to directory, so that it can be restored later. - Writes the default branch's values, listing the periods of those the - simulation calculated (see ``restore_simulation``). + Writes the simulation's branch values, including inherited values, and + lists the periods it calculated (see ``restore_simulation``). """ parent_directory = os.path.abspath(os.path.join(directory, os.pardir)) if not os.path.isdir(parent_directory): # To deal with reforms @@ -103,9 +104,22 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): lists_calculated_values = os.path.exists( os.path.join(entities_dump_dir, DUMP_FORMAT_FILE) ) + # Inputs first, then every calculated value under one later number. The + # dump does not say what each value was calculated from (nor what was read + # without being kept), so an input set for any variable on a branch of the + # restored simulation drops them all. for variable in variables_to_restore: - if _restore_holder(simulation, variable, directory): - lists_calculated_values = True + _restore_holder(simulation, variable, directory, inputs=True) + calculated_number = next_sequence_number() + restored = sum( + _restore_holder( + simulation, variable, directory, calculated_number=calculated_number + ) + for variable in variables_to_restore + ) + if restored: + simulation._get_store_history().record_unknown_sources(calculated_number) + lists_calculated_values = True if variables_to_restore and not lists_calculated_values: warnings.warn( f"The simulation dump in {directory} does not say which of its " @@ -123,13 +137,18 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): def _dump_holder(holder, directory): disk_storage = holder.create_disk_storage(directory, preserve=True) + branch_name = holder.simulation.branch_name derived_periods = set() - for period in holder.get_known_periods(): - value = holder.get_array(period) - disk_storage.put(value, period) + for period in dict.fromkeys(holder.get_known_periods()): + # What the simulation itself reads: on a branch, its own value, else + # its nearest ancestor's or the default one. + stored_on = holder._branch_storing(period, branch_name) + if stored_on is None: + continue + disk_storage.put(holder._get_array_from_storage(period, stored_on), period) # Read the mark of exactly the value dumped: the same period on the - # same branch as ``get_array``. - if holder.is_derived(period): + # same branch. + if holder.is_derived(period, branch_name): derived_periods.add(str(period)) if derived_periods: path = os.path.join(disk_storage.storage_dir, DERIVED_PERIODS_FILE) @@ -188,10 +207,14 @@ def _restore_entity(population, directory): return person_count -def _restore_holder(simulation, variable, directory): - """Restore one variable's values: those its ``DERIVED_PERIODS_FILE`` - lists as calculated values, every other one as an input. Return whether - it lists any.""" +def _restore_holder( + simulation, variable, directory, inputs=False, calculated_number=None +): + """Restore one variable's inputs or calculated values and return the count. + + Inputs are registered as supplied values; calculated values share the + later ``calculated_number`` because a dump does not keep their sources. + """ storage_dir = os.path.join(directory, variable) is_variable_eternal = ( simulation.tax_benefit_system.get_variable(variable).definition_period @@ -203,20 +226,31 @@ def _restore_holder(simulation, variable, directory): disk_storage.restore() holder = simulation.get_holder(variable) - + # Dumped before derived values were recorded, there is no such file and + # every value counts as an input, as it did then. 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()) + restored = 0 for period in disk_storage.get_known_periods(): + derived = str(period) in derived_periods + if derived == inputs: + continue + restored += 1 value = disk_storage.get(period) - if str(period) in derived_periods: - holder.put_in_cache(value, period, derived=True) + if derived: + holder._set( + period, + value, + derived=True, + sequence_number=calculated_number, + ) else: _restore_input(simulation, holder, period, value) - return bool(derived_periods) + return restored def _restore_input(simulation, holder, period, value): diff --git a/tests/core/test_branch_helper_invalidation.py b/tests/core/test_branch_helper_invalidation.py new file mode 100644 index 000000000..ba66b7847 --- /dev/null +++ b/tests/core/test_branch_helper_invalidation.py @@ -0,0 +1,137 @@ +"""Branch dependency invalidation composes with input-helper overlap cleanup.""" + +from collections import Counter +import warnings + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem, entities +from policyengine_core.experimental import MemoryConfig +from policyengine_core.holders import ( + set_input_dispatch_by_period, + set_input_divide_by_period, +) +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables import Variable + + +@pytest.mark.parametrize( + "helper,year_input", + [ + pytest.param(set_input_divide_by_period, [120.0, 240.0], id="divide"), + pytest.param(set_input_dispatch_by_period, [10.0, 20.0], id="dispatch"), + ], +) +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +@pytest.mark.parametrize("year_string", ["2024", "0025"]) +def test_branch_helper_drops_inherited_aggregates_and_dependents( + helper, year_input, on_disk, year_string +): + calls = Counter() + + class helper_flow(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + set_input = helper + label = "Monthly input handled over a year" + + class helper_result(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "Annual result depending on the monthly input aggregate" + + def formula(person, period): + calls[("result", person.simulation.branch_name)] += 1 + return person("helper_flow", period) * 2 + + class helper_independent(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "Independent result calculated before the input was stored" + + def formula(person, period): + calls[("independent", person.simulation.branch_name)] += 1 + return person.filled_array(7.0) + + system = CountryTaxBenefitSystem() + variables = (helper_flow, helper_result, helper_independent) + system.add_variables(*variables) + + def build(): + simulation = SimulationBuilder().build_default_simulation(system, count=2) + if on_disk: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + # Holders take their storage configuration when first created. + for variable in variables: + population = simulation.get_variable_population(variable.__name__) + population._holders.pop(variable.__name__, None) + return simulation + + year = periods.period(year_string) + january = year.first_month + months = [january.offset(index) for index in range(12)] + parent = build() + parent.calculate("helper_independent", year) + parent.set_input("helper_flow", january, [12.0, 24.0]) + parent_months = [january] + if year_string == "0025": + # Early-year fallback getters have a separate legacy parser issue. + # Supply every month so this case exercises history and input replay. + for month in months[1:]: + parent.set_input("helper_flow", month, [0.0, 0.0]) + parent_months = months + parent_inputs = {("helper_flow", "default", month) for month in parent_months} + np.testing.assert_array_equal(parent.calculate("helper_flow", year), [12, 24]) + np.testing.assert_array_equal(parent.calculate("helper_result", year), [24, 48]) + + branch = parent.get_branch("changed") + branch.set_input("helper_flow", year, np.asarray(year_input)) + fresh = build() + fresh.set_input("helper_flow", year, np.asarray(year_input)) + + for name in ("helper_flow", "helper_result"): + np.testing.assert_array_equal( + branch.calculate(name, year), fresh.calculate(name, year) + ) + for month in months: + np.testing.assert_array_equal(branch.calculate("helper_flow", month), [10, 20]) + np.testing.assert_array_equal(branch.calculate("helper_independent", year), [7, 7]) + # The round-2 calculation frames keep independent inherited caches. + assert calls[("independent", "changed")] == 0 + settled_calls = calls.copy() + for _ in range(2): + np.testing.assert_array_equal( + branch.calculate("helper_result", year), [240, 480] + ) + assert calls == settled_calls + + np.testing.assert_array_equal(parent.calculate("helper_flow", january), [12, 24]) + np.testing.assert_array_equal(parent.calculate("helper_flow", year), [12, 24]) + np.testing.assert_array_equal(parent.calculate("helper_result", year), [24, 48]) + + branch_inputs = {("helper_flow", "changed", month) for month in months} + assert branch_inputs | parent_inputs == branch._user_input_keys + assert parent._user_input_keys == parent_inputs + tier = "disk" if on_disk else "memory" + holder = branch.get_holder("helper_flow") + for month in months: + assert holder._stores_user_input(month, "changed", tier) + + branch._invalidate_all_caches() + + assert branch._user_input_keys == branch_inputs | parent_inputs + for month in months: + assert holder._stores_user_input(month, "changed", tier) + np.testing.assert_array_equal(branch.calculate("helper_flow", month), [10, 20]) + np.testing.assert_array_equal(branch.calculate("helper_result", year), [240, 480]) + np.testing.assert_array_equal(parent.calculate("helper_result", year), [24, 48]) + replayed_calls = calls.copy() + branch.calculate("helper_result", year) + assert calls == replayed_calls diff --git a/tests/core/test_branch_input_fixed_point.py b/tests/core/test_branch_input_fixed_point.py new file mode 100644 index 000000000..b51f4964e --- /dev/null +++ b/tests/core/test_branch_input_fixed_point.py @@ -0,0 +1,444 @@ +"""Repeated branch resets settle without forgetting real input changes. + +The fixed-point result must match a fresh branch given its override before +calculating. Equal results alone do not prove convergence: the branch's input +mutations must also settle, and changes outside the formula's context cannot +justify keeping a stale result. +""" + +from __future__ import annotations + +import contextvars +import threading + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem, entities +from policyengine_core.experimental import MemoryConfig +from policyengine_core.simulations import SimulationBuilder +import policyengine_core.simulations.simulation as simulation_module +from policyengine_core.variables import Variable + + +def _variable(name, definition_period, formula=None): + attributes = { + "value_type": float, + "entity": entities.Person, + "definition_period": definition_period, + "label": name, + } + if formula is not None: + attributes["formula"] = formula + return type(name, (Variable,), attributes) + + +def _in_thread(call): + values = [] + errors = [] + + def target(): + try: + values.append(call()) + except BaseException as error: + errors.append(error) + + context = contextvars.Context() + thread = threading.Thread(target=context.run, args=(target,)) + thread.start() + thread.join() + if errors: + raise errors[0] + return values[0] + + +def _simulation( + values, + definition_period=periods.YEAR, + delete_derived=False, + alternating=False, + own_input=False, + thread_operation=None, + masked_inputs=False, + recreate_branch=False, + cache_mode="memory", + branch_depth=1, +): + values = np.asarray(values, dtype=np.float32) + calls = [] + + def base(person, period): + source = "fp_thread_source" if thread_operation == "before_read" else "fp_x" + return np.abs(person(source, period)) / 2 + 5 + + def cap(person, period): + return np.full(person.count, 1000.0) + + def capped(person, period): + return np.minimum(person("fp_x", period), person("fp_cap", period)) + + def outer(person, period): + calls.append(str(period)) + simulation = person.simulation + # The branch inherits a record of the input it will replace. + person("fp_cap", period) + branch = simulation.get_branch("persistent") + for depth in range(1, branch_depth): + branch = branch.get_branch(f"nested_{depth}") + branch.delete_arrays("fp_cap") + if delete_derived: + branch.delete_arrays("fp_capped") + if thread_operation == "before_read": + _in_thread(lambda: branch.set_input("fp_thread_source", period, values + 1)) + + def read_base(): + return branch.calculate("fp_base", period) + + maximum = _in_thread(read_base) if thread_operation == "read" else read_base() + if alternating: + maximum = np.full(person.count, 1.0 + len(calls) % 2) + if masked_inputs: + maximum = np.ma.array(np.ones(person.count), mask=bool(len(calls) % 2)) + + def change(): + return branch.set_input("fp_cap", period, maximum) + + if thread_operation == "change": + _in_thread(change) + else: + change() + result = branch.calculate("fp_capped", period) + if masked_inputs: + result = np.asarray(result) * 0 + if recreate_branch: + del simulation.branches["persistent"] + if own_input: + simulation.set_input("fp_outer", period, np.full(person.count, 9.0)) + return result + + system = CountryTaxBenefitSystem() + system.add_variables( + _variable("fp_x", definition_period), + _variable("fp_thread_source", definition_period), + _variable("fp_base", definition_period, base), + _variable("fp_cap", definition_period, cap), + _variable("fp_capped", definition_period, capped), + _variable("fp_outer", definition_period, outer), + ) + simulation = SimulationBuilder().build_default_simulation(system, count=len(values)) + if cache_mode == "disk": + simulation.memory_config = MemoryConfig(0) + elif cache_mode == "drop": + simulation.memory_config = MemoryConfig(1, variables_to_drop=["fp_outer"]) + elif cache_mode == "blacklist": + system.cache_blacklist = ["fp_outer"] + simulation.opt_out_cache = True + if cache_mode in ("disk", "drop"): + # Holders capture MemoryConfig when constructed. Build them before + # setting inputs, including holders later copied into branches. + for population in simulation.populations.values(): + population._holders = {} + for variable in system.variables: + simulation.get_holder(variable) + input_periods = ( + ["2020"] + if definition_period == periods.YEAR + else [f"2020-{month:02d}" for month in range(1, 13)] + ) + for input_period in input_periods: + simulation.set_input("fp_x", input_period, values) + if thread_operation == "before_read": + # Its history follows cap's, so writing it preserves the cap + # dependency that the formula's later contextual set invalidates. + simulation.calculate("fp_cap", input_period) + simulation.set_input("fp_thread_source", input_period, values) + return simulation, calls + + +def _input_first_result(values): + fresh, _ = _simulation(values) + maximum = fresh.calculate("fp_base", "2020") + branch = fresh.get_branch("input_first") + branch.set_input("fp_cap", "2020", maximum) + return branch.calculate("fp_capped", "2020") + + +def _assert_settled(simulation, calls, expected, kept_in_holder=True): + result = simulation.calculate("fp_outer", "2020") + np.testing.assert_array_equal(result, expected) + assert len(calls) == 2 + stored = simulation.get_array("fp_outer", "2020") + if kept_in_holder: + np.testing.assert_array_equal(stored, expected) + else: + assert stored is None + completed = len(calls) + np.testing.assert_array_equal(simulation.calculate("fp_outer", "2020"), expected) + assert len(calls) == completed + + +@pytest.mark.parametrize("cache_mode", ["disk", "drop", "blacklist"]) +def test_fixed_point_respects_holder_storage_restrictions(cache_mode): + values = [0.0, 40.0] + simulation, calls = _simulation(values, cache_mode=cache_mode) + holder = simulation.get_holder("fp_outer") + assert holder._do_not_store == (cache_mode == "drop") + if cache_mode == "disk": + assert holder._on_disk_storable + _assert_settled( + simulation, + calls, + _input_first_result(values), + kept_in_holder=cache_mode == "disk", + ) + if cache_mode == "disk": + assert holder._memory_storage.get("2020") is None + np.testing.assert_array_equal( + holder._disk_storage.get("2020"), _input_first_result(values) + ) + + +@pytest.mark.parametrize("delete_derived", [False, True]) +@pytest.mark.parametrize("recreate_branch", [False, True]) +def test_branch_delete_read_set_reaches_a_cached_fixed_point( + delete_derived, recreate_branch +): + values = [-20.0, 0.0, 40.0] + simulation, calls = _simulation( + values, delete_derived=delete_derived, recreate_branch=recreate_branch + ) + _assert_settled(simulation, calls, _input_first_result(values)) + + +@pytest.mark.parametrize("branch_depth", [1, 2, 3]) +def test_reused_branch_settles_in_two_attempts(branch_depth): + values = [-20.0, 0.0, 40.0] + simulation, calls = _simulation(values, branch_depth=branch_depth) + _assert_settled(simulation, calls, _input_first_result(values)) + + +def test_replacing_reused_branch_registration_keeps_saved_snapshot_distinct(): + simulation, _ = _simulation([1.0]) + calls = [] + saved = simulation.get_branch("persistent") + saved.set_input("fp_x", "2020", np.array([0.0])) + saved.set_input("fp_cap", "2020", np.array([5.0])) + del simulation.branches["persistent"] + + def outer(person, period): + calls.append(str(period)) + simulation = person.simulation + person("fp_cap", period) + branch = simulation.get_branch("persistent") + old_cap = branch.calculate("fp_cap", period) + branch.delete_arrays("fp_cap") + maximum = branch.calculate("fp_base", period) + branch.set_input("fp_cap", period, maximum) + # The repeated transition swaps the registered reader to an older + # snapshot with different inputs. Its next result must still be read. + if np.all(old_cap <= maximum): + simulation.branches["persistent"] = saved + return np.minimum(branch.calculate("fp_x", period), 1) + + simulation.tax_benefit_system.add_variable( + _variable("fp_registration_outer", periods.YEAR, outer) + ) + assert simulation.calculate("fp_registration_outer", "2020").tolist() == [0.0] + assert len(calls) == 4 + assert simulation.get_array("fp_registration_outer", "2020").tolist() == [0.0] + assert simulation.calculate("fp_registration_outer", "2020").tolist() == [0.0] + assert len(calls) == 4 + + +def test_fixed_point_without_fast_cache_attribute(): + values = [0.0, 40.0] + simulation, calls = _simulation(values) + del simulation._fast_cache + _assert_settled(simulation, calls, _input_first_result(values)) + + +def test_repeated_result_with_alternating_inputs_is_not_a_fixed_point(): + # Both caps produce zero, but the mutations differ on every attempt. + simulation, calls = _simulation([0.0], alternating=True) + budget = simulation_module._RERUNS_AFTER_INPUT_CHANGE + 1 + for calculation in (1, 2): + assert simulation.calculate("fp_outer", "2020").tolist() == [0.0] + assert len(calls) == calculation * budget + assert simulation.get_array("fp_outer", "2020") is None + + +def test_masked_inputs_cannot_certify_a_fixed_point_from_backing_bytes(): + simulation, calls = _simulation([0.0], masked_inputs=True) + assert simulation.calculate("fp_outer", "2020").tolist() == [0.0] + assert len(calls) == simulation_module._RERUNS_AFTER_INPUT_CHANGE + 1 + assert simulation.get_array("fp_outer", "2020") is None + + +@pytest.mark.parametrize("thread_operation", ["read", "change"]) +def test_untracked_thread_operation_cannot_certify_a_fixed_point(thread_operation): + simulation, calls = _simulation([40.0], thread_operation=thread_operation) + np.testing.assert_array_equal( + simulation.calculate("fp_outer", "2020"), _input_first_result([40.0]) + ) + assert len(calls) == simulation_module._RERUNS_AFTER_INPUT_CHANGE + 1 + assert simulation.get_array("fp_outer", "2020") is None + + +def test_untracked_thread_write_before_first_read_cannot_certify_a_fixed_point(): + # The write's epoch precedes every read, so checking intervening drops + # alone cannot establish that the formula observed every input mutation. + simulation, calls = _simulation([40.0], thread_operation="before_read") + np.testing.assert_array_equal( + simulation.calculate("fp_outer", "2020"), _input_first_result([41.0]) + ) + assert len(calls) == simulation_module._RERUNS_AFTER_INPUT_CHANGE + 1 + assert simulation.get_array("fp_outer", "2020") is None + + +def test_distinct_clones_with_the_same_branch_path_do_not_prove_a_fixed_point(): + # A direct branch clone retains its parent's branch name and ancestry. + # One indistinguishable set(path, x, 0) per run changes a different clone: + # the first two results equal 1, but the third becomes 0. + clones = [] + calls = [] + + def outer(person, period): + calls.append(str(period)) + previous = [clone.calculate("fp_clone_x", period) for clone in clones] + largest = int(np.argmax([value[0] for value in previous])) + clones[largest].set_input("fp_clone_x", period, np.zeros(person.count)) + return clones[0].calculate("fp_clone_x", period) + + system = CountryTaxBenefitSystem() + system.add_variables( + _variable("fp_clone_x", periods.YEAR), + _variable("fp_clone_outer", periods.YEAR, outer), + ) + simulation = SimulationBuilder().build_default_simulation(system) + simulation.set_input("fp_clone_x", "2020", np.array([0.0])) + original = simulation.get_branch("same_path") + clones.extend([original, original.clone(), original.clone()]) + for clone, value in zip(clones, (1.0, 2.0, 3.0)): + clone.set_input("fp_clone_x", "2020", np.array([value])) + + assert simulation.calculate("fp_clone_outer", "2020").tolist() == [0.0] + assert len(calls) == 4 + assert simulation.get_array("fp_clone_outer", "2020").tolist() == [0.0] + assert simulation.calculate("fp_clone_outer", "2020").tolist() == [0.0] + assert len(calls) == 4 + + +def test_rotating_same_path_clones_do_not_prove_a_fixed_point_between_attempts(): + # Only one clone is read and changed per attempt. The next clone comes + # from an input, so the first two identical set(path, x, 0) transitions + # and results must not hide the third clone's different state. + clones = [] + calls = [] + + def outer(person, period): + calls.append(str(period)) + simulation = person.simulation + branch = simulation.get_branch("same_path") + previous = branch.calculate("fp_rotate_x", period) + next_clone = int(branch.calculate("fp_rotate_next", period)[0]) + branch.set_input("fp_rotate_x", period, np.zeros(person.count)) + simulation.branches["same_path"] = clones[next_clone] + return np.minimum(previous, 1) + + system = CountryTaxBenefitSystem() + system.add_variables( + _variable("fp_rotate_x", periods.YEAR), + _variable("fp_rotate_next", periods.YEAR), + _variable("fp_rotate_outer", periods.YEAR, outer), + ) + simulation = SimulationBuilder().build_default_simulation(system) + simulation.set_input("fp_rotate_x", "2020", np.array([0.0])) + original = simulation.get_branch("same_path") + clones.extend([original, original.clone(), original.clone()]) + for next_clone, (clone, value) in enumerate(zip(clones, (1.0, 2.0, 0.0)), 1): + clone.set_input("fp_rotate_x", "2020", np.array([value])) + clone.set_input("fp_rotate_next", "2020", np.array([next_clone % 3])) + + assert simulation.calculate("fp_rotate_outer", "2020").tolist() == [0.0] + assert len(calls) == 3 + assert simulation.get_array("fp_rotate_outer", "2020").tolist() == [0.0] + assert simulation.calculate("fp_rotate_outer", "2020").tolist() == [0.0] + assert len(calls) == 3 + + +def test_own_period_input_wins_over_a_branch_reset_formula_result(): + simulation, calls = _simulation([40.0], own_input=True) + assert simulation.calculate("fp_outer", "2020").tolist() == [9.0] + holder = simulation.get_holder("fp_outer") + assert holder._is_input(periods.period("2020"), "default") + assert holder.get_array("2020").tolist() == [9.0] + completed = len(calls) + assert simulation.calculate("fp_outer", "2020").tolist() == [9.0] + assert len(calls) == completed + + +def test_calculate_add_keeps_a_sum_after_persistent_branch_resets_settle(): + values = [0.0, 40.0] + simulation, calls = _simulation(values, definition_period=periods.MONTH) + expected = 12 * _input_first_result(values) + np.testing.assert_array_equal( + simulation.calculate_add("fp_outer", "2020"), expected + ) + assert len(calls) == 24 + np.testing.assert_array_equal(simulation.get_array("fp_outer", "2020"), expected) + completed = len(calls) + np.testing.assert_array_equal(simulation.calculate("fp_outer", "2020"), expected) + assert len(calls) == completed + + +@pytest.mark.parametrize("saved_recreations", [False, True]) +def test_rotating_read_only_clones_do_not_prove_a_fixed_point_between_attempts( + saved_recreations, +): + # The mutations and first two results match, but they read different + # direct clones. The third reader's zero must not be hidden by caching. + simulation, _ = _simulation([0.0]) + clones = [] + calls = [] + + def outer(person, period): + calls.append(str(period)) + simulation = person.simulation + branch = simulation.get_branch("same_path") + previous = branch.calculate("fp_rotate_x", period) + next_clone = int(branch.calculate("fp_rotate_next", period)[0]) + simulation.branches["same_path"] = clones[next_clone] + person("fp_cap", period) + persistent = simulation.get_branch("persistent") + persistent.delete_arrays("fp_cap") + maximum = persistent.calculate("fp_base", period) + persistent.set_input("fp_cap", period, maximum) + return np.minimum(previous, 1) + + system = simulation.tax_benefit_system + system.add_variables( + _variable("fp_rotate_x", periods.YEAR), + _variable("fp_rotate_next", periods.YEAR), + _variable("fp_rotate_outer", periods.YEAR, outer), + ) + simulation.set_input("fp_rotate_x", "2020", np.array([0.0])) + original = simulation.get_branch("same_path") + clones.append(original) + for _ in range(2): + if saved_recreations: + del simulation.branches["same_path"] + clones.append(simulation.get_branch("same_path")) + else: + clones.append(original.clone()) + for next_clone, (clone, value) in enumerate(zip(clones, (1.0, 2.0, 0.0)), 1): + clone.set_input("fp_rotate_x", "2020", np.array([value])) + clone.set_input("fp_rotate_next", "2020", np.array([min(next_clone, 2)])) + simulation.branches["same_path"] = original + + assert simulation.calculate("fp_rotate_outer", "2020").tolist() == [0.0] + assert 3 <= len(calls) <= 4 + assert simulation.get_array("fp_rotate_outer", "2020").tolist() == [0.0] + completed = len(calls) + assert simulation.calculate("fp_rotate_outer", "2020").tolist() == [0.0] + assert len(calls) == completed diff --git a/tests/core/test_branch_input_fixed_point_property.py b/tests/core/test_branch_input_fixed_point_property.py new file mode 100644 index 000000000..2d5358132 --- /dev/null +++ b/tests/core/test_branch_input_fixed_point_property.py @@ -0,0 +1,72 @@ +"""Differential coverage for repeated branch resets, including recreated branches.""" + +import pytest + +pytest.importorskip("hypothesis") +from hypothesis import given, settings, strategies as st + +from tests.core.test_branch_input_fixed_point import ( + _assert_settled, + _input_first_result, + _simulation, +) + + +@given( + values=st.lists( + st.floats(min_value=-500, max_value=500, allow_nan=False, width=32), + min_size=1, + max_size=5, + ), + delete_derived=st.booleans(), + recreate_branch=st.booleans(), + cache_mode=st.sampled_from(["memory", "disk", "drop", "blacklist"]), +) +@settings(max_examples=40, deadline=None) +def test_branch_reset_fixed_point_matches_inputs_first( + values, delete_derived, recreate_branch, cache_mode +): + simulation, calls = _simulation( + values, + delete_derived=delete_derived, + recreate_branch=recreate_branch, + cache_mode=cache_mode, + ) + _assert_settled( + simulation, + calls, + _input_first_result(values), + kept_in_holder=cache_mode in ("memory", "disk"), + ) + + +@given( + values=st.lists( + st.floats(min_value=-500, max_value=500, allow_nan=False, width=32), + min_size=1, + max_size=5, + ), + branch_depth=st.integers(min_value=1, max_value=3), + cache_mode=st.sampled_from(["memory", "disk", "drop", "blacklist"]), +) +@settings(max_examples=40, deadline=None) +def test_reused_and_recreated_branch_paths_have_equal_results_and_work( + values, branch_depth, cache_mode +): + # Persisting a branch path and recreating it must both match inputs-first + # values, settle in exactly two attempts, and reuse the settled result + # while respecting holder storage restrictions. + expected = _input_first_result(values) + for recreate_branch in (False, True): + simulation, calls = _simulation( + values, + branch_depth=branch_depth, + recreate_branch=recreate_branch, + cache_mode=cache_mode, + ) + _assert_settled( + simulation, + calls, + expected, + kept_in_holder=cache_mode in ("memory", "disk"), + ) diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py new file mode 100644 index 000000000..ec86d82ba --- /dev/null +++ b/tests/core/test_branch_input_invalidation.py @@ -0,0 +1,2599 @@ +"""``set_input`` on a branch drops what may depend on the value it replaces. + +A branch starts with every value its parent holds. ``set_input`` on a branch +used to store the new value and keep everything calculated from the old one, +so a value the parent calculated before branching answered for the branch, +whatever the branch's inputs said, and the branch's results depended on what +had been calculated before it was created. + +Every stored value now carries a sequence number from one process-wide +counter, and a value is stored after everything it was calculated from. Each +simulation keeps a history of the first stores its values may have been +calculated from: a branch copies its parent's, and a formula that calculates +in another simulation takes in that simulation's history. An input set on a +branch drops the branch's non-input values stored at or after the earliest +recorded store, or derivation, of the overridden variable for an overlapping +period (``StoreHistory.earliest_dependency``), and the records the drop made +obsolete. + +The invariant tested throughout (and, over random sequences of calculations, +branches and inputs, by ``test_branch_input_invalidation_property.py``): +whatever the branch's family calculated before, and in whatever order, a +calculation in a branch equals the same calculation in a new simulation given +the branch's inputs, its own and those it inherited when it was created, +before calculating anything. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.data_storage import InMemoryStorage +from policyengine_core.data_storage.store_history import ( + StoreHistory, + next_sequence_number, +) +from policyengine_core.experimental import MemoryConfig +from policyengine_core.model_api import Reform +from policyengine_core.simulations import SimulationBuilder +import policyengine_core.simulations.simulation as simulation_module +from tests.fixtures.branch_input_invalidation import ( + CARRY_OVER_SYSTEM, + FORMULA_RUNS, + ROOT_INPUTS, + branching_simulation, + synthetic_simulation, +) + +JANUARY = periods.period("2017-01") +JANUARY_2022 = periods.period("2022-01") + +SITUATION = { + "persons": { + "a": {"birth": {"ETERNITY": "1980-01-01"}, "salary": {"2017-01": 4000}}, + "b": {"birth": {"ETERNITY": "1985-01-01"}, "salary": {"2017-01": 1000}}, + }, + "households": { + "h": { + "parents": ["a", "b"], + "accommodation_size": {"2017-01": 80}, + "housing_occupancy_status": {"2017-01": "tenant"}, + "rent": {"2017-01": 900}, + } + }, +} + + +def _build(tax_benefit_system, situation=SITUATION): + return SimulationBuilder().build_from_entities(tax_benefit_system, situation) + + +def _with_salary(tax_benefit_system, salary, period=JANUARY): + """A new simulation given ``salary`` before calculating anything.""" + simulation = _build(tax_benefit_system) + simulation.set_input("salary", period, np.asarray(salary, dtype=float)) + return simulation + + +def _stored_keys(simulation, variable): + return set(simulation.get_holder(variable)._memory_storage._arrays) + + +# ----- Examples on the country template ----- # + + +def test_branch_input_reaches_a_value_the_parent_calculated_first( + tax_benefit_system, +): + simulation = _build(tax_benefit_system) + simulation.calculate("income_tax", JANUARY) + simulation.calculate("disposable_income", JANUARY) + + branch = simulation.get_branch("raise") + branch.set_input("salary", JANUARY, np.array([5000.0, 1000.0])) + + fresh = _with_salary(tax_benefit_system, [5000.0, 1000.0]) + for variable in ("income_tax", "disposable_income", "household_income"): + assert np.array_equal( + branch.calculate(variable, JANUARY), fresh.calculate(variable, JANUARY) + ), variable + # The parent keeps its own values. + assert np.array_equal( + simulation.calculate("income_tax", JANUARY), + _build(tax_benefit_system).calculate("income_tax", JANUARY), + ) + + +def test_branch_keeps_values_stored_before_the_input_first_existed( + tax_benefit_system, +): + """Values stored before the overridden variable was first stored stay.""" + situation = { + "persons": {"a": {"birth": {"ETERNITY": "1980-01-01"}}}, + "households": { + "h": { + "parents": ["a"], + "accommodation_size": {"2017-01": 80}, + "housing_occupancy_status": {"2017-01": "tenant"}, + } + }, + } + simulation = _build(tax_benefit_system, situation) + # housing_tax is stored before salary has any value. + simulation.calculate("housing_tax", "2017") + simulation.set_input("salary", JANUARY, np.array([3000.0])) + simulation.calculate("income_tax", JANUARY) + + branch = simulation.get_branch("raise") + branch.set_input("salary", JANUARY, np.array([4000.0])) + + assert branch.calculate("income_tax", JANUARY) == pytest.approx( + 4000 * tax_benefit_system.parameters(JANUARY).taxes.income_tax_rate + ) + assert _stored_keys(branch, "housing_tax") == {"default:2017"} + + +def test_branch_input_with_no_history_drops_nothing(tax_benefit_system): + """The usual case costs nothing: no value can depend on a value never stored.""" + simulation = _build(tax_benefit_system) + simulation.calculate("disposable_income", JANUARY) + branch = simulation.get_branch("later") + before = { + variable: _stored_keys(branch, variable) + for variable in ("disposable_income", "income_tax", "salary") + } + + # salary has never been stored, read or derived for 2018-03. + dropped = branch._drop_values_that_may_depend_on( + "salary", periods.period("2018-03") + ) + branch.set_input("salary", "2018-03", np.array([1.0, 2.0])) + + assert dropped == 0 + for variable, keys in before.items(): + assert _stored_keys(branch, variable) >= keys, variable + + +def test_branch_keeps_inputs_of_the_dataset_and_of_ancestor_branches( + tax_benefit_system, +): + simulation = _build(tax_benefit_system) + simulation.calculate("disposable_income", JANUARY) + parent_branch = simulation.get_branch("parent") + parent_branch.set_input("rent", JANUARY, np.array([500.0])) + parent_branch.calculate("housing_allowance", JANUARY) + + child = parent_branch.get_branch("child") + child.set_input("salary", JANUARY, np.array([0.0, 0.0])) + + assert "default:2017-01" in _stored_keys(child, "accommodation_size") + assert "parent:2017-01" in _stored_keys(child, "rent") + assert child.calculate("rent", JANUARY) == 500.0 + fresh = _with_salary(tax_benefit_system, [0.0, 0.0]) + fresh.set_input("rent", JANUARY, np.array([500.0])) + for variable in ("housing_allowance", "household_income", "income_tax"): + assert np.array_equal( + child.calculate(variable, JANUARY), fresh.calculate(variable, JANUARY) + ), variable + + +def test_input_for_a_year_drops_values_calculated_from_a_month( + tax_benefit_system, +): + simulation = _build(tax_benefit_system) + simulation.calculate("income_tax", JANUARY) + + branch = simulation.get_branch("annual") + # salary divides a yearly input over the months the branch has not set. + branch.set_input("salary", "2017", np.array([12000.0, 24000.0])) + + rate = tax_benefit_system.parameters(JANUARY).taxes.income_tax_rate + assert np.allclose( + branch.calculate("income_tax", JANUARY), [1000 * rate, 2000 * rate] + ) + + +def test_input_set_again_on_a_branch_drops_what_the_first_input_fed( + tax_benefit_system, +): + simulation = _build(tax_benefit_system) + branch = simulation.get_branch("raise") + branch.set_input("salary", JANUARY, np.array([5000.0, 1000.0])) + branch.calculate("disposable_income", JANUARY) + branch.set_input("salary", JANUARY, np.array([7000.0, 1000.0])) + + fresh = _with_salary(tax_benefit_system, [7000.0, 1000.0]) + assert np.array_equal( + branch.calculate("disposable_income", JANUARY), + fresh.calculate("disposable_income", JANUARY), + ) + + +def test_set_input_on_a_simulation_that_is_not_a_branch_drops_nothing( + tax_benefit_system, +): + """Unchanged: a root simulation's later inputs do not reach its cache.""" + simulation = _build(tax_benefit_system) + cached = simulation.calculate("income_tax", JANUARY).copy() + simulation.set_input("salary", JANUARY, np.array([9000.0, 9000.0])) + + assert np.array_equal(simulation.calculate("income_tax", JANUARY), cached) + + +def test_existing_child_branches_keep_their_values(tax_benefit_system): + simulation = _build(tax_benefit_system) + branch = simulation.get_branch("branch") + branch.calculate("income_tax", JANUARY) + child = branch.get_branch("child") + child_value = child.calculate("income_tax", JANUARY).copy() + + branch.set_input("salary", JANUARY, np.array([0.0, 0.0])) + + assert np.array_equal(child.calculate("income_tax", JANUARY), child_value) + assert np.array_equal(branch.calculate("income_tax", JANUARY), [0.0, 0.0]) + + +def test_disk_storage_drops_calculated_values_and_keeps_inputs(tax_benefit_system): + situation = { + "persons": {"a": {"birth": {"ETERNITY": "1980-01-01"}}}, + "households": {"h": {"parents": ["a"]}}, + } + simulation = _build(tax_benefit_system, situation) + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + for variable in ("rent", "housing_allowance"): + holder = simulation.get_holder(variable) + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + simulation.set_input("rent", JANUARY, np.array([700.0])) + simulation.calculate("housing_allowance", JANUARY) + assert simulation.get_holder("housing_allowance")._disk_storage._files + + branch = simulation.get_branch("cheaper") + branch.set_input("rent", JANUARY, np.array([300.0])) + + dropped = branch.get_holder("housing_allowance")._disk_storage._files == {} + fresh = _build(tax_benefit_system, situation) + fresh.set_input("rent", JANUARY, np.array([300.0])) + assert np.array_equal( + branch.calculate("housing_allowance", JANUARY), + fresh.calculate("housing_allowance", JANUARY), + ) + assert dropped + rent_disk = branch.get_holder("rent")._disk_storage + assert set(rent_disk._files) == {"default_2017-01", "cheaper_2017-01"} + assert set(rent_disk._files) - rent_disk._derived == { + "default_2017-01", + "cheaper_2017-01", + } + + +def test_value_served_from_the_macro_cache_is_tracked(tax_benefit_system, monkeypatch): + """A macro-cache read is not stored, but what reads it is.""" + import policyengine_core.simulations.simulation as simulation_module + + class CachedIncomeTax: + def __init__(self, tax_benefit_system): + pass + + def set_cache_path(self, *args): + pass + + def get_cache_path(self): + return type("Path", (), {"exists": lambda self: True})() + + def get_cache_value(self, path): + return np.array([100.0, 100.0], dtype=np.float32) + + monkeypatch.setattr(simulation_module, "SimulationMacroCache", CachedIncomeTax) + simulation = _build(tax_benefit_system) + simulation.macro_cache_read = True + simulation.dataset = type( + "Dataset", (), {"file_path": simulation_module.Path("cache"), "name": "d"} + )() + monkeypatch.setattr( + type(simulation), + "check_macro_cache", + lambda self, variable, period: variable == "income_tax", + ) + simulation.calculate("disposable_income", JANUARY) + assert not _stored_keys(simulation, "income_tax") + + branch = simulation.get_branch("branch") + branch.set_input("income_tax", JANUARY, np.array([0.0, 0.0])) + + assert not _stored_keys(branch, "disposable_income") + fresh = _build(tax_benefit_system) + fresh.set_input("income_tax", JANUARY, np.array([0.0, 0.0])) + assert np.array_equal( + branch.calculate("disposable_income", JANUARY), + fresh.calculate("disposable_income", JANUARY), + ) + + +# ----- drop_computed_arrays ----- # + + +def test_drop_computed_arrays_keeps_only_inputs(tax_benefit_system): + simulation = _build(tax_benefit_system) + simulation.calculate("total_taxes", JANUARY) + branch = simulation.get_branch("branch") + branch.set_input("rent", JANUARY, np.array([100.0])) + branch.calculate("total_benefits", JANUARY) + + dropped = branch.drop_computed_arrays() + + assert dropped > 0 + remaining = { + (name, key) + for population in branch.populations.values() + for name, holder in population._holders.items() + for key in holder._memory_storage._arrays + } + inputs = { + (name, key) + for population in branch.populations.values() + for name, holder in population._holders.items() + for key in holder._memory_storage._inputs + } + assert remaining == inputs + assert {("rent", "branch:2017-01"), ("salary", "default:2017-01")} <= remaining + assert not _stored_keys(branch, "total_taxes") + assert not _stored_keys(branch, "total_benefits") + assert np.array_equal(branch.calculate("rent", JANUARY), [100.0]) + # The parent is untouched. + assert _stored_keys(simulation, "total_taxes") == {"default:2017-01"} + + +def test_drop_computed_arrays_after_changing_a_branch_parameter(tax_benefit_system): + """Parameters are not inputs: a branch with other parameters drops its values.""" + simulation = _build(tax_benefit_system) + simulation.calculate("income_tax", JANUARY) + + branch = simulation.get_branch("higher_rate", clone_system=True) + branch.tax_benefit_system.parameters.taxes.income_tax_rate.update( + period="2017", value=0.5 + ) + stale = branch.calculate("income_tax", JANUARY).copy() + branch.drop_computed_arrays() + + assert np.array_equal(branch.calculate("income_tax", JANUARY), [2000.0, 500.0]) + assert not np.array_equal(stale, [2000.0, 500.0]) + + +def test_apply_reform_keeps_branch_inputs(tax_benefit_system): + simulation = _build(tax_benefit_system) + branch = simulation.get_branch("branch") + branch.set_input("salary", JANUARY, np.array([10.0, 20.0])) + branch.calculate("income_tax", JANUARY) + + class NoOp(Reform): + def apply(self): + pass + + simulation.apply_reform(NoOp) + + assert _stored_keys(branch, "income_tax") == set() + assert np.array_equal(branch.calculate("salary", JANUARY), [10.0, 20.0]) + + +def test_subsample_starts_a_new_store_history(): + """``subsample`` replaces every value, so earlier stores no longer count.""" + from policyengine_core.country_template import Microsimulation + + simulation = Microsimulation() + simulation.calculate("income_tax", JANUARY_2022) + last_before = next_sequence_number() + simulation.subsample(n=3, seed="store-history", time_period="2022") + + recorded = [ + number + for stored in simulation._store_history._first_stored.values() + for number in stored.values() + ] + assert recorded and min(recorded) > last_before + + # A branch's input still reaches what the parent calculated first. + salary = np.asarray(simulation.calculate("salary", JANUARY_2022)) + 1000 + before = simulation.get_branch("before") + before.set_input("salary", JANUARY_2022, salary) + expected = np.asarray( + before.calculate("social_security_contribution", JANUARY_2022) + ) + simulation.calculate("social_security_contribution", JANUARY_2022) + after = simulation.get_branch("after") + after.set_input("salary", JANUARY_2022, salary) + assert np.array_equal( + np.asarray(after.calculate("social_security_contribution", JANUARY_2022)), + expected, + ) + + +# ----- Storage and history ----- # + + +def test_storage_records_sequence_numbers_and_inputs_through_clone(): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0]), periods.period("2020")) + first = storage._sequence_numbers["default:2020"] + storage.put(np.array([2.0]), periods.period("2021"), derived=True) + clone = storage.clone() + + assert clone._sequence_numbers == storage._sequence_numbers + assert clone._sequence_numbers["default:2021"] > first + assert set(clone._inputs) == {"default:2020"} + + assert clone.drop_computed(since=first) == 1 + assert set(clone._arrays) == {"default:2020"} + assert set(storage._arrays) == {"default:2020", "default:2021"} + + +def test_storage_drops_only_values_at_or_after_since(): + storage = InMemoryStorage(is_eternal=False) + numbers = [] + for year in (2020, 2021, 2022): + number = next_sequence_number() + storage.put( + np.array([0.0]), periods.period(year), derived=True, sequence_number=number + ) + numbers.append(number) + + assert storage.drop_computed(since=numbers[1]) == 2 + assert set(storage._arrays) == {"default:2020"} + storage.delete(periods.period(2020)) + assert storage._sequence_numbers == {} + + +def test_history_matches_overlapping_periods_and_derived_values(): + history = StoreHistory() + history.record_store("v", periods.period("2020-03"), 10) + history.record_store("v", periods.period("2020-03"), 99) + history.record_store("v", periods.period("2021"), 20) + + assert history.earliest_dependency("v", periods.period("2020")) == 10 + assert history.earliest_dependency("v", periods.period("2020-04")) is None + assert history.earliest_dependency("v", periods.period("2021-06")) == 20 + assert history.earliest_dependency("w", periods.period("2020")) is None + + history.record_derived("v", 15) + assert history.earliest_dependency("v", periods.period("2020-04")) == 15 + history.record_store("e", periods.period(periods.ETERNITY), 5) + assert history.earliest_dependency("e", periods.period("1999-01")) == 5 + + +# ----- Derived values and branches inside formulas, on a synthetic system ----- # + + +def test_value_carried_over_from_an_earlier_input_is_dropped(): + simulation = synthetic_simulation(ROOT_INPUTS, system=CARRY_OVER_SYSTEM) + simulation.calculate("p_cond", "2015") # p_c 2015 carried over from 2012 + + branch = simulation.get_branch("branch") + branch.set_input("p_c", "2014", np.array([70.0, 80.0, 90.0])) + + fresh = synthetic_simulation( + {**ROOT_INPUTS, ("p_c", "2014"): (70.0, 80.0, 90.0)}, + system=CARRY_OVER_SYSTEM, + ) + assert np.array_equal(fresh.calculate("p_c", "2015"), [70.0, 80.0, 90.0]) + for variable in ("p_c", "p_cond"): + assert np.array_equal( + branch.calculate(variable, "2015"), fresh.calculate(variable, "2015") + ), variable + + +def test_uprated_value_is_dropped(): + simulation = synthetic_simulation(ROOT_INPUTS) + simulation.calculate("p_prod", "2015") # p_up 2015 uprated from 2012 + + branch = simulation.get_branch("branch") + branch.set_input("p_up", "2014", np.array([1.0, 2.0, 3.0])) + + fresh = synthetic_simulation({**ROOT_INPUTS, ("p_up", "2014"): (1.0, 2.0, 3.0)}) + assert np.allclose( + branch.calculate("p_prod", "2015"), fresh.calculate("p_prod", "2015") + ) + + +def test_value_calculated_only_inside_another_branch_is_tracked(): + """A formula's own branch stores the overridden variable; its result is dropped. + + ``p_inner`` calculates ``p_inner_only`` in a branch it deletes when done, + so the parent never stores ``p_inner_only`` itself. It takes in that + branch's history when its calculation returns, and the parent's later + branches copy it. + """ + simulation = synthetic_simulation(ROOT_INPUTS) + simulation.calculate("p_inner", "2013") + assert not simulation.get_holder("p_inner_only")._memory_storage._arrays + + branch = simulation.get_branch("branch") + branch.set_input("p_inner_only", "2013", np.array([1.0, 1.0, 1.0])) + + fresh = synthetic_simulation( + {**ROOT_INPUTS, ("p_inner_only", "2013"): (1.0, 1.0, 1.0)} + ) + assert np.array_equal( + branch.calculate("p_inner", "2013"), fresh.calculate("p_inner", "2013") + ) + + +def test_value_a_holder_does_not_keep_is_tracked(): + """``variables_to_drop`` values are not stored, but what reads them is.""" + config = MemoryConfig(max_memory_occupation=1, variables_to_drop=["p_sum"]) + simulation = synthetic_simulation(ROOT_INPUTS, memory_config=config) + simulation.calculate("p_prod", "2013") + assert not simulation.get_holder("p_sum")._memory_storage._arrays + + branch = simulation.get_branch("branch") + branch.set_input("p_sum", "2013", np.array([0.0, 0.0, 0.0])) + + assert np.array_equal(branch.calculate("p_prod", "2013"), [0.0, 0.0, 0.0]) + + +def test_value_a_spiral_defaults_is_tracked(): + """The default a spiral returns is not stored, but what reads it is.""" + simulation = synthetic_simulation(ROOT_INPUTS) + branch = simulation.get_branch("branch") + branch.calculate("p_spiral", "2015") + deepest = min( + periods.period(key.split(":", 1)[1]) + for key in branch.get_holder("p_spiral")._memory_storage._arrays + ).last_year + + branch.set_input("p_spiral", deepest, np.array([100.0, 100.0, 100.0])) + + fresh = synthetic_simulation( + {**ROOT_INPUTS, ("p_spiral", str(deepest)): (100.0,) * 3} + ) + assert np.array_equal( + branch.calculate("p_spiral", "2015"), fresh.calculate("p_spiral", "2015") + ) + + +# ----- Values that reach a simulation without being stored there ----- # + + +def test_macro_cache_files_a_branch_wrote_are_not_read_after_its_input_changes( + tmp_path, +): + """Macro-cache files are keyed by branch and period, not by inputs.""" + from policyengine_core.data.dataset import Dataset + from policyengine_core.entities import build_entity + from policyengine_core.simulations import Simulation + from policyengine_core.taxbenefitsystems import TaxBenefitSystem + from policyengine_core.variables import Variable + + person = build_entity("person", "persons", "Person", is_person=True) + + class source(Variable): + label = "Source" + value_type = float + entity = person + definition_period = periods.YEAR + + class result(Variable): + label = "Result" + value_type = float + entity = person + definition_period = periods.YEAR + exhaustive_parameter_dependencies = [] + + def formula(person, period): + return 2 * person("source", period) + + class OnePerson(Dataset): + name = "one_person" + label = "One person" + file_path = tmp_path / "one_person.h5" + data_format = Dataset.ARRAYS + time_period = "2022" + + def generate(self): + self.save_dataset({"person_id": np.array([0]), "source": np.array([1.0])}) + + system = TaxBenefitSystem([person]) + system.add_variables(source, result) + parent = Simulation(tax_benefit_system=system, dataset=OnePerson) + parent.macro_cache_read = True + branch = parent.get_branch("test") + assert branch.calculate("result", "2022").tolist() == [2.0] # writes the file + + branch.set_input("source", "2022", np.array([5.0])) + + assert branch.calculate("result", "2022").tolist() == [10.0] + + +def test_blacklisted_value_is_tracked(): + """With ``opt_out_cache``, blacklisted values are not stored; what reads them is.""" + simulation = synthetic_simulation(ROOT_INPUTS, opt_out_cache=True) + simulation.calculate("p_prod", "2013") + assert not simulation.get_holder("p_sum")._memory_storage._arrays + + branch = simulation.get_branch("branch") + branch.set_input("p_sum", "2013", np.array([0.0, 0.0, 0.0])) + + assert np.array_equal(branch.calculate("p_prod", "2013"), [0.0, 0.0, 0.0]) + + +def test_value_written_straight_into_storage_counts_as_a_dependency( + tax_benefit_system, +): + """A value stored without ``put`` has no number, so anything may depend on it.""" + simulation = _build(tax_benefit_system) + holder = simulation.get_holder("salary") + holder._memory_storage._arrays["default:2017-02"] = np.array([4000.0, 0.0]) + simulation.calculate("income_tax", "2017-02") + + branch = simulation.get_branch("branch") + branch.set_input("salary", "2017-02", np.array([0.0, 0.0])) + + assert np.array_equal(branch.calculate("income_tax", "2017-02"), [0.0, 0.0]) + + +def test_cached_value_another_simulation_returns_is_tracked(): + """A formula that reads another simulation's cached value takes in its history.""" + simulation = branching_simulation() + # Cached in the persistent branch before any formula of the parent reads it. + simulation.get_branch("persistent").calculate("tax", "2020") + simulation.calculate("from_persistent_branch", "2020") # a cache hit there + + branch = simulation.get_branch("branch") + branch.set_input("tax", "2020", np.zeros(3)) + + # A new simulation given the input first creates its persistent branch + # from itself, so the branch's input reaches it. + assert np.array_equal( + branch.calculate("from_persistent_branch", "2020"), np.zeros(3) + ) + + +def test_history_is_taken_in_again_after_a_drop_forgets_it(): + """A drop forgets records; reading the other simulation again restores them.""" + simulation = branching_simulation() + branch = simulation.get_branch("branch") + branch.calculate("from_persistent_branch", "2020") + branch.drop_computed_arrays() # forgets what the persistent branch stored + branch.calculate("from_persistent_branch", "2020") # a cache hit there + + child = branch.get_branch("child") + child.set_input("tax", "2020", np.zeros(3)) + + assert np.array_equal( + child.calculate("from_persistent_branch", "2020"), np.zeros(3) + ) + + +def _one_person_system(*variables): + from policyengine_core.country_template import CountryTaxBenefitSystem + + system = CountryTaxBenefitSystem() + system.add_variables(*variables) + return system + + +def _yearly_variable(name, formula=None): + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + namespace = dict( + value_type=float, + entity=entities.Person, + definition_period=periods.YEAR, + label=name, + ) + if formula is not None: + namespace["formula"] = formula + return type(name, (Variable,), namespace) + + +@pytest.mark.parametrize("drop", ["drop_computed_arrays", "set_input"]) +def test_drop_while_a_formula_runs_keeps_what_it_read_recorded(drop): + """A formula still holding a value it read keeps that value's record.""" + + def result(person, period): + source = person("source", period) + if drop == "drop_computed_arrays": + person.simulation.drop_computed_arrays() + else: + person.simulation.set_input("trigger", period, np.ones(person.count)) + return source * 2 + + system = _one_person_system( + _yearly_variable("source", lambda person, period: np.ones(person.count)), + _yearly_variable("trigger", lambda person, period: np.zeros(person.count)), + _yearly_variable("result", result), + ) + simulation = SimulationBuilder().build_default_simulation(system) + simulation.calculate("trigger", "2020") # stored before ``source`` + branch = simulation.get_branch("branch") + assert branch.calculate("result", "2020").tolist() == [2.0] + + branch.set_input("source", "2020", np.array([3.0])) + + assert branch.calculate("result", "2020").tolist() == [6.0] + + +def test_disk_restore_reads_the_latest_value_of_each_key(tmp_path, monkeypatch): + """Sequence numbers restart in each process; restore does not go by them.""" + import itertools + + import policyengine_core.data_storage.store_history as store_history + from policyengine_core.data_storage import OnDiskStorage + + writer = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + for _ in range(5): + writer.put(np.array([1.0]), periods.period("2020")) + # A later process: its counter restarts. + monkeypatch.setattr(store_history, "_sequence", itertools.count(1)) + later = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + later.put(np.array([2.0]), periods.period("2020")) + + reader = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + reader.restore() + + assert reader.get(periods.period("2020")).tolist() == [2.0] + + +def test_uprated_value_another_simulation_returns_is_tracked(): + """Taking in another simulation's history includes its uprated values.""" + inputs = {("p_up", "2012"): (100.0, 100.0, 100.0)} + simulation = synthetic_simulation(inputs) + simulation.calculate("p_imported_up", "2015") # uprated in a temporary branch + + branch = simulation.get_branch("branch") + branch.set_input("p_up", "2014", np.full(3, 300.0)) + + fresh = synthetic_simulation({**inputs, ("p_up", "2014"): (300.0,) * 3}) + assert np.allclose( + branch.calculate("p_imported_up", "2015"), + fresh.calculate("p_imported_up", "2015"), + ) + + +def test_value_a_formula_calculates_in_a_worker_thread_is_tracked(): + """A thread started without the formula's context still hands its history up.""" + from concurrent.futures import ThreadPoolExecutor + + def result(person, period): + simulation = person.simulation + child = simulation.get_branch("worker") + try: + with ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(child.calculate, "value", period).result() * 2 + finally: + del simulation.branches["worker"] + + system = _one_person_system( + _yearly_variable("value", lambda person, period: np.zeros(person.count)), + _yearly_variable("result", result), + ) + simulation = SimulationBuilder().build_default_simulation(system) + simulation.calculate("result", "2020") + + branch = simulation.get_branch("branch") + branch.set_input("value", "2020", np.array([10.0])) + + assert branch.calculate("result", "2020").tolist() == [20.0] + + +def test_failed_prerequisite_request_does_not_satisfy_the_gate(): + def prerequisite(person, period): + raise RuntimeError("the prerequisite failed") + + dependent = _yearly_variable( + "dependent", lambda person, period: np.full(person.count, 42.0) + ) + dependent.requires_computation_after = "prerequisite" + system = _one_person_system( + _yearly_variable("prerequisite", prerequisite), dependent + ) + simulation = SimulationBuilder().build_default_simulation(system) + with pytest.raises(RuntimeError): + simulation.calculate("prerequisite", "2020") + + with pytest.raises(ValueError, match="requires prerequisite"): + simulation.calculate("dependent", "2020") + + +def test_result_calculated_across_its_own_input_change_is_calculated_again(): + """A formula that read an input and then replaced it runs again from the new one.""" + + def result(person, period): + source = person("source", period) + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + return source * 2 + + system = _one_person_system( + _yearly_variable("source", lambda person, period: np.ones(person.count)), + _yearly_variable("result", result), + ) + branch = SimulationBuilder().build_default_simulation(system).get_branch("b") + assert branch.calculate("result", "2020").tolist() == [6.0] + + assert branch.calculate("result", "2020").tolist() == [6.0] + + +def test_result_calculated_again_is_kept(): + """A result calculated again after the input change its formula made is cached.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + class flag(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "flag" + + class source(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "source" + + def formula_2021(person, period): + if np.any(person("flag", "2021") == 0): + person.simulation.set_input("flag", "2021", np.ones(person.count)) + return np.full(person.count, 3.0) + + def formula_2022(person, period): + return None + + system = _one_person_system(flag, source) + system.auto_carry_over_input_variables = True + root = SimulationBuilder().build_default_simulation(system) + root.set_input("flag", "2021", np.array([0.0])) + branch = root.get_branch("branch") + + assert branch.calculate("source", "2021").tolist() == [3.0] + assert _stored_keys(branch, "source") == {"branch:2021"} + # Only inputs carry over, so a later year reads the default either way. + assert branch.calculate("source", "2022").tolist() == [0.0] + + +def test_result_is_calculated_again_until_no_input_changes(): + """Each run may change another input it read; the kept result follows the last.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + def flag(name): + return type( + name, + (Variable,), + dict( + value_type=float, + entity=entities.Person, + definition_period=periods.YEAR, + label=name, + ), + ) + + class source(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "source" + + def formula_2021(person, period): + for name in ("first_flag", "second_flag"): + if np.any(person(name, "2021") == 0): + person.simulation.set_input(name, "2021", np.ones(person.count)) + break + return np.full(person.count, 3.0) + + def formula_2022(person, period): + return None + + system = _one_person_system(flag("first_flag"), flag("second_flag"), source) + system.auto_carry_over_input_variables = True + root = SimulationBuilder().build_default_simulation(system) + root.set_input("first_flag", "2021", np.array([0.0])) + root.set_input("second_flag", "2021", np.array([0.0])) + branch = root.get_branch("branch") + + assert branch.calculate("source", "2021").tolist() == [3.0] + assert _stored_keys(branch, "source") == {"branch:2021"} + # Only inputs carry over, so a later year reads the default either way. + assert branch.calculate("source", "2022").tolist() == [0.0] + + +def test_direct_sum_is_taken_again_when_a_term_changes_an_earlier_input(): + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + class source(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "source" + + class result(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "result" + + def formula(person, period): + if str(period) == "2020-02": + january = person("source", "2020-01") + if np.any(january == 1): + person.simulation.set_input( + "source", "2020-01", np.full(person.count, 3.0) + ) + return person("source", period) + + system = _one_person_system(source, result) + root = SimulationBuilder().build_default_simulation(system) + for month in periods.period("2020").get_subperiods(periods.MONTH): + root.set_input("source", month, np.array([1.0])) + branch = root.get_branch("branch") + + assert branch.calculate_add("result", "2020").tolist() == [14.0] + + +def test_input_a_formula_sets_wins_over_a_carried_over_default(): + """A formula returning None after setting its own period's input.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + class source(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "source" + + def formula_2020(person, period): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + return None + + system = _one_person_system(source) + system.auto_carry_over_input_variables = True + root = SimulationBuilder().build_default_simulation(system) + root.set_input("source", "2021", np.array([9.0])) # later: no carry-over + branch = root.get_branch("branch") + + assert branch.calculate("source", "2020").tolist() == [3.0] + + +def test_input_set_under_the_default_key_during_a_branch_calculation_wins(): + """Holder.set_input stores under "default", which the branch reads.""" + + def source(person, period): + person.simulation.get_holder("source").set_input( + period, np.full(person.count, 3.0) + ) + return np.ones(person.count) + + system = _one_person_system(_yearly_variable("source", source)) + branch = SimulationBuilder().build_default_simulation(system).get_branch("b") + + assert branch.calculate("source", "2020").tolist() == [3.0] + assert branch.calculate("source", "2020").tolist() == [3.0] + + +def test_direct_divide_keeps_ignoring_an_earlier_monthly_input(): + """Only an input set during the calculation wins; one stored before does not.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + def store_whole_period(holder, period, array): + holder._set(period, array) + + class flag(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "flag" + + class source(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "source" + set_input = store_whole_period + + def formula(person, period): + if np.any(person("flag", "2020") == 0): + person.simulation.set_input("flag", "2020", np.ones(person.count)) + return np.full(person.count, 120.0) + + system = _one_person_system(flag, source) + + def branch(final_inputs): + root = SimulationBuilder().build_default_simulation(system) + root.set_input("flag", "2020", np.array([0.0])) + root.set_input("source", "2020-01", np.array([99.0])) + child = root.get_branch("b") + if final_inputs: + child.set_input("flag", "2020", np.array([1.0])) + return child + + fresh = branch(final_inputs=True).calculate_divide("source", "2020-01") + assert fresh.tolist() == [10.0] + assert branch(final_inputs=False).calculate_divide( + "source", "2020-01" + ).tolist() == [10.0] + + +def test_reruns_do_not_multiply_through_nested_calculations(): + """Only the outermost calculation runs again; the inner ones run with it.""" + from collections import Counter + + calls = Counter() + + def leaf(person, period): + calls["leaf"] += 1 + value = person("source", period) + person.simulation.set_input("source", period, value + 1) # every run + return value + + def level(name, inner): + def formula(person, period): + calls[name] += 1 + return person(inner, period) + + return formula + + system = _one_person_system( + _yearly_variable("source"), + _yearly_variable("level1", leaf), + _yearly_variable("level2", level("level2", "level1")), + _yearly_variable("level3", level("level3", "level2")), + ) + root = SimulationBuilder().build_default_simulation(system) + root.set_input("source", "2020", np.array([0.0])) + branch = root.get_branch("branch") + branch.calculate("level3", "2020") + + reruns = simulation_module._RERUNS_AFTER_INPUT_CHANGE + assert calls == {"leaf": reruns + 1, "level2": reruns + 1, "level3": reruns + 1} + assert branch._calculations_in_flight == 0 + + +def test_input_handler_calling_calculate_directly_still_drops(): + """A private _calculate (here a carry-over, no formula) also counts as calculating.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + def two_months(holder, period, array): + holder._set(period.first_month, array) + holder.simulation._calculate("source", periods.period("2020-03")) + holder._set(period.first_month.offset(1), array) + + class source(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "source" + set_input = two_months + + system = _one_person_system(source) + system.auto_carry_over_input_variables = True + + def branch(handler_input): + root = SimulationBuilder().build_default_simulation(system) + root.set_input("source", "2020-01", np.array([1.0])) + root.set_input("source", "2020-02", np.array([1.0])) + child = root.get_branch("branch") + if handler_input: + child.set_input("source", "2020", np.array([3.0])) + else: + child.set_input("source", "2020-01", np.array([3.0])) + child.set_input("source", "2020-02", np.array([3.0])) + return child + + assert branch(handler_input=False).calculate("source", "2020-03").tolist() == [3.0] + assert branch(handler_input=True).calculate("source", "2020-03").tolist() == [3.0] + + +@pytest.mark.parametrize("mode", ["calculate", "add", "divide"]) +def test_result_refused_in_one_simulation_is_not_kept_by_its_caller_in_another(mode): + """A parent formula calling into a branch whose formula changes its own input.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + def variable(name, formula=None, definition_period=periods.YEAR): + namespace = dict( + value_type=float, + entity=entities.Person, + definition_period=definition_period, + label=name, + ) + if formula is not None: + namespace["formula"] = formula + return type(name, (Variable,), namespace) + + def leaf(person, period): + previous = person("source", period.this_year) + if np.any(previous != 3.0): + person.simulation.set_input( + "source", period.this_year, np.full(person.count, 3.0) + ) + return previous * 2 + + def from_child(person, period): + worker = person.simulation.get_branch("worker") + if mode == "add": + return worker.calculate_add("leaf", period) + if mode == "divide": + return worker.calculate_divide("leaf", period.first_month) + return worker.calculate("leaf", period) + + def child_outer(person, period): + return person.simulation.parent_branch.calculate("from_child", period) + + def result_with(input_first): + def result(person, period): + simulation = person.simulation + worker = simulation.get_branch("worker") + try: + if input_first: + worker.set_input("source", period, np.full(person.count, 3.0)) + return worker.calculate("child_outer", period) + finally: + del simulation.branches["worker"] + + return result + + def run(input_first): + system = _one_person_system( + variable("source"), + variable("leaf", leaf, periods.MONTH if mode == "add" else periods.YEAR), + variable("from_child", from_child), + variable("child_outer", child_outer), + variable("result", result_with(input_first)), + ) + root = SimulationBuilder().build_default_simulation(system) + root.set_input("source", "2020", np.array([1.0])) + return ( + root.calculate("result", "2020").tolist(), + root.calculate("from_child", "2020").tolist(), + ) + + assert run(input_first=False) == run(input_first=True) + + +def test_input_stored_with_an_earlier_number_still_wins(): + """Inputs are numbered when stored, whatever number a caller passes.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + def january_only(holder, period, array): + holder._set(period.first_month, array, sequence_number=1) + + class source(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "source" + set_input = january_only + + def formula(person, period): + person.simulation.set_input("source", "2020", np.full(person.count, 42.0)) + return np.ones(person.count) + + system = _one_person_system(source) + branch = SimulationBuilder().build_default_simulation(system).get_branch("b") + + assert branch.calculate("source", "2020-01").tolist() == [42.0] + assert branch.get_holder("source")._is_input(periods.period("2020-01"), "b") + + +def test_direct_sum_runs_its_terms_again_only_as_a_whole(): + from collections import Counter + + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + calls = Counter() + + class source(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "source" + + class counted(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "counted" + + def formula(person, period): + calls["counted"] += 1 + value = person("source", period) + person.simulation.set_input("source", period, value + 1) # every run + return value + + system = _one_person_system(source, counted) + root = SimulationBuilder().build_default_simulation(system) + root.set_input("source", "2020-01", np.array([0.0])) + branch = root.get_branch("branch") + branch.calculate_add("counted", "2020-01") + + assert calls["counted"] == simulation_module._RERUNS_AFTER_INPUT_CHANGE + 1 + + +def test_input_change_in_an_unrelated_simulation_leaves_this_one_caching(): + """A settled result from another family's simulation is kept.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + def leaf(person, period): + if np.any(person("source", period) != 3): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + return person("source", period) * 2 + + other_root = SimulationBuilder().build_default_simulation( + _one_person_system(_yearly_variable("source"), _yearly_variable("leaf", leaf)) + ) + other_root.set_input("source", "2020", np.array([1.0])) + other = other_root.get_branch("worker") + + class result(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "result" + + def formula_2020(person, period): + return other.calculate("leaf", "2020") + + def formula_2021(person, period): + return None + + system = _one_person_system(result) + system.auto_carry_over_input_variables = True + branch = SimulationBuilder().build_default_simulation(system).get_branch("consumer") + + assert branch.calculate("result", "2020").tolist() == [6.0] + assert _stored_keys(branch, "result") == {"consumer:2020"} + # Only inputs carry over, so a later year reads the default either way. + assert branch.calculate("result", "2021").tolist() == [0.0] + + +def test_refused_result_is_not_kept_by_a_caller_in_another_family(): + """A worker (its own family) calls back into the consumer, which calls the worker again.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + def variable(name, formula=None): + namespace = dict( + value_type=float, + entity=entities.Person, + definition_period=periods.YEAR, + label=name, + ) + if formula is not None: + namespace["formula"] = formula + return type(name, (Variable,), namespace) + + def run(input_first): + holders = {} + + def leaf(person, period): + previous = person("source", period) + if np.any(previous != 3.0): + person.simulation.set_input( + "source", period, np.full(person.count, 3.0) + ) + return previous * 2 + + def worker_outer(person, period): + return holders["consumer"].calculate("from_worker", period) + + worker_root = SimulationBuilder().build_default_simulation( + _one_person_system( + variable("source"), + variable("leaf", leaf), + variable("worker_outer", worker_outer), + ) + ) + worker_root.set_input("source", "2020", np.array([1.0])) + worker = worker_root.get_branch("worker") + if input_first: + worker.set_input("source", "2020", np.array([3.0])) + + consumer = ( + SimulationBuilder() + .build_default_simulation( + _one_person_system( + variable( + "from_worker", + lambda person, period: worker.calculate("leaf", period), + ), + variable( + "result", + lambda person, period: worker.calculate("worker_outer", period), + ), + ) + ) + .get_branch("consumer") + ) + holders["consumer"] = consumer + return ( + consumer.calculate("result", "2020").tolist(), + consumer.calculate("from_worker", "2020").tolist(), + ) + + assert run(input_first=False) == run(input_first=True) == ([6.0], [6.0]) + + +def test_caller_in_the_same_family_runs_again_after_the_change(): + """A parent read its branch, then called a branch formula that changed the branch's input.""" + + def run(input_first): + def changes_source(person, period): + if np.any(person("source", period) != 3.0): + person.simulation.set_input( + "source", period, np.full(person.count, 3.0) + ) + return np.zeros(person.count) + + def result(person, period): + branch = person.simulation.branches["kept"] + doubled = branch.calculate("doubled", period) # read before the change + branch.calculate("changes_source", period) + return doubled + + system = _one_person_system( + _yearly_variable("source"), + _yearly_variable( + "doubled", lambda person, period: person("source", period) * 2 + ), + _yearly_variable("changes_source", changes_source), + _yearly_variable("result", result), + ) + root = SimulationBuilder().build_default_simulation(system) + root.set_input("source", "2020", np.array([1.0])) + kept = root.get_branch("kept") + if input_first: + kept.set_input("source", "2020", np.array([3.0])) + first = root.calculate("result", "2020").tolist() + return first, root.get_array("result", "2020").tolist() + + # The first run read 2 before the change, so the formula ran again. + assert run(input_first=False) == run(input_first=True) == ([6.0], [6.0]) + + +def _worker_and_consumer(leaf, auto_carry_over=False, worker_inputs=()): + """A worker family whose ``worker_outer`` calls back into a consumer's ``from_worker``.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + def variable(name, formula=None): + namespace = dict( + value_type=float, + entity=entities.Person, + definition_period=periods.YEAR, + label=name, + ) + if formula is not None: + namespace["formula"] = formula + return type(name, (Variable,), namespace) + + def run(input_first): + holders = {} + worker_system = _one_person_system( + variable("source"), + variable("leaf", leaf), + variable( + "worker_outer", + lambda person, period: holders["consumer"].calculate( + "from_worker", period + ), + ), + ) + worker_system.auto_carry_over_input_variables = auto_carry_over + worker_root = SimulationBuilder().build_default_simulation(worker_system) + worker_root.set_input("source", "2020", np.array([1.0])) + for variable_name, period, value in worker_inputs: + worker_root.set_input(variable_name, period, np.array([value])) + worker = worker_root.get_branch("worker") + if input_first: + worker.set_input("source", "2020", np.array([3.0])) + consumer = ( + SimulationBuilder() + .build_default_simulation( + _one_person_system( + variable( + "from_worker", + lambda person, period: worker.calculate("leaf", period), + ), + variable( + "result", + lambda person, period: worker.calculate("worker_outer", period), + ), + ) + ) + .get_branch("consumer") + ) + holders["consumer"] = consumer + first = consumer.calculate("result", "2020").tolist() + return first, consumer.calculate("result", "2020").tolist() + + return run + + +def test_refused_result_that_returns_early_still_taints_its_caller(): + """A carry-over default returned before any cache decision, after an input change.""" + + def leaf(person, period): + if np.any(person("source", period) != 3): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + return None # carry-over finds only 2021: the default, early + return np.full(person.count, 6.0) + + run = _worker_and_consumer( + leaf, auto_carry_over=True, worker_inputs=[("leaf", "2021", 9.0)] + ) + + assert run(input_first=False) == run(input_first=True) == ([6.0], [6.0]) + + +def test_caller_that_catches_an_error_after_an_input_change_runs_again(): + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + def leaf(person, period): + if np.any(person("source", period) != 3): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + raise ValueError("changed") + return np.full(person.count, 6.0) + + worker_root = SimulationBuilder().build_default_simulation( + _one_person_system(_yearly_variable("source"), _yearly_variable("leaf", leaf)) + ) + worker_root.set_input("source", "2020", np.array([1.0])) + worker = worker_root.get_branch("worker") + + def result(person, period): + try: + return worker.calculate("leaf", period) + except ValueError: + return np.zeros(person.count) + + consumer = ( + SimulationBuilder() + .build_default_simulation( + _one_person_system(_yearly_variable("result", result)) + ) + .get_branch("consumer") + ) + + # The first run fell back to 0 after the change, so the formula ran again. + assert consumer.calculate("result", "2020").tolist() == [6.0] + assert consumer.get_array("result", "2020").tolist() == [6.0] + + +@pytest.mark.parametrize("worker_kind", ["other_family", "own_branch_in_a_thread"]) +def test_earlier_read_is_read_again_after_the_worker_changes_its_input(worker_kind): + """A formula that read a value before its simulation's input changed runs again.""" + import threading + + def changes_source(person, period): + if np.any(person("source", period) != 3.0): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + return np.zeros(person.count) + + worker_variables = ( + _yearly_variable("source"), + _yearly_variable( + "doubled", lambda person, period: person("source", period) * 2 + ), + _yearly_variable("changes_source", changes_source), + ) + + if worker_kind == "other_family": + worker_root = SimulationBuilder().build_default_simulation( + _one_person_system(*worker_variables) + ) + worker_root.set_input("source", "2020", np.array([1.0])) + worker = worker_root.get_branch("worker") + + def result(person, period): + doubled = worker.calculate("doubled", period) + worker.calculate("changes_source", period) + return doubled + + simulation = SimulationBuilder().build_default_simulation( + _one_person_system(_yearly_variable("result", result)) + ) + else: + + def result(person, period): + branch = person.simulation.branches["kept"] + doubled = branch.calculate("doubled", period) + thread = threading.Thread( + target=branch.calculate, args=("changes_source", period) + ) + thread.start() + thread.join() + return doubled + + simulation = SimulationBuilder().build_default_simulation( + _one_person_system(*worker_variables, _yearly_variable("result", result)) + ) + simulation.set_input("source", "2020", np.array([1.0])) + simulation.get_branch("kept") + + assert simulation.calculate("result", "2020").tolist() == [6.0] + assert simulation.get_array("result", "2020").tolist() == [6.0] + + +def test_settled_call_into_own_branch_keeps_the_result(): + """The branch's change settles before returning: nothing the caller holds is old.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + def leaf(person, period): + if np.any(person("source", period) != 3): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + return person("source", period) * 2 + + class result(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "result" + + def formula_2020(person, period): + return person.simulation.branches["worker"].calculate("leaf", "2020") + + def formula_2021(person, period): + return None + + system = _one_person_system( + _yearly_variable("source"), _yearly_variable("leaf", leaf), result + ) + system.auto_carry_over_input_variables = True + root = SimulationBuilder().build_default_simulation(system) + root.set_input("source", "2020", np.array([1.0])) + root.get_branch("worker") + + assert root.calculate("result", "2020").tolist() == [6.0] + assert _stored_keys(root, "result") == {"default:2020"} + # Only inputs carry over, so a later year reads the default either way. + assert root.calculate("result", "2021").tolist() == [0.0] + + +def test_input_a_formula_sets_for_its_own_period_wins(): + """As it would had the input been set before the calculation.""" + + def source(person, period): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + return np.ones(person.count) + + system = _one_person_system(_yearly_variable("source", source)) + branch = SimulationBuilder().build_default_simulation(system).get_branch("b") + + assert branch.calculate("source", "2020").tolist() == [3.0] + assert branch.calculate("source", "2020").tolist() == [3.0] + assert branch.get_holder("source")._is_input(periods.period("2020"), "b") + + +def test_input_handler_that_calculates_between_its_stores(): + """Values a handler calculates before its last store are dropped after it.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + def two_months(holder, period, array): + holder._set(period.first_month, array) + holder.simulation.calculate("total", period) + holder._set(period.first_month.offset(1), array) + + class source(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "source" + set_input = two_months + + class total(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "total" + + def formula(person, period): + first = period.first_month + return person("source", first) + person("source", first.offset(1)) + + system = _one_person_system(source, total) + root = SimulationBuilder().build_default_simulation(system) + root.set_input("source", "2020-01", np.array([1.0])) + root.set_input("source", "2020-02", np.array([1.0])) + branch = root.get_branch("branch") + branch.set_input("source", "2020", np.array([3.0])) + + assert branch.calculate("total", "2020").tolist() == [6.0] + + +def test_input_handler_that_fails_after_calculating(): + """The inputs it stored stay, so what it calculated before them drops.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + def two_months_then_fail(holder, period, array): + holder._set(period.first_month, array) + holder.simulation.calculate("total", period) + holder._set(period.first_month.offset(1), array) + raise RuntimeError("handler failed") + + class source(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "source" + set_input = two_months_then_fail + + class total(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "total" + + def formula(person, period): + first = period.first_month + return person("source", first) + person("source", first.offset(1)) + + system = _one_person_system(source, total) + root = SimulationBuilder().build_default_simulation(system) + root.set_input("source", "2020-01", np.array([1.0])) + root.set_input("source", "2020-02", np.array([1.0])) + branch = root.get_branch("branch") + with pytest.raises(RuntimeError, match="handler failed"): + branch.set_input("source", "2020", np.array([3.0])) + + assert branch.calculate("total", "2020").tolist() == [6.0] + + +def test_branch_input_stops_macro_reads_before_any_read(monkeypatch): + """A macro file for the branch's name may come from a branch without the input.""" + import policyengine_core.simulations.simulation as simulation_module + + class CachedResult: + def __init__(self, tax_benefit_system): + pass + + def set_cache_path(self, *args): + pass + + def get_cache_path(self): + return type("Path", (), {"exists": lambda self: True})() + + def get_cache_value(self, path): + return np.array([2.0]) # written by another simulation's "override" + + def set_cache_value(self, path, value): + pass + + monkeypatch.setattr(simulation_module, "SimulationMacroCache", CachedResult) + system = _one_person_system( + _yearly_variable("source", lambda person, period: np.ones(person.count)), + _yearly_variable("result", lambda person, period: person("source", period) * 2), + _yearly_variable("total", lambda person, period: person("result", period) * 2), + ) + simulation = SimulationBuilder().build_default_simulation(system) + system.data_modified = False + simulation.macro_cache_read = True + simulation.dataset = type( + "Dataset", (), {"file_path": simulation_module.Path("cache"), "name": "d"} + )() + monkeypatch.setattr( + type(simulation), + "check_macro_cache", + lambda self, variable, period: variable == "result", + ) + + branch = simulation.get_branch("override") + branch.set_input("source", "2020", np.array([3.0])) # before any calculation + + assert branch.calculate("total", "2020").tolist() == [12.0] + + +def test_result_obsolete_after_its_own_input_change_is_not_written_to_the_macro_cache( + monkeypatch, +): + import policyengine_core.simulations.simulation as simulation_module + + written = [] + + class RecordingCache: + def __init__(self, tax_benefit_system): + pass + + def set_cache_path(self, *args): + pass + + def get_cache_path(self): + return type("Path", (), {"exists": lambda self: False})() + + def get_cache_value(self, path): + return None + + def set_cache_value(self, path, value): + written.append(np.array(value).tolist()) + + def result(person, period): + source = person("source", period) + if np.any(source == 1): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + return source * 2 + + monkeypatch.setattr(simulation_module, "SimulationMacroCache", RecordingCache) + system = _one_person_system( + _yearly_variable("source", lambda person, period: np.ones(person.count)), + _yearly_variable("result", result), + ) + simulation = SimulationBuilder().build_default_simulation(system) + simulation.dataset = type( + "Dataset", (), {"file_path": simulation_module.Path("cache"), "name": "d"} + )() + monkeypatch.setattr( + type(simulation), + "check_macro_cache", + lambda self, variable, period: variable == "result", + ) + branch = simulation.get_branch("b") + + assert branch.calculate("result", "2020").tolist() == [6.0] + assert [2.0] not in written # the first run's result, from the old input + + +def test_macro_cache_read_counts_as_depending_on_every_input(monkeypatch): + """What a macro-cache value was calculated from was never calculated here.""" + import policyengine_core.simulations.simulation as simulation_module + + class CachedResult: + def __init__(self, tax_benefit_system): + pass + + def set_cache_path(self, *args): + pass + + def get_cache_path(self): + return type("Path", (), {"exists": lambda self: True})() + + def get_cache_value(self, path): + return np.array([2.0]) + + def set_cache_value(self, path, value): + pass + + monkeypatch.setattr(simulation_module, "SimulationMacroCache", CachedResult) + system = _one_person_system( + _yearly_variable("source", lambda person, period: np.ones(person.count)), + _yearly_variable("result", lambda person, period: person("source", period) * 2), + _yearly_variable("total", lambda person, period: person("result", period) * 2), + ) + simulation = SimulationBuilder().build_default_simulation(system) + system.data_modified = False # adding variables marks the system modified + simulation.macro_cache_read = True + simulation.dataset = type( + "Dataset", (), {"file_path": simulation_module.Path("cache"), "name": "d"} + )() + monkeypatch.setattr( + type(simulation), + "check_macro_cache", + lambda self, variable, period: variable == "result", + ) + assert simulation.calculate("total", "2020").tolist() == [4.0] + assert not simulation.get_holder("source").get_known_periods() # never calculated + + branch = simulation.get_branch("branch") + branch.set_input("source", "2020", np.array([3.0])) # never calculated here + + assert branch.calculate("total", "2020").tolist() == [12.0] + + +def test_history_hand_back_failure_still_ends_the_calculation(monkeypatch): + def fail(self): + raise RuntimeError("hand-back failed") + + system = _one_person_system( + _yearly_variable("source", lambda person, period: np.ones(person.count)) + ) + simulation = SimulationBuilder().build_default_simulation(system) + monkeypatch.setattr(type(simulation), "_share_store_history_with_caller", fail) + + with pytest.raises(RuntimeError, match="hand-back failed"): + simulation.calculate("source", "2020") + + assert simulation._calculations_in_flight == 0 + assert not simulation.tracer.stack + + +def test_failed_calculation_in_another_simulation_hands_its_history_back(): + """Whether a calculation fails can depend on what it read.""" + + def checked(person, period): + source = person("source", period) + if (source == 0).any(): + raise ValueError("source is zero") + return source * 2 + + def result(person, period): + simulation = person.simulation + child = simulation.get_branch("child") + try: + return child.calculate("checked", period) + except ValueError: + return np.full(person.count, 7.0) + finally: + del simulation.branches["child"] + + system = _one_person_system( + _yearly_variable("source", lambda person, period: np.zeros(person.count)), + _yearly_variable("checked", checked), + _yearly_variable("result", result), + ) + simulation = SimulationBuilder().build_default_simulation(system) + assert simulation.calculate("result", "2020").tolist() == [7.0] + + branch = simulation.get_branch("branch") + branch.set_input("source", "2020", np.array([3.0])) + + assert branch.calculate("result", "2020").tolist() == [6.0] + + +def test_restored_simulation_drops_restored_values_an_input_may_have_fed(tmp_path): + from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, + ) + + system = _one_person_system( + _yearly_variable("source"), + _yearly_variable("result", lambda person, period: person("source", period) * 2), + ) + simulation = SimulationBuilder().build_default_simulation(system) + simulation.set_input("source", "2020", np.array([1.0])) + simulation.calculate("result", "2020") + dump_simulation(simulation, str(tmp_path / "dump")) + + restored = restore_simulation(str(tmp_path / "dump"), system) + branch = restored.get_branch("branch") + branch.set_input("source", "2020", np.array([3.0])) + + assert branch.calculate("result", "2020").tolist() == [6.0] + assert restored.calculate("source", "2020").tolist() == [1.0] # still an input + + +@pytest.mark.parametrize("read", ["carried_over", "not_kept"]) +def test_restored_values_drop_without_records_of_what_they_read(tmp_path, read): + """A dump keeps values, not what each was calculated from.""" + from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, + ) + + system = _one_person_system( + _yearly_variable("source"), + _yearly_variable( + "result", lambda person, period: person("source", period) * 2 + 7 + ), + ) + simulation = SimulationBuilder().build_default_simulation(system) + if read == "carried_over": + system.auto_carry_over_input_variables = True + simulation.set_input("source", "2020", np.array([1.0])) + override, fresh_inputs = "2021", {"2020": 1.0, "2021": 3.0} + else: + system.cache_blacklist = {"source"} + simulation.opt_out_cache = True + override, fresh_inputs = "2022", {"2022": 3.0} + simulation.calculate("result", "2022") + dump_simulation(simulation, str(tmp_path / "dump")) + + restored = restore_simulation(str(tmp_path / "dump"), system) + branch = restored.get_branch("branch") + branch.set_input("source", override, np.array([3.0])) + + fresh = SimulationBuilder().build_default_simulation(system) + for period, value in fresh_inputs.items(): + fresh.set_input("source", period, np.array([value])) + assert ( + branch.calculate("result", "2022").tolist() + == fresh.calculate("result", "2022").tolist() + ) + + +def test_dump_of_a_branch_holds_the_branch_values(tmp_path): + from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, + ) + + system = _one_person_system( + _yearly_variable("source"), + _yearly_variable("result", lambda person, period: person("source", period) * 2), + ) + simulation = SimulationBuilder().build_default_simulation(system) + simulation.set_input("source", "2020", np.array([1.0])) + simulation.set_input("source", "2021", np.array([5.0])) + branch = simulation.get_branch("branch") + branch.set_input("source", "2020", np.array([3.0])) + branch.calculate("result", "2020") + dump_simulation(branch, str(tmp_path / "dump")) + + restored = restore_simulation(str(tmp_path / "dump"), system) + + assert restored.calculate("source", "2020").tolist() == [3.0] + assert restored.calculate("source", "2021").tolist() == [5.0] # inherited + assert restored.calculate("result", "2020").tolist() == [6.0] + holder = restored.get_holder("source") + assert holder._is_input(periods.period("2020")) + assert holder._is_input(periods.period("2021")) + + +def test_merge_reads_again_what_was_recorded_while_it_merged(): + """A record added to the source during a merge (another thread) is not skipped.""" + source = StoreHistory() + source.record_store("a", periods.period("2020"), 1) + destination = StoreHistory() + original = destination.record_store + calls = [] + + def record_store(variable_name, period, sequence_number): + if not calls: # the other thread records while this merge runs + source.record_store("v", periods.period("2020"), 2) + calls.append(variable_name) + original(variable_name, period, sequence_number) + + destination.record_store = record_store + destination.merge(source) + destination.merge(source) + + assert destination.earliest_dependency("v", periods.period("2020")) == 2 + + +def test_history_pickled_before_journals_still_records(): + import pickle + + history = StoreHistory() + history.record_store("v", periods.period("2020"), 1) + del history._journal, history._generation # as pickled before journals + restored = pickle.loads(pickle.dumps(history)) + + restored.record_store("w", periods.period("2020"), 2) + + assert restored.earliest_dependency("w", periods.period("2020")) == 2 + + +# ----- Only what may depend on the input is dropped ----- # + + +def test_input_on_one_branch_does_not_make_another_drop(tax_benefit_system): + simulation = _build(tax_benefit_system) + sibling = simulation.get_branch("sibling") + sibling.set_input("rent", "2017-02", np.array([500.0])) + sibling.calculate("housing_allowance", "2017-02") + simulation.calculate("disposable_income", JANUARY) + held = { + (name, key) + for name in ("disposable_income", "income_tax") + for key in _stored_keys(simulation, name) + } + + branch = simulation.get_branch("branch") + dropped = branch._drop_values_that_may_depend_on("rent", periods.period("2017-02")) + + assert dropped == 0 + assert held <= { + (name, key) + for name in ("disposable_income", "income_tax") + for key in _stored_keys(branch, name) + } + + +def test_branches_that_choose_between_overrides_calculate_shared_values_once(): + """Values the overridden variable cannot reach are calculated once per arm. + + ``choose`` compares ``tax`` in two branches that set it, as + policyengine-us itemization does; ``agi`` does not depend on it. + """ + # A marginal-rate branch after the parent has calculated everything. + simulation = branching_simulation() + simulation.calculate("net", "2020") + FORMULA_RUNS["agi"] = 0 + rate = simulation.calculate("marginal_rate", "2020") + assert FORMULA_RUNS["agi"] == 1 + assert np.allclose(rate, 0.85) + + # A baseline branch calculating after the reform arm has. + FORMULA_RUNS["agi"] = 0 + simulation = branching_simulation() + baseline = simulation.get_branch("baseline") + simulation.calculate("net", "2020") + baseline.calculate("net", "2020") + assert FORMULA_RUNS["agi"] == 2 + + fresh = branching_simulation() + assert np.array_equal( + baseline.calculate("net", "2020"), fresh.calculate("net", "2020") + ) + + +# ----- Other paths ----- # + + +def test_prerequisite_requested_before_a_branch_input_still_counts(): + """``requires_computation_after`` holds if the prerequisite was requested, even if dropped.""" + from policyengine_core.entities import build_entity + from policyengine_core.taxbenefitsystems import TaxBenefitSystem + from policyengine_core.variables import Variable + + person = build_entity("person", "persons", "Person", is_person=True) + + class source(Variable): + label = "Source" + value_type = float + entity = person + definition_period = periods.YEAR + + class prerequisite(Variable): + label = "Prerequisite" + value_type = float + entity = person + definition_period = periods.YEAR + + def formula(person, period): + return person("source", period) * 0 + 1 + + class result(Variable): + label = "Result" + value_type = float + entity = person + definition_period = periods.YEAR + requires_computation_after = "prerequisite" + + def formula(person, period): + return person("source", period) * 2 + + system = TaxBenefitSystem([person]) + system.add_variables(source, prerequisite, result) + simulation = SimulationBuilder().build_default_simulation(system) + simulation.set_input("source", "2022", np.array([1.0])) + simulation.calculate("prerequisite", "2022") + simulation.calculate("result", "2022") + + branch = simulation.get_branch("branch") + branch.set_input("source", "2022", np.array([2.0])) + + assert branch.calculate("result", "2022").tolist() == [4.0] + + +def test_values_a_custom_input_handler_calculates_are_not_inputs(): + from policyengine_core.entities import build_entity + from policyengine_core.taxbenefitsystems import TaxBenefitSystem + from policyengine_core.variables import Variable + + person = build_entity("person", "persons", "Person", is_person=True) + + def store_first_month_then_calculate(holder, period, array): + holder._set(period.first_month, array) + holder.simulation.calculate("dependent", period) + + class dispatched(Variable): + label = "Dispatched" + value_type = float + entity = person + definition_period = periods.MONTH + set_input = store_first_month_then_calculate + + class dependent(Variable): + label = "Dependent" + value_type = float + entity = person + definition_period = periods.YEAR + + def formula(person, period): + return 2 * person("dispatched", period.first_month) + + system = TaxBenefitSystem([person]) + system.add_variables(dispatched, dependent) + branch = SimulationBuilder().build_default_simulation(system).get_branch("branch") + branch.set_input("dispatched", "2020", np.array([3.0])) + assert branch.calculate("dependent", "2020").tolist() == [6.0] + + branch.set_input("dispatched", "2020-01", np.array([4.0])) + + assert branch.calculate("dependent", "2020").tolist() == [8.0] + assert not branch.get_holder("dependent")._memory_storage._inputs + + +def test_year_input_after_calculating_one_of_its_months(): + """The drop comes before a yearly input is divided, so only input months count.""" + simulation = synthetic_simulation(ROOT_INPUTS) + branch = simulation.get_branch("branch") + branch.calculate("p_month", "2015-03") # 2015 has no p_m input: default + + branch.set_input("p_m", "2015", np.array([120.0, 120.0, 120.0])) + + fresh = synthetic_simulation( + {**ROOT_INPUTS, **{("p_m", f"2015-{m:02d}"): (10.0,) * 3 for m in range(1, 13)}} + ) + assert np.allclose( + branch.calculate("p_month", "2015-03"), fresh.calculate("p_month", "2015-03") + ) + + +def test_derivative_keeps_inputs_set_after_construction(): + """``derivative`` keeps every input of the simulation it differentiates.""" + simulation = synthetic_simulation({("p_a", "2013"): (1.0, 1.0, 1.0)}) + branch = simulation.get_branch("branch") + branch.set_input("p_b", "2013", np.array([9.0, 9.0, 9.0])) + + assert np.allclose(branch.derivative("p_sum", "p_b", "2013"), 2.0) + assert np.array_equal(branch.calculate("p_b", "2013"), [9.0, 9.0, 9.0]) + + +def test_disk_branch_keeps_its_values_when_its_parent_recalculates(): + """Each disk store writes its own file, so a recalculation leaves children's files.""" + config = MemoryConfig(max_memory_occupation=0) + simulation = synthetic_simulation(ROOT_INPUTS, memory_config=config) + branch = simulation.get_branch("branch") + before = branch.calculate("p_sum", "2013").copy() + child = branch.get_branch("child") + + branch.set_input("p_a", "2013", np.zeros(3)) + branch.calculate("p_sum", "2013") # recalculated: a new file + + assert np.array_equal(child.calculate("p_sum", "2013"), before) + + +def test_disk_branches_with_the_same_name_keep_their_own_values(): + config = MemoryConfig(max_memory_occupation=0) + simulation = synthetic_simulation(ROOT_INPUTS, memory_config=config) + first = simulation.get_branch("a").get_branch("leaf") + second = simulation.get_branch("b").get_branch("leaf") + first.set_input("p_a", "2013", np.full(3, 3.0)) + second.set_input("p_a", "2013", np.full(3, 4.0)) + # A name reused after its branch is forgotten. + old = simulation.get_branch("reused") + old.set_input("p_a", "2013", np.full(3, 7.0)) + del simulation.branches["reused"] + new = simulation.get_branch("reused") + new.set_input("p_a", "2013", np.full(3, 8.0)) + + assert np.array_equal(first.calculate("p_a", "2013"), [3.0] * 3) + assert np.array_equal(second.calculate("p_a", "2013"), [4.0] * 3) + assert np.array_equal(old.calculate("p_a", "2013"), [7.0] * 3) + assert np.array_equal(new.calculate("p_a", "2013"), [8.0] * 3) + + +def test_values_unpickled_from_another_process_stay_earlier(monkeypatch): + """Sequence numbers restart in each process; unpickling moves this one past them.""" + import itertools + import pickle + + import policyengine_core.data_storage.store_history as store_history + + storage = InMemoryStorage(is_eternal=False) + history = StoreHistory() + for _ in range(100): + next_sequence_number() + storage.put(np.array([1.0]), periods.period("2020")) + history.record_store("v", periods.period("2020"), next_sequence_number()) + source = StoreHistory() # kept alive: its weak read position must not stop pickling + history.merge(source) + assert len(history._merged) == 1 + payloads = pickle.dumps(storage), pickle.dumps(history) + + monkeypatch.setattr(store_history, "_sequence", itertools.count(1)) # a new process + restored_storage, restored_history = (pickle.loads(p) for p in payloads) + restored_storage.put(np.array([2.0]), periods.period("2021")) + + numbers = restored_storage._sequence_numbers + assert numbers["default:2021"] > numbers["default:2020"] + assert next_sequence_number() > restored_history.earliest_dependency( + "v", periods.period("2020") + ) + + +def test_storage_pickled_before_stores_were_numbered_still_stores(): + import pickle + + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0]), periods.period("2020"), derived=True) + del storage._sequence_numbers # as pickled before + restored = pickle.loads(pickle.dumps(storage)) + + restored.put(np.array([2.0]), periods.period("2021"), derived=True) + + assert set(restored._sequence_numbers) == {"default:2021"} + assert restored.has_unnumbered_values() # 2020: anything may depend on it + assert restored.drop_computed(since=10**18) == 1 # unnumbered counts as later + assert set(restored._arrays) == {"default:2021"} + + +# ----- Reads outside a frame, reads before settling, idle simulations ----- # + + +def _in_thread(call, copy_context=False): + """Run ``call`` in a thread started with an empty (or a copied) context.""" + import contextvars + import threading + + returned, failures = [], [] + + def target(): + try: + returned.append(call()) + except BaseException as error: # re-raised in the caller + failures.append(error) + + context = contextvars.copy_context() if copy_context else contextvars.Context() + thread = threading.Thread(target=context.run, args=(target,)) + thread.start() + thread.join() + if failures: + raise failures[0] + return returned[0] + + +@pytest.mark.parametrize("change_in_a_thread", [False, True]) +@pytest.mark.parametrize("copy_context", [False, True]) +def test_value_read_from_own_branch_in_a_thread_is_read_again_after_a_change( + change_in_a_thread, copy_context +): + """A read in a thread started without the formula's context still counts.""" + + def changes_source(person, period): + if np.any(person("source", period) != 3): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + return np.zeros(person.count) + + def result(person, period): + worker = person.simulation.branches["worker"] + doubled = _in_thread(lambda: worker.calculate("doubled", period), copy_context) + + def change(): + return worker.calculate("changes_source", period) + + if change_in_a_thread: + _in_thread(change, copy_context) + else: + change() + return doubled + + system = _one_person_system( + _yearly_variable("source"), + _yearly_variable( + "doubled", lambda person, period: person("source", period) * 2 + ), + _yearly_variable("changes_source", changes_source), + _yearly_variable("result", result), + ) + + def run(input_first): + root = SimulationBuilder().build_default_simulation(system) + root.set_input("source", "2020", np.array([1.0])) + worker = root.get_branch("worker") + if input_first: + worker.set_input("source", "2020", np.array([3.0])) + first = root.calculate("result", "2020").tolist() + return first, root.get_array("result", "2020").tolist() + + assert run(input_first=False) == run(input_first=True) == ([6.0], [6.0]) + + +def test_fast_cache_hit_in_a_thread_counts_as_a_read(): + """A value already calculated, read again from a thread without the context.""" + + def changes_source(person, period): + if np.any(person("source", period) != 3): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + return np.zeros(person.count) + + def result(person, period): + worker = person.simulation.branches["worker"] + doubled = _in_thread(lambda: worker.calculate("doubled", period)) + worker.calculate("changes_source", period) + return doubled + + system = _one_person_system( + _yearly_variable("source"), + _yearly_variable( + "doubled", lambda person, period: person("source", period) * 2 + ), + _yearly_variable("changes_source", changes_source), + _yearly_variable("result", result), + ) + root = SimulationBuilder().build_default_simulation(system) + root.set_input("source", "2020", np.array([1.0])) + worker = root.get_branch("worker") + assert worker.calculate("doubled", "2020").tolist() == [2.0] # now a fast hit + + assert root.calculate("result", "2020").tolist() == [6.0] + assert root.get_array("result", "2020").tolist() == [6.0] + + +def _consumer_of(worker, result_2020): + """A simulation in another family whose ``result`` for 2020 reads ``worker``.""" + from policyengine_core.country_template import entities + from policyengine_core.variables import Variable + + class result(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "result" + + def formula_2020(person, period): + return result_2020(person, period) + + return SimulationBuilder().build_default_simulation(_one_person_system(result)) + + +def _worker(leaf, input_first): + worker_root = SimulationBuilder().build_default_simulation( + _one_person_system(_yearly_variable("source"), _yearly_variable("leaf", leaf)) + ) + worker_root.set_input("source", "2020", np.array([1.0])) + worker = worker_root.get_branch("worker") + if input_first: + worker.set_input("source", "2020", np.array([3.0])) + return worker + + +@pytest.mark.parametrize("input_first", [False, True]) +def test_value_recovered_after_a_failed_call_is_kept(input_first): + """The call that failed read the old input; the retry's value is current.""" + + def leaf(person, period): + if np.any(person("source", period) != 3): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + raise ValueError("input was normalised; retry") + return person("source", period) * 2 + + worker = _worker(leaf, input_first) + + def result(person, period): + try: + return worker.calculate("leaf", period) + except ValueError: + return worker.calculate("leaf", period) + + consumer = _consumer_of(worker, result) + + assert consumer.calculate("result", "2020").tolist() == [6.0] + assert consumer.get_array("result", "2020").tolist() == [6.0] + + +@pytest.mark.parametrize("input_first", [False, True]) +def test_value_read_after_the_callee_settles_is_kept(input_first): + """A formula checks the worker, asks it to settle, then returns its value.""" + + def normalize(person, period): + if np.any(person("source", period) != 3): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + return np.zeros(person.count) + + def leaf(person, period): + return person("source", period) * 2 + + worker = _worker(leaf, input_first) + worker.tax_benefit_system.add_variable(_yearly_variable("normalize", normalize)) + + def result(person, period): + if np.any(worker.calculate("source", period) != 3): + worker.calculate("normalize", period) + return worker.calculate("leaf", period) + + consumer = _consumer_of(worker, result) + + assert consumer.calculate("result", "2020").tolist() == [6.0] + assert consumer.get_array("result", "2020").tolist() == [6.0] + + +@pytest.mark.parametrize("input_first", [False, True]) +def test_input_set_directly_on_an_idle_simulation_after_reading_it(input_first): + """The worker has finished calculating when the formula replaces its input.""" + worker = _worker(lambda person, period: person("source", period) * 2, input_first) + + def result(person, period): + earlier = worker.calculate("leaf", period) + worker.set_input("source", period, np.full(person.count, 3.0)) + return earlier + + consumer = _consumer_of(worker, result) + + assert consumer.calculate("result", "2020").tolist() == [6.0] + assert consumer.get_array("result", "2020").tolist() == [6.0] + assert worker.calculate("leaf", "2020").tolist() == [6.0] + + +# ----- Setting the input a branch already reads ----- # + + +def test_setting_the_input_a_branch_already_reads_drops_nothing(): + simulation = synthetic_simulation(ROOT_INPUTS) + branch = simulation.get_branch("branch") + branch.calculate("p_sum", "2013") + held = _stored_keys(branch, "p_sum") + epoch = getattr(branch, "_input_epoch", 0) + + branch.set_input("p_a", "2013", np.asarray(ROOT_INPUTS[("p_a", "2013")])) + + assert _stored_keys(branch, "p_sum") == held + assert getattr(branch, "_input_epoch", 0) == epoch + branch.set_input("p_a", "2013", np.asarray(ROOT_INPUTS[("p_a", "2013")]) + 1) + assert not _stored_keys(branch, "p_sum") & held + assert getattr(branch, "_input_epoch", 0) > epoch + + +def test_input_equal_to_a_calculated_value_still_drops(): + """An input where a value was calculated changes which periods hold inputs.""" + simulation = synthetic_simulation(ROOT_INPUTS) + branch = simulation.get_branch("branch") + calculated = branch.calculate("p_sum", "2013").copy() + branch.calculate("p_prod", "2013") + + branch.set_input("p_sum", "2013", calculated) + + assert _stored_keys(branch, "p_prod") == set() + + +def test_input_a_helper_spreads_over_months_still_drops(): + simulation = synthetic_simulation(ROOT_INPUTS) + branch = simulation.get_branch("branch") + yearly = sum( + np.asarray(ROOT_INPUTS[("p_m", f"2013-{m:02d}")]) for m in range(1, 13) + ) + branch.calculate("p_month", "2013-07") + epoch = getattr(branch, "_input_epoch", 0) + + branch.set_input("p_m", "2013", yearly) + + assert getattr(branch, "_input_epoch", 0) > epoch + + +def test_unchanged_branch_input_still_stops_macro_reads(monkeypatch): + """The file for the branch's name may come from another simulation's branch.""" + + class OtherBranchesFile: + def __init__(self, tax_benefit_system): + pass + + def set_cache_path(self, *args): + pass + + def get_cache_path(self): + return type("Path", (), {"exists": lambda self: True})() + + def get_cache_value(self, path): + return np.array([10.0]) # another simulation's "override", source 5 + + def set_cache_value(self, path, value): + pass + + monkeypatch.setattr(simulation_module, "SimulationMacroCache", OtherBranchesFile) + system = _one_person_system( + _yearly_variable("source"), + _yearly_variable("result", lambda person, period: person("source", period) * 2), + _yearly_variable("total", lambda person, period: person("result", period) * 2), + ) + simulation = SimulationBuilder().build_default_simulation(system) + simulation.set_input("source", "2020", np.array([1.0])) + system.data_modified = False + simulation.macro_cache_read = True + simulation.dataset = type( + "Dataset", (), {"file_path": simulation_module.Path("cache"), "name": "d"} + )() + monkeypatch.setattr( + type(simulation), + "check_macro_cache", + lambda self, variable, period: variable == "result", + ) + + branch = simulation.get_branch("override") + branch.set_input("source", "2020", np.array([1.0])) # the input it already reads + + assert branch.calculate("total", "2020").tolist() == [4.0] + + +def test_negative_zero_is_not_the_input_zero(): + """Equal as numbers, but a formula dividing by it tells them apart.""" + system = _one_person_system( + _yearly_variable("source"), + _yearly_variable( + "reciprocal", + lambda person, period: np.sign(1 / person("source", period)), + ), + ) + simulation = SimulationBuilder().build_default_simulation(system) + simulation.set_input("source", "2020", np.array([0.0])) + branch = simulation.get_branch("branch") + with np.errstate(divide="ignore"): + assert branch.calculate("reciprocal", "2020").tolist() == [1.0] + branch.set_input("source", "2020", np.array([-0.0])) + assert branch.calculate("reciprocal", "2020").tolist() == [-1.0] + + +def test_input_given_to_a_helper_drops_even_if_equal_to_a_stored_year(): + """The helper spreads it over months, which a yearly key does not show.""" + from policyengine_core.country_template import entities + from policyengine_core.holders import set_input_divide_by_period + from policyengine_core.variables import Variable + + class m(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "m" + set_input = set_input_divide_by_period + + class r(Variable): + value_type = float + entity = entities.Person + definition_period = periods.MONTH + label = "r" + + def formula(person, period): + return person("m", period) * 2 + + def build(): + system = _one_person_system(m, r) + system.auto_carry_over_input_variables = True + return SimulationBuilder().build_default_simulation(system) + + root = build() + # An input stored for the year itself, not through the helper. + root.get_holder("m").put_in_cache(np.array([12.0]), periods.period("2020")) + branch = root.get_branch("branch") + assert branch.calculate("r", "2020-03").tolist() == [24.0] # carried from 2020 + + branch.set_input("m", "2020", np.array([12.0])) # divided: 1 a month + + fresh = build() + fresh.set_input("m", "2020", np.array([12.0])) + assert branch.calculate("r", "2020-03").tolist() == [2.0] + assert fresh.calculate("r", "2020-03").tolist() == [2.0] + + +@pytest.mark.parametrize("input_first", [False, True]) +def test_callback_that_changes_its_input_taints_the_simulation_it_returns_to( + input_first, +): + """A branch formula calls a worker whose formula calls back into the branch.""" + + def changes_and_doubles(person, period): + previous = person("source", period) + if np.any(previous != 3): + person.simulation.set_input("source", period, np.full(person.count, 3.0)) + return previous * 2 + + simulations = {} + worker = ( + SimulationBuilder() + .build_default_simulation( + _one_person_system( + _yearly_variable( + "calls_back", + lambda person, period: simulations["branch"].calculate( + "changes_and_doubles", period + ), + ) + ) + ) + .get_branch("worker") + ) + root = SimulationBuilder().build_default_simulation( + _one_person_system( + _yearly_variable("source"), + _yearly_variable("changes_and_doubles", changes_and_doubles), + _yearly_variable( + "result", lambda person, period: worker.calculate("calls_back", period) + ), + ) + ) + root.set_input("source", "2020", np.array([1.0])) + branch = simulations["branch"] = root.get_branch("branch") + if input_first: + branch.set_input("source", "2020", np.array([3.0])) + + assert branch.calculate("result", "2020").tolist() == [6.0] + assert worker.calculate("calls_back", "2020").tolist() == [6.0] diff --git a/tests/core/test_branch_input_invalidation_property.py b/tests/core/test_branch_input_invalidation_property.py new file mode 100644 index 000000000..17372d277 --- /dev/null +++ b/tests/core/test_branch_input_invalidation_property.py @@ -0,0 +1,338 @@ +"""Branch results do not depend on what was calculated before. + +Random sequences of calculations, branches (including reused and forgotten +names), inputs (some setting again the value already read), +``drop_computed_arrays`` and dumps restored as new simulations run on a +synthetic system, with +values held in memory, on disk, not kept, or blacklisted. Every calculation +in a branch must equal the same calculation in a new simulation given the +branch's inputs, its own and those it inherited when it was created, before +calculating anything. The inputs a branch should hold are modelled here, not +read back from the branch. ``test_branch_input_invalidation.py`` pins the +same behaviour with examples. +""" + +from __future__ import annotations + +import os +import shutil +import tempfile + +import numpy as np +import pytest + +# The country-package smoke job installs no dev dependencies. +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from policyengine_core import periods +from policyengine_core.experimental import MemoryConfig +from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, +) +from tests.fixtures.branch_input_invalidation import ( + PEOPLE, + ROOT_INPUTS, + SYNTHETIC_SYSTEM, + YEARS, + synthetic_simulation, +) + +SETTABLE = [ + "p_a", + "p_b", + "p_up", + "p_c", + "p_m", + "p_sum", + "p_prod", + "p_switch", + "p_inner_only", +] +CALCULABLE = [ + "p_sum", + "p_prod", + "p_cond", + "p_lag", + "p_months", + "p_inner", + "p_inner_only", + "p_imported_up", + "p_up", + "p_c", +] +MONTHS = ["2013-01", "2013-07", "2014-12", "2015-03"] +# How the root simulation (and so every branch) stores values. +MODES = { + "memory": dict(), + "disk": dict(memory_config=lambda: MemoryConfig(max_memory_occupation=0)), + "not_kept": dict( + memory_config=lambda: MemoryConfig( + max_memory_occupation=1, variables_to_drop=["p_sum", "p_inner_only"] + ) + ), + # The synthetic system blacklists p_sum and p_inner_only. + "blacklist": dict(opt_out_cache=True), +} + +# Indices pick a simulation modulo how many exist (-1: the newest); small ones +# keep most operations on the root and the first few branches, where they +# interact. A branch is named "x" or "y" (so two lineages can share a name, +# and a forgotten name can be reused) or given a name of its own (None). +simulation_index = st.integers(0, 3) +branch_name = st.sampled_from(["x", "y", None, None]) +values = st.tuples(*[st.integers(0, 100).map(float)] * PEOPLE) +calculate = st.one_of( + st.tuples( + st.just("calculate"), + simulation_index, + st.sampled_from(CALCULABLE), + st.sampled_from(YEARS), + ), + st.tuples( + st.just("calculate"), + simulation_index, + st.just("p_month"), + st.sampled_from(MONTHS), + ), +) +set_input = st.one_of( + st.tuples( + st.just("set"), + simulation_index, + st.sampled_from(SETTABLE), + st.sampled_from(YEARS), + values, + ), + st.tuples( + st.just("set"), + simulation_index, + st.sampled_from(["p_m", "p_month"]), + st.sampled_from(MONTHS), + values, + ), +) +single_operation = st.one_of( + calculate, + calculate, + set_input, + set_input, + st.tuples(st.just("branch"), simulation_index, branch_name), + st.tuples(st.just("forget"), simulation_index), + st.tuples(st.just("drop"), simulation_index), + # Set again, on a branch, an input it reads, with the value it reads. + st.tuples(st.just("reset"), simulation_index, st.integers(0, 50)), + # Dump a simulation (a branch, say) and restore it as a new one. + st.tuples(st.just("dump"), simulation_index), +) +# Pairs of a calculated value and an input it depends on, each reaching it a +# different way: through a formula's own branch, uprating from an earlier +# year (directly, and through a formula's own branch), a month of a year, a +# value the holder may not keep, a lagged year, and a year summed from months. +DEPENDENCIES = [ + ("p_inner", "2013", "p_inner_only", "2013"), + ("p_prod", "2015", "p_up", "2014"), + ("p_imported_up", "2015", "p_up", "2014"), + ("p_month", "2013-07", "p_m", "2013"), + ("p_month", "2015-03", "p_m", "2015"), + ("p_prod", "2013", "p_sum", "2013"), + ("p_lag", "2014", "p_a", "2013"), + ("p_cond", "2014", "p_b", "2014"), + ("p_months", "2014", "p_m", "2014-12"), +] +# Calculate a value, branch, override an input it depends on in the new +# branch (index -1), and calculate the value again there; or branch first and +# calculate the value in the branch before and after overriding the input. +override_after_calculating = st.tuples( + simulation_index, st.sampled_from(DEPENDENCIES), values, st.booleans() +).map( + lambda drawn: ( + [ + ("calculate", drawn[0], drawn[1][0], drawn[1][1]), + ("branch", drawn[0], None), + ("set", -1, drawn[1][2], drawn[1][3], drawn[2]), + ("calculate", -1, drawn[1][0], drawn[1][1]), + ] + if drawn[3] + else [ + ("branch", drawn[0], None), + ("calculate", -1, drawn[1][0], drawn[1][1]), + ("set", -1, drawn[1][2], drawn[1][3], drawn[2]), + ("calculate", -1, drawn[1][0], drawn[1][1]), + ] + ) +) +# Calculate the value in a new branch, set again there the input it depends +# on, with the value the branch reads, and calculate the value again. +reset_after_calculating = st.tuples( + simulation_index, st.sampled_from(DEPENDENCIES) +).map( + lambda drawn: [ + ("branch", drawn[0], None), + ("calculate", -1, drawn[1][0], drawn[1][1]), + ("reset", -1, (drawn[1][2], drawn[1][3])), + ("calculate", -1, drawn[1][0], drawn[1][1]), + ] +) +# The same through a dump: calculate the value, dump that simulation, restore +# it, branch from the restored one, override the input and calculate again. +override_after_restoring = st.tuples( + simulation_index, st.sampled_from(DEPENDENCIES), values +).map( + lambda drawn: [ + ("calculate", drawn[0], drawn[1][0], drawn[1][1]), + ("dump", drawn[0]), + ("branch", -1, None), + ("set", -1, drawn[1][2], drawn[1][3], drawn[2]), + ("calculate", -1, drawn[1][0], drawn[1][1]), + ] +) +operations = st.lists( + st.one_of( + single_operation.map(lambda operation: [operation]), + override_after_calculating, + reset_after_calculating, + override_after_restoring, + ), + min_size=1, + max_size=15, +).map(lambda chunks: [operation for chunk in chunks for operation in chunk]) + + +def _stored_inputs(own_inputs, variable, period, value): + """What ``set_input(variable, period, value)`` stores on a branch, by its own rule. + + Monthly ``p_m`` divides a yearly value between the months the branch has + not set as inputs itself (what is left after the months it has set). + Returns ``None`` when every month is already set (an error unless the + totals match, so the program skips it). + """ + period = periods.period(period) + variable_period = SYNTHETIC_SYSTEM.get_variable(variable).definition_period + if variable_period != periods.MONTH or period.unit == periods.MONTH: + return {(variable, str(period)): tuple(value)} + months = [str(month) for month in period.get_subperiods(periods.MONTH)] + unset = [month for month in months if (variable, month) not in own_inputs] + if not unset: + return None + remaining = np.asarray(value, dtype=float) - sum( + ( + np.asarray(own_inputs[(variable, month)]) + for month in months + if month not in unset + ), + np.zeros(PEOPLE), + ) + return {(variable, month): tuple(remaining / len(unset)) for month in unset} + + +def _run(program, mode): + """Run ``program``; return each calculation with the inputs it should reflect. + + A branch's inputs are those of the simulation it was created from, as + they were then, and those set on it since. A simulation restored from a + dump has the inputs the dumped one had. + """ + options = MODES[mode] + root = synthetic_simulation( + ROOT_INPUTS, + memory_config=options.get("memory_config", lambda: None)(), + opt_out_cache=options.get("opt_out_cache", False), + ) + simulations = [root] + inputs = [dict(ROOT_INPUTS)] + own_inputs = [{}] + children = {} # (parent index, name) -> index, while the parent keeps it + roots = {0} # the root and simulations restored from dumps + results = [] + for operation in program: + kind, index = operation[0], operation[1] % len(simulations) + if operation[1] == -1: + index = len(simulations) - 1 # The newest simulation. + simulation = simulations[index] + if kind == "branch": + name = operation[2] or f"b{len(simulations)}" + if (index, name) in children or name == simulation.branch_name: + continue # get_branch would return an existing simulation. + simulations.append(simulation.get_branch(name)) + inputs.append(dict(inputs[index])) + own_inputs.append({}) + children[(index, name)] = len(simulations) - 1 + elif kind == "forget": + # The parent forgets a branch (formulas delete theirs); the branch + # object keeps working, and its name can be reused. + names = sorted(name for parent, name in children if parent == index) + if names: + del simulation.branches[names[0]] + del children[(index, names[0])] + elif kind == "set": + if index in roots: + continue # Inputs on a root after calculating are out of scope. + _, _, variable, period, value = operation + stored = _stored_inputs(own_inputs[index], variable, period, value) + if stored is None: + continue + simulation.set_input(variable, period, np.asarray(value)) + inputs[index].update(stored) + own_inputs[index].update(stored) + elif kind == "reset": + if index in roots: + continue + # Only inputs stored for the variable's own period unit: a helper + # spreading a year over months stores something else. + keys = sorted( + key + for key in inputs[index] + if periods.period(key[1]).unit + == SYNTHETIC_SYSTEM.get_variable(key[0]).definition_period + ) + selector = operation[2] + if isinstance(selector, tuple): + keys = [selector] if selector in keys else [] + if not keys: + continue + key = keys[selector % len(keys)] if isinstance(selector, int) else keys[0] + simulation.set_input(key[0], key[1], np.asarray(inputs[index][key])) + own_inputs[index][key] = inputs[index][key] + elif kind == "calculate": + _, _, variable, period = operation + value = np.array(simulation.calculate(variable, period), copy=True) + results.append((dict(inputs[index]), variable, period, value)) + elif kind == "dump": + directory = tempfile.mkdtemp(prefix="policyengine-branch-dump-") + try: + dump_simulation(simulation, directory) + restored = restore_simulation(directory, SYNTHETIC_SYSTEM) + finally: + shutil.rmtree(directory, ignore_errors=True) + simulations.append(restored) + inputs.append(dict(inputs[index])) + own_inputs.append({}) + roots.add(len(simulations) - 1) + else: + simulation.drop_computed_arrays() + return results + + +# CI runs a fixed set of programs; set POLICYENGINE_BRANCH_PROPERTY_EXAMPLES +# to explore new ones (e.g. 5000 for a soak). +_SOAK_EXAMPLES = os.environ.get("POLICYENGINE_BRANCH_PROPERTY_EXAMPLES") + + +@hypothesis.settings( + max_examples=int(_SOAK_EXAMPLES or 200), + derandomize=not _SOAK_EXAMPLES, + database=None, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given(program=operations, mode=st.sampled_from(sorted(MODES))) +def test_branch_calculations_match_a_simulation_given_its_inputs_first(program, mode): + for branch_inputs, variable, period, value in _run(program, mode): + expected = synthetic_simulation(branch_inputs).calculate(variable, period) + # Uprating chains, and months divided from a year, may round + # differently in float32 depending on what was calculated first. + np.testing.assert_allclose( + value, expected, rtol=1e-5, atol=1e-4, err_msg=f"{variable} {period}" + ) diff --git a/tests/core/test_branch_scoped_delete_arrays.py b/tests/core/test_branch_scoped_delete_arrays.py index 5db33dd90..8f697c1d1 100644 --- a/tests/core/test_branch_scoped_delete_arrays.py +++ b/tests/core/test_branch_scoped_delete_arrays.py @@ -5,6 +5,8 @@ import gc import os +from pathlib import Path + import numpy as np from policyengine_core.country_template import CountryTaxBenefitSystem @@ -148,7 +150,9 @@ def test_child_disk_delete_keeps_parent_view_intact(tmp_path): np.testing.assert_array_equal(inherited_value, [3_000.0]) assert disk_key in parent_storage._files assert disk_key not in child_storage._files - assert (tmp_path / "salary" / f"{disk_key}.npy").is_file() + # Each store writes its own file; the parent's is still there. + assert Path(parent_storage._files[disk_key]).is_file() + assert Path(parent_storage._files[disk_key]).parent == tmp_path / "salary" np.testing.assert_array_equal( parent_holder._disk_storage.get(PERIOD, "default"), [3_000.0], diff --git a/tests/core/test_branch_shared_arrays.py b/tests/core/test_branch_shared_arrays.py index 61aa398c8..9da7a9334 100644 --- a/tests/core/test_branch_shared_arrays.py +++ b/tests/core/test_branch_shared_arrays.py @@ -128,11 +128,16 @@ def test_branch_copies_only_what_it_reads(tax_benefit_system): branch.calculate("income_tax", JANUARY) still_shared = shared_keys(branch) + held = set(stored_arrays(branch)) assert still_shared < inherited - # ``income_tax`` reads ``salary`` only, so nothing else was copied. + # ``income_tax`` reads ``salary`` only, so nothing else was copied: each + # other array is still shared, or was dropped by the input (values + # calculated after ``salary`` was first stored). for variable in ("rent", "accommodation_size", "housing_tax", "basic_income"): keys = {key for key in inherited if key[0] == variable} - assert keys and keys <= still_shared, variable + assert keys and all(key in still_shared or key not in held for key in keys), ( + variable + ) branch_arrays = stored_arrays(branch) for key in still_shared: assert np.shares_memory(branch_arrays[key], parent_arrays[key]), key @@ -439,11 +444,9 @@ def test_set_input_on_branch_leaves_parent(tax_benefit_system): branch.set_input("salary", JANUARY, np.array([10_000.0, 0.0, 0.0, 0.0])) branch_tax = branch.calculate("income_tax", JANUARY) - # ``income_tax`` was cached before branching, so the branch keeps the - # parent's value, exactly as when the branch held a copy of it. - assert np.array_equal(branch_tax, parent_tax) - branch.delete_arrays("income_tax", JANUARY) - assert branch.calculate("income_tax", JANUARY)[0] > parent_tax[0] + # The input drops the ``income_tax`` the branch inherited, so the branch + # calculates it from the new salary. + assert branch_tax[0] > parent_tax[0] _assert_unchanged(simulation, snapshot) assert np.array_equal(simulation.calculate("income_tax", JANUARY), parent_tax) diff --git a/tests/core/test_on_disk_storage_clones.py b/tests/core/test_on_disk_storage_clones.py index dc245397e..911a7d3e2 100644 --- a/tests/core/test_on_disk_storage_clones.py +++ b/tests/core/test_on_disk_storage_clones.py @@ -228,6 +228,28 @@ def test_a_copied_storage_and_its_source_keep_their_own_values(storage, copier): np.testing.assert_array_equal(storage.get("2012"), [3]) +@COPIERS +@pytest.mark.parametrize("copy_first", [True, False], ids=["copy", "source"]) +def test_a_copied_storage_and_its_source_drop_by_their_own_store_order( + storage, copier, copy_first +): + storage.put(np.array([1.0]), "2012", derived=True) + original_number = storage._sequence_numbers["default_2012"] + copied = copier(storage) + writer, reader = (copied, storage) if copy_first else (storage, copied) + + writer.put(np.array([2.0]), "2012", derived=True) + replacement_number = writer._sequence_numbers["default_2012"] + + assert replacement_number > original_number + assert reader._sequence_numbers["default_2012"] == original_number + assert writer.drop_computed(since=replacement_number) == 1 + assert not reader.has_unnumbered_values() + assert reader.drop_computed(since=replacement_number) == 0 + np.testing.assert_array_equal(reader.get("2012"), [1]) + assert reader.drop_computed(since=original_number) == 1 + + @COPIERS @pytest.mark.parametrize("copy_first", [True, False], ids=["copy", "source"]) def test_a_copied_storage_and_its_source_keep_their_own_values_for_a_new_key( diff --git a/tests/core/test_restore_input_registry_property.py b/tests/core/test_restore_input_registry_property.py index 5905c7453..4d33f3d63 100644 --- a/tests/core/test_restore_input_registry_property.py +++ b/tests/core/test_restore_input_registry_property.py @@ -186,6 +186,10 @@ def _input_keys(simulation): ], requests=[], ) +@hypothesis.example( + steps=[("input", ("eligible", "2012", [False, False]))], + requests=[], +) def test_restored_simulation_keeps_and_drops_what_the_dumped_one_would( system, steps, requests ): @@ -195,8 +199,22 @@ def test_restored_simulation_keeps_and_drops_what_the_dumped_one_would( kind = None original_put = InMemoryStorage.put - def observe_put(storage, value, period, branch_name="default", derived=False): - result = original_put(storage, value, period, branch_name, derived=derived) + def observe_put( + storage, + value, + period, + branch_name="default", + derived=False, + sequence_number=None, + ): + result = original_put( + storage, + value, + period, + branch_name, + derived=derived, + sequence_number=sequence_number, + ) if branch_name != "default" or kind not in ("input", "cache", "calculate"): return result # Observe successful writes to this simulation's leaf storage, not diff --git a/tests/core/test_stale_frame_cache.py b/tests/core/test_stale_frame_cache.py new file mode 100644 index 000000000..b59053725 --- /dev/null +++ b/tests/core/test_stale_frame_cache.py @@ -0,0 +1,768 @@ +"""A stale caller must not disable caching for unrelated nested formulas.""" + +from collections import Counter +import contextvars +import threading + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem, entities +from policyengine_core.simulations import SimulationBuilder +import policyengine_core.simulations.simulation as simulation_module +from policyengine_core.variables import Variable + + +def _variable(name, definition_period, formula=None): + attributes = { + "value_type": float, + "entity": entities.Person, + "definition_period": definition_period, + "label": name, + } + if formula is not None: + attributes["formula"] = formula + return type(name, (Variable,), attributes) + + +def _simulation(definition_period, formulas): + system = CountryTaxBenefitSystem() + system.add_variables( + _variable("sf_source", definition_period), + *[ + _variable(name, definition_period, formula) + for name, formula in formulas.items() + ], + ) + simulation = SimulationBuilder().build_default_simulation(system) + input_periods = ( + ["2020"] + if definition_period == periods.YEAR + else [f"2020-{month:02d}" for month in range(1, 13)] + ) + for input_period in input_periods: + simulation.set_input("sf_source", input_period, np.array([1.0])) + return simulation + + +def _in_thread(call): + returned = [] + failures = [] + + def target(): + try: + returned.append(call()) + except BaseException as error: + failures.append(error) + + context = contextvars.Context() + thread = threading.Thread(target=context.run, args=(target,)) + thread.start() + thread.join() + if failures: + raise failures[0] + return returned[0] + + +@pytest.mark.parametrize("definition_period", [periods.YEAR, periods.MONTH]) +def test_unrelated_suffix_caches_while_its_stale_caller_retries(definition_period): + calls = Counter() + + def doubled(person, period): + return person("sf_source", period) * 2 + + def reader(person, period): + calls["reader"] += 1 + branch = person.simulation.get_branch("worker") + earlier = branch.calculate("sf_doubled", period) + if np.any(earlier != 6): + branch.set_input("sf_source", period, np.array([3.0])) + return earlier + + def independent(person, period): + calls["independent"] += 1 + return person("sf_source", period) + 10 + + def outer(person, period): + calls["outer"] += 1 + earlier = person("sf_reader", period) + # The reader made the caller stale. This value only uses the + # caller's own unchanged input, so it can still be kept. + return earlier + person("sf_independent", period) + + simulation = _simulation( + definition_period, + { + "sf_doubled": doubled, + "sf_reader": reader, + "sf_independent": independent, + "sf_outer": outer, + }, + ) + count = 1 if definition_period == periods.YEAR else 12 + calculate = ( + simulation.calculate + if definition_period == periods.YEAR + else simulation.calculate_add + ) + result = calculate("sf_outer", "2020") + # A new calculation with the final branch input first returns 6 + 11 + # per period. The dependent reader and caller must both run again. + np.testing.assert_array_equal(result, [count * 17.0]) + assert calls == {"outer": count * 2, "reader": count * 2, "independent": count} + np.testing.assert_array_equal(simulation.get_array("sf_outer", "2020"), result) + + completed = calls.copy() + np.testing.assert_array_equal(calculate("sf_outer", "2020"), result) + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), result) + assert calls == completed + + +def test_ordered_local_credits_cache_shared_dependencies_during_a_stale_attempt(): + calls = Counter() + + def credit_formula(index): + def formula(person, period): + calls[f"credit_{index}"] += 1 + return person("sf_source", period) + sum( + person(f"sf_credit_{previous}", period) for previous in range(index) + ) + + return formula + + def outer(person, period): + calls["outer"] += 1 + worker = person.simulation.get_branch("worker") + previous = worker.calculate("sf_source", period) + if np.any(previous != 3): + worker.set_input("sf_source", period, np.array([3.0])) + # Every credit reads all its predecessors. Reusing those local + # values keeps this dependency graph from expanding on stale retries. + return previous + person("sf_credit_5", period) + + simulation = _simulation( + periods.YEAR, + { + **{f"sf_credit_{index}": credit_formula(index) for index in range(6)}, + "sf_outer": outer, + }, + ) + result = simulation.calculate("sf_outer", "2020") + np.testing.assert_array_equal(result, [35.0]) + assert calls == {"outer": 2, **{f"credit_{index}": 1 for index in range(6)}} + np.testing.assert_array_equal(simulation.get_array("sf_outer", "2020"), result) + completed = calls.copy() + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), result) + np.testing.assert_array_equal(simulation.calculate("sf_credit_5", "2020"), [32.0]) + assert calls == completed + + +@pytest.mark.parametrize("depth", [1, 6, 12, 20]) +@pytest.mark.parametrize("mediated_foreign_read", [False, True]) +@pytest.mark.parametrize("replace_after_credits", [False, True]) +def test_foreign_read_credits_take_linear_work_during_a_stale_attempt( + depth, + mediated_foreign_read, + replace_after_credits, +): + calls = Counter() + + def foreign_source(person, period): + calls["foreign_source"] += 1 + return person.simulation.get_branch("worker").calculate("sf_source", period) + + def credit_formula(index): + def formula(person, period): + calls[f"credit_{index}"] += 1 + source = ( + person("sf_foreign_source", period) + if mediated_foreign_read + else person.simulation.get_branch("worker").calculate( + "sf_source", period + ) + ) + return source + sum( + person(f"sf_credit_{previous}", period) for previous in range(index) + ) + + return formula + + def outer(person, period): + calls["outer"] += 1 + worker = person.simulation.get_branch("worker") + previous = worker.calculate("sf_source", period) + first = np.any(previous == 1) + if first: + worker.set_input( + "sf_source", + period, + np.array([2.0 if replace_after_credits else 3.0]), + ) + # Every credit reads all its predecessors. Refusing all foreign-read + # reuse in this stale attempt expands depth n into 2**(n - 1) calls. + result = person(f"sf_credit_{depth - 1}", period) + if first and replace_after_credits: + np.testing.assert_array_equal(result, [2.0 * 2 ** (depth - 1)]) + worker.set_input("sf_source", period, np.array([3.0])) + # Reused values must retain their foreign dependencies, including + # dependencies reached through a locally calculated source. + result = person(f"sf_credit_{depth - 1}", period) + np.testing.assert_array_equal(result, [3.0 * 2 ** (depth - 1)]) + return result + + simulation = _simulation( + periods.YEAR, + { + "sf_foreign_source": foreign_source, + **{f"sf_credit_{index}": credit_formula(index) for index in range(depth)}, + "sf_outer": outer, + }, + ) + expected = [3.0 * 2 ** (depth - 1)] + result = simulation.calculate("sf_outer", "2020") + np.testing.assert_array_equal(result, expected) + assert calls["outer"] == 2 + # At most one evaluation per credit and settled input state in each + # attempt: the stale graph, an optional changed graph, and the clean retry. + graphs = 3 if replace_after_credits else 2 + credit_calls = sum(calls[f"credit_{index}"] for index in range(depth)) + assert credit_calls <= graphs * depth + assert max(calls[f"credit_{index}"] for index in range(depth)) <= graphs + assert calls["foreign_source"] <= (graphs if mediated_foreign_read else 0) + np.testing.assert_array_equal(simulation.get_array("sf_outer", "2020"), result) + + completed = calls.copy() + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), expected) + np.testing.assert_array_equal( + simulation.calculate(f"sf_credit_{depth - 1}", "2020"), expected + ) + assert calls == completed + + +@pytest.mark.parametrize("mutation_in_thread", [False, True]) +@pytest.mark.parametrize("raw_cache_write", [False, True]) +def test_stale_attempt_reuse_observes_raw_holder_changes_between_siblings( + mutation_in_thread, + raw_cache_write, +): + calls = Counter() + expected = [4.0 if raw_cache_write else 1.0] + + def suffix(person, period): + calls["suffix"] += 1 + return person.simulation.get_branch("worker").calculate("sf_source", period) + + def outer(person, period): + calls["outer"] += 1 + worker = person.simulation.get_branch("worker") + worker.calculate("sf_source", period) + worker.set_input("sf_source", period, np.array([2.0])) + np.testing.assert_array_equal(person("sf_suffix", period), [2.0]) + counters = worker._input_epoch, worker._inputs_set + + def change(): + holder = worker.get_holder("sf_source") + if raw_cache_write: + holder.put_in_cache(np.array([4.0]), period, worker.branch_name) + else: + # Removing only the worker override exposes its inherited 1. + holder.delete_arrays(period, worker.branch_name) + + if mutation_in_thread: + # The memo-producing child has exited. Context-free activity + # must invalidate its reuse even with no restricted frame open. + assert not simulation_module._restricted_frames + _in_thread(change) + else: + change() + assert (worker._input_epoch, worker._inputs_set) == counters + result = person("sf_suffix", period) + np.testing.assert_array_equal(result, expected) + return result + + simulation = _simulation(periods.YEAR, {"sf_suffix": suffix, "sf_outer": outer}) + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), expected) + # Context-free effects can exhaust the retry budget; every attempt must + # nevertheless recalculate the suffix after the counter-preserving change. + assert 2 <= calls["outer"] <= simulation_module._RERUNS_AFTER_INPUT_CHANGE + 1 + assert calls["suffix"] == 2 * calls["outer"] + + +def test_stale_attempt_reuse_revalidates_foreign_branch_registration(): + calls = Counter() + replacement = None + + def suffix(person, period): + calls["suffix"] += 1 + return person.simulation.get_branch("worker").calculate("sf_source", period) + + def outer(person, period): + calls["outer"] += 1 + trigger = person.simulation.get_branch("trigger") + previous = trigger.calculate("sf_source", period) + first = np.any(previous == 1) + if first: + trigger.set_input("sf_source", period, np.array([3.0])) + np.testing.assert_array_equal(person("sf_suffix", period), [1.0]) + original = person.simulation.get_branch("worker") + assert (original._input_epoch, original._inputs_set) == ( + replacement._input_epoch, + replacement._inputs_set, + ) + # Replacing a registration does not change the old snapshot's + # counters. A memoized named read must still follow the new branch. + person.simulation.branches["worker"] = replacement + result = person("sf_suffix", period) + np.testing.assert_array_equal(result, [4.0]) + return result + + simulation = _simulation(periods.YEAR, {"sf_suffix": suffix, "sf_outer": outer}) + original = simulation.get_branch("worker") + replacement = original.clone() + replacement.get_holder("sf_source").put_in_cache( + np.array([4.0]), periods.period("2020"), replacement.branch_name + ) + assert (original._input_epoch, original._inputs_set) == ( + replacement._input_epoch, + replacement._inputs_set, + ) + + result = simulation.calculate("sf_outer", "2020") + np.testing.assert_array_equal(result, [4.0]) + assert calls == {"outer": 2, "suffix": 3} + np.testing.assert_array_equal(simulation.get_array("sf_outer", "2020"), result) + completed = calls.copy() + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), result) + assert calls == completed + + +@pytest.mark.parametrize("registration_depth", [1, 2]) +@pytest.mark.parametrize("prime_fast_cache", [False, True]) +@pytest.mark.parametrize("read_route", ["direct", "memo", "foreign_formula"]) +@pytest.mark.parametrize("aliased_registration", [False, True]) +def test_registration_replaced_inside_memo_producer_preserves_the_next_read( + registration_depth, + prime_fast_cache, + read_route, + aliased_registration, +): + calls = Counter() + suffix_reads = [] + replacement = None + + def worker(simulation): + branch = simulation.get_branch("worker") + return branch.get_branch("leaf") if registration_depth == 2 else branch + + def replace_registration(): + simulation.branches["worker"] = replacement + + def foreign_value(person, period): + result = person("sf_source", period) + # Replacement inside the foreign calculation must retain the + # registration that was present when that calculation began. + replace_registration() + return result + + def reader(person, period): + calls["reader"] += 1 + return worker(person.simulation).calculate("sf_source", period) + + def suffix(person, period): + calls["suffix"] += 1 + if read_route == "memo": + # The outer formula already memoized this reader. The fresh + # suffix frame must receive its registration dependencies too. + result = person("sf_reader", period) + else: + name = ( + "sf_foreign_value" if read_route == "foreign_formula" else "sf_source" + ) + result = worker(person.simulation).calculate(name, period) + suffix_reads.append(float(result[0])) + if read_route != "foreign_formula": + replace_registration() + return result + + def outer(person, period): + calls["outer"] += 1 + trigger = person.simulation.get_branch("trigger") + if np.any(trigger.calculate("sf_source", period) == 1): + trigger.set_input("sf_source", period, np.array([3.0])) + if read_route == "memo": + person("sf_reader", period) + person("sf_suffix", period) + result = person("sf_suffix", period) + # A wrong memoized value becomes a supplied input and survives + # the outer retry, so checking only the final retry is insufficient. + sink.set_input("sf_source", period, result) + return sink.calculate("sf_source", period) + + simulation = _simulation( + periods.YEAR, + { + "sf_foreign_value": foreign_value, + "sf_reader": reader, + "sf_suffix": suffix, + "sf_outer": outer, + }, + ) + original = simulation.get_branch("other" if aliased_registration else "worker") + if aliased_registration: + # The branch object's own registration remains valid. Only the + # actual name resolved by worker() is replaced during the formula. + simulation.branches["worker"] = original + original_worker = worker(simulation) + replacement = original.clone() + replacement_worker = ( + replacement.get_branch("leaf") if registration_depth == 2 else replacement + ) + replacement_worker.get_holder("sf_source").put_in_cache( + np.array([4.0]), periods.period("2020"), replacement_worker.branch_name + ) + assert (original_worker._input_epoch, original_worker._inputs_set) == ( + replacement_worker._input_epoch, + replacement_worker._inputs_set, + ) + if prime_fast_cache: + original_worker.calculate("sf_source", "2020") + sink = simulation.get_branch("sink") + + result = simulation.calculate("sf_outer", "2020") + np.testing.assert_array_equal(result, [4.0]) + assert suffix_reads == [1.0, 4.0] + assert calls["outer"] == 2 + assert calls["suffix"] == 2 + if aliased_registration: + assert simulation.branches["other"] is original + if read_route == "memo": + assert calls["reader"] == 2 + np.testing.assert_array_equal(sink.calculate("sf_source", "2020"), result) + np.testing.assert_array_equal(simulation.get_array("sf_outer", "2020"), result) + np.testing.assert_array_equal(original_worker.calculate("sf_source", "2020"), [1.0]) + completed = calls.copy() + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), result) + assert calls == completed + + +def test_stale_attempt_reuse_keeps_direct_foreign_calculations_in_their_simulation(): + calls = Counter() + + def value(person, period): + calls[f"value_{person.simulation.branch_name}"] += 1 + return person("sf_source", period) * 10 + + def suffix(person, period): + calls["suffix"] += 1 + # Custom input handlers can call _calculate without opening a frame + # in that simulation. Its result still belongs to the foreign worker. + return person.simulation.get_branch("worker")._calculate("sf_value", period) + + def outer(person, period): + calls["outer"] += 1 + worker = person.simulation.get_branch("worker") + previous = worker.calculate("sf_source", period) + if np.any(previous != 3): + worker.set_input("sf_source", period, np.array([3.0])) + foreign = person("sf_suffix", period) + local = person("sf_value", period) + # Assert inside the stale attempt so its later retry cannot hide a + # worker result incorrectly memoized under the root's variable name. + np.testing.assert_array_equal(foreign, [30.0]) + np.testing.assert_array_equal(local, [10.0]) + return foreign + local + + simulation = _simulation( + periods.YEAR, + {"sf_value": value, "sf_suffix": suffix, "sf_outer": outer}, + ) + simulation.get_branch("worker") + result = simulation.calculate("sf_outer", "2020") + np.testing.assert_array_equal(result, [40.0]) + assert calls == {"outer": 2, "suffix": 2, "value_worker": 2, "value_default": 1} + completed = calls.copy() + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), result) + assert calls == completed + + +@pytest.mark.parametrize("suffix_branch", ["worker", "other"]) +def test_foreign_suffix_is_not_cached_before_its_stale_caller_settles(suffix_branch): + calls = Counter() + + def suffix(person, period): + calls["suffix"] += 1 + return person.simulation.get_branch(suffix_branch).calculate( + "sf_source", period + ) + + def outer(person, period): + calls["outer"] += 1 + worker = person.simulation.get_branch("worker") + previous = worker.calculate("sf_source", period) + first = np.any(previous == 1) + if first: + worker.set_input("sf_source", period, np.array([2.0])) + # A new frame for this suffix is locally clean, but its caller + # already read the replaced worker input. Keeping the suffix now + # would hide the later mutation from the caller's next attempt. + result = person("sf_suffix", period) + if first: + person.simulation.get_branch(suffix_branch).set_input( + "sf_source", period, np.array([3.0]) + ) + worker.set_input("sf_source", period, np.array([3.0])) + return result + + simulation = _simulation(periods.YEAR, {"sf_suffix": suffix, "sf_outer": outer}) + result = simulation.calculate("sf_outer", "2020") + np.testing.assert_array_equal(result, [3.0]) + assert calls == {"outer": 2, "suffix": 2} + np.testing.assert_array_equal(simulation.get_array("sf_outer", "2020"), result) + + completed = calls.copy() + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), result) + assert calls == completed + + +@pytest.mark.parametrize("write_in_thread", [False, True]) +@pytest.mark.parametrize("worker_family", ["own", "unrelated"]) +def test_foreign_writing_suffix_repeats_its_effects_while_independent_work_caches( + write_in_thread, + worker_family, +): + calls = Counter() + unrelated_worker = ( + _simulation(periods.YEAR, {}).get_branch("worker") + if worker_family == "unrelated" + else None + ) + + def get_worker(simulation): + return ( + unrelated_worker + if unrelated_worker is not None + else simulation.get_branch("worker") + ) + + def independent(person, period): + calls["independent"] += 1 + return person("sf_source", period) + 10 + + def suffix(person, period): + calls["suffix"] += 1 + worker = get_worker(person.simulation) + + def change(): + worker.set_input("sf_source", period, np.array([3.0])) + + if write_in_thread: + _in_thread(change) + else: + change() + # This formula writes a foreign input without reading a foreign + # value. Its effect must still repeat on the caller's retry, while + # the independent calculation beneath it can safely stay cached. + return person("sf_independent", period) + + def outer(person, period): + calls["outer"] += 1 + worker = get_worker(person.simulation) + previous = worker.calculate("sf_source", period) + worker.set_input("sf_source", period, np.array([2.0])) + return previous + person("sf_suffix", period) + + simulation = _simulation( + periods.YEAR, + { + "sf_independent": independent, + "sf_suffix": suffix, + "sf_outer": outer, + }, + ) + attempts = ( + simulation_module._RERUNS_AFTER_INPUT_CHANGE + 1 if write_in_thread else 3 + ) + result = simulation.calculate("sf_outer", "2020") + np.testing.assert_array_equal(result, [14.0]) + assert calls == {"outer": attempts, "suffix": attempts, "independent": 1} + np.testing.assert_array_equal( + get_worker(simulation).calculate("sf_source", "2020"), [3.0] + ) + if write_in_thread: + # A context-free write cannot certify convergence even though each + # transition ends at the same input. The original shared-frame rule + # exhausts the budget and leaves the result unkept in this case. + assert simulation.get_array("sf_outer", "2020") is None + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), result) + assert calls == { + "outer": attempts * 2, + "suffix": attempts * 2, + "independent": 1, + } + else: + # The first attempt returns the old input; the next two repeat the + # same complete 2 -> 3 transition and result and prove a fixed point. + np.testing.assert_array_equal(simulation.get_array("sf_outer", "2020"), result) + completed = calls.copy() + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), result) + assert calls == completed + + +def test_branch_creating_suffix_recreates_registration_while_independent_work_caches(): + def run(input_first): + calls = Counter() + created = [] + + def independent(person, period): + calls["independent"] += 1 + return person("sf_source", period) + 10 + + def suffix(person, period): + calls["suffix"] += 1 + # Creating a branch is the only foreign effect here: no + # foreign value is read and no input is changed by this suffix. + created.append(person.simulation.get_branch("scratch")) + return person("sf_independent", period) + + def outer(person, period): + calls["outer"] += 1 + worker = person.simulation.get_branch("worker") + previous = worker.calculate("sf_source", period) + worker.set_input("sf_source", period, np.array([3.0])) + result = previous + person("sf_suffix", period) + # Each attempt requires the suffix's creation, even after an + # earlier attempt removed its registration. A cached suffix + # would omit that effect and leave no branch to remove here. + assert person.simulation.branches.pop("scratch") is created[-1] + return result + + simulation = _simulation( + periods.YEAR, + { + "sf_independent": independent, + "sf_suffix": suffix, + "sf_outer": outer, + }, + ) + if input_first: + simulation.get_branch("worker").set_input( + "sf_source", "2020", np.array([3.0]) + ) + result = simulation.calculate("sf_outer", "2020") + attempts = 1 if input_first else 2 + assert calls == {"outer": attempts, "suffix": attempts, "independent": 1} + assert len(created) == attempts + if not input_first: + assert created[0] is not created[1] + np.testing.assert_array_equal(simulation.get_array("sf_outer", "2020"), result) + completed = calls.copy() + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), result) + assert calls == completed + return result + + np.testing.assert_array_equal(run(input_first=False), run(input_first=True)) + + +@pytest.mark.parametrize("prime_fast_cache", [False, True]) +def test_unrelated_thread_read_suffix_is_read_again_after_its_caller_changes_input( + prime_fast_cache, +): + calls = Counter() + unrelated_worker = _simulation( + periods.YEAR, + {"sf_doubled": lambda person, period: person("sf_source", period) * 2}, + ).get_branch("worker") + if prime_fast_cache: + unrelated_worker.calculate("sf_doubled", "2020") + + def suffix(person, period): + calls["suffix"] += 1 + # A context-free read from another family cannot reach this frame + # through ancestry, so it must still refuse the optional cache. + return _in_thread(lambda: unrelated_worker.calculate("sf_doubled", period)) + + def outer(person, period): + calls["outer"] += 1 + worker = person.simulation.get_branch("worker") + previous = worker.calculate("sf_source", period) + first = np.any(previous == 1) + if first: + worker.set_input("sf_source", period, np.array([2.0])) + result = person("sf_suffix", period) + if first: + unrelated_worker.set_input("sf_source", period, np.array([3.0])) + worker.set_input("sf_source", period, np.array([3.0])) + return result + + simulation = _simulation(periods.YEAR, {"sf_suffix": suffix, "sf_outer": outer}) + result = simulation.calculate("sf_outer", "2020") + np.testing.assert_array_equal(result, [6.0]) + assert calls == {"outer": 2, "suffix": 2} + np.testing.assert_array_equal(simulation.get_array("sf_outer", "2020"), result) + completed = calls.copy() + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), result) + assert calls == completed + + +@pytest.mark.parametrize("raw_cache_write", [False, True]) +def test_unrelated_unread_thread_mutation_does_not_prevent_caller_convergence( + raw_cache_write, +): + calls = Counter() + unrelated_worker = _simulation(periods.YEAR, {}).get_branch("worker") + + def independent(person, period): + calls["independent"] += 1 + return person("sf_source", period) + 10 + + def suffix(person, period): + calls["suffix"] += 1 + + def change(): + if raw_cache_write: + unrelated_worker.get_holder("sf_source").put_in_cache( + np.array([5.0]), period, unrelated_worker.branch_name + ) + else: + unrelated_worker.delete_arrays("sf_source") + unrelated_worker.set_input("sf_source", period, np.array([5.0])) + + _in_thread(change) + return person("sf_independent", period) + + def outer(person, period): + calls["outer"] += 1 + worker = person.simulation.get_branch("worker") + previous = worker.calculate("sf_source", period) + worker.set_input("sf_source", period, np.array([2.0])) + result = person("sf_suffix", period) + worker.set_input("sf_source", period, np.array([3.0])) + return previous + result + + simulation = _simulation( + periods.YEAR, + { + "sf_independent": independent, + "sf_suffix": suffix, + "sf_outer": outer, + }, + ) + result = simulation.calculate("sf_outer", "2020") + np.testing.assert_array_equal(result, [14.0]) + # The optional suffix cache cannot skip its threaded effects, but the + # caller's observed 2 -> 3 transition still settles on the third attempt. + # Unread changes in another family do not taint that convergence proof. + assert calls == {"outer": 3, "suffix": 3, "independent": 1} + if raw_cache_write: + # A raw cache write is observable activity, but is still excluded + # from the record of supplied inputs used for replay and export. + assert ( + "sf_source", + unrelated_worker.branch_name, + periods.period("2020"), + ) not in unrelated_worker._user_input_keys + np.testing.assert_array_equal(simulation.get_array("sf_outer", "2020"), result) + completed = calls.copy() + np.testing.assert_array_equal(simulation.calculate("sf_outer", "2020"), result) + assert calls == completed diff --git a/tests/fixtures/branch_input_invalidation.py b/tests/fixtures/branch_input_invalidation.py new file mode 100644 index 000000000..506adac75 --- /dev/null +++ b/tests/fixtures/branch_input_invalidation.py @@ -0,0 +1,291 @@ +"""A synthetic system for derived values and branches inside formulas.""" + +import atexit +import shutil +import tempfile +import weakref + +import numpy as np + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem, entities +from policyengine_core.holders import set_input_divide_by_period +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables import Variable + + +def _yearly(name, formula=None, **attributes): + namespace = dict( + value_type=float, + entity=entities.Person, + definition_period=periods.YEAR, + label=name, + **attributes, + ) + if formula is not None: + namespace["formula"] = formula + return type(name, (Variable,), namespace) + + +def _inner_branch_formula(person, period): + """Calculate a value in a branch that exists only while this formula runs.""" + simulation = person.simulation + inner = simulation.get_branch("inner") + try: + inner.set_input("p_switch", period, np.ones(person.count)) + inner_value = inner.calculate("p_inner_only", period) + finally: + del simulation.branches["inner"] + return inner_value + person("p_sum", period) + + +def _imported_uprating_formula(person, period): + """An uprated value, calculated in a branch that exists only while this runs.""" + simulation = person.simulation + temporary = simulation.get_branch("temporary_uprating") + try: + return temporary.calculate("p_up", period) * 2 + finally: + del simulation.branches["temporary_uprating"] + + +def _monthly_formula(person, period): + return ( + person("p_a", period) + + person("p_b", period.this_year) / 12 + + person("p_m", period) + ) + + +def _monthly(name, formula=None, **attributes): + namespace = dict( + value_type=float, + entity=entities.Person, + definition_period=periods.MONTH, + label=name, + **attributes, + ) + if formula is not None: + namespace["formula"] = formula + return type(name, (Variable,), namespace) + + +SYNTHETIC_VARIABLES = [ + # Inputs. p_up is uprated by a parameter that changes every year from + # 2012 to 2015; p_c is given for 2012 only, and carried over where the + # system carries inputs over. + _yearly("p_a"), + _yearly("p_b"), + _yearly("p_up", uprating="taxes.income_tax_rate"), + _yearly("p_c"), + _yearly( + "p_sum", + lambda person, period: person("p_a", period) + 2 * person("p_b", period), + ), + _yearly( + "p_prod", + lambda person, period: ( + person("p_sum", period) * (1 + person("p_up", period) / 100) + ), + ), + _yearly( + "p_cond", + lambda person, period: np.where( + person("p_sum", period) > 60, + person("p_prod", period), + person("p_c", period), + ), + ), + _yearly( + "p_lag", + lambda person, period: ( + person("p_sum", period.last_year) + person("p_a", period) + ), + ), + # A monthly input; an input for a year is divided between its months. + _monthly("p_m", set_input=set_input_divide_by_period), + _monthly("p_month", _monthly_formula), + _yearly("p_months", lambda person, period: person("p_month", period)), + _yearly("p_switch", lambda person, period: np.zeros(person.count)), + _yearly( + "p_inner_only", + lambda person, period: 3 * person("p_a", period) + person("p_switch", period), + ), + _yearly("p_inner", _inner_branch_formula), + _yearly("p_imported_up", _imported_uprating_formula), + # Reads itself a year earlier, back until core's spiral detection gives + # up and returns the default. + _yearly( + "p_spiral", lambda person, period: person("p_spiral", period.last_year) + 1 + ), +] + + +def _synthetic_system(carry_over): + class System(CountryTaxBenefitSystem): + auto_carry_over_input_variables = carry_over + + system = System() + system.add_variables(*SYNTHETIC_VARIABLES) + # Not stored by simulations that opt out of the cache (``opt_out_cache``). + system.cache_blacklist = {"p_sum", "p_inner_only"} + return system + + +# Carry-over is left out of the property test's random programs: core's +# carry-over itself depends on calculation order (a later period calculated +# first stops an earlier input being carried into the years between), in any +# simulation. +SYNTHETIC_SYSTEM = _synthetic_system(carry_over=False) +CARRY_OVER_SYSTEM = _synthetic_system(carry_over=True) + +PEOPLE = 3 +YEARS = ["2012", "2013", "2014", "2015"] + + +_RUN_DIRECTORY = None + + +def _run_directory(): + """A directory for this process's simulation directories, removed at exit. + + A run killed before it can remove them leaves this one directory behind, + not one for each simulation still alive. + """ + global _RUN_DIRECTORY + if _RUN_DIRECTORY is None: + _RUN_DIRECTORY = tempfile.mkdtemp(prefix="policyengine-branch-tests-") + atexit.register(shutil.rmtree, _RUN_DIRECTORY, ignore_errors=True) + return _RUN_DIRECTORY + + +def synthetic_simulation( + inputs, memory_config=None, system=SYNTHETIC_SYSTEM, opt_out_cache=False +): + """A simulation of ``PEOPLE`` people given ``inputs`` before anything else.""" + simulation = SimulationBuilder().build_default_simulation(system, count=PEOPLE) + simulation.opt_out_cache = opt_out_cache + if memory_config is not None: + # A directory of its own, removed once the simulation is + # garbage-collected, or at interpreter exit if it is still alive then + # (a system keeps the last simulation built with it). Its branches + # keep it alive (``parent_branch``), so the directory lasts while any + # of them can read it. Not a pytest ``tmp_path``: the property test + # builds simulations in each of its examples. + directory = tempfile.mkdtemp(dir=_run_directory()) + weakref.finalize(simulation, shutil.rmtree, directory, ignore_errors=True) + simulation._data_storage_dir = directory + simulation.memory_config = memory_config + # Holders read the memory configuration when created; nothing is + # stored yet, so create them again. Create them all here, so every + # branch's holders are copies of these: a holder first created in a + # branch would own (and remove on deletion) a directory the whole + # family's disk storage shares. + for population in simulation.populations.values(): + population._holders = {} + for variable in system.variables: + disk_storage = simulation.get_holder(variable)._disk_storage + # The directory is removed as a whole, above. A disk storage + # removing its own subdirectory too, when collected with the + # simulation, could find it already gone. + if disk_storage is not None: + disk_storage.preserve_storage_dir = True + for (variable, period), values in sorted( + inputs.items(), + key=lambda item: periods.key_period_size(periods.period(item[0][1])), + ): + simulation.set_input(variable, period, np.asarray(values, dtype=float)) + return simulation + + +ROOT_INPUTS = { + **{("p_a", year): (10.0 * i, 20.0, 35.0 + i) for i, year in enumerate(YEARS)}, + **{("p_b", year): (5.0, 15.0 + i, 25.0) for i, year in enumerate(YEARS)}, + ("p_up", "2012"): (100.0, 200.0, 300.0), + ("p_c", "2012"): (7.0, 8.0, 9.0), + **{ + ("p_m", f"{year}-{month:02d}"): (1.0 * month, 2.0, 0.5 * month) + for year in (2013, 2014) + for month in range(1, 13) + }, +} + + +# ----- A system whose formulas choose between branches, as itemization does ----- # + +# How many times each counted formula ran. +FORMULA_RUNS = {"agi": 0} + + +def _agi(person, period): + FORMULA_RUNS["agi"] += 1 + return person("earn", period) * 1.0 + + +def _tax_if(choice): + """Tax with ``choose`` set to ``choice``, from a branch deleted afterwards.""" + + def formula(person, period): + simulation = person.simulation + name = f"choose_{choice}" + branch = simulation.get_branch(name) + try: + branch.set_input("choose", period, np.full(person.count, float(choice))) + return branch.calculate("tax", period) + finally: + del simulation.branches[name] + + return formula + + +def _marginal_rate(person, period): + """Net income with one more unit of earnings, from a branch, as MTRs are.""" + simulation = person.simulation + branch = simulation.get_branch("raise") + try: + branch.set_input("earn", period, person("earn", period) + 1) + return branch.calculate("net", period) - person("net", period) + finally: + del simulation.branches["raise"] + + +def _from_persistent_branch(person, period): + """Read a value from a branch kept between calls.""" + return person.simulation.get_branch("persistent").calculate("tax", period) + + +BRANCHING_VARIABLES = [ + _yearly("earn"), + _yearly("agi", _agi), + _yearly( + "tax", + lambda person, period: ( + person("agi", period) * (0.2 - 0.05 * person("choose", period)) + ), + ), + _yearly("tax_if_chosen", _tax_if(1)), + _yearly("tax_if_not_chosen", _tax_if(0)), + _yearly( + "choose", + lambda person, period: ( + (person("tax_if_chosen", period) < person("tax_if_not_chosen", period)) + * 1.0 + ), + ), + _yearly( + "net", lambda person, period: person("agi", period) - person("tax", period) + ), + _yearly("marginal_rate", _marginal_rate), + _yearly("from_persistent_branch", _from_persistent_branch), +] + +BRANCHING_SYSTEM = CountryTaxBenefitSystem() +BRANCHING_SYSTEM.add_variables(*BRANCHING_VARIABLES) + + +def branching_simulation(earn=(10.0, 20.0, 30.0), period="2020"): + simulation = SimulationBuilder().build_default_simulation( + BRANCHING_SYSTEM, count=len(earn) + ) + simulation.set_input("earn", period, np.asarray(earn)) + return simulation