Use household-level fallback SPM units

src/microplex_us/pipelines/us.py

@@ -7062,30 +7062,27 @@ def _assign_family_and_spm_units(self, persons: pd.DataFrame) -> pd.DataFrame:
         next_spm_unit_id = 0
 
         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})
             if primary_mask.any():
                 primary_family_id = next_family_id
-                primary_spm_id = next_spm_unit_id
                 next_family_id += 1
-                next_spm_unit_id += 1
             else:
                 primary_family_id = None
-                primary_spm_id = None
 
             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,
                 }:
                     family_ids[int(row.name)] = primary_family_id
-                    spm_unit_ids[int(row.name)] = primary_spm_id
                     continue
 
                 family_ids[int(row.name)] = next_family_id
-                spm_unit_ids[int(row.name)] = next_spm_unit_id
                 next_family_id += 1
-                next_spm_unit_id += 1
 
         result["family_id"] = result.index.map(family_ids).astype(np.int64)
         result["spm_unit_id"] = result.index.map(spm_unit_ids).astype(np.int64)

tests/pipelines/test_us.py

@@ -788,6 +788,34 @@ def test_build_policyengine_entity_tables(self, persons, households):
             {"SINGLE", "JOINT", "SEPARATE", "HEAD_OF_HOUSEHOLD", "SURVIVING_SPOUSE"}
         )
 
+    def test_build_policyengine_entity_tables_uses_household_level_spm_fallback(
+        self,
+    ):
+        pipeline = USMicroplexPipeline(USMicroplexBuildConfig())
+        population = pd.DataFrame(
+            {
+                "person_id": [1, 2, 3, 4],
+                "household_id": [10, 10, 10, 10],
+                "weight": [1.0, 1.0, 1.0, 1.0],
+                "age": [45, 43, 12, 30],
+                "income": [60_000.0, 15_000.0, 0.0, 20_000.0],
+                "relationship_to_head": [0, 1, 2, 3],
+                "marital_status": [1, 1, 7, 7],
+                "state_fips": [6, 6, 6, 6],
+                "tenure": [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 person_rows["spm_unit_id"].nunique() == 1
+        assert len(tables.families) == 2
+        assert person_rows["family_id"].nunique() == 2
+        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_derives_tax_input_columns(self):
         pipeline = USMicroplexPipeline(USMicroplexBuildConfig())
         population = pd.DataFrame(