From 1fcfce2b78105fbe4ebe6b416e0d76d848774546 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 29 Sep 2026 21:15:38 -0400 Subject: [PATCH 1/6] Make branch-measured marginal tax rates match fresh re-simulation marginal_tax_rate and its relatives measure a rate by raising one adult's earnings in a branch that keeps only the simulation's listed inputs. Two defects made the branch disagree with a fresh simulation of the same household: - Simulation moves employment_income, self-employment income, weekly hours and long-term gains onto pre-response variables, and backfills state_code from state_code_str, after policyengine-core records input_variables. Only Microsimulation patched the list, so household branches erased the moved inputs. An Ohio parent earning $28,000 then drew $7,475 of Ohio Works First in the branch and reported a rate of -4.70; fresh simulations give 0.527. - Branches raised the employment_income aggregate, which programs that read pre-response income (TANF's earned income) never see, and left the base's cached monthly values in place. Microsimulation rates therefore missed TANF phase-outs. The moves and the input list now live in shared helpers used by both wrappers, and one branch builder raises the variable that holds each input for every branch-measured rate. marginal_tax_rate_on_capital_gains raised the capital_gains aggregate, which taxed the increase at ordinary rates; it now raises long-term gains. Co-Authored-By: Claude Opus 5.5 --- .../fix-mtr-branch-divergence.fixed.md | 1 + policyengine_us/system.py | 144 ++---- .../tests/core/test_marginal_rate_branches.py | 457 ++++++++++++++++++ .../baseline/household/marginal_tax_rate.yaml | 174 +++++++ .../behavioral_response_measurements.py | 10 + .../gov/simulation/capital_gains_responses.py | 30 +- .../variables/household/marginal_tax_rate.py | 66 +-- .../household/marginal_tax_rate_helpers.py | 155 ++++-- ...inal_tax_rate_including_health_benefits.py | 72 +-- 9 files changed, 845 insertions(+), 264 deletions(-) create mode 100644 changelog.d/fix-mtr-branch-divergence.fixed.md create mode 100644 policyengine_us/tests/core/test_marginal_rate_branches.py diff --git a/changelog.d/fix-mtr-branch-divergence.fixed.md b/changelog.d/fix-mtr-branch-divergence.fixed.md new file mode 100644 index 00000000000..973c462cd4a --- /dev/null +++ b/changelog.d/fix-mtr-branch-divergence.fixed.md @@ -0,0 +1 @@ +Make branch-measured marginal tax rates match fresh re-simulation. Household simulations now list the inputs they move onto pre-response variables (`employment_income_before_lsr`, `weekly_hours_worked_before_lsr`, `long_term_capital_gains_before_response`, `state_code`) in `input_variables`, so rate branches no longer erase them. The branches also raise those inputs rather than the aggregates, so programs that read pre-response income, such as TANF, see the raise. `marginal_tax_rate_on_capital_gains` now raises long-term gains, so preferential rates apply to the increase. diff --git a/policyengine_us/system.py b/policyengine_us/system.py index 237914bcb3d..dbd0cb0d7bc 100644 --- a/policyengine_us/system.py +++ b/policyengine_us/system.py @@ -22,6 +22,9 @@ backdate_parameters, ) from policyengine_us.reforms import create_structural_reforms_from_parameters +from policyengine_us.variables.gov.simulation.behavioral_response_measurements import ( + PRE_RESPONSE_INPUTS, +) from policyengine_core.parameters.operations.homogenize_parameters import ( homogenize_parameter_structures, ) @@ -174,6 +177,48 @@ def clone(self): system = CountryTaxBenefitSystem() +def _move_inputs_to_pre_response_variables(simulation): + """Store aggregate inputs on the pre-response variables they derive from.""" + for aggregate, pre_response in PRE_RESPONSE_INPUTS.items(): + holder = simulation.get_holder(aggregate) + for known_period in holder.get_known_periods(): + array = holder.get_array(known_period) + simulation.set_input(pre_response, known_period, array) + holder.delete_arrays(known_period) + + +def _sync_input_variables(simulation): + """List the variables that hold inputs after the moves and backfill. + + Core records ``input_variables`` before the pre-response moves and the + ``state_code`` backfill run, so the recorded list still names the emptied + aggregates and ``state_code_str`` and omits the variables that now hold + their values. Branch-based variables such as ``marginal_tax_rate`` delete + every variable outside this list, so a stale list erases those inputs in + the branch: a wage-earning parent's TANF then counts no earnings. + """ + moved = ( + *PRE_RESPONSE_INPUTS, + *PRE_RESPONSE_INPUTS.values(), + "state_code_str", + "state_code", + ) + + def holds_input(variable): + return len(simulation.get_holder(variable).get_known_periods()) > 0 + + input_variables = [ + variable + for variable in simulation.input_variables + if variable not in moved or holds_input(variable) + ] + simulation.input_variables = input_variables + [ + variable + for variable in moved + if variable not in input_variables and holds_input(variable) + ] + + def _backfill_state_code_from_str(simulation): """Backfill the ``state_code`` enum from a ``state_code_str`` input. @@ -240,46 +285,10 @@ def __init__(self, *args, **kwargs): if reform is not None: self.apply_reform(reform) - # Labor supply responses - - employment_income = self.get_holder("employment_income") - for known_period in employment_income.get_known_periods(): - array = employment_income.get_array(known_period) - self.set_input("employment_income_before_lsr", known_period, array) - employment_income.delete_arrays(known_period) - - self_employment_income = self.get_holder("self_employment_income") - for known_period in self_employment_income.get_known_periods(): - array = self_employment_income.get_array(known_period) - self.set_input("self_employment_income_before_lsr", known_period, array) - self_employment_income.delete_arrays(known_period) - - sstb_self_employment_income = self.get_holder("sstb_self_employment_income") - for known_period in sstb_self_employment_income.get_known_periods(): - array = sstb_self_employment_income.get_array(known_period) - self.set_input( - "sstb_self_employment_income_before_lsr", known_period, array - ) - sstb_self_employment_income.delete_arrays(known_period) - - weekly_hours = self.get_holder("weekly_hours_worked") - for known_period in weekly_hours.get_known_periods(): - array = weekly_hours.get_array(known_period) - self.set_input("weekly_hours_worked_before_lsr", known_period, array) - weekly_hours.delete_arrays(known_period) - - # Capital gains responses - - cg_holder = self.get_holder("long_term_capital_gains") - for known_period in cg_holder.get_known_periods(): - array = cg_holder.get_array(known_period) - self.set_input( - "long_term_capital_gains_before_response", known_period, array - ) - cg_holder.delete_arrays(known_period) - + _move_inputs_to_pre_response_variables(self) # Geography backfill: state_code_str-only input -> state_code enum. _backfill_state_code_from_str(self) + _sync_input_variables(self) def _download_or_explain(dataset_str, download): @@ -500,67 +509,12 @@ def __init__(self, *args, **kwargs): if reform is not None: self.apply_reform(reform) - # Labor supply responses - - employment_income = self.get_holder("employment_income") - for known_period in employment_income.get_known_periods(): - array = employment_income.get_array(known_period) - self.set_input("employment_income_before_lsr", known_period, array) - employment_income.delete_arrays(known_period) - - self_employment_income = self.get_holder("self_employment_income") - for known_period in self_employment_income.get_known_periods(): - array = self_employment_income.get_array(known_period) - self.set_input("self_employment_income_before_lsr", known_period, array) - self_employment_income.delete_arrays(known_period) - - sstb_self_employment_income = self.get_holder("sstb_self_employment_income") - for known_period in sstb_self_employment_income.get_known_periods(): - array = sstb_self_employment_income.get_array(known_period) - self.set_input( - "sstb_self_employment_income_before_lsr", known_period, array - ) - sstb_self_employment_income.delete_arrays(known_period) - - weekly_hours = self.get_holder("weekly_hours_worked") - for known_period in weekly_hours.get_known_periods(): - array = weekly_hours.get_array(known_period) - self.set_input("weekly_hours_worked_before_lsr", known_period, array) - weekly_hours.delete_arrays(known_period) - - # Capital gains responses - - cg_holder = self.get_holder("long_term_capital_gains") - for known_period in cg_holder.get_known_periods(): - array = cg_holder.get_array(known_period) - self.set_input( - "long_term_capital_gains_before_response", known_period, array - ) - cg_holder.delete_arrays(known_period) - + _move_inputs_to_pre_response_variables(self) # Geography backfill: state_code_str-only input -> state_code enum. # Datasets carry state_fips, so this only fires for situations that # explicitly supply state_code_str. _backfill_state_code_from_str(self) - - self.input_variables = [ - variable - for variable in self.input_variables - if variable - not in [ - "employment_income", - "self_employment_income", - "sstb_self_employment_income", - "weekly_hours_worked", - "capital_gains", - ] - ] + [ - "employment_income_before_lsr", - "self_employment_income_before_lsr", - "sstb_self_employment_income_before_lsr", - "weekly_hours_worked_before_lsr", - "long_term_capital_gains_before_response", - ] + _sync_input_variables(self) class IndividualSim(CoreIndividualSim): # Deprecated diff --git a/policyengine_us/tests/core/test_marginal_rate_branches.py b/policyengine_us/tests/core/test_marginal_rate_branches.py new file mode 100644 index 00000000000..d41bb1ff0e4 --- /dev/null +++ b/policyengine_us/tests/core/test_marginal_rate_branches.py @@ -0,0 +1,457 @@ +"""Branch-measured marginal rates must equal fresh re-simulation. + +``marginal_tax_rate`` raises one adult's earnings by ``delta`` in a branch of +the simulation and compares household net income. The branch must reproduce +an independent simulation of the same household with the raised earnings, so +for every household and every adult the rate measures: + +1. ``marginal_tax_rate == 1 - (net income at earnings + delta - net income) + / delta``, with net income from two fresh simulations; the same holds for + ``marginal_tax_rate_including_health_benefits`` with health benefits + counted, for ``federal_``, ``state_`` and ``fica_marginal_tax_rate`` as + ``(tax at earnings + delta - tax) / delta``, and for + ``marginal_tax_rate_on_capital_gains`` with long-term gains raised instead. +2. People the rate does not measure (children, adults past the cap) get 0. +3. Measuring a rate leaves the simulation's own results and branches as they + were. +4. Right after construction, ``input_variables`` lists exactly the variables + that hold values. + +Invariant 4 is what the branches rely on. The ``Simulation`` and +``Microsimulation`` wrappers move ``employment_income`` and similar inputs onto +pre-response variables (``employment_income_before_lsr`` ...) and backfill +``state_code`` from ``state_code_str`` after policyengine-core records +``input_variables``. The branches keep only listed variables, so a stale list +erased the moved inputs: an Ohio parent's branch counted no earnings for Ohio +Works First and reported a marginal tax rate of -4.70 where fresh simulations +give 0.527. + +These tests use the Python wrappers because the YAML runner skips the moves. +""" + +import copy + +import numpy as np +import pandas as pd +import pytest + +from policyengine_us import Microsimulation, Simulation +from policyengine_us.data.dataset_schema import USSingleYearDataset +from policyengine_us.variables.gov.simulation.behavioral_response_measurements import ( + PRE_RESPONSE_INPUTS, +) + +YEAR = 2026 +DELTA = 1_000 +TOLERANCE = 1e-3 +STATES = ( + "AL AK AZ AR CA CO CT DE DC FL GA HI ID IL IN IA KS KY LA ME MD MA MI MN MS " + "MO MT NE NV NH NJ NM NY NC ND OH OK OR PA RI SC SD TN TX UT VT VA WA WV WI WY" +).split() +EARNINGS_RATES = ( + "marginal_tax_rate", + "marginal_tax_rate_including_health_benefits", + "federal_marginal_tax_rate", + "state_marginal_tax_rate", + "fica_marginal_tax_rate", +) +COMPONENT_TAXES = { + "federal_marginal_tax_rate": "income_tax", + "state_marginal_tax_rate": "state_income_tax", + "fica_marginal_tax_rate": "employee_payroll_tax", +} + + +def situation(households, geography="state_code"): + """Build one situation holding several independent households. + + Each household is ``(state, adults, children)``: ``adults`` maps person + names to their yearly inputs and ``children`` lists ages. Two adults form a + married couple filing one return. + """ + result = { + key: {} + for key in ( + "people", + "tax_units", + "spm_units", + "families", + "marital_units", + "households", + ) + } + for h, (state, adults, children) in enumerate(households): + members = [] + adult_ids = [] + for name, inputs in adults.items(): + person = f"h{h}_{name}" + members.append(person) + adult_ids.append(person) + result["people"][person] = { + variable: {YEAR: value} for variable, value in inputs.items() + } + for i, age in enumerate(children): + person = f"h{h}_child{i}" + members.append(person) + result["people"][person] = { + "age": {YEAR: age}, + "is_tax_unit_dependent": {YEAR: True}, + } + result["marital_units"][f"{person}_marital_unit"] = {"members": [person]} + if len(adult_ids) == 2: + result["marital_units"][f"h{h}_couple"] = {"members": adult_ids} + else: + for person in adult_ids: + result["marital_units"][f"{person}_marital_unit"] = { + "members": [person] + } + for group in ("tax_units", "spm_units", "families"): + result[group][f"h{h}_{group}"] = {"members": members} + result["households"][f"h{h}"] = { + "members": members, + geography: {YEAR: state}, + } + return result + + +def raise_earnings(base_situation, simulation, adult_index): + """Raise the earnings of each household's ``adult_index``-th earner. + + Splits ``DELTA`` across wages, non-SSTB and SSTB self-employment income in + proportion to the person's positive earnings, wages when there are none; + the split the rates document. + """ + ranks = simulation.calculate("adult_earnings_index", YEAR) + wages = simulation.calculate("employment_income", YEAR) + self_employment = simulation.calculate("self_employment_income", YEAR) + sstb = simulation.calculate("sstb_self_employment_income", YEAR) + positive = [np.maximum(x, 0) for x in (wages, self_employment, sstb)] + total = sum(positive) + raised = copy.deepcopy(base_situation) + for i, person in enumerate(base_situation["people"]): + if ranks[i] != adult_index: + continue + shares = ( + [x[i] / total[i] for x in positive] if total[i] > 0 else [1.0, 0.0, 0.0] + ) + for variable, base, share in zip( + ( + "employment_income", + "self_employment_income", + "sstb_self_employment_income", + ), + (wages, self_employment, sstb), + shares, + ): + raised["people"][person][variable] = {YEAR: float(base[i] + DELTA * share)} + return raised, ranks == adult_index + + +def raise_long_term_gains(base_situation, simulation, adult_index): + ranks = simulation.calculate("adult_index_cg", YEAR) + gains = simulation.calculate("long_term_capital_gains", YEAR) + raised = copy.deepcopy(base_situation) + for i, person in enumerate(base_situation["people"]): + if ranks[i] == adult_index: + raised["people"][person]["long_term_capital_gains"] = { + YEAR: float(gains[i] + DELTA) + } + return raised, ranks == adult_index + + +def household_values(simulation, variable): + return simulation.calculate(variable, YEAR, map_to="person") + + +def head_taxes(simulation, variable): + """Tax-unit tax summed over each household's tax units, per person.""" + is_head = simulation.calculate("is_tax_unit_head", YEAR) + person_tax = simulation.calculate(variable, YEAR, map_to="person") * is_head + household = simulation.populations["household"] + return household.project(household.sum(person_tax)) + + +def fresh_earnings_rates(base_situation, simulation): + """Each earnings rate as a difference between fresh simulations.""" + adult_count = simulation.tax_benefit_system.parameters( + YEAR + ).simulation.marginal_tax_rate_adults + base = { + "marginal_tax_rate": household_values(simulation, "household_net_income"), + "marginal_tax_rate_including_health_benefits": household_values( + simulation, "household_net_income_including_health_benefits" + ), + **{rate: head_taxes(simulation, tax) for rate, tax in COMPONENT_TAXES.items()}, + } + expected = {rate: np.zeros(len(base_situation["people"])) for rate in base} + for adult_index in range(1, adult_count + 1): + raised_situation, mask = raise_earnings(base_situation, simulation, adult_index) + raised = Simulation(situation=raised_situation) + alt = { + "marginal_tax_rate": household_values(raised, "household_net_income"), + "marginal_tax_rate_including_health_benefits": household_values( + raised, "household_net_income_including_health_benefits" + ), + **{rate: head_taxes(raised, tax) for rate, tax in COMPONENT_TAXES.items()}, + } + for rate in expected: + change = (alt[rate] - base[rate]) / DELTA + if rate in COMPONENT_TAXES: + value = change + else: + value = 1 - change + expected[rate] = np.where(mask, value, expected[rate]) + return expected + + +def fresh_capital_gains_rate(base_situation, simulation): + base = household_values(simulation, "household_net_income") + expected = np.zeros(len(base_situation["people"])) + for adult_index in (1, 2): + raised_situation, mask = raise_long_term_gains( + base_situation, simulation, adult_index + ) + alt = household_values( + Simulation(situation=raised_situation), "household_net_income" + ) + expected = np.where(mask, 1 - (alt - base) / DELTA, expected) + return expected + + +def assert_rates_match(simulation, expected, rates): + for rate in rates: + np.testing.assert_allclose( + simulation.calculate(rate, YEAR), + expected[rate], + atol=TOLERANCE, + err_msg=rate, + ) + + +OHIO_PARENT = ( + "OH", + {"parent": {"age": 30, "employment_income": 28_000}}, + [4, 8], +) + + +def test_ohio_parent_marginal_tax_rate_matches_fresh_simulations(): + """The reported case: fresh simulations give 0.527, the branch gave -4.70.""" + base_situation = situation([OHIO_PARENT]) + simulation = Simulation(situation=base_situation) + raised_situation = copy.deepcopy(base_situation) + raised_situation["people"]["h0_parent"]["employment_income"] = {YEAR: 29_000} + base = simulation.calculate("household_net_income", YEAR)[0] + raised = Simulation(situation=raised_situation).calculate( + "household_net_income", YEAR + )[0] + + rate = simulation.calculate("marginal_tax_rate", YEAR) + + assert rate[0] == pytest.approx(1 - (raised - base) / DELTA, abs=TOLERANCE) + assert rate[0] == pytest.approx(0.5271, abs=TOLERANCE) + assert list(rate[1:]) == [0, 0] + + +def test_simulation_lists_the_variables_holding_inputs(): + """Invariant 4 for household simulations, across every moved input.""" + base_situation = situation( + [ + ( + "OH", + { + "worker": { + "age": 45, + "employment_income": 30_000, + "self_employment_income": 5_000, + "sstb_self_employment_income": 2_000, + "weekly_hours_worked": 35, + "long_term_capital_gains": 4_000, + } + }, + [], + ) + ], + geography="state_code_str", + ) + simulation = Simulation(situation=base_situation) + holding_values = { + variable + for variable in simulation.tax_benefit_system.variables + if simulation.get_holder(variable).get_known_periods() + } + + assert set(simulation.input_variables) == holding_values + assert set(PRE_RESPONSE_INPUTS.values()) | {"state_code"} <= holding_values + assert not (set(PRE_RESPONSE_INPUTS) | {"state_code_str"}) & holding_values + + +def _ohio_parent_dataset(wages): + """One Ohio parent of two children, as an entity-level dataset.""" + person = pd.DataFrame( + { + "person_id": [1, 2, 3], + "person_household_id": [1, 1, 1], + "person_tax_unit_id": [1, 1, 1], + "person_spm_unit_id": [1, 1, 1], + "person_family_id": [1, 1, 1], + "person_marital_unit_id": [1, 2, 3], + "age": [30.0, 4.0, 8.0], + "employment_income": [float(wages), 0.0, 0.0], + "weekly_hours_worked": [30.0, 0.0, 0.0], + "long_term_capital_gains": [500.0, 0.0, 0.0], + } + ) + household = pd.DataFrame( + { + "household_id": [1], + "state_fips": np.asarray([39], dtype="int64"), + "household_weight": [1.0], + } + ) + return USSingleYearDataset( + person=person, + household=household, + tax_unit=pd.DataFrame({"tax_unit_id": [1]}), + spm_unit=pd.DataFrame({"spm_unit_id": [1]}), + family=pd.DataFrame({"family_id": [1]}), + marital_unit=pd.DataFrame({"marital_unit_id": [1, 2, 3]}), + time_period=YEAR, + ) + + +def test_microsimulation_marginal_tax_rate_matches_fresh_microsimulations(): + """Invariants 1 and 4 for dataset-backed simulations. + + Before the fix, ``Microsimulation`` listed the moved inputs, so its branch + kept the parent's pre-response wages and raised only the + ``employment_income`` aggregate: Ohio Works First, which counts + pre-response wages, never saw the raise. + """ + simulation = Microsimulation(dataset=_ohio_parent_dataset(6_000)) + raised = Microsimulation(dataset=_ohio_parent_dataset(7_000)) + holding_values = { + variable + for variable in simulation.tax_benefit_system.variables + if simulation.get_holder(variable).get_known_periods() + } + assert set(simulation.input_variables) == holding_values + assert "long_term_capital_gains" not in simulation.input_variables + + base = simulation.calculate("household_net_income", YEAR).values[0] + alt = raised.calculate("household_net_income", YEAR).values[0] + assert ( + simulation.calculate("tanf", YEAR).values[0] + > raised.calculate("tanf", YEAR).values[0] + ) + + rate = simulation.calculate("marginal_tax_rate", YEAR).values + + assert rate[0] == pytest.approx(1 - (alt - base) / DELTA, abs=TOLERANCE) + + +def test_marginal_rates_leave_the_simulation_unchanged(): + """Invariant 3: branches neither leak into nor linger on the simulation.""" + simulation = Simulation(situation=situation([OHIO_PARENT])) + watched = ( + ("household_net_income", YEAR), + ("tanf", f"{YEAR}-06"), + ("employment_income_before_lsr", f"{YEAR}-06"), + ) + before = [simulation.calculate(*key).copy() for key in watched] + branches = set(simulation.branches) + + for rate in (*EARNINGS_RATES, "marginal_tax_rate_on_capital_gains"): + simulation.calculate(rate, YEAR) + + # The model keeps its own branches, such as itemizing; rates add none. + assert set(simulation.branches) == branches + for key, value in zip(watched, before): + np.testing.assert_array_equal(simulation.calculate(*key), value) + + +def test_state_code_str_household_keeps_its_state_in_the_branch(): + """A household located only by state_code_str keeps that state.""" + base_situation = situation([OHIO_PARENT], geography="state_code_str") + simulation = Simulation(situation=base_situation) + + assert_rates_match( + simulation, + fresh_earnings_rates(base_situation, simulation), + ("marginal_tax_rate", "marginal_tax_rate_including_health_benefits"), + ) + + +def test_capital_gains_rate_applies_preferential_rates(): + """The capital gains rate raises long-term gains, which are taxed at 15%. + + A Texas filer with $60,000 of wages and $15,000 of long-term gains has + $43,900 of ordinary taxable income after the $16,100 standard deduction, + so the gains fill the 0% bracket to $49,450 and the rest is taxed at 15%. + Texas has no income tax and the filer is below the net investment income + tax threshold, so an extra dollar of gains costs 15 cents. + """ + base_situation = situation( + [ + ( + "TX", + { + "filer": { + "age": 45, + "employment_income": 60_000, + "long_term_capital_gains": 15_000, + } + }, + [], + ) + ] + ) + simulation = Simulation(situation=base_situation) + + rate = simulation.calculate("marginal_tax_rate_on_capital_gains", YEAR) + + assert rate[0] == pytest.approx(0.15, abs=TOLERANCE) + assert rate[0] == pytest.approx( + fresh_capital_gains_rate(base_situation, simulation)[0], abs=TOLERANCE + ) + + +def random_households(seed, count): + """Households drawn to reach every moved input and many programs.""" + rng = np.random.default_rng(seed) + households = [] + for _ in range(count): + adults = {} + for name in ("head", "spouse")[: rng.integers(1, 3)]: + age = int(rng.integers(19, 80)) + inputs = {"age": age} + wage_band = rng.choice([0, 12_000, 35_000, 120_000]) + if wage_band: + inputs["employment_income"] = float(rng.uniform(0.2, 1.0) * wage_band) + if rng.random() < 0.7: + inputs["weekly_hours_worked"] = float(rng.uniform(5, 50)) + if rng.random() < 0.25: + inputs["self_employment_income"] = float(rng.uniform(-5_000, 40_000)) + if rng.random() < 0.1: + inputs["sstb_self_employment_income"] = float(rng.uniform(0, 60_000)) + if rng.random() < 0.2: + inputs["long_term_capital_gains"] = float(rng.uniform(0, 50_000)) + if age >= 62 and rng.random() < 0.7: + inputs["social_security_retirement"] = float(rng.uniform(5_000, 30_000)) + adults[name] = inputs + children = [int(age) for age in rng.integers(0, 18, rng.integers(0, 4))] + households.append((str(rng.choice(STATES)), adults, children)) + return households + + +@pytest.mark.parametrize("seed", [20260928]) +def test_marginal_rates_match_fresh_simulations_for_random_households(seed): + """Invariants 1 and 2 over randomly drawn households in every state.""" + base_situation = situation(random_households(seed, 24)) + simulation = Simulation(situation=base_situation) + expected = fresh_earnings_rates(base_situation, simulation) + expected["marginal_tax_rate_on_capital_gains"] = fresh_capital_gains_rate( + base_situation, simulation + ) + + assert_rates_match(simulation, expected, expected) diff --git a/policyengine_us/tests/policy/baseline/household/marginal_tax_rate.yaml b/policyengine_us/tests/policy/baseline/household/marginal_tax_rate.yaml index c6fc58a09cb..9a93492d0ac 100644 --- a/policyengine_us/tests/policy/baseline/household/marginal_tax_rate.yaml +++ b/policyengine_us/tests/policy/baseline/household/marginal_tax_rate.yaml @@ -88,3 +88,177 @@ state_fips: 48 # TX output: marginal_tax_rate: [0.3165, -0.0765] + +# A single Ohio parent of children aged 4 and 8. The marginal tax rate must +# equal the change in household net income between the two fresh simulations +# above it: 1 - (alt - base) / 1_000. Ohio Works First counts +# employment_income_before_lsr as earned income, so the rate's branch must +# raise that input, not only the employment_income aggregate. + +- name: Ohio parent earning $6,000 while on Ohio Works First, OH-6000-base + absolute_error_margin: 0.01 + period: 2026 + input: + people: + parent: + age: 30 + employment_income_before_lsr: 6_000 + child1: + age: 4 + is_tax_unit_dependent: true + child2: + age: 8 + is_tax_unit_dependent: true + tax_units: + tax_unit: + members: [parent, child1, child2] + spm_units: + spm_unit: + members: [parent, child1, child2] + households: + household: + members: [parent, child1, child2] + state_code: OH + output: + tanf: 5_974.60 + household_net_income: 23_317.57 + +- name: Ohio parent earning $7,000 while on Ohio Works First, OH-6000-alt + absolute_error_margin: 0.01 + period: 2026 + input: + people: + parent: + age: 30 + employment_income_before_lsr: 7_000 + child1: + age: 4 + is_tax_unit_dependent: true + child2: + age: 8 + is_tax_unit_dependent: true + tax_units: + tax_unit: + members: [parent, child1, child2] + spm_units: + spm_unit: + members: [parent, child1, child2] + households: + household: + members: [parent, child1, child2] + state_code: OH + output: + tanf: 5_474.60 + household_net_income: 24_198.07 + +- name: Ohio Works First phase-out enters the marginal tax rate, OH-6000-mtr + absolute_error_margin: 0.0001 + period: 2026 + input: + people: + parent: + age: 30 + employment_income_before_lsr: 6_000 + child1: + age: 4 + is_tax_unit_dependent: true + child2: + age: 8 + is_tax_unit_dependent: true + tax_units: + tax_unit: + members: [parent, child1, child2] + spm_units: + spm_unit: + members: [parent, child1, child2] + households: + household: + members: [parent, child1, child2] + state_code: OH + output: + # 1 - (24_198.07 - 23_317.57) / 1_000 + marginal_tax_rate: [0.1195, 0, 0] + +- name: Ohio parent earning $28,000, OH-28000-base + absolute_error_margin: 0.01 + period: 2026 + input: + people: + parent: + age: 30 + employment_income_before_lsr: 28_000 + child1: + age: 4 + is_tax_unit_dependent: true + child2: + age: 8 + is_tax_unit_dependent: true + tax_units: + tax_unit: + members: [parent, child1, child2] + spm_units: + spm_unit: + members: [parent, child1, child2] + households: + household: + members: [parent, child1, child2] + state_code: OH + output: + tanf: 0 + household_net_income: 41_096.40 + +- name: Ohio parent earning $29,000, OH-28000-alt + absolute_error_margin: 0.01 + period: 2026 + input: + people: + parent: + age: 30 + employment_income_before_lsr: 29_000 + child1: + age: 4 + is_tax_unit_dependent: true + child2: + age: 8 + is_tax_unit_dependent: true + tax_units: + tax_unit: + members: [parent, child1, child2] + spm_units: + spm_unit: + members: [parent, child1, child2] + households: + household: + members: [parent, child1, child2] + state_code: OH + output: + tanf: 0 + household_net_income: 41_569.30 + +- name: Ohio parent above the Ohio Works First limit, OH-28000-mtr + absolute_error_margin: 0.0001 + period: 2026 + input: + people: + parent: + age: 30 + employment_income_before_lsr: 28_000 + child1: + age: 4 + is_tax_unit_dependent: true + child2: + age: 8 + is_tax_unit_dependent: true + tax_units: + tax_unit: + members: [parent, child1, child2] + spm_units: + spm_unit: + members: [parent, child1, child2] + households: + household: + members: [parent, child1, child2] + state_code: OH + output: + # 1 - (41_569.30 - 41_096.40) / 1_000 + marginal_tax_rate: [0.5271, 0, 0] diff --git a/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py b/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py index 34503291a64..7c817149366 100644 --- a/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py +++ b/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py @@ -11,6 +11,16 @@ "sstb_self_employment_income_behavioral_response", "capital_gains_behavioral_response", ) +# Situations and datasets name these aggregates, but each aggregate adds a +# behavioral response to a pre-response input. Simulations store the value on +# the pre-response input and derive the aggregate from it. +PRE_RESPONSE_INPUTS = { + "employment_income": "employment_income_before_lsr", + "self_employment_income": "self_employment_income_before_lsr", + "sstb_self_employment_income": "sstb_self_employment_income_before_lsr", + "weekly_hours_worked": "weekly_hours_worked_before_lsr", + "long_term_capital_gains": "long_term_capital_gains_before_response", +} BEHAVIORAL_RESPONSE_INPUT_VARIABLES = ( "employment_income_before_lsr", "self_employment_income_before_lsr", diff --git a/policyengine_us/variables/gov/simulation/capital_gains_responses.py b/policyengine_us/variables/gov/simulation/capital_gains_responses.py index d39a372f76a..d23a1cb86a3 100644 --- a/policyengine_us/variables/gov/simulation/capital_gains_responses.py +++ b/policyengine_us/variables/gov/simulation/capital_gains_responses.py @@ -3,6 +3,9 @@ calculate_relative_capital_gains_mtr_change, get_behavioral_response_measurements, ) +from policyengine_us.variables.household.marginal_tax_rate_helpers import ( + create_perturbed_branch, +) class relative_capital_gains_mtr_change(Variable): @@ -98,28 +101,23 @@ def formula(person, period, parameters): # pragma: no cover simulation = person.simulation DELTA = 1_000 adult_index_values = person("adult_index_cg", period) + household_net_income = person.household("household_net_income", period) for adult_index in [1, 2]: - alt_simulation = simulation.get_branch(f"adult_{adult_index}_cg_rise") + branch_name = f"adult_{adult_index}_cg_rise" mask = adult_index_values == adult_index - for variable in simulation.tax_benefit_system.variables: - variable_data = simulation.tax_benefit_system.variables[variable] - if ( - variable not in simulation.input_variables - and not variable_data.is_input_variable() - ): - alt_simulation.delete_arrays(variable) - alt_simulation.set_input( - "capital_gains", + # Raise long-term gains, the gains the behavioral response scales, + # so the preferential rates apply to the increase. + alt_simulation = create_perturbed_branch( + simulation, period, - person("capital_gains", period) + mask * DELTA, + branch_name, + {"long_term_capital_gains": mask * DELTA}, ) - alt_person = alt_simulation.person - household_net_income = person.household("household_net_income", period) - household_net_income_higher_earnings = alt_person.household( + household_net_income_higher_gains = alt_simulation.person.household( "household_net_income", period ) - increase = household_net_income_higher_earnings - household_net_income + increase = household_net_income_higher_gains - household_net_income mtr_values += where(mask, 1 - increase / DELTA, 0) - del simulation.branches[f"adult_{adult_index}_cg_rise"] + del simulation.branches[branch_name] return mtr_values diff --git a/policyengine_us/variables/household/marginal_tax_rate.py b/policyengine_us/variables/household/marginal_tax_rate.py index a0abd052d58..94af0487932 100644 --- a/policyengine_us/variables/household/marginal_tax_rate.py +++ b/policyengine_us/variables/household/marginal_tax_rate.py @@ -1,4 +1,7 @@ from policyengine_us.model_api import * +from policyengine_us.variables.household.marginal_tax_rate_helpers import ( + marginal_earnings_change, +) class marginal_tax_rate(Variable): @@ -12,63 +15,14 @@ class marginal_tax_rate(Variable): unit = "/1" def formula(person, period, parameters): - netinc_base = person.household("household_net_income", period) - delta = parameters(period).simulation.marginal_tax_rate_delta - adult_count = parameters(period).simulation.marginal_tax_rate_adults - sim = person.simulation - mtr_values = np.zeros(person.count, dtype=np.float32) - adult_indexes = person("adult_earnings_index", period) - employment_income = person("employment_income", period) - self_employment_income = person("self_employment_income", period) - sstb_self_employment_income = person("sstb_self_employment_income", period) - emp_self_emp_ratio = person("emp_self_emp_ratio", period) - positive_self_employment_income = max_(0, self_employment_income) - positive_sstb_self_employment_income = max_(0, sstb_self_employment_income) - positive_self_employment_total = ( - positive_self_employment_income + positive_sstb_self_employment_income + increase, measured = marginal_earnings_change( + person, + period, + parameters, + lambda population: population.household("household_net_income", period), + "mtr", ) - non_sstb_share = where( - positive_self_employment_total > 0, - positive_self_employment_income / positive_self_employment_total, - 1, - ) - sstb_share = where( - positive_self_employment_total > 0, - positive_sstb_self_employment_income / positive_self_employment_total, - 0, - ) - - for adult_index in range(1, 1 + adult_count): - alt_sim = sim.get_branch(f"mtr_for_adult_{adult_index}") - for variable in sim.tax_benefit_system.variables: - if ( - variable not in sim.input_variables - or variable == "employment_income" - ): - alt_sim.delete_arrays(variable) - mask = adult_index == adult_indexes - alt_sim.set_input( - "employment_income", - period, - employment_income + mask * delta * emp_self_emp_ratio, - ) - self_employment_delta = mask * delta * (1 - emp_self_emp_ratio) - alt_sim.set_input( - "self_employment_income", - period, - self_employment_income + self_employment_delta * non_sstb_share, - ) - alt_sim.set_input( - "sstb_self_employment_income", - period, - sstb_self_employment_income + self_employment_delta * sstb_share, - ) - alt_person = alt_sim.person - netinc_alt = alt_person.household("household_net_income", period) - increase = netinc_alt - netinc_base - mtr_values += where(mask, 1 - increase / delta, 0) - del sim.branches[f"mtr_for_adult_{adult_index}"] - return mtr_values + return where(measured, 1 - increase, 0) class adult_index(Variable): diff --git a/policyengine_us/variables/household/marginal_tax_rate_helpers.py b/policyengine_us/variables/household/marginal_tax_rate_helpers.py index 3ac38a5a532..360da33407d 100644 --- a/policyengine_us/variables/household/marginal_tax_rate_helpers.py +++ b/policyengine_us/variables/household/marginal_tax_rate_helpers.py @@ -1,59 +1,140 @@ from policyengine_us.model_api import * +from policyengine_us.variables.gov.simulation.behavioral_response_measurements import ( + PRE_RESPONSE_INPUTS, +) -def compute_component_mtr(person, period, parameters, tax_variable, branch_prefix): - """Compute the marginal tax rate for a specific tax component. +def create_perturbed_branch(simulation, period, branch_name, increments): + """Branch in which each variable in ``increments`` rises by its increment. + + The branch keeps the simulation's inputs and recomputes everything else, + so it matches a fresh simulation of the same households with the raised + inputs. Simulations store an aggregate such as ``employment_income`` on + its pre-response input (``employment_income_before_lsr``), and some + programs read that input directly; TANF's earned income does. Raising the + aggregate itself would hide the increase from those programs, so the + increase goes to the variable that holds the input. + + Args: + simulation: The simulation to branch from. + period: The period whose inputs rise. + branch_name: Name of the branch; the caller deletes it when done. + increments: Person-level increase for each raised variable. + + Returns: + The branch simulation. + """ + input_variables = set(simulation.input_variables) + branch = simulation.get_branch(branch_name) + for variable in simulation.tax_benefit_system.variables: + if variable not in input_variables: + branch.delete_arrays(variable) + for variable, increment in increments.items(): + if variable not in input_variables: + variable = PRE_RESPONSE_INPUTS.get(variable, variable) + value = simulation.calculate(variable, period) + increment + # The branch still holds the base values cached for the months of + # the period, which monthly programs such as TANF read. + branch.delete_arrays(variable, period) + branch.set_input(variable, period, value) + return branch + + +def earnings_increments(person, period, increase): + """Split an earnings increase across wages and self-employment income. - Uses the same counterfactual branch pattern as marginal_tax_rate: - perturbs earnings by delta and measures the change in the given - tax_unit-level tax variable, aggregated to the household level - for consistency with how marginal_tax_rate uses household_net_income. + Follows each person's current mix of wages, non-SSTB and SSTB + self-employment income, putting the increase on wages when the person + has no positive earnings. + """ + self_employment_income = max_(0, person("self_employment_income", period)) + sstb_self_employment_income = max_(0, person("sstb_self_employment_income", period)) + total_self_employment_income = self_employment_income + sstb_self_employment_income + non_sstb_share = np.divide( + self_employment_income, + total_self_employment_income, + out=np.ones_like(total_self_employment_income, dtype=np.float64), + where=total_self_employment_income > 0, + ) + wage_share = person("emp_self_emp_ratio", period) + self_employment_increase = increase * (1 - wage_share) + return { + "employment_income": increase * wage_share, + "self_employment_income": self_employment_increase * non_sstb_share, + "sstb_self_employment_income": self_employment_increase * (1 - non_sstb_share), + } + + +def marginal_earnings_change(person, period, parameters, measure, branch_prefix): + """Change in ``measure`` per dollar of additional earnings, for each adult. + + For each of the first ``simulation.marginal_tax_rate_adults`` adults in + every household, ranked by market income, a branch raises that adult's + earnings by ``simulation.marginal_tax_rate_delta`` and records how much + ``measure`` changes per dollar. Args: person: The person entity. period: The simulation period. parameters: The parameter tree. - tax_variable: Name of the tax_unit variable to measure - (e.g. "income_tax", "state_income_tax", "employee_payroll_tax"). + measure: Function of a person population returning a person-level + array, e.g. the person's household net income. branch_prefix: Unique prefix for branch names to avoid conflicts. Returns: - Array of marginal tax rates per person. + A tuple of the change per dollar for each measured adult (zero for + everyone else) and a mask of the measured adults. """ - is_head = person("is_tax_unit_head", period) - base_tax = person.household.sum(person.tax_unit(tax_variable, period) * is_head) delta = parameters(period).simulation.marginal_tax_rate_delta adult_count = parameters(period).simulation.marginal_tax_rate_adults - sim = person.simulation - mtr_values = np.zeros(person.count, dtype=np.float32) + simulation = person.simulation + base = measure(person) adult_indexes = person("adult_earnings_index", period) - employment_income = person("employment_income", period) - self_employment_income = person("self_employment_income", period) - emp_self_emp_ratio = person("emp_self_emp_ratio", period) - + change = np.zeros(person.count, dtype=np.float32) + measured = np.zeros(person.count, dtype=bool) for adult_index in range(1, 1 + adult_count): branch_name = f"{branch_prefix}_for_adult_{adult_index}" - alt_sim = sim.get_branch(branch_name) - for variable in sim.tax_benefit_system.variables: - if variable not in sim.input_variables or variable == "employment_income": - alt_sim.delete_arrays(variable) mask = adult_index == adult_indexes - alt_sim.set_input( - "employment_income", - period, - employment_income + mask * delta * emp_self_emp_ratio, - ) - alt_sim.set_input( - "self_employment_income", + branch = create_perturbed_branch( + simulation, period, - self_employment_income + mask * delta * (1 - emp_self_emp_ratio), + branch_name, + earnings_increments(person, period, mask * delta), ) - alt_person = alt_sim.person - alt_is_head = alt_person("is_tax_unit_head", period) - alt_tax = alt_person.household.sum( - alt_person.tax_unit(tax_variable, period) * alt_is_head + alt = measure(branch.person) + change += where(mask, (alt - base) / delta, 0) + measured |= mask + del simulation.branches[branch_name] + return change, measured + + +def compute_component_mtr(person, period, parameters, tax_variable, branch_prefix): + """Compute the marginal tax rate for a specific tax component. + + Uses the same counterfactual branch as marginal_tax_rate: raises earnings + by delta and measures the change in the given tax_unit-level tax + variable, aggregated to the household level for consistency with how + marginal_tax_rate uses household_net_income. + + Args: + person: The person entity. + period: The simulation period. + parameters: The parameter tree. + tax_variable: Name of the tax_unit variable to measure + (e.g. "income_tax", "state_income_tax", "employee_payroll_tax"). + branch_prefix: Unique prefix for branch names to avoid conflicts. + + Returns: + Array of marginal tax rates per person. + """ + + def household_tax(population): + is_head = population("is_tax_unit_head", period) + return population.household.sum( + population.tax_unit(tax_variable, period) * is_head ) - increase = alt_tax - base_tax - mtr_values += where(mask, increase / delta, 0) - del sim.branches[branch_name] - return mtr_values + + change, _ = marginal_earnings_change( + person, period, parameters, household_tax, branch_prefix + ) + return change diff --git a/policyengine_us/variables/household/marginal_tax_rate_including_health_benefits.py b/policyengine_us/variables/household/marginal_tax_rate_including_health_benefits.py index 02fe206ca49..9a59c10c590 100644 --- a/policyengine_us/variables/household/marginal_tax_rate_including_health_benefits.py +++ b/policyengine_us/variables/household/marginal_tax_rate_including_health_benefits.py @@ -1,4 +1,7 @@ from policyengine_us.model_api import * +from policyengine_us.variables.household.marginal_tax_rate_helpers import ( + marginal_earnings_change, +) class marginal_tax_rate_including_health_benefits(Variable): @@ -12,64 +15,13 @@ class marginal_tax_rate_including_health_benefits(Variable): unit = "/1" def formula(person, period, parameters): - netinc_base = person.household( - "household_net_income_including_health_benefits", period - ) - delta = parameters(period).simulation.marginal_tax_rate_delta - adult_count = parameters(period).simulation.marginal_tax_rate_adults - sim = person.simulation - mtr_values = np.zeros(person.count, dtype=np.float32) - adult_indexes = person("adult_earnings_index", period) - employment_income = person("employment_income", period) - self_employment_income = person("self_employment_income", period) - sstb_self_employment_income = person("sstb_self_employment_income", period) - emp_self_emp_ratio = person("emp_self_emp_ratio", period) - positive_self_employment_income = max_(0, self_employment_income) - positive_sstb_self_employment_income = max_(0, sstb_self_employment_income) - positive_self_employment_total = ( - positive_self_employment_income + positive_sstb_self_employment_income - ) - non_sstb_share = where( - positive_self_employment_total > 0, - positive_self_employment_income / positive_self_employment_total, - 1, - ) - sstb_share = where( - positive_self_employment_total > 0, - positive_sstb_self_employment_income / positive_self_employment_total, - 0, - ) - - for adult_index in range(1, 1 + adult_count): - alt_sim = sim.get_branch(f"mtr_for_adult_{adult_index}") - for variable in sim.tax_benefit_system.variables: - if ( - variable not in sim.input_variables - or variable == "employment_income" - ): - alt_sim.delete_arrays(variable) - mask = adult_index == adult_indexes - alt_sim.set_input( - "employment_income", - period, - employment_income + mask * delta * emp_self_emp_ratio, - ) - self_employment_delta = mask * delta * (1 - emp_self_emp_ratio) - alt_sim.set_input( - "self_employment_income", - period, - self_employment_income + self_employment_delta * non_sstb_share, - ) - alt_sim.set_input( - "sstb_self_employment_income", - period, - sstb_self_employment_income + self_employment_delta * sstb_share, - ) - alt_person = alt_sim.person - netinc_alt = alt_person.household( + increase, measured = marginal_earnings_change( + person, + period, + parameters, + lambda population: population.household( "household_net_income_including_health_benefits", period - ) - increase = netinc_alt - netinc_base - mtr_values += where(mask, 1 - increase / delta, 0) - del sim.branches[f"mtr_for_adult_{adult_index}"] - return mtr_values + ), + "mtr_including_health_benefits", + ) + return where(measured, 1 - increase, 0) From 792e7d4dbb074a90b5f238efe1b627652bc9e922 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 06:06:07 -0400 Subject: [PATCH 2/6] Hold behavioral responses in rate branches; branch the baseline after moving inputs Rates measured in a reformed simulation with labor supply or capital gains elasticities now hold the responses at their simulated values. Since the branches raise pre-response inputs (employment_income_before_lsr, long_term_capital_gains_before_response), recomputing the responses there scaled them with the raise: a Texas worker's $1,000 raise became $971 of wages (rate 0.3166, fresh 0.2965), and a 20% capital gains rate read 0.3307. create_perturbed_branch, which every earnings and capital gains rate uses, now sets the three earnings responses, the capital gains response and the weekly hours response to their values in the simulation; neutralized ones, as in the behavioral response measurement branches, stay zero. Core branches a reformed simulation's baseline while constructing it, before the wrappers move inputs. That branch kept employment_income and weekly_hours_worked as inputs with empty pre-response variables and the stale input_variables, so its TANF counted no earnings and the baseline rate that labor supply responses compare against erased pre-response wages: an Ohio Works First parent got a -$416 response to a top-rate reform that does not touch them. Both wrappers now branch the baseline again after the moves, as core's subsample does. Adds tests for both (invariants 5 and 6) and corrects the #9653 comment on the component helper's adult loop. Co-Authored-By: Claude Opus 5.5 --- .../fix-mtr-branch-divergence.fixed.md | 2 +- policyengine_us/system.py | 28 ++ .../tests/core/test_marginal_rate_branches.py | 310 ++++++++++++++++-- .../baseline/household/marginal_tax_rate.yaml | 3 +- .../household/marginal_tax_rate_helpers.py | 28 +- 5 files changed, 334 insertions(+), 37 deletions(-) diff --git a/changelog.d/fix-mtr-branch-divergence.fixed.md b/changelog.d/fix-mtr-branch-divergence.fixed.md index 973c462cd4a..03bb8d35f0d 100644 --- a/changelog.d/fix-mtr-branch-divergence.fixed.md +++ b/changelog.d/fix-mtr-branch-divergence.fixed.md @@ -1 +1 @@ -Make branch-measured marginal tax rates match fresh re-simulation. Household simulations now list the inputs they move onto pre-response variables (`employment_income_before_lsr`, `weekly_hours_worked_before_lsr`, `long_term_capital_gains_before_response`, `state_code`) in `input_variables`, so rate branches no longer erase them. The branches also raise those inputs rather than the aggregates, so programs that read pre-response income, such as TANF, see the raise. `marginal_tax_rate_on_capital_gains` now raises long-term gains, so preferential rates apply to the increase. +Make branch-measured marginal tax rates match fresh re-simulation. Household simulations now list the inputs they move onto pre-response variables (`employment_income_before_lsr`, `weekly_hours_worked_before_lsr`, `long_term_capital_gains_before_response`, `state_code`) in `input_variables`, so rate branches no longer erase them. The branches also raise those inputs rather than the aggregates, so programs that read pre-response income, such as TANF, see the raise, and they hold a reform's behavioral responses at their simulated values. `marginal_tax_rate_on_capital_gains` now raises long-term gains, so preferential rates apply to the increase. A reformed simulation's baseline branch now has the moved inputs too, so the baseline rates that labor supply responses compare against no longer erase pre-response wages. diff --git a/policyengine_us/system.py b/policyengine_us/system.py index dbd0cb0d7bc..ae14e3e2246 100644 --- a/policyengine_us/system.py +++ b/policyengine_us/system.py @@ -219,6 +219,32 @@ def holds_input(variable): ] +def _rebuild_baseline_branch(simulation): + """Branch a reformed simulation's baseline again, after the moves. + + Core branches the baseline while it constructs a reformed simulation, + before the wrappers move inputs and backfill ``state_code``. The branch + therefore kept ``employment_income`` and ``weekly_hours_worked`` as + inputs, left ``employment_income_before_lsr`` and + ``weekly_hours_worked_before_lsr`` empty, and kept the stale + ``input_variables``. Its TANF counted no earnings, and the baseline + marginal tax rate that labor supply responses compare against erased + pre-response wages. Branching again from the simulation, as core's + ``subsample`` does, gives the baseline the same inputs as the reform. + """ + baseline = simulation.branches.pop("baseline", None) + if baseline is None: + return + # A branch copies its parent's attributes, and a baseline whose own + # baseline is set would compute behavioral responses. + simulation.baseline = None + rebuilt = simulation.get_branch("baseline") + rebuilt.trace = simulation.trace + rebuilt.tracer = simulation.tracer + rebuilt.tax_benefit_system = baseline.tax_benefit_system + simulation.baseline = rebuilt + + def _backfill_state_code_from_str(simulation): """Backfill the ``state_code`` enum from a ``state_code_str`` input. @@ -289,6 +315,7 @@ def __init__(self, *args, **kwargs): # Geography backfill: state_code_str-only input -> state_code enum. _backfill_state_code_from_str(self) _sync_input_variables(self) + _rebuild_baseline_branch(self) def _download_or_explain(dataset_str, download): @@ -515,6 +542,7 @@ def __init__(self, *args, **kwargs): # explicitly supply state_code_str. _backfill_state_code_from_str(self) _sync_input_variables(self) + _rebuild_baseline_branch(self) class IndividualSim(CoreIndividualSim): # Deprecated diff --git a/policyengine_us/tests/core/test_marginal_rate_branches.py b/policyengine_us/tests/core/test_marginal_rate_branches.py index d41bb1ff0e4..ac3a9b80aae 100644 --- a/policyengine_us/tests/core/test_marginal_rate_branches.py +++ b/policyengine_us/tests/core/test_marginal_rate_branches.py @@ -16,6 +16,12 @@ were. 4. Right after construction, ``input_variables`` lists exactly the variables that hold values. +5. In a reformed simulation with behavioral responses, the rates hold the + responses at their simulated values: invariant 1 holds against fresh + simulations that raise the pre-response input and keep every response. +6. A reformed simulation's baseline branch, and the baseline and reform + branches that measure behavioral responses, equal fresh simulations of the + baseline and of the reform without responses. Invariant 4 is what the branches rely on. The ``Simulation`` and ``Microsimulation`` wrappers move ``employment_income`` and similar inputs onto @@ -34,12 +40,18 @@ import numpy as np import pandas as pd import pytest +from policyengine_core.periods import period as make_period +from policyengine_core.reforms import Reform from policyengine_us import Microsimulation, Simulation from policyengine_us.data.dataset_schema import USSingleYearDataset from policyengine_us.variables.gov.simulation.behavioral_response_measurements import ( + BEHAVIORAL_RESPONSE_CACHE_ATTR, PRE_RESPONSE_INPUTS, ) +from policyengine_us.variables.household.marginal_tax_rate_helpers import ( + HELD_BEHAVIORAL_RESPONSE_VARIABLES, +) YEAR = 2026 DELTA = 1_000 @@ -60,6 +72,27 @@ "state_marginal_tax_rate": "state_income_tax", "fica_marginal_tax_rate": "employee_payroll_tax", } +DATE_RANGE = "2020-01-01.2100-12-31" +# The reforms of tests/microsimulation/test_lsr_cg_interaction.py: a top +# rate 10 points higher, CBO's labor supply elasticities and a capital gains +# realization elasticity. +TOP_RATE = {"gov.irs.income.bracket.rates.7": {DATE_RANGE: 0.47}} +CBO_LABOR_SUPPLY = { + "gov.simulation.labor_supply_responses.elasticities.income": {DATE_RANGE: -0.05}, + "gov.simulation.labor_supply_responses.elasticities.substitution.all": { + DATE_RANGE: 0.25 + }, +} +CAPITAL_GAINS_RESPONSE = { + "gov.simulation.capital_gains_responses.elasticity": {DATE_RANGE: -0.62} +} + + +def reform(*parameter_changes): + changes = {} + for change in parameter_changes: + changes.update(change) + return Reform.from_dict(changes, country_id="us") def situation(households, geography="state_code"): @@ -114,19 +147,27 @@ def situation(households, geography="state_code"): return result +EARNINGS = ( + "employment_income", + "self_employment_income", + "sstb_self_employment_income", +) + + def raise_earnings(base_situation, simulation, adult_index): """Raise the earnings of each household's ``adult_index``-th earner. Splits ``DELTA`` across wages, non-SSTB and SSTB self-employment income in proportion to the person's positive earnings, wages when there are none; - the split the rates document. + the split the rates document. The situation's earnings are the + pre-response inputs, so the raise goes on those. """ ranks = simulation.calculate("adult_earnings_index", YEAR) - wages = simulation.calculate("employment_income", YEAR) - self_employment = simulation.calculate("self_employment_income", YEAR) - sstb = simulation.calculate("sstb_self_employment_income", YEAR) - positive = [np.maximum(x, 0) for x in (wages, self_employment, sstb)] + positive = [np.maximum(simulation.calculate(x, YEAR), 0) for x in EARNINGS] total = sum(positive) + pre_response = [ + simulation.calculate(PRE_RESPONSE_INPUTS[x], YEAR) for x in EARNINGS + ] raised = copy.deepcopy(base_situation) for i, person in enumerate(base_situation["people"]): if ranks[i] != adult_index: @@ -134,22 +175,14 @@ def raise_earnings(base_situation, simulation, adult_index): shares = ( [x[i] / total[i] for x in positive] if total[i] > 0 else [1.0, 0.0, 0.0] ) - for variable, base, share in zip( - ( - "employment_income", - "self_employment_income", - "sstb_self_employment_income", - ), - (wages, self_employment, sstb), - shares, - ): + for variable, base, share in zip(EARNINGS, pre_response, shares): raised["people"][person][variable] = {YEAR: float(base[i] + DELTA * share)} return raised, ranks == adult_index def raise_long_term_gains(base_situation, simulation, adult_index): ranks = simulation.calculate("adult_index_cg", YEAR) - gains = simulation.calculate("long_term_capital_gains", YEAR) + gains = simulation.calculate(PRE_RESPONSE_INPUTS["long_term_capital_gains"], YEAR) raised = copy.deepcopy(base_situation) for i, person in enumerate(base_situation["people"]): if ranks[i] == adult_index: @@ -159,6 +192,31 @@ def raise_long_term_gains(base_situation, simulation, adult_index): return raised, ranks == adult_index +def hold_responses(raised_situation, simulation): + """Give every person the behavioral responses they have in ``simulation``.""" + responses = { + variable: simulation.calculate(variable, YEAR) + for variable in HELD_BEHAVIORAL_RESPONSE_VARIABLES + } + for i, person in enumerate(raised_situation["people"]): + for variable, values in responses.items(): + raised_situation["people"][person][variable] = {YEAR: float(values[i])} + return raised_situation + + +def fresh_simulation(raised_situation, simulation, policy): + """A fresh simulation of the raised households under ``policy``. + + Under a reform, the fresh simulation keeps the responses ``simulation`` + has, so the difference measures the rate at fixed behavior. + """ + if policy is None: + return Simulation(situation=raised_situation) + return Simulation( + situation=hold_responses(raised_situation, simulation), reform=policy + ) + + def household_values(simulation, variable): return simulation.calculate(variable, YEAR, map_to="person") @@ -171,7 +229,7 @@ def head_taxes(simulation, variable): return household.project(household.sum(person_tax)) -def fresh_earnings_rates(base_situation, simulation): +def fresh_earnings_rates(base_situation, simulation, policy=None): """Each earnings rate as a difference between fresh simulations.""" adult_count = simulation.tax_benefit_system.parameters( YEAR @@ -186,7 +244,7 @@ def fresh_earnings_rates(base_situation, simulation): expected = {rate: np.zeros(len(base_situation["people"])) for rate in base} for adult_index in range(1, adult_count + 1): raised_situation, mask = raise_earnings(base_situation, simulation, adult_index) - raised = Simulation(situation=raised_situation) + raised = fresh_simulation(raised_situation, simulation, policy) alt = { "marginal_tax_rate": household_values(raised, "household_net_income"), "marginal_tax_rate_including_health_benefits": household_values( @@ -204,7 +262,7 @@ def fresh_earnings_rates(base_situation, simulation): return expected -def fresh_capital_gains_rate(base_situation, simulation): +def fresh_capital_gains_rate(base_situation, simulation, policy=None): base = household_values(simulation, "household_net_income") expected = np.zeros(len(base_situation["people"])) for adult_index in (1, 2): @@ -212,7 +270,8 @@ def fresh_capital_gains_rate(base_situation, simulation): base_situation, simulation, adult_index ) alt = household_values( - Simulation(situation=raised_situation), "household_net_income" + fresh_simulation(raised_situation, simulation, policy), + "household_net_income", ) expected = np.where(mask, 1 - (alt - base) / DELTA, expected) return expected @@ -233,6 +292,30 @@ def assert_rates_match(simulation, expected, rates): {"parent": {"age": 30, "employment_income": 28_000}}, [4, 8], ) +# The Ohio Works First recipient of _ohio_parent_dataset(6_000). +OHIO_WORKS_FIRST_PARENT = ( + "OH", + { + "parent": { + "age": 30, + "employment_income": 6_000, + "weekly_hours_worked": 30, + "long_term_capital_gains": 500, + } + }, + [4, 8], +) +TEXAS_GAINS_FILER = ( + "TX", + { + "filer": { + "age": 45, + "employment_income": 60_000, + "long_term_capital_gains": 15_000, + } + }, + [], +) def test_ohio_parent_marginal_tax_rate_matches_fresh_simulations(): @@ -391,21 +474,7 @@ def test_capital_gains_rate_applies_preferential_rates(): Texas has no income tax and the filer is below the net investment income tax threshold, so an extra dollar of gains costs 15 cents. """ - base_situation = situation( - [ - ( - "TX", - { - "filer": { - "age": 45, - "employment_income": 60_000, - "long_term_capital_gains": 15_000, - } - }, - [], - ) - ] - ) + base_situation = situation([TEXAS_GAINS_FILER]) simulation = Simulation(situation=base_situation) rate = simulation.calculate("marginal_tax_rate_on_capital_gains", YEAR) @@ -455,3 +524,176 @@ def test_marginal_rates_match_fresh_simulations_for_random_households(seed): ) assert_rates_match(simulation, expected, expected) + + +# Households whose taxes 10-point higher 12% and 22% rates raise, so CBO's +# elasticities give them labor supply responses. +RESPONDING_HOUSEHOLDS = [ + ("TX", {"worker": {"age": 40, "employment_income": 50_000}}, []), + ( + "CA", + { + "head": { + "age": 45, + "employment_income": 70_000, + "self_employment_income": 20_000, + "weekly_hours_worked": 40, + }, + "spouse": { + "age": 43, + "sstb_self_employment_income": 30_000, + "weekly_hours_worked": 25, + }, + }, + [10], + ), + ( + "NY", + { + "retiree": { + "age": 67, + "employment_income": 35_000, + "social_security_retirement": 18_000, + } + }, + [], + ), + OHIO_WORKS_FIRST_PARENT, +] +MIDDLE_RATES = { + "gov.irs.income.bracket.rates.2": {DATE_RANGE: 0.22}, + "gov.irs.income.bracket.rates.3": {DATE_RANGE: 0.32}, +} + + +def test_rates_in_a_dynamic_reform_hold_behavioral_responses(): + """Invariant 5 for the earnings rates. + + Under CBO's elasticities, 10-point higher 12% and 22% rates cut the Texas + worker's earnings by $1,429 (substitution elasticity 0.25 times their + after-tax wage change, times $50,000). Before the fix the branch + recomputed that response on the raised pre-response earnings, so $1,000 + more earnings became about $971 of wages and their rate read 0.3166 + where fresh simulations at fixed behavior give 0.2965. + """ + policy = reform(MIDDLE_RATES, CBO_LABOR_SUPPLY) + base_situation = situation(RESPONDING_HOUSEHOLDS) + simulation = Simulation(situation=base_situation, reform=policy) + responses = simulation.calculate("labor_supply_behavioral_response", YEAR) + assert (np.abs(responses) > 100).sum() >= 4 + + assert_rates_match( + simulation, + fresh_earnings_rates(base_situation, simulation, policy), + EARNINGS_RATES, + ) + + +def test_capital_gains_rate_in_a_dynamic_reform_holds_the_response(): + """Invariant 5 for the capital gains rate. + + A 20% rate on the Texas filer's gains (see + test_capital_gains_rate_applies_preferential_rates) and a -0.62 + elasticity cut their $15,000 of gains by + 15,000 * (exp(-0.62 * ln(0.20 / 0.15)) - 1) = -$2,450. Before the fix + the branch recomputed the response on $16,000, so $1,000 more gains + became $837 and the rate read 0.3307 instead of 0.20. + """ + policy = reform( + {"gov.irs.capital_gains.rates.2": {DATE_RANGE: 0.20}}, + CAPITAL_GAINS_RESPONSE, + ) + base_situation = situation([TEXAS_GAINS_FILER]) + simulation = Simulation(situation=base_situation, reform=policy) + response = simulation.calculate("capital_gains_behavioral_response", YEAR) + assert response[0] == pytest.approx(-2_450.4, abs=0.5) + + rate = simulation.calculate("marginal_tax_rate_on_capital_gains", YEAR) + + assert rate[0] == pytest.approx(0.20, abs=TOLERANCE) + assert rate[0] == pytest.approx( + fresh_capital_gains_rate(base_situation, simulation, policy)[0], + abs=TOLERANCE, + ) + + +def values(simulation, variable): + return np.asarray(simulation.calculate(variable, YEAR, map_to="person")) + + +@pytest.mark.parametrize("wrapper", [Simulation, Microsimulation]) +def test_reformed_simulations_branch_the_baseline_after_moving_inputs(wrapper): + """Invariant 6 for the baseline branch of either wrapper. + + Core branches the baseline while it constructs a reformed simulation, + before the wrappers move inputs. The branch kept the Ohio Works First + parent's wages on employment_income and none on + employment_income_before_lsr, so its TANF counted no earnings ($7,474.60 + instead of $5,974.60) and its marginal tax rate read -0.2335 instead of + 0.1195. + """ + policy = reform(TOP_RATE, CBO_LABOR_SUPPLY) + if wrapper is Simulation: + base_situation = situation([OHIO_WORKS_FIRST_PARENT]) + reformed = Simulation(situation=base_situation, reform=policy) + fresh = Simulation(situation=base_situation) + else: + reformed = Microsimulation(dataset=_ohio_parent_dataset(6_000), reform=policy) + fresh = Microsimulation(dataset=_ohio_parent_dataset(6_000)) + baseline = reformed.baseline + + assert baseline is reformed.branches["baseline"] + assert baseline.baseline is None + assert baseline.tax_benefit_system is not reformed.tax_benefit_system + assert set(baseline.input_variables) == set(fresh.input_variables) + for variable in ( + "employment_income_before_lsr", + "weekly_hours_worked_before_lsr", + "long_term_capital_gains", + "tanf", + "household_net_income", + "marginal_tax_rate", + ): + np.testing.assert_allclose( + values(baseline, variable), + values(fresh, variable), + atol=TOLERANCE, + err_msg=variable, + ) + + +def test_behavioral_response_measurements_match_fresh_simulations(): + """Invariant 6 for the branches that measure behavioral responses. + + The top-rate reform changes no tax or benefit of these households, so + none of them responds. Before the fix, the baseline measurement branch + erased the Ohio Works First parent's pre-response wages in its rate + branch. It measured a -0.2335 baseline rate against the reform's 0.1195 + and a $24,365 baseline net income against $23,318, so the parent + responded to a reform that does not touch them with $416 less earnings. + """ + policy = reform(TOP_RATE, CBO_LABOR_SUPPLY) + base_situation = situation(RESPONDING_HOUSEHOLDS) + simulation = Simulation(situation=base_situation, reform=policy) + response = simulation.calculate("labor_supply_behavioral_response", YEAR) + measurements = getattr(simulation, BEHAVIORAL_RESPONSE_CACHE_ATTR)[ + str(make_period(YEAR)) + ] + fresh = { + "baseline": Simulation(situation=base_situation), + "reform": Simulation(situation=base_situation, reform=reform(TOP_RATE)), + } + + for name, comparison in fresh.items(): + for measurement, variable in ( + ("mtr", "marginal_tax_rate"), + ("capital_gains_mtr", "marginal_tax_rate_on_capital_gains"), + ("net_income", "household_net_income"), + ): + np.testing.assert_allclose( + measurements[f"{name}_{measurement}"], + values(comparison, variable), + atol=TOLERANCE, + err_msg=f"{name}_{measurement}", + ) + np.testing.assert_allclose(response, 0, atol=TOLERANCE) diff --git a/policyengine_us/tests/policy/baseline/household/marginal_tax_rate.yaml b/policyengine_us/tests/policy/baseline/household/marginal_tax_rate.yaml index 6af586b7879..379312b84e6 100644 --- a/policyengine_us/tests/policy/baseline/household/marginal_tax_rate.yaml +++ b/policyengine_us/tests/policy/baseline/household/marginal_tax_rate.yaml @@ -93,7 +93,8 @@ # earnings less the $14,600 standard deduction, so the $90,000 earner is in the # 22% bracket and the $60,000 and $30,000 earners are in the 12% bracket; each # also pays 7.65% employee payroll tax. The federal and FICA components check the -# shared component helper, which has its own copy of the adult loop. +# component helper, which runs the same adult loop (marginal_earnings_change) as +# marginal_tax_rate. # The same household with the cap raised to 3 is in # tests/core/test_marginal_tax_rate_coverage.py, which avoids another cached # tax-benefit system in the batched YAML runner. diff --git a/policyengine_us/variables/household/marginal_tax_rate_helpers.py b/policyengine_us/variables/household/marginal_tax_rate_helpers.py index 360da33407d..8d3f75a9004 100644 --- a/policyengine_us/variables/household/marginal_tax_rate_helpers.py +++ b/policyengine_us/variables/household/marginal_tax_rate_helpers.py @@ -1,8 +1,17 @@ from policyengine_us.model_api import * from policyengine_us.variables.gov.simulation.behavioral_response_measurements import ( + NEUTRALIZED_BEHAVIORAL_RESPONSE_VARIABLES, PRE_RESPONSE_INPUTS, ) +# Behavioral responses to a reform, which each aggregate in +# PRE_RESPONSE_INPUTS adds to its pre-response input. A perturbed branch holds +# them at their values in the simulation. +HELD_BEHAVIORAL_RESPONSE_VARIABLES = ( + *NEUTRALIZED_BEHAVIORAL_RESPONSE_VARIABLES, + "weekly_hours_worked_behavioural_response", +) + def create_perturbed_branch(simulation, period, branch_name, increments): """Branch in which each variable in ``increments`` rises by its increment. @@ -15,6 +24,15 @@ def create_perturbed_branch(simulation, period, branch_name, increments): aggregate itself would hide the increase from those programs, so the increase goes to the variable that holds the input. + The branch holds behavioral responses to a reform at their values in the + simulation, so a rate measured in a reformed simulation with labor supply + or capital gains elasticities excludes any response to the increase + itself. Recomputing them would scale the response with the raised + pre-response input (a substitution response of -2.5% of earnings turns a + $1,000 raise into $975 of wages). Responses neutralized in the + simulation, as in the behavioral response measurement branches, are zero + in both. + Args: simulation: The simulation to branch from. period: The period whose inputs rise. @@ -25,10 +43,18 @@ def create_perturbed_branch(simulation, period, branch_name, increments): The branch simulation. """ input_variables = set(simulation.input_variables) + variables = simulation.tax_benefit_system.variables + responses = { + variable: simulation.person(variable, period) + for variable in HELD_BEHAVIORAL_RESPONSE_VARIABLES + if not variables[variable].is_neutralized + } branch = simulation.get_branch(branch_name) - for variable in simulation.tax_benefit_system.variables: + for variable in variables: if variable not in input_variables: branch.delete_arrays(variable) + for variable, response in responses.items(): + branch.set_input(variable, period, response) for variable, increment in increments.items(): if variable not in input_variables: variable = PRE_RESPONSE_INPUTS.get(variable, variable) From bd04830fdb3859c39cef47b99736403a63c023f6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 08:10:30 -0400 Subject: [PATCH 3/6] Rebind restored baseline holders and tracing --- .../fix-mtr-branch-divergence.fixed.md | 2 ++ policyengine_us/system.py | 4 ++++ policyengine_us/tests/core/test_spm_system.py | 20 +++++++++++++++++++ 3 files changed, 26 insertions(+) diff --git a/changelog.d/fix-mtr-branch-divergence.fixed.md b/changelog.d/fix-mtr-branch-divergence.fixed.md index 03bb8d35f0d..5a131c346bc 100644 --- a/changelog.d/fix-mtr-branch-divergence.fixed.md +++ b/changelog.d/fix-mtr-branch-divergence.fixed.md @@ -1 +1,3 @@ Make branch-measured marginal tax rates match fresh re-simulation. Household simulations now list the inputs they move onto pre-response variables (`employment_income_before_lsr`, `weekly_hours_worked_before_lsr`, `long_term_capital_gains_before_response`, `state_code`) in `input_variables`, so rate branches no longer erase them. The branches also raise those inputs rather than the aggregates, so programs that read pre-response income, such as TANF, see the raise, and they hold a reform's behavioral responses at their simulated values. `marginal_tax_rate_on_capital_gains` now raises long-term gains, so preferential rates apply to the increase. A reformed simulation's baseline branch now has the moved inputs too, so the baseline rates that labor supply responses compare against no longer erase pre-response wages. + +Restoring a rebuilt baseline's policy also rebinds its holders, removes reform-only inputs, and preserves tracing isolation, so neutralized reform variables cannot corrupt baseline comparisons. diff --git a/policyengine_us/system.py b/policyengine_us/system.py index ae14e3e2246..65b14abb929 100644 --- a/policyengine_us/system.py +++ b/policyengine_us/system.py @@ -242,6 +242,10 @@ def _rebuild_baseline_branch(simulation): rebuilt.trace = simulation.trace rebuilt.tracer = simulation.tracer rebuilt.tax_benefit_system = baseline.tax_benefit_system + # Cloning bound holders to reform variables. Restoring policy must also + # restore those bindings and remove inputs that exist only in the reform. + rebuilt._rebind_holders() + rebuilt._isolate_parameter_tracing() simulation.baseline = rebuilt diff --git a/policyengine_us/tests/core/test_spm_system.py b/policyengine_us/tests/core/test_spm_system.py index 1fa4f4240b8..50a99429447 100644 --- a/policyengine_us/tests/core/test_spm_system.py +++ b/policyengine_us/tests/core/test_spm_system.py @@ -134,6 +134,26 @@ def test_reform_baseline_and_clones_calculate_original_tax(simulation_type): assert clone.tax_benefit_system.simulation is clone +def test_rebuilt_baseline_preserves_tracing_isolation(): + situation = single_person_situation() + situation["people"]["person"]["employment_income"] = {2024: 50_000} + changed = Simulation(situation=situation, reform=UserReform, trace=True) + baseline = changed.baseline + policy = baseline.tax_benefit_system + + assert policy.simulation is baseline + assert baseline.trace + assert baseline.tracer is changed.tracer + assert policy.parameters is not changed.tax_benefit_system.parameters + assert policy.parameters.trace + assert policy.parameters.tracer is baseline.tracer + assert policy.parameters.branch_name == "baseline" + assert system.parameters.trace is False + assert system.parameters.tracer is None + np.testing.assert_array_equal(baseline.calculate("income_tax", 2024), 4_016) + np.testing.assert_array_equal(changed.calculate("income_tax", 2024), 0) + + def _parameter_fingerprint(system): """Digest every authored parameter value, so a shared-tree edit is visible.""" digest = hashlib.sha256() From 37bb63197728f0d481b0782609321872f7bd0a29 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 08:35:26 -0400 Subject: [PATCH 4/6] Neutralize hours while measuring behavioral responses --- .../fix-mtr-branch-divergence.fixed.md | 2 + .../test_combined_behavioral_responses.py | 132 ++++++++++++++++++ .../behavioral_response_measurements.py | 4 + .../household/marginal_tax_rate_helpers.py | 5 +- 4 files changed, 139 insertions(+), 4 deletions(-) create mode 100644 policyengine_us/tests/core/test_combined_behavioral_responses.py diff --git a/changelog.d/fix-mtr-branch-divergence.fixed.md b/changelog.d/fix-mtr-branch-divergence.fixed.md index 5a131c346bc..b6e6d2758c1 100644 --- a/changelog.d/fix-mtr-branch-divergence.fixed.md +++ b/changelog.d/fix-mtr-branch-divergence.fixed.md @@ -1,3 +1,5 @@ Make branch-measured marginal tax rates match fresh re-simulation. Household simulations now list the inputs they move onto pre-response variables (`employment_income_before_lsr`, `weekly_hours_worked_before_lsr`, `long_term_capital_gains_before_response`, `state_code`) in `input_variables`, so rate branches no longer erase them. The branches also raise those inputs rather than the aggregates, so programs that read pre-response income, such as TANF, see the raise, and they hold a reform's behavioral responses at their simulated values. `marginal_tax_rate_on_capital_gains` now raises long-term gains, so preferential rates apply to the increase. A reformed simulation's baseline branch now has the moved inputs too, so the baseline rates that labor supply responses compare against no longer erase pre-response wages. Restoring a rebuilt baseline's policy also rebinds its holders, removes reform-only inputs, and preserves tracing isolation, so neutralized reform variables cannot corrupt baseline comparisons. + +Behavioral-response measurements also neutralize hours responses, avoiding recursion when capital gains are queried before labor supply with both elasticities enabled. diff --git a/policyengine_us/tests/core/test_combined_behavioral_responses.py b/policyengine_us/tests/core/test_combined_behavioral_responses.py new file mode 100644 index 00000000000..10e8b367eab --- /dev/null +++ b/policyengine_us/tests/core/test_combined_behavioral_responses.py @@ -0,0 +1,132 @@ +"""Combined response measurements must work from a cold cache in either order.""" + +import numpy as np +import pytest +from policyengine_core.periods import period as make_period +from policyengine_core.reforms import Reform + +from policyengine_us import Simulation +from policyengine_us.variables.gov.simulation.behavioral_response_measurements import ( + BEHAVIORAL_RESPONSE_CACHE_ATTR, + BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, + BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, +) + +YEAR = 2026 +RESPONSES = ( + "capital_gains_behavioral_response", + "labor_supply_behavioral_response", +) + + +def combined_simulation(fixed_net_income=False): + date_range = "2020-01-01.2100-12-31" + policy = Reform.from_dict( + { + "gov.irs.income.bracket.rates.3": {date_range: 0.32}, + "gov.irs.capital_gains.rates.2": {date_range: 0.20}, + "gov.simulation.labor_supply_responses.elasticities.income": { + date_range: -0.05 + }, + "gov.simulation.labor_supply_responses.elasticities.substitution.all": { + date_range: 0.25 + }, + "gov.simulation.capital_gains_responses.elasticity": {date_range: -0.62}, + }, + country_id="us", + ) + situation = { + "people": { + "worker": { + "age": {YEAR: 40}, + "employment_income": {YEAR: 60_000}, + "long_term_capital_gains": {YEAR: 15_000}, + "weekly_hours_worked": {YEAR: 40}, + } + }, + **{ + entity: {"unit": {"members": ["worker"]}} + for entity in ("tax_units", "spm_units", "families", "marital_units") + }, + "households": { + "household": { + "members": ["worker"], + "state_code": {YEAR: "TX"}, + } + }, + } + if fixed_net_income: + situation["households"]["household"]["household_net_income"] = {YEAR: 60_000} + return Simulation(situation=situation, reform=policy) + + +@pytest.mark.parametrize( + "query_order", + [RESPONSES, RESPONSES[::-1]], + ids=["capital-gains-first", "labor-supply-first"], +) +def test_combined_responses_measure_from_a_cold_cache_in_either_order(query_order): + """CG-first measurement must not recursively calculate the hours response. + + Supplying household net income bypasses program calculations involving + hours, isolating the eager hours calculation in the marginal-rate helper. + The fixed income also makes both responses zero despite nonzero + elasticities: neither the measured income nor the measured marginal rates + change under the reform. + """ + simulation = combined_simulation(fixed_net_income=True) + period_key = str(make_period(YEAR)) + assert period_key not in getattr(simulation, BEHAVIORAL_RESPONSE_CACHE_ATTR, {}) + + for response in query_order: + np.testing.assert_array_equal(simulation.calculate(response, YEAR), [0]) + + np.testing.assert_array_equal( + simulation.calculate("weekly_hours_worked_behavioural_response", YEAR), [0] + ) + np.testing.assert_array_equal( + simulation.calculate("weekly_hours_worked", YEAR), [40] + ) + assert period_key in getattr(simulation, BEHAVIORAL_RESPONSE_CACHE_ATTR) + assert BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH not in simulation.branches + assert ( + BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH + not in simulation.baseline.branches + ) + + +def test_nonzero_combined_responses_are_independent_of_query_order(): + """Shared measurements give the same income and hours changes in both orders.""" + results = [] + measurements = [] + period_key = str(make_period(YEAR)) + for query_order in (RESPONSES, RESPONSES[::-1]): + simulation = combined_simulation() + assert period_key not in getattr(simulation, BEHAVIORAL_RESPONSE_CACHE_ATTR, {}) + response_values = { + variable: simulation.calculate(variable, YEAR) for variable in query_order + } + response_values["weekly_hours_worked_behavioural_response"] = ( + simulation.calculate("weekly_hours_worked_behavioural_response", YEAR) + ) + assert response_values["capital_gains_behavioral_response"][0] == pytest.approx( + -2_450.4, abs=0.5 + ) + assert abs(response_values["labor_supply_behavioral_response"][0]) > 1 + assert abs(response_values["weekly_hours_worked_behavioural_response"][0]) > 0 + results.append(response_values) + measurements.append( + getattr(simulation, BEHAVIORAL_RESPONSE_CACHE_ATTR)[period_key] + ) + + for variable in results[0]: + np.testing.assert_allclose( + results[0][variable], results[1][variable], atol=1e-3, err_msg=variable + ) + for measurement in measurements[0]: + np.testing.assert_allclose( + measurements[0][measurement], + measurements[1][measurement], + atol=1e-3, + err_msg=measurement, + ) diff --git a/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py b/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py index 7c817149366..c67059fcfd4 100644 --- a/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py +++ b/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py @@ -10,6 +10,10 @@ "self_employment_income_behavioral_response", "sstb_self_employment_income_behavioral_response", "capital_gains_behavioral_response", + # Marginal-rate branches also hold hours responses. Neutralizing hours + # here keeps measurements independent of the order responses are queried: + # a capital-gains-first query must not restart labor-supply measurements. + "weekly_hours_worked_behavioural_response", ) # Situations and datasets name these aggregates, but each aggregate adds a # behavioral response to a pre-response input. Simulations store the value on diff --git a/policyengine_us/variables/household/marginal_tax_rate_helpers.py b/policyengine_us/variables/household/marginal_tax_rate_helpers.py index 8d3f75a9004..e7059055428 100644 --- a/policyengine_us/variables/household/marginal_tax_rate_helpers.py +++ b/policyengine_us/variables/household/marginal_tax_rate_helpers.py @@ -7,10 +7,7 @@ # Behavioral responses to a reform, which each aggregate in # PRE_RESPONSE_INPUTS adds to its pre-response input. A perturbed branch holds # them at their values in the simulation. -HELD_BEHAVIORAL_RESPONSE_VARIABLES = ( - *NEUTRALIZED_BEHAVIORAL_RESPONSE_VARIABLES, - "weekly_hours_worked_behavioural_response", -) +HELD_BEHAVIORAL_RESPONSE_VARIABLES = NEUTRALIZED_BEHAVIORAL_RESPONSE_VARIABLES def create_perturbed_branch(simulation, period, branch_name, increments): From bc17da2fdfae1824170225571e7a925449594f66 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 15:55:07 -0400 Subject: [PATCH 5/6] Update Ohio net-income snapshots for 2026 school meals --- .../policy/baseline/household/marginal_tax_rate.yaml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/policyengine_us/tests/policy/baseline/household/marginal_tax_rate.yaml b/policyengine_us/tests/policy/baseline/household/marginal_tax_rate.yaml index 379312b84e6..a3bbd71491e 100644 --- a/policyengine_us/tests/policy/baseline/household/marginal_tax_rate.yaml +++ b/policyengine_us/tests/policy/baseline/household/marginal_tax_rate.yaml @@ -168,7 +168,7 @@ state_code: OH output: tanf: 5_974.60 - household_net_income: 23_317.57 + household_net_income: 23_294.93 - name: Ohio parent earning $7,000 while on Ohio Works First, OH-6000-alt absolute_error_margin: 0.01 @@ -196,7 +196,7 @@ state_code: OH output: tanf: 5_474.60 - household_net_income: 24_198.07 + household_net_income: 24_175.43 - name: Ohio Works First phase-out enters the marginal tax rate, OH-6000-mtr absolute_error_margin: 0.0001 @@ -223,7 +223,7 @@ members: [parent, child1, child2] state_code: OH output: - # 1 - (24_198.07 - 23_317.57) / 1_000 + # 1 - (24_175.43 - 23_294.93) / 1_000 marginal_tax_rate: [0.1195, 0, 0] - name: Ohio parent earning $28,000, OH-28000-base @@ -252,7 +252,7 @@ state_code: OH output: tanf: 0 - household_net_income: 41_096.40 + household_net_income: 41_073.76 - name: Ohio parent earning $29,000, OH-28000-alt absolute_error_margin: 0.01 @@ -280,7 +280,7 @@ state_code: OH output: tanf: 0 - household_net_income: 41_569.30 + household_net_income: 41_546.66 - name: Ohio parent above the Ohio Works First limit, OH-28000-mtr absolute_error_margin: 0.0001 @@ -307,5 +307,5 @@ members: [parent, child1, child2] state_code: OH output: - # 1 - (41_569.30 - 41_096.40) / 1_000 + # 1 - (41_546.66 - 41_073.76) / 1_000 marginal_tax_rate: [0.5271, 0, 0] From a729eecd02f9121f0a0a2712268d44e2af7fe952 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:08:49 -0400 Subject: [PATCH 6/6] Preserve explicit branch inputs in marginal-rate counterfactuals --- .../fix-mtr-branch-divergence.fixed.md | 2 + policyengine_us/spm.py | 19 ++ .../test_combined_behavioral_responses.py | 87 ++++++++ .../tests/core/test_marginal_rate_branches.py | 200 ++++++++++++++++++ .../household/marginal_tax_rate_helpers.py | 2 + 5 files changed, 310 insertions(+) diff --git a/changelog.d/fix-mtr-branch-divergence.fixed.md b/changelog.d/fix-mtr-branch-divergence.fixed.md index b6e6d2758c1..261d00b4b2c 100644 --- a/changelog.d/fix-mtr-branch-divergence.fixed.md +++ b/changelog.d/fix-mtr-branch-divergence.fixed.md @@ -3,3 +3,5 @@ Make branch-measured marginal tax rates match fresh re-simulation. Household sim Restoring a rebuilt baseline's policy also rebinds its holders, removes reform-only inputs, and preserves tracing isolation, so neutralized reform variables cannot corrupt baseline comparisons. Behavioral-response measurements also neutralize hours responses, avoiding recursion when capital gains are queried before labor supply with both elasticities enabled. + +Inputs explicitly set after construction, including identification inputs on a measuring branch, now remain inputs in its raised branch. This prevents earnings raises from restoring child or ACA premium tax credits denied in the base simulation, and also preserves branch inputs in health-inclusive rates and behavioral-response measurements. diff --git a/policyengine_us/spm.py b/policyengine_us/spm.py index 367cb91d880..bebca648fe7 100644 --- a/policyengine_us/spm.py +++ b/policyengine_us/spm.py @@ -970,6 +970,25 @@ def set_input(self, variable_name, period, value): "Use primitive inputs and retain observed outputs under report-only names." ) result = super().set_input(variable_name, period, value) + # Core snapshots input_variables during construction, but later + # explicit inputs must also survive counterfactual branches. Keep + # only stored inputs; neutralized and expired writes are ignored. + input_variables = getattr(self, "input_variables", None) + if input_variables is not None and variable_name not in input_variables: + holder = self.get_holder(variable_name) + variable = holder.variable + input_period = periods_.period(period) + if ( + not variable.is_neutralized + and (variable.end is None or input_period.start.date <= variable.end) + and any( + branch_name == self.branch_name + for branch_name, _ in holder.get_known_branch_periods() + ) + ): + # Core shallow-copies metadata when branching. Assign a new + # list so this input cannot change another branch's registry. + self.input_variables = [*input_variables, variable_name] if variable_name == "county_fips": self._record_county_input_types(period) return result diff --git a/policyengine_us/tests/core/test_combined_behavioral_responses.py b/policyengine_us/tests/core/test_combined_behavioral_responses.py index 10e8b367eab..d92edb84ae6 100644 --- a/policyengine_us/tests/core/test_combined_behavioral_responses.py +++ b/policyengine_us/tests/core/test_combined_behavioral_responses.py @@ -5,11 +5,13 @@ from policyengine_core.periods import period as make_period from policyengine_core.reforms import Reform +import policyengine_us.variables.household.marginal_tax_rate_helpers as marginal_rate_helpers from policyengine_us import Simulation from policyengine_us.variables.gov.simulation.behavioral_response_measurements import ( BEHAVIORAL_RESPONSE_CACHE_ATTR, BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, + get_behavioral_response_measurements, ) YEAR = 2026 @@ -60,6 +62,91 @@ def combined_simulation(fixed_net_income=False): return Simulation(situation=situation, reform=policy) +def test_combined_measurements_keep_inputs_set_on_baseline_and_reform(monkeypatch): + """Measurement branches and their raised branches retain identification. + + Fixed net income skips benefit calculations while real measurement and + rate branches expose the input-preservation invariant. Independent Texas + federal and health MTR tests cover the actual CTC policy consequences. + The one-person setup and fixed income limit additional CI cost. + """ + identification = np.array([False]) + simulation = combined_simulation(fixed_net_income=True) + for current in (simulation, simulation.baseline): + current.set_input("has_itin", YEAR, identification) + caller_arrays = [ + ( + current, + { + variable: current.get_array(variable, YEAR).copy() + for variable in ( + "has_itin", + "employment_income_before_lsr", + "long_term_capital_gains_before_response", + "weekly_hours_worked_before_lsr", + "household_net_income", + ) + }, + ) + for current in (simulation, simulation.baseline) + ] + observed = {} + create_branch = marginal_rate_helpers.create_perturbed_branch + + def observe_raised_branch(parent, period, branch_name, increments): + branch = create_branch(parent, period, branch_name, increments) + assert branch.parent_branch is parent + assert branch.branch_name == branch_name + np.testing.assert_array_equal( + parent.get_array("has_itin", YEAR), identification + ) + kind = ( + "capital_gains" if "long_term_capital_gains" in increments else "earnings" + ) + value = branch.get_array("has_itin", YEAR) + observed.setdefault((parent.branch_name, kind), []).append( + None if value is None else value.copy() + ) + return branch + + # Earnings imports the helper module; Core dynamically loads the capital + # gains formula, whose helper binding lives in its actual globals. + monkeypatch.setattr( + marginal_rate_helpers, "create_perturbed_branch", observe_raised_branch + ) + for current in (simulation, simulation.baseline): + formula = current.tax_benefit_system.variables[ + "marginal_tax_rate_on_capital_gains" + ].get_formula(make_period(YEAR)) + monkeypatch.setitem( + formula.__globals__, "create_perturbed_branch", observe_raised_branch + ) + measurements = get_behavioral_response_measurements( + simulation.person, make_period(YEAR) + ) + expected_paths = { + (parent, kind) + for parent in ( + BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, + BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH, + ) + for kind in ("earnings", "capital_gains") + } + assert set(observed) == expected_paths + for path, arrays in observed.items(): + for value in arrays: + np.testing.assert_array_equal(value, identification, err_msg=str(path)) + for name in ("baseline", "reform"): + np.testing.assert_array_equal(measurements[f"{name}_net_income"], [60_000]) + np.testing.assert_array_equal(measurements[f"{name}_mtr"], [1]) + np.testing.assert_array_equal(measurements[f"{name}_capital_gains_mtr"], [1]) + for current, arrays in caller_arrays: + for variable, value in arrays.items(): + np.testing.assert_array_equal(current.get_array(variable, YEAR), value) + assert BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH not in current.branches + assert BASELINE_BEHAVIORAL_RESPONSE_MEASUREMENT_BRANCH not in current.branches + + @pytest.mark.parametrize( "query_order", [RESPONSES, RESPONSES[::-1]], diff --git a/policyengine_us/tests/core/test_marginal_rate_branches.py b/policyengine_us/tests/core/test_marginal_rate_branches.py index ac3a9b80aae..34612fdd2cb 100644 --- a/policyengine_us/tests/core/test_marginal_rate_branches.py +++ b/policyengine_us/tests/core/test_marginal_rate_branches.py @@ -22,6 +22,8 @@ 6. A reformed simulation's baseline branch, and the baseline and reform branches that measure behavioral responses, equal fresh simulations of the baseline and of the reform without responses. +7. Inputs explicitly set on a measuring branch survive its raised branch, + with only the raised earnings changed; calculated caches are recomputed. Invariant 4 is what the branches rely on. The ``Simulation`` and ``Microsimulation`` wrappers move ``employment_income`` and similar inputs onto @@ -36,10 +38,13 @@ """ import copy +from itertools import count import numpy as np import pandas as pd import pytest +from hypothesis import HealthCheck, Phase, given, settings +from hypothesis import strategies as st from policyengine_core.periods import period as make_period from policyengine_core.reforms import Reform @@ -51,6 +56,7 @@ ) from policyengine_us.variables.household.marginal_tax_rate_helpers import ( HELD_BEHAVIORAL_RESPONSE_VARIABLES, + create_perturbed_branch, ) YEAR = 2026 @@ -316,6 +322,200 @@ def assert_rates_match(simulation, expected, rates): }, [], ) +BRANCH_NAMES = count() + + +@pytest.fixture(scope="module") +def branch_input_simulation(): + """A read-only source model; each case mutates only its own branches.""" + return Simulation( + situation=situation( + [("TX", {"parent": {"age": 30, "employment_income": 80_000}}, [4, 8])] + ) + ) + + +@pytest.fixture(scope="module") +def branch_identification_rates(): + """Fresh references share the same identification inputs at both wages.""" + base_situation = situation( + [("TX", {"parent": {"age": 30, "employment_income": 80_000}}, [4, 8])] + ) + for person, has_itin in zip(base_situation["people"].values(), [False, True, True]): + person["has_itin"] = {YEAR: has_itin} + raised_situation = copy.deepcopy(base_situation) + raised_situation["people"]["h0_parent"]["employment_income"] = {YEAR: 81_000} + base = Simulation(situation=base_situation) + raised = Simulation(situation=raised_situation) + # Identification denies the $4,400 CTC at both earnings levels. + for reference in (base, raised): + np.testing.assert_array_equal(reference.calculate("ctc", YEAR), [0]) + return { + "federal_marginal_tax_rate": ( + raised.calculate("income_tax", YEAR)[0] + - base.calculate("income_tax", YEAR)[0] + ) + / DELTA, + "marginal_tax_rate_including_health_benefits": 1 + - ( + raised.calculate("household_net_income_including_health_benefits", YEAR)[0] + - base.calculate("household_net_income_including_health_benefits", YEAR)[0] + ) + / DELTA, + } + + +@pytest.mark.parametrize( + "rate", + ["federal_marginal_tax_rate", "marginal_tax_rate_including_health_benefits"], +) +def test_rates_keep_identification_inputs_set_on_the_measuring_branch( + branch_input_simulation, branch_identification_rates, rate +): + """A $1,000 raise cannot restore CTC denied by the parent's branch input. + + Core does not add branch set_input calls to input_variables. Keeping only + that construction-time list erased has_itin from the raised branch, so + the federal rate included a spurious $4,400 credit increase. + """ + simulation = branch_input_simulation + branch = simulation.get_branch(f"identification_{rate}") + branch.set_input("has_itin", YEAR, [False, True, True]) + try: + np.testing.assert_array_equal(branch.calculate("ctc", YEAR), [0]) + actual = branch.calculate(rate, YEAR) + assert actual[0] == pytest.approx( + branch_identification_rates[rate], abs=TOLERANCE + ) + np.testing.assert_array_equal(actual[1:], [0, 0]) + raised = create_perturbed_branch( + branch, + make_period(YEAR), + "identification_ctc_with_raise", + {"employment_income": np.array([DELTA, 0, 0])}, + ) + np.testing.assert_array_equal( + raised.get_array("has_itin", YEAR), [False, True, True] + ) + np.testing.assert_array_equal(raised.calculate("ctc", YEAR), [0]) + finally: + simulation.branches.pop(branch.branch_name) + + +@settings( + max_examples=8, + derandomize=True, + database=None, + deadline=None, + # Cloning real-model branches makes failure shrinking expensive. The + # retained-array assertion already identifies the lost input and value. + phases=[Phase.generate], + suppress_health_check=[HealthCheck.too_slow], +) +@given( + tin=st.booleans(), + state=st.sampled_from(["TX", "OH", "CA"]), + zip_code=st.sampled_from(["77001", "44101", "90001"]), + interest=st.integers(1, 20_000), + incarcerated=st.booleans(), + head=st.booleans(), + earnings=st.integers(0, 100_000), + increment=st.integers(1, 2_000), + nested=st.booleans(), +) +def test_raised_branch_preserves_explicit_inputs_and_isolation( + branch_input_simulation, + tin, + state, + zip_code, + interest, + incarcerated, + head, + earnings, + increment, + nested, +): + """Invariant 7 across input types, periods and one nested branch. + + All examples use branches of one read-only model, and calculate only the + earnings aggregate; policy outputs are covered by the fresh references. + """ + simulation = branch_input_simulation + name = f"input_property_{next(BRANCH_NAMES)}" + parent = simulation.get_branch(name) + sibling = simulation.get_branch(f"{name}_sibling") + original_inputs = list(simulation.input_variables) + original_earnings = simulation.get_array( + "employment_income_before_lsr", YEAR + ).copy() + parent.set_input("has_itin", YEAR, [tin, True, True]) + sibling.set_input("has_itin", YEAR, [not tin, True, True]) + measuring = parent.get_branch(f"{name}_nested") if nested else parent + explicit = ( + ("has_itin", YEAR, [tin, True, True]), + ("state_code", YEAR, [state]), + ("zip_code", YEAR, [zip_code]), + ("taxable_interest_income", YEAR, [interest, 0, 0]), + ("is_incarcerated", f"{YEAR}-06", [incarcerated, False, False]), + ("is_household_head", YEAR, [head, False, False]), + ("employment_income_before_lsr", YEAR, [earnings, 0, 0]), + ) + try: + for variable, input_period, value in explicit[1:]: + measuring.set_input(variable, input_period, value) + expected = { + (variable, input_period): measuring.get_array(variable, input_period).copy() + for variable, input_period, _ in explicit + } + cached_earnings = measuring.calculate("employment_income", YEAR).copy() + raised = create_perturbed_branch( + measuring, + make_period(YEAR), + f"{name}_raised", + {"employment_income": np.array([increment, 0, 0])}, + ) + + # employment_income is calculated from the explicitly supplied + # pre-response earnings, so the raised branch must clear its cache. + assert raised.get_array("employment_income", YEAR) is None + for (variable, input_period), value in expected.items(): + raised_value = ( + value + [increment, 0, 0] + if variable == "employment_income_before_lsr" + else value + ) + np.testing.assert_array_equal( + raised.get_array(variable, input_period), + raised_value, + err_msg=f"{variable} {input_period}", + ) + np.testing.assert_array_equal( + measuring.get_array(variable, input_period), value, err_msg=variable + ) + np.testing.assert_array_equal( + raised.calculate("employment_income", YEAR), + cached_earnings + [increment, 0, 0], + ) + np.testing.assert_array_equal( + measuring.get_array("employment_income", YEAR), cached_earnings + ) + # Eternal inputs retain their value when read through another year. + np.testing.assert_array_equal( + raised.get_array("is_household_head", YEAR - 1), + expected[("is_household_head", YEAR)], + ) + np.testing.assert_array_equal( + sibling.get_array("has_itin", YEAR), [not tin, True, True] + ) + np.testing.assert_array_equal( + simulation.get_array("employment_income_before_lsr", YEAR), + original_earnings, + ) + assert simulation.get_array("has_itin", YEAR) is None + assert simulation.input_variables == original_inputs + finally: + simulation.branches.pop(name) + simulation.branches.pop(sibling.branch_name) def test_ohio_parent_marginal_tax_rate_matches_fresh_simulations(): diff --git a/policyengine_us/variables/household/marginal_tax_rate_helpers.py b/policyengine_us/variables/household/marginal_tax_rate_helpers.py index e7059055428..6b56cbf3c02 100644 --- a/policyengine_us/variables/household/marginal_tax_rate_helpers.py +++ b/policyengine_us/variables/household/marginal_tax_rate_helpers.py @@ -39,6 +39,8 @@ def create_perturbed_branch(simulation, period, branch_name, increments): Returns: The branch simulation. """ + # US wrappers also register explicit set_input calls after construction, + # including identification and other inputs set on the measuring branch. input_variables = set(simulation.input_variables) variables = simulation.tax_benefit_system.variables responses = {