Use family relationship for PE family units

src/microplex_us/pipelines/us.py

@@ -7064,7 +7064,7 @@ def _assign_family_and_spm_units(self, persons: pd.DataFrame) -> pd.DataFrame:
         for _, household_persons in result.groupby("household_id", sort=False):
             household_spm_id = next_spm_unit_id
             next_spm_unit_id += 1
-            primary_mask = household_persons["relationship_to_head"].isin({0, 1, 2})
+            primary_mask = self._primary_family_member_mask(household_persons)
             if primary_mask.any():
                 primary_family_id = next_family_id
                 next_family_id += 1
@@ -7073,11 +7073,7 @@ def _assign_family_and_spm_units(self, persons: pd.DataFrame) -> pd.DataFrame:
 
             for _, row in household_persons.iterrows():
                 spm_unit_ids[int(row.name)] = household_spm_id
-                if primary_family_id is not None and row["relationship_to_head"] in {
-                    0,
-                    1,
-                    2,
-                }:
+                if primary_family_id is not None and bool(primary_mask.loc[row.name]):
                     family_ids[int(row.name)] = primary_family_id
                     continue
 
@@ -7088,6 +7084,27 @@ def _assign_family_and_spm_units(self, persons: pd.DataFrame) -> pd.DataFrame:
         result["spm_unit_id"] = result.index.map(spm_unit_ids).astype(np.int64)
         return result
 
+    def _primary_family_member_mask(
+        self,
+        household_persons: pd.DataFrame,
+    ) -> pd.Series:
+        """Identify people who belong to the household's primary family."""
+
+        relationship_primary = household_persons["relationship_to_head"].isin(
+            {0, 1, 2}
+        )
+        if "family_relationship" not in household_persons.columns:
+            return relationship_primary
+
+        family_relationship = pd.to_numeric(
+            household_persons["family_relationship"],
+            errors="coerce",
+        )
+        # CPS A_FAMREL is a family-membership code: 0 means not in a family;
+        # positive values are reference person, spouse, child, or other relative.
+        family_member = family_relationship.isin({1, 2, 3, 4})
+        return relationship_primary | family_member
+
     def _assign_marital_units(
         self,
         persons: pd.DataFrame,

tests/pipelines/test_us.py

@@ -816,6 +816,33 @@ def test_build_policyengine_entity_tables_uses_household_level_spm_fallback(
         assert person_rows.loc[:2, "family_id"].nunique() == 1
         assert person_rows.loc[3, "family_id"] != person_rows.loc[0, "family_id"]
 
+    def test_build_policyengine_entity_tables_uses_family_relationship_for_family_units(
+        self,
+    ):
+        pipeline = USMicroplexPipeline(USMicroplexBuildConfig())
+        population = pd.DataFrame(
+            {
+                "person_id": [1, 2, 3, 4, 5],
+                "household_id": [10, 10, 10, 10, 10],
+                "weight": [1.0, 1.0, 1.0, 1.0, 1.0],
+                "age": [45, 43, 12, 70, 30],
+                "income": [60_000.0, 15_000.0, 0.0, 5_000.0, 20_000.0],
+                "relationship_to_head": [0, 1, 2, 3, 3],
+                "family_relationship": [1, 2, 3, 4, 0],
+                "marital_status": [1, 1, 7, 4, 7],
+                "state_fips": [6, 6, 6, 6, 6],
+                "tenure": [1, 1, 1, 1, 1],
+            }
+        )
+
+        tables = pipeline.build_policyengine_entity_tables(population)
+        person_rows = tables.persons.sort_values("person_id").reset_index(drop=True)
+
+        assert len(tables.spm_units) == 1
+        assert len(tables.families) == 2
+        assert person_rows.loc[:3, "family_id"].nunique() == 1
+        assert person_rows.loc[4, "family_id"] != person_rows.loc[0, "family_id"]
+
     def test_build_policyengine_entity_tables_derives_tax_input_columns(self):
         pipeline = USMicroplexPipeline(USMicroplexBuildConfig())
         population = pd.DataFrame(