Migrate GREG calibration to svy (#1135)

Author: MaxGhenis

.github/workflows/pr.yaml

@@ -132,6 +132,23 @@ jobs:
         env:
           CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
 
+  calibration-extra-tests:
+    name: Calibration extra tests
+    runs-on: ubuntu-latest
+    needs: [check-fork, lint]
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.14"
+      - uses: astral-sh/setup-uv@v8.1.0
+      - run: uv sync --dev --extra calibration
+      - name: Run svy GREG adapter regression
+        run: >
+          uv run --extra calibration pytest
+          tests/unit/test_long_term_calibration_contract.py::test_greg_calibrator_hits_linear_controls_with_svy
+          -v
+
   bundle-release-manifest-contract:
     name: Validate bundle release manifest contract
     runs-on: ubuntu-latest

changelog.d/1129.changed.md

@@ -0,0 +1 @@
+Migrate long-term GREG calibration from archived samplics to svy.

policyengine_us_data/datasets/cps/long_term/assess_publishable_horizon.py

@@ -12,6 +12,7 @@
 from policyengine_us import Microsimulation
 
 from calibration import build_calibration_audit, calibrate_weights
+from calibration import GregCalibrator
 from calibration_profiles import (
     approximate_window_for_year,
     classify_calibration_quality,
@@ -35,12 +36,6 @@
     set_long_term_target_source,
 )
 
-try:
-    from samplics.weighting import SampleWeight
-except ImportError:  # pragma: no cover - only needed for greg profiles
-    SampleWeight = None
-
-
 DEFAULT_BASE_DATASET_PATH = (
     "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5"
 )
@@ -90,12 +85,7 @@ def parse_years(raw: str) -> list[int]:
 def maybe_build_calibrator(method: str):
     if method != "greg":
         return None
-    if SampleWeight is None:
-        raise ImportError(
-            "samplics is required for GREG calibration. "
-            "Install with: pip install policyengine-us-data[calibration]"
-        )
-    return SampleWeight()
+    return GregCalibrator()
 
 
 def benchmark_tob_values(

policyengine_us_data/datasets/cps/long_term/calibration.py

@@ -3,6 +3,66 @@
 from scipy import optimize, sparse
 
 
+class GregCalibrator:
+    """Small adapter around svy's GREG calibration workflow."""
+
+    _base_weight_column = "_policyengine_base_weight"
+    _calibrated_weight_column = "_policyengine_greg_weight"
+
+    def __init__(self):
+        try:
+            import polars as pl
+            import svy
+        except ImportError as e:  # pragma: no cover - exercised without extra
+            raise ImportError(
+                "svy is required for GREG calibration. "
+                "Install with: pip install policyengine-us-data[calibration]"
+            ) from e
+
+        self._pl = pl
+        self._svy = svy
+
+    def calibrate(self, *, samp_weight, aux_vars, control):
+        control = {str(name): float(target) for name, target in control.items()}
+        aux_df = self._auxiliary_dataframe(aux_vars, list(control))
+        aux_df[self._base_weight_column] = np.asarray(samp_weight, dtype=float)
+
+        sample = self._svy.Sample(
+            self._pl.from_pandas(aux_df),
+            design=self._svy.Design(wgt=self._base_weight_column),
+        )
+        sample.weighting.calibrate(
+            controls=control,
+            wgt_name=self._calibrated_weight_column,
+        )
+        return (
+            sample.data.get_column(self._calibrated_weight_column)
+            .to_numpy()
+            .astype(float)
+        )
+
+    def _auxiliary_dataframe(self, aux_vars, control_names):
+        if isinstance(aux_vars, pd.DataFrame):
+            aux_df = aux_vars.copy()
+            aux_df.columns = [str(column) for column in aux_df.columns]
+            return aux_df
+
+        if sparse.issparse(aux_vars):
+            aux_array = aux_vars.toarray()
+        else:
+            aux_array = np.asarray(aux_vars)
+
+        if aux_array.ndim == 1:
+            aux_array = aux_array.reshape(-1, 1)
+
+        if aux_array.shape[1] != len(control_names):
+            raise ValueError(
+                "aux_vars column count must match the number of GREG controls"
+            )
+
+        return pd.DataFrame(aux_array.astype(float), columns=control_names)
+
+
 def _pct_error(achieved, target):
     if target == 0:
         return 0.0 if achieved == 0 else float("inf")
@@ -106,10 +166,10 @@ def calibrate_greg(
     n_ages=86,
 ):
     """
-    Calibrate weights using GREG method via samplics.
+    Calibrate weights using GREG method via svy.
 
     Args:
-        calibrator: SampleWeight instance from samplics
+        calibrator: GregCalibrator instance
         X: Design matrix (n_households x n_ages)
         y_target: Target age distribution
         baseline_weights: Initial household weights

policyengine_us_data/datasets/cps/long_term/run_household_projection.py

@@ -56,6 +56,7 @@
     set_long_term_target_source,
 )
 from calibration import (
+    GregCalibrator,
     build_calibration_audit,
     build_clone_donor_component_weight_concentration_audit,
     build_clone_donor_family_weight_concentration_audit,
@@ -787,14 +788,7 @@ def _compose_reforms(*reforms):
     from ssa_data import load_hi_tob_projections, load_oasdi_tob_projections
 
 if USE_GREG:
-    try:
-        from samplics.weighting import SampleWeight
-    except ImportError:
-        raise ImportError(
-            "samplics is required for GREG calibration. "
-            "Install with: pip install policyengine-us-data[calibration]"
-        )
-    calibrator = SampleWeight()
+    calibrator = GregCalibrator()
 else:
     calibrator = None
 

pyproject.toml

@@ -54,7 +54,7 @@ dependencies = [
 
 [project.optional-dependencies]
 calibration = [
-    "samplics",
+    "svy>=0.18.2",
 ]
 l0 = [
     "l0-python",

tests/unit/test_long_term_calibration_contract.py

@@ -21,6 +21,7 @@
 )
 from policyengine_us_data.datasets.cps.long_term import run_long_term_production
 from policyengine_us_data.datasets.cps.long_term.calibration import (
+    GregCalibrator,
     assess_nonnegative_feasibility,
     build_calibration_audit,
     calibrate_entropy,
@@ -1107,6 +1108,29 @@ def test_strict_greg_failure_raises():
         )
 
 
+def test_greg_calibrator_hits_linear_controls_with_svy():
+    pytest.importorskip("svy")
+
+    X = np.array([[1.0, 0.0], [0.0, 1.0]])
+    y_target = np.array([2.0, 3.0])
+    baseline_weights = np.array([1.0, 1.0])
+
+    weights, iterations, audit = calibrate_weights(
+        X=X,
+        y_target=y_target,
+        baseline_weights=baseline_weights,
+        method="greg",
+        calibrator=GregCalibrator(),
+        n_ages=2,
+        allow_fallback_to_ipf=False,
+    )
+
+    assert iterations == 1
+    assert audit["method_used"] == "greg"
+    assert audit["fell_back_to_ipf"] is False
+    np.testing.assert_allclose(X.T @ weights, y_target)
+
+
 def test_build_calibration_audit_reports_constraint_error():
     X = np.array([[1.0, 0.0], [0.0, 1.0]])
     y_target = np.array([1.0, 1.0])