diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index 4f32a6638..b54bf42dd 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 ---------- @@ -134,6 +139,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 +166,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 @@ -229,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 for now." + ) + # Create working copy df = data.copy() @@ -1101,6 +1116,7 @@ def _compute_cluster_psi_sums( # ---- 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 treated_units = df_1[unit].values treated_times = df_1[time].values @@ -1116,6 +1132,9 @@ def _compute_cluster_psi_sums( w_total = float(np.sum(weights)) + 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) @@ -1127,8 +1146,6 @@ def _compute_cluster_psi_sums( 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): @@ -1513,6 +1530,69 @@ 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() + + # 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 + skip_preperiods = False + 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 not balanced_cohorts: + skip_preperiods = True # No cohorts qualify — skip entirely + + 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 + + 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, + ) + event_study_effects.update(pre_effects) + # Collect horizons with Proposition 5 violations prop5_horizons = [] @@ -1748,9 +1828,138 @@ 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, + 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). + + 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 + 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 — 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}" + 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 + 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 + 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(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), + "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 + + n_leads = len(lead_cols) + gamma = coefficients[:n_leads] + V_gamma = vcov[:n_leads, :n_leads] + + # Build per-horizon effects + effects = {} + 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 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, + "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). @@ -1782,7 +1991,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], @@ -1808,7 +2016,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) @@ -1821,49 +2028,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 + # 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, ) - 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 - - # 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: @@ -1874,6 +2045,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) @@ -1881,10 +2053,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, @@ -1910,6 +2079,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..e1623884d 100644 --- a/diff_diff/imputation_bootstrap.py +++ b/diff_diff/imputation_bootstrap.py @@ -174,18 +174,20 @@ def _precompute_bootstrap_psi( df_1 = df.loc[omega_1_mask] rel_times = df_1["_rel_time"].values + 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] + # 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 + for h in event_study_effects: if event_study_effects[h].get("n_obs", 0) == 0: continue @@ -193,6 +195,12 @@ def _precompute_bootstrap_psi( continue if not np.isfinite(event_study_effects[h].get("effect", np.nan)): continue + + # Skip pre-period horizons — their SEs come from Test 1 + # lead regression, not bootstrap + if h < -self.anticipation: + continue + 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..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 ---------- @@ -139,6 +144,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 +165,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 +1022,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 +1619,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..7d1270b7e 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -836,6 +836,22 @@ 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`, Test 1 / Equation 9):* + +Pre-period coefficients reuse the existing pre-trend test machinery (BJS Equation 9): +``` +Y_it = alpha_i + beta_t [+ X'_it * delta] + sum_h gamma_h * W_it(h) + epsilon_it +``` +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) +- 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. - **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 +941,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..0e059bf08 --- /dev/null +++ b/tests/test_pretrends_event_study.py @@ -0,0 +1,812 @@ +""" +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 TestPretrends_ContractTests: + """Verify ImputationDiD pre-period coefficients match pretrend_test() lead coefficients.""" + + def test_imputation_pretrends_match_pretrend_test(self): + """ImputationDiD pretrends=True effects match pretrend_test().lead_coefficients.""" + data = generate_test_data(seed=77) + + est = ImputationDiD(pretrends=True) + results = est.fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + # Get pretrend_test lead coefficients (called on results object) + pt = results.pretrend_test() + lead_coefs = pt["lead_coefficients"] + + # 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( + es_effect, + coef, + rtol=1e-10, + err_msg=f"h={h}: event study effect != pretrend_test lead coefficient", + ) + + +# ============================================================================= +# 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_bootstrap_post_only(self): + """ImputationDiD pretrends=True + bootstrap: bootstrap updates post SEs only. + + 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) + + # Without bootstrap + results_no_boot = ImputationDiD(pretrends=True).fit( + data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + + # 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). + + 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", + ) + + 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"