Use national calibration preset in pipeline (#947)

Author: MaxGhenis

changelog.d/945.fixed

@@ -0,0 +1 @@
+Use the national calibration preset L0 penalty in the Modal pipeline.

modal_app/pipeline.py

@@ -109,6 +109,8 @@
     or "policyengine-us-data-pipeline"
 )
 
+NATIONAL_FIT_LAMBDA_L0 = 1e-4
+
 hf_secret = modal.Secret.from_name("huggingface-token")
 gcp_secret = modal.Secret.from_name("gcp-credentials")
 
@@ -1170,7 +1172,7 @@ def run_pipeline(
             "epochs": national_epochs,
             "target_config": "policyengine_us_data/calibration/target_config.yaml",
             "beta": 0.65,
-            "lambda_l0": 2e-2,
+            "lambda_l0": NATIONAL_FIT_LAMBDA_L0,
             "lambda_l2": 1e-12,
             "log_freq": 100,
             "skip_national": skip_national,
@@ -1257,7 +1259,7 @@ def run_pipeline(
                     volume_package_path=vol_path,
                     target_config=target_cfg,
                     beta=0.65,
-                    lambda_l0=2e-2,
+                    lambda_l0=NATIONAL_FIT_LAMBDA_L0,
                     lambda_l2=1e-12,
                     log_freq=100,
                 )

tests/unit/test_pipeline.py

@@ -10,6 +10,7 @@
 modal = pytest.importorskip("modal")
 
 from modal_app.pipeline import (  # noqa: E402
+    NATIONAL_FIT_LAMBDA_L0,
     _build_diagnostics_upload_script,
     _calibration_package_parameters,
     _run_required_promotion_subprocess,
@@ -37,7 +38,7 @@ def test_calibration_package_parameters_track_matrix_mode():
     )
 
     assert params["chunked_matrix"] is True
-    assert params["workers"] is None
+    assert "workers" not in params
     assert params["chunk_size"] == 10_000
     assert params["parallel_matrix"] is True
     assert params["num_matrix_workers"] == 25
@@ -57,9 +58,13 @@ def test_calibration_package_parameters_ignore_unused_matrix_options():
 
     assert params["chunked_matrix"] is False
     assert params["workers"] == 50
-    assert params["chunk_size"] is None
+    assert "chunk_size" not in params
     assert params["parallel_matrix"] is False
-    assert params["num_matrix_workers"] is None
+    assert "num_matrix_workers" not in params
+
+
+def test_national_fit_lambda_matches_national_preset():
+    assert NATIONAL_FIT_LAMBDA_L0 == pytest.approx(1e-4)
 
 
 class TestRunMetadata: