Fix h5 files by saving calibration geography artifact, and model fit resume function (#708)

* Save geography artifacts and add calibration resume/checkpoint support

Calibration now persists geography_assignment.npz alongside weights so
that downstream publish and worker steps use the exact same geography
instead of regenerating it randomly. Adds --resume-from and
--checkpoint-output flags to unified_calibration for continuing fits
from a saved checkpoint or warm-starting from weights. Also gitignores
*.csv.gz to prevent accidental commits of cached ORG data.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Run ruff format

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Add changelog fragment for PR 708

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Stub l0 module in test so patch works without l0-python installed

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Add self-employment and SSN card type count targets to calibration config

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Distinguish ITIN holders from SSN holders in CPS data

Fix calibration crash on string constraint variables (ssn_card_type) by
falling back from float32 cast when values are non-numeric.

Impute ITIN status for undocumented (code-0) persons: select tax units
with code-0 earners via weighted random sampling targeting 4.4M ITIN
returns (IRS NTA), then mark all code-0 members of those units. Updates
has_tin = (ssn_card_type != 0) | has_itin_number so ITIN holders
correctly qualify for ODC ($500 credit).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Fix PR 708 checkpoint and ID regressions

* Handle string ID fields in PUF cloning

* Fold taxpayer ID imputation into calibration resume PR

* Fix PUF subsample logging format

* fixes

* Expand PR 708 changelog summary

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Max Ghenis <mghenis@gmail.com>

Author: 3 people

.gitignore

@@ -4,6 +4,7 @@
 **/*.h5
 **/*.npy
 **/*.csv
+**/*.csv.gz
 **/_build
 **/*.pkl
 **/*.db

changelog.d/708.added

@@ -0,0 +1 @@
+Save calibration geography as a pipeline artifact, add ``--resume-from`` and checkpoint support for long-running calibration fits, and fix resume/artifact handling in the remote calibration pipeline. This also adds conservative CPS taxpayer-ID outputs (``has_tin``, ``has_valid_ssn``, and a temporary ``has_itin`` compatibility alias), plus string-valued constraint handling needed for ID-target calibration.

docs/calibration.md

@@ -23,6 +23,12 @@ python -m policyengine_us_data.calibration.unified_calibration \
   --package-path storage/calibration/calibration_package.pkl \
   --epochs 500 --device cuda
 
+# Resume a previous fit for 500 more epochs:
+python -m policyengine_us_data.calibration.unified_calibration \
+  --package-path storage/calibration/calibration_package.pkl \
+  --resume-from storage/calibration/calibration_weights.npy \
+  --epochs 500 --device cuda
+
 # Full pipeline with PUF (build + fit in one shot):
 make calibrate
 ```
@@ -88,6 +94,30 @@ python -m policyengine_us_data.calibration.unified_calibration \
 You can re-run Step 2 as many times as you want with different hyperparameters. The expensive matrix
 build only happens once.
 
+Every fit now also writes a checkpoint next to the weights output
+(`calibration_weights.checkpoint.pt` by default). To continue the same fit,
+pass `--resume-from` with the weights file or checkpoint path. If a sibling
+checkpoint exists next to the weights file, it is used automatically so the
+L0 gate state is restored as well.
+
+```bash
+python -m policyengine_us_data.calibration.unified_calibration \
+  --package-path storage/calibration/calibration_package.pkl \
+  --epochs 2000 \
+  --beta 0.65 \
+  --lambda-l0 1e-4 \
+  --lambda-l2 1e-12 \
+  --log-freq 500 \
+  --target-config policyengine_us_data/calibration/target_config.yaml \
+  --device cpu \
+  --output policyengine_us_data/storage/calibration/national/weights.npy \
+  --resume-from policyengine_us_data/storage/calibration/national/weights.npy
+```
+
+When `--resume-from` points to a checkpoint, `--epochs` means additional epochs
+to run beyond the saved checkpoint epoch count. If only a `.npy` weights file
+exists, the run warm-starts from those weights.
+
 ### 2. Full pipeline with PUF
 
 Adding `--puf-dataset` doubles the record count (~24K base records x 430 clones = ~10.3M columns) by

modal_app/local_area.py

@@ -334,6 +334,8 @@ def build_areas_worker(
         "--output-dir",
         str(output_dir),
     ]
+    if "geography" in calibration_inputs:
+        worker_cmd.extend(["--geography-path", calibration_inputs["geography"]])
     if "n_clones" in calibration_inputs:
         worker_cmd.extend(["--n-clones", str(calibration_inputs["n_clones"])])
     if "seed" in calibration_inputs:
@@ -659,6 +661,7 @@ def coordinate_publish(
         Path(f"/pipeline/artifacts/{run_id}") if run_id else Path("/pipeline/artifacts")
     )
     weights_path = artifacts / "calibration_weights.npy"
+    geography_path = artifacts / "geography_assignment.npz"
     db_path = artifacts / "policy_data.db"
     dataset_path = artifacts / "source_imputed_stratified_extended_cps.h5"
     config_json_path = artifacts / "unified_run_config.json"
@@ -678,6 +681,7 @@ def coordinate_publish(
 
     calibration_inputs = {
         "weights": str(weights_path),
+        "geography": str(geography_path),
         "dataset": str(dataset_path),
         "database": str(db_path),
         "n_clones": n_clones,
@@ -943,6 +947,7 @@ def coordinate_national_publish(
         Path(f"/pipeline/artifacts/{run_id}") if run_id else Path("/pipeline/artifacts")
     )
     weights_path = artifacts / "national_calibration_weights.npy"
+    geography_path = artifacts / "national_geography_assignment.npz"
     db_path = artifacts / "policy_data.db"
     dataset_path = artifacts / "source_imputed_stratified_extended_cps.h5"
     config_json_path = artifacts / "national_unified_run_config.json"
@@ -962,6 +967,7 @@ def coordinate_national_publish(
 
     calibration_inputs = {
         "weights": str(weights_path),
+        "geography": str(geography_path),
         "dataset": str(dataset_path),
         "database": str(db_path),
         "n_clones": n_clones,
@@ -972,6 +978,7 @@ def coordinate_national_publish(
         artifacts,
         filename_remap={
             "calibration_weights.npy": "national_calibration_weights.npy",
+            "geography_assignment.npz": "national_geography_assignment.npz",
         },
     )
     run_dir = staging_dir / run_id

modal_app/pipeline.py

@@ -832,6 +832,11 @@ def run_pipeline(
                     BytesIO(regional_result["weights"]),
                     f"{artifacts_rel}/calibration_weights.npy",
                 )
+                if regional_result.get("geography"):
+                    batch.put_file(
+                        BytesIO(regional_result["geography"]),
+                        f"{artifacts_rel}/geography_assignment.npz",
+                    )
                 if regional_result.get("config"):
                     batch.put_file(
                         BytesIO(regional_result["config"]),
@@ -856,6 +861,11 @@ def run_pipeline(
                         BytesIO(national_result["weights"]),
                         f"{artifacts_rel}/national_calibration_weights.npy",
                     )
+                    if national_result.get("geography"):
+                        batch.put_file(
+                            BytesIO(national_result["geography"]),
+                            f"{artifacts_rel}/national_geography_assignment.npz",
+                        )
                     if national_result.get("config"):
                         batch.put_file(
                             BytesIO(national_result["config"]),