Require SIPP tip allocation flags

Author: MaxGhenis

policyengine_us_data/calibration/source_impute.py

@@ -43,8 +43,8 @@
     SIPP_ASSET_ALLOCATION_COLUMNS,
     SIPP_ASSET_TARGET_ALLOCATION_COLUMNS,
     SIPP_ASSET_TARGET_SOURCE_COLUMNS,
-    SIPP_TIP_ALLOCATION_COLUMNS,
     SIPP_TIP_AMOUNT_COLUMNS,
+    SIPP_TIP_AMOUNT_TO_ALLOCATION_COLUMN,
     SIPP_VEHICLE_TARGET_ALLOCATION_COLUMNS,
     SSI_DISABILITY_MODEL_VARIABLE,
     VEHICLE_MODEL_PREDICTORS,
@@ -652,6 +652,14 @@ def _impute_sipp(
     tip_amount_columns = [
         column for column in SIPP_TIP_AMOUNT_COLUMNS if column in sipp_df
     ]
+    tip_allocation_columns = [
+        SIPP_TIP_AMOUNT_TO_ALLOCATION_COLUMN[column] for column in tip_amount_columns
+    ]
+    require_columns_present(
+        sipp_df.columns,
+        tip_allocation_columns,
+        source_name="SIPP slim tip donor file",
+    )
     sipp_df["tip_income"] = sipp_df[tip_amount_columns].fillna(0).sum(axis=1) * 12
     sipp_df["employment_income"] = sipp_df.TPTOTINC * 12
     sipp_df["age"] = sipp_df.TAGE
@@ -679,7 +687,7 @@ def _impute_sipp(
         sipp_df,
         targets=["tip_income"],
         target_source_columns={"tip_income": tip_amount_columns},
-        target_allocation_flag_columns={"tip_income": SIPP_TIP_ALLOCATION_COLUMNS},
+        target_allocation_flag_columns={"tip_income": tip_allocation_columns},
         require_nonmissing_source=False,
     )
 

policyengine_us_data/datasets/sipp/sipp.py

@@ -11,16 +11,18 @@
 )
 from policyengine_us_data.utils.source_quality import (
     filter_observed_source_rows,
+    require_columns_present,
     sipp_allocation_flag_for,
     target_observed_source_masks,
 )
 
 
 SIPP_JOB_OCCUPATION_COLUMNS = [f"TJB{i}_OCC" for i in range(1, 8)]
 SIPP_TIP_AMOUNT_COLUMNS = [f"TJB{i}_TXAMT" for i in range(1, 8)]
-SIPP_TIP_ALLOCATION_COLUMNS = [
-    sipp_allocation_flag_for(column) for column in SIPP_TIP_AMOUNT_COLUMNS
-]
+SIPP_TIP_AMOUNT_TO_ALLOCATION_COLUMN = {
+    column: sipp_allocation_flag_for(column) for column in SIPP_TIP_AMOUNT_COLUMNS
+}
+SIPP_TIP_ALLOCATION_COLUMNS = list(SIPP_TIP_AMOUNT_TO_ALLOCATION_COLUMN.values())
 TIP_MODEL_PREDICTORS = [
     "employment_income",
     "age",
@@ -124,6 +126,14 @@ def train_tip_model():
     # AJB*_TXAMT Census allocation flags (small ints 0/1/2 indicating
     # imputation status) and added them to the dollar totals.
     tip_amount_columns = [column for column in SIPP_TIP_AMOUNT_COLUMNS if column in df]
+    tip_allocation_columns = [
+        SIPP_TIP_AMOUNT_TO_ALLOCATION_COLUMN[column] for column in tip_amount_columns
+    ]
+    require_columns_present(
+        df.columns,
+        tip_allocation_columns,
+        source_name="SIPP tip donor file",
+    )
     df["tip_income"] = df[tip_amount_columns].fillna(0).sum(axis=1) * 12
     df["employment_income"] = df.TPTOTINC * 12
     df["is_under_18"] = (df.TAGE < 18) & (df.MONTHCODE == 12)
@@ -159,7 +169,7 @@ def train_tip_model():
         df,
         targets=["tip_income"],
         target_source_columns={"tip_income": tip_amount_columns},
-        target_allocation_flag_columns={"tip_income": SIPP_TIP_ALLOCATION_COLUMNS},
+        target_allocation_flag_columns={"tip_income": tip_allocation_columns},
         require_nonmissing_source=False,
     )
 

tests/unit/calibration/test_source_impute.py

@@ -6,6 +6,7 @@
 import numpy as np
 import pandas as pd
 import huggingface_hub
+import pytest
 
 from policyengine_us_data.calibration import source_impute
 from policyengine_us_data.calibration.source_impute import (
@@ -346,7 +347,7 @@ def test_calibration_sipp_tip_counts_use_reference_month(self, monkeypatch):
         }
         for column in source_impute.SIPP_TIP_AMOUNT_COLUMNS:
             columns[column] = [0.0, 10.0, 0.0, 5.0]
-        for column in source_impute.SIPP_TIP_ALLOCATION_COLUMNS:
+        for column in source_impute.SIPP_TIP_AMOUNT_TO_ALLOCATION_COLUMN.values():
             columns[column] = [0, 0, 0, 0]
         for column in source_impute.SIPP_JOB_OCCUPATION_COLUMNS:
             columns[column] = [0, 0, 0, 0]
@@ -390,6 +391,25 @@ def predict(self, X_test):
         np.testing.assert_array_equal(household_one["count_under_18"], [1, 1])
         np.testing.assert_array_equal(household_one["count_under_6"], [0, 0])
 
+    def test_calibration_sipp_tip_requires_allocation_flags(self, monkeypatch):
+        monkeypatch.setattr(
+            huggingface_hub,
+            "hf_hub_download",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            source_impute.pd,
+            "read_csv",
+            lambda *args, **kwargs: pd.DataFrame({"TJB1_TXAMT": [10.0]}),
+        )
+
+        with pytest.raises(KeyError, match="AJB1_TXAMT"):
+            _impute_sipp(
+                data=_make_data_dict(n_persons=4),
+                state_fips=np.array([1, 1], dtype=np.int32),
+                time_period=2024,
+            )
+
     def test_impute_org_exists(self):
         assert callable(_impute_org)
 

tests/unit/datasets/test_sipp_tip_columns.py

@@ -6,8 +6,10 @@
 """
 
 import pandas as pd
+import pytest
 
 from policyengine_us_data.datasets.sipp.sipp import SIPP_TIP_AMOUNT_COLUMNS
+import policyengine_us_data.datasets.sipp.sipp as sipp_module
 
 
 def test_tip_regex_matches_dollar_amounts_only():
@@ -50,3 +52,17 @@ def test_tip_sum_excludes_allocation_flags():
         df[df.columns[df.columns.str.contains("TXAMT")]].fillna(0).sum(axis=1)
     )
     assert list(buggy_tip_income_monthly) == [151.0, 278.0]
+
+
+def test_train_tip_model_requires_allocation_flags_for_present_tip_columns(
+    monkeypatch,
+):
+    monkeypatch.setattr(sipp_module, "hf_hub_download", lambda *args, **kwargs: None)
+    monkeypatch.setattr(
+        sipp_module.pd,
+        "read_csv",
+        lambda *args, **kwargs: pd.DataFrame({"TJB1_TXAMT": [10.0]}),
+    )
+
+    with pytest.raises(KeyError, match="AJB1_TXAMT"):
+        sipp_module.train_tip_model()