Round-6 CI P1: narrow analytical guard to within-group-varying PSU only

The round-5 analytical guard in _survey_se_from_group_if was over-
scoped: it fired on any panel with leaked centered mass, including
PSU-within-group-constant regimes (PSU=group auto-inject, strictly-
coarser-PSU-within-group-constant). But the docs said those regimes
are the intended workaround. Panels with terminal missingness and
PSU=group are the documented supported path and must not raise.

Fix: add an analytical dispatcher inside `_survey_se_from_group_if`
mirroring the bootstrap's `_psu_varies_within_group` routing. When
PSU is within-group-constant on the eligible subset, fall back to
the legacy group-level allocator (which uses U_centered[g] directly
via the row-sum identity and does not trigger the sentinel-mass
guard). Only when PSU actually varies within group does the cell
allocator run — and only then can the sentinel-mass guard fire.

Byte-identity: under PSU=group the row-sum identity makes the cell
and group allocators statistically equivalent, but the legacy
branch was the one in use before PR #323 introduced the cell
allocator. Rerouting to it under within-group-constant regimes is
a regression-free fallback (it's what PR #323 aimed to preserve
via byte-identity in the first place).

Test coverage: added `test_fit_succeeds_on_terminal_missingness_with_psu_group`
— same fixture as the varying-PSU failure regression but with
auto-inject `psu=<group>`, asserts both `n_bootstrap=0` and
`n_bootstrap > 0` succeed with finite SE.

Also updated the cell-level bootstrap `ValueError` text (and the
_unroll_target_to_cells docstring) to no longer advertise
`n_bootstrap=0` as a workaround — both paths now fail consistently
on the varying-PSU carve-out, and the correct workaround is
pre-processing or using an explicit `psu=<group_col>`.

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

Author: igerber

diff_diff/chaisemartin_dhaultfoeuille.py

@@ -5863,6 +5863,38 @@ def _survey_se_from_group_if(
         and U_centered_per_period.size > 0
     )
 
+    # When the cell allocator would apply but PSU is within-group-
+    # constant (PSU=group, strictly-coarser PSU within-group-constant,
+    # or the auto-inject default), the cell and group allocators are
+    # equivalent at PSU-level aggregation via the row-sum identity
+    # `sum_{c in g} u_cell[c] == u_centered[g]`. Prefer the legacy
+    # group-level path in that regime: it sidesteps the sentinel-
+    # mass guard (below) that would otherwise fire spuriously on
+    # terminally-missing panels whose PSU structure does not require
+    # cell-level resolution. Matches the bootstrap dispatcher's
+    # routing rule (`_compute_dcdh_bootstrap` + `_psu_varies_within_group`).
+    if use_cell_allocator:
+        psu_arr = getattr(resolved, "psu", None)
+        if psu_arr is None:
+            use_cell_allocator = False
+        else:
+            psu_eff = np.asarray(psu_arr)[pos_mask]
+            eligible_set = set(eligible_groups)
+            elig_row_mask = np.array(
+                [g in eligible_set for g in gids_eff], dtype=bool
+            )
+            if elig_row_mask.any():
+                psu_varies_within = bool(
+                    pd.DataFrame({
+                        "g": gids_eff[elig_row_mask],
+                        "p": psu_eff[elig_row_mask],
+                    }).groupby("g")["p"].nunique().gt(1).any()
+                )
+                if not psu_varies_within:
+                    use_cell_allocator = False
+            else:
+                use_cell_allocator = False
+
     if use_cell_allocator:
         tids_eff = np.asarray(time_ids)[pos_mask]
         # Map row's group to an index in eligible_groups (−1 when the

diff_diff/chaisemartin_dhaultfoeuille_bootstrap.py

@@ -614,9 +614,13 @@ def _unroll_target_to_cells(
     mass at that sentinel cell. The cell-level bootstrap cannot
     allocate that mass to any PSU (the cell has no positive-weight
     obs), so silently dropping it would under-weight the group's
-    bootstrap contribution. The conservative guard rejects the
-    combination and points users to ``n_bootstrap=0`` (analytical
-    TSL) as the documented alternative for such panels.
+    bootstrap contribution. The analytical TSL path shares the same
+    cell-period allocator and fires a matching guard in
+    ``_survey_se_from_group_if``, so both paths reject this regime
+    consistently. Documented workarounds: pre-process the panel
+    (drop late-exit groups or trim to a balanced sub-panel), or use
+    an explicit ``psu=<group_col>`` so both analytical and bootstrap
+    paths route through the legacy group-level allocator.
 
     Returns ``(u_cell, psu_cell)`` of shape
     ``(n_valid_cells_in_target,)`` each.
@@ -658,10 +662,14 @@ def _unroll_target_to_cells(
             "within-group-varying PSU: `_cohort_recenter_per_period` "
             "subtracts column means across the full period grid, so a "
             "group with no observation at period t acquires non-zero "
-            "centered mass there, which the cell-level bootstrap "
-            "cannot allocate to any PSU. Use `n_bootstrap=0` to fall "
-            "back to analytical TSL variance (which supports this "
-            "panel regime)."
+            "centered mass there, which the cell-level allocator "
+            "cannot allocate to any PSU. The analytical TSL path "
+            "(`_survey_se_from_group_if`) fires a matching guard on "
+            "the same regime, so both paths reject this panel "
+            "consistently. Pre-process the panel to remove terminal "
+            "missingness (drop late-exit groups or trim to a balanced "
+            "sub-panel), or use an explicit `psu=<group_col>` so both "
+            "paths route through the legacy group-level allocator."
         )
     return flat_u[mask], flat_psu[mask].astype(np.int64, copy=False)
 

tests/test_survey_dcdh.py

@@ -2186,6 +2186,62 @@ def test_fit_raises_on_terminal_missingness_with_varying_psu(self):
                     survey_design=sd, L_max=1,
                 )
 
+    def test_fit_succeeds_on_terminal_missingness_with_psu_group(self):
+        """Companion regression: the terminal-missingness + varying-PSU
+        sentinel-mass guard must NOT fire when PSU is within-group-
+        constant. Same fixture as the varying-PSU test above but with
+        `psu=<group column>` (auto-inject default) — both analytical
+        and bootstrap paths route through the legacy group-level
+        allocator (the analytical dispatcher in
+        `_survey_se_from_group_if` falls back to the group-level
+        allocator when PSU does not vary within group; the bootstrap
+        dispatcher in `_compute_dcdh_bootstrap` does the same). Fit
+        must succeed with finite SE.
+        """
+        rows = []
+        for g in range(10):
+            if g < 5:
+                d_pattern = [0, 0, 0, 1, 1, 1]
+            elif g < 8:
+                d_pattern = [1, 1, 1, 1, 0, 0]
+            else:
+                d_pattern = [0, 0, 0, 0, 0, 0]
+            for t in range(6):
+                if g == 2 and t >= 4:
+                    continue
+                d = d_pattern[t]
+                y = float(g) + 0.1 * t + 1.0 * d
+                rows.append({
+                    "group": int(g),
+                    "period": int(t),
+                    "treatment": int(d),
+                    "outcome": y,
+                    "pw": 1.0,
+                })
+        df_ = pd.DataFrame(rows)
+        # Auto-inject: no explicit `psu` → `SurveyDesign` falls back to
+        # `psu=<group_col>` at fit() time. Within-group-constant.
+        sd = SurveyDesign(weights="pw")
+        import warnings as _w
+        for n_boot in (0, 50):
+            with _w.catch_warnings():
+                _w.simplefilter("ignore")
+                res = ChaisemartinDHaultfoeuille(
+                    n_bootstrap=n_boot, seed=1,
+                ).fit(
+                    df_, outcome="outcome", group="group",
+                    time="period", treatment="treatment",
+                    survey_design=sd, L_max=1,
+                )
+            assert np.isfinite(res.overall_att), (
+                f"n_bootstrap={n_boot}: overall_att must be finite "
+                f"under PSU=group + terminal missingness."
+            )
+            assert np.isfinite(res.overall_se) and res.overall_se >= 0.0, (
+                f"n_bootstrap={n_boot}: overall_se must be finite "
+                f"under PSU=group + terminal missingness."
+            )
+
     def test_bootstrap_dense_codes_under_singleton_baseline_excluded_group(self):
         """Regression for P0 #2: when a group is singleton-baseline-
         excluded (e.g., an always-treated group whose baseline D=1