Preserve modifier inputs across direct callbacks and measurements

changelog.d/fix-scenario-modifier-inputs-2666.fixed.md

@@ -0,0 +1,3 @@
+Retain every surviving explicit input through `reset_calculations`, including ordinary `set_input` calls made after construction outside any modifier or scenario. A reset alone no longer undoes those overrides: use `delete_arrays(variable, period)`, then reset or recalculate to restore the variable's formula or default. Retain these inputs through dynamics, marginal-rate calculations and exports, including input-only variables absent at construction.
+
+Refresh dependent values before the next outermost read after every top-level `set_input` on a built simulation, even without a modifier. Preserve selective invalidation for internal formula and framework writes and keep independent baselines. Use collision-safe temporary branches with cleanup on failure for all three marginal rates, and avoid duplicate cleanup after a successful construction-time modifier. Require PolicyEngine Core 3.33.0 or later for the documented public `drop_computed_arrays` method.

docs/book/usage/scenarios.md

@@ -159,6 +159,55 @@ A sum keeps the timing of each part's modifier. A simulation runs the modifiers
 
 Beyond simple parameter changes, you sometimes need scenarios that make structural modifications to how the tax-benefit system works. Let's walk through how to create these, using examples similar to those in PolicyEngine UK's scenarios module.
 
+### Setting inputs in a simulation modifier
+
+When you pass a scenario to `Simulation` or `Microsimulation`, its modifier runs
+after data load by default. Values it stores with `set_input` become inputs in
+`simulation.input_variables`. They survive `reset_calculations`, marginal-rate
+branches and `apply_dynamics`, and dataset exports include them as inputs.
+The simulation drops calculated values after the modifier finishes, so later
+calculations use the inputs the modifier left behind. Calculating a variable
+inside the modifier does not make that variable an input.
+Employment marginal-rate calculations also take the completed scenario's
+inputs when the modifier calculated a marginal rate before changing an input.
+
+An input remains fixed when a marginal-rate calculation changes earnings. A
+modifier that adds a fixed amount to Child Benefit therefore keeps that amount
+through the earnings step. For a policy that should vary with earnings, use
+parameter changes or a formula reform; setting its calculated value once cannot
+represent that dependence.
+
+This input retention also applies to ordinary `set_input` calls on an
+already-built simulation, including calls outside any modifier or scenario.
+`reset_calculations` retains every surviving explicit input visible to that
+simulation, including inputs absent at construction, and all three marginal-rate
+calculations retain them too. To undo an override, call
+`simulation.delete_arrays(variable, period)`, then reset or recalculate so the
+variable's formula or default runs again. A reset alone no longer undoes an
+input supplied after construction.
+
+Every top-level `set_input` on a built simulation queues dependent-cache cleanup
+before its next outermost read. This includes simulations with no modifier: a
+caller can read Child Benefit, add an amount and then obtain fresh income and
+response values. Inputs stored only on another clone or branch do not become
+the parent's inputs. Formula and framework writes during a calculation keep
+their existing selective invalidation.
+
+Direct edits of the parameter tree require a separate array-cache boundary.
+After calling a parameter's `update` method on a built simulation, call
+`simulation.drop_computed_arrays()` before reading its results. This public
+method requires PolicyEngine Core 3.33.0 or later, the model's minimum version.
+In a wrapper
+that runs a policy callback, late dynamic parameter updates and a dynamic
+callback, clear computed arrays after each operation in that order. This
+preserves the input values and lets the next operation see the updated policy.
+A wrapper must also use a country-model release containing the native input
+retention fix; array cleanup alone does not add it to an older release.
+
+Modifiers with `applied_before_data_load=True` run before the simulation has its
+data. They can change parameters or model structure at that stage; modifiers
+that read or set household inputs belong after data load.
+
 ### Creating a two-child limit repeal scenario
 
 The two-child limit restricts Universal Credit child elements to the first two children. Here's how to build a scenario that removes this restriction:

policyengine_uk/simulation.py

@@ -2,6 +2,7 @@
 import os
 import copy
 import warnings
+from contextlib import contextmanager, nullcontext
 from importlib.metadata import version
 from typing import Dict, Optional, Union, Type, List
 
@@ -57,6 +58,10 @@
 from policyengine_uk.utils.dependencies import get_variable_dependencies
 from policyengine_uk.utils.data_source import built_from_data
 from policyengine_uk.utils.input_dataframe_export import input_dataframe
+from policyengine_uk.utils.measurement_branches import (
+    drop_calculated_values,
+    retained_input_variables,
+)
 from policyengine_uk.utils.supplied_inputs import (
     SUPPLIED_INPUT_VARIABLES,
     copy_model_values,
@@ -113,7 +118,7 @@ def _canonical_parameter_paths(changes: dict) -> dict:
 
 
 class _InputKeyLog(set):
-    """A simulation's record of its inputs while Simulation.apply_reform runs.
+    """A simulation's record of inputs stored by a reform or scenario modifier.
 
     policyengine-core records each input in the simulation's
     ``_user_input_keys`` as (variable, branch, period). ``Holder._set`` adds
@@ -124,10 +129,10 @@ class _InputKeyLog(set):
     variable's end, one for a neutralized variable, or one for a sub-period
     that a ``set_input`` helper skips because it already holds an input.
     This set also adds each key to every set in ``logs``: one for each
-    reform being applied to the simulation, the innermost last.
+    operation being applied to the simulation, the innermost last.
 
     A copy or pickle of it is a plain set, so a simulation deep-copied or
-    pickled while a reform runs logs nothing. A shallow copy of the
+    pickled while an operation runs logs nothing. A shallow copy of the
     simulation shares its record, as it shares its holders.
     """
 
@@ -144,11 +149,11 @@ def __reduce__(self):
         return set, (list(self),)
 
 
-# The logs (see _InputKeyLog) of the reforms Simulation.apply_reform is
-# applying, by the id of the simulation it applies them to, while it does.
+# The logs (see _InputKeyLog) of the reforms or scenario modifiers running,
+# by the id of the simulation they apply to, while they do.
 # Kept outside the simulation, because core's clone copies its __dict__:
 # a clone or branch made while a reform runs must not log for it.
-_reform_logs: Dict[int, list] = {}
+_input_key_logs: Dict[int, list] = {}
 
 
 # Cache for fully-loaded, uprated, enum-pre-encoded multi-year datasets keyed
@@ -284,6 +289,7 @@ def __init__(
         self._cgt_input_provenance = {}
         self._cgt_defaulted_periods = set()
         self._cgt_calculation_depth = 0
+        self._external_input_change_pending = False
         self._cgt_operation_inputs = set()
         self._cgt_warned_unused = False
         self._cgt_warning_enabled = False
@@ -421,6 +427,10 @@ def __init__(
                 capital_gains_projection=capital_gains_projection,
             )
         else:
+            # Finish framework input invalidation before copying the baseline.
+            if self._external_input_change_pending:
+                drop_calculated_values(self)
+                self._external_input_change_pending = False
             self.baseline = self.clone()
         self.baseline._cgt_warning_enabled = False
         self.baseline._defer_uc_health_warning = False
@@ -433,7 +443,10 @@ def __init__(
             if self._construction_reform_state is not None:
                 self._construction_reform_state.refresh_framework(self)
             if after_data_load.simulation_modifier is not None:
-                after_data_load.simulation_modifier(self)
+                with self._preserve_stored_inputs():
+                    after_data_load.simulation_modifier(self)
+                drop_calculated_values(self)
+                self._external_input_change_pending = False
         self._refresh_capital_gains_projection()
         self._capture_cgt_supplied_inputs()
         self._cgt_warning_enabled = True
@@ -456,49 +469,32 @@ def _refresh_capital_gains_projection(self) -> None:
         add_capital_gains_projection(self.tax_benefit_system.parameters, projection)
         self.tax_benefit_system.reset_parameter_caches()
 
-    def apply_reform(self, reform: dict | tuple | Type[Reform]) -> None:
-        """Apply a runtime reform while keeping generated CGT assumptions coherent."""
-        # Reforms that set inputs (disable_simulated_benefits, adjust_budgets)
-        # reach the simulation through its tax-benefit system's link. A clone
-        # or branch may share its parent's system, so link the system to this
-        # simulation while the reform is applied, then restore the link.
-        system = self.tax_benefit_system
-        linked = getattr(system, "simulation", None)
-        system.simulation = self
-        # Log the inputs the reform stores (see _InputKeyLog and the
-        # _user_input_keys property). Within another reform applied to this
-        # simulation (a tuple's parts, or a reform that applies one), log
-        # them for that one too.
-        logs = _reform_logs.setdefault(id(self), [])
+    @contextmanager
+    def _record_stored_inputs(self):
+        """Record actual stores, including overwrites and custom input routes."""
+        logs = _input_key_logs.setdefault(id(self), [])
         stored = set()
         logs.append(stored)
         self._user_input_keys = getattr(self, "_user_input_keys", None) or set()
         try:
-            super().apply_reform(reform)
-            self._refresh_capital_gains_projection()
+            yield stored
         finally:
             logs.pop()
             if not logs:
-                del _reform_logs[id(self)]
-            # A plain set again, unless an enclosing reform is still logging.
+                del _input_key_logs[id(self)]
+            # A plain set again, unless an enclosing operation is still logging.
             self._user_input_keys = set(self._user_input_keys)
-            if linked is not None:
-                system.simulation = linked
-        # input_variables is taken when the simulation is built.
-        # reset_calculations deletes the values of every variable it does not
-        # name (as core's derivative did before policyengine-core 3.33.0, from
-        # which it keeps every input), and the marginal rate formulas those of
-        # every such variable that has a formula. A reform applied afterwards
-        # (a Reform class passed as reform=, or one applied to a clone or
-        # branch) can store inputs: new ones, or ones that replace an input
-        # set since the build. Name every variable with an input the reform
-        # stored that this simulation still holds, under a branch name it
-        # reads and a reset deletes, or a reset would delete what the reform
-        # stored and recalculate the formula. input_variables names whole
-        # variables, so a reset then keeps every input of the variable, any
-        # set since the build included. A reform that raises names nothing it
-        # stored itself (a reform it applied first has named its own). Assign
-        # a new list: a clone or branch shares its parent's.
+
+    @contextmanager
+    def _preserve_stored_inputs(self):
+        """Name inputs an operation stores and retains, for later resets."""
+        with self._record_stored_inputs() as stored:
+            yield
+        # Preserve the existing construction/reform input-name contract.
+        # Later direct stores are also retained by retained_input_variables;
+        # calculated values do not become inputs. An operation that raises
+        # names nothing it stored itself (a reform it applied first has named
+        # its own). Assign a new list: a clone or branch shares its parent's.
         input_variables = getattr(self, "input_variables", None)
         if input_variables is not None:
             branches = set(self._get_visible_branch_names())
@@ -511,13 +507,30 @@ def apply_reform(self, reform: dict | tuple | Type[Reform]) -> None:
             if added:
                 self.input_variables = list(input_variables) + added
 
+    def apply_reform(self, reform: dict | tuple | Type[Reform]) -> None:
+        """Apply a runtime reform while keeping generated CGT assumptions coherent."""
+        # Reforms that set inputs (disable_simulated_benefits, adjust_budgets)
+        # reach the simulation through its tax-benefit system's link. A clone
+        # or branch may share its parent's system, so link the system to this
+        # simulation while the reform is applied, then restore the link.
+        system = self.tax_benefit_system
+        linked = getattr(system, "simulation", None)
+        system.simulation = self
+        try:
+            with self._preserve_stored_inputs():
+                super().apply_reform(reform)
+                self._refresh_capital_gains_projection()
+        finally:
+            if linked is not None:
+                system.simulation = linked
+
     @property
     def _user_input_keys(self) -> set:
         # policyengine-core's record of the simulation's inputs, kept in
-        # __dict__ as core keeps it. While apply_reform applies a reform to
+        # __dict__ as core keeps it. While a reform or scenario modifier runs on
         # this simulation, a set assigned here becomes a log (see
         # _InputKeyLog), so a record core replaces mid-reform still logs: a
-        # branch's, when a reform applied to its parent rebuilds it, or one
+        # branch's, when an operation on its parent rebuilds it, or one
         # rebuilt when the system is rebound.
         try:
             return self.__dict__["_user_input_keys"]
@@ -526,7 +539,7 @@ def _user_input_keys(self) -> set:
 
     @_user_input_keys.setter
     def _user_input_keys(self, keys) -> None:
-        logs = _reform_logs.get(id(self))
+        logs = _input_key_logs.get(id(self))
         if logs and keys is not None:
             keys = _InputKeyLog(keys, logs)
         self.__dict__["_user_input_keys"] = keys
@@ -571,10 +584,31 @@ def set_input(self, variable_name: str, period: Period, value) -> None:
         # stores a value for, and only those. set_model_value records its
         # value as the model's after this call.
         records = model_value_input_records(self, variable_name)
+        # This depth covers every calculate call, including nongains
+        # variables. A formula can write to a new branch whose own depth is
+        # zero while an ancestor is calculating. Core already invalidates
+        # that branch's affected dependencies; keep its unrelated caches.
+        # Direct edits with every ancestor idle must refresh dependent values
+        # before the next outermost calculation.
+        calculation = self
+        while calculation is not None and not getattr(
+            calculation, "_cgt_calculation_depth", 0
+        ):
+            calculation = getattr(calculation, "parent_branch", None)
+        external = (
+            getattr(self, "input_variables", None) is not None and calculation is None
+        )
+        stored = set()
         try:
-            super().set_input(variable_name, period, value)
+            recording = (
+                self._record_stored_inputs() if external else nullcontext(stored)
+            )
+            with recording as stored:
+                super().set_input(variable_name, period, value)
         finally:
             forget_replaced_model_values(self, variable_name, records)
+            if external and stored:
+                self._external_input_change_pending = True
         # Inputs do not carry back, so give the simulation its identifiers
         # for the year the input starts in if it holds them only for later
         # years. See ``policyengine_uk.data.input_dataframe.hold_identifiers_in``.
@@ -844,8 +878,10 @@ def _dataset_cache_key(self, source: str):
         )
 
     def reset_calculations(self):
+        retained = retained_input_variables(self)
+        self._external_input_change_pending = False
         for variable in self.tax_benefit_system.variables:
-            if variable not in self.input_variables:
+            if variable not in retained:
                 self.delete_arrays(variable)
 
     def get_known_variables(self):
@@ -1412,6 +1448,9 @@ def calculate(
         map_to: str = None,
         decode_enums: bool = False,
     ):
+        if self._cgt_calculation_depth == 0 and self._external_input_change_pending:
+            self._external_input_change_pending = False
+            drop_calculated_values(self)
         if variable_name is None:
             return self.calculate_all()
         tracer: SimpleTracer = self.tracer

policyengine_uk/tests/code_health/test_cached_arrays_not_written_in_place.py

@@ -550,17 +550,33 @@ def _earning_couple_with_gains() -> dict:
         "capital_gains_default_response",
     ],
 )
-def test_branching_runs_with_read_only_cache(read_only_cache, case):
+def test_branching_runs_with_read_only_cache(read_only_cache, case, monkeypatch):
     from policyengine_uk import Simulation
     from policyengine_uk.model_api import Scenario
 
     if case == "marginal_rates":
         sim = Simulation(situation=_earning_couple_with_gains())
+        owned = []
+        original_get_branch = Simulation.get_branch
+
+        def capture(owner, name="branch", *args, **kwargs):
+            branch = original_get_branch(owner, name, *args, **kwargs)
+            if owner is sim and name.startswith(
+                ("adult_1_pay_rise", "adult_2_pay_rise")
+            ):
+                owned.append((name, branch))
+            return branch
+
+        monkeypatch.setattr(Simulation, "get_branch", capture)
         for year in YEARS:
             sim.calculate("marginal_tax_rate", year)
             sim.calculate("marginal_tax_rate_on_capital_gains", year)
-        # The branches ran: both earners face a rate, and so do the gains.
-        assert {"adult_1_pay_rise", "adult_2_pay_rise"} <= set(sim.branches)
+        # Both scratch branches ran, and each registration was detached.
+        assert {"adult_1_pay_rise", "adult_2_pay_rise"} <= {
+            name for name, branch in owned
+        }
+        assert all(sim.branches.get(name) is not branch for name, branch in owned)
+        # Both earners face a rate, and so do the gains.
         assert (sim.calculate("marginal_tax_rate", 2026)[:2] > 0).all()
         assert sim.calculate("marginal_tax_rate_on_capital_gains", 2026)[0] > 0
         return