Add calibration package checkpointing, target config, and hyperparameter CLI

[baogorek]

Fixes #533
Fixes #534
Fixes #558
Fixes #559

Summary

• Calibration package checkpointing: --build-only saves the expensive matrix build as a pickle, --package-path loads it for fast re-fitting with different hyperparameters or target sets
• Target config YAML: Declarative exclusion rules (target_config.yaml) replace hardcoded target filtering; checked-in config reproduces the junkyard's 22 excluded groups
• Hyperparameter CLI flags: --beta, --lambda-l2, --learning-rate are now tunable from the command line and Modal runner
• Modal runner improvements: Streaming subprocess output, support for new flags
• Documentation: docs/calibration.md covers all workflows (single-pass, build-then-fit, package re-filtering, Modal, portable fitting)
• At-large district naming fix: H5 filenames for at-large districts now use XX-01 (conventional 1-based) instead of XX-00
• GCS staging fix: GCS uploads moved from staging phase to promotion phase, so both GCS and HuggingFace are updated together during promote

Note: This branch includes commits from #537 (PUF impute) since the calibration pipeline depends on that work. The calibration-specific changes are in the top commit.

Test plan

• pytest policyengine_us_data/tests/test_calibration/test_unified_calibration.py — CLI arg parsing tests
• pytest policyengine_us_data/tests/test_calibration/test_target_config.py — target config filtering + package round-trip tests
• Manual: make calibrate-build produces package, --package-path loads it and fits

🤖 Generated with Claude Code

[juaristi22]

Minor comments, but generally LGTM, I was also able to run the calibration job in modal (after removing the ellipsis in unified_calibration.py)!

Small note: if im not mistaken this pr addressess issue #534. Seems like #310 was referenced in it as something that would be addressed together, but this pr does not save the calibration_log.csv among its outputs. Do we want to add it at this point?

[juaristi22]

A couple questions on recent changes

1. Post-cloning PUF imputation remove, is this permanent?

Commit 49a1f66 ("Remove redundant --puf-dataset flag, add national targets") removed the ability to run PUF cloning inside the calibration pipeline, with the rationale that PUF cloning already happens upstream in extended_cps.py.

However, PR #516 specifically designed the pipeline so that PUF + QRF imputation runs after cloning and geography assignment, so that each clone gets geographically-informed imputations (with state_fips as a QRF predictor). As Max described it:

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.

With the current flow (extended_cps.py runs PUF once on base records → calibration pipeline clones 10x), all clones of the same household share identical PUF-imputed values, losing that variability benefit.

Are we planning to bring back the post-cloning PUF imputation once the calibration pipeline is stabilized? Or has the approach changed?

2. All target variable precomputation moved to county level, is this worth the large increase in computation?

In commit 02f8ad0, _build_county_values was introduced to handle county-dependent variables (specifically aca_ptc, since marketplace premiums vary by county). It ran alongside the existing _build_state_values which handled everything else via 51 state-level simulations — a two-tier design gated by COUNTY_DEPENDENT_VARS = {"aca_ptc"}.

Then in commit 40fb389 (a "checkpoint"), COUNTY_DEPENDENT_VARS was removed and all target variable precomputation was moved to _build_county_values. _build_state_values was demoted to only computing constraint variables.

This means the matrix builder now runs ~1,000-2,000 county-level simulations (one per unique county in the geography assignment) instead of 51 state-level simulations for variables like snap, household_count, etc. that don't depend on county. This is a ~40x increase in compute cost with no accuracy benefit for those variables.

Was this an intentional simplification, or a debugging shortcut that could be reverted to the two-tier approach? Restoring COUNTY_DEPENDENT_VARS and routing only aca_ptc (and any future county-dependent vars) through county precomputation would significantly reduce matrix build time.

[baogorek]

@juaristi22 thank you for your thoughtful and excellent comments

1. By the time this PR is done, I do want the calibration_log.csv to be saved once again! Actually, I do get it in this workflow that I run after make data:

 # 1. Build matrix locally (no GPU needed)
  source ~/envs/sep/bin/activate
  python -m policyengine_us_data.calibration.unified_calibration \
      --build-only \
      --skip-source-impute \
      --target-config policyengine_us_data/calibration/target_config.yaml \
      --package-output /tmp/calibration_package.pkl

  # 2. Push package to Modal volume
  modal volume put calibration-data /tmp/calibration_package.pkl calibration_package.pkl --force

  # 3. Fit on GPU from package
  modal run modal_app/remote_calibration_runner.py \
      --branch calibration-pipeline-improvements \
      --gpu T4 \
      --package-volume \
      --epochs 1000 \
      --beta 0.65 \
      --lambda-l0 1e-7 \
      --lambda-l2 1e-8 \
      --log-freq 500 \
      --target-config policyengine_us_data/calibration/target_config.yaml

That will fit the model on modal and drop a calibration_log.csv right on your local drive. I know one of the Issues was about actually storing it in an archive, and maybe that should be out of scope given this PR's complexity.

2. You're absolutely right about the counties and I've now made the per-county computation optional. I know that aca_ptc's formula does involve the county, so we may have to live with not getting perfect matches on from X * w to sim.calculate("aca_ptc").sum(). But I saw the perfect ratios with your code so I know we can do it.

3. Yes, you're right that I lost a bit of the vision with not imputing new PUF values for every clone. Given the brutal difficulty with getting X * w to match sim.calculate().sum() and the speed (which is still slow even after taking out the counties). From Claude:

   • Would it make X*w consistency harder? Yes, significantly. Right now the matrix builder can precompute values once per state and reuse them across all 436 clones — because every clone of household #5000 has the same underlying values regardless of which state it's assigned to. The only things that change per clone are geography inputs and takeup draws.
   • Issue created: Restore post-cloning PUF QRF re-imputation for geographic tax variation #560