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..5e4b0b5d203 --- /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, so rate branches keep them (`marginal_tax_rate` with `long_term_capital_gains` read about 13 on Core 3.33), and branches raise those inputs, so programs that read pre-response income, such as TANF, see the raise. diff --git a/policyengine_us/spm.py b/policyengine_us/spm.py index 9fe2f296d70..94d547c331e 100644 --- a/policyengine_us/spm.py +++ b/policyengine_us/spm.py @@ -1138,6 +1138,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/system.py b/policyengine_us/system.py index c97ca49fa4f..65a99683f52 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, ) @@ -209,6 +212,78 @@ 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) + delete_input_values(holder, 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 _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 + # 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 + + def _backfill_state_code_from_str(simulation): """Backfill the ``state_code`` enum from a ``state_code_str`` input. @@ -280,47 +355,11 @@ def __init__(self, *args, **kwargs): self.apply_reform(reform) migrate_social_security_inputs(self) - - # 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) - delete_input_values(employment_income, 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) - delete_input_values(self_employment_income, 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 - ) - delete_input_values(sstb_self_employment_income, 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) - delete_input_values(weekly_hours, 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 - ) - delete_input_values(cg_holder, 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) + _rebuild_baseline_branch(self) def _download_or_explain(dataset_str, download): @@ -542,68 +581,13 @@ def __init__(self, *args, **kwargs): self.apply_reform(reform) migrate_social_security_inputs(self) - - # 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) - delete_input_values(employment_income, 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) - delete_input_values(self_employment_income, 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 - ) - delete_input_values(sstb_self_employment_income, 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) - delete_input_values(weekly_hours, 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 - ) - delete_input_values(cg_holder, 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) + _rebuild_baseline_branch(self) class IndividualSim(CoreIndividualSim): # Deprecated diff --git a/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py b/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py index b285cd8143f..cab3c2a1b89 100644 --- a/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py +++ b/policyengine_us/tests/core/test_behavioral_response_measurement_copies.py @@ -257,12 +257,21 @@ def test_migrated_input_records_are_private_and_explicit_overrides_survive(): for original, (migrated, value) in migrated_inputs.items(): assert (migrated, source.branch_name, period) in source._user_input_keys assert (original, source.branch_name, period) not in source._user_input_keys - # Baseline was copied before migration and must retain its raw inputs. + # The baseline is rebuilt after migration, so both arms retain the + # same pre-response inputs and forget the emptied aggregate inputs. assert ( - original, + migrated, source.branch_name, period, ) in source.baseline._user_input_keys + assert ( + original, + source.branch_name, + period, + ) not in source.baseline._user_input_keys + np.testing.assert_array_equal( + source.baseline.get_array(migrated, YEAR), [value] + ) np.testing.assert_array_equal( source.baseline.calculate(original, YEAR), [value] ) @@ -378,9 +387,8 @@ def test_baseline_measurement_branch_is_baseline_policy_over_the_inputs(): def test_baseline_measurement_branch_reads_a_branchs_own_inputs(): - # A reform household's baseline keeps the situation's employment_income as - # an override from construction, so a branch with other earnings cannot be - # measured from it. + # A reform household's baseline keeps the situation's original earnings, + # so a branch with other earnings cannot be measured from it. simulation = make_simulation(situation=household(50_000), reform=RATE_CUT) branch = simulation.get_branch("other_earnings") branch.set_input("employment_income_before_lsr", YEAR, np.array([100_000.0])) 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..d92edb84ae6 --- /dev/null +++ b/policyengine_us/tests/core/test_combined_behavioral_responses.py @@ -0,0 +1,219 @@ +"""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 + +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 +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) + + +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]], + 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/tests/core/test_marginal_rate_branches.py b/policyengine_us/tests/core/test_marginal_rate_branches.py new file mode 100644 index 00000000000..ee0780d5fb9 --- /dev/null +++ b/policyengine_us/tests/core/test_marginal_rate_branches.py @@ -0,0 +1,1090 @@ +"""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. +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. +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 +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 +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 + +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, + create_perturbed_branch, +) + +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", +} +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"): + """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 + + +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 situation's earnings are the + pre-response inputs, so the raise goes on those. + """ + ranks = simulation.calculate("adult_earnings_index", YEAR) + 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: + 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(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(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: + raised["people"][person]["long_term_capital_gains"] = { + YEAR: float(gains[i] + DELTA) + } + 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") + + +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, policy=None): + """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 = fresh_simulation(raised_situation, simulation, policy) + 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, policy=None): + 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( + fresh_simulation(raised_situation, simulation, policy), + "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], +) +# 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, + } + }, + [], +) +BRANCH_NAMES = count() + + +@pytest.fixture(scope="module") +def aggregate_input_period_references(): + """Read-only Ohio references for raises to pre-response and aggregate wages.""" + base_situation = situation([OHIO_WORKS_FIRST_PARENT]) + simulation = Simulation(situation=base_situation) + base = float(simulation.calculate("household_net_income", YEAR)[0]) + raised_situation = copy.deepcopy(base_situation) + raised_situation["people"]["h0_parent"]["employment_income"] = {YEAR: 7_000} + raised = Simulation(situation=raised_situation) + aggregate_raised = Simulation(situation=base_situation) + aggregate_raised.set_input("employment_income", YEAR, [7_000, 0, 0]) + references = {} + for input_year, reference in ((YEAR + 1, raised), (YEAR, aggregate_raised)): + references[input_year] = ( + 1 + - (float(reference.calculate("household_net_income", YEAR)[0]) - base) + / DELTA, + reference.calculate("tanf", YEAR)[0], + ) + return simulation, references + + +@pytest.mark.parametrize("input_year", [YEAR + 1, YEAR]) +def test_aggregate_input_only_changes_the_perturbation_target_in_its_period( + aggregate_input_period_references, input_year +): + """Next year's aggregate input cannot hide this year's raise from TANF. + + An aggregate supplied for the measured year must still be raised directly. + Both cases preserve their supplied input and leave the source unchanged. + """ + simulation, references = aggregate_input_period_references + branch = simulation.get_branch(f"aggregate_input_{input_year}") + supplied = [10_000 if input_year != YEAR else 6_000, 0, 0] + branch.set_input("employment_income", input_year, supplied) + try: + rate = branch.calculate("marginal_tax_rate", YEAR) + raised = create_perturbed_branch( + branch, + make_period(YEAR), + "aggregate_input_raised", + {"employment_income": np.array([DELTA, 0, 0])}, + ) + tanf = raised.calculate("tanf", YEAR)[0] + expected_rate, expected_tanf = references[input_year] + print( + f"input_year={input_year}: rate={rate[0]:.9f}, TANF={tanf:.2f}; " + f"fresh rate={expected_rate:.9f}, TANF={expected_tanf:.2f}" + ) + assert rate[0] == pytest.approx(expected_rate, abs=TOLERANCE) + np.testing.assert_array_equal(rate[1:], [0, 0]) + assert tanf == pytest.approx(expected_tanf, abs=0.01) + np.testing.assert_array_equal( + raised.calculate("employment_income_before_lsr", YEAR), + [6_000 if input_year == YEAR else 7_000, 0, 0], + ) + np.testing.assert_array_equal( + raised.calculate("employment_income", YEAR), [7_000, 0, 0] + ) + np.testing.assert_array_equal( + branch.get_array("employment_income", input_year), supplied + ) + if input_year != YEAR: + np.testing.assert_array_equal( + raised.get_array("employment_income", input_year), supplied + ) + for unchanged in (simulation, branch): + np.testing.assert_array_equal( + unchanged.get_array("employment_income_before_lsr", YEAR), + [6_000, 0, 0], + ) + assert "employment_income" not in simulation.input_variables + finally: + simulation.branches.pop(branch.branch_name) + + +@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) + + +def test_raised_branch_recomputes_ira_cache_from_another_input_period( + branch_input_simulation, +): + """A supplied next-year IRA amount must not freeze this year's allocation. + + With $1,000 of compensation and equal desired contributions, each IRA + receives $500. A $1,000 earnings raise increases each generated amount + to $1,000, while the explicitly supplied $2,000 Roth contribution for + the following year stays fixed. Reuse the read-only model through two + isolated branches; only the IRA allocation needs calculating. + """ + simulation = branch_input_simulation + name = f"ira_input_period_{next(BRANCH_NAMES)}" + parent = simulation.get_branch(name) + fresh = simulation.get_branch(f"{name}_fresh") + + def set_current_inputs(branch, earnings): + branch.set_input("employment_income_before_lsr", YEAR, [earnings, 0, 0]) + for variable in ( + "traditional_ira_contributions_desired", + "roth_ira_contributions_desired", + ): + branch.set_input(variable, YEAR, [5_000, 0, 0]) + + try: + set_current_inputs(parent, 1_000) + np.testing.assert_array_equal( + parent.calculate("roth_ira_contributions", YEAR), [500, 0, 0] + ) + parent.set_input("roth_ira_contributions", YEAR + 1, [2_000, 0, 0]) + + set_current_inputs(fresh, 2_000) + fresh.set_input("roth_ira_contributions", YEAR + 1, [2_000, 0, 0]) + raised = create_perturbed_branch( + parent, + make_period(YEAR), + f"{name}_raised", + {"employment_income": np.array([DELTA, 0, 0])}, + ) + + for variable in ("roth_ira_contributions", "traditional_ira_contributions"): + expected = fresh.calculate(variable, YEAR) + np.testing.assert_array_equal(expected, [1_000, 0, 0]) + np.testing.assert_array_equal(raised.calculate(variable, YEAR), expected) + for branch in (parent, raised, fresh): + np.testing.assert_array_equal( + branch.get_array("roth_ira_contributions", YEAR + 1), [2_000, 0, 0] + ) + np.testing.assert_array_equal( + parent.get_array("roth_ira_contributions", YEAR), [500, 0, 0] + ) + finally: + simulation.branches.pop(name) + simulation.branches.pop(fresh.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(): + """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_long_term_gains_keep_their_value_in_the_marginal_rate_branch(): + """Issue 10522: on Core 3.33 the rate read 13.05 with long-term gains. + + The branch erased long_term_capital_gains_before_response, which holds the + supplied gains, so they fell to $0 when the branch raised earnings. The + filer's $103,900 of ordinary taxable income ($120,000 less the $16,100 + standard deduction) stays in the 22% bracket after a $1,000 raise, with + 7.65% of payroll tax, and the gains stack on top in the 15% bracket. + """ + filer = {"age": 40, "employment_income": 120_000} + wages_only = Simulation(situation=situation([("TX", {"filer": filer}, [])])) + with_gains = Simulation( + situation=situation( + [("TX", {"filer": {**filer, "long_term_capital_gains": 15_000}}, [])] + ) + ) + + assert wages_only.calculate("marginal_tax_rate", YEAR)[0] == pytest.approx( + 0.2965, abs=TOLERANCE + ) + assert with_gains.calculate("marginal_tax_rate", YEAR)[0] == pytest.approx( + 0.2965, abs=TOLERANCE + ) + assert with_gains.calculate("long_term_capital_gains", YEAR)[0] == 15_000 + + +def test_marginal_rates_along_axes_match_fresh_simulations(): + """API requests sweep earnings with axes, which repeat each household.""" + wages = [14_000, 21_000, 28_000] + state, _, children = OHIO_PARENT + swept = situation([(state, {"parent": {"age": 30}}, children)]) + swept["axes"] = [ + [ + { + "name": "employment_income", + "count": 3, + "min": wages[0], + "max": wages[-1], + "period": YEAR, + } + ] + ] + + rates = Simulation(situation=swept).calculate("marginal_tax_rate", YEAR) + + def net_income(earnings): + household = ( + state, + {"parent": {"age": 30, "employment_income": earnings}}, + children, + ) + return Simulation(situation=situation([household])).calculate( + "household_net_income", YEAR + )[0] + + expected = [1 - (net_income(w + DELTA) - net_income(w)) / DELTA for w in wages] + np.testing.assert_allclose(rates[::3], expected, atol=TOLERANCE) + + +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([TEXAS_GAINS_FILER]) + 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) + + +# 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.1) + + 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/core/test_spm_system.py b/policyengine_us/tests/core/test_spm_system.py index dd3b80f96da..ee343a4c1b4 100644 --- a/policyengine_us/tests/core/test_spm_system.py +++ b/policyengine_us/tests/core/test_spm_system.py @@ -137,6 +137,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() 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 8ad352eb148..a3bbd71491e 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. @@ -134,3 +135,177 @@ marginal_tax_rate_computed: [true, true, false] federal_marginal_tax_rate: [0.22, 0.12, 0] fica_marginal_tax_rate: [0.0765, 0.0765, 0] + +# 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_294.93 + +- 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_175.43 + +- 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_175.43 - 23_294.93) / 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_073.76 + +- 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_546.66 + +- 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_546.66 - 41_073.76) / 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 269f0bf8dfd..780033b204e 100644 --- a/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py +++ b/policyengine_us/variables/gov/simulation/behavioral_response_measurements.py @@ -17,6 +17,16 @@ # it, and its measurements, again, wherever a program reads hours. "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 +# 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 5c85a12dd99..1d1e384a435 100644 --- a/policyengine_us/variables/gov/simulation/capital_gains_responses.py +++ b/policyengine_us/variables/gov/simulation/capital_gains_responses.py @@ -1,12 +1,12 @@ from policyengine_us.model_api import * -from policyengine_us.tools.period_branch import drop_inherited_values, has_input -from policyengine_us.variables.household.marginal_tax_rate_helpers import ( - user_set_variables, -) +from policyengine_us.tools.period_branch import has_input from policyengine_us.variables.gov.simulation.behavioral_response_measurements import ( calculate_relative_capital_gains_mtr_change, get_behavioral_response_measurements, ) +from policyengine_us.variables.household.marginal_tax_rate_helpers import ( + create_perturbed_branch, +) from policyengine_us.variables.gov.simulation.capital_gains_response_forms import ( CAPITAL_GAINS_RESPONSE_FORM_VARIABLES, capital_gains_response_factor, @@ -181,7 +181,6 @@ def formula(person, period, parameters): # pragma: no cover simulation = person.simulation adult_index_values = person("adult_index_cg", period) long_term_capital_gains = person("long_term_capital_gains", period) - capital_gains = person("capital_gains", period) # capital_gains adds short- and long-term gains, and adjusted gross # income reads it. A total given as an input, at construction or # later, is stored as given rather than added up again. @@ -191,38 +190,21 @@ def formula(person, period, parameters): # pragma: no cover # marginal_tax_rate's do; a branch made first measured its own. household_net_income = person.household("household_net_income", period) rise = capital_gains_mtr_rise(household_net_income, long_term_capital_gains) - variables = simulation.tax_benefit_system.variables - inputs = user_set_variables(simulation) - inputs_with_formulas = [ - variable - for variable in inputs - if variable in variables and not variables[variable].is_input_variable() - ] 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, variable_data in variables.items(): - if variable not in inputs and not variable_data.is_input_variable(): - alt_simulation.delete_arrays(variable) - # A variable given as an input for one period keeps only the - # periods it was given for; what was calculated for the others, - # this one included, is calculated again. - drop_inherited_values(alt_simulation, inputs_with_formulas) # Raise long-term gains, which net capital gain and so the # preferential rates read. Raising capital_gains instead reached # adjusted gross income but not net capital gain, so it measured # the rate on ordinary income except where taxable income was - # below net capital gain. The override holds any behavioral + # below net capital gain. The branch holds any behavioral # response at its value here. - alt_simulation.set_input( - "long_term_capital_gains", - period, - long_term_capital_gains + mask * rise, - ) + increments = {"long_term_capital_gains": mask * rise} if capital_gains_is_input: - alt_simulation.set_input( - "capital_gains", period, capital_gains + mask * rise - ) + increments["capital_gains"] = mask * rise + alt_simulation = create_perturbed_branch( + simulation, period, branch_name, increments + ) alt_person = alt_simulation.person # Divide by the rise as stored, after rounding to float32. stored_rise = ( @@ -234,5 +216,5 @@ def formula(person, period, parameters): # pragma: no cover increase = household_net_income_higher_gains - household_net_income mtr_values += where(mask, 1 - increase / where(mask, stored_rise, 1), 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 e6f38c8c647..503da259295 100644 --- a/policyengine_us/variables/household/marginal_tax_rate.py +++ b/policyengine_us/variables/household/marginal_tax_rate.py @@ -1,6 +1,6 @@ from policyengine_us.model_api import * from policyengine_us.variables.household.marginal_tax_rate_helpers import ( - user_set_variables, + marginal_earnings_change, ) @@ -18,61 +18,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, - ) - - inputs = user_set_variables(sim) - 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 inputs 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 6a105fc2a3a..8cb95a033f6 100644 --- a/policyengine_us/variables/household/marginal_tax_rate_helpers.py +++ b/policyengine_us/variables/household/marginal_tax_rate_helpers.py @@ -1,21 +1,32 @@ from policyengine_us.model_api import * +from policyengine_core.periods import period as make_period +from policyengine_us.tools.period_branch import drop_inherited_values +from policyengine_us.variables.gov.simulation.behavioral_response_measurements import ( + NEUTRALIZED_BEHAVIORAL_RESPONSE_VARIABLES, + PRE_RESPONSE_INPUTS, +) + from policyengine_us.variables.household.expense.retirement._ira_supplied_contributions import ( IRA_CONTRIBUTION_VARIABLES, ) -def user_set_variables(sim): +def user_set_variables(sim, period=None): """Names of the variables a counterfactual branch must keep. - ``sim.input_variables`` lists inputs given when the simulation was built. + ``sim.input_variables`` lists construction inputs and explicit later inputs + registered by the US wrappers. The IRA contribution variables are also kept when the caller set them later through ``set_input`` (core records those in ``_user_input_keys``): ``ira_contribution_limit`` reads them as supplied, so a branch that deleted them would regenerate a different IRA allocation alongside the - earnings bump. Other later ``set_input`` calls are not kept, because - model formulas also call ``set_input`` internally (the SSA revenue - formulas set ``tax_unit_taxable_social_security``), and a branch must - recompute those. + earnings bump. Other Core-recorded keys are not added here: formulas also call + ``holder.set_input`` internally (the SSA revenue formulas set + ``tax_unit_taxable_social_security``), and a branch must recompute those. + + With ``period``, only inputs supplied for that period count (eternal + inputs apply in every period). Without it, return names across all input + periods so cache cleanup can preserve other-year inputs too. """ names = set(sim.input_variables) keys = getattr(sim, "_user_input_keys", None) or () @@ -28,63 +39,181 @@ def user_set_variables(sim): for name, branch, _period in keys if name in IRA_CONTRIBUTION_VARIABLES and branch in branches ) - return names + if period is None: + return names + period = make_period(period) + return { + name + for name, branch, input_period in keys + if name in names + and branch in branches + and ( + input_period == period + or sim.tax_benefit_system.variables[name].definition_period == ETERNITY + ) + } -def compute_component_mtr(person, period, parameters, tax_variable, branch_prefix): - """Compute the marginal tax rate for a specific tax component. +# 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 - 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. + +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. + + 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. + branch_name: Name of the branch; the caller deletes it when done. + increments: Person-level increase for each raised variable. + + 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 = user_set_variables(simulation) + period_input_variables = user_set_variables(simulation, period) + 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 variables: + if variable not in input_variables: + branch.delete_arrays(variable) + # The name filter above must retain inputs for every supplied period. + # Keep only their actual input keys, rather than other-period caches. + drop_inherited_values(branch) + for variable, response in responses.items(): + branch.set_input(variable, period, response) + for variable, increment in increments.items(): + if variable not in period_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. + + 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) - - inputs = user_set_variables(sim) + 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 inputs 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 02616ff8eb7..84007070652 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,6 +1,6 @@ from policyengine_us.model_api import * from policyengine_us.variables.household.marginal_tax_rate_helpers import ( - user_set_variables, + marginal_earnings_change, ) @@ -18,62 +18,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, - ) - - inputs = user_set_variables(sim) - 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 inputs 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)