Fix cluster-as-PSU nesting across strata and defer FPC validation for late PSU injection from PR #218 review (round 16)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: igerber

diff_diff/survey.py

@@ -197,12 +197,10 @@ def resolve(self, data: pd.DataFrame) -> "ResolvedSurveyDesign":
                                 f"({n_psu_h}) in stratum {h}. FPC must be >= n_PSU."
                             )
                     else:
-                        n_h = np.sum(mask_h)
-                        if fpc_h < n_h:
-                            raise ValueError(
-                                f"FPC ({fpc_h}) is less than the number of observations "
-                                f"({n_h}) in stratum {h}. FPC must be >= n_obs."
-                            )
+                        # No PSU declared yet — clusters may be injected later
+                        # as effective PSUs, so skip per-obs FPC validation here.
+                        # FPC will be applied at the PSU level in compute_survey_vcov.
+                        pass
             elif psu_arr is not None:
                 # No strata: require FPC is a single constant value
                 if len(np.unique(fpc_arr)) > 1:
@@ -459,8 +457,15 @@ def _inject_cluster_as_psu(resolved, cluster_ids):
             "when used as effective PSUs for survey variance estimation."
         )
 
-    # Factorize cluster_ids for consistent integer encoding
-    codes, uniques = pd.factorize(cluster_ids)
+    # When strata are present, make cluster IDs unique within strata
+    # (same nesting logic as SurveyDesign.resolve() with nest=True)
+    if resolved.strata is not None:
+        combined = np.array(
+            [f"{s}_{c}" for s, c in zip(resolved.strata, cluster_ids)]
+        )
+        codes, uniques = pd.factorize(combined)
+    else:
+        codes, uniques = pd.factorize(cluster_ids)
     n_clusters = len(uniques)
 
     return replace(resolved, psu=codes, n_psu=n_clusters)

tests/test_survey.py

@@ -2872,3 +2872,47 @@ def test_cluster_as_psu_with_na_rejected(self):
         cluster_ids = np.array([0, 1, np.nan, 0])
         with pytest.raises(ValueError, match="Cluster IDs contain missing"):
             _inject_cluster_as_psu(resolved, cluster_ids)
+
+
+class TestRound16Fixes:
+    """Tests for PR #218 review round 16: cluster-as-PSU nesting and FPC."""
+
+    def test_injected_cluster_nested_in_strata(self):
+        """Injected cluster IDs with repeated labels across strata get unique codes."""
+        from diff_diff.survey import _inject_cluster_as_psu
+
+        # 2 strata, cluster "1" appears in both → should produce 4 unique PSUs
+        strata = np.array([0, 0, 0, 0, 1, 1, 1, 1])
+        resolved = ResolvedSurveyDesign(
+            weights=np.ones(8),
+            weight_type="pweight",
+            strata=strata,
+            psu=None,
+            fpc=None,
+            n_strata=2,
+            n_psu=0,
+            lonely_psu="remove",
+        )
+        cluster_ids = np.array([1, 1, 2, 2, 1, 1, 2, 2])  # labels repeat across strata
+        result = _inject_cluster_as_psu(resolved, cluster_ids)
+        # Should produce 4 unique PSUs (2 per stratum), not 2
+        assert result.n_psu == 4
+        # df_survey = n_psu - n_strata = 4 - 2 = 2
+        assert result.df_survey == 2
+
+    def test_fpc_with_strata_no_psu_accepted(self):
+        """FPC + strata (no PSU) is accepted — clusters may be injected later."""
+        df = pd.DataFrame(
+            {
+                "y": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
+                "w": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
+                "strat": [0, 0, 0, 1, 1, 1],
+                "pop": [10.0, 10.0, 10.0, 20.0, 20.0, 20.0],
+            }
+        )
+        sd = SurveyDesign(
+            weights="w", weight_type="pweight", strata="strat", fpc="pop"
+        )
+        # Should not raise — FPC validation defers when no PSU declared
+        resolved = sd.resolve(df)
+        assert resolved.fpc is not None