Add trust fund revenue variables with LSR recursion fix and tier separation

changelog_entry.yaml

@@ -0,0 +1,8 @@
+- bump: minor
+  changes:
+    added:
+      - Trust fund revenue variables (tob_revenue_total, tob_revenue_oasdi, tob_revenue_medicare_hi) using exact branching methodology
+      - Tier 1 and tier 2 taxable Social Security variables for proper OASDI vs Medicare HI allocation
+      - LSR recursion guard to prevent infinite loops when branches calculate variables
+    fixed:
+      - Labor supply behavioral response infinite recursion bug

policyengine_us/tests/policy/baseline/gov/ssa/revenue/tob_revenue_total.yaml

@@ -0,0 +1,19 @@
+- name: TOB revenue - single retiree with SS and wages
+  period: 2024
+  absolute_error_margin: 100
+  input:
+    people:
+      person1:
+        age: 67
+        social_security: 30_000
+        employment_income: 50_000
+    tax_units:
+      tax_unit:
+        members: [person1]
+        filing_status: SINGLE
+    households:
+      household:
+        members: [person1]
+  output:
+    # Actual value is around $4,240
+    tob_revenue_total: 4240

policyengine_us/variables/gov/irs/income/taxable_income/adjusted_gross_income/irs_gross_income/social_security/taxable_social_security_tier_1.py

@@ -0,0 +1,58 @@
+from policyengine_us.model_api import *
+
+
+class taxable_social_security_tier_1(Variable):
+    value_type = float
+    entity = TaxUnit
+    definition_period = YEAR
+    label = "Taxable Social Security (tier 1)"
+    documentation = "Taxable Social Security from 0-50% taxation tier, credited to OASDI trust funds"
+    unit = USD
+    reference = "https://www.law.cornell.edu/uscode/text/26/86#a_1"
+
+    def formula(tax_unit, period, parameters):
+        p = parameters(period).gov.irs.social_security.taxability
+        gross_ss = tax_unit("tax_unit_social_security", period)
+        combined_income = tax_unit(
+            "tax_unit_combined_income_for_social_security_taxability", period
+        )
+        filing_status = tax_unit("filing_status", period)
+        status = filing_status.possible_values
+        separate = filing_status == status.SEPARATE
+        cohabitating = tax_unit("cohabitating_spouses", period)
+
+        base_amount = where(
+            separate & cohabitating,
+            p.threshold.base.separate_cohabitating,
+            p.threshold.base.main[filing_status],
+        )
+        adjusted_base_amount = where(
+            separate & cohabitating,
+            p.threshold.adjusted_base.separate_cohabitating,
+            p.threshold.adjusted_base.main[filing_status],
+        )
+
+        under_first_threshold = combined_income < base_amount
+        under_second_threshold = combined_income < adjusted_base_amount
+
+        combined_income_excess = tax_unit(
+            "tax_unit_ss_combined_income_excess", period
+        )
+
+        # Tier 1 amount (IRC §86(a)(1))
+        amount_under_paragraph_1 = min_(
+            p.rate.base.benefit_cap * gross_ss,
+            p.rate.base.excess * combined_income_excess,
+        )
+
+        # Bracket amount when in tier 2 (IRC §86(a)(2)(A)(ii))
+        bracket_amount = min_(
+            amount_under_paragraph_1,
+            p.rate.additional.bracket * (adjusted_base_amount - base_amount),
+        )
+
+        return select(
+            [under_first_threshold, under_second_threshold],
+            [0, amount_under_paragraph_1],
+            default=bracket_amount,
+        )

policyengine_us/variables/gov/irs/income/taxable_income/adjusted_gross_income/irs_gross_income/social_security/taxable_social_security_tier_2.py

@@ -0,0 +1,15 @@
+from policyengine_us.model_api import *
+
+
+class taxable_social_security_tier_2(Variable):
+    value_type = float
+    entity = TaxUnit
+    definition_period = YEAR
+    label = "Taxable Social Security (tier 2)"
+    documentation = "Taxable Social Security from 50-85% taxation tier, credited to Medicare HI trust fund"
+    unit = USD
+    reference = "https://www.law.cornell.edu/uscode/text/26/86#a_2"
+    subtracts = [
+        "tax_unit_taxable_social_security",
+        "taxable_social_security_tier_1",
+    ]

policyengine_us/variables/gov/simulation/labor_supply_response/labor_supply_behavioral_response.py

@@ -16,53 +16,68 @@ def formula(person, period, parameters):
         if p.elasticities.income == 0 and p.elasticities.substitution.all == 0:
             return 0
 
-        measurement_branch = simulation.get_branch(
-            "lsr_measurement", clone_system=True
-        )  # A branch without LSRs
-        baseline_branch = simulation.get_branch("baseline").get_branch(
-            "baseline_lsr_measurement", clone_system=True
-        )  # Already created by default
-        baseline_branch.tax_benefit_system.parameters.simulation = (
-            measurement_branch.tax_benefit_system.parameters.simulation
-        )
+        # Guard against re-entry (prevents recursion when branches calculate variables)
+        if (
+            hasattr(simulation, "_lsr_calculating")
+            and simulation._lsr_calculating
+        ):
+            return 0
 
-        # (system with LSRs) <- (system without LSRs used to calculate LSRs)
-        #                      |
-        #                      * -(baseline system without LSRs used to calculate LSRs)
+        # Mark that we're calculating LSR
+        simulation._lsr_calculating = True
 
-        for branch in [measurement_branch, baseline_branch]:
-            branch.tax_benefit_system.neutralize_variable(
-                "employment_income_behavioral_response"
-            )
-            branch.tax_benefit_system.neutralize_variable(
-                "self_employment_income_behavioral_response"
-            )
-            branch.set_input(
-                "employment_income_before_lsr",
-                period,
-                person("employment_income_before_lsr", period),
+        try:
+            measurement_branch = simulation.get_branch(
+                "lsr_measurement", clone_system=True
+            )  # A branch without LSRs
+            baseline_branch = simulation.get_branch("baseline").get_branch(
+                "baseline_lsr_measurement", clone_system=True
+            )  # Already created by default
+            baseline_branch.tax_benefit_system.parameters.simulation = (
+                measurement_branch.tax_benefit_system.parameters.simulation
             )
-            branch.set_input(
-                "self_employment_income_before_lsr",
+
+            # (system with LSRs) <- (system without LSRs used to calculate LSRs)
+            #                      |
+            #                      * -(baseline system without LSRs used to calculate LSRs)
+
+            for branch in [measurement_branch, baseline_branch]:
+                branch.tax_benefit_system.neutralize_variable(
+                    "employment_income_behavioral_response"
+                )
+                branch.tax_benefit_system.neutralize_variable(
+                    "self_employment_income_behavioral_response"
+                )
+                branch.set_input(
+                    "employment_income_before_lsr",
+                    period,
+                    person("employment_income_before_lsr", period),
+                )
+                branch.set_input(
+                    "self_employment_income_before_lsr",
+                    period,
+                    person("self_employment_income_before_lsr", period),
+                )
+
+            response = add(
+                person,
                 period,
-                person("self_employment_income_before_lsr", period),
+                [
+                    "income_elasticity_lsr",
+                    "substitution_elasticity_lsr",
+                ],
             )
+            simulation = person.simulation
+            del simulation.branches["baseline"].branches[
+                "baseline_lsr_measurement"
+            ]
+            del simulation.branches["lsr_measurement"]
 
-        response = add(
-            person,
-            period,
-            [
-                "income_elasticity_lsr",
-                "substitution_elasticity_lsr",
-            ],
-        )
-        simulation = person.simulation
-        del simulation.branches["baseline"].branches[
-            "baseline_lsr_measurement"
-        ]
-        del simulation.branches["lsr_measurement"]
+            simulation.macro_cache_read = False
+            simulation.macro_cache_write = False
 
-        simulation.macro_cache_read = False
-        simulation.macro_cache_write = False
+            return response
 
-        return response
+        finally:
+            # Clear the re-entry guard
+            simulation._lsr_calculating = False

policyengine_us/variables/gov/ssa/revenue/tob_revenue_medicare_hi.py

@@ -0,0 +1,31 @@
+from policyengine_us.model_api import *
+
+
+class tob_revenue_medicare_hi(Variable):
+    value_type = float
+    entity = TaxUnit
+    definition_period = YEAR
+    label = "Medicare HI trust fund revenue from SS benefit taxation (tier 2)"
+    documentation = "Tax revenue from tier 2 (50-85%) Social Security benefit taxation credited to Medicare HI trust fund"
+    unit = USD
+
+    def formula(tax_unit, period, parameters):
+        """
+        Calculate Medicare HI trust fund revenue from tier 2 SS taxation.
+
+        Allocates total TOB revenue to Medicare HI based on tier 2's proportion
+        of total taxable SS.
+        """
+        # Get total TOB revenue
+        total_tob = tax_unit("tob_revenue_total", period)
+
+        # Get tier amounts
+        tier1 = tax_unit("taxable_social_security_tier_1", period)
+        tier2 = tax_unit("taxable_social_security_tier_2", period)
+        total_taxable = tier1 + tier2
+
+        # Allocate total TOB based on tier 2 proportion
+        # Use where to handle division by zero
+        medicare_share = where(total_taxable > 0, tier2 / total_taxable, 0)
+
+        return total_tob * medicare_share

policyengine_us/variables/gov/ssa/revenue/tob_revenue_oasdi.py

@@ -0,0 +1,31 @@
+from policyengine_us.model_api import *
+
+
+class tob_revenue_oasdi(Variable):
+    value_type = float
+    entity = TaxUnit
+    definition_period = YEAR
+    label = "OASDI trust fund revenue from SS benefit taxation (tier 1)"
+    documentation = "Tax revenue from tier 1 (0-50%) Social Security benefit taxation credited to OASDI trust funds"
+    unit = USD
+
+    def formula(tax_unit, period, parameters):
+        """
+        Calculate OASDI trust fund revenue from tier 1 SS taxation.
+
+        Allocates total TOB revenue to OASDI based on tier 1's proportion
+        of total taxable SS.
+        """
+        # Get total TOB revenue
+        total_tob = tax_unit("tob_revenue_total", period)
+
+        # Get tier amounts
+        tier1 = tax_unit("taxable_social_security_tier_1", period)
+        tier2 = tax_unit("taxable_social_security_tier_2", period)
+        total_taxable = tier1 + tier2
+
+        # Allocate total TOB based on tier 1 proportion
+        # Use where to handle division by zero
+        oasdi_share = where(total_taxable > 0, tier1 / total_taxable, 0)
+
+        return total_tob * oasdi_share

policyengine_us/variables/gov/ssa/revenue/tob_revenue_total.py

@@ -0,0 +1,44 @@
+from policyengine_us.model_api import *
+
+
+class tob_revenue_total(Variable):
+    value_type = float
+    entity = TaxUnit
+    definition_period = YEAR
+    label = "Total trust fund revenue from SS benefit taxation"
+    documentation = "Tax revenue from taxation of Social Security benefits using branching methodology"
+    unit = USD
+
+    def formula(tax_unit, period, parameters):
+        """
+        Calculate trust fund revenue using branching + neutralization.
+
+        This is the CORRECT way to isolate TOB revenue, superior to the
+        average effective tax rate approximation.
+        """
+        sim = tax_unit.simulation
+
+        # Calculate income tax WITH taxable SS
+        income_tax_with = tax_unit("income_tax", period)
+
+        # Create branch and neutralize taxable SS
+        branch = sim.get_branch("tob_calc", clone_system=True)
+        branch.tax_benefit_system.neutralize_variable(
+            "tax_unit_taxable_social_security"
+        )
+
+        # Delete all calculated variables to force recalculation
+        for var_name in list(branch.tax_benefit_system.variables.keys()):
+            if var_name not in branch.input_variables:
+                try:
+                    branch.delete_arrays(var_name)
+                except:
+                    pass
+
+        # Recalculate income tax without taxable SS
+        income_tax_without = branch.tax_unit("income_tax", period)
+
+        # Clean up branch
+        del sim.branches["tob_calc"]
+
+        return income_tax_with - income_tax_without