diff --git a/CHANGELOG.md b/CHANGELOG.md index a669e3c4..7e9169f3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- **`LPDiD` (Local Projections Difference-in-Differences; Dube, Girardi, Jordà & Taylor 2025, + *J. Applied Econometrics* 40(5):741-758).** Per-horizon long-difference OLS + (`y_{i,t+h} − y_{i,t−1}`) on a clean-control sample (newly-treated + not-yet-treated) with + calendar-time fixed effects and no unit FE, so the default variance-weighted estimand has + strictly non-negative weights (no TWFE negative-weighting). Options: `reweight=True` + (equally-weighted ATT, numerically equivalent to Callaway-Sant'Anna), premean-differenced base + period (`pmd`), `no_composition` (fixed post-treatment composition), pooled pre/post estimands, + outcome/first-difference lag controls (`ylags`/`dylags`), and a regression-adjustment covariate + path (ImputationDiD/BJS-family influence-function cluster variance). Cluster-robust SEs at the + unit level by default. This release implements the **absorbing-treatment** path; per-unit + interior time gaps are handled by calendar-correct feature construction. Non-absorbing + treatment, survey-design support, and external R-package parity are tracked follow-ups. - **`placebo_group_test` gained an optional `treatment` parameter.** When supplied, units that are ever real-treated are dropped before the placebo so it runs on never-treated units only (the uncontaminated design); without it, behavior is unchanged and the caller must pass diff --git a/README.md b/README.md index 3a3a26b1..026fc96d 100644 --- a/README.md +++ b/README.md @@ -117,6 +117,7 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`. - [TROP](https://diff-diff.readthedocs.io/en/stable/api/trop.html) - Triply Robust Panel estimator (Athey et al. 2025) with nuclear norm factor adjustment - [StaggeredTripleDifference](https://diff-diff.readthedocs.io/en/stable/api/staggered.html#staggeredtripledifference) - Ortiz-Villavicencio & Sant'Anna (2025) staggered DDD with group-time ATT - [WooldridgeDiD](https://diff-diff.readthedocs.io/en/stable/api/wooldridge_etwfe.html) - Wooldridge (2023, 2025) ETWFE: saturated OLS, logit/Poisson QMLE (ASF-based ATT). Alias `ETWFE`. +- [LPDiD](https://diff-diff.readthedocs.io/en/stable/api/lpdid.html) - Dube, Girardi, Jorda & Taylor (2025) Local Projections DiD: per-horizon long-difference event study on clean controls (no negative weighting), variance- or equally-weighted ATT, for absorbing treatment - [BaconDecomposition](https://diff-diff.readthedocs.io/en/stable/api/bacon.html) - Goodman-Bacon (2021) decomposition for diagnosing TWFE bias in staggered settings ## Diagnostics & Sensitivity diff --git a/diff_diff/__init__.py b/diff_diff/__init__.py index 886889df..3796754e 100644 --- a/diff_diff/__init__.py +++ b/diff_diff/__init__.py @@ -230,6 +230,8 @@ from diff_diff.synthetic_control_results import SyntheticControlResults from diff_diff.wooldridge import WooldridgeDiD from diff_diff.wooldridge_results import WooldridgeDiDResults +from diff_diff.lpdid import LPDiD +from diff_diff.lpdid_results import LPDiDResults from diff_diff.utils import ( WildBootstrapResults, check_parallel_trends, @@ -381,6 +383,9 @@ "WooldridgeDiD", "WooldridgeDiDResults", "ETWFE", + # LPDiD (Local Projections DiD) + "LPDiD", + "LPDiDResults", # Visualization "plot_bacon", "plot_event_study", diff --git a/diff_diff/guides/llms-full.txt b/diff_diff/guides/llms-full.txt index 17d23121..c63441b7 100644 --- a/diff_diff/guides/llms-full.txt +++ b/diff_diff/guides/llms-full.txt @@ -895,6 +895,56 @@ results = edid.fit(data, outcome='y', unit='id', time='t', results.print_summary() ``` +### LPDiD + +Local Projections DiD (Dube, Girardi, Jorda & Taylor 2025). Estimates a separate OLS at each event-time horizon of a long difference (`y_{i,t+h} - y_{i,t-1}`) on the treatment-switch indicator plus calendar-time fixed effects (no unit FE), restricted to a flexible "clean control" sample of newly-treated and not-yet-treated units. Excluding already-treated units from the control group removes the negative-weighting bias of naive TWFE, so the default (variance-weighted) estimand has strictly non-negative weights. `reweight=True` yields the equally-weighted ATT (numerically equivalent to Callaway-Sant'Anna); covariates then enter via regression adjustment. Standard errors on the default/weighted path are cluster-robust at the unit level (the paper specifies no SE; matches Stata `lpdid` `vce(cluster unit)`); the regression-adjustment covariate path (`reweight=True`) instead reports an influence-function cluster variance (ImputationDiD/BJS family). Scope: binary, absorbing treatment (rejects panels where treatment turns off). + +```python +LPDiD( + pre_window: int = 2, # Number of pre-treatment horizons (placebos) + post_window: int = 0, # Number of post-treatment horizons + control_group: str = "clean", # "clean" (not-yet-treated) or "never_treated" + reweight: bool = False, # True -> equally-weighted ATT (== Callaway-Sant'Anna); False -> variance-weighted + no_composition: bool = False, # Hold the post-treatment composition fixed across post horizons + pmd: str | int | None = None, # Base period: None=first-lag (t-1), "max"=premean over all pretreatment periods, int=last-k premean + alpha: float = 0.05, + cluster: str | None = None, # Cluster column for cluster-robust SEs; defaults to the unit identifier + rank_deficient_action: str = "warn", # "warn", "error", or "silent" +) +``` + +**fit() parameters:** + +```python +lpdid.fit( + data: pd.DataFrame, + outcome: str, + unit: str, + time: str, + treatment: str, # Binary, absorbing treatment indicator (0/1) + covariates: list[str] = None, # Direct inclusion (reweight=False) or regression adjustment (reweight=True) + ylags: int = 0, # Lagged-outcome controls + dylags: int = 0, # Lagged first-difference controls + absorb: list[str] = None, # Additional absorbed fixed-effect columns + post_pooled: int | tuple = None, # Pooled post-window horizons (int or (start, end)) + pre_pooled: int | tuple = None, # Pooled pre-window horizons (int or (start, end)) + only_event: bool = False, # Compute only the event-study table + only_pooled: bool = False, # Compute only the pooled pre/post table +) -> LPDiDResults +``` + +**Usage:** + +```python +from diff_diff import LPDiD + +lp = LPDiD(pre_window=5, post_window=10) +results = lp.fit(data, outcome='y', unit='id', time='t', treatment='treated') +results.print_summary() +print(results.event_study) # per-horizon coefficients +print(results.pooled) # pooled pre (placebo) / post (ATT) rows +``` + ### TROP Triply Robust Panel estimator (Athey, Imbens, Qu & Viviano 2025). Combines nuclear norm regularization, distance-based unit weights, and time decay weights. diff --git a/diff_diff/guides/llms.txt b/diff_diff/guides/llms.txt index 54ea1593..92abb8f2 100644 --- a/diff_diff/guides/llms.txt +++ b/diff_diff/guides/llms.txt @@ -2,7 +2,7 @@ > A Python library for Difference-in-Differences (DiD) causal inference analysis. Provides sklearn-like estimators with statsmodels-style summary output for econometric analysis. -diff-diff offers 18 estimators covering basic 2x2 DiD, modern staggered adoption methods, reversible (non-absorbing) treatments, advanced panel estimators, nonlinear models, and diagnostic tools. It supports robust and cluster-robust standard errors, wild cluster bootstrap, formula and column-name interfaces, fixed effects (dummy and absorbed), complex survey designs (strata/PSU/FPC, replicate weights, design-based variance), and publication-ready output. The optional Rust backend accelerates compute-intensive estimators like Synthetic DiD and TROP. +diff-diff offers 19 estimators covering basic 2x2 DiD, modern staggered adoption methods, reversible (non-absorbing) treatments, advanced panel estimators, nonlinear models, and diagnostic tools. It supports robust and cluster-robust standard errors, wild cluster bootstrap, formula and column-name interfaces, fixed effects (dummy and absorbed), complex survey designs (strata/PSU/FPC, replicate weights, design-based variance), and publication-ready output. The optional Rust backend accelerates compute-intensive estimators like Synthetic DiD and TROP. - Install: `pip install diff-diff` - License: MIT @@ -69,6 +69,7 @@ Full practitioner guide: call `diff_diff.get_llm_guide("practitioner")` - [TROP](https://diff-diff.readthedocs.io/en/stable/api/trop.html): Triply Robust Panel estimator (Athey et al. 2025) with nuclear norm factor adjustment - [StaggeredTripleDifference](https://diff-diff.readthedocs.io/en/stable/api/staggered.html#staggeredtripledifference): Ortiz-Villavicencio & Sant'Anna (2025) staggered DDD with group-time ATT - [WooldridgeDiD](https://diff-diff.readthedocs.io/en/stable/api/wooldridge_etwfe.html): Wooldridge (2023, 2025) ETWFE — saturated OLS, logit/Poisson QMLE (ASF-based ATT). Alias: ETWFE +- [LPDiD](https://diff-diff.readthedocs.io/en/stable/api/lpdid.html): Dube, Girardi, Jorda & Taylor (2025) Local Projections DiD: per-horizon long-difference event study on clean controls (no negative weighting); variance- or equally-weighted ATT, premean differencing, pooled pre/post, fast. Absorbing treatment. - [BaconDecomposition](https://diff-diff.readthedocs.io/en/stable/api/bacon.html): Goodman-Bacon (2021) decomposition for diagnosing TWFE bias in staggered settings ## Diagnostics and Sensitivity Analysis diff --git a/diff_diff/lpdid.py b/diff_diff/lpdid.py new file mode 100644 index 00000000..66db69b6 --- /dev/null +++ b/diff_diff/lpdid.py @@ -0,0 +1,1205 @@ +import warnings +from typing import Any, Dict, Iterable, Optional, Union + +import numpy as np +import pandas as pd + +from diff_diff.linalg import _rank_guarded_inv, solve_ols +from diff_diff.lpdid_results import LPDiDResults +from diff_diff.utils import safe_inference + +__all__ = ["LPDiD", "LPDiDResults"] + + +class LPDiD: + def __init__( + self, + pre_window: int = 2, + post_window: int = 0, + control_group: str = "clean", + reweight: bool = False, + no_composition: bool = False, + pmd: Optional[Union[str, int]] = None, + alpha: float = 0.05, + cluster: Optional[str] = None, + rank_deficient_action: str = "warn", + ): + self.pre_window = pre_window + self.post_window = post_window + self.control_group = control_group + self.reweight = reweight + self.no_composition = no_composition + self.pmd = pmd + self.alpha = alpha + self.cluster = cluster + self.rank_deficient_action = rank_deficient_action + self._validate_params() + self.is_fitted_ = False + self.results_: Optional[LPDiDResults] = None + + def _validate_params(self) -> None: + for _name, _val in (("pre_window", self.pre_window), ("post_window", self.post_window)): + if not isinstance(_val, int) or isinstance(_val, bool) or _val < 0: + raise ValueError(f"{_name} must be a non-negative integer") + if ( + isinstance(self.alpha, bool) + or not isinstance(self.alpha, (int, float)) + or not (0.0 < float(self.alpha) < 1.0) + ): + raise ValueError("alpha must be a float in (0, 1)") + if self.control_group not in ("clean", "never_treated"): + raise ValueError("control_group must be 'clean' or 'never_treated'") + if self.rank_deficient_action not in ("warn", "error", "silent"): + raise ValueError("rank_deficient_action must be 'warn', 'error', or 'silent'") + if self.pmd is not None and not ( + self.pmd == "max" + or (isinstance(self.pmd, int) and not isinstance(self.pmd, bool) and self.pmd > 0) + ): + raise ValueError("pmd must be None, 'max', or a positive integer") + + def _rhs_column_names(self, covariates=None, ylags=0, dylags=0): + rhs_columns = list(covariates or []) + rhs_columns.extend([f"_y_lag_{lag}" for lag in range(1, ylags + 1)]) + rhs_columns.extend([f"_dy_lag_{lag}" for lag in range(1, dylags + 1)]) + return rhs_columns + + def _prepare_panel( + self, + data, + outcome, + unit, + time, + treatment, + cluster, + covariates=None, + ylags=0, + dylags=0, + absorb=None, + ): + selected_columns = list( + dict.fromkeys( + [unit, time, outcome, treatment, cluster, *(covariates or []), *(absorb or [])] + ) + ) + panel = data[selected_columns].copy() + panel = panel.sort_values([unit, time]).reset_index(drop=True) + + if panel.duplicated([unit, time]).any(): + raise ValueError("LPDiD requires unique unit-time observations") + + treated_numeric = pd.to_numeric(panel[treatment], errors="coerce") + if treated_numeric.isna().any() or not treated_numeric.isin([0, 1]).all(): + raise ValueError( + "treatment must contain binary numeric 0/1 values with no missing data" + ) + + panel["_treated"] = treated_numeric.astype(int) + panel["_cluster"] = panel[cluster] + if panel["_cluster"].isna().any(): + raise ValueError( + f"cluster column '{cluster}' contains missing values; LPDiD cannot form " + "cluster-robust standard errors with missing cluster labels (the affected " + "rows would silently drop from the variance)." + ) + + # Absorbing-path validation and entry detection run on the OBSERVED rows, + # BEFORE any calendar reindex below (the absorbing fill would otherwise make + # the monotonicity check trivially pass, and gap rows carry NaN clusters). + treated_cummax = panel.groupby(unit)["_treated"].cummax() + if (treated_cummax > panel["_treated"]).any(): + raise ValueError( + "LPDiD currently requires an absorbing treatment path " + "(once treated, always treated)" + ) + # Entry = first OBSERVED treated period (documented convention; an unobserved + # pre-onset gap is unknowable). For an absorbing path this is min(t | D=1). + first_treat = panel.loc[panel["_treated"].eq(1)].groupby(unit)[time].min() + + # LP-DiD's per-unit features (outcome lags, first differences, premean + # baselines) are CALENDAR quantities (t-1, t-k, t+h). Computing them with + # row-order ops (shift/diff/rolling) silently equates "previous observed row" + # with "calendar t-1", which is wrong when a unit has an interior time gap. + # Reindex each unit to its complete interior calendar grid so every row-order + # op is calendar-correct, compute the features on the grid, then restrict back + # to the observed rows so the synthetic NaN gap rows never enter a regression. + # A gap-free panel skips this entirely and is bit-identical to before. + span = panel.groupby(unit)[time].agg(["min", "max", "nunique"]) + has_gap = bool((span["nunique"] != (span["max"] - span["min"] + 1)).any()) + if has_gap: + panel["_observed"] = True + grid = pd.concat( + [ + pd.DataFrame({unit: u, time: np.arange(int(lo), int(hi) + 1)}) + for u, lo, hi in zip(span.index, span["min"], span["max"]) + ], + ignore_index=True, + ) + panel = grid.merge(panel, on=[unit, time], how="left") + panel["_observed"] = panel["_observed"].fillna(False).astype(bool) + panel = panel.sort_values([unit, time]).reset_index(drop=True) + panel["_first_treat"] = panel[unit].map(first_treat).astype(float).fillna(np.inf) + # Absorbing fill: treatment is fully determined by the entry period, so the + # filled gap rows are consistent and observed rows are reproduced exactly. + panel["_treated"] = (panel[time] >= panel["_first_treat"]).astype(int) + panel["_entry"] = (panel[time] == panel["_first_treat"]).astype(float) + else: + panel["_first_treat"] = panel[unit].map(first_treat).astype(float).fillna(np.inf) + panel["_entry"] = (panel[time] == panel["_first_treat"]).astype(float) + + # Premean ("max") baseline = mean of all AVAILABLE strictly-prior outcomes. + # It must NOT depend on the base row's own outcome y_t: PMD replaces the + # t-1 baseline with the premean of prior periods, and the long difference is + # y_{t+h} - premean, so a base row with a missing current outcome (but + # observed priors and target) stays identified. fillna(0) before the + # cumulative sum makes the numerator the strictly-prior non-missing sum even + # when y_t (or an interior period) is missing; the denominator counts + # strictly-prior non-missing outcomes (not rows). Bit-identical to a plain + # cumsum when no outcome is NaN. + _filled_outcome = panel[outcome].fillna(0.0) + outcome_history_sum = _filled_outcome.groupby(panel[unit]).cumsum() - _filled_outcome + prior_nonnull = panel[outcome].notna().astype(int) + history_count = prior_nonnull.groupby(panel[unit]).cumsum() - prior_nonnull + panel["_pmd_all_baseline"] = outcome_history_sum / history_count.replace(0, np.nan) + + lagged_outcome = panel.groupby(unit)[outcome].shift(1) + if isinstance(self.pmd, int): + panel["_pmd_k_baseline"] = ( + lagged_outcome.groupby(panel[unit]) + .rolling(window=self.pmd, min_periods=self.pmd) + .mean() + .reset_index(level=0, drop=True) + ) + + for lag in range(1, ylags + 1): + panel[f"_y_lag_{lag}"] = panel.groupby(unit)[outcome].shift(lag) + + panel["_dy_current"] = panel.groupby(unit)[outcome].diff() + for lag in range(1, dylags + 1): + panel[f"_dy_lag_{lag}"] = panel.groupby(unit)["_dy_current"].shift(lag) + + if has_gap: + # Drop the synthetic gap rows. The features above are now calendar-correct + # on the observed rows (a lag/difference spanning a gap is NaN, so the + # observation fails closed via the downstream dropna), and no NaN-outcome / + # NaN-cluster phantom row reaches estimation or the reweight denominators. + panel = ( + panel.loc[panel["_observed"]] + .drop(columns="_observed") + .sort_values([unit, time]) + .reset_index(drop=True) + ) + return panel + + def _baseline_column(self): + if self.pmd == "max": + return "_pmd_all_baseline" + if isinstance(self.pmd, int): + return "_pmd_k_baseline" + return "_baseline_outcome" + + def _clean_control_mask(self, panel: pd.DataFrame, *, time: str, horizon: int) -> pd.Series: + if self.control_group == "never_treated": + return panel["_treated"].eq(0) & np.isinf(panel["_first_treat"]) + + control_mask = panel["_treated"].eq(0) & panel[time].lt(panel["_first_treat"]) + if horizon >= 0: + control_mask &= (panel[time] + horizon).lt(panel["_first_treat"]) + return control_mask + + def _common_clean_sample_indicator( + self, panel: pd.DataFrame, *, unit: str, time: str, outcome: str, max_post_horizon: int + ) -> pd.Series: + common_sample = panel["_entry"].eq(1.0) | self._clean_control_mask( + panel, + time=time, + horizon=max_post_horizon, + ) + # Fixed composition requires the active baseline AND every post-treatment + # target outcome (h = 0..max_post_horizon) to be NON-MISSING for each base + # observation -- not merely that the target row exists. This keeps the + # realized post sample fixed across all post horizons under any missingness + # encoding (absent rows OR present-but-NaN outcomes). + outcome_by_key = panel.set_index([unit, time])[outcome] + + def _value_available(horizon: int) -> pd.Series: + keys = pd.MultiIndex.from_arrays( + [panel[unit].to_numpy(), (panel[time] + horizon).to_numpy()] + ) + return pd.Series(outcome_by_key.reindex(keys).notna().to_numpy(), index=panel.index) + + if self.pmd is None: + available = _value_available(-1) # t-1 baseline outcome + else: + available = panel[self._baseline_column()].notna() + for h in range(0, max_post_horizon + 1): + available &= _value_available(h) + return common_sample & available + + def _rw_weights_from_sample(self, sample: pd.DataFrame) -> pd.Series: + """Equal-weighting weights from the REALIZED estimation sample. + + Computed after all row drops and clean-control restrictions, so the + per-event-time denominator matches the regression's actual risk set. + Computing from the pre-drop panel would silently change the estimand + (and break the Callaway-Sant'Anna equivalence) on unbalanced panels. + For each event time, weight = N_clean_control_sample / N_control. + """ + if sample.empty: + return pd.Series(dtype=float) + group_stats = sample.groupby("_event_time")["_entry"].agg(["sum", "count"]) + treated_counts = group_stats["sum"] + control_counts = group_stats["count"] - treated_counts + + valid = (treated_counts > 0) & (control_counts > 0) + if not valid.any(): + return pd.Series(dtype=float) + + return (group_stats.loc[valid, "count"] / control_counts.loc[valid]).astype(float) + + def _build_horizon_sample( + self, + panel, + *, + outcome, + unit, + time, + horizon, + covariates=None, + ylags=0, + dylags=0, + absorb=None, + apply_no_composition: bool = True, + ): + rhs_columns = self._rhs_column_names(covariates=covariates, ylags=ylags, dylags=dylags) + base_columns = [ + unit, + time, + "_treated", + "_entry", + "_first_treat", + "_cluster", + "_common_event_ok", + *rhs_columns, + *(absorb or []), + ] + if self.pmd == "max": + base_columns.append("_pmd_all_baseline") + elif isinstance(self.pmd, int): + base_columns.append("_pmd_k_baseline") + + base = panel[base_columns].copy() + base["_baseline_time"] = base[time] - 1 + base["_target_time"] = base[time] + horizon + + outcomes = panel[[unit, time, outcome]].copy() + + baseline = outcomes.rename(columns={time: "_baseline_time", outcome: "_baseline_outcome"}) + target = outcomes.rename(columns={time: "_target_time", outcome: "_target_outcome"}) + + sample = base.merge(baseline, on=[unit, "_baseline_time"], how="left") + sample = sample.merge(target, on=[unit, "_target_time"], how="left") + baseline_column = self._baseline_column() + # Require the ACTIVE baseline column only: under PMD the long difference + # uses the premean baseline, so a missing exact t-1 outcome must not drop + # an otherwise-identified observation (matters on unbalanced panels). + required_columns = [baseline_column, "_target_outcome", *rhs_columns, *(absorb or [])] + sample = sample.dropna(subset=required_columns).copy() + + treated_mask = sample["_entry"].eq(1.0) + if self.control_group == "never_treated": + control_mask = sample["_entry"].eq(0.0) & np.isinf(sample["_first_treat"]) + else: + control_mask = self._clean_control_mask(sample, time=time, horizon=horizon) + + sample = sample.loc[treated_mask | control_mask].copy() + # Fixed composition is a POST-treatment contract: apply it only to post + # horizons; pre-treatment placebos use whatever pre-period data exists. + if self.no_composition and apply_no_composition and horizon >= 0: + sample = sample.loc[sample["_common_event_ok"]].copy() + sample["horizon"] = horizon + sample["_event_time"] = sample[time] + sample["_long_diff"] = sample["_target_outcome"] - sample[baseline_column] + return sample[ + [ + "horizon", + "_event_time", + "_long_diff", + "_entry", + "_cluster", + *rhs_columns, + *(absorb or []), + ] + ] + + def _sample_is_identified(self, sample: pd.DataFrame) -> bool: + return len(sample) > 0 and sample["_entry"].nunique() >= 2 + + def _build_feature_frame( + self, + sample: pd.DataFrame, + *, + rhs_columns=None, + absorb_columns=None, + include_time_fe: bool = True, + time_levels=None, + absorb_levels=None, + ) -> pd.DataFrame: + rhs_columns = list(rhs_columns or []) + absorb_columns = list(absorb_columns or []) + feature_blocks = [] + + for col in rhs_columns: + values = pd.to_numeric(sample[col], errors="coerce") + if values.isna().any(): + raise ValueError( + f"LPDiD requires numeric covariate-style columns, got invalid values in '{col}'" + ) + feature_blocks.append( + pd.DataFrame({col: values.to_numpy(dtype=float)}, index=sample.index) + ) + + if include_time_fe: + time_categories = list( + time_levels if time_levels is not None else pd.unique(sample["_event_time"]) + ) + time_dummies = pd.get_dummies( + pd.Categorical(sample["_event_time"], categories=time_categories), + prefix="time", + drop_first=True, + dtype=float, + ) + if not time_dummies.empty: + time_dummies.index = sample.index + feature_blocks.append(time_dummies) + + for col in absorb_columns: + categories = list( + absorb_levels[col] if absorb_levels is not None else pd.unique(sample[col]) + ) + dummies = pd.get_dummies( + pd.Categorical(sample[col], categories=categories), + prefix=col, + drop_first=True, + dtype=float, + ) + if not dummies.empty: + dummies.index = sample.index + feature_blocks.append(dummies) + + if not feature_blocks: + return pd.DataFrame(index=sample.index) + return pd.concat(feature_blocks, axis=1) + + def _estimate_regression_adjustment_sample( + self, + sample: pd.DataFrame, + *, + response_column: str = "_long_diff", + include_time_fe: bool = True, + rhs_columns=None, + absorb_columns=None, + ) -> Dict[str, Optional[float]]: + rhs_columns = list(rhs_columns or []) + absorb_columns = list(absorb_columns or []) + dropna_columns = [*rhs_columns, *absorb_columns] + if dropna_columns: + sample = sample.dropna(subset=dropna_columns).copy() + + n_obs = int(len(sample)) + empty_result = { + "coefficient": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_low": np.nan, + "conf_high": np.nan, + "n_obs": n_obs, + "n_clusters": np.nan, + } + if n_obs == 0 or sample["_entry"].nunique() < 2: + return empty_result + + controls = sample.loc[sample["_entry"].eq(0.0)].copy() + treated = sample.loc[sample["_entry"].eq(1.0)].copy() + if controls.empty or treated.empty: + return empty_result + + # The RA counterfactual for a treated observation is the predicted long + # difference from the clean-control regression at that event time. An + # event time with treated units but no clean control has an unidentified + # time fixed effect, so those treated observations cannot be imputed: + # drop them (and surface the drop) rather than extrapolate off a + # rank-deficient fit. + if include_time_fe: + control_event_times = set(controls["_event_time"].unique()) + identified = treated["_event_time"].isin(control_event_times) + if not bool(identified.all()): + n_drop = int((~identified).sum()) + warnings.warn( + f"LPDiD regression adjustment: dropped {n_drop} treated observation(s) " + "at event time(s) with no clean control (counterfactual unidentified).", + UserWarning, + stacklevel=2, + ) + treated = treated.loc[identified].copy() + if treated.empty: + return empty_result + sample = pd.concat([controls, treated]) + n_obs = int(len(sample)) + + # Absorbed-factor overlap: a treated observation whose absorbed level is + # absent from the clean controls has an all-zero control dummy for that + # level, so its counterfactual would be extrapolated through an unidentified + # coefficient. Drop those treated observations (and surface the drop) rather + # than impute off a non-identified fit. Mirrors the event-time check above. + if absorb_columns: + unsupported = pd.Series(False, index=treated.index) + for col in absorb_columns: + control_levels = set(controls[col].unique()) + unsupported = unsupported | ~treated[col].isin(control_levels) + if bool(unsupported.any()): + n_drop = int(unsupported.sum()) + warnings.warn( + f"LPDiD regression adjustment: dropped {n_drop} treated observation(s) " + "with an absorbed-factor level absent from the clean controls " + "(counterfactual unidentified -- no overlap).", + UserWarning, + stacklevel=2, + ) + treated = treated.loc[~unsupported].copy() + if treated.empty: + return empty_result + sample = pd.concat([controls, treated]) + n_obs = int(len(sample)) + + time_levels = list(pd.unique(sample["_event_time"])) if include_time_fe else None + absorb_levels = {col: list(pd.unique(sample[col])) for col in absorb_columns} + control_features = self._build_feature_frame( + controls, + rhs_columns=rhs_columns, + absorb_columns=absorb_columns, + include_time_fe=include_time_fe, + time_levels=time_levels, + absorb_levels=absorb_levels, + ) + treated_features = self._build_feature_frame( + treated, + rhs_columns=rhs_columns, + absorb_columns=absorb_columns, + include_time_fe=include_time_fe, + time_levels=time_levels, + absorb_levels=absorb_levels, + ) + + control_design = np.column_stack( + [np.ones(len(controls), dtype=float), control_features.to_numpy(dtype=float)] + ) + column_names = ["intercept", *control_features.columns.tolist()] + control_coef, _, _ = solve_ols( + control_design, + controls[response_column].to_numpy(dtype=float), + return_vcov=False, + rank_deficient_action=self.rank_deficient_action, + column_names=column_names, + ) + # solve_ols sets DROPPED redundant-nuisance coefficients to NaN under + # rank_deficient_action="warn"/"silent" (a constant/duplicate covariate, a + # collinear absorbed level, or lag collinearity). The dropped column's + # contribution is absorbed by the retained collinear column(s), so it acts + # as 0 for prediction/residuals; without this zero-fill the NaN would + # propagate through every prediction and NaN an otherwise-identified ATT. + # ("error" still raises inside solve_ols before returning.) + control_coef = np.where(np.isfinite(control_coef), control_coef, 0.0) + + treated_design = np.column_stack( + [np.ones(len(treated), dtype=float), treated_features.to_numpy(dtype=float)] + ) + untreated_prediction = treated_design @ control_coef + treated_residual = treated[response_column].to_numpy(dtype=float) - untreated_prediction + + effect = float(treated_residual.mean()) + se = np.nan + cluster_ids = sample["_cluster"].to_numpy() + n_clusters = len(pd.unique(cluster_ids)) + if n_clusters >= 2 and np.isfinite(effect): + n_total = len(sample) + n_treated = len(treated) + q0_inv, _, _ = _rank_guarded_inv(control_design.T @ control_design) + mu_treated = treated_design.mean(axis=0) + control_projection = control_design @ (q0_inv @ mu_treated) + + phi = np.zeros(n_total, dtype=float) + treated_mask = sample["_entry"].to_numpy(dtype=float) == 1.0 + phi[treated_mask] = (n_total / n_treated) * (treated_residual - effect) + control_residual = ( + controls[response_column].to_numpy(dtype=float) - control_design @ control_coef + ) + phi[~treated_mask] = -n_total * control_projection * control_residual + + cluster_scores = pd.Series(phi).groupby(cluster_ids).sum().to_numpy(dtype=float) + vcov_scalar = float(cluster_scores @ cluster_scores) / float(n_total**2) + if np.isfinite(vcov_scalar) and vcov_scalar >= 0: + se = float(np.sqrt(vcov_scalar)) + + df = n_clusters - 1 if n_clusters > 1 else None + t_stat, p_value, conf_int = safe_inference(effect, se, alpha=self.alpha, df=df) + return { + "coefficient": effect, + "se": se, + "t_stat": t_stat, + "p_value": p_value, + "conf_low": conf_int[0], + "conf_high": conf_int[1], + "n_obs": n_obs, + "n_clusters": n_clusters, + } + + def _estimate_sample( + self, + sample: pd.DataFrame, + *, + response_column: str = "_long_diff", + include_time_fe: bool = True, + rhs_columns=None, + absorb_columns=None, + weight_column: Optional[str] = None, + ) -> Dict[str, Optional[float]]: + rhs_columns = list(rhs_columns or []) + absorb_columns = list(absorb_columns or []) + dropna_columns = [*rhs_columns, *absorb_columns] + if weight_column is not None: + dropna_columns.append(weight_column) + if dropna_columns: + sample = sample.dropna(subset=dropna_columns).copy() + + n_obs = int(len(sample)) + empty_result = { + "coefficient": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_low": np.nan, + "conf_high": np.nan, + "n_obs": n_obs, + "n_clusters": np.nan, + } + + if n_obs == 0 or sample["_entry"].nunique() < 2: + return empty_result + + if include_time_fe: + # Clean-control support: a treated observation at an event time with no + # clean control has a time fixed effect collinear with the treatment + # indicator. The rank handler could then drop that time dummy and + # identify the treatment effect off invalid cross-event-time + # comparisons, so drop those unsupported treated observations (and + # surface the drop) rather than emit a spurious estimate. Mirrors the + # regression-adjustment path's event-time identification check. + control_event_times = set(sample.loc[sample["_entry"].eq(0.0), "_event_time"].unique()) + unsupported = sample["_entry"].eq(1.0) & ~sample["_event_time"].isin( + control_event_times + ) + if bool(unsupported.any()): + n_drop = int(unsupported.sum()) + warnings.warn( + f"LPDiD: dropped {n_drop} treated observation(s) at event time(s) with no " + "clean control (the treatment effect is unidentified at that event time).", + UserWarning, + stacklevel=2, + ) + sample = sample.loc[~unsupported].copy() + n_obs = int(len(sample)) + if n_obs == 0 or sample["_entry"].nunique() < 2: + return {**empty_result, "n_obs": n_obs} + + design_columns = [ + np.ones(n_obs, dtype=float), + sample["_entry"].to_numpy(dtype=float), + ] + column_names = ["intercept", "treatment_entry"] + + for col in rhs_columns: + values = pd.to_numeric(sample[col], errors="coerce") + if values.isna().any(): + raise ValueError( + f"LPDiD requires numeric covariate-style columns, got invalid values in '{col}'" + ) + design_columns.append(values.to_numpy(dtype=float)) + column_names.append(col) + + if include_time_fe: + time_dummies = pd.get_dummies( + sample["_event_time"], + prefix="time", + drop_first=True, + dtype=float, + ) + if not time_dummies.empty: + design_columns.append(time_dummies.to_numpy(dtype=float)) + column_names.extend(time_dummies.columns.tolist()) + + for col in absorb_columns: + dummies = pd.get_dummies(sample[col], prefix=col, drop_first=True, dtype=float) + if not dummies.empty: + design_columns.append(dummies.to_numpy(dtype=float)) + column_names.extend(dummies.columns.tolist()) + + design = np.column_stack(design_columns) + response = sample[response_column].to_numpy(dtype=float) + cluster_ids = sample["_cluster"].to_numpy() + weights = None if weight_column is None else sample[weight_column].to_numpy(dtype=float) + if n_obs <= design.shape[1]: + coef, _, _ = solve_ols( + design, + response, + return_vcov=False, + rank_deficient_action=self.rank_deficient_action, + column_names=column_names, + weights=weights, + ) + return { + "coefficient": float(coef[1]), + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_low": np.nan, + "conf_high": np.nan, + "n_obs": n_obs, + "n_clusters": len(pd.unique(cluster_ids)), + } + + use_cluster_vcov = len(pd.unique(cluster_ids)) >= 2 + vcov = None + if use_cluster_vcov: + try: + coef, _, vcov = solve_ols( + design, + response, + cluster_ids=cluster_ids, + return_vcov=True, + rank_deficient_action=self.rank_deficient_action, + column_names=column_names, + weights=weights, + ) + except (ValueError, ZeroDivisionError): + coef, _, _ = solve_ols( + design, + response, + return_vcov=False, + rank_deficient_action=self.rank_deficient_action, + column_names=column_names, + weights=weights, + ) + else: + coef, _, _ = solve_ols( + design, + response, + return_vcov=False, + rank_deficient_action=self.rank_deficient_action, + column_names=column_names, + weights=weights, + ) + + effect = float(coef[1]) + se = np.nan + if vcov is not None and vcov.shape[0] > 1 and np.isfinite(vcov[1, 1]) and vcov[1, 1] >= 0: + se = float(np.sqrt(vcov[1, 1])) + + n_clusters = len(pd.unique(cluster_ids)) + df = n_clusters - 1 if vcov is not None and n_clusters > 1 else None + t_stat, p_value, conf_int = safe_inference(effect, se, alpha=self.alpha, df=df) + return { + "coefficient": effect, + "se": se, + "t_stat": t_stat, + "p_value": p_value, + "conf_low": conf_int[0], + "conf_high": conf_int[1], + "n_obs": n_obs, + "n_clusters": n_clusters, + } + + def _estimate_horizon( + self, + panel, + *, + outcome, + unit, + time, + horizon, + covariates=None, + ylags=0, + dylags=0, + absorb=None, + ): + rhs_columns = self._rhs_column_names(covariates=covariates, ylags=ylags, dylags=dylags) + sample = self._build_horizon_sample( + panel, + outcome=outcome, + unit=unit, + time=time, + horizon=horizon, + covariates=covariates, + ylags=ylags, + dylags=dylags, + absorb=absorb, + ) + ra_required = self.reweight and bool(rhs_columns or absorb) + if ra_required: + return self._estimate_regression_adjustment_sample( + sample, + rhs_columns=rhs_columns, + absorb_columns=absorb, + ) + if self.reweight: + weight_map = self._rw_weights_from_sample(sample) + sample["_rw_event_weight"] = sample["_event_time"].map(weight_map) + return self._estimate_sample( + sample, + rhs_columns=rhs_columns, + absorb_columns=absorb, + weight_column="_rw_event_weight" if self.reweight else None, + ) + + def _build_pooled_sample( + self, + panel, + *, + outcome, + unit, + time, + horizons, + kind: str, + covariates=None, + ylags=0, + dylags=0, + absorb=None, + ): + rhs_columns = self._rhs_column_names(covariates=covariates, ylags=ylags, dylags=dylags) + base_columns = [ + unit, + time, + "_treated", + "_entry", + "_first_treat", + "_cluster", + "_common_pooled_ok", + *rhs_columns, + *(absorb or []), + ] + if self.pmd == "max": + base_columns.append("_pmd_all_baseline") + elif isinstance(self.pmd, int): + base_columns.append("_pmd_k_baseline") + + base = panel[base_columns].copy() + base["_baseline_time"] = base[time] - 1 + + outcomes = panel[[unit, time, outcome]].copy() + baseline = outcomes.rename(columns={time: "_baseline_time", outcome: "_baseline_outcome"}) + sample = base.merge(baseline, on=[unit, "_baseline_time"], how="left") + + target_columns = [] + for idx, horizon in enumerate(horizons): + target_time_col = f"_target_time_{idx}" + target_outcome_col = f"_target_outcome_{idx}" + sample[target_time_col] = sample[time] + horizon + target = outcomes.rename(columns={time: target_time_col, outcome: target_outcome_col}) + sample = sample.merge(target, on=[unit, target_time_col], how="left") + target_columns.append(target_outcome_col) + + baseline_column = self._baseline_column() + # Require the ACTIVE baseline column only (see _build_horizon_sample). + required_columns = [baseline_column, *target_columns, *rhs_columns, *(absorb or [])] + sample = sample.dropna(subset=required_columns).copy() + + treated_mask = sample["_entry"].eq(1.0) + if self.control_group == "never_treated": + control_mask = sample["_entry"].eq(0.0) & np.isinf(sample["_first_treat"]) + else: + control_horizon = max(horizons) if kind == "post" else 0 + control_mask = self._clean_control_mask(sample, time=time, horizon=control_horizon) + + sample = sample.loc[treated_mask | control_mask].copy() + if self.no_composition and kind == "post": + sample = sample.loc[sample["_common_pooled_ok"]].copy() + sample["_event_time"] = sample[time] + sample["_pooled_diff"] = sample[target_columns].mean(axis=1) - sample[baseline_column] + return sample[ + ["_event_time", "_pooled_diff", "_entry", "_cluster", *rhs_columns, *(absorb or [])] + ] + + def _estimate_window( + self, + panel, + *, + outcome, + unit, + time, + horizons: Iterable[int], + kind: str, + covariates=None, + ylags=0, + dylags=0, + absorb=None, + ): + horizons = list(horizons) + if not horizons: + raise ValueError( + f"pooled {kind} window is empty (no horizons to pool); a pooled pre " + "window requires pre_window >= 2. Use only_event=True or widen the window." + ) + unidentified_horizons = [ + horizon + for horizon in horizons + if not self._sample_is_identified( + self._build_horizon_sample( + panel, + outcome=outcome, + unit=unit, + time=time, + horizon=horizon, + covariates=covariates, + ylags=ylags, + dylags=dylags, + absorb=absorb, + # Base identification only; pooled estimation applies its own + # (pooled-window) composition mask via _common_pooled_ok. + apply_no_composition=False, + ) + ) + ] + if unidentified_horizons: + raise ValueError(f"unidentified pooled {kind} horizons: {unidentified_horizons}") + + rhs_columns = self._rhs_column_names(covariates=covariates, ylags=ylags, dylags=dylags) + pooled_sample = self._build_pooled_sample( + panel, + outcome=outcome, + unit=unit, + time=time, + horizons=horizons, + kind=kind, + covariates=covariates, + ylags=ylags, + dylags=dylags, + absorb=absorb, + ) + if pooled_sample.empty: + raise ValueError(f"pooled {kind} window did not contain any horizons") + ra_required = self.reweight and bool(rhs_columns or absorb) + if ra_required: + return self._estimate_regression_adjustment_sample( + pooled_sample, + response_column="_pooled_diff", + include_time_fe=True, + rhs_columns=rhs_columns, + absorb_columns=absorb, + ) + if self.reweight: + weight_map = self._rw_weights_from_sample(pooled_sample) + pooled_sample["_rw_pooled_weight"] = pooled_sample["_event_time"].map(weight_map) + return self._estimate_sample( + pooled_sample, + response_column="_pooled_diff", + include_time_fe=True, + rhs_columns=rhs_columns, + absorb_columns=absorb, + weight_column="_rw_pooled_weight" if self.reweight else None, + ) + + def _resolve_pooled_horizons(self, pooled, *, kind): + if isinstance(pooled, bool): + raise ValueError(f"{kind}_pooled must be None, an int, or a length-2 tuple (not bool)") + if kind == "pre": + default = list(range(-self.pre_window, -1)) + if isinstance(pooled, int): + horizons = list(range(-pooled, -1)) + elif pooled is None: + horizons = default + else: + default = list(range(0, self.post_window + 1)) + if isinstance(pooled, int): + horizons = list(range(0, pooled + 1)) + elif pooled is None: + horizons = default + + if isinstance(pooled, tuple) and len(pooled) == 2: + start, end = pooled + horizons = list(range(start, end + 1)) + elif not (pooled is None or isinstance(pooled, int)): + raise ValueError(f"{kind}_pooled must be None, an int, or a length-2 tuple") + + if kind == "pre": + # Exclude h=-1: it is the fixed base-period reference (coefficient 0), + # not an estimable placebo, so it cannot enter a pooled pre window. + supported_horizons = set(range(-self.pre_window, -1)) + else: + supported_horizons = set(range(0, self.post_window + 1)) + + invalid_horizons = [horizon for horizon in horizons if horizon not in supported_horizons] + if invalid_horizons: + raise ValueError( + f"Requested pooled {kind} horizons {invalid_horizons} fall outside the supported {kind} " + f"window {sorted(supported_horizons)}" + ) + + return horizons + + def fit( + self, + data, + outcome, + unit, + time, + treatment, + covariates=None, + ylags=0, + dylags=0, + absorb=None, + post_pooled=None, + pre_pooled=None, + only_event=False, + only_pooled=False, + ): + self.results_ = None + self.is_fitted_ = False + self._fit_meta = None + + # Validate covariate/absorb shape and reserved names BEFORE building the + # required-columns list, so a string (e.g. absorb="region") raises the + # precise "must be a list" error rather than an iterate-characters + # missing-column error, and a name colliding with an LPDiD working column + # is rejected up front. + for _arg_name, _arg in (("covariates", covariates), ("absorb", absorb)): + if isinstance(_arg, str): + raise ValueError(f"{_arg_name} must be a list of column names, not a string") + for _col in _arg or []: + if isinstance(_col, str) and (_col.startswith("_") or _col == "horizon"): + raise ValueError( + f"{_arg_name} column '{_col}' collides with an LPDiD reserved/internal " + "name (names starting with '_' or named 'horizon' are reserved by the " + "estimator's working columns); rename it." + ) + + required = [outcome, unit, time, treatment] + if covariates: + required.extend(covariates) + if absorb: + required.extend(absorb) + if self.cluster: + required.append(self.cluster) + missing = [col for col in required if col not in data.columns] + if missing: + raise ValueError(f"Missing columns: {missing}") + if only_event and only_pooled: + raise ValueError("only_event and only_pooled cannot both be True") + if not isinstance(ylags, int) or isinstance(ylags, bool) or ylags < 0: + raise ValueError("ylags must be a non-negative integer") + if not isinstance(dylags, int) or isinstance(dylags, bool) or dylags < 0: + raise ValueError("dylags must be a non-negative integer") + if not pd.api.types.is_numeric_dtype(data[time]): + raise ValueError( + "LPDiD requires a numeric `time` column with integer-spaced periods: " + "long differences use t-1 / t+h arithmetic on the time labels. Map " + "irregular or datetime periods to consecutive integers first." + ) + _unique_times = np.sort(pd.unique(data[time])) + if not np.allclose(_unique_times, np.round(_unique_times)): + raise ValueError( + "LPDiD requires integer-valued `time` labels (e.g. 0, 1, 2 or 2000, " + "2001, ...): the per-unit calendar grid is built with t-1 / t+h integer " + "arithmetic. Map fractional or datetime periods to consecutive integers first." + ) + if len(_unique_times) > 1 and not np.allclose(np.diff(_unique_times), 1.0): + raise ValueError( + "LPDiD requires integer-spaced `time` periods (consecutive, spaced by 1); " + "the global time grid is irregular (e.g. 2000, 2002, ...). Remap periods " + "to consecutive integers first so t-1 / t+h horizons are well defined." + ) + if (covariates or absorb or ylags or dylags) and not self.reweight: + warnings.warn( + "LPDiD: covariate-style controls (covariates, outcome lags `ylags`, " + "first-difference lags `dylags`, and absorbed factors) are entering the " + "long-difference regression by direct inclusion (reweight=False). Per " + "Dube, Girardi, Jorda & Taylor (2025) online Appendix B.2.2, direct " + "inclusion preserves the non-negative LP-DiD weighting result only under " + "linear and homogeneous control effects (Assumption 6 / Appendix B.2.1); " + "otherwise the implied weights are not guaranteed non-negative. For the " + "recommended regression-adjustment path, set reweight=True " + "(Appendix B.2.2 / Section 4.1).", + UserWarning, + stacklevel=2, + ) + + cluster = self.cluster or unit + pre_horizons = self._resolve_pooled_horizons(pre_pooled, kind="pre") + post_horizons = self._resolve_pooled_horizons(post_pooled, kind="post") + panel = self._prepare_panel( + data, + outcome=outcome, + unit=unit, + time=time, + treatment=treatment, + cluster=cluster, + covariates=covariates, + ylags=ylags, + dylags=dylags, + absorb=absorb, + ) + panel["_common_event_ok"] = self._common_clean_sample_indicator( + panel, + unit=unit, + time=time, + outcome=outcome, + max_post_horizon=self.post_window, + ) + panel["_common_pooled_ok"] = self._common_clean_sample_indicator( + panel, + unit=unit, + time=time, + outcome=outcome, + max_post_horizon=max(post_horizons) if post_horizons else 0, + ) + treatment_by_unit = panel.groupby(unit)["_treated"].max() + event_study = None + pooled = None + + if not only_pooled: + event_rows = [] + for horizon in range(-self.pre_window, self.post_window + 1): + if horizon == -1: + estimate = { + "coefficient": 0.0, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_low": np.nan, + "conf_high": np.nan, + "n_obs": np.nan, + } + else: + estimate = self._estimate_horizon( + panel, + outcome=outcome, + unit=unit, + time=time, + horizon=horizon, + covariates=covariates, + ylags=ylags, + dylags=dylags, + absorb=absorb, + ) + event_rows.append( + { + "horizon": horizon, + **estimate, + } + ) + event_study = pd.DataFrame(event_rows) + + if not only_event: + pre_estimate = self._estimate_window( + panel, + outcome=outcome, + unit=unit, + time=time, + horizons=pre_horizons, + kind="pre", + covariates=covariates, + ylags=ylags, + dylags=dylags, + absorb=absorb, + ) + post_estimate = self._estimate_window( + panel, + outcome=outcome, + unit=unit, + time=time, + horizons=post_horizons, + kind="post", + covariates=covariates, + ylags=ylags, + dylags=dylags, + absorb=absorb, + ) + pooled = pd.DataFrame( + [ + { + "window": "pre", + **pre_estimate, + }, + { + "window": "post", + **post_estimate, + }, + ] + ) + + # Headline G = realized clusters in the pooled-post (headline ATT) sample + # when computed; otherwise the panel-level unit-cluster count. Per-horizon + # rows carry their own realized n_clusters in the event_study/pooled tables. + headline_n_clusters = int(panel["_cluster"].nunique()) + if pooled is not None: + _post = pooled.loc[pooled["window"] == "post", "n_clusters"] + if not _post.empty and pd.notna(_post.iloc[0]): + headline_n_clusters = int(_post.iloc[0]) + + self.results_ = LPDiDResults( + event_study=event_study, + pooled=pooled, + n_obs=len(data), + n_treated_units=int(treatment_by_unit.gt(0).sum()), + n_control_units=int(treatment_by_unit.eq(0).sum()), + pre_window=self.pre_window, + post_window=self.post_window, + control_group=self.control_group, + reweight=self.reweight, + no_composition=self.no_composition, + pmd=self.pmd, + alpha=self.alpha, + cluster_name=cluster, + n_clusters=headline_n_clusters, + vcov_type=( + "if_cluster" + if (self.reweight and bool(covariates or ylags or dylags or absorb)) + else "hc1" + ), + rank_deficient_action=self.rank_deficient_action, + covariates=list(covariates) if covariates else None, + absorb=list(absorb) if absorb else None, + ylags=ylags, + dylags=dylags, + ) + self._fit_meta = {"cluster": cluster, "outcome": outcome, "unit": unit, "time": time} + self.is_fitted_ = True + return self.results_ + + def get_params(self) -> Dict[str, Any]: + return { + "pre_window": self.pre_window, + "post_window": self.post_window, + "control_group": self.control_group, + "reweight": self.reweight, + "no_composition": self.no_composition, + "pmd": self.pmd, + "alpha": self.alpha, + "cluster": self.cluster, + "rank_deficient_action": self.rank_deficient_action, + } + + def set_params(self, **params: Any) -> "LPDiD": + previous_values = {} + for key, value in params.items(): + if hasattr(self, key): + previous_values[key] = getattr(self, key) + setattr(self, key, value) + else: + raise ValueError(f"Unknown parameter: {key}") + try: + self._validate_params() + except ValueError: + for key, value in previous_values.items(): + setattr(self, key, value) + raise + return self diff --git a/diff_diff/lpdid_results.py b/diff_diff/lpdid_results.py new file mode 100644 index 00000000..2ec79733 --- /dev/null +++ b/diff_diff/lpdid_results.py @@ -0,0 +1,256 @@ +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import pandas as pd + + +@dataclass +class LPDiDResults: + """Results container for the :class:`~diff_diff.lpdid.LPDiD` estimator. + + Holds the per-horizon ``event_study`` table and the ``pooled`` pre/post + table (each a :class:`pandas.DataFrame` with ``coefficient``, ``se``, + ``t_stat``, ``p_value``, ``conf_low``, ``conf_high``, ``n_obs``, + ``n_clusters`` columns). The headline ATT is the pooled ``post`` row. + + ``n_control_units`` counts **never-treated** units only (the library-wide + field convention, surfaced as "Never-treated units" in ``summary()``); under + ``control_group="clean"`` the realized control pool at each horizon also + includes not-yet-treated cohorts, whose per-horizon counts live in the + ``n_obs`` / ``n_clusters`` columns of the tables. + """ + + event_study: Optional[pd.DataFrame] + pooled: Optional[pd.DataFrame] + n_obs: int + n_treated_units: int + n_control_units: int + pre_window: int + post_window: int + control_group: str + reweight: bool + no_composition: bool + pmd: Optional[Union[str, int]] + alpha: float = 0.05 + cluster_name: Optional[str] = None + n_clusters: Optional[int] = None + vcov_type: str = "hc1" + rank_deficient_action: str = "warn" + covariates: Optional[List[str]] = None + absorb: Optional[List[str]] = None + ylags: int = 0 + dylags: int = 0 + + # ------------------------------------------------------------------ + # internal helpers + # ------------------------------------------------------------------ + @property + def estimand(self) -> str: + return "equally-weighted ATT" if self.reweight else "variance-weighted ATT" + + def _base_period_label(self) -> str: + if self.pmd == "max": + return "premean (all available pretreatment periods)" + if isinstance(self.pmd, int) and not isinstance(self.pmd, bool): + return f"premean (last {self.pmd} pretreatment periods)" + return "first-lag (t-1)" + + def _pooled_row(self, window: str) -> Optional[pd.Series]: + if self.pooled is None: + return None + match = self.pooled.loc[self.pooled["window"] == window] + if match.empty: + return None + return match.iloc[0] + + # ------------------------------------------------------------------ + # headline inference aliases (over the pooled `post` row) + # ------------------------------------------------------------------ + @property + def att(self) -> float: + row = self._pooled_row("post") + return float(row["coefficient"]) if row is not None else float("nan") + + @property + def se(self) -> float: + row = self._pooled_row("post") + return float(row["se"]) if row is not None else float("nan") + + @property + def t_stat(self) -> float: + row = self._pooled_row("post") + return float(row["t_stat"]) if row is not None else float("nan") + + @property + def p_value(self) -> float: + row = self._pooled_row("post") + return float(row["p_value"]) if row is not None else float("nan") + + @property + def conf_int(self) -> Tuple[float, float]: + row = self._pooled_row("post") + if row is None: + return (float("nan"), float("nan")) + return (float(row["conf_low"]), float(row["conf_high"])) + + # ------------------------------------------------------------------ + # serialization + # ------------------------------------------------------------------ + def to_dataframe(self, level: str = "event") -> pd.DataFrame: + if level == "event": + if self.event_study is None: + raise ValueError("event_study dataframe was not computed") + return self.event_study.copy() + if level == "pooled": + if self.pooled is None: + raise ValueError("pooled dataframe was not computed") + return self.pooled.copy() + raise ValueError("level must be 'event' or 'pooled'") + + def to_dict(self) -> Dict[str, Any]: + pre = self._pooled_row("pre") + ci = self.conf_int + result: Dict[str, Any] = { + "att": self.att, + "se": self.se, + "t_stat": self.t_stat, + "p_value": self.p_value, + "conf_int_lower": ci[0], + "conf_int_upper": ci[1], + "pre_att": float(pre["coefficient"]) if pre is not None else float("nan"), + "pre_se": float(pre["se"]) if pre is not None else float("nan"), + "n_obs": self.n_obs, + "n_treated_units": self.n_treated_units, + "n_control_units": self.n_control_units, + "pre_window": self.pre_window, + "post_window": self.post_window, + "control_group": self.control_group, + "reweight": self.reweight, + "no_composition": self.no_composition, + "pmd": self.pmd, + "estimand": self.estimand, + "alpha": self.alpha, + "vcov_type": self.vcov_type, + "rank_deficient_action": self.rank_deficient_action, + "ylags": self.ylags, + "dylags": self.dylags, + "covariates": self.covariates, + "absorb": self.absorb, + } + if self.cluster_name is not None: + result["cluster_name"] = self.cluster_name + if self.n_clusters is not None: + result["n_clusters"] = self.n_clusters + result["inference_method"] = "cluster_robust" + return result + + # ------------------------------------------------------------------ + # text summary + # ------------------------------------------------------------------ + def summary(self) -> str: + from diff_diff.results import _format_vcov_label, _get_significance_stars + + # Confidence intervals in the event_study / pooled tables are computed at + # fit time using ``self.alpha``; the displayed level must match them, so + # summary() does not accept an alpha override (it would relabel without + # recomputing the intervals). + ci_pct = int(round((1 - self.alpha) * 100)) + width = 88 + bar = "=" * width + dash = "-" * width + + def _fmt(x: Any, nd: int = 4) -> str: + try: + xf = float(x) + except (TypeError, ValueError): + return "" + return "" if np.isnan(xf) else f"{xf:.{nd}f}" + + lines: List[str] = [ + bar, + "Local Projections DiD (Dube, Girardi, Jorda & Taylor 2025) Results".center(width), + bar, + f"Observations: {self.n_obs} Treated units: {self.n_treated_units}" + f" Never-treated units: {self.n_control_units}", + f"Estimand: {self.estimand} Control group: {self.control_group}", + f"Base period: {self._base_period_label()} No composition: {self.no_composition}", + ] + if self.covariates or self.absorb or self.ylags or self.dylags: + cov_path = "regression-adjustment" if self.reweight else "direct inclusion" + lag_bits = [] + if self.ylags: + lag_bits.append(f"ylags={self.ylags}") + if self.dylags: + lag_bits.append(f"dylags={self.dylags}") + lag_str = (" " + ", ".join(lag_bits)) if lag_bits else "" + lines.append( + f"Covariates: {self.covariates or []} Absorb: {self.absorb or []}" + f"{lag_str} ({cov_path})" + ) + if self.vcov_type == "if_cluster": + # Regression-adjustment path: influence-function cluster variance + # (ImputationDiD/BJS family), not an OLS CR1 sandwich. + g = f", G={self.n_clusters}" if self.n_clusters else "" + vcov_label = f"Influence-function cluster-robust at {self.cluster_name}{g}" + else: + vcov_label = _format_vcov_label( + self.vcov_type, + cluster_name=self.cluster_name, + n_clusters=self.n_clusters, + n_obs=self.n_obs, + ) + if vcov_label: + lines.append(f"Std. errors: {vcov_label}") + + header = ( + f"{'':>8} {'Estimate':>10} {'Std.Err':>10} {'t':>8} {'P>|t|':>8}" + f" [{ci_pct}% Conf. Int.]" + ) + + def _table(df: pd.DataFrame, key: str) -> List[str]: + rows: List[str] = [dash, header, dash] + for _, r in df.iterrows(): + label = r[key] + if key == "horizon" and int(r[key]) == -1: + rows.append(f"{int(label):>8} {'0.0000':>10} {'(reference)':>10}") + continue + p = r["p_value"] + stars = "" if pd.isna(p) else _get_significance_stars(float(p)) + label_str = f"{int(label):>8}" if key == "horizon" else f"{str(label):>8}" + rows.append( + f"{label_str} {_fmt(r['coefficient']):>10} {_fmt(r['se']):>10}" + f" {_fmt(r['t_stat'], 2):>8} {_fmt(r['p_value'], 3):>8}" + f" [{_fmt(r['conf_low']):>9}, {_fmt(r['conf_high']):>9}] {stars}" + ) + return rows + + if self.event_study is not None: + lines.append("") + lines.append("Event study (relative horizon):") + lines.extend(_table(self.event_study, "horizon")) + if self.pooled is not None: + lines.append("") + lines.append("Pooled (pre = placebo, post = ATT):") + lines.extend(_table(self.pooled, "window")) + + lines.append(bar) + lines.append("Signif. codes: *** p<0.001, ** p<0.01, * p<0.05") + return "\n".join(lines) + + def print_summary(self) -> None: + print(self.summary()) + + def __repr__(self) -> str: + cluster = f", cluster={self.cluster_name}, G={self.n_clusters}" if self.cluster_name else "" + att = self.att + se = self.se + att_s = "nan" if np.isnan(att) else f"{att:.4f}" + se_s = "nan" if np.isnan(se) else f"{se:.4f}" + return ( + "LPDiDResults(" + f"estimand={'reweight' if self.reweight else 'variance-weighted'}, " + f"post_ATT={att_s}, SE={se_s}, " + f"pre_window={self.pre_window}, post_window={self.post_window}, " + f"control_group={self.control_group!r}{cluster})" + ) diff --git a/docs/api/index.rst b/docs/api/index.rst index 426e55f0..145ba611 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -30,6 +30,7 @@ Core estimator classes for DiD analysis: diff_diff.TwoStageDiD diff_diff.SpilloverDiD diff_diff.WooldridgeDiD + diff_diff.LPDiD diff_diff.BaconDecomposition diff_diff.StaggeredTripleDifference @@ -70,6 +71,7 @@ Result containers returned by estimators: diff_diff.SpilloverDiDResults diff_diff.BaconDecompositionResults diff_diff.wooldridge_results.WooldridgeDiDResults + diff_diff.lpdid_results.LPDiDResults diff_diff.Comparison2x2 diff_diff.StaggeredTripleDiffResults diff_diff.TWFEWeightsResult @@ -318,6 +320,7 @@ Estimators two_stage spillover wooldridge_etwfe + lpdid bacon Infrastructure diff --git a/docs/api/lpdid.rst b/docs/api/lpdid.rst new file mode 100644 index 00000000..e6970161 --- /dev/null +++ b/docs/api/lpdid.rst @@ -0,0 +1,158 @@ +Local Projections Difference-in-Differences +=========================================== + +Local Projections DiD (LP-DiD) estimator for staggered, absorbing-treatment +event studies, from Dube, Girardi, Jordà & Taylor (2025). + +LP-DiD estimates a separate regression at each event-time horizon ``h`` of a +long difference of the outcome (``y_{i,t+h} - y_{i,t-1}``) on the +treatment-switch indicator, restricted to a flexible "clean control" sample of +newly-treated observations and not-yet-treated controls. Excluding +already-treated units from the control group removes the negative-weighting +bias of naive two-way fixed effects, so the default (variance-weighted) +estimand is a strictly non-negatively-weighted average of cohort effects. + +.. note:: + + This release implements the **absorbing-treatment main path**: treatment is + binary and, once switched on, stays on. The estimator rejects panels where a + unit's treatment turns off. Non-absorbing (switch on/off) treatment and + survey-design support are planned follow-ups. Covariates and absorbed fixed + effects are supported; under ``reweight=False`` they enter by direct + inclusion, which preserves the non-negative weighting result only under + homogeneous covariate effects (online Appendix B.2.2) — the + regression-adjustment path (``reweight=True``) is preferred for + covariate-adjusted designs (it does not auto-switch; the default remains + ``reweight=False``, which emits the warning). The ``time`` column must be + numeric with integer-spaced periods (long differences use ``t-1`` / ``t+h`` + arithmetic on the labels); map irregular or datetime periods to consecutive + integers first. See ``docs/methodology/REGISTRY.md`` for the full contract. + +**When to use LPDiD:** + +- Staggered, **absorbing** adoption where you want a fast, transparent, + regression-based event study free of negative weighting +- You want both a dynamic event-study path and a single pooled pre/post ATT +- You want to flexibly choose the pretreatment base period (first-lag or + premean-differenced) or hold the post-treatment sample composition fixed across post horizons +- You want an estimator that is numerically equivalent to Callaway-Sant'Anna + (reweighted) or to a Cengiz et al. (2019)-style stacked regression + (variance-weighted), but much faster + +**Reference:** Dube, A., Girardi, D., Jordà, Ò., & Taylor, A. M. (2025). A Local +Projections Approach to Difference-in-Differences. *Journal of Applied +Econometrics*, 40(5), 741-758. + +.. module:: diff_diff.lpdid + +LPDiD +----- + +Main estimator class for Local Projections Difference-in-Differences. + +.. autoclass:: diff_diff.LPDiD + :no-index: + :members: + :undoc-members: + :show-inheritance: + + .. rubric:: Methods + + .. autosummary:: + + ~LPDiD.fit + ~LPDiD.get_params + ~LPDiD.set_params + +LPDiDResults +------------ + +Results container for LP-DiD estimation (event-study and pooled tables). + +.. autoclass:: diff_diff.lpdid_results.LPDiDResults + :no-index: + :members: + :undoc-members: + :show-inheritance: + + .. rubric:: Methods + + .. autosummary:: + + ~LPDiDResults.summary + ~LPDiDResults.print_summary + ~LPDiDResults.to_dataframe + ~LPDiDResults.to_dict + +Example Usage +------------- + +Basic usage (LP-DiD takes a binary, absorbing treatment indicator):: + + from diff_diff import LPDiD, generate_staggered_data + + data = generate_staggered_data(n_units=300, n_periods=12, + cohort_periods=[4, 7, 10], seed=42) + # Binary absorbing indicator: 1 from a unit's first treated period onward. + data["treated"] = (data["period"] >= data["first_treat"]).astype(int) + + lp = LPDiD(pre_window=5, post_window=4) + results = lp.fit(data, outcome="outcome", unit="unit", + time="period", treatment="treated") + results.print_summary() + print(results.event_study) # per-horizon coefficients + print(results.pooled) # pooled pre (placebo) / post (ATT) rows + +Equally-weighted ATT (numerically equivalent to Callaway-Sant'Anna):: + + lp_rw = LPDiD(pre_window=5, post_window=4, reweight=True) + results_rw = lp_rw.fit(data, outcome="outcome", unit="unit", + time="period", treatment="treated") + print(f"Variance-weighted ATT: {results.att:.4f} (SE={results.se:.4f})") + print(f"Equally-weighted ATT: {results_rw.att:.4f} (SE={results_rw.se:.4f})") + +Premean-differenced base period and fixed-composition sample:: + + lp_pmd = LPDiD(pre_window=5, post_window=4, pmd="max", no_composition=True) + results_pmd = lp_pmd.fit(data, outcome="outcome", unit="unit", + time="period", treatment="treated") + +Comparison with Other Staggered Estimators +------------------------------------------ + +.. list-table:: + :header-rows: 1 + :widths: 20 27 27 26 + + * - Feature + - LPDiD + - CallawaySantAnna + - ImputationDiD + * - Approach + - Per-horizon long-difference LP regression on clean controls + - Separate 2x2 DiD aggregation + - Impute Y(0) via FE model + * - Treatment + - Binary, absorbing (this release) + - Binary, absorbing + - Binary, absorbing + * - Default estimand + - Variance-weighted ATT (non-negative weights) + - Equally-weighted ATT + - Equally-weighted ATT + * - Equivalences + - Reweighted == CS; variance-weighted == Cengiz (2019)-style stacking (not ``diff_diff.StackedDiD``, which is Wing et al. 2024 Q-weights); PMD single-cohort == BJS + - Baseline + - == reweighted PMD LP-DiD (single cohort) + * - Covariates + - Supported (regression adjustment preferred; direct inclusion under homogeneity) + - Supported (OR, IPW, DR) + - Supported + * - Inference + - Cluster-robust at unit (default) + - Multiplier bootstrap + - Influence-function cluster variance + * - Speed + - Very fast (stack of small OLS fits) + - Slower (pairwise group-time) + - Fast diff --git a/docs/choosing_estimator.rst b/docs/choosing_estimator.rst index 15df86cf..2fd1223d 100644 --- a/docs/choosing_estimator.rst +++ b/docs/choosing_estimator.rst @@ -129,6 +129,10 @@ Quick Reference - Nonlinear outcomes or saturated OLS - Conditional parallel trends - OLS: direct coefficients; logit/Poisson: ASF-based ATT + * - ``LPDiD`` + - Fast staggered (absorbing) event studies without negative weighting + - Parallel trends, no anticipation; absorbing treatment + - Event-study path + pooled pre/post ATT * - ``BaconDecomposition`` - TWFE diagnostic - (diagnostic tool) diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index c0408168..c6e53ef8 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -62,6 +62,9 @@ groups: wooldridge: - diff_diff/wooldridge.py - diff_diff/wooldridge_results.py + lpdid: + - diff_diff/lpdid.py + - diff_diff/lpdid_results.py visualization: - diff_diff/visualization/__init__.py - diff_diff/visualization/_common.py @@ -590,6 +593,30 @@ sources: - path: docs/choosing_estimator.rst type: user_guide + # ── LPDiD (lpdid group) ──────────────────────────────────────────── + + diff_diff/lpdid.py: + drift_risk: medium + docs: + - path: docs/methodology/REGISTRY.md + section: "LPDiD" + type: methodology + - path: docs/api/lpdid.rst + type: api_reference + - path: README.md + section: "Estimators (one-line catalog entry)" + type: user_guide + - path: docs/references.rst + type: user_guide + - path: diff_diff/guides/llms-full.txt + section: "LPDiD" + type: user_guide + - path: diff_diff/guides/llms.txt + section: "Estimators" + type: user_guide + - path: docs/choosing_estimator.rst + type: user_guide + # ── TROP (trop group) ────────────────────────────────────────────── diff_diff/trop.py: diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 20647cab..f54fe684 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -1829,7 +1829,7 @@ Weights are **always non-negative** (the central result). Via Frisch-Waugh-Lovel ### Edge Cases -- **Composition effects (Section 3.6):** the treated/clean-control set can change across horizons. `no_composition` tightens the clean-control condition to `D_{i,t+H}=0` at all horizons (and excludes cohorts with `p_g > T-H` to fix the treated set). Costs statistical power. +- **Composition effects (Section 3.6):** the treated/clean-control set can change across horizons. `no_composition` tightens the clean-control condition to `D_{i,t+H}=0` at all horizons (and excludes cohorts with `p_g > T-H` to fix the treated set). Costs statistical power. **Implementation note:** this fixes the *post-treatment* composition (every post horizon shares the same realized sample, even on unbalanced panels); pre-treatment placebo horizons use whatever pre-period data is available. Reweighting denominators and the regression-adjustment counterfactual are computed from the realized post-drop sample (not the pre-drop panel) so they stay consistent with the regression's risk set on unbalanced panels. - **Bias-variance (Sections 3.3, 5.3):** variance weighting (default) -> lower variance, some bias; equal weighting (`reweight`) -> unbiased, higher variance. Variance won at short horizons, equal at long horizons in the paper's simulation. - **PMD vs first-lag (Section 3.4):** PMD gains efficiency under low autocorrelation but can amplify bias if PT holds only in some pretreatment periods; first-lag relies on weaker PT (Marcus & Sant'Anna 2021). Choose the base period ex-ante. - **Covariate-weight positivity (online Appendix B.2):** direct covariate inclusion keeps non-negative weights ONLY under linear + homogeneous covariate effects (B.2.1; main-text Assumption 6); in the general case (B.2.2) weights are not guaranteed positive -> prefer the RA covariate path (the direct path should carry a homogeneity-assumption warning). @@ -1837,20 +1837,28 @@ Weights are **always non-negative** (the central result). Via Frisch-Waugh-Lovel ### Deviations from the paper / from R / library extensions -*To be populated in PR-B (source + tests). Anticipated entries: (1) the analytical/cluster SE convention (paper specifies none - implementation choice vs Stata `lpdid` `vce(cluster unit)`); (2) any RA-path influence-function variance; (3) the `pmd="max"` / integer-`k` panel-start edge behavior vs the package; (4) absorbing-only scope in the first release, with non-absorbing (Section 4.2) deferred.* +The paper specifies no standard-error formula (Section 1 defers to "standard, well-understood techniques"); the reference Stata `lpdid` uses `vce(cluster unit)`. The entries below document diff-diff's inference and scope choices. + +1. **Note:** Standard errors are **cluster-robust at the unit level by default** - `cluster=None` auto-clusters at the unit identifier and the results record `cluster_name`/`n_clusters` - with a `t(G-1)` reference distribution (G = realized clusters in each horizon's clean-control sample). Matches Stata `lpdid` `vce(cluster unit)`; the paper prescribes no SE. +2. **Note:** The regression-adjustment (RA) covariate path (`reweight=True` with covariates/absorb) reports an **influence-function cluster variance** `sum_c (sum_{i in c} psi_i)^2 / n^2`, in the same family as `ImputationDiD`'s Theorem-3 / BJS variance (see "IF-based variance estimators vs analytical-sandwich estimators" above). Its single Gram inversion is routed through `linalg._rank_guarded_inv` (finite SE on the identified subspace under near-collinearity; NaN at rank 0). Unlike the default/weighted `solve_ols` `hc1`-cluster path - which applies the `(G/(G-1))*((n-1)/(n-k))` finite-sample factor - the RA IF variance currently carries **no finite-sample factor** (the ImputationDiD convention), while both paths share the `t(G-1)` reference. This RA-vs-default small-sample-scaling asymmetry is documented here and will be validated against the reference R packages and reconciled in the R-parity follow-up (PR-B2). +3. **Note:** Direct covariate inclusion (`reweight=False` with covariates/absorb) emits a `UserWarning`: per online Appendix B.2.2 it preserves the non-negative LP-DiD weighting result only under linear and homogeneous covariate effects, so the regression-adjustment path (`reweight=True`) is preferred. +4. **Deviation from R:** Scope - this release implements the **absorbing-treatment main path only** (the estimator raises on non-absorbing input). Non-absorbing treatment (Section 4.2) and survey-design support are deferred to later PRs; the reference `lpdid` packages support non-absorbing treatment. +5. **Note:** LP-DiD's per-unit quantities (outcome lags `ylags`, first-difference lags `dylags`, integer-`pmd` premean baselines, treatment-entry detection) are **calendar** quantities (`t-1`, `t-k`), so the estimator requires integer-valued, globally consecutive `time` labels. A unit with an **interior time gap** is handled by reindexing that unit to its complete interior calendar grid `[min_t, max_t]`, computing the features on the grid, then **restricting back to the observed rows** - so a lag/first-difference spanning a gap is NaN and the observation fails closed (never the previous-*observed* row), and no synthetic gap row enters a regression. A gap-free panel skips this entirely and is bit-identical. **Entry = first OBSERVED treated period** (`min(t | D_it=1)`): an unobserved pre-onset gap cannot move a cohort earlier, the only well-defined convention when the true switch falls in an unobserved period. +6. **Note (pooled estimand):** The pooled pre/post ATT (the headline `results.att` is the pooled-post row) is the **unit-equal-weighted average of each unit-event-time's mean long difference** over the window - `mean_h(y_{i,t+h}) - baseline_{i,t}`, one observation per (unit, event-time), regressed on the treatment-switch indicator with event-time fixed effects on the **fixed-composition** sample (only units observing *every* pooled target, with clean controls required through `max(h)`). This equals the mean of the per-horizon event-study coefficients on a balanced panel; under cross-horizon composition changes it differs from the authors' horizon-**stacked** pooled regression (each (unit, event-time, horizon) long difference as a separate observation). It is therefore reported as a fixed-composition pooled ATT, and exact parity with the reference `lpdid` pooled specification is validated/reconciled in the R-parity follow-up (PR-B2). ### Implementation Checklist -- [ ] Per-horizon long-difference OLS with time FE, no unit FE; `h=-1` reference fixed at 0 (PR-B) -- [ ] Clean-control sample restriction (absorbing: `D_{i,t+h}=0`) (PR-B) -- [ ] Variance-weighted (default) + reweighted (equal-weight) estimands (PR-B) -- [ ] Regression-adjustment covariate path (recommended) + direct-inclusion path with homogeneity warning (PR-B) -- [ ] PMD base period; pooled pre/post estimands (PR-B) -- [ ] `no_composition` option (PR-B) -- [ ] Cluster-robust SE at unit level by default; NaN-consistent inference via `safe_inference` (PR-B) -- [ ] `LPDiDResults` with `summary()` / `to_dict()` / cluster metadata (PR-B) -- [ ] Layered tests: analytical DGPs + cross-estimator equivalence (CS / BJS / Stacked / DiD) + self-generated R-parity (PR-B) -- [ ] doc-deps.yaml mapping for `diff_diff/lpdid.py` + `lpdid_results.py`; llms.txt catalog entry (PR-B, test-enforced) +- [x] Per-horizon long-difference OLS with time FE, no unit FE; `h=-1` reference fixed at 0 (PR-B1) +- [x] Clean-control sample restriction (absorbing: `D_{i,t+h}=0`) (PR-B1) +- [x] Variance-weighted (default) + reweighted (equal-weight) estimands (PR-B1) +- [x] Regression-adjustment covariate path (recommended) + direct-inclusion path with homogeneity warning (PR-B1) +- [x] PMD base period; pooled pre/post estimands (PR-B1) +- [x] `no_composition` option (PR-B1) +- [x] Cluster-robust SE at unit level by default; NaN-consistent inference via `safe_inference` (PR-B1) +- [x] `LPDiDResults` with `summary()` / `to_dict()` / cluster metadata (PR-B1) +- [x] doc-deps.yaml mapping for `diff_diff/lpdid.py` + `lpdid_results.py`; llms.txt / llms-full.txt catalog entries (PR-B1, test-enforced) +- [x] B1 pure-Python tests: analytical DGPs + cross-estimator equivalence (CS / BJS / DiD; Cengiz-stacked dropped, documented) + unbalanced / interior-gap / RA-overlap / pmd-missing edge cases (PR-B1) +- [ ] B2: self-generated R-parity (authors' `danielegirardi/lpdid` + `alexCardazzi/lpdid` cross-check) - [ ] Non-absorbing extension (Section 4.2) - deferred to a later PR - [ ] Survey-design support - deferred to a later PR diff --git a/tests/_dgp_utils.py b/tests/_dgp_utils.py index ff5a8b1d..a3d2f28e 100644 --- a/tests/_dgp_utils.py +++ b/tests/_dgp_utils.py @@ -283,3 +283,69 @@ def generate_butts_staggered_dgp( } ) return pd.DataFrame(rows) + + +def make_lpdid_panel( + *, + cohorts=(5, 8), + n_per_cohort: int = 20, + n_never: int = 30, + n_periods: int = 12, + tau: Callable[[int], float] = lambda k: 1.0 + 0.5 * k, + unit_fe_sd: float = 0.5, + time_trend: float = 0.5, + error_sd: float = 0.2, + heterogeneous: bool = True, + seed: int = 20260628, +) -> pd.DataFrame: + """Staggered, absorbing-treatment panel for LP-DiD tests. + + Emits columns the two estimator families both need: + + - ``treat``: binary absorbing indicator ``1[t >= first_treat]`` (LPDiD's + ``treatment=`` input; never-treated units are always 0). + - ``first_treat``: onset period; ``np.inf`` for never-treated (the repo + convention consumed by CallawaySantAnna / ImputationDiD / StackedDiD). + + Untreated potential outcome is ``alpha_i + time_trend * t + noise`` so + parallel trends and no-anticipation hold by construction. The dynamic + effect at event time ``k = t - g`` is ``tau(k)`` for ``k >= 0`` (zero + otherwise), scaled per cohort when ``heterogeneous=True`` so different + cohorts have genuinely different effect paths. + """ + rng = np.random.default_rng(seed) + rows = [] + uid = 0 + for ci, g in enumerate(cohorts): + scale = (1.0 + 0.3 * ci) if heterogeneous else 1.0 + for _ in range(n_per_cohort): + uid += 1 + alpha = rng.normal(0.0, unit_fe_sd) + for t in range(n_periods): + y0 = alpha + time_trend * t + rng.normal(0.0, error_sd) + k = t - g + effect = scale * tau(k) if k >= 0 else 0.0 + rows.append( + { + "unit": uid, + "time": t, + "y": y0 + effect, + "treat": int(t >= g), + "first_treat": float(g), + } + ) + for _ in range(n_never): + uid += 1 + alpha = rng.normal(0.0, unit_fe_sd) + for t in range(n_periods): + y0 = alpha + time_trend * t + rng.normal(0.0, error_sd) + rows.append( + { + "unit": uid, + "time": t, + "y": y0, + "treat": 0, + "first_treat": np.inf, + } + ) + return pd.DataFrame(rows) diff --git a/tests/test_guides.py b/tests/test_guides.py index fadc85fc..63891aaf 100644 --- a/tests/test_guides.py +++ b/tests/test_guides.py @@ -770,3 +770,54 @@ def test_llms_full_stacked_fit_documents_covariates(self): ) # balance= must be documented somewhere in the section (constructor param) assert "balance" in section + + +class TestLLMsFullLPDiDCoverage: + """Pin the LPDiD section of llms-full.txt to the real API. + + Adding a public parameter to LPDiD.__init__ or LPDiD.fit() requires updating + diff_diff/guides/llms-full.txt — these tests catch drift. + """ + + def _lpdid_section(self): + text = get_llm_guide("full") + start = text.index("### LPDiD") + nxt = text.index("\n### ", start + 1) + return text[start:nxt] + + def test_llms_full_has_lpdid_section(self): + assert "### LPDiD" in get_llm_guide("full") + + def test_llms_full_lpdid_constructor_signature_matches_real_api(self): + import inspect + + from diff_diff import LPDiD + + sig_params = set(inspect.signature(LPDiD.__init__).parameters) + sig_params.discard("self") + section = self._lpdid_section() + block_start = section.index("LPDiD(") + block_end = section.index("\n)", block_start) + ctor_block = section[block_start:block_end] + for param in sig_params: + assert f"{param}:" in ctor_block or f"{param} " in ctor_block, ( + f"LPDiD constructor block in llms-full.txt is missing the real " + f"public parameter {param!r} (adding a public param requires updating " + f"the guide)." + ) + + def test_llms_full_lpdid_fit_signature_matches_real_api(self): + import inspect + + from diff_diff import LPDiD + + sig_params = set(inspect.signature(LPDiD.fit).parameters) + sig_params.discard("self") + section = self._lpdid_section() + fit_start = section.index("lpdid.fit(") + fit_block = section[fit_start : section.index(") -> ", fit_start)] + for param in sig_params: + assert f"{param}:" in fit_block or f"{param} " in fit_block, ( + f"LPDiD.fit() block in llms-full.txt is missing the real public " + f"parameter {param!r} (adding a public param requires updating the guide)." + ) diff --git a/tests/test_lpdid.py b/tests/test_lpdid.py new file mode 100644 index 00000000..c8ba4f62 --- /dev/null +++ b/tests/test_lpdid.py @@ -0,0 +1,1326 @@ +"""Tests for the LPDiD (Local Projections DiD) estimator. + +Validation strategy (B1, pure-Python, no R required): + +- ``TestLPDiDAPI`` / ``TestLPDiDEdgeCases``: parameter validation, idempotence, + NaN-consistent inference, absorbing-path enforcement. +- ``TestLPDiDFormula``: small hand-built panels with analytically-known + coefficients (true effect derivable by hand). +- ``TestLPDiDMethodology``: DGP-recovery on panels with a known dynamic effect + path, plus the clean-control / weighting properties. +- ``TestLPDiDCrossEstimator``: the equivalences proved in Dube, Girardi, Jordà & + Taylor (2025), validated against estimators already in the library + (point estimates only; SEs are anchored separately by the B2 R-parity layer). + +External R-parity (authors' ``danielegirardi/lpdid`` + ``alexCardazzi/lpdid``) +will live in ``tests/test_methodology_lpdid.py`` (added in PR-B2). +""" + +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from _dgp_utils import make_lpdid_panel # noqa: E402 + +from diff_diff import ( # noqa: E402 + CallawaySantAnna, + ImputationDiD, + LPDiD, + LPDiDResults, +) + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- +def _make_linear_panel(units): + """Build a panel from per-unit specs (treat_start=99 => never-treated).""" + rows = [] + for spec in units: + for time in spec["times"]: + rows.append( + { + "unit": spec["unit"], + "time": time, + "y": spec["y"][time], + "treat": int(time >= spec["treat_start"]), + **{ + k: v + for k, v in spec.items() + if k not in {"unit", "times", "y", "treat_start"} + }, + } + ) + return pd.DataFrame(rows) + + +def _event_coef(res, horizon=0): + es = res.event_study + return es.loc[es["horizon"] == horizon, "coefficient"].iloc[0] + + +def _event_row(res, horizon=0): + es = res.event_study + return es.loc[es["horizon"] == horizon].iloc[0] + + +# =========================================================================== +# API / validation +# =========================================================================== +class TestLPDiDAPI: + def test_get_params_round_trip(self): + est = LPDiD(pre_window=4, post_window=6, reweight=True, no_composition=True) + params = est.get_params() + assert params["pre_window"] == 4 + assert params["post_window"] == 6 + assert params["reweight"] is True + assert params["no_composition"] is True + + def test_set_params_updates_attributes(self): + est = LPDiD() + returned = est.set_params(pre_window=5, control_group="never_treated") + assert returned is est + assert est.pre_window == 5 + assert est.control_group == "never_treated" + + def test_rejects_invalid_control_group(self): + with pytest.raises(ValueError, match="control_group"): + LPDiD(control_group="bad") + + def test_rejects_invalid_rank_deficient_action(self): + with pytest.raises(ValueError, match="rank_deficient_action"): + LPDiD(rank_deficient_action="bad") + + def test_rejects_invalid_pmd_value(self): + with pytest.raises(ValueError, match="pmd"): + LPDiD(pmd="bad") + + def test_rejects_negative_window(self): + with pytest.raises(ValueError, match="pre_window"): + LPDiD(pre_window=-1) + with pytest.raises(ValueError, match="post_window"): + LPDiD(post_window=-2) + with pytest.raises(ValueError, match="pre_window"): + LPDiD(pre_window=True) # bool is not a valid int window + + def test_rejects_invalid_alpha(self): + for bad in (0.0, 1.0, 1.5, -0.1, True): + with pytest.raises(ValueError, match="alpha"): + LPDiD(alpha=bad) + + def test_rejects_non_numeric_time(self): + df = pd.DataFrame( + { + "unit": [1, 1, 2, 2], + "time": ["a", "b", "a", "b"], + "y": [1.0, 2, 1, 1], + "treat": [0, 1, 0, 0], + } + ) + with pytest.raises(ValueError, match="numeric `time`"): + LPDiD().fit(df, outcome="y", unit="unit", time="time", treatment="treat") + + def test_rejects_irregular_time_grid(self): + df = pd.DataFrame( + { + "unit": [1, 1, 1, 2, 2, 2], + "time": [2000, 2002, 2004] * 2, # spacing 2, not 1 + "y": [1.0, 2, 3, 1, 1, 1], + "treat": [0, 1, 1, 0, 0, 0], + } + ) + with pytest.raises(ValueError, match="integer-spaced"): + LPDiD(pre_window=1, post_window=1).fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + + def test_rejects_string_covariates(self): + df = pd.DataFrame( + { + "unit": [1, 1, 2, 2], + "time": [0, 1, 0, 1], + "y": [1.0, 2, 1, 1], + "treat": [0, 1, 0, 0], + "x": [1.0, 2, 3, 4], + } + ) + with pytest.raises(ValueError, match="covariates must be a list"): + LPDiD().fit( + df, outcome="y", unit="unit", time="time", treatment="treat", covariates="x" + ) + + def test_rejects_reserved_internal_column_names(self): + # a covariate/absorb colliding with an LPDiD working column (e.g. "_entry") + # must be rejected, not silently overwrite the internal column. + df = pd.DataFrame( + { + "unit": [1, 1, 2, 2], + "time": [0, 1, 0, 1], + "y": [1.0, 2, 1, 1], + "treat": [0, 1, 0, 0], + "_entry": [0.0, 0.0, 0.0, 0.0], + } + ) + with pytest.raises(ValueError, match="reserved"): + LPDiD().fit( + df, outcome="y", unit="unit", time="time", treatment="treat", covariates=["_entry"] + ) + + def test_rejects_missing_cluster_labels(self): + df = pd.DataFrame( + { + "unit": [1, 1, 1, 2, 2, 2], + "time": [0, 1, 2] * 2, + "y": [1.0, 3, 5, 1, 1, 1], + "treat": [0, 1, 1, 0, 0, 0], + "cl": [1, 1, 1, np.nan, 2, 2], + } + ) + with pytest.raises(ValueError, match="missing values"): + LPDiD(pre_window=1, post_window=1, cluster="cl").fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + + def test_set_params_rejects_unknown_key(self): + with pytest.raises(ValueError, match="Unknown parameter"): + LPDiD().set_params(nonexistent_param=1) + + def test_set_params_is_transactional(self): + # An invalid update must roll back ALL fields (snapshot + restore). + est = LPDiD(pre_window=2) + with pytest.raises(ValueError, match="control_group"): + est.set_params(pre_window=7, control_group="bad") + assert est.pre_window == 2 # rolled back, not left at 7 + + def test_requires_core_columns(self): + df = pd.DataFrame({"y": [1.0], "id": [1], "t": [0]}) + with pytest.raises(ValueError, match="Missing columns"): + LPDiD().fit(df, outcome="y", unit="id", time="t", treatment="treat") + + def test_rejects_only_event_and_only_pooled_together(self): + df = pd.DataFrame({"y": [1.0], "id": [1], "t": [0], "treat": [0]}) + with pytest.raises(ValueError, match="only_event"): + LPDiD().fit( + df, + outcome="y", + unit="id", + time="t", + treatment="treat", + only_event=True, + only_pooled=True, + ) + + def test_rejects_non_numeric_treatment_values(self): + df = pd.DataFrame( + { + "unit": [1, 1, 2, 2], + "time": [0, 1, 0, 1], + "y": [1, 2, 1, 1], + "treat": [0, "treated", 0, 0], + } + ) + with pytest.raises(ValueError, match="binary numeric"): + LPDiD().fit(df, outcome="y", unit="unit", time="time", treatment="treat") + + def test_rejects_non_binary_treatment_values(self): + df = pd.DataFrame( + {"unit": [1, 1, 2, 2], "time": [0, 1, 0, 1], "y": [1, 2, 1, 1], "treat": [0, 2, 0, 0]} + ) + with pytest.raises(ValueError, match="binary numeric"): + LPDiD().fit(df, outcome="y", unit="unit", time="time", treatment="treat") + + def test_event_dataframe_contains_reference_and_requested_horizons(self): + df = pd.DataFrame( + { + "unit": [1, 1, 1, 2, 2, 2, 3, 3, 3], + "time": [0, 1, 2] * 3, + "y": [1, 2, 4, 1, 1, 1, 2, 2, 2], + "treat": [0, 1, 1, 0, 0, 0, 0, 0, 0], + } + ) + res = LPDiD(pre_window=2, post_window=1).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + assert list(res.event_study["horizon"]) == [-2, -1, 0, 1] + + def test_pooled_dataframe_has_pre_and_post_rows(self): + df = pd.DataFrame( + { + "unit": [1, 1, 1, 1, 2, 2, 2, 2], + "time": [0, 1, 2, 3] * 2, + "y": [1, 2, 4, 6, 1, 1, 1, 1], + "treat": [0, 0, 1, 1, 0, 0, 0, 0], + } + ) + res = LPDiD(pre_window=2, post_window=1).fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert set(res.pooled["window"]) == {"pre", "post"} + + def test_rejects_pooled_horizon_outside_supported_window(self): + df = pd.DataFrame( + { + "unit": [1, 1, 1, 2, 2, 2], + "time": [0, 1, 2, 0, 1, 2], + "y": [1, 3, 5, 1, 1, 1], + "treat": [0, 1, 1, 0, 0, 0], + } + ) + with pytest.raises(ValueError, match="outside the supported pre window"): + LPDiD(pre_window=2, post_window=1).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", pre_pooled=(-3, -1) + ) + + def test_results_to_dataframe_and_repr(self): + df = pd.DataFrame({"horizon": [0], "coefficient": [1.0], "se": [0.1]}) + results = LPDiDResults( + event_study=df, + pooled=None, + n_obs=10, + n_treated_units=4, + n_control_units=6, + pre_window=2, + post_window=0, + control_group="clean", + reweight=False, + no_composition=False, + pmd=None, + ) + assert results.to_dataframe(level="event").equals(df) + with pytest.raises(ValueError, match="not computed"): + results.to_dataframe(level="pooled") + with pytest.raises(ValueError, match="level must be"): + results.to_dataframe(level="bad") + assert "LPDiDResults" in repr(results) + + def test_results_summary_and_to_dict_run(self): + df = make_lpdid_panel(cohorts=(4,), n_per_cohort=15, n_never=15, n_periods=8, seed=11) + res = LPDiD(pre_window=2, post_window=2).fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert isinstance(res.summary(), str) + assert "Local Projections DiD" in res.summary() + d = res.to_dict() + assert d["inference_method"] == "cluster_robust" + assert d["cluster_name"] == "unit" + assert np.isfinite(d["att"]) + + def test_clears_fitted_state_after_failed_refit(self): + good = pd.DataFrame( + { + "unit": [1, 1, 1, 2, 2, 2], + "time": [0, 1, 2, 0, 1, 2], + "y": [1, 3, 5, 1, 1, 1], + "treat": [0, 1, 1, 0, 0, 0], + } + ) + bad = good.copy() + bad["treat"] = [0, "bad", 1, 0, 0, 0] + est = LPDiD(pre_window=2, post_window=1) + est.fit(good, outcome="y", unit="unit", time="time", treatment="treat", only_event=True) + with pytest.raises(ValueError, match="binary numeric"): + est.fit(bad, outcome="y", unit="unit", time="time", treatment="treat") + assert est.is_fitted_ is False + assert est.results_ is None + + def test_fit_is_idempotent_on_config(self): + # Re-fitting with the same config must give identical results (fit + # writes no config-mutating state). + df = make_lpdid_panel(cohorts=(4,), n_per_cohort=12, n_never=12, n_periods=8, seed=7) + est = LPDiD(pre_window=2, post_window=2) + r1 = est.fit(df, outcome="y", unit="unit", time="time", treatment="treat") + params_after = est.get_params() + r2 = est.fit(df, outcome="y", unit="unit", time="time", treatment="treat") + assert est.get_params() == params_after + pd.testing.assert_frame_equal(r1.event_study, r2.event_study) + + +# =========================================================================== +# Analytical closed-form coefficients (true effect derivable by hand) +# =========================================================================== +class TestLPDiDFormula: + def test_event_estimation_controls_for_calendar_time(self): + df = pd.DataFrame( + { + "unit": ["c1"] * 3 + + ["c2"] * 3 + + ["t1"] * 3 + + ["t2a"] * 3 + + ["t2b"] * 3 + + ["t2c"] * 3, + "time": [0, 1, 2] * 6, + "y": [0, 0, 100, 0, 0, 100, 0, 5, 105, 0, 0, 105, 0, 0, 105, 0, 0, 105], + "treat": [0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1], + } + ) + res = LPDiD(pre_window=2, post_window=0).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + assert _event_coef(res, 0) == pytest.approx(5.0, abs=1e-8) + + def test_pmd_max_uses_mean_of_all_available_pre_periods(self): + df = _make_linear_panel( + [ + { + "unit": "t1", + "times": [0, 1, 2, 3], + "treat_start": 3, + "y": {0: 0, 1: 2, 2: 4, 3: 6}, + }, + { + "unit": "t2", + "times": [0, 1, 2, 3], + "treat_start": 3, + "y": {0: 0, 1: 2, 2: 4, 3: 6}, + }, + { + "unit": "c1", + "times": [0, 1, 2, 3], + "treat_start": 99, + "y": {0: 0, 1: 1, 2: 2, 3: 3}, + }, + { + "unit": "c2", + "times": [0, 1, 2, 3], + "treat_start": 99, + "y": {0: 0, 1: 1, 2: 2, 3: 3}, + }, + ] + ) + standard = LPDiD(pre_window=2, post_window=0).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + pmd_max = LPDiD(pre_window=2, post_window=0, pmd="max").fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + assert _event_coef(standard) == pytest.approx(1.0, abs=1e-8) + assert _event_coef(pmd_max) == pytest.approx(2.0, abs=1e-8) + + def test_pmd_integer_uses_last_k_pre_periods(self): + df = _make_linear_panel( + [ + { + "unit": "t1", + "times": [0, 1, 2, 3], + "treat_start": 3, + "y": {0: 0, 1: 2, 2: 4, 3: 6}, + }, + { + "unit": "t2", + "times": [0, 1, 2, 3], + "treat_start": 3, + "y": {0: 0, 1: 2, 2: 4, 3: 6}, + }, + { + "unit": "c1", + "times": [0, 1, 2, 3], + "treat_start": 99, + "y": {0: 0, 1: 1, 2: 2, 3: 3}, + }, + { + "unit": "c2", + "times": [0, 1, 2, 3], + "treat_start": 99, + "y": {0: 0, 1: 1, 2: 2, 3: 3}, + }, + ] + ) + pmd_two = LPDiD(pre_window=2, post_window=0, pmd=2).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + assert _event_coef(pmd_two) == pytest.approx(1.5, abs=1e-8) + + def test_covariates_and_absorb_remove_confounding_bias(self): + df = _make_linear_panel( + [ + { + "unit": "tA1", + "times": [0, 1, 2], + "treat_start": 1, + "y": {0: 0, 1: 6, 2: 12}, + "x1": 2, + "region": "A", + }, + { + "unit": "tA2", + "times": [0, 1, 2], + "treat_start": 1, + "y": {0: 0, 1: 5, 2: 10}, + "x1": 1, + "region": "A", + }, + { + "unit": "tB1", + "times": [0, 1, 2], + "treat_start": 1, + "y": {0: 0, 1: -1, 2: -2}, + "x1": 1, + "region": "B", + }, + { + "unit": "cA1", + "times": [0, 1, 2], + "treat_start": 99, + "y": {0: 0, 1: 4, 2: 8}, + "x1": 0, + "region": "A", + }, + { + "unit": "cB1", + "times": [0, 1, 2], + "treat_start": 99, + "y": {0: 0, 1: -1, 2: -2}, + "x1": 1, + "region": "B", + }, + { + "unit": "cB2", + "times": [0, 1, 2], + "treat_start": 99, + "y": {0: 0, 1: -2, 2: -4}, + "x1": 0, + "region": "B", + }, + ] + ) + uncontrolled = LPDiD(pre_window=2, post_window=0).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + controlled = LPDiD(pre_window=2, post_window=0).fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1"], + absorb=["region"], + only_event=True, + ) + assert _event_coef(uncontrolled) > 2.5 + assert abs(_event_coef(controlled)) < 1e-8 + + def test_ylags_remove_lagged_outcome_bias(self): + df = _make_linear_panel( + [ + { + "unit": "t1", + "times": [0, 1, 2, 3], + "treat_start": 2, + "y": {0: 8, 1: 4, 2: 2, 3: 1}, + }, + { + "unit": "t2", + "times": [0, 1, 2, 3], + "treat_start": 2, + "y": {0: 6, 1: 3, 2: 1.5, 3: 0.75}, + }, + { + "unit": "c1", + "times": [0, 1, 2, 3], + "treat_start": 99, + "y": {0: 4, 1: 2, 2: 1, 3: 0.5}, + }, + { + "unit": "c2", + "times": [0, 1, 2, 3], + "treat_start": 99, + "y": {0: 2, 1: 1, 2: 0.5, 3: 0.25}, + }, + ] + ) + uncontrolled = LPDiD(pre_window=2, post_window=0).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + with_ylag = LPDiD(pre_window=2, post_window=0).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", ylags=1, only_event=True + ) + assert _event_coef(uncontrolled) < -0.5 + assert abs(_event_coef(with_ylag)) < 1e-8 + + def test_dylags_remove_lagged_difference_bias(self): + df = _make_linear_panel( + [ + { + "unit": "t1", + "times": [0, 1, 2, 3], + "treat_start": 3, + "y": {0: 0, 1: 0, 2: 4, 3: 8}, + }, + { + "unit": "t2", + "times": [0, 1, 2, 3], + "treat_start": 3, + "y": {0: 0, 1: 0, 2: 3, 3: 6}, + }, + { + "unit": "c1", + "times": [0, 1, 2, 3], + "treat_start": 99, + "y": {0: 0, 1: 0, 2: 2, 3: 4}, + }, + { + "unit": "c2", + "times": [0, 1, 2, 3], + "treat_start": 99, + "y": {0: 0, 1: 0, 2: 1, 3: 2}, + }, + ] + ) + uncontrolled = LPDiD(pre_window=2, post_window=0).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + with_dylag = LPDiD(pre_window=2, post_window=0).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", dylags=1, only_event=True + ) + assert _event_coef(uncontrolled) > 1.5 + assert abs(_event_coef(with_dylag)) < 1e-8 + + +# =========================================================================== +# DGP recovery + methodological properties +# =========================================================================== +class TestLPDiDMethodology: + def test_single_cohort_exact_dynamic_recovery(self): + # Noiseless single-cohort DGP: LP-DiD recovers tau_h = 1 + 0.5h exactly, + # and pre-treatment placebos are zero. + df = make_lpdid_panel( + cohorts=(5,), + n_per_cohort=20, + n_never=20, + n_periods=12, + tau=lambda k: 1.0 + 0.5 * k, + unit_fe_sd=0.0, + error_sd=0.0, + heterogeneous=False, + seed=1, + ) + res = LPDiD(pre_window=3, post_window=4).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + for h in range(0, 5): + assert _event_coef(res, h) == pytest.approx(1.0 + 0.5 * h, abs=1e-8) + for h in (-3, -2): + assert _event_coef(res, h) == pytest.approx(0.0, abs=1e-8) + + def test_clean_control_rule_defeats_already_treated_contamination(self): + # An early cohort with a large ongoing effect would bias a naive + # long-difference that used it as a control; LP-DiD's clean-control + # restriction excludes it and recovers the late cohort's true effect. + df = make_lpdid_panel( + cohorts=(2, 8), + n_per_cohort=25, + n_never=25, + n_periods=12, + tau=lambda k: 5.0, + unit_fe_sd=0.0, + error_sd=0.0, + heterogeneous=False, + seed=2, + ) + res = LPDiD(pre_window=2, post_window=1).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + # Homogeneous effect 5.0 across cohorts -> h=0 effect is exactly 5.0. + assert _event_coef(res, 0) == pytest.approx(5.0, abs=1e-8) + + def test_pooled_post_equals_mean_of_event_study(self): + df = make_lpdid_panel( + cohorts=(5,), + n_per_cohort=20, + n_never=20, + n_periods=12, + tau=lambda k: 1.0 + 0.5 * k, + unit_fe_sd=0.0, + error_sd=0.0, + heterogeneous=False, + seed=3, + ) + res = LPDiD(pre_window=2, post_window=3).fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + es = res.event_study + post_mean = es.loc[es["horizon"].between(0, 3), "coefficient"].mean() + post_pooled = res.pooled.loc[res.pooled["window"] == "post", "coefficient"].iloc[0] + assert post_pooled == pytest.approx(post_mean, abs=1e-8) + + def test_variance_and_equal_weighting_differ_under_heterogeneity(self): + # With heterogeneous cohort effects and unequal cohort sizes, the + # variance-weighted (default) and equally-weighted (reweight) ATTs differ. + df = make_lpdid_panel( + cohorts=(4, 8), n_per_cohort=20, n_never=20, n_periods=12, heterogeneous=True, seed=4 + ) + vw = LPDiD(pre_window=2, post_window=2).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + ew = LPDiD(pre_window=2, post_window=2, reweight=True).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + assert _event_coef(vw, 2) != pytest.approx(_event_coef(ew, 2), abs=1e-6) + + def test_cluster_se_finite_and_nan_consistent(self): + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=20, n_never=20, n_periods=10, seed=5) + res = LPDiD(pre_window=2, post_window=2).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + row0 = _event_row(res, 0) + assert np.isfinite(row0["se"]) and row0["se"] > 0 + assert np.isfinite(row0["t_stat"]) and np.isfinite(row0["conf_low"]) + + def test_single_cluster_yields_nan_inference(self): + # All obs in one cluster -> cluster vcov undefined -> NaN-consistent. + df = make_lpdid_panel(cohorts=(2,), n_per_cohort=4, n_never=4, n_periods=4, seed=6) + df["only"] = 1 + res = LPDiD(pre_window=1, post_window=1, cluster="only").fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + row0 = _event_row(res, 0) + for col in ("se", "t_stat", "p_value", "conf_low", "conf_high"): + assert pd.isna(row0[col]) + + +# =========================================================================== +# Cross-estimator equivalences proved in Dube et al. (2025) +# =========================================================================== +class TestLPDiDCrossEstimator: + def test_ce1_2x2_h0_equals_first_difference_did(self): + # 2 periods, treated-at-1 vs never: LP-DiD h=0 == the closed-form 2x2 + # DiD = mean(dy | treated) - mean(dy | control) [paper Appendix A.1]. + df = make_lpdid_panel( + cohorts=(1,), + n_per_cohort=60, + n_never=60, + n_periods=2, + tau=lambda k: 3.0, + heterogeneous=False, + seed=10, + ) + lp = LPDiD(pre_window=1, post_window=0).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + wide = df.pivot(index="unit", columns="time", values="y") + dy = wide[1] - wide[0] + treated_units = df.loc[df["first_treat"] == 1, "unit"].unique() + is_t = wide.index.isin(treated_units) + did_2x2 = dy[is_t].mean() - dy[~is_t].mean() + assert _event_coef(lp, 0) == pytest.approx(did_2x2, abs=1e-9) + + def test_ce2_reweighted_equals_callaway_santanna(self): + # Reweighted LP-DiD == Callaway-Sant'Anna event-study effects + # [paper Section 3.7]. Match CS to LP-DiD's not-yet-treated controls. + df = make_lpdid_panel(cohorts=(5, 8), n_per_cohort=25, n_never=30, n_periods=12, seed=21) + lp = LPDiD(pre_window=3, post_window=3, reweight=True).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + cs = CallawaySantAnna(control_group="not_yet_treated").fit( + df, + outcome="y", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + for h in range(0, 4): + assert _event_coef(lp, h) == pytest.approx( + cs.event_study_effects[h]["effect"], abs=1e-9 + ) + + def test_ce3_pmd_single_cohort_equals_bjs_imputation(self): + # PMD LP-DiD (k=t-1, "max") with a SINGLE treated cohort == BJS + # imputation [paper Section 3.4, footnotes 10-11; only single-cohort + # is an exact equality, multi-cohort is only "very similar"]. + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=40, n_never=40, n_periods=12, seed=31) + assert df.loc[np.isfinite(df["first_treat"]), "first_treat"].nunique() == 1 + lp = LPDiD(pre_window=3, post_window=3, pmd="max").fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + bjs = ImputationDiD().fit( + df, + outcome="y", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + for h in range(0, 4): + assert _event_coef(lp, h) == pytest.approx( + bjs.event_study_effects[h]["effect"], abs=1e-6 + ) + + # CE-4 (variance-weighted LP-DiD == stacked regression) is intentionally + # omitted: the paper's equivalence is to the Cengiz et al. (2019) stacking + # scheme, whereas diff-diff's `StackedDiD` implements Wing, Freedman & + # Hollingsworth (2024) with Q-weights — a different (corrected) scheme that + # does not numerically coincide. See REGISTRY `## LPDiD`. + + +# =========================================================================== +# Edge cases +# =========================================================================== +class TestLPDiDEdgeCases: + def test_rejects_non_absorbing_treatment(self): + # Treatment that turns off must raise (absorbing-path scope). + df = pd.DataFrame( + { + "unit": [1, 1, 1, 2, 2, 2], + "time": [0, 1, 2, 0, 1, 2], + "y": [1.0, 2, 3, 1, 1, 1], + "treat": [0, 1, 0, 0, 0, 0], + } + ) + with pytest.raises(ValueError, match="absorbing"): + LPDiD(pre_window=1, post_window=1).fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + + def test_no_composition_drops_controls(self): + df = _make_linear_panel( + [ + { + "unit": "t1", + "times": [0, 1, 2, 3], + "treat_start": 1, + "y": {0: 0, 1: 1, 2: 2, 3: 3}, + }, + { + "unit": "t2", + "times": [0, 1, 2, 3], + "treat_start": 2, + "y": {0: 0, 1: 0, 2: 2, 3: 4}, + }, + { + "unit": "c1", + "times": [0, 1, 2, 3], + "treat_start": 99, + "y": {0: 0, 1: 0, 2: 0, 3: 0}, + }, + ] + ) + unrestricted = LPDiD(pre_window=2, post_window=2).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + common = LPDiD(pre_window=2, post_window=2, no_composition=True).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + u0 = unrestricted.event_study.loc[unrestricted.event_study["horizon"] == 0, "n_obs"].iloc[0] + c0 = common.event_study.loc[common.event_study["horizon"] == 0, "n_obs"].iloc[0] + assert c0 < u0 + + +class TestLPDiDUnbalanced: + """Unbalanced-panel correctness (review round 1: reweight denominators, RA + identification, no_composition all computed from the realized sample).""" + + def _unbalanced(self): + df = make_lpdid_panel(cohorts=(5, 8), n_per_cohort=25, n_never=30, n_periods=12, seed=21) + drop = ( + (df["first_treat"] == np.inf) & (df["time"].isin([9, 10, 11])) & (df["unit"] % 3 == 0) + ) + return df.loc[~drop].reset_index(drop=True) + + def test_reweight_matches_cs_on_unbalanced(self): + # Equal-weighting denominators must come from the realized (post-drop) + # sample, else the Callaway-Sant'Anna equivalence breaks on missing rows. + ub = self._unbalanced() + lp = LPDiD(pre_window=3, post_window=3, reweight=True).fit( + ub, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + cs = CallawaySantAnna(control_group="not_yet_treated").fit( + ub, + outcome="y", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + for h in range(0, 4): + assert _event_coef(lp, h) == pytest.approx( + cs.event_study_effects[h]["effect"], abs=1e-9 + ) + + def test_ra_nan_consistent_on_uncontrolled_event_time(self): + # An event time with treated units but no clean control is unidentified: + # the RA path drops those treated (with a warning) and returns jointly + # NaN inference, never a coef=NaN / se=0 mismatch. + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=10, n_never=10, n_periods=10, seed=3) + df["x"] = np.arange(len(df)) % 3 + df = df.loc[~((df["first_treat"] == np.inf) & (df["time"] == 5))].reset_index(drop=True) + with pytest.warns(UserWarning, match="no clean control"): + res = LPDiD(pre_window=2, post_window=2, reweight=True).fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x"], + only_event=True, + ) + row0 = _event_row(res, 0) + assert pd.isna(row0["coefficient"]) and pd.isna(row0["se"]) + for col in ("t_stat", "p_value", "conf_low", "conf_high"): + assert pd.isna(row0[col]) + + def test_no_composition_holds_post_horizons_fixed(self): + # no_composition fixes the POST-treatment composition: every post horizon + # shares the same realized n_obs even on an unbalanced panel. + ub = self._unbalanced() + res = LPDiD(pre_window=3, post_window=3, no_composition=True).fit( + ub, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + post = res.event_study.loc[res.event_study["horizon"].between(0, 3), "n_obs"] + assert post.nunique() == 1 + + def test_pmd_retains_treated_missing_exact_t_minus_1(self): + # Under PMD the long difference uses the premean baseline, so a treated + # observation missing the exact t-1 outcome (but with earlier pre-data) + # must NOT be dropped. + df = make_lpdid_panel( + cohorts=(5,), + n_per_cohort=12, + n_never=12, + n_periods=10, + unit_fe_sd=0.0, + error_sd=0.0, + heterogeneous=False, + tau=lambda k: 2.0, + seed=1, + ) + treated = df.loc[df["first_treat"] == 5, "unit"].unique() + df = df.loc[~(df["unit"].isin(treated[:6]) & (df["time"] == 4))].reset_index(drop=True) + std = LPDiD(pre_window=2, post_window=2).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + pmd = LPDiD(pre_window=2, post_window=2, pmd="max").fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + std_n = std.event_study.loc[std.event_study["horizon"] == 0, "n_obs"].iloc[0] + pmd_n = pmd.event_study.loc[pmd.event_study["horizon"] == 0, "n_obs"].iloc[0] + assert pmd_n > std_n # PMD keeps missing-t-1 treated that first-lag drops + assert np.isfinite(_event_coef(pmd, 0)) + + def test_pre_window_too_small_for_pooled_raises(self): + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=8, n_never=8, n_periods=10, seed=2) + with pytest.raises(ValueError, match="pooled pre window is empty"): + LPDiD(pre_window=1, post_window=1).fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + # only_event=True does not need a pooled pre window + r = LPDiD(pre_window=1, post_window=1).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + assert list(r.event_study["horizon"]) == [-1, 0, 1] + + def test_pre_pooled_rejects_reference_horizon(self): + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=8, n_never=8, n_periods=10, seed=3) + with pytest.raises(ValueError, match="outside the supported pre window"): + LPDiD(pre_window=3, post_window=1).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", pre_pooled=(-2, -1) + ) + + def test_no_composition_post_fixed_under_nonmonotone_missingness(self): + # Missing an INTERMEDIATE target (t+1) while the max target (t+2) is + # observed must still give identical post-horizon n_obs (the common mask + # requires every post target, not just the maximum horizon). + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=20, n_never=20, n_periods=12, seed=7) + drop = (df["first_treat"] == np.inf) & (df["time"] == 6) & (df["unit"] % 4 == 0) + ub = df.loc[~drop].reset_index(drop=True) + res = LPDiD(pre_window=2, post_window=2, no_composition=True).fit( + ub, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + post = res.event_study.loc[res.event_study["horizon"].between(0, 2), "n_obs"] + assert post.nunique() == 1 + + def test_rank_deficient_action_propagates_to_results(self): + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=10, n_never=10, n_periods=8, seed=8) + res = LPDiD(pre_window=2, post_window=2, rank_deficient_action="silent").fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert res.rank_deficient_action == "silent" + assert res.to_dict()["rank_deficient_action"] == "silent" + + def test_ra_path_reports_if_cluster_variance_label(self): + # The regression-adjustment covariate path uses an influence-function + # cluster variance, not an OLS CR1 sandwich: the results metadata and the + # summary label must say so (not "hc1" / "CR1"). + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=12, n_never=12, n_periods=8, seed=9) + df["x"] = np.arange(len(df)) % 3 + ra = LPDiD(pre_window=2, post_window=2, reweight=True).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", covariates=["x"] + ) + assert ra.to_dict()["vcov_type"] == "if_cluster" + assert "Influence-function" in ra.summary() + # the default (non-RA) path stays hc1 / CR1 + base = LPDiD(pre_window=2, post_window=2).fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert base.to_dict()["vcov_type"] == "hc1" + assert "CR1 cluster-robust" in base.summary() + + def test_pmd_max_excludes_present_but_nan_pretreatment(self): + # A present-but-NaN pretreatment outcome must not deflate the premean + # baseline: the denominator counts non-missing prior outcomes, not rows. + rows = [] + for u in (1, 2): + for t, y in enumerate([10.0, np.nan, 30.0, 40.0, 50.0]): + rows.append({"unit": u, "time": t, "y": y, "treat": int(t >= 3)}) + for u, base in {3: 0, 4: 1, 5: 2, 6: 3}.items(): + for t in range(5): + rows.append({"unit": u, "time": t, "y": float(base + t), "treat": 0}) + df = pd.DataFrame(rows) + est = LPDiD(pre_window=2, post_window=1, pmd="max") + est.fit(df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True) + panel = est._prepare_panel( + df, outcome="y", unit="unit", time="time", treatment="treat", cluster="unit" + ) + baseline = panel.loc[(panel["unit"] == 1) & (panel["time"] == 3), "_pmd_all_baseline"].iloc[ + 0 + ] + assert baseline == pytest.approx(20.0, abs=1e-9) # mean(10, 30); NaN at t=1 excluded + + def test_pmd_max_retains_obs_with_missing_current_outcome(self): + # The premean baseline uses only PRIOR outcomes, never the base row's own + # y_t, so a treated entry whose current-period outcome is missing (but + # priors and the h>0 target are observed) stays identified -- it must not + # be silently dropped (which would zero out the treatment variation). + rows = [] + for u in (1, 2): # treated, entry t=3, entry-period outcome y_3 missing + for t, y in enumerate([10.0, 11.0, 12.0, np.nan, 14.0, 15.0]): + rows.append({"unit": u, "time": t, "y": y, "treat": int(t >= 3)}) + for u in (3, 4, 5, 6): + for t in range(6): + rows.append({"unit": u, "time": t, "y": float(u + t), "treat": 0}) + df = pd.DataFrame(rows) + res = LPDiD(pre_window=2, post_window=2, pmd="max").fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + h1 = res.event_study[res.event_study["horizon"] == 1].iloc[0] + assert np.isfinite(h1["coefficient"]) # treated entry retained -> estimable + assert h1["n_obs"] > 0 + + def test_pooled_post_is_mean_of_event_study_balanced(self): + # Documented pooled estimand: the pooled-post ATT is the unit-equal-weighted + # mean of the per-horizon event-study coefficients. On a balanced single- + # cohort panel (consistent clean-control sample across horizons) this is + # exact; it pins the fixed-composition pooled definition (REGISTRY Note 6). + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=15, n_never=15, n_periods=12, seed=3) + res = LPDiD(pre_window=2, post_window=3).fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + es_post = res.event_study[res.event_study["horizon"] >= 0]["coefficient"] + pooled_post = res.pooled.loc[res.pooled["window"] == "post", "coefficient"].iloc[0] + assert pooled_post == pytest.approx(es_post.mean(), rel=1e-6) + + def test_ra_absorb_overlap_drops_unsupported_treated(self): + # RA path (reweight=True) with an absorbed factor: a treated unit whose + # absorbed level has NO clean-control support is dropped with a warning, + # not extrapolated through a non-identified all-zero dummy. + rows = [] + for u, region in [(1, "A"), (2, "A"), (7, "B"), (8, "B")]: # region B = treated-only + for t in range(5): + rows.append( + { + "unit": u, + "time": t, + "y": float(t + (2 if t >= 2 else 0)), + "treat": int(t >= 2), + "region": region, + } + ) + for u in (3, 4, 5, 6): # clean controls only in region A + for t in range(5): + rows.append({"unit": u, "time": t, "y": float(t), "treat": 0, "region": "A"}) + df = pd.DataFrame(rows) + with pytest.warns(UserWarning, match="absorbed-factor level absent"): + res = LPDiD(pre_window=1, post_window=1, reweight=True).fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + absorb=["region"], + only_event=True, + ) + # region-A treated remain identified -> finite estimate (no extrapolation) + h0 = res.event_study.loc[res.event_study["horizon"] == 0, "coefficient"].iloc[0] + assert np.isfinite(h0) + + @staticmethod + def _unsupported_event_time_panel(): + # cohorts {2, 4}, NO never-treated: cohort B (entry t=4) has no clean control + # (every other unit treated since t=2). Cohort B gets a HUGE effect (99) so a + # spurious cross-time identification would visibly contaminate the estimate; + # cohort A's true effect is +3. + rows = [] + for u in (1, 2): + for t in range(6): + rows.append( + { + "unit": u, + "time": t, + "y": float(t + (3 if t >= 2 else 0)), + "treat": int(t >= 2), + } + ) + for u in (3, 4): + for t in range(6): + rows.append( + { + "unit": u, + "time": t, + "y": float(t + (99 if t >= 4 else 0)), + "treat": int(t >= 4), + } + ) + return pd.DataFrame(rows) + + def test_default_path_drops_unsupported_event_time(self): + # Default (reweight=False) OLS path: a treated event-time with NO clean + # control must be dropped, not identified via a collinear time-FE drop -- + # even under rank_deficient_action="silent". + df = self._unsupported_event_time_panel() + with pytest.warns(UserWarning, match="no clean control"): + res = LPDiD( + pre_window=1, post_window=1, control_group="clean", rank_deficient_action="silent" + ).fit(df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True) + # cohort A (+3) only; the unidentified cohort B (+99) is dropped, so the + # estimate is NOT pulled toward 99. + h0 = res.event_study.loc[res.event_study["horizon"] == 0, "coefficient"].iloc[0] + assert h0 == pytest.approx(3.0, abs=1e-6) + + def test_pooled_path_drops_unsupported_event_time(self): + # The same enforcement must apply on the pooled path (pre_window>=2 so the + # pooled-pre window is non-empty). + df = self._unsupported_event_time_panel() + with pytest.warns(UserWarning, match="no clean control"): + res = LPDiD(pre_window=2, post_window=1, control_group="clean").fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + pooled_post = res.pooled.loc[res.pooled["window"] == "post", "coefficient"].iloc[0] + assert pooled_post == pytest.approx(3.0, abs=1e-6) # cohort A only, not pulled to 99 + + def test_ra_path_drops_redundant_covariate_without_nan(self): + # RA path: a redundant nuisance column (duplicate / constant covariate) is + # dropped by solve_ols (its coefficient becomes NaN); the ATT must stay + # finite and EQUAL the no-redundant-column fit -- the dropped coefficient + # acts as 0, not NaN (its effect is absorbed by the retained columns). + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=15, n_never=15, n_periods=10, seed=3) + df["x"] = np.arange(len(df)) % 4 + df["xdup"] = df["x"] # exact duplicate -> redundant + df["xc"] = 7.0 # constant -> collinear with the intercept + kw = dict(outcome="y", unit="unit", time="time", treatment="treat", only_event=True) + ref = LPDiD(pre_window=2, post_window=2, reweight=True, rank_deficient_action="silent").fit( + df, covariates=["x"], **kw + ) + red = LPDiD(pre_window=2, post_window=2, reweight=True, rank_deficient_action="silent").fit( + df, covariates=["x", "xdup", "xc"], **kw + ) + ref_h0 = ref.event_study.loc[ref.event_study["horizon"] == 0, "coefficient"].iloc[0] + red_h0 = red.event_study.loc[red.event_study["horizon"] == 0, "coefficient"].iloc[0] + assert np.isfinite(red_h0) + assert red_h0 == pytest.approx(ref_h0, abs=1e-9) + # rank_deficient_action="error" still raises on the redundant nuisance design + with pytest.raises(ValueError): + LPDiD(pre_window=2, post_window=2, reweight=True, rank_deficient_action="error").fit( + df, covariates=["x", "xdup"], **kw + ) + + def test_ylags_dylags_trigger_direct_inclusion_warning(self): + # Outcome/first-difference lags are direct-included controls under + # reweight=False, so they fire the same homogeneity warning as covariates. + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=10, n_never=10, n_periods=10, seed=4) + with pytest.warns(UserWarning, match="covariate-style controls"): + LPDiD(pre_window=2, post_window=2).fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + ylags=1, + only_event=True, + ) + with pytest.warns(UserWarning, match="covariate-style controls"): + LPDiD(pre_window=2, post_window=2).fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + dylags=1, + only_event=True, + ) + + def test_no_composition_post_fixed_with_present_but_nan_outcome(self): + # Fixed composition must hold even when a target row exists but its + # outcome is NaN (value-based availability, not row existence). + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=20, n_never=20, n_periods=12, seed=7) + df.loc[(df["first_treat"] == np.inf) & (df["time"] == 6) & (df["unit"] % 4 == 0), "y"] = ( + np.nan + ) + res = LPDiD(pre_window=2, post_window=2, no_composition=True).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", only_event=True + ) + post = res.event_study.loc[res.event_study["horizon"].between(0, 2), "n_obs"] + assert post.nunique() == 1 + + def test_pooled_rejects_bool_window(self): + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=8, n_never=8, n_periods=10, seed=2) + with pytest.raises(ValueError, match="not bool"): + LPDiD(pre_window=3, post_window=1).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", pre_pooled=True + ) + + def test_ylags_dylags_recorded_in_results(self): + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=10, n_never=10, n_periods=10, seed=4) + res = LPDiD(pre_window=2, post_window=2, reweight=True).fit( + df, outcome="y", unit="unit", time="time", treatment="treat", ylags=2, dylags=1 + ) + assert res.ylags == 2 and res.dylags == 1 + d = res.to_dict() + assert d["ylags"] == 2 and d["dylags"] == 1 + + def test_headline_n_clusters_matches_pooled_post(self): + # The results-level n_clusters (the summary "G") is the realized cluster + # count of the pooled-post headline row, not the full-panel count. + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=15, n_never=15, n_periods=10, seed=5) + res = LPDiD(pre_window=2, post_window=2).fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + post_g = res.pooled.loc[res.pooled["window"] == "post", "n_clusters"].iloc[0] + assert res.n_clusters == int(post_g) + + +class TestLPDiDInteriorGaps: + """Interior per-unit time gaps must be handled by calendar-correct feature + construction (reindex each unit to its interior grid), NOT by the previous- + observed-row semantics of ``shift()``/``diff()``/``rolling()``.""" + + @staticmethod + def _gapped_panel(): + # unit 1 = never-treated control with an INTERIOR gap at t=2 (observed 0,1,3,4); + # units 2-4 = regular controls; units 5-6 = treated with a +2 jump from t=2. + rows = [(1, 0, 10.0, 0), (1, 1, 11.0, 0), (1, 3, 13.0, 0), (1, 4, 14.0, 0)] + for u in (2, 3, 4): + for t in range(5): + rows.append((u, t, float(u + t), 0)) + for u in (5, 6): + for t in range(5): + rows.append((u, t, float(10 + u + t + (2.0 if t >= 2 else 0.0)), int(t >= 2))) + return pd.DataFrame(rows, columns=["unit", "time", "y", "treat"]) + + def test_lag_features_are_calendar_correct_across_gap(self): + df = self._gapped_panel() + est = LPDiD(pre_window=1, post_window=2) + panel = est._prepare_panel( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + cluster="unit", + ylags=1, + dylags=1, + ) + u1 = panel[panel["unit"] == 1].set_index("time") + # calendar t-1 = 2 is a gap -> NaN, NOT the previous-observed value at t=1 + assert np.isnan(u1.loc[3, "_y_lag_1"]) + assert np.isnan(u1.loc[3, "_dy_current"]) + # rows whose calendar t-1 IS observed are unaffected + assert u1.loc[4, "_y_lag_1"] == 13.0 + assert u1.loc[1, "_y_lag_1"] == 10.0 + + def test_pmd_int_baseline_fails_closed_across_gap(self): + df = self._gapped_panel() + est = LPDiD(pre_window=1, post_window=2, pmd=2) + panel = est._prepare_panel( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + cluster="unit", + ) + u1 = panel[panel["unit"] == 1].set_index("time") + # the 2-period premean window at t=4 spans the t=2 gap -> NaN (not last-2-observed) + assert np.isnan(u1.loc[4, "_pmd_k_baseline"]) + + def test_gap_rows_dropped_and_no_nan_cluster(self): + df = self._gapped_panel() + est = LPDiD(pre_window=1, post_window=2) + panel = est._prepare_panel( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + cluster="unit", + ) + # only the 4 observed rows of unit 1 survive; the synthetic t=2 gap row is dropped + assert int((panel["unit"] == 1).sum()) == 4 + assert len(panel) == len(df) + # no phantom NaN-cluster row can reach the variance computation + assert not panel["_cluster"].isna().any() + + def test_fit_on_gapped_panel_recovers_effect_with_finite_clusters(self): + df = self._gapped_panel() + res = LPDiD(pre_window=1, post_window=2).fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + only_event=True, + ) + post = res.event_study[res.event_study["horizon"] >= 0] + assert np.isfinite(post["coefficient"]).all() + assert (post["n_clusters"].dropna() >= 1).all() # no NaN-cluster inflation + assert np.allclose(post["coefficient"], 2.0, atol=1e-6) # +2 recovered despite the gap + + def test_regular_panel_takes_early_out(self): + # a contiguous panel returns exactly its input rows (no reindex, no fabricated rows) + df = make_lpdid_panel(cohorts=(5,), n_per_cohort=8, n_never=8, n_periods=10, seed=1) + est = LPDiD(pre_window=2, post_window=2) + panel = est._prepare_panel( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + cluster="unit", + ) + assert len(panel) == len(df) + + def test_entry_is_first_observed_treated_with_pre_onset_gap(self): + # unit 1 observed 0,1,4,5 (gap at 2-3), treated from t=4 -> entry = first OBSERVED + # treated period = 4 (the unobservable pre-onset gap cannot move it earlier). + rows = [(1, 0, 1.0, 0), (1, 1, 1.0, 0), (1, 4, 5.0, 1), (1, 5, 6.0, 1)] + for u in (2, 3): + for t in range(6): + rows.append((u, t, float(t), 0)) + df = pd.DataFrame(rows, columns=["unit", "time", "y", "treat"]) + est = LPDiD(pre_window=1, post_window=1) + panel = est._prepare_panel( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + cluster="unit", + ) + assert panel.loc[panel["unit"] == 1, "_first_treat"].iloc[0] == 4.0 + ent = panel[(panel["unit"] == 1) & (panel["_entry"] == 1.0)] + assert ent["time"].tolist() == [4] + + def test_rejects_fractional_time_labels(self): + df = pd.DataFrame( + { + "unit": [1, 1, 2, 2], + "time": [0.5, 1.5, 0.5, 1.5], + "y": [1.0, 2, 1, 1], + "treat": [0, 1, 0, 0], + } + ) + with pytest.raises(ValueError, match="integer-valued"): + LPDiD().fit(df, outcome="y", unit="unit", time="time", treatment="treat")