Move FPC >= n_PSU validation to compute_survey_vcov for effective PSU structure from PR #218 review (round 17)

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

Author: igerber

diff_diff/survey.py

@@ -169,14 +169,10 @@ def resolve(self, data: pd.DataFrame) -> "ResolvedSurveyDesign":
             if np.any(np.isnan(fpc_arr)) or np.any(~np.isfinite(fpc_arr)):
                 raise ValueError("FPC values must be finite and non-NaN")
 
-            # FPC requires survey structure (psu or strata)
-            if self.psu is None and self.strata is None:
-                raise ValueError(
-                    "FPC requires either psu or strata to be specified. "
-                    "FPC alone without survey structure is not supported."
-                )
-
-            # Validate FPC >= n_h per stratum
+            # Validate FPC structure (constant within strata, positive).
+            # FPC >= n_PSU validation is deferred to compute_survey_vcov()
+            # where the final effective PSU structure is known (after
+            # cluster-as-PSU injection and implicit per-obs PSU fallback).
             if strata_arr is not None:
                 for h in np.unique(strata_arr):
                     mask_h = strata_arr == h
@@ -188,27 +184,23 @@ def resolve(self, data: pd.DataFrame) -> "ResolvedSurveyDesign":
                             f"Stratum {h} has values: {np.unique(fpc_vals)}"
                         )
                     fpc_h = fpc_vals[0]
-                    # Validate FPC >= number of sampled PSUs (not obs count)
+                    # Validate FPC >= n_PSU when explicit PSU is declared
                     if psu_arr is not None:
                         n_psu_h = len(np.unique(psu_arr[mask_h]))
                         if fpc_h < n_psu_h:
                             raise ValueError(
                                 f"FPC ({fpc_h}) is less than the number of PSUs "
                                 f"({n_psu_h}) in stratum {h}. FPC must be >= n_PSU."
                             )
-                    else:
-                        # 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:
+            else:
                 # No strata: require FPC is a single constant value
                 if len(np.unique(fpc_arr)) > 1:
                     raise ValueError(
                         "FPC values must be constant when no strata are specified. "
                         f"Found {len(np.unique(fpc_arr))} distinct values."
                     )
-                if fpc_arr[0] < n_psu:
+                # Validate FPC >= n_PSU when explicit PSU is declared
+                if psu_arr is not None and fpc_arr[0] < n_psu:
                     raise ValueError(
                         f"FPC ({fpc_arr[0]}) is less than the number of PSUs "
                         f"({n_psu}). FPC must be >= number of PSUs."
@@ -544,6 +536,11 @@ def compute_survey_vcov(
             f_h = 0.0  # No FPC
             if resolved.fpc is not None:
                 N_h = resolved.fpc[0]
+                if N_h < n_psu:
+                    raise ValueError(
+                        f"FPC ({N_h}) is less than the number of effective PSUs "
+                        f"({n_psu}). FPC must be >= n_PSU."
+                    )
                 f_h = n_psu / N_h
                 if f_h >= 1.0:
                     legitimate_zero_count += 1
@@ -599,6 +596,11 @@ def compute_survey_vcov(
             f_h = 0.0
             if resolved.fpc is not None:
                 N_h = resolved.fpc[mask_h][0]
+                if N_h < n_psu_h:
+                    raise ValueError(
+                        f"FPC ({N_h}) is less than the number of effective PSUs "
+                        f"({n_psu_h}) in stratum. FPC must be >= n_PSU."
+                    )
                 f_h = n_psu_h / N_h
                 if f_h >= 1.0:
                     legitimate_zero_count += 1

tests/test_survey.py

@@ -1232,8 +1232,8 @@ def test_fpc_nonconstant_within_stratum(self):
         with pytest.raises(ValueError, match="constant within each stratum"):
             sd.resolve(df)
 
-    def test_fpc_only_without_psu_raises(self):
-        """FPC without psu or strata raises ValueError (P1)."""
+    def test_fpc_only_without_psu_resolves(self):
+        """FPC without psu or strata resolves — validated later against effective PSUs."""
         n = 10
         df = pd.DataFrame(
             {
@@ -1243,8 +1243,8 @@ def test_fpc_only_without_psu_raises(self):
             }
         )
         sd = SurveyDesign(weights="w", fpc="fpc")
-        with pytest.raises(ValueError, match="FPC requires either psu or strata"):
-            sd.resolve(df)
+        resolved = sd.resolve(df)
+        assert resolved.fpc is not None
 
     def test_weighted_within_transform_matches_explicit_wls(self):
         """Weighted within_transform + OLS matches explicit WLS with dummies (P0-2).
@@ -2901,7 +2901,7 @@ def test_injected_cluster_nested_in_strata(self):
         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."""
+        """FPC + strata (no PSU) resolves — FPC validated later against effective PSUs."""
         df = pd.DataFrame(
             {
                 "y": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
@@ -2913,6 +2913,46 @@ def test_fpc_with_strata_no_psu_accepted(self):
         sd = SurveyDesign(
             weights="w", weight_type="pweight", strata="strat", fpc="pop"
         )
-        # Should not raise — FPC validation defers when no PSU declared
+        # Should not raise at resolve time — FPC >= n_PSU validated at vcov time
         resolved = sd.resolve(df)
         assert resolved.fpc is not None
+
+    def test_fpc_alone_no_strata_no_psu_accepted(self):
+        """FPC alone (no PSU/strata) resolves — clusters may be injected later."""
+        df = pd.DataFrame(
+            {
+                "y": [1.0, 2.0, 3.0, 4.0],
+                "w": [1.0, 1.0, 1.0, 1.0],
+                "pop": [100.0, 100.0, 100.0, 100.0],
+            }
+        )
+        sd = SurveyDesign(weights="w", weight_type="pweight", fpc="pop")
+        resolved = sd.resolve(df)
+        assert resolved.fpc is not None
+
+    def test_fpc_lt_effective_npsu_rejected_at_vcov(self):
+        """FPC < effective n_PSU is rejected at compute_survey_vcov time."""
+        np.random.seed(42)
+        n = 12
+        strata = np.repeat([0, 1], 6)
+        psu = np.tile(np.arange(3), 4)  # 3 PSUs per stratum
+
+        X = np.column_stack([np.ones(n), np.random.randn(n)])
+        residuals = np.random.randn(n)
+        weights = np.ones(n)
+
+        # FPC = 2 per stratum, but we have 3 PSUs → invalid
+        fpc = np.array([2.0] * n)
+
+        resolved = ResolvedSurveyDesign(
+            weights=weights,
+            weight_type="pweight",
+            strata=strata,
+            psu=psu,
+            fpc=fpc,
+            n_strata=2,
+            n_psu=6,
+            lonely_psu="remove",
+        )
+        with pytest.raises(ValueError, match="FPC.*less than.*effective PSUs"):
+            compute_survey_vcov(X, residuals, resolved=resolved)