Reorder calibration pipeline: clone before PUF, geography before QRF

[MaxGhenis]

Summary

Restructures the unified calibration pipeline so geography is assigned before PUF imputation, enabling state FIPS to be used as a QRF predictor. This improves state-level accuracy of imputed tax variables.

New pipeline flow

Raw CPS (~55K HH)
  → Clone 10x → Assign geography (census block → state, county, CD)
  → PUF clone (2x, second half gets PUF tax variables)
  → QRF impute (with state_fips as predictor)
  → PE simulate (taxes, benefits) per clone
  → Build sparse calibration matrix (~18K targets)
  → L0-regularized calibration

Key changes

• New module puf_impute.py: Extracts PUF clone + QRF imputation from extended_cps.py into a standalone module. DEMOGRAPHIC_PREDICTORS now includes state_fips.
• clone_and_assign.py: Default n_clones reduced from 130 to 10 (fewer clones × more records per clone after PUF doubling). Added double_geography_for_puf() helper to duplicate geography arrays when PUF cloning doubles records.
• unified_calibration.py: Rewritten to implement the new pipeline order. New CLI flags --puf-dataset and --skip-puf. Default input changed from extended_cps_2024.h5 to cps_2024_full.h5 (raw CPS before any PUF processing).
• Pipeline comparison diagram added to docs/pipeline-comparison.png.

Testing

• 8 new tests for puf_impute.py (clone logic with skip_qrf=True)
• 3 new tests for double_geography_for_puf()
• All existing tests updated for new defaults and parameters
• End-to-end run verified with --skip-puf --n-clones 2 --epochs 3 on real CPS data (55K HH, 37K targets, ~2 min)

Test plan

• All 154 unit tests pass
• Black formatting passes (26.1.0)
• End-to-end --skip-puf path works with real data
• End-to-end PUF path (requires PUF dataset access)
• Compare calibration accuracy vs old pipeline

🤖 Generated with Claude Code

[MaxGhenis]

Full pipeline reorder incoming

I'm restructuring the unified calibration pipeline to move PUF cloning after geography assignment, so each geographic clone gets its own state-informed PUF imputations via QRF.

New pipeline:

1. Raw CPS (~55K HH)
2. Clone 10× (v1) / 100× (v2)
3. Assign geography (census block → state, county, CD)
4. PUF clone (2×) — duplicate each HH, QRF impute PUF vars on one half only, with state as predictor
5. QRF impute remaining variables (with state as predictor)
6. PE simulate (taxes + benefits)
7. Build unified sparse matrix (37K+ DB-driven targets)
8. L0 calibration → two weight sets (national + local)

Key benefits:

• Each geographic clone gets geographically-appropriate PUF tax imputations instead of identical national-average ones duplicated everywhere
• State is now a QRF predictor — California clones get California-like tax distributions
• v1 (10×): 1.1M records for fast iteration
• v2 (100×): 11M records for production

Updated pipeline comparison diagram attached below.

[MaxGhenis]

Updated pipeline diagram

I'm implementing the full pipeline reorder in this PR. The key changes:

1. Clone before PUF match — each geographic clone gets its own state-informed PUF imputation via QRF
2. QRF imputation moves after geography — state/county/CD become predictors
3. Single unified L0 solve with all DB-driven targets
4. v1 (10x) for fast iteration, v2 (100x) for production

[MaxGhenis]

Lambda sweep results (directional — needs PUF re-run)

Ran a lambda_l0 sweep on v1 (10 clones × 55K HH = 558K records, 100 epochs each). Caveat: these runs used --skip-puf so record count is half what production will be (558K vs 1.1M for v1, 5.6M vs 11.2M for v2). Results are directional for lambda selection but absolute numbers will shift.

Sweep results (v1, 100 epochs, no PUF)

lambda_l0	Sparsity	Mean error	Active weights
1e-5	32.1%	60.8%	378K / 558K
1e-4	23.9%	61.1%	424K / 558K
1e-3	28.4%	61.0%	399K / 558K
1e-2	53.1%	60.3%	261K / 558K
1e-1	89.5%	60.9%	58K / 558K

Key finding: sparsity is free

Mean error is ~60-61% across all lambdas. Pruning 90% of weights doesn't hurt fit quality at 100 epochs. This makes sense — with 558K weights fitting 40K targets, most weights are redundant.

The original run used lambda_l0=1e-8 with INIT_KEEP_PROB=0.999, which produced 0% sparsity. The library default is lambda_l0=0.01, init_keep_prob=0.5.

Extrapolation to diagram targets

Per the pipeline diagram, we want two weight sets:

• Local weights (~77% sparse): lambda_l0 ≈ 1e-2 should reach this at 500 epochs (53% at 100, accelerating)
• National weights (~99% sparse): need lambda_l0 ≈ 0.5–1.0 (1e-1 gives 89.5% at 100 epochs)

What didn't work

1. v2 with PUF (11.2M records) — first attempt died (OOM at ~15GB RSS). Second attempt ran without PUF (5.6M records) and reached epoch 350/500 at 72% mean error before being killed.
2. PUF is required — the no-PUF runs are not usable for production.

Next steps for @baogorek

1. Re-run lambda sweep WITH PUF on v1 (1.1M records, should be ~3-4GB per run, feasible locally)
2. v2 with PUF (11.2M records) needs either:
   • Target filtering via --stratum-groups to reduce matrix size (you mentioned the full 40K target set causes the fit to fall apart)
   • Running on Modal with more RAM
   • Or both
3. Lambda tuning: Based on this sweep, recommend starting with lambda_l0=0.01 for local and lambda_l0=0.5 for national, with init_keep_prob=0.5 (library default instead of 0.999)
4. Epoch count: 100 epochs was enough to see sparsity trends. For production quality, 500-1000 epochs per your earlier benchmarks.

Full sweep log at lambda_sweep_v1.log in repo root. Diagnostics for each lambda were saved (last one overwrites previous — could add lambda to filename in future runs).