diff --git a/CHANGELOG.md b/CHANGELOG.md index e8b69541..5a24fe8a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -60,6 +60,128 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 route on bootstrapped CS fits instead of the deprecated fit-time kwarg. The sibling estimators' (EfficientDiD/ImputationDiD/TwoStageDiD/ContinuousDiD) bootstrapped recompute gates are unchanged. +- **`LWDiD` (Lee & Wooldridge 2025, 2026 rolling-transformation DiD).** Unit-specific + demean/detrend (plus quarterly `demeanq`/`detrendq`) converts panel data to + cross-sectional transformed outcomes; supports common timing and staggered + adoption with never-treated / not-yet-treated controls, + `estimation_method` in `'reg'`/`'ipw'`/`'dr'`/`'psm'`, analytical + (`vcov_type` in `'classical'`/`'hc1'`/`'hc2'`/`'hc3'`) and cluster-robust + (constructor `cluster=`) inference, multiplier bootstrap, wild cluster + bootstrap, and randomization inference. Common-timing fits expose the same + post-fit event-study surface as staggered ones — + `results.aggregate('event_study')` returns per-period effects keyed by + event time relative to treatment onset, so no separate per-period fit + option exists. + +### Changed +- **`LWDiD` API canonicalized to the v4 vocabulary agreed in PR #588's review** + (renames relative to the PR's earlier review rounds; nothing here was ever + released): `estimator=` -> `estimation_method=` with values `'ra'` -> `'reg'` + and `'ipwra'` -> `'dr'`; `vce=` -> `vcov_type=` with no `'cluster'` value — + cluster-robust (CR1) inference activates via the constructor `cluster=` + column instead; `bootstrap_seed=` -> `seed=` (default `None`); + `trim_threshold=` -> `pscore_trim=`. `vcov_type='hc3'` is computed through + the shared `diff_diff.linalg` HC machinery used by the other estimators, + and the `'hc0'`/`'hc4'` values are removed from the surface. + Unit-constancy validation is centralized and applies uniformly to + covariates and the cluster column across all estimation paths. + +### Removed +- **`LWDiD` pre-v4 review-round surface** (never released): the `LW` alias, + the functional `lwdid()` wrapper, the `lwdid_trend_diagnostics` module + (including `recommend_transformation`), and the `overall_att` / + `period_effects` result fields together with the `period_specific` fit + option — per-period effects are served by the post-fit + `results.aggregate('event_study')` surface instead. + +### Fixed +- **`LWDiD` maintainer fix wave** (post-acceptance validation campaign: 43 + execution-verified findings, all resolved): + - Estimand: the `tau_omega` composite is complete-case with FIXED cohort + weights (treated units without a finite own-cohort post average and + controls not observing every surviving cohort's post window are + dropped with warnings and counters; the pre-fix code zero-filled + missing control entries and silently reweighted the treated side); + with any drops, `.att` is the influence-weighted cohort-mass point on + every variance route and the composite is exposed as + `att_tau_omega_complete_case`. `demeanq`/`detrendq` overall ATTs now + aggregate seasonal cohort ATTs (the composite silently substituted + the non-seasonal transforms); the seasonal transforms fail closed on + insufficient pre-periods and unobserved seasons. + - Inference: one reference-distribution policy per surface (one-cell + aggregates use the cell's residual t; multi-cell unclustered stay + large-sample; clustered use contributing-cluster G-1); sub-samples + with <2 clusters fail closed; the common-timing bootstrap resamples + positionally (row order/index labels no longer corrupt SEs), honors + `cluster=` via whole-cluster resampling, and reports the df it used; + `wild_cluster_bootstrap` was rebuilt on the house WCR engine + (test-inversion CI, CR1 se, strict-exceedance p; the intercept-only + null model, ULP tie handling, and the G=2 zero-SE escape are gone; + API: `n_bootstrap`/`alpha`, result fields renamed). The result-level + `wild_cluster_bootstrap()`/`randomization_test()` methods now REPLAY + the fitted estimation sample and exact RA design (no data arguments; + RI recomputes the treated covariate mean per permutation; the + replayed statistic is asserted equal to `.att` before caching) - the + prior signatures accepted arbitrary arrays and could cache p-values + for a different estimand than the fitted ATT. The seeded bootstrap + draws identical streams for every `n_jobs`, resamples only units + surviving the transformation, and preserves fail-closed NaN inference + when fewer than 2 effective clusters survive; `aggregate(balance_e=)` + is rejected (was silently ignored); both RI/WCR fit through the + rank-aware shared solver (a duplicated treatment column previously + yielded a finite minimum-norm ATT). + - Contracts: `vcov_type` is restricted to values with real behavior + (`ipw`/`dr`/`psm` accept `hc1` only; `cluster=` composes only with + `hc1`; `psm`+`cluster` rejected); NaN covariates/clusters are rejected + explicitly; PSM calipers never average out-of-caliper controls; + cohort encodings are normalized once (`inf`/beyond-window recode with + warnings, negative cohorts raise, validator and `fit()` agree); + sensitivity results are NaN-honest (`baseline_pvalue`, NaN + `significant_05` for failed specs, full-frame pre-validation, unknown + kwargs raise); staggered sample metadata counts contributing units; + rank-deficient designs rebuild the influence function on kept columns. + - Tutorial: the contribution's `27_lwdid.ipynb` is WITHDRAWN (its + Walmart narrative was built on a fabricated common onset with jobs + figures inconsistent with the staggered estimate; a fresh notebook + is a tracked follow-up). + - Shared surfaces: `hc3` escapes closed across siblings (DiD/MP-DiD + `absorb=` now full-dummy-routes hc3 like hc2; TWFE no longer crashes + misleadingly; SpilloverDiD rejects hc3 at construction with its own + reason) plus a structural roster guard. `LWDiDResults.to_latex()` and + the `lwdid_exceptions` shim removed (unreleased API). +- **`LWDiD` review-round fixes** (staggered contract and inference tightenings): + - Staggered classical/HC SEs now come from the joint influence function + across cohort-time cells (the LW 2026 eq. 7.19 pooled-regression basis), + accounting for correlation among cohort effects that share controls + instead of assuming independence. + - On unbalanced panels the overall ATT point estimate is unified so a + variance selection never moves it (gated to `rolling` in + `'demean'`/`'detrend'` with `control_group='never_treated'`, + `estimation_method='reg'`, and no covariates). Superseded in detail by + the maintainer fix wave below: the composite `tau_omega` is now + complete-case with fixed cohort weights, reported as `.att` only when + no unit is dropped, and the quarterly variants now aggregate SEASONAL + cohort ATTs on every variance route. + - t-test degrees of freedom are computed from one design-based rule across + common-timing and staggered paths instead of two inconsistent ones. + - All-eventually-treated panels under `control_group='not_yet_treated'` + raise `ValueError` instead of silently truncating the sample; staggered + `covariates` must be unit-constant, time-varying columns raise + `ValueError`. + - Randomization inference uses the inclusive Phipson-Smyth rule + p = (c+1)/(B+1) and counts ties as extreme (`>=`), so p is never 0 and + an all-tie permutation distribution yields p = 1.0. + - `estimation_method='dr'` without covariates warns (`UserWarning`) that it + reduces to regression adjustment instead of silently doing so. + - `sensitivity_analysis` gains a `not_estimable` robustness level (with a + warning) when the ratio cannot be computed — including the zero-baseline + case — instead of mislabeling it. + - `to_dict()` output is fully JSON-native, including datetime/Period + cohort and time labels (ISO-8601 / period strings, NaT -> None). + - Staggered fits accept datetime64 and Period time scales, and panels + mixing the two time families are rejected in both directions with a + clear `ValueError`; cluster variable equal to the unit column no longer + raises a spurious column-lookup error. ## [3.9.1] - 2026-08-17 diff --git a/DEFERRED.md b/DEFERRED.md index 32ec54b2..94f8d415 100644 --- a/DEFERRED.md +++ b/DEFERRED.md @@ -20,6 +20,8 @@ provenance and AI-review deviation-documentation: a row here (or in | Issue | Location | PR | Priority | |-------|----------|----|----------| +| LWDiD PSM matching variance: implement the Abadie-Imbens (2006) matching variance (matched-control reuse + first-stage matching uncertainty) so `estimation_method='psm'` can report valid inference instead of the current fail-closed NaN tuple (point retained; naive var(diffs)/n was invalid under with-replacement reuse) | `diff_diff/lwdid.py` | #588 | Low | +| LWDiD cohort-relative sensitivity exclusions: `robustness_pre_periods` / `sensitivity_no_anticipation` currently reject multi-cohort staggered inputs because their exclusion windows are defined relative to the EARLIEST adoption (later cohorts' own pre periods fall inside the global post window and survive every restriction, mislabeling the specification). Supporting staggered inputs needs per-cohort window semantics (exclude the last k periods of each cohort's own `t < g` window before its transformation), which the current row-subset design cannot express — a per-cohort masking derivation + its aggregation contract | `diff_diff/lwdid_sensitivity.py` | #588 | Low | | HonestDiD non-chronological declared partitions (native `MultiPeriodDiDResults` route): the Rambachan-Roth restriction matrices are built POSITIONALLY over the concatenated declared pre/post lists assuming one chronological boundary, but the native route accepts non-suffix `post_periods` / non-last-pre references and returns bounds whose restriction system does not match the Registry equations (pre-existing; surfaced by the Phase 3(a) calendar-route review, which fails closed instead). Fix = transform the declared partition into boundary form where a valid mapping exists, else reject on the native route too - needs the restriction-geometry derivation. REGISTRY HonestDiD Note records the limitation. | `diff_diff/honest_did.py` | 3(a) | Medium | | `PlaceboTests` `boundary_gap` — a permutation randomization-inference margin (SE-audit item (b)); NOT computed anywhere in code today, so this is a new feature + result field, not a coverage lock. **User-locked 2026-07-09: defer until a derivation/paper source exists** — do not design or implement from scratch. | `tests/test_methodology_placebo.py`, `diff_diff/diagnostics.py` | SE-audit | Low | | TwoStageDiD honest/pretrends container admission DEFERRED (decision revised from "widen" during the 2(b) PR-3b plan review): analytical fits carry the joint Gardner-GMM event-study covariance (M-092), but the pre-period coefficients are stage-1 residual MEANS — the reference horizon is dropped from the no-intercept Stage-2 design and the zero anchor row appended mechanically — not contrasts against a reference period, while HonestDiD's Δ^RM/Δ^SD arithmetic hard-codes the `delta_0 = 0` normalization into its boundary/bridge constraints. Admission needs either a Stage-2 re-estimation with the reference horizon in the design or a derived residual-to-reference normalization mapping (+ its variance transform). Both consumers' TypeErrors state the deferral; see the REGISTRY TwoStageDiD Note (d). | `diff_diff/honest_did.py`, `diff_diff/pretrends.py`, `diff_diff/two_stage_aggregation.py` | 2(b) PR-3b | Low | diff --git a/README.md b/README.md index f5f7478c..556b8b5c 100644 --- a/README.md +++ b/README.md @@ -120,6 +120,7 @@ Full guide: `diff_diff.get_llm_guide("practitioner")`. - [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 or non-absorbing (reversible) treatment - [ChangesInChanges](https://diff-diff.readthedocs.io/en/stable/api/changes_in_changes.html) - Athey & Imbens (2006) nonlinear/distributional DiD for the 2x2 design: full counterfactual distribution and quantile treatment effects via CDF transformation, plus the QDiD comparison estimator via `method="qdid"`; bootstrap inference; R qte parity. Alias `CiC` +- [LWDiD](https://diff-diff.readthedocs.io/en/stable/api/lwdid.html) - Lee & Wooldridge (2025, 2026) rolling-transformation DiD: unit-specific demean/detrend converts panel to cross-section, staggered adoption, `estimation_method` in `reg`/`ipw`/`dr`/`psm` (the papers' RA/IPW/IPWRA plus propensity-score matching), exact small-N inference on the classical collapsed regression - [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/TODO.md b/TODO.md index 69d9482b..aaec09c1 100644 --- a/TODO.md +++ b/TODO.md @@ -21,6 +21,9 @@ Related tracking surfaces: | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| +| hc2/hc2_bm floor `1 - h_ii` at 1e-10 in the shared leverage meat, fabricating finite (if inflated) variances for leverage-one observations - hc3 now fails closed there (LWDiD fix wave) but the pre-existing hc2 family behavior is released surface; decide fail-closed vs keep-floor for hc2/hc2_bm | `diff_diff/linalg.py` | #588 | Quick | Low | +| Numeric between-period cohorts (e.g. `first_treat=4.5` with integer times) are rejected by LWDiD while CallawaySantAnna estimates them and LWDiD's own datetime/Period cohorts map to the next observed period — close the dtype asymmetry by adopting the next-observed-period mapping for numeric cohorts too (contract documented in REGISTRY cohort-encodings Note + `docs/api/lwdid.rst` Input Contract). Lands only after PR #588 merges | `diff_diff/lwdid.py` | #588 | Quick | Low | +| Implement the LW 2026 eq. 7.9/7.10 unit-average cohort estimand (regress per-unit post-average transformed outcomes on `[1, D_g]` vs never-treated) as an alternative to the documented cell-mass `cohort_effects` convention (REGISTRY within-cohort aggregation Note; the two differ on unbalanced panels, where cell-mass weights units by observed post periods). Needs the 7.10 regression + its covariance on the NT path. Lands only after PR #588 merges | `diff_diff/lwdid_staggered.py` | #588 | Quick | Low | | Expose cell-mass overall ATT (Stata `Post_avg` convention; = CS-simple on balanced panels) as an aggregate extra on LWDiD results — the fit's `.att` is the paper's `tau_omega` (cohort-mean-then-treated-weight, eq. 7.18); the authors' large-N display uses cell-mass weighting instead, and both are legitimate estimands (see the REGISTRY LWDiD Aggregation note). Lands only after PR #588 merges | `diff_diff/lwdid_results.py` | #588 | Quick | Low | | Post-fit `aggregate()` for the staggered DDD container: `StaggeredTripleDiffResults` carries no `AggregationMixin`, which is why the phase-3(b) merge had to carry fit-time `aggregate=`/`balance_e=` onto the surviving `TripleDifference` (rows M-140/M-141) as the ONE documented exception to the section-6 aggregate-postfit program. Porting the container onto the M-122 aggregation contract retires both rows; note the bootstrapped-fit recompute levels will need replay or a fail-closed relay — solved for CS via the BootstrapReplaySpec state replay (the container port can adopt the same mechanism); EfficientDiD/ImputationDiD/TwoStageDiD/ContinuousDiD still track theirs. Until it lands, the DDD docs deliberately keep teaching the fit-time kwarg (the canonical route there) | `diff_diff/staggered_triple_diff_results.py`, `diff_diff/aggregation.py`, `docs/api/triple_diff.rst`, `docs/tutorials/08_triple_diff.ipynb` | 3(b) | Heavy | Medium | | Staggered-DDD power support: `simulate_power`/`simulate_mde`/`simulate_sample_size` now REJECT a staggered-configured `TripleDifference` (both registered DDD generators emit 2x2x2 data and fit with `(group, partition, post)`, so a staggered config would be simulated under the wrong design). Support needs a staggered DDD DGP profile plus fit-kwargs builder, and a decision on whether the mode is selected by profile or by the estimator's own config | `diff_diff/power.py` | 3(b) | Mid | Low | @@ -66,6 +69,8 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| +| Author a fresh LWDiD tutorial notebook (the contribution's tutorial 27 was WITHDRAWN from PR #588 during CI review: its Walmart empirical narrative was built on a fabricated common onset with jobs translations that did not match the staggered estimate, and per-sentence repair failed across 5 review rounds). Build via the numbers-locked workflow: prototype in scripts, lock the numbers, then assemble and execute once; teach the staggered fit as the causal specification and keep any pooled contrast explicitly diagnostic | `docs/tutorials/` | #588 | Medium | +| LWDiD staggered fit recomputes cohort-wide work across surfaces: each cohort's rolling transformation is built once per `fit_staggered` cohort loop, but unit-level post summaries, control-eligibility sets, and the composite path's `ydot_by_cohort` are re-derived independently by the cell, aggregate, and `_composite_regression_aggregation` paths — runtime grows ~cohorts × panel size. Cache the per-cohort transformation + post summaries in one pass and reuse; add a many-cohort unbalanced-panel benchmark | `diff_diff/lwdid_staggered.py`, `diff_diff/lwdid.py` | #588 | Mid | Low | | Reuse the demeaner's factorized codes in `absorbed_fe_rank`/`absorbed_fe_cr1_k_increment` instead of re-factorizing: at 186k rows the rank helper adds ~1.9 ms per absorbed fit (7.7% of the fastest Rust-served TWFE fit) and the K_reference increment ~3.2 ms per clustered-hc1 absorbed fit (~13%; see `docs/performance-plan.md` "Component-aware absorbed-FE rank"), and the helpers and `demean_by_groups` factorize the same group columns. Threading the codes through the call sites halves the factorize work; the `connected_components` call itself is ~1.1 ms. Deliberately not done in the correctness PRs. | `diff_diff/utils.py` | #variance-inventory | Quick | Low | | `EfficientDiD` conditional path: the largest remaining O(n) stage is the sieve/nuisance construction outside the tiled pass (~9s at 10k). (The `_ridge_solve_weights` Python-prep shave landed 2026-07-07 — the `omega_stack[rest]` fancy-index copy and tail scatter are skipped when no row is zero-masked, byte-identical outputs; the `zero_mask` abs scan itself remains, needed for correctness.) | `efficient_did_covariates.py` | CS-scaling | Mid | Low | | `_rq_fit` LP assembly is dense (`A_eq = [X, I, -I]` with dense identity blocks, rebuilt per cell fit): a `scipy.sparse` construction would cut memory and likely HiGHS time for large cells / bootstrap-heavy covariate CiC/QDiD fits. CAVEAT before doing it: a different matrix representation can change HiGHS's vertex selection at degenerate/tied QR optima - end-to-end covariate goldens are tie-selection-gated (fine), but the `qr_cases` tight coefficient matches may shift to the equal-loss branch; re-run the parity suite and re-calibrate if needed. | `diff_diff/changes_in_changes.py::_rq_fit` | covariates PR | Quick | Low | diff --git a/diff_diff/__init__.py b/diff_diff/__init__.py index 0834c3e7..d3a7fa14 100644 --- a/diff_diff/__init__.py +++ b/diff_diff/__init__.py @@ -160,6 +160,8 @@ ) from diff_diff.lpdid import LPDiD from diff_diff.lpdid_results import LPDiDResults +from diff_diff.lwdid import LWDiD +from diff_diff.lwdid_results import LWDiDResults from diff_diff.mmm import ( MeridianROIPrior, meridian_calibration_mask, @@ -459,6 +461,9 @@ def __getattr__(name: str) -> _Any: # LPDiD (Local Projections DiD) "LPDiD", "LPDiDResults", + # LWDiD (Lee & Wooldridge rolling transformation DiD) + "LWDiD", + "LWDiDResults", # Visualization "plot_bacon", "plot_event_study", diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index 521add24..87f3b4a5 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -101,7 +101,7 @@ class DifferenceInDifferences(BaseEstimator): ``vcov_type``: with ``"hc1"`` dispatches to CR1 (Liang-Zeger); with ``"hc2_bm"`` dispatches to CR2 Bell-McCaffrey (Pustejovsky-Tipton 2018 symmetric-sqrt + Satterthwaite DOF). - vcov_type : {"classical", "hc1", "hc2", "hc2_bm", "conley"}, optional + vcov_type : {"classical", "hc1", "hc2", "hc2_bm", "hc3", "conley"}, optional Variance-covariance family. Defaults to the ``robust`` alias. - ``"classical"``: non-robust OLS SEs, ``sigma_hat^2 * (X'X)^{-1}``. @@ -116,6 +116,12 @@ class DifferenceInDifferences(BaseEstimator): post-period-average ATT (see ``_compute_cr2_bm_contrast_dof`` in ``linalg.py`` and the REGISTRY.md note). Weighted CR2-BM (``survey_design=`` paths) is a separate gate. + - ``"hc3"``: jackknife-style leverage correction, meat + ``e_i^2 / (1 - h_ii)^2`` (one-way only; errors with ``cluster=``). + A leverage-one observation has no defined HC3 variance and the + vcov fails closed (warning + NaN inference) rather than flooring + ``1 - h_ii``. With ``absorb=``, routes through the full-dummy + design like hc2. - ``"conley"``: Conley 1999 spatial-HAC sandwich. Pass ``conley_coords=(lat_col, lon_col)``, ``conley_cutoff_km=``, and ``conley_lag_cutoff=`` on the constructor; pass @@ -519,7 +525,7 @@ def fit( # explicit hc2 request would still change the result surface # (full-dummy coefficients vs absorbed reduced fit) despite the # "has no effect" warning. - if absorb and not _replicate_vcov_remap and self.vcov_type in ("hc2", "hc2_bm"): + if absorb and not _replicate_vcov_remap and self.vcov_type in ("hc2", "hc2_bm", "hc3"): fixed_effects = list(fixed_effects or []) + list(absorb) absorb = None absorbed_vars = [] @@ -1669,7 +1675,7 @@ def _fit_event_study_core( # fixed_effects: the fixed_effects= path builds the full-dummy # design and solves WLS directly, with no within-transform step. # Route on the EFFECTIVE vcov family (see DifferenceInDifferences). - if absorb and not _replicate_vcov_remap_mp and self.vcov_type in ("hc2", "hc2_bm"): + if absorb and not _replicate_vcov_remap_mp and self.vcov_type in ("hc2", "hc2_bm", "hc3"): fixed_effects = list(fixed_effects or []) + list(absorb) absorb = None n_absorbed_effects = 0 @@ -2499,7 +2505,7 @@ class MultiPeriodDiD(DifferenceInDifferences): ``linalg.py``; matches clubSandwich's ``Wald_test(test="HTZ")$df_denom`` at atol=1e-10). Weighted CR2-BM (``survey_design=``) is a separate, still-gated path. - vcov_type : {"classical", "hc1", "hc2", "hc2_bm", "conley"}, optional + vcov_type : {"classical", "hc1", "hc2", "hc2_bm", "hc3", "conley"}, optional Variance-covariance family. Defaults to the ``robust`` alias. - ``"classical"``: non-robust OLS SEs, ``sigma_hat^2 * (X'X)^{-1}``. @@ -2513,6 +2519,12 @@ class MultiPeriodDiD(DifferenceInDifferences): CR2 cluster-robust with a Bell-McCaffrey Satterthwaite contrast DOF on the post-period average (see ``cluster`` above for parity details). Weighted CR2-BM (``survey_design=``) is still gated. + - ``"hc3"``: jackknife-style leverage correction, meat + ``e_i^2 / (1 - h_ii)^2`` (one-way only; errors with ``cluster=``). + A leverage-one observation has no defined HC3 variance and the + vcov fails closed (warning + NaN inference) rather than flooring + ``1 - h_ii``. With ``absorb=``, routes through the full-dummy + design like hc2. - ``"conley"``: Conley 1999 spatial-HAC sandwich via the panel block-decomposed form (matches R ``conleyreg`` with ``lag_cutoff > 0``). Pass ``conley_coords=(lat_col, lon_col)``, diff --git a/diff_diff/guides/llms-full.txt b/diff_diff/guides/llms-full.txt index 8f3c4eb0..c16a7a42 100644 --- a/diff_diff/guides/llms-full.txt +++ b/diff_diff/guides/llms-full.txt @@ -52,7 +52,7 @@ Basic 2x2 Difference-in-Differences estimator. ```python DifferenceInDifferences( - vcov_type: str | None = None, # Variance family: "hc1" (default), "classical", "hc2", "hc2_bm", "conley" + vcov_type: str | None = None, # Variance family: "hc1" (default), "classical", "hc2", "hc2_bm", "hc3", "conley" cluster: str | None = None, # Column for cluster-robust SEs alpha: float = 0.05, # Significance level inference: str = "analytical", # "analytical" or "wild_bootstrap" (wild_bootstrap requires cluster=) diff --git a/diff_diff/guides/llms.txt b/diff_diff/guides/llms.txt index 8369d7a5..01832f5d 100644 --- a/diff_diff/guides/llms.txt +++ b/diff_diff/guides/llms.txt @@ -80,6 +80,7 @@ The site is organized into 5 sections, each with a landing page: - [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 by default; non-absorbing (reversible) treatment via `non_absorbing="first_entry"` (Eq. 12) or `"effect_stabilization"` (Eq. 13, window `L`). Complex-survey designs (pweight + stratified-PSU TSL SEs) on the default path via `fit(survey_design=...)`. - [ChangesInChanges](https://diff-diff.readthedocs.io/en/stable/api/changes_in_changes.html): Athey & Imbens (2006) nonlinear/distributional DiD for the 2x2 design: recovers the treated group's full counterfactual outcome distribution and quantile treatment effects (ATT + QTE grid) via the CDF transformation `F_10(F_00^{-1}(F_01(y)))`; invariant to monotone outcome transformations (unconditional fits; the covariate QR branch is not); bootstrap inference (panel or repeated cross-section resampling); point parity with R `qte::CiC()`, including its covariate branch (`covariates=` -> per-cell linear quantile regression, Melly-Santangelo-style conditional CiC). Continuous outcomes, numeric covariates. Alias `CiC`. - [QDiD](https://diff-diff.readthedocs.io/en/stable/api/changes_in_changes.html): **Deprecated 3.9, removed 4.0 - use `ChangesInChanges(method="qdid")`.** Athey & Imbens (2006) quantile DiD comparison estimator (additive quantile-by-quantile DiD, matching R `qte::QDiD()` including its covariate branch via `covariates=`); same bootstrap machinery as ChangesInChanges. The paper recommends CiC over QDiD (scale-dependent model with testable restrictions; a non-monotonicity warning fires when violated - unconditional fits only, the covariate-path counterfactual quantile curve is monotone by construction). +- [LWDiD](https://diff-diff.readthedocs.io/en/stable/api/lwdid.html): Lee & Wooldridge (2025, 2026) rolling-transformation DiD — unit-specific demean/detrend converts panel to cross-section, supports staggered adoption with flexible control groups. Signature: `LWDiD(rolling='demean', estimation_method='reg', vcov_type='hc1', cluster=None, control_group='not_yet_treated', alpha=0.05, n_bootstrap=0, seed=None, pscore_trim=0.01, n_neighbors=1, caliper=None, with_replacement=True, n_jobs=1).fit(data, outcome, unit, time, treatment, first_treat=None, covariates=None)`. `estimation_method` values: `reg` (papers' RA), `ipw`, `dr` (papers' IPWRA, doubly robust), `psm`; `vcov_type` values: `classical`/`hc1`/`hc2`/`hc3` for `reg`; `ipw`/`dr` accept `hc1` only (influence-function variance); `psm` accepts `hc1` as configuration only - PSM inference is unavailable (NaN) pending an Abadie-Imbens matching variance; cluster-robust inference via the constructor's `cluster=` column (hc1/CR1 only, not a `vcov_type` value; rejected for `psm`). Per-period effects: post-fit `results.aggregate('event_study')`. - [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/linalg.py b/diff_diff/linalg.py index a0fd18a4..2546f1c2 100644 --- a/diff_diff/linalg.py +++ b/diff_diff/linalg.py @@ -1161,7 +1161,7 @@ def solve_ols( Type of weights: "pweight" (inverse selection probability), "fweight" (frequency), or "aweight" (inverse variance). Affects variance estimation but not coefficient computation. - vcov_type : {"classical", "hc1", "hc2", "hc2_bm", "conley"}, default "hc1" + vcov_type : {"classical", "hc1", "hc2", "hc2_bm", "hc3", "conley"}, default "hc1" Variance-covariance family forwarded to :func:`compute_robust_vcov`: - ``"classical"``: non-robust OLS SE, ``sigma_hat^2 * (X'X)^{-1}``. @@ -1177,6 +1177,11 @@ def solve_ols( ``fweight`` raise ``NotImplementedError`` (port matches the ``pweight`` convention only; aweight/fweight derivations are a separate methodology task). + - ``"hc3"``: jackknife-style leverage correction, meat + ``e_i^2 / (1 - h_ii)^2``. One-way only; raises with + ``cluster_ids``. An observation with leverage ``h_ii ~ 1`` has no + defined HC3 variance and the vcov fails closed (warning + NaN) + rather than flooring ``1 - h_ii``. - ``"conley"``: Conley (1999) spatial-HAC sandwich. Requires ``conley_coords`` (n × 2 array) and ``conley_cutoff_km`` (positive bandwidth, no default per Conley 1999 Section 5's sensitivity-grid @@ -1913,7 +1918,7 @@ def _solve_ols_numpy( return coefficients, residuals, vcov -_VALID_VCOV_TYPES = frozenset({"classical", "hc1", "hc2", "hc2_bm", "conley"}) +_VALID_VCOV_TYPES = frozenset({"classical", "hc1", "hc2", "hc2_bm", "hc3", "conley"}) def _validate_vcov_args( @@ -1936,7 +1941,7 @@ def _validate_vcov_args( ValueError If ``vcov_type`` is not in the allowed set, or if ``cluster_ids`` is combined with a ``vcov_type`` that is one-way only (``classical``, - ``hc2``). + ``hc2``, ``hc3``). NotImplementedError If ``vcov_type == "conley"`` is combined with ``weights`` (regardless of ``weight_type``: weighted Conley is not implemented on the @@ -1956,7 +1961,7 @@ def _validate_vcov_args( # Mirrored K_reference-adjustment contract for direct compute_robust_vcov # / kernel callers (solve_ols routes enforce it at its own front door). _validate_cluster_k_adjustment(cluster_k_adjustment, cluster_ids, vcov_type) - if vcov_type in ("classical", "hc2") and cluster_ids is not None: + if vcov_type in ("classical", "hc2", "hc3") and cluster_ids is not None: msg = { "classical": ( "classical SEs are one-way only; pass vcov_type='hc1' or " @@ -1965,6 +1970,10 @@ def _validate_vcov_args( "hc2": ( "hc2 is one-way only. Use vcov_type='hc2_bm' for " "cluster-robust Bell-McCaffrey." ), + "hc3": ( + "hc3 is one-way only. Use vcov_type='hc1' (CR1) or " + "'hc2_bm' (CR2 Bell-McCaffrey) for cluster-robust." + ), }[vcov_type] raise ValueError(msg) # Weighted Bell-McCaffrey (both one-way and cluster) is now supported via @@ -2068,7 +2077,7 @@ def resolve_vcov_type( ``"hc1"`` and ``robust=False`` to ``"classical"``. - If ``vcov_type`` is supplied: it must be one of the values in the module-level ``_VALID_VCOV_TYPES`` set, namely - ``{"classical", "hc1", "hc2", "hc2_bm", "conley"}``. + ``{"classical", "hc1", "hc2", "hc2_bm", "hc3", "conley"}``. - If ``robust=False`` is supplied together with a non-``"classical"`` ``vcov_type``, raise ``ValueError`` - the combination is ambiguous. @@ -2086,7 +2095,8 @@ def resolve_vcov_type( Returns ------- str - One of ``"classical"``, ``"hc1"``, ``"hc2"``, ``"hc2_bm"``, ``"conley"``. + One of ``"classical"``, ``"hc1"``, ``"hc2"``, ``"hc2_bm"``, + ``"hc3"``, ``"conley"``. Raises ------ @@ -2134,7 +2144,7 @@ def compute_robust_vcov( conley_lag_cutoff: Optional[int] = None, ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: """ - Compute variance-covariance matrix under one of five `vcov_type` variants. + Compute variance-covariance matrix under one of six `vcov_type` variants. Uses the sandwich estimator: (X'X)^{-1} * meat * (X'X)^{-1}, with the meat matrix determined by the ``vcov_type`` dispatch: @@ -2149,6 +2159,9 @@ def compute_robust_vcov( ``sum_i (u_i^2 / (1 - h_ii)) x_i x_i'`` where ``h_ii`` are hat-matrix diagonals. No DOF adjustment beyond ``n - k``. One-way only; errors with ``cluster_ids``. + - ``"hc3"``: jackknife-style leverage correction, meat + ``sum_i (u_i^2 / (1 - h_ii)^2) x_i x_i'`` (matches ``sandwich::vcovHC`` + type="HC3": no DOF factor). One-way only; errors with ``cluster_ids``. - ``"hc2_bm"``: one-way HC2 meat plus Imbens-Kolesar (2016) Bell-McCaffrey Satterthwaite degrees of freedom per coefficient when ``cluster_ids`` is ``None``. When ``cluster_ids`` is supplied, dispatches to the @@ -2195,8 +2208,8 @@ def compute_robust_vcov( Weight type: "pweight", "fweight", or "aweight". vcov_type : str, default "hc1" One of ``"classical"``, ``"hc1"``, ``"hc2"``, ``"hc2_bm"``, - ``"conley"`` (see top-level docstring above for the dispatch - contract). + ``"hc3"``, ``"conley"`` (see top-level docstring above for the + dispatch contract). conley_coords : ndarray of shape (n, 2), optional, keyword-only Required when ``vcov_type="conley"``. Two-column array of ``[lat, lon]`` (degrees, for ``conley_metric="haversine"``) or @@ -2221,8 +2234,9 @@ def compute_robust_vcov( return_dof : bool, default False When True, returns ``(vcov, dof_vec)`` tuple. ``dof_vec`` is a length-k array of per-coefficient degrees of freedom. For ``classical``, - ``hc1``, ``hc2``: every element is ``n_eff - k``. For ``hc2_bm`` - one-way: Imbens-Kolesar (2016) Satterthwaite DOF per contrast. + ``hc1``, ``hc2``, ``hc3``: every element is ``n_eff - k``. For + ``hc2_bm`` one-way: Imbens-Kolesar (2016) Satterthwaite DOF per + contrast. cluster_k_adjustment : int, default 0, keyword-only Signed K_reference adjustment added to the visible column count in the CLUSTERED CR1 finite-sample factor only (absorbed FE not nested @@ -3550,9 +3564,9 @@ def _compute_robust_vcov_numpy( return vcov_cr2 # ------------------------------------------------------------------ - # HC2 / HC2+BM one-way (no cluster). + # HC2 / HC2+BM / HC3 one-way (no cluster). # ------------------------------------------------------------------ - if vcov_type in ("hc2", "hc2_bm"): + if vcov_type in ("hc2", "hc2_bm", "hc3"): # cluster path handled above; here cluster_ids is None by construction. # **Weighted hc2_bm one-way**: clubSandwich's CR2 with singleton clusters # uses the bias-corrected adjustment `A_i = 1 / sqrt(G_i)` where @@ -3580,8 +3594,50 @@ def _compute_robust_vcov_numpy( if return_dof: return vcov_cr2, dof_cr2 return vcov_cr2 - h_diag = _compute_hat_diagonals(X, bread_matrix, weights=weights) + # fweight semantics are REPLICATED DATA (Registry: integer counts, + # df = sum(w) - k, HC1 expansion parity): each replicate row's + # leverage in the expanded design is x_i'(X'WX)^{-1}x_i WITHOUT the + # w multiplier, so the leverage denominator uses the unweighted + # quadratic form against the weighted bread. The WLS-hat convention + # (w * quadform, R sandwich::vcovHC) applies to aweight/pweight + # only (review round 6: the weighted hat under fweight produced + # HC2/HC3 variances up to ~5x the literal np.repeat expansion). + h_diag = _compute_hat_diagonals( + X, + bread_matrix, + weights=None if weight_type == "fweight" else weights, + ) + # Leverage-one observations make the HC3 leave-one-out residual + # undefined (and HC2 nearly so): flooring 1 - h_ii would fabricate + # an arbitrary finite variance for a perfectly-leveraged point + # (e.g. a single treated unit under [1, D]). HC3 fails closed with + # a NaN vcov instead (LWDiD fix-wave review finding); HC2/HC2-BM + # keep their long-standing floor behavior (released surface; + # pre-existing, tracked separately). This check runs BEFORE the + # generic over-one HC1 fallback below (round-10 review: numerically + # over-one leverage previously escaped into an HC1 result still + # labeled hc3 - h >= 1 - 1e-8 covers h > 1 + 1e-6 entirely). + if vcov_type == "hc3" and np.any(h_diag >= 1.0 - 1e-8): + n_lev1 = int(np.sum(h_diag >= 1.0 - 1e-8)) + warnings.warn( + f"HC3 variance is undefined: {n_lev1} observation(s) have " + f"hat-matrix leverage ~1 (a perfectly-leveraged design, " + f"e.g. a single treated unit). Returning NaN vcov; use " + f"vcov_type='classical' exact inference or add treated " + f"units.", + UserWarning, + stacklevel=3, + ) + nan_vcov = np.full((X.shape[1], X.shape[1]), np.nan) + if return_dof: + # Contract: a length-k DOF vector (round-24 review: None + # broke direct consumers indexing the result); NaN keeps + # the fail-closed semantics through safe_inference. + return nan_vcov, np.full(X.shape[1], np.nan) + return nan_vcov if np.any(h_diag > 1.0 + 1e-6): + # hc2/hc2_bm only: hc3 designs with over-one leverage are + # already caught by the fail-closed guard above. warnings.warn( f"Hat-matrix diagonal exceeds 1 (max={h_diag.max():.6f}); " "the design is near-singular. Falling back to HC1.", @@ -3598,26 +3654,29 @@ def _compute_robust_vcov_numpy( return_dof=return_dof, ) one_minus_h = np.maximum(1.0 - h_diag, 1e-10) - # HC2 meat: sum_i (u_i^2 / (1 - h_ii)) x_i x_i', with pweight scaling - # matching the HC1 convention (w_i * u_i / sqrt(1 - h_ii) as score). + # HC2 meat: sum_i (u_i^2 / (1 - h_ii)) x_i x_i'; HC3 squares the + # leverage denominator (jackknife-style, sandwich::vcovHC type="HC3"). + # pweight scaling matches the HC1 convention (w_i * u_i / sqrt(denom) + # as score). + lev_denom = one_minus_h**2 if vcov_type == "hc3" else one_minus_h if weights is not None and weight_type == "fweight": - factor = weights * (residuals**2) / one_minus_h + factor = weights * (residuals**2) / lev_denom meat = X.T @ (X * factor[:, np.newaxis]) elif weights is not None and weight_type == "pweight": - # pweight scores carry w in the score, so meat = sum (w u / sqrt(1-h))^2 x x' - scaled = weights * residuals / np.sqrt(one_minus_h) + # pweight scores carry w in the score, so meat = sum (w u / sqrt(denom))^2 x x' + scaled = weights * residuals / np.sqrt(lev_denom) scores_hc2 = X * scaled[:, np.newaxis] meat = scores_hc2.T @ scores_hc2 else: - # aweight / unweighted: meat = sum_i (u_i^2 / (1 - h_ii)) x_i x_i' - factor = (residuals**2) / one_minus_h + # aweight / unweighted: meat = sum_i (u_i^2 / denom_i) x_i x_i' + factor = (residuals**2) / lev_denom # Zero out zero-weight rows under aweight (subpopulation invariance) if weights is not None and np.any(weights == 0): factor = factor * (weights > 0) meat = X.T @ (X * factor[:, np.newaxis]) - # Sandwich without DOF adjustment for HC2 (matches sandwich::vcovHC - # type="HC2" convention: no (n/(n-k)) factor). + # Sandwich without DOF adjustment for HC2/HC3 (matches sandwich::vcovHC + # type="HC2"/"HC3" convention: no (n/(n-k)) factor). try: temp = np.linalg.solve(bread_matrix, meat) vcov = np.linalg.solve(bread_matrix, temp.T).T @@ -3632,7 +3691,7 @@ def _compute_robust_vcov_numpy( if not return_dof: return vcov - if vcov_type == "hc2": + if vcov_type in ("hc2", "hc3"): dof_vec = np.full(k, n_eff - k, dtype=np.float64) else: # hc2_bm dof_vec = _compute_bm_dof_oneway(X, bread_matrix, h_diag, weights=weights) @@ -4324,7 +4383,7 @@ class LinearRegression: Resolved survey design for Taylor Series Linearization variance estimation. When provided, weights and weight_type are canonicalized from this object. - vcov_type : {"classical", "hc1", "hc2", "hc2_bm", "conley"}, optional + vcov_type : {"classical", "hc1", "hc2", "hc2_bm", "hc3", "conley"}, optional Variance-covariance family. Defaults to the ``robust`` alias (``robust=True`` -> ``"hc1"``, ``robust=False`` -> ``"classical"``). Passing an explicit ``vcov_type`` overrides ``robust`` unless the diff --git a/diff_diff/lwdid.py b/diff_diff/lwdid.py new file mode 100644 index 00000000..14bc6efd --- /dev/null +++ b/diff_diff/lwdid.py @@ -0,0 +1,4503 @@ +"""LWDiD: Lee & Wooldridge (2025, 2026) rolling-transformation DiD. + +Converts panel DiD into cross-sectional estimation via unit-specific +rolling transformations of the outcome variable. Supports common timing +and staggered adoption designs with regression adjustment ('reg'), +inverse probability weighting ('ipw'), doubly robust ('dr'), and +propensity score matching ('psm') estimation. + +References +---------- +Lee, S. J. & Wooldridge, J. M. (2025). "A Simple Transformation Approach + to Difference-in-Differences Estimation for Panel Data." SSRN 4516518. +Lee, S. J. & Wooldridge, J. M. (2026). "Simple Approaches to Inference + with Difference-in-Differences Estimators with Small Cross-Sectional + Sample Sizes." SSRN 5325686. +""" + +from __future__ import annotations + +import warnings +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import pandas as pd +from scipy import linalg as scipy_linalg + +from diff_diff._base import BaseEstimator +from diff_diff.linalg import _detect_rank_deficiency, solve_logit, solve_ols +from diff_diff.lwdid_results import LWDiDResults +from diff_diff.utils import safe_inference, validate_binary + +_VALID_ROLLING = ("demean", "detrend", "demeanq", "detrendq") +_VALID_ESTIMATION_METHODS = ("reg", "ipw", "dr", "psm") +_VALID_VCOV_TYPES = ("classical", "hc1", "hc2", "hc3") +_VALID_CONTROL_GROUPS = ("never_treated", "not_yet_treated") + +# Propensity score trimming bounds for numerical stability +#: Column names written into internal estimation/plotting frames. A user +#: role column with one of these names would be silently overwritten +#: (e.g. cluster='_treat' turned the cluster labels into the treatment +#: regressor), so they are rejected up front by _validate_inputs. +_RESERVED_INTERNAL_COLUMNS = frozenset( + { + "_treat", + "_ydot", + "_ydot_avg", + "_ever_treated", + "_boot_unit", + "_lwdid_time_pos", + "_lwdid_cohort_pos", + "_lwdid_season", + } +) + + +def _normalize_cohorts( + cohort_series: pd.Series, + *, + max_time: Any, +) -> Tuple[pd.Series, int, int]: + """Canonicalize NUMERIC-encoded cohort values to the house convention. + + Never-treated is encoded as ``NaN`` or ``0`` downstream. Two additional + encodings are recoded here with a warning (no silent reinterpretation): + + - ``np.inf`` -> ``0`` (CallawaySantAnna convention: CS recodes exactly + ``0``/``inf``; the NaN limb accepted downstream is an LWDiD-only + extension needed for datetime scales, documented in REGISTRY). + - finite ``g > max_time`` (beyond-window) -> ``0``: a unit that never + switches on inside the observed window is within-sample + never-treated. This is a documented DEVIATION from CS (which keeps + finite cohorts out of never-treated); under + ``control_group='never_treated'`` recoded units join the control + pool and contribute no pre-period event cells. + + Negative finite values and ``-inf`` are nonsensical cohort encodings + and raise ``ValueError`` (never silently classified). + + Parameters + ---------- + cohort_series : pd.Series + Row-level cohort column, already numeric (datetime/Period panels + must be encoded to integer positions first — see + ``_encode_staggered_time_scale``). + max_time : scalar + Largest observed time value on the same numeric scale. + + Returns + ------- + tuple + ``(normalized_series, n_inf_rows_recoded, n_beyond_rows_recoded)``. + """ + values = pd.to_numeric(cohort_series, errors="raise") + finite = np.isfinite(values.to_numpy(dtype=float, na_value=np.nan)) + negative = finite & (values.to_numpy(dtype=float, na_value=np.nan) < 0) + neg_inf = np.isneginf(values.to_numpy(dtype=float, na_value=np.nan)) + if negative.any() or neg_inf.any(): + bad = sorted(pd.unique(values[negative | neg_inf]).tolist()) + raise ValueError( + f"Cohort column contains negative value(s) {bad[:5]}: cohort " + f"values must be 0/NaN (never-treated), np.inf (recoded to " + f"never-treated), or an observed treatment period." + ) + inf_mask = np.isposinf(values.to_numpy(dtype=float, na_value=np.nan)) + n_inf = int(inf_mask.sum()) + if n_inf: + warnings.warn( + f"first_treat=inf found on {n_inf} row(s); recoding to 0 " + f"(never-treated). Use first_treat=0 to suppress this warning.", + UserWarning, + stacklevel=3, + ) + beyond_mask = finite & (values.to_numpy(dtype=float, na_value=np.nan) > 0) + beyond_mask &= values.to_numpy(dtype=float, na_value=np.nan) > float(max_time) + n_beyond = int(beyond_mask.sum()) + if n_beyond: + bad_vals = sorted(pd.unique(values[beyond_mask]).tolist()) + warnings.warn( + f"Cohort value(s) {bad_vals[:5]} exceed the last observed period " + f"({max_time}); units in these cohorts never switch on within " + f"the sample and are recoded to never-treated (0). This deviates " + f"from CallawaySantAnna, which keeps finite cohorts out of " + f"never-treated; see docs/methodology/REGISTRY.md (LWDiD).", + UserWarning, + stacklevel=3, + ) + if n_inf or n_beyond: + values = values.copy() + values[inf_mask | beyond_mask] = 0 + return values, n_inf, n_beyond + + +def _check_treatment_design( + df: pd.DataFrame, + unit: str, + time: str, + treatment: str, + first_treat: Optional[str] = None, +) -> None: + """Validate the treatment design in a single vectorized pass. + + One sort + groupby covers three checks: + + 1. Absorbing treatment: within each unit the sequence D_it must be + non-decreasing over time (once treated, always treated). + 2. Common timing (``first_treat is None``): every treated unit must + first switch to D_it = 1 in the same period; heterogeneous onsets + require the staggered interface (``first_treat`` cohort column). + 3. Staggered (``first_treat`` given): over each unit's OBSERVED rows, + the treatment indicator must equal ``1[t >= g_i]`` exactly — no + D_it = 1 before the cohort value and no D_it = 0 at or after it; + units with cohort NaN/0 (never treated) must have no D_it = 1 + rows. The row at ``t == g_i`` itself may be unobserved (unbalanced + panels with a missing onset row are accepted). Finite positive + cohort values must be members of the observed time support — + numeric between-period cohorts are rejected (datetime/Period + cohorts are mapped to the next observed period by the encoding + step before this check). + + Precondition: on the staggered path the cohort column is expected to + be already normalized+encoded (``_encode_staggered_time_scale`` + + ``_normalize_cohorts``) so that never-treated is NaN/0 and + beyond-window/inf sentinels no longer occur. + + Parameters + ---------- + df : pd.DataFrame + Panel data in long format. + unit, time, treatment : str + Column names of the unit identifier, time period, and binary + treatment indicator. + first_treat : str or None, default None + Cohort (first-treatment-time) column for staggered designs. + + Raises + ------ + ValueError + If any applicable design check fails. + """ + cols = [unit, time, treatment] + if first_treat is not None: + cols.append(first_treat) + ordered = df[cols].sort_values([unit, time], kind="stable") + + # (1) Absorbing treatment: within-unit first difference must never + # be negative (a 1 -> 0 switch). + diffs = ordered.groupby(unit, sort=False)[treatment].diff() + non_absorbing = (diffs < 0).to_numpy() + if non_absorbing.any(): + bad_units = pd.unique(ordered.loc[non_absorbing, unit]) + preview = ", ".join(repr(u) for u in bad_units[:5]) + suffix = "" if len(bad_units) <= 5 else f", ... ({len(bad_units)} units total)" + raise ValueError( + f"Non-absorbing treatment detected for unit(s) {preview}{suffix}: " + f"treatment switches from 1 to 0. LWDiD requires absorbing treatment." + ) + + # First observed treatment period per treated unit (rows are already + # time-sorted within unit, so first() is the onset). + treated_rows = ordered.loc[ordered[treatment] == 1] + onset = treated_rows.groupby(unit, sort=False)[time].first() + + if first_treat is None: + # (2) Common timing: every ever-treated unit's observed rows must + # satisfy D_it = 1[t >= S] for the single global onset + # S = min(first observed treated period). Comparing first OBSERVED + # treated rows directly would falsely reject a unit whose t = S + # row is simply missing (round-8 review; the staggered branch + # already permits an unobserved onset row). A genuinely + # heterogeneous unit (true onset S' > S) has an observed + # UNTREATED row at t >= S and is rejected. + if len(onset) == 0: + return + onset_s = onset.min() + ever_treated = set(onset.index) + ever_row = ordered[unit].isin(ever_treated).to_numpy() + late_zero = ( + ever_row + & (ordered[time] >= onset_s).to_numpy() + & (ordered[treatment].to_numpy(dtype=float) == 0) + ) + if late_zero.any(): + bad_units = sorted(pd.unique(ordered.loc[late_zero, unit]).tolist()) + preview = ", ".join(repr(u) for u in bad_units[:5]) + suffix = "" if len(bad_units) <= 5 else f", ... ({len(bad_units)} units total)" + raise ValueError( + f"Treated unit(s) {preview}{suffix} have untreated observed " + f"rows at or after the common onset {onset_s!r}: treated " + f"units have heterogeneous first-treatment periods, but no " + f"cohort column was given. Common-timing LWDiD requires a single " + f"treatment onset; pass first_treat= to use the staggered " + f"(cohort) interface." + ) + return + + # (3a) Support membership: finite positive cohorts must be observed + # time values. Post-normalization, beyond-window sentinels no longer + # occur, so this targets exactly numeric BETWEEN-period cohorts + # (e.g. g=4.5 with observed times {4, 5}). + cohort_by_unit = ordered.groupby(unit, sort=False)[first_treat].first() + observed_times = pd.Index(pd.unique(ordered[time])) + positive = cohort_by_unit.notna() & (cohort_by_unit > 0) + off_support = positive & ~cohort_by_unit.isin(observed_times) + if off_support.to_numpy().any(): + bad_vals = sorted(pd.unique(cohort_by_unit[off_support]).tolist()) + raise ValueError( + f"Cohort value(s) {bad_vals[:5]} in column '{first_treat}' are " + f"not observed time periods. Numeric between-period cohorts are " + f"not supported; use an observed period value (datetime/Period " + f"cohorts are mapped to the next observed period automatically)." + ) + + # (3b) Per unit, over OBSERVED rows, D_it must equal 1[t >= g_i]: + # never-treated units (cohort NaN/0) have no D=1 rows; treated-cohort + # units have no D=1 before g and no D=0 at/after g. The onset row + # itself may be unobserved (unbalanced panels are accepted). + g_by_row = ordered[unit].map(cohort_by_unit) + g_arr = g_by_row.to_numpy(dtype=float, na_value=np.nan) + never_row = np.isnan(g_arr) | (g_arr == 0) + d_arr = ordered[treatment].to_numpy(dtype=float) + t_arr = ordered[time].to_numpy(dtype=float) + with np.errstate(invalid="ignore"): + premature = ~never_row & (t_arr < g_arr) & (d_arr == 1) + untreated_post = ~never_row & (t_arr >= g_arr) & (d_arr == 0) + never_treated_rows = never_row & (d_arr == 1) + violation = premature | untreated_post | never_treated_rows + if violation.any(): + bad_units = pd.unique(ordered.loc[violation, unit]) + preview = ", ".join(repr(u) for u in bad_units[:5]) + suffix = "" if len(bad_units) <= 5 else f", ... ({len(bad_units)} units total)" + raise ValueError( + f"Treatment column '{treatment}' is inconsistent with cohort " + f"column '{first_treat}' for unit(s) {preview}{suffix}: over " + f"observed rows, treatment must equal 1[t >= cohort] — no " + f"treatment=1 before the cohort period, no treatment=0 at or " + f"after it, and never-treated units (cohort NaN or 0) must have " + f"no treatment=1 rows." + ) + + +def _encode_staggered_time_scale( + df: pd.DataFrame, + time: str, + first_treat: str, +) -> Tuple[pd.DataFrame, str, str, Optional[Dict[str, Dict[int, Any]]]]: + """Re-encode datetime/Period time scales as integer positions. + + The staggered machinery relies on integer time semantics: cohort + eligibility comparisons (``g > 0``, ``g > t``), event times ``t - g``, + and the never-treated sentinel 0. Datetime and Period panels are + therefore mapped onto the ordered support of observed time values -- + the k-th observed period becomes position k (1-based, so 0 stays free + for the never-treated sentinel, coded NaT in datetime panels). Cohort + values between observed periods map to the next observed position and + cohorts beyond the window map to T + 1, preserving the onset + consistency checks. Numeric panels are returned unchanged. + + Parameters + ---------- + df : pd.DataFrame + Panel data (a private copy owned by the caller; encoded position + columns are added in place). + time, first_treat : str + Time and cohort column names. + + Returns + ------- + tuple + ``(df, time_column, cohort_column, label_maps)`` where + ``label_maps`` is None for numeric panels and otherwise maps + integer positions back to the original time/cohort labels. + + Raises + ------ + ValueError + If the two columns do not share the same time family: exactly one + of them is date-like, one is datetime64 while the other is Period + (either direction), or both are Period with different frequencies. + """ + time_is_datetime = pd.api.types.is_datetime64_any_dtype(df[time]) + cohort_is_datetime = pd.api.types.is_datetime64_any_dtype(df[first_treat]) + time_is_period = isinstance(df[time].dtype, pd.PeriodDtype) + cohort_is_period = isinstance(df[first_treat].dtype, pd.PeriodDtype) + if not (time_is_datetime or time_is_period or cohort_is_datetime or cohort_is_period): + return df, time, first_treat, None + # datetime64 and Period are distinct time families: position lookups + # (searchsorted, dict membership) crash inside pandas when mixed, so + # both directions are rejected up front with the documented ValueError. + if time_is_datetime != cohort_is_datetime or time_is_period != cohort_is_period: + raise ValueError( + f"Columns '{time}' (time) and '{first_treat}' (first_treat) must " + f"share the same time scale; got dtypes {df[time].dtype} and " + f"{df[first_treat].dtype}. Encode both as datetime64, both as " + f"Period with the same frequency, or both as numeric." + ) + if time_is_period and df[time].dtype.freq != df[first_treat].dtype.freq: + raise ValueError( + f"Columns '{time}' (time) and '{first_treat}' (first_treat) are " + f"Period columns with different frequencies ({df[time].dtype} vs " + f"{df[first_treat].dtype}). Convert them to a common frequency " + f"before fitting." + ) + + support = pd.Index(pd.unique(df[time])).sort_values() + time_pos: Dict[Any, int] = {value: index + 1 for index, value in enumerate(support)} + + def cohort_pos(value: Any) -> int: + if value in time_pos: + return time_pos[value] + # Between observed periods -> next observed position; beyond the + # window -> T + 1 (vacuously consistent, like numeric cohorts + # past the last observed period). + return int(support.searchsorted(value, side="left")) + 1 + + cohort_map: Dict[Any, int] = { + value: cohort_pos(value) + for value in pd.Index(pd.unique(df[first_treat])) + if pd.notna(value) + } + df["_lwdid_time_pos"] = df[time].map(time_pos).astype(int) + df["_lwdid_cohort_pos"] = df[first_treat].map(cohort_map).fillna(0).astype(int) + # Preserve the CALENDAR season before time values become dense + # positions: the seasonal transforms' numeric fallback derives quarter + # as (t - 1) % 4 + 1, which relabels every season after a globally + # missing calendar period (review round 3 P0 - silent seasonal + # mixing). The q-variant transforms prefer this column when present. + if time_is_datetime: + df["_lwdid_season"] = df[time].dt.quarter.to_numpy() + else: + df["_lwdid_season"] = np.array([value.quarter for value in df[time]]) + time_reverse = {position: value for value, position in time_pos.items()} + # Cohort positions relabel to the CANONICAL observed period at that + # position (round-24 review: reversing cohort_map let two raw + # between-period labels mapping to the same onset collide, with the + # surviving label depending on input row order); off-support + # positions (beyond-window T+1) keep a deterministic raw label and + # are normalized to never-treated downstream anyway. + label_maps = { + "time": time_reverse, + "cohort": { + position: time_reverse.get(position, value) + for value, position in sorted(cohort_map.items(), key=lambda kv: str(kv[0])) + }, + } + return df, "_lwdid_time_pos", "_lwdid_cohort_pos", label_maps + + +def _relabel_staggered_results( + results: LWDiDResults, + label_maps: Dict[str, Dict[int, Any]], +) -> LWDiDResults: + """Map integer time positions in staggered results back to original labels. + + Cohort and calendar-time keys (and the nested ``'cohort'``/``'time'`` + entries) are restored to the user's datetime/Period labels. Relative + event times remain integers: they are position differences on the + ordered time support. + """ + time_labels = label_maps["time"] + cohort_labels = label_maps["cohort"] + if results.cohort_effects is not None: + relabeled_cohorts: Dict[Any, Dict[str, Any]] = {} + for g, info in results.cohort_effects.items(): + info["cohort"] = cohort_labels.get(info["cohort"], info["cohort"]) + relabeled_cohorts[cohort_labels.get(g, g)] = info + results.cohort_effects = relabeled_cohorts + if results.cohort_time_effects is not None: + relabeled_cells: Dict[Any, Dict[str, Any]] = {} + for (g, t), info in results.cohort_time_effects.items(): + info["cohort"] = cohort_labels.get(info["cohort"], info["cohort"]) + info["time"] = time_labels.get(info["time"], info["time"]) + relabeled_cells[(cohort_labels.get(g, g), time_labels.get(t, t))] = info + results.cohort_time_effects = relabeled_cells + return results + + +class LWDiD(BaseEstimator): + """Lee & Wooldridge rolling-transformation DiD estimator. + + Parameters + ---------- + rolling : {'demean', 'detrend', 'demeanq', 'detrendq'}, default 'demean' + Unit-specific transformation method. + 'demean': subtract pre-treatment mean + 'detrend': subtract pre-treatment linear trend + 'demeanq': subtract unit-specific seasonal (quarterly) means + 'detrendq': subtract unit-specific linear trend + seasonal effects + estimation_method : {'reg', 'ipw', 'dr', 'psm'}, default 'reg' + Treatment effect estimation method. + 'reg': regression adjustment (OLS) + 'ipw': inverse probability weighting + 'dr': doubly robust (augmented IPW) + 'psm': propensity score matching (1:n_neighbors nearest-neighbor, + 1:1 by default); + POINT ESTIMATES ONLY - inference is NaN pending an + Abadie-Imbens matching variance (see DEFERRED.md) + vcov_type : {'classical', 'hc1', 'hc2', 'hc3'}, default 'hc1' + Variance-covariance estimator. + 'hc2': leverage-corrected (u_i^2 / (1-h_ii)) + 'hc3': jackknife-style leverage correction (u_i^2 / (1-h_ii)^2) + The full set applies to ``estimation_method='reg'`` only: + 'ipw'/'dr'/'psm' accept 'hc1' alone (the influence-function + variance on ipw/dr; psm reports NaN inference - see below). + Cluster-robust (CR1) inference activates via the ``cluster=`` + parameter, not through a ``vcov_type`` value, composes only with + 'hc1', and is rejected for 'psm'. + cluster : str or None, default None + Column name for cluster-robust (CR1) standard errors. When set, + clustered inference is active for the whole fit. + control_group : {'never_treated', 'not_yet_treated'}, default 'not_yet_treated' + Control group definition for staggered designs. Both options + require never-treated units: 'never_treated' needs at least two, + and 'not_yet_treated' needs at least one so that every cohort-time + cell keeps a valid control pool. Panels where all units are + eventually treated are rejected with a ValueError. + alpha : float, default 0.05 + Significance level for confidence intervals. + n_bootstrap : int, default 0 + Number of bootstrap replications (0 = analytical inference). + seed : int or None, default None + Random seed for bootstrap inference. + pscore_trim : float, default 0.01 + Propensity score trimming threshold. Scores below this value + or above (1 - pscore_trim) are clipped. Used by IPW/DR/PSM. + n_neighbors : int, default 1 + Number of nearest neighbors for PSM matching. + caliper : float or None, default None + Maximum allowable distance for PSM matches. Unmatched treated + units (no control within caliper) receive NaN. + with_replacement : bool, default True + Whether PSM matching is done with replacement. + n_jobs : int, default 1 + Execution parallelism for the common-timing bootstrap + (ThreadPoolExecutor when > 1; experimental). Purely an execution + setting: seeded bootstrap draws are identical for every value + (per-replicate SeedSequence streams), so it never affects any + reported number and is not stored in result provenance. + + Notes + ----- + **Parameter mapping from lwdid-py to diff-diff:** + + The standalone ``lwdid-py`` package (``from lwdid import lwdid``) uses a + functional interface with separate ``d`` (ever-treated indicator) and + ``post`` (post-period indicator) columns. In diff-diff, the ``treatment`` + column is the time-varying binary indicator ``D_i * post_t``—i.e., the + product of the two lwdid-py columns. + + .. code-block:: python + + # lwdid-py (functional API): + lwdid(data, y='y', d='d', ivar='unit', tvar='time', post='post', + rolling='demean', estimator='ra', vce=None) + + # Equivalent in diff-diff (class-based API): + LWDiD(rolling='demean', estimation_method='reg', + vcov_type='classical').fit( + data, outcome='y', unit='unit', time='time', treatment='treat') + # where data['treat'] == data['d'] * data['post'] + + Parameter correspondence: + + ================= ================== ==================================== + lwdid-py diff-diff Notes + ================= ================== ==================================== + y outcome Outcome column name + d + post treatment Binary D_it (ever-treated × post) + ivar unit Unit identifier + tvar time Time variable + gvar first_treat Cohort (first treatment period) + rolling rolling Same values + estimator='ra' estimation_method 'ra' -> 'reg', 'ipwra' -> 'dr' + vce=None vcov_type Homoskedastic == 'classical' + vce='hc1' vcov_type='hc1' Heteroskedasticity-robust + vce='cluster' cluster= Constructor cluster= parameter + cluster_var cluster Cluster variable name + controls covariates fit() covariates= parameter + control_group control_group Same values + ================= ================== ==================================== + + **Results mapping:** + + ================== ========================= ============================== + lwdid-py diff-diff Notes + ================== ========================= ============================== + result.att result.att ATT point estimate + result.se_att result.se Standard error + result.t_stat result.t_stat t-statistic + result.pvalue result.p_value p-value (note underscore) + result.ci_lower result.conf_int[0] CI lower bound + result.ci_upper result.conf_int[1] CI upper bound + result.nobs result.n_obs Number of observations + result.n_treated result.n_treated Treated units + result.n_control result.n_control Control units + result.vce_type result.vcov_type Variance family + result.cluster_var result.cluster_name Cluster variable name + result.n_clusters result.n_clusters Number of clusters + ================== ========================= ============================== + + Examples + -------- + >>> import numpy as np, pandas as pd + >>> from diff_diff.lwdid import LWDiD + >>> from diff_diff import generate_staggered_data + >>> data = generate_staggered_data(n_units=100, n_periods=8, seed=0) + >>> model = LWDiD(rolling='demean', estimation_method='reg') + >>> result = model.fit(data, outcome='outcome', unit='unit', + ... time='period', treatment='treated', + ... first_treat='first_treat') + >>> result.att != 0 + True + """ + + def __init__( + self, + rolling: str = "demean", + estimation_method: str = "reg", + vcov_type: str = "hc1", + cluster: Optional[str] = None, + control_group: str = "not_yet_treated", + alpha: float = 0.05, + n_bootstrap: int = 0, + seed: Optional[int] = None, + # Engineering parameters: + pscore_trim: float = 0.01, + n_neighbors: int = 1, + caliper: Optional[float] = None, + with_replacement: bool = True, + n_jobs: int = 1, + ) -> None: + # Validate rolling + if rolling not in _VALID_ROLLING: + raise ValueError(f"rolling must be one of {_VALID_ROLLING}, got '{rolling}'") + # Validate estimation_method + if estimation_method not in _VALID_ESTIMATION_METHODS: + raise ValueError( + f"estimation_method must be one of {_VALID_ESTIMATION_METHODS}, " + f"got '{estimation_method}'" + ) + # Validate vcov_type ('cluster' is retired as a MODE value: clustering + # activates via the cluster= constructor parameter) + if vcov_type == "cluster": + raise ValueError( + "vcov_type='cluster' is retired; pass the cluster= constructor " + "parameter (column name) to activate cluster-robust inference." + ) + if vcov_type not in _VALID_VCOV_TYPES: + raise ValueError(f"vcov_type must be one of {_VALID_VCOV_TYPES}, got '{vcov_type}'") + self._validate_vcov_config(vcov_type, estimation_method, cluster) + # Validate control_group + if control_group not in _VALID_CONTROL_GROUPS: + raise ValueError( + f"control_group must be one of {_VALID_CONTROL_GROUPS}, " f"got '{control_group}'" + ) + # Validate alpha + if ( + isinstance(alpha, bool) + or not isinstance(alpha, (int, float, np.integer, np.floating)) + or not np.isfinite(alpha) + or not (0 < alpha < 1) + ): + # Round-22 review: a one-element array passed the range check + # and failed later inside inference with a raw TypeError. + raise ValueError(f"alpha must be a scalar in (0, 1), got {alpha!r}") + alpha = float(alpha) + # Validate n_bootstrap (0 = analytical; a bootstrap needs >= 2 + # replicates for a sample standard deviation - review finding: + # n_bootstrap=1 was accepted and produced NaN downstream) + if not isinstance(n_bootstrap, (int, np.integer)) or n_bootstrap < 0: + raise ValueError(f"n_bootstrap must be a non-negative integer, " f"got {n_bootstrap}") + if n_bootstrap == 1: + raise ValueError( + "n_bootstrap must be 0 (analytical inference) or >= 2 (a " + "bootstrap standard deviation needs at least 2 replicates)." + ) + + self.rolling = rolling + self.estimation_method = estimation_method + self.vcov_type = vcov_type + self.cluster = cluster + self.control_group = control_group + self.alpha = alpha + self.n_bootstrap = int(n_bootstrap) + self.seed = seed + + # Engineering parameters (validated, never silently coerced - + # review finding: fractional n_neighbors truncated, strings became + # with_replacement=True, negative calipers matched nothing) + if not isinstance(pscore_trim, (int, float, np.integer, np.floating)) or isinstance( + pscore_trim, bool + ): + raise ValueError(f"pscore_trim must be a number, got {pscore_trim!r}") + self.pscore_trim = float(pscore_trim) + if not np.isfinite(self.pscore_trim) or not (0.0 < self.pscore_trim < 0.5): + raise ValueError("pscore_trim must be between 0 and 0.5") + if not isinstance(n_neighbors, (int, np.integer)) or isinstance(n_neighbors, bool): + raise ValueError(f"n_neighbors must be an integer, got {n_neighbors!r}") + self.n_neighbors = int(n_neighbors) + if self.n_neighbors < 1: + raise ValueError("n_neighbors must be >= 1") + if caliper is not None: + if not isinstance(caliper, (int, float, np.integer, np.floating)) or isinstance( + caliper, bool + ): + raise ValueError(f"caliper must be a positive number or None, got {caliper!r}") + caliper = float(caliper) + if not np.isfinite(caliper) or caliper <= 0: + raise ValueError(f"caliper must be a positive finite number, got {caliper}") + self.caliper = caliper + if not isinstance(with_replacement, (bool, np.bool_)): + raise ValueError(f"with_replacement must be a boolean, got {with_replacement!r}") + self.with_replacement = bool(with_replacement) + if isinstance(n_jobs, bool) or not isinstance(n_jobs, (int, np.integer)) or n_jobs < 1: + raise ValueError(f"n_jobs must be a positive integer, got {n_jobs!r}") + self.n_jobs = int(n_jobs) + + def fit( + self, + data: pd.DataFrame, + outcome: str, + unit: str, + time: str, + treatment: str, + first_treat: Optional[str] = None, + covariates: Optional[List[str]] = None, + ) -> LWDiDResults: + """Fit the LWDiD estimator. + + Parameters + ---------- + data : pd.DataFrame + Panel dataset in long format. + outcome : str + Column name of the outcome variable. + unit : str + Column name of the unit identifier. + time : str + Column name of the time period variable. + treatment : str + Column name of the binary treatment indicator (0/1). + first_treat : str, optional + Column name of the first-treatment-time (cohort) variable. + If None, assumes common timing (all treated units adopt + treatment simultaneously). + covariates : list of str, optional + Column names for control variables (covariates). Every LWDiD + path requires unit-constant (time-invariant) covariates; + time-varying columns raise a ValueError. The same + unit-constancy contract applies to the constructor's + ``cluster=`` column. + + Returns + ------- + LWDiDResults + Object containing ATT estimates, standard errors, and + inference results. + + Raises + ------ + ValueError + If required columns are missing, treatment is not binary, + or panel structure is invalid. + """ + # --- Input validation --- + df = data.copy() + cluster = self.cluster + # Re-check the vcov configuration at fit time (set_params probe + # re-init covers most mutations; this closes direct-attribute edits). + self._validate_vcov_config(self.vcov_type, self.estimation_method, cluster) + self._validate_inputs(df, outcome, unit, time, treatment, first_treat, cluster, covariates) + if cluster is not None: + from diff_diff.linalg import effective_cluster_count + + n_cl = effective_cluster_count(df[cluster].to_numpy()) + if n_cl < 2: + raise ValueError( + f"cluster='{cluster}' has {n_cl} effective cluster(s); " + f"cluster-robust inference requires at least 2." + ) + + # Validate treatment is binary + validate_binary(df[treatment].values, treatment) + + # Datetime/Period time scales are re-encoded as integer positions + # before design validation: the staggered checks and estimation + # compare cohorts against the never-treated sentinel 0 and build + # event times as t - g, which are undefined for datetime values. + label_maps = None + if first_treat is not None: + df, time, first_treat, label_maps = _encode_staggered_time_scale(df, time, first_treat) + # Cohort normalization runs AFTER encoding (numeric positions + # only — datetime beyond-window cohorts arrive as T+1 and are + # caught by the g > max_time rule) and BEFORE the design check + # (which requires canonical never-treated encodings). + df[first_treat], _, _ = _normalize_cohorts(df[first_treat], max_time=df[time].max()) + + # Unified treatment-design validation (absorbing + timing + # consistency) covering both dispatch paths + _check_treatment_design(df, unit, time, treatment, first_treat) + + # Normalize covariates + if covariates is None: + covariates = [] + + if self.estimation_method == "psm" and self.n_bootstrap > 0: + # Review findings (rounds 1-3): the staggered multiplier + # bootstrap silently did nothing for PSM (no influence-function + # representation), and the common-timing unit bootstrap + # replaced the documented fail-closed NaN inference with a + # naive pairs-bootstrap SE - invalid for nearest-neighbor + # matching with replacement (Abadie & Imbens 2008, "On the + # Failure of the Bootstrap for Matching Estimators"). + raise ValueError( + "estimation_method='psm' does not support n_bootstrap > 0: " + "matching has no influence-function representation for the " + "staggered multiplier bootstrap, and the standard bootstrap " + "is invalid for nearest-neighbor matching estimators " + "(Abadie & Imbens 2008). Use n_bootstrap=0, or " + "estimation_method='dr'." + ) + if self.estimation_method == "psm" and not covariates: + # Review round 3: without covariates there is no propensity + # model to match on; the silent delegation to regression + # adjustment returned a finite OLS SE while the results + # metadata reported method 'psm' under its documented + # fail-closed NaN-inference contract. + raise ValueError( + "estimation_method='psm' requires covariates: without them " + "there is no propensity score to match on (PSM would reduce " + "to a difference in means). Use estimation_method='reg', or " + "supply covariates." + ) + + # Dispatch to common timing or staggered + if first_treat is None: + if isinstance(df[time].dtype, pd.CategoricalDtype) and df[time].dtype.ordered: + # Ordered categoricals declare the chronology; encode to + # codes so every comparison/sort respects it (round-20). + df[time] = df[time].cat.codes.astype(int) + self._validate_common_time_scale(df, time) + return self._fit_common_timing(df, outcome, unit, time, treatment, cluster, covariates) + from diff_diff.lwdid_staggered import fit_staggered + + results = fit_staggered(self, df, outcome, unit, time, first_treat, cluster, covariates) + if label_maps is not None: + _relabel_staggered_results(results, label_maps) + return results + + def _validate_common_time_scale(self, df: pd.DataFrame, time: str) -> None: + """Common-timing time-scale contract, SHARED by fit() and + get_transformation_diagnostics() (round-18 review: diagnostics + bypassed these checks and reached the transforms' raw float + conversion errors).""" + if self.rolling in ("detrend", "detrendq") and isinstance(df[time].dtype, pd.PeriodDtype): + # Period values cannot be cast to float for the unit trend + # design (review finding: validation accepted PeriodDtype + # but the transform raised a raw TypeError). datetime64 + # works (nanosecond ordinals). + raise ValueError( + f"rolling='{self.rolling}' does not support a Period " + f"time column on the common-timing path; convert with " + f".dt.to_timestamp() or encode the time column " + f"numerically." + ) + if self.rolling != "demean" and not ( + pd.api.types.is_numeric_dtype(df[time]) + or pd.api.types.is_datetime64_any_dtype(df[time]) + or isinstance(df[time].dtype, pd.PeriodDtype) + ): + # Campaign finding: detrend/demeanq/detrendq cast the time + # column to float for the trend/quarter design and raised a + # raw numpy conversion error on string times (while demean, + # which never touches the time values, succeeded). + raise ValueError( + f"rolling='{self.rolling}' requires a numeric or " + f"datetime/Period time column (the unit-specific trend/" + f"seasonal design uses the time values); column " + f"'{time}' has dtype {df[time].dtype}. Encode the time " + f"column numerically, or use rolling='demean'." + ) + if self.rolling == "demean" and not ( + pd.api.types.is_numeric_dtype(df[time]) + or pd.api.types.is_datetime64_any_dtype(df[time]) + or isinstance(df[time].dtype, pd.PeriodDtype) + ): + # Round-20 review: plain object labels sort LEXICOGRAPHICALLY + # ('Q10' < 'Q2'), silently corrupting the pre/post partition + # and event-time positions; an ordered categorical declares + # the chronology explicitly and is encoded to its codes. + dtype = df[time].dtype + if not (isinstance(dtype, pd.CategoricalDtype) and dtype.ordered): + raise ValueError( + f"rolling='demean' with a non-numeric, non-datetime " + f"time column requires an ORDERED categorical (the " + f"chronology cannot be inferred from labels - " + f"lexicographic order breaks at e.g. 'Q10' < 'Q2'). " + f"Column '{time}' has dtype {dtype}. Use " + f"pd.Categorical(values, categories=..., ordered=True) " + f"or encode the time column numerically." + ) + + def get_transformation_diagnostics( + self, + data: pd.DataFrame, + outcome: str, + unit: str, + time: str, + treatment: str, + first_treat: Optional[str] = None, + ) -> Dict[str, Any]: + """Run the transformation step and return diagnostics without full estimation. + + This is useful for inspecting pre-treatment fit quality before running + the full estimator. + + Parameters + ---------- + data : pd.DataFrame + Panel data. + outcome : str + Name of the outcome column. + unit : str + Name of the unit identifier column. + time : str + Name of the time period column. + treatment : str + Name of the treatment indicator column. + first_treat : str or None, default None + Name of the first-treatment-time (cohort) column, for + staggered designs. + + Returns + ------- + dict + Common timing: transformation diagnostics (see _transform_* + docstrings). Staggered: per-cohort diagnostics organized as + ``{'method': ..., 'design': 'staggered', 'by_cohort': {g: + diagnostics_g}}`` where each cohort g uses its own pre-period + definition ``time < g`` and the same unit subset as estimation + (cohort-g treated units plus the control superset implied by + ``control_group``). + """ + df = data.copy() + + # Same front-door validation as fit() (review round 3: diagnostics + # previously accepted designs fit() rejects - non-binary treatment, + # reserved-name collisions, duplicate panels, incoherent cohorts). + self._validate_inputs(df, outcome, unit, time, treatment, first_treat, None, None) + validate_binary(df[treatment].values, treatment) + + if first_treat is not None: + # Staggered: each cohort g has its own pre-period t < g, + # mirroring _transform_for_cohort in estimation. Datetime and + # Period panels use the same integer-position encoding as fit(). + df, time, first_treat, label_maps = _encode_staggered_time_scale(df, time, first_treat) + # Same cohort normalization as fit(): inf and beyond-window + # cohorts are recoded to never-treated here too, so the + # diagnostics iterate the same cohort set estimation uses. + df[first_treat], _, _ = _normalize_cohorts(df[first_treat], max_time=df[time].max()) + _check_treatment_design(df, unit, time, treatment, first_treat) + cohort_by_unit = df.drop_duplicates(subset=[unit], keep="first").set_index(unit)[ + first_treat + ] + never_mask = cohort_by_unit.isna() | (cohort_by_unit == 0) + never_units = cohort_by_unit.index[never_mask].to_list() + treated_cohorts = sorted( + value for value in pd.unique(df[first_treat]) if pd.notna(value) and value > 0 + ) + if not treated_cohorts: + # Round-20 review: an all-never-treated panel returned an + # empty {'by_cohort': {}} that read as successful + # diagnostics; fit_staggered rejects the same input. + raise ValueError("No treated cohorts found.") + by_cohort: Dict[Any, Dict[str, Any]] = {} + for g in treated_cohorts: + treated_units = cohort_by_unit.index[cohort_by_unit == g].to_list() + if self.control_group == "never_treated": + control_superset = never_units + else: + later = cohort_by_unit.index[cohort_by_unit > g].to_list() + control_superset = never_units + later + relevant_units = list(dict.fromkeys(treated_units + control_superset)) + cohort_frame = df.loc[df[unit].isin(relevant_units)].copy() + pre_mask = cohort_frame[time] < g + by_cohort[g] = self._run_transformation_diagnostics( + cohort_frame, outcome, unit, time, pre_mask + ) + if label_maps is not None: + cohort_labels = label_maps["cohort"] + by_cohort = {cohort_labels.get(g, g): value for g, value in by_cohort.items()} + return { + "method": self.rolling, + "design": "staggered", + "by_cohort": by_cohort, + } + + # Common timing: partition at the single onset S, same as + # _fit_common_timing (round-9 review: the per-period max(D) rule + # here still classified a controls-only post period as pre). + if isinstance(df[time].dtype, pd.CategoricalDtype) and df[time].dtype.ordered: + df[time] = df[time].cat.codes.astype(int) + self._validate_common_time_scale(df, time) + _check_treatment_design(df, unit, time, treatment, None) + treated_times = df.loc[df[treatment] == 1, time] + if len(treated_times) == 0: + raise ValueError( + "No post-treatment periods found. At least one period " + "with some treatment=1 is required." + ) + pre_mask = df[time] < treated_times.min() + return self._run_transformation_diagnostics(df, outcome, unit, time, pre_mask) + + def _run_transformation_diagnostics( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + pre_mask: Union[pd.Series, np.ndarray], + ) -> Dict[str, Any]: + """Dispatch to the configured transformation with diagnostics enabled.""" + if self.rolling == "demean": + _, diagnostics = self._transform_demean( + df, outcome, unit, pre_mask, return_diagnostics=True + ) + elif self.rolling == "detrend": + _, diagnostics = self._transform_detrend( + df, outcome, unit, time, pre_mask, return_diagnostics=True + ) + elif self.rolling == "demeanq": + _, diagnostics = self._transform_demeanq( + df, outcome, unit, time, pre_mask, return_diagnostics=True + ) + elif self.rolling == "detrendq": + _, diagnostics = self._transform_detrendq( + df, outcome, unit, time, pre_mask, return_diagnostics=True + ) + else: + _, diagnostics = self._transform_detrend( + df, outcome, unit, time, pre_mask, return_diagnostics=True + ) + + return diagnostics + + def _validate_inputs( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + treatment: str, + cohort: Optional[str], + cluster: Optional[str], + controls: Optional[List[str]], + ) -> None: + """Validate that all required columns exist and data is valid. + + Parameters + ---------- + df : pd.DataFrame + The input dataframe. + outcome, unit, time, treatment : str + Required column names. + cohort, cluster : str or None + Optional column names. + controls : list of str or None + Optional control variable column names. + + Raises + ------ + ValueError + If any specified column is not in the dataframe. + """ + required_cols = [outcome, unit, time, treatment] + if cohort is not None: + required_cols.append(cohort) + if cluster is not None: + required_cols.append(cluster) + if controls: + required_cols.extend(controls) + + # Internal working columns are written into the estimation frames; + # a user role column bearing one of these names is silently + # overwritten (review round 2: cluster='_treat' reported the + # cluster labels' coefficient as the ATT). + reserved = _RESERVED_INTERNAL_COLUMNS.intersection(required_cols) + if reserved: + raise ValueError( + f"Column name(s) {sorted(reserved)} are reserved for LWDiD " + f"internal use and cannot be supplied as outcome, unit, " + f"time, treatment, first_treat, cluster, or covariate " + f"columns. Rename the column(s) before fitting." + ) + core_roles = { + "outcome": outcome, + "unit": unit, + "time": time, + "treatment": treatment, + } + if cohort is not None: + core_roles["first_treat"] = cohort + seen: Dict[str, str] = {} + for role, name in core_roles.items(): + if name in seen: + raise ValueError( + f"Column '{name}' was supplied as both '{seen[name]}' and " + f"'{role}'; each role requires a distinct column." + ) + seen[name] = role + overlap = set(controls or []).intersection(core_roles.values()) + if overlap: + raise ValueError( + f"Covariate column(s) {sorted(overlap)} are already supplied " + f"as outcome/unit/time/treatment/first_treat columns." + ) + if controls and len(set(controls)) != len(controls): + duplicated = sorted({c for c in controls if controls.count(c) > 1}) + raise ValueError( + f"Covariate list contains duplicate column(s): {duplicated} " + f"(a repeated covariate makes the design matrix rank-" + f"deficient by construction)." + ) + + missing = [c for c in required_cols if c not in df.columns] + if missing: + raise ValueError(f"Columns not found in data: {missing}") + + # Check for NaN in key columns + for col in [outcome, unit, time, treatment]: + if df[col].isna().any(): + raise ValueError( + f"Column '{col}' contains missing values. " + f"Please handle missing data before fitting." + ) + # Round-9 review: Inf outcomes passed the NaN check and were + # silently np.isfinite-filtered inside staggered cells (changing + # the estimation sample with no warning); non-numeric outcomes + # crashed with raw conversion errors. + try: + outcome_values = df[outcome].to_numpy(dtype=float) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Outcome column '{outcome}' is not numeric (dtype " + f"{df[outcome].dtype}); encode it numerically before fitting." + ) from exc + n_nonfinite_y = int((~np.isfinite(outcome_values)).sum()) + if n_nonfinite_y > 0: + raise ValueError( + f"Outcome column '{outcome}' contains {n_nonfinite_y} " + f"non-finite value(s) (Inf). LWDiD does not silently drop " + f"outcome rows; remove or recode them before fitting." + ) + # Numeric TIME values must be finite too (round-14 review: +/-Inf + # passed the NaN check and reached event-time arithmetic, raising + # a raw OverflowError). Datetime/Period/ordered-label columns are + # untouched. + if pd.api.types.is_numeric_dtype(df[time]): + time_values = df[time].to_numpy(dtype=float) + n_nonfinite_t = int((~np.isfinite(time_values)).sum()) + if n_nonfinite_t > 0: + raise ValueError( + f"Time column '{time}' contains {n_nonfinite_t} " + f"non-finite value(s) (Inf). Time periods must be " + f"finite; remove or recode them before fitting." + ) + + # Check panel structure: each unit-time pair should be unique + duplicates = df.duplicated(subset=[unit, time], keep=False) + if duplicates.any(): + n_dup = duplicates.sum() + raise ValueError( + f"Panel is not balanced: {n_dup} duplicate " + f"unit-time observations found. Each (unit, time) " + f"pair must be unique." + ) + + # Panel balance check + obs_per_unit = df.groupby(unit)[time].nunique() + if obs_per_unit.nunique() > 1: + n_short = (obs_per_unit < obs_per_unit.max()).sum() + warnings.warn( + f"Unbalanced panel: {n_short} of {obs_per_unit.shape[0]} units have " + f"fewer than {obs_per_unit.max()} time periods. LWDiD assumes balanced " + "panels for optimal performance.", + UserWarning, + stacklevel=2, + ) + + # Unit-constancy contracts, shared by BOTH dispatch paths: LWDiD + # collapses the panel to one row per unit, reading unit-level + # covariate and cluster values. A time-varying column would make + # the estimate depend on the row order of the input frame (and, + # in staggered designs, silently pull post-treatment covariate + # values into the cohort-time cells), so it is rejected here. + for column in controls or []: + n_missing = int(df[column].isna().sum()) + if n_missing > 0: + # Campaign finding: NaN covariates silently dropped units + # on the cell paths while poisoning the common-timing OLS + # into a NaN ATT - unify by rejecting up front. + raise ValueError( + f"Covariate '{column}' contains {n_missing} missing " + f"value(s). LWDiD does not silently drop or impute " + f"covariate rows; remove or impute them before fitting." + ) + # Round-8 review: Inf passed the NaN check and was silently + # filtered per staggered cell (changing the estimation sample) + # or crashed inside the solver on the common-timing path; + # non-numeric covariates crashed with a raw conversion error. + try: + column_values = df[column].to_numpy(dtype=float) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Covariate '{column}' is not numeric (dtype " + f"{df[column].dtype}); encode it numerically before " + f"fitting." + ) from exc + n_nonfinite = int((~np.isfinite(column_values)).sum()) + if n_nonfinite > 0: + raise ValueError( + f"Covariate '{column}' contains {n_nonfinite} " + f"non-finite value(s) (Inf). LWDiD does not silently " + f"drop covariate rows; remove or recode them before " + f"fitting." + ) + varying = df.groupby(unit)[column].nunique(dropna=False) + if (varying > 1).any(): + raise ValueError( + f"Covariate '{column}' is not unit-constant; time-varying " + "covariates are not supported by LWDiD. Aggregate the " + "column to one value per unit (e.g. its pre-treatment " + "value) before fitting." + ) + if cluster is not None and cluster != unit: + n_missing = int(df[cluster].isna().sum()) + if n_missing > 0: + raise ValueError( + f"Cluster column '{cluster}' contains {n_missing} missing " + f"value(s); every observation must belong to a cluster." + ) + varying = df.groupby(unit)[cluster].nunique(dropna=False) + if (varying > 1).any(): + raise ValueError( + f"Cluster column '{cluster}' is not unit-constant; each " + "unit must belong to exactly one cluster. Assign one " + "cluster value per unit before fitting." + ) + + def _fit_common_timing( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + treatment: str, + cluster: Optional[str], + controls: List[str], + ) -> LWDiDResults: + """Estimate ATT under common treatment timing. + + All treated units adopt treatment at the same time period. + + Parameters + ---------- + df : pd.DataFrame + Panel data. + outcome : str + Outcome variable column. + unit : str + Unit identifier column. + time : str + Time period column. + treatment : str + Binary treatment indicator column. + cluster : str or None + Cluster variable for cluster-robust SEs. + controls : list of str + Control variable columns. + + Returns + ------- + LWDiDResults + Estimation results. + """ + # Treatment-design validation (absorbing + common timing) is + # performed by _check_treatment_design in fit() before dispatch. + + # Step 1: Partition the calendar support at the single adoption + # period S (validated by _check_treatment_design): pre = t < S, + # post = t >= S. Round-8 review: the previous per-period + # `groupby(time)[treatment].max()` partition silently classified a + # post period with no observed TREATED rows (controls only) as + # pre-treatment, contaminating the rolling pre window and biasing + # a zero-effect trend panel to ATT ~ 0.75. + treated_times = df.loc[df[treatment] == 1, time] + if len(treated_times) == 0: + raise ValueError( + "No post-treatment periods found. At least one period " + "with some treatment=1 is required." + ) + onset_s = treated_times.min() + support = sorted(pd.unique(df[time])) + pre_periods = [t for t in support if t < onset_s] + post_periods = [t for t in support if t >= onset_s] + + if len(pre_periods) == 0: + raise ValueError( + "No pre-treatment periods found. At least one period " + "with all treatment=0 is required." + ) + if len(post_periods) == 0: + raise ValueError( + "No post-treatment periods found. At least one period " + "with some treatment=1 is required." + ) + + # Identify treated and control units + unit_ever_treated = df.groupby(unit)[treatment].max() + treated_units = unit_ever_treated[unit_ever_treated == 1].index.tolist() + control_units = unit_ever_treated[unit_ever_treated == 0].index.tolist() + treated_set = set(treated_units) + + if len(treated_units) == 0: + raise ValueError("No treated units found in the data.") + if len(control_units) == 0: + raise ValueError( + "No control units found. At least one never-treated " "unit is required." + ) + + # Step 2: Apply transformation + pre_mask = df[time].isin(pre_periods) + + if self.rolling == "demean": + df = self._transform_demean(df, outcome, unit, pre_mask) + elif self.rolling == "detrend": + df = self._transform_detrend(df, outcome, unit, time, pre_mask) + elif self.rolling == "demeanq": + df = self._transform_demeanq(df, outcome, unit, time, pre_mask) + elif self.rolling == "detrendq": + df = self._transform_detrendq(df, outcome, unit, time, pre_mask) + else: + df = self._transform_detrend(df, outcome, unit, time, pre_mask) + + # Per-period event-study surface (LW 2026 eq. 2.20): each post + # period is one small cross-sectional regression on the transformed + # outcome, so the surface is populated at fit time, exactly like + # the staggered path. + ( + event_effects, + reference_periods, + event_vcov, + event_vcov_index, + event_study_df, + cband_method, + cband_crit_value, + cband_n_bootstrap, + ) = self._common_timing_event_study( + df, unit, time, cluster, controls, post_periods, treated_set + ) + + # Step 3: Take post-treatment cross-section of transformed outcomes + # Average transformed outcome over the FIXED post window per unit. + # The estimand is the paper's fixed-window post average (LW 2026, + # denominator T - S + 1): units not observing EVERY post period are + # dropped as complete cases with a warning (round-6 review: the + # previous per-unit mean over whichever post periods a unit + # observed let calendar composition masquerade as treatment effect + # - a zero-effect panel where treated units observed one extra + # post period reported that period's trend as ATT). + post_mask = df[time].isin(post_periods) + post_df = df.loc[post_mask].copy() + + # Reindex over EVERY panel unit so zero-post-row units are counted + # as incomplete too (round-10 review: they were absent from the + # counts, silently vanished in the merge, and the documented + # fixed-window drop warning never fired for them). + all_panel_units = pd.Index(pd.unique(df[unit])) + post_counts = ( + post_df.loc[np.isfinite(post_df["_ydot"])] + .groupby(unit)["_ydot"] + .size() + .reindex(all_panel_units, fill_value=0) + ) + complete_units = set(post_counts.index[post_counts == len(post_periods)]) + n_incomplete = int((post_counts < len(post_periods)).sum()) + if n_incomplete > 0: + warnings.warn( + f"LWDiD: {n_incomplete} unit(s) dropped from the collapsed " + f"cross-section for lacking a finite transformed outcome in " + f"every post-treatment period (missing rows or failed " + f"transformation): the headline ATT is the fixed-window " + f"post average over complete cases. See " + f"docs/methodology/REGISTRY.md (LWDiD).", + UserWarning, + stacklevel=2, + ) + post_df = post_df.loc[post_df[unit].isin(complete_units)] + + # Compute unit-level average of transformed outcome in post periods + unit_post_avg = post_df.groupby(unit)["_ydot"].mean().reset_index() + unit_post_avg.columns = [unit, "_ydot_avg"] + + # Build cross-sectional dataset + # Take first observation per unit for controls + cs_df = df.drop_duplicates(subset=[unit], keep="first")[[unit] + controls].copy() + # Treatment indicator: 1 if unit is ever-treated + cs_df["_treat"] = cs_df[unit].isin(treated_set).astype(float) + if cluster is not None: + # Get cluster from original data + if cluster == unit: + cs_df[cluster] = cs_df[unit] + else: + cluster_map = df.drop_duplicates(subset=[unit], keep="first").set_index(unit)[ + cluster + ] + cs_df[cluster] = cs_df[unit].map(cluster_map) + + cs_df = cs_df.merge(unit_post_avg, on=unit, how="inner") + + # After merge, drop units whose transformation produced NaN + n_before_drop = len(cs_df) + cs_df = cs_df.dropna(subset=["_ydot_avg"]) + n_dropped = n_before_drop - len(cs_df) + if n_dropped > 0 and len(cs_df) > 0: + warnings.warn( + f"LWDiD: {n_dropped} unit(s) dropped due to NaN transformed outcomes " + f"(insufficient pre-treatment periods for '{self.rolling}' transformation).", + UserWarning, + stacklevel=2, + ) + if len(cs_df) == 0: + nan = float("nan") + warnings.warn( + f"All units have NaN transformed outcomes for rolling='{self.rolling}'. " + "Likely insufficient pre-treatment periods. Cannot estimate ATT.", + UserWarning, + stacklevel=2, + ) + return LWDiDResults( + att=nan, + se=nan, + t_stat=nan, + p_value=nan, + conf_int=(nan, nan), + n_obs=0, + n_treated=0, + n_control=0, + rolling=self.rolling, + estimation_method=self.estimation_method, + vcov_type=self.vcov_type, + control_group=self.control_group, + n_bootstrap=self.n_bootstrap, + seed=self.seed, + cluster_name=cluster, + pscore_trim=( + self.pscore_trim if self.estimation_method in ("ipw", "dr", "psm") else None + ), + psm_config=( + { + "pscore_trim": self.pscore_trim, + "n_neighbors": self.n_neighbors, + "caliper": self.caliper, + "with_replacement": self.with_replacement, + } + if self.estimation_method == "psm" + else None + ), + alpha=self.alpha, + event_study_effects=event_effects, + event_study_vcov=event_vcov, + event_study_vcov_index=event_vcov_index, + event_study_df=event_study_df, + reference_periods=reference_periods, + cband_method=cband_method, + cband_crit_value=cband_crit_value, + cband_n_bootstrap=cband_n_bootstrap, + ) + + # Step 4: Estimate ATT + y = cs_df["_ydot_avg"].values.astype(np.float64) + treat = cs_df["_treat"].values.astype(np.float64) + n_obs = len(y) + n_treated = int(treat.sum()) + n_control = n_obs - n_treated + + if n_treated == 0 or n_control == 0: + # Round-6 review: the raw-panel arm counts are checked before + # the transformation, but complete-case/NaN drops can empty an + # arm - the pre-fix code dispatched a one-arm design into the + # estimators (rank warnings, NaN arithmetic, and the IPW family + # entered propensity fitting with an empty group). + raise ValueError( + f"After the transformation and complete-case drops, the " + f"collapsed cross-section has {n_treated} treated and " + f"{n_control} control unit(s); estimation requires at " + f"least one of each. Likely insufficient pre-treatment " + f"periods or incomplete post-period coverage for one arm " + f"under rolling='{self.rolling}'." + ) + + # Guard: if transformation produced all-NaN outcomes, return NaN result + if np.all(np.isnan(y)): + warnings.warn( + f"All transformed outcomes are NaN (likely insufficient " + f"pre-treatment periods for '{self.rolling}' transformation). " + f"Cannot estimate ATT.", + UserWarning, + stacklevel=2, + ) + nan = float("nan") + return LWDiDResults( + att=nan, + se=nan, + t_stat=nan, + p_value=nan, + conf_int=(nan, nan), + n_obs=n_obs, + n_treated=n_treated, + n_control=n_control, + rolling=self.rolling, + estimation_method=self.estimation_method, + vcov_type=self.vcov_type, + control_group=self.control_group, + n_bootstrap=self.n_bootstrap, + seed=self.seed, + cluster_name=cluster, + pscore_trim=( + self.pscore_trim if self.estimation_method in ("ipw", "dr", "psm") else None + ), + psm_config=( + { + "pscore_trim": self.pscore_trim, + "n_neighbors": self.n_neighbors, + "caliper": self.caliper, + "with_replacement": self.with_replacement, + } + if self.estimation_method == "psm" + else None + ), + alpha=self.alpha, + event_study_effects=event_effects, + event_study_vcov=event_vcov, + event_study_vcov_index=event_vcov_index, + event_study_df=event_study_df, + reference_periods=reference_periods, + cband_method=cband_method, + cband_crit_value=cband_crit_value, + cband_n_bootstrap=cband_n_bootstrap, + ) + + # Build controls matrix + controls_matrix = None + if controls: + controls_matrix = cs_df[controls].values.astype(np.float64) + + # Get cluster ids (clustered inference activates via the cluster= + # constructor parameter) + cluster_ids = None + collapsed_single_cluster = False + if cluster is not None: + cluster_ids = cs_df[cluster].values + if len(np.unique(cluster_ids)) < 2: + collapsed_single_cluster = True + # The NaN-transformation dropna can reduce the collapsed + # cross-section below 2 clusters even when the raw panel + # passed the fit-level guard - fail closed rather than let + # a single-cluster CR1 SE through on roundoff. + warnings.warn( + "LWDiD: after transformation drops, the collapsed " + "cross-section contains fewer than 2 clusters; " + "cluster-robust inference is not identified. The point " + "estimate is retained with NaN inference.", + UserWarning, + stacklevel=2, + ) + cluster_ids = None # estimate the POINT unclustered + + # Estimate + att, se, coefs, vcov, n_params, _ = self._dispatch_estimator( + y, treat, controls_matrix, cluster_ids, n_obs + ) + if collapsed_single_cluster: + se = np.nan # fail-closed (warned above); point retained + + # Step 5: Compute inference + # n_params is the fitted design's parameter count, so the residual + # df is design-coherent (LW 2026 Section 2): T_{N-2} without + # controls, T_{N-K-2} for the plain design and T_{N-2K-2} when the + # treatment-covariate interaction is active. + if self.estimation_method == "psm": + # PSM inference is NaN by contract and uses no residual df + # (round-20 review: the exact-OLS residual-df guard below + # rejected valid point-only matching fits whose nominal + # propensity width exhausted an OLS df count PSM never uses). + df_dof = None + elif n_obs < 3 or n_obs - n_params <= 0: + # Registry small-sample guards (N >= 3; N > K + 2 with controls): + # coercing the residual df to 1 fabricated exact inference on + # invalid designs, and N=2 reached sse/(n-k) division by zero + # (review finding). + raise ValueError( + f"Invalid exact-inference design: {n_obs} collapsed " + f"observation(s) with {n_params} fitted parameter(s). LWDiD " + f"requires at least 3 cross-sectional units and a positive " + f"residual df (N > K + 2 with controls)." + ) + else: + df_dof = n_obs - n_params + + # Issue 3: Cluster-robust inference uses df = G-1 + if cluster_ids is not None: + df_dof = max(int(len(np.unique(cluster_ids))) - 1, 1) + elif collapsed_single_cluster: + df_dof = 0 # safe_inference fails the tuple closed + + # Scale-equivariant degenerate-SE guard, same rule as the + # staggered/event surfaces (round-21 review: an exactly fitted + # panel produced se ~ 1e-16 and t ~ 1e16 on the common headline). + from diff_diff.lwdid_staggered import _guard_standard_error + + se = _guard_standard_error(att, se, scale=float(np.max(np.abs(y))) if len(y) else 0.0) + + t_stat, p_value, conf_int = safe_inference(att, se, alpha=self.alpha, df=df_dof) + + # Step 6: Bootstrap if requested + inference_basis = None + if self.n_bootstrap > 0 and collapsed_single_cluster: + # Round-5 review: the unconditional bootstrap call overwrote + # the single-effective-cluster fail-closed NaN inference with a + # finite (near-zero) SE built from the raw cluster map. The + # fail-closed state wins; the earlier warning already fired. + warnings.warn( + "LWDiD: bootstrap skipped - fewer than 2 effective clusters " + "survive the transformation, so clustered inference is not " + "identified (the NaN inference tuple is retained).", + UserWarning, + stacklevel=2, + ) + elif self.n_bootstrap > 0: + att, se, t_stat, p_value, conf_int, df_dof = self._bootstrap( + df, + outcome, + unit, + time, + treatment, + cluster, + controls, + pre_periods, + post_periods, + treated_units, + control_units, + ) + # Provenance (review round 3): the headline se/p/CI now come + # from the resampling bootstrap while params/vcov remain the + # analytical regression quantities - record which family backs + # the headline so consumers (and summary()) can tell. + inference_basis = "cluster_bootstrap" if cluster is not None else "unit_bootstrap" + + result = LWDiDResults( + inference_basis=inference_basis, + att=att, + se=se, + t_stat=t_stat, + p_value=p_value, + conf_int=conf_int, + n_obs=n_obs, + n_treated=n_treated, + n_control=n_control, + rolling=self.rolling, + estimation_method=self.estimation_method, + vcov_type=self.vcov_type, + alpha=self.alpha, + cluster_name=cluster if cluster_ids is not None else None, + control_group=self.control_group, + n_bootstrap=self.n_bootstrap, + seed=self.seed, + pscore_trim=( + self.pscore_trim if self.estimation_method in ("ipw", "dr", "psm") else None + ), + psm_config=( + { + "pscore_trim": self.pscore_trim, + "n_neighbors": self.n_neighbors, + "caliper": self.caliper, + "with_replacement": self.with_replacement, + } + if self.estimation_method == "psm" + else None + ), + n_clusters=int(len(np.unique(cluster_ids))) if cluster_ids is not None else None, + cohort_effects=None, + params=coefs, + vcov=vcov, + df_inference=df_dof, + event_study_effects=event_effects, + event_study_vcov=event_vcov, + event_study_vcov_index=event_vcov_index, + event_study_df=event_study_df, + reference_periods=reference_periods, + cband_method=cband_method, + cband_crit_value=cband_crit_value, + cband_n_bootstrap=cband_n_bootstrap, + ) + + # Fit-time replay spec for the post-fit advanced-inference methods + # (round-5 review: they previously accepted arbitrary caller arrays + # and a non-interacted design, caching p-values for a DIFFERENT + # estimand than .att on covariate-unbalanced RA fits). + object.__setattr__( + result, + "_replay_spec", + { + "y": y.copy(), + "treatment": treat.copy(), + "controls": controls_matrix.copy() if controls_matrix is not None else None, + "cluster_ids": cluster_ids.copy() if cluster_ids is not None else None, + }, + ) + + # Final safety net: warn if result has NaN ATT + if np.isnan(result.att): + warnings.warn( + f"LWDiD estimation returned NaN ATT. This typically indicates " + f"insufficient data for the '{self.rolling}' transformation or " + f"numerical issues in estimation. Check your data structure and " + f"consider using a simpler transformation (e.g., rolling='demean').", + UserWarning, + stacklevel=2, + ) + + return result + + def _common_timing_event_study( + self, + df: pd.DataFrame, + unit: str, + time: str, + cluster: Optional[str], + controls: List[str], + post_periods: List[Any], + treated_set: set, + ) -> Tuple[ + Dict[int, Dict[str, Any]], + Tuple[int, ...], + Optional[np.ndarray], + Optional[np.ndarray], + Dict[int, Any], + Optional[str], + Optional[float], + Optional[int], + ]: + """Per-period event-study surface for a common-timing fit. + + LW (2026) eq. (2.20): after the rolling transformation, the effect + for post period t is the coefficient on D in the cross-sectional + regression of the transformed outcome at t -- numerically identical + to a standard DiD on the subset panel {1, ..., S-1, t}. Each post + period is therefore one small regression run through the same + ``_dispatch_estimator`` path as the overall ATT, and the surface + follows the staggered storage contract (integer event-time keys, + position-difference convention). + + Reference anchors are the transformation's nominal anchors + (``-1`` for demean/demeanq, ``-2, -1`` for detrend/detrendq) + restricted to genuinely observed relative times: an unobserved + anchor is never synthesized. + + Returns + ------- + tuple + ``(event_effects, reference_periods, event_vcov, + event_vcov_index, event_study_df, cband_method, + cband_crit_value, cband_n_bootstrap)``. + """ + from diff_diff.lwdid_staggered import ( + _guard_standard_error, + compute_event_study_bands, + ) + + # Event-time convention, shared with the staggered path (round-9 + # review: the interfaces previously disagreed on gapped numeric + # calendars - common used ordered-support positions while + # staggered used arithmetic t - g): NUMERIC calendars use the + # Registry's arithmetic r = t - S (validated integral so distinct + # horizons can never merge under the integer storage keys); + # datetime/Period calendars use position differences (matching + # _encode_staggered_time_scale, which encodes them to positions + # before the staggered machinery runs). + all_times = sorted(pd.unique(df[time])) + onset_s = min(post_periods) + if not pd.api.types.is_numeric_dtype(df[time]): + # datetime/Period (position-encoded on the staggered path) and + # ordered string labels (demean-only contract): position + # differences on the ordered support. + time_pos = {value: index for index, value in enumerate(all_times)} + g_pos = time_pos[onset_s] + relative_of = {t: int(time_pos[t] - g_pos) for t in all_times} + else: + relative_of = {} + for t in all_times: + rel = float(t) - float(onset_s) + if abs(rel - round(rel)) > 1e-9: + raise ValueError( + f"Event time t - S = {rel!r} for period {t!r} is not " + f"an integer: the event-study surface stores integer " + f"event-time keys and cannot represent fractional " + f"horizons without silently merging them. Encode the " + f"time column as consecutive integer periods or as " + f"datetime/Period values." + ) + relative_of[t] = int(round(rel)) + nominal_anchors = (-1,) if self.rolling in ("demean", "demeanq") else (-2, -1) + observed_relative = set(relative_of.values()) + reference_periods = tuple(r for r in nominal_anchors if r in observed_relative) + + unit_rows = df.drop_duplicates(subset=[unit], keep="first").set_index(unit) + all_units = unit_rows.index.to_list() + unit_to_index = {value: index for index, value in enumerate(all_units)} + global_cluster_ids = None + if cluster is not None: + if cluster == unit: + global_cluster_ids = unit_rows.index.to_numpy() + else: + global_cluster_ids = unit_rows.loc[all_units, cluster].to_numpy() + + event_effects: Dict[int, Dict[str, Any]] = {} + event_influence: Dict[int, np.ndarray] = {} + skipped: List[Tuple[int, str]] = [] + for t in post_periods: + relative_time = relative_of[t] + columns = [unit, "_ydot"] + controls + if cluster is not None and cluster not in columns: + columns.append(cluster) + cell = df.loc[df[time] == t, columns].drop_duplicates(subset=[unit], keep="first") + finite = np.isfinite(cell["_ydot"].to_numpy(dtype=float)) + if controls: + finite &= np.all(np.isfinite(cell[controls].to_numpy(dtype=float)), axis=1) + cell = cell.loc[finite].copy() + treatment_vec = cell[unit].isin(treated_set).to_numpy(dtype=float) + n_treated = int(treatment_vec.sum()) + n_control = int(len(treatment_vec) - n_treated) + if n_treated == 0 or n_control == 0: + skipped.append((relative_time, "zero_treated_control")) + continue + + y = cell["_ydot"].to_numpy(dtype=float) + controls_matrix = cell[controls].to_numpy(dtype=float) if controls else None + cluster_ids = None + single_cluster_period = False + if cluster is not None: + cluster_ids = cell[cluster].to_numpy() + if len(np.unique(cluster_ids)) < 2: + warnings.warn( + "LWDiD: a common-timing event-study period cell " + "contains fewer than 2 clusters; its cluster-robust " + "inference is not identified (point retained, " + "inference NaN).", + UserWarning, + stacklevel=2, + ) + single_cluster_period = True + cluster_ids = None # estimate the POINT unclustered + try: + att, se, _, _, n_params, influence = self._dispatch_estimator( + y, treatment_vec, controls_matrix, cluster_ids, len(cell) + ) + except ValueError as exc: + if "Invalid exact-inference design" in str(exc): + # Non-estimable period cell (Registry: NaN, not a + # mid-fit raise; only the OVERALL design raises). + skipped.append((relative_time, "insufficient_sample")) + continue + raise + if not np.isfinite(att): + skipped.append((relative_time, "non_finite_estimate")) + continue + + se = _guard_standard_error(att, se, scale=float(np.max(np.abs(y)))) + if single_cluster_period: + se = np.nan # fail-closed (warned above); point retained + influence = None + if cluster_ids is not None: + df_event = max(len(np.unique(cluster_ids)) - 1, 1) + else: + # Raw residual df: safe_inference fails the tuple closed + # when df <= 0 (no fabricated df=1 - review finding). + df_event = len(cell) - n_params + t_stat, p_value, conf_int = safe_inference(att, se, alpha=self.alpha, df=df_event) + event_effects[relative_time] = { + "effect": float(att), + "se": se, + "t_stat": t_stat, + "p_value": p_value, + "conf_int": conf_int, + "n_treated": n_treated, + "n_control": n_control, + "n_cells": 1, + "df": df_event, + } + if influence is not None and np.isfinite(se): + global_influence = np.zeros(len(all_units), dtype=float) + for local_index, unit_value in enumerate(cell[unit].to_list()): + global_influence[unit_to_index[unit_value]] = influence[local_index] + event_influence[relative_time] = global_influence + + if skipped: + preview = ", ".join(f"r={r}: {reason}" for r, reason in skipped[:6]) + suffix = "" if len(skipped) <= 6 else f"; plus {len(skipped) - 6} more" + warnings.warn( + f"LWDiD skipped {len(skipped)} per-period effect(s): {preview}{suffix}. " + "The event-study surface omits these event times.", + UserWarning, + stacklevel=2, + ) + + ( + event_vcov, + event_vcov_index, + cband_method, + cband_crit_value, + cband_n_bootstrap, + ) = compute_event_study_bands(self, event_effects, event_influence, global_cluster_ids) + event_study_df = { + label: value["df"] + for label, value in event_effects.items() + if value.get("df") is not None + } + return ( + event_effects, + reference_periods, + event_vcov, + event_vcov_index, + event_study_df, + cband_method, + cband_crit_value, + cband_n_bootstrap, + ) + + def _composite_regression_aggregation( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + cohort: str, + ) -> Tuple[float, float, int, int, int, float, Dict[Any, int]]: + """Compute tau_omega via composite outcome regression (LW 2026 Eq 7.18/7.19). + + For staggered designs, constructs a composite outcome vector: + - Treated units in cohort g: use their cohort's transformed outcome + - Never-treated units: weighted average of all cohort transformations + Then runs a single cross-sectional OLS: y_composite ~ [1, D_ever_treated] + + Parameters + ---------- + df : pd.DataFrame + Full panel data. + outcome : str + Outcome variable column. + unit : str + Unit identifier column. + time : str + Time period column. + cohort : str + Cohort (first treatment time) column. + + Returns + ------- + att : float + ATT from composite regression coefficient on D. + se : float + Classical OLS SE from composite regression. + dof : int + Degrees of freedom (n_complete_case_units - 2). + n_treated_dropped : int + Treated units dropped by the complete-case resolution. + n_controls_dropped : int + Control units dropped by the complete-case resolution. + """ + # Step 1: Identify cohorts and unit membership + fy = df.groupby(unit)[cohort].first() + cohorts = sorted([g for g in fy.unique() if g > 0 and not np.isnan(g)]) + n_treat = int((fy > 0).sum()) + + if n_treat == 0: + return np.nan, np.nan, 0, 0, 0, 0.0, {} + + # Step 2: For each cohort g, compute per-unit post-average transformed outcome + # using cohort g's pre-period for ALL units + ydot_by_cohort: Dict[Any, pd.Series] = {} + for g in cohorts: + # pre_mask: periods < g (i.e., time <= g-1) + pre_mask_g = df[time] < g + post_mask_g = df[time] >= g + + # Apply transformation to full dataset. The composite (tau_omega) + # estimand is defined for the plain demean/detrend transforms + # only; the routing gate in lwdid_staggered restricts rolling, + # and this raise keeps a future gate change from silently + # substituting a non-seasonal transform for a q-variant again + # (campaign finding: demeanq/detrendq were mapped to + # demean/detrend here, moving the point by ~8% silently). + if self.rolling == "demean": + df_transformed = self._transform_demean(df, outcome, unit, pre_mask_g) + elif self.rolling == "detrend": + df_transformed = self._transform_detrend(df, outcome, unit, time, pre_mask_g) + else: + raise ValueError( + f"Internal error: composite (tau_omega) aggregation is " + f"only defined for rolling in ('demean', 'detrend'); got " + f"{self.rolling!r}. The routing gate should not have " + f"dispatched here." + ) + + # Per-unit average of transformed outcome in post-periods + # (>= g). Completeness semantics (ADJUDICATED, pinned by the + # acceptance suite's frozen reference oracle): a unit + # contributes cohort g's component iff its OBSERVED post-g + # rows yield a finite average - partial post windows are + # averaged over the observed rows, symmetrically for treated + # and control units. A stricter every-period window rule was + # considered in review round 8 and NOT adopted: it changes the + # estimand the acceptance oracle pins; the unbalanced + # composition caveat is documented in REGISTRY. + post_data = df_transformed.loc[post_mask_g] # type: ignore[union-attr] + unit_avg_g = post_data.groupby(unit)["_ydot"].mean() + ydot_by_cohort[g] = unit_avg_g + + # Step 3: Complete-case resolution (deterministic, one-directional). + # Fixed cohort weights omega_g = N_g / N_treat are defined on the + # ESTIMATION sample: units that cannot contribute their required + # transformed outcomes are dropped with a warning (never silently + # zero-filled or asymmetrically reweighted by the OLS finite mask). + all_units = fy.index + + # 3.1: Treated units must have a finite post-window average for + # their OWN cohort (missing post rows or a NaN transform both + # count - the pre-fix code silently NaN'd these out of the OLS, + # implicitly reweighting the treated side). + surviving_treated: List[Any] = [] + n_treated_dropped = 0 + for u in all_units: + g_u = fy[u] + if not (g_u > 0): + continue + value = ydot_by_cohort[g_u].get(u, np.nan) + if np.isfinite(value): + surviving_treated.append(u) + else: + n_treated_dropped += 1 + if n_treated_dropped: + warnings.warn( + f"LWDiD tau_omega composite: dropped {n_treated_dropped} " + f"treated unit(s) with no finite post-window transformed " + f"outcome for their cohort (complete-case estimation; cohort " + f"weights are recomputed on the surviving sample).", + UserWarning, + stacklevel=3, + ) + + # 3.2: Recompute cohort masses on the surviving treated sample. + fy_surviving = fy.loc[surviving_treated] + cohort_sizes = {g: int((fy_surviving == g).sum()) for g in cohorts} + weighted_cohorts = [g for g in cohorts if cohort_sizes[g] > 0] + n_treat_cc = len(surviving_treated) + + # 3.3: Control units must observe every surviving-weight cohort's + # post window with a finite transformed outcome (the pre-fix code + # injected a literal 0.0 for missing entries, biasing the + # composite control mean toward zero). + control_units = [u for u in all_units if not (fy[u] > 0)] + surviving_controls: List[Any] = [] + n_controls_dropped = 0 + for u in control_units: + vals = [ydot_by_cohort[g].get(u, np.nan) for g in weighted_cohorts] + if vals and np.all(np.isfinite(vals)): + surviving_controls.append(u) + else: + n_controls_dropped += 1 + if n_controls_dropped: + warnings.warn( + f"LWDiD tau_omega composite: dropped {n_controls_dropped} " + f"control unit(s) not observing every treated cohort's post " + f"window with a finite transformed outcome (complete-case " + f"estimation with fixed cohort weights).", + UserWarning, + stacklevel=3, + ) + + # 3.4: Empty-arm fail-closed guard (the pre-drop n_treat check + # does not cover drops emptying an arm). + if n_treat_cc == 0 or not surviving_controls: + warnings.warn( + "LWDiD tau_omega composite: complete-case filtering left an " + "empty treated or control arm; the composite ATT and its " + "inference are NaN.", + UserWarning, + stacklevel=3, + ) + return np.nan, np.nan, 0, n_treated_dropped, n_controls_dropped, 0.0, dict(cohort_sizes) + + # Step 4: Assemble composite outcome vector on the complete-case + # sample (finite by construction). + included = surviving_treated + surviving_controls + n = len(included) + y_composite = np.empty(n, dtype=np.float64) + d_ever_treated = np.empty(n, dtype=np.float64) + for i, u in enumerate(included): + g_u = fy[u] + if g_u > 0: + y_composite[i] = float(ydot_by_cohort[g_u][u]) + d_ever_treated[i] = 1.0 + else: + weighted_sum = 0.0 + for g in weighted_cohorts: + w_g = cohort_sizes[g] / n_treat_cc + weighted_sum += w_g * float(ydot_by_cohort[g][u]) + y_composite[i] = weighted_sum + d_ever_treated[i] = 0.0 + + if n < 3: + return np.nan, np.nan, 0, n_treated_dropped, n_controls_dropped, 0.0, dict(cohort_sizes) + + # Step 5: Single OLS regression y_composite ~ [1, D] via the house + # linalg engine (classical SE from the same regression). + X = np.column_stack([np.ones(n, dtype=np.float64), d_ever_treated]) + coefs, _, vcov = solve_ols(X, y_composite, return_vcov=True, vcov_type="classical") + att = float(coefs[1]) + dof = n - 2 + if vcov is not None and np.isfinite(vcov[1, 1]): + se = float(np.sqrt(max(vcov[1, 1], 0.0))) + else: + se = np.nan + + # Data scale for the degenerate-SE guard (scale-equivariant + # roundoff reference - see _guard_standard_error). + y_scale = float(np.max(np.abs(y_composite))) if len(y_composite) else 0.0 + # Survivor cohort masses (round-12 review: the drops-route + # aggregation needs these - raw masses left dropped treated + # units in the cohort weights). + return att, se, dof, n_treated_dropped, n_controls_dropped, y_scale, dict(cohort_sizes) + + def _transform_demean( + self, + df: pd.DataFrame, + outcome_col: str, + unit_col: str, + pre_mask: Union[pd.Series, np.ndarray], + return_diagnostics: bool = False, + ) -> Union[pd.DataFrame, Tuple[pd.DataFrame, Dict[str, Any]]]: + """Apply unit-specific demeaning transformation. + + For each unit, compute the mean of the outcome in pre-treatment + periods, then subtract that mean from ALL periods (pre and post). + + Parameters + ---------- + df : pd.DataFrame + Panel data. + outcome_col : str + Name of the outcome column. + unit_col : str + Name of the unit identifier column. + pre_mask : Series or ndarray of bool + Boolean mask indicating pre-treatment observations. + return_diagnostics : bool, default False + If True, return (df, diagnostics) tuple instead of just df. + + Returns + ------- + pd.DataFrame or (pd.DataFrame, dict) + Input data with '_ydot' column containing demeaned outcomes. + If return_diagnostics=True, also returns diagnostics dict. + """ + df = df.copy() + + # Compute pre-treatment mean for each unit + pre_df = df.loc[pre_mask, [unit_col, outcome_col]] + pre_means = pre_df.groupby(unit_col)[outcome_col].mean() + + # Collect per-unit diagnostics if requested + per_unit: Dict[Any, Dict[str, Any]] = {} + if return_diagnostics: + pre_stds = pre_df.groupby(unit_col)[outcome_col].std() + pre_counts = pre_df.groupby(unit_col)[outcome_col].count() + post_mask_inv = ~pre_mask + post_df = df.loc[post_mask_inv, [unit_col, outcome_col]] + post_counts = post_df.groupby(unit_col)[outcome_col].count() + all_units = df[unit_col].unique() + for uid in all_units: + has_pre = uid in pre_means.index + info: Dict[str, Any] = { + "pre_mean": float(pre_means[uid]) if has_pre else float("nan"), + "pre_n_periods": int(pre_counts.get(uid, 0)), + "pre_std": float(pre_stds.get(uid, float("nan"))), + "post_n_periods": int(post_counts.get(uid, 0)), + "valid": has_pre, + } + per_unit[uid] = info + + # Map pre-means back to all observations + unit_means = df[unit_col].map(pre_means) + + # Check for units with no pre-treatment obs (shouldn't happen + # after validation, but guard defensively) + no_pre = unit_means.isna() + if no_pre.any(): + n_missing = df.loc[no_pre, unit_col].nunique() + warnings.warn( + f"{n_missing} unit(s) have no pre-treatment observations. " + f"Their transformed outcomes will be NaN.", + UserWarning, + stacklevel=2, + ) + + # Subtract pre-treatment mean from all periods + df["_ydot"] = df[outcome_col].values - unit_means.values + + if return_diagnostics: + valid_units = [uid for uid, info in per_unit.items() if info["valid"]] + n_valid = len(valid_units) + n_total = len(per_unit) + pre_period_counts = [per_unit[uid]["pre_n_periods"] for uid in valid_units] + diagnostics: Dict[str, Any] = { + "method": "demean", + "description": "\u0232_{i,pre} subtracted from all periods (Procedure 2.1, Eq 2.12)", + "per_unit": per_unit, + "summary": { + "n_units_total": n_total, + "n_units_valid": n_valid, + "n_units_dropped": n_total - n_valid, + "mean_pre_periods": ( + float(np.mean(pre_period_counts)) if pre_period_counts else 0.0 + ), + "min_pre_periods": int(np.min(pre_period_counts)) if pre_period_counts else 0, + "max_pre_periods": int(np.max(pre_period_counts)) if pre_period_counts else 0, + }, + } + return df, diagnostics + + return df + + def _transform_detrend( + self, + df: pd.DataFrame, + outcome_col: str, + unit_col: str, + time_col: str, + pre_mask: Union[pd.Series, np.ndarray], + return_diagnostics: bool = False, + ) -> Union[pd.DataFrame, Tuple[pd.DataFrame, Dict[str, Any]]]: + """Apply unit-specific linear detrending transformation. + + For each unit, fit y = alpha + beta*t on pre-treatment periods + using scipy.linalg.lstsq, then subtract the fitted trend from + ALL periods. + + Parameters + ---------- + df : pd.DataFrame + Panel data. + outcome_col : str + Name of the outcome column. + unit_col : str + Name of the unit identifier column. + time_col : str + Name of the time period column. + pre_mask : Series or ndarray of bool + Boolean mask indicating pre-treatment observations. + return_diagnostics : bool, default False + If True, return (df, diagnostics) tuple instead of just df. + + Returns + ------- + pd.DataFrame or (pd.DataFrame, dict) + Input data with '_ydot' column containing detrended outcomes. + If return_diagnostics=True, also returns diagnostics dict. + """ + df = df.copy() + df["_ydot"] = np.nan + + # Pre-extract numpy arrays to avoid repeated df.loc[] overhead + unit_arr = df[unit_col].values + time_arr = df[time_col].values.astype(np.float64) + y_arr = df[outcome_col].values.astype(np.float64) + pre_arr = pre_mask.values if hasattr(pre_mask, "values") else np.asarray(pre_mask) + + units = df[unit_col].unique() + per_unit: Dict[Any, Dict[str, Any]] = {} + ydot_out = np.full(len(df), np.nan) + + for uid in units: + mask_u = unit_arr == uid + idx_u = np.where(mask_u)[0] + t_u = time_arr[idx_u] + y_u = y_arr[idx_u] + pre_u = pre_arr[idx_u] + + # Pre-treatment data for this unit + pre_sel = pre_u.astype(bool) + n_pre = int(pre_sel.sum()) + + if n_pre < 2: + warnings.warn( + f"Unit {uid}: detrend requires at least 2 " + f"pre-treatment periods, found {n_pre}. " + f"Transformed outcome set to NaN.", + UserWarning, + stacklevel=2, + ) + if return_diagnostics: + per_unit[uid] = { + "alpha": float("nan"), + "beta": float("nan"), + "pre_n_periods": n_pre, + "residual_std": float("nan"), + "r_squared": float("nan"), + "valid": False, + } + continue + + # Extract pre-treatment time and outcome + t_pre = t_u[pre_sel] + y_pre = y_u[pre_sel] + + # Center time for numerical stability + t_mean = t_pre.mean() + t_pre_centered = t_pre - t_mean + + # Build design matrix [intercept, centered_time] + X_pre = np.column_stack( + [ + np.ones(n_pre, dtype=np.float64), + t_pre_centered, + ] + ) + + # Solve via scipy.linalg.lstsq + result = scipy_linalg.lstsq(X_pre, y_pre, cond=None) + coefs = result[0] # [alpha, beta] + + # Check for valid coefficients + if not np.all(np.isfinite(coefs)): + warnings.warn( + f"Unit {uid}: detrending produced non-finite " + f"coefficients. Transformed outcome set to NaN.", + UserWarning, + stacklevel=2, + ) + if return_diagnostics: + per_unit[uid] = { + "alpha": float("nan"), + "beta": float("nan"), + "pre_n_periods": n_pre, + "residual_std": float("nan"), + "r_squared": float("nan"), + "valid": False, + } + continue + + # Predict on ALL periods for this unit (using same centering) + t_all_centered = t_u - t_mean + y_hat = coefs[0] + coefs[1] * t_all_centered + + # Residuals = outcome - fitted trend + ydot_out[idx_u] = y_u - y_hat + + # Collect diagnostics for this unit + if return_diagnostics: + y_hat_pre = X_pre @ coefs + residuals_pre = y_pre - y_hat_pre + residual_std = float(np.std(residuals_pre, ddof=2)) if n_pre > 2 else float("nan") + ss_res = float(np.sum(residuals_pre**2)) + ss_tot = float(np.sum((y_pre - y_pre.mean()) ** 2)) + r_squared = 1.0 - ss_res / ss_tot if ss_tot > 0 else float("nan") + per_unit[uid] = { + "alpha": float(coefs[0]), + "beta": float(coefs[1]), + "pre_n_periods": n_pre, + "residual_std": residual_std, + "r_squared": r_squared, + "valid": True, + } + + df["_ydot"] = ydot_out + + if return_diagnostics: + valid_units = [uid for uid, info in per_unit.items() if info["valid"]] + n_valid = len(valid_units) + n_total = len(per_unit) + betas = [per_unit[uid]["beta"] for uid in valid_units] + r2s = [ + per_unit[uid]["r_squared"] + for uid in valid_units + if np.isfinite(per_unit[uid]["r_squared"]) + ] + diagnostics: Dict[str, Any] = { + "method": "detrend", + "description": "Y_{it} - (\u03b1\u0302_i + \u03b2\u0302_i * t) based on pre-treatment OLS (Procedure 3.1)", + "per_unit": per_unit, + "summary": { + "n_units_total": n_total, + "n_units_valid": n_valid, + "n_units_dropped": n_total - n_valid, + "mean_beta": float(np.mean(betas)) if betas else float("nan"), + "std_beta": float(np.std(betas)) if betas else float("nan"), + "mean_r_squared": float(np.mean(r2s)) if r2s else float("nan"), + }, + } + return df, diagnostics + + return df + + def _transform_demeanq( + self, + df: pd.DataFrame, + outcome_col: str, + unit_col: str, + time_col: str, + pre_mask: Union[pd.Series, np.ndarray], + return_diagnostics: bool = False, + ) -> Union[pd.DataFrame, Tuple[pd.DataFrame, Dict[str, Any]]]: + """Apply unit-specific seasonal (quarterly) demeaning transformation. + + For each unit, fit Y on [1, Q2, Q3, Q4] dummies using pre-treatment + periods only, then subtract fitted values from ALL periods. + Quarter is determined by time_col % 4. + + Parameters + ---------- + df : pd.DataFrame + Panel data. + outcome_col : str + Name of the outcome column. + unit_col : str + Name of the unit identifier column. + time_col : str + Name of the time period column. + pre_mask : Series or ndarray of bool + Boolean mask indicating pre-treatment observations. + return_diagnostics : bool, default False + If True, return (df, diagnostics) tuple instead of just df. + + Returns + ------- + pd.DataFrame or (pd.DataFrame, dict) + Input data with '_ydot' column containing seasonally-demeaned outcomes. + If return_diagnostics=True, also returns diagnostics dict. + """ + df = df.copy() + df["_ydot"] = np.nan + + # Determine quarter from time column (0-indexed modulo 4 → 1-4). + # Encoded staggered frames carry the CALENDAR season in + # _lwdid_season (the time column holds dense positions there, and + # (pos - 1) % 4 + 1 would relabel seasons after a calendar gap). + t_series = df[time_col] + if "_lwdid_season" in df.columns: + quarters = df["_lwdid_season"].to_numpy() + elif pd.api.types.is_datetime64_any_dtype(t_series): + quarters = t_series.dt.quarter.to_numpy() + elif hasattr(t_series.iloc[0], "quarter"): + quarters = np.array([v.quarter for v in t_series]) + else: + t_vals = t_series.to_numpy() + quarters = (t_vals.astype(np.int64) - 1) % 4 + 1 + + # Pre-extract numpy arrays to avoid repeated df.loc[] overhead + unit_arr = df[unit_col].values + y_arr = df[outcome_col].values.astype(np.float64) + pre_arr = pre_mask.values if hasattr(pre_mask, "values") else np.asarray(pre_mask) + + units = df[unit_col].unique() + per_unit: Dict[Any, Dict[str, Any]] = {} + ydot_out = np.full(len(df), np.nan) + + for uid in units: + mask_u = unit_arr == uid + idx_u = np.where(mask_u)[0] + y_u = y_arr[idx_u] + q_u = quarters[idx_u] + pre_u = pre_arr[idx_u].astype(bool) + + # Pre-treatment data for this unit + n_pre = int(pre_u.sum()) + + # Need at least as many pre-obs as parameters (intercept + up to 3 dummies) + q_pre = q_u[pre_u] + observed_seasons = sorted(np.unique(q_pre)) + n_params = len(observed_seasons) # = 1 intercept + (n_seasons-1) dummies + + if n_pre < n_params: + warnings.warn( + f"Unit {uid}: demeanq requires at least as many pre-treatment " + f"observations as seasonal parameters ({n_params}), " + f"found {n_pre}. Transformed outcome set to NaN.", + UserWarning, + stacklevel=2, + ) + if return_diagnostics: + per_unit[uid] = { + "intercept": float("nan"), + "seasonal_effects": {}, + "pre_n_periods": n_pre, + "valid": False, + } + continue + + # Build seasonal dummy design matrix for pre-treatment + y_pre = y_u[pre_u] + + # Create dummies: drop first category (reference) + X_pre_parts = [np.ones(n_pre, dtype=np.float64)] + for s in observed_seasons[1:]: + X_pre_parts.append((q_pre == s).astype(np.float64)) + X_pre = np.column_stack(X_pre_parts) + + # Solve via scipy.linalg.lstsq + result = scipy_linalg.lstsq(X_pre, y_pre, cond=None) + coefs = result[0] + + if not np.all(np.isfinite(coefs)): + warnings.warn( + f"Unit {uid}: demeanq produced non-finite " + f"coefficients. Transformed outcome set to NaN.", + UserWarning, + stacklevel=2, + ) + if return_diagnostics: + per_unit[uid] = { + "intercept": float("nan"), + "seasonal_effects": {}, + "pre_n_periods": n_pre, + "valid": False, + } + continue + + # Season coverage: a quarter never observed in the unit's + # pre-period has no estimated effect; predicting it at the + # reference-season level is a silent out-of-support + # extrapolation (campaign finding) -> warn + NaN the unit. + unobserved_seasons = sorted(set(q_u.tolist()) - set(observed_seasons)) + if unobserved_seasons: + warnings.warn( + f"Unit {uid}: demeanq cannot predict quarter(s) " + f"{unobserved_seasons} that never appear in the unit's " + f"pre-treatment periods. Transformed outcome set to NaN.", + UserWarning, + stacklevel=2, + ) + if return_diagnostics: + per_unit[uid] = { + "intercept": float("nan"), + "seasonal_effects": {}, + "pre_n_periods": n_pre, + "valid": False, + } + continue + + # Predict on ALL periods for this unit + n_all = len(q_u) + X_all_parts = [np.ones(n_all, dtype=np.float64)] + for s in observed_seasons[1:]: + X_all_parts.append((q_u == s).astype(np.float64)) + X_all = np.column_stack(X_all_parts) + y_hat = X_all @ coefs + + # Residuals + ydot_out[idx_u] = y_u - y_hat + + # Collect diagnostics for this unit + if return_diagnostics: + seasonal_effects = { + int(s): float(coefs[idx + 1]) for idx, s in enumerate(observed_seasons[1:]) + } + per_unit[uid] = { + "intercept": float(coefs[0]), + "seasonal_effects": seasonal_effects, + "pre_n_periods": n_pre, + "valid": True, + } + + df["_ydot"] = ydot_out + + if return_diagnostics: + valid_units = [uid for uid, info in per_unit.items() if info["valid"]] + n_valid = len(valid_units) + n_total = len(per_unit) + diagnostics: Dict[str, Any] = { + "method": "demeanq", + "description": "Remove unit-specific seasonal (quarterly) fixed effects from pre-treatment", + "per_unit": per_unit, + "summary": { + "n_units_total": n_total, + "n_units_valid": n_valid, + "n_units_dropped": n_total - n_valid, + }, + } + return df, diagnostics + + return df + + def _transform_detrendq( + self, + df: pd.DataFrame, + outcome_col: str, + unit_col: str, + time_col: str, + pre_mask: Union[pd.Series, np.ndarray], + return_diagnostics: bool = False, + ) -> Union[pd.DataFrame, Tuple[pd.DataFrame, Dict[str, Any]]]: + """Apply unit-specific linear detrending with seasonal adjustment. + + For each unit, fit Y on [1, t, Q2, Q3, Q4] using pre-treatment + periods only, then subtract fitted values from ALL periods. + Quarter is determined by time_col % 4. + + Parameters + ---------- + df : pd.DataFrame + Panel data. + outcome_col : str + Name of the outcome column. + unit_col : str + Name of the unit identifier column. + time_col : str + Name of the time period column. + pre_mask : Series or ndarray of bool + Boolean mask indicating pre-treatment observations. + return_diagnostics : bool, default False + If True, return (df, diagnostics) tuple instead of just df. + + Returns + ------- + pd.DataFrame or (pd.DataFrame, dict) + Input data with '_ydot' column containing detrended+seasonally-adjusted outcomes. + If return_diagnostics=True, also returns diagnostics dict. + """ + df = df.copy() + df["_ydot"] = np.nan + + # Determine quarter from time column. Encoded staggered frames + # carry the CALENDAR season in _lwdid_season (see _transform_demeanq). + t_series = df[time_col] + if "_lwdid_season" in df.columns: + quarters = df["_lwdid_season"].to_numpy() + elif pd.api.types.is_datetime64_any_dtype(t_series): + quarters = t_series.dt.quarter.to_numpy() + elif hasattr(t_series.iloc[0], "quarter"): + quarters = np.array([v.quarter for v in t_series]) + else: + t_vals = t_series.to_numpy() + quarters = (t_vals.astype(np.int64) - 1) % 4 + 1 + + # Pre-extract numpy arrays to avoid repeated df.loc[] overhead + unit_arr = df[unit_col].values + time_arr = df[time_col].values.astype(np.float64) + y_arr = df[outcome_col].values.astype(np.float64) + pre_arr = pre_mask.values if hasattr(pre_mask, "values") else np.asarray(pre_mask) + + units = df[unit_col].unique() + per_unit: Dict[Any, Dict[str, Any]] = {} + ydot_out = np.full(len(df), np.nan) + + for uid in units: + mask_u = unit_arr == uid + idx_u = np.where(mask_u)[0] + t_u = time_arr[idx_u] + y_u = y_arr[idx_u] + q_u = quarters[idx_u] + pre_u = pre_arr[idx_u].astype(bool) + + # Pre-treatment data for this unit + n_pre = int(pre_u.sum()) + + if n_pre < 2: + warnings.warn( + f"Unit {uid}: detrendq requires at least 2 " + f"pre-treatment periods, found {n_pre}. " + f"Transformed outcome set to NaN.", + UserWarning, + stacklevel=2, + ) + if return_diagnostics: + per_unit[uid] = { + "alpha": float("nan"), + "beta": float("nan"), + "seasonal_effects": {}, + "pre_n_periods": n_pre, + "valid": False, + } + continue + + # Check seasonal parameters + q_pre = q_u[pre_u] + t_pre = t_u[pre_u] + observed_seasons = sorted(np.unique(q_pre)) + # Parameters: intercept + slope + (n_seasons - 1) dummies + n_params = 1 + len(observed_seasons) + + y_pre = y_u[pre_u] + + # Center time for numerical stability + t_mean = t_pre.mean() + t_pre_centered = t_pre - t_mean + + # Insufficient pre-observations for the seasonal model: fail + # closed like demeanq (warn + NaN the unit). The pre-fix code + # silently fit intercept+trend only (per-unit detrend) while + # the fit still reported rolling='detrendq' - with quarterly + # data and <= 5 pre-periods EVERY unit fell back, making the + # whole fit numerically identical to detrend with no trace + # (campaign finding). + if n_pre < n_params: + warnings.warn( + f"Unit {uid}: detrendq requires at least as many " + f"pre-treatment observations as seasonal parameters " + f"({n_params}), found {n_pre}. Transformed outcome set " + f"to NaN.", + UserWarning, + stacklevel=2, + ) + if return_diagnostics: + per_unit[uid] = { + "alpha": float("nan"), + "beta": float("nan"), + "seasonal_effects": {}, + "pre_n_periods": n_pre, + "valid": False, + } + continue + + # Season coverage: quarters unobserved in the pre-period have + # no estimated effect; predicting them at the reference-season + # level is silent extrapolation (campaign finding). + unobserved_seasons = sorted(set(q_u.tolist()) - set(observed_seasons)) + if unobserved_seasons: + warnings.warn( + f"Unit {uid}: detrendq cannot predict quarter(s) " + f"{unobserved_seasons} that never appear in the unit's " + f"pre-treatment periods. Transformed outcome set to NaN.", + UserWarning, + stacklevel=2, + ) + if return_diagnostics: + per_unit[uid] = { + "alpha": float("nan"), + "beta": float("nan"), + "seasonal_effects": {}, + "pre_n_periods": n_pre, + "valid": False, + } + continue + + use_seasonal = True + # Build design matrix: [1, t_centered, Q2, Q3, Q4] + X_pre_parts = [ + np.ones(n_pre, dtype=np.float64), + t_pre_centered, + ] + for s in observed_seasons[1:]: + X_pre_parts.append((q_pre == s).astype(np.float64)) + X_pre = np.column_stack(X_pre_parts) + + # Solve via scipy.linalg.lstsq + result = scipy_linalg.lstsq(X_pre, y_pre, cond=None) + coefs = result[0] + + if not np.all(np.isfinite(coefs)): + warnings.warn( + f"Unit {uid}: detrendq produced non-finite " + f"coefficients. Transformed outcome set to NaN.", + UserWarning, + stacklevel=2, + ) + if return_diagnostics: + per_unit[uid] = { + "alpha": float("nan"), + "beta": float("nan"), + "seasonal_effects": {}, + "pre_n_periods": n_pre, + "valid": False, + } + continue + + # Predict on ALL periods for this unit + t_all_centered = t_u - t_mean + n_all = len(t_u) + + X_all_parts = [ + np.ones(n_all, dtype=np.float64), + t_all_centered, + ] + if use_seasonal: + for s in observed_seasons[1:]: + X_all_parts.append((q_u == s).astype(np.float64)) + X_all = np.column_stack(X_all_parts) + y_hat = X_all @ coefs + + # Residuals + ydot_out[idx_u] = y_u - y_hat + + # Collect diagnostics for this unit + if return_diagnostics: + if use_seasonal: + seasonal_effects = { + int(s): float(coefs[idx + 2]) for idx, s in enumerate(observed_seasons[1:]) + } + else: + seasonal_effects = {} + per_unit[uid] = { + "alpha": float(coefs[0]), + "beta": float(coefs[1]), + "seasonal_effects": seasonal_effects, + "pre_n_periods": n_pre, + "valid": True, + } + + df["_ydot"] = ydot_out + + if return_diagnostics: + valid_units = [uid for uid, info in per_unit.items() if info["valid"]] + n_valid = len(valid_units) + n_total = len(per_unit) + diagnostics: Dict[str, Any] = { + "method": "detrendq", + "description": "Remove unit-specific trend + seasonal effects (α̂_i + β̂_i*t + Σγ̂_q*Q_q)", + "per_unit": per_unit, + "summary": { + "n_units_total": n_total, + "n_units_valid": n_valid, + "n_units_dropped": n_total - n_valid, + }, + } + return df, diagnostics + + return df + + def _dispatch_estimator( + self, + y: np.ndarray, + treatment: np.ndarray, + controls_matrix: Optional[np.ndarray], + cluster_ids: Optional[np.ndarray], + n_obs: int, + ) -> Tuple[ + float, + float, + Optional[np.ndarray], + Optional[np.ndarray], + int, + Optional[np.ndarray], + ]: + """Dispatch estimation to the appropriate method based on self.estimation_method. + + This is the central routing function that maps the user's estimation-method + choice to the corresponding implementation. After unit-specific rolling transformation + converts the panel into a cross-sectional dataset, this method applies the + chosen treatment-effect estimator to obtain the ATT. + + Corresponds to Step 2 of the Lee & Wooldridge (2025, 2026) procedure: + after computing Ẏ_{ir} (transformed outcome), apply reg/ipw/dr/psm + to the cross-section {(Ẏ_{ir}, D_i, X_i)}. + + Parameters + ---------- + y : np.ndarray of shape (n,) + Transformed outcome variable (\u1e8e_{ir} in paper notation). + This is the post-transformation average residual for each unit. + treatment : np.ndarray of shape (n,) + Binary treatment indicator (D_i). 1 = treated, 0 = control. + controls_matrix : np.ndarray of shape (n, K) or None + Covariate matrix (X_i). None if no controls specified. + Used for regression adjustment, propensity score, and matching. + cluster_ids : np.ndarray of shape (n,) or None + Cluster identifiers for cluster-robust variance estimation. + None unless the cluster= constructor parameter is set. + n_obs : int + Number of cross-sectional observations (units). + + Returns + ------- + tuple of (att, se, coefs, vcov, n_params, influence) + att : float + Estimated average treatment effect on the treated (\u03c4\u0302 in paper). + se : float + Standard error of the ATT estimate. + coefs : np.ndarray or None + Full coefficient vector from the regression (RA/IPW paths). + None for PSM. + vcov : np.ndarray or None + Variance-covariance matrix of coefficients. + None for PSM. + n_params : int + Number of parameters in the fitted design, used for the + residual degrees of freedom: df = N - n_params. For the reg + path this is design-coherent (LW 2026 Section 2): N - 2 + without controls, N - K - 2 for the plain design and + N - 2K - 2 when the treatment-covariate interaction is + active. + influence : np.ndarray or None + Observation-aligned influence contributions. Their unit-level + or cluster-level norm reproduces ``se`` and they can therefore + be combined across staggered cohort-time cells without assuming + independence. Matching returns ``None``. + + Raises + ------ + ValueError + If self.estimation_method is not in {'reg', 'ipw', 'dr', 'psm'}. + (Should not occur if __init__ validation passed.) + + Notes + ----- + Routing logic: + - 'reg' → _estimate_reg(): OLS of Ẏ on [1, D, X, D*(X-X̄₁)] + per Equation 3.3 in Lee & Wooldridge (2025) + - 'ipw' → _estimate_ipw(): Inverse probability weighting via + logit propensity score, Hajek-style normalization + - 'dr' → _estimate_dr(): Doubly-robust augmented IPW + combining outcome model and propensity weighting + - 'psm' → _estimate_psm(): Nearest-neighbor propensity score + matching (1:n with optional caliper) + + When controls_matrix is None, IPW/DR/PSM fall back to regression + adjustment (simple difference in means) with a warning. + + The variance family (self.vcov_type) determines which variance + estimator is used: + - 'classical': homoskedastic OLS variance + - 'hc1': HC1 (White) heteroskedasticity-robust + - 'hc2': HC2 (leverage-adjusted) + - 'hc3': HC3 (jackknife-style leverage adjustment) + Cluster-robust (Liang-Zeger CR1) inference activates via the + cluster= constructor parameter. + + References + ---------- + Lee, S. & Wooldridge, J. M. (2025). "A Simple Transformation Approach + to Difference-in-Differences Estimation for Panel Data." + Procedure 3.1, Equation 3.3. + Lee, S. J. & Wooldridge, J. M. (2026). "Simple Approaches to + Inference with Difference-in-Differences Estimators with + Small Cross-Sectional Sample Sizes." Procedure 2.1. + """ + if self.estimation_method == "reg": + return self._estimate_reg(y, treatment, controls_matrix, cluster_ids, n_obs) + elif self.estimation_method == "ipw": + return self._estimate_ipw(y, treatment, controls_matrix, cluster_ids, n_obs) + elif self.estimation_method == "psm": + return self._estimate_psm(y, treatment, controls_matrix, cluster_ids, n_obs) + else: # dr + return self._estimate_dr(y, treatment, controls_matrix, cluster_ids, n_obs) + + @staticmethod + def _finalize_influence( + influence: np.ndarray, + se: float, + ) -> Optional[np.ndarray]: + """Drop influence contributions that cannot support joint inference.""" + if not np.isfinite(se) or se <= 0 or not np.all(np.isfinite(influence)): + return None + return influence + + def _ols_treatment_influence( + self, + X: np.ndarray, + xtx_inv: np.ndarray, + residuals: np.ndarray, + n_obs: int, + n_params: int, + cluster_ids: Optional[np.ndarray], + coef_index: int = 1, + ) -> np.ndarray: + r"""Influence contributions for the OLS treatment coefficient. + + ``coef_index`` is the treatment coefficient's position in the + (possibly rank-reduced) design actually passed in - the caller + remaps it when solve_ols dropped columns (fix-wave WS8: computing + the bread/leverage on the full-width design while only the df + moved left the influence function describing a min-norm fit). + + The asymptotically linear representation of :math:`\hat\tau` is + :math:`\psi_i = e_2' (X'X)^{-1} x_i \varepsilon_i`. Each variance + estimator is a reweighting of those contributions, so applying the + estimator's own weights here makes the sum of squared contributions + (summed within clusters when clustering) reproduce the reported + standard error exactly, while preserving the per-unit structure that + cross-cell covariance needs. Every branch, including classical, + keeps the residual-based direction: replacing residuals by their + homoskedastic magnitude (``sigma * basis``) would fabricate + covariance between staggered cells that merely share control units. + """ + basis = X @ xtx_inv[:, coef_index] + dof = max(n_obs - n_params, 1) + psi = basis * residuals + + if cluster_ids is not None: + n_clusters = len(np.unique(cluster_ids)) + if n_clusters > 1: + cr1 = (n_clusters / (n_clusters - 1)) * ((n_obs - 1) / dof) + return psi * float(np.sqrt(cr1)) + # Degenerate single-cluster fallback that solve_ols resolves + # to hc1. + return psi * float(np.sqrt(n_obs / dof)) + + if self.vcov_type == "classical": + # Homoskedastic magnitude: sum_i sigma^2 (x_i' a)^2 = sigma^2 + # (X'X)^{-1}_22, the textbook OLS variance. The contributions are + # the residual-based psi rescaled to that magnitude, so a single + # cell reproduces the classical SE exactly while cross-cell + # products retain the unit-level residual dependence. + sigma = float(np.sqrt(float(residuals @ residuals) / dof)) + target = sigma * float(np.sqrt(float(basis @ basis))) + norm = float(np.sqrt(float(psi @ psi))) + if norm > 0.0 and np.isfinite(norm): + return psi * (target / norm) + return psi + + if self.vcov_type in ("hc2", "hc3"): + raw_leverage = np.sum((X @ xtx_inv) * X, axis=1) + if self.vcov_type in ("hc2", "hc3") and np.any(raw_leverage >= 1.0 - 1e-8): + # Match the shared linalg fail-closed contract (round-10 + # review: clipping fabricated a finite HC3 influence + # vector for a design whose HC3 vcov is NaN, so aggregate + # inference disagreed with the cell's own). + return np.full_like(psi, np.nan) + leverage = np.clip(raw_leverage, 0.0, 1.0 - 1e-10) + if self.vcov_type == "hc2": + return psi / np.sqrt(1.0 - leverage) + return psi / (1.0 - leverage) + + # hc1 + return psi * float(np.sqrt(n_obs / dof)) + + def _moment_influence( + self, + psi_full: np.ndarray, + n_obs: int, + cluster_ids: Optional[np.ndarray], + ) -> np.ndarray: + """Rescale a semiparametric influence function to ATT scale. + + Mirrors the variance formulas used by the IPW/IPWRA paths, so the + sum of squared contributions reproduces their reported variance. + """ + if cluster_ids is not None: + n_clusters = len(np.unique(cluster_ids)) + if n_clusters > 1: + return psi_full * float(np.sqrt(n_clusters / (n_clusters - 1))) / n_obs + return (psi_full - float(np.mean(psi_full))) / float(np.sqrt(n_obs * (n_obs - 1))) + + def _estimate_reg( + self, + y: np.ndarray, + treatment: np.ndarray, + controls_matrix: Optional[np.ndarray], + cluster_ids: Optional[np.ndarray], + n_obs: int, + ) -> Tuple[ + float, + float, + Optional[np.ndarray], + Optional[np.ndarray], + int, + Optional[np.ndarray], + ]: + """Estimate ATT via regression adjustment (OLS). + + Fits y = alpha + tau*D + X*beta + D*(X - X_bar_1)*gamma + epsilon + and returns tau as the ATT estimate (LW2025 Equation 3.3). + + The interaction term D*(X - X_bar_1) allows covariate effects to + differ between treated and control groups. It is only included when + both N_treated > K+1 and N_control > K+1. + + Parameters + ---------- + y : ndarray of shape (n,) + Transformed outcome. + treatment : ndarray of shape (n,) + Binary treatment indicator. + controls_matrix : ndarray of shape (n, p) or None + Control variables. + cluster_ids : ndarray of shape (n,) or None + Cluster identifiers for cluster-robust SEs. + n_obs : int + Number of observations. + + Returns + ------- + att : float + Treatment effect coefficient. + se : float + Standard error of treatment coefficient. + coefs : ndarray + Full coefficient vector. + vcov : ndarray or None + Variance-covariance matrix. + n_params : int + Number of parameters in the regression. + """ + # Build design matrix: [intercept, treatment, controls, interaction] + parts = [np.ones((n_obs, 1)), treatment.reshape(-1, 1)] + if controls_matrix is not None: + parts.append(controls_matrix) + # Add D*(X - X_bar_1) interaction term when sample sizes permit + # (LW2025 Eq 3.3: requires N_0 > K+1 and N_1 > K+1). K is the + # IDENTIFIED control dimension (round-11 review: the nominal + # column count let a perfectly collinear control flip the gate + # off and silently change the ATT while adding no information). + K = int( + _detect_rank_deficiency(np.column_stack([np.ones(n_obs), controls_matrix]))[0] - 1 + ) + treated_mask = treatment == 1 + n_treated = int(treated_mask.sum()) + n_control = n_obs - n_treated + if n_treated > K + 1 and n_control > K + 1: + X_bar_1 = controls_matrix[treated_mask].mean(axis=0) + interaction = treatment.reshape(-1, 1) * (controls_matrix - X_bar_1) + parts.append(interaction) + X = np.hstack(parts) + # EFFECTIVE design rank on the column-equilibrated matrix + # (round-13 review: the nominal width rejected designs whose + # redundant columns the rank-aware solver drops, e.g. N=4 with + # [1, D, x, 2x] has rank 3 and one residual df; equilibration + # keeps the rank decision scale-invariant). + # Round-17 review: matrix_rank's looser default tolerance + # disagreed with solve_ols's scale-invariant pivoted-QR 1e-7 + # convention on near-collinear columns, so the gate and the fit + # could select different designs - use the SHARED detector. + rank_eff = int(_detect_rank_deficiency(X)[0]) + if X.shape[0] < 3 or X.shape[0] - rank_eff <= 0: + # Registry small-sample guards (N >= 3; positive residual df, + # i.e. N > K + 2 with controls / N > 2K + 2 interacted): the + # shared classical vcov divides by n - k, so an exactly- + # saturated design reached ZeroDivisionError (review finding). + raise ValueError( + f"Invalid exact-inference design: {X.shape[0]} " + f"observation(s) with {rank_eff} identified parameter(s). " + f"LWDiD requires at least 3 cross-sectional units and a " + f"positive residual df (N > K + 2 with controls)." + ) + + # Determine vcov_type for solve_ols (hc3 routes through the shared + # linalg backend; clustered fits resolve to CR1 via cluster_ids) + vcov_type = self._resolve_vcov_type() + + # Call solve_ols + coefs, residuals, vcov = solve_ols( + X, + y, + cluster_ids=cluster_ids, + return_vcov=True, + vcov_type=vcov_type, + ) + + # ATT = coefficient on treatment (index 1) + att = float(coefs[1]) + # SE from vcov diagonal + if vcov is not None and np.isfinite(vcov[1, 1]): + se = float(np.sqrt(max(vcov[1, 1], 0.0))) + else: + se = np.nan + + # Return the fitted design's EFFECTIVE parameter count so callers + # compute a design-coherent residual df: N - 2 without controls, + # N - K - 2 for the plain design (1, D, X), N - 2K - 2 when the + # interaction D*(X - X_bar_1) is active (LW 2026 Section 2) - and, + # under rank deficiency, the KEPT-column count (fix-wave WS8: the + # nominal count understated the df and the full-width pinv bread + # broke the IF == solve_ols SE identity). + nan_mask = np.isnan(coefs) + n_params_effective = int(np.sum(~nan_mask)) + if nan_mask.any(): + if nan_mask[1]: + # The treatment column itself was pivoted out: the ATT is + # unidentified (solve_ols already emitted the rank warning). + return np.nan, np.nan, coefs, vcov, n_params_effective, None + kept = np.flatnonzero(~nan_mask) + X_used = X[:, kept] + coef_index = int(np.flatnonzero(kept == 1)[0]) + else: + X_used = X + coef_index = 1 + # Scale-equilibrated bread (round-13 review: the raw-Gram pinv + # silently dropped low-scale directions at large covariate units - + # cell coefficients/SEs from solve_ols were invariant while the + # reconstructed influence, and therefore every aggregate SE and + # multiplier-bootstrap input, was not). With column scales D, + # (X'X)^{-1} = D^{-1} (Xs'Xs)^{-1} D^{-1} for Xs = X D^{-1}. + used_scales = np.linalg.norm(X_used, axis=0) + used_scales[used_scales == 0] = 1.0 + X_scaled = X_used / used_scales + xtx_inv = np.linalg.pinv(X_scaled.T @ X_scaled) / np.outer(used_scales, used_scales) + if self.vcov_type == "hc2" and cluster_ids is None: + # Round-21 review: the shared hc2 kernel keeps its RELEASED + # 1 - h floor (tracked separately), but the NEW LWDiD surface + # must not report a fabricated finite variance for a + # perfectly-leveraged design - fail closed HERE, mirroring + # hc3 (point retained, inference NaN). + leverage_used = np.sum((X_used @ xtx_inv) * X_used, axis=1) + if np.any(leverage_used >= 1.0 - 1e-8): + n_lev1 = int(np.sum(leverage_used >= 1.0 - 1e-8)) + warnings.warn( + f"HC2 variance is undefined for this design: {n_lev1} " + f"observation(s) have hat-matrix leverage ~1 (e.g. a " + f"single treated unit). Returning NaN inference (point " + f"retained); use vcov_type='classical' exact inference " + f"or add treated units.", + UserWarning, + stacklevel=2, + ) + se = np.nan + influence = self._finalize_influence( + self._ols_treatment_influence( + X_used, + xtx_inv, + residuals, + n_obs, + n_params_effective, + cluster_ids, + coef_index=coef_index, + ), + se, + ) + return att, se, coefs, vcov, n_params_effective, influence + + def _estimate_ipw( + self, + y: np.ndarray, + treatment: np.ndarray, + controls_matrix: Optional[np.ndarray], + cluster_ids: Optional[np.ndarray], + n_obs: int, + ) -> Tuple[ + float, + float, + Optional[np.ndarray], + Optional[np.ndarray], + int, + Optional[np.ndarray], + ]: + """Estimate ATT via inverse probability weighting. + + Uses propensity scores to reweight control observations. + + Parameters + ---------- + y : ndarray of shape (n,) + Transformed outcome. + treatment : ndarray of shape (n,) + Binary treatment indicator. + controls_matrix : ndarray of shape (n, p) or None + Covariates for propensity score model. + cluster_ids : ndarray of shape (n,) or None + Cluster identifiers. + n_obs : int + Number of observations. + + Returns + ------- + att : float + IPW-estimated ATT. + se : float + Standard error. + coefs : ndarray or None + Not returned for IPW (None). + vcov : ndarray or None + Not returned for IPW (None). + n_params : int + Number of parameters in the underlying regression. + """ + if controls_matrix is None or controls_matrix.shape[1] == 0: + # Without covariates, IPW reduces to simple difference + # in means (propensity score is constant) + warnings.warn( + "IPW without control variables reduces to a simple " + "difference in means. Consider using estimation_method='reg'.", + UserWarning, + stacklevel=2, + ) + return self._estimate_reg( + y, treatment, None, cluster_ids, n_obs + ) # returns 5-tuple including n_params + + # Step 1: Estimate propensity score via logit + # solve_logit adds intercept automatically + coefs_logit, probs = solve_logit(controls_matrix, treatment) + + # Rank/convergence handling (round-11 review): the shared solver + # marks DROPPED collinear columns with NaN coefficients while the + # fitted probabilities remain valid - that reduced-rank fit stays + # an IPW fit (the pre-fix code silently substituted regression + # adjustment under ipw provenance). Only a genuinely failed solve + # (non-finite probabilities) falls back. + kept_ps = np.isfinite(coefs_logit) + if not kept_ps.all(): + if np.all(np.isfinite(probs)): + warnings.warn( + f"Propensity model is rank-deficient: " + f"{int((~kept_ps).sum())} collinear column(s) dropped; " + f"continuing IPW with the reduced-rank propensity fit.", + UserWarning, + stacklevel=2, + ) + else: + warnings.warn( + "Logistic regression did not converge (non-finite " + "probabilities). Falling back to 'reg' estimation. " + "Consider standardizing controls.", + UserWarning, + stacklevel=2, + ) + return self._estimate_reg(y, treatment, controls_matrix, cluster_ids, n_obs) + + # Convergence check: complete/quasi-complete separation + if np.any(probs < 1e-8) or np.any(probs > 1 - 1e-8): + warnings.warn( + "Possible complete separation detected in propensity score model. " + "Some predicted probabilities are near 0 or 1. " + "Results may be unreliable.", + UserWarning, + stacklevel=2, + ) + + # Step 2: Trim propensity scores to [pscore_trim, 1 - pscore_trim] + trim_lo, trim_hi = self.pscore_trim, 1.0 - self.pscore_trim + n_trimmed = int((probs < trim_lo).sum() + (probs > trim_hi).sum()) + if n_trimmed > 0: + warnings.warn( + f"LWDiD: {n_trimmed} observation(s) had propensity scores trimmed " + f"to [{self.pscore_trim:.3f}, {1-self.pscore_trim:.3f}].", + UserWarning, + stacklevel=2, + ) + probs_raw = probs + probs = np.clip(probs_raw, trim_lo, trim_hi) + # Observations at the clip boundary have ZERO weight-derivative in + # gamma (round-11 review: using clipped probabilities in the logit + # score/Hessian broke the estimating-equation linearization - the + # score at the MLE is ~0 in the RAW fitted probabilities only). + unclipped = (probs_raw > trim_lo) & (probs_raw < trim_hi) + + # Step 3: Compute IPW weights + # For treated: weight = 1 + # For control: weight = p(x) / (1 - p(x)) + # Normalized so control weights sum to n_treated + ipw_weights = np.where( + treatment == 1, + 1.0, + probs / (1.0 - probs), + ) + + # Normalize weights: treated get weight 1/n_treated, + # control weights normalized to sum to 1 + treat_mask = treatment == 1 + ctrl_mask = treatment == 0 + + w_ctrl_sum = ipw_weights[ctrl_mask].sum() + if w_ctrl_sum <= 0: + warnings.warn( + "IPW control weights sum to zero. Falling back to " "unweighted 'reg' estimation.", + UserWarning, + stacklevel=2, + ) + return self._estimate_reg(y, treatment, controls_matrix, cluster_ids, n_obs) + + # Hajek-style ATT estimator + att_treated = y[treat_mask].mean() + att_control = np.sum(ipw_weights[ctrl_mask] * y[ctrl_mask]) / w_ctrl_sum + att = float(att_treated - att_control) + + # Step 4: Compute SE via semiparametric influence function + # (Lunceford & Davidian 2004 AIPW form - the documented, + # adjudicated alternative to the papers'/Stata package's stacked + # E.3/E.4 form; see the REGISTRY IPWRA-variance Note). + # The full IF consists of the Hajek main term plus a propensity score + # estimation uncertainty correction. + n_treated_f = float(treat_mask.sum()) + p_bar = n_treated_f / n_obs # P(D=1) estimate + + # --- Hajek influence function (main term) --- + w_ctrl = ipw_weights[ctrl_mask] # p/(1-p) for controls + + psi_ht = np.zeros(n_obs) + psi_ht[treat_mask] = (y[treat_mask] - att_treated) / p_bar + psi_ht[ctrl_mask] = -w_ctrl * (y[ctrl_mask] - att_control) / p_bar + + # --- Propensity score estimation uncertainty correction --- + # Design matrix with intercept (solve_logit adds intercept internally, + # so we reconstruct it here for the IF computation), restricted to + # the KEPT (identified) propensity columns under rank deficiency. + X_ps = np.column_stack([np.ones(n_obs), controls_matrix])[:, kept_ps] + + # Logit score: S_i = (D_i - p_i) * X_i, at the RAW fitted + # probabilities (the score of the actual MLE; clipped probabilities + # are a weighting choice, not the estimating equation). + S_gamma = (treatment - probs_raw)[:, np.newaxis] * X_ps + + # Logit Hessian: H = -(1/n) * X' diag(p*(1-p)) X (raw fit) + W_ps = probs_raw * (1 - probs_raw) + H_gamma = -(X_ps.T * W_ps) @ X_ps / n_obs + try: + H_gamma_inv = np.linalg.inv(H_gamma) + except np.linalg.LinAlgError: + H_gamma_inv = np.linalg.pinv(H_gamma) + + # Sensitivity: dATT/dgamma + # dw/dgamma_i = w_i * X_i (logit chain rule) for UNCLIPPED + # observations; a clipped weight is locally constant in gamma. + # dATT/dgamma = -(1/w_sum) * sum_ctrl(w_i * X_i * (Y_i - mu_0)) + # The (Y_i - mu_0) centering comes from the quotient rule for the + # Hajek estimator (d/dgamma of Sigma(wY)/Sigma(w)) and ensures + # translation invariance of the resulting SE. + dw_dgamma_ctrl = (w_ctrl * unclipped[ctrl_mask])[:, np.newaxis] * X_ps[ctrl_mask] + Y_ctrl_centered = y[ctrl_mask] - att_control + dATT_dgamma = -(dw_dgamma_ctrl * Y_ctrl_centered[:, np.newaxis]).sum(axis=0) / ( + n_obs * p_bar + ) + + # PS adjustment: psi_adj_i = (S_i @ H^{-1}) @ dATT_dgamma + ps_adjustment = (S_gamma @ H_gamma_inv.T) @ dATT_dgamma + + # Full IF = main term - PS correction + psi_full = psi_ht - ps_adjustment + + # --- Variance estimation --- + if cluster_ids is not None: + cluster_df = pd.DataFrame({"psi": psi_full, "cluster": cluster_ids}) + cluster_sums = cluster_df.groupby("cluster")["psi"].sum().values + n_clusters = len(cluster_sums) + if n_clusters <= 1: + warnings.warn( + "Only 1 cluster found; falling back to non-clustered " + "variance for IPW influence function.", + UserWarning, + stacklevel=2, + ) + var_att = float(np.var(psi_full, ddof=1) / n_obs) + else: + var_att = float( + (n_clusters / (n_clusters - 1)) * np.sum(cluster_sums**2) / n_obs**2 + ) + else: + var_att = float(np.var(psi_full, ddof=1) / n_obs) + + se = float(np.sqrt(max(var_att, 0.0))) + + # n_params: IDENTIFIED propensity-model rank (round-12 review: + # the nominal count let a redundant control shrink residual df). + n_params = int(kept_ps.sum()) + influence = self._finalize_influence( + self._moment_influence(psi_full, n_obs, cluster_ids), se + ) + return att, se, None, None, n_params, influence + + def _estimate_psm( + self, + y: np.ndarray, + treatment: np.ndarray, + controls_matrix: Optional[np.ndarray], + cluster_ids: Optional[np.ndarray], + n_obs: int, + ) -> Tuple[ + float, + float, + Optional[np.ndarray], + Optional[np.ndarray], + int, + Optional[np.ndarray], + ]: + """Estimate ATT via propensity score matching. + + For each treated unit, find the nearest control unit(s) by + propensity score (1:n_neighbors nearest-neighbor matching, 1:1 by + default, with replacement by default), then compute ATT as the + average difference between treated and matched controls. + + Parameters + ---------- + y : ndarray of shape (n,) + Transformed outcome. + treatment : ndarray of shape (n,) + Binary treatment indicator. + controls_matrix : ndarray of shape (n, p) or None + Covariates for propensity score model. + cluster_ids : ndarray of shape (n,) or None + Cluster identifiers. + n_obs : int + Number of observations. + + Returns + ------- + att : float + PSM-estimated ATT. + se : float + Always NaN (fail-closed): no valid matching variance is + implemented (the naive matched-pairs formula ignored control + reuse and first-stage uncertainty; an Abadie-Imbens variance + is tracked in DEFERRED.md). + coefs : ndarray or None + Not returned for PSM (None). + vcov : ndarray or None + Not returned for PSM (None). + n_params : int + Effective number of parameters. + """ + if controls_matrix is None or controls_matrix.shape[1] == 0: + # Unreachable from fit() (config guard rejects covariate-less + # PSM); kept as defense in depth for direct callers. + raise ValueError( + "estimation_method='psm' requires covariates: without them " + "there is no propensity score to match on. Use " + "estimation_method='reg', or supply covariates." + ) + + treat_mask = treatment == 1 + ctrl_mask = treatment == 0 + n_treated = int(treat_mask.sum()) + n_control = int(ctrl_mask.sum()) + + if n_treated == 0 or n_control == 0: + warnings.warn( + "PSM estimation failed: no treated or no control units available. " + "Returning NaN results.", + UserWarning, + stacklevel=2, + ) + return np.nan, np.nan, None, None, 2, None + + # Step 1: Estimate propensity score via logit + coefs_logit, probs = solve_logit(controls_matrix, treatment) + + # Rank/convergence handling (round-19 review; mirrors ipw/dr): + # NaN coefficients with FINITE probabilities are a reduced-rank + # propensity fit - matching needs only the probabilities, so PSM + # continues (the pre-fix path substituted a regression-adjustment + # point under psm provenance). Only genuinely failed solves + # (non-finite probabilities) fall back, fail-closed. + kept_ps_match = np.isfinite(coefs_logit) + if not kept_ps_match.all(): + if np.all(np.isfinite(probs)): + warnings.warn( + f"Propensity model is rank-deficient: " + f"{int((~kept_ps_match).sum())} collinear column(s) " + f"dropped; continuing PSM with the reduced-rank " + f"propensity fit.", + UserWarning, + stacklevel=2, + ) + else: + # Review round 3: the pre-fix fallback returned the + # regression point WITH its finite OLS inference while the + # results metadata still said 'psm'. Point retained, + # inference NaN. + warnings.warn( + "Logistic regression did not converge (non-finite " + "probabilities); the point estimate falls back to " + "regression adjustment, and inference is NaN under the " + "PSM fail-closed contract. Consider standardizing " + "controls or using estimation_method='reg'.", + UserWarning, + stacklevel=2, + ) + att_fb, _, _, _, n_params_fb, _ = self._estimate_reg( + y, treatment, controls_matrix, cluster_ids, n_obs + ) + return att_fb, np.nan, None, None, n_params_fb, None + + # Convergence check: complete/quasi-complete separation + if np.any(probs < 1e-8) or np.any(probs > 1 - 1e-8): + warnings.warn( + "Possible complete separation detected in propensity score model. " + "Some predicted probabilities are near 0 or 1. " + "Results may be unreliable.", + UserWarning, + stacklevel=2, + ) + + # Step 2: Trim propensity scores to [pscore_trim, 1 - pscore_trim] + trim_lo, trim_hi = self.pscore_trim, 1.0 - self.pscore_trim + n_trimmed = int((probs < trim_lo).sum() + (probs > trim_hi).sum()) + if n_trimmed > 0: + warnings.warn( + f"LWDiD: {n_trimmed} observation(s) had propensity scores trimmed " + f"to [{self.pscore_trim:.3f}, {1-self.pscore_trim:.3f}].", + UserWarning, + stacklevel=2, + ) + probs = np.clip(probs, trim_lo, trim_hi) + + # Step 3: Nearest-neighbor matching (with replacement) + p_treated = probs[treat_mask] + p_control = probs[ctrl_mask] + y_treated = y[treat_mask] + y_control = y[ctrl_mask] + + # For each treated unit, find n_neighbors nearest controls + matched_y_control = np.empty(n_treated) + available_mask = np.ones(n_control, dtype=bool) + n_partial_matches = 0 + + for i in range(n_treated): + valid_control_idx = np.where(available_mask)[0] + if len(valid_control_idx) == 0: + matched_y_control[i] = np.nan + continue + + distances = np.abs(p_treated[i] - p_control[valid_control_idx]) + + if self.caliper is not None: + within_caliper = distances <= self.caliper + if not within_caliper.any(): + matched_y_control[i] = np.nan + continue + distances = np.where(within_caliper, distances, np.inf) + + nearest_local = np.argsort(distances)[: self.n_neighbors] + # Caliper contract: only within-caliper controls may be + # averaged. argsort places np.inf (out-of-caliper) LAST but + # still returns it, so a partial shortfall (>=1 but + # < n_neighbors controls inside the caliper) used to average + # arbitrarily distant controls into the counterfactual + # (campaign finding: deterministic ATT of -49 vs the correct + # caliper-respecting 1.0 on the repro fixture). + nearest_local = nearest_local[np.isfinite(distances[nearest_local])] + if len(nearest_local) < self.n_neighbors: + n_partial_matches += 1 + nearest_global = valid_control_idx[nearest_local] + matched_y_control[i] = y_control[nearest_global].mean() + + if not self.with_replacement: + available_mask[nearest_global] = False + + # Step 4: Compute ATT = mean(Y_treated - Y_matched_control) + # Exclude NaN matches (from caliper) + valid_matches = np.isfinite(matched_y_control) + n_unmatched = int(np.isnan(matched_y_control).sum()) + if n_unmatched > 0: + warnings.warn( + f"LWDiD PSM: {n_unmatched} treated unit(s) could not be matched " + f"within caliper={self.caliper}. ATT computed from {n_treated - n_unmatched} matches.", + UserWarning, + stacklevel=2, + ) + if n_partial_matches > 0: + warnings.warn( + f"LWDiD PSM: {n_partial_matches} treated unit(s) had fewer " + f"than n_neighbors={self.n_neighbors} control(s) within " + f"caliper={self.caliper}; their matches average the " + f"within-caliper control(s) only.", + UserWarning, + stacklevel=2, + ) + if not valid_matches.any(): + warnings.warn( + "PSM estimation failed: no valid matches found (all exceeded caliper). " + "Returning NaN results.", + UserWarning, + stacklevel=2, + ) + return np.nan, np.nan, None, None, 2, None + diffs = y_treated[valid_matches] - matched_y_control[valid_matches] + att = float(np.mean(diffs)) + + # Step 5: Inference fails closed (review finding). The former + # sqrt(var(diffs)/n) treated matched differences as INDEPENDENT - + # with replacement matching a control can appear in many treated + # counterfactuals, so their common uncertainty cancels out of that + # formula - and it omits the propensity/matching first-stage + # uncertainty entirely. A valid matching variance (Abadie-Imbens) + # is tracked in DEFERRED.md; until it lands the point is retained + # and the inference tuple is NaN (same convention as the staggered + # 'unavailable_matching' basis). + warnings.warn( + "LWDiD PSM: no valid matching variance estimator is implemented " + "(the naive var(diffs)/n formula ignores matched-control reuse " + "and first-stage matching uncertainty). The ATT point estimate " + "is reported with NaN inference; use estimation_method='dr' for " + "a doubly robust alternative with valid inference.", + UserWarning, + stacklevel=2, + ) + se = np.nan + + # Effective n_params: intercept + controls (for propensity model) + n_params = 1 + controls_matrix.shape[1] + return att, se, None, None, n_params, None + + def _estimate_dr( + self, + y: np.ndarray, + treatment: np.ndarray, + controls_matrix: Optional[np.ndarray], + cluster_ids: Optional[np.ndarray], + n_obs: int, + ) -> Tuple[ + float, + float, + Optional[np.ndarray], + Optional[np.ndarray], + int, + Optional[np.ndarray], + ]: + """Estimate ATT via augmented IPW (doubly robust). + + Combines regression adjustment with inverse probability weighting + for double robustness. + + Parameters + ---------- + y : ndarray of shape (n,) + Transformed outcome. + treatment : ndarray of shape (n,) + Binary treatment indicator. + controls_matrix : ndarray of shape (n, p) or None + Covariates. + cluster_ids : ndarray of shape (n,) or None + Cluster identifiers. + n_obs : int + Number of observations. + + Returns + ------- + att : float + Doubly-robust ATT estimate. + se : float + Standard error. + coefs : ndarray or None + Not returned for DR (None). + vcov : ndarray or None + Not returned for DR (None). + n_params : int + Effective number of parameters for df computation. + + Notes + ----- + **Variance form (paper mapping).** The reported SE and the influence + function consumed by the multiplier bootstrap use the AIPW efficient + influence function (Lunceford & Davidian 2004) — NOT the stacked + M-estimator form of Lee & Wooldridge (2026) Appendix E.3 that the + authors' Stata package implements. This is a documented, + independently anchored alternative, adjudicated in PR #588's final + round: the AIPW EIF is the standard doubly-robust influence function + in the causal-inference literature and is anchored by the RA config's + bootstrap-SE parity gate against the Stata golden plus the suite's + analytical/bootstrap cross-path pins. Measured on the Walmart + application (2026-08-16): DR point estimates agree with the authors' + package to ~1e-3, while DR (IPWRA) multiplier-bootstrap SEs diverge + systematically by ~15%; the RA config's SEs agree within Monte-Carlo + bounds. See the LWDiD IPWRA-variance note in + ``docs/methodology/REGISTRY.md``; implementing the E.3 stacked form + remains an available follow-up if package-form SE parity is preferred. + """ + if controls_matrix is None or controls_matrix.shape[1] == 0: + # Without covariates, DR reduces to regression adjustment. + # Say so, matching the routing docstring and the ipw branch + # (no-silent-failures contract). + warnings.warn( + "DR (doubly robust) without control variables reduces to " + "regression adjustment. Consider using " + "estimation_method='reg'.", + UserWarning, + stacklevel=2, + ) + return self._estimate_reg( + y, treatment, None, cluster_ids, n_obs + ) # returns 5-tuple including n_params + + treat_mask = treatment == 1 + ctrl_mask = treatment == 0 + n_treated = int(treat_mask.sum()) + n_control = int(ctrl_mask.sum()) + + # Step 1: Get propensity scores + coefs_logit, probs = solve_logit(controls_matrix, treatment) + + # Rank/convergence handling (round-11 review; mirrors _estimate_ipw): + # NaN coefficients with finite probabilities = a reduced-rank + # propensity fit that remains a DR fit; only non-finite + # probabilities fall back to regression adjustment. + kept_ps = np.isfinite(coefs_logit) + if not kept_ps.all(): + if np.all(np.isfinite(probs)): + warnings.warn( + f"Propensity model is rank-deficient: " + f"{int((~kept_ps).sum())} collinear column(s) dropped; " + f"continuing DR with the reduced-rank propensity fit.", + UserWarning, + stacklevel=2, + ) + else: + warnings.warn( + "Logistic regression did not converge (non-finite " + "probabilities). Falling back to 'reg' estimation. " + "Consider standardizing controls.", + UserWarning, + stacklevel=2, + ) + return self._estimate_reg(y, treatment, controls_matrix, cluster_ids, n_obs) + + # Convergence check: complete/quasi-complete separation + if np.any(probs < 1e-8) or np.any(probs > 1 - 1e-8): + warnings.warn( + "Possible complete separation detected in propensity score model. " + "Some predicted probabilities are near 0 or 1. " + "Results may be unreliable.", + UserWarning, + stacklevel=2, + ) + + trim_lo_dr, trim_hi_dr = self.pscore_trim, 1.0 - self.pscore_trim + n_trimmed_dr = int((probs < trim_lo_dr).sum() + (probs > trim_hi_dr).sum()) + if n_trimmed_dr > 0: + warnings.warn( + f"LWDiD: {n_trimmed_dr} observation(s) had propensity scores trimmed " + f"to [{self.pscore_trim:.3f}, {1-self.pscore_trim:.3f}].", + UserWarning, + stacklevel=2, + ) + probs_raw = probs + probs = np.clip(probs_raw, self.pscore_trim, 1.0 - self.pscore_trim) + # Zero weight-derivative at the clip boundary; score/Hessian use + # the RAW fitted probabilities (round-11 review; see _estimate_ipw). + unclipped = (probs_raw > trim_lo_dr) & (probs_raw < trim_hi_dr) + + # Step 2: Fit outcome model on control units only using WLS with IPW weights + # This matches the Stata/lwdid-py reference: outcome model is fitted on + # controls with weights w_i = p(X_i)/(1-p(X_i)) to target ATT. + X_ctrl = np.column_stack([np.ones(n_control), controls_matrix[ctrl_mask]]) + y_ctrl = y[ctrl_mask] + + # IPW weights for control units + ipw_ctrl = probs[ctrl_mask] / (1.0 - probs[ctrl_mask]) + ipw_ctrl_sum = ipw_ctrl.sum() + + if ipw_ctrl_sum <= 0: + # Fall back to RA if IPW weights degenerate + return self._estimate_reg( + y, treatment, controls_matrix, cluster_ids, n_obs + ) # returns 5-tuple including n_params + + # WLS via sqrt(w) transformation through the shared RANK-AWARE + # solver (round-12 review: the raw inv/pinv Gram on the nominal + # columns was not scale-invariant - an exactly redundant + # 1e12-rescaled duplicate changed the DR SE by ~2.5x). Dropped + # collinear columns get NaN coefficients; the identified mask is + # reused for prediction and every outcome-model IF term. + sqrt_w = np.sqrt(ipw_ctrl) + X_ctrl_w = X_ctrl * sqrt_w[:, np.newaxis] + y_ctrl_w = y_ctrl * sqrt_w + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # rank warning surfaced below + coefs_full, _, _ = solve_ols(X_ctrl_w, y_ctrl_w) + kept_om = np.isfinite(coefs_full) + if not kept_om.all(): + warnings.warn( + f"DR outcome model is rank-deficient: " + f"{int((~kept_om).sum())} collinear column(s) dropped; " + f"continuing with the identified outcome design.", + UserWarning, + stacklevel=2, + ) + coefs_outcome = coefs_full[kept_om] + + # Predict counterfactual for all units (identified columns only) + X_all = np.column_stack([np.ones(n_obs), controls_matrix])[:, kept_om] + mu_0 = X_all @ coefs_outcome + + # Step 3: Compute AIPW/IPWRA estimator (Hajek normalization) + # ATT = mean_{D=1}(Y - mu_0) - sum_{D=0}[w*(Y-mu_0)] / sum_{D=0}(w) + resid = y - mu_0 + resid_ctrl = resid[ctrl_mask] + + # Treated component + att_treated_part = resid[treat_mask].mean() + + # Control component (Hajek: divide by sum of weights) + weights_sum = ipw_ctrl_sum + att_ctrl_part = np.sum(ipw_ctrl * resid_ctrl) / weights_sum + + att = float(att_treated_part - att_ctrl_part) + + # Step 4: Compute SE via full semiparametric influence function + # The IPWRA IF consists of 3 components (Cattaneo 2010, Lunceford & Davidian 2004): + # 1. Hajek main term (plug-in IF) + # 2. Propensity score estimation uncertainty correction + # 3. Outcome model estimation uncertainty correction + n_treated_f = float(n_treated) + p_bar = n_treated_f / n_obs # P(D=1) estimate + + # Control term (Hajek weighted mean of control residuals) + control_term = att_ctrl_part # = sum(w*resid_C) / sum(w) + + # ================================================================ + # Component 1: Hajek influence function (main term) + # Hajek linearization for ATT = mean_T(resid) - sum_C(w*resid)/sum_C(w) + # ================================================================ + psi = np.zeros(n_obs) + psi[treat_mask] = (resid[treat_mask] - att) / p_bar + psi[ctrl_mask] = -ipw_ctrl * (resid_ctrl - control_term) / weights_sum * n_obs + + # ================================================================ + # Component 2: Propensity score estimation uncertainty correction + # S_gamma_i = (D_i - p_i) * X_i (logit score) + # H_gamma = -(1/n) * X' diag(p*(1-p)) X (logit Hessian) + # dATT/dgamma = -sum_C[dw/dgamma * (resid - B)] / sum_C(w) + # ================================================================ + X_ps = np.column_stack([np.ones(n_obs), controls_matrix])[:, kept_ps] + + # Logit score (RAW fitted probabilities - the actual MLE's + # estimating equation; clipping is a weighting choice) + S_gamma = (treatment - probs_raw)[:, np.newaxis] * X_ps + + # Logit Hessian (raw fit) + W_ps = probs_raw * (1 - probs_raw) + H_gamma = -(X_ps.T * W_ps) @ X_ps / n_obs + try: + H_gamma_inv = np.linalg.inv(H_gamma) + except np.linalg.LinAlgError: + H_gamma_inv = np.linalg.pinv(H_gamma) + + # Sensitivity of ATT to propensity score parameters + # dw/dgamma_i = w_i * X_i for UNCLIPPED observations (a clipped + # weight is locally constant in gamma); chain through the Hajek + # control term + r_minus_B = resid_ctrl - control_term + dw_dgamma_ctrl = (ipw_ctrl * unclipped[ctrl_mask])[:, np.newaxis] * X_ps[ctrl_mask] + dATT_dgamma = -(dw_dgamma_ctrl * r_minus_B[:, np.newaxis]).sum(axis=0) / weights_sum + + # PS adjustment + ps_adjustment = (S_gamma @ H_gamma_inv.T) @ dATT_dgamma + + # ================================================================ + # Component 3: Outcome model estimation uncertainty correction + # The outcome model is WLS fitted on controls with IPW weights: + # E[Y|X, D=0] fitted by WLS with w_i = p/(1-p). + # S_beta_i = w_i * resid_i * X_i * I(D_i=0) (WLS score) + # H_beta = -(1/n) * X_ctrl' diag(w) X_ctrl (WLS Hessian) + # dATT/dbeta = -mean_T(X_i) + sum_C(w_i*X_i) / sum_C(w) + # ================================================================ + X_om = np.column_stack([np.ones(n_obs), controls_matrix])[:, kept_om] + X_ctrl_om = X_om[ctrl_mask] + + # WLS score (nonzero only for control units) + S_beta = np.zeros((n_obs, X_om.shape[1])) + S_beta[ctrl_mask] = ipw_ctrl[:, np.newaxis] * resid_ctrl[:, np.newaxis] * X_ctrl_om + + # WLS Hessian: H_beta = -(1/n) * X_ctrl' diag(w) X_ctrl + H_beta = -(X_ctrl_om.T * ipw_ctrl) @ X_ctrl_om / n_obs + try: + H_beta_inv = np.linalg.inv(H_beta) + except np.linalg.LinAlgError: + H_beta_inv = np.linalg.pinv(H_beta) + + # Sensitivity of ATT to outcome model parameters + # dATT/dbeta = -mean_T(X_i) + weighted_mean_C(X_i) + X_bar_treated = X_om[treat_mask].mean(axis=0) + X_bar_ctrl_w = (ipw_ctrl[:, np.newaxis] * X_ctrl_om).sum(axis=0) / weights_sum + dATT_dbeta = -X_bar_treated + X_bar_ctrl_w + + # Outcome model adjustment + om_adjustment = (S_beta @ H_beta_inv.T) @ dATT_dbeta + + # ================================================================ + # Combine: full IF = main - PS correction - outcome correction + # ================================================================ + psi_full = psi - ps_adjustment - om_adjustment + + # --- Variance estimation --- + if cluster_ids is not None: + # Cluster-robust: sum phi within clusters, then outer product + cluster_df = pd.DataFrame({"psi": psi_full, "cluster": cluster_ids}) + cluster_sums = cluster_df.groupby("cluster")["psi"].sum().values + n_clusters = len(cluster_sums) + if n_clusters <= 1: + warnings.warn( + "Only 1 cluster found; falling back to non-clustered " + "variance for DR influence function.", + UserWarning, + stacklevel=2, + ) + var_att = float(np.var(psi_full, ddof=1) / n_obs) + else: + var_att = float( + (n_clusters / (n_clusters - 1)) * np.sum(cluster_sums**2) / n_obs**2 + ) + else: + var_att = float(np.var(psi_full, ddof=1) / n_obs) + + se = float(np.sqrt(max(var_att, 0.0))) + + # Effective n_params: treatment dimension + the IDENTIFIED + # outcome-model rank (round-12 review: the nominal count let an + # exactly redundant control shrink residual df and move + # p-values/CIs while ATT and SE were unchanged). + n_params = 1 + int(kept_om.sum()) + influence = self._finalize_influence( + self._moment_influence(psi_full, n_obs, cluster_ids), se + ) + return att, se, None, None, n_params, influence + + @staticmethod + def _validate_vcov_config(vcov_type, estimation_method, cluster) -> None: + """Config-only vcov coherence checks (called from __init__ AND fit). + + Accepted sets (campaign finding: vcov_type was silently inert for + ipw/dr/psm - the influence-function / matching variance is always + used there, so only the value whose behavior is real is accepted): + + - ``reg``: {classical, hc1, hc2, hc3} + - ``ipw`` / ``dr``: {hc1} only (the default; implemented as the + heteroskedasticity-robust influence-function sandwich) + - ``psm``: {hc1} only, and ``cluster=`` is rejected (the matching + SE has no clustered form - pre-fix it presented a non-clustered + SE under cluster-robust metadata) + + ``cluster=`` composes ONLY with hc1 (for any method): pre-fix, + classical/hc2/hc3 + cluster were silently remapped to CR1 while + the results object kept the requested label. + """ + if estimation_method in ("ipw", "dr") and vcov_type != "hc1": + raise ValueError( + f"estimation_method='{estimation_method}' supports " + f"vcov_type='hc1' only (the influence-function sandwich; " + f"heteroskedasticity-robust by construction). Got " + f"vcov_type='{vcov_type}', which would be silently inert." + ) + if estimation_method == "psm": + if vcov_type != "hc1": + raise ValueError( + f"estimation_method='psm' accepts vcov_type='hc1' only " + f"(the accepted configuration; matching inference is " + f"currently unavailable and reported as NaN - see " + f"DEFERRED.md). Got vcov_type='{vcov_type}', which would " + f"be silently inert." + ) + if cluster is not None: + raise ValueError( + "estimation_method='psm' does not support cluster=: the " + "matching SE has no cluster-robust form. Use " + "estimation_method='dr' for a doubly robust alternative " + "with clustered inference." + ) + if cluster is not None and vcov_type not in ("hc1",): + raise ValueError( + f"cluster= composes only with vcov_type='hc1' (CR1); got " + f"vcov_type='{vcov_type}'. Cluster-robust leverage-corrected " + f"families are not implemented for LWDiD." + ) + + def _resolve_vcov_type(self) -> str: + """Map the requested variance family to a solve_ols vcov_type. + + Returns + ------- + str + The vcov_type string compatible with solve_ols. When the + cluster= constructor parameter is set, cluster-robust (CR1) + inference is requested via hc1 plus cluster_ids. + """ + # Post fix-wave WS6, the config validator guarantees cluster only + # composes with hc1, so the requested family IS the resolved family + # on every path (no silent remap can occur). + if self.cluster is not None: + assert self.vcov_type == "hc1", "validator invariant violated" + return self.vcov_type + + def _bootstrap( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + treatment: str, + cluster: Optional[str], + controls: List[str], + pre_periods: List[Any], + post_periods: List[Any], + treated_units: List[Any], + control_units: List[Any], + ) -> Tuple[float, float, float, float, Tuple[float, float], int]: + """Compute bootstrap standard errors. + + Uses unit-level block bootstrap for panel data; when ``cluster`` is + set, whole clusters are resampled instead (a cluster may contain + both treated and control units, so replicates with an empty arm are + counted as failed). + + Parameters + ---------- + df : pd.DataFrame + Full panel data. + outcome : str + Outcome column name. + unit : str + Unit identifier column name. + time : str + Time period column name. + treatment : str + Treatment indicator column name. + cluster : str or None + Cluster column name. + controls : list of str + Control variable column names. + pre_periods : list + Pre-treatment period values. + post_periods : list + Post-treatment period values. + treated_units : list + Treated unit identifiers. + control_units : list + Control unit identifiers. + + Returns + ------- + att : float + Point estimate from full sample. + se : float + Bootstrap standard error. + t_stat : float + t-statistic. + p_value : float + Two-sided p-value. + conf_int : tuple of float + Confidence interval (lower, upper). + df_used : int + Degrees of freedom the p-value/CI actually used (G-1 when + clustered, N-k otherwise) - stored as ``df_inference``. + """ + # Full-sample estimate + treated_set = set(treated_units) + pre_mask = df[time].isin(pre_periods) + if self.rolling == "demean": + df_t = self._transform_demean(df, outcome, unit, pre_mask) + elif self.rolling == "detrend": + df_t = self._transform_detrend(df, outcome, unit, time, pre_mask) + elif self.rolling == "demeanq": + df_t = self._transform_demeanq(df, outcome, unit, time, pre_mask) + elif self.rolling == "detrendq": + df_t = self._transform_detrendq(df, outcome, unit, time, pre_mask) + else: + df_t = self._transform_detrend(df, outcome, unit, time, pre_mask) + + post_mask = df_t[time].isin(post_periods) # type: ignore[union-attr, call-overload] + post_df = df_t.loc[post_mask] # type: ignore[union-attr] + # Same fixed-window complete-case rule as _fit_common_timing (the + # bootstrap must estimate the same estimand as the point path). + post_counts = post_df.loc[np.isfinite(post_df["_ydot"])].groupby(unit)["_ydot"].size() + complete_units = set(post_counts.index[post_counts == len(post_periods)]) + post_df = post_df.loc[post_df[unit].isin(complete_units)] + unit_post_avg = post_df.groupby(unit)["_ydot"].mean() + + cs_df = df.drop_duplicates(subset=[unit], keep="first")[[unit] + controls].copy() + cs_df["_treat"] = cs_df[unit].isin(treated_set).astype(float) + cs_df["_ydot_avg"] = cs_df[unit].map(unit_post_avg) + cs_df = cs_df.dropna(subset=["_ydot_avg"]) + + y_full = cs_df["_ydot_avg"].values.astype(np.float64) + treat_full = cs_df["_treat"].values.astype(np.float64) + controls_mat = cs_df[controls].values.astype(np.float64) if controls else None + + att_full, _, _, _, n_params_full, _ = self._dispatch_estimator( + y_full, treat_full, controls_mat, None, len(y_full) + ) + + # Bootstrap replications. Resampling level: units (default) or + # whole CLUSTERS when cluster= is set (campaign finding: the + # cluster parameter was silently ignored here, producing an iid + # unit bootstrap labeled as clustered). The resampling population + # is restricted to units SURVIVING the transformation and finite + # filters (round-5 review: the raw-panel population let a raw + # cluster map claim G clusters when fewer contribute). + surviving = set(cs_df[unit]) + treated_arr = np.array([u for u in treated_units if u in surviving]) + control_arr = np.array([u for u in control_units if u in surviving]) + n_treated = len(treated_arr) + n_control = len(control_arr) + unit_counts = df.groupby(unit).size().to_dict() + + # Positional row map computed ONCE (campaign finding: the previous + # code collected index LABELS then fetched rows POSITIONALLY via + # .iloc - crashing on non-default indexes and silently resampling + # the wrong rows when labels were permuted relative to positions). + unit_col_arr = df[unit].to_numpy() + all_unit_ids = np.concatenate([treated_arr, control_arr]) + unit_positions = {u: np.flatnonzero(unit_col_arr == u) for u in all_unit_ids} + + cluster_draw: Optional[Dict[Any, np.ndarray]] = None + if cluster is not None: + # cluster is unit-constant (validated); map cluster -> units. + cluster_by_unit = df.drop_duplicates(subset=[unit], keep="first").set_index(unit)[ + cluster + ] + cluster_lists: Dict[Any, List[Any]] = {} + for u in all_unit_ids: + cluster_lists.setdefault(cluster_by_unit[u], []).append(u) + cluster_draw = {cl: np.asarray(us) for cl, us in cluster_lists.items()} + if len(cluster_draw) < 2: + # Fewer than 2 effective clusters survive the + # transformation: clustered bootstrap inference is not + # identified (round-5 review - the raw cluster map + # previously produced a near-zero SE here). Point + # retained, inference NaN. + warnings.warn( + "LWDiD bootstrap: fewer than 2 effective clusters " + "survive the transformation; clustered bootstrap " + "inference is not identified (NaN).", + UserWarning, + stacklevel=2, + ) + return att_full, np.nan, np.nan, np.nan, (np.nan, np.nan), 0 + treated_set_all = set(treated_units) + + def _draw_units(rng_b: np.random.Generator) -> np.ndarray: + if cluster_draw is not None: + # Cluster-level draws (no treated/control stratification: a + # cluster may contain both arms). A draw collapsing onto a + # single distinct cluster carries no between-cluster + # variation - counted as a failed replicate (round-5 + # review), signalled by an empty draw. + cluster_keys = list(cluster_draw) + picks = rng_b.choice(len(cluster_keys), size=len(cluster_keys), replace=True) + if len(set(picks.tolist())) < 2: + return np.array([], dtype=object) + return np.concatenate([cluster_draw[cluster_keys[i]] for i in picks]) + boot_treated = rng_b.choice(treated_arr, size=n_treated, replace=True) + boot_control = rng_b.choice(control_arr, size=n_control, replace=True) + return np.concatenate([boot_treated, boot_control]) + + def _replicate_att(boot_units: np.ndarray) -> float: + """Estimate one bootstrap replicate (shared serial/parallel).""" + if boot_units.size == 0: + return np.nan # single-distinct-cluster draw (failed) + boot_indices = np.concatenate([unit_positions[u] for u in boot_units]) + boot_df = df.iloc[boot_indices].copy() + # Occurrence-specific synthetic unit ids keep duplicate draws + # distinct through the transform step. + repeat_counts = [unit_counts[u] for u in boot_units] + boot_df["_boot_unit"] = np.repeat(np.arange(len(boot_units)), repeat_counts) + + # Treatment is EVER-TREATED MEMBERSHIP of the source unit, not + # the collapsed row's time-varying D (with unsorted input the + # drop_duplicates(keep="first") row is arbitrary - typically + # pre-treatment, which would zero the treatment vector). + boot_treat_vec = np.array( + [1.0 if u in treated_set_all else 0.0 for u in boot_units], dtype=np.float64 + ) + if boot_treat_vec.sum() == 0 or boot_treat_vec.sum() == len(boot_units): + return np.nan # invalid replicate: an arm is empty + + # Apply transformation + pre_mask_b = boot_df[time].isin(pre_periods) + if self.rolling == "demean": + boot_df = self._transform_demean(boot_df, outcome, "_boot_unit", pre_mask_b) + elif self.rolling == "detrend": + boot_df = self._transform_detrend(boot_df, outcome, "_boot_unit", time, pre_mask_b) + elif self.rolling == "demeanq": + boot_df = self._transform_demeanq(boot_df, outcome, "_boot_unit", time, pre_mask_b) + elif self.rolling == "detrendq": + boot_df = self._transform_detrendq(boot_df, outcome, "_boot_unit", time, pre_mask_b) + else: + boot_df = self._transform_detrend(boot_df, outcome, "_boot_unit", time, pre_mask_b) + + # Cross-sectional estimate + post_mask_b = boot_df[time].isin(post_periods) # type: ignore[union-attr, call-overload] + post_b = boot_df.loc[post_mask_b] # type: ignore[union-attr] + unit_avg_b = post_b.groupby("_boot_unit")["_ydot"].mean() + + first_rows = boot_df.drop_duplicates(subset=["_boot_unit"], keep="first") # type: ignore[union-attr] + cs_b = first_rows[["_boot_unit"]].copy() + if controls: + for c in controls: + cs_b[c] = first_rows[c].values + + cs_b["_treat"] = cs_b["_boot_unit"].map( + dict(zip(range(len(boot_units)), boot_treat_vec)) + ) + cs_b["_ydot_avg"] = cs_b["_boot_unit"].map(unit_avg_b) + cs_b = cs_b.dropna(subset=["_ydot_avg"]) + + if len(cs_b) < 3: + return np.nan + + y_b = cs_b["_ydot_avg"].values.astype(np.float64) + treat_b = cs_b["_treat"].values.astype(np.float64) + ctrl_b = cs_b[controls].values.astype(np.float64) if controls else None + + try: + att_b, _, _, _, _, _ = self._dispatch_estimator( + y_b, treat_b, ctrl_b, None, len(y_b) + ) + return float(att_b) + except (np.linalg.LinAlgError, ValueError): + return np.nan + + # Per-replicate RNG streams via SeedSequence spawning, IDENTICAL for + # every n_jobs: replicate b always draws from child stream b, so a + # seeded fit is reproducible regardless of the execution mode + # (review round 2: the serial path consumed one sequential stream + # while the parallel path spawned, so the same seed produced + # different bootstrap SEs across n_jobs). seed=None still draws + # fresh OS entropy (non-deterministic). + seed_seq = np.random.SeedSequence(self.seed) + child_seqs = seed_seq.spawn(self.n_bootstrap) + boot_unit_samples = [ + _draw_units(np.random.default_rng(child_seqs[b])) for b in range(self.n_bootstrap) + ] + + if self.n_jobs == 1: + # --- Serial path --- + boot_atts = np.array([_replicate_att(sample) for sample in boot_unit_samples]) + else: + # --- Parallel path (n_jobs > 1) --- + from concurrent.futures import ThreadPoolExecutor + + warnings.warn( + "Parallel bootstrap (n_jobs > 1) is experimental. " + "ThreadPoolExecutor is used; speedup depends on " + "GIL-releasing operations in numpy/scipy.", + UserWarning, + stacklevel=2, + ) + + with ThreadPoolExecutor(max_workers=self.n_jobs) as executor: + boot_atts = np.array(list(executor.map(_replicate_att, boot_unit_samples))) + + # Compute bootstrap SE + n_failed = int(np.isnan(boot_atts).sum()) + if n_failed > 0: + warnings.warn( + f"LWDiD bootstrap: {n_failed}/{self.n_bootstrap} replication(s) failed " + f"(returned NaN). Results based on {self.n_bootstrap - n_failed} valid replications.", + UserWarning, + stacklevel=2, + ) + valid_boots = boot_atts[np.isfinite(boot_atts)] + if len(valid_boots) < 2: + se = np.nan + else: + from diff_diff.lwdid_staggered import _guard_standard_error + + se = _guard_standard_error( + att_full, + float(np.std(valid_boots, ddof=1)), + scale=float(np.max(np.abs(y_full))) if len(y_full) else 0.0, + ) + + # df matches the analytical path's rule (campaign finding: the + # reported df_inference was G-1 under cluster= while the bootstrap + # p-value used N-k): G-1 when clustered, N-k otherwise. The df + # actually used is returned so the caller can store it. + if cluster_draw is not None: + df_used = max(len(cluster_draw) - 1, 1) + else: + df_used = len(y_full) - n_params_full + t_stat, p_value, conf_int = safe_inference(att_full, se, alpha=self.alpha, df=df_used) + + return att_full, se, t_stat, p_value, conf_int, df_used + + def __repr__(self) -> str: + """Return string representation of the estimator.""" + params = self.get_params() + params_str = ", ".join(f"{k}={v!r}" for k, v in params.items()) + return f"LWDiD({params_str})" + + +def validate_staggered_data(data, unit, time, cohort) -> Dict[str, Any]: + """Validate panel data structure for staggered DiD estimation. + + Checks (using the same never-treated definition as ``fit()``: + cohort ``NaN``/``NaT``, ``0``, ``np.inf`` (recoded), or a finite value + beyond the last observed period (recoded)): + + - Panel is complete (all unit×time combinations exist) + - Cohort is time-invariant within units (missing values included, + matching ``fit_staggered``'s ``nunique(dropna=False)`` check) + - At least one never-treated unit exists + - At least one treated cohort remains after normalization + - Time and cohort columns share the same time family + + Parameters + ---------- + data : pd.DataFrame + Panel dataset. + unit : str + Unit identifier column name. + time : str + Time period column name. + cohort : str + Cohort column name (``0``/``NaN``/``NaT`` = never-treated; + ``np.inf`` and beyond-window values are recoded to never-treated + with a warning, as in ``fit()``). + + Returns + ------- + dict + Validation results with keys: 'valid', 'warnings', 'errors', + 'n_units', 'n_periods', 'n_cohorts', 'n_never_treated'. + """ + + df = data.copy() + + results: dict[str, Any] = {"valid": True, "warnings": [], "errors": []} + + # Check required columns exist + for col in [unit, time, cohort]: + if col not in df.columns: + results["valid"] = False + results["errors"].append(f"Column '{col}' not found in data") + return results + + # Dtype coherence + normalization (dtype-aware; mirrors fit()'s + # encode-then-normalize pipeline without building position maps). + # Round-8 review: datetime64 and Period are DISTINCT families (mixing + # them crashes pandas position lookups), so both directions and + # Period-frequency mismatches are rejected exactly like + # _encode_staggered_time_scale - a single "datelike" flag previously + # let a datetime time column meet a Period cohort column in an + # invalid cross-family comparison below. + cohort_is_datetime = pd.api.types.is_datetime64_any_dtype(df[cohort]) + time_is_datetime = pd.api.types.is_datetime64_any_dtype(df[time]) + cohort_is_period = isinstance(df[cohort].dtype, pd.PeriodDtype) + time_is_period = isinstance(df[time].dtype, pd.PeriodDtype) + cohort_datelike = cohort_is_datetime or cohort_is_period + if time_is_datetime != cohort_is_datetime or time_is_period != cohort_is_period: + results["valid"] = False + results["errors"].append( + f"Columns '{time}' (time) and '{cohort}' (cohort) must share the " + f"same time scale; got dtypes {df[time].dtype} and {df[cohort].dtype}. " + f"Encode both as datetime64, both as Period with the same " + f"frequency, or both as numeric." + ) + return results + if time_is_period and df[time].dtype.freq != df[cohort].dtype.freq: + results["valid"] = False + results["errors"].append( + f"Columns '{time}' (time) and '{cohort}' (cohort) are Period " + f"columns with different frequencies ({df[time].dtype} vs " + f"{df[cohort].dtype}). Convert them to a common frequency." + ) + return results + if cohort_datelike: + # Datetime/Period: NaT = never-treated; beyond-window recodes to + # NaT via a same-dtype comparison (no numeric sentinel exists). + max_time = df[time].max() + beyond = df[cohort].notna() & (df[cohort] > max_time) + if beyond.to_numpy().any(): + results["warnings"].append( + f"{int(beyond.sum())} row(s) have cohort values beyond the " + f"last observed period ({max_time}); treated as never-treated." + ) + df.loc[beyond, cohort] = pd.NaT + never_mask_series = df[cohort].isna() + treated_vals = df.loc[~never_mask_series, cohort] + else: + try: + df[cohort], _, _ = _normalize_cohorts(df[cohort], max_time=df[time].max()) + except ValueError as exc: + results["valid"] = False + results["errors"].append(str(exc)) + return results + never_mask_series = df[cohort].isna() | (df[cohort] == 0) + treated_vals = df.loc[~never_mask_series, cohort] + + # Check cohort time-invariance (missing values included, matching + # fit_staggered's nunique(dropna=False) — a unit mixing NaT/NaN with a + # finite cohort must fail here, not later inside fit). + cohort_per_unit = df.groupby(unit)[cohort].nunique(dropna=False) + varying = cohort_per_unit[cohort_per_unit > 1] + if len(varying) > 0: + results["valid"] = False + results["errors"].append(f"{len(varying)} units have time-varying cohort values") + + # Check for never-treated. All-eventually-treated panels are rejected + # by fit() for staggered designs, so mirror that hard-error here + # instead of reporting a valid-with-warning contradiction. + never_treated = df.loc[never_mask_series, unit].nunique() + if never_treated == 0: + results["valid"] = False + results["errors"].append( + "No never-treated units found (cohort NaN/NaT or 0); staggered " + "LWDiD estimation requires a never-treated control group." + ) + + # Treated-cohort coherence: if normalization recoded every cohort, + # fit_staggered would raise "No treated cohorts found." — report the + # same failure here instead of valid-with-zero-cohorts. + n_cohorts = int(treated_vals.nunique()) + if n_cohorts == 0: + results["valid"] = False + results["errors"].append("No treated cohorts found.") + + # Duplicate (unit, time) cells are INVALID (fit() rejects them), and + # a duplicate can mask a missing cell in the row-count balance check + # below (round-11 review). + n_dup_cells = int(df.duplicated(subset=[unit, time]).sum()) + if n_dup_cells > 0: + results["valid"] = False + results["errors"].append( + f"{n_dup_cells} duplicate (unit, time) observation(s); each " f"pair must be unique." + ) + + # Check panel balance + n_units = df[unit].nunique() + n_times = df[time].nunique() + expected_rows = n_units * n_times + if len(df) - n_dup_cells != expected_rows: + results["warnings"].append( + f"Unbalanced panel: {len(df) - n_dup_cells} distinct cell(s) vs " + f"{expected_rows} expected" + ) + + # Missing unit/time values are ERRORS (fit() rejects the same frame; + # round-23 review: warning-only let 'valid: True' disagree with fit). + # Cohort NaN/NaT stays a documented never-treated encoding. + for col in [unit, time]: + n_missing = df[col].isna().sum() + if n_missing > 0: + results["valid"] = False + results["errors"].append(f"{n_missing} missing values in '{col}'") + + results["n_units"] = n_units + results["n_periods"] = n_times + results["n_cohorts"] = n_cohorts + results["n_never_treated"] = never_treated + + return results + + +def is_never_treated(data, unit, cohort, time=None) -> np.ndarray: + """Identify never-treated units in staggered design. + + Parameters + ---------- + data : pd.DataFrame + Panel dataset. + unit : str + Unit identifier column name. + cohort : str + Cohort column name. Never-treated encodings: ``0``, ``NaN``/``NaT``, + and ``np.inf``. + time : str or None, default None + Time column name. When provided, beyond-window classification also + applies (dtype-aware): a finite cohort value greater than the last + observed period counts as never-treated, matching ``fit()``'s + normalization. Without it, only the sentinel encodings above are + classified. + + Returns + ------- + np.ndarray of bool + True for never-treated units (one entry per unique unit). + """ + unit_cohort = data.groupby(unit)[cohort].first() + never = (unit_cohort == 0) | unit_cohort.isna() + if pd.api.types.is_numeric_dtype(unit_cohort): + never |= np.isposinf(unit_cohort.to_numpy(dtype=float, na_value=np.nan)) + if time is not None: + max_time = data[time].max() + never |= unit_cohort.notna() & (unit_cohort > max_time) + return np.asarray(never) diff --git a/diff_diff/lwdid_randomization.py b/diff_diff/lwdid_randomization.py new file mode 100644 index 00000000..96d751ad --- /dev/null +++ b/diff_diff/lwdid_randomization.py @@ -0,0 +1,508 @@ +"""Randomization inference for LWDiD estimator. + +Implements Fisher's randomization inference under the sharp null +hypothesis H0: τ_i = 0 for all i (no individual treatment effect). + +References +---------- +Fisher, R. A. (1935). The Design of Experiments. +Lee, S. J. & Wooldridge, J. M. (2026). "Simple Approaches to Inference + with Difference-in-Differences Estimators with Small Cross-Sectional + Sample Sizes." SSRN 5325686 (randomization inference for small-N + designs; implemented per the authors'-package inclusive convention). +""" + +import warnings +from dataclasses import dataclass +from typing import Optional + +import numpy as np + +from diff_diff.linalg import solve_ols + + +@dataclass +class RandomizationResult: + """Result container for randomization inference. + + Attributes + ---------- + pvalue : float + Two-sided p-value from the randomization distribution. + att_observed : float + Observed ATT estimate from the original data. + att_distribution : np.ndarray + Array of ATT estimates from randomization replications (includes NaN + for failed replications). + n_reps : int + Total number of replications requested. + n_valid : int + Number of valid (non-degenerate) replications used for p-value. + n_failed : int + Number of failed or degenerate replications. + failure_rate : float + Proportion of replications that failed (n_failed / n_reps). + method : str + Resampling method used: always 'permutation'. + seed : int or None + Random seed used for reproducibility. + n_dropped : int + Observations dropped for non-finite y before estimation (warned). + """ + + pvalue: float + att_observed: float + att_distribution: np.ndarray + n_reps: int + n_valid: int + n_failed: int + failure_rate: float + method: str + seed: Optional[int] + #: Observations dropped for non-finite y before estimation (warned). + n_dropped: int = 0 + + +def _validate_inputs( + y: np.ndarray, + treatment: np.ndarray, + controls: Optional[np.ndarray], + n_reps: int, + method: str, +) -> None: + """Validate inputs for randomization inference. + + Raises + ------ + ValueError + If any validation check fails. + """ + if ( + n_reps is None + or isinstance(n_reps, bool) + or not isinstance(n_reps, (int, np.integer)) + or n_reps < 10 + ): + # Round-24 review: the valid-replication floor below is + # max(10, ...), so n_reps < 10 can NEVER satisfy it - reject up + # front instead of failing after the permutation loop. + raise ValueError( + f"n_reps must be an integer >= 10 (the reliable-inference " + f"floor requires at least 10 valid replications), got {n_reps!r}" + ) + + if method == "bootstrap": + # Review finding: resampling treatment labels WITH replacement + # changes the treated count and is not the complete-randomization + # assignment mechanism of Fisher randomization inference - it was + # presented under the Fisher umbrella without a specified + # assignment design. The mode is removed (LWDiD is unreleased). + raise ValueError( + "method='bootstrap' has been removed: resampling treatment " + "labels with replacement is not Fisher randomization inference " + "(it changes the treated count and has no specified assignment " + "mechanism). Use method='permutation'." + ) + if method != "permutation": + raise ValueError(f"method must be 'permutation', got '{method}'") + + if y.ndim != 1: + raise ValueError(f"y must be a 1-d array, got shape {y.shape}") + + if treatment.ndim != 1: + raise ValueError(f"treatment must be a 1-d array, got shape {treatment.shape}") + + if len(y) == 0: + raise ValueError("y must not be empty.") + + if len(y) != len(treatment): + raise ValueError( + f"y and treatment must have the same length, " f"got {len(y)} and {len(treatment)}" + ) + + n = len(y) + if n < 3: + raise ValueError(f"Sample size too small for randomization inference: N={n}") + + if not np.all((treatment == 0) | (treatment == 1)): + raise ValueError( + "treatment must be binary (0 or 1). " + f"Got values in [{treatment.min()}, {treatment.max()}]." + ) + + n1 = int(treatment.sum()) + if n1 == 0 or n1 == n: + raise ValueError( + "Treatment variable is constant (all treated or all control). " + "Randomization inference requires variation in treatment." + ) + + if controls is not None: + if controls.ndim == 1: + controls = controls.reshape(-1, 1) + if controls.shape[0] != n: + raise ValueError(f"controls must have {n} rows, got {controls.shape[0]}") + if not np.all(np.isfinite(controls)): + raise ValueError( + "controls contains non-finite values (NaN or Inf). " + "Please remove or impute missing values before calling " + "randomization_inference()." + ) + + +def _build_design( + y: np.ndarray, + treatment: np.ndarray, + controls: np.ndarray, + design: str, +) -> np.ndarray: + """Build the regression design for a given treatment assignment. + + ``'linear'`` is the generic ``[1, D, X]`` covariate-adjusted contrast. + ``'ra_interacted'`` is the LWDiD RA design ``[1, D, X, D(X - Xbar_1)]`` + (LW eq. E.1) with the treated covariate mean RECOMPUTED for the given + assignment - required so each permutation tests the same estimator the + fit reported (round-5 review). + """ + n = len(y) + if design == "ra_interacted": + xbar1 = controls[treatment == 1].mean(axis=0) + return np.column_stack( + [np.ones(n), treatment, controls, treatment[:, None] * (controls - xbar1)] + ) + return np.column_stack([np.ones(n), treatment, controls]) + + +def _compute_observed_att( + y: np.ndarray, + treatment: np.ndarray, + controls: Optional[np.ndarray], + design: str = "linear", +) -> float: + """Compute the observed ATT from the data. + + When controls are present, uses OLS via the shared solve_ols. + Otherwise computes the simple mean difference. + """ + if controls is None: + mask1 = treatment == 1 + return float(y[mask1].mean() - y[~mask1].mean()) + + if controls.ndim == 1: + controls = controls.reshape(-1, 1) + X = _build_design(y, treatment, controls, design) + # Rank-aware shared solver (round-4 review: lstsq returned a finite + # minimum-norm treatment coefficient when a control duplicated the + # treatment column, so RI tested an unidentified statistic). + coefs, _, _ = solve_ols(X, y) + if not np.isfinite(coefs[1]): + raise ValueError( + "The treatment coefficient is not identified: the design is " + "rank-deficient and the shared solver dropped the treatment " + "column (e.g. a control collinear with treatment). Remove the " + "collinear control(s) before running randomization inference." + ) + return float(coefs[1]) + + +def _fast_path( + y: np.ndarray, + treatment: np.ndarray, + n_reps: int, + method: str, + rng: np.random.Generator, +) -> np.ndarray: + """Fast path: no controls, direct mean-difference computation. + + Returns + ------- + att_dist : ndarray of shape (n_reps,) + Randomization distribution of ATT. Failed reps contain NaN. + """ + n = len(y) + att_dist = np.empty(n_reps) + + for b in range(n_reps): + if method == "permutation": + d_b = rng.permutation(treatment) + else: + d_b = rng.choice(treatment, size=n, replace=True) + + n1_b = d_b.sum() + if n1_b == 0 or n1_b == n: + att_dist[b] = np.nan + continue + + mask1 = d_b == 1 + att_dist[b] = y[mask1].mean() - y[~mask1].mean() + + return att_dist + + +def _slow_path( + y: np.ndarray, + treatment: np.ndarray, + controls: np.ndarray, + n_reps: int, + method: str, + rng: np.random.Generator, + design: str = "linear", +) -> np.ndarray: + """Slow path: with controls, OLS via pre-allocated design matrix. + + The design matrix is pre-allocated and only the treatment column + (column 1) is updated per replication. This avoids repeated memory + allocation and keeps the cost to O(N*K) per iteration. + + Returns + ------- + att_dist : ndarray of shape (n_reps,) + Randomization distribution of ATT. Failed reps contain NaN. + """ + n = len(y) + if controls.ndim == 1: + controls = controls.reshape(-1, 1) + + att_dist = np.empty(n_reps) + + for b in range(n_reps): + if method == "permutation": + d_b = rng.permutation(treatment) + else: + d_b = rng.choice(treatment, size=n, replace=True) + + n1_b = d_b.sum() + if n1_b == 0 or n1_b == n: + att_dist[b] = np.nan + continue + + # The design is rebuilt per assignment: under 'ra_interacted' the + # treated covariate mean (and the interaction columns) depend on + # the drawn assignment (round-5 review - the pre-fix code updated + # only the treatment column of a fixed [1, D, X] matrix). + X = _build_design(y, d_b, controls, design) + + try: + with warnings.catch_warnings(): + # Rank warnings per draw would flood; a dropped treatment + # coefficient is recorded as a failed replication (NaN) + # and surfaced through the failed-rep accounting. + warnings.simplefilter("ignore") + coefs, _, _ = solve_ols(X, y) + att_dist[b] = coefs[1] if np.isfinite(coefs[1]) else np.nan + except (np.linalg.LinAlgError, ValueError): + att_dist[b] = np.nan + + return att_dist + + +def _compute_pvalue(att_dist: np.ndarray, att_obs: float) -> tuple: + """Compute two-sided p-value from randomization distribution. + + Uses the formula: p = (sum(|ATT*| >= |ATT_obs|) + 1) / (n_valid + 1) + following Phipson & Smyth (2010). The non-strict inequality counts + replications at least as extreme as the observed statistic, so a + fully tied distribution (e.g. constant outcome) yields p = 1.0, + while the +1 in numerator and denominator accounts for the observed + statistic itself and guarantees p > 0. + + Returns + ------- + pvalue : float + n_valid : int + n_failed : int + """ + valid_mask = np.isfinite(att_dist) + n_valid = int(valid_mask.sum()) + n_failed = len(att_dist) - n_valid + + if n_valid == 0: + return 1.0, 0, n_failed + + valid_atts = att_dist[valid_mask] + pvalue = float((np.sum(np.abs(valid_atts) >= np.abs(att_obs)) + 1) / (n_valid + 1)) + return pvalue, n_valid, n_failed + + +def randomization_inference( + y: np.ndarray, + treatment: np.ndarray, + controls: Optional[np.ndarray] = None, + n_reps: int = 1000, + method: str = "permutation", + seed: Optional[int] = None, + design: str = "linear", +) -> RandomizationResult: + """Fisher randomization inference for testing zero treatment effect. + + Tests the sharp null hypothesis H0: τ_i = 0 for all i by permuting + treatment labels and computing a Monte Carlo p-value + as the proportion of resampled test statistics at least as extreme as + the observed statistic. + + Parameters + ---------- + y : ndarray of shape (n,) + Transformed outcome variable. + treatment : ndarray of shape (n,) + Binary treatment indicator (0/1). + controls : ndarray of shape (n, K) or None, optional + Control variables to include in the regression model. When None, + ATT is computed as a simple mean difference (fast path). When + provided, ATT is estimated via OLS with controls (slow path). + n_reps : int, default 1000 + Number of randomization replications for computing the p-value. + method : {'permutation'}, default 'permutation' + Resampling method: + + - 'permutation': Classical Fisher randomization inference. Permutes + treatment labels without replacement, preserving the original + number of treated and control units. + + seed : int or None, optional + Random seed for reproducibility. + design : {'linear', 'ra_interacted'}, default 'linear' + Regression design used for the covariate-adjusted statistic. + ``'linear'`` fits ``[1, D, X]``. ``'ra_interacted'`` fits the LWDiD + RA design ``[1, D, X, D(X - Xbar_1)]`` and RECOMPUTES the treated + covariate mean for every permuted assignment, so the permuted + statistic is the same estimator as the observed one (used by + ``LWDiDResults.randomization_test`` to match the fitted ATT). + + Returns + ------- + RandomizationResult + Dataclass containing p-value, observed ATT, randomization + distribution, and diagnostic information. + + Raises + ------ + ValueError + If inputs are invalid, sample size is too small, treatment is + constant, or insufficient valid replications are produced. + + Notes + ----- + The p-value is computed as: + + p = (sum(|ATT*| >= |ATT_obs|) + 1) / (n_valid + 1) + + following Phipson & Smyth (2010). The non-strict inequality counts + replications at least as extreme as the observed statistic (standard + randomization-test convention, so a degenerate all-tie distribution + yields p = 1.0), while the +1 ensures the p-value is strictly + positive and provides valid finite-sample inference. + + When controls are absent, ATT is computed directly as the difference + in means between treated and control groups. With controls, a + pre-allocated design matrix is refit through the shared rank-aware + ``solve_ols`` solver (draws whose treatment coefficient is dropped + count as failed replications). + + Examples + -------- + >>> import numpy as np + >>> from diff_diff.lwdid_randomization import randomization_inference + >>> rng = np.random.default_rng(42) + >>> y = rng.normal(0, 1, 100) + >>> y[:30] += 2.0 + >>> treatment = np.zeros(100); treatment[:30] = 1.0 + >>> r = randomization_inference(y, treatment, n_reps=999, seed=0) + >>> r.pvalue < 0.05 + True + """ + # ------------------------------------------------------------------ + # Input validation + # ------------------------------------------------------------------ + y = np.asarray(y, dtype=np.float64) + treatment = np.asarray(treatment, dtype=np.float64) + + if controls is not None: + controls = np.asarray(controls, dtype=np.float64) + if controls.ndim == 1: + controls = controls.reshape(-1, 1) + + # Shape/length coherence FIRST (round-18 review: applying the finite + # mask before these checks turned a mismatched treatment/controls + # length into a raw boolean-index IndexError instead of the + # documented ValueError). + if y.ndim != 1: + raise ValueError(f"y must be a 1-d array, got shape {y.shape}") + if treatment.ndim != 1: + raise ValueError(f"treatment must be a 1-d array, got shape {treatment.shape}") + if len(y) != len(treatment): + raise ValueError( + f"y and treatment must have the same length, got {len(y)} and {len(treatment)}" + ) + if controls is not None and controls.shape[0] != len(y): + raise ValueError(f"controls must have {len(y)} rows, got {controls.shape[0]}") + + # Drop observations with non-finite y WITH a warning (campaign + # finding: silent drops) and record the count on the result. + n_dropped = 0 + if len(y) > 0: + finite_mask = np.isfinite(y) + if not finite_mask.all(): + n_dropped = int((~finite_mask).sum()) + warnings.warn( + f"randomization_inference: dropped {n_dropped} observation(s) " + f"with non-finite y before estimation.", + UserWarning, + stacklevel=2, + ) + y = y[finite_mask] + treatment = treatment[finite_mask] + if controls is not None: + controls = controls[finite_mask] + + _validate_inputs(y, treatment, controls, n_reps, method) + + # ------------------------------------------------------------------ + # Compute observed ATT + # ------------------------------------------------------------------ + if design not in ("linear", "ra_interacted"): + raise ValueError(f"design must be 'linear' or 'ra_interacted', got {design!r}") + if design == "ra_interacted" and controls is None: + raise ValueError( + "design='ra_interacted' requires controls (the design is [1, D, X, D(X - Xbar_1)])." + ) + + att_obs = _compute_observed_att(y, treatment, controls, design) + + # ------------------------------------------------------------------ + # Generate randomization distribution + # ------------------------------------------------------------------ + rng = np.random.default_rng(seed) + + if controls is None: + att_dist = _fast_path(y, treatment, n_reps, method, rng) + else: + att_dist = _slow_path(y, treatment, controls, n_reps, method, rng, design) + + # ------------------------------------------------------------------ + # Compute p-value and diagnostics + # ------------------------------------------------------------------ + pvalue, n_valid, n_failed = _compute_pvalue(att_dist, att_obs) + failure_rate = n_failed / n_reps + + # Error if too few valid replications + if n_valid < max(10, int(0.1 * n_reps)): + raise ValueError( + f"Insufficient valid replications for reliable inference: " + f"{n_valid}/{n_reps} valid (failure rate {failure_rate:.1%}). " + f"Increase n_reps, or check for near-constant treatment/" + f"outcome configurations that make draws degenerate." + ) + + return RandomizationResult( + pvalue=pvalue, + att_observed=att_obs, + att_distribution=att_dist, + n_reps=n_reps, + n_valid=n_valid, + n_failed=n_failed, + failure_rate=failure_rate, + method=method, + seed=seed, + n_dropped=n_dropped, + ) diff --git a/diff_diff/lwdid_results.py b/diff_diff/lwdid_results.py new file mode 100644 index 00000000..cc80e1be --- /dev/null +++ b/diff_diff/lwdid_results.py @@ -0,0 +1,826 @@ +"""Results class for the LWDiD (Lee & Wooldridge 2025, 2026) estimator.""" + +from __future__ import annotations + +import datetime +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import pandas as pd + +from diff_diff.aggregation import AggregationMixin, AggregationResult +from diff_diff.results_base import BaseResults, EventStudyResults + + +# How the overall staggered standard error was obtained. Cohort effects that +# share control units are correlated, so the basis is reported rather than +# left implicit. +def _as_float(value: Any) -> float: + """Coerce an optional numeric cell entry to float, mapping None to NaN.""" + return np.nan if value is None else float(value) + + +_INFERENCE_BASIS_LABELS = { + "composite_regression": "composite regression (LW 2026 eq. 7.18/7.19)", + "joint_influence_function": "joint influence function across cohort-time cells", + "unavailable_matching": "unavailable (matching has no influence function)", + "unavailable_degenerate_cells": "unavailable (degenerate cohort-time cells)", + "unit_bootstrap": "unit-resampling bootstrap (params/vcov remain analytical)", + "cluster_bootstrap": "cluster-resampling bootstrap (params/vcov remain analytical)", +} + + +def _json_native_key(key: Any) -> Any: + """Convert a numpy scalar or datetime-like dict key to its native equivalent.""" + if isinstance(key, np.bool_): + return bool(key) + if isinstance(key, np.integer): + return int(key) + if isinstance(key, np.floating): + return float(key) + # pd.NaT is datetime-like but has no meaningful isoformat; keep the + # same convention as _to_json_native (NaT -> None) for consistency. + if key is pd.NaT: + return None + if isinstance(key, (datetime.date, datetime.datetime)): + # covers pd.Timestamp (subclass of datetime.datetime) + return key.isoformat() + if isinstance(key, np.datetime64): + return pd.Timestamp(key).isoformat() + if isinstance(key, pd.Period): + return str(key) # e.g. "2020Q1", preserves frequency semantics + return key + + +def _to_json_native(obj: Any) -> Any: + """Recursively convert numpy types to JSON-serializable Python natives. + + numpy scalars become int/float/bool, ndarrays become nested lists, + and dict/list/tuple containers are converted element-wise (dict keys + included). NaN/inf floats are kept as-is (float semantics preserved). + Datetime-like values (datetime.date/datetime.datetime incl. pd.Timestamp, + np.datetime64) become ISO-8601 strings; pd.Period becomes str (e.g. + "2020Q1") to preserve frequency semantics; pd.NaT becomes None so the + output is always json.dumps-able. + """ + if isinstance(obj, np.bool_): + return bool(obj) + if isinstance(obj, np.integer): + return int(obj) + if isinstance(obj, np.floating): + return float(obj) + if obj is pd.NaT: + return None + if isinstance(obj, (datetime.date, datetime.datetime)): + # covers pd.Timestamp (subclass of datetime.datetime) + return obj.isoformat() + if isinstance(obj, np.datetime64): + if pd.isna(obj): + return None + return pd.Timestamp(obj).isoformat() + if isinstance(obj, pd.Period): + return str(obj) # e.g. "2020Q1", preserves frequency semantics + if isinstance(obj, np.ndarray): + return [_to_json_native(v) for v in obj.tolist()] + if isinstance(obj, dict): + return {_json_native_key(k): _to_json_native(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_to_json_native(v) for v in obj] + return obj + + +@dataclass +class LWDiDResults(BaseResults, AggregationMixin): + """Results from LWDiD.fit(). + + Follows the diff-diff standard results interface. Holds the headline ATT + estimate and inference for the common-timing case, or per-cohort effects + and an overall weighted ATT for the staggered case. + + Parameters + ---------- + att : float + Average treatment effect on the treated. + se : float + Standard error of the ATT estimate. + t_stat : float + t-statistic (att / se). + p_value : float + Two-sided p-value. + conf_int : tuple of float + (lower, upper) confidence interval at level ``1 - alpha``. + n_obs : int + Total observations used in estimation. + n_treated : int + Number of treated units. + n_control : int + Number of control units. + rolling : str + Transformation method used ('demean', 'detrend', 'demeanq', or 'detrendq'). + estimation_method : str + Estimation method ('reg', 'ipw', 'dr', or 'psm'). + vcov_type : str + Variance family ('classical', 'hc1', 'hc2', or 'hc3'). + alpha : float + Significance level used for confidence intervals. + df_inference : int or None + Degrees of freedom used for t-distribution inference. + cluster_name : str or None + Name of the cluster variable, if clustered. + n_clusters : int or None + Number of clusters, if clustered. + cohort_effects : dict or None + Per-cohort ATT results for staggered designs. + params : ndarray or None + All coefficient estimates from the regression. + bse : ndarray or None + All standard errors from the regression. + vcov : ndarray or None + Variance-covariance matrix. + """ + + # ------------------------------------------------------------------ # + # Core inference fields # + # ------------------------------------------------------------------ # + att: float + se: float + t_stat: float + p_value: float + conf_int: Tuple[float, float] + + # ------------------------------------------------------------------ # + # Sample information # + # ------------------------------------------------------------------ # + n_obs: int + n_treated: int + n_control: int + + # ------------------------------------------------------------------ # + # Method metadata # + # ------------------------------------------------------------------ # + rolling: str + estimation_method: str + vcov_type: str + alpha: float + df_inference: Optional[int] = None + cluster_name: Optional[str] = None + n_clusters: Optional[int] = None + + # ------------------------------------------------------------------ # + # Fit provenance (estimand/inference-affecting configuration - review # + # finding: serialized results could not reconstruct what was fitted) # + # ------------------------------------------------------------------ # + control_group: Optional[str] = None + n_bootstrap: int = 0 + seed: Optional[int] = None + #: Propensity-score trim bound used by the ipw/dr/psm paths (None for + #: estimation_method='reg', where no propensity model is fitted). + pscore_trim: Optional[float] = None + #: PSM matching settings (None unless estimation_method='psm'): + #: {'pscore_trim', 'n_neighbors', 'caliper', 'with_replacement'} + psm_config: Optional[Dict[str, Any]] = None + + # ------------------------------------------------------------------ # + # Staggered-specific (optional) # + # ------------------------------------------------------------------ # + cohort_effects: Optional[Dict[Any, Dict]] = field(default=None, repr=False) + cohort_time_effects: Optional[Dict[Tuple[Any, Any], Dict]] = field(default=None, repr=False) + inference_basis: Optional[str] = None + #: Complete-case tau_omega composite point, exposed as a diagnostic when + #: complete-case drops prevented it from being ``.att`` (None otherwise). + att_tau_omega_complete_case: Optional[float] = None + #: Treated / control units dropped by the tau_omega complete-case + #: resolution (0 when the composite path did not run or dropped none). + n_composite_treated_dropped: int = 0 + n_composite_controls_dropped: int = 0 + + # ------------------------------------------------------------------ # + # Event study (Appendix D) fields # + # ------------------------------------------------------------------ # + event_study_effects: Optional[Dict[int, Dict]] = field(default=None, repr=False) + event_study_vcov: Optional[np.ndarray] = field(default=None, repr=False) + event_study_vcov_index: Optional[np.ndarray] = field(default=None, repr=False) + event_study_df: Optional[Dict[int, float]] = field(default=None, repr=False) + reference_periods: Tuple[int, ...] = field(default_factory=tuple, repr=False) + cband_method: Optional[str] = field(default=None, repr=False) + cband_crit_value: Optional[float] = field(default=None, repr=False) + cband_n_bootstrap: Optional[int] = field(default=None, repr=False) + + # ------------------------------------------------------------------ # + # Full regression output (optional) # + # ------------------------------------------------------------------ # + params: Optional[np.ndarray] = field(default=None, repr=False) + bse: Optional[np.ndarray] = field(default=None, repr=False) + vcov: Optional[np.ndarray] = field(default=None, repr=False) + + # ------------------------------------------------------------------ # + # Cached RI/WCB results (optional) # + # ------------------------------------------------------------------ # + _ri_result: Optional[Any] = field(default=None, repr=False) + _wcb_result: Optional[Any] = field(default=None, repr=False) + + # ------------------------------------------------------------------ # + # Properties # + # ------------------------------------------------------------------ # + @property + def pvalue(self) -> float: + """Alias for p_value (diff-diff API convention).""" + return self.p_value + + @property + def ci(self) -> Tuple[float, float]: + """Alias for conf_int (diff-diff API convention).""" + return self.conf_int + + @property + def is_staggered(self) -> bool: + """Whether this result comes from a staggered adoption design.""" + return self.cohort_effects is not None + + #: ``simple`` reports the estimand ``fit()`` already computed; it never + #: recombines cohort effects, which would silently swap the composite + #: regression's joint inference for a cohort-independence assumption. + _AGGREGATE_SUPPORTED = ("simple", "event_study", "group") + #: balance_e is REJECTED (round-5 review: it was accepted but ignored, + #: silently returning the unbalanced event-study surface): LWDiD stores + #: no per-cohort estimation kit from which a balanced-cohort sample and + #: its joint influence-function covariance could be recomputed post fit. + _AGGREGATE_BALANCE_E_TYPES = () + + def _aggregate_validate_weights(self, weights: Optional[str]) -> None: + if weights is not None: + raise ValueError( + "LWDiDResults.aggregate() does not accept a weights selector " + f"(got {weights!r}); LWDiD weights cohort-time cells by their " + "treated mass, which is fixed by the estimator." + ) + + def _aggregate_compute( + self, + level: str, + *, + weights: Optional[str], + balance_e: Optional[int], + ) -> Any: + if level == "group" and not self.is_staggered: + raise ValueError( + "aggregate('group') is only available for staggered fits; a " + "common-timing design has a single treatment cohort, so " + "there is no group dimension to aggregate over." + ) + + if level == "simple": + ci = self.conf_int + return AggregationResult( + level="simple", + label=np.array(["overall"], dtype=object), + target=np.array(["att"], dtype=object), + att=np.array([self.att], dtype=float), + se=np.array([self.se], dtype=float), + t_stat=np.array([self.t_stat], dtype=float), + p_value=np.array([self.p_value], dtype=float), + conf_int_lower=np.array([ci[0]], dtype=float), + conf_int_upper=np.array([ci[1]], dtype=float), + n=np.array([float(self.n_treated)], dtype=float), + df=np.array( + [np.nan if self.df_inference is None else float(self.df_inference)], + dtype=float, + ), + alpha=self.alpha, + n_kind="units", + weight=np.array([1.0], dtype=float), + estimator="LWDiD", + ) + + if level == "group": + cohorts = list(self.cohort_effects or {}) + effects = [self.cohort_effects[g] for g in cohorts] # type: ignore[index] + + def _column(key: str, default: float = np.nan) -> np.ndarray: + return np.array([_as_float(e.get(key, default)) for e in effects], dtype=float) + + bounds = [e.get("conf_int", (np.nan, np.nan)) for e in effects] + return AggregationResult( + level="group", + label=np.array(cohorts, dtype=object), + target=np.array(["att"] * len(cohorts), dtype=object), + att=_column("att"), + se=_column("se"), + t_stat=_column("t_stat"), + p_value=_column("p_value"), + conf_int_lower=np.array([_as_float(b[0]) for b in bounds], dtype=float), + conf_int_upper=np.array([_as_float(b[1]) for b in bounds], dtype=float), + n=_column("n_treated"), + df=_column("df"), + alpha=self.alpha, + n_kind="units", + weight=_column("weight"), + estimator="LWDiD", + ) + + if level == "event_study": + es_effects = self.event_study_effects or {} + reference_periods = set(self.reference_periods or ()) + labels = sorted(set(es_effects) | reference_periods) + rows = [es_effects.get(label, {}) for label in labels] + is_reference = np.array([label in reference_periods for label in labels], dtype=bool) + att = np.array( + [ + row.get("effect", 0.0 if reference else np.nan) + for row, reference in zip(rows, is_reference) + ], + dtype=float, + ) + se = np.array([row.get("se", np.nan) for row in rows], dtype=float) + t_stat = np.array([row.get("t_stat", np.nan) for row in rows], dtype=float) + p_value = np.array([row.get("p_value", np.nan) for row in rows], dtype=float) + ci_lower = np.array( + [row.get("conf_int", (np.nan, np.nan))[0] for row in rows], dtype=float + ) + ci_upper = np.array( + [row.get("conf_int", (np.nan, np.nan))[1] for row in rows], dtype=float + ) + n = np.array([row.get("n_treated", np.nan) for row in rows], dtype=float) + cband_lower = np.array( + [row.get("cband_conf_int", (np.nan, np.nan))[0] for row in rows], dtype=float + ) + cband_upper = np.array( + [row.get("cband_conf_int", (np.nan, np.nan))[1] for row in rows], dtype=float + ) + has_band = any(np.isfinite(cband_lower) & np.isfinite(cband_upper)) + vcov = self.event_study_vcov if self.event_study_vcov is not None else None + vcov_index = ( + self.event_study_vcov_index if self.event_study_vcov_index is not None else None + ) + has_vcov = vcov is not None and vcov_index is not None and len(vcov_index) > 0 + df = None + if self.event_study_df is not None: + df = np.array([self.event_study_df.get(label, np.nan) for label in labels]) + return EventStudyResults( + event_time=np.array(labels), + att=att, + se=se, + t_stat=t_stat, + p_value=p_value, + conf_int_lower=ci_lower, + conf_int_upper=ci_upper, + is_reference=is_reference, + n=n, + n_kind="units", + time_scale="relative", + event_time_convention="e0_first_treated", + vcov=vcov if has_vcov else None, + vcov_index=vcov_index if has_vcov else None, + cband_lower=cband_lower if has_band else None, + cband_upper=cband_upper if has_band else None, + cband_crit_value=self.cband_crit_value, + alpha=self.alpha, + source="LWDiDResults", + df=df, + # Scalar-df provenance, mirroring results_base's resolution + # rule: no survey notion, so the bare df_inference carrier. + df_survey=None if self.df_inference is None else float(self.df_inference), + ) + + raise ValueError(f"Unsupported aggregation method: {level!r}") + + # ------------------------------------------------------------------ # + # Serialization # + # ------------------------------------------------------------------ # + def to_dataframe(self) -> pd.DataFrame: + """Convert results to a pandas DataFrame. + + Returns + ------- + pd.DataFrame + For common timing: a single-row DataFrame. + For staggered: one row per cohort plus an "Overall" row. + """ + if not self.is_staggered: + rows: List[Dict[str, Any]] = [ + { + "term": "ATT", + "att": self.att, + "se": self.se, + "t_stat": self.t_stat, + "p_value": self.p_value, + "ci_lower": self.conf_int[0], + "ci_upper": self.conf_int[1], + "n_obs": self.n_obs, + "n_treated": self.n_treated, + "n_control": self.n_control, + "rolling": self.rolling, + "estimation_method": self.estimation_method, + "vcov_type": self.vcov_type, + } + ] + return pd.DataFrame(rows) + + rows_stag: List[Dict[str, Any]] = [] + for cohort, eff in self.cohort_effects.items(): # type: ignore[union-attr] + ci = eff.get("conf_int", (np.nan, np.nan)) + n_t = eff.get("n_treated", 0) + n_c = eff.get("n_control", 0) + rows_stag.append( + { + "cohort": cohort, + "att": eff.get("att", np.nan), + "se": eff.get("se", np.nan), + "t_stat": eff.get("t_stat", np.nan), + "p_value": eff.get("p_value", np.nan), + "ci_lower": ci[0] if ci else np.nan, + "ci_upper": ci[1] if ci else np.nan, + "n_treated": n_t, + "n_control": n_c, + "rolling": self.rolling, + "estimation_method": self.estimation_method, + "vcov_type": self.vcov_type, + } + ) + # Append overall row + rows_stag.append( + { + "cohort": "Overall", + "att": self.att, + "se": self.se, + "t_stat": self.t_stat, + "p_value": self.p_value, + "ci_lower": self.conf_int[0], + "ci_upper": self.conf_int[1], + "n_treated": self.n_treated, + "n_control": self.n_control, + "rolling": self.rolling, + "estimation_method": self.estimation_method, + "vcov_type": self.vcov_type, + } + ) + return pd.DataFrame(rows_stag) + + def to_dict(self) -> Dict[str, Any]: + """Convert results to a JSON-serializable dictionary. + + Returns + ------- + dict + All scalar results and metadata. Arrays are converted to lists + and numpy scalars (including nested dict values and keys) to + native Python types, so ``json.dumps(result.to_dict())`` works + directly. + """ + result: Dict[str, Any] = { + "att": self.att, + "se": self.se, + "t_stat": self.t_stat, + "p_value": self.p_value, + "conf_int_lower": self.conf_int[0], + "conf_int_upper": self.conf_int[1], + "n_obs": self.n_obs, + "n_treated": self.n_treated, + "n_control": self.n_control, + "rolling": self.rolling, + "estimation_method": self.estimation_method, + "vcov_type": self.vcov_type, + "alpha": self.alpha, + } + 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 + if self.cohort_effects is not None: + result["cohort_effects"] = {str(k): v for k, v in self.cohort_effects.items()} + if self.cohort_time_effects is not None: + result["cohort_time_effects"] = { + f"{g},{t}": value for (g, t), value in self.cohort_time_effects.items() + } + if self.inference_basis is not None: + result["inference_basis"] = self.inference_basis + if self.df_inference is not None: + result["df_inference"] = self.df_inference + if self.control_group is not None: + result["control_group"] = self.control_group + result["n_bootstrap"] = self.n_bootstrap + if self.seed is not None: + result["seed"] = self.seed + if self.pscore_trim is not None: + result["pscore_trim"] = self.pscore_trim + if self.psm_config is not None: + result["psm_config"] = dict(self.psm_config) + if self.att_tau_omega_complete_case is not None: + result["att_tau_omega_complete_case"] = self.att_tau_omega_complete_case + if self.n_composite_treated_dropped or self.n_composite_controls_dropped: + result["n_composite_treated_dropped"] = self.n_composite_treated_dropped + result["n_composite_controls_dropped"] = self.n_composite_controls_dropped + if self.params is not None: + result["params"] = self.params.tolist() + if self.bse is not None: + result["bse"] = self.bse.tolist() + if self.event_study_effects is not None: + result["event_study_effects"] = {str(k): v for k, v in self.event_study_effects.items()} + result["reference_periods"] = list(self.reference_periods) + result["cband_method"] = self.cband_method + result["cband_crit_value"] = self.cband_crit_value + result["cband_n_bootstrap"] = self.cband_n_bootstrap + return _to_json_native(result) + + # ------------------------------------------------------------------ # + # Aggregation # + # ------------------------------------------------------------------ # + def to_csv(self, path: str) -> None: + """Export results to CSV file. + + Parameters + ---------- + path : str + File path for the CSV output. + """ + self.to_dataframe().to_csv(path, index=False) + + # ------------------------------------------------------------------ # + # Text summary # + # ------------------------------------------------------------------ # + def summary(self) -> str: + """Formatted text summary of results. + + Returns + ------- + str + Human-readable summary table. + """ + from diff_diff.results import _format_vcov_label, _get_significance_stars + + 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, + "Lee & Wooldridge DiD (LWDiD) Results".center(width), + bar, + f"Observations: {self.n_obs} " + f"Treated units: {self.n_treated} " + f"Control units: {self.n_control}", + f"Rolling: {self.rolling} " + f"Method: {self.estimation_method} " + f"Alpha: {self.alpha}", + ] + + # Variance label + 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 for results table + header = ( + f"{'':>12} {'Estimate':>10} {'Std.Err':>10} {'t':>8} " + f"{'P>|t|':>8} [{ci_pct}% Conf. Int.]" + ) + + # Main ATT row + lines.append("") + if self.is_staggered: + lines.append("Cohort-level effects:") + lines.append(dash) + lines.append(header) + lines.append(dash) + for cohort, eff in self.cohort_effects.items(): # type: ignore[union-attr] + ci = eff.get("conf_int", (np.nan, np.nan)) + p = eff.get("p_value", np.nan) + stars = "" if np.isnan(p) else _get_significance_stars(float(p)) + label = f"G={cohort}" + lines.append( + f"{label:>12} {_fmt(eff.get('att')):>10} " + f"{_fmt(eff.get('se')):>10} " + f"{_fmt(eff.get('t_stat'), 2):>8} " + f"{_fmt(p, 3):>8} " + f"[{_fmt(ci[0]):>9}, {_fmt(ci[1]):>9}] {stars}" + ) + lines.append(dash) + # Overall ATT + stars = _get_significance_stars(self.p_value) if not np.isnan(self.p_value) else "" + lines.append( + f"{'Overall ATT':>12} {_fmt(self.att):>10} " + f"{_fmt(self.se):>10} " + f"{_fmt(self.t_stat, 2):>8} " + f"{_fmt(self.p_value, 3):>8} " + f"[{_fmt(self.conf_int[0]):>9}, {_fmt(self.conf_int[1]):>9}] {stars}" + ) + else: + lines.append("ATT estimate:") + lines.append(dash) + lines.append(header) + lines.append(dash) + stars = _get_significance_stars(self.p_value) if not np.isnan(self.p_value) else "" + lines.append( + f"{'ATT':>12} {_fmt(self.att):>10} " + f"{_fmt(self.se):>10} " + f"{_fmt(self.t_stat, 2):>8} " + f"{_fmt(self.p_value, 3):>8} " + f"[{_fmt(self.conf_int[0]):>9}, {_fmt(self.conf_int[1]):>9}] {stars}" + ) + + lines.append(bar) + if self.inference_basis is not None: + label = _INFERENCE_BASIS_LABELS.get(self.inference_basis, self.inference_basis) + lines.append(f"Overall inference: {label}") + lines.append("Signif. codes: *** p<0.001, ** p<0.01, * p<0.05") + return "\n".join(lines) + + def print_summary(self) -> None: + """Print the formatted summary to stdout.""" + print(self.summary()) + + # ================================================================ + # Advanced inference and diagnostics (delegate to standalone modules) + # ================================================================ + + @property + def ri_pvalue(self): + """Randomization inference p-value (None if not computed).""" + if self._ri_result is not None: + return self._ri_result.pvalue + return None + + @property + def bootstrap_pvalue(self): + """Wild cluster bootstrap p-value (None if not computed).""" + if self._wcb_result is not None: + return self._wcb_result.p_value + return None + + def _replay_arrays(self, method_name): + """Fitted-sample arrays for the post-fit advanced-inference methods. + + Round-5 review: these methods previously accepted arbitrary caller + arrays and fit a non-interacted design, so the cached p-values + could describe a DIFFERENT estimand than ``.att`` (measured on a + covariate-unbalanced RA fit: fitted 3.98 vs tested 3.26). They now + REPLAY the fit-time collapsed cross-section and the exact RA + design; no data arguments are accepted. + """ + spec = getattr(self, "_replay_spec", None) + if spec is None: + raise ValueError( + f"{method_name} replays the fitted common-timing estimation " + "sample, which this results object does not carry (staggered " + "and degenerate fits are not supported). Use the standalone " + "module function with explicit arrays instead." + ) + if self.estimation_method != "reg": + raise ValueError( + f"{method_name} replays the fitted RA regression and is only " + f"defined for estimation_method='reg' (got " + f"'{self.estimation_method}'): re-estimating the " + f"{self.estimation_method} estimator per draw is not " + "implemented. Use the standalone module function on arrays " + "of your choosing (a generic [1, D, X] contrast, NOT the " + "fitted estimand)." + ) + return spec + + @staticmethod + def _fit_used_interactions(treatment, controls): + """Mirror _estimate_reg's LW eq. 3.3 gate: interactions require + N_1 > K+1 and N_0 > K+1; otherwise the fit used plain (1, D, X) + (round-7 review: the replay always interacted, so small-arm fits' + replayed statistic mismatched .att and the coherence assert made + their post-fit inference unusable).""" + if controls is None: + return False + n_treated = int((treatment == 1).sum()) + n_control = len(treatment) - n_treated + # IDENTIFIED control dimension, mirroring _estimate_reg exactly + # (round-11 review: the nominal column count diverged from the + # fit's gate under collinear controls). + from diff_diff.linalg import _detect_rank_deficiency + + k = int( + _detect_rank_deficiency(np.column_stack([np.ones(len(treatment)), controls]))[0] - 1 + ) + return n_treated > k + 1 and n_control > k + 1 + + @classmethod + def _replay_controls(cls, treatment, controls): + """The exact auxiliary columns the fitted RA regression used: + [X, D*(X - Xbar_1)] when the interaction gate held, plain X + otherwise (LW eq. E.1 / eq. 3.3).""" + if controls is None: + return None + if not cls._fit_used_interactions(treatment, controls): + return controls + xbar1 = controls[treatment == 1].mean(axis=0) + return np.column_stack([controls, treatment[:, None] * (controls - xbar1)]) + + def _assert_replay_coherent(self, observed, method_name): + if not np.isclose(observed, self.att, rtol=1e-8, atol=1e-10): + raise RuntimeError( + f"{method_name}: the replayed observed ATT ({observed!r}) " + f"does not match the fitted .att ({self.att!r}); refusing to " + "cache inference for a different estimand (fail closed)." + ) + + def wild_cluster_bootstrap( + self, + *, + n_bootstrap=999, + weight_type="rademacher", + alpha=None, + seed=None, + ): + """Run wild cluster bootstrap inference on the fitted estimation sample. + + Replays the fit-time collapsed cross-section and the exact fitted + RA design (intercept, treatment, covariates, and the treatment- + centered interactions) through the house WCR engine + (test-inversion CI, CR1 se, strict-exceedance p-value); the + observed coefficient is asserted equal to ``.att`` before caching. + Requires a clustered (``cluster=``), common-timing, + ``estimation_method='reg'`` fit. ``alpha=None`` inherits the + fitted confidence level. Result is cached and accessible via the + ``bootstrap_pvalue`` property. + """ + from diff_diff.lwdid_wild_bootstrap import wild_cluster_bootstrap as _wcb + + spec = self._replay_arrays("wild_cluster_bootstrap") + if spec["cluster_ids"] is None: + raise ValueError( + "wild_cluster_bootstrap requires a clustered fit: construct " + "the estimator with cluster= and refit." + ) + result = _wcb( + spec["y"], + spec["treatment"], + spec["cluster_ids"], + self._replay_controls(spec["treatment"], spec["controls"]), + n_bootstrap=n_bootstrap, + weight_type=weight_type, + alpha=self.alpha if alpha is None else alpha, + seed=seed, + ) + self._assert_replay_coherent(result.att, "wild_cluster_bootstrap") + object.__setattr__(self, "_wcb_result", result) + return result + + def randomization_test(self, *, n_reps=1000, method="permutation", seed=None): + """Run Fisher randomization inference on the fitted estimation sample. + + Replays the fit-time collapsed cross-section; with covariates the + RA design's treated covariate mean and interaction columns are + RECOMPUTED for every permuted assignment (``design='ra_interacted'``), + so each permutation tests the same estimator the fit reported. The + observed statistic is asserted equal to ``.att`` before caching. + Requires a common-timing ``estimation_method='reg'`` fit. Result is + cached and accessible via the ``ri_pvalue`` property. + """ + from diff_diff.lwdid_randomization import randomization_inference as _ri + + spec = self._replay_arrays("randomization_test") + # Match the fitted design exactly: 'ra_interacted' only when the + # fit's interaction gate held (round-7 review). NOTE the permuted + # draws under the plain design keep plain (1, D, X) too - each + # permutation tests the same estimator the fit reported. + design = ( + "ra_interacted" + if self._fit_used_interactions(spec["treatment"], spec["controls"]) + else "linear" + ) + result = _ri( + spec["y"], + spec["treatment"], + spec["controls"], + n_reps=n_reps, + method=method, + seed=seed, + design=design, + ) + self._assert_replay_coherent(result.att_observed, "randomization_test") + object.__setattr__(self, "_ri_result", result) + return result + + # ------------------------------------------------------------------ # + # Repr # + # ------------------------------------------------------------------ # + def __repr__(self) -> str: + cluster = f", cluster={self.cluster_name}, G={self.n_clusters}" if self.cluster_name else "" + att_s = "nan" if np.isnan(self.att) else f"{self.att:.4f}" + se_s = "nan" if np.isnan(self.se) else f"{self.se:.4f}" + stag = ", staggered=True" if self.is_staggered else "" + return ( + f"LWDiDResults(" + f"ATT={att_s}, SE={se_s}, " + f"rolling={self.rolling!r}, estimation_method={self.estimation_method!r}, " + f"vcov_type={self.vcov_type!r}{cluster}{stag})" + ) diff --git a/diff_diff/lwdid_sensitivity.py b/diff_diff/lwdid_sensitivity.py new file mode 100644 index 00000000..2f64da11 --- /dev/null +++ b/diff_diff/lwdid_sensitivity.py @@ -0,0 +1,907 @@ +"""Sensitivity analysis for LWDiD estimator. + +Assesses robustness of ATT estimates along the two axes with direct +theoretical grounding in Lee & Wooldridge (2025, 2026): +- Pre-period selection sensitivity (T0-robustness) +- No-anticipation assumption sensitivity + +Classification thresholds (a diff-diff LIBRARY HEURISTIC — the Lee & +Wooldridge papers recommend the diagnostics but define no categorical +robustness scale; see docs/methodology/REGISTRY.md, LWDiD): + sensitivity_ratio < 10% → 'highly_robust' + 10% ≤ ratio < 25% → 'moderately_robust' + 25% ≤ ratio < 50% → 'sensitive' + ratio ≥ 50% → 'highly_sensitive' + ratio is NaN → 'not_estimable' (baseline ATT non-finite or + fewer than two specifications produced finite + estimates; robustness cannot be assessed) + +References +---------- +Lee, S. J. & Wooldridge, J. M. (2025). "A Simple Transformation Approach + to Difference-in-Differences Estimation for Panel Data." SSRN 4516518. +Lee, S. J. & Wooldridge, J. M. (2026). "Simple Approaches to Inference + with Difference-in-Differences Estimators with Small Cross-Sectional + Sample Sizes." SSRN 5325686. +""" + +from __future__ import annotations + +import warnings +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import numpy as np +import pandas as pd + +# ============================================================================= +# Constants +# ============================================================================= + +_ROBUSTNESS_THRESHOLDS = { + "highly_robust": 0.10, + "moderately_robust": 0.25, + "sensitive": 0.50, +} + + +# ============================================================================= +# Data Classes +# ============================================================================= + + +@dataclass +class SpecificationResult: + """Result from a single specification in sensitivity analysis. + + Attributes + ---------- + label : str + Human-readable label describing this specification. + rolling : str + Transformation method used ('demean' or 'detrend'). + estimation_method : str + Estimation method used ('reg', 'ipw', 'dr'). + n_pre_periods : int + Number of pre-treatment periods used. -1 if all periods used. + att : float + Average treatment effect on the treated. + se : float + Standard error of ATT. + pvalue : float + Two-sided p-value for testing H0: ATT = 0. + """ + + label: str + rolling: str + estimation_method: str + n_pre_periods: int + att: float + se: float + pvalue: float + #: Fitted confidence interval endpoints (round-21 review: plots must + #: render the fitted interval, not a normal-theory +/-1.96*SE). + conf_int: Optional[Tuple[float, float]] = None + + @property + def is_significant(self) -> float: + """1.0 / 0.0 for a decidable 5%-level test, NaN when the p-value + is missing or non-finite (a failed specification must never + publish "not significant" - fix-wave WS10).""" + if self.pvalue is None or not np.isfinite(self.pvalue): + return float("nan") + return float(self.pvalue < 0.05) + + def to_dict(self) -> dict: + """Convert to dictionary for DataFrame construction.""" + return { + "label": self.label, + "rolling": self.rolling, + "estimation_method": self.estimation_method, + "n_pre_periods": self.n_pre_periods, + "att": self.att, + "se": self.se, + "pvalue": self.pvalue, + "conf_int": self.conf_int, + "significant_05": self.is_significant, + } + + +@dataclass +class SensitivityResult: + """Result of comprehensive sensitivity analysis. + + Attributes + ---------- + specifications : List[SpecificationResult] + Results from each non-baseline specification. + baseline_att : float + ATT from the baseline specification. + baseline_se : float + Standard error from the baseline specification. + sensitivity_ratio : float + (max_att - min_att) / |baseline_att|, measuring estimate instability. + NaN when robustness cannot be assessed (non-finite baseline ATT or + fewer than two finite estimates). + robustness_level : str + Categorical assessment: 'highly_robust', 'moderately_robust', + 'sensitive', 'highly_sensitive', or 'not_estimable'. The + 'not_estimable' level indicates the sensitivity ratio is NaN + because too few specifications produced finite estimates. + n_specifications : int + Total number of specifications tested (including baseline). + """ + + specifications: List[SpecificationResult] + baseline_att: float + baseline_se: float + sensitivity_ratio: float + robustness_level: str + n_specifications: int + #: Baseline specification's p-value (None only on legacy construction; + #: NaN when the baseline fit failed). + baseline_pvalue: Optional[float] = None + + def summary(self) -> str: + """Return a formatted summary of sensitivity analysis results. + + Returns + ------- + str + Multi-line string summarizing the sensitivity analysis. + """ + lines = [ + "=" * 60, + "LWDiD Sensitivity Analysis Summary", + "=" * 60, + f"Baseline ATT: {self.baseline_att:.6f}", + f"Baseline SE: {self.baseline_se:.6f}", + f"Sensitivity Ratio: {self.sensitivity_ratio:.4f} " + f"({self.sensitivity_ratio * 100:.1f}%)", + f"Robustness Level: {self.robustness_level}", + f"N Specifications: {self.n_specifications}", + "-" * 60, + ] + + if self.specifications: + lines.append(f"{'Label':<25} {'ATT':>10} {'SE':>10} {'p-value':>10}") + lines.append("-" * 60) + for spec in self.specifications: + lines.append( + f"{spec.label:<25} {spec.att:>10.6f} " f"{spec.se:>10.6f} {spec.pvalue:>10.4f}" + ) + else: + lines.append("No alternative specifications computed.") + + lines.append("=" * 60) + return "\n".join(lines) + + def to_dataframe(self) -> pd.DataFrame: + """Convert all specification results to a DataFrame. + + Returns + ------- + pd.DataFrame + DataFrame with columns: label, rolling, estimation_method, + n_pre_periods, att, se, pvalue, significant_05. + """ + baseline_p = self.baseline_pvalue + if baseline_p is None or not np.isfinite(baseline_p): + baseline_sig = float("nan") + baseline_p = float("nan") if baseline_p is None else baseline_p + else: + baseline_sig = float(baseline_p < 0.05) + rows = [ + { + "label": "baseline", + "rolling": "", + "estimation_method": "", + "n_pre_periods": -1, + "att": self.baseline_att, + "se": self.baseline_se, + "pvalue": baseline_p, + "significant_05": baseline_sig, + } + ] + for spec in self.specifications: + rows.append(spec.to_dict()) + return pd.DataFrame(rows) + + def __repr__(self) -> str: + return ( + f"SensitivityResult(baseline_att={self.baseline_att:.4f}, " + f"ratio={self.sensitivity_ratio:.4f}, " + f"level='{self.robustness_level}', " + f"n_specs={self.n_specifications})" + ) + + +# ============================================================================= +# Helper Functions +# ============================================================================= + + +def _classify_robustness(ratio: float) -> str: + """Classify sensitivity ratio into robustness level. + + Parameters + ---------- + ratio : float + Sensitivity ratio (range / |baseline|). NaN indicates the ratio + could not be estimated. + + Returns + ------- + str + One of 'highly_robust', 'moderately_robust', 'sensitive', + 'highly_sensitive', or 'not_estimable' (when ratio is NaN). + """ + if np.isnan(ratio): + return "not_estimable" + if ratio < _ROBUSTNESS_THRESHOLDS["highly_robust"]: + return "highly_robust" + elif ratio < _ROBUSTNESS_THRESHOLDS["moderately_robust"]: + return "moderately_robust" + elif ratio < _ROBUSTNESS_THRESHOLDS["sensitive"]: + return "sensitive" + else: + return "highly_sensitive" + + +def _compute_sensitivity_ratio(baseline_att: float, all_atts: List[float]) -> float: + """Compute sensitivity ratio from ATT estimates. + + Parameters + ---------- + baseline_att : float + Baseline ATT estimate. + all_atts : list of float + All ATT estimates including baseline. + + Returns + ------- + float + Sensitivity ratio: (max - min) / |baseline|. NaN when the baseline + ATT is non-finite, fewer than two estimates are finite, or the + baseline is (numerically) zero -- the relative ratio is undefined + there, so robustness cannot be assessed (classified as + 'not_estimable', never 'highly_robust'). + """ + if not np.isfinite(baseline_att): + return float(np.nan) + finite_atts = [a for a in all_atts if np.isfinite(a)] + if len(finite_atts) <= 1: + return float(np.nan) + if abs(baseline_att) < 1e-10: + return float(np.nan) + return (max(finite_atts) - min(finite_atts)) / abs(baseline_att) + + +def _fit_single_spec( + data: pd.DataFrame, + outcome: str, + unit: str, + time: str, + treatment: str, + cohort: Optional[str], + rolling: str, + estimation_method: str, + vcov_type: str, + cluster: Optional[str], + controls: Optional[List[str]], + control_group: str = "not_yet_treated", + raise_errors: bool = False, +) -> Tuple[float, float, float, Tuple[float, float]]: + """Fit a single LWDiD specification and return + (att, se, pvalue, conf_int). + + Column existence is validated eagerly: missing columns raise + ValueError instead of being silently converted to NaN. Only + data-dependent failures of the fit itself (ValueError from a + degenerate specification, e.g. no remaining pre-periods, or a + LinAlgError from a singular design) are mapped to (nan, nan, nan); + any other exception is a programming error and propagates. + + ``raise_errors=True`` (the BASELINE fit in both public helpers, on + the full frame) propagates every fit error: a full-frame failure is + a configuration/support problem (e.g. covariate-free PSM), not a + restricted-specification non-estimability, and must not be reported + as ``not_estimable`` (round-10 review). + """ + from diff_diff.lwdid import LWDiD + + required = { + "outcome": outcome, + "unit": unit, + "time": time, + "treatment": treatment, + } + if cohort is not None: + required["cohort"] = cohort + if cluster is not None: + required["cluster"] = cluster + missing = [f"{role}={name!r}" for role, name in required.items() if name not in data.columns] + if controls is not None: + missing.extend(f"control={c!r}" for c in controls if c not in data.columns) + if missing: + raise ValueError( + f"Column(s) not found in data for sensitivity analysis: {', '.join(missing)}" + ) + + est = LWDiD( + rolling=rolling, + estimation_method=estimation_method, + vcov_type=vcov_type, + cluster=cluster, + control_group=control_group, + ) + try: + res = est.fit( + data, + outcome=outcome, + unit=unit, + time=time, + treatment=treatment, + first_treat=cohort, + covariates=controls, + ) + return res.att, res.se, res.p_value, res.conf_int + except (ValueError, np.linalg.LinAlgError): + if raise_errors: + raise + return np.nan, np.nan, np.nan, (np.nan, np.nan) + + +def _prevalidate_frame(data, outcome, unit, time, treatment, cohort, cluster, controls) -> None: + """Run LWDiD's shared input validation on the full frame (raises). + + Includes the treatment-design check (absorbing treatment, common- + timing onset homogeneity, D_it/cohort consistency) with the same + encode-then-normalize ordering as ``fit()`` - round-7 review: without + it, a structurally invalid design (e.g. a 1 -> 0 treatment reversal) + was swallowed by the per-spec ValueError handler and reported as + ``robustness_level='not_estimable'`` instead of raising. + """ + from diff_diff.lwdid import ( + LWDiD, + _check_treatment_design, + _encode_staggered_time_scale, + _normalize_cohorts, + ) + from diff_diff.utils import validate_binary + + probe = LWDiD(cluster=cluster) + frame = data.copy() + probe._validate_inputs( + frame, outcome, unit, time, treatment, cohort, cluster, list(controls or []) + ) + validate_binary(frame[treatment].values, treatment) + time_col, cohort_col = time, cohort + if cohort is not None: + frame, time_col, cohort_col, _ = _encode_staggered_time_scale(frame, time, cohort) + with warnings.catch_warnings(): + # fit() re-normalizes and re-warns; suppress the duplicate here. + warnings.simplefilter("ignore") + frame[cohort_col], _, _ = _normalize_cohorts( + frame[cohort_col], max_time=frame[time_col].max() + ) + _check_treatment_design(frame, unit, time_col, treatment, cohort_col) + if cohort is not None: + # Multi-cohort rejection on the NORMALIZED cohorts (round-10 + # review: counting raw values rejected valid single-cohort + # designs whose beyond-window/inf encodings normalize to + # never-treated). + values = frame[cohort_col].to_numpy(dtype=float) + treated_cohorts = np.unique(values[np.isfinite(values) & (values > 0)]) + if len(treated_cohorts) > 1: + raise ValueError( + f"Sensitivity analyses currently support a single treated " + f"cohort; found {len(treated_cohorts)} distinct cohorts in " + f"'{cohort}' (after never-treated normalization). " + f"Pre-period exclusions are defined relative to the " + f"earliest adoption, which would mislabel the samples used " + f"for later cohorts' transformations. Run the analysis per " + f"cohort, or see DEFERRED.md (cohort-relative sensitivity " + f"exclusions)." + ) + + +def _get_pre_periods(data: pd.DataFrame, time: str, treatment: str) -> np.ndarray: + """Identify pre-treatment periods from the data. + + Parameters + ---------- + data : pd.DataFrame + Panel dataset. + time : str + Time column name. + treatment : str + Treatment indicator column name. + + Returns + ------- + np.ndarray + Sorted array of pre-treatment period values. + """ + # Partition at the single onset S = min(observed treated period), + # matching fit()'s calendar rule (round-9 review: the former + # any-unit-treated rule classified a controls-only post period as + # pre-treatment, so exclusions could remove a POST period while + # labeling it an excluded pre-period). + all_periods = np.sort(data[time].unique()) + treated_times = data.loc[data[treatment] == 1, time] + if len(treated_times) == 0: + return all_periods + onset_s = treated_times.min() + return np.array([p for p in all_periods if p < onset_s]) + + +# ============================================================================= +# Public API: robustness_pre_periods +# ============================================================================= + + +def robustness_pre_periods( + data: pd.DataFrame, + outcome: str = None, + unit: str = None, + time: str = None, + treatment: str = None, + cohort: Optional[str] = None, + rolling: str = "demean", + estimation_method: str = "reg", + vcov_type: str = "hc1", + cluster: Optional[str] = None, + controls: Optional[List[str]] = None, + control_group: str = "not_yet_treated", + k_min: int = 2, + k_max: Optional[int] = None, + # lwdid-py compatible aliases + y: Optional[str] = None, + ivar: Optional[str] = None, + tvar: Optional[str] = None, + d: Optional[str] = None, + gvar: Optional[str] = None, + **kwargs, +) -> SensitivityResult: + """Assess sensitivity of ATT to number of pre-treatment periods used. + + For each k in range(k_min, k_max+1), restricts the data to use only + the last k pre-treatment periods for rolling transformation, then fits + LWDiD and collects the ATT estimate. + + Parameters + ---------- + data : pd.DataFrame + Panel dataset in long format. + outcome : str + Outcome column name. (alias: y) + unit : str + Unit identifier column name. (alias: ivar) + time : str + Time period column name. (alias: tvar) + treatment : str + Binary treatment indicator column name. (alias: d) + cohort : str, optional + Cohort variable for staggered designs. (alias: gvar) At most ONE + distinct treated cohort is supported: the exclusion windows are + defined relative to the earliest adoption, which would mislabel + later cohorts' transformation samples (multi-cohort inputs raise + ValueError; see DEFERRED.md, cohort-relative exclusions). + rolling : str, default 'demean' + Transformation method. + estimation_method : str, default 'reg' + Estimation method. + vcov_type : str, default 'hc1' + Variance-covariance family. + cluster : str, optional + Cluster variable for standard errors. + controls : list of str, optional + Control variable column names. + k_min : int, default 2 + Minimum number of pre-treatment periods to test. + k_max : int, optional + Maximum number of pre-treatment periods. If None, uses all available. + + Returns + ------- + SensitivityResult + Sensitivity analysis result with per-specification ATT estimates + and overall robustness classification. + """ + if kwargs: + raise TypeError( + f"robustness_pre_periods() got unexpected keyword argument(s): " f"{sorted(kwargs)}" + ) + # Resolve lwdid-py aliases + outcome = outcome or y + unit = unit or ivar + time = time or tvar + treatment = treatment or d + cohort = cohort or gvar + + # Validate required params + if outcome is None: + raise ValueError("'outcome' (or 'y') parameter is required") + if unit is None: + raise ValueError("'unit' (or 'ivar') parameter is required") + if time is None: + raise ValueError("'time' (or 'tvar') parameter is required") + if treatment is None: + raise ValueError("'treatment' (or 'd') parameter is required") + + # Pre-validate the FULL frame once so genuine specification errors + # (missing/NaN key columns, non-binary treatment, malformed panels) + # RAISE here instead of being swallowed as per-spec "failed fits" + # inside _fit_single_spec (campaign finding: a string covariate's + # ValueError became a silent NaN spec). + _prevalidate_frame(data, outcome, unit, time, treatment, cohort, cluster, controls) + + for name, value in (("k_min", k_min), ("k_max", k_max)): + if value is None: + continue + if isinstance(value, bool) or not isinstance(value, (int, np.integer)) or value < 1: + raise ValueError(f"{name} must be a positive integer; got {value!r}.") + + pre_periods = _get_pre_periods(data, time, treatment) + n_pre = len(pre_periods) + + if k_max is None: + k_max = n_pre + + k_max = min(k_max, n_pre) + # Transformation-aware minimum (round-16 review: the former + # unconditional max(k_min, 2) SILENTLY dropped an explicitly + # requested, methodologically valid k=1 demeaning specification - + # demeaning needs one pre-period, detrending two). + min_required = 1 if rolling in ("demean", "demeanq") else 2 + if k_min < min_required: + raise ValueError( + f"k_min={k_min} is below the minimum pre-period requirement " + f"for rolling='{rolling}' ({min_required}; detrending needs " + f"two pre-periods for its rank condition)." + ) + + if k_min > k_max: + warnings.warn( + f"k_min ({k_min}) > k_max ({k_max}). " + "Insufficient pre-treatment periods for robustness analysis.", + UserWarning, + stacklevel=2, + ) + # Return degenerate result with baseline only + att, se, pval, spec_ci = _fit_single_spec( + data, + outcome, + unit, + time, + treatment, + cohort, + rolling, + estimation_method, + vcov_type, + cluster, + controls, + control_group=control_group, + raise_errors=True, + ) + degenerate_ratio = _compute_sensitivity_ratio(att, [att]) + return SensitivityResult( + specifications=[], + baseline_att=att, + baseline_se=se, + sensitivity_ratio=degenerate_ratio, + robustness_level=_classify_robustness(degenerate_ratio), + n_specifications=1, + baseline_pvalue=pval, + ) + + # Baseline: use all pre-periods + baseline_att, baseline_se, baseline_pval, _baseline_ci = _fit_single_spec( + data, + outcome, + unit, + time, + treatment, + cohort, + rolling, + estimation_method, + vcov_type, + cluster, + controls, + control_group=control_group, + raise_errors=True, + ) + + # ALL observed periods >= S are post (round-9 review: the any-unit- + # treated rule dropped controls-only post periods from the subset). + _onset_s = data.loc[data[treatment] == 1, time].min() + post_periods = np.sort(data.loc[data[time] >= _onset_s, time].unique()) + + specs: List[SpecificationResult] = [] + + for k in range(k_min, k_max + 1): + if k == n_pre: + # Same as baseline, skip + continue + + # Keep only the last k pre-periods + all post-periods + keep_pre = pre_periods[-k:] + keep_periods = np.concatenate([keep_pre, post_periods]) + subset = data[data[time].isin(keep_periods)].copy() + + att, se, pval, spec_ci = _fit_single_spec( + subset, + outcome, + unit, + time, + treatment, + cohort, + rolling, + estimation_method, + vcov_type, + cluster, + controls, + control_group=control_group, + ) + + specs.append( + SpecificationResult( + label=f"k={k}_pre_periods", + rolling=rolling, + estimation_method=estimation_method, + n_pre_periods=k, + att=att, + se=se, + pvalue=pval, + conf_int=spec_ci, + ) + ) + + # Compute sensitivity ratio + all_atts = [baseline_att] + [s.att for s in specs] + ratio = _compute_sensitivity_ratio(baseline_att, all_atts) + level = _classify_robustness(ratio) + + if level == "not_estimable": + warnings.warn( + "Sensitivity ratio could not be estimated: baseline ATT is " + "non-finite or fewer than two specifications produced finite " + "estimates. Robustness to pre-period selection cannot be " + "assessed.", + UserWarning, + stacklevel=2, + ) + elif level in ("sensitive", "highly_sensitive"): + warnings.warn( + f"ATT estimates are {level} to pre-period selection " + f"(ratio={ratio:.3f}). Consider investigating data structure.", + UserWarning, + stacklevel=2, + ) + + return SensitivityResult( + specifications=specs, + baseline_att=baseline_att, + baseline_se=baseline_se, + sensitivity_ratio=ratio, + robustness_level=level, + n_specifications=len(specs) + 1, + baseline_pvalue=baseline_pval, + ) + + +# ============================================================================= +# Public API: sensitivity_no_anticipation +# ============================================================================= + + +def sensitivity_no_anticipation( + data: pd.DataFrame, + outcome: str = None, + unit: str = None, + time: str = None, + treatment: str = None, + cohort: Optional[str] = None, + exclude_periods: Optional[List[int]] = None, + rolling: str = "demean", + estimation_method: str = "reg", + vcov_type: str = "hc1", + cluster: Optional[str] = None, + controls: Optional[List[str]] = None, + control_group: str = "not_yet_treated", + # lwdid-py compatible aliases + y: Optional[str] = None, + ivar: Optional[str] = None, + tvar: Optional[str] = None, + d: Optional[str] = None, + gvar: Optional[str] = None, + **kwargs, +) -> SensitivityResult: + """Assess sensitivity to potential anticipation effects. + + For each n_exclude in exclude_periods, drops the last n_exclude + pre-treatment periods and re-estimates LWDiD. If ATT changes + substantially when excluding periods just before treatment, + this suggests anticipation effects may be present. + + Parameters + ---------- + data : pd.DataFrame + Panel dataset in long format. + outcome : str + Outcome column name. (alias: y) + unit : str + Unit identifier column name. (alias: ivar) + time : str + Time period column name. (alias: tvar) + treatment : str + Binary treatment indicator column name. (alias: d) + cohort : str, optional + Cohort variable for staggered designs. (alias: gvar) At most ONE + distinct treated cohort is supported: the exclusion windows are + defined relative to the earliest adoption, which would mislabel + later cohorts' transformation samples (multi-cohort inputs raise + ValueError; see DEFERRED.md, cohort-relative exclusions). + exclude_periods : list of int, optional + Number of pre-treatment periods to exclude in each test. + Default is [1, 2, 3]. + rolling : str, default 'demean' + Transformation method. + estimation_method : str, default 'reg' + Estimation method. + vcov_type : str, default 'hc1' + Variance-covariance family. + cluster : str, optional + Cluster variable for standard errors. + controls : list of str, optional + Control variable column names. + + Returns + ------- + SensitivityResult + Sensitivity result with per-exclusion ATT estimates and + overall robustness classification. + """ + if kwargs: + raise TypeError( + f"sensitivity_no_anticipation() got unexpected keyword argument(s): " + f"{sorted(kwargs)}" + ) + # Resolve lwdid-py aliases + outcome = outcome or y + unit = unit or ivar + time = time or tvar + treatment = treatment or d + cohort = cohort or gvar + + # Validate required params + if outcome is None: + raise ValueError("'outcome' (or 'y') parameter is required") + if unit is None: + raise ValueError("'unit' (or 'ivar') parameter is required") + if time is None: + raise ValueError("'time' (or 'tvar') parameter is required") + if treatment is None: + raise ValueError("'treatment' (or 'd') parameter is required") + + _prevalidate_frame(data, outcome, unit, time, treatment, cohort, cluster, controls) + + if exclude_periods is None: + exclude_periods = [1, 2, 3] + validated_exclusions: List[int] = [] + for value in exclude_periods: + # Round-3 review: exclude 0 sliced pre_periods[:-0] == EMPTY + # (dropping every pre-period instead of none); negative values + # selected the wrong window; bool is an int subclass. + if isinstance(value, bool) or not isinstance(value, (int, np.integer)) or value < 1: + raise ValueError( + f"exclude_periods entries must be positive integers " + f"(number of trailing pre-periods to drop); got {value!r}." + ) + validated_exclusions.append(int(value)) + if len(set(validated_exclusions)) != len(validated_exclusions): + raise ValueError(f"exclude_periods contains duplicate entries: {exclude_periods!r}.") + exclude_periods = validated_exclusions + + pre_periods = _get_pre_periods(data, time, treatment) + n_pre = len(pre_periods) + + # Baseline: no exclusion + baseline_att, baseline_se, baseline_pval, _baseline_ci = _fit_single_spec( + data, + outcome, + unit, + time, + treatment, + cohort, + rolling, + estimation_method, + vcov_type, + cluster, + controls, + control_group=control_group, + raise_errors=True, + ) + + # ALL observed periods >= S are post (round-9 review: the any-unit- + # treated rule dropped controls-only post periods from the subset). + _onset_s = data.loc[data[treatment] == 1, time].min() + post_periods = np.sort(data.loc[data[time] >= _onset_s, time].unique()) + + specs: List[SpecificationResult] = [] + + for n_exclude in exclude_periods: + if n_exclude >= n_pre: + warnings.warn( + f"Cannot exclude {n_exclude} periods with only {n_pre} " + "pre-treatment periods. Skipping.", + UserWarning, + stacklevel=2, + ) + continue + + # Exclude the last n_exclude pre-periods + remaining_pre = pre_periods[:-n_exclude] + keep_periods = np.concatenate([remaining_pre, post_periods]) + subset = data[data[time].isin(keep_periods)].copy() + + att, se, pval, spec_ci = _fit_single_spec( + subset, + outcome, + unit, + time, + treatment, + cohort, + rolling, + estimation_method, + vcov_type, + cluster, + controls, + control_group=control_group, + ) + + specs.append( + SpecificationResult( + label=f"exclude_{n_exclude}_periods", + rolling=rolling, + estimation_method=estimation_method, + n_pre_periods=n_pre - n_exclude, + att=att, + se=se, + pvalue=pval, + conf_int=spec_ci, + ) + ) + + # Compute sensitivity ratio + all_atts = [baseline_att] + [s.att for s in specs] + ratio = _compute_sensitivity_ratio(baseline_att, all_atts) + level = _classify_robustness(ratio) + + if level == "not_estimable": + warnings.warn( + "Sensitivity ratio could not be estimated: baseline ATT is " + "non-finite or fewer than two specifications produced finite " + "estimates. Robustness to anticipation exclusions cannot be " + "assessed.", + UserWarning, + stacklevel=2, + ) + elif level in ("sensitive", "highly_sensitive"): + warnings.warn( + f"ATT estimates are {level} to anticipation exclusions " + f"(ratio={ratio:.3f}). Potential anticipation effects detected.", + UserWarning, + stacklevel=2, + ) + + return SensitivityResult( + specifications=specs, + baseline_att=baseline_att, + baseline_se=baseline_se, + sensitivity_ratio=ratio, + robustness_level=level, + n_specifications=len(specs) + 1, + baseline_pvalue=baseline_pval, + ) diff --git a/diff_diff/lwdid_staggered.py b/diff_diff/lwdid_staggered.py new file mode 100644 index 00000000..3dff9998 --- /dev/null +++ b/diff_diff/lwdid_staggered.py @@ -0,0 +1,817 @@ +"""Cohort-time estimation and joint aggregation for LWDiD.""" + +from __future__ import annotations + +import warnings +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import pandas as pd + +from diff_diff.lwdid_results import LWDiDResults +from diff_diff.utils import safe_inference + +CellKey = Tuple[Any, Any] + + +def _guard_standard_error(effect: float, se: float, scale: float = 0.0) -> float: + """Return NaN for numerically degenerate finite standard errors. + + The tolerance is RELATIVE to the problem's magnitude - the larger of + the effect and the caller-supplied data ``scale`` (e.g. the cell's + max |transformed outcome|). Round-6 review: the former + ``max(1, |effect|)`` floor made inference depend on the outcome's + UNITS - rescaling a valid fit by 1e-10 turned its finite SE into NaN + while the t-statistic is scale-invariant. Because every reference + (effect, scale, se) scales linearly with the outcome, the decision is + scale-equivariant; an exactly-fitting design (residuals at roundoff + of the DATA scale, e.g. the pure-trend zero-effect panel, or the G=2 + exactly-identified case that reported t ~ 5e15 pre-guard) still fails + closed because its se is roundoff RELATIVE TO ``scale``. Non-positive + and non-finite SEs are always rejected. + """ + tolerance = np.sqrt(np.finfo(float).eps) * max(abs(effect), abs(scale)) + if not np.isfinite(se) or se <= tolerance or se <= 0.0: + return np.nan + return float(se) + + +def _effective_influence( + influence: np.ndarray, + cluster_ids: Optional[np.ndarray], +) -> np.ndarray: + if cluster_ids is None: + return influence + frame = pd.DataFrame({"cluster": cluster_ids}) + columns = [] + for index in range(influence.shape[1]): + frame["value"] = influence[:, index] + columns.append(frame.groupby("cluster", sort=False)["value"].sum().to_numpy()) + return np.column_stack(columns) + + +def _combine_influence( + keys: List[CellKey], + weights: np.ndarray, + cell_influence: Dict[CellKey, np.ndarray], + n_units: int, +) -> Optional[np.ndarray]: + if any(key not in cell_influence for key in keys): + return None + combined = np.zeros(n_units, dtype=float) + for key, weight in zip(keys, weights): + combined += float(weight) * cell_influence[key] + return combined + + +def _inference_from_influence( + effect: float, + influence: Optional[np.ndarray], + alpha: float, + cluster_ids: Optional[np.ndarray], + *, + df_unclustered: Optional[int] = None, + contributing_mask: Optional[np.ndarray] = None, +) -> Tuple[float, float, float, Tuple[float, float], Optional[int]]: + """Aggregate-level inference from a combined influence vector. + + Reference-distribution policy (fix-wave WS6): clustered aggregates use + ``G - 1`` where G counts the clusters CONTRIBUTING to the aggregate + (``contributing_mask``; clusters supplying no estimated cell must not + inflate the df); unclustered aggregates composed of EXACTLY ONE cell + use that cell's residual df (``df_unclustered`` - matching the + common-timing rules, so a single-post-period staggered fit and the + common-timing fit of the same data agree); unclustered multi-cell + aggregates keep the large-sample normal reference (units recur across + cells with overlapping influence functions, so no residual-df pooling + is valid - documented in REGISTRY). + """ + if influence is None: + return np.nan, np.nan, np.nan, (np.nan, np.nan), None + effective = _effective_influence(influence[:, None], cluster_ids)[:, 0] + se = _guard_standard_error(effect, float(np.sqrt(np.sum(effective**2)))) + if not np.isfinite(se): + return np.nan, np.nan, np.nan, (np.nan, np.nan), None + if cluster_ids is not None: + ids = cluster_ids if contributing_mask is None else cluster_ids[contributing_mask] + df: Optional[int] = max(len(np.unique(ids)) - 1, 1) + else: + df = df_unclustered + t_stat, p_value, conf_int = safe_inference(effect, se, alpha=alpha, df=df) + return se, t_stat, p_value, conf_int, df + + +def _empty_cell( + g: Any, + t: Any, + reason: str, + n_treated: int = 0, + n_control: int = 0, +) -> Dict[str, Any]: + return { + "cohort": g, + "time": t, + "relative_time": t - g, + "att": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "n_treated": n_treated, + "n_control": n_control, + "df": None, + "skip_reason": reason, + "inference_status": "not_estimable", + } + + +def _transform_for_cohort( + estimator: Any, + frame: pd.DataFrame, + outcome: str, + unit: str, + time: str, + g: Any, +) -> pd.DataFrame: + pre_mask = frame[time] < g + if estimator.rolling == "demean": + return estimator._transform_demean(frame, outcome, unit, pre_mask) + if estimator.rolling == "detrend": + return estimator._transform_detrend(frame, outcome, unit, time, pre_mask) + if estimator.rolling == "demeanq": + return estimator._transform_demeanq(frame, outcome, unit, time, pre_mask) + return estimator._transform_detrendq(frame, outcome, unit, time, pre_mask) + + +def compute_event_study_bands( + estimator: Any, + event_effects: Dict[int, Dict[str, Any]], + event_influence: Dict[int, np.ndarray], + cluster_ids: Optional[np.ndarray], +) -> Tuple[ + Optional[np.ndarray], + Optional[np.ndarray], + Optional[str], + Optional[float], + Optional[int], +]: + """Analytical event-study covariance plus optional multiplier bootstrap. + + Shared by the staggered and common-timing paths. The analytical + covariance of the event-study effects is the cross-product of their + effective (cluster-summed) influence columns. When + ``estimator.n_bootstrap > 0`` the Rademacher multiplier bootstrap + replaces the analytical per-event SEs in ``event_effects`` in place, + attaches sup-t simultaneous ``cband_conf_int`` bounds, and suppresses + the (now inconsistent) analytical covariance. + + Returns + ------- + tuple + ``(event_vcov, event_vcov_index, cband_method, cband_crit_value, + cband_n_bootstrap)``. + """ + # Defense in depth (round-16 review): a label whose accepted SE is + # non-finite must not contribute a covariance row - its inference is + # NaN and a 0.0 diagonal would present it as known without + # uncertainty. + event_labels = sorted( + label + for label in event_influence + if np.isfinite(event_effects.get(label, {}).get("se", np.nan)) + ) + event_vcov = None + event_vcov_index = None + cband_method = None + cband_crit_value = None + cband_n_bootstrap = None + if event_labels: + influence_matrix = np.column_stack([event_influence[label] for label in event_labels]) + effective = _effective_influence(influence_matrix, cluster_ids) + event_vcov = effective.T @ effective + event_vcov_index = np.array(event_labels) + if estimator.n_bootstrap > 0: + rng = np.random.default_rng(estimator.seed) + centered = effective - effective.mean(axis=0, keepdims=True) + multipliers = rng.choice([-1.0, 1.0], size=(estimator.n_bootstrap, centered.shape[0])) + draws = multipliers @ centered + bootstrap_se = np.std(draws, axis=0, ddof=1) + valid = np.isfinite(bootstrap_se) & (bootstrap_se > 0) + if valid.any(): + sup_t = np.max(np.abs(draws[:, valid]) / bootstrap_se[valid], axis=1) + cband_crit_value = float(np.quantile(sup_t, 1 - estimator.alpha)) + cband_method = "multiplier_bootstrap_sup_t" + cband_n_bootstrap = estimator.n_bootstrap + invalid_labels = [label for index, label in enumerate(event_labels) if not valid[index]] + if invalid_labels: + # Fail closed: a requested-bootstrap cell whose draws are + # degenerate must not silently keep its analytical SE (an + # undocumented mixture of inference families - review + # round 2). Point retained, inference NaN. + warnings.warn( + f"Multiplier bootstrap produced degenerate draws for " + f"event time(s) {invalid_labels}; their inference is " + f"set to NaN (points retained) rather than silently " + f"reverting to analytical standard errors.", + UserWarning, + stacklevel=3, + ) + for index, label in enumerate(event_labels): + row = event_effects[label] + if not valid[index]: + row["se"] = float("nan") + row["t_stat"], row["p_value"], row["conf_int"] = safe_inference( + row["effect"], float("nan"), alpha=estimator.alpha + ) + row["inference_status"] = "degenerate_bootstrap" + continue + row["se"] = float(bootstrap_se[index]) + row["t_stat"], row["p_value"], row["conf_int"] = safe_inference( + row["effect"], row["se"], alpha=estimator.alpha, df=row.get("df") + ) + if cband_crit_value is not None: + row["cband_conf_int"] = ( + row["effect"] - cband_crit_value * row["se"], + row["effect"] + cband_crit_value * row["se"], + ) + # Bootstrap SEs replace the analytical diagonal, so do not expose + # an inconsistent analytical covariance matrix. + event_vcov = None + event_vcov_index = None + return event_vcov, event_vcov_index, cband_method, cband_crit_value, cband_n_bootstrap + + +def fit_staggered( + estimator: Any, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + cohort: str, + cluster: Optional[str], + controls: List[str], +) -> LWDiDResults: + """Estimate all supported cohort-time cells and aggregate them jointly.""" + varying = df.groupby(unit)[cohort].nunique(dropna=False) + if (varying > 1).any(): + raise ValueError( + f"Cohort must be time-invariant. Found {int((varying > 1).sum())} " + "unit(s) with varying cohort." + ) + # Unit-constancy of covariate and cluster columns is enforced by the + # shared fit() validation layer (LWDiD._validate_inputs), which covers + # this path and the common-timing path alike. + + unit_rows = df.drop_duplicates(subset=[unit], keep="first").set_index(unit) + all_units = unit_rows.index.to_list() + unit_to_index = {value: index for index, value in enumerate(all_units)} + cohort_by_unit = unit_rows[cohort] + never_mask = cohort_by_unit.isna() | (cohort_by_unit == 0) + never_units = cohort_by_unit.index[never_mask].to_list() + treated_cohorts = sorted( + value for value in pd.unique(df[cohort]) if pd.notna(value) and value > 0 + ) + if not treated_cohorts: + raise ValueError("No treated cohorts found.") + if estimator.control_group == "never_treated" and len(never_units) < 2: + raise ValueError( + "control_group='never_treated' requires at least 2 never-treated " + f"units for valid estimation; found {len(never_units)}." + ) + if estimator.control_group == "not_yet_treated" and not never_units: + raise ValueError( + "All units are eventually treated: control_group='not_yet_treated' " + "requires at least one never-treated unit (or an explicit " + "reference cohort, which is not supported). Without one, the " + "latest cohort-time cells have no valid control group and the " + "estimand would be silently truncated." + ) + + all_times = sorted(pd.unique(df[time])) + observed_time_set = set(all_times) + # Integer event-time contract (round-9 review): aggregation stores + # effects under int(t - g), so a fractional horizon (numeric calendar + # with non-integer spacing relative to a cohort) would silently MERGE + # distinct event times, overwriting estimates and covariance entries. + # Datetime/Period panels are already position-encoded (integral by + # construction); numeric panels are validated here and fail closed. + for g in treated_cohorts: + for t in all_times: + rel = float(t) - float(g) + if abs(rel - round(rel)) > 1e-9: + raise ValueError( + f"Event time t - g = {rel!r} (period {t!r}, cohort " + f"{g!r}) is not an integer: the event-study surface " + f"stores integer event-time keys and cannot represent " + f"fractional horizons without silently merging them. " + f"Encode the time and cohort columns as consecutive " + f"integer periods or as datetime/Period values." + ) + reference_periods = (-1,) if estimator.rolling in ("demean", "demeanq") else (-2, -1) + global_cluster_ids = None + if cluster is not None: + if cluster == unit: + # The unit column was consumed by set_index; read it from the index. + global_cluster_ids = unit_rows.index.to_numpy() + else: + global_cluster_ids = unit_rows.loc[all_units, cluster].to_numpy() + + cell_effects: Dict[CellKey, Dict[str, Any]] = {} + cell_influence: Dict[CellKey, np.ndarray] = {} + cell_members: Dict[CellKey, np.ndarray] = {} + single_cluster_cells: List[CellKey] = [] + skipped: List[Tuple[Any, Any, str]] = [] + cohort_sizes: Dict[Any, int] = {} + + for g in treated_cohorts: + treated_units = cohort_by_unit.index[cohort_by_unit == g].to_list() + cohort_sizes[g] = len(treated_units) + if estimator.control_group == "never_treated": + control_superset = never_units + else: + later = cohort_by_unit.index[cohort_by_unit > g].to_list() + control_superset = never_units + later + relevant_units = list(dict.fromkeys(treated_units + control_superset)) + cohort_frame = df.loc[df[unit].isin(relevant_units)].copy() + n_pre_periods = len([value for value in all_times if value < g]) + required_pre = 2 if estimator.rolling in ("detrend", "detrendq") else 1 + if n_pre_periods < required_pre: + for t in all_times: + if (t - g) not in reference_periods: + key = (g, t) + cell_effects[key] = _empty_cell(g, t, "insufficient_pre_periods") + skipped.append((g, t, "insufficient_pre_periods")) + continue + + transformed = _transform_for_cohort(estimator, cohort_frame, outcome, unit, time, g) + for t in all_times: + relative_time = t - g + if relative_time in reference_periods: + continue + key = (g, t) + if estimator.control_group == "never_treated": + valid_controls = set(never_units) + else: + threshold = max(g, t) + valid_controls = set(never_units) + valid_controls.update(cohort_by_unit.index[cohort_by_unit > threshold].to_list()) + + sample_units = set(treated_units) | valid_controls + columns = [unit, "_ydot"] + controls + if cluster is not None and cluster not in columns: + columns.append(cluster) + cell = transformed.loc[ + (transformed[time] == t) & transformed[unit].isin(sample_units), columns + ].drop_duplicates(subset=[unit], keep="first") + finite = np.isfinite(cell["_ydot"].to_numpy(dtype=float)) + if controls: + finite &= np.all(np.isfinite(cell[controls].to_numpy(dtype=float)), axis=1) + cell = cell.loc[finite].copy() + treatment = cell[unit].isin(treated_units).to_numpy(dtype=float) + n_treated = int(treatment.sum()) + n_control = int(len(treatment) - n_treated) + if n_treated == 0 or n_control == 0: + cell_effects[key] = _empty_cell(g, t, "zero_treated_control", n_treated, n_control) + skipped.append((g, t, "zero_treated_control")) + continue + if estimator.control_group == "never_treated" and n_control < 2: + # Registry: the NT-only design requires at least 2 + # never-treated controls. The raw-unit guard runs pre-fit, + # but transformation drops / unbalanced availability can + # leave a single control in a cell (round-4 review) - mark + # it non-estimable rather than estimate on one control. + cell_effects[key] = _empty_cell( + g, t, "insufficient_never_treated_controls", n_treated, n_control + ) + skipped.append((g, t, "insufficient_never_treated_controls")) + continue + + y = cell["_ydot"].to_numpy(dtype=float) + controls_matrix = cell[controls].to_numpy(dtype=float) if controls else None + cluster_ids = None + cell_single_cluster = False + if cluster is not None: + cluster_ids = cell[cluster].to_numpy() + if len(np.unique(cluster_ids)) < 2: + # Cluster ids are re-derived per cell, so a cell whose + # units share one cluster can exist with G >= 2 + # globally. Estimate the POINT unclustered and fail the + # inference closed below (campaign finding: ipw/dr + # silently fell back to unclustered variance here under + # a CR1 label; reg raised mid-fit). + cell_single_cluster = True + cluster_ids = None + try: + att, se, _, _, n_params, influence = estimator._dispatch_estimator( + y, treatment, controls_matrix, cluster_ids, len(cell) + ) + except ValueError as exc: + if "Invalid exact-inference design" in str(exc): + # Non-estimable cell (Registry: NaN, not a mid-fit raise). + cell_effects[key] = _empty_cell( + g, t, "insufficient_sample", n_treated, n_control + ) + skipped.append((g, t, "insufficient_sample")) + continue + raise + if not np.isfinite(att): + cell_effects[key] = _empty_cell(g, t, "non_finite_estimate", n_treated, n_control) + skipped.append((g, t, "non_finite_estimate")) + continue + + se = _guard_standard_error(att, se, scale=float(np.max(np.abs(y)))) + if cell_single_cluster: + # Fail closed: point retained, inference NaN; aggregates + # that include this cell inherit NaN inference (deliberate + # - see the aggregated warning below and REGISTRY). + single_cluster_cells.append(key) + se = np.nan + influence = None + if cluster_ids is not None: + df_cell = max(len(np.unique(cluster_ids)) - 1, 1) + else: + # n_params is the fitted design's parameter count, so the + # residual df is design-coherent for every method. + # Raw residual df: safe_inference fails the tuple closed + # when df <= 0 (no fabricated df=1 - review finding). + df_cell = len(cell) - n_params + t_stat, p_value, conf_int = safe_inference(att, se, alpha=estimator.alpha, df=df_cell) + cell_effects[key] = { + "cohort": g, + "time": t, + "relative_time": relative_time, + "att": float(att), + "se": se, + "t_stat": t_stat, + "p_value": p_value, + "conf_int": conf_int, + "n_treated": n_treated, + "n_control": n_control, + "df": df_cell, + "skip_reason": None, + "inference_status": "ok" if np.isfinite(se) else "degenerate", + } + member_mask = np.zeros(len(all_units), dtype=bool) + for unit_value in cell[unit].unique(): + member_mask[unit_to_index[unit_value]] = True + cell_members[key] = member_mask + if influence is not None and np.isfinite(se): + global_influence = np.zeros(len(all_units), dtype=float) + for local_index, unit_value in enumerate(cell[unit].to_list()): + global_influence[unit_to_index[unit_value]] = influence[local_index] + cell_influence[key] = global_influence + + if single_cluster_cells: + listed = ", ".join(str(k) for k in single_cluster_cells[:6]) + suffix = ( + "" if len(single_cluster_cells) <= 6 else f"; plus {len(single_cluster_cells) - 6} more" + ) + warnings.warn( + f"LWDiD: cohort-time cell(s) {listed}{suffix} contain fewer than " + "2 clusters, so their cluster-robust inference is not identified. " + "Cell points are retained with NaN inference; any aggregate that " + "includes such a cell reports NaN inference as well (fail-closed).", + UserWarning, + stacklevel=2, + ) + if skipped: + preview = ", ".join(f"({g}, {t}): {reason}" for g, t, reason in skipped[:6]) + suffix = "" if len(skipped) <= 6 else f"; plus {len(skipped) - 6} more" + warnings.warn( + f"LWDiD skipped {len(skipped)} unsupported cohort-time cell(s): " f"{preview}{suffix}.", + UserWarning, + stacklevel=2, + ) + + cohort_effects: Dict[Any, Dict[str, Any]] = {} + cohort_influence: Dict[Any, np.ndarray] = {} + # Treated-unit positions per cohort, for CONTRIBUTING-mass weighting + # (round-23 review: cohort_sizes counts RAW cohort members, so a + # treated unit contributing to no estimable post cell still raised + # its cohort's overall weight on non-tau_omega paths). + treated_positions_by_cohort: Dict[Any, np.ndarray] = {} + for g in treated_cohorts: + positions = np.zeros(len(all_units), dtype=bool) + for u in cohort_by_unit.index[cohort_by_unit == g]: + positions[unit_to_index[u]] = True + treated_positions_by_cohort[g] = positions + contributing_sizes: Dict[Any, int] = {} + for g in treated_cohorts: + keys = [ + key + for key, value in cell_effects.items() + if key[0] == g and key[1] >= g and np.isfinite(value["att"]) + ] + if not keys: + continue + # Within-cohort CELL-MASS convention (documented deviation from + # LW 2026 eq. 7.10 on unbalanced panels - see the REGISTRY + # within-cohort aggregation Note): cells weight by contributing + # treated mass, the same convention as the WATT(r) event axis + # (E.1) and the package's Post_avg display; equals the eq. 7.10 + # unit-average estimand on balanced NT designs. + masses = np.array([cell_effects[key]["n_treated"] for key in keys], dtype=float) + weights = masses / masses.sum() + effect = float(np.dot(weights, [cell_effects[key]["att"] for key in keys])) + influence = _combine_influence(keys, weights, cell_influence, len(all_units)) + mask = np.zeros(len(all_units), dtype=bool) + for key in keys: + mask |= cell_members.get(key, False) + se, t_stat, p_value, conf_int, df_group = _inference_from_influence( + effect, + influence, + estimator.alpha, + global_cluster_ids, + df_unclustered=(cell_effects[keys[0]]["df"] if len(keys) == 1 else None), + contributing_mask=mask, + ) + contributing_sizes[g] = int((mask & treated_positions_by_cohort[g]).sum()) + cohort_effects[g] = { + "cohort": g, + "att": effect, + "se": se, + "t_stat": t_stat, + "p_value": p_value, + "conf_int": conf_int, + "n_treated": contributing_sizes[g], + "n_control": max(cell_effects[key]["n_control"] for key in keys), + "n_cells": len(keys), + "df": df_group, + } + if influence is not None: + cohort_influence[g] = influence + + if not cohort_effects: + raise ValueError("No supported post-treatment cohort-time cells were estimable.") + + valid_cohorts = list(cohort_effects) + overall_keys = [ + key + for key, value in cell_effects.items() + if key[0] in cohort_effects and key[1] >= key[0] and np.isfinite(value["att"]) + ] + overall_cluster_mask = np.zeros(len(all_units), dtype=bool) + for key in overall_keys: + overall_cluster_mask |= cell_members.get(key, False) + # Overall masses from treated units CONTRIBUTING to each cohort's + # estimable post cells (the Registry's contributing-sample rule; raw + # cohort membership previously weighted non-contributing units in). + cohort_masses = np.array([contributing_sizes[g] for g in valid_cohorts], dtype=float) + cohort_weights = cohort_masses / cohort_masses.sum() + for g, weight in zip(valid_cohorts, cohort_weights): + cohort_effects[g]["weight"] = float(weight) + overall_effect = float( + np.dot(cohort_weights, [cohort_effects[g]["att"] for g in valid_cohorts]) + ) + # LW 2026 (7.16)/(7.18): with never-treated controls, regression + # adjustment and no covariates, the overall estimand is tau_omega -- + # the coefficient on D in the composite-outcome cross-sectional + # regression, which averages each unit's transformed outcome over its + # OBSERVED post periods. The composite is defined for the plain + # demean/detrend transforms only (the q variants have no seasonal + # composite; their overall is the cohort-mass-weighted average of + # seasonal cohort ATTs). + tau_omega_config = ( + estimator.control_group == "never_treated" + and estimator.estimation_method == "reg" + and not controls + and estimator.rolling in ("demean", "detrend") + ) + use_composite = tau_omega_config and estimator.vcov_type == "classical" and cluster is None + # Complete-case resolution: the composite is computed ONCE for every + # tau_omega-eligible configuration. With ZERO complete-case drops the + # composite point is reported on BOTH vcov routes (status quo: the + # classical route pairs it with the composite's own SE, hc1/clustered + # with the joint-IF SE -- the documented approximation in REGISTRY). + # With ANY drops, `.att` is the IF-weighted cohort-mass point on ALL + # routes (the same point under every vcov setting -- a variance + # selection must never move the point) and the complete-case composite + # is exposed as the diagnostic `att_tau_omega_complete_case`. + att_tau_omega_complete_case: Optional[float] = None + n_composite_treated_dropped = 0 + n_composite_controls_dropped = 0 + composite_is_att = False + comp_att = comp_se = np.nan + comp_df = 0 + comp_scale = 0.0 + comp_surviving_sizes: Dict[Any, int] = {} + if tau_omega_config: + ( + comp_att, + comp_se, + comp_df, + n_composite_treated_dropped, + n_composite_controls_dropped, + comp_scale, + comp_surviving_sizes, + ) = estimator._composite_regression_aggregation(df, outcome, unit, time, cohort) + composite_drops = n_composite_treated_dropped + n_composite_controls_dropped + if composite_drops == 0 and np.isfinite(comp_att): + composite_is_att = True + else: + att_tau_omega_complete_case = float(comp_att) if np.isfinite(comp_att) else None + warnings.warn( + "LWDiD: the tau_omega composite required complete-case " + f"drops ({n_composite_treated_dropped} treated, " + f"{n_composite_controls_dropped} control unit(s)) on this " + "unbalanced panel, so `.att` reports the influence-weighted " + "cohort-mass point (identical under every vcov setting) " + "instead of tau_omega. The complete-case composite is " + "available as `att_tau_omega_complete_case`. See " + "docs/methodology/REGISTRY.md (LWDiD).", + UserWarning, + stacklevel=2, + ) + if tau_omega_config and not composite_is_att and comp_surviving_sizes: + # DROPS route: the Registry complete-case rule fixes cohort masses + # on the SURVIVING treated sample (round-12 review: the raw masses + # still weighted dropped treated units into `.att` and its + # combined influence function). + surviving_masses = np.array( + [comp_surviving_sizes.get(g, 0) for g in valid_cohorts], dtype=float + ) + if surviving_masses.sum() > 0: + cohort_weights = surviving_masses / surviving_masses.sum() + for g, weight in zip(valid_cohorts, cohort_weights): + cohort_effects[g]["weight"] = float(weight) + overall_effect = float( + np.dot(cohort_weights, [cohort_effects[g]["att"] for g in valid_cohorts]) + ) + if composite_is_att: + overall_effect = float(comp_att) + if use_composite and composite_is_att: + overall_se = _guard_standard_error(overall_effect, comp_se, scale=comp_scale) + overall_df = comp_df + inference_basis = "composite_regression" + else: + overall_influence = None + missing = [g for g in valid_cohorts if g not in cohort_influence] + if not missing: + overall_influence = sum( + float(weight) * cohort_influence[g] + for g, weight in zip(valid_cohorts, cohort_weights) + ) + overall_se, _, _, _, overall_df = _inference_from_influence( + overall_effect, + overall_influence, + estimator.alpha, + global_cluster_ids, + df_unclustered=( + cell_effects[overall_keys[0]]["df"] if len(overall_keys) == 1 else None + ), + contributing_mask=overall_cluster_mask, + ) + if overall_influence is not None: + inference_basis = "joint_influence_function" + elif estimator.estimation_method == "psm": + inference_basis = "unavailable_matching" + warnings.warn( + "LWDiD: propensity-score matching has no influence-function " + "representation, so cohort effects cannot be combined without " + "assuming independence. Overall inference is reported as NaN; " + "use estimation_method='dr' for a doubly robust alternative " + "with valid joint inference.", + UserWarning, + stacklevel=2, + ) + else: + inference_basis = "unavailable_degenerate_cells" + listed = ", ".join(str(g) for g in missing) + warnings.warn( + f"LWDiD: cohort(s) {listed} contain cohort-time cells with a " + "degenerate or non-finite standard error, so no joint influence " + "function is available. Overall inference is reported as NaN.", + UserWarning, + stacklevel=2, + ) + overall_t, overall_p, overall_ci = safe_inference( + overall_effect, overall_se, alpha=estimator.alpha, df=overall_df + ) + + event_effects: Dict[int, Dict[str, Any]] = {} + event_influence: Dict[int, np.ndarray] = {} + for relative_time in sorted({value["relative_time"] for value in cell_effects.values()}): + keys = [ + key + for key, value in cell_effects.items() + if value["relative_time"] == relative_time and np.isfinite(value["att"]) + ] + if not keys: + continue + masses = np.array([cell_effects[key]["n_treated"] for key in keys], dtype=float) + weights = masses / masses.sum() + effect = float(np.dot(weights, [cell_effects[key]["att"] for key in keys])) + influence = _combine_influence(keys, weights, cell_influence, len(all_units)) + mask = np.zeros(len(all_units), dtype=bool) + for key in keys: + mask |= cell_members.get(key, False) + se, t_stat, p_value, conf_int, df_event = _inference_from_influence( + effect, + influence, + estimator.alpha, + global_cluster_ids, + df_unclustered=(cell_effects[keys[0]]["df"] if len(keys) == 1 else None), + contributing_mask=mask, + ) + event_effects[int(relative_time)] = { + "effect": effect, + "se": se, + "t_stat": t_stat, + "p_value": p_value, + "conf_int": conf_int, + "n_treated": int(masses.sum()), + "n_cells": len(keys), + "df": df_event, + } + if influence is not None and np.isfinite(se): + # Round-16 review: storing a degenerate-influence column + # (NaN-inference row) exposed a 0.0 covariance diagonal for a + # row whose se/t/p/CI are all NaN. Matches the common-timing + # guard (influence is not None AND finite se). + event_influence[int(relative_time)] = influence + + ( + event_vcov, + event_vcov_index, + cband_method, + cband_crit_value, + cband_n_bootstrap, + ) = compute_event_study_bands(estimator, event_effects, event_influence, global_cluster_ids) + + # Sample metadata describes the CELL-ESTIMATION sample: units actually + # contributing to at least one estimated cell (campaign finding: the + # previous counts covered every input unit, overstating the sample + # whenever cell-level drops occurred). Complete-case composite drops + # are reported separately via n_composite_*_dropped. + contributing_units_mask = np.zeros(len(all_units), dtype=bool) + for member_mask in cell_members.values(): + contributing_units_mask |= member_mask + contributing_index = np.flatnonzero(contributing_units_mask) + contributing_ids = [all_units[i] for i in contributing_index] + never_set = set(never_units) + n_contrib_control = sum(1 for u in contributing_ids if u in never_set) + n_contrib_treated = int((~never_mask.loc[contributing_ids]).sum()) + result = LWDiDResults( + att=float(overall_effect), + se=float(overall_se), + t_stat=overall_t, + p_value=overall_p, + conf_int=overall_ci, + n_obs=int(contributing_units_mask.sum()), + n_treated=n_contrib_treated, + n_control=n_contrib_control, + rolling=estimator.rolling, + estimation_method=estimator.estimation_method, + vcov_type=estimator.vcov_type, + alpha=estimator.alpha, + df_inference=overall_df, + cluster_name=cluster, + control_group=estimator.control_group, + n_bootstrap=estimator.n_bootstrap, + seed=estimator.seed, + pscore_trim=( + estimator.pscore_trim if estimator.estimation_method in ("ipw", "dr", "psm") else None + ), + psm_config=( + { + "pscore_trim": estimator.pscore_trim, + "n_neighbors": estimator.n_neighbors, + "caliper": estimator.caliper, + "with_replacement": estimator.with_replacement, + } + if estimator.estimation_method == "psm" + else None + ), + n_clusters=( + len(np.unique(global_cluster_ids[overall_cluster_mask])) + if global_cluster_ids is not None + else None + ), + cohort_effects=cohort_effects, + cohort_time_effects=cell_effects, + inference_basis=inference_basis, + att_tau_omega_complete_case=att_tau_omega_complete_case, + n_composite_treated_dropped=n_composite_treated_dropped, + n_composite_controls_dropped=n_composite_controls_dropped, + event_study_effects=event_effects, + event_study_vcov=event_vcov, + event_study_vcov_index=event_vcov_index, + event_study_df={ + label: value["df"] + for label, value in event_effects.items() + if value.get("df") is not None + }, + # Only OBSERVED anchors are emitted (Registry: the zero-valued + # is_reference rows are a display convention for anchors that + # exist in the panel; round-4 review - a numeric time gap could + # otherwise synthesize a zero effect at a nonexistent event time). + reference_periods=tuple( + r + for r in reference_periods + if any((g + r) in observed_time_set for g in treated_cohorts) + ), + cband_method=cband_method, + cband_crit_value=cband_crit_value, + cband_n_bootstrap=cband_n_bootstrap, + ) + return result diff --git a/diff_diff/lwdid_visualization.py b/diff_diff/lwdid_visualization.py new file mode 100644 index 00000000..ec85c0f6 --- /dev/null +++ b/diff_diff/lwdid_visualization.py @@ -0,0 +1,303 @@ +"""Visualization methods for LWDiD results. + +Provides plotting functions for cohort trends, event studies, +sensitivity analysis, and bootstrap distributions. + +Requires matplotlib (optional dependency). If not installed, +raises ImportError with installation instructions. + +Note +---- +All plot functions return a matplotlib Figure object without closing it. +In batch/loop usage, call ``plt.close(fig)`` after saving or displaying +each figure to avoid memory accumulation. +""" + +from typing import Any, Optional + +import numpy as np +import pandas as pd + + +def _require_matplotlib(): + try: + import matplotlib.pyplot as plt + + return plt + except ImportError: + raise ImportError( + "matplotlib is required for LWDiD visualization. " + "Install with: pip install matplotlib" + ) + + +def plot_cohort_trends( + data: pd.DataFrame, + outcome: str, + unit: str, + time: str, + treatment: str, + cohort: Optional[str] = None, + title: Optional[str] = None, + figsize: tuple = (10, 6), + show_ci: bool = True, + ax=None, +): + """Plot pre/post outcome trajectories by treatment group or cohort. + + Without ``cohort=``, shows average outcomes over time for the + ever-treated vs control groups. With ``cohort=`` (round-23 review: + the parameter was previously accepted but silently ignored), one + trajectory is drawn PER treated cohort (never-treated encodings + 0/NaN form the control line) with a per-cohort onset marker. + Optional confidence bands in both modes. + """ + plt = _require_matplotlib() + + if ax is None: + fig, ax = plt.subplots(figsize=figsize) + else: + fig = ax.get_figure() + + data = data.copy() + if cohort is not None: + cohort_by_unit = data.drop_duplicates(subset=[unit], keep="first").set_index(unit)[cohort] + never_mask = cohort_by_unit.isna() | (cohort_by_unit == 0) + data["_plot_group"] = data[unit].map( + { + u: ("Control" if never_mask[u] else f"Cohort {cohort_by_unit[u]}") + for u in cohort_by_unit.index + } + ) + group_order = sorted({g for g in data["_plot_group"].unique() if g != "Control"}) + ( + ["Control"] if (data["_plot_group"] == "Control").any() else [] + ) + else: + treated_units = data.loc[data[treatment] == 1, unit].unique() + data["_plot_group"] = np.where(data[unit].isin(treated_units), "Treated", "Control") + group_order = ["Treated", "Control"] + + group_means = ( + data.groupby([time, "_plot_group"])[outcome].agg(["mean", "std", "count"]).reset_index() + ) + group_means["se"] = group_means["std"] / np.sqrt(group_means["count"]) + + colors = plt.rcParams["axes.prop_cycle"].by_key().get("color", ["steelblue", "coral"]) + for i, label in enumerate(group_order): + gdf = group_means[group_means["_plot_group"] == label] + color = "coral" if label == "Control" else colors[i % len(colors)] + ax.plot(gdf[time], gdf["mean"], "o-", label=label, color=color) + if show_ci: + ax.fill_between( + gdf[time], + gdf["mean"] - 1.96 * gdf["se"], + gdf["mean"] + 1.96 * gdf["se"], + alpha=0.15, + color=color, + ) + + # Onset markers: one per cohort when cohort= is given, else the + # common onset. Datetime/Period/string time columns cannot take + # `- 0.5` (raw TypeError pre-fix); draw AT the onset for non-numeric + # scales, offset by half a period for numeric ones. + def _onset_x(value): + return value - 0.5 if pd.api.types.is_numeric_dtype(data[time]) else value + + if cohort is not None: + onsets = sorted({v for v in cohort_by_unit.dropna().unique() if not (pd.isna(v) or v == 0)}) + for i, g in enumerate(onsets): + ax.axvline( + _onset_x(g), + color="gray", + linestyle="--", + alpha=0.7, + label="Cohort onsets" if i == 0 else None, + ) + else: + treated_times = data.loc[data[treatment] == 1, time] + if len(treated_times) > 0: + ax.axvline( + _onset_x(treated_times.min()), + color="gray", + linestyle="--", + alpha=0.7, + label="Treatment onset", + ) + + ax.set_xlabel("Time") + ax.set_ylabel(outcome) + ax.set_title(title or "LWDiD: Cohort Trends") + ax.legend() + ax.grid(True, alpha=0.3) + + return fig + + +def plot_event_study( + results: Any, + title: Optional[str] = None, + figsize: tuple = (10, 6), + ax=None, +): + """Plot event-study estimates from a fitted LWDiD result. + + Consumes the unified post-fit event-study surface + (``results.event_study_effects``, keyed by event time relative to + first treatment, with ``reference_periods`` anchored at zero). + Both staggered and common-timing fits populate this surface at fit + time; a result without one raises a clear error instead of plotting. + + Parameters + ---------- + results : LWDiDResults + Fitted result whose event-study surface is populated. + title : str or None + Plot title. + figsize : tuple + Figure size (ignored when ``ax`` is supplied). + ax : matplotlib Axes or None + Axes to draw on; a new figure is created when None. + + Raises + ------ + ValueError + If ``results`` carries no populated event-study surface (e.g. a + degenerate fit with no estimable post period). + """ + effects = getattr(results, "event_study_effects", None) + if not effects: + raise ValueError( + "plot_event_study requires a fitted result with a populated " + "event-study surface (results.event_study_effects); this fit " + "does not carry one (no estimable post-period effects)." + ) + + plt = _require_matplotlib() + + if ax is None: + fig, ax = plt.subplots(figsize=figsize) + else: + fig = ax.get_figure() + + reference_periods = set(getattr(results, "reference_periods", ()) or ()) + event_times = sorted(set(effects) | reference_periods) + atts = [] + err_lo = [] + err_hi = [] + for r in event_times: + if r in reference_periods and r not in effects: + atts.append(0.0) + err_lo.append(0.0) + err_hi.append(0.0) + continue + row = effects[r] + att_r = row.get("effect", np.nan) + atts.append(att_r) + # Round-21 review: render the FITTED interval endpoints (t-based + # per-row df, fitted alpha, cband/bootstrap intervals), not a + # fabricated normal-theory +/-1.96*SE. Preference: simultaneous + # cband when present, else the stored conf_int. House rule: OMIT + # the interval when unavailable/non-finite (a zero-length bar + # would render an inference-unavailable effect as infinitely + # precise); reference periods keep their deliberate zero bars. + interval = row.get("cband_conf_int") or row.get("conf_int") + if interval is not None and np.all(np.isfinite(interval)) and np.isfinite(att_r): + err_lo.append(att_r - float(interval[0])) + err_hi.append(float(interval[1]) - att_r) + else: + err_lo.append(np.nan) + err_hi.append(np.nan) + + yerr = np.vstack([err_lo, err_hi]) + ax.errorbar(event_times, atts, yerr=yerr, fmt="o-", capsize=3, color="steelblue") + ax.axhline(0, color="gray", linestyle="--", alpha=0.5) + ax.set_xlabel("Event time") + ax.set_ylabel("ATT") + ax.set_title(title or "LWDiD: Event-Study Effects") + ax.grid(True, alpha=0.3) + + return fig + + +def plot_sensitivity( + sensitivity_result, + title: Optional[str] = None, + figsize: tuple = (10, 6), + ax=None, +): + """Plot sensitivity analysis results. + + Shows ATT estimates across different specifications with + confidence bands, highlighting the baseline estimate. + """ + plt = _require_matplotlib() + + if ax is None: + fig, ax = plt.subplots(figsize=figsize) + else: + fig = ax.get_figure() + + specs = sensitivity_result.specifications + x = range(len(specs)) + atts = [s.att for s in specs] + labels = [s.label for s in specs] + + # Round-21 review: render each specification's FITTED interval + # endpoints when stored; failed/legacy specs show the point only + # (no fabricated normal-theory interval). + err_lo = [] + err_hi = [] + for s_ in specs: + interval = getattr(s_, "conf_int", None) + if interval is not None and np.all(np.isfinite(interval)) and np.isfinite(s_.att): + err_lo.append(s_.att - float(interval[0])) + err_hi.append(float(interval[1]) - s_.att) + else: + err_lo.append(np.nan) + err_hi.append(np.nan) + yerr = np.vstack([err_lo, err_hi]) + ax.errorbar(x, atts, yerr=yerr, fmt="o", capsize=3, color="steelblue") + ax.axhline( + sensitivity_result.baseline_att, + color="red", + linestyle="--", + alpha=0.7, + label="Baseline ATT", + ) + ax.set_xticks(list(x)) + ax.set_xticklabels(labels, rotation=45, ha="right") + ax.set_ylabel("ATT") + ax.set_title( + title or f"Sensitivity Analysis (robustness: {sensitivity_result.robustness_level})" + ) + ax.legend() + ax.grid(True, alpha=0.3) + plt.tight_layout() + + return fig + + +def plot_bootstrap_distribution( + t_stats: np.ndarray, + t_observed: float, + title: Optional[str] = None, + figsize: tuple = (8, 5), + ax=None, +): + """Plot bootstrap t-statistic distribution with observed value.""" + plt = _require_matplotlib() + + if ax is None: + fig, ax = plt.subplots(figsize=figsize) + else: + fig = ax.get_figure() + + ax.hist(t_stats, bins=50, density=True, alpha=0.7, color="steelblue", edgecolor="white") + ax.axvline(t_observed, color="red", linewidth=2, label=f"t_obs = {t_observed:.3f}") + ax.axvline(-t_observed, color="red", linewidth=2, linestyle="--", alpha=0.5) + ax.set_xlabel("t-statistic") + ax.set_ylabel("Density") + ax.set_title(title or "Wild Cluster Bootstrap Distribution") + ax.legend() + + return fig diff --git a/diff_diff/lwdid_wild_bootstrap.py b/diff_diff/lwdid_wild_bootstrap.py new file mode 100644 index 00000000..ce1e520a --- /dev/null +++ b/diff_diff/lwdid_wild_bootstrap.py @@ -0,0 +1,360 @@ +"""Wild cluster bootstrap for inference with few clusters. + +Thin LWDiD-facing wrapper over the house Wild Cluster Restricted (WCR) +bootstrap engine (:func:`diff_diff.utils.wild_bootstrap_se`, matched to R's +``fwildclusterboot::boottest``): the null is genuinely imposed by dropping +the treatment column from the restricted model (controls retained), the CI +is obtained by inverting the bootstrap test, and Rademacher weights are +fully enumerated automatically when ``2**G <= n_bootstrap`` and ``G <= 20``. + +The wild cluster bootstrap is recommended when: + +- Number of clusters G < 30 +- Cluster sizes are unbalanced +- Few treated clusters + +P-value convention: the house strict-exceedance count with a ~1e-9 relative +tie guard and a documented zero-p floor at ``1/(n_valid + 1)`` when that +floor is below ``alpha`` (a deliberate, documented departure from +``boottest`` — see ``diff_diff/utils.py``). This differs from the +randomization-inference module's inclusive Phipson-Smyth rule; both are +documented in ``docs/methodology/REGISTRY.md``. + +References +---------- +Cameron, A. C., Gelbach, J. B., & Miller, D. L. (2008). Bootstrap-based +improvements for inference with clustered errors. *Review of Economics +and Statistics*, 90(3), 414-427. + +Roodman, D., MacKinnon, J. G., Nielsen, M. O., & Webb, M. D. (2019). Fast +and wild: Bootstrap inference in Stata using boottest. *The Stata +Journal*, 19(1), 4-60. + +Webb, M. D. (2014). Reworking wild bootstrap based inference for clustered +errors. *Queen's Economics Department Working Paper*, No. 1315. +""" + +from __future__ import annotations + +import warnings +from dataclasses import dataclass, field +from typing import Optional + +import numpy as np + +from diff_diff.linalg import solve_ols +from diff_diff.utils import wild_bootstrap_se + +_VALID_WEIGHT_TYPES = ("rademacher", "mammen", "webb") + + +@dataclass +class WildClusterBootstrapResult: + """Result of wild cluster bootstrap inference. + + Attributes + ---------- + att : float + Point estimate of the average treatment effect on the treated + (coefficient on the treatment column of the unrestricted OLS). + se : float + Analytical cluster-robust (CR1) standard error of ``att``. The + studentized bootstrap drives the p-value and CI; this is not a + rescaled bootstrap dispersion. + t_stat_original : float + Studentized statistic of the original estimate, ``att / se``. + p_value : float + Wild cluster bootstrap p-value (two-tailed; house convention — + strict exceedance with tie guard and documented zero-p floor). + ci_lower : float + Lower bound of the confidence interval (by test inversion). + ci_upper : float + Upper bound of the confidence interval (by test inversion). + n_clusters : int + Number of clusters in the (post-drop) data. + n_bootstrap : int + Number of bootstrap replications actually performed (equals + ``2**n_clusters`` under automatic full enumeration). + weight_type : str + Weight distribution used ('rademacher', 'mammen', or 'webb'). + alpha : float + Significance level used for the CI. + bootstrap_distribution : np.ndarray or None + Bootstrap t* distribution (finite-filtered, so its length may be + below ``n_bootstrap``); ``None`` when the degenerate guard fired. + n_dropped : int + Observations dropped for non-finite ``y`` (warned). + """ + + att: float + se: float + t_stat_original: float + p_value: float + ci_lower: float + ci_upper: float + n_clusters: int + n_bootstrap: int + weight_type: str + alpha: float + bootstrap_distribution: Optional[np.ndarray] = field(repr=False, default=None) + n_dropped: int = 0 + + def summary(self) -> str: + """Return a human-readable summary string.""" + sig = ( + "***" + if self.p_value < 0.01 + else "**" if self.p_value < 0.05 else "*" if self.p_value < 0.1 else "" + ) + level = int(round((1 - self.alpha) * 100)) + return ( + f"Wild Cluster Bootstrap Results\n" + f"{'=' * 50}\n" + f"ATT: {self.att:.4f} {sig}\n" + f"Cluster-robust (CR1) SE: {self.se:.4f}\n" + f"{level}% CI (test inversion): [{self.ci_lower:.4f}, {self.ci_upper:.4f}]\n" + f"P-value: {self.p_value:.4f}\n" + f"N clusters: {self.n_clusters}\n" + f"N bootstrap reps: {self.n_bootstrap}\n" + f"Weight type: {self.weight_type}\n" + f"{'=' * 50}" + ) + + +def wild_cluster_bootstrap( + y: np.ndarray, + treatment: np.ndarray, + cluster_ids: np.ndarray, + controls: Optional[np.ndarray] = None, + *, + n_bootstrap: int = 999, + weight_type: str = "rademacher", + alpha: float = 0.05, + seed: Optional[int] = None, +) -> WildClusterBootstrapResult: + """Perform wild cluster restricted bootstrap inference (CGM 2008). + + Delegates to the house engine :func:`diff_diff.utils.wild_bootstrap_se` + (``fwildclusterboot::boottest``-matched): the null is imposed by + re-estimating with the treatment column dropped while KEEPING the + controls (the earlier module-local implementation fit an intercept-only + restricted model, dumping covariate signal into the bootstrap + residuals), the CI is obtained by test inversion, and Rademacher full + enumeration engages automatically at ``2**G <= n_bootstrap``. + + Parameters + ---------- + y : np.ndarray, shape (N,) + Outcome variable. Non-finite entries are dropped with a warning + (see ``n_dropped`` on the result). + treatment : np.ndarray, shape (N,) + Binary treatment indicator (0/1). + cluster_ids : np.ndarray, shape (N,) + Cluster membership for each observation. + controls : np.ndarray or None, shape (N, p) + Optional matrix of control variables. Non-finite entries raise + ``ValueError`` (impute or remove before calling). + n_bootstrap : int, default 999 + Number of bootstrap replications (reported as ``2**G`` when full + enumeration engages). + weight_type : str, default 'rademacher' + Bootstrap weight distribution: 'rademacher', 'mammen', or 'webb'. + alpha : float, default 0.05 + Significance level for the test-inversion confidence interval. + seed : int or None, default None + Random seed for reproducibility. + + Returns + ------- + WildClusterBootstrapResult + Point estimate with CR1 SE, test-inversion CI, bootstrap p-value, + and the finite-filtered t* distribution. + + Raises + ------ + ValueError + On incompatible shapes, an invalid ``weight_type``, non-finite + controls, or fewer than 2 clusters. + + Examples + -------- + >>> import numpy as np + >>> from diff_diff.lwdid_wild_bootstrap import wild_cluster_bootstrap + >>> rng = np.random.default_rng(42) + >>> n = 200 + >>> y = rng.normal(0, 1, n) + >>> y[:50] += 1.5 + >>> treatment = np.zeros(n); treatment[:50] = 1.0 + >>> cluster_ids = np.repeat(np.arange(20), 10) + >>> result = wild_cluster_bootstrap(y, treatment, cluster_ids, seed=123) + >>> print(f"ATT={result.att:.3f}, p={result.p_value:.3f}") + ATT=1.662, p=0.001 + """ + # ----- Input validation ----- + y = np.asarray(y, dtype=np.float64).ravel() + treatment = np.asarray(treatment, dtype=np.float64).ravel() + cluster_ids = np.asarray(cluster_ids).ravel() + + N = len(y) + if N == 0: + raise ValueError("y must not be empty.") + if len(treatment) != N: + raise ValueError(f"Length mismatch: y has {N} obs but treatment has {len(treatment)}.") + if len(cluster_ids) != N: + raise ValueError(f"Length mismatch: y has {N} obs but cluster_ids has {len(cluster_ids)}.") + if not np.all((treatment == 0) | (treatment == 1)): + raise ValueError( + "treatment must be binary (0 or 1). " + f"Got values in [{treatment.min()}, {treatment.max()}]." + ) + if treatment.sum() == 0: + raise ValueError("No treated observations (treatment is all zeros).") + if treatment.sum() == N: + raise ValueError("No control observations (treatment is all ones).") + + if controls is not None: + controls = np.asarray(controls, dtype=np.float64) + if controls.ndim == 1: + controls = controls.reshape(-1, 1) + if controls.shape[0] != N: + raise ValueError(f"Controls have {controls.shape[0]} rows but y has {N} obs.") + if not np.all(np.isfinite(controls)): + raise ValueError( + "controls contains non-finite values (NaN or Inf). " + "Please remove or impute missing values before calling " + "wild_cluster_bootstrap()." + ) + + if np.issubdtype(cluster_ids.dtype, np.floating) and not np.all(np.isfinite(cluster_ids)): + raise ValueError( + "cluster_ids contains non-finite values (NaN or Inf). " + "Cluster identifiers must be valid for all observations." + ) + if weight_type not in _VALID_WEIGHT_TYPES: + raise ValueError( + f"Unknown weight_type '{weight_type}'. Must be one of: {_VALID_WEIGHT_TYPES}" + ) + if ( + isinstance(alpha, bool) + or not isinstance(alpha, (int, float, np.integer, np.floating)) + or not np.isfinite(alpha) + or not (0.0 < alpha < 1.0) + ): + raise ValueError(f"alpha must be a scalar in (0, 1), got {alpha!r}.") + if ( + isinstance(n_bootstrap, bool) + or not isinstance(n_bootstrap, (int, np.integer)) + or n_bootstrap < 2 + ): + # Round-4 review: one draw cannot estimate a bootstrap dispersion + # or support test inversion (matches the estimator's 0-or->=2 rule). + raise ValueError(f"n_bootstrap must be an integer >= 2, got {n_bootstrap!r}.") + + # Drop non-finite y WITH a warning (campaign finding: silent drops). + finite_mask = np.isfinite(y) + n_dropped = int((~finite_mask).sum()) + if n_dropped: + warnings.warn( + f"wild_cluster_bootstrap: dropped {n_dropped} observation(s) " + f"with non-finite y before estimation.", + UserWarning, + stacklevel=2, + ) + y = y[finite_mask] + treatment = treatment[finite_mask] + cluster_ids = cluster_ids[finite_mask] + if controls is not None: + controls = controls[finite_mask] + N = len(y) + if N == 0: + raise ValueError("All observations have non-finite y values.") + if treatment.sum() == 0: + raise ValueError("After dropping non-finite y, no treated observations remain.") + if treatment.sum() == N: + raise ValueError("After dropping non-finite y, no control observations remain.") + + unique_clusters = np.unique(cluster_ids) + G = len(unique_clusters) + if G < 2: + raise ValueError(f"Need at least 2 clusters for wild cluster bootstrap, got {G}.") + + # Design matrix: [intercept, treatment, controls...]; treatment at 1. + parts = [np.ones(N, dtype=np.float64), treatment] + if controls is not None: + parts.extend(controls[:, j] for j in range(controls.shape[1])) + X = np.column_stack(parts) + + # Exactly-identified degenerate design guard (fires BEFORE delegating; + # the shared helper stays byte-identical so its R-parity goldens cannot + # move). With cluster-invariant treatment and G small enough that OLS + # fits every cluster-arm mean exactly, all cluster scores are ~0 and + # BLAS roundoff yields a tiny-positive SE instead of 0 - pre-fix this + # reported t ~ 5e15 with p = 0.25 (below the attainable G=2 floor of + # 0.5). Point retained; inference NaN (house fail-closed pattern). + # Rank-aware fit through the shared solver (round-4 review: lstsq + # returned a finite minimum-norm treatment coefficient when a control + # duplicated the treatment column, so an unidentified ATT was reported + # with finite bootstrap inference). + beta_hat, resid, _ = solve_ols(X, y) + if not np.isfinite(beta_hat[1]): + raise ValueError( + "The treatment coefficient is not identified: the design is " + "rank-deficient and the shared solver dropped the treatment " + "column (e.g. a control collinear with treatment). Remove the " + "collinear control(s) before bootstrapping." + ) + att_point = float(beta_hat[1]) + scores = np.array([X[cluster_ids == cl].T @ resid[cluster_ids == cl] for cl in unique_clusters]) + score_scale = float(np.abs(X.T @ np.abs(resid)).max()) + if score_scale > 0 and float(np.abs(scores).max()) <= 1e-10 * score_scale: + warnings.warn( + "wild_cluster_bootstrap: the cluster-level scores are exactly " + "zero (exactly-identified design, e.g. cluster-invariant " + "treatment with as many parameters as cluster-arm means): the " + "cluster-robust variance is not identified. The point estimate " + "is retained; SE, p-value, and CI are NaN.", + UserWarning, + stacklevel=2, + ) + return WildClusterBootstrapResult( + att=att_point, + se=np.nan, + t_stat_original=np.nan, + p_value=np.nan, + ci_lower=np.nan, + ci_upper=np.nan, + n_clusters=G, + n_bootstrap=0, + weight_type=weight_type, + alpha=alpha, + bootstrap_distribution=None, + n_dropped=n_dropped, + ) + + house = wild_bootstrap_se( + X, + y, + resid, + cluster_ids, + 1, + n_bootstrap=n_bootstrap, + weight_type=weight_type, + alpha=alpha, + seed=seed, + return_distribution=True, + ) + + return WildClusterBootstrapResult( + att=att_point, + se=float(house.se), + t_stat_original=float(house.t_stat_original), + p_value=float(house.p_value), + ci_lower=float(house.ci_lower), + ci_upper=float(house.ci_upper), + n_clusters=int(house.n_clusters), + n_bootstrap=int(house.n_bootstrap), + weight_type=weight_type, + alpha=alpha, + bootstrap_distribution=house.bootstrap_distribution, + n_dropped=n_dropped, + ) diff --git a/diff_diff/results.py b/diff_diff/results.py index 20a7993e..202d2c3b 100644 --- a/diff_diff/results.py +++ b/diff_diff/results.py @@ -71,6 +71,8 @@ def _format_vcov_label( return "HC1 heteroskedasticity-robust" if vcov_type == "hc2": return "HC2 leverage-corrected" + if vcov_type == "hc3": + return "HC3 jackknife-style leverage-corrected" if vcov_type == "hc2_bm": if cluster_name: suffix = f", G={n_clusters}" if n_clusters else "" diff --git a/diff_diff/spillover.py b/diff_diff/spillover.py index 506d4c65..5cf6a671 100644 --- a/diff_diff/spillover.py +++ b/diff_diff/spillover.py @@ -1760,6 +1760,29 @@ def __init__( f"rank_deficient_action must be 'warn', 'error', or 'silent', " f"got '{rank_deficient_action}'" ) + # Never-supported leverage families fail at construction (not fit): + # hc2/hc2_bm need per-coefficient Bell-McCaffrey/CR2 dof the stage-2 + # path does not supply; hc3 is simply not implemented for the + # two-stage spillover variance. Other values keep fit-time checks. + if vcov_type in ("hc2", "hc2_bm"): + raise NotImplementedError( + f"SpilloverDiD does not yet support vcov_type='{vcov_type}'. " + "The current stage-2 inference uses a generic residual df " + "(n - effective_rank) for t-distribution lookups, but " + "hc2 / hc2_bm require per-coefficient Bell-McCaffrey / CR2 " + "degrees of freedom for correct p-values and CIs. Routing " + "stage 2 through LinearRegression (which supplies the " + "per-coefficient DOF metadata) is queued as a follow-up " + "extension. Use vcov_type='hc1' or 'conley', or " + "leave default; combine with cluster= for CR1." + ) + if vcov_type == "hc3": + raise NotImplementedError( + "SpilloverDiD does not support vcov_type='hc3': hc3 is not " + "implemented for the two-stage spillover variance. Use " + "vcov_type='hc1' or 'conley', or leave default; combine " + "with cluster= for CR1." + ) self.rings = rings self.d_bar = d_bar self.vcov_type = vcov_type @@ -2259,6 +2282,15 @@ def fit( "path mirroring TwoStageDiD._fit_untreated_model is queued as " "a follow-up extension. See DEFERRED.md." ) + if self.vcov_type == "hc3": + # Defense-in-depth mirror of the constructor guard (set_params + # bypass); hc3 has its own reason distinct from hc2/hc2_bm. + raise NotImplementedError( + "SpilloverDiD does not support vcov_type='hc3': hc3 is not " + "implemented for the two-stage spillover variance. Use " + "vcov_type='hc1' or 'conley', or leave default; combine " + "with cluster= for CR1." + ) if self.vcov_type in ("hc2", "hc2_bm"): raise NotImplementedError( f"SpilloverDiD does not yet support vcov_type='{self.vcov_type}'. " diff --git a/diff_diff/twfe.py b/diff_diff/twfe.py index c88632be..a7e365e3 100644 --- a/diff_diff/twfe.py +++ b/diff_diff/twfe.py @@ -54,11 +54,11 @@ class TwoWayFixedEffects(DifferenceInDifferences): DifferenceInDifferences where cluster=None means no clustering. **Exception (one-way analytical):** when - ``vcov_type in {"classical", "hc2"}`` is explicit AND + ``vcov_type in {"classical", "hc2", "hc3"}`` is explicit AND ``inference="analytical"``, the unit auto-cluster is dropped because these families are by construction one-way only and the - validator rejects ``cluster_ids + classical`` / ``cluster_ids + - hc2``. The user's explicit one-way choice wins over the TWFE + validator rejects ``cluster_ids`` with these one-way families. The + user's explicit one-way choice wins over the TWFE default. Under ``inference="wild_bootstrap"`` the auto-cluster is preserved regardless of ``vcov_type`` (the bootstrap uses the cluster structure to resample residuals). On ``hc2_bm`` the @@ -117,7 +117,7 @@ class TwoWayFixedEffects(DifferenceInDifferences): Because TWFE's within-transformation preserves coefficients but not the hat matrix, HC2 leverage and CR2 Bell-McCaffrey corrections on the demeaned design would produce wrong small-sample SEs. When - ``vcov_type in {"hc2","hc2_bm"}``, TWFE bypasses the within-transform + ``vcov_type in {"hc2","hc2_bm","hc3"}``, TWFE bypasses the within-transform and builds the full-dummy design ``[intercept, treated×post, covariates, unit_dummies, time_dummies]`` directly, so the leverage correction and BM DOF compute on the full FE projection. Under this @@ -126,8 +126,8 @@ class TwoWayFixedEffects(DifferenceInDifferences): full-dummy fit rather than the within-transformed reduced fit; the ATT coefficient, its SE, and analytical inference are unchanged. Auto-cluster-at-unit is preserved on ``hc2_bm`` (routes to CR2-BM at - unit) and on ``hc2`` + ``wild_bootstrap``; dropped on explicit ``hc2`` - + ``analytical`` to match the one-way contract. **This wording applies + unit) and on ``hc2``/``hc3`` + ``wild_bootstrap``; dropped on explicit + ``hc2``/``hc3`` + ``analytical`` to match the one-way contract. **This wording applies to the non-survey analytical path**: under ``survey_design=`` with no explicit ``cluster=``, TWFE intentionally keeps the documented implicit-PSU path (auto-cluster is NOT injected into the survey PSU @@ -330,7 +330,7 @@ def fit( # type: ignore[override] # with explicit unit + time dummies routes through ``solve_ols``'s # full-design hat matrix. HC1/CR1 paths remain on the demeaned # design (no leverage term). - use_full_dummy = self.vcov_type in ("hc2", "hc2_bm") + use_full_dummy = self.vcov_type in ("hc2", "hc2_bm", "hc3") # Phase 2 panel block-decomposed Conley (matches R conleyreg). # FWL composability: the within-transformed scores S = X_demeaned * @@ -443,7 +443,7 @@ def fit( # type: ignore[override] if self.cluster is not None: cluster_var: Optional[str] = self.cluster elif ( - self.vcov_type in ("classical", "hc2") + self.vcov_type in ("classical", "hc2", "hc3") and self._vcov_type_explicit and self.inference == "analytical" ): @@ -508,7 +508,7 @@ def fit( # type: ignore[override] f"~{_design_entries * 8 / 1e9:.2f} GB). For panels with " f"many units/periods, consider vcov_type='hc1' (within-" "transform path; no leverage term, lower memory) unless " - "small-sample HC2/HC2-BM inference is required.", + "leverage-corrected HC2/HC2-BM/HC3 inference is required.", UserWarning, stacklevel=2, ) @@ -1116,7 +1116,7 @@ def _fit_event_study( # rule): the auto-cluster is never injected as a survey PSU; # only user-explicit cluster= becomes one. cluster_override = None - elif self.vcov_type in ("classical", "hc2") and self._vcov_type_explicit: + elif self.vcov_type in ("classical", "hc2", "hc3") and self._vcov_type_explicit: # The explicit one-way analytical exception (mirrors the static # cluster_var block; inference is always analytical here - wild # raised above). @@ -1134,7 +1134,7 @@ def _fit_event_study( assert unit_resolved is not None absorb_arg: Optional[List[str]] = [unit_resolved] if ( - self.vcov_type in ("hc2", "hc2_bm") + self.vcov_type in ("hc2", "hc2_bm", "hc3") and unit_resolved in data.columns and time in data.columns ): @@ -1159,11 +1159,11 @@ def _fit_event_study( f"{_design_cols} full-dummy design " f"(~{_design_entries / 1e6:.1f}M float64 entries, " f"~{_design_entries * 8 / 1e9:.2f} GB) for the " - "leverage-corrected HC2/HC2-BM path. For panels with " + "leverage-corrected HC2/HC2-BM/HC3 path. For panels with " "many units, consider vcov_type='hc1' (absorbed " "within-transform path; no leverage term, lower " - "memory) unless small-sample HC2/HC2-BM inference is " - "required.", + "memory) unless leverage-corrected HC2/HC2-BM/HC3 inference " + "is required.", UserWarning, stacklevel=3, ) diff --git a/docs/api/_autosummary/diff_diff.LWDiD.rst b/docs/api/_autosummary/diff_diff.LWDiD.rst new file mode 100644 index 00000000..1028473c --- /dev/null +++ b/docs/api/_autosummary/diff_diff.LWDiD.rst @@ -0,0 +1,22 @@ +diff\_diff.LWDiD +================ + +.. currentmodule:: diff_diff + +.. autoclass:: LWDiD + :no-members: + + + .. rubric:: Methods + + .. autosummary:: + + ~LWDiD.__init__ + ~LWDiD.fit + ~LWDiD.get_params + ~LWDiD.get_transformation_diagnostics + ~LWDiD.set_params + + + + diff --git a/docs/api/_autosummary/diff_diff.lwdid_results.LWDiDResults.rst b/docs/api/_autosummary/diff_diff.lwdid_results.LWDiDResults.rst new file mode 100644 index 00000000..1cf221ea --- /dev/null +++ b/docs/api/_autosummary/diff_diff.lwdid_results.LWDiDResults.rst @@ -0,0 +1,73 @@ +diff\_diff.lwdid\_results.LWDiDResults +====================================== + +.. currentmodule:: diff_diff.lwdid_results + +.. autoclass:: LWDiDResults + :no-members: + + + .. rubric:: Methods + + .. autosummary:: + + ~LWDiDResults.__init__ + ~LWDiDResults.aggregate + ~LWDiDResults.print_summary + ~LWDiDResults.randomization_test + ~LWDiDResults.summary + ~LWDiDResults.to_csv + ~LWDiDResults.to_dataframe + ~LWDiDResults.to_dict + ~LWDiDResults.wild_cluster_bootstrap + + + + + .. rubric:: Attributes + + .. autosummary:: + + ~LWDiDResults.bootstrap_pvalue + ~LWDiDResults.bse + ~LWDiDResults.cband_crit_value + ~LWDiDResults.cband_method + ~LWDiDResults.cband_n_bootstrap + ~LWDiDResults.att_tau_omega_complete_case + ~LWDiDResults.ci + ~LWDiDResults.cluster_name + ~LWDiDResults.cohort_effects + ~LWDiDResults.cohort_time_effects + ~LWDiDResults.control_group + ~LWDiDResults.df_inference + ~LWDiDResults.event_study_df + ~LWDiDResults.event_study_effects + ~LWDiDResults.event_study_vcov + ~LWDiDResults.event_study_vcov_index + ~LWDiDResults.inference_basis + ~LWDiDResults.is_staggered + ~LWDiDResults.n_bootstrap + ~LWDiDResults.n_clusters + ~LWDiDResults.n_composite_controls_dropped + ~LWDiDResults.n_composite_treated_dropped + ~LWDiDResults.params + ~LWDiDResults.pscore_trim + ~LWDiDResults.psm_config + ~LWDiDResults.pvalue + ~LWDiDResults.seed + ~LWDiDResults.ri_pvalue + ~LWDiDResults.vcov + ~LWDiDResults.att + ~LWDiDResults.se + ~LWDiDResults.t_stat + ~LWDiDResults.p_value + ~LWDiDResults.conf_int + ~LWDiDResults.n_obs + ~LWDiDResults.n_treated + ~LWDiDResults.n_control + ~LWDiDResults.rolling + ~LWDiDResults.estimation_method + ~LWDiDResults.vcov_type + ~LWDiDResults.alpha + ~LWDiDResults.reference_periods + diff --git a/docs/api/index.rst b/docs/api/index.rst index 807ec873..26c61aef 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -34,6 +34,7 @@ regression discontinuity, and the Goodman-Bacon decomposition diagnostic: diff_diff.LPDiD diff_diff.ChangesInChanges diff_diff.QDiD + diff_diff.LWDiD diff_diff.BaconDecomposition diff_diff.StaggeredTripleDifference diff_diff.RegressionDiscontinuity @@ -77,6 +78,7 @@ Result containers returned by estimators: diff_diff.wooldridge_results.WooldridgeDiDResults diff_diff.lpdid_results.LPDiDResults diff_diff.changes_in_changes_results.ChangesInChangesResults + diff_diff.lwdid_results.LWDiDResults diff_diff.Comparison2x2 diff_diff.StaggeredTripleDiffResults diff_diff.TWFEWeightsResult @@ -354,6 +356,7 @@ Estimators wooldridge_etwfe lpdid changes_in_changes + lwdid bacon Infrastructure diff --git a/docs/api/lwdid.rst b/docs/api/lwdid.rst new file mode 100644 index 00000000..1ff174a1 --- /dev/null +++ b/docs/api/lwdid.rst @@ -0,0 +1,603 @@ +LWDiD — Lee & Wooldridge Rolling Transformation DiD +==================================================== + +A simple transformation approach to Difference-in-Differences estimation +that converts panel data into cross-sectional regressions (Lee & Wooldridge +2025, 2026). + +The key insight from the Lee & Wooldridge papers is that, under parallel +trends and no anticipation, a unit-specific time-series transformation of +the outcome eliminates the need for two-way fixed effects entirely. For +each unit *i* with treatment onset at period *S*, Procedure 2.1 (LW 2026) +computes the pre-treatment mean: + +.. math:: + + \bar{Y}_{i,\text{pre}} = \frac{1}{S-1} \sum_{t=1}^{S-1} Y_{it} + +and forms the transformed outcome: + +.. math:: + + \dot{Y}_{it} = Y_{it} - \bar{Y}_{i,\text{pre}}, \quad t = S, \ldots, T + \qquad \text{(Equation 2.12, LW 2026)} + +Under Assumption CPTC (Conditional Parallel Trends, Common Timing; +Equation 2.10, LW 2025), this transformation +removes unit-specific fixed effects, and the ATT is identified as the +coefficient on the treatment indicator in a cross-sectional regression of +:math:`\dot{Y}_{it}` on :math:`D_i` and covariates. Because the panel +problem is reduced to a cross section, *any* treatment effect estimator — +regression adjustment (RA), inverse probability weighting (IPW), doubly +robust IPWRA, or propensity-score matching — can be applied without +negative weighting, heterogeneity bias, or "bad comparisons" between +already-treated cohorts. + +A second contribution (LW 2026) demonstrates that this representation +enables *exact* small-sample inference: under homoskedastic normality of +the cross-sectional error, the t-statistic follows an exact +:math:`\mathcal{T}_{N-K-2}` distribution — valid even with a single +treated unit (:math:`N_1 = 1`). When :math:`T_0` or :math:`T_1` is large, +the central limit theorem across time justifies the normality assumption +without requiring a large cross section. + +.. note:: + + **Why rolling transformation works.** The parallel trends assumption + (Equation 2.15, LW 2026) implies that :math:`\Delta\bar{Y}_i(0)` + — the difference between post-treatment and pre-treatment means of + control potential outcomes — is mean-independent of the treatment + indicator :math:`D_i`. This is precisely the unconfoundedness condition + needed for cross-sectional treatment effect estimation. The + transformation eliminates *both* unit-specific levels (via demeaning) + and unit-specific linear trends (via detrending), weakening the + standard parallel trends assumption to one that allows heterogeneous + pre-intervention dynamics. + +.. module:: diff_diff.lwdid + +Methodology +----------- + +**Procedure 2.1 — Unit-Specific Demeaning (LW 2026, Section 2)** + +For common timing with intervention at period *S*: + +1. Compute the pre-treatment mean for each unit: + :math:`\bar{Y}_{i,\text{pre}} = \frac{1}{S-1}\sum_{t=1}^{S-1} Y_{it}` + +2. Obtain the transformed outcome (out-of-sample residuals): + + .. math:: + + \dot{Y}_{it} = Y_{it} - \bar{Y}_{i,\text{pre}}, \quad t = S, \ldots, T + +3. Estimate the ATT from the cross-sectional regression (Equation 2.13, LW 2026): + + .. math:: + + \dot{Y}_{it} \text{ on } 1,\; D_i, \quad i = 1, \ldots, N + +The coefficient on :math:`D_i` identifies the ATT for period *t*. + +**Procedure 3.1 — Unit-Specific Detrending (LW 2025, Section 5; LW 2026, Section 3)** + +When parallel trends may fail but unit-specific *linear* trends capture +the pre-intervention dynamics (Assumption CHT, LW 2025): + +1. For each unit *i*, regress on a constant and time over pre-treatment + periods: + + .. math:: + + Y_{it} \text{ on } 1,\; t, \quad t = 1, \ldots, S-1 + \qquad \text{(Equation 3.1, LW 2026)} + + obtaining fitted values :math:`\hat{A}_i + \hat{B}_i \cdot t`. + +2. Compute the detrended outcome: + + .. math:: + + \ddot{Y}_{it} = Y_{it} - \hat{A}_i - \hat{B}_i \cdot t, \quad t = S, \ldots, T + \qquad \text{(Equation 3.2, LW 2026)} + +3. Estimate the ATT from: + + .. math:: + + \ddot{Y}_{it} \text{ on } 1,\; D_i, \quad i = 1, \ldots, N + \qquad \text{(Equation 3.4, LW 2026)} + +Detrending removes unit-specific intercepts :math:`\alpha_i` *and* linear +trends :math:`\beta_i t`, thus relaxing the parallel trends assumption to +allow differential pre-intervention growth rates across units (Procedure +5.1, LW 2025). This is the key advantage over Callaway & Sant'Anna (2021), +who do not accommodate heterogeneous trends. + +**Procedure 4.1 — Staggered Interventions (LW 2025, Section 4)** + +For staggered adoption with cohort *g* (first treatment period) and +calendar time *r*: + +1. Compute the cohort-specific transformed outcome: + + .. math:: + + \dot{Y}_{irg} = Y_{ir} - \frac{1}{g-1}\sum_{s=1}^{g-1} Y_{is} + \equiv Y_{ir} - \bar{Y}_{i,\text{pre}(g)} + \qquad \text{(Equation 4.11, LW 2025)} + +2. Select the control group: units not yet treated by period *r*, + i.e., cohorts :math:`\{r+1, \ldots, T, \infty\}`. + +3. Apply any TE estimator (RA, IPW, IPWRA, matching) to the cross section + :math:`\{(\dot{Y}_{irg}, D_{ig}, \mathbf{X}_i)\}` restricted to the + treated cohort *g* plus control units. + +Under Assumptions CNAS (conditional no anticipation, Equation 4.4) and +CPTS (conditional parallel trends, Equation 4.6), the cohort assignment +is unconfounded with respect to the transformed outcome (Theorem 4.1). + +**Regression Adjustment with Interactions (Equation 3.3, LW 2025)** + +When both :math:`N_0` and :math:`N_1` are sufficiently large, full +regression adjustment includes covariate interactions: + +.. math:: + + \dot{Y}_{ir} = \beta_0 + \beta_1 D_i + \beta_2' \mathbf{X}_i + + \beta_3' D_i(\mathbf{X}_i - \bar{\mathbf{X}}_1) + u_i + +where :math:`\bar{\mathbf{X}}_1 = N_1^{-1}\sum_{i} D_i \mathbf{X}_i` is +the mean of covariates over treated units. The ATT is :math:`\hat{\beta}_1`. +This is equivalent to separate regressions for treated and control groups +(Equation 3.3, LW 2025). + +Key Assumptions +--------------- + +.. important:: + + The LWDiD estimator requires the following assumptions for identification: + + **Assumption CPTC — Conditional Parallel Trends, Common Timing** + (Equation 2.10, LW 2025): + + .. math:: + + E[Y_{it}(0) - Y_{i1}(0) \mid D_i, \mathbf{X}_i] + = E[Y_{it}(0) - Y_{i1}(0) \mid \mathbf{X}_i], \quad t = 2, \ldots, T + + The *trend* in control potential outcomes is independent of treatment + assignment conditional on covariates. Note this is weaker than + unconditional parallel trends — assignment can be correlated with + *levels* :math:`Y_{i1}(0)`, but not with *trends*. + + **Assumption NAC — No Anticipation, Common Timing** (Equation 2.7, LW 2025): + + .. math:: + + E[Y_{it}(1) - Y_{it}(0) \mid D_i = 1] = 0, \quad t = 1, \ldots, S-1 + + Treatment effects are zero on average before the intervention. + + **Assumption CPTS — Conditional PT, Staggered** (Equation 4.6, LW 2025): + + .. math:: + + E[Y_t(\infty) - Y_1(\infty) \mid \mathbf{D}, \mathbf{X}] + = E[Y_t(\infty) - Y_1(\infty) \mid \mathbf{X}], \quad t = 2, \ldots, T + + Trends in the never-treated state are independent of the full vector + of cohort assignments, enabling use of not-yet-treated units as controls. + + **Conditional Heterogeneous Trends** (Assumption CHT, Equation 5.3, + LW 2025): When using ``detrend``, the parallel trends assumption is + relaxed to allow unit-specific linear trends + :math:`\eta_g \cdot t` that vary by cohort. Detrending removes these + heterogeneous trends, restoring unconfoundedness. + +Small-Sample Inference +---------------------- + +A distinctive feature of the LW approach (LW 2026, Section 2) is the +availability of *exact* inference. Under the classical linear model +assumptions on the cross-sectional regression: + +.. math:: + + U_i \mid D_i \sim \text{Normal}(0, \sigma_U^2) + \qquad \text{(Equation 2.9, LW 2026)} + +the t-statistic follows an exact Student-t distribution: + +.. math:: + + \frac{\hat{\tau}_{DD} - \tau}{\text{se}(\hat{\tau}_{DD})} + \sim \mathcal{T}_{N-2} + \qquad \text{(Equation 2.10, LW 2026)} + +This holds even with :math:`N_1 = 1` (single treated unit), where the +t-statistic is interpretable as a *studentized residual* — testing whether +the treated unit is an "outlier" relative to the controls (LW 2026, +Section 2.1). + +When :math:`N` is not too small, the HC3 heteroskedasticity-robust +standard error (Davidson & MacKinnon, 1993) provides reliable inference +without the homoskedasticity assumption, as shown by Simonsohn (2021). + +**Randomization inference** is also supported: under the sharp null of +zero treatment effects, permutation of :math:`D_i` yields Monte Carlo +p-values without requiring normality (LW 2026, the small-sample +inference paper). Validity is conditional on the assignment mechanism the +permutation encodes — complete randomization of the treatment labels +(the treated count is held fixed); the implementation follows the +authors' package convention (inclusive Phipson-Smyth counting; see the +methodology registry's RI Note). + +**HC3 caveat** — HC3 requires the leverage of every observation to be +bounded away from one; a perfectly-leveraged design (e.g. a single +treated unit) has no defined HC3 variance and fails closed with a +warning and NaN inference. Use classical exact inference there. + +**PSM inference** — ``estimation_method='psm'`` reports the matched ATT +point estimate with NaN inference: no valid matching variance estimator +is currently implemented (the naive matched-pairs formula ignores +matched-control reuse and first-stage matching uncertainty; an +Abadie-Imbens variance is tracked in ``DEFERRED.md``). The contract is +enforced on every route: PSM requires ``covariates`` (there is no +propensity score without them), rejects ``n_bootstrap > 0`` in BOTH +common-timing and staggered designs (the standard bootstrap is invalid +for nearest-neighbor matching estimators — Abadie & Imbens 2008), and a +propensity-model failure falls back to a regression-adjustment POINT with +NaN inference rather than finite OLS standard errors. Use +``estimation_method='dr'`` for valid inference. + +LWDiD +------ + +Main estimator class. + +.. autoclass:: diff_diff.LWDiD + :no-index: + :members: + :undoc-members: + :show-inheritance: + :inherited-members: + + .. rubric:: Methods + + .. autosummary:: + + ~LWDiD.fit + ~LWDiD.get_params + ~LWDiD.set_params + +LWDiDResults +------------ + +Results container returned by :meth:`~diff_diff.LWDiD.fit`. + +.. autoclass:: diff_diff.lwdid_results.LWDiDResults + :no-index: + :members: + :undoc-members: + :show-inheritance: + + .. rubric:: Methods + + .. autosummary:: + + ~LWDiDResults.summary + ~LWDiDResults.print_summary + ~LWDiDResults.to_dataframe + ~LWDiDResults.to_dict + +Input Contract +-------------- + +:meth:`~diff_diff.LWDiD.fit` validates the treatment design before any +transformation is applied. Eight requirements are enforced: + +- **Absorbing treatment** — within each unit the ``treatment`` indicator + must be non-decreasing over time: once a unit switches from 0 to 1 it + must remain treated. Units that revert to 0 raise ``ValueError``. +- **Common timing** — when ``first_treat`` is not supplied, all treated + units must first switch on in the same period. Heterogeneous onsets + are rejected with a ``ValueError`` pointing to the staggered interface + (pass ``first_treat``). +- **Staggered consistency** — when ``first_treat`` is supplied, the + ``treatment`` indicator must satisfy :math:`D_{it} = 1[t \ge g_i]` over + each unit's OBSERVED rows, where :math:`g_i` is the unit's + first-treatment period; the row at :math:`t = g_i` itself may be + unobserved (unbalanced panels with a missing onset row are accepted). + Units that are never treated must have no treated rows. +- **Never-treated encodings** — ``first_treat`` coded ``0``, ``NaN``/ + ``NaT``, or ``np.inf`` means never-treated; ``inf`` and finite cohorts + BEYOND the last observed period are recoded to never-treated with a + warning (beyond-window units never switch on inside the sample). + Negative cohorts raise. NUMERIC cohorts strictly between observed + periods are rejected; datetime/Period cohorts map to the next observed + period — a dtype-dependent contract documented in the methodology + registry. +- **Variance configuration** — ``estimation_method='reg'`` accepts + ``vcov_type`` in ``{'classical', 'hc1', 'hc2', 'hc3'}``; ``'ipw'``/ + ``'dr'``/``'psm'`` accept ``'hc1'`` only (the influence-function / + matching variance is always used on those paths); ``cluster=`` composes + only with ``'hc1'`` (CR1) and is rejected for ``'psm'``. +- **Never-treated units under not-yet-treated control** — when + ``first_treat`` is supplied and ``control_group='not_yet_treated'``, + at least one never-treated unit (``first_treat`` coded NaN or 0) must + be present. A panel in which every unit is eventually treated raises + ``ValueError`` rather than silently truncating the estimation sample. +- **Unit-constant covariates** — ``covariates`` (and a non-unit + ``cluster=`` column) must be constant within each unit on BOTH timing + paths; time-varying columns raise ``ValueError`` (LWDiD collapses the + panel to one row per unit, so a time-varying value would make the + estimate depend on row order). +- **Distinct, non-reserved column names** — the core role columns + (outcome/unit/time/treatment/``first_treat``) must be pairwise + distinct, covariates may not repeat a core role, and no role column + may use an LWDiD-internal working name (``_treat``, ``_ydot``, + ``_ydot_avg``, ``_ever_treated``, ``_boot_unit``, ``_lwdid_time_pos``, + ``_lwdid_cohort_pos``, ``_lwdid_season``) — a collision would silently overwrite the + internal column (e.g. ``cluster='_treat'`` previously reported the + cluster labels' coefficient as the ATT). ``cluster=`` equal to the + unit column remains supported. + +.. note:: + + **Bootstrap scope and reproducibility.** In common-timing fits + ``n_bootstrap`` activates a unit-resampling (or cluster-resampling, + under ``cluster=``) bootstrap for the overall ATT; the headline + se/p-value/CI then come from the bootstrap while ``params``/``vcov`` + remain the analytical regression quantities, recorded via + ``inference_basis`` (``'unit_bootstrap'``/``'cluster_bootstrap'``) and + rendered by ``summary()``. The per-replicate RNG streams are + ``SeedSequence``-spawned identically for every ``n_jobs``, so a seeded + fit reproduces exactly across serial and parallel execution. In STAGGERED fits + ``n_bootstrap`` governs the event-study multiplier bootstrap only + (sup-t simultaneous bands); the overall and cohort aggregates keep + analytical influence-function inference, with the per-surface basis + recorded on the results object (``cband_method``, + ``cband_n_bootstrap``, ``inference_basis``). Event cells whose + multiplier draws are degenerate fail closed (point retained, NaN + inference) rather than silently reverting to analytical standard + errors. The common-timing event-study surface covers post periods + only; pre-treatment placebo cells are produced by the staggered path + (pass ``first_treat=``, which for a single cohort matches the + common-timing regression on single-post-period panels). + +.. note:: + + :meth:`~diff_diff.lwdid_results.LWDiDResults.to_dict` returns only + JSON-native types: numpy scalars and arrays are converted to Python + ints/floats/bools and lists, and datetime-like labels (Timestamp, + Period) become strings, so ``json.dumps(result.to_dict())`` works + directly. + +Example Usage +------------- + +**Basic demeaning with regression adjustment (Procedure 2.1):** + +.. code-block:: python + + import pandas as pd + from diff_diff import LWDiD, generate_staggered_data + + # Generate staggered panel data; the 'treated' column is the binary + # indicator D_it = 1[period >= first_treat] (0 for never-treated units) + data = generate_staggered_data(n_units=200, n_periods=10, + cohort_periods=[4, 7], seed=42) + + # Procedure 2.1: demean + reg estimates the ATT via cross-sectional OLS + # on the transformed outcome Y_dot = Y_post - Y_bar_pre + lw = LWDiD(rolling="demean", estimation_method="reg", vcov_type="hc1") + results = lw.fit(data, outcome="outcome", unit="unit", + time="period", treatment="treated", + first_treat="first_treat") + results.print_summary() + +**Doubly-robust IPWRA estimation (Procedure 3.1, Step 2):** + +.. code-block:: python + + # DR (IPWRA) combines propensity score weighting with regression adjustment + # on the transformed outcome — doubly robust as in Wooldridge (2007). + # Cluster-robust inference activates via the constructor's cluster= parameter. + data["state"] = data["unit"] % 40 # cluster identifier + lw_dr = LWDiD(rolling="demean", estimation_method="dr", cluster="state") + results_dr = lw_dr.fit(data, outcome="outcome", unit="unit", + time="period", treatment="treated", + first_treat="first_treat") + print(f"ATT: {results_dr.att:.4f} (SE={results_dr.se:.4f})") + +**Staggered adoption with detrending (Procedure 4.1 + 5.1):** + +.. code-block:: python + + # Detrending removes unit-specific linear trends before estimation, + # relaxing parallel trends to allow heterogeneous pre-intervention dynamics + lw_stag = LWDiD(rolling="detrend", control_group="never_treated") + results_stag = lw_stag.fit(data, outcome="outcome", unit="unit", + time="period", treatment="treated", + first_treat="first_treat") + # Cohort-specific ATT(g) estimates (Equations 7.2/7.10, LW 2026). + # Aggregation convention: each cohort's estimable cells weight by + # their contributing treated mass (cell-mass, matching the WATT(r) + # axis); this equals the eq. 7.10 unit-average estimand on balanced + # never-treated designs and deviates on unbalanced panels (see the + # methodology registry's within-cohort aggregation Note). + df_cohorts = results_stag.to_dataframe() + print(df_cohorts) + +**Robustness check — demean vs detrend (informal pre-test for trend +sensitivity):** + +.. code-block:: python + + # Comparing demean vs detrend provides a specification robustness check. + # If results differ substantially, it suggests unit-specific trends matter + # (see LW 2025, Section 6 — Walmart application, Figure 1 panels b vs c) + for transform in ("demean", "detrend"): + lw_check = LWDiD(rolling=transform, estimation_method="dr", vcov_type="hc1") + res = lw_check.fit(data, outcome="outcome", unit="unit", + time="period", treatment="treated", + first_treat="first_treat") + print(f"{transform}: ATT={res.att:.4f} (SE={res.se:.4f})") + +Wild cluster bootstrap +---------------------- + +``diff_diff.lwdid_wild_bootstrap.wild_cluster_bootstrap(y, treatment, +cluster_ids, controls=None, *, n_bootstrap=999, weight_type='rademacher', +alpha=0.05, seed=None)`` provides few-cluster inference on a collapsed +cross-section. It delegates to the house Wild Cluster Restricted engine +(:func:`diff_diff.wild_bootstrap_se`, matched to R's +``fwildclusterboot::boottest``): the null is imposed by dropping the +treatment column while keeping the controls, the confidence interval is +obtained by test inversion, and Rademacher weights are fully enumerated +automatically when :math:`2^G \le` ``n_bootstrap``. The result carries +``att``, the analytical CR1 ``se``, ``t_stat_original``, ``p_value`` +(strict-exceedance house convention), the test-inversion +``ci_lower``/``ci_upper``, ``n_clusters``, ``n_bootstrap``, +``weight_type``, ``alpha``, the finite-filtered ``bootstrap_distribution`` +(``None`` when the degenerate-design guard fires), and ``n_dropped`` +(non-finite outcome rows dropped with a warning). Exactly-identified +designs (cluster-invariant treatment with zero cluster scores) fail +closed: the point estimate is retained with NaN inference. + +The RESULT-LEVEL methods ``LWDiDResults.wild_cluster_bootstrap()`` and +``LWDiDResults.randomization_test()`` take no data arguments: they REPLAY +the fit-time collapsed cross-section and the exact fitted RA design +(including the treatment-centered covariate interactions; randomization +draws recompute the treated covariate mean per assignment), and assert +the replayed statistic equals ``.att`` before caching a p-value — so +``bootstrap_pvalue``/``ri_pvalue`` always describe the fitted estimand. +They are defined for common-timing ``estimation_method='reg'`` fits +(``wild_cluster_bootstrap`` additionally requires a ``cluster=`` fit); +use the standalone module functions above for generic arrays. + + +Empirical Applications +---------------------- + +The Lee & Wooldridge papers validate the methodology with two empirical +studies: + +- **California Proposition 99** (LW 2026, Section 6): With a single treated + state (:math:`N_1 = 1`) and 38 control states, Procedure 3.1 + (unit-specific detrending) achieves an excellent pre-treatment fit and + yields a per-period treatment trajectory that grows over time — from + :math:`\hat{\tau}_{1989} = -0.043` (SE = 0.059) to + :math:`\hat{\tau}_{2000} = -0.403` (SE = 0.152). The exact-inference + p-value (0.021) is valid under the conditional-normality and + homoskedasticity assumptions — it tests the treatment-effect null, not + those assumptions themselves; randomization inference (below) is a + robustness check that does not require normality, conditional on the + complete-randomization assignment mechanism. (The paper's printed + randomization-inference p-value of 0.020 is not reproducible with the + authors' own package, which implements the inclusive Phipson-Smyth rule + and converges to ~0.051 at 100k replications — see the methodology + registry's RI Note; the implementation follows the package convention.) This demonstrates the + method works with as few as one treated unit. + +- **Walmart minimum-wage study** (LW 2025, Section 6): A balanced panel of + 1,280 counties over 23 years, with staggered Walmart openings. The + rolling IPWRA estimator with detrending (Procedure 5.1) reveals that + county-level linear trends are critical: the CS (2021) estimate of 5.4% + employment increase shrinks to 3.2% (SE = 0.5%) once heterogeneous + trends are removed — the latter consistent with Basker's (2005) estimate + of 150–300 new retail jobs per Walmart store. + +- **Castle doctrine laws** (LW 2026, Section 7.2): A staggered rollout + across 21 states (2005–2009), with 29 never-treated controls. The + aggregated ATT :math:`\hat{\tau}_\omega = 0.092` (9.2% increase in + homicides) is obtained from a single cross-sectional regression + (Equation 7.19, LW 2026), with the HC3 t-statistic of 1.50. + +Estimator Comparison +-------------------- + +.. list-table:: LWDiD vs. CallawaySantAnna vs. WooldridgeDiD + :header-rows: 1 + :widths: 20 27 27 26 + + * - Feature + - LWDiD + - CallawaySantAnna + - WooldridgeDiD + * - Approach + - Unit-specific transform → cross-sectional TE estimation + - Long-difference :math:`Y_{it} - Y_{i,g-1}` (Eq. 4.13, LW 2025) + - Single saturated POLS/TWFE regression + * - Pre-treatment info + - All periods :math:`\{1,\ldots,g-1\}` (rolling average) + - Only period :math:`g-1` (long difference) + - All periods (full regression) + * - Key identification + - Unconfoundedness of :math:`D_i` w.r.t. :math:`\dot{Y}(0)` (Thm 4.1) + - PT on first differences + - Mundlak-style cohort×time interactions + * - Estimators + - RA, IPW, IPWRA, PSM, matching + - OR, IPW, DR + - OLS, Poisson, Logit + * - Heterogeneous trends + - Yes (detrend, Procedure 5.1) + - No + - No + * - Exact small-N inference + - Yes (:math:`\mathcal{T}_{N-2}` under CLM, Eq. 2.10 LW 2026) + - No (requires large N) + - No (requires large N) + * - Doubly robust + - Yes (IPWRA) + - Yes (DR) + - No (single equation) + * - Efficiency (common timing) + - BLUE + asymptotically efficient (Theorem 3.1, LW 2025) + - Less efficient (uses only :math:`g-1`) + - Equivalent to LW RA (Theorem 3.1) + +Restrictions +------------ + +.. warning:: + + The following restrictions apply to the current implementation: + +- **At least 2 pre-treatment observations per unit for detrend** — the ``detrend`` transformation + fits a unit-specific linear trend on pre-treatment observations; units + with fewer than 2 pre-treatment periods cannot be detrended and are + dropped with a ``UserWarning``. +- **Binary absorbing treatment** — the ``treatment`` column must be a binary + indicator that switches from 0 to 1 and stays on. Non-binary or + non-absorbing treatment raises ``ValueError``. +- **PSM matching** — when ``estimation_method='psm'``, unmatched treated units + (no control within ``caliper``) receive NaN and are excluded from the + ATT. A ``UserWarning`` reports the count of dropped treated units. +- **Propensity score trimming** — IPW/DR clip estimated propensity scores + to ``[pscore_trim, 1 - pscore_trim]`` (default 0.01/0.99) for + numerical stability. Extreme scores indicate poor overlap (violation of + Assumption OVLS, Equation 4.10, LW 2025). +- **Per-period effects** — per-period (event-study) effects live on the + unified post-fit surface: call ``results.aggregate('event_study')`` on + the fitted :class:`~diff_diff.lwdid_results.LWDiDResults`. +- **Not-yet-treated control** — when ``control_group='not_yet_treated'``, + the set of valid controls for cohort *g* at time *r* comprises units + with :math:`D_{i,r+1} + \cdots + D_{iT} + D_{i\infty} = 1` + (Equation 4.12, LW 2025). This excludes already-treated cohorts, + preventing "bad comparisons." + +.. seealso:: + + :class:`~diff_diff.CallawaySantAnna` + Propensity-score reweighting using long differences (Equation 4.13, LW 2025). + :class:`~diff_diff.WooldridgeDiD` + Mundlak-style saturated regression — equivalent to RA under LWDiD for + common timing (Theorem 3.1, LW 2025). + :class:`~diff_diff.ImputationDiD` + FE imputation approach (Borusyak, Jaravel & Spiess 2024). diff --git a/docs/choosing_estimator.rst b/docs/choosing_estimator.rst index 7ab24c00..229b7c7a 100644 --- a/docs/choosing_estimator.rst +++ b/docs/choosing_estimator.rst @@ -614,6 +614,43 @@ exponential unit distance weights, and time decay weights with LOOCV tuning. TROP is computationally intensive. Use ``method='global'`` for faster estimation at the cost of some flexibility vs. ``method='local'``. +LWDiD (Lee & Wooldridge) +~~~~~~~~~~~~~~~~~~~~~~~~ + +**When to use**: Panel data where unit-specific rolling transformations +(demeaning or detrending) can remove pre-treatment heterogeneity, combined +with flexible cross-sectional treatment effect estimation (RA, IPW, IPWRA, +or PSM). Particularly suited when you want a transformation-based +alternative to propensity-score reweighting under staggered adoption. + +**Key features**: + +- Converts panel DiD into cross-sectional estimation via unit-specific + transformations (demean or detrend) applied to pre-treatment outcomes +- Supports both common timing and staggered adoption designs + (never-treated / not-yet-treated controls) +- Doubly-robust estimation (``estimation_method='dr'``) with + influence-function inference (``vcov_type='hc1'``; the ``reg`` path + additionally offers classical/HC2/HC3); cluster-robust inference via + the constructor's ``cluster=`` parameter +- Built-in specification robustness: compare demean vs detrend as an + informal pre-test for sensitivity to trend assumptions + +**vs TWFE**: LWDiD explicitly handles heterogeneous treatment effects; +the transformation removes unit fixed effects prior to estimation, avoiding +the negative-weighting problem under treatment effect heterogeneity. + +**vs Callaway-Sant'Anna**: LWDiD uses rolling transformations rather than +propensity-score reweighting for staggered designs, offering a different +identification strategy with analytical (non-bootstrap) inference. + +**Example**:: + + from diff_diff import LWDiD + est = LWDiD(rolling='demean', estimation_method='dr', cluster='state') + results = est.fit(data, outcome='y', unit='id', time='time', + treatment='treated', first_treat='first_treat') + Bacon Decomposition ~~~~~~~~~~~~~~~~~~~ diff --git a/docs/dev-status.md b/docs/dev-status.md index 3a8bc55d..444f4c6c 100644 --- a/docs/dev-status.md +++ b/docs/dev-status.md @@ -18,6 +18,7 @@ Target: ideally < 1000 lines per module; modules ≥3000 lines are candidates fo | `had.py` | 4906 | Consider splitting (continuous / mass-point / event-study / survey paths) | | `had_pretests.py` | 4769 | Consider splitting (Stute / Yatchew / QUG / joint pretests) | | `diagnostic_report.py` | 4135 | Consider splitting (per-method renderers + provenance) | +| `lwdid.py` | 3925 | Consider splitting (PR #588 fix wave; validation + transforms + 4 cross-sectional estimators + bootstrap orchestration — the transform contracts and estimator dispatch are the natural seams) | | `spillover.py` | 3655 | Consider splitting | | `two_stage.py` | 2430 | Monitor — exited the splitting band when the M-022 aggregate() migration extracted the Stage-2/GMM engine into `two_stage_aggregation.py` | | `power.py` | 3488 | Consider splitting (power analysis + MDE + sample size) | @@ -67,8 +68,10 @@ Target: ideally < 1000 lines per module; modules ≥3000 lines are candidates fo ## Standard Error Consistency `vcov_type` has subsumed the previously-proposed `se_type` knob. `DifferenceInDifferences` -and `TwoWayFixedEffects` accept `vcov_type ∈ {classical, hc1, hc2, hc2_bm, conley}` -(the validated set in `linalg.py::_VALID_VCOV_TYPES`); cluster-robust variance comes from +and `TwoWayFixedEffects` accept `vcov_type ∈ {classical, hc1, hc2, hc2_bm, hc3, conley}` +(the validated set in `linalg.py::_VALID_VCOV_TYPES`); `hc3` is one-way only (it cannot +be combined with `cluster=`) and applies the jackknife-style leverage correction matching +`sandwich::vcovHC(type = "HC3")`; cluster-robust variance comes from `cluster=` alongside the heteroscedasticity kind (`hc1+cluster` ⇒ CR1 Liang-Zeger; `hc2_bm+cluster` ⇒ CR2 Bell-McCaffrey, including the weighted WLS-CR2 port; the N>1 absorbed-FE + weights composition is supported via iterative alternating-projection demeaning, #586); diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index 9c019c33..a2e9cd59 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -80,6 +80,14 @@ groups: changes_in_changes: - diff_diff/changes_in_changes.py - diff_diff/changes_in_changes_results.py + lwdid: + - diff_diff/lwdid.py + - diff_diff/lwdid_results.py + - diff_diff/lwdid_wild_bootstrap.py + - diff_diff/lwdid_randomization.py + - diff_diff/lwdid_sensitivity.py + - diff_diff/lwdid_visualization.py + - diff_diff/lwdid_staggered.py visualization: - diff_diff/visualization/__init__.py - diff_diff/visualization/_common.py @@ -852,6 +860,57 @@ sources: - path: docs/migration-4.0.md type: user_guide + # ── LWDiD (lwdid group) ─────────────────────────────────────────── + + diff_diff/lwdid_wild_bootstrap.py: + drift_risk: medium + docs: + - path: docs/api/lwdid.rst + type: api_reference + + diff_diff/lwdid_randomization.py: + drift_risk: medium + docs: + - path: docs/api/lwdid.rst + type: api_reference + + diff_diff/lwdid_sensitivity.py: + drift_risk: medium + docs: + - path: docs/api/lwdid.rst + type: api_reference + + diff_diff/lwdid_visualization.py: + drift_risk: low + docs: + - path: docs/api/lwdid.rst + type: api_reference + + diff_diff/lwdid_staggered.py: + drift_risk: medium + docs: + - path: docs/api/lwdid.rst + type: api_reference + - path: docs/methodology/REGISTRY.md + section: "LWDiD" + type: methodology + + diff_diff/lwdid.py: + drift_risk: medium + docs: + - path: docs/api/lwdid.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.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/index.rst b/docs/index.rst index 9a18f87f..f8116e66 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -183,6 +183,8 @@ Supported Estimators - Wooldridge (2023, 2025) extended TWFE (ETWFE) via saturated OLS or QMLE * - :class:`~diff_diff.ChangesInChanges` - Athey & Imbens (2006) distributional DiD with quantile treatment effects + * - :class:`~diff_diff.LWDiD` + - Lee & Wooldridge (2025, 2026) rolling-transformation DiD; ``rolling='detrend'`` handles heterogeneous linear trends * - :class:`~diff_diff.QDiD` - Quantile DiD comparison estimator applying DiD quantile-by-quantile (deprecated 3.9 - use :class:`~diff_diff.ChangesInChanges` with ``method="qdid"``) * - :class:`~diff_diff.RegressionDiscontinuity` diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 55e5fec6..3dfb9788 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -673,7 +673,7 @@ distinction governs which `vcov_type` values an estimator can accept: **Analytical-sandwich estimators** fit a single (or per-cohort) linear regression and derive variance via `solve_ols(..., vcov_type=...)`, returning a sandwich `(X'X)^{-1} M (X'X)^{-1}` whose meat `M` is parameterized by -`vcov_type ∈ {classical, hc1, hc2, hc2_bm}` (plus `conley` for spatial-HAC). +`vcov_type ∈ {classical, hc1, hc2, hc2_bm, hc3}` (plus `conley` for spatial-HAC). `hc3` entered the shared vocabulary with the LWDiD contribution; estimators that do not implement it reject it explicitly (structural roster guard in `tests/test_estimators_vcov_type.py`). Examples: `DifferenceInDifferences`, `MultiPeriodDiD`, `TwoWayFixedEffects`, `SunAbraham`, `StackedDiD`, `WooldridgeDiD`, `LinearRegression`. The full `vcov_type` contract is methodologically applicable because every family has @@ -2524,12 +2524,13 @@ Full maintainer paper reviews (equation-level detail, replication targets): `doc - **Note:** Registry entry authored with the paper reviews ahead of the implementation (PR #588, third-party contribution under maintainer revision). Checklist boxes are unchecked until the implementation lands; the implementation-specific Notes in this section (control pools, RI convention, overall conventions, Sec 4.3 rejection, API conformance) are the maintainer pre-pass of that finalization — remaining edge-case notes close with the merge. - **Note:** Maintainer validation suite: `tests/test_methodology_lwdid.py` (import-skip-gated until `diff_diff.lwdid` exists) is the ACCEPTANCE BAR for PR #588's final round: it is written against the agreed post-rename API and intentionally fails on the pre-rename head; the round is complete when it passes unmodified. The earlier xfail scaffolding is retired. Goldens: `benchmarks/data/lwdid_walmart_eventstudy_golden.json` (printed Tables A4/A5; point halves consumed, printed SE halves superseded), `benchmarks/data/real/castle_lw_subset.csv`, and `benchmarks/data/lwdid_stata_golden.json` (full-precision authors'-package parity, `benchmarks/stata/generate_lwdid_golden.do`; schema checked ungated by `tests/test_lwdid_stata_golden_schema.py`). -- **Note:** Pre-release API-conformance decisions (locked 2026-08-15/16; no deprecation-ledger rows are owed because LWDiD never shipped on main): `estimation_method=` with house-CS values `reg`/`ipw`/`dr` plus the LWDiD-only `psm`; `vcov_type=` (name locked; value set finalized in the contributor's round within the house `linalg.py` vocabulary — `hc3` retained as the paper-mandated extra via a shared linalg extension, the `"cluster"` mode value retired in favor of clustering-by-`cluster=`, `hc0`/`hc4`/`conley` not offered); `seed=` with default `None`; `pscore_trim=`; `cluster=` on the constructor; the `LW` alias, module-level `lwdid()` wrapper, `lwdid_trend_diagnostics` module, and `overall_att`/`period_effects`/`period_specific` surfaces retired (per-period effects live on the unified post-fit event-study surface). +- **Note:** Pre-release API-conformance decisions (locked 2026-08-15/16; no deprecation-ledger rows are owed because LWDiD never shipped on main): `estimation_method=` with house-CS values `reg`/`ipw`/`dr` plus the LWDiD-only `psm`; `vcov_type=` (name locked; accepted sets finalized in the maintainer fix wave: `reg` accepts `{classical, hc1, hc2, hc3}` — `hc3` retained as the paper-mandated extra via a shared linalg extension; `ipw`/`dr` accept `hc1` only, because the influence-function variance is always used on those paths and any other value would be silently inert; `psm` accepts `hc1` as its only configuration value while its inference is UNAVAILABLE (NaN; the Abadie-Imbens matching variance is a DEFERRED item); `cluster=` composes only with `hc1` (CR1) and is rejected for `psm`; the `"cluster"` mode value retired in favor of clustering-by-`cluster=`, `hc0`/`hc4`/`conley` not offered); `seed=` with default `None`; `pscore_trim=`; `cluster=` on the constructor; the `LW` alias, module-level `lwdid()` wrapper, `lwdid_trend_diagnostics` module, and `overall_att`/`period_effects`/`period_specific` surfaces retired (per-period effects live on the unified post-fit event-study surface). **Key implementation requirements:** *Assumption checks / warnings:* - Treatment is absorbing (no reversibility), common timing (`1 < S <= T`) or staggered (cohorts `g in {S,...,T,infinity}`, mutually exclusive and exhaustive); at least one pre-treatment period. +- **Note (cohort encodings — maintainer fix wave, one shared normalizer):** never-treated is `0`, `NaN`/`NaT`, or `np.inf`; `inf` is recoded to never-treated with a warning (the CS convention recodes exactly `0`/`inf` — the NaN/NaT limb is an LWDiD-only extension required by datetime/Period scales and is a documented CS deviation), and a finite cohort BEYOND the last observed period is also recoded to never-treated with a warning — a second documented CS deviation (CS keeps finite cohorts out of never-treated): under `control_group='never_treated'` such units join the control pool and contribute no pre-period event cells; the alternative leaves the all-eventually-treated guard incoherent. Negative/`-inf` cohorts raise. Numeric BETWEEN-period cohorts are rejected (an observed-support membership check with an explicit message), while datetime/Period cohorts map to the next observed period — a dtype-dependent input contract stated in `docs/api/lwdid.rst`; numeric parity with CS's between-period acceptance is a TODO row. The treatment design check requires `D_it == 1[t >= g_i]` over each unit's OBSERVED rows; the onset row itself may be unobserved (unbalanced panels with a missing onset row are accepted). - Common timing: **NAC** (no anticipation, eq. (2.7)), **CPTC** (conditional parallel trends, eq. (2.10)), **OVLC** (overlap, eqs. (2.14)-(2.15)); Theorem 2.1 identifies `tau_r`, `r = S,...,T`. - Staggered: **CNAS** (eq. (4.4); X-conditioning droppable with NT-only controls), **CPTS** (eq. (4.6)), **OVLS** (eq. (4.10), control pool `A_{r+1} = D_{r+1} + ... + D_T + D_infinity`); Theorem 4.1. - Heterogeneous linear trends: **CHT** (eq. (5.3)) — demeaning inconsistent under CHT; unit-specific detrending (Procedure 5.1) restores consistency under CNAS + CHT + OVLS. @@ -2549,22 +2550,58 @@ Event-study/placebo transformations over ALL periods (Appendix D): demeaning (D. - IPWRA (workhorse): logit propensity score per cell + WLS with weights `w = D + (1-D) p/(1-p)`; IPW = special case without the outcome-regression component. - **Note:** API vocabulary (as specified for the PR #588 implementation, in review): `estimation_method="reg"` = the paper's RA, `"ipw"` = IPW, `"dr"` = IPWRA (house CallawaySantAnna values; the doubly-robust option), `"psm"` = propensity-score matching, an LWDiD-only extra with no CS analog. - Control pool at (g, r): `A_{r+1} = 1` (never-treated + not-yet-treated) by default; NT-only optional. Pre-treatment placebo cells use the Appendix D.3 rule `A_{g,t} = {G = g} ∪ {G = 0} ∪ {G > max(g,t)}`. +- **Note (replicating the paper's staggered numbers):** the implementation's default is `control_group='not_yet_treated'`, matching OVLS (eq. (4.10)) as stated in the text. The paper's *printed* staggered results, however, are computed against the never-treated pool only, so **reproducing them requires passing `control_group='never_treated'` explicitly**. This is not a discrepancy in either direction — it is a sample-definition choice that the text leaves to the analyst while the applications fix it to NT-only. Every staggered replication golden in `tests/test_methodology_lwdid.py` (castle `tau_omega`, and the composite-regression reference) therefore passes `control_group="never_treated"`; a default-pool fit yields different, equally valid estimates because the (g,t) cells draw on a strictly larger control sample. *Aggregation:* - Event-study: `WATT(r) = sum_{g in G_r} omega_{g,r} ATT(g, g+r)` with `omega_{g,r}` = (treated units of cohort g contributing at event time r) / (total treated units contributing at event time r) - the operative definition per LW 2025 Appendix E.1, required under unbalanced panels where a cohort's contributing count at r can differ from `N_g`. In balanced panels this simplifies to `N_g / N_{G_r}` (Sec. 6.2/D.3). Aggregated influence function `IF_{i,r} = sum_g omega_{g,r} IF_{i,g,g+r}`. - Overall: composite-outcome single regression (LW 2026 eqs. (7.18)-(7.19)) — `tau_omega` with cohort-share weights `omega_g = N_g / N_treat`; automatically accounts for correlation among per-cohort effects and supports exact small-N inference. -- **Note:** The authors use TWO overall conventions across their own modes (measured 2026-08-15 against Stata `lwdid` v2.4.2): the small-N composite regression targets `tau_omega` (cohort-mean-then-treated-weight, eq. (7.18)) while the large-N display's `Post_avg` is the cell-mass (treated-count-per-cell) weighted average of the post ATT(g,t) — which equals the house CallawaySantAnna `"simple"` convention ON BALANCED PANELS (verified to the last digit on Walmart; CS-simple prefers fixed unit-cohort mass under unbalanced panels, so the equivalence carries that qualification). The implementation's `.att` is `tau_omega` — what the papers print, and the estimand the composite regression's inference is built for; a `vcov_type` selection must never move the point estimate (cross-path pins in the validation suite, including an unbalanced-panel pin). `aggregate("simple")` view-relays the fit per the house aggregation contract; exposing the cell-mass overall as an aggregate extra is a TODO row. +- **Note:** The authors use TWO overall conventions across their own modes (measured 2026-08-15 against Stata `lwdid` v2.4.2): the small-N composite regression targets `tau_omega` (cohort-mean-then-treated-weight, eq. (7.18)) while the large-N display's `Post_avg` is the cell-mass (treated-count-per-cell) weighted average of the post ATT(g,t) — which equals the house CallawaySantAnna `"simple"` convention ON BALANCED PANELS (verified to the last digit on Walmart; CS-simple prefers fixed unit-cohort mass under unbalanced panels, so the equivalence carries that qualification). The implementation's `.att` is `tau_omega` — what the papers print — WHENEVER the complete-case resolution drops no unit (all balanced panels included); with drops, `.att` is the influence-weighted cohort-mass point on every route and `tau_omega` moves to the `att_tau_omega_complete_case` diagnostic (see the routing Note below). A `vcov_type` selection never moves the point estimate in either stratum (cross-path pins in the validation suite, including unbalanced-panel pins). `aggregate("simple")` view-relays the fit per the house aggregation contract; exposing the cell-mass overall as an aggregate extra is a TODO row. +- **Note (overall estimand routing — maintainer fix wave, supersedes the round-4 unification Note):** the composite (`tau_omega`) surface is COMPLETE-CASE with FIXED cohort weights: treated units without a finite own-cohort post-window average are dropped with a warning and the cohort masses `omega_g = N_g / N_treat` are RECOMPUTED on the surviving treated sample; control units must observe every surviving-weight cohort's post window with a finite transformed outcome (warn + drop otherwise; the pre-fix code zero-filled missing control entries and let the OLS finite mask silently reweight the treated side). Routing is vcov-INVARIANT in every stratum: with ZERO complete-case drops, `.att` is the composite `tau_omega` on every vcov route (classical pairs it with the composite's own `T_{n-2}` SE; hc1/clustered with the joint-IF SE — a documented approximation, the two weightings coincide only under balance); with ANY drops, `.att` is the influence-weighted cohort-mass point on ALL routes with route-matched IF SEs (internally coherent pairs), and the complete-case composite is exposed as the `att_tau_omega_complete_case` diagnostic with `n_composite_treated_dropped`/`n_composite_controls_dropped` counters and a warning. Gate: `rolling in {'demean', 'detrend'}` + `control_group='never_treated'` + `estimation_method='reg'` + no covariates. +- **Note (seasonal-variant overall boundary — maintainer fix wave):** for `demeanq`/`detrendq` the overall `.att` is the cohort-mass-weighted average of SEASONAL cohort ATTs, identical under every vcov setting; no seasonal composite `tau_omega` is defined or implemented (the composite raises if a q-variant reaches it). The pre-fix behavior — routing q-mode `classical` fits through a composite built on the NON-seasonal transforms — matched neither this estimand nor any seasonal one and moved the point with the vcov selection. *Standard errors:* - Large-N default: influence-function **multiplier bootstrap** (LW 2025 Algorithm 1): IFs per (g,t) from E.2 (RA, finite-sample exact), E.3 (IPWRA, stacked M-estimator with logit-score correction), E.4 (IPW, `psi - Gamma' IF_gamma` correction); centered IFs; **unit-level Rademacher multipliers** (one draw per unit across all cells — unit clustering by construction); sup-t simultaneous bands over the event-study path; B = 999 in the paper's application; anchor periods excluded from the raw effect set (the public `EventStudyResults` surface still emits OBSERVED anchors as zero-valued `is_reference` rows — that is a display convention, not a contradiction of the exclusion). - **Note:** Inference DEFAULT (as specified for the PR #588 implementation, in review): analytical vcov with the bootstrap as opt-in (`n_bootstrap=999` for the paper's setting) — the house convention (CallawaySantAnna ships the same analytical default), while this bullet describes the paper's recommended large-N inference. - Small-N exact (LW 2026): usual OLS SE on the collapsed cross-sectional regression with exact `T_{N-2}` / `T_{N-K-2}` reference distribution; valid down to `N = 3` and a single treated unit (`N1 = 1` — the t statistic is the studentized residual; same for `N_g = 1` per cohort in (7.8)/(7.10)). - **Note (controlled exact inference is design-coherent):** LW 2026 Section 2 sanctions two controlled collapsed regressions — plain controls `(1, D, X)` with exact reference `T_{N-K-2}`, and interacted controls `(1, D, X, D(X - Xbar_1))` whose design rank implies `T_{N-2K-2}` (with the per-group guards `N0 > K+1` and `N1 > K+1`). Either is admissible for `vcov_type="classical"`; mixing them (fitting one design while reporting the other design's df) yields anti-conservative exact p-values and is a defect. The validation suite pins coherence: the reported p-value must use the residual df of whichever design reproduces the fitted point estimate. +- **Note (per-surface reference distributions — maintainer fix wave):** the reference distribution follows the surface, pinned by the validation suite: the composite `tau_omega` under `vcov_type='classical'` reports its own regression's `T_{n-2}`; an influence-function aggregate composed of EXACTLY ONE cell uses that cell's residual df (so a single-post-period staggered fit matches the common-timing fit identically — resolving a pre-fix asymmetry where the same one-cell design switched between t and normal references on the dispatch path); multi-cell unclustered IF aggregates use the large-sample normal reference (units recur across cells with overlapping influence functions, so no residual-df pooling is valid); clustered aggregates use `G - 1` where G counts the clusters CONTRIBUTING to the aggregate's estimated cells (clusters supplying no cell must not inflate the df or `n_clusters`). Sub-samples (staggered cells, the collapsed common-timing cross-section, event-study period cells) with fewer than 2 clusters fail closed: point retained, inference NaN, and any aggregate including such a cell inherits NaN inference. +- **Note (wild-bootstrap convention — maintainer fix wave):** `wild_cluster_bootstrap` delegates to the house WCR engine (`diff_diff.utils.wild_bootstrap_se`, `fwildclusterboot::boottest`-matched): the null is imposed by dropping the treatment column while KEEPING controls, the CI is test-inverted, and the p-value uses strict-exceedance counting with a ~1e-9 tie guard AND a documented zero-p floor at `1/(n_valid+1)` when that floor is below alpha (a deliberate, documented departure from boottest recorded in `diff_diff/utils.py`). This differs from the randomization-inference module's INCLUSIVE Phipson-Smyth rule (see the RI Note above) — different procedures, both documented. Exactly-identified degenerate designs (cluster-invariant treatment with zero cluster scores, e.g. the canonical G=2 two-group case) fail closed in the wrapper: point retained, se/t/p/CI NaN. +- **Note (review round: enforced small-sample + degenerate-design guards):** the Registry's sample-size guards are ENFORCED - an overall common-timing collapsed design with `N < 3` or non-positive residual df raises an informative ValueError (the pre-fix code coerced df to 1 or hit a raw ZeroDivisionError); non-estimable cohort-time/event cells are marked NaN with a skip reason instead of raising. HC3 with a leverage-one observation (e.g. a single treated unit) has no defined variance and fails closed (warning + NaN inference tuple) in the shared linalg meat rather than flooring `1 - h_ii`. `estimation_method='psm'` reports NaN inference (point retained): the naive matched-pairs variance ignored matched-control reuse and first-stage uncertainty; an Abadie-Imbens matching variance is a DEFERRED item, and staggered PSM rejects `n_bootstrap > 0` (matching has no influence function for the multiplier bootstrap). Fisher randomization inference offers `method='permutation'` only (the former with-replacement 'bootstrap' label resampling had no assignment-mechanism justification and is removed). `n_bootstrap` accepts 0 or >= 2. Fit provenance (`control_group`, `n_bootstrap`, `seed`, PSM settings) is stored on the results object. +- **Note (review round 2: contract + inference-scope guards):** user role columns (outcome/unit/time/treatment/first_treat/cluster/covariates) may not use LWDiD's reserved internal working names (`_treat`, `_ydot`, `_ydot_avg`, `_ever_treated`, `_boot_unit`, `_lwdid_time_pos`, `_lwdid_cohort_pos`) and core roles must be pairwise distinct — a collision previously overwrote the internal column silently (e.g. `cluster='_treat'` reported the cluster labels' coefficient as the ATT). The common-timing unit bootstrap draws per-replicate `SeedSequence`-spawned streams identically for every `n_jobs`, so a seeded fit is reproducible across execution modes. In STAGGERED fits `n_bootstrap` governs the event-study multiplier bootstrap ONLY (LW 2025 Algorithm 1 is defined over the event-study influence path); the overall/cohort aggregates keep analytical influence-function inference — per-surface provenance is recorded via `cband_method`/`cband_n_bootstrap` and `inference_basis`. An event cell whose multiplier draws are degenerate (zero/non-finite bootstrap SD) fails closed — point retained, inference NaN, `inference_status='degenerate_bootstrap'` — instead of silently reverting to its analytical SE. The common-timing event-study surface covers POST periods only (`WATT(r)`, `r >= 0`); pre-treatment placebo cells are a staggered-path surface (Appendix D pools), available for a single cohort by passing `first_treat=`. Sensitivity analyses (`robustness_pre_periods`, `sensitivity_no_anticipation`) accept a SINGLE treated cohort only: their exclusion windows are defined relative to the earliest adoption, which would mislabel later cohorts' transformation samples (cohort-relative exclusions are a DEFERRED item). The `robustness_level` classification (10%/25%/50% sensitivity-ratio cuts) is a diff-diff library heuristic — the papers recommend the underlying diagnostics but define no categorical scale. `pscore_trim` joins the stored fit provenance on the ipw/dr/psm paths. +- **Note (review round 3: calendar seasons, PSM contract closure, bootstrap provenance):** encoded staggered panels (datetime/Period) now carry the CALENDAR quarter through the dense-position encoding (`_lwdid_season`), and the seasonal transforms prefer it — the pre-fix numeric fallback `(position - 1) % 4 + 1` silently relabeled every season after a globally missing calendar period (execution-verified: a gapped quarterly panel with treated/control-differential seasonality biased a zero-effect staggered demeanq ATT to ~0.12); numeric time columns keep the documented `(t - 1) % 4 + 1` user contract. The PSM fail-closed contract is enforced on every route: covariate-less PSM is rejected (no propensity score to match on — the pre-fix delegation returned finite OLS inference under method 'psm'), `n_bootstrap > 0` is rejected in BOTH timing modes (the standard bootstrap is invalid for nearest-neighbor matching estimators, Abadie & Imbens 2008, Econometrica — the pre-fix common-timing path replaced the NaN contract with a naive pairs-bootstrap SE), and a non-converged propensity model returns the regression-adjustment POINT with NaN inference. Common-timing bootstrap fits record `inference_basis` (`unit_bootstrap`/`cluster_bootstrap`) because the headline se/p/CI are bootstrap while `params`/`vcov` remain analytical; `summary()` renders the basis on all fits. `get_transformation_diagnostics` runs the same front-door validation as `fit()` (binary treatment, reserved names, panel duplicates, treatment-design coherence). Sensitivity parameters are strictly validated (`exclude_periods` positive unique ints — `0` previously sliced away EVERY pre-period via `[:-0]`; `k_min`/`k_max` positive ints); duplicate covariate names are rejected. +- **Note (review round 4: rank-aware advanced inference + cell guards):** the randomization-inference and wild-cluster-bootstrap modules fit through the shared rank-aware `solve_ols` (the pre-fix `np.linalg.lstsq` returned a finite MINIMUM-NORM treatment coefficient when a control duplicated the treatment column — execution-verified: a true ATT of 2.0 reported as 0.877 with finite p-values in both modules); an unidentified observed treatment coefficient raises, and a permutation draw whose treatment coefficient is dropped counts as a failed replication. The standalone wild wrapper requires an integer `n_bootstrap >= 2` (one draw supports neither a dispersion estimate nor test inversion), `n_reps` must be a non-boolean positive integer, and `LWDiDResults.wild_cluster_bootstrap` inherits the fitted `alpha` (override by argument). Staggered cells under `control_group='never_treated'` are non-estimable when transformation/finite drops leave fewer than 2 never-treated controls (the Registry's NT-only minimum, previously checked only on raw units pre-fit). Staggered `reference_periods` emit only OBSERVED anchors — an anchor `r` appears iff some treated cohort g has calendar period `g + r` in the panel (previously a numeric time gap synthesized a zero-valued reference row at a nonexistent event time). +- **Note (review round 5: post-fit inference replays the fitted estimand):** `LWDiDResults.wild_cluster_bootstrap()` and `.randomization_test()` REPLAY the fit-time collapsed cross-section and the exact RA design `[1, D, X, D(X - Xbar_1)]` — no data arguments are accepted, the RI permutations RECOMPUTE the treated covariate mean (and interaction columns) per assignment (`design='ra_interacted'` on the standalone function), and the replayed observed statistic is asserted equal to `.att` before any p-value is cached (fail-closed RuntimeError on mismatch). The pre-fix methods accepted arbitrary caller arrays and a non-interacted `[1, D, X]` design, caching p-values for a DIFFERENT estimand (execution-verified: fitted ATT 3.98 vs tested 3.26 on a covariate-unbalanced RA fit). Both methods are defined for common-timing `estimation_method='reg'` fits only (re-estimating ipw/dr/psm per draw is not implemented — informative rejection); the standalone module functions keep their generic array contracts. The common-timing bootstrap resamples from units SURVIVING the transformation/finite filters, treats a draw collapsing onto one distinct cluster as failed, and fails closed (point retained, NaN inference) when fewer than 2 effective clusters survive — a requested bootstrap no longer overwrites the single-effective-cluster fail-closed state (pre-fix: SE ~ 2e-16 from the raw cluster map). `aggregate(balance_e=)` is REJECTED: LWDiD stores no per-cohort estimation kit from which a balanced-cohort sample and covariance could be recomputed post fit (pre-fix the argument was accepted and silently ignored). +- **Note (review round 6: fixed-window estimand, scale-equivariant guards, fweight leverage):** the common-timing headline ATT is the paper's FIXED-WINDOW post average (LW 2026, denominator `T - S + 1`): units lacking a finite transformed outcome in every post period are dropped as complete cases with a warning, on the point path and the bootstrap alike (pre-fix, each unit averaged whichever post periods it observed, so calendar composition could masquerade as treatment effect — execution class: zero-effect panel, treated units observing one extra post period reported that period's trend as the ATT). If the drops empty a treatment arm, fit raises informatively instead of dispatching a one-arm design. The degenerate-SE guard is SCALE-EQUIVARIANT (`se <= sqrt(eps) * |effect|`, plus non-positive/non-finite): the former `max(1, |effect|)` floor NaN'd valid inference when the outcome was rescaled by 1e-10 while the t-statistic is scale-invariant. Constructor numerics are strictly validated (`pscore_trim` real/non-bool/finite before conversion; `n_jobs` rejects booleans). `n_jobs` is deliberately EXCLUDED from result provenance: since the round-2 SeedSequence fix the seeded bootstrap draws are identical for every `n_jobs`, so it is pure execution configuration with no effect on any reported number. `df_inference` serializes in `to_dict()`. +- **Note (round 6, shared linalg — fweight leverage for HC2/HC3):** under `weight_type='fweight'` the HC2/HC3 leverage denominator uses each replicate row's UNWEIGHTED quadratic form against the weighted bread, `h_i = x_i'(X'WX)^{-1}x_i` — frequency weights mean replicated data (integer counts, `df = sum(w) - k`, HC1 expansion parity), and this makes compressed HC2/HC3 exactly equal literal `np.repeat` expansion (oracle-pinned). The WLS-hat convention `h_i = w_i x_i'(X'WX)^{-1}x_i` (R `sandwich::vcovHC`) continues to apply to aweight/pweight. Pre-fix, the weighted hat under fweight produced HC2/HC3 variances up to ~5x the expansion (hc3 was introduced by this fix wave; hc2's fweight surface predates it and carried the same mismatch — both now expansion-exact, unpinned by any golden). +- **Note (review round 7: replay follows the fitted design; sensitivity runs the design check):** the post-fit replay mirrors `_estimate_reg`'s LW eq. 3.3 interaction gate (`N_1 > K+1` AND `N_0 > K+1`): small-arm fits use the plain `(1, D, X)` design and their replayed RI/WCR statistic matches `.att` (pre-fix the replay always interacted, so the round-5 coherence assert made small-arm fits' post-fit inference unusable — the fail-closed backstop working as designed, now with the correct design selected). `_prevalidate_frame` in the sensitivity helpers runs the full treatment-design check (absorbing treatment, common-timing onset homogeneity, D_it/cohort consistency, with fit's encode-then-normalize ordering), so structural design violations RAISE instead of being swallowed by the per-spec ValueError handler as `not_estimable`. +- **Note (review round 8: calendar partition from S, tau_omega window counts, aweight leverage-family convention):** the common-timing pre/post partition derives from the SINGLE adoption period `S = min(observed treated period)`: `pre = {t < S}`, `post = {t >= S}` (the pre-fix per-period `max(D)` partition classified a post period with no observed treated rows as PRE-treatment, contaminating the rolling pre window — execution-verified: zero-effect trend panel biased to ATT 0.75); the common-timing design check likewise validates `D_it = 1[t >= S]` over observed rows, so a unit whose `t = S` row is missing is accepted (matching the staggered branch) while genuinely heterogeneous onsets still raise. The `tau_omega` complete-case semantics are CLARIFIED, not changed (round-8 reviewer proposed full-window per-period counting; NOT adopted): a unit contributes cohort g's component iff its OBSERVED post-g rows yield a finite average — partial post windows are averaged over observed rows, symmetrically for treated and control units. This is the adjudicated WS1 design pinned byte-frozen by the acceptance suite's independent reference oracle (`_complete_case_tau_omega_reference`) and its zero-drop metadata test; changing to every-period counting would change the estimand those tests pin. CAVEAT (documented): on unbalanced panels with time trends, differential post-period availability enters the composite through the observed-window averages — complete-case drops fire only when a required window is entirely missing/non-finite. Covariates must be numeric and FINITE at the front door (Inf passed the NaN check and was silently cell-filtered); `validate_staggered_data` rejects datetime64-vs-Period mixtures and Period-frequency mismatches exactly like the encoding step. **aweight + hc2/hc3 refutation (round 8):** the reviewer's proposed `w^2` score meat contradicts the documented aweight convention (this section, Weight Type Effects: "aweights use unweighted meat ... matches Stata convention" — known-heteroskedasticity WLS leaves ~homoskedastic errors) — hc2's aweight surface is RELEASED behavior retained byte-identical from main, and hc3 follows the same family branch (unweighted meat with the WLS-hat leverage). The convention is deliberate and documented, not a defect. +- **Note (review round 9: one event-time convention, onset partition propagated, finite outcomes):** event-time labels follow ONE convention across the common and staggered interfaces: NUMERIC calendars use the Registry's arithmetic `r = t - g` (validated INTEGRAL — a fractional horizon raises instead of silently merging under the integer storage keys, which previously overwrote distinct horizons' estimates and covariance entries via `int(t - g)`); datetime/Period calendars use position differences on the ordered support (they are position-encoded before the staggered machinery). Pre-fix, the common interface used positional labels for ALL dtypes, so a gapped numeric calendar got different event keys per interface ({0,1} vs {0,2} on {1,2,4,6} with onset 4). The round-8 onset partition (`pre = {t < S}`) is propagated to `get_transformation_diagnostics` and the sensitivity helpers' pre/post sets (a controls-only post period is post everywhere; sensitivity subsets retain every `t >= S` period). Outcomes must be numeric and FINITE at the front door in both timing modes (Inf previously passed the NaN check and was silently np.isfinite-filtered inside staggered cells, changing the estimation sample without warning). +- **Note (review round 10: guard ordering + sensitivity coherence):** hc3's undefined-leverage fail-closed check runs BEFORE the generic over-one HC1 fallback (numerically over-one leverage previously escaped into an HC1 result still labeled hc3), and the LWDiD hc3 influence vector fails closed to NaN under the same condition instead of clipping (aggregate inference matches the cell's NaN vcov; the NaN influence drops the cell from joint aggregation). The sensitivity helpers count treated cohorts on the NORMALIZED frame (beyond-window/inf encodings no longer masquerade as extra cohorts), their BASELINE full-frame fit propagates every fit error (a configuration/support failure such as covariate-free PSM raises instead of reporting `not_estimable`; only restricted-subset fits map failures to NaN specs), and zero-post-row units are counted by the fixed-window drop warning (previously they vanished silently in the merge). +- **Note (review round 11: propensity linearization, reduced-rank propensity fits, identified-rank gate):** the IPW/DR influence functions build the logit score and Hessian from the RAW fitted probabilities (the actual MLE's estimating equation — its score is ~0 at the fit; the pre-fix code used the CLIPPED probabilities, breaking the linearization whenever `pscore_trim` fired), and the weight-derivative `dw/dgamma` is ZERO for clipped observations (a clipped weight is locally constant in gamma); the clipped probabilities remain the WEIGHTING choice for the point estimator. A rank-deficient propensity model (NaN logit coefficients from dropped collinear columns, finite probabilities) CONTINUES as an IPW/DR fit on the reduced-rank propensity (score/Hessian on the kept columns) — the pre-fix code silently substituted regression adjustment under ipw/dr provenance; only a genuinely failed solve (non-finite probabilities) falls back, with its warning. The RA interaction gate (`N_1 > K+1` and `N_0 > K+1`, eq. 3.3) counts the IDENTIFIED control dimension (matrix rank), mirrored exactly by the post-fit replay — a perfectly collinear control previously flipped the gate and changed the ATT while adding no information. `validate_staggered_data` marks duplicate `(unit, time)` cells invalid (a duplicate could mask a missing cell in the row-count balance check). Tutorial 27 re-executed against the final code: the single-treated California HC3 example now TEACHES the leverage-one fail-closed boundary (classical exact-t / RI are the small-N tools), and the CS-efficiency comparison states the paper-faithful serial-correlation trade-off instead of a dominance claim. +- **Note (review round 12: rank-aware DR nuisances, identified parameter counts, survivor cohort masses):** the DR outcome WLS fits through the shared rank-aware solver and every outcome-model influence term (prediction, `S_beta`, `H_beta`, `dATT/dbeta`) uses the IDENTIFIED column mask (the pre-fix raw `inv`/`pinv` Gram was not scale-equilibrated — an exactly redundant 1e12-rescaled duplicate changed the DR SE by ~2.5x); IPW/DR report the identified propensity/outcome ranks as `n_params` (nominal counts let a redundant control shrink residual df and move p-values/CIs). On the tau_omega DROPS route, `.att` and its combined influence function weight cohorts by the SURVIVING cohort masses returned by the composite helper (the Registry complete-case rule; raw masses previously kept dropped treated units in the weights) — pinned by an independent survivor-mass oracle. +- **Note (review round 13: scale-equilibrated influence bread, effective-rank guard):** the RA influence reconstruction inverts the COLUMN-EQUILIBRATED Gram and unscales (`(X'X)^{-1} = D^{-1}(Xs'Xs)^{-1}D^{-1}`) — the pre-fix raw-Gram pinv silently dropped low-scale directions at large covariate units, so cell ATT/SE (from the equilibrated `solve_ols`) were unit-invariant while every AGGREGATE SE/p/CI and the multiplier-bootstrap inputs were not (execution class: rescaling one covariate by 1e7 moved the overall SE from 0.128 to 0.028 with no warning). Aggregate-inference unit-invariance is pinned across the overall and event-study surfaces. The exact-inference small-sample guard uses the EFFECTIVE (equilibrated) design rank, so a redundant-column design with positive effective residual df fits while a genuinely saturated design still raises. docs/index.rst and the practitioner tree scope the heterogeneous-trends claim to `rolling='detrend'` and describe PSM as point-estimation-only. +- **Note (review round 14):** numeric TIME columns must be finite at the front door (`+/-Inf` previously passed the NaN check and raised a raw OverflowError in event-time arithmetic); datetime/Period/ordered-label time columns are unaffected. The Prop-99 api-docs passage no longer claims the exact-inference p-value "validates the normality assumption" (it tests the treatment-effect null under those assumptions; RI is the assumption-free robustness check). +- **Note (review round 16):** a staggered event row whose accepted SE is non-finite contributes NO column to the analytical event-study covariance (its influence is not stored and `compute_event_study_bands` filters defensively) — previously a NaN-inference row could expose a 0.0 covariance diagonal, presenting it as known without uncertainty (the common-timing path already guarded on finite SE). `robustness_pre_periods` honors `k_min=1` for demeaning (the former unconditional `max(k_min, 2)` silently dropped a valid one-pre-period spec) and rejects sub-minimum `k_min` for detrending explicitly. The degenerate all-NaN-transform early returns carry full fit provenance (`cluster_name`, `psm_config`). +- **Note (CI review rounds, tutorial withdrawal):** the contribution's tutorial (`27_lwdid.ipynb`) was WITHDRAWN from the PR after five CI review rounds whose only remaining findings were its empirical narrative — the Walmart "common-timing" example fabricated a shared 1986 onset for cohorts first treated as late as 1999 and translated the pooled contrast into hires figures inconsistent with the staggered estimate (0.0109 × 6,589 ≈ 72 jobs, CI including zero, vs a claimed 150–300). Earlier tutorial-related clauses in the round notes above are historical. A replacement notebook is a tracked follow-up (TODO.md) using the numbers-locked authoring workflow. +- **Note (within-cohort aggregation — cell-mass convention, documented deviation from eq. 7.10 on unbalanced panels):** `cohort_effects[g]` (and `aggregate('group')`) aggregate cohort g's estimable cohort-time cells weighted by each cell's contributing TREATED mass (`n_treated`) — the same contributing-treated-unit convention as the WATT(r) event-time axis (LW 2025 E.1) and the authors' package's cell-mass `Post_avg` display (golden-pinned). On balanced never-treated designs this equals the LW 2026 eq. 7.9/7.10 estimand (regress unit-level post averages on `[1, D_g]`); on UNBALANCED panels the two differ — under cell-mass weighting a treated unit observing more post periods carries proportionally more weight within its cohort, while eq. 7.10 weights units equally via their own post averages. The eq. 7.10 unit-average cohort estimand is a tracked follow-up (TODO.md, alongside the cell-mass overall-ATT row). Pinned by an unbalanced-panel oracle that distinguishes the two weightings — the same fixture also shows the surfaces answering different documented estimands: `.att` on the tau_omega route is the eq. 7.18 COMPOSITE built from unit post-averages (there exactly the eq. 7.10 value), while `cohort_effects` reports the cell-mass aggregate. +- **Note (review round 17):** the RA interaction gate, the small-sample effective-rank guard, and the post-fit replay mirror all use the SHARED solver's pivoted-QR rank detector (`_detect_rank_deficiency`, scale-invariant 1e-7 convention) — `np.linalg.matrix_rank`'s looser default tolerance previously disagreed with the solver on NEAR-collinear controls (`x2 = x + 1e-10`), so the gate could count a direction the solver drops, turn the interacted design off, and change the ATT relative to the identified single-control fit. Near-collinear invariance + replay coherence pinned. +- **Note (review round 18):** the common-timing time-scale contract (Period rejected for detrend/detrendq; trend/seasonal transforms require numeric/datetime/Period time) lives in one shared validator called by BOTH `fit()` and `get_transformation_diagnostics()` (diagnostics previously reached the transforms' raw float-conversion errors). `randomization_inference` validates array shapes/lengths BEFORE the non-finite-outcome filter (a mismatched length combined with a non-finite y previously raised a raw boolean-index IndexError). RI citations point at the LW 2026 small-sample paper (the 2025 Section-5 reference concerned detrending, not RI), and the api-docs no longer call RI "assumption-free" (it does not require normality, conditional on the complete-randomization assignment mechanism). +- **Note (review round 19: family-consistent multiplier contributions — deliberate, externally validated):** the influence contributions feeding the event-study multiplier bootstrap are NORMALIZED TO THE REQUESTED ANALYTICAL VARIANCE FAMILY (classical: per-cell scalar rescale to the classical magnitude; hc1/CR1: the small-sample factor; hc2/hc3: leverage adjustment), not the raw Appendix E.2 contributions. Consequences: per-cell SCALAR adjustments (classical/hc1/CR1) leave the sup-t critical value INVARIANT (draws and SEs scale together and the normalized statistic cancels the factor) while the per-event bootstrap SEs report magnitudes consistent with the requested family rather than the raw asymptotic form — a deliberate coherence choice, so a fit's analytical and bootstrap surfaces answer in the same family. External validation: the RA/hc1 configuration's multiplier-bootstrap SEs are gated against the AUTHORS' Stata package's high-B multiplier bootstrap within the Monte-Carlo bound (acceptance suite, `test_walmart_eventstudy_se_vs_stata`). PSM continues under a rank-deficient (finite-probability) propensity fit exactly like ipw/dr — matching needs only the probabilities — with the regression-point fail-closed fallback reserved for genuinely non-finite propensity fits. +- **Note (review round 20):** point-only PSM is EXEMPT from the common-timing exact-OLS residual-df guard (its inference is NaN by contract and `df_inference=None`; the guard previously rejected valid matching fits whose nominal propensity width exhausted an OLS df count PSM never uses — the staggered path already retained the point). `get_transformation_diagnostics` rejects an all-never-treated staggered panel ("No treated cohorts found", matching `fit_staggered`) instead of returning an empty `by_cohort` that read as success. Non-numeric, non-datetime common-timing time columns must be ORDERED CATEGORICALS (encoded to their codes before any comparison; plain object labels are rejected — lexicographic order breaks at 'Q10' < 'Q2', silently corrupting the pre/post partition), replacing the former plain-string demean acceptance; declared-order chronology pinned on a Q1..Q10 zero-effect trend panel. +- **Note (review round 21):** the NEW LWDiD surface fails closed for `vcov_type='hc2'` at leverage-one designs (warning + NaN inference, point retained), mirroring hc3 — the SHARED hc2/hc2_bm kernel keeps its released `1 - h` floor for the pre-existing estimators pending the tracked family decision (TODO row), so the fabricated-variance path is unreachable from LWDiD while released surfaces are unchanged. The common-timing HEADLINE and unit-bootstrap SEs run through the same scale-equivariant degenerate-SE guard as the staggered/event surfaces (an exactly fitted panel previously reported se ~ 1e-16 with t ~ 1e16). LWDiD plots render the FITTED interval endpoints (per-row t/df, fitted alpha, cband when present; sensitivity specs now carry `conf_int`) instead of a fabricated normal-theory `+/-1.96*SE`; inference-unavailable rows keep the omit-interval rule. PSM docstrings describe 1:`n_neighbors` matching (1:1 default). +- **Note (review round 23):** staggered overall/cohort masses count treated units CONTRIBUTING to each cohort's estimable post cells on every route (a raw cohort member with no estimable cell no longer raises its cohort's weight in `.att`, its combined influence function, or `cohort_effects[g]['n_treated']` — previously only the tau_omega drops route recomputed masses); pinned by a contributing-unit oracle (4/5 vs the raw 4/8 weighting on a mostly-unobserved cohort). `plot_cohort_trends(cohort=)` is IMPLEMENTED (one trajectory per treated cohort, never-treated control line, per-cohort onset markers — previously the parameter was accepted and silently ignored). `validate_staggered_data` marks missing unit/time values as ERRORS (fit rejects the same frame; warning-only let `valid: True` disagree with fit). Input-contract docs state the unit-constant covariate/cluster rule applies to BOTH timing paths and add `_lwdid_season` to the reserved-name list. +- **Note (review round 24, final):** verdict "Looks good — no unmitigated P0 or P1 findings"; the three P2 nits are resolved: `randomization_inference` requires `n_reps >= 10` up front (the reliable-inference floor made smaller values fail after the permutation loop with a misleading hint); the HC3 leverage-one fail-closed path honors the `return_dof` contract with a length-k NaN vector (was `None`); datetime cohort positions relabel to the CANONICAL observed period (two raw between-period labels mapping to the same onset previously collided with row-order-dependent survivor). +- **Note (round-2 refutations, evidence-anchored):** two reviewer claims were checked and REFUTED by execution: (1) pre-treatment placebo transformations — the implementation applies one per-cohort transformation over the full `t < g` pre window with anchor exclusions and the D.3 placebo control pools, and matches the authors' Stata `lwdid` 2.4.2 at full precision (~1e-9) on every Walmart placebo cell `r in [-22, -3]` (the fail-closed label-set gate pins the surface), so the horizon-specific future-window reading is not what the reference implementation does; (2) the IPW influence function's `p_bar = n_1/n` normalization (Lunceford-Davidian linearization of the Hajek ATT) was compared against the proposed finite-sample `B_hat = sum_ctrl(w)/n` variant by Monte Carlo (400 reps, strong propensity heterogeneity): the variants are first-order equivalent and `B_hat` calibrated no better (SE/SD 0.865 vs 0.873), so the implemented convention stands. - Alternatives: HC3 when there are "at least a handful" of treated units; randomization inference for the sharp null (two-sided p = c / #permutations; Stata `lwdid` `ri` option); higher-level clustering and Conley SHAC SEs for larger cross sections (LW 2026 Sec. 8.2, citing Abadie-Athey-Imbens-Wooldridge 2023). -- **Note (IPWRA variance forms, measured divergence):** the PR #588 implementation's IPW/IPWRA influence functions are AIPW/Lunceford-Davidian-style, NOT the papers' E.2-E.4 stacked forms that the authors' Stata package implements. Measured 2026-08-16 on the Walmart application: the IPWRA multiplier-bootstrap SEs diverge SYSTEMATICALLY from the package's (~15% at event-time level, far beyond Monte-Carlo bounds), while the RA config's SEs agree within the MC bound and IPWRA POINT estimates agree to ~1e-3 (logit-optimizer paths). The validation suite therefore gates bootstrap-SE parity on the RA config only; the Stata IPWRA SE columns are committed as provenance in `lwdid_stata_golden.json`, and the E.3-form adjudication (implement the stacked IF, or document-and-anchor the AIPW alternative) is a required item of the contribution's final round — the checklist's E.2/E.3/E.4 box stays unchecked until it resolves. +- **Note (IPWRA variance forms, measured divergence — ADJUDICATED, final round):** the PR #588 implementation's IPW/IPWRA influence functions are AIPW/Lunceford-Davidian-style, NOT the papers' E.2-E.4 stacked forms that the authors' Stata package implements. Measured 2026-08-16 on the Walmart application: the IPWRA multiplier-bootstrap SEs diverge SYSTEMATICALLY from the package's (~15% at event-time level, far beyond Monte-Carlo bounds), while the RA config's SEs agree within the MC bound and IPWRA POINT estimates agree to ~1e-3 (logit-optimizer paths). The validation suite therefore gates bootstrap-SE parity on the RA config only; the Stata IPWRA SE columns are committed as provenance in `lwdid_stata_golden.json`. **Adjudication (2026-08-16, maintainer offered both routes):** the DR (IPWRA) variance KEEPS the AIPW influence-function form as a documented, independently anchored alternative to the papers' E.3 stacked form. The AIPW efficient influence function (Lunceford & Davidian 2004) is the standard doubly-robust EIF in the causal-inference literature and carries its own anchors: the RA config's bootstrap-SE parity gate against the Stata golden, the ~1e-3 point-estimate agreement on every Walmart IPWRA config, and the suite's analytical/bootstrap cross-path pins. Implementing the E.3 stacked IF (and then gating the four remaining Walmart IPWRA SE columns against the golden) remains available as a follow-up if package-form SE parity is later preferred; the checklist's E.2/E.3/E.4 row records the adjudicated scope. - **Note (RI convention):** the paper states `p = c / #permutations`, but the authors' own package (Stata `lwdid` v2.4.2, measured 2026-08-15) implements the INCLUSIVE Phipson-Smyth rule — Monte-Carlo shuffles of the treatment vector with `p = (#{|coef| >= |b0|} + 1) / (reps + 1)`, ties counted as extreme — converging to ~0.0508 on the Prop 99 detrend application at 100k reps. The paper's printed RI p = 0.020 is NOT reproducible with the package (~4.5 binomial SD away) and is recorded as an as-printed discrepancy; the implementation and the validation goldens follow the package convention. The paper reviews (`docs/methodology/papers/lee-wooldridge-2026-review.md`) remain paper-faithful and state the paper's c/N convention as printed. - **Note:** Conley SHAC SEs (listed above as a paper alternative) are NOT offered by the implementation — LWDiD exposes no spatial-coordinate inputs; the vcov design keeps to the house `linalg.py` vocabulary. +*Clustering-level guidance (advisory):* +- Choosing the clustering level for the collapsed cross-sectional regression follows the Cameron & Miller (2015) rule of thumb: cluster at the highest aggregation level that still has enough clusters (G >= 20); if every candidate level has G < 20, use the level with the most clusters and prefer wild cluster bootstrap (Webb weights) over analytical cluster-robust SEs. +- Sensitivity check: compare wild-cluster-bootstrap SEs across candidate levels (unit, state, region, ...); if the max/min SE ratio exceeds ~2x, results are sensitive to the clustering choice and the coarser level should be reported alongside a caveat. +- This guidance is documentation-level only: run `wild_cluster_bootstrap` (in `diff_diff.lwdid_wild_bootstrap`) per candidate level and compare. A dedicated `diagnose_clustering` helper module was removed as out of scope for the estimator API. + *Edge cases:* - Anchor periods: event-study omits `r = -1` (demeaning) / `r = -2, -1` (detrending); bootstrap excludes them. (Raw-effect exclusion; the public results surface emits observed anchors as `is_reference` rows — see the Standard errors note.) - All units eventually treated (LW 2025 Sec. 4.3): drop `D_infinity`; effects defined relative to the last cohort; no effect estimable for the last cohort. @@ -2593,21 +2630,21 @@ Event-study/placebo transformations over ALL periods (Appendix D): demeaning (D. - Walmart entry (LW 2025 Tables A4/A5, 1,277 counties): per-relative-period WATT(r) with SEs for r = 0..13. **Requirements checklist:** -- [ ] Rolling demeaning (3.2)/(4.11) using ALL pre-g periods; detrending (5.6)/(D.2) via unit OLS with out-of-sample residuals -- [ ] Minimum pre-period enforcement (>= 1 demeaning / >= 2 detrending); failing cells dropped with warning -- [ ] Control pools: NT + NYT (`A_{r+1} = 1`) default, NT-only option (`N_infinity >= 2` guard); placebo cells per D.3 rule `G > max(g,t)` -- [ ] RA (E.1) with treated-cohort-centered interactions; IPWRA (logit + WLS); IPW special case -- [ ] Influence functions per E.2/E.3/E.4 including first-stage logit-score corrections; IFs centered -- [ ] WATT(r) event-study aggregation with contributing-treated-unit weights (E.1 definition; cohort-size weights `N_g / N_{G_r}` only as the balanced-panel simplification); anchor periods excluded (r = -1 / r = -2,-1) -- [ ] Algorithm 1 multiplier bootstrap: unit-level Rademacher, sup-t simultaneous bands -- [ ] Composite-outcome overall aggregation (7.18)/(7.19) with cohort-share weights -- [ ] Exact-t inference: `T_{N-2}` / `T_{N-K-2}`, valid to `N = 3`, `N1 = 1`, `N_g = 1`; sample-size guards enforced -- [ ] HC3 alternative; randomization inference (paper convention p = c / #permutations — implemented per the authors'-package inclusive convention, see the RI note); higher-level clustering per the vcov design (SHAC/Conley not offered, see note) -- [ ] Anticipation-robustness period dropping; seasonal dummies in the transformation step +- [x] Rolling demeaning (3.2)/(4.11) using ALL pre-g periods; detrending (5.6)/(D.2) via unit OLS with out-of-sample residuals +- [x] Minimum pre-period enforcement (>= 1 demeaning / >= 2 detrending); failing cells dropped with warning +- [x] Control pools: NT + NYT (`A_{r+1} = 1`) default, NT-only option (`N_infinity >= 2` guard); placebo cells per D.3 rule `G > max(g,t)` +- [x] RA (E.1) with treated-cohort-centered interactions; IPWRA (logit + WLS); IPW special case +- [x] Influence functions per E.2/E.3/E.4 including first-stage logit-score corrections; IFs centered — adjudicated scope: RA follows E.2; IPW/IPWRA ship the AIPW-EIF form as a documented, independently anchored alternative to E.3/E.4 (see the IPWRA-variance note; E.3 stacked form remains an available follow-up) +- [x] WATT(r) event-study aggregation with contributing-treated-unit weights (E.1 definition; cohort-size weights `N_g / N_{G_r}` only as the balanced-panel simplification); anchor periods excluded (r = -1 / r = -2,-1) +- [x] Algorithm 1 multiplier bootstrap: unit-level Rademacher, sup-t simultaneous bands +- [x] Composite-outcome overall aggregation (7.18)/(7.19) with cohort-share weights +- [x] Exact-t inference: `T_{N-2}` / `T_{N-K-2}`, valid to `N = 3`, `N1 = 1`, `N_g = 1`; sample-size guards enforced +- [x] HC3 alternative; randomization inference (paper convention p = c / #permutations — implemented per the authors'-package inclusive convention, see the RI note); higher-level clustering per the vcov design (SHAC/Conley not offered, see note) +- [x] Anticipation-robustness period dropping; seasonal dummies in the transformation step - [ ] All-eventually-treated (Sec. 4.3) support — deferred by decision (implementation rejects; see the Edge cases note and the DEFERRED.md row) -- [ ] Unbalanced-panel (Sec. 4.4) support -- [ ] Common-timing no-covariate case reproduces plain DiD (3.4); Theorem 3.1 pooled-OLS equivalence (cross-estimator test vs `DifferenceInDifferences` / ETWFE at r = g) -- [ ] Prop 99 / castle-laws / Walmart replication targets pinned as tests +- [x] Unbalanced-panel (Sec. 4.4) support +- [x] Common-timing no-covariate case reproduces plain DiD (3.4); Theorem 3.1 pooled-OLS equivalence (cross-estimator test vs `DifferenceInDifferences` / ETWFE at r = g) +- [x] Prop 99 / castle-laws / Walmart replication targets pinned as tests --- @@ -5725,7 +5762,7 @@ Degrees of freedom for the t-distribution lookup use `ResolvedSurveyDesign.df_su - `survey_design=` for `vcov_type ∈ {"hc1"}` (plus `cluster=` for CR1) SHIPPED in Wave E.1 — see "Variance (Wave E.1)" subsection below. Threads Hájek-normalized survey weights through stage-1 FE estimation, gamma_hat solve, eps construction, and bread inversion; aggregates the Wave D Psi to PSU totals and routes through the audited `_compute_stratified_meat_from_psu_scores` Binder TSL meat helper. `vcov_type="conley"` combined with `survey_design=` SHIPPED in Wave E.2 for cross-sectional Conley (`conley_lag_cutoff = 0`) — see "Variance (Wave E.2)" subsection below (stratified-Conley sandwich on PSU totals). Wave E.2 follow-up adds the panel-block composition (`conley_lag_cutoff > 0`) via spatial + serial Bartlett HAC — see "Variance (Wave E.2 follow-up)" subsection below. `SurveyDesign.subpopulation()` and warn-and-drop full-design retention via zero-pad scores SHIPPED in Wave E.3 — see "Variance (Wave E.3)" subsection below (matches R `survey::svyrecvar(subset())` + in-library precedent at `imputation.py:2175-2183` and `prep.py:1401-1432`). Replicate-weight variance (BRR / Fay / JK1 / JKn / SDR) raises `NotImplementedError` — Gerber (2026) Appendix A notes the IF-reweighting shortcut does NOT apply to TwoStageDiD-class estimators because `gamma_hat` is weight-sensitive; correct support requires per-replicate full re-fit and is queued as a follow-up. - `covariates=` raises `NotImplementedError` — Gardner-style stage-1 residualization not yet wired through; planned follow-up. - `ring_method="count"` not exposed — only the nearest-treated-ring specification. -- `vcov_type` ∈ {`"hc2"`, `"hc2_bm"`, `"classical"`} raises `NotImplementedError` — `hc2`/`hc2_bm` because current stage-2 inference uses generic residual df rather than per-coefficient Bell-McCaffrey / CR2 DOF; `classical` because the Wave D Gardner GMM first-stage correction has not been derived for the classical homoskedastic variance (different meat structure `sigma_hat^2 * (X_10' X_10)` vs the Wave D IF outer product `Psi' Psi`). Use `"hc1"` or `"conley"`, or pair with `cluster=` for CR1 — all three apply the Wave D GMM correction. +- `vcov_type` ∈ {`"hc2"`, `"hc2_bm"`, `"hc3"`, `"classical"`} raises `NotImplementedError` — `hc2`/`hc2_bm` (at construction) because current stage-2 inference uses generic residual df rather than per-coefficient Bell-McCaffrey / CR2 DOF; `hc3` (at construction) because it is not implemented for the two-stage spillover variance; `classical` (at fit) because the Wave D Gardner GMM first-stage correction has not been derived for the classical homoskedastic variance (different meat structure `sigma_hat^2 * (X_10' X_10)` vs the Wave D IF outer product `Psi' Psi`). Use `"hc1"` or `"conley"`, or pair with `cluster=` for CR1 — all three apply the Wave D GMM correction. - **`rings[0]` must equal 0** — the partition must cover treated locations (`d_it = 0` belongs to Ring 1). Rings starting at a nonzero inner edge would leave units in `0 <= d_it < rings[0]` as exposed-but-unmodeled, silently biasing the estimator. Validator rejects such inputs. - **Balanced panel required (Wave B MVP)** — every unit must observe every period. An unbalanced (unit, time) Ω₀ bipartite graph can produce disconnected FE components and unidentified stage-1 residuals on treated rows. Exact graph-connectivity-based identification (which would relax this to a strictly weaker condition) is queued as a follow-up extension. Validator rejects unbalanced inputs. - **One row per `(unit, time)` cell required** — duplicate cells silently re-weight stage-1 FE estimation AND stage-2 OLS. Validator rejects duplicate cells. diff --git a/docs/methodology/variance-conventions.md b/docs/methodology/variance-conventions.md index 0d28c953..1d08045f 100644 --- a/docs/methodology/variance-conventions.md +++ b/docs/methodology/variance-conventions.md @@ -137,7 +137,7 @@ output). since the M-021 migration — and carries the K_reference increment there (+6, the [time, unit] no-intercept increment on df_0 — pinned via expected_adjustment on its matrix row). -- **L4 — hc2/hc2_bm** (leverage / Satterthwaite DOF — no CR1 factor), +- **L4 — hc2/hc2_bm/hc3** (leverage / Satterthwaite DOF — no CR1 factor; hc3 squares the leverage denominator and is one-way only; under `fweight` the leverage is each replicate row's UNWEIGHTED quadratic form against the weighted bread — frequency weights are replicated data, and compressed HC2/HC3 equal literal expansion exactly — while aweight/pweight keep the WLS-hat `w_i x_i'(X'WX)^{-1}x_i` convention), **survey TSL** (n_PSU - n_strata over the full design), and **Wooldridge cohort_trends full-dummy** (documented opt-in landing on the L1 convention). **conley** is out of this matrix by decision: the spatial-HAC diff --git a/docs/practitioner_decision_tree.rst b/docs/practitioner_decision_tree.rst index 0cb670aa..2f623c13 100644 --- a/docs/practitioner_decision_tree.rst +++ b/docs/practitioner_decision_tree.rst @@ -520,6 +520,19 @@ staggered approaches, Local Projections DiD, Stacked DiD, Efficient DiD, Triple Difference, TROP, Changes-in-Changes for distributional/quantile effects, and more. The six scenarios above cover the most common business use cases. +- **Want rolling-transformation approach?** → :class:`~diff_diff.LWDiD` (Lee & Wooldridge 2025, 2026) + + Converts panel data into cross-sectional estimation via unit-specific demeaning + or detrending of pre-treatment outcomes (``rolling='detrend'`` for + heterogeneous linear trends). Supports RA (``vcov_type`` in + ``classical``/``hc1``/``hc2``/``hc3``), IPW, and IPWRA estimators + (influence-function variance, ``hc1``) with cluster-robust inference + via ``cluster=`` on those paths; PSM provides point estimates only + (inference is NaN pending an Abadie-Imbens matching variance, and + ``cluster=`` is rejected). Works for both common + timing and staggered adoption designs. Compare ``rolling='demean'`` vs + ``rolling='detrend'`` as a built-in specification robustness check. + For the full academic decision tree with all estimators, see :doc:`choosing_estimator`. diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst index e91c47da..1945deb2 100644 --- a/docs/tutorials/index.rst +++ b/docs/tutorials/index.rst @@ -238,6 +238,7 @@ Modern estimators for designs the basic toolkit cannot handle. Sharp and fuzzy RD from plot to estimate, when a naive cutoff comparison overstates the effect fivefold. + .. toctree:: :maxdepth: 1 :caption: Advanced Methods diff --git a/tests/conftest.py b/tests/conftest.py index 06c54118..f1c6086a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -261,3 +261,9 @@ def assert_nan_inference(inference_dict): ci = inference_dict["conf_int"] assert np.isnan(ci[0]), f"ci_lower should be NaN when SE={se}, got {ci[0]}" assert np.isnan(ci[1]), f"ci_upper should be NaN when SE={se}, got {ci[1]}" + + +@pytest.fixture +def require_lwdid(): + """Skip test if lwdid package not installed (optional for equivalence tests).""" + pytest.importorskip("lwdid", reason="lwdid package required for equivalence tests") diff --git a/tests/test_estimators_vcov_type.py b/tests/test_estimators_vcov_type.py index b494a765..9ba7709d 100644 --- a/tests/test_estimators_vcov_type.py +++ b/tests/test_estimators_vcov_type.py @@ -14,6 +14,7 @@ from __future__ import annotations +import inspect import warnings import numpy as np @@ -74,7 +75,7 @@ def test_robust_false_explicit_hc2_raises(self): def test_unknown_vcov_type_raises(self): with pytest.raises(ValueError, match="vcov_type must be one of"): - DifferenceInDifferences(vcov_type="hc3") + DifferenceInDifferences(vcov_type="hc9") def test_hc0_not_accepted(self): for bad in ("hc0", "HC1", "CR2", "cr1", "hc2+bm"): @@ -132,7 +133,7 @@ def test_set_params_robust_only_rederives_vcov_type(self): def test_set_params_invalid_vcov_type_rejected(self): est = DifferenceInDifferences() with pytest.raises(ValueError, match="vcov_type must be one of"): - est.set_params(vcov_type="hc3") + est.set_params(vcov_type="hc9") def test_set_params_robust_true_then_back_to_hc1(self): """robust=True after construction restores hc1 when no explicit vcov_type.""" @@ -3204,3 +3205,119 @@ def test_validation_accepts_normal_everywhere(self): for ctor in (DifferenceInDifferences, TwoWayFixedEffects, MultiPeriodDiD): est = ctor(df_convention="normal") assert est.get_params()["df_convention"] == "normal" + + +class TestHC3SharedSurfaceHardening: + """LWDiD fix-wave WS7: the PR's linalg widening admitted "hc3" to the + shared vcov vocabulary, but pre-existing per-estimator guard lists let + it escape. Campaign finding (execution-verified): hc3 + absorb= computed + leverage on the within-transformed reduced design, silently understating + DiD/MP-DiD SEs ~11% vs the full-dummy computation; TWFE crashed with a + misleading "hc3 is one-way only ... cluster-robust" error for users who + never passed cluster=. + """ + + def _staggered_fe_panel(self, seed: int = 20260819) -> pd.DataFrame: + rng = np.random.default_rng(seed) + rows = [] + for u in range(24): + alpha = rng.normal(0, 1.0) + treated_unit = u < 12 + for t in range(6): + treat = int(treated_unit and t >= 3) + y = alpha + 0.3 * t + 1.5 * treat + rng.normal(0, 1.0) + rows.append( + { + "unit": u, + "time": t, + "treated": int(treated_unit), + "post": int(t >= 3), + "y": y, + } + ) + return pd.DataFrame(rows) + + def test_did_hc3_absorb_matches_full_dummy(self): + # hc3 must route absorb= through the full-dummy design exactly like + # hc2/hc2_bm (leverage families need the FULL FE projection). + df = self._staggered_fe_panel() + r_absorb = DifferenceInDifferences(vcov_type="hc3").fit( + df, outcome="y", treatment="treated", post="post", absorb=["unit"] + ) + r_fe = DifferenceInDifferences(vcov_type="hc3").fit( + df, outcome="y", treatment="treated", post="post", fixed_effects=["unit"] + ) + np.testing.assert_allclose(r_absorb.att, r_fe.att, rtol=1e-10) + np.testing.assert_allclose(r_absorb.se, r_fe.se, rtol=1e-10) + + def test_mpd_hc3_absorb_matches_full_dummy(self): + df = self._staggered_fe_panel() + kw = dict(outcome="y", treatment="treated", time="time", post_periods=[3, 4, 5]) + r_absorb = MultiPeriodDiD(vcov_type="hc3").fit(df, absorb=["unit"], **kw) + r_fe = MultiPeriodDiD(vcov_type="hc3").fit(df, fixed_effects=["unit"], **kw) + for period in r_absorb.period_effects: + np.testing.assert_allclose( + r_absorb.period_effects[period].effect, + r_fe.period_effects[period].effect, + rtol=1e-10, + ) + np.testing.assert_allclose( + r_absorb.period_effects[period].se, + r_fe.period_effects[period].se, + rtol=1e-10, + ) + + def test_twfe_hc3_fits_without_misleading_cluster_error(self): + # Pre-fix: explicit hc3 kept the unit auto-cluster, and solve_ols + # raised "hc3 is one-way only ... for cluster-robust" although the + # user never passed cluster=. + df = self._staggered_fe_panel() + df["treat_it"] = df["treated"] * df["post"] + res = TwoWayFixedEffects(vcov_type="hc3").fit( + df, outcome="y", treatment="treat_it", unit="unit", time="time" + ) + assert np.isfinite(res.att) and np.isfinite(res.se) and res.se > 0 + + def test_roster_every_vcov_estimator_supports_or_rejects_hc3(self): + # Structural guard: the next vcov-vocabulary widening must not + # escape a sibling's hardcoded list. Every BaseEstimator-roster + # class that accepts the linalg vcov vocabulary (constructs with + # vcov_type="hc1") must either be on the known-support allowlist or + # reject vcov_type="hc3" with an informative error at construction. + from tests.test_base_estimator import DEFAULT_KWARGS, MIXIN_CLASSES + + # Foreign vcov_type vocabularies (not the linalg family namespace): + # RDDensityTest uses {jackknife, plugin}; RegressionDiscontinuity + # uses the rdrobust vce namespace. + FOREIGN_VOCAB = {"RDDensityTest", "RegressionDiscontinuity"} + HC3_SUPPORTED = { + "DifferenceInDifferences", + "TwoWayFixedEffects", + "MultiPeriodDiD", + "LWDiD", # reg path only; ipw/dr/psm restricted by its own validator + } + checked = [] + for cls in MIXIN_CLASSES: + name = cls.__name__ + if name in FOREIGN_VOCAB: + continue + base_kwargs = dict(DEFAULT_KWARGS.get(name, {})) + sig = inspect.signature(cls.__init__) + if "vcov_type" not in sig.parameters: + continue + try: + cls(vcov_type="hc1", **base_kwargs) + except (ValueError, NotImplementedError, TypeError): + # hc1 itself not accepted at construction -> out of scope + continue + checked.append(name) + if name in HC3_SUPPORTED: + cls(vcov_type="hc3", **base_kwargs) # must construct cleanly + else: + with pytest.raises((ValueError, NotImplementedError)) as exc_info: + cls(vcov_type="hc3", **base_kwargs) + assert "hc3" in str( + exc_info.value + ), f"{name} rejected hc3 without naming it: {exc_info.value}" + # The guard must actually be exercising a meaningful roster. + assert len(checked) >= 8, f"roster unexpectedly small: {checked}" diff --git a/tests/test_linalg_hc2_bm.py b/tests/test_linalg_hc2_bm.py index ebb01d7d..95f9b0ef 100644 --- a/tests/test_linalg_hc2_bm.py +++ b/tests/test_linalg_hc2_bm.py @@ -369,13 +369,14 @@ def test_unknown_vcov_type_raises(self, small_ols_dataset): X, y = small_ols_dataset _, resid, _ = _fit_unweighted(X, y) with pytest.raises(ValueError, match="vcov_type must be one of"): - compute_robust_vcov(X, resid, vcov_type="hc3") + compute_robust_vcov(X, resid, vcov_type="hc9") def test_hc0_not_accepted(self, small_ols_dataset): - """HC0/HC3/CR0 are out of scope for Phase 1a.""" + """HC0/CR0 are out of scope for Phase 1a (HC3 joined the valid set + for the LWDiD canonical-vocabulary rename).""" X, y = small_ols_dataset _, resid, _ = _fit_unweighted(X, y) - for bad in ("hc0", "hc3", "cr0"): + for bad in ("hc0", "cr0"): with pytest.raises(ValueError, match="vcov_type must be one of"): compute_robust_vcov(X, resid, vcov_type=bad) diff --git a/tests/test_lwdid.py b/tests/test_lwdid.py new file mode 100644 index 00000000..4a8514d9 --- /dev/null +++ b/tests/test_lwdid.py @@ -0,0 +1,4603 @@ +"""Tests for LWDiD estimator (Lee & Wooldridge 2025, 2026).""" + +import json +import warnings + +import numpy as np +import pandas as pd +import pytest + +from diff_diff import LWDiD, LWDiDResults + +# ─── Test Data Generators ─────────────────────────────────────────────────── + + +def _make_common_timing_panel( + n_treated=30, + n_control=50, + n_pre=5, + n_post=3, + true_att=2.0, + seed=42, +): + """Generate balanced common-timing panel with known ATT. + + Pre-treatment periods: 1..n_pre (treatment=0 for all) + Post-treatment periods: n_pre+1..n_pre+n_post (treatment=1 for treated) + """ + rng = np.random.default_rng(seed) + n_units = n_treated + n_control + n_periods = n_pre + n_post + + rows = [] + for i in range(n_units): + is_treated = i < n_treated + unit_fe = rng.normal(0, 1) + for t in range(1, n_periods + 1): + time_trend = 0.3 * t + noise = rng.normal(0, 0.5) + post = 1 if t > n_pre else 0 + treat = 1 if (is_treated and post) else 0 + y = unit_fe + time_trend + noise + (true_att if treat else 0) + rows.append( + { + "unit": i, + "time": t, + "y": y, + "treat": treat, + } + ) + return pd.DataFrame(rows) + + +def _make_staggered_panel( + n_units=120, + n_periods=10, + n_cohorts=3, + true_att=1.5, + seed=42, +): + """Generate staggered adoption panel with multiple cohorts. + + Cohort assignment: + - First ~1/4 units: never-treated (cohort=0) + - Remaining units split across n_cohorts with treatment times spread. + """ + rng = np.random.default_rng(seed) + n_never = n_units // 4 + n_per_cohort = (n_units - n_never) // n_cohorts + + # Cohort adoption times (spread across middle periods) + cohort_times = [3 + i * 2 for i in range(n_cohorts)] + + rows = [] + uid = 0 + for i in range(n_never): + unit_fe = rng.normal(0, 1) + for t in range(1, n_periods + 1): + y = unit_fe + 0.2 * t + rng.normal(0, 0.5) + rows.append( + { + "unit": uid, + "time": t, + "y": y, + "treat": 0, + "cohort": 0, + } + ) + uid += 1 + + for c_idx, g in enumerate(cohort_times): + for i in range(n_per_cohort): + unit_fe = rng.normal(0, 1) + for t in range(1, n_periods + 1): + post = 1 if t >= g else 0 + treat = post # treated once cohort adopts + effect = true_att * post + y = unit_fe + 0.2 * t + rng.normal(0, 0.5) + effect + rows.append( + { + "unit": uid, + "time": t, + "y": y, + "treat": treat, + "cohort": g, + } + ) + uid += 1 + + return pd.DataFrame(rows) + + +# ─── Parameter Interface Tests ────────────────────────────────────────────── + + +class TestLWDiDParams: + """Test parameter setting, getting, and validation.""" + + def test_get_params_returns_all(self): + est = LWDiD(rolling="demean", estimation_method="reg", vcov_type="hc1") + params = est.get_params() + assert "rolling" in params + assert "estimation_method" in params + assert "vcov_type" in params + assert "control_group" in params + assert "alpha" in params + assert "n_bootstrap" in params + assert params["rolling"] == "demean" + assert params["estimation_method"] == "reg" + assert params["vcov_type"] == "hc1" + + def test_set_params_modifies(self): + est = LWDiD() + est.set_params(rolling="detrend") + assert est.rolling == "detrend" + + def test_set_params_returns_self(self): + est = LWDiD() + ret = est.set_params(estimation_method="ipw") + assert ret is est + + def test_invalid_rolling_raises(self): + with pytest.raises(ValueError, match="rolling"): + LWDiD(rolling="invalid") + + def test_invalid_estimation_method_raises(self): + with pytest.raises(ValueError, match="estimation_method"): + LWDiD(estimation_method="invalid") + + def test_invalid_vcov_type_raises(self): + with pytest.raises(ValueError, match="vcov_type"): + LWDiD(vcov_type="invalid") + + def test_invalid_control_group_raises(self): + with pytest.raises(ValueError, match="control_group"): + LWDiD(control_group="invalid") + + def test_invalid_alpha_raises(self): + with pytest.raises(ValueError, match="alpha"): + LWDiD(alpha=0.0) + with pytest.raises(ValueError, match="alpha"): + LWDiD(alpha=1.0) + + def test_invalid_n_bootstrap_raises(self): + with pytest.raises(ValueError, match="n_bootstrap"): + LWDiD(n_bootstrap=-1) + + def test_LW_alias_removed(self): + import diff_diff + + assert not hasattr(diff_diff, "LW") + assert "LW" not in diff_diff.__all__ + + def test_default_params(self): + est = LWDiD() + assert est.rolling == "demean" + assert est.estimation_method == "reg" + assert est.vcov_type == "hc1" + assert est.control_group == "not_yet_treated" + assert est.alpha == 0.05 + assert est.n_bootstrap == 0 + + def test_repr(self): + est = LWDiD(rolling="demean", estimation_method="reg") + r = repr(est) + assert "LWDiD" in r + assert "demean" in r + assert "reg" in r + + def test_set_params_invalid_key_raises(self): + est = LWDiD() + with pytest.raises(ValueError, match="Unknown parameter"): + est.set_params(bad_param="x") + + +# ─── Input Validation Tests ───────────────────────────────────────────────── + + +class TestLWDiDInputValidation: + """Test input data validation.""" + + def test_missing_column_raises(self): + df = pd.DataFrame({"unit": [1], "time": [1], "y": [1.0]}) + with pytest.raises(ValueError, match="Columns not found"): + LWDiD().fit(df, outcome="y", unit="unit", time="time", treatment="treat") + + def test_nan_in_outcome_raises(self): + df = pd.DataFrame( + { + "unit": [1, 1, 2, 2], + "time": [1, 2, 1, 2], + "y": [1.0, np.nan, 2.0, 3.0], + "treat": [0, 1, 0, 0], + } + ) + with pytest.raises(ValueError, match="missing values"): + LWDiD().fit(df, outcome="y", unit="unit", time="time", treatment="treat") + + def test_nan_in_treatment_raises(self): + df = pd.DataFrame( + { + "unit": [1, 1, 2, 2], + "time": [1, 2, 1, 2], + "y": [1.0, 2.0, 2.0, 3.0], + "treat": [0, np.nan, 0, 0], + } + ) + with pytest.raises(ValueError, match="missing values"): + LWDiD().fit(df, outcome="y", unit="unit", time="time", treatment="treat") + + def test_duplicate_unit_time_raises(self): + df = pd.DataFrame( + { + "unit": [1, 1, 1, 2], + "time": [1, 1, 2, 1], + "y": [1.0, 1.5, 2.0, 3.0], + "treat": [0, 0, 1, 0], + } + ) + with pytest.raises(ValueError, match="duplicate"): + LWDiD().fit(df, outcome="y", unit="unit", time="time", treatment="treat") + + def test_non_binary_treatment_raises(self): + df = pd.DataFrame( + { + "unit": [1, 1, 2, 2], + "time": [1, 2, 1, 2], + "y": [1.0, 2.0, 3.0, 4.0], + "treat": [0, 2, 0, 0], # not binary + } + ) + with pytest.raises(ValueError): + LWDiD().fit(df, outcome="y", unit="unit", time="time", treatment="treat") + + def test_vcov_type_cluster_rejected(self): + with pytest.raises(ValueError, match="cluster"): + LWDiD(vcov_type="cluster") + + def test_no_treated_units_raises(self): + df = pd.DataFrame( + { + "unit": [1, 1, 2, 2], + "time": [1, 2, 1, 2], + "y": [1.0, 2.0, 3.0, 4.0], + "treat": [0, 0, 0, 0], + } + ) + with pytest.raises(ValueError, match="[Nn]o treated|[Nn]o post"): + LWDiD().fit(df, outcome="y", unit="unit", time="time", treatment="treat") + + def test_no_control_units_raises(self): + df = pd.DataFrame( + { + "unit": [1, 1, 2, 2], + "time": [1, 2, 1, 2], + "y": [1.0, 2.0, 3.0, 4.0], + "treat": [0, 1, 0, 1], + } + ) + with pytest.raises(ValueError, match="[Nn]o control"): + LWDiD().fit(df, outcome="y", unit="unit", time="time", treatment="treat") + + +# ─── Treatment Design Validation Tests ────────────────────────────────────── + + +def _make_design_panel(cohort_map, n_periods=5, seed=7): + """Small panel (len(cohort_map) units x n_periods) with D_it = 1[t >= g_i]. + + cohort_map: {unit_id: g} with g=0 for never-treated. Returns columns + unit/time/y/treat/cohort so tests can freely corrupt treat or cohort. + """ + rng = np.random.default_rng(seed) + rows = [] + for uid, g in cohort_map.items(): + for t in range(1, n_periods + 1): + treat = int(g > 0 and t >= g) + rows.append( + { + "unit": uid, + "time": t, + "y": rng.normal(0, 0.5) + 1.5 * treat, + "treat": treat, + "cohort": g, + } + ) + return pd.DataFrame(rows) + + +class TestTreatmentDesignValidation: + """Unified vectorized design checks (_check_treatment_design).""" + + @staticmethod + def _cohorts(n_treated_3=5, n_treated_4=5, n_never=10): + cohorts = {} + uid = 0 + for _ in range(n_treated_3): + cohorts[uid] = 3 + uid += 1 + for _ in range(n_treated_4): + cohorts[uid] = 4 + uid += 1 + for _ in range(n_never): + cohorts[uid] = 0 + uid += 1 + return cohorts + + # ── (a) absorbing treatment ── + + def test_non_absorbing_common_timing_raises(self): + panel = _make_design_panel({u: (3 if u < 8 else 0) for u in range(20)}) + # unit 0 switches back to 0 at the last period + panel.loc[(panel["unit"] == 0) & (panel["time"] == 5), "treat"] = 0 + with pytest.raises(ValueError, match="Non-absorbing"): + LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + + def test_non_absorbing_staggered_raises(self): + panel = _make_design_panel(self._cohorts()) + panel.loc[(panel["unit"] == 0) & (panel["time"] == 5), "treat"] = 0 + with pytest.raises(ValueError, match="Non-absorbing"): + LWDiD().fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + + def test_non_absorbing_unsorted_input_raises(self): + """Detection must not depend on the input row order.""" + panel = _make_design_panel({u: (3 if u < 8 else 0) for u in range(20)}) + panel.loc[(panel["unit"] == 0) & (panel["time"] == 4), "treat"] = 0 + shuffled = panel.sample(frac=1.0, random_state=0).reset_index(drop=True) + with pytest.raises(ValueError, match="Non-absorbing"): + LWDiD().fit(shuffled, outcome="y", unit="unit", time="time", treatment="treat") + + # ── (b) common timing: unique onset ── + + def test_heterogeneous_onset_without_cohort_raises(self): + cohorts = {u: 3 for u in range(5)} + cohorts.update({u: 4 for u in range(5, 10)}) + cohorts.update({u: 0 for u in range(10, 20)}) + panel = _make_design_panel(cohorts) + with pytest.raises(ValueError, match="heterogeneous first-treatment"): + LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + + def test_common_timing_valid_passes(self): + panel = _make_design_panel({u: (3 if u < 8 else 0) for u in range(20)}) + res = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + assert np.isfinite(res.att) + + # ── (c) staggered: onset == cohort ── + + def test_onset_cohort_mismatch_raises(self): + panel = _make_design_panel(self._cohorts()) + # unit 0 (cohort 3) starts treatment one period early + panel.loc[(panel["unit"] == 0) & (panel["time"] == 2), "treat"] = 1 + with pytest.raises(ValueError, match="inconsistent with cohort"): + LWDiD().fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + + def test_never_treated_with_treatment_rows_raises(self): + panel = _make_design_panel(self._cohorts()) + # unit 19 is never-treated by cohort but has a treatment=1 row + panel.loc[(panel["unit"] == 19) & (panel["time"] == 5), "treat"] = 1 + with pytest.raises(ValueError, match="inconsistent with cohort"): + LWDiD().fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + + def test_cohort_in_window_never_switching_on_raises(self): + panel = _make_design_panel(self._cohorts()) + # unit 0 keeps cohort=3 but never actually switches on + panel.loc[panel["unit"] == 0, "treat"] = 0 + with pytest.raises(ValueError, match="no\\s+treatment=1 rows"): + LWDiD().fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + + def test_staggered_valid_passes(self): + panel = _make_design_panel(self._cohorts()) + res = LWDiD(control_group="never_treated").fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert np.isfinite(res.att) + + def test_staggered_nan_cohort_never_treated_passes(self): + """Never-treated encoded as NaN cohort is a valid design.""" + panel = _make_design_panel(self._cohorts()) + panel["cohort"] = panel["cohort"].replace(0, np.nan) + res = LWDiD(control_group="never_treated").fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert np.isfinite(res.att) + + def test_cohort_beyond_window_recoded_to_never_treated(self): + """Beyond-window cohorts are recoded to never-treated by the + normalizer (with a warning), then pass the design check as + never-treated units. The design check itself now documents a + normalized-input precondition, so direct callers normalize first. + """ + from diff_diff.lwdid import _check_treatment_design, _normalize_cohorts + + cohorts = self._cohorts() + cohorts[0] = 9 # beyond n_periods=5: all treat rows are 0 + panel = _make_design_panel(cohorts) + with pytest.warns(UserWarning, match="exceed the last observed period"): + panel["cohort"], n_inf, n_beyond = _normalize_cohorts( + panel["cohort"], max_time=panel["time"].max() + ) + assert n_inf == 0 and n_beyond > 0 + assert (panel.loc[panel["unit"] == 0, "cohort"] == 0).all() + # Must not raise: the recoded unit is never-treated with no D=1 rows + _check_treatment_design(panel, "unit", "time", "treat", "cohort") + + +# ─── Transformation Tests ─────────────────────────────────────────────────── + + +class TestLWDiDTransformations: + """Test that rolling transformations are correctly applied.""" + + def test_demean_subtracts_pre_mean(self): + """Construct simple 3-unit panel where pre-means are known. + + (Fix-wave update: the former 2-unit fixture is an INVALID exact + design - 2 collapsed observations for 2 parameters - which the + Registry small-sample guard now rejects; a second control keeps + the hand-computed arithmetic with a positive residual df.) + """ + # Unit 0 (control): y = [2, 4, 6] -> pre_mean = 3, post ydot = 3 + # Unit 2 (control): y = [4, 6, 8] -> pre_mean = 5, post ydot = 3 + # Unit 1 (treated): y = [1, 3, 10] -> pre_mean = 2, post ydot = 8 + df = pd.DataFrame( + { + "unit": [0, 0, 0, 1, 1, 1, 2, 2, 2], + "time": [1, 2, 3, 1, 2, 3, 1, 2, 3], + "y": [2.0, 4.0, 6.0, 1.0, 3.0, 10.0, 4.0, 6.0, 8.0], + "treat": [0, 0, 0, 0, 0, 1, 0, 0, 0], + } + ) + res = LWDiD(rolling="demean", estimation_method="reg").fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + # The method demeaned using pre-treatment periods (time 1,2) + # Unit 0: pre_mean = 3, post (time 3) ydot = 6-3 = 3 + # Unit 1: pre_mean = 2, post (time 3) ydot = 10-2 = 8 + # ATT = 8 - 3 = 5 (treatment effect + any trend difference) + assert isinstance(res, LWDiDResults) + assert np.isfinite(res.att) + + def test_detrend_removes_linear_trend(self): + """Construct unit with perfect linear trend y = 1 + 2*t. + + After detrend, residuals should be ~0 in pre-period. + """ + # Need at least 2 pre periods for detrend; a third unit keeps the + # collapsed design valid (fix-wave Registry small-sample guard). + # Units 0/2 (controls): y = 1 + 2*t (unit 2 offset by +2) + # Unit 1 (treated): y = 1 + 2*t in pre, + 5 in post + df = pd.DataFrame( + { + "unit": [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2], + "time": [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4], + "y": [3.0, 5.0, 7.0, 9.0, 3.0, 5.0, 12.0, 14.0, 5.0, 7.0, 9.0, 11.0], + "treat": [0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0], + } + ) + res = LWDiD(rolling="detrend", estimation_method="reg").fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert isinstance(res, LWDiDResults) + # Detrended control should be ~0, detrended treated should show effect + assert res.att > 0 + + def test_transform_preserves_treatment_effect(self): + """After demean, the treatment effect should still be visible.""" + panel = _make_common_timing_panel(true_att=5.0, seed=123) + res = LWDiD(rolling="demean", estimation_method="reg").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + # True ATT is 5.0, estimate should be positive and in range + assert res.att > 2.0 + + +# ─── Common Timing Tests ──────────────────────────────────────────────────── + + +class TestLWDiDCommonTiming: + """Test common-timing estimation paths.""" + + @pytest.fixture + def panel(self): + return _make_common_timing_panel(true_att=2.0) + + def test_ra_returns_results(self, panel): + est = LWDiD(rolling="demean", estimation_method="reg") + res = est.fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + assert isinstance(res, LWDiDResults) + + def test_ra_demean_positive_att(self, panel): + res = LWDiD(rolling="demean", estimation_method="reg").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert res.att > 0 # True ATT is 2.0 + + def test_ra_detrend_positive_att(self, panel): + res = LWDiD(rolling="detrend", estimation_method="reg").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert res.att > 0 + + def test_ra_att_close_to_truth(self, panel): + """RA demean should recover ATT near 2.0 with enough data.""" + res = LWDiD(rolling="demean", estimation_method="reg").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + # Allow generous tolerance due to small sample noise + assert 0.5 < res.att < 4.0 + + def test_ipw_positive_att(self, panel): + """IPW needs controls for propensity score.""" + panel_with_x = panel.copy() + rng = np.random.default_rng(0) + units = panel_with_x["unit"].unique() + xmap = dict(zip(units, rng.normal(size=len(units)))) + panel_with_x["x1"] = panel_with_x["unit"].map(xmap) # unit-constant + res = LWDiD(rolling="demean", estimation_method="ipw").fit( + panel_with_x, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1"], + ) + assert res.att > 0 + + def test_dr_positive_att(self, panel): + """DR (doubly robust) should recover positive ATT.""" + panel_with_x = panel.copy() + rng = np.random.default_rng(0) + units = panel_with_x["unit"].unique() + xmap = dict(zip(units, rng.normal(size=len(units)))) + panel_with_x["x1"] = panel_with_x["unit"].map(xmap) # unit-constant + res = LWDiD(rolling="demean", estimation_method="dr").fit( + panel_with_x, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1"], + ) + assert res.att > 0 + + def test_hc1_se_positive(self, panel): + res = LWDiD(vcov_type="hc1").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert res.se > 0 + + def test_classical_se_positive(self, panel): + res = LWDiD(vcov_type="classical").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert res.se > 0 + + def test_cluster_robust_se(self, panel): + """Cluster-robust SE should be positive.""" + # Create a cluster variable (group units into clusters) + panel_cl = panel.copy() + panel_cl["cluster_id"] = panel_cl["unit"] % 10 + res = LWDiD(cluster="cluster_id").fit( + panel_cl, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert res.se > 0 + + def test_n_obs_n_treated_n_control(self, panel): + """Sample sizes should be consistent.""" + res = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + assert res.n_treated == 30 + assert res.n_control == 50 + assert res.n_obs == 80 + + def test_result_not_staggered(self, panel): + res = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + assert not res.is_staggered + assert res.cohort_effects is None + + def test_params_stored(self, panel): + """RA should store coefficient vector.""" + res = LWDiD(rolling="demean", estimation_method="reg").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert res.params is not None + assert len(res.params) >= 2 # intercept + treatment + + def test_vcov_stored(self, panel): + """RA should store vcov matrix.""" + res = LWDiD(rolling="demean", estimation_method="reg").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert res.vcov is not None + assert res.vcov.shape[0] == res.vcov.shape[1] + + def test_controls_improve_precision(self): + """Adding relevant controls should reduce SE (most cases).""" + rng = np.random.default_rng(99) + panel = _make_common_timing_panel(n_treated=50, n_control=100, seed=99) + # Add control correlated with outcome + unit_map = {} + for uid in panel["unit"].unique(): + unit_map[uid] = rng.normal(0, 2) + panel["x_corr"] = panel["unit"].map(unit_map) + + res_no_ctrl = LWDiD(estimation_method="reg").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + res_ctrl = LWDiD(estimation_method="reg").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat", covariates=["x_corr"] + ) + # Both should produce finite results + assert np.isfinite(res_no_ctrl.se) + assert np.isfinite(res_ctrl.se) + + +# ─── Staggered Design Tests ───────────────────────────────────────────────── + + +class TestLWDiDStaggered: + """Test staggered adoption designs.""" + + @pytest.fixture + def stag_panel(self): + return _make_staggered_panel(true_att=1.5) + + def test_staggered_never_treated(self, stag_panel): + res = LWDiD(control_group="never_treated").fit( + stag_panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert isinstance(res, LWDiDResults) + assert res.cohort_effects is not None + + def test_staggered_not_yet_treated(self, stag_panel): + res = LWDiD(control_group="not_yet_treated").fit( + stag_panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert res.att is not None + assert np.isfinite(res.att) + + def test_cohort_effects_populated(self, stag_panel): + res = LWDiD().fit( + stag_panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert res.cohort_effects is not None + assert len(res.cohort_effects) > 0 + + def test_staggered_att_positive(self, stag_panel): + """Overall ATT should be positive (true_att=1.5).""" + res = LWDiD(control_group="never_treated").fit( + stag_panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert res.att > 0 + + def test_staggered_is_staggered(self, stag_panel): + res = LWDiD().fit( + stag_panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert res.is_staggered + + def test_staggered_se_positive(self, stag_panel): + res = LWDiD(control_group="never_treated").fit( + stag_panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert res.se > 0 + + def test_staggered_detrend(self, stag_panel): + """Detrend should also work for staggered.""" + res = LWDiD(rolling="detrend", control_group="never_treated").fit( + stag_panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert isinstance(res, LWDiDResults) + assert res.att > 0 + + def test_staggered_cluster_equals_unit_column(self, stag_panel): + """Regression: cluster= the unit column must not raise KeyError. + + The unit column is consumed by set_index inside the staggered + engine, so looking it up as a regular column used to crash when + cluster == unit (the most common by-unit clustering spelling). + """ + res = LWDiD(cluster="unit", control_group="never_treated").fit( + stag_panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert np.isfinite(res.att) + assert np.isfinite(res.se) and res.se > 0 + + # An explicit copy of the unit column under a different name must + # give exactly the same estimates. + copied = stag_panel.copy() + copied["cluster_id"] = copied["unit"] + res_copy = LWDiD(cluster="cluster_id", control_group="never_treated").fit( + copied, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert res.att == res_copy.att + assert res.se == res_copy.se + + def test_no_treated_cohorts_raises(self): + """All cohort=0 should raise.""" + df = pd.DataFrame( + { + "unit": [1, 1, 2, 2], + "time": [1, 2, 1, 2], + "y": [1.0, 2.0, 3.0, 4.0], + "treat": [0, 0, 0, 0], + "cohort": [0, 0, 0, 0], + } + ) + with pytest.raises(ValueError, match="[Nn]o treated cohort"): + LWDiD().fit( + df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="cohort" + ) + + def test_never_treated_required_when_specified(self): + """control_group='never_treated' requires at least one cohort=0 unit.""" + # All units are in cohort 3 (treated) + df = pd.DataFrame( + { + "unit": [1, 1, 2, 2, 3, 3], + "time": [1, 2, 1, 2, 1, 2], + "y": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + "treat": [0, 1, 0, 1, 0, 0], + "cohort": [2, 2, 2, 2, 3, 3], + } + ) + with pytest.raises(ValueError, match="never-treated"): + LWDiD(control_group="never_treated").fit( + df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="cohort" + ) + + +# ─── Results Container Tests ──────────────────────────────────────────────── + + +class TestLWDiDResults: + """Test the LWDiDResults dataclass interface.""" + + @pytest.fixture + def result(self): + panel = _make_common_timing_panel() + return LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + + def test_inference_consistency(self, result): + """t_stat ≈ att / se.""" + if result.se > 0 and np.isfinite(result.se): + np.testing.assert_allclose(result.t_stat, result.att / result.se, rtol=1e-10) + + def test_conf_int_bounds(self, result): + """CI should bracket ATT.""" + lo, hi = result.conf_int + assert lo < result.att < hi + + def test_conf_int_symmetric(self, result): + """CI should be symmetric around ATT (normal-based).""" + lo, hi = result.conf_int + half_width_lo = result.att - lo + half_width_hi = hi - result.att + np.testing.assert_allclose(half_width_lo, half_width_hi, rtol=1e-10) + + def test_p_value_range(self, result): + """p-value should be in [0, 1].""" + assert 0 <= result.p_value <= 1 + + def test_summary_contains_fields(self, result): + s = result.summary() + assert "ATT" in s or "att" in s.lower() + assert "LWDiD" in s + + def test_to_dataframe(self, result): + df = result.to_dataframe() + assert isinstance(df, pd.DataFrame) + assert len(df) >= 1 + assert "att" in df.columns + + def test_to_dict_serializable(self, result): + """to_dict() should produce JSON-serializable output.""" + d = result.to_dict() + json.dumps(d, default=str) + + def test_to_dict_contains_keys(self, result): + d = result.to_dict() + assert "att" in d + assert "se" in d + assert "rolling" in d + assert "estimation_method" in d + + def test_repr_informative(self, result): + r = repr(result) + assert "LWDiDResults" in r + assert "ATT" in r + + def test_rolling_metadata(self, result): + assert result.rolling == "demean" + assert result.estimation_method == "reg" + assert result.vcov_type == "hc1" + assert result.alpha == 0.05 + + def test_nan_inference_when_se_zero(self): + """Direct construction with se=0 should give NaN inference.""" + res = LWDiDResults( + att=1.0, + se=0.0, + t_stat=float("nan"), + p_value=float("nan"), + conf_int=(float("nan"), float("nan")), + n_obs=100, + n_treated=30, + n_control=70, + rolling="demean", + estimation_method="reg", + vcov_type="hc1", + alpha=0.05, + ) + assert np.isnan(res.t_stat) + assert np.isnan(res.p_value) + assert np.isnan(res.conf_int[0]) + assert np.isnan(res.conf_int[1]) + + +# ─── Different VCE Comparisons ────────────────────────────────────────────── + + +class TestLWDiDVCEComparisons: + """Compare VCE methods produce different but finite SEs.""" + + @pytest.fixture + def panel(self): + return _make_common_timing_panel(n_treated=40, n_control=80, seed=77) + + def test_hc1_vs_classical(self, panel): + res_cl = LWDiD(vcov_type="classical").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + res_hc1 = LWDiD(vcov_type="hc1").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + # ATTs should be the same (same point estimate) + np.testing.assert_allclose(res_cl.att, res_hc1.att, atol=1e-12) + # SEs differ + assert res_cl.se > 0 + assert res_hc1.se > 0 + + def test_cluster_vs_hc1(self, panel): + panel_cl = panel.copy() + panel_cl["cluster_id"] = panel_cl["unit"] % 10 + res_hc1 = LWDiD(vcov_type="hc1").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + res_cl = LWDiD(cluster="cluster_id").fit( + panel_cl, outcome="y", unit="unit", time="time", treatment="treat" + ) + # Point estimates should be identical + np.testing.assert_allclose(res_hc1.att, res_cl.att, atol=1e-12) + # Both SEs positive + assert res_cl.se > 0 + assert res_hc1.se > 0 + + +# ─── Estimator Consistency Tests ──────────────────────────────────────────── + + +class TestLWDiDEstimatorConsistency: + """Test that different estimators produce consistent results.""" + + @pytest.fixture + def panel_with_controls(self): + panel = _make_common_timing_panel(n_treated=50, n_control=100, seed=55) + rng = np.random.default_rng(55) + units = panel["unit"].unique() + xmap = dict(zip(units, rng.normal(size=len(units)))) + panel["x1"] = panel["unit"].map(xmap) # unit-constant + return panel + + def test_ra_ipw_same_sign(self, panel_with_controls): + """RA and IPW should give same-sign ATT.""" + res_ra = LWDiD(estimation_method="reg").fit( + panel_with_controls, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1"], + ) + res_ipw = LWDiD(estimation_method="ipw").fit( + panel_with_controls, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1"], + ) + assert np.sign(res_ra.att) == np.sign(res_ipw.att) + + def test_reg_dr_same_sign(self, panel_with_controls): + """Regression adjustment and DR should give same-sign ATT.""" + res_ra = LWDiD(estimation_method="reg").fit( + panel_with_controls, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1"], + ) + res_dr = LWDiD(estimation_method="dr").fit( + panel_with_controls, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1"], + ) + assert np.sign(res_ra.att) == np.sign(res_dr.att) + + def test_ipw_without_controls_warns(self): + """IPW without controls should warn and behave like RA.""" + panel = _make_common_timing_panel(seed=88) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + res = LWDiD(estimation_method="ipw").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + # Should produce a warning about no controls + ipw_warnings = [x for x in w if "IPW" in str(x.message)] + assert len(ipw_warnings) > 0 + assert np.isfinite(res.att) + + +# ─── Cohort-Time Cell Support (issue #734) ────────────────────────────────── + + +def _make_eligibility_panel(seed=7): + """Staggered panel with distinct cohort sizes. + + Sizes are distinct so a cell's control count identifies the eligible + pool uniquely: never-treated 2, cohort 3 has 4 units, cohort 8 has 3, + cohort 10 has 5. + """ + sizes = {0: 2, 3: 4, 8: 3, 10: 5} + rng = np.random.default_rng(seed) + rows = [] + uid = 0 + for g, size in sizes.items(): + for _ in range(size): + unit_fe = rng.normal() + for t in range(1, 13): + treated = g > 0 and t >= g + rows.append( + { + "unit": uid, + "time": t, + "cohort": g, + "treat": int(treated), + "y": (unit_fe + 0.3 * t + rng.normal(0, 0.5) + (2.0 if treated else 0.0)), + } + ) + uid += 1 + return pd.DataFrame(rows), sizes + + +def _make_trend_only_panel(shift=None): + """The issue #734 reproduction: a pure common time trend, zero effect. + + Cohort 3 (5 units) and cohort 5 (5 units) over t = 1..6, plus two + never-treated units observed only through t = 4. No control is + available from t = 5 on: cohort 3 loses every control there and + cohort 5 never has a post-treatment control. + """ + rows = [] + for unit in range(12): + cohort = 3 if unit < 5 else (5 if unit < 10 else 0) + last_period = 4 if cohort == 0 else 6 + for time in range(1, last_period + 1): + y = float(time) + if shift is not None: + y += shift(time) + rows.append( + { + "unit": unit, + "time": time, + "cohort": cohort, + "treat": int(cohort > 0 and time >= cohort), + "y": y, + } + ) + return pd.DataFrame(rows) + + +def _fit_trend_only(data): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + res = LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="classical", + control_group="not_yet_treated", + ).fit( + data, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + return res, [str(x.message) for x in caught] + + +class TestCohortTimeCellSupport: + """Per-(g, t) cells with calendar-time-specific control eligibility. + + The estimand is built from cohort-time cells whose control pool is + A_{g,t} = {G = g} u {G = 0} u {G > max(g, t)} (LW 2026 Sec. 7). Applying + eligibility as a unit-level filter and then averaging each unit's + transformed outcomes over unequal calendar windows produces a non-zero + ATT under a pure common time trend, which is the defect these tests pin. + """ + + def test_later_cohort_eligibility_is_period_specific(self): + """A later cohort is a valid control at r = 3 but not at r = 5.""" + panel, sizes = _make_eligibility_panel() + res = LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="classical", + control_group="not_yet_treated", + ).fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + cells = res.cohort_time_effects + + # r = 3 is calendar t = 6: cohorts 8 and 10 are both still untreated. + at_r3 = cells[(3, 6)] + assert at_r3["n_treated"] == sizes[3] + assert at_r3["n_control"] == sizes[0] + sizes[8] + sizes[10] + + # r = 5 is calendar t = 8: cohort 8 is treated by then and drops out. + at_r5 = cells[(3, 8)] + assert at_r5["n_treated"] == sizes[3] + assert at_r5["n_control"] == sizes[0] + sizes[10] + + # By t = 10 only the never-treated remain eligible. + assert cells[(3, 10)]["n_control"] == sizes[0] + + def test_eligibility_matches_formula_for_every_cell(self): + """Every cohort-3 cell's control count equals |A_{3,t}| - |G = 3|.""" + panel, sizes = _make_eligibility_panel() + res = LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="classical", + control_group="not_yet_treated", + ).fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + for t in range(3, 13): + expected = sizes[0] + sum(size for g, size in sizes.items() if g > 0 and g > max(3, t)) + assert res.cohort_time_effects[(3, t)]["n_control"] == expected, t + + def test_common_time_trend_yields_zero_att(self): + """A pure time trend with no treatment effect must estimate zero.""" + res, _ = _fit_trend_only(_make_trend_only_panel()) + assert abs(res.att) < 1e-10 + + @pytest.mark.parametrize( + "shift", + [ + lambda t: 100.0, + lambda t: 0.5 * t**2, + lambda t: (-1.0) ** t * 3.0, + ], + ids=["level", "quadratic", "sawtooth"], + ) + def test_common_time_shift_leaves_att_unchanged(self, shift): + """Adding any time-only h(t) to every unit cannot move the ATT.""" + base, _ = _fit_trend_only(_make_trend_only_panel()) + shifted, _ = _fit_trend_only(_make_trend_only_panel(shift=shift)) + assert abs(shifted.att - base.att) < 1e-10 + + def test_unsupported_cells_are_reported(self): + """Cells with an empty control pool are recorded and warned about.""" + res, messages = _fit_trend_only(_make_trend_only_panel()) + + # Cohort 3 keeps no controls from t = 5 onward. + for t in (5, 6): + cell = res.cohort_time_effects[(3, t)] + assert cell["skip_reason"] == "zero_treated_control" + assert cell["inference_status"] == "not_estimable" + assert np.isnan(cell["att"]) + + assert any("skipped" in m and "unsupported" in m for m in messages) + + def test_cohort_without_any_supported_cell_is_dropped(self): + """Cohort 5 has no eligible post-treatment control and is dropped.""" + res, _ = _fit_trend_only(_make_trend_only_panel()) + assert 5 not in res.cohort_effects + assert all( + res.cohort_time_effects[key]["skip_reason"] == "zero_treated_control" + for key in res.cohort_time_effects + if key[0] == 5 and key[1] >= 5 + ) + + def test_degenerate_standard_errors_are_not_reported(self): + """An exactly-fitting design must not report a ~0 SE as inference.""" + res, messages = _fit_trend_only(_make_trend_only_panel()) + assert np.isnan(res.se) + assert np.isnan(res.p_value) + assert res.inference_basis == "unavailable_degenerate_cells" + assert any("degenerate or non-finite standard error" in m for m in messages) + + def test_supported_design_reports_joint_influence_inference(self): + """A well-identified staggered panel still gets finite inference.""" + panel, _ = _make_eligibility_panel() + res = LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="hc1", + control_group="not_yet_treated", + ).fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert res.inference_basis == "joint_influence_function" + assert np.isfinite(res.se) and res.se > 0 + assert res.att == pytest.approx(2.0, abs=0.5) + + +# ─── Joint Influence-Function Inference (issue #735) ──────────────────────── + + +def _cluster_sums(values, ids): + frame = pd.DataFrame({"value": values, "cluster": ids}) + return frame.groupby("cluster", sort=False)["value"].sum().to_numpy() + + +def _make_shared_control_panel(seed=101, n_never=40, per_cohort=20, cohorts=(5, 7, 9)): + """Staggered panel whose cohorts all draw on the same never-treated pool.""" + rng = np.random.default_rng(seed) + rows = [] + uid = 0 + for g in (0,) + tuple(cohorts): + size = n_never if g == 0 else per_cohort + for _ in range(size): + unit_fe = rng.normal() + for t in range(1, 13): + treated = g > 0 and t >= g + rows.append( + { + "unit": uid, + "time": t, + "cohort": g, + "treat": int(treated), + "y": (unit_fe + 0.2 * t + rng.normal(0, 0.7) + (1.5 if treated else 0.0)), + } + ) + uid += 1 + return pd.DataFrame(rows) + + +class TestInfluenceFunctionReconciliation: + """Each estimator returns the influence function behind its own SE. + + Cohort effects that share control units are not independent, so the + staggered aggregation combines per-cell influence functions rather than + summing marginal variances. That is only sound if a single cell's + influence function reproduces that cell's standard error exactly, which + is the identity pinned here: the contributions are the estimator's own + asymptotically linear representation reweighted by the variance + estimator, not a proxy rescaled to hit a target. + """ + + @pytest.fixture(scope="class") + def sample(self): + rng = np.random.default_rng(11) + n = 200 + controls = rng.normal(size=(n, 2)) + index = 0.6 * controls[:, 0] - 0.4 * controls[:, 1] + treatment = (rng.uniform(size=n) < 1 / (1 + np.exp(-index))).astype(float) + y = 1.0 + 2.0 * treatment + controls @ np.array([0.5, -0.3]) + rng.normal(0, 1.2, size=n) + clusters = rng.integers(0, 12, size=n) + return y, treatment, controls, clusters, n + + # Fix-wave WS6: ipw/dr accept vcov_type='hc1' ONLY (the IF sandwich); + # other families were silently inert and are now rejected at + # construction, so the reconciliation grid enumerates real configs. + @pytest.mark.parametrize( + "estimation_method,vcov", + [ + ("reg", "classical"), + ("reg", "hc1"), + ("reg", "hc2"), + ("reg", "hc3"), + ("reg", "cluster"), + ("ipw", "hc1"), + ("ipw", "cluster"), + ("dr", "hc1"), + ("dr", "cluster"), + ], + ) + def test_influence_reproduces_standard_error(self, sample, estimation_method, vcov): + y, treatment, controls, clusters, n = sample + cluster_ids = clusters if vcov == "cluster" else None + if vcov == "cluster": + est = LWDiD(estimation_method=estimation_method, cluster="cl") + else: + est = LWDiD(estimation_method=estimation_method, vcov_type=vcov) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + _att, se, _, _, _, influence = getattr(est, f"_estimate_{estimation_method}")( + y, treatment, controls, cluster_ids, n + ) + assert influence is not None + effective = influence if cluster_ids is None else _cluster_sums(influence, cluster_ids) + assert float(np.sqrt(np.sum(effective**2))) == pytest.approx(se, rel=1e-10) + + @pytest.mark.parametrize("estimation_method", ["ipw", "dr", "psm"]) + @pytest.mark.parametrize("vcov", ["classical", "hc2", "hc3"]) + def test_inert_vcov_values_rejected(self, estimation_method, vcov): + with pytest.raises(ValueError, match="silently inert"): + LWDiD(estimation_method=estimation_method, vcov_type=vcov) + + def test_cluster_composes_only_with_hc1(self): + with pytest.raises(ValueError, match="composes only with vcov_type='hc1'"): + LWDiD(estimation_method="reg", vcov_type="hc3", cluster="cl") + + def test_psm_cluster_rejected(self): + with pytest.raises(ValueError, match="psm.*does not support cluster"): + LWDiD(estimation_method="psm", cluster="cl") + + @pytest.mark.parametrize("vcov", ["classical", "hc1", "hc2", "hc3", "cluster"]) + def test_influence_reproduces_standard_error_without_controls(self, sample, vcov): + """The regression design matrix drops the interaction block without controls.""" + y, treatment, _controls, clusters, n = sample + cluster_ids = clusters if vcov == "cluster" else None + if vcov == "cluster": + est = LWDiD(estimation_method="reg", cluster="cl") + else: + est = LWDiD(estimation_method="reg", vcov_type=vcov) + _att, se, _, _, _, influence = est._estimate_reg(y, treatment, None, cluster_ids, n) + effective = influence if cluster_ids is None else _cluster_sums(influence, cluster_ids) + assert float(np.sqrt(np.sum(effective**2))) == pytest.approx(se, rel=1e-10) + + def test_matching_reports_no_influence_function(self, sample): + """PSM has no influence-function representation and must say so.""" + y, treatment, controls, _clusters, n = sample + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + *_, influence = LWDiD(estimation_method="psm")._estimate_psm( + y, treatment, controls, None, n + ) + assert influence is None + + def test_staggered_psm_reports_unavailable_basis(self): + """Overall PSM inference is NaN rather than an independence guess.""" + panel = _make_shared_control_panel(per_cohort=15, n_never=30) + rng = np.random.default_rng(5) + x_by_unit = pd.Series( + rng.normal(size=panel["unit"].nunique()), index=panel["unit"].unique() + ) + panel["x1"] = panel["unit"].map(x_by_unit) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + res = LWDiD( + rolling="demean", + estimation_method="psm", + control_group="never_treated", + ).fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + covariates=["x1"], + ) + assert res.inference_basis == "unavailable_matching" + assert np.isnan(res.se) + assert any("matching" in str(w.message) for w in caught) + + +class TestStaggeredJointInference: + """Overall staggered inference accounts for shared control units. + + Cohorts estimated against a common never-treated pool are positively + correlated. Summing marginal cohort variances therefore understates the + overall standard error; combining influence functions does not. + """ + + @pytest.fixture(scope="class") + def fitted(self): + panel = _make_shared_control_panel() + res = LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="hc1", + control_group="never_treated", + ).fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + return panel, res + + def test_reports_joint_influence_basis(self, fitted): + _panel, res = fitted + assert res.inference_basis == "joint_influence_function" + assert np.isfinite(res.se) and res.se > 0 + + def test_wider_than_independence_assumption(self, fitted): + """The independence formula is the specific thing being corrected.""" + _panel, res = fitted + cohort_se = np.array([v["se"] for v in res.cohort_effects.values()]) + weights = np.array([v["weight"] for v in res.cohort_effects.values()]) + independence_se = float(np.sqrt(np.sum(weights**2 * cohort_se**2))) + assert res.se > independence_se + + @pytest.mark.slow + def test_matches_unit_cluster_bootstrap(self, fitted, ci_params): + """Concordance with a unit-level bootstrap, which needs no + independence assumption. The independence formula misses by ~24% on + this design; the joint influence function lands within 10%.""" + panel, res = fitted + units = panel["unit"].unique() + blocks = {u: g for u, g in panel.groupby("unit")} + rng = np.random.default_rng(2024) + draws = [] + for _ in range(ci_params.bootstrap(300, min_n=60)): + picked = rng.choice(units, size=len(units), replace=True) + frames = [] + for new_id, u in enumerate(picked): + block = blocks[u].copy() + block["unit"] = new_id + frames.append(block) + sample = pd.concat(frames, ignore_index=True) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + try: + att = ( + LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="hc1", + control_group="never_treated", + ) + .fit( + sample, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + .att + ) + except ValueError: + continue + if np.isfinite(att): + draws.append(att) + + bootstrap_se = float(np.std(np.array(draws), ddof=1)) + assert bootstrap_se == pytest.approx(res.se, rel=0.10) + + +# ─── Post-Fit Aggregation Contract (issues #732, #733) ────────────────────── + + +class TestAggregationContract: + """``aggregate()`` reports the fit; it never re-derives inference. + + A staggered fit already chooses an inference basis - the composite + regression where the paper's theory applies, joint influence functions + otherwise. Recomputing an overall ATT from marginal cohort effects would + substitute a cohort-independence assumption for that basis and quietly + report a different standard error for the same estimand. + """ + + @pytest.fixture(scope="class") + def staggered(self): + return _make_shared_control_panel(n_never=30, per_cohort=15) + + @pytest.fixture(scope="class") + def composite_fit(self, staggered): + """The composite-regression path (never-treated + RA + classical).""" + return LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="classical", + control_group="never_treated", + ).fit( + staggered, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + + def test_uses_composite_regression(self, composite_fit): + assert composite_fit.inference_basis == "composite_regression" + + def test_simple_preserves_the_fitted_result(self, composite_fit): + """Exact agreement, including the finite-sample degrees of freedom.""" + agg = composite_fit.aggregate("simple") + assert agg.att[0] == composite_fit.att + assert agg.se[0] == composite_fit.se + assert agg.t_stat[0] == composite_fit.t_stat + assert agg.p_value[0] == composite_fit.p_value + assert agg.conf_int_lower[0] == composite_fit.conf_int[0] + assert agg.conf_int_upper[0] == composite_fit.conf_int[1] + assert agg.df[0] == composite_fit.df_inference + assert agg.alpha == composite_fit.alpha + + @pytest.mark.parametrize("vcov", ["classical", "hc1"]) + @pytest.mark.parametrize("control_group", ["never_treated", "not_yet_treated"]) + def test_simple_preserves_every_inference_basis(self, staggered, vcov, control_group): + """Holds off the composite path too, not just where it is gated on.""" + res = LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type=vcov, + control_group=control_group, + ).fit( + staggered, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + agg = res.aggregate("simple") + assert agg.att[0] == res.att + assert agg.se[0] == res.se + if res.df_inference is None: + assert np.isnan(agg.df[0]) + else: + assert agg.df[0] == res.df_inference + + def test_group_reports_cohort_effects_with_weights(self, composite_fit): + agg = composite_fit.aggregate("group") + assert agg.level == "group" + assert list(agg.label) == list(composite_fit.cohort_effects) + assert agg.weight is not None + assert float(np.nansum(agg.weight)) == pytest.approx(1.0) + for i, cohort in enumerate(agg.label): + assert agg.att[i] == composite_fit.cohort_effects[cohort]["att"] + + def test_group_dataframe_matches_shared_schema(self, composite_fit): + from diff_diff.aggregation import AGGREGATION_SCHEMA + + frame = composite_fit.aggregate("group").to_dataframe() + assert tuple(frame.columns) == AGGREGATION_SCHEMA + + def test_event_study_returns_shared_container(self, staggered): + from diff_diff.results_base import EVENT_STUDY_SCHEMA, EventStudyResults + + res = LWDiD(rolling="demean", estimation_method="reg", n_bootstrap=199, seed=7).fit( + staggered, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + es = res.aggregate("event_study") + assert isinstance(es, EventStudyResults) + frame = es.to_dataframe() + assert tuple(frame.columns) == EVENT_STUDY_SCHEMA + + # The anchor period is carried as a reference row, not dropped. + assert list(frame.loc[frame["is_reference"], "event_time"]) == [-1] + assert frame.loc[frame["is_reference"], "att"].tolist() == [0.0] + assert es.cband_lower is not None + assert es.cband_crit_value > 0 + + def test_event_study_serialises_through_to_dict(self, staggered): + res = LWDiD(rolling="demean", estimation_method="reg", n_bootstrap=199, seed=7).fit( + staggered, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + payload = res.to_dict() + assert payload["reference_periods"] == [-1] + assert payload["cband_method"] == "multiplier_bootstrap_sup_t" + assert payload["cband_n_bootstrap"] == 199 + assert payload["inference_basis"] == res.inference_basis + assert set(payload["event_study_effects"]) == {str(r) for r in res.event_study_effects} + + def test_unsupported_type_names_the_supported_set(self, composite_fit): + with pytest.raises(ValueError, match="Unsupported aggregation type"): + composite_fit.aggregate("overall") + with pytest.raises(ValueError, match="'simple', 'event_study', 'group'"): + composite_fit.aggregate("calendar") + + def test_weights_selector_is_rejected(self, composite_fit): + with pytest.raises(ValueError, match="does not accept a weights selector"): + composite_fit.aggregate("simple", weights="cell") + + def test_balance_e_is_rejected_off_event_study(self, composite_fit): + with pytest.raises(ValueError, match="balance_e"): + composite_fit.aggregate("simple", balance_e=2) + + def test_common_timing_aggregate_simple_and_group(self): + """Guard relaxation: simple relays the fit on common timing, while + group still raises (there is no cohort dimension).""" + panel = _make_common_timing_panel(seed=3) + res = LWDiD(rolling="demean").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + agg = res.aggregate("simple") + frame = agg.to_dataframe() + assert frame["att"].iloc[0] == res.att + assert frame["se"].iloc[0] == res.se + with pytest.raises(ValueError, match="only available for staggered"): + res.aggregate("group") + + def test_fit_time_aggregate_is_gone(self): + """Aggregation is post-fit only: fit() no longer takes aggregate.""" + panel = _make_shared_control_panel(n_never=20, per_cohort=10) + with pytest.raises(TypeError, match="aggregate"): + LWDiD(rolling="demean").fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + aggregate="group", + ) + + +# ─── PR #588 review: statistical-core fixes ──────────────────────────────────── + + +class TestClassicalJointInference: + """Classical joint covariance is built from residual-based influence. + + The former ``sigma * basis`` contributions gave every shared control + unit a non-zero cross-cell product regardless of its actual outcome + draw, fabricating correlation between cohort-time cells and inflating + the classical joint SE roughly two-fold against a unit-level bootstrap. + """ + + @pytest.fixture(scope="class") + def fitted(self): + panel = _make_shared_control_panel(seed=303) + res = LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="classical", + control_group="not_yet_treated", + ).fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + return panel, res + + def test_reports_joint_influence_basis(self, fitted): + _panel, res = fitted + assert res.inference_basis == "joint_influence_function" + assert np.isfinite(res.se) and res.se > 0 + + @pytest.mark.slow + def test_matches_unit_level_bootstrap(self, fitted, ci_params): + """Shared not-yet-treated controls: the classical joint/overall SE + must agree with a unit-level bootstrap that assumes no independence. + The sigma * basis contributions missed by ~2x on this design.""" + panel, res = fitted + units = panel["unit"].unique() + blocks = {u: g for u, g in panel.groupby("unit")} + rng = np.random.default_rng(588) + draws = [] + for _ in range(ci_params.bootstrap(400, min_n=60)): + picked = rng.choice(units, size=len(units), replace=True) + frames = [] + for new_id, u in enumerate(picked): + block = blocks[u].copy() + block["unit"] = new_id + frames.append(block) + sample = pd.concat(frames, ignore_index=True) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + try: + att = ( + LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="classical", + control_group="not_yet_treated", + ) + .fit( + sample, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + .att + ) + except ValueError: + continue + if np.isfinite(att): + draws.append(att) + + bootstrap_se = float(np.std(np.array(draws), ddof=1)) + assert res.se == pytest.approx(bootstrap_se, rel=0.3) + + def test_event_study_simultaneous_band_is_sane(self): + """The sup-t band exists and is at least as wide as pointwise CIs.""" + panel = _make_shared_control_panel(seed=303) + res = LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="classical", + control_group="not_yet_treated", + n_bootstrap=199, + seed=7, + ).fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert res.cband_method == "multiplier_bootstrap_sup_t" + assert np.isfinite(res.cband_crit_value) and res.cband_crit_value > 0 + tolerance = 1e-12 + for row in res.event_study_effects.values(): + if "cband_conf_int" not in row: + continue + lo, hi = row["cband_conf_int"] + assert np.isfinite(lo) and np.isfinite(hi) and lo < hi + assert lo <= row["conf_int"][0] + tolerance + assert hi >= row["conf_int"][1] - tolerance + + +class TestAllEventuallyTreatedRejection: + """No never-treated units + not_yet_treated controls is rejected. + + The final-period cohort-time cells of such designs have an empty + control pool, so estimating them would silently truncate the estimand + (e.g. cohorts {3, 5} over T = 5 lose (3, 5) and (5, 5), dropping event + time 2 entirely). + """ + + @staticmethod + def _all_treated_panel(cohorts=(3, 5), n_periods=5, per_cohort=6, seed=11): + rng = np.random.default_rng(seed) + rows = [] + uid = 0 + for g in cohorts: + for _ in range(per_cohort): + unit_fe = rng.normal() + for t in range(1, n_periods + 1): + treated = t >= g + rows.append( + { + "unit": uid, + "time": t, + "cohort": g, + "treat": int(treated), + "y": unit_fe + 0.3 * t + rng.normal(0, 0.4) + float(treated), + } + ) + uid += 1 + return pd.DataFrame(rows) + + def test_all_eventually_treated_raises(self): + panel = self._all_treated_panel() + with pytest.raises(ValueError, match="eventually treated"): + LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="hc1", + control_group="not_yet_treated", + ).fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + + def test_design_with_never_treated_runs_complete(self): + """A regular staggered design estimates every relative event time.""" + panel = _make_shared_control_panel(seed=101, cohorts=(5, 7, 9)) + res = LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="hc1", + control_group="not_yet_treated", + ).fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + expected = {t - g for g in (5, 7, 9) for t in range(1, 13)} - {-1} + assert set(res.event_study_effects) == expected + assert all(np.isfinite(row["effect"]) for row in res.event_study_effects.values()) + assert np.isfinite(res.att) and np.isfinite(res.se) + + +class TestStaggeredCovariateConstancy: + """Staggered LWDiD only supports unit-constant covariates. + + Cohort-time cells read covariates at each calendar time, so a column + that changes after treatment would silently move the ATT; such columns + are rejected up front (matching the lwdid-py reference behaviour). + """ + + @staticmethod + def _panel_with_covariate(time_varying): + panel = _make_shared_control_panel(seed=17, n_never=20, per_cohort=10) + rng = np.random.default_rng(23) + x_by_unit = pd.Series( + rng.normal(size=panel["unit"].nunique()), index=panel["unit"].unique() + ) + panel["x1"] = panel["unit"].map(x_by_unit) + if time_varying: + # Post-treatment shift: constant pre-treatment, jumps at adoption. + panel["x1"] += 0.5 * panel["treat"] + return panel + + def test_post_treatment_varying_covariate_raises(self): + panel = self._panel_with_covariate(time_varying=True) + with pytest.raises(ValueError, match="not unit-constant"): + LWDiD(rolling="demean", estimation_method="reg", vcov_type="hc1").fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + covariates=["x1"], + ) + + def test_unit_constant_covariate_estimates(self): + panel = self._panel_with_covariate(time_varying=False) + res = LWDiD(rolling="demean", estimation_method="reg", vcov_type="hc1").fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + covariates=["x1"], + ) + assert np.isfinite(res.att) + assert res.att == pytest.approx(1.5, abs=0.5) + + +# ─── Datetime/Period Time Scale Tests ─────────────────────────────────────── + + +class TestDatetimeTimeScale: + """Staggered fits on datetime64/Period panels via integer-position encoding.""" + + @staticmethod + def _datetime_panel(): + """Numeric staggered panel plus a quarterly datetime relabeling.""" + numeric = _make_staggered_panel(seed=42) + date_map = { + t: pd.Timestamp("2000-01-01") + pd.DateOffset(months=3 * (t - 1)) + for t in sorted(numeric["time"].unique()) + } + panel = numeric.copy() + panel["date"] = panel["time"].map(date_map) + panel["adopt"] = panel["cohort"].map(lambda g: date_map[g] if g > 0 else pd.NaT) + return numeric, panel, date_map + + def test_datetime_staggered_matches_numeric(self): + numeric, panel, date_map = self._datetime_panel() + model = LWDiD(rolling="demean", estimation_method="reg", vcov_type="hc1") + res_num = model.fit( + numeric, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + res_dt = model.fit( + panel, + outcome="y", + unit="unit", + time="date", + treatment="treat", + first_treat="adopt", + ) + assert res_dt.att == pytest.approx(res_num.att) + assert res_dt.se == pytest.approx(res_num.se) + # Cohort keys are restored to the original datetime labels + expected_cohorts = {date_map[g] for g in res_num.cohort_effects} + assert set(res_dt.cohort_effects) == expected_cohorts + for g, info in res_dt.cohort_effects.items(): + assert info["cohort"] == g + # Cohort-time cells carry datetime labels with integer event times + for (g, t), info in res_dt.cohort_time_effects.items(): + assert isinstance(g, pd.Timestamp) and isinstance(t, pd.Timestamp) + assert info["cohort"] == g and info["time"] == t + assert int(info["relative_time"]) == info["relative_time"] + # Event-study labels stay integer position differences + assert list(res_dt.event_study_effects) == list(res_num.event_study_effects) + for label, row in res_num.event_study_effects.items(): + assert res_dt.event_study_effects[label]["effect"] == pytest.approx(row["effect"]) + + def test_period_dtype_staggered_fits(self): + numeric, panel, _ = self._datetime_panel() + panel["date"] = panel["date"].dt.to_period("Q") + panel["adopt"] = pd.PeriodIndex(panel["adopt"], freq="Q") + model = LWDiD(rolling="demean", estimation_method="reg", vcov_type="hc1") + res_num = model.fit( + numeric, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + res_p = model.fit( + panel, + outcome="y", + unit="unit", + time="date", + treatment="treat", + first_treat="adopt", + ) + assert res_p.att == pytest.approx(res_num.att) + assert all(isinstance(g, pd.Period) for g in res_p.cohort_effects) + + def test_mixed_time_scales_raise(self): + _, panel, _ = self._datetime_panel() + with pytest.raises(ValueError, match="same time scale"): + LWDiD(rolling="demean").fit( + panel, + outcome="y", + unit="unit", + time="date", + treatment="treat", + first_treat="cohort", + ) + + def test_datetime_all_eventually_treated_rejected(self): + """The all-eventually-treated guard must also fire on datetime panels.""" + _, panel, _ = self._datetime_panel() + eventually = panel.loc[panel["adopt"].notna()] + with pytest.raises(ValueError, match="eventually treated"): + LWDiD(rolling="demean", control_group="not_yet_treated").fit( + eventually, + outcome="y", + unit="unit", + time="date", + treatment="treat", + first_treat="adopt", + ) + + def test_datetime_time_varying_covariate_rejected(self): + """The covariate constancy guard must also fire on datetime panels.""" + _, panel, _ = self._datetime_panel() + rng = np.random.default_rng(0) + panel["x1"] = rng.normal(size=len(panel)) + with pytest.raises(ValueError, match="not unit-constant"): + LWDiD(rolling="demean").fit( + panel, + outcome="y", + unit="unit", + time="date", + treatment="treat", + first_treat="adopt", + covariates=["x1"], + ) + + def test_datetime_transformation_diagnostics_keys(self): + _, panel, date_map = self._datetime_panel() + diagnostics = LWDiD(rolling="demean").get_transformation_diagnostics( + panel, + outcome="y", + unit="unit", + time="date", + treatment="treat", + first_treat="adopt", + ) + assert diagnostics["design"] == "staggered" + assert all(isinstance(g, pd.Timestamp) for g in diagnostics["by_cohort"]) + + +class TestCohortNormalization: + """LWDiD fix-wave WS9: one shared cohort normalizer (`_normalize_cohorts`) + applied after time-scale encoding and before the design check, making + every downstream never-treated predicate coherent. Campaign findings + (execution-verified): first_treat=inf was iterated as a real cohort that + consumed tau_omega weight mass; beyond-window cohorts distorted the + composite; unbalanced panels missing the onset row were falsely + rejected; validator and fit() disagreed on the never-treated encoding. + """ + + def _cohorts(self): + # 8 treated across two cohorts + 12 never-treated + return {u: (3 if u < 4 else (4 if u < 8 else 0)) for u in range(20)} + + def test_inf_cohort_recoded_to_never_treated(self): + cohorts = self._cohorts() + cohorts[19] = np.inf + panel = _make_design_panel(cohorts) + est = LWDiD(rolling="demean", estimation_method="reg", control_group="never_treated") + with pytest.warns(UserWarning, match="first_treat=inf"): + res = est.fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert np.isfinite(res.att) + # inf never appears as a cohort anywhere in the results + assert all(np.isfinite(g) and g > 0 for g in res.cohort_effects) + + def test_beyond_window_cohort_recoded_and_counts_as_control(self): + cohorts = self._cohorts() + cohorts[19] = 9 # beyond n_periods=5 + panel = _make_design_panel(cohorts) + est = LWDiD(rolling="demean", estimation_method="reg", control_group="never_treated") + with pytest.warns(UserWarning, match="exceed the last observed period"): + res = est.fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert np.isfinite(res.att) + assert all(g <= 5 for g in res.cohort_effects) + + def test_negative_cohort_rejected(self): + cohorts = self._cohorts() + cohorts[19] = -2 + panel = _make_design_panel(cohorts) + # treat rows for a negative cohort: 1[t >= -2] would be all-1; keep 0 + panel.loc[panel["unit"] == 19, "treat"] = 0 + est = LWDiD(rolling="demean", estimation_method="reg") + with pytest.raises(ValueError, match="negative"): + est.fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + + def test_between_period_numeric_cohort_rejected_with_clear_message(self): + cohorts = self._cohorts() + cohorts[0] = 3.5 # between observed periods 3 and 4 + panel = _make_design_panel(cohorts) + # D_it = 1[t >= 3.5] -> treated at t=4,5 + panel.loc[panel["unit"] == 0, "treat"] = ( + panel.loc[panel["unit"] == 0, "time"] >= 3.5 + ).astype(int) + est = LWDiD(rolling="demean", estimation_method="reg") + with pytest.raises(ValueError, match="not observed time periods"): + est.fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + + def test_unobserved_onset_row_accepted(self): + # Campaign finding: requiring the onset row itself to be observed + # falsely rejected valid unbalanced panels. + panel = _make_design_panel(self._cohorts()) + drop_mask = (panel["unit"] == 0) & (panel["time"] == 3) # unit 0's onset row + panel = panel.loc[~drop_mask].reset_index(drop=True) + est = LWDiD(rolling="demean", estimation_method="reg") + res = est.fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert np.isfinite(res.att) + + def test_untreated_observed_row_after_onset_rejected(self): + # Review-caught hole: a first-observed-treated-row inequality alone + # would ACCEPT a unit whose observed post-onset rows are all D=0 + # (no D=1 rows anywhere, onset row unobserved). The equality + # predicate D_it == 1[t >= g_i] over observed rows must reject it. + panel = _make_design_panel(self._cohorts()) + u0 = panel["unit"] == 0 # cohort 3 + panel = panel.loc[~(u0 & (panel["time"] == 3))] # onset row unobserved + panel.loc[panel["unit"] == 0, "treat"] = 0 # observed t=4,5 stay D=0 + est = LWDiD(rolling="demean", estimation_method="reg") + with pytest.raises(ValueError, match="1\\[t >= cohort\\]"): + est.fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + + def test_validator_accepts_nan_coded_never_treated(self): + # Validator/fit split: fit() accepts NaN never-treated; the + # validator previously required cohort==0 and rejected it. + from diff_diff.lwdid import validate_staggered_data + + panel = _make_design_panel(self._cohorts()) + panel["cohort"] = panel["cohort"].astype(float).replace(0.0, np.nan) + out = validate_staggered_data(panel, "unit", "time", "cohort") + assert out["valid"], out["errors"] + assert out["n_never_treated"] == 12 + assert out["n_cohorts"] == 2 + + def test_validator_flags_nat_mixed_with_finite_cohort(self): + # nunique() excluded missing values, so a unit mixing NaT/NaN with + # a finite cohort passed validation then raised inside fit(). + from diff_diff.lwdid import validate_staggered_data + + panel = _make_design_panel(self._cohorts()) + panel["cohort"] = panel["cohort"].astype(float) + mix = (panel["unit"] == 0) & (panel["time"] == 1) + panel.loc[mix, "cohort"] = np.nan + out = validate_staggered_data(panel, "unit", "time", "cohort") + assert not out["valid"] + assert any("time-varying cohort" in e for e in out["errors"]) + + def test_validator_reports_no_treated_cohorts(self): + from diff_diff.lwdid import validate_staggered_data + + panel = _make_design_panel({u: 0 for u in range(6)}) + out = validate_staggered_data(panel, "unit", "time", "cohort") + assert not out["valid"] + assert any("No treated cohorts found" in e for e in out["errors"]) + + def test_validator_handles_datetime_cohorts_without_raw_errors(self): + # Previously df[cohort] > 0 raised a raw pandas TypeError on + # datetime cohorts and df[cohort] == 0 silently reported "no + # never-treated units". + from diff_diff.lwdid import validate_staggered_data + + base = _make_design_panel(self._cohorts()) + time_map = {t: pd.Timestamp(f"2020-0{t}-01") for t in range(1, 6)} + panel = base.assign( + time=base["time"].map(time_map), + cohort=base["cohort"].map(lambda g: time_map.get(g, pd.NaT)), + ) + out = validate_staggered_data(panel, "unit", "time", "cohort") + assert out["valid"], out["errors"] + assert out["n_never_treated"] == 12 + assert out["n_cohorts"] == 2 + + def test_is_never_treated_time_aware(self): + from diff_diff.lwdid import is_never_treated + + cohorts = self._cohorts() + cohorts[18] = np.inf + cohorts[19] = 9 # beyond the window + panel = _make_design_panel(cohorts) + panel.loc[panel["unit"].isin([18, 19]), "treat"] = 0 + base = is_never_treated(panel, "unit", "cohort") + aware = is_never_treated(panel, "unit", "cohort", time="time") + units = panel.drop_duplicates("unit")["unit"].to_numpy() + base_map = dict(zip(units, base)) + aware_map = dict(zip(units, aware)) + assert base_map[18] and aware_map[18] # inf is never-treated either way + assert not base_map[19] # beyond-window needs the time support + assert aware_map[19] + + def test_diagnostics_iterate_normalized_cohorts_only(self): + cohorts = self._cohorts() + cohorts[19] = np.inf + panel = _make_design_panel(cohorts) + est = LWDiD(rolling="demean", estimation_method="reg") + with pytest.warns(UserWarning, match="first_treat=inf"): + diag = est.get_transformation_diagnostics( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + ) + assert set(diag["by_cohort"]) == {3, 4} + + +class TestSeasonalTransformFailClosed: + """LWDiD fix-wave WS2: the seasonal transforms fail closed. Campaign + findings: detrendq silently fit intercept+trend (plain detrend) per + unit when pre-periods < seasonal parameter count - with quarterly data + and <=5 pre-periods EVERY unit fell back, so the whole fit was + numerically identical to rolling='detrend' while reporting 'detrendq'; + both q transforms silently extrapolated quarters unobserved in the + pre-period at the reference-season level. + """ + + @staticmethod + def _quarterly_common_panel(n_pre, t_max=12, n_units=30, seed=5): + rng = np.random.default_rng(seed) + season = np.array([1.0, -0.5, 0.8, -1.3]) + rows = [] + onset = n_pre + 1 + for u in range(n_units): + alpha = rng.normal() + treated = u < n_units // 2 + for t in range(1, t_max + 1): + d = int(treated and t >= onset) + y = alpha + season[(t - 1) % 4] + 0.05 * t + rng.normal(scale=0.3) + 1.2 * d + rows.append(dict(unit=u, time=t, treat=d, y=y)) + return pd.DataFrame(rows) + + def test_detrendq_insufficient_pre_fails_closed_not_silent_detrend(self): + # 4 pre-periods cover all 4 seasons -> n_params = 1 + 1 + 3 = 5 > 4: + # every unit is unidentified. Pre-fix this silently produced the + # detrend numbers; now the fit warns and the ATT is NaN with a + # consistent inference tuple. + df = self._quarterly_common_panel(n_pre=4) + est = LWDiD(rolling="detrendq", estimation_method="reg") + with pytest.warns(UserWarning, match="detrendq requires at least"): + res = est.fit(df, outcome="y", unit="unit", time="time", treatment="treat") + assert np.isnan(res.att) + from tests.conftest import assert_nan_inference + + assert_nan_inference( + {"se": res.se, "t_stat": res.t_stat, "p_value": res.p_value, "conf_int": res.conf_int} + ) + + def test_detrendq_identified_differs_from_detrend(self): + # With enough pre-periods the seasonal fit is identified and must + # NOT equal plain detrend on a seasonal DGP. + df = self._quarterly_common_panel(n_pre=8, t_max=16) + kw = dict(outcome="y", unit="unit", time="time", treatment="treat") + rq = LWDiD(rolling="detrendq", estimation_method="reg").fit(df, **kw) + rp = LWDiD(rolling="detrend", estimation_method="reg").fit(df, **kw) + assert np.isfinite(rq.att) + assert abs(rq.att - rp.att) > 1e-8 + + @pytest.mark.parametrize("rolling", ["demeanq", "detrendq"]) + def test_unobserved_pre_season_fails_closed(self, rolling): + # Pre-period covers quarters 1-3 only; post includes quarter 4 -> + # out-of-support prediction must warn + NaN, never extrapolate. + df = self._quarterly_common_panel(n_pre=7, t_max=8) + # onset at t=8 (quarter 4); pre t=1..7 covers quarters 1,2,3,4? + # t=1..7 -> quarters 1,2,3,4,1,2,3: quarter 4 IS observed. Drop + # every pre row in quarter 4 instead. + df = df.loc[~((df["time"] == 4))].reset_index(drop=True) + est = LWDiD(rolling=rolling, estimation_method="reg") + with pytest.warns(UserWarning, match="cannot predict quarter"): + res = est.fit(df, outcome="y", unit="unit", time="time", treatment="treat") + assert np.isnan(res.att) + + +class TestBootstrapIntegrity: + """LWDiD fix-wave WS3: common-timing bootstrap resampling integrity. + + Campaign findings (execution-verified): the bootstrap collected index + LABELS but fetched rows POSITIONALLY (raw IndexError on offset indexes; + silently doubled SEs on row-shuffled frames); cluster= was accepted but + dead inside _bootstrap (iid unit bootstrap labeled clustered); the + reported df_inference was G-1 while the bootstrap p-value used N-k. + """ + + @staticmethod + def _common_panel(n_units=40, t_max=6, onset=4, seed=11, n_clusters=8): + rng = np.random.default_rng(seed) + rows = [] + for u in range(n_units): + alpha = rng.normal() + cl = u % n_clusters + cl_shock = np.sin(cl) # cluster-correlated level + treated = u < n_units // 2 + for t in range(1, t_max + 1): + d = int(treated and t >= onset) + y = alpha + cl_shock + 0.1 * t + rng.normal(scale=0.4) + 1.3 * d + rows.append(dict(unit=u, time=t, treat=d, y=y, cl=cl)) + return pd.DataFrame(rows) + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_bootstrap_invariant_to_index_labels_and_row_order(self): + df = self._common_panel() + + def est(): + return LWDiD(rolling="demean", estimation_method="reg", n_bootstrap=60, seed=1) + + base = est().fit(df, **self.KW) + shifted = est().fit(df.set_axis(df.index + 1000), **self.KW) # offset labels + shuffled = df.sample(frac=1.0, random_state=5) # permuted labels + res_shuffled = est().fit(shuffled, **self.KW) + assert np.isfinite(base.se) + np.testing.assert_allclose(shifted.att, base.att, rtol=0, atol=1e-12) + np.testing.assert_allclose(shifted.se, base.se, rtol=0, atol=1e-12) + np.testing.assert_allclose(res_shuffled.att, base.att, rtol=0, atol=1e-12) + # Same seed + same units resampled -> the SE must not move with + # row order (pre-fix it more than doubled). + np.testing.assert_allclose(res_shuffled.se, base.se, rtol=1e-10) + + def test_cluster_bootstrap_resamples_clusters(self, ci_params): + df = self._common_panel() + n_boot = ci_params.bootstrap(120) + iid = LWDiD(rolling="demean", estimation_method="reg", n_bootstrap=n_boot, seed=7).fit( + df, **self.KW + ) + clustered = LWDiD( + rolling="demean", estimation_method="reg", n_bootstrap=n_boot, seed=7, cluster="cl" + ).fit(df, **self.KW) + # Points identical (resampling never moves the full-sample point) + np.testing.assert_allclose(clustered.att, iid.att, rtol=0, atol=1e-12) + # SEs genuinely differ on a cluster-correlated DGP + assert np.isfinite(clustered.se) and clustered.se > 0 + assert abs(clustered.se - iid.se) / iid.se > 1e-3 + # df matches the analytical clustered rule (G-1), not N-k + assert clustered.df_inference == 8 - 1 + assert clustered.cluster_name == "cl" + assert clustered.n_clusters == 8 + + def test_cluster_bootstrap_concords_with_analytical_cr1(self, ci_params): + df = self._common_panel(n_units=80, n_clusters=16) + n_boot = ci_params.bootstrap(300, min_n=199) + analytical = LWDiD(rolling="demean", estimation_method="reg", cluster="cl").fit( + df, **self.KW + ) + boot = LWDiD( + rolling="demean", + estimation_method="reg", + cluster="cl", + n_bootstrap=n_boot, + seed=13, + ).fit(df, **self.KW) + threshold = 0.40 if n_boot < 100 else 0.15 + assert abs(boot.se - analytical.se) / analytical.se < threshold, (boot.se, analytical.se) + + +class TestPSMCaliperContract: + """LWDiD fix-wave WS5 (campaign finding, deterministic repro): with a + caliper and n_neighbors > 1, argsort kept selecting np.inf-distance + (out-of-caliper) controls whenever fewer than n_neighbors controls fell + inside the caliper, silently averaging arbitrarily distant controls + into the counterfactual (ATT -49 vs the correct 1.0 on this fixture). + """ + + @staticmethod + def _fixture(): + # 2 treated (pscore ~0.64-0.69), 1 near control (~0.67, ydot=0 + # effect scale), 3 far controls (pscore ~0) whose transformed + # outcome is +100. + rows = [] + units = [ + ("t1", 1, 5.0, 1.0), + ("t2", 1, 4.6, 1.0), + ("c_near", 0, 4.8, 0.0), + ("c_far1", 0, -9.0, 100.0), + ("c_far2", 0, -9.4, 100.0), + ("c_far3", 0, -9.8, 100.0), + ] + for name, d, x, post_shift in units: + for t in (1, 2): + y = 1.0 if d and t == 2 else 0.0 + y += post_shift if t == 2 else 0.0 + rows.append(dict(unit=name, time=t, treat=d * int(t == 2), y=y, x=x)) + return pd.DataFrame(rows) + + def test_partial_caliper_shortfall_averages_survivors_only(self): + df = self._fixture() + est = LWDiD( + rolling="demean", + estimation_method="psm", + n_neighbors=2, + caliper=0.05, + ) + with pytest.warns(UserWarning, match="fewer than n_neighbors"): + res = est.fit( + df, outcome="y", unit="unit", time="time", treatment="treat", covariates=["x"] + ) + # Demeaned outcomes: treated ydot = 2.0, near control ydot = 0, + # far controls ydot = +100. Caliper-respecting match (c_near only): + # ATT = 2.0. Contaminated pre-fix value: 2.0 - (0+100)/2 = -48. + np.testing.assert_allclose(res.att, 2.0, atol=1e-10) + + +class TestSilentDataHandling: + """LWDiD fix-wave WS8 (campaign findings): NaN covariates silently + dropped units on cell paths while poisoning the common-timing OLS; + staggered n_obs/n_treated counted every input unit regardless of cell + drops; rank-deficient designs used the NOMINAL parameter count for the + df and a full-width pinv bread, breaking the IF == solve_ols SE + identity the docstring claims. + """ + + @staticmethod + def _staggered_panel(seed=31): + rng = np.random.default_rng(seed) + rows = [] + for u in range(24): + g = 3 if u < 5 else (4 if u < 10 else 0) + alpha = rng.normal() + x = rng.normal() + for t in range(1, 7): + d = int(g > 0 and t >= g) + y = alpha + 0.4 * x + rng.normal(scale=0.4) + 1.5 * d + rows.append(dict(unit=u, time=t, first=g, treat=d, y=y, x=x)) + return pd.DataFrame(rows) + + def test_nan_covariate_rejected_explicitly(self): + df = self._staggered_panel() + df.loc[3, "x"] = np.nan + with pytest.raises(ValueError, match="missing value"): + LWDiD(rolling="demean", estimation_method="reg").fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="first", + covariates=["x"], + ) + + def test_nan_cluster_rejected_explicitly(self): + df = self._staggered_panel() + df["cl"] = df["unit"] % 4 + df["cl"] = df["cl"].astype(float) + df.loc[df["unit"] == 2, "cl"] = np.nan + with pytest.raises(ValueError, match="Cluster column .* missing"): + LWDiD(rolling="demean", estimation_method="reg", cluster="cl").fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="first", + ) + + def test_rank_deficient_covariate_if_reproduces_solve_ols_se(self): + # A NON-trailing collinear covariate is dropped by solve_ols; the + # influence function must be rebuilt on the kept columns so its + # norm still reproduces the reported SE (docstring identity). + rng = np.random.default_rng(37) + n = 200 + controls = rng.normal(size=(n, 3)) + controls[:, 0] = 2.0 * controls[:, 2] + 1.0 # column 0 collinear + treatment = (rng.uniform(size=n) < 0.4).astype(float) + y = 1.0 + 2.0 * treatment + controls[:, 1] * 0.5 + rng.normal(size=n) + est = LWDiD(estimation_method="reg", vcov_type="hc1") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + att, se, coefs, _, n_params, influence = est._estimate_reg( + y, treatment, controls, None, n + ) + assert np.isfinite(att) and np.isfinite(se) + assert np.isnan(coefs).any() # a column really was dropped + assert n_params == int(np.sum(~np.isnan(coefs))) + assert influence is not None + assert float(np.sqrt(np.sum(influence**2))) == pytest.approx(se, rel=1e-10) + + def test_treatment_column_dropped_yields_nan_att(self): + # If the treatment column itself is pivoted out (collinear with a + # control), the ATT is unidentified: NaN point + no influence. + rng = np.random.default_rng(41) + n = 120 + treatment = (rng.uniform(size=n) < 0.5).astype(float) + controls = np.column_stack([treatment * 3.0, rng.normal(size=n)]) + y = 1.0 + rng.normal(size=n) + est = LWDiD(estimation_method="reg", vcov_type="hc1") + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + att, se, _, _, _, influence = est._estimate_reg(y, treatment, controls, None, n) + # Either the treatment or its collinear twin is dropped; if the + # treatment survives the ATT is finite - accept both resolutions + # but NEVER a finite ATT with se=0-style inference. + if np.isnan(att): + assert np.isnan(se) and influence is None + else: + assert np.isfinite(se) and se > 0 + + def test_staggered_metadata_counts_contributing_units(self): + # Under rolling='detrend', a unit with a single pre-period row has + # NaN transformed outcomes in EVERY cell (per-unit trend needs >= 2 + # pre points), so it is dropped from every cell's finite filter and + # contributes nothing - the estimation-sample metadata must not + # count it (campaign finding: n_obs/n_treated covered every input + # unit regardless of cell drops). + df = self._staggered_panel() + rng = np.random.default_rng(5) + extra = [ + dict(unit=99, time=t, first=0, treat=0, y=rng.normal(), x=0.0) for t in (1, 4, 5, 6) + ] # ONE pre row (t=1) + post rows + df = pd.concat([df, pd.DataFrame(extra)], ignore_index=True) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="detrend", estimation_method="reg").fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="first", + ) + assert res.n_obs == 24 # unit 99 contributed to no estimated cell + assert res.n_control == 14 + assert res.n_treated == 10 + + +class TestResultsPolish: + """LWDiD fix-wave WS10 result-surface pins.""" + + def test_to_latex_removed(self): + from diff_diff.lwdid_results import LWDiDResults + + assert not hasattr(LWDiDResults, "to_latex") + + def test_staggered_to_dataframe_carries_config_columns(self): + rng = np.random.default_rng(19) + rows = [] + for u in range(20): + g = 3 if u < 8 else 0 + alpha = rng.normal() + for t in range(1, 6): + d = int(g > 0 and t >= g) + rows.append(dict(unit=u, time=t, first=g, treat=d, y=alpha + rng.normal() + d)) + df = pd.DataFrame(rows) + res = LWDiD(rolling="demean", estimation_method="reg").fit( + df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="first" + ) + frame = res.to_dataframe() + for col in ("rolling", "estimation_method", "vcov_type"): + assert col in frame.columns + assert frame[col].nunique() == 1 + + def test_detrend_degenerate_cohort_composite_is_graceful(self): + # Campaign finding: a cohort with < 2 pre-periods crashed the + # composite with a raw LinAlgError under detrend + never_treated. + # The complete-case machinery now drops it with warnings. + rng = np.random.default_rng(2) + rows = [] + for u in range(20): + g = 2 if u < 4 else (5 if u < 9 else 0) # cohort 2: ONE pre period + alpha = rng.normal() + for t in range(1, 8): + d = int(g > 0 and t >= g) + rows.append( + dict(unit=u, time=t, first=g, treat=d, y=alpha + 0.1 * t + rng.normal() + d) + ) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD( + rolling="detrend", + estimation_method="reg", + vcov_type="classical", + control_group="never_treated", + ).fit(df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="first") + assert np.isfinite(res.att) + assert res.n_composite_treated_dropped == 4 + + +class TestNonNumericTimeContract: + """LWDiD fix-wave (campaign finding): string time columns made + detrend/demeanq/detrendq raise raw numpy conversion errors while + demean succeeded - now an informative ValueError states the contract. + """ + + @staticmethod + def _string_time_panel(): + rng = np.random.default_rng(1) + rows = [] + labels = ["Q1", "Q2", "Q3", "Q4", "Q5", "Q6"] + for u in range(10): + for i, t in enumerate(labels): + d = int(u < 5 and i >= 3) + rows.append(dict(unit=u, time=t, treat=d, y=rng.normal() + d)) + return pd.DataFrame(rows) + + def test_demean_accepts_ordered_categorical_time(self): + # Round-20 review: plain string labels sort lexicographically + # ('Q10' < 'Q2'), so the chronology must be DECLARED - demean + # accepts an ordered categorical and rejects plain object labels. + df = self._string_time_panel() + labels = ["Q1", "Q2", "Q3", "Q4", "Q5", "Q6"] + df_cat = df.assign(time=pd.Categorical(df["time"], categories=labels, ordered=True)) + res = LWDiD(rolling="demean", estimation_method="reg").fit( + df_cat, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert np.isfinite(res.att) + + def test_demean_rejects_plain_string_time(self): + df = self._string_time_panel() + with pytest.raises(ValueError, match="ORDERED categorical"): + LWDiD(rolling="demean", estimation_method="reg").fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + + @pytest.mark.parametrize("rolling", ["detrend", "demeanq", "detrendq"]) + def test_trend_seasonal_transforms_reject_string_time_informatively(self, rolling): + df = self._string_time_panel() + with pytest.raises(ValueError, match="numeric or datetime"): + LWDiD(rolling=rolling, estimation_method="reg").fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + + +class TestReviewRound1Guards: + """Local-review round 1 (fix wave): execution-verified guards. + + - HC3 with a leverage-one observation (single treated unit) fabricated + finite inference via a 1e-10 floor on 1 - h_ii + - N=2 / N=3,K=1 collapsed designs hit ZeroDivisionError or a coerced + df=1 instead of the Registry's small-sample guards + - n_bootstrap=1 was accepted; staggered PSM + bootstrap silently no-oped + - PSM's naive matched-pairs SE ignored control reuse and first-stage + uncertainty (now fail-closed NaN inference, point retained) + - the common-timing single-cluster fallback warned then raised + - Period time crashed detrend with a raw TypeError + """ + + @staticmethod + def _panel(n_units=12, n_treated=1, t_max=6, onset=4, seed=0, **cols): + rng = np.random.default_rng(seed) + rows = [] + for u in range(n_units): + treated = u < n_treated + for t in range(1, t_max + 1): + d = int(treated and t >= onset) + row = dict(unit=u, time=t, treat=d, y=rng.normal() + d) + for k, fn in cols.items(): + row[k] = fn(u) + rows.append(row) + return pd.DataFrame(rows) + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_hc3_leverage_one_fails_closed(self): + df = self._panel(n_units=12, n_treated=1) + with pytest.warns(UserWarning, match="HC3 variance is undefined"): + res = LWDiD(rolling="demean", estimation_method="reg", vcov_type="hc3").fit( + df, **self.KW + ) + assert np.isfinite(res.att) + from tests.conftest import assert_nan_inference + + assert_nan_inference( + {"se": res.se, "t_stat": res.t_stat, "p_value": res.p_value, "conf_int": res.conf_int} + ) + + def test_invalid_exact_designs_rejected(self): + df2 = self._panel(n_units=2, n_treated=1) + with pytest.raises(ValueError, match="Invalid exact-inference design"): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + LWDiD(rolling="demean", estimation_method="reg", vcov_type="classical").fit( + df2, **self.KW + ) + df3 = self._panel(n_units=3, n_treated=1, x=lambda u: float(u)) + with pytest.raises(ValueError, match="Invalid exact-inference design"): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + LWDiD(rolling="demean", estimation_method="reg", vcov_type="classical").fit( + df3, covariates=["x"], **self.KW + ) + + def test_n_bootstrap_one_rejected(self): + with pytest.raises(ValueError, match="n_bootstrap must be 0"): + LWDiD(n_bootstrap=1) + + def test_staggered_psm_bootstrap_rejected(self): + df = self._panel(n_units=16, n_treated=6) + df["first"] = np.where(df["unit"] < 6, 4, 0) + est = LWDiD(estimation_method="psm", n_bootstrap=50) + with pytest.raises(ValueError, match="psm.*does not support n_bootstrap"): + est.fit(df, first_treat="first", **self.KW) + + def test_psm_inference_fails_closed_point_retained(self): + df = self._panel(n_units=20, n_treated=8, x=lambda u: float(u % 4)) + with pytest.warns(UserWarning, match="no valid matching variance"): + res = LWDiD(rolling="demean", estimation_method="psm").fit( + df, covariates=["x"], **self.KW + ) + assert np.isfinite(res.att) + assert np.isnan(res.se) and np.isnan(res.p_value) + assert res.psm_config is not None + assert res.psm_config["n_neighbors"] == 1 + + def test_matching_params_strictly_validated(self): + with pytest.raises(ValueError, match="n_neighbors must be an integer"): + LWDiD(n_neighbors=1.5) + with pytest.raises(ValueError, match="with_replacement must be a boolean"): + LWDiD(with_replacement="yes") + with pytest.raises(ValueError, match="caliper must be a positive"): + LWDiD(caliper=-0.1) + + def test_common_single_cluster_post_drop_point_retained(self): + rng = np.random.default_rng(0) + rows = [] + for u in range(10): + for t in range(1, 7): + d = int(u < 5 and t >= 4) + rows.append(dict(unit=u, time=t, treat=d, y=rng.normal() + d, cl=0 if u < 9 else 1)) + df = pd.DataFrame(rows) + df = df.loc[~((df.unit == 9) & (df.time.isin([2, 3])))] + with pytest.warns(UserWarning, match="fewer than 2 clusters"): + res = LWDiD(rolling="detrend", estimation_method="reg", cluster="cl").fit(df, **self.KW) + assert np.isfinite(res.att) + assert np.isnan(res.se) + + def test_period_time_detrend_rejected_informatively(self): + rng = np.random.default_rng(1) + times = pd.period_range("2020Q1", periods=8, freq="Q") + rows = [] + for u in range(10): + for i, t in enumerate(times): + d = int(u < 5 and i >= 5) + rows.append(dict(unit=u, time=t, treat=d, y=rng.normal() + d)) + df = pd.DataFrame(rows) + with pytest.raises(ValueError, match="does not support a Period"): + LWDiD(rolling="detrend", estimation_method="reg").fit(df, **self.KW) + + def test_provenance_fields_round_trip(self): + df = self._panel(n_units=12, n_treated=5) + res = LWDiD( + rolling="demean", + estimation_method="reg", + control_group="never_treated", + n_bootstrap=0, + seed=7, + ).fit(df, **self.KW) + assert res.control_group == "never_treated" + assert res.n_bootstrap == 0 + assert res.seed == 7 + assert res.psm_config is None + d = res.to_dict() + assert d["control_group"] == "never_treated" + assert d["seed"] == 7 + + +class TestReviewRound2Guards: + """Local-review round 2: execution-verified guards. + + - cluster='_treat' with numeric labels silently reported the cluster + labels' coefficient as the ATT (reserved-name collision) + - the same seed produced different bootstrap SEs across n_jobs (the + serial path consumed one sequential RNG stream while the parallel + path spawned per-replicate streams) + - pscore_trim was absent from ipw/dr result provenance + - an event cell with degenerate multiplier-bootstrap draws silently + kept its analytical SE (undocumented mixture of inference families) + """ + + @staticmethod + def _panel(n_units=12, n_treated=6, t_max=6, onset=4, seed=0, **cols): + rng = np.random.default_rng(seed) + rows = [] + for u in range(n_units): + treated = u < n_treated + for t in range(1, t_max + 1): + d = int(treated and t >= onset) + row = dict(unit=u, time=t, treat=d, y=rng.normal() + 2 * d + 0.5 * t) + for k, fn in cols.items(): + row[k] = fn(u) + rows.append(row) + return pd.DataFrame(rows) + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_reserved_internal_names_rejected(self): + df = self._panel(cl=lambda u: float(u % 3)).rename(columns={"cl": "_treat"}) + with pytest.raises(ValueError, match="reserved for LWDiD internal use"): + LWDiD(rolling="demean", cluster="_treat").fit(df, **self.KW) + df2 = self._panel(x=lambda u: float(u)).rename(columns={"x": "_ydot"}) + with pytest.raises(ValueError, match="reserved for LWDiD internal use"): + LWDiD(rolling="demean").fit(df2, covariates=["_ydot"], **self.KW) + df3 = self._panel().rename(columns={"unit": "_boot_unit"}) + with pytest.raises(ValueError, match="reserved for LWDiD internal use"): + LWDiD(rolling="demean").fit( + df3, outcome="y", unit="_boot_unit", time="time", treatment="treat" + ) + + def test_duplicate_role_columns_rejected(self): + df = self._panel() + with pytest.raises(ValueError, match="distinct column"): + LWDiD(rolling="demean").fit(df, outcome="y", unit="unit", time="time", treatment="y") + df2 = self._panel(x=lambda u: float(u)) + with pytest.raises(ValueError, match="already supplied"): + LWDiD(rolling="demean").fit(df2, covariates=["y"], **self.KW) + # cluster == unit stays supported (documented intentional case) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean", cluster="unit").fit(df, **self.KW) + assert np.isfinite(res.att) + + def test_seeded_bootstrap_invariant_to_n_jobs(self): + df = self._panel(n_units=20, n_treated=10) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r1 = LWDiD(rolling="demean", n_bootstrap=49, seed=7, n_jobs=1).fit(df, **self.KW) + r2 = LWDiD(rolling="demean", n_bootstrap=49, seed=7, n_jobs=2).fit(df, **self.KW) + assert r1.se == r2.se + assert r1.att == r2.att + + def test_pscore_trim_provenance(self): + df = self._panel(x=lambda u: float(u % 4)) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ipw = LWDiD(rolling="demean", estimation_method="ipw", pscore_trim=0.02).fit( + df, covariates=["x"], **self.KW + ) + reg = LWDiD(rolling="demean", estimation_method="reg").fit(df, **self.KW) + assert ipw.pscore_trim == 0.02 + assert ipw.to_dict()["pscore_trim"] == 0.02 + assert reg.pscore_trim is None + assert "pscore_trim" not in reg.to_dict() + + def test_degenerate_bootstrap_event_cell_fails_closed(self): + from types import SimpleNamespace + + from diff_diff.lwdid_staggered import compute_event_study_bands + + rng = np.random.default_rng(0) + estimator = SimpleNamespace(n_bootstrap=199, seed=3, alpha=0.05) + event_effects = { + 0: { + "effect": 1.0, + "se": 0.2, + "t_stat": 5.0, + "p_value": 0.0, + "conf_int": (0.6, 1.4), + "df": None, + }, + 1: { + "effect": 0.5, + "se": 0.1, + "t_stat": 5.0, + "p_value": 0.0, + "conf_int": (0.3, 0.7), + "df": None, + }, + } + event_influence = { + 0: np.zeros(30), # degenerate: zero influence column + 1: rng.normal(size=30), + } + with pytest.warns(UserWarning, match="degenerate draws"): + compute_event_study_bands(estimator, event_effects, event_influence, None) + assert np.isnan(event_effects[0]["se"]) + assert np.isnan(event_effects[0]["p_value"]) + assert event_effects[0]["inference_status"] == "degenerate_bootstrap" + assert event_effects[0]["effect"] == 1.0 # point retained + assert np.isfinite(event_effects[1]["se"]) # valid cell bootstrapped + assert "cband_conf_int" in event_effects[1] + + +class TestReviewRound3Guards: + """Local-review round 3: execution-verified guards. + + - encoded staggered panels fed dense POSITIONS to the seasonal + transforms' (t-1)%4+1 fallback, so a globally missing calendar + quarter silently relabeled every later season (probe: zero-effect + gapped quarterly panel with treated/control-differential + seasonality biased demeanq ATT to ~0.12) + - PSM bypassed its fail-closed NaN-inference contract via the + covariate-less delegation, the common-timing pairs bootstrap, and + the logit-failure regression fallback + - common-timing bootstrap fits reported bootstrap headline inference + with no provenance while params/vcov stayed analytical + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + @staticmethod + def _gapped_quarterly(as_period=True): + rng = np.random.default_rng(5) + periods = pd.period_range("2018Q1", "2023Q4", freq="Q") + periods = periods[periods != pd.Period("2020Q3", freq="Q")] + seas = {1: 2.0, 2: -1.0, 3: 0.5, 4: -1.5} + onset = pd.Period("2022Q1", freq="Q") + rows = [] + for u in range(24): + treated_unit = u < 12 + amp = 3.0 if treated_unit else 1.0 + for p in periods: + d = int(treated_unit and p >= onset) + y = 1.0 + amp * seas[p.quarter] + rng.normal(0, 0.1) + rows.append(dict(unit=u, p=p, y=y, treat=d, g=onset if treated_unit else pd.NaT)) + df = pd.DataFrame(rows) + if as_period: + df["time"] = pd.PeriodIndex(df["p"], freq="Q") + df["gv"] = pd.PeriodIndex(df["g"], freq="Q") + else: + # ordinal numeric encoding preserves calendar-quarter identity + # under (t-1)%4+1 (Q ordinals advance one per quarter), so the + # numeric path is the correct-season oracle + df["time"] = pd.PeriodIndex(df["p"], freq="Q").map(lambda v: v.ordinal + 1) + df["gv"] = [pd.Period(v, freq="Q").ordinal + 1 if pd.notna(v) else 0 for v in df["g"]] + return df.drop(columns=["p", "g"]) + + def test_gapped_calendar_seasonal_parity(self): + # Period path (encoded to dense positions) must match the + # ordinal-numeric oracle on the SAME gapped data. demeanq is + # trend-free, so seasonal-grouping parity is exact. detrendq + # additionally uses the time values as its trend coordinate + # (dense positions on the encoded path vs calendar ordinals on + # the numeric path - a documented encoding choice), so its pin is + # the no-seasonal-leakage bound, not bitwise parity. + got, want = {}, {} + for rolling in ("demeanq", "detrendq"): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r_period = LWDiD(rolling=rolling).fit( + self._gapped_quarterly(True), first_treat="gv", **self.KW + ) + r_numeric = LWDiD(rolling=rolling).fit( + self._gapped_quarterly(False), first_treat="gv", **self.KW + ) + got[rolling], want[rolling] = r_period.att, r_numeric.att + np.testing.assert_allclose(got["demeanq"], want["demeanq"], rtol=1e-10) + for rolling in got: + # zero-effect DGP with 3x treated seasonal amplitude: the + # pre-fix position-modulo labeling biased this to ~0.12 + assert abs(got[rolling]) < 0.06, rolling + + @staticmethod + def _panel(n_units=16, x=True): + rng = np.random.default_rng(0) + rows = [] + for u in range(n_units): + for t in range(1, 7): + d = 1 if (u < n_units // 2 and t >= 4) else 0 + row = dict(unit=u, time=t, treat=d, y=1 + 0.5 * t + 2 * d + rng.normal(0, 0.5)) + if x: + row["x"] = float(u % 4) + rows.append(row) + return pd.DataFrame(rows) + + def test_psm_requires_covariates(self): + with pytest.raises(ValueError, match="requires covariates"): + LWDiD(rolling="demean", estimation_method="psm").fit(self._panel(x=False), **self.KW) + + def test_psm_bootstrap_rejected_common_timing(self): + with pytest.raises(ValueError, match="does not support n_bootstrap"): + LWDiD(rolling="demean", estimation_method="psm", n_bootstrap=50).fit( + self._panel(), covariates=["x"], **self.KW + ) + + def test_psm_logit_failure_fails_closed(self, monkeypatch): + import diff_diff.lwdid as lwdid_mod + + est = LWDiD(rolling="demean", estimation_method="psm") + rng = np.random.default_rng(1) + y = rng.normal(size=20) + treatment = np.array([1.0] * 8 + [0.0] * 12) + controls = rng.normal(size=(20, 1)) + # Non-finite PROBABILITIES = genuine solver failure (round 19: + # NaN coefs with finite probs is now the reduced-rank + # continuation path, not the fallback). + monkeypatch.setattr( + lwdid_mod, + "solve_logit", + lambda X, d: (np.array([np.nan, np.nan]), np.full(len(d), np.nan)), + ) + with pytest.warns(UserWarning, match="PSM fail-closed"): + att, se, coefs, vcov, _, influence = est._estimate_psm(y, treatment, controls, None, 20) + assert np.isfinite(att) + assert np.isnan(se) + assert coefs is None and vcov is None and influence is None + + def test_bootstrap_inference_basis_provenance(self): + df = self._panel(x=False) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + plain = LWDiD(rolling="demean").fit(df, **self.KW) + boot = LWDiD(rolling="demean", n_bootstrap=49, seed=3).fit(df, **self.KW) + assert plain.inference_basis is None + assert boot.inference_basis == "unit_bootstrap" + assert "bootstrap" in boot.summary() + assert boot.to_dict()["inference_basis"] == "unit_bootstrap" + df["cl"] = df["unit"] % 4 + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + cboot = LWDiD(rolling="demean", n_bootstrap=49, seed=3, cluster="cl").fit(df, **self.KW) + assert cboot.inference_basis == "cluster_bootstrap" + + def test_diagnostics_shares_fit_validation(self): + df = self._panel(x=False) + bad = df.copy() + bad["treat"] = bad["treat"] * 2 # non-binary + with pytest.raises(ValueError): + LWDiD(rolling="demean").get_transformation_diagnostics( + bad, outcome="y", unit="unit", time="time", treatment="treat" + ) + dup = pd.concat([df, df.iloc[:6]]) # duplicate unit-time rows + with pytest.raises(ValueError, match="duplicate"): + LWDiD(rolling="demean").get_transformation_diagnostics( + dup, outcome="y", unit="unit", time="time", treatment="treat" + ) + + def test_duplicate_covariates_rejected(self): + with pytest.raises(ValueError, match="duplicate column"): + LWDiD(rolling="demean").fit(self._panel(), covariates=["x", "x"], **self.KW) + + +class TestReviewRound4Guards: + """Local-review round 4: execution-verified guards. + + - RI and the WCR wrapper fit via np.linalg.lstsq, so a control + duplicating treatment returned a finite MINIMUM-NORM ATT (probe: + true ATT 2.0 reported as 0.877 with finite p-values in both) + - staggered NT-only cells could estimate on a single surviving + control after transformation drops + - unobserved staggered anchors were synthesized as zero-valued + reference rows + - the standalone WCR wrapper accepted n_bootstrap=1; the result-level + wrapper ignored the fitted alpha + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + @staticmethod + def _arrays(seed=0, n=40): + rng = np.random.default_rng(seed) + treat = np.array([1.0] * (n // 2) + [0.0] * (n // 2)) + y = 2.0 * treat + rng.normal(0, 1, n) + cl = np.arange(n) % 8 + return y, treat, cl, rng + + def test_collinear_control_identified_in_ri_and_wcb(self): + # Pre-fix, lstsq split the effect across the duplicate columns + # (minimum-norm: true ATT 2.0 reported as ~0.88). The shared + # rank-aware solver pivots, keeps the treatment column, drops the + # duplicate control, and reports the IDENTIFIED ATT (here exactly + # the difference in means) with a rank warning. The raise branch + # remains as a backstop should the treatment column itself be + # pivoted out. + from diff_diff.lwdid_randomization import randomization_inference + from diff_diff.lwdid_wild_bootstrap import wild_cluster_bootstrap + + y, treat, cl, rng = self._arrays() + dup = treat.reshape(-1, 1).copy() + truth = y[treat == 1].mean() - y[treat == 0].mean() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + ri = randomization_inference(y, treat, controls=dup, n_reps=49, seed=3) + wb = wild_cluster_bootstrap(y, treat, cl, controls=dup, n_bootstrap=49, seed=3) + np.testing.assert_allclose(ri.att_observed, truth, rtol=1e-12) + np.testing.assert_allclose(wb.att, truth, rtol=1e-12) + assert any("rank" in str(x.message).lower() for x in caught) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + # valid controls: studentization stays coherent + x = rng.normal(size=(len(y), 1)) + wb2 = wild_cluster_bootstrap(y, treat, cl, controls=x, n_bootstrap=49, seed=3) + np.testing.assert_allclose(wb2.t_stat_original, wb2.att / wb2.se) + + def test_wcb_n_bootstrap_validation(self): + from diff_diff.lwdid_wild_bootstrap import wild_cluster_bootstrap + + y, treat, cl, _ = self._arrays() + for bad in (1, 0, -3, 2.5, True): + with pytest.raises(ValueError, match="integer >= 2"): + wild_cluster_bootstrap(y, treat, cl, n_bootstrap=bad) + + def test_ri_n_reps_validation(self): + from diff_diff.lwdid_randomization import randomization_inference + + y, treat, _, _ = self._arrays() + for bad in (0, -1, 99.5, True): + with pytest.raises(ValueError, match="n_reps must be"): + randomization_inference(y, treat, n_reps=bad) + + def test_results_wcb_inherits_fitted_alpha(self): + rng = np.random.default_rng(0) + rows = [] + for u in range(16): + for t in range(1, 7): + d = 1 if (u < 8 and t >= 4) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=1 + 2 * d + rng.normal(0, 0.5))) + df = pd.DataFrame(rows) + df["cl"] = df["unit"] % 4 + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean", alpha=0.10, cluster="cl").fit(df, **self.KW) + wb = res.wild_cluster_bootstrap(n_bootstrap=49, seed=1) + assert wb.alpha == 0.10 + wb2 = res.wild_cluster_bootstrap(n_bootstrap=49, seed=1, alpha=0.05) + assert wb2.alpha == 0.05 + + def test_nt_only_cell_needs_two_surviving_controls(self): + rng = np.random.default_rng(0) + rows = [] + for u in range(8): + g = 4 if u < 6 else 0 # 6 treated, exactly 2 never-treated + for t in range(1, 7): + d = int(g > 0 and t >= g) + y = 1 + 0.5 * t + d + rng.normal(0, 0.3) + if u == 6 and t == 5: + y = np.nan # one NT control loses its t=5 outcome + rows.append(dict(unit=u, time=t, treat=d, g=g, y=y)) + df = pd.DataFrame(rows).dropna(subset=[]).copy() + df.loc[(df.unit == 6) & (df.time == 5), "y"] = np.nan + df = df.dropna(subset=["y"]) if False else df + # NaN y raises in validation; drop the row instead (unbalanced panel) + df = df[~((df.unit == 6) & (df.time == 5))] + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean", control_group="never_treated").fit( + df, first_treat="g", **self.KW + ) + cell = res.cohort_time_effects[(4, 5)] + assert cell["skip_reason"] == "insufficient_never_treated_controls" + assert np.isnan(cell["att"]) + # other post cells still estimated with both controls + assert np.isfinite(res.cohort_time_effects[(4, 4)]["att"]) + + def test_unobserved_anchor_not_synthesized(self): + rng = np.random.default_rng(0) + rows = [] + times = [1, 2, 3, 5, 6] # time 4 (= g-1 anchor for g=5) missing + for u in range(10): + g = 5 if u < 5 else 0 + for t in times: + d = int(g > 0 and t >= g) + rows.append( + dict(unit=u, time=t, treat=d, g=g, y=1 + 0.3 * t + d + rng.normal(0, 0.3)) + ) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean", control_group="never_treated").fit( + df, first_treat="g", **self.KW + ) + assert res.reference_periods == () # anchor r=-1 unobserved -> not emitted + + +class TestReviewRound5Guards: + """Local-review round 5: execution-verified guards. + + - post-fit WCR/RI accepted arbitrary arrays + a non-interacted design, + caching p-values for a DIFFERENT estimand than .att (probe: fitted + 3.98 vs tested 3.26 on a covariate-unbalanced RA fit) - now replay + the fit spec (pinned in test_lwdid_wild_bootstrap.py) + - a requested bootstrap overwrote the single-effective-cluster + fail-closed NaN inference with a near-zero SE from the raw cluster + map + - aggregate(balance_e=) was accepted but silently ignored + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_bootstrap_preserves_single_cluster_fail_closed(self): + rng = np.random.default_rng(0) + rows = [] + for u in range(12): + cl = 0 if u < 6 else 1 + # cluster 1 units observe only ONE pre period: detrend needs 2, + # so their transformed outcomes are NaN and the whole cluster + # drops from the collapsed cross-section + times = range(3, 7) if cl == 1 else range(1, 7) + for t in times: + d = 1 if (u % 6 < 3 and t >= 4) else 0 + rows.append( + dict(unit=u, time=t, treat=d, cl=cl, y=1 + 0.5 * t + 2 * d + rng.normal(0, 0.3)) + ) + df = pd.DataFrame(rows) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + res = LWDiD(rolling="detrend", cluster="cl", n_bootstrap=49, seed=1).fit(df, **self.KW) + assert np.isfinite(res.att) + from tests.conftest import assert_nan_inference + + assert_nan_inference( + {"se": res.se, "t_stat": res.t_stat, "p_value": res.p_value, "conf_int": res.conf_int} + ) + assert any("bootstrap skipped" in str(x.message) for x in caught) + assert res.inference_basis is None # no bootstrap ran + + def test_balance_e_rejected(self): + rng = np.random.default_rng(0) + rows = [] + for u in range(16): + g = 4 if u < 4 else (5 if u < 8 else 0) + for t in range(1, 8): + d = int(g > 0 and t >= g) + rows.append( + dict(unit=u, time=t, treat=d, g=g, y=1 + 0.3 * t + d + rng.normal(0, 0.3)) + ) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean").fit(df, first_treat="g", **self.KW) + with pytest.raises((ValueError, TypeError), match="balance_e"): + res.aggregate("event_study", balance_e=1) + # without balance_e the aggregation still works + assert res.aggregate("event_study") is not None + + +class TestReviewRound6Guards: + """Local-review round 6: execution-verified guards. + + - fweight + HC2/HC3 used the WLS-hat (weighted) leverage, so the + compressed variance was up to ~5x the literal np.repeat expansion + (fweights are replicated data by definition) + - the common-timing headline averaged whichever post periods each + unit observed, letting calendar composition masquerade as ATT + - the degenerate-SE guard's max(1, |effect|) floor NaN'd valid + inference under outcome rescaling + - complete-case drops could empty an arm and dispatch a one-arm design + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_fweight_hc2_hc3_expansion_parity(self): + from diff_diff.linalg import solve_ols + + X = np.column_stack([np.ones(5), np.arange(5.0)]) + y = np.array([1.0, 2.2, 2.9, 4.1, 5.3]) + w = np.array([3.0, 1.0, 2.0, 1.0, 4.0]) + Xe = np.repeat(X, w.astype(int), axis=0) + ye = np.repeat(y, w.astype(int)) + for vt in ("hc2", "hc3"): + _, _, v_c = solve_ols( + X, y, return_vcov=True, vcov_type=vt, weights=w, weight_type="fweight" + ) + _, _, v_e = solve_ols(Xe, ye, return_vcov=True, vcov_type=vt) + np.testing.assert_allclose(np.diag(v_c), np.diag(v_e), rtol=1e-12, err_msg=vt) + + def test_fixed_window_complete_case_headline(self): + # Zero-effect panel, Y_it = t: half the controls miss the last + # post period. Pre-fix their shorter post average biased the + # headline; complete-case drops them (warned) and ATT ~ 0. + rows = [] + for u in range(12): + treated = u < 6 + t_max = 6 if (treated or u < 9) else 5 # controls 9-11 miss t=6 + for t in range(1, t_max + 1): + d = 1 if (treated and t >= 4) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=float(t))) + df = pd.DataFrame(rows) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + res = LWDiD(rolling="demean").fit(df, **self.KW) + assert any("fixed-window" in str(x.message) for x in caught) + np.testing.assert_allclose(res.att, 0.0, atol=1e-10) + assert res.n_control == 3 # complete controls only + + def test_se_guard_scale_equivariant(self): + rng = np.random.default_rng(0) + rows = [] + for u in range(16): + g = 4 if u < 8 else 0 + for t in range(1, 7): + d = int(g > 0 and t >= g) + rows.append( + dict(unit=u, time=t, treat=d, g=g, y=1 + 0.3 * t + d + rng.normal(0, 0.3)) + ) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r1 = LWDiD(rolling="demean").fit(df, first_treat="g", **self.KW) + df2 = df.assign(y=df["y"] * 1e-10) + r2 = LWDiD(rolling="demean").fit(df2, first_treat="g", **self.KW) + # t-statistic is invariant to outcome rescaling + np.testing.assert_allclose(r2.t_stat, r1.t_stat, rtol=1e-8) + np.testing.assert_allclose(r2.att, r1.att * 1e-10, rtol=1e-8) + assert np.isfinite(r2.se) + + def test_empty_arm_after_drops_raises(self): + # All treated units observe only one pre period -> detrend NaNs + # every treated unit; pre-fix a one-arm design was dispatched. + rng = np.random.default_rng(0) + rows = [] + for u in range(8): + treated = u < 4 + times = range(3, 7) if treated else range(1, 7) + for t in times: + d = 1 if (treated and t >= 4) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=1 + 0.5 * t + rng.normal(0, 0.3))) + df = pd.DataFrame(rows) + with pytest.raises(ValueError, match="at\\s+least one of each"): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + LWDiD(rolling="detrend").fit(df, **self.KW) + + def test_constructor_numeric_validation(self): + with pytest.raises(ValueError, match="pscore_trim"): + LWDiD(pscore_trim=True) + with pytest.raises(ValueError, match="pscore_trim"): + LWDiD(pscore_trim="0.1") + with pytest.raises(ValueError, match="n_jobs"): + LWDiD(n_jobs=True) + + def test_df_inference_serializes(self): + rng = np.random.default_rng(0) + rows = [] + for u in range(12): + for t in range(1, 7): + d = 1 if (u < 6 and t >= 4) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=1 + 2 * d + rng.normal(0, 0.4))) + res = LWDiD(rolling="demean").fit(pd.DataFrame(rows), **self.KW) + d = res.to_dict() + assert d["df_inference"] == res.df_inference + + +class TestReviewRound7Guards: + """Local-review round 7: execution-verified guards. + + - the post-fit replay always rebuilt the covariate interactions, but + the fit uses plain (1, D, X) when an arm has N <= K+1 (LW eq. 3.3 + gate) - small-arm fits' replayed statistic mismatched .att and the + round-5 coherence assert made their inference unusable + - the sensitivity helpers swallowed treatment-design violations + (absorbing/onset/cohort-consistency) as not_estimable specs + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_small_arm_plain_design_replay(self): + # K=1 covariate, exactly 2 treated units: n_treated <= K+1, so the + # fit uses the plain design; the replay must follow it. + rng = np.random.default_rng(3) + rows = [] + for u in range(12): + treated = u < 2 + x = float(u % 4) + for t in range(1, 7): + d = 1 if (treated and t >= 4) else 0 + rows.append( + dict( + unit=u, + time=t, + treat=d, + x=x, + cl=u % 4, + y=1 + 0.4 * x + 2 * d + rng.normal(0, 0.3), + ) + ) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean", cluster="cl").fit(df, covariates=["x"], **self.KW) + ri = res.randomization_test(n_reps=99, seed=1) + wb = res.wild_cluster_bootstrap(n_bootstrap=49, seed=1) + np.testing.assert_allclose(ri.att_observed, res.att, rtol=1e-10) + np.testing.assert_allclose(wb.att, res.att, rtol=1e-10) + + def test_sensitivity_rejects_design_violations(self): + from diff_diff.lwdid_sensitivity import ( + robustness_pre_periods, + sensitivity_no_anticipation, + ) + + rng = np.random.default_rng(0) + rows = [] + for u in range(10): + for t in range(1, 9): + d = 1 if (u < 5 and t >= 6) else 0 + if u == 0 and t == 7: + d = 0 # 1 -> 0 reversal: non-absorbing treatment + rows.append(dict(unit=u, time=t, treat=d, y=rng.normal() + d)) + df = pd.DataFrame(rows) + for fn in (robustness_pre_periods, sensitivity_no_anticipation): + with pytest.raises(ValueError, match="absorbing|revert"): + fn(df, outcome="y", unit="unit", time="time", treatment="treat") + # heterogeneous onsets without first_treat: also a raise, not + # a silent not_estimable + rows2 = [] + for u in range(10): + onset = 5 if u < 3 else (6 if u < 5 else 99) + for t in range(1, 9): + d = 1 if (u < 5 and t >= onset) else 0 + rows2.append(dict(unit=u, time=t, treat=d, y=rng.normal() + d)) + df2 = pd.DataFrame(rows2) + with pytest.raises(ValueError, match="common timing|first_treat|onset"): + robustness_pre_periods(df2, outcome="y", unit="unit", time="time", treatment="treat") + + +class TestReviewRound8Guards: + """Local-review round 8: execution-verified guards. + + - the per-period max(D) partition classified a post period with no + observed treated rows as PRE-treatment (zero-effect trend probe: + ATT 0.75); the partition now derives from the single onset S + - the common-timing onset check rejected units missing their t = S + row as heterogeneous timing (the staggered branch permits it) + - the tau_omega completeness check accepted any finite post average, + so a control observing part of a cohort's window survived + - Inf covariates passed the NaN check; non-numeric covariates crashed + with raw conversion errors + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + @staticmethod + def _trend_panel(drop=lambda u, t: False, n_units=12, onset=4, t_max=6, effect=0.0): + rows = [] + for u in range(n_units): + treated = u < n_units // 2 + for t in range(1, t_max + 1): + if drop(u, t) and treated: + continue + d = 1 if (treated and t >= onset) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=float(t) + effect * d)) + return pd.DataFrame(rows) + + def test_controls_only_post_period_not_misclassified(self): + # 2 of 6 treated units miss post period t=5: pre-fix that period + # was classified as pre (contaminating the pre window, ATT 0.75 on + # this zero-effect trend); now it stays post and the incomplete + # treated units are complete-case dropped. + df = self._trend_panel(drop=lambda u, t: u < 2 and t == 5) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean").fit(df, **self.KW) + np.testing.assert_allclose(res.att, 0.0, atol=1e-10) + assert res.n_treated == 4 + # ALL treated missing the period -> no fixed-window comparison + df_all = self._trend_panel(drop=lambda u, t: t == 5) + with pytest.raises(ValueError, match="at\\s+least one of each"): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + LWDiD(rolling="demean").fit(df_all, **self.KW) + + def test_missing_onset_row_accepted_common_timing(self): + df = self._trend_panel(drop=lambda u, t: u == 0 and t == 4, effect=2.0) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean").fit(df, **self.KW) + np.testing.assert_allclose(res.att, 2.0, atol=1e-10) + # genuinely heterogeneous onsets still rejected + rows = [] + for u in range(8): + onset = 4 if u < 2 else (5 if u < 4 else 99) + for t in range(1, 7): + d = 1 if (u < 4 and t >= onset) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=float(t))) + with pytest.raises(ValueError, match="heterogeneous|common onset"): + LWDiD(rolling="demean").fit(pd.DataFrame(rows), **self.KW) + + def test_tau_omega_partial_window_semantics_pinned(self): + # ADJUDICATED (round 8): completeness = a finite average over the + # OBSERVED post-g rows, symmetric across arms - a control missing + # one window period is RETAINED (its component averages observed + # rows); a control missing the ENTIRE window is dropped. The + # acceptance suite's frozen reference oracle pins the same rule. + def build(missing): + rows = [] + for u in range(12): + g = 5 if u < 6 else 0 + for t in range(1, 9): + if u == 11 and t in missing: + continue + d = int(g > 0 and t >= g) + rows.append(dict(unit=u, time=t, treat=d, g=g, y=float(t))) + return pd.DataFrame(rows) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + part = LWDiD( + rolling="demean", control_group="never_treated", vcov_type="classical" + ).fit(build({8}), first_treat="g", **self.KW) + whole = LWDiD( + rolling="demean", control_group="never_treated", vcov_type="classical" + ).fit(build({5, 6, 7, 8}), first_treat="g", **self.KW) + assert part.n_composite_controls_dropped == 0 # partial window retained + assert whole.n_composite_controls_dropped == 1 # entire window missing + # Documented composition caveat, pinned exactly: with y = t and + # demean (cohort-5 pre-mean 2.5), complete units average ydot + # over t=5..8 (=4.0) while the partial control averages t=5..7 + # (=3.5), so tau_omega = 4.0 - (5*4.0 + 3.5)/6 = 1/12. + np.testing.assert_allclose(part.att, 1.0 / 12.0, atol=1e-10) + np.testing.assert_allclose(whole.att, 0.0, atol=1e-10) + + def test_nonfinite_and_nonnumeric_covariates_rejected(self): + df = self._trend_panel() + df["x"] = 1.0 + df.loc[df.index[3], "x"] = np.inf + with pytest.raises(ValueError, match="non-finite"): + LWDiD(rolling="demean").fit(df, covariates=["x"], **self.KW) + df["x2"] = "a" + with pytest.raises(ValueError, match="not numeric"): + df_ok = df.assign(x=1.0) + LWDiD(rolling="demean").fit(df_ok, covariates=["x2"], **self.KW) + + def test_validate_staggered_data_rejects_mixed_families(self): + from diff_diff.lwdid import validate_staggered_data + + rows = [] + for u in range(6): + g = pd.Period("2020Q1", freq="Q") if u < 3 else pd.NaT + for i, ts in enumerate(pd.date_range("2019-01-01", periods=6, freq="QS")): + d = int(u < 3 and i >= 4) + rows.append(dict(unit=u, time=ts, treat=d, g=g, y=float(i))) + df = pd.DataFrame(rows) + df["g"] = pd.PeriodIndex(df["g"], freq="Q") + out = validate_staggered_data(df, unit="unit", time="time", cohort="g") + assert out["valid"] is False + assert any("same time scale" in e for e in out["errors"]) + + +class TestReviewRound9Guards: + """Local-review round 9: execution-verified guards. + + - staggered aggregation stored effects under int(t - g), silently + merging distinct fractional horizons; the common interface used + positional labels while staggered numeric used arithmetic, so the + same gapped design got different event keys per interface + - the round-8 onset partition was not propagated to diagnostics and + the sensitivity helpers + - Inf outcomes passed the NaN check and were silently cell-filtered + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_fractional_horizons_fail_closed(self): + rows = [] + times = [0.5, 1.0, 1.5, 2.0, 2.5] + for u in range(10): + g = 1.5 if u < 5 else 0 + for t in times: + d = int(g > 0 and t >= g) + rows.append(dict(unit=u, time=t, treat=d, g=g, y=float(t) + d)) + df = pd.DataFrame(rows) + with pytest.raises(ValueError, match="not an integer"): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + LWDiD(rolling="demean").fit(df, first_treat="g", **self.KW) + + def test_gapped_calendar_common_staggered_label_parity(self): + # {1, 2, 4, 6} with onset 4: both interfaces must label events by + # arithmetic t - g on numeric calendars (pre-fix: common reported + # {0, 1} positional while staggered reported {0, 2}). + rows = [] + for u in range(12): + treated = u < 6 + for t in (1, 2, 4, 6): + d = 1 if (treated and t >= 4) else 0 + rows.append( + dict(unit=u, time=t, treat=d, g=4 if treated else 0, y=float(t) + 2 * d) + ) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + common = LWDiD(rolling="demean").fit(df, **self.KW) + stag = LWDiD(rolling="demean", control_group="never_treated").fit( + df, first_treat="g", **self.KW + ) + common_post = sorted(common.event_study_effects) + stag_post = sorted(k for k, v in stag.event_study_effects.items() if k >= 0) + assert common_post == [0, 2] + assert stag_post == [0, 2] + + def test_diagnostics_and_sensitivity_use_onset_partition(self): + from diff_diff.lwdid_sensitivity import _get_pre_periods + + # controls-only post period t=5 (all treated rows missing there) + rows = [] + for u in range(12): + treated = u < 6 + for t in range(1, 7): + if treated and t == 5: + continue + d = 1 if (treated and t >= 4) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=float(t))) + df = pd.DataFrame(rows) + pre = _get_pre_periods(df, "time", "treat") + assert list(pre) == [1, 2, 3] # t=5 stays POST despite no treated rows + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + diag = LWDiD(rolling="demean").get_transformation_diagnostics( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + # control units' pre window excludes t=5: with y = t the pre mean + # over {1,2,3} is 2.0 for every unit + assert diag is not None + + def test_inf_outcome_rejected(self): + rows = [] + for u in range(8): + for t in range(1, 7): + d = 1 if (u < 4 and t >= 4) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=1.0 + d)) + df = pd.DataFrame(rows) + df.loc[df.index[5], "y"] = np.inf + with pytest.raises(ValueError, match="non-finite"): + LWDiD(rolling="demean").fit(df, **self.KW) + df["g"] = np.where(df["unit"] < 4, 4, 0) + with pytest.raises(ValueError, match="non-finite"): + LWDiD(rolling="demean").fit(df, first_treat="g", **self.KW) + + +class TestReviewRound10Guards: + """Local-review round 10: execution-verified guards. + + - the generic over-one-leverage HC1 fallback ran BEFORE hc3's + fail-closed check, so numerically over-one designs got an HC1 + result still labeled hc3 (and a clipped hc3 influence vector) + - the sensitivity multi-cohort count used RAW cohorts, rejecting + valid single-cohort designs with beyond-window encodings + - zero-post-row units evaded the fixed-window drop warning + - baseline sensitivity fits swallowed config errors as not_estimable + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_hc3_over_one_leverage_fails_closed(self): + from diff_diff.linalg import compute_robust_vcov + + rng = np.random.default_rng(0) + # near-duplicate rows -> numerically over-one leverage is hard to + # force deterministically; drive the guard directly with h >= 1 + X = np.column_stack([np.ones(4), np.array([0.0, 0.0, 0.0, 1.0])]) + y = np.array([1.0, 1.1, 0.9, 5.0]) + resid = y - X @ np.linalg.lstsq(X, y, rcond=None)[0] + with pytest.warns(UserWarning, match="HC3 variance is undefined"): + v = compute_robust_vcov(X, resid, vcov_type="hc3") + assert np.all(np.isnan(v)) + del rng + + def test_sensitivity_accepts_beyond_window_single_cohort(self): + from diff_diff.lwdid_sensitivity import robustness_pre_periods + + rng = np.random.default_rng(0) + rows = [] + for u in range(16): + # one real cohort (5); 4 units carry a beyond-window encoding + # (99) that normalizes to never-treated; rest never-treated + g = 5 if u < 6 else (99 if u < 10 else 0) + for t in range(1, 10): + d = int(g == 5 and t >= 5) + rows.append(dict(unit=u, time=t, treat=d, g=g, y=rng.normal() + d)) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = robustness_pre_periods( + df, outcome="y", unit="unit", time="time", treatment="treat", cohort="g" + ) + assert np.isfinite(res.baseline_att) + + def test_zero_post_unit_counted_in_drop_warning(self): + rows = [] + for u in range(12): + treated = u < 6 + t_range = range(1, 4) if u == 11 else range(1, 7) # unit 11: pre rows only + for t in t_range: + d = 1 if (treated and t >= 4) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=float(t) + 2 * d)) + df = pd.DataFrame(rows) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + res = LWDiD(rolling="demean").fit(df, **self.KW) + assert any("fixed-window" in str(x.message) for x in caught) + assert res.n_control == 5 # unit 11 dropped and accounted for + + def test_sensitivity_baseline_config_errors_raise(self): + from diff_diff.lwdid_sensitivity import ( + robustness_pre_periods, + sensitivity_no_anticipation, + ) + + rng = np.random.default_rng(0) + rows = [] + for u in range(12): + for t in range(1, 9): + d = 1 if (u < 6 and t >= 6) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=rng.normal() + d)) + df = pd.DataFrame(rows) + for fn in (robustness_pre_periods, sensitivity_no_anticipation): + with pytest.raises(ValueError, match="requires covariates"): + fn( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + estimation_method="psm", + ) + + +class TestReviewRound11Guards: + """Local-review round 11: execution-verified guards. + + - the IPW/DR logit score/Hessian used CLIPPED propensities, breaking + the estimating-equation linearization whenever trimming fired (the + MLE's score is ~0 in the RAW fitted probabilities only), and clipped + observations kept a nonzero weight-derivative + - NaN logit coefficients from a rank-deficient (collinear) propensity + model were treated as non-convergence and silently substituted + regression adjustment under ipw/dr provenance + - the RA interaction gate counted NOMINAL covariate columns, so a + perfectly collinear control flipped the eq. 3.3 design off and + changed the ATT + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + @staticmethod + def _panel(x_fn, n_units=24, extra=None): + rng = np.random.default_rng(0) + rows = [] + for u in range(n_units): + treated = u < n_units // 2 + x = x_fn(u) + for t in range(1, 7): + d = 1 if (treated and t >= 4) else 0 + row = dict(unit=u, time=t, treat=d, x=x, y=1 + 0.4 * x + 2 * d + rng.normal(0, 0.3)) + if extra is not None: + row["x2"] = extra(x) + rows.append(row) + return pd.DataFrame(rows) + + def test_rank_deficient_propensity_stays_ipw(self): + # x2 = 2x: the logit drops a column (NaN coef) but the fitted + # probabilities are valid - the fit must REMAIN IPW (pre-fix it + # silently became regression adjustment under ipw provenance). + df_full = self._panel(lambda u: float(u % 5)) + df_dup = self._panel(lambda u: float(u % 5), extra=lambda x: 2.0 * x) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + r_dup = LWDiD(rolling="demean", estimation_method="ipw").fit( + df_dup, covariates=["x", "x2"], **self.KW + ) + r_ipw = LWDiD(rolling="demean", estimation_method="ipw").fit( + df_full, covariates=["x"], **self.KW + ) + r_reg = LWDiD(rolling="demean", estimation_method="reg").fit( + df_dup, covariates=["x", "x2"], **self.KW + ) + assert any("reduced-rank propensity" in str(x.message) for x in caught) + # the duplicated-column IPW fit equals the identified IPW fit, + # NOT the regression-adjustment fit + np.testing.assert_allclose(r_dup.att, r_ipw.att, rtol=1e-10) + assert abs(r_dup.att - r_reg.att) > 1e-12 or abs(r_dup.se - r_reg.se) > 1e-12 + + def test_redundant_control_does_not_change_ra_estimand(self): + df_full = self._panel(lambda u: float(u % 3), n_units=8) + df_dup = self._panel(lambda u: float(u % 3), n_units=8, extra=lambda x: 2.0 * x) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + base = LWDiD(rolling="demean", estimation_method="reg").fit( + df_full, covariates=["x"], **self.KW + ) + dup = LWDiD(rolling="demean", estimation_method="reg").fit( + df_dup, covariates=["x", "x2"], **self.KW + ) + # identical identified design -> identical ATT (pre-fix: the + # nominal K flipped the interaction gate and moved the point) + np.testing.assert_allclose(dup.att, base.att, rtol=1e-10) + + def test_trimmed_propensity_score_uses_raw_fit(self): + # Strong-heterogeneity DGP that activates trimming: the logit + # score at the MLE, as constructed by the IF code path, must be + # ~0 (raw probabilities), not the clipped-probability residual. + from diff_diff.linalg import solve_logit + + rng = np.random.default_rng(3) + n = 300 + x = rng.normal(0, 2.5, n) + p = 1 / (1 + np.exp(-2.5 * x)) + d = (rng.random(n) < p).astype(float) + X = x.reshape(-1, 1) + coefs, probs_raw = solve_logit(X, d) + trim_lo, trim_hi = 0.01, 0.99 + assert ((probs_raw < trim_lo) | (probs_raw > trim_hi)).any() # trimming active + X_ps = np.column_stack([np.ones(n), X]) + score_raw = ((d - probs_raw)[:, None] * X_ps).sum(axis=0) + probs_clipped = np.clip(probs_raw, trim_lo, trim_hi) + score_clipped = ((d - probs_clipped)[:, None] * X_ps).sum(axis=0) + assert np.abs(score_raw).max() < 1e-6 # MLE estimating equation + assert np.abs(score_clipped).max() > 1e-2 # the pre-fix construction + + +class TestReviewRound12Guards: + """Local-review round 12: execution-verified guards. + + - the DR outcome WLS inverted the raw nominal Gram (not rank-aware / + scale-equilibrated): an exactly redundant 1e12-rescaled duplicate + changed the DR SE by ~2.5x + - IPW/DR returned NOMINAL parameter counts, so a redundant control + shrank residual df and moved p-values/CIs with ATT/SE unchanged + - the drops-route staggered aggregation weighted cohorts by RAW + masses, keeping dropped treated units in the cohort weights + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + @staticmethod + def _panel(extra_scale=None, n_units=24): + rng = np.random.default_rng(0) + rows = [] + for u in range(n_units): + treated = u < n_units // 2 + x = float(u % 5) + for t in range(1, 7): + d = 1 if (treated and t >= 4) else 0 + row = dict(unit=u, time=t, treat=d, x=x, y=1 + 0.4 * x + 2 * d + rng.normal(0, 0.3)) + if extra_scale is not None: + row["x2"] = extra_scale * x + rows.append(row) + return pd.DataFrame(rows) + + @pytest.mark.parametrize("method", ["ipw", "dr"]) + @pytest.mark.parametrize("scale", [2.0, 1e12]) + def test_redundant_control_full_inference_invariance(self, method, scale): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + base = LWDiD(rolling="demean", estimation_method=method).fit( + self._panel(), covariates=["x"], **self.KW + ) + dup = LWDiD(rolling="demean", estimation_method=method).fit( + self._panel(extra_scale=scale), covariates=["x", "x2"], **self.KW + ) + np.testing.assert_allclose(dup.att, base.att, rtol=1e-8, err_msg=f"{method} att") + np.testing.assert_allclose(dup.se, base.se, rtol=1e-8, err_msg=f"{method} se") + assert dup.df_inference == base.df_inference, f"{method} df" + np.testing.assert_allclose(dup.p_value, base.p_value, rtol=1e-8) + np.testing.assert_allclose(dup.conf_int, base.conf_int, rtol=1e-8) + + def test_drops_route_uses_survivor_cohort_masses(self): + # Independent oracle: cohorts {3: 4 units, 5: 4 units}; ONE + # cohort-5 treated unit observes only t=1..4 (missing its own post + # window entirely) -> dropped. Survivor masses 4/7 and 3/7 must + # weight the cohort effects (raw masses would use 4/8, 4/8). + rng = np.random.default_rng(1) + rows = [] + uid = 0 + spec = [(0, 8, None), (3, 4, None), (5, 3, None), (5, 1, (1, 2, 3, 4))] + for g, n, keep in spec: + for _ in range(n): + alpha = rng.normal() + for t in range(1, 7): + if keep is not None and t not in keep: + continue + d = int(g > 0 and t >= g) + y = alpha + 0.2 * t + rng.normal(scale=0.3) + (1.5 + 0.4 * (g == 5)) * d + rows.append(dict(unit=uid, time=t, treat=d, g=g, y=y)) + uid += 1 + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean", control_group="never_treated", vcov_type="classical").fit( + df, first_treat="g", **self.KW + ) + assert res.n_composite_treated_dropped == 1 + att3 = res.cohort_effects[3]["att"] + att5 = res.cohort_effects[5]["att"] + expected = (4.0 * att3 + 3.0 * att5) / 7.0 # SURVIVOR masses + raw_weighted = (4.0 * att3 + 4.0 * att5) / 8.0 + np.testing.assert_allclose(res.att, expected, rtol=1e-12) + assert abs(res.att - raw_weighted) > 1e-6 # distinguishes the rules + + +class TestReviewRound13Guards: + """Local-review round 13: execution-verified guards. + + - the RA influence bread was rebuilt with a RAW-Gram pinv after + solve_ols's scale-equilibrated fit: at large covariate units the + pinv silently dropped low-scale directions, so cell ATT/SE were + invariant while every AGGREGATE SE/p/CI (and the multiplier- + bootstrap inputs) depended on covariate units + - the exact-inference guard counted nominal columns, rejecting + redundant-column designs with positive effective residual df + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_aggregate_inference_invariant_to_covariate_units(self): + rng = np.random.default_rng(2) + + def build(scale): + rows = [] + for u in range(20): + g = 4 if u < 5 else (5 if u < 10 else 0) + x = float(u % 4) * scale + for t in range(1, 8): + d = int(g > 0 and t >= g) + rows.append( + dict( + unit=u, + time=t, + treat=d, + g=g, + x=x, + y=1 + 0.2 * t + 0.3 * (x / scale) + 1.5 * d + rng.normal(0, 0.3), + ) + ) + return pd.DataFrame(rows) + + df1 = build(1.0) + rng = np.random.default_rng(2) # same noise stream + df2 = build(10.0**7.25) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r1 = LWDiD(rolling="demean").fit(df1, first_treat="g", covariates=["x"], **self.KW) + r2 = LWDiD(rolling="demean").fit(df2, first_treat="g", covariates=["x"], **self.KW) + np.testing.assert_allclose(r2.att, r1.att, rtol=1e-8) + np.testing.assert_allclose(r2.se, r1.se, rtol=1e-6) # aggregate IF SE + np.testing.assert_allclose(r2.p_value, r1.p_value, rtol=1e-5, atol=1e-300) + for k in r1.event_study_effects: + np.testing.assert_allclose( + r2.event_study_effects[k]["se"], + r1.event_study_effects[k]["se"], + rtol=1e-6, + err_msg=f"event {k}", + ) + + def test_redundant_column_small_sample_fits(self): + # 4 collapsed units, design [1, D, x, 2x]: effective rank 3, + # residual df 1 -> must FIT (pre-fix: nominal width 4 raised). + rows = [] + for u in range(4): + x = float(u) + for t in range(1, 5): + d = 1 if (u < 2 and t >= 3) else 0 + rows.append( + dict( + unit=u, + time=t, + treat=d, + x=x, + x2=2.0 * x, + y=1 + 0.5 * t + 0.3 * x + 2 * d + 0.01 * u * t, + ) + ) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean").fit(df, covariates=["x", "x2"], **self.KW) + assert np.isfinite(res.att) + # a genuinely saturated full-rank design still raises + rows2 = [] + for u in range(3): + for t in range(1, 5): + d = 1 if (u < 1 and t >= 3) else 0 + rows2.append( + dict(unit=u, time=t, treat=d, x=float(u**2), y=1 + 0.5 * t + 2 * d + 0.01 * u) + ) + with pytest.raises(ValueError, match="Invalid exact-inference design"): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + LWDiD(rolling="demean").fit(pd.DataFrame(rows2), covariates=["x"], **self.KW) + + +class TestReviewRound14Guards: + """Local-review round 14: Inf time values raised a raw OverflowError + in event-time arithmetic instead of a validation error.""" + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + @pytest.mark.parametrize("bad", [np.inf, -np.inf]) + def test_nonfinite_time_rejected(self, bad): + rows = [] + for u in range(8): + for t in range(1, 7): + d = 1 if (u < 4 and t >= 4) else 0 + rows.append(dict(unit=u, time=float(t), treat=d, y=1.0 + d)) + df = pd.DataFrame(rows) + df.loc[df.index[2], "time"] = bad + with pytest.raises(ValueError, match="Time column .* non-finite"): + LWDiD(rolling="demean").fit(df, **self.KW) + df["g"] = np.where(df["unit"] < 4, 4.0, 0.0) + with pytest.raises(ValueError, match="Time column .* non-finite"): + LWDiD(rolling="demean").fit(df, first_treat="g", **self.KW) + + +class TestReviewRound16Guards: + """Local-review round 16: execution-verified guards. + + - a degenerate staggered event row (NaN inference) still contributed + a 0.0-diagonal column to the ANALYTICAL event-study covariance + - k_min=1 was silently clamped to 2, dropping a valid demeaning spec + - the degenerate early-return results lost psm_config / cluster_name + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_degenerate_event_row_excluded_from_analytical_vcov(self): + from types import SimpleNamespace + + from diff_diff.lwdid_staggered import compute_event_study_bands + + rng = np.random.default_rng(0) + estimator = SimpleNamespace(n_bootstrap=0, seed=None, alpha=0.05) + event_effects = { + 0: { + "effect": 1.0, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "df": None, + }, + 1: { + "effect": 0.5, + "se": 0.1, + "t_stat": 5.0, + "p_value": 0.0, + "conf_int": (0.3, 0.7), + "df": None, + }, + } + event_influence = {0: np.zeros(30), 1: rng.normal(size=30)} + vcov, index, *_ = compute_event_study_bands(estimator, event_effects, event_influence, None) + assert list(index) == [1] # NaN-inference row excluded + assert vcov.shape == (1, 1) and np.isfinite(vcov[0, 0]) + + def test_k_min_one_honored_for_demean(self): + from diff_diff.lwdid_sensitivity import robustness_pre_periods + + rng = np.random.default_rng(0) + rows = [] + for u in range(12): + for t in range(1, 9): + d = 1 if (u < 6 and t >= 6) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=rng.normal() + d)) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = robustness_pre_periods( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + rolling="demean", + k_min=1, + ) + labels = [s.label for s in res.specifications] + assert "k=1_pre_periods" in labels + with pytest.raises(ValueError, match="minimum pre-period requirement"): + robustness_pre_periods( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + rolling="detrend", + k_min=1, + ) + + def test_degenerate_return_keeps_psm_and_cluster_provenance(self): + # detrend with a single pre-period: every transformed outcome is + # NaN -> the degenerate early return must still carry the fit + # configuration. + rows = [] + for u in range(8): + for t in range(2, 7): # one pre-period (t=2), onset t=3 + d = 1 if (u < 4 and t >= 3) else 0 + rows.append(dict(unit=u, time=t, treat=d, cl=u % 3, x=float(u % 4), y=1.0 + d)) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res_psm = LWDiD(rolling="detrend", estimation_method="psm", caliper=0.5).fit( + df, covariates=["x"], **self.KW + ) + res_cl = LWDiD(rolling="detrend", cluster="cl").fit(df, **self.KW) + assert np.isnan(res_psm.att) + assert res_psm.psm_config is not None + assert res_psm.psm_config["caliper"] == 0.5 + assert np.isnan(res_cl.att) + assert res_cl.cluster_name == "cl" + + +class TestReviewRound17Guards: + """Local-review round 17: the RA gates used matrix_rank's looser + default tolerance, disagreeing with solve_ols's pivoted-QR 1e-7 + convention on NEAR-collinear controls (x2 = x + 1e-10): the gate + could count two identified controls, turn the interacted design off, + and move the ATT while the solver fit the identified single-control + model.""" + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + @staticmethod + def _panel(near=False, n_units=8): + rng = np.random.default_rng(4) + rows = [] + for u in range(n_units): + treated = u < n_units // 2 + x = float(u % 4) + for t in range(1, 7): + d = 1 if (treated and t >= 4) else 0 + row = dict(unit=u, time=t, treat=d, x=x, y=1 + 0.4 * x + 2 * d + rng.normal(0, 0.3)) + if near: + row["x2"] = x + 1e-10 + rows.append(row) + return pd.DataFrame(rows) + + def test_near_collinear_control_matches_identified_design(self): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + base = LWDiD(rolling="demean", estimation_method="reg").fit( + self._panel(), covariates=["x"], **self.KW + ) + near = LWDiD(rolling="demean", estimation_method="reg").fit( + self._panel(near=True), covariates=["x", "x2"], **self.KW + ) + # same identified design under the SHARED rank convention: the + # near-duplicate is dropped, the interaction gate stays on, and + # the ATT matches the single-control fit + np.testing.assert_allclose(near.att, base.att, rtol=1e-6) + # replay coherence (the mirror uses the same shared detector) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ri = near.randomization_test(n_reps=49, seed=1) + np.testing.assert_allclose(ri.att_observed, near.att, rtol=1e-10) + + +class TestReviewRound18Guards: + """Local-review round 18: diagnostics bypassed the common-timing + time-scale checks; RI applied the finite mask before shape checks.""" + + def test_diagnostics_reject_period_detrend_and_string_time(self): + est = LWDiD(rolling="detrend") + periods = pd.period_range("2020Q1", periods=6, freq="Q") + rows = [] + for u in range(6): + for i, p in enumerate(periods): + d = 1 if (u < 3 and i >= 4) else 0 + rows.append(dict(unit=u, time=p, treat=d, y=float(i) + d)) + df = pd.DataFrame(rows) + df["time"] = pd.PeriodIndex(df["time"], freq="Q") + with pytest.raises(ValueError, match="Period"): + est.get_transformation_diagnostics( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + df2 = df.assign(time=[f"P{i%6}" for i in range(len(df))]) + with pytest.raises(ValueError, match="numeric or"): + est.get_transformation_diagnostics( + df2, outcome="y", unit="unit", time="time", treatment="treat" + ) + + def test_ri_shape_checks_precede_finite_mask(self): + from diff_diff.lwdid_randomization import randomization_inference + + y = np.array([1.0, np.inf, 2.0, 3.0]) + with pytest.raises(ValueError, match="same length"): + randomization_inference(y, np.array([1.0, 0.0]), n_reps=9) + with pytest.raises(ValueError, match="controls must have"): + randomization_inference( + y, + np.array([1.0, 0.0, 1.0, 0.0]), + controls=np.zeros((2, 1)), + n_reps=9, + ) + + +class TestReviewRound19Guards: + """Local-review round 19: PSM treated a rank-deficient (finite- + probability) propensity fit as non-convergence and substituted a + regression-adjustment point under psm provenance (ipw/dr already + continued reduced-rank).""" + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + @staticmethod + def _panel(extra=False): + rng = np.random.default_rng(0) + rows = [] + for u in range(24): + treated = u < 12 + x = float(u % 5) + for t in range(1, 7): + d = 1 if (treated and t >= 4) else 0 + row = dict(unit=u, time=t, treat=d, x=x, y=1 + 0.4 * x + 2 * d + rng.normal(0, 0.3)) + if extra: + row["x2"] = 2.0 * x + rows.append(row) + return pd.DataFrame(rows) + + def test_psm_continues_reduced_rank(self): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + base = LWDiD(rolling="demean", estimation_method="psm").fit( + self._panel(), covariates=["x"], **self.KW + ) + dup = LWDiD(rolling="demean", estimation_method="psm").fit( + self._panel(extra=True), covariates=["x", "x2"], **self.KW + ) + assert any("reduced-rank" in str(x.message) for x in caught) + # identical propensity fit -> identical matches -> identical ATT + np.testing.assert_allclose(dup.att, base.att, rtol=1e-10) + from tests.conftest import assert_nan_inference + + assert_nan_inference( + {"se": dup.se, "t_stat": dup.t_stat, "p_value": dup.p_value, "conf_int": dup.conf_int} + ) + + +class TestReviewRound20Guards: + """Local-review round 20: PSM was rejected by the exact-OLS df guard; + staggered diagnostics returned an empty success on all-never-treated + panels; ordered chronology enforced for non-numeric common time.""" + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_psm_point_only_exempt_from_residual_df_guard(self): + # 2 treated + 2 controls with 3 redundant unit-constant + # covariates: nominal width exhausts an OLS df count PSM never + # uses - the point-only matching fit must still run. + rng = np.random.default_rng(0) + rows = [] + for u in range(4): + x = float(u) + for t in range(1, 7): + d = 1 if (u < 2 and t >= 4) else 0 + rows.append( + dict( + unit=u, + time=t, + treat=d, + x=x, + x2=2 * x, + x3=3 * x, + y=1 + 0.3 * x + 2 * d + rng.normal(0, 0.2), + ) + ) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean", estimation_method="psm").fit( + df, covariates=["x", "x2", "x3"], **self.KW + ) + assert np.isfinite(res.att) + assert res.df_inference is None + from tests.conftest import assert_nan_inference + + assert_nan_inference( + {"se": res.se, "t_stat": res.t_stat, "p_value": res.p_value, "conf_int": res.conf_int} + ) + + def test_diagnostics_reject_all_never_treated(self): + rows = [] + for u in range(6): + for t in range(1, 7): + rows.append(dict(unit=u, time=t, treat=0, g=0, y=float(t))) + df = pd.DataFrame(rows) + with pytest.raises(ValueError, match="No treated cohorts"): + LWDiD(rolling="demean").get_transformation_diagnostics( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="g", + ) + + def test_ordered_categorical_nonlexicographic_chronology(self): + # 'Q10' sorts before 'Q2' lexicographically; the declared order + # must win (zero-effect trend panel -> att 0 under the correct + # chronology). + labels = [f"Q{i}" for i in range(1, 11)] # Q1..Q10 + rng = np.random.default_rng(0) + rows = [] + for u in range(10): + for i, lab in enumerate(labels): + d = 1 if (u < 5 and i >= 7) else 0 + rows.append(dict(unit=u, time=lab, treat=d, y=float(i))) + df = pd.DataFrame(rows) + df["time"] = pd.Categorical(df["time"], categories=labels, ordered=True) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean").fit(df, **self.KW) + np.testing.assert_allclose(res.att, 0.0, atol=1e-12) + del rng + + +class TestReviewRound21Guards: + """Local-review round 21: hc2 fabricated finite inference at leverage + one on the new LWDiD surface; the common-timing headline bypassed the + degenerate-SE guard; plots rendered +/-1.96*SE instead of the fitted + intervals.""" + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_hc2_leverage_one_fails_closed_on_lwdid(self): + rng = np.random.default_rng(0) + rows = [] + for u in range(12): + for t in range(1, 7): + d = 1 if (u < 1 and t >= 4) else 0 # single treated unit + rows.append(dict(unit=u, time=t, treat=d, y=rng.normal() + d)) + df = pd.DataFrame(rows) + with pytest.warns(UserWarning, match="HC2 variance is undefined"): + res = LWDiD(rolling="demean", vcov_type="hc2").fit(df, **self.KW) + assert np.isfinite(res.att) + from tests.conftest import assert_nan_inference + + assert_nan_inference( + {"se": res.se, "t_stat": res.t_stat, "p_value": res.p_value, "conf_int": res.conf_int} + ) + + @pytest.mark.parametrize("vcov", ["classical", "hc1"]) + def test_exact_fit_headline_fails_closed(self, vcov): + # y = t exactly: the collapsed regression fits exactly, so the SE + # is roundoff of zero - pre-fix t ~ 1e16 was reported. + rows = [] + for u in range(6): + for t in range(1, 7): + d = 1 if (u < 3 and t >= 4) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=float(t))) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean", vcov_type=vcov).fit(df, **self.KW) + np.testing.assert_allclose(res.att, 0.0, atol=1e-12) + from tests.conftest import assert_nan_inference + + assert_nan_inference( + {"se": res.se, "t_stat": res.t_stat, "p_value": res.p_value, "conf_int": res.conf_int} + ) + + def test_event_plot_uses_fitted_interval_endpoints(self): + pytest.importorskip("matplotlib") + import matplotlib + + matplotlib.use("Agg") + from diff_diff.lwdid_visualization import plot_event_study + + rng = np.random.default_rng(0) + rows = [] + for u in range(14): + g = 4 if u < 7 else 0 + for t in range(1, 8): + d = int(g > 0 and t >= g) + rows.append( + dict(unit=u, time=t, treat=d, g=g, y=1 + 0.2 * t + d + rng.normal(0, 0.4)) + ) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean", alpha=0.10).fit(df, first_treat="g", **self.KW) + fig = plot_event_study(res) + ax = fig.axes[0] + # the rendered whiskers must match the FITTED interval endpoints + # (alpha=0.10 t-intervals), not +/-1.96*SE + seg_ys = sorted( + y + for coll in ax.collections + for seg in coll.get_segments() + for y in (seg[0][1], seg[-1][1]) + ) + row = res.event_study_effects[max(res.event_study_effects)] + lo, hi = row["conf_int"] + assert any(abs(y - lo) < 1e-9 for y in seg_ys) + assert any(abs(y - hi) < 1e-9 for y in seg_ys) + naive = 1.96 * row["se"] + fitted_half = row["effect"] - lo + assert abs(naive - fitted_half) > 1e-6 # the two conventions differ here + import matplotlib.pyplot as plt + + plt.close(fig) + + +class TestReviewRound22Guards: + """Local-review round 22: array-valued alpha passed construction and + failed later with a raw TypeError.""" + + @pytest.mark.parametrize( + "bad", [np.array([0.05]), "0.05", None, complex(0.05), np.nan, np.inf, True] + ) + def test_alpha_scalar_validation(self, bad): + with pytest.raises((ValueError, TypeError)): + LWDiD(alpha=bad) + from diff_diff.lwdid_wild_bootstrap import wild_cluster_bootstrap + + y = np.random.default_rng(0).normal(size=20) + d = np.array([1.0] * 10 + [0.0] * 10) + cl = np.arange(20) % 5 + with pytest.raises((ValueError, TypeError)): + wild_cluster_bootstrap(y, d, cl, alpha=bad, n_bootstrap=19) + + +class TestReviewRound23Guards: + """Local-review round 23: raw cohort masses weighted non-contributing + treated units into staggered overall aggregates outside the tau_omega + route; plot_cohort_trends silently ignored cohort=; validator + reported missing unit/time as warnings only.""" + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_overall_weights_use_contributing_treated_units(self): + # Two cohorts under NOT_YET_TREATED (non-tau_omega route): 3 of 4 + # cohort-5 treated units observe NO post rows -> only 1 + # contributes. Overall masses must be 4 (cohort 3) and 1 + # (cohort 5), not the raw 4 and 4. + rng = np.random.default_rng(2) + rows = [] + uid = 0 + spec = [(0, 8, None), (3, 4, None), (5, 1, None), (5, 3, (1, 2, 3, 4))] + for g, n, keep in spec: + for _ in range(n): + alpha = rng.normal() + for t in range(1, 7): + if keep is not None and t not in keep: + continue + d = int(g > 0 and t >= g) + rows.append( + dict( + unit=uid, + time=t, + treat=d, + g=g, + y=alpha + 0.2 * t + 1.5 * d + rng.normal(0, 0.3), + ) + ) + uid += 1 + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean", control_group="not_yet_treated").fit( + df, first_treat="g", **self.KW + ) + assert res.cohort_effects[3]["n_treated"] == 4 + assert res.cohort_effects[5]["n_treated"] == 1 # contributing only + att3 = res.cohort_effects[3]["att"] + att5 = res.cohort_effects[5]["att"] + expected = (4.0 * att3 + 1.0 * att5) / 5.0 + raw = (4.0 * att3 + 4.0 * att5) / 8.0 + np.testing.assert_allclose(res.att, expected, rtol=1e-12) + assert abs(res.att - raw) > 1e-9 + + def test_plot_cohort_trends_honors_cohort(self): + pytest.importorskip("matplotlib") + import matplotlib + + matplotlib.use("Agg") + from diff_diff.lwdid_visualization import plot_cohort_trends + + rng = np.random.default_rng(0) + rows = [] + for u in range(12): + g = 3 if u < 4 else (5 if u < 8 else 0) + for t in range(1, 7): + d = int(g > 0 and t >= g) + rows.append(dict(unit=u, time=t, treat=d, g=g, y=rng.normal() + d)) + df = pd.DataFrame(rows) + fig = plot_cohort_trends( + df, outcome="y", unit="unit", time="time", treatment="treat", cohort="g" + ) + labels = [line.get_label() for ax in fig.axes for line in ax.get_lines()] + assert any("Cohort 3" in lab for lab in labels) + assert any("Cohort 5" in lab for lab in labels) + assert any(lab == "Control" for lab in labels) + import matplotlib.pyplot as plt + + plt.close(fig) + + def test_validator_missing_unit_time_invalid(self): + from diff_diff.lwdid import validate_staggered_data + + rows = [] + for u in range(6): + g = 4 if u < 3 else 0 + for t in range(1, 7): + rows.append(dict(unit=u, time=t, g=g, y=1.0)) + df = pd.DataFrame(rows) + df.loc[df.index[3], "unit"] = np.nan + out = validate_staggered_data(df, unit="unit", time="time", cohort="g") + assert out["valid"] is False + assert any("missing values" in e for e in out["errors"]) + + +class TestReviewRound24Guards: + """Local-review round 24 P2s: unusable small n_reps rejected up + front; HC3 fail-closed keeps the length-k DOF contract; datetime + cohort relabeling is canonical (collision/row-order independent).""" + + def test_ri_small_n_reps_rejected_up_front(self): + from diff_diff.lwdid_randomization import randomization_inference + + y = np.random.default_rng(0).normal(size=20) + d = np.array([1.0] * 10 + [0.0] * 10) + with pytest.raises(ValueError, match="integer >= 10"): + randomization_inference(y, d, n_reps=9) + res = randomization_inference(y, d, n_reps=10, seed=0) + assert 0 < res.pvalue <= 1 + + def test_hc3_fail_closed_dof_vector(self): + from diff_diff.linalg import compute_robust_vcov + + X = np.column_stack([np.ones(4), np.array([0.0, 0.0, 0.0, 1.0])]) + y = np.array([1.0, 1.1, 0.9, 5.0]) + resid = y - X @ np.linalg.lstsq(X, y, rcond=None)[0] + with pytest.warns(UserWarning, match="HC3 variance is undefined"): + vcov, dof = compute_robust_vcov(X, resid, vcov_type="hc3", return_dof=True) + assert np.all(np.isnan(vcov)) + assert dof.shape == (2,) and np.all(np.isnan(dof)) + + def test_datetime_cohort_relabel_canonical(self): + # Two raw between-period cohort dates map to the SAME observed + # onset; the reported cohort key must be the canonical observed + # period regardless of row order. + rng = np.random.default_rng(0) + times = pd.date_range("2020-01-01", periods=6, freq="QS") + onset = times[4] + raw_a = onset - pd.Timedelta(days=10) + raw_b = onset - pd.Timedelta(days=20) + + def build(order): + rows = [] + for idx, u in enumerate(order): + g = {0: raw_a, 1: raw_b}.get(u, pd.NaT) if u < 2 else pd.NaT + for i, ts in enumerate(times): + d = int(u < 2 and ts >= onset) + rows.append(dict(unit=u, time=ts, treat=d, g=g, y=rng.normal() + d)) + return pd.DataFrame(rows) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r1 = LWDiD(rolling="demean").fit( + build([0, 1, 2, 3, 4, 5]), + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="g", + ) + r2 = LWDiD(rolling="demean").fit( + build([1, 0, 2, 3, 4, 5]), + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="g", + ) + assert list(r1.cohort_effects) == [onset] + assert list(r2.cohort_effects) == [onset] + + +class TestReviewRoundCI7Guards: + """CI review round 7: pin the DOCUMENTED within-cohort cell-mass + convention on an unbalanced panel where it provably differs from the + LW 2026 eq. 7.10 unit-average estimand (a treated unit observing + more post periods carries more cell-mass weight).""" + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat") + + def test_cohort_effect_cell_mass_oracle_unbalanced(self): + # One cohort (g=4), 2 treated units: unit 0 observes post {4,5,6}, + # unit 1 observes post {4} only. 6 never-treated controls observe + # everything. Deterministic outcomes. + rows = [] + for u in range(8): + g = 4 if u < 2 else 0 + times = range(1, 7) + for t in times: + if u == 1 and t > 4: + continue # unit 1 misses post periods 5, 6 + d = int(g > 0 and t >= g) + # unit-specific level + zero noise; treated add u-dependent effect + y = 10.0 * u + 0.0 * t + (2.0 + 3.0 * u) * d + rows.append(dict(unit=u, time=t, treat=d, g=g, y=y)) + df = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean", control_group="never_treated", vcov_type="hc1").fit( + df, first_treat="g", **self.KW + ) + # Independent oracle. Demean: pre-mean = level (zero trend), so + # ydot = effect for treated rows, 0 for controls. + # Cell ATTs: t=4 has units {0,1} -> mean(2, 5) = 3.5, n_treated=2; + # t=5, t=6 have unit 0 only -> 2.0, n_treated=1. + # CELL-MASS cohort effect = (2*3.5 + 1*2 + 1*2) / 4 = 2.75. + # eq. 7.10 UNIT-AVERAGE estimand: unit post-averages are 2.0 + # (unit 0) and 5.0 (unit 1) -> cohort effect 3.5. The documented + # convention is cell-mass. + cell_mass = res.cohort_effects[4]["att"] + np.testing.assert_allclose(cell_mass, 2.75, atol=1e-10) + assert abs(cell_mass - 3.5) > 0.5 # distinguishes eq. 7.10 + # .att on this NT/reg route is the tau_omega COMPOSITE (7.18), + # built from unit post-averages - here exactly the eq. 7.10 + # unit-average value (3.5). The two surfaces answer different, + # separately documented estimands on unbalanced panels. + np.testing.assert_allclose(res.att, 3.5, atol=1e-10) diff --git a/tests/test_lwdid_diagnostics.py b/tests/test_lwdid_diagnostics.py new file mode 100644 index 00000000..3def486f --- /dev/null +++ b/tests/test_lwdid_diagnostics.py @@ -0,0 +1,504 @@ +"""Tests for LWDiD diagnostics output and mathematical correctness. + +Verifies: +1. _dispatch_estimator routing and return structure +2. Transformation diagnostics (get_transformation_diagnostics) +3. Mathematical correctness against Lee & Wooldridge (2025, 2026) formulas +4. Backward compatibility (existing fit() behavior unchanged) +""" + +import numpy as np +import pandas as pd +import pytest + +from diff_diff import LWDiD + +# ============================================================ +# Fixtures +# ============================================================ + + +@pytest.fixture +def simple_panel(): + """Simple balanced panel: 40 units, 8 periods, treatment at t=5.""" + rng = np.random.default_rng(42) + records = [] + for i in range(40): + d = int(i < 15) + for t in range(1, 9): + y = 1.0 + 0.3 * i / 40 + 0.1 * t + rng.normal(0, 0.3) + post = int(t > 4) + if d and post: + y += 2.0 + records.append({"unit": i, "time": t, "y": y, "treat": d * post}) + return pd.DataFrame(records) + + +@pytest.fixture +def panel_with_controls(): + """Panel with covariate X.""" + rng = np.random.default_rng(123) + records = [] + for i in range(60): + d = int(i < 20) + x1 = rng.normal() + d * 0.3 + for t in range(1, 9): + y = 1.0 + 0.5 * x1 + 0.1 * t + rng.normal(0, 0.3) + post = int(t > 4) + if d and post: + y += 2.0 + records.append({"unit": i, "time": t, "y": y, "treat": d * post, "x1": x1}) + return pd.DataFrame(records) + + +@pytest.fixture +def quarterly_panel(): + """Panel with 16 periods (4 years of quarterly data).""" + rng = np.random.default_rng(99) + records = [] + for i in range(50): + d = int(i < 18) + for t in range(1, 17): + q = (t - 1) % 4 + 1 + seasonal = 0.5 * (q == 4) - 0.3 * (q == 1) + y = 2.0 + 0.05 * t + seasonal + rng.normal(0, 0.2) + post = int(t > 8) + if d and post: + y += 1.5 + records.append({"unit": i, "time": t, "y": y, "treat": d * post}) + return pd.DataFrame(records) + + +# ============================================================ +# Class 1: _dispatch_estimator behavior verification +# ============================================================ + + +class TestDispatchEstimator: + """Verify _dispatch_estimator routing and return structure.""" + + def test_ra_returns_valid_result(self, simple_panel): + """RA path returns valid ATT estimate.""" + est = LWDiD(rolling="demean", estimation_method="reg") + res = est.fit(simple_panel, outcome="y", unit="unit", time="time", treatment="treat") + # Verify ATT is finite and reasonable + assert np.isfinite(res.att) + assert 1.0 < res.att < 3.0 # true ATT = 2.0 + + def test_ipw_returns_valid_result(self, panel_with_controls): + """IPW path returns valid results with controls.""" + est = LWDiD(rolling="demean", estimation_method="ipw") + res = est.fit( + panel_with_controls, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1"], + ) + assert np.isfinite(res.att) + assert np.isfinite(res.se) + + def test_dr_returns_valid_result(self, panel_with_controls): + """DR path returns valid doubly-robust results.""" + est = LWDiD(rolling="demean", estimation_method="dr") + res = est.fit( + panel_with_controls, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1"], + ) + assert np.isfinite(res.att) + + def test_psm_returns_valid_result(self, panel_with_controls): + """PSM path returns valid matched results.""" + est = LWDiD(rolling="demean", estimation_method="psm") + res = est.fit( + panel_with_controls, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1"], + ) + assert np.isfinite(res.att) + + def test_all_estimators_same_data_give_reasonable_att(self, panel_with_controls): + """All 4 estimators should give ATT in [1.0, 3.0] for true ATT=2.0.""" + for est_name in ["reg", "ipw", "dr", "psm"]: + est = LWDiD(rolling="demean", estimation_method=est_name) + res = est.fit( + panel_with_controls, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1"], + ) + assert 1.0 < res.att < 3.0, f"{est_name} ATT={res.att} outside [1,3]" + + def test_ipw_without_controls_still_works(self, simple_panel): + """IPW without controls still produces a result.""" + import warnings + + est = LWDiD(estimation_method="ipw") + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + res = est.fit(simple_panel, outcome="y", unit="unit", time="time", treatment="treat") + assert np.isfinite(res.att) + + +# ============================================================ +# Class 2: Transformation diagnostics +# ============================================================ + + +class TestTransformationDiagnostics: + """Verify get_transformation_diagnostics() output structure and values.""" + + def test_demean_diagnostics_structure(self, simple_panel): + """Demean diagnostics has correct structure.""" + est = LWDiD(rolling="demean") + diag = est.get_transformation_diagnostics( + simple_panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert diag["method"] == "demean" + assert "per_unit" in diag + assert "summary" in diag + assert len(diag["per_unit"]) == 40 # 40 units + # Check per-unit fields + first_unit = list(diag["per_unit"].values())[0] + assert "pre_mean" in first_unit + assert "pre_n_periods" in first_unit + assert "pre_std" in first_unit + assert "valid" in first_unit + # Check summary fields + assert "n_units_total" in diag["summary"] + assert "n_units_valid" in diag["summary"] + + def test_detrend_diagnostics_structure(self, simple_panel): + """Detrend diagnostics has correct structure with alpha/beta.""" + est = LWDiD(rolling="detrend") + diag = est.get_transformation_diagnostics( + simple_panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert diag["method"] == "detrend" + first_unit = list(diag["per_unit"].values())[0] + assert "alpha" in first_unit + assert "beta" in first_unit + assert "r_squared" in first_unit + + def test_demeanq_diagnostics_structure(self, quarterly_panel): + """Demeanq diagnostics has seasonal effects.""" + est = LWDiD(rolling="demeanq") + diag = est.get_transformation_diagnostics( + quarterly_panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert diag["method"] == "demeanq" + first_unit = list(diag["per_unit"].values())[0] + assert "intercept" in first_unit + assert "seasonal_effects" in first_unit + + def test_detrendq_diagnostics_structure(self, quarterly_panel): + """Detrendq diagnostics has trend + seasonal.""" + est = LWDiD(rolling="detrendq") + diag = est.get_transformation_diagnostics( + quarterly_panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert diag["method"] == "detrendq" + first_unit = list(diag["per_unit"].values())[0] + assert "alpha" in first_unit + assert "beta" in first_unit + assert "seasonal_effects" in first_unit + + def test_diagnostics_does_not_affect_estimation(self, simple_panel): + """get_transformation_diagnostics does not change fit() results.""" + est = LWDiD(rolling="detrend") + # Get diagnostics first + est.get_transformation_diagnostics( + simple_panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + # Then fit + res = est.fit(simple_panel, outcome="y", unit="unit", time="time", treatment="treat") + # Should still be correct + assert np.isfinite(res.att) + assert 1.0 < res.att < 3.0 + + +# ============================================================ +# Class 2b: Per-cohort transformation diagnostics (staggered) +# ============================================================ + + +def _make_staggered_diag_panel(): + """Deterministic staggered panel: y = 10*unit + t, 6 units, 6 periods. + + Units 0-1: cohort g=3; units 2-3: cohort g=5; units 4-5: never (g=0). + """ + records = [] + cohorts = {0: 3, 1: 3, 2: 5, 3: 5, 4: 0, 5: 0} + for i, g in cohorts.items(): + for t in range(1, 7): + records.append( + { + "unit": i, + "time": t, + "y": 10.0 * i + t, + "treat": int(g > 0 and t >= g), + "cohort": g, + } + ) + return pd.DataFrame(records) + + +class TestStaggeredPerCohortDiagnostics: + """Staggered diagnostics use each cohort's own pre-period t < g.""" + + def test_by_cohort_structure(self): + """Top-level dict is organized by cohort keys g.""" + df = _make_staggered_diag_panel() + est = LWDiD(rolling="demean") + diag = est.get_transformation_diagnostics( + df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="cohort" + ) + assert diag["method"] == "demean" + assert diag["design"] == "staggered" + assert set(diag["by_cohort"].keys()) == {3, 5} + # Each per-cohort entry keeps the _transform_* diagnostics contract + for g in (3, 5): + assert diag["by_cohort"][g]["method"] == "demean" + assert "per_unit" in diag["by_cohort"][g] + assert "summary" in diag["by_cohort"][g] + + def test_per_cohort_pre_means_hand_computed(self): + """Ȳ_{i,pre} uses t < g per cohort: mean over its own pre-window. + + y_it = 10*i + t, so for cohort g=3 (pre t=1,2): Ȳ = 10*i + 1.5; + for cohort g=5 (pre t=1..4): Ȳ = 10*i + 2.5. + """ + df = _make_staggered_diag_panel() + est = LWDiD(rolling="demean") + diag = est.get_transformation_diagnostics( + df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="cohort" + ) + g3 = diag["by_cohort"][3]["per_unit"] + g5 = diag["by_cohort"][5]["per_unit"] + # Cohort 3 treated units: 2 pre-periods (t=1,2) + np.testing.assert_allclose(g3[0]["pre_mean"], 1.5, atol=1e-10) + np.testing.assert_allclose(g3[1]["pre_mean"], 11.5, atol=1e-10) + assert g3[0]["pre_n_periods"] == 2 + # Cohort 5 treated units: 4 pre-periods (t=1..4) + np.testing.assert_allclose(g5[2]["pre_mean"], 22.5, atol=1e-10) + np.testing.assert_allclose(g5[3]["pre_mean"], 32.5, atol=1e-10) + assert g5[2]["pre_n_periods"] == 4 + # Same never-treated unit gets a different pre-window per cohort + np.testing.assert_allclose(g3[4]["pre_mean"], 41.5, atol=1e-10) + np.testing.assert_allclose(g5[4]["pre_mean"], 42.5, atol=1e-10) + + def test_control_group_determines_unit_subset(self): + """Diagnostics mirror the estimation unit subset per cohort.""" + df = _make_staggered_diag_panel() + # not_yet_treated: cohort 3's frame includes later cohort 5 units + diag_nyt = LWDiD(control_group="not_yet_treated").get_transformation_diagnostics( + df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="cohort" + ) + assert set(diag_nyt["by_cohort"][3]["per_unit"].keys()) == {0, 1, 2, 3, 4, 5} + # never_treated: cohort 3's frame excludes cohort 5 units + diag_nt = LWDiD(control_group="never_treated").get_transformation_diagnostics( + df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="cohort" + ) + assert set(diag_nt["by_cohort"][3]["per_unit"].keys()) == {0, 1, 4, 5} + assert set(diag_nt["by_cohort"][5]["per_unit"].keys()) == {2, 3, 4, 5} + + def test_detrend_per_cohort_slope_hand_computed(self): + """β̂_i from pre-period OLS is 1.0 for y = 10*i + t in every cohort.""" + df = _make_staggered_diag_panel() + est = LWDiD(rolling="detrend") + diag = est.get_transformation_diagnostics( + df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="cohort" + ) + for g in (3, 5): + for info in diag["by_cohort"][g]["per_unit"].values(): + np.testing.assert_allclose(info["beta"], 1.0, atol=1e-10) + np.testing.assert_allclose(info["r_squared"], 1.0, atol=1e-10) + + +# ============================================================ +# Class 3: Mathematical correctness (Lee & Wooldridge formulas) +# ============================================================ + + +class TestMathematicalCorrectness: + """Verify mathematical formulas against hand-computed values. + + Reference: Lee & Wooldridge (2025), Procedures 2.1 and 3.1. + """ + + def test_demean_formula_hand_computed(self): + """Verify Ȳ_{i,pre} = (1/(S-1)) * Σ_{t=1}^{S-1} Y_{it}. + + Per Procedure 2.1: pre-treatment mean subtracted from all periods. + """ + # Construct tiny known dataset: 3 units, 4 periods, treatment at t=3 + # All units are treated so pre_mask = (treat == 0) → t=1,2 for all + df = pd.DataFrame( + { + "unit": [0] * 4 + [1] * 4 + [2] * 4, + "time": [1, 2, 3, 4] * 3, + "y": [ + 2.0, + 4.0, + 10.0, + 12.0, # unit 0: pre_mean = (2+4)/2 = 3.0 + 1.0, + 3.0, + 8.0, + 10.0, # unit 1: pre_mean = (1+3)/2 = 2.0 + 3.0, + 5.0, + 6.0, + 7.0, # unit 2: pre_mean = (3+5)/2 = 4.0 + ], + "treat": [0, 0, 1, 1] * 3, # all units treated at t=3 + } + ) + est = LWDiD(rolling="demean") + diag = est.get_transformation_diagnostics( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + # Verify pre-treatment means + np.testing.assert_allclose(diag["per_unit"][0]["pre_mean"], 3.0, atol=1e-10) + np.testing.assert_allclose(diag["per_unit"][1]["pre_mean"], 2.0, atol=1e-10) + np.testing.assert_allclose(diag["per_unit"][2]["pre_mean"], 4.0, atol=1e-10) + + def test_detrend_formula_hand_computed(self): + """Verify α̂_i, β̂_i from pre-treatment OLS: Y_{it} = α + β*t + ε. + + Per Procedure 3.1: unit-specific linear trend removed. + """ + # Unit with perfect linear trend: Y = 1 + 2*t + # Pre periods: t=1→3, t=2→5, t=3→7 + # OLS fit with centered time: Y = α + β*(t - t_mean) + # t_mean = 2.0, so t_centered = [-1, 0, 1] + # Y = [3, 5, 7] => perfect fit: α=5 (at t_centered=0), β=2 + df = pd.DataFrame( + { + "unit": [0] * 6, + "time": [1, 2, 3, 4, 5, 6], + "y": [3.0, 5.0, 7.0, 20.0, 22.0, 24.0], # post has treatment effect + "treat": [0, 0, 0, 1, 1, 1], + } + ) + est = LWDiD(rolling="detrend") + diag = est.get_transformation_diagnostics( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + unit_diag = diag["per_unit"][0] + # Beta (slope) should be 2.0 — invariant to centering + np.testing.assert_allclose(unit_diag["beta"], 2.0, atol=1e-10) + # Alpha is intercept at centered origin: Y at t_centered=0 = Y at t=2 = 5.0 + np.testing.assert_allclose(unit_diag["alpha"], 5.0, atol=1e-10) + # R^2 should be 1.0 for perfect linear fit + np.testing.assert_allclose(unit_diag["r_squared"], 1.0, atol=1e-10) + + def test_degrees_of_freedom_formula(self, simple_panel): + """Verify df = N - K - 2 per paper Section 2.4. + + Without controls: df = N - 0 - 2 = N - 2 + """ + est = LWDiD(rolling="demean", estimation_method="reg") + res = est.fit(simple_panel, outcome="y", unit="unit", time="time", treatment="treat") + # N = 40 units, K = 0 controls → df = 40 - 0 - 2 = 38 + assert res.df_inference == 38 + + def test_ra_interaction_term_present(self, panel_with_controls): + """Verify RA includes interaction per Eq 3.3. + + Design matrix should include [1, D, X, D*(X-X̄₁)] when controls present. + """ + est = LWDiD(rolling="demean", estimation_method="reg") + res = est.fit( + panel_with_controls, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1"], + ) + assert np.isfinite(res.att) + assert np.isfinite(res.se) + + def test_cluster_uses_g_minus_1_df(self, simple_panel): + """Verify cluster-robust uses df = G - 1.""" + est = LWDiD(rolling="demean", cluster="unit") + res = est.fit(simple_panel, outcome="y", unit="unit", time="time", treatment="treat") + # G = 40 units as clusters → df = 39 + assert res.df_inference == 39 + + def test_t_stat_equals_att_over_se(self, simple_panel): + """Verify t_stat = att / se (basic algebra check).""" + est = LWDiD() + res = est.fit(simple_panel, outcome="y", unit="unit", time="time", treatment="treat") + if np.isfinite(res.t_stat) and np.isfinite(res.se) and res.se > 0: + np.testing.assert_allclose(res.t_stat, res.att / res.se, rtol=1e-10) + + def test_confidence_interval_symmetric(self, simple_panel): + """Verify CI is symmetric around ATT.""" + est = LWDiD() + res = est.fit(simple_panel, outcome="y", unit="unit", time="time", treatment="treat") + ci_lower, ci_upper = res.conf_int + midpoint = (ci_lower + ci_upper) / 2 + np.testing.assert_allclose(midpoint, res.att, atol=1e-10) + + +# ============================================================ +# Class 4: Backward compatibility +# ============================================================ + + +class TestBackwardCompatibility: + """Ensure existing fit() behavior is preserved.""" + + def test_fit_unchanged_demean(self, simple_panel): + """fit() with demean gives correct result.""" + est = LWDiD(rolling="demean") + res = est.fit(simple_panel, outcome="y", unit="unit", time="time", treatment="treat") + assert isinstance(res.att, float) + assert np.isfinite(res.att) + assert 1.0 < res.att < 3.0 + + def test_fit_unchanged_detrend(self, simple_panel): + """fit() with detrend gives correct result.""" + est = LWDiD(rolling="detrend") + res = est.fit(simple_panel, outcome="y", unit="unit", time="time", treatment="treat") + assert np.isfinite(res.att) + + def test_fit_unchanged_staggered(self): + """Staggered fit still works correctly.""" + rng = np.random.default_rng(42) + records = [] + for i in range(90): + g = [0, 4, 7][i % 3] + for t in range(1, 10): + y = 1.0 + 0.05 * t + rng.normal(0, 0.2) + if g > 0 and t >= g: + y += 1.5 + records.append( + {"unit": i, "time": t, "y": y, "treat": int(g > 0 and t >= g), "cohort": g} + ) + df = pd.DataFrame(records) + est = LWDiD(control_group="never_treated") + res = est.fit( + df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="cohort" + ) + assert np.isfinite(res.att) + assert 1.0 < res.att < 2.5 + + def test_bootstrap_unchanged(self, simple_panel): + """Bootstrap still works after transform changes.""" + est = LWDiD(n_bootstrap=20) + res = est.fit(simple_panel, outcome="y", unit="unit", time="time", treatment="treat") + assert np.isfinite(res.att) + assert np.isfinite(res.se) diff --git a/tests/test_lwdid_equivalence.py b/tests/test_lwdid_equivalence.py new file mode 100644 index 00000000..69e4dfee --- /dev/null +++ b/tests/test_lwdid_equivalence.py @@ -0,0 +1,531 @@ +"""Numerical equivalence tests: diff-diff LWDiD vs lwdid-py reference. + +These tests require lwdid>=0.2.2 (optional dev dependency). +Run with: pytest tests/test_lwdid_equivalence.py -v +Skipped automatically if lwdid is not installed. + +Tolerance standards (per Lee & Wooldridge paper precision requirements): +- RA + classical/HC1: atol=1e-10 (direct matrix inversion, deterministic) +- RA + cluster: atol=1e-8 (grouping introduces floating-point reassociation) +- IPW/IPWRA: atol=1e-6 (logit optimization path may differ) +- PSM: atol=1e-4 (matching tie-breaking may differ) +- Staggered aggregation: atol=1e-6 (multi-layer aggregation) +""" + +import numpy as np +import pandas as pd +import pytest + +# ============================================================ +# Test Data Generators (deterministic, shared between both packages) +# ============================================================ + + +def _generate_common_timing_panel(n=100, T=8, post_start=6, true_att=2.0, n_controls=1, seed=42): + """Generate balanced panel for common-timing tests. + + Produces columns compatible with BOTH lwdid-py and diff-diff APIs: + - unit: unit identifier + - time: time period (1..T) + - y: outcome variable + - treat: unit-level treatment indicator (time-invariant) + - post: post-treatment indicator (0 in pre, 1 in post) + - d: treatment status per obs (treat * post) + - x1: a covariate + """ + rng = np.random.default_rng(seed) + n_treated = n // 3 + + rows = [] + for i in range(n): + is_treated = i < n_treated + unit_fe = rng.normal(0, 2) + trend_slope = rng.normal(0.3, 0.1) + x1 = rng.normal() + int(is_treated) * 0.3 + for t in range(1, T + 1): + time_trend = trend_slope * t + noise = rng.normal(0, 0.3) + is_post = int(t >= post_start) + treatment_effect = true_att if (is_treated and is_post) else 0.0 + y = unit_fe + time_trend + noise + treatment_effect + 0.5 * x1 + rows.append( + { + "unit": i, + "time": t, + "y": y, + "treat": int(is_treated), + "post": is_post, + "d": int(is_treated and bool(is_post)), + "x1": x1, + } + ) + + return pd.DataFrame(rows) + + +def _generate_staggered_panel(n=120, T=10, seed=42): + """Generate staggered adoption panel. + + Produces columns compatible with BOTH packages: + - unit: unit identifier + - time: time period (1..T) + - y: outcome variable + - treat: current treatment status (0/1) + - cohort: first treatment time (0 = never-treated) + - gvar: cohort var for lwdid-py (NaN for never-treated) + - x1: a covariate + """ + rng = np.random.default_rng(seed) + cohorts = [0, 4, 6, 8] # 0 = never-treated + true_att = 1.5 + + rows = [] + for i in range(n): + g = cohorts[i % len(cohorts)] + unit_fe = rng.normal(0, 2) + x1 = rng.normal() + for t in range(1, T + 1): + is_post = int(g > 0 and t >= g) + effect = true_att * is_post + y = unit_fe + 0.2 * t + rng.normal(0, 0.2) + effect + rows.append( + { + "unit": i, + "time": t, + "y": y, + "treat": is_post, + "d": int(g > 0), + "post": is_post, + "cohort": g, + "gvar": g if g > 0 else np.nan, + "x1": x1, + } + ) + + return pd.DataFrame(rows) + + +# ============================================================ +# Helper functions to run both packages +# ============================================================ + + +def _run_lwdid_py_common(df, rolling, estimator, vce, controls=None, cluster_var=None): + """Run lwdid-py on common-timing panel.""" + from lwdid import lwdid as lwdid_func + + kwargs = dict( + data=df.copy(), + y="y", + d="treat", + ivar="unit", + tvar="time", + post="post", + rolling=rolling, + estimator=estimator, + verbose="quiet", + ) + if vce is not None: + if vce == "cluster": + kwargs["vce"] = "cluster" + kwargs["cluster_var"] = cluster_var or "unit" + else: + kwargs["vce"] = vce + if controls: + kwargs["controls"] = controls + return lwdid_func(**kwargs) + + +def _run_diff_diff_common(df, rolling, estimator, vce, controls=None, cluster=None): + """Run diff-diff LWDiD on common-timing panel. + + The estimator/vce spec tokens follow lwdid-py vocabulary; they are + mapped to diff-diff's canonical estimation_method/vcov_type here. + """ + from diff_diff import LWDiD + + method_map = {"ra": "reg", "ipwra": "dr"} + vce_map = {"robust": "hc1", "ols": "classical", "cluster": "hc1"} + dd_vcov = vce_map.get(vce, vce) if vce else "classical" + + model = LWDiD( + rolling=rolling, + estimation_method=method_map.get(estimator, estimator), + vcov_type=dd_vcov, + cluster=cluster, + ) + return model.fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="d", + covariates=controls, + ) + + +def _run_lwdid_py_staggered( + df, rolling, estimator, vce, control_group, controls=None, cluster_var=None +): + """Run lwdid-py on staggered panel. + + Returns (result, actual_control_group_used) tuple because lwdid-py may + auto-switch from 'not_yet_treated' to 'never_treated' when aggregate='cohort'. + + Aggregation basis: we explicitly request aggregate="overall" so that + lwdid-py estimates the overall ATT from a single pooled cross-section + regression, the basis recommended by Lee & Wooldridge (2026, eq. 7.19), + which "automatically accounts for the correlations among the tau_g". + lwdid-py's default aggregate="cohort" instead combines per-cohort SEs via + sqrt(sum(w^2 * SE^2)), which assumes independence across cohort estimates + and therefore understates the overall SE. diff-diff's joint influence + function SE matches the eq. 7.19 pooled-regression basis (and Stata + lwdid.ado), so "overall" is the correct reference for equivalence. + """ + import warnings + + from lwdid import lwdid as lwdid_func + + kwargs = dict( + data=df.copy(), + y="y", + gvar="gvar", + ivar="unit", + tvar="time", + rolling=rolling, + estimator=estimator, + control_group=control_group, + aggregate="overall", + verbose="quiet", + ) + if vce is not None: + if vce == "cluster": + kwargs["vce"] = "cluster" + kwargs["cluster_var"] = cluster_var or "unit" + else: + kwargs["vce"] = vce + if controls: + kwargs["controls"] = controls + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = lwdid_func(**kwargs) + actual_cg = getattr(result, "control_group_used", control_group) + return result, actual_cg + + +def _run_diff_diff_staggered( + df, rolling, estimator, vce, control_group, controls=None, cluster=None +): + """Run diff-diff LWDiD on staggered panel (lwdid-py spec tokens mapped + to canonical estimation_method/vcov_type).""" + from diff_diff import LWDiD + + method_map = {"ra": "reg", "ipwra": "dr"} + vce_map = {"robust": "hc1", "ols": "classical", "cluster": "hc1"} + dd_vcov = vce_map.get(vce, vce) if vce else "classical" + + model = LWDiD( + rolling=rolling, + estimation_method=method_map.get(estimator, estimator), + vcov_type=dd_vcov, + cluster=cluster, + control_group=control_group, + ) + return model.fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="cohort", + covariates=controls, + ) + + +# ============================================================ +# Parametrized Equivalence Matrix: Common Timing +# ============================================================ + + +COMMON_TIMING_CONFIGS = [ + # (rolling, estimator, vce, use_controls, atol, description) + ("demean", "ra", None, False, 1e-10, "demean+RA+classical, no controls"), + ("demean", "ra", "hc1", False, 1e-10, "demean+RA+HC1, no controls"), + ("demean", "ra", None, True, 1e-10, "demean+RA+classical, with controls"), + ("demean", "ra", "hc1", True, 1e-10, "demean+RA+HC1, with controls"), + ("demean", "ra", "cluster", False, 1e-8, "demean+RA+cluster"), + ("demean", "ra", "cluster", True, 1e-8, "demean+RA+cluster, with controls"), + ("detrend", "ra", None, False, 1e-10, "detrend+RA+classical"), + ("detrend", "ra", "hc1", False, 1e-10, "detrend+RA+HC1"), + ("detrend", "ra", "hc1", True, 1e-10, "detrend+RA+HC1, with controls"), + ("detrend", "ra", "cluster", False, 1e-8, "detrend+RA+cluster"), + ("demean", "ipw", "hc1", True, 0.05, "demean+IPW+HC1"), + ("demean", "ipwra", "hc1", True, 0.01, "demean+IPWRA+HC1"), + ("detrend", "ipw", "hc1", True, 0.05, "detrend+IPW+HC1"), + ("detrend", "ipwra", "hc1", True, 0.01, "detrend+IPWRA+HC1"), +] + + +@pytest.mark.parametrize( + "rolling,estimator,vce,use_controls,atol,desc", + COMMON_TIMING_CONFIGS, + ids=[c[-1] for c in COMMON_TIMING_CONFIGS], +) +def test_equivalence_common_timing( + rolling, estimator, vce, use_controls, atol, desc, require_lwdid +): + """Verify numerical equivalence against lwdid-py for common timing.""" + + df = _generate_common_timing_panel(seed=42) + + # --- lwdid-py reference --- + controls_py = ["x1"] if use_controls else None + cluster_py = "unit" if vce == "cluster" else None + + ref = _run_lwdid_py_common( + df, rolling, estimator, vce, controls=controls_py, cluster_var=cluster_py + ) + + # --- diff-diff native --- + dd = _run_diff_diff_common( + df, rolling, estimator, vce, controls=controls_py, cluster=cluster_py + ) + + # --- Compare --- + np.testing.assert_allclose(dd.att, ref.att, atol=atol, err_msg=f"ATT mismatch [{desc}]") + # SE comparison + if np.isfinite(ref.se_att) and ref.se_att > 0: + np.testing.assert_allclose(dd.se, ref.se_att, atol=atol, err_msg=f"SE mismatch [{desc}]") + # t-stat comparison (use rtol for IPW/IPWRA since t-stats are large + # and differences compound from ATT+SE optimization path divergence) + if hasattr(ref, "t_stat") and np.isfinite(ref.t_stat): + if hasattr(dd, "t_stat") and np.isfinite(dd.t_stat): + t_rtol = 0.25 if estimator in ("ipw", "ipwra") else 1e-3 + np.testing.assert_allclose( + dd.t_stat, ref.t_stat, rtol=t_rtol, err_msg=f"t-stat mismatch [{desc}]" + ) + + +# ============================================================ +# Parametrized Equivalence Matrix: Staggered +# ============================================================ + + +STAGGERED_CONFIGS = [ + # (rolling, estimator, vce, control_group, controls, atol) + ("demean", "ra", "cluster", "never_treated", None, 1e-8), + ("demean", "ra", "cluster", "not_yet_treated", None, 1e-8), + ("detrend", "ra", "cluster", "never_treated", None, 1e-8), + ("demean", "ra", "hc1", "never_treated", None, 1e-8), + ("demean", "ra", "hc1", "not_yet_treated", None, 1e-8), + ("demean", "ipw", "cluster", "not_yet_treated", ["x1"], 0.01), + ("demean", "ipwra", "cluster", "not_yet_treated", ["x1"], 0.01), + ("demean", "ipw", "hc1", "never_treated", ["x1"], 0.01), + ("demean", "ipwra", "hc1", "never_treated", ["x1"], 0.01), +] + + +@pytest.mark.parametrize( + "rolling,estimator,vce,control_group,controls,atol", + STAGGERED_CONFIGS, + ids=[f"{r}+{e}+{v}+{cg}" for r, e, v, cg, _, _ in STAGGERED_CONFIGS], +) +def test_equivalence_staggered( + rolling, estimator, vce, control_group, controls, atol, require_lwdid +): + """Verify numerical equivalence against lwdid-py for staggered designs.""" + df = _generate_staggered_panel(seed=42) + + cluster_var = "unit" if vce == "cluster" else None + + # --- lwdid-py reference --- + # lwdid-py may auto-switch 'not_yet_treated' -> 'never_treated' + # when aggregate='cohort' (default). Use actual control group for fair comparison. + ref, actual_cg = _run_lwdid_py_staggered( + df, rolling, estimator, vce, control_group, controls=controls, cluster_var=cluster_var + ) + + # --- diff-diff native (use the control group lwdid-py actually used) --- + dd = _run_diff_diff_staggered( + df, rolling, estimator, vce, actual_cg, controls=controls, cluster=cluster_var + ) + + # --- Compare overall ATT --- + np.testing.assert_allclose( + dd.att, + ref.att, + atol=atol, + err_msg=f"Staggered ATT mismatch [{rolling}/{estimator}/{vce}/{control_group}]", + ) + # SE comparison: both sides use the LW 2026 eq. 7.19 pooled-regression + # basis (lwdid-py aggregate="overall" vs diff-diff joint influence + # function). rtol=0.01 absorbs the small difference in where the HC1 + # dof correction is applied (per-cell vs overall regression). IPW-family + # estimators get a looser rtol since the logit optimization path differs. + if np.isfinite(ref.se_att) and ref.se_att > 0: + se_rtol = 0.05 if estimator in ("ipw", "ipwra") else 0.01 + np.testing.assert_allclose( + dd.se, + ref.se_att, + rtol=se_rtol, + err_msg=f"Staggered SE mismatch [{rolling}/{estimator}/{vce}/{control_group}]", + ) + + +# ============================================================ +# Multi-seed robustness +# ============================================================ + + +@pytest.mark.parametrize("seed", [1, 7, 42, 99, 123]) +def test_equivalence_multi_seed(seed, require_lwdid): + """Verify equivalence holds across multiple random seeds.""" + df = _generate_common_timing_panel(seed=seed) + + ref = _run_lwdid_py_common(df, "demean", "ra", "hc1") + dd = _run_diff_diff_common(df, "demean", "ra", "hc1") + + np.testing.assert_allclose(dd.att, ref.att, atol=1e-10, err_msg=f"Seed {seed} ATT mismatch") + if np.isfinite(ref.se_att) and ref.se_att > 0: + np.testing.assert_allclose( + dd.se, ref.se_att, atol=1e-10, err_msg=f"Seed {seed} SE mismatch" + ) + + +@pytest.mark.parametrize("seed", [0, 1, 42, 99, 123]) +def test_equivalence_detrend_multiseed(seed, require_lwdid): + """Detrend+RA path across multiple seeds.""" + df = _generate_common_timing_panel(seed=seed) + + ref = _run_lwdid_py_common(df, "detrend", "ra", "hc1") + dd = _run_diff_diff_common(df, "detrend", "ra", "hc1") + + np.testing.assert_allclose( + dd.att, ref.att, atol=1e-10, err_msg=f"Detrend ATT mismatch at seed={seed}" + ) + + +@pytest.mark.parametrize("seed", [0, 42, 99]) +def test_equivalence_staggered_multiseed(seed, require_lwdid): + """Staggered RA+demean across multiple seeds.""" + df = _generate_staggered_panel(seed=seed) + + ref, actual_cg = _run_lwdid_py_staggered(df, "demean", "ra", "hc1", "never_treated") + dd = _run_diff_diff_staggered(df, "demean", "ra", "hc1", actual_cg) + + np.testing.assert_allclose( + dd.att, ref.att, atol=1e-8, err_msg=f"Staggered ATT mismatch at seed={seed}" + ) + + +# ============================================================ +# Transformation intermediate values +# ============================================================ + + +def test_transformed_outcomes_match(require_lwdid): + """Verify that transformed Y values match between implementations. + + Since we cannot easily access internal transformed data from lwdid-py, + we verify through ATT (which is a direct function of the transformed + outcomes) at machine-epsilon tolerance. + """ + df = _generate_common_timing_panel(seed=42) + + for rolling in ["demean", "detrend"]: + ref = _run_lwdid_py_common(df, rolling, "ra", None) + dd = _run_diff_diff_common(df, rolling, "ra", None) + np.testing.assert_allclose( + dd.att, ref.att, atol=1e-10, err_msg=f"{rolling} transform mismatch" + ) + + +# ============================================================ +# Inference Equivalence +# ============================================================ + + +def test_equivalence_t_stat_and_pvalue(require_lwdid): + """t-stat and p-value should match between implementations.""" + df = _generate_common_timing_panel(seed=42) + + ref = _run_lwdid_py_common(df, "demean", "ra", "hc1") + dd = _run_diff_diff_common(df, "demean", "ra", "hc1") + + # t-stat + if hasattr(ref, "t_stat") and np.isfinite(ref.t_stat): + np.testing.assert_allclose(dd.t_stat, ref.t_stat, rtol=1e-3, err_msg="t-stat mismatch") + + # p-value + if hasattr(ref, "pvalue") and np.isfinite(ref.pvalue): + np.testing.assert_allclose(dd.p_value, ref.pvalue, rtol=1e-2, err_msg="p-value mismatch") + + +def test_equivalence_confidence_interval(require_lwdid): + """CI bounds should match between implementations.""" + df = _generate_common_timing_panel(seed=42) + + ref = _run_lwdid_py_common(df, "demean", "ra", "hc1") + dd = _run_diff_diff_common(df, "demean", "ra", "hc1") + + if hasattr(ref, "ci_lower") and np.isfinite(ref.ci_lower): + np.testing.assert_allclose( + dd.conf_int[0], ref.ci_lower, rtol=1e-3, err_msg="CI lower mismatch" + ) + if hasattr(ref, "ci_upper") and np.isfinite(ref.ci_upper): + np.testing.assert_allclose( + dd.conf_int[1], ref.ci_upper, rtol=1e-3, err_msg="CI upper mismatch" + ) + + +# ============================================================ +# Sample Size Equivalence +# ============================================================ + + +def test_equivalence_sample_sizes(require_lwdid): + """n_treated and n_control should match.""" + df = _generate_common_timing_panel(seed=42) + + ref = _run_lwdid_py_common(df, "demean", "ra", "hc1") + dd = _run_diff_diff_common(df, "demean", "ra", "hc1") + + assert dd.n_treated == ref.n_treated + assert dd.n_control == ref.n_control + + +# ============================================================ +# Edge Case Equivalence +# ============================================================ + + +def test_equivalence_single_post_period(require_lwdid): + """Single post-treatment period should still match.""" + df = _generate_common_timing_panel(n=80, T=6, post_start=6, seed=42) + + ref = _run_lwdid_py_common(df, "demean", "ra", "hc1") + dd = _run_diff_diff_common(df, "demean", "ra", "hc1") + + np.testing.assert_allclose(dd.att, ref.att, atol=1e-10) + + +def test_equivalence_many_periods(require_lwdid): + """Many pre/post periods should still match.""" + df = _generate_common_timing_panel(n=80, T=18, post_start=10, seed=42) + + ref = _run_lwdid_py_common(df, "demean", "ra", "hc1") + dd = _run_diff_diff_common(df, "demean", "ra", "hc1") + + np.testing.assert_allclose(dd.att, ref.att, atol=1e-10) + + +def test_equivalence_large_sample(require_lwdid): + """Larger sample size should maintain equivalence.""" + df = _generate_common_timing_panel(n=500, T=8, post_start=6, seed=42) + + ref = _run_lwdid_py_common(df, "demean", "ra", "hc1") + dd = _run_diff_diff_common(df, "demean", "ra", "hc1") + + np.testing.assert_allclose(dd.att, ref.att, atol=1e-10) + if np.isfinite(ref.se_att) and ref.se_att > 0: + np.testing.assert_allclose(dd.se, ref.se_att, atol=1e-10) diff --git a/tests/test_lwdid_numerics.py b/tests/test_lwdid_numerics.py new file mode 100644 index 00000000..768e8dec --- /dev/null +++ b/tests/test_lwdid_numerics.py @@ -0,0 +1,474 @@ +"""Numerical precision and edge case tests for LWDiD.""" + +import time +import warnings + +import numpy as np +import pandas as pd + +from diff_diff import LWDiD, LWDiDResults + +# ─── Data Helpers ─────────────────────────────────────────────────────────── + + +def _make_common_timing_panel( + n_treated=30, + n_control=50, + n_pre=5, + n_post=3, + true_att=2.0, + seed=42, +): + """Generate balanced common-timing panel with known ATT.""" + rng = np.random.default_rng(seed) + n_units = n_treated + n_control + n_periods = n_pre + n_post + + rows = [] + for i in range(n_units): + is_treated = i < n_treated + unit_fe = rng.normal(0, 1) + for t in range(1, n_periods + 1): + time_trend = 0.3 * t + noise = rng.normal(0, 0.5) + post = 1 if t > n_pre else 0 + treat = 1 if (is_treated and post) else 0 + y = unit_fe + time_trend + noise + (true_att if treat else 0) + rows.append( + { + "unit": i, + "time": t, + "y": y, + "treat": treat, + } + ) + return pd.DataFrame(rows) + + +def _make_large_panel(n_units=1000, n_periods=20, seed=42): + """Large panel for performance testing.""" + rng = np.random.default_rng(seed) + n_treated = n_units // 3 + n_pre = n_periods // 2 + + unit_ids = np.repeat(np.arange(n_units), n_periods) + time_ids = np.tile(np.arange(1, n_periods + 1), n_units) + + is_treated = (unit_ids < n_treated).astype(float) + is_post = (time_ids > n_pre).astype(float) + treat = is_treated * is_post + + # Unit FEs + time trend + noise + treatment effect + unit_fes = rng.normal(0, 2, size=n_units) + y = unit_fes[unit_ids] + 0.3 * time_ids + rng.normal(0, 0.5, size=len(unit_ids)) + 2.0 * treat + + return pd.DataFrame( + { + "unit": unit_ids, + "time": time_ids, + "y": y, + "treat": treat.astype(int), + } + ) + + +# ─── Hand-Computed ATT Tests ─────────────────────────────────────────────── + + +class TestLWDiDHandComputed: + """Tests where ATT can be computed by hand.""" + + def test_hand_computed_att_3units(self): + """3 units, 4 periods, hand-computable ATT. + + Unit 0 (control): y = [1, 2, 3, 4], pre_mean = 1.5 + demeaned post: [3-1.5, 4-1.5] = [1.5, 2.5] → avg = 2.0 + Unit 1 (control): y = [2, 4, 6, 8], pre_mean = 3 + demeaned post: [6-3, 8-3] = [3, 5] → avg = 4.0 + Unit 2 (treated): y = [1, 3, 10, 12], pre_mean = 2 + demeaned post: [10-2, 12-2] = [8, 10] → avg = 9.0 + + Cross-section: control_mean = (2.0 + 4.0)/2 = 3.0 + treated_mean = 9.0 + ATT = 9.0 - 3.0 = 6.0 + + But RA is y = alpha + tau*D, so: + Intercept = mean of controls = 3.0 + tau = mean(treated) - mean(controls) = 9.0 - 3.0 = 6.0 + """ + df = pd.DataFrame( + { + "unit": [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2], + "time": [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4], + "y": [1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0, 1.0, 3.0, 10.0, 12.0], + "treat": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1], + } + ) + res = LWDiD(rolling="demean", estimation_method="reg", vcov_type="classical").fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + np.testing.assert_allclose(res.att, 6.0, atol=1e-10) + + def test_hand_computed_att_zero_effect(self): + """When treatment effect is exactly 0, ATT should be ~0. + + Both treated and controls have same DGP: y = unit_fe + t. + """ + df = pd.DataFrame( + { + "unit": [0, 0, 0, 1, 1, 1, 2, 2, 2], + "time": [1, 2, 3, 1, 2, 3, 1, 2, 3], + "y": [1.0, 2.0, 3.0, 2.0, 3.0, 4.0, 3.0, 4.0, 5.0], + "treat": [0, 0, 0, 0, 0, 0, 0, 0, 1], + } + ) + res = LWDiD(rolling="demean", estimation_method="reg", vcov_type="classical").fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + # All units have pre_mean = 1.5, 2.5, 3.5 + # Post demeaned: control = [3-1.5, 4-2.5] = [1.5, 1.5] avg=1.5 + # Treated: 5-3.5 = 1.5 + # ATT = 1.5 - 1.5 = 0 + np.testing.assert_allclose(res.att, 0.0, atol=1e-10) + + def test_detrend_perfect_linear_zero_effect(self): + """Perfect linear trend, no treatment effect → ATT = 0. + + All units follow y = a_i + b_i * t with no treatment effect. + After detrending, residuals are 0 everywhere. + """ + df = pd.DataFrame( + { + "unit": [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2], + "time": [1, 2, 3, 4, 1, 2, 3, 4, 1, 2, 3, 4], + "y": [1.0, 2.0, 3.0, 4.0, 2.0, 4.0, 6.0, 8.0, 0.0, 1.0, 2.0, 3.0], + "treat": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1], + } + ) + res = LWDiD(rolling="detrend", estimation_method="reg", vcov_type="classical").fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + np.testing.assert_allclose(res.att, 0.0, atol=1e-10) + + def test_detrend_with_known_effect(self): + """Linear trend + constant treatment effect. + + Controls: y = a_i + t (perfectly linear) + Treated: y = a_i + t in pre, y = a_i + t + 3 in post + After detrend, control residuals = 0, treated residuals = 3. + ATT = 3 - 0 = 3. + """ + df = pd.DataFrame( + { + "unit": [0] * 4 + [1] * 4 + [2] * 4 + [3] * 4, + "time": [1, 2, 3, 4] * 4, + "y": [ + 2.0, + 3.0, + 4.0, + 5.0, # control 0: y = 1 + t + 3.0, + 4.0, + 5.0, + 6.0, # control 1: y = 2 + t + 4.0, + 5.0, + 6.0, + 7.0, # control 2: y = 3 + t + 2.0, + 3.0, + 7.0, + 8.0, # treated: y = 1 + t + 3*post + ], + "treat": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 1, + 1, + ], + } + ) + res = LWDiD(rolling="detrend", estimation_method="reg", vcov_type="classical").fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + np.testing.assert_allclose(res.att, 3.0, atol=1e-10) + + +# ─── Numerical Precision Tests ────────────────────────────────────────────── + + +class TestLWDiDNumericalPrecision: + """Test numerical stability with challenging data configurations.""" + + def test_collinear_controls_handled(self): + """Rank-deficient design matrix should not crash.""" + panel = _make_common_timing_panel(seed=11) + # Add duplicate (unit-constant) control column + rng = np.random.default_rng(11) + units = panel["unit"].unique() + xmap = dict(zip(units, rng.normal(size=len(units)))) + panel["x1"] = panel["unit"].map(xmap) + panel["x2"] = panel["x1"] # perfectly collinear + + # Should produce a result (possibly with warning), not crash + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(estimation_method="reg").fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1", "x2"], + ) + assert np.isfinite(res.att) + + def test_near_singular_design(self): + """Near-singular design should still produce finite estimate.""" + rng = np.random.default_rng(22) + panel = _make_common_timing_panel(seed=22) + # Add nearly collinear (unit-constant) controls + units = panel["unit"].unique() + xmap = dict(zip(units, rng.normal(size=len(units)))) + emap = dict(zip(units, rng.normal(0, 1e-8, size=len(units)))) + panel["x1"] = panel["unit"].map(xmap) + panel["x2"] = panel["x1"] + panel["unit"].map(emap) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(estimation_method="reg").fit( + panel, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x1", "x2"], + ) + assert np.isfinite(res.att) + + def test_zero_variance_outcome_handled(self): + """Constant outcome should be handled gracefully.""" + df = pd.DataFrame( + { + "unit": [0, 0, 0, 1, 1, 1, 2, 2, 2], + "time": [1, 2, 3, 1, 2, 3, 1, 2, 3], + "y": [5.0] * 9, # constant outcome + "treat": [0, 0, 0, 0, 0, 0, 0, 0, 1], + } + ) + # Should not crash; ATT should be 0 or NaN + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(rolling="demean", vcov_type="classical").fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + # With constant outcome, demeaned values are all 0, ATT = 0 + assert res.att == 0.0 or np.isnan(res.att) + + def test_single_treated_unit(self): + """Only 1 treated unit should still produce a result.""" + df = pd.DataFrame( + { + "unit": [0, 0, 0, 1, 1, 1, 2, 2, 2], + "time": [1, 2, 3, 1, 2, 3, 1, 2, 3], + "y": [1.0, 2.0, 3.0, 2.0, 3.0, 4.0, 1.0, 2.0, 8.0], + "treat": [0, 0, 0, 0, 0, 0, 0, 0, 1], + } + ) + res = LWDiD(rolling="demean", vcov_type="classical").fit( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert isinstance(res, LWDiDResults) + assert np.isfinite(res.att) + assert res.n_treated == 1 + + def test_large_outcome_values(self): + """Large outcome values should not cause overflow.""" + panel = _make_common_timing_panel(seed=33) + panel["y"] = panel["y"] * 1e8 + + res = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + assert np.isfinite(res.att) + assert np.isfinite(res.se) + + def test_small_outcome_values(self): + """Small outcome values should not underflow.""" + panel = _make_common_timing_panel(seed=44) + panel["y"] = panel["y"] * 1e-8 + + res = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + assert np.isfinite(res.att) + + def test_negative_outcomes(self): + """Negative outcomes should work fine.""" + panel = _make_common_timing_panel(seed=55) + panel["y"] = panel["y"] - 100 # shift all negative + + res = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + assert np.isfinite(res.att) + # ATT should still be positive (shift doesn't affect demeaned values) + assert res.att > 0 + + +# ─── Performance Tests ────────────────────────────────────────────────────── + + +class TestLWDiDPerformance: + """Test that estimation completes in reasonable time.""" + + def test_large_panel_performance(self): + """1000 units × 20 periods should complete in reasonable time.""" + panel = _make_large_panel(n_units=1000, n_periods=20) + start = time.time() + res = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + elapsed = time.time() - start + assert elapsed < 30 # Should complete in < 30 seconds + assert np.isfinite(res.att) + + def test_moderate_staggered_performance(self): + """200 units × 10 periods staggered should be fast.""" + from tests.test_lwdid import _make_staggered_panel + + panel = _make_staggered_panel(n_units=200, n_periods=10, seed=77) + start = time.time() + res = LWDiD(control_group="never_treated").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat", first_treat="cohort" + ) + elapsed = time.time() - start + assert elapsed < 30 + assert np.isfinite(res.att) + + +# ─── VCE Consistency Tests ────────────────────────────────────────────────── + + +class TestLWDiDVCEConsistency: + """Test variance-covariance estimation properties.""" + + def test_hc1_se_positive(self): + """HC1 SE must be strictly positive when ATT is identified.""" + panel = _make_common_timing_panel() + res = LWDiD(vcov_type="hc1").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert res.se > 0 + + def test_cluster_se_invariant_to_row_order(self): + """Shuffling rows should not change cluster-robust SE.""" + panel = _make_common_timing_panel(seed=66) + panel["cluster_id"] = panel["unit"] % 10 + + # Fit on original order + res1 = LWDiD(cluster="cluster_id").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + + # Shuffle rows + panel_shuffled = panel.sample(frac=1, random_state=99).reset_index(drop=True) + res2 = LWDiD(cluster="cluster_id").fit( + panel_shuffled, + outcome="y", + unit="unit", + time="time", + treatment="treat", + ) + + np.testing.assert_allclose(res1.att, res2.att, atol=1e-12) + np.testing.assert_allclose(res1.se, res2.se, atol=1e-12) + + def test_vcov_symmetric(self): + """VCE matrix must be symmetric.""" + panel = _make_common_timing_panel() + res = LWDiD(vcov_type="hc1").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + if res.vcov is not None: + np.testing.assert_allclose(res.vcov, res.vcov.T, atol=1e-14) + + def test_vcov_positive_semidefinite(self): + """VCE matrix diagonal should be non-negative.""" + panel = _make_common_timing_panel() + res = LWDiD(vcov_type="hc1").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + if res.vcov is not None: + diag = np.diag(res.vcov) + assert np.all(diag >= -1e-15) # allow small numerical error + + def test_se_consistent_with_vcov(self): + """SE should equal sqrt(vcov[1,1]) for the treatment coefficient.""" + panel = _make_common_timing_panel() + res = LWDiD(vcov_type="hc1", estimation_method="reg").fit( + panel, outcome="y", unit="unit", time="time", treatment="treat" + ) + if res.vcov is not None: + expected_se = np.sqrt(max(res.vcov[1, 1], 0.0)) + np.testing.assert_allclose(res.se, expected_se, atol=1e-14) + + +# ─── Determinism Tests ────────────────────────────────────────────────────── + + +class TestLWDiDDeterminism: + """Test that results are deterministic (same input → same output).""" + + def test_same_data_same_result(self): + """Running twice on same data gives identical results.""" + panel = _make_common_timing_panel(seed=42) + + res1 = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + res2 = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + + assert res1.att == res2.att + assert res1.se == res2.se + assert res1.t_stat == res2.t_stat + + def test_copy_invariance(self): + """Deep copy of data should give same results.""" + panel = _make_common_timing_panel(seed=42) + panel_copy = panel.copy(deep=True) + + res1 = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + res2 = LWDiD().fit(panel_copy, outcome="y", unit="unit", time="time", treatment="treat") + + assert res1.att == res2.att + assert res1.se == res2.se + + +# ─── Multiple Post-Period Aggregation ─────────────────────────────────────── + + +class TestLWDiDPostPeriodAggregation: + """Test that multiple post-periods are correctly averaged.""" + + def test_single_post_period(self): + """Single post period = no averaging needed.""" + panel = _make_common_timing_panel(n_pre=5, n_post=1, seed=42) + res = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + assert np.isfinite(res.att) + + def test_many_post_periods(self): + """Many post periods should be averaged correctly.""" + panel = _make_common_timing_panel(n_pre=3, n_post=10, seed=42) + res = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + assert np.isfinite(res.att) + assert res.att > 0 # True ATT = 2.0 + + def test_more_pre_than_post(self): + """Many pre periods, few post.""" + panel = _make_common_timing_panel(n_pre=10, n_post=2, seed=42) + res = LWDiD().fit(panel, outcome="y", unit="unit", time="time", treatment="treat") + assert np.isfinite(res.att) + assert res.att > 0 diff --git a/tests/test_lwdid_randomization_inference.py b/tests/test_lwdid_randomization_inference.py new file mode 100644 index 00000000..9b35d3f0 --- /dev/null +++ b/tests/test_lwdid_randomization_inference.py @@ -0,0 +1,210 @@ +"""Tests for lwdid_randomization module.""" + +import numpy as np +import pytest + +from diff_diff.lwdid_randomization import ( + _compute_pvalue, + randomization_inference, +) + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def cross_section_data(): + rng = np.random.default_rng(42) + n = 100 + y = np.concatenate([rng.normal(2, 0.5, 30), rng.normal(0, 0.5, 70)]) + treatment = np.array([1.0] * 30 + [0.0] * 70) + cluster_ids = np.repeat(np.arange(20), 5) + controls = rng.normal(0, 1, (n, 2)) + return y, treatment, cluster_ids, controls + + +# --------------------------------------------------------------------------- +# Result fields +# --------------------------------------------------------------------------- + + +class TestRandomizationResultFields: + """Test that RandomizationResult has all expected fields.""" + + def test_result_fields_present(self, cross_section_data): + y, treatment, _, _ = cross_section_data + r = randomization_inference(y, treatment, n_reps=200, seed=0) + assert hasattr(r, "pvalue") + assert hasattr(r, "att_observed") + assert hasattr(r, "att_distribution") + assert hasattr(r, "n_reps") + assert hasattr(r, "n_valid") + assert hasattr(r, "n_failed") + assert hasattr(r, "failure_rate") + assert hasattr(r, "method") + assert hasattr(r, "seed") + + def test_result_types(self, cross_section_data): + y, treatment, _, _ = cross_section_data + r = randomization_inference(y, treatment, n_reps=200, seed=0) + assert isinstance(r.pvalue, float) + assert isinstance(r.att_observed, float) + assert isinstance(r.att_distribution, np.ndarray) + assert isinstance(r.n_reps, int) + assert isinstance(r.n_valid, int) + assert isinstance(r.n_failed, int) + assert isinstance(r.failure_rate, float) + assert isinstance(r.method, str) + + +# --------------------------------------------------------------------------- +# Permutation preserves N_treated +# --------------------------------------------------------------------------- + + +class TestPermutationPreservation: + """Permutation should preserve number of treated units.""" + + def test_permutation_preserves_n_treated(self, cross_section_data): + y, treatment, _, _ = cross_section_data + r = randomization_inference(y, treatment, method="permutation", n_reps=500, seed=0) + # With permutation, no draws are degenerate + assert r.n_failed == 0 + assert r.failure_rate == 0.0 + + def test_bootstrap_method_removed(self, cross_section_data): + # Fix-wave review finding: label resampling WITH replacement is not + # Fisher randomization inference; the mode is removed. + y, treatment, *_ = cross_section_data + with pytest.raises(ValueError, match="method='bootstrap' has been removed"): + randomization_inference(y, treatment, method="bootstrap", n_reps=100, seed=0) + + def test_pvalue_in_0_1_permutation(self, cross_section_data): + y, treatment, _, _ = cross_section_data + r = randomization_inference(y, treatment, method="permutation", n_reps=500, seed=42) + assert 0.0 <= r.pvalue <= 1.0 + + def test_clear_treatment_effect_detected(self, cross_section_data): + """With a clear treatment effect, p-value should be small.""" + y, treatment, _, _ = cross_section_data + r = randomization_inference(y, treatment, method="permutation", n_reps=999, seed=0) + assert r.pvalue < 0.05 + + +# --------------------------------------------------------------------------- +# With and without controls +# --------------------------------------------------------------------------- + + +class TestControls: + """Test with and without control variables.""" + + def test_without_controls(self, cross_section_data): + y, treatment, _, _ = cross_section_data + r = randomization_inference(y, treatment, n_reps=200, seed=0) + assert np.isfinite(r.att_observed) + assert r.n_valid > 0 + + def test_with_controls(self, cross_section_data): + y, treatment, _, controls = cross_section_data + r = randomization_inference(y, treatment, controls=controls, n_reps=200, seed=0) + assert np.isfinite(r.att_observed) + assert r.n_valid > 0 + + +# --------------------------------------------------------------------------- +# Degenerate data handling +# --------------------------------------------------------------------------- + + +class TestDegenerateData: + """Test handling of degenerate inputs.""" + + def test_all_treated_raises(self): + y = np.array([1.0, 2.0, 3.0, 4.0]) + treatment = np.array([1.0, 1.0, 1.0, 1.0]) + with pytest.raises(ValueError): + randomization_inference(y, treatment, n_reps=100) + + def test_all_control_raises(self): + y = np.array([1.0, 2.0, 3.0, 4.0]) + treatment = np.array([0.0, 0.0, 0.0, 0.0]) + with pytest.raises(ValueError): + randomization_inference(y, treatment, n_reps=100) + + def test_too_small_sample_raises(self): + y = np.array([1.0, 2.0]) + treatment = np.array([1.0, 0.0]) + with pytest.raises(ValueError): + randomization_inference(y, treatment, n_reps=100) + + def test_invalid_method_raises(self, cross_section_data): + y, treatment, _, _ = cross_section_data + with pytest.raises(ValueError): + randomization_inference(y, treatment, method="invalid", n_reps=100) + + +# --------------------------------------------------------------------------- +# Seed reproducibility +# --------------------------------------------------------------------------- + + +class TestSeedReproducibility: + """Test that seed produces reproducible results.""" + + def test_same_seed_same_result(self, cross_section_data): + y, treatment, _, _ = cross_section_data + r1 = randomization_inference(y, treatment, n_reps=200, seed=123) + r2 = randomization_inference(y, treatment, n_reps=200, seed=123) + assert r1.pvalue == r2.pvalue + np.testing.assert_array_equal(r1.att_distribution, r2.att_distribution) + + def test_different_seed_different_result(self, cross_section_data): + y, treatment, _, _ = cross_section_data + r1 = randomization_inference(y, treatment, n_reps=200, seed=1) + r2 = randomization_inference(y, treatment, n_reps=200, seed=2) + # Distributions should differ (extremely unlikely to be equal) + assert not np.array_equal(r1.att_distribution, r2.att_distribution) + + +# --------------------------------------------------------------------------- +# Tie handling ('at least as extreme' convention) +# --------------------------------------------------------------------------- + + +class TestTieHandling: + """Ties must count as 'at least as extreme' (>=), not strictly greater.""" + + def test_constant_outcome_all_ties_pvalue_is_one(self): + """Constant outcome: every permutation ATT ties with the observed + ATT (all zero), so the two-sided p-value must be exactly 1.0.""" + y = np.full(40, 3.0) + treatment = np.array([1.0] * 15 + [0.0] * 25) + r = randomization_inference(y, treatment, method="permutation", n_reps=999, seed=0) + assert r.pvalue == 1.0 + + def test_compute_pvalue_full_tie_distribution(self): + """All replications tied with the observed statistic -> p == 1.0.""" + att_dist = np.zeros(999) + pvalue, n_valid, n_failed = _compute_pvalue(att_dist, att_obs=0.0) + assert pvalue == 1.0 + assert n_valid == 999 + assert n_failed == 0 + + def test_compute_pvalue_half_tie_distribution(self): + """Half the replications tie in absolute value, the rest are less + extreme: p = (n_tied + 1) / (n_valid + 1) under the >= rule.""" + att_dist = np.concatenate([np.full(50, 1.0), np.full(49, 0.0)]) + pvalue, n_valid, _ = _compute_pvalue(att_dist, att_obs=-1.0) + assert n_valid == 99 + assert pvalue == pytest.approx((50 + 1) / (99 + 1)) + + def test_discrete_outcome_pvalue_near_theoretical(self): + """Binary outcome with a coarse permutation distribution: the exact + randomization p-value is 1/3 (2 of 6 assignments are at least as + extreme), so the Monte Carlo p should be close to that.""" + y = np.array([1.0, 1.0, 0.0, 0.0]) + treatment = np.array([1.0, 1.0, 0.0, 0.0]) + r = randomization_inference(y, treatment, method="permutation", n_reps=999, seed=42) + assert abs(r.pvalue - 1.0 / 3.0) < 0.05 diff --git a/tests/test_lwdid_results_serialization.py b/tests/test_lwdid_results_serialization.py new file mode 100644 index 00000000..93ff3509 --- /dev/null +++ b/tests/test_lwdid_results_serialization.py @@ -0,0 +1,177 @@ +"""Tests for JSON serialization of LWDiDResults.to_dict(). + +Regression tests for the shawcharles review finding that ``to_dict()`` +leaked numpy scalar types and arrays into nested dicts, so +``json.dumps(result.to_dict())`` raised TypeError. +""" + +import json + +import numpy as np +import pandas as pd +import pytest + +from diff_diff import LWDiD, generate_staggered_data +from diff_diff.lwdid_results import _json_native_key, _to_json_native + + +def _make_common_timing_panel(n_treated=20, n_control=30, n_pre=4, n_post=3, seed=11): + rng = np.random.default_rng(seed) + rows = [] + for i in range(n_treated + n_control): + is_treated = i < n_treated + unit_fe = rng.normal(0, 1) + for t in range(1, n_pre + n_post + 1): + post = t > n_pre + treat = 1 if (is_treated and post) else 0 + y = unit_fe + 0.3 * t + rng.normal(0, 0.5) + 2.0 * treat + rows.append({"unit": i, "time": t, "y": y, "treat": treat}) + return pd.DataFrame(rows) + + +@pytest.fixture(scope="module") +def staggered_data(): + return generate_staggered_data(n_units=120, n_periods=8, seed=3) + + +class TestToDictJsonSerializable: + """json.dumps(result.to_dict()) must succeed for every result flavor.""" + + def test_common_timing_roundtrip(self): + data = _make_common_timing_panel() + result = LWDiD(rolling="demean", estimation_method="reg", vcov_type="hc1").fit( + data, outcome="y", unit="unit", time="time", treatment="treat" + ) + payload = result.to_dict() + roundtrip = json.loads(json.dumps(payload)) + assert roundtrip["att"] == pytest.approx(result.att) + assert roundtrip["se"] == pytest.approx(result.se) + assert roundtrip["n_obs"] == result.n_obs + + def test_staggered_roundtrip(self, staggered_data): + result = LWDiD(rolling="demean", estimation_method="reg", vcov_type="hc1").fit( + staggered_data, + outcome="outcome", + unit="unit", + time="period", + treatment="treated", + first_treat="first_treat", + ) + payload = result.to_dict() + roundtrip = json.loads(json.dumps(payload)) + assert roundtrip["att"] == pytest.approx(result.att) + # Nested cohort dicts must contain only native types + for key, info in roundtrip["cohort_effects"].items(): + assert isinstance(key, str) + assert info["att"] == pytest.approx(result.cohort_effects[int(key)]["att"]) + assert set(roundtrip["cohort_time_effects"]) == { + f"{g},{t}" for (g, t) in result.cohort_time_effects + } + + def test_event_study_roundtrip(self, staggered_data): + result = LWDiD(rolling="demean", estimation_method="reg", n_bootstrap=99, seed=5).fit( + staggered_data, + outcome="outcome", + unit="unit", + time="period", + treatment="treated", + first_treat="first_treat", + ) + payload = result.to_dict() + roundtrip = json.loads(json.dumps(payload)) + assert "event_study_effects" in roundtrip + for key, info in roundtrip["event_study_effects"].items(): + expected = result.event_study_effects[int(key)] + assert info["effect"] == pytest.approx(expected["effect"]) + assert isinstance(info["conf_int"], list) + assert roundtrip["reference_periods"] == list(result.reference_periods) + + +def _relabel_staggered_datetime(data): + """Relabel an integer staggered panel with quarterly Timestamps.""" + date_map = { + t: pd.Timestamp("2000-01-01") + pd.DateOffset(months=3 * (int(t) - 1)) + for t in sorted(data["period"].unique()) + } + panel = data.copy() + panel["date"] = panel["period"].map(date_map) + panel["adopt"] = panel["first_treat"].map(lambda g: date_map[g] if g > 0 else pd.NaT) + return panel + + +class TestDatetimeLabelsJsonSerializable: + """Datetime/Period cohort and time labels must serialize to JSON strings. + + Regression tests: after ``_relabel_staggered_results`` restores datetime + labels, nested ``info["cohort"]``/``info["time"]`` entries were + pd.Timestamp/pd.Period objects and ``json.dumps(result.to_dict())`` + raised TypeError. + """ + + def test_datetime_staggered_roundtrip(self, staggered_data): + panel = _relabel_staggered_datetime(staggered_data) + result = LWDiD(rolling="demean", estimation_method="reg", vcov_type="hc1").fit( + panel, + outcome="outcome", + unit="unit", + time="date", + treatment="treated", + first_treat="adopt", + ) + payload = result.to_dict() + roundtrip = json.loads(json.dumps(payload)) + assert roundtrip["att"] == pytest.approx(result.att) + # Nested cohort/time labels must be ISO-8601 strings + for key, info in roundtrip["cohort_effects"].items(): + assert isinstance(key, str) + assert isinstance(info["cohort"], str) + assert pd.Timestamp(info["cohort"]) in result.cohort_effects + for info in roundtrip["cohort_time_effects"].values(): + assert isinstance(info["cohort"], str) + assert isinstance(info["time"], str) + pd.Timestamp(info["time"]) # parses back without error + + def test_period_staggered_roundtrip(self, staggered_data): + panel = _relabel_staggered_datetime(staggered_data) + panel["date"] = panel["date"].dt.to_period("Q") + panel["adopt"] = pd.PeriodIndex(panel["adopt"], freq="Q") + result = LWDiD(rolling="demean", estimation_method="reg", vcov_type="hc1").fit( + panel, + outcome="outcome", + unit="unit", + time="date", + treatment="treated", + first_treat="adopt", + ) + payload = result.to_dict() + roundtrip = json.loads(json.dumps(payload)) + assert roundtrip["att"] == pytest.approx(result.att) + # Period labels keep their frequency semantics, e.g. "2000Q1" + for info in roundtrip["cohort_effects"].values(): + assert isinstance(info["cohort"], str) + assert pd.Period(info["cohort"], freq="Q") in result.cohort_effects + + +class TestNaTSerializationContract: + """NaT values must map to None so the payload stays json.dumps-able. + + Direct unit coverage for the NaT branches of the private helpers: + the branch is unreachable through ``to_dict()`` in the current design + (never-treated cohorts are dropped before relabeling), so the contract + is pinned here explicitly. + """ + + def test_nat_maps_to_none(self): + assert _to_json_native(pd.NaT) is None + assert _to_json_native(np.datetime64("NaT")) is None + assert _json_native_key(pd.NaT) is None + # A nested dict containing NaT values must be json.dumps-able + payload = { + "cohorts": { + pd.Timestamp("2000-01-01"): {"adopt": pd.NaT}, + "never": [pd.NaT, np.datetime64("NaT")], + } + } + roundtrip = json.loads(json.dumps(_to_json_native(payload))) + assert roundtrip["cohorts"]["2000-01-01T00:00:00"]["adopt"] is None + assert roundtrip["cohorts"]["never"] == [None, None] diff --git a/tests/test_lwdid_sensitivity.py b/tests/test_lwdid_sensitivity.py new file mode 100644 index 00000000..827972ca --- /dev/null +++ b/tests/test_lwdid_sensitivity.py @@ -0,0 +1,397 @@ +"""Tests for lwdid_sensitivity module.""" + +import warnings + +import numpy as np +import pandas as pd +import pytest + +from diff_diff.lwdid_sensitivity import ( + _classify_robustness, + _compute_sensitivity_ratio, + robustness_pre_periods, + sensitivity_no_anticipation, +) + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def panel_data(): + rng = np.random.default_rng(42) + records = [] + for i in range(80): + d = int(i < 25) + for t in range(1, 9): + y = 1.0 + 0.1 * t + rng.normal(0, 0.3) + if d and t > 4: + y += 2.0 + records.append({"unit": i, "time": t, "y": y, "treat": d * int(t > 4)}) + return pd.DataFrame(records) + + +# --------------------------------------------------------------------------- +# SensitivityResult fields +# --------------------------------------------------------------------------- + + +class TestSensitivityResultFields: + """Test SensitivityResult dataclass has all expected fields.""" + + def test_result_fields_present(self, panel_data): + r = robustness_pre_periods( + panel_data, + outcome="y", + unit="unit", + time="time", + treatment="treat", + ) + assert hasattr(r, "specifications") + assert hasattr(r, "baseline_att") + assert hasattr(r, "baseline_se") + assert hasattr(r, "sensitivity_ratio") + assert hasattr(r, "robustness_level") + assert hasattr(r, "n_specifications") + + def test_result_types(self, panel_data): + r = robustness_pre_periods( + panel_data, + outcome="y", + unit="unit", + time="time", + treatment="treat", + ) + assert isinstance(r.specifications, list) + assert isinstance(r.baseline_att, float) + assert isinstance(r.baseline_se, float) + assert isinstance(r.sensitivity_ratio, float) + assert isinstance(r.robustness_level, str) + assert isinstance(r.n_specifications, int) + + +# --------------------------------------------------------------------------- +# Robustness level valid +# --------------------------------------------------------------------------- + + +class TestRobustnessLevel: + """Test robustness_level is a valid classification.""" + + VALID_LEVELS = {"highly_robust", "moderately_robust", "sensitive", "highly_sensitive"} + + def test_robustness_level_valid(self, panel_data): + r = robustness_pre_periods( + panel_data, + outcome="y", + unit="unit", + time="time", + treatment="treat", + ) + assert r.robustness_level in self.VALID_LEVELS + + def test_classify_robustness_helper(self): + assert _classify_robustness(0.05) == "highly_robust" + assert _classify_robustness(0.15) == "moderately_robust" + assert _classify_robustness(0.35) == "sensitive" + assert _classify_robustness(0.60) == "highly_sensitive" + + +# --------------------------------------------------------------------------- +# Sensitivity ratio non-negative +# --------------------------------------------------------------------------- + + +class TestSensitivityRatio: + """Test sensitivity_ratio is non-negative.""" + + def test_ratio_non_negative(self, panel_data): + r = robustness_pre_periods( + panel_data, + outcome="y", + unit="unit", + time="time", + treatment="treat", + ) + assert r.sensitivity_ratio >= 0.0 + + def test_compute_sensitivity_ratio_helper(self): + assert _compute_sensitivity_ratio(2.0, [2.0, 2.1, 1.9]) == pytest.approx(0.1) + # Single finite estimate: robustness cannot be assessed + assert np.isnan(_compute_sensitivity_ratio(2.0, [2.0])) + # Near-zero baseline: the ratio is undefined -> not estimable (NaN) + assert np.isnan(_compute_sensitivity_ratio(1e-15, [1e-15, 0.5])) + + +# --------------------------------------------------------------------------- +# Specifications list populated +# --------------------------------------------------------------------------- + + +class TestSpecifications: + """Test specifications list is populated.""" + + def test_specs_populated(self, panel_data): + r = robustness_pre_periods( + panel_data, + outcome="y", + unit="unit", + time="time", + treatment="treat", + ) + # Should have at least 1 specification + assert len(r.specifications) >= 1 + assert r.n_specifications >= 2 # baseline + at least 1 alternative + + def test_spec_has_expected_attributes(self, panel_data): + r = robustness_pre_periods( + panel_data, + outcome="y", + unit="unit", + time="time", + treatment="treat", + ) + if r.specifications: + spec = r.specifications[0] + assert hasattr(spec, "label") + assert hasattr(spec, "rolling") + assert hasattr(spec, "estimation_method") + assert hasattr(spec, "att") + assert hasattr(spec, "se") + assert hasattr(spec, "pvalue") + + +# --------------------------------------------------------------------------- +# to_dataframe() +# --------------------------------------------------------------------------- + + +class TestToDataframe: + """Test to_dataframe() returns a DataFrame.""" + + def test_to_dataframe_returns_df(self, panel_data): + r = robustness_pre_periods( + panel_data, + outcome="y", + unit="unit", + time="time", + treatment="treat", + ) + df = r.to_dataframe() + assert isinstance(df, pd.DataFrame) + assert len(df) >= 1 + assert "att" in df.columns + assert "label" in df.columns + + def test_summary_returns_string(self, panel_data): + r = robustness_pre_periods( + panel_data, + outcome="y", + unit="unit", + time="time", + treatment="treat", + ) + s = r.summary() + assert isinstance(s, str) + assert "Sensitivity" in s + + +# --------------------------------------------------------------------------- +# not_estimable classification and failure reporting +# --------------------------------------------------------------------------- + + +class TestNotEstimable: + """Failed fits must be reported as 'not_estimable', never as robust.""" + + def test_nan_baseline_ratio_is_nan(self): + assert np.isnan(_compute_sensitivity_ratio(np.nan, [np.nan, 1.0, 2.0])) + + def test_classify_nan_ratio_not_estimable(self): + assert _classify_robustness(float("nan")) == "not_estimable" + + @staticmethod + def _unestimable_but_valid_panel(): + # Data fit() ACCEPTS but cannot estimate: rolling='detrendq' with + # 4 pre-periods covering all 4 seasons -> every unit is seasonal- + # unidentified (warn + NaN ATT on every spec). + rng = np.random.default_rng(3) + records = [] + for i in range(20): + d = int(i < 8) + for t in range(1, 9): + records.append({"unit": i, "time": t, "y": rng.normal(), "treat": d * int(t >= 5)}) + return pd.DataFrame(records) + + def test_all_specs_fail_reports_not_estimable(self): + """Every spec unestimable (fit accepts the data but NaNs) -> the + result must be 'not_estimable' with a NaN ratio, not + 'highly_robust'.""" + df = self._unestimable_but_valid_panel() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r = robustness_pre_periods( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + rolling="detrendq", + ) + assert r.robustness_level == "not_estimable" + assert np.isnan(r.sensitivity_ratio) + + def test_data_fit_would_reject_raises_not_silently_swallowed(self): + # Fix-wave WS10 (campaign finding): genuine specification errors + # were swallowed as per-spec 'failed fits'. Data that LWDiD.fit() + # itself rejects (all-NaN outcome) must RAISE from the sensitivity + # helpers too. + records = [] + for i in range(20): + d = int(i < 8) + for t in range(1, 7): + records.append({"unit": i, "time": t, "y": np.nan, "treat": d * int(t > 3)}) + df = pd.DataFrame(records) + with pytest.raises(ValueError): + robustness_pre_periods(df, outcome="y", unit="unit", time="time", treatment="treat") + with pytest.raises(ValueError): + sensitivity_no_anticipation( + df, outcome="y", unit="unit", time="time", treatment="treat" + ) + + def test_all_specs_fail_no_anticipation_not_estimable(self): + df = self._unestimable_but_valid_panel() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + r = sensitivity_no_anticipation( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + rolling="detrendq", + ) + assert r.robustness_level == "not_estimable" + assert np.isnan(r.sensitivity_ratio) + + def test_missing_outcome_column_raises(self, panel_data): + with pytest.raises(ValueError, match="not found in data"): + robustness_pre_periods( + panel_data, + outcome="no_such_column", + unit="unit", + time="time", + treatment="treat", + ) + + def test_missing_control_column_raises(self, panel_data): + with pytest.raises(ValueError, match="not found in data"): + robustness_pre_periods( + panel_data, + outcome="y", + unit="unit", + time="time", + treatment="treat", + controls=["no_such_control"], + ) + + +class TestMultiCohortRejection: + """Round-2 finding: pre-period windows are earliest-adoption-relative, + so multi-cohort staggered inputs would mislabel later cohorts' + transformation samples. They fail closed; a single treated cohort is + exactly the global rule and stays supported.""" + + @staticmethod + def _staggered(n_cohorts=2): + rng = np.random.default_rng(0) + rows = [] + onsets = [5, 7][:n_cohorts] + for u in range(16): + if u < 4 * n_cohorts: + g = onsets[u % n_cohorts] + else: + g = 0 + for t in range(1, 10): + d = int(g > 0 and t >= g) + rows.append(dict(unit=u, time=t, treat=d, g=g, y=rng.normal() + d)) + return pd.DataFrame(rows) + + def test_multi_cohort_rejected_both_functions(self): + df = self._staggered(n_cohorts=2) + with pytest.raises(ValueError, match="single treated\\s+cohort"): + robustness_pre_periods( + df, outcome="y", unit="unit", time="time", treatment="treat", cohort="g" + ) + with pytest.raises(ValueError, match="single treated\\s+cohort"): + sensitivity_no_anticipation( + df, outcome="y", unit="unit", time="time", treatment="treat", cohort="g" + ) + + def test_single_cohort_still_supported(self): + df = self._staggered(n_cohorts=1) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = robustness_pre_periods( + df, outcome="y", unit="unit", time="time", treatment="treat", cohort="g" + ) + assert np.isfinite(res.baseline_att) + + +class TestParameterValidation: + """Round-3: strict validation of exclusion/window parameters + (exclude_periods=0 previously sliced pre_periods[:-0] == EMPTY, + silently dropping every pre-period).""" + + @staticmethod + def _panel(): + rng = np.random.default_rng(0) + rows = [] + for u in range(12): + for t in range(1, 9): + d = 1 if (u < 6 and t >= 6) else 0 + rows.append(dict(unit=u, time=t, treat=d, y=rng.normal() + d)) + return pd.DataFrame(rows) + + def test_exclude_periods_zero_rejected(self): + with pytest.raises(ValueError, match="positive integers"): + sensitivity_no_anticipation( + self._panel(), + outcome="y", + unit="unit", + time="time", + treatment="treat", + exclude_periods=[0], + ) + + def test_exclude_periods_duplicates_and_types_rejected(self): + with pytest.raises(ValueError, match="duplicate"): + sensitivity_no_anticipation( + self._panel(), + outcome="y", + unit="unit", + time="time", + treatment="treat", + exclude_periods=[1, 1], + ) + with pytest.raises(ValueError, match="positive integers"): + sensitivity_no_anticipation( + self._panel(), + outcome="y", + unit="unit", + time="time", + treatment="treat", + exclude_periods=[True], + ) + + def test_k_bounds_validated(self): + with pytest.raises(ValueError, match="k_min must be"): + robustness_pre_periods( + self._panel(), + outcome="y", + unit="unit", + time="time", + treatment="treat", + k_min=2.5, + ) diff --git a/tests/test_lwdid_visualization.py b/tests/test_lwdid_visualization.py new file mode 100644 index 00000000..0145b10e --- /dev/null +++ b/tests/test_lwdid_visualization.py @@ -0,0 +1,190 @@ +"""Tests for lwdid_visualization module.""" + +from unittest.mock import patch + +import numpy as np +import pandas as pd +import pytest + +from diff_diff.lwdid_visualization import ( + _require_matplotlib, + plot_bootstrap_distribution, + plot_cohort_trends, + plot_event_study, + plot_sensitivity, +) + +# --------------------------------------------------------------------------- +# Importability +# --------------------------------------------------------------------------- + + +class TestImportability: + """Test that all visualization functions are importable.""" + + def test_plot_cohort_trends_importable(self): + assert callable(plot_cohort_trends) + + def test_plot_event_study_importable(self): + assert callable(plot_event_study) + + def test_plot_sensitivity_importable(self): + assert callable(plot_sensitivity) + + def test_plot_bootstrap_distribution_importable(self): + assert callable(plot_bootstrap_distribution) + + def test_require_matplotlib_importable(self): + assert callable(_require_matplotlib) + + +# --------------------------------------------------------------------------- +# _require_matplotlib error handling +# --------------------------------------------------------------------------- + + +class TestRequireMatplotlib: + """Test _require_matplotlib raises proper error if no matplotlib.""" + + def test_raises_visualization_error_when_no_matplotlib(self): + """Mock ImportError to simulate missing matplotlib.""" + import builtins + + real_import = builtins.__import__ + + def mock_import(name, *args, **kwargs): + if name == "matplotlib.pyplot" or name == "matplotlib": + raise ImportError("No module named 'matplotlib'") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=mock_import): + with pytest.raises(ImportError, match="matplotlib"): + _require_matplotlib() + + +# --------------------------------------------------------------------------- +# Plot functions return Figure when matplotlib available +# --------------------------------------------------------------------------- + + +class TestPlotFunctions: + """Test plot functions return Figure when matplotlib is available.""" + + @pytest.fixture + def panel_data(self): + rng = np.random.default_rng(42) + records = [] + for i in range(80): + d = int(i < 25) + for t in range(1, 9): + y = 1.0 + 0.1 * t + rng.normal(0, 0.3) + if d and t > 4: + y += 2.0 + records.append({"unit": i, "time": t, "y": y, "treat": d * int(t > 4)}) + return pd.DataFrame(records) + + def test_plot_cohort_trends_returns_figure(self, panel_data): + pytest.importorskip("matplotlib") + import matplotlib.pyplot as plt + + fig = plot_cohort_trends( + panel_data, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert fig is not None + assert hasattr(fig, "savefig") # duck-type check for Figure + plt.close(fig) + + def test_plot_event_study_returns_figure(self, panel_data): + pytest.importorskip("matplotlib") + import matplotlib.pyplot as plt + + from diff_diff import LWDiD + + # Staggered fit populates the event-study surface + staggered = panel_data.copy() + staggered["first_treat"] = np.where(staggered["unit"] < 25, 5, 0) + res = LWDiD(rolling="demean").fit( + staggered, + outcome="y", + unit="unit", + time="time", + treatment="treat", + first_treat="first_treat", + ) + assert res.event_study_effects + + fig = plot_event_study(res) + assert fig is not None + assert hasattr(fig, "savefig") + plt.close(fig) + + def test_plot_event_study_common_timing_returns_figure(self, panel_data): + """Common-timing fits now populate the per-period event-study + surface at fit time, so plot_event_study works on them too.""" + pytest.importorskip("matplotlib") + import matplotlib.pyplot as plt + + from diff_diff import LWDiD + + res = LWDiD(rolling="demean").fit( + panel_data, outcome="y", unit="unit", time="time", treatment="treat" + ) + assert res.event_study_effects + + fig = plot_event_study(res) + assert fig is not None + assert hasattr(fig, "savefig") + plt.close(fig) + + def test_plot_bootstrap_distribution_returns_figure(self): + pytest.importorskip("matplotlib") + import matplotlib.pyplot as plt + + t_stats = np.random.default_rng(0).normal(0, 1, 500) + fig = plot_bootstrap_distribution(t_stats, t_observed=2.5) + assert fig is not None + assert hasattr(fig, "savefig") + plt.close(fig) + + +class TestPlottingConventions: + """Fix-wave WS10 pins: NaN-SE effects plot the point and OMIT the + interval (never a zero-length bar); cohort-trend plots accept datetime + time columns (the onset marker no longer computes 'Timestamp - 0.5'). + """ + + def test_event_study_nan_se_omits_interval(self): + matplotlib = pytest.importorskip("matplotlib") + + matplotlib.use("Agg") + from types import SimpleNamespace + + from diff_diff.lwdid_visualization import plot_event_study + + results = SimpleNamespace( + event_study_effects={ + -2: {"effect": 0.1, "se": 0.05}, + 0: {"effect": 1.0, "se": float("nan")}, # inference unavailable + 1: {"effect": 1.2, "se": 0.07}, + }, + reference_periods=(-1,), + ) + fig = plot_event_study(results) + assert fig is not None + + def test_cohort_trends_accepts_datetime_time(self): + matplotlib = pytest.importorskip("matplotlib") + + matplotlib.use("Agg") + from diff_diff.lwdid_visualization import plot_cohort_trends + + rng = np.random.default_rng(3) + rows = [] + times = pd.date_range("2020-01-01", periods=6, freq="MS") + for u in range(8): + for i, t in enumerate(times): + d = int(u < 4 and i >= 3) + rows.append(dict(unit=u, time=t, treat=d, y=rng.normal() + d, first=0)) + df = pd.DataFrame(rows) + fig = plot_cohort_trends(df, outcome="y", unit="unit", time="time", treatment="treat") + assert fig is not None diff --git a/tests/test_lwdid_wild_bootstrap.py b/tests/test_lwdid_wild_bootstrap.py new file mode 100644 index 00000000..e6f0ff4c --- /dev/null +++ b/tests/test_lwdid_wild_bootstrap.py @@ -0,0 +1,300 @@ +"""Tests for lwdid_wild_bootstrap module (house-engine wrapper API). + +Rewritten in the LWDiD fix wave (WS4): wild_cluster_bootstrap now delegates +to the house WCR engine ``diff_diff.utils.wild_bootstrap_se``. The former +module-local implementation carried three execution-verified defects +(1-ULP tie handling below the attainable p floor, an intercept-only +restricted model that dropped controls from the null DGP, and a G=2 +zero-SE roundoff escape reporting t~5e15 with p=0.25). +""" + +import warnings + +import numpy as np +import pandas as pd +import pytest + +from diff_diff.lwdid_wild_bootstrap import ( + WildClusterBootstrapResult, + wild_cluster_bootstrap, +) +from diff_diff.utils import wild_bootstrap_se + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def cross_section_data(): + rng = np.random.default_rng(42) + n = 100 + y = np.concatenate([rng.normal(2, 0.5, 30), rng.normal(0, 0.5, 70)]) + treatment = np.array([1.0] * 30 + [0.0] * 70) + cluster_ids = np.repeat(np.arange(20), 5) + controls = rng.normal(0, 1, (n, 2)) + return y, treatment, cluster_ids, controls + + +# --------------------------------------------------------------------------- +# Result schema (house-aligned) +# --------------------------------------------------------------------------- + + +class TestResultSchema: + def test_schema_fields(self, cross_section_data): + y, d, cl, _ = cross_section_data + r = wild_cluster_bootstrap(y, d, cl, seed=1, n_bootstrap=99) + assert isinstance(r, WildClusterBootstrapResult) + assert np.isfinite(r.att) + assert np.isfinite(r.se) and r.se > 0 + assert np.isfinite(r.t_stat_original) + assert 0.0 <= r.p_value <= 1.0 + assert r.ci_lower <= r.ci_upper + assert r.n_clusters == 20 + assert r.n_bootstrap >= 99 + assert r.weight_type == "rademacher" + assert r.alpha == 0.05 + assert r.bootstrap_distribution is not None + assert len(r.bootstrap_distribution) <= r.n_bootstrap + assert r.n_dropped == 0 + # Retired fields are gone (API break, LWDiD unreleased) + assert not hasattr(r, "se_bootstrap") + assert not hasattr(r, "pvalue") + assert not hasattr(r, "n_reps") + assert not hasattr(r, "t_stats") + + def test_summary_returns_string(self, cross_section_data): + y, d, cl, _ = cross_section_data + r = wild_cluster_bootstrap(y, d, cl, seed=1, n_bootstrap=99) + s = r.summary() + assert isinstance(s, str) + assert "Wild Cluster Bootstrap" in s + assert "CR1" in s + + def test_matches_house_engine_exactly(self, cross_section_data): + # The wrapper is a thin adapter: same X construction, same engine, + # same numbers as calling wild_bootstrap_se directly. + y, d, cl, controls = cross_section_data + r = wild_cluster_bootstrap(y, d, cl, controls, seed=7, n_bootstrap=199) + X = np.column_stack([np.ones(len(y)), d, controls]) + beta, *_ = np.linalg.lstsq(X, y, rcond=None) + house = wild_bootstrap_se( + X, y, y - X @ beta, cl, 1, n_bootstrap=199, seed=7, return_distribution=True + ) + np.testing.assert_allclose(r.se, house.se, rtol=0, atol=0) + np.testing.assert_allclose(r.p_value, house.p_value, rtol=0, atol=0) + np.testing.assert_allclose(r.ci_lower, house.ci_lower, rtol=0, atol=0) + np.testing.assert_allclose(r.ci_upper, house.ci_upper, rtol=0, atol=0) + np.testing.assert_allclose(r.att, beta[1], rtol=1e-12) + + +class TestWeightTypes: + @pytest.mark.parametrize("wt", ["rademacher", "mammen", "webb"]) + def test_weight_types_run(self, cross_section_data, wt): + y, d, cl, _ = cross_section_data + r = wild_cluster_bootstrap(y, d, cl, weight_type=wt, seed=3, n_bootstrap=99) + assert 0.0 <= r.p_value <= 1.0 + assert r.weight_type == wt + + def test_invalid_weight_type_raises(self, cross_section_data): + y, d, cl, _ = cross_section_data + with pytest.raises(ValueError, match="weight_type"): + wild_cluster_bootstrap(y, d, cl, weight_type="gaussian") + + +class TestStatisticalProperties: + def test_null_imposition_keeps_controls(self): + # Campaign finding: the old restricted model was intercept-only, + # dumping covariate signal into the bootstrap residuals (Monte + # Carlo size 12.5% vs nominal 5% with a treatment-correlated + # control). The house engine drops ONLY the treatment column; on a + # null DGP with a strong treatment-correlated control the test must + # not over-reject. + rng = np.random.default_rng(0) + rejections = 0 + n_sims = 40 + for _ in range(n_sims): + G = 12 + cl = np.repeat(np.arange(G), 10) + d = (cl < 4).astype(float) + x = 2.0 * d + rng.normal(size=cl.size) + y = 1.0 + 1.5 * x + rng.normal(size=cl.size) # no treatment effect + r = wild_cluster_bootstrap( + y, d, cl, controls=x.reshape(-1, 1), n_bootstrap=199, seed=int(rng.integers(1e6)) + ) + rejections += int(r.p_value < 0.05) + # Binomial(40, 0.05): P(X >= 9) < 1e-4 + assert rejections <= 8, rejections + + def test_full_enumeration_deterministic(self): + rng = np.random.default_rng(5) + G = 8 + cl = np.repeat(np.arange(G), 6) + d = (cl < 3).astype(float) + y = 0.5 * d + rng.normal(size=cl.size) + r1 = wild_cluster_bootstrap(y, d, cl, n_bootstrap=999, seed=1) + r2 = wild_cluster_bootstrap(y, d, cl, n_bootstrap=999, seed=2) + # 2**8 = 256 <= 999 -> full enumeration, independent of the seed + assert r1.n_bootstrap == 256 and r2.n_bootstrap == 256 + assert r1.p_value == r2.p_value + + def test_enumeration_p_is_exact_atom(self): + # The campaign's 1-ULP tie finding (reported p below the attainable + # floor of the OLD percentile-t enumeration) is resolved by + # ADOPTION of the house WCR convention, whose tie handling is + # pinned by the house R-parity goldens (tests/test_wild_bootstrap). + # Contract here: under enumeration the p-value is an exact atom + # k/2**G of the deterministic distribution. + rng = np.random.default_rng(9) + G = 4 + cl = np.repeat(np.arange(G), 8) + d = (cl < 2).astype(float) + y = 3.0 * d + rng.normal(scale=0.2, size=cl.size) + with pytest.warns(UserWarning, match="fewer than 5 clusters"): + r = wild_cluster_bootstrap(y, d, cl, n_bootstrap=999, seed=11) + assert r.n_bootstrap == 16 + k = r.p_value * 16 + np.testing.assert_allclose(k, round(k), atol=1e-12) + + +class TestDegenerateDesigns: + def test_g2_exactly_identified_fails_closed(self): + # Campaign finding: the canonical two-cluster design (cluster- + # invariant treatment) has cluster scores exactly ~0; BLAS roundoff + # gave a tiny-positive SE, t ~ 5e15, and p = 0.25 (below the G=2 + # attainable floor of 0.5). Point retained; inference NaN. + rng = np.random.default_rng(2) + cl = np.repeat([0, 1], 12) + d = (cl == 0).astype(float) + y = 1.0 + 0.8 * d + rng.normal(scale=0.5, size=cl.size) + with pytest.warns(UserWarning, match="not identified"): + r = wild_cluster_bootstrap(y, d, cl, n_bootstrap=99, seed=4) + assert np.isfinite(r.att) + assert np.isnan(r.se) and np.isnan(r.p_value) + assert np.isnan(r.ci_lower) and np.isnan(r.ci_upper) + assert r.bootstrap_distribution is None + + def test_single_cluster_rejected(self): + rng = np.random.default_rng(3) + y = rng.normal(size=20) + d = np.r_[np.ones(10), np.zeros(10)] + cl = np.zeros(20) + with pytest.raises(ValueError, match="at least 2 clusters"): + wild_cluster_bootstrap(y, d, cl) + + +class TestInputContracts: + def test_nonfinite_y_dropped_with_warning_and_counted(self): + rng = np.random.default_rng(6) + G = 10 + cl = np.repeat(np.arange(G), 8) + d = (cl < 4).astype(float) + y = 1.0 * d + rng.normal(size=cl.size) + y[3] = np.nan + y[40] = np.inf + with pytest.warns(UserWarning, match="dropped 2 observation"): + r = wild_cluster_bootstrap(y, d, cl, n_bootstrap=99, seed=8) + assert r.n_dropped == 2 + assert np.isfinite(r.p_value) + + def test_nonfinite_controls_raise(self, cross_section_data): + y, d, cl, controls = cross_section_data + controls = controls.copy() + controls[0, 0] = np.nan + with pytest.raises(ValueError, match="controls contains non-finite"): + wild_cluster_bootstrap(y, d, cl, controls) + + def test_retired_parameters_rejected(self, cross_section_data): + y, d, cl, _ = cross_section_data + with pytest.raises(TypeError): + wild_cluster_bootstrap(y, d, cl, impose_null=False) + with pytest.raises(TypeError): + wild_cluster_bootstrap(y, d, cl, full_enumeration=True) + with pytest.raises(TypeError): + wild_cluster_bootstrap(y, d, cl, n_reps=99) + with pytest.raises(TypeError): + wild_cluster_bootstrap(y, d, cl, ci_level=0.9) + + +class TestResultsConvenienceMethods: + """LWDiDResults.wild_cluster_bootstrap() / .randomization_test(). + + Round-5 review: these REPLAY the fitted estimation sample and RA + design (no data arguments) and assert their observed statistic equals + ``.att`` before caching - previously they accepted arbitrary caller + arrays and a non-interacted design, so the cached p-values could + describe a different estimand than the fitted ATT. + """ + + @staticmethod + def _fitted_results(cluster=None, covariate=False): + from diff_diff import LWDiD + + rng = np.random.default_rng(42) + records = [] + for i in range(60): + d = int(i < 20) + x = float(i % 4) + (1.5 if d else 0.0) # treatment-unbalanced + for t in range(1, 7): + y = 1.0 + 0.1 * t + 0.4 * x + rng.normal(0, 0.3) + if d and t > 3: + y += 2.0 + 0.5 * x + records.append( + {"unit": i, "time": t, "y": y, "treat": d * int(t > 3), "x": x, "cl": i % 12} + ) + df = pd.DataFrame(records) + est = LWDiD(cluster=cluster) + return est.fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x"] if covariate else None, + ) + + def test_results_wild_cluster_bootstrap_replays_fit(self): + res = self._fitted_results(cluster="cl", covariate=True) + wcb = res.wild_cluster_bootstrap(n_bootstrap=99, seed=42) + # coherence: the replayed observed ATT IS the fitted ATT + np.testing.assert_allclose(wcb.att, res.att, rtol=1e-10) + assert 0 <= wcb.p_value <= 1 + assert res.bootstrap_pvalue == wcb.p_value + + def test_results_wcb_requires_clustered_fit(self): + res = self._fitted_results(cluster=None) + with pytest.raises(ValueError, match="requires a clustered fit"): + res.wild_cluster_bootstrap(n_bootstrap=99, seed=42) + + def test_results_randomization_test_replays_fit(self): + res = self._fitted_results(covariate=True) + ri = res.randomization_test(n_reps=199, seed=42) + np.testing.assert_allclose(ri.att_observed, res.att, rtol=1e-10) + assert 0 <= ri.pvalue <= 1 + assert res.ri_pvalue == ri.pvalue + + def test_results_methods_reject_non_reg_fits(self): + from diff_diff import LWDiD + + rng = np.random.default_rng(0) + records = [] + for i in range(40): + d = int(i < 20) + x = float(i % 5) + for t in range(1, 7): + y = 1.0 + 0.2 * x + rng.normal(0, 0.3) + (2.0 if d and t > 3 else 0.0) + records.append({"unit": i, "time": t, "y": y, "treat": d * int(t > 3), "x": x}) + df = pd.DataFrame(records) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD(estimation_method="ipw").fit( + df, + outcome="y", + unit="unit", + time="time", + treatment="treat", + covariates=["x"], + ) + with pytest.raises(ValueError, match="only\\s+defined for estimation_method='reg'"): + res.randomization_test(n_reps=99, seed=0) diff --git a/tests/test_methodology_lwdid.py b/tests/test_methodology_lwdid.py index 5116738f..8bdc7bf3 100644 --- a/tests/test_methodology_lwdid.py +++ b/tests/test_methodology_lwdid.py @@ -99,13 +99,12 @@ reason="LWDiD estimator not yet on main (arrives via PR #588)", ) -from diff_diff.lwdid import LWDiD # noqa: E402 - from diff_diff import ( # noqa: E402 DifferenceInDifferences, # noqa: E402 load_prop99, load_walmart, ) +from diff_diff.lwdid import LWDiD # noqa: E402 # --------------------------------------------------------------------------- # Published replication targets (LW 2026; see module docstring for provenance) @@ -1714,3 +1713,350 @@ def test_detrending_with_two_pre_periods_works(self): df, outcome="y", unit="unit", time="time", treatment="treat" ) assert np.isfinite(res.att) + + +# --------------------------------------------------------------------------- +# Fix-wave WS1: complete-case fixed-weight tau_omega (campaign findings: +# 0.0-injection for controls missing a cohort's post window; silent +# treated-side reweighting through the finite mask; composite gate applying +# the NON-seasonal transform to demeanq/detrendq overall ATTs) +# --------------------------------------------------------------------------- + + +def _synthetic_drops_staggered(seed=42): + """Deterministic unbalanced panel with GENUINE complete-case drops. + + Cohorts {3, 5} over t=1..6. One control unit observes only t=1..4 + (missing cohort-5's entire post window -> dropped from the composite + control side) and one cohort-5 treated unit observes only t=1..4 + (missing its OWN post window -> dropped from the treated side, with the + cohort masses recomputed on the survivors). + """ + rng = np.random.default_rng(seed) + rows = [] + uid = 0 + spec = [(0, 40, None), (0, 1, (1, 2, 3, 4)), (3, 20, None), (5, 14, None), (5, 1, (1, 2, 3, 4))] + for g, n, keep in spec: + for _ in range(n): + alpha = rng.normal() + for t in range(1, 7): + if keep is not None and t not in keep: + continue + d = int(g > 0 and t >= g) + y = alpha + 0.2 * t + rng.normal(scale=0.4) + (1.5 + 0.4 * (g == 5)) * d + rows.append(dict(unit=uid, time=t, first=g, treat=d, y=y)) + uid += 1 + return pd.DataFrame(rows) + + +def _complete_case_tau_omega_reference(df): + """From-scratch complete-case fixed-weight tau_omega (demean). + + Fixed cohort weights omega_g = N_g / N_treat defined on the ESTIMATION + sample: treated units without a finite own-cohort post average are + dropped and the masses recomputed; control units must observe every + surviving-weight cohort's post window. + """ + fy = df.groupby("unit")["first"].first() + cohorts = sorted(set(fy[fy > 0])) + ydot = {} + for g in cohorts: + pre_mean = df.loc[df["time"] < g].groupby("unit")["y"].mean() + post = df.loc[df["time"] >= g].copy() + post["_ydot"] = post["y"] - post["unit"].map(pre_mean) + ydot[g] = post.groupby("unit")["_ydot"].mean() + surviving_treated = [ + u for u in fy.index if fy[u] > 0 and np.isfinite(ydot[fy[u]].get(u, np.nan)) + ] + fy_cc = fy.loc[surviving_treated] + sizes = {g: int((fy_cc == g).sum()) for g in cohorts} + weighted = [g for g in cohorts if sizes[g] > 0] + n_treat_cc = len(surviving_treated) + controls = [ + u + for u in fy.index + if not fy[u] > 0 and all(np.isfinite(ydot[g].get(u, np.nan)) for g in weighted) + ] + y, d = [], [] + for u in surviving_treated: + y.append(float(ydot[fy[u]][u])) + d.append(1.0) + for u in controls: + y.append(sum(sizes[g] / n_treat_cc * float(ydot[g][u]) for g in weighted)) + d.append(0.0) + y_arr, d_arr = np.asarray(y), np.asarray(d) + return float(y_arr[d_arr == 1].mean() - y_arr[d_arr == 0].mean()) + + +class TestTauOmegaCompleteCase: + """WS1: fixed-weight complete-case tau_omega with vcov-invariant routing.""" + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat", first_treat="first") + + def _fit(self, df, **overrides): + params = dict( + rolling="demean", + estimation_method="reg", + vcov_type="classical", + control_group="never_treated", + ) + params.update(overrides) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return LWDiD(**params).fit(df, **self.KW) + + def test_two_sided_complete_case_oracle(self): + # Independent arithmetic oracle exercising BOTH a dropped treated + # unit AND a dropped control unit (the route-consistency test + # cannot see a wrong weighting: both routes share the composite). + df = _synthetic_drops_staggered() + res = self._fit(df) + assert res.n_composite_treated_dropped == 1 + assert res.n_composite_controls_dropped == 1 + expected = _complete_case_tau_omega_reference(df) + np.testing.assert_allclose(res.att_tau_omega_complete_case, expected, atol=1e-10, rtol=0) + + def test_drops_route_reports_if_weighted_point_with_warning(self): + df = _synthetic_drops_staggered() + with pytest.warns(UserWarning, match="complete-case"): + res = LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="classical", + control_group="never_treated", + ).fit(df, **self.KW) + # .att equals the cohort-mass IF-weighted point on the drops route + expected = sum(v["weight"] * v["att"] for v in res.cohort_effects.values()) + np.testing.assert_allclose(res.att, expected, atol=1e-10, rtol=0) + assert res.inference_basis == "joint_influence_function" + assert np.isfinite(res.se) and res.se > 0 + + def test_vcov_invariance_with_drops(self): + # Pre-fix hazard class: a variance selection moving the point. + df = _synthetic_drops_staggered() + atts = [self._fit(df, vcov_type=v).att for v in ("classical", "hc1")] + np.testing.assert_allclose(atts[0], atts[1], atol=1e-10, rtol=0) + + def test_zero_drop_panel_keeps_status_quo_metadata(self): + df = _synthetic_unbalanced_staggered() # unbalanced but zero drops + res = self._fit(df) + assert res.n_composite_treated_dropped == 0 + assert res.n_composite_controls_dropped == 0 + assert res.att_tau_omega_complete_case is None + assert res.inference_basis == "composite_regression" + + def test_classical_if_se_concords_with_unit_bootstrap(self, ci_params): + # With drops, classical pairs the IF point with the IF SE; pin the + # pair against an external unit-resampling bootstrap. + df = _synthetic_drops_staggered() + res = self._fit(df) + rng = np.random.default_rng(42) + units = df["unit"].unique() + n_boot = ci_params.bootstrap(300, min_n=199) + draws = [] + for _ in range(n_boot): + picks = rng.choice(units, size=len(units), replace=True) + frames = [] + for j, u in enumerate(picks): + block = df.loc[df["unit"] == u].copy() + block["unit"] = j + frames.append(block) + bs = pd.concat(frames, ignore_index=True) + try: + draws.append(self._fit(bs).att) + except ValueError: + continue + boot_se = float(np.std([v for v in draws if np.isfinite(v)], ddof=1)) + threshold = 0.40 if n_boot < 100 else 0.15 + assert abs(res.se - boot_se) / boot_se < threshold, (res.se, boot_se) + + +class TestSeasonalOverallRouting: + """WS1 gate coherence: demeanq/detrendq never enter the tau_omega + composite (which is defined for the plain transforms only). Pre-fix, + vcov_type='classical' silently swapped the NON-seasonal transform into + the staggered q-mode overall ATT (~8% shift on seasonal DGPs) while + hc1 aggregated seasonal cohort ATTs - a vcov selection moved the point. + """ + + KW = dict(outcome="y", unit="unit", time="time", treatment="treat", first_treat="first") + + @staticmethod + def _quarterly_panel(seed=3): + rng = np.random.default_rng(seed) + rows = [] + uid = 0 + season = np.array([1.2, -0.6, 0.9, -1.5]) + for g, n in [(0, 30), (9, 12), (11, 10)]: + for _ in range(n): + alpha = rng.normal() + amp = rng.uniform(0.5, 1.5) + for t in range(1, 17): + d = int(g > 0 and t >= g) + y = ( + alpha + + amp * season[(t - 1) % 4] + + 0.1 * t + + rng.normal(scale=0.3) + + 1.5 * d + ) + rows.append(dict(unit=uid, time=t, first=g, treat=d, y=y)) + uid += 1 + return pd.DataFrame(rows) + + def _fit(self, df, **overrides): + params = dict( + rolling="demeanq", + estimation_method="reg", + vcov_type="classical", + control_group="never_treated", + ) + params.update(overrides) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return LWDiD(**params).fit(df, **self.KW) + + def test_qmode_overall_is_seasonal_cohort_mass_average(self): + df = self._quarterly_panel() + res = self._fit(df) + expected = sum(v["weight"] * v["att"] for v in res.cohort_effects.values()) + np.testing.assert_allclose(res.att, expected, atol=1e-10, rtol=0) + assert res.inference_basis == "joint_influence_function" + + def test_qmode_vcov_never_moves_the_point(self): + df = self._quarterly_panel() + atts = [self._fit(df, vcov_type=v).att for v in ("classical", "hc1")] + np.testing.assert_allclose(atts[0], atts[1], atol=1e-12, rtol=0) + + def test_qmode_differs_from_plain_demean_on_seasonal_dgp(self): + # The seasonal adjustment must actually matter on this DGP - + # guards against the fix regressing into a silent transform swap. + df = self._quarterly_panel() + att_q = self._fit(df).att + att_plain = self._fit(df, rolling="demean").att + assert abs(att_q - att_plain) > 1e-6 + + def test_composite_raises_if_q_variant_reaches_it(self): + df = self._quarterly_panel() + est = LWDiD(rolling="demeanq", estimation_method="reg") + with pytest.raises(ValueError, match="only defined for rolling"): + est._composite_regression_aggregation(df, "y", "unit", "time", "first") + + +# --------------------------------------------------------------------------- +# Fix-wave WS6: vcov contract + per-surface reference-distribution policy +# --------------------------------------------------------------------------- + + +class TestInferenceDispatchPolicy: + """Campaign finding: identical single-cohort designs produced p-values + differing by ~34 orders of magnitude depending on whether first_treat + was passed (t vs normal dispatch). Policy now: an aggregate composed of + EXACTLY ONE cell uses that cell's residual df (the common-timing rule); + multi-cell aggregates keep the large-sample reference; clustered + aggregates use G-1 over CONTRIBUTING clusters. + """ + + @staticmethod + def _single_post_panel(seed=17): + # T_post = 1: the staggered fit has exactly one estimable cell. + rng = np.random.default_rng(seed) + rows = [] + for u in range(24): + alpha = rng.normal() + treated = u < 10 + for t in range(1, 6): + d = int(treated and t >= 5) + y = alpha + 0.1 * t + rng.normal(scale=0.4) + 1.4 * d + rows.append(dict(unit=u, time=t, first=5 if treated else 0, treat=d, y=y)) + return pd.DataFrame(rows) + + def test_single_post_period_staggered_matches_common_timing(self): + df = self._single_post_panel() + kw_common = dict(outcome="y", unit="unit", time="time", treatment="treat") + est = dict(rolling="demean", estimation_method="reg", vcov_type="hc1") + rc = LWDiD(**est).fit(df, **kw_common) + rs = LWDiD(**est).fit(df, first_treat="first", **kw_common) + np.testing.assert_allclose(rs.att, rc.att, rtol=1e-10) + np.testing.assert_allclose(rs.se, rc.se, rtol=1e-10) + assert rs.df_inference == rc.df_inference # same residual t reference + np.testing.assert_allclose(rs.p_value, rc.p_value, rtol=1e-8) + + def test_cell_with_single_cluster_fails_closed_and_propagates(self): + # Cluster ids are re-derived per cell; a cell whose units share one + # cluster must NaN its inference (point retained) and any aggregate + # including it inherits NaN inference - deliberate fail-closed. + rng = np.random.default_rng(23) + rows = [] + uid = 0 + # cohort-4 treated units all in cluster 0; controls span clusters + for g, n, cl_fn in [(0, 12, lambda u: 1 + (u % 4)), (4, 6, lambda u: 0)]: + for _ in range(n): + alpha = rng.normal() + for t in range(1, 7): + d = int(g > 0 and t >= g) + y = alpha + rng.normal(scale=0.4) + 1.2 * d + rows.append(dict(unit=uid, time=t, first=g, treat=d, y=y, cl=cl_fn(uid))) + uid += 1 + df = pd.DataFrame(rows) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + res = LWDiD( + rolling="demean", + estimation_method="reg", + cluster="cl", + control_group="never_treated", + ).fit(df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="first") + # Cells contain treated (cluster 0) + controls (clusters 1-4): G=5 + # per cell, so this design is NOT degenerate; flip to a truly + # degenerate one below. + assert np.isfinite(res.att) + del caught + + def test_degenerate_single_cluster_cells_nan_inference(self): + rng = np.random.default_rng(29) + rows = [] + uid = 0 + # EVERY unit in one cluster: per-cell G=1 while a second, empty + # cluster never contributes -> global guard passes via a control + # unit parked alone in cluster 1 with NaN-free data but the cells + # all draw from cluster 0. + for g, n in [(0, 10), (4, 5)]: + for _ in range(n): + alpha = rng.normal() + for t in range(1, 7): + d = int(g > 0 and t >= g) + y = alpha + rng.normal(scale=0.4) + 1.2 * d + rows.append(dict(unit=uid, time=t, first=g, treat=d, y=y, cl=0)) + uid += 1 + # one extra never-treated unit in its own cluster, excluded from + # cells by control_group='never_treated'? No - it IS a control. + # Give it data only in pre-periods so it drops from post cells. + for t in range(1, 4): + rows.append(dict(unit=uid, time=t, first=0, treat=0, y=rng.normal(), cl=1)) + df = pd.DataFrame(rows) + with pytest.warns(UserWarning): + res = LWDiD( + rolling="demean", + estimation_method="reg", + cluster="cl", + control_group="never_treated", + ).fit(df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="first") + assert np.isfinite(res.att) # point retained + assert np.isnan(res.se) # fail-closed propagation + assert np.isnan(res.p_value) + + def test_qmode_multicell_aggregate_keeps_normal_reference(self): + # Multi-cell unclustered aggregates: large-sample reference + # (df_inference is None), documented - not a pooled residual df. + df = TestSeasonalOverallRouting._quarterly_panel() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + res = LWDiD( + rolling="demean", + estimation_method="reg", + vcov_type="hc1", + control_group="never_treated", + ).fit(df, outcome="y", unit="unit", time="time", treatment="treat", first_treat="first") + assert res.inference_basis == "joint_influence_function" + assert res.df_inference is None diff --git a/tests/test_naming_guard.py b/tests/test_naming_guard.py index bc41306b..174dea26 100644 --- a/tests/test_naming_guard.py +++ b/tests/test_naming_guard.py @@ -488,6 +488,8 @@ def _build_rowed_index(): "HeterogeneousAdoptionDiD.fit[time]", "ImputationDiD.fit[time]", "LPDiD.fit[time]", + "LWDiD.fit[time]", + "LWDiD.get_transformation_diagnostics[time]", "SpilloverDiD.fit[time]", "StackedDiD.fit[time]", "SunAbraham.fit[time]", @@ -1068,6 +1070,10 @@ def _token_family_code_refs(tok): ("time", "docs/methodology/papers/wooldridge-2023-review.md"): ( "canonical calendar column prose in the shipped-API description, not the M-030 overload" ), + ("time", "diff_diff/lwdid_sensitivity.py"): ( + "internal refits pass the canonical calendar column through to " + "LWDiD.fit[time] (rule-1), not the M-030 overload" + ), ("cohort", "docs/methodology/papers/borusyak-jaravel-spiess-2024-review.md"): ( "ImputationDiD partition-value prose, not the Wooldridge fit[cohort] kwarg" ), diff --git a/tests/test_spillover.py b/tests/test_spillover.py index f10334b4..f1d724f5 100644 --- a/tests/test_spillover.py +++ b/tests/test_spillover.py @@ -2280,16 +2280,20 @@ def test_rings_starting_above_zero_raises(self): class TestSpilloverDiDHC2NotSupported: """vcov_type='hc2' and 'hc2_bm' require per-coefficient BM/CR2 DOF - that the inline stage-2 inference doesn't provide. Round-8 codex - review caught that we'd silently return wrong p-values/CIs. + that the inline stage-2 inference doesn't provide; hc3 is not + implemented for the two-stage spillover variance at all. All three + now fail closed at CONSTRUCTION (LWDiD fix-wave hc3 hardening) so a + never-supported family cannot silently reach fit-time state. """ @pytest.mark.parametrize("vcov_type", ["hc2", "hc2_bm"]) def test_hc2_paths_raise_not_implemented(self, vcov_type): - df = _make_butts_2period_dgp(seed=42) - est = SpilloverDiD(rings=[0.0, 100.0], conley_coords=("lat", "lon"), vcov_type=vcov_type) with pytest.raises(NotImplementedError, match="hc2"): - est.fit(df, outcome="y", unit="unit", time="time", treatment="D") + SpilloverDiD(rings=[0.0, 100.0], conley_coords=("lat", "lon"), vcov_type=vcov_type) + + def test_hc3_raises_not_implemented_with_own_reason(self): + with pytest.raises(NotImplementedError, match="hc3.*two-stage spillover"): + SpilloverDiD(rings=[0.0, 100.0], conley_coords=("lat", "lon"), vcov_type="hc3") class TestSpilloverDiDRankDeficientActionValidation: