Format rebased CPS clone imputation changes

Author: MaxGhenis

policyengine_us_data/datasets/cps/extended_cps.py

@@ -22,9 +22,11 @@
 
 logger = logging.getLogger(__name__)
 
+
 def _supports_structural_mortgage_inputs() -> bool:
     return has_policyengine_us_variables(*STRUCTURAL_MORTGAGE_VARIABLES)
 
+
 # CPS-only categorical features to donor-impute onto the PUF clone half.
 # These drive subgroup analysis and occupation-based logic, so naive donor
 # duplication dilutes the relationship between the clone's PUF-imputed
@@ -197,7 +199,9 @@ def _clone_half_person_values(data: dict, variable: str, time_period: int):
             continue
         entity_half = len(entity_ids) // 2
         clone_entity_ids = entity_ids[entity_half:]
-        clone_person_entity_ids = data[person_entity_id_var][time_period][n_persons_half:]
+        clone_person_entity_ids = data[person_entity_id_var][time_period][
+            n_persons_half:
+        ]
         value_map = dict(zip(clone_entity_ids, values[entity_half:]))
         return np.array([value_map[idx] for idx in clone_person_entity_ids])
 
@@ -270,7 +274,9 @@ def _impute_clone_cps_features(
         CPS_CLONE_FEATURE_PREDICTORS + CPS_CLONE_FEATURE_VARIABLES
     )
     available_outputs = [
-        variable for variable in CPS_CLONE_FEATURE_VARIABLES if variable in X_train.columns
+        variable
+        for variable in CPS_CLONE_FEATURE_VARIABLES
+        if variable in X_train.columns
     ]
     if not available_outputs:
         n_half = len(data["person_id"][time_period]) // 2
@@ -323,8 +329,12 @@ def _impute_clone_cps_features(
         predictions.loc[test_mask, available_outputs] = donor_outputs.to_numpy()
 
     if "detailed_occupation_recode" in predictions:
-        occupation_codes = predictions["detailed_occupation_recode"].astype(float).to_numpy()
-        for column, values in _derive_overtime_occupation_inputs(occupation_codes).items():
+        occupation_codes = (
+            predictions["detailed_occupation_recode"].astype(float).to_numpy()
+        )
+        for column, values in _derive_overtime_occupation_inputs(
+            occupation_codes
+        ).items():
             predictions[column] = values
 
     return predictions

policyengine_us_data/tests/test_extended_cps.py

@@ -364,9 +364,7 @@ def calculate_dataframe(self, columns):
             "employment_income": {
                 tp: np.array([20_000, 35_000, 90_000, 150_000], dtype=np.float32)
             },
-            "self_employment_income": {
-                tp: np.zeros(4, dtype=np.float32)
-            },
+            "self_employment_income": {tp: np.zeros(4, dtype=np.float32)},
             "social_security": {tp: np.zeros(4, dtype=np.float32)},
             "is_tax_unit_head": {tp: np.ones(4, dtype=bool)},
             "is_tax_unit_spouse": {tp: np.zeros(4, dtype=bool)},
@@ -387,9 +385,7 @@ def calculate_dataframe(self, columns):
         assert result["tax_unit_is_joint"].tolist() == [0, 1]
 
     def test_derive_overtime_occupation_inputs(self):
-        derived = _derive_overtime_occupation_inputs(
-            np.array([53, 52, 8, 41, 1, 99])
-        )
+        derived = _derive_overtime_occupation_inputs(np.array([53, 52, 8, 41, 1, 99]))
 
         assert derived["has_never_worked"].tolist() == [
             True,
@@ -477,7 +473,9 @@ def calculate_dataframe(self, columns):
             "is_tax_unit_head": {tp: np.array([1, 0, 1, 0], dtype=bool)},
             "is_tax_unit_spouse": {tp: np.zeros(4, dtype=bool)},
             "is_tax_unit_dependent": {tp: np.array([0, 1, 0, 1], dtype=bool)},
-            "employment_income": {tp: np.array([95_000, 0, 97_000, 0], dtype=np.float32)},
+            "employment_income": {
+                tp: np.array([95_000, 0, 97_000, 0], dtype=np.float32)
+            },
             "self_employment_income": {tp: np.zeros(4, dtype=np.float32)},
             "social_security": {tp: np.zeros(4, dtype=np.float32)},
         }