diff --git a/changelog.d/rebind-holders-after-reform.fixed.md b/changelog.d/rebind-holders-after-reform.fixed.md new file mode 100644 index 000000000..1f4da5a92 --- /dev/null +++ b/changelog.d/rebind-holders-after-reform.fixed.md @@ -0,0 +1 @@ +Simulation.apply_reform now rebinds existing holders to each branch's system before clearing cached calculations. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index de5824c59..9b1ffc3a4 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -932,6 +932,8 @@ def apply_reform(self, reform: Union[tuple, Reform]): if isinstance(reform, dict): reform = Reform.from_dict(reform) reform.apply(self.tax_benefit_system) + # Refresh holders before input replay looks them up in the new system. + self._bind_to_tax_benefit_system() # Invalidate every cached value so the simulation recomputes under # the new tax-benefit system. Previously ``apply_reform`` left # ``_fast_cache``, in-memory holder storage, and on-disk holder @@ -1318,6 +1320,11 @@ def _bind_to_tax_benefit_system(self) -> None: holder was built for the reform's definition in a way a new variable cannot change: in another entity's population, or with storage for (or not for) an ``ETERNITY`` variable. + + 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. """ system = self.tax_benefit_system if system is None: @@ -1346,7 +1353,11 @@ def _bind_to_tax_benefit_system(self) -> None: holder.variable = variable def kept(name): - return name in variables and name not in dropped + # A reform may supply a new entity's input before we drop the old holder. + return name in variables and ( + name not in dropped + or name in self.populations[variables[name].entity.key]._holders + ) if getattr(self, "_user_input_keys", None) is not None: self._user_input_keys = { @@ -1354,6 +1365,8 @@ def kept(name): } if getattr(self, "input_variables", None) is not None: self.input_variables = [name for name in self.input_variables if kept(name)] + for branch in self.branches.values(): + branch._bind_to_tax_benefit_system() def create_shortcuts(self) -> None: for _key, population in self.populations.items(): diff --git a/tests/core/test_apply_reform_holders.py b/tests/core/test_apply_reform_holders.py new file mode 100644 index 000000000..888d24a69 --- /dev/null +++ b/tests/core/test_apply_reform_holders.py @@ -0,0 +1,295 @@ +"""Reforms refresh existing holders without discarding compatible inputs.""" + +import numpy as np +import pytest + +from policyengine_core.errors import VariableNotFoundError +from policyengine_core.parameters import ParameterNode +from policyengine_core.periods import ETERNITY, YEAR, period +from policyengine_core.reforms import Reform +from policyengine_core.variables import Variable +from tests.fixtures.reform_holders import ( + Household, + Person, + build, + change_default, + store_on_disk, +) + +PERIOD = "2025" + + +def assert_bound(simulation): + for population in simulation.populations.values(): + for name, holder in population._holders.items(): + assert holder.variable is simulation.tax_benefit_system.variables[name] + assert holder.variable.entity.key == population.entity.key + assert holder._memory_storage.is_eternal == ( + holder.variable.definition_period == ETERNITY + ) + + +def test_defined_for_default_does_not_depend_on_other_people(): + results = [] + for mask in ([False], [False, True]): + simulation = build(mask) + old_holder = simulation.get_holder("award") + simulation.apply_reform(change_default(20)) + results.append(simulation.calculate("award", PERIOD).tolist()) + assert simulation.get_holder("award") is old_holder + + assert results == [[20], [20, 100]] + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_replaced_variable_keeps_recorded_inputs(on_disk): + simulation = build() + if on_disk: + store_on_disk(simulation, "award") + simulation.set_input("award", PERIOD, [7, 9]) + holder = simulation.get_holder("award") + records = simulation._user_input_keys.copy() + simulation.input_variables = ["award", "eligible"] + + simulation.apply_reform(change_default(20)) + + assert simulation.get_holder("award") is holder + assert_bound(simulation) + assert simulation._user_input_keys == records + assert simulation.input_variables == ["award", "eligible"] + np.testing.assert_array_equal(simulation.calculate("award", PERIOD), [7, 9]) + np.testing.assert_array_equal(holder._get_array_from_storage(PERIOD), [7, 9]) + + +class NeutralizeAward(Reform): + def apply(self): + self.neutralize_variable("award") + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_neutralized_variable_keeps_input_storage_but_reads_its_default(on_disk): + simulation = build() + if on_disk: + store_on_disk(simulation, "award") + simulation.set_input("award", PERIOD, [7, 9]) + records = simulation._user_input_keys.copy() + + simulation.apply_reform(NeutralizeAward) + + holder = simulation.get_holder("award") + assert_bound(simulation) + assert holder.variable.is_neutralized + assert simulation._user_input_keys == records + np.testing.assert_array_equal(holder._get_array_from_storage(PERIOD), [7, 9]) + np.testing.assert_array_equal(simulation.calculate("award", PERIOD), [0, 0]) + + +class added(Variable): + value_type = float + entity = Person + definition_period = YEAR + default_value = 13 + label = "Added amount" + + +class AddAndRemove(Reform): + def apply(self): + self.add_variable(added) + del self.variables["annual_value"] + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_added_variable_is_lazy_and_removed_variable_drops_holder_and_inputs(on_disk): + simulation = build() + if on_disk: + store_on_disk(simulation, "annual_value") + simulation.set_input("annual_value", PERIOD, [7, 9]) + simulation.input_variables = ["annual_value", "eligible"] + + simulation.apply_reform(AddAndRemove) + + assert_bound(simulation) + assert "annual_value" not in simulation.persons._holders + assert "annual_value" not in simulation.input_variables + assert all(key[0] != "annual_value" for key in simulation._user_input_keys) + with pytest.raises(VariableNotFoundError): + simulation.get_holder("annual_value") + assert "added" not in simulation.persons._holders + np.testing.assert_array_equal(simulation.calculate("added", PERIOD), [13, 13]) + assert_bound(simulation) + + +def test_tuple_reform_rebinds_after_each_reform(): + class CheckFirstReform(Reform): + def apply(self): + assert ( + self.simulation.get_holder("award").variable is self.variables["award"] + ) + change_default(30).apply(self) + + simulation = build([False]) + simulation.apply_reform((change_default(20), CheckFirstReform)) + + assert_bound(simulation) + np.testing.assert_array_equal(simulation.calculate("award", PERIOD), [30]) + + +def test_dict_reform_keeps_inputs_and_rebinds_stale_holders(): + simulation = build([False]) + simulation.tax_benefit_system.parameters = ParameterNode( + "", data={"amount": {"values": {"2025-01-01": {"value": 1}}}} + ) + simulation.set_input("income", PERIOD, [7]) + change_default(20).apply(simulation.tax_benefit_system) + + simulation.apply_reform({"amount": {"2025": 2}}) + + assert_bound(simulation) + assert simulation.tax_benefit_system.parameters("2025-01-01").amount == 2 + np.testing.assert_array_equal(simulation.calculate("income", PERIOD), [7]) + np.testing.assert_array_equal(simulation.calculate("award", PERIOD), [20]) + + +def test_existing_baseline_keeps_its_own_system(): + simulation = build([False]) + baseline = simulation.get_branch("baseline", clone_system=True) + simulation.baseline = baseline + baseline_system = baseline.tax_benefit_system + np.testing.assert_array_equal(baseline.calculate("award", PERIOD), [0]) + + simulation.apply_reform(change_default(20)) + + assert simulation.baseline is baseline + assert baseline.tax_benefit_system is baseline_system + assert_bound(simulation) + assert_bound(baseline) + np.testing.assert_array_equal(baseline.calculate("award", PERIOD), [0]) + np.testing.assert_array_equal(simulation.calculate("award", PERIOD), [20]) + + +def test_existing_branches_rebind_to_their_own_systems(): + simulation = build([False]) + shared = simulation.get_branch("shared") + nested = shared.get_branch("nested") + independent = simulation.get_branch("independent", clone_system=True) + independent.apply_reform(change_default(40)) + for branch in (simulation, shared, nested, independent): + branch.set_input("income", PERIOD, [7]) + branch.calculate("award", PERIOD) + + simulation.apply_reform(change_default(20)) + + for branch, expected in ( + (simulation, 20), + (shared, 20), + (nested, 20), + (independent, 40), + ): + assert_bound(branch) + np.testing.assert_array_equal(branch.calculate("award", PERIOD), [expected]) + np.testing.assert_array_equal(branch.calculate("income", PERIOD), [7]) + + +class MoveAnnualToHousehold(Reform): + def apply(self): + class annual_value(Variable): + entity = Household + default_value = 17 + + self.update_variable(annual_value) + + +def test_entity_change_drops_orphaned_holder_and_input_records(): + simulation = build() + simulation.set_input("annual_value", PERIOD, [7, 9]) + simulation.input_variables = ["annual_value", "eligible"] + + simulation.apply_reform(MoveAnnualToHousehold) + + assert "annual_value" not in simulation.persons._holders + assert "annual_value" not in simulation.input_variables + assert all(key[0] != "annual_value" for key in simulation._user_input_keys) + np.testing.assert_array_equal(simulation.calculate("annual_value", PERIOD), [17]) + assert_bound(simulation) + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_entity_change_keeps_inputs_supplied_to_the_new_holder(on_disk): + class MoveAndSupply(Reform): + def apply(self): + MoveAnnualToHousehold.apply(self) + self.simulation.set_input("annual_value", PERIOD, [31]) + + simulation = build() + if on_disk: + store_on_disk(simulation, "annual_value") + simulation.set_input("annual_value", PERIOD, [7, 9]) + simulation.input_variables = ["annual_value", "eligible"] + + simulation.apply_reform(MoveAndSupply) + + assert "annual_value" not in simulation.persons._holders + assert "annual_value" in simulation.input_variables + assert ("annual_value", "default", period(PERIOD)) in simulation._user_input_keys + np.testing.assert_array_equal(simulation.calculate("annual_value", PERIOD), [31]) + assert_bound(simulation) + + +class MakeAnnual(Reform): + def apply(self): + class eternal_value(Variable): + definition_period = YEAR + + def formula(person, period): + return person.filled_array(period.start.year) + + self.update_variable(eternal_value) + + +class MakeEternal(Reform): + def apply(self): + class annual_value(Variable): + definition_period = ETERNITY + + self.update_variable(annual_value) + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_eternity_to_year_drops_incompatible_storage_and_inputs(on_disk): + simulation = build() + if on_disk: + store_on_disk(simulation, "eternal_value") + simulation.set_input("eternal_value", PERIOD, [7, 9]) + simulation.input_variables = ["eternal_value", "eligible"] + + simulation.apply_reform(MakeAnnual) + + np.testing.assert_array_equal( + simulation.calculate("eternal_value", "2025"), [2025, 2025] + ) + np.testing.assert_array_equal( + simulation.calculate("eternal_value", "2026"), [2026, 2026] + ) + assert "eternal_value" not in simulation.input_variables + assert all(key[0] != "eternal_value" for key in simulation._user_input_keys) + assert_bound(simulation) + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_year_to_eternity_uses_eternal_storage_for_new_inputs(on_disk): + simulation = build() + if on_disk: + store_on_disk(simulation, "annual_value") + simulation.set_input("annual_value", PERIOD, [7, 9]) + simulation.apply_reform(MakeEternal) + + assert all(key[0] != "annual_value" for key in simulation._user_input_keys) + simulation.set_input("annual_value", PERIOD, [11, 12]) + + assert_bound(simulation) + np.testing.assert_array_equal( + simulation.calculate("annual_value", "2025"), [11, 12] + ) + np.testing.assert_array_equal( + simulation.calculate("annual_value", "2026"), [11, 12] + ) diff --git a/tests/core/test_apply_reform_holders_property.py b/tests/core/test_apply_reform_holders_property.py new file mode 100644 index 000000000..be7b498e6 --- /dev/null +++ b/tests/core/test_apply_reform_holders_property.py @@ -0,0 +1,226 @@ +"""Check holder identity, defaults and recorded inputs after a reform.""" + +import numpy as np +import pytest + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +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 + +pytestmark = pytest.mark.filterwarnings( + "ignore:Memory configuration:policyengine_core.warnings.MemoryConfigWarning" +) + +PROPERTY_SETTINGS = hypothesis.settings(max_examples=50, deadline=None, database=None) +MASKS = st.lists(st.booleans(), min_size=1, max_size=6) +DEFAULTS = st.integers(min_value=-1_000, max_value=1_000) +REPLACED = ("award", "annual_value", "group_value") + + +def assert_holder_identity(simulation): + for population in simulation.populations.values(): + for name, holder in population._holders.items(): + assert holder.variable is simulation.tax_benefit_system.variables[name] + + +def expected_award(mask, default): + return np.where(mask, 100.0, float(default)) + + +def change_several_defaults(default): + class ChangeSeveralDefaults(Reform): + def apply(self): + for offset, name in enumerate(REPLACED): + attributes = {"default_value": float(default + offset)} + if name == "award": + attributes["defined_for"] = "eligible" + self.update_variable(type(name, (Variable,), attributes)) + + return ChangeSeveralDefaults + + +@PROPERTY_SETTINGS +@hypothesis.given( + mask=MASKS, + default=DEFAULTS, + order=st.permutations(REPLACED), + warm=st.booleans(), + branches=st.sampled_from(("none", "shared", "independent", "both")), +) +@hypothesis.example( + mask=[False, True], + default=7, + order=("group_value", "award", "annual_value"), + warm=True, + branches="both", +) +def test_every_existing_holder_uses_its_system_variable( + mask, default, order, warm, branches +): + simulation = build(mask) + for name in REPLACED: + population = simulation.get_variable_population(name) + del population._holders[name] + for name in order: + simulation.get_holder(name) + if warm: + simulation.calculate(name, YEAR) + related = [] + if branches in ("shared", "both"): + related.append((simulation.get_branch("shared"), default)) + if branches in ("independent", "both"): + related.append((simulation.get_branch("independent", clone_system=True), 0)) + + simulation.apply_reform(change_several_defaults(default)) + + for member in [simulation, *(branch for branch, _ in related)]: + assert_holder_identity(member) + for member, expected_default in [(simulation, default), *related]: + np.testing.assert_array_equal( + member.calculate("award", YEAR), expected_award(mask, expected_default) + ) + np.testing.assert_array_equal( + simulation.calculate("annual_value", YEAR), + np.full(len(mask), default + 1), + ) + np.testing.assert_array_equal( + simulation.calculate("group_value", YEAR), [default + 2] + ) + + +@PROPERTY_SETTINGS +@hypothesis.given(mask=MASKS, default=DEFAULTS, warm=st.booleans()) +@hypothesis.example(mask=[False, False], default=7, warm=False) +@hypothesis.example(mask=[False, True], default=7, warm=True) +def test_person_result_does_not_depend_on_other_people(mask, default, warm): + combined = build(mask) + if warm: + combined.calculate("award", YEAR) + reform = change_default(default) + combined.apply_reform(reform) + together = combined.calculate("award", YEAR) + alone = [] + for eligible in mask: + individual = build([eligible]) + if warm: + individual.calculate("award", YEAR) + individual.apply_reform(reform) + alone.append(individual.calculate("award", YEAR)[0]) + + np.testing.assert_array_equal(together, expected_award(mask, default)) + np.testing.assert_array_equal(alone, expected_award(mask, default)) + np.testing.assert_array_equal(together, alone) + + +@PROPERTY_SETTINGS +@hypothesis.given(mask=MASKS, default=DEFAULTS, warm=st.booleans()) +@hypothesis.example(mask=[False, False], default=7, warm=True) +@hypothesis.example(mask=[False, True], default=7, warm=False) +def test_reform_before_or_after_holder_creation_returns_same_values( + mask, default, warm +): + before_holder = build(mask) + del before_holder.person._holders["award"] + after_holder = build(mask) + if warm: + after_holder.calculate("award", YEAR) + reform = change_default(default) + + before_holder.apply_reform(reform) + after_holder.apply_reform(reform) + + expected = expected_award(mask, default) + np.testing.assert_array_equal(before_holder.calculate("award", YEAR), expected) + np.testing.assert_array_equal(after_holder.calculate("award", YEAR), expected) + assert_holder_identity(before_holder) + assert_holder_identity(after_holder) + + +INPUT_ROWS = st.lists( + st.tuples( + st.booleans(), + st.integers(min_value=-10_000, max_value=10_000), + st.integers(min_value=-10_000, max_value=10_000), + ), + min_size=1, + max_size=6, +) + + +@PROPERTY_SETTINGS +@hypothesis.given( + rows=INPUT_ROWS, + defaults=st.lists(DEFAULTS, min_size=1, max_size=4), + on_disk=st.booleans(), + branch=st.booleans(), +) +@hypothesis.example( + rows=[(False, 12, 34), (True, -5, 9)], + defaults=[7, 11], + on_disk=True, + branch=True, +) +def test_unchanged_recorded_inputs_survive_reforms(rows, defaults, on_disk, branch): + mask, incomes, eternal = zip(*rows) + simulation = build(mask) + if on_disk: + store_on_disk(simulation, "income", "annual_value", "eternal_value") + supplied = { + "income": np.asarray(incomes, dtype=np.float32), + "annual_value": np.asarray(incomes, dtype=np.float32) + 1, + "eternal_value": np.asarray(eternal, dtype=np.float32), + } + for name, values in supplied.items(): + simulation.set_input(name, YEAR, values) + if on_disk: + holder = simulation.get_holder(name) + assert holder._disk_storage.get(periods.period(YEAR)) is not None + assert holder._memory_storage.get(periods.period(YEAR)) is None + members = [simulation] + if branch: + members.append(simulation.get_branch("existing")) + input_keys = [set(member._user_input_keys) for member in members] + + for default in defaults: + simulation.apply_reform(change_default(default)) + for member, keys in zip(members, input_keys): + assert member._user_input_keys == keys + for name, expected in supplied.items(): + np.testing.assert_array_equal(member.calculate(name, YEAR), expected) + np.testing.assert_array_equal( + member.calculate("award", YEAR), expected_award(mask, default) + ) + assert_holder_identity(member) + + +@PROPERTY_SETTINGS +@hypothesis.given(mask=MASKS, default=DEFAULTS, warm=st.booleans()) +@hypothesis.example(mask=[False, False], default=7, warm=True) +@hypothesis.example(mask=[False, True], default=7, warm=False) +def test_applying_same_reform_object_twice_returns_same_values(mask, default, warm): + once = build(mask) + twice = build(mask) + inputs = np.arange(len(mask), dtype=np.float32) + 5 + for simulation in (once, twice): + simulation.set_input("income", YEAR, inputs) + if warm: + simulation.calculate("award", YEAR) + reform = change_default(default) + + once.apply_reform(reform) + twice.apply_reform(reform) + twice.calculate("award", YEAR) + twice.apply_reform(reform) + + expected = expected_award(mask, default) + for simulation in (once, twice): + np.testing.assert_array_equal(simulation.calculate("award", YEAR), expected) + np.testing.assert_array_equal(simulation.calculate("income", YEAR), inputs) + assert_holder_identity(simulation) + np.testing.assert_array_equal( + twice.calculate("award", YEAR), once.calculate("award", YEAR) + ) diff --git a/tests/fixtures/reform_holders.py b/tests/fixtures/reform_holders.py new file mode 100644 index 000000000..65687bcab --- /dev/null +++ b/tests/fixtures/reform_holders.py @@ -0,0 +1,115 @@ +"""Build small simulations for reforms applied after holders exist.""" + +from policyengine_core import periods +from policyengine_core.entities import Entity, GroupEntity +from policyengine_core.experimental import MemoryConfig +from policyengine_core.parameters import ParameterNode +from policyengine_core.reforms import Reform +from policyengine_core.simulations import Simulation +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.variables import Variable + +YEAR = "2025" +Person = Entity("person", "persons", "Person", "") +Household = GroupEntity( + "household", + "households", + "Household", + "", + roles=[{"key": "member", "plural": "members"}], +) + + +class eligible(Variable): + value_type = bool + entity = Person + definition_period = periods.YEAR + label = "Eligibility" + + +class award(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + default_value = 0 + defined_for = "eligible" + label = "Award" + + def formula(person, period): + return person.filled_array(100) + + +class income(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Income" + + +class annual_value(Variable): + value_type = float + entity = Person + definition_period = periods.YEAR + label = "Annual value" + + +class eternal_value(Variable): + value_type = float + entity = Person + definition_period = periods.ETERNITY + label = "Eternal value" + + +class group_value(Variable): + value_type = float + entity = Household + definition_period = periods.YEAR + label = "Group value" + + +def build(mask=(False, True)): + """Put everyone in one household and record each person's eligibility.""" + system = TaxBenefitSystem([Person, Household]) + system.add_variables( + eligible, award, income, annual_value, eternal_value, group_value + ) + system.parameters = ParameterNode( + "", data={"amount": {"values": {"2025-01-01": 1}}} + ) + names = [f"p{index}" for index in range(len(mask))] + return Simulation( + tax_benefit_system=system, + situation={ + "persons": { + name: {"eligible": {YEAR: bool(value)}} + for name, value in zip(names, mask) + }, + "households": {"h": {"members": names}}, + }, + ) + + +def change_default(value): + """Return a reform that replaces the award's default with ``value``.""" + + class ChangeDefault(Reform): + def apply(self): + self.update_variable( + type( + "award", + (Variable,), + {"default_value": float(value), "defined_for": "eligible"}, + ) + ) + + return ChangeDefault + + +def store_on_disk(simulation, *names): + """Store future writes to the named holders on disk.""" + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + for name in names: + holder = simulation.get_holder(name) + if holder._disk_storage is None: + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True