Add calibration package checkpointing, target config, and hyperparameter CLI

[baogorek]

Fixes #533
Fixes #534
Fixes #558
Fixes #559
Fixes #562
Fixes #599

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

[juaristi22]

PR #538 Review: Calibration Pipeline Improvements

Solid architectural improvements — checkpointing, declarative target config, and CLI hyperparameter tuning are well-designed. Approving with notes below for follow-up.

---

Critical (recommend fixing before or shortly after merge)

1. cd_geoid undefined in single-cd mode
stacked_dataset_builder.py:160 — calibrated_cds = sorted(set(cd_geoid)) will crash with NameError. This code path appears untested.

2. docs/build_h5.md documents wrong function signature
The documented build_h5 signature shows blocks, cds_to_calibrate, rerandomize_takeup parameters that don't exist in the actual function, which takes a geography: GeographyAssignment object instead. Anyone following this doc will hit immediate errors.

3. n_clones calculated from unique CDs, not actual clone count
unified_calibration.py:1292 — n_clones=len(sorted(set(geography_info["cd_geoid"]...))) uses the number of unique congressional districts (~435) rather than the actual clone count (430). These are close but semantically wrong and could silently affect H5 output.

4. validate-staging workflow never gates on failures
local_area_publish.yaml — The validation job posts a summary but has no step that exits non-zero on sanity failures. A run with 100% failures would show green CI.

---

High

5. Target Config docs describe only exclude: but the active config uses include:
docs/calibration.md Target Config section exclusively documents exclude: format with examples, but target_config.yaml uses include: mode. Confusing for contributors.

6. modal_app/README.md line 184 claims make calibrate-modal passes --prebuilt-matrices — it does not; only --push-results is passed.

7. builtins.print monkey-patch never restored when --log-freq is used
unified_calibration.py:615-746 — The finally block restoring builtins.print is only on the else branch.

8. Pickle deserialization without integrity checks
unified_calibration.py:457, remote_calibration_runner.py:373,940 — Consider checksum verification before loading packages from shared Modal volumes.

---

Medium

• _collect_outputs (remote_calibration_runner.py:90) crashes with unhelpful TypeError if OUTPUT_PATH: marker is missing from subprocess output.
• Makefile calibrate-both / stage-all-h5s use & + wait — if the first backgrounded process fails but the second succeeds last, wait returns 0.
• build-h5s hardcodes --n-clones 430 vs the 436 default elsewhere.
• promote target $(eval VERSION) silently becomes empty if the Python command errors.
• Workflows use Python 3.13 but project targets 3.11.

---

Low / Code Quality

• ~350 lines of GPU function duplication in remote_calibration_runner.py (5x fit_weights_* + 5x fit_from_package_* differing only in gpu param). A factory function would eliminate this.
• include rules in apply_target_config have no dedicated test.
• No integration test for --package-path or --build-only + --package-output flows.
• Tracked binary w_district_calibration.npy should be git rm --cached since it's now covered by the broader .gitignore.

---

Overall: great work on the pipeline consolidation. The items above are mostly documentation drift and edge-case hardening — nothing blocking.