From 14837f784b52d731101dd2e322c7727895d08753 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 6 Nov 2026 21:52:04 -0500 Subject: [PATCH] Keep holders on their system's current definitions Population.get_holder now gives an existing holder its system's current variable definition when the definition fits the holder's storage, so a reform that restores a neutralized variable can set its input inside apply. Simulation._bind_to_tax_benefit_system replaces input_variables and the input records only when it drops an entry, and the code that builds a reform simulation's baseline gives the baseline its own list. Co-Authored-By: Claude Opus 5.5 --- ...holders-follow-system-definitions.fixed.md | 1 + policyengine_core/holders/holder.py | 21 +++ policyengine_core/populations/population.py | 7 +- policyengine_core/simulations/simulation.py | 31 ++++- tests/core/test_apply_reform_holders.py | 121 ++++++++++++++++++ .../test_apply_reform_holders_property.py | 82 +++++++++++- tests/core/test_reform_baseline_system.py | 9 ++ 7 files changed, 263 insertions(+), 9 deletions(-) create mode 100644 changelog.d/holders-follow-system-definitions.fixed.md diff --git a/changelog.d/holders-follow-system-definitions.fixed.md b/changelog.d/holders-follow-system-definitions.fixed.md new file mode 100644 index 000000000..73cc56aa4 --- /dev/null +++ b/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. diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 4b3c117fb..7e0b7c4fd 100644 --- a/policyengine_core/holders/holder.py +++ b/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. diff --git a/policyengine_core/populations/population.py b/policyengine_core/populations/population.py index b4c99123f..63488ac78 100644 --- a/policyengine_core/populations/population.py +++ b/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 diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 9b1ffc3a4..4c8e6d51d 100644 --- a/policyengine_core/simulations/simulation.py +++ b/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 diff --git a/tests/core/test_apply_reform_holders.py b/tests/core/test_apply_reform_holders.py index 888d24a69..04861ce73 100644 --- a/tests/core/test_apply_reform_holders.py +++ b/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") diff --git a/tests/core/test_apply_reform_holders_property.py b/tests/core/test_apply_reform_holders_property.py index be7b498e6..d787346dc 100644 --- a/tests/core/test_apply_reform_holders_property.py +++ b/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) diff --git a/tests/core/test_reform_baseline_system.py b/tests/core/test_reform_baseline_system.py index 9b1fde11f..519223374 100644 --- a/tests/core/test_reform_baseline_system.py +++ b/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