Make set_input on a branch drop values calculated from the input it replaces

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.

changelog.d/branch-set-input-invalidation.changed.md

@@ -0,0 +1 @@
+Branch input changes prevent caching stale calculations and rerun outermost `calculate` and `calculate_add` calls up to ten times, with exact repeated input transitions and results in other branches able to settle under the usual cache restrictions; unchanged direct inputs skip invalidation, branch input stores disable macro-cache reads, `apply_reform` preserves surviving supplied input tiers, branch input helpers share overlapping-cache cleanup with ordinary input helpers, and restored calculated values are invalidated conservatively.

changelog.d/branch-set-input-invalidation.fixed.md

@@ -0,0 +1 @@
+Branch input changes invalidate derived values while preserving supplied inputs, independent caches, guarded retry effects, fixed-point settlement, branch dumps and disk invalidation order, with read-only foreign calculations reused within each attempt only while their captured input and branch-registration dependencies remain current.

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

policyengine_core/data_storage/in_memory_storage.py

@@ -1,9 +1,14 @@
-from typing import Dict, FrozenSet, Set, Union
+from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Union
 
 import numpy
 from numpy.typing import ArrayLike
 
 from policyengine_core import periods
+from policyengine_core.data_storage.store_history import (
+    _period_from_storage_key,
+    advance_sequence_past,
+    next_sequence_number,
+)
 from policyengine_core.periods import Period
 
 
@@ -83,6 +88,10 @@ class InMemoryStorage:
 
     def __init__(self, is_eternal: bool):
         self._arrays = {}
+        # When each array was stored (see ``store_history``). A number
+        # describes the stored value, as its input mark does, so ``clone``
+        # copies the numbers with the arrays.
+        self._sequence_numbers: Dict[str, int] = {}
         self.is_eternal = is_eternal
 
     def __getstate__(self) -> dict:
@@ -100,7 +109,13 @@ def __setstate__(self, state: dict) -> None:
             inputs = set(state.get("_arrays", {})).difference(derived)
             if inputs:
                 state["_inputs"] = inputs
+        # Pickled before stores were numbered: no value has a number.
+        state.setdefault("_sequence_numbers", {})
         self.__dict__.update(state)
+        # Numbers from the process that pickled this storage must stay below
+        # those of stores made after unpickling it.
+        if self._sequence_numbers:
+            advance_sequence_past(max(self._sequence_numbers.values()))
 
     def clone(self, share_arrays: bool = False) -> "InMemoryStorage":
         """Copy this storage.
@@ -152,6 +167,7 @@ def clone(self, share_arrays: bool = False) -> "InMemoryStorage":
                 held = inputs.intersection(clone._arrays)
                 if held:
                     clone._inputs = frozenset(held)
+        clone._sequence_numbers = dict(self._sequence_numbers)
         return clone
 
     def get(self, period: Period, branch_name: str = "default") -> ArrayLike:
@@ -260,7 +276,14 @@ def put(
         period: Period,
         branch_name: str = "default",
         derived: bool = False,
+        sequence_number: Optional[int] = None,
     ) -> None:
+        """Store ``value`` for ``period`` on ``branch_name``.
+
+        ``derived`` marks a value the simulation calculated rather than took
+        as input (see ``_inputs``), and ``sequence_number`` records when the
+        value was stored (a new number by default); see :meth:`drop_computed`.
+        """
         if self.is_eternal:
             period = periods.period(periods.ETERNITY)
         period = periods.period(period)
@@ -287,6 +310,58 @@ def put(
             self._unmark_input(key)
         else:
             self._mark_input(key)
+        self._sequence_numbers[key] = (
+            next_sequence_number() if sequence_number is None else sequence_number
+        )
+
+    def drop_computed(self, *, since: Optional[int] = None) -> int:
+        """Delete stored values that are not inputs, and return how many.
+
+        With ``since``, only values stored with that sequence number or a
+        later one are deleted (a value with no recorded number counts as
+        later). Inputs, which ``put`` received without ``derived``, are kept
+        whatever their number.
+        """
+        inputs = self._inputs
+        numbers = self._sequence_numbers
+        dropped = [
+            key
+            for key in self._arrays
+            if key not in inputs and (since is None or numbers.get(key, since) >= since)
+        ]
+        for key in dropped:
+            del self._arrays[key]
+            numbers.pop(key, None)
+            self._stop_sharing(key)
+        return len(dropped)
+
+    def inputs_since(self, since: Optional[int] = None) -> List[Tuple[Period, int]]:
+        """The period and number of each input stored at ``since`` or later (or ever)."""
+        numbers = self._sequence_numbers
+        return [
+            (_period_from_storage_key(key.split(":", 1)[1]), numbers[key])
+            for key in self._inputs
+            if key in self._arrays
+            and key in numbers
+            and (since is None or numbers[key] >= since)
+        ]
+
+    def mark_derived_except(self, keys: Set[str]) -> None:
+        """Mark every stored value derived except those stored for ``keys``."""
+        for key in [key for key in self._inputs if key not in keys]:
+            self._unmark_input(key)
+
+    def has_unnumbered_values(self) -> bool:
+        """Whether a value was stored without ``put`` (so without a number)."""
+        return any(key not in self._sequence_numbers for key in self._arrays)
+
+    def _forget_dropped_numbers(self) -> None:
+        """Forget the numbers of keys ``_arrays`` no longer has."""
+        numbers = self._sequence_numbers
+        if not numbers.keys() <= self._arrays.keys():
+            self._sequence_numbers = {
+                key: number for key, number in numbers.items() if key in self._arrays
+            }
 
     def delete(self, period: Period = None, branch_name: str = "default") -> None:
         if period is None:
@@ -300,6 +375,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None:
             }
             self._stop_sharing_dropped_keys()
             self._unmark_dropped_keys()
+            self._forget_dropped_numbers()
             return
 
         if self.is_eternal:
@@ -319,6 +395,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None:
         }
         self._stop_sharing_dropped_keys()
         self._unmark_dropped_keys()
+        self._forget_dropped_numbers()
 
     def get_known_periods(self) -> list:
         # Split on the first colon only: an anchored period's string form

policyengine_core/data_storage/on_disk_storage.py

@@ -2,11 +2,17 @@
 import shutil
 import tempfile
 import weakref
+from typing import Dict, List, Optional, Set, Tuple
 
 import numpy
 from numpy.typing import ArrayLike
 
 from policyengine_core import periods
+from policyengine_core.data_storage.store_history import (
+    _period_from_storage_key,
+    advance_sequence_past,
+    next_sequence_number,
+)
 from policyengine_core.data_storage.storage_directory import (
     directory_containing,
     path_key,
@@ -91,6 +97,8 @@ def _start(self, storage_dir: str, is_eternal: bool, parent_directory) -> None:
         # File keys stored with ``put(..., derived=True)``; see
         # ``InMemoryStorage``.
         self._derived = set()
+        # When each file key was stored; see ``InMemoryStorage``.
+        self._sequence_numbers: Dict[str, int] = {}
         # For each key, the file this storage last wrote for it, while it has
         # not shared that file since: the only files of its family ``put``
         # writes over (see ``_path_to_write``). Kept when the key is deleted,
@@ -204,7 +212,15 @@ def __setstate__(self, state: dict) -> None:
             state["_storage_dir_owner"] = owner
         if state.get("_parent_directory") is None:
             state["_parent_directory"] = directory_containing(state["storage_dir"])
+        # Nor numbers: pickled before stores were numbered. Copies have their
+        # own numbers, like their file index: replacing or dropping a key
+        # through one must not change when the other's value was stored.
+        state["_sequence_numbers"] = dict(state.get("_sequence_numbers", {}))
         self.__dict__.update(state)
+        # Numbers from the process that pickled this storage must stay below
+        # those of stores made after unpickling it.
+        if self._sequence_numbers:
+            advance_sequence_past(max(self._sequence_numbers.values()))
 
     def clone(self) -> "OnDiskStorage":
         """Create a private metadata view over this storage directory.
@@ -225,6 +241,7 @@ def clone(self) -> "OnDiskStorage":
         clone._files = self._files.copy()
         clone._enums = self._enums.copy()
         clone._derived = set(self._derived)
+        clone._sequence_numbers = dict(self._sequence_numbers)
         # One set of read-back paths for the family, so a clone made before
         # ``restore`` still writes over none of them after a fork.
         clone._restored_paths = self._restored_paths
@@ -276,7 +293,13 @@ def put(
         period: Period,
         branch_name: str = "default",
         derived: bool = False,
+        sequence_number: Optional[int] = None,
     ) -> None:
+        """Store ``value`` for ``period`` on ``branch_name``.
+
+        ``derived`` and ``sequence_number`` are as in
+        :meth:`InMemoryStorage.put`.
+        """
         if self.is_eternal:
             period = periods.period(periods.ETERNITY)
         period = periods.period(period)
@@ -296,6 +319,59 @@ def put(
             self._derived.add(filename)
         else:
             self._derived.discard(filename)
+        self._sequence_numbers[filename] = (
+            next_sequence_number() if sequence_number is None else sequence_number
+        )
+
+    def drop_computed(self, *, since: Optional[int] = None) -> int:
+        """Forget stored values that are not inputs, and return how many.
+
+        As :meth:`InMemoryStorage.drop_computed`. The files stay on disk
+        (clones may still read them) until the storage directory is removed;
+        writing a dropped key again reuses this storage's own file for it,
+        as after ``delete``.
+        """
+        numbers = self._sequence_numbers
+        dropped = [
+            key
+            for key in self._files
+            if key in self._derived
+            and (since is None or numbers.get(key, since) >= since)
+        ]
+        for key in dropped:
+            del self._files[key]
+            numbers.pop(key, None)
+        self._derived.difference_update(dropped)
+        return len(dropped)
+
+    def inputs_since(self, since: Optional[int] = None) -> List[Tuple[Period, int]]:
+        """As :meth:`InMemoryStorage.inputs_since`."""
+        numbers = self._sequence_numbers
+        # Period strings contain no "_"; branch names may.
+        return [
+            (_period_from_storage_key(key.rsplit("_", 1)[1]), numbers[key])
+            for key in self._files
+            if key not in self._derived
+            and key in numbers
+            and (since is None or numbers[key] >= since)
+        ]
+
+    def mark_derived_except(self, keys: Set[str]) -> None:
+        """As :meth:`InMemoryStorage.mark_derived_except`."""
+        self._derived.update(key for key in self._files if key not in keys)
+
+    def has_unnumbered_values(self) -> bool:
+        """As :meth:`InMemoryStorage.has_unnumbered_values`."""
+        return any(key not in self._sequence_numbers for key in self._files)
+
+    def _forget_dropped_keys(self) -> None:
+        """Forget the derived marks and numbers of keys ``_files`` no longer has."""
+        self._derived.intersection_update(self._files)
+        numbers = self._sequence_numbers
+        if not numbers.keys() <= self._files.keys():
+            self._sequence_numbers = {
+                key: number for key, number in numbers.items() if key in self._files
+            }
 
     def _path_to_write(self, filename: str) -> str:
         """The path ``put`` writes the value for key ``filename`` to.
@@ -348,7 +424,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None:
                 for period_item, value in self._files.items()
                 if not period_item.startswith(branch_prefix)
             }
-            self._derived.intersection_update(self._files)
+            self._forget_dropped_keys()
             return
 
         if self.is_eternal:
@@ -361,7 +437,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None:
                 for period_item, value in self._files.items()
                 if not period_item == f"{branch_name}_{period}"
             }
-            self._derived.intersection_update(self._files)
+            self._forget_dropped_keys()
 
     def get_known_periods(self) -> list:
         return list([periods.period(x.split("_")[1]) for x in self._files.keys()])
@@ -393,8 +469,9 @@ def restore(self) -> None:
         are unchanged.
         """
         self._files = files = {}
-        # Files read back from a directory carry no derived marks.
+        # Files read back from a directory carry no derived marks or numbers.
         self._derived = set()
+        self._sequence_numbers = {}
         # Restore self._files from content of storage_dir.
         for filename in os.listdir(self.storage_dir):
             if not filename.endswith(".npy"):