diff --git a/changelog.d/fix-branch-override-shadowing.fixed.md b/changelog.d/fix-branch-override-shadowing.fixed.md new file mode 100644 index 00000000000..b183cc39fd7 --- /dev/null +++ b/changelog.d/fix-branch-override-shadowing.fixed.md @@ -0,0 +1 @@ +Limit the non-refundable Child Tax Credit by the actual tax liability, SALT deduction included (26 U.S.C. 26(a)), instead of a recomputation without SALT that applied or not depending on which variables were calculated first; and make the itemization, Delaware and Virginia EITC, Idaho aged or disabled, Missouri TANF caretaker and Medicaid SSI-supplement comparison branches calculate under their overridden inputs even when the simulation has already calculated those inputs, and again for each period. diff --git a/policyengine_us/tests/core/test_override_branches.py b/policyengine_us/tests/core/test_override_branches.py new file mode 100644 index 00000000000..07434effd78 --- /dev/null +++ b/policyengine_us/tests/core/test_override_branches.py @@ -0,0 +1,347 @@ +"""Formula branches calculate under their overrides, whatever was cached first. + +Several formulas compare a tax unit's liability under alternative choices by +calculating it in a branch with one input overridden: itemizing or not +(``tax_unit_itemizes``), Delaware and Virginia EITC refundability, the Idaho +aged or disabled credit or deduction, Missouri TANF caretaker inclusion and +Medicaid for SSI state supplements. A policyengine-core branch starts as a copy +of every array its parent has cached, and setting an input on it clears none +of them, so a value the parent calculated from the old input answers for the +branch. ``get_override_branch`` drops the copied values when the parent has +already calculated the overridden input, and creates the branch again for each +period. + +The non-refundable CTC used to be limited by the tax liability recomputed +without the SALT deduction, in a "no_salt" branch. The branch usually inherited +the liability with SALT, so the CTC depended on which variables a caller asked +for first. 26 U.S.C. 26(a) limits the credit by the actual tax liability, SALT +deduction included, and ``ctc_limiting_tax_liability`` now reads it directly. + +A seeded sample of households shares one simulation; the tests check: + +1. The CTC-limiting liability is income tax before credits less the other + non-refundable credits, and the non-refundable and refundable parts add up + to the credit. +2. Every reported variable is the same whichever variable is calculated first. +3. Each comparison branch equals a fresh simulation that sets the overridden + input before anything is calculated (a differential test of the branch + against the reference path). +4. The same holds when the parent has already calculated the overridden + input, and when an earlier year was calculated first. +""" + +import numpy as np +import pytest + +from policyengine_us import Simulation +from policyengine_us.tools.override_branch import ( + drop_inherited_values, + get_override_branch, +) + +YEAR = 2026 +SEED = 20261001 +N = 48 +STATES = ["CA", "NY", "VA", "DE", "ID", "NJ", "MA", "IL", "TX", "OR"] +REPORTED = [ + "income_tax", + "income_tax_before_credits", + "ctc_limiting_tax_liability", + "non_refundable_ctc", + "refundable_ctc", + "tax_unit_itemizes", + "state_income_tax", + "household_net_income", +] + + +def _sample(): + rng = np.random.default_rng(SEED) + households = [] + for i in range(N): + married = bool(rng.random() < 0.6) + children = int(rng.integers(0, 4)) + earnings = float(np.round(np.exp(rng.uniform(np.log(8_000), np.log(300_000))))) + households.append( + dict( + state=STATES[i % len(STATES)], + married=married, + children=children, + earnings=earnings, + spouse_share=float(rng.choice([0, 0.2, 0.5])) if married else 0, + mortgage=float(rng.choice([0, 8_000, 18_000, 30_000])), + property_tax=float(rng.choice([0, 3_000, 9_000, 16_000])), + charity=float(rng.choice([0, 2_000, 9_000])), + aged_parent=bool(rng.random() < 0.15), + ) + ) + return households + + +def _situation(households, years=(YEAR,), itemizes=None): + people, units, marital = {}, {}, {} + for i, h in enumerate(households): + members = [] + + def add(name, **inputs): + people[name] = {k: {y: v for y in years} for k, v in inputs.items()} + members.append(name) + + head = f"h{i}" + earnings = h["earnings"] + add( + head, + age=40, + employment_income=earnings * (1 - h["spouse_share"]), + deductible_mortgage_interest=h["mortgage"], + real_estate_taxes=h["property_tax"], + charitable_cash_donations=h["charity"], + ) + marital[f"mu_{head}"] = {"members": [head]} + if h["married"]: + add(f"s{i}", age=38, employment_income=earnings * h["spouse_share"]) + marital[f"mu_{head}"]["members"].append(f"s{i}") + for c in range(h["children"]): + add(f"c{i}_{c}", age=3 + 4 * c, is_tax_unit_dependent=True) + marital[f"mu_c{i}_{c}"] = {"members": [f"c{i}_{c}"]} + if h["aged_parent"]: + add( + f"p{i}", + age=74, + is_tax_unit_dependent=True, + share_of_care_and_support_costs_paid_by_tax_filer=1.0, + ) + marital[f"mu_p{i}"] = {"members": [f"p{i}"]} + units[i] = members + tax_units = {f"tu{i}": {"members": m} for i, m in units.items()} + if itemizes is not None: + for i, unit in enumerate(tax_units.values()): + unit["tax_unit_itemizes"] = {y: bool(itemizes[i]) for y in years} + return { + "people": people, + "tax_units": tax_units, + "spm_units": {f"spm{i}": {"members": m} for i, m in units.items()}, + "families": {f"fam{i}": {"members": m} for i, m in units.items()}, + "marital_units": marital, + "households": { + f"hh{i}": { + "members": m, + "state_code": {y: households[i]["state"] for y in years}, + } + for i, m in units.items() + }, + } + + +@pytest.fixture(scope="module") +def households(): + return _sample() + + +@pytest.fixture(scope="module") +def baseline(households): + simulation = Simulation(situation=_situation(households)) + return {v: simulation.calculate(v, YEAR) for v in REPORTED} + + +def test_ctc_limit_is_actual_liability_less_other_credits(households): + simulation = Simulation(situation=_situation(households)) + limit = simulation.calculate("ctc_limiting_tax_liability", YEAR) + before_credits = simulation.calculate("income_tax_before_credits", YEAR) + credits = simulation.tax_benefit_system.parameters(YEAR).gov.irs.credits + other = sum( + simulation.calculate(credit, YEAR) + for credit in credits.non_refundable + if credit != "non_refundable_ctc" + ) + np.testing.assert_allclose(limit, np.maximum(0, before_credits - other)) + ctc = simulation.calculate("ctc", YEAR) + non_refundable = simulation.calculate("non_refundable_ctc", YEAR) + refundable = simulation.calculate("refundable_ctc", YEAR) + # non_refundable_ctc is the credit less its refundable part; the tax + # limit applies when non-refundable credits are capped. + assert (refundable <= ctc + 0.01).all() + np.testing.assert_allclose(non_refundable + refundable, ctc, atol=0.01) + # The sample includes itemizers with SALT whose CTC the limit binds. + salt = simulation.calculate("salt_deduction", YEAR) + itemizes = simulation.calculate("tax_unit_itemizes", YEAR) + assert (itemizes & (salt > 0) & (limit < ctc)).any() + + +@pytest.mark.parametrize( + "first", + [ + "income_tax", + "refundable_ctc", + "ctc_value", + "state_income_tax", + "tax_liability_if_itemizing", + "spm_unit_net_income", + ], +) +def test_results_do_not_depend_on_calculation_order(households, baseline, first): + simulation = Simulation(situation=_situation(households)) + simulation.calculate(first, YEAR) + for variable in REPORTED: + np.testing.assert_array_equal( + simulation.calculate(variable, YEAR), + baseline[variable], + err_msg=f"{variable} changes when {first} is calculated first", + ) + + +def _fresh(households, variable, overrides, year=YEAR): + simulation = Simulation(situation=_situation(households, years=(year,))) + for name, value in overrides.items(): + simulation.set_input(name, year, value) + return simulation.calculate(variable, year) + + +def test_itemization_branches_match_fresh_simulations(households): + simulation = Simulation(situation=_situation(households)) + simulation.calculate("household_net_income", YEAR) + n = len(households) + for comparison, itemizes in ( + ("tax_liability_if_itemizing", True), + ("tax_liability_if_not_itemizing", False), + ): + np.testing.assert_allclose( + simulation.calculate(comparison, YEAR), + _fresh(households, "income_tax", {"tax_unit_itemizes": [itemizes] * n}), + atol=0.01, + err_msg=comparison, + ) + + +@pytest.mark.parametrize( + "comparison,variable,override,value", + [ + ( + "de_income_tax_if_claiming_refundable_eitc", + "de_income_tax", + "de_claims_refundable_eitc", + True, + ), + ( + "de_income_tax_if_claiming_non_refundable_eitc", + "de_income_tax", + "de_claims_refundable_eitc", + False, + ), + ( + "va_income_tax_if_claiming_refundable_eitc", + "va_income_tax", + "va_claims_refundable_eitc", + True, + ), + ( + "va_income_tax_if_claiming_non_refundable_eitc", + "va_income_tax", + "va_claims_refundable_eitc", + False, + ), + ( + "id_income_tax_if_receiving_aged_or_disabled_credit", + "id_income_tax", + "id_receives_aged_or_disabled_credit", + True, + ), + ( + "id_income_tax_if_receiving_aged_or_disabled_deduction", + "id_income_tax", + "id_receives_aged_or_disabled_credit", + False, + ), + ], +) +def test_state_choice_branches_match_fresh_simulations( + households, comparison, variable, override, value +): + simulation = Simulation(situation=_situation(households)) + simulation.calculate("household_net_income", YEAR) + np.testing.assert_allclose( + simulation.calculate(comparison, YEAR), + _fresh(households, variable, {override: [value] * len(households)}), + atol=0.01, + ) + + +def test_branch_after_parent_calculated_the_overridden_input(households): + # tax_unit_itemizes is an input here, so the parent calculates income tax + # without the itemizing branch; the branch is created afterwards and must + # not answer with the parent's income tax. + n = len(households) + situation = _situation(households, itemizes=[False] * n) + simulation = Simulation(situation=situation) + not_itemizing = simulation.calculate("income_tax", YEAR) + itemizing = simulation.calculate("tax_liability_if_itemizing", YEAR) + expected = _fresh(households, "income_tax", {"tax_unit_itemizes": [True] * n}) + np.testing.assert_allclose(itemizing, expected, atol=0.01) + assert not np.allclose(itemizing, not_itemizing) + + +def test_later_year_branches_match_single_year_simulation(households): + years = (2025, YEAR) + simulation = Simulation(situation=_situation(households, years=years)) + simulation.calculate("tax_liability_if_itemizing", 2025) + simulation.calculate("income_tax", 2025) + fresh = Simulation(situation=_situation(households, years=years)) + for variable in REPORTED + ["tax_liability_if_itemizing"]: + np.testing.assert_array_equal( + simulation.calculate(variable, YEAR), + fresh.calculate(variable, YEAR), + err_msg=f"{variable} for {YEAR} depends on 2025 being calculated first", + ) + assert simulation.branches["itemizing"].branch_period.start.year == YEAR + + +def test_get_override_branch_reuses_within_period_and_recreates_otherwise( + households, +): + simulation = Simulation(situation=_situation(households, years=(2025, YEAR))) + n = len(households) + ones = np.ones(n, dtype=bool) + first = get_override_branch( + simulation, "test_branch", YEAR, {"tax_unit_itemizes": ones} + ) + assert ( + get_override_branch( + simulation, "test_branch", YEAR, {"tax_unit_itemizes": ones} + ) + is first + ) + other_inputs = get_override_branch( + simulation, "test_branch", YEAR, {"tax_unit_itemizes": ~ones} + ) + assert other_inputs is not first + other_period = get_override_branch( + simulation, "test_branch", 2025, {"tax_unit_itemizes": ones} + ) + assert other_period is not other_inputs + assert simulation.branches["test_branch"] is other_period + + +def test_drop_inherited_values_keeps_inputs_only(households): + n = len(households) + itemizes = np.ones(n, dtype=bool) + simulation = Simulation(situation=_situation(households)) + simulation.calculate("income_tax", YEAR) + # A plain branch with an input of its own, and a branch nested in it. + parent = simulation.get_branch("parent_with_input") + parent.set_input("tax_unit_itemizes", YEAR, itemizes) + child = parent.get_branch("child") + drop_inherited_values(child) + # Inputs survive: the situation's and the one set on the parent branch. + np.testing.assert_array_equal( + child.get_array("employment_income", YEAR), + simulation.get_array("employment_income", YEAR), + ) + np.testing.assert_array_equal(child.get_array("tax_unit_itemizes", YEAR), itemizes) + # Calculated values are gone, and are calculated again from the inputs. + assert child.get_array("income_tax", YEAR) is None + assert child.get_array("taxable_income", YEAR) is None + np.testing.assert_allclose( + child.calculate("taxable_income", YEAR), + _fresh(households, "taxable_income", {"tax_unit_itemizes": itemizes}), + atol=0.01, + ) diff --git a/policyengine_us/tests/test_ctc_itemizing_branch_cycle.py b/policyengine_us/tests/test_ctc_itemizing_branch_cycle.py index 9d797ce7b43..94ac4a349da 100644 --- a/policyengine_us/tests/test_ctc_itemizing_branch_cycle.py +++ b/policyengine_us/tests/test_ctc_itemizing_branch_cycle.py @@ -15,9 +15,9 @@ consumers (e.g. policyengine.py's household-impact integration tests) on `policyengine-core >= 3.24`. -The fix in `ctc_limiting_tax_liability.py` propagates the parent's -`tax_unit_itemizes` value to the no_salt child branch so the -`tax_unit_itemizes` formula is never re-entered there. +`ctc_limiting_tax_liability` no longer creates the no_salt branch: it reads +the branch's own `income_tax_before_credits`, which the itemizing branch +calculates with `tax_unit_itemizes` set, so the chain cannot re-enter it. """ import numpy as np diff --git a/policyengine_us/tools/general.py b/policyengine_us/tools/general.py index 7ad450690e3..0696875e350 100644 --- a/policyengine_us/tools/general.py +++ b/policyengine_us/tools/general.py @@ -1,6 +1,10 @@ from policyengine_core.model_api import * from policyengine_us.entities import * from policyengine_us.tools.branched_simulation import BranchedSimulation +from policyengine_us.tools.override_branch import ( + get_branch_for_period, + get_override_branch, +) from pathlib import Path import pandas as pd from policyengine_us.typing import Formula diff --git a/policyengine_us/tools/override_branch.py b/policyengine_us/tools/override_branch.py new file mode 100644 index 00000000000..d48368c3647 --- /dev/null +++ b/policyengine_us/tools/override_branch.py @@ -0,0 +1,152 @@ +"""Branches that calculate one period under overridden inputs. + +A policyengine-core branch starts as a copy of every array its parent has +cached, and ``set_input`` on the branch stores the new value without clearing +anything that was calculated from the old one. A branch therefore answers any +variable its parent has already calculated with the parent's value, whatever +the branch's inputs say. + +``get_override_branch`` makes the override reach every value the branch +calculates: + +- A branch serves one period. Asked for another period, it is created again + from the parent as the parent stands then, as a simulation calculating only + that period would create it. +- When the branch is created, it keeps the parent's cache only if the parent + has no value yet for any overridden variable and period. A cached value + cannot have been calculated from a value that did not exist, so the + parent's cache is then safe to share; this is the usual case, where a + formula branches while its parent is still calculating the variable the + branch overrides. Otherwise the branch drops every array it copied except + inputs, and calculates the rest itself. +- A branch reused within its period with different inputs is created again. +""" + +from typing import Dict, Tuple, Union + +import numpy as np +from policyengine_core.periods import Period +from policyengine_core.periods import period as to_period +from policyengine_core.simulations import Simulation + +Override = Union[np.ndarray, Tuple[Period, np.ndarray]] + + +def _overrides(period: Period, inputs: Dict[str, Override]): + for variable, value in inputs.items(): + if isinstance(value, tuple): + input_period, value = value + yield variable, to_period(input_period), np.asarray(value) + else: + yield variable, period, np.asarray(value) + + +def _is_known(simulation: Simulation, variable: str, period: Period) -> bool: + holder = simulation.get_holder(variable) + if holder.variable.is_neutralized: + return False + return holder.get_array(period, simulation.branch_name) is not None + + +def drop_inherited_values(branch: Simulation) -> None: + """Delete every array ``branch`` holds except inputs. + + Inputs are the variables the simulation was built with and each value set + through ``set_input`` on a branch this one reads. + """ + input_variables = set(branch.input_variables) + user_input_keys = getattr(branch, "_user_input_keys", set()) + visible_branches = set(branch._get_visible_branch_names()) + for population in branch.populations.values(): + for name, holder in population._holders.items(): + if name in input_variables: + continue + for branch_name, known_period in holder.get_known_branch_periods(): + if ( + branch_name in visible_branches + and (name, branch_name, known_period) in user_input_keys + ): + continue + # Exact key: ``Holder.delete_arrays`` would also delete any + # input stored at a sub-period of ``known_period``. + key = f"{branch_name}:{known_period}" + holder._memory_storage._arrays.pop(key, None) + if holder._disk_storage is not None: + holder._disk_storage._files.pop( + f"{branch_name}_{known_period}", None + ) + branch._fast_cache = {} + + +def get_override_branch( + simulation: Simulation, + name: str, + period: Period, + inputs: Dict[str, Override], +) -> Simulation: + """Return ``simulation``'s branch ``name`` for ``period``, with ``inputs`` set. + + ``inputs`` maps each overridden variable to its value for ``period``, or + to a ``(period, value)`` pair for another period (e.g. a month). + """ + overrides = list(_overrides(period, inputs)) + if name == simulation.branch_name: + # Already inside this branch (core's get_branch returns the + # simulation itself): only the inputs need setting. + for variable, input_period, value in overrides: + simulation.set_input(variable, input_period, value) + return simulation + branch = simulation.branches.get(name) + if branch is not None and ( + getattr(branch, "branch_period", None) != period + or not _same_overrides(branch, overrides) + ): + del simulation.branches[name] + branch = None + if branch is None: + parent_knows_override = any( + _is_known(simulation, variable, input_period) + for variable, input_period, _ in overrides + ) + branch = simulation.get_branch(name) + branch.branch_period = period + if parent_knows_override: + drop_inherited_values(branch) + for variable, input_period, value in overrides: + branch.set_input(variable, input_period, value) + branch.branch_overrides = overrides + return branch + + +def _same_overrides(branch: Simulation, overrides) -> bool: + previous = getattr(branch, "branch_overrides", None) + if previous is None or len(previous) != len(overrides): + return False + return all( + variable == previous_variable + and input_period == previous_period + and np.array_equal(value, previous_value) + for (variable, input_period, value), ( + previous_variable, + previous_period, + previous_value, + ) in zip(overrides, previous) + ) + + +def get_branch_for_period( + simulation: Simulation, name: str, period: Period +) -> Simulation: + """Return ``simulation``'s branch ``name`` for ``period``, created per period. + + For branches that change the tax-benefit system rather than inputs: the + caller swaps the system and deletes the variables it recalculates. + """ + if name == simulation.branch_name: + return simulation + branch = simulation.branches.get(name) + if branch is not None and getattr(branch, "branch_period", None) != period: + del simulation.branches[name] + branch = simulation.get_branch(name) + branch.branch_period = period + return branch diff --git a/policyengine_us/variables/gov/hhs/medicaid/medicaid_enrolled_for_ssi_state_supplement.py b/policyengine_us/variables/gov/hhs/medicaid/medicaid_enrolled_for_ssi_state_supplement.py index e9256eca27c..d02f1625c4e 100644 --- a/policyengine_us/variables/gov/hhs/medicaid/medicaid_enrolled_for_ssi_state_supplement.py +++ b/policyengine_us/variables/gov/hhs/medicaid/medicaid_enrolled_for_ssi_state_supplement.py @@ -35,12 +35,17 @@ def formula(person, period, parameters): # age/blindness/disability independently of SNAP/TANF. Keep all other # Medicaid conditions, take-up and supplied inputs in a private branch. branch_name = f"{simulation.branch_name}_ssi_state_supplement_medicaid_{period}" - branch = simulation.get_branch(branch_name) try: - branch.set_input( - "medicaid_community_engagement_pass_through_eligible", - period.first_month, - np.zeros(person.count, dtype=bool), + branch = get_override_branch( + simulation, + branch_name, + period, + { + "medicaid_community_engagement_pass_through_eligible": ( + period.first_month, + np.zeros(person.count, dtype=bool), + ) + }, ) return branch.calculate("medicaid_enrolled", period) finally: diff --git a/policyengine_us/variables/gov/irs/credits/ctc/refundable/ctc_limiting_tax_liability.py b/policyengine_us/variables/gov/irs/credits/ctc/refundable/ctc_limiting_tax_liability.py index 796ce992c1d..afd5737e9a4 100644 --- a/policyengine_us/variables/gov/irs/credits/ctc/refundable/ctc_limiting_tax_liability.py +++ b/policyengine_us/variables/gov/irs/credits/ctc/refundable/ctc_limiting_tax_liability.py @@ -6,28 +6,26 @@ class ctc_limiting_tax_liability(Variable): entity = TaxUnit label = "CTC-limiting tax liability" unit = USD - documentation = "The tax liability used to determine the maximum amount of the non-refundable CTC. Excludes SALT from all calculations (this is an inaccuracy required to avoid circular dependencies)." + documentation = ( + "The tax liability that limits the non-refundable Child Tax Credit: " + "income tax before credits (regular tax plus alternative minimum tax), " + "less the other non-refundable credits." + ) definition_period = YEAR + reference = ( + # 26 U.S.C. 26(a): credits in this subpart are limited to regular tax + # liability (26(b)(1): the tax imposed by chapter 1) plus the tax + # imposed by section 55(a). + "https://www.law.cornell.edu/uscode/text/26/26#a", + # 2025 Schedule 8812 instructions, Credit Limit Worksheet A: line 1 is + # the amount from Form 1040 line 18. + "https://www.irs.gov/instructions/i1040s8", + ) def formula(tax_unit, period, parameters): - simulation = tax_unit.simulation - no_salt_branch = simulation.get_branch("no_salt") - no_salt_branch.set_input("salt_deduction", period, np.zeros(tax_unit.count)) - # Propagate the parent's itemization determination so the - # no_salt branch doesn't re-enter - # `tax_unit_itemizes` -> `tax_liability_if_itemizing` -> - # `income_tax` -> `refundable_ctc`, which forms a cycle - # (issue #8059). The parent's value has already been computed - # by the time we get here: either set as input on the - # itemizing / not_itemizing branch, or computed and cached on - # the top-level sim before `refundable_ctc` was reached (the - # `income_tax_before_credits` branch of - # `income_tax_before_refundable_credits` runs first). - itemizes = tax_unit("tax_unit_itemizes", period) - no_salt_branch.set_input("tax_unit_itemizes", period, itemizes) - tax_liability_before_credits = no_salt_branch.calculate( - "income_tax_before_credits", period - ) + # The tax on taxable income reflects every itemized deduction, + # including state and local taxes. + tax_liability_before_credits = tax_unit("income_tax_before_credits", period) non_refundable_credits = parameters(period).gov.irs.credits.non_refundable non_refundable_credits_ex_ctc = [ x for x in non_refundable_credits if x != "non_refundable_ctc" diff --git a/policyengine_us/variables/gov/irs/income/taxable_income/deductions/tax_liability_if_itemizing.py b/policyengine_us/variables/gov/irs/income/taxable_income/deductions/tax_liability_if_itemizing.py index 775823dd9c3..bbfb5c097ba 100644 --- a/policyengine_us/variables/gov/irs/income/taxable_income/deductions/tax_liability_if_itemizing.py +++ b/policyengine_us/variables/gov/irs/income/taxable_income/deductions/tax_liability_if_itemizing.py @@ -10,8 +10,10 @@ class tax_liability_if_itemizing(Variable): def formula(tax_unit, period, parameters): simulation = tax_unit.simulation - itemized_branch = simulation.get_branch("itemizing") - itemized_branch.set_input( - "tax_unit_itemizes", period, np.ones((tax_unit.count,), dtype=bool) + itemized_branch = get_override_branch( + simulation, + "itemizing", + period, + {"tax_unit_itemizes": np.ones((tax_unit.count,), dtype=bool)}, ) return itemized_branch.calculate("income_tax", period) diff --git a/policyengine_us/variables/gov/irs/income/taxable_income/deductions/tax_liability_if_not_itemizing.py b/policyengine_us/variables/gov/irs/income/taxable_income/deductions/tax_liability_if_not_itemizing.py index eca0caec429..07279c5639b 100644 --- a/policyengine_us/variables/gov/irs/income/taxable_income/deductions/tax_liability_if_not_itemizing.py +++ b/policyengine_us/variables/gov/irs/income/taxable_income/deductions/tax_liability_if_not_itemizing.py @@ -11,10 +11,10 @@ class tax_liability_if_not_itemizing(Variable): def formula(tax_unit, period, parameters): simulation = tax_unit.simulation - non_itemized_branch = simulation.get_branch("not_itemizing") - non_itemized_branch.set_input( - "tax_unit_itemizes", + non_itemized_branch = get_override_branch( + simulation, + "not_itemizing", period, - np.zeros((tax_unit.count,), dtype=bool), + {"tax_unit_itemizes": np.zeros((tax_unit.count,), dtype=bool)}, ) return non_itemized_branch.calculate("income_tax", period) diff --git a/policyengine_us/variables/gov/states/al/tax/income/deductions/federal_income_tax/al_federal_income_tax_deduction.py b/policyengine_us/variables/gov/states/al/tax/income/deductions/federal_income_tax/al_federal_income_tax_deduction.py index e8e93ffa639..46a0b93b1c8 100644 --- a/policyengine_us/variables/gov/states/al/tax/income/deductions/federal_income_tax/al_federal_income_tax_deduction.py +++ b/policyengine_us/variables/gov/states/al/tax/income/deductions/federal_income_tax/al_federal_income_tax_deduction.py @@ -52,7 +52,7 @@ def formula(tax_unit, period, parameters): actual_non_refundable_cdcc = tax_unit("cdcc", period) * (not cdcc_is_refundable) simulation = tax_unit.simulation - branch = simulation.get_branch("al_2020_irc") + branch = get_branch_for_period(simulation, "al_2020_irc", period) branch.tax_benefit_system = get_2020_irc_tbs(simulation.tax_benefit_system) for variable in branch.tax_benefit_system.variables: if any(key in variable for key in ("ctc", "cdcc", "eitc")): diff --git a/policyengine_us/variables/gov/states/de/tax/income/credits/eitc/refundability_calculation/de_income_tax_if_claiming_non_refundable_eitc.py b/policyengine_us/variables/gov/states/de/tax/income/credits/eitc/refundability_calculation/de_income_tax_if_claiming_non_refundable_eitc.py index 94782f0ddac..cccd12d0820 100644 --- a/policyengine_us/variables/gov/states/de/tax/income/credits/eitc/refundability_calculation/de_income_tax_if_claiming_non_refundable_eitc.py +++ b/policyengine_us/variables/gov/states/de/tax/income/credits/eitc/refundability_calculation/de_income_tax_if_claiming_non_refundable_eitc.py @@ -11,10 +11,10 @@ class de_income_tax_if_claiming_non_refundable_eitc(Variable): def formula(tax_unit, period, parameters): simulation = tax_unit.simulation - non_refundable_branch = simulation.get_branch("de_non_refundable_eitc") - non_refundable_branch.set_input( - "de_claims_refundable_eitc", + non_refundable_branch = get_override_branch( + simulation, + "de_non_refundable_eitc", period, - np.zeros((tax_unit.count,), dtype=bool), + {"de_claims_refundable_eitc": np.zeros((tax_unit.count,), dtype=bool)}, ) return non_refundable_branch.calculate("de_income_tax", period) diff --git a/policyengine_us/variables/gov/states/de/tax/income/credits/eitc/refundability_calculation/de_income_tax_if_claiming_refundable_eitc.py b/policyengine_us/variables/gov/states/de/tax/income/credits/eitc/refundability_calculation/de_income_tax_if_claiming_refundable_eitc.py index cec2dfa4451..df3161e0ded 100644 --- a/policyengine_us/variables/gov/states/de/tax/income/credits/eitc/refundability_calculation/de_income_tax_if_claiming_refundable_eitc.py +++ b/policyengine_us/variables/gov/states/de/tax/income/credits/eitc/refundability_calculation/de_income_tax_if_claiming_refundable_eitc.py @@ -11,10 +11,10 @@ class de_income_tax_if_claiming_refundable_eitc(Variable): def formula(tax_unit, period, parameters): simulation = tax_unit.simulation - refundable_branch = simulation.get_branch("de_refundable_eitc") - refundable_branch.set_input( - "de_claims_refundable_eitc", + refundable_branch = get_override_branch( + simulation, + "de_refundable_eitc", period, - np.ones((tax_unit.count,), dtype=bool), + {"de_claims_refundable_eitc": np.ones((tax_unit.count,), dtype=bool)}, ) return refundable_branch.calculate("de_income_tax", period) diff --git a/policyengine_us/variables/gov/states/id/tax/income/id_income_tax_if_receiving_aged_or_disabled_credit.py b/policyengine_us/variables/gov/states/id/tax/income/id_income_tax_if_receiving_aged_or_disabled_credit.py index c024963bf58..cc8c3b9f9d4 100644 --- a/policyengine_us/variables/gov/states/id/tax/income/id_income_tax_if_receiving_aged_or_disabled_credit.py +++ b/policyengine_us/variables/gov/states/id/tax/income/id_income_tax_if_receiving_aged_or_disabled_credit.py @@ -13,10 +13,14 @@ class id_income_tax_if_receiving_aged_or_disabled_credit(Variable): def formula(tax_unit, period, parameters): simulation = tax_unit.simulation - branch = simulation.get_branch("id_receives_aged_or_disabled_credit_branch") - branch.set_input( - "id_receives_aged_or_disabled_credit", + branch = get_override_branch( + simulation, + "id_receives_aged_or_disabled_credit_branch", period, - np.ones((tax_unit.count,), dtype=bool), + { + "id_receives_aged_or_disabled_credit": np.ones( + (tax_unit.count,), dtype=bool + ) + }, ) return branch.calculate("id_income_tax", period) diff --git a/policyengine_us/variables/gov/states/id/tax/income/id_income_tax_if_receiving_aged_or_disabled_deduction.py b/policyengine_us/variables/gov/states/id/tax/income/id_income_tax_if_receiving_aged_or_disabled_deduction.py index 259be96edbe..8f55c47de8a 100644 --- a/policyengine_us/variables/gov/states/id/tax/income/id_income_tax_if_receiving_aged_or_disabled_deduction.py +++ b/policyengine_us/variables/gov/states/id/tax/income/id_income_tax_if_receiving_aged_or_disabled_deduction.py @@ -13,10 +13,14 @@ class id_income_tax_if_receiving_aged_or_disabled_deduction(Variable): def formula(tax_unit, period, parameters): simulation = tax_unit.simulation - branch = simulation.get_branch("id_receives_aged_or_disabled_deduction_branch") - branch.set_input( - "id_receives_aged_or_disabled_credit", + branch = get_override_branch( + simulation, + "id_receives_aged_or_disabled_deduction_branch", period, - np.zeros((tax_unit.count,), dtype=bool), + { + "id_receives_aged_or_disabled_credit": np.zeros( + (tax_unit.count,), dtype=bool + ) + }, ) return branch.calculate("id_income_tax", period) diff --git a/policyengine_us/variables/gov/states/mo/dss/tanf/assistance_unit/non_parent_caretaker/mo_tanf_if_non_parent_caretaker_excluded.py b/policyengine_us/variables/gov/states/mo/dss/tanf/assistance_unit/non_parent_caretaker/mo_tanf_if_non_parent_caretaker_excluded.py index c1e0900e7c5..a2cfb749a93 100644 --- a/policyengine_us/variables/gov/states/mo/dss/tanf/assistance_unit/non_parent_caretaker/mo_tanf_if_non_parent_caretaker_excluded.py +++ b/policyengine_us/variables/gov/states/mo/dss/tanf/assistance_unit/non_parent_caretaker/mo_tanf_if_non_parent_caretaker_excluded.py @@ -19,12 +19,16 @@ def formula(spm_unit, period, parameters): # mo_tanf_if_non_parent_caretaker_included. simulation = spm_unit.simulation branch_name = f"{simulation.branch_name}_mo_tanf_npcr_excluded_{period}" - branch = simulation.get_branch(branch_name) try: - branch.set_input( - "mo_tanf_non_parent_caretaker_included", + branch = get_override_branch( + simulation, + branch_name, period, - np.zeros(spm_unit.count, dtype=bool), + { + "mo_tanf_non_parent_caretaker_included": np.zeros( + spm_unit.count, dtype=bool + ) + }, ) return branch.calculate("mo_tanf", period) finally: diff --git a/policyengine_us/variables/gov/states/mo/dss/tanf/assistance_unit/non_parent_caretaker/mo_tanf_if_non_parent_caretaker_included.py b/policyengine_us/variables/gov/states/mo/dss/tanf/assistance_unit/non_parent_caretaker/mo_tanf_if_non_parent_caretaker_included.py index 75d99502077..3ffd42c9ac6 100644 --- a/policyengine_us/variables/gov/states/mo/dss/tanf/assistance_unit/non_parent_caretaker/mo_tanf_if_non_parent_caretaker_included.py +++ b/policyengine_us/variables/gov/states/mo/dss/tanf/assistance_unit/non_parent_caretaker/mo_tanf_if_non_parent_caretaker_included.py @@ -20,9 +20,13 @@ def formula(spm_unit, period, parameters): needy = spm_unit("mo_tanf_non_parent_caretaker_needy", period) simulation = spm_unit.simulation branch_name = f"{simulation.branch_name}_mo_tanf_npcr_included_{period}" - branch = simulation.get_branch(branch_name) try: - branch.set_input("mo_tanf_non_parent_caretaker_included", period, needy) + branch = get_override_branch( + simulation, + branch_name, + period, + {"mo_tanf_non_parent_caretaker_included": needy}, + ) return branch.calculate("mo_tanf", period) finally: # A branch clones cached arrays; drop it once the grant is read. diff --git a/policyengine_us/variables/gov/states/ny/tax/income/credits/ctc/ny_ctc_pre_2024.py b/policyengine_us/variables/gov/states/ny/tax/income/credits/ctc/ny_ctc_pre_2024.py index 0688b6389bb..c331a3336b8 100644 --- a/policyengine_us/variables/gov/states/ny/tax/income/credits/ctc/ny_ctc_pre_2024.py +++ b/policyengine_us/variables/gov/states/ny/tax/income/credits/ctc/ny_ctc_pre_2024.py @@ -29,7 +29,7 @@ def formula(tax_unit, period, parameters): # Initialize pre-TCJA CTC branch with cached pinned parameters # (one clone per process; see tools/pinned_tbs.py, issue #8114). simulation = tax_unit.simulation - pre_tcja_ctc = simulation.get_branch("pre_tcja_ctc") + pre_tcja_ctc = get_branch_for_period(simulation, "pre_tcja_ctc", period) pre_tcja_ctc.tax_benefit_system = get_pre_tcja_ctc_tbs( simulation.tax_benefit_system ) diff --git a/policyengine_us/variables/gov/states/ny/tax/income/credits/ctc/ny_ctc_pre_2024_eligible.py b/policyengine_us/variables/gov/states/ny/tax/income/credits/ctc/ny_ctc_pre_2024_eligible.py index 3364c58fd18..c7c3b398e65 100644 --- a/policyengine_us/variables/gov/states/ny/tax/income/credits/ctc/ny_ctc_pre_2024_eligible.py +++ b/policyengine_us/variables/gov/states/ny/tax/income/credits/ctc/ny_ctc_pre_2024_eligible.py @@ -28,7 +28,7 @@ def formula(tax_unit, period, parameters): # Initialize pre-TCJA CTC branch for eligibility check with # cached pinned parameters (see tools/pinned_tbs.py, issue #8114). simulation = tax_unit.simulation - pre_tcja_ctc = simulation.get_branch("pre_tcja_ctc") + pre_tcja_ctc = get_branch_for_period(simulation, "pre_tcja_ctc", period) pre_tcja_ctc.tax_benefit_system = get_pre_tcja_ctc_tbs( simulation.tax_benefit_system ) diff --git a/policyengine_us/variables/gov/states/ny/tax/income/credits/ny_eitc.py b/policyengine_us/variables/gov/states/ny/tax/income/credits/ny_eitc.py index a45b7d86eb2..20c38b23cdb 100644 --- a/policyengine_us/variables/gov/states/ny/tax/income/credits/ny_eitc.py +++ b/policyengine_us/variables/gov/states/ny/tax/income/credits/ny_eitc.py @@ -23,7 +23,7 @@ def formula(tax_unit, period, parameters): # does not conform. Recompute the federal EITC with # pre-ARPA (2020) parameter values. simulation = tax_unit.simulation - branch = simulation.get_branch("ny_pre_arpa_eitc") + branch = get_branch_for_period(simulation, "ny_pre_arpa_eitc", period) branch.tax_benefit_system = get_pre_arpa_eitc_tbs( simulation.tax_benefit_system ) diff --git a/policyengine_us/variables/gov/states/va/tax/income/credits/eitc/refundability_calculation/va_income_tax_if_claiming_non_refundable_eitc.py b/policyengine_us/variables/gov/states/va/tax/income/credits/eitc/refundability_calculation/va_income_tax_if_claiming_non_refundable_eitc.py index cef69a3c26c..2de32bddeb0 100644 --- a/policyengine_us/variables/gov/states/va/tax/income/credits/eitc/refundability_calculation/va_income_tax_if_claiming_non_refundable_eitc.py +++ b/policyengine_us/variables/gov/states/va/tax/income/credits/eitc/refundability_calculation/va_income_tax_if_claiming_non_refundable_eitc.py @@ -11,10 +11,10 @@ class va_income_tax_if_claiming_non_refundable_eitc(Variable): def formula(tax_unit, period, parameters): simulation = tax_unit.simulation - non_refundable_branch = simulation.get_branch("va_non_refundable_eitc") - non_refundable_branch.set_input( - "va_claims_refundable_eitc", + non_refundable_branch = get_override_branch( + simulation, + "va_non_refundable_eitc", period, - np.zeros((tax_unit.count,), dtype=bool), + {"va_claims_refundable_eitc": np.zeros((tax_unit.count,), dtype=bool)}, ) return non_refundable_branch.calculate("va_income_tax", period) diff --git a/policyengine_us/variables/gov/states/va/tax/income/credits/eitc/refundability_calculation/va_income_tax_if_claiming_refundable_eitc.py b/policyengine_us/variables/gov/states/va/tax/income/credits/eitc/refundability_calculation/va_income_tax_if_claiming_refundable_eitc.py index b510db847a2..ff7ea5ee356 100644 --- a/policyengine_us/variables/gov/states/va/tax/income/credits/eitc/refundability_calculation/va_income_tax_if_claiming_refundable_eitc.py +++ b/policyengine_us/variables/gov/states/va/tax/income/credits/eitc/refundability_calculation/va_income_tax_if_claiming_refundable_eitc.py @@ -11,10 +11,10 @@ class va_income_tax_if_claiming_refundable_eitc(Variable): def formula(tax_unit, period, parameters): simulation = tax_unit.simulation - refundable_branch = simulation.get_branch("va_refundable_eitc") - refundable_branch.set_input( - "va_claims_refundable_eitc", + refundable_branch = get_override_branch( + simulation, + "va_refundable_eitc", period, - np.ones((tax_unit.count,), dtype=bool), + {"va_claims_refundable_eitc": np.ones((tax_unit.count,), dtype=bool)}, ) return refundable_branch.calculate("va_income_tax", period)