Filter invalid SIPP imputation weights

Author: MaxGhenis

changelog.d/sipp-positive-weights.fixed

@@ -0,0 +1 @@
+Filter non-positive SIPP donor weights before fitting source imputation models.

policyengine_us_data/calibration/source_impute.py

@@ -82,6 +82,7 @@
 from policyengine_us_data.pipeline_schema import PipelineNode
 from policyengine_us_data.utils.source_quality import (
     cap_training_sample,
+    filter_positive_finite_weight_rows,
     require_columns_present,
     target_observed_source_masks,
 )
@@ -710,6 +711,12 @@ def _impute_sipp(
         "household_weight",
     ]
     tip_train = sipp_df[tip_cols].dropna()
+    tip_train, tip_target_filters = filter_positive_finite_weight_rows(
+        tip_train,
+        weight_col="household_weight",
+        target_filters=tip_target_filters,
+        context_name="SIPP source tip donor",
+    )
     tip_train, tip_target_filters = cap_training_sample(
         tip_train,
         max_train_samples=10_000,
@@ -849,6 +856,12 @@ def _impute_sipp(
             target_source_columns=SIPP_ASSET_TARGET_SOURCE_COLUMNS,
             target_allocation_flag_columns=SIPP_ASSET_TARGET_ALLOCATION_COLUMNS,
         )
+        asset_train, asset_target_filters = filter_positive_finite_weight_rows(
+            asset_train,
+            weight_col="household_weight",
+            target_filters=asset_target_filters,
+            context_name="SIPP source asset donor",
+        )
         asset_train, asset_target_filters = cap_training_sample(
             asset_train,
             max_train_samples=20_000,
@@ -1013,6 +1026,12 @@ def _impute_sipp(
             targets=vehicle_vars,
             target_allocation_flag_columns=SIPP_VEHICLE_TARGET_ALLOCATION_COLUMNS,
         )
+        vehicle_train, vehicle_target_filters = filter_positive_finite_weight_rows(
+            vehicle_train,
+            weight_col="household_weight",
+            target_filters=vehicle_target_filters,
+            context_name="SIPP source vehicle donor",
+        )
         vehicle_train, vehicle_target_filters = cap_training_sample(
             vehicle_train,
             max_train_samples=20_000,

policyengine_us_data/datasets/sipp/sipp.py

@@ -11,6 +11,7 @@
 )
 from policyengine_us_data.utils.source_quality import (
     cap_training_sample,
+    filter_positive_finite_weight_rows,
     filter_observed_source_rows,
     require_columns_present,
     sipp_allocation_flag_for,
@@ -188,6 +189,12 @@ def train_tip_model():
     ]
 
     sipp = sipp[~sipp.isna().any(axis=1)]
+    sipp, tip_target_filters = filter_positive_finite_weight_rows(
+        sipp,
+        weight_col="household_weight",
+        target_filters=tip_target_filters,
+        context_name="SIPP tip donor",
+    )
     sipp, tip_target_filters = cap_training_sample(
         sipp,
         max_train_samples=10_000,
@@ -652,6 +659,12 @@ def train_asset_model():
         target_source_columns=SIPP_ASSET_TARGET_SOURCE_COLUMNS,
         target_allocation_flag_columns=SIPP_ASSET_TARGET_ALLOCATION_COLUMNS,
     )
+    sipp, asset_target_filters = filter_positive_finite_weight_rows(
+        sipp,
+        weight_col="household_weight",
+        target_filters=asset_target_filters,
+        context_name="SIPP asset donor",
+    )
     sipp, asset_target_filters = cap_training_sample(
         sipp,
         max_train_samples=20_000,
@@ -839,6 +852,12 @@ def train_vehicle_model():
         targets=vehicle_vars,
         target_allocation_flag_columns=SIPP_VEHICLE_TARGET_ALLOCATION_COLUMNS,
     )
+    sipp, vehicle_target_filters = filter_positive_finite_weight_rows(
+        sipp,
+        weight_col="household_weight",
+        target_filters=vehicle_target_filters,
+        context_name="SIPP vehicle donor",
+    )
     sipp, vehicle_target_filters = cap_training_sample(
         sipp,
         max_train_samples=20_000,

policyengine_us_data/utils/source_quality.py

@@ -225,3 +225,52 @@ def cap_training_sample(
         for target, mask in filters.items()
     }
     return sampled_df, sampled_filters
+
+
+def filter_positive_finite_weight_rows(
+    df: pd.DataFrame,
+    *,
+    weight_col: str,
+    target_filters: Mapping[str, pd.Series] | None = None,
+    context_name: str = "donor training frame",
+) -> tuple[pd.DataFrame, dict[str, pd.Series]]:
+    """Drop rows whose fit weight cannot be passed to microimpute."""
+    if weight_col not in df:
+        raise KeyError(f"{context_name} is missing weight column {weight_col!r}")
+
+    filters = {}
+    for target, mask in (target_filters or {}).items():
+        aligned = mask.reindex(df.index)
+        if aligned.isna().any():
+            raise ValueError(f"target_filters[{target!r}] contains missing values")
+        filters[target] = aligned.astype(bool)
+
+    weights = pd.to_numeric(df[weight_col], errors="coerce")
+    valid_weight = np.isfinite(weights) & weights.gt(0)
+    dropped = int((~valid_weight).sum())
+    if dropped:
+        logger.info(
+            "Dropped %d/%d %s rows with non-positive or non-finite %s",
+            dropped,
+            len(df),
+            context_name,
+            weight_col,
+        )
+
+    filtered_df = df.loc[valid_weight].copy().reset_index(drop=True)
+    filtered_filters = {
+        target: pd.Series(
+            mask.loc[valid_weight].to_numpy(dtype=bool),
+            index=filtered_df.index,
+        )
+        for target, mask in filters.items()
+    }
+
+    for target, mask in filtered_filters.items():
+        if not mask.any():
+            raise ValueError(
+                f"No observed donor rows with positive finite {weight_col} "
+                f"available for {target}"
+            )
+
+    return filtered_df, filtered_filters

tests/unit/datasets/test_sipp_tip_columns.py

@@ -5,10 +5,14 @@
 to explicit `TJB*_TXAMT` dollar-amount columns only.
 """
 
+import numpy as np
 import pandas as pd
 import pytest
 
-from policyengine_us_data.datasets.sipp.sipp import SIPP_TIP_AMOUNT_COLUMNS
+from policyengine_us_data.datasets.sipp.sipp import (
+    SIPP_JOB_OCCUPATION_COLUMNS,
+    SIPP_TIP_AMOUNT_COLUMNS,
+)
 import policyengine_us_data.datasets.sipp.sipp as sipp_module
 
 
@@ -66,3 +70,47 @@ def test_train_tip_model_requires_allocation_flags_for_present_tip_columns(
 
     with pytest.raises(KeyError, match="AJB1_TXAMT"):
         sipp_module.train_tip_model()
+
+
+def test_train_tip_model_drops_non_positive_weights(monkeypatch):
+    monkeypatch.setattr(sipp_module, "hf_hub_download", lambda *args, **kwargs: None)
+
+    data = {
+        "SSUID": [1, 2, 3, 4],
+        "MONTHCODE": [12, 12, 12, 12],
+        "TAGE": [30, 31, 32, 33],
+        "WPFINWGT": [100.0, 0.0, -5.0, 200.0],
+        "TPTOTINC": [1_000.0, 2_000.0, 3_000.0, 4_000.0],
+        "TJB1_TXAMT": [10.0, 20.0, 30.0, 40.0],
+        "AJB1_TXAMT": [0, 0, 0, 0],
+    }
+    for column in SIPP_JOB_OCCUPATION_COLUMNS:
+        data[column] = [0, 0, 0, 0]
+    monkeypatch.setattr(
+        sipp_module.pd,
+        "read_csv",
+        lambda *args, **kwargs: pd.DataFrame(data),
+    )
+
+    captured = {}
+
+    class FakeQRF:
+        def fit(
+            self,
+            *,
+            X_train,
+            predictors,
+            imputed_variables,
+            target_filters,
+            weight_col,
+        ):
+            captured["weights"] = X_train[weight_col].to_numpy()
+            captured["target_filter"] = target_filters["tip_income"].to_numpy()
+            return self
+
+    monkeypatch.setattr(sipp_module, "QRF", FakeQRF)
+
+    sipp_module.train_tip_model()
+
+    np.testing.assert_array_equal(captured["weights"], [100.0, 200.0])
+    np.testing.assert_array_equal(captured["target_filter"], [True, True])

tests/unit/test_source_quality.py

@@ -3,6 +3,7 @@
 
 from policyengine_us_data.utils.source_quality import (
     cap_training_sample,
+    filter_positive_finite_weight_rows,
     observed_source_mask,
     require_columns_present,
     sipp_allocation_flag_for,
@@ -201,3 +202,54 @@ def test_cap_training_sample_rejects_misaligned_filters():
         raise AssertionError("Expected misaligned target filters to fail")
 
     assert "target_filters['value']" in message
+
+
+def test_filter_positive_finite_weight_rows_reindexes_target_filters():
+    df = pd.DataFrame(
+        {
+            "value": [10, 20, 30, 40, 50],
+            "household_weight": [1.0, 0.0, np.nan, np.inf, 5.0],
+        },
+        index=[10, 11, 12, 13, 14],
+    )
+    filters = {
+        "value": pd.Series(
+            [True, True, False, True, True],
+            index=df.index,
+        )
+    }
+
+    filtered, filtered_filters = filter_positive_finite_weight_rows(
+        df,
+        weight_col="household_weight",
+        target_filters=filters,
+        context_name="unit-test donor",
+    )
+
+    assert filtered["value"].tolist() == [10, 50]
+    assert filtered.index.tolist() == [0, 1]
+    np.testing.assert_array_equal(filtered_filters["value"].values, [True, True])
+    assert filtered_filters["value"].index.tolist() == [0, 1]
+
+
+def test_filter_positive_finite_weight_rows_requires_observed_target_rows():
+    df = pd.DataFrame(
+        {
+            "value": [10, 20],
+            "household_weight": [0.0, 1.0],
+        }
+    )
+    filters = {"value": pd.Series([True, False], index=df.index)}
+
+    try:
+        filter_positive_finite_weight_rows(
+            df,
+            weight_col="household_weight",
+            target_filters=filters,
+        )
+    except ValueError as error:
+        message = str(error)
+    else:
+        raise AssertionError("Expected all invalid observed weights to fail")
+
+    assert "No observed donor rows with positive finite household_weight" in message