Add PUF + source impute modules, fix AGI ceiling (issue #530)

[baogorek]

Summary

Fixes #530 — extracts PUF cloning and source imputation into standalone modules, replacing inline logic in extended_cps.py.

• puf_impute.py — PUF clone + batched QRF using demographic predictors. Uses stratified subsample (20K target) that force-includes top 0.5% by AGI, preserving the high-income tail. Includes weeks_unemployed imputation for the PUF half.
• source_impute.py — ACS re-imputation with state_fips as a QRF predictor for regional variation; SIPP and SCF imputation using demographic and financial predictors only (surveys lack state identifiers).
• extended_cps.py refactor — 443 → 75 lines, delegates to puf_clone_dataset().
• unified_calibration.py — New --puf-dataset, --skip-puf, --skip-source-impute flags; _build_puf_cloned_dataset() helper; pipeline: geography → source impute → PUF clone → double_geography_for_puf() → takeup rerandomize → matrix build → L0 calibrate.
• 21 new tests for both modules (mock data, skip_qrf=True path).

Also fixes #539 — sets the environment variable DC_STATEHOOD=1 so that -us recognizes DC as a state, including it in every file that pulls all states throughout the data pipeline.

Notes

I'm seeing a $400M max AGI in ExtendedCPS_2024, per the manual test below.

Test plan

• pytest policyengine_us_data/tests/test_calibration/ — 63/63 pass (21 new + 42 existing)
• black . -l 79 --check — formatting clean
• Manual: ExtendedCPS_2024().generate() and verify max AGI > $10M

🤖 Generated with Claude Code

[MaxGhenis]

Why assign and impute from random states for the SCF and SIPP?

[MaxGhenis]

Can you mark as draft until we can merge

[baogorek]

@MaxGhenis good call out on the random state simulation for SIPP and SCF. The argument was "predictor consistency" but they were already different predictor sets. So the state simulation is out now.

Another thing I added was stratified sampling of the top .5% incomes in the PUF (and uniformly sampling the rest) for the purposes of imputation. Ignoring the weights shouldn't be a problem statistically since we just want a good P(Y | X) and thus we just need a wide range of Xs (we're not sampling proportional to residual size or something).

To push things along, I also added you to as a reviewer to the prerequisite PR #531 alongside Pavel. With your development, my cherry picking and refinements, and then Pavel's feedback, I think it's probably pretty close. But it is a large and very important PR.

[baogorek]

I'm putting this back in draft status until this issue is solved. #539

DC has no representation!

[baogorek]

Thanks @juaristi22 for getting this ready to go. I had claude already going on it and here's what it said about your fix:

Ha — she took a completely different (and arguably cleaner) approach. Instead of patching each script individually, she set DC_STATEHOOD=1 in
  storage/__init__.py, which tells the us library to include DC in STATES_AND_TERRITORIES globally. That means us.states.lookup("DC") works and DC shows up in
   the state loop — no per-script fixes needed. 

Nice work.

Update: I did add one more commit to make it explicit: + [us.states.DC] in the state lists and a getattr fallback in fetch_block_population. Works regardless of import order. (So at least my claude session wasn't totally wasted)

[MaxGhenis]

All CI checks pass. Code reviewed and fixes applied: corrected failing test (state_fips intentionally excluded from PUF predictors), fixed _person_state_fips fallback for unequal household sizes, accurate docstrings for state predictor usage, extracted tenure_type helper, narrowed exception handling, cleaned up code.

[MaxGhenis]

Post-review fixes applied in two commits before merge:

Bug fixes & accuracy:

• Fixed failing test_demographic_predictors_includes_state — state_fips is intentionally excluded from PUF DEMOGRAPHIC_PREDICTORS (PUF has no state identifier). Test now asserts exclusion.
• Fixed _person_state_fips fallback to handle unequal household sizes (was silently producing wrong-length arrays).
• Updated source_impute.py module docstring and log messages — only ACS uses state_fips as a predictor; SIPP/SCF lack state identifiers.
• Narrowed except Exception in SIPP asset imputation to specific types (FileNotFoundError, KeyError, ValueError, OSError).

Code cleanup:

• Extracted duplicate tenure_type string→int mapping into _encode_tenure_type helper.
• Replaced unnecessary getattr(args, ...) with direct attribute access in main().
• Removed unused imports (time, pytest, Microsimulation, PUF constants).
• Removed redundant person_id branch (already covered by _id check).
• Inlined single-use local aliases, merged implicit string concatenations.

All 31 calibration tests pass.