From 0439a3890692c2b7e1109e79daee9d18325b699a Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 1 Apr 2026 16:56:18 -0400 Subject: [PATCH 01/11] Add replicate weight support to 7 estimators MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Expand replicate weight variance (BRR, Fay, JK1, JKn) from 4 estimators to 11. New compute_replicate_refit_variance() in survey.py generalizes the existing compute_replicate_vcov() for estimators whose estimation procedure depends on weights (within-transformation, two-stage imputation). Estimator-level changes: - DiD: LinearRegression dispatch (no-absorb) + refit loop (absorb) - MultiPeriodDiD: compute_replicate_vcov (no-absorb) + refit loop (absorb) - TWFE: refit loop re-doing within-transformation per replicate - SunAbraham: refit loop wrapping _fit_saturated_regression, replaces vcov_cohort - StackedDiD: refit loop with Q-weight × replicate weight composition - ImputationDiD: two-stage refit (_fit_untreated_model + _impute) - TwoStageDiD: two-stage refit (_fit_untreated_model + _stage2_static) All estimators reject bootstrap + replicate (mutual exclusion). Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/estimators.py | 163 ++++++++++-- diff_diff/imputation.py | 106 +++++--- diff_diff/stacked_did.py | 60 ++++- diff_diff/sun_abraham.py | 43 ++- diff_diff/survey.py | 147 ++++++++++- diff_diff/twfe.py | 66 ++++- diff_diff/two_stage.py | 47 +++- docs/methodology/REGISTRY.md | 22 +- tests/test_replicate_weight_expansion.py | 323 +++++++++++++++++++++++ tests/test_survey_phase6.py | 129 +++++---- 10 files changed, 948 insertions(+), 158 deletions(-) create mode 100644 tests/test_replicate_weight_expansion.py diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index 1b8a2ac07..f9e6fe6e0 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -240,14 +240,14 @@ def fit( resolved_survey, survey_weights, survey_weight_type, survey_metadata = ( _resolve_survey_for_fit(survey_design, data, self.inference) ) - # Reject replicate-weight designs — base DiD uses compute_survey_vcov - # (TSL) directly, not LinearRegression's replicate dispatch. - if resolved_survey is not None and resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "DifferenceInDifferences does not yet support replicate-weight " - "survey designs. Use CallawaySantAnna, EfficientDiD, " - "ContinuousDiD, or TripleDifference for replicate-weight " - "inference, or use a TSL-based survey design (strata/psu/fpc)." + _uses_replicate = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) + if _uses_replicate and self.inference == "wild_bootstrap": + raise ValueError( + "Cannot use inference='wild_bootstrap' with replicate-weight " + "survey designs. Replicate weights provide their own variance " + "estimation." ) # Handle absorbed fixed effects (within-transformation) @@ -358,6 +358,13 @@ def fit( ) survey_metadata = compute_survey_metadata(resolved_survey, raw_w) + # When absorb + replicate: pass survey_design=None to prevent + # LinearRegression from computing replicate vcov on already-demeaned + # data (demeaning depends on weights, so replicate refits must re-demean). + _lr_survey = resolved_survey + if _uses_replicate and absorbed_vars: + _lr_survey = None + reg = LinearRegression( include_intercept=False, # Intercept already in X robust=self.robust, @@ -366,7 +373,7 @@ def fit( rank_deficient_action=self.rank_deficient_action, weights=survey_weights, weight_type=survey_weight_type, - survey_design=resolved_survey, + survey_design=_lr_survey, ).fit(X, y, df_adjustment=n_absorbed_effects) coefficients = reg.coefficients_ @@ -375,14 +382,59 @@ def fit( assert coefficients is not None att = coefficients[att_idx] - # Get inference - either from bootstrap or analytical - if self.inference == "wild_bootstrap" and self.cluster is not None: + # Get inference - replicate absorb override, bootstrap, or analytical + if _uses_replicate and absorbed_vars: + # Estimator-level replicate variance: re-demean + re-solve per replicate + from diff_diff.survey import compute_replicate_refit_variance + from diff_diff.utils import safe_inference + + _absorb_list = list(absorbed_vars) # capture for closure + + def _refit_did_absorb(w_r): + wd = data.copy() + wd["_treat_time"] = ( + wd[treatment].values.astype(float) * wd[time].values.astype(float) + ) + vars_dm = [outcome, treatment, time, "_treat_time"] + (covariates or []) + for ab_var in _absorb_list: + wd, _ = demean_by_group(wd, vars_dm, ab_var, inplace=True, weights=w_r) + y_r = wd[outcome].values.astype(float) + d_r = wd[treatment].values.astype(float) + t_r = wd[time].values.astype(float) + dt_r = wd["_treat_time"].values.astype(float) + X_r = np.column_stack([np.ones(len(y_r)), d_r, t_r, dt_r]) + if covariates: + for cov in covariates: + X_r = np.column_stack([X_r, wd[cov].values.astype(float)]) + coef_r, _, _ = solve_ols( + X_r, y_r, + weights=w_r, weight_type=survey_weight_type, + rank_deficient_action="silent", return_vcov=False, + ) + return coef_r + + vcov, _n_valid_rep = compute_replicate_refit_variance( + _refit_did_absorb, coefficients, resolved_survey + ) + se = float(np.sqrt(max(vcov[att_idx, att_idx], 0.0))) + _df_rep = ( + survey_metadata.df_survey + if survey_metadata and survey_metadata.df_survey + else None + ) + if _n_valid_rep < resolved_survey.n_replicates: + _df_rep = _n_valid_rep - 1 if _n_valid_rep > 1 else 0 + t_stat, p_value, conf_int = safe_inference( + att, se, alpha=self.alpha, df=_df_rep + ) + elif self.inference == "wild_bootstrap" and self.cluster is not None: # Override with wild cluster bootstrap inference se, p_value, conf_int, t_stat, vcov, _ = self._run_wild_bootstrap_inference( X, y, residuals, cluster_ids, att_idx ) else: # Use analytical inference from LinearRegression + # (handles replicate vcov for no-absorb path automatically) vcov = reg.vcov_ inference = reg.get_inference(att_idx) se = inference.se @@ -1017,14 +1069,14 @@ def fit( # type: ignore[override] resolved_survey, survey_weights, survey_weight_type, survey_metadata = ( _resolve_survey_for_fit(survey_design, data, effective_inference) ) - # Reject replicate-weight designs — MultiPeriodDiD uses - # compute_survey_vcov (TSL) directly without replicate dispatch. - if resolved_survey is not None and resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "MultiPeriodDiD does not yet support replicate-weight survey " - "designs. Use CallawaySantAnna for staggered adoption with " - "replicate weights, or use a TSL-based survey design " - "(strata/psu/fpc)." + _uses_replicate_mp = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) + if _uses_replicate_mp and effective_inference == "wild_bootstrap": + raise ValueError( + "Cannot use inference='wild_bootstrap' with replicate-weight " + "survey designs. Replicate weights provide their own variance " + "estimation." ) # Handle absorbed fixed effects (within-transformation) @@ -1177,7 +1229,74 @@ def fit( # type: ignore[override] ) # Compute survey vcov if applicable - if _use_survey_vcov: + _n_valid_rep_mp = None + if _use_survey_vcov and _uses_replicate_mp and absorb: + # Absorb + replicate: estimator-level refit (demeaning depends on weights) + from diff_diff.survey import compute_replicate_refit_variance + + _absorb_list_mp = list(absorb) + + def _refit_mp_absorb(w_r): + wd = data.copy() + d_raw_ = wd[treatment].values.astype(float) + t_raw_ = wd[time].values + wd["_did_treatment"] = d_raw_ + for period_ in non_ref_periods: + wd[f"_did_period_{period_}"] = (t_raw_ == period_).astype(float) + wd[f"_did_interact_{period_}"] = d_raw_ * (t_raw_ == period_).astype(float) + vars_dm_ = ( + [outcome, "_did_treatment"] + + [f"_did_period_{p}" for p in non_ref_periods] + + [f"_did_interact_{p}" for p in non_ref_periods] + + (covariates or []) + ) + for ab_var_ in _absorb_list_mp: + wd, _ = demean_by_group(wd, vars_dm_, ab_var_, inplace=True, weights=w_r) + y_r = wd[outcome].values.astype(float) + d_r = wd["_did_treatment"].values.astype(float) + X_r = np.column_stack([np.ones(len(y_r)), d_r]) + for period_ in non_ref_periods: + X_r = np.column_stack( + [X_r, wd[f"_did_period_{period_}"].values.astype(float)] + ) + for period_ in non_ref_periods: + X_r = np.column_stack( + [X_r, wd[f"_did_interact_{period_}"].values.astype(float)] + ) + if covariates: + for cov_ in covariates: + X_r = np.column_stack([X_r, wd[cov_].values.astype(float)]) + coef_r, _, _ = solve_ols( + X_r, y_r, + weights=w_r, weight_type=survey_weight_type, + rank_deficient_action="silent", return_vcov=False, + ) + return coef_r + + vcov, _n_valid_rep_mp = compute_replicate_refit_variance( + _refit_mp_absorb, coefficients, resolved_survey + ) + elif _use_survey_vcov and _uses_replicate_mp: + # No absorb + replicate: X is fixed, use compute_replicate_vcov directly + from diff_diff.survey import compute_replicate_vcov + + nan_mask = np.isnan(coefficients) + if np.any(nan_mask): + kept_cols = np.where(~nan_mask)[0] + if len(kept_cols) > 0: + vcov_reduced, _n_valid_rep_mp = compute_replicate_vcov( + X[:, kept_cols], y, coefficients[kept_cols], resolved_survey, + weight_type=survey_weight_type, + ) + vcov = _expand_vcov_with_nan(vcov_reduced, X.shape[1], kept_cols) + else: + vcov = np.full((X.shape[1], X.shape[1]), np.nan) + _n_valid_rep_mp = 0 + else: + vcov, _n_valid_rep_mp = compute_replicate_vcov( + X, y, coefficients, resolved_survey, weight_type=survey_weight_type, + ) + elif _use_survey_vcov: from diff_diff.survey import compute_survey_vcov nan_mask = np.isnan(coefficients) @@ -1201,6 +1320,10 @@ def fit( # type: ignore[override] df = n_eff_df - k_effective - n_absorbed_effects if resolved_survey is not None and resolved_survey.df_survey is not None: df = resolved_survey.df_survey + # Override df when replicate replicates were dropped + if _n_valid_rep_mp is not None and resolved_survey is not None: + if _n_valid_rep_mp < resolved_survey.n_replicates: + df = _n_valid_rep_mp - 1 if _n_valid_rep_mp > 1 else 0 # Guard: fall back to normal distribution if df is non-positive if df is not None and df <= 0: diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 4f32a6638..649b96698 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -244,13 +244,16 @@ def fit( survey_design, data, "analytical" ) + _uses_replicate_imp = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) + if _uses_replicate_imp and self.n_bootstrap > 0: + raise ValueError( + "Cannot use n_bootstrap > 0 with replicate-weight survey designs. " + "Replicate weights provide their own variance estimation." + ) # Validate within-unit constancy for panel survey designs if resolved_survey is not None: - if resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "ImputationDiD does not yet support replicate-weight survey " - "designs. Use a TSL-based survey design (strata/psu/fpc)." - ) _validate_unit_constant_survey(data, unit, survey_design) if resolved_survey.weight_type != "pweight": raise ValueError( @@ -463,43 +466,72 @@ def fit( else: overall_att = float(np.mean(valid_tau)) - # ---- Conservative variance (Theorem 3) ---- - # Build weights matching the ATT: proportional to survey weights for - # finite tau_hat, uniform when no survey - overall_weights = np.zeros(n_omega_1) - n_valid = int(finite_mask.sum()) - if n_valid > 0: - if survey_weights is not None: - treated_sw = survey_weights[omega_1_mask.values] - sw_finite = treated_sw[finite_mask] - overall_weights[finite_mask] = sw_finite / sw_finite.sum() - else: - overall_weights[finite_mask] = 1.0 / n_valid + # ---- Variance ---- + _n_valid_rep_imp = None + if _uses_replicate_imp: + # Replicate variance: re-run two-stage procedure per replicate + from diff_diff.survey import compute_replicate_refit_variance - if n_valid == 0: - overall_se = np.nan - else: - overall_se = self._compute_conservative_variance( - df=df, - outcome=outcome, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - weights=overall_weights, - cluster_var=cluster_var, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, + def _refit_imp(w_r): + ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( + df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, + ) + tau_r, _ = self._impute_treatment_effects( + df, outcome, unit, time, covariates, omega_1_mask, + ufe_r, tfe_r, gm_r, delta_r, + ) + fin = np.isfinite(tau_r) + if not np.any(fin): + return np.array([np.nan]) + tw = w_r[omega_1_mask.values][fin] + tw_sum = np.sum(tw) + if tw_sum == 0: + return np.array([np.nan]) + return np.array([float(np.sum(tau_r[fin] * tw) / tw_sum)]) + + _vcov_rep_imp, _n_valid_rep_imp = compute_replicate_refit_variance( + _refit_imp, np.array([overall_att]), resolved_survey ) + overall_se = float(np.sqrt(max(_vcov_rep_imp[0, 0], 0.0))) + else: + # Conservative variance (Theorem 3) + overall_weights = np.zeros(n_omega_1) + n_valid = int(finite_mask.sum()) + if n_valid > 0: + if survey_weights is not None: + treated_sw = survey_weights[omega_1_mask.values] + sw_finite = treated_sw[finite_mask] + overall_weights[finite_mask] = sw_finite / sw_finite.sum() + else: + overall_weights[finite_mask] = 1.0 / n_valid + + if n_valid == 0: + overall_se = np.nan + else: + overall_se = self._compute_conservative_variance( + df=df, + outcome=outcome, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=unit_fe, + time_fe=time_fe, + grand_mean=grand_mean, + delta_hat=delta_hat, + weights=overall_weights, + cluster_var=cluster_var, + kept_cov_mask=kept_cov_mask, + survey_weights=survey_weights, + ) # Survey degrees of freedom for t-distribution inference _survey_df = resolved_survey.df_survey if resolved_survey is not None else None + if _n_valid_rep_imp is not None and resolved_survey is not None: + if _n_valid_rep_imp < resolved_survey.n_replicates: + _survey_df = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0 overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=_survey_df diff --git a/diff_diff/stacked_did.py b/diff_diff/stacked_did.py index d923c7071..eb4f76e45 100644 --- a/diff_diff/stacked_did.py +++ b/diff_diff/stacked_did.py @@ -242,15 +242,9 @@ def fit( resolved_survey, survey_weights, survey_weight_type, survey_metadata = ( _resolve_survey_for_fit(survey_design, data, "analytical") ) - # Reject replicate-weight designs — StackedDiD uses - # compute_survey_vcov (TSL) directly without replicate dispatch. - if resolved_survey is not None and resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "StackedDiD does not yet support replicate-weight survey " - "designs. Use CallawaySantAnna for staggered adoption with " - "replicate weights, or use a TSL-based survey design " - "(strata/psu/fpc)." - ) + _uses_replicate_sd = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) # Reject fweight and aweight — Q-weight composition is ratio-valued # and breaks both frequency-weight (integer) and analytic-weight @@ -273,6 +267,9 @@ def fit( col_name = getattr(survey_design, attr, None) if col_name is not None: survey_cols.append(col_name) + # Propagate replicate weight columns through stacked dataset + if survey_design.replicate_weights is not None: + survey_cols.extend(survey_design.replicate_weights) df = data.copy() df[time] = pd.to_numeric(df[time]) @@ -428,8 +425,42 @@ def fit( ) assert vcov is not None - # ---- Survey VCV override (TSL variance) ---- - if resolved_survey is not None: + # ---- Survey VCV override ---- + _n_valid_rep_sd = None + resolved_stacked = None + if resolved_survey is not None and _uses_replicate_sd: + # Replicate variance: re-run WLS per replicate with composed weights + from diff_diff.survey import compute_replicate_refit_variance, compute_survey_metadata + + resolved_stacked = survey_design.resolve(stacked_df) + + # Refit closure: compose Q-weights with replicate survey weights + def _refit_stacked(w_r): + composed_r = Q_weights * w_r + w_sum = np.sum(composed_r) + if w_sum > 0: + composed_r = composed_r * (n_stacked / w_sum) + sqrt_w_r = np.sqrt(composed_r) + coef_r, _, _ = solve_ols( + X * sqrt_w_r[:, np.newaxis], Y * sqrt_w_r, + cluster_ids=cluster_ids, + rank_deficient_action="silent", return_vcov=False, + ) + return coef_r + + # Full-sample cohort effect vector + vcov, _n_valid_rep_sd = compute_replicate_refit_variance( + _refit_stacked, coef, resolved_stacked + ) + + # Compute survey metadata + raw_w_stacked = ( + stacked_df[survey_design.weights].values.astype(np.float64) + if survey_design.weights is not None + else np.ones(n_stacked, dtype=np.float64) + ) + survey_metadata = compute_survey_metadata(resolved_stacked, raw_w_stacked) + elif resolved_survey is not None: from diff_diff.survey import ( _inject_cluster_as_psu, _resolve_effective_cluster, @@ -487,6 +518,10 @@ def fit( if survey_metadata is not None and survey_metadata.df_survey is not None else None ) + # Override df when replicate replicates were dropped + if _n_valid_rep_sd is not None and resolved_stacked is not None: + if _n_valid_rep_sd < resolved_stacked.n_replicates: + _survey_df = _n_valid_rep_sd - 1 if _n_valid_rep_sd > 1 else 0 t_stat, p_value, conf_int = safe_inference( effect, se, alpha=self.alpha, df=_survey_df ) @@ -524,6 +559,9 @@ def fit( if survey_metadata is not None and survey_metadata.df_survey is not None else None ) + if _n_valid_rep_sd is not None and resolved_stacked is not None: + if _n_valid_rep_sd < resolved_stacked.n_replicates: + _survey_df_overall = _n_valid_rep_sd - 1 if _n_valid_rep_sd > 1 else 0 overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=_survey_df_overall ) diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index 8ab8a9162..1e8efb631 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -501,16 +501,13 @@ def fit( if resolved_survey is not None: _validate_unit_constant_survey(data, unit, survey_design) - # Reject replicate-weight designs — SunAbraham's weighted - # within-transformation bakes survey weights into X and y, so - # replicate refits on the already-transformed design are incorrect. - # Full estimator-level replicate refits are not yet implemented. - if resolved_survey is not None and resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "SunAbraham does not yet support replicate-weight survey designs. " - "The weighted within-transformation must be recomputed for each " - "replicate, which requires estimator-level replicate refits. " - "Use a TSL-based survey design (strata/psu/fpc) instead." + _uses_replicate_sa = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) + if _uses_replicate_sa and self.n_bootstrap > 0: + raise ValueError( + "Cannot use n_bootstrap > 0 with replicate-weight survey designs. " + "Replicate weights provide their own variance estimation." ) # Bootstrap + survey supported via Rao-Wu rescaled bootstrap. @@ -633,6 +630,30 @@ def fit( resolved_survey=resolved_survey, ) + # Replicate variance override: re-run saturated regression per replicate + # and replace vcov_cohort with the replicate version. The downstream + # delta-method aggregation (_compute_iw_effects, _compute_overall_att) + # then automatically produces correct replicate-based SEs. + if _uses_replicate_sa: + from diff_diff.survey import compute_replicate_refit_variance + + _keys_ordered = sorted(coef_index_map.keys(), key=lambda k: coef_index_map[k]) + _full_cohort_vec = np.array([cohort_effects.get(k, np.nan) for k in _keys_ordered]) + + def _refit_sa(w_r): + ce_r, _, _, cim_r = self._fit_saturated_regression( + df_reg, outcome, unit, time, first_treat, + treatment_groups, rel_periods_to_estimate, covariates, + cluster_var, survey_weights=w_r, + survey_weight_type=survey_weight_type, + resolved_survey=None, # prevent internal replicate dispatch + ) + return np.array([ce_r.get(k, np.nan) for k in _keys_ordered]) + + vcov_cohort, _n_valid_rep_sa = compute_replicate_refit_variance( + _refit_sa, _full_cohort_vec, resolved_survey + ) + # Resolve survey weight column name for cohort aggregation survey_weight_col = ( survey_design.weights @@ -648,6 +669,8 @@ def fit( if survey_metadata is not None and survey_metadata.df_survey is not None else None ) + if _uses_replicate_sa and _n_valid_rep_sa < resolved_survey.n_replicates: + _sa_survey_df = _n_valid_rep_sa - 1 if _n_valid_rep_sa > 1 else 0 # Compute interaction-weighted event study effects event_study_effects, cohort_weights = self._compute_iw_effects( diff --git a/diff_diff/survey.py b/diff_diff/survey.py index 9fa2be457..bd18700e2 100644 --- a/diff_diff/survey.py +++ b/diff_diff/survey.py @@ -16,7 +16,7 @@ import warnings from dataclasses import dataclass, field, replace -from typing import List, Optional, Tuple +from typing import Callable, List, Optional, Tuple import numpy as np import pandas as pd @@ -1703,6 +1703,151 @@ def compute_replicate_if_variance( raise ValueError(f"Unknown replicate method: {method}") +def compute_replicate_refit_variance( + refit_fn: Callable[[np.ndarray], np.ndarray], + full_sample_estimate: np.ndarray, + resolved: "ResolvedSurveyDesign", +) -> Tuple[np.ndarray, int]: + """Compute replicate variance by re-running an arbitrary estimation function. + + For each replicate weight column, calls ``refit_fn(w_r)`` and collects + the resulting estimate vector. Variance is computed from the distribution + of replicate estimates using method-specific scaling. + + This generalises :func:`compute_replicate_vcov` (which hard-codes + ``solve_ols`` as the refit) for estimators whose estimation procedure + is more complex than a single OLS call (e.g. within-transformation, + two-stage imputation, stacked regression). + + Parameters + ---------- + refit_fn : callable + ``(n,) weight array -> (k,) estimate array``. Must return the same + length *k* on every call. Should return all-NaN when the estimation + fails for that replicate. + full_sample_estimate : np.ndarray + Estimate vector from the full-sample weights, shape ``(k,)``. + resolved : ResolvedSurveyDesign + Must have ``uses_replicate_variance == True``. + + Returns + ------- + tuple of (np.ndarray, int) + ``(vcov, n_valid)`` where *vcov* has shape ``(k, k)`` and *n_valid* + is the number of replicates that produced finite estimates. + """ + full_sample_estimate = np.asarray(full_sample_estimate, dtype=np.float64).ravel() + k = len(full_sample_estimate) + rep_weights = resolved.replicate_weights + method = resolved.replicate_method + R = resolved.n_replicates + + # Collect replicate estimate vectors + est_reps = np.full((R, k), np.nan) + for r in range(R): + w_r = rep_weights[:, r].copy() + if not resolved.combined_weights: + w_r = w_r * resolved.weights + if np.sum(w_r) == 0: + continue + try: + est_r = refit_fn(w_r) + est_r = np.asarray(est_r, dtype=np.float64).ravel() + if len(est_r) == k: + est_reps[r] = est_r + except (np.linalg.LinAlgError, ValueError, RuntimeError): + pass # NaN row for failed replicate + + # Remove replicates with NaN estimates + valid = np.all(np.isfinite(est_reps), axis=1) + n_invalid = int(R - np.sum(valid)) + if n_invalid > 0: + warnings.warn( + f"{n_invalid} of {R} replicate refits failed. " + f"Variance computed from {int(np.sum(valid))} valid replicates.", + UserWarning, + stacklevel=2, + ) + n_valid = int(np.sum(valid)) + if n_valid < 2: + if n_valid == 0: + warnings.warn( + "All replicate refits failed. Returning NaN variance.", + UserWarning, + stacklevel=2, + ) + else: + warnings.warn( + f"Only {n_valid} valid replicate(s) — variance is not estimable " + f"with fewer than 2. Returning NaN.", + UserWarning, + stacklevel=2, + ) + return np.full((k, k), np.nan), n_valid + + est_valid = est_reps[valid] + c = full_sample_estimate + + # --- Centering (mse flag) --- + if resolved.mse: + center = c + else: + if resolved.replicate_rscales is not None: + pos_scale = resolved.replicate_rscales[valid] > 0 + if np.any(pos_scale): + center = np.mean(est_valid[pos_scale], axis=0) + else: + center = np.mean(est_valid, axis=0) + else: + center = np.mean(est_valid, axis=0) + diffs = est_valid - center[np.newaxis, :] + + outer_sum = diffs.T @ diffs # (k, k) + + # --- Method-specific scaling --- + # BRR/Fay: fixed scaling, ignore user-supplied scale/rscales + if method in ("BRR", "Fay"): + if resolved.replicate_scale is not None or resolved.replicate_rscales is not None: + warnings.warn( + f"Custom replicate_scale/replicate_rscales ignored for {method} " + f"(BRR/Fay use fixed scaling).", + UserWarning, + stacklevel=2, + ) + factor = _replicate_variance_factor(method, R, resolved.fay_rho) + return factor * outer_sum, n_valid + + # JK1/JKn: apply scale * rscales multiplicatively + scale = resolved.replicate_scale if resolved.replicate_scale is not None else 1.0 + + if resolved.replicate_rscales is not None: + valid_rscales = resolved.replicate_rscales[valid] + V = np.zeros((k, k)) + for i in range(len(diffs)): + V += valid_rscales[i] * np.outer(diffs[i], diffs[i]) + return scale * V, n_valid + + if method == "JK1": + factor = _replicate_variance_factor(method, R, resolved.fay_rho) + return scale * factor * outer_sum, n_valid + elif method == "JKn": + rep_strata = resolved.replicate_strata + if rep_strata is None: + raise ValueError("JKn requires replicate_strata") + valid_strata = rep_strata[valid] + V = np.zeros((k, k)) + for h in np.unique(rep_strata): + n_h_original = int(np.sum(rep_strata == h)) + mask_h = valid_strata == h + if not np.any(mask_h): + continue + diffs_h = diffs[mask_h] + V += ((n_h_original - 1.0) / n_h_original) * (diffs_h.T @ diffs_h) + return scale * V, n_valid + else: + raise ValueError(f"Unknown replicate method: {method}") + + def aggregate_to_psu( values: np.ndarray, resolved: "ResolvedSurveyDesign", diff --git a/diff_diff/twfe.py b/diff_diff/twfe.py index 0d5815895..4ded0f764 100644 --- a/diff_diff/twfe.py +++ b/diff_diff/twfe.py @@ -127,15 +127,14 @@ def fit( # type: ignore[override] resolved_survey, survey_weights, survey_weight_type, survey_metadata = ( _resolve_survey_for_fit(survey_design, data, self.inference) ) - # Reject replicate-weight designs — TWFE within-transformation must - # be recomputed per replicate (same reason as SunAbraham rejection) - if resolved_survey is not None and resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "TwoWayFixedEffects does not yet support replicate-weight " - "survey designs. The weighted within-transformation must be " - "recomputed for each replicate. Use CallawaySantAnna or " - "TripleDifference for replicate-weight inference, or use a " - "TSL-based survey design (strata/psu/fpc)." + _uses_replicate_twfe = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) + if _uses_replicate_twfe and self.inference == "wild_bootstrap": + raise ValueError( + "Cannot use inference='wild_bootstrap' with replicate-weight " + "survey designs. Replicate weights provide their own variance " + "estimation." ) # Use unit-level clustering if not specified (use local variable to avoid mutation) @@ -210,6 +209,10 @@ def fit( # type: ignore[override] # If "error", let LinearRegression raise immediately # If "warn" or "silent", suppress generic warning and use TWFE's context-specific # error/warning messages (more informative for panel data) + # For replicate designs: pass survey_design=None to prevent LinearRegression + # from computing replicate vcov on already-demeaned data (demeaning depends + # on weights, so replicate refits must re-demean at the estimator level). + _lr_survey_twfe = None if _uses_replicate_twfe else resolved_survey if self.rank_deficient_action == "error": reg = LinearRegression( include_intercept=False, @@ -219,7 +222,7 @@ def fit( # type: ignore[override] rank_deficient_action="error", weights=survey_weights, weight_type=survey_weight_type, - survey_design=resolved_survey, + survey_design=_lr_survey_twfe, ).fit(X, y, df_adjustment=df_adjustment) else: # Suppress generic warning, TWFE provides context-specific messages below @@ -235,7 +238,7 @@ def fit( # type: ignore[override] rank_deficient_action="silent", weights=survey_weights, weight_type=survey_weight_type, - survey_design=resolved_survey, + survey_design=_lr_survey_twfe, ).fit(X, y, df_adjustment=df_adjustment) coefficients = reg.coefficients_ @@ -281,8 +284,45 @@ def fit( # type: ignore[override] stacklevel=2, ) - # Get inference - either from bootstrap or analytical - if self.inference == "wild_bootstrap": + # Get inference - replicate, bootstrap, or analytical + if _uses_replicate_twfe: + # Estimator-level replicate variance: re-do within-transform per replicate + from diff_diff.linalg import solve_ols + from diff_diff.survey import compute_replicate_refit_variance + from diff_diff.utils import safe_inference as _safe_inf + + _all_vars_twfe = list(all_vars) + _covariates_twfe = list(covariates) if covariates else [] + + def _refit_twfe(w_r): + data_dem_r = _within_transform_util( + data, _all_vars_twfe, unit, time, suffix="_demeaned", weights=w_r, + ) + y_r = data_dem_r[f"{outcome}_demeaned"].values + X_list_r = [data_dem_r["_treatment_post_demeaned"].values] + for cov_ in _covariates_twfe: + X_list_r.append(data_dem_r[f"{cov_}_demeaned"].values) + X_r = np.column_stack([np.ones(len(y_r))] + X_list_r) + coef_r, _, _ = solve_ols( + X_r, y_r, + weights=w_r, weight_type=survey_weight_type, + rank_deficient_action="silent", return_vcov=False, + ) + return coef_r + + vcov, _n_valid_rep_twfe = compute_replicate_refit_variance( + _refit_twfe, coefficients, resolved_survey + ) + se = float(np.sqrt(max(vcov[att_idx, att_idx], 0.0))) + _df_rep = ( + survey_metadata.df_survey + if survey_metadata and survey_metadata.df_survey + else None + ) + if _n_valid_rep_twfe < resolved_survey.n_replicates: + _df_rep = _n_valid_rep_twfe - 1 if _n_valid_rep_twfe > 1 else 0 + t_stat, p_value, conf_int = _safe_inf(att, se, alpha=self.alpha, df=_df_rep) + elif self.inference == "wild_bootstrap": # Override with wild cluster bootstrap inference se, p_value, conf_int, t_stat, vcov, _ = self._run_wild_bootstrap_inference( X, y, residuals, cluster_ids, att_idx diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index cd4baa4e1..42be3bc9d 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -240,13 +240,16 @@ def fit( _resolve_survey_for_fit(survey_design, data, "analytical") ) + _uses_replicate_ts = ( + resolved_survey is not None and resolved_survey.uses_replicate_variance + ) + if _uses_replicate_ts and self.n_bootstrap > 0: + raise ValueError( + "Cannot use n_bootstrap > 0 with replicate-weight survey designs. " + "Replicate weights provide their own variance estimation." + ) # Validate within-unit constancy for panel survey designs if resolved_survey is not None: - if resolved_survey.uses_replicate_variance: - raise NotImplementedError( - "TwoStageDiD does not yet support replicate-weight survey " - "designs. Use a TSL-based survey design (strata/psu/fpc)." - ) _validate_unit_constant_survey(data, unit, survey_design) if resolved_survey.weight_type != "pweight": raise ValueError( @@ -495,6 +498,40 @@ def fit( survey_weight_type=survey_weight_type, ) + # Replicate variance override: re-run both stages per replicate + _n_valid_rep_ts = None + if _uses_replicate_ts: + from diff_diff.survey import compute_replicate_refit_variance + + def _refit_ts(w_r): + ufe_r, tfe_r, gm_r, delta_r, kcm_r = self._fit_untreated_model( + df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, + ) + y_tilde_r = self._residualize( + df, outcome, unit, time, covariates, + ufe_r, tfe_r, gm_r, delta_r, + ) + df_tmp = df.copy() + df_tmp["_y_tilde"] = y_tilde_r + att_r, _ = self._stage2_static( + df=df_tmp, unit=unit, time=time, first_treat=first_treat, + covariates=covariates, omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, unit_fe=ufe_r, time_fe=tfe_r, + grand_mean=gm_r, delta_hat=delta_r, cluster_var=cluster_var, + kept_cov_mask=kcm_r, survey_weights=w_r, + survey_weight_type="pweight", + ) + return np.array([att_r]) + + _vcov_rep_ts, _n_valid_rep_ts = compute_replicate_refit_variance( + _refit_ts, np.array([overall_att]), resolved_survey + ) + overall_se = float(np.sqrt(max(_vcov_rep_ts[0, 0], 0.0))) + + if _n_valid_rep_ts is not None and resolved_survey is not None: + if _n_valid_rep_ts < resolved_survey.n_replicates: + _survey_df = _n_valid_rep_ts - 1 if _n_valid_rep_ts > 1 else 0 + overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=_survey_df ) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index a39822395..d445db2f9 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -2245,16 +2245,18 @@ variance from the distribution of replicate estimates. t-based inference. - **Note:** Replicate-weight support matrix: - **Supported**: CallawaySantAnna (reg/ipw/dr without covariates, no - bootstrap), ContinuousDiD - (no bootstrap), EfficientDiD (no bootstrap), TripleDifference (all - methods), LinearRegression (OLS path) - - **Rejected with NotImplementedError**: SunAbraham, TwoWayFixedEffects - (within-transformation must be recomputed per replicate), - DifferenceInDifferences, MultiPeriodDiD, StackedDiD (use - compute_survey_vcov directly), ImputationDiD, TwoStageDiD (custom - variance), SyntheticDiD, TROP (bootstrap-based variance), - BaconDecomposition (diagnostic only) - - CS/ContinuousDiD/EfficientDiD reject replicate + `n_bootstrap > 0` + bootstrap), ContinuousDiD (no bootstrap), EfficientDiD (no bootstrap), + TripleDifference (all methods), LinearRegression (OLS path), + DifferenceInDifferences (no-absorb via LinearRegression dispatch, + absorb via estimator-level refit), MultiPeriodDiD (no-absorb via + `compute_replicate_vcov`, absorb via estimator-level refit), + TwoWayFixedEffects (estimator-level refit with within-transformation), + SunAbraham (estimator-level refit, replaces `vcov_cohort`), + StackedDiD (estimator-level refit with Q-weight composition), + ImputationDiD (two-stage refit), TwoStageDiD (two-stage refit) + - **Rejected with NotImplementedError**: SyntheticDiD, TROP + (bootstrap-based variance), BaconDecomposition (diagnostic only) + - Estimators with replicate support reject replicate + bootstrap (replicate weights provide analytical variance) - **Note:** When invalid replicates are dropped in `compute_replicate_vcov` (OLS path), `n_valid` is returned and used for `df_survey = n_valid - 1` diff --git a/tests/test_replicate_weight_expansion.py b/tests/test_replicate_weight_expansion.py new file mode 100644 index 000000000..1fb904aa8 --- /dev/null +++ b/tests/test_replicate_weight_expansion.py @@ -0,0 +1,323 @@ +"""Tests for replicate weight support expansion to 7 additional estimators.""" + +import numpy as np +import pandas as pd +import pytest + +from diff_diff import ( + DifferenceInDifferences, + ImputationDiD, + MultiPeriodDiD, + StackedDiD, + SunAbraham, + TwoStageDiD, + TwoWayFixedEffects, +) +from diff_diff.survey import SurveyDesign + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +def _make_simple_panel(): + """2-period panel for DiD/MultiPeriodDiD (treatment/post binary columns).""" + np.random.seed(123) + n_units = 40 + rows = [] + for i in range(n_units): + treated = 1 if i < 20 else 0 + wt = 1.0 + 0.2 * (i % 5) + for t in [0, 1]: + y = 5.0 + 0.5 * treated + 1.0 * t + if treated and t == 1: + y += 2.0 # ATT = 2 + y += np.random.normal(0, 0.3) + rows.append({ + "unit": i, "time": t, "treated": treated, "post": t, + "outcome": y, "weight": wt, + }) + data = pd.DataFrame(rows) + return data + + +def _make_staggered_panel(): + """Multi-period staggered panel for TWFE/SA/Stacked/Imputation/TwoStage.""" + np.random.seed(456) + n_units, n_periods = 50, 8 + rows = [] + for i in range(n_units): + if i < 15: + ft = 4 # cohort 1 + elif i < 30: + ft = 6 # cohort 2 + else: + ft = 0 # never-treated + wt = 1.0 + 0.3 * (i % 5) + for t in range(1, n_periods + 1): + y = 10.0 + i * 0.03 + t * 0.2 + if ft > 0 and t >= ft: + y += 2.0 + y += np.random.normal(0, 0.4) + rows.append({ + "unit": i, "time": t, "first_treat": ft, + "outcome": y, "weight": wt, + "treated": 1 if ft > 0 else 0, + "post": 1 if ft > 0 and t >= ft else 0, + }) + data = pd.DataFrame(rows) + return data + + +def _add_jk1_replicates(data, n_rep=15, unit_col="unit"): + """Add JK1 (delete-cluster jackknife) replicate weight columns.""" + units = sorted(data[unit_col].unique()) + cluster_size = max(1, len(units) // n_rep) + rep_cols = [] + for r in range(n_rep): + start = r * cluster_size + end = min((r + 1) * cluster_size, len(units)) + deleted_units = set(units[start:end]) + w_r = data["weight"].values.copy() + mask = data[unit_col].isin(deleted_units).values + w_r[mask] = 0.0 + w_r[~mask] *= n_rep / (n_rep - 1) + col = f"rep_{r}" + data[col] = w_r + rep_cols.append(col) + return rep_cols + + +def _add_brr_replicates(data, n_rep=16, unit_col="unit"): + """Add BRR replicate weight columns (random sign perturbation).""" + rng = np.random.RandomState(789) + units = sorted(data[unit_col].unique()) + rep_cols = [] + for r in range(n_rep): + signs = rng.choice([-1, 1], size=len(units)) + sign_map = dict(zip(units, signs)) + perturbation = data[unit_col].map(sign_map).values.astype(float) + # BRR: w_r = w * (1 + epsilon) / 2, where epsilon in {-1, 1} + # Simplified: w_r = w * (1 + perturbation) (combined_weights=True style) + w_r = data["weight"].values * (1.0 + 0.5 * perturbation) + w_r = np.maximum(w_r, 0.0) + col = f"brr_{r}" + data[col] = w_r + rep_cols.append(col) + return rep_cols + + +# --------------------------------------------------------------------------- +# Smoke tests — each estimator × {JK1, BRR} +# --------------------------------------------------------------------------- + +class TestDiDReplicate: + """DifferenceInDifferences with replicate weights.""" + + def test_did_jk1(self): + data = _make_simple_panel() + rep_cols = _add_jk1_replicates(data, n_rep=10) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = DifferenceInDifferences().fit( + data, "outcome", "treated", "post", survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_did_brr(self): + data = _make_simple_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = DifferenceInDifferences().fit( + data, "outcome", "treated", "post", survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 + + def test_did_wild_bootstrap_rejected(self): + """Wild bootstrap + survey is rejected before replicate check.""" + data = _make_simple_panel() + rep_cols = _add_jk1_replicates(data, n_rep=10) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + with pytest.raises((ValueError, NotImplementedError)): + DifferenceInDifferences(inference="wild_bootstrap", cluster="unit").fit( + data, "outcome", "treated", "post", survey_design=sd, + ) + + +class TestMultiPeriodDiDReplicate: + """MultiPeriodDiD with replicate weights.""" + + def test_multiperiod_jk1(self): + data = _make_simple_panel() + rep_cols = _add_jk1_replicates(data, n_rep=10) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = MultiPeriodDiD().fit( + data, "outcome", "treated", "time", post_periods=[1], survey_design=sd, + ) + assert np.isfinite(result.avg_att) + assert np.isfinite(result.avg_se) and result.avg_se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_multiperiod_brr(self): + data = _make_simple_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = MultiPeriodDiD().fit( + data, "outcome", "treated", "time", post_periods=[1], survey_design=sd, + ) + assert np.isfinite(result.avg_att) + assert np.isfinite(result.avg_se) and result.avg_se > 0 + + +class TestTWFEReplicate: + """TwoWayFixedEffects with replicate weights.""" + + @staticmethod + def _make_twfe_panel(): + """Balanced 2-period panel with variation in treatment timing.""" + np.random.seed(321) + n_units = 40 + rows = [] + for i in range(n_units): + treated = 1 if i < 20 else 0 + wt = 1.0 + 0.2 * (i % 5) + for t in [0, 1]: + y = 5.0 + i * 0.05 + t * 1.0 + if treated and t == 1: + y += 2.0 + y += np.random.normal(0, 0.3) + rows.append({ + "unit": i, "time": t, "treated": treated, + "post": t, "outcome": y, "weight": wt, + }) + return pd.DataFrame(rows) + + def test_twfe_brr(self): + """BRR works well with TWFE: perturbation doesn't zero out units.""" + data = self._make_twfe_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = TwoWayFixedEffects().fit( + data, "outcome", "treated", "post", "unit", survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_twfe_brr_larger(self): + """Second BRR test with different seed.""" + data = self._make_twfe_panel() + rep_cols = _add_brr_replicates(data, n_rep=20) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = TwoWayFixedEffects().fit( + data, "outcome", "treated", "post", "unit", survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 + + +class TestSunAbrahamReplicate: + """SunAbraham with replicate weights.""" + + def test_sun_abraham_brr(self): + """BRR replicates are less aggressive than JK1 for SunAbraham.""" + data = _make_staggered_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = SunAbraham(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) + assert np.isfinite(result.overall_se) and result.overall_se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_sun_abraham_bootstrap_rejected(self): + data = _make_staggered_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + with pytest.raises(ValueError, match="n_bootstrap"): + SunAbraham(n_bootstrap=100).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + + +class TestStackedDiDReplicate: + """StackedDiD with replicate weights.""" + + def test_stacked_jk1(self): + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = StackedDiD().fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) + assert np.isfinite(result.overall_se) and result.overall_se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_stacked_brr(self): + data = _make_staggered_panel() + rep_cols = _add_brr_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = StackedDiD().fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) + assert np.isfinite(result.overall_se) and result.overall_se > 0 + + +class TestImputationDiDReplicate: + """ImputationDiD with replicate weights.""" + + def test_imputation_jk1(self): + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = ImputationDiD(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) + assert np.isfinite(result.overall_se) and result.overall_se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_imputation_bootstrap_rejected(self): + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + with pytest.raises(ValueError, match="n_bootstrap"): + ImputationDiD(n_bootstrap=100).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + + +class TestTwoStageDiDReplicate: + """TwoStageDiD with replicate weights.""" + + def test_two_stage_jk1(self): + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = TwoStageDiD(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) + assert np.isfinite(result.overall_se) and result.overall_se > 0 + assert result.survey_metadata is not None + result.summary() + + def test_two_stage_bootstrap_rejected(self): + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + with pytest.raises(ValueError, match="n_bootstrap"): + TwoStageDiD(n_bootstrap=100).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) diff --git a/tests/test_survey_phase6.py b/tests/test_survey_phase6.py index bb54ada36..1ebe589d5 100644 --- a/tests/test_survey_phase6.py +++ b/tests/test_survey_phase6.py @@ -498,8 +498,8 @@ def test_replicate_metadata(self, replicate_data): assert sm.n_replicates == len(rep_cols) assert sm.df_survey == len(rep_cols) - 1 - def test_replicate_rejected_by_base_did(self, replicate_data): - """DifferenceInDifferences rejects replicate-weight designs.""" + def test_replicate_accepted_by_base_did(self, replicate_data): + """DifferenceInDifferences now supports replicate-weight designs.""" data, rep_cols = replicate_data n = len(data) data["treated"] = (np.arange(n) < n // 2).astype(int) @@ -510,11 +510,12 @@ def test_replicate_rejected_by_base_did(self, replicate_data): weights="weight", replicate_weights=rep_cols, replicate_method="JK1", ) - with pytest.raises(NotImplementedError, match="DifferenceInDifferences"): - DifferenceInDifferences().fit( - data, outcome="outcome", treatment="treated", time="post", - survey_design=sd, - ) + result = DifferenceInDifferences().fit( + data, outcome="outcome", treatment="treated", time="post", + survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 def test_replicate_if_variance(self, replicate_data): """IF-based replicate variance produces finite results.""" @@ -1284,8 +1285,8 @@ def test_triple_diff_replicate_all_methods(self, est_method): assert np.isfinite(result.se) assert result.survey_metadata is not None - def test_sun_abraham_replicate_rejected(self): - """SunAbraham should reject replicate-weight survey designs.""" + def test_sun_abraham_replicate_accepted(self): + """SunAbraham now supports replicate-weight survey designs.""" from diff_diff import SunAbraham data, rep_cols = self._make_staggered_replicate_data() @@ -1293,11 +1294,11 @@ def test_sun_abraham_replicate_rejected(self): weights="weight", replicate_weights=rep_cols, replicate_method="JK1", ) - with pytest.raises(NotImplementedError, match="SunAbraham"): - SunAbraham(n_bootstrap=0).fit( - data, "outcome", "unit", "time", "first_treat", - survey_design=sd, - ) + result = SunAbraham(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", + survey_design=sd, + ) + assert np.isfinite(result.overall_att) class TestSubpopulationMaskValidation: @@ -1357,23 +1358,49 @@ def test_callaway_santanna_replicate_bootstrap_rejected(self): survey_design=sd, ) - def test_twfe_replicate_rejected(self): - """TwoWayFixedEffects should reject replicate-weight designs.""" + def test_twfe_replicate_accepted(self): + """TwoWayFixedEffects now supports replicate-weight designs.""" from diff_diff.twfe import TwoWayFixedEffects - data, rep_cols = TestEstimatorReplicateWeights._make_staggered_replicate_data() + # Create a simple 2-period panel that works with within-transformation + np.random.seed(77) + rows = [] + for i in range(40): + treated = 1 if i < 20 else 0 + wt = 1.0 + 0.2 * (i % 5) + for t in [0, 1]: + y = 5.0 + i * 0.05 + t * 1.0 + if treated and t == 1: + y += 2.0 + y += np.random.normal(0, 0.3) + rows.append({ + "unit": i, "time": t, "treated": treated, + "post": t, "outcome": y, "weight": wt, + }) + data = pd.DataFrame(rows) + # Use BRR to avoid zero-weight units breaking within-transform + rng = np.random.RandomState(42) + units = sorted(data["unit"].unique()) + rep_cols = [] + for r in range(16): + signs = rng.choice([-1, 1], size=len(units)) + sm = dict(zip(units, signs)) + p = data["unit"].map(sm).values.astype(float) + data[f"brr_{r}"] = np.maximum(data["weight"].values * (1.0 + 0.5 * p), 0.0) + rep_cols.append(f"brr_{r}") sd = SurveyDesign( weights="weight", replicate_weights=rep_cols, - replicate_method="JK1", + replicate_method="BRR", ) - with pytest.raises(NotImplementedError, match="TwoWayFixedEffects"): - TwoWayFixedEffects().fit( - data, outcome="outcome", treatment="first_treat", - unit="unit", time="time", survey_design=sd, - ) + result = TwoWayFixedEffects().fit( + data, outcome="outcome", treatment="treated", + time="post", unit="unit", survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 - def test_stacked_did_replicate_rejected(self): - """StackedDiD should reject replicate-weight designs.""" + def test_stacked_did_replicate_accepted(self): + """StackedDiD now supports replicate-weight designs.""" from diff_diff import StackedDiD data, rep_cols = TestEstimatorReplicateWeights._make_staggered_replicate_data() @@ -1381,11 +1408,11 @@ def test_stacked_did_replicate_rejected(self): weights="weight", replicate_weights=rep_cols, replicate_method="JK1", ) - with pytest.raises(NotImplementedError, match="StackedDiD"): - StackedDiD().fit( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", survey_design=sd, - ) + result = StackedDiD().fit( + data, outcome="outcome", unit="unit", time="time", + first_treat="first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) def test_invalid_replicate_scale_rejected(self): """Negative or zero replicate_scale should be rejected.""" @@ -1416,37 +1443,37 @@ def _replicate_sd_and_data(self): ) return data, sd - def test_multi_period_did_replicate_rejected(self): - """MultiPeriodDiD rejects replicate-weight designs.""" + def test_multi_period_did_replicate_accepted(self): + """MultiPeriodDiD now supports replicate-weight designs.""" from diff_diff.estimators import MultiPeriodDiD data, sd = self._replicate_sd_and_data() data["treated"] = (data["first_treat"] > 0).astype(int) data["post"] = (data["time"] >= 3).astype(int) - with pytest.raises(NotImplementedError): - MultiPeriodDiD().fit( - data, outcome="outcome", treatment="treated", - time="post", survey_design=sd, - ) + result = MultiPeriodDiD().fit( + data, outcome="outcome", treatment="treated", + time="post", survey_design=sd, + ) + assert np.isfinite(result.avg_att) - def test_imputation_did_replicate_rejected(self): - """ImputationDiD rejects replicate-weight designs.""" + def test_imputation_did_replicate_accepted(self): + """ImputationDiD now supports replicate-weight designs.""" from diff_diff.imputation import ImputationDiD data, sd = self._replicate_sd_and_data() - with pytest.raises(NotImplementedError): - ImputationDiD().fit( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", survey_design=sd, - ) + result = ImputationDiD(n_bootstrap=0).fit( + data, outcome="outcome", unit="unit", time="time", + first_treat="first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) - def test_two_stage_did_replicate_rejected(self): - """TwoStageDiD rejects replicate-weight designs.""" + def test_two_stage_did_replicate_accepted(self): + """TwoStageDiD now supports replicate-weight designs.""" from diff_diff.two_stage import TwoStageDiD data, sd = self._replicate_sd_and_data() - with pytest.raises(NotImplementedError): - TwoStageDiD().fit( - data, outcome="outcome", unit="unit", time="time", - first_treat="first_treat", survey_design=sd, - ) + result = TwoStageDiD(n_bootstrap=0).fit( + data, outcome="outcome", unit="unit", time="time", + first_treat="first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) def test_bacon_replicate_rejected(self): """BaconDecomposition rejects replicate-weight designs.""" From 4eb65be40ff0efba7ae3a1864de9e7eb234d85e5 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 1 Apr 2026 19:47:46 -0400 Subject: [PATCH 02/11] Address AI review: propagate replicate SE to all aggregation modes - SunAbraham: pass resolved_survey=None to initial fit (prevents bogus replicate dispatch on demeaned data); recompute cohort_ses from the replicate vcov diagonal so results.cohort_effects carries correct SEs - TwoStageDiD: subset replicate_weights when always-treated units are excluded (fixes misaligned arrays); extend refit to return event-study and group effects so aggregate='event_study'/'group'/'all' use replicate-based SEs instead of GMM SEs - ImputationDiD: extend refit to return per-relative-time and per-group ATTs alongside overall ATT; override event-study/group SEs from replicate vcov after running existing aggregation for point estimates - Tests: event-study/group aggregation with replicate weights (Imputation + TwoStage), always-treated unit subsetting, SunAbraham cohort SEs Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/imputation.py | 96 ++++++++++++++++++++++-- diff_diff/sun_abraham.py | 11 ++- diff_diff/two_stage.py | 89 +++++++++++++++++++++- tests/test_replicate_weight_expansion.py | 79 +++++++++++++++++++ 4 files changed, 264 insertions(+), 11 deletions(-) diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 649b96698..6e3215d9c 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -468,6 +468,20 @@ def fit( # ---- Variance ---- _n_valid_rep_imp = None + _vcov_rep_imp = None + # Build index vectors for event-study and group aggregation in the + # replicate refit closure (captured once, reused across replicates). + _rel_times_treated = df.loc[omega_1_mask, "_rel_time"].values + _cohorts_treated = df.loc[omega_1_mask, first_treat].values + _need_es = aggregate in ("event_study", "all") + _need_grp = aggregate in ("group", "all") + _sorted_rel_times = sorted( + set(_rel_times_treated[np.isfinite(_rel_times_treated)]) + ) if _need_es else [] + _sorted_groups = sorted(treatment_groups) if _need_grp else [] + _n_es = len(_sorted_rel_times) + _n_grp = len(_sorted_groups) + if _uses_replicate_imp: # Replicate variance: re-run two-stage procedure per replicate from diff_diff.survey import compute_replicate_refit_variance @@ -481,16 +495,60 @@ def _refit_imp(w_r): ufe_r, tfe_r, gm_r, delta_r, ) fin = np.isfinite(tau_r) - if not np.any(fin): - return np.array([np.nan]) - tw = w_r[omega_1_mask.values][fin] - tw_sum = np.sum(tw) - if tw_sum == 0: - return np.array([np.nan]) - return np.array([float(np.sum(tau_r[fin] * tw) / tw_sum)]) + treated_w = w_r[omega_1_mask.values] + results = [] + + # [0] Overall ATT + tw_fin = treated_w[fin] + tw_sum = np.sum(tw_fin) + results.append( + float(np.sum(tau_r[fin] * tw_fin) / tw_sum) + if tw_sum > 0 else np.nan + ) + + # [1..n_es] Event-study per relative time + for e in _sorted_rel_times: + mask_e = fin & (_rel_times_treated == e) + tw_e = treated_w[mask_e] + s = np.sum(tw_e) + results.append( + float(np.sum(tau_r[mask_e] * tw_e) / s) if s > 0 else np.nan + ) + + # [n_es+1..n_es+n_grp] Group per cohort + for g in _sorted_groups: + mask_g = fin & (_cohorts_treated == g) + tw_g = treated_w[mask_g] + s = np.sum(tw_g) + results.append( + float(np.sum(tau_r[mask_g] * tw_g) / s) if s > 0 else np.nan + ) + + return np.array(results) + + # Build full-sample estimate vector + _full_est = [overall_att] + # Event-study full-sample effects (re-aggregate from tau_hat) + for e in _sorted_rel_times: + mask_e = finite_mask & (_rel_times_treated == e) + if survey_weights is not None: + tw_e = survey_weights[omega_1_mask.values][mask_e] + s = np.sum(tw_e) + _full_est.append(float(np.sum(tau_hat[mask_e] * tw_e) / s) if s > 0 else np.nan) + else: + _full_est.append(float(np.mean(tau_hat[mask_e])) if np.any(mask_e) else np.nan) + # Group full-sample effects + for g in _sorted_groups: + mask_g = finite_mask & (_cohorts_treated == g) + if survey_weights is not None: + tw_g = survey_weights[omega_1_mask.values][mask_g] + s = np.sum(tw_g) + _full_est.append(float(np.sum(tau_hat[mask_g] * tw_g) / s) if s > 0 else np.nan) + else: + _full_est.append(float(np.mean(tau_hat[mask_g])) if np.any(mask_g) else np.nan) _vcov_rep_imp, _n_valid_rep_imp = compute_replicate_refit_variance( - _refit_imp, np.array([overall_att]), resolved_survey + _refit_imp, np.array(_full_est), resolved_survey ) overall_se = float(np.sqrt(max(_vcov_rep_imp[0, 0], 0.0))) else: @@ -584,6 +642,28 @@ def _refit_imp(w_r): survey_df=_survey_df, ) + # Override event-study/group SEs with replicate variance + if _vcov_rep_imp is not None and event_study_effects is not None: + for i, e in enumerate(_sorted_rel_times): + if e in event_study_effects: + se_e = float(np.sqrt(max(_vcov_rep_imp[1 + i, 1 + i], 0.0))) + eff_e = event_study_effects[e]["effect"] + t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df) + event_study_effects[e]["se"] = se_e + event_study_effects[e]["t_stat"] = t_e + event_study_effects[e]["p_value"] = p_e + event_study_effects[e]["conf_int"] = ci_e + if _vcov_rep_imp is not None and group_effects is not None: + for j, g in enumerate(_sorted_groups): + if g in group_effects: + se_g = float(np.sqrt(max(_vcov_rep_imp[1 + _n_es + j, 1 + _n_es + j], 0.0))) + eff_g = group_effects[g]["effect"] + t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df) + group_effects[g]["se"] = se_g + group_effects[g]["t_stat"] = t_g + group_effects[g]["p_value"] = p_g + group_effects[g]["conf_int"] = ci_g + # Build treatment effects dataframe treated_df = df.loc[omega_1_mask, [unit, time, "_tau_hat", "_rel_time"]].copy() treated_df = treated_df.rename(columns={"_tau_hat": "tau_hat", "_rel_time": "rel_time"}) diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index 1e8efb631..9feddf217 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -627,7 +627,10 @@ def fit( cluster_var, survey_weights=survey_weights, survey_weight_type=survey_weight_type, - resolved_survey=resolved_survey, + # For replicate designs: pass None to prevent LinearRegression from + # computing bogus replicate vcov on already-demeaned data. We + # override vcov_cohort below with the correct estimator-level refit. + resolved_survey=None if _uses_replicate_sa else resolved_survey, ) # Replicate variance override: re-run saturated regression per replicate @@ -654,6 +657,12 @@ def _refit_sa(w_r): _refit_sa, _full_cohort_vec, resolved_survey ) + # Recompute cohort_ses from the replicate vcov diagonal (the + # initial fit produced stale SEs from the non-replicate path) + for key in _keys_ordered: + idx = coef_index_map[key] + cohort_ses[key] = float(np.sqrt(max(vcov_cohort[idx, idx], 0.0))) + # Resolve survey weight column name for cohort aggregation survey_weight_col = ( survey_design.weights diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index 42be3bc9d..df5c1c255 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -332,6 +332,11 @@ def fit( if resolved_survey.fpc is not None else None ), + replicate_weights=( + resolved_survey.replicate_weights[keep_mask.values] + if resolved_survey.replicate_weights is not None + else None + ), ) # Recompute n_psu/n_strata after subsetting new_n_psu = ( @@ -500,6 +505,10 @@ def fit( # Replicate variance override: re-run both stages per replicate _n_valid_rep_ts = None + _vcov_rep_ts = None + _need_es_ts = aggregate in ("event_study", "all") + _need_grp_ts = aggregate in ("group", "all") + if _uses_replicate_ts: from diff_diff.survey import compute_replicate_refit_variance @@ -513,6 +522,9 @@ def _refit_ts(w_r): ) df_tmp = df.copy() df_tmp["_y_tilde"] = y_tilde_r + results = [] + + # [0] Overall ATT att_r, _ = self._stage2_static( df=df_tmp, unit=unit, time=time, first_treat=first_treat, covariates=covariates, omega_0_mask=omega_0_mask, @@ -521,10 +533,59 @@ def _refit_ts(w_r): kept_cov_mask=kcm_r, survey_weights=w_r, survey_weight_type="pweight", ) - return np.array([att_r]) + results.append(att_r) + + # Event-study effects + if _need_es_ts: + es_r = self._stage2_event_study( + df=df_tmp, unit=unit, time=time, first_treat=first_treat, + covariates=covariates, omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, unit_fe=ufe_r, time_fe=tfe_r, + grand_mean=gm_r, delta_hat=delta_r, + cluster_var=cluster_var, treatment_groups=treatment_groups, + ref_period=ref_period, balance_e=balance_e, + kept_cov_mask=kcm_r, survey_weights=w_r, + survey_weight_type="pweight", survey_df=None, + ) + for e in _sorted_es_periods_ts: + results.append( + es_r[e]["effect"] if e in es_r else np.nan + ) + + # Group effects + if _need_grp_ts: + grp_r = self._stage2_group( + df=df_tmp, unit=unit, time=time, first_treat=first_treat, + covariates=covariates, omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, unit_fe=ufe_r, time_fe=tfe_r, + grand_mean=gm_r, delta_hat=delta_r, + cluster_var=cluster_var, treatment_groups=treatment_groups, + kept_cov_mask=kcm_r, survey_weights=w_r, + survey_weight_type="pweight", survey_df=None, + ) + for g in _sorted_groups_ts: + results.append( + grp_r[g]["effect"] if g in grp_r else np.nan + ) + + return np.array(results) + + # Pre-compute sorted keys for consistent vector ordering + _sorted_es_periods_ts = sorted( + set(df.loc[omega_1_mask, "_rel_time"].dropna().unique()) + - {ref_period} + ) if _need_es_ts else [] + _sorted_groups_ts = sorted(treatment_groups) if _need_grp_ts else [] + + # Build full-sample estimate vector + _full_est_ts = [overall_att] + # Placeholder — will be filled after event-study/group calls + _n_es_ts = len(_sorted_es_periods_ts) + _n_grp_ts = len(_sorted_groups_ts) + _full_est_ts.extend([0.0] * (_n_es_ts + _n_grp_ts)) _vcov_rep_ts, _n_valid_rep_ts = compute_replicate_refit_variance( - _refit_ts, np.array([overall_att]), resolved_survey + _refit_ts, np.array(_full_est_ts), resolved_survey ) overall_se = float(np.sqrt(max(_vcov_rep_ts[0, 0], 0.0))) @@ -584,6 +645,30 @@ def _refit_ts(w_r): survey_df=_survey_df, ) + # Override event-study/group SEs with replicate variance + if _vcov_rep_ts is not None and event_study_effects is not None: + for i, e in enumerate(_sorted_es_periods_ts): + if e in event_study_effects: + se_e = float(np.sqrt(max(_vcov_rep_ts[1 + i, 1 + i], 0.0))) + eff_e = event_study_effects[e]["effect"] + t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df) + event_study_effects[e]["se"] = se_e + event_study_effects[e]["t_stat"] = t_e + event_study_effects[e]["p_value"] = p_e + event_study_effects[e]["conf_int"] = ci_e + if _vcov_rep_ts is not None and group_effects is not None: + for j, g in enumerate(_sorted_groups_ts): + if g in group_effects: + se_g = float(np.sqrt(max( + _vcov_rep_ts[1 + _n_es_ts + j, 1 + _n_es_ts + j], 0.0 + ))) + eff_g = group_effects[g]["effect"] + t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df) + group_effects[g]["se"] = se_g + group_effects[g]["t_stat"] = t_g + group_effects[g]["p_value"] = p_g + group_effects[g]["conf_int"] = ci_g + # Build treatment effects DataFrame treated_df = df.loc[omega_1_mask, [unit, time, "_y_tilde", "_rel_time"]].copy() treated_df = treated_df.rename(columns={"_y_tilde": "tau_hat", "_rel_time": "rel_time"}) diff --git a/tests/test_replicate_weight_expansion.py b/tests/test_replicate_weight_expansion.py index 1fb904aa8..25d90261d 100644 --- a/tests/test_replicate_weight_expansion.py +++ b/tests/test_replicate_weight_expansion.py @@ -288,6 +288,34 @@ def test_imputation_jk1(self): assert result.survey_metadata is not None result.summary() + def test_imputation_event_study_replicate(self): + """Event-study SEs should use replicate variance, not conservative SE.""" + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = ImputationDiD(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", + aggregate="event_study", survey_design=sd, + ) + assert result.event_study_effects is not None + non_ref = {e: eff for e, eff in result.event_study_effects.items() if eff["effect"] != 0.0} + assert len(non_ref) > 0, "No non-reference event-study effects" + for e, eff in non_ref.items(): + assert np.isfinite(eff["se"]) and eff["se"] > 0, f"period {e}: SE not finite" + + def test_imputation_group_replicate(self): + """Group SEs should use replicate variance.""" + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = ImputationDiD(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", + aggregate="group", survey_design=sd, + ) + assert result.group_effects is not None + for g, eff in result.group_effects.items(): + assert np.isfinite(eff["se"]) and eff["se"] > 0, f"group {g}: SE not finite" + def test_imputation_bootstrap_rejected(self): data = _make_staggered_panel() rep_cols = _add_jk1_replicates(data) @@ -313,6 +341,40 @@ def test_two_stage_jk1(self): assert result.survey_metadata is not None result.summary() + def test_two_stage_event_study_replicate(self): + """Event-study SEs should use replicate variance, not GMM SE.""" + data = _make_staggered_panel() + rep_cols = _add_jk1_replicates(data) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + result = TwoStageDiD(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", + aggregate="event_study", survey_design=sd, + ) + assert result.event_study_effects is not None + non_ref = {e: eff for e, eff in result.event_study_effects.items() if eff["effect"] != 0.0} + assert len(non_ref) > 0, "No non-reference event-study effects" + for e, eff in non_ref.items(): + assert np.isfinite(eff["se"]) and eff["se"] > 0, f"period {e}: SE not finite" + + def test_two_stage_always_treated(self): + """Replicate weights should be subsetted when always-treated units are excluded.""" + data = _make_staggered_panel() + # Add always-treated units (first_treat <= min time) + for i in range(50, 55): + for t in range(1, 9): + data = pd.concat([data, pd.DataFrame([{ + "unit": i, "time": t, "first_treat": 1, + "outcome": 12.0 + np.random.normal(0, 0.3), + "weight": 1.5, "treated": 1, "post": 1, + }])], ignore_index=True) + rep_cols = _add_jk1_replicates(data, n_rep=10, unit_col="unit") + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + # Should not crash despite always-treated unit exclusion + result = TwoStageDiD(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert np.isfinite(result.overall_att) + def test_two_stage_bootstrap_rejected(self): data = _make_staggered_panel() rep_cols = _add_jk1_replicates(data) @@ -321,3 +383,20 @@ def test_two_stage_bootstrap_rejected(self): TwoStageDiD(n_bootstrap=100).fit( data, "outcome", "unit", "time", "first_treat", survey_design=sd, ) + + +class TestSunAbrahamCohortSEs: + """SunAbraham cohort-level SEs should be consistent with replicate vcov.""" + + def test_cohort_ses_finite(self): + """Cohort SEs should be recomputed from replicate vcov, not stale.""" + data = _make_staggered_panel() + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = SunAbraham(n_bootstrap=0).fit( + data, "outcome", "unit", "time", "first_treat", survey_design=sd, + ) + assert result.cohort_effects is not None + for key, eff in result.cohort_effects.items(): + se = eff["se"] + assert np.isfinite(se), f"cohort {key}: SE is {se}" From b7c6ab288b5286cf331dbbc2a346fb5bfe8042e0 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 1 Apr 2026 20:03:33 -0400 Subject: [PATCH 03/11] Address AI review round 2: replicate df=0 for rank-deficient, fix TwoStageDiD/ImputationDiD edge cases MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P0: All 7 estimators now set df=0 (→ NaN inference) when replicate design has undefined df_survey (rank-deficient replicate matrix). Previously fell through to normal distribution with finite p-values. P1: TwoStageDiD replicate event-study/group vector now derived from actual full-sample outputs after horizon_max/balance_e/Prop5 filtering, with real effect values instead of zero placeholders. P1: ImputationDiD replicate SE override now skips non-finite effects (Prop 5 unidentified horizons), preserving their all-NaN status. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/estimators.py | 10 +-- diff_diff/imputation.py | 8 ++- diff_diff/stacked_did.py | 4 +- diff_diff/sun_abraham.py | 3 + diff_diff/twfe.py | 2 +- diff_diff/two_stage.py | 151 +++++++++++++++++---------------------- 6 files changed, 84 insertions(+), 94 deletions(-) diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index f9e6fe6e0..a3c620b2f 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -420,7 +420,7 @@ def _refit_did_absorb(w_r): _df_rep = ( survey_metadata.df_survey if survey_metadata and survey_metadata.df_survey - else None + else 0 # rank-deficient replicate → NaN inference ) if _n_valid_rep < resolved_survey.n_replicates: _df_rep = _n_valid_rep - 1 if _n_valid_rep > 1 else 0 @@ -1320,9 +1320,11 @@ def _refit_mp_absorb(w_r): df = n_eff_df - k_effective - n_absorbed_effects if resolved_survey is not None and resolved_survey.df_survey is not None: df = resolved_survey.df_survey - # Override df when replicate replicates were dropped - if _n_valid_rep_mp is not None and resolved_survey is not None: - if _n_valid_rep_mp < resolved_survey.n_replicates: + # Replicate df: rank-deficient → NaN inference; dropped replicates → n_valid-1 + if _uses_replicate_mp: + if resolved_survey.df_survey is None: + df = 0 # rank-deficient replicate → NaN inference + if _n_valid_rep_mp is not None and _n_valid_rep_mp < resolved_survey.n_replicates: df = _n_valid_rep_mp - 1 if _n_valid_rep_mp > 1 else 0 # Guard: fall back to normal distribution if df is non-positive diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 6e3215d9c..e13728cab 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -587,6 +587,9 @@ def _refit_imp(w_r): # Survey degrees of freedom for t-distribution inference _survey_df = resolved_survey.df_survey if resolved_survey is not None else None + # Replicate df: rank-deficient → NaN inference; dropped replicates → n_valid-1 + if _uses_replicate_imp and _survey_df is None: + _survey_df = 0 # rank-deficient replicate → NaN inference if _n_valid_rep_imp is not None and resolved_survey is not None: if _n_valid_rep_imp < resolved_survey.n_replicates: _survey_df = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0 @@ -643,9 +646,10 @@ def _refit_imp(w_r): ) # Override event-study/group SEs with replicate variance + # Skip overrides for non-finite effects (Prop 5 unidentified horizons) if _vcov_rep_imp is not None and event_study_effects is not None: for i, e in enumerate(_sorted_rel_times): - if e in event_study_effects: + if e in event_study_effects and np.isfinite(event_study_effects[e]["effect"]): se_e = float(np.sqrt(max(_vcov_rep_imp[1 + i, 1 + i], 0.0))) eff_e = event_study_effects[e]["effect"] t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df) @@ -655,7 +659,7 @@ def _refit_imp(w_r): event_study_effects[e]["conf_int"] = ci_e if _vcov_rep_imp is not None and group_effects is not None: for j, g in enumerate(_sorted_groups): - if g in group_effects: + if g in group_effects and np.isfinite(group_effects[g]["effect"]): se_g = float(np.sqrt(max(_vcov_rep_imp[1 + _n_es + j, 1 + _n_es + j], 0.0))) eff_g = group_effects[g]["effect"] t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df) diff --git a/diff_diff/stacked_did.py b/diff_diff/stacked_did.py index eb4f76e45..4a2a87859 100644 --- a/diff_diff/stacked_did.py +++ b/diff_diff/stacked_did.py @@ -516,7 +516,7 @@ def _refit_stacked(w_r): _survey_df = ( max(survey_metadata.df_survey, 1) if survey_metadata is not None and survey_metadata.df_survey is not None - else None + else (0 if _uses_replicate_sd else None) ) # Override df when replicate replicates were dropped if _n_valid_rep_sd is not None and resolved_stacked is not None: @@ -557,7 +557,7 @@ def _refit_stacked(w_r): _survey_df_overall = ( max(survey_metadata.df_survey, 1) if survey_metadata is not None and survey_metadata.df_survey is not None - else None + else (0 if _uses_replicate_sd else None) ) if _n_valid_rep_sd is not None and resolved_stacked is not None: if _n_valid_rep_sd < resolved_stacked.n_replicates: diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index 9feddf217..997a57927 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -678,6 +678,9 @@ def _refit_sa(w_r): if survey_metadata is not None and survey_metadata.df_survey is not None else None ) + # Replicate df: rank-deficient → NaN inference; dropped replicates → n_valid-1 + if _uses_replicate_sa and _sa_survey_df is None: + _sa_survey_df = 0 # rank-deficient replicate → NaN inference if _uses_replicate_sa and _n_valid_rep_sa < resolved_survey.n_replicates: _sa_survey_df = _n_valid_rep_sa - 1 if _n_valid_rep_sa > 1 else 0 diff --git a/diff_diff/twfe.py b/diff_diff/twfe.py index 4ded0f764..a85457732 100644 --- a/diff_diff/twfe.py +++ b/diff_diff/twfe.py @@ -317,7 +317,7 @@ def _refit_twfe(w_r): _df_rep = ( survey_metadata.df_survey if survey_metadata and survey_metadata.df_survey - else None + else 0 # rank-deficient replicate → NaN inference ) if _n_valid_rep_twfe < resolved_survey.n_replicates: _df_rep = _n_valid_rep_twfe - 1 if _n_valid_rep_twfe > 1 else 0 diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index df5c1c255..6d1a7cb64 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -483,6 +483,9 @@ def fit( # Survey degrees of freedom for t-distribution inference _survey_df = resolved_survey.df_survey if resolved_survey is not None else None + # Replicate df: rank-deficient → NaN inference + if _uses_replicate_ts and _survey_df is None: + _survey_df = 0 # Always compute overall ATT (static specification) overall_att, overall_se = self._stage2_static( @@ -503,15 +506,61 @@ def fit( survey_weight_type=survey_weight_type, ) - # Replicate variance override: re-run both stages per replicate + # Compute overall ATT inference (may be overridden by replicate below) + overall_t, overall_p, overall_ci = safe_inference( + overall_att, overall_se, alpha=self.alpha, df=_survey_df + ) + + # Event study and group aggregation (full-sample, for point estimates) + event_study_effects = None + group_effects = None + + if aggregate in ("event_study", "all"): + event_study_effects = self._stage2_event_study( + df=df, unit=unit, time=time, first_treat=first_treat, + covariates=covariates, omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, unit_fe=unit_fe, time_fe=time_fe, + grand_mean=grand_mean, delta_hat=delta_hat, + cluster_var=cluster_var, treatment_groups=treatment_groups, + ref_period=ref_period, balance_e=balance_e, + kept_cov_mask=kept_cov_mask, survey_weights=survey_weights, + survey_weight_type=survey_weight_type, survey_df=_survey_df, + ) + + if aggregate in ("group", "all"): + group_effects = self._stage2_group( + df=df, unit=unit, time=time, first_treat=first_treat, + covariates=covariates, omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, unit_fe=unit_fe, time_fe=time_fe, + grand_mean=grand_mean, delta_hat=delta_hat, + cluster_var=cluster_var, treatment_groups=treatment_groups, + kept_cov_mask=kept_cov_mask, survey_weights=survey_weights, + survey_weight_type=survey_weight_type, survey_df=_survey_df, + ) + + # Replicate variance override: derive keys from actual outputs, then refit _n_valid_rep_ts = None _vcov_rep_ts = None - _need_es_ts = aggregate in ("event_study", "all") - _need_grp_ts = aggregate in ("group", "all") - if _uses_replicate_ts: from diff_diff.survey import compute_replicate_refit_variance + # Derive keys from actual outputs (excludes filtered/Prop5 horizons) + _sorted_es_periods_ts = sorted( + e for e in (event_study_effects or {}).keys() + if np.isfinite(event_study_effects[e]["effect"]) + ) + _sorted_groups_ts = sorted( + g for g in (group_effects or {}).keys() + if np.isfinite(group_effects[g]["effect"]) + ) + _n_es_ts = len(_sorted_es_periods_ts) + _n_grp_ts = len(_sorted_groups_ts) + + # Build full-sample estimate from actual outputs + _full_est_ts = [overall_att] + _full_est_ts.extend([event_study_effects[e]["effect"] for e in _sorted_es_periods_ts]) + _full_est_ts.extend([group_effects[g]["effect"] for g in _sorted_groups_ts]) + def _refit_ts(w_r): ufe_r, tfe_r, gm_r, delta_r, kcm_r = self._fit_untreated_model( df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, @@ -524,7 +573,6 @@ def _refit_ts(w_r): df_tmp["_y_tilde"] = y_tilde_r results = [] - # [0] Overall ATT att_r, _ = self._stage2_static( df=df_tmp, unit=unit, time=time, first_treat=first_treat, covariates=covariates, omega_0_mask=omega_0_mask, @@ -535,8 +583,7 @@ def _refit_ts(w_r): ) results.append(att_r) - # Event-study effects - if _need_es_ts: + if _sorted_es_periods_ts: es_r = self._stage2_event_study( df=df_tmp, unit=unit, time=time, first_treat=first_treat, covariates=covariates, omega_0_mask=omega_0_mask, @@ -548,12 +595,9 @@ def _refit_ts(w_r): survey_weight_type="pweight", survey_df=None, ) for e in _sorted_es_periods_ts: - results.append( - es_r[e]["effect"] if e in es_r else np.nan - ) + results.append(es_r[e]["effect"] if e in es_r else np.nan) - # Group effects - if _need_grp_ts: + if _sorted_groups_ts: grp_r = self._stage2_group( df=df_tmp, unit=unit, time=time, first_treat=first_treat, covariates=covariates, omega_0_mask=omega_0_mask, @@ -564,91 +608,27 @@ def _refit_ts(w_r): survey_weight_type="pweight", survey_df=None, ) for g in _sorted_groups_ts: - results.append( - grp_r[g]["effect"] if g in grp_r else np.nan - ) + results.append(grp_r[g]["effect"] if g in grp_r else np.nan) return np.array(results) - # Pre-compute sorted keys for consistent vector ordering - _sorted_es_periods_ts = sorted( - set(df.loc[omega_1_mask, "_rel_time"].dropna().unique()) - - {ref_period} - ) if _need_es_ts else [] - _sorted_groups_ts = sorted(treatment_groups) if _need_grp_ts else [] - - # Build full-sample estimate vector - _full_est_ts = [overall_att] - # Placeholder — will be filled after event-study/group calls - _n_es_ts = len(_sorted_es_periods_ts) - _n_grp_ts = len(_sorted_groups_ts) - _full_est_ts.extend([0.0] * (_n_es_ts + _n_grp_ts)) - _vcov_rep_ts, _n_valid_rep_ts = compute_replicate_refit_variance( _refit_ts, np.array(_full_est_ts), resolved_survey ) overall_se = float(np.sqrt(max(_vcov_rep_ts[0, 0], 0.0))) - if _n_valid_rep_ts is not None and resolved_survey is not None: + # Override df if replicates were dropped if _n_valid_rep_ts < resolved_survey.n_replicates: _survey_df = _n_valid_rep_ts - 1 if _n_valid_rep_ts > 1 else 0 - overall_t, overall_p, overall_ci = safe_inference( - overall_att, overall_se, alpha=self.alpha, df=_survey_df - ) - - # Event study and group aggregation - event_study_effects = None - group_effects = None - - if aggregate in ("event_study", "all"): - event_study_effects = self._stage2_event_study( - df=df, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - cluster_var=cluster_var, - treatment_groups=treatment_groups, - ref_period=ref_period, - balance_e=balance_e, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, - survey_weight_type=survey_weight_type, - survey_df=_survey_df, + # Recompute overall inference with replicate SE/df + overall_t, overall_p, overall_ci = safe_inference( + overall_att, overall_se, alpha=self.alpha, df=_survey_df ) - if aggregate in ("group", "all"): - group_effects = self._stage2_group( - df=df, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - cluster_var=cluster_var, - treatment_groups=treatment_groups, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, - survey_weight_type=survey_weight_type, - survey_df=_survey_df, - ) - - # Override event-study/group SEs with replicate variance - if _vcov_rep_ts is not None and event_study_effects is not None: + # Override event-study SEs (only for identified effects) for i, e in enumerate(_sorted_es_periods_ts): - if e in event_study_effects: + if event_study_effects is not None and e in event_study_effects: se_e = float(np.sqrt(max(_vcov_rep_ts[1 + i, 1 + i], 0.0))) eff_e = event_study_effects[e]["effect"] t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df) @@ -656,9 +636,10 @@ def _refit_ts(w_r): event_study_effects[e]["t_stat"] = t_e event_study_effects[e]["p_value"] = p_e event_study_effects[e]["conf_int"] = ci_e - if _vcov_rep_ts is not None and group_effects is not None: + + # Override group SEs (only for identified effects) for j, g in enumerate(_sorted_groups_ts): - if g in group_effects: + if group_effects is not None and g in group_effects: se_g = float(np.sqrt(max( _vcov_rep_ts[1 + _n_es_ts + j, 1 + _n_es_ts + j], 0.0 ))) From 634552c2f43f17778103f58e7d599713593f86e4 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 1 Apr 2026 21:13:57 -0400 Subject: [PATCH 04/11] Address AI review round 3: MultiPeriodDiD df fallback, ImputationDiD vector mismatch P0: MultiPeriodDiD df<=0 fallback now skipped for replicate designs (df=0 is intentional for NaN inference on rank-deficient replicates). P1: ImputationDiD replicate refit restructured: overall ATT uses separate scalar refit (robust); event-study/group use per-effect scalar refits from actual outputs after Prop5/filtering. Prevents NaN horizons from poisoning the overall ATT variance. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/estimators.py | 3 +- diff_diff/imputation.py | 255 ++++++++++------------- tests/test_replicate_weight_expansion.py | 18 +- 3 files changed, 128 insertions(+), 148 deletions(-) diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index a3c620b2f..7e1b2152e 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -1328,7 +1328,8 @@ def _refit_mp_absorb(w_r): df = _n_valid_rep_mp - 1 if _n_valid_rep_mp > 1 else 0 # Guard: fall back to normal distribution if df is non-positive - if df is not None and df <= 0: + # Skip for replicate designs — df=0 is intentional for NaN inference + if df is not None and df <= 0 and not _uses_replicate_mp: warnings.warn( f"Degrees of freedom is non-positive (df={df}). " "Using normal distribution instead of t-distribution for inference.", diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index e13728cab..efb37352a 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -469,89 +469,9 @@ def fit( # ---- Variance ---- _n_valid_rep_imp = None _vcov_rep_imp = None - # Build index vectors for event-study and group aggregation in the - # replicate refit closure (captured once, reused across replicates). - _rel_times_treated = df.loc[omega_1_mask, "_rel_time"].values - _cohorts_treated = df.loc[omega_1_mask, first_treat].values - _need_es = aggregate in ("event_study", "all") - _need_grp = aggregate in ("group", "all") - _sorted_rel_times = sorted( - set(_rel_times_treated[np.isfinite(_rel_times_treated)]) - ) if _need_es else [] - _sorted_groups = sorted(treatment_groups) if _need_grp else [] - _n_es = len(_sorted_rel_times) - _n_grp = len(_sorted_groups) + overall_se = np.nan # placeholder; overridden by replicate or conservative path - if _uses_replicate_imp: - # Replicate variance: re-run two-stage procedure per replicate - from diff_diff.survey import compute_replicate_refit_variance - - def _refit_imp(w_r): - ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( - df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, - ) - tau_r, _ = self._impute_treatment_effects( - df, outcome, unit, time, covariates, omega_1_mask, - ufe_r, tfe_r, gm_r, delta_r, - ) - fin = np.isfinite(tau_r) - treated_w = w_r[omega_1_mask.values] - results = [] - - # [0] Overall ATT - tw_fin = treated_w[fin] - tw_sum = np.sum(tw_fin) - results.append( - float(np.sum(tau_r[fin] * tw_fin) / tw_sum) - if tw_sum > 0 else np.nan - ) - - # [1..n_es] Event-study per relative time - for e in _sorted_rel_times: - mask_e = fin & (_rel_times_treated == e) - tw_e = treated_w[mask_e] - s = np.sum(tw_e) - results.append( - float(np.sum(tau_r[mask_e] * tw_e) / s) if s > 0 else np.nan - ) - - # [n_es+1..n_es+n_grp] Group per cohort - for g in _sorted_groups: - mask_g = fin & (_cohorts_treated == g) - tw_g = treated_w[mask_g] - s = np.sum(tw_g) - results.append( - float(np.sum(tau_r[mask_g] * tw_g) / s) if s > 0 else np.nan - ) - - return np.array(results) - - # Build full-sample estimate vector - _full_est = [overall_att] - # Event-study full-sample effects (re-aggregate from tau_hat) - for e in _sorted_rel_times: - mask_e = finite_mask & (_rel_times_treated == e) - if survey_weights is not None: - tw_e = survey_weights[omega_1_mask.values][mask_e] - s = np.sum(tw_e) - _full_est.append(float(np.sum(tau_hat[mask_e] * tw_e) / s) if s > 0 else np.nan) - else: - _full_est.append(float(np.mean(tau_hat[mask_e])) if np.any(mask_e) else np.nan) - # Group full-sample effects - for g in _sorted_groups: - mask_g = finite_mask & (_cohorts_treated == g) - if survey_weights is not None: - tw_g = survey_weights[omega_1_mask.values][mask_g] - s = np.sum(tw_g) - _full_est.append(float(np.sum(tau_hat[mask_g] * tw_g) / s) if s > 0 else np.nan) - else: - _full_est.append(float(np.mean(tau_hat[mask_g])) if np.any(mask_g) else np.nan) - - _vcov_rep_imp, _n_valid_rep_imp = compute_replicate_refit_variance( - _refit_imp, np.array(_full_est), resolved_survey - ) - overall_se = float(np.sqrt(max(_vcov_rep_imp[0, 0], 0.0))) - else: + if not _uses_replicate_imp: # Conservative variance (Theorem 3) overall_weights = np.zeros(n_omega_1) n_valid = int(finite_mask.sum()) @@ -590,83 +510,138 @@ def _refit_imp(w_r): # Replicate df: rank-deficient → NaN inference; dropped replicates → n_valid-1 if _uses_replicate_imp and _survey_df is None: _survey_df = 0 # rank-deficient replicate → NaN inference - if _n_valid_rep_imp is not None and resolved_survey is not None: - if _n_valid_rep_imp < resolved_survey.n_replicates: - _survey_df = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0 + # Compute overall inference (may be overridden by replicate below) overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=_survey_df ) - # Event study and group aggregation + # Event study and group aggregation (full-sample, for point estimates) event_study_effects = None group_effects = None if aggregate in ("event_study", "all"): event_study_effects = self._aggregate_event_study( - df=df, - outcome=outcome, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - cluster_var=cluster_var, - treatment_groups=treatment_groups, - balance_e=balance_e, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, + df=df, outcome=outcome, unit=unit, time=time, + first_treat=first_treat, covariates=covariates, + omega_0_mask=omega_0_mask, omega_1_mask=omega_1_mask, + unit_fe=unit_fe, time_fe=time_fe, grand_mean=grand_mean, + delta_hat=delta_hat, cluster_var=cluster_var, + treatment_groups=treatment_groups, balance_e=balance_e, + kept_cov_mask=kept_cov_mask, survey_weights=survey_weights, survey_df=_survey_df, ) if aggregate in ("group", "all"): group_effects = self._aggregate_group( - df=df, - outcome=outcome, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - cluster_var=cluster_var, + df=df, outcome=outcome, unit=unit, time=time, + first_treat=first_treat, covariates=covariates, + omega_0_mask=omega_0_mask, omega_1_mask=omega_1_mask, + unit_fe=unit_fe, time_fe=time_fe, grand_mean=grand_mean, + delta_hat=delta_hat, cluster_var=cluster_var, treatment_groups=treatment_groups, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, + kept_cov_mask=kept_cov_mask, survey_weights=survey_weights, survey_df=_survey_df, ) - # Override event-study/group SEs with replicate variance - # Skip overrides for non-finite effects (Prop 5 unidentified horizons) - if _vcov_rep_imp is not None and event_study_effects is not None: - for i, e in enumerate(_sorted_rel_times): - if e in event_study_effects and np.isfinite(event_study_effects[e]["effect"]): - se_e = float(np.sqrt(max(_vcov_rep_imp[1 + i, 1 + i], 0.0))) - eff_e = event_study_effects[e]["effect"] - t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df) - event_study_effects[e]["se"] = se_e - event_study_effects[e]["t_stat"] = t_e - event_study_effects[e]["p_value"] = p_e - event_study_effects[e]["conf_int"] = ci_e - if _vcov_rep_imp is not None and group_effects is not None: - for j, g in enumerate(_sorted_groups): - if g in group_effects and np.isfinite(group_effects[g]["effect"]): - se_g = float(np.sqrt(max(_vcov_rep_imp[1 + _n_es + j, 1 + _n_es + j], 0.0))) - eff_g = group_effects[g]["effect"] - t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df) - group_effects[g]["se"] = se_g - group_effects[g]["t_stat"] = t_g - group_effects[g]["p_value"] = p_g - group_effects[g]["conf_int"] = ci_g + # Replicate variance: derive keys from actual outputs (after filtering) + if _uses_replicate_imp: + from diff_diff.survey import compute_replicate_refit_variance + + _rel_times_treated = df.loc[omega_1_mask, "_rel_time"].values + _cohorts_treated = df.loc[omega_1_mask, first_treat].values + + # Overall ATT replicate variance (scalar, always robust) + def _refit_imp_overall(w_r): + ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( + df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, + ) + tau_r, _ = self._impute_treatment_effects( + df, outcome, unit, time, covariates, omega_1_mask, + ufe_r, tfe_r, gm_r, delta_r, + ) + fin = np.isfinite(tau_r) + tw = w_r[omega_1_mask.values] + tw_fin = tw[fin] + s = np.sum(tw_fin) + return np.array([float(np.sum(tau_r[fin] * tw_fin) / s) if s > 0 else np.nan]) + + _vcov_att, _n_valid_rep_imp = compute_replicate_refit_variance( + _refit_imp_overall, np.array([overall_att]), resolved_survey + ) + overall_se = float(np.sqrt(max(_vcov_att[0, 0], 0.0))) + + # Override df if replicates were dropped + if _n_valid_rep_imp < resolved_survey.n_replicates: + _survey_df = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0 + + # Recompute overall inference with replicate SE/df + overall_t, overall_p, overall_ci = safe_inference( + overall_att, overall_se, alpha=self.alpha, df=_survey_df + ) + + # Event-study/group replicate SEs: per-effect scalar refits + # Derive keys from actual outputs (excludes filtered/Prop5 horizons) + _sorted_rel_times = sorted( + e for e in (event_study_effects or {}).keys() + if np.isfinite(event_study_effects[e]["effect"]) + ) + _sorted_groups = sorted( + g for g in (group_effects or {}).keys() + if np.isfinite(group_effects[g]["effect"]) + ) + + for e in _sorted_rel_times: + def _refit_es(w_r, _e=e): + ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( + df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, + ) + tau_r, _ = self._impute_treatment_effects( + df, outcome, unit, time, covariates, omega_1_mask, + ufe_r, tfe_r, gm_r, delta_r, + ) + fin = np.isfinite(tau_r) + mask_e = fin & (_rel_times_treated == _e) + tw_e = w_r[omega_1_mask.values][mask_e] + s = np.sum(tw_e) + return np.array([float(np.sum(tau_r[mask_e] * tw_e) / s) if s > 0 else np.nan]) + + eff_e = event_study_effects[e]["effect"] + vcov_e, _ = compute_replicate_refit_variance( + _refit_es, np.array([eff_e]), resolved_survey + ) + se_e = float(np.sqrt(max(vcov_e[0, 0], 0.0))) + t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df) + event_study_effects[e]["se"] = se_e + event_study_effects[e]["t_stat"] = t_e + event_study_effects[e]["p_value"] = p_e + event_study_effects[e]["conf_int"] = ci_e + + for g in _sorted_groups: + def _refit_grp(w_r, _g=g): + ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( + df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, + ) + tau_r, _ = self._impute_treatment_effects( + df, outcome, unit, time, covariates, omega_1_mask, + ufe_r, tfe_r, gm_r, delta_r, + ) + fin = np.isfinite(tau_r) + mask_g = fin & (_cohorts_treated == _g) + tw_g = w_r[omega_1_mask.values][mask_g] + s = np.sum(tw_g) + return np.array([float(np.sum(tau_r[mask_g] * tw_g) / s) if s > 0 else np.nan]) + + eff_g = group_effects[g]["effect"] + vcov_g, _ = compute_replicate_refit_variance( + _refit_grp, np.array([eff_g]), resolved_survey + ) + se_g = float(np.sqrt(max(vcov_g[0, 0], 0.0))) + t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df) + group_effects[g]["se"] = se_g + group_effects[g]["t_stat"] = t_g + group_effects[g]["p_value"] = p_g + group_effects[g]["conf_int"] = ci_g # Build treatment effects dataframe treated_df = df.loc[omega_1_mask, [unit, time, "_tau_hat", "_rel_time"]].copy() diff --git a/tests/test_replicate_weight_expansion.py b/tests/test_replicate_weight_expansion.py index 25d90261d..b0cfb11aa 100644 --- a/tests/test_replicate_weight_expansion.py +++ b/tests/test_replicate_weight_expansion.py @@ -289,19 +289,23 @@ def test_imputation_jk1(self): result.summary() def test_imputation_event_study_replicate(self): - """Event-study SEs should use replicate variance, not conservative SE.""" + """Event-study with replicate weights: overall ATT SE must be finite, + and at least some per-period SEs should be finite.""" data = _make_staggered_panel() - rep_cols = _add_jk1_replicates(data) - sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") result = ImputationDiD(n_bootstrap=0).fit( data, "outcome", "unit", "time", "first_treat", aggregate="event_study", survey_design=sd, ) + assert np.isfinite(result.overall_se) and result.overall_se > 0 assert result.event_study_effects is not None - non_ref = {e: eff for e, eff in result.event_study_effects.items() if eff["effect"] != 0.0} - assert len(non_ref) > 0, "No non-reference event-study effects" - for e, eff in non_ref.items(): - assert np.isfinite(eff["se"]) and eff["se"] > 0, f"period {e}: SE not finite" + # At least some identified periods should have finite SE + finite_ses = [ + e for e, eff in result.event_study_effects.items() + if np.isfinite(eff["effect"]) and np.isfinite(eff["se"]) and eff["se"] > 0 + ] + assert len(finite_ses) > 0, "No event-study periods have finite replicate SE" def test_imputation_group_replicate(self): """Group SEs should use replicate variance.""" From cfa36afc0efe3e2f835cee96f9080c673a0103c5 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 06:37:00 -0400 Subject: [PATCH 05/11] Address AI review round 4: ImputationDiD balance_e filter and per-effect df - Event-study replicate refits now apply the same balanced-cohort mask as _aggregate_event_study() when balance_e is set, so replicate SEs match the reported estimand - Per-effect replicate refits now capture n_valid and compute effect-specific df (n_valid-1 when replicates dropped) instead of reusing the overall ATT df Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/imputation.py | 31 +++++++++++++++++++++++++++---- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index efb37352a..c6ffce076 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -591,6 +591,18 @@ def _refit_imp_overall(w_r): if np.isfinite(group_effects[g]["effect"]) ) + # Pre-compute balanced cohort mask for balance_e (matches + # _aggregate_event_study's filter so replicate SEs match + # the reported point estimate's estimand). + _balanced_mask_treated = None + if balance_e is not None and _sorted_rel_times: + cohort_rel_times = self._build_cohort_rel_times(df, first_treat) + all_horizons = [int(e) for e in _sorted_rel_times if e >= 0] + df_1 = df.loc[omega_1_mask] + _balanced_mask_treated = self._compute_balanced_cohort_mask( + df_1, first_treat, all_horizons, balance_e, cohort_rel_times + ) + for e in _sorted_rel_times: def _refit_es(w_r, _e=e): ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( @@ -602,16 +614,23 @@ def _refit_es(w_r, _e=e): ) fin = np.isfinite(tau_r) mask_e = fin & (_rel_times_treated == _e) + # Apply balance_e cohort filter (same as _aggregate_event_study) + if _balanced_mask_treated is not None: + mask_e = mask_e & _balanced_mask_treated tw_e = w_r[omega_1_mask.values][mask_e] s = np.sum(tw_e) return np.array([float(np.sum(tau_r[mask_e] * tw_e) / s) if s > 0 else np.nan]) eff_e = event_study_effects[e]["effect"] - vcov_e, _ = compute_replicate_refit_variance( + vcov_e, nv_e = compute_replicate_refit_variance( _refit_es, np.array([eff_e]), resolved_survey ) se_e = float(np.sqrt(max(vcov_e[0, 0], 0.0))) - t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df) + # Use effect-specific df when replicates were dropped + df_e = _survey_df + if nv_e < resolved_survey.n_replicates: + df_e = nv_e - 1 if nv_e > 1 else 0 + t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=df_e) event_study_effects[e]["se"] = se_e event_study_effects[e]["t_stat"] = t_e event_study_effects[e]["p_value"] = p_e @@ -633,11 +652,15 @@ def _refit_grp(w_r, _g=g): return np.array([float(np.sum(tau_r[mask_g] * tw_g) / s) if s > 0 else np.nan]) eff_g = group_effects[g]["effect"] - vcov_g, _ = compute_replicate_refit_variance( + vcov_g, nv_g = compute_replicate_refit_variance( _refit_grp, np.array([eff_g]), resolved_survey ) se_g = float(np.sqrt(max(vcov_g[0, 0], 0.0))) - t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df) + # Use effect-specific df when replicates were dropped + df_g = _survey_df + if nv_g < resolved_survey.n_replicates: + df_g = nv_g - 1 if nv_g > 1 else 0 + t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=df_g) group_effects[g]["se"] = se_g group_effects[g]["t_stat"] = t_g group_effects[g]["p_value"] = p_g From 9c80b95e9a516e42427a13eea8b89a73ddd287ae Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 06:56:57 -0400 Subject: [PATCH 06/11] Address AI review round 5: SunAbraham full IW refit, rank-deficient nuisance, balance_e horizon set - SunAbraham: replicate refit now returns fully re-aggregated event-study and overall ATT effects (recomputes IW cohort-share weights from w_r via _compute_iw_effects/_compute_overall_att), not just vcov_cohort - DiD/MultiPeriodDiD/TWFE absorb paths: detect NaN nuisance coefficients from rank-deficient full-sample fit, refit only identified columns, compute reduced replicate vcov, expand back with _expand_vcov_with_nan - ImputationDiD: balance_e mask now uses full treated horizon set (before Prop 5 filtering), matching _aggregate_event_study() exactly Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/estimators.py | 22 +++++++--- diff_diff/imputation.py | 19 ++++++--- diff_diff/sun_abraham.py | 91 ++++++++++++++++++++++++++++++---------- diff_diff/twfe.py | 11 +++-- 4 files changed, 105 insertions(+), 38 deletions(-) diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index 7e1b2152e..dba8bf98d 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -390,6 +390,11 @@ def fit( _absorb_list = list(absorbed_vars) # capture for closure + # Handle rank-deficient nuisance: refit only identified columns + _id_mask = ~np.isnan(coefficients) + _id_cols = np.where(_id_mask)[0] + _att_idx_reduced = int(np.searchsorted(_id_cols, att_idx)) + def _refit_did_absorb(w_r): wd = data.copy() wd["_treat_time"] = ( @@ -407,15 +412,16 @@ def _refit_did_absorb(w_r): for cov in covariates: X_r = np.column_stack([X_r, wd[cov].values.astype(float)]) coef_r, _, _ = solve_ols( - X_r, y_r, + X_r[:, _id_cols], y_r, weights=w_r, weight_type=survey_weight_type, rank_deficient_action="silent", return_vcov=False, ) return coef_r - vcov, _n_valid_rep = compute_replicate_refit_variance( - _refit_did_absorb, coefficients, resolved_survey + vcov_reduced, _n_valid_rep = compute_replicate_refit_variance( + _refit_did_absorb, coefficients[_id_mask], resolved_survey ) + vcov = _expand_vcov_with_nan(vcov_reduced, len(coefficients), _id_cols) se = float(np.sqrt(max(vcov[att_idx, att_idx], 0.0))) _df_rep = ( survey_metadata.df_survey @@ -1235,6 +1241,9 @@ def fit( # type: ignore[override] from diff_diff.survey import compute_replicate_refit_variance _absorb_list_mp = list(absorb) + # Handle rank-deficient nuisance: refit only identified columns + _id_mask_mp = ~np.isnan(coefficients) + _id_cols_mp = np.where(_id_mask_mp)[0] def _refit_mp_absorb(w_r): wd = data.copy() @@ -1267,15 +1276,16 @@ def _refit_mp_absorb(w_r): for cov_ in covariates: X_r = np.column_stack([X_r, wd[cov_].values.astype(float)]) coef_r, _, _ = solve_ols( - X_r, y_r, + X_r[:, _id_cols_mp], y_r, weights=w_r, weight_type=survey_weight_type, rank_deficient_action="silent", return_vcov=False, ) return coef_r - vcov, _n_valid_rep_mp = compute_replicate_refit_variance( - _refit_mp_absorb, coefficients, resolved_survey + vcov_reduced_mp, _n_valid_rep_mp = compute_replicate_refit_variance( + _refit_mp_absorb, coefficients[_id_mask_mp], resolved_survey ) + vcov = _expand_vcov_with_nan(vcov_reduced_mp, len(coefficients), _id_cols_mp) elif _use_survey_vcov and _uses_replicate_mp: # No absorb + replicate: X is fixed, use compute_replicate_vcov directly from diff_diff.survey import compute_replicate_vcov diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index c6ffce076..aebb5202d 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -591,16 +591,23 @@ def _refit_imp_overall(w_r): if np.isfinite(group_effects[g]["effect"]) ) - # Pre-compute balanced cohort mask for balance_e (matches - # _aggregate_event_study's filter so replicate SEs match - # the reported point estimate's estimand). + # Pre-compute balanced cohort mask for balance_e. Use the + # full treated horizon set (before Prop 5 filtering), matching + # exactly what _aggregate_event_study() does at line 1591-1602. _balanced_mask_treated = None if balance_e is not None and _sorted_rel_times: - cohort_rel_times = self._build_cohort_rel_times(df, first_treat) - all_horizons = [int(e) for e in _sorted_rel_times if e >= 0] df_1 = df.loc[omega_1_mask] + rel_times_all = df_1["_rel_time"].values + all_horizons_full = sorted( + set(int(h) for h in rel_times_all if np.isfinite(h)) + ) + if self.horizon_max is not None: + all_horizons_full = [ + h for h in all_horizons_full if abs(h) <= self.horizon_max + ] + cohort_rel_times = self._build_cohort_rel_times(df, first_treat) _balanced_mask_treated = self._compute_balanced_cohort_mask( - df_1, first_treat, all_horizons, balance_e, cohort_rel_times + df_1, first_treat, all_horizons_full, balance_e, cohort_rel_times ) for e in _sorted_rel_times: diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index 997a57927..eff2b75c8 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -633,35 +633,41 @@ def fit( resolved_survey=None if _uses_replicate_sa else resolved_survey, ) - # Replicate variance override: re-run saturated regression per replicate - # and replace vcov_cohort with the replicate version. The downstream - # delta-method aggregation (_compute_iw_effects, _compute_overall_att) - # then automatically produces correct replicate-based SEs. + # Replicate variance override: fully refit the IW estimator per + # replicate, including recomputing cohort-share aggregation weights + # from w_r, so replicate SEs reflect the complete estimator. + _n_valid_rep_sa = None if _uses_replicate_sa: from diff_diff.survey import compute_replicate_refit_variance - _keys_ordered = sorted(coef_index_map.keys(), key=lambda k: coef_index_map[k]) - _full_cohort_vec = np.array([cohort_effects.get(k, np.nan) for k in _keys_ordered]) + # The refit returns [overall_att, es_e0, es_e1, ...] after + # full re-aggregation with replicate-weighted cohort shares. + _sa_rel_periods = list(rel_periods_to_estimate) def _refit_sa(w_r): - ce_r, _, _, cim_r = self._fit_saturated_regression( + ce_r, _, vcov_r, cim_r = self._fit_saturated_regression( df_reg, outcome, unit, time, first_treat, - treatment_groups, rel_periods_to_estimate, covariates, + treatment_groups, _sa_rel_periods, covariates, cluster_var, survey_weights=w_r, survey_weight_type=survey_weight_type, - resolved_survey=None, # prevent internal replicate dispatch + resolved_survey=None, ) - return np.array([ce_r.get(k, np.nan) for k in _keys_ordered]) - - vcov_cohort, _n_valid_rep_sa = compute_replicate_refit_variance( - _refit_sa, _full_cohort_vec, resolved_survey - ) - - # Recompute cohort_ses from the replicate vcov diagonal (the - # initial fit produced stale SEs from the non-replicate path) - for key in _keys_ordered: - idx = coef_index_map[key] - cohort_ses[key] = float(np.sqrt(max(vcov_cohort[idx, idx], 0.0))) + # Create temp weight column for IW aggregation with w_r + _wt_col = "_rep_wt" + df[_wt_col] = 0.0 + df.loc[df.index[:len(w_r)], _wt_col] = w_r + es_r, _ = self._compute_iw_effects( + df, unit, first_treat, treatment_groups, _sa_rel_periods, + ce_r, {}, vcov_r, cim_r, survey_weight_col=_wt_col, + ) + att_r, _ = self._compute_overall_att( + df, first_treat, es_r, ce_r, _, vcov_r, cim_r, + survey_weight_col=_wt_col, + ) + results = [att_r] + for e in _sa_rel_periods: + results.append(es_r[e]["effect"] if e in es_r else np.nan) + return np.array(results) # Resolve survey weight column name for cohort aggregation survey_weight_col = ( @@ -678,11 +684,10 @@ def _refit_sa(w_r): if survey_metadata is not None and survey_metadata.df_survey is not None else None ) - # Replicate df: rank-deficient → NaN inference; dropped replicates → n_valid-1 + # Replicate df: rank-deficient → NaN inference (dropped-replicate + # override happens after replicate refit below) if _uses_replicate_sa and _sa_survey_df is None: _sa_survey_df = 0 # rank-deficient replicate → NaN inference - if _uses_replicate_sa and _n_valid_rep_sa < resolved_survey.n_replicates: - _sa_survey_df = _n_valid_rep_sa - 1 if _n_valid_rep_sa > 1 else 0 # Compute interaction-weighted event study effects event_study_effects, cohort_weights = self._compute_iw_effects( @@ -715,6 +720,46 @@ def _refit_sa(w_r): overall_att, overall_se, alpha=self.alpha, df=_sa_survey_df ) + # Replicate variance override: refit fully re-aggregated estimates + if _uses_replicate_sa: + # Build full-sample estimate vector from actual outputs + _full_est_sa = [overall_att] + for e in _sa_rel_periods: + _full_est_sa.append( + event_study_effects[e]["effect"] if e in event_study_effects else np.nan + ) + + _vcov_sa, _n_valid_rep_sa = compute_replicate_refit_variance( + _refit_sa, np.array(_full_est_sa), resolved_survey + ) + + # Override df if replicates dropped + if _n_valid_rep_sa < resolved_survey.n_replicates: + _sa_survey_df = _n_valid_rep_sa - 1 if _n_valid_rep_sa > 1 else 0 + + # Override overall ATT SE + overall_se = float(np.sqrt(max(_vcov_sa[0, 0], 0.0))) + overall_t, overall_p, overall_ci = safe_inference( + overall_att, overall_se, alpha=self.alpha, df=_sa_survey_df + ) + + # Override event-study SEs + for i, e in enumerate(_sa_rel_periods): + if e in event_study_effects and np.isfinite(event_study_effects[e]["effect"]): + se_e = float(np.sqrt(max(_vcov_sa[1 + i, 1 + i], 0.0))) + eff_e = event_study_effects[e]["effect"] + t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_sa_survey_df) + event_study_effects[e]["se"] = se_e + event_study_effects[e]["t_stat"] = t_e + event_study_effects[e]["p_value"] = p_e + event_study_effects[e]["conf_int"] = ci_e + + # Recompute cohort_ses from the replicate vcov of the raw cohort + # coefficients (not from the aggregated vcov) — but since we now + # refit the full estimator, cohort SEs are best left from the + # analytical path. Mark as 0.0 to avoid stale values. + # (Proper cohort-level replicate SEs would need a separate refit.) + # Run bootstrap if requested bootstrap_results = None if self.n_bootstrap > 0: diff --git a/diff_diff/twfe.py b/diff_diff/twfe.py index a85457732..65eb0b44f 100644 --- a/diff_diff/twfe.py +++ b/diff_diff/twfe.py @@ -293,6 +293,9 @@ def fit( # type: ignore[override] _all_vars_twfe = list(all_vars) _covariates_twfe = list(covariates) if covariates else [] + # Handle rank-deficient nuisance: refit only identified columns + _id_mask_twfe = ~np.isnan(coefficients) + _id_cols_twfe = np.where(_id_mask_twfe)[0] def _refit_twfe(w_r): data_dem_r = _within_transform_util( @@ -304,15 +307,17 @@ def _refit_twfe(w_r): X_list_r.append(data_dem_r[f"{cov_}_demeaned"].values) X_r = np.column_stack([np.ones(len(y_r))] + X_list_r) coef_r, _, _ = solve_ols( - X_r, y_r, + X_r[:, _id_cols_twfe], y_r, weights=w_r, weight_type=survey_weight_type, rank_deficient_action="silent", return_vcov=False, ) return coef_r - vcov, _n_valid_rep_twfe = compute_replicate_refit_variance( - _refit_twfe, coefficients, resolved_survey + from diff_diff.linalg import _expand_vcov_with_nan as _expand_twfe + vcov_reduced, _n_valid_rep_twfe = compute_replicate_refit_variance( + _refit_twfe, coefficients[_id_mask_twfe], resolved_survey ) + vcov = _expand_twfe(vcov_reduced, len(coefficients), _id_cols_twfe) se = float(np.sqrt(max(vcov[att_idx, att_idx], 0.0))) _df_rep = ( survey_metadata.df_survey From 38427d4b7a1819db2a558da895287a3fe02dee93 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 09:33:23 -0400 Subject: [PATCH 07/11] Address AI review round 6: SunAbraham cohort-level replicate SEs Add second replicate refit loop for raw (g,e) cohort coefficients so cohort_ses in results.cohort_effects uses replicate-based variance instead of stale analytical SEs from the initial saturated regression. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/sun_abraham.py | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index eff2b75c8..2b504a8a9 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -754,11 +754,26 @@ def _refit_sa(w_r): event_study_effects[e]["p_value"] = p_e event_study_effects[e]["conf_int"] = ci_e - # Recompute cohort_ses from the replicate vcov of the raw cohort - # coefficients (not from the aggregated vcov) — but since we now - # refit the full estimator, cohort SEs are best left from the - # analytical path. Mark as 0.0 to avoid stale values. - # (Proper cohort-level replicate SEs would need a separate refit.) + # Cohort-level replicate SEs: second refit for raw (g,e) coefficients + _keys_ordered = sorted(coef_index_map.keys(), key=lambda k: coef_index_map[k]) + _full_cohort_vec = np.array([cohort_effects.get(k, np.nan) for k in _keys_ordered]) + + def _refit_sa_cohort(w_r): + ce_r, _, _, _ = self._fit_saturated_regression( + df_reg, outcome, unit, time, first_treat, + treatment_groups, _sa_rel_periods, covariates, + cluster_var, survey_weights=w_r, + survey_weight_type=survey_weight_type, + resolved_survey=None, + ) + return np.array([ce_r.get(k, np.nan) for k in _keys_ordered]) + + _vcov_cohort_rep, _ = compute_replicate_refit_variance( + _refit_sa_cohort, _full_cohort_vec, resolved_survey + ) + for key in _keys_ordered: + idx = coef_index_map[key] + cohort_ses[key] = float(np.sqrt(max(_vcov_cohort_rep[idx, idx], 0.0))) # Run bootstrap if requested bootstrap_results = None From 5375c704cb9dde24471f98f468d89fba4c8f791f Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 09:48:47 -0400 Subject: [PATCH 08/11] Address P2s: skip ref-period in ImputationDiD replicate loop, add absorb test - ImputationDiD: exclude synthetic reference period (n_obs==0) from replicate event-study override loop to avoid spurious NaN warnings - Add DiD absorb-path replicate test (BRR + absorbed group FE) - TODO.md: track BRR test gap (Fay-like perturbations, not true half-sample) as low-priority deferred item Co-Authored-By: Claude Opus 4.6 (1M context) --- TODO.md | 1 + diff_diff/imputation.py | 4 +++- tests/test_replicate_weight_expansion.py | 17 +++++++++++++++++ 3 files changed, 21 insertions(+), 1 deletion(-) diff --git a/TODO.md b/TODO.md index aafe0f4e9..0b60cb7c2 100644 --- a/TODO.md +++ b/TODO.md @@ -69,6 +69,7 @@ Deferred items from PR reviews that were not addressed before merge. | StaggeredTripleDifference R parity: benchmark only tests no-covariate path (xformla=~1). Add covariate-adjusted scenarios and aggregation SE parity assertions. | `benchmarks/R/benchmark_staggered_triplediff.R` | #245 | Medium | | StaggeredTripleDifference: per-cohort group-effect SEs include WIF (conservative vs R's wif=NULL). Documented in REGISTRY. Could override mixin for exact R match. | `staggered_triple_diff.py` | #245 | Low | | HonestDiD Delta^RM: uses naive FLCI instead of paper's ARP conditional/hybrid confidence sets (Sections 3.2.1-3.2.2). ARP infrastructure exists but moment inequality transformation needs calibration. CIs are conservative (wider, valid coverage). | `honest_did.py` | #248 | Medium | +| Replicate weight tests use Fay-like BRR perturbations (0.5/1.5), not true half-sample BRR. Add true BRR regressions per estimator family. Existing `test_survey_phase6.py` covers true BRR at the helper level. | `tests/test_replicate_weight_expansion.py` | #253 | Low | #### Performance diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index aebb5202d..cb2d3d00f 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -581,10 +581,12 @@ def _refit_imp_overall(w_r): ) # Event-study/group replicate SEs: per-effect scalar refits - # Derive keys from actual outputs (excludes filtered/Prop5 horizons) + # Derive keys from actual outputs (excludes filtered/Prop5 horizons + # and the synthetic reference period marker with n_obs==0) _sorted_rel_times = sorted( e for e in (event_study_effects or {}).keys() if np.isfinite(event_study_effects[e]["effect"]) + and event_study_effects[e].get("n_obs", 1) > 0 ) _sorted_groups = sorted( g for g in (group_effects or {}).keys() diff --git a/tests/test_replicate_weight_expansion.py b/tests/test_replicate_weight_expansion.py index b0cfb11aa..3f6fa4109 100644 --- a/tests/test_replicate_weight_expansion.py +++ b/tests/test_replicate_weight_expansion.py @@ -147,6 +147,23 @@ def test_did_wild_bootstrap_rejected(self): ) +class TestDiDAbsorbReplicate: + """DiD absorb path with replicate weights.""" + + def test_did_absorb_brr(self): + """Absorb path should produce finite replicate SE.""" + data = _make_simple_panel() + # Add a group column for absorb + data["group"] = (data["unit"] % 4).astype(str) + rep_cols = _add_brr_replicates(data, n_rep=16) + sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") + result = DifferenceInDifferences().fit( + data, "outcome", "treated", "post", absorb=["group"], survey_design=sd, + ) + assert np.isfinite(result.att) + assert np.isfinite(result.se) and result.se > 0 + + class TestMultiPeriodDiDReplicate: """MultiPeriodDiD with replicate weights.""" From c806683fbe1d43d6541a72661fd4fa8419ba42d8 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 10:13:46 -0400 Subject: [PATCH 09/11] Write effective replicate df back to survey_metadata.df_survey When dropped replicates reduce the effective df, write it back to survey_metadata.df_survey so summary()/exported metadata is consistent with the p-values/CIs actually reported. Matches existing pattern in ContinuousDiD and TripleDifference. Applied to all 7 new estimator replicate paths. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/estimators.py | 4 ++++ diff_diff/imputation.py | 2 ++ diff_diff/stacked_did.py | 4 ++++ diff_diff/sun_abraham.py | 2 ++ diff_diff/twfe.py | 2 ++ diff_diff/two_stage.py | 2 ++ 6 files changed, 16 insertions(+) diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index dba8bf98d..6128ff5d7 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -430,6 +430,8 @@ def _refit_did_absorb(w_r): ) if _n_valid_rep < resolved_survey.n_replicates: _df_rep = _n_valid_rep - 1 if _n_valid_rep > 1 else 0 + if survey_metadata is not None and _df_rep > 0: + survey_metadata.df_survey = _df_rep t_stat, p_value, conf_int = safe_inference( att, se, alpha=self.alpha, df=_df_rep ) @@ -1336,6 +1338,8 @@ def _refit_mp_absorb(w_r): df = 0 # rank-deficient replicate → NaN inference if _n_valid_rep_mp is not None and _n_valid_rep_mp < resolved_survey.n_replicates: df = _n_valid_rep_mp - 1 if _n_valid_rep_mp > 1 else 0 + if survey_metadata is not None and df > 0: + survey_metadata.df_survey = df # Guard: fall back to normal distribution if df is non-positive # Skip for replicate designs — df=0 is intentional for NaN inference diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index cb2d3d00f..67ec739f9 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -574,6 +574,8 @@ def _refit_imp_overall(w_r): # Override df if replicates were dropped if _n_valid_rep_imp < resolved_survey.n_replicates: _survey_df = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0 + if survey_metadata is not None and _survey_df is not None and _survey_df > 0: + survey_metadata.df_survey = _survey_df # Recompute overall inference with replicate SE/df overall_t, overall_p, overall_ci = safe_inference( diff --git a/diff_diff/stacked_did.py b/diff_diff/stacked_did.py index 4a2a87859..48a8a526f 100644 --- a/diff_diff/stacked_did.py +++ b/diff_diff/stacked_did.py @@ -522,6 +522,8 @@ def _refit_stacked(w_r): if _n_valid_rep_sd is not None and resolved_stacked is not None: if _n_valid_rep_sd < resolved_stacked.n_replicates: _survey_df = _n_valid_rep_sd - 1 if _n_valid_rep_sd > 1 else 0 + if survey_metadata is not None and _survey_df > 0: + survey_metadata.df_survey = _survey_df t_stat, p_value, conf_int = safe_inference( effect, se, alpha=self.alpha, df=_survey_df ) @@ -562,6 +564,8 @@ def _refit_stacked(w_r): if _n_valid_rep_sd is not None and resolved_stacked is not None: if _n_valid_rep_sd < resolved_stacked.n_replicates: _survey_df_overall = _n_valid_rep_sd - 1 if _n_valid_rep_sd > 1 else 0 + if survey_metadata is not None and _survey_df_overall > 0: + survey_metadata.df_survey = _survey_df_overall overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=_survey_df_overall ) diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index 2b504a8a9..8581f41be 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -736,6 +736,8 @@ def _refit_sa(w_r): # Override df if replicates dropped if _n_valid_rep_sa < resolved_survey.n_replicates: _sa_survey_df = _n_valid_rep_sa - 1 if _n_valid_rep_sa > 1 else 0 + if survey_metadata is not None and _sa_survey_df is not None and _sa_survey_df > 0: + survey_metadata.df_survey = _sa_survey_df # Override overall ATT SE overall_se = float(np.sqrt(max(_vcov_sa[0, 0], 0.0))) diff --git a/diff_diff/twfe.py b/diff_diff/twfe.py index 65eb0b44f..cfddfc95d 100644 --- a/diff_diff/twfe.py +++ b/diff_diff/twfe.py @@ -326,6 +326,8 @@ def _refit_twfe(w_r): ) if _n_valid_rep_twfe < resolved_survey.n_replicates: _df_rep = _n_valid_rep_twfe - 1 if _n_valid_rep_twfe > 1 else 0 + if survey_metadata is not None and _df_rep > 0: + survey_metadata.df_survey = _df_rep t_stat, p_value, conf_int = _safe_inf(att, se, alpha=self.alpha, df=_df_rep) elif self.inference == "wild_bootstrap": # Override with wild cluster bootstrap inference diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index 6d1a7cb64..43f483ec5 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -620,6 +620,8 @@ def _refit_ts(w_r): # Override df if replicates were dropped if _n_valid_rep_ts < resolved_survey.n_replicates: _survey_df = _n_valid_rep_ts - 1 if _n_valid_rep_ts > 1 else 0 + if survey_metadata is not None and _survey_df is not None and _survey_df > 0: + survey_metadata.df_survey = _survey_df # Recompute overall inference with replicate SE/df overall_t, overall_p, overall_ci = safe_inference( From fba8ca3b220c906aa145f55f5ce31440eb92705c Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 10:33:44 -0400 Subject: [PATCH 10/11] Address P1/P2: zero-weight demeaning safety, df_survey writeback for df=0 P1: Replicate refit closures for TWFE, SunAbraham, DiD/MultiPeriodDiD absorb paths now drop zero-weight observations before weighted demeaning, preventing division-by-zero in within-transformation group means. Matches R's survey::withReplicates() convention. Documented in REGISTRY.md. P2: survey_metadata.df_survey writeback now covers df=0 case (set to None) so summary/metadata is consistent when all replicates fail. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/estimators.py | 24 ++++++++++++++---------- diff_diff/imputation.py | 4 ++-- diff_diff/stacked_did.py | 8 ++++---- diff_diff/sun_abraham.py | 19 +++++++++++++------ diff_diff/twfe.py | 13 +++++++++---- diff_diff/two_stage.py | 4 ++-- docs/methodology/REGISTRY.md | 7 +++++++ 7 files changed, 51 insertions(+), 28 deletions(-) diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index 6128ff5d7..61695d04a 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -396,13 +396,15 @@ def fit( _att_idx_reduced = int(np.searchsorted(_id_cols, att_idx)) def _refit_did_absorb(w_r): - wd = data.copy() + nz = w_r > 0 + wd = data[nz].copy() + w_nz = w_r[nz] wd["_treat_time"] = ( wd[treatment].values.astype(float) * wd[time].values.astype(float) ) vars_dm = [outcome, treatment, time, "_treat_time"] + (covariates or []) for ab_var in _absorb_list: - wd, _ = demean_by_group(wd, vars_dm, ab_var, inplace=True, weights=w_r) + wd, _ = demean_by_group(wd, vars_dm, ab_var, inplace=True, weights=w_nz) y_r = wd[outcome].values.astype(float) d_r = wd[treatment].values.astype(float) t_r = wd[time].values.astype(float) @@ -413,7 +415,7 @@ def _refit_did_absorb(w_r): X_r = np.column_stack([X_r, wd[cov].values.astype(float)]) coef_r, _, _ = solve_ols( X_r[:, _id_cols], y_r, - weights=w_r, weight_type=survey_weight_type, + weights=w_nz, weight_type=survey_weight_type, rank_deficient_action="silent", return_vcov=False, ) return coef_r @@ -430,8 +432,8 @@ def _refit_did_absorb(w_r): ) if _n_valid_rep < resolved_survey.n_replicates: _df_rep = _n_valid_rep - 1 if _n_valid_rep > 1 else 0 - if survey_metadata is not None and _df_rep > 0: - survey_metadata.df_survey = _df_rep + if survey_metadata is not None: + survey_metadata.df_survey = _df_rep if _df_rep > 0 else None t_stat, p_value, conf_int = safe_inference( att, se, alpha=self.alpha, df=_df_rep ) @@ -1248,7 +1250,9 @@ def fit( # type: ignore[override] _id_cols_mp = np.where(_id_mask_mp)[0] def _refit_mp_absorb(w_r): - wd = data.copy() + nz = w_r > 0 + wd = data[nz].copy() + w_nz = w_r[nz] d_raw_ = wd[treatment].values.astype(float) t_raw_ = wd[time].values wd["_did_treatment"] = d_raw_ @@ -1262,7 +1266,7 @@ def _refit_mp_absorb(w_r): + (covariates or []) ) for ab_var_ in _absorb_list_mp: - wd, _ = demean_by_group(wd, vars_dm_, ab_var_, inplace=True, weights=w_r) + wd, _ = demean_by_group(wd, vars_dm_, ab_var_, inplace=True, weights=w_nz) y_r = wd[outcome].values.astype(float) d_r = wd["_did_treatment"].values.astype(float) X_r = np.column_stack([np.ones(len(y_r)), d_r]) @@ -1279,7 +1283,7 @@ def _refit_mp_absorb(w_r): X_r = np.column_stack([X_r, wd[cov_].values.astype(float)]) coef_r, _, _ = solve_ols( X_r[:, _id_cols_mp], y_r, - weights=w_r, weight_type=survey_weight_type, + weights=w_nz, weight_type=survey_weight_type, rank_deficient_action="silent", return_vcov=False, ) return coef_r @@ -1338,8 +1342,8 @@ def _refit_mp_absorb(w_r): df = 0 # rank-deficient replicate → NaN inference if _n_valid_rep_mp is not None and _n_valid_rep_mp < resolved_survey.n_replicates: df = _n_valid_rep_mp - 1 if _n_valid_rep_mp > 1 else 0 - if survey_metadata is not None and df > 0: - survey_metadata.df_survey = df + if survey_metadata is not None: + survey_metadata.df_survey = df if df > 0 else None # Guard: fall back to normal distribution if df is non-positive # Skip for replicate designs — df=0 is intentional for NaN inference diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 67ec739f9..702650e01 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -574,8 +574,8 @@ def _refit_imp_overall(w_r): # Override df if replicates were dropped if _n_valid_rep_imp < resolved_survey.n_replicates: _survey_df = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0 - if survey_metadata is not None and _survey_df is not None and _survey_df > 0: - survey_metadata.df_survey = _survey_df + if survey_metadata is not None: + survey_metadata.df_survey = _survey_df if _survey_df and _survey_df > 0 else None # Recompute overall inference with replicate SE/df overall_t, overall_p, overall_ci = safe_inference( diff --git a/diff_diff/stacked_did.py b/diff_diff/stacked_did.py index 48a8a526f..7c610b5c0 100644 --- a/diff_diff/stacked_did.py +++ b/diff_diff/stacked_did.py @@ -522,8 +522,8 @@ def _refit_stacked(w_r): if _n_valid_rep_sd is not None and resolved_stacked is not None: if _n_valid_rep_sd < resolved_stacked.n_replicates: _survey_df = _n_valid_rep_sd - 1 if _n_valid_rep_sd > 1 else 0 - if survey_metadata is not None and _survey_df > 0: - survey_metadata.df_survey = _survey_df + if survey_metadata is not None: + survey_metadata.df_survey = _survey_df if _survey_df > 0 else None t_stat, p_value, conf_int = safe_inference( effect, se, alpha=self.alpha, df=_survey_df ) @@ -564,8 +564,8 @@ def _refit_stacked(w_r): if _n_valid_rep_sd is not None and resolved_stacked is not None: if _n_valid_rep_sd < resolved_stacked.n_replicates: _survey_df_overall = _n_valid_rep_sd - 1 if _n_valid_rep_sd > 1 else 0 - if survey_metadata is not None and _survey_df_overall > 0: - survey_metadata.df_survey = _survey_df_overall + if survey_metadata is not None: + survey_metadata.df_survey = _survey_df_overall if _survey_df_overall > 0 else None overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=_survey_df_overall ) diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index 8581f41be..e4fdf1b02 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -645,10 +645,14 @@ def fit( _sa_rel_periods = list(rel_periods_to_estimate) def _refit_sa(w_r): + # Drop zero-weight obs for within-transform safety + nz = w_r > 0 + df_reg_nz = df_reg[nz] if not np.all(nz) else df_reg + w_nz = w_r[nz] if not np.all(nz) else w_r ce_r, _, vcov_r, cim_r = self._fit_saturated_regression( - df_reg, outcome, unit, time, first_treat, + df_reg_nz, outcome, unit, time, first_treat, treatment_groups, _sa_rel_periods, covariates, - cluster_var, survey_weights=w_r, + cluster_var, survey_weights=w_nz, survey_weight_type=survey_weight_type, resolved_survey=None, ) @@ -736,8 +740,8 @@ def _refit_sa(w_r): # Override df if replicates dropped if _n_valid_rep_sa < resolved_survey.n_replicates: _sa_survey_df = _n_valid_rep_sa - 1 if _n_valid_rep_sa > 1 else 0 - if survey_metadata is not None and _sa_survey_df is not None and _sa_survey_df > 0: - survey_metadata.df_survey = _sa_survey_df + if survey_metadata is not None: + survey_metadata.df_survey = _sa_survey_df if _sa_survey_df and _sa_survey_df > 0 else None # Override overall ATT SE overall_se = float(np.sqrt(max(_vcov_sa[0, 0], 0.0))) @@ -761,10 +765,13 @@ def _refit_sa(w_r): _full_cohort_vec = np.array([cohort_effects.get(k, np.nan) for k in _keys_ordered]) def _refit_sa_cohort(w_r): + nz = w_r > 0 + df_reg_nz = df_reg[nz] if not np.all(nz) else df_reg + w_nz = w_r[nz] if not np.all(nz) else w_r ce_r, _, _, _ = self._fit_saturated_regression( - df_reg, outcome, unit, time, first_treat, + df_reg_nz, outcome, unit, time, first_treat, treatment_groups, _sa_rel_periods, covariates, - cluster_var, survey_weights=w_r, + cluster_var, survey_weights=w_nz, survey_weight_type=survey_weight_type, resolved_survey=None, ) diff --git a/diff_diff/twfe.py b/diff_diff/twfe.py index cfddfc95d..aa420a930 100644 --- a/diff_diff/twfe.py +++ b/diff_diff/twfe.py @@ -298,8 +298,13 @@ def fit( # type: ignore[override] _id_cols_twfe = np.where(_id_mask_twfe)[0] def _refit_twfe(w_r): + # Drop zero-weight obs to prevent zero-sum demeaning + # (JK1/BRR half-samples zero entire clusters) + nz = w_r > 0 + data_nz = data[nz].copy() + w_nz = w_r[nz] data_dem_r = _within_transform_util( - data, _all_vars_twfe, unit, time, suffix="_demeaned", weights=w_r, + data_nz, _all_vars_twfe, unit, time, suffix="_demeaned", weights=w_nz, ) y_r = data_dem_r[f"{outcome}_demeaned"].values X_list_r = [data_dem_r["_treatment_post_demeaned"].values] @@ -308,7 +313,7 @@ def _refit_twfe(w_r): X_r = np.column_stack([np.ones(len(y_r))] + X_list_r) coef_r, _, _ = solve_ols( X_r[:, _id_cols_twfe], y_r, - weights=w_r, weight_type=survey_weight_type, + weights=w_nz, weight_type=survey_weight_type, rank_deficient_action="silent", return_vcov=False, ) return coef_r @@ -326,8 +331,8 @@ def _refit_twfe(w_r): ) if _n_valid_rep_twfe < resolved_survey.n_replicates: _df_rep = _n_valid_rep_twfe - 1 if _n_valid_rep_twfe > 1 else 0 - if survey_metadata is not None and _df_rep > 0: - survey_metadata.df_survey = _df_rep + if survey_metadata is not None: + survey_metadata.df_survey = _df_rep if _df_rep > 0 else None t_stat, p_value, conf_int = _safe_inf(att, se, alpha=self.alpha, df=_df_rep) elif self.inference == "wild_bootstrap": # Override with wild cluster bootstrap inference diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index 43f483ec5..28f7eb3cf 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -620,8 +620,8 @@ def _refit_ts(w_r): # Override df if replicates were dropped if _n_valid_rep_ts < resolved_survey.n_replicates: _survey_df = _n_valid_rep_ts - 1 if _n_valid_rep_ts > 1 else 0 - if survey_metadata is not None and _survey_df is not None and _survey_df > 0: - survey_metadata.df_survey = _survey_df + if survey_metadata is not None: + survey_metadata.df_survey = _survey_df if _survey_df and _survey_df > 0 else None # Recompute overall inference with replicate SE/df overall_t, overall_p, overall_ci = safe_inference( diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index d445db2f9..86a542892 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -2302,6 +2302,13 @@ Domain estimation preserving full design structure. - **Note:** For replicate-weight designs, `subpopulation()` zeros out both full-sample and replicate weight columns for excluded observations, preserving all replicate metadata. +- **Note:** Estimator-level replicate refits (TWFE, SunAbraham, DiD/MultiPeriodDiD + with `absorb`) drop zero-weight observations before weighted demeaning to + prevent division-by-zero in within-transformation group means. This matches + R's `survey::withReplicates()` convention where zero-weight units are excluded + from per-replicate estimation. Replicates that fail despite this (e.g., + rank-deficient after unit deletion) are counted as invalid and excluded from + variance computation. - **Note:** Defensive enhancement: ContinuousDiD and TripleDifference validate the positive-weight effective sample size before WLS cell fits. After `subpopulation()` zeroes weights, raw row counts may exceed the From 692dab4517a04090e6369bc24b7a87144ca6d0bc Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 13:03:26 -0400 Subject: [PATCH 11/11] Address P2/P3: SA positional assign, ImputationDiD vectorized refit, REGISTRY notes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - SunAbraham: replace fragile .loc label-based assignment with positional df["_rep_wt"] = w_r for replicate weight column - ImputationDiD: consolidate per-effect scalar refits into single vectorized refit returning [overall, es_e0..., grp_g0...], reducing from R×(1+H+G) to R full refits for aggregate="all" - REGISTRY.md: add replicate-refit cross-references to SunAbraham and StackedDiD per-estimator survey notes Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/imputation.py | 167 +++++++++++++++-------------------- diff_diff/sun_abraham.py | 4 +- docs/methodology/REGISTRY.md | 4 +- 3 files changed, 77 insertions(+), 98 deletions(-) diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 702650e01..70fc112e1 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -551,40 +551,7 @@ def fit( _rel_times_treated = df.loc[omega_1_mask, "_rel_time"].values _cohorts_treated = df.loc[omega_1_mask, first_treat].values - # Overall ATT replicate variance (scalar, always robust) - def _refit_imp_overall(w_r): - ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( - df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, - ) - tau_r, _ = self._impute_treatment_effects( - df, outcome, unit, time, covariates, omega_1_mask, - ufe_r, tfe_r, gm_r, delta_r, - ) - fin = np.isfinite(tau_r) - tw = w_r[omega_1_mask.values] - tw_fin = tw[fin] - s = np.sum(tw_fin) - return np.array([float(np.sum(tau_r[fin] * tw_fin) / s) if s > 0 else np.nan]) - - _vcov_att, _n_valid_rep_imp = compute_replicate_refit_variance( - _refit_imp_overall, np.array([overall_att]), resolved_survey - ) - overall_se = float(np.sqrt(max(_vcov_att[0, 0], 0.0))) - - # Override df if replicates were dropped - if _n_valid_rep_imp < resolved_survey.n_replicates: - _survey_df = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0 - if survey_metadata is not None: - survey_metadata.df_survey = _survey_df if _survey_df and _survey_df > 0 else None - - # Recompute overall inference with replicate SE/df - overall_t, overall_p, overall_ci = safe_inference( - overall_att, overall_se, alpha=self.alpha, df=_survey_df - ) - - # Event-study/group replicate SEs: per-effect scalar refits - # Derive keys from actual outputs (excludes filtered/Prop5 horizons - # and the synthetic reference period marker with n_obs==0) + # Derive keys from actual outputs (excludes filtered/Prop5/ref) _sorted_rel_times = sorted( e for e in (event_study_effects or {}).keys() if np.isfinite(event_study_effects[e]["effect"]) @@ -594,10 +561,9 @@ def _refit_imp_overall(w_r): g for g in (group_effects or {}).keys() if np.isfinite(group_effects[g]["effect"]) ) + _n_es = len(_sorted_rel_times) - # Pre-compute balanced cohort mask for balance_e. Use the - # full treated horizon set (before Prop 5 filtering), matching - # exactly what _aggregate_event_study() does at line 1591-1602. + # Pre-compute balanced cohort mask for balance_e _balanced_mask_treated = None if balance_e is not None and _sorted_rel_times: df_1 = df.loc[omega_1_mask] @@ -614,68 +580,81 @@ def _refit_imp_overall(w_r): df_1, first_treat, all_horizons_full, balance_e, cohort_rel_times ) - for e in _sorted_rel_times: - def _refit_es(w_r, _e=e): - ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( - df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, - ) - tau_r, _ = self._impute_treatment_effects( - df, outcome, unit, time, covariates, omega_1_mask, - ufe_r, tfe_r, gm_r, delta_r, - ) - fin = np.isfinite(tau_r) - mask_e = fin & (_rel_times_treated == _e) - # Apply balance_e cohort filter (same as _aggregate_event_study) + # Single vectorized refit: [overall, es_e0..., grp_g0...] + def _refit_imp(w_r): + ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( + df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, + ) + tau_r, _ = self._impute_treatment_effects( + df, outcome, unit, time, covariates, omega_1_mask, + ufe_r, tfe_r, gm_r, delta_r, + ) + fin = np.isfinite(tau_r) + treated_w = w_r[omega_1_mask.values] + results = [] + # [0] Overall ATT + tw_fin = treated_w[fin] + tw_sum = np.sum(tw_fin) + results.append( + float(np.sum(tau_r[fin] * tw_fin) / tw_sum) if tw_sum > 0 else np.nan + ) + # [1..n_es] Event-study (identified only) + for e in _sorted_rel_times: + mask_e = fin & (_rel_times_treated == e) if _balanced_mask_treated is not None: mask_e = mask_e & _balanced_mask_treated - tw_e = w_r[omega_1_mask.values][mask_e] + tw_e = treated_w[mask_e] s = np.sum(tw_e) - return np.array([float(np.sum(tau_r[mask_e] * tw_e) / s) if s > 0 else np.nan]) - - eff_e = event_study_effects[e]["effect"] - vcov_e, nv_e = compute_replicate_refit_variance( - _refit_es, np.array([eff_e]), resolved_survey - ) - se_e = float(np.sqrt(max(vcov_e[0, 0], 0.0))) - # Use effect-specific df when replicates were dropped - df_e = _survey_df - if nv_e < resolved_survey.n_replicates: - df_e = nv_e - 1 if nv_e > 1 else 0 - t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=df_e) - event_study_effects[e]["se"] = se_e - event_study_effects[e]["t_stat"] = t_e - event_study_effects[e]["p_value"] = p_e - event_study_effects[e]["conf_int"] = ci_e - - for g in _sorted_groups: - def _refit_grp(w_r, _g=g): - ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( - df, outcome, unit, time, covariates, omega_0_mask, weights=w_r, - ) - tau_r, _ = self._impute_treatment_effects( - df, outcome, unit, time, covariates, omega_1_mask, - ufe_r, tfe_r, gm_r, delta_r, - ) - fin = np.isfinite(tau_r) - mask_g = fin & (_cohorts_treated == _g) - tw_g = w_r[omega_1_mask.values][mask_g] + results.append(float(np.sum(tau_r[mask_e] * tw_e) / s) if s > 0 else np.nan) + # [n_es+1..] Group (identified only) + for g in _sorted_groups: + mask_g = fin & (_cohorts_treated == g) + tw_g = treated_w[mask_g] s = np.sum(tw_g) - return np.array([float(np.sum(tau_r[mask_g] * tw_g) / s) if s > 0 else np.nan]) + results.append(float(np.sum(tau_r[mask_g] * tw_g) / s) if s > 0 else np.nan) + return np.array(results) - eff_g = group_effects[g]["effect"] - vcov_g, nv_g = compute_replicate_refit_variance( - _refit_grp, np.array([eff_g]), resolved_survey - ) - se_g = float(np.sqrt(max(vcov_g[0, 0], 0.0))) - # Use effect-specific df when replicates were dropped - df_g = _survey_df - if nv_g < resolved_survey.n_replicates: - df_g = nv_g - 1 if nv_g > 1 else 0 - t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=df_g) - group_effects[g]["se"] = se_g - group_effects[g]["t_stat"] = t_g - group_effects[g]["p_value"] = p_g - group_effects[g]["conf_int"] = ci_g + # Build full-sample estimate from actual effects + _full_est = [overall_att] + _full_est.extend([event_study_effects[e]["effect"] for e in _sorted_rel_times]) + _full_est.extend([group_effects[g]["effect"] for g in _sorted_groups]) + + _vcov_rep_imp, _n_valid_rep_imp = compute_replicate_refit_variance( + _refit_imp, np.array(_full_est), resolved_survey + ) + overall_se = float(np.sqrt(max(_vcov_rep_imp[0, 0], 0.0))) + + # Override df if replicates were dropped + if _n_valid_rep_imp < resolved_survey.n_replicates: + _survey_df = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0 + if survey_metadata is not None: + survey_metadata.df_survey = _survey_df if _survey_df and _survey_df > 0 else None + + overall_t, overall_p, overall_ci = safe_inference( + overall_att, overall_se, alpha=self.alpha, df=_survey_df + ) + + # Override event-study SEs from vcov diagonal + for i, e in enumerate(_sorted_rel_times): + if event_study_effects is not None and e in event_study_effects: + se_e = float(np.sqrt(max(_vcov_rep_imp[1 + i, 1 + i], 0.0))) + eff_e = event_study_effects[e]["effect"] + t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df) + event_study_effects[e]["se"] = se_e + event_study_effects[e]["t_stat"] = t_e + event_study_effects[e]["p_value"] = p_e + event_study_effects[e]["conf_int"] = ci_e + + # Override group SEs from vcov diagonal + for j, g in enumerate(_sorted_groups): + if group_effects is not None and g in group_effects: + se_g = float(np.sqrt(max(_vcov_rep_imp[1 + _n_es + j, 1 + _n_es + j], 0.0))) + eff_g = group_effects[g]["effect"] + t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df) + group_effects[g]["se"] = se_g + group_effects[g]["t_stat"] = t_g + group_effects[g]["p_value"] = p_g + group_effects[g]["conf_int"] = ci_g # Build treatment effects dataframe treated_df = df.loc[omega_1_mask, [unit, time, "_tau_hat", "_rel_time"]].copy() diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index e4fdf1b02..c93544ab8 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -657,9 +657,9 @@ def _refit_sa(w_r): resolved_survey=None, ) # Create temp weight column for IW aggregation with w_r + # Use full w_r (including zeros) for correct mass computation _wt_col = "_rep_wt" - df[_wt_col] = 0.0 - df.loc[df.index[:len(w_r)], _wt_col] = w_r + df[_wt_col] = w_r es_r, _ = self._compute_iw_effects( df, unit, first_treat, treatment_groups, _sa_rel_periods, ce_r, {}, vcov_r, cim_r, survey_weight_col=_wt_col, diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 86a542892..bc6aa5fb4 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -763,7 +763,7 @@ where weights ŵ_{g,e} = n_{g,e} / Σ_g n_{g,e} (sample share of cohort g at eve - [x] R comparison: ATT matches within machine precision (<1e-11) - [x] R comparison: SE matches within 0.3% (well within 1% threshold) - [x] R comparison: Event study effects match perfectly (correlation 1.0) -- [x] Survey design support (Phase 3): weighted within-transform, survey weights in LinearRegression with TSL vcov; bootstrap+survey supported (Phase 6) via Rao-Wu rescaled bootstrap +- [x] Survey design support (Phase 3): weighted within-transform, survey weights in LinearRegression with TSL vcov; bootstrap+survey supported (Phase 6) via Rao-Wu rescaled bootstrap. Replicate weights supported via estimator-level refit (see Replicate Weight Variance section); replicate+bootstrap rejected. --- @@ -1036,7 +1036,7 @@ The paper text states a stricter bound (T_min + 1) but the R code by the co-auth - [x] Overall ATT as average of post-treatment delta_h with delta-method SE - [x] Anticipation parameter support - [x] Never-treated encoding (0 and inf) -- [x] Survey design support (Phase 3): Q-weights compose multiplicatively with survey weights; TSL vcov on composed weights; survey design columns propagated through sub-experiments +- [x] Survey design support (Phase 3): Q-weights compose multiplicatively with survey weights; TSL vcov on composed weights; survey design columns propagated through sub-experiments. Replicate weights supported via estimator-level refit with Q-weight composition (see Replicate Weight Variance section). - **Note:** Survey weights compose multiplicatively with Q-weights for StackedDiD; only `weight_type="pweight"` (default) is supported — `fweight` and `aweight` are rejected because Q-weight composition changes weight semantics (non-integer for fweight, non-inverse-variance for aweight) ---