From 9d35b2950c0cd570e6393aea7e2cf3c6e5866515 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 11:54:56 -0400 Subject: [PATCH 1/2] Bind a subsampled reform's rebuilt baseline arm to baseline policy Core's subsample rebuilds the baseline arm as a branch of the reform and then hands it the replaced arm's policy. Before core 3.32.21 it does not rebind the branch's holders, so the baseline read the reform's neutralized income_tax ($0 instead of $4,016), and the arm kept the replaced arm as its own baseline. No core version points the policy back at the new arm or primes its private parameter-tracing root, so a traced baseline arm wrote its tracer onto the shared parameter tree (core 3.32.8). Clear the arm's baseline, rebind its holders and isolate its tracing, as construction does. Adds a regression test for both wrappers and their clones and a Hypothesis property that subsampling never changes which policy a holder belongs to. Found in review of #9698 (P3-1). Co-Authored-By: Claude Opus 5.5 --- .../fix-subsample-baseline-rebind.fixed.md | 1 + policyengine_us/spm.py | 24 ++- policyengine_us/tests/core/test_spm_system.py | 186 ++++++++++++++++++ 3 files changed, 207 insertions(+), 4 deletions(-) create mode 100644 changelog.d/fix-subsample-baseline-rebind.fixed.md diff --git a/changelog.d/fix-subsample-baseline-rebind.fixed.md b/changelog.d/fix-subsample-baseline-rebind.fixed.md new file mode 100644 index 00000000000..486cc17ba85 --- /dev/null +++ b/changelog.d/fix-subsample-baseline-rebind.fixed.md @@ -0,0 +1 @@ +A reform simulation that is subsampled now rebuilds its baseline arm bound to baseline policy: holders, populations, inputs and parameter tracing belong to the baseline again, and the arm no longer keeps the replaced arm as its own baseline. diff --git a/policyengine_us/spm.py b/policyengine_us/spm.py index 367cb91d880..d6e37a95534 100644 --- a/policyengine_us/spm.py +++ b/policyengine_us/spm.py @@ -949,10 +949,26 @@ def _simulation_family(self): def subsample(self, *args, **kwargs): result = super().subsample(*args, **kwargs) - # Core rebuilds the baseline arm from the subsampled population, then - # hands it the previous arm's policy and so its provider, which - # recorded the counties the previous arm read. Record what every - # simulation in the family reads now. + baseline = self.branches.get("baseline") + if baseline is not None: + # Core rebuilds the baseline arm by branching this simulation, so + # the new arm starts out as a copy of this one: its holders hold + # this simulation's variables, a reform's neutralized or replaced + # ones included, and its own ``baseline`` is the arm it replaces, + # on the population from before subsampling. Core then hands it + # the replaced arm's policy; before 3.32.21 it does not rebind the + # holders to that policy, and it never claims the policy for the + # new arm or primes the policy's private parameter-tracing root, + # which are this package's state. Give it what construction gives + # a baseline arm: no baseline of its own, so it computes no + # behavioural responses, holders bound to its own policy, and that + # policy's tracing primed for this arm. + baseline.baseline = None + baseline._rebind_holders() + baseline._isolate_parameter_tracing() + # Core hands the rebuilt baseline arm the previous arm's policy and so + # its provider, which recorded the counties the previous arm read. + # Record what every simulation in the family reads now. self._record_county_input_types() return result diff --git a/policyengine_us/tests/core/test_spm_system.py b/policyengine_us/tests/core/test_spm_system.py index 1fa4f4240b8..e9cd7a5580c 100644 --- a/policyengine_us/tests/core/test_spm_system.py +++ b/policyengine_us/tests/core/test_spm_system.py @@ -12,6 +12,8 @@ import numpy as np import pandas as pd import pytest +from hypothesis import HealthCheck, example, given, settings +from hypothesis import strategies as st from policyengine_core.reforms import Reform from policyengine_core.periods import YEAR from policyengine_core.variables import Variable @@ -665,6 +667,190 @@ def test_subsampling_records_the_counties_each_arm_then_reads(): assert np.all(arm.calculate("spm_unit_spm_threshold", 2024) > 0) +def earner_dataset(): + """Two single $50,000 earners, each in a household of their own.""" + source = small_dataset() + source.person["age"] = [40, 40] + source.person["employment_income"] = [50_000, 50_000] + return source + + +def assert_bound_to_own_policy(simulation): + """Each population, holder and recorded input belongs to this arm's policy.""" + policy = simulation.tax_benefit_system + assert policy.simulation is simulation + for population in simulation.populations.values(): + assert population.entity._tax_benefit_system is policy + for name, holder in population._holders.items(): + assert holder.variable is policy.variables.get(name), name + assert holder.simulation is simulation + assert {key[0] for key in simulation._user_input_keys} <= set(policy.variables) + assert set(simulation.input_variables) <= set(policy.variables) + + +@pytest.mark.parametrize("simulation_type", [Simulation, Microsimulation]) +def test_subsampled_reform_baseline_and_clones_calculate_original_tax(simulation_type): + """Core rebuilds a subsampled reform's baseline arm as a branch of the reform. + + The branch starts out bound to the reform. Until it was bound to the + baseline policy core hands it, its holders held the reform's neutralized + ``income_tax`` (core before 3.32.21), that policy still named the replaced + arm as its simulation, and the arm's own ``baseline`` was the replaced arm. + Tracing turned on after construction then reached the baseline policy's + parameter tree unprimed: the tree it shares with every ordinary + simulation, which core before 3.32.21 marks with the request's tracer. + """ + changed = simulation_type( + dataset=earner_dataset(), + reform=UserReform, + spm={"geography_kind": "national"}, + ) + reform_policy = changed.tax_benefit_system + baseline_policy = changed.baseline.tax_benefit_system + changed.trace = True + + changed.subsample(2) + + baseline = changed.baseline + np.testing.assert_array_equal(baseline.calculate("income_tax", 2024), 4_016) + np.testing.assert_array_equal(changed.calculate("income_tax", 2024), 0) + np.testing.assert_array_equal( + baseline.calculate("household_market_income", 2024), 50_000 + ) + np.testing.assert_array_equal( + changed.calculate("household_market_income", 2024), 123 + ) + assert baseline is changed.branches["baseline"] + assert baseline.baseline is None + assert changed.tax_benefit_system is reform_policy + assert baseline.tax_benefit_system is baseline_policy + for arm in (changed, baseline): + assert_bound_to_own_policy(arm) + + # The baseline arm is traced in the reform's tracer, through its own root. + root = baseline_policy.parameters + assert baseline.trace + assert baseline.tracer is changed.tracer + assert root is not system.parameters + assert root.trace + assert root.tracer is baseline.tracer + assert root.branch_name == "baseline" + assert system.parameters.trace is False + assert system.parameters.tracer is None + + for original, expected in ((changed, 0), (baseline, 4_016)): + clone = original.clone() + np.testing.assert_array_equal(clone.calculate("income_tax", 2024), expected) + clone.delete_arrays("income_tax") + np.testing.assert_array_equal(clone.calculate("income_tax", 2024), expected) + assert_bound_to_own_policy(clone) + + +class reform_only_income(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "income defined only under a reform" + + def formula(person, period, parameters): + return person.filled_array(1) + + +class AddReformOnlyVariable(Reform): + def apply(self): + self.add_variable(reform_only_income) + self.neutralize_variable("income_tax") + + +SUBSAMPLED_REFORMS = { + "replaces variables": UserReform, + "adds a variable": AddReformOnlyVariable, + "changes a parameter": Reform.from_dict( + {"gov.irs.deductions.standard.amount.SINGLE": {"2024": 0}} + ), +} +# What an arm may have calculated before subsampling, so that it holds it. +CALCULATED_BEFORE_SUBSAMPLING = ( + ("reform", "household_market_income"), + ("reform", "reform_only_income"), + ("baseline", "household_market_income"), + ("baseline", "income_tax"), +) + + +@settings( + max_examples=2, + deadline=None, + derandomize=True, + suppress_health_check=[HealthCheck.too_slow], +) +@given( + reform=st.sampled_from(sorted(SUBSAMPLED_REFORMS)), + calculated=st.lists( + st.sampled_from(CALCULATED_BEFORE_SUBSAMPLING), unique=True, max_size=3 + ), + draws=st.lists( + st.tuples(st.integers(1, 3), st.integers(0, 2**16), st.booleans()), + min_size=1, + max_size=2, + ), +) +@example( + reform="adds a variable", + calculated=[("reform", "reform_only_income"), ("baseline", "income_tax")], + draws=[(1, 0, True), (2, 1, False)], +) +def test_subsampling_never_changes_which_policy_a_holder_belongs_to( + reform, calculated, draws +): + """Invariant: subsampling binds each arm's holders to that arm's own policy. + + However often a reform simulation is subsampled, each arm keeps the policy + it had, and every population, holder and recorded input in either arm + belongs to that policy. The rebuilt baseline arm has no baseline of its + own and holds nothing that only the reform defines, including what the + reform arm calculated, and core therefore exported as an input, before + subsampling. + + This compares no values: values the reform arm calculated before + subsampling become inputs of the rebuilt population, in both arms. That + is core's export, not a binding, and is not this invariant. + """ + changed = Microsimulation( + dataset=earner_dataset(), + # Reforming a supplied system clones it instead of rebuilding policy + # from source, which keeps each example to a few seconds. + tax_benefit_system=system, + reform=SUBSAMPLED_REFORMS[reform], + spm={"geography_kind": "national"}, + ) + policies = { + "reform": changed.tax_benefit_system, + "baseline": changed.baseline.tax_benefit_system, + } + for arm, variable in calculated: + simulation = changed if arm == "reform" else changed.baseline + if variable in simulation.tax_benefit_system.variables: + simulation.calculate(variable, 2024) + + for n, seed, quantize_weights in draws: + changed.subsample(n, seed=seed, quantize_weights=quantize_weights) + + baseline = changed.branches["baseline"] + assert changed.baseline is baseline + assert baseline.baseline is None + assert changed.tax_benefit_system is policies["reform"] + assert baseline.tax_benefit_system is policies["baseline"] + for simulation in (changed, baseline): + assert_bound_to_own_policy(simulation) + if reform == "adds a variable" and ("reform", "reform_only_income") in ( + calculated + ): + # The reform-only value came through the rebuild as an input, so + # the baseline branch started out holding it. + assert changed.person._holders["reform_only_income"].get_known_periods() + + def test_a_national_selection_records_no_county_input_types(): """Nothing will ask that provider for a county, so nothing is scanned.""" source = small_dataset() From 22183cb3b9f7597d07edf5384b1d3b6d8a4e1892 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 15:28:57 -0400 Subject: [PATCH 2/2] Drop holders a baseline arm cannot use, at construction and after subsampling A reform's baseline arm is a branch of the reform, so its holders were built for the reform's variables. Rebinding one to the baseline's variable kept the holder's population and storage, so a reform that moved a variable to another entity, or into or out of ETERNITY, left the baseline arm with storage its own variable cannot use (core 3.32.8; core 3.32.21 drops these holders itself). Drop them with their recorded inputs, as core 3.32.21's binder does, on both baseline-arm paths. apply_reform keeps a simulation's own inputs, as core's does. No shipped reform changes a variable's entity or ETERNITY-ness. Tests now also check that each holder suits its variable, and the property test's structural reform redefines is_household_head as yearly. The regression test reforms the supplied system instead of rebuilding policy from source. Addresses review r1 of #9994. Co-Authored-By: Claude Opus 5.5 --- policyengine_us/spm.py | 38 ++++++++++++ policyengine_us/tests/core/test_spm_system.py | 61 ++++++++++++++----- 2 files changed, 83 insertions(+), 16 deletions(-) diff --git a/policyengine_us/spm.py b/policyengine_us/spm.py index d6e37a95534..d5c540e0a96 100644 --- a/policyengine_us/spm.py +++ b/policyengine_us/spm.py @@ -678,6 +678,9 @@ class SPMSimulationMixin: def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # Core switches the baseline system after cloning its populations. + baseline = self.branches.get("baseline") + if baseline is not None: + baseline._drop_holders_built_for_another_policy() self._rebind_holders() # Core sets ``trace`` before this simulation owns its policy state. self._isolate_parameter_tracing() @@ -738,6 +741,40 @@ def _rebind_holders(self): for branch in self.branches.values(): branch._rebind_holders() + def _drop_holders_built_for_another_policy(self): + """Drop holders whose storage this simulation's policy cannot use. + + A reform's baseline arm is a branch of the reform, so its holders were + built for the reform's variables. Rebinding a holder to the baseline's + variable keeps the holder's population and storage. A holder in + another entity's population, or with storage for (or not for) an + ``ETERNITY`` variable, therefore cannot hold the baseline's + definition. Drop it with its recorded inputs, as core 3.32.21 and + later do in ``_bind_to_tax_benefit_system``; core 3.32.8 keeps it. + + Only a baseline arm does this. ``apply_reform`` keeps a simulation's + own inputs, as core's does. + """ + variables = self.tax_benefit_system.variables + dropped = set() + for population in self.populations.values(): + for name, holder in list(population._holders.items()): + variable = variables.get(name) + if variable is not None and ( + variable.entity.key != population.entity.key + or (variable.definition_period == periods_.ETERNITY) + != holder._memory_storage.is_eternal + ): + del population._holders[name] + dropped.add(name) + if dropped: + self._user_input_keys = { + key for key in self._user_input_keys if key[0] not in dropped + } + self.input_variables = [ + name for name in self.input_variables if name not in dropped + ] + def apply_reform(self, reform): policy = self.tax_benefit_system detaching = getattr(policy, "detaching_parameters", None) @@ -964,6 +1001,7 @@ def subsample(self, *args, **kwargs): # behavioural responses, holders bound to its own policy, and that # policy's tracing primed for this arm. baseline.baseline = None + baseline._drop_holders_built_for_another_policy() baseline._rebind_holders() baseline._isolate_parameter_tracing() # Core hands the rebuilt baseline arm the previous arm's policy and so diff --git a/policyengine_us/tests/core/test_spm_system.py b/policyengine_us/tests/core/test_spm_system.py index e9cd7a5580c..388dc55cd84 100644 --- a/policyengine_us/tests/core/test_spm_system.py +++ b/policyengine_us/tests/core/test_spm_system.py @@ -15,7 +15,7 @@ from hypothesis import HealthCheck, example, given, settings from hypothesis import strategies as st from policyengine_core.reforms import Reform -from policyengine_core.periods import YEAR +from policyengine_core.periods import ETERNITY, YEAR from policyengine_core.variables import Variable from spm_calculator.errors import SPMInputError from spm_calculator.policyengine_adapter import FORMULA_OWNED_INPUTS @@ -676,14 +676,24 @@ def earner_dataset(): def assert_bound_to_own_policy(simulation): - """Each population, holder and recorded input belongs to this arm's policy.""" + """Each population, holder and recorded input belongs to this arm's policy. + + A holder belongs to a variable only if it also suits it: it sits in that + variable's entity's population, with storage for an ``ETERNITY`` variable + exactly when the variable is one. + """ policy = simulation.tax_benefit_system assert policy.simulation is simulation for population in simulation.populations.values(): assert population.entity._tax_benefit_system is policy for name, holder in population._holders.items(): - assert holder.variable is policy.variables.get(name), name + variable = policy.variables.get(name) + assert holder.variable is variable, name assert holder.simulation is simulation + assert variable.entity.key == population.entity.key, name + assert holder._memory_storage.is_eternal == ( + variable.definition_period == ETERNITY + ), name assert {key[0] for key in simulation._user_input_keys} <= set(policy.variables) assert set(simulation.input_variables) <= set(policy.variables) @@ -702,6 +712,9 @@ def test_subsampled_reform_baseline_and_clones_calculate_original_tax(simulation """ changed = simulation_type( dataset=earner_dataset(), + # Reforming a supplied system clones it instead of rebuilding policy + # from source; the constructor tests above cover building it. + tax_benefit_system=system, reform=UserReform, spm={"geography_kind": "national"}, ) @@ -756,15 +769,22 @@ def formula(person, period, parameters): return person.filled_array(1) -class AddReformOnlyVariable(Reform): +class is_household_head(Variable): + # Each year's input instead of one for all time: a holder the reform arm + # stores this in has storage the baseline's ETERNITY variable cannot use. + definition_period = YEAR + + +class AddAndRedefineVariables(Reform): def apply(self): self.add_variable(reform_only_income) + self.update_variable(is_household_head) self.neutralize_variable("income_tax") SUBSAMPLED_REFORMS = { "replaces variables": UserReform, - "adds a variable": AddReformOnlyVariable, + "adds and redefines variables": AddAndRedefineVariables, "changes a parameter": Reform.from_dict( {"gov.irs.deductions.standard.amount.SINGLE": {"2024": 0}} ), @@ -773,6 +793,7 @@ def apply(self): CALCULATED_BEFORE_SUBSAMPLING = ( ("reform", "household_market_income"), ("reform", "reform_only_income"), + ("reform", "is_household_head"), ("baseline", "household_market_income"), ("baseline", "income_tax"), ) @@ -796,8 +817,12 @@ def apply(self): ), ) @example( - reform="adds a variable", - calculated=[("reform", "reform_only_income"), ("baseline", "income_tax")], + reform="adds and redefines variables", + calculated=[ + ("reform", "reform_only_income"), + ("reform", "is_household_head"), + ("baseline", "income_tax"), + ], draws=[(1, 0, True), (2, 1, False)], ) def test_subsampling_never_changes_which_policy_a_holder_belongs_to( @@ -808,9 +833,11 @@ def test_subsampling_never_changes_which_policy_a_holder_belongs_to( However often a reform simulation is subsampled, each arm keeps the policy it had, and every population, holder and recorded input in either arm belongs to that policy. The rebuilt baseline arm has no baseline of its - own and holds nothing that only the reform defines, including what the - reform arm calculated, and core therefore exported as an input, before - subsampling. + own, holds nothing that only the reform defines, and keeps no holder the + reform built for its own definition of a variable: here, yearly storage + for the baseline's ``ETERNITY`` ``is_household_head``. That includes what + the reform arm calculated, and core therefore exported as an input, + before subsampling. This compares no values: values the reform arm calculated before subsampling become inputs of the rebuilt population, in both arms. That @@ -843,12 +870,14 @@ def test_subsampling_never_changes_which_policy_a_holder_belongs_to( assert baseline.tax_benefit_system is policies["baseline"] for simulation in (changed, baseline): assert_bound_to_own_policy(simulation) - if reform == "adds a variable" and ("reform", "reform_only_income") in ( - calculated - ): - # The reform-only value came through the rebuild as an input, so - # the baseline branch started out holding it. - assert changed.person._holders["reform_only_income"].get_known_periods() + if reform == "adds and redefines variables": + # What the reform arm held came through the rebuild as inputs, so + # the baseline branch started out holding the reform's holders. + for variable in ("reform_only_income", "is_household_head"): + if ("reform", variable) in calculated: + holder = changed.person._holders[variable] + assert holder.get_known_periods() + assert not holder._memory_storage.is_eternal def test_a_national_selection_records_no_county_input_types():