Improve SSI disability imputation predictors

Author: MaxGhenis

policyengine_us_data/calibration/source_impute.py

@@ -46,6 +46,7 @@
     SIPP_TIP_AMOUNT_COLUMNS,
     SIPP_TIP_AMOUNT_TO_ALLOCATION_COLUMN,
     SIPP_VEHICLE_TARGET_ALLOCATION_COLUMNS,
+    SSI_DISABILITY_DIFFICULTY_PREDICTORS,
     SSI_DISABILITY_MODEL_VARIABLE,
     VEHICLE_MODEL_PREDICTORS,
     build_vehicle_training_frame,
@@ -902,7 +903,7 @@ def _impute_sipp(
                 "rental_income",
                 "age",
                 "is_male",
-                "is_disabled",
+                *SSI_DISABILITY_DIFFICULTY_PREDICTORS,
                 "social_security_disability",
                 "disability_benefits",
             ],
@@ -930,7 +931,7 @@ def _impute_sipp(
             "interest_income",
             "dividend_income",
             "rental_income",
-            "is_disabled",
+            *SSI_DISABILITY_DIFFICULTY_PREDICTORS,
             "social_security_disability",
         ]:
             if var not in cps_ssi_df.columns:

policyengine_us_data/datasets/cps/cps.py

@@ -130,6 +130,15 @@
     ),
 }
 
+CPS_SSI_DISABILITY_DIFFICULTY_COLUMNS = {
+    "difficulty_dressing_or_bathing": "PEDISDRS",
+    "difficulty_hearing": "PEDISEAR",
+    "difficulty_seeing": "PEDISEYE",
+    "difficulty_doing_errands": "PEDISOUT",
+    "difficulty_walking_or_climbing_stairs": "PEDISPHY",
+    "difficulty_remembering_or_making_decisions": "PEDISREM",
+}
+
 # Census CPS ASEC 2024 technical documentation, PERRP:
 # https://www2.census.gov/programs-surveys/cps/techdocs/cpsmar24.pdf
 PERRP_UNMARRIED_PARTNER_OF_HOUSEHOLD_HEAD_CODES = {
@@ -1076,8 +1085,11 @@ def add_personal_variables(cps: h5py.File, person: DataFrame) -> None:
     # "Is...blind or does...have serious difficulty seeing even when Wearing
     #  glasses?" 1 -> Yes
     cps["is_blind"] = person.PEDISEYE == 1
-    DISABILITY_FLAGS = ["PEDIS" + i for i in ["DRS", "EAR", "EYE", "OUT", "PHY", "REM"]]
-    cps["is_disabled"] = (person[DISABILITY_FLAGS] == 1).any(axis=1)
+    for variable, cps_column in CPS_SSI_DISABILITY_DIFFICULTY_COLUMNS.items():
+        cps[variable] = person[cps_column] == 1
+    cps["is_disabled"] = np.column_stack(
+        [cps[variable] for variable in CPS_SSI_DISABILITY_DIFFICULTY_COLUMNS]
+    ).any(axis=1)
 
     def children_per_parent(col: str) -> pd.DataFrame:
         """Calculate number of children in the household using parental
@@ -2719,6 +2731,7 @@ def add_tips(self, cps: h5py.File):
     cps["bond_assets"] = asset_predictions.bond_assets.values
 
     from policyengine_us_data.datasets.sipp import (
+        SSI_DISABILITY_DIFFICULTY_PREDICTORS,
         SSI_DISABILITY_MODEL_VARIABLE,
         get_ssi_disability_model,
         predict_ssi_disability_criteria,
@@ -2727,7 +2740,7 @@ def add_tips(self, cps: h5py.File):
 
     n_persons = len(cps)
     for variable in [
-        "is_disabled",
+        *SSI_DISABILITY_DIFFICULTY_PREDICTORS,
         "social_security_disability",
     ]:
         cps[variable] = np.asarray(

policyengine_us_data/datasets/cps/extended_cps.py

@@ -9,6 +9,11 @@
 from policyengine_us_data.calibration.formulaic_inputs import (
     FORMULAIC_SPM_INPUTS_TO_DROP,
 )
+from policyengine_us_data.calibration.puf_impute import (
+    CLONE_ORIGIN_FLAGS,
+    IMPUTED_VARIABLES,
+    OVERRIDDEN_IMPUTED_VARIABLES,
+)
 from policyengine_us_data.datasets.cps.cps import (
     CPS,
     CPS_2024,
@@ -91,6 +96,8 @@ def _supports_structural_mortgage_inputs() -> bool:
 if has_policyengine_us_variables("treasury_tipped_occupation_code"):
     CPS_CLONE_FEATURE_VARIABLES.append("treasury_tipped_occupation_code")
 
+PUF_IMPUTED_VARIABLES = set(IMPUTED_VARIABLES) | set(OVERRIDDEN_IMPUTED_VARIABLES)
+
 # Predictors used to rematch CPS features onto the PUF clone half.
 # These are all available on the CPS half and on the doubled extended CPS.
 CPS_CLONE_FEATURE_PREDICTORS = [
@@ -208,6 +215,27 @@ def _supports_structural_mortgage_inputs() -> bool:
 # Set for O(1) lookup in the splice loop.
 _CPS_ONLY_SET = set(CPS_ONLY_IMPUTED_VARIABLES)
 
+_CLONE_REFRESH_GEOGRAPHY_VARIABLES = {
+    "block_geoid",
+    "cbsa_code",
+    "congressional_district_geoid",
+    "county",
+    "county_fips",
+    "place_fips",
+    "puma",
+    "sldl",
+    "sldu",
+    "state_fips",
+    "tract_geoid",
+    "vtd",
+    "zcta",
+    "zip_code",
+}
+
+_CLONE_REFRESH_ANCHOR_VARIABLES = {
+    "age",
+}
+
 # Predictors used for the second-stage CPS-only imputation: demographics
 # plus key income variables that were already imputed from PUF data.
 CPS_STAGE2_DEMOGRAPHIC_PREDICTORS = [
@@ -259,6 +287,93 @@ def _clone_half_person_values(data: dict, variable: str, time_period: int):
     return None
 
 
+def _first_half_person_values(data: dict, variable: str, time_period: int):
+    """Return original-CPS-half values for person-level variables."""
+    if variable not in data:
+        return None
+
+    values = data[variable][time_period]
+    n_persons = len(data["person_id"][time_period])
+    if len(values) != n_persons:
+        return None
+
+    return np.asarray(values[: n_persons // 2])
+
+
+def _is_structural_clone_variable(variable: str) -> bool:
+    """Return whether a variable should remain copied, not rematched."""
+    return (
+        variable.endswith("_id")
+        or variable.endswith("_weight")
+        or variable in _CLONE_REFRESH_GEOGRAPHY_VARIABLES
+        or variable in CLONE_ORIGIN_FLAGS.values()
+        or variable in _CLONE_REFRESH_ANCHOR_VARIABLES
+        or variable in _STAGE2_COMPUTED_PREDICTORS
+    )
+
+
+def _cps_clone_feature_variables_for_data(
+    data: dict,
+    time_period: int,
+) -> list[str]:
+    """Return person-level CPS-only fields to donor-rematch onto PUF clones.
+
+    The PUF clone starts as a literal copy of each CPS donor, then selected
+    tax/income fields are replaced with PUF-imputed values. Any remaining
+    person-level CPS-only field should be refreshed from CPS donors unless it
+    is structural, a PUF-imputed field, or a QRF-handled CPS-only output.
+    """
+    result = []
+    seen = set()
+    explicit_clone_features = set(CPS_CLONE_FEATURE_VARIABLES)
+    for variable in [*CPS_CLONE_FEATURE_VARIABLES, *data.keys()]:
+        if variable in seen:
+            continue
+        seen.add(variable)
+        if variable in PUF_IMPUTED_VARIABLES or variable in _CPS_ONLY_SET:
+            continue
+        is_explicit_clone_feature = variable in explicit_clone_features
+        if not is_explicit_clone_feature and _is_structural_clone_variable(variable):
+            continue
+        if (
+            not is_explicit_clone_feature
+            and _first_half_person_values(data, variable, time_period) is None
+        ):
+            continue
+        result.append(variable)
+    return result
+
+
+def _build_cps_train_frame(
+    cps_sim,
+    data: dict,
+    time_period: int,
+    variables: list[str],
+) -> pd.DataFrame:
+    """Build original-CPS-half training values from PE or stored data."""
+    tbs = getattr(cps_sim, "tax_benefit_system", None)
+    if tbs is None:
+        calculable_variables = variables
+    else:
+        calculable_variables = [
+            variable for variable in variables if variable in tbs.variables
+        ]
+    if calculable_variables:
+        train = cps_sim.calculate_dataframe(calculable_variables).copy()
+    else:
+        n_half = len(data["person_id"][time_period]) // 2
+        train = pd.DataFrame(index=np.arange(n_half))
+
+    for variable in variables:
+        if variable in train.columns:
+            continue
+        values = _first_half_person_values(data, variable, time_period)
+        if values is not None:
+            train[variable] = values
+
+    return train
+
+
 def _build_clone_test_frame(
     cps_sim,
     data: dict,
@@ -321,13 +436,15 @@ def _impute_clone_cps_features(
     from sklearn.neighbors import NearestNeighbors
 
     cps_sim = Microsimulation(dataset=dataset_path)
-    X_train = cps_sim.calculate_dataframe(
-        CPS_CLONE_FEATURE_PREDICTORS + CPS_CLONE_FEATURE_VARIABLES
+    feature_variables = _cps_clone_feature_variables_for_data(data, time_period)
+    X_train = _build_cps_train_frame(
+        cps_sim,
+        data,
+        time_period,
+        CPS_CLONE_FEATURE_PREDICTORS + feature_variables,
     )
     available_outputs = [
-        variable
-        for variable in CPS_CLONE_FEATURE_VARIABLES
-        if variable in X_train.columns
+        variable for variable in feature_variables if variable in X_train.columns
     ]
     if not available_outputs:
         n_half = len(data["person_id"][time_period]) // 2

policyengine_us_data/datasets/sipp/__init__.py

@@ -5,6 +5,7 @@
     get_tip_model,
     train_asset_model,
     get_asset_model,
+    SSI_DISABILITY_DIFFICULTY_PREDICTORS,
     SSI_DISABILITY_MODEL_PREDICTORS,
     SSI_DISABILITY_MODEL_VARIABLE,
     apply_ssi_disability_signal_screen,
@@ -27,6 +28,7 @@
     "get_tip_model",
     "train_asset_model",
     "get_asset_model",
+    "SSI_DISABILITY_DIFFICULTY_PREDICTORS",
     "SSI_DISABILITY_MODEL_PREDICTORS",
     "SSI_DISABILITY_MODEL_VARIABLE",
     "apply_ssi_disability_signal_screen",

policyengine_us_data/datasets/sipp/sipp.py

@@ -48,6 +48,24 @@
 
 SSI_DISABILITY_MODEL_VARIABLE = "meets_ssi_disability_criteria"
 
+SSI_DISABILITY_DIFFICULTY_PREDICTORS = [
+    "difficulty_dressing_or_bathing",
+    "difficulty_hearing",
+    "difficulty_seeing",
+    "difficulty_doing_errands",
+    "difficulty_walking_or_climbing_stairs",
+    "difficulty_remembering_or_making_decisions",
+]
+
+SIPP_SSI_DISABILITY_DIFFICULTY_COLUMNS = {
+    "difficulty_dressing_or_bathing": "ESELFCARE",
+    "difficulty_hearing": "EHEARING",
+    "difficulty_seeing": "ESEEING",
+    "difficulty_doing_errands": "EERRANDS",
+    "difficulty_walking_or_climbing_stairs": "EAMBULAT",
+    "difficulty_remembering_or_making_decisions": "ECOGNIT",
+}
+
 SSI_DISABILITY_MODEL_PREDICTORS = [
     "age",
     "is_female",
@@ -60,7 +78,7 @@
     "stock_assets",
     "bond_assets",
     "count_under_18",
-    "is_disabled",
+    *SSI_DISABILITY_DIFFICULTY_PREDICTORS,
     "social_security_disability",
     "has_disability_income",
 ]
@@ -356,6 +374,7 @@ def get_tip_model() -> QRF:
             "ENJ_NOWRK3",
             "ESSRSN2YN",
             "ESSI_BRSN",
+            *SIPP_SSI_DISABILITY_DIFFICULTY_COLUMNS.values(),
             *SSI_DISABILITY_INCOME_AMOUNT_COLUMNS,
             *SSI_DISABILITY_LABEL_ALLOCATION_COLUMNS,
         ]
@@ -432,6 +451,11 @@ def _yes(df: pd.DataFrame, column: str) -> pd.Series:
     return values.fillna(0).astype(float).eq(1)
 
 
+def _add_ssi_disability_difficulty_predictors(df: pd.DataFrame) -> None:
+    for predictor, source_column in SIPP_SSI_DISABILITY_DIFFICULTY_COLUMNS.items():
+        df[predictor] = _yes(df, source_column)
+
+
 def _ssi_financial_candidate_mask(
     df: pd.DataFrame, time_period: int = 2024
 ) -> pd.Series:
@@ -503,14 +527,15 @@ def build_ssi_disability_training_frame(
         if column in df:
             disability_income_amount += df[column].fillna(0)
 
-    df["is_disabled"] = (
-        _yes(df, "RDIS_ALT")
-        | _yes(df, "RDIS")
-        | _yes(df, "EDISABL")
-        | _yes(df, "EHLTHCOND")
-        | _yes(df, "ENJ_NOWRK3")
+    _add_ssi_disability_difficulty_predictors(df)
+    social_security_amount = (
+        df["TSSSAMT"] if "TSSSAMT" in df else pd.Series(0.0, index=df.index)
+    )
+    df["social_security_disability"] = np.where(
+        _yes(df, "ESSRSN2YN"),
+        social_security_amount.fillna(0).astype(float) * 12,
+        0.0,
     )
-    df["social_security_disability"] = _yes(df, "ESSRSN2YN")
     df["has_disability_income"] = _yes(df, "EDISANY") | disability_income_amount.gt(0)
 
     received_ssi = _yes(df, "RSSI_YRYN")
@@ -570,13 +595,13 @@ def _coerce_ssi_disability_signal(values) -> np.ndarray:
 
 def apply_ssi_disability_signal_screen(
     meets_ssi_disability_criteria: np.ndarray,
-    is_disabled: np.ndarray,
+    disability_difficulty_signal: np.ndarray,
     social_security_disability: np.ndarray,
     has_disability_income: np.ndarray,
 ) -> np.ndarray:
     """Require at least one observed disability signal before accepting imputation."""
     disability_signal = (
-        _coerce_ssi_disability_signal(is_disabled)
+        _coerce_ssi_disability_signal(disability_difficulty_signal)
         | _coerce_ssi_disability_signal(social_security_disability)
         | _coerce_ssi_disability_signal(has_disability_income)
     )
@@ -617,6 +642,16 @@ def coerce_ssi_disability_predictions(values) -> np.ndarray:
     return normalized.isin(["true", "1", "yes"]).to_numpy(dtype=bool)
 
 
+def _ssi_disability_difficulty_signal(receiver: pd.DataFrame) -> np.ndarray:
+    difficulty_signals = [
+        _coerce_ssi_disability_signal(receiver[predictor])
+        for predictor in SSI_DISABILITY_DIFFICULTY_PREDICTORS
+    ]
+    if not difficulty_signals:
+        return np.zeros(len(receiver), dtype=bool)
+    return np.column_stack(difficulty_signals).any(axis=1)
+
+
 def predict_ssi_disability_criteria(model, receiver_df: pd.DataFrame) -> np.ndarray:
     """Predict SSI disability criteria before applying dynamic policy screens."""
     receiver = prepare_ssi_disability_receiver(receiver_df)
@@ -626,7 +661,7 @@ def predict_ssi_disability_criteria(model, receiver_df: pd.DataFrame) -> np.ndar
     )
     return apply_ssi_disability_signal_screen(
         meets_ssi_disability_criteria,
-        receiver["is_disabled"],
+        _ssi_disability_difficulty_signal(receiver),
         receiver["social_security_disability"],
         receiver["has_disability_income"],
     )