Fix CI review R1: SE divisor scaling, zero-weight cells, dead helper

- P1-A: Scale U_centered by 1/divisor before survey IF expansion.
  dCDH IFs are numerator-scale (U.sum() == N_S * DID_M), but
  compute_survey_if_variance() expects estimator-scale psi.
- P1-B: Zero-weight cells (w_gt <= 0) now treated as absent by
  setting n_gt=0, preventing NaN propagation into estimates.
- P2: Add SE-pinning test (uniform weights + PSU=group matches
  plug-in SE) and zero-weight cell exclusion test.
- P3: Delete unused _validate_group_constant_survey() from survey.py
  that contradicted the supported within-group variation contract.

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

Author: igerber

diff_diff/chaisemartin_dhaultfoeuille.py

@@ -227,6 +227,12 @@ def _validate_and_aggregate_to_cells(
         )
         cell["y_gt"] = cell["_wy_sum"] / cell["w_gt"]
         cell = cell.drop(columns=["_wy_sum"])
+        # Zero-weight cells: treat as absent so downstream presence
+        # logic (N_mat > 0) correctly excludes them.
+        zero_w_mask = cell["w_gt"] <= 0
+        if zero_w_mask.any():
+            cell.loc[zero_w_mask, "n_gt"] = 0
+            cell.loc[zero_w_mask, "y_gt"] = 0.0
         df.drop(columns=["_w_", "_wy_"], inplace=True)
     else:
         cell = df.groupby([group, time], as_index=False).agg(
@@ -4548,8 +4554,14 @@ def _compute_se(
         return _plugin_se(U_centered=U_centered, divisor=divisor)
     if eligible_groups is None:
         return _plugin_se(U_centered=U_centered, divisor=divisor)
+    if divisor <= 0:
+        return float("nan")
+    # dCDH IFs are numerator-scale (U.sum() == N_S * DID_M).
+    # compute_survey_if_variance() expects estimator-scale psi.
+    # Scale by 1/divisor to normalize before survey expansion.
+    U_scaled = U_centered / divisor
     return _survey_se_from_group_if(
-        U_centered=U_centered,
+        U_centered=U_scaled,
         eligible_groups=eligible_groups,
         obs_survey_info=obs_survey_info,
     )

diff_diff/survey.py

@@ -911,47 +911,6 @@ def _validate_unit_constant_survey(data, unit_col, survey_design):
                 )
 
 
-def _validate_group_constant_survey(data, group_col, survey_design):
-    """Validate that survey design columns are constant within groups.
-
-    The dCDH estimator aggregates to ``(group, time)`` cells and then
-    works at the group level. Survey columns (weights, strata, PSU)
-    must not vary within groups for the IF expansion and survey variance
-    to be well-defined.
-
-    Parameters
-    ----------
-    data : pd.DataFrame
-        Input data (pre-aggregation).
-    group_col : str
-        Group identifier column name.
-    survey_design : SurveyDesign
-        Survey design specification (uses attribute names, not resolved arrays).
-
-    Raises
-    ------
-    ValueError
-        If any survey column varies within groups.
-    """
-    cols_to_check = [
-        survey_design.weights,
-        survey_design.strata,
-        survey_design.psu,
-        survey_design.fpc,
-    ]
-    for col in cols_to_check:
-        if col is not None and col in data.columns:
-            n_unique = data.groupby(group_col)[col].nunique()
-            varying_groups = n_unique[n_unique > 1]
-            if len(varying_groups) > 0:
-                raise ValueError(
-                    f"Survey column '{col}' varies within groups "
-                    f"(found {len(varying_groups)} groups with multiple values). "
-                    f"dCDH survey support requires survey design columns to be "
-                    f"constant within groups."
-                )
-
-
 def _resolve_pweight_only(resolved_survey, estimator_name):
     """Guard: reject non-pweight and strata/PSU/FPC for pweight-only estimators.
 

tests/test_survey_dcdh.py

@@ -377,3 +377,68 @@ def test_bootstrap_survey_emits_warning(self, data_with_survey):
                 treatment="treatment",
                 survey_design=sd,
             )
+
+
+# ── Test: SE scale pinning ──────────────────────────────────────────
+
+
+class TestSEScalePinning:
+    """Survey SE with uniform weights and no strata/PSU must match plug-in SE."""
+
+    def test_uniform_survey_se_matches_plugin(self, base_data):
+        """Pins the divisor normalization: uniform survey SE with group-level
+        PSU clustering should be close to plug-in SE.
+
+        Without PSU clustering, survey treats each observation as independent
+        (N_obs observations), while plug-in treats each group as independent
+        (N_groups). Clustering at the group level aligns the two.
+        """
+        df = base_data.copy()
+        df["pw"] = 1.0
+        sd = SurveyDesign(weights="pw", psu="group")
+
+        r_plain = ChaisemartinDHaultfoeuille(seed=1).fit(
+            base_data, outcome="outcome", group="group",
+            time="period", treatment="treatment",
+        )
+        r_survey = ChaisemartinDHaultfoeuille(seed=1).fit(
+            df, outcome="outcome", group="group",
+            time="period", treatment="treatment",
+            survey_design=sd,
+        )
+        # With PSU=group and uniform weights, survey SE should be
+        # close to plug-in SE (both assume group-level independence).
+        # Small-sample corrections (n/(n-1)) cause minor differences.
+        if np.isfinite(r_plain.overall_se) and np.isfinite(r_survey.overall_se):
+            assert r_plain.overall_se == pytest.approx(
+                r_survey.overall_se, rel=0.15
+            ), (
+                f"Survey SE ({r_survey.overall_se:.6f}) should be close to "
+                f"plug-in SE ({r_plain.overall_se:.6f}) with uniform weights "
+                f"and PSU=group"
+            )
+
+
+# ── Test: Zero-weight cells ─────────────────────────────────────────
+
+
+class TestZeroWeightCells:
+
+    def test_zero_weight_cell_excluded(self, base_data):
+        """A cell with zero survey weight is treated as absent."""
+        df = base_data.copy()
+        df["pw"] = 1.0
+        # Zero out weight for one group at one period
+        target_group = df["group"].unique()[0]
+        target_period = df["period"].unique()[1]
+        mask = (df["group"] == target_group) & (df["period"] == target_period)
+        df.loc[mask, "pw"] = 0.0
+        sd = SurveyDesign(weights="pw")
+
+        # Should not raise; the zero-weight cell is just absent
+        result = ChaisemartinDHaultfoeuille(seed=1).fit(
+            df, outcome="outcome", group="group",
+            time="period", treatment="treatment",
+            survey_design=sd,
+        )
+        assert np.isfinite(result.overall_att)