Calibrate retirement contributions: targets, SS reconciliation, and QRF imputation

[PavelMakarchuk]

Summary

Combines three related pieces of work for retirement contribution calibration (issue #561):

1. Calibration targets (original PR #554)

• Correct traditional IRA target from $25B to $13.2B (SOI 1304)
• Split 401(k) into traditional $482.7B and Roth $85.2B (BEA/FRED × Vanguard share)
• Add Roth IRA target at $35.0B (SOI Tables 5 & 6)
• Add self-employed pension ALD target at $29.5B (SOI 1304)
• Fix RETCB_VAL allocation to use proportional split (traditional/Roth 401k + Roth IRA)

2. SS sub-component reconciliation (from PR #552)

• After PUF imputation overwrites social_security_retirement, reconcile DI/survivor/dependent sub-components
• QRF-based share prediction with age-heuristic fallback
• Drop formula variables from the dataset to avoid stale values

3. QRF imputation for retirement contributions on PUF clone half (new)

Problem: When building the Extended CPS, PUF clone records get PUF-imputed income via QRF, but retirement contribution variables were just blindly duplicated from the CPS donor. This means a PUF clone with $0 wages could have $50k in 401(k) contributions — there's no model linking contributions to income.

Solution: Train a QRF on the CPS half (which has realistic income↔contribution relationships) and predict retirement contributions onto the PUF clone half using PUF-imputed income as input. Follows the exact pattern of the existing _impute_weeks_unemployed().

Variables modeled:

• traditional_401k_contributions
• roth_401k_contributions
• traditional_ira_contributions
• roth_ira_contributions
• self_employed_pension_contributions

Predictors (13 total):

• Demographics (from CPS sim): age, is_male, tax_unit_is_joint, tax_unit_count_dependents, is_tax_unit_head, is_tax_unit_spouse, is_tax_unit_dependent
• Income (from PUF imputations): employment_income, self_employment_income, taxable_interest_income, qualified_dividend_income, taxable_pension_income, social_security

Post-prediction constraints:

• Non-negativity on all values
• 401(k) capped at year-specific limit ($23K + $7.5K catch-up for age ≥ 50 in 2024)
• IRA capped at year-specific limit ($7K + $1K catch-up for age ≥ 50 in 2024)
• 401(k) zeroed out for records with $0 employment income
• SE pension zeroed out for records with $0 self-employment income

Data flow: CPS half keeps original values; PUF half gets CPS-trained QRF predictions. The traditional_ira_contributions was moved out of IMPUTED_VARIABLES (PUF-based QRF) into the CPS-trained model for consistency with the other retirement variables.

Test plan

• 36 unit tests for retirement QRF imputation (constants, limits, constraints, routing, predictor validation)
• 12 existing puf_impute tests pass
• SS reconciliation tests pass
• All tests passing locally (48/48 relevant tests)
• Full Extended CPS build on Modal to verify aggregate weighted totals against calibration targets

Validation

A post-build validation script is provided at validation/validate_retirement_imputation.py. After a full build, run:

python validation/validate_retirement_imputation.py

It checks:

1. Hard constraints: no negative values, no cap violations, $0 contributions for $0 income records
2. Aggregate comparisons: weighted sums vs calibration targets (traditional 401k $482.7B, Roth 401k $85.2B, traditional IRA $13.2B, Roth IRA $35.0B, SE pension $29.5B)
3. CPS vs PUF half distributions: mean values and participation rates per half

Files changed

File	Change
policyengine_us_data/calibration/puf_impute.py	Core implementation: constants, _get_retirement_limits(), _impute_retirement_contributions(), routing in puf_clone_dataset()
policyengine_us_data/tests/test_calibration/test_retirement_imputation.py	36 comprehensive tests
validation/validate_retirement_imputation.py	Post-build validation script
conftest.py (root)	Mock microimpute for Python < 3.12 test environments
policyengine_us_data/tests/test_calibration/conftest.py	Same mock for calibration test directory

Closes #561

🤖 Generated with Claude Code

[baogorek]

Looks good, just please update this one bullet point in the PR body (if this is indeed correct):

Says:

• Remove roth_ira_contributions target ($39B) — structurally $0 due to CPS allocation bug where traditional IRA always exhausts the IRA limit first

But:
code actually keeps roth_ira_contributions and updates it from $39B to $35.0B in both loss.py and etl_national_targets.py (which it can do be