Map SOI retirement contribution targets

MaxGhenis · MaxGhenis · commit 7e22bb9eb619 · 2026-05-21T21:52:35.000-04:00
diff --git a/src/microplex_us/targets/arch.py b/src/microplex_us/targets/arch.py
@@ -396,6 +396,18 @@
         "tip_income_taxpayers",
         "COUNT",
     ),
+    "irs_soi.form_w2_401k_elective_deferrals": (
+        "traditional_401k_contributions",
+        "AMOUNT",
+    ),
+    "irs_soi.form_w2_designated_roth_401k_contributions": (
+        "roth_401k_contributions",
+        "AMOUNT",
+    ),
+    "irs_soi.payments_to_keogh_plan": (
+        "self_employed_pension_contribution_ald",
+        "AMOUNT",
+    ),
     "census_decennial.resident_population": ("population", "COUNT"),
     "census_decennial.occupied_housing_units": ("household_count", "COUNT"),
     "census_pep.resident_population": ("population", "COUNT"),
@@ -630,7 +642,9 @@
     "short_term_capital_gains": EntityType.PERSON,
     "proprietors_income_amount": EntityType.PERSON,
     "rental_income": EntityType.PERSON,
+    "roth_401k_contributions": EntityType.PERSON,
     "self_employment_income": EntityType.PERSON,
+    "self_employed_pension_contribution_ald": EntityType.TAX_UNIT,
     "salt_refund_income": EntityType.PERSON,
     "state_income_tax": EntityType.TAX_UNIT,
     "taxable_ira_distributions": EntityType.PERSON,
@@ -639,6 +653,7 @@
     "taxable_pension_income": EntityType.PERSON,
     "taxable_social_security": EntityType.PERSON,
     "tip_income": EntityType.PERSON,
+    "traditional_401k_contributions": EntityType.PERSON,
     "unemployment_compensation": EntityType.PERSON,
     "medicaid": EntityType.PERSON,
     "social_security": EntityType.PERSON,
@@ -763,8 +778,13 @@
     "tanf_family_count": "TANF family count",
     "tanf_recipient_count": "TANF recipient count",
     "tip_income": "Tip income",
+    "traditional_401k_contributions": "Traditional 401(k) contributions",
     "traditional_ira_contributions": "Traditional IRA contributions",
+    "roth_401k_contributions": "Roth 401(k) contributions",
     "roth_ira_contributions": "Roth IRA contributions",
+    "self_employed_pension_contribution_ald": (
+        "Self-employed pension contribution ALD"
+    ),
 }
 
 ARCH_AGI_BRACKET_LABELS = {
@@ -791,11 +811,16 @@
     "interest_deduction": "interest_paid_deduction_amount",
     "net_capital_gains": "net_capital_gains_amount",
     "real_estate_taxes": "real_estate_taxes_amount",
+    "roth_401k_contributions": "roth_401k_contributions",
+    "self_employed_pension_contribution_ald": (
+        "self_employed_pension_contribution_ald"
+    ),
     "total_self_employment_income": "schedule_c_income_amount",
     "taxable_ira_distributions": "taxable_ira_distributions_amount",
     "taxable_pension_income": "taxable_pension_income_amount",
     "taxable_social_security": "taxable_social_security_amount",
     "tip_income": "tip_income",
+    "traditional_401k_contributions": "traditional_401k_contributions",
     "unemployment_compensation": "unemployment_compensation_amount",
 }
 
@@ -898,6 +923,8 @@
     "population": "Census Population Estimates Program Vintage 2024 age-sex files",
     "real_estate_taxes": "IRS SOI itemized deduction tables or ACS state files",
     "roth_ira_contributions": "IRS SOI IRA contribution tables",
+    "roth_401k_contributions": "IRS SOI Form W-2 Statistics Table 4.B",
+    "self_employed_pension_contribution_ald": "IRS SOI Publication 1304 Table 1.4",
     "state_individual_income_tax_collections": (
         "Census State Tax Collections item T40"
     ),
@@ -919,6 +946,7 @@
     "tanf_recipient_count": "ACF TANF Caseload Data",
     "tip_income": "IRS SOI Form W-2 Statistics",
     "traditional_ira_contributions": "IRS SOI IRA contribution tables",
+    "traditional_401k_contributions": "IRS SOI Form W-2 Statistics Table 4.B",
     "taxable_ira_distributions": "IRS SOI IRA accumulation/distribution tables",
     "taxable_pension_income": "IRS SOI Publication 1304 Table 1.4",
     "taxable_social_security": "IRS SOI Publication 1304 Table 1.4",
diff --git a/tests/targets/test_arch.py b/tests/targets/test_arch.py
@@ -3449,6 +3449,52 @@ def test_arch_target_gap_queue_points_energy_subsidy_households_to_liheap(
     assert row.gap_category == "ready_primary_loader"
 
 
+def test_arch_target_gap_queue_points_retirement_contributions_to_soi(
+    tmp_path,
+):
+    db_path = tmp_path / "arch_targets.db"
+    _create_arch_targets_db(db_path)
+
+    provider = ArchSQLiteTargetProvider(db_path)
+    report = summarize_arch_target_gap_queue(
+        provider,
+        period=2024,
+        profile_name="custom",
+        target_cells=(
+            PolicyEngineUSTargetCell(
+                "traditional_401k_contributions",
+                geo_level="national",
+            ),
+            PolicyEngineUSTargetCell("roth_401k_contributions", geo_level="national"),
+            PolicyEngineUSTargetCell(
+                "self_employed_pension_contribution_ald",
+                geo_level="national",
+            ),
+        ),
+    )
+
+    rows_by_variable = {row.variable: row for row in report.rows}
+    traditional = rows_by_variable["traditional_401k_contributions"]
+    roth = rows_by_variable["roth_401k_contributions"]
+    self_employed = rows_by_variable["self_employed_pension_contribution_ald"]
+
+    assert {row.expected_source for row in report.rows} == {"IRS_SOI"}
+    assert traditional.expected_source_table == "IRS SOI Form W-2 Statistics Table 4.B"
+    assert traditional.expected_arch_variable == "traditional_401k_contributions"
+    assert traditional.expected_entity == "person"
+    assert roth.expected_source_table == "IRS SOI Form W-2 Statistics Table 4.B"
+    assert roth.expected_arch_variable == "roth_401k_contributions"
+    assert roth.expected_entity == "person"
+    assert self_employed.expected_source_table == (
+        "IRS SOI Publication 1304 Table 1.4"
+    )
+    assert self_employed.expected_arch_variable == (
+        "self_employed_pension_contribution_ald"
+    )
+    assert self_employed.expected_entity == "tax_unit"
+    assert {row.gap_category for row in report.rows} == {"ready_primary_loader"}
+
+
 def test_arch_target_gap_queue_points_agi_person_counts_to_soi_table_2(tmp_path):
     db_path = tmp_path / "arch_targets.db"
     _create_arch_targets_db(db_path)
diff --git a/tests/targets/test_arch_facts.py b/tests/targets/test_arch_facts.py
@@ -1343,6 +1343,36 @@ def test_arch_consumer_fact_jsonl_provider_maps_us_admin_source_families(
             value=526,
             unit="usd",
         ),
+        _consumer_fact(
+            "w2-traditional-401k",
+            concept="irs_soi.form_w2_401k_elective_deferrals",
+            domain="form_w2_items",
+            source_name="irs_soi",
+            source_table="Form W-2 Statistics Table 4.B",
+            period={"type": "tax_year", "value": 2024},
+            value=277_859_181_000,
+            unit="usd",
+        ),
+        _consumer_fact(
+            "w2-roth-401k",
+            concept="irs_soi.form_w2_designated_roth_401k_contributions",
+            domain="form_w2_items",
+            source_name="irs_soi",
+            source_table="Form W-2 Statistics Table 4.B",
+            period={"type": "tax_year", "value": 2024},
+            value=32_302_509_000,
+            unit="usd",
+        ),
+        _consumer_fact(
+            "soi-keogh",
+            concept="irs_soi.payments_to_keogh_plan",
+            domain="all_individual_income_tax_returns",
+            source_name="irs_soi",
+            source_table="Publication 1304 Table 1.4",
+            period={"type": "tax_year", "value": 2024},
+            value=30_130_848_000,
+            unit="usd",
+        ),
     ]
     consumer_jsonl.write_text(
         "\n".join(json.dumps(row, sort_keys=True) for row in rows) + "\n"
@@ -1398,11 +1428,26 @@ def test_arch_consumer_fact_jsonl_provider_maps_us_admin_source_families(
                 "geo_level": "national",
                 "domain_variable": "age",
             },
+            {
+                "variable": "traditional_401k_contributions",
+                "geo_level": "national",
+                "domain_variable": None,
+            },
+            {
+                "variable": "roth_401k_contributions",
+                "geo_level": "national",
+                "domain_variable": None,
+            },
+            {
+                "variable": "self_employed_pension_contribution_ald",
+                "geo_level": "national",
+                "domain_variable": None,
+            },
         ),
     )
 
-    assert report.target_cell_count == 12
-    assert report.covered_cell_count == 12
+    assert report.target_cell_count == 15
+    assert report.covered_cell_count == 15
 
     target_set = provider.load_target_set(TargetQuery(period=2024))
     targets_by_arch_variable = {
@@ -1442,6 +1487,17 @@ def test_arch_consumer_fact_jsonl_provider_maps_us_admin_source_families(
         == "social_security_retirement"
     )
     assert targets_by_arch_variable["ssi_payments"].measure == "ssi"
+    traditional_401k = targets_by_arch_variable["traditional_401k_contributions"]
+    assert traditional_401k.measure == "traditional_401k_contributions"
+    assert traditional_401k.entity.value == "person"
+    roth_401k = targets_by_arch_variable["roth_401k_contributions"]
+    assert roth_401k.measure == "roth_401k_contributions"
+    assert roth_401k.entity.value == "person"
+    self_employed_pension = targets_by_arch_variable[
+        "self_employed_pension_contribution_ald"
+    ]
+    assert self_employed_pension.measure == "self_employed_pension_contribution_ald"
+    assert self_employed_pension.entity.value == "tax_unit"
     assert "aca_average_monthly_aptc" not in targets_by_arch_variable
 
 
