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 with state_fips as predictor. Uses stratified subsample (20K target) that force-includes top 0.5% by AGI, preserving the high-income tail. Includes weeks_unemployed imputation for PUF half.
• source_impute.py — ACS re-imputation with state predictor. SIPP/SCF imputation without state (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)