From 5c24a747cbf90655a6e4d2951c23b939271fb8c0 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 1 Apr 2026 14:57:37 -0400 Subject: [PATCH 1/7] Add pre-period event study coefficients to ImputationDiD and TwoStageDiD Add pretrends=True parameter to both estimators. When enabled, event study includes pre-treatment horizons (should be ~0 under parallel trends) for visual pre-trends assessment, matching CallawaySantAnna and StackedDiD. ImputationDiD: computes tau_hat for pre-treatment observations of eventually-treated units, aggregates by relative time, uses extended Theorem 3 variance (pre-period targets contribute both directly and via FE correction). TwoStageDiD: extends Stage 2 design matrix with pre-period dummies; GMM sandwich handles naturally. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/imputation.py | 430 ++++++++++++++++++++++-- diff_diff/imputation_bootstrap.py | 85 +++++ diff_diff/two_stage.py | 12 +- diff_diff/two_stage_bootstrap.py | 7 +- docs/methodology/REGISTRY.md | 14 + tests/test_pretrends_event_study.py | 492 ++++++++++++++++++++++++++++ 6 files changed, 1007 insertions(+), 33 deletions(-) create mode 100644 tests/test_pretrends_event_study.py diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 4f32a6638..6449785c2 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -134,6 +134,7 @@ def __init__( rank_deficient_action: str = "warn", horizon_max: Optional[int] = None, aux_partition: str = "cohort_horizon", + pretrends: bool = False, ): if rank_deficient_action not in ("warn", "error", "silent"): raise ValueError( @@ -160,6 +161,7 @@ def __init__( self.rank_deficient_action = rank_deficient_action self.horizon_max = horizon_max self.aux_partition = aux_partition + self.pretrends = pretrends self.is_fitted_ = False self.results_: Optional[ImputationDiDResults] = None @@ -448,6 +450,25 @@ def fit( df["_tau_hat"] = np.nan df.loc[omega_1_mask, "_tau_hat"] = tau_hat + # Pre-period tau_hat for eventually-treated units (pretrends feature) + omega_pre_mask = None + if self.pretrends and aggregate in ("event_study", "all"): + omega_pre_mask = ~df["_never_treated"] & ~df["_treated"] + if omega_pre_mask.any(): + tau_hat_pre, _ = self._impute_treatment_effects( + df, + outcome, + unit, + time, + covariates, + omega_pre_mask, + unit_fe, + time_fe, + grand_mean, + delta_hat, + ) + df.loc[omega_pre_mask, "_tau_hat"] = tau_hat_pre + # ---- Step 3: Aggregate ---- # Always compute overall ATT (simple aggregation) finite_mask = np.isfinite(tau_hat) @@ -1077,12 +1098,20 @@ def _compute_cluster_psi_sums( cluster_var: str, kept_cov_mask: Optional[np.ndarray] = None, survey_weights_0: Optional[np.ndarray] = None, + preperiod_weights: Optional[np.ndarray] = None, ) -> Tuple[np.ndarray, np.ndarray]: """ Compute cluster-level influence function sums (Theorem 3). psi_i = sum_t v_it * epsilon_tilde_it, summed within each cluster. + Parameters + ---------- + preperiod_weights : np.ndarray, optional + Shape (n_0,). Non-zero for pre-period target observations in omega_0. + When provided, these observations contribute both directly (via w_it) + and indirectly (via FE correction) to the influence function. + Returns ------- cluster_psi_sums : np.ndarray @@ -1095,26 +1124,50 @@ def _compute_cluster_psi_sums( n_0 = len(df_0) n_1 = len(df_1) + is_preperiod = preperiod_weights is not None + # ---- Compute v_it for treated observations ---- - v_treated = weights.copy() + if is_preperiod: + v_treated = np.zeros(n_1) # omega_1 contributes nothing + else: + v_treated = weights.copy() # ---- Compute v_it for untreated observations ---- if covariates is None or len(covariates) == 0: # FE-only case: closed-form - treated_units = df_1[unit].values - treated_times = df_1[time].values + # Build w_by_unit, w_by_time, w_total from the target weights + if is_preperiod: + # Target weights are on omega_0 observations + untreated_units = df_0[unit].values + untreated_times = df_0[time].values + + w_by_unit: Dict[Any, float] = {} + for j in range(n_0): + u = untreated_units[j] + w_by_unit[u] = w_by_unit.get(u, 0.0) + preperiod_weights[j] + + w_by_time: Dict[Any, float] = {} + for j in range(n_0): + t = untreated_times[j] + w_by_time[t] = w_by_time.get(t, 0.0) + preperiod_weights[j] + + w_total = float(np.sum(preperiod_weights)) + else: + # Target weights are on omega_1 observations + treated_units = df_1[unit].values + treated_times = df_1[time].values - w_by_unit: Dict[Any, float] = {} - for i_idx in range(n_1): - u = treated_units[i_idx] - w_by_unit[u] = w_by_unit.get(u, 0.0) + weights[i_idx] + w_by_unit = {} + for i_idx in range(n_1): + u = treated_units[i_idx] + w_by_unit[u] = w_by_unit.get(u, 0.0) + weights[i_idx] - w_by_time: Dict[Any, float] = {} - for i_idx in range(n_1): - t = treated_times[i_idx] - w_by_time[t] = w_by_time.get(t, 0.0) + weights[i_idx] + w_by_time = {} + for i_idx in range(n_1): + t = treated_times[i_idx] + w_by_time[t] = w_by_time.get(t, 0.0) + weights[i_idx] - w_total = float(np.sum(weights)) + w_total = float(np.sum(weights)) # Use survey-weighted sums for untreated denominators when present if survey_weights_0 is not None: @@ -1123,12 +1176,13 @@ def _compute_cluster_psi_sums( n0_by_time = sw0_series.groupby(df_0[time]).sum().to_dict() n0_denom = float(np.sum(survey_weights_0)) else: + if not is_preperiod: + untreated_units = df_0[unit].values + untreated_times = df_0[time].values n0_by_unit = df_0.groupby(unit).size().to_dict() n0_by_time = df_0.groupby(time).size().to_dict() n0_denom = n_0 - untreated_units = df_0[unit].values - untreated_times = df_0[time].values v_untreated = np.zeros(n_0) for j in range(n_0): @@ -1142,19 +1196,35 @@ def _compute_cluster_psi_sums( # WLS projection requires per-obs survey weight factor if survey_weights_0 is not None: base_v *= survey_weights_0[j] - v_untreated[j] = base_v + # For preperiod targets, add direct weight contribution + if is_preperiod: + v_untreated[j] = preperiod_weights[j] + base_v + else: + v_untreated[j] = base_v else: - v_untreated = self._compute_v_untreated_with_covariates( - df_0, - df_1, - unit, - time, - covariates, - weights, - delta_hat, - kept_cov_mask=kept_cov_mask, - survey_weights_0=survey_weights_0, - ) + if is_preperiod: + v_untreated = self._compute_v_untreated_with_covariates_preperiod( + df_0, + unit, + time, + covariates, + preperiod_weights, + delta_hat, + kept_cov_mask=kept_cov_mask, + survey_weights_0=survey_weights_0, + ) + else: + v_untreated = self._compute_v_untreated_with_covariates( + df_0, + df_1, + unit, + time, + covariates, + weights, + delta_hat, + kept_cov_mask=kept_cov_mask, + survey_weights_0=survey_weights_0, + ) # ---- Compute auxiliary model residuals (Equation 8) ---- epsilon_treated = self._compute_auxiliary_residuals_treated( @@ -1174,6 +1244,25 @@ def _compute_cluster_psi_sums( df_0, outcome, unit, time, covariates, unit_fe, time_fe, grand_mean, delta_hat ) + # Override epsilon for pre-period target observations with auxiliary residuals + if is_preperiod: + target_mask_in_0 = preperiod_weights > 0 + if target_mask_in_0.any(): + aux_resid = self._compute_auxiliary_residuals_preperiod( + df_0.loc[df_0.index[target_mask_in_0]], + outcome, + unit, + time, + first_treat, + covariates, + unit_fe, + time_fe, + grand_mean, + delta_hat, + v_untreated[target_mask_in_0], + ) + epsilon_untreated[target_mask_in_0] = aux_resid + # ---- psi_it = v_it * epsilon_tilde_it ---- v_all = np.empty(len(df)) v_all[omega_1_mask.values] = v_treated @@ -1213,6 +1302,7 @@ def _compute_conservative_variance( cluster_var: str, kept_cov_mask: Optional[np.ndarray] = None, survey_weights: Optional[np.ndarray] = None, + preperiod_weights: Optional[np.ndarray] = None, ) -> float: """ Compute conservative clustered variance (Theorem 3, Equation 7). @@ -1225,6 +1315,10 @@ def _compute_conservative_variance( survey_weights : np.ndarray, optional Full-panel survey weights. When provided, untreated denominators in v_it use survey-weighted sums instead of raw counts. + preperiod_weights : np.ndarray, optional + Aggregation weights for pre-period target observations in omega_0. + Shape: (n_untreated,). When provided, these observations are both + in the aggregation target and in the FE estimation sample. Returns ------- @@ -1249,6 +1343,7 @@ def _compute_conservative_variance( cluster_var=cluster_var, kept_cov_mask=kept_cov_mask, survey_weights_0=sw_0, + preperiod_weights=preperiod_weights, ) sigma_sq = float((cluster_psi_sums**2).sum()) return np.sqrt(max(sigma_sq, 0.0)) @@ -1443,6 +1538,148 @@ def _compute_residuals_untreated( return df_0[outcome].values - y_hat + def _compute_auxiliary_residuals_preperiod( + self, + df_pre: pd.DataFrame, + outcome: str, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + v_pre: np.ndarray, + ) -> np.ndarray: + """ + Compute auxiliary residuals for pre-period target observations. + + Analogous to _compute_auxiliary_residuals_treated but operates on + pre-treatment observations of eventually-treated units in omega_0. + + epsilon_tilde_it = tau_hat_pre_it - tau_tilde_g + where tau_tilde_g = sum(v_it * tau_hat_pre_it) / sum(v_it) within group g. + """ + n_pre = len(df_pre) + + # Compute tau_hat_pre = Y - Y_hat(0) (same as Step 1 residual) + alpha_i = df_pre[unit].map(unit_fe).values.astype(float) + beta_t = df_pre[time].map(time_fe).values.astype(float) + y_hat_0 = grand_mean + alpha_i + beta_t + + if delta_hat is not None and covariates: + y_hat_0 = y_hat_0 + np.dot(df_pre[covariates].values, delta_hat) + + tau_hat_pre = df_pre[outcome].values - y_hat_0 + + # Partition by aux_partition and compute group means + if self.aux_partition == "cohort_horizon": + group_keys = list( + zip(df_pre[first_treat].values, df_pre["_rel_time"].values) + ) + elif self.aux_partition == "cohort": + group_keys = list(df_pre[first_treat].values) + elif self.aux_partition == "horizon": + group_keys = list(df_pre["_rel_time"].values) + else: + group_keys = list(range(n_pre)) + + group_series = pd.Series(group_keys, index=df_pre.index) + tau_series = pd.Series(tau_hat_pre, index=df_pre.index) + v_series = pd.Series(v_pre, index=df_pre.index) + + weighted_tau_sum = (v_series * tau_series).groupby(group_series).sum() + weight_sum = v_series.groupby(group_series).sum() + + zero_weight_groups = weight_sum.abs() < 1e-15 + if zero_weight_groups.any(): + simple_means = tau_series.groupby(group_series).mean() + tau_tilde_map = weighted_tau_sum / weight_sum + tau_tilde_map = tau_tilde_map.where(~zero_weight_groups, simple_means) + else: + tau_tilde_map = weighted_tau_sum / weight_sum + + tau_tilde = group_series.map(tau_tilde_map).values + return tau_hat_pre - tau_tilde + + def _compute_v_untreated_with_covariates_preperiod( + self, + df_0: pd.DataFrame, + unit: str, + time: str, + covariates: List[str], + preperiod_weights: np.ndarray, + delta_hat: Optional[np.ndarray], + kept_cov_mask: Optional[np.ndarray] = None, + survey_weights_0: Optional[np.ndarray] = None, + ) -> np.ndarray: + """ + Compute v_it for untreated observations when covariates present (preperiod case). + + v = preperiod_weights - [W_0] A_0 (A_0'[W]A_0)^{-1} A_0' preperiod_weights + This is (I - P_0) * preperiod_weights where P_0 is the hat matrix. + """ + if kept_cov_mask is not None and not np.all(kept_cov_mask): + covariates = [c for c, k in zip(covariates, kept_cov_mask) if k] + + units_0 = df_0[unit].values + times_0 = df_0[time].values + + all_units = np.unique(units_0) + all_times = np.unique(times_0) + unit_to_idx = {u: i for i, u in enumerate(all_units)} + time_to_idx = {t: i for i, t in enumerate(all_times)} + n_units = len(all_units) + n_fe_cols = (n_units - 1) + (len(all_times) - 1) + + n_0 = len(df_0) + n_cov = len(covariates) + + # Build A_0 sparse + u_indices = np.array([unit_to_idx[u] for u in units_0]) + u_mask = u_indices > 0 + u_rows = np.arange(n_0)[u_mask] + u_cols = u_indices[u_mask] - 1 + + t_indices = np.array([time_to_idx[t] for t in times_0]) + t_mask = t_indices > 0 + t_rows = np.arange(n_0)[t_mask] + t_cols = (n_units - 1) + t_indices[t_mask] - 1 + + rows = np.concatenate([u_rows, t_rows]) + cols = np.concatenate([u_cols, t_cols]) + data = np.ones(len(rows)) + + A_fe = sparse.csr_matrix((data, (rows, cols)), shape=(n_0, n_fe_cols)) + + if n_cov > 0: + A_cov = sparse.csr_matrix(df_0[covariates].values) + A_0 = sparse.hstack([A_fe, A_cov], format="csr") + else: + A_0 = A_fe + + # Compute A_0' preperiod_weights + A0_w = A_0.T @ preperiod_weights # shape (p,) + + # Solve (A_0'[W]A_0) z = A_0' preperiod_weights + if survey_weights_0 is not None: + A0tA0 = A_0.T @ A_0.multiply(survey_weights_0[:, None]) + else: + A0tA0 = A_0.T @ A_0 + try: + z = spsolve(A0tA0.tocsc(), A0_w) + except Exception: + A0tA0_dense = A0tA0.toarray() + z, _, _, _ = np.linalg.lstsq(A0tA0_dense, A0_w, rcond=None) + + # v = preperiod_weights - [W_0] A_0 z + projection = A_0 @ z + if survey_weights_0 is not None: + projection = projection * survey_weights_0 + v_untreated = preperiod_weights - projection + return v_untreated + # ========================================================================= # Aggregation # ========================================================================= @@ -1473,8 +1710,25 @@ def _aggregate_event_study( tau_hat = df["_tau_hat"].loc[omega_1_mask].values rel_times = df_1["_rel_time"].values - # Get all horizons - all_horizons = sorted(set(int(h) for h in rel_times if np.isfinite(h))) + # Get post-treatment horizons from omega_1 + post_horizons = sorted(set(int(h) for h in rel_times if np.isfinite(h))) + + # Pre-period horizons (pretrends feature) + pre_horizons: List[int] = [] + df_pre = None + tau_hat_pre = None + rel_times_pre = None + if self.pretrends: + omega_pre_mask = ~df["_never_treated"] & ~df["_treated"] + if omega_pre_mask.any(): + df_pre = df.loc[omega_pre_mask] + tau_hat_pre = df["_tau_hat"].loc[omega_pre_mask].values + rel_times_pre = df_pre["_rel_time"].values + pre_horizons = sorted( + set(int(h) for h in rel_times_pre if np.isfinite(h)) + ) + + all_horizons = sorted(set(post_horizons + pre_horizons)) # Apply horizon_max filter if self.horizon_max is not None: @@ -1489,8 +1743,20 @@ def _aggregate_event_study( ), index=df_1.index, ) + if df_pre is not None: + balanced_mask_pre = pd.Series( + self._compute_balanced_cohort_mask( + df_pre, first_treat, all_horizons, balance_e, cohort_rel_times + ), + index=df_pre.index, + ) + else: + balanced_mask_pre = None else: balanced_mask = pd.Series(True, index=df_1.index) + balanced_mask_pre = ( + pd.Series(True, index=df_pre.index) if df_pre is not None else None + ) # Check Proposition 5: no never-treated units has_never_treated = df["_never_treated"].any() @@ -1520,6 +1786,113 @@ def _aggregate_event_study( if h == ref_period: continue + # ---- Pre-period horizon (from omega_pre observations) ---- + if h < -self.anticipation: + if rel_times_pre is None or balanced_mask_pre is None: + continue + h_mask_pre = (rel_times_pre == h) & balanced_mask_pre.values + n_h = int(h_mask_pre.sum()) + if n_h == 0: + continue + + tau_h = tau_hat_pre[h_mask_pre] + finite_h = np.isfinite(tau_h) + valid_tau = tau_h[finite_h] + + if len(valid_tau) == 0: + event_study_effects[h] = { + "effect": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "n_obs": n_h, + } + continue + + # Effect: survey-weighted or simple mean + if survey_weights is not None: + pre_sw = survey_weights[ + (~df["_never_treated"] & ~df["_treated"]).values + ] + sw_h = pre_sw[h_mask_pre] + sw_valid = sw_h[finite_h] + effect = float(np.average(valid_tau, weights=sw_valid)) + else: + effect = float(np.mean(valid_tau)) + + # SE via conservative variance with preperiod_weights + n_0 = int(omega_0_mask.sum()) + preperiod_weights_h = np.zeros(n_0) + # Map h_mask_pre (indices in df_pre) to indices in omega_0 + omega_0_indices = np.where(omega_0_mask.values)[0] + pre_in_0 = (~df["_never_treated"] & ~df["_treated"]).values + # df_pre indices map to omega_0 positions + pre_positions_in_0 = np.zeros(len(df), dtype=bool) + pre_positions_in_0[df_pre.index[h_mask_pre]] = True + pre_in_0_mask = pre_positions_in_0[omega_0_indices] + + if survey_weights is not None: + sw_0 = survey_weights[omega_0_mask.values] + sw_target = sw_0[pre_in_0_mask] + finite_in_target = np.isfinite( + df["_tau_hat"].values[omega_0_indices[pre_in_0_mask]] + ) + sw_finite = sw_target[finite_in_target] + if sw_finite.sum() > 0: + target_indices = np.where(pre_in_0_mask)[0] + finite_indices = target_indices[finite_in_target] + preperiod_weights_h[finite_indices] = ( + sw_finite / sw_finite.sum() + ) + else: + n_finite = int(np.isfinite( + df["_tau_hat"].values[omega_0_indices[pre_in_0_mask]] + ).sum()) + if n_finite > 0: + tau_at_target = df["_tau_hat"].values[ + omega_0_indices[pre_in_0_mask] + ] + finite_target = np.isfinite(tau_at_target) + target_indices = np.where(pre_in_0_mask)[0] + finite_indices = target_indices[finite_target] + preperiod_weights_h[finite_indices] = 1.0 / n_finite + + 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=np.zeros(len(tau_hat)), + cluster_var=cluster_var, + kept_cov_mask=kept_cov_mask, + survey_weights=survey_weights, + preperiod_weights=preperiod_weights_h, + ) + + t_stat, p_value, conf_int = safe_inference( + effect, se, alpha=self.alpha, df=survey_df + ) + + event_study_effects[h] = { + "effect": effect, + "se": se, + "t_stat": t_stat, + "p_value": p_value, + "conf_int": conf_int, + "n_obs": n_h, + } + continue + + # ---- Post-treatment horizon (from omega_1 observations) ---- # Select treated obs at this horizon from balanced cohorts h_mask = (rel_times == h) & balanced_mask.values n_h = int(h_mask.sum()) @@ -1910,6 +2283,7 @@ def get_params(self) -> Dict[str, Any]: "rank_deficient_action": self.rank_deficient_action, "horizon_max": self.horizon_max, "aux_partition": self.aux_partition, + "pretrends": self.pretrends, } def set_params(self, **params) -> "ImputationDiD": diff --git a/diff_diff/imputation_bootstrap.py b/diff_diff/imputation_bootstrap.py index f3f1a04ea..d1f09046b 100644 --- a/diff_diff/imputation_bootstrap.py +++ b/diff_diff/imputation_bootstrap.py @@ -186,6 +186,37 @@ def _precompute_bootstrap_psi( ) ref_period = -1 - self.anticipation + + # Pre-period data for pretrends bootstrap + pre_rel_times = None + pre_tau_hat = None + pre_balanced_mask = None + n_0 = int(omega_0_mask.sum()) + if self.pretrends: + omega_pre_mask = ~df["_never_treated"] & ~df["_treated"] + if omega_pre_mask.any(): + df_pre = df.loc[omega_pre_mask] + pre_rel_times = df_pre["_rel_time"].values + pre_tau_hat = df["_tau_hat"].loc[omega_pre_mask].values + if balance_e is not None: + all_h_pre = sorted( + set(int(h) for h in pre_rel_times if np.isfinite(h)) + ) + if self.horizon_max is not None: + all_h_pre = [ + h for h in all_h_pre if abs(h) <= self.horizon_max + ] + cohort_rel_times_pre = self._build_cohort_rel_times( + df, first_treat + ) + pre_balanced_mask = self._compute_balanced_cohort_mask( + df_pre, + first_treat, + all_h_pre, + balance_e, + cohort_rel_times_pre, + ) + for h in event_study_effects: if event_study_effects[h].get("n_obs", 0) == 0: continue @@ -193,6 +224,60 @@ def _precompute_bootstrap_psi( continue if not np.isfinite(event_study_effects[h].get("effect", np.nan)): continue + + # ---- Pre-period horizon ---- + if h < -self.anticipation and pre_rel_times is not None: + h_mask_pre = pre_rel_times == h + if pre_balanced_mask is not None: + h_mask_pre = h_mask_pre & pre_balanced_mask + + # Build preperiod_weights_h on omega_0 + omega_0_indices = np.where(omega_0_mask.values)[0] + pre_positions_in_0 = np.zeros(len(df), dtype=bool) + df_pre_subset = df.loc[ + (~df["_never_treated"]) & (~df["_treated"]) + ] + pre_positions_in_0[df_pre_subset.index[h_mask_pre]] = True + pre_in_0_mask = pre_positions_in_0[omega_0_indices] + + preperiod_weights_h = np.zeros(n_0) + if survey_weights_0 is not None: + sw_0 = survey_weights_0 + sw_target = sw_0[pre_in_0_mask] + tau_at_target = df["_tau_hat"].values[ + omega_0_indices[pre_in_0_mask] + ] + finite_in_target = np.isfinite(tau_at_target) + sw_finite = sw_target[finite_in_target] + if sw_finite.sum() > 0: + target_idx = np.where(pre_in_0_mask)[0] + fin_idx = target_idx[finite_in_target] + preperiod_weights_h[fin_idx] = ( + sw_finite / sw_finite.sum() + ) + else: + continue + else: + tau_at_target = df["_tau_hat"].values[ + omega_0_indices[pre_in_0_mask] + ] + finite_target = np.isfinite(tau_at_target) + n_finite = int(finite_target.sum()) + if n_finite == 0: + continue + target_idx = np.where(pre_in_0_mask)[0] + fin_idx = target_idx[finite_target] + preperiod_weights_h[fin_idx] = 1.0 / n_finite + + psi_h, _ = self._compute_cluster_psi_sums( + **common, + weights=np.zeros(len(tau_hat)), + preperiod_weights=preperiod_weights_h, + ) + result["event_study"][h] = psi_h + continue + + # ---- Post-treatment horizon ---- h_mask = rel_times == h if balanced_mask is not None: h_mask = h_mask & balanced_mask diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index cd4baa4e1..47eada2f3 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -139,6 +139,7 @@ def __init__( seed: Optional[int] = None, rank_deficient_action: str = "warn", horizon_max: Optional[int] = None, + pretrends: bool = False, ): if rank_deficient_action not in ("warn", "error", "silent"): raise ValueError( @@ -159,6 +160,7 @@ def __init__( self.seed = seed self.rank_deficient_action = rank_deficient_action self.horizon_max = horizon_max + self.pretrends = pretrends self.is_fitted_ = False self.results_: Optional[TwoStageDiDResults] = None @@ -1015,9 +1017,12 @@ def _stage2_event_study( rel_times = df["_rel_time"].values n = len(df) - # Get all horizons from treated observations - treated_rel = rel_times[omega_1_mask.values] - all_horizons = sorted(set(int(h) for h in treated_rel if np.isfinite(h))) + # Get all horizons — include pre-periods when pretrends=True + if self.pretrends: + evt_rel = rel_times[~df["_never_treated"].values] + else: + evt_rel = rel_times[omega_1_mask.values] + all_horizons = sorted(set(int(h) for h in evt_rel if np.isfinite(h))) # Apply horizon_max filter if self.horizon_max is not None: @@ -1609,6 +1614,7 @@ def get_params(self) -> Dict[str, Any]: "seed": self.seed, "rank_deficient_action": self.rank_deficient_action, "horizon_max": self.horizon_max, + "pretrends": self.pretrends, } def set_params(self, **params) -> "TwoStageDiD": diff --git a/diff_diff/two_stage_bootstrap.py b/diff_diff/two_stage_bootstrap.py index f4d69816a..3effef4ed 100644 --- a/diff_diff/two_stage_bootstrap.py +++ b/diff_diff/two_stage_bootstrap.py @@ -321,8 +321,11 @@ def _run_bootstrap( if original_event_study and aggregate in ("event_study", "all"): # Recompute S scores for event study specification rel_times = df["_rel_time"].values - treated_rel = rel_times[omega_1_mask.values] - all_horizons = sorted(set(int(h) for h in treated_rel if np.isfinite(h))) + if self.pretrends: + evt_rel = rel_times[~df["_never_treated"].values] + else: + evt_rel = rel_times[omega_1_mask.values] + all_horizons = sorted(set(int(h) for h in evt_rel if np.isfinite(h))) if self.horizon_max is not None: all_horizons = [h for h in all_horizons if abs(h) <= self.horizon_max] diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index a39822395..23d844060 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -836,6 +836,19 @@ Y_it = alpha_i + beta_t [+ X'_it * delta] + W'_it * gamma + epsilon_it - Test `gamma = 0` via cluster-robust Wald F-test - Independent of treatment effect estimation (Proposition 9) +*Pre-period event study coefficients (`pretrends=True`, BJS Section 4.2):* + +For pre-treatment observation (i,t) of eventually-treated unit i at relative time `K_it = t - E_i < -anticipation`: +``` +tau_hat_pre(i,t) = Y_it - alpha_hat_i - beta_hat_t [- X'_it * delta_hat] +theta_h = (1/n_h) * sum_{(i,t): K_it = h} tau_hat_pre(i,t) +``` +- Under parallel trends (Assumption 1), `E[theta_h] = 0` for all `h < -anticipation` +- Reference period `h = -1 - anticipation` normalized to zero (not estimated) +- Variance via Theorem 3 extension: pre-period targets are in `Omega_0`, so `v_it = w_it + FE_correction` (both direct aggregation weight and indirect FE estimation correction) +- Only affects event study aggregation; overall ATT and group aggregation unchanged +- Proposition 5 does not apply to pre-period horizons (always identified) + *Edge cases:* - **Unbalanced panels:** FE estimated via iterative alternating projection (Gauss-Seidel), equivalent to OLS with unit+time dummies. Converges in O(max_iter) passes; typically 5-20 iterations for unbalanced panels, 1-2 for balanced. One-pass demeaning is only exact for balanced panels. - **No never-treated units (Proposition 5):** Long-run effects at horizons `h >= H_bar` are not identified. Set to NaN with warning listing affected horizons. @@ -925,6 +938,7 @@ Our implementation uses multiplier bootstrap on the GMM influence function: clus - **NaN y_tilde handling:** When Stage 1 FE are unidentified for some observations, the residualized outcome `y_tilde` is NaN. These observations are zeroed out (excluded) from the Stage 2 regression and variance computation, matching the treatment of unimputable observations in ImputationDiD. - **NaN inference for undefined statistics:** t_stat uses NaN when SE is non-finite or zero; p_value and CI also NaN. Matches CallawaySantAnna/ImputationDiD NaN convention. - **Event study aggregation:** Horizon-specific effects use the same two-stage procedure with horizon indicator dummies in Stage 2. Unidentified horizons (e.g., long-run effects without never-treated units, per Proposition 5 of Borusyak et al. 2024) produce NaN. +- **Pre-period event study coefficients (`pretrends=True`):** When enabled, the Stage 2 design matrix `X_2` includes pre-period relative-time dummies. Pre-period observations have `y_tilde = Step 1 residual` by construction. The GMM sandwich variance accounts for Stage 1 estimation error (Gardner 2022, Theorem 1). Only affects event study aggregation; overall ATT unchanged. - **balance_e with no qualifying cohorts:** If no cohorts have sufficient pre/post coverage for the requested `balance_e`, a warning is emitted and event study results contain only the reference period. - **No never-treated units (Proposition 5):** When there are no never-treated units and multiple treatment cohorts, horizons h >= h_bar (where h_bar = max(groups) - min(groups)) are unidentified per Proposition 5 of Borusyak et al. (2024). These produce NaN inference with n_obs > 0 (treated observations exist but counterfactual is unidentified) and a warning listing affected horizons. Matches ImputationDiD behavior. Proposition 5 applies to event study horizons only, not cohort aggregation — a cohort whose treated obs all fall at Prop 5 horizons naturally gets n_obs=0 in group effects because all its y_tilde values are NaN. - **Zero-observation horizons after filtering:** When `balance_e` or NaN `y_tilde` filtering results in zero observations for some non-Prop-5 event study horizons, those horizons produce NaN for all inference fields (effect, SE, t-stat, p-value, CI) with n_obs=0. diff --git a/tests/test_pretrends_event_study.py b/tests/test_pretrends_event_study.py new file mode 100644 index 000000000..e13319f9e --- /dev/null +++ b/tests/test_pretrends_event_study.py @@ -0,0 +1,492 @@ +""" +Tests for pretrends=True feature in ImputationDiD and TwoStageDiD. + +Tests that pre-period event study coefficients are computed correctly +when pretrends=True is set on the estimator. +""" + +import numpy as np +import pandas as pd +import pytest + +from diff_diff.imputation import ImputationDiD +from diff_diff.two_stage import TwoStageDiD + + +def generate_test_data( + n_units: int = 100, + n_periods: int = 10, + treatment_effect: float = 2.0, + never_treated_frac: float = 0.3, + seed: int = 42, +) -> pd.DataFrame: + """Generate staggered adoption data for pretrends testing.""" + rng = np.random.default_rng(seed) + + units = np.repeat(np.arange(n_units), n_periods) + times = np.tile(np.arange(n_periods), n_units) + + n_never = int(n_units * never_treated_frac) + n_treated = n_units - n_never + + cohort_periods = np.array([3, 5, 7]) + first_treat = np.zeros(n_units, dtype=int) + if n_treated > 0: + cohort_assignments = rng.choice(len(cohort_periods), size=n_treated) + first_treat[n_never:] = cohort_periods[cohort_assignments] + + first_treat_expanded = np.repeat(first_treat, n_periods) + + unit_fe = rng.standard_normal(n_units) * 2.0 + time_fe = np.linspace(0, 1, n_periods) + + unit_fe_expanded = np.repeat(unit_fe, n_periods) + time_fe_expanded = np.tile(time_fe, n_units) + + post = (times >= first_treat_expanded) & (first_treat_expanded > 0) + relative_time = times - first_treat_expanded + dynamic_mult = 1 + 0.1 * np.maximum(relative_time, 0) + effect = treatment_effect * dynamic_mult + + outcomes = ( + unit_fe_expanded + + time_fe_expanded + + effect * post + + rng.standard_normal(len(units)) * 0.5 + ) + + return pd.DataFrame( + { + "unit": units, + "time": times, + "outcome": outcomes, + "first_treat": first_treat_expanded, + } + ) + + +# ============================================================================= +# ImputationDiD pretrends tests +# ============================================================================= + + +class TestImputationPretrends: + """Tests for ImputationDiD pretrends feature.""" + + def test_pretrends_includes_negative_horizons(self): + """Pre-period horizons appear in event study when pretrends=True.""" + data = generate_test_data() + est = ImputationDiD(pretrends=True) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + assert results.event_study_effects is not None + horizons = sorted(results.event_study_effects.keys()) + negative = [h for h in horizons if h < 0] + positive = [h for h in horizons if h >= 0] + assert len(negative) > 1, "Should have pre-period horizons" + assert len(positive) > 0, "Should have post-treatment horizons" + + def test_pretrends_coefficients_near_zero(self): + """Pre-period effects ~0 under parallel trends (no violation).""" + data = generate_test_data(treatment_effect=2.0, seed=99) + est = ImputationDiD(pretrends=True) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + ref_period = -1 + for h, eff in results.event_study_effects.items(): + if h >= 0 or h == ref_period: + continue + assert np.isfinite(eff["effect"]), f"h={h}: effect not finite" + assert abs(eff["effect"]) < 3 * eff["se"] + 0.5, ( + f"h={h}: pre-period effect {eff['effect']:.3f} too large" + ) + + def test_pretrends_se_finite_positive(self): + """All pre-period horizons have finite, positive SEs.""" + data = generate_test_data() + est = ImputationDiD(pretrends=True) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + ref_period = -1 + for h, eff in results.event_study_effects.items(): + if h == ref_period: + continue + if eff["n_obs"] == 0: + continue + assert np.isfinite(eff["se"]), f"h={h}: SE not finite" + assert eff["se"] > 0, f"h={h}: SE not positive" + + def test_reference_period_correct(self): + """Reference period h=-1 normalized to 0.""" + data = generate_test_data() + est = ImputationDiD(pretrends=True) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + assert -1 in results.event_study_effects + ref = results.event_study_effects[-1] + assert ref["effect"] == 0.0 + assert ref["n_obs"] == 0 + + def test_backward_compatibility(self): + """pretrends=False (default) gives identical results.""" + data = generate_test_data() + + results_default = ImputationDiD().fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + results_false = ImputationDiD(pretrends=False).fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + assert results_default.overall_att == results_false.overall_att + assert results_default.overall_se == results_false.overall_se + assert set(results_default.event_study_effects.keys()) == set( + results_false.event_study_effects.keys() + ) + + def test_post_treatment_invariance(self): + """Post-treatment effects identical with pretrends=True vs False.""" + data = generate_test_data() + + results_off = ImputationDiD().fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + results_on = ImputationDiD(pretrends=True).fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + # Overall ATT unchanged + assert results_on.overall_att == results_off.overall_att + assert results_on.overall_se == results_off.overall_se + + # Post-treatment event study effects unchanged + for h in results_off.event_study_effects: + assert h in results_on.event_study_effects + eff_off = results_off.event_study_effects[h] + eff_on = results_on.event_study_effects[h] + np.testing.assert_allclose( + eff_off["effect"], eff_on["effect"], rtol=1e-10 + ) + np.testing.assert_allclose(eff_off["se"], eff_on["se"], rtol=1e-10) + + def test_horizon_max_interaction(self): + """horizon_max limits both pre and post horizons.""" + data = generate_test_data() + est = ImputationDiD(pretrends=True, horizon_max=2) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + for h in results.event_study_effects: + assert abs(h) <= 2, f"h={h} exceeds horizon_max=2" + + def test_anticipation_interaction(self): + """With anticipation=1, reference shifts to h=-2.""" + data = generate_test_data() + est = ImputationDiD(pretrends=True, anticipation=1) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + assert -2 in results.event_study_effects + ref = results.event_study_effects[-2] + assert ref["effect"] == 0.0 + assert ref["n_obs"] == 0 + + def test_get_params_includes_pretrends(self): + """get_params includes pretrends parameter.""" + est = ImputationDiD(pretrends=True) + params = est.get_params() + assert "pretrends" in params + assert params["pretrends"] is True + + est2 = ImputationDiD() + assert est2.get_params()["pretrends"] is False + + def test_no_pretreatment_obs_graceful(self): + """All units treated at t=1 with pretrends=True: no error.""" + rng = np.random.default_rng(42) + n_units = 20 + n_periods = 5 + data = pd.DataFrame( + { + "unit": np.repeat(np.arange(n_units), n_periods), + "time": np.tile(np.arange(n_periods), n_units), + "outcome": rng.standard_normal(n_units * n_periods), + "first_treat": np.repeat( + np.ones(n_units, dtype=int), n_periods + ), + } + ) + + est = ImputationDiD(pretrends=True) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + assert results.event_study_effects is not None + + +# ============================================================================= +# TwoStageDiD pretrends tests +# ============================================================================= + + +class TestTwoStagePretrends: + """Tests for TwoStageDiD pretrends feature.""" + + def test_pretrends_includes_negative_horizons(self): + """Pre-period horizons appear in event study when pretrends=True.""" + data = generate_test_data() + est = TwoStageDiD(pretrends=True) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + assert results.event_study_effects is not None + horizons = sorted(results.event_study_effects.keys()) + negative = [h for h in horizons if h < 0] + positive = [h for h in horizons if h >= 0] + assert len(negative) > 1, "Should have pre-period horizons" + assert len(positive) > 0, "Should have post-treatment horizons" + + def test_pretrends_coefficients_near_zero(self): + """Pre-period effects ~0 under parallel trends.""" + data = generate_test_data(treatment_effect=2.0, seed=99) + est = TwoStageDiD(pretrends=True) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + ref_period = -1 + for h, eff in results.event_study_effects.items(): + if h >= 0 or h == ref_period: + continue + assert np.isfinite(eff["effect"]), f"h={h}: effect not finite" + assert abs(eff["effect"]) < 3 * eff["se"] + 0.5, ( + f"h={h}: pre-period effect {eff['effect']:.3f} too large" + ) + + def test_pretrends_se_finite_positive(self): + """All pre-period horizons have finite, positive SEs.""" + data = generate_test_data() + est = TwoStageDiD(pretrends=True) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + ref_period = -1 + for h, eff in results.event_study_effects.items(): + if h == ref_period: + continue + if eff["n_obs"] == 0: + continue + assert np.isfinite(eff["se"]), f"h={h}: SE not finite" + assert eff["se"] > 0, f"h={h}: SE not positive" + + def test_post_treatment_invariance(self): + """Post-treatment effects identical with pretrends=True vs False.""" + data = generate_test_data() + + results_off = TwoStageDiD().fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + results_on = TwoStageDiD(pretrends=True).fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + # Overall ATT unchanged + assert results_on.overall_att == results_off.overall_att + assert results_on.overall_se == results_off.overall_se + + # Post-treatment event study effects unchanged + for h in results_off.event_study_effects: + assert h in results_on.event_study_effects + eff_off = results_off.event_study_effects[h] + eff_on = results_on.event_study_effects[h] + np.testing.assert_allclose( + eff_off["effect"], eff_on["effect"], rtol=1e-10 + ) + np.testing.assert_allclose(eff_off["se"], eff_on["se"], rtol=1e-10) + + def test_get_params_includes_pretrends(self): + """get_params includes pretrends parameter.""" + est = TwoStageDiD(pretrends=True) + params = est.get_params() + assert "pretrends" in params + assert params["pretrends"] is True + + def test_horizon_max_interaction(self): + """horizon_max limits both pre and post horizons.""" + data = generate_test_data() + est = TwoStageDiD(pretrends=True, horizon_max=2) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + for h in results.event_study_effects: + assert abs(h) <= 2, f"h={h} exceeds horizon_max=2" + + def test_anticipation_interaction(self): + """With anticipation=1, reference shifts to h=-2.""" + data = generate_test_data() + est = TwoStageDiD(pretrends=True, anticipation=1) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + assert -2 in results.event_study_effects + ref = results.event_study_effects[-2] + assert ref["effect"] == 0.0 + assert ref["n_obs"] == 0 + + +# ============================================================================= +# Cross-estimator consistency +# ============================================================================= + + +class TestPretrendsCrossEstimator: + """Verify ImputationDiD and TwoStageDiD produce consistent pre-period effects.""" + + def test_point_estimates_close(self): + """ImputationDiD and TwoStageDiD pre-period effects should be similar.""" + data = generate_test_data(seed=77) + + imp = ImputationDiD(pretrends=True).fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + ts = TwoStageDiD(pretrends=True).fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + # Both should have negative horizons + imp_neg = { + h: v["effect"] + for h, v in imp.event_study_effects.items() + if h < -1 and v["n_obs"] > 0 + } + ts_neg = { + h: v["effect"] + for h, v in ts.event_study_effects.items() + if h < -1 and v["n_obs"] > 0 + } + + common_h = set(imp_neg.keys()) & set(ts_neg.keys()) + assert len(common_h) > 0, "No common pre-period horizons" + + for h in common_h: + # Point estimates should be numerically identical (same imputation) + np.testing.assert_allclose( + imp_neg[h], + ts_neg[h], + atol=1e-6, + err_msg=f"h={h}: point estimates differ", + ) From debddc59e03f6c1826323ceb2195cb155a6cd048 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 1 Apr 2026 15:17:04 -0400 Subject: [PATCH 2/7] Address AI review: fix P0 bootstrap balance_e mismatch, P1 survey regression P0: Bootstrap pre-period path now uses same union of pre+post horizons for balance_e cohort filtering as the analytical path. Previously used pre-only horizons which could include different cohorts. P1: Moved untreated_units/untreated_times initialization before the survey_weights_0 branch in _compute_cluster_psi_sums, fixing a regression where the survey-weighted FE-only variance path would crash. P2: Added regression tests for survey-weighted, covariates + pretrends, and balance_e + bootstrap combinations. Added pretrends parameter to public docstrings for both estimators. P3: Removed dead pre_tau_hat variable in bootstrap path. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/imputation.py | 13 ++- diff_diff/imputation_bootstrap.py | 64 ++++++-------- diff_diff/two_stage.py | 5 ++ tests/test_pretrends_event_study.py | 128 ++++++++++++++++++++++++++++ 4 files changed, 171 insertions(+), 39 deletions(-) diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 6449785c2..017673d2c 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -81,6 +81,11 @@ class ImputationDiD(ImputationDiDBootstrapMixin): - "cohort_horizon": Groups by cohort x relative time (tightest SEs) - "cohort": Groups by cohort only (more conservative) - "horizon": Groups by relative time only (more conservative) + pretrends : bool, default=False + If True, event study includes pre-treatment horizons for visual + pre-trends assessment. Pre-period effects should be ~0 under + parallel trends. Only affects event_study aggregation; overall + ATT and group aggregation are unchanged. Attributes ---------- @@ -1169,6 +1174,11 @@ def _compute_cluster_psi_sums( w_total = float(np.sum(weights)) + # Always initialize untreated unit/time arrays for the FE loop + if not is_preperiod: + untreated_units = df_0[unit].values + untreated_times = df_0[time].values + # Use survey-weighted sums for untreated denominators when present if survey_weights_0 is not None: sw0_series = pd.Series(survey_weights_0, index=df_0.index) @@ -1176,9 +1186,6 @@ def _compute_cluster_psi_sums( n0_by_time = sw0_series.groupby(df_0[time]).sum().to_dict() n0_denom = float(np.sum(survey_weights_0)) else: - if not is_preperiod: - untreated_units = df_0[unit].values - untreated_times = df_0[time].values n0_by_unit = df_0.groupby(unit).size().to_dict() n0_by_time = df_0.groupby(time).size().to_dict() n0_denom = n_0 diff --git a/diff_diff/imputation_bootstrap.py b/diff_diff/imputation_bootstrap.py index d1f09046b..d8f780c52 100644 --- a/diff_diff/imputation_bootstrap.py +++ b/diff_diff/imputation_bootstrap.py @@ -174,48 +174,43 @@ def _precompute_bootstrap_psi( df_1 = df.loc[omega_1_mask] rel_times = df_1["_rel_time"].values - # Balanced cohort mask (same logic as _aggregate_event_study) - balanced_mask = None - if balance_e is not None: - all_horizons = sorted(set(int(h) for h in rel_times if np.isfinite(h))) - if self.horizon_max is not None: - all_horizons = [h for h in all_horizons if abs(h) <= self.horizon_max] - cohort_rel_times = self._build_cohort_rel_times(df, first_treat) - balanced_mask = self._compute_balanced_cohort_mask( - df_1, first_treat, all_horizons, balance_e, cohort_rel_times - ) - - ref_period = -1 - self.anticipation + # Compute combined horizon set (pre + post) matching analytical path + post_horizons = sorted(set(int(h) for h in rel_times if np.isfinite(h))) # Pre-period data for pretrends bootstrap pre_rel_times = None - pre_tau_hat = None pre_balanced_mask = None n_0 = int(omega_0_mask.sum()) + df_pre = None + pre_horizons: list = [] if self.pretrends: omega_pre_mask = ~df["_never_treated"] & ~df["_treated"] if omega_pre_mask.any(): df_pre = df.loc[omega_pre_mask] pre_rel_times = df_pre["_rel_time"].values - pre_tau_hat = df["_tau_hat"].loc[omega_pre_mask].values - if balance_e is not None: - all_h_pre = sorted( - set(int(h) for h in pre_rel_times if np.isfinite(h)) - ) - if self.horizon_max is not None: - all_h_pre = [ - h for h in all_h_pre if abs(h) <= self.horizon_max - ] - cohort_rel_times_pre = self._build_cohort_rel_times( - df, first_treat - ) - pre_balanced_mask = self._compute_balanced_cohort_mask( - df_pre, - first_treat, - all_h_pre, - balance_e, - cohort_rel_times_pre, - ) + pre_horizons = sorted( + set(int(h) for h in pre_rel_times if np.isfinite(h)) + ) + + # Use union of pre + post horizons for balance_e (matches analytical path) + all_horizons = sorted(set(post_horizons + pre_horizons)) + if self.horizon_max is not None: + all_horizons = [h for h in all_horizons if abs(h) <= self.horizon_max] + + # Balanced cohort mask (same logic as _aggregate_event_study) + balanced_mask = None + if balance_e is not None: + cohort_rel_times = self._build_cohort_rel_times(df, first_treat) + balanced_mask = self._compute_balanced_cohort_mask( + df_1, first_treat, all_horizons, balance_e, cohort_rel_times + ) + if df_pre is not None: + pre_balanced_mask = self._compute_balanced_cohort_mask( + df_pre, first_treat, all_horizons, balance_e, + cohort_rel_times, + ) + + ref_period = -1 - self.anticipation for h in event_study_effects: if event_study_effects[h].get("n_obs", 0) == 0: @@ -234,10 +229,7 @@ def _precompute_bootstrap_psi( # Build preperiod_weights_h on omega_0 omega_0_indices = np.where(omega_0_mask.values)[0] pre_positions_in_0 = np.zeros(len(df), dtype=bool) - df_pre_subset = df.loc[ - (~df["_never_treated"]) & (~df["_treated"]) - ] - pre_positions_in_0[df_pre_subset.index[h_mask_pre]] = True + pre_positions_in_0[df_pre.index[h_mask_pre]] = True pre_in_0_mask = pre_positions_in_0[omega_0_indices] preperiod_weights_h = np.zeros(n_0) diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index 47eada2f3..4d3b8a6ee 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -87,6 +87,11 @@ class TwoStageDiD(TwoStageDiDBootstrapMixin): horizon_max : int, optional Maximum event-study horizon. If set, event study effects are only computed for |h| <= horizon_max. + pretrends : bool, default=False + If True, event study includes pre-treatment horizons for visual + pre-trends assessment. Pre-period effects should be ~0 under + parallel trends. Only affects event_study aggregation; overall + ATT and group aggregation are unchanged. Attributes ---------- diff --git a/tests/test_pretrends_event_study.py b/tests/test_pretrends_event_study.py index e13319f9e..65fdfdc63 100644 --- a/tests/test_pretrends_event_study.py +++ b/tests/test_pretrends_event_study.py @@ -490,3 +490,131 @@ def test_point_estimates_close(self): atol=1e-6, err_msg=f"h={h}: point estimates differ", ) + + +# ============================================================================= +# Regression tests for P0/P1 review findings +# ============================================================================= + + +class TestPretrends_Regressions: + """Regression tests for bugs found in AI code review.""" + + def test_imputation_survey_weighted_no_covariates(self): + """ImputationDiD with survey weights, no covariates, no pretrends. + + Regression for P1: untreated_units/untreated_times uninitialized + in the survey-weighted FE-only variance path. + """ + from diff_diff.survey import SurveyDesign + + data = generate_test_data(seed=42) + rng = np.random.default_rng(42) + n_units = data["unit"].nunique() + unit_weights = rng.uniform(0.5, 2.0, n_units) + data["weight"] = data["unit"].map( + dict(enumerate(unit_weights)) + ) + + sd = SurveyDesign(weights="weight") + est = ImputationDiD() + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + survey_design=sd, + ) + + assert np.isfinite(results.overall_att) + assert np.isfinite(results.overall_se) + assert results.overall_se > 0 + + def test_imputation_survey_weighted_event_study(self): + """ImputationDiD with survey weights and event study aggregation.""" + from diff_diff.survey import SurveyDesign + + data = generate_test_data(seed=42) + rng = np.random.default_rng(42) + n_units = data["unit"].nunique() + unit_weights = rng.uniform(0.5, 2.0, n_units) + data["weight"] = data["unit"].map( + dict(enumerate(unit_weights)) + ) + + sd = SurveyDesign(weights="weight") + est = ImputationDiD() + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + survey_design=sd, + ) + + for h, eff in results.event_study_effects.items(): + if eff["n_obs"] == 0: + continue + assert np.isfinite(eff["se"]), f"h={h}: SE not finite" + + def test_imputation_pretrends_with_covariates(self): + """ImputationDiD pretrends=True with covariates. + + Tests _compute_v_untreated_with_covariates_preperiod(). + """ + data = generate_test_data(seed=42) + rng = np.random.default_rng(42) + data["x1"] = rng.standard_normal(len(data)) + + est = ImputationDiD(pretrends=True) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + covariates=["x1"], + aggregate="event_study", + ) + + negative = { + h: v + for h, v in results.event_study_effects.items() + if h < -1 and v["n_obs"] > 0 + } + assert len(negative) > 0, "Should have pre-period horizons" + for h, eff in negative.items(): + assert np.isfinite(eff["se"]), f"h={h}: SE not finite with covariates" + assert eff["se"] > 0, f"h={h}: SE not positive with covariates" + + def test_imputation_pretrends_balance_e_bootstrap(self): + """ImputationDiD pretrends=True + balance_e + bootstrap. + + Regression for P0: bootstrap must use the same balance_e cohort + filter (union of pre + post horizons) as the analytical path. + """ + data = generate_test_data(seed=42) + est = ImputationDiD(pretrends=True, n_bootstrap=50, seed=42) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + balance_e=1, + ) + + assert results.event_study_effects is not None + # Verify pre-period horizons have bootstrap inference + negative = { + h: v + for h, v in results.event_study_effects.items() + if h < -1 and v["n_obs"] > 0 + } + for h, eff in negative.items(): + assert np.isfinite(eff["se"]), f"h={h}: SE not finite" + assert eff["se"] > 0, f"h={h}: SE not positive" From 9d8025d1f3175eb200b86dad1bc77e5acd88bbc4 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 1 Apr 2026 16:26:47 -0400 Subject: [PATCH 3/7] Fix P0: use positional indexing for pre-period target-to-omega_0 mapping Replace label-based df_pre.index[h_mask_pre] with positional np.where(omega_pre_mask.values)[0][h_mask_pre] in both analytical and bootstrap paths. The previous code used DataFrame labels as positional array indices, which corrupts pre-period SEs/CIs when users pass DataFrames with non-default index (gapped, permuted, or filtered). Add regression tests verifying identical inference across RangeIndex, permuted, and gapped index types for both analytical and bootstrap paths. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/imputation.py | 10 ++- diff_diff/imputation_bootstrap.py | 8 +- tests/test_pretrends_event_study.py | 113 ++++++++++++++++++++++++++++ 3 files changed, 125 insertions(+), 6 deletions(-) diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 017673d2c..84bd63ca2 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -1831,12 +1831,14 @@ def _aggregate_event_study( # SE via conservative variance with preperiod_weights n_0 = int(omega_0_mask.sum()) preperiod_weights_h = np.zeros(n_0) - # Map h_mask_pre (indices in df_pre) to indices in omega_0 + # Positional mapping: pre-period positions in df → omega_0 omega_0_indices = np.where(omega_0_mask.values)[0] - pre_in_0 = (~df["_never_treated"] & ~df["_treated"]).values - # df_pre indices map to omega_0 positions + omega_pre_positions = np.where( + (~df["_never_treated"] & ~df["_treated"]).values + )[0] + pre_positions_in_df = omega_pre_positions[h_mask_pre] pre_positions_in_0 = np.zeros(len(df), dtype=bool) - pre_positions_in_0[df_pre.index[h_mask_pre]] = True + pre_positions_in_0[pre_positions_in_df] = True pre_in_0_mask = pre_positions_in_0[omega_0_indices] if survey_weights is not None: diff --git a/diff_diff/imputation_bootstrap.py b/diff_diff/imputation_bootstrap.py index d8f780c52..b698bdcc0 100644 --- a/diff_diff/imputation_bootstrap.py +++ b/diff_diff/imputation_bootstrap.py @@ -226,10 +226,14 @@ def _precompute_bootstrap_psi( if pre_balanced_mask is not None: h_mask_pre = h_mask_pre & pre_balanced_mask - # Build preperiod_weights_h on omega_0 + # Build preperiod_weights_h on omega_0 (positional mapping) omega_0_indices = np.where(omega_0_mask.values)[0] + omega_pre_positions = np.where( + (~df["_never_treated"] & ~df["_treated"]).values + )[0] + pre_positions_in_df = omega_pre_positions[h_mask_pre] pre_positions_in_0 = np.zeros(len(df), dtype=bool) - pre_positions_in_0[df_pre.index[h_mask_pre]] = True + pre_positions_in_0[pre_positions_in_df] = True pre_in_0_mask = pre_positions_in_0[omega_0_indices] preperiod_weights_h = np.zeros(n_0) diff --git a/tests/test_pretrends_event_study.py b/tests/test_pretrends_event_study.py index 65fdfdc63..e706c94f5 100644 --- a/tests/test_pretrends_event_study.py +++ b/tests/test_pretrends_event_study.py @@ -618,3 +618,116 @@ def test_imputation_pretrends_balance_e_bootstrap(self): for h, eff in negative.items(): assert np.isfinite(eff["se"]), f"h={h}: SE not finite" assert eff["se"] > 0, f"h={h}: SE not positive" + + def test_nondefault_index_analytical(self): + """Pre-period inference identical across index types (analytical). + + Regression for P0: label-based indexing in pre-period target mapping + corrupts inference when DataFrame has non-default index. + """ + data = generate_test_data(seed=42) + + # Baseline: default RangeIndex + results_default = ImputationDiD(pretrends=True).fit( + data.copy(), + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + # Permuted integer index + rng = np.random.default_rng(42) + data_perm = data.copy() + data_perm.index = rng.permutation(len(data_perm)) + results_perm = ImputationDiD(pretrends=True).fit( + data_perm, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + # Gapped/nonconsecutive index + data_gap = data.copy() + data_gap.index = data_gap.index * 3 + 100 + results_gap = ImputationDiD(pretrends=True).fit( + data_gap, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + # All pre-period effects and SEs must match + for h in results_default.event_study_effects: + if results_default.event_study_effects[h]["n_obs"] == 0: + continue + eff_d = results_default.event_study_effects[h] + eff_p = results_perm.event_study_effects[h] + eff_g = results_gap.event_study_effects[h] + np.testing.assert_allclose( + eff_d["effect"], eff_p["effect"], rtol=1e-10, + err_msg=f"h={h}: permuted index changed effect", + ) + np.testing.assert_allclose( + eff_d["se"], eff_p["se"], rtol=1e-10, + err_msg=f"h={h}: permuted index changed SE", + ) + np.testing.assert_allclose( + eff_d["effect"], eff_g["effect"], rtol=1e-10, + err_msg=f"h={h}: gapped index changed effect", + ) + np.testing.assert_allclose( + eff_d["se"], eff_g["se"], rtol=1e-10, + err_msg=f"h={h}: gapped index changed SE", + ) + + def test_nondefault_index_bootstrap(self): + """Pre-period bootstrap inference identical across index types. + + Regression for P0: same indexing bug in bootstrap path. + """ + data = generate_test_data(seed=42) + + results_default = ImputationDiD( + pretrends=True, n_bootstrap=50, seed=99 + ).fit( + data.copy(), + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + rng = np.random.default_rng(42) + data_perm = data.copy() + data_perm.index = rng.permutation(len(data_perm)) + results_perm = ImputationDiD( + pretrends=True, n_bootstrap=50, seed=99 + ).fit( + data_perm, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + for h in results_default.event_study_effects: + if results_default.event_study_effects[h]["n_obs"] == 0: + continue + eff_d = results_default.event_study_effects[h] + eff_p = results_perm.event_study_effects[h] + np.testing.assert_allclose( + eff_d["effect"], eff_p["effect"], rtol=1e-10, + err_msg=f"h={h}: permuted index changed effect", + ) + np.testing.assert_allclose( + eff_d["se"], eff_p["se"], rtol=1e-10, + err_msg=f"h={h}: permuted index changed bootstrap SE", + ) From b028e9845ede93458e3470c06530291bc82ca4b8 Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 1 Apr 2026 20:06:45 -0400 Subject: [PATCH 4/7] Reimplement ImputationDiD pre-period coefficients using BJS Test 1 Address CI review P1: replace placebo-style residual averages with BJS Test 1 lead regression (Equation 9). Pre-period coefficients now come from within-transformed OLS on Omega_0 with cluster-robust SEs, consistent with pretrend_test(). This matches the cited methodology. - Add _compute_lead_coefficients() helper reused by both _pretrend_test() and _aggregate_event_study() - Remove placebo machinery: preperiod_weights, auxiliary residuals for pre-periods, covariate projection for pre-periods - Bootstrap no longer updates pre-period SEs (they come from lead regression) - Update REGISTRY.md to reference Test 1 / Equation 9 - Replace cross-estimator parity test with contract test against pretrend_test().lead_coefficients Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/imputation.py | 604 +++++++--------------------- diff_diff/imputation_bootstrap.py | 81 +--- docs/methodology/REGISTRY.md | 17 +- tests/test_pretrends_event_study.py | 102 ++--- 4 files changed, 220 insertions(+), 584 deletions(-) diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 84bd63ca2..472eb0362 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -455,25 +455,6 @@ def fit( df["_tau_hat"] = np.nan df.loc[omega_1_mask, "_tau_hat"] = tau_hat - # Pre-period tau_hat for eventually-treated units (pretrends feature) - omega_pre_mask = None - if self.pretrends and aggregate in ("event_study", "all"): - omega_pre_mask = ~df["_never_treated"] & ~df["_treated"] - if omega_pre_mask.any(): - tau_hat_pre, _ = self._impute_treatment_effects( - df, - outcome, - unit, - time, - covariates, - omega_pre_mask, - unit_fe, - time_fe, - grand_mean, - delta_hat, - ) - df.loc[omega_pre_mask, "_tau_hat"] = tau_hat_pre - # ---- Step 3: Aggregate ---- # Always compute overall ATT (simple aggregation) finite_mask = np.isfinite(tau_hat) @@ -1103,20 +1084,12 @@ def _compute_cluster_psi_sums( cluster_var: str, kept_cov_mask: Optional[np.ndarray] = None, survey_weights_0: Optional[np.ndarray] = None, - preperiod_weights: Optional[np.ndarray] = None, ) -> Tuple[np.ndarray, np.ndarray]: """ Compute cluster-level influence function sums (Theorem 3). psi_i = sum_t v_it * epsilon_tilde_it, summed within each cluster. - Parameters - ---------- - preperiod_weights : np.ndarray, optional - Shape (n_0,). Non-zero for pre-period target observations in omega_0. - When provided, these observations contribute both directly (via w_it) - and indirectly (via FE correction) to the influence function. - Returns ------- cluster_psi_sums : np.ndarray @@ -1129,55 +1102,30 @@ def _compute_cluster_psi_sums( n_0 = len(df_0) n_1 = len(df_1) - is_preperiod = preperiod_weights is not None - # ---- Compute v_it for treated observations ---- - if is_preperiod: - v_treated = np.zeros(n_1) # omega_1 contributes nothing - else: - v_treated = weights.copy() + v_treated = weights.copy() # ---- Compute v_it for untreated observations ---- if covariates is None or len(covariates) == 0: # FE-only case: closed-form # Build w_by_unit, w_by_time, w_total from the target weights - if is_preperiod: - # Target weights are on omega_0 observations - untreated_units = df_0[unit].values - untreated_times = df_0[time].values - - w_by_unit: Dict[Any, float] = {} - for j in range(n_0): - u = untreated_units[j] - w_by_unit[u] = w_by_unit.get(u, 0.0) + preperiod_weights[j] - - w_by_time: Dict[Any, float] = {} - for j in range(n_0): - t = untreated_times[j] - w_by_time[t] = w_by_time.get(t, 0.0) + preperiod_weights[j] - - w_total = float(np.sum(preperiod_weights)) - else: - # Target weights are on omega_1 observations - treated_units = df_1[unit].values - treated_times = df_1[time].values + treated_units = df_1[unit].values + treated_times = df_1[time].values - w_by_unit = {} - for i_idx in range(n_1): - u = treated_units[i_idx] - w_by_unit[u] = w_by_unit.get(u, 0.0) + weights[i_idx] + w_by_unit: Dict[Any, float] = {} + for i_idx in range(n_1): + u = treated_units[i_idx] + w_by_unit[u] = w_by_unit.get(u, 0.0) + weights[i_idx] - w_by_time = {} - for i_idx in range(n_1): - t = treated_times[i_idx] - w_by_time[t] = w_by_time.get(t, 0.0) + weights[i_idx] + w_by_time: Dict[Any, float] = {} + for i_idx in range(n_1): + t = treated_times[i_idx] + w_by_time[t] = w_by_time.get(t, 0.0) + weights[i_idx] - w_total = float(np.sum(weights)) + w_total = float(np.sum(weights)) - # Always initialize untreated unit/time arrays for the FE loop - if not is_preperiod: - untreated_units = df_0[unit].values - untreated_times = df_0[time].values + untreated_units = df_0[unit].values + untreated_times = df_0[time].values # Use survey-weighted sums for untreated denominators when present if survey_weights_0 is not None: @@ -1203,35 +1151,19 @@ def _compute_cluster_psi_sums( # WLS projection requires per-obs survey weight factor if survey_weights_0 is not None: base_v *= survey_weights_0[j] - # For preperiod targets, add direct weight contribution - if is_preperiod: - v_untreated[j] = preperiod_weights[j] + base_v - else: - v_untreated[j] = base_v + v_untreated[j] = base_v else: - if is_preperiod: - v_untreated = self._compute_v_untreated_with_covariates_preperiod( - df_0, - unit, - time, - covariates, - preperiod_weights, - delta_hat, - kept_cov_mask=kept_cov_mask, - survey_weights_0=survey_weights_0, - ) - else: - v_untreated = self._compute_v_untreated_with_covariates( - df_0, - df_1, - unit, - time, - covariates, - weights, - delta_hat, - kept_cov_mask=kept_cov_mask, - survey_weights_0=survey_weights_0, - ) + v_untreated = self._compute_v_untreated_with_covariates( + df_0, + df_1, + unit, + time, + covariates, + weights, + delta_hat, + kept_cov_mask=kept_cov_mask, + survey_weights_0=survey_weights_0, + ) # ---- Compute auxiliary model residuals (Equation 8) ---- epsilon_treated = self._compute_auxiliary_residuals_treated( @@ -1251,25 +1183,6 @@ def _compute_cluster_psi_sums( df_0, outcome, unit, time, covariates, unit_fe, time_fe, grand_mean, delta_hat ) - # Override epsilon for pre-period target observations with auxiliary residuals - if is_preperiod: - target_mask_in_0 = preperiod_weights > 0 - if target_mask_in_0.any(): - aux_resid = self._compute_auxiliary_residuals_preperiod( - df_0.loc[df_0.index[target_mask_in_0]], - outcome, - unit, - time, - first_treat, - covariates, - unit_fe, - time_fe, - grand_mean, - delta_hat, - v_untreated[target_mask_in_0], - ) - epsilon_untreated[target_mask_in_0] = aux_resid - # ---- psi_it = v_it * epsilon_tilde_it ---- v_all = np.empty(len(df)) v_all[omega_1_mask.values] = v_treated @@ -1309,7 +1222,6 @@ def _compute_conservative_variance( cluster_var: str, kept_cov_mask: Optional[np.ndarray] = None, survey_weights: Optional[np.ndarray] = None, - preperiod_weights: Optional[np.ndarray] = None, ) -> float: """ Compute conservative clustered variance (Theorem 3, Equation 7). @@ -1322,10 +1234,6 @@ def _compute_conservative_variance( survey_weights : np.ndarray, optional Full-panel survey weights. When provided, untreated denominators in v_it use survey-weighted sums instead of raw counts. - preperiod_weights : np.ndarray, optional - Aggregation weights for pre-period target observations in omega_0. - Shape: (n_untreated,). When provided, these observations are both - in the aggregation target and in the FE estimation sample. Returns ------- @@ -1350,7 +1258,6 @@ def _compute_conservative_variance( cluster_var=cluster_var, kept_cov_mask=kept_cov_mask, survey_weights_0=sw_0, - preperiod_weights=preperiod_weights, ) sigma_sq = float((cluster_psi_sums**2).sum()) return np.sqrt(max(sigma_sq, 0.0)) @@ -1545,148 +1452,6 @@ def _compute_residuals_untreated( return df_0[outcome].values - y_hat - def _compute_auxiliary_residuals_preperiod( - self, - df_pre: pd.DataFrame, - outcome: str, - unit: str, - time: str, - first_treat: str, - covariates: Optional[List[str]], - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - grand_mean: float, - delta_hat: Optional[np.ndarray], - v_pre: np.ndarray, - ) -> np.ndarray: - """ - Compute auxiliary residuals for pre-period target observations. - - Analogous to _compute_auxiliary_residuals_treated but operates on - pre-treatment observations of eventually-treated units in omega_0. - - epsilon_tilde_it = tau_hat_pre_it - tau_tilde_g - where tau_tilde_g = sum(v_it * tau_hat_pre_it) / sum(v_it) within group g. - """ - n_pre = len(df_pre) - - # Compute tau_hat_pre = Y - Y_hat(0) (same as Step 1 residual) - alpha_i = df_pre[unit].map(unit_fe).values.astype(float) - beta_t = df_pre[time].map(time_fe).values.astype(float) - y_hat_0 = grand_mean + alpha_i + beta_t - - if delta_hat is not None and covariates: - y_hat_0 = y_hat_0 + np.dot(df_pre[covariates].values, delta_hat) - - tau_hat_pre = df_pre[outcome].values - y_hat_0 - - # Partition by aux_partition and compute group means - if self.aux_partition == "cohort_horizon": - group_keys = list( - zip(df_pre[first_treat].values, df_pre["_rel_time"].values) - ) - elif self.aux_partition == "cohort": - group_keys = list(df_pre[first_treat].values) - elif self.aux_partition == "horizon": - group_keys = list(df_pre["_rel_time"].values) - else: - group_keys = list(range(n_pre)) - - group_series = pd.Series(group_keys, index=df_pre.index) - tau_series = pd.Series(tau_hat_pre, index=df_pre.index) - v_series = pd.Series(v_pre, index=df_pre.index) - - weighted_tau_sum = (v_series * tau_series).groupby(group_series).sum() - weight_sum = v_series.groupby(group_series).sum() - - zero_weight_groups = weight_sum.abs() < 1e-15 - if zero_weight_groups.any(): - simple_means = tau_series.groupby(group_series).mean() - tau_tilde_map = weighted_tau_sum / weight_sum - tau_tilde_map = tau_tilde_map.where(~zero_weight_groups, simple_means) - else: - tau_tilde_map = weighted_tau_sum / weight_sum - - tau_tilde = group_series.map(tau_tilde_map).values - return tau_hat_pre - tau_tilde - - def _compute_v_untreated_with_covariates_preperiod( - self, - df_0: pd.DataFrame, - unit: str, - time: str, - covariates: List[str], - preperiod_weights: np.ndarray, - delta_hat: Optional[np.ndarray], - kept_cov_mask: Optional[np.ndarray] = None, - survey_weights_0: Optional[np.ndarray] = None, - ) -> np.ndarray: - """ - Compute v_it for untreated observations when covariates present (preperiod case). - - v = preperiod_weights - [W_0] A_0 (A_0'[W]A_0)^{-1} A_0' preperiod_weights - This is (I - P_0) * preperiod_weights where P_0 is the hat matrix. - """ - if kept_cov_mask is not None and not np.all(kept_cov_mask): - covariates = [c for c, k in zip(covariates, kept_cov_mask) if k] - - units_0 = df_0[unit].values - times_0 = df_0[time].values - - all_units = np.unique(units_0) - all_times = np.unique(times_0) - unit_to_idx = {u: i for i, u in enumerate(all_units)} - time_to_idx = {t: i for i, t in enumerate(all_times)} - n_units = len(all_units) - n_fe_cols = (n_units - 1) + (len(all_times) - 1) - - n_0 = len(df_0) - n_cov = len(covariates) - - # Build A_0 sparse - u_indices = np.array([unit_to_idx[u] for u in units_0]) - u_mask = u_indices > 0 - u_rows = np.arange(n_0)[u_mask] - u_cols = u_indices[u_mask] - 1 - - t_indices = np.array([time_to_idx[t] for t in times_0]) - t_mask = t_indices > 0 - t_rows = np.arange(n_0)[t_mask] - t_cols = (n_units - 1) + t_indices[t_mask] - 1 - - rows = np.concatenate([u_rows, t_rows]) - cols = np.concatenate([u_cols, t_cols]) - data = np.ones(len(rows)) - - A_fe = sparse.csr_matrix((data, (rows, cols)), shape=(n_0, n_fe_cols)) - - if n_cov > 0: - A_cov = sparse.csr_matrix(df_0[covariates].values) - A_0 = sparse.hstack([A_fe, A_cov], format="csr") - else: - A_0 = A_fe - - # Compute A_0' preperiod_weights - A0_w = A_0.T @ preperiod_weights # shape (p,) - - # Solve (A_0'[W]A_0) z = A_0' preperiod_weights - if survey_weights_0 is not None: - A0tA0 = A_0.T @ A_0.multiply(survey_weights_0[:, None]) - else: - A0tA0 = A_0.T @ A_0 - try: - z = spsolve(A0tA0.tocsc(), A0_w) - except Exception: - A0tA0_dense = A0tA0.toarray() - z, _, _, _ = np.linalg.lstsq(A0tA0_dense, A0_w, rcond=None) - - # v = preperiod_weights - [W_0] A_0 z - projection = A_0 @ z - if survey_weights_0 is not None: - projection = projection * survey_weights_0 - v_untreated = preperiod_weights - projection - return v_untreated - # ========================================================================= # Aggregation # ========================================================================= @@ -1717,25 +1482,8 @@ def _aggregate_event_study( tau_hat = df["_tau_hat"].loc[omega_1_mask].values rel_times = df_1["_rel_time"].values - # Get post-treatment horizons from omega_1 - post_horizons = sorted(set(int(h) for h in rel_times if np.isfinite(h))) - - # Pre-period horizons (pretrends feature) - pre_horizons: List[int] = [] - df_pre = None - tau_hat_pre = None - rel_times_pre = None - if self.pretrends: - omega_pre_mask = ~df["_never_treated"] & ~df["_treated"] - if omega_pre_mask.any(): - df_pre = df.loc[omega_pre_mask] - tau_hat_pre = df["_tau_hat"].loc[omega_pre_mask].values - rel_times_pre = df_pre["_rel_time"].values - pre_horizons = sorted( - set(int(h) for h in rel_times_pre if np.isfinite(h)) - ) - - all_horizons = sorted(set(post_horizons + pre_horizons)) + # Get all horizons + all_horizons = sorted(set(int(h) for h in rel_times if np.isfinite(h))) # Apply horizon_max filter if self.horizon_max is not None: @@ -1750,20 +1498,8 @@ def _aggregate_event_study( ), index=df_1.index, ) - if df_pre is not None: - balanced_mask_pre = pd.Series( - self._compute_balanced_cohort_mask( - df_pre, first_treat, all_horizons, balance_e, cohort_rel_times - ), - index=df_pre.index, - ) - else: - balanced_mask_pre = None else: balanced_mask = pd.Series(True, index=df_1.index) - balanced_mask_pre = ( - pd.Series(True, index=df_pre.index) if df_pre is not None else None - ) # Check Proposition 5: no never-treated units has_never_treated = df["_never_treated"].any() @@ -1786,6 +1522,27 @@ def _aggregate_event_study( "n_obs": 0, } + # Pre-period coefficients via BJS Test 1 lead regression + if self.pretrends: + df_0 = df.loc[omega_0_mask].copy() + rel_time_0 = np.where( + ~df_0["_never_treated"], + df_0[time] - df_0[first_treat], + np.nan, + ) + pre_rel_times = sorted( + set(int(h) for h in rel_time_0 if np.isfinite(h) and h < -self.anticipation) + ) + pre_rel_times = [h for h in pre_rel_times if h != ref_period] + if self.horizon_max is not None: + pre_rel_times = [h for h in pre_rel_times if abs(h) <= self.horizon_max] + if pre_rel_times: + pre_effects, _, _ = self._compute_lead_coefficients( + df_0, outcome, unit, time, first_treat, covariates, + cluster_var, pre_rel_times, alpha=self.alpha, + ) + event_study_effects.update(pre_effects) + # Collect horizons with Proposition 5 violations prop5_horizons = [] @@ -1793,115 +1550,6 @@ def _aggregate_event_study( if h == ref_period: continue - # ---- Pre-period horizon (from omega_pre observations) ---- - if h < -self.anticipation: - if rel_times_pre is None or balanced_mask_pre is None: - continue - h_mask_pre = (rel_times_pre == h) & balanced_mask_pre.values - n_h = int(h_mask_pre.sum()) - if n_h == 0: - continue - - tau_h = tau_hat_pre[h_mask_pre] - finite_h = np.isfinite(tau_h) - valid_tau = tau_h[finite_h] - - if len(valid_tau) == 0: - event_study_effects[h] = { - "effect": np.nan, - "se": np.nan, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (np.nan, np.nan), - "n_obs": n_h, - } - continue - - # Effect: survey-weighted or simple mean - if survey_weights is not None: - pre_sw = survey_weights[ - (~df["_never_treated"] & ~df["_treated"]).values - ] - sw_h = pre_sw[h_mask_pre] - sw_valid = sw_h[finite_h] - effect = float(np.average(valid_tau, weights=sw_valid)) - else: - effect = float(np.mean(valid_tau)) - - # SE via conservative variance with preperiod_weights - n_0 = int(omega_0_mask.sum()) - preperiod_weights_h = np.zeros(n_0) - # Positional mapping: pre-period positions in df → omega_0 - omega_0_indices = np.where(omega_0_mask.values)[0] - omega_pre_positions = np.where( - (~df["_never_treated"] & ~df["_treated"]).values - )[0] - pre_positions_in_df = omega_pre_positions[h_mask_pre] - pre_positions_in_0 = np.zeros(len(df), dtype=bool) - pre_positions_in_0[pre_positions_in_df] = True - pre_in_0_mask = pre_positions_in_0[omega_0_indices] - - if survey_weights is not None: - sw_0 = survey_weights[omega_0_mask.values] - sw_target = sw_0[pre_in_0_mask] - finite_in_target = np.isfinite( - df["_tau_hat"].values[omega_0_indices[pre_in_0_mask]] - ) - sw_finite = sw_target[finite_in_target] - if sw_finite.sum() > 0: - target_indices = np.where(pre_in_0_mask)[0] - finite_indices = target_indices[finite_in_target] - preperiod_weights_h[finite_indices] = ( - sw_finite / sw_finite.sum() - ) - else: - n_finite = int(np.isfinite( - df["_tau_hat"].values[omega_0_indices[pre_in_0_mask]] - ).sum()) - if n_finite > 0: - tau_at_target = df["_tau_hat"].values[ - omega_0_indices[pre_in_0_mask] - ] - finite_target = np.isfinite(tau_at_target) - target_indices = np.where(pre_in_0_mask)[0] - finite_indices = target_indices[finite_target] - preperiod_weights_h[finite_indices] = 1.0 / n_finite - - 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=np.zeros(len(tau_hat)), - cluster_var=cluster_var, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, - preperiod_weights=preperiod_weights_h, - ) - - t_stat, p_value, conf_int = safe_inference( - effect, se, alpha=self.alpha, df=survey_df - ) - - event_study_effects[h] = { - "effect": effect, - "se": se, - "t_stat": t_stat, - "p_value": p_value, - "conf_int": conf_int, - "n_obs": n_h, - } - continue - - # ---- Post-treatment horizon (from omega_1 observations) ---- # Select treated obs at this horizon from balanced cohorts h_mask = (rel_times == h) & balanced_mask.values n_h = int(h_mask.sum()) @@ -2130,9 +1778,107 @@ def _aggregate_group( return group_effects # ========================================================================= - # Pre-trend test (Equation 9) + # Pre-trend test (Equation 9) & pre-period lead coefficients # ========================================================================= + def _compute_lead_coefficients( + self, + df_0: pd.DataFrame, + outcome: str, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + cluster_var: str, + pre_rel_times: List[int], + alpha: float = 0.05, + ) -> Tuple[Dict[int, Dict[str, Any]], np.ndarray, np.ndarray]: + """ + Compute pre-period lead coefficients via within-transformed OLS (Test 1). + + Adds lead indicator dummies W_it(h) = 1[K_it = h] to the untreated + model and estimates their coefficients with cluster-robust SEs. + + Returns + ------- + effects : dict + Per-horizon event_study_effects entries. + gamma : ndarray + Lead coefficient vector. + V_gamma : ndarray + Sub-VCV matrix for lead coefficients. + """ + rel_time_0 = np.where( + ~df_0["_never_treated"], + df_0[time] - df_0[first_treat], + np.nan, + ) + + # Build lead indicators + lead_cols = [] + for h in pre_rel_times: + col_name = f"_lead_{h}" + df_0[col_name] = (rel_time_0 == h).astype(float) + lead_cols.append(col_name) + + # Within-transform via iterative demeaning + y_dm = self._iterative_demean( + df_0[outcome].values, df_0[unit].values, df_0[time].values, df_0.index + ) + + all_x_cols = lead_cols[:] + if covariates: + all_x_cols.extend(covariates) + + X_dm = np.column_stack( + [ + self._iterative_demean( + df_0[col].values, df_0[unit].values, df_0[time].values, df_0.index + ) + for col in all_x_cols + ] + ) + + # OLS with cluster-robust SEs + cluster_ids = df_0[cluster_var].values + result = solve_ols( + X_dm, + y_dm, + cluster_ids=cluster_ids, + return_vcov=True, + rank_deficient_action=self.rank_deficient_action, + column_names=all_x_cols, + ) + coefficients = result[0] + vcov = result[2] + assert vcov is not None + + n_leads = len(lead_cols) + gamma = coefficients[:n_leads] + V_gamma = vcov[:n_leads, :n_leads] + + # Build per-horizon effects + effects: Dict[int, Dict[str, Any]] = {} + for j, h in enumerate(pre_rel_times): + effect = float(gamma[j]) + se = float(np.sqrt(max(V_gamma[j, j], 0.0))) + n_obs = int((rel_time_0 == h).sum()) + t_stat, p_value, conf_int = safe_inference(effect, se, alpha=alpha) + effects[h] = { + "effect": effect, + "se": se, + "t_stat": t_stat, + "p_value": p_value, + "conf_int": conf_int, + "n_obs": n_obs, + } + + # Clean up temporary columns + for col in lead_cols: + df_0.drop(columns=col, inplace=True) + + return effects, gamma, V_gamma + def _pretrend_test(self, n_leads: Optional[int] = None) -> Dict[str, Any]: """ Run pre-trend test (Equation 9). @@ -2164,7 +1910,6 @@ def _pretrend_test(self, n_leads: Optional[int] = None) -> Dict[str, Any]: df_0 = df.loc[omega_0_mask].copy() # Compute relative time for untreated obs - # For not-yet-treated units in their pre-treatment periods rel_time_0 = np.where( ~df_0["_never_treated"], df_0[time] - df_0[first_treat], @@ -2190,7 +1935,6 @@ def _pretrend_test(self, n_leads: Optional[int] = None) -> Dict[str, Any]: pre_rel_times = [h for h in pre_rel_times if h != ref] if n_leads is not None: - # Take the n_leads periods closest to treatment pre_rel_times = sorted(pre_rel_times, reverse=True)[:n_leads] pre_rel_times = sorted(pre_rel_times) @@ -2203,49 +1947,13 @@ def _pretrend_test(self, n_leads: Optional[int] = None) -> Dict[str, Any]: "lead_coefficients": {}, } - # Build lead indicators - lead_cols = [] - for h in pre_rel_times: - col_name = f"_lead_{h}" - df_0[col_name] = ((rel_time_0 == h)).astype(float) - lead_cols.append(col_name) - - # Within-transform via iterative demeaning (exact for unbalanced panels) - y_dm = self._iterative_demean( - df_0[outcome].values, df_0[unit].values, df_0[time].values, df_0.index - ) - - all_x_cols = lead_cols[:] - if covariates: - all_x_cols.extend(covariates) - - X_dm = np.column_stack( - [ - self._iterative_demean( - df_0[col].values, df_0[unit].values, df_0[time].values, df_0.index - ) - for col in all_x_cols - ] - ) - - # OLS with cluster-robust SEs - cluster_ids = df_0[cluster_var].values - result = solve_ols( - X_dm, - y_dm, - cluster_ids=cluster_ids, - return_vcov=True, - rank_deficient_action=self.rank_deficient_action, - column_names=all_x_cols, + # Use shared lead coefficient computation + effects, gamma, V_gamma = self._compute_lead_coefficients( + df_0, outcome, unit, time, first_treat, covariates, + cluster_var, pre_rel_times, alpha=self.alpha, ) - coefficients = result[0] - vcov = result[2] - assert vcov is not None - # Extract lead coefficients and their sub-VCV - n_leads_actual = len(lead_cols) - gamma = coefficients[:n_leads_actual] - V_gamma = vcov[:n_leads_actual, :n_leads_actual] + n_leads_actual = len(pre_rel_times) # Wald F-test: F = (gamma' V^{-1} gamma) / n_leads try: @@ -2256,6 +1964,7 @@ def _pretrend_test(self, n_leads: Optional[int] = None) -> Dict[str, Any]: f_stat = np.nan # P-value from F distribution + cluster_ids = df_0[cluster_var].values if np.isfinite(f_stat) and f_stat >= 0: n_clusters = len(np.unique(cluster_ids)) df_denom = max(n_clusters - 1, 1) @@ -2263,10 +1972,7 @@ def _pretrend_test(self, n_leads: Optional[int] = None) -> Dict[str, Any]: else: p_value = np.nan - # Store lead coefficients - lead_coefficients = {} - for j, h in enumerate(pre_rel_times): - lead_coefficients[h] = float(gamma[j]) + lead_coefficients = {h: effects[h]["effect"] for h in pre_rel_times} return { "f_stat": f_stat, diff --git a/diff_diff/imputation_bootstrap.py b/diff_diff/imputation_bootstrap.py index b698bdcc0..e1623884d 100644 --- a/diff_diff/imputation_bootstrap.py +++ b/diff_diff/imputation_bootstrap.py @@ -174,26 +174,7 @@ def _precompute_bootstrap_psi( df_1 = df.loc[omega_1_mask] rel_times = df_1["_rel_time"].values - # Compute combined horizon set (pre + post) matching analytical path - post_horizons = sorted(set(int(h) for h in rel_times if np.isfinite(h))) - - # Pre-period data for pretrends bootstrap - pre_rel_times = None - pre_balanced_mask = None - n_0 = int(omega_0_mask.sum()) - df_pre = None - pre_horizons: list = [] - if self.pretrends: - omega_pre_mask = ~df["_never_treated"] & ~df["_treated"] - if omega_pre_mask.any(): - df_pre = df.loc[omega_pre_mask] - pre_rel_times = df_pre["_rel_time"].values - pre_horizons = sorted( - set(int(h) for h in pre_rel_times if np.isfinite(h)) - ) - - # Use union of pre + post horizons for balance_e (matches analytical path) - all_horizons = sorted(set(post_horizons + pre_horizons)) + all_horizons = sorted(set(int(h) for h in rel_times if np.isfinite(h))) if self.horizon_max is not None: all_horizons = [h for h in all_horizons if abs(h) <= self.horizon_max] @@ -204,11 +185,6 @@ def _precompute_bootstrap_psi( balanced_mask = self._compute_balanced_cohort_mask( df_1, first_treat, all_horizons, balance_e, cohort_rel_times ) - if df_pre is not None: - pre_balanced_mask = self._compute_balanced_cohort_mask( - df_pre, first_treat, all_horizons, balance_e, - cohort_rel_times, - ) ref_period = -1 - self.anticipation @@ -220,60 +196,11 @@ def _precompute_bootstrap_psi( if not np.isfinite(event_study_effects[h].get("effect", np.nan)): continue - # ---- Pre-period horizon ---- - if h < -self.anticipation and pre_rel_times is not None: - h_mask_pre = pre_rel_times == h - if pre_balanced_mask is not None: - h_mask_pre = h_mask_pre & pre_balanced_mask - - # Build preperiod_weights_h on omega_0 (positional mapping) - omega_0_indices = np.where(omega_0_mask.values)[0] - omega_pre_positions = np.where( - (~df["_never_treated"] & ~df["_treated"]).values - )[0] - pre_positions_in_df = omega_pre_positions[h_mask_pre] - pre_positions_in_0 = np.zeros(len(df), dtype=bool) - pre_positions_in_0[pre_positions_in_df] = True - pre_in_0_mask = pre_positions_in_0[omega_0_indices] - - preperiod_weights_h = np.zeros(n_0) - if survey_weights_0 is not None: - sw_0 = survey_weights_0 - sw_target = sw_0[pre_in_0_mask] - tau_at_target = df["_tau_hat"].values[ - omega_0_indices[pre_in_0_mask] - ] - finite_in_target = np.isfinite(tau_at_target) - sw_finite = sw_target[finite_in_target] - if sw_finite.sum() > 0: - target_idx = np.where(pre_in_0_mask)[0] - fin_idx = target_idx[finite_in_target] - preperiod_weights_h[fin_idx] = ( - sw_finite / sw_finite.sum() - ) - else: - continue - else: - tau_at_target = df["_tau_hat"].values[ - omega_0_indices[pre_in_0_mask] - ] - finite_target = np.isfinite(tau_at_target) - n_finite = int(finite_target.sum()) - if n_finite == 0: - continue - target_idx = np.where(pre_in_0_mask)[0] - fin_idx = target_idx[finite_target] - preperiod_weights_h[fin_idx] = 1.0 / n_finite - - psi_h, _ = self._compute_cluster_psi_sums( - **common, - weights=np.zeros(len(tau_hat)), - preperiod_weights=preperiod_weights_h, - ) - result["event_study"][h] = psi_h + # Skip pre-period horizons — their SEs come from Test 1 + # lead regression, not bootstrap + if h < -self.anticipation: continue - # ---- Post-treatment horizon ---- h_mask = rel_times == h if balanced_mask is not None: h_mask = h_mask & balanced_mask diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 23d844060..69b09a1d4 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -836,18 +836,19 @@ Y_it = alpha_i + beta_t [+ X'_it * delta] + W'_it * gamma + epsilon_it - Test `gamma = 0` via cluster-robust Wald F-test - Independent of treatment effect estimation (Proposition 9) -*Pre-period event study coefficients (`pretrends=True`, BJS Section 4.2):* +*Pre-period event study coefficients (`pretrends=True`, Test 1 / Equation 9):* -For pre-treatment observation (i,t) of eventually-treated unit i at relative time `K_it = t - E_i < -anticipation`: +Pre-period coefficients reuse the existing pre-trend test machinery (BJS Equation 9): ``` -tau_hat_pre(i,t) = Y_it - alpha_hat_i - beta_hat_t [- X'_it * delta_hat] -theta_h = (1/n_h) * sum_{(i,t): K_it = h} tau_hat_pre(i,t) +Y_it = alpha_i + beta_t [+ X'_it * delta] + sum_h gamma_h * W_it(h) + epsilon_it ``` -- Under parallel trends (Assumption 1), `E[theta_h] = 0` for all `h < -anticipation` -- Reference period `h = -1 - anticipation` normalized to zero (not estimated) -- Variance via Theorem 3 extension: pre-period targets are in `Omega_0`, so `v_it = w_it + FE_correction` (both direct aggregation weight and indirect FE estimation correction) +where `W_it(h) = 1[K_it = h]` are lead indicators, estimated on `Omega_0` only. +- `gamma_h` are the pre-period event study coefficients with cluster-robust SEs +- Under parallel trends (Assumption 1), `gamma_h = 0` for all `h < -anticipation` +- Reference period `h = -1 - anticipation` is the omitted category (normalized to zero) +- SEs from cluster-robust Wald variance (consistent with `pretrend_test()`) +- Bootstrap does not update pre-period SEs (they are from the lead regression) - Only affects event study aggregation; overall ATT and group aggregation unchanged -- Proposition 5 does not apply to pre-period horizons (always identified) *Edge cases:* - **Unbalanced panels:** FE estimated via iterative alternating projection (Gauss-Seidel), equivalent to OLS with unit+time dummies. Converges in O(max_iter) passes; typically 5-20 iterations for unbalanced panels, 1-2 for balanced. One-pass demeaning is only exact for balanced panels. diff --git a/tests/test_pretrends_event_study.py b/tests/test_pretrends_event_study.py index e706c94f5..8f80e816a 100644 --- a/tests/test_pretrends_event_study.py +++ b/tests/test_pretrends_event_study.py @@ -443,22 +443,15 @@ def test_anticipation_interaction(self): # ============================================================================= -class TestPretrendsCrossEstimator: - """Verify ImputationDiD and TwoStageDiD produce consistent pre-period effects.""" +class TestPretrends_ContractTests: + """Verify ImputationDiD pre-period coefficients match pretrend_test() lead coefficients.""" - def test_point_estimates_close(self): - """ImputationDiD and TwoStageDiD pre-period effects should be similar.""" + def test_imputation_pretrends_match_pretrend_test(self): + """ImputationDiD pretrends=True effects match pretrend_test().lead_coefficients.""" data = generate_test_data(seed=77) - imp = ImputationDiD(pretrends=True).fit( - data, - outcome="outcome", - unit="unit", - time="time", - first_treat="first_treat", - aggregate="event_study", - ) - ts = TwoStageDiD(pretrends=True).fit( + est = ImputationDiD(pretrends=True) + results = est.fit( data, outcome="outcome", unit="unit", @@ -467,28 +460,21 @@ def test_point_estimates_close(self): aggregate="event_study", ) - # Both should have negative horizons - imp_neg = { - h: v["effect"] - for h, v in imp.event_study_effects.items() - if h < -1 and v["n_obs"] > 0 - } - ts_neg = { - h: v["effect"] - for h, v in ts.event_study_effects.items() - if h < -1 and v["n_obs"] > 0 - } - - common_h = set(imp_neg.keys()) & set(ts_neg.keys()) - assert len(common_h) > 0, "No common pre-period horizons" + # Get pretrend_test lead coefficients (called on results object) + pt = results.pretrend_test() + lead_coefs = pt["lead_coefficients"] - for h in common_h: - # Point estimates should be numerically identical (same imputation) + # Every lead coefficient should match the event study pre-period effect + for h, coef in lead_coefs.items(): + assert h in results.event_study_effects, ( + f"h={h}: lead coefficient not in event study" + ) + es_effect = results.event_study_effects[h]["effect"] np.testing.assert_allclose( - imp_neg[h], - ts_neg[h], - atol=1e-6, - err_msg=f"h={h}: point estimates differ", + es_effect, + coef, + rtol=1e-10, + err_msg=f"h={h}: event study effect != pretrend_test lead coefficient", ) @@ -590,34 +576,50 @@ def test_imputation_pretrends_with_covariates(self): assert np.isfinite(eff["se"]), f"h={h}: SE not finite with covariates" assert eff["se"] > 0, f"h={h}: SE not positive with covariates" - def test_imputation_pretrends_balance_e_bootstrap(self): - """ImputationDiD pretrends=True + balance_e + bootstrap. + def test_imputation_pretrends_bootstrap_post_only(self): + """ImputationDiD pretrends=True + bootstrap: bootstrap updates post SEs only. - Regression for P0: bootstrap must use the same balance_e cohort - filter (union of pre + post horizons) as the analytical path. + Pre-period SEs come from Test 1 lead regression (cluster-robust), + not from bootstrap. Verify bootstrap doesn't break pre-period inference. """ data = generate_test_data(seed=42) - est = ImputationDiD(pretrends=True, n_bootstrap=50, seed=42) - results = est.fit( + + # Without bootstrap + results_no_boot = ImputationDiD(pretrends=True).fit( data, outcome="outcome", unit="unit", time="time", first_treat="first_treat", aggregate="event_study", - balance_e=1, ) - assert results.event_study_effects is not None - # Verify pre-period horizons have bootstrap inference - negative = { - h: v - for h, v in results.event_study_effects.items() - if h < -1 and v["n_obs"] > 0 - } - for h, eff in negative.items(): - assert np.isfinite(eff["se"]), f"h={h}: SE not finite" - assert eff["se"] > 0, f"h={h}: SE not positive" + # With bootstrap + results_boot = ImputationDiD( + pretrends=True, n_bootstrap=50, seed=42 + ).fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + # Pre-period effects and SEs should be identical (bootstrap doesn't touch them) + for h in results_no_boot.event_study_effects: + if h >= 0 or results_no_boot.event_study_effects[h]["n_obs"] == 0: + continue + eff_nb = results_no_boot.event_study_effects[h] + eff_b = results_boot.event_study_effects[h] + np.testing.assert_allclose( + eff_nb["effect"], eff_b["effect"], rtol=1e-10, + err_msg=f"h={h}: bootstrap changed pre-period effect", + ) + np.testing.assert_allclose( + eff_nb["se"], eff_b["se"], rtol=1e-10, + err_msg=f"h={h}: bootstrap changed pre-period SE", + ) def test_nondefault_index_analytical(self): """Pre-period inference identical across index types (analytical). From a26301835618ae5101577e785f4d9cdd2ed5fcba Mon Sep 17 00:00:00 2001 From: igerber Date: Wed, 1 Apr 2026 21:13:30 -0400 Subject: [PATCH 5/7] Address CI review: survey guard, balance_e propagation, regression tests P0: Raise NotImplementedError when pretrends=True + survey_design (lead regression is not yet survey-aware, matching pretrend_test() behavior). P1: Thread balance_e through pre-period lead regression by filtering Omega_0 to balanced cohorts before computing leads. Handle rank-deficient lead regression gracefully (returns NaN effects). P1: Add regression tests for pretrends + survey_design (NotImplementedError), pretrends + balance_e, and TwoStageDiD pretrends + bootstrap. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/imputation.py | 95 +++++++++++++++++++++-------- tests/test_pretrends_event_study.py | 75 +++++++++++++++++++++++ 2 files changed, 145 insertions(+), 25 deletions(-) diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 472eb0362..ae8a992d4 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -236,6 +236,14 @@ def fit( if missing: raise ValueError(f"Missing columns: {missing}") + if self.pretrends and survey_design is not None and aggregate in ("event_study", "all"): + raise NotImplementedError( + "pretrends=True is not yet compatible with survey_design. " + "The pre-period lead regression uses unweighted demeaning, " + "which does not account for survey weights. Use pretrends=False " + "with survey_design, or use pretrend_test() for pre-trend assessment." + ) + # Create working copy df = data.copy() @@ -1525,23 +1533,43 @@ def _aggregate_event_study( # Pre-period coefficients via BJS Test 1 lead regression if self.pretrends: df_0 = df.loc[omega_0_mask].copy() - rel_time_0 = np.where( - ~df_0["_never_treated"], - df_0[time] - df_0[first_treat], - np.nan, - ) - pre_rel_times = sorted( - set(int(h) for h in rel_time_0 if np.isfinite(h) and h < -self.anticipation) - ) - pre_rel_times = [h for h in pre_rel_times if h != ref_period] - if self.horizon_max is not None: - pre_rel_times = [h for h in pre_rel_times if abs(h) <= self.horizon_max] - if pre_rel_times: - pre_effects, _, _ = self._compute_lead_coefficients( - df_0, outcome, unit, time, first_treat, covariates, - cluster_var, pre_rel_times, alpha=self.alpha, + + # Apply balance_e restriction to pre-period lead regression + # so negative horizons use the same balanced cohort set + if balance_e is not None: + cohort_rel_times_0 = self._build_cohort_rel_times(df, first_treat) + balanced_cohorts = set() + if all_horizons: + max_h = max(all_horizons) + required_range = set(range(-balance_e, max_h + 1)) + for g, horizons in cohort_rel_times_0.items(): + if required_range.issubset(horizons): + balanced_cohorts.add(g) + if balanced_cohorts: + balance_mask_0 = df_0[first_treat].isin(balanced_cohorts) + df_0 = df_0.loc[balance_mask_0].copy() + else: + # No cohorts qualify — skip pre-period coefficients + df_0 = None + + if df_0 is not None and len(df_0) > 0: + rel_time_0 = np.where( + ~df_0["_never_treated"], + df_0[time] - df_0[first_treat], + np.nan, ) - event_study_effects.update(pre_effects) + pre_rel_times = sorted( + set(int(h) for h in rel_time_0 if np.isfinite(h) and h < -self.anticipation) + ) + pre_rel_times = [h for h in pre_rel_times if h != ref_period] + if self.horizon_max is not None: + pre_rel_times = [h for h in pre_rel_times if abs(h) <= self.horizon_max] + if pre_rel_times: + pre_effects, _, _ = self._compute_lead_coefficients( + df_0, outcome, unit, time, first_treat, covariates, + cluster_var, pre_rel_times, alpha=self.alpha, + ) + event_study_effects.update(pre_effects) # Collect horizons with Proposition 5 violations prop5_horizons = [] @@ -1841,14 +1869,31 @@ def _compute_lead_coefficients( # OLS with cluster-robust SEs cluster_ids = df_0[cluster_var].values - result = solve_ols( - X_dm, - y_dm, - cluster_ids=cluster_ids, - return_vcov=True, - rank_deficient_action=self.rank_deficient_action, - column_names=all_x_cols, - ) + try: + result = solve_ols( + X_dm, + y_dm, + cluster_ids=cluster_ids, + return_vcov=True, + rank_deficient_action=self.rank_deficient_action, + column_names=all_x_cols, + ) + except (IndexError, np.linalg.LinAlgError): + # All lead columns dropped (rank deficient after demeaning) + effects: Dict[int, Dict[str, Any]] = {} + for h in pre_rel_times: + n_obs = int((rel_time_0 == h).sum()) + effects[h] = { + "effect": np.nan, "se": np.nan, "t_stat": np.nan, + "p_value": np.nan, "conf_int": (np.nan, np.nan), + "n_obs": n_obs, + } + for col in lead_cols: + df_0.drop(columns=col, inplace=True) + return effects, np.full(len(pre_rel_times), np.nan), np.full( + (len(pre_rel_times), len(pre_rel_times)), np.nan + ) + coefficients = result[0] vcov = result[2] assert vcov is not None @@ -1858,7 +1903,7 @@ def _compute_lead_coefficients( V_gamma = vcov[:n_leads, :n_leads] # Build per-horizon effects - effects: Dict[int, Dict[str, Any]] = {} + effects = {} for j, h in enumerate(pre_rel_times): effect = float(gamma[j]) se = float(np.sqrt(max(V_gamma[j, j], 0.0))) diff --git a/tests/test_pretrends_event_study.py b/tests/test_pretrends_event_study.py index 8f80e816a..5173d80b3 100644 --- a/tests/test_pretrends_event_study.py +++ b/tests/test_pretrends_event_study.py @@ -733,3 +733,78 @@ def test_nondefault_index_bootstrap(self): eff_d["se"], eff_p["se"], rtol=1e-10, err_msg=f"h={h}: permuted index changed bootstrap SE", ) + + def test_imputation_pretrends_survey_raises(self): + """pretrends=True + survey_design raises NotImplementedError.""" + from diff_diff.survey import SurveyDesign + + data = generate_test_data(seed=42) + n_units = data["unit"].nunique() + unit_weights = np.random.default_rng(42).uniform(0.5, 2.0, n_units) + data["weight"] = data["unit"].map(dict(enumerate(unit_weights))) + + sd = SurveyDesign(weights="weight") + est = ImputationDiD(pretrends=True) + with pytest.raises(NotImplementedError, match="pretrends=True is not yet compatible"): + est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + survey_design=sd, + ) + + def test_imputation_pretrends_balance_e(self): + """balance_e restricts pre-period lead regression to balanced cohorts.""" + data = generate_test_data(seed=42) + + results_no_be = ImputationDiD(pretrends=True).fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + results_be = ImputationDiD(pretrends=True).fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + balance_e=1, + ) + + # Both should have pre-period horizons + neg_no_be = [h for h in results_no_be.event_study_effects if h < -1] + assert len(neg_no_be) > 0 + # balance_e restricts cohorts; pre-period coefficients may be NaN + # if the restricted sample is rank-deficient after demeaning + # The key invariant: the method runs without error + assert results_be.event_study_effects is not None + + def test_two_stage_pretrends_bootstrap(self): + """TwoStageDiD pretrends=True with bootstrap produces finite pre-period SEs.""" + data = generate_test_data(seed=42) + est = TwoStageDiD(pretrends=True, n_bootstrap=50, seed=42) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + negative = { + h: v + for h, v in results.event_study_effects.items() + if h < -1 and v["n_obs"] > 0 + } + assert len(negative) > 0, "Should have pre-period horizons" + for h, eff in negative.items(): + assert np.isfinite(eff["se"]), f"h={h}: SE not finite" + assert eff["se"] > 0, f"h={h}: SE not positive" From 2af1b65ce2cc346fb40cb5651e17624d93bf4d22 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 06:37:27 -0400 Subject: [PATCH 6/7] Fix balance_e: keep full Omega_0 for lead regression, restrict only lead indicators P1: balance_e pre-period path was dropping never-treated controls from Omega_0 before the lead regression. Now keeps the full untreated sample for within-transformation and restricts only which cohorts' lead indicators are activated (via balanced_cohorts parameter to _compute_lead_coefficients). P1: Add **Note:** in REGISTRY.md documenting pretrends + survey_design NotImplementedError. Fix error message to not suggest pretrend_test(). Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/imputation.py | 78 ++++++++++++++++++----------- docs/methodology/REGISTRY.md | 2 + tests/test_pretrends_event_study.py | 4 +- 3 files changed, 54 insertions(+), 30 deletions(-) diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index ae8a992d4..4ffbf694d 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -241,7 +241,7 @@ def fit( "pretrends=True is not yet compatible with survey_design. " "The pre-period lead regression uses unweighted demeaning, " "which does not account for survey weights. Use pretrends=False " - "with survey_design, or use pretrend_test() for pre-trend assessment." + "with survey_design for now." ) # Create working copy @@ -1534,8 +1534,11 @@ def _aggregate_event_study( if self.pretrends: df_0 = df.loc[omega_0_mask].copy() - # Apply balance_e restriction to pre-period lead regression - # so negative horizons use the same balanced cohort set + # Determine which cohorts' lead indicators to include. + # balance_e restricts which cohorts contribute lead dummies, + # but the full Omega_0 sample (including never-treated controls) + # is kept for the within-transformed OLS (BJS Test 1, Equation 9). + balanced_cohorts = None if balance_e is not None: cohort_rel_times_0 = self._build_cohort_rel_times(df, first_treat) balanced_cohorts = set() @@ -1545,31 +1548,35 @@ def _aggregate_event_study( for g, horizons in cohort_rel_times_0.items(): if required_range.issubset(horizons): balanced_cohorts.add(g) - if balanced_cohorts: - balance_mask_0 = df_0[first_treat].isin(balanced_cohorts) - df_0 = df_0.loc[balance_mask_0].copy() - else: - # No cohorts qualify — skip pre-period coefficients - df_0 = None - - if df_0 is not None and len(df_0) > 0: - rel_time_0 = np.where( - ~df_0["_never_treated"], - df_0[time] - df_0[first_treat], - np.nan, - ) - pre_rel_times = sorted( - set(int(h) for h in rel_time_0 if np.isfinite(h) and h < -self.anticipation) + if not balanced_cohorts: + balanced_cohorts = None # No cohorts qualify — skip + + rel_time_0 = np.where( + ~df_0["_never_treated"], + df_0[time] - df_0[first_treat], + np.nan, + ) + + # When balance_e is set, only include leads from balanced cohorts + if balanced_cohorts is not None: + is_balanced = df_0[first_treat].isin(balanced_cohorts).values + rel_time_for_leads = np.where(is_balanced, rel_time_0, np.nan) + else: + rel_time_for_leads = rel_time_0 + + pre_rel_times = sorted( + set(int(h) for h in rel_time_for_leads if np.isfinite(h) and h < -self.anticipation) + ) + pre_rel_times = [h for h in pre_rel_times if h != ref_period] + if self.horizon_max is not None: + pre_rel_times = [h for h in pre_rel_times if abs(h) <= self.horizon_max] + if pre_rel_times: + pre_effects, _, _ = self._compute_lead_coefficients( + df_0, outcome, unit, time, first_treat, covariates, + cluster_var, pre_rel_times, alpha=self.alpha, + balanced_cohorts=balanced_cohorts, ) - pre_rel_times = [h for h in pre_rel_times if h != ref_period] - if self.horizon_max is not None: - pre_rel_times = [h for h in pre_rel_times if abs(h) <= self.horizon_max] - if pre_rel_times: - pre_effects, _, _ = self._compute_lead_coefficients( - df_0, outcome, unit, time, first_treat, covariates, - cluster_var, pre_rel_times, alpha=self.alpha, - ) - event_study_effects.update(pre_effects) + event_study_effects.update(pre_effects) # Collect horizons with Proposition 5 violations prop5_horizons = [] @@ -1820,6 +1827,7 @@ def _compute_lead_coefficients( cluster_var: str, pre_rel_times: List[int], alpha: float = 0.05, + balanced_cohorts: Optional[set] = None, ) -> Tuple[Dict[int, Dict[str, Any]], np.ndarray, np.ndarray]: """ Compute pre-period lead coefficients via within-transformed OLS (Test 1). @@ -1827,6 +1835,10 @@ def _compute_lead_coefficients( Adds lead indicator dummies W_it(h) = 1[K_it = h] to the untreated model and estimates their coefficients with cluster-robust SEs. + The full Omega_0 sample (including never-treated controls) is always + used for within-transformation. When balanced_cohorts is provided, + lead indicators are restricted to observations from those cohorts only. + Returns ------- effects : dict @@ -1842,11 +1854,19 @@ def _compute_lead_coefficients( np.nan, ) - # Build lead indicators + # Build lead indicators — restrict to balanced cohorts if specified + if balanced_cohorts is not None: + is_balanced = df_0[first_treat].isin(balanced_cohorts).values + else: + is_balanced = None + lead_cols = [] for h in pre_rel_times: col_name = f"_lead_{h}" - df_0[col_name] = (rel_time_0 == h).astype(float) + indicator = (rel_time_0 == h).astype(float) + if is_balanced is not None: + indicator = indicator * is_balanced # zero out non-balanced cohorts + df_0[col_name] = indicator lead_cols.append(col_name) # Within-transform via iterative demeaning diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 69b09a1d4..7d1270b7e 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -848,7 +848,9 @@ where `W_it(h) = 1[K_it = h]` are lead indicators, estimated on `Omega_0` only. - Reference period `h = -1 - anticipation` is the omitted category (normalized to zero) - SEs from cluster-robust Wald variance (consistent with `pretrend_test()`) - Bootstrap does not update pre-period SEs (they are from the lead regression) +- When `balance_e` is set, lead indicators are restricted to balanced cohorts; the full Omega_0 sample (including never-treated) is kept for within-transformation - Only affects event study aggregation; overall ATT and group aggregation unchanged +- **Note:** `pretrends=True` with `survey_design` raises `NotImplementedError` because the lead regression uses unweighted demeaning (same limitation as `pretrend_test()`) *Edge cases:* - **Unbalanced panels:** FE estimated via iterative alternating projection (Gauss-Seidel), equivalent to OLS with unit+time dummies. Converges in O(max_iter) passes; typically 5-20 iterations for unbalanced panels, 1-2 for balanced. One-pass demeaning is only exact for balanced panels. diff --git a/tests/test_pretrends_event_study.py b/tests/test_pretrends_event_study.py index 5173d80b3..0e059bf08 100644 --- a/tests/test_pretrends_event_study.py +++ b/tests/test_pretrends_event_study.py @@ -745,7 +745,9 @@ def test_imputation_pretrends_survey_raises(self): sd = SurveyDesign(weights="weight") est = ImputationDiD(pretrends=True) - with pytest.raises(NotImplementedError, match="pretrends=True is not yet compatible"): + with pytest.raises( + NotImplementedError, match="pretrends=True is not yet compatible" + ): est.fit( data, outcome="outcome", From 0b65daa6038df51a9217c86ed78263757c7d9d95 Mon Sep 17 00:00:00 2001 From: igerber Date: Thu, 2 Apr 2026 06:54:51 -0400 Subject: [PATCH 7/7] Fix balance_e zero-qualifying-cohort fallback in pretrends path When no cohort satisfies balance_e, skip pre-period coefficient generation entirely instead of falling back to unrestricted leads. Also fix n_obs to use the lead indicator (which respects balanced_cohorts restriction) instead of the unrestricted relative-time count. Co-Authored-By: Claude Opus 4.6 (1M context) --- diff_diff/imputation.py | 68 +++++++++++++++++++++++++---------------- 1 file changed, 42 insertions(+), 26 deletions(-) diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 4ffbf694d..b54bf42dd 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -1539,6 +1539,7 @@ def _aggregate_event_study( # but the full Omega_0 sample (including never-treated controls) # is kept for the within-transformed OLS (BJS Test 1, Equation 9). balanced_cohorts = None + skip_preperiods = False if balance_e is not None: cohort_rel_times_0 = self._build_cohort_rel_times(df, first_treat) balanced_cohorts = set() @@ -1549,34 +1550,48 @@ def _aggregate_event_study( if required_range.issubset(horizons): balanced_cohorts.add(g) if not balanced_cohorts: - balanced_cohorts = None # No cohorts qualify — skip + skip_preperiods = True # No cohorts qualify — skip entirely - rel_time_0 = np.where( - ~df_0["_never_treated"], - df_0[time] - df_0[first_treat], - np.nan, - ) + if not skip_preperiods: + rel_time_0 = np.where( + ~df_0["_never_treated"], + df_0[time] - df_0[first_treat], + np.nan, + ) - # When balance_e is set, only include leads from balanced cohorts - if balanced_cohorts is not None: - is_balanced = df_0[first_treat].isin(balanced_cohorts).values - rel_time_for_leads = np.where(is_balanced, rel_time_0, np.nan) - else: - rel_time_for_leads = rel_time_0 + # When balance_e is set, only include leads from balanced cohorts + if balanced_cohorts is not None: + is_balanced = df_0[first_treat].isin(balanced_cohorts).values + rel_time_for_leads = np.where(is_balanced, rel_time_0, np.nan) + else: + rel_time_for_leads = rel_time_0 - pre_rel_times = sorted( - set(int(h) for h in rel_time_for_leads if np.isfinite(h) and h < -self.anticipation) - ) - pre_rel_times = [h for h in pre_rel_times if h != ref_period] - if self.horizon_max is not None: - pre_rel_times = [h for h in pre_rel_times if abs(h) <= self.horizon_max] - if pre_rel_times: - pre_effects, _, _ = self._compute_lead_coefficients( - df_0, outcome, unit, time, first_treat, covariates, - cluster_var, pre_rel_times, alpha=self.alpha, - balanced_cohorts=balanced_cohorts, + pre_rel_times = sorted( + set( + int(h) + for h in rel_time_for_leads + if np.isfinite(h) and h < -self.anticipation + ) ) - event_study_effects.update(pre_effects) + pre_rel_times = [h for h in pre_rel_times if h != ref_period] + if self.horizon_max is not None: + pre_rel_times = [ + h for h in pre_rel_times if abs(h) <= self.horizon_max + ] + if pre_rel_times: + pre_effects, _, _ = self._compute_lead_coefficients( + df_0, + outcome, + unit, + time, + first_treat, + covariates, + cluster_var, + pre_rel_times, + alpha=self.alpha, + balanced_cohorts=balanced_cohorts, + ) + event_study_effects.update(pre_effects) # Collect horizons with Proposition 5 violations prop5_horizons = [] @@ -1902,7 +1917,7 @@ def _compute_lead_coefficients( # All lead columns dropped (rank deficient after demeaning) effects: Dict[int, Dict[str, Any]] = {} for h in pre_rel_times: - n_obs = int((rel_time_0 == h).sum()) + n_obs = int(df_0[f"_lead_{h}"].sum()) effects[h] = { "effect": np.nan, "se": np.nan, "t_stat": np.nan, "p_value": np.nan, "conf_int": (np.nan, np.nan), @@ -1927,7 +1942,8 @@ def _compute_lead_coefficients( for j, h in enumerate(pre_rel_times): effect = float(gamma[j]) se = float(np.sqrt(max(V_gamma[j, j], 0.0))) - n_obs = int((rel_time_0 == h).sum()) + # n_obs from the lead indicator (respects balanced_cohorts restriction) + n_obs = int(df_0[f"_lead_{h}"].sum()) t_stat, p_value, conf_int = safe_inference(effect, se, alpha=alpha) effects[h] = { "effect": effect,