Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
igerber merged 1 commit into from
Jul 3, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion TODO.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 |
Expand Down Expand Up @@ -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 |
Expand Down
127 changes: 46 additions & 81 deletions diff_diff/continuous_did.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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)))
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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),
}
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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 "
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
27 changes: 4 additions & 23 deletions diff_diff/efficient_did.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions diff_diff/stacked_did.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading
Loading