Extend dCDH heterogeneity SE to cell-period allocator

Lifts the NotImplementedError gate for heterogeneity= + within-group-
varying PSU/strata under survey designs. The heterogeneity WLS
coefficient IF psi_g is now attributed in full to the (g, out_idx)
post-period cell and expanded to observation level as
psi_i = psi_g * (w_i / W_{g, out_idx}) — the DID_l single-cell
convention shipped in PR #323. Under PSU=group the per-obs
distribution differs from the legacy psi_i = psi_g * (w_i / W_g)
expansion, but the PSU-level aggregate telescopes to psi_g in both
paths, so Binder TSL variance and Rao-Wu replicate variance are
byte-identical under PSU=group; under within-group-varying PSU,
mass lands in the post-period PSU of the transition.

Tests: flipped the gating test to assert all five inference fields
finite; added PSU-level byte-identity unit test constructing both
psi_obs arrays and asserting compute_survey_if_variance agreement
within ULP; added nest=True + varying-strata + heterogeneity
smoke test (newly-unblocked regime); added multi-horizon smoke test;
added slow-tier MC null-coverage test (500 reps, within-group-
varying PSU, empirical 95% coverage inside [0.925, 0.975]).

n_bootstrap > 0 + within-group-varying PSU remains gated (follow-up
PR). Updated REGISTRY.md heterogeneity Note + Survey IF expansion
Note scope-limitations paragraph; updated _compute_heterogeneity_test
docstring + the stale legacy-allocator comment in
_survey_se_from_group_if; added CHANGELOG Changed entry.

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

Author: igerber

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 - Add Zenodo DOI badge to README; upgrade the BibTeX citation block with the concept DOI (`10.5281/zenodo.19646175`) and list author as Isaac Gerber (matching `CITATION.cff`). Add `doi:` and `identifiers:` entries (concept + versioned) to `CITATION.cff`. DOI was minted by Zenodo when v3.1.3 was released.
+- **`ChaisemartinDHaultfoeuille` heterogeneity + within-group-varying PSU/strata now supported** - `fit(heterogeneity=..., survey_design=...)` no longer raises `NotImplementedError` when the resolved design's PSU or strata vary across the cells of a group. The heterogeneity WLS coefficient IF is now expanded to observation level via the cell-period allocator (`ψ_i = ψ_g * (w_i / W_{g, out_idx})` on the post-period cell), consistent with the DID_l post-period single-cell convention shipped in v3.1.x. Under PSU=group the PSU-level Binder TSL variance and Rao-Wu replicate variance are byte-identical to the previous release (the PSU-level aggregate telescopes to `ψ_g` in both expansions); under within-group-varying PSU, mass lands in the post-period PSU of the transition. `n_bootstrap > 0` combined with within-group-varying PSU remains gated with `NotImplementedError` — the PSU-level Hall-Mammen wild bootstrap still uses the legacy group-level PSU map and will be extended in a follow-up PR.
 
 ## [3.1.3] - 2026-04-18
 

diff_diff/chaisemartin_dhaultfoeuille.py

@@ -665,10 +665,7 @@ def fit(
             clustered bootstrap. **Out-of-scope combinations raise
             ``NotImplementedError``**: (a) replicate weights with
             ``n_bootstrap > 0`` (replicate variance is closed-form;
-            bootstrap would double-count variance); (b)
-            ``heterogeneity=`` with PSU/strata that vary within group
-            (heterogeneity WLS still uses the legacy group-level IF
-            expansion; follow-up PR extends it); (c) ``n_bootstrap >
+            bootstrap would double-count variance); (b) ``n_bootstrap >
             0`` with PSU that varies within group (PSU-level bootstrap
             still uses the legacy group-level PSU map; follow-up PR
             extends it). See REGISTRY.md
@@ -828,39 +825,25 @@ def fit(
 
             # Cell-period IF allocator contract: strata and PSU must be
             # constant within each (g, t) cell, a strict relaxation of
-            # the previous within-group constancy rule. Two out-of-scope
-            # combinations are gated with NotImplementedError until the
-            # corresponding follow-up PRs extend them:
-            #   - heterogeneity= + within-group-varying PSU/strata
-            #     (PR 3: cell-period allocator for the WLS psi_obs)
+            # the previous within-group constancy rule. One out-of-scope
+            # combination remains gated with NotImplementedError until
+            # the corresponding follow-up PR extends it:
             #   - n_bootstrap > 0 + within-group-varying PSU
             #     (PR 4: cell-level Hall-Mammen wild bootstrap)
-            strata_varies, psu_varies = _strata_psu_vary_within_group(
+            _, psu_varies = _strata_psu_vary_within_group(
                 resolved_survey, data, group, survey_weights,
             )
-            if strata_varies or psu_varies:
-                if heterogeneity is not None:
-                    raise NotImplementedError(
-                        "heterogeneity= is not supported under a survey "
-                        "design whose PSU or strata vary within group. "
-                        "The heterogeneity WLS path uses the legacy "
-                        "group-level IF expansion and will be extended "
-                        "to the cell-period allocator in a follow-up "
-                        "PR. For now, either (a) set heterogeneity=None, "
-                        "or (b) collapse PSU/strata to be constant "
-                        "within each group."
-                    )
-                if self.n_bootstrap > 0:
-                    raise NotImplementedError(
-                        "n_bootstrap > 0 is not supported under a "
-                        "survey design whose PSU varies within group. "
-                        "The PSU-level Hall-Mammen wild bootstrap uses "
-                        "the legacy group-level PSU map and will be "
-                        "extended to cell-level PSU in a follow-up PR. "
-                        "For now, use n_bootstrap=0 (analytical TSL "
-                        "variance, which fully supports within-group-"
-                        "varying PSU via the cell-period allocator)."
-                    )
+            if psu_varies and self.n_bootstrap > 0:
+                raise NotImplementedError(
+                    "n_bootstrap > 0 is not supported under a "
+                    "survey design whose PSU varies within group. "
+                    "The PSU-level Hall-Mammen wild bootstrap uses "
+                    "the legacy group-level PSU map and will be "
+                    "extended to cell-level PSU in a follow-up PR. "
+                    "For now, use n_bootstrap=0 (analytical TSL "
+                    "variance, which fully supports within-group-"
+                    "varying PSU via the cell-period allocator)."
+                )
             _validate_cell_constant_strata_psu(
                 resolved_survey, data, group, time, survey_weights,
             )
@@ -3734,15 +3717,29 @@ def _compute_heterogeneity_test(
         Required when ``obs_survey_info`` is supplied.
     obs_survey_info : dict, optional
         Observation-level survey info with keys ``group_ids`` (raw per-row
-        group labels), ``weights`` (per-row survey weights), and ``resolved``
-        (ResolvedSurveyDesign). When provided, the regression uses WLS with
-        per-group weights W_g = sum of obs survey weights. SE is computed
-        via Binder TSL IF expansion through ``compute_survey_if_variance``
-        by default; under a replicate-weight design (BRR/Fay/JK1/JKn/SDR),
-        dispatches to ``compute_replicate_if_variance`` for Rao-Wu-style
-        variance. The effective df for t-critical values follows the
-        site-level ``min(df_s, n_valid_het - 1)`` rule and the helper
-        mutates ``replicate_n_valid_list`` so the final
+        group labels), ``time_ids`` (raw per-row period labels),
+        ``weights`` (per-row survey weights), ``resolved``
+        (ResolvedSurveyDesign), and ``periods`` (sorted canonical period
+        array matching ``Y_mat``'s column order). When provided, the
+        regression uses WLS with per-group weights
+        ``W_g = sum of obs survey weights in group g``. The group-level
+        WLS coefficient IF is
+        ``ψ_g = inv(X'WX)[1,:] @ x_g * W_g * r_g``. Under the cell-period
+        allocator, ``ψ_g`` is attributed in full to the ``(g, out_idx)``
+        post-period cell and expanded to observation level as
+        ``ψ_i = ψ_g * (w_i / W_{g, out_idx})`` for obs in that cell,
+        zero elsewhere. Under PSU=group this produces the same PSU-level
+        aggregate as the legacy ``ψ_i = ψ_g * (w_i / W_g)`` expansion,
+        so Binder TSL variance (and Rao-Wu replicate variance) are
+        byte-identical to the pre-cell-period release; under
+        within-group-varying PSU, mass lands in the post-period PSU of
+        the transition. SE is computed through
+        ``compute_survey_if_variance`` by default; under a
+        replicate-weight design (BRR/Fay/JK1/JKn/SDR), dispatches to
+        ``compute_replicate_if_variance`` for Rao-Wu-style variance. The
+        effective df for t-critical values follows the site-level
+        ``min(df_s, n_valid_het - 1)`` rule and the helper mutates
+        ``replicate_n_valid_list`` so the final
         ``_effective_df_survey(...)`` sees this site's n_valid.
     replicate_n_valid_list : list[int], optional
         Shared accumulator for replicate-weight ``n_valid`` counts across
@@ -3931,15 +3928,35 @@ def _compute_heterogeneity_test(
                 XtWX_inv = np.linalg.pinv(XtWX)
                 psi_g = (XtWX_inv[1, :] @ design.T) * W_elig * r_g  # (n_eligible,)
 
-            # Expand to obs level: ψ_i = ψ_g * (w_i / W_g) for i in group g.
-            psi_obs = np.zeros(len(obs_w_raw))
+            # Cell-period allocator: attribute ψ_g in full to the
+            # (g, out_idx) post-period cell (DID_l single-cell convention,
+            # see REGISTRY.md ChaisemartinDHaultfoeuille survey IF
+            # expansion Note). Observation-level expansion:
+            #   ψ_i = ψ_g * (w_i / W_{g, out_idx})
+            # for obs in (g, out_idx), zero elsewhere. Under PSU=group the
+            # per-observation distribution differs from the legacy
+            # ψ_i = ψ_g * (w_i / W_g) path, but the PSU-level aggregate
+            # telescopes to the same ψ_g — so compute_survey_if_variance
+            # and compute_replicate_if_variance (both aggregate to PSU
+            # first) produce byte-identical variance. Under within-group-
+            # varying PSU, mass lands in the post-period PSU of the
+            # transition, which is what Binder TSL needs.
+            obs_tids = np.asarray(obs_survey_info["time_ids"])
+            periods_arr = np.asarray(obs_survey_info["periods"])
+            psi_obs = np.zeros(len(obs_w_raw), dtype=np.float64)
             for e_idx, g_idx in enumerate(eligible):
                 gid = gid_list[g_idx]
-                mask_g = (obs_gids_raw == gid) & valid
-                w_sum_g = obs_w_raw[mask_g].sum()
-                if w_sum_g > 0:
-                    psi_obs[mask_g] = psi_g[e_idx] * (
-                        obs_w_raw[mask_g] / w_sum_g
+                out_idx = first_switch_idx[g_idx] - 1 + l_h
+                t_val_out = periods_arr[out_idx]
+                mask_cell = (
+                    (obs_gids_raw == gid)
+                    & (obs_tids == t_val_out)
+                    & valid
+                )
+                w_cell = obs_w_raw[mask_cell].sum()
+                if w_cell > 0:
+                    psi_obs[mask_cell] = psi_g[e_idx] * (
+                        obs_w_raw[mask_cell] / w_cell
                     )
 
             # Dispatch: replicate-weight variance (BRR/Fay/JK1/JKn/SDR)
@@ -5766,10 +5783,12 @@ def _survey_se_from_group_if(
     else:
         # Legacy group-level allocator (no per-period attribution
         # provided, or time/period info unavailable). Preserved for
-        # paths that haven't threaded per-period attribution through
-        # yet (e.g., the heterogeneity psi_obs construction in
-        # _compute_heterogeneity_test — gated to within-group-constant
-        # PSU in Stage 2 per PR 2 scope).
+        # defensive fallback and for unit tests that exercise the
+        # legacy allocator. No current caller in fit() uses this
+        # branch — ATT / joiners / leavers / placebos all thread
+        # U_centered_per_period, and heterogeneity (as of PR 3)
+        # constructs its own cell-period psi_obs and calls
+        # compute_survey_if_variance directly.
         group_to_u = {gid: U_centered[idx] for idx, gid in enumerate(eligible_groups)}
         u_obs_eff = np.array([group_to_u.get(gid, 0.0) for gid in gids_eff])
         unique_gids, inverse = np.unique(gids_eff, return_inverse=True)