refactor(survey): consolidate panel-to-unit collapse (ContinuousDiD/EfficientDiD)

Extract the row_idx->unit survey-design collapse hand-rolled at four sites
(continuous_did.py event-study/overall/bootstrap SE + efficient_did.py
build-once) into two diff_diff/survey.py helpers:

- ResolvedSurveyDesign.subset_to_units_by_row_idx(row_idx, unit_weights=None):
  folds the index-and-recount preamble around the existing subset_to_units.
- build_unit_first_row_index(unit_values, unit_order): first-panel-row index
  per unit (replaces ContinuousDiD's slow df.iterrows() build and
  EfficientDiD's inline .values scan).

Bit-identical: same np.unique + subset_to_units + arrays; survey-weighted SEs,
df_survey, and design-effect metadata are unchanged (oracle test locks it vs
the old inline logic). Removes ~95 lines of duplicated collapse.

StackedDiD is deliberately left on its own path (control units are duplicated
across sub-experiments, so it re-resolves at stacked granularity rather than
collapsing to one row per unit) with a clarifying comment; the residual
stacked-specific dedup is parked in TODO.md.

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

Author: igerber

CHANGELOG.md

@@ -34,6 +34,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   still defers cluster with a warning).
 
 ### Changed
+- **Internal: consolidated the panel-to-unit survey-design collapse** shared by `ContinuousDiD`
+  and `EfficientDiD` into two `diff_diff/survey.py` helpers —
+  `ResolvedSurveyDesign.subset_to_units_by_row_idx` (folds the index-and-recount preamble around
+  the existing `subset_to_units`) and `build_unit_first_row_index` (first-panel-row index per
+  unit). Removes four hand-rolled copies of the collapse plus a slow `df.iterrows()` index build.
+  No behavior change: survey-weighted SEs, `df_survey`, and design-effect metadata are
+  bit-identical (locked by an oracle test against the old inline logic).
 - **`DiagnosticReport` now routes `SpilloverDiDResults`.** Previously a fitted
   `SpilloverDiD` (Butts 2021) result matched no `_APPLICABILITY` / `_PT_METHOD` entry, so
   `DiagnosticReport(spillover_result)` reported every check as `not_applicable`. It now routes

TODO.md

@@ -28,7 +28,6 @@ The `Origin` column (Actionable tables) and the `PR` column (Deferred tables) bo
 
 | Issue | Location | Origin | Effort | Priority |
 |-------|----------|--------|--------|----------|
-| Survey-design resolution / collapse patterns are inconsistent across panel estimators — `ContinuousDiD` rebuilds unit-level design in SE code, `EfficientDiD` builds once in `fit()`, `StackedDiD` re-resolves on stacked data. Extract shared helpers for panel-to-unit collapse, post-filter re-resolution, and metadata recomputation. | `continuous_did.py`, `efficient_did.py`, `stacked_did.py` | #226 | Mid | Low |
 | `SyntheticControl` remaining ADH-2015 §4 items: the regression-weight `W^reg = X_0'(X_0 X_0')^{-1} X_1` extrapolation diagnostic (flag implied OLS weights outside `[0,1]`) and sparse-SC subset search (`l < J`, holding `V` fixed). LOO, in-time placebo, CV `V`-selection, and inverse-variance `V` have landed; these two are the deferred tail. | `synthetic_control.py`, `synthetic_control_results.py` | ADH-2015 | Mid | Low |
 | `SyntheticControl` conformal (CWZ 2021) extensions: (a) one-sided / signed-`t` variants (§7); (b) covariates in the conformal proxy (`X_jt`, eqs 4/6 — current proxy is outcomes-only); (c) AR / innovation-permutation path (Lemmas 5-7) for time-series proxies. The joint test, pointwise CIs, and average-effect CI have landed. | `conformal.py`, `synthetic_control_results.py` | CWZ-2021 | Heavy | Low |
 | `ContinuousDiD` CGBS-2024 extensions (matches R `contdid` v0.1.0 deferral set): (a) `covariates=` kwarg; (b) discrete-treatment saturated regression (integer dose currently warned, not routed to per-level coefficients); (c) lowest-dose-as-control per Remark 3.1 when `P(D=0)=0`. REGISTRY `## ContinuousDiD` → Implementation Checklist marks these `[ ]`. | `continuous_did.py` | CGBS-2024 | Heavy | Low |
@@ -123,6 +122,7 @@ Doable in principle, but no current caller and/or explicitly out of paper scope.
 
 | Issue | Location | PR | Priority |
 |-------|----------|----|----------|
+| `StackedDiD` survey re-resolution intra-file dedup (raw-weight extraction ×3, compose-normalize ×3, resolve-on-stacked ×2). The cross-estimator ContinuousDiD/EfficientDiD panel-to-unit collapse consolidation LANDED (#226 shared helpers `ResolvedSurveyDesign.subset_to_units_by_row_idx` / `build_unit_first_row_index`); StackedDiD is deliberately NOT on that path (control units are duplicated across sub-experiments, so it re-resolves at stacked granularity rather than collapsing to one row per unit). The residual is stacked-specific, low value, and touches the numerically-sensitive composed-weight renormalization. Post-filter re-resolution / metadata-recompute unification across the three estimators was assessed and is not warranted — they use genuinely different mechanisms and already delegate to shared `_resolve_survey_for_fit` / `compute_survey_metadata`. | `stacked_did.py` | #226 | Low |
 | dCDH parity-test SE/CI assertions only cover pure-direction scenarios; mixed-direction SE comparison is structurally apples-to-oranges (cell-count vs obs-count weighting). | `test_chaisemartin_dhaultfoeuille_parity.py` | #294 | Low |
 | `HeterogeneousAdoptionDiD` survey-design API consolidation (**scheduled: next minor bump**): drop the deprecated `survey=` / `weights=` kwargs on all 8 HAD surfaces; only `survey_design=` remains. Also fold the legacy back-end `weights=` routing into the unified `_resolve_survey_for_fit` path. DeprecationWarning has shipped; the removal is ~50 LoC gated on the semver bump. | `had.py`, `had_pretests.py` | next minor | Medium |
 | `HeterogeneousAdoptionDiD` joint cross-horizon covariance / sup-t bands: per-horizon SEs use independent sandwiches (paper-faithful pointwise CIs per Pierce-Schott Fig 2). Follow-ups (low demand): IF-based stacking for joint cross-horizon inference; analytical H×H covariance on the weighted ES path; a sup-t band on the unweighted ES path. | `had.py::_fit_event_study` | Phase 2b / 4.5 B | Low |

diff_diff/continuous_did.py

@@ -31,9 +31,9 @@
 )
 from diff_diff.linalg import _rank_guarded_inv, solve_ols
 from diff_diff.survey import (
-    ResolvedSurveyDesign,
     _resolve_survey_for_fit,
     _validate_unit_constant_survey,
+    build_unit_first_row_index,
     compute_survey_vcov,
 )
 from diff_diff.utils import safe_inference
@@ -413,8 +413,7 @@ def fit(
 
         # Filter out NaN cells (e.g., from zero effective survey mass)
         gt_results = {
-            gt: r for gt, r in gt_results.items()
-            if np.isfinite(r.get("att_glob", np.nan))
+            gt: r for gt, r in gt_results.items() if np.isfinite(r.get("att_glob", np.nan))
         }
 
         if len(gt_results) == 0:
@@ -573,9 +572,12 @@ def fit(
                 # Survey df for t-distribution inference (unit-level, not panel-level)
                 _survey_df = analytic.get("df_survey")
                 # Guard: replicate design with undefined df → NaN inference
-                if (_survey_df is None and resolved_survey is not None
-                        and hasattr(resolved_survey, 'uses_replicate_variance')
-                        and resolved_survey.uses_replicate_variance):
+                if (
+                    _survey_df is None
+                    and resolved_survey is not None
+                    and hasattr(resolved_survey, "uses_replicate_variance")
+                    and resolved_survey.uses_replicate_variance
+                ):
                     _survey_df = 0
 
                 # Recompute survey_metadata from unit-level design so reported
@@ -589,8 +591,7 @@ def fit(
 
                 # Propagate replicate df override to survey_metadata for display
                 # (but not the df=0 sentinel — keep metadata as None for undefined df)
-                if (_survey_df is not None and _survey_df != 0
-                        and survey_metadata is not None):
+                if _survey_df is not None and _survey_df != 0 and survey_metadata is not None:
                     if survey_metadata.df_survey != _survey_df:
                         survey_metadata.df_survey = _survey_df
 
@@ -624,30 +625,8 @@ def fit(
                     unit_resolved_es = None
                     if resolved_survey is not None:
                         row_idx = precomp["unit_first_panel_row"]
-                        uw = (
-                            precomp.get("unit_survey_weights")
-                            if precomp.get("unit_survey_weights") is not None
-                            else np.ones(n_units)
-                        )
-                        us = (
-                            resolved_survey.strata[row_idx]
-                            if resolved_survey.strata is not None
-                            else None
-                        )
-                        up = (
-                            resolved_survey.psu[row_idx]
-                            if resolved_survey.psu is not None
-                            else None
-                        )
-                        uf = (
-                            resolved_survey.fpc[row_idx]
-                            if resolved_survey.fpc is not None
-                            else None
-                        )
-                        n_strata_u = len(np.unique(us)) if us is not None else 0
-                        n_psu_u = len(np.unique(up)) if up is not None else 0
-                        unit_resolved_es = resolved_survey.subset_to_units(
-                            row_idx, uw, us, up, uf, n_strata_u, n_psu_u,
+                        unit_resolved_es = resolved_survey.subset_to_units_by_row_idx(
+                            row_idx, unit_weights=precomp.get("unit_survey_weights")
                         )
 
                     for e_val, info_e in event_study_effects.items():
@@ -711,13 +690,21 @@ def fit(
 
                                 # Score-scale: psi = w * if_es (matches TSL bread)
                                 psi_es = unit_resolved_es.weights * if_es
-                                variance, _nv = compute_replicate_if_variance(psi_es, unit_resolved_es)
-                                es_se = float(np.sqrt(max(variance, 0.0))) if np.isfinite(variance) else np.nan
+                                variance, _nv = compute_replicate_if_variance(
+                                    psi_es, unit_resolved_es
+                                )
+                                es_se = (
+                                    float(np.sqrt(max(variance, 0.0)))
+                                    if np.isfinite(variance)
+                                    else np.nan
+                                )
                             else:
                                 X_ones_es = np.ones((n_units, 1))
                                 tsl_scale_es = float(unit_resolved_es.weights.sum())
                                 if_es_tsl = if_es * tsl_scale_es
-                                vcov_es = compute_survey_vcov(X_ones_es, if_es_tsl, unit_resolved_es)
+                                vcov_es = compute_survey_vcov(
+                                    X_ones_es, if_es_tsl, unit_resolved_es
+                                )
                                 es_se = float(np.sqrt(np.abs(vcov_es[0, 0])))
                         else:
                             es_se = float(np.sqrt(np.sum(if_es**2)))
@@ -831,15 +818,11 @@ def _precompute_structures(
             unit_cohorts[i] = unit_first.loc[u, first_treat]
             dose_vector[i] = unit_first.loc[u, dose]
 
-        # Build unit-to-first-panel-row mapping (for subsetting panel-level arrays)
-        # This maps each unit index to the positional index of its first row in df.
-        unit_first_panel_row = np.zeros(n_units, dtype=int)
-        seen_units: set = set()
-        for pos_idx, (_, row) in enumerate(df.iterrows()):
-            u = row[unit]
-            if u not in seen_units:
-                seen_units.add(u)
-                unit_first_panel_row[unit_to_idx[u]] = pos_idx
+        # Build unit-to-first-panel-row mapping (for subsetting panel-level
+        # arrays): the positional index of each unit's first row in df, aligned
+        # to ``all_units`` (== ``unit_to_idx`` order since
+        # ``unit_to_idx = {u: i for i, u in enumerate(all_units)}``).
+        unit_first_panel_row = build_unit_first_row_index(df[unit].values, all_units)
 
         # Per-unit survey weights (take first obs per unit from panel data)
         unit_survey_weights = None
@@ -949,8 +932,10 @@ def _compute_dose_response_gt(
             # Guard against zero effective mass (e.g., after subpopulation)
             if np.sum(w_treated) <= 0 or np.sum(w_control) <= 0:
                 return {
-                    "att_glob": np.nan, "acrt_glob": np.nan,
-                    "n_treated": 0, "n_control": 0,
+                    "att_glob": np.nan,
+                    "acrt_glob": np.nan,
+                    "n_treated": 0,
+                    "n_control": 0,
                     "att_d": np.full(len(dvals), np.nan),
                     "acrt_d": np.full(len(dvals), np.nan),
                 }
@@ -1293,23 +1278,8 @@ def _compute_analytical_se(
             # but influence functions are unit-level (n_units). Build a unit-level
             # ResolvedSurveyDesign by subsetting to one obs per unit.
             row_idx = precomp["unit_first_panel_row"]
-            unit_weights = precomp.get("unit_survey_weights")
-            if unit_weights is None:
-                unit_weights = np.ones(n_units)
-
-            unit_strata = (
-                resolved_survey.strata[row_idx] if resolved_survey.strata is not None else None
-            )
-            unit_psu = resolved_survey.psu[row_idx] if resolved_survey.psu is not None else None
-            unit_fpc = resolved_survey.fpc[row_idx] if resolved_survey.fpc is not None else None
-
-            # Count unique strata/PSU in the unit-level subset
-            n_strata_unit = len(np.unique(unit_strata)) if unit_strata is not None else 0
-            n_psu_unit = len(np.unique(unit_psu)) if unit_psu is not None else 0
-
-            unit_resolved = resolved_survey.subset_to_units(
-                row_idx, unit_weights, unit_strata, unit_psu, unit_fpc,
-                n_strata_unit, n_psu_unit,
+            unit_resolved = resolved_survey.subset_to_units_by_row_idx(
+                row_idx, unit_weights=precomp.get("unit_survey_weights")
             )
 
             X_ones = np.ones((n_units, 1))
@@ -1370,7 +1340,11 @@ def _rep_se(if_vals):
 
         # Return unit-level survey df and resolved design for metadata recomputation
         # Only override with n_valid-based df when replicates were actually dropped
-        if resolved_survey is not None and hasattr(resolved_survey, 'uses_replicate_variance') and resolved_survey.uses_replicate_variance:
+        if (
+            resolved_survey is not None
+            and hasattr(resolved_survey, "uses_replicate_variance")
+            and resolved_survey.uses_replicate_variance
+        ):
             if _rep_n_valid < unit_resolved.n_replicates:
                 unit_df_survey = _rep_n_valid - 1 if _rep_n_valid > 1 else None
             else:
@@ -1415,7 +1389,11 @@ def _run_bootstrap(
 
         # Reject replicate-weight designs for bootstrap — replicate variance
         # is an analytical alternative to bootstrap, not compatible with it
-        if resolved_survey is not None and hasattr(resolved_survey, "uses_replicate_variance") and resolved_survey.uses_replicate_variance:
+        if (
+            resolved_survey is not None
+            and hasattr(resolved_survey, "uses_replicate_variance")
+            and resolved_survey.uses_replicate_variance
+        ):
             raise NotImplementedError(
                 "ContinuousDiD bootstrap (n_bootstrap > 0) is not supported "
                 "with replicate-weight survey designs. Replicate weights provide "
@@ -1429,22 +1407,9 @@ def _run_bootstrap(
         # Build unit-level ResolvedSurveyDesign for survey-aware bootstrap
         unit_resolved = None
         if resolved_survey is not None:
-            from diff_diff.survey import ResolvedSurveyDesign
-
             row_idx = precomp["unit_first_panel_row"]
-            unit_weights = precomp.get("unit_survey_weights")
-            if unit_weights is None:
-                unit_weights = np.ones(n_units)
-            unit_strata = (
-                resolved_survey.strata[row_idx] if resolved_survey.strata is not None else None
-            )
-            unit_psu = resolved_survey.psu[row_idx] if resolved_survey.psu is not None else None
-            unit_fpc = resolved_survey.fpc[row_idx] if resolved_survey.fpc is not None else None
-            n_strata_u = len(np.unique(unit_strata)) if unit_strata is not None else 0
-            n_psu_u = len(np.unique(unit_psu)) if unit_psu is not None else 0
-            unit_resolved = resolved_survey.subset_to_units(
-                row_idx, unit_weights, unit_strata, unit_psu, unit_fpc,
-                n_strata_u, n_psu_u,
+            unit_resolved = resolved_survey.subset_to_units_by_row_idx(
+                row_idx, unit_weights=precomp.get("unit_survey_weights")
             )
 
         # Generate bootstrap weights — PSU-level when survey design is present
@@ -1682,7 +1647,7 @@ def _bootstrap_gt_cell(gt, info):
                     boot_es[e],
                     alpha=self.alpha,
                     context=f"event study e={e}",
-                    )
+                )
                 es_se[e] = se_e
                 es_ci[e] = ci_e
                 es_p[e] = p_e

diff_diff/efficient_did.py

@@ -521,7 +521,7 @@ def fit(
             )
 
         # Resolve survey design if provided
-        from diff_diff.survey import _resolve_survey_for_fit
+        from diff_diff.survey import _resolve_survey_for_fit, build_unit_first_row_index
 
         resolved_survey, survey_weights, survey_weight_type, survey_metadata = (
             _resolve_survey_for_fit(survey_design, data, "analytical")
@@ -655,33 +655,14 @@ def fit(
         # order.  The previous approach (groupby cumcount == 0) yielded
         # first-appearance order which can differ from sorted order when the
         # input DataFrame is not pre-sorted by unit.
-        first_pos: Dict[Any, int] = {}
-        for i, u in enumerate(df[unit].values):
-            if u not in first_pos:
-                first_pos[u] = i
-        self._unit_first_panel_row = np.array([first_pos[u] for u in all_units])
+        self._unit_first_panel_row = build_unit_first_row_index(df[unit].values, all_units)
 
         # Build unit-level ResolvedSurveyDesign once (avoids repeated
         # construction in _compute_survey_eif_se and ensures consistent
         # unit-level df for safe_inference t-distribution).
         if resolved_survey is not None:
-            row_idx = self._unit_first_panel_row
-            unit_weights_s = resolved_survey.weights[row_idx]
-            unit_strata = (
-                resolved_survey.strata[row_idx] if resolved_survey.strata is not None else None
-            )
-            unit_psu = resolved_survey.psu[row_idx] if resolved_survey.psu is not None else None
-            unit_fpc = resolved_survey.fpc[row_idx] if resolved_survey.fpc is not None else None
-            n_strata_u = len(np.unique(unit_strata)) if unit_strata is not None else 0
-            n_psu_u = len(np.unique(unit_psu)) if unit_psu is not None else 0
-            self._unit_resolved_survey = resolved_survey.subset_to_units(
-                row_idx,
-                unit_weights_s,
-                unit_strata,
-                unit_psu,
-                unit_fpc,
-                n_strata_u,
-                n_psu_u,
+            self._unit_resolved_survey = resolved_survey.subset_to_units_by_row_idx(
+                self._unit_first_panel_row
             )
             # Use unit-level df (not panel-level) for t-distribution
             self._survey_df = self._unit_resolved_survey.df_survey

diff_diff/stacked_did.py

@@ -773,6 +773,10 @@ def _refit_stacked(w_r):
 
             # Re-resolve survey design on the stacked data so that strata/PSU
             # arrays have the correct length for TSL variance estimation.
+            # (Unlike ContinuousDiD/EfficientDiD, StackedDiD does NOT collapse
+            # to unit level via ``subset_to_units_by_row_idx``: control units are
+            # duplicated across sub-experiments, so the design must be resolved
+            # at the stacked granularity, not one row per original unit.)
             resolved_stacked = survey_design.resolve(stacked_df)
 
             # Create a copy with composed weights (normalized to sum=n_stacked)