From 11b3fbb68f82404ded0fc130902f1ba6946df0cb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 1 Oct 2026 14:59:02 -0400 Subject: [PATCH 01/42] Drop what a branch input may have fed; add drop_computed_arrays A branch starts with every value its parent holds, and set_input on a branch used to keep everything calculated from the value it replaced, so a value the parent calculated before branching answered for the branch whatever its inputs said. Every stored array now carries a sequence number from one process-wide counter (store order), and storages record which values are inputs; both are copied with clones. A simulation family shares a StoreHistory of the first store of each variable and period (including values a holder does not keep) and of the first uprated or carried-over value of each variable. set_input on a branch drops the branch's non-input values stored at or after the earliest of those for an overlapping period; when there is none, nothing is dropped. Simulation.drop_computed_arrays() drops every non-input value, for branches whose policy changes; _invalidate_all_caches now uses the same input flags. Co-Authored-By: Claude Opus 5.5 --- .../branch-set-input-invalidation.added.md | 1 + .../branch-set-input-invalidation.fixed.md | 1 + docs/_toc.yml | 1 + docs/usage/branches.md | 93 +++ .../data_storage/in_memory_storage.py | 62 +- .../data_storage/on_disk_storage.py | 57 +- .../data_storage/store_history.py | 88 ++ policyengine_core/holders/holder.py | 48 +- policyengine_core/simulations/simulation.py | 195 +++-- pyproject.toml | 1 + tests/core/test_branch_input_invalidation.py | 785 ++++++++++++++++++ uv.lock | 96 ++- 12 files changed, 1325 insertions(+), 103 deletions(-) create mode 100644 changelog.d/branch-set-input-invalidation.added.md create mode 100644 changelog.d/branch-set-input-invalidation.fixed.md create mode 100644 docs/usage/branches.md create mode 100644 policyengine_core/data_storage/store_history.py create mode 100644 tests/core/test_branch_input_invalidation.py 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.fixed.md b/changelog.d/branch-set-input-invalidation.fixed.md new file mode 100644 index 000000000..b7916d200 --- /dev/null +++ b/changelog.d/branch-set-input-invalidation.fixed.md @@ -0,0 +1 @@ +Simulation.set_input on a branch now drops the values the branch holds that may have been calculated from the value it replaces (tracked with a sequence number on every stored array), so a value the parent calculated before branching no longer shadows the branch's input. 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..6bfc095a3 --- /dev/null +++ b/docs/usage/branches.md @@ -0,0 +1,93 @@ +# 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 + +`get_branch` returns a branch holding every value its parent holds at that +moment: inputs and calculated values alike. What the parent stores after that +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. +A simulation and all its branches also share a record of the first time each +variable was stored (or calculated without being kept) for each period, and of +the first time one of its values was uprated or carried over from another +period. 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 of: + +- the first time the variable was stored for any period that shares a day with + `period`, anywhere in the family; and +- the first time a value of the variable was uprated or carried over. + +If neither has happened, the branch drops nothing. That is the usual case: a +formula creates the branch while it is still calculating the variable the +branch overrides, so nothing has stored that variable yet. + +### Why this is sound + +A formula's result is stored after everything the formula read has been +stored, so a calculated value carries a larger number than each stored value it +was calculated from, directly or through other calculated values. A value that +depends on the overridden variable at an overlapping period was therefore +stored after the first such value existed. A value derived from another period +of the same variable (uprating, carry-over) depends on which periods hold +values at all, which is why the first such derivation also counts. The record +covers the whole family because a formula can calculate in one branch and +return the result to another (`branch.calculate` inside a formula), and +because a branch may be deleted after its results were stored elsewhere. + +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 test what is stored.** A formula that reads + `holder.get_known_periods()` or `simulation.get_array()` and acts on whether + a value exists depends on that value without its number being recorded. +- **Only the family is tracked.** A value calculated in a separate simulation + (one not created with `get_branch`) and passed in is not covered. +- **Inputs on the root simulation drop nothing.** `set_input` on a simulation + that is not a branch keeps its earlier behaviour: values it already + calculated stay. +- **Existing child branches keep their values.** An input set on a branch does + not reach branches already created from it. + +## Dropping calculated values + +`simulation.drop_computed_arrays()` deletes every value the simulation holds +except inputs (the dataset or situation it was built from, and values set with +`set_input` on it or on the simulations it was created from), and returns how +many arrays it deleted. 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 a52bdceb1..0c82d6a9e 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -1,9 +1,10 @@ -from typing import Dict, Union +from typing import Dict, Optional, Set, Union import numpy from numpy.typing import ArrayLike from policyengine_core import periods +from policyengine_core.data_storage.store_history import next_sequence_number from policyengine_core.periods import Period @@ -17,11 +18,18 @@ class InMemoryStorage: def __init__(self, is_eternal: bool): self._arrays = {} + # When each array was stored (see ``store_history``), and which were + # stored as inputs rather than calculated. Both describe the stored + # value, so ``clone`` copies them with the arrays. + self._sequence_numbers: Dict[str, int] = {} + self._input_keys: Set[str] = set() self.is_eternal = is_eternal def clone(self) -> "InMemoryStorage": clone = InMemoryStorage(self.is_eternal) clone._arrays = {period: array.copy() for period, array in self._arrays.items()} + clone._sequence_numbers = dict(self._sequence_numbers) + clone._input_keys = set(self._input_keys) return clone def get(self, period: Period, branch_name: str = "default") -> ArrayLike: @@ -34,8 +42,19 @@ def get(self, period: Period, branch_name: str = "default") -> ArrayLike: return values def put( - self, value: ArrayLike, period: Period, branch_name: str = "default" + self, + value: ArrayLike, + period: Period, + branch_name: str = "default", + sequence_number: Optional[int] = None, + is_input: bool = False, ) -> None: + """Store ``value`` for ``period`` on ``branch_name``. + + ``sequence_number`` records when the value was stored (a new number + by default), and ``is_input`` whether it is an input rather than a + calculated value; see :meth:`drop_computed`. + """ if self.is_eternal: period = periods.period(periods.ETERNITY) period = periods.period(period) @@ -55,7 +74,42 @@ def put( f"Cannot cache period {period} anchored mid-month: its " "string form is lossy (see policyengine-core#526)." ) - self._arrays[f"{branch_name}:{period}"] = value + key = f"{branch_name}:{period}" + self._arrays[key] = value + self._sequence_numbers[key] = ( + next_sequence_number() if sequence_number is None else sequence_number + ) + if is_input: + self._input_keys.add(key) + else: + self._input_keys.discard(key) + + 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 with ``is_input``, are kept + whatever their number. + """ + dropped = [ + key + for key in self._arrays + if key not in self._input_keys + and (since is None or self._sequence_numbers.get(key, since) >= since) + ] + for key in dropped: + del self._arrays[key] + self._sequence_numbers.pop(key, None) + return len(dropped) + + def _forget_deleted_keys(self) -> None: + self._sequence_numbers = { + key: number + for key, number in self._sequence_numbers.items() + if key in self._arrays + } + self._input_keys.intersection_update(self._arrays) def delete(self, period: Period = None, branch_name: str = "default") -> None: if period is None: @@ -67,6 +121,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: for period_item, value in self._arrays.items() if not period_item.startswith(branch_prefix) } + self._forget_deleted_keys() return if self.is_eternal: @@ -84,6 +139,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: and period.contains(periods.period(period_item.split(":", 1)[1])) ) } + self._forget_deleted_keys() def get_known_periods(self) -> list: # Split on the first colon only: an anchored period's string form diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 3563c805d..0977d0bd9 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -1,10 +1,12 @@ import os import shutil +from typing import Dict, Optional, Set import numpy from numpy.typing import ArrayLike from policyengine_core import periods +from policyengine_core.data_storage.store_history import next_sequence_number from policyengine_core.enums import EnumArray from policyengine_core.periods import Period @@ -22,6 +24,10 @@ def __init__( ): self._files = {} self._enums = {} + # As in ``InMemoryStorage``: when each file was stored, and which were + # stored as inputs. + self._sequence_numbers: Dict[str, int] = {} + self._input_keys: Set[str] = set() self.is_eternal = is_eternal self.preserve_storage_dir = preserve_storage_dir self.storage_dir = storage_dir @@ -43,6 +49,8 @@ def clone(self) -> "OnDiskStorage": ) clone._files = self._files.copy() clone._enums = self._enums.copy() + clone._sequence_numbers = dict(self._sequence_numbers) + clone._input_keys = set(self._input_keys) clone._storage_dir_owner = getattr(self, "_storage_dir_owner", self) return clone @@ -64,8 +72,18 @@ def get(self, period: Period, branch_name: str = "default") -> ArrayLike: return self._decode_file(values) def put( - self, value: ArrayLike, period: Period, branch_name: str = "default" + self, + value: ArrayLike, + period: Period, + branch_name: str = "default", + sequence_number: Optional[int] = None, + is_input: bool = False, ) -> None: + """Store ``value`` for ``period`` on ``branch_name``. + + ``sequence_number`` and ``is_input`` are as in + :meth:`InMemoryStorage.put`. + """ if self.is_eternal: period = periods.period(periods.ETERNITY) period = periods.period(period) @@ -77,6 +95,39 @@ def put( value = value.view(numpy.ndarray) numpy.save(path, value) self._files[filename] = path + self._sequence_numbers[filename] = ( + next_sequence_number() if sequence_number is None else sequence_number + ) + if is_input: + self._input_keys.add(filename) + else: + self._input_keys.discard(filename) + + 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 + (other views of this directory may still use them) until the storage + directory is removed. + """ + dropped = [ + key + for key in self._files + if key not in self._input_keys + and (since is None or self._sequence_numbers.get(key, since) >= since) + ] + for key in dropped: + del self._files[key] + self._sequence_numbers.pop(key, None) + return len(dropped) + + def _forget_deleted_keys(self) -> None: + self._sequence_numbers = { + key: number + for key, number in self._sequence_numbers.items() + if key in self._files + } + self._input_keys.intersection_update(self._files) def delete(self, period: Period = None, branch_name: str = "default") -> None: if period is None: @@ -89,6 +140,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._forget_deleted_keys() return if self.is_eternal: @@ -101,6 +153,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._forget_deleted_keys() def get_known_periods(self) -> list: return list([periods.period(x.split("_")[1]) for x in self._files.keys()]) @@ -113,6 +166,8 @@ def get_known_branch_periods(self) -> list: def restore(self) -> None: self._files = files = {} + self._sequence_numbers = {} + self._input_keys = set() # Restore self._files from content of storage_dir. for filename in os.listdir(self.storage_dir): if not filename.endswith(".npy"): diff --git a/policyengine_core/data_storage/store_history.py b/policyengine_core/data_storage/store_history.py new file mode 100644 index 000000000..a1725fdcd --- /dev/null +++ b/policyengine_core/data_storage/store_history.py @@ -0,0 +1,88 @@ +"""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 +from typing import Dict, Optional + +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 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: + """When each variable's values were first stored in a family of simulations. + + A simulation and the branches created from it (``get_branch``, at any + depth) share one history. For each variable it records: + + - for each period, the sequence number of the first value stored or + calculated for that period anywhere in the family, 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 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 a value). Such a value depends on which other periods + hold values, so an input set for any period can change it. + + Deleting or replacing a stored value does not change the history: a value + calculated from the deleted or replaced one keeps its number. + """ + + def __init__(self): + self._first_stored: Dict[str, Dict[Period, int]] = {} + self._first_derived: Dict[str, int] = {} + + def record_store(self, variable_name: str, period: Period, sequence_number: int): + stored = self._first_stored.setdefault(variable_name, {}) + if period not in stored: + stored[period] = sequence_number + + def record_derived(self, variable_name: str, sequence_number: int): + if variable_name not in self._first_derived: + self._first_derived[variable_name] = sequence_number + + 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 first store of the variable for any period that + shares a day with ``period``, or the first time one of its values was + derived from its other periods, whichever came first. ``None`` means + nothing in the family has read or stored the variable for an + overlapping period nor derived a value of it, so no stored value can + depend on its 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) + return min(candidates) if candidates else None diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 52ae81295..c686ea286 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -8,6 +8,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 @@ -98,6 +99,18 @@ def delete_arrays( if self._disk_storage: self._disk_storage.delete(period, branch_name) + def drop_computed(self, since: 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. + """ + dropped = self._memory_storage.drop_computed(since) + if self._disk_storage is not None: + dropped += self._disk_storage.drop_computed(since) + return dropped + def _get_array_from_storage( self, period: Period, branch_name: str = "default" ) -> ArrayLike: @@ -257,6 +270,9 @@ 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``). + simulation._drop_values_that_may_depend_on(self.variable.name, period) if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() if not hasattr(simulation, "_user_input_contexts"): @@ -354,10 +370,17 @@ def _set( >= self.simulation.memory_config.max_memory_occupation_pc ) - if should_store_on_disk: - self._disk_storage.put(value, period, branch_name) - else: - self._memory_storage.put(value, period, branch_name) + sequence_number = next_sequence_number() + is_input = bool(user_input_contexts) + storage = self._disk_storage if should_store_on_disk else self._memory_storage + storage.put( + value, + period, + branch_name, + sequence_number=sequence_number, + is_input=is_input, + ) + self._record_store(period, sequence_number) if user_input_contexts: if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() @@ -366,18 +389,27 @@ def _set( def put_in_cache( self, value: ArrayLike, period: Period, branch_name: str = "default" ) -> None: - 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 self._set(period, value, branch_name) + def _record_store(self, period: Period, sequence_number: int) -> None: + simulation = getattr(self, "simulation", None) + if simulation is None: + return + if self.variable.definition_period == periods.ETERNITY: + period = periods.period(periods.ETERNITY) + simulation._get_store_history().record_store( + self.variable.name, periods.period(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 566407a59..4c0a0ab89 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1,6 +1,5 @@ import hashlib import tempfile -from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union import numpy as np @@ -11,6 +10,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.entities.entity import Entity from policyengine_core.enums import Enum, EnumArray from policyengine_core.errors import CycleError, SpiralError @@ -88,18 +91,6 @@ def _uprating_index_value(parameter, instant) -> Optional[float]: ) -@dataclass(frozen=True) -class PreservedUserInput: - variable_name: str - branch_name: str - period: Period - value: object - storage: str - disk_key: Optional[str] = None - disk_file: Optional[str] = None - disk_enum: object = None - - class Simulation: """ Represents a simulation, and handles the calculation logic @@ -181,6 +172,9 @@ def __init__( # post-``apply_reform`` cache wipe would also wipe the dataset the # simulation was loaded from. self._user_input_keys: set[tuple[str, str, Period]] = set() + # When each variable's values were first stored, shared with every + # branch created from this simulation (see ``set_input``). + self._store_history = StoreHistory() self.debug: bool = False self.trace: bool = trace self.tracer: SimpleTracer = SimpleTracer() if not trace else FullTracer() @@ -305,83 +299,15 @@ def _invalidate_all_caches(self) -> None: Called after ``apply_reform`` and any other operation that changes the tax-benefit system underneath an already-calculated simulation. - Every (variable, branch, period) that was populated via - ``set_input`` is preserved — those are source data, not stale - formula output — so a structural reform applied after dataset - load doesn't silently discard the dataset. Everything else - (formula outputs, cached short-path results, on-disk caches) is - wiped so the next ``calculate`` recomputes under the new - tax-benefit system. + Every value stored through ``set_input`` is preserved — those are + source data, not stale formula output — so a structural reform + applied after dataset load doesn't silently discard the dataset. + Everything else (formula outputs, cached short-path results, on-disk + caches) is wiped, here and in every branch, 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 - # hand-building keys, since ETERNITY variables canonicalize every - # period to the single ETERNITY storage key. - preserved: list[PreservedUserInput] = [] - user_input_keys = getattr(self, "_user_input_keys", None) or set() - for variable_name, branch_name, period in user_input_keys: - holder = self.get_holder(variable_name) - stored_value = holder._memory_storage.get(period, branch_name) - if stored_value is not None: - preserved.append( - PreservedUserInput( - variable_name=variable_name, - branch_name=branch_name, - period=period, - value=stored_value, - storage="memory", - ) - ) - continue - if holder._disk_storage is not None: - disk_period = ( - periods.period(periods.ETERNITY) - if holder._disk_storage.is_eternal - else periods.period(period) - ) - disk_key = f"{branch_name}_{disk_period}" - disk_file = holder._disk_storage._files.get(disk_key) - if disk_file is not None: - preserved.append( - PreservedUserInput( - variable_name=variable_name, - branch_name=branch_name, - period=period, - value=None, - storage="disk", - disk_key=disk_key, - disk_file=disk_file, - disk_enum=holder._disk_storage._enums.get(disk_file), - ) - ) - # Iterate only over holders that already exist on each population — - # lazy-creating a holder for every variable in the tax-benefit - # system (thousands in policyengine-us) inflated the cost of - # ``apply_reform`` from milliseconds to seconds and broke the - # YAML full-suite on downstream repos. Untouched variables have - # no holder and therefore nothing to wipe. - for population in self.populations.values(): - for holder in population._holders.values(): - holder._memory_storage._arrays = {} - if holder._disk_storage is not None: - holder._disk_storage._files = {} - # 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 - if user_input.disk_enum is not None: - holder._disk_storage._enums[user_input.disk_file] = ( - user_input.disk_enum - ) - else: - holder._memory_storage.put( - user_input.value, - user_input.period, - user_input.branch_name, - ) + self._drop_computed() for branch in self.branches.values(): branch._invalidate_all_caches() @@ -864,6 +790,16 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # 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 = [ @@ -1364,6 +1300,24 @@ 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, and inputs set on it or, + for a branch, on the simulations it was created from. Every other + value is calculated again when next requested. + + 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. + """ + 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 @@ -1397,6 +1351,31 @@ 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 (see + :mod:`policyengine_core.data_storage.store_history`). So the branch + drops each value it holds, other than an input, that was stored at + or after the earliest of: the first time anything in this + simulation's family (the simulation the branches were created from, + and all its branches) stored or calculated ``variable_name`` for a + period that shares a day with ``period``; and the first time a value + of ``variable_name`` was derived from its other periods (uprated, + carried over, or given the default for lack of a formula result). If + neither has happened, nothing can depend on the value and nothing is + dropped; this is the usual case, where a formula creates the branch + while still calculating the variable it overrides. + + 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``) rather than calculating it; + and values calculated in a different simulation family. Inputs set + on a simulation that is not a branch drop nothing, as before. """ period = periods.period(period) if self.start_instant is None or self.start_instant > period.start: @@ -1411,6 +1390,38 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non if _fast_cache is not None: _fast_cache.pop((variable_name, period), None) + def _get_store_history(self) -> StoreHistory: + history = getattr(self, "_store_history", None) + if history is None: + history = self._store_history = StoreHistory() + return history + + 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 + since = self._get_store_history().earliest_dependency( + variable_name, periods.period(period) + ) + if since is None: + return 0 + return self._drop_computed(since) + + def _drop_computed(self, since: Optional[int] = None) -> int: + 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 = {} + return dropped + def get_variable_population(self, variable_name: str) -> Population: variable = self.tax_benefit_system.get_variable( variable_name, check_existence=True @@ -1485,6 +1496,12 @@ def get_branch( ) -> "Simulation": """Create a clone of this simulation, whose calculations are traced in the original. + The branch starts with the values this simulation holds. 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. @@ -1496,6 +1513,8 @@ def get_branch( return self if name in self.branches: return self.branches[name] + # The branch shares this simulation's store history. + self._get_store_history() branch = self.clone(clone_tax_benefit_system=clone_system) self.branches[name] = branch branch.branch_name = name diff --git a/pyproject.toml b/pyproject.toml index 4f102bfb5..abc4bb467 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ dev = [ "ruff>=0.9.0", "coverage", "furo", + "hypothesis>=6,<7", "jupyter-book<1", "markupsafe==2.0.1", "mypy==0.991", diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py new file mode 100644 index 000000000..1827a4bbc --- /dev/null +++ b/tests/core/test_branch_input_invalidation.py @@ -0,0 +1,785 @@ +"""``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. An +input set on a branch drops the branch's non-input values stored at or after +the first time anything in the family stored, or derived, the overridden +variable for an overlapping period (``StoreHistory.earliest_dependency``). + +The invariant tested throughout (and by the property test at the end, over +random sequences of calculations, branches and inputs): 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 hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem, entities +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.holders import set_input_divide_by_period +from policyengine_core.model_api import Reform +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.variables import Variable + +JANUARY = periods.period("2017-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 _stored_keys(branch, "housing_tax") == {"default:2017"} + assert _stored_keys(branch, "income_tax") == set() + assert branch.calculate("income_tax", JANUARY) == pytest.approx( + 4000 * tax_benefit_system.parameters(JANUARY).taxes.income_tax_rate + ) + + +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])) + + assert branch.get_holder("housing_allowance")._disk_storage._files == {} + rent_disk = branch.get_holder("rent")._disk_storage + assert set(rent_disk._files) == {"default_2017-01", "cheaper_2017-01"} + assert rent_disk._input_keys == {"default_2017-01", "cheaper_2017-01"} + 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), + ) + + +# ----- 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._input_keys + } + 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]) + + +# ----- 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"), is_input=True) + first = storage._sequence_numbers["default:2020"] + storage.put(np.array([2.0]), periods.period("2021")) + clone = storage.clone() + + assert clone._sequence_numbers == storage._sequence_numbers + assert clone._sequence_numbers["default:2021"] > first + assert clone._input_keys == {"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), 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 + + +# ----- A synthetic system for derived values and branches inside formulas ----- # + + +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 _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), +] + + +def _synthetic_system(carry_over): + class System(CountryTaxBenefitSystem): + auto_carry_over_input_variables = carry_over + + system = System() + system.add_variables(*SYNTHETIC_VARIABLES) + return system + + +# Carry-over is left out of the random programs below: 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"] + + +def _synthetic(inputs, memory_config=None, system=SYNTHETIC_SYSTEM): + """A simulation of ``PEOPLE`` people given ``inputs`` before anything else.""" + simulation = SimulationBuilder().build_default_simulation(system, count=PEOPLE) + if memory_config is not None: + 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: + simulation.get_holder(variable) + 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) + }, +} + + +def test_value_carried_over_from_an_earlier_input_is_dropped(): + simulation = _synthetic(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( + {**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(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({**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 and its later branches never store ``p_inner_only``. The + history the whole family shares still records it. + """ + simulation = _synthetic(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({**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(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]) + + +# ----- Property: branch results do not depend on what was calculated before ----- # + +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_up", + "p_c", +] +MONTHS = ["2013-01", "2013-07", "2014-12", "2015-03"] +MEMORY_CONFIGS = { + "memory": lambda: None, + "disk": lambda: MemoryConfig(max_memory_occupation=0), + "not_kept": lambda: MemoryConfig( + max_memory_occupation=1, variables_to_drop=["p_sum", "p_inner_only"] + ), +} + +# 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. +simulation_index = st.integers(0, 3) +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), + st.tuples(st.just("drop"), 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, 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_month", "2013-07", "p_m", "2013"), + ("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]), + ("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]), + ("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]), + ] + ) +) +operations = st.lists( + st.one_of( + single_operation.map(lambda operation: [operation]), + override_after_calculating, + ), + min_size=1, + max_size=15, +).map(lambda chunks: [operation for chunk in chunks for operation in chunk]) + + +def _stored_periods(variable, period): + """The periods ``set_input`` stores ``variable`` at, given ``period``.""" + period = periods.period(period) + if SYNTHETIC_SYSTEM.get_variable(variable).definition_period == periods.MONTH: + return [str(month) for month in period.get_subperiods(periods.MONTH)] + return [str(period)] + + +def _own_inputs(simulation, variable, period): + """The values stored on the branch itself for ``variable`` over ``period``.""" + holder = simulation.get_holder(variable) + return { + (variable, stored_period): tuple( + holder._get_array_from_storage(stored_period, simulation.branch_name) + ) + for stored_period in _stored_periods(variable, period) + } + + +def _run(program, memory_config): + """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, as stored (an input for a + year is divided between the months the branch has not set itself). + """ + root = _synthetic(ROOT_INPUTS, memory_config=memory_config) + simulations = [root] + inputs = [dict(ROOT_INPUTS)] + own_inputs = [set()] + 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": + simulations.append(simulation.get_branch(f"b{len(simulations)}")) + inputs.append(dict(inputs[index])) + own_inputs.append(set()) + elif kind == "set": + if index == 0: + continue # Inputs on the root after calculating are out of scope. + _, _, variable, period, value = operation + stored = {(variable, p) for p in _stored_periods(variable, period)} + if len(stored) > 1 and stored <= own_inputs[index]: + # Dividing a year between months the branch has all set + # itself is an error unless the totals match. + continue + simulation.set_input(variable, period, np.asarray(value)) + stored = _own_inputs(simulation, variable, period) + inputs[index].update(stored) + own_inputs[index].update(stored) + elif kind == "calculate": + _, _, variable, period = operation + value = np.array(simulation.calculate(variable, period), copy=True) + results.append((dict(inputs[index]), variable, period, value)) + else: + simulation.drop_computed_arrays() + return results + + +@settings( + max_examples=500, + deadline=None, + suppress_health_check=[HealthCheck.too_slow], +) +@given(program=operations, mode=st.sampled_from(sorted(MEMORY_CONFIGS))) +def test_branch_calculations_match_a_simulation_given_its_inputs_first(program, mode): + for branch_inputs, variable, period, value in _run(program, MEMORY_CONFIGS[mode]()): + expected = _synthetic(branch_inputs).calculate(variable, period) + # Uprating chains may round differently in float32 depending on which + # earlier periods were calculated first. + np.testing.assert_allclose( + value, expected, rtol=1e-5, err_msg=f"{variable} {period}" + ) diff --git a/uv.lock b/uv.lock index 2608e8504..222deb407 100644 --- a/uv.lock +++ b/uv.lock @@ -130,7 +130,7 @@ name = "cffi" version = "2.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser" }, + { name = "pycparser", marker = "implementation_name != 'PyPy'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } wheels = [ @@ -682,6 +682,94 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d5/ae/2f6d96b4e6c5478d87d606a1934b5d436c4a2bce6bb7c6fdece891c128e3/huggingface_hub-1.4.1-py3-none-any.whl", hash = "sha256:9931d075fb7a79af5abc487106414ec5fba2c0ae86104c0c62fd6cae38873d18", size = 553326, upload-time = "2026-02-06T09:20:00.728Z" }, ] +[[package]] +name = "hypothesis" +version = "6.168.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "sortedcontainers" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/09/b7/13118bbc45d6d8b9d04e2de779e2a4ff23145ea39efa692b994fb874ca72/hypothesis-6.168.3.tar.gz", hash = "sha256:a43388f9067678fef6e13bdff325b6cfa6961a590498bb37f7ff31589c83bc75", size = 511022, upload-time = "2026-09-28T05:20:58.499Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2b/cd/746a2fc1e5e5ba34f07ea6ee1ad07525a3dde5ff4836db3a7980af49a891/hypothesis-6.168.3-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:d20972ca134e652e9928ecd200966a8a2857adc2812c1d74b12a872e5cd503eb", size = 791536, upload-time = "2026-09-28T05:18:57.162Z" }, + { url = "https://files.pythonhosted.org/packages/c3/0f/7a158e377b69556c8e12c25c8fd811d0103e54c24a12dab1f3b9819f202e/hypothesis-6.168.3-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:eafbec09d3e87d13d8242411d1f5868b5e879f1bd6d95e233e9ff80c527bea1c", size = 787313, upload-time = "2026-09-28T05:20:18.122Z" }, + { url = "https://files.pythonhosted.org/packages/45/c6/7df9104ee359e8fcd77791781c47e70794964ca69673665e0de2a6590f4c/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:73c5627497968cc62d140e9fde1ea21e12cac7b64dc419fea8286cd34ff1ad4e", size = 1124036, upload-time = "2026-09-28T05:20:04.372Z" }, + { url = "https://files.pythonhosted.org/packages/92/0f/8a61715404a73b9a82cb1f976803428d603cdea976a9c30266c72ad6a7ce/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:29dc56e6dc6eeb0aeb08ac463f847279bde2336c4786faf7ee0af4b93cd031a7", size = 1147875, upload-time = "2026-09-28T05:19:15.495Z" }, + { url = "https://files.pythonhosted.org/packages/e4/3f/2b16e95cc3b9a069afa23830a012aefe377ce4a457a227e5a5026eb827ff/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:753bb501f8560d2e321ed3b62596b4496c56f15668b0a0669231ccc3e6c4e80d", size = 1149489, upload-time = "2026-09-28T05:19:18.886Z" }, + { url = "https://files.pythonhosted.org/packages/f6/e7/d7cf6dc068bb02732b2a6e38a35a0dcb7f2137608198fc79d740f69d197c/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:71ab606c472449cb872ef2a7acaec679ee4f651b7f0a477a04ebc85d8933ef8c", size = 1191992, upload-time = "2026-09-28T05:19:21.83Z" }, + { url = "https://files.pythonhosted.org/packages/5f/26/f1e5b25dec15998e8375d722825ec3dae1075cbe7cb1f74221b2312a2e33/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:26928956c54748e4dfa587333741ab750246123b93484955646ff316cca27eef", size = 1169911, upload-time = "2026-09-28T05:19:27.996Z" }, + { url = "https://files.pythonhosted.org/packages/03/36/901234a49147e6bf8a5b9e54b775706a223c2b52befea1106ea2aaf15d48/hypothesis-6.168.3-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:0bdfc53c041b61c854fa3761735736b991bc2bddafee45a361cfe4d43027c1ef", size = 1129406, upload-time = "2026-09-28T05:18:42.632Z" }, + { url = "https://files.pythonhosted.org/packages/55/c1/5fc9ff91ec9fba6bb527833386c66812524b5cc75fe7d4a4f9122ee1c420/hypothesis-6.168.3-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:650528e2b1e2a45e95c2df624d4b9364b8ac5a027ba84e1eea2bc5398dd01dcd", size = 1160399, upload-time = "2026-09-28T05:20:23.911Z" }, + { url = "https://files.pythonhosted.org/packages/9b/d2/49ad1ef5c55547c8abfd0fc571b3493806dda530fd5cb3d8263dbad86a72/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:209dc54cdb1b4d6d7020ad8d09e44c49756361b7183adacea4d6f5705085b595", size = 1299906, upload-time = "2026-09-28T05:18:58.417Z" }, + { url = "https://files.pythonhosted.org/packages/0d/5f/f98a26094c4f462c4215807bed0f9a6b5508b27a3c329bcba1bc8adc7ac7/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:af8ca98cd11dc7f9427bb90483abcd9688d4df2e8e63893a30a2423b027ebb11", size = 1425538, upload-time = "2026-09-28T05:20:49.074Z" }, + { url = "https://files.pythonhosted.org/packages/d0/b8/c5b4406e3a6591e51a42f85469351d1d824e0a54b84167aec9b2353759fe/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:f8122cfdd0bc0ba843063effa22ac310bdebdb3a1319bf1e63f14baa119c72ab", size = 1377084, upload-time = "2026-09-28T05:20:06.618Z" }, + { url = "https://files.pythonhosted.org/packages/e0/01/0d84a8ea469d024c15f602799c0b8410fbe03d99fdb0370449b33c3c916f/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_ppc64le.whl", hash = "sha256:4bedbb379eab34f792af7ee9a05aae04e9c08bbb52e0f90f5a2110d4fe4b2fbd", size = 1281162, upload-time = "2026-09-28T05:19:42.768Z" }, + { url = "https://files.pythonhosted.org/packages/22/b5/ea6435038de1a795f005cb603531e8ccc053febc50180f3ab55583c9c60e/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_riscv64.whl", hash = "sha256:5a953115b9f5c95133ab2d04efffeec96e5658c3207923dca7285f7db3e6bbef", size = 1300433, upload-time = "2026-09-28T05:18:55.634Z" }, + { url = "https://files.pythonhosted.org/packages/6c/78/fa6c77d9f64b6d592e69e582d83cbce7d7e56897faa4762885a2123b0166/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:6173558e676ad25ed1e20507fa4024a0d816dd90f715f77006e5a28f19109026", size = 1336273, upload-time = "2026-09-28T05:19:08.273Z" }, + { url = "https://files.pythonhosted.org/packages/f1/56/4acd0d778bab818c9ea336fcba768bf43bc70db7f84cc373579298df4641/hypothesis-6.168.3-cp310-abi3-win32.whl", hash = "sha256:dc66390fb12d80585aa9222bf538ce8b7aa22cf5d118250647355a1c9e8f62f4", size = 678211, upload-time = "2026-09-28T05:19:23.355Z" }, + { url = "https://files.pythonhosted.org/packages/b0/db/0a02b146ad1c30f9716362f2e1a379dfd6b0f7be03ac1e22c9da24c9e2a9/hypothesis-6.168.3-cp310-abi3-win_amd64.whl", hash = "sha256:92325b276360fe86c5bf71a568c0d53a6d140b0de36dcf029f9164a17803bb24", size = 684906, upload-time = "2026-09-28T05:18:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/a3/c0/958deaf726848f96f52250740bf39f13f476b068e22f73e2b358eaa07532/hypothesis-6.168.3-cp310-abi3-win_arm64.whl", hash = "sha256:3cf6f1eeaf41cd8d60cf1f88fde905ca1dd77c906929a507d6ac7f66f2ccba2a", size = 683292, upload-time = "2026-09-28T05:19:57.404Z" }, + { url = "https://files.pythonhosted.org/packages/f4/ab/443208c2988a24eaafbfbd6791b897ad9a8a1dd7f870d5b455d7d15a28ad/hypothesis-6.168.3-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:9029775dc2e25e3e8e596c4306926f0079158353631c7c52c285bbc5c8073c2b", size = 792239, upload-time = "2026-09-28T05:19:59.144Z" }, + { url = "https://files.pythonhosted.org/packages/cf/7b/619c1a76c52de2d3cfd68a2f9876c09464d63097db5bf86a28b62dc37bb1/hypothesis-6.168.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:0e1d91530cefdb9e0b6467c203eb509c61a11d70480a37db59cb53a2f9913102", size = 788113, upload-time = "2026-09-28T05:19:17.145Z" }, + { url = "https://files.pythonhosted.org/packages/ea/dd/33841e8833105f850214a486d83740ba52041892ca94e687a05829c0ce83/hypothesis-6.168.3-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0df00cbe8133aa220308d11fa28e3add996cf5102a39b279c1d711015fd9114b", size = 1124241, upload-time = "2026-09-28T05:18:37.589Z" }, + { url = "https://files.pythonhosted.org/packages/4b/43/c84a7900feced6a420e824ab61d316e4f59a6887d991bc143ecc294a09ed/hypothesis-6.168.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd7f75a2e23288ee82ee965a952473d9c5447c2cbc1afc94be09d2402201774e", size = 1170406, upload-time = "2026-09-28T05:19:46.365Z" }, + { url = "https://files.pythonhosted.org/packages/47/5d/6e2a32b9c3020af7a06b00fe446a2c5155ae76ac468c0b68d7188e524002/hypothesis-6.168.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:bc2b1c37631f2231dd0d077225aa5908fc2bd34a4c3faf475053b0bc2eb24c95", size = 1300339, upload-time = "2026-09-28T05:20:31.968Z" }, + { url = "https://files.pythonhosted.org/packages/f9/78/19f3c9eb9a5c215b9fbde57be18ef7c1c78afa509b460d4e8f86026d741a/hypothesis-6.168.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:73960a58f6efc8cbc57d8ad0647955f7c5b25f3d562e3eaa57afa1c69894f7aa", size = 1336484, upload-time = "2026-09-28T05:19:02.841Z" }, + { url = "https://files.pythonhosted.org/packages/88/bb/d46c4519e989e022667705d8a56e83a09ad4ea50eae07ca4bb0f4708dd23/hypothesis-6.168.3-cp311-cp311-win_amd64.whl", hash = "sha256:3f1122223759acc0fe5301c1ae505d762ee61af91db7b191a791ccdbcd965462", size = 684646, upload-time = "2026-09-28T05:20:00.842Z" }, + { url = "https://files.pythonhosted.org/packages/7a/56/44a17266bcf1add5fba7c3df5c3920396a5e924d917fb0534a6f9e2b9b9a/hypothesis-6.168.3-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:bff12e036b67abc5ded682f724f4d7afaff5223190270a0763695f218c079068", size = 793313, upload-time = "2026-09-28T05:20:22.045Z" }, + { url = "https://files.pythonhosted.org/packages/d8/18/1a25b11e9133cb54c02ee6cf99a61296ad461aabc7aae3399e1d0ff1f095/hypothesis-6.168.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c73c6188056e6dc110260e39d141c69c3d02f23e3ba1d8610c7a8b6f2510ec96", size = 784855, upload-time = "2026-09-28T05:19:34.496Z" }, + { url = "https://files.pythonhosted.org/packages/8b/da/728d3210f0b0089337f37206a100d31ef546d1f99839e709f0a9ca443f8e/hypothesis-6.168.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bb1063cb794765097b8c0add598bd34a904f44dc15db81275dd1f0ed1ac567ed", size = 1123043, upload-time = "2026-09-28T05:19:05.468Z" }, + { url = "https://files.pythonhosted.org/packages/53/86/bfb9ec344a8b3ef2fa6f59bb3d30c5b3a7d86d594107670539c34ee53d1b/hypothesis-6.168.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9b2d47f5d9060b049039bef0b267e88480f6468f95ee394e514e84a81ccdc5f5", size = 1169121, upload-time = "2026-09-28T05:20:40.512Z" }, + { url = "https://files.pythonhosted.org/packages/3b/e0/5a76dbafbcaf8fff075f09377cddefa17289ef05dbf51878dde88385b966/hypothesis-6.168.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8367f623a98cbf90f33fd2a43b100671b152b5975bc343a578290945f52f7018", size = 1298749, upload-time = "2026-09-28T05:19:32.883Z" }, + { url = "https://files.pythonhosted.org/packages/28/15/0e65af96b353f2c4bdc283c53ed0980e5ac5009afe97ef3b23c77b364664/hypothesis-6.168.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fd81241f4cd76fb18d481e87b9688b30a0cb5c30841730513323ba1c7aba1837", size = 1335271, upload-time = "2026-09-28T05:20:54.005Z" }, + { url = "https://files.pythonhosted.org/packages/f5/c5/37eab2155664aa544df321f6d3940b89c6702d79a72fdc98d7b5a2286f8c/hypothesis-6.168.3-cp312-cp312-win_amd64.whl", hash = "sha256:377438de53afb94347d9845b7d90db5905c6d2a8d612108e9938e0f80503bb90", size = 682211, upload-time = "2026-09-28T05:18:51.206Z" }, + { url = "https://files.pythonhosted.org/packages/0c/a2/6787da846d929e52fc3344d803299c45782dbeac287528af08380e984bf8/hypothesis-6.168.3-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:1b230a850de63334c16654a34a2d547e0179d36b9071d4439b3e7237f6d077e7", size = 793254, upload-time = "2026-09-28T05:19:36.294Z" }, + { url = "https://files.pythonhosted.org/packages/89/ba/2893f5ca42501f4d562cba3229c8994c3a3e5fac66ef60c1e0e0b5220a5d/hypothesis-6.168.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d63b0226cd3e0d8bdd97c3384b22a21934ed4d53246c1c93575dada616672499", size = 784699, upload-time = "2026-09-28T05:18:54.194Z" }, + { url = "https://files.pythonhosted.org/packages/b7/aa/7d7349daf75b71f6f35876f8de115e974c5c04d600e86df1b2779b0883ab/hypothesis-6.168.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0369f5df055f96e117ab12e5f249668ff144731ab5280bc7a205fdf81f990b89", size = 1123016, upload-time = "2026-09-28T05:20:36.372Z" }, + { url = "https://files.pythonhosted.org/packages/00/f0/7774e1ea072708ea46cb25c4aeee5f9978b6c271764809069e8a6789858e/hypothesis-6.168.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f076bcd0f77fdcdb7797826c879099573d03a02e228ebe649ea917081b962ac0", size = 1168989, upload-time = "2026-09-28T05:20:08.44Z" }, + { url = "https://files.pythonhosted.org/packages/a7/7d/113992abed9efbd7944e3a496a382da58152dce126fddf6419ff31a0a57d/hypothesis-6.168.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:b823ba1fcec8da730f29316d010b06d3f7e0c3828dcf91020e24c55e7d24652a", size = 1298626, upload-time = "2026-09-28T05:19:40.986Z" }, + { url = "https://files.pythonhosted.org/packages/d5/65/3659fa5e733027e5b37a27486e57f4053535e40853d23501c422bb275043/hypothesis-6.168.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dd2849c269d674e4618590f3b48d443bd4c06b5aef3d8d869086c2e6d213d248", size = 1335184, upload-time = "2026-09-28T05:20:20.087Z" }, + { url = "https://files.pythonhosted.org/packages/b4/6c/35aab2221b5ea65125340f236e1a765a9e9f3b28d89a389ea0033fed29f7/hypothesis-6.168.3-cp313-cp313-win_amd64.whl", hash = "sha256:3ef7d26f5789e691401d5f87eafed9bd2763f0dbe47af6d6b66012a509404766", size = 682173, upload-time = "2026-09-28T05:19:12.762Z" }, + { url = "https://files.pythonhosted.org/packages/6e/79/27b0cb56ff5d2bc92458bf6fcebdb1f0dc01f57f043a178e9f9d74498169/hypothesis-6.168.3-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:e2d4c68729a13df9af4998d2652cfb5d541c5609b88c306880a5dfb284938ec2", size = 793287, upload-time = "2026-09-28T05:20:38.522Z" }, + { url = "https://files.pythonhosted.org/packages/e8/bd/5342f95c3bc36586777ef8cb8eba87ac7acf0184210327fb0e6d8654e295/hypothesis-6.168.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:b345f818083ec99966a43ca4f7b38feb62c6920ce28572bc2bde4948a8b7eaba", size = 784857, upload-time = "2026-09-28T05:20:42.609Z" }, + { url = "https://files.pythonhosted.org/packages/4a/5f/fc774241518a0588d362680a3084275058aab86d6bc02a2c61e1073142d0/hypothesis-6.168.3-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0608c610fc002978fc5de0f471770d8817e9455cb8649a47983c9f8bccfc1897", size = 1123330, upload-time = "2026-09-28T05:19:29.735Z" }, + { url = "https://files.pythonhosted.org/packages/b0/e6/131f16775a3dca5f4fe27f0d6ad9a9851600098a54a872d163f6ebf0b664/hypothesis-6.168.3-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9dcf6448b1ecc37f2b23f2d1b3ddfc9ff6b6910614c15a642dc82f419b96037b", size = 1169178, upload-time = "2026-09-28T05:19:53.436Z" }, + { url = "https://files.pythonhosted.org/packages/06/61/c9f5bff8b73321c9fd12f2d69666c6e4861b97b60f5f24fdd6ac293007fd/hypothesis-6.168.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ae4f9f094041dcce5119ebd7bab71062743ab02b6e654d056b370beda78c19e2", size = 1299127, upload-time = "2026-09-28T05:20:14.486Z" }, + { url = "https://files.pythonhosted.org/packages/e4/6d/d4618f8ab12dd76c4d58ea052252759aa2bf0c4e71f4b502ba5d577e0b9f/hypothesis-6.168.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d479985fe73af97badfb72dc6d20c6a353e486a36f6d065f600026e0cf954a86", size = 1335390, upload-time = "2026-09-28T05:20:10.326Z" }, + { url = "https://files.pythonhosted.org/packages/95/b6/727161e17cc297df1aeea945de05c532033b66751625bec192bc5b5ff707/hypothesis-6.168.3-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:785e2c45f8c08e274b4bf1ccae97f1a4e09407e9a68e80790cfab1d17a4a45fa", size = 624291, upload-time = "2026-09-28T05:18:36.376Z" }, + { url = "https://files.pythonhosted.org/packages/77/23/2f9b506392a0b8712440bcb781abc5713be9ffac1303b6d10f3669050882/hypothesis-6.168.3-cp314-cp314-win_amd64.whl", hash = "sha256:320920b1e3dae8611eee8a03d063cf2187446f2a17c38cfb8a7fc1466f71eee2", size = 682064, upload-time = "2026-09-28T05:20:16.336Z" }, + { url = "https://files.pythonhosted.org/packages/50/93/efabfd95eb2b69c0c9fa1e1c83c8290aa2715d46baeb5a7764faf09c133a/hypothesis-6.168.3-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:35380baa981108a7f60c4eab71e46acd8d8f58440520346a1a6aba06dca7e074", size = 791875, upload-time = "2026-09-28T05:19:37.789Z" }, + { url = "https://files.pythonhosted.org/packages/1b/fc/2a0ada1623a9a33048bbf9f00a7c92513398c6b8865cba5efb854262ce30/hypothesis-6.168.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:f0aaaed00438fa6673d856aaee12a4a8afbde82a9f3c5ff487cacfb064239afe", size = 783412, upload-time = "2026-09-28T05:18:48.266Z" }, + { url = "https://files.pythonhosted.org/packages/41/31/73c615d37eeb12da0209555612c32ced628d860267c23cfd2691c6207aea/hypothesis-6.168.3-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f27e6df1576bf838e7f484d4ae5cba92114497ca30afb917056bf1ea33e95675", size = 1121662, upload-time = "2026-09-28T05:19:26.225Z" }, + { url = "https://files.pythonhosted.org/packages/92/54/1893344a3b8bbf83fa7c9bb7e96f6836580213854f46671cb749f7b4706c/hypothesis-6.168.3-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3bc85014577982ec6d2e266edc5cd6e7a7d3674c648791ca983c67a52f89a0c4", size = 1167761, upload-time = "2026-09-28T05:18:49.85Z" }, + { url = "https://files.pythonhosted.org/packages/26/97/a61d82968febd25a0f08ffbbfabfaeead66f68fc5fa73f3c9ece1c36fc8f/hypothesis-6.168.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:ccfc29505aa1cdcc254cb9cd701fd0811fef5480addd97d4127a00df023410f7", size = 1297309, upload-time = "2026-09-28T05:19:51.711Z" }, + { url = "https://files.pythonhosted.org/packages/bc/01/fe8e4cf6d39d02f4efadaa351427fa13722086cfc54a2da3579741c5177a/hypothesis-6.168.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:32d0699566aaa93f9e97a44705d7164386f1e91de78d277f53bada35e78bcc96", size = 1334260, upload-time = "2026-09-28T05:19:00.121Z" }, + { url = "https://files.pythonhosted.org/packages/39/6a/09177ebc62f4778f94379ecade6a51299d3d8c4f4b75ba636bf3a0202364/hypothesis-6.168.3-cp314-cp314t-win_amd64.whl", hash = "sha256:28d88fa174ecbd4ecbd7bb290f06d0db084a3971c3f511ae2830b65b5e25500f", size = 681992, upload-time = "2026-09-28T05:19:48.043Z" }, + { url = "https://files.pythonhosted.org/packages/f9/f6/890bf33d608cd63348d3146a3ca83362fbf09e589c5e5230a00c403070f4/hypothesis-6.168.3-cp315-abi3.abi3t-macosx_10_12_x86_64.whl", hash = "sha256:61f5782d807b1e6aa5c9beef1054cb2037e7d1add48a64972cd6ee447778781f", size = 791231, upload-time = "2026-09-28T05:18:43.998Z" }, + { url = "https://files.pythonhosted.org/packages/03/dd/fc36b204f8aa7437c676576d9437f0422de46aa8a729e6cd5bbb0224e759/hypothesis-6.168.3-cp315-abi3.abi3t-macosx_11_0_arm64.whl", hash = "sha256:7aacf3cf40c7ce8f9e4924d5b57bc0b068beafdfed2347bcf160a28b4cfbba7b", size = 783171, upload-time = "2026-09-28T05:19:20.353Z" }, + { url = "https://files.pythonhosted.org/packages/5c/d2/3cd087d577db67c404991d7723e64d7cd75570cdc82a0f7329aed6f2086c/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:01768a03a30dc54df7fe457c0b34016c84d598fa00d5965ededab93ba4eb3408", size = 1121152, upload-time = "2026-09-28T05:18:39.095Z" }, + { url = "https://files.pythonhosted.org/packages/42/6f/a05a68cc66e7a22701c50bf94fb9ddeb36f4f5d0de3cc130066709eda621/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:1a8a4ffc6c6e6e577f2bfbcfebf7cffbb310283ba2532c729570a0a752cebaaa", size = 1144069, upload-time = "2026-09-28T05:20:44.705Z" }, + { url = "https://files.pythonhosted.org/packages/4a/f4/fdd7a093fb2fb3a0ce24256c92ad5bf67936f4669bd06d5ceae4298150ec/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b987d73eba95183a7e59cca6d1925c588aa4307922d9852cc2b4d282a2ae4128", size = 1146692, upload-time = "2026-09-28T05:20:25.83Z" }, + { url = "https://files.pythonhosted.org/packages/f1/c8/6dbd4377e935505ae4fc8ee4c7b18c69994c4bbaca015447c470218bdbee/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d4569c39bd97d9573e946429ed676f3b55a7c8ed80920addd67d384a47d59b38", size = 1189480, upload-time = "2026-09-28T05:20:46.734Z" }, + { url = "https://files.pythonhosted.org/packages/ba/bf/ff288b496b690000d2686dcfa7f67855e4c5a6dddb464ed8e4880be83bb2/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6042b8707a4b25bbbfe20b110258fb7549a68951e5525608a8ce06091039e9fc", size = 1167109, upload-time = "2026-09-28T05:19:24.758Z" }, + { url = "https://files.pythonhosted.org/packages/d9/d1/3fee2bc29fc747fabbb27e174515592e390809ca465d3bbbabbcfa3235a8/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_31_riscv64.whl", hash = "sha256:f413629de94d38a7a2ad259697a6752526143cb49c2e2d7bdba19381c80693f6", size = 1126823, upload-time = "2026-09-28T05:18:52.636Z" }, + { url = "https://files.pythonhosted.org/packages/58/6a/585294392fa6a6d9719446a042a777ba98ecca265d9e48cbedc10c88df0f/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:f071737e4e775bebba07e319eb645a880d1e0186b4d24bad23255e849d86d483", size = 1155773, upload-time = "2026-09-28T05:20:51.268Z" }, + { url = "https://files.pythonhosted.org/packages/cd/4d/fa283ff79996debf1ff08593e01f8b773eb8211b19b3d6d5d491317a65bf/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_aarch64.whl", hash = "sha256:f59a3912858d0609c26054aa1c474847c797937f5e0e1c77f0d3f08032e50f09", size = 1296671, upload-time = "2026-09-28T05:18:46.937Z" }, + { url = "https://files.pythonhosted.org/packages/42/05/98f9c2f628da5afabb6c3e8df5d2d299d2d50bf306b14a6e26d3954c4951/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_armv7l.whl", hash = "sha256:f09c05a23a8025dd5cad07c2ed299a46778e72d49ff0dd5bf353bcc0f7dc1580", size = 1422001, upload-time = "2026-09-28T05:19:04.188Z" }, + { url = "https://files.pythonhosted.org/packages/b8/c9/f2109ade29a7ec284d55bca328d784743e8b13e321576fd34ee7e65f99af/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_i686.whl", hash = "sha256:44ace770bda3a0301739fc5a413c764de790c739df1f1a7218dd049f8594d9f5", size = 1374182, upload-time = "2026-09-28T05:20:27.764Z" }, + { url = "https://files.pythonhosted.org/packages/ce/6d/e05d5f72441564a3bebc71fa155deafd0fd3b015d6014ec8a00edf6e42bb/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_ppc64le.whl", hash = "sha256:03b131043608f94a2578a2a896a1a72092acb7079815eef5b8513b08447e0e62", size = 1278402, upload-time = "2026-09-28T05:20:56.332Z" }, + { url = "https://files.pythonhosted.org/packages/96/d7/6988a7f1f69c5c530687dce03a32c157e2d878e3a2e32af9269ce016b78a/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_riscv64.whl", hash = "sha256:e9784aca26eddfe99b03a0292320db742cc8a74200ef864fdce949f526f973cd", size = 1297770, upload-time = "2026-09-28T05:19:11.387Z" }, + { url = "https://files.pythonhosted.org/packages/a2/28/5bb82b60b836bd2329e1fe01ad94efc2b14dfa806f8e50ed14b795d78770/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_x86_64.whl", hash = "sha256:2b52ac363096232bebc2add117e9178d91f4248f4cdb919fd1026b5f86a4bb16", size = 1333977, upload-time = "2026-09-28T05:20:02.483Z" }, + { url = "https://files.pythonhosted.org/packages/8e/7d/d419841b8f65481ea1a50c4ba36f2670d48b8c7a51e4e8586089564cf7f6/hypothesis-6.168.3-cp315-abi3.abi3t-win32.whl", hash = "sha256:d28e3a6b511a74ce37df5274b51f36c2b274fea7365e00b16f0c81e22acd5957", size = 675394, upload-time = "2026-09-28T05:19:55.237Z" }, + { url = "https://files.pythonhosted.org/packages/d8/d2/1de6a2ad100e44817f2e9a8e8ba3eeaf4769621f4f87f2d4966c4971fcc4/hypothesis-6.168.3-cp315-abi3.abi3t-win_amd64.whl", hash = "sha256:7b9638789548361a57d984f56619ac694a914c328181d911271618409261ff4a", size = 681699, upload-time = "2026-09-28T05:19:09.957Z" }, + { url = "https://files.pythonhosted.org/packages/c0/79/be3370fd02734d6b1d950183ae9224580d19fa8d356b95348d60f1e3eeb7/hypothesis-6.168.3-cp315-abi3.abi3t-win_arm64.whl", hash = "sha256:65d78e4357ec48ed2c67825f06740ee3599be4cfe770a6092bed07108d679ac5", size = 679849, upload-time = "2026-09-28T05:19:06.884Z" }, + { url = "https://files.pythonhosted.org/packages/3f/0e/27960db1d45a94497e1681fb9d7b9140b44da4e655ba20f0d82f95d10230/hypothesis-6.168.3-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:101531b4ccf7fa12965a6b295d10b64c2f8a8fab61e1f9b18c55335d8fb02575", size = 793135, upload-time = "2026-09-28T05:19:14.056Z" }, + { url = "https://files.pythonhosted.org/packages/9a/cb/be08379286c1925fcd651b92ed835f470a83b859b5c484b3431da9d851f5/hypothesis-6.168.3-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:41e0120814de5c3a6b58d8cb27cb11b8873ab33fba130cd09855a27f5a12a1ca", size = 788995, upload-time = "2026-09-28T05:20:34.164Z" }, + { url = "https://files.pythonhosted.org/packages/38/7b/6d885d139b2ed76e2f4908cf10eae56acc73bf05cbef286357c4559c9a04/hypothesis-6.168.3-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9d120009d145697909f2f0c532f2d568558ddd6990707f8cba6576bfb9770696", size = 1124974, upload-time = "2026-09-28T05:19:49.903Z" }, + { url = "https://files.pythonhosted.org/packages/05/cc/b302ae809b48254830cae57801d0af0b66783df515b5dd56b9f21c8dfa86/hypothesis-6.168.3-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:818b3d09bba60ef90463e47a4bc87275dba1e8aeb184d3102888440ca7c7837d", size = 1171967, upload-time = "2026-09-28T05:20:12.295Z" }, + { url = "https://files.pythonhosted.org/packages/2c/ab/51031080a6600f5b07adf5f9dfa117345c2adda07a5c24bda523e70bd1e9/hypothesis-6.168.3-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:9e3c36eba80e089ecfa28ec08049723223cb41d3cfe4d9988169cf98440e402b", size = 685641, upload-time = "2026-09-28T05:19:44.609Z" }, +] + [[package]] name = "idna" version = "3.11" @@ -1324,7 +1412,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess" }, + { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [ @@ -1364,7 +1452,7 @@ wheels = [ [[package]] name = "policyengine-core" -version = "3.32.9" +version = "3.32.11" source = { editable = "." } dependencies = [ { name = "dpath" }, @@ -1390,6 +1478,7 @@ dev = [ { name = "build" }, { name = "coverage" }, { name = "furo" }, + { name = "hypothesis" }, { name = "jupyter-book" }, { name = "markupsafe" }, { name = "mypy" }, @@ -1414,6 +1503,7 @@ requires-dist = [ { name = "furo", marker = "extra == 'dev'" }, { name = "h5py", specifier = ">=3,<4" }, { name = "huggingface-hub", specifier = ">=0.25.1" }, + { name = "hypothesis", marker = "extra == 'dev'", specifier = ">=6,<7" }, { name = "ipython", specifier = ">=8,<9" }, { name = "jupyter-book", marker = "extra == 'dev'", specifier = "<1" }, { name = "markupsafe", marker = "extra == 'dev'", specifier = "==2.0.1" }, From e8fe03e64cefaf859425d8908781a8faf117cc76 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 1 Oct 2026 15:06:09 -0400 Subject: [PATCH 02/42] Start a new store history when subsample rebuilds the simulation Co-Authored-By: Claude Opus 5.5 --- policyengine_core/simulations/simulation.py | 5 ++- tests/core/test_branch_input_invalidation.py | 33 ++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 4c0a0ab89..88cde4b85 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1920,8 +1920,11 @@ def subsample( df = subset_df - # Update the dataset and rebuild the simulation + # Update the dataset and rebuild the simulation. Nothing stored before + # survives the rebuild, so start a new store history before the + # rebuilt inputs are recorded in it. self.dataset = Dataset.from_dataframe(df, self.dataset.time_period) + self._store_history = StoreHistory() self.build_from_dataset() # Purge ``_fast_cache`` entries populated by ``to_input_dataframe`` diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index 1827a4bbc..dde6d5390 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -41,6 +41,7 @@ from policyengine_core.variables import Variable JANUARY = periods.period("2017-01") +JANUARY_2022 = periods.period("2022-01") SITUATION = { "persons": { @@ -324,6 +325,38 @@ def apply(self): 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 ----- # From 56f19b14d5fad3c24dd7805a6571acd19d0623dc Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 1 Oct 2026 15:18:21 -0400 Subject: [PATCH 03/42] Record values served from the macro cache or a spiral's default Neither is stored, but values calculated from them are, so set_input on a branch must see them in the store history. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/simulations/simulation.py | 5 ++ tests/core/test_branch_input_invalidation.py | 66 ++++++++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 88cde4b85..c3e80f646 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -734,6 +734,9 @@ 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``). + holder._record_store(period, next_sequence_number()) return value if variable.requires_computation_after is not None: @@ -880,6 +883,8 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: 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() diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index dde6d5390..db93497e4 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -259,6 +259,49 @@ def test_disk_storage_drops_calculated_values_and_keeps_inputs(tax_benefit_syste ) +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 ----- # @@ -498,6 +541,11 @@ def _monthly(name, formula=None, **attributes): lambda person, period: 3 * person("p_a", period) + person("p_switch", period), ), _yearly("p_inner", _inner_branch_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 + ), ] @@ -619,6 +667,24 @@ def test_value_a_holder_does_not_keep_is_tracked(): 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(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({**ROOT_INPUTS, ("p_spiral", str(deepest)): (100.0,) * 3}) + assert np.array_equal( + branch.calculate("p_spiral", "2015"), fresh.calculate("p_spiral", "2015") + ) + + # ----- Property: branch results do not depend on what was calculated before ----- # SETTABLE = [ From 3499695fd4235100927381f78c9cdd767494fb89 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 1 Oct 2026 18:25:00 -0400 Subject: [PATCH 04/42] Move the Hypothesis property into its own module that skips without Hypothesis The country-package smoke job installs no dev dependencies, so the module-level hypothesis import failed its collection. The synthetic system the examples and the property share moves to tests/fixtures/branch_input_invalidation.py. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_branch_input_invalidation.py | 395 ++---------------- ...test_branch_input_invalidation_property.py | 224 ++++++++++ tests/fixtures/branch_input_invalidation.py | 160 +++++++ 3 files changed, 409 insertions(+), 370 deletions(-) create mode 100644 tests/core/test_branch_input_invalidation_property.py create mode 100644 tests/fixtures/branch_input_invalidation.py diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index db93497e4..0a2806c00 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -12,33 +12,33 @@ the first time anything in the family stored, or derived, the overridden variable for an overlapping period (``StoreHistory.earliest_dependency``). -The invariant tested throughout (and by the property test at the end, over -random sequences of calculations, branches and inputs): 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. +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 hypothesis import HealthCheck, given, settings -from hypothesis import strategies as st from policyengine_core import periods -from policyengine_core.country_template import CountryTaxBenefitSystem, entities 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.holders import set_input_divide_by_period from policyengine_core.model_api import Reform from policyengine_core.simulations import SimulationBuilder -from policyengine_core.variables import Variable +from tests.fixtures.branch_input_invalidation import ( + CARRY_OVER_SYSTEM, + ROOT_INPUTS, + synthetic_simulation, +) JANUARY = periods.period("2017-01") JANUARY_2022 = periods.period("2022-01") @@ -450,167 +450,17 @@ def test_history_matches_overlapping_periods_and_derived_values(): assert history.earliest_dependency("e", periods.period("1999-01")) == 5 -# ----- A synthetic system for derived values and branches inside formulas ----- # - - -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 _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), - # 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) - return system - - -# Carry-over is left out of the random programs below: 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"] - - -def _synthetic(inputs, memory_config=None, system=SYNTHETIC_SYSTEM): - """A simulation of ``PEOPLE`` people given ``inputs`` before anything else.""" - simulation = SimulationBuilder().build_default_simulation(system, count=PEOPLE) - if memory_config is not None: - 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: - simulation.get_holder(variable) - 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) - }, -} +# ----- Derived values and branches inside formulas, on a synthetic system ----- # def test_value_carried_over_from_an_earlier_input_is_dropped(): - simulation = _synthetic(ROOT_INPUTS, system=CARRY_OVER_SYSTEM) + 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( + fresh = synthetic_simulation( {**ROOT_INPUTS, ("p_c", "2014"): (70.0, 80.0, 90.0)}, system=CARRY_OVER_SYSTEM, ) @@ -622,13 +472,13 @@ def test_value_carried_over_from_an_earlier_input_is_dropped(): def test_uprated_value_is_dropped(): - simulation = _synthetic(ROOT_INPUTS) + 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({**ROOT_INPUTS, ("p_up", "2014"): (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") ) @@ -641,14 +491,16 @@ def test_value_calculated_only_inside_another_branch_is_tracked(): so the parent and its later branches never store ``p_inner_only``. The history the whole family shares still records it. """ - simulation = _synthetic(ROOT_INPUTS) + 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({**ROOT_INPUTS, ("p_inner_only", "2013"): (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") ) @@ -657,7 +509,7 @@ def test_value_calculated_only_inside_another_branch_is_tracked(): 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(ROOT_INPUTS, memory_config=config) + simulation = synthetic_simulation(ROOT_INPUTS, memory_config=config) simulation.calculate("p_prod", "2013") assert not simulation.get_holder("p_sum")._memory_storage._arrays @@ -669,7 +521,7 @@ def test_value_a_holder_does_not_keep_is_tracked(): def test_value_a_spiral_defaults_is_tracked(): """The default a spiral returns is not stored, but what reads it is.""" - simulation = _synthetic(ROOT_INPUTS) + simulation = synthetic_simulation(ROOT_INPUTS) branch = simulation.get_branch("branch") branch.calculate("p_spiral", "2015") deepest = min( @@ -679,206 +531,9 @@ def test_value_a_spiral_defaults_is_tracked(): branch.set_input("p_spiral", deepest, np.array([100.0, 100.0, 100.0])) - fresh = _synthetic({**ROOT_INPUTS, ("p_spiral", str(deepest)): (100.0,) * 3}) + 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") ) - - -# ----- Property: branch results do not depend on what was calculated before ----- # - -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_up", - "p_c", -] -MONTHS = ["2013-01", "2013-07", "2014-12", "2015-03"] -MEMORY_CONFIGS = { - "memory": lambda: None, - "disk": lambda: MemoryConfig(max_memory_occupation=0), - "not_kept": lambda: MemoryConfig( - max_memory_occupation=1, variables_to_drop=["p_sum", "p_inner_only"] - ), -} - -# 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. -simulation_index = st.integers(0, 3) -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), - st.tuples(st.just("drop"), 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, 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_month", "2013-07", "p_m", "2013"), - ("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]), - ("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]), - ("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]), - ] - ) -) -operations = st.lists( - st.one_of( - single_operation.map(lambda operation: [operation]), - override_after_calculating, - ), - min_size=1, - max_size=15, -).map(lambda chunks: [operation for chunk in chunks for operation in chunk]) - - -def _stored_periods(variable, period): - """The periods ``set_input`` stores ``variable`` at, given ``period``.""" - period = periods.period(period) - if SYNTHETIC_SYSTEM.get_variable(variable).definition_period == periods.MONTH: - return [str(month) for month in period.get_subperiods(periods.MONTH)] - return [str(period)] - - -def _own_inputs(simulation, variable, period): - """The values stored on the branch itself for ``variable`` over ``period``.""" - holder = simulation.get_holder(variable) - return { - (variable, stored_period): tuple( - holder._get_array_from_storage(stored_period, simulation.branch_name) - ) - for stored_period in _stored_periods(variable, period) - } - - -def _run(program, memory_config): - """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, as stored (an input for a - year is divided between the months the branch has not set itself). - """ - root = _synthetic(ROOT_INPUTS, memory_config=memory_config) - simulations = [root] - inputs = [dict(ROOT_INPUTS)] - own_inputs = [set()] - 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": - simulations.append(simulation.get_branch(f"b{len(simulations)}")) - inputs.append(dict(inputs[index])) - own_inputs.append(set()) - elif kind == "set": - if index == 0: - continue # Inputs on the root after calculating are out of scope. - _, _, variable, period, value = operation - stored = {(variable, p) for p in _stored_periods(variable, period)} - if len(stored) > 1 and stored <= own_inputs[index]: - # Dividing a year between months the branch has all set - # itself is an error unless the totals match. - continue - simulation.set_input(variable, period, np.asarray(value)) - stored = _own_inputs(simulation, variable, period) - inputs[index].update(stored) - own_inputs[index].update(stored) - elif kind == "calculate": - _, _, variable, period = operation - value = np.array(simulation.calculate(variable, period), copy=True) - results.append((dict(inputs[index]), variable, period, value)) - else: - simulation.drop_computed_arrays() - return results - - -@settings( - max_examples=500, - deadline=None, - suppress_health_check=[HealthCheck.too_slow], -) -@given(program=operations, mode=st.sampled_from(sorted(MEMORY_CONFIGS))) -def test_branch_calculations_match_a_simulation_given_its_inputs_first(program, mode): - for branch_inputs, variable, period, value in _run(program, MEMORY_CONFIGS[mode]()): - expected = _synthetic(branch_inputs).calculate(variable, period) - # Uprating chains may round differently in float32 depending on which - # earlier periods were calculated first. - np.testing.assert_allclose( - value, expected, rtol=1e-5, err_msg=f"{variable} {period}" - ) 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..f4da315c8 --- /dev/null +++ b/tests/core/test_branch_input_invalidation_property.py @@ -0,0 +1,224 @@ +"""Branch results do not depend on what was calculated before. + +Random sequences of calculations, branches, inputs and +``drop_computed_arrays`` run on a synthetic system, with values held in +memory, on disk, or not kept. 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. +``test_branch_input_invalidation.py`` pins the same behaviour with examples. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +# The country-package smoke job installs no dev dependencies. +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from policyengine_core import periods +from policyengine_core.experimental import MemoryConfig +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_up", + "p_c", +] +MONTHS = ["2013-01", "2013-07", "2014-12", "2015-03"] +MEMORY_CONFIGS = { + "memory": lambda: None, + "disk": lambda: MemoryConfig(max_memory_occupation=0), + "not_kept": lambda: MemoryConfig( + max_memory_occupation=1, variables_to_drop=["p_sum", "p_inner_only"] + ), +} + +# 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. +simulation_index = st.integers(0, 3) +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), + st.tuples(st.just("drop"), 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, 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_month", "2013-07", "p_m", "2013"), + ("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]), + ("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]), + ("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]), + ] + ) +) +operations = st.lists( + st.one_of( + single_operation.map(lambda operation: [operation]), + override_after_calculating, + ), + min_size=1, + max_size=15, +).map(lambda chunks: [operation for chunk in chunks for operation in chunk]) + + +def _stored_periods(variable, period): + """The periods ``set_input`` stores ``variable`` at, given ``period``.""" + period = periods.period(period) + if SYNTHETIC_SYSTEM.get_variable(variable).definition_period == periods.MONTH: + return [str(month) for month in period.get_subperiods(periods.MONTH)] + return [str(period)] + + +def _own_inputs(simulation, variable, period): + """The values stored on the branch itself for ``variable`` over ``period``.""" + holder = simulation.get_holder(variable) + return { + (variable, stored_period): tuple( + holder._get_array_from_storage(stored_period, simulation.branch_name) + ) + for stored_period in _stored_periods(variable, period) + } + + +def _run(program, memory_config): + """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, as stored (an input for a + year is divided between the months the branch has not set itself). + """ + root = synthetic_simulation(ROOT_INPUTS, memory_config=memory_config) + simulations = [root] + inputs = [dict(ROOT_INPUTS)] + own_inputs = [set()] + 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": + simulations.append(simulation.get_branch(f"b{len(simulations)}")) + inputs.append(dict(inputs[index])) + own_inputs.append(set()) + elif kind == "set": + if index == 0: + continue # Inputs on the root after calculating are out of scope. + _, _, variable, period, value = operation + stored = {(variable, p) for p in _stored_periods(variable, period)} + if len(stored) > 1 and stored <= own_inputs[index]: + # Dividing a year between months the branch has all set + # itself is an error unless the totals match. + continue + simulation.set_input(variable, period, np.asarray(value)) + stored = _own_inputs(simulation, variable, period) + inputs[index].update(stored) + own_inputs[index].update(stored) + elif kind == "calculate": + _, _, variable, period = operation + value = np.array(simulation.calculate(variable, period), copy=True) + results.append((dict(inputs[index]), variable, period, value)) + else: + simulation.drop_computed_arrays() + return results + + +@hypothesis.settings( + max_examples=500, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given(program=operations, mode=st.sampled_from(sorted(MEMORY_CONFIGS))) +def test_branch_calculations_match_a_simulation_given_its_inputs_first(program, mode): + for branch_inputs, variable, period, value in _run(program, MEMORY_CONFIGS[mode]()): + expected = synthetic_simulation(branch_inputs).calculate(variable, period) + # Uprating chains may round differently in float32 depending on which + # earlier periods were calculated first. + np.testing.assert_allclose( + value, expected, rtol=1e-5, err_msg=f"{variable} {period}" + ) diff --git a/tests/fixtures/branch_input_invalidation.py b/tests/fixtures/branch_input_invalidation.py new file mode 100644 index 000000000..5fe035344 --- /dev/null +++ b/tests/fixtures/branch_input_invalidation.py @@ -0,0 +1,160 @@ +"""A synthetic system for derived values and branches inside formulas.""" + +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 _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), + # 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) + 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"] + + +def synthetic_simulation(inputs, memory_config=None, system=SYNTHETIC_SYSTEM): + """A simulation of ``PEOPLE`` people given ``inputs`` before anything else.""" + simulation = SimulationBuilder().build_default_simulation(system, count=PEOPLE) + if memory_config is not None: + 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: + simulation.get_holder(variable) + 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) + }, +} From f826d57a9bc46354b3d8f1b52d0d2cd181091024 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 1 Oct 2026 18:41:11 -0400 Subject: [PATCH 05/42] Keep a store history per simulation; fix what two reviews found Each simulation now keeps its own store history: a branch copies its parent's, and a formula that calculates in another simulation takes in that simulation's history when calculate returns (fast-cache hits included). A drop prunes the records it made obsolete and records the kept inputs again. An unrelated store in a sibling branch, or in the other arm of a reform/baseline pair, no longer lowers a branch's threshold, so MTR-style and baseline branches calculate shared values once. Also from the soundness and code reviews: - a branch whose input dropped values stops reading macro-cache files (they are keyed by branch and period, not inputs); drop_computed_arrays too; - requires_computation_after counts a prerequisite requested before a drop; - values a custom set_input handler calculates are not inputs; - disk storage writes a new file per store, so a recalculation no longer changes what child branches or same-named branches read; - a value written into storage without a number counts as a dependency; - storages pickled before numbering still store (__setstate__); - derivative keeps inputs set after the simulation was built; - property test: an independent model of divided inputs, reused and forgotten branch names, a blacklist mode, derandomized CI examples with an environment soak; examples for each finding; - docs and docstrings say what the code does, with the remaining limits. Co-Authored-By: Claude Opus 5.5 --- .../branch-set-input-invalidation.changed.md | 1 + .../branch-set-input-invalidation.fixed.md | 2 +- docs/usage/branches.md | 96 +++-- .../data_storage/in_memory_storage.py | 23 +- .../data_storage/on_disk_storage.py | 44 ++- .../data_storage/store_history.py | 107 ++++-- policyengine_core/holders/holder.py | 34 +- policyengine_core/simulations/simulation.py | 124 ++++-- tests/core/test_branch_input_invalidation.py | 354 +++++++++++++++++- ...test_branch_input_invalidation_property.py | 139 ++++--- .../core/test_branch_scoped_delete_arrays.py | 6 +- tests/fixtures/branch_input_invalidation.py | 94 ++++- 12 files changed, 859 insertions(+), 165 deletions(-) create mode 100644 changelog.d/branch-set-input-invalidation.changed.md 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..8273a4824 --- /dev/null +++ b/changelog.d/branch-set-input-invalidation.changed.md @@ -0,0 +1 @@ +Disk storage writes a new file for every store, so recalculating a value no longer changes what branches sharing the directory read; apply_reform keeps inputs by a flag on each stored array instead of replaying _user_input_keys, and values a custom set_input handler calculates are no longer treated as inputs. diff --git a/changelog.d/branch-set-input-invalidation.fixed.md b/changelog.d/branch-set-input-invalidation.fixed.md index b7916d200..26adbe967 100644 --- a/changelog.d/branch-set-input-invalidation.fixed.md +++ b/changelog.d/branch-set-input-invalidation.fixed.md @@ -1 +1 @@ -Simulation.set_input on a branch now drops the values the branch holds that may have been calculated from the value it replaces (tracked with a sequence number on every stored array), so a value the parent calculated before branching no longer shadows the branch's input. +Simulation.set_input on a branch now drops the values the branch holds that may have been calculated from the value it replaces (every stored array carries a sequence number, and each simulation records what its values may have been calculated from), so a value the parent calculated before branching no longer shadows the branch's input. Simulation.derivative keeps inputs set after the simulation was built. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 6bfc095a3..ee39b8840 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -13,11 +13,12 @@ tax_if_itemizing = branch.calculate("income_tax", 2026) ## What a branch starts with -`get_branch` returns a branch holding every value its parent holds at that -moment: inputs and calculated values alike. What the parent stores after that -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. +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 @@ -29,33 +30,47 @@ 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. -A simulation and all its branches also share a record of the first time each -variable was stored (or calculated without being kept) for each period, and of -the first time one of its values was uprated or carried over from another -period. 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 of: - -- the first time the variable was stored for any period that shares a day with - `period`, anywhere in the family; and -- the first time a value of the variable was uprated or carried over. - -If neither has happened, the branch drops nothing. That is the usual case: a +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; +- 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 (a formula that branches, sets an input and + calculates in the branch hands back values calculated there). + +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. + +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, so nothing has stored that variable yet. +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 without going through `set_input` or a +calculation 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 has been -stored, so a calculated value carries a larger number than each stored value it -was calculated from, directly or through other calculated values. A value that -depends on the overridden variable at an overlapping period was therefore -stored after the first such value existed. A value derived from another period -of the same variable (uprating, carry-over) depends on which periods hold -values at all, which is why the first such derivation also counts. The record -covers the whole family because a formula can calculate in one branch and -return the result to another (`branch.calculate` inside a formula), and -because a branch may be deleted after its results were stored elsewhere. +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, or handed +back by another simulation's `calculate`, and in each case what it was +calculated from is in the simulation's history. So any value that depends on +the overridden variable at an overlapping period was stored after the earliest +record of it. Uprating and carry-over read which periods hold values at all, +which is why the first uprated or carried-over value also counts. 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: @@ -68,23 +83,32 @@ depend on the input is calculated again) but not less, within these limits: 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 test what is stored.** A formula that reads - `holder.get_known_periods()` or `simulation.get_array()` and acts on whether - a value exists depends on that value without its number being recorded. -- **Only the family is tracked.** A value calculated in a separate simulation - (one not created with `get_branch`) and passed in is not covered. +- **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. +- **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 drop nothing.** `set_input` on a simulation that is not a branch keeps its earlier behaviour: values it already calculated stay. - **Existing child branches keep their values.** An input set on a branch does not reach branches already created from it. +Two related behaviours: a branch whose input dropped values stops reading +macro-cache files, which are keyed by branch and period but not by inputs; and +`requires_computation_after` is satisfied by a prerequisite requested before a +drop removed its values. + ## Dropping calculated values `simulation.drop_computed_arrays()` deletes every value the simulation holds -except inputs (the dataset or situation it was built from, and values set with -`set_input` on it or on the simulations it was created from), and returns how -many arrays it deleted. Use it on a branch after changing its policy: +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 a custom `set_input` handler calculates are not inputs. Use it on a +branch after changing its policy: ```python branch = simulation.get_branch("pre_reform_rules", clone_system=True) diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 0c82d6a9e..8769fd1e1 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional, Set, Union +from typing import Dict, List, Optional, Set, Tuple, Union import numpy from numpy.typing import ArrayLike @@ -25,6 +25,12 @@ def __init__(self, is_eternal: bool): self._input_keys: Set[str] = set() self.is_eternal = is_eternal + def __setstate__(self, state: dict) -> None: + # Storages pickled before stores were numbered have neither record. + state.setdefault("_sequence_numbers", {}) + state.setdefault("_input_keys", set()) + self.__dict__.update(state) + def clone(self) -> "InMemoryStorage": clone = InMemoryStorage(self.is_eternal) clone._arrays = {period: array.copy() for period, array in self._arrays.items()} @@ -84,7 +90,7 @@ def put( else: self._input_keys.discard(key) - def drop_computed(self, since: Optional[int] = None) -> int: + 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 @@ -103,6 +109,19 @@ def drop_computed(self, since: Optional[int] = None) -> int: self._sequence_numbers.pop(key, None) 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).""" + return [ + (periods.period(key.split(":", 1)[1]), self._sequence_numbers[key]) + for key in self._input_keys + if key in self._sequence_numbers + and (since is None or self._sequence_numbers[key] >= since) + ] + + 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_deleted_keys(self) -> None: self._sequence_numbers = { key: number diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 0977d0bd9..5411ed70e 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -1,6 +1,6 @@ import os import shutil -from typing import Dict, Optional, Set +from typing import Dict, List, Optional, Set, Tuple import numpy from numpy.typing import ArrayLike @@ -32,6 +32,12 @@ def __init__( self.preserve_storage_dir = preserve_storage_dir self.storage_dir = storage_dir + def __setstate__(self, state: dict) -> None: + # Storages pickled before stores were numbered have neither record. + state.setdefault("_sequence_numbers", {}) + state.setdefault("_input_keys", set()) + self.__dict__.update(state) + def clone(self) -> "OnDiskStorage": """Create a private metadata view over this storage directory. @@ -89,21 +95,23 @@ def put( period = periods.period(period) filename = f"{branch_name}_{period}" - path = os.path.join(self.storage_dir, filename) + ".npy" + if sequence_number is None: + sequence_number = next_sequence_number() + # A new file for every store: clones share this directory and may + # still map an earlier file for the same key. + path = os.path.join(self.storage_dir, f"{filename}.{sequence_number}") + ".npy" if isinstance(value, EnumArray): self._enums[path] = value.possible_values value = value.view(numpy.ndarray) numpy.save(path, value) self._files[filename] = path - self._sequence_numbers[filename] = ( - next_sequence_number() if sequence_number is None else sequence_number - ) + self._sequence_numbers[filename] = sequence_number if is_input: self._input_keys.add(filename) else: self._input_keys.discard(filename) - def drop_computed(self, since: Optional[int] = None) -> int: + 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 @@ -121,6 +129,20 @@ def drop_computed(self, since: Optional[int] = None) -> int: self._sequence_numbers.pop(key, None) return len(dropped) + def inputs_since(self, since: Optional[int] = None) -> List[Tuple[Period, int]]: + """As :meth:`InMemoryStorage.inputs_since`.""" + # Period strings contain no "_"; branch names may. + return [ + (periods.period(key.rsplit("_", 1)[1]), self._sequence_numbers[key]) + for key in self._input_keys + if key in self._sequence_numbers + and (since is None or self._sequence_numbers[key] >= since) + ] + + 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_deleted_keys(self) -> None: self._sequence_numbers = { key: number @@ -169,12 +191,20 @@ def restore(self) -> None: self._sequence_numbers = {} self._input_keys = set() # Restore self._files from content of storage_dir. + latest = {} for filename in os.listdir(self.storage_dir): if not filename.endswith(".npy"): continue path = os.path.join(self.storage_dir, filename) filename_core = filename.rsplit(".", 1)[0] - files[filename_core] = path + # Files are named "..npy" (each store writes + # a new file); older dumps are ".npy". Keep each key's latest. + key, _, number = filename_core.rpartition(".") + if not (key and number.isdigit()): + key, number = filename_core, "0" + if key not in latest or int(number) > latest[key]: + latest[key] = int(number) + files[key] = path def __del__(self) -> None: if self.preserve_storage_dir: diff --git a/policyengine_core/data_storage/store_history.py b/policyengine_core/data_storage/store_history.py index a1725fdcd..bdae4491f 100644 --- a/policyengine_core/data_storage/store_history.py +++ b/policyengine_core/data_storage/store_history.py @@ -18,6 +18,7 @@ from policyengine_core.periods import Period _sequence = itertools.count(1) +_history_ids = itertools.count(1) def next_sequence_number() -> int: @@ -33,47 +34,109 @@ def periods_overlap(first: Period, second: Period) -> bool: class StoreHistory: - """When each variable's values were first stored in a family of simulations. + """What the values a simulation holds may have been calculated from. - A simulation and the branches created from it (``get_branch``, at any - depth) share one history. For each variable it records: + 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 anywhere in the family, 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 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 a value). Such a value depends on which other periods - hold values, so an input set for any period can change it. - - Deleting or replacing a stored value does not change the history: a value - calculated from the deleted or replaced one keeps its number. + 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 a value). Such a value depends on which other periods hold + values, so an input set for any period can change it. + + 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] = {} + # Changes on every new record, so a merge can skip a history it has + # already taken in unchanged. + self._id = next(_history_ids) + self._version = 0 + self._merged_versions: Dict[int, int] = {} + + 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._merged_versions = dict(self._merged_versions) + return new def record_store(self, variable_name: str, period: Period, sequence_number: int): stored = self._first_stored.setdefault(variable_name, {}) - if period not in stored: + recorded = stored.get(period) + if recorded is None or sequence_number < recorded: stored[period] = sequence_number + self._version += 1 def record_derived(self, variable_name: str, sequence_number: int): - if variable_name not in self._first_derived: + recorded = self._first_derived.get(variable_name) + if recorded is None or sequence_number < recorded: self._first_derived[variable_name] = sequence_number + self._version += 1 + + def merge(self, other: "StoreHistory") -> None: + """Take in ``other``'s records, keeping the earlier number of each.""" + if other is self or self._merged_versions.get(other._id) == other._version: + return + for variable_name, stored in other._first_stored.items(): + for period, sequence_number in stored.items(): + self.record_store(variable_name, period, sequence_number) + for variable_name, sequence_number in other._first_derived.items(): + self.record_derived(variable_name, sequence_number) + self._merged_versions[other._id] = other._version + + 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 + } + self._version += 1 + # Merging a history again must restore what was forgotten here. + self._merged_versions = {} 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 first store of the variable for any period that - shares a day with ``period``, or the first time one of its values was - derived from its other periods, whichever came first. ``None`` means - nothing in the family has read or stored the variable for an - overlapping period nor derived a value of it, so no stored value can - depend on its value at ``period``. + This is the earliest recorded store of the variable for any period + that shares a day with ``period``, or the earliest recorded derivation + of one of its values from other periods, whichever came first. + ``None`` means no value the simulation holds can depend on the + variable's value at ``period``. """ candidates = [ sequence_number diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index c686ea286..a5e42fe2f 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -1,6 +1,6 @@ import os import warnings -from typing import TYPE_CHECKING, Any, List, Tuple +from typing import TYPE_CHECKING, Any, List, Optional, Tuple import numpy import psutil @@ -99,18 +99,31 @@ def delete_arrays( if self._disk_storage: self._disk_storage.delete(period, branch_name) - def drop_computed(self, since: int = None) -> int: + 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. """ - dropped = self._memory_storage.drop_computed(since) + dropped = self._memory_storage.drop_computed(since=since) if self._disk_storage is not None: - dropped += self._disk_storage.drop_computed(since) + 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 _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: @@ -351,10 +364,16 @@ def _set( value: ArrayLike, branch_name: str = "default", validate_nan: bool = False, + is_input: Optional[bool] = None, ) -> None: simulation = getattr(self, "simulation", None) user_input_contexts = getattr(simulation, "_user_input_contexts", None) - if user_input_contexts and branch_name == "default": + # A value is an input when stored while ``set_input`` runs, unless the + # caller says otherwise: ``put_in_cache`` stores calculated values, + # including those a custom ``set_input`` handler calculates. + if is_input is None: + is_input = bool(user_input_contexts) + if is_input and user_input_contexts and branch_name == "default": branch_name = user_input_contexts[-1] value = self._to_array(value, validate_nan=validate_nan) if self.variable.definition_period != periods.ETERNITY: @@ -371,7 +390,6 @@ def _set( ) sequence_number = next_sequence_number() - is_input = bool(user_input_contexts) storage = self._disk_storage if should_store_on_disk else self._memory_storage storage.put( value, @@ -381,7 +399,7 @@ def _set( is_input=is_input, ) self._record_store(period, sequence_number) - if user_input_contexts: + if is_input and simulation is not None: if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() simulation._user_input_keys.add((self.variable.name, branch_name, period)) @@ -398,7 +416,7 @@ def put_in_cache( self._record_store(period, next_sequence_number()) return - self._set(period, value, branch_name) + self._set(period, value, branch_name, is_input=False) def _record_store(self, period: Period, sequence_number: int) -> None: simulation = getattr(self, "simulation", None) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index c3e80f646..dc28b5a56 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1,5 +1,6 @@ import hashlib import tempfile +from contextvars import ContextVar from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union import numpy as np @@ -91,6 +92,16 @@ def _uprating_index_value(parameter, instant) -> Optional[float]: ) +# 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 +) + + class Simulation: """ Represents a simulation, and handles the calculation logic @@ -558,6 +569,7 @@ 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() return _cached self.tracer.record_calculation_start(variable_name, period, self.branch_name) @@ -577,6 +589,7 @@ def calculate( variable_name ).entity.key result = self.map_result(result, source_entity, map_to) + self._share_store_history_with_caller() return result finally: self.tracer.record_calculation_end() @@ -692,6 +705,7 @@ 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.") + self._get_requested_variables().add(variable_name) population = self.get_variable_population(variable_name) holder = population.get_holder(variable_name) variable = self.tax_benefit_system.get_variable( @@ -747,7 +761,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}" ) @@ -789,7 +812,11 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # 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: @@ -1309,9 +1336,11 @@ 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, and inputs set on it or, - for a branch, on the simulations it was created from. Every other - value is calculated again when next requested. + 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; @@ -1321,6 +1350,8 @@ def drop_computed_arrays(self) -> int: Returns: int: The number of arrays deleted. """ + # Recalculate rather than read macro-cache files written before. + self.macro_cache_read = False return self._drop_computed() def get_known_periods(self, variable: str) -> List[Period]: @@ -1361,26 +1392,35 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non 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 (see - :mod:`policyengine_core.data_storage.store_history`). So the branch - drops each value it holds, other than an input, that was stored at - or after the earliest of: the first time anything in this - simulation's family (the simulation the branches were created from, - and all its branches) stored or calculated ``variable_name`` for a - period that shares a day with ``period``; and the first time a value - of ``variable_name`` was derived from its other periods (uprated, - carried over, or given the default for lack of a formula result). If - neither has happened, nothing can depend on the value and nothing is - dropped; this is the usual case, where a formula creates the branch - while still calculating the variable it overrides. + 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). If + there is neither, 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. + + After such a drop the branch stops reading macro-cache files, which + are keyed by branch and period but not by inputs. 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``) rather than calculating it; - and values calculated in a different simulation family. Inputs set - on a simulation that is not a branch drop nothing, as before. + (``get_known_periods``, ``get_array``) or read another simulation's + storage directly, rather than calculating; and branches a formula + keeps between calls, which hold what their parent held when they + were created. Inputs set on a simulation that is not a branch drop + nothing, as before, and branches already created from the branch + keep their values. """ period = periods.period(period) if self.start_instant is None or self.start_instant > period.start: @@ -1401,6 +1441,18 @@ def _get_store_history(self) -> StoreHistory: history = self._store_history = StoreHistory() return history + def _get_requested_variables(self) -> set: + requested = getattr(self, "_requested_variables", None) + 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 and caller is not self: + caller._get_store_history().merge(self._get_store_history()) + def _drop_values_that_may_depend_on( self, variable_name: str, period: Period ) -> int: @@ -1414,17 +1466,32 @@ def _drop_values_that_may_depend_on( 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 + # Macro-cache files are keyed by branch and period, not by inputs, so + # those this branch wrote may hold values calculated from the old one. + self.macro_cache_read = False 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) + dropped += holder._drop_computed(since) # The fast cache can also hold values a holder does not keep. self._fast_cache = {} + # 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: @@ -1493,6 +1560,10 @@ 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()) return new @@ -1501,8 +1572,9 @@ def get_branch( ) -> "Simulation": """Create a clone of this simulation, whose calculations are traced in the original. - The branch starts with the values this simulation holds. An input set - on the branch drops those that may have been calculated from the + 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. @@ -1518,8 +1590,6 @@ def get_branch( return self if name in self.branches: return self.branches[name] - # The branch shares this simulation's store history. - self._get_store_history() branch = self.clone(clone_tax_benefit_system=clone_system) self.branches[name] = branch branch.branch_name = name @@ -1551,9 +1621,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) diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index 0a2806c00..8b2b8dcaf 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -7,10 +7,14 @@ 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. An -input set on a branch drops the branch's non-input values stored at or after -the first time anything in the family stored, or derived, the overridden -variable for an overlapping period (``StoreHistory.earliest_dependency``). +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``): @@ -36,7 +40,9 @@ from policyengine_core.simulations import SimulationBuilder from tests.fixtures.branch_input_invalidation import ( CARRY_OVER_SYSTEM, + FORMULA_RUNS, ROOT_INPUTS, + branching_simulation, synthetic_simulation, ) @@ -122,11 +128,10 @@ def test_branch_keeps_values_stored_before_the_input_first_existed( branch = simulation.get_branch("raise") branch.set_input("salary", JANUARY, np.array([4000.0])) - assert _stored_keys(branch, "housing_tax") == {"default:2017"} - assert _stored_keys(branch, "income_tax") == set() 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): @@ -247,16 +252,17 @@ def test_disk_storage_drops_calculated_values_and_keeps_inputs(tax_benefit_syste branch = simulation.get_branch("cheaper") branch.set_input("rent", JANUARY, np.array([300.0])) - assert branch.get_holder("housing_allowance")._disk_storage._files == {} - rent_disk = branch.get_holder("rent")._disk_storage - assert set(rent_disk._files) == {"default_2017-01", "cheaper_2017-01"} - assert rent_disk._input_keys == {"default_2017-01", "cheaper_2017-01"} + 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 rent_disk._input_keys == {"default_2017-01", "cheaper_2017-01"} def test_value_served_from_the_macro_cache_is_tracked(tax_benefit_system, monkeypatch): @@ -537,3 +543,331 @@ def test_value_a_spiral_defaults_is_tracked(): 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) + ) + + +# ----- 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._input_keys + + +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_storage_pickled_before_stores_were_numbered_still_stores(): + import pickle + + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0]), periods.period("2020")) + del storage._sequence_numbers, storage._input_keys # as pickled before + restored = pickle.loads(pickle.dumps(storage)) + + restored.put(np.array([2.0]), periods.period("2021"), is_input=True) + + assert restored._input_keys == {"default:2021"} diff --git a/tests/core/test_branch_input_invalidation_property.py b/tests/core/test_branch_input_invalidation_property.py index f4da315c8..890094733 100644 --- a/tests/core/test_branch_input_invalidation_property.py +++ b/tests/core/test_branch_input_invalidation_property.py @@ -1,15 +1,19 @@ """Branch results do not depend on what was calculated before. -Random sequences of calculations, branches, inputs and -``drop_computed_arrays`` run on a synthetic system, with values held in -memory, on disk, or not kept. 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. -``test_branch_input_invalidation.py`` pins the same behaviour with examples. +Random sequences of calculations, branches (including reused and forgotten +names), inputs and ``drop_computed_arrays`` 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 numpy as np import pytest @@ -50,18 +54,25 @@ "p_c", ] MONTHS = ["2013-01", "2013-07", "2014-12", "2015-03"] -MEMORY_CONFIGS = { - "memory": lambda: None, - "disk": lambda: MemoryConfig(max_memory_occupation=0), - "not_kept": lambda: MemoryConfig( - max_memory_occupation=1, variables_to_drop=["p_sum", "p_inner_only"] +# 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. +# 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( @@ -98,7 +109,8 @@ calculate, set_input, set_input, - st.tuples(st.just("branch"), simulation_index), + st.tuples(st.just("branch"), simulation_index, branch_name), + st.tuples(st.just("forget"), simulation_index), st.tuples(st.just("drop"), simulation_index), ) # Pairs of a calculated value and an input it depends on, each reaching it a @@ -109,6 +121,7 @@ ("p_inner", "2013", "p_inner_only", "2013"), ("p_prod", "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"), @@ -123,13 +136,13 @@ lambda drawn: ( [ ("calculate", drawn[0], drawn[1][0], drawn[1][1]), - ("branch", drawn[0]), + ("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]), + ("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]), @@ -146,36 +159,49 @@ ).map(lambda chunks: [operation for chunk in chunks for operation in chunk]) -def _stored_periods(variable, period): - """The periods ``set_input`` stores ``variable`` at, given ``period``.""" - period = periods.period(period) - if SYNTHETIC_SYSTEM.get_variable(variable).definition_period == periods.MONTH: - return [str(month) for month in period.get_subperiods(periods.MONTH)] - return [str(period)] - +def _stored_inputs(own_inputs, variable, period, value): + """What ``set_input(variable, period, value)`` stores on a branch, by its own rule. -def _own_inputs(simulation, variable, period): - """The values stored on the branch itself for ``variable`` over ``period``.""" - holder = simulation.get_holder(variable) - return { - (variable, stored_period): tuple( - holder._get_array_from_storage(stored_period, simulation.branch_name) - ) - for stored_period in _stored_periods(variable, period) - } + 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, memory_config): +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, as stored (an input for a - year is divided between the months the branch has not set itself). + they were then, and those set on it since. """ - root = synthetic_simulation(ROOT_INPUTS, memory_config=memory_config) + 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 = [set()] + own_inputs = [{}] + children = {} # (parent index, name) -> index, while the parent keeps it results = [] for operation in program: kind, index = operation[0], operation[1] % len(simulations) @@ -183,20 +209,28 @@ def _run(program, memory_config): index = len(simulations) - 1 # The newest simulation. simulation = simulations[index] if kind == "branch": - simulations.append(simulation.get_branch(f"b{len(simulations)}")) + 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(set()) + 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 == 0: continue # Inputs on the root after calculating are out of scope. _, _, variable, period, value = operation - stored = {(variable, p) for p in _stored_periods(variable, period)} - if len(stored) > 1 and stored <= own_inputs[index]: - # Dividing a year between months the branch has all set - # itself is an error unless the totals match. + stored = _stored_inputs(own_inputs[index], variable, period, value) + if stored is None: continue simulation.set_input(variable, period, np.asarray(value)) - stored = _own_inputs(simulation, variable, period) inputs[index].update(stored) own_inputs[index].update(stored) elif kind == "calculate": @@ -208,17 +242,24 @@ def _run(program, memory_config): 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=500, + 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(MEMORY_CONFIGS))) +@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, MEMORY_CONFIGS[mode]()): + for branch_inputs, variable, period, value in _run(program, mode): expected = synthetic_simulation(branch_inputs).calculate(variable, period) - # Uprating chains may round differently in float32 depending on which - # earlier periods were calculated first. + # 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, err_msg=f"{variable} {period}" + 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 5014bc62b..2dcb77935 100644 --- a/tests/core/test_branch_scoped_delete_arrays.py +++ b/tests/core/test_branch_scoped_delete_arrays.py @@ -4,6 +4,8 @@ import gc +from pathlib import Path + import numpy as np from policyengine_core.country_template import CountryTaxBenefitSystem @@ -149,7 +151,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/fixtures/branch_input_invalidation.py b/tests/fixtures/branch_input_invalidation.py index 5fe035344..5cce65eb3 100644 --- a/tests/fixtures/branch_input_invalidation.py +++ b/tests/fixtures/branch_input_invalidation.py @@ -1,5 +1,7 @@ """A synthetic system for derived values and branches inside formulas.""" +import tempfile + import numpy as np from policyengine_core import periods @@ -111,6 +113,8 @@ class System(CountryTaxBenefitSystem): 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 @@ -125,10 +129,18 @@ class System(CountryTaxBenefitSystem): YEARS = ["2012", "2013", "2014", "2015"] -def synthetic_simulation(inputs, memory_config=None, system=SYNTHETIC_SYSTEM): +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: + # Each holder's disk storage removes its directory, and this one + # once empty, when garbage-collected. + simulation._data_storage_dir = tempfile.mkdtemp( + prefix="policyengine-branch-tests-" + ) 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 @@ -158,3 +170,83 @@ def synthetic_simulation(inputs, memory_config=None, system=SYNTHETIC_SYSTEM): 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 From 5a18526b2a904f22290eab7f1877835a81ac64a1 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 00:33:01 -0400 Subject: [PATCH 06/42] Keep records while a formula runs; merge incrementally; survive pickling From the second soundness review (in progress): - A drop while a formula is still running in the simulation (its own set_input, or drop_computed_arrays) no longer prunes the history: the formula may hold values it read before the drop and store a result from them. Calculations in flight are counted per simulation; a clone starts at zero. - Merging is incremental: each history keeps a journal of its changes since it was created or pruned, so repeated calls into a simulation whose history keeps growing no longer re-read every record (quadratic before). - The merge map is keyed weakly, so temporary branches' entries go with them, and is left out of pickles (weak references do not pickle). - Unpickling a storage or history advances the process's sequence counter past the numbers it carries; disk restore keeps each key's most recently written file (numbers restart in every process). - Tests: drops inside a running formula, an uprated value handed back by a temporary branch (also in the property's scenarios), cross-process unpickling, restore order. Docs: the in-flight rule and the thread limit. Co-Authored-By: Claude Opus 5.5 --- docs/usage/branches.md | 7 +- .../data_storage/in_memory_storage.py | 9 +- .../data_storage/on_disk_storage.py | 18 ++- .../data_storage/store_history.py | 83 +++++++++--- policyengine_core/simulations/simulation.py | 19 ++- tests/core/test_branch_input_invalidation.py | 121 ++++++++++++++++++ ...test_branch_input_invalidation_property.py | 6 +- tests/fixtures/branch_input_invalidation.py | 11 ++ 8 files changed, 245 insertions(+), 29 deletions(-) diff --git a/docs/usage/branches.md b/docs/usage/branches.md index ee39b8840..0fd9afc33 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -48,7 +48,9 @@ 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. +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. 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 @@ -87,6 +89,9 @@ depend on the input is calculated again) but not less, within these limits: `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. +- **Calculations in other threads.** A formula that calls `calculate` 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. - **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. diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 8769fd1e1..22d5a1adc 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -4,7 +4,10 @@ from numpy.typing import ArrayLike from policyengine_core import periods -from policyengine_core.data_storage.store_history import next_sequence_number +from policyengine_core.data_storage.store_history import ( + advance_sequence_past, + next_sequence_number, +) from policyengine_core.periods import Period @@ -30,6 +33,10 @@ def __setstate__(self, state: dict) -> None: state.setdefault("_sequence_numbers", {}) state.setdefault("_input_keys", set()) 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) -> "InMemoryStorage": clone = InMemoryStorage(self.is_eternal) diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 5411ed70e..bcba9287e 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -6,7 +6,10 @@ from numpy.typing import ArrayLike from policyengine_core import periods -from policyengine_core.data_storage.store_history import next_sequence_number +from policyengine_core.data_storage.store_history import ( + advance_sequence_past, + next_sequence_number, +) from policyengine_core.enums import EnumArray from policyengine_core.periods import Period @@ -37,6 +40,10 @@ def __setstate__(self, state: dict) -> None: state.setdefault("_sequence_numbers", {}) state.setdefault("_input_keys", set()) 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. @@ -198,12 +205,15 @@ def restore(self) -> None: path = os.path.join(self.storage_dir, filename) filename_core = filename.rsplit(".", 1)[0] # Files are named "..npy" (each store writes - # a new file); older dumps are ".npy". Keep each key's latest. + # a new file); older dumps are ".npy". Keep each key's most + # recently written file: sequence numbers restart in every + # process, so they only order files one process wrote. key, _, number = filename_core.rpartition(".") if not (key and number.isdigit()): key, number = filename_core, "0" - if key not in latest or int(number) > latest[key]: - latest[key] = int(number) + order = (os.stat(path).st_mtime_ns, int(number)) + if key not in latest or order > latest[key]: + latest[key] = order files[key] = path def __del__(self) -> None: diff --git a/policyengine_core/data_storage/store_history.py b/policyengine_core/data_storage/store_history.py index bdae4491f..f7b7291e7 100644 --- a/policyengine_core/data_storage/store_history.py +++ b/policyengine_core/data_storage/store_history.py @@ -12,13 +12,13 @@ """ import itertools -from typing import Dict, Optional +import weakref +from typing import Dict, List, Optional, Tuple from policyengine_core import periods from policyengine_core.periods import Period _sequence = itertools.count(1) -_history_ids = itertools.count(1) def next_sequence_number() -> int: @@ -26,6 +26,17 @@ def next_sequence_number() -> int: 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 periods_overlap(first: Period, second: Period) -> bool: """Whether two periods share at least one day.""" if periods.ETERNITY in (first.unit, second.unit): @@ -61,11 +72,34 @@ class StoreHistory: def __init__(self): self._first_stored: Dict[str, Dict[Period, int]] = {} self._first_derived: Dict[str, int] = {} - # Changes on every new record, so a merge can skip a history it has - # already taken in unchanged. - self._id = next(_history_ids) - self._version = 0 - self._merged_versions: Dict[int, int] = {} + # 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). + self._journal: List[Tuple[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: + 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 numbers: + advance_sequence_past(max(numbers)) def copy(self) -> "StoreHistory": new = StoreHistory() @@ -74,7 +108,7 @@ def copy(self) -> "StoreHistory": for variable_name, stored in self._first_stored.items() } new._first_derived = dict(self._first_derived) - new._merged_versions = dict(self._merged_versions) + new._merged = weakref.WeakKeyDictionary(self._merged) return new def record_store(self, variable_name: str, period: Period, sequence_number: int): @@ -82,24 +116,36 @@ def record_store(self, variable_name: str, period: Period, sequence_number: int) recorded = stored.get(period) if recorded is None or sequence_number < recorded: stored[period] = sequence_number - self._version += 1 + 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._version += 1 + self._journal.append((variable_name, None, sequence_number)) def merge(self, other: "StoreHistory") -> None: """Take in ``other``'s records, keeping the earlier number of each.""" - if other is self or self._merged_versions.get(other._id) == other._version: + if other is self: return - for variable_name, stored in other._first_stored.items(): - for period, sequence_number in stored.items(): + read = self._merged.get(other) + if read is not None and read[0] == other._generation: + changes = other._journal[read[1] :] + else: + changes = [ + (variable_name, period, sequence_number) + for variable_name, stored in other._first_stored.items() + for period, sequence_number in stored.items() + ] + [ + (variable_name, None, sequence_number) + for variable_name, sequence_number in other._first_derived.items() + ] + for variable_name, period, sequence_number in changes: + if period is None: + self.record_derived(variable_name, sequence_number) + else: self.record_store(variable_name, period, sequence_number) - for variable_name, sequence_number in other._first_derived.items(): - self.record_derived(variable_name, sequence_number) - self._merged_versions[other._id] = other._version + self._merged[other] = (other._generation, len(other._journal)) def prune(self, since: Optional[int] = None) -> None: """Forget the records numbered ``since`` or later (all, without ``since``). @@ -125,9 +171,10 @@ def prune(self, since: Optional[int] = None) -> None: for variable_name, sequence_number in self._first_derived.items() if since is not None and sequence_number < since } - self._version += 1 + self._journal = [] + self._generation += 1 # Merging a history again must restore what was forgotten here. - self._merged_versions = {} + 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``. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index dc28b5a56..18fd6ffda 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -579,6 +579,10 @@ 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``). + self._calculations_in_flight = getattr(self, "_calculations_in_flight", 0) + 1 try: result = self._calculate(variable_name, period) if isinstance(result, EnumArray) and decode_enums: @@ -592,6 +596,7 @@ def calculate( self._share_store_history_with_caller() return result finally: + self._calculations_in_flight -= 1 self.tracer.record_calculation_end() self.purge_cache_of_invalid_values() @@ -1416,9 +1421,10 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non 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; and branches a formula - keeps between calls, which hold what their parent held when they - were created. Inputs set on a simulation that is not a branch drop + storage directly, rather than calculating; ``calculate`` called from + a thread a formula starts without copying its context; and branches a + formula keeps between calls, which hold what their parent held when + they were created. Inputs set on a simulation that is not a branch drop nothing, as before, and branches already created from the branch keep their values. """ @@ -1485,6 +1491,11 @@ def _drop_computed(self, since: Optional[int] = None) -> int: dropped += holder._drop_computed(since) # The fast cache can also hold values a holder does not keep. self._fast_cache = {} + if getattr(self, "_calculations_in_flight", 0): + # A formula running here may still hold values calculated from + # what the records describe, and store a result from them: 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. @@ -1564,6 +1575,8 @@ def clone( # 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 return new diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index 8b2b8dcaf..bfd9ece4d 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -658,6 +658,100 @@ def test_history_is_taken_in_again_after_a_drop_forgets_it(): ) +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_file_of_each_key(tmp_path): + """Sequence numbers restart in each process; restore goes by write time.""" + import itertools + import os + + 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")) + first = writer._files["default_2020"] + saved = store_history._sequence + store_history._sequence = itertools.count(1) # a later process + try: + later = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + later.put(np.array([2.0]), periods.period("2020")) + stat = os.stat(first) + os.utime(first, ns=(stat.st_atime_ns, stat.st_mtime_ns - 10**9)) + finally: + store_history._sequence = saved + + 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"), + ) + + # ----- Only what may depend on the input is dropped ----- # @@ -860,6 +954,33 @@ def test_disk_branches_with_the_same_name_keep_their_own_values(): 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()) + history.merge(StoreHistory()) # weak references must not stop pickling + 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 diff --git a/tests/core/test_branch_input_invalidation_property.py b/tests/core/test_branch_input_invalidation_property.py index 890094733..8fbd16cc6 100644 --- a/tests/core/test_branch_input_invalidation_property.py +++ b/tests/core/test_branch_input_invalidation_property.py @@ -50,6 +50,7 @@ "p_months", "p_inner", "p_inner_only", + "p_imported_up", "p_up", "p_c", ] @@ -115,11 +116,12 @@ ) # 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, a month of a year, a value the holder may not keep, a lagged year, -# and a year summed from months. +# 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"), diff --git a/tests/fixtures/branch_input_invalidation.py b/tests/fixtures/branch_input_invalidation.py index 5cce65eb3..06f737cc3 100644 --- a/tests/fixtures/branch_input_invalidation.py +++ b/tests/fixtures/branch_input_invalidation.py @@ -36,6 +36,16 @@ def _inner_branch_formula(person, period): 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) @@ -99,6 +109,7 @@ def _monthly(name, formula=None, **attributes): 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( From 4b70aac16db6585b5648cd0e63ec53e2b85d1662 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 00:43:33 -0400 Subject: [PATCH 07/42] Hand branch history to waiting ancestors; tokened disk files; gate on success From the second soundness review: - A branch calculated from a thread with no formula context (a worker a formula starts without copying its context) hands its store history to every ancestor with a calculation running, so the waiting formula's result keeps the dependency. Unrelated simulations are not touched. - Disk files carry a per-process token as well as the sequence number: numbers restart in each process, so a later process could otherwise overwrite a file a restored view still maps. - requires_computation_after counts a prerequisite only once its calculation succeeded. - Docs and comments: the soundness argument states records for overlapping periods (summed and divided values), disk files stay until the directory is removed, the thread limit covers only unrelated simulations, and the history is copied into branches, not shared. Co-Authored-By: Claude Opus 5.5 --- docs/usage/branches.md | 26 ++++--- .../data_storage/on_disk_storage.py | 43 +++++++---- policyengine_core/simulations/simulation.py | 33 ++++++--- tests/core/test_branch_input_invalidation.py | 71 +++++++++++++++++++ 4 files changed, 141 insertions(+), 32 deletions(-) diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 0fd9afc33..37f55d50f 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -41,7 +41,9 @@ from another period. The history records: - 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 (a formula that branches, sets an input and - calculates in the branch hands back values calculated there). + 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 @@ -68,10 +70,12 @@ 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, or handed -back by another simulation's `calculate`, and in each case what it was -calculated from is in the simulation's history. So any value that depends on -the overridden variable at an overlapping period was stored after the earliest -record of it. Uprating and carry-over read which periods hold values at all, +back by another simulation's `calculate`. In each case, 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). So +any value that depends on the overridden variable at an overlapping period was +stored after the earliest record of it. Uprating and carry-over read which periods hold values at all, which is why the first uprated or carried-over value also counts. The rule can drop more than it needs to (a value stored later that does not @@ -89,9 +93,11 @@ depend on the input is calculated again) but not less, within these limits: `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. -- **Calculations in other threads.** A formula that calls `calculate` 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. +- **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. Its own branches are + covered: their history goes to every ancestor with a calculation running. - **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. @@ -101,6 +107,10 @@ depend on the input is calculated again) but not less, within these limits: - **Existing child branches keep their values.** An input set on a branch does not reach branches already created from it. +With disk storage (`MemoryConfig`), every store writes a new file, so a value +recalculated in one simulation does not change a file another simulation in +the family still maps; files stay until the storage directory is removed. + Two related behaviours: a branch whose input dropped values stops reading macro-cache files, which are keyed by branch and period but not by inputs; and `requires_computation_after` is satisfied by a prerequisite requested before a diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index bcba9287e..90d5c43be 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -1,5 +1,6 @@ import os import shutil +import uuid from typing import Dict, List, Optional, Set, Tuple import numpy @@ -11,6 +12,10 @@ next_sequence_number, ) from policyengine_core.enums import EnumArray + +# Distinguishes the files this process writes from other processes' files in +# the same directory, whose sequence numbers may repeat this process's. +_PROCESS_TOKEN = uuid.uuid4().hex[:12] from policyengine_core.periods import Period @@ -49,11 +54,11 @@ def clone(self) -> "OnDiskStorage": """Create a private metadata view over this storage directory. The file and enum mappings are copied so deleting or rewiring entries - through the clone does not mutate the source storage. The underlying - ``.npy`` files remain shared: writing the same ``{branch}_{period}`` - key from two views targets the same path and can overwrite the file. - Clones retain the original cleanup owner so the shared directory stays - alive, but never own cleanup themselves. + through the clone does not mutate the source storage. The directory + is shared, but every store writes a new file, so writing a key + through one view leaves the file another view maps. Files stay until + the directory is removed. Clones retain the original cleanup owner so + the shared directory stays alive, but never own cleanup themselves. """ clone = OnDiskStorage( self.storage_dir, @@ -105,8 +110,14 @@ def put( if sequence_number is None: sequence_number = next_sequence_number() # A new file for every store: clones share this directory and may - # still map an earlier file for the same key. - path = os.path.join(self.storage_dir, f"{filename}.{sequence_number}") + ".npy" + # still map an earlier file for the same key. The process token keeps + # files from processes whose counters restarted apart. + path = ( + os.path.join( + self.storage_dir, f"{filename}.{_PROCESS_TOKEN}.{sequence_number}" + ) + + ".npy" + ) if isinstance(value, EnumArray): self._enums[path] = value.possible_values value = value.view(numpy.ndarray) @@ -204,14 +215,16 @@ def restore(self) -> None: continue path = os.path.join(self.storage_dir, filename) filename_core = filename.rsplit(".", 1)[0] - # Files are named "..npy" (each store writes - # a new file); older dumps are ".npy". Keep each key's most - # recently written file: sequence numbers restart in every - # process, so they only order files one process wrote. - key, _, number = filename_core.rpartition(".") - if not (key and number.isdigit()): - key, number = filename_core, "0" - order = (os.stat(path).st_mtime_ns, int(number)) + # Files are named "...npy" + # (each store writes a new file); older dumps are ".npy". + # Keep each key's most recently written file: sequence numbers + # restart in every process, so they only order one process's files. + parts = filename_core.rsplit(".", 2) + if len(parts) == 3 and parts[2].isdigit(): + key, number = parts[0], int(parts[2]) + else: + key, number = filename_core, 0 + order = (os.stat(path).st_mtime_ns, number) if key not in latest or order > latest[key]: latest[key] = order files[key] = path diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 18fd6ffda..4cb0b9392 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -183,8 +183,8 @@ def __init__( # post-``apply_reform`` cache wipe would also wipe the dataset the # simulation was loaded from. self._user_input_keys: set[tuple[str, str, Period]] = set() - # When each variable's values were first stored, shared with every - # branch created from this simulation (see ``set_input``). + # 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 @@ -585,6 +585,9 @@ def calculate( self._calculations_in_flight = getattr(self, "_calculations_in_flight", 0) + 1 try: result = self._calculate(variable_name, period) + # Satisfies ``requires_computation_after`` from now on, even if a + # branch input later drops the values. + self._get_requested_variables().add(variable_name) if isinstance(result, EnumArray) and decode_enums: result = result.decode_to_str() self.tracer.record_calculation_result(result) @@ -710,7 +713,6 @@ 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.") - self._get_requested_variables().add(variable_name) population = self.get_variable_population(variable_name) holder = population.get_holder(variable_name) variable = self.tax_benefit_system.get_variable( @@ -1421,10 +1423,11 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non 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; ``calculate`` called from - a thread a formula starts without copying its context; and branches a - formula keeps between calls, which hold what their parent held when - they were created. Inputs set on a simulation that is not a branch drop + 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. Inputs set on a simulation that is not a branch drop nothing, as before, and branches already created from the branch keep their values. """ @@ -1456,8 +1459,20 @@ def _get_requested_variables(self) -> set: 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 and caller is not self: - caller._get_store_history().merge(self._get_store_history()) + 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 getattr(ancestor, "_calculations_in_flight", 0): + 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 diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index bfd9ece4d..7d4166c34 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -752,6 +752,77 @@ def test_uprated_value_another_simulation_returns_is_tracked(): ) +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_disk_files_from_another_process_are_not_overwritten(tmp_path, monkeypatch): + """Processes write their own files even when their sequence numbers repeat.""" + import itertools + + import policyengine_core.data_storage.on_disk_storage as on_disk_storage + import policyengine_core.data_storage.store_history as store_history + from policyengine_core.data_storage import OnDiskStorage + + monkeypatch.setattr(store_history, "_sequence", itertools.count(1)) + monkeypatch.setattr(on_disk_storage, "_PROCESS_TOKEN", "first") + writer = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + writer.put(np.array([1.0]), periods.period("2020")) + snapshot = writer.clone() + + # A later process: its counter restarts. + monkeypatch.setattr(store_history, "_sequence", itertools.count(1)) + monkeypatch.setattr(on_disk_storage, "_PROCESS_TOKEN", "second") + reader = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + reader.restore() + reader.put(np.array([99.0]), periods.period("2020")) + + assert snapshot.get(periods.period("2020")).tolist() == [1.0] + reader.restore() + assert reader.get(periods.period("2020")).tolist() == [99.0] + + # ----- Only what may depend on the input is dropped ----- # From 32d5bc435a4183ffba77da49c4d749a26a5a52b5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 01:26:28 -0400 Subject: [PATCH 08/42] Close the third review's gaps: in-flight results, failures, dumps, disk names From the round-3 soundness review: - A calculation still running when an input set on its simulation dropped values returns its result but does not cache it (an input epoch), so the next calculation uses the new input. - A calculation that raises still hands its store history to the calling formula: whether it fails can depend on what it read. - A merge records how far it read the source before reading, so a record added meanwhile (another thread) is read next time. - dump_simulation records each variable's input periods (inputs.txt); restore_simulation restores inputs as inputs, then every other value under one later number, so a branch input drops what it may have fed. Older dumps restore every value as an input, as before. - Disk restore reads the previous `..npy` names, prefers this process's file on a timestamp tie, and a forked child gets its own token. - StoreHistory pickles from before journals load with empty journals. - Tests for each; docs updated. Co-Authored-By: Claude Opus 5.5 --- docs/usage/branches.md | 25 ++- .../data_storage/on_disk_storage.py | 30 +++- .../data_storage/store_history.py | 17 +- policyengine_core/holders/holder.py | 14 +- policyengine_core/simulations/simulation.py | 39 ++++- policyengine_core/tools/simulation_dumper.py | 44 ++++- tests/core/test_branch_input_invalidation.py | 159 +++++++++++++++++- 7 files changed, 292 insertions(+), 36 deletions(-) diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 37f55d50f..69a5dda73 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -40,8 +40,8 @@ from another period. The history records: macro cache, and the default a spiral returns; - 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 (a formula that branches, sets an input and - calculates in the branch hands back values calculated there); a branch + `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. @@ -52,7 +52,8 @@ earliest recorded store of the variable for a period that shares a day with 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. +variables, a value it read before the drop, so the records stay, and the +results of the calculations that were running are returned but not kept. 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 @@ -107,9 +108,21 @@ depend on the input is calculated again) but not less, within these limits: - **Existing child branches keep their values.** An input set on a branch does not reach branches already created from it. -With disk storage (`MemoryConfig`), every store writes a new file, so a value -recalculated in one simulation does not change a file another simulation in -the family still maps; files stay until the storage directory is removed. +With disk storage (`MemoryConfig`), every store writes a new file, named with +the sequence number and a token for the process (a forked child gets its own), +so a value recalculated in one simulation does not change a file another +simulation, or another process, still maps; files stay until the storage +directory is removed. `OnDiskStorage.restore` takes each key's most recently +written file, and on a timestamp tie the current process's own; between two +other processes' files written within one clock tick it cannot tell which came +last. + +A simulation dump (`dump_simulation`) records which values were inputs. +`restore_simulation` restores those as inputs and every other value as +calculated under one number, since the dump does not say what each was +calculated from: an input set on a branch of the restored simulation drops all +of them that are not inputs. A dump written before inputs were recorded is +restored with every value as an input, as before, so such values never drop. Two related behaviours: a branch whose input dropped values stops reading macro-cache files, which are keyed by branch and period but not by inputs; and diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 90d5c43be..78233fb5e 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -16,6 +16,16 @@ # Distinguishes the files this process writes from other processes' files in # the same directory, whose sequence numbers may repeat this process's. _PROCESS_TOKEN = uuid.uuid4().hex[:12] + + +def _new_process_token() -> None: + global _PROCESS_TOKEN + _PROCESS_TOKEN = uuid.uuid4().hex[:12] + + +# A forked child inherits the parent's token and counter: give it its own token. +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_new_process_token) from policyengine_core.periods import Period @@ -219,12 +229,20 @@ def restore(self) -> None: # (each store writes a new file); older dumps are ".npy". # Keep each key's most recently written file: sequence numbers # restart in every process, so they only order one process's files. - parts = filename_core.rsplit(".", 2) - if len(parts) == 3 and parts[2].isdigit(): - key, number = parts[0], int(parts[2]) - else: - key, number = filename_core, 0 - order = (os.stat(path).st_mtime_ns, number) + key, token, number = filename_core, None, 0 + stem, _, last = filename_core.rpartition(".") + if stem and last.isdigit(): + key, number = stem, int(last) + stem, _, middle = key.rpartition(".") + if ( + stem + and len(middle) == 12 + and all(c in "0123456789abcdef" for c in middle) + ): + key, token = stem, middle + # On a timestamp tie (a coarse filesystem clock), this process's + # own file is the later one; numbers only order one process's. + order = (os.stat(path).st_mtime_ns, token == _PROCESS_TOKEN, number) if key not in latest or order > latest[key]: latest[key] = order files[key] = path diff --git a/policyengine_core/data_storage/store_history.py b/policyengine_core/data_storage/store_history.py index f7b7291e7..140976faf 100644 --- a/policyengine_core/data_storage/store_history.py +++ b/policyengine_core/data_storage/store_history.py @@ -91,6 +91,8 @@ def __getstate__(self) -> dict: return state def __setstate__(self, state: dict) -> None: + state.setdefault("_journal", []) + state.setdefault("_generation", 0) self.__dict__.update(state) self._merged = weakref.WeakKeyDictionary() numbers = [ @@ -128,24 +130,27 @@ 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] == other._generation: - changes = other._journal[read[1] :] + 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 other._first_stored.items() - for period, sequence_number in stored.items() + 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 other._first_derived.items() + for variable_name, sequence_number in list(other._first_derived.items()) ] for variable_name, period, sequence_number in changes: if period is None: self.record_derived(variable_name, sequence_number) else: self.record_store(variable_name, period, sequence_number) - self._merged[other] = (other._generation, len(other._journal)) + self._merged[other] = (generation, end) def prune(self, since: Optional[int] = None) -> None: """Forget the records numbered ``since`` or later (all, without ``since``). diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 3b00c6430..37999bcf4 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -126,6 +126,16 @@ def _record_inputs(self, since: Optional[int] = None) -> 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 the value stored for ``period`` on ``branch_name`` is an input.""" + if self.variable.definition_period == periods.ETERNITY: + period = periods.period(periods.ETERNITY) + period = periods.period(period) + return f"{branch_name}:{period}" in self._memory_storage._input_keys or ( + self._disk_storage is not None + and f"{branch_name}_{period}" in self._disk_storage._input_keys + ) + def _has_unnumbered_values(self) -> bool: return self._memory_storage.has_unnumbered_values() or ( self._disk_storage is not None @@ -377,6 +387,7 @@ def _set( branch_name: str = "default", validate_nan: bool = False, is_input: Optional[bool] = None, + sequence_number: Optional[int] = None, ) -> None: simulation = getattr(self, "simulation", None) user_input_contexts = getattr(simulation, "_user_input_contexts", None) @@ -401,7 +412,8 @@ def _set( >= self.simulation.memory_config.max_memory_occupation_pc ) - sequence_number = next_sequence_number() + if sequence_number is None: + sequence_number = next_sequence_number() storage = self._disk_storage if should_store_on_disk else self._memory_storage storage.put( value, diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 243a21ae8..8eca4bd33 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -624,9 +624,11 @@ def calculate( variable_name ).entity.key result = self.map_result(result, source_entity, map_to) - self._share_store_history_with_caller() return result finally: + # 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() self._calculations_in_flight -= 1 self.tracer.record_calculation_end() self.purge_cache_of_invalid_values() @@ -741,6 +743,7 @@ 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_epoch = getattr(self, "_input_epoch", 0) population = self.get_variable_population(variable_name) holder = population.get_holder(variable_name) variable = self.tax_benefit_system.get_variable( @@ -839,7 +842,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if np.all(~mask): array = holder.default_array() array = self._cast_formula_result(array, variable) - holder.put_in_cache(array, period, self.branch_name) + self._cache_result(holder, array, period, input_epoch) return array array = None @@ -941,7 +944,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: array = EnumArray(array, variable.possible_values) array = self._cast_formula_result(array, variable) - holder.put_in_cache(array, period, self.branch_name) + self._cache_result(holder, array, period, input_epoch) except SpiralError: array = holder.default_array() @@ -964,11 +967,28 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if is_cache_available: smc.set_cache_value(cache_path, array) - if hasattr(self, "_fast_cache"): + if hasattr(self, "_fast_cache") and input_epoch == getattr( + self, "_input_epoch", 0 + ): self._fast_cache[(variable_name, period)] = array return array + def _cache_result( + self, holder: Holder, array: ArrayLike, period: Period, input_epoch: int + ) -> None: + """Cache a calculated value, unless inputs changed while it was calculated. + + 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``). + """ + if input_epoch == getattr(self, "_input_epoch", 0): + holder.put_in_cache(array, period, self.branch_name) + else: + # Not kept, but whatever reads it is stored after it all the same. + holder._record_store(period, next_sequence_number()) + 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: @@ -1018,12 +1038,13 @@ def calculate_add( ) ) + input_epoch = getattr(self, "_input_epoch", 0) result = sum( self.calculate(variable_name, sub_period) for sub_period in period.get_subperiods(variable.definition_period) ) holder = self.get_holder(variable.name) - holder.put_in_cache(result, period, self.branch_name) + self._cache_result(holder, result, period, input_epoch) return result def calculate_divide( @@ -1053,10 +1074,11 @@ def calculate_divide( ) if period.unit == periods.MONTH: + input_epoch = getattr(self, "_input_epoch", 0) computation_period = period.this_year result = self.calculate(variable_name, period=computation_period) / 12.0 holder = self.get_holder(variable.name) - holder.put_in_cache(result, period, self.branch_name) + self._cache_result(holder, result, period, input_epoch) return result elif period.unit == periods.YEAR: return self.calculate(variable_name, period) @@ -1536,8 +1558,9 @@ def _drop_computed(self, since: Optional[int] = None) -> int: self._fast_cache = {} if getattr(self, "_calculations_in_flight", 0): # A formula running here may still hold values calculated from - # what the records describe, and store a result from them: keep - # the records. + # what the records describe: keep the records, and keep none of + # the results the running calculations return (``_cache_result``). + self._input_epoch = getattr(self, "_input_epoch", 0) + 1 return dropped # Nothing the simulation still holds was calculated from what the # records numbered ``since`` or later describe, except the inputs it diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index c3db0c4fa..724757c5d 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -6,9 +6,13 @@ 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 +# Next to each variable's arrays: the periods whose values were inputs. +_INPUTS_FILE = "inputs.txt" + def dump_simulation(simulation, directory): """ @@ -55,20 +59,33 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): _restore_entity(population, entities_dump_dir) population.count = person_count - variables_to_restore = ( + variables_to_restore = [ variable for variable in os.listdir(directory) if variable != "__entities__" - ) + ] + # Inputs first, then every calculated value under one later number: the + # dump does not say what each value was calculated from, so any input set + # on a branch of the restored simulation counts it as possibly dependent. for variable in variables_to_restore: - _restore_holder(simulation, variable, directory) + _restore_holder(simulation, variable, directory, inputs=True) + calculated_number = next_sequence_number() + for variable in variables_to_restore: + _restore_holder( + simulation, variable, directory, calculated_number=calculated_number + ) return simulation def _dump_holder(holder, directory): disk_storage = holder.create_disk_storage(directory, preserve=True) + inputs = [] for period in holder.get_known_periods(): value = holder.get_array(period) disk_storage.put(value, period) + if holder._is_input(period): + inputs.append(str(period)) + with open(os.path.join(disk_storage.storage_dir, _INPUTS_FILE), "w") as file: + file.write("\n".join(inputs)) def _dump_entity(population, directory): @@ -122,7 +139,9 @@ def _restore_entity(population, directory): return person_count -def _restore_holder(simulation, variable, directory): +def _restore_holder( + simulation, variable, directory, inputs=False, calculated_number=None +): storage_dir = os.path.join(directory, variable) is_variable_eternal = ( simulation.tax_benefit_system.get_variable(variable).definition_period @@ -134,7 +153,22 @@ def _restore_holder(simulation, variable, directory): disk_storage.restore() holder = simulation.get_holder(variable) + inputs_path = os.path.join(storage_dir, _INPUTS_FILE) + if os.path.exists(inputs_path): + with open(inputs_path) as file: + input_periods = set(file.read().split()) + else: + # Dumped before inputs were recorded: keep every value, as inputs. + input_periods = None for period in disk_storage.get_known_periods(): + is_input = input_periods is None or str(period) in input_periods + if is_input != inputs: + continue value = disk_storage.get(period) - holder.put_in_cache(value, period) + holder._set( + period, + value, + is_input=is_input, + sequence_number=None if is_input else calculated_number, + ) diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index 7d4166c34..215f9d65d 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -494,8 +494,9 @@ 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 and its later branches never store ``p_inner_only``. The - history the whole family shares still records it. + 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") @@ -806,14 +807,14 @@ def test_disk_files_from_another_process_are_not_overwritten(tmp_path, monkeypat from policyengine_core.data_storage import OnDiskStorage monkeypatch.setattr(store_history, "_sequence", itertools.count(1)) - monkeypatch.setattr(on_disk_storage, "_PROCESS_TOKEN", "first") + monkeypatch.setattr(on_disk_storage, "_PROCESS_TOKEN", "aaaaaaaaaaaa") writer = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) writer.put(np.array([1.0]), periods.period("2020")) snapshot = writer.clone() # A later process: its counter restarts. monkeypatch.setattr(store_history, "_sequence", itertools.count(1)) - monkeypatch.setattr(on_disk_storage, "_PROCESS_TOKEN", "second") + monkeypatch.setattr(on_disk_storage, "_PROCESS_TOKEN", "bbbbbbbbbbbb") reader = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) reader.restore() reader.put(np.array([99.0]), periods.period("2020")) @@ -823,6 +824,156 @@ def test_disk_files_from_another_process_are_not_overwritten(tmp_path, monkeypat assert reader.get(periods.period("2020")).tolist() == [99.0] +def test_result_calculated_across_its_own_input_change_is_not_kept(): + """A formula that read an input and then replaced it returns, but does not cache.""" + + 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() == [2.0] + + assert branch.calculate("result", "2020").tolist() == [6.0] + + +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 + + +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 + + +def test_disk_restore_reads_older_file_names_and_breaks_time_ties(tmp_path): + import os + + import policyengine_core.data_storage.on_disk_storage as on_disk_storage + from policyengine_core.data_storage import OnDiskStorage + + # A file named before process tokens, and one from another process. + np.save(tmp_path / "default_2020.4.npy", np.array([1.0])) + np.save(tmp_path / "default_2020.cccccccccccc.999999999999.npy", np.array([2.0])) + storage = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + storage.put(np.array([3.0]), periods.period("2020")) # this process, number lower + for path in tmp_path.glob("*.npy"): + os.utime(path, ns=(10**18, 10**18)) # a coarse clock: every time ties + + reader = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + reader.restore() + + assert set(reader._files) == {"default_2020"} # one key, old name included + assert reader.get(periods.period("2020")).tolist() == [3.0] + assert on_disk_storage._PROCESS_TOKEN in reader._files["default_2020"] + + +def test_forked_process_gets_its_own_disk_file_token(): + import multiprocessing + + import policyengine_core.data_storage.on_disk_storage as on_disk_storage + + if "fork" not in multiprocessing.get_all_start_methods(): + pytest.skip("no fork on this platform") + context = multiprocessing.get_context("fork") + with context.Pool(1) as pool: + child_token = pool.apply(_process_token) + + assert child_token != on_disk_storage._PROCESS_TOKEN + + +def _process_token(): + import policyengine_core.data_storage.on_disk_storage as on_disk_storage + + return on_disk_storage._PROCESS_TOKEN + + # ----- Only what may depend on the input is dropped ----- # From bff21a79ac9717331f4c73598cc1a639597c7369 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 01:29:25 -0400 Subject: [PATCH 09/42] Test pickling with a live weak read position; tidy disk storage imports The pickle test merged a history that died at once, so the weak map was empty when pickled. Keep the source alive so the test covers the case the __getstate__ exists for. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/data_storage/on_disk_storage.py | 2 +- tests/core/test_branch_input_invalidation.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 78233fb5e..4330e6d99 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -12,6 +12,7 @@ next_sequence_number, ) from policyengine_core.enums import EnumArray +from policyengine_core.periods import Period # Distinguishes the files this process writes from other processes' files in # the same directory, whose sequence numbers may repeat this process's. @@ -26,7 +27,6 @@ def _new_process_token() -> None: # A forked child inherits the parent's token and counter: give it its own token. if hasattr(os, "register_at_fork"): os.register_at_fork(after_in_child=_new_process_token) -from policyengine_core.periods import Period class OnDiskStorage: diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index 215f9d65d..02e316a88 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -1189,7 +1189,9 @@ def test_values_unpickled_from_another_process_stay_earlier(monkeypatch): next_sequence_number() storage.put(np.array([1.0]), periods.period("2020")) history.record_store("v", periods.period("2020"), next_sequence_number()) - history.merge(StoreHistory()) # weak references must not stop pickling + 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 From 9b42e030f0dfdb11d0ed1f931b84787cf6b12731 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 01:49:39 -0400 Subject: [PATCH 10/42] Fork the token test's child directly instead of through a process pool A process pool forked from a multi-threaded process (the test process can have threads by then) can deadlock; the child now only writes its token to a pipe and exits. Still fails without the at-fork hook. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_branch_input_invalidation.py | 31 ++++++++++++-------- 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index 02e316a88..892bdf994 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -955,23 +955,30 @@ def test_disk_restore_reads_older_file_names_and_breaks_time_ties(tmp_path): def test_forked_process_gets_its_own_disk_file_token(): - import multiprocessing + import os + import warnings import policyengine_core.data_storage.on_disk_storage as on_disk_storage - if "fork" not in multiprocessing.get_all_start_methods(): + if not hasattr(os, "fork"): pytest.skip("no fork on this platform") - context = multiprocessing.get_context("fork") - with context.Pool(1) as pool: - child_token = pool.apply(_process_token) - - assert child_token != on_disk_storage._PROCESS_TOKEN - - -def _process_token(): - import policyengine_core.data_storage.on_disk_storage as on_disk_storage + # A bare fork whose child only writes to a pipe: a process pool forked + # from a multi-threaded test process can deadlock. + read, write = os.pipe() + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + pid = os.fork() + if pid == 0: + try: + os.write(write, on_disk_storage._PROCESS_TOKEN.encode()) + finally: + os._exit(0) + os.close(write) + child_token = os.read(read, 64).decode() + os.close(read) + os.waitpid(pid, 0) - return on_disk_storage._PROCESS_TOKEN + assert child_token and child_token != on_disk_storage._PROCESS_TOKEN # ----- Only what may depend on the input is dropped ----- # From 2b31dcb76d3b174249038026f26369519a364d26 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 02:08:24 -0400 Subject: [PATCH 11/42] Close the fourth review's gaps: refused results, dumps, handlers, macro reads From the round-4 soundness review (on bff21a79): - A result refused because an input changed while it was calculated left its period unknown, so later carry-over found nothing (0 instead of 3). calculate now runs such a calculation once more from the new inputs and keeps that result; a second change still returns without keeping. - An input a formula sets for the very period it calculates is the result and is no longer overwritten by the formula's return value. - A dump does not say what each value was calculated from, nor what was read without being kept, and a macro-cache value's sources were never calculated here: the store history now records the number from which values may depend on anything (restored values, the first macro-cache read), so an input for any variable on a branch drops them. - A custom set_input handler that calculates between its own stores: the branch drops again, by the same rule, once the handler returns. - dump_simulation on a branch dumps the values the branch reads (its own and inherited) instead of the default branch's, which could not be restored. - calculate ends the calculation (in-flight count, tracer) even if handing its history back fails. - The property test also dumps and restores simulations, including a chunk that calculates, dumps, branches the restored simulation and overrides an input; it catches a history without the new record. Co-Authored-By: Claude Opus 5.5 --- .../branch-set-input-invalidation.changed.md | 2 +- .../branch-set-input-invalidation.fixed.md | 2 +- docs/usage/branches.md | 57 +++-- .../data_storage/store_history.py | 44 +++- policyengine_core/holders/holder.py | 30 ++- policyengine_core/simulations/simulation.py | 96 +++++--- policyengine_core/tools/simulation_dumper.py | 29 ++- tests/core/test_branch_input_invalidation.py | 223 +++++++++++++++++- ...test_branch_input_invalidation_property.py | 44 +++- 9 files changed, 455 insertions(+), 72 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.changed.md b/changelog.d/branch-set-input-invalidation.changed.md index 8273a4824..78a326820 100644 --- a/changelog.d/branch-set-input-invalidation.changed.md +++ b/changelog.d/branch-set-input-invalidation.changed.md @@ -1 +1 @@ -Disk storage writes a new file for every store, so recalculating a value no longer changes what branches sharing the directory read; apply_reform keeps inputs by a flag on each stored array instead of replaying _user_input_keys, and values a custom set_input handler calculates are no longer treated as inputs. +Disk storage writes a new file for every store, so recalculating a value no longer changes what branches sharing the directory read; apply_reform keeps inputs by a flag on each stored array instead of replaying _user_input_keys; values a custom set_input handler calculates are no longer treated as inputs; dump_simulation records which values were inputs (inputs.txt), and restore_simulation restores only those as inputs; a calculation whose formula changes an input it read is run once more from the new input. diff --git a/changelog.d/branch-set-input-invalidation.fixed.md b/changelog.d/branch-set-input-invalidation.fixed.md index 26adbe967..146bf2378 100644 --- a/changelog.d/branch-set-input-invalidation.fixed.md +++ b/changelog.d/branch-set-input-invalidation.fixed.md @@ -1 +1 @@ -Simulation.set_input on a branch now drops the values the branch holds that may have been calculated from the value it replaces (every stored array carries a sequence number, and each simulation records what its values may have been calculated from), so a value the parent calculated before branching no longer shadows the branch's input. Simulation.derivative keeps inputs set after the simulation was built. +Simulation.set_input on a branch now drops the values the branch holds that may have been calculated from the value it replaces (every stored array carries a sequence number, and each simulation records what its values may have been calculated from), so a value the parent calculated before branching no longer shadows the branch's input. Simulation.derivative keeps inputs set after the simulation was built. An input a formula sets for the period it is calculating is no longer overwritten by the formula's result, and dump_simulation on a branch dumps the branch's own values (it read the default branch's, so the dump could not be restored). diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 69a5dda73..673c6bdc4 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -38,6 +38,10 @@ 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 @@ -52,8 +56,18 @@ earliest recorded store of the variable for a period that shares a day with 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, and the -results of the calculations that were running are returned but not kept. +variables, a value it read before the drop, so the records stay. A +calculation that was running then (one whose formula sets an input, say) may +have read the replaced value, so its result is not kept; `calculate` runs it +once more from the new inputs and keeps that result, so uprating and +carry-over find its period as they would had the input come first. If the +second run changes an input again, its result is returned but not kept. An +input set for the very period being calculated, while it is calculated, is +the result, as it would be had it been set first. + +A custom `set_input` handler that calculates values between its own stores +calculates them from inputs it has not yet replaced, so the branch drops +again, by the same rule, once the handler returns. 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 @@ -70,14 +84,20 @@ every calculated value. 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, or handed -back by another simulation's `calculate`. In each case, 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). So -any value that depends on the overridden variable at an overlapping period was -stored after the earliest record of it. Uprating and carry-over read which periods hold values at all, -which is why the first uprated or carried-over value also counts. +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 values at all, 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. 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: @@ -117,15 +137,20 @@ written file, and on a timestamp tie the current process's own; between two other processes' files written within one clock tick it cannot tell which came last. -A simulation dump (`dump_simulation`) records which values were inputs. +A simulation dump (`dump_simulation`) holds the values the simulation reads +(on a branch, its own and those it inherited) and records which were inputs. `restore_simulation` restores those as inputs and every other value as -calculated under one number, since the dump does not say what each was -calculated from: an input set on a branch of the restored simulation drops all -of them that are not inputs. A dump written before inputs were recorded is -restored with every value as an input, as before, so such values never drop. +calculated 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 inputs were recorded is restored with every value as an +input, as before, so such values never drop. Two related behaviours: a branch whose input dropped values stops reading -macro-cache files, which are keyed by branch and period but not by inputs; and +macro-cache files, which are keyed by branch and period but not by inputs +(and a branch that read one drops, on any input, everything calculated from +the first read on); and `requires_computation_after` is satisfied by a prerequisite requested before a drop removed its values. diff --git a/policyengine_core/data_storage/store_history.py b/policyengine_core/data_storage/store_history.py index 140976faf..2640cef21 100644 --- a/policyengine_core/data_storage/store_history.py +++ b/policyengine_core/data_storage/store_history.py @@ -60,6 +60,13 @@ class StoreHistory: from holds a value). Such a value depends on which other periods hold values, 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 @@ -72,10 +79,13 @@ class StoreHistory: 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). - self._journal: List[Tuple[str, Optional[Period], int]] = [] + # 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 @@ -93,6 +103,7 @@ def __getstate__(self) -> dict: 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 = [ @@ -100,6 +111,8 @@ def __setstate__(self, state: dict) -> None: 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)) @@ -110,6 +123,7 @@ def copy(self) -> "StoreHistory": 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 @@ -126,6 +140,13 @@ def record_derived(self, variable_name: str, sequence_number: int): 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: @@ -145,8 +166,13 @@ def merge(self, other: "StoreHistory") -> None: (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 period is None: + 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) @@ -176,6 +202,11 @@ def prune(self, since: Optional[int] = None) -> None: 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. @@ -185,8 +216,9 @@ def earliest_dependency(self, variable_name: str, period: Period) -> Optional[in """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``, or the earliest recorded derivation - of one of its values from other periods, whichever came first. + 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``. """ @@ -200,4 +232,6 @@ def earliest_dependency(self, variable_name: str, period: Period) -> Optional[in 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/holder.py b/policyengine_core/holders/holder.py index 37999bcf4..7fd9aa9fd 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -190,6 +190,24 @@ def get_array(self, period: Period, branch_name: str = "default") -> ArrayLike: if default_value is not None: return default_value + def _branch_holding(self, period: Period, branch_name: str) -> Optional[str]: + """The branch whose value :meth:`get_array` returns for ``period``, if any. + + As there: ``branch_name`` itself, else its nearest ancestor, else + ``default``. + """ + names = [branch_name] + if branch_name != "default": + parent = getattr(self.simulation, "parent_branch", None) + while parent is not None: + names.append(parent.branch_name) + parent = getattr(parent, "parent_branch", None) + names.append("default") + for name in names: + if self._get_array_from_storage(period, name) is not None: + return name + return None + def get_memory_usage(self) -> dict: """ Get data about the virtual memory usage of the holder. @@ -308,6 +326,9 @@ def set_input( # On a branch, drop what may have been calculated from the value # this input replaces (see ``Simulation.set_input``). 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 = getattr(simulation, "_inputs_set", 0) + 1 if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() if not hasattr(simulation, "_user_input_contexts"): @@ -318,7 +339,14 @@ def set_input( self.variable.set_input and period.unit != self.variable.definition_period ): - return self.variable.set_input(self, period, array) + self.variable.set_input(self, period, array) + if simulation is not None: + # The handler may have calculated values between its + # stores, from inputs it had not yet replaced. + simulation._drop_values_that_may_depend_on( + self.variable.name, period + ) + return return self._set(period, array, branch_name, validate_nan=True) finally: if simulation is not None: diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 8eca4bd33..500cd4a0b 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1,7 +1,7 @@ import hashlib import tempfile from contextvars import ContextVar -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union import numpy as np import pandas as pd @@ -612,7 +612,15 @@ def calculate( # from (see ``_drop_computed``). self._calculations_in_flight = getattr(self, "_calculations_in_flight", 0) + 1 try: + input_epoch = getattr(self, "_input_epoch", 0) result = self._calculate(variable_name, period) + if input_epoch != getattr(self, "_input_epoch", 0): + # An input set here while it ran (by a formula, say) dropped + # values, so the result was not kept (``_cache_result``). + # Calculate it once more from the new inputs and keep it, as + # a simulation given those inputs first would: later + # uprating and carry-over look for the periods held. + result = self._calculate(variable_name, period) # Satisfies ``requires_computation_after`` from now on, even if a # branch input later drops the values. self._get_requested_variables().add(variable_name) @@ -626,12 +634,14 @@ def calculate( result = self.map_result(result, source_entity, map_to) return result finally: - # 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() - self._calculations_in_flight -= 1 - 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 + self.tracer.record_calculation_end() + self.purge_cache_of_invalid_values() def map_result( self, @@ -743,7 +753,7 @@ 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_epoch = getattr(self, "_input_epoch", 0) + input_state = self._input_state() population = self.get_variable_population(variable_name) holder = population.get_holder(variable_name) variable = self.tax_benefit_system.get_variable( @@ -788,7 +798,11 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if value is not None: # Served without being stored: record it, so values # calculated from it count as later (see ``set_input``). - holder._record_store(period, next_sequence_number()) + # 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: @@ -842,8 +856,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if np.all(~mask): array = holder.default_array() array = self._cast_formula_result(array, variable) - self._cache_result(holder, array, period, input_epoch) - return array + return self._cache_result(holder, array, period, input_state) array = None @@ -944,7 +957,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: array = EnumArray(array, variable.possible_values) array = self._cast_formula_result(array, variable) - self._cache_result(holder, array, period, input_epoch) + array = self._cache_result(holder, array, period, input_state) except SpiralError: array = holder.default_array() @@ -967,27 +980,42 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: if is_cache_available: smc.set_cache_value(cache_path, array) - if hasattr(self, "_fast_cache") and input_epoch == getattr( - self, "_input_epoch", 0 - ): + if hasattr(self, "_fast_cache") and input_state == self._input_state(): self._fast_cache[(variable_name, period)] = array return array - def _cache_result( - self, holder: Holder, array: ArrayLike, period: Period, input_epoch: int - ) -> None: - """Cache a calculated value, unless inputs changed while it was calculated. + def _input_state(self) -> Tuple[int, int]: + """How many drops ran during calculations, and inputs were set, here.""" + return getattr(self, "_input_epoch", 0), getattr(self, "_inputs_set", 0) - 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``). + def _cache_result( + self, + holder: Holder, + array: ArrayLike, + period: Period, + input_state: Tuple[int, int], + ) -> ArrayLike: + """Cache a calculated value, and return the value to use for it. + + ``input_state`` is :meth:`_input_state` 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``). """ - if input_epoch == getattr(self, "_input_epoch", 0): + epoch, inputs_set = input_state + if inputs_set != getattr(self, "_inputs_set", 0) and holder._is_input( + period, self.branch_name + ): + return holder.get_array(period, self.branch_name) + if epoch == getattr(self, "_input_epoch", 0): holder.put_in_cache(array, period, self.branch_name) else: # Not kept, but whatever reads it is stored after it all the same. holder._record_store(period, next_sequence_number()) + 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 @@ -1038,14 +1066,13 @@ def calculate_add( ) ) - input_epoch = getattr(self, "_input_epoch", 0) + input_state = self._input_state() result = sum( self.calculate(variable_name, sub_period) for sub_period in period.get_subperiods(variable.definition_period) ) holder = self.get_holder(variable.name) - self._cache_result(holder, result, period, input_epoch) - return result + return self._cache_result(holder, result, period, input_state) def calculate_divide( self, @@ -1074,12 +1101,11 @@ def calculate_divide( ) if period.unit == periods.MONTH: - input_epoch = getattr(self, "_input_epoch", 0) + input_state = self._input_state() computation_period = period.this_year result = self.calculate(variable_name, period=computation_period) / 12.0 holder = self.get_holder(variable.name) - self._cache_result(holder, result, period, input_epoch) - return result + return self._cache_result(holder, result, period, input_state) elif period.unit == periods.YEAR: return self.calculate(variable_name, period) @@ -1457,16 +1483,20 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non 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). If - there is neither, nothing it holds depends on the value and nothing - is dropped. That is the case when a formula creates the branch while + 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. After such a drop the branch stops reading macro-cache files, which - are keyed by branch and period but not by inputs. + are keyed by branch and period but not by inputs. A calculation that + was running in the branch when its input changed is run once more + from the new inputs, and an input set for the very period it + calculates is its result. What this does not track: a branch given a different tax-benefit system or parameters (call :meth:`drop_computed_arrays` on it); diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index 724757c5d..4c26b89a1 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -62,27 +62,37 @@ def restore_simulation(directory, tax_benefit_system, **kwargs): variables_to_restore = [ variable for variable in os.listdir(directory) if variable != "__entities__" ] - # Inputs first, then every calculated value under one later number: the - # dump does not say what each value was calculated from, so any input set - # on a branch of the restored simulation counts it as possibly dependent. + # 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: _restore_holder(simulation, variable, directory, inputs=True) calculated_number = next_sequence_number() - for variable in variables_to_restore: + 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) return simulation def _dump_holder(holder, directory): disk_storage = holder.create_disk_storage(directory, preserve=True) + branch_name = holder.simulation.branch_name inputs = [] - for period in holder.get_known_periods(): - value = holder.get_array(period) - disk_storage.put(value, period) - if holder._is_input(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_holding(period, branch_name) + if stored_on is None: + continue + disk_storage.put(holder._get_array_from_storage(period, stored_on), period) + if holder._is_input(period, stored_on): inputs.append(str(period)) with open(os.path.join(disk_storage.storage_dir, _INPUTS_FILE), "w") as file: file.write("\n".join(inputs)) @@ -161,10 +171,12 @@ def _restore_holder( # Dumped before inputs were recorded: keep every value, as inputs. input_periods = None + restored = 0 for period in disk_storage.get_known_periods(): is_input = input_periods is None or str(period) in input_periods if is_input != inputs: continue + restored += 1 value = disk_storage.get(period) holder._set( period, @@ -172,3 +184,4 @@ def _restore_holder( is_input=is_input, sequence_number=None if is_input else calculated_number, ) + return restored diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index 892bdf994..8369b3605 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -824,8 +824,8 @@ def test_disk_files_from_another_process_are_not_overwritten(tmp_path, monkeypat assert reader.get(periods.period("2020")).tolist() == [99.0] -def test_result_calculated_across_its_own_input_change_is_not_kept(): - """A formula that read an input and then replaced it returns, but does not cache.""" +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) @@ -837,11 +837,161 @@ def result(person, period): _yearly_variable("result", result), ) branch = SimulationBuilder().build_default_simulation(system).get_branch("b") - assert branch.calculate("result", "2020").tolist() == [2.0] + assert branch.calculate("result", "2020").tolist() == [6.0] assert branch.calculate("result", "2020").tolist() == [6.0] +def test_result_calculated_again_is_kept_for_carry_over(): + """A result that was not kept would leave its period unknown to carry-over.""" + 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 branch.calculate("source", "2022").tolist() == [3.0] # carried over + + +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_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.""" @@ -898,6 +1048,73 @@ def test_restored_simulation_drops_restored_values_an_input_may_have_fed(tmp_pat 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() diff --git a/tests/core/test_branch_input_invalidation_property.py b/tests/core/test_branch_input_invalidation_property.py index 8fbd16cc6..c76d7c448 100644 --- a/tests/core/test_branch_input_invalidation_property.py +++ b/tests/core/test_branch_input_invalidation_property.py @@ -1,7 +1,8 @@ """Branch results do not depend on what was calculated before. Random sequences of calculations, branches (including reused and forgotten -names), inputs and ``drop_computed_arrays`` run on a synthetic system, with +names), inputs, ``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 @@ -13,6 +14,8 @@ from __future__ import annotations import os +import shutil +import tempfile import numpy as np import pytest @@ -23,6 +26,10 @@ 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, @@ -113,6 +120,8 @@ st.tuples(st.just("branch"), simulation_index, branch_name), st.tuples(st.just("forget"), simulation_index), st.tuples(st.just("drop"), simulation_index), + # 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 @@ -151,10 +160,24 @@ ] ) ) +# 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, + override_after_restoring, ), min_size=1, max_size=15, @@ -192,7 +215,8 @@ 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. + 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( @@ -204,6 +228,7 @@ def _run(program, mode): 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) @@ -226,8 +251,8 @@ def _run(program, mode): del simulation.branches[names[0]] del children[(index, names[0])] elif kind == "set": - if index == 0: - continue # Inputs on the root after calculating are out of scope. + 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: @@ -239,6 +264,17 @@ def _run(program, mode): _, _, 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 From 3075a6cb0b2d88032ab0a67c98cc7cac0c431fb3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 02:27:14 -0400 Subject: [PATCH 12/42] DIAG: report traced parameter trees on Windows (to be reverted) Co-Authored-By: Claude Opus 5.5 --- tests/conftest.py | 29 +++++++++++++++++++++++++++++ tests/fixtures/taxbenefitsystems.py | 12 ++++++++++-- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 6ab38421a..848e9a3b1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,8 +2,37 @@ This module directs PyTest to include certain global fixtures when running tests. """ +import gc +import sys + +import pytest + pytest_plugins = [ "tests.fixtures.entities", "tests.fixtures.simulations", "tests.fixtures.taxbenefitsystems", ] + +# TEMPORARY DIAGNOSTIC (to be reverted): which test first leaves a traced +# parameter tree behind on Windows. +_DIAG_TRACED = set() + + +@pytest.fixture(autouse=True) +def _diag_parameter_trace(request): + yield + if sys.platform != "win32": + return + from policyengine_core.parameters import ParameterNode + + for obj in gc.get_objects(): + if ( + type(obj) is ParameterNode + and obj.name == "" + and obj.trace + and id(obj) not in _DIAG_TRACED + ): + _DIAG_TRACED.add(id(obj)) + sys.__stderr__.write( + f"\nDIAG traced root {id(obj):x} after {request.node.nodeid}\n" + ) diff --git a/tests/fixtures/taxbenefitsystems.py b/tests/fixtures/taxbenefitsystems.py index af8b1a8db..6c8404468 100644 --- a/tests/fixtures/taxbenefitsystems.py +++ b/tests/fixtures/taxbenefitsystems.py @@ -4,8 +4,16 @@ @pytest.fixture(scope="module") -def tax_benefit_system(): - return CountryTaxBenefitSystem() +def tax_benefit_system(request): + system = CountryTaxBenefitSystem() + import sys + + if sys.platform == "win32": # TEMPORARY DIAGNOSTIC (to be reverted) + sys.__stderr__.write( + f"\nDIAG fixture {request.module.__name__} system {id(system):x} " + f"params {id(system.parameters):x} trace {system.parameters.trace}\n" + ) + return system @pytest.fixture From acc5f58e51f6dac01a78d0ef76c46177b684d828 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 02:40:38 -0400 Subject: [PATCH 13/42] Give the disk-restore test's later process its own token, and age all earlier files It faked a later process by restarting the counter only, so with the same token its first store could reuse the writer's first file name; and it aged only the writer's last file, so on a coarse file clock (Windows) the other four could tie with the later write. Passes with every write given the same timestamp. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_branch_input_invalidation.py | 23 ++++++++++---------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index 8369b3605..04a3bb9b3 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -709,27 +709,28 @@ def result(person, period): assert branch.calculate("result", "2020").tolist() == [6.0] -def test_disk_restore_reads_the_latest_file_of_each_key(tmp_path): +def test_disk_restore_reads_the_latest_file_of_each_key(tmp_path, monkeypatch): """Sequence numbers restart in each process; restore goes by write time.""" import itertools import os + import policyengine_core.data_storage.on_disk_storage as on_disk_storage import policyengine_core.data_storage.store_history as store_history from policyengine_core.data_storage import OnDiskStorage + monkeypatch.setattr(on_disk_storage, "_PROCESS_TOKEN", "aaaaaaaaaaaa") writer = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) for _ in range(5): writer.put(np.array([1.0]), periods.period("2020")) - first = writer._files["default_2020"] - saved = store_history._sequence - store_history._sequence = itertools.count(1) # a later process - try: - later = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) - later.put(np.array([2.0]), periods.period("2020")) - stat = os.stat(first) - os.utime(first, ns=(stat.st_atime_ns, stat.st_mtime_ns - 10**9)) - finally: - store_history._sequence = saved + # A later process: its counter restarts, and it has its own token. + monkeypatch.setattr(store_history, "_sequence", itertools.count(1)) + monkeypatch.setattr(on_disk_storage, "_PROCESS_TOKEN", "bbbbbbbbbbbb") + later = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) + later.put(np.array([2.0]), periods.period("2020")) + # The writer's files clearly earlier, even on a coarse file clock. + for path in tmp_path.glob("default_2020.aaaaaaaaaaaa.*.npy"): + stat = os.stat(path) + os.utime(path, ns=(stat.st_atime_ns, stat.st_mtime_ns - 10**9)) reader = OnDiskStorage(str(tmp_path), preserve_storage_dir=True) reader.restore() From a81adb43892c9d9c4573d6ca563357017d30a224 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 02:40:38 -0400 Subject: [PATCH 14/42] DIAG: write the diagnostic through the terminal reporter (to be reverted) Co-Authored-By: Claude Opus 5.5 --- tests/conftest.py | 15 ++++++++++++--- tests/fixtures/taxbenefitsystems.py | 13 +++++++++---- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 848e9a3b1..58641114e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,6 +3,7 @@ """ import gc +import os import sys import pytest @@ -21,7 +22,7 @@ @pytest.fixture(autouse=True) def _diag_parameter_trace(request): yield - if sys.platform != "win32": + if sys.platform != "win32" and not os.environ.get("DIAG_LOCAL"): return from policyengine_core.parameters import ParameterNode @@ -33,6 +34,14 @@ def _diag_parameter_trace(request): and id(obj) not in _DIAG_TRACED ): _DIAG_TRACED.add(id(obj)) - sys.__stderr__.write( - f"\nDIAG traced root {id(obj):x} after {request.node.nodeid}\n" + _diag_write( + request.config, + f"DIAG traced root {id(obj):x} after {request.node.nodeid}", ) + + +def _diag_write(config, line): + # Bypasses output capture, so it shows for passing tests too. + reporter = config.pluginmanager.get_plugin("terminalreporter") + if reporter is not None: + reporter.write_line(line) diff --git a/tests/fixtures/taxbenefitsystems.py b/tests/fixtures/taxbenefitsystems.py index 6c8404468..f66bf7491 100644 --- a/tests/fixtures/taxbenefitsystems.py +++ b/tests/fixtures/taxbenefitsystems.py @@ -8,10 +8,15 @@ def tax_benefit_system(request): system = CountryTaxBenefitSystem() import sys - if sys.platform == "win32": # TEMPORARY DIAGNOSTIC (to be reverted) - sys.__stderr__.write( - f"\nDIAG fixture {request.module.__name__} system {id(system):x} " - f"params {id(system.parameters):x} trace {system.parameters.trace}\n" + if sys.platform == "win32" or __import__("os").environ.get( + "DIAG_LOCAL" + ): # TEMPORARY DIAGNOSTIC (to be reverted) + from tests.conftest import _diag_write + + _diag_write( + request.config, + f"DIAG fixture {request.module.__name__} system {id(system):x} " + f"params {id(system.parameters):x} trace {system.parameters.trace}", ) return system From 1d6221f1c48f495942e73321723fa4aca4d1671d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 02:46:32 -0400 Subject: [PATCH 15/42] DIAG: tolerate half-built parameter nodes (to be reverted) Co-Authored-By: Claude Opus 5.5 --- tests/conftest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 58641114e..34bb736e4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -29,8 +29,8 @@ def _diag_parameter_trace(request): for obj in gc.get_objects(): if ( type(obj) is ParameterNode - and obj.name == "" - and obj.trace + and getattr(obj, "name", None) == "" + and getattr(obj, "trace", False) and id(obj) not in _DIAG_TRACED ): _DIAG_TRACED.add(id(obj)) From 10b1d01d071804e6b6bcd8adaffd1507355e8684 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 02:54:27 -0400 Subject: [PATCH 16/42] Revert the Windows diagnostics (3075a6cb, a81adb43, 1d6221f1) They found the cause: test_disk_restore_reads_the_latest_file_of_each_key failed on Windows' coarse file clock (fixed in acc5f58e), and its rerun by pytest-rerunfailures dropped the module's setup-state entry without running its finalizers, so the module-scoped tax_benefit_system leaked into later modules; test_branch_shared_arrays traced it, and test_parameters and test_reforms then got a traced parameter tree. Co-Authored-By: Claude Opus 5.5 --- tests/conftest.py | 38 ----------------------------- tests/fixtures/taxbenefitsystems.py | 17 ++----------- 2 files changed, 2 insertions(+), 53 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 34bb736e4..6ab38421a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,46 +2,8 @@ This module directs PyTest to include certain global fixtures when running tests. """ -import gc -import os -import sys - -import pytest - pytest_plugins = [ "tests.fixtures.entities", "tests.fixtures.simulations", "tests.fixtures.taxbenefitsystems", ] - -# TEMPORARY DIAGNOSTIC (to be reverted): which test first leaves a traced -# parameter tree behind on Windows. -_DIAG_TRACED = set() - - -@pytest.fixture(autouse=True) -def _diag_parameter_trace(request): - yield - if sys.platform != "win32" and not os.environ.get("DIAG_LOCAL"): - return - from policyengine_core.parameters import ParameterNode - - for obj in gc.get_objects(): - if ( - type(obj) is ParameterNode - and getattr(obj, "name", None) == "" - and getattr(obj, "trace", False) - and id(obj) not in _DIAG_TRACED - ): - _DIAG_TRACED.add(id(obj)) - _diag_write( - request.config, - f"DIAG traced root {id(obj):x} after {request.node.nodeid}", - ) - - -def _diag_write(config, line): - # Bypasses output capture, so it shows for passing tests too. - reporter = config.pluginmanager.get_plugin("terminalreporter") - if reporter is not None: - reporter.write_line(line) diff --git a/tests/fixtures/taxbenefitsystems.py b/tests/fixtures/taxbenefitsystems.py index f66bf7491..af8b1a8db 100644 --- a/tests/fixtures/taxbenefitsystems.py +++ b/tests/fixtures/taxbenefitsystems.py @@ -4,21 +4,8 @@ @pytest.fixture(scope="module") -def tax_benefit_system(request): - system = CountryTaxBenefitSystem() - import sys - - if sys.platform == "win32" or __import__("os").environ.get( - "DIAG_LOCAL" - ): # TEMPORARY DIAGNOSTIC (to be reverted) - from tests.conftest import _diag_write - - _diag_write( - request.config, - f"DIAG fixture {request.module.__name__} system {id(system):x} " - f"params {id(system.parameters):x} trace {system.parameters.trace}", - ) - return system +def tax_benefit_system(): + return CountryTaxBenefitSystem() @pytest.fixture From 9f8d513c2cbb4bf6a9a275385575c91fc9708648 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 03:19:09 -0400 Subject: [PATCH 17/42] Close the fifth review's gaps: macro reads, input wins, sums, failing handlers From the round-5 soundness review (on 2b31dcb7); all were also on master: - A branch stops reading macro-cache files as soon as an input is set on it, not only when the input drops something: a file for its name may come from another simulation's branch of the same name, read before any record exists. - An input set during a calculation for its own period wins under any branch name the simulation reads (Holder.set_input stores under "default"), and also when the formula returns None and core would fall back to the default for want of an earlier period. - A direct calculate_add sums again when a term changed an input while summing (an earlier term could be obsolete). - A custom set_input handler that calculated and then raised: the branch still drops again (the inputs it stored stay). The drop after a handler now runs only if the handler calculated anything. - calculate runs a calculation again until a run changes no input, at most ten times (was once), so a formula changing a second input on its rerun no longer leaves the period unknown to carry-over. - Docs: the above, and that disk-backed branches whose names contain "_" cannot be dumped (OnDiskStorage key parsing, as on master; #552 fixes it). Co-Authored-By: Claude Opus 5.5 --- .../branch-set-input-invalidation.changed.md | 2 +- docs/usage/branches.md | 34 +-- policyengine_core/holders/holder.py | 22 +- policyengine_core/simulations/simulation.py | 87 ++++++-- tests/core/test_branch_input_invalidation.py | 201 ++++++++++++++++++ 5 files changed, 303 insertions(+), 43 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.changed.md b/changelog.d/branch-set-input-invalidation.changed.md index 78a326820..a4a16ceef 100644 --- a/changelog.d/branch-set-input-invalidation.changed.md +++ b/changelog.d/branch-set-input-invalidation.changed.md @@ -1 +1 @@ -Disk storage writes a new file for every store, so recalculating a value no longer changes what branches sharing the directory read; apply_reform keeps inputs by a flag on each stored array instead of replaying _user_input_keys; values a custom set_input handler calculates are no longer treated as inputs; dump_simulation records which values were inputs (inputs.txt), and restore_simulation restores only those as inputs; a calculation whose formula changes an input it read is run once more from the new input. +Disk storage writes a new file for every store, so recalculating a value no longer changes what branches sharing the directory read; apply_reform keeps inputs by a flag on each stored array instead of replaying _user_input_keys; values a custom set_input handler calculates are no longer treated as inputs; dump_simulation records which values were inputs (inputs.txt), and restore_simulation restores only those as inputs; a calculation whose formula changes an input it read is run again from the new inputs until a run changes none (at most ten times); a branch stops reading macro-cache files once an input is set on it. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 673c6bdc4..170d6e286 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -58,16 +58,19 @@ 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. A calculation that was running then (one whose formula sets an input, say) may -have read the replaced value, so its result is not kept; `calculate` runs it -once more from the new inputs and keeps that result, so uprating and -carry-over find its period as they would had the input come first. If the -second run changes an input again, its result is returned but not kept. An -input set for the very period being calculated, while it is calculated, is -the result, as it would be had it been set first. +have read the replaced value, so its result is not kept; `calculate` (and a +direct `calculate_add`) runs it again from the new inputs until a run changes +no input, and keeps that result, so uprating and carry-over find its period as +they would had the inputs come first. A formula that changes an input on every +run stops after ten reruns; its last result is returned but not kept. An input +set for the very period being calculated, while it is calculated (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. A custom `set_input` handler that calculates values between its own stores -calculates them from inputs it has not yet replaced, so the branch drops -again, by the same rule, once the handler returns. +calculates them from inputs it has not yet replaced, so if it calculated +anything, 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 @@ -145,12 +148,15 @@ 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 inputs were recorded is restored with every value as an -input, as before, so such values never drop. - -Two related behaviours: a branch whose input dropped values stops reading -macro-cache files, which are keyed by branch and period but not by inputs -(and a branch that read one drops, on any input, everything calculated from -the first read on); and +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, 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, everything calculated from the first read on; and `requires_computation_after` is satisfied by a prerequisite requested before a drop removed its values. diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 7fd9aa9fd..49b4059d3 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -339,14 +339,20 @@ def set_input( self.variable.set_input and period.unit != self.variable.definition_period ): - self.variable.set_input(self, period, array) - if simulation is not None: - # The handler may have calculated values between its - # stores, from inputs it had not yet replaced. - simulation._drop_values_that_may_depend_on( - self.variable.name, period - ) - return + started = getattr(simulation, "_calculations_started", 0) + try: + return self.variable.set_input(self, period, array) + finally: + if ( + simulation is not None + and getattr(simulation, "_calculations_started", 0) != 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: diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 500cd4a0b..92ada8f83 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -129,6 +129,12 @@ def _uprating_index_value(parameter, instant) -> Optional[float]: "_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 + class Simulation: """ @@ -611,15 +617,21 @@ def calculate( # local variables; a drop meanwhile must not forget what they came # from (see ``_drop_computed``). self._calculations_in_flight = getattr(self, "_calculations_in_flight", 0) + 1 + # Lets a custom ``set_input`` handler tell whether it calculated. + self._calculations_started = getattr(self, "_calculations_started", 0) + 1 try: input_epoch = getattr(self, "_input_epoch", 0) result = self._calculate(variable_name, period) - if input_epoch != getattr(self, "_input_epoch", 0): + for _ in range(_RERUNS_AFTER_INPUT_CHANGE): + if input_epoch == getattr(self, "_input_epoch", 0): + break # An input set here while it ran (by a formula, say) dropped # values, so the result was not kept (``_cache_result``). - # Calculate it once more from the new inputs and keep it, as - # a simulation given those inputs first would: later - # uprating and carry-over look for the periods held. + # Calculate it again from the new inputs, until a run changes + # none, and keep that result, as a simulation given those + # inputs first would: later uprating and carry-over look for + # the periods held. + input_epoch = getattr(self, "_input_epoch", 0) result = self._calculate(variable_name, period) # Satisfies ``requires_computation_after`` from now on, even if a # branch input later drops the values. @@ -937,6 +949,13 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # win over a later "2024-06" monthly value (bug H1). last_known_period = max(known_periods, key=lambda p: p.start) if last_known_period.start > period.start: + stored_input = ( + self._input_set_meanwhile(holder, period, input_state) + if input_state[1] != getattr(self, "_inputs_set", 0) + else None + ) + if stored_input is not None: + return stored_input return holder.default_array() # Pass branch_name through so auto-carry-over respects the # active branch instead of reaching for the "default" @@ -989,6 +1008,21 @@ def _input_state(self) -> Tuple[int, int]: """How many drops ran during calculations, and inputs were set, here.""" return getattr(self, "_input_epoch", 0), getattr(self, "_inputs_set", 0) + def _input_set_meanwhile( + self, holder: Holder, period: Period, input_state: Tuple[int, int] + ) -> Optional[ArrayLike]: + """The input this simulation reads for ``period``, if one was set during the calculation. + + Only called once an input was set since ``input_state``. A value the + simulation could read for ``period`` when the calculation began + would have been returned without calculating, so an input it reads + now was set meanwhile, under any branch name it reads. + """ + stored_on = holder._branch_holding(period, self.branch_name) + if stored_on is not None and holder._is_input(period, stored_on): + return holder._get_array_from_storage(period, stored_on) + return None + def _cache_result( self, holder: Holder, @@ -1006,10 +1040,10 @@ def _cache_result( result is returned but not kept (see ``_drop_computed``). """ epoch, inputs_set = input_state - if inputs_set != getattr(self, "_inputs_set", 0) and holder._is_input( - period, self.branch_name - ): - return holder.get_array(period, self.branch_name) + if inputs_set != getattr(self, "_inputs_set", 0): + stored_input = self._input_set_meanwhile(holder, period, input_state) + if stored_input is not None: + return stored_input if epoch == getattr(self, "_input_epoch", 0): holder.put_in_cache(array, period, self.branch_name) else: @@ -1066,11 +1100,21 @@ def calculate_add( ) ) + def total(): + return sum( + self.calculate(variable_name, sub_period) + for sub_period in period.get_subperiods(variable.definition_period) + ) + input_state = self._input_state() - result = sum( - self.calculate(variable_name, sub_period) - for sub_period in period.get_subperiods(variable.definition_period) - ) + result = total() + for _ in range(_RERUNS_AFTER_INPUT_CHANGE): + if input_state[0] == getattr(self, "_input_epoch", 0): + break + # An input changed while summing: earlier terms may be obsolete + # (see ``calculate``). Sum again. + input_state = self._input_state() + result = total() holder = self.get_holder(variable.name) return self._cache_result(holder, result, period, input_state) @@ -1492,11 +1536,12 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non calculated from the variable; then only what came back from there, and what was calculated after, is dropped. - After such a drop the branch stops reading macro-cache files, which - are keyed by branch and period but not by inputs. A calculation that - was running in the branch when its input changed is run once more - from the new inputs, and an input set for the very period it - calculates is its result. + 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. + A calculation that was running in the branch when its input changed + is run again from the new inputs until a run changes none (at most + ten times; after that its result is returned but not kept), and an + input set for the very period it calculates is its result. What this does not track: a branch given a different tax-benefit system or parameters (call :meth:`drop_computed_arrays` on it); @@ -1564,6 +1609,11 @@ def _drop_values_that_may_depend_on( """ if getattr(self, "parent_branch", None) is None: return 0 + # Macro-cache files are keyed by branch and period, not by inputs, so + # a file for this branch's 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 has read any. + self.macro_cache_read = False since = self._get_store_history().earliest_dependency( variable_name, periods.period(period) ) @@ -1573,9 +1623,6 @@ def _drop_values_that_may_depend_on( since = 0 if since is None: return 0 - # Macro-cache files are keyed by branch and period, not by inputs, so - # those this branch wrote may hold values calculated from the old one. - self.macro_cache_read = False return self._drop_computed(since) def _drop_computed(self, since: Optional[int] = None) -> int: diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index 04a3bb9b3..c3b84d28a 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -878,6 +878,124 @@ def formula_2022(person, period): assert branch.calculate("source", "2022").tolist() == [3.0] # carried over +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 branch.calculate("source", "2022").tolist() == [3.0] # carried over + + +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_input_a_formula_sets_for_its_own_period_wins(): """As it would had the input been set before the calculation.""" @@ -930,6 +1048,89 @@ def formula(person, period): 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_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 From ed843df9d953c916a0c712fda32a099ff93c9cd9 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 03:50:52 -0400 Subject: [PATCH 18/42] Close the sixth review's gaps: input wins only if stored meanwhile, one rerun budget From the round-6 soundness review (on 9f8d513c): - Direct calculate_divide/calculate_add could return an input stored before the calculation (they do not look in storage first), after any unrelated input change: 99 instead of 10, 100 instead of 12. An input now wins only if its sequence number is after the calculation began. - Only the outermost calculation in a simulation (calculate or a direct calculate_add) runs again after an input change; inner ones run with it. The ten-rerun budget no longer multiplies through nesting (11/121/1,331 leaf runs at depths 1/2/3, now 11 each). - A result calculated across an input change is no longer written to the macro cache (only kept results are), so a branch of the same name elsewhere cannot read it (also on master). - The post-handler drop counts _calculate, not only calculate, as calculating (a handler's private carry-over calculation was missed). - Docs: the above, and the rerun cap's consequence for carry-over. Co-Authored-By: Claude Opus 5.5 --- docs/usage/branches.md | 24 +-- policyengine_core/holders/holder.py | 12 ++ policyengine_core/simulations/simulation.py | 68 +++++--- tests/core/test_branch_input_invalidation.py | 165 +++++++++++++++++++ 4 files changed, 233 insertions(+), 36 deletions(-) diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 170d6e286..7b6495818 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -58,19 +58,23 @@ 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. A calculation that was running then (one whose formula sets an input, say) may -have read the replaced value, so its result is not kept; `calculate` (and a -direct `calculate_add`) runs it again from the new inputs until a run changes -no input, and keeps that result, so uprating and carry-over find its period as -they would had the inputs come first. A formula that changes an input on every -run stops after ten reruns; its last result is returned but not kept. An input -set for the very period being calculated, while it is calculated (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. +have read the replaced value, so its result is kept neither in storage nor in +the macro cache. The outermost calculation running in the simulation (a +`calculate`, or a direct `calculate_add`) then runs again from the new inputs, +inner calculations included, until a run changes no input, and keeps that +result, so uprating and carry-over find its period as they would had the inputs +come first. After ten reruns it stops: a formula that keeps changing inputs +returns its last result without keeping it, and a later uprating or carry-over +may then not find that period. 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. 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, the branch drops again, by the same rule, once the handler returns -or raises (the inputs it stored before raising stay). +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 diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 49b4059d3..964f37ada 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -136,6 +136,18 @@ def _is_input(self, period: Period, branch_name: str = "default") -> bool: and f"{branch_name}_{period}" in self._disk_storage._input_keys ) + 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 _has_unnumbered_values(self) -> bool: return self._memory_storage.has_unnumbered_values() or ( self._disk_storage is not None diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 92ada8f83..6ff35904d 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -616,13 +616,14 @@ def calculate( # 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 here runs again after an input + # change: running it again runs the inner ones again too. + outermost = not getattr(self, "_calculations_in_flight", 0) self._calculations_in_flight = getattr(self, "_calculations_in_flight", 0) + 1 - # Lets a custom ``set_input`` handler tell whether it calculated. - self._calculations_started = getattr(self, "_calculations_started", 0) + 1 try: input_epoch = getattr(self, "_input_epoch", 0) result = self._calculate(variable_name, period) - for _ in range(_RERUNS_AFTER_INPUT_CHANGE): + for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): if input_epoch == getattr(self, "_input_epoch", 0): break # An input set here while it ran (by a formula, say) dropped @@ -765,7 +766,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._input_state() + input_state = self._calculation_start() + # Lets a custom ``set_input`` handler tell whether it calculated. + self._calculations_started = getattr(self, "_calculations_started", 0) + 1 population = self.get_variable_population(variable_name) holder = population.get_holder(variable_name) variable = self.tax_benefit_system.get_variable( @@ -854,7 +857,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 input_state[:2] == self._input_state(): smc.set_cache_value(cache_path, values) return values @@ -950,7 +953,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: last_known_period = max(known_periods, key=lambda p: p.start) if last_known_period.start > period.start: stored_input = ( - self._input_set_meanwhile(holder, period, input_state) + self._input_set_meanwhile(holder, period, input_state[2]) if input_state[1] != getattr(self, "_inputs_set", 0) else None ) @@ -996,10 +999,13 @@ 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 = input_state[:2] == self._input_state() + if is_cache_available and unchanged: smc.set_cache_value(cache_path, array) - if hasattr(self, "_fast_cache") and input_state == self._input_state(): + if hasattr(self, "_fast_cache") and unchanged: self._fast_cache[(variable_name, period)] = array return array @@ -1008,18 +1014,25 @@ def _input_state(self) -> Tuple[int, int]: """How many drops ran during calculations, and inputs were set, here.""" return getattr(self, "_input_epoch", 0), getattr(self, "_inputs_set", 0) + def _calculation_start(self) -> Tuple[int, int, int]: + """:meth:`_input_state` when a calculation begins, and a sequence number then.""" + return (*self._input_state(), next_sequence_number()) + def _input_set_meanwhile( - self, holder: Holder, period: Period, input_state: Tuple[int, int] + self, holder: Holder, period: Period, started_at: int ) -> Optional[ArrayLike]: - """The input this simulation reads for ``period``, if one was set during the calculation. + """The input this simulation reads for ``period``, if it was stored after ``started_at``. - Only called once an input was set since ``input_state``. A value the - simulation could read for ``period`` when the calculation began - would have been returned without calculating, so an input it reads - now was set meanwhile, under any branch name it reads. + 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_holding(period, self.branch_name) - if stored_on is not None and holder._is_input(period, stored_on): + 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 @@ -1028,20 +1041,20 @@ def _cache_result( holder: Holder, array: ArrayLike, period: Period, - input_state: Tuple[int, int], + input_state: Tuple[int, int, int], ) -> ArrayLike: """Cache a calculated value, and return the value to use for it. - ``input_state`` is :meth:`_input_state` when the calculation began. + ``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 = input_state + epoch, inputs_set, started_at = input_state if inputs_set != getattr(self, "_inputs_set", 0): - stored_input = self._input_set_meanwhile(holder, period, input_state) + stored_input = self._input_set_meanwhile(holder, period, started_at) if stored_input is not None: return stored_input if epoch == getattr(self, "_input_epoch", 0): @@ -1106,14 +1119,16 @@ def total(): for sub_period in period.get_subperiods(variable.definition_period) ) - input_state = self._input_state() + # As in ``calculate``, only an outermost sum runs again. + outermost = not getattr(self, "_calculations_in_flight", 0) + input_state = self._calculation_start() result = total() - for _ in range(_RERUNS_AFTER_INPUT_CHANGE): + for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): if input_state[0] == getattr(self, "_input_epoch", 0): break # An input changed while summing: earlier terms may be obsolete # (see ``calculate``). Sum again. - input_state = self._input_state() + input_state = self._calculation_start() result = total() holder = self.get_holder(variable.name) return self._cache_result(holder, result, period, input_state) @@ -1145,7 +1160,7 @@ def calculate_divide( ) if period.unit == periods.MONTH: - input_state = self._input_state() + input_state = self._calculation_start() computation_period = period.this_year result = self.calculate(variable_name, period=computation_period) / 12.0 holder = self.get_holder(variable.name) @@ -1539,9 +1554,10 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non 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. A calculation that was running in the branch when its input changed - is run again from the new inputs until a run changes none (at most - ten times; after that its result is returned but not kept), and an - input set for the very period it calculates is its result. + is not kept; the outermost one is run again from the new inputs until + a run changes none (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. What this does not track: a branch given a different tax-benefit system or parameters (call :meth:`drop_computed_arrays` on it); diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index c3b84d28a..da3627b64 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -38,6 +38,7 @@ 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, @@ -996,6 +997,121 @@ def source(person, period): 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] + + def test_input_a_formula_sets_for_its_own_period_wins(): """As it would had the input been set before the calculation.""" @@ -1131,6 +1247,55 @@ def set_cache_value(self, path, value): 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 From 116b01818dba55d7858a5fb3fe2a713ea33cd9d5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 04:23:04 -0400 Subject: [PATCH 19/42] Close the seventh review's gaps: refused results stay unkept across simulations From the round-7 soundness review (on ed843df9): - A parent formula calling into a branch whose formula calls back into the parent and then changes the branch's input: the parent kept the branch's refused inner result (2 instead of 6; also through calculate_add and calculate_divide). Drops made while calculations run now also bump a process-wide count, and no result calculated across one, in any simulation, is kept. Which calculation runs again is still decided per simulation: the outermost one in the simulation whose input changed, so a temporary branch settles locally and nesting does not multiply reruns. - A direct calculate_add counts as in flight while it sums, so its terms do not each get their own reruns (121 runs, now 11). - Inputs are numbered when stored, even if the caller passes a number (a private _set with an earlier number lost its input to "input wins"). - Tests for each; docs. Co-Authored-By: Claude Opus 5.5 --- .../branch-set-input-invalidation.changed.md | 2 +- docs/usage/branches.md | 19 ++- policyengine_core/holders/holder.py | 4 +- policyengine_core/simulations/simulation.py | 56 +++++--- tests/core/test_branch_input_invalidation.py | 128 ++++++++++++++++++ 5 files changed, 182 insertions(+), 27 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.changed.md b/changelog.d/branch-set-input-invalidation.changed.md index a4a16ceef..59dd4d4c5 100644 --- a/changelog.d/branch-set-input-invalidation.changed.md +++ b/changelog.d/branch-set-input-invalidation.changed.md @@ -1 +1 @@ -Disk storage writes a new file for every store, so recalculating a value no longer changes what branches sharing the directory read; apply_reform keeps inputs by a flag on each stored array instead of replaying _user_input_keys; values a custom set_input handler calculates are no longer treated as inputs; dump_simulation records which values were inputs (inputs.txt), and restore_simulation restores only those as inputs; a calculation whose formula changes an input it read is run again from the new inputs until a run changes none (at most ten times); a branch stops reading macro-cache files once an input is set on it. +Disk storage writes a new file for every store, so recalculating a value no longer changes what branches sharing the directory read; apply_reform keeps inputs by a flag on each stored array instead of replaying _user_input_keys; values a custom set_input handler calculates are no longer treated as inputs; dump_simulation records which values were inputs (inputs.txt), and restore_simulation restores only those as inputs; a calculation running when an input set on its simulation drops values is not cached (nor is anything calculated meanwhile in other simulations), and the outermost one in that simulation is run again from the new inputs until a run changes none (at most ten times); a branch stops reading macro-cache files once an input is set on it. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 7b6495818..d3350ba5c 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -59,13 +59,18 @@ 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. A calculation that was running 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. The outermost calculation running in the simulation (a -`calculate`, or a direct `calculate_add`) then runs again from the new inputs, -inner calculations included, until a run changes no input, and keeps that -result, so uprating and carry-over find its period as they would had the inputs -come first. After ten reruns it stops: a formula that keeps changing inputs -returns its last result without keeping it, and a later uprating or carry-over -may then not find that period. An input stored for the very period being +the macro cache; nor is any result calculated meanwhile in another simulation, +which may have called into this one (a parent formula calculating in a branch +whose formula calls back into the parent, say). The outermost calculation +running in the simulation whose input changed (a `calculate`, or a direct +`calculate_add`, whose terms run within it) then runs again from the new +inputs, inner calculations included, until a run changes no input, and keeps +that result, so uprating and carry-over find its period as they would had the +inputs come first. After ten reruns it stops: a formula that keeps changing +inputs returns its last result without keeping it, and a later uprating or +carry-over may then not find that period. The budget is per simulation whose +input changes, so calculations nested across several such simulations can +rerun more. 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. diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 964f37ada..f5a13ee69 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -458,7 +458,9 @@ def _set( >= self.simulation.memory_config.max_memory_occupation_pc ) - if sequence_number is None: + if sequence_number is None or is_input: + # 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( diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 6ff35904d..503ddd8dc 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -135,6 +135,14 @@ def _uprating_index_value(parameter, instant) -> Optional[float]: # returned but not kept. _RERUNS_AFTER_INPUT_CHANGE = 10 +# Counts drops made while calculations were running (``_drop_computed``), in +# any simulation. A result calculated across one, wherever it happened, may +# come from the replaced value (a formula in one simulation can calculate in +# another, which calls back into the first), so it is not kept. Each +# simulation also counts its own (``_input_epoch``): those decide which +# calculation runs again. +_input_epoch_anywhere = [0] + class Simulation: """ @@ -616,8 +624,8 @@ def calculate( # 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 here runs again after an input - # change: running it again runs the inner ones again too. + # Only the outermost calculation in this simulation runs again after + # an input change here: running it again runs the inner ones too. outermost = not getattr(self, "_calculations_in_flight", 0) self._calculations_in_flight = getattr(self, "_calculations_in_flight", 0) + 1 try: @@ -1012,7 +1020,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: def _input_state(self) -> Tuple[int, int]: """How many drops ran during calculations, and inputs were set, here.""" - return getattr(self, "_input_epoch", 0), getattr(self, "_inputs_set", 0) + return _input_epoch_anywhere[0], getattr(self, "_inputs_set", 0) def _calculation_start(self) -> Tuple[int, int, int]: """:meth:`_input_state` when a calculation begins, and a sequence number then.""" @@ -1057,7 +1065,7 @@ def _cache_result( stored_input = self._input_set_meanwhile(holder, period, started_at) if stored_input is not None: return stored_input - if epoch == getattr(self, "_input_epoch", 0): + if epoch == _input_epoch_anywhere[0]: holder.put_in_cache(array, period, self.branch_name) else: # Not kept, but whatever reads it is stored after it all the same. @@ -1119,17 +1127,25 @@ def total(): for sub_period in period.get_subperiods(variable.definition_period) ) - # As in ``calculate``, only an outermost sum runs again. + # 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 getattr(self, "_calculations_in_flight", 0) - input_state = self._calculation_start() - result = total() - for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): - if input_state[0] == getattr(self, "_input_epoch", 0): - break - # An input changed while summing: earlier terms may be obsolete - # (see ``calculate``). Sum again. + self._calculations_in_flight = getattr(self, "_calculations_in_flight", 0) + 1 + try: + input_epoch = getattr(self, "_input_epoch", 0) input_state = self._calculation_start() result = total() + for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): + if input_epoch == getattr(self, "_input_epoch", 0): + break + # An input changed while summing: earlier terms may be + # obsolete (see ``calculate``). Sum again. + input_epoch = getattr(self, "_input_epoch", 0) + input_state = self._calculation_start() + result = total() + finally: + self._calculations_in_flight -= 1 holder = self.get_holder(variable.name) return self._cache_result(holder, result, period, input_state) @@ -1553,11 +1569,12 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non 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. - A calculation that was running in the branch when its input changed - is not kept; the outermost one is run again from the new inputs until - a run changes none (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. + A calculation that was running when its input changed, in the branch + or in a simulation calling into it, is not kept; the outermost one in + the branch is run again from the new inputs until a run changes none + (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. What this does not track: a branch given a different tax-benefit system or parameters (call :meth:`drop_computed_arrays` on it); @@ -1652,8 +1669,11 @@ def _drop_computed(self, since: Optional[int] = None) -> int: if getattr(self, "_calculations_in_flight", 0): # A formula running here may still hold values calculated from # what the records describe: keep the records, and keep none of - # the results the running calculations return (``_cache_result``). + # the results running calculations return, here or in the + # simulations calling into this one (``_cache_result``); the + # outermost one here runs again (``calculate``). self._input_epoch = getattr(self, "_input_epoch", 0) + 1 + _input_epoch_anywhere[0] += 1 return dropped # Nothing the simulation still holds was calculated from what the # records numbered ``since`` or later describe, except the inputs it diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index da3627b64..908a69f23 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -1112,6 +1112,134 @@ def branch(handler_input): 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_a_formula_sets_for_its_own_period_wins(): """As it would had the input been set before the calculation.""" From 359c97fb3983f52d95ccdf5497617f5ab905520e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 07:57:50 -0400 Subject: [PATCH 20/42] Close the eighth review's regression: scope "not kept" to callers and family From the round-8 soundness review (on 116b0181): refusing to cache any result calculated across an in-flight drop anywhere in the process let an unrelated simulation's input change leave a period uncached, so a later carry-over returned 0 instead of 6 (master, 9f8d513c and ed843df9 give 6). A result calculated across an input change is now kept out of caches only where it can matter: - in the simulation whose input changed (its own epoch, as before); - in any calculation that received a refused result, directly or through other calculations, in any simulation (a per-frame flag in _calculation_frames, passed to the caller when a frame returns); - in calculations of the same family still running in this context, which may hold values read before the drop (their cache epoch is bumped). Unrelated simulations and other threads keep caching. Reruns stay with the outermost calculation in the simulation whose input changed; calculate_add and calculate_divide run in frames too. Docs: the scope, and that values a formula already read from a branch before a formula there changed the branch's input stay read (its result is returned but not kept). Tests isolate each part (taint across families, the family bump, the unrelated family); new mutants M38/M39 are caught. Co-Authored-By: Claude Opus 5.5 --- .../branch-set-input-invalidation.changed.md | 2 +- docs/usage/branches.md | 14 +- policyengine_core/simulations/simulation.py | 101 +++++++++---- tests/core/test_branch_input_invalidation.py | 141 ++++++++++++++++++ 4 files changed, 227 insertions(+), 31 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.changed.md b/changelog.d/branch-set-input-invalidation.changed.md index 59dd4d4c5..1fdef0f1a 100644 --- a/changelog.d/branch-set-input-invalidation.changed.md +++ b/changelog.d/branch-set-input-invalidation.changed.md @@ -1 +1 @@ -Disk storage writes a new file for every store, so recalculating a value no longer changes what branches sharing the directory read; apply_reform keeps inputs by a flag on each stored array instead of replaying _user_input_keys; values a custom set_input handler calculates are no longer treated as inputs; dump_simulation records which values were inputs (inputs.txt), and restore_simulation restores only those as inputs; a calculation running when an input set on its simulation drops values is not cached (nor is anything calculated meanwhile in other simulations), and the outermost one in that simulation is run again from the new inputs until a run changes none (at most ten times); a branch stops reading macro-cache files once an input is set on it. +Disk storage writes a new file for every store, so recalculating a value no longer changes what branches sharing the directory read; apply_reform keeps inputs by a flag on each stored array instead of replaying _user_input_keys; values a custom set_input handler calculates are no longer treated as inputs; dump_simulation records which values were inputs (inputs.txt), and restore_simulation restores only those as inputs; a calculation running when an input set on its simulation drops values is not cached (nor is anything calculated from it in other simulations, or still running in its family), and the outermost one in that simulation is run again from the new inputs until a run changes none (at most ten times); a branch stops reading macro-cache files once an input is set on it. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index d3350ba5c..04f45fc60 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -59,9 +59,12 @@ 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. A calculation that was running 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; nor is any result calculated meanwhile in another simulation, -which may have called into this one (a parent formula calculating in a branch -whose formula calls back into the parent, say). The outermost calculation +the macro cache. Neither is any result calculated from it: a calculation in any +simulation that received it, directly or through other calculations (a parent +formula calculating in a branch whose formula calls back into the parent, say), +or one still running in this context in another simulation of the same family, +which may hold values it read from the simulation before the drop. Unrelated +simulations, and other threads, keep caching. The outermost calculation running in the simulation whose input changed (a `calculate`, or a direct `calculate_add`, whose terms run within it) then runs again from the new inputs, inner calculations included, until a run changes no input, and keeps @@ -139,6 +142,11 @@ depend on the input is calculated again) but not less, within these limits: calculated stay. - **Existing child branches keep their values.** An input set on a branch does not reach branches already created from it. +- **Values already read stay read.** A formula that read a value from a branch + and then calculates there a formula that changes that branch's input keeps + what it read before the change. Its result is returned but not kept (it is + in the same family), and it is not run again: running it again would create + and read its branches the same way. With disk storage (`MemoryConfig`), every store writes a new file, named with the sequence number and a token for the process (a forked child gets its own), diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 503ddd8dc..286876294 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1,5 +1,6 @@ import hashlib import tempfile +from contextlib import contextmanager from contextvars import ContextVar from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union @@ -135,13 +136,32 @@ def _uprating_index_value(parameter, instant) -> Optional[float]: # returned but not kept. _RERUNS_AFTER_INPUT_CHANGE = 10 -# Counts drops made while calculations were running (``_drop_computed``), in -# any simulation. A result calculated across one, wherever it happened, may -# come from the replaced value (a formula in one simulation can calculate in -# another, which calls back into the first), so it is not kept. Each -# simulation also counts its own (``_input_epoch``): those decide which -# calculation runs again. -_input_epoch_anywhere = [0] +# The calculations running in this context, innermost last, as (simulation, +# frame) pairs. A frame notes whether its result was not kept, or whether a +# calculation it called returned a result that was not kept: either way what +# it returns may come from a replaced value, so its caller does not keep its +# own result either (see ``_cache_result``). +_calculation_frames: ContextVar[tuple] = ContextVar("_calculation_frames", default=()) + + +def _family_root(simulation: "Simulation") -> "Simulation": + while getattr(simulation, "parent_branch", None) is not None: + simulation = simulation.parent_branch + return simulation + + +@contextmanager +def _calculation_frame(simulation: "Simulation"): + """Run a calculation in a new frame of ``_calculation_frames``.""" + frames = _calculation_frames.get() + frame = {"received": False, "refused": False} + token = _calculation_frames.set(frames + ((simulation, frame),)) + try: + yield frame + finally: + _calculation_frames.reset(token) + if frames and (frame["received"] or frame["refused"]): + frames[-1][1]["received"] = True class Simulation: @@ -628,6 +648,8 @@ def calculate( # an input change here: running it again runs the inner ones too. outermost = not getattr(self, "_calculations_in_flight", 0) self._calculations_in_flight = getattr(self, "_calculations_in_flight", 0) + 1 + frame_context = _calculation_frame(self) + frame = frame_context.__enter__() try: input_epoch = getattr(self, "_input_epoch", 0) result = self._calculate(variable_name, period) @@ -641,6 +663,7 @@ def calculate( # inputs first would: later uprating and carry-over look for # the periods held. input_epoch = getattr(self, "_input_epoch", 0) + frame["received"] = frame["refused"] = False result = self._calculate(variable_name, period) # Satisfies ``requires_computation_after`` from now on, even if a # branch input later drops the values. @@ -661,6 +684,7 @@ def calculate( self._share_store_history_with_caller() finally: self._calculations_in_flight -= 1 + frame_context.__exit__(None, None, None) self.tracer.record_calculation_end() self.purge_cache_of_invalid_values() @@ -865,7 +889,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 and input_state[:2] == self._input_state(): + if is_cache_available and self._may_keep(input_state): smc.set_cache_value(cache_path, values) return values @@ -1009,7 +1033,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: # Neither cache keeps a result an input change may have made obsolete # (see ``_cache_result``). - unchanged = input_state[:2] == self._input_state() + unchanged = self._may_keep(input_state) if is_cache_available and unchanged: smc.set_cache_value(cache_path, array) @@ -1020,7 +1044,14 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: def _input_state(self) -> Tuple[int, int]: """How many drops ran during calculations, and inputs were set, here.""" - return _input_epoch_anywhere[0], getattr(self, "_inputs_set", 0) + return getattr(self, "_cache_epoch", 0), getattr(self, "_inputs_set", 0) + + def _may_keep(self, input_state: Tuple[int, int, int]) -> bool: + """Whether a result calculated since ``input_state`` may be kept (see ``_cache_result``).""" + frames = _calculation_frames.get() + return input_state[:2] == self._input_state() and not ( + frames and frames[-1][1]["received"] + ) def _calculation_start(self) -> Tuple[int, int, int]: """:meth:`_input_state` when a calculation begins, and a sequence number then.""" @@ -1065,11 +1096,15 @@ def _cache_result( stored_input = self._input_set_meanwhile(holder, period, started_at) if stored_input is not None: return stored_input - if epoch == _input_epoch_anywhere[0]: + frames = _calculation_frames.get() + received = bool(frames) and frames[-1][1]["received"] + if epoch == getattr(self, "_cache_epoch", 0) and not received: holder.put_in_cache(array, period, self.branch_name) else: # Not kept, but whatever reads it is stored after it all the same. holder._record_store(period, next_sequence_number()) + if frames: + frames[-1][1]["refused"] = True return array def purge_cache_of_invalid_values(self) -> None: @@ -1133,21 +1168,23 @@ def total(): outermost = not getattr(self, "_calculations_in_flight", 0) self._calculations_in_flight = getattr(self, "_calculations_in_flight", 0) + 1 try: - input_epoch = getattr(self, "_input_epoch", 0) - input_state = self._calculation_start() - result = total() - for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): - if input_epoch == getattr(self, "_input_epoch", 0): - break - # An input changed while summing: earlier terms may be - # obsolete (see ``calculate``). Sum again. + with _calculation_frame(self) as frame: input_epoch = getattr(self, "_input_epoch", 0) input_state = self._calculation_start() result = total() + for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): + if input_epoch == getattr(self, "_input_epoch", 0): + break + # An input changed while summing: earlier terms may be + # obsolete (see ``calculate``). Sum again. + input_epoch = getattr(self, "_input_epoch", 0) + frame["received"] = frame["refused"] = False + input_state = self._calculation_start() + result = total() + holder = self.get_holder(variable.name) + return self._cache_result(holder, result, period, input_state) finally: self._calculations_in_flight -= 1 - holder = self.get_holder(variable.name) - return self._cache_result(holder, result, period, input_state) def calculate_divide( self, @@ -1176,11 +1213,12 @@ def calculate_divide( ) if period.unit == periods.MONTH: - input_state = self._calculation_start() - computation_period = period.this_year - result = self.calculate(variable_name, period=computation_period) / 12.0 - holder = self.get_holder(variable.name) - return self._cache_result(holder, result, period, input_state) + with _calculation_frame(self): + input_state = self._calculation_start() + computation_period = period.this_year + result = self.calculate(variable_name, period=computation_period) / 12.0 + holder = self.get_holder(variable.name) + return self._cache_result(holder, result, period, input_state) elif period.unit == periods.YEAR: return self.calculate(variable_name, period) @@ -1673,7 +1711,16 @@ def _drop_computed(self, since: Optional[int] = None) -> int: # simulations calling into this one (``_cache_result``); the # outermost one here runs again (``calculate``). self._input_epoch = getattr(self, "_input_epoch", 0) + 1 - _input_epoch_anywhere[0] += 1 + self._cache_epoch = getattr(self, "_cache_epoch", 0) + 1 + # Calculations of this family running in this context, in other + # simulations, may hold values read from this one before the + # drop (a parent formula that read its branch, then called into + # it): keep none of their results either. Other families get + # only what calculations return here (``_calculation_frames``). + root = _family_root(self) + for simulation, _ in _calculation_frames.get(): + if simulation is not self and _family_root(simulation) is root: + simulation._cache_epoch = getattr(simulation, "_cache_epoch", 0) + 1 return dropped # Nothing the simulation still holds was calculated from what the # records numbered ``since`` or later describe, except the inputs it diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index 908a69f23..1e2047d72 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -1240,6 +1240,147 @@ def formula(person, period): 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, so carry-over finds it.""" + 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 branch.calculate("result", "2021").tolist() == [6.0] # carried over + + +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_does_not_keep_what_it_read_before_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.calculate("result", "2020").tolist() + + live_first, live_next = run(input_first=False) + assert live_first == [2.0] # documented: what it read stays read + assert live_next == run(input_first=True)[1] == [6.0] # but it was not kept + + def test_input_a_formula_sets_for_its_own_period_wins(): """As it would had the input been set before the calculation.""" From 94f21309961b7dc5df5b16458241d5e5a7de1ae6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 09:52:20 -0400 Subject: [PATCH 21/42] Close the ninth review's gaps: refuse by what each calculation read From the round-9 soundness review (on 359c97fb): - A calculation that changed its simulation's input and then left before any cache decision (a carry-over default returned early, an exception) did not tell its caller, which kept an obsolete result (0 instead of 6). - A formula that read another simulation (another family, a root clone, its own branch from a plain thread) and then calculated there something that changed that simulation's input kept its old result (2 on the second request instead of 6). - The family bump refused a settled, correct result, so a later carry-over found no period (0 instead of 6; master gives 6). One rule replaces the taint flag and the family bump: each calculation frame records, for every other simulation it got a value from (directly or through its callees), that simulation's _input_epoch when the value's calculation began; it keeps its result only if neither its own simulation's epoch nor any of those has changed. A callee hands its reads to its caller however it ends (value, early default, error), its own simulation counted at its start; a rerun that settles restarts the count, so a settled call leaves its caller clean. Fast-cache hits from another simulation are reads too. Docs updated ("values already read stay read" now covers any simulation and caught errors). Tests for each witness; mutants M39 (no reads) and M41 (a read counted from return, not start) are caught. Co-Authored-By: Claude Opus 5.5 --- .../branch-set-input-invalidation.changed.md | 2 +- docs/usage/branches.md | 28 ++- policyengine_core/simulations/simulation.py | 84 +++++--- tests/core/test_branch_input_invalidation.py | 198 ++++++++++++++++++ 4 files changed, 270 insertions(+), 42 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.changed.md b/changelog.d/branch-set-input-invalidation.changed.md index 1fdef0f1a..d1dc7ef0c 100644 --- a/changelog.d/branch-set-input-invalidation.changed.md +++ b/changelog.d/branch-set-input-invalidation.changed.md @@ -1 +1 @@ -Disk storage writes a new file for every store, so recalculating a value no longer changes what branches sharing the directory read; apply_reform keeps inputs by a flag on each stored array instead of replaying _user_input_keys; values a custom set_input handler calculates are no longer treated as inputs; dump_simulation records which values were inputs (inputs.txt), and restore_simulation restores only those as inputs; a calculation running when an input set on its simulation drops values is not cached (nor is anything calculated from it in other simulations, or still running in its family), and the outermost one in that simulation is run again from the new inputs until a run changes none (at most ten times); a branch stops reading macro-cache files once an input is set on it. +Disk storage writes a new file for every store, so recalculating a value no longer changes what branches sharing the directory read; apply_reform keeps inputs by a flag on each stored array instead of replaying _user_input_keys; values a custom set_input handler calculates are no longer treated as inputs; dump_simulation records which values were inputs (inputs.txt), and restore_simulation restores only those as inputs; a calculation running when an input set on its simulation drops values is not cached (nor is anything in another simulation calculated from a value read before the change), and the outermost one in that simulation is run again from the new inputs until a run changes none (at most ten times); a branch stops reading macro-cache files once an input is set on it. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 04f45fc60..df3e45d45 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -59,12 +59,17 @@ 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. A calculation that was running 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 it: a calculation in any -simulation that received it, directly or through other calculations (a parent -formula calculating in a branch whose formula calls back into the parent, say), -or one still running in this context in another simulation of the same family, -which may hold values it read from the simulation before the drop. Unrelated -simulations, and other threads, keep caching. The outermost calculation +the macro cache. Neither is any result calculated from it, 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 +such input changes when the value's calculation began, and keeps its own +result only if none has changed since. 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. A +calculation whose call into another simulation settled before returning keeps +its result, and unrelated simulations and other threads keep caching. The +outermost calculation running in the simulation whose input changed (a `calculate`, or a direct `calculate_add`, whose terms run within it) then runs again from the new inputs, inner calculations included, until a run changes no input, and keeps @@ -142,11 +147,12 @@ depend on the input is calculated again) but not less, within these limits: calculated stay. - **Existing child branches keep their values.** An input set on a branch does not reach branches already created from it. -- **Values already read stay read.** A formula that read a value from a branch - and then calculates there a formula that changes that branch's input keeps - what it read before the change. Its result is returned but not kept (it is - in the same family), and it is not run again: running it again would create - and read its branches the same way. +- **Values already read stay read.** A formula that read a value from another + simulation and then calculates there a formula that changes that + simulation's input keeps what it read before the change. Its result is + returned but not kept, and it is not run again: running it again would + create and read its branches the same way. Likewise a formula that catches + an error raised after such a change returns its own fallback, unkept. With disk storage (`MemoryConfig`), every store writes a new file, named with the sequence number and a token for the process (a forked child gets its own), diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 286876294..b9143d086 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -137,31 +137,60 @@ def _uprating_index_value(parameter, instant) -> Optional[float]: _RERUNS_AFTER_INPUT_CHANGE = 10 # The calculations running in this context, innermost last, as (simulation, -# frame) pairs. A frame notes whether its result was not kept, or whether a -# calculation it called returned a result that was not kept: either way what -# it returns may come from a replaced value, so its caller does not keep its -# own result either (see ``_cache_result``). +# frame) pairs. A frame notes ``start``, its simulation's ``_input_epoch`` when +# it (last) began, and ``reads``: 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 while calculations ran there), the value +# may come from a replaced one, and so may what the frame calculates from it, +# which is then not kept (``_frame_is_stale``). _calculation_frames: ContextVar[tuple] = ContextVar("_calculation_frames", default=()) -def _family_root(simulation: "Simulation") -> "Simulation": - while getattr(simulation, "parent_branch", None) is not None: - simulation = simulation.parent_branch - return simulation +def _new_frame(simulation: "Simulation") -> dict: + return {"start": getattr(simulation, "_input_epoch", 0), "reads": {}} + + +def _frame_is_stale(simulation: "Simulation", frame: dict) -> bool: + """Whether what the frame calculates may come from a value since replaced.""" + return getattr(simulation, "_input_epoch", 0) != frame["start"] or any( + getattr(other, "_input_epoch", 0) != epoch + for other, epoch in frame["reads"].items() + ) + + +def _hand_to_caller(frames: tuple, simulation: "Simulation", frame: dict) -> None: + """Pass what a calculation in ``simulation`` read on to the calling frame. + + The calculation's own simulation counts at its ``start``: if an input + there changed while it ran, its caller's result is stale too. + """ + if not frames: + return + caller_simulation, caller = frames[-1] + reads = caller["reads"] + for other, epoch in frame["reads"].items(): + if other is not caller_simulation and other not in reads: + reads[other] = epoch + if simulation is not caller_simulation and simulation not in reads: + reads[simulation] = frame["start"] @contextmanager def _calculation_frame(simulation: "Simulation"): - """Run a calculation in a new frame of ``_calculation_frames``.""" + """Run a calculation in a new frame of ``_calculation_frames``. + + Whatever way it ends (a value, an early default, an error), its caller + learns what it read and whether that may since have been replaced. + """ frames = _calculation_frames.get() - frame = {"received": False, "refused": False} + frame = _new_frame(simulation) token = _calculation_frames.set(frames + ((simulation, frame),)) try: yield frame finally: _calculation_frames.reset(token) - if frames and (frame["received"] or frame["refused"]): - frames[-1][1]["received"] = True + _hand_to_caller(frames, simulation, frame) class Simulation: @@ -632,6 +661,11 @@ def calculate( _cached = _fast_cache.get(_fast_key) if _cached is not None: self._share_store_history_with_caller() + frames = _calculation_frames.get() + if frames and frames[-1][0] is not self: + frames[-1][1]["reads"].setdefault( + self, getattr(self, "_input_epoch", 0) + ) return _cached self.tracer.record_calculation_start(variable_name, period, self.branch_name) @@ -663,7 +697,7 @@ def calculate( # inputs first would: later uprating and carry-over look for # the periods held. input_epoch = getattr(self, "_input_epoch", 0) - frame["received"] = frame["refused"] = False + frame.update(_new_frame(self)) result = self._calculate(variable_name, period) # Satisfies ``requires_computation_after`` from now on, even if a # branch input later drops the values. @@ -1044,13 +1078,13 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: def _input_state(self) -> Tuple[int, int]: """How many drops ran during calculations, and inputs were set, here.""" - return getattr(self, "_cache_epoch", 0), getattr(self, "_inputs_set", 0) + return getattr(self, "_input_epoch", 0), getattr(self, "_inputs_set", 0) def _may_keep(self, input_state: Tuple[int, int, int]) -> bool: """Whether a result calculated since ``input_state`` may be kept (see ``_cache_result``).""" frames = _calculation_frames.get() return input_state[:2] == self._input_state() and not ( - frames and frames[-1][1]["received"] + frames and _frame_is_stale(*frames[-1]) ) def _calculation_start(self) -> Tuple[int, int, int]: @@ -1097,14 +1131,12 @@ def _cache_result( if stored_input is not None: return stored_input frames = _calculation_frames.get() - received = bool(frames) and frames[-1][1]["received"] - if epoch == getattr(self, "_cache_epoch", 0) and not received: + stale = bool(frames) and _frame_is_stale(*frames[-1]) + if epoch == getattr(self, "_input_epoch", 0) and not stale: holder.put_in_cache(array, period, self.branch_name) else: # Not kept, but whatever reads it is stored after it all the same. holder._record_store(period, next_sequence_number()) - if frames: - frames[-1][1]["refused"] = True return array def purge_cache_of_invalid_values(self) -> None: @@ -1178,7 +1210,7 @@ def total(): # An input changed while summing: earlier terms may be # obsolete (see ``calculate``). Sum again. input_epoch = getattr(self, "_input_epoch", 0) - frame["received"] = frame["refused"] = False + frame.update(_new_frame(self)) input_state = self._calculation_start() result = total() holder = self.get_holder(variable.name) @@ -1710,17 +1742,9 @@ def _drop_computed(self, since: Optional[int] = None) -> int: # the results running calculations return, here or in the # simulations calling into this one (``_cache_result``); the # outermost one here runs again (``calculate``). + # Calculations elsewhere that already read from this simulation + # see the change through their frames' ``reads``. self._input_epoch = getattr(self, "_input_epoch", 0) + 1 - self._cache_epoch = getattr(self, "_cache_epoch", 0) + 1 - # Calculations of this family running in this context, in other - # simulations, may hold values read from this one before the - # drop (a parent formula that read its branch, then called into - # it): keep none of their results either. Other families get - # only what calculations return here (``_calculation_frames``). - root = _family_root(self) - for simulation, _ in _calculation_frames.get(): - if simulation is not self and _family_root(simulation) is root: - simulation._cache_epoch = getattr(simulation, "_cache_epoch", 0) + 1 return dropped # Nothing the simulation still holds was calculated from what the # records numbered ``since`` or later describe, except the inputs it diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index 1e2047d72..5650e02ea 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -1381,6 +1381,204 @@ def result(person, period): assert live_next == run(input_first=True)[1] == [6.0] # but it was not kept +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_does_not_keep_its_result(): + 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") + ) + + assert consumer.calculate("result", "2020").tolist() == [0.0] # its own fallback + assert consumer.calculate("result", "2020").tolist() == [6.0] # not kept + + +@pytest.mark.parametrize("worker_kind", ["other_family", "own_branch_in_a_thread"]) +def test_earlier_read_is_not_kept_after_the_worker_changes_its_input(worker_kind): + """What a formula read before stays read, but its result is not kept.""" + 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() == [2.0] # read before + assert simulation.calculate("result", "2020").tolist() == [6.0] # not kept + + +def test_settled_call_into_own_branch_keeps_the_result_for_carry_over(): + """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 root.calculate("result", "2021").tolist() == [6.0] # carried over + + def test_input_a_formula_sets_for_its_own_period_wins(): """As it would had the input been set before the calculation.""" From 23971e3b0deb0929e986dee5935801c1bc9f764a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 03:19:41 -0400 Subject: [PATCH 22/42] Close the tenth review's gaps: thread reads, settled reads, idle changes 1. A value read in a thread started without the formula's context now counts as a read: with no calling frame, a calculation's reads go to every frame open in its simulation's ancestors (each simulation lists its open frames). Before, a parent that got its branch's value that way kept it after the branch's input changed. 2. The outermost calculation in a simulation now runs again while its frame is stale, not only while its own input epoch changed, so a formula that read a value before another simulation's input changed gets the current one and keeps it. Before, a correct settled result went unkept and an earlier read stayed in the returned value. 3. Every drop counts as an input change of the simulation, also when nothing is running there, so a formula that read a simulation and then set an input on it directly does not keep what it read. 4. Setting a branch input to the value the branch already reads there, as an input, drops nothing (not for a value calculated there, nor through a set_input helper). Without this, a formula that sets the same input on every run would rerun until the budget ran out. The property test now also sets inputs again with the value read. Docs describe the new rules, and the limit that a value kept in one simulation stays when another's input changes after its calculation ended. Co-Authored-By: Claude Opus 5.5 --- .../branch-set-input-invalidation.changed.md | 2 +- docs/usage/branches.md | 77 +++-- .../data_storage/in_memory_storage.py | 3 +- policyengine_core/holders/holder.py | 34 ++- policyengine_core/simulations/simulation.py | 102 ++++--- tests/core/test_branch_input_invalidation.py | 266 +++++++++++++++++- ...test_branch_input_invalidation_property.py | 39 ++- 7 files changed, 437 insertions(+), 86 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.changed.md b/changelog.d/branch-set-input-invalidation.changed.md index 4a554437b..70e5a7073 100644 --- a/changelog.d/branch-set-input-invalidation.changed.md +++ b/changelog.d/branch-set-input-invalidation.changed.md @@ -1 +1 @@ -apply_reform drops every stored value set_input did not store, in place, instead of wiping the storages and replaying _user_input_keys; restore_simulation restores calculated values after the inputs, under one later number; a calculation running when an input set on its simulation drops values is not cached (nor is anything in another simulation calculated from a value read before the change), and the outermost one in that simulation is run again from the new inputs until a run changes none (at most ten times); a branch stops reading macro-cache files once an input is set on it. +apply_reform drops every stored value set_input did not store, in place, instead of wiping the storages and replaying _user_input_keys; restore_simulation restores calculated values after the inputs, under one later number; a calculation that may have read a value a branch input dropped (one running in that branch, or in any simulation that got a value from it before the change, also from a thread started without the formula's context) is not cached, and the outermost such calculation in each simulation runs again until it reads nothing changed since (at most ten times); setting a branch input to the value the branch already reads there drops nothing; a branch stops reading macro-cache files once an input is set on it. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 7816a5aea..45b215ef6 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -56,27 +56,35 @@ earliest recorded store of the variable for a period that shares a day with 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. A -calculation that was running 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 it, 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 -such input changes when the value's calculation began, and keeps its own -result only if none has changed since. 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. A -calculation whose call into another simulation settled before returning keeps -its result, and unrelated simulations and other threads keep caching. The -outermost calculation -running in the simulation whose input changed (a `calculate`, or a direct -`calculate_add`, whose terms run within it) then runs again from the new -inputs, inner calculations included, until a run changes no input, and keeps -that result. After ten reruns it stops: a formula that keeps changing inputs -returns its last result without keeping it. The budget is per simulation whose -input changes, so calculations nested across several such simulations can -rerun more. An input stored for the very period being +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. 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 input +set again with the value the branch already reads for that very period (not +through a `set_input` helper that spreads it over other periods) 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. @@ -116,7 +124,10 @@ 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. +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: @@ -136,8 +147,9 @@ depend on the input is calculated again) but not less, within these limits: - **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. Its own branches are - covered: their history goes to every ancestor with a calculation running. + 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. - **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. @@ -146,12 +158,17 @@ depend on the input is calculated again) but not less, within these limits: calculated stay. - **Existing child branches keep their values.** An input set on a branch does not reach branches already created from it. -- **Values already read stay read.** A formula that read a value from another - simulation and then calculates there a formula that changes that - simulation's input keeps what it read before the change. Its result is - returned but not kept, and it is not run again: running it again would - create and read its branches the same way. Likewise a formula that catches - an error raised after such a change returns its own fallback, unkept. +- **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. +- **Formulas that create a branch each run.** A formula that creates a branch, + reads it, and then changes its input runs again until the budget is spent + (each run reads a new branch from before the change), then returns its last + result without keeping it. With disk storage (`MemoryConfig`), a value dropped from a branch only leaves its storage's index: the file stays until the storage directory is removed, diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 3f82a2677..20790d477 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -326,8 +326,7 @@ def drop_computed(self, *, since: Optional[int] = None) -> int: dropped = [ key for key in self._arrays - if key not in inputs - and (since is None or numbers.get(key, since) >= since) + if key not in inputs and (since is None or numbers.get(key, since) >= since) ] for key in dropped: del self._arrays[key] diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 6bd0d467b..d708dbe78 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -134,6 +134,32 @@ def _is_input(self, period: Period, branch_name: str = "default") -> bool: return not storage.is_derived(period, branch_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). + """ + 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: + return bool( + numpy.array_equal( + self._get_array_from_storage(period, stored_on), + numpy.asarray(array), + ) + ) + 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]: @@ -330,8 +356,12 @@ def set_input( 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``). - simulation._drop_values_that_may_depend_on(self.variable.name, period) + # this input replaces (see ``Simulation.set_input``), unless it + # replaces the same input. + if getattr( + simulation, "parent_branch", None + ) is None or 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 = getattr(simulation, "_inputs_set", 0) + 1 diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index e0b2a9350..05ba8d632 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -195,9 +195,11 @@ def _latest_input_key(period: Period, definition_period: str) -> tuple: # it (last) began, and ``reads``: 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 while calculations ran there), the value -# may come from a replaced one, and so may what the frame calculates from it, -# which is then not kept (``_frame_is_stale``). +# 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 (``_frame_is_stale``). Each simulation also lists the frames open +# in it, in any context (``_open_frames``), so that a value read in a thread +# with no frame reaches the frames waiting for it (``_hand_to_caller``). _calculation_frames: ContextVar[tuple] = ContextVar("_calculation_frames", default=()) @@ -207,9 +209,10 @@ def _new_frame(simulation: "Simulation") -> dict: def _frame_is_stale(simulation: "Simulation", frame: dict) -> bool: """Whether what the frame calculates may come from a value since replaced.""" + # A snapshot: a thread handing a read to this frame may add one meanwhile. return getattr(simulation, "_input_epoch", 0) != frame["start"] or any( getattr(other, "_input_epoch", 0) != epoch - for other, epoch in frame["reads"].items() + for other, epoch in tuple(frame["reads"].items()) ) @@ -217,17 +220,45 @@ def _hand_to_caller(frames: tuple, simulation: "Simulation", frame: dict) -> Non """Pass what a calculation in ``simulation`` read on to the calling frame. The calculation's own simulation counts at its ``start``: if an input - there changed while it ran, its caller's result is stale too. + there changed while it ran, its caller's result is stale too. With no + calling frame in this context, the reads go to the frames open in the + simulation's ancestors (``_hand_to_waiting_ancestors``). """ + reads = tuple(frame["reads"].items()) if not frames: + _hand_to_waiting_ancestors(simulation, frame["start"], reads) return caller_simulation, caller = frames[-1] - reads = caller["reads"] - for other, epoch in frame["reads"].items(): - if other is not caller_simulation and other not in reads: - reads[other] = epoch - if simulation is not caller_simulation and simulation not in reads: - reads[simulation] = frame["start"] + 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. + """ + ancestor = getattr(simulation, "parent_branch", None) + while ancestor is not None: + for open_frame in tuple(getattr(ancestor, "_open_frames", {}).values()): + 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) @contextmanager @@ -240,9 +271,12 @@ def _calculation_frame(simulation: "Simulation"): frames = _calculation_frames.get() frame = _new_frame(simulation) token = _calculation_frames.set(frames + ((simulation, frame),)) + open_frames = simulation.__dict__.setdefault("_open_frames", {}) + open_frames[id(frame)] = frame try: yield frame finally: + open_frames.pop(id(frame), None) _calculation_frames.reset(token) _hand_to_caller(frames, simulation, frame) @@ -479,9 +513,7 @@ def _invalidate_all_caches(self) -> None: # no holder and therefore nothing to wipe. for population in self.populations.values(): for holder in population._holders.values(): - holder._mark_derived_except( - user_inputs.get(holder.variable.name, ()) - ) + holder._mark_derived_except(user_inputs.get(holder.variable.name, ())) self._drop_computed() for branch in self.branches.values(): branch._invalidate_all_caches() @@ -735,10 +767,11 @@ def calculate( if _cached is not None: self._share_store_history_with_caller() frames = _calculation_frames.get() - if frames and frames[-1][0] is not self: - frames[-1][1]["reads"].setdefault( - self, getattr(self, "_input_epoch", 0) - ) + epoch = getattr(self, "_input_epoch", 0) + if not frames: + _hand_to_waiting_ancestors(self, epoch, ()) + elif frames[-1][0] is not self: + frames[-1][1]["reads"].setdefault(self, epoch) return _cached self.tracer.record_calculation_start(variable_name, period, self.branch_name) @@ -758,18 +791,15 @@ def calculate( frame_context = _calculation_frame(self) frame = frame_context.__enter__() try: - input_epoch = getattr(self, "_input_epoch", 0) result = self._calculate(variable_name, period) for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): - if input_epoch == getattr(self, "_input_epoch", 0): + if not _frame_is_stale(self, frame): break - # An input set here while it ran (by a formula, say) 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: later uprating and carry-over look for - # the periods held. - input_epoch = getattr(self, "_input_epoch", 0) + # 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.update(_new_frame(self)) result = self._calculate(variable_name, period) # Satisfies ``requires_computation_after`` from now on, even if a @@ -1327,15 +1357,13 @@ def total(): self._calculations_in_flight = getattr(self, "_calculations_in_flight", 0) + 1 try: with _calculation_frame(self) as frame: - input_epoch = getattr(self, "_input_epoch", 0) input_state = self._calculation_start() result = total() for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): - if input_epoch == getattr(self, "_input_epoch", 0): + if not _frame_is_stale(self, frame): break # An input changed while summing: earlier terms may be # obsolete (see ``calculate``). Sum again. - input_epoch = getattr(self, "_input_epoch", 0) frame.update(_new_frame(self)) input_state = self._calculation_start() result = total() @@ -1884,15 +1912,14 @@ def _drop_computed(self, since: Optional[int] = None) -> int: 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 = getattr(self, "_input_epoch", 0) + 1 if getattr(self, "_calculations_in_flight", 0): # A formula running here may still hold values calculated from - # what the records describe: keep the records, and keep none of - # the results running calculations return, here or in the - # simulations calling into this one (``_cache_result``); the - # outermost one here runs again (``calculate``). - # Calculations elsewhere that already read from this simulation - # see the change through their frames' ``reads``. - self._input_epoch = getattr(self, "_input_epoch", 0) + 1 + # 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 @@ -1997,6 +2024,7 @@ def clone( 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 = {} return new diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index e55db005c..830a422b5 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -1320,7 +1320,7 @@ def worker_outer(person, period): assert run(input_first=False) == run(input_first=True) == ([6.0], [6.0]) -def test_caller_in_the_same_family_does_not_keep_what_it_read_before_the_change(): +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): @@ -1351,11 +1351,10 @@ def result(person, period): if input_first: kept.set_input("source", "2020", np.array([3.0])) first = root.calculate("result", "2020").tolist() - return first, root.calculate("result", "2020").tolist() + return first, root.get_array("result", "2020").tolist() - live_first, live_next = run(input_first=False) - assert live_first == [2.0] # documented: what it read stays read - assert live_next == run(input_first=True)[1] == [6.0] # but it was not kept + # 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=()): @@ -1433,7 +1432,7 @@ def leaf(person, period): assert run(input_first=False) == run(input_first=True) == ([6.0], [6.0]) -def test_caller_that_catches_an_error_after_an_input_change_does_not_keep_its_result(): +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 @@ -1463,13 +1462,14 @@ def result(person, period): .get_branch("consumer") ) - assert consumer.calculate("result", "2020").tolist() == [0.0] # its own fallback - assert consumer.calculate("result", "2020").tolist() == [6.0] # not kept + # 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_not_kept_after_the_worker_changes_its_input(worker_kind): - """What a formula read before stays read, but its result is not kept.""" +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): @@ -1518,8 +1518,8 @@ def result(person, period): simulation.set_input("source", "2020", np.array([1.0])) simulation.get_branch("kept") - assert simulation.calculate("result", "2020").tolist() == [2.0] # read before - assert simulation.calculate("result", "2020").tolist() == [6.0] # not 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(): @@ -2207,3 +2207,245 @@ def test_storage_pickled_before_stores_were_numbered_still_stores(): 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 diff --git a/tests/core/test_branch_input_invalidation_property.py b/tests/core/test_branch_input_invalidation_property.py index c76d7c448..17372d277 100644 --- a/tests/core/test_branch_input_invalidation_property.py +++ b/tests/core/test_branch_input_invalidation_property.py @@ -1,8 +1,9 @@ """Branch results do not depend on what was calculated before. Random sequences of calculations, branches (including reused and forgotten -names), inputs, ``drop_computed_arrays`` and dumps restored as new -simulations run on a synthetic system, with +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 @@ -120,6 +121,8 @@ 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), ) @@ -160,6 +163,18 @@ ] ) ) +# 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( @@ -177,6 +192,7 @@ st.one_of( single_operation.map(lambda operation: [operation]), override_after_calculating, + reset_after_calculating, override_after_restoring, ), min_size=1, @@ -260,6 +276,25 @@ def _run(program, mode): 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) From 789d7f05c3da493f586930c74c1c831a0dec1573 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 03:19:41 -0400 Subject: [PATCH 23/42] Remove synthetic_simulation's temporary folders when done (US + core hub patch) Each disk-backed synthetic simulation gets a folder inside one folder per test process, removed when the simulation is collected or at exit, and its pre-created disk storages leave removal to that, so a killed run leaves one folder, not one per live simulation. Patch from the US + core hub (policyengine-core#585 notes). Co-Authored-By: Claude Opus 5.5 --- tests/fixtures/branch_input_invalidation.py | 40 +++++++++++++++++---- 1 file changed, 34 insertions(+), 6 deletions(-) diff --git a/tests/fixtures/branch_input_invalidation.py b/tests/fixtures/branch_input_invalidation.py index 06f737cc3..506adac75 100644 --- a/tests/fixtures/branch_input_invalidation.py +++ b/tests/fixtures/branch_input_invalidation.py @@ -1,6 +1,9 @@ """A synthetic system for derived values and branches inside formulas.""" +import atexit +import shutil import tempfile +import weakref import numpy as np @@ -140,6 +143,22 @@ class System(CountryTaxBenefitSystem): 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 ): @@ -147,11 +166,15 @@ def synthetic_simulation( simulation = SimulationBuilder().build_default_simulation(system, count=PEOPLE) simulation.opt_out_cache = opt_out_cache if memory_config is not None: - # Each holder's disk storage removes its directory, and this one - # once empty, when garbage-collected. - simulation._data_storage_dir = tempfile.mkdtemp( - prefix="policyengine-branch-tests-" - ) + # 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 @@ -161,7 +184,12 @@ def synthetic_simulation( for population in simulation.populations.values(): population._holders = {} for variable in system.variables: - simulation.get_holder(variable) + 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])), From e586d0bb154d9803c8a272c7928f675982b8ec99 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 03:28:29 -0400 Subject: [PATCH 24/42] Bring two docstrings in line with the merged rules StoreHistory: uprating and carry-over read which periods hold inputs. Simulation.set_input: drops count as input changes, stale calculations anywhere run again, unchanged inputs drop nothing, and values kept elsewhere stay. Co-Authored-By: Claude Opus 5.5 --- .../data_storage/store_history.py | 4 ++-- policyengine_core/simulations/simulation.py | 21 ++++++++++++------- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/policyengine_core/data_storage/store_history.py b/policyengine_core/data_storage/store_history.py index 2640cef21..787b9762e 100644 --- a/policyengine_core/data_storage/store_history.py +++ b/policyengine_core/data_storage/store_history.py @@ -57,8 +57,8 @@ class StoreHistory: - 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 a value). Such a value depends on which other periods hold - values, so an input set for any period can change it. + 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`): diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 05ba8d632..1164cb556 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1800,12 +1800,15 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non 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. - A calculation that was running when its input changed, in the branch - or in a simulation calling into it, is not kept; the outermost one in - the branch is run again from the new inputs until a run changes none - (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. + 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); @@ -1817,8 +1820,10 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non starts without copying its context; and branches a formula keeps between calls, which hold what their parent held when they were created. Inputs set on a simulation that is not a branch drop - nothing, as before, and branches already created from the branch - keep their values. + nothing, as before; 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: From 7084f56b14a5772e41f8db78e0e35b369e1186e6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 04:20:29 -0400 Subject: [PATCH 25/42] Keep macro reads off on unchanged inputs; cut read tracking's per-call cost Review of 789d7f05 (round 11, read-only): - The unchanged-input short-circuit also skipped the branch's macro-cache switch-off, so a branch given the input it already read kept reading files another simulation's same-named branch may have written. Holder. set_input now turns macro reads off on every branch input. - -0.0 no longer counts as the input 0.0 (1 / x tells them apart). - The changelog states the short-circuit's conditions. Cost: a single policyengine-us household (household_net_income, 2026) took about 20% longer than on master, all in per-calculation bookkeeping. Now: - a calculation nested in another in the same simulation shares that one's frame (it began no later and notes at least the same reads, each at its first epoch, so sharing only keeps less); a frame opens only at the outermost calculation in a simulation or on entering another one; - frames are a slotted class, their own context manager and registry key; - counters are class attributes, read without getattr defaults. The household now takes about 4-5% longer than on master. Co-Authored-By: Claude Opus 5.5 --- .../branch-set-input-invalidation.changed.md | 2 +- docs/usage/branches.md | 62 ++--- policyengine_core/holders/holder.py | 37 ++- policyengine_core/simulations/simulation.py | 235 +++++++++++------- tests/core/test_branch_input_invalidation.py | 62 +++++ 5 files changed, 264 insertions(+), 134 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.changed.md b/changelog.d/branch-set-input-invalidation.changed.md index 70e5a7073..ae1d62d69 100644 --- a/changelog.d/branch-set-input-invalidation.changed.md +++ b/changelog.d/branch-set-input-invalidation.changed.md @@ -1 +1 @@ -apply_reform drops every stored value set_input did not store, in place, instead of wiping the storages and replaying _user_input_keys; restore_simulation restores calculated values after the inputs, under one later number; a calculation that may have read a value a branch input dropped (one running in that branch, or in any simulation that got a value from it before the change, also from a thread started without the formula's context) is not cached, and the outermost such calculation in each simulation runs again until it reads nothing changed since (at most ten times); setting a branch input to the value the branch already reads there drops nothing; a branch stops reading macro-cache files once an input is set on it. +apply_reform drops every stored value set_input did not store, in place, instead of wiping the storages and replaying _user_input_keys; restore_simulation restores calculated values after the inputs, under one later number; a calculation that may have read a value a branch input dropped (one running in that branch, or in any simulation that got a value from it before the change, also from a thread started without the formula's context) is not cached, and the outermost such calculation in each simulation runs again until it reads nothing changed since (at most ten times); setting a branch input to the input the branch already reads for that period (not through a set_input helper) drops nothing; a branch stops reading macro-cache files once an input is set on it. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 45b215ef6..c1d996f6a 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -59,35 +59,36 @@ 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 +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. 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 input -set again with the value the branch already reads for that very period (not -through a `set_input` helper that spreads it over other periods) 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. +only if none has changed since. (A calculation nested in another in the same +simulation shares that one's record, which can only make it keep less.) 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 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. 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 @@ -188,10 +189,11 @@ 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, 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, everything calculated from the first read on; and +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 requested before a drop removed its values. diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index d708dbe78..a8e1ef5be 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -143,7 +143,8 @@ def _input_unchanged( 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). + 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 @@ -151,11 +152,11 @@ def _input_unchanged( 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( - self._get_array_from_storage(period, stored_on), - numpy.asarray(array), - ) + 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 @@ -357,14 +358,21 @@ def set_input( 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. - if getattr( - simulation, "parent_branch", None - ) is None or not self._input_unchanged(period, array, branch_name): + # 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 may come from another + # simulation's branch. + 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 = getattr(simulation, "_inputs_set", 0) + 1 + simulation._inputs_set += 1 if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() if not hasattr(simulation, "_user_input_contexts"): @@ -381,7 +389,7 @@ def set_input( finally: if ( simulation is not None - and getattr(simulation, "_calculations_started", 0) != started + and simulation._calculations_started != started ): # The handler calculated, maybe between its stores, # from inputs it had not yet replaced (also if it @@ -542,10 +550,13 @@ def _record_store(self, period: Period, sequence_number: int) -> None: simulation = getattr(self, "simulation", None) if simulation is None: return - if self.variable.definition_period == periods.ETERNITY: + 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( - self.variable.name, periods.period(period), sequence_number + variable.name, period, sequence_number ) def default_array(self) -> ArrayLike: diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 1164cb556..d43a71345 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1,6 +1,6 @@ import hashlib import tempfile -from contextlib import contextmanager +from contextlib import nullcontext from contextvars import ContextVar from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union @@ -190,51 +190,112 @@ def _latest_input_key(period: Period, definition_period: str) -> tuple: # returned but not kept. _RERUNS_AFTER_INPUT_CHANGE = 10 -# The calculations running in this context, innermost last, as (simulation, -# frame) pairs. A frame notes ``start``, its simulation's ``_input_epoch`` when -# it (last) began, and ``reads``: 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 (``_frame_is_stale``). Each simulation also lists the frames open -# in it, in any context (``_open_frames``), so that a value read in a thread -# with no frame reaches the frames waiting for it (``_hand_to_caller``). +# The calculations running in this context, innermost last (``_Frame``). _calculation_frames: ContextVar[tuple] = ContextVar("_calculation_frames", default=()) -def _new_frame(simulation: "Simulation") -> dict: - return {"start": getattr(simulation, "_input_epoch", 0), "reads": {}} +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", "outer", "token", "registry") -def _frame_is_stale(simulation: "Simulation", frame: dict) -> bool: - """Whether what the frame calculates may come from a value since replaced.""" - # A snapshot: a thread handing a read to this frame may add one meanwhile. - return getattr(simulation, "_input_epoch", 0) != frame["start"] or any( - getattr(other, "_input_epoch", 0) != epoch - for other, epoch in tuple(frame["reads"].items()) - ) + def __init__(self, simulation: "Simulation"): + self.simulation = simulation + + def restart(self) -> None: + """Begin again, as when the calculation runs again.""" + self.start = self.simulation._input_epoch + 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 __enter__(self) -> "_Frame": + simulation = self.simulation + self.start = simulation._input_epoch + self.reads = {} + self.outer = outer = _calculation_frames.get() + 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 + return self + + def __exit__(self, *exc_info) -> bool: + self.registry.discard(self) + _calculation_frames.reset(self.token) + _hand_to_caller(self.outer, self) + return False -def _hand_to_caller(frames: tuple, simulation: "Simulation", frame: dict) -> None: - """Pass what a calculation in ``simulation`` read on to the calling frame. +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: 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 in a simulation has its + own, to run again from (``Simulation.calculate``). + """ + if outermost: + return _Frame(simulation) + frames = _calculation_frames.get() + if not frames or frames[-1].simulation is not simulation: + return _Frame(simulation) + 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, the reads go to the frames open in the - simulation's ancestors (``_hand_to_waiting_ancestors``). + calling frame in this context (``outer`` empty), the reads go to the + frames open in the simulation's ancestors (``_hand_to_waiting_ancestors``). """ - reads = tuple(frame["reads"].items()) - if not frames: - _hand_to_waiting_ancestors(simulation, frame["start"], reads) + 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_simulation, caller = frames[-1] - caller_reads = caller["reads"] + caller = outer[-1] + 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"] + caller_reads[simulation] = frame.start def _hand_to_waiting_ancestors( @@ -252,8 +313,8 @@ def _hand_to_waiting_ancestors( """ ancestor = getattr(simulation, "parent_branch", None) while ancestor is not None: - for open_frame in tuple(getattr(ancestor, "_open_frames", {}).values()): - open_reads = open_frame["reads"] + for open_frame in tuple(getattr(ancestor, "_open_frames", None) or ()): + open_reads = open_frame.reads for other, epoch in reads: if other is not ancestor: open_reads.setdefault(other, epoch) @@ -261,26 +322,6 @@ def _hand_to_waiting_ancestors( ancestor = getattr(ancestor, "parent_branch", None) -@contextmanager -def _calculation_frame(simulation: "Simulation"): - """Run a calculation in a new frame of ``_calculation_frames``. - - Whatever way it ends (a value, an early default, an error), its caller - learns what it read and whether that may since have been replaced. - """ - frames = _calculation_frames.get() - frame = _new_frame(simulation) - token = _calculation_frames.set(frames + ((simulation, frame),)) - open_frames = simulation.__dict__.setdefault("_open_frames", {}) - open_frames[id(frame)] = frame - try: - yield frame - finally: - open_frames.pop(id(frame), None) - _calculation_frames.reset(token) - _hand_to_caller(frames, simulation, frame) - - class Simulation: """ Represents a simulation, and handles the calculation logic @@ -289,6 +330,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.""" @@ -767,11 +819,10 @@ def calculate( if _cached is not None: self._share_store_history_with_caller() frames = _calculation_frames.get() - epoch = getattr(self, "_input_epoch", 0) if not frames: - _hand_to_waiting_ancestors(self, epoch, ()) - elif frames[-1][0] is not self: - frames[-1][1]["reads"].setdefault(self, epoch) + _hand_to_waiting_ancestors(self, self._input_epoch, ()) + elif frames[-1].simulation is not self: + frames[-1].reads.setdefault(self, self._input_epoch) return _cached self.tracer.record_calculation_start(variable_name, period, self.branch_name) @@ -786,25 +837,29 @@ def calculate( # 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 getattr(self, "_calculations_in_flight", 0) - self._calculations_in_flight = getattr(self, "_calculations_in_flight", 0) + 1 - frame_context = _calculation_frame(self) - frame = frame_context.__enter__() + 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) for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): - if not _frame_is_stale(self, frame): + if not frame.is_stale(): break # 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.update(_new_frame(self)) + frame.restart() result = self._calculate(variable_name, period) # Satisfies ``requires_computation_after`` from now on, even if a # branch input later drops the values. - self._get_requested_variables().add(variable_name) + 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) @@ -821,7 +876,8 @@ def calculate( self._share_store_history_with_caller() finally: self._calculations_in_flight -= 1 - frame_context.__exit__(None, None, None) + if frame is not None: + frame.__exit__(None, None, None) self.tracer.record_calculation_end() self.purge_cache_of_invalid_values() @@ -937,7 +993,7 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: 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 = getattr(self, "_calculations_started", 0) + 1 + 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( @@ -1231,20 +1287,17 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike: return array - def _input_state(self) -> Tuple[int, int]: - """How many drops ran during calculations, and inputs were set, here.""" - return getattr(self, "_input_epoch", 0), getattr(self, "_inputs_set", 0) - 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 input_state[:2] == self._input_state() and not ( - frames and _frame_is_stale(*frames[-1]) - ) + return not (frames and frames[-1].is_stale()) def _calculation_start(self) -> Tuple[int, int, int]: - """:meth:`_input_state` when a calculation begins, and a sequence number then.""" - return (*self._input_state(), next_sequence_number()) + """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 @@ -1281,13 +1334,13 @@ def _cache_result( result is returned but not kept (see ``_drop_computed``). """ epoch, inputs_set, started_at = input_state - if inputs_set != getattr(self, "_inputs_set", 0): + 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() - stale = bool(frames) and _frame_is_stale(*frames[-1]) - if epoch == getattr(self, "_input_epoch", 0) and not stale: + stale = bool(frames) and frames[-1].is_stale() + if epoch == self._input_epoch and not stale: 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. @@ -1353,18 +1406,19 @@ def total(): # 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 getattr(self, "_calculations_in_flight", 0) - self._calculations_in_flight = getattr(self, "_calculations_in_flight", 0) + 1 + outermost = not self._calculations_in_flight + self._calculations_in_flight += 1 + frame = _frame_for(self, outermost) try: - with _calculation_frame(self) as frame: + with frame if frame is not None else nullcontext(): input_state = self._calculation_start() result = total() for _ in range(_RERUNS_AFTER_INPUT_CHANGE if outermost else 0): - if not _frame_is_stale(self, frame): + if not frame.is_stale(): break # An input changed while summing: earlier terms may be # obsolete (see ``calculate``). Sum again. - frame.update(_new_frame(self)) + frame.restart() input_state = self._calculation_start() result = total() # Cache only a sum over several sub-periods, as derived from @@ -1404,7 +1458,8 @@ def calculate_divide( ) if period.unit == periods.MONTH: - with _calculation_frame(self): + 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 @@ -1854,13 +1909,13 @@ def set_input(self, variable_name: str, period: Period, value: ArrayLike) -> Non del _fast_cache[key] def _get_store_history(self) -> StoreHistory: - history = getattr(self, "_store_history", None) + history = self._store_history if history is None: history = self._store_history = StoreHistory() return history def _get_requested_variables(self) -> set: - requested = getattr(self, "_requested_variables", None) + requested = self._requested_variables if requested is None: requested = self._requested_variables = set() return requested @@ -1879,7 +1934,7 @@ def _share_store_history_with_caller(self) -> None: # with a calculation running (more records only make drops broader). ancestor = getattr(self, "parent_branch", None) while ancestor is not None: - if getattr(ancestor, "_calculations_in_flight", 0): + if ancestor._calculations_in_flight: ancestor._get_store_history().merge(self._get_store_history()) ancestor = getattr(ancestor, "parent_branch", None) @@ -1919,10 +1974,10 @@ def _drop_computed(self, since: Optional[int] = None) -> int: 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 + # from it: none of them keeps its result (``_Frame.is_stale``), and # the outermost one in each simulation runs again (``calculate``). - self._input_epoch = getattr(self, "_input_epoch", 0) + 1 - if getattr(self, "_calculations_in_flight", 0): + self._input_epoch = self._input_epoch + 1 + if self._calculations_in_flight: # A formula running here may still hold values calculated from # what the records describe: keep the records. return dropped @@ -2029,7 +2084,7 @@ def clone( 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 = {} + new._open_frames = set() return new diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index 830a422b5..c06e9628e 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -2449,3 +2449,65 @@ def test_input_a_helper_spreads_over_months_still_drops(): 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] From 23c98a7f8010a1a69ca3f65f18b2519f49aea0b0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 04:22:23 -0400 Subject: [PATCH 26/42] Test a helper's input over a stored year, and a callback that changes its input Two mutants of the round-10 mechanisms survived every test: an equal input given to a set_input helper skipping the drop, and a call into another simulation sharing its caller's frame. Each test fails under its mutant. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_branch_input_invalidation.py | 86 ++++++++++++++++++++ 1 file changed, 86 insertions(+) diff --git a/tests/core/test_branch_input_invalidation.py b/tests/core/test_branch_input_invalidation.py index c06e9628e..ec86d82ba 100644 --- a/tests/core/test_branch_input_invalidation.py +++ b/tests/core/test_branch_input_invalidation.py @@ -2511,3 +2511,89 @@ def test_negative_zero_is_not_the_input_zero(): 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] From 78fb8d4b68a987ecd174fce45f1ca63facf3e567 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 04:49:19 -0400 Subject: [PATCH 27/42] Stop macro reads in one place: Holder.set_input on every branch input _drop_values_that_may_depend_on switched them off too, but it only runs from Holder.set_input after that has; the mutation check could not tell the second switch-off from none. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/holders/holder.py | 5 +++-- policyengine_core/simulations/simulation.py | 6 +----- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index a8e1ef5be..8e6deda99 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -360,8 +360,9 @@ def set_input( # 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 may come from another - # simulation's branch. + # 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: diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index d43a71345..aa6f8792d 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1948,11 +1948,7 @@ def _drop_values_that_may_depend_on( """ if getattr(self, "parent_branch", None) is None: return 0 - # Macro-cache files are keyed by branch and period, not by inputs, so - # a file for this branch's 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 has read any. - self.macro_cache_read = False + # (``Holder.set_input`` has already stopped macro-cache reads here.) since = self._get_store_history().earliest_dependency( variable_name, periods.period(period) ) From 0e73dc83b30da96b2c42be71eafb53455b4c4a61 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 11:05:16 -0400 Subject: [PATCH 28/42] Keep copied disk storage invalidation order independent --- .../branch-set-input-invalidation.fixed.md | 2 +- .../data_storage/on_disk_storage.py | 6 +++-- tests/core/test_on_disk_storage_clones.py | 22 +++++++++++++++++++ 3 files changed, 27 insertions(+), 3 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.fixed.md b/changelog.d/branch-set-input-invalidation.fixed.md index 146bf2378..fcd147b9c 100644 --- a/changelog.d/branch-set-input-invalidation.fixed.md +++ b/changelog.d/branch-set-input-invalidation.fixed.md @@ -1 +1 @@ -Simulation.set_input on a branch now drops the values the branch holds that may have been calculated from the value it replaces (every stored array carries a sequence number, and each simulation records what its values may have been calculated from), so a value the parent calculated before branching no longer shadows the branch's input. Simulation.derivative keeps inputs set after the simulation was built. An input a formula sets for the period it is calculating is no longer overwritten by the formula's result, and dump_simulation on a branch dumps the branch's own values (it read the default branch's, so the dump could not be restored). +Branch input changes invalidate derived values instead of inheriting stale results, derivatives preserve later inputs, formulas retain own-period inputs, branch dumps preserve branch values, and copied disk storages retain independent invalidation order. diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 54ca4050d..169c7f1f5 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -206,8 +206,10 @@ 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. - state.setdefault("_sequence_numbers", {}) + # 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. 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( From 9b9552d0cb313f2232b40f803069bb2dce9f50b4 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 11:05:49 -0400 Subject: [PATCH 29/42] Settle repeated branch input transitions without caching stale reads --- .../branch-set-input-invalidation.changed.md | 2 +- docs/usage/branches.md | 74 +++- policyengine_core/holders/holder.py | 17 + policyengine_core/simulations/simulation.py | 201 ++++++++++- tests/core/test_branch_input_fixed_point.py | 326 ++++++++++++++++++ .../test_branch_input_fixed_point_property.py | 40 +++ 6 files changed, 641 insertions(+), 19 deletions(-) create mode 100644 tests/core/test_branch_input_fixed_point.py create mode 100644 tests/core/test_branch_input_fixed_point_property.py diff --git a/changelog.d/branch-set-input-invalidation.changed.md b/changelog.d/branch-set-input-invalidation.changed.md index ae1d62d69..5c884104d 100644 --- a/changelog.d/branch-set-input-invalidation.changed.md +++ b/changelog.d/branch-set-input-invalidation.changed.md @@ -1 +1 @@ -apply_reform drops every stored value set_input did not store, in place, instead of wiping the storages and replaying _user_input_keys; restore_simulation restores calculated values after the inputs, under one later number; a calculation that may have read a value a branch input dropped (one running in that branch, or in any simulation that got a value from it before the change, also from a thread started without the formula's context) is not cached, and the outermost such calculation in each simulation runs again until it reads nothing changed since (at most ten times); setting a branch input to the input the branch already reads for that period (not through a set_input helper) drops nothing; a branch stops reading macro-cache files once an input is set on it. +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 stored inputs in place, and restored calculated values are invalidated conservatively. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index c1d996f6a..c6bdd7193 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -65,7 +65,8 @@ 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. (A calculation nested in another in the same -simulation shares that one's record, which can only make it keep less.) So a +simulation shares that one's record, which can only make it keep less and can +cause extra inner recalculations when formulas call between simulations.) 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 @@ -82,7 +83,32 @@ 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 input set again with the value the branch already reads for that very +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 by its path of names from the same root simulation, so a +branch recreated by `get_branch` under the same name on each attempt can settle +too. Direct clones retain separate identities even when their names match. Using +multiple distinct branches at the same path within one attempt prevents this +shortcut. + +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 @@ -151,12 +177,17 @@ depend on the input is calculated again) but not less, within these limits: 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 drop nothing.** `set_input` on a simulation - that is not a branch keeps its earlier behaviour: values it already - calculated stay. + that is not a branch keeps values it already calculated. 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 @@ -166,15 +197,26 @@ depend on the input is calculated again) but not less, within these limits: 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. -- **Formulas that create a branch each run.** A formula that creates a branch, - reads it, and then changes its input runs again until the budget is spent - (each run reads a new branch from before the change), then returns its last - result without keeping it. +- **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`), a value dropped from a branch only leaves -its storage's index: the file stays until the storage directory is removed, -since a clone of the storage may still read it, and storing the key again -writes a file of the branch's own (see `OnDiskStorage._path_to_write`). +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 @@ -194,8 +236,12 @@ 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 requested before a -drop removed its values. +`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 diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 945b13874..a14acb5ec 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -107,6 +107,12 @@ def delete_arrays( If ``period`` is not ``None``, only remove all values for any period included in period (e.g. if period is "2017", values for "2017-01", "2017-07", etc. would be removed) """ + 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) @@ -508,6 +514,17 @@ def _set( sequence_number=sequence_number, ) self._record_store(period, sequence_number) + if not derived and simulation is not None: + from policyengine_core.simulations.simulation import _note_input_change + + stored_period = ( + periods.period(periods.ETERNITY) + if self.variable.definition_period == periods.ETERNITY + else periods.period(period) + ) + _note_input_change( + simulation, "set", self.variable.name, stored_period, branch_name, value + ) if user_input_contexts: if not hasattr(simulation, "_user_input_keys"): simulation._user_input_keys = set() diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index c2c635296..d15551946 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1,7 +1,7 @@ import hashlib import os -import types from contextlib import nullcontext +import types from contextvars import ContextVar from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union @@ -219,7 +219,19 @@ class _Frame: not a generator or a dict: one runs for every calculation. """ - __slots__ = ("simulation", "start", "reads", "outer", "token", "registry") + __slots__ = ( + "simulation", + "start", + "reads", + "outer", + "token", + "registry", + "changes", + "drop_epochs", + "untracked_changes", + "inputs_set", + "branch_identities", + ) def __init__(self, simulation: "Simulation"): self.simulation = simulation @@ -227,7 +239,59 @@ def __init__(self, simulation: "Simulation"): 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 = {} + self.changes = [] + self.drop_epochs = {} + self.untracked_changes = False + self.branch_identities = {} + + 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 + return tuple(self.changes), 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.""" @@ -241,8 +305,7 @@ def is_stale(self) -> bool: def __enter__(self) -> "_Frame": simulation = self.simulation - self.start = simulation._input_epoch - self.reads = {} + self.restart() self.outer = outer = _calculation_frames.get() self.token = _calculation_frames.set(outer + (self,)) registry = simulation._open_frames @@ -294,6 +357,7 @@ def _hand_to_caller(outer: tuple, frame: _Frame) -> None: _hand_to_waiting_ancestors(simulation, frame.start, reads) return caller = outer[-1] + caller.untracked_changes |= frame.untracked_changes caller_simulation = caller.simulation caller_reads = caller.reads for other, epoch in reads: @@ -319,6 +383,7 @@ def _hand_to_waiting_ancestors( 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: @@ -327,6 +392,86 @@ def _hand_to_waiting_ancestors( 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_path(simulation: "Simulation") -> 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 + ): + 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 + 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, + _branch_path(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. + """ + 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) + + class Simulation: """ Represents a simulation, and handles the calculation logic @@ -611,6 +756,8 @@ def _invalidate_all_caches(self) -> None: for population in self.populations.values(): for holder in population._holders.values(): holder._mark_derived_except(user_inputs.get(holder.variable.name, ())) + for frame in _calculation_frames.get(): + frame.untracked_changes = True self._drop_computed() for branch in self.branches.values(): branch._invalidate_all_caches() @@ -977,9 +1124,37 @@ def calculate( 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): + 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 @@ -1546,9 +1721,15 @@ def total(): 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() @@ -1929,6 +2110,8 @@ def drop_computed_arrays(self) -> int: """ # 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]: @@ -2106,6 +2289,11 @@ def _drop_computed(self, since: Optional[int] = None) -> int: # 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. @@ -2281,6 +2469,10 @@ def clone( # 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 branch shares its parent's baseline: formulas that run in a # branch read the parent's baseline values through it. That holds for @@ -2374,6 +2566,7 @@ def get_branch( self.branches[name] = branch branch.branch_name = name branch.parent_branch = self + branch._fixed_point_branch = True if self.trace: branch.trace = True branch.tracer = self.tracer 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..723ee7441 --- /dev/null +++ b/tests/core/test_branch_input_fixed_point.py @@ -0,0 +1,326 @@ +"""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", +): + 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") + 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 + 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 1 < len(calls) <= 3 + 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) + _assert_settled( + simulation, + calls, + _input_first_result(values), + kept_in_holder=cache_mode == "disk", + ) + + +@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)) + + +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 12 < len(calls) <= 36 + 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 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..69d835c9d --- /dev/null +++ b/tests/core/test_branch_input_fixed_point_property.py @@ -0,0 +1,40 @@ +"""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"), + ) From 1a2576c9ced20e93c9e0b4cdad212d925f5050c1 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 11:25:06 -0400 Subject: [PATCH 30/42] Keep fixed-point readers and saved branch snapshots distinct --- docs/usage/branches.md | 8 ++- policyengine_core/simulations/simulation.py | 28 +++++++-- tests/core/test_branch_input_fixed_point.py | 68 +++++++++++++++++++++ 3 files changed, 96 insertions(+), 8 deletions(-) diff --git a/docs/usage/branches.md b/docs/usage/branches.md index c6bdd7193..4fd5c7996 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -91,9 +91,11 @@ 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 by its path of names from the same root simulation, so a -branch recreated by `get_branch` under the same name on each attempt can settle -too. Direct clones retain separate identities even when their names match. Using +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. Existing branches 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. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index d15551946..f138d4079 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -231,6 +231,7 @@ class _Frame: "untracked_changes", "inputs_set", "branch_identities", + "created_branches", ) def __init__(self, simulation: "Simulation"): @@ -245,6 +246,7 @@ def restart(self) -> None: self.drop_epochs = {} self.untracked_changes = False self.branch_identities = {} + self.created_branches = set() def note_branch(self, simulation: "Simulation") -> None: """A name path must refer to just one simulation within an attempt. @@ -287,7 +289,17 @@ def transition_signature(self, result: ArrayLike) -> Optional[tuple]: value = _value_signature(result) if value is None: return None - return tuple(self.changes), value + # A saved branch snapshot can share a name with a later recreation. + # Only creations observed in this attempt may compare by name; + # existing branches, including read-only ones, keep their identities. + changes = tuple( + (operation, _branch_path(other, self.created_branches), *details) + for operation, other, *details in self.changes + ) + reads = frozenset( + _branch_path(other, self.created_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.""" @@ -407,7 +419,9 @@ def _value_signature(value: ArrayLike) -> Optional[tuple]: return array.dtype.str, array.shape, enum, array.tobytes() -def _branch_path(simulation: "Simulation") -> tuple: +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 @@ -415,8 +429,10 @@ def _branch_path(simulation: "Simulation") -> tuple: 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 + 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 @@ -443,7 +459,7 @@ def _note_input_change( return change = ( operation, - _branch_path(simulation), + simulation, branch_name, variable, period, @@ -2567,6 +2583,8 @@ def get_branch( branch.branch_name = name branch.parent_branch = self branch._fixed_point_branch = True + for frame in _calculation_frames.get(): + frame.created_branches.add(branch) if self.trace: branch.trace = True branch.tracer = self.tracer diff --git a/tests/core/test_branch_input_fixed_point.py b/tests/core/test_branch_input_fixed_point.py index 723ee7441..876e8de40 100644 --- a/tests/core/test_branch_input_fixed_point.py +++ b/tests/core/test_branch_input_fixed_point.py @@ -131,6 +131,13 @@ def change(): 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 @@ -172,12 +179,21 @@ def _assert_settled(simulation, calls, expected, kept_in_holder=True): 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]) @@ -324,3 +340,55 @@ def test_calculate_add_keeps_a_sum_after_persistent_branch_resets_settle(): 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 From 542bb5b47face73e8d43d8ebc77025b7e4cffa51 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 19:51:01 -0400 Subject: [PATCH 31/42] Retain registered branch paths across calculation retries --- .../branch-set-input-invalidation.fixed.md | 2 +- docs/usage/branches.md | 13 +++-- policyengine_core/simulations/simulation.py | 36 +++++++++++--- tests/core/test_branch_input_fixed_point.py | 47 ++++++++++++++++++- .../test_branch_input_fixed_point_property.py | 32 +++++++++++++ 5 files changed, 117 insertions(+), 13 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.fixed.md b/changelog.d/branch-set-input-invalidation.fixed.md index fcd147b9c..cde6c2529 100644 --- a/changelog.d/branch-set-input-invalidation.fixed.md +++ b/changelog.d/branch-set-input-invalidation.fixed.md @@ -1 +1 @@ -Branch input changes invalidate derived values instead of inheriting stale results, derivatives preserve later inputs, formulas retain own-period inputs, branch dumps preserve branch values, and copied disk storages retain independent invalidation order. +Branch input changes invalidate derived values instead of inheriting stale results, repeated resets on reused branches settle in two attempts without conflating saved snapshots, derivatives preserve later inputs, formulas retain own-period inputs, branch dumps preserve branch values, and copied disk storages retain independent invalidation order. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 4fd5c7996..dda654601 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -66,7 +66,8 @@ 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. (A calculation nested in another in the same simulation shares that one's record, which can only make it keep less and can -cause extra inner recalculations when formulas call between simulations.) So a +cause extra inner recalculations when formulas call between simulations: a +stale attempt leaves nested results uncached, so retries repeat that work.) 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 @@ -94,10 +95,14 @@ 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. Existing branches and direct clones -retain separate identities even when their names match. Using +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. +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 diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index f138d4079..8d19bfc00 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -232,10 +232,13 @@ class _Frame: "inputs_set", "branch_identities", "created_branches", + "carried_branches", ) def __init__(self, simulation: "Simulation"): self.simulation = simulation + self.created_branches = {} + self.carried_branches = {} def restart(self) -> None: """Begin again, as when the calculation runs again.""" @@ -246,7 +249,12 @@ def restart(self) -> None: self.drop_epochs = {} self.untracked_changes = False self.branch_identities = {} - self.created_branches = set() + # 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. @@ -290,14 +298,19 @@ def transition_signature(self, result: ArrayLike) -> Optional[tuple]: if value is None: return None # A saved branch snapshot can share a name with a later recreation. - # Only creations observed in this attempt may compare by name; - # existing branches, including read-only ones, keep their identities. + # 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, self.created_branches), *details) + (operation, _branch_path(other, comparable_branches), *details) for operation, other, *details in self.changes ) reads = frozenset( - _branch_path(other, self.created_branches) for other in self.reads + _branch_path(other, comparable_branches) for other in self.reads ) return changes, reads, value @@ -419,6 +432,17 @@ def _value_signature(value: ArrayLike) -> Optional[tuple]: return array.dtype.str, array.shape, enum, array.tobytes() +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: @@ -2584,7 +2608,7 @@ def get_branch( branch.parent_branch = self branch._fixed_point_branch = True for frame in _calculation_frames.get(): - frame.created_branches.add(branch) + frame.created_branches[branch] = (self, name) if self.trace: branch.trace = True branch.tracer = self.tracer diff --git a/tests/core/test_branch_input_fixed_point.py b/tests/core/test_branch_input_fixed_point.py index 876e8de40..d369c24af 100644 --- a/tests/core/test_branch_input_fixed_point.py +++ b/tests/core/test_branch_input_fixed_point.py @@ -63,6 +63,7 @@ def _simulation( masked_inputs=False, recreate_branch=False, cache_mode="memory", + branch_depth=1, ): values = np.asarray(values, dtype=np.float32) calls = [] @@ -83,6 +84,8 @@ def outer(person, period): # 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") @@ -164,7 +167,7 @@ def _input_first_result(values): def _assert_settled(simulation, calls, expected, kept_in_holder=True): result = simulation.calculate("fp_outer", "2020") np.testing.assert_array_equal(result, expected) - assert 1 < len(calls) <= 3 + assert len(calls) == 2 stored = simulation.get_array("fp_outer", "2020") if kept_in_holder: np.testing.assert_array_equal(stored, expected) @@ -208,6 +211,46 @@ def test_branch_delete_read_set_reaches_a_cached_fixed_point( _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_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) @@ -335,7 +378,7 @@ def test_calculate_add_keeps_a_sum_after_persistent_branch_resets_settle(): np.testing.assert_array_equal( simulation.calculate_add("fp_outer", "2020"), expected ) - assert 12 < len(calls) <= 36 + 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) diff --git a/tests/core/test_branch_input_fixed_point_property.py b/tests/core/test_branch_input_fixed_point_property.py index 69d835c9d..2d5358132 100644 --- a/tests/core/test_branch_input_fixed_point_property.py +++ b/tests/core/test_branch_input_fixed_point_property.py @@ -38,3 +38,35 @@ def test_branch_reset_fixed_point_matches_inputs_first( _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"), + ) From 5f5e987828e080dd90d04d4adce73f8160658734 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 20:05:33 -0400 Subject: [PATCH 32/42] Guard fixed-point writes to optional fast caches --- changelog.d/branch-set-input-invalidation.fixed.md | 2 +- policyengine_core/simulations/simulation.py | 3 ++- tests/core/test_branch_input_fixed_point.py | 7 +++++++ 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.fixed.md b/changelog.d/branch-set-input-invalidation.fixed.md index cde6c2529..4f8f0aac9 100644 --- a/changelog.d/branch-set-input-invalidation.fixed.md +++ b/changelog.d/branch-set-input-invalidation.fixed.md @@ -1 +1 @@ -Branch input changes invalidate derived values instead of inheriting stale results, repeated resets on reused branches settle in two attempts without conflating saved snapshots, derivatives preserve later inputs, formulas retain own-period inputs, branch dumps preserve branch values, and copied disk storages retain independent invalidation order. +Branch input changes invalidate derived values instead of inheriting stale results, repeated resets on reused branches settle in two attempts without conflating saved snapshots, fixed-point calculations support missing optional fast caches, derivatives preserve later inputs, formulas retain own-period inputs, branch dumps preserve branch values, and copied disk storages retain independent invalidation order. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 8d19bfc00..9f81ec25e 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1182,7 +1182,8 @@ def calculate( keep_state, ) if self._may_keep(keep_state): - self._fast_cache[(variable_name, period)] = result + 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( diff --git a/tests/core/test_branch_input_fixed_point.py b/tests/core/test_branch_input_fixed_point.py index d369c24af..b51f4964e 100644 --- a/tests/core/test_branch_input_fixed_point.py +++ b/tests/core/test_branch_input_fixed_point.py @@ -251,6 +251,13 @@ def outer(person, period): 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) From 09ae8604c58cb92738ae9b99afefe7742cf3985b Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 02:40:22 -0400 Subject: [PATCH 33/42] Keep independent calculations cached during stale branch retries --- .../branch-set-input-invalidation.fixed.md | 2 +- docs/usage/branches.md | 14 +- policyengine_core/simulations/simulation.py | 75 +++- tests/core/test_stale_frame_cache.py | 385 ++++++++++++++++++ 4 files changed, 459 insertions(+), 17 deletions(-) create mode 100644 tests/core/test_stale_frame_cache.py diff --git a/changelog.d/branch-set-input-invalidation.fixed.md b/changelog.d/branch-set-input-invalidation.fixed.md index 4f8f0aac9..196699707 100644 --- a/changelog.d/branch-set-input-invalidation.fixed.md +++ b/changelog.d/branch-set-input-invalidation.fixed.md @@ -1 +1 @@ -Branch input changes invalidate derived values instead of inheriting stale results, repeated resets on reused branches settle in two attempts without conflating saved snapshots, fixed-point calculations support missing optional fast caches, derivatives preserve later inputs, formulas retain own-period inputs, branch dumps preserve branch values, and copied disk storages retain independent invalidation order. +Branch input changes invalidate derived values instead of inheriting stale results, independent nested calculations retain their caches during stale retries without hiding foreign reads or effects, repeated resets on reused branches settle in two attempts without conflating saved snapshots, fixed-point calculations support missing optional fast caches, derivatives preserve later inputs, formulas retain own-period inputs, branch dumps preserve branch values, and copied disk storages retain independent invalidation order. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index dda654601..6a6a52d23 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -64,12 +64,14 @@ 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. (A calculation nested in another in the same -simulation shares that one's record, which can only make it keep less and can -cause extra inner recalculations when formulas call between simulations: a -stale attempt leaves nested results uncached, so retries repeat that work.) 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 +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, +create branches, or observe untracked activity remain unkept so the caller's +retry observes the same reads and repeats those operations. 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 diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 9f81ec25e..85586a0f1 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -3,6 +3,7 @@ from contextlib import nullcontext import types from contextvars import ContextVar +from threading import Lock from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union import numpy as np @@ -198,6 +199,19 @@ def _latest_input_key(period: Period, definition_period: str) -> tuple: # 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() + + +def _note_unobserved_activity() -> None: + if _restricted_frames: + with _restricted_frames_lock: + for frame in _restricted_frames: + frame.unobserved_activity = True + class _Frame: """A calculation running in ``simulation``. @@ -233,12 +247,15 @@ class _Frame: "branch_identities", "created_branches", "carried_branches", + "stale_parent", + "unobserved_activity", ) - def __init__(self, simulation: "Simulation"): + 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.""" @@ -248,6 +265,7 @@ def restart(self) -> None: self.changes = [] self.drop_epochs = {} self.untracked_changes = False + self.unobserved_activity = False self.branch_identities = {} # Retain only creations whose original registration still names # the same object. Saved or replaced snapshots keep their identities. @@ -328,6 +346,28 @@ def is_stale(self) -> bool: 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 __enter__(self) -> "_Frame": simulation = self.simulation self.restart() @@ -339,9 +379,15 @@ def __enter__(self) -> "_Frame": 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) @@ -353,17 +399,21 @@ def _frame_for(simulation: "Simulation", outermost: bool) -> Optional[_Frame]: the frame it runs in. A calculation nested in another in the same simulation shares that one's - frame: 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 in a simulation has its - own, to run again from (``Simulation.calculate``). + 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 @@ -405,6 +455,7 @@ def _hand_to_waiting_ancestors( ``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 ()): @@ -506,6 +557,7 @@ def _taint_waiting_frames(simulation: "Simulation") -> None: 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 @@ -1641,7 +1693,7 @@ def _may_keep(self, input_state: Tuple[int, int, int]) -> bool: if input_state[0] != self._input_epoch or input_state[1] != self._inputs_set: return False frames = _calculation_frames.get() - return not (frames and frames[-1].is_stale()) + 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``) @@ -1688,8 +1740,8 @@ def _cache_result( if stored_input is not None: return stored_input frames = _calculation_frames.get() - stale = bool(frames) and frames[-1].is_stale() - if epoch == self._input_epoch and not stale: + 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. @@ -2608,7 +2660,10 @@ def get_branch( branch.branch_name = name branch.parent_branch = self branch._fixed_point_branch = True - for frame in _calculation_frames.get(): + frames = _calculation_frames.get() + if not frames: + _note_unobserved_activity() + for frame in frames: frame.created_branches[branch] = (self, name) if self.trace: branch.trace = True diff --git a/tests/core/test_stale_frame_cache.py b/tests/core/test_stale_frame_cache.py new file mode 100644 index 000000000..c682cebd1 --- /dev/null +++ b/tests/core/test_stale_frame_cache.py @@ -0,0 +1,385 @@ +"""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 + + +@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 + + +def test_unrelated_unread_thread_mutation_does_not_prevent_caller_convergence(): + 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(): + 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} + 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 From 3e14ed6ed68faf341661eb8248b0db9e13722c4a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 03:31:43 -0400 Subject: [PATCH 34/42] Cover shared credit dependencies during stale retries --- tests/core/test_stale_frame_cache.py | 39 ++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/core/test_stale_frame_cache.py b/tests/core/test_stale_frame_cache.py index c682cebd1..481e77f88 100644 --- a/tests/core/test_stale_frame_cache.py +++ b/tests/core/test_stale_frame_cache.py @@ -119,6 +119,45 @@ def outer(person, period): 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("suffix_branch", ["worker", "other"]) def test_foreign_suffix_is_not_cached_before_its_stale_caller_settles(suffix_branch): calls = Counter() From 0853d64bf766b243dd5108da7584c0258d251b51 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 06:42:32 -0400 Subject: [PATCH 35/42] Isolate early-year branch replay regression from fallback parsing --- tests/core/test_branch_helper_invalidation.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/core/test_branch_helper_invalidation.py b/tests/core/test_branch_helper_invalidation.py index 0488a5587..ba66b7847 100644 --- a/tests/core/test_branch_helper_invalidation.py +++ b/tests/core/test_branch_helper_invalidation.py @@ -80,6 +80,14 @@ def build(): 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]) @@ -109,9 +117,8 @@ def build(): np.testing.assert_array_equal(parent.calculate("helper_result", year), [24, 48]) branch_inputs = {("helper_flow", "changed", month) for month in months} - parent_input = ("helper_flow", "default", january) - assert branch_inputs | {parent_input} == branch._user_input_keys - assert parent._user_input_keys == {parent_input} + 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: @@ -119,7 +126,7 @@ def build(): branch._invalidate_all_caches() - assert branch._user_input_keys == branch_inputs | {parent_input} + 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]) From 77c29bf918761cfb07af1d79964a32605413d3b3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 07:04:50 -0400 Subject: [PATCH 36/42] Keep raw cache writes observable during guarded retries --- .../branch-set-input-invalidation.fixed.md | 2 ++ docs/usage/branches.md | 15 ++++++++----- policyengine_core/holders/holder.py | 12 ++++++++-- tests/core/test_stale_frame_cache.py | 22 ++++++++++++++++--- 4 files changed, 41 insertions(+), 10 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.fixed.md b/changelog.d/branch-set-input-invalidation.fixed.md index 196699707..f996dc390 100644 --- a/changelog.d/branch-set-input-invalidation.fixed.md +++ b/changelog.d/branch-set-input-invalidation.fixed.md @@ -1 +1,3 @@ Branch input changes invalidate derived values instead of inheriting stale results, independent nested calculations retain their caches during stale retries without hiding foreign reads or effects, repeated resets on reused branches settle in two attempts without conflating saved snapshots, fixed-point calculations support missing optional fast caches, derivatives preserve later inputs, formulas retain own-period inputs, branch dumps preserve branch values, and copied disk storages retain independent invalidation order. + +Raw holder cache writes remain observable during guarded stale retries without being replayed or exported as supplied inputs. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 6190779a1..4b79c95cf 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -68,8 +68,12 @@ 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, -create branches, or observe untracked activity remain unkept so the caller's -retry observes the same reads and repeats those operations. So a parent formula +write raw cache values, create branches, or observe untracked activity remain +unkept so the caller's retry observes the same reads and repeats those operations. +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 @@ -265,9 +269,10 @@ branches for comparisons that need both the earlier and the later 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; also values cached with -`Holder.put_in_cache` without `derived=True`), and returns how many arrays it -deleted. Values a custom `set_input` handler calculates are not inputs. Use it on a +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 diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index a6947423a..91c52b8a8 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -779,12 +779,20 @@ def _set( "disk" if should_store_on_disk else "memory", is_input=is_input, ) - if is_input and simulation is not None: + 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, branch_name, value + 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() # Record the period as storage keys the value (eternity for an diff --git a/tests/core/test_stale_frame_cache.py b/tests/core/test_stale_frame_cache.py index 481e77f88..6369e0e2c 100644 --- a/tests/core/test_stale_frame_cache.py +++ b/tests/core/test_stale_frame_cache.py @@ -377,7 +377,10 @@ def outer(person, period): assert calls == completed -def test_unrelated_unread_thread_mutation_does_not_prevent_caller_convergence(): +@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") @@ -389,8 +392,13 @@ def suffix(person, period): calls["suffix"] += 1 def change(): - unrelated_worker.delete_arrays("sf_source") - unrelated_worker.set_input("sf_source", period, np.array([5.0])) + 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) @@ -418,6 +426,14 @@ def outer(person, period): # 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) From 880ab0d81aa9d0169650025be91acbfc3c15cf8d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 07:09:22 -0400 Subject: [PATCH 37/42] Keep branch invalidation changelog concise --- changelog.d/branch-set-input-invalidation.fixed.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.fixed.md b/changelog.d/branch-set-input-invalidation.fixed.md index f996dc390..79e4ba811 100644 --- a/changelog.d/branch-set-input-invalidation.fixed.md +++ b/changelog.d/branch-set-input-invalidation.fixed.md @@ -1,3 +1 @@ -Branch input changes invalidate derived values instead of inheriting stale results, independent nested calculations retain their caches during stale retries without hiding foreign reads or effects, repeated resets on reused branches settle in two attempts without conflating saved snapshots, fixed-point calculations support missing optional fast caches, derivatives preserve later inputs, formulas retain own-period inputs, branch dumps preserve branch values, and copied disk storages retain independent invalidation order. - -Raw holder cache writes remain observable during guarded stale retries without being replayed or exported as supplied inputs. +Branch input changes invalidate derived values while preserving supplied inputs, independent caches, guarded retry effects including raw cache writes, fixed-point settlement, branch dump values, and separate disk invalidation order. From 041dcdad87ddc8c35c6cb52a8062387766c30376 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 07:36:11 -0400 Subject: [PATCH 38/42] Clarify root simulation helper cache cleanup --- docs/usage/branches.md | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 4b79c95cf..3208690e5 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -205,10 +205,12 @@ depend on the input is calculated again) but not less, within these limits: - **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 drop nothing.** `set_input` on a simulation - that is not a branch keeps values it already calculated. 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. +- **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 From 73b90c0fb4728e7abaae63b143eebad8a5cf93b5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 09:59:21 -0400 Subject: [PATCH 39/42] Forward store sequence numbers in restore property observer --- .../test_restore_input_registry_property.py | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) 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 From bbb7b6fb26faa68f6f09c49b5c615c41208c2dfd Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 18:30:08 -0400 Subject: [PATCH 40/42] Reuse stable foreign reads within stale calculation attempts --- .../branch-set-input-invalidation.fixed.md | 2 +- docs/usage/branches.md | 22 +- policyengine_core/simulations/simulation.py | 109 ++++++++- tests/core/test_stale_frame_cache.py | 219 ++++++++++++++++++ 4 files changed, 346 insertions(+), 6 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.fixed.md b/changelog.d/branch-set-input-invalidation.fixed.md index 79e4ba811..30f960ec9 100644 --- a/changelog.d/branch-set-input-invalidation.fixed.md +++ b/changelog.d/branch-set-input-invalidation.fixed.md @@ -1 +1 @@ -Branch input changes invalidate derived values while preserving supplied inputs, independent caches, guarded retry effects including raw cache writes, fixed-point settlement, branch dump values, and separate disk invalidation order. +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. diff --git a/docs/usage/branches.md b/docs/usage/branches.md index 3208690e5..43db30f5f 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -69,7 +69,27 @@ 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 -unkept so the caller's retry observes the same reads and repeats those operations. +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. It also passes 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 diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 11136a5fb..9cefff1fe 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -205,13 +205,44 @@ def _latest_input_key(period: Period, definition_period: str) -> tuple: # 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: - if _restricted_frames: - with _restricted_frames_lock: - for frame in _restricted_frames: - frame.unobserved_activity = True + 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: @@ -250,6 +281,7 @@ class _Frame: "carried_branches", "stale_parent", "unobserved_activity", + "attempt_cache", ) def __init__(self, simulation: "Simulation", stale_parent: bool = False): @@ -268,6 +300,9 @@ def restart(self) -> None: 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( @@ -369,10 +404,58 @@ def may_keep(self) -> bool: 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() + ): + return + registrations = set() + for other in self.reads: + while getattr(other, "parent_branch", None) is not None: + parent = other.parent_branch + name = other.branch_name + if parent.branches.get(name) is other: + registrations.add((parent, name, other)) + other = parent + 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(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) + # 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: @@ -528,6 +611,10 @@ def _note_input_change( 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: @@ -1524,6 +1611,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: @@ -1866,6 +1959,13 @@ def _cache_result( 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: @@ -2821,6 +2921,7 @@ def get_branch( 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 diff --git a/tests/core/test_stale_frame_cache.py b/tests/core/test_stale_frame_cache.py index 6369e0e2c..76ff782e4 100644 --- a/tests/core/test_stale_frame_cache.py +++ b/tests/core/test_stale_frame_cache.py @@ -158,6 +158,225 @@ def outer(person, period): 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 + + +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() From 26019baa2296ec9e9135faf09c7b8e792f22a0e5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 20:53:40 -0400 Subject: [PATCH 41/42] Capture branch registrations before memo-producing formulas run --- .../branch-set-input-invalidation.fixed.md | 2 +- docs/usage/branches.md | 6 +- policyengine_core/simulations/simulation.py | 37 +++++-- tests/core/test_stale_frame_cache.py | 101 ++++++++++++++++++ 4 files changed, 135 insertions(+), 11 deletions(-) diff --git a/changelog.d/branch-set-input-invalidation.fixed.md b/changelog.d/branch-set-input-invalidation.fixed.md index 30f960ec9..390794265 100644 --- a/changelog.d/branch-set-input-invalidation.fixed.md +++ b/changelog.d/branch-set-input-invalidation.fixed.md @@ -1 +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. +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/usage/branches.md b/docs/usage/branches.md index 43db30f5f..bebc38435 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -74,7 +74,11 @@ 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. It also passes the +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. 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 diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 9cefff1fe..473d18600 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -269,6 +269,7 @@ class _Frame: "simulation", "start", "reads", + "registrations", "outer", "token", "registry", @@ -295,6 +296,9 @@ def restart(self) -> None: 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 @@ -414,16 +418,12 @@ def remember(self, variable: str, period: Period, array: ArrayLike) -> None: 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 - registrations = set() - for other in self.reads: - while getattr(other, "parent_branch", None) is not None: - parent = other.parent_branch - name = other.branch_name - if parent.branches.get(name) is other: - registrations.add((parent, name, other)) - other = parent self.attempt_cache[(self.simulation, variable, period)] = _AttemptResult( array, self.start, @@ -431,7 +431,7 @@ def remember(self, variable: str, period: Period, array: ArrayLike) -> None: tuple( (other, epoch, other._inputs_set) for other, epoch in self.reads.items() ), - tuple(registrations), + tuple(self.registrations), _unobserved_activity_generation, ) @@ -446,6 +446,7 @@ def recall(self, variable: str, period: Period) -> Optional[ArrayLike]: 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 @@ -517,6 +518,7 @@ def _hand_to_caller(outer: tuple, frame: _Frame) -> None: 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: @@ -567,6 +569,22 @@ def _value_signature(value: ArrayLike) -> Optional[tuple]: 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 { @@ -1388,6 +1406,7 @@ def calculate( _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) diff --git a/tests/core/test_stale_frame_cache.py b/tests/core/test_stale_frame_cache.py index 76ff782e4..7e0c214c7 100644 --- a/tests/core/test_stale_frame_cache.py +++ b/tests/core/test_stale_frame_cache.py @@ -337,6 +337,107 @@ def outer(person, period): 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"]) +def test_registration_replaced_inside_memo_producer_preserves_the_next_read( + registration_depth, + prime_fast_cache, + read_route, +): + 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("worker") + 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 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() From 8c5a48a88443679bd1f6c17ed20e96490f7f51e8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 22:01:50 -0400 Subject: [PATCH 42/42] Guard attempt reuse with the branch name actually resolved --- docs/usage/branches.md | 10 ++++++---- policyengine_core/simulations/simulation.py | 8 +++++++- tests/core/test_stale_frame_cache.py | 10 +++++++++- 3 files changed, 22 insertions(+), 6 deletions(-) diff --git a/docs/usage/branches.md b/docs/usage/branches.md index bebc38435..1465c01e7 100644 --- a/docs/usage/branches.md +++ b/docs/usage/branches.md @@ -76,10 +76,12 @@ 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. 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 +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. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 473d18600..e4628e244 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -2924,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: diff --git a/tests/core/test_stale_frame_cache.py b/tests/core/test_stale_frame_cache.py index 7e0c214c7..b59053725 100644 --- a/tests/core/test_stale_frame_cache.py +++ b/tests/core/test_stale_frame_cache.py @@ -340,10 +340,12 @@ def outer(person, period): @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 = [] @@ -406,7 +408,11 @@ def outer(person, period): "sf_outer": outer, }, ) - original = simulation.get_branch("worker") + 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 = ( @@ -428,6 +434,8 @@ def outer(person, period): 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)