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 @@
+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.

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 (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).

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

docs/usage/branches.md

@@ -0,0 +1,199 @@
+# Branches
+
+A branch is a copy of a simulation that calculates under different inputs or
+policy without changing the simulation it came from. Formulas use branches to
+compare alternatives (for example, tax liability if itemizing), and marginal
+tax rates use them to recalculate with slightly higher earnings.
+
+```python
+branch = simulation.get_branch("itemizing")
+branch.set_input("tax_unit_itemizes", 2026, itemizes)
+tax_if_itemizing = branch.calculate("income_tax", 2026)
+```
+
+## What a branch starts with
+
+The first `get_branch(name)` call returns a new branch holding every value its
+parent holds at that moment: inputs and calculated values alike. Later calls
+with the same name return that branch as it is. What the parent stores after
+the branch is created does not reach the branch, and what the branch stores
+does not reach the parent. A branch of a branch reads its own values first,
+then its parent's, then those of the simulation at the root.
+
+## Inputs set on a branch
+
+`set_input` on a branch stores the new value and drops every value the branch
+holds that may have been calculated from the value it replaces. What the branch
+calculates next therefore uses the input, as a simulation given the input
+before calculating anything would, whatever its parent had calculated before
+the branch was created.
+
+To decide what to drop, every stored value carries a sequence number from one
+counter shared by the whole process, so a larger number means a later store.
+Each simulation also keeps a store history: for each variable and period, the
+number of the first store its values may have been calculated from, and for
+each variable, the first time one of its values was uprated or carried over
+from another period. The history records:
+
+- the simulation's own stores, including values a holder calculates but does
+  not keep (`variables_to_drop`, the cache blacklist), values read from the
+  macro cache, and the default a spiral returns;
+- the number from which values may have been calculated from anything: the
+  first value read from the macro cache (what it was calculated from was
+  never calculated here) and the values restored from a dump (see below), so
+  an input for any variable drops them;
+- a copy of its parent's history, taken when the branch is created;
+- the history of any simulation its formulas calculate in, taken in each time
+  `calculate` there returns or raises (a formula that branches, sets an input
+  and calculates in the branch hands back values calculated there); a branch
+  calculated from a thread with no formula context hands its history to every
+  ancestor with a calculation running.
+
+When `set_input(variable, period, value)` is called on a branch, the branch
+drops each value it holds, other than an input, whose number is at least the
+earliest recorded store of the variable for a period that shares a day with
+`period`, or the earliest recorded uprated or carried-over value of the
+variable. It then forgets the records numbered from there on (its remaining
+values were all stored earlier) and records again the inputs it keeps, unless a
+formula is still running in the branch: such a formula may hold, in its own
+variables, a value it read before the drop, so the records stay. 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, 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.
+
+A custom `set_input` handler that calculates values between its own stores
+calculates them from inputs it has not yet replaced, so if it calculated
+anything (through `calculate`, `calculate_add` or `_calculate`), the branch
+drops again, by the same rule, once the handler returns or raises (the inputs
+it stored before raising stay).
+
+If there is no such record, the branch drops nothing. That is the case when a
+formula creates the branch while it is still calculating the variable the
+branch overrides, unless another branch the formula created for the same
+comparison already handed back a value calculated from that variable: then the
+branch drops what came back and what was calculated after it.
+
+A value written into a holder's storage 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 was stored or
+recorded, so a calculated value carries a larger number than each value it was
+calculated from, directly or through other calculated values. Every value a
+simulation holds was calculated there, inherited from its parent, handed back
+by another simulation's `calculate`, or restored from a dump. In the first
+three cases, for every variable it was calculated from, the simulation's
+history holds a record of that variable for an overlapping period numbered no
+later than the value read (for a value summed or divided from other periods,
+the record may be of those periods), unless it was calculated from a
+macro-cache read; restored values and values calculated from a macro-cache
+read are covered by the record that values from a number on may depend on
+anything. So any value that depends on the overridden variable at an
+overlapping period was stored after the earliest record that applies. Uprating
+and carry-over read which periods hold 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:
+
+- **Policy changes are not tracked.** A branch whose tax-benefit system or
+  parameters differ from its parent's still holds the parent's values
+  calculated under the parent's policy. Call `branch.drop_computed_arrays()`
+  after changing them.
+- **Formulas must not write into arrays they read.** A formula that changes a
+  cached array in place (`array += x`) changes a value after its sequence
+  number was assigned, so values calculated from it may be dropped too late or
+  not at all.
+- **Formulas should calculate, not inspect storage.** A formula that reads
+  `holder.get_known_periods()`, `simulation.get_array()` or another
+  simulation's storage directly, and acts on what it finds, depends on values
+  that are not recorded.
+- **Unrelated simulations in other threads.** A formula that calculates in a
+  simulation other than its own branches, from a thread it starts without
+  copying its context (`contextvars.copy_context`, which `asyncio.to_thread`
+  does), does not take in that simulation's history. 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.
+- **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.
+- **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),
+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`) 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 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. 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.
+
+## Dropping calculated values
+
+`simulation.drop_computed_arrays()` deletes every value the simulation holds
+except inputs (the dataset or situation it was built from, values set with
+`set_input` on it, and, for a branch, values set on the simulations it was
+created from before it was created), and returns how many arrays it deleted.
+Values 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)
+branch.tax_benefit_system.parameters.gov.some_rate.update(period="2026", value=0.2)
+branch.drop_computed_arrays()
+```

policyengine_core/data_storage/in_memory_storage.py

@@ -1,9 +1,13 @@
-from typing import Dict, Union
+from typing import Dict, 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 (
+    advance_sequence_past,
+    next_sequence_number,
+)
 from policyengine_core.periods import Period
 
 
@@ -47,8 +51,23 @@ def __init__(self, is_eternal: bool):
         # with a copy the first time it is read. A key left here after code
         # outside this class empties ``_arrays`` costs one extra copy at most.
         self._shared = set()
+        # 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 __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)
+        # 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.
 
@@ -83,6 +102,8 @@ def clone(self, share_arrays: bool = False) -> "InMemoryStorage":
                     clone._arrays[key] = array.copy()
         else:
             clone._arrays = {key: array.copy() for key, 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:
@@ -102,8 +123,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)
@@ -126,6 +158,54 @@ def put(
         key = f"{branch_name}:{period}"
         self._arrays[key] = value
         self._shared.discard(key)
+        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)
+            self._shared.discard(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)."""
+        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
+            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:
@@ -138,6 +218,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None:
                 if not period_item.startswith(branch_prefix)
             }
             self._shared.intersection_update(self._arrays)
+            self._forget_deleted_keys()
             return
 
         if self.is_eternal:
@@ -156,6 +237,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None:
             )
         }
         self._shared.intersection_update(self._arrays)
+        self._forget_deleted_keys()
 
     def get_known_periods(self) -> list:
         # Split on the first colon only: an anchored period's string form