Fix SSI disability label source filtering

Author: MaxGhenis

policyengine_us_data/datasets/sipp/sipp.py

@@ -12,7 +12,7 @@
 from policyengine_us_data.utils.source_quality import (
     cap_training_sample,
     filter_positive_finite_weight_rows,
-    filter_observed_source_rows,
+    observed_source_mask,
     require_columns_present,
     sipp_allocation_flag_for,
     target_observed_source_masks,
@@ -48,7 +48,7 @@
 
 SSI_DISABILITY_CRITERIA_VARIABLE = "meets_ssi_disability_criteria"
 SSI_DISABILITY_MODEL_VARIABLE = SSI_DISABILITY_CRITERIA_VARIABLE
-SSI_DISABILITY_MODEL_VERSION = 5
+SSI_DISABILITY_MODEL_VERSION = 6
 SSI_DISABILITY_EXPORT_VARIABLES = (SSI_DISABILITY_CRITERIA_VARIABLE,)
 
 # These six CPS/SIPP difficulty items are construction-time predictors for the
@@ -464,6 +464,30 @@ def _add_ssi_disability_difficulty_predictors(df: pd.DataFrame) -> None:
         df[predictor] = _yes(df, source_column)
 
 
+def _observed_ssi_disability_label_mask(
+    df: pd.DataFrame, received_ssi: pd.Series
+) -> pd.Series:
+    ssi_receipt_observed = observed_source_mask(
+        df,
+        source_columns=["RSSI_YRYN"],
+        allocation_flag_columns=[sipp_allocation_flag_for("RSSI_YRYN")],
+    )
+    ssi_receipt_observed &= pd.to_numeric(
+        df.get("RSSI_YRYN", pd.Series(np.nan, index=df.index)),
+        errors="coerce",
+    ).isin([1, 2])
+    ssi_reason_observed = observed_source_mask(
+        df,
+        source_columns=["ESSI_BRSN"],
+        allocation_flag_columns=[sipp_allocation_flag_for("ESSI_BRSN")],
+    )
+    ssi_reason_observed &= pd.to_numeric(
+        df.get("ESSI_BRSN", pd.Series(np.nan, index=df.index)),
+        errors="coerce",
+    ).isin([1, 2])
+    return ssi_receipt_observed & (~received_ssi | ssi_reason_observed)
+
+
 def _ssi_financial_candidate_mask(
     df: pd.DataFrame, time_period: int = 2024
 ) -> pd.Series:
@@ -571,12 +595,7 @@ def build_ssi_disability_training_frame(
     df["ssi_disability_training_candidate"] = (financial_candidate & under_65) | df[
         SSI_DISABILITY_CRITERIA_VARIABLE
     ]
-    df = filter_observed_source_rows(
-        df,
-        target_name=SSI_DISABILITY_CRITERIA_VARIABLE,
-        source_columns=SSI_DISABILITY_LABEL_SOURCE_COLUMNS,
-        allocation_flag_columns=SSI_DISABILITY_LABEL_ALLOCATION_COLUMNS,
-    )
+    df = df.loc[_observed_ssi_disability_label_mask(df, received_ssi)].copy()
 
     columns = SSI_DISABILITY_MODEL_PREDICTORS + [
         SSI_DISABILITY_CRITERIA_VARIABLE,

tests/unit/datasets/test_sipp_ssi_disability.py

@@ -101,9 +101,9 @@ def test_ssi_disability_predictors_use_six_comparable_difficulty_items():
 
 
 def test_ssi_disability_model_cache_version_tracks_predictor_schema():
-    assert SSI_DISABILITY_MODEL_VERSION == 5
+    assert SSI_DISABILITY_MODEL_VERSION == 6
     assert _ssi_disability_model_path(2024).name == (
-        "ssi_disability_criteria_v5_2024.pkl"
+        "ssi_disability_criteria_v6_2024.pkl"
     )
 
 
@@ -132,6 +132,35 @@ def test_build_ssi_disability_training_frame_excludes_allocated_label_source():
     )
 
 
+def test_build_ssi_disability_training_frame_keeps_non_ssi_without_reason_source():
+    frame = _base_sipp_frame()
+    frame["ASSI_YRYN"] = 0
+    frame["ASSI_BRSN"] = 3
+
+    result = build_ssi_disability_training_frame(frame)
+
+    assert len(result) == 2
+    np.testing.assert_array_equal(
+        result[SSI_DISABILITY_MODEL_VARIABLE].values,
+        np.array([False, False]),
+    )
+
+
+def test_build_ssi_disability_training_frame_excludes_ssi_with_missing_reason_source():
+    frame = _base_sipp_frame()
+    frame.loc[0, "ESSI_BRSN"] = -9
+    frame["ASSI_YRYN"] = 0
+    frame["ASSI_BRSN"] = 0
+
+    result = build_ssi_disability_training_frame(frame)
+
+    assert len(result) == 3
+    np.testing.assert_array_equal(
+        result[SSI_DISABILITY_MODEL_VARIABLE].values,
+        np.array([False, False, False]),
+    )
+
+
 def test_prepare_ssi_disability_receiver_fills_missing_predictors():
     result = prepare_ssi_disability_receiver(
         pd.DataFrame(