diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d29500b..1a129a3c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,55 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- **EfficientDiD post-fit `aggregate()` via a lazy recompute kit** + (v4 program 2(b) PR-3a; ledger row [M-023] flips to shimmed, new row + [M-120] claims the reserved balance_e slot, rows [M-092]/[M-093] + amended). + - `EfficientDiDResults.aggregate('event_study'/'group'/'simple', + balance_e=)` RECOMPUTES the aggregations post-fit from a retained + `AggregationKit` (per-(g,t) EIF dict + O(n_units) bookkeeping + + the post-overall `df_survey` snapshot) - no refit, no estimator + ref; post-fit results match the fit-time surfaces at 1e-14 across + plain/cluster/survey-TSL/replicate fits and every `balance_e`. + `aggregate('simple')` relays the stored overall row bit-exact + (`target="att"`, `n` = disjoint treated+control units, `df` = the + provenance-exact snapshot); `aggregate('group')` relays per-row + `df_used` provenance captured at each row's `safe_inference` call. + - `fit(aggregate=, balance_e=)` is DEPRECATED (3.9, removed in 4.0): + supplying either (any value, `None` included) emits one joint + FutureWarning; the deprecated path still computes and stores the + requested surfaces exactly as before. Unknown `aggregate` strings + keep their silent legacy no-op on the deprecated path; the post-fit + successor fails closed on unknown types (behavior improvement). + `EfficientDiD`'s balance rule is the anchor-horizon rule - the + same rule CallawaySantAnna uses. + - Bootstrapped fits (`n_bootstrap > 0`) fail closed on `aggregate()` + at every level (CS parity); the deprecated fit-time aggregation + remains the supported bootstrapped route, and its group rows now + clear the analytical `df_used` provenance under the bootstrap + override. Exact post-fit bootstrap replay is a tracked TODO row. + - The aggregation methods moved verbatim to the new + `diff_diff/efficient_did_aggregation.py` (importable by both the + estimator and the results module); `efficient_did.py` drops below + the 2000-line monitoring band. The kit always retains the EIF dict: + `store_eif` now governs only the public `influence_functions` field + (memory-contract change; a kit opt-out knob is a DEFERRED row). + - PT-Post reference provenance: the new membership-gated + `EfficientDiDResults.reference_period` property (SunAbraham rule - + never synthesized when the anchor cell was not estimated) marks the + materialized mechanical zero anchor `is_reference` in the + `aggregate('event_study')` container and corrects + `plot_event_study`'s inferred reference (previously the `-1` + fallback) on PT-Post `anticipation>0` fits. PT-All containers stay + reference-free (every row is a genuine estimate). + - honest/pretrends container admission is NOT widened to + EfficientDiD-sourced containers BY DESIGN (no joint event-study + covariance; both terminal TypeErrors now state the rejection + explicitly - see the REGISTRY EfficientDiD Note and row [M-093]). + - `hausman_pretest` no longer passes the deprecated `aggregate=` + internally (zero FutureWarnings), and the practitioner guidance for + EfficientDiD gains a post-fit aggregation step with a bootstrap + carve-out. - **StackedDiD post-fit `aggregate()` + honest/pretrends container admission** (v4 program 2(b) PR-2; ledger row [M-024] flips to shimmed, row [M-093] amended a second time). diff --git a/DEFERRED.md b/DEFERRED.md index 864a43d1..e1e002cf 100644 --- a/DEFERRED.md +++ b/DEFERRED.md @@ -26,7 +26,7 @@ provenance and AI-review deviation-documentation: a row here (or in | dCDH: survey cell-period allocator's post-period attribution is a library convention, not derived from the observation-level survey linearization. MC coverage is empirically close to nominal; a formal derivation (or covariance-aware two-cell alternative) is deferred. Documented in REGISTRY survey IF expansion Note. | `chaisemartin_dhaultfoeuille.py`, REGISTRY | #408 | Medium | | dCDH by_path: survey-aware backward-horizon (`placebo + predict_het + survey_design`) raises `NotImplementedError` (and `_compute_heterogeneity_test` warn-and-skips to forward-horizon-only heterogeneity) — the Binder TSL cell-period allocator's REGISTRY justification is tied to post-period attribution; backward horizons would put ψ_g mass on a pre-period cell. Needs the pre-period cell allocator derived. | `chaisemartin_dhaultfoeuille.py`, REGISTRY | follow-up | Medium | | **HonestDiD Δ^RM ARP confidence sets** : uses a naive FLCI instead of the paper's ARP conditional/hybrid sets (Sections 3.2.1-3.2.2). ARP infrastructure exists but the moment-inequality transformation needs calibration; CIs are conservative (valid coverage). Sub-items folded here: improved C-LF via direct optimization instead of grid search (`honest_did.py:947`); hybrid inference methods; event-study-specific bounds per post-period; simulation-based power analysis for honest bounds. | `honest_did.py` | #248 | Medium | -| **Conley `vcov_type` for IF / GMM estimators** (consolidates 8 near-identical rows). No reference implementation exists for any of these spatial-HAC × influence-function/GMM compositions; each was rejected at `__init__` with a deferral pointer here. SunAbraham + WooldridgeDiD-OLS conley are already supported (within-transform via `solve_ols`) and are not part of this row. Per estimator:
• `CallawaySantAnna` — Conley kernel × per-(g,t) IF aggregation (`staggered.py`).
• `TripleDifference` — × the 3-pairwise-DiD IF decomposition `w3·IF_3 + w2·IF_2 - w1·IF_1` (`triple_diff.py`).
• `ImputationDiD` — × Theorem-3 per-unit IF `sigma_sq = (cluster_psi_sums**2).sum()` (`imputation.py`).
• `EfficientDiD` — × per-unit EIF `_compute_se_from_eif` (`efficient_did.py`).
• `TwoStageDiD` — thread into the GMM sandwich meat `_compute_gmm_variance`; the SpilloverDiD `_compute_gmm_corrected_meat` machinery could be adapted to score `S_g = gamma_hat' c_g - X'_{2g} eps_{2g}` but two-stage-GMM × Conley has no reference (`two_stage.py`).
• `StackedDiD` — **methodology-blocked, not plumbing**: the stacked design replicates each control unit across sub-experiments, so Conley's distance matrix sees same-unit copies at distance 0 (`K(0)=1`); needs a per-stack spatial identifier (`stacked_did.py`).
• `SyntheticDiD` — uses `variance_method ∈ {bootstrap, jackknife, placebo}`, no analytical sandwich for Conley to plug into; needs an analytical-sandwich path or a spatial-block bootstrap (Politis-Romano 1994) (`synthetic_did.py`).
• Conley + survey weights / `survey_design` — score-reweighting is mechanical but the PSU×spatial-kernel interaction and replicate-weight spatial variance are non-trivial (Bertanha-Imbens 2014 covers cluster-sample, not Conley); raises `NotImplementedError` at the linalg validator (`linalg.py::_validate_vcov_args`). | (per sub-item) | follow-up · Phase 1b · Phase 5 | Low-Med | +| **Conley `vcov_type` for IF / GMM estimators** (consolidates 8 near-identical rows). No reference implementation exists for any of these spatial-HAC × influence-function/GMM compositions; each was rejected at `__init__` with a deferral pointer here. SunAbraham + WooldridgeDiD-OLS conley are already supported (within-transform via `solve_ols`) and are not part of this row. Per estimator:
• `CallawaySantAnna` — Conley kernel × per-(g,t) IF aggregation (`staggered.py`).
• `TripleDifference` — × the 3-pairwise-DiD IF decomposition `w3·IF_3 + w2·IF_2 - w1·IF_1` (`triple_diff.py`).
• `ImputationDiD` — × Theorem-3 per-unit IF `sigma_sq = (cluster_psi_sums**2).sum()` (`imputation.py`).
• `EfficientDiD` — × per-unit EIF `_compute_se_from_eif` (`efficient_did_aggregation.py`).
• `TwoStageDiD` — thread into the GMM sandwich meat `_compute_gmm_variance`; the SpilloverDiD `_compute_gmm_corrected_meat` machinery could be adapted to score `S_g = gamma_hat' c_g - X'_{2g} eps_{2g}` but two-stage-GMM × Conley has no reference (`two_stage.py`).
• `StackedDiD` — **methodology-blocked, not plumbing**: the stacked design replicates each control unit across sub-experiments, so Conley's distance matrix sees same-unit copies at distance 0 (`K(0)=1`); needs a per-stack spatial identifier (`stacked_did.py`).
• `SyntheticDiD` — uses `variance_method ∈ {bootstrap, jackknife, placebo}`, no analytical sandwich for Conley to plug into; needs an analytical-sandwich path or a spatial-block bootstrap (Politis-Romano 1994) (`synthetic_did.py`).
• Conley + survey weights / `survey_design` — score-reweighting is mechanical but the PSU×spatial-kernel interaction and replicate-weight spatial variance are non-trivial (Bertanha-Imbens 2014 covers cluster-sample, not Conley); raises `NotImplementedError` at the linalg validator (`linalg.py::_validate_vcov_args`). | (per sub-item) | follow-up · Phase 1b · Phase 5 | Low-Med | | `HeterogeneousAdoptionDiD` Phase 4.5 C still-open: (a) **replicate-weight designs** (BRR/Fay/JK1/JKn/SDR) — per-replicate weight-ratio rescaling for the OLS-on-residuals refit isn't covered by the multiplier-bootstrap composition; each linearity-family helper raises `NotImplementedError` on replicate weights. (b) **`lonely_psu='adjust'` + singleton-strata** on the Stute family — the pseudo-stratum centering transform isn't derived for the Stute CvM functional. | `had_pretests.py` | Phase 4.5 C | Low | | `HeterogeneousAdoptionDiD` mass-point `vcov_type in {hc2, hc2_bm}` raises `NotImplementedError` — OLS leverage `x_i'(X'X)^{-1}x_i` is wrong for 2SLS; needs the `x_i'(Z'X)^{-1}(...)(X'Z)^{-1}x_i` correction plus an R/Stata (`ivreg2 small robust`) parity anchor. | `had.py::_fit_mass_point_2sls` | Phase 2a | Medium | | `HeterogeneousAdoptionDiD` `trends_lin × survey_design`: per-group linear-trend slope under survey weighting is not derived from the paper. Raises `NotImplementedError` across all 3 `trends_lin` surfaces. | `had.py`, `had_pretests.py` | #389 | Low | @@ -69,6 +69,7 @@ For survey-specific limitations (`NotImplementedError` paths), see the | Issue | Location | PR | Priority | |-------|----------|----|----------| +| EfficientDiD kit retention removed the `store_eif=False` memory opt-out: the private aggregation kit always retains the per-(g,t) EIF dict since 3.9 (M-023) to power post-fit `aggregate()`; a `store_kit=False` constructor knob would restore an opt-out at the cost of post-fit aggregation on that result | `diff_diff/efficient_did.py` | 2(b) PR-3a | Low | | MMM interop: result-derived scaling for the exporters - let a fitted result produce the scoped incremental outcome + SE directly (via the post-4.0 `results.aggregate()` layer, where the estimator owns its aggregation weights/balance/survey masses), so callers need not hand-scale the ATT. v1 is explicit-in by design; this is the seamless follow-up. | `diff_diff/mmm.py`, `docs/v4-design.md` | mmm-interop | Low | | `SyntheticControl` fit-snapshot residency (`_SyntheticControlFitSnapshot`) — **investigated 2026-07-07, parked**: the snapshot ALIASES the fit's own working pivots (zero extra construction cost); the retained residency implements the documented freeze contract (post-fit mutation of estimator inputs must not change `in_space_placebo()` / `leave_one_out()` / conformal output on an already-returned results object, and `__getstate__` already excludes it from pickles). A compact array representation saves only pandas overhead (the float panel dominates); releasing residency needs new API surface (`release`/opt-out flag) or a freeze-contract change. Revisit on user demand for very large donor panels. | `synthetic_control.py`, `synthetic_control_results.py` | follow-up | Low | | Stratified survey-PSU multiplier-weight draw-tiling — **investigated 2026-07-07, parked**: the stratified generator (`generate_survey_multiplier_weights_batch`) consumes ONE sequential rng stream stratum-major (`rng.choice(size=(n_bootstrap, n_h))` per stratum, then lonely-PSU pooling), so draw-chunked assembly CANNOT reproduce the stream bit-identically (contra the old row's parenthetical) — it would need per-stratum generator state skipping (PCG64.advance + per-weight-type variate accounting; fragile) or a stream-layout change (MC-level SE changes → baseline/golden recapture + REGISTRY note). Stratified designs have few PSUs, so the full `(n_bootstrap × n_psu)` matrix rarely matters; unstratified (the large-`n_units` case) is already tiled. Revisit only if a large-PSU stratified design hits memory, as a documented stream change. | `diff_diff/bootstrap_chunking.py::iter_survey_multiplier_weight_blocks` | follow-up | Low | diff --git a/README.md b/README.md index c42568f7..67c8cc48 100644 --- a/README.md +++ b/README.md @@ -91,7 +91,7 @@ For rigorous DiD analysis, follow these 8 steps. Skipping diagnostic steps produ 4. **Choose estimator** - staggered adoption -> CS/SA/BJS (NOT plain TWFE); few treated units -> SDiD; factor confounding -> TROP; simple 2x2 -> DiD. Run `BaconDecomposition` to diagnose TWFE bias. 5. **Estimate** - `estimator.fit(data, ...)`. Always print the cluster count first and choose inference method based on the result (cluster-robust if >= 50 clusters, wild bootstrap if fewer). 6. **Sensitivity analysis** - `compute_honest_did(results)` for bounds under PT violations (MultiPeriodDiD, CS, or dCDH natively; a StackedDiD `results.aggregate('event_study')` container also admits - needs `kappa_pre >= 2`), `run_all_placebo_tests()` for 2x2 falsification, specification comparisons for staggered designs. -7. **Heterogeneity** - CS: `results.aggregate('group')`/`'event_study'` (post-fit, no refit); SA: `results.event_study_effects` / `to_dataframe(level='cohort')`; Stacked: `results.aggregate('event_study')`/`'simple'` post-fit views (surface always computed since 3.9); subgroup re-estimation. +7. **Heterogeneity** - CS: `results.aggregate('group')`/`'event_study'` (post-fit, no refit); SA: `results.event_study_effects` / `to_dataframe(level='cohort')`; Stacked: `results.aggregate('event_study')`/`'simple'` post-fit views (surface always computed since 3.9); EDiD: `results.aggregate(...)` post-fit from retained EIFs (3.9); subgroup re-estimation. 8. **Robustness** - compare 2-3 estimators (CS vs SA vs BJS), report with and without covariates (shows whether conditioning drives identification), present pre-trends and sensitivity bounds. Full guide: `diff_diff.get_llm_guide("practitioner")`. diff --git a/TODO.md b/TODO.md index 5d74bb32..445133bb 100644 --- a/TODO.md +++ b/TODO.md @@ -21,6 +21,8 @@ Related tracking surfaces: | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| +| EfficientDiD `aggregate()` on bootstrapped fits fails closed (M-023); wiring `BootstrapReplaySpec` (or retaining the n_bootstrap x n_gt draw matrix materialized at fit) would enable exact post-fit replay of percentile inference | `diff_diff/efficient_did_results.py`, `diff_diff/aggregation.py` | 2(b) PR-3a | Mid | Low | +| EfficientDiD is the outstanding M-092 event-study df-provenance hole: the container's per-row df is all-NaN even on survey fits where a finite `_survey_df` governed the p-values (the container-level scalar `df_survey` IS exposed - the hole is the PER-ROW column only; no event_study_df/df_inference field; pre-existing, NOT a regression of the M-023 PR - today's builder output is identical). The kit now retains the scalar; threading it into the per-row channel is a contained follow-up | `diff_diff/efficient_did_results.py`, `diff_diff/results_base.py` | 2(b) PR-3a | Quick | Low | | practitioner `step_name="heterogeneity"` producer-side collisions: three OTHER estimators' advice steps reuse the key with non-heterogeneity labels (`:975` ContinuousDiD dose-response, `:1022` Triple placebo-group, `:1413` LPDiD WAS arrays), so DiagnosticReport's heterogeneity completion silently drops that unrelated advice from `next_steps` via `_filter_steps` - the same latent collision fixed for StackedDiD in M-024 (renamed to `sub_experiment_balance`). Renaming these changes those estimators' report output; audit + rename with per-estimator pins. | `diff_diff/practitioner.py` | 2(b) PR-2 review R9 | Quick | Low | | PreTrendsPower `violation='linear'` on CS `base_period='varying'` input targets the wrong alternative: `δ_pre = M · \|t\|` assumes level coefficients against a common reference, but varying-base pre-treatment effects are consecutive-period comparisons (constant increments under a linear trend). Both CS-sourced routes now WARN (REGISTRY PreTrendsPower Note), and universal-base GAPPED grids fail closed via the `reference_event_times` common-reference guard; what remains is the varying-base resolution - either transforming the violation vector through each coefficient's actual base mapping (needs per-horizon base provenance) or requiring `base_period='universal'` for the linear benchmark - a per-estimator methodology decision with a hand-calculated linear-violation gate | `diff_diff/pretrends.py` | 2(b) PR-1 R5 | Mid | Medium | | `run_placebo_test`/`run_all_placebo_tests` expose ONE overloaded `time` param (calendar for timing/group tests, 0/1 post dummy forwarded into the M-137/M-138 callees - works only because a two-period 0/1 calendar column is both); redesign the wrapper signatures (dual `time=`/`post=` params or a split) so each callee gets an honestly-named argument | `diff_diff/diagnostics.py` | 2(c)-ii PR-B | Mid | Medium | diff --git a/diff_diff/efficient_did.py b/diff_diff/efficient_did.py index 7d83c21e..9ee259fa 100644 --- a/diff_diff/efficient_did.py +++ b/diff_diff/efficient_did.py @@ -29,6 +29,13 @@ import pandas as pd from diff_diff._base import BaseEstimator +from diff_diff._deprecation import NOT_SUPPLIED +from diff_diff.aggregation import AggregationKit +from diff_diff.efficient_did_aggregation import ( + _cluster_aggregate, + _compute_se_from_eif, + _EfficientAggregationMixin, +) from diff_diff.efficient_did_bootstrap import ( EDiDBootstrapResults, EfficientDiDBootstrapMixin, @@ -94,57 +101,72 @@ def _validate_and_build_cluster_mapping( return indices, n_clusters -def _cluster_aggregate( - eif_mat: np.ndarray, - cluster_indices: np.ndarray, - n_clusters: int, -) -> np.ndarray: - """Sum EIF values within clusters and center. - - Parameters - ---------- - eif_mat : ndarray, shape (n_units,) or (n_units, k) - EIF values — 1-D for a single estimand, 2-D for multiple. - cluster_indices : ndarray, shape (n_units,) - Integer cluster assignment per unit. - n_clusters : int - Number of unique clusters. - - Returns - ------- - ndarray, shape (n_clusters,) or (n_clusters, k) - Centered cluster-level sums. - """ - if eif_mat.ndim == 1: - sums = np.bincount(cluster_indices, weights=eif_mat, minlength=n_clusters).astype(float) - else: - sums = np.column_stack( - [ - np.bincount(cluster_indices, weights=eif_mat[:, j], minlength=n_clusters) - for j in range(eif_mat.shape[1]) - ] - ).astype(float) - return sums - sums.mean(axis=0) - - -def _compute_se_from_eif( - eif: np.ndarray, +def _build_edid_aggregation_kit( + eif_by_gt: Dict[Tuple[Any, Any], np.ndarray], + group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]], + treatment_groups: List[Any], + time_periods: List[Any], + pt_assumption: str, + n_units_total: int, n_units: int, - cluster_indices: Optional[np.ndarray] = None, - n_clusters: Optional[int] = None, -) -> float: - """SE from EIF values, optionally with cluster-robust correction. - - Without clusters: ``sqrt(mean(EIF^2) / n)``. - With clusters: Liang-Zeger sandwich — aggregate EIF within clusters, - center, and apply G/(G-1) small-sample correction. + cohort_fractions: Dict[float, float], + unit_cohorts: np.ndarray, + cluster_indices: Optional[np.ndarray], + n_clusters: Optional[int], + unit_level_weights: Optional[np.ndarray], + resolved_survey_unit: Optional["ResolvedSurveyDesign"], + df_survey: Optional[float], + alpha: float, + anticipation: int, +) -> Optional[AggregationKit]: + """Bundle the retained EIF payload + bookkeeping for post-fit aggregate(). + + Reference bundling only — zero array copies, zero extra computation, so + plain-fit timing is unchanged; the dominant memory payload is the + per-(g,t) EIF dict, O(n_units x n_gt) (full enumeration in the + ``docs/methodology/REGISTRY.md`` EfficientDiD M-023 Note). + + ``df_survey`` MUST be the post-overall snapshot captured in ``fit()`` + immediately after the overall inference and before the ES/group gates — + NOT a read of ``estimator._survey_df`` at kit-build time, which on a + deprecated ``fit(aggregate="group"/"all")`` replicate fit is the + post-group value (``_compute_survey_eif_se`` can set it to ``None`` on + a degenerate replicate design). The snapshot is the exact value every + fit-time aggregation seeded from, so post-fit recompute replays it. """ - if cluster_indices is not None and n_clusters is not None: - centered = _cluster_aggregate(eif, cluster_indices, n_clusters) - correction = n_clusters / (n_clusters - 1) if n_clusters > 1 else 1.0 - var = correction * np.sum(centered**2) / (n_units**2) - return float(np.sqrt(max(var, 0.0))) - return float(np.sqrt(np.mean(eif**2) / n_units)) + if not eif_by_gt: + # Unreachable after fit()'s empty-effects raise; kept for the CS + # guard shape (a kit with nothing to re-aggregate is not attached). + return None + return AggregationKit( + bookkeeping={ + # PRIVATE SNAPSHOTS of the aggregation inputs (CI review P0): + # aggregate() must never read the MUTABLE public result fields - + # a user edit of results.group_time_effects/groups would + # otherwise mix altered point estimates with the retained + # fit-time EIF variance, yielding plausible-but-invalid + # inference. Row values are scalars/tuples, so per-row dict + # copies suffice. alpha/anticipation are already kit fields. + "group_time_effects": {gt: dict(row) for gt, row in group_time_effects.items()}, + "treatment_groups": list(treatment_groups), + "time_periods": list(time_periods), + "pt_assumption": pt_assumption, + "n_units_total": n_units_total, + "n_units": n_units, + "cohort_fractions": cohort_fractions, + "unit_cohorts": unit_cohorts, + "cluster_indices": cluster_indices, + "n_clusters": n_clusters, + "unit_level_weights": unit_level_weights, + "resolved_survey_unit": resolved_survey_unit, + "df_survey": df_survey, + }, + influence=eif_by_gt, + alpha=alpha, + anticipation=anticipation, + cband=False, + bootstrap=None, + ) def _hausman_quadratic_form( @@ -210,7 +232,7 @@ def _hausman_quadratic_form( return H, effective_rank, p_value, n_negative, True -class EfficientDiD(EfficientDiDBootstrapMixin, BaseEstimator): +class EfficientDiD(EfficientDiDBootstrapMixin, _EfficientAggregationMixin, BaseEstimator): """Efficient DiD estimator (Chen, Sant'Anna & Xie 2025). Without covariates, achieves the semiparametric efficiency bound for @@ -310,8 +332,9 @@ class EfficientDiD(EfficientDiDBootstrapMixin, BaseEstimator): >>> from diff_diff import EfficientDiD >>> edid = EfficientDiD(pt_assumption="all") >>> results = edid.fit(data, outcome="y", unit="id", time="t", - ... first_treat="first_treat", aggregate="all") + ... first_treat="first_treat") >>> results.print_summary() + >>> results.aggregate("event_study").summary() # post-fit aggregation """ def __init__( @@ -447,8 +470,8 @@ def fit( time: str, first_treat: str, covariates: Optional[List[str]] = None, - aggregate: Optional[str] = None, - balance_e: Optional[int] = None, + aggregate: Any = NOT_SUPPLIED, + balance_e: Any = NOT_SUPPLIED, survey_design: Optional[Any] = None, store_eif: bool = False, ) -> EfficientDiDResults: @@ -472,19 +495,37 @@ def fit( When provided, uses the doubly robust path (outcome regression + propensity score ratios). aggregate : str, optional - ``None``, ``"simple"``, ``"event_study"``, ``"group"``, or - ``"all"``. + DEPRECATED (3.9, removed in 4.0, row M-023): supplying any + value (``None`` included) emits a FutureWarning; the + deprecated path still computes and stores the requested + surfaces exactly as before (``None``, ``"simple"``, + ``"event_study"``, ``"group"``, or ``"all"``; unknown strings + act like ``None``, unchanged). Aggregate as a post-fit step + instead: ``results.aggregate('event_study')`` / + ``.aggregate('group')`` / ``.aggregate('simple')``. On + bootstrapped fits (``n_bootstrap > 0``) post-fit + ``aggregate()`` fails closed — the deprecated fit-time path + remains the supported route for a bootstrapped aggregated + surface. balance_e : int, optional - Balance event study at this relative period. + DEPRECATED (3.9, removed in 4.0, row M-120): moves onto + post-fit ``aggregate()`` — + ``results.aggregate('event_study', balance_e=2)``. EDiD's + balance rule is the ANCHOR-HORIZON rule (keep cohorts with a + finite effect at ``e == balance_e``), the same rule + CallawaySantAnna uses. survey_design : SurveyDesign, optional Survey design specification for design-based inference. Applies survey weights to all means, covariances, and cohort fractions, and uses Taylor Series Linearization for SE estimation. Cannot be combined with ``cluster``. store_eif : bool, default False - Store per-(g,t) EIF vectors in the results object. Used - internally by :meth:`hausman_pretest`; not needed for - normal usage. + Expose per-(g,t) EIF vectors on the PUBLIC + ``influence_functions`` results field. Used internally by + :meth:`hausman_pretest`. Since 3.9 (row M-023) the private + aggregation kit ALWAYS retains the per-(g,t) EIF dict to power + post-fit ``results.aggregate()`` — ``store_eif`` governs only + the public field, no longer the retention itself. Returns ------- @@ -496,6 +537,35 @@ def fit( Missing columns, unbalanced panel, non-absorbing treatment, or PT-Post without a never-treated group. """ + # M-023/M-120 deprecation shim (CS-style joint warning): a plain + # fit() never warns; supplying EITHER param with ANY value (None + # included) warns once, then the legacy routing below runs + # unchanged - the deprecated path returns exactly the numbers it + # always did (no new value validation; unknown strings still act + # like None). The post-fit successor validates its own vocabulary. + _deprecated_passed = [ + n + for n, v in (("aggregate", aggregate), ("balance_e", balance_e)) + if v is not NOT_SUPPLIED + ] + if _deprecated_passed: + _args = " / ".join(f"{n}=" for n in _deprecated_passed) + warnings.warn( + f"EfficientDiD.fit({_args}) is deprecated and will be " + "removed in 4.0. Fit once, then aggregate as a post-fit " + "step: results = EfficientDiD().fit(...); " + "results.aggregate('event_study') / .aggregate('group') / " + ".aggregate('simple'). balance_e moves onto aggregate() " + "alongside it: results.aggregate('event_study', " + "balance_e=2).", + FutureWarning, + stacklevel=2, + ) + if aggregate is NOT_SUPPLIED: + aggregate = None + if balance_e is NOT_SUPPLIED: + balance_e = None + self._validate_params() if self.cluster is not None and survey_design is not None: @@ -1200,6 +1270,15 @@ def _finalize_cell(g: Any, att_gt: float, eif_vals: np.ndarray) -> Dict[str, Any overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=self._survey_df ) + # M-023 kit snapshot: the post-overall ``_survey_df`` is the exact + # value fit-time ES/group aggregation seeds from (the gates below run + # next), and post-fit aggregate() must replay from the same seed. + # Captured HERE, not at fit end: ``_aggregate_by_group`` can mutate + # ``self._survey_df`` through ``_compute_survey_eif_se`` on degenerate + # replicate designs, so a fit-end read on a deprecated + # ``fit(aggregate="group"/"all")`` would seed recompute from the + # post-group value and break post-fit/fit-time inertness. + _survey_df_post_overall = self._survey_df event_study_effects = None group_effects = None @@ -1307,6 +1386,11 @@ def _finalize_cell(g: Any, att_gt: float, eif_vals: np.ndarray) -> Dict[str, Any eff = float(group_effects[g]["effect"]) se = float(group_effects[g]["se"]) group_effects[g]["t_stat"] = safe_inference(eff, se, alpha=self.alpha)[0] + # Percentile-bootstrap inference has no analytical df; + # clear the provenance key the analytical pass recorded + # (the CS precedent) so bootstrap rows never publish an + # analytical survey df beside percentile p/CI. + group_effects[g]["df_used"] = None # ----- Build results ----- self.results_ = EfficientDiDResults( @@ -1348,7 +1432,14 @@ def _finalize_cell(g: Any, att_gt: float, eif_vals: np.ndarray) -> Dict[str, Any else (n_clusters if self.cluster is not None else None) ), vcov_type=self.vcov_type, - influence_functions=eif_by_gt if store_eif else None, + # Independent COPY, never the kit's arrays: the aggregation kit + # below retains eif_by_gt by reference as the canonical payload + # for post-fit aggregate(), so an aliased public diagnostic + # would let a user mutation silently corrupt recomputed + # SEs/p-values/CIs (M-023 review pin). + influence_functions=( + {gt: arr.copy() for gt, arr in eif_by_gt.items()} if store_eif else None + ), bootstrap_results=bootstrap_results, estimation_path="dr" if use_covariates else "nocov", sieve_k_max=self.sieve_k_max, @@ -1362,6 +1453,28 @@ def _finalize_cell(g: Any, att_gt: float, eif_vals: np.ndarray) -> Dict[str, Any else None ), ) + # Attach the post-fit aggregation kit (M-023). Built HERE because + # the EIF dict and the unit-level bookkeeping are fit() locals that + # do not otherwise survive the call; ``results.aggregate()`` + # recomputes from this retained payload without an estimator ref. + self.results_._aggregation_kit = _build_edid_aggregation_kit( + eif_by_gt=eif_by_gt, + group_time_effects=group_time_effects, + treatment_groups=treatment_groups, + time_periods=time_periods, + pt_assumption=self.pt_assumption, + n_units_total=n_treated_units + n_control_units, + n_units=n_units, + cohort_fractions=cohort_fractions, + unit_cohorts=unit_cohorts, + cluster_indices=unit_cluster_indices, + n_clusters=n_clusters, + unit_level_weights=self._unit_level_weights, + resolved_survey_unit=self._unit_resolved_survey, + df_survey=_survey_df_post_overall, + alpha=self.alpha, + anticipation=self.anticipation, + ) self.is_fitted_ = True return self.results_ @@ -1385,400 +1498,6 @@ def _recompute_unit_survey_metadata(self, panel_metadata): return meta return panel_metadata - # -- Survey SE helpers ---------------------------------------------------- - - def _compute_survey_eif_se(self, eif_vals: np.ndarray) -> float: - """Compute SE from EIF scores using Taylor Series Linearization. - - Uses the pre-built unit-level ``_unit_resolved_survey`` constructed - once in ``fit()``, ensuring consistent unit-level arrays and - avoiding repeated subsetting of panel-level survey data. - """ - # Built once in fit() before any call lands here (see docstring). - assert self._unit_resolved_survey is not None - if self._unit_resolved_survey.uses_replicate_variance: - from diff_diff.survey import compute_replicate_if_variance - - # Score-scale IFs to match TSL bread: psi = w * eif / sum(w) - w = self._unit_resolved_survey.weights - psi_scaled = w * eif_vals / w.sum() - variance, n_valid = compute_replicate_if_variance( - psi_scaled, self._unit_resolved_survey - ) - # Update survey df to reflect effective replicate count - if n_valid < self._unit_resolved_survey.n_replicates: - self._survey_df = n_valid - 1 if n_valid > 1 else None - return float(np.sqrt(max(variance, 0.0))) if np.isfinite(variance) else np.nan - - from diff_diff.survey import compute_survey_vcov - - X_ones = np.ones((len(eif_vals), 1)) - vcov = compute_survey_vcov(X_ones, eif_vals, self._unit_resolved_survey) - return float(np.sqrt(np.abs(vcov[0, 0]))) - - def _eif_se( - self, - eif_vals: np.ndarray, - n_units: int, - cluster_indices: Optional[np.ndarray] = None, - n_clusters: Optional[int] = None, - ) -> float: - """Compute SE from aggregated EIF scores. - - Dispatches to survey TSL when ``_unit_resolved_survey`` is set - (during fit), otherwise uses cluster-robust or standard formula. - """ - if self._unit_resolved_survey is not None: - return self._compute_survey_eif_se(eif_vals) - return _compute_se_from_eif(eif_vals, n_units, cluster_indices, n_clusters) - - # -- Aggregation helpers -------------------------------------------------- - - def _compute_wif_contribution( - self, - keepers: List[Tuple], - effects: np.ndarray, - unit_cohorts: np.ndarray, - cohort_fractions: Dict[float, float], - n_units: int, - unit_weights: Optional[np.ndarray] = None, - ) -> np.ndarray: - """Compute weight influence function correction (O(1) scale, matching EIF). - - This accounts for uncertainty in cohort-size aggregation weights. - Matches R's ``did`` package WIF formula (staggered_aggregation.py:282-309), - adapted to EDiD's EIF scale. - - Parameters - ---------- - keepers : list of (g, t) tuples - Post-treatment group-time pairs included in aggregation. - effects : ndarray, shape (n_keepers,) - ATT estimates for each keeper. - unit_cohorts : ndarray, shape (n_units,) - Cohort assignment for each unit (0 = never-treated). - cohort_fractions : dict - ``{cohort: n_cohort / n}`` for each cohort. - n_units : int - Total number of units. - unit_weights : ndarray, shape (n_units,), optional - Survey weights at the unit level. When provided, uses the - survey-weighted WIF formula: IF_i(p_g) = (w_i * 1{G_i=g} - pg_k). - - Returns - ------- - ndarray, shape (n_units,) - WIF contribution at O(1) scale, additive with ``agg_eif``. - """ - groups_for_keepers = np.array([g for (g, t) in keepers]) - pg_keepers = np.array([cohort_fractions.get(g, 0.0) for g, t in keepers]) - sum_pg = pg_keepers.sum() - if sum_pg == 0: - return np.zeros(n_units) - - indicator = (unit_cohorts[:, None] == groups_for_keepers[None, :]).astype(float) - - if unit_weights is not None: - # Survey-weighted WIF (matches staggered_aggregation.py:392-401): - # IF_i(p_g) = (w_i * 1{G_i=g} - pg_k), NOT (1{G_i=g} - pg_k) - weighted_indicator = indicator * unit_weights[:, None] - indicator_diff = weighted_indicator - pg_keepers - indicator_sum = np.sum(indicator_diff, axis=1) - else: - indicator_diff = indicator - pg_keepers - indicator_sum = np.sum(indicator_diff, axis=1) - - with np.errstate(divide="ignore", invalid="ignore", over="ignore"): - if1 = indicator_diff / sum_pg - if2 = np.outer(indicator_sum, pg_keepers) / sum_pg**2 - wif_matrix = if1 - if2 - wif_contrib = wif_matrix @ effects - return wif_contrib # O(1) scale, same as agg_eif - - def _aggregate_overall( - self, - group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]], - eif_by_gt: Dict[Tuple[Any, Any], np.ndarray], - n_units: int, - cohort_fractions: Dict[float, float], - unit_cohorts: np.ndarray, - cluster_indices: Optional[np.ndarray] = None, - n_clusters: Optional[int] = None, - ) -> Tuple[float, float]: - """Compute overall ATT with WIF-adjusted SE. - - Parameters - ---------- - group_time_effects : dict - Group-time ATT estimates. - eif_by_gt : dict - Per-unit EIF values for each (g, t). - n_units : int - Total number of units. - cohort_fractions : dict - Cohort size fractions. - unit_cohorts : ndarray, shape (n_units,) - Cohort assignment for each unit. - """ - # Filter to post-treatment effects - keepers = [ - (g, t) - for (g, t) in group_time_effects - if t >= g - self.anticipation and np.isfinite(group_time_effects[(g, t)]["effect"]) - ] - if not keepers: - return np.nan, np.nan - - # Cohort-size weights - pg = np.array([cohort_fractions.get(g, 0.0) for (g, _) in keepers]) - total_pg = pg.sum() - if total_pg == 0: - return np.nan, np.nan - w = pg / total_pg - - effects = np.array([group_time_effects[gt]["effect"] for gt in keepers]) - overall_att = float(np.sum(w * effects)) - - # Aggregate EIF - agg_eif = np.zeros(n_units) - for k, gt in enumerate(keepers): - agg_eif += w[k] * eif_by_gt[gt] - - # WIF correction: accounts for uncertainty in cohort-size weights - wif = self._compute_wif_contribution( - keepers, - effects, - unit_cohorts, - cohort_fractions, - n_units, - unit_weights=self._unit_level_weights, - ) - # Compute SE: survey path uses score-level psi to avoid double-weighting - # (compute_survey_vcov applies w_i internally, which would double-weight - # the survey-weighted WIF term). Dispatch replicate vs TSL. - if self._unit_resolved_survey is not None: - uw = self._unit_level_weights - # Set together with _unit_resolved_survey in fit(). - assert uw is not None - total_w = float(np.sum(uw)) - psi_total = uw * agg_eif / total_w + wif / total_w - - if ( - hasattr(self._unit_resolved_survey, "uses_replicate_variance") - and self._unit_resolved_survey.uses_replicate_variance - ): - from diff_diff.survey import compute_replicate_if_variance - - variance, _ = compute_replicate_if_variance(psi_total, self._unit_resolved_survey) - else: - from diff_diff.survey import compute_survey_if_variance - - variance = compute_survey_if_variance(psi_total, self._unit_resolved_survey) - se = float(np.sqrt(max(variance, 0.0))) if np.isfinite(variance) else np.nan - else: - agg_eif_total = agg_eif + wif - se = self._eif_se(agg_eif_total, n_units, cluster_indices, n_clusters) - - return overall_att, se - - def _aggregate_event_study( - self, - group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]], - eif_by_gt: Dict[Tuple[Any, Any], np.ndarray], - n_units: int, - cohort_fractions: Dict[float, float], - treatment_groups: List[Any], - time_periods: List[Any], - balance_e: Optional[int] = None, - unit_cohorts: Optional[np.ndarray] = None, - cluster_indices: Optional[np.ndarray] = None, - n_clusters: Optional[int] = None, - ) -> Dict[int, Dict[str, Any]]: - """Aggregate ATT(g,t) by relative time e = t - g. - - Parameters - ---------- - group_time_effects : dict - Group-time ATT estimates. - eif_by_gt : dict - Per-unit EIF values for each (g, t). - n_units : int - Total number of units. - cohort_fractions : dict - Cohort size fractions. - treatment_groups : list - Treatment cohort identifiers. - time_periods : list - All time periods. - balance_e : int, optional - Balance event study at this relative period. - unit_cohorts : ndarray, optional - Cohort assignment for each unit (for WIF correction). - """ - # Organize by relative time - effects_by_e: Dict[int, List[Tuple[Tuple[Any, Any], float, float]]] = {} - for (g, t), data in group_time_effects.items(): - if not np.isfinite(data["effect"]): - continue - e = int(t - g) - if e not in effects_by_e: - effects_by_e[e] = [] - effects_by_e[e].append(((g, t), data["effect"], cohort_fractions.get(g, 0.0))) - - # Balance if requested - if balance_e is not None: - groups_at_e = {gt[0] for gt, _, _ in effects_by_e.get(balance_e, [])} - balanced: Dict[int, List[Tuple[Tuple[Any, Any], float, float]]] = {} - for (g, t), data in group_time_effects.items(): - if not np.isfinite(data["effect"]): - continue - if g in groups_at_e: - e = int(t - g) - if e not in balanced: - balanced[e] = [] - balanced[e].append(((g, t), data["effect"], cohort_fractions.get(g, 0.0))) - effects_by_e = balanced - - if balance_e is not None and not effects_by_e: - warnings.warn( - f"balance_e={balance_e}: no cohort has a finite effect at the " - "anchor horizon. Event study will be empty.", - UserWarning, - stacklevel=2, - ) - - result: Dict[int, Dict[str, Any]] = {} - for e, elist in sorted(effects_by_e.items()): - gt_pairs = [x[0] for x in elist] - effs = np.array([x[1] for x in elist]) - pgs = np.array([x[2] for x in elist]) - total_pg = pgs.sum() - w = pgs / total_pg if total_pg > 0 else np.ones(len(pgs)) / len(pgs) - - agg_eff = float(np.sum(w * effs)) - - # Aggregate EIF - agg_eif = np.zeros(n_units) - for k, gt in enumerate(gt_pairs): - agg_eif += w[k] * eif_by_gt[gt] - - # WIF correction for event-study aggregation - wif_e = np.zeros(n_units) - if unit_cohorts is not None: - es_keepers = [(g, t) for (g, t) in gt_pairs] - es_effects = effs - wif_e = self._compute_wif_contribution( - es_keepers, - es_effects, - unit_cohorts, - cohort_fractions, - n_units, - unit_weights=self._unit_level_weights, - ) - - if self._unit_resolved_survey is not None: - uw = self._unit_level_weights - # Set together with _unit_resolved_survey in fit(). - assert uw is not None - total_w = float(np.sum(uw)) - psi_total = uw * agg_eif / total_w + wif_e / total_w - - if ( - hasattr(self._unit_resolved_survey, "uses_replicate_variance") - and self._unit_resolved_survey.uses_replicate_variance - ): - from diff_diff.survey import compute_replicate_if_variance - - variance, _ = compute_replicate_if_variance( - psi_total, self._unit_resolved_survey - ) - else: - from diff_diff.survey import compute_survey_if_variance - - variance = compute_survey_if_variance(psi_total, self._unit_resolved_survey) - agg_se = float(np.sqrt(max(variance, 0.0))) if np.isfinite(variance) else np.nan - else: - agg_eif = agg_eif + wif_e - agg_se = self._eif_se(agg_eif, n_units, cluster_indices, n_clusters) - - t_stat, p_val, ci = safe_inference( - agg_eff, agg_se, alpha=self.alpha, df=self._survey_df - ) - result[e] = { - "effect": agg_eff, - "se": agg_se, - "t_stat": t_stat, - "p_value": p_val, - "conf_int": ci, - "n_groups": len(elist), - } - - return result - - def _aggregate_by_group( - self, - group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]], - eif_by_gt: Dict[Tuple[Any, Any], np.ndarray], - n_units: int, - cohort_fractions: Dict[float, float], - treatment_groups: List[Any], - unit_cohorts: Optional[np.ndarray] = None, - cluster_indices: Optional[np.ndarray] = None, - n_clusters: Optional[int] = None, - ) -> Dict[Any, Dict[str, Any]]: - """Aggregate ATT(g,t) by treatment cohort. - - Parameters - ---------- - group_time_effects : dict - Group-time ATT estimates. - eif_by_gt : dict - Per-unit EIF values for each (g, t). - n_units : int - Total number of units. - cohort_fractions : dict - Cohort size fractions. - treatment_groups : list - Treatment cohort identifiers. - unit_cohorts : ndarray, optional - Cohort assignment for each unit (unused — group aggregation - uses equal weights, not cohort-size weights). - """ - result: Dict[Any, Dict[str, Any]] = {} - for g in treatment_groups: - g_gts = [ - (gg, t) - for (gg, t) in group_time_effects - if gg == g - and t >= g - self.anticipation - and np.isfinite(group_time_effects[(gg, t)]["effect"]) - ] - if not g_gts: - continue - - effs = np.array([group_time_effects[gt]["effect"] for gt in g_gts]) - w = np.ones(len(effs)) / len(effs) - agg_eff = float(np.sum(w * effs)) - - agg_eif = np.zeros(n_units) - for k, gt in enumerate(g_gts): - agg_eif += w[k] * eif_by_gt[gt] - agg_se = self._eif_se(agg_eif, n_units, cluster_indices, n_clusters) - - t_stat, p_val, ci = safe_inference( - agg_eff, agg_se, alpha=self.alpha, df=self._survey_df - ) - result[g] = { - "effect": agg_eff, - "se": agg_se, - "t_stat": t_stat, - "p_value": p_val, - "conf_int": ci, - "n_periods": len(g_gts), - } - - return result - def summary(self) -> str: """Get summary of estimation results.""" if not self.is_fitted_: @@ -1849,7 +1568,6 @@ def hausman_pretest( time=time, first_treat=first_treat, covariates=covariates, - aggregate=None, ) edid_all = cls(pt_assumption="all", alpha=alpha, **common_kwargs) diff --git a/diff_diff/efficient_did_aggregation.py b/diff_diff/efficient_did_aggregation.py new file mode 100644 index 00000000..54a6015c --- /dev/null +++ b/diff_diff/efficient_did_aggregation.py @@ -0,0 +1,521 @@ +""" +Aggregation mixin for the EfficientDiD estimator (Chen, Sant'Anna & Xie 2025). + +Extracted from ``diff_diff/efficient_did.py`` with the M-023 post-fit +``aggregate()`` migration so the aggregation methods are importable by BOTH +the estimator (``efficient_did.py``) and the results module +(``efficient_did_results.py``) without an import cycle — ``efficient_did.py`` +imports ``efficient_did_results.py``, so the results module can never import +the estimator. This mirrors the CallawaySantAnna layout +(``staggered_aggregation.py``). + +Contents: + +- ``_cluster_aggregate`` / ``_compute_se_from_eif`` — module-level variance + helpers (moved verbatim; ``efficient_did.py`` re-imports both). +- ``_EfficientAggregationMixin`` — the six estimator methods that compute + the overall / event-study / group aggregations from the per-(g,t) EIF + dict plus O(n_units) bookkeeping. ``EfficientDiD`` inherits it for the + fit-time path, and ``EfficientDiDResults._aggregate_compute`` runs the + same methods post-fit on a throwaway ``_EDiDKitAggregator`` host built + from the retained :class:`~diff_diff.aggregation.AggregationKit` — which + is what keeps ``aggregate()`` off an ``_estimator_ref``. + +The numerical content of every function in this module is byte-identical to +its pre-extraction form, with ONE additive exception recorded in the M-023 +ledger notes: ``_aggregate_by_group`` records a per-row ``df_used`` key (the +``self._survey_df`` value at that row's ``safe_inference`` call) so the +post-fit group relay can publish exact per-row df provenance. +""" + +import warnings +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +import numpy as np + +from diff_diff.utils import safe_inference + +if TYPE_CHECKING: + from diff_diff.survey import ResolvedSurveyDesign + + +def _cluster_aggregate( + eif_mat: np.ndarray, + cluster_indices: np.ndarray, + n_clusters: int, +) -> np.ndarray: + """Sum EIF values within clusters and center. + + Parameters + ---------- + eif_mat : ndarray, shape (n_units,) or (n_units, k) + EIF values — 1-D for a single estimand, 2-D for multiple. + cluster_indices : ndarray, shape (n_units,) + Integer cluster assignment per unit. + n_clusters : int + Number of unique clusters. + + Returns + ------- + ndarray, shape (n_clusters,) or (n_clusters, k) + Centered cluster-level sums. + """ + if eif_mat.ndim == 1: + sums = np.bincount(cluster_indices, weights=eif_mat, minlength=n_clusters).astype(float) + else: + sums = np.column_stack( + [ + np.bincount(cluster_indices, weights=eif_mat[:, j], minlength=n_clusters) + for j in range(eif_mat.shape[1]) + ] + ).astype(float) + return sums - sums.mean(axis=0) + + +def _compute_se_from_eif( + eif: np.ndarray, + n_units: int, + cluster_indices: Optional[np.ndarray] = None, + n_clusters: Optional[int] = None, +) -> float: + """SE from EIF values, optionally with cluster-robust correction. + + Without clusters: ``sqrt(mean(EIF^2) / n)``. + With clusters: Liang-Zeger sandwich — aggregate EIF within clusters, + center, and apply G/(G-1) small-sample correction. + """ + if cluster_indices is not None and n_clusters is not None: + centered = _cluster_aggregate(eif, cluster_indices, n_clusters) + correction = n_clusters / (n_clusters - 1) if n_clusters > 1 else 1.0 + var = correction * np.sum(centered**2) / (n_units**2) + return float(np.sqrt(max(var, 0.0))) + return float(np.sqrt(np.mean(eif**2) / n_units)) + + +class _EfficientAggregationMixin: + """EIF-based aggregation methods shared by fit-time and post-fit paths. + + Not intended for standalone use. A host class must expose exactly the + five attributes declared below — this is the ``_EDiDKitAggregator`` + contract (``efficient_did_results.py``): the post-fit path constructs a + fresh throwaway host per ``aggregate()`` call, so the ONE mutation these + methods perform (``_compute_survey_eif_se`` writes ``self._survey_df`` + when a degenerate replicate design drops replicates) lands on the + throwaway, never on the retained kit — preserving the aggregate() + immutability contract. + """ + + # Typed host-attribute contract (mypy attr-defined; the + # staggered_aggregation.py precedent). Values are read-only here + # except _survey_df (see class docstring). + alpha: float + anticipation: int + _survey_df: Optional[float] + _unit_resolved_survey: Optional["ResolvedSurveyDesign"] + _unit_level_weights: Optional[np.ndarray] + + # -- Survey SE helpers ---------------------------------------------------- + + def _compute_survey_eif_se(self, eif_vals: np.ndarray) -> float: + """Compute SE from EIF scores using Taylor Series Linearization. + + Uses the pre-built unit-level ``_unit_resolved_survey`` constructed + once in ``fit()`` (or carried by the post-fit aggregation kit), + ensuring consistent unit-level arrays and avoiding repeated + subsetting of panel-level survey data. + """ + # Built once in fit() before any call lands here (see docstring). + assert self._unit_resolved_survey is not None + if self._unit_resolved_survey.uses_replicate_variance: + from diff_diff.survey import compute_replicate_if_variance + + # Score-scale IFs to match TSL bread: psi = w * eif / sum(w) + w = self._unit_resolved_survey.weights + psi_scaled = w * eif_vals / w.sum() + variance, n_valid = compute_replicate_if_variance( + psi_scaled, self._unit_resolved_survey + ) + # Update survey df to reflect effective replicate count + if n_valid < self._unit_resolved_survey.n_replicates: + self._survey_df = n_valid - 1 if n_valid > 1 else None + return float(np.sqrt(max(variance, 0.0))) if np.isfinite(variance) else np.nan + + from diff_diff.survey import compute_survey_vcov + + X_ones = np.ones((len(eif_vals), 1)) + vcov = compute_survey_vcov(X_ones, eif_vals, self._unit_resolved_survey) + return float(np.sqrt(np.abs(vcov[0, 0]))) + + def _eif_se( + self, + eif_vals: np.ndarray, + n_units: int, + cluster_indices: Optional[np.ndarray] = None, + n_clusters: Optional[int] = None, + ) -> float: + """Compute SE from aggregated EIF scores. + + Dispatches to survey TSL when ``_unit_resolved_survey`` is set + (during fit, or via the post-fit kit), otherwise uses + cluster-robust or standard formula. + """ + if self._unit_resolved_survey is not None: + return self._compute_survey_eif_se(eif_vals) + return _compute_se_from_eif(eif_vals, n_units, cluster_indices, n_clusters) + + # -- Aggregation helpers -------------------------------------------------- + + def _compute_wif_contribution( + self, + keepers: List[Tuple], + effects: np.ndarray, + unit_cohorts: np.ndarray, + cohort_fractions: Dict[float, float], + n_units: int, + unit_weights: Optional[np.ndarray] = None, + ) -> np.ndarray: + """Compute weight influence function correction (O(1) scale, matching EIF). + + This accounts for uncertainty in cohort-size aggregation weights. + Matches R's ``did`` package WIF formula (staggered_aggregation.py:282-309), + adapted to EDiD's EIF scale. + + Parameters + ---------- + keepers : list of (g, t) tuples + Post-treatment group-time pairs included in aggregation. + effects : ndarray, shape (n_keepers,) + ATT estimates for each keeper. + unit_cohorts : ndarray, shape (n_units,) + Cohort assignment for each unit (0 = never-treated). + cohort_fractions : dict + ``{cohort: n_cohort / n}`` for each cohort. + n_units : int + Total number of units. + unit_weights : ndarray, shape (n_units,), optional + Survey weights at the unit level. When provided, uses the + survey-weighted WIF formula: IF_i(p_g) = (w_i * 1{G_i=g} - pg_k). + + Returns + ------- + ndarray, shape (n_units,) + WIF contribution at O(1) scale, additive with ``agg_eif``. + """ + groups_for_keepers = np.array([g for (g, t) in keepers]) + pg_keepers = np.array([cohort_fractions.get(g, 0.0) for g, t in keepers]) + sum_pg = pg_keepers.sum() + if sum_pg == 0: + return np.zeros(n_units) + + indicator = (unit_cohorts[:, None] == groups_for_keepers[None, :]).astype(float) + + if unit_weights is not None: + # Survey-weighted WIF (matches staggered_aggregation.py:392-401): + # IF_i(p_g) = (w_i * 1{G_i=g} - pg_k), NOT (1{G_i=g} - pg_k) + weighted_indicator = indicator * unit_weights[:, None] + indicator_diff = weighted_indicator - pg_keepers + indicator_sum = np.sum(indicator_diff, axis=1) + else: + indicator_diff = indicator - pg_keepers + indicator_sum = np.sum(indicator_diff, axis=1) + + with np.errstate(divide="ignore", invalid="ignore", over="ignore"): + if1 = indicator_diff / sum_pg + if2 = np.outer(indicator_sum, pg_keepers) / sum_pg**2 + wif_matrix = if1 - if2 + wif_contrib = wif_matrix @ effects + return wif_contrib # O(1) scale, same as agg_eif + + def _aggregate_overall( + self, + group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]], + eif_by_gt: Dict[Tuple[Any, Any], np.ndarray], + n_units: int, + cohort_fractions: Dict[float, float], + unit_cohorts: np.ndarray, + cluster_indices: Optional[np.ndarray] = None, + n_clusters: Optional[int] = None, + ) -> Tuple[float, float]: + """Compute overall ATT with WIF-adjusted SE. + + Parameters + ---------- + group_time_effects : dict + Group-time ATT estimates. + eif_by_gt : dict + Per-unit EIF values for each (g, t). + n_units : int + Total number of units. + cohort_fractions : dict + Cohort size fractions. + unit_cohorts : ndarray, shape (n_units,) + Cohort assignment for each unit. + """ + # Filter to post-treatment effects + keepers = [ + (g, t) + for (g, t) in group_time_effects + if t >= g - self.anticipation and np.isfinite(group_time_effects[(g, t)]["effect"]) + ] + if not keepers: + return np.nan, np.nan + + # Cohort-size weights + pg = np.array([cohort_fractions.get(g, 0.0) for (g, _) in keepers]) + total_pg = pg.sum() + if total_pg == 0: + return np.nan, np.nan + w = pg / total_pg + + effects = np.array([group_time_effects[gt]["effect"] for gt in keepers]) + overall_att = float(np.sum(w * effects)) + + # Aggregate EIF + agg_eif = np.zeros(n_units) + for k, gt in enumerate(keepers): + agg_eif += w[k] * eif_by_gt[gt] + + # WIF correction: accounts for uncertainty in cohort-size weights + wif = self._compute_wif_contribution( + keepers, + effects, + unit_cohorts, + cohort_fractions, + n_units, + unit_weights=self._unit_level_weights, + ) + # Compute SE: survey path uses score-level psi to avoid double-weighting + # (compute_survey_vcov applies w_i internally, which would double-weight + # the survey-weighted WIF term). Dispatch replicate vs TSL. + if self._unit_resolved_survey is not None: + uw = self._unit_level_weights + # Set together with _unit_resolved_survey in fit(). + assert uw is not None + total_w = float(np.sum(uw)) + psi_total = uw * agg_eif / total_w + wif / total_w + + if ( + hasattr(self._unit_resolved_survey, "uses_replicate_variance") + and self._unit_resolved_survey.uses_replicate_variance + ): + from diff_diff.survey import compute_replicate_if_variance + + variance, _ = compute_replicate_if_variance(psi_total, self._unit_resolved_survey) + else: + from diff_diff.survey import compute_survey_if_variance + + variance = compute_survey_if_variance(psi_total, self._unit_resolved_survey) + se = float(np.sqrt(max(variance, 0.0))) if np.isfinite(variance) else np.nan + else: + agg_eif_total = agg_eif + wif + se = self._eif_se(agg_eif_total, n_units, cluster_indices, n_clusters) + + return overall_att, se + + def _aggregate_event_study( + self, + group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]], + eif_by_gt: Dict[Tuple[Any, Any], np.ndarray], + n_units: int, + cohort_fractions: Dict[float, float], + treatment_groups: List[Any], + time_periods: List[Any], + balance_e: Optional[int] = None, + unit_cohorts: Optional[np.ndarray] = None, + cluster_indices: Optional[np.ndarray] = None, + n_clusters: Optional[int] = None, + ) -> Dict[int, Dict[str, Any]]: + """Aggregate ATT(g,t) by relative time e = t - g. + + Parameters + ---------- + group_time_effects : dict + Group-time ATT estimates. + eif_by_gt : dict + Per-unit EIF values for each (g, t). + n_units : int + Total number of units. + cohort_fractions : dict + Cohort size fractions. + treatment_groups : list + Treatment cohort identifiers. + time_periods : list + All time periods. + balance_e : int, optional + Balance event study at this relative period. + unit_cohorts : ndarray, optional + Cohort assignment for each unit (for WIF correction). + """ + # Organize by relative time + effects_by_e: Dict[int, List[Tuple[Tuple[Any, Any], float, float]]] = {} + for (g, t), data in group_time_effects.items(): + if not np.isfinite(data["effect"]): + continue + e = int(t - g) + if e not in effects_by_e: + effects_by_e[e] = [] + effects_by_e[e].append(((g, t), data["effect"], cohort_fractions.get(g, 0.0))) + + # Balance if requested + if balance_e is not None: + groups_at_e = {gt[0] for gt, _, _ in effects_by_e.get(balance_e, [])} + balanced: Dict[int, List[Tuple[Tuple[Any, Any], float, float]]] = {} + for (g, t), data in group_time_effects.items(): + if not np.isfinite(data["effect"]): + continue + if g in groups_at_e: + e = int(t - g) + if e not in balanced: + balanced[e] = [] + balanced[e].append(((g, t), data["effect"], cohort_fractions.get(g, 0.0))) + effects_by_e = balanced + + if balance_e is not None and not effects_by_e: + warnings.warn( + f"balance_e={balance_e}: no cohort has a finite effect at the " + "anchor horizon. Event study will be empty.", + UserWarning, + stacklevel=2, + ) + + result: Dict[int, Dict[str, Any]] = {} + for e, elist in sorted(effects_by_e.items()): + gt_pairs = [x[0] for x in elist] + effs = np.array([x[1] for x in elist]) + pgs = np.array([x[2] for x in elist]) + total_pg = pgs.sum() + w = pgs / total_pg if total_pg > 0 else np.ones(len(pgs)) / len(pgs) + + agg_eff = float(np.sum(w * effs)) + + # Aggregate EIF + agg_eif = np.zeros(n_units) + for k, gt in enumerate(gt_pairs): + agg_eif += w[k] * eif_by_gt[gt] + + # WIF correction for event-study aggregation + wif_e = np.zeros(n_units) + if unit_cohorts is not None: + es_keepers = [(g, t) for (g, t) in gt_pairs] + es_effects = effs + wif_e = self._compute_wif_contribution( + es_keepers, + es_effects, + unit_cohorts, + cohort_fractions, + n_units, + unit_weights=self._unit_level_weights, + ) + + if self._unit_resolved_survey is not None: + uw = self._unit_level_weights + # Set together with _unit_resolved_survey in fit(). + assert uw is not None + total_w = float(np.sum(uw)) + psi_total = uw * agg_eif / total_w + wif_e / total_w + + if ( + hasattr(self._unit_resolved_survey, "uses_replicate_variance") + and self._unit_resolved_survey.uses_replicate_variance + ): + from diff_diff.survey import compute_replicate_if_variance + + variance, _ = compute_replicate_if_variance( + psi_total, self._unit_resolved_survey + ) + else: + from diff_diff.survey import compute_survey_if_variance + + variance = compute_survey_if_variance(psi_total, self._unit_resolved_survey) + agg_se = float(np.sqrt(max(variance, 0.0))) if np.isfinite(variance) else np.nan + else: + agg_eif = agg_eif + wif_e + agg_se = self._eif_se(agg_eif, n_units, cluster_indices, n_clusters) + + t_stat, p_val, ci = safe_inference( + agg_eff, agg_se, alpha=self.alpha, df=self._survey_df + ) + result[e] = { + "effect": agg_eff, + "se": agg_se, + "t_stat": t_stat, + "p_value": p_val, + "conf_int": ci, + "n_groups": len(elist), + } + + return result + + def _aggregate_by_group( + self, + group_time_effects: Dict[Tuple[Any, Any], Dict[str, Any]], + eif_by_gt: Dict[Tuple[Any, Any], np.ndarray], + n_units: int, + cohort_fractions: Dict[float, float], + treatment_groups: List[Any], + unit_cohorts: Optional[np.ndarray] = None, + cluster_indices: Optional[np.ndarray] = None, + n_clusters: Optional[int] = None, + ) -> Dict[Any, Dict[str, Any]]: + """Aggregate ATT(g,t) by treatment cohort. + + Parameters + ---------- + group_time_effects : dict + Group-time ATT estimates. + eif_by_gt : dict + Per-unit EIF values for each (g, t). + n_units : int + Total number of units. + cohort_fractions : dict + Cohort size fractions. + treatment_groups : list + Treatment cohort identifiers. + unit_cohorts : ndarray, optional + Cohort assignment for each unit (unused — group aggregation + uses equal weights, not cohort-size weights). + + Notes + ----- + Each row dict records ``df_used`` — the ``self._survey_df`` value at + that row's ``safe_inference`` call (M-023: exact per-row df + provenance for the post-fit group relay). In every constructible + fit all rows share one value (replicate tightening completes during + the per-cell estimation loop), but capture-at-use is exact by + construction and robust to any future path that could diverge. + The key is additive to the public row-dict schema. + """ + result: Dict[Any, Dict[str, Any]] = {} + for g in treatment_groups: + g_gts = [ + (gg, t) + for (gg, t) in group_time_effects + if gg == g + and t >= g - self.anticipation + and np.isfinite(group_time_effects[(gg, t)]["effect"]) + ] + if not g_gts: + continue + + effs = np.array([group_time_effects[gt]["effect"] for gt in g_gts]) + w = np.ones(len(effs)) / len(effs) + agg_eff = float(np.sum(w * effs)) + + agg_eif = np.zeros(n_units) + for k, gt in enumerate(g_gts): + agg_eif += w[k] * eif_by_gt[gt] + agg_se = self._eif_se(agg_eif, n_units, cluster_indices, n_clusters) + + df_used = self._survey_df + t_stat, p_val, ci = safe_inference(agg_eff, agg_se, alpha=self.alpha, df=df_used) + result[g] = { + "effect": agg_eff, + "se": agg_se, + "t_stat": t_stat, + "p_value": p_val, + "conf_int": ci, + "n_periods": len(g_gts), + "df_used": df_used, + } + + return result diff --git a/diff_diff/efficient_did_results.py b/diff_diff/efficient_did_results.py index 958695dc..004f60f0 100644 --- a/diff_diff/efficient_did_results.py +++ b/diff_diff/efficient_did_results.py @@ -5,14 +5,17 @@ to_dataframe(), and significance properties. """ +import dataclasses from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd +from diff_diff.aggregation import AggregationMixin, AggregationResult +from diff_diff.efficient_did_aggregation import _EfficientAggregationMixin from diff_diff.results import _format_survey_block, _get_significance_stars -from diff_diff.results_base import BaseResults +from diff_diff.results_base import BaseResults, build_event_study_surface if TYPE_CHECKING: from diff_diff.efficient_did_bootstrap import EDiDBootstrapResults @@ -52,8 +55,38 @@ def __repr__(self) -> str: ) +class _EDiDKitAggregator(_EfficientAggregationMixin): + """Throwaway host for post-fit EIF re-aggregation (M-023). + + Exposes exactly the five attributes the extracted + ``_EfficientAggregationMixin`` methods read (its typed host contract): + ``alpha``, ``anticipation``, ``_survey_df``, ``_unit_resolved_survey``, + ``_unit_level_weights``. A FRESH instance is built per ``aggregate()`` + call, which is what contains the one mutation the mixin performs - + ``_compute_survey_eif_se`` writes ``self._survey_df`` when a degenerate + replicate design drops replicates - on the throwaway rather than the + retained kit, preserving the aggregate() immutability contract. This + is also what keeps ``aggregate()`` off an ``_estimator_ref`` (the CS + ``_KitAggregator`` precedent). + """ + + def __init__( + self, + alpha: float, + anticipation: int, + survey_df: Optional[float], + resolved_survey_unit: Optional[Any], + unit_level_weights: Optional["np.ndarray"], + ) -> None: + self.alpha = alpha + self.anticipation = anticipation + self._survey_df = survey_df + self._unit_resolved_survey = resolved_survey_unit + self._unit_level_weights = unit_level_weights + + @dataclass -class EfficientDiDResults(BaseResults): +class EfficientDiDResults(BaseResults, AggregationMixin): """ Results from Efficient DiD (Chen, Sant'Anna & Xie 2025) estimation. @@ -120,6 +153,9 @@ class EfficientDiDResults(BaseResults): ``{(g, t): ndarray(n_units,)}`` — per-unit EIF values for each group-time cell. Only populated when ``store_eif=True`` in :meth:`~EfficientDiD.fit` (used internally by ``hausman_pretest``). + Since 3.9 (row M-023) the private aggregation kit ALWAYS retains + the same per-(g,t) EIF dict to power post-fit + :meth:`aggregate` — ``store_eif`` governs only this public field. bootstrap_results : EDiDBootstrapResults, optional Bootstrap inference results. estimation_path : str @@ -191,6 +227,15 @@ class EfficientDiDResults(BaseResults): omega_ridge: float = 0.0 # Survey design metadata (SurveyMetadata instance from diff_diff.survey) survey_metadata: Optional[Any] = field(default=None) + # Post-fit aggregation kit (M-023) - declared LAST: the generated + # __init__'s positional indexes are public API (CS precedent). + _aggregation_kit: Optional[Any] = field(default=None, repr=False, compare=False) + + # Post-fit aggregate() hooks (M-023). Plain class attributes (no + # annotation) so they never enter dataclasses.fields; the mixin's + # ClassVar-annotated defaults document the contract. balance_e keeps + # the mixin default ("event_study",) - CS precedent, do not redeclare. + _AGGREGATE_SUPPORTED = ("simple", "event_study", "group") # --- Inference-field aliases (balance/external-adapter compatibility) --- @property @@ -213,6 +258,191 @@ def p_value(self) -> float: def t_stat(self) -> float: return self.overall_t_stat + # --- Post-fit aggregation (M-023) ----------------------------------- + + @property + def reference_period(self) -> Optional[int]: + """Reference event time of the materialized PT-Post anchor, or None. + + MEMBERSHIP-GATED (the SunAbraham rule): returns + ``-1 - anticipation`` only when this fit ran under + ``pt_assumption="post"`` AND that event time is actually present in + ``event_study_effects`` — under PT-Post the per-cohort baseline + cell ``(g, g-1-anticipation)`` is estimated as a mechanical zero + whenever it is not the panel's first period, and the builder marks + that materialized row ``is_reference``. When the anchor cell was + never estimated (every surviving cohort baselined at the first + period) the property is None, so no reference row is ever + synthesized. Under ``pt_assumption="all"`` there is no reference + row (universal first-period baseline; every row is a genuine + estimate) and the property is None. + """ + if self.pt_assumption != "post": + return None + if not self.event_study_effects: + return None + ref = -1 - int(self.anticipation) + return ref if ref in self.event_study_effects else None + + def _aggregate_compute( + self, level: str, *, weights: Optional[str], balance_e: Optional[int] + ) -> Any: + kit = self._aggregation_kit + if kit is None: + raise ValueError( + "This EfficientDiDResults carries no aggregation kit - it is " + "attached by EfficientDiD.fit(), so a result unpickled from " + "an older release will not have one. Re-fit with " + "diff-diff >= 3.9 to aggregate post-fit." + ) + if self.bootstrap_results is not None: + raise NotImplementedError( + "aggregate() is not yet available on a bootstrapped fit " + "(n_bootstrap > 0): the per-horizon bootstrap draws are not " + "retained, so post-fit re-aggregation cannot replay " + "percentile inference and analytical inference would " + "misrepresent the fit. Re-fit with the aggregation you " + "need, or use n_bootstrap=0." + ) + bk = dict(kit.bookkeeping) + agg = _EDiDKitAggregator( + alpha=kit.alpha, + anticipation=kit.anticipation, + survey_df=bk["df_survey"], + resolved_survey_unit=bk["resolved_survey_unit"], + unit_level_weights=bk["unit_level_weights"], + ) + if level == "simple": + return self._aggregate_simple_result(kit) + if level == "group": + effects = agg._aggregate_by_group( + bk["group_time_effects"], + kit.influence, + bk["n_units"], + bk["cohort_fractions"], + bk["treatment_groups"], + unit_cohorts=bk["unit_cohorts"], + cluster_indices=bk["cluster_indices"], + n_clusters=bk["n_clusters"], + ) + return self._group_effects_to_aggregation(effects, kit) + # level == "event_study" (the mixin validated the vocabulary) + es = agg._aggregate_event_study( + bk["group_time_effects"], + kit.influence, + bk["n_units"], + bk["cohort_fractions"], + bk["treatment_groups"], + bk["time_periods"], + balance_e, + unit_cohorts=bk["unit_cohorts"], + cluster_indices=bk["cluster_indices"], + n_clusters=bk["n_clusters"], + ) + # Carrier + shared builder: EDiD is a _from_relative_dict producer, + # so the recomputed dict rides the same route as the fit-time + # surface. The carrier's survey_metadata is a COPY whose df_survey + # is the kit's post-overall snapshot - _recompute_unit_survey_ + # metadata's is-not-None guard can leave the raw metadata at the + # resolved design's finite value when the governing df degenerated + # to None (n_valid <= 1), and the copy keeps the container honest + # there (in non-degenerate fits the two agree and the copy is + # inert). Copying also keeps the parent's metadata unmutated. + meta = self.survey_metadata + if meta is not None: + meta = dataclasses.replace(meta, df_survey=bk["df_survey"]) + # The carrier's PROVENANCE fields also come from the kit snapshots + # (CI review R2): the reference_period property and the container's + # alpha/anticipation must reflect the FIT's regime, not a possibly + # mutated public field - a PT-All fit whose public pt_assumption + # was flipped to "post" would otherwise mark the genuine e=-1 + # estimate as a reference row and zero it. + carrier = dataclasses.replace( + self, + event_study_effects=es, + survey_metadata=meta, + pt_assumption=bk["pt_assumption"], + anticipation=kit.anticipation, + alpha=kit.alpha, + ) + return build_event_study_surface(carrier) + + def _aggregate_simple_result(self, kit: Any) -> AggregationResult: + """One-row relay of the stored overall inference (bit-exact). + + ``n = n_treated_units + n_control_units`` with ``n_kind="units"``: + EDiD's treated and control unit sets are DISJOINT by construction + (``last_cohort`` trimming reassigns the last cohort to control + BEFORE the counts), so a true disjoint total exists - the CS + convention applies, unlike StackedDiD's overlapping-sets carve-out. + + ``df`` is the kit's post-overall ``df_survey`` snapshot - the very + value fit passed to ``safe_inference`` for the overall row - NOT + ``resolve_inference_df(self)``: ``survey_metadata.df_survey`` can + diverge from the governing df in the degenerate replicate case + (``n_valid <= 1`` sets the working df to None while the metadata + keeps the resolved design's finite value), and the snapshot is + provenance-exact in every state. None → all-NaN df column; + the replicate-undefined 0-sentinel row NaNs out via post_init. + """ + df_val = kit.bookkeeping["df_survey"] + return AggregationResult( + level="simple", + label=np.array(["overall"], dtype=object), + target=np.array(["att"], dtype=object), + att=np.array([self.overall_att], dtype=float), + se=np.array([self.overall_se], dtype=float), + t_stat=np.array([self.overall_t_stat], dtype=float), + p_value=np.array([self.overall_p_value], dtype=float), + conf_int_lower=np.array([self.overall_conf_int[0]], dtype=float), + conf_int_upper=np.array([self.overall_conf_int[1]], dtype=float), + n=np.array([kit.bookkeeping["n_units_total"]], dtype=float), + df=df_val, + alpha=kit.alpha, + n_kind="units", + weight=np.array([1.0], dtype=float), + estimator=type(self).__name__.replace("Results", ""), + ) + + def _group_effects_to_aggregation( + self, effects: Dict[Any, Dict[str, Any]], kit: Any + ) -> AggregationResult: + """Per-cohort AggregationResult from the recomputed group dict. + + ``df`` relays the PER-ROW ``df_used`` array the aggregation + recorded at each row's ``safe_inference`` call - exact by + construction (in every constructible fit all rows share one value; + capture-at-use is robust regardless). ``weight=None``: rows use + equal within-cohort weights and carry no cross-cohort mass, so a + weight column would fabricate one (the CS rationale). + ``n_kind="cells"``: ``n_periods`` counts contributing (g,t) cells. + """ + labels = list(effects.keys()) + df_arr = np.array( + [ + (np.nan if effects[g].get("df_used") is None else float(effects[g]["df_used"])) + for g in labels + ], + dtype=float, + ) + return AggregationResult( + level="group", + label=np.array(labels, dtype=object), + target=np.array(["att"] * len(labels), dtype=object), + att=np.array([effects[g]["effect"] for g in labels], dtype=float), + se=np.array([effects[g]["se"] for g in labels], dtype=float), + t_stat=np.array([effects[g]["t_stat"] for g in labels], dtype=float), + p_value=np.array([effects[g]["p_value"] for g in labels], dtype=float), + conf_int_lower=np.array([effects[g]["conf_int"][0] for g in labels], dtype=float), + conf_int_upper=np.array([effects[g]["conf_int"][1] for g in labels], dtype=float), + n=np.array([effects[g]["n_periods"] for g in labels], dtype=float), + df=df_arr, + alpha=kit.alpha, + n_kind="cells", + weight=None, + estimator=type(self).__name__.replace("Results", ""), + ) + def __repr__(self) -> str: sig = _get_significance_stars(self.overall_p_value) path = "DR" if self.estimation_path == "dr" else "nocov" @@ -398,7 +628,14 @@ def to_dataframe(self, level: str = "group_time") -> pd.DataFrame: elif level == "event_study": if self.event_study_effects is None: - raise ValueError("Event study effects not computed. Use aggregate='event_study'.") + raise ValueError( + "Event study effects not computed at fit time. Use " + "results.aggregate('event_study') for the post-fit " + "event-study container (on bootstrapped fits, re-fit " + "with n_bootstrap=0 or use the deprecated fit-time " + "aggregate=); a result unpickled from a pre-3.9 " + "release carries no aggregation kit and must be refit." + ) rows = [] for rel_t, data in sorted(self.event_study_effects.items()): rows.append( @@ -416,7 +653,14 @@ def to_dataframe(self, level: str = "group_time") -> pd.DataFrame: elif level == "group": if self.group_effects is None: - raise ValueError("Group effects not computed. Use aggregate='group'.") + raise ValueError( + "Group effects not computed at fit time. Use " + "results.aggregate('group') for the post-fit group " + "container (on bootstrapped fits, re-fit with " + "n_bootstrap=0 or use the deprecated fit-time " + "aggregate=); a result unpickled from a pre-3.9 " + "release carries no aggregation kit and must be refit." + ) rows = [] for group, data in sorted(self.group_effects.items()): rows.append( diff --git a/diff_diff/guides/llms-full.txt b/diff_diff/guides/llms-full.txt index bb924895..15bfb4e4 100644 --- a/diff_diff/guides/llms-full.txt +++ b/diff_diff/guides/llms-full.txt @@ -1082,8 +1082,8 @@ edid.fit( time: str, first_treat: str, covariates: list[str] = None, # Time-invariant unit-level covariates; uses doubly-robust sieve path when non-None - aggregate: str = None, # None, "simple", "event_study", "group", or "all" - balance_e: int = None, + aggregate: Any = NOT_SUPPLIED, # DEPRECATED (M-023): warns, still works - use results.aggregate() post-fit + balance_e: Any = NOT_SUPPLIED, # DEPRECATED (M-120): moves onto results.aggregate(balance_e=) ) -> EfficientDiDResults ``` @@ -1094,8 +1094,12 @@ from diff_diff import EfficientDiD edid = EfficientDiD(pt_assumption="all") results = edid.fit(data, outcome='y', unit='id', time='t', - first_treat='first_treat', aggregate='all') + first_treat='first_treat') results.print_summary() +# Aggregate post-fit (recomputed from retained EIFs; bootstrapped fits raise): +es = results.aggregate('event_study') # EventStudyResults container +grp = results.aggregate('group') # per-cohort AggregationResult +print(grp.to_dataframe()) ``` ### LPDiD @@ -1761,8 +1765,8 @@ ImputationDiDResults, TwoStageDiDResults, StackedDiDResults, and EfficientDiDRes | `overall_t_stat` | `float` | T-statistic | | `overall_p_value` | `float` | P-value | | `overall_conf_int` | `tuple[float, float]` | CI | -| `event_study_effects` | `dict[int, dict]` | Event study effects (if aggregate includes event_study; ALWAYS populated for StackedDiD since 3.9 - row M-024) | -| `group_effects` | `dict` | Group-level effects (if aggregate includes group) | +| `event_study_effects` | `dict[int, dict]` | Event study effects (if the DEPRECATED fit-time aggregate included event_study; ALWAYS populated for StackedDiD since 3.9 - row M-024; post-fit `aggregate('event_study')` returns a CONTAINER, it does not populate this field) | +| `group_effects` | `dict` | Group-level effects (if the deprecated fit-time aggregate included group; post-fit `aggregate('group')` returns an AggregationResult container) | | `groups` | `list` | Treatment cohorts | | `time_periods` | `list` | All time periods | | `n_obs` | `int` | Number of observations | @@ -1780,7 +1784,7 @@ Each event study effect dict contains: `effect`, `se`, `t_stat`, `p_value`, `con | `n_clusters` | `int | None` | Number of effective clusters; `None` under survey designs and `None` under EfficientDiD's default unclustered fit | | `df_convention` | `str | None` | On the knob-carrying containers (`StackedDiDResults`, `ImputationDiDResults`, `WooldridgeDiDResults`, `LPDiDResults`, `SunAbrahamResults`): the configured df convention (3.9 / M-127); `StackedDiDResults` additionally carries `inference_df` (the overall-ATT df actually used) | -**Methods:** `summary()`, `print_summary()`, `to_dataframe()`, `to_dict()` (flat dict of headline aliases + `vcov_type` + conditional `cluster_name`/`n_clusters`/`n_bootstrap`/`inference_method`) +**Methods:** `summary()`, `print_summary()`, `to_dataframe()`, `to_dict()` (flat dict of headline aliases + `vcov_type` + conditional `cluster_name`/`n_clusters`/`n_bootstrap`/`inference_method`); `aggregate(type, weights=None, *, balance_e=None)` on the shipped post-fit adopters (`StackedDiDResults` views since 3.9/M-024; `EfficientDiDResults` recomputes from retained EIFs since 3.9/M-023 - bootstrapped EDiD fits raise) ### ContinuousDiDResults diff --git a/diff_diff/guides/llms-practitioner.txt b/diff_diff/guides/llms-practitioner.txt index 433e8811..e0fb3ab1 100644 --- a/diff_diff/guides/llms-practitioner.txt +++ b/diff_diff/guides/llms-practitioner.txt @@ -417,18 +417,16 @@ estimate is contaminated by forbidden comparisons (see Step 4). Aggregate treatment effects may mask important variation. -### For CallawaySantAnna (and estimators with `aggregate` parameter) +### For CallawaySantAnna and EfficientDiD (post-fit `aggregate()`) ```python -# Group-level effects — do early vs late adopters differ? -results = cs.fit(data, ..., aggregate='group') -print(results.overall_att) # Overall ATT -print(results.group_effects) # Per-cohort ATTs - -# Event study — how does the effect evolve over time? -results = cs.fit(data, ..., aggregate='event_study') -print(results.event_study_effects) # Per-relative-period ATTs - -# All aggregations at once +# Post-fit, no refit needed (fit-time aggregate= is deprecated since 3.9 +# for CS/StackedDiD/EfficientDiD; removed in 4.0): +results = cs.fit(data, ...) +print(results.overall_att) # Overall ATT +print(results.aggregate('group').to_dataframe()) # Per-cohort ATTs +print(results.aggregate('event_study').to_dataframe()) # Dynamic effects +# EXCEPTION: a BOOTSTRAPPED CS or EfficientDiD fit's aggregate() raises — +# use the deprecated fit-time aggregation there: results = cs.fit(data, ..., aggregate='all') ``` diff --git a/diff_diff/guides/llms.txt b/diff_diff/guides/llms.txt index 14301b3a..fd097fbd 100644 --- a/diff_diff/guides/llms.txt +++ b/diff_diff/guides/llms.txt @@ -21,7 +21,7 @@ diagnostic steps produces unreliable results. 4. **Choose estimator** — staggered adoption → CS/SA/BJS (NOT plain TWFE); few treated units → SDiD; factor confounding → TROP; simple 2x2 → DiD. Run `BaconDecomposition` to diagnose TWFE bias. 5. **Estimate** — `estimator.fit(data, ...)`. Always print the cluster count first and choose inference method based on the result (cluster-robust if >= 50 clusters, wild bootstrap if fewer). 6. **Sensitivity analysis** — `compute_honest_did(results)` for bounds under PT violations (MultiPeriodDiD, CS, or dCDH natively; a StackedDiD `results.aggregate('event_study')` container also admits - needs `kappa_pre >= 2` so estimated pre-periods exist), `run_all_placebo_tests()` for 2x2 falsification, specification comparisons for staggered designs. -7. **Heterogeneity** — CS: `results.aggregate('group')`/`.aggregate('event_study')` post-fit, no refit (fit-time `aggregate=`/`balance_e=` are deprecated since 3.9, removed in 4.0; `compute_honest_did` / `compute_pretrends_power` / `plot_event_study` all accept the post-fit `results.aggregate('event_study')` container directly; EXCEPTION: a BOOTSTRAPPED CS fit's `aggregate()` raises — use the fit-time aggregation there); dCDH: `results.aggregate('event_study')`/`.aggregate('simple')` post-fit views (bootstrap fits included — pure views); SA: `results.event_study_effects`/`to_dataframe(level='cohort')`; Stacked: `results.aggregate('event_study')`/`.aggregate('simple')` post-fit views (the surface is ALWAYS computed at fit since 3.9 - row M-024 - and the container admits into `compute_honest_did`/`compute_pretrends_power` with `kappa_pre >= 2`); subgroup re-estimation. +7. **Heterogeneity** — CS: `results.aggregate('group')`/`.aggregate('event_study')` post-fit, no refit (fit-time `aggregate=`/`balance_e=` are deprecated since 3.9, removed in 4.0; `compute_honest_did` / `compute_pretrends_power` / `plot_event_study` all accept the post-fit `results.aggregate('event_study')` container directly; EXCEPTION: a BOOTSTRAPPED CS fit's `aggregate()` raises — use the fit-time aggregation there); dCDH: `results.aggregate('event_study')`/`.aggregate('simple')` post-fit views (bootstrap fits included — pure views); SA: `results.event_study_effects`/`to_dataframe(level='cohort')`; Stacked: `results.aggregate('event_study')`/`.aggregate('simple')` post-fit views (the surface is ALWAYS computed at fit since 3.9 - row M-024 - and the container admits into `compute_honest_did`/`compute_pretrends_power` with `kappa_pre >= 2`); EDiD: `results.aggregate('event_study')`/`.aggregate('group')`/`.aggregate('simple')` post-fit, RECOMPUTED from retained EIFs (3.9, row M-023; fit-time `aggregate=`/`balance_e=` deprecated; bootstrapped EDiD fits raise - use the fit-time aggregation there; EDiD containers are NOT admitted into honest/pretrends - no joint ES covariance); subgroup re-estimation. 8. **Robustness** — compare 2-3 estimators (CS vs SA vs BJS), MUST report with and without covariates (shows whether conditioning drives identification), present pre-trends and sensitivity bounds. After estimation, call `practitioner_next_steps(results)` for context-aware diff --git a/diff_diff/honest_did.py b/diff_diff/honest_did.py index 9cc953ea..3c004d8e 100644 --- a/diff_diff/honest_did.py +++ b/diff_diff/honest_did.py @@ -608,8 +608,11 @@ def _extract_container_params( f"(got source={surface.source!r}). For other estimators pass " "the native results object where supported " "(MultiPeriodDiDResults, CallawaySantAnnaResults, or " - "ChaisemartinDHaultfoeuilleResults); container admission for " - "further producers arrives with their own aggregate() " + "ChaisemartinDHaultfoeuilleResults); " + "EfficientDiDResults containers are rejected BY DESIGN " + "(their surface carries no joint event-study covariance - " + "see the REGISTRY EfficientDiD Note); container admission " + "for further producers arrives with their own aggregate() " "migrations." ) _producer = surface.source.replace("Results", "") diff --git a/diff_diff/practitioner.py b/diff_diff/practitioner.py index 4aba5aa0..010d3085 100644 --- a/diff_diff/practitioner.py +++ b/diff_diff/practitioner.py @@ -890,6 +890,34 @@ def _handle_efficient(results: Any): # as a mismatched-step-name bug. step_name="parallel_trends", ), + _step( + baker_step=7, + label="Aggregate treatment-effect heterogeneity post-fit", + why=( + "EfficientDiD aggregates post-fit from retained EIFs " "(M-023) - no refit needed." + if getattr(results, "bootstrap_results", None) is None + else "This fit is BOOTSTRAPPED, and post-fit aggregate() " + "raises on bootstrap fits - refit with the deprecated " + "fit-time aggregation (or n_bootstrap=0) to obtain the " + "aggregated surfaces." + ), + code=( + "# Aggregate post-fit - no refit needed:\n" + "print(results.aggregate('group').to_dataframe()) # Per-cohort ATTs\n" + "print(results.aggregate('event_study').to_dataframe()) # Dynamic effects" + if getattr(results, "bootstrap_results", None) is None + else "# Bootstrap fit: aggregate at fit time (deprecated kwarg):\n" + "results = edid.fit(data, ..., aggregate='all')\n" + "print(results.group_effects) # Per-cohort ATTs\n" + "print(results.event_study_effects) # Dynamic effects" + ), + priority="medium", + # NON-STEPS key (the M-024 "sub_experiment_balance" lesson): + # a STEPS-vocabulary name would let _filter_steps suppress + # this guidance whenever a same-named DiagnosticReport check + # completes, which never runs this aggregation. + step_name="aggregation", + ), _robustness_compare_step("CS, SA, or BJS"), _covariates_step(), ] diff --git a/diff_diff/pretrends.py b/diff_diff/pretrends.py index e6f1d3d8..ea5eca53 100644 --- a/diff_diff/pretrends.py +++ b/diff_diff/pretrends.py @@ -1445,8 +1445,11 @@ def _extract_container_pre_period_params( f"(got source={surface.source!r}). For other estimators " "pass the native results object where supported " "(MultiPeriodDiDResults, CallawaySantAnnaResults, or " - "SunAbrahamResults); container admission for further " - "producers arrives with their own aggregate() migrations." + "SunAbrahamResults); EfficientDiDResults containers " + "are rejected BY DESIGN (their surface carries no joint " + "event-study covariance - see the REGISTRY EfficientDiD " + "Note); container admission for further producers " + "arrives with their own aggregate() migrations." ) _producer = surface.source.replace("Results", "") if surface.time_scale != "relative": diff --git a/diff_diff/results_base.py b/diff_diff/results_base.py index 92c52d03..4b0adc44 100644 --- a/diff_diff/results_base.py +++ b/diff_diff/results_base.py @@ -619,7 +619,7 @@ def summary(self, alpha: Optional[float] = None) -> str: # materialize the surface (row M-024). "StackedDiDResults": "re-fit with diff-diff >= 3.9, which always computes the surface", "StaggeredTripleDiffResults": "refit with aggregate='event_study' (or 'all')", - "EfficientDiDResults": "refit with aggregate='event_study' (or 'all')", + "EfficientDiDResults": "call results.aggregate('event_study') (on a bootstrapped fit, re-fit with n_bootstrap=0 or the deprecated fit-time aggregate=)", "ContinuousDiDResults": "refit with aggregate='eventstudy' (or 'all')", "WooldridgeDiDResults": "call results.aggregate('event_study') first", "SpilloverDiDResults": "refit with event_study=True", diff --git a/docs/api/_autosummary/diff_diff.EfficientDiDResults.rst b/docs/api/_autosummary/diff_diff.EfficientDiDResults.rst index 8ff6ce9a..a74f7f7f 100644 --- a/docs/api/_autosummary/diff_diff.EfficientDiDResults.rst +++ b/docs/api/_autosummary/diff_diff.EfficientDiDResults.rst @@ -12,6 +12,7 @@ .. autosummary:: ~EfficientDiDResults.__init__ + ~EfficientDiDResults.aggregate ~EfficientDiDResults.print_summary ~EfficientDiDResults.summary ~EfficientDiDResults.to_dataframe @@ -47,6 +48,7 @@ ~EfficientDiDResults.p_value ~EfficientDiDResults.pt_assumption ~EfficientDiDResults.ratio_clip + ~EfficientDiDResults.reference_period ~EfficientDiDResults.se ~EfficientDiDResults.seed ~EfficientDiDResults.sieve_criterion diff --git a/docs/api/efficient_did.rst b/docs/api/efficient_did.rst index 3c2681b3..09bb6484 100644 --- a/docs/api/efficient_did.rst +++ b/docs/api/efficient_did.rst @@ -82,10 +82,17 @@ Results container for Efficient DiD estimation. .. autosummary:: + ~EfficientDiDResults.aggregate ~EfficientDiDResults.summary ~EfficientDiDResults.print_summary ~EfficientDiDResults.to_dataframe + .. rubric:: Attributes + + .. autosummary:: + + ~EfficientDiDResults.reference_period + EDiDBootstrapResults -------------------- @@ -109,16 +116,18 @@ Basic usage:: edid = EfficientDiD(pt_assumption="all") results = edid.fit(data, outcome='outcome', unit='unit', - time='period', first_treat='first_treat', - aggregate='all') + time='period', first_treat='first_treat') results.print_summary() + # Aggregate post-fit - recomputed from retained EIFs, no refit: + print(results.aggregate('event_study').to_dataframe()) + print(results.aggregate('group').to_dataframe()) + PT-Post mode (matches CS for post-treatment ATT):: edid_post = EfficientDiD(pt_assumption="post") results_post = edid_post.fit(data, outcome='outcome', unit='unit', - time='period', first_treat='first_treat', - aggregate='all') + time='period', first_treat='first_treat') print(f"PT-All ATT: {results.overall_att:.4f} (SE={results.overall_se:.4f})") print(f"PT-Post ATT: {results_post.overall_att:.4f} (SE={results_post.overall_se:.4f})") @@ -126,8 +135,7 @@ Bootstrap inference:: edid_boot = EfficientDiD(pt_assumption="all", n_bootstrap=999, seed=42) results_boot = edid_boot.fit(data, outcome='outcome', unit='unit', - time='period', first_treat='first_treat', - aggregate='all') + time='period', first_treat='first_treat') print(f"Bootstrap SE: {results_boot.overall_se:.4f}") print(f"Bootstrap CI: [{results_boot.overall_conf_int[0]:.4f}, " f"{results_boot.overall_conf_int[1]:.4f}]") diff --git a/docs/choosing_estimator.rst b/docs/choosing_estimator.rst index d447cae2..0fa6f0f8 100644 --- a/docs/choosing_estimator.rst +++ b/docs/choosing_estimator.rst @@ -471,8 +471,7 @@ Use :class:`~diff_diff.EfficientDiD` when: edid = EfficientDiD(pt_assumption="all") # or "post" for post-treatment CS match results = edid.fit(data, outcome='y', unit='unit_id', - time='period', first_treat='first_treat', - aggregate='all') + time='period', first_treat='first_treat') results.print_summary() Sun-Abraham diff --git a/docs/dev-status.md b/docs/dev-status.md index ab99e6c6..f94cddc6 100644 --- a/docs/dev-status.md +++ b/docs/dev-status.md @@ -36,7 +36,7 @@ Target: ideally < 1000 lines per module; modules ≥3000 lines are candidates fo | `triple_diff.py` | 2231 | Monitor | | `wooldridge.py` | 2192 | Monitor | | `practitioner.py` | 2113 | Monitor — grew with per-estimator handlers (was 1511 on 2026-07-13) | -| `efficient_did.py` | 2083 | Monitor | +| `efficient_did.py` | 1729 | Acceptable — dropped below 2000 when the M-023 aggregate() migration extracted the aggregation mixin into `efficient_did_aggregation.py` (~520 lines, below this table's floor) | | `chaisemartin_dhaultfoeuille_results.py` | 2004 | Monitor | | `results.py` | 1948 | Acceptable | | `_rdrobust_port.py` | 1913 | Acceptable | diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index fb0d5daa..5ff3663d 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -44,6 +44,7 @@ groups: - diff_diff/two_stage_results.py efficient_did: - diff_diff/efficient_did.py + - diff_diff/efficient_did_aggregation.py - diff_diff/efficient_did_bootstrap.py - diff_diff/efficient_did_covariates.py - diff_diff/efficient_did_results.py diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index cea440ba..54062223 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -619,8 +619,8 @@ IFs inline; `ImputationDiD` computes the Theorem 3 conservative variance (`sigma_sq = (cluster_psi_sums**2).sum()`) directly from per-cluster influence-function sums; `EfficientDiD` aggregates per-unit EIF within clusters, centers, and applies the standard `G/(G-1)` correction -(`_cluster_aggregate` + `_compute_se_from_eif` at -`diff_diff/efficient_did.py:79-127`). The CR1 Liang-Zeger algebra on the +(`_cluster_aggregate` + `_compute_se_from_eif` in +`diff_diff/efficient_did_aggregation.py`). The CR1 Liang-Zeger algebra on the IF is Williams (2000) / Hansen (2007) in all four cases — no new methodology is introduced. @@ -1519,9 +1519,10 @@ where `q_{g,e} = pi_g / sum_{g' in G_{trt,e}} pi_{g'}`. - **Note:** Cluster-robust SEs use the standard Liang-Zeger clustered sandwich estimator applied to EIF values: aggregate EIF within clusters, center, and compute variance with G/(G-1) small-sample correction. Cluster bootstrap generates multiplier weights at the cluster level (all units in a cluster share the same weight). Analytical clustered SEs are the default when `cluster` is set; cluster bootstrap is opt-in via `n_bootstrap > 0`. - **Note:** Hausman pretest operates on the post-treatment event-study vector ES(e) per Theorem A.1. Both PT-All and PT-Post fits are aggregated to ES(e) using cohort-size weights before computing the test statistic H = delta' V^{-1} delta where delta = ES_post - ES_all and V = Cov(ES_post) - Cov(ES_all). Covariance is computed from aggregated ES(e)-level EIF values. The variance-difference matrix V is inverted via Moore-Penrose pseudoinverse to handle finite-sample non-positive-definiteness. Effective rank of V (number of positive eigenvalues) is used as degrees of freedom. - **Note:** Last-cohort-as-control (`control_group="last_cohort"`) reclassifies the latest treatment cohort as pseudo-never-treated and drops time periods at `t >= last_g - anticipation`, excluding anticipation-contaminated periods from the pseudo-control's pre-treatment window. This is distinct from CallawaySantAnna's `not_yet_treated` option which dynamically selects not-yet-treated units per (g,t) pair. -- **Note:** `vcov_type` is permanently narrow to `{"hc1"}` per the Chen-Sant'Anna-Xie (2025) EIF-based variance achieving the semiparametric efficiency bound. Analytical-sandwich families `{classical, hc2, hc2_bm}` are rejected at `__init__` — the per-unit EIF aggregation has no equivalent single design matrix on which hat-matrix leverage or Bell-McCaffrey Satterthwaite DOF can be defined. `cluster=` invokes Liang-Zeger CR1 on cluster-aggregated EIF (`_compute_se_from_eif` with `cluster_indices` at `diff_diff/efficient_did.py:124-127`); `survey_design=` invokes TSL on the combined IF (`_compute_survey_eif_se` at `diff_diff/efficient_did.py:1151-1176`). `vcov_type='conley'` deferred to the EfficientDiD Conley follow-up row in DEFERRED.md. +- **Note:** `vcov_type` is permanently narrow to `{"hc1"}` per the Chen-Sant'Anna-Xie (2025) EIF-based variance achieving the semiparametric efficiency bound. Analytical-sandwich families `{classical, hc2, hc2_bm}` are rejected at `__init__` — the per-unit EIF aggregation has no equivalent single design matrix on which hat-matrix leverage or Bell-McCaffrey Satterthwaite DOF can be defined. `cluster=` invokes Liang-Zeger CR1 on cluster-aggregated EIF (`_compute_se_from_eif` with `cluster_indices`); `survey_design=` invokes TSL on the combined IF (`_compute_survey_eif_se`); both live in `diff_diff/efficient_did_aggregation.py` since the M-023 post-fit aggregate() extraction. `vcov_type='conley'` deferred to the EfficientDiD Conley follow-up row in DEFERRED.md. - **Note:** Default `cluster=None` (no survey design) renders summary label "HC1 heteroskedasticity-robust" because the per-unit EIF SE `sqrt(mean(EIF²)/n)` is methodologically HC1-style (no Liang-Zeger G/(G-1) finite-sample correction). `EfficientDiDResults.cluster_name` and `n_clusters` stay None under unclustered fits. This diverges from `ImputationDiD` which auto-clusters at unit per Borusyak-Jaravel-Spiess (2024) Theorem 3 — there the default summary renders the CR1 unit-clustered label. - **Note:** `set_params(vcov_type=bad)` raises immediately on EVERY estimator: since the shared `BaseEstimator` mixin (`diff_diff/_base.py`, v4 2(c)-i), `set_params` validates transactionally by constructor probe re-init, so it enforces exactly `__init__`'s validation, eagerly, library-wide. The former split — EfficientDiD eager vs `ImputationDiD`/`TripleDifference`/`CallawaySantAnna` (and six more: SunAbraham, StackedDiD, StaggeredTripleDifference, SpilloverDiD, TROP, PreTrendsPower) accepting constructor-rejected values until `fit()` — is retired; the fit-time re-validation layers remain as a second check against DIRECT attribute mutation (`est.vcov_type = ...`), which no setter can see. +- **Note (post-fit aggregate() - rows M-023/M-120):** `fit(aggregate=, balance_e=)` is deprecated (3.9; removed 4.0; joint FutureWarning, warn-and-still-work) in favor of post-fit `EfficientDiDResults.aggregate(type, balance_e=)` - a LAZY RECOMPUTING KIT (the CallawaySantAnna class, not a StackedDiD/dCDH view relay): `fit()` computes nothing extra, the results object retains an `AggregationKit`, and `aggregate('event_study'/'group', balance_e=)` re-runs the extracted `_EfficientAggregationMixin` aggregators on a throwaway host while `aggregate('simple')` relays the stored overall row bit-exact. (a) RETAINED BUFFERS (memory contract; phrased as maxima - optional design fields stay None when unsupplied): the per-(g,t) EIF dict, O(n_units x n_gt), the dominant payload - retained on EVERY fit regardless of `store_eif`, which since 3.9 governs only the public `influence_functions` field; `unit_cohorts` (cohort labels), `unit_level_weights`, factorized cluster codes - O(n_units) each; on ordinary (TSL) survey fits the unit-level `ResolvedSurveyDesign` adds `weights` plus, where supplied, `strata`/`psu`/`fpc` (factorized int codes / float values, never raw labels) - up to four O(n_units) arrays; on replicate designs it adds the O(n_units x n_replicates) replicate matrix plus, where supplied, `replicate_strata`/`replicate_rscales` (O(n_replicates)); per-row dict SNAPSHOTS of `group_time_effects` plus copies of the `groups`/`time_periods` lists and the scalar `pt_assumption`/`n_treated+n_control` provenance (aggregate() recomputes exclusively from these private snapshots, never from the mutable public result fields - a user edit of the public rows cannot mix altered point estimates with the retained EIF variance); scalars `n_units`, `cohort_fractions`, and the POST-OVERALL `df_survey` snapshot (captured after the overall inference and before the ES/group gates: the group pass can degenerate the working df to None on replicate designs with `n_valid <= 1`, and every fit-time aggregation seeds from the post-overall value, so recompute replays the exact seed). The data-minimization guarantee is scoped to unit identifiers - no unit-label container is retained. (b) `balance_e` uses the ANCHOR-HORIZON rule (keep cohorts with a finite effect at `e == balance_e`, then retain all their horizons) - the SAME rule CallawaySantAnna uses, divergent only from ImputationDiD/TwoStageDiD's balanced-window rule; an anchor no cohort reaches warns and yields a legal zero-row container. (c) BOOTSTRAP fits fail closed for ALL levels: ES/group because per-horizon draws are not retained (exact-replay wiring is a TODO row), and 'simple' for uniform CS-parity conservatism BY DECISION - the stored overall row could relay (as dCDH's view-only aggregate permits), but EDiD keeps one uniform bootstrap rule so no level publishes analytical-provenance fields beside percentile inference; the same rule clears the fit-time group rows' `df_used` key under bootstrap. (d) CONTAINER ADMISSION into `compute_honest_did`/`compute_pretrends_power` is REJECTED BY DESIGN (both terminal TypeErrors state it): the primary ground is the absent joint event-study covariance (container `vcov=None`, all-NaN per-row df - the scalar `df_survey` channel is the container's only df provenance; the per-row hole is the tracked M-092-completion TODO row). Reference semantics are regime-dependent: under `pt_assumption="all"` there is NO reference row (universal first-period baseline; e=-1 is a genuine estimate); under `"post"` the per-cohort baseline cell is materialized as a mechanical zero anchor at `e = -1 - anticipation` whenever it is not the panel's first period, and the MEMBERSHIP-GATED `reference_period` property (the SunAbraham rule - never synthesized when the anchor cell was not estimated) marks it `is_reference` in the container and corrects `plot_event_study`'s inferred reference (previously the `-1` fallback) on PT-Post `anticipation>0` fits. (e) 'simple' relay conventions: `target="att"`, `n = n_treated_units + n_control_units` with `n_kind="units"` (DISJOINT by construction - `last_cohort` trimming reassigns before the counts, so a true total exists, unlike StackedDiD's overlapping sets), `df` = the post-overall snapshot (provenance-exact where `survey_metadata.df_survey` can diverge in the degenerate replicate state); 'group' relay: `n_kind="cells"`, `weight=None` (equal within-cohort weights, no cross-cohort mass), per-row `df_used` array captured at each row's `safe_inference` call (exact by construction; a stated divergence from CS's conservative-min scalar broadcast); 'event_study' rides the shared `_from_relative_dict` builder via a carrier whose `survey_metadata` copy carries the snapshot `df_survey`. --- diff --git a/docs/tutorials/15_efficient_did.ipynb b/docs/tutorials/15_efficient_did.ipynb index f5d8fea3..02f0c2c0 100644 --- a/docs/tutorials/15_efficient_did.ipynb +++ b/docs/tutorials/15_efficient_did.ipynb @@ -127,7 +127,7 @@ "source": [ "edid = EfficientDiD(pt_assumption=\"all\")\n", "results = edid.fit(data, outcome='outcome', unit='unit', time='period',\n", - " first_treat='first_treat', aggregate='all')\n", + " first_treat='first_treat')\n", "results.print_summary()" ] }, @@ -177,11 +177,11 @@ "# Fit under both assumptions\n", "results_all = EfficientDiD(pt_assumption=\"all\").fit(\n", " data, outcome='outcome', unit='unit', time='period',\n", - " first_treat='first_treat', aggregate='all')\n", + " first_treat='first_treat')\n", "\n", "results_post = EfficientDiD(pt_assumption=\"post\").fit(\n", " data, outcome='outcome', unit='unit', time='period',\n", - " first_treat='first_treat', aggregate='all')\n", + " first_treat='first_treat')\n", "\n", "# Compare with Callaway-Sant'Anna\n", "results_cs = CallawaySantAnna().fit(\n", @@ -269,9 +269,13 @@ "source": [ "edid_es = EfficientDiD(pt_assumption=\"all\")\n", "results_es = edid_es.fit(data, outcome='outcome', unit='unit', time='period',\n", - " first_treat='first_treat', aggregate='event_study')\n", + " first_treat='first_treat')\n", "\n", - "es_df = results_es.to_dataframe(level='event_study')\n", + "# Post-fit aggregation (3.9): recomputed from the retained EIFs - no refit.\n", + "# The container's to_dataframe() uses the unified schema\n", + "# (event_time / att / se / ... / is_reference).\n", + "es = results_es.aggregate('event_study')\n", + "es_df = es.to_dataframe()\n", "es_df" ] }, @@ -284,15 +288,15 @@ "source": [ "if HAS_MATPLOTLIB:\n", " fig, ax = plt.subplots(figsize=(10, 6))\n", - " ax.errorbar(es_df['relative_period'], es_df['effect'],\n", - " yerr=[es_df['effect'] - es_df['conf_int_lower'],\n", - " es_df['conf_int_upper'] - es_df['effect']],\n", + " ax.errorbar(es_df['event_time'], es_df['att'],\n", + " yerr=[es_df['att'] - es_df['conf_int_lower'],\n", + " es_df['conf_int_upper'] - es_df['att']],\n", " fmt='o-', capsize=4, color='steelblue', label='EDiD (PT-All)')\n", " ax.axhline(y=0, color='black', linestyle='--', linewidth=0.8)\n", " ax.axvline(x=-0.5, color='red', linestyle=':', linewidth=0.8, label='Treatment onset')\n", - " ax.set_xlabel('Relative Period (e = t - g)')\n", + " ax.set_xlabel('Event Time (e = t - g)')\n", " ax.set_ylabel('Effect')\n", - " ax.set_title('Efficient DiD Event Study')\n", + " ax.set_title('Efficient DiD Event Study (PT-All)')\n", " ax.legend()\n", " plt.tight_layout()\n", " plt.show()\n", @@ -317,7 +321,8 @@ "metadata": {}, "outputs": [], "source": [ - "grp_df = results.to_dataframe(level='group')\n", + "# Post-fit per-cohort aggregation (AggregationResult container):\n", + "grp_df = results.aggregate('group').to_dataframe()\n", "grp_df" ] }, @@ -346,7 +351,7 @@ "# Analytical vs bootstrap inference\n", "results_boot = EfficientDiD(pt_assumption=\"all\", n_bootstrap=499, seed=42).fit(\n", " data, outcome='outcome', unit='unit', time='period',\n", - " first_treat='first_treat', aggregate='all')\n", + " first_treat='first_treat')\n", "\n", "print(\"Analytical vs Bootstrap Inference\")\n", "print(\"=\" * 70)\n", @@ -428,11 +433,11 @@ "source": [ "r_no_antic = EfficientDiD(pt_assumption=\"all\").fit(\n", " data, outcome='outcome', unit='unit', time='period',\n", - " first_treat='first_treat', aggregate='all')\n", + " first_treat='first_treat')\n", "\n", "r_antic = EfficientDiD(pt_assumption=\"all\", anticipation=1).fit(\n", " data, outcome='outcome', unit='unit', time='period',\n", - " first_treat='first_treat', aggregate='all')\n", + " first_treat='first_treat')\n", "\n", "print(\"Anticipation Comparison\")\n", "print(\"=\" * 55)\n", @@ -470,7 +475,7 @@ "source": [ "edid_r = EfficientDiD(pt_assumption=\"all\").fit(\n", " data, outcome='outcome', unit='unit', time='period',\n", - " first_treat='first_treat', aggregate='all')\n", + " first_treat='first_treat')\n", "cs_r = CallawaySantAnna().fit(\n", " data, outcome='outcome', unit='unit', time='period',\n", " first_treat='first_treat')\n", @@ -497,19 +502,25 @@ "metadata": {}, "outputs": [], "source": [ - "# Side-by-side event study comparison\n", + "# Side-by-side event study comparison.\n", + "# EfficientDiD and CallawaySantAnna aggregate POST-FIT (3.9);\n", + "# ImputationDiD keeps the fit-time kwarg until its own migration.\n", "edid_es_r = EfficientDiD(pt_assumption=\"all\").fit(\n", " data, outcome='outcome', unit='unit', time='period',\n", - " first_treat='first_treat', aggregate='event_study')\n", + " first_treat='first_treat')\n", "cs_es_r = CallawaySantAnna().fit(\n", " data, outcome='outcome', unit='unit', time='period',\n", - " first_treat='first_treat', aggregate='event_study')\n", + " first_treat='first_treat')\n", "imp_es_r = ImputationDiD().fit(\n", " data, outcome='outcome', unit='unit', time='period',\n", " first_treat='first_treat', aggregate='event_study')\n", "\n", - "edid_es_df = edid_es_r.to_dataframe(level='event_study')\n", - "cs_es_df = cs_es_r.to_dataframe(level='event_study')\n", + "# The post-fit containers use the unified schema (event_time/att);\n", + "# rename to the native column names so one plotting block serves all\n", + "# three estimators.\n", + "_ren = {\"event_time\": \"relative_period\", \"att\": \"effect\"}\n", + "edid_es_df = edid_es_r.aggregate('event_study').to_dataframe().rename(columns=_ren)\n", + "cs_es_df = cs_es_r.aggregate('event_study').to_dataframe().rename(columns=_ren)\n", "imp_es_df = imp_es_r.to_dataframe(level='event_study')\n", "\n", "if HAS_MATPLOTLIB:\n", @@ -558,8 +569,9 @@ "6. **Multiplier bootstrap** provides robust inference with Rademacher, Mammen, or Webb weights\n", "7. **Condition numbers** flag potentially unstable weight matrices\n", "8. **Anticipation** shifts the effective treatment boundary for pre-treatment effects\n", - "9. Phase 1 is **no-covariates only** --- Phase 2 will add covariate support\n", - "10. When in doubt, run both EDiD and CS --- if ATTs agree, report EDiD for tighter CIs\n", + "9. **Covariates are supported** via the doubly robust path (sieve outcome regressions + propensity ratios) --- pass `covariates=[...]` to `fit()`\n", + "10. **Aggregate post-fit** (3.9): `results.aggregate('event_study'/'group'/'simple', balance_e=)` recomputes from retained EIFs --- no refit; fit-time `aggregate=`/`balance_e=` are deprecated, and bootstrapped fits keep the fit-time route\n", + "11. When in doubt, run both EDiD and CS --- if ATTs agree, report EDiD for tighter CIs\n", "\n", "**Parameter reference:**\n", "\n", diff --git a/docs/v4-deprecations.yaml b/docs/v4-deprecations.yaml index 74061748..c2faec89 100644 --- a/docs/v4-deprecations.yaml +++ b/docs/v4-deprecations.yaml @@ -261,11 +261,12 @@ rows: introduced_in: "3.9" deprecated_in: "3.9" removed_in: "4.0" - status: planned - phase: 2 + status: shimmed + phase: 5 warning: FutureWarning - code_refs: [diff_diff/efficient_did.py] - notes: "balance_e moves too." + test_ref: tests/test_aggregate_contract.py + code_refs: [diff_diff/efficient_did.py, diff_diff/efficient_did_aggregation.py, diff_diff/efficient_did_results.py, diff_diff/aggregation.py, diff_diff/results_base.py, diff_diff/honest_did.py, diff_diff/pretrends.py, diff_diff/practitioner.py, diff_diff/guides/llms-practitioner.txt] + notes: "Shimmed in 3.9: fit(aggregate=) warns via the shared NOT_SUPPLIED sentinel (a plain fit() never warns; supplying ANY value, None included, warns - CS-style joint warning with balance_e [M-120], warn-and-still-work since the params genuinely worked). NO fit-time value validation existed and none is added: unknown strings silently act like None on the deprecated path, unchanged; the post-fit successor fails closed on unknown types via the mixin vocabulary - a behavior improvement. The successor is a LAZY RECOMPUTING KIT (the CallawaySantAnna class, not a view relay): fit() computes nothing extra, the results object retains an AggregationKit referencing the per-(g,t) EIF dict (O(n_units x n_gt) dominant payload; full buffer enumeration in the REGISTRY EfficientDiD Note) plus O(n_units) bookkeeping, PRIVATE per-row snapshots of group_time_effects/groups/time_periods plus the pt_assumption/alpha/anticipation/n-total provenance (recompute and the ES carrier never read the mutable public fields), and the POST-OVERALL df_survey snapshot (captured before the ES/group gates - the group pass can degenerate the working df on replicate designs, and every aggregation seeds from the post-overall value), and aggregate('event_study'/'group', balance_e=) recomputes on demand while 'simple' relays the stored overall row bit-exact (n = treated+control units - disjoint by construction; df = the snapshot, provenance-exact where survey_metadata.df_survey can diverge in the degenerate n_valid<=1 replicate state). SUPPORTED SUBSET simple|event_study|group; calendar and 'all' fail closed via the mixin; weights= rejected. Bootstrap fits fail closed for ALL levels - ES/group because per-horizon draws are not retained (exact-replay wiring is a TODO row), 'simple' for uniform CS-parity conservatism BY DECISION (the stored overall row could relay, as dCDH's M-026 permits for its view-only aggregate, but EDiD keeps one uniform bootstrap rule so no level publishes analytical-provenance fields beside percentile inference). Group rows record per-row df_used at each safe_inference call (additive public row-dict key; the fit-time bootstrap override clears it - CS precedent); the post-fit group relay publishes the per-row array, a stated divergence from CS's conservative-min scalar broadcast. PT-Post reference provenance: the membership-gated reference_period property (SunAbraham rule - never synthesized when the anchor cell was not estimated) marks the materialized mechanical zero anchor is_reference in the container, and plot_event_study's inferred reference correctly shifts to -1-anticipation on PT-Post anticipation>0 fits. store_eif now governs only the public influence_functions field - the kit ALWAYS retains the EIF dict (memory-contract change; a store_kit opt-out is a DEFERRED row). Container admission NOT widened to EfficientDiD (see M-093); balance_e moves as its own row [M-120]." - id: M-024 kind: param group: aggregate-postfit @@ -985,7 +986,7 @@ rows: phase: 2 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). Amended pre-cut (same test_ref, introduced_in 3.9 not yet released): df became PER-ROW (one entry per event time, the df each stored p/CI actually used; joins the pinned to_dataframe schema) and StackedDiD/TwoStageDiD persist their internal full ES VCVs (event_study_vcov/_index/_df container fields; mode-gated for TwoStageDiD bootstrap/replicate). Completed pre-cut by the remaining producer channels: SunAbraham (per-event dict) and de Chaisemartin-D'Haultfoeuille (scalar) event_study_df, plus LPDiD pooled_df for the headline pre/post windows - every producer whose inference records a df now exposes it. Amended pre-cut a second time (2(b) PR-1, with M-026): three optional PROVENANCE fields appended last - base_period, anticipation, and df_survey (the fit's resolved SCALAR inference df beside the per-row channel: survey_metadata.df_survey with replicate-undefined mapping to the 0.0 fail-closed sentinel, else df_inference, else None - the per-row df column cannot encode that sentinel because __post_init__ NaNs it wherever p is non-finite) - threaded by the builders (the _empty_surface early return included) so the container consumers (HonestDiD's universal-base check, PreTrendsPower's anticipation cutoff, the honest df extraction) read fit-faithful values instead of dropping them. Amended pre-cut a third time (same PR): a fourth provenance field reference_event_times (also new on CallawaySantAnnaResults, computed at fit under base_period=universal) - the DISTINCT per-cohort positional-base event times, the common-reference signal is_reference cannot carry on gapped grids where a cohort's base overlaps another cohort's estimated horizon; HonestDiD and PreTrendsPower fail closed on more than one entry, on BOTH input routes (REGISTRY HonestDiD common-reference-guard Note)." + 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). Amended pre-cut (same test_ref, introduced_in 3.9 not yet released): df became PER-ROW (one entry per event time, the df each stored p/CI actually used; joins the pinned to_dataframe schema) and StackedDiD/TwoStageDiD persist their internal full ES VCVs (event_study_vcov/_index/_df container fields; mode-gated for TwoStageDiD bootstrap/replicate). Completed pre-cut by the remaining producer channels: SunAbraham (per-event dict) and de Chaisemartin-D'Haultfoeuille (scalar) event_study_df, plus LPDiD pooled_df for the headline pre/post windows - every producer whose inference records a df now exposes it. Amended pre-cut a second time (2(b) PR-1, with M-026): three optional PROVENANCE fields appended last - base_period, anticipation, and df_survey (the fit's resolved SCALAR inference df beside the per-row channel: survey_metadata.df_survey with replicate-undefined mapping to the 0.0 fail-closed sentinel, else df_inference, else None - the per-row df column cannot encode that sentinel because __post_init__ NaNs it wherever p is non-finite) - threaded by the builders (the _empty_surface early return included) so the container consumers (HonestDiD's universal-base check, PreTrendsPower's anticipation cutoff, the honest df extraction) read fit-faithful values instead of dropping them. Amended pre-cut a third time (same PR): a fourth provenance field reference_event_times (also new on CallawaySantAnnaResults, computed at fit under base_period=universal) - the DISTINCT per-cohort positional-base event times, the common-reference signal is_reference cannot carry on gapped grids where a cohort's base overlaps another cohort's estimated horizon; HonestDiD and PreTrendsPower fail closed on more than one entry, on BOTH input routes (REGISTRY HonestDiD common-reference-guard Note). SCOPE QUALIFIER (2(b) PR-3a, with M-023): the 'every producer whose inference records a df now exposes it' sentence is satisfied on EfficientDiD through the SCALAR df_survey channel only - its newly public aggregate('event_study') container has no per-row df source (no event_study_df/df_inference field; the per-row column is all-NaN, contract-permitted for a producer that records none); threading the retained kit scalar into the per-row channel is the TODO.md M-092-completion row." - id: M-093 kind: behavior group: results-contract @@ -996,7 +997,7 @@ rows: 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, diff_diff/honest_did.py, diff_diff/pretrends.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. PARTIAL PRE-CUT DELIVERY (2(b) PR-1, with M-026): the consumer half shipped for CS-SOURCED containers - compute_honest_did, compute_pretrends_power (both with SOURCE-SCOPED admission: source == CallawaySantAnnaResults only, rejecting dCDH l1 containers BY DESIGN since their placebo semantics need honest_did's native branch, and every other producer pending its own aggregate() migration) and plot_event_study / plot_honest_event_study (no source scoping - plotting is label-faithful). Admission widening is each later shim PR's methodology decision, not automatic. SECOND PRE-CUT AMENDMENT (2(b) PR-2, with M-024): admission widened to STACKED-SOURCED containers in compute_honest_did + compute_pretrends_power (source in {CallawaySantAnnaResults, StackedDiDResults}; kappa_pre >= 2 required for estimated pre-periods; honest additionally needs a non-singular full retained event-study covariance (pre+post sub-block); withheld-inference rows admitted with a source-scoped warning; the seven producer-derived guard messages now derive the producer from surface.source). Source scoping continues for all other producers. This row stays planned: its transition remains the 4.0 sentinel retirement + full-producer schema enforcement." + 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. PARTIAL PRE-CUT DELIVERY (2(b) PR-1, with M-026): the consumer half shipped for CS-SOURCED containers - compute_honest_did, compute_pretrends_power (both with SOURCE-SCOPED admission: source == CallawaySantAnnaResults only, rejecting dCDH l1 containers BY DESIGN since their placebo semantics need honest_did's native branch, and every other producer pending its own aggregate() migration) and plot_event_study / plot_honest_event_study (no source scoping - plotting is label-faithful). Admission widening is each later shim PR's methodology decision, not automatic. SECOND PRE-CUT AMENDMENT (2(b) PR-2, with M-024): admission widened to STACKED-SOURCED containers in compute_honest_did + compute_pretrends_power (source in {CallawaySantAnnaResults, StackedDiDResults}; kappa_pre >= 2 required for estimated pre-periods; honest additionally needs a non-singular full retained event-study covariance (pre+post sub-block); withheld-inference rows admitted with a source-scoped warning; the seven producer-derived guard messages now derive the producer from surface.source). Source scoping continues for all other producers. THIRD PRE-CUT AMENDMENT (2(b) PR-3a, with M-023): admission NOT widened to EfficientDiD-sourced containers BY DESIGN (the dCDH precedent) - the PRIMARY ground is the absent joint event-study covariance (container vcov=None, all-NaN per-row df); reference semantics are regime-dependent (PT-All: no reference row, universal first-period baseline; PT-Post: a materialized mechanical zero anchor at e=-1-anticipation when estimated, marked via the membership-gated reference_period property - never fabricated when absent); both terminal TypeErrors now state the EfficientDiD rejection explicitly; source scoping continues to reject. This row stays planned: its transition remains the 4.0 sentinel retirement + full-producer schema enforcement." # ---- Behavior policies (schema-tracked, spec-governed; no reality probe) - - id: M-080 @@ -1357,9 +1358,12 @@ rows: notes: "Fourth site of the 'robust' drop - spec section 7 says the flag dies 'everywhere it exists' but [M-045]..[M-047] enumerated only three. LinearRegression is a top-level export with its OWN __init__ (the TWFE/MultiPeriodDiD hits inherit DifferenceInDifferences.__init__ and ride [M-045]). Redundant with vcov_type, same as its siblings." # ---- Phase 2b PR 1: the post-fit aggregation surface --------------------- - # Ids M-116 and M-118..M-121 are intentionally UNUSED: they were drafted for - # the later 2b PRs (HAD rename, the other three balance_e sites, Wooldridge) - # and are reserved rather than reassigned, since ids are never reused. + # Reserved-id pool for the later 2b PRs (drafted with M-117): M-120 was + # claimed by the EfficientDiD balance_e row (2b PR-3a, the third balance_e + # slot matching the M-021/M-022/M-023 ordering). Ids M-116, M-118, M-119 + # and M-121 remain intentionally UNUSED (HAD rename, the Imputation and + # TwoStage balance_e sites, Wooldridge) - reserved rather than reassigned, + # since ids are never reused. - id: M-117 kind: param group: aggregate-postfit @@ -1373,7 +1377,21 @@ rows: warning: FutureWarning test_ref: tests/test_aggregate_contract.py code_refs: [diff_diff/staggered.py, diff_diff/staggered_results.py, diff_diff/aggregation.py] - notes: "balance_e moves from fit() onto aggregate() with [M-020]. Previously tracked only as the prose 'balance_e moves to aggregate() in the same PR' inside M-020's notes, which no test could assert - the same un-rowed-obligation class the gating-completeness amendment closed. Applies to event-study aggregation ONLY (the shipped code threads it nowhere else), so aggregate(type='simple'|'group', balance_e=...) raises rather than silently ignoring it. The other three balance_e sites (ImputationDiD, TwoStageDiD, EfficientDiD) get their own rows in the PRs that migrate them." + notes: "balance_e moves from fit() onto aggregate() with [M-020]. Previously tracked only as the prose 'balance_e moves to aggregate() in the same PR' inside M-020's notes, which no test could assert - the same un-rowed-obligation class the gating-completeness amendment closed. Applies to event-study aggregation ONLY (the shipped code threads it nowhere else), so aggregate(type='simple'|'group', balance_e=...) raises rather than silently ignoring it. The other three balance_e sites (ImputationDiD, TwoStageDiD, EfficientDiD) get their own rows in the PRs that migrate them - EfficientDiD's is [M-120]." + - id: M-120 + kind: param + group: aggregate-postfit + old: "diff_diff:EfficientDiD.fit[balance_e]" + new: "diff_diff:EfficientDiDResults.aggregate[balance_e]" + introduced_in: "3.9" + deprecated_in: "3.9" + removed_in: "4.0" + status: shimmed + phase: 5 + warning: FutureWarning + test_ref: tests/test_aggregate_contract.py + code_refs: [diff_diff/efficient_did.py, diff_diff/efficient_did_aggregation.py, diff_diff/efficient_did_results.py, diff_diff/aggregation.py] + notes: "balance_e moves from fit() onto aggregate() with [M-023] (joint FutureWarning; the M-117 twin). Applies to event-study aggregation only (the mixin default vocabulary), so aggregate(type='simple'|'group', balance_e=...) raises. EfficientDiD's balance rule is the ANCHOR-HORIZON rule - keep cohorts with a finite effect at e == balance_e, then retain all their horizons - the SAME rule CallawaySantAnna uses, divergent only from ImputationDiD/TwoStageDiD's balanced-window rule (their own rows document theirs). An anchor no cohort reaches warns (UserWarning) and yields a legal zero-row EventStudyResults container." - id: M-126 kind: behavior group: variance-cr1-convergence diff --git a/docs/v4-design.md b/docs/v4-design.md index 928a0958..cea81717 100644 --- a/docs/v4-design.md +++ b/docs/v4-design.md @@ -674,7 +674,7 @@ above; anything only one PR cares about stays in that PR's plan.** | 2: contract foundations | 3.9 | (a) results base + unified event-study representation [M-092] + to_dict completion + the Diagnostic marker base on the diagnostic result roster [M-091] (section 3.5); (b) `aggregate()` + fit(aggregate=) shims [M-020..M-027] (M-020's shim already shipped); (c) param renames [M-030..M-047] [M-084] [M-086..M-089] + their results-field mirrors [M-094] [M-095] (section 8 rule 9) + the public-function completeness sweep [M-097..M-113] (section 8 rule 10) + the dCDH results mirror [M-114] + the fourth `robust` site [M-115] + the 2(c)-ii missed-rename amendments [M-136..M-138] (LPDiD `level` value; the two post-dummy diagnostics params) + BaseEstimator mixin + ContinuousDiD covariates move; (d) alias introduction [M-062] (the Spillover introduction is cancelled [M-063]) + the alias-diet `__getattr__` warning shim [M-135] + wrapper deprecations [M-070..M-077] + the two inference-surface policies: `n_bootstrap` semantic unification [M-081] and the wild-cluster-bootstrap roster guard [M-096]; shipped insertions (all done): the aggregate contract [M-122], the ETWFE reference-period family [M-123] [M-124] [M-125], and the variance-consolidation program [M-126] [M-127] | | 3: merges | 3.9 | (a) TWFE event-study mode [M-010] + EventStudy warn [M-060] + the fit `time`->`post` rename [M-082] (gates: section 4.1's equivalence/divergence/pooled-parity test triple); (b) TripleDifference facade [M-013] + the SDDD alias [M-064]; (c) CiC method= [M-015] | | 4: release + soak | 3.9 cut | Migration guide written (skeleton: section 10); maintainer cuts 3.9; maint/3.8 rule active | -| 5: enforcement | 4.0 | Removals [M-010..M-015, M-020..M-027, M-030, M-032..M-047 old names, M-060, M-061, M-064, M-070..M-077, M-084, M-086..M-089, M-001..M-003, M-117] + the alias diet [M-132]..[M-134] + the amendment's old names [M-094] [M-095] [M-097..M-115] [M-136..M-138] (incl. their consumer migrations and the `clean_control` serialized reporting key); M-031's old `time` name persists as the merged class's calendar column, so it is deliberately absent from the removal roster (its 4.0 enforcement is the M-085 behavior entry below); property window: [M-016] property-flips at 4.0 (removal at 5.0); storage flips [M-050..M-058]; default policies [M-004..M-006, M-128..M-131, M-080]; merged-class behavior enforcements [M-083] [M-085]; warning retirement [M-007]; fastpath go/no-go [M-008]; diagnostic-family docs/roster reorganization [M-090]; sentinel retirement [M-093]; docs/llms.txt/README refresh | +| 5: enforcement | 4.0 | Removals [M-010..M-015, M-020..M-027, M-030, M-032..M-047 old names, M-060, M-061, M-064, M-070..M-077, M-084, M-086..M-089, M-001..M-003, M-117, M-120] + the alias diet [M-132]..[M-134] + the amendment's old names [M-094] [M-095] [M-097..M-115] [M-136..M-138] (incl. their consumer migrations and the `clean_control` serialized reporting key); M-031's old `time` name persists as the merged class's calendar column, so it is deliberately absent from the removal roster (its 4.0 enforcement is the M-085 behavior entry below); property window: [M-016] property-flips at 4.0 (removal at 5.0); storage flips [M-050..M-058]; default policies [M-004..M-006, M-128..M-131, M-080]; merged-class behavior enforcements [M-083] [M-085]; warning retirement [M-007]; fastpath go/no-go [M-008]; diagnostic-family docs/roster reorganization [M-090]; sentinel retirement [M-093]; docs/llms.txt/README refresh | | 6: front door | 4.1 | `event_study(data, outcome, unit, time, first_treat, estimator=...)` comparison entry point over the staggered family (sketch only; specified in its own plan) | Citation semantic for the table: a cell may cite a row whose current `phase` diff --git a/tests/test_aggregate_contract.py b/tests/test_aggregate_contract.py index d30720ee..06eb7d11 100644 --- a/tests/test_aggregate_contract.py +++ b/tests/test_aggregate_contract.py @@ -1110,6 +1110,7 @@ def test_mixin_hooks_are_not_dataclass_fields(self): assert "CallawaySantAnnaResults" in checked assert "ChaisemartinDHaultfoeuilleResults" in checked assert "StackedDiDResults" in checked + assert "EfficientDiDResults" in checked # --------------------------------------------------------------------------- # @@ -1463,3 +1464,484 @@ def test_weights_rejected(self, stacked_fitted): def test_unsupported_types_fail_closed(self, stacked_fitted, bad): with pytest.raises(ValueError, match="Unsupported aggregation type"): stacked_fitted.aggregate(bad) + + +# --------------------------------------------------------------------------- # +# EfficientDiD (rows M-023/M-120): fit(aggregate=/balance_e=) shim + the +# lazy-KIT recompute aggregate() (the CallawaySantAnna class) +# --------------------------------------------------------------------------- # + +EFFICIENT_KW = dict(outcome="outcome", unit="unit", time="period", first_treat="first_treat") + + +def _efficient_panel(seed=42, n_units=120, n_periods=8): + from diff_diff.prep_dgp import generate_staggered_data + + return generate_staggered_data( + n_units=n_units, n_periods=n_periods, cohort_periods=[4, 6], seed=seed + ) + + +def _efficient_clustered_panel(seed=11): + # Units nested 3-per-cluster (the sibling test_efficient_did.py shape). + d = _efficient_panel(seed=seed) + d = d.copy() + d["cl"] = (d["unit"] // 3).astype(int) + return d + + +def _efficient_survey_panel(seed=5, replicate=False, degenerate=None): + """Unit-constant pweights; optionally a JK replicate design. + + degenerate: None | "dropped" (2 all-zero replicate columns -> + n_valid = n_replicates - 2) | "undefined" (all but one column all-zero + -> n_valid <= 1 -> working df None, the load-bearing degenerate state). + """ + import numpy as np + + d = _efficient_panel(seed=seed).copy() + rng = np.random.default_rng(seed) + wmap = {u: rng.uniform(0.5, 2.0) for u in d["unit"].unique()} + d["w"] = d["unit"].map(wmap) + rep_cols = [] + if replicate: + n_rep = 8 + for r in range(n_rep): + col = f"rw{r}" + rep_cols.append(col) + if degenerate == "dropped" and r >= n_rep - 2: + d[col] = 0.0 + elif degenerate == "undefined" and r >= 1: + d[col] = 0.0 + else: + jitter = {u: rng.uniform(0.0, 2.0) for u in d["unit"].unique()} + d[col] = d["unit"].map(jitter) * d["w"] + return d, rep_cols + + +def _efficient_survey_design(rep_cols=None): + from diff_diff import SurveyDesign + + if rep_cols: + return SurveyDesign(weights="w", replicate_weights=rep_cols, replicate_method="JK1") + return SurveyDesign(weights="w") + + +def _fit_efficient(data, *, est_kw=None, **fit_kw): + from diff_diff import EfficientDiD + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return EfficientDiD(**(est_kw or {})).fit(data, **EFFICIENT_KW, **fit_kw) + + +@pytest.fixture(scope="module") +def efficient_panel(): + return _efficient_panel() + + +@pytest.fixture(scope="module") +def efficient_fitted(efficient_panel): + """Plain fit - computes NOTHING extra; the kit powers aggregate().""" + return _fit_efficient(efficient_panel) + + +@pytest.fixture(scope="module") +def efficient_fit_time(efficient_panel): + """Deprecated fit-time aggregate="all" - the inertness reference.""" + return _fit_efficient(efficient_panel, aggregate="all") + + +class TestEfficientShim: + def test_plain_fit_does_not_warn(self, efficient_panel): + from diff_diff import EfficientDiD + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + EfficientDiD().fit(efficient_panel, **EFFICIENT_KW) + assert [w for w in caught if issubclass(w.category, FutureWarning)] == [] + + def test_aggregate_kwarg_warns_even_at_none(self, efficient_panel): + from diff_diff import EfficientDiD + + with pytest.warns(FutureWarning, match=r"EfficientDiD\.fit\(aggregate=\)"): + EfficientDiD().fit(efficient_panel, **EFFICIENT_KW, aggregate=None) + + def test_balance_e_kwarg_warns_alone(self, efficient_panel): + from diff_diff import EfficientDiD + + with pytest.warns(FutureWarning, match=r"EfficientDiD\.fit\(balance_e=\)"): + EfficientDiD().fit(efficient_panel, **EFFICIENT_KW, balance_e=None) + + def test_joint_supply_warns_once_naming_both(self, efficient_panel): + from diff_diff import EfficientDiD + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + EfficientDiD().fit( + efficient_panel, **EFFICIENT_KW, aggregate="event_study", balance_e=0 + ) + fw = [w for w in caught if issubclass(w.category, FutureWarning)] + assert len(fw) == 1 + msg = str(fw[0].message) + assert "aggregate=" in msg and "balance_e=" in msg + + def test_unknown_string_still_acts_like_none(self, efficient_panel): + # The legacy path performs NO value validation - an unknown string + # behaves exactly like a plain fit (documented unchanged behavior; + # the post-fit mixin is what fails closed on unknown types). + res = _fit_efficient(efficient_panel, aggregate="nonsense") + plain = _fit_efficient(efficient_panel) + assert res.event_study_effects is None and res.group_effects is None + assert res.overall_att == plain.overall_att + + def test_warn_and_still_work(self, efficient_panel, efficient_fitted, efficient_fit_time): + # The deprecated path still populates the fit-time surfaces AND + # equals the plain fit's post-fit recompute. + assert efficient_fit_time.event_study_effects is not None + assert efficient_fit_time.group_effects is not None + es = efficient_fitted.aggregate("event_study") + for e, row in efficient_fit_time.event_study_effects.items(): + i = list(es.event_time).index(e) + np.testing.assert_allclose(row["effect"], es.att[i], rtol=1e-14) + + +class TestEfficientAggregate: + def _assert_es_matches_fit_time(self, res, fit_time, balance_e=None, skip_reference=False): + es = res.aggregate("event_study", balance_e=balance_e) + native = fit_time.event_study_effects + ref_marked = set() + if skip_reference: + ref_marked = {int(e) for e, m in zip(es.event_time, es.is_reference) if m} + assert sorted(int(e) for e in es.event_time) == sorted(native) + for e, row in native.items(): + i = list(es.event_time).index(e) + if int(e) in ref_marked: + # The container's reference normalization NaNs se/t/p/CI/n + # and forces att 0.0 on the marked row - compare the NATIVE + # dict instead (its anchor must be an exact mechanical zero). + assert row["effect"] == 0.0 and row["se"] == 0.0 + assert es.att[i] == 0.0 + assert np.isnan(es.se[i]) and np.isnan(es.n[i]) + continue + np.testing.assert_allclose(row["effect"], es.att[i], rtol=1e-14) + np.testing.assert_allclose(row["se"], es.se[i], rtol=1e-14) + if np.isfinite(row["p_value"]): + np.testing.assert_allclose(row["p_value"], es.p_value[i], rtol=1e-14) + return es + + def test_event_study_inertness(self, efficient_fitted, efficient_fit_time): + self._assert_es_matches_fit_time(efficient_fitted, efficient_fit_time) + + def test_group_inertness(self, efficient_fitted, efficient_fit_time): + grp = efficient_fitted.aggregate("group") + for i, g in enumerate(grp.label): + row = efficient_fit_time.group_effects[g] + np.testing.assert_allclose(row["effect"], grp.att[i], rtol=1e-14) + np.testing.assert_allclose(row["se"], grp.se[i], rtol=1e-14) + assert grp.n_kind == "cells" + assert grp.weight is None + + def test_simple_relay_bit_exact(self, efficient_fitted): + simple = efficient_fitted.aggregate("simple") + assert simple.att[0] == efficient_fitted.overall_att + assert simple.se[0] == efficient_fitted.overall_se + assert simple.t_stat[0] == efficient_fitted.overall_t_stat + assert simple.p_value[0] == efficient_fitted.overall_p_value + assert simple.conf_int_lower[0] == efficient_fitted.overall_conf_int[0] + assert simple.conf_int_upper[0] == efficient_fitted.overall_conf_int[1] + assert simple.target[0] == "att" + # Disjoint treated+control total (last_cohort trimming reassigns + # BEFORE the counts, so a true total exists - unlike StackedDiD). + assert simple.n[0] == float( + efficient_fitted.n_treated_units + efficient_fitted.n_control_units + ) + assert simple.n_kind == "units" + # Plain fit: the kit's post-overall df snapshot is None -> NaN df. + assert np.isnan(simple.df[0]) + + @pytest.mark.parametrize("balance_e", [0, 1, 2]) + def test_balance_e_inertness(self, efficient_panel, efficient_fitted, balance_e): + fit_time = _fit_efficient(efficient_panel, aggregate="event_study", balance_e=balance_e) + self._assert_es_matches_fit_time(efficient_fitted, fit_time, balance_e=balance_e) + + def test_cluster_inertness(self): + d = _efficient_clustered_panel() + plain = _fit_efficient(d, est_kw={"cluster": "cl"}) + fit_time = _fit_efficient(d, est_kw={"cluster": "cl"}, aggregate="all") + self._assert_es_matches_fit_time(plain, fit_time) + grp = plain.aggregate("group") + for i, g in enumerate(grp.label): + np.testing.assert_allclose(fit_time.group_effects[g]["se"], grp.se[i], rtol=1e-14) + + def test_survey_tsl_inertness_and_df(self): + d, _ = _efficient_survey_panel() + sd = _efficient_survey_design() + plain = _fit_efficient(d, survey_design=sd) + fit_time = _fit_efficient(d, survey_design=sd, aggregate="all") + es = self._assert_es_matches_fit_time(plain, fit_time) + # TSL survey: the post-overall snapshot is the finite survey df and + # is provenance-exact on every row. + assert es.df_survey == plain._aggregation_kit.bookkeeping["df_survey"] + simple = plain.aggregate("simple") + assert simple.df[0] == plain._aggregation_kit.bookkeeping["df_survey"] + grp = plain.aggregate("group") + # Non-circular oracle: safe_inference(att, se, df=row.df) reproduces + # each stored group row's inference (a stale df would mismatch p). + from diff_diff.utils import safe_inference + + for i in range(len(grp.label)): + t, p, ci = safe_inference( + float(grp.att[i]), + float(grp.se[i]), + alpha=plain.alpha, + df=float(grp.df[i]), + ) + np.testing.assert_allclose(t, grp.t_stat[i], rtol=1e-12) + np.testing.assert_allclose(p, grp.p_value[i], rtol=1e-12) + + def test_replicate_inertness_healthy_and_dropped(self): + from diff_diff.utils import safe_inference + + for degenerate, expected_df in ((None, 7.0), ("dropped", 5.0)): + d, rep_cols = _efficient_survey_panel(replicate=True, degenerate=degenerate) + sd = _efficient_survey_design(rep_cols) + plain = _fit_efficient(d, survey_design=sd) + fit_time = _fit_efficient(d, survey_design=sd, aggregate="all") + # Inertness against the SAME object's stored fit-time surface. + self._assert_es_matches_fit_time(plain, fit_time) + grp = plain.aggregate("group") + # Cross-FIT df contrast: all group rows share n_valid - 1. + assert set(np.unique(grp.df[np.isfinite(grp.df)])) <= {expected_df} + for i in range(len(grp.label)): + if not np.isfinite(grp.df[i]): + continue + t, p, ci = safe_inference( + float(grp.att[i]), + float(grp.se[i]), + alpha=plain.alpha, + df=float(grp.df[i]), + ) + np.testing.assert_allclose(p, grp.p_value[i], rtol=1e-12) + + def test_replicate_undefined_df_degenerate(self): + # n_valid <= 1: the working df degenerates to None mid-fit. The + # post-overall snapshot carries that state; the container's scalar + # df_survey resolves to the 0.0 replicate-undefined sentinel (the + # shared resolver ladder), NOT the raw survey_metadata value the + # is-not-None guard leaves finite - this arm is the discriminator + # for the carrier's metadata copy. + d, rep_cols = _efficient_survey_panel(replicate=True, degenerate="undefined") + sd = _efficient_survey_design(rep_cols) + plain = _fit_efficient(d, survey_design=sd) + assert plain._aggregation_kit.bookkeeping["df_survey"] is None + es = plain.aggregate("event_study") + assert es.df_survey == 0.0 + fit_time = _fit_efficient(d, survey_design=sd, aggregate="all") + self._assert_es_matches_fit_time(plain, fit_time) + + def test_pt_post_inertness_and_reference(self, efficient_panel): + plain = _fit_efficient(efficient_panel, est_kw={"pt_assumption": "post"}) + fit_time = _fit_efficient( + efficient_panel, est_kw={"pt_assumption": "post"}, aggregate="all" + ) + es = self._assert_es_matches_fit_time(plain, fit_time, skip_reference=True) + # Exactly one is_reference row, at the materialized mechanical anchor. + marked = [int(e) for e, m in zip(es.event_time, es.is_reference) if m] + assert marked == [-1] + # The NATIVE anchor is an exact mechanical zero BEFORE marking (the + # explicit-reference branch has no value check; post_init would + # silently rewrite a nonzero anchor, so the native pin is the guard). + assert fit_time.event_study_effects[-1]["effect"] == 0.0 + assert fit_time.event_study_effects[-1]["se"] == 0.0 + assert fit_time.event_study_effects[-1]["n_groups"] > 0 + + def test_pt_post_absent_anchor_not_synthesized(self): + # Single cohort baselined at the panel's first period: the anchor + # cell is never estimated, the membership gate returns None, and the + # container has NO reference row - nothing synthesized (the + # no-fabrication rule; results_base synthesis branch NOT taken). + import pandas as pd + + rng = np.random.default_rng(3) + rows = [] + for u in range(40): + coh = 1 if u < 20 else 0 + for t in range(6): + rows.append( + { + "unit": u, + "period": t, + "first_treat": coh, + "outcome": rng.normal() + (1.5 if coh == 1 and t >= 1 else 0.0), + } + ) + d = pd.DataFrame(rows) + plain = _fit_efficient(d, est_kw={"pt_assumption": "post"}) + assert plain.reference_period is None + es = plain.aggregate("event_study") + assert not es.is_reference.any() + assert -1 not in {int(e) for e in es.event_time} + + def test_pt_all_no_reference_row(self, efficient_fitted): + assert efficient_fitted.reference_period is None + es = efficient_fitted.aggregate("event_study") + assert not es.is_reference.any() + + def test_retention_no_dataframe_no_unit_labels(self, efficient_panel): + import pickle + + import pandas as pd + + d = efficient_panel.copy() + sentinel = {u: f"SENTINEL-ID-{u}@example.invalid" for u in d["unit"].unique()} + d["unit"] = d["unit"].map(sentinel) + res = _fit_efficient(d) + kit = res._aggregation_kit + assert kit is not None + for v in kit.bookkeeping.values(): + assert not isinstance(v, pd.DataFrame) + assert not isinstance(kit.influence, pd.DataFrame) + blob = pickle.dumps(res) + assert b"SENTINEL-ID" not in blob + # Pickle round-trip: aggregate() still works and matches. + res2 = pickle.loads(blob) + np.testing.assert_allclose( + res.aggregate("group").att, res2.aggregate("group").att, rtol=1e-15 + ) + + def test_immutability(self, efficient_fitted): + MUTATED = ("event_study_effects", "group_effects") + before = {f: getattr(efficient_fitted, f) for f in MUTATED} + for level in ("event_study", "group", "simple", "event_study", "group"): + efficient_fitted.aggregate(level) + for f in MUTATED: + assert getattr(efficient_fitted, f) is before[f] + # Repeated calls agree (no kit mutation between calls). + a = efficient_fitted.aggregate("group") + b = efficient_fitted.aggregate("group") + np.testing.assert_array_equal(a.att, b.att) + np.testing.assert_array_equal(a.df, b.df) + + def test_survey_metadata_not_mutated(self): + d, rep_cols = _efficient_survey_panel(replicate=True) + sd = _efficient_survey_design(rep_cols) + res = _fit_efficient(d, survey_design=sd) + before = res.survey_metadata.df_survey + res.aggregate("event_study") + res.aggregate("group") + assert res.survey_metadata.df_survey == before + + def test_bootstrap_fails_closed_all_levels(self, efficient_panel): + res = _fit_efficient(efficient_panel, est_kw={"n_bootstrap": 20, "seed": 1}) + for level in ("simple", "event_study", "group"): + with pytest.raises(NotImplementedError, match="bootstrap"): + res.aggregate(level) + + def test_bootstrap_fit_time_group_rows_clear_df_used(self, efficient_panel): + res = _fit_efficient( + efficient_panel, + est_kw={"n_bootstrap": 20, "seed": 1}, + aggregate="group", + ) + for row in res.group_effects.values(): + assert row.get("df_used") is None + + def test_weights_rejected(self, efficient_fitted): + with pytest.raises(ValueError, match="does not accept a weights selector"): + efficient_fitted.aggregate("event_study", weights="cell") + + @pytest.mark.parametrize("level", ["simple", "group"]) + def test_balance_e_rejected_on_non_event_study(self, efficient_fitted, level): + with pytest.raises(ValueError, match="balance_e is not used"): + efficient_fitted.aggregate(level, balance_e=1) + + @pytest.mark.parametrize("bad", ["calendar", "all", "nonsense"]) + def test_unsupported_types_fail_closed(self, efficient_fitted, bad): + with pytest.raises(ValueError, match="Unsupported aggregation type"): + efficient_fitted.aggregate(bad) + + def test_legacy_no_kit_raises(self, efficient_fitted): + import copy + + legacy = copy.copy(efficient_fitted) + object.__setattr__(legacy, "_aggregation_kit", None) + with pytest.raises(ValueError, match="aggregation kit"): + legacy.aggregate("event_study") + + def test_public_influence_functions_isolated_from_kit(self, efficient_panel): + # store_eif=True exposes the PUBLIC influence_functions diagnostic; + # it must be an independent COPY of the kit's canonical EIF payload, + # else a user mutation of the public field silently corrupts + # recomputed post-fit SEs/p-values/CIs (review pin, M-023). + res = _fit_efficient(efficient_panel, store_eif=True) + kit = res._aggregation_kit + assert res.influence_functions is not kit.influence + for gt, arr in res.influence_functions.items(): + assert not np.shares_memory(arr, kit.influence[gt]) + before = res.aggregate("event_study").se.copy() + first = next(iter(res.influence_functions)) + res.influence_functions[first][:] = 0.0 + after = res.aggregate("event_study").se + np.testing.assert_array_equal(before, after) + + def test_public_aggregation_inputs_isolated_from_kit(self, efficient_panel): + # CI review P0: aggregate() must recompute exclusively from the + # kit's PRIVATE snapshots - mutating the public group_time_effects + # rows or the groups/time_periods lists after fit must not change + # post-fit aggregation output (else altered point estimates mix + # with retained fit-time EIF variance: plausible-but-invalid + # inference). + res = _fit_efficient(efficient_panel) + g_before = res.aggregate("group") + es_before = res.aggregate("event_study") + post_gt = next(gt for gt in res.group_time_effects if gt[1] >= gt[0]) + res.group_time_effects[post_gt]["effect"] = 999.0 + res.groups.pop() + res.time_periods.pop() + # Provenance mutations (CI review R2): a flipped pt_assumption must + # not turn the genuine e=-1 estimate into a zeroed reference row, + # and alpha/anticipation edits must not relabel the containers. + res.pt_assumption = "post" + res.anticipation = 3 + res.alpha = 0.5 + g_after = res.aggregate("group") + es_after = res.aggregate("event_study") + s_after = res.aggregate("simple") + np.testing.assert_array_equal(g_before.att, g_after.att) + np.testing.assert_array_equal(g_before.se, g_after.se) + assert list(g_before.label) == list(g_after.label) + np.testing.assert_array_equal(es_before.att, es_after.att) + np.testing.assert_array_equal(es_before.se, es_after.se) + assert not es_after.is_reference.any() + assert es_after.alpha == es_before.alpha + assert s_after.alpha == 0.05 and float(s_after.n[0]) == float( + res.n_treated_units + res.n_control_units + ) + + def test_zero_row_balance_e_surface(self, efficient_fitted): + # An anchor no cohort reaches: warns, returns a LEGAL 0-row surface. + with pytest.warns(UserWarning, match="anchor horizon"): + es = efficient_fitted.aggregate("event_study", balance_e=99) + assert len(es.event_time) == 0 + + +class TestEfficientInternalCallers: + def test_hausman_pretest_emits_no_future_warning(self, efficient_panel): + # hausman_pretest refits internally; its fit_kwargs no longer pass + # aggregate=, so the shim must not fire (regression pin for the + # internal-caller cleanup - without it a re-added kwarg would warn + # spuriously from every hausman_pretest / DiagnosticReport._pt_hausman + # call with no test failing). + from diff_diff import EfficientDiD + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + EfficientDiD.hausman_pretest( + efficient_panel, + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + ) + assert [w for w in caught if issubclass(w.category, FutureWarning)] == [] diff --git a/tests/test_efficient_did.py b/tests/test_efficient_did.py index d2f058ac..8bfce357 100644 --- a/tests/test_efficient_did.py +++ b/tests/test_efficient_did.py @@ -212,25 +212,28 @@ class TestAggregation: def test_event_study_aggregation(self): df = _make_simple_panel() - result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="event_study") - assert result.event_study_effects is not None + result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") + es = result.aggregate("event_study") + assert es is not None # Should have pre and post-treatment event times - keys = sorted(result.event_study_effects.keys()) + keys = sorted(int(e) for e in es.event_time) assert any(e < 0 for e in keys), "Should have pre-treatment event times" assert any(e >= 0 for e in keys), "Should have post-treatment event times" def test_group_aggregation(self): df = _make_staggered_panel() - result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="group") - assert result.group_effects is not None - assert 3.0 in result.group_effects - assert 5.0 in result.group_effects + result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") + grp = result.aggregate("group") + assert grp is not None + labels = [float(g) for g in grp.label] + assert 3.0 in labels + assert 5.0 in labels def test_aggregate_all(self): df = _make_staggered_panel() - result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="all") - assert result.event_study_effects is not None - assert result.group_effects is not None + result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") + assert result.aggregate("event_study") is not None + assert result.aggregate("group") is not None class TestValidation: @@ -320,23 +323,24 @@ class TestOutputFormats: def test_summary_and_dataframe(self): df = _make_simple_panel() - result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="all") + result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") # summary() returns a string s = result.summary() assert isinstance(s, str) assert "Efficient DiD" in s - # to_dataframe at different levels + # to_dataframe at the native group-time level df_gt = result.to_dataframe("group_time") assert isinstance(df_gt, pd.DataFrame) assert "effect" in df_gt.columns - df_es = result.to_dataframe("event_study") - assert "relative_period" in df_es.columns + # aggregated tables come from the post-fit containers + df_es = result.aggregate("event_study").to_dataframe() + assert "event_time" in df_es.columns - df_g = result.to_dataframe("group") - assert "group" in df_g.columns + df_g = result.aggregate("group").to_dataframe() + assert "label" in df_g.columns def test_to_dataframe_raises_without_aggregation(self): df = _make_simple_panel() @@ -383,20 +387,20 @@ class TestPretreatment: def test_pretreatment_placebo_near_zero(self): """Under correct PT, pre-treatment ATT(g,t) for t < g should be near 0.""" df = _make_simple_panel(n_units=200, effect=2.0, sigma=0.3) - result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="event_study") + result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") + es = result.aggregate("event_study") # Check pre-treatment effects are near zero - for e, d in result.event_study_effects.items(): + for e, a in zip(es.event_time, es.att): if e < 0: - assert ( - abs(d["effect"]) < 1.0 - ), f"Pre-treatment effect at e={e} is {d['effect']:.4f}, expected ~0" + assert abs(a) < 1.0, f"Pre-treatment effect at e={e} is {a:.4f}, expected ~0" def test_pretreatment_in_event_study(self): """Placebo effects should appear with negative event-time keys.""" df = _make_simple_panel(n_periods=6, treat_period=3) - result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="event_study") - assert result.event_study_effects is not None - neg_keys = [e for e in result.event_study_effects if e < 0] + result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") + es = result.aggregate("event_study") + assert es is not None + neg_keys = [e for e in es.event_time if e < 0] assert len(neg_keys) > 0, "Should have negative event-time keys" def test_pretreatment_detects_violation(self): @@ -422,9 +426,10 @@ def test_pretreatment_detects_violation(self): "y": y, } ) - result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="event_study") + result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") + es = result.aggregate("event_study") # Pre-treatment effects should be significantly non-zero - pre_effects = [d["effect"] for e, d in result.event_study_effects.items() if e < 0] + pre_effects = [a for e, a in zip(es.event_time, es.att) if e < 0] assert any( abs(e) > 0.1 for e in pre_effects ), f"Pre-trend should be detected; pre effects: {pre_effects}" @@ -1460,23 +1465,20 @@ def test_clustered_se_at_least_as_large(self): def test_clustered_aggregate_event_study(self): """Clustered SE with aggregate='event_study' should produce finite results.""" df = self._make_clustered_panel(n_clusters=60, units_per_cluster=3) - result = EfficientDiD(cluster="cluster_id").fit( - df, "y", "unit", "time", "first_treat", aggregate="event_study" - ) - assert result.event_study_effects is not None - for e, d in result.event_study_effects.items(): - assert np.isfinite(d["se"]) + result = EfficientDiD(cluster="cluster_id").fit(df, "y", "unit", "time", "first_treat") + es = result.aggregate("event_study") + assert es is not None + assert np.isfinite(es.se[~es.is_reference]).all() def test_clustered_aggregate_all(self): """Clustered SE with aggregate='all' should produce finite results.""" df = self._make_clustered_panel(n_clusters=60, units_per_cluster=3) - result = EfficientDiD(cluster="cluster_id").fit( - df, "y", "unit", "time", "first_treat", aggregate="all" - ) - assert result.event_study_effects is not None - assert result.group_effects is not None - for g, d in result.group_effects.items(): - assert np.isfinite(d["se"]) + result = EfficientDiD(cluster="cluster_id").fit(df, "y", "unit", "time", "first_treat") + es = result.aggregate("event_study") + grp = result.aggregate("group") + assert es is not None + assert grp is not None + assert np.isfinite(grp.se).all() def test_cluster_bootstrap(self, ci_params): """Cluster bootstrap should produce finite inference.""" @@ -1561,9 +1563,12 @@ def test_clustered_bootstrap_aggregate_all(self, ci_params): """Clustered bootstrap with aggregate='all' should produce finite results.""" n_boot = ci_params.bootstrap(99) df = self._make_clustered_panel(n_clusters=60, units_per_cluster=3) - result = EfficientDiD(cluster="cluster_id", n_bootstrap=n_boot, seed=42).fit( - df, "y", "unit", "time", "first_treat", aggregate="all" - ) + # Bootstrapped fits keep the fit-time kwarg: post-fit aggregate() + # fails closed on n_bootstrap > 0. + with pytest.warns(FutureWarning): + result = EfficientDiD(cluster="cluster_id", n_bootstrap=n_boot, seed=42).fit( + df, "y", "unit", "time", "first_treat", aggregate="all" + ) assert result.event_study_effects is not None assert result.group_effects is not None for e, d in result.event_study_effects.items(): @@ -1773,12 +1778,12 @@ def test_last_cohort_aggregate_event_study(self): effects={3: 2.0, 5: 1.5, 7: 1.0}, ) result = EfficientDiD(control_group="last_cohort").fit( - df, "y", "unit", "time", "first_treat", aggregate="event_study" + df, "y", "unit", "time", "first_treat" ) - assert result.event_study_effects is not None + es = result.aggregate("event_study") + assert es is not None assert 7 not in result.groups - for e, d in result.event_study_effects.items(): - assert np.isfinite(d["effect"]) + assert np.isfinite(es.att).all() def test_last_cohort_aggregate_all(self): """last_cohort with aggregate='all' should produce finite results.""" @@ -1789,14 +1794,16 @@ def test_last_cohort_aggregate_all(self): effects={3: 2.0, 5: 1.5, 7: 1.0}, ) result = EfficientDiD(control_group="last_cohort").fit( - df, "y", "unit", "time", "first_treat", aggregate="all" + df, "y", "unit", "time", "first_treat" ) - assert result.event_study_effects is not None - assert result.group_effects is not None + es = result.aggregate("event_study") + grp = result.aggregate("group") + assert es is not None + assert grp is not None assert 7 not in result.groups - for g, d in result.group_effects.items(): + for g, a in zip(grp.label, grp.att): assert g != 7 - assert np.isfinite(d["effect"]) + assert np.isfinite(a) def test_last_cohort_bootstrap(self, ci_params): """last_cohort with bootstrap should produce finite inference.""" @@ -1821,32 +1828,27 @@ class TestBalanceE: def test_balance_e_basic(self): """balance_e restricts event study to cohorts present at anchor horizon.""" df = _make_staggered_panel(n_per_group=80, n_control=80, groups=(3, 5)) - result = EfficientDiD().fit( - df, - "y", - "unit", - "time", - "first_treat", - aggregate="event_study", - balance_e=0, - ) - assert result.event_study_effects is not None - for e, d in result.event_study_effects.items(): - assert np.isfinite(d["effect"]) + result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") + es = result.aggregate("event_study", balance_e=0) + assert es is not None + assert np.isfinite(es.att).all() def test_balance_e_with_bootstrap(self, ci_params): """Bootstrap balance_e should produce finite SEs.""" n_boot = ci_params.bootstrap(99) df = _make_staggered_panel(n_per_group=80, n_control=80, groups=(3, 5)) - result = EfficientDiD(n_bootstrap=n_boot, seed=42).fit( - df, - "y", - "unit", - "time", - "first_treat", - aggregate="event_study", - balance_e=0, - ) + # Bootstrapped fits keep the fit-time kwargs: post-fit aggregate() + # fails closed on n_bootstrap > 0. + with pytest.warns(FutureWarning): + result = EfficientDiD(n_bootstrap=n_boot, seed=42).fit( + df, + "y", + "unit", + "time", + "first_treat", + aggregate="event_study", + balance_e=0, + ) assert result.event_study_effects is not None for e, d in result.event_study_effects.items(): if np.isfinite(d["effect"]): @@ -1914,9 +1916,12 @@ def test_bootstrap_se_finite(self, ci_params): def test_bootstrap_with_aggregation(self, ci_params): n_boot = ci_params.bootstrap(99) df = _make_simple_panel() - result = EfficientDiD(n_bootstrap=n_boot, seed=42).fit( - df, "y", "unit", "time", "first_treat", aggregate="all" - ) + # Bootstrapped fits keep the fit-time kwarg: post-fit aggregate() + # fails closed on n_bootstrap > 0. + with pytest.warns(FutureWarning): + result = EfficientDiD(n_bootstrap=n_boot, seed=42).fit( + df, "y", "unit", "time", "first_treat", aggregate="all" + ) assert result.bootstrap_results is not None if result.event_study_effects: for e, d in result.event_study_effects.items(): @@ -1948,7 +1953,7 @@ class TestSimulationValidation: def test_synthetic_staggered_unbiased(self): """Single run at rho=0, verify ATT estimates near true values.""" df = _make_compustat_dgp(rho=0.0, seed=42) - result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="all") + result = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") # Check individual ATT(g,t) estimates # ATT(5,5) should be near 0.154 @@ -2479,12 +2484,11 @@ def test_covariates_aggregate_event_study(self): "time", "first_treat", covariates=["x1"], - aggregate="event_study", ) - assert result.event_study_effects is not None - assert len(result.event_study_effects) > 0 - for e, eff in result.event_study_effects.items(): - assert np.isfinite(eff["effect"]) + es = result.aggregate("event_study") + assert es is not None + assert len(es.event_time) > 0 + assert np.isfinite(es.att).all() def test_covariates_aggregate_group(self): df = _make_covariate_panel() @@ -2495,10 +2499,10 @@ def test_covariates_aggregate_group(self): "time", "first_treat", covariates=["x1"], - aggregate="group", ) - assert result.group_effects is not None - assert len(result.group_effects) > 0 + grp = result.aggregate("group") + assert grp is not None + assert len(grp.label) > 0 def test_covariates_aggregate_all(self): df = _make_covariate_panel() @@ -2509,10 +2513,9 @@ def test_covariates_aggregate_all(self): "time", "first_treat", covariates=["x1"], - aggregate="all", ) - assert result.event_study_effects is not None - assert result.group_effects is not None + assert result.aggregate("event_study") is not None + assert result.aggregate("group") is not None assert np.isfinite(result.overall_att) @@ -2639,15 +2642,18 @@ def test_bootstrap_with_covariates_smoke(self): def test_covariates_pt_all_bootstrap(self): """PT-All + bootstrap + covariates end-to-end.""" df = _make_covariate_panel(n_units=300) - result = EfficientDiD(pt_assumption="all", n_bootstrap=99, seed=42).fit( - df, - "y", - "unit", - "time", - "first_treat", - covariates=["x1"], - aggregate="all", - ) + # Bootstrapped fits keep the fit-time kwarg: post-fit aggregate() + # fails closed on n_bootstrap > 0. + with pytest.warns(FutureWarning): + result = EfficientDiD(pt_assumption="all", n_bootstrap=99, seed=42).fit( + df, + "y", + "unit", + "time", + "first_treat", + covariates=["x1"], + aggregate="all", + ) assert result.bootstrap_results is not None assert result.event_study_effects is not None assert result.group_effects is not None @@ -2971,8 +2977,12 @@ def test_default_hc1_bit_equal_baseline(self, aggregate): first_treat="first_treat", aggregate=aggregate, ) - r_default = EfficientDiD().fit(**common) - r_explicit = EfficientDiD(vcov_type="hc1").fit(**common) + # Subject is bit equality, not the deprecated fit-time shim - keep + # all four aggregate arms and silence the FutureWarning. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + r_default = EfficientDiD().fit(**common) + r_explicit = EfficientDiD(vcov_type="hc1").fit(**common) assert r_default.overall_att == r_explicit.overall_att assert r_default.overall_se == r_explicit.overall_se @@ -2989,8 +2999,10 @@ def test_cluster_hc1_bit_equal_baseline(self, aggregate): first_treat="first_treat", aggregate=aggregate, ) - r_default = EfficientDiD(cluster="state").fit(**common) - r_explicit = EfficientDiD(cluster="state", vcov_type="hc1").fit(**common) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + r_default = EfficientDiD(cluster="state").fit(**common) + r_explicit = EfficientDiD(cluster="state", vcov_type="hc1").fit(**common) assert r_default.overall_att == r_explicit.overall_att assert r_default.overall_se == r_explicit.overall_se @@ -3009,8 +3021,10 @@ def test_survey_tsl_hc1_bit_equal_baseline(self, aggregate): aggregate=aggregate, survey_design=design, ) - r_default = EfficientDiD().fit(**common) - r_explicit = EfficientDiD(vcov_type="hc1").fit(**common) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + r_default = EfficientDiD().fit(**common) + r_explicit = EfficientDiD(vcov_type="hc1").fit(**common) assert r_default.overall_att == r_explicit.overall_att assert r_default.overall_se == r_explicit.overall_se @@ -3034,8 +3048,10 @@ def test_survey_replicate_hc1_bit_equal_baseline(self, aggregate): survey_design=design, aggregate=aggregate, ) - r_default = EfficientDiD().fit(**common) - r_explicit = EfficientDiD(vcov_type="hc1").fit(**common) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + r_default = EfficientDiD().fit(**common) + r_explicit = EfficientDiD(vcov_type="hc1").fit(**common) assert r_default.overall_att == r_explicit.overall_att assert r_default.overall_se == r_explicit.overall_se # Per-horizon / per-group SE override branches must also agree under @@ -3067,10 +3083,12 @@ def test_bootstrap_cluster_hc1_bit_equal(self, ci_params): first_treat="first_treat", aggregate="all", ) - r_default = EfficientDiD(cluster="state", n_bootstrap=n_boot, seed=11).fit(**common) - r_explicit = EfficientDiD( - cluster="state", n_bootstrap=n_boot, seed=11, vcov_type="hc1" - ).fit(**common) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + r_default = EfficientDiD(cluster="state", n_bootstrap=n_boot, seed=11).fit(**common) + r_explicit = EfficientDiD( + cluster="state", n_bootstrap=n_boot, seed=11, vcov_type="hc1" + ).fit(**common) assert r_default.bootstrap_results is not None assert r_explicit.bootstrap_results is not None assert ( @@ -3101,8 +3119,10 @@ def test_bootstrap_survey_hc1_bit_equal(self, ci_params): survey_design=design, aggregate="all", ) - r_default = EfficientDiD(n_bootstrap=n_boot, seed=23).fit(**common) - r_explicit = EfficientDiD(n_bootstrap=n_boot, seed=23, vcov_type="hc1").fit(**common) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + r_default = EfficientDiD(n_bootstrap=n_boot, seed=23).fit(**common) + r_explicit = EfficientDiD(n_bootstrap=n_boot, seed=23, vcov_type="hc1").fit(**common) assert r_default.bootstrap_results is not None assert r_explicit.bootstrap_results is not None assert ( @@ -3196,7 +3216,6 @@ def test_dr_path_hc1_bit_equal(self): time="time", first_treat="first_treat", covariates=["x1"], - aggregate="event_study", ) r_default = EfficientDiD().fit(**common) r_explicit = EfficientDiD(vcov_type="hc1").fit(**common) diff --git a/tests/test_efficient_did_validation.py b/tests/test_efficient_did_validation.py index a20321eb..07909bed 100644 --- a/tests/test_efficient_did_validation.py +++ b/tests/test_efficient_did_validation.py @@ -8,6 +8,7 @@ "Efficient Difference-in-Differences and Event Study Estimators." """ +import warnings from pathlib import Path import numpy as np @@ -95,6 +96,21 @@ def _compute_es_avg(result): return np.mean(list(es.values())) +def _attach_post_fit_event_study(res): + """Recompute the event-study surface post-fit and hand it to the + mapping-consuming helpers/assertions in this module (``_compute_es_avg`` + and the dict-iterating tests). Post-fit ``aggregate()`` recomputes from + the retained kit and relays the same numbers as the deprecated fit-time + ``aggregate=`` surface. + """ + es = res.aggregate("event_study") + res.event_study_effects = { + int(e): {"effect": float(a), "se": float(s)} + for e, a, s in zip(es.event_time, es.att, es.se) + } + return res + + _TRUE_ES_AVG_COMPUSTAT = true_es_avg() @@ -116,8 +132,8 @@ def _run_mc_simulation(n_sims, rho, seed=1000, also_cs=False): unit="unit", time="time", first_treat="first_treat", - aggregate="all", ) + _attach_post_fit_event_study(res) edid_estimates.append(_compute_es_avg(res)) edid_overall_att.append(res.overall_att) edid_overall_se.append(res.overall_se) @@ -125,14 +141,18 @@ def _run_mc_simulation(n_sims, rho, seed=1000, also_cs=False): if also_cs: cs = CallawaySantAnna(control_group="never_treated") - cs_res = cs.fit( - data, - outcome="y", - unit="unit", - time="time", - first_treat="first_treat", - aggregate="event_study", - ) + # CS fit-time aggregate= is deprecated too; migrating CS is out + # of scope here, so silence its shim warning. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + cs_res = cs.fit( + data, + outcome="y", + unit="unit", + time="time", + first_treat="first_treat", + aggregate="event_study", + ) cs_estimates_list.append(_compute_es_avg(cs_res)) return { @@ -161,14 +181,14 @@ def hrs_data(): def edid_hrs_result(hrs_data): """Fit EDiD on HRS data (shared across tests).""" edid = EfficientDiD(pt_assumption="all") - return edid.fit( + res = edid.fit( hrs_data, outcome="outcome", unit="unit", time="time", first_treat="first_treat", - aggregate="all", ) + return _attach_post_fit_event_study(res) class TestHRSReplication: diff --git a/tests/test_event_study_consumers.py b/tests/test_event_study_consumers.py index c593b98d..6eb6acb9 100644 --- a/tests/test_event_study_consumers.py +++ b/tests/test_event_study_consumers.py @@ -1917,3 +1917,37 @@ def test_provenance_ladder(self, stacked_surface): ) ssurf = sres.aggregate("event_study") assert ssurf.df_survey == float(sres.survey_metadata.df_survey) + + +# --------------------------------------------------------------------------- # +# EfficientDiD containers (rows M-023/M-093): REJECTED BY DESIGN +# --------------------------------------------------------------------------- # + + +class TestEfficientContainerRejection: + """A REAL EfficientDiD post-fit container is rejected by both consumers. + + The pinned hand-built rejection (test_non_cs_e0_source_rejected) cannot + detect a missing message edit - "EfficientDiDResults" already matches + today's messages via the got source={...!r} interpolation - so this test + pins the NEW by-design clause text on a genuine aggregate() container. + """ + + def test_real_efficient_container_rejected_by_design(self): + from diff_diff import EfficientDiD + from diff_diff.prep_dgp import generate_staggered_data + + d = generate_staggered_data(n_units=80, n_periods=8, cohort_periods=[4, 6], seed=9) + res = EfficientDiD().fit( + d, + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + ) + surface = res.aggregate("event_study") + assert surface.source == "EfficientDiDResults" + with pytest.raises(TypeError, match="rejected BY DESIGN"): + compute_honest_did(surface, M=1.0) + with pytest.raises(TypeError, match="rejected BY DESIGN"): + compute_pretrends_power(surface, M=1.0) diff --git a/tests/test_guides.py b/tests/test_guides.py index bc934a6b..9f362121 100644 --- a/tests/test_guides.py +++ b/tests/test_guides.py @@ -862,3 +862,31 @@ def test_llms_full_lpdid_fit_signature_matches_real_api(self): f"LPDiD.fit() block in llms-full.txt is missing the real public " f"parameter {param!r} (adding a public param requires updating the guide)." ) + + +class TestLLMsFullEfficientDiDShimLine: + """M-023/M-120: the documented EfficientDiD fit signature must carry the + sentinel defaults + deprecation markers (the M-024 precedent) - no other + pin covers these lines, so they could silently go stale.""" + + def _efficient_section(self): + text = get_llm_guide("full") + start = text.index("### EfficientDiD") + nxt = text.index("\n### ", start + 1) + return text[start:nxt] + + def test_llms_full_efficient_fit_aggregate_line_documents_shim(self): + section = self._efficient_section() + fit_start = section.index("edid.fit(") + fit_block = section[fit_start : section.index("\n)", fit_start)] + agg_line = next( + line for line in fit_block.splitlines() if line.strip().startswith("aggregate") + ) + assert "NOT_SUPPLIED" in agg_line + assert "DEPRECATED (M-023)" in agg_line + assert "results.aggregate()" in agg_line + bal_line = next( + line for line in fit_block.splitlines() if line.strip().startswith("balance_e") + ) + assert "NOT_SUPPLIED" in bal_line + assert "DEPRECATED (M-120)" in bal_line diff --git a/tests/test_methodology_efficient_did.py b/tests/test_methodology_efficient_did.py index 5275f79b..4454120e 100644 --- a/tests/test_methodology_efficient_did.py +++ b/tests/test_methodology_efficient_did.py @@ -407,13 +407,13 @@ def test_fit_level_twin(self): ) with warnings.catch_warnings(): warnings.simplefilter("ignore") - r_gram = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="all") + r_gram = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") orig = ed_mod._omega_star_nocov_gram ed_mod._omega_star_nocov_gram = compute_omega_star_nocov try: with warnings.catch_warnings(): warnings.simplefilter("ignore") - r_loop = EfficientDiD().fit(df, "y", "unit", "time", "first_treat", aggregate="all") + r_loop = EfficientDiD().fit(df, "y", "unit", "time", "first_treat") finally: ed_mod._omega_star_nocov_gram = orig assert r_gram.overall_att == pytest.approx(r_loop.overall_att, rel=1e-10) @@ -1114,7 +1114,7 @@ def test_covariate_dr_path_last_cohort_with_anticipation(self): warnings.simplefilter("ignore") res = EfficientDiD( pt_assumption="all", control_group="last_cohort", anticipation=1 - ).fit(df, "y", "unit", "time", "first_treat", covariates=["x1"], aggregate="all") + ).fit(df, "y", "unit", "time", "first_treat", covariates=["x1"]) # anticipation=1, last_g=7 -> effective last cohort at 6 -> time_periods 1..5 assert max(res.time_periods) == 5 assert 7 not in res.groups # last cohort reclassified as pseudo-control @@ -1153,12 +1153,12 @@ def test_covariate_path_beats_forced_linear_under_nonlinear_nuisance(self, ci_pa warnings.simplefilter("ignore") a_sieve = ( EfficientDiD(pt_assumption="all") - .fit(df, "y", "unit", "time", "first_treat", covariates=["x1"], aggregate="all") + .fit(df, "y", "unit", "time", "first_treat", covariates=["x1"]) .overall_att ) a_linear = ( EfficientDiD(pt_assumption="all", sieve_k_max=1) - .fit(df, "y", "unit", "time", "first_treat", covariates=["x1"], aggregate="all") + .fit(df, "y", "unit", "time", "first_treat", covariates=["x1"]) .overall_att ) if np.isfinite(a_sieve): @@ -1386,7 +1386,6 @@ def fit(df): "first_treat", covariates=["x1"], survey_design=SurveyDesign(weights="w"), - aggregate="all", ) r_filtered = fit(base) @@ -1467,7 +1466,6 @@ def fit(df): "first_treat", covariates=["x1"], survey_design=SurveyDesign(weights="w"), - aggregate="all", ) r_filtered = fit(base) diff --git a/tests/test_naming_guard.py b/tests/test_naming_guard.py index 16278dc1..314d735f 100644 --- a/tests/test_naming_guard.py +++ b/tests/test_naming_guard.py @@ -48,7 +48,7 @@ FUNCTION rows additionally get a bare-call AST lane (an unqualified ``bacon_decompose(...)`` call has no dot, no quotes, no ``=``). Accepted limitations: rows whose old token equals the new token (the - API-move family M-020..M-027, M-084, M-117) are skipped - token lanes + API-move family M-020..M-027, M-084, M-117, M-120) are skipped - token lanes cannot distinguish old surface from new; their known readers are recorded in ``code_refs`` anyway. ``param-value`` rows (M-086) are excluded the same way. Rows whose ``deprecated_in`` window has not opened defer via the diff --git a/tests/test_practitioner.py b/tests/test_practitioner.py index 2faf79c0..4d88d9f4 100644 --- a/tests/test_practitioner.py +++ b/tests/test_practitioner.py @@ -535,6 +535,47 @@ def test_hausman_snippet_uses_classmethod(self, mock_efficient_results): assert len(hausman_steps) > 0 assert "hausman_pretest" in hausman_steps[0]["code"] + def _agg_step(self, output): + return [ + s + for s in output["next_steps"] + if "Aggregate treatment-effect heterogeneity" in s["label"] + ] + + def test_aggregation_step_post_fit_branch(self, mock_efficient_results): + # Analytical fit (no bootstrap_results attr on the mock -> None + # branch): the guidance recommends post-fit aggregate() (M-023). + output = practitioner_next_steps(mock_efficient_results, verbose=False) + steps = self._agg_step(output) + assert len(steps) == 1 + assert "results.aggregate('group')" in steps[0]["code"] + assert "no refit needed" in steps[0]["why"] + + def test_aggregation_step_bootstrap_branch(self, mock_efficient_results): + # Bootstrapped fit: post-fit aggregate() fails closed, so the + # guidance routes through the deprecated fit-time aggregation. + mock_efficient_results.bootstrap_results = object() + output = practitioner_next_steps(mock_efficient_results, verbose=False) + steps = self._agg_step(output) + assert len(steps) == 1 + assert "BOOTSTRAPPED" in steps[0]["why"] + assert "aggregate='all'" in steps[0]["code"] + + def test_aggregation_step_name_is_non_steps_key(self, mock_efficient_results): + # The "aggregation" key is deliberately OUTSIDE the STEPS + # completion vocabulary (the M-024 step-name-collision lesson): + # no DiagnosticReport check can auto-suppress this guidance, and + # a completed heterogeneity check must not swallow it. + from diff_diff.practitioner import STEPS + + assert "aggregation" not in STEPS + output = practitioner_next_steps( + mock_efficient_results, + completed_steps=["heterogeneity"], + verbose=False, + ) + assert len(self._agg_step(output)) == 1 + # --------------------------------------------------------------------------- # Tests: unknown result type fallback diff --git a/tests/test_survey_phase3.py b/tests/test_survey_phase3.py index 6e7748e4..acc83ec3 100644 --- a/tests/test_survey_phase3.py +++ b/tests/test_survey_phase3.py @@ -822,14 +822,14 @@ def test_survey_event_study_aggregation(self, staggered_survey_data): "unit", "time", "first_treat", - aggregate="event_study", survey_design=sd, ) - assert result.event_study_effects is not None - for e, eff in result.event_study_effects.items(): - assert np.isfinite(eff["effect"]) - assert np.isfinite(eff["se"]) - assert eff["se"] > 0 + es = result.aggregate("event_study") + assert es is not None + for att, se in zip(es.att, es.se): + assert np.isfinite(att) + assert np.isfinite(se) + assert se > 0 def test_survey_group_aggregation(self, staggered_survey_data): """EfficientDiD survey with aggregate='group' produces finite results.""" @@ -842,13 +842,13 @@ def test_survey_group_aggregation(self, staggered_survey_data): "unit", "time", "first_treat", - aggregate="group", survey_design=sd, ) - assert result.group_effects is not None - for g, eff in result.group_effects.items(): - assert np.isfinite(eff["effect"]) - assert np.isfinite(eff["se"]) + grp = result.aggregate("group") + assert grp is not None + for att, se in zip(grp.att, grp.se): + assert np.isfinite(att) + assert np.isfinite(se) def test_survey_all_aggregation(self, staggered_survey_data): """EfficientDiD survey with aggregate='all' produces finite results.""" @@ -861,11 +861,10 @@ def test_survey_all_aggregation(self, staggered_survey_data): "unit", "time", "first_treat", - aggregate="all", survey_design=sd, ) - assert result.event_study_effects is not None - assert result.group_effects is not None + assert result.aggregate("event_study") is not None + assert result.aggregate("group") is not None assert np.isfinite(result.overall_att) assert np.isfinite(result.overall_se) @@ -1040,17 +1039,18 @@ def test_aggregation_with_survey(self, cov_survey_data): "time", "first_treat", covariates=["x1"], - aggregate="all", survey_design=sd, ) - assert result.event_study_effects is not None - assert result.group_effects is not None - for _, eff in result.event_study_effects.items(): - assert np.isfinite(eff["effect"]) - assert np.isfinite(eff["se"]) - for _, eff in result.group_effects.items(): - assert np.isfinite(eff["effect"]) - assert np.isfinite(eff["se"]) + es = result.aggregate("event_study") + grp = result.aggregate("group") + assert es is not None + assert grp is not None + for att, se in zip(es.att, es.se): + assert np.isfinite(att) + assert np.isfinite(se) + for att, se in zip(grp.att, grp.se): + assert np.isfinite(att) + assert np.isfinite(se) def test_bootstrap_covariates_survey(self, cov_survey_data): """Bootstrap + covariates + survey produces finite results.""" @@ -1230,17 +1230,17 @@ def test_replicate_weight_aggregation(self): "time", "first_treat", covariates=["x1"], - aggregate="event_study", survey_design=sd, ) assert np.isfinite(result.overall_att) assert np.isfinite(result.overall_se) assert result.overall_se > 0 - assert result.event_study_effects is not None - for _, eff in result.event_study_effects.items(): - assert np.isfinite(eff["effect"]) - assert np.isfinite(eff["se"]) - assert eff["se"] > 0 + es = result.aggregate("event_study") + assert es is not None + for att, se in zip(es.att, es.se): + assert np.isfinite(att) + assert np.isfinite(se) + assert se > 0 # ============================================================================= diff --git a/tests/test_v4_matrix.py b/tests/test_v4_matrix.py index 9a1508fd..7132705d 100644 --- a/tests/test_v4_matrix.py +++ b/tests/test_v4_matrix.py @@ -122,11 +122,12 @@ # warning mechanism; EDiD was initially slated but RETAINED - it is the CSX # paper's own estimator label) = 114, plus the 2(c)-ii missed-rename # amendments (M-136 LPDiD level value; M-137/M-138 diagnostics time->post) -# = 117. +# = 117, plus 2b PR-3a's EfficientDiD balance_e row (M-120, claimed from the +# reserved pool) = 118. # 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 = 117 +ROW_COUNT_FLOOR = 118 # 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 @@ -150,9 +151,10 @@ # M-135 = the module __getattr__ FutureWarning mechanism that carries their # 3.9 warning, since the surviving target classes have no shim; EDiD was # initially slated but retained - the CSX paper's own estimator label). -# M-116 and -# M-118..M-121 are reserved for the later 2b PRs, not deleted - ids are -# never reused, so a gap here is intentional. +# (120,120) = +# 2b PR-3a's EfficientDiD balance_e move (claimed from the reserved pool). +# M-116, M-118, M-119 and M-121 remain reserved for the later 2b PRs, not +# deleted - ids are never reused, so a gap here is intentional. _INITIAL_ID_RANGES = [ (1, 8), (10, 16), @@ -166,6 +168,7 @@ (94, 96), (97, 115), (117, 117), + (120, 120), (122, 124), (125, 125), (126, 126), @@ -570,13 +573,13 @@ def test_initial_ids_never_deleted(): """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. - Extends as rows ship (117 as of the 2(c)-ii missed-rename amendments: + Extends as rows ship (118 as of 2b PR-3a's EfficientDiD balance_e row: Phase 1 + diagnostic-family + M-092/M-093 + M-094..M-096 + the M-097..M-115 - public-function completeness sweep + M-117/M-122 + M-123/M-124 + M-125 + - M-126 + M-127..M-131 + M-132..M-135 + M-136..M-138).""" + public-function completeness sweep + M-117/M-120/M-122 + M-123/M-124 + + M-125 + M-126 + M-127..M-131 + M-132..M-135 + M-136..M-138).""" missing = sorted(EXPECTED_INITIAL_IDS - set(_ROW_IDS)) assert not missing, f"ledger rows deleted (ids are permanent): {missing}" - assert len(EXPECTED_INITIAL_IDS) == 117 + assert len(EXPECTED_INITIAL_IDS) == 118 def test_version_tuple_pads_to_three_components(): diff --git a/tests/test_visualization.py b/tests/test_visualization.py index 475832e8..dc23fd69 100644 --- a/tests/test_visualization.py +++ b/tests/test_visualization.py @@ -745,3 +745,39 @@ def test_full_workflow_callaway_santanna(self): assert ax is not None plt.close() + + +class TestEfficientDiDPlotReference: + """M-023: the membership-gated ``reference_period`` property corrects the + plotted reference on PT-Post EfficientDiD fits. + + Before M-023 the native extraction fell back to ``-1`` for EDiD; under + ``pt_assumption="post"`` with ``anticipation=1`` the materialized + mechanical anchor sits at ``e = -2``, and the property now feeds the + extractor the true value. PT-All fits have no anchor (the property is + None) and keep the legacy ``-1`` fallback unchanged. + """ + + def test_pt_post_anticipation_reference_shifts(self): + pytest.importorskip("matplotlib") + import warnings + + from diff_diff import EfficientDiD + from diff_diff.prep_dgp import generate_staggered_data + from diff_diff.visualization._event_study import _extract_plot_data + + d = generate_staggered_data(n_units=80, n_periods=8, cohort_periods=[4, 6], seed=9) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + res = EfficientDiD(pt_assumption="post", anticipation=1).fit( + d, + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + aggregate="event_study", + ) + assert res.reference_period == -2 + extracted = _extract_plot_data(res, None, None, None, None) + assert extracted[5] == -2 # reference_period slot + assert extracted[6] is True # inferred, not caller-supplied