Use tipped occupation status in SIPP tip imputation

Author: MaxGhenis

policyengine_us_data/calibration/source_impute.py

@@ -26,6 +26,10 @@
 
 import numpy as np
 import pandas as pd
+from policyengine_us_data.datasets.cps.tipped_occupation import (
+    derive_any_treasury_tipped_occupation_code,
+    derive_is_tipped_occupation,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -68,6 +72,7 @@
     "age",
     "count_under_18",
     "count_under_6",
+    "is_tipped_occupation",
 ]
 
 SIPP_ASSETS_PREDICTORS = [
@@ -97,6 +102,8 @@
     "NONE": 0,
 }
 
+SIPP_JOB_OCCUPATION_COLUMNS = [f"TJB{i}_OCC" for i in range(1, 8)]
+
 
 def _encode_tenure_type(df: pd.DataFrame) -> pd.DataFrame:
     """Convert tenure_type enum strings to numeric codes."""
@@ -362,6 +369,14 @@ def _impute_sipp(
     sipp_df["age"] = sipp_df.TAGE
     sipp_df["household_weight"] = sipp_df.WPFINWGT
     sipp_df["household_id"] = sipp_df.SSUID
+    sipp_df["treasury_tipped_occupation_code"] = (
+        derive_any_treasury_tipped_occupation_code(
+            sipp_df[SIPP_JOB_OCCUPATION_COLUMNS]
+        )
+    )
+    sipp_df["is_tipped_occupation"] = derive_is_tipped_occupation(
+        sipp_df["treasury_tipped_occupation_code"]
+    )
 
     sipp_df["is_under_18"] = sipp_df.TAGE < 18
     sipp_df["is_under_6"] = sipp_df.TAGE < 6
@@ -379,6 +394,7 @@ def _impute_sipp(
         "count_under_18",
         "count_under_6",
         "age",
+        "is_tipped_occupation",
         "household_weight",
     ]
     tip_train = sipp_df[tip_cols].dropna()
@@ -409,6 +425,12 @@ def _impute_sipp(
     else:
         cps_tip_df["count_under_18"] = 0.0
         cps_tip_df["count_under_6"] = 0.0
+    if "treasury_tipped_occupation_code" in data:
+        cps_tip_df["is_tipped_occupation"] = derive_is_tipped_occupation(
+            data["treasury_tipped_occupation_code"][time_period]
+        ).astype(np.float32)
+    else:
+        cps_tip_df["is_tipped_occupation"] = 0.0
 
     qrf = QRF()
     logger.info(

policyengine_us_data/datasets/cps/cps.py

@@ -21,6 +21,7 @@
 from policyengine_us_data.utils.randomness import seeded_rng
 from policyengine_us_data.datasets.cps.tipped_occupation import (
     derive_treasury_tipped_occupation_code,
+    derive_is_tipped_occupation,
 )
 
 
@@ -1800,6 +1801,9 @@ def add_tips(self, cps: h5py.File):
         .values
     )
     cps = pd.DataFrame(cps)
+    cps["is_tipped_occupation"] = derive_is_tipped_occupation(
+        cps["treasury_tipped_occupation_code"]
+    )
 
     # Impute tips
 

policyengine_us_data/datasets/cps/tipped_occupation.py

@@ -79,3 +79,32 @@ def derive_treasury_tipped_occupation_code(
     return (
         values.map(CENSUS_OCCUPATION_CODE_TO_TTOC).fillna(0).astype(np.int16).to_numpy()
     )
+
+
+def derive_any_treasury_tipped_occupation_code(
+    occupation_columns: pd.DataFrame,
+) -> np.ndarray:
+    """Collapse multiple job occupation columns to one person-level tipped code."""
+
+    if occupation_columns.shape[1] == 0:
+        return np.zeros(len(occupation_columns), dtype=np.int16)
+
+    mapped_columns = [
+        derive_treasury_tipped_occupation_code(occupation_columns[column])
+        for column in occupation_columns.columns
+    ]
+    return np.column_stack(mapped_columns).max(axis=1).astype(np.int16)
+
+
+def derive_is_tipped_occupation(
+    treasury_tipped_occupation_codes: pd.Series | np.ndarray,
+) -> np.ndarray:
+    """Return a boolean indicator for whether any Treasury tipped code is present."""
+
+    return (
+        pd.Series(treasury_tipped_occupation_codes, copy=False)
+        .fillna(0)
+        .astype(np.int16)
+        .gt(0)
+        .to_numpy()
+    )

policyengine_us_data/datasets/sipp/sipp.py

@@ -4,6 +4,20 @@
 from policyengine_us_data.storage import STORAGE_FOLDER
 import pickle
 from huggingface_hub import hf_hub_download
+from policyengine_us_data.datasets.cps.tipped_occupation import (
+    derive_any_treasury_tipped_occupation_code,
+    derive_is_tipped_occupation,
+)
+
+
+SIPP_JOB_OCCUPATION_COLUMNS = [f"TJB{i}_OCC" for i in range(1, 8)]
+TIP_MODEL_PREDICTORS = [
+    "employment_income",
+    "age",
+    "count_under_18",
+    "count_under_6",
+    "is_tipped_occupation",
+]
 
 
 def train_tip_model():
@@ -79,6 +93,12 @@ def train_tip_model():
     df["household_weight"] = df.WPFINWGT
     df["household_id"] = df.SSUID
     df["age"] = df.TAGE
+    df["treasury_tipped_occupation_code"] = derive_any_treasury_tipped_occupation_code(
+        df[SIPP_JOB_OCCUPATION_COLUMNS]
+    )
+    df["is_tipped_occupation"] = derive_is_tipped_occupation(
+        df["treasury_tipped_occupation_code"]
+    )
 
     sipp = df[
         [
@@ -88,6 +108,7 @@ def train_tip_model():
             "count_under_18",
             "count_under_6",
             "age",
+            "is_tipped_occupation",
             "household_weight",
         ]
     ]
@@ -107,20 +128,15 @@ def train_tip_model():
 
     model = model.fit(
         X_train=sipp,
-        predictors=[
-            "employment_income",
-            "age",
-            "count_under_18",
-            "count_under_6",
-        ],
+        predictors=TIP_MODEL_PREDICTORS,
         imputed_variables=["tip_income"],
     )
 
     return model
 
 
 def get_tip_model() -> QRF:
-    model_path = STORAGE_FOLDER / "tips.pkl"
+    model_path = STORAGE_FOLDER / "tips_tipped_occ_v2.pkl"
 
     if not model_path.exists():
         model = train_tip_model()

policyengine_us_data/tests/test_calibration/test_source_impute.py

@@ -4,6 +4,7 @@
 """
 
 import numpy as np
+import pandas as pd
 
 from policyengine_us_data.calibration.source_impute import (
     ACS_IMPUTED_VARIABLES,
@@ -20,6 +21,10 @@
     _person_state_fips,
     impute_source_variables,
 )
+from policyengine_us_data.datasets.cps.tipped_occupation import (
+    derive_any_treasury_tipped_occupation_code,
+    derive_is_tipped_occupation,
+)
 
 
 def _make_data_dict(n_persons=20, time_period=2024):
@@ -41,6 +46,9 @@ def _make_data_dict(n_persons=20, time_period=2024):
         "employment_income": {
             time_period: rng.uniform(0, 100000, n_persons).astype(np.float32),
         },
+        "treasury_tipped_occupation_code": {
+            time_period: np.zeros(n_persons, dtype=np.int16),
+        },
         "rent": {time_period: np.zeros(n_persons)},
         "real_estate_taxes": {time_period: np.zeros(n_persons)},
         "tip_income": {time_period: np.zeros(n_persons)},
@@ -85,6 +93,9 @@ def test_acs_uses_state(self):
     def test_sipp_tips_has_income(self):
         assert "employment_income" in SIPP_TIPS_PREDICTORS
 
+    def test_sipp_tips_uses_tipped_occupation_status(self):
+        assert "is_tipped_occupation" in SIPP_TIPS_PREDICTORS
+
     def test_sipp_assets_has_income(self):
         assert "employment_income" in SIPP_ASSETS_PREDICTORS
 
@@ -205,3 +216,19 @@ def test_impute_sipp_exists(self):
 
     def test_impute_scf_exists(self):
         assert callable(_impute_scf)
+
+
+class TestTippedOccupationHelpers:
+    def test_derive_any_treasury_tipped_occupation_code(self):
+        occupations = pd.DataFrame(
+            {
+                "TJB1_OCC": [4040, 1021, np.nan],
+                "TJB2_OCC": [np.nan, 4110, 9620],
+            }
+        )
+        derived = derive_any_treasury_tipped_occupation_code(occupations)
+        np.testing.assert_array_equal(derived, np.array([101, 102, 809]))
+
+    def test_derive_is_tipped_occupation(self):
+        derived = derive_is_tipped_occupation(np.array([0, 101, 809]))
+        np.testing.assert_array_equal(derived, np.array([False, True, True]))