diff --git a/CHANGELOG.md b/CHANGELOG.md index 701d213aa..f13c13c8c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,8 +18,28 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 7 event-study horizons, confirming the no-finite-sample-factor convention (REGISTRY LPDiD Deviation 2). This is the repo's first Stata parity arm — `teffects` is native (no SSC dependency) and goldens are committed, so CI needs no Stata. +- **4.0 API-unification program, Phase 2 results-contract foundation.** Three new + public symbols in `diff_diff.results_base`: `BaseResults` (shared estimator-results + base, behaviorally inert in 3.9), `Diagnostic` (marker base applied to the diagnostic + result roster - Bacon, RDPlot, HonestDiD, pre-trends/power, placebo, HAD pretests, + DiagnosticReport results - which carry no inference row), and `EventStudyResults` (the + unified per-event-time representation, with package-internal builders over all 14 + event-study producers; public exposure rides `aggregate(type="event_study")` in a + later Phase 2 PR). Seven estimator results classes gained `to_dict()` + (CallawaySantAnna, StackedDiD, ContinuousDiD, SunAbraham, Wooldridge, + ChaisemartinDHaultfoeuille, Bacon). Ledger row M-091 flips to `done`; M-092 (unified + event-study surface) and M-093 (4.0 sentinel retirement, phase 5) are added. + **Estimator equations, weighting, variance, and numerical output are unchanged.** ### Changed +- **Diagnostic input validation in the report consumers (4.0 program Phase 2).** + `BusinessReport`, `practitioner_next_steps`, and `DiagnosticReport` now route + diagnostics by the new `Diagnostic` marker instead of by result-class name. Passing a + marked non-Bacon diagnostic (e.g. `HonestDiDResults`, `PlaceboTestResults`) as the + primary estimator input now raises `TypeError` - previously `BusinessReport` + special-cased only Bacon by name, and `DiagnosticReport` silently produced a + zero-check report. Estimator results are unaffected; Bacon's existing read-out is + retained. - **Internal tracking docs reorganized (no library behavior change).** `TODO.md` is now the actionable backlog only; blocked / parked work and won't-fix decisions moved to the new root-level `DEFERRED.md` (deferral & decision registry, same blocker sections plus diff --git a/TODO.md b/TODO.md index 33ed91d3f..21a738b89 100644 --- a/TODO.md +++ b/TODO.md @@ -23,6 +23,7 @@ Related tracking surfaces: |-------|----------|--------|--------|----------| | `SyntheticControl` conformal (CWZ 2021) AR / innovation-permutation path (Lemmas 5-7) for time-series proxies — the residual-permutation shortcut is only valid for time-permutation-invariant proxies (SC/Lasso/DiD); an AR proxy needs innovation permutation. | `diff_diff/conformal.py`, `diff_diff/synthetic_control_results.py` | CWZ-2021 | Heavy | Low | | `ContinuousDiD` CGBS-2024 remaining extensions (earlier phases — `covariates=` reg/dr, `treatment_type="discrete"`, single-cohort `control_group="lowest_dose"` with estimand `ATT(d)−ATT(d_L)` — are already supported; see REGISTRY Note #7). Remaining (all deferred `NotImplementedError`, documented): `estimation_method="ipw"` on the dose curve (scalar-adjustment / degenerate); `covariates=` × `survey_design=` (weighted OR + weighted nuisance IF); multi-cohort **heterogeneous-support** discrete aggregation (support-aware: average each dose only over the cohorts that observe it); **multi-cohort `lowest_dose`** (within-cohort `d_L` reference + support-aware cross-cohort aggregation); and **`covariates=` × `lowest_dose`** (conditional-PT-relative-to-`d_L` estimand). Single-cohort / 2-period / shared-support multi-cohort are supported. | `continuous_did.py` | CGBS-2024 | Heavy | Low | +| `EventStudyResults` full-vcov + df round-trip for more producers (4.0 program Phase 2 follow-up). StackedDiD and TwoStageDiD compute a full event-study VCV (regression / Gardner GMM) internally but retain only marginal SEs on their result containers, so their unified surfaces carry `vcov=None`; persist the event-time submatrix + ordered horizon index on those containers and thread it in. Likewise thread producer-specific inference df into `EventStudyResults.df` (MPD `inference_df`, explicit-survey CS `survey_metadata.df_survey`, HC2-BM horizon-specific contrast df). The adapter is faithful to what containers expose today; this is additive container work, not a defect. | `diff_diff/results_base.py`, `diff_diff/stacked_did_results.py`, `diff_diff/two_stage_results.py` | #M-092 | Mid | Low | ### Performance diff --git a/diff_diff/__init__.py b/diff_diff/__init__.py index a526901e7..fc8f8a15c 100644 --- a/diff_diff/__init__.py +++ b/diff_diff/__init__.py @@ -228,6 +228,11 @@ SpilloverDiDResults, # re-export SyntheticDiDResults, ) +from diff_diff.results_base import ( + BaseResults, + Diagnostic, + EventStudyResults, +) from diff_diff.spillover import ( SpilloverDiD, ) @@ -577,6 +582,10 @@ "MeridianROIPrior", # LLM guide accessor "get_llm_guide", + # Results-contract foundations (4.0 program Phase 2, spec sections 3.5/5) + "BaseResults", + "Diagnostic", + "EventStudyResults", ] # Agent-facing entrypoints surface first in dir(diff_diff). LLM agents diff --git a/diff_diff/_reporting_helpers.py b/diff_diff/_reporting_helpers.py index e1800d1b0..7ce1422c3 100644 --- a/diff_diff/_reporting_helpers.py +++ b/diff_diff/_reporting_helpers.py @@ -21,6 +21,8 @@ from typing import Any, Dict +from diff_diff.results_base import Diagnostic + def describe_target_parameter(results: Any) -> Dict[str, Any]: """Return the target-parameter block for a fitted result. @@ -655,6 +657,26 @@ def describe_target_parameter(results: Any) -> Dict[str, Any]: "reference": "Butts (2021); REGISTRY.md Sec. SpilloverDiD", } + # Marked diagnostic results (spec section 3.5): no causal-effect + # inference row exists, so no ATT-style target parameter applies. + # The named ``BaconDecompositionResults`` branch above returns first + # (its block documents the TWFE coefficient being decomposed). This + # branch is defensive depth for direct callers: BR / DR reject + # marked non-Bacon diagnostics before ever calling this helper. + if isinstance(results, Diagnostic): + return { + "name": f"Diagnostic result ({name})", + "definition": ( + f"``{name}`` is a diagnostic result: it assesses a design, " + "an identifying assumption, or robustness, and carries no " + "causal-effect inference row. There is no target parameter; " + "interpret it alongside the primary estimator's results." + ), + "aggregation": "diagnostic", + "headline_attribute": "", + "reference": "docs/v4-design.md Sec. 3.5", + } + # Default: unrecognized result class. Fall through with a neutral # block — agents / downstream consumers can still dispatch on # ``aggregation="unknown"`` and fall back to generic ATT narration. diff --git a/diff_diff/bacon.py b/diff_diff/bacon.py index eaba1014e..ca5f5926e 100644 --- a/diff_diff/bacon.py +++ b/diff_diff/bacon.py @@ -18,6 +18,7 @@ import pandas as pd from diff_diff.results import _format_survey_block +from diff_diff.results_base import Diagnostic from diff_diff.utils import pre_demean_norms, snap_absorbed_regressors from diff_diff.utils import within_transform as _within_transform_util @@ -76,7 +77,7 @@ def __repr__(self) -> str: @dataclass -class BaconDecompositionResults: +class BaconDecompositionResults(Diagnostic): """ Results from Goodman-Bacon decomposition of TWFE. @@ -270,6 +271,36 @@ def _total_weight(self) -> float: """Calculate total weight (should be 1.0).""" return sum(c.weight for c in self.comparisons) + def to_dict(self) -> Dict[str, Any]: + """ + Convert the decomposition headline to a dictionary. + + Returns + ------- + Dict[str, Any] + TWFE estimate, per-comparison-type weights and averages, and + counts. This is a diagnostic container - there is no causal + inference row. The per-comparison table is available via + ``to_dataframe()``. + """ + result = { + "twfe_estimate": self.twfe_estimate, + "total_weight_treated_vs_never": self.total_weight_treated_vs_never, + "total_weight_earlier_vs_later": self.total_weight_earlier_vs_later, + "total_weight_later_vs_earlier": self.total_weight_later_vs_earlier, + "weighted_avg_treated_vs_never": self.weighted_avg_treated_vs_never, + "weighted_avg_earlier_vs_later": self.weighted_avg_earlier_vs_later, + "weighted_avg_later_vs_earlier": self.weighted_avg_later_vs_earlier, + "n_timing_groups": self.n_timing_groups, + "n_never_treated": self.n_never_treated, + "n_obs": self.n_obs, + "n_comparisons": len(self.comparisons), + "decomposition_error": self.decomposition_error, + } + if self.n_always_treated_remapped: + result["n_always_treated_remapped"] = self.n_always_treated_remapped + return result + def to_dataframe(self) -> pd.DataFrame: """ Convert comparisons to a DataFrame. diff --git a/diff_diff/business_report.py b/diff_diff/business_report.py index 19c967146..40fdbf58f 100644 --- a/diff_diff/business_report.py +++ b/diff_diff/business_report.py @@ -44,6 +44,7 @@ from diff_diff._reporting_helpers import describe_target_parameter from diff_diff.diagnostic_report import DiagnosticReport, DiagnosticReportResults +from diff_diff.results_base import Diagnostic BUSINESS_REPORT_SCHEMA_VERSION = "2.0" @@ -92,9 +93,11 @@ class BusinessReport: Parameters ---------- results : Any - A fitted diff-diff results object. Any of the 16 result types is - accepted. ``BaconDecompositionResults`` is not a valid input — Bacon - is a diagnostic, not an estimator; use ``DiagnosticReport`` for that. + A fitted diff-diff ESTIMATOR results object. Diagnostic results + (anything subclassing ``diff_diff.Diagnostic``, e.g. + ``BaconDecompositionResults``, ``HonestDiDResults``) are rejected + by type — pass diagnostics via ``diagnostics=`` or + ``DiagnosticReport(precomputed=...)`` instead. outcome_label : str, optional Stakeholder-friendly outcome name (e.g. ``"Revenue per user"``). outcome_unit : str, optional @@ -174,11 +177,22 @@ def __init__( survey_design: Optional[Any] = None, precomputed: Optional[Dict[str, Any]] = None, ): - if type(results).__name__ == "BaconDecompositionResults": + # Marked diagnostic results are rejected BY TYPE (spec section + # 3.5, ledger row M-091): BusinessReport's primary input is a + # fitted ESTIMATOR result carrying the canonical inference row. + if isinstance(results, Diagnostic): + if type(results).__name__ == "BaconDecompositionResults": + raise TypeError( + "BaconDecompositionResults is a diagnostic, not an estimator; " + "wrap the underlying estimator with BusinessReport and pass the " + "Bacon object to DiagnosticReport(precomputed={'bacon': ...})." + ) raise TypeError( - "BaconDecompositionResults is a diagnostic, not an estimator; " - "wrap the underlying estimator with BusinessReport and pass the " - "Bacon object to DiagnosticReport(precomputed={'bacon': ...})." + f"{type(results).__name__} is a diagnostic result, not an " + "estimator result; BusinessReport takes the fitted " + "estimator's results as its primary input. Pass diagnostic " + "objects via the diagnostics= parameter (as a " + "DiagnosticReport) or interpret them alongside the report." ) if diagnostics is not None and not isinstance( diff --git a/diff_diff/chaisemartin_dhaultfoeuille_results.py b/diff_diff/chaisemartin_dhaultfoeuille_results.py index b9372f7db..e7536f373 100644 --- a/diff_diff/chaisemartin_dhaultfoeuille_results.py +++ b/diff_diff/chaisemartin_dhaultfoeuille_results.py @@ -32,6 +32,7 @@ import pandas as pd from diff_diff.results import _get_significance_stars +from diff_diff.results_base import BaseResults __all__ = [ "ChaisemartinDHaultfoeuilleResults", @@ -180,7 +181,7 @@ class DCDHBootstrapResults: @dataclass -class ChaisemartinDHaultfoeuilleResults: +class ChaisemartinDHaultfoeuilleResults(BaseResults): """ Results from de Chaisemartin-D'Haultfoeuille (dCDH) Phase 1 estimation. @@ -1504,6 +1505,42 @@ def _render_honest_did_section(self, lines: List[str], width: int, thin: str) -> # to_dataframe # ------------------------------------------------------------------ + def to_dict(self) -> Dict[str, Any]: + """ + Convert headline results to a dictionary. + + Returns + ------- + Dict[str, Any] + Canonical inference row plus scalar metadata; joiner / leaver / + placebo decompositions are included only when available. + Detailed tables are available via ``to_dataframe(level=...)``. + """ + result = { + "att": self.att, + "se": self.se, + "t_stat": self.t_stat, + "p_value": self.p_value, + "conf_int_lower": self.overall_conf_int[0], + "conf_int_upper": self.overall_conf_int[1], + "n_obs": self.n_obs, + "n_treated_obs": self.n_treated_obs, + "n_switcher_cells": self.n_switcher_cells, + "n_cohorts": self.n_cohorts, + "L_max": self.L_max, + "alpha": self.alpha, + } + if self.joiners_available: + result["joiners_att"] = self.joiners_att + result["joiners_se"] = self.joiners_se + if self.leavers_available: + result["leavers_att"] = self.leavers_att + result["leavers_se"] = self.leavers_se + if self.placebo_available: + result["placebo_effect"] = self.placebo_effect + result["placebo_se"] = self.placebo_se + return result + def to_dataframe(self, level: str = "overall") -> pd.DataFrame: """ Convert results to a DataFrame at the requested level of aggregation. diff --git a/diff_diff/changes_in_changes_results.py b/diff_diff/changes_in_changes_results.py index 9baae82af..72fd7e373 100644 --- a/diff_diff/changes_in_changes_results.py +++ b/diff_diff/changes_in_changes_results.py @@ -6,6 +6,8 @@ import numpy as np import pandas as pd +from diff_diff.results_base import BaseResults + _ESTIMATOR_TITLES = { "cic": "Changes-in-Changes (Athey & Imbens 2006) Results", "qdid": "Quantile Difference-in-Differences (QDiD) Results", @@ -13,7 +15,7 @@ @dataclass -class ChangesInChangesResults: +class ChangesInChangesResults(BaseResults): """Results for :class:`~diff_diff.changes_in_changes.ChangesInChanges` and :class:`~diff_diff.changes_in_changes.QDiD`. diff --git a/diff_diff/continuous_did_results.py b/diff_diff/continuous_did_results.py index e4a93446c..e609f32cf 100644 --- a/diff_diff/continuous_did_results.py +++ b/diff_diff/continuous_did_results.py @@ -12,6 +12,7 @@ import pandas as pd from diff_diff.results import _format_survey_block, _get_significance_stars +from diff_diff.results_base import BaseResults __all__ = ["ContinuousDiDResults", "DoseResponseCurve"] @@ -78,7 +79,7 @@ def to_dataframe(self) -> pd.DataFrame: @dataclass -class ContinuousDiDResults: +class ContinuousDiDResults(BaseResults): """ Results from Continuous Difference-in-Differences estimation. @@ -369,6 +370,44 @@ def print_summary(self, alpha: Optional[float] = None) -> None: """Print summary to stdout.""" print(self.summary(alpha)) + def to_dict(self) -> Dict[str, Any]: + """ + Convert headline results to a dictionary. + + Returns + ------- + Dict[str, Any] + Canonical ATT inference row, the ACRT companion estimand, and + scalar metadata. Detailed dose-response / event-study tables + are available via ``to_dataframe(level=...)``. + """ + return { + "att": self.att, + "se": self.se, + "t_stat": self.t_stat, + "p_value": self.p_value, + "conf_int_lower": self.overall_att_conf_int[0], + "conf_int_upper": self.overall_att_conf_int[1], + "acrt": self.overall_acrt, + "acrt_se": self.overall_acrt_se, + "acrt_t_stat": self.overall_acrt_t_stat, + "acrt_p_value": self.overall_acrt_p_value, + "acrt_conf_int_lower": self.overall_acrt_conf_int[0], + "acrt_conf_int_upper": self.overall_acrt_conf_int[1], + "n_obs": self.n_obs, + "n_treated_units": self.n_treated_units, + "n_control_units": self.n_control_units, + "control_group": self.control_group, + "treatment_type": self.treatment_type, + "estimation_method": self.estimation_method, + "degree": self.degree, + "num_knots": self.num_knots, + "base_period": self.base_period, + "anticipation": self.anticipation, + "n_bootstrap": self.n_bootstrap, + "alpha": self.alpha, + } + def to_dataframe(self, level: str = "dose_response") -> pd.DataFrame: """ Convert results to DataFrame. diff --git a/diff_diff/diagnostic_report.py b/diff_diff/diagnostic_report.py index c9a4261d3..5527595ca 100644 --- a/diff_diff/diagnostic_report.py +++ b/diff_diff/diagnostic_report.py @@ -39,6 +39,7 @@ import pandas as pd from diff_diff._reporting_helpers import describe_target_parameter # noqa: E402 (top-level import) +from diff_diff.results_base import Diagnostic DIAGNOSTIC_REPORT_SCHEMA_VERSION = "2.0" @@ -253,7 +254,7 @@ @dataclass(frozen=True) -class DiagnosticReportResults: +class DiagnosticReportResults(Diagnostic): """Frozen container holding the outcome of a ``DiagnosticReport.run_all()`` call. Attributes @@ -277,6 +278,30 @@ class DiagnosticReportResults: skipped_checks: Dict[str, str] = field(default_factory=dict) warnings: Tuple[str, ...] = () + def to_dict(self) -> Dict[str, Any]: + """Return the AI-legible structured schema.""" + return self.schema + + def summary(self) -> str: + """Return a short plain-English paragraph.""" + return self.interpretation + + def to_dataframe(self) -> pd.DataFrame: + """Return one row per check with status and headline metric.""" + rows = [] + for check in _CHECK_NAMES: + section_key = "estimator_native_diagnostics" if check == "estimator_native" else check + section = self.schema.get(section_key, {}) + rows.append( + { + "check": check, + "status": section.get("status"), + "headline": _check_headline(check, section), + "reason": section.get("reason"), + } + ) + return pd.DataFrame(rows) + class DiagnosticReport: """Run the standard diff-diff diagnostic battery on a fitted result. @@ -284,9 +309,14 @@ class DiagnosticReport: Parameters ---------- results : Any - A fitted diff-diff results object (e.g. ``CallawaySantAnnaResults``, - ``DiDResults``, ``SyntheticDiDResults``). Any registered result type - in the library is accepted. + A fitted diff-diff ESTIMATOR results object (e.g. + ``CallawaySantAnnaResults``, ``DiDResults``, + ``SyntheticDiDResults``). ``BaconDecompositionResults`` is also + accepted (its dedicated Bacon read-out). Any OTHER marked + diagnostic result (subclassing ``diff_diff.Diagnostic``, e.g. + ``HonestDiDResults``, ``PlaceboTestResults``) is rejected with + ``TypeError`` — supply such objects via ``precomputed=`` where + supported, or interpret them directly. data : pandas.DataFrame, optional The underlying panel. Required for checks that need raw data (2x2 parallel-trends check on ``DiDResults``; Bacon-from-scratch when @@ -375,6 +405,22 @@ def __init__( outcome_label: Optional[str] = None, treatment_label: Optional[str] = None, ): + # Marked diagnostic results (spec section 3.5, ledger row M-091) + # are rejected BY TYPE — except Bacon, whose dedicated read-out + # is retained. Before the marker, such inputs silently produced + # a zero-check report (empty applicability); now they raise. + if ( + isinstance(results, Diagnostic) + and type(results).__name__ != "BaconDecompositionResults" + ): + raise TypeError( + f"{type(results).__name__} is a diagnostic result, not an " + "estimator result; DiagnosticReport runs its battery on a " + "fitted estimator's results. Pass the estimator result " + "here and supply diagnostic objects via precomputed= where " + "supported (e.g. precomputed={'bacon': ...}), or interpret " + "the diagnostic directly via its summary()/to_dataframe()." + ) self._results = results self._data = data self._outcome = outcome @@ -521,11 +567,11 @@ def run_all(self) -> DiagnosticReportResults: def to_dict(self) -> Dict[str, Any]: """Return the AI-legible structured schema.""" - return self.run_all().schema + return self.run_all().to_dict() def summary(self) -> str: """Return a short plain-English paragraph.""" - return self.run_all().interpretation + return self.run_all().summary() def full_report(self) -> str: """Return the multi-section markdown report.""" @@ -537,20 +583,7 @@ def export_markdown(self) -> str: def to_dataframe(self) -> pd.DataFrame: """Return one row per check with status and headline metric.""" - schema = self.to_dict() - rows = [] - for check in _CHECK_NAMES: - section_key = "estimator_native_diagnostics" if check == "estimator_native" else check - section = schema.get(section_key, {}) - rows.append( - { - "check": check, - "status": section.get("status"), - "headline": _check_headline(check, section), - "reason": section.get("reason"), - } - ) - return pd.DataFrame(rows) + return self.run_all().to_dataframe() @property def applicable_checks(self) -> Tuple[str, ...]: diff --git a/diff_diff/diagnostics.py b/diff_diff/diagnostics.py index 97be9bb36..a193d769f 100644 --- a/diff_diff/diagnostics.py +++ b/diff_diff/diagnostics.py @@ -19,11 +19,12 @@ from diff_diff.estimators import DifferenceInDifferences from diff_diff.results import _get_significance_stars +from diff_diff.results_base import Diagnostic from diff_diff.utils import safe_inference, validate_binary @dataclass -class PlaceboTestResults: +class PlaceboTestResults(Diagnostic): """ Results from a placebo test for DiD assumption validation. diff --git a/diff_diff/efficient_did_results.py b/diff_diff/efficient_did_results.py index 12627b697..958695dc8 100644 --- a/diff_diff/efficient_did_results.py +++ b/diff_diff/efficient_did_results.py @@ -12,6 +12,7 @@ import pandas as pd from diff_diff.results import _format_survey_block, _get_significance_stars +from diff_diff.results_base import BaseResults if TYPE_CHECKING: from diff_diff.efficient_did_bootstrap import EDiDBootstrapResults @@ -52,7 +53,7 @@ def __repr__(self) -> str: @dataclass -class EfficientDiDResults: +class EfficientDiDResults(BaseResults): """ Results from Efficient DiD (Chen, Sant'Anna & Xie 2025) estimation. diff --git a/diff_diff/had.py b/diff_diff/had.py index b87989c95..081953cdf 100644 --- a/diff_diff/had.py +++ b/diff_diff/had.py @@ -81,6 +81,7 @@ BiasCorrectedFit, bias_corrected_local_linear, ) +from diff_diff.results_base import BaseResults from diff_diff.survey import ( SurveyMetadata, compute_survey_metadata, @@ -197,7 +198,7 @@ def _json_safe_filter_info( @dataclass -class HeterogeneousAdoptionDiDResults: +class HeterogeneousAdoptionDiDResults(BaseResults): """Estimator output for :class:`HeterogeneousAdoptionDiD`. NaN-safe inference: the three downstream fields ``t_stat``, @@ -547,7 +548,7 @@ def to_dataframe(self) -> pd.DataFrame: @dataclass -class HeterogeneousAdoptionDiDEventStudyResults: +class HeterogeneousAdoptionDiDEventStudyResults(BaseResults): """Event-study results for :class:`HeterogeneousAdoptionDiD` (Phase 2b). Per-horizon arrays align with ``event_times`` by index; all per-horizon diff --git a/diff_diff/had_pretests.py b/diff_diff/had_pretests.py index 0ccada77c..49f87c411 100644 --- a/diff_diff/had_pretests.py +++ b/diff_diff/had_pretests.py @@ -88,6 +88,7 @@ _validate_had_panel, _validate_had_panel_event_study, ) +from diff_diff.results_base import Diagnostic from diff_diff.survey import ( SurveyDesign, make_pweight_design, @@ -144,7 +145,7 @@ @dataclass -class QUGTestResults: +class QUGTestResults(Diagnostic): """Result of :func:`qug_test` (paper Theorem 4). The QUG test rejects ``H_0: d_lower = 0`` when the order-statistic @@ -238,7 +239,7 @@ def to_dataframe(self) -> pd.DataFrame: @dataclass -class StuteTestResults: +class StuteTestResults(Diagnostic): """Result of :func:`stute_test` (paper Appendix D). The Stute test rejects the null that ``E[ΔY | D_2]`` is linear in @@ -323,7 +324,7 @@ def to_dataframe(self) -> pd.DataFrame: @dataclass -class YatchewTestResults: +class YatchewTestResults(Diagnostic): """Result of :func:`yatchew_hr_test` (paper Theorem 7 / Equation 29). Heteroskedasticity-robust specification test using Yatchew's @@ -446,7 +447,7 @@ def to_dataframe(self) -> pd.DataFrame: @dataclass -class StuteJointResult: +class StuteJointResult(Diagnostic): """Result of :func:`stute_joint_pretest` (joint Cramer-von Mises across horizons). Aggregates the per-horizon Stute (1997) CvM statistic into a joint @@ -610,7 +611,7 @@ def to_dataframe(self) -> pd.DataFrame: @dataclass -class HADPretestReport: +class HADPretestReport(Diagnostic): """Composite output of :func:`did_had_pretest_workflow`. Two dispatch shapes, distinguished by :attr:`aggregate`: diff --git a/diff_diff/honest_did.py b/diff_diff/honest_did.py index ea6ec7bab..3e2edc4a3 100644 --- a/diff_diff/honest_did.py +++ b/diff_diff/honest_did.py @@ -28,6 +28,7 @@ from diff_diff.results import ( MultiPeriodDiDResults, ) +from diff_diff.results_base import Diagnostic from diff_diff.utils import _get_critical_value # ============================================================================= @@ -156,7 +157,7 @@ def __repr__(self) -> str: @dataclass -class HonestDiDResults: +class HonestDiDResults(Diagnostic): """ Results from Honest DiD sensitivity analysis. @@ -377,7 +378,7 @@ def to_dataframe(self) -> pd.DataFrame: @dataclass -class SensitivityResults: +class SensitivityResults(Diagnostic): """ Results from sensitivity analysis over a grid of M values. diff --git a/diff_diff/imputation_results.py b/diff_diff/imputation_results.py index 762619d7c..3305caf93 100644 --- a/diff_diff/imputation_results.py +++ b/diff_diff/imputation_results.py @@ -12,6 +12,7 @@ import pandas as pd from diff_diff.results import _format_survey_block, _get_significance_stars +from diff_diff.results_base import BaseResults __all__ = [ "ImputationBootstrapResults", @@ -74,7 +75,7 @@ class ImputationBootstrapResults: @dataclass -class ImputationDiDResults: +class ImputationDiDResults(BaseResults): """ Results from Borusyak-Jaravel-Spiess (2024) imputation DiD estimation. diff --git a/diff_diff/lpdid_results.py b/diff_diff/lpdid_results.py index 5303ab88b..49c1f142e 100644 --- a/diff_diff/lpdid_results.py +++ b/diff_diff/lpdid_results.py @@ -4,9 +4,11 @@ import numpy as np import pandas as pd +from diff_diff.results_base import BaseResults + @dataclass -class LPDiDResults: +class LPDiDResults(BaseResults): """Results container for the :class:`~diff_diff.lpdid.LPDiD` estimator. Holds the per-horizon ``event_study`` table and the ``pooled`` pre/post diff --git a/diff_diff/power.py b/diff_diff/power.py index 007a53ddd..e644eefd8 100644 --- a/diff_diff/power.py +++ b/diff_diff/power.py @@ -28,6 +28,8 @@ import pandas as pd from scipy import stats +from diff_diff.results_base import Diagnostic + # Maximum sample size returned when effect is too small to detect # (e.g., zero effect or extremely small relative to noise) MAX_SAMPLE_SIZE = 2**31 - 1 @@ -980,7 +982,7 @@ def _ddd_panel_profile() -> "_EstimatorProfile": @dataclass -class PowerResults: +class PowerResults(Diagnostic): """ Results from analytical power analysis. @@ -1126,7 +1128,7 @@ def to_dataframe(self) -> pd.DataFrame: @dataclass -class SimulationPowerResults: +class SimulationPowerResults(Diagnostic): """ Results from simulation-based power analysis. @@ -2597,7 +2599,7 @@ def simulate_power( @dataclass -class SimulationMDEResults: +class SimulationMDEResults(Diagnostic): """ Results from simulation-based minimum detectable effect search. @@ -2694,7 +2696,7 @@ def to_dataframe(self) -> pd.DataFrame: @dataclass -class SimulationSampleSizeResults: +class SimulationSampleSizeResults(Diagnostic): """ Results from simulation-based sample size search. diff --git a/diff_diff/practitioner.py b/diff_diff/practitioner.py index 8bd0b78ae..11f07389d 100644 --- a/diff_diff/practitioner.py +++ b/diff_diff/practitioner.py @@ -10,6 +10,8 @@ import math from typing import Any, Dict, List, Optional, Set +from diff_diff.results_base import Diagnostic + # --------------------------------------------------------------------------- # Valid step names (Baker et al. 8-step framework) # --------------------------------------------------------------------------- @@ -108,15 +110,27 @@ def practitioner_next_steps( if unknown: raise ValueError(f"Unknown step names: {unknown}. Valid names: {sorted(STEPS)}") - # Estimation is always complete if we have a results object - completed.add("estimation") - type_name = type(results).__name__ - handler = _HANDLERS.get(type_name, _handle_generic) + # Marked diagnostic results route through diagnostic-specific + # handling (spec section 3.5, ledger row M-091) instead of the + # unknown-result estimator fallback. Bacon stays on its name-keyed + # handler with its existing framing. + diagnostic_input = isinstance(results, Diagnostic) and type_name not in _HANDLERS + + if not diagnostic_input: + # Estimation is always complete if we have an estimator results + # object; a diagnostic input carries no estimation of its own. + completed.add("estimation") + + handler = _HANDLERS.get(type_name) + if handler is None: + handler = _handle_diagnostic if diagnostic_input else _handle_generic steps, warnings = handler(results) # Prepend Steps 1-2 (pre-estimation reasoning) to every handler's output. # These are always relevant and filterable via completed_steps. + # Diagnostic inputs skip the estimator framing entirely - Steps 1-2 + # define an estimation target the diagnostic does not carry. pre_estimation = [ _step( baker_step=1, @@ -150,13 +164,18 @@ def practitioner_next_steps( if type_name == "ChangesInChangesResults": pre_estimation[1] = _cic_assumptions_step(results) - steps = pre_estimation + steps + if not diagnostic_input: + steps = pre_estimation + steps # Filter out completed steps steps = _filter_steps(steps, completed) output = { - "estimator": _estimator_display(type_name, results), + "estimator": ( + f"{type_name} (diagnostic result)" + if diagnostic_input + else _estimator_display(type_name, results) + ), "completed": sorted(completed), "next_steps": steps, "warnings": warnings, @@ -1918,6 +1937,44 @@ def _handle_cic(results: Any): return steps, warnings +def _handle_diagnostic(results: Any): + """Marked diagnostic results (``diff_diff.Diagnostic``). + + Diagnostic containers assess a design, an assumption, or robustness + and carry no causal inference row, so the estimator-style fallback + (parallel-trends test, sensitivity, estimator comparison) does not + apply. Route the user to interpreting the diagnostic alongside the + primary estimator fit. + """ + name = type(results).__name__ + steps = [ + _step( + baker_step=2, + label="Interpret the diagnostic alongside the primary fit", + why=( + f"{name} is a diagnostic result - it assesses a design, an " + "identifying assumption, or robustness, and carries no " + "causal-effect estimate. Read its summary() next to the " + "primary estimator's results before drawing conclusions." + ), + code="print(diagnostic.summary())\ndiagnostic.to_dataframe()", + step_name="assumptions", + ), + _step( + baker_step=8, + label="Get next steps from the estimator result", + why=( + "practitioner_next_steps() tailors its checklist to the " + "ESTIMATOR result; pass the fitted estimator's results " + "object for estimator-specific guidance." + ), + code="practitioner_next_steps(results) # the estimator's results", + step_name="estimation", + ), + ] + return steps, [] + + def _handle_generic(results: Any): """Fallback for unknown result types.""" steps = [ diff --git a/diff_diff/pretrends.py b/diff_diff/pretrends.py index 0cb20cdef..11a3d44f9 100644 --- a/diff_diff/pretrends.py +++ b/diff_diff/pretrends.py @@ -34,6 +34,7 @@ from scipy import optimize, stats from diff_diff.results import MultiPeriodDiDResults +from diff_diff.results_base import Diagnostic def _compute_nis_acceptance_prob( @@ -243,7 +244,7 @@ def _extract_event_study_vcov_subblock( @dataclass -class PreTrendsPowerResults: +class PreTrendsPowerResults(Diagnostic): """ Results from pre-trends power analysis. @@ -614,7 +615,7 @@ def power_at(self, M: float) -> float: @dataclass -class PreTrendsPowerCurve: +class PreTrendsPowerCurve(Diagnostic): """ Power curve across violation magnitudes. @@ -650,6 +651,25 @@ class PreTrendsPowerCurve: def __repr__(self) -> str: return f"PreTrendsPowerCurve(n_points={len(self.M_values)}, " f"mdv={self.mdv:.4f})" + def summary(self) -> str: + """Return a formatted summary of the power curve.""" + lines = [ + "Pre-Trends Power Curve", + "=" * 60, + f"{'Grid points:':<28} {len(self.M_values)}", + f"{'Violation type:':<28} {self.violation_type}", + f"{'Pretest form:':<28} {self.pretest_form}", + f"{'Significance level (alpha):':<28} {self.alpha}", + f"{'Target power:':<28} {self.target_power}", + f"{'Minimum detectable violation:':<28} {self.mdv:.4f}", + "-" * 60, + f"{'M':>12} {'Power':>10}", + ] + for m, p in zip(self.M_values, self.powers): + lines.append(f"{m:>12.4f} {p:>10.4f}") + lines.append("=" * 60) + return "\n".join(lines) + def to_dataframe(self) -> pd.DataFrame: """Convert to DataFrame with M, power, and pretest_form columns.""" return pd.DataFrame( diff --git a/diff_diff/rdd.py b/diff_diff/rdd.py index f6dcef0bf..202ae8169 100644 --- a/diff_diff/rdd.py +++ b/diff_diff/rdd.py @@ -119,6 +119,7 @@ rdbwselect, rdrobust_fit, ) +from diff_diff.results_base import BaseResults from diff_diff.utils import safe_inference, validate_covariate_names __all__ = [ @@ -134,7 +135,7 @@ def _json_safe(value: Any) -> Any: @dataclass -class RegressionDiscontinuityResults: +class RegressionDiscontinuityResults(BaseResults): """Results of a regression discontinuity fit (sharp or fuzzy; the ``estimand`` field names which one, and ``first_stage*`` fields are populated on fuzzy fits only). diff --git a/diff_diff/rdplot.py b/diff_diff/rdplot.py index 35c4a58b4..865a1731a 100644 --- a/diff_diff/rdplot.py +++ b/diff_diff/rdplot.py @@ -44,6 +44,7 @@ qrXXinv, rdrobust_kweight, ) +from .results_base import Diagnostic from .utils import validate_covariate_names __all__ = ["RDPlot", "RDPlotResult"] @@ -142,7 +143,7 @@ def _resolve_pair( @dataclass -class RDPlotResult: +class RDPlotResult(Diagnostic): """Quantities of a data-driven RD plot (field names follow R's rdplot return list; ``vars_bins`` columns carry R's ``rdplot_*`` names). diff --git a/diff_diff/results.py b/diff_diff/results.py index 5be4fc37b..bb39bf46e 100644 --- a/diff_diff/results.py +++ b/diff_diff/results.py @@ -11,6 +11,8 @@ import numpy as np import pandas as pd +from diff_diff.results_base import BaseResults + def _format_survey_block(sm, width: int) -> list: """Format survey design metadata block for summary() output. @@ -91,7 +93,7 @@ def _format_vcov_label( @dataclass -class DiDResults: +class DiDResults(BaseResults): """ Results from a Difference-in-Differences estimation. @@ -627,7 +629,7 @@ def significance_stars(self) -> str: @dataclass -class MultiPeriodDiDResults: +class MultiPeriodDiDResults(BaseResults): """ Results from a Multi-Period Difference-in-Differences estimation. @@ -1073,7 +1075,7 @@ def __post_init__(self): @dataclass -class SyntheticDiDResults: +class SyntheticDiDResults(BaseResults): """ Results from a Synthetic Difference-in-Differences estimation. diff --git a/diff_diff/results_base.py b/diff_diff/results_base.py new file mode 100644 index 000000000..6bb451f61 --- /dev/null +++ b/diff_diff/results_base.py @@ -0,0 +1,949 @@ +"""Shared result-contract foundations (4.0 program, Phase 2 PR (a)). + +This module is a deliberate LEAF: it imports numpy/pandas only, never other +``diff_diff`` modules, so every results/diagnostic/consumer module can import +it without cycles. + +Three public symbols (see ``docs/v4-design.md``): + +- :class:`BaseResults` - the shared base for estimator result containers + (spec section 5, results contract). +- :class:`Diagnostic` - the marker base for diagnostic result containers + (spec section 3.5, ledger row M-091). +- :class:`EventStudyResults` - the unified event-study representation + (spec section 5, ledger row M-092). + +``build_event_study_surface`` is package-internal: Phase 2 PR (b) wires it +into ``results.aggregate(type="event_study")``, and the Phase 3 merged +TwoWayFixedEffects event-study mode returns the container natively. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import pandas as pd + +__all__ = ["BaseResults", "Diagnostic", "EventStudyResults"] + + +class Diagnostic: + """Marker base for diagnostic RESULT containers (spec section 3.5). + + diff-diff distinguishes three object kinds: estimators, estimator + results, and diagnostics. Exactly one bit is load-bearing: an + estimator's result carries a causal-effect inference row (the canonical + quintet ``att``/``se``/``t_stat``/``p_value``/``conf_int``); a + diagnostic's result does NOT - it assesses a design, an identifying + assumption, or robustness. + + Marked classes expose ``summary()`` and ``to_dataframe()`` and are + exempt from the quintet BY TYPE, so consumers route with + ``isinstance(result, Diagnostic)`` instead of result-class-name + special-casing. The contract is enforced by the roster test in + ``tests/test_diagnostic_marker.py`` (ledger row M-091), not by the type + system - this class intentionally has no methods. + """ + + __slots__ = () + + +class BaseResults: + """Shared base for ESTIMATOR result containers (spec section 5). + + Behaviorally inert in 3.9 - a TRANSITIONAL marker base that makes the + results contract checkable and anchors the 4.0 storage flip. It does not + yet enforce a uniform runtime protocol: most subclasses are scalar + estimator results (one canonical quintet), but :class:`EventStudyResults` + is intentionally vector-valued, and a few inherited classes still expose + ``summary()`` without an ``alpha`` argument. Full protocol uniformity + (uniform ``summary(alpha=None)``, native canonical storage) is enforced in + later Phase 2 / 4.0 PRs. The target contract every ESTIMATOR-results + subclass converges to: + + - **Canonical quintet.** ``att``, ``se``, ``t_stat``, ``p_value``, + ``conf_int`` bound to ONE coherent inference row (native fields or + property aliases over legacy ``overall_*``/``avg_*`` storage; the + storage flips to native canonical fields at 4.0, ledger rows + M-050..M-058, with legacy names living on as FutureWarning + properties until 5.0). + - **Serialization.** ``summary(alpha=None)``, ``to_dict()`` (canonical + key names only - deprecated names never leak into serialized + output), and ``to_dataframe(level=...)`` where multiple views exist. + - **Pickle migration.** Classes whose stored field names change ship + ``__setstate__`` migration following the + ``SyntheticDiDResults.__setstate__`` precedent + (``diff_diff/results.py``) so 3.x pickles load under 4.0. + - **Planned estimand self-description hook.** The per-class headline + semantics (name / definition / aggregation / headline attribute) + currently live in ``_reporting_helpers.describe_target_parameter``, + keyed by result-class name. A later PR lifts that block onto this + base so results self-describe their target parameter (the MMM + exporter's verified allowlist is the first consumer); nothing here + constrains that lift. + + Diagnostic result containers (spec section 3.5) do NOT inherit this + base - they are marked with :class:`Diagnostic` and are exempt from the + quintet by type. + """ + + __slots__ = () + + +def _json_safe_label(value: Any) -> Any: + """Convert an event-time label to a JSON-serializable form. + + Calendar labels may be ``pandas.Timestamp`` / ``Period`` / ``datetime``; + ``.tolist()`` preserves those objects, which ``json.dumps`` cannot + serialize. Datetimes become ISO strings, pandas ``Period`` its ``str``, + numpy scalars their Python value; everything else passes through. + """ + if value is None: + return None + if type(value).__name__ == "Period": # pandas.Period (no isoformat) + return str(value) + if hasattr(value, "isoformat"): # datetime / date / pandas.Timestamp + return value.isoformat() + if hasattr(value, "item"): # numpy scalar + return value.item() + return value + + +#: Pinned column schema of ``EventStudyResults.to_dataframe()`` - identical +#: for every producer (spec section 5; enforced by +#: ``tests/test_event_study_surface.py``, ledger row M-092). +EVENT_STUDY_SCHEMA: Tuple[str, ...] = ( + "event_time", + "att", + "se", + "t_stat", + "p_value", + "conf_int_lower", + "conf_int_upper", + "cband_lower", + "cband_upper", + "n", + "is_reference", +) + + +@dataclass +class EventStudyResults(BaseResults): + """Unified event-study representation (spec section 5, row M-092). + + ONE representation for per-event-time effects across all estimators. + Columnar numpy arrays index-aligned to ``event_time`` (the + ``HeterogeneousAdoptionDiDEventStudyResults`` precedent). Values are + copied bit-exactly from each estimator's native surface - never + recomputed - except the mandated reference-row normalization + (``att=0.0``, inference NaN). + + Parameters + ---------- + event_time : np.ndarray + Sorted estimator-native event-time labels, NEVER renumbered. + Relative producers use their own origin (see + ``event_time_convention``); the pre-4.0 MultiPeriodDiD surface is + calendar-keyed (``time_scale="calendar"``) and may carry object + dtype (str/datetime period labels). + att, se, t_stat, p_value : np.ndarray + Canonical per-event-time inference columns. On the reference row + (if any): ``att == 0.0`` and ``se``/``t_stat``/``p_value`` are NaN. + conf_int_lower, conf_int_upper : np.ndarray + Confidence-interval bounds at ``alpha`` (NaN on the reference row). + is_reference : np.ndarray + Boolean; the EXPLICIT reference-period marking. This column - not + any count sentinel - is the sole consumer-facing signal. Usually one + True entry, but MULTIPLE are legal (CallawaySantAnna universal base + on a gapped grid carries one per cohort's positional base), and ZERO + when the estimator omits no baseline (e.g. HAD, Wooldridge). Use + ``reference_periods`` for the general case; ``reference_period`` is + the single-reference convenience scalar. + n : np.ndarray + Per-event-time count as float, NaN where the producer records none + (and on the reference row - no estimation happened there). + n_kind : str or None + Semantic of ``n`` for this producer: ``"groups"`` (a group-level + count - cohorts for CallawaySantAnna/SunAbraham, eligible switcher + groups per horizon for de Chaisemartin-D'Haultfoeuille with + ``L_max >= 1``), ``"switcher_cells"`` (dCDH legacy ``L_max is None`` + path: switching ``(g, t)`` cells, where one group may contribute + several), ``"obs"`` (observations), ``"clusters"``, or None when no + count is recorded. Never conflate these units. + reference_period : Any or None + Convenience scalar echo of the marked row's ``event_time`` label when + there is EXACTLY ONE reference row; None when there are zero or + several. Use ``is_reference`` (or the ``reference_periods`` property) + for the general case - some estimators (CallawaySantAnna universal + base on a gapped grid) carry multiple reference-only horizons. + time_scale : str + ``"relative"`` or ``"calendar"``. + event_time_convention : str or None + Origin documentation for relative scales: ``"e0_first_treated"`` + (e = t - g; first treated period at e=0) or ``"l1_first_switch"`` + (de Chaisemartin-D'Haultfoeuille: instantaneous effect at l=1, + placebos at negative keys). Horizons are documented, not + renumbered. + vcov : np.ndarray or None + Full event-study variance-covariance matrix where the RESULT + CONTAINER exposes one (today: CallawaySantAnna, SunAbraham, + MultiPeriodDiD), ordered by ``vcov_index``. StackedDiD and + TwoStageDiD compute a full VCV internally but retain only marginal + SEs on their result containers, so their surfaces carry ``vcov=None`` + until that retention is added (tracked in TODO.md). + vcov_index : np.ndarray or None + ``event_time`` labels labelling ``vcov``'s rows/columns (explicit + ordering for HonestDiD / PreTrendsPower consumption). + cband_lower, cband_upper : np.ndarray or None + Simultaneous confidence-band bounds where computed; None when the + producer has none (``to_dataframe`` then emits NaN columns - the + schema never changes). + cband_crit_value : float or None + Critical value of the simultaneous band, where computed. + alpha : float + Significance level of the stored intervals. + source : str or None + Producing results-class name (provenance). + df : float or None + Inference degrees of freedom threaded from the source, where the + producer exposes one. + """ + + event_time: np.ndarray + att: np.ndarray + se: np.ndarray + t_stat: np.ndarray + p_value: np.ndarray + conf_int_lower: np.ndarray + conf_int_upper: np.ndarray + is_reference: np.ndarray + n: np.ndarray + n_kind: Optional[str] = None + reference_period: Optional[Any] = None + time_scale: str = "relative" + event_time_convention: Optional[str] = None + vcov: Optional[np.ndarray] = field(default=None, repr=False) + vcov_index: Optional[np.ndarray] = field(default=None, repr=False) + cband_lower: Optional[np.ndarray] = None + cband_upper: Optional[np.ndarray] = None + cband_crit_value: Optional[float] = None + alpha: float = 0.05 + source: Optional[str] = None + df: Optional[float] = None + + _ARRAY_FIELDS = ( + "att", + "se", + "t_stat", + "p_value", + "conf_int_lower", + "conf_int_upper", + "n", + ) + + def __post_init__(self) -> None: + # Coerce with a COPY (np.array, not np.asarray): reference-row + # normalization below writes in place, so the container must never + # mutate a caller-owned buffer (and read-only inputs must not fail). + # event_time keeps its native dtype (object labels legal in calendar + # mode - no numeric assumptions anywhere below). + self.event_time = np.array(self.event_time) + for name in self._ARRAY_FIELDS: + setattr(self, name, np.array(getattr(self, name), dtype=float)) + self.is_reference = np.array(self.is_reference, dtype=bool) + + if self.event_time.ndim != 1: + raise ValueError( + f"EventStudyResults event_time must be one-dimensional; " + f"got ndim={self.event_time.ndim}." + ) + # A zero-row surface is legal: an event study that was REQUESTED but + # has no estimable horizons (e.g. EfficientDiD's balance_e removing + # every cohort) is a valid degenerate result, distinct from one that + # was never requested (which raises upstream in the builder). + n_rows = self.event_time.shape[0] + for name in self._ARRAY_FIELDS + ("is_reference",): + arr = getattr(self, name) + if arr.shape != (n_rows,): + raise ValueError( + f"EventStudyResults field '{name}' has shape {arr.shape}; " + f"expected ({n_rows},) to align with event_time." + ) + + # Multiple reference rows are legal: CallawaySantAnna with + # base_period="universal" on a gapped period grid materializes each + # cohort's positional base as its own reference-only horizon (all + # n_groups==0), so a valid surface can carry several. `is_reference` + # is the general marker; `reference_period` is a convenience scalar + # DERIVED from it - authoritative only when there is exactly one + # reference, None otherwise - so a caller-supplied value can never + # disagree with `is_reference` / `reference_periods`. + n_ref = int(self.is_reference.sum()) + if n_ref == 1: + ref_label = self.event_time[self.is_reference][0] + # Plain Python scalar so to_dict() stays JSON-serializable. + self.reference_period = ref_label.item() if hasattr(ref_label, "item") else ref_label + else: + self.reference_period = None + + # Simultaneous-band bounds are all-or-nothing (a lower without an upper + # is meaningless). + if (self.cband_lower is None) != (self.cband_upper is None): + raise ValueError( + "EventStudyResults requires cband_lower and cband_upper together, " "or neither." + ) + for name in ("cband_lower", "cband_upper"): + arr = getattr(self, name) + if arr is not None: + arr = np.array(arr, dtype=float) # copy: normalized in place below + if arr.shape != (n_rows,): + raise ValueError( + f"EventStudyResults field '{name}' has shape {arr.shape}; " + f"expected ({n_rows},)." + ) + setattr(self, name, arr) + + # Reference rows are normalization anchors: att is exactly 0.0 and all + # inference / count / band fields are undefined (NaN). Enforce this on + # EVERY container, not just builder output, so a direct public + # construction can never expose finite inference on a reference row. + if self.is_reference.any(): + self.att[self.is_reference] = 0.0 + for name in ("se", "t_stat", "p_value", "conf_int_lower", "conf_int_upper", "n"): + getattr(self, name)[self.is_reference] = np.nan + for name in ("cband_lower", "cband_upper"): + arr = getattr(self, name) + if arr is not None: + arr[self.is_reference] = np.nan + + if (self.vcov is None) != (self.vcov_index is None): + raise ValueError( + "EventStudyResults requires vcov and vcov_index together " + "(explicit ordering) or neither." + ) + if self.vcov is not None: + self.vcov = np.asarray(self.vcov, dtype=float) + self.vcov_index = np.asarray(self.vcov_index) + k = self.vcov_index.shape[0] + if self.vcov.shape != (k, k): + raise ValueError( + f"EventStudyResults vcov has shape {self.vcov.shape}; " + f"expected ({k}, {k}) to match vcov_index." + ) + event_labels = set(self.event_time.tolist()) + missing = [lbl for lbl in self.vcov_index.tolist() if lbl not in event_labels] + if missing: + raise ValueError( + f"EventStudyResults vcov_index entries {missing} are not " "event_time labels." + ) + + if not 0.0 < self.alpha < 1.0: + raise ValueError(f"alpha must be in (0, 1); got {self.alpha}.") + + @property + def reference_periods(self) -> List[Any]: + """All reference-row ``event_time`` labels (JSON-safe scalars). + + The general accessor for the normalization anchors: one entry for the + common single-reference case, several for CallawaySantAnna universal + base on a gapped grid, empty when the estimator omits no baseline. + """ + return [_json_safe_label(v) for v in self.event_time[self.is_reference].tolist()] + + # ------------------------------------------------------------------ + # Serialization contract + # ------------------------------------------------------------------ + def to_dataframe(self) -> pd.DataFrame: + """Return the pinned per-event-time table (``EVENT_STUDY_SCHEMA``). + + The column schema is identical for every producer; band columns are + NaN-filled when the producer computes no simultaneous band. + """ + n_rows = self.event_time.shape[0] + nan_col = np.full(n_rows, np.nan) + return pd.DataFrame( + { + "event_time": self.event_time, + "att": self.att, + "se": self.se, + "t_stat": self.t_stat, + "p_value": self.p_value, + "conf_int_lower": self.conf_int_lower, + "conf_int_upper": self.conf_int_upper, + "cband_lower": self.cband_lower if self.cband_lower is not None else nan_col, + "cband_upper": self.cband_upper if self.cband_upper is not None else nan_col, + "n": self.n, + "is_reference": self.is_reference, + }, + columns=list(EVENT_STUDY_SCHEMA), + ) + + def to_dict(self) -> Dict[str, Any]: + """Return a JSON-friendly dict (columns as lists, plus metadata). + + Calendar event-time labels (``pandas.Timestamp`` / ``Period`` / + ``datetime``) are converted to ISO strings so the result is + JSON-serializable; relative-time integer labels pass through. + """ + out: Dict[str, Any] = { + "event_time": [_json_safe_label(v) for v in self.event_time.tolist()], + "att": self.att.tolist(), + "se": self.se.tolist(), + "t_stat": self.t_stat.tolist(), + "p_value": self.p_value.tolist(), + "conf_int_lower": self.conf_int_lower.tolist(), + "conf_int_upper": self.conf_int_upper.tolist(), + "cband_lower": self.cband_lower.tolist() if self.cband_lower is not None else None, + "cband_upper": self.cband_upper.tolist() if self.cband_upper is not None else None, + "n": self.n.tolist(), + "is_reference": self.is_reference.tolist(), + "n_kind": self.n_kind, + "reference_period": _json_safe_label(self.reference_period), + "reference_periods": self.reference_periods, + "time_scale": self.time_scale, + "event_time_convention": self.event_time_convention, + # Full event-study vcov + its ordered index (labels JSON-safed), + # so covariance-aware consumers (HonestDiD / PreTrendsPower) can + # round-trip through the serialized surface. + "vcov": self.vcov.tolist() if self.vcov is not None else None, + "vcov_index": ( + [_json_safe_label(v) for v in self.vcov_index.tolist()] + if self.vcov_index is not None + else None + ), + "cband_crit_value": self.cband_crit_value, + "alpha": self.alpha, + "source": self.source, + "df": self.df, + } + return out + + def summary(self, alpha: Optional[float] = None) -> str: + """Return a formatted event-study table. + + Parameters + ---------- + alpha : float, optional + Accepted for signature uniformity (spec section 5). The stored + confidence columns were computed at fit time; passing a value + different from the stored ``alpha`` raises rather than silently + recomputing or mislabeling - re-aggregate at the desired level + instead. + """ + if alpha is not None and alpha != self.alpha: + raise ValueError( + f"This event-study surface stores intervals computed at " + f"alpha={self.alpha}; re-aggregate to obtain alpha={alpha} " + "intervals (summary() never recomputes stored inference)." + ) + ci_pct = int(round((1 - self.alpha) * 100)) + lines = [ + "Event-Study Effects", + "=" * 78, + ] + meta_bits = [] + if self.source: + meta_bits.append(f"source: {self.source}") + meta_bits.append(f"time scale: {self.time_scale}") + if self.event_time_convention: + meta_bits.append(f"convention: {self.event_time_convention}") + if self.n_kind: + meta_bits.append(f"n counts: {self.n_kind}") + lines.append(" ".join(meta_bits)) + lines.append("-" * 78) + lines.append( + f"{'Event time':>12} {'ATT':>10} {'SE':>10} {'t':>8} " + f"{'P>|t|':>8} {f'[{ci_pct}% CI]':>21}" + ) + for i in range(self.event_time.shape[0]): + label = f"{self.event_time[i]}" + if self.is_reference[i]: + lines.append(f"{label:>12} {0.0:>10.4f} {'(reference)':>{50}}") + continue + lines.append( + f"{label:>12} {self.att[i]:>10.4f} {self.se[i]:>10.4f} " + f"{self.t_stat[i]:>8.3f} {self.p_value[i]:>8.3f} " + f"[{self.conf_int_lower[i]:>9.4f}, {self.conf_int_upper[i]:>9.4f}]" + ) + lines.append("=" * 78) + return "\n".join(lines) + + +# ---------------------------------------------------------------------- +# Builders: native representation -> EventStudyResults (package-internal) +# ---------------------------------------------------------------------- + +#: Per-producer remediation for an absent event-study surface. +_ABSENT_SURFACE_HINTS: Dict[str, str] = { + "CallawaySantAnnaResults": "refit with aggregate='event_study' (or 'all')", + "ImputationDiDResults": "refit with aggregate='event_study' (or 'all')", + "TwoStageDiDResults": "refit with aggregate='event_study' (or 'all')", + "StackedDiDResults": "refit with aggregate='event_study'", + "StaggeredTripleDiffResults": "refit with aggregate='event_study' (or 'all')", + "EfficientDiDResults": "refit with aggregate='event_study' (or 'all')", + "ContinuousDiDResults": "refit with aggregate='eventstudy' (or 'all')", + "WooldridgeDiDResults": "call results.aggregate('event') first", + "SpilloverDiDResults": "refit with event_study=True", + "ChaisemartinDHaultfoeuilleResults": "refit with L_max >= 1", + "LPDiDResults": "refit with only_pooled=False", +} + + +def _absent(results: Any) -> ValueError: + name = type(results).__name__ + hint = _ABSENT_SURFACE_HINTS.get(name, "request event-study aggregation at fit") + return ValueError(f"{name} carries no event-study surface - {hint}.") + + +def _empty_surface(results: Any) -> EventStudyResults: + """Zero-row surface for a requested-but-empty event study.""" + empty_f = np.empty(0, dtype=float) + return EventStudyResults( + event_time=np.empty(0), + att=empty_f.copy(), + se=empty_f.copy(), + t_stat=empty_f.copy(), + p_value=empty_f.copy(), + conf_int_lower=empty_f.copy(), + conf_int_upper=empty_f.copy(), + is_reference=np.empty(0, dtype=bool), + n=empty_f.copy(), + time_scale="relative", + alpha=getattr(results, "alpha", 0.05), + source=type(results).__name__, + ) + + +def _materialize_cband( + pairs: Dict[int, Tuple[float, float]], n_rows: int +) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]: + """Turn sparse per-row band tuples into aligned arrays (None if empty).""" + if not pairs: + return None, None + lo = np.full(n_rows, np.nan) + hi = np.full(n_rows, np.nan) + for i, (lo_i, hi_i) in pairs.items(): + lo[i] = lo_i + hi[i] = hi_i + return lo, hi + + +def _from_relative_dict(results: Any) -> EventStudyResults: + """Adapter for ``event_study_effects: Dict[int, Dict]`` producers. + + Covers CallawaySantAnna, SunAbraham, ImputationDiD, TwoStageDiD, + StackedDiD, SpilloverDiD, ContinuousDiD, EfficientDiD, Wooldridge + (``att`` inner key normalized here), and StaggeredTripleDiff. + + Reference resolution is PRODUCER-AWARE, because a zero count is not a + universal reference marker: + + 1. If the producer names a ``reference_period`` that is PRESENT in the + effects dict (SpilloverDiD materializes a zero row), that row is THE + reference. Other zero-count rows - SpilloverDiD emits genuinely + non-estimable rectangular horizons with ``n_obs == 0`` and NaN + inference - are preserved as NON-reference NaN rows. + 2. If the producer names a ``reference_period`` that is ABSENT (an + OMITTED baseline, e.g. SunAbraham's ``e = -1 - anticipation``), the + anchor row is synthesized ONLY when the estimator reports it was + genuinely observed (``reference_observed``) - never invented on a + gapped grid where that period was never in the data. + 3. Otherwise the in-band zero-count sentinel marks the reference: a + zero-count row is a reference iff its effect is exactly 0 (the + normalization anchor). A NaN-effect zero-count row is a genuinely + non-estimable horizon (TwoStageDiD emits effect=NaN, n_obs=0 for a + horizon whose observations are all filtered) - preserved, never + normalized. CallawaySantAnna universal base can mark several such + rows; ContinuousDiD / Wooldridge carry no count key and yield none. + """ + effects = getattr(results, "event_study_effects", None) + if effects is None: + raise _absent(results) # not requested + if len(effects) == 0: + # Requested but no estimable horizons (e.g. EfficientDiD balance_e + # eliminated every cohort). A zero-row surface, not an error. + return _empty_surface(results) + + explicit_ref = getattr(results, "reference_period", None) + + # Synthesize an OMITTED-but-OBSERVED baseline before building arrays. + effects = dict(effects) + synthesized_ref = None + if ( + explicit_ref is not None + and explicit_ref not in effects + and getattr(results, "reference_observed", False) + ): + synthesized_ref = explicit_ref + effects[explicit_ref] = { + "effect": 0.0, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + } + + keys = sorted(effects.keys()) + n_rows = len(keys) + att = np.empty(n_rows) + se = np.empty(n_rows) + t_stat = np.empty(n_rows) + p_value = np.empty(n_rows) + ci_lo = np.empty(n_rows) + ci_hi = np.empty(n_rows) + n = np.full(n_rows, np.nan) + zero_count = np.zeros(n_rows, dtype=bool) + cband_pairs: Dict[int, Tuple[float, float]] = {} + + n_kind: Optional[str] = None + for i, k in enumerate(keys): + row = effects[k] + att[i] = row["att"] if "att" in row else row["effect"] + se[i] = row.get("se", np.nan) + t_stat[i] = row.get("t_stat", np.nan) + p_value[i] = row.get("p_value", np.nan) + ci = row.get("conf_int") + if ci is not None: + ci_lo[i], ci_hi[i] = float(ci[0]), float(ci[1]) + else: + ci_lo[i] = ci_hi[i] = np.nan + if "n_groups" in row: + n_kind = "groups" + n[i] = float(row["n_groups"]) + zero_count[i] = row["n_groups"] == 0 + elif "n_obs" in row: + n_kind = "obs" + n[i] = float(row["n_obs"]) + zero_count[i] = row["n_obs"] == 0 + + cband = row.get("cband_conf_int") + if cband is not None: + cband_pairs[i] = (float(cband[0]), float(cband[1])) + + event_time = np.asarray(keys) + ref_present = explicit_ref is not None and ( + synthesized_ref is not None or explicit_ref in set(event_time.tolist()) + ) + if ref_present: + # Named reference that is materialized (Spillover) or synthesized + # (SunAbraham observed): mark exactly that row. + is_ref = event_time == explicit_ref + elif explicit_ref is not None: + # Named reference that is absent and NOT synthesized (SunAbraham on a + # gapped grid where the anchor was never observed): no reference row. + is_ref = np.zeros(n_rows, dtype=bool) + else: + # In-band zero-count sentinel. A zero-count row is the REFERENCE only + # when its effect is exactly 0 (the normalization anchor). A NaN-effect + # zero-count row is a genuinely non-estimable horizon - preserved. A + # zero-count row with a finite NONZERO effect is malformed - fail + # loudly rather than silently rewrite it. + finite = np.isfinite(att) + bad = zero_count & finite & (att != 0.0) + if bad.any(): + raise ValueError( + f"{type(results).__name__} event-study surface has a zero-count " + f"row with a finite nonzero effect at {event_time[bad].tolist()}; " + "a reference/normalization row must have effect 0.0." + ) + is_ref = zero_count & finite + + cband_lo, cband_hi = _materialize_cband(cband_pairs, n_rows) + + vcov = getattr(results, "event_study_vcov", None) + vcov_index = getattr(results, "event_study_vcov_index", None) + if vcov is None or vcov_index is None: + vcov = vcov_index = None + + return EventStudyResults( + event_time=event_time, + att=att, + se=se, + t_stat=t_stat, + p_value=p_value, + conf_int_lower=ci_lo, + conf_int_upper=ci_hi, + is_reference=is_ref, + n=n, + n_kind=n_kind, + time_scale="relative", + event_time_convention="e0_first_treated", + vcov=np.asarray(vcov, dtype=float) if vcov is not None else None, + vcov_index=np.asarray(vcov_index) if vcov_index is not None else None, + cband_lower=cband_lo, + cband_upper=cband_hi, + cband_crit_value=getattr(results, "cband_crit_value", None), + alpha=getattr(results, "alpha", 0.05), + source=type(results).__name__, + df=getattr(results, "df_inference", 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 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. + """ + 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()) + all_periods = sorted(periods + [ref]) if ref is not None else 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) + n = np.full(n_rows, np.nan) + is_ref = np.zeros(n_rows, dtype=bool) + + for i, p in enumerate(all_periods): + 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]) + + 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: + idx = [interaction_indices[p] for p in covered] + vcov_sub = np.asarray(full_vcov, dtype=float)[np.ix_(idx, idx)] + vcov_index_arr = np.asarray(covered) + + 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, + is_reference=is_ref, + n=n, + n_kind=None, + reference_period=ref, + time_scale="calendar", + event_time_convention=None, + vcov=vcov_sub, + vcov_index=vcov_index_arr, + alpha=getattr(results, "alpha", 0.05), + source=type(results).__name__, + ) + + +def _from_lpdid(results: Any) -> EventStudyResults: + """Adapter for LPDiD's per-horizon ``event_study`` DataFrame. + + Native columns ``horizon``/``coefficient``/``conf_low``/``conf_high`` + map onto the canonical names; the materialized ``horizon == -1`` base + row (coefficient 0.0, NaN inference) translates to the marked + reference row. The per-horizon ``n_clusters`` column stays on the + native frame only (``n_kind="obs"``). + """ + frame = getattr(results, "event_study", None) + if frame is None or len(frame) == 0: + raise _absent(results) + + frame = frame.sort_values("horizon") + horizons = frame["horizon"].to_numpy() + # copy=True: pandas may hand back a read-only view, and the reference + # normalization writes into these arrays. + att = np.array(frame["coefficient"], dtype=float) + se = np.array(frame["se"], dtype=float) + t_stat = np.array(frame["t_stat"], dtype=float) + p_value = np.array(frame["p_value"], dtype=float) + ci_lo = np.array(frame["conf_low"], dtype=float) + ci_hi = np.array(frame["conf_high"], dtype=float) + n = np.array(frame["n_obs"], dtype=float) + is_ref = horizons == -1 + + return EventStudyResults( + event_time=horizons, + att=att, + se=se, + t_stat=t_stat, + p_value=p_value, + conf_int_lower=ci_lo, + conf_int_upper=ci_hi, + is_reference=is_ref, + n=n, + n_kind="obs", + time_scale="relative", + event_time_convention="e0_first_treated", + alpha=getattr(results, "alpha", 0.05), + source=type(results).__name__, + ) + + +def _from_dcdh(results: Any) -> EventStudyResults: + """Adapter for de Chaisemartin-D'Haultfoeuille's split dicts. + + Merges ``placebo_event_study`` (negative keys), a synthesized + reference at 0, and ``event_study_effects`` (post horizons l >= 1, + l=1 = instantaneous effect) - the same merge the event-study plotter + performs. Horizons keep the paper's convention + (``event_time_convention="l1_first_switch"``); they are documented, + never renumbered. + """ + effects = getattr(results, "event_study_effects", None) + if not effects: + raise _absent(results) + placebos = getattr(results, "placebo_event_study", None) or {} + + keys: List[Any] = sorted(placebos.keys()) + [0] + sorted(effects.keys()) + n_rows = len(keys) + 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) + n = np.full(n_rows, np.nan) + is_ref = np.zeros(n_rows, dtype=bool) + cband_pairs: Dict[int, Tuple[float, float]] = {} + + for i, k in enumerate(keys): + if k == 0: + is_ref[i] = True + continue + row = placebos[k] if k < 0 else effects[k] + att[i] = row["effect"] + se[i] = row.get("se", np.nan) + t_stat[i] = row.get("t_stat", np.nan) + p_value[i] = row.get("p_value", np.nan) + ci = row.get("conf_int") + if ci is not None: + ci_lo[i], ci_hi[i] = float(ci[0]), float(ci[1]) + if "n_obs" in row: + n[i] = float(row["n_obs"]) + cband = row.get("cband_conf_int") + if cband is not None: + cband_pairs[i] = (float(cband[0]), float(cband[1])) + + cband_lo, cband_hi = _materialize_cband(cband_pairs, n_rows) + + # Simultaneous-band critical value lives on sup_t_bands (populated only on + # a bootstrap fit); carry it so the band's provenance is not lost. + sup_t = getattr(results, "sup_t_bands", None) + cband_crit = sup_t.get("crit_value") if isinstance(sup_t, dict) else None + + # dCDH stores its count under the legacy "n_obs" key, but the UNIT is + # L_max-dependent (never observations): with L_max >= 1 it is N_l, the + # number of eligible switcher GROUPS per horizon; with L_max is None it is + # N_S, the number of switching (g, t) CELLS (one group can contribute + # several). Label each accurately so downstream sample-size logic cannot + # conflate them. + l_max = getattr(results, "L_max", None) + dcdh_n_kind = "groups" if (l_max is not None and l_max >= 1) else "switcher_cells" + + return EventStudyResults( + event_time=np.asarray(keys), + att=att, + se=se, + t_stat=t_stat, + p_value=p_value, + conf_int_lower=ci_lo, + conf_int_upper=ci_hi, + is_reference=is_ref, + n=n, + n_kind=dcdh_n_kind, + time_scale="relative", + event_time_convention="l1_first_switch", + cband_lower=cband_lo, + cband_upper=cband_hi, + cband_crit_value=cband_crit, + alpha=getattr(results, "alpha", 0.05), + source=type(results).__name__, + ) + + +def _from_had(results: Any) -> EventStudyResults: + """Adapter for HAD's columnar event-study container (near-passthrough). + + The anchor horizon e = -1 is excluded from the native surface by + design (trivially zero, WAS not identified there), so this producer + legitimately has no reference row. + """ + event_times = getattr(results, "event_times", None) + if event_times is None or len(event_times) == 0: + raise _absent(results) + + n_rows = len(event_times) + n_obs = getattr(results, "n_obs_per_horizon", None) + cband_lo = getattr(results, "cband_low", None) + cband_hi = getattr(results, "cband_high", None) + + return EventStudyResults( + event_time=np.array(event_times), + att=np.array(results.att, dtype=float), + se=np.array(results.se, dtype=float), + t_stat=np.array(results.t_stat, dtype=float), + p_value=np.array(results.p_value, dtype=float), + conf_int_lower=np.array(results.conf_int_low, dtype=float), + conf_int_upper=np.array(results.conf_int_high, dtype=float), + is_reference=np.zeros(n_rows, dtype=bool), + n=(np.asarray(n_obs, dtype=float) if n_obs is not None else np.full(n_rows, np.nan)), + n_kind="obs" if n_obs is not None else None, + time_scale="relative", + event_time_convention="e0_first_treated", + cband_lower=np.asarray(cband_lo, dtype=float) if cband_lo is not None else None, + cband_upper=np.asarray(cband_hi, dtype=float) if cband_hi is not None else None, + cband_crit_value=getattr(results, "cband_crit_value", None), + alpha=getattr(results, "alpha", 0.05), + source=type(results).__name__, + ) + + +def build_event_study_surface(results: Any) -> EventStudyResults: + """Build the unified event-study surface from a fitted results object. + + Package-internal (spec section 5 / row M-092): public exposure rides + ``results.aggregate(type="event_study")`` (Phase 2 PR (b)); the Phase 3 + merged TwoWayFixedEffects event-study mode returns the container + natively. Values are copied bit-exactly from the native surface; + an absent surface raises ``ValueError`` naming the call that produces + it (never a silent empty container). + """ + # 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. + 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"): + return _from_had(results) + if hasattr(results, "period_effects") and hasattr(results, "interaction_indices"): + return _from_mpd(results) + if hasattr(results, "pooled") and hasattr(results, "event_study"): + return _from_lpdid(results) + if hasattr(results, "event_study_effects"): + # SunAbraham (omitted-but-observed baseline) and SpilloverDiD + # (materialized reference) both flow through here via reference_period. + return _from_relative_dict(results) + raise TypeError( + f"{type(results).__name__} does not expose an event-study surface. " + "Supported producers: CallawaySantAnna, SunAbraham, ImputationDiD, " + "TwoStageDiD, StackedDiD, SpilloverDiD, ContinuousDiD, EfficientDiD, " + "WooldridgeDiD, StaggeredTripleDifference, MultiPeriodDiD, LPDiD, " + "ChaisemartinDHaultfoeuille, and HeterogeneousAdoptionDiD " + "event-study results." + ) diff --git a/diff_diff/stacked_did_results.py b/diff_diff/stacked_did_results.py index 52a60ddf5..246021ba3 100644 --- a/diff_diff/stacked_did_results.py +++ b/diff_diff/stacked_did_results.py @@ -12,6 +12,7 @@ import pandas as pd from diff_diff.results import _format_survey_block, _get_significance_stars +from diff_diff.results_base import BaseResults __all__ = [ "StackedDiDResults", @@ -19,7 +20,7 @@ @dataclass -class StackedDiDResults: +class StackedDiDResults(BaseResults): """ Results from Stacked DiD estimation (Wing, Freedman & Hollingsworth 2024). @@ -341,6 +342,43 @@ def print_summary(self, alpha: Optional[float] = None) -> None: """Print summary to stdout.""" print(self.summary(alpha)) + def to_dict(self) -> Dict[str, Any]: + """ + Convert headline results to a dictionary. + + Returns + ------- + Dict[str, Any] + Canonical inference row plus scalar metadata. Detailed + event-study / group tables are available via + ``to_dataframe(level=...)``. + """ + result = { + "att": self.att, + "se": self.se, + "t_stat": self.t_stat, + "p_value": self.p_value, + "conf_int_lower": self.overall_conf_int[0], + "conf_int_upper": self.overall_conf_int[1], + "n_obs": self.n_obs, + "n_stacked_obs": self.n_stacked_obs, + "n_sub_experiments": self.n_sub_experiments, + "n_treated_units": self.n_treated_units, + "n_control_units": self.n_control_units, + "kappa_pre": self.kappa_pre, + "kappa_post": self.kappa_post, + "weighting": self.weighting, + "clean_control": self.clean_control, + "anticipation": self.anticipation, + "alpha": self.alpha, + "vcov_type": self.vcov_type, + } + if self.cluster_name is not None: + result["cluster_name"] = self.cluster_name + if self.n_clusters is not None: + result["n_clusters"] = self.n_clusters + return result + def to_dataframe(self, level: str = "event_study") -> pd.DataFrame: """ Convert results to DataFrame. diff --git a/diff_diff/staggered_results.py b/diff_diff/staggered_results.py index fde82418b..dadf04498 100644 --- a/diff_diff/staggered_results.py +++ b/diff_diff/staggered_results.py @@ -12,6 +12,7 @@ import pandas as pd from diff_diff.results import _format_survey_block, _get_significance_stars +from diff_diff.results_base import BaseResults if TYPE_CHECKING: from diff_diff.staggered_bootstrap import CSBootstrapResults @@ -66,7 +67,7 @@ def significance_stars(self) -> str: @dataclass -class CallawaySantAnnaResults: +class CallawaySantAnnaResults(BaseResults): """ Results from Callaway-Sant'Anna (2021) staggered DiD estimation. @@ -422,6 +423,40 @@ def print_summary(self, alpha: Optional[float] = None) -> None: """Print summary to stdout.""" print(self.summary(alpha)) + def to_dict(self) -> Dict[str, Any]: + """ + Convert headline results to a dictionary. + + Returns + ------- + Dict[str, Any] + Canonical inference row plus scalar metadata. Detailed + group-time / event-study tables are available via + ``to_dataframe(level=...)``. + """ + result = { + "att": self.att, + "se": self.se, + "t_stat": self.t_stat, + "p_value": self.p_value, + "conf_int_lower": self.overall_conf_int[0], + "conf_int_upper": self.overall_conf_int[1], + "n_obs": self.n_obs, + "n_treated_units": self.n_treated_units, + "n_control_units": self.n_control_units, + "control_group": self.control_group, + "base_period": self.base_period, + "anticipation": self.anticipation, + "panel": self.panel, + "alpha": self.alpha, + "vcov_type": self.vcov_type, + } + if self.cluster_name is not None: + result["cluster_name"] = self.cluster_name + if self.n_clusters is not None: + result["n_clusters"] = self.n_clusters + return result + def to_dataframe(self, level: str = "group_time") -> pd.DataFrame: """ Convert results to DataFrame. diff --git a/diff_diff/staggered_triple_diff_results.py b/diff_diff/staggered_triple_diff_results.py index 2c7c434fc..f2942cbd0 100644 --- a/diff_diff/staggered_triple_diff_results.py +++ b/diff_diff/staggered_triple_diff_results.py @@ -12,13 +12,14 @@ import pandas as pd from diff_diff.results import _format_survey_block, _get_significance_stars +from diff_diff.results_base import BaseResults if TYPE_CHECKING: from diff_diff.staggered_bootstrap import CSBootstrapResults @dataclass -class StaggeredTripleDiffResults: +class StaggeredTripleDiffResults(BaseResults): """ Results from Staggered Triple Difference (DDD) estimation. diff --git a/diff_diff/sun_abraham.py b/diff_diff/sun_abraham.py index a366f3966..79bca4874 100644 --- a/diff_diff/sun_abraham.py +++ b/diff_diff/sun_abraham.py @@ -22,6 +22,7 @@ from diff_diff.survey import ResolvedSurveyDesign, SurveyDesign from diff_diff.linalg import LinearRegression from diff_diff.results import _format_survey_block, _get_significance_stars +from diff_diff.results_base import BaseResults from diff_diff.utils import ( pre_demean_norms, safe_inference, @@ -33,7 +34,7 @@ @dataclass -class SunAbrahamResults: +class SunAbrahamResults(BaseResults): """ Results from Sun-Abraham (2021) interaction-weighted estimation. @@ -127,6 +128,17 @@ class SunAbrahamResults: # summary label). Both None on non-conley fits. conley_lag_cutoff: Optional[int] = None cluster_name: Optional[str] = None + # The normalization reference relative time (e = -1 - anticipation, the + # omitted category of the saturated regression) and whether it was + # GENUINELY OBSERVED in the panel. The reference is excluded from + # event_study_effects either because it is the omitted baseline (observed) + # OR because no cohort has an observation there (unobserved, on a gapped + # grid) - the two are indistinguishable from the estimated keys alone, so + # the unified event-study surface synthesizes the anchor row only when + # reference_observed is True. Defaults are conservative (no synthesis) for + # externally / legacy-constructed results. + reference_period: Optional[int] = None + reference_observed: bool = False # --- Inference-field aliases (balance/external-adapter compatibility) --- @property @@ -280,6 +292,38 @@ def print_summary(self, alpha: Optional[float] = None) -> None: """Print summary to stdout.""" print(self.summary(alpha)) + def to_dict(self) -> Dict[str, Any]: + """ + Convert headline results to a dictionary. + + Returns + ------- + Dict[str, Any] + Canonical inference row plus scalar metadata. Detailed + event-study / cohort tables are available via + ``to_dataframe(level=...)``. + """ + result = { + "att": self.att, + "se": self.se, + "t_stat": self.t_stat, + "p_value": self.p_value, + "conf_int_lower": self.overall_conf_int[0], + "conf_int_upper": self.overall_conf_int[1], + "n_obs": self.n_obs, + "n_treated_units": self.n_treated_units, + "n_control_units": self.n_control_units, + "control_group": self.control_group, + "anticipation": self.anticipation, + "alpha": self.alpha, + "vcov_type": self.vcov_type, + } + if self.cluster_name is not None: + result["cluster_name"] = self.cluster_name + if self.conley_lag_cutoff is not None: + result["conley_lag_cutoff"] = self.conley_lag_cutoff + return result + def to_dataframe(self, level: str = "event_study") -> pd.DataFrame: """ Convert results to DataFrame. @@ -835,6 +879,10 @@ def fit( # Reference period: last pre-treatment period (typically -1) self._reference_period = -1 - self.anticipation + # Whether that anchor was GENUINELY OBSERVED (vs a gap on an + # unbalanced grid). The unified event-study surface synthesizes the + # reference row only when it was observed. + self._reference_observed = self._reference_period in all_rel_times # Get relative periods to estimate (excluding reference) rel_periods_to_estimate = [ @@ -1328,6 +1376,8 @@ def _refit_sa_cohort(w_r): event_study_vcov_index=es_vcov_index, conley_lag_cutoff=(self.conley_lag_cutoff if self.vcov_type == "conley" else None), cluster_name=(self.cluster if self.vcov_type == "conley" else None), + reference_period=self._reference_period, + reference_observed=self._reference_observed, ) self.is_fitted_ = True diff --git a/diff_diff/synthetic_control_results.py b/diff_diff/synthetic_control_results.py index 74e2a3057..d5e7ce342 100644 --- a/diff_diff/synthetic_control_results.py +++ b/diff_diff/synthetic_control_results.py @@ -22,6 +22,7 @@ import pandas as pd from diff_diff.results import _format_survey_block, _get_significance_stars +from diff_diff.results_base import BaseResults __all__ = ["SyntheticControlResults"] @@ -126,7 +127,7 @@ def _warn_conformal_ci_status(res: Dict[str, Any], method_name: str) -> None: @dataclass -class SyntheticControlResults: +class SyntheticControlResults(BaseResults): """ Results from a classic Synthetic Control Method (SCM) estimation. diff --git a/diff_diff/triple_diff.py b/diff_diff/triple_diff.py index 3990c2d46..1d08f5ee6 100644 --- a/diff_diff/triple_diff.py +++ b/diff_diff/triple_diff.py @@ -36,6 +36,7 @@ from diff_diff.linalg import _rank_guarded_inv, solve_logit, solve_ols from diff_diff.results import _format_survey_block, _get_significance_stars +from diff_diff.results_base import BaseResults from diff_diff.utils import safe_inference if TYPE_CHECKING: @@ -50,7 +51,7 @@ @dataclass -class TripleDifferenceResults: +class TripleDifferenceResults(BaseResults): """ Results from Triple Difference (DDD) estimation. diff --git a/diff_diff/trop_results.py b/diff_diff/trop_results.py index d95f09cea..8cc9a4781 100644 --- a/diff_diff/trop_results.py +++ b/diff_diff/trop_results.py @@ -17,6 +17,7 @@ from typing_extensions import TypedDict from diff_diff.results import _format_survey_block, _get_significance_stars +from diff_diff.results_base import BaseResults __all__ = [ "_LAMBDA_INF", @@ -64,7 +65,7 @@ class _PrecomputedStructures(TypedDict): @dataclass -class TROPResults: +class TROPResults(BaseResults): """ Results from a Triply Robust Panel (TROP) estimation. diff --git a/diff_diff/two_stage_results.py b/diff_diff/two_stage_results.py index af2407874..422f7a3ef 100644 --- a/diff_diff/two_stage_results.py +++ b/diff_diff/two_stage_results.py @@ -12,6 +12,7 @@ import pandas as pd from diff_diff.results import _format_survey_block, _get_significance_stars +from diff_diff.results_base import BaseResults __all__ = [ "TwoStageBootstrapResults", @@ -76,7 +77,7 @@ class TwoStageBootstrapResults: @dataclass -class TwoStageDiDResults: +class TwoStageDiDResults(BaseResults): """ Results from Gardner (2022) two-stage DiD estimation. diff --git a/diff_diff/wooldridge_results.py b/diff_diff/wooldridge_results.py index a290239b3..55ed701e0 100644 --- a/diff_diff/wooldridge_results.py +++ b/diff_diff/wooldridge_results.py @@ -9,11 +9,12 @@ import numpy as np import pandas as pd +from diff_diff.results_base import BaseResults from diff_diff.utils import safe_inference @dataclass -class WooldridgeDiDResults: +class WooldridgeDiDResults(BaseResults): """Results from WooldridgeDiD.fit(). Core output is ``group_time_effects``: a dict keyed by (cohort_g, time_t) @@ -707,6 +708,41 @@ def _fmt_row(label: str, att: float, se: float, t: float, p: float, ci: Tuple) - lines.append("=" * 70) return "\n".join(lines) + def to_dict(self) -> Dict[str, Any]: + """ + Convert headline results to a dictionary. + + Returns + ------- + Dict[str, Any] + Canonical inference row plus scalar metadata. Detailed + group-time / aggregated tables are available via + ``to_dataframe(aggregation=...)``. + """ + result = { + "att": self.att, + "se": self.se, + "t_stat": self.t_stat, + "p_value": self.p_value, + "conf_int_lower": self.overall_conf_int[0], + "conf_int_upper": self.overall_conf_int[1], + "method": self.method, + "control_group": self.control_group, + "n_obs": self.n_obs, + "n_treated_units": self.n_treated_units, + "n_control_units": self.n_control_units, + "anticipation": self.anticipation, + "alpha": self.alpha, + "vcov_type": self.vcov_type, + } + if self.cluster_name is not None: + result["cluster_name"] = self.cluster_name + if self.n_clusters is not None: + result["n_clusters"] = self.n_clusters + if self.conley_lag_cutoff is not None: + result["conley_lag_cutoff"] = self.conley_lag_cutoff + return result + def to_dataframe(self, aggregation: str = "event") -> pd.DataFrame: """Export aggregated effects to a DataFrame. diff --git a/docs/api/_autosummary/diff_diff.BaconDecompositionResults.rst b/docs/api/_autosummary/diff_diff.BaconDecompositionResults.rst index a4afa3cde..d0e57f490 100644 --- a/docs/api/_autosummary/diff_diff.BaconDecompositionResults.rst +++ b/docs/api/_autosummary/diff_diff.BaconDecompositionResults.rst @@ -16,6 +16,7 @@ ~BaconDecompositionResults.print_summary ~BaconDecompositionResults.summary ~BaconDecompositionResults.to_dataframe + ~BaconDecompositionResults.to_dict ~BaconDecompositionResults.weight_by_type diff --git a/docs/api/_autosummary/diff_diff.BaseResults.rst b/docs/api/_autosummary/diff_diff.BaseResults.rst new file mode 100644 index 000000000..2b1b1a492 --- /dev/null +++ b/docs/api/_autosummary/diff_diff.BaseResults.rst @@ -0,0 +1,18 @@ +diff\_diff.BaseResults +====================== + +.. currentmodule:: diff_diff + +.. autoclass:: BaseResults + :no-members: + + + .. rubric:: Methods + + .. autosummary:: + + ~BaseResults.__init__ + + + + diff --git a/docs/api/_autosummary/diff_diff.CallawaySantAnnaResults.rst b/docs/api/_autosummary/diff_diff.CallawaySantAnnaResults.rst index 12aaa22d7..4c0b19664 100644 --- a/docs/api/_autosummary/diff_diff.CallawaySantAnnaResults.rst +++ b/docs/api/_autosummary/diff_diff.CallawaySantAnnaResults.rst @@ -16,6 +16,7 @@ ~CallawaySantAnnaResults.print_summary ~CallawaySantAnnaResults.summary ~CallawaySantAnnaResults.to_dataframe + ~CallawaySantAnnaResults.to_dict diff --git a/docs/api/_autosummary/diff_diff.ChaisemartinDHaultfoeuilleResults.rst b/docs/api/_autosummary/diff_diff.ChaisemartinDHaultfoeuilleResults.rst index 36475b352..85cebec58 100644 --- a/docs/api/_autosummary/diff_diff.ChaisemartinDHaultfoeuilleResults.rst +++ b/docs/api/_autosummary/diff_diff.ChaisemartinDHaultfoeuilleResults.rst @@ -15,6 +15,7 @@ ~ChaisemartinDHaultfoeuilleResults.print_summary ~ChaisemartinDHaultfoeuilleResults.summary ~ChaisemartinDHaultfoeuilleResults.to_dataframe + ~ChaisemartinDHaultfoeuilleResults.to_dict diff --git a/docs/api/_autosummary/diff_diff.ContinuousDiDResults.rst b/docs/api/_autosummary/diff_diff.ContinuousDiDResults.rst index 2d4609550..c2a10b072 100644 --- a/docs/api/_autosummary/diff_diff.ContinuousDiDResults.rst +++ b/docs/api/_autosummary/diff_diff.ContinuousDiDResults.rst @@ -15,6 +15,7 @@ ~ContinuousDiDResults.print_summary ~ContinuousDiDResults.summary ~ContinuousDiDResults.to_dataframe + ~ContinuousDiDResults.to_dict diff --git a/docs/api/_autosummary/diff_diff.Diagnostic.rst b/docs/api/_autosummary/diff_diff.Diagnostic.rst new file mode 100644 index 000000000..429a407ce --- /dev/null +++ b/docs/api/_autosummary/diff_diff.Diagnostic.rst @@ -0,0 +1,18 @@ +diff\_diff.Diagnostic +===================== + +.. currentmodule:: diff_diff + +.. autoclass:: Diagnostic + :no-members: + + + .. rubric:: Methods + + .. autosummary:: + + ~Diagnostic.__init__ + + + + diff --git a/docs/api/_autosummary/diff_diff.DiagnosticReportResults.rst b/docs/api/_autosummary/diff_diff.DiagnosticReportResults.rst index b16d7a514..274f3a6fd 100644 --- a/docs/api/_autosummary/diff_diff.DiagnosticReportResults.rst +++ b/docs/api/_autosummary/diff_diff.DiagnosticReportResults.rst @@ -12,6 +12,9 @@ .. autosummary:: ~DiagnosticReportResults.__init__ + ~DiagnosticReportResults.summary + ~DiagnosticReportResults.to_dataframe + ~DiagnosticReportResults.to_dict diff --git a/docs/api/_autosummary/diff_diff.EventStudyResults.rst b/docs/api/_autosummary/diff_diff.EventStudyResults.rst new file mode 100644 index 000000000..e820809be --- /dev/null +++ b/docs/api/_autosummary/diff_diff.EventStudyResults.rst @@ -0,0 +1,48 @@ +diff\_diff.EventStudyResults +============================ + +.. currentmodule:: diff_diff + +.. autoclass:: EventStudyResults + :no-members: + + + .. rubric:: Methods + + .. autosummary:: + + ~EventStudyResults.__init__ + ~EventStudyResults.summary + ~EventStudyResults.to_dataframe + ~EventStudyResults.to_dict + + + + + .. rubric:: Attributes + + .. autosummary:: + + ~EventStudyResults.alpha + ~EventStudyResults.cband_crit_value + ~EventStudyResults.cband_lower + ~EventStudyResults.cband_upper + ~EventStudyResults.df + ~EventStudyResults.event_time_convention + ~EventStudyResults.n_kind + ~EventStudyResults.reference_period + ~EventStudyResults.reference_periods + ~EventStudyResults.source + ~EventStudyResults.time_scale + ~EventStudyResults.vcov + ~EventStudyResults.vcov_index + ~EventStudyResults.event_time + ~EventStudyResults.att + ~EventStudyResults.se + ~EventStudyResults.t_stat + ~EventStudyResults.p_value + ~EventStudyResults.conf_int_lower + ~EventStudyResults.conf_int_upper + ~EventStudyResults.is_reference + ~EventStudyResults.n + diff --git a/docs/api/_autosummary/diff_diff.PreTrendsPowerCurve.rst b/docs/api/_autosummary/diff_diff.PreTrendsPowerCurve.rst index ed681361e..ec220f089 100644 --- a/docs/api/_autosummary/diff_diff.PreTrendsPowerCurve.rst +++ b/docs/api/_autosummary/diff_diff.PreTrendsPowerCurve.rst @@ -13,6 +13,7 @@ ~PreTrendsPowerCurve.__init__ ~PreTrendsPowerCurve.plot + ~PreTrendsPowerCurve.summary ~PreTrendsPowerCurve.to_dataframe diff --git a/docs/api/_autosummary/diff_diff.StackedDiDResults.rst b/docs/api/_autosummary/diff_diff.StackedDiDResults.rst index 9134d07a9..f1362af12 100644 --- a/docs/api/_autosummary/diff_diff.StackedDiDResults.rst +++ b/docs/api/_autosummary/diff_diff.StackedDiDResults.rst @@ -15,6 +15,7 @@ ~StackedDiDResults.print_summary ~StackedDiDResults.summary ~StackedDiDResults.to_dataframe + ~StackedDiDResults.to_dict diff --git a/docs/api/_autosummary/diff_diff.SunAbrahamResults.rst b/docs/api/_autosummary/diff_diff.SunAbrahamResults.rst index 3af9b3263..527b7c8ce 100644 --- a/docs/api/_autosummary/diff_diff.SunAbrahamResults.rst +++ b/docs/api/_autosummary/diff_diff.SunAbrahamResults.rst @@ -15,6 +15,7 @@ ~SunAbrahamResults.print_summary ~SunAbrahamResults.summary ~SunAbrahamResults.to_dataframe + ~SunAbrahamResults.to_dict diff --git a/docs/api/_autosummary/diff_diff.wooldridge_results.WooldridgeDiDResults.rst b/docs/api/_autosummary/diff_diff.wooldridge_results.WooldridgeDiDResults.rst index 8bc82818d..144daa655 100644 --- a/docs/api/_autosummary/diff_diff.wooldridge_results.WooldridgeDiDResults.rst +++ b/docs/api/_autosummary/diff_diff.wooldridge_results.WooldridgeDiDResults.rst @@ -16,6 +16,7 @@ ~WooldridgeDiDResults.plot_event_study ~WooldridgeDiDResults.summary ~WooldridgeDiDResults.to_dataframe + ~WooldridgeDiDResults.to_dict diff --git a/docs/api/index.rst b/docs/api/index.rst index 6393c4849..fd6d32ae1 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -82,6 +82,9 @@ Result containers returned by estimators: diff_diff.TWFEWeightsResult diff_diff.RegressionDiscontinuityResults diff_diff.RDPlotResult + diff_diff.BaseResults + diff_diff.Diagnostic + diff_diff.EventStudyResults Visualization ------------- diff --git a/docs/api/results.rst b/docs/api/results.rst index 4545ef664..36c4c0466 100644 --- a/docs/api/results.rst +++ b/docs/api/results.rst @@ -90,3 +90,28 @@ Results from SyntheticDiD estimation. ~SyntheticDiDResults.att ~SyntheticDiDResults.unit_weights ~SyntheticDiDResults.time_weights + +Results-contract foundations +---------------------------- + +Shared bases introduced by the 4.0 API-unification program. +``BaseResults`` is the estimator-results base; +``Diagnostic`` marks diagnostic result containers (which carry no +inference row); ``EventStudyResults`` is the unified per-event-time +representation. + +.. autoclass:: diff_diff.BaseResults + :no-index: + :members: + :show-inheritance: + +.. autoclass:: diff_diff.Diagnostic + :no-index: + :members: + :show-inheritance: + +.. autoclass:: diff_diff.EventStudyResults + :no-index: + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index 586c2da9f..638e9a76d 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -1047,6 +1047,16 @@ sources: type: methodology note: "SyntheticDiDResults hosts validation diagnostics (LOO, weight concentration, in-time placebo, zeta sensitivity)" + diff_diff/results_base.py: + drift_risk: low + docs: + - path: docs/api/results.rst + type: api_reference + - path: docs/v4-design.md + section: "5. Results contract" + type: design_spec + note: "BaseResults / Diagnostic marker / EventStudyResults contract (4.0 program Phase 2; ledger rows M-091/M-092/M-093)" + diff_diff/bootstrap_utils.py: drift_risk: medium docs: diff --git a/docs/v4-deprecations.yaml b/docs/v4-deprecations.yaml index 3b3b25aa2..52ec2623f 100644 --- a/docs/v4-deprecations.yaml +++ b/docs/v4-deprecations.yaml @@ -904,7 +904,7 @@ rows: removed_in: null status: planned phase: 5 - code_refs: [README.md, diff_diff/guides/llms.txt, docs/api/index.rst, docs/index.rst, diff_diff/bacon.py] + code_refs: [README.md, diff_diff/guides/llms.txt, docs/api/index.rst, docs/index.rst, diff_diff/bacon.py, tests/test_docs_ia.py] notes: "4.0 diagnostic-family docs/roster reorganization (spec section 3.5): Bacon re-homed OUT of the API estimator roster (misfiled today); RDPlot (already documented as a diagnostic) consolidated under the unified family grouping; README/llms.txt/API/alias groupings (incl. docs/api/index.rst AND the docs homepage docs/index.rst, both listing Bacon as an estimator today) and estimator-count claims split into estimators + diagnostics; the flip's terminal test asserts the estimator/diagnostic grouping on the named surfaces. The Phase 2 marker introduction is gated separately by [M-091]." - id: M-091 kind: behavior @@ -914,10 +914,35 @@ rows: introduced_in: "3.9" deprecated_in: null removed_in: null - status: planned + status: done + phase: 2 + test_ref: tests/test_diagnostic_marker.py + code_refs: [diff_diff/results_base.py, diff_diff/bacon.py, diff_diff/rdplot.py, diff_diff/honest_did.py, diff_diff/pretrends.py, diff_diff/power.py, diff_diff/diagnostics.py, diff_diff/had_pretests.py, diff_diff/diagnostic_report.py, diff_diff/business_report.py, diff_diff/practitioner.py, diff_diff/_reporting_helpers.py, diff_diff/__init__.py] + notes: "Phase 2 Diagnostic marker base on the RESULT roster (spec section 3.5): BaconDecompositionResults, RDPlotResult, HonestDiDResults + SensitivityResults, PreTrendsPowerResults + PreTrendsPowerCurve, PowerResults + Simulation*Results, PlaceboTestResults, and the HAD pretest containers (QUGTestResults, StuteTestResults, YatchewTestResults, StuteJointResult, HADPretestReport), and DiagnosticReportResults (serialization pair moved from the builder onto the container). Flip to done ALSO requires consumer-propagation tests: BusinessReport rejects marked diagnostics as primary estimator input BY TYPE (not by name); practitioner_next_steps routes marked diagnostics through diagnostic handling (not the unknown-result estimator fallback); and DiagnosticReport routes by the marker - at least one non-Bacon marked result handled as a diagnostic, never via estimator fallback, with Bacon's existing read-out retained and tested. Flip to done requires the dedicated roster test (isinstance + summary/to_dataframe on every member; estimator results NOT marked). Raw-dict functions and TWFEWeightsResult are docs-family only (narrowed, spec section 3.5). introduced_in gates the 3.9 cut: the marker cannot be silently skipped. Marker lives in diff_diff/results_base.py; done in this introducing diff." + - id: M-092 + kind: behavior + group: results-contract + old: "diff_diff:EventStudyResults" + new: null + introduced_in: "3.9" + deprecated_in: null + removed_in: null + status: done phase: 2 - code_refs: [diff_diff/bacon.py, diff_diff/rdplot.py, diff_diff/honest_did.py, diff_diff/pretrends.py, diff_diff/power.py, diff_diff/diagnostics.py, diff_diff/had_pretests.py, diff_diff/diagnostic_report.py, diff_diff/business_report.py, diff_diff/practitioner.py, diff_diff/_reporting_helpers.py, diff_diff/__init__.py] - notes: "Phase 2 Diagnostic marker base on the RESULT roster (spec section 3.5): BaconDecompositionResults, RDPlotResult, HonestDiDResults + SensitivityResults, PreTrendsPowerResults + PreTrendsPowerCurve, PowerResults + Simulation*Results, PlaceboTestResults, and the HAD pretest containers (QUGTestResults, StuteTestResults, YatchewTestResults, StuteJointResult, HADPretestReport), and DiagnosticReportResults (serialization pair moves/delegates from the builder onto the container). Flip to done ALSO requires consumer-propagation tests: BusinessReport rejects marked diagnostics as primary estimator input BY TYPE (not by name); practitioner_next_steps routes marked diagnostics through diagnostic handling (not the unknown-result estimator fallback); and DiagnosticReport routes by the marker - at least one non-Bacon marked result handled as a diagnostic, never via estimator fallback, with Bacon's existing read-out retained and tested. Flip to done requires the dedicated roster test (isinstance + summary/to_dataframe on every member; estimator results NOT marked). Raw-dict functions and TWFEWeightsResult are docs-family only (narrowed, spec section 3.5). introduced_in gates the 3.9 cut: the marker cannot be silently skipped." + test_ref: tests/test_event_study_surface.py + code_refs: [diff_diff/results_base.py, diff_diff/__init__.py] + notes: "Phase 2 unified event-study representation (spec section 5): EventStudyResults container + builders for the 14 producers (CallawaySantAnna, SunAbraham, ImputationDiD, TwoStageDiD, StackedDiD, SpilloverDiD, ContinuousDiD, EfficientDiD, WooldridgeDiD, StaggeredTripleDifference, MultiPeriodDiD, LPDiD, ChaisemartinDHaultfoeuille, HeterogeneousAdoptionDiD). Canonical quintet columns, explicit is_reference marking (successor to the retiring sentinels [M-093]), vcov+vcov_index ordering, cband columns, event_time_convention metadata. Public exposure rides aggregate(type='event_study') in Phase 2 PR (b); merged TWFE returns it in Phase 3 [M-010]. introduced_in gates the 3.9 cut, mirroring [M-091]. Born done in this introducing diff (builder is package-internal; the class is exported)." + - id: M-093 + kind: behavior + group: results-contract + old: "diff_diff:CallawaySantAnnaResults.event_study_effects" + new: "diff_diff:EventStudyResults" + deprecated_in: "4.0" + removed_in: null + status: planned + phase: 5 + code_refs: [diff_diff/staggered_results.py, diff_diff/sun_abraham.py, diff_diff/imputation_results.py, diff_diff/two_stage_results.py, diff_diff/stacked_did_results.py, diff_diff/efficient_did_results.py, diff_diff/continuous_did_results.py, diff_diff/wooldridge_results.py, diff_diff/chaisemartin_dhaultfoeuille_results.py, diff_diff/lpdid_results.py, diff_diff/staggered_triple_diff_results.py, diff_diff/results.py, diff_diff/had.py, diff_diff/visualization/_event_study.py] + notes: "4.0 sentinel retirement + schema enforcement (spec section 5): the n_groups==0 / n_obs==0 reference-row sentinels retire; every estimator's to_dataframe(level='event_study') emits the [M-092] column schema; the plotter / HonestDiD / PreTrendsPower consume the unified surface. Thirteen ES-carrying source modules enumerated file-by-file (results.py covers MultiPeriodDiD + SpilloverDiD) plus the plotter. behavior-at-done requires test_ref." # ---- Behavior policies (schema-tracked, spec-governed; no reality probe) - - id: M-080 diff --git a/docs/v4-design.md b/docs/v4-design.md index c87289904..a025a0a3a 100644 --- a/docs/v4-design.md +++ b/docs/v4-design.md @@ -371,14 +371,17 @@ the seven classes currently missing `to_dict()` gain it in Phase 2 (additive). **Unified event-study representation.** ONE representation for event-study -effects across all estimators (today there are five). The unified surface is -specified in the Phase 2 results-base PR plan (per the section 9 boundary -rule - only that PR cares about the container's internals); this document -pins the requirements it must satisfy: per-event-time canonical quintet rows, -explicit reference-period marking (no sentinel-value conventions - the -n_groups==0 / n_obs==0 sentinels are retired), the event-study vcov exposed -uniformly where computed, and `to_dataframe(level="event_study")` emitting -identical column schemas from every estimator. +effects across all estimators (today the fourteen event-study producers use +four incompatible native container shapes). The `EventStudyResults` +container plus its per-producer builders ship additively in Phase 2 [M-092] +(the builder is package-internal; public exposure rides +`aggregate(type="event_study")` in the same phase, and the merged TWFE returns +it natively in Phase 3). This document pins the requirements it satisfies: +per-event-time canonical quintet rows, explicit reference-period marking (via +an `is_reference` column - no sentinel-value conventions; the +n_groups==0 / n_obs==0 sentinels are retired at 4.0 [M-093]), the event-study +vcov exposed uniformly where computed, and `to_dataframe(level="event_study")` +emitting identical column schemas from every estimator. **Pickle migration.** Renamed-field classes ship `__setstate__` migration following the existing `SyntheticDiDResults.__setstate__` precedent @@ -610,9 +613,12 @@ forever - a removed symbol resurrecting is a test failure. the same object as its target, so the deprecation warning rides the parent class row (schema-enforced). Top-level `diff_diff:Name` class/function rows and alias rows also assert `__all__` membership consistent with their - status (stale `import *` entries fail). The initial 75 row ids are a - committed snapshot in the enforcement test: ids are never deleted or - reused, and the test fails if any snapshot id disappears. + status (stale `import *` entries fail). The shipped row ids are a + committed snapshot in the enforcement test (77 as of Phase 2a: Phase 1 + + the diagnostic-family amendment + the M-092/M-093 results-contract rows; + the snapshot extends by a new id range in the same diff that appends + rows): ids are never deleted or reused, and the test fails if any + snapshot id disappears. **Cross-row migration rule.** Removing a symbol requires migrating, in the same diff, every other row whose locators or `code_refs` reference it (e.g. diff --git a/tests/helpers/results_foundation.py b/tests/helpers/results_foundation.py new file mode 100644 index 000000000..e8fec2a65 --- /dev/null +++ b/tests/helpers/results_foundation.py @@ -0,0 +1,235 @@ +"""Shared factories for the Phase 2 results-contract test files. + +Two kinds of factories: + +- ``make_constructed_diagnostics()`` - directly constructed instances of + every class-backed diagnostic roster member EXCEPT + ``BaconDecompositionResults`` (which is produced by a real fit in the + test files - its container is populated by the decomposition and its + ``summary()`` renders the comparison table). Synthetic values are + internally consistent so ``summary()`` / ``to_dataframe()`` render. +- Small-fit panel builders shared by the event-study surface tests and + the serialization tests (see ``staggered_panel()`` and friends). + +Kept in ``tests/helpers`` (on ``sys.path`` via ``tests/conftest.py``) so +multiple test modules can import the same factories without duplicating +panel constructions. +""" + +from __future__ import annotations + +from typing import Any, Dict + +import numpy as np +import pandas as pd + +import diff_diff +from diff_diff.diagnostic_report import DiagnosticReportResults + + +def make_constructed_diagnostics() -> Dict[str, Any]: + """Return {class_name: instance} for the direct-construction roster.""" + rng_bins = pd.DataFrame( + { + "rdplot_id": [1, 2], + "rdplot_mean_bin": [-0.5, 0.5], + "rdplot_mean_x": [-0.5, 0.5], + "rdplot_mean_y": [1.0, 2.0], + "rdplot_N": [10, 10], + "rdplot_ci_l": [0.8, 1.8], + "rdplot_ci_r": [1.2, 2.2], + } + ) + poly = pd.DataFrame({"rdplot_x": [-1.0, 0.0, 1.0], "rdplot_y": [0.9, 1.4, 2.1]}) + coef = pd.DataFrame({"side": ["left", "right"], "coef_0": [1.0, 2.0]}) + + qug = diff_diff.QUGTestResults( + t_stat=1.2, + p_value=0.23, + reject=False, + alpha=0.05, + critical_value=1.96, + n_obs=50, + n_excluded_zero=0, + d_order_1=0.1, + d_order_2=0.2, + ) + stute = diff_diff.StuteTestResults( + cvm_stat=0.4, + p_value=0.31, + reject=False, + alpha=0.05, + n_bootstrap=99, + n_obs=50, + seed=42, + ) + yatchew = diff_diff.YatchewTestResults( + t_stat_hr=0.8, + p_value=0.42, + reject=False, + alpha=0.05, + critical_value=1.64, + sigma2_lin=1.1, + sigma2_diff=1.0, + sigma2_W=0.9, + n_obs=50, + ) + + instances: Dict[str, Any] = { + "RDPlotResult": diff_diff.RDPlotResult( + coef=coef, + vars_bins=rng_bins, + vars_poly=poly, + J=(2.0, 2.0), + J_IMSE=(2.0, 2.0), + J_MV=(3.0, 3.0), + scale=(1.0, 1.0), + rscale=(1.0, 1.0), + bin_avg=(5.0, 5.0), + bin_med=(5.0, 5.0), + p=4, + cutoff=0.0, + h=(1.0, 1.0), + N=(10, 10), + N_h=(10, 10), + binselect="esmv", + kernel_type="Uniform", + ci_level=95.0, + ci_requested=False, + ), + "HonestDiDResults": diff_diff.HonestDiDResults( + lb=0.1, + ub=0.9, + ci_lb=0.05, + ci_ub=0.95, + M=1.0, + method="relative_magnitude", + original_estimate=0.5, + original_se=0.2, + ), + "SensitivityResults": diff_diff.SensitivityResults( + M_values=np.array([0.5, 1.0]), + bounds=[(0.2, 0.8), (0.1, 0.9)], + robust_cis=[(0.1, 0.9), (0.0, 1.0)], + breakdown_M=1.0, + method="relative_magnitude", + original_estimate=0.5, + original_se=0.2, + ), + "PreTrendsPowerResults": diff_diff.PreTrendsPowerResults( + power=0.62, + mdv=0.8, + violation_magnitude=0.5, + violation_type="linear", + alpha=0.05, + target_power=0.8, + n_pre_periods=3, + test_statistic=2.1, + critical_value=7.8, + noncentrality=3.3, + pre_period_effects=np.array([0.01, -0.02, 0.03]), + pre_period_ses=np.array([0.05, 0.05, 0.05]), + vcov=np.eye(3) * 0.0025, + ), + "PreTrendsPowerCurve": diff_diff.PreTrendsPowerCurve( + M_values=np.array([0.0, 0.5, 1.0]), + powers=np.array([0.05, 0.44, 0.91]), + mdv=0.8, + alpha=0.05, + target_power=0.8, + violation_type="linear", + ), + "PowerResults": diff_diff.PowerResults( + power=0.8, + mde=0.25, + required_n=120, + effect_size=0.3, + alpha=0.05, + alternative="two-sided", + n_treated=60, + n_control=60, + n_pre=4, + n_post=4, + sigma=1.0, + ), + "SimulationPowerResults": diff_diff.SimulationPowerResults( + power=0.78, + power_se=0.04, + power_ci=(0.7, 0.86), + rejection_rate=0.78, + mean_estimate=0.29, + std_estimate=0.11, + mean_se=0.1, + coverage=0.94, + n_simulations=100, + effect_sizes=[0.3], + powers=[0.78], + true_effect=0.3, + alpha=0.05, + estimator_name="DifferenceInDifferences", + ), + "SimulationMDEResults": diff_diff.SimulationMDEResults( + mde=0.31, + power_at_mde=0.8, + target_power=0.8, + alpha=0.05, + n_units=100, + n_simulations_per_step=50, + n_steps=6, + search_path=[{"effect": 0.3, "power": 0.75}], + estimator_name="DifferenceInDifferences", + ), + "SimulationSampleSizeResults": diff_diff.SimulationSampleSizeResults( + required_n=140, + power_at_n=0.81, + target_power=0.8, + alpha=0.05, + effect_size=0.3, + n_simulations_per_step=50, + n_steps=6, + search_path=[{"n": 140, "power": 0.81}], + estimator_name="DifferenceInDifferences", + ), + "PlaceboTestResults": diff_diff.PlaceboTestResults( + test_type="placebo_treatment", + placebo_effect=0.02, + se=0.05, + t_stat=0.4, + p_value=0.69, + conf_int=(-0.08, 0.12), + n_obs=200, + is_significant=False, + ), + "QUGTestResults": qug, + "StuteTestResults": stute, + "YatchewTestResults": yatchew, + "StuteJointResult": diff_diff.StuteJointResult( + cvm_stat_joint=0.6, + p_value=0.27, + reject=False, + alpha=0.05, + horizon_labels=["e0", "e1"], + per_horizon_stats={"e0": 0.3, "e1": 0.3}, + n_bootstrap=99, + n_obs=50, + n_horizons=2, + seed=42, + null_form="linear", + exact_linear_short_circuited=False, + ), + "HADPretestReport": diff_diff.HADPretestReport( + qug=qug, + stute=stute, + yatchew=yatchew, + all_pass=True, + verdict="pass", + alpha=0.05, + n_obs=50, + ), + "DiagnosticReportResults": DiagnosticReportResults( + schema={"schema_version": "2.0", "estimator": "DiDResults"}, + interpretation="All applicable checks passed.", + applicable_checks=("parallel_trends",), + ), + } + return instances diff --git a/tests/test_diagnostic_marker.py b/tests/test_diagnostic_marker.py new file mode 100644 index 000000000..936bf29e9 --- /dev/null +++ b/tests/test_diagnostic_marker.py @@ -0,0 +1,226 @@ +"""Diagnostic-marker roster + consumer-propagation contract (ledger M-091). + +This file is the ``test_ref`` for ledger row M-091 (spec section 3.5). It +enforces three things: + +1. Every class-backed diagnostic RESULT container subclasses + :class:`diff_diff.Diagnostic`, exposes ``summary()`` / ``to_dataframe()``, + and is exempt from the estimator quintet by type. +2. Estimator results are NOT marked; ``TWFEWeightsResult`` (a docs-family + member, narrowed out of the type contract) is neither. +3. The three consumers route by the marker: BusinessReport and + DiagnosticReport reject marked non-Bacon primaries by type, + ``practitioner_next_steps`` routes marked diagnostics through + diagnostic-specific handling, and Bacon retains its existing read-out. +""" + +import numpy as np +import pandas as pd +import pytest +from results_foundation import make_constructed_diagnostics + +import diff_diff +from diff_diff import BaconDecomposition, Diagnostic +from diff_diff._reporting_helpers import describe_target_parameter +from diff_diff.business_report import BusinessReport +from diff_diff.diagnostic_report import DiagnosticReport, DiagnosticReportResults +from diff_diff.practitioner import practitioner_next_steps + +# The class-backed diagnostic roster (spec section 3.5). Bacon is handled +# separately: it is produced by a real fit below (its container is only +# meaningful when populated by the decomposition). +DIAGNOSTIC_ROSTER = [ + "BaconDecompositionResults", + "RDPlotResult", + "HonestDiDResults", + "SensitivityResults", + "PreTrendsPowerResults", + "PreTrendsPowerCurve", + "PowerResults", + "SimulationPowerResults", + "SimulationMDEResults", + "SimulationSampleSizeResults", + "PlaceboTestResults", + "QUGTestResults", + "StuteTestResults", + "YatchewTestResults", + "StuteJointResult", + "HADPretestReport", + "DiagnosticReportResults", +] + +# Representative ESTIMATOR results: marked with BaseResults, never Diagnostic. +ESTIMATOR_SAMPLE = [ + "DiDResults", + "CallawaySantAnnaResults", + "LPDiDResults", + "HeterogeneousAdoptionDiDResults", + "HeterogeneousAdoptionDiDEventStudyResults", + "WooldridgeDiDResults", +] + + +def _staggered_panel(seed=7): + """Small staggered panel that BaconDecomposition accepts.""" + rng = np.random.RandomState(seed) + n_units, n_periods = 12, 6 + units = np.repeat(np.arange(n_units), n_periods) + times = np.tile(np.arange(n_periods), n_units) + first_treat = np.repeat(np.array([0, 0, 0, 0, 3, 3, 3, 3, 4, 4, 4, 4]), n_periods) + post = (times >= first_treat) & (first_treat > 0) + unit_fe = np.repeat(rng.randn(n_units) * 2, n_periods) + time_fe = np.tile(np.linspace(0, 1, n_periods), n_units) + outcome = unit_fe + time_fe + 1.5 * post + rng.randn(n_units * n_periods) * 0.5 + return pd.DataFrame( + { + "unit": units, + "time": times, + "outcome": outcome, + "first_treat": first_treat.astype(int), + "treated": post.astype(int), + } + ) + + +@pytest.fixture(scope="module") +def bacon_result(): + data = _staggered_panel() + return BaconDecomposition().fit( + data, outcome="outcome", unit="unit", time="time", first_treat="first_treat" + ) + + +@pytest.fixture(scope="module") +def constructed(): + return make_constructed_diagnostics() + + +def _roster_instance(name, constructed, bacon_result): + if name == "BaconDecompositionResults": + return bacon_result + return constructed[name] + + +# --------------------------------------------------------------------------- +# 1. Roster contract +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("name", DIAGNOSTIC_ROSTER) +def test_roster_member_is_marked_diagnostic(name, constructed, bacon_result): + obj = _roster_instance(name, constructed, bacon_result) + assert isinstance(obj, Diagnostic), f"{name} must subclass Diagnostic" + + +@pytest.mark.parametrize("name", DIAGNOSTIC_ROSTER) +def test_roster_member_exposes_serialization_pair(name, constructed, bacon_result): + obj = _roster_instance(name, constructed, bacon_result) + summary = obj.summary() + assert isinstance(summary, str) and summary, f"{name}.summary() must be a str" + frame = obj.to_dataframe() + assert isinstance(frame, pd.DataFrame), f"{name}.to_dataframe() must be a DataFrame" + + +@pytest.mark.parametrize("name", DIAGNOSTIC_ROSTER) +def test_roster_member_has_no_quintet(name, constructed, bacon_result): + # Diagnostics are exempt from the estimator quintet BY TYPE. A diagnostic + # exposing all five canonical inference names would blur the boundary the + # marker exists to enforce (PlaceboTestResults carries se/t_stat/p_value + # but its estimate is `placebo_effect`, not `att` - so no full quintet). + obj = _roster_instance(name, constructed, bacon_result) + quintet = ("att", "se", "t_stat", "p_value", "conf_int") + assert not all( + hasattr(obj, q) for q in quintet + ), f"{name} exposes the full estimator quintet; diagnostics must not." + + +def test_roster_matches_marked_classes_in_namespace(): + # Guard against a roster member silently losing its marker: every name in + # DIAGNOSTIC_ROSTER resolves to a Diagnostic subclass. + import diff_diff.diagnostic_report as drm + + for name in DIAGNOSTIC_ROSTER: + cls = getattr(diff_diff, name, None) or getattr(drm, name) + assert issubclass(cls, Diagnostic), name + + +# --------------------------------------------------------------------------- +# 2. Negative roster: estimators are not diagnostics +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("name", ESTIMATOR_SAMPLE) +def test_estimator_results_not_marked(name): + cls = getattr(diff_diff, name) + assert issubclass(cls, diff_diff.BaseResults), f"{name} must be BaseResults" + assert not issubclass(cls, Diagnostic), f"{name} must NOT be Diagnostic" + + +def test_twfe_weights_result_is_neither(): + # Narrowed out of the type contract (docs-family only): a plain __slots__ + # class, not a marked container. + cls = diff_diff.TWFEWeightsResult + assert not issubclass(cls, Diagnostic) + assert not issubclass(cls, diff_diff.BaseResults) + + +def test_frozen_diagnostic_constructs(constructed): + # DiagnosticReportResults is the only frozen roster member; pin that + # frozen-dataclass + non-dataclass marker inheritance is safe. + drr = constructed["DiagnosticReportResults"] + assert isinstance(drr, DiagnosticReportResults) + assert isinstance(drr, Diagnostic) + with pytest.raises((AttributeError, TypeError)): + drr.schema = {} # frozen + + +# --------------------------------------------------------------------------- +# 3. Consumer propagation +# --------------------------------------------------------------------------- +def test_business_report_rejects_bacon_by_type(bacon_result): + with pytest.raises(TypeError, match="diagnostic"): + BusinessReport(bacon_result) + + +def test_business_report_rejects_non_bacon_marked(constructed): + with pytest.raises(TypeError, match="diagnostic"): + BusinessReport(constructed["HonestDiDResults"]) + + +def test_diagnostic_report_rejects_non_bacon_marked(constructed): + with pytest.raises(TypeError, match="diagnostic"): + DiagnosticReport(constructed["PreTrendsPowerCurve"]) + + +def test_diagnostic_report_retains_bacon_readout(bacon_result): + # Bacon is NOT rejected: it keeps its dedicated read-out. + report = DiagnosticReport(bacon_result) + schema = report.to_dict() + assert "bacon" in schema + assert schema["bacon"].get("status") is not None + + +def test_practitioner_routes_marked_through_diagnostic_handler(constructed): + out = practitioner_next_steps(constructed["PreTrendsPowerCurve"], verbose=False) + assert out["estimator"].endswith("(diagnostic result)") + labels = [s["label"] for s in out["next_steps"]] + # Estimator framing (Steps 1-2) is skipped for a diagnostic input. + assert "Define target parameter" not in labels + assert "estimation" not in out["completed"] + + +def test_practitioner_bacon_keeps_name_keyed_handler(bacon_result): + # Bacon stays on _handle_bacon with its estimator-selection framing. + out = practitioner_next_steps(bacon_result, verbose=False) + labels = [s["label"] for s in out["next_steps"]] + assert any("heterogeneity-robust estimator" in lbl for lbl in labels) + + +def test_describe_target_parameter_marks_diagnostic_kind(constructed): + # Defensive-depth branch for direct callers. + block = describe_target_parameter(constructed["PreTrendsPowerResults"]) + assert block["aggregation"] == "diagnostic" + assert block["headline_attribute"] == "" + + +def test_describe_target_parameter_bacon_still_twfe(bacon_result): + # The named Bacon branch wins over the generic diagnostic branch. + block = describe_target_parameter(bacon_result) + assert block["aggregation"] == "twfe" + assert block["headline_attribute"] == "twfe_estimate" diff --git a/tests/test_event_study_surface.py b/tests/test_event_study_surface.py new file mode 100644 index 000000000..acf3ba962 --- /dev/null +++ b/tests/test_event_study_surface.py @@ -0,0 +1,1121 @@ +"""Unified event-study representation (ledger M-092, spec section 5). + +This file is the ``test_ref`` for ledger row M-092. It covers: + +- ``EventStudyResults`` container validation, reference marking, the pinned + ``to_dataframe`` schema, and the summary-alpha contract. +- ``build_event_study_surface`` over all FOURTEEN producers: identical + column schema, bit-exact quintet vs the native representation on + non-reference rows, explicit ``is_reference`` marking (sentinels never + leaking into the container schema), per-producer expected reference-row + count, and absent-surface ValueErrors. + +Fits are small and analytical (no bootstrap is needed for any producer's +event-study surface); fixtures are module-scoped. +""" + +import warnings + +import numpy as np +import pandas as pd +import pytest + +from diff_diff.results_base import ( + EVENT_STUDY_SCHEMA, + EventStudyResults, + build_event_study_surface, +) +from diff_diff.spillover import SpilloverDiD +from diff_diff.sun_abraham import SunAbraham + +# =========================================================================== +# Container unit tests (hand-built instances - no fits) +# =========================================================================== + + +def _tiny_surface(**overrides): + kwargs = dict( + event_time=np.array([-1, 0, 1]), + att=np.array([0.0, 0.5, 0.6]), + se=np.array([np.nan, 0.1, 0.1]), + t_stat=np.array([np.nan, 5.0, 6.0]), + p_value=np.array([np.nan, 0.0, 0.0]), + conf_int_lower=np.array([np.nan, 0.3, 0.4]), + conf_int_upper=np.array([np.nan, 0.7, 0.8]), + is_reference=np.array([True, False, False]), + n=np.array([np.nan, 10.0, 10.0]), + n_kind="groups", + ) + kwargs.update(overrides) + return EventStudyResults(**kwargs) + + +def test_to_dataframe_schema_is_pinned(): + df = _tiny_surface().to_dataframe() + assert tuple(df.columns) == EVENT_STUDY_SCHEMA + + +def test_reference_row_inferred_and_marked(): + surface = _tiny_surface() + assert surface.reference_period == -1 + assert surface.is_reference.sum() == 1 + + +def test_empty_surface_allowed(): + # A requested-but-empty event study (no estimable horizons) is a valid + # zero-row surface with the pinned schema, not an error. + surface = EventStudyResults( + event_time=np.array([]), + att=np.array([]), + se=np.array([]), + t_stat=np.array([]), + p_value=np.array([]), + conf_int_lower=np.array([]), + conf_int_upper=np.array([]), + is_reference=np.array([], dtype=bool), + n=np.array([]), + ) + assert surface.reference_period is None + assert surface.reference_periods == [] + df = surface.to_dataframe() + assert tuple(df.columns) == EVENT_STUDY_SCHEMA + assert len(df) == 0 + assert isinstance(surface.summary(), str) + + +def test_multiple_reference_rows_allowed(): + # CallawaySantAnna universal base on a gapped grid materializes several + # reference-only horizons; the container must represent them. + surface = _tiny_surface(is_reference=np.array([True, True, False])) + assert int(surface.is_reference.sum()) == 2 + # reference_period (single-scalar convenience) is None when not exactly one. + assert surface.reference_period is None + assert surface.reference_periods == [-1, 0] + + +def test_length_mismatch_rejected(): + with pytest.raises(ValueError, match="align with event_time"): + _tiny_surface(att=np.array([0.0, 0.5])) + + +def test_vcov_requires_index(): + with pytest.raises(ValueError, match="vcov and vcov_index together"): + _tiny_surface(vcov=np.eye(3)) + + +def test_summary_alpha_mismatch_raises(): + surface = _tiny_surface(alpha=0.05) + with pytest.raises(ValueError, match="re-aggregate"): + surface.summary(alpha=0.10) + # None and the stored alpha are both fine. + assert isinstance(surface.summary(), str) + assert isinstance(surface.summary(alpha=0.05), str) + + +def test_object_dtype_calendar_event_time(): + # Calendar mode may carry string/period labels - no numeric assumptions. + surface = _tiny_surface( + event_time=np.array(["2018", "2019", "2020"], dtype=object), + is_reference=np.array([True, False, False]), + time_scale="calendar", + ) + df = surface.to_dataframe() + assert list(df["event_time"]) == ["2018", "2019", "2020"] + assert surface.reference_period == "2018" + + +def test_to_dict_json_safe_timestamp_labels(): + # pandas.Timestamp calendar labels must serialize (json.dumps would raise + # on raw Timestamp objects from .tolist()). + import json + + surface = _tiny_surface( + event_time=np.array( + [pd.Timestamp("2018-01-01"), pd.Timestamp("2019-01-01"), pd.Timestamp("2020-01-01")], + dtype=object, + ), + is_reference=np.array([True, False, False]), + time_scale="calendar", + ) + d = surface.to_dict() + json.dumps(d) # must not raise + assert d["event_time"][0] == "2018-01-01T00:00:00" + assert d["reference_period"] == "2018-01-01T00:00:00" + + +def test_to_dict_json_safe_period_labels(): + import json + + surface = _tiny_surface( + event_time=np.array( + [pd.Period("2018", "Y"), pd.Period("2019", "Y"), pd.Period("2020", "Y")], + dtype=object, + ), + is_reference=np.array([True, False, False]), + time_scale="calendar", + ) + d = surface.to_dict() + json.dumps(d) # must not raise + assert d["event_time"][0] == "2018" + assert d["reference_period"] == "2018" + + +# =========================================================================== +# Producer builders (small analytical fits) +# =========================================================================== + + +@pytest.fixture(scope="module") +def surfaces(): + """Build every producer's surface once; return {label: (native, surface)}.""" + out = {} + + from diff_diff import ( + CallawaySantAnna, + EfficientDiD, + LPDiD, + MultiPeriodDiD, + StackedDiD, + StaggeredTripleDifference, + WooldridgeDiD, + generate_staggered_ddd_data, + ) + from diff_diff.chaisemartin_dhaultfoeuille import ChaisemartinDHaultfoeuille + from diff_diff.continuous_did import ContinuousDiD + from diff_diff.datasets import load_mpdta + from diff_diff.imputation import ImputationDiD + from diff_diff.prep import generate_staggered_data as prep_staggered + from diff_diff.prep_dgp import ( + generate_continuous_did_data, + generate_reversible_did_data, + ) + from diff_diff.prep_dgp import generate_staggered_data as dgp_staggered + from diff_diff.spillover import SpilloverDiD + from diff_diff.sun_abraham import SunAbraham + from diff_diff.two_stage import TwoStageDiD + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + + # 1. CallawaySantAnna + cs_data = prep_staggered( + n_units=80, n_periods=8, cohort_periods=[4], treatment_effect=2.0, seed=42 + ) + cs = CallawaySantAnna(n_bootstrap=0).fit( + cs_data, + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="event_study", + ) + out["CallawaySantAnna"] = (cs, build_event_study_surface(cs)) + + # 2. SunAbraham (always-on) + sa_data = _sa_panel() + sa = SunAbraham().fit( + sa_data, outcome="outcome", unit="unit", time="time", first_treat="first_treat" + ) + out["SunAbraham"] = (sa, build_event_study_surface(sa)) + + # 3. ImputationDiD + imp_data = _bjs_panel() + imp = ImputationDiD().fit( + imp_data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + out["ImputationDiD"] = (imp, build_event_study_surface(imp)) + + # 4. TwoStageDiD (same panel shape) + ts = TwoStageDiD().fit( + imp_data, + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) + out["TwoStageDiD"] = (ts, build_event_study_surface(ts)) + + # 5. StackedDiD + st_data = dgp_staggered( + n_units=120, + n_periods=12, + cohort_periods=[4, 6, 8], + never_treated_frac=0.3, + treatment_effect=5.0, + dynamic_effects=True, + seed=42, + ) + st = StackedDiD(kappa_pre=2, kappa_post=2).fit( + st_data, + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="event_study", + ) + out["StackedDiD"] = (st, build_event_study_surface(st)) + + # 6. SpilloverDiD + from tests._dgp_utils import generate_butts_staggered_dgp + + sp_data = generate_butts_staggered_dgp(seed=42) + sp = SpilloverDiD( + rings=[0.0, 50.0, 200.0], + d_bar=200.0, + conley_coords=("lat", "lon"), + conley_metric="haversine", + conley_cutoff_km=200.0, + conley_lag_cutoff=0, + vcov_type="hc1", + event_study=True, + horizon_max=2, + anticipation=0, + ).fit(sp_data, outcome="y", unit="unit", time="time", first_treat="first_treat") + out["SpilloverDiD"] = (sp, build_event_study_surface(sp)) + + # 7. ContinuousDiD + cd_data = generate_continuous_did_data( + n_units=200, n_periods=5, cohort_periods=[2, 4], seed=42, noise_sd=0.5 + ) + cd = ContinuousDiD(control_group="not_yet_treated", n_bootstrap=0).fit( + cd_data, "outcome", "unit", "period", "first_treat", "dose", aggregate="eventstudy" + ) + out["ContinuousDiD"] = (cd, build_event_study_surface(cd)) + + # 8. EfficientDiD + ed_data = _efficient_panel() + ed = EfficientDiD().fit( + ed_data, "y", "unit", "time", "first_treat", aggregate="event_study" + ) + out["EfficientDiD"] = (ed, build_event_study_surface(ed)) + + # 9. WooldridgeDiD (post-fit aggregate) + mp = load_mpdta() + wd = WooldridgeDiD(control_group="never_treated").fit( + mp, outcome="lemp", unit="countyreal", time="year", cohort="first_treat" + ) + wd.aggregate("event") + out["WooldridgeDiD"] = (wd, build_event_study_surface(wd)) + + # 10. StaggeredTripleDifference + sddd_data = generate_staggered_ddd_data(n_units=300, treatment_effect=3.0, seed=42) + sddd = StaggeredTripleDifference().fit( + sddd_data, + "outcome", + "unit", + "period", + "first_treat", + "eligibility", + aggregate="event_study", + ) + out["StaggeredTripleDifference"] = (sddd, build_event_study_surface(sddd)) + + # 11. MultiPeriodDiD + mpd_data = _mpd_panel() + mpd = MultiPeriodDiD().fit( + mpd_data, + outcome="outcome", + treatment="treated", + time="period", + post_periods=[3, 4, 5], + reference_period=2, + ) + out["MultiPeriodDiD"] = (mpd, build_event_study_surface(mpd)) + + # 12. LPDiD + lp_data = _lpdid_panel() + lp = LPDiD(pre_window=2, post_window=2).fit( + lp_data, outcome="y", unit="unit", time="time", treatment="treat" + ) + out["LPDiD"] = (lp, build_event_study_surface(lp)) + + # 13. ChaisemartinDHaultfoeuille (event + placebo) + dcdh_data = generate_reversible_did_data( + n_groups=50, n_periods=10, pattern="joiners_only", seed=42 + ) + dcdh = ChaisemartinDHaultfoeuille(twfe_diagnostic=False).fit( + dcdh_data, + outcome="outcome", + group="group", + time="period", + treatment="treatment", + L_max=3, + ) + out["ChaisemartinDHaultfoeuille"] = (dcdh, build_event_study_surface(dcdh)) + + # 14. HeterogeneousAdoptionDiD event study + had_data, had_es = _had_event_study() + out["HeterogeneousAdoptionDiD"] = (had_es, build_event_study_surface(had_es)) + + return out + + +ALL_PRODUCERS = [ + "CallawaySantAnna", + "SunAbraham", + "ImputationDiD", + "TwoStageDiD", + "StackedDiD", + "SpilloverDiD", + "ContinuousDiD", + "EfficientDiD", + "WooldridgeDiD", + "StaggeredTripleDifference", + "MultiPeriodDiD", + "LPDiD", + "ChaisemartinDHaultfoeuille", + "HeterogeneousAdoptionDiD", +] + + +def _expected_reference_labels(label, native): + """The exact event_time(s) the surface should mark as reference. + + Reference resolution is producer-specific, NOT a count heuristic (a zero + count is not a universal reference marker - SpilloverDiD uses n_obs=0 for + both its genuine reference and its non-estimable rectangular horizons): + + - Explicit reference_period wins (SpilloverDiD, MultiPeriodDiD). + - Structurally-omitted baselines are synthesized (SunAbraham at + -1-anticipation; dCDH at 0). + - LPDiD's horizon==-1 base row. + - Wooldridge / HAD omit no baseline (no reference). + - The remaining relative-dict producers mark EVERY zero-count sentinel + 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"): + return [] if native.reference_period is None else [native.reference_period] + if label == "ChaisemartinDHaultfoeuille": + return [0] + if label == "SunAbraham": + return [-1 - (getattr(native, "anticipation", 0) or 0)] + if label == "LPDiD": + return [-1] if (native.event_study["horizon"] == -1).any() else [] + if label in ("WooldridgeDiD", "HeterogeneousAdoptionDiD"): + return [] + # Count-sentinel producers: a zero-count row is the reference only when its + # effect is finite (0.0); a NaN-effect zero-count row is non-estimable. + return sorted( + e + for e, row in native.event_study_effects.items() + if (row.get("n_groups", 1) == 0 or row.get("n_obs", 1) == 0) + and np.isfinite(row.get("att", row.get("effect", np.nan))) + ) + + +def test_all_producers_built(surfaces): + assert set(surfaces) == set(ALL_PRODUCERS) + + +@pytest.mark.parametrize("label", ALL_PRODUCERS) +def test_identical_schema(label, surfaces): + _, surface = surfaces[label] + assert tuple(surface.to_dataframe().columns) == EVENT_STUDY_SCHEMA + + +@pytest.mark.parametrize("label", ALL_PRODUCERS) +def test_reference_labels_match_native(label, surfaces): + # is_reference marks EXACTLY the producer-specific reference label(s) - + # never an arbitrary zero-count row (the count heuristic would mislabel + # SpilloverDiD's non-estimable horizons and miss SunAbraham's anchor). + native, surface = surfaces[label] + expected = _expected_reference_labels(label, native) + marked = sorted(surface.event_time[surface.is_reference].tolist()) + assert marked == expected, f"{label}: marked {marked}, expected {expected}" + + +@pytest.mark.parametrize("label", ALL_PRODUCERS) +def test_reference_row_values_normalized(label, surfaces): + _, surface = surfaces[label] + ref = surface.is_reference + if ref.any(): + assert surface.att[ref][0] == 0.0 + assert np.isnan(surface.se[ref][0]) + assert np.isnan(surface.conf_int_lower[ref][0]) + + +@pytest.mark.parametrize("label", ALL_PRODUCERS) +def test_no_sentinel_in_schema(label, surfaces): + # The retiring sentinels (n_groups==0 / n_obs==0) must not surface as a + # zero count on the marked reference row: it is NaN there. + _, surface = surfaces[label] + ref = surface.is_reference + if ref.any(): + assert np.isnan(surface.n[ref][0]) + + +@pytest.mark.parametrize("label", ALL_PRODUCERS) +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. + native, surface = surfaces[label] + non_ref = ~surface.is_reference + native_map = _native_event_map(label, native) + checked = 0 + for i in np.where(non_ref)[0]: + e = surface.event_time[i] + if e not in native_map: + continue + exp = native_map[e] + assert surface.att[i] == pytest.approx(exp["att"], nan_ok=True) + assert surface.se[i] == pytest.approx(exp["se"], nan_ok=True) + assert surface.t_stat[i] == pytest.approx(exp["t_stat"], nan_ok=True) + assert surface.p_value[i] == pytest.approx(exp["p_value"], nan_ok=True) + assert surface.conf_int_lower[i] == pytest.approx(exp["ci_lo"], nan_ok=True) + assert surface.conf_int_upper[i] == pytest.approx(exp["ci_hi"], nan_ok=True) + checked += 1 + assert checked > 0, f"{label}: no non-reference rows compared" + + +def test_cs_vcov_alignment(surfaces): + native, surface = surfaces["CallawaySantAnna"] + assert surface.vcov is not None + assert surface.vcov_index is not None + # vcov_index labels are a subset of event_time labels. + assert set(surface.vcov_index.tolist()).issubset(set(surface.event_time.tolist())) + # Same ordered index the native surface exposes. + np.testing.assert_array_equal(np.asarray(native.event_study_vcov_index), surface.vcov_index) + + +def test_mpd_vcov_subblock(surfaces): + native, surface = surfaces["MultiPeriodDiD"] + # MPD supplies a full vcov + interaction_indices; the surface exposes an + # ordered sub-block. Presence (not exact values) is the contract here. + if native.vcov is not None and native.interaction_indices: + assert surface.vcov is not None + assert surface.vcov.shape[0] == surface.vcov_index.shape[0] + + +def test_dcdh_convention_and_placebo_merge(surfaces): + native, surface = surfaces["ChaisemartinDHaultfoeuille"] + assert surface.event_time_convention == "l1_first_switch" + times = surface.event_time.tolist() + assert 0 in times # synthesized reference + assert any(t < 0 for t in times) # placebo horizons merged in + assert any(t >= 1 for t in times) # post horizons + + +def test_dcdh_n_kind_is_groups_not_obs(surfaces): + # dCDH stores N_l (eligible switcher GROUPS) under its legacy "n_obs" key. + # The unified surface must label it "groups", never "obs" - a consumer + # doing sample-size logic would otherwise misread switcher groups as + # observations. + native, surface = surfaces["ChaisemartinDHaultfoeuille"] + assert surface.n_kind == "groups" + df = surface.to_dataframe() + # Non-reference n values equal the native N_l (from event_study_effects / + # placebo_event_study), not an observation count. + native_n = {} + for k, row in native.event_study_effects.items(): + native_n[k] = row.get("n_obs") + for k, row in (native.placebo_event_study or {}).items(): + native_n[k] = row.get("n_obs") + for _, r in df[~df["is_reference"]].iterrows(): + e = r["event_time"] + if e in native_n and native_n[e] is not None: + assert r["n"] == float(native_n[e]) + + +def test_dcdh_l_max_none_count_kind_is_switcher_cells(): + # Legacy single-horizon path (L_max=None): the count under "n_obs" is N_S, + # the number of switching (g,t) CELLS - one group can contribute several - + # so n_kind is "switcher_cells", NOT "groups" and NOT "obs". + class _FakeDCDHLegacy: + alpha = 0.05 + L_max = None + placebo_event_study = None + event_study_effects = { + 1: { + "effect": 1.0, + "se": 0.2, + "t_stat": 5.0, + "p_value": 0.0, + "conf_int": (0.6, 1.4), + "n_obs": 30, # N_S switching cells, NOT observations or groups + } + } + sup_t_bands = None + + surface = build_event_study_surface(_FakeDCDHLegacy()) + assert surface.n_kind == "switcher_cells" + + +def test_spillover_oversized_horizon_single_reference(): + # Regression for the count-heuristic bug: with horizon_max=4 SpilloverDiD + # emits BOTH a genuine reference (reference_period, n_obs=0, coef=0) and a + # non-estimable rectangular horizon (n_obs=0, coef=NaN). Only the former + # is the reference; the latter must survive as a non-reference NaN row. + from tests._dgp_utils import generate_butts_staggered_dgp + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + sp = SpilloverDiD( + rings=[0.0, 50.0, 200.0], + d_bar=200.0, + conley_coords=("lat", "lon"), + conley_metric="haversine", + conley_cutoff_km=200.0, + conley_lag_cutoff=0, + vcov_type="hc1", + event_study=True, + horizon_max=4, + anticipation=0, + ).fit( + generate_butts_staggered_dgp(seed=42), + outcome="y", + unit="unit", + time="time", + first_treat="first_treat", + ) + surface = build_event_study_surface(sp) + df = surface.to_dataframe() + # Native has >=2 zero-count rows; exactly one is the reference. + zero_native = [k for k, r in sp.event_study_effects.items() if r.get("n_obs") == 0] + assert len(zero_native) >= 2 + marked = df[df["is_reference"]]["event_time"].tolist() + assert marked == [sp.reference_period] + # A non-reference zero-count horizon keeps its NaN att and n=0. + other_zero = next(k for k in zero_native if k != sp.reference_period) + row = df[df["event_time"] == other_zero].iloc[0] + assert not bool(row["is_reference"]) + assert np.isnan(row["att"]) + assert row["n"] == 0.0 + + +def test_sun_abraham_synthesizes_anticipation_reference(): + # SunAbraham omits its baseline; positive anticipation shifts the + # synthesized reference to e = -1 - anticipation. + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + sa = SunAbraham(anticipation=1).fit( + _sa_panel(n_periods=10), + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + ) + surface = build_event_study_surface(sa) + marked = surface.event_time[surface.is_reference].tolist() + assert marked == [-2] # -1 - anticipation + ref_idx = int(np.where(surface.is_reference)[0][0]) + assert surface.att[ref_idx] == 0.0 + assert np.isnan(surface.se[ref_idx]) + + +def test_sun_abraham_unobserved_reference_not_fabricated(): + # SunAbraham records reference_period = -1 - anticipation even when that + # period was never observed (reference_observed=False). The adapter must + # NOT invent the anchor row. + class _FakeSA: + alpha = 0.05 + anticipation = 0 + reference_period = -1 + reference_observed = False + event_study_effects = { + 0: {"effect": 1.0, "se": 0.1, "n_groups": 5}, + 1: {"effect": 1.2, "se": 0.1, "n_groups": 5}, + 2: {"effect": 1.3, "se": 0.1, "n_groups": 5}, + } + + surface = build_event_study_surface(_FakeSA()) + assert int(surface.is_reference.sum()) == 0 + assert -1 not in surface.event_time.tolist() + + +def test_sun_abraham_gapped_grid_unobserved_reference(): + # Reviewer's case: native keys {-2, 0}. -1 falls BETWEEN them, so a + # range-based guard would wrongly synthesize it. reference_observed=False + # (the -1 gap was never in the data) => no reference row. + class _FakeSAGap: + alpha = 0.05 + anticipation = 0 + reference_period = -1 + reference_observed = False + event_study_effects = { + -2: {"effect": -0.1, "se": 0.1, "n_groups": 5}, + 0: {"effect": 1.0, "se": 0.1, "n_groups": 5}, + } + + surface = build_event_study_surface(_FakeSAGap()) + assert int(surface.is_reference.sum()) == 0 + assert -1 not in surface.event_time.tolist() + + +def test_sun_abraham_gapped_grid_observed_reference_synthesized(): + # Same keys {-2, 0} but the -1 anchor WAS observed (omitted baseline): + # synthesize the reference row. + class _FakeSAObs: + alpha = 0.05 + anticipation = 0 + reference_period = -1 + reference_observed = True + event_study_effects = { + -2: {"effect": -0.1, "se": 0.1, "n_groups": 5}, + 0: {"effect": 1.0, "se": 0.1, "n_groups": 5}, + } + + surface = build_event_study_surface(_FakeSAObs()) + marked = surface.event_time[surface.is_reference].tolist() + assert marked == [-1] + ref_idx = int(np.where(surface.is_reference)[0][0]) + assert surface.att[ref_idx] == 0.0 + assert np.isnan(surface.se[ref_idx]) + + +def test_cs_universal_gapped_multiple_references(): + # Regression: CS base_period="universal" on a gapped grid ({1,3,6}, + # cohorts {3,6}) materializes each cohort's positional base as its own + # reference-only horizon (all n_groups==0). The surface must represent + # all of them - the count path marks every sentinel, and the container + # allows multiple references. + from diff_diff import CallawaySantAnna + + rng = np.random.RandomState(3) + rows = [] + for u in range(40): + g = 3 if u < 15 else (6 if u < 30 else 0) + ufe = rng.randn() * 2 + for t in (1, 3, 6): + post = 1 if (g > 0 and t >= g) else 0 + rows.append( + { + "unit": u, + "period": t, + "outcome": ufe + 0.3 * t + 2.0 * post + rng.randn() * 0.5, + "first_treat": g, + } + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + cs = CallawaySantAnna(n_bootstrap=0, base_period="universal").fit( + pd.DataFrame(rows), + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="event_study", + ) + native_refs = sorted(e for e, r in cs.event_study_effects.items() if r.get("n_groups") == 0) + assert len(native_refs) >= 2 # the case only matters when >1 + surface = build_event_study_surface(cs) + marked = sorted(surface.event_time[surface.is_reference].tolist()) + assert marked == native_refs + assert surface.reference_period is None # not exactly one + assert surface.reference_periods == native_refs + # Every reference row is normalized (att=0, NaN inference). + df = surface.to_dataframe() + ref_df = df[df["is_reference"]] + assert (ref_df["att"] == 0.0).all() + assert ref_df["se"].isna().all() + + +def test_requested_but_empty_event_study_builds_zero_row_surface(): + # event_study_effects == {} means REQUESTED but no estimable horizons + # (EfficientDiD balance_e removing every cohort) - a zero-row surface, + # distinct from None (never requested -> raises). + class _RequestedEmpty: + alpha = 0.05 + event_study_effects: dict = {} + + surface = build_event_study_surface(_RequestedEmpty()) + assert len(surface.to_dataframe()) == 0 + assert surface.reference_periods == [] + + +def test_absent_surface_distinguished_from_empty(): + # None (attribute present but None / missing) still raises the refit hint. + class _NotRequested: + alpha = 0.05 + event_study_effects = None + + with pytest.raises(ValueError, match="no event-study surface"): + build_event_study_surface(_NotRequested()) + + +def test_zero_count_finite_nonzero_effect_raises(): + # A zero-count row with a finite NONZERO effect is malformed (a reference + # must be exactly 0) - fail loudly rather than silently rewrite to 0. + class _Malformed: + alpha = 0.05 + event_study_effects = { + 0: {"effect": 1.0, "se": 0.1, "n_obs": 40}, + 1: {"effect": 3.3, "se": np.nan, "n_obs": 0}, # nonzero + zero count + } + + with pytest.raises(ValueError, match="finite nonzero effect"): + build_event_study_surface(_Malformed()) + + +def test_zero_count_nan_horizon_is_not_a_reference(): + # TwoStageDiD emits effect=NaN, n_obs=0 for an estimated horizon whose + # observations are all filtered (two_stage.py:2669-2681) - distinct from + # its effect=0.0, n_obs=0 reference. The count-sentinel adapter path must + # mark ONLY the finite-effect (0.0) row as reference and preserve the + # NaN-effect horizon as a non-reference NaN row (never normalize it to 0). + class _FakeTwoStage: + alpha = 0.05 + event_study_effects = { + -1: { # the reference: effect 0.0, n_obs 0 + "effect": 0.0, + "se": 0.0, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (0.0, 0.0), + "n_obs": 0, + }, + 0: { # an identified horizon + "effect": 1.5, + "se": 0.2, + "t_stat": 7.5, + "p_value": 0.0, + "conf_int": (1.1, 1.9), + "n_obs": 40, + }, + 1: { # non-estimable: effect NaN, n_obs 0 (NOT a reference) + "effect": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "n_obs": 0, + }, + } + + surface = build_event_study_surface(_FakeTwoStage()) + df = surface.to_dataframe() + marked = sorted(surface.event_time[surface.is_reference].tolist()) + assert marked == [-1], "only the finite-effect zero-count row is the reference" + # The NaN-effect zero-count horizon survives untouched. + row = df[df["event_time"] == 1].iloc[0] + assert not bool(row["is_reference"]) + assert np.isnan(row["att"]) + assert row["n"] == 0.0 + + +def test_reference_period_scalar_forced_none_when_not_unique(): + # A caller cannot leave reference_period disagreeing with reference_periods: + # the scalar is authoritative only when there is exactly one reference. + surface = _tiny_surface( + is_reference=np.array([True, True, False]), + reference_period=99, # bogus caller-supplied scalar + ) + assert surface.reference_period is None + assert surface.reference_periods == [-1, 0] + + +def test_single_reference_scalar_derived_over_supplied(): + # With exactly one reference, a mismatched caller-supplied scalar is + # overridden by the marked label - never retained. + surface = _tiny_surface( + is_reference=np.array([False, True, False]), # marks event_time 0 + reference_period=99, + ) + assert surface.reference_period == 0 + + +def test_reference_row_inference_normalized_by_container(): + # The public container enforces the reference-row contract itself: finite + # inference on a marked row is normalized to att=0 / NaN, not trusted. + surface = _tiny_surface( + att=np.array([5.0, 0.5, 0.6]), # bogus finite att on the reference + se=np.array([9.9, 0.1, 0.1]), + n=np.array([99.0, 10.0, 10.0]), + is_reference=np.array([True, False, False]), + ) + assert surface.att[0] == 0.0 + assert np.isnan(surface.se[0]) + assert np.isnan(surface.t_stat[0]) + assert np.isnan(surface.n[0]) + + +def test_cband_requires_both_bounds(): + with pytest.raises(ValueError, match="cband_lower and cband_upper together"): + _tiny_surface(cband_lower=np.array([np.nan, 0.2, 0.3])) + + +def test_container_does_not_mutate_caller_arrays(): + # __post_init__ normalizes reference rows in place; it must copy first so + # the caller's own arrays are never overwritten. + att = np.array([5.0, 0.5, 0.6]) # finite att on the reference row (index 0) + se = np.array([9.9, 0.1, 0.1]) + _tiny_surface(att=att, se=se, is_reference=np.array([True, False, False])) + assert att[0] == 5.0 # untouched + assert se[0] == 9.9 + + +def test_container_accepts_read_only_arrays(): + att = np.array([0.0, 0.5, 0.6]) + att.setflags(write=False) + se = np.array([np.nan, 0.1, 0.1]) + se.setflags(write=False) + surface = _tiny_surface(att=att, se=se, is_reference=np.array([True, False, False])) + assert surface.att[0] == 0.0 # constructed without error + + +def test_vcov_round_trips_through_to_dict(): + import json + + vcov = np.array([[0.04, 0.01], [0.01, 0.09]]) + surface = _tiny_surface(vcov=vcov, vcov_index=np.array([0, 1])) + d = surface.to_dict() + json.dumps(d) # must not raise + np.testing.assert_array_equal(np.asarray(d["vcov"]), vcov) + assert d["vcov_index"] == [0, 1] + + +def test_dcdh_crit_value_carried_from_sup_t_bands(): + # Constructed dCDH-shaped result: cband_crit_value is copied from + # sup_t_bands["crit_value"] (bootstrap-only in practice). + class _FakeDCDH: + alpha = 0.05 + placebo_event_study = {-1: {"effect": 0.1, "se": 0.05, "n_obs": 40}} + event_study_effects = { + 1: { + "effect": 1.0, + "se": 0.2, + "t_stat": 5.0, + "p_value": 0.0, + "conf_int": (0.6, 1.4), + "n_obs": 50, + "cband_conf_int": (0.4, 1.6), + } + } + sup_t_bands = {"crit_value": 2.71, "alpha": 0.05} + + surface = build_event_study_surface(_FakeDCDH()) + assert surface.cband_crit_value == 2.71 + assert surface.event_time_convention == "l1_first_switch" + + +def test_had_carries_cband_and_crit_value(): + # HAD's simultaneous-band bounds must travel with their critical value. + class _FakeHAD: + alpha = 0.05 + event_times = np.array([0, 1, 2]) + att = np.array([1.0, 1.2, 1.3]) + se = np.array([0.2, 0.2, 0.2]) + t_stat = np.array([5.0, 6.0, 6.5]) + p_value = np.array([0.0, 0.0, 0.0]) + conf_int_low = np.array([0.6, 0.8, 0.9]) + conf_int_high = np.array([1.4, 1.6, 1.7]) + n_obs_per_horizon = np.array([40, 40, 40]) + cband_low = np.array([0.4, 0.6, 0.7]) + cband_high = np.array([1.6, 1.8, 1.9]) + cband_crit_value = 2.5 + + surface = build_event_study_surface(_FakeHAD()) + assert surface.cband_crit_value == 2.5 + np.testing.assert_array_equal(surface.cband_lower, _FakeHAD.cband_low) + np.testing.assert_array_equal(surface.cband_upper, _FakeHAD.cband_high) + + +def test_wooldridge_key_normalized(surfaces): + native, surface = surfaces["WooldridgeDiD"] + # Native uses inner key "att"; the surface exposes canonical att column. + assert "att" in surface.to_dataframe().columns + assert surface.n_kind is None # Wooldridge records no per-row count + + +def test_lpdid_n_kind_obs(surfaces): + _, surface = surfaces["LPDiD"] + assert surface.n_kind == "obs" + # n_clusters dropped from the unified view; the schema is the pinned one. + assert "n_clusters" not in surface.to_dataframe().columns + + +def test_absent_surface_raises(): + from diff_diff import CallawaySantAnna + from diff_diff.prep import generate_staggered_data as prep_staggered + + data = prep_staggered(n_units=60, n_periods=6, cohort_periods=[3], treatment_effect=2.0, seed=1) + # aggregate defaults to "simple": no event_study_effects populated. + cs = CallawaySantAnna(n_bootstrap=0).fit( + data, outcome="outcome", unit="unit", time="period", first_treat="first_treat" + ) + with pytest.raises(ValueError, match="event-study"): + build_event_study_surface(cs) + + +# =========================================================================== +# Panel builders (inlined from the existing suite; small + analytical) +# =========================================================================== + + +def _sa_panel(n_units=80, n_periods=8, n_cohorts=3, seed=42): + np.random.seed(seed) + units = np.repeat(np.arange(n_units), n_periods) + times = np.tile(np.arange(n_periods), n_units) + n_never = int(n_units * 0.3) + n_treated = n_units - n_never + cohort_periods = np.linspace(3, n_periods - 2, n_cohorts).astype(int) + first_treat = np.zeros(n_units) + if n_treated > 0: + first_treat[n_never:] = cohort_periods[ + np.random.choice(len(cohort_periods), size=n_treated) + ] + fte = np.repeat(first_treat, n_periods) + unit_fe = np.repeat(np.random.randn(n_units) * 2, n_periods) + time_fe = np.tile(np.linspace(0, 1, n_periods), n_units) + post = (times >= fte) & (fte > 0) + dyn = 2.0 * (1 + 0.1 * np.maximum(times - fte, 0)) + y = unit_fe + time_fe + dyn * post + np.random.randn(len(units)) * 0.5 + return pd.DataFrame( + {"unit": units, "time": times, "outcome": y, "first_treat": fte.astype(int)} + ) + + +def _bjs_panel(n_units=100, n_periods=10, seed=42): + rng = np.random.default_rng(seed) + units = np.repeat(np.arange(n_units), n_periods) + times = np.tile(np.arange(n_periods), n_units) + n_never = int(n_units * 0.3) + n_treated = n_units - n_never + cohort_periods = np.array([3, 5, 7]) + first_treat = np.zeros(n_units, dtype=int) + if n_treated > 0: + first_treat[n_never:] = cohort_periods[rng.choice(len(cohort_periods), size=n_treated)] + fte = np.repeat(first_treat, n_periods) + unit_fe = np.repeat(rng.standard_normal(n_units) * 2.0, n_periods) + time_fe = np.tile(np.linspace(0, 1, n_periods), n_units) + post = (times >= fte) & (fte > 0) + rel = times - fte + mult = 1 + 0.1 * np.maximum(rel, 0) + y = unit_fe + time_fe + (2.0 * mult) * post + rng.standard_normal(len(units)) * 0.5 + return pd.DataFrame({"unit": units, "time": times, "outcome": y, "first_treat": fte}) + + +def _efficient_panel(n_units=100, n_periods=5, n_treated=50, treat_period=3, seed=42): + rng = np.random.default_rng(seed) + units = np.repeat(np.arange(n_units), n_periods) + times = np.tile(np.arange(1, n_periods + 1), n_units) + ft = np.full(n_units, np.inf) + ft[:n_treated] = treat_period + ft_col = np.repeat(ft, n_periods) + unit_fe = np.repeat(rng.normal(0, 1, n_units), n_periods) + time_fe = np.tile(np.arange(1, n_periods + 1) * 0.5, n_units) + tau = np.where((ft_col < np.inf) & (times >= ft_col), 2.0, 0.0) + y = unit_fe + time_fe + tau + rng.normal(0, 0.5, len(units)) + return pd.DataFrame({"unit": units, "time": times, "first_treat": ft_col, "y": y}) + + +def _mpd_panel(n_units=100, n_periods=6, seed=42): + np.random.seed(seed) + rows = [] + 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) + rows.append({"unit": unit, "period": period, "treated": int(is_treated), "outcome": y}) + return pd.DataFrame(rows) + + +def _lpdid_panel(seed=7): + rng = np.random.default_rng(seed) + rows = [] + uid = 0 + for _ in range(12): + uid += 1 + alpha = rng.normal(0.0, 0.5) + for t in range(8): + y0 = alpha + 0.5 * t + rng.normal(0.0, 0.2) + k = t - 4 + effect = (1.0 + 0.5 * k) if k >= 0 else 0.0 + rows.append( + {"unit": uid, "time": t, "y": y0 + effect, "treat": int(t >= 4), "first_treat": 4.0} + ) + for _ in range(12): + uid += 1 + alpha = rng.normal(0.0, 0.5) + for t in range(8): + y0 = alpha + 0.5 * t + rng.normal(0.0, 0.2) + rows.append({"unit": uid, "time": t, "y": y0, "treat": 0, "first_treat": np.inf}) + return pd.DataFrame(rows) + + +def _had_event_study(): + from diff_diff.had import HeterogeneousAdoptionDiD + + rng = np.random.default_rng(0) + G = 300 + mass_n = int(0.3 * G) + d_at_F = np.concatenate([np.full(mass_n, 0.5), rng.uniform(0.5, 1.0, G - mass_n)]) + units = np.arange(G) + alpha_g = 0.5 * rng.standard_normal(G) + F, n_periods = 3, 5 + rows = [] + for g in units: + d_g = float(d_at_F[g]) + for t in range(1, n_periods + 1): + dose = d_g if t >= F else 0.0 + outcome = alpha_g[g] + 0.3 * dose + 0.1 * rng.standard_normal() + rows.append({"unit": g, "period": t, "dose": dose, "outcome": outcome}) + panel = pd.DataFrame(rows) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + es = HeterogeneousAdoptionDiD(design="mass_point").fit( + panel, "outcome", "dose", "period", "unit", aggregate="event_study" + ) + return panel, es + + +def _native_event_map(label, native): + """Map event_time -> full canonical quintet from the native representation.""" + out = {} + if label == "LPDiD": + for _, row in native.event_study.iterrows(): + out[row["horizon"]] = { + "att": row["coefficient"], + "se": row["se"], + "t_stat": row["t_stat"], + "p_value": row["p_value"], + "ci_lo": row["conf_low"], + "ci_hi": row["conf_high"], + } + elif label == "MultiPeriodDiD": + for period, pe in native.period_effects.items(): + out[period] = { + "att": pe.effect, + "se": pe.se, + "t_stat": pe.t_stat, + "p_value": pe.p_value, + "ci_lo": pe.conf_int[0], + "ci_hi": pe.conf_int[1], + } + elif label == "HeterogeneousAdoptionDiD": + for i, e in enumerate(native.event_times): + out[e] = { + "att": native.att[i], + "se": native.se[i], + "t_stat": native.t_stat[i], + "p_value": native.p_value[i], + "ci_lo": native.conf_int_low[i], + "ci_hi": native.conf_int_high[i], + } + else: # relative-dict + dCDH producers + source = dict(native.event_study_effects) + placebos = getattr(native, "placebo_event_study", None) + if placebos: + source.update(placebos) + for e, row in source.items(): + att = row["att"] if "att" in row else row["effect"] + ci = row.get("conf_int", (np.nan, np.nan)) + out[e] = { + "att": att, + "se": row.get("se", np.nan), + "t_stat": row.get("t_stat", np.nan), + "p_value": row.get("p_value", np.nan), + "ci_lo": ci[0], + "ci_hi": ci[1], + } + return out diff --git a/tests/test_results_serialization.py b/tests/test_results_serialization.py new file mode 100644 index 000000000..dba60b795 --- /dev/null +++ b/tests/test_results_serialization.py @@ -0,0 +1,284 @@ +"""Serialization contract for the Phase 2 results foundation (spec section 5). + +- The seven newly-added ``to_dict()`` methods emit canonical keys, values + equal to the property surface, and JSON-serializable scalars. +- Every estimator results class subclasses ``BaseResults``; no class is + both ``BaseResults`` and ``Diagnostic``; ``EventStudyResults`` is + ``BaseResults`` and not ``Diagnostic``. +""" + +import json +import warnings + +import numpy as np +import pytest + +import diff_diff +from diff_diff import BaseResults, Diagnostic, EventStudyResults + +# Estimator results classes that gained to_dict() in this PR. +NEW_TO_DICT = [ + "CallawaySantAnnaResults", + "StackedDiDResults", + "ContinuousDiDResults", + "SunAbrahamResults", + "WooldridgeDiDResults", + "ChaisemartinDHaultfoeuilleResults", + "BaconDecompositionResults", +] + +# Every public estimator results class (must be BaseResults, never Diagnostic). +ESTIMATOR_RESULTS = [ + "DiDResults", + "SpilloverDiDResults", + "MultiPeriodDiDResults", + "SyntheticDiDResults", + "CallawaySantAnnaResults", + "SunAbrahamResults", + "ImputationDiDResults", + "TwoStageDiDResults", + "StackedDiDResults", + "EfficientDiDResults", + "ContinuousDiDResults", + "WooldridgeDiDResults", + "ChaisemartinDHaultfoeuilleResults", + "LPDiDResults", + "StaggeredTripleDiffResults", + "TripleDifferenceResults", + "TROPResults", + "SyntheticControlResults", + "ChangesInChangesResults", + "HeterogeneousAdoptionDiDResults", + "HeterogeneousAdoptionDiDEventStudyResults", + "RegressionDiscontinuityResults", +] + +CANONICAL_QUINTET_KEYS = {"att", "se", "t_stat", "p_value", "conf_int_lower", "conf_int_upper"} + + +# --------------------------------------------------------------------------- +# to_dict() behavioral tests (real small fits) +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def fitted_results(): + """One fitted instance per to_dict-gaining estimator (except Bacon fits + its own staggered panel; all analytical).""" + out = {} + from diff_diff import ( + BaconDecomposition, + CallawaySantAnna, + StackedDiD, + WooldridgeDiD, + ) + from diff_diff.continuous_did import ContinuousDiD + from diff_diff.datasets import load_mpdta + from diff_diff.prep import generate_staggered_data as prep_staggered + from diff_diff.prep_dgp import ( + generate_continuous_did_data, + ) + from diff_diff.prep_dgp import generate_staggered_data as dgp_staggered + from diff_diff.sun_abraham import SunAbraham + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + + cs_data = prep_staggered( + n_units=80, n_periods=8, cohort_periods=[4], treatment_effect=2.0, seed=42 + ) + out["CallawaySantAnnaResults"] = CallawaySantAnna(n_bootstrap=0).fit( + cs_data, outcome="outcome", unit="unit", time="period", first_treat="first_treat" + ) + + out["SunAbrahamResults"] = SunAbraham().fit( + _sa_panel(), outcome="outcome", unit="unit", time="time", first_treat="first_treat" + ) + + st_data = dgp_staggered( + n_units=120, + n_periods=12, + cohort_periods=[4, 6, 8], + never_treated_frac=0.3, + treatment_effect=5.0, + dynamic_effects=True, + seed=42, + ) + out["StackedDiDResults"] = StackedDiD(kappa_pre=2, kappa_post=2).fit( + st_data, outcome="outcome", unit="unit", time="period", first_treat="first_treat" + ) + + cd_data = generate_continuous_did_data( + n_units=200, n_periods=5, cohort_periods=[2, 4], seed=42, noise_sd=0.5 + ) + out["ContinuousDiDResults"] = ContinuousDiD( + control_group="not_yet_treated", n_bootstrap=0 + ).fit(cd_data, "outcome", "unit", "period", "first_treat", "dose") + + out["WooldridgeDiDResults"] = WooldridgeDiD().fit( + load_mpdta(), outcome="lemp", unit="countyreal", time="year", cohort="first_treat" + ) + + out["ChaisemartinDHaultfoeuilleResults"] = _dcdh_fit() + + out["BaconDecompositionResults"] = BaconDecomposition().fit( + _bacon_panel(), + outcome="outcome", + unit="unit", + time="time", + first_treat="first_treat", + ) + return out + + +@pytest.mark.parametrize("name", NEW_TO_DICT) +def test_to_dict_returns_dict(name, fitted_results): + d = fitted_results[name].to_dict() + assert isinstance(d, dict) and d + + +@pytest.mark.parametrize("name", [n for n in NEW_TO_DICT if n != "BaconDecompositionResults"]) +def test_to_dict_has_canonical_quintet(name, fitted_results): + d = fitted_results[name].to_dict() + assert CANONICAL_QUINTET_KEYS.issubset(d.keys()), ( + f"{name}.to_dict() missing canonical keys: " f"{CANONICAL_QUINTET_KEYS - set(d.keys())}" + ) + + +@pytest.mark.parametrize("name", [n for n in NEW_TO_DICT if n != "BaconDecompositionResults"]) +def test_to_dict_values_match_property_surface(name, fitted_results): + res = fitted_results[name] + d = res.to_dict() + assert d["att"] == pytest.approx(res.att, nan_ok=True) + assert d["se"] == pytest.approx(res.se, nan_ok=True) + assert d["conf_int_lower"] == pytest.approx(res.conf_int[0], nan_ok=True) + assert d["conf_int_upper"] == pytest.approx(res.conf_int[1], nan_ok=True) + + +def test_bacon_to_dict_has_no_quintet(fitted_results): + d = fitted_results["BaconDecompositionResults"].to_dict() + assert "att" not in d + assert "twfe_estimate" in d + + +@pytest.mark.parametrize("name", NEW_TO_DICT) +def test_to_dict_no_deprecated_names(name, fitted_results): + # Serialization emits canonical names only (spec section 5): no + # overall_*/avg_* keys leak into to_dict output. + d = fitted_results[name].to_dict() + leaked = [k for k in d if k.startswith("overall_") or k.startswith("avg_")] + assert not leaked, f"{name}.to_dict() leaked deprecated keys: {leaked}" + + +@pytest.mark.parametrize("name", NEW_TO_DICT) +def test_to_dict_scalars_json_serializable(name, fitted_results): + d = fitted_results[name].to_dict() + # numpy scalar types are the common JSON offender; coerce and dump. + coerced = { + k: (v.item() if hasattr(v, "item") else v) + for k, v in d.items() + if not isinstance(v, (list, dict, np.ndarray)) + } + json.dumps(coerced) + + +# --------------------------------------------------------------------------- +# BaseResults / Diagnostic coverage (introspection - no fits) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("name", ESTIMATOR_RESULTS) +def test_estimator_result_is_base_results(name): + cls = getattr(diff_diff, name) + assert issubclass(cls, BaseResults), f"{name} must subclass BaseResults" + + +def test_no_class_is_both_base_and_diagnostic(): + # Scan every public results/diagnostic class in the namespace. + import diff_diff.diagnostic_report as drm + + both = [] + for name in dir(diff_diff): + obj = getattr(diff_diff, name) + if isinstance(obj, type) and issubclass(obj, (BaseResults, Diagnostic)): + if issubclass(obj, BaseResults) and issubclass(obj, Diagnostic): + both.append(name) + # DiagnosticReportResults lives in diagnostic_report; check it too. + drr = drm.DiagnosticReportResults + if issubclass(drr, BaseResults) and issubclass(drr, Diagnostic): + both.append("DiagnosticReportResults") + assert not both, f"classes are BOTH BaseResults and Diagnostic: {both}" + + +def test_event_study_results_is_base_not_diagnostic(): + assert issubclass(EventStudyResults, BaseResults) + assert not issubclass(EventStudyResults, Diagnostic) + + +# --------------------------------------------------------------------------- +# Panel builders (small, analytical) +# --------------------------------------------------------------------------- + + +def _sa_panel(n_units=80, n_periods=8, n_cohorts=3, seed=42): + import pandas as pd + + np.random.seed(seed) + units = np.repeat(np.arange(n_units), n_periods) + times = np.tile(np.arange(n_periods), n_units) + n_never = int(n_units * 0.3) + n_treated = n_units - n_never + cohort_periods = np.linspace(3, n_periods - 2, n_cohorts).astype(int) + first_treat = np.zeros(n_units) + if n_treated > 0: + first_treat[n_never:] = cohort_periods[ + np.random.choice(len(cohort_periods), size=n_treated) + ] + fte = np.repeat(first_treat, n_periods) + unit_fe = np.repeat(np.random.randn(n_units) * 2, n_periods) + time_fe = np.tile(np.linspace(0, 1, n_periods), n_units) + post = (times >= fte) & (fte > 0) + y = unit_fe + time_fe + 2.0 * post + np.random.randn(len(units)) * 0.5 + return pd.DataFrame( + {"unit": units, "time": times, "outcome": y, "first_treat": fte.astype(int)} + ) + + +def _bacon_panel(n_units=100, n_periods=10, n_cohorts=3, seed=42): + import pandas as pd + + np.random.seed(seed) + units = np.repeat(np.arange(n_units), n_periods) + times = np.tile(np.arange(n_periods), n_units) + n_never = int(n_units * 0.3) + n_treated = n_units - n_never + cohort_periods = np.linspace(3, n_periods - 2, n_cohorts).astype(int) + first_treat = np.zeros(n_units) + if n_treated > 0: + first_treat[n_never:] = cohort_periods[ + np.random.choice(len(cohort_periods), size=n_treated) + ] + fte = np.repeat(first_treat, n_periods) + unit_fe = np.repeat(np.random.randn(n_units) * 2, n_periods) + time_fe = np.tile(np.linspace(0, 1, n_periods), n_units) + post = (times >= fte) & (fte > 0) + y = unit_fe + time_fe + 2.0 * post + np.random.randn(len(units)) * 0.5 + return pd.DataFrame( + { + "unit": units, + "time": times, + "outcome": y, + "first_treat": fte.astype(int), + "treated": post.astype(int), + } + ) + + +def _dcdh_fit(): + from diff_diff.chaisemartin_dhaultfoeuille import ChaisemartinDHaultfoeuille + from diff_diff.prep_dgp import generate_reversible_did_data + + data = generate_reversible_did_data(n_groups=50, n_periods=8, pattern="joiners_only", seed=42) + return ChaisemartinDHaultfoeuille(twfe_diagnostic=False).fit( + data, outcome="outcome", group="group", time="period", treatment="treatment", L_max=1 + ) diff --git a/tests/test_v4_matrix.py b/tests/test_v4_matrix.py index 7d4907907..861385323 100644 --- a/tests/test_v4_matrix.py +++ b/tests/test_v4_matrix.py @@ -109,14 +109,29 @@ _FIELD_RE = re.compile(r"^ ([a-z_]+):\s*(.*?)\s*$") _MD_TOKEN_RE = re.compile(r"\[(M-\d{3})\]") -# Row-count floor: exactly the 75 rows shipped by the Phase 1 spec (incl. the diagnostic-family amendment). Ids are never reused and terminal -# rows are never deleted, so the ledger only grows - raise the floor when rows are added; a -# lower parse count means scanner/format drift or an illegal row deletion. -ROW_COUNT_FLOOR = 75 - -# Committed snapshot of the Phase 1 id set ("ids are never deleted or reused" contract - a -# delete-one-add-one edit keeps the count above the floor but trips this). Extend, never edit. -_INITIAL_ID_RANGES = [(1, 8), (10, 16), (20, 27), (30, 47), (50, 58), (60, 64), (70, 77), (80, 91)] +# Row-count floor: 75 rows from the Phase 1 spec + diagnostic-family amendment, +# plus the 2 Phase 2a results-contract rows (M-092/M-093) = 77. Ids are never +# reused and terminal rows are never deleted, so the ledger only grows - raise +# the floor when rows are added; a lower parse count means scanner/format drift +# or an illegal row deletion. +ROW_COUNT_FLOOR = 77 + +# Committed snapshot of the shipped id set ("ids are never deleted or reused" +# contract - a delete-one-add-one edit keeps the count above the floor but trips +# this). Extend with a NEW tuple, never edit an existing one. Ranges: +# (1,8)/(10,16)/(20,27)/(30,47)/(50,58)/(60,64)/(70,77)/(80,91) = Phase 1 + the +# diagnostic-family amendment; (92,93) = Phase 2a results-contract rows. +_INITIAL_ID_RANGES = [ + (1, 8), + (10, 16), + (20, 27), + (30, 47), + (50, 58), + (60, 64), + (70, 77), + (80, 91), + (92, 93), +] EXPECTED_INITIAL_IDS = frozenset( f"M-{n:03d}" for lo, hi in _INITIAL_ID_RANGES for n in range(lo, hi + 1) ) @@ -509,12 +524,13 @@ def test_due_rows_are_terminal(): def test_initial_ids_never_deleted(): - """The Phase 1 id set is immutable: ids are never deleted or reused (spec section 11). + """The shipped id set is immutable: ids are never deleted or reused (spec section 11). - ROW_COUNT_FLOOR alone would let a delete-one-add-one edit pass; this snapshot cannot.""" + ROW_COUNT_FLOOR alone would let a delete-one-add-one edit pass; this snapshot cannot. + Extends as rows ship (77 as of Phase 2a: Phase 1 + diagnostic-family + M-092/M-093).""" missing = sorted(EXPECTED_INITIAL_IDS - set(_ROW_IDS)) assert not missing, f"ledger rows deleted (ids are permanent): {missing}" - assert len(EXPECTED_INITIAL_IDS) == 75 + assert len(EXPECTED_INITIAL_IDS) == 77 def test_version_tuple_pads_to_three_components():