Keep holders on their system's current definitions

changelog.d/holders-follow-system-definitions.fixed.md

@@ -0,0 +1 @@
+Holders now follow their system's current variable definitions whenever a simulation looks them up, so a reform can set an input for a variable it restores, and `apply_reform` keeps `input_variables` unless it drops a name.

policyengine_core/holders/holder.py

@@ -48,6 +48,27 @@ def __init__(self, variable: "Variable", population: "Population"):
             if self.variable.name in self.simulation.memory_config.variables_to_drop:
                 self._do_not_store = True
 
+    def _follow_definition(self, variable: "Variable") -> None:
+        """Take ``variable`` as this holder's definition when its storage fits.
+
+        A definition for the same entity, with the same ``ETERNITY`` storage,
+        replaces the holder's own, so reads and writes follow it: a
+        neutralized definition reads its default and ignores new inputs, and
+        a restored one stores them again. A change of entity or of
+        ``ETERNITY`` storage needs a new holder, so this leaves the holder as
+        it is; ``Simulation.apply_reform`` drops such holders when it binds
+        them to the reformed system.
+        """
+        if variable is self.variable:
+            return
+        if variable.entity.key != self.variable.entity.key:
+            return
+        if (
+            variable.definition_period == periods.ETERNITY
+        ) != self._memory_storage.is_eternal:
+            return
+        self.variable = variable
+
     def clone(self, population: "Population", share_arrays: bool = False) -> "Holder":
         """
         Copy the holder just enough to be able to run a new simulation without modifying the original simulation.

policyengine_core/populations/population.py

@@ -168,10 +168,15 @@ def __call__(
 
     def get_holder(self, variable_name: str) -> Holder:
         self.entity.check_variable_defined_for_entity(variable_name)
+        variable = self.entity.get_variable(variable_name)
         holder = self._holders.get(variable_name)
         if holder:
+            # The system may have replaced the definition since the holder
+            # was made or last bound: a reform's ``apply`` that restores a
+            # neutralized variable and then sets its input, or code that
+            # changes a branch's system directly.
+            holder._follow_definition(variable)
             return holder
-        variable = self.entity.get_variable(variable_name)
         self._holders[variable_name] = holder = Holder(variable, self)
         return holder
 

policyengine_core/simulations/simulation.py

@@ -919,6 +919,10 @@ def __init__(
             # variables, so a variable the reform neutralized read as the
             # default in the baseline too.
             self.baseline._bind_to_tax_benefit_system()
+            # The branch shares this simulation's list of input variables,
+            # and the binding keeps a list it drops nothing from. The
+            # baseline runs under another system, so it keeps its own list.
+            self.baseline.input_variables = list(self.baseline.input_variables)
         else:
             self.baseline = None
 
@@ -1324,7 +1328,11 @@ def _bind_to_tax_benefit_system(self) -> None:
         After an in-place reform, compatible holders keep their storage and
         recorded inputs. Each existing branch binds to its own system, so
         shared-system branches follow the reform and a baseline keeps its
-        baseline variables.
+        baseline variables. This replaces ``input_variables`` and the
+        recorded inputs only when it drops an entry, so a list that a clone or
+        branch shares stays shared; the code that builds a baseline branch
+        gives it a list of its own. Between reforms, ``Population.get_holder``
+        keeps compatible holders on the current definitions.
         """
         system = self.tax_benefit_system
         if system is None:
@@ -1359,12 +1367,19 @@ def kept(name):
                 or name in self.populations[variables[name].entity.key]._holders
             )
 
-        if getattr(self, "_user_input_keys", None) is not None:
-            self._user_input_keys = {
-                key for key in self._user_input_keys if kept(key[0])
-            }
-        if getattr(self, "input_variables", None) is not None:
-            self.input_variables = [name for name in self.input_variables if kept(name)]
+        # Replace a record only when it loses a name, and never edit it in
+        # place: a clone or branch may share its parent's, and callers may
+        # hold it.
+        user_input_keys = getattr(self, "_user_input_keys", None)
+        if user_input_keys is not None and not all(
+            kept(key[0]) for key in user_input_keys
+        ):
+            self._user_input_keys = {key for key in user_input_keys if kept(key[0])}
+        input_variables = getattr(self, "input_variables", None)
+        if input_variables is not None and not all(
+            kept(name) for name in input_variables
+        ):
+            self.input_variables = [name for name in input_variables if kept(name)]
         for branch in self.branches.values():
             branch._bind_to_tax_benefit_system()
 
@@ -3658,6 +3673,8 @@ def retain_structure(name, values):
             baseline.tracer = self.tracer
             baseline.tax_benefit_system = baseline_tax_benefit_system
             baseline._bind_to_tax_benefit_system()
+            # As in ``__init__``, the baseline keeps its own list.
+            baseline.input_variables = list(baseline.input_variables)
             baseline.baseline = None
             if getattr(self, "baseline", None) is not None:
                 self.baseline = baseline

tests/core/test_apply_reform_holders.py

@@ -1,5 +1,7 @@
 """Reforms refresh existing holders without discarding compatible inputs."""
 
+import warnings
+
 import numpy as np
 import pytest
 
@@ -11,6 +13,7 @@
 from tests.fixtures.reform_holders import (
     Household,
     Person,
+    award,
     build,
     change_default,
     store_on_disk,
@@ -84,6 +87,120 @@ def test_neutralized_variable_keeps_input_storage_but_reads_its_default(on_disk)
     np.testing.assert_array_equal(simulation.calculate("award", PERIOD), [0, 0])
 
 
+class RestoreAward(Reform):
+    def apply(self):
+        self.replace_variable(award)
+
+
+def restore_and_set_award(values):
+    """Return a reform that restores the award, then sets and reads its input."""
+
+    class RestoreAndSetAward(Reform):
+        def apply(self):
+            self.replace_variable(award)
+            self.simulation.set_input("award", PERIOD, values)
+            np.testing.assert_array_equal(
+                self.simulation.calculate("award", PERIOD), values
+            )
+
+    return RestoreAndSetAward
+
+
+def neutralize_and_set_award(values):
+    """Return a reform that neutralizes the award, then sets its input."""
+
+    class NeutralizeAndSetAward(Reform):
+        def apply(self):
+            self.neutralize_variable("award")
+            self.simulation.set_input("award", PERIOD, values)
+
+    return NeutralizeAndSetAward
+
+
+def neutralized_input_warnings(caught):
+    return [warning for warning in caught if "neutralized" in str(warning.message)]
+
+
+@pytest.mark.parametrize("on_disk", [False, True])
+def test_reform_keeps_an_input_it_sets_after_restoring_a_neutralized_variable(
+    on_disk,
+):
+    simulation = build()
+    if on_disk:
+        store_on_disk(simulation, "award")
+    simulation.set_input("award", PERIOD, [7, 9])
+    simulation.apply_reform(NeutralizeAward)
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        simulation.apply_reform(restore_and_set_award([20, 30]))
+
+    assert not neutralized_input_warnings(caught)
+    assert_bound(simulation)
+    holder = simulation.get_holder("award")
+    np.testing.assert_array_equal(holder._get_array_from_storage(PERIOD), [20, 30])
+    np.testing.assert_array_equal(simulation.calculate("award", PERIOD), [20, 30])
+
+
+@pytest.mark.parametrize("on_disk", [False, True])
+def test_reform_ignores_an_input_it_sets_after_neutralizing_the_variable(on_disk):
+    # A holder made after the neutralization ignores the input, so one made
+    # before it does too.
+    simulation = build()
+    if on_disk:
+        store_on_disk(simulation, "award")
+    simulation.set_input("award", PERIOD, [7, 9])
+
+    with pytest.warns(Warning, match="neutralized"):
+        simulation.apply_reform(neutralize_and_set_award([1, 2]))
+
+    np.testing.assert_array_equal(simulation.calculate("award", PERIOD), [0, 0])
+    simulation.apply_reform(RestoreAward)
+    np.testing.assert_array_equal(simulation.calculate("award", PERIOD), [7, 9])
+
+
+def test_input_set_while_a_variable_is_neutralized_is_ignored():
+    simulation = build()
+    simulation.set_input("award", PERIOD, [7, 9])
+    simulation.apply_reform(NeutralizeAward)
+
+    with pytest.warns(Warning, match="neutralized"):
+        simulation.set_input("award", PERIOD, [55, 55])
+
+    simulation.apply_reform(RestoreAward)
+    np.testing.assert_array_equal(simulation.calculate("award", PERIOD), [7, 9])
+
+
+def test_holders_follow_a_branch_system_changed_outside_a_reform():
+    simulation = build()
+    simulation.set_input("award", PERIOD, [7, 9])
+    branch = simulation.get_branch("measured", clone_system=True)
+
+    branch.tax_benefit_system.neutralize_variable("award")
+
+    assert branch.get_holder("award").variable.is_neutralized
+    assert not simulation.get_holder("award").variable.is_neutralized
+    np.testing.assert_array_equal(branch.calculate("award", PERIOD), [0, 0])
+    with pytest.warns(Warning, match="neutralized"):
+        branch.set_input("award", PERIOD, [1, 1])
+    np.testing.assert_array_equal(simulation.calculate("award", PERIOD), [7, 9])
+
+
+def test_reform_that_drops_no_input_keeps_each_input_list():
+    simulation = build()
+    simulation.set_input("income", PERIOD, [7, 9])
+    names = simulation.input_variables = ["income", "eligible"]
+    branch = simulation.get_branch("other")
+    branch_names = branch.input_variables
+
+    simulation.apply_reform(change_default(20))
+
+    # Callers may hold either list, and a branch may share its parent's.
+    assert simulation.input_variables is names
+    assert branch.input_variables is branch_names
+    assert names == ["income", "eligible"]
+
+
 class added(Variable):
     value_type = float
     entity = Person
@@ -106,11 +223,15 @@ def test_added_variable_is_lazy_and_removed_variable_drops_holder_and_inputs(on_
     simulation.set_input("annual_value", PERIOD, [7, 9])
     simulation.input_variables = ["annual_value", "eligible"]
 
+    names = simulation.input_variables
+
     simulation.apply_reform(AddAndRemove)
 
     assert_bound(simulation)
     assert "annual_value" not in simulation.persons._holders
     assert "annual_value" not in simulation.input_variables
+    # The list may be shared, so the reform replaces it instead of editing it.
+    assert names == ["annual_value", "eligible"]
     assert all(key[0] != "annual_value" for key in simulation._user_input_keys)
     with pytest.raises(VariableNotFoundError):
         simulation.get_holder("annual_value")

tests/core/test_apply_reform_holders_property.py

@@ -1,5 +1,7 @@
 """Check holder identity, defaults and recorded inputs after a reform."""
 
+import warnings
+
 import numpy as np
 import pytest
 
@@ -9,7 +11,13 @@
 from policyengine_core import periods
 from policyengine_core.reforms import Reform
 from policyengine_core.variables import Variable
-from tests.fixtures.reform_holders import YEAR, build, change_default, store_on_disk
+from tests.fixtures.reform_holders import (
+    YEAR,
+    award,
+    build,
+    change_default,
+    store_on_disk,
+)
 
 pytestmark = pytest.mark.filterwarnings(
     "ignore:Memory configuration:policyengine_core.warnings.MemoryConfigWarning"
@@ -224,3 +232,75 @@ def test_applying_same_reform_object_twice_returns_same_values(mask, default, wa
     np.testing.assert_array_equal(
         twice.calculate("award", YEAR), once.calculate("award", YEAR)
     )
+
+
+AWARD_VALUES = st.lists(
+    st.integers(min_value=-1_000, max_value=1_000), min_size=2, max_size=2
+)
+NEUTRALIZATION_STEPS = st.lists(
+    st.one_of(
+        st.tuples(st.sampled_from(("neutralize", "restore")), st.none()),
+        st.tuples(
+            st.sampled_from(("write", "neutralize_then_write", "restore_then_write")),
+            AWARD_VALUES,
+        ),
+    ),
+    min_size=1,
+    max_size=8,
+)
+
+
+def award_reform(neutralize, values=None):
+    """Return a reform that neutralizes or restores the award, then may set it."""
+
+    class ChangeAward(Reform):
+        def apply(self):
+            if neutralize:
+                self.neutralize_variable("award")
+            else:
+                self.replace_variable(award)
+            if values is not None:
+                self.simulation.set_input("award", YEAR, values)
+
+    return ChangeAward
+
+
+@PROPERTY_SETTINGS
+@hypothesis.given(steps=NEUTRALIZATION_STEPS, on_disk=st.booleans())
+@hypothesis.example(
+    steps=[("neutralize", None), ("restore_then_write", [20, 30])], on_disk=False
+)
+@hypothesis.example(
+    steps=[("neutralize_then_write", [1, 2]), ("restore", None)], on_disk=True
+)
+@hypothesis.example(
+    steps=[("neutralize", None), ("write", [55, 55]), ("restore", None)],
+    on_disk=False,
+)
+def test_neutralized_award_reads_its_last_input_set_while_defined(steps, on_disk):
+    """A model of neutralization: an input counts only if set while defined.
+
+    While the award is neutralized it reads its default, 0, and ignores new
+    inputs. Restored, it reads the last input set while it was not
+    neutralized, whether a reform set it inside ``apply`` or code set it
+    outside one.
+    """
+    simulation = build()
+    if on_disk:
+        store_on_disk(simulation, "award")
+    simulation.set_input("award", YEAR, [7, 9])
+    neutralized, last_input = False, [7, 9]
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        for action, values in steps:
+            if action == "write":
+                simulation.set_input("award", YEAR, values)
+            else:
+                neutralize = action.startswith("neutralize")
+                simulation.apply_reform(award_reform(neutralize, values))
+                neutralized = neutralize
+            if values is not None and not neutralized:
+                last_input = values
+            expected = [0, 0] if neutralized else last_input
+            np.testing.assert_array_equal(simulation.calculate("award", YEAR), expected)
+            assert_holder_identity(simulation)

tests/core/test_reform_baseline_system.py

@@ -213,6 +213,15 @@ def test_the_baselines_record_of_inputs_is_its_own():
     assert "doubled_salary" in reform.tax_benefit_system.variables
 
 
+def test_subsample_gives_the_rebuilt_baseline_its_own_record_of_inputs():
+    simulation = Microsimulation(reform=neutralize_income_tax)
+
+    simulation.subsample(n=1, seed="baseline-system", time_period="2022")
+
+    assert simulation.baseline.input_variables is not simulation.input_variables
+    assert simulation.baseline._user_input_keys is not simulation._user_input_keys
+
+
 def test_subsample_leaves_the_baseline_without_a_baseline():
     simulation = Microsimulation(reform=neutralize_income_tax)
     assert simulation.baseline.baseline is None