From d9cba7f90763c7ca1a7407ea0085b6b80daacf89 Mon Sep 17 00:00:00 2001 From: igerber Date: Sat, 24 Oct 2026 12:22:10 -0400 Subject: [PATCH 1/2] refactor(v4): PR 4c-i of the 4.0 cut - decouple the TWFE event-study mode from the MultiPeriodDiD containers and migrate the MPD-fixture suites Cut-sequence item 4c-i (docs/v4-design.md section 9; item 4c split into 4c-i/4c-ii, user decision 2026-10-24). The TWFE event-study mode no longer constructs a MultiPeriodDiDResults: DifferenceInDifferences._fit_event_study_core returns a private _EventStudyCorePayload, TwoWayFixedEffects.fit(..., event_study=True) builds its EventStudyResults from it via results_base._from_event_study_payload, and MultiPeriodDiD.fit wraps the same payload into its container (estimators._mpd_results_from_payload). _from_mpd and the payload route share one calendar-surface builder (_calendar_surface_from_arrays, sorted-label contract enforced with an explicit raise), so the two paths cannot drift; _resolve_scalar_df_survey's metadata half is split out as _scalar_df_survey_from_metadata. No user-visible behavior change; mypy zero errors. Oracle (captured BEFORE the refactor): tests/_capture_v4_merge_mpd_oracles.py ran on the UNMODIFIED tree (68e383bd, python backend; the rust capture agreed at rtol 1e-12 / atol 1e-13) and wrote tests/fixtures/v4_twfe_event_study_oracles.json - every EventStudyResults field of the TWFE surface AND of the MultiPeriodDiD twin across eleven lanes (pooled default / explicit cluster / hc2_bm; within + cluster / hc2 / Conley / survey TSL / JK1 replicates; string labels; a non-suffix partition; a rank-deficient NaN post period). Lanes live in tests/helpers/twfe_es_oracle_lanes.py (shared by the script and the test; the fixture echoes each lane so an edited lane fails as lane drift). tests/test_twfe_event_study_decoupling.py pins the refactored surface against that capture (CI bounds at the scipy-drift tolerance 1e-9), the MPD twin against its capture, payload-vs-_from_mpd byte identity, the full-container TWFE-vs-MPD identity (exact pooled / 1e-12 within), the sorted-zip contract on an unsorted hand-built dict, the df_survey resolver split, and the post_period_average helper's MPD arithmetic incl. its NaN contract. Suite migration (ported = replaced; MPD-only tests stay for 4c-ii): honest_did, pretrends (+ a TWFE pre_periods= class; the MPD class is MPD-only because it promotes declared-post labels), methodology pretrends/honest_did (hand-built calendar surfaces, reference inserted so no coefficient is demoted), visualization (the native-MPD plot test stays as the period_effects-branch pin), plotly, Conley (unit= kept: Conley requires it and never auto-clusters), survey, replicate weights, survey_phase6, vcov_type (incl. a within+hc2 R-parity sibling class; the one-way HC2-BM lane goes to TODO.md), estimators (a TwoWayFixedEffectsEventStudy twin class; the class-scoped panels hoisted to module builders), the event-study surface roster (two native producers), variance-conventions (two twfe_es_* rows + regenerated table), and the DiagnosticReport/BusinessReport generic fixtures (static TWFE fit, legacy MPD fixture kept for the event-study-lane and dispatch pins). Ported avg_att/avg_se reads use tests/helpers/event_study_average.py (MPD's arithmetic); hand-built containers use tests/helpers/event_study_surface_builders.py. Spec: item 4c split; 4c-ii records the spec="within" benchmark relabel, the M-016 record type and the build_event_study_surface sniff reorder. Ledger: M-010/M-011 notes only (no status flips). TODO row for the HC2-BM lane. --- TODO.md | 1 + .../20261024-v4-pr4c-i-twfe-es-decoupling.md | 26 + diff_diff/estimators.py | 145 +- diff_diff/results_base.py | 287 ++- diff_diff/twfe.py | 29 +- diff_diff/visualization/_event_study.py | 5 +- docs/methodology/variance-conventions.md | 3 + docs/v4-deprecations.yaml | 4 +- docs/v4-design.md | 40 +- tests/_capture_v4_merge_mpd_oracles.py | 145 ++ tests/fixtures/README.md | 31 + .../fixtures/v4_twfe_event_study_oracles.json | 1903 +++++++++++++++++ tests/helpers/event_study_average.py | 52 + tests/helpers/event_study_surface_builders.py | 106 + tests/helpers/twfe_es_oracle_lanes.py | 223 ++ tests/test_business_report.py | 139 +- tests/test_conley_vcov.py | 111 +- tests/test_diagnostic_report.py | 81 +- tests/test_estimators.py | 299 ++- tests/test_estimators_vcov_type.py | 420 +++- tests/test_event_study_surface.py | 91 +- tests/test_honest_did.py | 361 ++-- tests/test_methodology_honest_did.py | 79 +- tests/test_methodology_pretrends.py | 42 +- tests/test_pretrends.py | 452 ++-- tests/test_replicate_weight_expansion.py | 71 +- tests/test_survey.py | 56 +- tests/test_survey_phase6.py | 17 +- tests/test_twfe_event_study_decoupling.py | 446 ++++ tests/test_variance_conventions.py | 46 + tests/test_visualization.py | 58 +- tests/test_visualization_plotly.py | 26 +- 32 files changed, 4755 insertions(+), 1040 deletions(-) create mode 100644 changelog.d/20261024-v4-pr4c-i-twfe-es-decoupling.md create mode 100644 tests/_capture_v4_merge_mpd_oracles.py create mode 100644 tests/fixtures/v4_twfe_event_study_oracles.json create mode 100644 tests/helpers/event_study_average.py create mode 100644 tests/helpers/event_study_surface_builders.py create mode 100644 tests/helpers/twfe_es_oracle_lanes.py create mode 100644 tests/test_twfe_event_study_decoupling.py diff --git a/TODO.md b/TODO.md index 7915ddda6..dd14937a2 100644 --- a/TODO.md +++ b/TODO.md @@ -86,6 +86,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| +| TWFE event-study one-way HC2-BM absorbed-FE R parity against the `mpd_absorbed_fe_did` golden (`clubSandwich::vcovCR(cluster=1:n, type="CR2")`): the within spec with hc2_bm auto-clusters at unit, so the one-way lane needs an explicit singleton-cluster column (`cluster=` + hc2_bm = CR2 with G=n, PT2018 §3.3) once the G=n CR2 path is verified; the hc2 lane is already pinned by `TestTWFEEventStudyAbsorbedFERParity`, and the MultiPeriodDiD lanes die at the 4.0 cut (PR 4c-ii) | `tests/test_estimators_vcov_type.py` | #v4-4c-i | Mid | Medium | | Re-validate the plan-review dual engine on `gpt-6.1-sol` with `tools/plan-review-eval`: re-pin every arm's composite `cli_version` to the live Claude Code + codex CLI pair, add a gpt-6.1-sol arm, run the harness against the gpt-5.6-sol baseline (the 2026-10 swap carried the model over from the code-review campaign 2 verdict without a plan-review campaign). | `tools/plan-review-eval/config/configs.json` | reviewer-eval gpt61sol PR-2 | Heavy | Low | | TWFE-weights tutorial section: `attgt_weights` (all three `type=` values side by side) + `decompose_twfe_weights` + `plot_twfe_weights` on `mpdta`, either as a section of `02_staggered_did.ipynb` or a new notebook registered in `docs/tutorials/index.rst`; follow-on PR by decision, the feature ships without it | `docs/tutorials/02_staggered_did.ipynb` | #812 | Mid | Medium | | Geo-experiments tutorial: the SDiD fits on the `generate_factor_data(...)` panel (conversions ~1,500, sigma-hat ~310) hit the Frank-Wolfe iteration cap - the placebo fit warns once and the bootstrap cross-check warns on 100/100 draws - yet the narrative presents the bootstrap SE as a paper-faithful robustness check. Pre-existing on `main` under BOTH backends (surfaced when the notebook outputs were refreshed for the shape-only fit diagnostic); the committed April outputs predate the aggregated non-convergence warning. Either rescale / re-tune the DGP (or raise `max_iter` / loosen `min_decrease`) so the refits converge and re-execute, or label the bootstrap result unreliable and drop it from the "Validity evidence" claim. | `docs/tutorials/18_geo_experiments.ipynb` | local codex review of the SDiD shape-only fit-RMSE PR | Mid | Low | diff --git a/changelog.d/20261024-v4-pr4c-i-twfe-es-decoupling.md b/changelog.d/20261024-v4-pr4c-i-twfe-es-decoupling.md new file mode 100644 index 000000000..56ac2d705 --- /dev/null +++ b/changelog.d/20261024-v4-pr4c-i-twfe-es-decoupling.md @@ -0,0 +1,26 @@ +### Internal +- **The TWFE event-study surface is built directly from the shared core's payload** + (4.0 cut, PR 4c-i; ledger rows `M-010` / `M-011`): `DifferenceInDifferences._fit_event_study_core` + returns a private `_EventStudyCorePayload`, `TwoWayFixedEffects.fit(..., event_study=True)` + builds its `EventStudyResults` from it via `results_base._from_event_study_payload`, and only + `MultiPeriodDiD.fit` still wraps the payload into `MultiPeriodDiDResults`. `_from_mpd` and the + payload route share one calendar-surface builder, so they cannot drift; a committed oracle + captured on the pre-refactor tree (`tests/fixtures/v4_twfe_event_study_oracles.json`, capture + script `tests/_capture_v4_merge_mpd_oracles.py`) pins every field of the TWFE surface and of + its MultiPeriodDiD twin across eleven variance lanes. No user-visible behavior changes. +- Spec item 4c is split into 4c-i (this PR) and 4c-ii (the removals); the 4c-ii bullet records + the `spec="within"` benchmark-arm relabel and the `M-016` view's record type (user decisions + 2026-10-24). Ledger `M-010` / `M-011` notes updated; no status flips. + +### Testing +- The MultiPeriodDiD-fixture suites (honest_did, pretrends, methodology pretrends/honest, + visualization, Conley, survey, replicate weights, vcov_type, df_convention, the estimators + event-study behaviour pins, the event-study surface roster, the DiagnosticReport / + BusinessReport generic fixtures) now run on `TwoWayFixedEffects().fit(..., event_study=True)` + (or a static TWFE fit for the report classes); MultiPeriodDiD's own tests, the MPD-only + designs (multi-`absorb`, `fixed_effects=`, wild bootstrap, container-field assertions, the + MPD-branch consumer tests) and the MPD-vs-TWFE parity oracles are unchanged until PR 4c-ii. + New `tests/test_twfe_event_study_decoupling.py`, shared helpers + `tests/helpers/twfe_es_oracle_lanes.py`, `event_study_average.py` and + `event_study_surface_builders.py`; the variance-conventions inventory gains the two TWFE + event-study rows; a TODO row tracks the one-way HC2-BM absorbed-FE R-parity lane. diff --git a/diff_diff/estimators.py b/diff_diff/estimators.py index c04a74696..912eb226e 100644 --- a/diff_diff/estimators.py +++ b/diff_diff/estimators.py @@ -28,6 +28,7 @@ solve_ols, ) from diff_diff.results import DiDResults, MultiPeriodDiDResults, PeriodEffect +from diff_diff.results_base import _EventStudyCorePayload from diff_diff.utils import ( WildBootstrapResults, absorbed_fe_cr1_k_increment, @@ -1336,9 +1337,16 @@ def _fit_event_study_core( cluster_override: Any = _USE_SELF_CLUSTER, estimator_name: str = "MultiPeriodDiD", _frame_offset: int = 0, - ) -> Tuple[MultiPeriodDiDResults, np.ndarray, np.ndarray]: + ) -> _EventStudyCorePayload: """Shared event-study estimation core (row M-010). + Returns a private :class:`~diff_diff.results_base._EventStudyCorePayload` + (PR 4c-i decoupling): ``MultiPeriodDiD.fit`` wraps it into the legacy + ``MultiPeriodDiDResults`` container via ``_mpd_results_from_payload``; + the TWFE event-study mode builds the unified ``EventStudyResults`` + surface from it directly, so no MultiPeriodDiD container is + constructed on that path. + The relocated body of ``MultiPeriodDiD.fit`` (verbatim through the 3.9 merge; docs/v4-design.md section 4.1): validation, the pooled event-study design build, OLS via ``solve_ols``, survey / @@ -2251,7 +2259,15 @@ def _refit_mp_absorb(w_r): _bm_dof_avg = float(_dof_all[-1]) # Extract period-specific treatment effects for ALL non-reference periods - period_effects = {} + # (index-aligned lists over non_ref_periods; the MultiPeriodDiD + # wrapper zips them into PeriodEffect records, the TWFE surface + # builder reads them as arrays). + pe_att: List[float] = [] + pe_se: List[float] = [] + pe_t: List[float] = [] + pe_p: List[float] = [] + pe_lo: List[float] = [] + pe_hi: List[float] = [] post_effect_values = [] post_effect_indices = [] # Per-period df PROVENANCE for the unified event-study surface @@ -2284,14 +2300,12 @@ def _refit_mp_absorb(w_r): else float("nan") ) - period_effects[period] = PeriodEffect( - period=period, - effect=effect, - se=se, - t_stat=t_stat, - p_value=p_value, - conf_int=conf_int, - ) + pe_att.append(effect) + pe_se.append(se) + pe_t.append(t_stat) + pe_p.append(p_value) + pe_lo.append(conf_int[0]) + pe_hi.append(conf_int[1]) if period in post_periods: post_effect_values.append(effect) @@ -2338,32 +2352,41 @@ def _refit_mp_absorb(w_r): n_treated = n_treated_raw n_control = n_control_raw - # Create coefficient dictionary (var_names uniqueness already enforced - # before the fit above). - coef_dict = {name: coef for name, coef in zip(var_names, coefficients)} - - # Store results - _core_results = MultiPeriodDiDResults( - period_effects=period_effects, + # Hand off to the builders (PR 4c-i): MultiPeriodDiD wraps this into + # its container; the TWFE event-study mode builds the surface from it. + return _EventStudyCorePayload( + periods=list(non_ref_periods), + att=np.asarray(pe_att, dtype=float), + se=np.asarray(pe_se, dtype=float), + t_stat=np.asarray(pe_t, dtype=float), + p_value=np.asarray(pe_p, dtype=float), + ci_lower=np.asarray(pe_lo, dtype=float), + ci_upper=np.asarray(pe_hi, dtype=float), + df_by_period=es_df_used, + reference_period=reference_period, + pre_periods=pre_periods, + post_periods=post_periods, + alpha=self.alpha, + coefficients=coefficients, + coef_names=list(var_names), + vcov=vcov, + interaction_indices=interaction_indices, + residuals=residuals, + fitted_values=fitted, + r_squared=r_squared, avg_att=avg_att, avg_se=avg_se, avg_t_stat=avg_t_stat, avg_p_value=avg_p_value, avg_conf_int=avg_conf_int, + inference_df=( + float(_avg_df) + if _avg_df is not None and np.isfinite(_avg_df) and _avg_df > 0 + else None + ), n_obs=len(y), n_treated=n_treated, n_control=n_control, - pre_periods=pre_periods, - post_periods=post_periods, - alpha=self.alpha, - coefficients=coef_dict, - vcov=vcov, - residuals=residuals, - fitted_values=fitted, - r_squared=r_squared, - reference_period=reference_period, - interaction_indices=interaction_indices, - survey_metadata=survey_metadata, # Report the family that actually produced the SE; may be the # remapped hc1 under a replicate-weight survey design. vcov_type=_fit_vcov_type, @@ -2373,15 +2396,59 @@ def _refit_mp_absorb(w_r): ), conley_lag_cutoff=(self.conley_lag_cutoff if _fit_vcov_type == "conley" else None), df_convention=self.df_convention, - inference_df=( - float(_avg_df) - if _avg_df is not None and np.isfinite(_avg_df) and _avg_df > 0 - else None - ), - event_study_df=es_df_used, + survey_metadata=survey_metadata, ) - return _core_results, coefficients, vcov + +def _mpd_results_from_payload(p: _EventStudyCorePayload) -> MultiPeriodDiDResults: + """Wrap the event-study core's payload into the legacy MultiPeriodDiD + container (deleted with the class at the 4.0 cut, PR 4c-ii). + + ``period_effects`` is built in ``p.periods`` order (the core's sorted + non-reference order, the same insertion order the core used before the + PR 4c-i decoupling); per-period scalars round-trip through float64 + arrays (value-identical; ``np.float64`` is a ``float`` subclass). + """ + period_effects = { + per: PeriodEffect( + period=per, + effect=p.att[i], + se=p.se[i], + t_stat=p.t_stat[i], + p_value=p.p_value[i], + conf_int=(p.ci_lower[i], p.ci_upper[i]), + ) + for i, per in enumerate(p.periods) + } + return MultiPeriodDiDResults( + period_effects=period_effects, + avg_att=p.avg_att, + avg_se=p.avg_se, + avg_t_stat=p.avg_t_stat, + avg_p_value=p.avg_p_value, + avg_conf_int=p.avg_conf_int, + n_obs=p.n_obs, + n_treated=p.n_treated, + n_control=p.n_control, + pre_periods=p.pre_periods, + post_periods=p.post_periods, + alpha=p.alpha, + coefficients=dict(zip(p.coef_names, p.coefficients)), + vcov=p.vcov, + residuals=p.residuals, + fitted_values=p.fitted_values, + r_squared=p.r_squared, + reference_period=p.reference_period, + interaction_indices=p.interaction_indices, + survey_metadata=p.survey_metadata, + vcov_type=p.vcov_type, + cluster_name=p.cluster_name, + n_clusters=p.n_clusters, + conley_lag_cutoff=p.conley_lag_cutoff, + df_convention=p.df_convention, + inference_df=p.inference_df, + event_study_df=p.df_by_period, + ) class MultiPeriodDiD(DifferenceInDifferences): @@ -2624,7 +2691,7 @@ def fit( # type: ignore[override] ) effective_inference = "analytical" - results, coefficients, vcov = self._fit_event_study_core( + payload = self._fit_event_study_core( data, outcome, treatment, @@ -2639,9 +2706,9 @@ def fit( # type: ignore[override] effective_inference, _frame_offset=1, ) - self.results_ = results - self._coefficients = coefficients - self._vcov = vcov + self.results_ = _mpd_results_from_payload(payload) + self._coefficients = payload.coefficients + self._vcov = payload.vcov self.is_fitted_ = True return self.results_ diff --git a/diff_diff/results_base.py b/diff_diff/results_base.py index 8b3e66147..474c8aa05 100644 --- a/diff_diff/results_base.py +++ b/diff_diff/results_base.py @@ -401,8 +401,9 @@ class EventStudyResults(BaseResults): #: need pre/post classification on a calendar surface (HonestDiD, #: PreTrendsPower, the plotter) read THIS field; pre-periods derive as #: the complement: non-reference rows whose ``event_time`` is not in - #: ``post_periods``. Threaded by ``_from_mpd`` (and the TWFE - #: event-study producer); None from every other builder. + #: ``post_periods``. Threaded by the calendar-surface builder (the + #: MultiPeriodDiD adapter and the TWFE event-study producer); None + #: from every other builder. post_periods: Optional[Tuple[Any, ...]] = None #: TWFE event-study design provenance (M-092 amendment #5, with row #: M-010): ``"within"`` (unit + time FE) or ``"pooled"`` (the @@ -886,19 +887,34 @@ def _resolve_scalar_df_survey(results: Any) -> Optional[float]: together is tracked in TODO.md (df-resolution / adapter-naming consolidation row). """ - sm = getattr(results, "survey_metadata", None) - if sm is not None: - df_survey = getattr(sm, "df_survey", None) - if df_survey is not None: - return float(df_survey) - if getattr(sm, "replicate_method", None) is not None: - return 0.0 + from_metadata = _scalar_df_survey_from_metadata(getattr(results, "survey_metadata", None)) + if from_metadata is not None: + return from_metadata df_inference = getattr(results, "df_inference", None) if df_inference is not None: return float(df_inference) return None +def _scalar_df_survey_from_metadata(sm: Any) -> Optional[float]: + """The survey-metadata half of :func:`_resolve_scalar_df_survey`. + + ``survey_metadata.df_survey`` where present; a replicate design whose df + is undefined maps to the ``0.0`` sentinel; otherwise None (the caller + may then consult a bare-``cluster=`` ``df_inference`` carrier). Shared + by the duck-typed adapters and the TWFE event-study payload builder, + whose payload carries the metadata object but no results container. + """ + if sm is None: + return None + df_survey = getattr(sm, "df_survey", None) + if df_survey is not None: + return float(df_survey) + if getattr(sm, "replicate_method", None) is not None: + return 0.0 + return None + + def _provenance_kwargs(results: Any) -> Dict[str, Any]: """Producer-provenance fields threaded onto the container. @@ -1119,30 +1135,121 @@ def _from_relative_dict(results: Any) -> EventStudyResults: ) -def _from_mpd(results: Any) -> EventStudyResults: - """Adapter for MultiPeriodDiD's calendar-keyed ``period_effects``. +@dataclass(frozen=True, eq=False) +class _EventStudyCorePayload: + """Producer-neutral output of ``DifferenceInDifferences._fit_event_study_core``. - The reference period is omitted from the native dict and recorded on - ``results.reference_period``; the marked row is synthesized here - (``att=0.0``, NaN inference). The event-study vcov is the sub-block of - ``results.vcov`` selected by ``results.interaction_indices`` - pure - indexing of stored entries, no recomputation. + Private hand-off between the shared event-study core and its two + builders: ``MultiPeriodDiD.fit`` wraps it into the legacy + ``MultiPeriodDiDResults`` container, and the TWFE event-study mode + builds the unified :class:`EventStudyResults` surface from it directly + (:func:`_from_event_study_payload`) - so the TWFE path no longer + constructs a MultiPeriodDiD container (ledger rows M-010 / M-011). + + Per-period arrays are index-aligned to ``periods``, the core's SORTED + non-reference iteration order. The ``avg_*`` block, ``inference_df`` and + the count / label fields exist only for the MultiPeriodDiD container and + are deleted with it at the 4.0 cut (PR 4c-ii). ``eq=False``: the + generated ``__eq__``/``__hash__`` would raise on the ndarray fields. """ - period_effects = getattr(results, "period_effects", None) - if not period_effects: - raise _absent(results) - ref = getattr(results, "reference_period", None) - periods = sorted(period_effects.keys()) + # Per-period inference (one entry per non-reference period, sorted). + periods: List[Any] + att: np.ndarray + se: np.ndarray + t_stat: np.ndarray + p_value: np.ndarray + ci_lower: np.ndarray + ci_upper: np.ndarray + df_by_period: Dict[Any, float] + # Partition / normalisation. + reference_period: Any + pre_periods: List[Any] + post_periods: List[Any] + alpha: float + # Full-design inference (the MultiPeriodDiD container + the surface's + # event-study vcov sub-block). + coefficients: np.ndarray + coef_names: List[str] + vcov: np.ndarray + interaction_indices: Dict[Any, int] + residuals: np.ndarray + fitted_values: np.ndarray + r_squared: float + # MultiPeriodDiD-only: the post-period average and its contrast df + # (die with the class at 4c-ii). + avg_att: float + avg_se: float + avg_t_stat: float + avg_p_value: float + avg_conf_int: Tuple[float, float] + inference_df: Optional[float] + # MultiPeriodDiD-only counts and labels. + n_obs: int + n_treated: int + n_control: int + vcov_type: Optional[str] + cluster_name: Optional[str] + n_clusters: Optional[int] + conley_lag_cutoff: Optional[int] + df_convention: Optional[str] + # Shared provenance (the survey df is resolved from it). + survey_metadata: Optional[Any] + + +def _calendar_surface_from_arrays( + *, + periods: Any, + ref: Any, + att: Any, + se: Any, + t_stat: Any, + p_value: Any, + ci_lo: Any, + ci_hi: Any, + full_vcov: Any, + interaction_indices: Any, + df_map: Any, + post_periods: Any, + alpha: float, + source: Optional[str], + df_survey: Optional[float], + estimation_spec: Optional[str] = None, + base_period: Any = None, + anticipation: Any = None, + reference_event_times: Any = None, +) -> EventStudyResults: + """Build the calendar-keyed surface from index-aligned per-period arrays. + + The one builder behind both :func:`_from_mpd` (which extracts the arrays + from a ``MultiPeriodDiDResults``) and :func:`_from_event_study_payload` + (which reads them off the core's payload), so the two paths cannot + drift. ``periods`` must already be SORTED (the caller supplies the + non-reference labels in sorted order and the arrays aligned to them); + an explicit raise guards that contract - a silent re-sort of the labels + alone would pair labels with the wrong values. The reference row is + inserted (``is_reference``, NaN inference); the event-study vcov is the + sub-block of ``full_vcov`` selected by ``interaction_indices`` - pure + indexing of stored entries, no recomputation; per-period ``df`` comes + STRICTLY from ``df_map`` (the per-period df actually handed to + ``safe_inference``), never from the post-average contrast df. + """ + periods = list(periods) + if periods != sorted(periods): + raise ValueError( + "_calendar_surface_from_arrays requires `periods` in sorted order " + f"(got {periods!r}); the per-period arrays are index-aligned to it." + ) all_periods = sorted(periods + [ref]) if ref is not None else periods + pos = {p: i for i, p in enumerate(periods)} n_rows = len(all_periods) - att = np.full(n_rows, np.nan) - se = np.full(n_rows, np.nan) - t_stat = np.full(n_rows, np.nan) - p_value = np.full(n_rows, np.nan) - ci_lo = np.full(n_rows, np.nan) - ci_hi = np.full(n_rows, np.nan) + att_out = np.full(n_rows, np.nan) + se_out = np.full(n_rows, np.nan) + t_out = np.full(n_rows, np.nan) + p_out = np.full(n_rows, np.nan) + lo_out = np.full(n_rows, np.nan) + hi_out = np.full(n_rows, np.nan) n = np.full(n_rows, np.nan) is_ref = np.zeros(n_rows, dtype=bool) @@ -1150,17 +1257,15 @@ def _from_mpd(results: Any) -> EventStudyResults: if ref is not None and p == ref: is_ref[i] = True continue - pe = period_effects[p] - att[i] = pe.effect - se[i] = pe.se - t_stat[i] = pe.t_stat - p_value[i] = pe.p_value - ci_lo[i], ci_hi[i] = float(pe.conf_int[0]), float(pe.conf_int[1]) + j = pos[p] + att_out[i] = att[j] + se_out[i] = se[j] + t_out[i] = t_stat[j] + p_out[i] = p_value[j] + lo_out[i], hi_out[i] = float(ci_lo[j]), float(ci_hi[j]) vcov_sub: Optional[np.ndarray] = None vcov_index_arr: Optional[np.ndarray] = None - full_vcov = getattr(results, "vcov", None) - interaction_indices = getattr(results, "interaction_indices", None) if full_vcov is not None and interaction_indices: covered = [p for p in periods if p in interaction_indices] if covered: @@ -1168,31 +1273,20 @@ def _from_mpd(results: Any) -> EventStudyResults: vcov_sub = np.asarray(full_vcov, dtype=float)[np.ix_(idx, idx)] vcov_index_arr = np.asarray(covered) - # Per-period df: STRICTLY the `event_study_df` channel ({period: df} - # actually passed to safe_inference - per-period BM DOF under hc2_bm). - # `inference_df` is deliberately NOT a fallback here: it stores the - # POST-AVERAGE contrast df, which is the wrong provenance for the - # per-period rows under hc2_bm. - df_map = getattr(results, "event_study_df", None) df_arg: Optional[np.ndarray] = None if isinstance(df_map, dict): df_arg = np.array([float(df_map[p]) if p in df_map else np.nan for p in all_periods]) - # Calendar-partition provenance (M-092 amendment #5): the producer's - # AUTHORITATIVE post_periods list - an arbitrary subset is legal on - # MultiPeriodDiD, so consumers must never re-derive the partition - # positionally from the reference row. - _mpd_post = getattr(results, "post_periods", None) - post_periods_arg: Optional[Tuple[Any, ...]] = tuple(_mpd_post) if _mpd_post else None + post_periods_arg: Optional[Tuple[Any, ...]] = tuple(post_periods) if post_periods else None return EventStudyResults( event_time=np.asarray(all_periods), - att=att, - se=se, - t_stat=t_stat, - p_value=p_value, - conf_int_lower=ci_lo, - conf_int_upper=ci_hi, + att=att_out, + se=se_out, + t_stat=t_out, + p_value=p_out, + conf_int_lower=lo_out, + conf_int_upper=hi_out, is_reference=is_ref, n=n, n_kind=None, @@ -1201,11 +1295,88 @@ def _from_mpd(results: Any) -> EventStudyResults: event_time_convention=None, vcov=vcov_sub, vcov_index=vcov_index_arr, - alpha=getattr(results, "alpha", 0.05), - source=type(results).__name__, + alpha=alpha, + source=source, df=df_arg, post_periods=post_periods_arg, - **_provenance_kwargs(results), + estimation_spec=estimation_spec, + base_period=base_period, + anticipation=anticipation, + df_survey=df_survey, + reference_event_times=( + tuple(reference_event_times) if reference_event_times is not None else None + ), + ) + + +def _from_mpd(results: Any) -> EventStudyResults: + """Adapter for MultiPeriodDiD's calendar-keyed ``period_effects``. + + The reference period is omitted from the native dict and recorded on + ``results.reference_period``; the marked row is synthesized by the + shared builder (``att=0.0``, NaN inference). The dict's insertion order + is irrelevant: the labels are sorted here and the value lists built in + that order before the hand-off. + """ + period_effects = getattr(results, "period_effects", None) + if not period_effects: + raise _absent(results) + + periods = sorted(period_effects.keys()) + pes = [period_effects[p] for p in periods] + provenance = _provenance_kwargs(results) + return _calendar_surface_from_arrays( + periods=periods, + ref=getattr(results, "reference_period", None), + att=[pe.effect for pe in pes], + se=[pe.se for pe in pes], + t_stat=[pe.t_stat for pe in pes], + p_value=[pe.p_value for pe in pes], + ci_lo=[pe.conf_int[0] for pe in pes], + ci_hi=[pe.conf_int[1] for pe in pes], + full_vcov=getattr(results, "vcov", None), + interaction_indices=getattr(results, "interaction_indices", None), + df_map=getattr(results, "event_study_df", None), + post_periods=getattr(results, "post_periods", None), + alpha=getattr(results, "alpha", 0.05), + source=type(results).__name__, + df_survey=provenance["df_survey"], + base_period=provenance["base_period"], + anticipation=provenance["anticipation"], + reference_event_times=provenance["reference_event_times"], + ) + + +def _from_event_study_payload( + payload: _EventStudyCorePayload, + *, + source: str, + estimation_spec: Optional[str] = None, +) -> EventStudyResults: + """Build the unified surface directly from the event-study core's payload. + + The TWFE event-study producer's path: no MultiPeriodDiD container is + constructed. Field for field this is what ``_from_mpd`` yields for the + MultiPeriodDiD twin of the same fit (pinned by + ``tests/test_twfe_event_study_decoupling.py``). + """ + return _calendar_surface_from_arrays( + periods=payload.periods, + ref=payload.reference_period, + att=payload.att, + se=payload.se, + t_stat=payload.t_stat, + p_value=payload.p_value, + ci_lo=payload.ci_lower, + ci_hi=payload.ci_upper, + full_vcov=payload.vcov, + interaction_indices=payload.interaction_indices, + df_map=payload.df_by_period, + post_periods=payload.post_periods, + alpha=payload.alpha, + source=source, + estimation_spec=estimation_spec, + df_survey=_scalar_df_survey_from_metadata(payload.survey_metadata), ) @@ -1415,7 +1586,9 @@ def build_event_study_surface(results: Any) -> EventStudyResults: """ # Order matters: dCDH carries both event_study_effects and the split # placebo dict; HAD's columnar container and MPD's period_effects have - # unique attributes; LPDiD's surface is a DataFrame field. + # unique attributes; LPDiD's surface is a DataFrame field. (The TWFE + # event-study producer never arrives here: it builds the surface + # natively from the core's payload via _from_event_study_payload.) if hasattr(results, "placebo_event_study") and hasattr(results, "event_study_effects"): return _from_dcdh(results) if hasattr(results, "event_times") and hasattr(results, "n_obs_per_horizon"): diff --git a/diff_diff/twfe.py b/diff_diff/twfe.py index ea037c3f0..4059b5967 100644 --- a/diff_diff/twfe.py +++ b/diff_diff/twfe.py @@ -2,7 +2,6 @@ Two-Way Fixed Effects estimator for panel Difference-in-Differences. """ -import dataclasses import warnings from typing import TYPE_CHECKING, Any, List, Optional, Union @@ -17,7 +16,7 @@ from diff_diff.estimators import DifferenceInDifferences from diff_diff.linalg import LinearRegression from diff_diff.results import DiDResults -from diff_diff.results_base import EventStudyResults, _from_mpd +from diff_diff.results_base import EventStudyResults, _from_event_study_payload from diff_diff.utils import ( absorbed_fe_cr1_k_increment, absorbed_fe_rank, @@ -1157,7 +1156,7 @@ def _fit_event_study( ) else: absorb_arg = None - results_internal, _, _ = self._fit_event_study_core( + payload = self._fit_event_study_core( data, outcome, treatment, @@ -1177,20 +1176,16 @@ def _fit_event_study( _frame_offset=2, ) - # Step 8: convert to the unified surface. The internal - # MultiPeriodDiDResults is a construction detail (the 4.0 removal - # PR decouples it - ledger row M-011); the surface carries the - # TWFE provenance: source, the authoritative calendar partition - # (threaded by _from_mpd), and the design spec. The legacy - # writer-only _coefficients/_vcov attrs are deliberately left - # untouched (static TWFE has never set them; nothing reads them). - # dataclasses.replace re-runs __post_init__, so the final provenance - # (source + estimation_spec) passes the container's own validation - # atomically instead of being patched on after construction. - surface = dataclasses.replace( - _from_mpd(results_internal), - source="TwoWayFixedEffects", - estimation_spec=spec, + # Step 8: build the unified surface directly from the core's payload + # (PR 4c-i decoupling, ledger rows M-010 / M-011: no MultiPeriodDiD + # container is constructed on this path). The surface carries the + # TWFE provenance - source, the authoritative calendar partition and + # the design spec - at construction, so it passes the container's + # own validation atomically. The legacy writer-only + # _coefficients/_vcov attrs are deliberately left untouched (static + # TWFE has never set them; nothing reads them). + surface = _from_event_study_payload( + payload, source="TwoWayFixedEffects", estimation_spec=spec ) self.results_ = surface self.is_fitted_ = True diff --git a/diff_diff/visualization/_event_study.py b/diff_diff/visualization/_event_study.py index fc79a61d8..8f3a4dd95 100644 --- a/diff_diff/visualization/_event_study.py +++ b/diff_diff/visualization/_event_study.py @@ -800,8 +800,9 @@ def _extract_plot_data( derived_post = [p for p in periods if p >= _post_start] elif surface.post_periods is not None: # Calendar labels with the AUTHORITATIVE partition - # provenance (row M-092 amendment #5, threaded by - # _from_mpd / the TWFE event-study producer): the declared + # provenance (row M-092 amendment #5, threaded by the + # calendar-surface builder behind the MultiPeriodDiD adapter + # and the TWFE event-study producer): the declared # post set may be an arbitrary subset (non-suffix splits # are legal on MultiPeriodDiD), so the positional # fallback below would mislabel it - always prefer the diff --git a/docs/methodology/variance-conventions.md b/docs/methodology/variance-conventions.md index 6f38beaa1..b077eee1b 100644 --- a/docs/methodology/variance-conventions.md +++ b/docs/methodology/variance-conventions.md @@ -33,6 +33,8 @@ produced wrong figures three separate times while this inventory was drafted). | `mpd_absorb_hc1_cluster_unit` | 11 | 290, 290, 290, 290, 290, 290 | **legitimate** | K_reference: 6 visible + rank(time given unit) = 5; equals the fixed_effects form's 70 - 59 (MPD absorb/fixed_effects equivalence) | | `mpd_fixed_effects_hc1_cluster_unit` | 11 | 290, 290, 290, 290, 290, 290 | **legitimate** | K_reference: 70 visible (incl. built-in period dummies, MPD's time-FE block) minus the 59 cluster-nested unit dummies = 11 — identical to the absorb form | | `mpd_plain_hc1_cluster_time` | 7 | 348, 348, 348, 348, 348, 348 | **legitimate** | the NESTED orientation of the built-in period dummies: 12 visible minus their rank 5 under a time cluster (under-subtraction is caught here; the unit-cluster rows catch over-) | +| `twfe_es_within_hc1_cluster_unit` | 11 | 290, 290, 290, 290, 290, 290 | **legitimate** | TWFE event-study, within spec (4.0 cut, PR 4c-i): the absorbed unit FE plus the built-in period dummies give the same K_reference as the mpd_absorb row - 6 visible + rank(time given unit) = 5 | +| `twfe_es_pooled_hc1_cluster_time` | 7 | 348, 348, 348, 348, 348, 348 | **legitimate** | TWFE event-study, pooled spec (the MultiPeriodDiD design; 4.0 cut, PR 4c-i): the NESTED orientation of the built-in period dummies under a time cluster - 12 visible minus their rank 5, as mpd_plain_hc1_cluster_time | | `lpdid_absorb_nested_cluster_grp` | 4, 4, 5, 5, 5, 6 | 1, 1, 1, 1, 1, 1 | **legitimate** | LPDiD absorb dummies nested in the cluster subtract their rank (adj -1 per horizon: region == grp here); _event_time stays counted (unit-level cluster does not nest time); G-1 tail df (L2) | | `imputation_default` | — | None | **legitimate** | L3: BJS imputation variance, not the shared CR1 sandwich | | `imputation_pretrends_event_study` | unpinned | 157, 157, 157, None, None, None, None, None | **legitimate** | pretrends lead regression: CR1 k = K_reference (D2 fixed); tail df converged in 3.9 (M-127): leads use t(residual df) via the df_convention knob (was silent normal theory); the None calls are the knob-independent BJS aggregates (L3) | @@ -42,6 +44,7 @@ produced wrong figures three separate times while this inventory was drafted). | `dml_did_bad_control` | — | None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None | **legitimate** | L3: Caetano et al. (2026) Eq. 11 augmented-score plug-in variance (same psi_bar - D*theta/pi centering as Chang Thm 2; per-unit influence function; normal-theory safe_inference on no-design fits; cluster= fits use the CR1 per-cell variance with df=df_survey t-inference) and the per-cell ATT_X diagnostic tuple sharing the same branch and df (one extra safe_inference call per retained cell) | | `dml_did_rcs` | — | None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None | **legitimate** | L3: Chang (2020) Thm 2 lambda-corrected augmented-score plug-in variance (Case 2; per-observation influence function; normal-theory safe_inference on no-design fits; survey_design=/cluster= fits use the design-based CR1/weighted-IF per-cell variance — replicate designs the IF-reweighting per-cell variance — with df=df_survey t-inference outside the shared CR1 sandwich) | | `duration_did` | — | None, None | **legitimate** | L3: Deaner & Ku (2026) whole-individual pooled bootstrap (centered absolute-deviation pointwise and simultaneous bands); no CR1 sandwich, normal-theory safe_inference gate only | + cr1_k is the sorted multiset of K_reference counts reaching the shared clustered CR1 denominator — visible columns + the signed cluster_k_adjustment (linalg._compute_robust_vcov_numpy with diff --git a/docs/v4-deprecations.yaml b/docs/v4-deprecations.yaml index 8fab79980..7e879ee77 100644 --- a/docs/v4-deprecations.yaml +++ b/docs/v4-deprecations.yaml @@ -153,7 +153,7 @@ rows: warning: FutureWarning test_ref: tests/test_v4_merge_mpd.py code_refs: [diff_diff/estimators.py, diff_diff/twfe.py, diff_diff/__init__.py, diff_diff/honest_did.py, diff_diff/pretrends.py, diff_diff/diagnostic_report.py, diff_diff/business_report.py, diff_diff/visualization/_event_study.py] - notes: "Unit-FE event study default; spec='within'|'pooled' opt-in (pooled = old MPD model, repeated cross-sections). Wild bootstrap in event-study mode = explicit error. SHIPPED 3.9 (Phase 3(a)): TWFE.fit(event_study=True, spec=, reference_period=, post_periods=) returns the unified EventStudyResults surface natively (source='TwoWayFixedEffects', estimation_spec + post_periods provenance); the shared core is DifferenceInDifferences._fit_event_study_core (the relocated MPD fit body, parameterized - MPD numerics/messages/warning-attribution bit-identical); ES calls pass time=/unit= as KEYWORDS (slot 4 stays post= through the M-082 window; section 4.1's snippet amended to keyword form). post_periods= is REQUIRED (non-empty) in ES mode (Phase 3(a) review R4, section 4.1 second amendment): the boundary is unobservable from time-invariant D_i, so MPD's midpoint default is a silent guess and is not carried over - MPD itself keeps it through 3.9. Day-one unit auto-cluster (user decision 2026-08-07) with static TWFE's carve-outs mirrored lane-for-lane: dropped on Conley (no implicit spatial x unit product kernel), never injected as a survey PSU (implicit-per-observation-PSU rule), dropped for explicit one-way analytical families; explicit cluster= passes through and matches MPD's explicit-cluster behavior (PSU injection included; a survey design carrying its own PSU keeps PSU precedence on every class). The wild ValueError and the no-legacy-reference-warning contract are live from the mode's 3.9 birth (the section 4.1 'at 4.0' clause describes when MPD's fallback dies with the class). MultiPeriodDiD.__init__ is the warning shim (forwarding *args/**kwargs + an import-time __signature__ mirror of DiD's constructor; set_params re-emits via its probe re-init; the static-type-checking loss for the deprecated class is the DEFERRED.md decision-record waiver). Consumer ports: HonestDiD + PreTrendsPower calendar container routes reconstruct the native-branch inputs from the post_periods provenance (HonestDiD geometry-scoped to chronological partitions - REGISTRY Note + DEFERRED row for the pre-existing native non-suffix geometry; PreTrendsPower plumbs reference_period through the gamma-unit helper incl. the string-label degradation); plot_event_study derives calendar splits from the provenance and shades per contiguous pre-run on both renderers; DiagnosticReport/BusinessReport explicitly reject the surface (TODO row for admission). 4.0 REMOVAL NOTE: the TWFE ES branch builds the surface via the internal MultiPeriodDiDResults container + _from_mpd - decouple when [M-011] removes the container." + notes: "Unit-FE event study default; spec='within'|'pooled' opt-in (pooled = old MPD model, repeated cross-sections). Wild bootstrap in event-study mode = explicit error. SHIPPED 3.9 (Phase 3(a)): TWFE.fit(event_study=True, spec=, reference_period=, post_periods=) returns the unified EventStudyResults surface natively (source='TwoWayFixedEffects', estimation_spec + post_periods provenance); the shared core is DifferenceInDifferences._fit_event_study_core (the relocated MPD fit body, parameterized - MPD numerics/messages/warning-attribution bit-identical); ES calls pass time=/unit= as KEYWORDS (slot 4 stays post= through the M-082 window; section 4.1's snippet amended to keyword form). post_periods= is REQUIRED (non-empty) in ES mode (Phase 3(a) review R4, section 4.1 second amendment): the boundary is unobservable from time-invariant D_i, so MPD's midpoint default is a silent guess and is not carried over - MPD itself keeps it through 3.9. Day-one unit auto-cluster (user decision 2026-08-07) with static TWFE's carve-outs mirrored lane-for-lane: dropped on Conley (no implicit spatial x unit product kernel), never injected as a survey PSU (implicit-per-observation-PSU rule), dropped for explicit one-way analytical families; explicit cluster= passes through and matches MPD's explicit-cluster behavior (PSU injection included; a survey design carrying its own PSU keeps PSU precedence on every class). The wild ValueError and the no-legacy-reference-warning contract are live from the mode's 3.9 birth (the section 4.1 'at 4.0' clause describes when MPD's fallback dies with the class). MultiPeriodDiD.__init__ is the warning shim (forwarding *args/**kwargs + an import-time __signature__ mirror of DiD's constructor; set_params re-emits via its probe re-init; the static-type-checking loss for the deprecated class is the DEFERRED.md decision-record waiver). Consumer ports: HonestDiD + PreTrendsPower calendar container routes reconstruct the native-branch inputs from the post_periods provenance (HonestDiD geometry-scoped to chronological partitions - REGISTRY Note + DEFERRED row for the pre-existing native non-suffix geometry; PreTrendsPower plumbs reference_period through the gamma-unit helper incl. the string-label degradation); plot_event_study derives calendar splits from the provenance and shades per contiguous pre-run on both renderers; DiagnosticReport/BusinessReport explicitly reject the surface (TODO row for admission). DECOUPLED in PR 4c-i (2026-10-24): the TWFE event-study surface is built directly from _fit_event_study_core's private payload (results_base._from_event_study_payload); MultiPeriodDiDResults/_from_mpd are now reached only through MultiPeriodDiD.fit and the duck-typed build_event_study_surface branch, which PR 4c-ii deletes with the class." - id: M-011 kind: class group: merge-mpd @@ -164,7 +164,7 @@ rows: status: planned phase: 5 code_refs: [diff_diff/results.py, diff_diff/__init__.py] - notes: "Deprecation rides parent [M-010] (returned only by MPD). Successor = unified event-study surface on the merged TWFE results. 4.0 REMOVAL NOTE: the TWFE event-study branch (twfe.py:_fit_event_study) constructs this container INTERNALLY and converts via results_base._from_mpd - the removal PR must decouple that construction (build the EventStudyResults directly from the core's payload) before deleting the class." + notes: "Deprecation rides parent [M-010] (returned only by MPD). Successor = unified event-study surface on the merged TWFE results. DECOUPLED in PR 4c-i (2026-10-24): twfe.py:_fit_event_study no longer constructs this container - the core returns a private payload and the surface is built from it (results_base._from_event_study_payload; bit-identity pinned by tests/test_twfe_event_study_decoupling.py against the pre-refactor capture tests/fixtures/v4_twfe_event_study_oracles.json). The container is now constructed only by MultiPeriodDiD.fit (estimators._mpd_results_from_payload); PR 4c-ii deletes both." - id: M-012 kind: class group: merge-mpd diff --git a/docs/v4-design.md b/docs/v4-design.md index 0b7160b86..0ef31d663 100644 --- a/docs/v4-design.md +++ b/docs/v4-design.md @@ -373,7 +373,11 @@ settings and unconditionally in the no-unit repeated-cross-sections form. **Results.** 3.9's `TWFE(event_study=True)` returns the unified event-study surface (section 5) from day one - no intermediate container churn. `period_effects` becomes a FutureWarning property view over that surface at -4.0, removed 5.0 [M-016]. The HonestDiD / PreTrendsPower integrations (which +4.0, removed 5.0 [M-016] (the view's values are a private frozen record +mirroring the former `PeriodEffect` fields - `period, effect, se, t_stat, +p_value, conf_int, is_significant, significance_stars` - keyed by calendar +period with the reference omitted; user decision 2026-10-24, lands in PR +4c-ii). The HonestDiD / PreTrendsPower integrations (which read `interaction_indices` off MultiPeriodDiDResults) are ported to the unified surface in the same Phase 3 PR. @@ -1093,13 +1097,33 @@ tests-first PR that flips no rows). (01/02/09 - the [M-082] alias, misused as the post dummy) migrated to `post=` (01) and `TwoWayFixedEffects().decompose(...).twfe_estimate` (02/09) so [M-083] touches no executed notebook; 02 cell 36 seeded.) - - 4c - the MultiPeriodDiD class, `MultiPeriodDiDResults`, `PeriodEffect` + - 4c-i - decouple the TWFE event-study mode from `MultiPeriodDiDResults` / + `_from_mpd` (`_fit_event_study_core` returns a private payload; the + MultiPeriodDiD wrapper and the TWFE surface are both built from it) and + migrate the MultiPeriodDiD-fixture test suites to + `TwoWayFixedEffects().fit(..., event_study=True)` while MultiPeriodDiD is + still the live oracle (the 3a -> 3b pattern; user decision 2026-10-24). + (SHIPPED, PR 4c-i, 2026-10-24: committed pre-refactor oracle + `tests/fixtures/v4_twfe_event_study_oracles.json` captured at 68e383bd; + MPD-only designs - multi-`absorb`, `fixed_effects=`, wild bootstrap, + container-field assertions, the MPD-branch consumer tests - stay on the + class until 4c-ii.) + - 4c-ii - the MultiPeriodDiD class, `MultiPeriodDiDResults`, `PeriodEffect` and `EventStudy`, with [M-007], the [M-016] property view on - `EventStudyResults`, the [M-083] raise, the remaining `filterwarnings` - entry, the MultiPeriodDiD benchmark arm (relabelled `spec="pooled"`), the - power registry, and the two markdown-only MultiPeriodDiD mentions left in - the notebooks (06 cell 34's registry table - not executed in CI - and 16 - cell 40's survey-support row, both true until the class goes). + `EventStudyResults` (its values are a private frozen record - `period, + effect, se, t_stat, p_value, conf_int, is_significant, + significance_stars` - keyed by calendar period with the reference omitted; + user decision 2026-10-24), the [M-083] raise, the remaining + `filterwarnings` entry, the MultiPeriodDiD benchmark arm (relabelled + `spec="within"`: the arm fits the unit-FE design, `absorb=["unit"]` vs R + `feols(... | unit)`, so the pooled relabel the spec first named would + change what the arm measures - user decision 2026-10-24), the power + registry, the two markdown-only MultiPeriodDiD mentions left in the + notebooks (06 cell 34's registry table - not executed in CI - and 16 cell + 40's survey-support row, both true until the class goes), and the + `build_event_study_surface` sniff reorder (an `isinstance(EventStudyResults)` + guard ahead of the `period_effects` duck check, so the [M-016] view's + FutureWarning cannot fire on a surface input). 5. Storage flip [M-050]..[M-058] + pickle migration + the `.att` sweep over docs/notebooks/guides/benchmarks (incl. a disposition for the un-rowed `CSBootstrapResults.overall_att_se`). @@ -1114,7 +1138,7 @@ tests-first PR that flips no rows). 10. Release: bump to 4.0.0 (closes the cut window), changelog compile, migration-guide retitle, due gate green. -Ordering constraints: 4a before 4b before 4c; 4 before 5 and 6; 5 before 6 +Ordering constraints: 4a before 4b before 4c-i before 4c-ii; 4 before 5 and 6; 5 before 6 (the unified DDD container is born canonical); 7 and 8 independent; 9 after 4/6; 10 last. Waves 2 and 3 are independent of each other and of 4. diff --git a/tests/_capture_v4_merge_mpd_oracles.py b/tests/_capture_v4_merge_mpd_oracles.py new file mode 100644 index 000000000..278da41ed --- /dev/null +++ b/tests/_capture_v4_merge_mpd_oracles.py @@ -0,0 +1,145 @@ +"""Capture the TWFE event-study decoupling oracles for +tests/test_twfe_event_study_decoupling.py (4.0 cut, PR 4c-i). + +Run on the UNMODIFIED tree (the commit named in ``provenance.captured_at_head`` +of ``tests/fixtures/v4_twfe_event_study_oracles.json``) BEFORE the TWFE +event-study mode is decoupled from the MultiPeriodDiD containers, and write to +stdout like the aggregate precedent:: + + DIFF_DIFF_BACKEND=python python tests/_capture_v4_merge_mpd_oracles.py \ + > tests/fixtures/v4_twfe_event_study_oracles.json + DIFF_DIFF_BACKEND=rust python tests/_capture_v4_merge_mpd_oracles.py \ + > /path/to/scratch/rust.json # cross-backend check only + +Why. PR 4c-i rewrites how ``TwoWayFixedEffects.fit(..., event_study=True)`` +builds its ``EventStudyResults`` (today: an internal ``MultiPeriodDiDResults`` +converted by ``results_base._from_mpd``; after: a private payload from +``_fit_event_study_core``). Every ``dataclasses.fields(EventStudyResults)`` +value of the TWFE surface and of the MultiPeriodDiD twin +(``_from_mpd(MultiPeriodDiD(...).fit(...))``) is recorded per lane, so the +refactor's output can be compared field by field against the pre-refactor +tree. PR 4c-ii deletes MultiPeriodDiD; re-capturing on a post-4c-ii tree is +NOT a substitute for the MPD-twin half (nothing independent remains). + +Panels and lanes come from ``tests/helpers/twfe_es_oracle_lanes.py`` (shared +with the test module); the fixture carries each lane's ``echo()`` so an edited +lane fails as lane drift, not as a numeric regression. +""" + +from __future__ import annotations + +import dataclasses +import json +import os +import subprocess +import sys +import warnings +from typing import Any, Dict + +import numpy as np + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, os.path.join(HERE, "helpers")) + +from twfe_es_oracle_lanes import LANES, PANELS # noqa: E402 + +from diff_diff import MultiPeriodDiD, TwoWayFixedEffects # noqa: E402 +from diff_diff.results_base import EventStudyResults, _from_mpd # noqa: E402 + + +def _encode_value(value: Any) -> Any: + """JSON-safe encoding of one EventStudyResults field value. + + ndarray -> {"dtype": , "values": [...]} with NaN as null; + tuple -> list; numpy scalars -> Python scalars; NaN floats -> null. + """ + if isinstance(value, np.ndarray): + return {"dtype": value.dtype.str, "values": _nan_to_none(value.tolist())} + if isinstance(value, (tuple, list)): + return [_encode_value(v) for v in value] + if isinstance(value, np.generic): + value = value.item() + if isinstance(value, float) and np.isnan(value): + return None + return value + + +def _nan_to_none(obj: Any) -> Any: + """Recursively map float NaN to None in (nested) lists from ``ndarray.tolist()``.""" + if isinstance(obj, list): + return [_nan_to_none(v) for v in obj] + if isinstance(obj, float) and np.isnan(obj): + return None + return obj + + +def encode_surface(surface: EventStudyResults) -> Dict[str, Any]: + return {f.name: _encode_value(getattr(surface, f.name)) for f in dataclasses.fields(surface)} + + +def capture_lane(name: str) -> Dict[str, Any]: + lane = LANES[name] + df = PANELS[lane.panel]() + sd = lane.survey() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # MPD's FutureWarning + rank-deficiency notices + twfe = TwoWayFixedEffects(**lane.twfe_ctor).fit( + df, event_study=True, survey_design=sd, **lane.twfe_fit + ) + mpd = _from_mpd(MultiPeriodDiD(**lane.mpd_ctor).fit(df, survey_design=sd, **lane.mpd_fit)) + assert isinstance(twfe, EventStudyResults) + return {"echo": lane.echo(), "twfe": encode_surface(twfe), "mpd_twin": encode_surface(mpd)} + + +def dump_fixture(obj: Any) -> str: + """Sorted keys, one indent level per nesting level, every leaf container + (a dict or list of scalars) on one line - the aggregate-oracle format.""" + + def _fmt(o: Any, ind: int = 0) -> str: + pad = " " * ind + if isinstance(o, dict): + if all(not isinstance(v, (dict, list)) for v in o.values()): + return json.dumps(o, sort_keys=True) + items = [f"{pad} {json.dumps(k)}: {_fmt(v, ind + 1)}" for k, v in sorted(o.items())] + return "{\n" + ",\n".join(items) + f"\n{pad}}}" + if isinstance(o, list): + if all(not isinstance(v, (dict, list)) for v in o): + return json.dumps(o) + items = [f"{pad} {_fmt(v, ind + 1)}" for v in o] + return "[\n" + ",\n".join(items) + f"\n{pad}]" + return json.dumps(o) + + return _fmt(obj) + "\n" + + +def _git_head() -> str: + try: + return subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=HERE, text=True, stderr=subprocess.DEVNULL + ).strip() + except Exception: # noqa: BLE001 + return "unknown" + + +if __name__ == "__main__": + import platform + + import pandas as pd + import scipy + + import diff_diff + from diff_diff.bootstrap_chunking import effective_weight_backend + + payload = { + "provenance": { + "captured_at_head": _git_head(), + "diff_diff_version": diff_diff.__version__, + "backend": effective_weight_backend(), + "numpy": np.__version__, + "scipy": scipy.__version__, + "pandas": pd.__version__, + "python": platform.python_version(), + }, + "lanes": {name: capture_lane(name) for name in LANES}, + } + sys.stdout.write(dump_fixture(payload)) diff --git a/tests/fixtures/README.md b/tests/fixtures/README.md index fc90e4186..ce71b9df6 100644 --- a/tests/fixtures/README.md +++ b/tests/fixtures/README.md @@ -65,3 +65,34 @@ against it, so a surface that mis-marks an estimated horizon as reference fails instead of having that row skipped. Re-capturing on a post-removal tree is NOT a substitute: the fit-time route is gone, so nothing independent remains to disagree with. + +## `v4_twfe_event_study_oracles.json` - TWFE event-study decoupling oracles (4.0 PR 4c-i) + +Numeric reference for `tests/test_twfe_event_study_decoupling.py`, captured by +`tests/_capture_v4_merge_mpd_oracles.py` on the UNMODIFIED tree at the commit +named in `provenance.captured_at_head` (`68e383bd`, origin/main at the PR 4c-i +branch cut) - i.e. while `TwoWayFixedEffects.fit(..., event_study=True)` still +built its surface through an internal `MultiPeriodDiDResults` and +`results_base._from_mpd`, and while `MultiPeriodDiD` still existed as the +independent twin of every lane. Layout: + +- `provenance` - capture SHA, `diff_diff.__version__`, weight backend, + numpy/scipy/pandas/python versions. +- `lanes.` - one entry per lane of + `tests/helpers/twfe_es_oracle_lanes.py::LANES` (pooled default / explicit + cluster / hc2_bm; within + explicit cluster / hc2 carve-out / Conley / + survey TSL / survey JK1 replicates; string period labels; a non-suffix + partition; a rank-deficient NaN post period). Each carries `echo` (the + lane's panel name, constructor and fit kwargs and survey label - the test + asserts it equals the live lane, so an edited lane fails as lane drift, not + as a numeric regression), `twfe` (every `dataclasses.fields(EventStudyResults)` + value of the TWFE event-study surface) and `mpd_twin` (the same for + `_from_mpd(MultiPeriodDiD(...).fit(...))`). + +Arrays are stored as `{"dtype": , "values": [...]}` (the +real dtype, e.g. `" Tuple[float, float]: + """Return ``(att_bar, se_bar)`` over ``es.post_periods``. + + ``att_bar`` is the equal-weight mean of the post-period ``att`` rows; + ``se_bar`` is ``sqrt(sum of the post-period vcov block) / n_post`` + (computed as ``sqrt(sum / n_post**2)``, MultiPeriodDiD's expression). + """ + if es.post_periods is None: + raise ValueError("post_period_average requires a surface with post_periods provenance") + post_rows = np.isin(es.event_time, list(es.post_periods)) + post_att = np.asarray(es.att, dtype=float)[post_rows] + if post_att.size == 0: + raise ValueError("post_period_average: no post-period rows on the surface") + if np.any(np.isnan(post_att)): + return float("nan"), float("nan") + att_bar = float(np.mean(post_att)) + if es.vcov is None or es.vcov_index is None: + return att_bar, float("nan") + post_cols = np.isin(es.vcov_index, list(es.post_periods)) + n_post = int(post_cols.sum()) + if n_post == 0: + return att_bar, float("nan") + avg_var = float(np.sum(np.asarray(es.vcov, dtype=float)[np.ix_(post_cols, post_cols)])) / ( + n_post**2 + ) + if np.isnan(avg_var) or avg_var < 0: + return att_bar, float("nan") + return att_bar, float(np.sqrt(avg_var)) diff --git a/tests/helpers/event_study_surface_builders.py b/tests/helpers/event_study_surface_builders.py new file mode 100644 index 000000000..7f1cb4d5d --- /dev/null +++ b/tests/helpers/event_study_surface_builders.py @@ -0,0 +1,106 @@ +"""Hand-built calendar ``EventStudyResults`` for consumer tests (4.0 PR 4c-i). + +Replaces the hand-built ``MultiPeriodDiDResults`` / ``PeriodEffect`` mocks of +the MultiPeriodDiD-fixture tests with the surface the TWFE event-study +producer emits: one row per calendar label, the reference row marked +``is_reference`` (NaN inference), ``time_scale="calendar"``, +``source="TwoWayFixedEffects"`` and the authoritative ``post_periods`` +partition. The calendar consumer route requires exactly one reference row, +which is the last pre-period, so callers supply ``reference`` explicitly. +""" + +from __future__ import annotations + +from typing import Any, Dict, Mapping, Optional, Sequence, Tuple, Union + +import numpy as np + +from diff_diff.results_base import EventStudyResults +from diff_diff.utils import safe_inference + +Row = Union[ + Tuple[float, float], Tuple[float, float, float, float, Tuple[float, float]], Dict[str, Any] +] + + +def _row_fields(row: Row, alpha: float = 0.05) -> Tuple[float, float, float, float, float, float]: + """(effect, se, t_stat, p_value, ci_lo, ci_hi) from a terse row spec. + + A 2-tuple ``(effect, se)`` fills t/p/CI through ``safe_inference`` (the + library's own helper, so se <= 0 / non-finite rows fail closed to all-NaN + inference and ``alpha`` sets the interval level); a 5-tuple + ``(effect, se, t_stat, p_value, (lo, hi))`` and a dict with those keys are + deliberate literals taken verbatim. + """ + if isinstance(row, dict): + lo, hi = row.get("conf_int", (np.nan, np.nan)) + return ( + float(row["effect"]), + float(row["se"]), + float(row.get("t_stat", np.nan)), + float(row.get("p_value", np.nan)), + float(lo), + float(hi), + ) + if len(row) == 2: + effect, se = float(row[0]), float(row[1]) + t, p, (lo, hi) = safe_inference(effect, se, alpha=alpha) + return effect, se, float(t), float(p), float(lo), float(hi) + effect, se, t, p, ci = row # type: ignore[misc] + return float(effect), float(se), float(t), float(p), float(ci[0]), float(ci[1]) + + +def calendar_surface( + period_effects: Mapping[Any, Row], + *, + reference: Any, + post_periods: Sequence[Any], + vcov: Optional[np.ndarray] = None, + vcov_index: Optional[Sequence[Any]] = None, + source: str = "TwoWayFixedEffects", + alpha: float = 0.05, + df_survey: Optional[float] = None, +) -> EventStudyResults: + """Build a calendar surface from ``{label: row}`` (non-reference labels only). + + ``vcov`` is the event-study sub-block over ``vcov_index`` (default: the + sorted non-reference labels, i.e. the same order as the rows). + """ + labels = sorted(period_effects) + if reference in labels: + raise ValueError("reference must not be among the estimated labels") + all_labels = sorted(labels + [reference]) + n = len(all_labels) + att = np.full(n, np.nan) + se = np.full(n, np.nan) + t = np.full(n, np.nan) + p = np.full(n, np.nan) + lo = np.full(n, np.nan) + hi = np.full(n, np.nan) + is_ref = np.zeros(n, dtype=bool) + for i, lab in enumerate(all_labels): + if lab == reference: + is_ref[i] = True + continue + att[i], se[i], t[i], p[i], lo[i], hi[i] = _row_fields(period_effects[lab], alpha=alpha) + if vcov is not None and vcov_index is None: + vcov_index = labels + return EventStudyResults( + event_time=np.asarray(all_labels), + att=att, + se=se, + t_stat=t, + p_value=p, + conf_int_lower=lo, + conf_int_upper=hi, + is_reference=is_ref, + n=np.full(n, np.nan), + reference_period=reference, + time_scale="calendar", + vcov=None if vcov is None else np.asarray(vcov, dtype=float), + vcov_index=None if vcov_index is None else np.asarray(list(vcov_index)), + alpha=alpha, + source=source, + post_periods=tuple(post_periods), + df_survey=df_survey, + ) diff --git a/tests/helpers/twfe_es_oracle_lanes.py b/tests/helpers/twfe_es_oracle_lanes.py new file mode 100644 index 000000000..cc50070be --- /dev/null +++ b/tests/helpers/twfe_es_oracle_lanes.py @@ -0,0 +1,223 @@ +"""Shared lane definitions for the TWFE event-study decoupling oracle (4.0 PR 4c-i). + +Single owner of the oracle PANELS and LANES used by BOTH +``tests/_capture_v4_merge_mpd_oracles.py`` (which writes +``tests/fixtures/v4_twfe_event_study_oracles.json`` on the unmodified tree) and +``tests/test_twfe_event_study_decoupling.py`` (which refits the lanes after the +refactor and compares). Keeping panel construction and lane kwargs here means +the tests rebuild exactly what the capture fitted; the fixture carries each +lane's ``echo()`` so an edited lane fails as "lane drift", not as a numeric +regression. + +Why these lanes: ``TwoWayFixedEffects.fit(..., event_study=True)`` and the +deprecated ``MultiPeriodDiD`` share ``_fit_event_study_core``; each lane pairs +a TWFE event-study fit with the MPD design that is its bit-exact (pooled) or +1e-12 (within) twin, covering every variance lane the core offers (HC1 +default, explicit cluster, hc2 carve-out, hc2_bm per-period BM df, Conley, +survey TSL, survey replicate, string labels, a non-suffix partition and a +rank-deficient NaN post period). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional + +import numpy as np +import pandas as pd + +from diff_diff import SurveyDesign, generate_did_data + +PANEL_SEED = 20261024 +NAN_PANEL_SEED = 42 +JK_REPLICATES = 15 +JK_COLUMNS: List[str] = [f"rep_{r}" for r in range(JK_REPLICATES)] + + +def build_panel() -> pd.DataFrame: + """The main oracle panel: ``generate_did_data`` plus the derived columns every + lane may need (coordinates for Conley, survey weights/strata/PSU, JK1 + replicate weights, a string copy of the period label).""" + df = generate_did_data( + n_units=60, + n_periods=8, + treatment_period=4, + treatment_effect=2.0, + seed=PANEL_SEED, + ).copy() + df["lat"] = 40.0 + 0.01 * df["unit"] + df["lon"] = -100.0 + 0.01 * (df["unit"] % 7) + rng = np.random.default_rng(1) + df["w"] = rng.uniform(0.5, 2.0, size=len(df)) # per-row sampling weight + df["stratum"] = df["unit"] % 4 + df["psu"] = df["unit"] + # JK1 delete-a-group-of-PSUs replicate weights: 60 PSUs in 15 groups of 4 + # (the tests/test_replicate_weight_expansion.py recipe); retained rows are + # rescaled by n_rep / (n_rep - 1). + units = sorted(df["unit"].unique()) + group = len(units) // JK_REPLICATES + for r, col in enumerate(JK_COLUMNS): + deleted = list(units[r * group : (r + 1) * group]) + w_r = df["w"].to_numpy(copy=True) + mask = df["unit"].isin(deleted).to_numpy() + w_r[mask] = 0.0 + w_r[~mask] *= JK_REPLICATES / (JK_REPLICATES - 1) + df[col] = w_r + df["period_s"] = "p" + df["period"].astype(str) + return df + + +def build_nan_post_panel() -> pd.DataFrame: + """The ``test_avg_att_nan_when_period_effect_nan`` design: 100 units x 6 + periods (treated = first half, effect 3.0 from period 3), then EVERY treated + row in post period 3 dropped so the ``treated:period_3`` interaction is a + zero column and that period's effect is NaN.""" + rng = np.random.RandomState(NAN_PANEL_SEED) + n_units = 100 + rows = [] + for unit in range(n_units): + is_treated = unit < n_units // 2 + unit_effect = rng.normal(0, 1) + for period in range(6): + y = 10.0 + unit_effect + period * 0.5 + if is_treated and period >= 3: + y += 3.0 + y += rng.normal(0, 0.5) + rows.append({"unit": unit, "period": period, "treated": int(is_treated), "outcome": y}) + df = pd.DataFrame(rows) + keep = ~((df["treated"] == 1) & (df["period"] == 3)) + return pd.DataFrame(df.loc[keep]).reset_index(drop=True) + + +PANELS: Dict[str, Callable[[], pd.DataFrame]] = { + "main": build_panel, + "nan_post": build_nan_post_panel, +} + + +def build_survey(label: Optional[str]) -> Optional[SurveyDesign]: + """Materialize a lane's survey design from its label (None -> no design).""" + if label is None: + return None + if label == "tsl:w/stratum/psu": + return SurveyDesign(weights="w", strata="stratum", psu="psu") + if label == "jk1:w": + return SurveyDesign(weights="w", replicate_weights=JK_COLUMNS, replicate_method="JK1") + raise ValueError(f"unknown survey label {label!r}") + + +@dataclass(frozen=True) +class Lane: + """One oracle lane: the TWFE event-study fit and its MultiPeriodDiD twin. + + Fit kwargs never contain objects; the survey design is built from + ``survey_label`` at fit time. ``panel`` names a ``PANELS`` entry. + """ + + panel: str + twfe_ctor: Dict[str, Any] + twfe_fit: Dict[str, Any] + mpd_ctor: Dict[str, Any] + mpd_fit: Dict[str, Any] + survey_label: Optional[str] = None + notes: str = field(default="", compare=False) + + def echo(self) -> Dict[str, Any]: + """JSON-safe record of everything that defines the lane (drift pin).""" + return { + "panel": self.panel, + "twfe_ctor": dict(self.twfe_ctor), + "twfe_fit": dict(self.twfe_fit), + "mpd_ctor": dict(self.mpd_ctor), + "mpd_fit": dict(self.mpd_fit), + "survey_label": self.survey_label, + } + + def survey(self) -> Optional[SurveyDesign]: + return build_survey(self.survey_label) + + +_COMMON = {"outcome": "outcome", "treatment": "treated", "time": "period"} +_PART = {"post_periods": [4, 5, 6, 7], "reference_period": 3} +_CONLEY_CTOR = { + "vcov_type": "conley", + "conley_coords": ["lat", "lon"], + "conley_cutoff_km": 50.0, + "conley_lag_cutoff": 2, +} + + +def _pooled(ctor: Dict[str, Any], **fit_extra: Any) -> Lane: + return Lane( + panel="main", + twfe_ctor=dict(ctor), + twfe_fit={**_COMMON, **_PART, "spec": "pooled", **fit_extra}, + mpd_ctor=dict(ctor), + mpd_fit={**_COMMON, **_PART, **fit_extra}, + ) + + +def _within(ctor: Dict[str, Any], survey_label: Optional[str] = None, **mpd_extra: Any) -> Lane: + return Lane( + panel="main", + twfe_ctor=dict(ctor), + twfe_fit={**_COMMON, **_PART, "spec": "within", "unit": "unit"}, + mpd_ctor=dict(ctor), + mpd_fit={**_COMMON, **_PART, "absorb": ["unit"], **mpd_extra}, + survey_label=survey_label, + ) + + +LANES: Dict[str, Lane] = { + "pooled_default": _pooled({}), + "pooled_cluster_unit": _pooled({"cluster": "unit"}), + "pooled_hc2_bm": _pooled({"vcov_type": "hc2_bm"}), + "within_cluster_unit": _within({"cluster": "unit"}), + "within_hc2": _within({"vcov_type": "hc2"}), + "within_conley": _within(_CONLEY_CTOR, unit="unit"), + "within_survey_tsl": _within({}, survey_label="tsl:w/stratum/psu"), + "within_replicate_jk1": _within({}, survey_label="jk1:w"), + "pooled_string_labels": Lane( + panel="main", + twfe_ctor={}, + twfe_fit={ + "outcome": "outcome", + "treatment": "treated", + "time": "period_s", + "spec": "pooled", + "post_periods": ["p4", "p5", "p6", "p7"], + "reference_period": "p3", + }, + mpd_ctor={}, + mpd_fit={ + "outcome": "outcome", + "treatment": "treated", + "time": "period_s", + "post_periods": ["p4", "p5", "p6", "p7"], + "reference_period": "p3", + }, + ), + "pooled_non_suffix": Lane( + panel="main", + twfe_ctor={}, + twfe_fit={**_COMMON, "spec": "pooled", "post_periods": [2, 5], "reference_period": 4}, + mpd_ctor={}, + mpd_fit={**_COMMON, "post_periods": [2, 5], "reference_period": 4}, + ), + "pooled_nan_post_period": Lane( + panel="nan_post", + twfe_ctor={}, + twfe_fit={**_COMMON, "spec": "pooled", "post_periods": [3, 4, 5], "reference_period": 2}, + mpd_ctor={}, + mpd_fit={**_COMMON, "post_periods": [3, 4, 5], "reference_period": 2}, + ), +} + +#: Pooled lanes whose post-period average is finite (the NaN lane is excluded: +#: ``nan == nan`` is False; its contract is pinned NaN-aware separately). +FINITE_POOLED_LANES = [ + name + for name, lane in LANES.items() + if lane.twfe_fit.get("spec") == "pooled" and name != "pooled_nan_post_period" +] +WITHIN_LANES = [name for name, lane in LANES.items() if lane.twfe_fit.get("spec") == "within"] diff --git a/tests/test_business_report.py b/tests/test_business_report.py index e17732e65..e5d5feb73 100644 --- a/tests/test_business_report.py +++ b/tests/test_business_report.py @@ -39,6 +39,7 @@ MultiPeriodDiD, SyntheticControl, SyntheticDiD, + TwoWayFixedEffects, generate_did_data, generate_factor_data, generate_staggered_data, @@ -77,7 +78,9 @@ def did_fit(): @pytest.fixture(scope="module") -def event_study_fit(): +def legacy_mpd_fit(): + """Native MultiPeriodDiD fit for the MPD-only tests (event-study lanes, the + name-string dispatch pins): dies with the class at the 4.0 cut (PR 4c-ii).""" df = generate_did_data(n_units=80, n_periods=8, treatment_effect=1.5, seed=7) es = MultiPeriodDiD().fit( df, @@ -90,6 +93,17 @@ def event_study_fit(): return es, df +@pytest.fixture(scope="module") +def twfe_fit(): + """Static TwoWayFixedEffects fit (a scalar DiDResults) on the same panel: the + successor fixture for the generic report tests (4.0 cut, PR 4c-i).""" + df = generate_did_data(n_units=80, n_periods=8, treatment_effect=1.5, seed=7) + fit = TwoWayFixedEffects().fit( + df, outcome="outcome", treatment="treated", unit="unit", post="post" + ) + return fit, df + + @pytest.fixture(scope="module") def cs_fit(): sdf = generate_staggered_data(n_units=100, n_periods=6, treatment_effect=1.5, seed=7) @@ -156,13 +170,13 @@ def edid_fit(): # Schema contract # --------------------------------------------------------------------------- class TestSchemaContract: - def test_top_level_keys(self, event_study_fit): - fit, _ = event_study_fit + def test_top_level_keys(self, twfe_fit): + fit, _ = twfe_fit br = BusinessReport(fit, auto_diagnostics=False) assert set(br.to_dict().keys()) == _BR_TOP_LEVEL_KEYS - def test_schema_version(self, event_study_fit): - fit, _ = event_study_fit + def test_schema_version(self, twfe_fit): + fit, _ = twfe_fit assert ( BusinessReport(fit, auto_diagnostics=False).to_dict()["schema_version"] == BUSINESS_REPORT_SCHEMA_VERSION @@ -191,22 +205,22 @@ def test_json_round_trip_sdid(self, sdid_fit): # BR ↔ DR integration # --------------------------------------------------------------------------- class TestDiagnosticsIntegration: - def test_auto_diagnostics_true_populates_diagnostics_block(self, event_study_fit): - fit, _ = event_study_fit + def test_auto_diagnostics_true_populates_diagnostics_block(self, twfe_fit): + fit, _ = twfe_fit br = BusinessReport(fit, auto_diagnostics=True) d = br.to_dict() assert d["diagnostics"]["status"] == "ran" assert "schema" in d["diagnostics"] - def test_auto_diagnostics_false_skips(self, event_study_fit): - fit, _ = event_study_fit + def test_auto_diagnostics_false_skips(self, twfe_fit): + fit, _ = twfe_fit br = BusinessReport(fit, auto_diagnostics=False) d = br.to_dict() assert d["diagnostics"]["status"] == "skipped" assert "auto_diagnostics=False" in d["diagnostics"]["reason"] - def test_explicit_diagnostics_results_takes_precedence(self, event_study_fit): - fit, _ = event_study_fit + def test_explicit_diagnostics_results_takes_precedence(self, twfe_fit): + fit, _ = twfe_fit dr = DiagnosticReport(fit) dr_results = dr.run_all() br = BusinessReport(fit, diagnostics=dr_results) @@ -215,14 +229,14 @@ def test_explicit_diagnostics_results_takes_precedence(self, event_study_fit): # Same dict identity shows the supplied results were used verbatim. assert d["diagnostics"]["schema"] is dr_results.schema - def test_explicit_diagnostics_report_runs(self, event_study_fit): - fit, _ = event_study_fit + def test_explicit_diagnostics_report_runs(self, twfe_fit): + fit, _ = twfe_fit dr = DiagnosticReport(fit) br = BusinessReport(fit, diagnostics=dr) assert br.to_dict()["diagnostics"]["status"] == "ran" - def test_diagnostics_wrong_type_raises(self, event_study_fit): - fit, _ = event_study_fit + def test_diagnostics_wrong_type_raises(self, twfe_fit): + fit, _ = twfe_fit with pytest.raises(TypeError): BusinessReport(fit, diagnostics="not a DR") # type: ignore[arg-type] @@ -231,8 +245,8 @@ def test_diagnostics_wrong_type_raises(self, event_study_fit): # HonestDiD passthrough # --------------------------------------------------------------------------- class TestHonestDiDPassthrough: - def test_supplied_sensitivity_is_not_recomputed(self, event_study_fit): - fit, _ = event_study_fit + def test_supplied_sensitivity_is_not_recomputed(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit class _FakeSens: M_values = np.array([0.5, 1.0]) @@ -336,8 +350,8 @@ def test_direction_none_uses_neutral_verb(self, cs_fit): class TestWarningsPassthrough: """Broad exception handling still records provenance in schema.warnings.""" - def test_diagnostic_error_surfaces_as_top_level_warning(self, event_study_fit): - fit, _ = event_study_fit + def test_diagnostic_error_surfaces_as_top_level_warning(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit def _raise(*args, **kwargs): raise RuntimeError("synthetic test failure") @@ -361,30 +375,30 @@ def test_high_significance_produces_strong_language(self, cs_fit): summary = br.summary() assert "strongly supported" in summary - def test_near_threshold_caveat(self, event_study_fit): + def test_near_threshold_caveat(self, twfe_fit): """Fabricate a p-value near 0.05 to exercise the significance-chasing guard.""" - fit, _ = event_study_fit + fit, _ = twfe_fit # Monkey-patch the result to land p_value in (0.04, 0.051). - original = fit.avg_p_value + original = fit.p_value try: - fit.avg_p_value = 0.045 + fit.p_value = 0.045 br = BusinessReport(fit, auto_diagnostics=False) caveats = br.caveats() topics = {c.get("topic") for c in caveats} assert "near_significance" in topics finally: - fit.avg_p_value = original + fit.p_value = original - def test_far_from_threshold_no_near_caveat(self, event_study_fit): - fit, _ = event_study_fit - original = fit.avg_p_value + def test_far_from_threshold_no_near_caveat(self, twfe_fit): + fit, _ = twfe_fit + original = fit.p_value try: - fit.avg_p_value = 0.010 + fit.p_value = 0.010 br = BusinessReport(fit, auto_diagnostics=False) topics = {c.get("topic") for c in br.caveats()} assert "near_significance" not in topics finally: - fit.avg_p_value = original + fit.p_value = original # --------------------------------------------------------------------------- @@ -393,8 +407,8 @@ def test_far_from_threshold_no_near_caveat(self, event_study_fit): class TestPreTrendsVerdictPhrasing: """Verdict and tier should flow through into schema AND phrasing.""" - def test_verdict_and_tier_surface_in_schema(self, event_study_fit): - fit, _ = event_study_fit + def test_verdict_and_tier_surface_in_schema(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit br = BusinessReport(fit, auto_diagnostics=True) pt = br.to_dict()["pre_trends"] # This fixture has a clear violation and an underpowered test — both set. @@ -405,8 +419,8 @@ def test_verdict_and_tier_surface_in_schema(self, event_study_fit): "clear_violation", } - def test_clear_violation_phrased_tentatively(self, event_study_fit): - fit, _ = event_study_fit + def test_clear_violation_phrased_tentatively(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit br = BusinessReport(fit, auto_diagnostics=True) if br.to_dict()["pre_trends"].get("verdict") == "clear_violation": summary = br.summary() @@ -445,8 +459,8 @@ def test_underpowered_phrasing_uses_hedge_language(self, cs_fit): # NaN ATT # --------------------------------------------------------------------------- class TestNaNATT: - def test_nan_att_produces_caveat_and_does_not_crash(self, event_study_fit): - fit, _ = event_study_fit + def test_nan_att_produces_caveat_and_does_not_crash(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit original = fit.avg_att try: fit.avg_att = float("nan") @@ -459,25 +473,40 @@ def test_nan_att_produces_caveat_and_does_not_crash(self, event_study_fit): finally: fit.avg_att = original + def test_nan_att_produces_caveat_and_does_not_crash_twfe(self, twfe_fit): + """Static-fit twin (4.0 cut, PR 4c-i): the headline field is `att`.""" + fit, _ = twfe_fit + original = fit.att + try: + fit.att = float("nan") + br = BusinessReport(fit, auto_diagnostics=False) + summary = br.summary() + caveats = br.caveats() + assert isinstance(summary, str) + assert any(c.get("topic") == "estimation_failure" for c in caveats) + assert br.to_dict()["headline"]["sign"] == "undefined" + finally: + fit.att = original + # --------------------------------------------------------------------------- # include_appendix toggle # --------------------------------------------------------------------------- class TestAppendix: - def test_include_appendix_true_embeds_summary(self, event_study_fit): - fit, _ = event_study_fit + def test_include_appendix_true_embeds_summary(self, twfe_fit): + fit, _ = twfe_fit br = BusinessReport(fit, auto_diagnostics=False, include_appendix=True) md = br.full_report() assert "## Technical Appendix" in md - def test_include_appendix_false_omits(self, event_study_fit): - fit, _ = event_study_fit + def test_include_appendix_false_omits(self, twfe_fit): + fit, _ = twfe_fit br = BusinessReport(fit, auto_diagnostics=False, include_appendix=False) md = br.full_report() assert "## Technical Appendix" not in md - def test_appendix_summary_failure_is_visible(self, event_study_fit): - fit, _ = event_study_fit + def test_appendix_summary_failure_is_visible(self, twfe_fit): + fit, _ = twfe_fit with patch.object(fit, "summary", side_effect=RuntimeError("internal path /tmp/private")): br = BusinessReport(fit, auto_diagnostics=False, include_appendix=True) md = br.full_report() @@ -486,8 +515,8 @@ def test_appendix_summary_failure_is_visible(self, event_study_fit): assert "Technical appendix unavailable" in md assert "RuntimeError" in md - def test_appendix_summary_failure_does_not_leak_exception_details(self, event_study_fit): - fit, _ = event_study_fit + def test_appendix_summary_failure_does_not_leak_exception_details(self, twfe_fit): + fit, _ = twfe_fit with patch.object(fit, "summary", side_effect=RuntimeError("secret-token-123")): br = BusinessReport(fit, auto_diagnostics=False, include_appendix=True) md = br.full_report() @@ -495,12 +524,12 @@ def test_appendix_summary_failure_does_not_leak_exception_details(self, event_st assert "secret-token-123" not in md assert "Traceback" not in md - def test_appendix_stringification_failure_is_visible(self, event_study_fit): + def test_appendix_stringification_failure_is_visible(self, twfe_fit): class _BadSummary: def __str__(self): raise ValueError("sensitive-render-context") - fit, _ = event_study_fit + fit, _ = twfe_fit with patch.object(fit, "summary", return_value=_BadSummary()): br = BusinessReport(fit, auto_diagnostics=False, include_appendix=True) md = br.full_report() @@ -534,9 +563,9 @@ def test_survey_absent_yields_null_survey_block(self, cs_fit): d = br.to_dict() assert d["sample"]["survey"] is None - def test_survey_present_populates_block(self, event_study_fit): + def test_survey_present_populates_block(self, twfe_fit): """Synthetically attach a survey_metadata shim and verify BR surfaces it.""" - fit, _ = event_study_fit + fit, _ = twfe_fit class _ShimMeta: weight_type = "pweight" @@ -573,14 +602,14 @@ class _ShimMeta: # Single-knob alpha # --------------------------------------------------------------------------- class TestAlphaKnob: - def test_alpha_equal_to_result_alpha_drives_ci_level(self, event_study_fit): + def test_alpha_equal_to_result_alpha_drives_ci_level(self, twfe_fit): """When caller's alpha matches the fit's native alpha, ``ci_level`` reflects that alpha (e.g., alpha=0.05 -> 95% CI).""" - fit, _ = event_study_fit + fit, _ = twfe_fit br = BusinessReport(fit, alpha=0.05, auto_diagnostics=False) assert br.to_dict()["headline"]["ci_level"] == 95 - def test_alpha_mismatch_preserves_fitted_ci_at_native_level(self, event_study_fit): + def test_alpha_mismatch_preserves_fitted_ci_at_native_level(self, twfe_fit): """Round-7 regression: a caller alpha that differs from the fit's native alpha must NOT recompute a z-based CI (the fit used t-based inference with a finite ``df`` that BR cannot reproduce from @@ -590,7 +619,7 @@ def test_alpha_mismatch_preserves_fitted_ci_at_native_level(self, event_study_fi """ import math - fit, _ = event_study_fit + fit, _ = twfe_fit br95 = BusinessReport(fit, alpha=0.05, auto_diagnostics=False) br90 = BusinessReport(fit, alpha=0.10, auto_diagnostics=False) h95 = br95.to_dict()["headline"] @@ -3806,8 +3835,8 @@ def _fake_single_m(M=1.5, ci_lb=1.0, ci_ub=3.0): alpha=0.05, ) - def test_full_report_does_not_claim_full_grid_for_single_m(self, event_study_fit): - fit, _ = event_study_fit + def test_full_report_does_not_claim_full_grid_for_single_m(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit br = BusinessReport(fit, honest_did_results=self._fake_single_m()) md = br.full_report() assert "robust across full grid" not in md @@ -3818,8 +3847,8 @@ def test_full_report_does_not_claim_full_grid_for_single_m(self, event_study_fit # Summary + full_report work across estimators # --------------------------------------------------------------------------- class TestAcrossEstimators: - def test_summary_nonempty_for_all(self, did_fit, event_study_fit, cs_fit, sdid_fit): - for fit, _ in (did_fit, event_study_fit, cs_fit, sdid_fit): + def test_summary_nonempty_for_all(self, did_fit, legacy_mpd_fit, cs_fit, sdid_fit): + for fit, _ in (did_fit, legacy_mpd_fit, cs_fit, sdid_fit): br = BusinessReport(fit, auto_diagnostics=False) s = br.summary() assert isinstance(s, str) diff --git a/tests/test_conley_vcov.py b/tests/test_conley_vcov.py index a934a15a2..2d3b64a85 100644 --- a/tests/test_conley_vcov.py +++ b/tests/test_conley_vcov.py @@ -1128,13 +1128,13 @@ def test_did_conley_unknown_cluster_column_raises(self, two_period_panel): unit="unit", ) - def test_mpd_conley_unknown_cluster_column_raises(self): - """MPD + Conley + cluster= raises a clear estimator- + def test_twfe_es_conley_unknown_cluster_column_raises(self): + """TWFE event-study + Conley + cluster= raises a clear estimator- level ValueError before `data[self.cluster]` access (combined-kernel path; codex CI R7 P1).""" import pandas as _pd - from diff_diff import MultiPeriodDiD + from diff_diff import TwoWayFixedEffects rng = np.random.default_rng(seed=83) rows = [] @@ -1154,7 +1154,7 @@ def test_mpd_conley_unknown_cluster_column_raises(self): ) df = _pd.DataFrame(rows) with pytest.raises(ValueError, match="Cluster column 'missing_region' not found"): - MultiPeriodDiD( + TwoWayFixedEffects( vcov_type="conley", cluster="missing_region", conley_coords=("lat", "lon"), @@ -1163,6 +1163,8 @@ def test_mpd_conley_unknown_cluster_column_raises(self): ).fit( df, outcome="y", + event_study=True, + spec="pooled", treatment="treated", time="time", unit="unit", @@ -1312,11 +1314,11 @@ def test_did_conley_missing_lag_cutoff_raises(self, two_period_panel): unit="unit", ) - def test_did_conley_matches_mpd_post_periods_1(self, two_period_panel): - """DiD + Conley on a 2-period panel matches MultiPeriodDiD with + def test_did_conley_matches_twfe_event_study_post_periods_1(self, two_period_panel): + """DiD + Conley on a 2-period panel matches the TWFE event-study mode with post_periods=[1], reference_period=0 on the same data (locks the - DiD wire-up correctness against the already-shipped MPD path).""" - from diff_diff import DifferenceInDifferences, MultiPeriodDiD + DiD wire-up correctness against the event-study core).""" + from diff_diff import DifferenceInDifferences, TwoWayFixedEffects df = two_period_panel.copy() kwargs = dict( @@ -1328,18 +1330,22 @@ def test_did_conley_matches_mpd_post_periods_1(self, two_period_panel): res_did = DifferenceInDifferences(**kwargs).fit( df, outcome="y", treatment="treated", post="time", unit="unit" ) - res_mpd = MultiPeriodDiD(**kwargs).fit( + res_es = TwoWayFixedEffects(**kwargs).fit( df, outcome="y", + event_study=True, + spec="pooled", treatment="treated", time="time", unit="unit", post_periods=[1], reference_period=0, ) - # MPD reports ATT for the single post period (1) - np.testing.assert_allclose(res_did.att, res_mpd.att, atol=1e-10) - np.testing.assert_allclose(res_did.se, res_mpd.se, atol=1e-10) + # The event study has exactly one non-reference row: post period 1. + (row,) = np.where(~res_es.is_reference)[0] + assert res_es.event_time[row] == 1 + np.testing.assert_allclose(res_did.att, res_es.att[row], atol=1e-10) + np.testing.assert_allclose(res_did.se, res_es.se[row], atol=1e-10) def test_did_conley_with_absorb_uses_raw_time_labels(self, two_period_panel, monkeypatch): """DiD + Conley + absorb=[] must feed the Conley helper the @@ -1389,8 +1395,8 @@ def _spy(*args, **kwargs): ) assert set(uniques.tolist()) == {0, 1}, f"Expected raw integer labels 0/1, got {uniques}" - def test_mpd_conley_with_absorb_uses_raw_coords_and_time(self, monkeypatch): - """MultiPeriodDiD + Conley + absorb=[] must feed the Conley + def test_twfe_es_within_conley_uses_raw_coords_and_time(self, monkeypatch): + """TWFE event-study (within spec) + Conley must feed the Conley helper the ORIGINAL coords/time/unit columns from `data`, not the absorb-demeaned `working_data`. If a user lists the `time` column (or coord columns) in `absorb`, working_data has those demeaned and @@ -1401,7 +1407,7 @@ def test_mpd_conley_with_absorb_uses_raw_coords_and_time(self, monkeypatch): import pandas as _pd import diff_diff.linalg as linalg_module - from diff_diff import MultiPeriodDiD + from diff_diff import TwoWayFixedEffects rng = np.random.default_rng(seed=71) n_units = 10 @@ -1438,7 +1444,7 @@ def _spy(*args, **kwargs): return orig(*args, **kwargs) monkeypatch.setattr(linalg_module, "_compute_conley_vcov", _spy) - MultiPeriodDiD( + TwoWayFixedEffects( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, @@ -1446,12 +1452,13 @@ def _spy(*args, **kwargs): ).fit( df, outcome="y", + event_study=True, + spec="within", treatment="treated", time="time", unit="unit", post_periods=[2, 3], reference_period=0, - absorb=["unit"], ) # Raw time labels span exactly 0..T-1 = 4 distinct values; demeaned # absorb would collapse to per-unit means → ~n_units distinct values. @@ -1475,11 +1482,11 @@ def _spy(*args, **kwargs): "absorb=['unit'] is leaking demeaned coords into the Conley helper." ) - def test_multi_period_did_with_conley_panel(self): - """Phase 2 MultiPeriodDiD + vcov_type='conley' uses the block-decomposed + def test_twfe_event_study_with_conley_panel(self): + """TWFE event-study mode + vcov_type='conley' uses the block-decomposed sandwich (matches R conleyreg). Verifies that finite SEs are produced when conley_lag_cutoff and unit are both supplied.""" - from diff_diff import MultiPeriodDiD + from diff_diff import TwoWayFixedEffects rng = np.random.default_rng(seed=13) n_units = 30 @@ -1496,7 +1503,7 @@ def test_multi_period_did_with_conley_panel(self): import pandas as pd df_mp = pd.DataFrame(rows) - res = MultiPeriodDiD( + res = TwoWayFixedEffects( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, @@ -1504,17 +1511,19 @@ def test_multi_period_did_with_conley_panel(self): ).fit( df_mp, outcome="y", + event_study=True, + spec="pooled", treatment="treated", time="time", post_periods=[2, 3], unit="unit", reference_period=1, ) - assert np.all(np.isfinite(res.vcov)), "MPD+Conley vcov must be finite" + assert np.all(np.isfinite(res.vcov)), "event-study + Conley vcov must be finite" - def test_multi_period_did_conley_missing_unit_raises(self): - """MPD + vcov_type='conley' without unit= at fit-time raises ValueError.""" - from diff_diff import MultiPeriodDiD + def test_twfe_event_study_conley_missing_unit_raises(self): + """TWFE event-study + vcov_type='conley' without unit= at fit-time raises ValueError.""" + from diff_diff import TwoWayFixedEffects rng = np.random.default_rng(seed=13) n_units = 20 @@ -1538,7 +1547,7 @@ def test_multi_period_did_conley_missing_unit_raises(self): df_mp = pd.DataFrame(rows) with pytest.raises(ValueError, match="unit="): - MultiPeriodDiD( + TwoWayFixedEffects( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, @@ -1546,16 +1555,18 @@ def test_multi_period_did_conley_missing_unit_raises(self): ).fit( df_mp, outcome="y", + event_study=True, + spec="pooled", treatment="treated", time="time", post_periods=[2], reference_period=1, ) - def test_multi_period_did_conley_missing_lag_cutoff_raises(self): + def test_twfe_event_study_conley_missing_lag_cutoff_raises(self): """MPD + vcov_type='conley' without conley_lag_cutoff raises ValueError (no defensible default per Conley 1999 Section 5).""" - from diff_diff import MultiPeriodDiD + from diff_diff import TwoWayFixedEffects rng = np.random.default_rng(seed=13) n_units = 20 @@ -1579,13 +1590,15 @@ def test_multi_period_did_conley_missing_lag_cutoff_raises(self): df_mp = pd.DataFrame(rows) with pytest.raises(ValueError, match="conley_lag_cutoff"): - MultiPeriodDiD( + TwoWayFixedEffects( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, ).fit( df_mp, outcome="y", + event_study=True, + spec="pooled", treatment="treated", time="time", post_periods=[2], @@ -1593,7 +1606,7 @@ def test_multi_period_did_conley_missing_lag_cutoff_raises(self): reference_period=1, ) - def test_multi_period_did_conley_with_survey_design_raises(self): + def test_twfe_event_study_conley_with_survey_design_raises(self): """MPD + vcov_type='conley' + survey_design raises NotImplementedError. Closes Codex P0: previously, MPD passed return_vcov=False to solve_ols @@ -1603,7 +1616,7 @@ def test_multi_period_did_conley_with_survey_design_raises(self): """ import pandas as pd - from diff_diff import MultiPeriodDiD + from diff_diff import TwoWayFixedEffects from diff_diff.survey import SurveyDesign rng = np.random.default_rng(seed=29) @@ -1631,7 +1644,7 @@ def test_multi_period_did_conley_with_survey_design_raises(self): # path; the guard must fire before solve_ols. sd_tsl = SurveyDesign(weights="weight", weight_type="pweight") with pytest.raises(NotImplementedError, match="survey_design"): - MultiPeriodDiD( + TwoWayFixedEffects( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, @@ -1639,6 +1652,8 @@ def test_multi_period_did_conley_with_survey_design_raises(self): ).fit( df_mp, outcome="y", + event_study=True, + spec="pooled", treatment="treated", time="time", post_periods=[2], @@ -1656,7 +1671,7 @@ def test_multi_period_did_conley_with_survey_design_raises(self): nest=True, ) with pytest.raises(NotImplementedError, match="survey_design"): - MultiPeriodDiD( + TwoWayFixedEffects( vcov_type="conley", conley_coords=("lat", "lon"), conley_cutoff_km=2000.0, @@ -1664,6 +1679,8 @@ def test_multi_period_did_conley_with_survey_design_raises(self): ).fit( df_mp, outcome="y", + event_study=True, + spec="pooled", treatment="treated", time="time", post_periods=[2], @@ -1672,7 +1689,7 @@ def test_multi_period_did_conley_with_survey_design_raises(self): survey_design=sd_psu, ) - def test_multi_period_did_conley_with_datetime64_time(self): + def test_twfe_event_study_conley_with_datetime64_time(self): """End-to-end MPD + vcov_type='conley' with datetime64 time labels. Closes Codex re-review P1: the wrapper must NOT coerce time to float64 before passing to _compute_conley_vcov; the helper normalizes to @@ -1681,7 +1698,7 @@ def test_multi_period_did_conley_with_datetime64_time(self): """ import pandas as pd - from diff_diff import MultiPeriodDiD + from diff_diff import TwoWayFixedEffects rng = np.random.default_rng(seed=37) n_units = 12 @@ -1711,18 +1728,22 @@ def test_multi_period_did_conley_with_datetime64_time(self): conley_cutoff_km=2000.0, conley_lag_cutoff=1, ) - res_int = MultiPeriodDiD(**kwargs).fit( + res_int = TwoWayFixedEffects(**kwargs).fit( df_mp, outcome="y", + event_study=True, + spec="pooled", treatment="treated", time="time_int", post_periods=[1, 2], unit="unit", reference_period=0, ) - res_dt = MultiPeriodDiD(**kwargs).fit( + res_dt = TwoWayFixedEffects(**kwargs).fit( df_mp, outcome="y", + event_study=True, + spec="pooled", treatment="treated", time="time_dt", post_periods=[date_labels[1], date_labels[2]], @@ -1782,14 +1803,14 @@ def test_multi_period_did_conley_to_dict_carries_lag_cutoff(self): assert d["vcov_type"] == "conley" assert d["conley_lag_cutoff"] == 2 - def test_multi_period_did_conley_missing_coords_raises(self): + def test_twfe_event_study_conley_missing_coords_raises(self): """MPD + vcov_type='conley' without conley_coords raises a clean ValueError instead of a raw TypeError on `self.conley_coords[0]`. Closes Codex P2 #1. """ import pandas as pd - from diff_diff import MultiPeriodDiD + from diff_diff import TwoWayFixedEffects rng = np.random.default_rng(seed=31) n_units = 10 @@ -1806,12 +1827,14 @@ def test_multi_period_did_conley_missing_coords_raises(self): ) df_mp = pd.DataFrame(rows) with pytest.raises(ValueError, match="conley_coords.*conley_cutoff_km"): - MultiPeriodDiD( + TwoWayFixedEffects( vcov_type="conley", conley_lag_cutoff=1, ).fit( df_mp, outcome="y", + event_study=True, + spec="pooled", treatment="treated", time="time", post_periods=[1], @@ -3106,17 +3129,17 @@ def test_mpd_combined_kernel_finite_se_and_cluster_name(self): d = res.to_dict() assert d.get("cluster_name") == "region" - def test_mpd_combined_kernel_time_varying_cluster_raises(self): + def test_twfe_es_combined_kernel_time_varying_cluster_raises(self): """MultiPeriodDiD + Conley + cluster= with a cluster that varies across periods within a unit raises ValueError (same time- invariance contract as the linalg validator). Codex CI R1 P1 #2.""" - from diff_diff import MultiPeriodDiD + from diff_diff import TwoWayFixedEffects df = self._multi_period_panel_with_region(n_units=10, T=3, seed=59) mask_violator = (df["unit"] == 2) & (df["time"] == 2) df.loc[mask_violator, "region"] = 77 with pytest.raises(ValueError, match="constant within each unit"): - MultiPeriodDiD( + TwoWayFixedEffects( vcov_type="conley", cluster="region", conley_coords=("lat", "lon"), @@ -3125,6 +3148,8 @@ def test_mpd_combined_kernel_time_varying_cluster_raises(self): ).fit( df, outcome="y", + event_study=True, + spec="pooled", treatment="treated", time="time", unit="unit", diff --git a/tests/test_diagnostic_report.py b/tests/test_diagnostic_report.py index 8f73f7ffd..b3c23693f 100644 --- a/tests/test_diagnostic_report.py +++ b/tests/test_diagnostic_report.py @@ -34,6 +34,7 @@ MultiPeriodDiD, SyntheticControl, SyntheticDiD, + TwoWayFixedEffects, generate_did_data, generate_factor_data, generate_staggered_data, @@ -86,7 +87,9 @@ def did_fit(): @pytest.fixture(scope="module") -def multi_period_fit(): +def legacy_mpd_fit(): + """Native MultiPeriodDiD fit for the MPD-only tests (event-study lanes, the + name-string dispatch pins): dies with the class at the 4.0 cut (PR 4c-ii).""" warnings.filterwarnings("ignore") df = generate_did_data(n_units=80, n_periods=8, treatment_effect=1.5, seed=7) es = MultiPeriodDiD().fit( @@ -100,6 +103,18 @@ def multi_period_fit(): return es, df +@pytest.fixture(scope="module") +def twfe_fit(): + """Static TwoWayFixedEffects fit (a scalar DiDResults) on the same panel: the + successor fixture for the generic report tests (4.0 cut, PR 4c-i).""" + warnings.filterwarnings("ignore") + df = generate_did_data(n_units=80, n_periods=8, treatment_effect=1.5, seed=7) + fit = TwoWayFixedEffects().fit( + df, outcome="outcome", treatment="treated", unit="unit", post="post" + ) + return fit, df + + @pytest.fixture(scope="module") def cs_fit(): warnings.filterwarnings("ignore") @@ -198,8 +213,8 @@ def test_every_top_level_key_present_did(self, did_fit): schema = dr.to_dict() assert set(schema.keys()) == _TOP_LEVEL_KEYS - def test_every_top_level_key_present_multiperiod(self, multi_period_fit): - fit, _ = multi_period_fit + def test_every_top_level_key_present_multiperiod(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit schema = DiagnosticReport(fit).to_dict() assert set(schema.keys()) == _TOP_LEVEL_KEYS @@ -220,8 +235,8 @@ def test_every_top_level_key_present_sdid(self, sdid_fit): schema = DiagnosticReport(fit).to_dict() assert set(schema.keys()) == _TOP_LEVEL_KEYS - def test_schema_version_constant(self, multi_period_fit): - fit, _ = multi_period_fit + def test_schema_version_constant(self, twfe_fit): + fit, _ = twfe_fit schema = DiagnosticReport(fit).to_dict() assert schema["schema_version"] == DIAGNOSTIC_REPORT_SCHEMA_VERSION assert DIAGNOSTIC_REPORT_SCHEMA_VERSION == "2.0" @@ -253,8 +268,8 @@ def test_all_statuses_use_closed_enum(self, cs_fit): section.get("status") in _STATUS_ENUM ), f"{key}.status = {section.get('status')!r} not in {_STATUS_ENUM}" - def test_json_round_trip_multiperiod(self, multi_period_fit): - fit, _ = multi_period_fit + def test_json_round_trip_twfe(self, twfe_fit): + fit, _ = twfe_fit dr = DiagnosticReport(fit) dumped = json.dumps(dr.to_dict()) assert len(dumped) > 0 @@ -312,8 +327,8 @@ def test_bacon_applicability_requires_all_column_kwargs(self, cs_fit): reason = dr.skipped_checks["bacon"] assert "outcome" in reason or "time" in reason or "unit" in reason - def test_multiperiod_runs_pt_and_power_and_sensitivity(self, multi_period_fit): - fit, _ = multi_period_fit + def test_multiperiod_runs_pt_and_power_and_sensitivity(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit dr = DiagnosticReport(fit) applicable = set(dr.applicable_checks) assert "parallel_trends" in applicable @@ -340,8 +355,8 @@ def test_sdid_has_estimator_native(self, sdid_fit): dr = DiagnosticReport(fit) assert "estimator_native" in dr.applicable_checks - def test_run_opt_outs_move_checks_to_skipped(self, multi_period_fit): - fit, _ = multi_period_fit + def test_run_opt_outs_move_checks_to_skipped(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit dr = DiagnosticReport(fit, run_sensitivity=False) assert "sensitivity" not in dr.applicable_checks assert dr.skipped_checks["sensitivity"].startswith("run_sensitivity=False") @@ -366,8 +381,8 @@ def test_placebo_is_reserved_and_skipped(self, did_fit): # Precomputed passthrough # --------------------------------------------------------------------------- class TestPrecomputed: - def test_precomputed_sensitivity_is_used_verbatim(self, multi_period_fit): - fit, _ = multi_period_fit + def test_precomputed_sensitivity_is_used_verbatim(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit # Construct a minimal SensitivityResults-shaped object the formatter recognizes. class _FakeSens: @@ -1177,8 +1192,8 @@ def test_placebo_is_always_skipped_not_not_applicable(self, did_fit): f"got {placebo['status']!r}" ) - def test_placebo_skipped_for_multiperiod_fit(self, multi_period_fit): - fit, _ = multi_period_fit + def test_placebo_skipped_for_multiperiod_fit(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit placebo = DiagnosticReport(fit).to_dict()["placebo"] assert placebo["status"] == "skipped" @@ -1731,18 +1746,18 @@ class TestPrecomputedValidation: """Regression for the P1 finding that ``precomputed=`` silently accepted keys that were never implemented. Unsupported keys now raise.""" - def test_unsupported_precomputed_key_raises(self, multi_period_fit): - fit, _ = multi_period_fit + def test_unsupported_precomputed_key_raises(self, twfe_fit): + fit, _ = twfe_fit with pytest.raises(ValueError, match="not implemented"): DiagnosticReport(fit, precomputed={"design_effect": object()}) - def test_supported_precomputed_keys_accepted(self, multi_period_fit): - fit, _ = multi_period_fit + def test_supported_precomputed_keys_accepted(self, twfe_fit): + fit, _ = twfe_fit # The four implemented keys should not raise at construction. DiagnosticReport(fit, precomputed={"parallel_trends": {"p_value": 0.5}}) - def test_mixed_supported_and_unsupported_raises(self, multi_period_fit): - fit, _ = multi_period_fit + def test_mixed_supported_and_unsupported_raises(self, twfe_fit): + fit, _ = twfe_fit with pytest.raises(ValueError, match="epv"): DiagnosticReport(fit, precomputed={"sensitivity": None, "epv": object()}) @@ -1768,8 +1783,8 @@ def _fake_single_m(self, M=1.5, ci_lb=1.0, ci_ub=3.0): alpha=0.05, ) - def test_dr_schema_preserves_single_m_marker(self, multi_period_fit): - fit, _ = multi_period_fit + def test_dr_schema_preserves_single_m_marker(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit dr = DiagnosticReport(fit, precomputed={"sensitivity": self._fake_single_m()}) sens = dr.to_dict()["sensitivity"] assert sens["status"] == "ran" @@ -1777,8 +1792,8 @@ def test_dr_schema_preserves_single_m_marker(self, multi_period_fit): assert sens["breakdown_M"] is None assert len(sens["grid"]) == 1 - def test_dr_summary_does_not_claim_full_grid_robustness(self, multi_period_fit): - fit, _ = multi_period_fit + def test_dr_summary_does_not_claim_full_grid_robustness(self, legacy_mpd_fit): + fit, _ = legacy_mpd_fit dr = DiagnosticReport(fit, precomputed={"sensitivity": self._fake_single_m()}) summary = dr.summary() assert "across the entire HonestDiD grid" not in summary @@ -1787,11 +1802,11 @@ def test_dr_summary_does_not_claim_full_grid_robustness(self, multi_period_fit): assert "single point checked" in summary assert "not a breakdown" in summary or "not a grid" in summary - def test_br_summary_does_not_claim_full_grid_robustness(self, multi_period_fit): + def test_br_summary_does_not_claim_full_grid_robustness(self, legacy_mpd_fit): """BR via honest_did_results= passthrough must not oversell a point check.""" from diff_diff import BusinessReport - fit, _ = multi_period_fit + fit, _ = legacy_mpd_fit br = BusinessReport(fit, honest_did_results=self._fake_single_m()) summary = br.summary() assert "full grid" not in summary @@ -2307,9 +2322,9 @@ def test_scm_native_marks_infeasible_placebo(self): # Error handling # --------------------------------------------------------------------------- class TestErrorHandling: - def test_sensitivity_error_does_not_break_report(self, multi_period_fit): + def test_sensitivity_error_does_not_break_report(self, legacy_mpd_fit): """A failing diagnostic records its error in the section; the report still renders.""" - fit, _ = multi_period_fit + fit, _ = legacy_mpd_fit def _raise(*args, **kwargs): raise RuntimeError("synthetic test failure") @@ -2478,16 +2493,16 @@ def test_dr_normal_fragile_keeps_multiplier(self): # Public result class # --------------------------------------------------------------------------- class TestDiagnosticReportResults: - def test_run_all_returns_dataclass(self, multi_period_fit): - fit, _ = multi_period_fit + def test_run_all_returns_dataclass(self, twfe_fit): + fit, _ = twfe_fit dr = DiagnosticReport(fit) results = dr.run_all() assert isinstance(results, DiagnosticReportResults) assert isinstance(results.applicable_checks, tuple) assert isinstance(results.schema, dict) - def test_run_all_is_idempotent(self, multi_period_fit): - fit, _ = multi_period_fit + def test_run_all_is_idempotent(self, twfe_fit): + fit, _ = twfe_fit dr = DiagnosticReport(fit) a = dr.run_all() b = dr.run_all() diff --git a/tests/test_estimators.py b/tests/test_estimators.py index b35c82687..ab4523a28 100644 --- a/tests/test_estimators.py +++ b/tests/test_estimators.py @@ -5,6 +5,7 @@ import numpy as np import pandas as pd import pytest +from event_study_average import post_period_average from diff_diff import ( DiDResults, @@ -1126,43 +1127,48 @@ def test_cluster_robust_se(self): assert results_cluster.se != results_no_cluster.se +def _multi_period_panel() -> pd.DataFrame: + """Panel data with multiple time periods and known ATT (6 periods, 3 pre / 3 post, + effect 3.0 from period 3). Module-level so the MultiPeriodDiD tests and their TWFE + event-study twins share it (4.0 cut, PR 4c-i).""" + np.random.seed(42) + n_units = 100 + n_periods = 6 + data = [] + for unit in range(n_units): + is_treated = unit < n_units // 2 + unit_effect = np.random.normal(0, 1) + for period in range(n_periods): + y = 10.0 + unit_effect + period * 0.5 + if is_treated and period >= 3: + y += 3.0 + y += np.random.normal(0, 0.5) + data.append({"unit": unit, "period": period, "treated": int(is_treated), "outcome": y}) + return pd.DataFrame(data) + + +def _event_study_panel() -> pd.DataFrame: + """Panel data with 6 periods, treatment at period 3 (no unit effects).""" + np.random.seed(42) + data = [] + for unit in range(100): + is_treated = unit < 50 + for period in range(6): + y = 10.0 + period * 0.5 + if is_treated and period >= 3: + y += 3.0 + y += np.random.normal(0, 0.5) + data.append({"unit": unit, "period": period, "treated": int(is_treated), "outcome": y}) + return pd.DataFrame(data) + + class TestMultiPeriodDiD: """Tests for MultiPeriodDiD estimator.""" @pytest.fixture def multi_period_data(self): """Create panel data with multiple time periods and known ATT.""" - np.random.seed(42) - n_units = 100 - n_periods = 6 # 3 pre-treatment, 3 post-treatment - - data = [] - for unit in range(n_units): - is_treated = unit < n_units // 2 - unit_effect = np.random.normal(0, 1) - - for period in range(n_periods): - # Common time trend - time_effect = period * 0.5 - - y = 10.0 + unit_effect + time_effect - - # Treatment effect: 3.0 in post-periods (periods 3, 4, 5) - if is_treated and period >= 3: - y += 3.0 - - y += np.random.normal(0, 0.5) - - data.append( - { - "unit": unit, - "period": period, - "treated": int(is_treated), - "outcome": y, - } - ) - - return pd.DataFrame(data) + return _multi_period_panel() @pytest.fixture def heterogeneous_effects_data(self): @@ -1922,30 +1928,202 @@ def test_rank_deficient_action_silent_no_warning(self, multi_period_data): ), "Collinear covariate coefficient should be NaN" -class TestMultiPeriodDiDEventStudy: - """Tests for MultiPeriodDiD full event-study specification (pre + post periods).""" +class TestTwoWayFixedEffectsEventStudy: + """TWFE event-study twins of the MultiPeriodDiD core-behaviour pins that had no + TWFE coverage (4.0 cut, PR 4c-i). The unit-driven validation twins keep + `unit=` (the warnings run only inside the core's `if unit is not None:` + block) and assert the warning or its absence, never SE values - the pooled + spec auto-clusters at unit when `unit=` is supplied.""" @pytest.fixture - def panel_data(self): - """Panel data with 6 periods, treatment at period 3.""" + def multi_period_panel(self): + return _multi_period_panel() + + @pytest.fixture + def event_study_panel(self): + return _event_study_panel() + + def _fit(self, data, **kw): + base = dict(outcome="outcome", treatment="treated", time="period", event_study=True) + base.update(kw) + return TwoWayFixedEffects().fit(data, **base) + + def test_reference_period_in_post_raises(self, event_study_panel): + with pytest.raises(ValueError, match="post-treatment period"): + self._fit(event_study_panel, spec="pooled", post_periods=[3, 4, 5], reference_period=4) + + def test_staggered_treatment_warning(self): np.random.seed(42) data = [] - for unit in range(100): - is_treated = unit < 50 - for period in range(6): - y = 10.0 + period * 0.5 - if is_treated and period >= 3: - y += 3.0 - y += np.random.normal(0, 0.5) + for unit in range(40): + if unit < 10: + treat_start = 3 + elif unit < 20: + treat_start = 5 + else: + treat_start = None + for period in range(8): + is_treated = treat_start is not None and period >= treat_start + y = 10.0 + period * 0.5 + (2.0 if is_treated else 0) + np.random.normal(0, 0.3) data.append( - { - "unit": unit, - "period": period, - "treated": int(is_treated), - "outcome": y, - } + {"unit": unit, "period": period, "treated": int(is_treated), "outcome": y} ) - return pd.DataFrame(data) + df = pd.DataFrame(data) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + self._fit( + df, spec="pooled", post_periods=[3, 4, 5, 6, 7], reference_period=2, unit="unit" + ) + assert any("staggered" in str(x.message).lower() for x in w) + + def test_single_pre_period_warns(self): + np.random.seed(42) + data = [] + for unit_id in range(40): + is_treated = unit_id < 20 + for period in range(3): + y = 10.0 + period * 0.5 + (2.0 if is_treated and period >= 1 else 0) + y += np.random.normal(0, 0.3) + data.append( + {"unit": unit_id, "period": period, "treated": int(is_treated), "outcome": y} + ) + df = pd.DataFrame(data) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + results = self._fit(df, spec="pooled", post_periods=[1, 2], reference_period=0) + assert any("Only one pre-treatment period" in str(x.message) for x in w) + att_bar, _ = post_period_average(results) + assert np.isfinite(att_bar) + + def test_treatment_reversal_warns(self): + np.random.seed(42) + data = [] + for unit_id in range(40): + for period in range(6): + is_treated = unit_id < 20 + if unit_id == 0: + d = 1 if 2 <= period <= 3 else 0 + elif is_treated: + d = 1 if period >= 3 else 0 + else: + d = 0 + y = 10.0 + period * 0.5 + (2.0 if d else 0) + np.random.normal(0, 0.3) + data.append({"unit": unit_id, "period": period, "treated": d, "outcome": y}) + df = pd.DataFrame(data) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + self._fit(df, spec="pooled", post_periods=[3, 4, 5], reference_period=2, unit="unit") + assert any("Treatment reversal" in str(x.message) for x in w) + + def test_time_varying_treatment_warning(self): + np.random.seed(42) + data = [] + for unit_id in range(40): + is_treated = unit_id < 20 + for period in range(6): + d = 1 if is_treated and period >= 3 else 0 + y = 10.0 + period * 0.5 + (2.0 if d else 0) + np.random.normal(0, 0.3) + data.append({"unit": unit_id, "period": period, "treated": d, "outcome": y}) + df = pd.DataFrame(data) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + results = self._fit( + df, spec="pooled", post_periods=[3, 4, 5], reference_period=2, unit="unit" + ) + dit_warnings = [x for x in w if "time-varying" in str(x.message).lower()] + assert dit_warnings, "Expected warning about time-varying treatment" + assert "ever_treated" in str(dit_warnings[0].message) + assert results is not None + assert np.sum(~results.is_reference) > 0 + + def test_staggered_no_false_positive_unbalanced(self): + np.random.seed(42) + data = [] + for unit_id in range(40): + is_treated = unit_id < 20 + start_period = 4 if (is_treated and unit_id < 5) else 0 + for period in range(start_period, 8): + d = 1 if is_treated and period >= 3 else 0 + y = 10.0 + period * 0.5 + (2.0 if d else 0) + np.random.normal(0, 0.3) + data.append({"unit": unit_id, "period": period, "treated": d, "outcome": y}) + df = pd.DataFrame(data) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + self._fit( + df, spec="pooled", post_periods=[3, 4, 5, 6, 7], reference_period=2, unit="unit" + ) + assert not any("staggered" in str(x.message).lower() for x in w) + + def _collinear(self, panel): + panel = panel.copy() + panel["collinear_cov"] = panel["treated"].copy() + return panel + + def test_rank_deficient_design_warns_and_keeps_period_rows_finite(self, multi_period_panel): + """A collinear covariate warns; the per-period quintet stays finite (the + dropped column is the covariate, not a period interaction).""" + df = self._collinear(multi_period_panel) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + results = self._fit( + df, + spec="pooled", + post_periods=[3, 4, 5], + reference_period=2, + covariates=["collinear_cov"], + ) + assert any( + "Rank-deficient" in str(x.message) or "collinear" in str(x.message).lower() for x in w + ) + rows = ~results.is_reference + for field in ("att", "se", "t_stat", "p_value"): + assert np.all(np.isfinite(getattr(results, field)[rows])), field + + def test_rank_deficient_action_error_raises(self, multi_period_panel): + df = self._collinear(multi_period_panel) + with pytest.raises(ValueError, match="rank-deficient"): + TwoWayFixedEffects(rank_deficient_action="error").fit( + df, + outcome="outcome", + treatment="treated", + time="period", + event_study=True, + spec="pooled", + post_periods=[3, 4, 5], + reference_period=2, + covariates=["collinear_cov"], + ) + + def test_rank_deficient_action_silent_no_warning(self, multi_period_panel): + df = self._collinear(multi_period_panel) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + results = TwoWayFixedEffects(rank_deficient_action="silent").fit( + df, + outcome="outcome", + treatment="treated", + time="period", + event_study=True, + spec="pooled", + post_periods=[3, 4, 5], + reference_period=2, + covariates=["collinear_cov"], + ) + assert not any( + "Rank-deficient" in str(x.message) or "rank-deficient" in str(x.message).lower() + for x in w + ), [str(x.message) for x in w] + rows = ~results.is_reference + assert np.all(np.isfinite(results.att[rows])) + + +class TestMultiPeriodDiDEventStudy: + """Tests for MultiPeriodDiD full event-study specification (pre + post periods).""" + + @pytest.fixture + def panel_data(self): + """Panel data with 6 periods, treatment at period 3.""" + return _event_study_panel() def test_default_reference_period_is_last_pre(self, panel_data): """Verify reference_period defaults to the last pre-period.""" @@ -3302,8 +3480,8 @@ def test_twfe_with_unbalanced_panel(self): df, outcome="outcome", treatment="post", unit="unit", post="period" ) - def test_multiperiod_with_sparse_data(self): - """Test MultiPeriodDiD with sparse data across periods.""" + def test_event_study_with_sparse_data(self): + """Test the TWFE event-study mode with sparse data across periods.""" np.random.seed(42) data = [] @@ -3333,14 +3511,23 @@ def test_multiperiod_with_sparse_data(self): df = pd.DataFrame(data) - mp_did = MultiPeriodDiD() - results = mp_did.fit( - df, outcome="outcome", treatment="treated", time="period", reference_period=1 + # Pooled event study (the former MultiPeriodDiD fit); the midpoint default + # over the observed periods 0-4 resolved to post [2, 3, 4]. + results = TwoWayFixedEffects().fit( + df, + outcome="outcome", + treatment="treated", + time="period", + event_study=True, + spec="pooled", + post_periods=[2, 3, 4], + reference_period=1, ) # Should produce valid results - assert np.isfinite(results.avg_att) - assert len(results.period_effects) > 0 + att_bar, _ = post_period_average(results) + assert np.isfinite(att_bar) + assert np.sum(~results.is_reference) > 0 class TestSingleTreatedUnit: diff --git a/tests/test_estimators_vcov_type.py b/tests/test_estimators_vcov_type.py index bffa48ede..9ada20d34 100644 --- a/tests/test_estimators_vcov_type.py +++ b/tests/test_estimators_vcov_type.py @@ -21,6 +21,7 @@ import numpy as np import pandas as pd import pytest +from event_study_average import post_period_average from diff_diff import SunAbraham, SurveyDesign from diff_diff.estimators import DifferenceInDifferences, MultiPeriodDiD @@ -58,8 +59,17 @@ def test_single_treated_observation_per_period(self, vcov_type, event_study): ) kwargs = dict(outcome="y", treatment="treated") if event_study: - cls = MultiPeriodDiD - kwargs.update(time="time", unit="unit", post_periods=[2, 3], reference_period=1) + # TWFE event-study mode, pooled design (4.0 cut, PR 4c-i). No unit=: + # hc2_bm is not in the auto-cluster carve-out, so a unit= would + # cluster the fit; the former MultiPeriodDiD arm never clustered. + cls = TwoWayFixedEffects + kwargs.update( + time="time", + event_study=True, + spec="pooled", + post_periods=[2, 3], + reference_period=1, + ) else: cls = DifferenceInDifferences kwargs.update(post="time") @@ -70,24 +80,37 @@ def test_single_treated_observation_per_period(self, vcov_type, event_study): warnings.simplefilter("always") result = cls(vcov_type=vcov_type).fit(data, **kwargs) family = "HC2-BM" if vcov_type == "hc2_bm" else "HC2" - # MPD's additional contrast-DOF calculation must not warn a second time. + # The core's additional contrast-DOF calculation must not warn a second time. assert sum(f"{family} variance is undefined" in str(w.message) for w in caught) == 1 - assert np.isnan(result.vcov).all() and np.isnan(result.se) - assert result.att == pytest.approx(baseline.att) - assert_nan_inference( - dict( - se=result.se, t_stat=result.t_stat, p_value=result.p_value, conf_int=result.conf_int - ) - ) - np.testing.assert_allclose(result.residuals, baseline.residuals, atol=1e-13) - np.testing.assert_allclose(result.fitted_values, baseline.fitted_values, atol=1e-13) if event_study: - for period, effect in result.period_effects.items(): - if period == result.reference_period: - continue - assert effect.effect == pytest.approx(baseline.period_effects[period].effect) - assert np.isnan(effect.se) - assert_nan_inference(vars(effect)) + # Per-period surface: every non-reference row fails closed. + assert np.isnan(result.vcov).all() + rows = np.where(~result.is_reference)[0] + assert len(rows) > 0 + np.testing.assert_allclose(result.att[rows], baseline.att[rows]) + for i in rows: + assert np.isnan(result.se[i]) + assert_nan_inference( + dict( + se=result.se[i], + t_stat=result.t_stat[i], + p_value=result.p_value[i], + conf_int=(result.conf_int_lower[i], result.conf_int_upper[i]), + ) + ) + else: + assert np.isnan(result.vcov).all() and np.isnan(result.se) + assert result.att == pytest.approx(baseline.att) + assert_nan_inference( + dict( + se=result.se, + t_stat=result.t_stat, + p_value=result.p_value, + conf_int=result.conf_int, + ) + ) + np.testing.assert_allclose(result.residuals, baseline.residuals, atol=1e-13) + np.testing.assert_allclose(result.fitted_values, baseline.fitted_values, atol=1e-13) # ============================================================================= @@ -367,8 +390,8 @@ def test_summary_includes_vcov_label_cr1(self): summary = res.summary() assert "CR1 cluster-robust at unit" in summary - def test_multi_period_fit_honors_classical(self): - """MultiPeriodDiD.fit with vcov_type='classical' produces non-robust SEs. + def test_event_study_fit_honors_classical(self): + """The TWFE event-study fit with vcov_type='classical' produces non-robust SEs. Regression test for the CI review finding: `MultiPeriodDiD` inherits `vcov_type` from the base class via get_params but its `fit()` path @@ -386,19 +409,34 @@ def test_multi_period_fit_honors_classical(self): rows.append({"unit": i, "time": t, "treated": treated, "y": y}) data = pd.DataFrame(rows) - r_hc1 = MultiPeriodDiD(vcov_type="hc1").fit( - data, outcome="y", treatment="treated", time="time" + _periods = sorted(data["time"].unique()) + _post = _periods[len(_periods) // 2 :] # the former MultiPeriodDiD midpoint default + r_hc1 = TwoWayFixedEffects(vcov_type="hc1").fit( + data, + outcome="y", + treatment="treated", + time="time", + event_study=True, + spec="pooled", + post_periods=_post, ) - r_classical = MultiPeriodDiD(vcov_type="classical").fit( - data, outcome="y", treatment="treated", time="time" + r_classical = TwoWayFixedEffects(vcov_type="classical").fit( + data, + outcome="y", + treatment="treated", + time="time", + event_study=True, + spec="pooled", + post_periods=_post, ) + rows = ~r_hc1.is_reference # Point estimates identical. - assert r_hc1.avg_att == pytest.approx(r_classical.avg_att, abs=1e-10) + np.testing.assert_allclose(r_hc1.att[rows], r_classical.att[rows], atol=1e-10) # SEs must differ — vcov_type actually changed the variance family. - assert r_hc1.avg_se != pytest.approx(r_classical.avg_se, abs=1e-10) + assert not np.allclose(r_hc1.se[rows], r_classical.se[rows], atol=1e-10) - def test_multi_period_cluster_plus_hc2_bm_produces_finite_inference(self): - """MultiPeriodDiD(cluster=..., vcov_type='hc2_bm') is now supported. + def test_event_study_cluster_plus_hc2_bm_produces_finite_inference(self): + """The TWFE event-study mode with cluster=... + vcov_type='hc2_bm' is supported. The cluster-aware CR2 Bell-McCaffrey contrast DOF for the post-period-average ATT is implemented via the new @@ -416,25 +454,30 @@ def test_multi_period_cluster_plus_hc2_bm_produces_finite_inference(self): rows.append({"unit": i, "time": t, "treated": treated, "y": y}) data = pd.DataFrame(rows) - res = MultiPeriodDiD(vcov_type="hc2_bm", cluster="unit").fit( - data, outcome="y", treatment="treated", time="time" - ) - # Headline contract: finite avg_att inference under cluster+hc2_bm. - assert np.isfinite(res.avg_att) - assert np.isfinite(res.avg_se) - assert np.isfinite(res.avg_t_stat) - assert np.isfinite(res.avg_p_value) - assert np.isfinite(res.avg_conf_int[0]) - assert np.isfinite(res.avg_conf_int[1]) - # Per-period inference should also be finite for the post period. - post_pe = res.period_effects[2] - assert np.isfinite(post_pe.effect) - assert np.isfinite(post_pe.se) - assert np.isfinite(post_pe.p_value) - - def test_multi_period_hc2_bm_nonfinite_bm_dof_fails_closed(self, monkeypatch): + _periods = sorted(data["time"].unique()) + _post = _periods[len(_periods) // 2 :] # the former MultiPeriodDiD midpoint default + res = TwoWayFixedEffects(vcov_type="hc2_bm", cluster="unit").fit( + data, + outcome="y", + treatment="treated", + time="time", + event_study=True, + spec="pooled", + post_periods=_post, + ) + # Headline contract: finite post-period-average inference under cluster+hc2_bm. + att_bar, se_bar = post_period_average(res) + assert np.isfinite(att_bar) and np.isfinite(se_bar) + # Per-period inference is finite on every non-reference row (incl. the post period 2). + rows = ~res.is_reference + assert np.all(np.isfinite(res.att[rows])) + assert np.all(np.isfinite(res.se[rows])) + assert np.all(np.isfinite(res.p_value[rows])) + assert 2 in res.event_time[rows] + + def test_event_study_hc2_bm_nonfinite_bm_dof_fails_closed(self, monkeypatch): """A non-finite (guard-suppressed) Bell-McCaffrey DOF must fail closed to - NaN inference on user-facing MPD contrasts, not fall back to the shared + NaN inference on the user-facing event-study rows, not fall back to the shared residual df. The `_cr2_bm_dof_inner` guard can NaN a per-coefficient DOF (high-leverage / collinear column); reporting finite t/p/CI for such a coefficient via the residual df would violate the joint-NaN inference @@ -460,23 +503,29 @@ def _nan_dof(*args, **kwargs): vcov, dof = _orig(*args, **kwargs) return vcov, np.full(np.asarray(dof).shape, np.nan) - # MPD's hc2_bm+cluster path imports this from linalg inside fit(), so the + # The core's hc2_bm+cluster path imports this from linalg inside fit(), so the # local import binds to the patched attribute at call time. monkeypatch.setattr(_linalg, "_compute_cr2_bm_vcov_and_dof", _nan_dof) - res = MultiPeriodDiD(vcov_type="hc2_bm", cluster="unit").fit( - data, outcome="y", treatment="treated", time="time" - ) - for p, pe in res.period_effects.items(): - assert np.isfinite(pe.se), f"period {p}: SE should remain finite" - assert np.isnan(pe.t_stat), f"period {p}: t_stat should be NaN" - assert np.isnan(pe.p_value), f"period {p}: p_value should be NaN" - assert np.isnan(pe.conf_int[0]) and np.isnan(pe.conf_int[1]) - # The post-period-average contrast fails closed too. - assert np.isfinite(res.avg_se), "avg SE should remain finite" - assert np.isnan(res.avg_t_stat) - assert np.isnan(res.avg_p_value) - assert np.isnan(res.avg_conf_int[0]) and np.isnan(res.avg_conf_int[1]) + _periods = sorted(data["time"].unique()) + _post = _periods[len(_periods) // 2 :] # the former MultiPeriodDiD midpoint default + res = TwoWayFixedEffects(vcov_type="hc2_bm", cluster="unit").fit( + data, + outcome="y", + treatment="treated", + time="time", + event_study=True, + spec="pooled", + post_periods=_post, + ) + for i in np.where(~res.is_reference)[0]: + p = res.event_time[i] + assert np.isfinite(res.se[i]), f"period {p}: SE should remain finite" + assert np.isnan(res.t_stat[i]), f"period {p}: t_stat should be NaN" + assert np.isnan(res.p_value[i]), f"period {p}: p_value should be NaN" + assert np.isnan(res.conf_int_lower[i]) and np.isnan(res.conf_int_upper[i]) + # The per-period df carries the NaN BM DOF (no silent residual-df fallback). + assert np.all(np.isnan(res.df[~res.is_reference])) def test_multi_period_cluster_hc2_bm_avg_att_uses_clubsandwich_dof(self): """MPD(cluster=..., hc2_bm) `avg_att` inference uses the new @@ -586,12 +635,12 @@ def _residual(df): expected_ci = (pe.effect - t_crit * pe.se, pe.effect + t_crit * pe.se) np.testing.assert_allclose(pe.conf_int, expected_ci, atol=1e-8, rtol=0) - def test_multi_period_fit_honors_hc2_bm(self): - """MultiPeriodDiD.fit with vcov_type='hc2_bm' uses Bell-McCaffrey DOF. + def test_event_study_fit_honors_hc2_bm(self): + """The TWFE event-study fit with vcov_type='hc2_bm' uses Bell-McCaffrey DOF. Checks two things: (a) fit completes without error on the hc2_bm path for the period-effect loop, and (b) the BM Satterthwaite DOF produces - a CI for avg_att with a finite width (non-degenerate case). + per-period CIs with a finite width (non-degenerate case). """ rng = np.random.default_rng(1919) n_units = 50 @@ -604,16 +653,24 @@ def test_multi_period_fit_honors_hc2_bm(self): rows.append({"unit": i, "time": t, "treated": treated, "y": y}) data = pd.DataFrame(rows) - r_hc2bm = MultiPeriodDiD(vcov_type="hc2_bm").fit( - data, outcome="y", treatment="treated", time="time" - ) - assert np.isfinite(r_hc2bm.avg_att) - assert np.isfinite(r_hc2bm.avg_se) - assert np.isfinite(r_hc2bm.avg_conf_int[0]) - assert np.isfinite(r_hc2bm.avg_conf_int[1]) - # CI width is finite and positive. - ci_width = r_hc2bm.avg_conf_int[1] - r_hc2bm.avg_conf_int[0] - assert ci_width > 0 + _periods = sorted(data["time"].unique()) + _post = _periods[len(_periods) // 2 :] # the former MultiPeriodDiD midpoint default + r_hc2bm = TwoWayFixedEffects(vcov_type="hc2_bm").fit( + data, + outcome="y", + treatment="treated", + time="time", + event_study=True, + spec="pooled", + post_periods=_post, + ) + rows = ~r_hc2bm.is_reference + assert np.all(np.isfinite(r_hc2bm.att[rows])) + assert np.all(np.isfinite(r_hc2bm.se[rows])) + assert np.all(np.isfinite(r_hc2bm.conf_int_lower[rows])) + assert np.all(np.isfinite(r_hc2bm.conf_int_upper[rows])) + # CI widths are finite and positive. + assert np.all((r_hc2bm.conf_int_upper[rows] - r_hc2bm.conf_int_lower[rows]) > 0) def test_twfe_hc2_and_hc2_bm_produce_finite_inference(self): """TWFE with vcov_type in {hc2, hc2_bm} now produces finite inference @@ -2317,8 +2374,8 @@ def _boom(*args, **kwargs): assert set(r.coefficients.keys()) == {"ATT"} -class TestMPDClusterHC2BMSharedPrecompute: - """`MultiPeriodDiD(cluster=..., vcov_type='hc2_bm')` builds the CR2 +class TestEventStudyClusterHC2BMSharedPrecompute: + """The TWFE event-study mode with cluster=... + vcov_type='hc2_bm' builds the CR2 Bell-McCaffrey precomputes ONCE, not twice. Mechanism guard for the perf dedup: vcov and the per-coefficient + @@ -2378,8 +2435,16 @@ def _counting(*args, **kwargs): monkeypatch.setattr(L, "_compute_cr2_bm_vcov_and_dof", _counting) with warnings.catch_warnings(): warnings.simplefilter("ignore") - res = MultiPeriodDiD(vcov_type="hc2_bm", cluster="unit").fit( - data, outcome="y", treatment="treated", time="time", unit="unit" + _periods = sorted(data["time"].unique()) + _post = _periods[len(_periods) // 2 :] # the former MultiPeriodDiD midpoint default + res = TwoWayFixedEffects(vcov_type="hc2_bm", cluster="unit").fit( + data, + outcome="y", + treatment="treated", + time="time", + event_study=True, + spec="pooled", + post_periods=_post, ) # One build: the shared core entered exactly once (was twice when @@ -2390,35 +2455,88 @@ def _counting(*args, **kwargs): "built more than once." ) # Inference is still finite (the dedup did not break the path). - assert np.isfinite(res.avg_att) and np.isfinite(res.avg_se) - assert np.isfinite(res.avg_p_value) + att_bar, se_bar = post_period_average(res) + assert np.isfinite(att_bar) and np.isfinite(se_bar) + assert np.all(np.isfinite(res.p_value[~res.is_reference])) def test_fit_is_reproducible(self): """Two independent fits (and a repeat fit of the same estimator) give - identical avg-ATT and per-period inference — determinism + the + identical per-period inference — determinism + the fit-does-not-mutate-config contract on the bypass path.""" data = self._balanced_panel() - est = MultiPeriodDiD(vcov_type="hc2_bm", cluster="unit") + kw = dict(outcome="y", treatment="treated", time="time", event_study=True, spec="pooled") + kw["post_periods"] = sorted(data["time"].unique())[len(data["time"].unique()) // 2 :] + est = TwoWayFixedEffects(vcov_type="hc2_bm", cluster="unit") with warnings.catch_warnings(): warnings.simplefilter("ignore") - r1 = est.fit(data, outcome="y", treatment="treated", time="time", unit="unit") + r1 = est.fit(data, **kw) # Repeat fit of the SAME object (config must be unmutated). - r2 = est.fit(data, outcome="y", treatment="treated", time="time", unit="unit") + r2 = est.fit(data, **kw) # Fresh object, same config. - r3 = MultiPeriodDiD(vcov_type="hc2_bm", cluster="unit").fit( - data, outcome="y", treatment="treated", time="time", unit="unit" - ) + r3 = TwoWayFixedEffects(vcov_type="hc2_bm", cluster="unit").fit(data, **kw) for r in (r2, r3): - assert r.avg_att == r1.avg_att - assert r.avg_se == r1.avg_se - assert r.avg_t_stat == r1.avg_t_stat - assert r.avg_p_value == r1.avg_p_value - assert set(r.period_effects) == set(r1.period_effects) - for p in r1.period_effects: - assert r.period_effects[p].effect == r1.period_effects[p].effect - assert r.period_effects[p].se == r1.period_effects[p].se + np.testing.assert_array_equal(r.event_time, r1.event_time) + for field in ("att", "se", "t_stat", "p_value", "vcov"): + a, b_ = getattr(r, field), getattr(r1, field) + np.testing.assert_array_equal(np.isnan(a), np.isnan(b_)) + np.testing.assert_array_equal(a[~np.isnan(a)], b_[~np.isnan(b_)]) + + +class TestTWFEEventStudyAbsorbedFERParity: + """R-parity for the TWFE event-study mode's within spec with the explicit + hc2 family (the classical/hc2/hc3 carve-out: no auto-cluster), against the + same ``mpd_absorbed_fe_did`` golden the MultiPeriodDiD absorb tests use + (4.0 cut, PR 4c-i). The one-way HC2-BM lane has no within twin (within + + hc2_bm auto-clusters at unit) and stays on the MultiPeriodDiD class until + PR 4c-ii; see TODO.md.""" + + def _load_golden(self): + import json + from pathlib import Path + + golden_path = ( + Path(__file__).parent.parent / "benchmarks" / "data" / "clubsandwich_cr2_golden.json" + ) + if not golden_path.exists(): + pytest.skip("Golden JSON not present; see benchmarks/R/generate_clubsandwich_golden.R") + with open(golden_path) as f: + golden = json.load(f) + if "mpd_absorbed_fe_did" not in golden: + pytest.skip("Golden JSON lacks the `mpd_absorbed_fe_did` scenario") + return golden["mpd_absorbed_fe_did"] + + def test_within_hc2_matches_sandwich_vcovhc(self): + """spec='within' + hc2 matches `lm() + sandwich::vcovHC(type="HC2")` + at 1e-10 on the target per-period interaction.""" + d = self._load_golden() + target_period = int(d["target_period"]) + target_coef = f"treated_period_{target_period}" + idx = d["coef_names"].index(target_coef) + expected_vcov = np.asarray(d["vcov_hc2"]).reshape(d["vcov_hc2_shape"]) + expected_se = float(np.sqrt(expected_vcov[idx, idx])) + expected_coef = float(d["coef"][idx]) + + data = pd.DataFrame( + {"unit": d["unit"], "period": d["period"], "treated": d["treated"], "y": d["y"]} + ) + # The golden has no post_periods key: the MPD side used the midpoint + # default over periods 1..5 (reference 1), which resolves to [3, 4, 5]. + res = TwoWayFixedEffects(vcov_type="hc2").fit( + data, + outcome="y", + treatment="treated", + unit="unit", + time="period", + event_study=True, + spec="within", + post_periods=[3, 4, 5], + reference_period=int(d["reference_period"]), + ) + (row,) = np.where(res.event_time == target_period)[0] + np.testing.assert_allclose(res.att[row], expected_coef, atol=1e-10) + np.testing.assert_allclose(res.se[row], expected_se, atol=1e-10) def _make_absorb_panel(seed: int = 11, n_units: int = 50, n_periods: int = 6) -> pd.DataFrame: @@ -2639,6 +2757,47 @@ def test_absorb_saturated_full_k_df_le_zero_nan_inference(self): ) assert np.isnan(r.avg_se), "saturated full-K design must yield NaN SE (fail-closed)" + def test_within_event_study_saturated_full_k_df_le_zero_nan_inference(self): + """TWFE event-study twin (spec='within' + classical) of the saturated + full-K fail-closed contract: every non-reference row's SE and inference + are NaN while the point estimates are still computed.""" + from tests.conftest import assert_nan_inference + + rng = np.random.default_rng(0) + rows = [] + for u in range(2): + tr = int(u >= 1) + for t in range(5): + rows.append( + { + "unit": u, + "time": t, + "treated": tr, + "y": float(rng.normal()) + 0.5 * tr * (t >= 1), + } + ) + df = pd.DataFrame(rows) + r = TwoWayFixedEffects(vcov_type="classical").fit( + df, + outcome="y", + treatment="treated", + unit="unit", + time="time", + event_study=True, + spec="within", + post_periods=[2, 3, 4], + ) + for i in np.where(~r.is_reference)[0]: + assert np.isnan(r.se[i]), "saturated full-K design must yield NaN SE (fail-closed)" + assert_nan_inference( + dict( + se=r.se[i], + t_stat=r.t_stat[i], + p_value=r.p_value[i], + conf_int=(r.conf_int_lower[i], r.conf_int_upper[i]), + ) + ) + class TestDfConvention: """`df_convention="cluster"` (opt-in Stata/fixest G-1 inference df on @@ -2724,22 +2883,30 @@ def test_twfe_inherits_knob(self): G = data["unit"].nunique() np.testing.assert_allclose(r1.p_value, 2 * stats.t.sf(abs(r1.t_stat), G - 1), rtol=1e-12) - def test_mpd_period_effects_follow_convention(self): + def test_event_study_period_effects_follow_convention(self): from scipy import stats data = self._clustered_panel() data["treated"] = data["group"] - kw = dict(outcome="y", treatment="treated", time="time", unit="unit") - m0 = MultiPeriodDiD(cluster="unit").fit(data, **kw) - m1 = MultiPeriodDiD(cluster="unit", df_convention="cluster").fit(data, **kw) + # Pooled event study with the explicit unit cluster (the former + # MultiPeriodDiD(cluster="unit") fit; periods 0-3 -> post [2, 3]). + kw = dict( + outcome="y", + treatment="treated", + time="time", + event_study=True, + spec="pooled", + post_periods=[2, 3], + ) + m0 = TwoWayFixedEffects(cluster="unit").fit(data, **kw) + m1 = TwoWayFixedEffects(cluster="unit", df_convention="cluster").fit(data, **kw) G = data["unit"].nunique() - for period, pe1 in m1.period_effects.items(): - pe0 = m0.period_effects[period] - assert pe0.se == pe1.se - if np.isfinite(pe1.p_value): + for i in np.where(~m1.is_reference)[0]: + assert m0.se[i] == m1.se[i] + if np.isfinite(m1.p_value[i]): np.testing.assert_allclose( - pe1.p_value, - 2 * stats.t.sf(abs(pe1.effect / pe1.se), G - 1), + m1.p_value[i], + 2 * stats.t.sf(abs(m1.att[i] / m1.se[i]), G - 1), rtol=1e-12, ) @@ -2890,17 +3057,23 @@ def test_get_inference_guard_fails_closed_at_one_cluster(self): assert np.isfinite(inf.se) # SE reported; only the reference- # distribution-dependent fields are NaN - def test_mpd_one_effective_cluster_fails_closed(self): - """MPD twin of the guard: a one-effective-cluster fit fails at the + def test_event_study_one_effective_cluster_fails_closed(self): + """Event-study twin of the guard: a one-effective-cluster fit fails at the CR1 vcov layer (the estimator surface has no raw weights= kwarg, so the weighted one-positive-cluster route is exercised at the linalg - level above; here a single cluster label locks the MPD route).""" + level above; here a single cluster label locks the event-study route).""" data = self._clustered_panel(n_units=4) data["treated"] = data["group"] data["one_cluster"] = 0 with pytest.raises(ValueError, match="at least 2 clusters"): - MultiPeriodDiD(cluster="one_cluster", df_convention="cluster").fit( - data, outcome="y", treatment="treated", time="time", unit="unit" + TwoWayFixedEffects(cluster="one_cluster", df_convention="cluster").fit( + data, + outcome="y", + treatment="treated", + time="time", + event_study=True, + spec="pooled", + post_periods=[2, 3], ) def test_nan_cluster_labels_rejected(self): @@ -2933,7 +3106,7 @@ def test_nan_cluster_labels_rejected(self): def test_nan_cluster_labels_rejected_all_backends(self): """Round-6 P0: every clustered backend route shares the front-door rejection — solve_ols(return_vcov=True) (the Rust HC1/CR1 dispatch - when available), CR2-BM (hc2_bm + cluster), and the MPD estimator + when available), CR2-BM (hc2_bm + cluster), and the event-study estimator route (which calls solve_ols directly).""" from diff_diff.linalg import compute_robust_vcov, solve_ols @@ -2954,8 +3127,14 @@ def test_nan_cluster_labels_rejected_all_backends(self): data["treated"] = data["group"] data.loc[data["unit"] == 0, "unit"] = np.nan with pytest.raises(ValueError, match="missing values"): - MultiPeriodDiD(cluster="unit").fit( - data, outcome="y", treatment="treated", time="time", unit="unit" + TwoWayFixedEffects(cluster="unit").fit( + data, + outcome="y", + treatment="treated", + time="time", + event_study=True, + spec="pooled", + post_periods=[2, 3], ) def test_validation_and_transactional_set_params(self): @@ -3066,6 +3245,23 @@ def test_twfe_and_mpd_normal(self): ) assert rm.inference_df is None assert rm.avg_p_value == pytest.approx(2 * stats.norm.sf(abs(rm.avg_t_stat)), rel=1e-13) + # TWFE event-study twin (pooled design): z-based p-values on every + # non-reference row and no finite df under the normal convention. + re_ = TwoWayFixedEffects(cluster="unit", df_convention="normal").fit( + data, + outcome="y", + treatment="group", + time="time", + event_study=True, + spec="pooled", + post_periods=[3, 4, 5], + reference_period=2, + ) + rows = ~re_.is_reference + np.testing.assert_allclose( + re_.p_value[rows], 2 * stats.norm.sf(np.abs(re_.t_stat[rows])), rtol=1e-13 + ) + assert np.all(np.isnan(re_.df[rows])) def test_linear_regression_normal_and_precedence(self): from scipy import stats diff --git a/tests/test_event_study_surface.py b/tests/test_event_study_surface.py index a2a2deebc..ff01c5f98 100644 --- a/tests/test_event_study_surface.py +++ b/tests/test_event_study_surface.py @@ -469,6 +469,34 @@ def surfaces(): ) out["MultiPeriodDiD"] = (mpd, build_event_study_surface(mpd)) + # 11b/11c. TWFE event-study mode (4.0 cut, PR 4c-i): the fit result IS + # the surface, so the tuple carries it twice (native == surface). + from diff_diff import TwoWayFixedEffects + + twfe_pooled = TwoWayFixedEffects().fit( + mpd_data, + outcome="outcome", + treatment="treated", + time="period", + event_study=True, + spec="pooled", + post_periods=[3, 4, 5], + reference_period=2, + ) + out["TwoWayFixedEffects_pooled"] = (twfe_pooled, twfe_pooled) + twfe_within = TwoWayFixedEffects().fit( + mpd_data, + outcome="outcome", + treatment="treated", + unit="unit", + time="period", + event_study=True, + spec="within", + post_periods=[3, 4, 5], + reference_period=2, + ) + out["TwoWayFixedEffects_within"] = (twfe_within, twfe_within) + # 12. LPDiD lp_data = _lpdid_panel() lp = LPDiD(pre_window=2, post_window=2).fit( @@ -512,12 +540,21 @@ def surfaces(): "WooldridgeDiD", "TripleDifference_staggered", "MultiPeriodDiD", + "TwoWayFixedEffects_pooled", + "TwoWayFixedEffects_within", "LPDiD", "ChaisemartinDHaultfoeuille", "HeterogeneousAdoptionDiD", ] +#: Producers whose native result IS the unified surface: no independent native +#: representation exists to compare against (their quintets are pinned against +#: the MultiPeriodDiD twin by tests/test_v4_merge_mpd.py and +#: tests/test_twfe_event_study_decoupling.py). +NATIVE_SURFACE_PRODUCERS = {"TwoWayFixedEffects_pooled", "TwoWayFixedEffects_within"} + + def _expected_reference_labels(label, native): """The exact event_time(s) the surface should mark as reference. @@ -534,7 +571,7 @@ def _expected_reference_labels(label, native): row (CallawaySantAnna universal base on a gapped grid can materialize several); CS varying and EfficientDiD estimate e=-1, so none. """ - if label in ("MultiPeriodDiD", "SpilloverDiD"): + if label in ("MultiPeriodDiD", "SpilloverDiD") or label in NATIVE_SURFACE_PRODUCERS: return [] if native.reference_period is None else [native.reference_period] if label == "ChaisemartinDHaultfoeuille": return [0] @@ -599,6 +636,12 @@ def test_no_sentinel_in_schema(label, surfaces): def test_bit_exact_quintet_nonreference(label, surfaces): # Full canonical quintet (att/se/t_stat/p_value/both CI bounds) on every # native non-reference row, across all fourteen producers. + if label in NATIVE_SURFACE_PRODUCERS: + pytest.skip( + "native EventStudyResults producer: the quintet is pinned against the " + "MultiPeriodDiD twin by tests/test_v4_merge_mpd.py and " + "tests/test_twfe_event_study_decoupling.py" + ) native, surface = surfaces[label] non_ref = ~surface.is_reference native_map = _native_event_map(label, native) @@ -637,6 +680,16 @@ def test_mpd_vcov_subblock(surfaces): assert surface.vcov.shape[0] == surface.vcov_index.shape[0] +@pytest.mark.parametrize("label", sorted(NATIVE_SURFACE_PRODUCERS)) +def test_twfe_event_study_vcov_subblock(label, surfaces): + """The TWFE event-study producer exposes the ordered event-study vcov + sub-block over its non-reference labels.""" + _, surface = surfaces[label] + assert surface.vcov is not None and surface.vcov_index is not None + assert surface.vcov.shape == (surface.vcov_index.shape[0], surface.vcov_index.shape[0]) + np.testing.assert_array_equal(surface.vcov_index, surface.event_time[~surface.is_reference]) + + def test_dcdh_convention_and_placebo_merge(surfaces): native, surface = surfaces["ChaisemartinDHaultfoeuille"] assert surface.event_time_convention == "l1_first_switch" @@ -1668,6 +1721,42 @@ def test_mpd_hc2_bm_per_period_df(): assert df_val == res.event_study_df[et] +def test_twfe_event_study_hc2_bm_per_period_df(): + """TWFE event-study twin: the surface's per-period ``df`` carries the + per-coefficient Bell-McCaffrey DOF (they differ across periods on an + unbalanced panel), not a broadcast scalar.""" + from diff_diff import TwoWayFixedEffects + + rng = np.random.default_rng(5) + rows = [] + for u in range(50): + tr = u % 2 + t_max = 8 - (u % 4) + for t in range(1, t_max + 1): + y = 1 + 0.1 * t + 0.3 * tr + (0.8 if tr and t >= 5 else 0.0) + rng.normal(0, 0.3) + rows.append({"unit": u, "time": t, "outcome": y, "treated": tr}) + data = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + # Periods 1..8: the former MultiPeriodDiD midpoint default resolved to [5, 6, 7, 8]. + res = TwoWayFixedEffects(vcov_type="hc2_bm", cluster="unit").fit( + data, + outcome="outcome", + treatment="treated", + time="time", + event_study=True, + spec="pooled", + post_periods=[5, 6, 7, 8], + ) + assert res.df is not None + finite = res.df[~res.is_reference] + finite = finite[np.isfinite(finite)] + assert finite.size > 0 + assert ( + len(set(np.round(finite, 6))) > 1 + ), "per-period BM DOFs should differ on an unbalanced panel" + + def test_lpdid_per_horizon_df(surfaces): native, surface = surfaces["LPDiD"] frame = native.event_study diff --git a/tests/test_honest_did.py b/tests/test_honest_did.py index ea071429a..931632810 100644 --- a/tests/test_honest_did.py +++ b/tests/test_honest_did.py @@ -9,9 +9,10 @@ import numpy as np import pandas as pd import pytest +from event_study_surface_builders import calendar_surface from legacy_surface import legacy_standin -from diff_diff import MultiPeriodDiD +from diff_diff import MultiPeriodDiD, TwoWayFixedEffects from diff_diff.honest_did import ( DeltaRM, DeltaSD, @@ -71,23 +72,8 @@ def simple_panel_data(): @pytest.fixture -def multiperiod_results(simple_panel_data): - """Fit MultiPeriodDiD and return results.""" - mp_did = MultiPeriodDiD() - results = mp_did.fit( - simple_panel_data, - outcome="outcome", - treatment="treated", - time="period", - post_periods=[4, 5, 6, 7], - reference_period=3, - ) - return results - - -@pytest.fixture -def mock_multiperiod_results(): - """Create mock MultiPeriodDiDResults for unit testing. +def legacy_mock_mpd_results(): + """Create mock MultiPeriodDiDResults for the MPD-branch tests (4.0 cut: dies with the class). Simulates a full event-study with pre-period and post-period effects. Reference period is 3 (last pre-period), so period_effects has @@ -150,6 +136,25 @@ def mock_multiperiod_results(): ) +@pytest.fixture +def mock_event_study_results(): + """The same event study as a hand-built calendar ``EventStudyResults`` - what + ``TwoWayFixedEffects().fit(..., event_study=True)`` emits (4.0 cut, PR 4c-i): + pre 0/1/2, reference 3, post 4-7; the event-study vcov is the 7x7 diagonal + sub-block (the MPD mock spreads it into a 20x20 full vcov).""" + rows = { + 0: (0.05, 0.4, 0.125, 0.90, (-0.73, 0.83)), + 1: (-0.02, 0.35, -0.057, 0.95, (-0.71, 0.67)), + 2: (0.08, 0.3, 0.267, 0.79, (-0.51, 0.67)), + 4: (5.0, 0.5, 10.0, 0.0001, (4.02, 5.98)), + 5: (5.2, 0.5, 10.4, 0.0001, (4.22, 6.18)), + 6: (4.8, 0.5, 9.6, 0.0001, (3.82, 5.78)), + 7: (5.0, 0.5, 10.0, 0.0001, (4.02, 5.98)), + } + vcov_diag = [0.4**2, 0.35**2, 0.3**2, 0.5**2, 0.5**2, 0.5**2, 0.5**2] + return calendar_surface(rows, reference=3, post_periods=[4, 5, 6, 7], vcov=np.diag(vcov_diag)) + + # ============================================================================= # Tests for Delta Restriction Classes # ============================================================================= @@ -298,10 +303,10 @@ def test_flci_point_identified(self): class TestParameterExtraction: """Tests for extracting parameters from results objects.""" - def test_extract_from_multiperiod(self, mock_multiperiod_results): + def test_extract_from_multiperiod(self, legacy_mock_mpd_results): """Test extraction from MultiPeriodDiDResults.""" beta_hat, sigma, num_pre, num_post, pre_periods, post_periods, _df = ( - _extract_event_study_params(mock_multiperiod_results) + _extract_event_study_params(legacy_mock_mpd_results) ) # 7 estimated effects: 3 pre (0,1,2) + 4 post (4,5,6,7), ref=3 excluded @@ -312,11 +317,11 @@ def test_extract_from_multiperiod(self, mock_multiperiod_results): assert post_periods == [4, 5, 6, 7] # Verify sub-VCV diagonal matches squared SEs - for i, period in enumerate(sorted(mock_multiperiod_results.period_effects.keys())): - pe = mock_multiperiod_results.period_effects[period] + for i, period in enumerate(sorted(legacy_mock_mpd_results.period_effects.keys())): + pe = legacy_mock_mpd_results.period_effects[period] assert sigma[i, i] == pytest.approx(pe.se**2, abs=1e-10) - def test_zero_se_period_dropped_from_multiperiod(self, mock_multiperiod_results): + def test_zero_se_period_dropped_from_multiperiod(self, legacy_mock_mpd_results): """se == 0 rows drop on the MPD branch (undefined inference). safe_inference treats se <= 0 as undefined inference; admitting @@ -326,10 +331,10 @@ def test_zero_se_period_dropped_from_multiperiod(self, mock_multiperiod_results) """ import dataclasses - pe0 = mock_multiperiod_results.period_effects[0] - mock_multiperiod_results.period_effects[0] = dataclasses.replace(pe0, se=0.0) + pe0 = legacy_mock_mpd_results.period_effects[0] + legacy_mock_mpd_results.period_effects[0] = dataclasses.replace(pe0, se=0.0) beta_hat, sigma, num_pre, num_post, pre_periods, post_periods, _df = ( - _extract_event_study_params(mock_multiperiod_results) + _extract_event_study_params(legacy_mock_mpd_results) ) assert num_pre == 2 and num_post == 4 assert len(beta_hat) == 6 @@ -345,7 +350,7 @@ def test_zero_se_period_dropped_from_multiperiod(self, mock_multiperiod_results) assert post_periods == [4, 5, 6, 7] assert len(pre_periods) == num_pre and len(post_periods) == num_post # End to end: the metadata a user sees relays the same lists. - h = compute_honest_did(mock_multiperiod_results, M=0.5) + h = compute_honest_did(legacy_mock_mpd_results, M=0.5) assert h.pre_periods_used == [1, 2] assert h.post_periods_used == [4, 5, 6, 7] @@ -354,9 +359,7 @@ def test_zero_se_period_dropped_from_multiperiod(self, mock_multiperiod_results) [1, 2, 4, 5], ids=["interior-pre", "ref-adjacent-pre", "first-post", "interior-post"], ) - def test_zero_se_breaking_grid_geometry_fails_closed( - self, mock_multiperiod_results, bad_period - ): + def test_zero_se_breaking_grid_geometry_fails_closed(self, legacy_mock_mpd_results, bad_period): """Interior / reference-adjacent zero-SE drops fail closed. The RR constraint builders index retained coefficients @@ -367,30 +370,30 @@ def test_zero_se_breaking_grid_geometry_fails_closed( """ import dataclasses - pe = mock_multiperiod_results.period_effects[bad_period] - mock_multiperiod_results.period_effects[bad_period] = dataclasses.replace(pe, se=0.0) + pe = legacy_mock_mpd_results.period_effects[bad_period] + legacy_mock_mpd_results.period_effects[bad_period] = dataclasses.replace(pe, se=0.0) with pytest.raises(ValueError, match="consecutive estimated horizons"): - _extract_event_study_params(mock_multiperiod_results) + _extract_event_study_params(legacy_mock_mpd_results) - def test_all_pre_periods_zero_se_fails_closed(self, mock_multiperiod_results): + def test_all_pre_periods_zero_se_fails_closed(self, legacy_mock_mpd_results): """Pin: an all-invalid pre block is rejected with a clear message (never a zero-dimensional restriction build).""" import dataclasses for p in (0, 1, 2): - pe = mock_multiperiod_results.period_effects[p] - mock_multiperiod_results.period_effects[p] = dataclasses.replace(pe, se=0.0) + pe = legacy_mock_mpd_results.period_effects[p] + legacy_mock_mpd_results.period_effects[p] = dataclasses.replace(pe, se=0.0) with pytest.raises(ValueError, match="No pre-period effects"): - _extract_event_study_params(mock_multiperiod_results) + _extract_event_study_params(legacy_mock_mpd_results) - def test_zero_se_trailing_post_dropped_ok(self, mock_multiperiod_results): + def test_zero_se_trailing_post_dropped_ok(self, legacy_mock_mpd_results): """A trailing post-period drop keeps valid positional geometry.""" import dataclasses - pe = mock_multiperiod_results.period_effects[7] - mock_multiperiod_results.period_effects[7] = dataclasses.replace(pe, se=0.0) + pe = legacy_mock_mpd_results.period_effects[7] + legacy_mock_mpd_results.period_effects[7] = dataclasses.replace(pe, se=0.0) beta_hat, sigma, num_pre, num_post, pre_p, post_p, _df = _extract_event_study_params( - mock_multiperiod_results + legacy_mock_mpd_results ) assert num_pre == 3 and num_post == 3 assert pre_p == [0, 1, 2] and post_p == [4, 5, 6] @@ -463,58 +466,58 @@ def test_set_params(self): assert honest.M == 2.0 assert honest.alpha == 0.1 - def test_fit_returns_results(self, mock_multiperiod_results): + def test_fit_returns_results(self, mock_event_study_results): """Test that fit returns HonestDiDResults.""" honest = HonestDiD(method="relative_magnitude", M=1.0) - results = honest.fit(mock_multiperiod_results) + results = honest.fit(mock_event_study_results) assert isinstance(results, HonestDiDResults) assert results.M == 1.0 assert results.method == "relative_magnitude" - def test_fit_negative_M_override_raises(self, mock_multiperiod_results): + def test_fit_negative_M_override_raises(self, mock_event_study_results): """A negative M passed to fit() (bypassing constructor validation) must raise, not be silently treated as +|M| by the FLCI cv abs().""" honest = HonestDiD(method="smoothness", M=0.0) with pytest.raises(ValueError, match="non-negative"): - honest.fit(mock_multiperiod_results, M=-0.1) + honest.fit(mock_event_study_results, M=-0.1) - def test_sensitivity_analysis_negative_M_grid_raises(self, mock_multiperiod_results): + def test_sensitivity_analysis_negative_M_grid_raises(self, mock_event_study_results): """A negative value in M_grid must raise (bypasses constructor validation).""" honest = HonestDiD(method="smoothness") with pytest.raises(ValueError, match="non-negative"): - honest.sensitivity_analysis(mock_multiperiod_results, M_grid=[0.0, -0.1]) + honest.sensitivity_analysis(mock_event_study_results, M_grid=[0.0, -0.1]) - def test_fit_smoothness(self, mock_multiperiod_results): + def test_fit_smoothness(self, mock_event_study_results): """Test fit with smoothness method.""" honest = HonestDiD(method="smoothness", M=0.0) - results = honest.fit(mock_multiperiod_results) + results = honest.fit(mock_event_study_results) assert isinstance(results, HonestDiDResults) assert results.method == "smoothness" assert results.ci_method == "FLCI" - def test_fit_combined(self, mock_multiperiod_results): + def test_fit_combined(self, mock_event_study_results): """Test fit with combined method.""" honest = HonestDiD(method="combined", M=1.0) - results = honest.fit(mock_multiperiod_results) + results = honest.fit(mock_event_study_results) assert isinstance(results, HonestDiDResults) assert results.method == "combined" - def test_fit_override_M(self, mock_multiperiod_results): + def test_fit_override_M(self, mock_event_study_results): """Test that M can be overridden in fit.""" honest = HonestDiD(method="relative_magnitude", M=1.0) - results = honest.fit(mock_multiperiod_results, M=2.0) + results = honest.fit(mock_event_study_results, M=2.0) assert results.M == 2.0 - def test_bounds_widen_with_M(self, mock_multiperiod_results): + def test_bounds_widen_with_M(self, mock_event_study_results): """Test that bounds widen as M increases.""" honest = HonestDiD(method="relative_magnitude") - results_small = honest.fit(mock_multiperiod_results, M=0.5) - results_large = honest.fit(mock_multiperiod_results, M=2.0) + results_small = honest.fit(mock_event_study_results, M=0.5) + results_large = honest.fit(mock_event_study_results, M=2.0) # Larger M should give wider bounds assert ( @@ -525,26 +528,26 @@ def test_bounds_widen_with_M(self, mock_multiperiod_results): class TestSensitivityAnalysis: """Tests for sensitivity analysis functionality.""" - def test_sensitivity_analysis_returns_results(self, mock_multiperiod_results): + def test_sensitivity_analysis_returns_results(self, mock_event_study_results): """Test that sensitivity_analysis returns SensitivityResults.""" honest = HonestDiD(method="relative_magnitude") - sensitivity = honest.sensitivity_analysis(mock_multiperiod_results) + sensitivity = honest.sensitivity_analysis(mock_event_study_results) assert isinstance(sensitivity, SensitivityResults) - def test_sensitivity_analysis_custom_grid(self, mock_multiperiod_results): + def test_sensitivity_analysis_custom_grid(self, mock_event_study_results): """Test sensitivity analysis with custom M grid.""" honest = HonestDiD(method="relative_magnitude") M_grid = [0, 0.5, 1.0, 1.5] - sensitivity = honest.sensitivity_analysis(mock_multiperiod_results, M_grid=M_grid) + sensitivity = honest.sensitivity_analysis(mock_event_study_results, M_grid=M_grid) assert len(sensitivity.M_values) == 4 np.testing.assert_array_equal(sensitivity.M_values, M_grid) - def test_sensitivity_analysis_bounds_list(self, mock_multiperiod_results): + def test_sensitivity_analysis_bounds_list(self, mock_event_study_results): """Test that bounds list has correct length.""" honest = HonestDiD(method="relative_magnitude") - sensitivity = honest.sensitivity_analysis(mock_multiperiod_results) + sensitivity = honest.sensitivity_analysis(mock_event_study_results) assert len(sensitivity.bounds) == len(sensitivity.M_values) assert len(sensitivity.robust_cis) == len(sensitivity.M_values) @@ -553,27 +556,27 @@ def test_sensitivity_analysis_bounds_list(self, mock_multiperiod_results): class TestBreakdownValue: """Tests for breakdown value computation.""" - def test_breakdown_value_type(self, mock_multiperiod_results): + def test_breakdown_value_type(self, mock_event_study_results): """Test that breakdown_value returns float or None.""" honest = HonestDiD(method="relative_magnitude") - breakdown = honest.breakdown_value(mock_multiperiod_results) + breakdown = honest.breakdown_value(mock_event_study_results) assert breakdown is None or isinstance(breakdown, float) - def test_breakdown_value_monotonic(self, mock_multiperiod_results): + def test_breakdown_value_monotonic(self, mock_event_study_results): """Test breakdown value properties.""" honest = HonestDiD(method="relative_magnitude") - breakdown = honest.breakdown_value(mock_multiperiod_results) + breakdown = honest.breakdown_value(mock_event_study_results) if breakdown is not None and breakdown > 0: # Before breakdown, should be significant - result_before = honest.fit(mock_multiperiod_results, M=breakdown * 0.9) + result_before = honest.fit(mock_event_study_results, M=breakdown * 0.9) assert result_before.is_significant # At/after breakdown, should not be significant # We call fit but don't assert on significance since # the binary search tolerance may not match exactly - honest.fit(mock_multiperiod_results, M=breakdown * 1.1) + honest.fit(mock_event_study_results, M=breakdown * 1.1) # ============================================================================= @@ -584,10 +587,10 @@ def test_breakdown_value_monotonic(self, mock_multiperiod_results): class TestHonestDiDResults: """Tests for HonestDiDResults dataclass.""" - def test_results_properties(self, mock_multiperiod_results): + def test_results_properties(self, mock_event_study_results): """Test HonestDiDResults properties.""" honest = HonestDiD(method="relative_magnitude", M=1.0) - results = honest.fit(mock_multiperiod_results) + results = honest.fit(mock_event_study_results) assert hasattr(results, "lb") assert hasattr(results, "ub") @@ -597,37 +600,37 @@ def test_results_properties(self, mock_multiperiod_results): assert hasattr(results, "identified_set_width") assert hasattr(results, "ci_width") - def test_results_is_significant(self, mock_multiperiod_results): + def test_results_is_significant(self, mock_event_study_results): """Test is_significant property.""" honest = HonestDiD(method="relative_magnitude", M=0.0) - results = honest.fit(mock_multiperiod_results) + results = honest.fit(mock_event_study_results) # With M=0 (parallel trends), should be significant for our data # Check that property is a bool assert isinstance(results.is_significant, bool) - def test_results_width_properties(self, mock_multiperiod_results): + def test_results_width_properties(self, mock_event_study_results): """Test width properties are non-negative.""" honest = HonestDiD(method="relative_magnitude", M=1.0) - results = honest.fit(mock_multiperiod_results) + results = honest.fit(mock_event_study_results) assert results.identified_set_width >= 0 assert results.ci_width >= 0 - def test_results_summary(self, mock_multiperiod_results): + def test_results_summary(self, mock_event_study_results): """Test summary method produces string.""" honest = HonestDiD(method="relative_magnitude", M=1.0) - results = honest.fit(mock_multiperiod_results) + results = honest.fit(mock_event_study_results) summary = results.summary() assert isinstance(summary, str) assert "Honest DiD" in summary assert "Rambachan" in summary - def test_results_to_dict(self, mock_multiperiod_results): + def test_results_to_dict(self, mock_event_study_results): """Test to_dict method.""" honest = HonestDiD(method="relative_magnitude", M=1.0) - results = honest.fit(mock_multiperiod_results) + results = honest.fit(mock_event_study_results) d = results.to_dict() assert isinstance(d, dict) @@ -636,10 +639,10 @@ def test_results_to_dict(self, mock_multiperiod_results): assert "M" in d assert "method" in d - def test_results_to_dataframe(self, mock_multiperiod_results): + def test_results_to_dataframe(self, mock_event_study_results): """Test to_dataframe method.""" honest = HonestDiD(method="relative_magnitude", M=1.0) - results = honest.fit(mock_multiperiod_results) + results = honest.fit(mock_event_study_results) df = results.to_dataframe() assert isinstance(df, pd.DataFrame) @@ -649,10 +652,10 @@ def test_results_to_dataframe(self, mock_multiperiod_results): class TestSensitivityResults: """Tests for SensitivityResults dataclass.""" - def test_sensitivity_results_to_dataframe(self, mock_multiperiod_results): + def test_sensitivity_results_to_dataframe(self, mock_event_study_results): """Test to_dataframe method.""" honest = HonestDiD(method="relative_magnitude") - sensitivity = honest.sensitivity_analysis(mock_multiperiod_results) + sensitivity = honest.sensitivity_analysis(mock_event_study_results) df = sensitivity.to_dataframe() assert isinstance(df, pd.DataFrame) @@ -660,10 +663,10 @@ def test_sensitivity_results_to_dataframe(self, mock_multiperiod_results): assert "lb" in df.columns assert "ci_lb" in df.columns - def test_sensitivity_results_summary(self, mock_multiperiod_results): + def test_sensitivity_results_summary(self, mock_event_study_results): """Test summary method.""" honest = HonestDiD(method="relative_magnitude") - sensitivity = honest.sensitivity_analysis(mock_multiperiod_results) + sensitivity = honest.sensitivity_analysis(mock_event_study_results) summary = sensitivity.summary() assert isinstance(summary, str) @@ -677,9 +680,9 @@ def test_sensitivity_results_summary(self, mock_multiperiod_results): class TestConvenienceFunctions: """Tests for convenience functions.""" - def test_compute_honest_did(self, mock_multiperiod_results): + def test_compute_honest_did(self, mock_event_study_results): """Test compute_honest_did function.""" - results = compute_honest_did(mock_multiperiod_results, method="relative_magnitude", M=1.0) + results = compute_honest_did(mock_event_study_results, method="relative_magnitude", M=1.0) assert isinstance(results, HonestDiDResults) assert results.M == 1.0 @@ -693,15 +696,16 @@ def test_compute_honest_did(self, mock_multiperiod_results): class TestIntegration: """Integration tests with real estimators.""" - def test_with_multiperiod_did(self, simple_panel_data): - """Test full pipeline with MultiPeriodDiD.""" - # Fit event study - mp_did = MultiPeriodDiD() - event_results = mp_did.fit( + def test_with_twfe_event_study(self, simple_panel_data): + """Test full pipeline with the TWFE event-study surface.""" + # Fit event study (TWFE event-study mode, pooled design = the former MPD fit) + event_results = TwoWayFixedEffects().fit( simple_panel_data, outcome="outcome", treatment="treated", time="period", + event_study=True, + spec="pooled", post_periods=[4, 5, 6, 7], reference_period=3, ) @@ -717,12 +721,13 @@ def test_with_multiperiod_did(self, simple_panel_data): def test_sensitivity_analysis_integration(self, simple_panel_data): """Test sensitivity analysis with real data.""" - mp_did = MultiPeriodDiD() - event_results = mp_did.fit( + event_results = TwoWayFixedEffects().fit( simple_panel_data, outcome="outcome", treatment="treated", time="period", + event_study=True, + spec="pooled", post_periods=[4, 5, 6, 7], reference_period=3, ) @@ -736,12 +741,13 @@ def test_sensitivity_analysis_integration(self, simple_panel_data): def test_smoothness_method_integration(self, simple_panel_data): """Test smoothness method with real data.""" - mp_did = MultiPeriodDiD() - event_results = mp_did.fit( + event_results = TwoWayFixedEffects().fit( simple_panel_data, outcome="outcome", treatment="treated", time="period", + event_study=True, + spec="pooled", post_periods=[4, 5, 6, 7], reference_period=3, ) @@ -802,23 +808,10 @@ def test_single_post_period_no_pre_effects_raises(self): HonestDiD requires pre-period coefficients for sensitivity analysis. A results object with only post-period effects is not usable. """ - period_effects = { - 4: PeriodEffect( - period=4, effect=5.0, se=0.5, t_stat=10.0, p_value=0.0001, conf_int=(4.02, 5.98) - ), - } - - results = MultiPeriodDiDResults( - period_effects=period_effects, - avg_att=5.0, - avg_se=0.5, - avg_t_stat=10.0, - avg_p_value=0.0001, - avg_conf_int=(4.02, 5.98), - n_obs=200, - n_treated=100, - n_control=100, - pre_periods=[0, 1, 2, 3], + # Only the post period is estimated; the reference (3) is the last pre. + results = calendar_surface( + {4: (5.0, 0.5, 10.0, 0.0001, (4.02, 5.98))}, + reference=3, post_periods=[4], vcov=np.array([[0.25]]), ) @@ -827,20 +820,20 @@ def test_single_post_period_no_pre_effects_raises(self): with pytest.raises(ValueError, match="No pre-period effects with finite"): honest.fit(results) - def test_m_zero_recovers_standard(self, mock_multiperiod_results): + def test_m_zero_recovers_standard(self, mock_event_study_results): """Test that M=0 gives tighter bounds.""" honest = HonestDiD(method="relative_magnitude") - results_0 = honest.fit(mock_multiperiod_results, M=0) - results_1 = honest.fit(mock_multiperiod_results, M=1) + results_0 = honest.fit(mock_event_study_results, M=0) + results_1 = honest.fit(mock_event_study_results, M=1) # M=0 should give tighter or equal bounds assert results_0.ci_width <= results_1.ci_width + 0.01 # Small tolerance - def test_very_large_M(self, mock_multiperiod_results): + def test_very_large_M(self, mock_event_study_results): """Test with very large M value.""" honest = HonestDiD(method="relative_magnitude", M=100) - results = honest.fit(mock_multiperiod_results) + results = honest.fit(mock_event_study_results) # Should still return valid results assert isinstance(results, HonestDiDResults) @@ -1017,108 +1010,76 @@ def test_honest_did_filters_nan_pre_period_effects(self): assert np.isfinite(max_viol) assert max_viol == pytest.approx(0.1, abs=1e-10) # only period 1's |effect| + def test_calendar_surface_filters_nan_pre_period_effects(self): + """The TWFE event-study surface twin of the MPD-branch test above: a NaN + pre-period row (rank-deficient design) is filtered by the calendar + route's `_extract_event_study_params`, not propagated. + + Only the extraction half is twinned: `_estimate_max_pre_violation` has + no surface branch (it falls through to a constant), so that half stays + on the MPD original (4.0 cut, PR 4c-i). + """ + nan_row = (np.nan, np.nan, np.nan, np.nan, (np.nan, np.nan)) + vcov_with_nan = np.full((4, 4), 0.0) + vcov_with_nan[0, :] = np.nan + vcov_with_nan[:, 0] = np.nan + vcov_with_nan[1, 1] = 0.09 + vcov_with_nan[2, 2] = 0.16 + vcov_with_nan[3, 3] = 0.16 + surface = calendar_surface( + { + 0: nan_row, + 1: (0.1, 0.3, 0.33, 0.74, (-0.49, 0.69)), + 3: (2.5, 0.4, 6.25, 0.0001, (1.72, 3.28)), + 4: (2.8, 0.4, 7.0, 0.0001, (2.02, 3.58)), + }, + reference=2, + post_periods=[3, 4], + vcov=vcov_with_nan, + ) + beta_hat, sigma, num_pre, num_post, pre_p, post_p, _df = _extract_event_study_params( + surface + ) + assert len(beta_hat) == 3 # periods 1, 3, 4 (period 0 filtered) + assert num_pre == 1 + assert num_post == 2 + assert np.all(np.isfinite(beta_hat)) + assert np.all(np.isfinite(sigma)) + def test_honest_did_all_pre_nan_raises(self): """HonestDiD should raise ValueError when all pre-period effects are NaN.""" - period_effects = { - 0: PeriodEffect( - period=0, - effect=np.nan, - se=np.nan, - t_stat=np.nan, - p_value=np.nan, - conf_int=(np.nan, np.nan), - ), - # Reference period (1) omitted - 2: PeriodEffect( - period=2, - effect=2.5, - se=0.4, - t_stat=6.25, - p_value=0.0001, - conf_int=(1.72, 3.28), - ), - } - - results = MultiPeriodDiDResults( - period_effects=period_effects, - avg_att=2.5, - avg_se=0.4, - avg_t_stat=6.25, - avg_p_value=0.0001, - avg_conf_int=(1.72, 3.28), - n_obs=200, - n_treated=100, - n_control=100, - pre_periods=[0, 1], + nan_row = (np.nan, np.nan, np.nan, np.nan, (np.nan, np.nan)) + vcov = np.full((2, 2), np.nan) + vcov[1, 1] = 0.16 + results = calendar_surface( + {0: nan_row, 2: (2.5, 0.4, 6.25, 0.0001, (1.72, 3.28))}, + reference=1, post_periods=[2], - vcov=np.diag([0.16]), - reference_period=1, - interaction_indices={0: 0, 2: 1}, + vcov=vcov, ) - with pytest.raises(ValueError, match="No pre-period effects with finite"): _extract_event_study_params(results) def test_honest_did_all_post_nan_raises(self): """HonestDiD should raise ValueError when all post-period effects are NaN. - When MultiPeriodDiD produces NaN for all post-period effects (e.g. from - severe rank deficiency), HonestDiD.fit() should raise rather than + When the event-study fit produces NaN for all post-period effects (e.g. + from severe rank deficiency), HonestDiD.fit() should raise rather than silently computing with an empty weight vector. """ - period_effects = { - 0: PeriodEffect( - period=0, - effect=0.1, - se=0.3, - t_stat=0.33, - p_value=0.74, - conf_int=(-0.49, 0.69), - ), - # Reference period (1) omitted - 2: PeriodEffect( - period=2, - effect=np.nan, - se=np.nan, - t_stat=np.nan, - p_value=np.nan, - conf_int=(np.nan, np.nan), - ), - 3: PeriodEffect( - period=3, - effect=np.nan, - se=np.nan, - t_stat=np.nan, - p_value=np.nan, - conf_int=(np.nan, np.nan), - ), - } - - interaction_indices = {0: 0, 2: 1, 3: 2} + nan_row = (np.nan, np.nan, np.nan, np.nan, (np.nan, np.nan)) vcov = np.full((3, 3), 0.0) vcov[0, 0] = 0.09 vcov[1, :] = np.nan vcov[:, 1] = np.nan vcov[2, :] = np.nan vcov[:, 2] = np.nan - - results = MultiPeriodDiDResults( - period_effects=period_effects, - avg_att=np.nan, - avg_se=np.nan, - avg_t_stat=np.nan, - avg_p_value=np.nan, - avg_conf_int=(np.nan, np.nan), - n_obs=300, - n_treated=150, - n_control=150, - pre_periods=[0, 1], + results = calendar_surface( + {0: (0.1, 0.3, 0.33, 0.74, (-0.49, 0.69)), 2: nan_row, 3: nan_row}, + reference=1, post_periods=[2, 3], vcov=vcov, - reference_period=1, - interaction_indices=interaction_indices, ) - honest = HonestDiD(method="relative_magnitude", M=1.0) with pytest.raises(ValueError, match="No post-period effects with finite"): honest.fit(results) @@ -1560,10 +1521,10 @@ def test_singular_vcov_few_psus_raises_on_both_routes(self, n_bootstrap): class TestVisualizationNoMatplotlib: """Tests for visualization that don't require rendering.""" - def test_sensitivity_results_has_plot_method(self, mock_multiperiod_results): + def test_sensitivity_results_has_plot_method(self, mock_event_study_results): """Test that SensitivityResults has plot method.""" honest = HonestDiD(method="relative_magnitude") - sensitivity = honest.sensitivity_analysis(mock_multiperiod_results) + sensitivity = honest.sensitivity_analysis(mock_event_study_results) assert hasattr(sensitivity, "plot") assert callable(sensitivity.plot) diff --git a/tests/test_methodology_honest_did.py b/tests/test_methodology_honest_did.py index 7a0efe598..d82837ade 100644 --- a/tests/test_methodology_honest_did.py +++ b/tests/test_methodology_honest_did.py @@ -419,37 +419,21 @@ def test_infeasible_smoothness_fit_returns_flci_with_empty_idset(self): returns the FLCI in exactly this case, and so do we (previously the fit NaN-propagated the whole result, silently yielding no inference). """ - from diff_diff.results import MultiPeriodDiDResults, PeriodEffect - - # Non-linear pre-trends: inconsistent with Delta^SD(M=0.0) - period_effects = { - 1: PeriodEffect( - period=1, effect=1.0, se=0.1, t_stat=10.0, p_value=0.0, conf_int=(0.8, 1.2) - ), - 2: PeriodEffect( - period=2, effect=0.0, se=0.1, t_stat=0.0, p_value=1.0, conf_int=(-0.2, 0.2) - ), - 3: PeriodEffect( - period=3, effect=1.0, se=0.1, t_stat=10.0, p_value=0.0, conf_int=(0.8, 1.2) - ), - 5: PeriodEffect( - period=5, effect=2.0, se=0.1, t_stat=20.0, p_value=0.0, conf_int=(1.8, 2.2) - ), - } - results = MultiPeriodDiDResults( - avg_att=2.0, - avg_se=0.1, - avg_t_stat=20.0, - avg_p_value=0.0, - avg_conf_int=(1.8, 2.2), - n_obs=500, - n_treated=250, - n_control=250, - period_effects=period_effects, - pre_periods=[1, 2, 3], + from event_study_surface_builders import calendar_surface + + # Non-linear pre-trends: inconsistent with Delta^SD(M=0.0). Pre 1-3 are + # estimated, the reference is the (absent) label 4 - inserted so no + # estimated coefficient is demoted - and 5 is the post period. + results = calendar_surface( + { + 1: (1.0, 0.1, 10.0, 0.0, (0.8, 1.2)), + 2: (0.0, 0.1, 0.0, 1.0, (-0.2, 0.2)), + 3: (1.0, 0.1, 10.0, 0.0, (0.8, 1.2)), + 5: (2.0, 0.1, 20.0, 0.0, (1.8, 2.2)), + }, + reference=4, post_periods=[5], vcov=np.eye(4) * 0.01, - interaction_indices={1: 0, 2: 1, 3: 2, 5: 3}, ) honest = HonestDiD(method="smoothness", M=0.0) @@ -688,34 +672,19 @@ class TestBreakdownValueMethodology: def test_breakdown_monotonicity(self): """If significant at M=k, should be significant at all M < k.""" - from diff_diff.results import MultiPeriodDiDResults, PeriodEffect - - # Use a weak effect so breakdown is reachable at moderate M - period_effects = { - 1: PeriodEffect( - period=1, effect=0.1, se=0.05, t_stat=2.0, p_value=0.05, conf_int=(0.0, 0.2) - ), - 2: PeriodEffect( - period=2, effect=0.05, se=0.05, t_stat=1.0, p_value=0.32, conf_int=(-0.05, 0.15) - ), - 4: PeriodEffect( - period=4, effect=0.15, se=0.05, t_stat=3.0, p_value=0.003, conf_int=(0.05, 0.25) - ), - } - results = MultiPeriodDiDResults( - avg_att=0.15, - avg_se=0.05, - avg_t_stat=3.0, - avg_p_value=0.003, - avg_conf_int=(0.05, 0.25), - n_obs=500, - n_treated=250, - n_control=250, - period_effects=period_effects, - pre_periods=[1, 2], + from event_study_surface_builders import calendar_surface + + # Use a weak effect so breakdown is reachable at moderate M. Pre 1-2, + # reference 3 (inserted), post 4. + results = calendar_surface( + { + 1: (0.1, 0.05, 2.0, 0.05, (0.0, 0.2)), + 2: (0.05, 0.05, 1.0, 0.32, (-0.05, 0.15)), + 4: (0.15, 0.05, 3.0, 0.003, (0.05, 0.25)), + }, + reference=3, post_periods=[4], vcov=np.eye(3) * 0.0025, - interaction_indices={1: 0, 2: 1, 4: 2}, ) honest = HonestDiD(method="smoothness") diff --git a/tests/test_methodology_pretrends.py b/tests/test_methodology_pretrends.py index ae4cb4dd4..8df273a4c 100644 --- a/tests/test_methodology_pretrends.py +++ b/tests/test_methodology_pretrends.py @@ -1218,50 +1218,30 @@ def test_irregular_grid_gamma_unit_end_to_end_matches_r(self): ``test_mdv_gamma_p_matches_r_slope_for_power`` does not. Locks the PR-B Step 4 γ-unit fix at the public API boundary. """ - from diff_diff.results import MultiPeriodDiDResults, PeriodEffect + from event_study_surface_builders import calendar_surface goldens = self._load_r_golden() fixture = goldens["irregular_pre_periods"] - pre_periods, _, _ = self._extract_python_params(fixture) panel = fixture["panel"] Sigma_full = np.atleast_2d(np.array(panel["Sigma"], dtype=float)) beta_hat = np.array(panel["beta_hat"], dtype=float) - all_periods = list(np.atleast_1d(np.array(panel["all_periods"], dtype=int))) - pre_period_ids = [int(p) for p in pre_periods] + all_periods = [int(p) for p in np.atleast_1d(np.array(panel["all_periods"], dtype=int))] post_period_ids = [ int(p) for p in np.atleast_1d(np.array(panel["post_periods"], dtype=int)) ] - # Build period_effects + interaction_indices over all_periods. + # The fixture's beta_hat + full Sigma as a hand-built TWFE event-study + # surface (4.0 cut, PR 4c-i): rows over all_periods, reference 0. ses_full = np.sqrt(np.diag(Sigma_full)) - period_effects = { - int(p): PeriodEffect( - period=int(p), - effect=float(beta_hat[i]), - se=float(ses_full[i]), - t_stat=0.0, - p_value=0.5, - conf_int=(0.0, 0.0), - ) - for i, p in enumerate(all_periods) - } - interaction_indices = {int(p): i for i, p in enumerate(all_periods)} - - mpd_results = MultiPeriodDiDResults( - period_effects=period_effects, - avg_att=0.0, - avg_se=0.2, - avg_t_stat=0.0, - avg_p_value=0.5, - avg_conf_int=(0.0, 0.0), - n_obs=1000, - n_treated=500, - n_control=500, - pre_periods=pre_period_ids, + mpd_results = calendar_surface( + { + int(p): (float(beta_hat[i]), float(ses_full[i]), 0.0, 0.5, (0.0, 0.0)) + for i, p in enumerate(all_periods) + }, + reference=0, post_periods=post_period_ids, - reference_period=0, vcov=Sigma_full, - interaction_indices=interaction_indices, + vcov_index=all_periods, ) # Target power 0.8 (second entry in r_gamma_p). diff --git a/tests/test_pretrends.py b/tests/test_pretrends.py index ebeff0d26..8b7b5cb8e 100644 --- a/tests/test_pretrends.py +++ b/tests/test_pretrends.py @@ -8,8 +8,9 @@ import numpy as np import pandas as pd import pytest +from event_study_surface_builders import calendar_surface -from diff_diff import MultiPeriodDiD +from diff_diff import TwoWayFixedEffects, generate_did_data from diff_diff.pretrends import ( PreTrendsPower, PreTrendsPowerCurve, @@ -62,85 +63,27 @@ def simple_panel_data(): @pytest.fixture -def multiperiod_results(simple_panel_data): - """Fit MultiPeriodDiD and return results.""" - mp_did = MultiPeriodDiD() - results = mp_did.fit( - simple_panel_data, - outcome="outcome", - treatment="treated", - time="period", - post_periods=[4, 5, 6, 7], - ) - return results - - -@pytest.fixture -def mock_multiperiod_results(): - """Create mock MultiPeriodDiDResults for unit testing. - - This fixture simulates event study results with: - - Pre-periods: 0, 1, 2, 3 (period 3 is reference, omitted from estimation) - - Post-periods: 4, 5, 6, 7 - - Estimated coefficients for periods 0, 1, 2 (pre) and 4, 5, 6, 7 (post) +def mock_event_study_results(): + """Hand-built calendar ``EventStudyResults`` for unit testing (4.0 cut, PR 4c-i: + what ``TwoWayFixedEffects().fit(..., event_study=True)`` emits). + + Simulates a full event study: pre-periods 0, 1, 2 (period 3 is the + reference, the ``is_reference`` row) and post-periods 4, 5, 6, 7; the + event-study vcov is the 7x7 diagonal of the squared SEs. With a declared + reference the calendar route uses the Roth relative-time weights + (``reference - period``), unlike the legacy MPD mock that carried no + reference. """ - # Pre-period effects (excluding reference period 3) - period_effects = { - 0: PeriodEffect( - period=0, effect=0.1, se=0.5, t_stat=0.2, p_value=0.84, conf_int=(-0.88, 1.08) - ), - 1: PeriodEffect( - period=1, effect=-0.05, se=0.5, t_stat=-0.1, p_value=0.92, conf_int=(-1.03, 0.93) - ), - 2: PeriodEffect( - period=2, effect=0.08, se=0.5, t_stat=0.16, p_value=0.87, conf_int=(-0.90, 1.06) - ), - # Period 3 is reference - not in period_effects - # Post-period effects - 4: PeriodEffect( - period=4, effect=5.0, se=0.5, t_stat=10.0, p_value=0.0001, conf_int=(4.02, 5.98) - ), - 5: PeriodEffect( - period=5, effect=5.2, se=0.5, t_stat=10.4, p_value=0.0001, conf_int=(4.22, 6.18) - ), - 6: PeriodEffect( - period=6, effect=4.8, se=0.5, t_stat=9.6, p_value=0.0001, conf_int=(3.82, 5.78) - ), - 7: PeriodEffect( - period=7, effect=5.0, se=0.5, t_stat=10.0, p_value=0.0001, conf_int=(4.02, 5.98) - ), - } - - # Coefficients for estimated periods (excludes reference period 3) - coefficients = { - "treated:period_0": 0.1, - "treated:period_1": -0.05, - "treated:period_2": 0.08, - "treated:period_4": 5.0, - "treated:period_5": 5.2, - "treated:period_6": 4.8, - "treated:period_7": 5.0, + rows = { + 0: (0.1, 0.5, 0.2, 0.84, (-0.88, 1.08)), + 1: (-0.05, 0.5, -0.1, 0.92, (-1.03, 0.93)), + 2: (0.08, 0.5, 0.16, 0.87, (-0.90, 1.06)), + 4: (5.0, 0.5, 10.0, 0.0001, (4.02, 5.98)), + 5: (5.2, 0.5, 10.4, 0.0001, (4.22, 6.18)), + 6: (4.8, 0.5, 9.6, 0.0001, (3.82, 5.78)), + 7: (5.0, 0.5, 10.0, 0.0001, (4.02, 5.98)), } - - # Create vcov matrix (diagonal for simplicity) - # 7 coefficients: 3 pre + 4 post - vcov = np.diag([0.25] * 7) - - return MultiPeriodDiDResults( - period_effects=period_effects, - avg_att=5.0, - avg_se=0.25, - avg_t_stat=20.0, - avg_p_value=0.0001, - avg_conf_int=(4.51, 5.49), - n_obs=800, - n_treated=400, - n_control=400, - pre_periods=[0, 1, 2, 3], # 4 pre-periods, but period 3 is reference - post_periods=[4, 5, 6, 7], - vcov=vcov, - coefficients=coefficients, - ) + return calendar_surface(rows, reference=3, post_periods=[4, 5, 6, 7], vcov=np.diag([0.25] * 7)) # ============================================================================= @@ -378,24 +321,24 @@ def test_mdv_increases_with_target_power(self): class TestPreTrendsPowerFit: """Tests for PreTrendsPower.fit() method.""" - def test_fit_returns_results(self, mock_multiperiod_results): + def test_fit_returns_results(self, mock_event_study_results): """Test that fit returns PreTrendsPowerResults.""" pt = PreTrendsPower() - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) assert isinstance(results, PreTrendsPowerResults) - def test_fit_custom_M(self, mock_multiperiod_results): + def test_fit_custom_M(self, mock_event_study_results): """Test fit with custom M value.""" pt = PreTrendsPower() - results = pt.fit(mock_multiperiod_results, M=1.0) + results = pt.fit(mock_event_study_results, M=1.0) assert results.violation_magnitude == 1.0 - def test_results_has_expected_attributes(self, mock_multiperiod_results): + def test_results_has_expected_attributes(self, mock_event_study_results): """Test that results have all expected attributes.""" pt = PreTrendsPower() - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) assert hasattr(results, "power") assert hasattr(results, "mdv") @@ -408,68 +351,60 @@ def test_results_has_expected_attributes(self, mock_multiperiod_results): assert hasattr(results, "critical_value") assert hasattr(results, "noncentrality") - def test_results_n_pre_periods(self, mock_multiperiod_results): + def test_results_n_pre_periods(self, mock_event_study_results): """Test that n_pre_periods matches estimated pre-periods (excluding reference).""" pt = PreTrendsPower() - results = pt.fit(mock_multiperiod_results) - - # n_pre_periods should be the number of estimated coefficients (3) - # not the total number of pre-periods (4), since period 3 is the reference - expected_n_pre = len( - [ - p - for p in mock_multiperiod_results.pre_periods - if f"treated:period_{p}" in mock_multiperiod_results.coefficients - ] + results = pt.fit(mock_event_study_results) + + # n_pre_periods should be the number of estimated pre coefficients (3): + # the non-reference rows outside the declared post partition. + es = mock_event_study_results + expected_n_pre = int( + np.sum(~es.is_reference & ~np.isin(es.event_time, list(es.post_periods))) ) assert results.n_pre_periods == expected_n_pre assert results.n_pre_periods == 3 # 4 pre-periods minus 1 reference -# ============================================================================= -# Tests for power_curve method -# ============================================================================= - - class TestPowerCurve: """Tests for power_curve method.""" - def test_power_curve_returns_curve(self, mock_multiperiod_results): + def test_power_curve_returns_curve(self, mock_event_study_results): """Test that power_curve returns PreTrendsPowerCurve.""" pt = PreTrendsPower() - curve = pt.power_curve(mock_multiperiod_results) + curve = pt.power_curve(mock_event_study_results) assert isinstance(curve, PreTrendsPowerCurve) - def test_power_curve_custom_grid(self, mock_multiperiod_results): + def test_power_curve_custom_grid(self, mock_event_study_results): """Test power_curve with custom M grid.""" pt = PreTrendsPower() M_grid = [0, 0.5, 1.0, 1.5, 2.0] - curve = pt.power_curve(mock_multiperiod_results, M_grid=M_grid) + curve = pt.power_curve(mock_event_study_results, M_grid=M_grid) assert len(curve.M_values) == len(M_grid) np.testing.assert_array_equal(curve.M_values, M_grid) - def test_power_curve_n_points(self, mock_multiperiod_results): + def test_power_curve_n_points(self, mock_event_study_results): """Test power_curve with custom n_points.""" pt = PreTrendsPower() - curve = pt.power_curve(mock_multiperiod_results, n_points=20) + curve = pt.power_curve(mock_event_study_results, n_points=20) assert len(curve.M_values) == 20 - def test_power_curve_monotonic(self, mock_multiperiod_results): + def test_power_curve_monotonic(self, mock_event_study_results): """Test that power curve is monotonically increasing.""" pt = PreTrendsPower() - curve = pt.power_curve(mock_multiperiod_results) + curve = pt.power_curve(mock_event_study_results) # Power should be non-decreasing diffs = np.diff(curve.powers) assert np.all(diffs >= -0.01) # Allow small numerical noise - def test_power_curve_to_dataframe(self, mock_multiperiod_results): + def test_power_curve_to_dataframe(self, mock_event_study_results): """Test power_curve to_dataframe method.""" pt = PreTrendsPower() - curve = pt.power_curve(mock_multiperiod_results) + curve = pt.power_curve(mock_event_study_results) df = curve.to_dataframe() assert isinstance(df, pd.DataFrame) @@ -485,20 +420,20 @@ def test_power_curve_to_dataframe(self, mock_multiperiod_results): class TestPreTrendsPowerResults: """Tests for PreTrendsPowerResults dataclass.""" - def test_results_summary(self, mock_multiperiod_results): + def test_results_summary(self, mock_event_study_results): """Test summary method produces string.""" pt = PreTrendsPower() - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) summary = results.summary() assert isinstance(summary, str) assert "Pre-Trends Power Analysis" in summary assert "Roth 2022" in summary - def test_results_to_dict(self, mock_multiperiod_results): + def test_results_to_dict(self, mock_event_study_results): """Test to_dict method.""" pt = PreTrendsPower() - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) d = results.to_dict() assert isinstance(d, dict) @@ -506,30 +441,30 @@ def test_results_to_dict(self, mock_multiperiod_results): assert "mdv" in d assert "violation_type" in d - def test_results_to_dataframe(self, mock_multiperiod_results): + def test_results_to_dataframe(self, mock_event_study_results): """Test to_dataframe method.""" pt = PreTrendsPower() - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) df = results.to_dataframe() assert isinstance(df, pd.DataFrame) assert len(df) == 1 - def test_is_informative_property(self, mock_multiperiod_results): + def test_is_informative_property(self, mock_event_study_results): """Test is_informative property.""" pt = PreTrendsPower() - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) assert isinstance(results.is_informative, bool) - def test_power_adequate_property(self, mock_multiperiod_results): + def test_power_adequate_property(self, mock_event_study_results): """Test power_adequate property.""" pt = PreTrendsPower() - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) assert isinstance(results.power_adequate, bool) - def test_power_at_works_for_custom_violation_type(self, mock_multiperiod_results): + def test_power_at_works_for_custom_violation_type(self, mock_event_study_results): """power_at(M) now works for custom violation type (PR-B Step 5). PR-A R18 added a NotImplementedError guard because @@ -542,7 +477,7 @@ def test_power_at_works_for_custom_violation_type(self, mock_multiperiod_results """ weights = np.array([0.1, 0.3, 0.6]) pt = PreTrendsPower(violation_type="custom", violation_weights=weights) - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) # No longer raises; returns a finite power value in [0, 1]. power = results.power_at(0.5) @@ -550,7 +485,7 @@ def test_power_at_works_for_custom_violation_type(self, mock_multiperiod_results assert 0.0 <= power <= 1.0 def test_power_at_raises_on_legacy_custom_result_without_weights( - self, mock_multiperiod_results + self, mock_event_study_results ): """power_at(M) still raises for old serialized results lacking violation_weights (backwards-compat guard). @@ -563,7 +498,7 @@ def test_power_at_raises_on_legacy_custom_result_without_weights( """ weights = np.array([0.1, 0.3, 0.6]) pt = PreTrendsPower(violation_type="custom", violation_weights=weights) - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) # Simulate a legacy-result scenario by clearing the persisted weights. results.violation_weights = None @@ -579,30 +514,30 @@ def test_power_at_raises_on_legacy_custom_result_without_weights( class TestConvenienceFunctions: """Tests for convenience functions.""" - def test_compute_pretrends_power(self, mock_multiperiod_results): + def test_compute_pretrends_power(self, mock_event_study_results): """Test compute_pretrends_power function.""" - results = compute_pretrends_power(mock_multiperiod_results) + results = compute_pretrends_power(mock_event_study_results) assert isinstance(results, PreTrendsPowerResults) - def test_compute_pretrends_power_custom_params(self, mock_multiperiod_results): + def test_compute_pretrends_power_custom_params(self, mock_event_study_results): """Test compute_pretrends_power with custom parameters.""" results = compute_pretrends_power( - mock_multiperiod_results, alpha=0.10, target_power=0.90, violation_type="constant" + mock_event_study_results, alpha=0.10, target_power=0.90, violation_type="constant" ) assert results.alpha == 0.10 assert results.target_power == 0.90 assert results.violation_type == "constant" - def test_compute_mdv(self, mock_multiperiod_results): + def test_compute_mdv(self, mock_event_study_results): """Test compute_mdv function.""" - mdv = compute_mdv(mock_multiperiod_results) + mdv = compute_mdv(mock_event_study_results) assert isinstance(mdv, float) assert mdv > 0 - def test_compute_pretrends_power_rejects_custom_violation_type(self, mock_multiperiod_results): + def test_compute_pretrends_power_rejects_custom_violation_type(self, mock_event_study_results): """compute_pretrends_power(..., violation_type='custom') without explicit ``violation_weights`` must raise ValueError. @@ -615,9 +550,9 @@ def test_compute_pretrends_power_rejects_custom_violation_type(self, mock_multip PreTrendsPower section + docs/methodology/papers/roth-2022-review.md. """ with pytest.raises(ValueError, match="violation_weights"): - compute_pretrends_power(mock_multiperiod_results, violation_type="custom") + compute_pretrends_power(mock_event_study_results, violation_type="custom") - def test_compute_mdv_rejects_custom_violation_type(self, mock_multiperiod_results): + def test_compute_mdv_rejects_custom_violation_type(self, mock_event_study_results): """compute_mdv(..., violation_type='custom') without ``violation_weights`` must raise ValueError. @@ -626,7 +561,7 @@ def test_compute_mdv_rejects_custom_violation_type(self, mock_multiperiod_result to the unsupplied-weights case rather than the entire custom path. """ with pytest.raises(ValueError, match="violation_weights"): - compute_mdv(mock_multiperiod_results, violation_type="custom") + compute_mdv(mock_event_study_results, violation_type="custom") # ============================================================================= @@ -663,37 +598,36 @@ def test_set_params(self): class TestIntegration: """Integration tests with event study results. - Note: These tests use mock results with pre-period coefficients. - MultiPeriodDiD by default only estimates post-period treatment effects, - so we use a fixture that simulates full event study results with - pre-period coefficients (excluding the reference period). + Note: These tests use a hand-built calendar surface with pre-period + coefficients (the reference period is the marked row), as the TWFE + event-study mode emits. """ - def test_with_multiperiod_results(self, mock_multiperiod_results): - """Test full pipeline with MultiPeriodDiDResults.""" + def test_with_event_study_results(self, mock_event_study_results): + """Test full pipeline with the calendar event-study surface.""" # Run pre-trends power analysis pt = PreTrendsPower() - power_results = pt.fit(mock_multiperiod_results) + power_results = pt.fit(mock_event_study_results) # Check results are reasonable assert power_results.power >= 0 assert power_results.power <= 1 assert power_results.mdv > 0 or np.isinf(power_results.mdv) - def test_power_curve_with_results(self, mock_multiperiod_results): + def test_power_curve_with_results(self, mock_event_study_results): """Test power curve with event study results.""" pt = PreTrendsPower() - curve = pt.power_curve(mock_multiperiod_results, n_points=10) + curve = pt.power_curve(mock_event_study_results, n_points=10) # Check curve properties assert len(curve.M_values) == 10 assert len(curve.powers) == 10 assert np.all((curve.powers >= 0) & (curve.powers <= 1)) - def test_sensitivity_to_honest_did(self, mock_multiperiod_results): + def test_sensitivity_to_honest_did(self, mock_event_study_results): """Test sensitivity_to_honest_did method.""" pt = PreTrendsPower() - sensitivity = pt.sensitivity_to_honest_did(mock_multiperiod_results) + sensitivity = pt.sensitivity_to_honest_did(mock_event_study_results) assert "mdv" in sensitivity assert "interpretation" in sensitivity @@ -708,36 +642,36 @@ def test_sensitivity_to_honest_did(self, mock_multiperiod_results): class TestViolationTypes: """Tests for different violation types.""" - def test_linear_violation(self, mock_multiperiod_results): + def test_linear_violation(self, mock_event_study_results): """Test power analysis with linear violation.""" pt = PreTrendsPower(violation_type="linear") - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) assert results.violation_type == "linear" - def test_constant_violation(self, mock_multiperiod_results): + def test_constant_violation(self, mock_event_study_results): """Test power analysis with constant violation.""" pt = PreTrendsPower(violation_type="constant") - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) assert results.violation_type == "constant" - def test_last_period_violation(self, mock_multiperiod_results): + def test_last_period_violation(self, mock_event_study_results): """Test power analysis with last_period violation.""" pt = PreTrendsPower(violation_type="last_period") - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) assert results.violation_type == "last_period" - def test_different_types_give_different_results(self, mock_multiperiod_results): + def test_different_types_give_different_results(self, mock_event_study_results): """Test that different violation types can give different MDV.""" pt_linear = PreTrendsPower(violation_type="linear") pt_constant = PreTrendsPower(violation_type="constant") pt_last = PreTrendsPower(violation_type="last_period") - mdv_linear = pt_linear.fit(mock_multiperiod_results).mdv - mdv_constant = pt_constant.fit(mock_multiperiod_results).mdv - mdv_last = pt_last.fit(mock_multiperiod_results).mdv + mdv_linear = pt_linear.fit(mock_event_study_results).mdv + mdv_constant = pt_constant.fit(mock_event_study_results).mdv + mdv_last = pt_last.fit(mock_event_study_results).mdv # All MDVs should be positive (or inf for degenerate cases) assert mdv_linear > 0 or np.isinf(mdv_linear) @@ -757,41 +691,20 @@ class TestEdgeCases: """Tests for edge cases.""" def test_single_pre_period(self): - """Test with single pre-period (excluding reference). + """Test with only one estimated pre-period. - This tests the case where there's only one estimated pre-period - coefficient. We have pre_periods=[2, 3] where period 3 is the - reference (excluded), leaving only period 2 as estimated. + The reference period is excluded from estimation, so it has no + coefficient: pre labels [2, 3] where period 3 is the reference leave + exactly one estimated pre-period coefficient. """ - period_effects = { - 2: PeriodEffect( - period=2, effect=0.1, se=0.5, t_stat=0.2, p_value=0.84, conf_int=(-0.88, 1.08) - ), - # Period 3 is reference - not estimated - 4: PeriodEffect( - period=4, effect=5.0, se=0.5, t_stat=10.0, p_value=0.0001, conf_int=(4.02, 5.98) - ), - } - - coefficients = { - "treated:period_2": 0.1, - "treated:period_4": 5.0, - } - - results = MultiPeriodDiDResults( - period_effects=period_effects, - avg_att=5.0, - avg_se=0.5, - avg_t_stat=10.0, - avg_p_value=0.0001, - avg_conf_int=(4.02, 5.98), - n_obs=200, - n_treated=100, - n_control=100, - pre_periods=[2, 3], # Period 3 is reference + results = calendar_surface( + { + 2: (0.1, 0.5, 0.2, 0.84, (-0.88, 1.08)), + 4: (5.0, 0.5, 10.0, 0.0001, (4.02, 5.98)), + }, + reference=3, post_periods=[4], vcov=np.array([[0.25, 0], [0, 0.25]]), - coefficients=coefficients, ) pt = PreTrendsPower() @@ -808,50 +721,18 @@ def test_many_pre_periods(self): n_pre_total = 10 n_pre_estimated = 9 # Excluding reference period - # Pre-period effects (excluding reference period 9) - period_effects = {} - for i in range(n_pre_estimated): - period_effects[i] = PeriodEffect( - period=i, - effect=0.05 * (i - 4), - se=0.5, - t_stat=0.1 * (i - 4), - p_value=0.92, - conf_int=(-0.88, 1.08), - ) - - # Post-period effects - for i in range(4): - period_effects[n_pre_total + i] = PeriodEffect( - period=n_pre_total + i, - effect=5.0, - se=0.5, - t_stat=10.0, - p_value=0.0001, - conf_int=(4.02, 5.98), - ) - - # Coefficients (excluding reference period 9) - coefficients = {} - for i in range(n_pre_estimated): - coefficients[f"treated:period_{i}"] = 0.05 * (i - 4) + rows = { + i: (0.05 * (i - 4), 0.5, 0.1 * (i - 4), 0.92, (-0.88, 1.08)) + for i in range(n_pre_estimated) + } for i in range(4): - coefficients[f"treated:period_{n_pre_total + i}"] = 5.0 + rows[n_pre_total + i] = (5.0, 0.5, 10.0, 0.0001, (4.02, 5.98)) - results = MultiPeriodDiDResults( - period_effects=period_effects, - avg_att=5.0, - avg_se=0.5, - avg_t_stat=10.0, - avg_p_value=0.0001, - avg_conf_int=(4.02, 5.98), - n_obs=200, - n_treated=100, - n_control=100, - pre_periods=list(range(n_pre_total)), # Includes reference period 9 + results = calendar_surface( + rows, + reference=n_pre_total - 1, post_periods=list(range(n_pre_total, n_pre_total + 4)), vcov=np.diag([0.25] * (n_pre_estimated + 4)), - coefficients=coefficients, ) pt = PreTrendsPower() @@ -910,10 +791,10 @@ def test_callaway_santanna_universal_base_period(self): class TestVisualization: """Tests for visualization methods (without rendering).""" - def test_power_curve_has_plot_method(self, mock_multiperiod_results): + def test_power_curve_has_plot_method(self, mock_event_study_results): """Test that PreTrendsPowerCurve has plot method.""" pt = PreTrendsPower() - curve = pt.power_curve(mock_multiperiod_results) + curve = pt.power_curve(mock_event_study_results) assert hasattr(curve, "plot") assert callable(curve.plot) @@ -927,73 +808,75 @@ def test_power_curve_has_plot_method(self, mock_multiperiod_results): class TestPreTrendsPowerResultsPowerAt: """Tests for the power_at method on PreTrendsPowerResults.""" - def test_power_at_basic(self, mock_multiperiod_results): + def test_power_at_basic(self, mock_event_study_results): """Test basic power_at functionality.""" pt = PreTrendsPower() - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) - # Compute power at different M values - power_1 = results.power_at(1.0) - power_2 = results.power_at(2.0) - power_5 = results.power_at(5.0) + # Compute power at different M values. With the declared reference the + # calendar route uses the Roth relative-time weights [3, 2, 1], so the + # legacy grid (1/2/5) saturates at 1.0 from M=2; use interior values. + power_1 = results.power_at(0.1) + power_2 = results.power_at(0.25) + power_5 = results.power_at(0.5) # Power should increase with M - assert power_1 < power_2 < power_5 + assert power_1 < power_2 < power_5 < 1 # Power should be between 0 and 1 assert 0 <= power_1 <= 1 assert 0 <= power_2 <= 1 assert 0 <= power_5 <= 1 - def test_power_at_zero(self, mock_multiperiod_results): + def test_power_at_zero(self, mock_event_study_results): """Test power_at with M=0 (should equal alpha under Wald form). See note on TestPowerComputation.test_power_at_zero_equals_alpha: the exact-equals-alpha property is Wald-specific. Pin Wald. """ pt = PreTrendsPower(alpha=0.05, pretest_form="wald") - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) power_0 = results.power_at(0.0) # At M=0, power should equal size (alpha) under Wald. assert np.isclose(power_0, 0.05, atol=0.01) - def test_power_at_matches_fit(self, mock_multiperiod_results): + def test_power_at_matches_fit(self, mock_event_study_results): """Test that power_at gives same result as fitting with that M.""" pt = PreTrendsPower() # Get results from fit - results1 = pt.fit(mock_multiperiod_results, M=2.0) + results1 = pt.fit(mock_event_study_results, M=2.0) # Get power from power_at method - results_base = pt.fit(mock_multiperiod_results) + results_base = pt.fit(mock_event_study_results) power_from_method = results_base.power_at(2.0) # Should be the same (or very close) assert np.isclose(results1.power, power_from_method, rtol=0.01) - def test_power_at_linear_weights(self, mock_multiperiod_results): + def test_power_at_linear_weights(self, mock_event_study_results): """Test power_at uses correct linear weights.""" pt = PreTrendsPower(violation_type="linear") - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) # Power_at should work without error power = results.power_at(1.0) assert 0 <= power <= 1 - def test_power_at_constant_weights(self, mock_multiperiod_results): + def test_power_at_constant_weights(self, mock_event_study_results): """Test power_at uses correct constant weights.""" pt = PreTrendsPower(violation_type="constant") - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) power = results.power_at(1.0) assert 0 <= power <= 1 - def test_power_at_last_period_weights(self, mock_multiperiod_results): + def test_power_at_last_period_weights(self, mock_event_study_results): """Test power_at uses correct last_period weights.""" pt = PreTrendsPower(violation_type="last_period") - results = pt.fit(mock_multiperiod_results) + results = pt.fit(mock_event_study_results) power = results.power_at(1.0) assert 0 <= power <= 1 @@ -1005,7 +888,14 @@ def test_power_at_last_period_weights(self, mock_multiperiod_results): class TestPrePeriodsParameter: - """Tests for the pre_periods parameter in fit and related methods.""" + """Tests for the pre_periods parameter in fit and related methods. + + MPD-only (4.0 cut, PR 4c-i): every test here promotes declared-POST labels + to pre via ``pre_periods=``, which only the MultiPeriodDiD branch accepts; + the TWFE calendar route validates the override against the eligible + declared-pre set (``TestPrePeriodsParameterTWFE`` below). Dies with the + class in PR 4c-ii. + """ @pytest.fixture def event_study_all_periods_results(self): @@ -1124,3 +1014,71 @@ def test_convenience_functions_with_pre_periods(self, event_study_all_periods_re event_study_all_periods_results, M=1.0, pre_periods=[0, 1, 2] ) assert results.n_pre_periods == 3 + + +@pytest.fixture(scope="module") +def twfe_long_pre_results(): + """TWFE pooled event study with a long pre window (declared pre 0-4, reference + 5, post 6-7) for the calendar route's ``pre_periods=`` override tests.""" + df = generate_did_data( + n_units=50, n_periods=8, treatment_period=6, treatment_effect=2.0, seed=20261024 + ) + return TwoWayFixedEffects().fit( + df, + outcome="outcome", + treatment="treated", + time="period", + event_study=True, + spec="pooled", + post_periods=[6, 7], + reference_period=5, + ) + + +class TestPrePeriodsParameterTWFE: + """``pre_periods=`` on the TWFE event-study surface: the override selects a + SUBSET of the eligible declared-pre labels (here 0-4; the reference 5 and + the declared post labels 6-7 are not eligible).""" + + def test_fit_with_explicit_pre_periods(self, twfe_long_pre_results): + pt = PreTrendsPower() + results = pt.fit(twfe_long_pre_results, pre_periods=[2, 3, 4]) + assert results.n_pre_periods == 3 + assert results.power >= 0 + assert results.mdv > 0 + + def test_pre_periods_overrides_results(self, twfe_long_pre_results): + pt = PreTrendsPower() + full = pt.fit(twfe_long_pre_results) + assert full.n_pre_periods == 5 # all declared pre labels 0-4 + results = pt.fit(twfe_long_pre_results, pre_periods=[3, 4]) + assert results.n_pre_periods == 2 + + def test_power_at_with_pre_periods(self, twfe_long_pre_results): + pt = PreTrendsPower() + power = pt.power_at(twfe_long_pre_results, M=1.0, pre_periods=[2, 3, 4]) + assert 0 <= power <= 1 + + def test_power_curve_with_pre_periods(self, twfe_long_pre_results): + pt = PreTrendsPower() + curve = pt.power_curve(twfe_long_pre_results, n_points=10, pre_periods=[2, 3, 4]) + assert len(curve.M_values) == 10 + assert len(curve.powers) == 10 + + def test_sensitivity_to_honest_did_with_pre_periods(self, twfe_long_pre_results): + pt = PreTrendsPower() + sensitivity = pt.sensitivity_to_honest_did(twfe_long_pre_results, pre_periods=[2, 3, 4]) + assert "mdv" in sensitivity + assert sensitivity["mdv"] > 0 + + def test_convenience_functions_with_pre_periods(self, twfe_long_pre_results): + mdv = compute_mdv(twfe_long_pre_results, pre_periods=[2, 3, 4]) + assert mdv > 0 + results = compute_pretrends_power(twfe_long_pre_results, M=1.0, pre_periods=[2, 3, 4]) + assert results.n_pre_periods == 3 + + def test_declared_post_label_is_rejected_as_pre(self, twfe_long_pre_results): + """The calendar route never promotes a declared post label to pre.""" + pt = PreTrendsPower() + with pytest.raises(ValueError): + pt.fit(twfe_long_pre_results, pre_periods=[4, 6]) diff --git a/tests/test_replicate_weight_expansion.py b/tests/test_replicate_weight_expansion.py index 24adb7d9a..f1640e193 100644 --- a/tests/test_replicate_weight_expansion.py +++ b/tests/test_replicate_weight_expansion.py @@ -5,12 +5,12 @@ import numpy as np import pandas as pd import pytest +from event_study_average import post_period_average from legacy_surface import es_dict, group_dict from diff_diff import ( DifferenceInDifferences, ImputationDiD, - MultiPeriodDiD, StackedDiD, SunAbraham, TwoStageDiD, @@ -196,40 +196,46 @@ def test_did_absorb_brr(self): assert np.isfinite(result.se) and result.se > 0 -class TestMultiPeriodDiDReplicate: - """MultiPeriodDiD with replicate weights.""" +class TestEventStudyReplicate: + """The TWFE event-study mode with replicate weights.""" - def test_multiperiod_jk1(self): + def test_event_study_jk1(self): data = _make_simple_panel() rep_cols = _add_jk1_replicates(data, n_rep=10) sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="JK1") - result = MultiPeriodDiD().fit( + result = TwoWayFixedEffects().fit( data, "outcome", "treated", - "time", + time="time", + event_study=True, + spec="pooled", post_periods=[1], survey_design=sd, ) - assert np.isfinite(result.avg_att) - assert np.isfinite(result.avg_se) and result.avg_se > 0 - assert result.survey_metadata is not None + att_bar, se_bar = post_period_average(result) + assert np.isfinite(att_bar) + assert np.isfinite(se_bar) and se_bar > 0 + assert result.df_survey is not None result.summary() - def test_multiperiod_brr(self): + def test_event_study_brr(self): data = _make_simple_panel() rep_cols = _add_brr_replicates(data, n_rep=16) sd = SurveyDesign(weights="weight", replicate_weights=rep_cols, replicate_method="BRR") - result = MultiPeriodDiD().fit( + result = TwoWayFixedEffects().fit( data, "outcome", "treated", - "time", + time="time", + event_study=True, + spec="pooled", post_periods=[1], survey_design=sd, ) - assert np.isfinite(result.avg_att) - assert np.isfinite(result.avg_se) and result.avg_se > 0 + att_bar, se_bar = post_period_average(result) + assert np.isfinite(att_bar) + assert np.isfinite(se_bar) and se_bar > 0 class TestTWFEReplicate: @@ -703,32 +709,37 @@ def test_did_default_vcov_stays_silent(self): ) assert np.isfinite(res.se) and res.se > 0 - def test_multiperiod_explicit_hc2_warns_and_matches_hc1(self): + def test_event_study_explicit_hc2_warns_and_matches_hc1(self): data = _make_simple_panel() rep_cols = _add_brr_replicates(data, n_rep=8) with pytest.warns(UserWarning, match="has no effect with replicate-weight"): - res = MultiPeriodDiD(vcov_type="hc2").fit( + res = TwoWayFixedEffects(vcov_type="hc2").fit( data, "outcome", "treated", - "time", + time="time", + event_study=True, + spec="pooled", post_periods=[1], survey_design=self._sd(rep_cols), ) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") - base = MultiPeriodDiD(vcov_type="hc1").fit( + base = TwoWayFixedEffects(vcov_type="hc1").fit( data, "outcome", "treated", - "time", + time="time", + event_study=True, + spec="pooled", post_periods=[1], survey_design=self._sd(rep_cols), ) # explicit hc1 must NOT trigger the replicate-override warning # (the fixture legitimately emits an unrelated pre-period warning). assert not any("has no effect with replicate-weight" in str(x.message) for x in w) - assert res.avg_att == base.avg_att and res.avg_se == base.avg_se + np.testing.assert_array_equal(res.att, base.att) + np.testing.assert_array_equal(res.se, base.se) def test_conley_keeps_own_survey_contract(self): """conley is excluded from the warn-and-remap: its survey-design @@ -922,27 +933,33 @@ def test_did_absorb_true_brr(self): assert np.isfinite(res.se) and res.se > 0 self._assert_point_invariance(res.att, base.att) - def test_multiperiod_true_brr(self): + def test_event_study_true_brr(self): data = _make_simple_panel() rep_cols = _add_true_brr_replicates(data) - res = MultiPeriodDiD().fit( + res = TwoWayFixedEffects().fit( data, "outcome", "treated", - "time", + time="time", + event_study=True, + spec="pooled", post_periods=[1], survey_design=self._sd(rep_cols), ) - base = MultiPeriodDiD().fit( + base = TwoWayFixedEffects().fit( data, "outcome", "treated", - "time", + time="time", + event_study=True, + spec="pooled", post_periods=[1], survey_design=SurveyDesign(weights="weight"), ) - assert np.isfinite(res.avg_se) and res.avg_se > 0 - self._assert_point_invariance(res.avg_att, base.avg_att) + res_att, res_se = post_period_average(res) + base_att, _ = post_period_average(base) + assert np.isfinite(res_se) and res_se > 0 + self._assert_point_invariance(res_att, base_att) def test_twfe_true_brr(self): # The family's dedicated 2-period panel (the staggered fixture's diff --git a/tests/test_survey.py b/tests/test_survey.py index 7fa9c0438..d226f7905 100644 --- a/tests/test_survey.py +++ b/tests/test_survey.py @@ -5,6 +5,7 @@ import numpy as np import pandas as pd import pytest +from event_study_average import post_period_average from diff_diff import ( DifferenceInDifferences, @@ -920,30 +921,33 @@ def test_twfe_with_survey_design(self, twfe_panel_data): assert result.survey_metadata is not None assert result.survey_metadata.weight_type == "pweight" - def test_multiperiod_with_survey_design(self, multiperiod_data): - """MultiPeriodDiD accepts and uses survey_design.""" - mpd = MultiPeriodDiD() + def test_twfe_event_study_with_survey_design(self, multiperiod_data): + """The TWFE event-study mode accepts and uses survey_design.""" + es = TwoWayFixedEffects() sd = SurveyDesign( weights="weight", strata="stratum", psu="psu", weight_type="pweight", ) - result = mpd.fit( + result = es.fit( multiperiod_data, outcome="outcome", treatment="treated", time="period", + event_study=True, + spec="pooled", post_periods=[3, 4, 5], survey_design=sd, ) + att_bar, se_bar = post_period_average(result) - assert np.isfinite(result.avg_att) - assert np.isfinite(result.avg_se) - assert result.avg_se > 0 + assert np.isfinite(att_bar) + assert np.isfinite(se_bar) + assert se_bar > 0 # Average ATT should be close to 2.5 - assert abs(result.avg_att - 2.5) < 1.5 + assert abs(att_bar - 2.5) < 1.5 def test_fweight_error_for_fractional(self): """Fractional fweights raise ValueError.""" @@ -2594,8 +2598,8 @@ def test_all_certainty_psu_zero_vcov(self): assert np.all(vcov == 0.0) assert not np.any(np.isnan(vcov)) - def test_multiperiod_fweight_df_rounding(self): - """MultiPeriodDiD uses rounded (not truncated) fweight df.""" + def test_twfe_event_study_fweight_df_rounding(self): + """The event-study core uses rounded (not truncated) fweight df.""" np.random.seed(907) n_units = 10 periods = [1, 2, 3] @@ -2615,22 +2619,25 @@ def test_multiperiod_fweight_df_rounding(self): df = pd.DataFrame(rows) sd = SurveyDesign(weights="fw", weight_type="fweight") - mpd = MultiPeriodDiD() + es = TwoWayFixedEffects() with warnings.catch_warnings(): warnings.simplefilter("ignore") - result = mpd.fit( + result = es.fit( df, outcome="outcome", treatment="treated", time="time", + event_study=True, + spec="pooled", post_periods=[3], survey_design=sd, ) + att_bar, se_bar = post_period_average(result) # Should produce valid results (not crash on df truncation) - assert np.isfinite(result.avg_att) - assert np.isfinite(result.avg_se) - assert result.avg_se > 0 + assert np.isfinite(att_bar) + assert np.isfinite(se_bar) + assert se_bar > 0 class TestRound10Fixes: @@ -2925,8 +2932,8 @@ def test_single_absorb_survey_allowed(self): # Should succeed without error assert np.isfinite(result.att) - def test_multiperiod_nonpositive_df_fallback(self): - """MultiPeriodDiD with df=0 falls back to normal distribution.""" + def test_twfe_event_study_nonpositive_df_fallback(self): + """The TWFE event-study mode with df=0 falls back to the normal distribution.""" np.random.seed(42) n = 40 # 4 strata, 1 PSU per stratum → n_PSU=4, n_strata=4, df_survey=0 @@ -2949,23 +2956,26 @@ def test_multiperiod_nonpositive_df_fallback(self): psu="cluster", lonely_psu="adjust", ) - mpd = MultiPeriodDiD() + es = TwoWayFixedEffects() with pytest.warns(UserWarning, match="non-positive"): - result = mpd.fit( + result = es.fit( df, outcome="outcome", treatment="treated", time="time", + event_study=True, + spec="pooled", post_periods=[2, 3], survey_design=sd, ) # Period effects with finite SE > 0 should have finite p-values # (normal distribution fallback, not NaN from t(df=0)) - for period, pe in result.period_effects.items(): - if np.isfinite(pe.se) and pe.se > 0: + for row in np.where(~result.is_reference)[0]: + se, p_value, period = result.se[row], result.p_value[row], result.event_time[row] + if np.isfinite(se) and se > 0: assert np.isfinite( - pe.p_value - ), f"Period {period}: finite SE={pe.se} but p_value={pe.p_value}" + p_value + ), f"Period {period}: finite SE={se} but p_value={p_value}" class TestRound13Fixes: diff --git a/tests/test_survey_phase6.py b/tests/test_survey_phase6.py index ca10cbc0e..c6cd7e532 100644 --- a/tests/test_survey_phase6.py +++ b/tests/test_survey_phase6.py @@ -1665,21 +1665,28 @@ def _replicate_sd_and_data(self): ) return data, sd - def test_multi_period_did_replicate_accepted(self): - """MultiPeriodDiD now supports replicate-weight designs.""" - from diff_diff.estimators import MultiPeriodDiD + def test_twfe_event_study_replicate_accepted(self): + """The TWFE event-study mode supports replicate-weight designs.""" + from event_study_average import post_period_average + + from diff_diff import TwoWayFixedEffects data, sd = self._replicate_sd_and_data() data["treated"] = (data["first_treat"] > 0).astype(int) data["post"] = (data["time"] >= 3).astype(int) - result = MultiPeriodDiD().fit( + # time="post" has labels {0, 1}: the former midpoint default resolved to [1]. + result = TwoWayFixedEffects().fit( data, outcome="outcome", treatment="treated", time="post", + event_study=True, + spec="pooled", + post_periods=[1], survey_design=sd, ) - assert np.isfinite(result.avg_att) + att_bar, _ = post_period_average(result) + assert np.isfinite(att_bar) def test_imputation_did_replicate_accepted(self): """ImputationDiD now supports replicate-weight designs.""" diff --git a/tests/test_twfe_event_study_decoupling.py b/tests/test_twfe_event_study_decoupling.py new file mode 100644 index 000000000..5af910d1e --- /dev/null +++ b/tests/test_twfe_event_study_decoupling.py @@ -0,0 +1,446 @@ +"""PR 4c-i (4.0 cut): the TWFE event-study surface is built from the core's +payload, not from a MultiPeriodDiD container - and nothing moved. + +Three layers of oracle: + +* the COMMITTED oracle ``tests/fixtures/v4_twfe_event_study_oracles.json``, + captured on the unmodified tree (``provenance.captured_at_head``) by + ``tests/_capture_v4_merge_mpd_oracles.py`` for every lane of + ``tests/helpers/twfe_es_oracle_lanes.py`` - the TWFE surface AND its + MultiPeriodDiD twin, every ``EventStudyResults`` field; +* the LIVE MultiPeriodDiD twin (``_from_mpd(MultiPeriodDiD(...).fit(...))``), + which exists until PR 4c-ii deletes the class; +* the payload itself: ``_from_event_study_payload`` and + ``_from_mpd(_mpd_results_from_payload(...))`` must be bit-identical on the + same payload (same inputs, same shared builder). + +Tolerances follow ``tests/test_aggregate_contract.py``: analytic arithmetic at +rtol 1e-12 / atol 1e-13, t-statistics at 2x (a ratio of two bounded columns), +confidence-interval bounds at rtol 1e-9 (scipy's ``t.ppf`` moved between 1.13 +and 1.14+). +""" + +from __future__ import annotations + +import dataclasses +import inspect +import json +import os +import warnings +from typing import Any, Dict + +import numpy as np +import pytest +from event_study_average import post_period_average +from event_study_surface_builders import calendar_surface +from twfe_es_oracle_lanes import FINITE_POOLED_LANES, LANES, PANELS, WITHIN_LANES, Lane + +import diff_diff.twfe +from diff_diff import MultiPeriodDiD, TwoWayFixedEffects +from diff_diff.estimators import _mpd_results_from_payload +from diff_diff.results import MultiPeriodDiDResults, PeriodEffect +from diff_diff.results_base import ( + EventStudyResults, + _EventStudyCorePayload, + _from_event_study_payload, + _from_mpd, + _resolve_scalar_df_survey, + _scalar_df_survey_from_metadata, +) +from diff_diff.survey import SurveyMetadata +from diff_diff.utils import safe_inference + +FIXTURE = os.path.join(os.path.dirname(__file__), "fixtures", "v4_twfe_event_study_oracles.json") + +ANALYTIC_TOL = dict(rtol=1e-12, atol=1e-13) +T_STAT_RTOL_FACTOR = 2.0 +CI_RTOL = 1e-9 +CI_FIELDS = {"conf_int_lower", "conf_int_upper"} +LABEL_FIELDS = {"event_time", "vcov_index", "is_reference"} +TUPLE_FIELDS = {"post_periods", "reference_event_times"} + + +def _load_oracle() -> Dict[str, Any]: + if not os.path.exists(FIXTURE): + pytest.skip("v4_twfe_event_study_oracles.json not present (see tests/fixtures/README.md)") + with open(FIXTURE) as fh: + return json.load(fh) + + +def _decode_array(stored: Dict[str, Any]) -> np.ndarray: + def _none_to_nan(obj: Any) -> Any: + if isinstance(obj, list): + return [_none_to_nan(v) for v in obj] + return np.nan if obj is None else obj + + return np.asarray(_none_to_nan(stored["values"]), dtype=np.dtype(stored["dtype"])) + + +def _fit_lane(lane: Lane): + df = PANELS[lane.panel]() + sd = lane.survey() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + twfe = TwoWayFixedEffects(**lane.twfe_ctor).fit( + df, event_study=True, survey_design=sd, **lane.twfe_fit + ) + mpd_native = MultiPeriodDiD(**lane.mpd_ctor).fit(df, survey_design=sd, **lane.mpd_fit) + assert isinstance(twfe, EventStudyResults) + return df, twfe, mpd_native + + +def _tol_for(field: str) -> Dict[str, float]: + if field in CI_FIELDS: + return dict(rtol=CI_RTOL, atol=ANALYTIC_TOL["atol"]) + if field == "t_stat": + return dict(rtol=T_STAT_RTOL_FACTOR * ANALYTIC_TOL["rtol"], atol=ANALYTIC_TOL["atol"]) + return dict(ANALYTIC_TOL) + + +def _assert_surface_matches_record(live: EventStudyResults, stored: Dict[str, Any]) -> None: + """Every ``dataclasses.fields(EventStudyResults)`` value against its record.""" + for f in dataclasses.fields(live): + name = f.name + value = getattr(live, name) + rec = stored[name] + if isinstance(rec, dict) and "dtype" in rec: + assert value is not None, name + expected = _decode_array(rec) + assert np.asarray(value).dtype.str == rec["dtype"], (name, np.asarray(value).dtype) + assert np.asarray(value).shape == expected.shape, name + if expected.dtype.kind == "f": + np.testing.assert_array_equal(np.isnan(value), np.isnan(expected), err_msg=name) + mask = ~np.isnan(expected) + np.testing.assert_allclose( + np.asarray(value)[mask], expected[mask], err_msg=name, **_tol_for(name) + ) + else: + np.testing.assert_array_equal(np.asarray(value), expected, err_msg=name) + elif name in TUPLE_FIELDS: + assert (tuple(rec) if rec is not None else None) == value, name + elif isinstance(rec, float) and isinstance(value, float): + np.testing.assert_allclose(value, rec, err_msg=name, **_tol_for(name)) + else: + assert rec == value, (name, rec, value) + + +def _assert_surfaces_equal(a: EventStudyResults, b: EventStudyResults, *, exact: bool, skip=()): + for f in dataclasses.fields(a): + name = f.name + if name in skip: + continue + va, vb = getattr(a, name), getattr(b, name) + if isinstance(va, np.ndarray) or isinstance(vb, np.ndarray): + assert va is not None and vb is not None, name + va, vb = np.asarray(va), np.asarray(vb) + assert va.dtype == vb.dtype, name + assert va.shape == vb.shape, name + if va.dtype.kind == "f": + np.testing.assert_array_equal(np.isnan(va), np.isnan(vb), err_msg=name) + m = ~np.isnan(va) + if exact: + np.testing.assert_array_equal(va[m], vb[m], err_msg=name) + else: + np.testing.assert_allclose(va[m], vb[m], err_msg=name, **_tol_for(name)) + else: + np.testing.assert_array_equal(va, vb, err_msg=name) + else: + assert va == vb, (name, va, vb) + + +# --------------------------------------------------------------------------- # +# Committed oracle (captured before the refactor) +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("name", sorted(LANES)) +def test_lane_echo_matches_committed_oracle(name): + """Lane drift is reported as lane drift, not as a numeric regression.""" + oracle = _load_oracle() + assert oracle["lanes"][name]["echo"] == json.loads(json.dumps(LANES[name].echo())) + + +@pytest.mark.parametrize("name", sorted(LANES)) +def test_surface_matches_committed_oracle(name): + """The refactored TWFE event-study surface reproduces the pre-refactor capture.""" + oracle = _load_oracle() + _, twfe, _ = _fit_lane(LANES[name]) + _assert_surface_matches_record(twfe, oracle["lanes"][name]["twfe"]) + + +@pytest.mark.parametrize("name", sorted(LANES)) +def test_mpd_twin_still_matches_committed_oracle(name): + """The refactor left MultiPeriodDiD's own surface untouched too.""" + oracle = _load_oracle() + _, _, mpd_native = _fit_lane(LANES[name]) + _assert_surface_matches_record(_from_mpd(mpd_native), oracle["lanes"][name]["mpd_twin"]) + + +# --------------------------------------------------------------------------- # +# Live oracles (MultiPeriodDiD exists until PR 4c-ii) +# --------------------------------------------------------------------------- # + + +def _payload_for(lane: Lane) -> _EventStudyCorePayload: + df = PANELS[lane.panel]() + fit = dict(lane.mpd_fit) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + est = MultiPeriodDiD(**lane.mpd_ctor) + payload = est._fit_event_study_core( + df, + fit["outcome"], + fit["treatment"], + fit["time"], + fit.get("post_periods"), + fit.get("covariates"), + fit.get("fixed_effects"), + fit.get("absorb"), + fit.get("reference_period"), + fit.get("unit"), + lane.survey(), + "analytical", + _frame_offset=1, + ) + assert isinstance(payload, _EventStudyCorePayload) + return payload + + +@pytest.mark.parametrize("name", sorted(LANES)) +def test_payload_surface_bit_identical_to_from_mpd(name): + """Same payload, same shared builder: the two routes must be byte-equal.""" + payload = _payload_for(LANES[name]) + a = _from_event_study_payload(payload, source="MultiPeriodDiDResults") + b = _from_mpd(_mpd_results_from_payload(payload)) + _assert_surfaces_equal(a, b, exact=True) + + +@pytest.mark.parametrize("name", sorted(LANES)) +def test_twfe_surface_full_container_matches_mpd_twin(name): + """TWFE fit vs the live MultiPeriodDiD twin over EVERY field but the provenance + labels: exact on the pooled lanes (same design), 1e-12 on the within lanes + (the within spec absorbs the unit FE where MPD snaps it).""" + lane = LANES[name] + _, twfe, mpd_native = _fit_lane(lane) + twin = _from_mpd(mpd_native) + _assert_surfaces_equal( + twfe, twin, exact=lane.twfe_fit["spec"] == "pooled", skip=("source", "estimation_spec") + ) + assert twfe.source == "TwoWayFixedEffects" + assert twfe.estimation_spec == lane.twfe_fit["spec"] + assert twin.source == "MultiPeriodDiDResults" + + +def test_core_returns_payload(): + df = PANELS["main"]() + payload = TwoWayFixedEffects()._fit_event_study_core( + df, + "outcome", + "treated", + "period", + [4, 5, 6, 7], + None, + None, + None, + 3, + None, + None, + "analytical", + include_treatment_main=True, + warn_legacy_reference_default=False, + estimator_name="TwoWayFixedEffects", + ) + assert isinstance(payload, _EventStudyCorePayload) + assert payload.periods == sorted(payload.periods) + assert 3 not in payload.periods + for arr in (payload.att, payload.se, payload.t_stat, payload.p_value, payload.ci_lower): + assert arr.shape == (len(payload.periods),) + assert set(payload.df_by_period) == set(payload.periods) + src = inspect.getsource(diff_diff.twfe) + assert "_from_mpd" not in src + assert "MultiPeriodDiDResults" not in src + assert "period_effects" not in src + + +def test_from_mpd_unsorted_dict_keeps_label_value_pairing(): + """The sorted-zip contract: an insertion order of {5, 6, 1, 2} must still pair + every label with ITS PeriodEffect and sort the vcov sub-block with it.""" + effects = {5: 0.5, 6: 0.6, 1: 0.1, 2: 0.2} + idx = {5: 2, 6: 3, 1: 0, 2: 1} + period_effects = { + p: PeriodEffect( + period=p, + effect=e, + se=0.1 * p, + t_stat=1.0, # deliberate literal: the test pins label/value pairing, not inference + p_value=0.5, + conf_int=(e - 1, e + 1), + ) + for p, e in effects.items() # insertion order 5, 6, 1, 2 + } + full_vcov = np.diag([0.01, 0.04, 0.25, 0.36]) + res = MultiPeriodDiDResults( + period_effects=period_effects, + avg_att=0.55, + avg_se=0.1, + avg_t_stat=5.5, + avg_p_value=0.0, + avg_conf_int=(0.3, 0.8), + n_obs=10, + n_treated=5, + n_control=5, + pre_periods=[1, 2, 3], + post_periods=[5, 6], + reference_period=3, + vcov=full_vcov, + interaction_indices=idx, + ) + s = _from_mpd(res) + np.testing.assert_array_equal(s.event_time, [1, 2, 3, 5, 6]) + for row, label in enumerate(s.event_time.tolist()): + if label == 3: + assert s.is_reference[row] + continue + assert s.att[row] == effects[label] + assert s.se[row] == 0.1 * label + assert (s.conf_int_lower[row], s.conf_int_upper[row]) == ( + effects[label] - 1, + effects[label] + 1, + ) + np.testing.assert_array_equal(s.vcov_index, [1, 2, 5, 6]) + np.testing.assert_array_equal(np.diag(s.vcov), [0.01, 0.04, 0.25, 0.36]) + + +def test_df_survey_metadata_split_equals_resolver(): + def _sm(**kw): + return SurveyMetadata( + weight_type="pweight", effective_n=10.0, design_effect=1.0, sum_weights=10.0, **kw + ) + + class _Res: + def __init__(self, sm): + self.survey_metadata = sm + + tsl = _sm(df_survey=56) + rep_defined = _sm(df_survey=14, replicate_method="JK1") + rep_undefined = _sm(df_survey=None, replicate_method="JK1") + for sm, expected in ((tsl, 56.0), (rep_defined, 14.0), (rep_undefined, 0.0), (None, None)): + assert _scalar_df_survey_from_metadata(sm) == expected + assert _resolve_scalar_df_survey(_Res(sm)) == expected + assert isinstance(_scalar_df_survey_from_metadata(tsl), float) + + +# --------------------------------------------------------------------------- # +# The post-period average helper (replaces avg_att / avg_se in ported tests) +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("name", sorted(FINITE_POOLED_LANES)) +def test_post_average_identity_matches_mpd_avg(name): + """`post_period_average` reproduces MultiPeriodDiD's avg_att / avg_se EXACTLY + on the pooled twin (the live pin that licenses the ported sites' rewrite).""" + _, twfe, mpd_native = _fit_lane(LANES[name]) + att_bar, se_bar = post_period_average(twfe) + assert att_bar == mpd_native.avg_att + assert se_bar == mpd_native.avg_se + + +def _surface(att, se, vcov, post, ref): + """Hand-built calendar surface for the helper's NaN lanes: per-row inference + comes from ``safe_inference`` (all-NaN on the reference row and on any + undefined-SE row), never from an inline ratio.""" + att = np.asarray(att, dtype=float).copy() + se = np.asarray(se, dtype=float).copy() + labels = np.arange(len(att)) + is_ref = labels == ref + att[is_ref] = np.nan + se[is_ref] = np.nan + t = np.full(len(att), np.nan) + p = np.full(len(att), np.nan) + lo = np.full(len(att), np.nan) + hi = np.full(len(att), np.nan) + for i in np.where(~is_ref)[0]: + t[i], p[i], (lo[i], hi[i]) = safe_inference(float(att[i]), float(se[i]), alpha=0.05) + return EventStudyResults( + event_time=labels, + att=att, + se=se, + t_stat=t, + p_value=p, + conf_int_lower=lo, + conf_int_upper=hi, + is_reference=is_ref, + n=np.full(len(att), np.nan), + time_scale="calendar", + source="TwoWayFixedEffects", + post_periods=tuple(post), + vcov=np.asarray(vcov, dtype=float), + vcov_index=labels[~is_ref], + ) + + +def test_post_period_average_nan_contract(): + # (d) healthy: hand arithmetic + s = _surface([0.1, 0.0, 1.0, 2.0], [0.1, 0.1, 0.2, 0.3], np.diag([0.01, 0.04, 0.09]), [2, 3], 1) + att_bar, se_bar = post_period_average(s) + assert att_bar == 1.5 + assert se_bar == pytest.approx(np.sqrt((0.04 + 0.09) / 4)) + # (a) one post att NaN -> (nan, nan) + s = _surface( + [0.1, 0.0, np.nan, 2.0], [0.1, 0.1, 0.2, 0.3], np.diag([0.01, 0.04, 0.09]), [2, 3], 1 + ) + att_bar, se_bar = post_period_average(s) + assert np.isnan(att_bar) and np.isnan(se_bar) + # (b) finite att but a NaN entry in the post vcov block -> finite att_bar, NaN se_bar + v = np.diag([0.01, 0.04, 0.09]).astype(float) + v[2, 2] = np.nan + s = _surface([0.1, 0.0, 1.0, 2.0], [0.1, 0.1, 0.2, 0.3], v, [2, 3], 1) + att_bar, se_bar = post_period_average(s) + assert att_bar == 1.5 and np.isnan(se_bar) + # (c) negative summed variance -> NaN se_bar, no RuntimeWarning + v = np.array([[0.01, 0.0, 0.0], [0.0, 0.04, -0.5], [0.0, -0.5, 0.09]]) + s = _surface([0.1, 0.0, 1.0, 2.0], [0.1, 0.1, 0.2, 0.3], v, [2, 3], 1) + with warnings.catch_warnings(): + warnings.simplefilter("error") + att_bar, se_bar = post_period_average(s) + assert att_bar == 1.5 and np.isnan(se_bar) + # (e) a REAL fit: the rank-deficient NaN post period, TWFE pooled vs MPD twin + lane = LANES["pooled_nan_post_period"] + _, twfe, mpd_native = _fit_lane(lane) + row3 = int(np.where(twfe.event_time == 3)[0][0]) + assert np.isnan(twfe.att[row3]) + att_bar, se_bar = post_period_average(twfe) + assert np.isnan(att_bar) and np.isnan(mpd_native.avg_att) + assert np.isnan(se_bar) and np.isnan(mpd_native.avg_se) + + +def test_within_lanes_are_the_absorb_twins(): + """Guards the lane table itself: every within lane pairs spec='within' + unit + with an MPD absorb=['unit'] fit (the 1e-12 twin the merge tests pin).""" + for name in WITHIN_LANES: + lane = LANES[name] + assert lane.twfe_fit["unit"] == "unit" + assert lane.mpd_fit["absorb"] == ["unit"] + + +def test_calendar_surface_builder_uses_safe_inference(): + """The hand-built-surface helper derives 2-tuple rows through safe_inference: + alpha drives the interval level and undefined SEs fail closed to all-NaN.""" + from scipy import stats + + s90 = calendar_surface( + {1: (1.0, 0.5), 3: (2.0, 0.25)}, reference=2, post_periods=[3], alpha=0.10 + ) + s95 = calendar_surface({1: (1.0, 0.5), 3: (2.0, 0.25)}, reference=2, post_periods=[3]) + row = int(np.where(s90.event_time == 3)[0][0]) + z90, z95 = stats.norm.ppf(0.95), stats.norm.ppf(0.975) + np.testing.assert_allclose(s90.conf_int_upper[row] - s90.att[row], z90 * 0.25, rtol=1e-12) + np.testing.assert_allclose(s95.conf_int_upper[row] - s95.att[row], z95 * 0.25, rtol=1e-12) + assert s90.conf_int_upper[row] < s95.conf_int_upper[row] + for bad_se in (0.0, -1.0, np.nan, np.inf): + s = calendar_surface({1: (1.0, bad_se), 3: (2.0, 0.25)}, reference=2, post_periods=[3]) + r = int(np.where(s.event_time == 1)[0][0]) + assert np.isnan(s.t_stat[r]) and np.isnan(s.p_value[r]) + assert np.isnan(s.conf_int_lower[r]) and np.isnan(s.conf_int_upper[r]) + ref = int(np.where(s95.event_time == 2)[0][0]) + assert s95.is_reference[ref] and np.isnan(s95.se[ref]) and np.isnan(s95.p_value[ref]) diff --git a/tests/test_variance_conventions.py b/tests/test_variance_conventions.py index cdc690470..75b4772ff 100644 --- a/tests/test_variance_conventions.py +++ b/tests/test_variance_conventions.py @@ -353,6 +353,52 @@ def snapshot(self): "subtraction is caught here; the unit-cluster rows catch over-)" ), ), + dict( + key="twfe_es_within_hc1_cluster_unit", + fit=lambda df: diff_diff.TwoWayFixedEffects(cluster="unit").fit( + df, + outcome="y", + treatment="grp", + unit="unit", + time="time", + event_study=True, + spec="within", + post_periods=[4, 5, 6], # the former MultiPeriodDiD midpoint default + reference_period=1, + ), + # tail_df: one entry per safe_inference call the shared core makes - + # 5 periods + the post-period average (the average block dies with + # MultiPeriodDiD at the 4.0 cut; re-pin then). + cr1_k=(11,), + tail_df=(290.0,) * 6, + status="legitimate", + reason=( + "TWFE event-study, within spec (4.0 cut, PR 4c-i): the absorbed unit FE " + "plus the built-in period dummies give the same K_reference as the " + "mpd_absorb row - 6 visible + rank(time given unit) = 5" + ), + ), + dict( + key="twfe_es_pooled_hc1_cluster_time", + fit=lambda df: diff_diff.TwoWayFixedEffects(cluster="time").fit( + df, + outcome="y", + treatment="grp", + time="time", + event_study=True, + spec="pooled", + post_periods=[4, 5, 6], # the former MultiPeriodDiD midpoint default + reference_period=1, + ), + cr1_k=(7,), + tail_df=(348.0,) * 6, + status="legitimate", + reason=( + "TWFE event-study, pooled spec (the MultiPeriodDiD design; 4.0 cut, PR " + "4c-i): the NESTED orientation of the built-in period dummies under a " + "time cluster - 12 visible minus their rank 5, as mpd_plain_hc1_cluster_time" + ), + ), dict( key="lpdid_absorb_nested_cluster_grp", fit=lambda df: diff_diff.LPDiD(pre_window=2, post_window=2, cluster="grp").fit( diff --git a/tests/test_visualization.py b/tests/test_visualization.py index 7d0c97de3..2133e5ed0 100644 --- a/tests/test_visualization.py +++ b/tests/test_visualization.py @@ -10,6 +10,7 @@ from diff_diff import ( CallawaySantAnna, MultiPeriodDiD, + TwoWayFixedEffects, plot_event_study, ) @@ -80,12 +81,31 @@ def generate_staggered_data( ) +@pytest.fixture +def twfe_es_results(): + """TWFE event-study surface (pooled design) on the multi-period panel - the + successor of the MultiPeriodDiD fixture for the plot-option tests.""" + data = generate_multi_period_data() + return TwoWayFixedEffects().fit( + data, + outcome="outcome", + treatment="treated", + time="period", + event_study=True, + spec="pooled", + post_periods=[2, 3], + reference_period=1, + ) + + class TestPlotEventStudy: """Tests for plot_event_study function.""" @pytest.fixture - def multi_period_results(self): - """Fixture for MultiPeriodDiD results.""" + def legacy_mpd_results(self): + """Fixture for native MultiPeriodDiD results - the only remaining pin of + plot_event_study's ``period_effects`` branch (4.0 cut, PR 4c-i; dies with + the class in PR 4c-ii).""" data = generate_multi_period_data() did = MultiPeriodDiD() return did.fit( @@ -115,11 +135,18 @@ def cs_results(self, cs_fit): """The post-fit event-study container - what plot_event_study consumes.""" return cs_fit.aggregate("event_study") - def test_plot_from_multi_period_results(self, multi_period_results): - """Test plotting from MultiPeriodDiD results.""" + def test_plot_from_multi_period_results(self, legacy_mpd_results): + """Test plotting from native MultiPeriodDiD results (the period_effects branch).""" + pytest.importorskip("matplotlib") + + ax = plot_event_study(legacy_mpd_results, show=False) + assert ax is not None + + def test_plot_from_twfe_event_study_results(self, twfe_es_results): + """Test plotting from the TWFE event-study surface.""" pytest.importorskip("matplotlib") - ax = plot_event_study(multi_period_results, show=False) + ax = plot_event_study(twfe_es_results, show=False) assert ax is not None def test_plot_from_cs_results(self, cs_results): @@ -260,12 +287,12 @@ def test_plot_from_dict(self): ax = plot_event_study(effects=effects, se=se, reference_period=0, show=False) assert ax is not None - def test_plot_customization(self, multi_period_results): + def test_plot_customization(self, twfe_es_results): """Test plot customization options.""" pytest.importorskip("matplotlib") ax = plot_event_study( - multi_period_results, + twfe_es_results, title="Custom Title", xlabel="Custom X", ylabel="Custom Y", @@ -279,20 +306,20 @@ def test_plot_customization(self, multi_period_results): assert ax.get_xlabel() == "Custom X" assert ax.get_ylabel() == "Custom Y" - def test_plot_no_zero_line(self, multi_period_results): + def test_plot_no_zero_line(self, twfe_es_results): """Test disabling zero line.""" pytest.importorskip("matplotlib") - ax = plot_event_study(multi_period_results, show_zero_line=False, show=False) + ax = plot_event_study(twfe_es_results, show_zero_line=False, show=False) assert ax is not None - def test_plot_with_existing_axes(self, multi_period_results): + def test_plot_with_existing_axes(self, twfe_es_results): """Test plotting on existing axes.""" matplotlib = pytest.importorskip("matplotlib") plt = matplotlib.pyplot fig, ax = plt.subplots() - returned_ax = plot_event_study(multi_period_results, ax=ax, show=False) + returned_ax = plot_event_study(twfe_es_results, ax=ax, show=False) assert returned_ax is ax plt.close() @@ -760,8 +787,8 @@ def test_plot_use_cband_false(self, cs_cband_results): class TestPlotEventStudyIntegration: """Integration tests for event study plotting.""" - def test_full_workflow_multi_period(self): - """Test full workflow with MultiPeriodDiD.""" + def test_full_workflow_twfe_event_study(self): + """Test full workflow with the TWFE event-study mode.""" matplotlib = pytest.importorskip("matplotlib") plt = matplotlib.pyplot @@ -769,12 +796,13 @@ def test_full_workflow_multi_period(self): data = generate_multi_period_data() # Fit model - did = MultiPeriodDiD() - results = did.fit( + results = TwoWayFixedEffects().fit( data, outcome="outcome", treatment="treated", time="period", + event_study=True, + spec="pooled", post_periods=[2, 3], reference_period=1, ) diff --git a/tests/test_visualization_plotly.py b/tests/test_visualization_plotly.py index 309ee053e..47cc306ea 100644 --- a/tests/test_visualization_plotly.py +++ b/tests/test_visualization_plotly.py @@ -8,6 +8,7 @@ import numpy as np import pytest +from event_study_surface_builders import calendar_surface go = pytest.importorskip("plotly.graph_objects") @@ -65,20 +66,21 @@ def honest_results(): results = MagicMock() results.alpha = 0.05 results.M = 1.0 + # Calendar labels 1-5: pre 1, 2; reference 3 (the last pre-period, as the + # calendar route requires); post 4, 5. The original results are a hand-built + # TWFE event-study surface (4.0 cut, PR 4c-i) whose STORED intervals the + # container branch plots as the original CIs. results.event_study_bounds = { - -2: {"ci_lb": -0.3, "ci_ub": 0.3}, - -1: {"ci_lb": -0.2, "ci_ub": 0.2}, - 0: {"ci_lb": 0.1, "ci_ub": 0.9}, - 1: {"ci_lb": 0.2, "ci_ub": 1.0}, + 1: {"ci_lb": -0.3, "ci_ub": 0.3}, + 2: {"ci_lb": -0.2, "ci_ub": 0.2}, + 4: {"ci_lb": 0.1, "ci_ub": 0.9}, + 5: {"ci_lb": 0.2, "ci_ub": 1.0}, } - original = MagicMock() - original.period_effects = { - -2: MagicMock(effect=0.05, se=0.1), - -1: MagicMock(effect=0.0, se=0.1), - 0: MagicMock(effect=0.5, se=0.12), - 1: MagicMock(effect=0.6, se=0.13), - } - results.original_results = original + results.original_results = calendar_surface( + {1: (0.05, 0.1), 2: (0.0, 0.1), 4: (0.5, 0.12), 5: (0.6, 0.13)}, + reference=3, + post_periods=[4, 5], + ) return results From aee1cc40d68c7174aacf9f0b926330ae2ca57bae Mon Sep 17 00:00:00 2001 From: igerber Date: Sun, 25 Oct 2026 10:53:12 -0400 Subject: [PATCH 2/2] test(v4): bound the oracle t_stat by its propagated att/se tolerance CI on every Linux and Windows leg failed the two committed-oracle tests on the rank-deficient pooled_nan_post_period lane: the near-zero period-0 effect (att -3.8e-3, se 0.29, t -0.013) drifted 1.6e-13 (OpenBLAS) and 2.4e-13 (Windows) in t_stat against the macOS Accelerate capture while att and se each passed at atol 1e-13. A ratio inherits atol / se in absolute terms, so a flat atol on t_stat is tighter than the inputs' own contract whenever se < 1. Replace the flat tolerance with the first-order propagated bound |d t| <= atol (1 + |t|) / se + 2 rtol |t| computed from the stored att/se (admits the observed CI values with 2x margin; rejects 1e-12 on the small row and 1e-10 on the t = 10 rows). Fixture README notes the rule. --- tests/fixtures/README.md | 6 ++- tests/test_twfe_event_study_decoupling.py | 46 ++++++++++++++++++++--- 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/tests/fixtures/README.md b/tests/fixtures/README.md index ce71b9df6..2538d2b43 100644 --- a/tests/fixtures/README.md +++ b/tests/fixtures/README.md @@ -93,6 +93,10 @@ Arrays are stored as `{"dtype": , "values": [...]}` (the real dtype, e.g. `" Dict[str, float]: if field in CI_FIELDS: return dict(rtol=CI_RTOL, atol=ANALYTIC_TOL["atol"]) - if field == "t_stat": - return dict(rtol=T_STAT_RTOL_FACTOR * ANALYTIC_TOL["rtol"], atol=ANALYTIC_TOL["atol"]) return dict(ANALYTIC_TOL) +def _assert_t_stat_within_propagated_bound( + live_t: np.ndarray, stored_t: np.ndarray, stored_att: np.ndarray, stored_se: np.ndarray +) -> None: + """``t_stat = att / se`` checked against the bound its inputs' tolerances imply. + + With ``|d att| <= atol + rtol |att|`` and ``|d se| <= atol + rtol |se|`` (the + ``ANALYTIC_TOL`` contract that ``att`` and ``se`` pass individually), first-order + propagation gives ``|d t| <= atol (1 + |t|) / se + 2 rtol |t|``. A flat + ``atol=1e-13`` on ``t`` is tighter than that whenever ``se < 1``: on the + rank-deficient ``pooled_nan_post_period`` lane the near-zero period-0 effect + (``att`` -3.8e-3, ``se`` 0.29, ``t`` -0.013) drifts 1.6e-13 (Linux OpenBLAS) and + 2.4e-13 (Windows) against the macOS Accelerate capture while ``att`` and ``se`` + both pass, exactly the ``atol / se`` amplification this bound allows for. + """ + mask = ~np.isnan(stored_t) + np.testing.assert_array_equal(np.isnan(live_t), ~mask, err_msg="t_stat") + atol, rtol = ANALYTIC_TOL["atol"], ANALYTIC_TOL["rtol"] + t, se = stored_t[mask], stored_se[mask] + bound = atol * (1.0 + np.abs(t)) / se + T_STAT_RTOL_FACTOR * rtol * np.abs(t) + diff = np.abs(live_t[mask] - t) + assert np.all(diff <= bound), ( + "t_stat outside the propagated att/se bound", + dict(diff=diff, bound=bound, live=live_t[mask], stored=t, att=stored_att[mask], se=se), + ) + + def _assert_surface_matches_record(live: EventStudyResults, stored: Dict[str, Any]) -> None: """Every ``dataclasses.fields(EventStudyResults)`` value against its record.""" for f in dataclasses.fields(live): @@ -108,7 +135,14 @@ def _assert_surface_matches_record(live: EventStudyResults, stored: Dict[str, An expected = _decode_array(rec) assert np.asarray(value).dtype.str == rec["dtype"], (name, np.asarray(value).dtype) assert np.asarray(value).shape == expected.shape, name - if expected.dtype.kind == "f": + if name == "t_stat": + _assert_t_stat_within_propagated_bound( + np.asarray(value, dtype=float), + expected, + _decode_array(stored["att"]), + _decode_array(stored["se"]), + ) + elif expected.dtype.kind == "f": np.testing.assert_array_equal(np.isnan(value), np.isnan(expected), err_msg=name) mask = ~np.isnan(expected) np.testing.assert_allclose(