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..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) @@ -949,10 +986,27 @@ 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._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 + # 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..388dc55cd84 100644 --- a/policyengine_us/tests/core/test_spm_system.py +++ b/policyengine_us/tests/core/test_spm_system.py @@ -12,8 +12,10 @@ 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.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 @@ -665,6 +667,219 @@ 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. + + 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(): + 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) + + +@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(), + # 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"}, + ) + 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 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 and redefines variables": AddAndRedefineVariables, + "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"), + ("reform", "is_household_head"), + ("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 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( + 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, 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 + 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 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(): """Nothing will ask that provider for a county, so nothing is scanned.""" source = small_dataset()