Address SSI disability review findings

Author: MaxGhenis

policyengine_us_data/datasets/cps/extended_cps.py

@@ -27,6 +27,12 @@
     ORG_IMPUTED_VARIABLES,
     apply_org_domain_constraints,
 )
+from policyengine_us_data.datasets.sipp import (
+    SSI_DISABILITY_MODEL_PREDICTORS,
+    SSI_DISABILITY_MODEL_VARIABLE,
+    get_ssi_disability_model,
+    predict_ssi_disability_criteria,
+)
 from policyengine_us_data.pipeline_metadata import pipeline_node
 from policyengine_us_data.pipeline_schema import PipelineNode
 from policyengine_us_data.datasets.puf import PUF, PUF_2024
@@ -185,7 +191,7 @@ def _supports_structural_mortgage_inputs() -> bool:
     "financial_assistance",
     "survivor_benefits",
     "disability_benefits",
-    "meets_ssi_disability_criteria",
+    SSI_DISABILITY_MODEL_VARIABLE,
     "strike_benefits",
     "receives_wic",
     # SPM variables
@@ -236,6 +242,20 @@ def _supports_structural_mortgage_inputs() -> bool:
     "age",
 }
 
+_CLONE_REFRESH_STRUCTURAL_ROLE_VARIABLES = {
+    "is_household_head",
+    "is_tax_unit_head",
+    "is_tax_unit_spouse",
+    "is_tax_unit_dependent",
+    "is_tax_unit_head_or_spouse",
+    "is_family_head",
+    "is_family_spouse",
+    "is_family_dependent",
+    "is_spm_unit_head",
+    "is_spm_unit_spouse",
+    "is_spm_unit_dependent",
+}
+
 # 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 = [
@@ -308,6 +328,7 @@ def _is_structural_clone_variable(variable: str) -> bool:
         or variable in _CLONE_REFRESH_GEOGRAPHY_VARIABLES
         or variable in CLONE_ORIGIN_FLAGS.values()
         or variable in _CLONE_REFRESH_ANCHOR_VARIABLES
+        or variable in _CLONE_REFRESH_STRUCTURAL_ROLE_VARIABLES
         or variable in _STAGE2_COMPUTED_PREDICTORS
     )
 
@@ -389,6 +410,41 @@ def _build_clone_test_frame(
     return X_test[predictors]
 
 
+def _build_ssi_disability_clone_receiver(
+    predictions: pd.DataFrame,
+    X_test: pd.DataFrame,
+    data: dict,
+    time_period: int,
+) -> pd.DataFrame:
+    """Build SIPP SSI disability model inputs for PUF clone records."""
+    n = len(X_test)
+    receiver = pd.DataFrame(index=X_test.index)
+    for predictor in SSI_DISABILITY_MODEL_PREDICTORS:
+        values = None
+        if (
+            predictor == "has_disability_income"
+            and "disability_benefits" in predictions
+        ):
+            values = predictions["disability_benefits"].to_numpy() > 0
+        elif predictor in predictions:
+            values = predictions[predictor].to_numpy()
+        elif predictor in X_test:
+            values = X_test[predictor].to_numpy()
+        else:
+            clone_values = _clone_half_person_values(data, predictor, time_period)
+            if clone_values is not None and len(clone_values) == n:
+                values = clone_values
+
+        if values is None and predictor == "is_female" and "is_male" in X_test:
+            values = ~X_test["is_male"].astype(bool).to_numpy()
+        if values is None:
+            values = np.zeros(n)
+
+        receiver[predictor] = values
+
+    return receiver
+
+
 def _prepare_knn_matrix(
     df: pd.DataFrame,
     reference: pd.DataFrame | None = None,
@@ -923,6 +979,18 @@ def _apply_post_processing(predictions, X_test, time_period, data):
                 "employer_sponsored_insurance_premiums",
             ] = 0
 
+    if SSI_DISABILITY_MODEL_VARIABLE in predictions.columns:
+        receiver = _build_ssi_disability_clone_receiver(
+            predictions,
+            X_test,
+            data,
+            time_period,
+        )
+        predictions[SSI_DISABILITY_MODEL_VARIABLE] = predict_ssi_disability_criteria(
+            get_ssi_disability_model(time_period=time_period),
+            receiver,
+        )
+
     return predictions
 
 

policyengine_us_data/datasets/sipp/sipp.py

@@ -47,6 +47,7 @@
 ]
 
 SSI_DISABILITY_MODEL_VARIABLE = "meets_ssi_disability_criteria"
+SSI_DISABILITY_MODEL_VERSION = 3
 
 SSI_DISABILITY_DIFFICULTY_PREDICTORS = [
     "difficulty_dressing_or_bathing",
@@ -818,7 +819,7 @@ def train_ssi_disability_model(time_period: int = 2024):
 
 def get_ssi_disability_model(time_period: int = 2024) -> QRF:
     """Get or train the SSI disability criteria imputation model."""
-    model_path = STORAGE_FOLDER / f"ssi_disability_criteria_v2_{time_period}.pkl"
+    model_path = _ssi_disability_model_path(time_period)
 
     if not model_path.exists():
         model = train_ssi_disability_model(time_period=time_period)
@@ -832,6 +833,13 @@ def get_ssi_disability_model(time_period: int = 2024) -> QRF:
     return model
 
 
+def _ssi_disability_model_path(time_period: int):
+    return (
+        STORAGE_FOLDER
+        / f"ssi_disability_criteria_v{SSI_DISABILITY_MODEL_VERSION}_{time_period}.pkl"
+    )
+
+
 def build_vehicle_training_frame() -> pd.DataFrame:
     """Build a household-level SIPP frame for vehicle asset imputation."""
     hf_hub_download(

tests/unit/datasets/test_sipp_ssi_disability.py

@@ -12,7 +12,11 @@
     preserve_under_65_ssi_disability_criteria,
     prepare_ssi_disability_receiver,
 )
-from policyengine_us_data.datasets.sipp.sipp import SSI_DISABILITY_COLUMNS
+from policyengine_us_data.datasets.sipp.sipp import (
+    SSI_DISABILITY_COLUMNS,
+    SSI_DISABILITY_MODEL_VERSION,
+    _ssi_disability_model_path,
+)
 
 
 def _base_sipp_frame() -> pd.DataFrame:
@@ -95,6 +99,13 @@ def test_ssi_disability_predictors_use_six_comparable_difficulty_items():
     assert "is_disabled" not in SSI_DISABILITY_MODEL_PREDICTORS
 
 
+def test_ssi_disability_model_cache_version_tracks_predictor_schema():
+    assert SSI_DISABILITY_MODEL_VERSION == 3
+    assert (
+        _ssi_disability_model_path(2024).name == "ssi_disability_criteria_v3_2024.pkl"
+    )
+
+
 def test_build_ssi_disability_training_frame_annualizes_ssdi_amount():
     frame = _base_sipp_frame().iloc[[2]].copy()
     frame["ESSRSN2YN"] = 1

tests/unit/test_extended_cps.py

@@ -29,6 +29,7 @@
     _load_raw_spm_capped_housing_subsidy,
     _apply_post_processing,
     _build_clone_test_frame,
+    _build_ssi_disability_clone_receiver,
     _cps_clone_feature_variables_for_data,
     _derive_overtime_occupation_inputs,
     _impute_clone_cps_features,
@@ -216,6 +217,8 @@ def test_clone_feature_candidates_include_person_level_cps_only_flags(self):
             "household_weight": {2024: np.array([1.0, 1.0, 0.0, 0.0])},
             "state_fips": {2024: np.array([6, 36, 6, 36])},
             "employment_income": {2024: np.array([1.0, 2.0, 3.0, 4.0])},
+            "is_household_head": {2024: np.array([True, True, True, True])},
+            "is_tax_unit_head": {2024: np.array([True, False, True, False])},
             "is_disabled": {2024: np.array([True, False, True, False])},
             "difficulty_hearing": {2024: np.array([False, True, False, True])},
             "meets_ssi_disability_criteria": {
@@ -232,6 +235,8 @@ def test_clone_feature_candidates_include_person_level_cps_only_flags(self):
         assert "household_weight" not in result
         assert "state_fips" not in result
         assert "employment_income" not in result
+        assert "is_household_head" not in result
+        assert "is_tax_unit_head" not in result
         assert "meets_ssi_disability_criteria" not in result
 
     def test_spm_threshold_is_formula_output_not_qrf_imputed(self):
@@ -881,6 +886,83 @@ def test_leaves_data_unchanged_when_pe_us_lacks_llc_inputs(self, monkeypatch):
 
 
 class TestStage2PostProcessing:
+    def test_ssi_disability_clone_receiver_uses_stage2_disability_benefits(self):
+        data = {
+            "person_id": {2024: np.array([1, 2, 101, 102])},
+            "difficulty_hearing": {2024: np.array([False, False, True, False])},
+        }
+        predictions = pd.DataFrame({"disability_benefits": [0.0, 500.0]})
+        x_test = pd.DataFrame(
+            {
+                "age": [40, 40],
+                "is_male": [False, True],
+                "employment_income": [0.0, 0.0],
+            }
+        )
+
+        result = _build_ssi_disability_clone_receiver(
+            predictions,
+            x_test,
+            data,
+            2024,
+        )
+
+        np.testing.assert_array_equal(result["difficulty_hearing"], [True, False])
+        np.testing.assert_array_equal(result["has_disability_income"], [False, True])
+        np.testing.assert_array_equal(result["is_female"], [True, False])
+
+    def test_post_processing_replaces_generic_ssi_disability_predictions(
+        self,
+        monkeypatch,
+    ):
+        class AlwaysTrueModel:
+            def predict(self, X_test):
+                return pd.DataFrame(
+                    {
+                        "meets_ssi_disability_criteria": np.ones(
+                            len(X_test),
+                            dtype=bool,
+                        )
+                    }
+                )
+
+        monkeypatch.setattr(
+            extended_cps_module,
+            "get_ssi_disability_model",
+            lambda time_period: AlwaysTrueModel(),
+        )
+        data = {
+            "person_id": {2024: np.arange(6)},
+            "difficulty_walking_or_climbing_stairs": {
+                2024: np.array([False, False, False, True, False, False])
+            },
+        }
+        predictions = pd.DataFrame(
+            {
+                "meets_ssi_disability_criteria": [False, False, True],
+                "disability_benefits": [0.0, 500.0, 0.0],
+            }
+        )
+        x_test = pd.DataFrame(
+            {
+                "age": [40, 40, 40],
+                "is_male": [False, True, False],
+                "employment_income": [0.0, 0.0, 0.0],
+            }
+        )
+
+        result = _apply_post_processing(
+            predictions=predictions,
+            X_test=x_test,
+            time_period=2024,
+            data=data,
+        )
+
+        np.testing.assert_array_equal(
+            result["meets_ssi_disability_criteria"],
+            np.array([True, True, False]),
+        )
+
     def test_splice_replaces_clone_half_ssi_disability_criteria(self, monkeypatch):
         import policyengine_us