Preserve SSI disability anchors during imputation

Author: MaxGhenis

policyengine_us_data/calibration/source_impute.py

@@ -41,6 +41,7 @@
     build_vehicle_training_frame,
     get_ssi_disability_model,
     predict_ssi_disability_criteria,
+    preserve_under_65_ssi_disability_criteria,
 )
 
 from policyengine_us_data.datasets.org import (
@@ -767,13 +768,16 @@ def _impute_sipp(
             ssi_disability_model,
             cps_ssi_df,
         )
-        if "ssi_reported" in data:
-            reported_under_65 = (data["ssi_reported"][time_period] > 0) & (
-                data["age"][time_period] < 65
-            )
-            meets_ssi_disability_criteria = (
-                meets_ssi_disability_criteria | reported_under_65
-            )
+        existing_meets_ssi_disability_criteria = data.get(
+            SSI_DISABILITY_MODEL_VARIABLE, {}
+        ).get(time_period)
+        ssi_reported = data.get("ssi_reported", {}).get(time_period)
+        meets_ssi_disability_criteria = preserve_under_65_ssi_disability_criteria(
+            meets_ssi_disability_criteria,
+            age=data["age"][time_period],
+            ssi_reported=ssi_reported,
+            existing_meets_ssi_disability_criteria=existing_meets_ssi_disability_criteria,
+        )
         data[SSI_DISABILITY_MODEL_VARIABLE] = {
             time_period: meets_ssi_disability_criteria
         }

policyengine_us_data/datasets/cps/cps.py

@@ -2652,6 +2652,7 @@ def add_tips(self, cps: h5py.File):
         SSI_DISABILITY_MODEL_VARIABLE,
         get_ssi_disability_model,
         predict_ssi_disability_criteria,
+        preserve_under_65_ssi_disability_criteria,
     )
 
     n_persons = len(cps)
@@ -2671,13 +2672,14 @@ def add_tips(self, cps: h5py.File):
         ssi_disability_model,
         cps,
     )
-    if "ssi_reported" in existing_data:
-        reported_under_65 = (np.asarray(existing_data["ssi_reported"]) > 0) & (
-            np.asarray(existing_data["age"]) < 65
-        )
-        meets_ssi_disability_criteria = (
-            meets_ssi_disability_criteria | reported_under_65
-        )
+    meets_ssi_disability_criteria = preserve_under_65_ssi_disability_criteria(
+        meets_ssi_disability_criteria,
+        age=existing_data.get("age", np.full(n_persons, 65)),
+        ssi_reported=existing_data.get("ssi_reported"),
+        existing_meets_ssi_disability_criteria=existing_data.get(
+            SSI_DISABILITY_MODEL_VARIABLE
+        ),
+    )
     cps[SSI_DISABILITY_MODEL_VARIABLE] = meets_ssi_disability_criteria
 
     from policyengine_us_data.datasets.sipp import get_vehicle_model

policyengine_us_data/datasets/sipp/__init__.py

@@ -9,6 +9,7 @@
     build_ssi_disability_training_frame,
     coerce_ssi_disability_predictions,
     predict_ssi_disability_criteria,
+    preserve_under_65_ssi_disability_criteria,
     prepare_ssi_disability_receiver,
     train_ssi_disability_model,
     get_ssi_disability_model,
@@ -28,6 +29,7 @@
     "build_ssi_disability_training_frame",
     "coerce_ssi_disability_predictions",
     "predict_ssi_disability_criteria",
+    "preserve_under_65_ssi_disability_criteria",
     "prepare_ssi_disability_receiver",
     "train_ssi_disability_model",
     "get_ssi_disability_model",

policyengine_us_data/datasets/sipp/sipp.py

@@ -427,6 +427,29 @@ def apply_ssi_disability_signal_screen(
     return np.asarray(meets_ssi_disability_criteria, dtype=bool) & disability_signal
 
 
+def preserve_under_65_ssi_disability_criteria(
+    meets_ssi_disability_criteria: np.ndarray,
+    age: np.ndarray,
+    ssi_reported: np.ndarray | None = None,
+    existing_meets_ssi_disability_criteria: np.ndarray | None = None,
+) -> np.ndarray:
+    """Preserve observed under-65 SSI disability criteria anchors."""
+    result = np.asarray(meets_ssi_disability_criteria, dtype=bool).copy()
+    under_65 = pd.Series(age).fillna(np.inf).astype(float).lt(65).to_numpy()
+
+    if ssi_reported is not None:
+        reported_ssi = pd.Series(ssi_reported).fillna(0).astype(float).gt(0).to_numpy()
+        result |= reported_ssi & under_65
+
+    if existing_meets_ssi_disability_criteria is not None:
+        result |= (
+            _coerce_ssi_disability_signal(existing_meets_ssi_disability_criteria)
+            & under_65
+        )
+
+    return result
+
+
 def coerce_ssi_disability_predictions(values) -> np.ndarray:
     """Convert classifier labels to booleans without treating 'False' as true."""
     series = pd.Series(values)

tests/integration/test_cps_generation.py

@@ -51,6 +51,7 @@ def calculate(self, variable_name):
                 "receives_wic": [False, False],
                 "hud_income_level": ["VERY_LOW"],
                 "spm_unit_tenure_type": ["RENTER"],
+                "is_eligible_for_housing_assistance": [True],
                 "tax_unit_child_dependents": [0],
                 "age_head": [40],
             }
@@ -223,9 +224,27 @@ def predict(self, X_test, mean_quantile):
                 }
             )
 
+    class FakeSsiDisabilityModel:
+        pass
+
+    def fake_predict_ssi_disability_criteria(model, receiver_df):
+        assert isinstance(model, FakeSsiDisabilityModel)
+        assert receiver_df["employment_income"].tolist() == [25_000.0, 30_000.0]
+        return np.array([True, False])
+
     monkeypatch.setattr(sipp_module, "get_tip_model", lambda: FakeTipModel())
     monkeypatch.setattr(sipp_module, "get_asset_model", lambda: FakeAssetModel())
     monkeypatch.setattr(sipp_module, "get_vehicle_model", lambda: FakeVehicleModel())
+    monkeypatch.setattr(
+        sipp_module,
+        "get_ssi_disability_model",
+        lambda: FakeSsiDisabilityModel(),
+    )
+    monkeypatch.setattr(
+        sipp_module,
+        "predict_ssi_disability_criteria",
+        fake_predict_ssi_disability_criteria,
+    )
 
     dataset = FakeDataset()
     add_tips(
@@ -245,6 +264,10 @@ def predict(self, X_test, mean_quantile):
         18_000.0,
         7_500.0,
     ]
+    assert dataset.saved_dataset["meets_ssi_disability_criteria"].tolist() == [
+        True,
+        False,
+    ]
 
 
 def test_add_rent_requests_person_level_frames(monkeypatch, tmp_path):

tests/unit/calibration/test_source_impute.py

@@ -22,6 +22,7 @@
     _impute_sipp,
     _person_state_fips,
     impute_source_variables,
+    preserve_under_65_ssi_disability_criteria,
 )
 from policyengine_us_data.datasets.cps.tipped_occupation import (
     derive_any_treasury_tipped_occupation_code,
@@ -249,6 +250,17 @@ def test_impute_org_exists(self):
     def test_impute_scf_exists(self):
         assert callable(_impute_scf)
 
+    def test_source_impute_preserves_existing_under_65_ssi_criteria(self):
+        fake_model_predictions = np.array([False, False, False])
+
+        result = preserve_under_65_ssi_disability_criteria(
+            fake_model_predictions,
+            age=np.array([40, 64, 70]),
+            existing_meets_ssi_disability_criteria=np.array([True, False, True]),
+        )
+
+        np.testing.assert_array_equal(result, np.array([True, False, False]))
+
 
 class TestTippedOccupationHelpers:
     def test_derive_any_treasury_tipped_occupation_code(self):

tests/unit/datasets/test_sipp_ssi_disability.py

@@ -8,6 +8,7 @@
     build_ssi_disability_training_frame,
     coerce_ssi_disability_predictions,
     predict_ssi_disability_criteria,
+    preserve_under_65_ssi_disability_criteria,
     prepare_ssi_disability_receiver,
 )
 
@@ -103,6 +104,17 @@ def test_apply_ssi_disability_signal_screen_treats_missing_as_false():
     np.testing.assert_array_equal(result, np.array([False, False, False]))
 
 
+def test_preserve_under_65_ssi_disability_criteria_keeps_observed_anchors():
+    result = preserve_under_65_ssi_disability_criteria(
+        np.array([False, False, False, False]),
+        age=np.array([40, 64, 70, 30]),
+        ssi_reported=np.array([0, 100, 100, np.nan]),
+        existing_meets_ssi_disability_criteria=np.array([True, False, True, np.nan]),
+    )
+
+    np.testing.assert_array_equal(result, np.array([True, True, False, False]))
+
+
 def test_coerce_ssi_disability_predictions_handles_string_false():
     result = coerce_ssi_disability_predictions(
         pd.Series(["False", "True", "0", "1", False, True, 0, 1])