diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a129a3c..14614067 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,71 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- **ImputationDiD + TwoStageDiD post-fit `aggregate()` via panel-backed + recompute kits** (v4 program 2(b) PR-3b; ledger rows [M-021]/[M-022] + flip to shimmed, new rows [M-118]/[M-119] claim the reserved balance_e + slots, rows [M-092]/[M-093]/[M-127] amended). + - `ImputationDiDResults.aggregate('event_study'/'group'/'simple', + balance_e=)` and `TwoStageDiDResults.aggregate(...)` RECOMPUTE the + aggregations post-fit from panel-backed kits - no main-estimator + refit, no live estimator reference (replicate-weight fits are the + exception in COST, not correctness: their inference replay re-runs + the per-replicate refits each call, the same work fit-time paid); + post-fit results match the fit-time surfaces at + 1e-14 across plain/covariate/cluster/survey-TSL/replicate/pretrends + fits and every `balance_e` (`balance_e` keeps each estimator's + BALANCED-WINDOW rule, divergent from CS/EfficientDiD's + anchor-horizon rule). + - `fit(aggregate=, balance_e=)` deprecates on both estimators (3.9, + removed 4.0) with the shared CS-style joint FutureWarning: a plain + `fit()` never warns; supplying ANY value (None included) warns once; + the deprecated path still returns exactly the numbers it always did. + The `imputation_did`/`two_stage_did` wrappers forward the sentinel, + so plain wrapper calls never warn. Unknown `aggregate` strings keep + silently acting like None on the deprecated path; the post-fit + successor fails closed on unknown types - a behavior improvement. + - MEMORY CONTRACT: ImputationDiD's kit references the SAME per-fit + objects `_fit_data` already retains for `pretrend_test()` - zero + marginal memory, pickles unchanged via memoization. TwoStageDiD's + kit is the FIRST panel retention on its results: a column-subset + copy of the working frame (O(n_obs); replicate designs additionally + retain the O(n_obs x R) replicate matrix). A `store_kit` opt-out is + tracked in DEFERRED.md. + - Bootstrapped fits fail closed for ALL aggregate() levels including + 'simple' (uniform CS/EfficientDiD parity; replay wiring is tracked + in TODO.md); a fit whose bootstrap FAILED aggregates normally. + TwoStage post-fit event-study containers reproduce the M-092 + contract exactly (joint GMM vcov + index + df on analytical fits; + vcov=None with the replayed df on replicate fits). + - Replicate-weight fits replay the extracted override helpers with a + LEVEL-MATCHED joint stack: `aggregate(L)` reproduces + `fit(aggregate=L)` exactly. Documented migration delta: on + degenerate replicate designs (a replicate NaNs one family's targets) + the deprecated `fit(aggregate=)` coupled the public OVERALL row to + the joint stack, so plain-fit overall se/CI/df can differ there - + each surface is self-consistent (pinned in the contract tests). + - ImputationDiD's M-127 df_convention inert-config warning predicate + is now REACHABILITY-based (post-fit `aggregate('event_study')` + reaches the pretrends lead inference on analytical fits, so + `pretrends=True` fits no longer warn on the deprecated + `aggregate='group'` shape; bootstrapped/replicate pretrends fits + with no reachable lead inference now do). + - honest/pretrends container admission NOT widened: ImputationDiD + rejected BY DESIGN (no joint event-study covariance); TwoStageDiD + DEFERRED pending a normalization derivation (its pre-period + coefficients are stage-1 residual means, not reference-normalized + contrasts, while HonestDiD's Delta arithmetic hard-codes + delta_0=0) - both terminal TypeErrors state the grounds; the + paper-gated derivation is tracked in DEFERRED.md. + - `imputation.py`/`two_stage.py` shed their aggregation + variance + engines into new `imputation_aggregation.py` / + `two_stage_aggregation.py` modules (verbatim moves; two_stage.py + exits the "consider splitting" band). + - Fixed (TwoStageDiD, pre-existing corner surfaced by the migration + review): an event study where EVERY non-reference horizon is + Proposition-5-unidentified now retains the all-NaN Prop-5 rows + (`n_obs > 0`) and emits the consolidated warning instead of + returning a reference-only surface. - **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] diff --git a/DEFERRED.md b/DEFERRED.md index e1e002cf..9144a69c 100644 --- a/DEFERRED.md +++ b/DEFERRED.md @@ -21,6 +21,7 @@ provenance and AI-review deviation-documentation: a row here (or in | Issue | Location | PR | Priority | |-------|----------|----|----------| | `PlaceboTests` `boundary_gap` — a permutation randomization-inference margin (SE-audit item (b)); NOT computed anywhere in code today, so this is a new feature + result field, not a coverage lock. **User-locked 2026-07-09: defer until a derivation/paper source exists** — do not design or implement from scratch. | `tests/test_methodology_placebo.py`, `diff_diff/diagnostics.py` | SE-audit | Low | +| TwoStageDiD honest/pretrends container admission DEFERRED (decision revised from "widen" during the 2(b) PR-3b plan review): analytical fits carry the joint Gardner-GMM event-study covariance (M-092), but the pre-period coefficients are stage-1 residual MEANS — the reference horizon is dropped from the no-intercept Stage-2 design and the zero anchor row appended mechanically — not contrasts against a reference period, while HonestDiD's Δ^RM/Δ^SD arithmetic hard-codes the `delta_0 = 0` normalization into its boundary/bridge constraints. Admission needs either a Stage-2 re-estimation with the reference horizon in the design or a derived residual-to-reference normalization mapping (+ its variance transform). Both consumers' TypeErrors state the deferral; see the REGISTRY TwoStageDiD Note (d). | `diff_diff/honest_did.py`, `diff_diff/pretrends.py`, `diff_diff/two_stage_aggregation.py` | 2(b) PR-3b | Low | | CBWSDID covariate balancing (`StackedDiD(balance="entropy")`) v1 supports only balanced event windows + `weighting="aggregate"`; unbalanced/ragged panels fail closed (unit-count vs observation-count corrector convention unresolved off balanced panels). Matching-based balancing and the repeated `0→1`/`1→0` episode extension are also deferred. Documented in REGISTRY StackedDiD "Covariate balancing (CBWSDID)" Notes. | `stacked_did.py`, `balancing.py`, REGISTRY | follow-up | Low | | dCDH: Phase-1 per-period placebo `DID_M^pl` has NaN SE (no IF derivation for the per-period aggregation path). Multi-horizon placebos (`L_max ≥ 1`) have valid SE. | `chaisemartin_dhaultfoeuille.py` | #294 | Low | | 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 | @@ -69,7 +70,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 | +| EfficientDiD/ImputationDiD/TwoStageDiD kit retention has no memory opt-out: EfficientDiD's kit always retains the per-(g,t) EIF dict (M-023), ImputationDiD's references the `_fit_data` panel objects (M-021, pre-existing retention), and TwoStageDiD's retains a column-subset working-frame copy (M-022, the first new panel retention); 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`, `diff_diff/imputation.py`, `diff_diff/two_stage.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/METHODOLOGY_REVIEW.md b/METHODOLOGY_REVIEW.md index e17fca4d..57253492 100644 --- a/METHODOLOGY_REVIEW.md +++ b/METHODOLOGY_REVIEW.md @@ -580,7 +580,7 @@ and covariate-adjusted specifications.) | Field | Value | |-------|-------| -| Module | `two_stage.py`, `two_stage_bootstrap.py` | +| Module | `two_stage.py`, `two_stage_aggregation.py`, `two_stage_bootstrap.py` | | Primary Reference | Gardner (2022), *Two-stage differences in differences*, arXiv:2207.05943 | | R Reference | `did2s` | | Status | **Complete** | diff --git a/README.md b/README.md index 67c8cc48..2e8728da 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); EDiD: `results.aggregate(...)` post-fit from retained EIFs (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); ImputationDiD/TwoStageDiD: `results.aggregate(...)` post-fit from panel-backed kits (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 445133bb..fc72170d 100644 --- a/TODO.md +++ b/TODO.md @@ -22,7 +22,9 @@ 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 | +| ImputationDiD/TwoStageDiD `aggregate()` on bootstrapped fits fails closed (M-021/M-022); ImputationDiD's per-target psi machinery makes seeded replay tractable (the panel-backed kit retains everything the psi precompute reads), TwoStageDiD's per-level GMM scores are function-locals and would need retention | `diff_diff/imputation_results.py`, `diff_diff/two_stage_results.py`, `diff_diff/aggregation.py` | 2(b) PR-3b | Mid | Low | +| diagnostic_report's ES-gated checks read the raw `event_study_effects` field, which post-fit `results.aggregate()` never populates - their remediation strings steer users to the deprecated fit-time kwarg (qualified "deprecated but functional until 4.0" since 2(b) PR-3b); teach the checks to consume a post-fit container (or recompute via the kit) before 4.0 removes the kwarg | `diff_diff/diagnostic_report.py` | 2(b) PR-3b | Mid | Medium | +| EfficientDiD and ImputationDiD are the outstanding M-092 event-study df-provenance holes: 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 kits now retain the scalar (ImputationDiD's since 2(b) PR-3b - same shape: scalar `df_survey` exposed, per-row column all-NaN, identical to its fit-time surface); threading it into the per-row channel is a contained follow-up | `diff_diff/efficient_did_results.py`, `diff_diff/imputation_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/_reporting_helpers.py b/diff_diff/_reporting_helpers.py index 9e95788d..647c1cf0 100644 --- a/diff_diff/_reporting_helpers.py +++ b/diff_diff/_reporting_helpers.py @@ -48,12 +48,13 @@ def describe_target_parameter(results: Any) -> Dict[str, Any]: ``feedback_verify_claims.md``). All wording choices are deliberate: - - ``ImputationDiD`` / ``TwoStageDiD``: the ``aggregate`` fit-time - kwarg controls which horizon / group tables get populated but - does NOT change ``overall_att``. The headline is always the - sample-mean overall ATT (per BJS 2024 Step 3 with - ``w_it = 1/N_1``); disambiguate via the event-study or group - aggregate if you need the horizon / group target. + - ``ImputationDiD`` / ``TwoStageDiD``: horizon / group tables come + from post-fit ``results.aggregate('event_study'/'group')`` since + 3.9 (the fit-time ``aggregate`` kwarg is deprecated) and never + change ``overall_att``. The headline is always the sample-mean + overall ATT (per BJS 2024 Step 3 with ``w_it = 1/N_1``); + disambiguate via the event-study or group aggregate if you need + the horizon / group target. - ``CallawaySantAnna``: ``overall_att`` is cohort-size-weighted across post-treatment ``ATT(g, t)`` cells regardless of the fit-time ``aggregate`` kwarg. The event-study / group @@ -167,11 +168,12 @@ def describe_target_parameter(results: Any) -> Dict[str, Any]: "tau_hat_it`` across treated observations, where " "``tau_hat_it = Y_it - Y_hat_it(0)`` and ``Y_hat_it(0)`` is " "imputed from a unit+time fixed-effects model fitted on " - "untreated observations only (BJS 2024 Step 3). The fit-time " - "``aggregate`` kwarg populates additional horizon / group " - "tables but does NOT change ``overall_att`` — for the " - "horizon or group estimand, consult the event-study or group " - "aggregate directly." + "untreated observations only (BJS 2024 Step 3). Post-fit " + "``results.aggregate('event_study'/'group')`` (3.9; the " + "fit-time ``aggregate`` kwarg is deprecated) yields the " + "horizon / group tables and does NOT change ``overall_att`` " + "— for the horizon or group estimand, consult the " + "event-study or group aggregate directly." ), "aggregation": "simple", "headline_attribute": "overall_att", @@ -188,8 +190,9 @@ def describe_target_parameter(results: Any) -> Dict[str, Any]: "the residualized outcome on the treatment indicator across " "treated observations. Point estimate is algebraically " "equivalent to Borusyak-Jaravel-Spiess imputation. As with " - "ImputationDiD, the fit-time ``aggregate`` kwarg populates " - "additional tables but does NOT change ``overall_att``." + "ImputationDiD, post-fit ``results.aggregate(...)`` (3.9) " + "yields the additional tables and does NOT change " + "``overall_att``." ), "aggregation": "simple", "headline_attribute": "overall_att", diff --git a/diff_diff/aggregation.py b/diff_diff/aggregation.py index 34bcf6a5..e9f55599 100644 --- a/diff_diff/aggregation.py +++ b/diff_diff/aggregation.py @@ -376,6 +376,19 @@ class AggregationKit: ``covariate_matrix``, ``obs_outcome``, ``obs_covariates``): re-aggregation reads only unit-level bookkeeping, so the source panel is never retained. + PANEL-BACKED EXCEPTION (ImputationDiD [M-021] and TwoStageDiD [M-022]): + their recompute is target-specific - a different ``balance_e`` re-masks + which observations enter each estimand and re-solves the variance from + the panel + FE model - so no compact influence payload can replace the + frame, and their kits' ``bookkeeping`` DOES retain panel objects. + ImputationDiD references the SAME per-fit objects ``_fit_data`` already + retains for ``pretrend_test()`` (zero marginal memory; pickles unchanged + via memoization). TwoStageDiD retains a column-subset copy of its + working frame - the first new panel retention, O(n_obs), documented in + its ledger row and REGISTRY Note. For those two, ``influence`` is empty + by design and the exclusion above applies to everything OUTSIDE the + enumerated bookkeeping payload. + Attributes ---------- bookkeeping : dict diff --git a/diff_diff/diagnostic_report.py b/diff_diff/diagnostic_report.py index bb7f0275..33c35b75 100644 --- a/diff_diff/diagnostic_report.py +++ b/diff_diff/diagnostic_report.py @@ -788,7 +788,10 @@ def _instance_skip_reason(self, check: str) -> Optional[str]: return ( "No pre-period event-study coefficients are exposed on " "this fit. For staggered estimators, re-fit with " - "aggregate='event_study' to populate event-study output." + "aggregate='event_study' to populate event-study " + "output (deprecated but functional until 4.0 - these " + "checks read the raw event_study_effects field, which " + "post-fit results.aggregate() does not populate)." ) # vcov is optional for the Bonferroni fallback. if method == "hausman": @@ -846,7 +849,10 @@ def _instance_skip_reason(self, check: str) -> Optional[str]: return ( "Pre-trends power needs either results.vcov or " "event_study_effects (from aggregate='event_study' on " - "staggered estimators); neither available." + "staggered estimators - deprecated but functional until " + "4.0; these checks read the raw field, which post-fit " + "results.aggregate() does not populate); neither " + "available." ) pre_coefs, _ = _collect_pre_period_coefs(r) if len(pre_coefs) < 2: diff --git a/diff_diff/guides/llms-full.txt b/diff_diff/guides/llms-full.txt index 15bfb4e4..b3ff9754 100644 --- a/diff_diff/guides/llms-full.txt +++ b/diff_diff/guides/llms-full.txt @@ -442,8 +442,8 @@ imp.fit( time: str, first_treat: str, covariates: list[str] = None, - aggregate: str = None, # None, "simple", "event_study", "group", or "all" - balance_e: int = None, + aggregate: Any = NOT_SUPPLIED, # DEPRECATED (M-021): warns, still works - use results.aggregate() post-fit + balance_e: Any = NOT_SUPPLIED, # DEPRECATED (M-118): moves onto results.aggregate(balance_e=) survey_design: SurveyDesign = None, # Optional design-based inference (pweight + analytical strata/PSU/FPC or replicate BRR/Fay/JK1/JKn/SDR) ) -> ImputationDiDResults ``` @@ -455,10 +455,10 @@ from diff_diff import ImputationDiD, plot_event_study est = ImputationDiD() results = est.fit(data, outcome='outcome', unit='unit', - time='period', first_treat='first_treat', - aggregate='event_study') + time='period', first_treat='first_treat') results.print_summary() -plot_event_study(results) +es = results.aggregate('event_study') # post-fit (M-021); balance_e= lives here too +plot_event_study(es) ``` ### TwoStageDiD @@ -491,8 +491,8 @@ ts.fit( time: str, first_treat: str, covariates: list[str] = None, - aggregate: str = None, # None, "simple", "event_study", "group", or "all" - balance_e: int = None, + aggregate: Any = NOT_SUPPLIED, # DEPRECATED (M-022): warns, still works - use results.aggregate() post-fit + balance_e: Any = NOT_SUPPLIED, # DEPRECATED (M-119): moves onto results.aggregate(balance_e=) ) -> TwoStageDiDResults ``` @@ -505,6 +505,7 @@ est = TwoStageDiD() results = est.fit(data, outcome='outcome', unit='unit', time='period', first_treat='first_treat') results.print_summary() +es = results.aggregate('event_study') # post-fit (M-022); balance_e= lives here too ``` ### SpilloverDiD @@ -1784,7 +1785,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`); `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) +**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; `ImputationDiDResults` and `TwoStageDiDResults` recompute from their PANEL-BACKED kits since 3.9/M-021/M-022 - bootstrapped fits raise on all three recomputing adopters) ### ContinuousDiDResults diff --git a/diff_diff/guides/llms-practitioner.txt b/diff_diff/guides/llms-practitioner.txt index e0fb3ab1..5434bf73 100644 --- a/diff_diff/guides/llms-practitioner.txt +++ b/diff_diff/guides/llms-practitioner.txt @@ -417,16 +417,19 @@ estimate is contaminated by forbidden comparisons (see Step 4). Aggregate treatment effects may mask important variation. -### For CallawaySantAnna and EfficientDiD (post-fit `aggregate()`) +### For CallawaySantAnna, EfficientDiD, ImputationDiD, TwoStageDiD (post-fit `aggregate()`) ```python # Post-fit, no refit needed (fit-time aggregate= is deprecated since 3.9 -# for CS/StackedDiD/EfficientDiD; removed in 4.0): +# for CS/StackedDiD/EfficientDiD/ImputationDiD/TwoStageDiD; 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: +# balance_e moves onto aggregate() too: +print(results.aggregate('event_study', balance_e=2).to_dataframe()) +# EXCEPTION: a BOOTSTRAPPED fit's aggregate() raises on the recomputing +# adopters (CS/EfficientDiD/ImputationDiD/TwoStageDiD) — 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 fd097fbd..4c1bb284 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`); 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. +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); BJS/Gardner: `results.aggregate('event_study')`/`.aggregate('group')`/`.aggregate('simple')` post-fit on ImputationDiD and TwoStageDiD too (3.9, rows M-021/M-022; recomputed from panel-backed kits, `balance_e=` on `aggregate('event_study')`; bootstrapped fits raise - use the deprecated fit-time aggregation there; their containers are not admitted into honest/pretrends - Imputation by design, TwoStage deferred pending a normalization derivation); 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 3c004d8e..7cc4d5fe 100644 --- a/diff_diff/honest_did.py +++ b/diff_diff/honest_did.py @@ -609,11 +609,16 @@ def _extract_container_params( "the native results object where supported " "(MultiPeriodDiDResults, CallawaySantAnnaResults, or " "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." + "EfficientDiDResults and ImputationDiDResults containers are " + "rejected BY DESIGN (their surfaces carry no joint event-study " + "covariance - per-horizon SEs only; see the REGISTRY " + "EfficientDiD and ImputationDiD Notes); TwoStageDiDResults " + "container admission is DEFERRED pending a normalization " + "derivation - analytical fits carry the joint Gardner-GMM " + "covariance, but the pre-period coefficients are stage-1 " + "residual means, not contrasts against a reference period as " + "HonestDiD's delta_0=0 arithmetic requires (see the REGISTRY " + "TwoStageDiD Note and DEFERRED.md)." ) _producer = surface.source.replace("Results", "") if surface.time_scale != "relative": diff --git a/diff_diff/imputation.py b/diff_diff/imputation.py index d8e2e8b2..2d06b2f7 100644 --- a/diff_diff/imputation.py +++ b/diff_diff/imputation.py @@ -23,30 +23,31 @@ IF-vs-sandwich taxonomy. """ +import dataclasses import warnings -from typing import TYPE_CHECKING, Any, Callable, Dict, List, NamedTuple, Optional, Set, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional import numpy as np import pandas as pd -from scipy import sparse, stats -from scipy.sparse.linalg import factorized as sparse_factorized +from scipy import stats from diff_diff._base import BaseEstimator -from diff_diff.imputation_bootstrap import ImputationDiDBootstrapMixin, _compute_target_weights +from diff_diff._deprecation import NOT_SUPPLIED +from diff_diff.aggregation import AggregationKit +from diff_diff.imputation_aggregation import ( # noqa: F401 (compat re-exports) + _compute_target_weights, + _ImputationAggregationMixin, + _lsmr_minnorm_normal_solve, + _LSMRUnconvergedError, + _UntreatedProjection, +) +from diff_diff.imputation_bootstrap import ImputationDiDBootstrapMixin from diff_diff.imputation_results import ( # noqa: F401 (re-export) ImputationBootstrapResults, ImputationDiDResults, ) -from diff_diff.linalg import solve_ols from diff_diff.utils import ( - _iterative_fe_solve, - absorbed_fe_cr1_k_increment, - absorbed_fe_rank, - demean_by_groups, - pre_demean_norms, - resolve_tail_df, safe_inference, - snap_absorbed_regressors, validate_df_convention, ) @@ -54,99 +55,12 @@ from diff_diff.survey import SurveyDesign -class _UntreatedProjection(NamedTuple): - """Cached, target-invariant pieces of the untreated imputation projection - ``v_untreated = -A_0 (A_0' [W] A_0)^{-1} A_1' w`` (BJS 2024 Theorem 3). - - Within a single ``fit()`` the untreated design (``df_0``/``df_1``, covariates, - survey weights) is identical across every estimand target (overall ATT, each - event-study horizon, each group, and the bootstrap precompute) -- only the - treated aggregation ``weights`` (the RHS ``A_1' w``) vary. So ``A_0``, ``A_1`` - and the factorization of ``A_0'[W]A_0`` are built once and reused across - targets (factorize-once / solve-many), mirroring the TwoStageDiD GMM-sandwich - ``sparse_factorized`` pattern. - """ - - A_0: sparse.csr_matrix - A_1: sparse.csr_matrix - # solver(rhs) -> z; None when the factorization was exactly singular (the - # solve path then routes to the sparse LSMR least-squares fallback). - solver: Optional[Callable[[np.ndarray], np.ndarray]] - A0tA0_csc: sparse.csc_matrix # retained for the LSMR fallback - survey_weights_0: Optional[np.ndarray] - singular: bool - - # ============================================================================= # Main Estimator # ============================================================================= -class _LSMRUnconvergedError(RuntimeError): - """LSMR failed to certify a solution on the singular-variance fallback. - - Raised (not returned as NaN) so the variance boundary can fail closed: - a NaN vector would be laundered into zeros by the missing-FE - ``nan_to_num`` in the psi product — producing a finite, WRONG variance — - whereas this exception is caught in ``_compute_conservative_variance`` - and converted to a NaN SE (the all-or-nothing NaN inference convention). - """ - - -def _lsmr_minnorm_normal_solve(A0tA0_csc, rhs: np.ndarray) -> np.ndarray: - """Least-squares solve of the (possibly singular) normal equations - ``(A_0'[W]A_0) z = rhs`` WITHOUT densifying the sparse matrix. - - Replaces the previous ``np.linalg.lstsq(A0tA0.toarray(), ...)`` fallback, - whose dense materialization scales ``O((U+T+K)^2)`` — an OOM risk on - large panels (the TODO row this resolves). ``scipy.sparse.linalg.lsmr`` - handles singular symmetric systems, converging to the minimum-norm - least-squares solution (the same solution family as ``lstsq``'s - pseudo-inverse solution). - - Solver choice cannot change the estimator output: any two least-squares - solutions differ by a ``null(A_0'[W]A_0) = null(sqrt(W) A_0)`` component, - which the downstream projection ``v_untreated = -[W_0] A_0 z`` - annihilates (unweighted: ``null = null(A_0)`` so ``A_0 z`` is invariant; - weighted: the weight multiplication zeroes exactly the rows where the - null component can be nonzero). Locked by the singular-system parity - test against a dense-lstsq oracle. - - CONVERGENCE IS VALIDATED (fail-closed): ``istop`` in ``{0, 1, 2, 4, 5}`` - means LSMR certified an (approximate) solution / least-squares solution - within ``atol``/``btol`` (4 and 5 are the machine-precision analogues of - 1 and 2 per SciPy's documentation); anything else (condition-limit stop, - max-iteration exhaustion) gets ONE retry with an uncapped condition - limit and a generous iteration budget, and if still uncertified raises - :class:`_LSMRUnconvergedError` — caught at the variance boundary and - converted to a NaN SE — rather than feeding a finite-but-unverified - solution into the Theorem 3 weights. - """ - import scipy.sparse.linalg as spla - - _certified = (0, 1, 2, 4, 5) - result = spla.lsmr(A0tA0_csc, rhs, atol=1e-14, btol=1e-14) - z, istop = result[0], int(result[1]) - if istop not in _certified or not np.all(np.isfinite(z)): - dim = A0tA0_csc.shape[0] - result = spla.lsmr( - A0tA0_csc, rhs, atol=1e-14, btol=1e-14, conlim=1e16, maxiter=max(50 * dim, 10_000) - ) - z, istop = result[0], int(result[1]) - if istop not in _certified or not np.all(np.isfinite(z)): - warnings.warn( - "ImputationDiD variance: the LSMR fallback solve of " - f"(A_0'[W]A_0) z = rhs did not converge (istop={istop}); " - "the affected variance is reported as NaN rather than from " - "an unverified solution.", - UserWarning, - stacklevel=3, - ) - raise _LSMRUnconvergedError(f"LSMR uncertified (istop={istop})") - return z - - -class ImputationDiD(ImputationDiDBootstrapMixin, BaseEstimator): +class ImputationDiD(ImputationDiDBootstrapMixin, _ImputationAggregationMixin, BaseEstimator): """ Borusyak-Jaravel-Spiess (2024) imputation DiD estimator. @@ -229,8 +143,9 @@ class ImputationDiD(ImputationDiDBootstrapMixin, BaseEstimator): df_convention : {"residual", "cluster", "normal"}, default "residual" Degrees-of-freedom convention for the PRETRENDS lead regression's per-lead t/p/CI — the one ImputationDiD surface running the shared - clustered CR1 sandwich (``pretrends=True`` with - ``aggregate="event_study"``/``"all"``). ``"residual"`` (default) uses + clustered CR1 sandwich (``pretrends=True``, surfaced fit-time via + the deprecated ``aggregate="event_study"``/``"all"`` or post-fit via + ``results.aggregate('event_study')``). ``"residual"`` (default) uses the lead regression's residual df (``n − k_kept − absorbed [time, unit] rank``) — the 3.9 fix: previously silent normal-theory z on plain clustered fits; ``"cluster"`` uses ``G − 1``; @@ -260,14 +175,14 @@ class ImputationDiD(ImputationDiDBootstrapMixin, BaseEstimator): ... time='time', first_treat='first_treat') >>> results.print_summary() - With event study: + With a post-fit event study (M-021): >>> est = ImputationDiD() >>> results = est.fit(data, outcome='outcome', unit='unit', - ... time='time', first_treat='first_treat', - ... aggregate='event_study') + ... time='time', first_treat='first_treat') + >>> es = results.aggregate('event_study') >>> from diff_diff import plot_event_study - >>> plot_event_study(results) + >>> plot_event_study(es) Notes ----- @@ -346,8 +261,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["SurveyDesign"] = None, ) -> ImputationDiDResults: """ @@ -369,11 +284,18 @@ def fit( covariates : list of str, optional List of covariate column names. aggregate : str, optional - Aggregation mode: None/"simple" (overall ATT only), - "event_study", "group", or "all". + DEPRECATED (3.9, removed in 4.0; row M-021): aggregate as a + post-fit step instead — ``results.aggregate('event_study')`` / + ``.aggregate('group')`` / ``.aggregate('simple')``. Supplying + ANY value (``None`` included) warns; the deprecated path still + works and returns exactly the numbers it always did + (fit-time mode: None/"simple" overall only, "event_study", + "group", or "all"). balance_e : int, optional - When computing event study, restrict to cohorts observed at all - relative times in [-balance_e, max_h]. + DEPRECATED (3.9, removed in 4.0; row M-118): moves onto + ``results.aggregate('event_study', balance_e=...)``. Restricts + the event study to cohorts observed at every relative time in + ``[-balance_e, max_h]`` (the balanced-window rule). survey_design : SurveyDesign, optional Survey design specification for design-based inference. Supports pweight only (aweight/fweight raise ValueError). Supports strata, @@ -391,6 +313,38 @@ def fit( ValueError If required columns are missing or data validation fails. """ + # M-021/M-118 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. + # The sentinel is normalized HERE, before every downstream read of + # ``aggregate`` (the pretrends+replicate gate and the df_convention + # reachability warning below both read it). + _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"ImputationDiD.fit({_args}) is deprecated and will be " + "removed in 4.0. Fit once, then aggregate as a post-fit " + "step: results = ImputationDiD().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 + # Re-validate vcov_type at fit-time: set_params validates eagerly # (BaseEstimator probe re-init), so this only catches DIRECT # attribute mutation (est.vcov_type = ...). @@ -422,19 +376,36 @@ def fit( ) # Inert-config warning (no-silent-failures): the df_convention knob - # moves only the pretrends lead regression's per-lead t/p/CI, which - # this fit surfaces only when pretrends=True with - # aggregate="event_study"/"all". (The post-fit ``pretrend_test()`` - # reaches the lead helper too, but consumes only gamma/V_gamma — its - # joint Wald F denominator is knob-independent.) An explicitly - # non-default value on any other configuration is a silent no-op. - if self.df_convention != "residual" and not ( - self.pretrends and aggregate in ("event_study", "all") - ): + # moves only the pretrends lead regression's per-lead t/p/CI. Since + # the M-021 post-fit migration that inference is REACHABLE from any + # analytical pretrends=True fit via results.aggregate('event_study'), + # so the predicate is reachability-based, not aggregate-keyed: + # reachable iff pretrends AND not replicate-weight (the gate above + # rejects fit-time ES and the post-fit path fails closed too) AND + # (the deprecated fit-time ES/all was supplied OR n_bootstrap <= 0 — + # a bootstrapped fit builds no ES surface and post-fit aggregate() + # fails closed on it; <= 0 because n_bootstrap is unvalidated and + # every bootstrap gate is `> 0`). Reachability-BASED, not exact: a + # fit whose bootstrap later FAILS (bootstrap_results=None) can still + # aggregate post-fit, so that corner warns spuriously — the warning + # fires before the bootstrap runs and cannot know. (The post-fit + # ``pretrend_test()`` reaches the lead helper too, but consumes only + # gamma/V_gamma — its joint Wald F denominator is knob-independent.) + _is_replicate_design = ( + survey_design is not None and survey_design.replicate_method is not None + ) + _lead_inference_reachable = ( + self.pretrends + and not _is_replicate_design + and (aggregate in ("event_study", "all") or self.n_bootstrap <= 0) + ) + if self.df_convention != "residual" and not _lead_inference_reachable: warnings.warn( f"df_convention={self.df_convention!r} affects only the " - "pretrends event-study per-lead inference (pretrends=True " - "with aggregate='event_study'/'all'); it has no effect on " + "pretrends event-study per-lead inference (pretrends=True, " + "surfaced fit-time via the deprecated aggregate=" + "'event_study'/'all' or post-fit via " + "results.aggregate('event_study')); it has no effect on " "this configuration.", UserWarning, stacklevel=2, @@ -767,6 +738,11 @@ def fit( if _uses_replicate_imp and _survey_df is None: _survey_df = 0 # rank-deficient replicate → NaN inference + # Kit df-provenance SEED (M-021): the exact value the analytical + # aggregators below receive, captured BEFORE the replicate override + # can rebind _survey_df — post-fit recompute must re-seed from it. + _survey_df_seed = _survey_df + # Compute overall inference (may be overridden by replicate below) overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=_survey_df @@ -825,129 +801,33 @@ def fit( # Replicate variance: derive keys from actual outputs (after filtering) if _uses_replicate_imp: - from diff_diff.survey import compute_replicate_refit_variance - - _rel_times_treated = df.loc[omega_1_mask, "_rel_time"].values - _cohorts_treated = df.loc[omega_1_mask, first_treat].values - - # Derive keys from actual outputs (excludes filtered/Prop5/ref) - _es_effects = event_study_effects or {} - _grp_effects = group_effects or {} - _sorted_rel_times = sorted( - e - for e in _es_effects.keys() - if np.isfinite(_es_effects[e]["effect"]) and _es_effects[e].get("n_obs", 1) > 0 - ) - _sorted_groups = sorted( - g for g in _grp_effects.keys() if np.isfinite(_grp_effects[g]["effect"]) - ) - _n_es = len(_sorted_rel_times) - - # Pre-compute balanced cohort mask for balance_e - _balanced_mask_treated = None - if balance_e is not None and _sorted_rel_times: - df_1 = df.loc[omega_1_mask] - rel_times_all = df_1["_rel_time"].values - all_horizons_full = sorted(set(int(h) for h in rel_times_all if np.isfinite(h))) - if self.horizon_max is not None: - all_horizons_full = [h for h in all_horizons_full if abs(h) <= self.horizon_max] - cohort_rel_times = self._build_cohort_rel_times(df, first_treat) - _balanced_mask_treated = self._compute_balanced_cohort_mask( - df_1, first_treat, all_horizons_full, balance_e, cohort_rel_times - ) - - # Single vectorized refit: [overall, es_e0..., grp_g0...] - def _refit_imp(w_r): - ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( - df, - outcome, - unit, - time, - covariates, - omega_0_mask, - weights=w_r, - ) - tau_r, _ = self._impute_treatment_effects( - df, - outcome, - unit, - time, - covariates, - omega_1_mask, - ufe_r, - tfe_r, - gm_r, - delta_r, - ) - fin = np.isfinite(tau_r) - treated_w = w_r[omega_1_mask.values] - results = [] - # [0] Overall ATT - tw_fin = treated_w[fin] - tw_sum = np.sum(tw_fin) - results.append( - float(np.sum(tau_r[fin] * tw_fin) / tw_sum) if tw_sum > 0 else np.nan - ) - # [1..n_es] Event-study (identified only) - for e in _sorted_rel_times: - mask_e = fin & (_rel_times_treated == e) - if _balanced_mask_treated is not None: - mask_e = mask_e & _balanced_mask_treated - tw_e = treated_w[mask_e] - s = np.sum(tw_e) - results.append(float(np.sum(tau_r[mask_e] * tw_e) / s) if s > 0 else np.nan) - # [n_es+1..] Group (identified only) - for g in _sorted_groups: - mask_g = fin & (_cohorts_treated == g) - tw_g = treated_w[mask_g] - s = np.sum(tw_g) - results.append(float(np.sum(tau_r[mask_g] * tw_g) / s) if s > 0 else np.nan) - return np.array(results) - - # Build full-sample estimate from actual effects - _full_est = [overall_att] - _full_est.extend([_es_effects[e]["effect"] for e in _sorted_rel_times]) - _full_est.extend([_grp_effects[g]["effect"] for g in _sorted_groups]) - - _vcov_rep_imp, _n_valid_rep_imp = compute_replicate_refit_variance( - _refit_imp, np.array(_full_est), resolved_survey + ( + _vcov_rep_imp, + _n_valid_rep_imp, + _survey_df, + ) = self._replicate_override_aggregates( + df=df, + outcome=outcome, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + resolved_survey=resolved_survey, + overall_att=overall_att, + event_study_effects=event_study_effects, + group_effects=group_effects, + balance_e=balance_e, + survey_df_seed=_survey_df, ) overall_se = float(np.sqrt(max(_vcov_rep_imp[0, 0], 0.0))) - - # Override df if replicates were dropped - # Replicate-refit path is only reached with a resolved design. - assert resolved_survey is not None - if _n_valid_rep_imp < resolved_survey.n_replicates: - _survey_df = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0 if survey_metadata is not None: survey_metadata.df_survey = _survey_df if _survey_df and _survey_df > 0 else None - overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=_survey_df ) - # Override event-study SEs from vcov diagonal - for i, e in enumerate(_sorted_rel_times): - if event_study_effects is not None and e in event_study_effects: - se_e = float(np.sqrt(max(_vcov_rep_imp[1 + i, 1 + i], 0.0))) - eff_e = event_study_effects[e]["effect"] - t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df) - event_study_effects[e]["se"] = se_e - event_study_effects[e]["t_stat"] = t_e - event_study_effects[e]["p_value"] = p_e - event_study_effects[e]["conf_int"] = ci_e - - # Override group SEs from vcov diagonal - for j, g in enumerate(_sorted_groups): - if group_effects is not None and g in group_effects: - se_g = float(np.sqrt(max(_vcov_rep_imp[1 + _n_es + j, 1 + _n_es + j], 0.0))) - eff_g = group_effects[g]["effect"] - t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df) - group_effects[g]["se"] = se_g - group_effects[g]["t_stat"] = t_g - group_effects[g]["p_value"] = p_g - group_effects[g]["conf_int"] = ci_g - # Build treatment effects dataframe treated_df = df.loc[omega_1_mask, [unit, time, "_tau_hat", "_rel_time"]].copy() treated_df = treated_df.rename(columns={"_tau_hat": "tau_hat", "_rel_time": "rel_time"}) @@ -1081,6 +961,10 @@ def _refit_imp(w_r): group_effects[g]["t_stat"] = safe_inference( eff_val, se_val, alpha=self.alpha )[0] + # Percentile inference replaced the analytical row — + # never publish an analytical df beside it (M-021; + # the EfficientDiD M-023 precedent). + group_effects[g]["df_used"] = None # Resolve cluster_name / n_clusters for Results metadata. # Suppress under ANY survey design (the survey block in summary() @@ -1133,6 +1017,29 @@ def _refit_imp(w_r): df_convention=self.df_convention, ) + # Attach the post-fit aggregation kit (M-021/M-118). Unconditional — + # including bootstrap fits, whose gate lives in _aggregate_compute + # (a FAILED bootstrap leaves bootstrap_results=None and the fit + # aggregates normally). + self.results_._aggregation_kit = _build_imputation_aggregation_kit( + fit_data=self._fit_data, + treatment_groups=treatment_groups, + overall_att=overall_att, + n_treated_obs=n_omega_1, + uses_replicate=_uses_replicate_imp, + survey_df_seed=_survey_df_seed, + survey_df_final=_survey_df, + survey_metadata=survey_metadata, + horizon_max=self.horizon_max, + pretrends=self.pretrends, + aux_partition=self.aux_partition, + leave_one_out=self.leave_one_out, + rank_deficient_action=self.rank_deficient_action, + df_convention=self.df_convention, + alpha=self.alpha, + anticipation=self.anticipation, + ) + self.is_fitted_ = True return self.results_ @@ -1140,1475 +1047,22 @@ def _refit_imp(w_r): # Step 1: OLS on untreated observations # ========================================================================= - def _iterative_fe( - self, - y: np.ndarray, - unit_vals: np.ndarray, - time_vals: np.ndarray, - idx: pd.Index, - max_iter: int = 10_000, - tol: float = 1e-10, - weights: Optional[np.ndarray] = None, - ) -> Tuple[Dict[Any, float], Dict[Any, float]]: - """ - Estimate unit and time FE via iterative alternating projection (Gauss-Seidel). - - Thin wrapper over the shared bincount solver - (``diff_diff.utils._iterative_fe_solve``): factorize unit/time once, - solve on integer codes, map the level arrays back to dicts. - Converges to the exact (W)LS solution for balanced and unbalanced - panels; balanced panels converge in 1-2 iterations. - - Parameters - ---------- - idx : pd.Index - Unused; retained for call-site stability. - weights : np.ndarray, optional - Survey weights (weighted group means ``sum(w*x)/sum(w)``). A - unit/period whose observations ALL carry zero weight has no - identifying contribution and gets ``NaN`` FE (its key is kept so - the rank-condition membership check still sees the group). - - Returns - ------- - unit_fe : dict - Mapping from unit -> unit fixed effect. - time_fe : dict - Mapping from time -> time fixed effect. - """ - unit_codes, unit_uniques = pd.factorize(unit_vals, sort=False) - time_codes, time_uniques = pd.factorize(time_vals, sort=False) - if (unit_codes < 0).any() or (time_codes < 0).any(): - raise ValueError( - "ImputationDiD: unit or time column contains NaN. Drop or " - "impute missing group keys before fitting." - ) - unit_fe_arr, time_fe_arr = _iterative_fe_solve( - np.asarray(y, dtype=np.float64), - unit_codes.astype(np.intp, copy=False), - time_codes.astype(np.intp, copy=False), - len(unit_uniques), - len(time_uniques), - weights=weights, - max_iter=max_iter, - tol=tol, - method_name="ImputationDiD iterative FE solver", - ) - unit_fe = dict(zip(unit_uniques, unit_fe_arr)) - time_fe = dict(zip(time_uniques, time_fe_arr)) - return unit_fe, time_fe - - @staticmethod - def _compute_balanced_cohort_mask( - df_treated: pd.DataFrame, - first_treat: str, - all_horizons: List[int], - balance_e: int, - cohort_rel_times: Dict[Any, Set[int]], - ) -> np.ndarray: - """Compute boolean mask selecting treated obs from balanced cohorts. - - A cohort is 'balanced' if it has observations at every relative time - in [-balance_e, max(all_horizons)]. - - Parameters - ---------- - df_treated : pd.DataFrame - Post-treatment observations (Omega_1). - first_treat : str - Column name for cohort identifier. - all_horizons : list of int - Post-treatment horizons in the event study. - balance_e : int - Number of pre-treatment periods to require. - cohort_rel_times : dict - Maps each cohort value to the set of all observed relative times - (including pre-treatment) from the full panel. Built by - _build_cohort_rel_times(). - """ - if not all_horizons: - return np.ones(len(df_treated), dtype=bool) - - max_h = max(all_horizons) - required_range = set(range(-balance_e, max_h + 1)) - - balanced_cohorts = set() - for g, horizons in cohort_rel_times.items(): - if required_range.issubset(horizons): - balanced_cohorts.add(g) - - return df_treated[first_treat].isin(balanced_cohorts).values - - @staticmethod - def _build_cohort_rel_times( - df: pd.DataFrame, - first_treat: str, - ) -> Dict[Any, Set[int]]: - """Build mapping of cohort -> set of observed relative times from full panel. - - Precondition: df must have '_never_treated' and '_rel_time' columns - (set by fit() before any aggregation calls). - """ - treated_mask = ~df["_never_treated"] - treated_df = df.loc[treated_mask] - result: Dict[Any, Set[int]] = {} - ft_vals = treated_df[first_treat].values - rt_vals = treated_df["_rel_time"].values - for i in range(len(treated_df)): - h = rt_vals[i] - if np.isfinite(h): - result.setdefault(ft_vals[i], set()).add(int(h)) - return result - - def _fit_untreated_model( - self, - df: pd.DataFrame, - outcome: str, - unit: str, - time: str, - covariates: Optional[List[str]], - omega_0_mask: pd.Series, - weights: Optional[np.ndarray] = None, - ) -> Tuple[ - Dict[Any, float], Dict[Any, float], float, Optional[np.ndarray], Optional[np.ndarray] - ]: - """ - Step 1: Estimate unit + time FE on untreated observations. - - Uses iterative alternating projection (Gauss-Seidel) to compute exact - OLS fixed effects for both balanced and unbalanced panels. For balanced - panels, converges in 1-2 iterations (identical to one-pass demeaning). - - Parameters - ---------- - weights : np.ndarray, optional - Full-panel survey weights (same length as df). The untreated subset - is extracted internally via omega_0_mask. When None, unweighted. - - Returns - ------- - unit_fe : dict - Unit fixed effects {unit_id: alpha_i}. - time_fe : dict - Time fixed effects {time_period: beta_t}. - grand_mean : float - Grand mean (0.0 — absorbed into iterative FE). - delta_hat : np.ndarray or None - Covariate coefficients (if covariates provided). - kept_cov_mask : np.ndarray or None - Boolean mask of shape (n_covariates,) indicating which covariates - have finite coefficients. None if no covariates. - """ - df_0 = df.loc[omega_0_mask] - w_0 = weights[omega_0_mask.values] if weights is not None else None - - if covariates is None or len(covariates) == 0: - # No covariates: estimate FE via iterative alternating projection - # (exact OLS for both balanced and unbalanced panels) - y = df_0[outcome].values.copy() - unit_fe, time_fe = self._iterative_fe( - y, df_0[unit].values, df_0[time].values, df_0.index, weights=w_0 - ) - # grand_mean = 0: iterative FE absorb the intercept - return unit_fe, time_fe, 0.0, None, None - - else: - # With covariates: iteratively demean Y and X, OLS for delta, - # then recover FE from covariate-adjusted outcome - y = df_0[outcome].values.copy() - X_raw = df_0[covariates].values.copy() - units = df_0[unit].values - times = df_0[time].values - - # Step A: within-transform Y and all X columns through the shared - # MAP engine (factorize-once + bincount + optional Rust kernel), - # one dispatch for every column. within_transform pins - # [unit, time]; [time, unit] here preserves the historical - # time-then-unit sweep order of the per-estimator loops. - narrow = df_0[[outcome, *covariates, time, unit]].copy() - demeaned, _ = demean_by_groups( - narrow, - [outcome, *covariates], - [time, unit], - inplace=True, - weights=w_0, - max_iter=10_000, - tol=1e-10, - ) - y_dm = demeaned[outcome].to_numpy(dtype=np.float64) - X_dm = demeaned[covariates].to_numpy(dtype=np.float64) - - # Step B: OLS for covariate coefficients on demeaned data - result = solve_ols( - X_dm, - y_dm, - return_vcov=False, - rank_deficient_action=self.rank_deficient_action, - column_names=covariates, - weights=w_0, - ) - delta_hat = result[0] - - # Mask of covariates with finite coefficients (before cleaning) - # Used to exclude rank-deficient covariates from variance design matrices - kept_cov_mask = np.isfinite(delta_hat) - - # Replace NaN coefficients with 0 for adjustment - # (rank-deficient covariates are dropped) - delta_hat_clean = np.where(np.isfinite(delta_hat), delta_hat, 0.0) - - # Step C: Recover FE from covariate-adjusted outcome using iterative FE - y_adj = y - np.dot(X_raw, delta_hat_clean) - unit_fe, time_fe = self._iterative_fe(y_adj, units, times, df_0.index, weights=w_0) - - # grand_mean = 0: iterative FE absorb the intercept - return unit_fe, time_fe, 0.0, delta_hat_clean, kept_cov_mask - # ========================================================================= # Step 2: Impute counterfactuals # ========================================================================= - def _impute_treatment_effects( - self, - df: pd.DataFrame, - outcome: str, - unit: str, - time: str, - covariates: Optional[List[str]], - omega_1_mask: pd.Series, - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - grand_mean: float, - delta_hat: Optional[np.ndarray], - ) -> Tuple[np.ndarray, np.ndarray]: - """ - Step 2: Impute Y(0) for treated observations and compute tau_hat. - - Returns - ------- - tau_hat : np.ndarray - Imputed treatment effects for each treated observation. - y_hat_0 : np.ndarray - Imputed counterfactual Y(0). - """ - df_1 = df.loc[omega_1_mask] - - # Look up unit and time FE - alpha_i = df_1[unit].map(unit_fe).values - beta_t = df_1[time].map(time_fe).values - - # Handle missing FE (set to NaN) - alpha_i = np.where(pd.isna(alpha_i), np.nan, alpha_i).astype(float) - beta_t = np.where(pd.isna(beta_t), np.nan, beta_t).astype(float) - - y_hat_0 = grand_mean + alpha_i + beta_t - - if delta_hat is not None and covariates: - X_1 = df_1[covariates].values - y_hat_0 = y_hat_0 + np.dot(X_1, delta_hat) - - tau_hat = df_1[outcome].values - y_hat_0 - - return tau_hat, y_hat_0 - # ========================================================================= # Conservative Variance (Theorem 3) # ========================================================================= - def _compute_cluster_psi_sums( - self, - df: pd.DataFrame, - outcome: str, - unit: str, - time: str, - first_treat: str, - covariates: Optional[List[str]], - omega_0_mask: pd.Series, - omega_1_mask: pd.Series, - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - grand_mean: float, - delta_hat: Optional[np.ndarray], - weights: np.ndarray, - cluster_var: str, - kept_cov_mask: Optional[np.ndarray] = None, - survey_weights_0: Optional[np.ndarray] = None, - proj_cache: Optional[Dict[Any, _UntreatedProjection]] = None, - ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: - """ - Compute cluster-level influence function sums (Theorem 3). - - psi_i = sum_t v_it * epsilon_tilde_it, summed within each cluster. - - Returns - ------- - cluster_psi_sums : np.ndarray - Array of cluster-level psi sums. - cluster_ids_unique : np.ndarray - Unique cluster identifiers (matching order of psi sums). - """ - df_0 = df.loc[omega_0_mask] - df_1 = df.loc[omega_1_mask] - - # ---- Compute v_it for treated observations ---- - v_treated = weights.copy() - - # ---- Compute v_it for untreated observations ---- - # Exact two-way-FE imputation projection - # v_untreated = -A_0 (A_0' [W] A_0)^{-1} A_1' w_treated (Theorem 3 / the - # implied weights of Supplementary Proposition A3), used for BOTH the - # FE-only and the covariate case. The earlier FE-only closed form - # -(w_i/n0_i + w_t/n0_t - w/N_0) is exact only for a *balanced* untreated - # panel; Omega_0 is generically unbalanced in staggered designs (treated - # observations are removed), which biased the analytical SE downward - # (~27% on the parity panel). The projection matches R `didimputation` - # exactly -- see tests/test_methodology_imputation.py::TestImputationDiDParityR. - # Build the target-invariant projection design + factorization once per - # fit() (cached in proj_cache), then solve only the target-specific RHS. - # survey_weights is DELIBERATELY excluded from the key: the cache is a - # fit-LOCAL dict, and within one fit() survey_weights is a single fixed - # object, so the masks deterministically map to one sw_0 = - # survey_weights[omega_0_mask]. The masks + covariates + kept_cov_mask - # therefore FULLY identify the design (sw_0 itself is a fresh-sliced array - # per call -- keying on its id() would miss every time and balloon the - # cache to 1+H+G full A_0/A_1/factorization entries). id()-keys are safe: - # the masks are fit() locals alive for the whole fit and the cache is a - # fit-local dict, so no cross-fit leak / id reuse. - cov_list = covariates if covariates is not None else [] - ctx: Optional[_UntreatedProjection] = None - if proj_cache is not None: - key = ( - id(omega_0_mask), - id(omega_1_mask), - tuple(cov_list), - kept_cov_mask.tobytes() if kept_cov_mask is not None else None, - ) - ctx = proj_cache.get(key) - if ctx is None: - ctx = self._build_untreated_projection( - df_0, - df_1, - unit, - time, - cov_list, - kept_cov_mask=kept_cov_mask, - survey_weights_0=survey_weights_0, - ) - if proj_cache is not None: - proj_cache[key] = ctx - v_untreated = self._solve_untreated_v(ctx, weights) - - # ---- Compute auxiliary model residuals (Equation 8) ---- - epsilon_treated = self._compute_auxiliary_residuals_treated( - df_1, - outcome, - unit, - time, - first_treat, - covariates, - unit_fe, - time_fe, - grand_mean, - delta_hat, - v_treated, - ) - epsilon_untreated = self._compute_residuals_untreated( - df_0, outcome, unit, time, covariates, unit_fe, time_fe, grand_mean, delta_hat - ) - - # ---- psi_it = v_it * epsilon_tilde_it ---- - v_all = np.empty(len(df)) - v_all[omega_1_mask.values] = v_treated - v_all[omega_0_mask.values] = v_untreated - - eps_all = np.empty(len(df)) - eps_all[omega_1_mask.values] = epsilon_treated - eps_all[omega_0_mask.values] = epsilon_untreated - - ve_product = v_all * eps_all - # NaN eps from missing FE (rank condition violation). Zero their variance - # contribution — matches R's did_imputation which drops unimputable obs. - np.nan_to_num(ve_product, copy=False, nan=0.0) - - # Sum within clusters - cluster_ids = df[cluster_var].values - ve_series = pd.Series(ve_product, index=df.index) - cluster_sums = ve_series.groupby(cluster_ids).sum() - - return cluster_sums.values, cluster_sums.index.values, ve_product - - def _compute_conservative_variance( - self, - df: pd.DataFrame, - outcome: str, - unit: str, - time: str, - first_treat: str, - covariates: Optional[List[str]], - omega_0_mask: pd.Series, - omega_1_mask: pd.Series, - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - grand_mean: float, - delta_hat: Optional[np.ndarray], - weights: np.ndarray, - cluster_var: str, - kept_cov_mask: Optional[np.ndarray] = None, - survey_weights: Optional[np.ndarray] = None, - resolved_survey=None, - proj_cache: Optional[Dict[Any, _UntreatedProjection]] = None, - ) -> float: - """ - Compute conservative clustered variance (Theorem 3, Equation 7). - - Parameters - ---------- - weights : np.ndarray - Aggregation weights w_it for treated observations. - Shape: (n_treated,), must sum to 1. - survey_weights : np.ndarray, optional - Full-panel survey weights. When provided, they enter the untreated - v_it WLS projection (weighted normal equations plus the left - per-observation weight factor) and the design-based variance path. - resolved_survey : ResolvedSurveyDesign, optional - When provided, uses design-based variance via - ``compute_survey_if_variance()`` (supports strata, PSU, FPC). - - Returns - ------- - float - Standard error. - """ - sw_0 = survey_weights[omega_0_mask.values] if survey_weights is not None else None - try: - cluster_psi_sums, _, ve_product = self._compute_cluster_psi_sums( - df=df, - outcome=outcome, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - weights=weights, - cluster_var=cluster_var, - kept_cov_mask=kept_cov_mask, - survey_weights_0=sw_0, - proj_cache=proj_cache, - ) - except _LSMRUnconvergedError: - # Solver failure is GLOBAL (the untreated projection is invalid), - # unlike per-observation missing-FE NaNs — fail the whole SE - # closed instead of letting nan_to_num launder it to zeros. - return np.nan - - if resolved_survey is not None: - # Design-based variance with strata/PSU/FPC support - from diff_diff.survey import compute_survey_if_variance - - variance = compute_survey_if_variance(ve_product, resolved_survey) - if np.isnan(variance): - return np.nan - return np.sqrt(max(variance, 0.0)) - - sigma_sq = float((cluster_psi_sums**2).sum()) - return np.sqrt(max(sigma_sq, 0.0)) - - def _build_untreated_projection( - self, - df_0: pd.DataFrame, - df_1: pd.DataFrame, - unit: str, - time: str, - covariates: List[str], - kept_cov_mask: Optional[np.ndarray] = None, - survey_weights_0: Optional[np.ndarray] = None, - ) -> _UntreatedProjection: - """ - Build the target-INVARIANT pieces of the exact imputation projection - ``v_untreated = -A_0 (A_0' [W] A_0)^{-1} A_1' w_treated`` and factorize the - normal-equations matrix once. The result is cached per ``fit()`` (see - ``_compute_cluster_psi_sums``) and reused across all estimand targets; - only the target-specific RHS ``A_1' w`` is solved per target in - ``_solve_untreated_v``. - - This is the GENERAL path -- used for both the FE-only and the covariate - cases (an empty ``covariates`` list builds a pure two-way-FE design; - ``n_cov == 0`` is the FE-only path). When survey_weights_0 is provided, - uses the weighted normal equations ``A_0' W A_0`` (the per-observation - survey weight is reapplied to the solved v in ``_solve_untreated_v``). - - Uses scipy.sparse for FE dummy columns to reduce memory from O(N*(U+T)) - to O(N) for the FE portion. An exactly singular ``A_0'[W]A_0`` makes - ``sparse_factorized`` raise ``RuntimeError``; we emit a UserWarning (once - per fit) and record ``singular=True`` so the solve routes to the sparse - LSMR least-squares fallback (no dense materialization; see - :func:`_lsmr_minnorm_normal_solve`). - """ - # Exclude rank-deficient covariates from design matrices - if kept_cov_mask is not None and not np.all(kept_cov_mask): - covariates = [c for c, k in zip(covariates, kept_cov_mask) if k] - - units_0 = df_0[unit].values - times_0 = df_0[time].values - units_1 = df_1[unit].values - times_1 = df_1[time].values - - all_units = np.unique(np.concatenate([units_0, units_1])) - all_times = np.unique(np.concatenate([times_0, times_1])) - unit_to_idx = {u: i for i, u in enumerate(all_units)} - time_to_idx = {t: i for i, t in enumerate(all_times)} - n_units = len(all_units) - n_times = len(all_times) - n_cov = len(covariates) - # Two-way FE design = all unit dummies (their sum spans the intercept) + - # time dummies dropping the first (identification). Dropping the first - # unit dummy too -- with no intercept column -- would omit the baseline - # level dimension and project onto a space one rank short of the true - # two-way-FE span, biasing the imputation weights (and hence the SE). - n_fe_cols = n_units + (n_times - 1) - - def _build_A_sparse(df_sub, unit_vals, time_vals): - n = len(df_sub) - - # Unit dummies — keep ALL (together they span the intercept). - u_indices = np.array([unit_to_idx[u] for u in unit_vals]) - u_rows = np.arange(n) - u_cols = u_indices - - # Time dummies (drop first) — vectorized - t_indices = np.array([time_to_idx[t] for t in time_vals]) - t_mask = t_indices > 0 - t_rows = np.arange(n)[t_mask] - t_cols = n_units + (t_indices[t_mask] - 1) - - rows = np.concatenate([u_rows, t_rows]) - cols = np.concatenate([u_cols, t_cols]) - data = np.ones(len(rows)) - - A_fe = sparse.csr_matrix((data, (rows, cols)), shape=(n, n_fe_cols)) - - # Covariates (dense, typically few columns) - if n_cov > 0: - A_cov = sparse.csr_matrix(df_sub[covariates].values) - A = sparse.hstack([A_fe, A_cov], format="csr") - else: - A = A_fe - - return A - - A_0 = _build_A_sparse(df_0, units_0, times_0) - A_1 = _build_A_sparse(df_1, units_1, times_1) - - # Form (A_0' [W] A_0). When survey weights present, use the weighted - # normal equations A_0' W A_0. - if survey_weights_0 is not None: - A0tA0_sparse = A_0.T @ A_0.multiply(survey_weights_0[:, None]) - else: - A0tA0_sparse = A_0.T @ A_0 # stays sparse - A0tA0_csc = A0tA0_sparse.tocsc() - - # Factorize once (factorize-once / solve-many). An exactly singular - # matrix makes sparse_factorized raise RuntimeError -- the same condition - # that previously surfaced as spsolve's MatrixRankWarning -> non-finite - # solution. Warn once and fall back to the sparse LSMR least-squares - # solve per target (no dense materialization). (The factorized path is - # bit-identical to the prior per-target spsolve for a single dense - # RHS -- both use the SuperLU simple driver with the same defaults.) - try: - solver: Optional[Callable[[np.ndarray], np.ndarray]] = sparse_factorized(A0tA0_csc) - singular = False - except RuntimeError as exc: - # Silent-failure audit axis C: emit a UserWarning on fallback instead - # of swallowing the error. Keep the "sparse LSMR" substring (asserted - # by tests). - warnings.warn( - "ImputationDiD variance: sparse factorization of (A_0' [W] A_0) " - f"failed ({type(exc).__name__}); falling back to a sparse LSMR " - "least-squares solve (no dense materialization). This may " - "indicate a rank-deficient or near-singular normal-equations " - "matrix and variance estimates may be less reliable.", - UserWarning, - stacklevel=2, - ) - solver = None - singular = True - - return _UntreatedProjection( - A_0=A_0, - A_1=A_1, - solver=solver, - A0tA0_csc=A0tA0_csc, - survey_weights_0=survey_weights_0, - singular=singular, - ) - - def _solve_untreated_v(self, ctx: _UntreatedProjection, weights: np.ndarray) -> np.ndarray: - """ - Solve the target-SPECIFIC RHS of the untreated imputation projection using - the cached design + factorization in ``ctx``: - ``v_untreated = -[W_0] A_0 (A_0'[W]A_0)^{-1} A_1' w_treated``. - """ - A1_w = ctx.A_1.T @ weights # (p,) - - if ctx.singular: - # Factorization was singular at build time (warned once already). - z = _lsmr_minnorm_normal_solve(ctx.A0tA0_csc, A1_w) - else: - assert ctx.solver is not None - z = ctx.solver(A1_w) - if not np.all(np.isfinite(z)): - # Defensive, target-specific: a non-finite solve on an otherwise - # factorizable matrix routes this RHS to the LSMR fallback. Warn per - # target (silent-failure audit axis C) -- distinct from the - # once-per-fit build-time singular warning. - warnings.warn( - "ImputationDiD variance: sparse solve of (A_0' [W] A_0) z = " - "A_1' w returned a non-finite solution; falling back to a " - "sparse LSMR least-squares solve for this target. Variance " - "estimates may be less reliable.", - UserWarning, - stacklevel=2, - ) - z = _lsmr_minnorm_normal_solve(ctx.A0tA0_csc, A1_w) - - # v_untreated = -[W_0] A_0 z (WLS projection requires per-obs weight) - v_untreated = -(ctx.A_0 @ z) - if ctx.survey_weights_0 is not None: - v_untreated = v_untreated * ctx.survey_weights_0 - return v_untreated - - def _compute_auxiliary_residuals_treated( - self, - df_1: pd.DataFrame, - outcome: str, - unit: str, - time: str, - first_treat: str, - covariates: Optional[List[str]], - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - grand_mean: float, - delta_hat: Optional[np.ndarray], - v_treated: np.ndarray, - ) -> np.ndarray: - """ - Compute auxiliary residuals for treated obs (Theorem 3, Equation 8). - - Implements the paper's *unit-clustered* group aggregator (Borusyak, - Jaravel & Spiess 2024, eq. 8, p. 3272), which minimizes the excess - variance of the conservative estimator under a within-group - constant-effect auxiliary model (Supplementary Appendix A.8): - - tau_tilde_g = sum_i (sum_{t in G_g,i} v_it)(sum_{t in G_g,i} v_it * tau_hat_it) - ---------------------------------------------------------------- - sum_i (sum_{t in G_g,i} v_it)^2 - - i.e. for each unit i form the within-unit weight sum a_{i,g} and the - within-unit weighted-effect sum b_{i,g} over the unit's observations in - group g, then combine across units. At the default cohort x event-time - partition (<=1 obs/unit/group) this reduces to sum(v^2 * tau_hat) / - sum(v^2) -- the form the R `didimputation` package implements -- and - equals the naive observation-level mean sum(v * tau_hat) / sum(v) only - when within-group weights are uniform. Under coarser `cohort` / `horizon` - partitions (a unit contributes several observations to a group) or - non-uniform v_it (e.g. survey weights) the two genuinely differ. - - epsilon_tilde_it = Y_it - alpha_i - beta_t [- X'delta] - tau_tilde_g - """ - n_1 = len(df_1) - - # Compute base residuals (Y - Y_hat(0) = tau_hat) - # NaN for missing FE (consistent with _impute_treatment_effects) - alpha_i = df_1[unit].map(unit_fe).values.astype(float) # NaN for missing - beta_t = df_1[time].map(time_fe).values.astype(float) # NaN for missing - y_hat_0 = grand_mean + alpha_i + beta_t - - if delta_hat is not None and covariates: - y_hat_0 = y_hat_0 + np.dot(df_1[covariates].values, delta_hat) - - tau_hat = df_1[outcome].values - y_hat_0 - - # Partition Omega_1 into groups G_g - if self.aux_partition == "cohort_horizon": - group_keys = list(zip(df_1[first_treat].values, df_1["_rel_time"].values)) - elif self.aux_partition == "cohort": - group_keys = list(df_1[first_treat].values) - elif self.aux_partition == "horizon": - group_keys = list(df_1["_rel_time"].values) - else: - group_keys = list(range(n_1)) # each obs is its own group - - # Factorize group keys to integer codes (robust to tuple-valued keys). - group_codes = pd.factorize(pd.Series(group_keys), sort=False)[0] - gc_series = pd.Series(group_codes, index=df_1.index) - tau_series = pd.Series(tau_hat, index=df_1.index) - - # Unit-clustered Equation 8. Only v_it != 0 observations contribute: a - # zero-weight row adds exactly 0 to both a_{i,g} and b_{i,g}, so dropping - # it is exact for finite tau_hat AND avoids letting an unimputable row - # (NaN tau_hat, which always carries v_it == 0 by construction in - # _compute_target_weights) poison its whole group via 0 * NaN = NaN. The - # previous observation-level pandas sum relied on skipna to drop them. - contrib = (v_treated != 0.0) & np.isfinite(tau_hat) - loo_factor: Optional[pd.Series] = None - n_single_loo = 0 - if contrib.any(): - inner = pd.DataFrame( - { - "g": group_codes[contrib], - "u": df_1[unit].values[contrib], - "v": v_treated[contrib], - "vt": v_treated[contrib] * tau_hat[contrib], - } - ) - # Per (group, unit): a_{i,g} = sum v_it, b_{i,g} = sum v_it * tau_hat - per_unit = inner.groupby(["g", "u"], sort=False).agg(a=("v", "sum"), b=("vt", "sum")) - # Per group: numerator sum_i a*b, denominator sum_i a^2 - per_group = ( - per_unit.assign(ab=per_unit["a"] * per_unit["b"], a2=per_unit["a"] ** 2) - .groupby(level="g") - .agg(num=("ab", "sum"), den=("a2", "sum")) - ) - den_ok = per_group["den"].abs() >= 1e-15 - tau_tilde_map = (per_group["num"] / per_group["den"]).where(den_ok) - # BJS 2024 App. A.9 leave-one-out refinement: rescale each treated - # residual by 1/(1 - v_ig^2 / sum_j v_jg^2) (== the direct-LOO tau_tilde - # exactly, at the per-unit cluster sum). Reuses a_{i,g} = per_unit['a'] - # and sum_j v_jg^2 = per_group['den']; applied to epsilon_treated below. - if self.leave_one_out: - loo_factor, n_single_loo = self._leave_one_out_factor(per_unit, per_group) - else: - tau_tilde_map = pd.Series(dtype=float) - - tau_tilde_per_obs = gc_series.map(tau_tilde_map) - - # Groups with no contributing (v_it != 0, finite tau_hat) observations -- - # e.g. off-target horizons in an event-study SE -- are a variance no-op - # (psi_g = sum_t v_it * eps_tilde_it = 0 there regardless of tau_tilde_g), - # so fall back to the unweighted group mean of tau_hat for a finite value. - if tau_tilde_per_obs.isna().any(): - simple_means = tau_series.groupby(gc_series).mean() - tau_tilde_per_obs = tau_tilde_per_obs.fillna(gc_series.map(simple_means)) - - tau_tilde = tau_tilde_per_obs.values - - # Auxiliary residuals - epsilon_treated = tau_hat - tau_tilde - - # Leave-one-out rescale (BJS 2024 App. A.9): map each treated obs to its - # (group, unit) factor and inflate the residual. Non-contributing rows - # (v_it == 0, psi == 0 anyway) and single-positive-weight-unit groups - # (LOO undefined, fn. 51) keep factor 1.0. - if self.leave_one_out and loo_factor is not None: - obs_index = pd.MultiIndex.from_arrays( - [group_codes, df_1[unit].values], names=["g", "u"] - ) - factor_per_obs = loo_factor.reindex(obs_index).to_numpy(dtype=float) - factor_per_obs = np.where(np.isfinite(factor_per_obs), factor_per_obs, 1.0) - epsilon_treated = epsilon_treated * factor_per_obs - if n_single_loo > 0: - warnings.warn( - f"leave_one_out=True: {n_single_loo} auxiliary group(s) have a single " - f"positive-weight unit, where the leave-one-out variance is undefined " - f"(Borusyak, Jaravel & Spiess 2024, Supp. App. A.9 fn. 51); those groups " - f"keep the non-leave-out residual. A coarser aux_partition reduces " - f"singleton groups.", - UserWarning, - stacklevel=2, - ) - - return epsilon_treated - - def _compute_residuals_untreated( - self, - df_0: pd.DataFrame, - outcome: str, - unit: str, - time: str, - covariates: Optional[List[str]], - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - grand_mean: float, - delta_hat: Optional[np.ndarray], - ) -> np.ndarray: - """Compute Step 1 residuals for untreated observations.""" - # Preserve NaN for any missing FE, symmetric with the treated path in - # _compute_auxiliary_residuals_treated. On valid data this is inert -- - # every untreated observation's unit and period appear in the Step 1 FE - # dicts (the dicts are estimated FROM Omega_0) -- but it stops a missing - # FE from silently becoming a 0 residual, which would mask a rank- - # condition logic error. Any NaN is zeroed downstream in the variance - # product (np.nan_to_num), exactly like the treated path. - alpha_i = df_0[unit].map(unit_fe).values.astype(float) - beta_t = df_0[time].map(time_fe).values.astype(float) - y_hat = grand_mean + alpha_i + beta_t - - if delta_hat is not None and covariates: - y_hat = y_hat + np.dot(df_0[covariates].values, delta_hat) - - return df_0[outcome].values - y_hat - # ========================================================================= # Aggregation # ========================================================================= - def _aggregate_event_study( - self, - df: pd.DataFrame, - outcome: str, - unit: str, - time: str, - first_treat: str, - covariates: Optional[List[str]], - omega_0_mask: pd.Series, - omega_1_mask: pd.Series, - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - grand_mean: float, - delta_hat: Optional[np.ndarray], - cluster_var: str, - treatment_groups: List[Any], - balance_e: Optional[int] = None, - kept_cov_mask: Optional[np.ndarray] = None, - survey_weights: Optional[np.ndarray] = None, - survey_df: Optional[int] = None, - resolved_survey=None, - proj_cache: Optional[Dict[Any, _UntreatedProjection]] = None, - ) -> Dict[int, Dict[str, Any]]: - """Aggregate treatment effects by event-study horizon.""" - df_1 = df.loc[omega_1_mask] - tau_hat = df["_tau_hat"].loc[omega_1_mask].values - rel_times = df_1["_rel_time"].values - - # Get all horizons - all_horizons = sorted(set(int(h) for h in rel_times if np.isfinite(h))) - - # Apply horizon_max filter - if self.horizon_max is not None: - all_horizons = [h for h in all_horizons if abs(h) <= self.horizon_max] - - # Apply balance_e filter - if balance_e is not None: - cohort_rel_times = self._build_cohort_rel_times(df, first_treat) - balanced_mask = pd.Series( - self._compute_balanced_cohort_mask( - df_1, first_treat, all_horizons, balance_e, cohort_rel_times - ), - index=df_1.index, - ) - else: - balanced_mask = pd.Series(True, index=df_1.index) - - # Check Proposition 5: no never-treated units - has_never_treated = df["_never_treated"].any() - h_bar = np.inf - if not has_never_treated and len(treatment_groups) > 1: - h_bar = max(treatment_groups) - min(treatment_groups) - - # Reference period - ref_period = -1 - self.anticipation - - event_study_effects: Dict[int, Dict[str, Any]] = {} - - # Add reference period marker - event_study_effects[ref_period] = { - "effect": 0.0, - "se": 0.0, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (0.0, 0.0), - "n_obs": 0, - } - - # Pre-period coefficients via BJS Test 1 lead regression - if self.pretrends: - df_0 = df.loc[omega_0_mask].copy() - - # Determine which cohorts' lead indicators to include. - # balance_e restricts which cohorts contribute lead dummies, - # but the full Omega_0 sample (including never-treated controls) - # is kept for the within-transformed OLS (BJS Test 1, Equation 9). - balanced_cohorts = None - skip_preperiods = False - if balance_e is not None: - cohort_rel_times_0 = self._build_cohort_rel_times(df, first_treat) - balanced_cohorts = set() - if all_horizons: - max_h = max(all_horizons) - required_range = set(range(-balance_e, max_h + 1)) - for g, horizons in cohort_rel_times_0.items(): - if required_range.issubset(horizons): - balanced_cohorts.add(g) - if not balanced_cohorts: - skip_preperiods = True # No cohorts qualify — skip entirely - - if not skip_preperiods: - rel_time_0 = np.where( - ~df_0["_never_treated"], - df_0[time] - df_0[first_treat], - np.nan, - ) - - # When balance_e is set, only include leads from balanced cohorts - if balanced_cohorts is not None: - is_balanced = df_0[first_treat].isin(balanced_cohorts).values - rel_time_for_leads = np.where(is_balanced, rel_time_0, np.nan) - else: - rel_time_for_leads = rel_time_0 - - pre_rel_times = sorted( - set( - int(h) - for h in rel_time_for_leads - if np.isfinite(h) and h < -self.anticipation - ) - ) - pre_rel_times = [h for h in pre_rel_times if h != ref_period] - if self.horizon_max is not None: - pre_rel_times = [h for h in pre_rel_times if abs(h) <= self.horizon_max] - if pre_rel_times: - # Survey pretrends: pass full design (subpopulation approach) - _sw_0_pre = None - _rs_full_pre = None - _n_full_pre = None - _o0_idx_pre = None - if survey_weights is not None and resolved_survey is not None: - _sw_0_pre = survey_weights[omega_0_mask.values] - _rs_full_pre = resolved_survey - _n_full_pre = len(df) - _o0_idx_pre = np.where(omega_0_mask.values)[0] - _survey_df_pre = ( - resolved_survey.df_survey if resolved_survey is not None else None - ) - pre_effects, _, _ = self._compute_lead_coefficients( - df_0, - outcome, - unit, - time, - first_treat, - covariates, - cluster_var, - pre_rel_times, - alpha=self.alpha, - balanced_cohorts=balanced_cohorts, - survey_weights_0=_sw_0_pre, - resolved_survey_full=_rs_full_pre, - n_obs_full=_n_full_pre, - omega_0_indices=_o0_idx_pre, - survey_df=_survey_df_pre, - ) - event_study_effects.update(pre_effects) - - # Collect horizons with Proposition 5 violations - prop5_horizons = [] - - for h in all_horizons: - if h == ref_period: - continue - - # Select treated obs at this horizon from balanced cohorts - h_mask = (rel_times == h) & balanced_mask.values - n_h = int(h_mask.sum()) - - if n_h == 0: - continue - - # Proposition 5 check - if not has_never_treated and h >= h_bar: - prop5_horizons.append(h) - event_study_effects[h] = { - "effect": np.nan, - "se": np.nan, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (np.nan, np.nan), - "n_obs": n_h, - } - continue - - tau_h = tau_hat[h_mask] - finite_h = np.isfinite(tau_h) - valid_tau = tau_h[finite_h] - - if len(valid_tau) == 0: - event_study_effects[h] = { - "effect": np.nan, - "se": np.nan, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (np.nan, np.nan), - "n_obs": n_h, - } - continue - - # Survey-weighted or simple mean for per-horizon effect - if survey_weights is not None: - treated_sw = survey_weights[omega_1_mask.values] - sw_h = treated_sw[h_mask] - sw_valid = sw_h[finite_h] - effect = float(np.average(valid_tau, weights=sw_valid)) - else: - effect = float(np.mean(valid_tau)) - - # Compute SE via conservative variance with horizon-specific weights - # When survey, aggregation weights are proportional to survey weights - if survey_weights is not None: - treated_sw = survey_weights[omega_1_mask.values] - n_1 = len(tau_hat) - weights_h = np.zeros(n_1) - sw_h = treated_sw[h_mask] - finite_in_h = np.isfinite(tau_h) - sw_finite = sw_h[finite_in_h] - # Set weights proportional to survey weights, summing to 1 - if sw_finite.sum() > 0: - h_indices = np.where(h_mask)[0] - finite_indices = h_indices[finite_in_h] - weights_h[finite_indices] = sw_finite / sw_finite.sum() - n_valid = int(finite_in_h.sum()) - else: - weights_h, n_valid = _compute_target_weights(tau_hat, h_mask) - - se = self._compute_conservative_variance( - df=df, - outcome=outcome, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - weights=weights_h, - cluster_var=cluster_var, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, - resolved_survey=resolved_survey, - proj_cache=proj_cache, - ) - - t_stat, p_value, conf_int = safe_inference(effect, se, alpha=self.alpha, df=survey_df) - - event_study_effects[h] = { - "effect": effect, - "se": se, - "t_stat": t_stat, - "p_value": p_value, - "conf_int": conf_int, - "n_obs": n_h, - } - - # Proposition 5 warning - if prop5_horizons: - warnings.warn( - f"Horizons {prop5_horizons} are not identified without " - f"never-treated units (Proposition 5). Set to NaN.", - UserWarning, - stacklevel=3, - ) - - # Check for empty result set after filtering - real_effects = [ - h for h, v in event_study_effects.items() if h != ref_period and v.get("n_obs", 0) > 0 - ] - if len(real_effects) == 0: - filter_info = [] - if balance_e is not None: - filter_info.append(f"balance_e={balance_e}") - if self.horizon_max is not None: - filter_info.append(f"horizon_max={self.horizon_max}") - filter_str = " and ".join(filter_info) if filter_info else "filters" - warnings.warn( - f"Event study aggregation produced no horizons with observations " - f"after applying {filter_str}. The result contains only the " - f"reference period marker. Consider relaxing filter parameters.", - UserWarning, - stacklevel=3, - ) - - return event_study_effects - - def _aggregate_group( - self, - df: pd.DataFrame, - outcome: str, - unit: str, - time: str, - first_treat: str, - covariates: Optional[List[str]], - omega_0_mask: pd.Series, - omega_1_mask: pd.Series, - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - grand_mean: float, - delta_hat: Optional[np.ndarray], - cluster_var: str, - treatment_groups: List[Any], - kept_cov_mask: Optional[np.ndarray] = None, - survey_weights: Optional[np.ndarray] = None, - survey_df: Optional[int] = None, - resolved_survey=None, - proj_cache: Optional[Dict[Any, _UntreatedProjection]] = None, - ) -> Dict[Any, Dict[str, Any]]: - """Aggregate treatment effects by cohort.""" - df_1 = df.loc[omega_1_mask] - tau_hat = df["_tau_hat"].loc[omega_1_mask].values - cohorts = df_1[first_treat].values - - group_effects: Dict[Any, Dict[str, Any]] = {} - - for g in treatment_groups: - g_mask = cohorts == g - n_g = int(g_mask.sum()) - - if n_g == 0: - continue - - tau_g = tau_hat[g_mask] - finite_g = np.isfinite(tau_g) - valid_tau = tau_g[finite_g] - - if len(valid_tau) == 0: - group_effects[g] = { - "effect": np.nan, - "se": np.nan, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (np.nan, np.nan), - "n_obs": n_g, - } - continue - - # Survey-weighted or simple mean for per-group effect - if survey_weights is not None: - treated_sw = survey_weights[omega_1_mask.values] - sw_g = treated_sw[g_mask] - sw_valid = sw_g[finite_g] - effect = float(np.average(valid_tau, weights=sw_valid)) - else: - effect = float(np.mean(valid_tau)) - - # Compute SE with group-specific weights - # When survey, aggregation weights proportional to survey weights - if survey_weights is not None: - treated_sw = survey_weights[omega_1_mask.values] - n_1 = len(tau_hat) - weights_g = np.zeros(n_1) - sw_g = treated_sw[g_mask] - sw_finite = sw_g[finite_g] - if sw_finite.sum() > 0: - g_indices = np.where(g_mask)[0] - finite_indices = g_indices[finite_g] - weights_g[finite_indices] = sw_finite / sw_finite.sum() - else: - weights_g, _ = _compute_target_weights(tau_hat, g_mask) - - se = self._compute_conservative_variance( - df=df, - outcome=outcome, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=unit_fe, - time_fe=time_fe, - grand_mean=grand_mean, - delta_hat=delta_hat, - weights=weights_g, - cluster_var=cluster_var, - kept_cov_mask=kept_cov_mask, - survey_weights=survey_weights, - resolved_survey=resolved_survey, - proj_cache=proj_cache, - ) - - t_stat, p_value, conf_int = safe_inference(effect, se, alpha=self.alpha, df=survey_df) - - group_effects[g] = { - "effect": effect, - "se": se, - "t_stat": t_stat, - "p_value": p_value, - "conf_int": conf_int, - "n_obs": n_g, - } - - return group_effects - # ========================================================================= # Pre-trend test (Equation 9) & pre-period lead coefficients # ========================================================================= - def _compute_lead_coefficients( - self, - df_0: pd.DataFrame, - outcome: str, - unit: str, - time: str, - first_treat: str, - covariates: Optional[List[str]], - cluster_var: str, - pre_rel_times: List[int], - alpha: float = 0.05, - balanced_cohorts: Optional[set] = None, - survey_weights_0: Optional[np.ndarray] = None, - resolved_survey_full=None, - n_obs_full: Optional[int] = None, - omega_0_indices: Optional[np.ndarray] = None, - survey_df: Optional[int] = None, - ) -> Tuple[Dict[int, Dict[str, Any]], np.ndarray, np.ndarray]: - """ - Compute pre-period lead coefficients via within-transformed OLS (Test 1). - - Adds lead indicator dummies W_it(h) = 1[K_it = h] to the untreated - model and estimates their coefficients. Uses cluster-robust SEs by - default, or design-based survey VCV when ``resolved_survey_full`` - is provided (subpopulation approach: scores zero-padded to full - panel length to preserve PSU/strata structure). - - The full Omega_0 sample (including never-treated controls) is always - used for within-transformation. When balanced_cohorts is provided, - lead indicators are restricted to observations from those cohorts only. - - Returns - ------- - effects : dict - Per-horizon event_study_effects entries. - gamma : ndarray - Lead coefficient vector. - V_gamma : ndarray - Sub-VCV matrix for lead coefficients. - """ - rel_time_0 = np.where( - ~df_0["_never_treated"], - df_0[time] - df_0[first_treat], - np.nan, - ) - - # Build lead indicators — restrict to balanced cohorts if specified - if balanced_cohorts is not None: - is_balanced = df_0[first_treat].isin(balanced_cohorts).values - else: - is_balanced = None - - lead_cols = [] - for h in pre_rel_times: - col_name = f"_lead_{h}" - indicator = (rel_time_0 == h).astype(float) - if is_balanced is not None: - indicator = indicator * is_balanced # zero out non-balanced cohorts - df_0[col_name] = indicator - lead_cols.append(col_name) - - all_x_cols = lead_cols[:] - if covariates: - all_x_cols.extend(covariates) - - # Within-transform through the shared MAP engine (survey-weighted when - # present), one dispatch for outcome + leads + covariates. Demean into - # a narrow copy: df_0's raw lead indicators must survive for the - # per-horizon n_obs counts below. within_transform pins [unit, time]; - # [time, unit] here preserves the historical time-then-unit sweep order. - narrow = df_0[[outcome, *all_x_cols, time, unit]].copy() - _pre_norms = pre_demean_norms(narrow, all_x_cols, weights=survey_weights_0) - demeaned, _ = demean_by_groups( - narrow, - [outcome, *all_x_cols], - [time, unit], - inplace=True, - weights=survey_weights_0, - max_iter=10_000, - tol=1e-10, - ) - # FE-spanned regressors demean to numerical junk, not exact zero; - # snap them so rank handling drops them deterministically (NaN - # coefficient for that horizon) instead of the junk direction - # perturbing the identified lead coefficients. Lead indicators are - # the most plausible FE-spanned regressors here: with a single - # (balanced-restricted) cohort a lead h collapses to a calendar-time - # dummy on Omega_0, which lies exactly in the span of the absorbed - # time FE. - snap_absorbed_regressors( - demeaned, - all_x_cols, - _pre_norms, - absorbed_desc="unit and time fixed effects (pretrends lead model)", - group_vars=[time, unit], - rank_deficient_action=self.rank_deficient_action, - display_names={f"_lead_{h}": f"lead[{h}]" for h in pre_rel_times}, - weights=survey_weights_0, - ) - y_dm = demeaned[outcome].to_numpy(dtype=np.float64) - X_dm = demeaned[all_x_cols].to_numpy(dtype=np.float64) - - # OLS for point estimates + VCV. When survey VCV will replace the - # cluster-robust VCV, skip cluster_ids to avoid errors on domains - # with few PSUs (the cluster-robust VCV is discarded anyway). - cluster_ids = df_0[cluster_var].values - _ols_weights = survey_weights_0 - _ols_weight_type = "pweight" if survey_weights_0 is not None else None - _use_survey_vcov = resolved_survey_full is not None - # Clustered-CR1 K_reference increment (variance-conventions.md D2): - # the demeaned lead design carries NO intercept column, so the - # absorbed constant contributes the +1 term. Survey fits pass 0 - # (cluster_ids is None there and the survey vcov replaces CR1). - _cr1_k_adj_imp = 0 - if not _use_survey_vcov: - _cr1_k_adj_imp = absorbed_fe_cr1_k_increment( - df_0, - [time, unit], - cluster_ids, - has_intercept_col=False, - weights=survey_weights_0, - ) - try: - result = solve_ols( - X_dm, - y_dm, - weights=_ols_weights, - weight_type=_ols_weight_type, - cluster_ids=None if _use_survey_vcov else cluster_ids, - cluster_k_adjustment=_cr1_k_adj_imp, - return_vcov=True, - rank_deficient_action=self.rank_deficient_action, - column_names=all_x_cols, - ) - except (IndexError, np.linalg.LinAlgError): - # All lead columns dropped (rank deficient after demeaning) - effects: Dict[int, Dict[str, Any]] = {} - for h in pre_rel_times: - n_obs = int(df_0[f"_lead_{h}"].sum()) - effects[h] = { - "effect": np.nan, - "se": np.nan, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (np.nan, np.nan), - "n_obs": n_obs, - } - for col in lead_cols: - df_0.drop(columns=col, inplace=True) - return ( - effects, - np.full(len(pre_rel_times), np.nan), - np.full((len(pre_rel_times), len(pre_rel_times)), np.nan), - ) - - coefficients = result[0] - vcov = result[2] - assert vcov is not None - - # Replace cluster-robust VCV with survey design-based VCV. - # Use the FULL survey design (subpopulation approach): zero-pad - # the Omega_0 scores back to full-panel length so PSU/strata - # structure is preserved for variance estimation. - if resolved_survey_full is not None: - from diff_diff.survey import compute_survey_vcov - - # Use residuals from solve_ols (safe for rank-deficient fits). - residuals_0 = result[1] - - # Reduce to kept (finite-coefficient) columns for VCV - kept_mask = np.isfinite(coefficients) - if np.all(kept_mask): - X_for_vcov = X_dm - res_for_vcov = residuals_0 - else: - X_for_vcov = X_dm[:, kept_mask] - res_for_vcov = residuals_0 - - # Zero-pad to full panel length (subpopulation approach): - # observations outside Omega_0 contribute zero to the score, - # but preserve PSU/strata structure for design-based variance. - # The survey full-design path always supplies the full obs count. - assert n_obs_full is not None - n_full_obs = n_obs_full - k_vcov = X_for_vcov.shape[1] - X_full = np.zeros((n_full_obs, k_vcov), dtype=np.float64) - res_full = np.zeros(n_full_obs, dtype=np.float64) - X_full[omega_0_indices] = X_for_vcov - res_full[omega_0_indices] = res_for_vcov - - vcov_kept = compute_survey_vcov(X_full, res_full, resolved_survey_full) - - if not np.all(kept_mask): - # Expand back: NaN rows/cols for dropped columns - n_coef = len(coefficients) - vcov = np.full((n_coef, n_coef), np.nan) - kept_idx = np.where(kept_mask)[0] - vcov[np.ix_(kept_idx, kept_idx)] = vcov_kept - else: - vcov = vcov_kept - - n_leads = len(lead_cols) - gamma = coefficients[:n_leads] - V_gamma = vcov[:n_leads, :n_leads] - - # Resolve the per-lead tail df: the full-design survey df keeps - # precedence on survey fits; otherwise resolve through the - # df_convention knob (3.9 fix: previously df=None → silent - # normal-theory z on plain clustered fits). Residual df = n − - # k_kept − absorbed [time, unit] rank (the demeaned lead design - # carries no intercept column); the "cluster" G counts - # positive-weight clusters on the SAME raw cluster ids the CR1 - # vcov partitioned on. NOTE this resolution serves BOTH callers — - # the fit-time pretrends event-study path and the post-fit public - # ``pretrend_test()`` (which consumes only gamma/V_gamma; its - # joint Wald F denominator is knob-independent). - if _use_survey_vcov: - _df = survey_df - else: - _k_kept_imp = int(np.count_nonzero(np.isfinite(coefficients))) - _df_res_imp = float( - len(df_0) - - _k_kept_imp - - absorbed_fe_rank( - df_0, - [time, unit], - has_intercept_col=False, - weights=survey_weights_0, - ) - ) - from diff_diff.linalg import effective_cluster_count - - _df = resolve_tail_df( - self.df_convention, - residual_df=_df_res_imp, - n_clusters=effective_cluster_count(cluster_ids, survey_weights_0), - ) - - # Build per-horizon effects - effects = {} - for j, h in enumerate(pre_rel_times): - effect = float(gamma[j]) - se = float(np.sqrt(max(V_gamma[j, j], 0.0))) - # n_obs from the lead indicator (respects balanced_cohorts restriction) - n_obs = int(df_0[f"_lead_{h}"].sum()) - t_stat, p_value, conf_int = safe_inference(effect, se, alpha=alpha, df=_df) - effects[h] = { - "effect": effect, - "se": se, - "t_stat": t_stat, - "p_value": p_value, - "conf_int": conf_int, - "n_obs": n_obs, - } - - # Clean up temporary columns - for col in lead_cols: - df_0.drop(columns=col, inplace=True) - - return effects, gamma, V_gamma - def _pretrend_test(self, n_leads: Optional[int] = None) -> Dict[str, Any]: """ Run pre-trend test (Equation 9). @@ -2763,63 +1217,6 @@ def _validate_leave_one_out(leave_one_out: Any) -> None: if not isinstance(leave_one_out, bool): raise TypeError(f"leave_one_out must be a bool, got {type(leave_one_out).__name__}") - @staticmethod - def _leave_one_out_factor( - per_unit: pd.DataFrame, per_group: pd.DataFrame - ) -> Tuple[pd.Series, int]: - """Per-(group, unit) leave-one-out residual-rescale factor (BJS 2024 A.9). - - ``factor_{g,i} = 1 / (1 - v_ig**2 / sum_j v_jg**2)`` with - ``v_ig = per_unit['a']`` and ``sum_j v_jg**2 = per_group['den']``. This - rescale of ``epsilon_tilde_it`` reproduces the direct leave-one-out - aggregate ``tau_tilde_it^LO`` exactly at the per-unit cluster sum - ``psi_i = sum_t v_it * epsilon_tilde_it`` (App. A.9). A group with a - single positive-weight unit has ``v_ig**2 == sum_j v_jg**2`` so the - factor diverges (LOO undefined, App. A.9 fn. 51); those groups fall back - to ``1.0`` (non-LOO). A genuinely unit-dominated but >=2-unit group keeps - its large finite factor -- that is the paper's intended inflation. - - Returns - ------- - (factor : pd.Series indexed like ``per_unit`` (g, u), n_single_unit_groups : int) - """ - a = per_unit["a"].to_numpy(dtype=float) - sq = a**2 - g_level = per_unit.index.get_level_values("g") - u_level = per_unit.index.get_level_values("u") - den = per_group["den"].reindex(g_level).to_numpy(dtype=float) # D_g per (g,u) - - # A group is "singleton" for LOO (App. A.9 fn. 51) when fewer than two - # units carry positive squared weight -- covers a true 1-unit group AND - # the effective-singleton case (>=2 rows, only one with a_ig != 0). - pos_per_group = pd.Series(sq > 0.0, index=per_unit.index).groupby(level="g").sum() - single_groups = pos_per_group.index[pos_per_group < 2] - is_single = pd.Index(g_level).isin(single_groups) - - den_ok = np.abs(den) >= 1e-15 - # factor = D_g / (D_g - v_ig^2) = D_g / sum_{j!=i} v_jg^2. Compute the - # leave-one-out denominator as the sum of the OTHER units' squared - # weights -- NOT as D_g - v_ig^2 after forming the ratio: for a genuinely - # dominated (but >=2-unit) group the subtraction loses precision (and can - # cancel to 0/negative) in float64 -- a finite-but-wrong or silently - # non-LOO factor. The fast subtraction is accurate away from the - # near-cancellation boundary; wherever the leave-one-out mass is a tiny - # fraction of D (relative loss of >~1e-6), recompute it exactly as the - # drop-then-sum of the OTHER units' squared weights. At most one unit per - # group can be that dominant, so the recompute stays O(units). - other_mass = den - sq - suspect = (~is_single) & den_ok & (other_mass <= 1e-6 * den) - if suspect.any(): - sq_series = pd.Series(sq, index=per_unit.index) - for pos in np.nonzero(suspect)[0]: - grp = sq_series.xs(g_level[pos], level="g") - other_mass[pos] = float(grp.drop(u_level[pos]).sum()) - # Fall back to non-LOO (factor 1.0) only where LOO is genuinely undefined: - # a singleton group (fn. 51), a degenerate den, or no other positive mass. - fallback = is_single | ~den_ok | (other_mass <= 0.0) - factor = np.where(fallback, 1.0, den / np.where(fallback, 1.0, other_mass)) - return pd.Series(factor, index=per_unit.index), int(len(single_groups)) - @staticmethod def _validate_vcov_type(vcov_type: str) -> None: """Validate ``vcov_type`` membership against ImputationDiD's @@ -2876,6 +1273,104 @@ def print_summary(self) -> None: print(self.summary()) +# ============================================================================= +# Post-fit aggregation kit (M-021/M-118) +# ============================================================================= + + +def _build_imputation_aggregation_kit( + *, + fit_data: Dict[str, Any], + treatment_groups: List[Any], + overall_att: float, + n_treated_obs: int, + uses_replicate: bool, + survey_df_seed: Optional[int], + survey_df_final: Optional[int], + survey_metadata: Optional[Any], + horizon_max: Optional[int], + pretrends: bool, + aux_partition: str, + leave_one_out: bool, + rank_deficient_action: str, + df_convention: str, + alpha: float, + anticipation: int, +) -> AggregationKit: + """Build the PANEL-BACKED post-fit aggregation kit (rows M-021/M-118). + + ImputationDiD's event-study/group aggregation is a target-specific + Theorem-3 recompute from the working panel + untreated FE model — no + compact influence payload can honor a different ``balance_e`` post-fit + — so ``bookkeeping`` holds REFERENCES to the SAME per-fit objects + ``self._fit_data`` already retains for ``pretrend_test()`` (passed in + as ``fit_data``; zero marginal memory, and pickles are unchanged via + memoization — ``_estimator_ref`` already ships these objects). fit() + rebinds a fresh ``_fit_data`` dict + a fresh working frame per call, + so kits from different fits never alias. + + Value SNAPSHOTS (not refs) isolate recompute from public-field + mutation: ``treatment_groups`` is copied (fit hands the same list + object to the results' public cohort list), scalar config is copied + by value, and + ``survey_metadata`` is a ``dataclasses.replace`` copy (the ES carrier + builds from the kit copy, never the mutable public field). The three + df-provenance channels: ``survey_df_seed`` (what the analytical + aggregators received — recompute re-seeds from it), + ``survey_df_final`` (what the stored overall inference received — + the 'simple' relay's df), and the metadata copy's own ``df_survey`` + (the fit-final container channel). + + ``influence`` is EMPTY BY DESIGN: the recompute is panel-backed, not + IF-payload-backed, and the fit-local projection cache holds + unpicklable factorizations — each ``aggregate()`` call rebuilds a + call-local ``proj_cache``. + """ + bookkeeping: Dict[str, Any] = { + # Panel-backed refs (the _fit_data objects) + "df": fit_data["df"], + "outcome": fit_data["outcome"], + "unit": fit_data["unit"], + "time": fit_data["time"], + "first_treat": fit_data["first_treat"], + "covariates": (list(fit_data["covariates"]) if fit_data["covariates"] else None), + "omega_0_mask": fit_data["omega_0_mask"], + "omega_1_mask": fit_data["omega_1_mask"], + "cluster_var": fit_data["cluster_var"], + "unit_fe": fit_data["unit_fe"], + "time_fe": fit_data["time_fe"], + "grand_mean": fit_data["grand_mean"], + "delta_hat": fit_data["delta_hat"], + "kept_cov_mask": fit_data["kept_cov_mask"], + "resolved_survey": fit_data["resolved_survey"], + "survey_weights": fit_data["survey_weights"], + # Value snapshots (isolation from public-field / estimator mutation) + "treatment_groups": list(treatment_groups), + "overall_att": float(overall_att), + "n_treated_obs": int(n_treated_obs), + "uses_replicate": bool(uses_replicate), + "horizon_max": horizon_max, + "pretrends": pretrends, + "aux_partition": aux_partition, + "leave_one_out": leave_one_out, + "rank_deficient_action": rank_deficient_action, + "df_convention": df_convention, + "survey_df_seed": survey_df_seed, + "survey_df_final": survey_df_final, + "survey_metadata": ( + dataclasses.replace(survey_metadata) if survey_metadata is not None else None + ), + } + return AggregationKit( + bookkeeping=bookkeeping, + influence={}, + alpha=alpha, + anticipation=anticipation, + cband=False, # no simultaneous-band concept on this estimator + bootstrap=None, # replay not wired; results.aggregate() fails closed + ) + + # ============================================================================= # Convenience function # ============================================================================= @@ -2888,8 +1383,8 @@ def imputation_did( 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["SurveyDesign"] = None, vcov_type: str = "hc1", **kwargs, @@ -2914,9 +1409,14 @@ def imputation_did( covariates : list of str, optional Covariate column names. aggregate : str, optional - Aggregation mode: None, "simple", "event_study", "group", "all". + DEPRECATED (3.9, removed in 4.0; row M-021): forwarded to ``fit()``, + which warns — aggregate post-fit via + ``results.aggregate('event_study')`` instead. A plain wrapper call + (kwarg not supplied) never warns. balance_e : int, optional - Balance event study to cohorts observed at all relative times. + DEPRECATED (3.9, removed in 4.0; row M-118): forwarded to ``fit()``, + which warns — moves onto ``results.aggregate('event_study', + balance_e=...)``. survey_design : SurveyDesign, optional Survey design specification for design-based inference. Supports pweight only (aweight/fweight raise ValueError). Supports strata, @@ -2943,9 +1443,9 @@ def imputation_did( -------- >>> from diff_diff import imputation_did, generate_staggered_data >>> data = generate_staggered_data(seed=42) - >>> results = imputation_did(data, 'outcome', 'unit', 'time', 'first_treat', - ... aggregate='event_study') + >>> results = imputation_did(data, 'outcome', 'unit', 'time', 'first_treat') >>> results.print_summary() + >>> results.aggregate('event_study').summary() # post-fit aggregation """ est = ImputationDiD(vcov_type=vcov_type, **kwargs) return est.fit( diff --git a/diff_diff/imputation_aggregation.py b/diff_diff/imputation_aggregation.py new file mode 100644 index 00000000..a9feaebc --- /dev/null +++ b/diff_diff/imputation_aggregation.py @@ -0,0 +1,1858 @@ +"""Aggregation + Theorem-3 variance engine for the BJS imputation estimator. + +Extracted verbatim from ``diff_diff/imputation.py`` for the M-021/M-118 +post-fit ``aggregate()`` migration: ``diff_diff/imputation.py`` imports +``imputation_results.py`` (and ``imputation_bootstrap.py`` imports it too), +so the results module can import neither -- the shared machinery lives here, +an import-leaf module both sides can reach (the +``efficient_did_aggregation.py`` precedent). + +Contents: + +- module helpers ``_compute_target_weights`` (lifted from + ``imputation_bootstrap.py`` -- the bootstrap module re-imports it; moving + it here keeps the module graph acyclic), ``_UntreatedProjection``, + ``_LSMRUnconvergedError`` and ``_lsmr_minnorm_normal_solve``; +- :class:`_ImputationAggregationMixin` -- the event-study / group + aggregators, the Theorem-3 conservative-variance stack they recompute + through, the pretrends lead regression, and the replicate-weight + inference override replay. Inherited by ``ImputationDiD`` (fit-time + behavior byte-identical) and hosted post-fit by the throwaway + ``_ImputationKitAggregator`` (``imputation_results.py``). +""" + +import warnings +from typing import Any, Callable, Dict, List, NamedTuple, Optional, Set, Tuple + +import numpy as np +import pandas as pd +from scipy import sparse +from scipy.sparse.linalg import factorized as sparse_factorized + +from diff_diff.linalg import solve_ols +from diff_diff.utils import ( + _iterative_fe_solve, + absorbed_fe_cr1_k_increment, + absorbed_fe_rank, + demean_by_groups, + pre_demean_norms, + resolve_tail_df, + safe_inference, + snap_absorbed_regressors, +) + + +def _compute_target_weights( + tau_hat: np.ndarray, + target_mask: np.ndarray, +) -> "tuple[np.ndarray, int]": + """ + Equal weights for finite tau_hat observations within target_mask. + + Used by both aggregation and bootstrap paths to avoid weight logic + duplication. + + Parameters + ---------- + tau_hat : np.ndarray + Per-observation treatment effects (may contain NaN). + target_mask : np.ndarray + Boolean mask selecting the target subset within tau_hat. + + Returns + ------- + weights : np.ndarray + Weight array (same length as tau_hat). 1/n_valid for finite + observations in target_mask, 0 elsewhere. + n_valid : int + Number of finite observations in the target subset. + """ + finite_target = np.isfinite(tau_hat) & target_mask + n_valid = int(finite_target.sum()) + weights = np.zeros(len(tau_hat)) + if n_valid > 0: + weights[np.where(finite_target)[0]] = 1.0 / n_valid + return weights, n_valid + + +class _UntreatedProjection(NamedTuple): + """Cached, target-invariant pieces of the untreated imputation projection + ``v_untreated = -A_0 (A_0' [W] A_0)^{-1} A_1' w`` (BJS 2024 Theorem 3). + + Within a single ``fit()`` the untreated design (``df_0``/``df_1``, covariates, + survey weights) is identical across every estimand target (overall ATT, each + event-study horizon, each group, and the bootstrap precompute) -- only the + treated aggregation ``weights`` (the RHS ``A_1' w``) vary. So ``A_0``, ``A_1`` + and the factorization of ``A_0'[W]A_0`` are built once and reused across + targets (factorize-once / solve-many), mirroring the TwoStageDiD GMM-sandwich + ``sparse_factorized`` pattern. + """ + + A_0: sparse.csr_matrix + A_1: sparse.csr_matrix + # solver(rhs) -> z; None when the factorization was exactly singular (the + # solve path then routes to the sparse LSMR least-squares fallback). + solver: Optional[Callable[[np.ndarray], np.ndarray]] + A0tA0_csc: sparse.csc_matrix # retained for the LSMR fallback + survey_weights_0: Optional[np.ndarray] + singular: bool + + +class _LSMRUnconvergedError(RuntimeError): + """LSMR failed to certify a solution on the singular-variance fallback. + + Raised (not returned as NaN) so the variance boundary can fail closed: + a NaN vector would be laundered into zeros by the missing-FE + ``nan_to_num`` in the psi product — producing a finite, WRONG variance — + whereas this exception is caught in ``_compute_conservative_variance`` + and converted to a NaN SE (the all-or-nothing NaN inference convention). + """ + + +def _lsmr_minnorm_normal_solve(A0tA0_csc, rhs: np.ndarray) -> np.ndarray: + """Least-squares solve of the (possibly singular) normal equations + ``(A_0'[W]A_0) z = rhs`` WITHOUT densifying the sparse matrix. + + Replaces the previous ``np.linalg.lstsq(A0tA0.toarray(), ...)`` fallback, + whose dense materialization scales ``O((U+T+K)^2)`` — an OOM risk on + large panels (the TODO row this resolves). ``scipy.sparse.linalg.lsmr`` + handles singular symmetric systems, converging to the minimum-norm + least-squares solution (the same solution family as ``lstsq``'s + pseudo-inverse solution). + + Solver choice cannot change the estimator output: any two least-squares + solutions differ by a ``null(A_0'[W]A_0) = null(sqrt(W) A_0)`` component, + which the downstream projection ``v_untreated = -[W_0] A_0 z`` + annihilates (unweighted: ``null = null(A_0)`` so ``A_0 z`` is invariant; + weighted: the weight multiplication zeroes exactly the rows where the + null component can be nonzero). Locked by the singular-system parity + test against a dense-lstsq oracle. + + CONVERGENCE IS VALIDATED (fail-closed): ``istop`` in ``{0, 1, 2, 4, 5}`` + means LSMR certified an (approximate) solution / least-squares solution + within ``atol``/``btol`` (4 and 5 are the machine-precision analogues of + 1 and 2 per SciPy's documentation); anything else (condition-limit stop, + max-iteration exhaustion) gets ONE retry with an uncapped condition + limit and a generous iteration budget, and if still uncertified raises + :class:`_LSMRUnconvergedError` — caught at the variance boundary and + converted to a NaN SE — rather than feeding a finite-but-unverified + solution into the Theorem 3 weights. + """ + import scipy.sparse.linalg as spla + + _certified = (0, 1, 2, 4, 5) + result = spla.lsmr(A0tA0_csc, rhs, atol=1e-14, btol=1e-14) + z, istop = result[0], int(result[1]) + if istop not in _certified or not np.all(np.isfinite(z)): + dim = A0tA0_csc.shape[0] + result = spla.lsmr( + A0tA0_csc, rhs, atol=1e-14, btol=1e-14, conlim=1e16, maxiter=max(50 * dim, 10_000) + ) + z, istop = result[0], int(result[1]) + if istop not in _certified or not np.all(np.isfinite(z)): + warnings.warn( + "ImputationDiD variance: the LSMR fallback solve of " + f"(A_0'[W]A_0) z = rhs did not converge (istop={istop}); " + "the affected variance is reported as NaN rather than from " + "an unverified solution.", + UserWarning, + stacklevel=3, + ) + raise _LSMRUnconvergedError(f"LSMR uncertified (istop={istop})") + return z + + +class _ImputationAggregationMixin: + """Shared aggregation/variance methods (moved verbatim from ``ImputationDiD``). + + HOST-ATTRIBUTE CONTRACT -- the complete ``self.`` surface the moved + methods read (typed class-level declarations, not docstring prose: + ``mypy diff_diff`` at zero errors needs the attributes declared on the + mixin for both hosts). Zero methods WRITE to ``self`` -- the post-fit + throwaway host exists for estimator-mutation isolation only. + """ + + alpha: float + anticipation: int + horizon_max: Optional[int] + pretrends: bool + aux_partition: str + leave_one_out: bool + rank_deficient_action: str + df_convention: str + + def _iterative_fe( + self, + y: np.ndarray, + unit_vals: np.ndarray, + time_vals: np.ndarray, + idx: pd.Index, + max_iter: int = 10_000, + tol: float = 1e-10, + weights: Optional[np.ndarray] = None, + ) -> Tuple[Dict[Any, float], Dict[Any, float]]: + """ + Estimate unit and time FE via iterative alternating projection (Gauss-Seidel). + + Thin wrapper over the shared bincount solver + (``diff_diff.utils._iterative_fe_solve``): factorize unit/time once, + solve on integer codes, map the level arrays back to dicts. + Converges to the exact (W)LS solution for balanced and unbalanced + panels; balanced panels converge in 1-2 iterations. + + Parameters + ---------- + idx : pd.Index + Unused; retained for call-site stability. + weights : np.ndarray, optional + Survey weights (weighted group means ``sum(w*x)/sum(w)``). A + unit/period whose observations ALL carry zero weight has no + identifying contribution and gets ``NaN`` FE (its key is kept so + the rank-condition membership check still sees the group). + + Returns + ------- + unit_fe : dict + Mapping from unit -> unit fixed effect. + time_fe : dict + Mapping from time -> time fixed effect. + """ + unit_codes, unit_uniques = pd.factorize(unit_vals, sort=False) + time_codes, time_uniques = pd.factorize(time_vals, sort=False) + if (unit_codes < 0).any() or (time_codes < 0).any(): + raise ValueError( + "ImputationDiD: unit or time column contains NaN. Drop or " + "impute missing group keys before fitting." + ) + unit_fe_arr, time_fe_arr = _iterative_fe_solve( + np.asarray(y, dtype=np.float64), + unit_codes.astype(np.intp, copy=False), + time_codes.astype(np.intp, copy=False), + len(unit_uniques), + len(time_uniques), + weights=weights, + max_iter=max_iter, + tol=tol, + method_name="ImputationDiD iterative FE solver", + ) + unit_fe = dict(zip(unit_uniques, unit_fe_arr)) + time_fe = dict(zip(time_uniques, time_fe_arr)) + return unit_fe, time_fe + + @staticmethod + def _compute_balanced_cohort_mask( + df_treated: pd.DataFrame, + first_treat: str, + all_horizons: List[int], + balance_e: int, + cohort_rel_times: Dict[Any, Set[int]], + ) -> np.ndarray: + """Compute boolean mask selecting treated obs from balanced cohorts. + + A cohort is 'balanced' if it has observations at every relative time + in [-balance_e, max(all_horizons)]. + + Parameters + ---------- + df_treated : pd.DataFrame + Post-treatment observations (Omega_1). + first_treat : str + Column name for cohort identifier. + all_horizons : list of int + Post-treatment horizons in the event study. + balance_e : int + Number of pre-treatment periods to require. + cohort_rel_times : dict + Maps each cohort value to the set of all observed relative times + (including pre-treatment) from the full panel. Built by + _build_cohort_rel_times(). + """ + if not all_horizons: + return np.ones(len(df_treated), dtype=bool) + + max_h = max(all_horizons) + required_range = set(range(-balance_e, max_h + 1)) + + balanced_cohorts = set() + for g, horizons in cohort_rel_times.items(): + if required_range.issubset(horizons): + balanced_cohorts.add(g) + + return df_treated[first_treat].isin(balanced_cohorts).values + + @staticmethod + def _build_cohort_rel_times( + df: pd.DataFrame, + first_treat: str, + ) -> Dict[Any, Set[int]]: + """Build mapping of cohort -> set of observed relative times from full panel. + + Precondition: df must have '_never_treated' and '_rel_time' columns + (set by fit() before any aggregation calls). + """ + treated_mask = ~df["_never_treated"] + treated_df = df.loc[treated_mask] + result: Dict[Any, Set[int]] = {} + ft_vals = treated_df[first_treat].values + rt_vals = treated_df["_rel_time"].values + for i in range(len(treated_df)): + h = rt_vals[i] + if np.isfinite(h): + result.setdefault(ft_vals[i], set()).add(int(h)) + return result + + def _fit_untreated_model( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + weights: Optional[np.ndarray] = None, + ) -> Tuple[ + Dict[Any, float], Dict[Any, float], float, Optional[np.ndarray], Optional[np.ndarray] + ]: + """ + Step 1: Estimate unit + time FE on untreated observations. + + Uses iterative alternating projection (Gauss-Seidel) to compute exact + OLS fixed effects for both balanced and unbalanced panels. For balanced + panels, converges in 1-2 iterations (identical to one-pass demeaning). + + Parameters + ---------- + weights : np.ndarray, optional + Full-panel survey weights (same length as df). The untreated subset + is extracted internally via omega_0_mask. When None, unweighted. + + Returns + ------- + unit_fe : dict + Unit fixed effects {unit_id: alpha_i}. + time_fe : dict + Time fixed effects {time_period: beta_t}. + grand_mean : float + Grand mean (0.0 — absorbed into iterative FE). + delta_hat : np.ndarray or None + Covariate coefficients (if covariates provided). + kept_cov_mask : np.ndarray or None + Boolean mask of shape (n_covariates,) indicating which covariates + have finite coefficients. None if no covariates. + """ + df_0 = df.loc[omega_0_mask] + w_0 = weights[omega_0_mask.values] if weights is not None else None + + if covariates is None or len(covariates) == 0: + # No covariates: estimate FE via iterative alternating projection + # (exact OLS for both balanced and unbalanced panels) + y = df_0[outcome].values.copy() + unit_fe, time_fe = self._iterative_fe( + y, df_0[unit].values, df_0[time].values, df_0.index, weights=w_0 + ) + # grand_mean = 0: iterative FE absorb the intercept + return unit_fe, time_fe, 0.0, None, None + + else: + # With covariates: iteratively demean Y and X, OLS for delta, + # then recover FE from covariate-adjusted outcome + y = df_0[outcome].values.copy() + X_raw = df_0[covariates].values.copy() + units = df_0[unit].values + times = df_0[time].values + + # Step A: within-transform Y and all X columns through the shared + # MAP engine (factorize-once + bincount + optional Rust kernel), + # one dispatch for every column. within_transform pins + # [unit, time]; [time, unit] here preserves the historical + # time-then-unit sweep order of the per-estimator loops. + narrow = df_0[[outcome, *covariates, time, unit]].copy() + demeaned, _ = demean_by_groups( + narrow, + [outcome, *covariates], + [time, unit], + inplace=True, + weights=w_0, + max_iter=10_000, + tol=1e-10, + ) + y_dm = demeaned[outcome].to_numpy(dtype=np.float64) + X_dm = demeaned[covariates].to_numpy(dtype=np.float64) + + # Step B: OLS for covariate coefficients on demeaned data + result = solve_ols( + X_dm, + y_dm, + return_vcov=False, + rank_deficient_action=self.rank_deficient_action, + column_names=covariates, + weights=w_0, + ) + delta_hat = result[0] + + # Mask of covariates with finite coefficients (before cleaning) + # Used to exclude rank-deficient covariates from variance design matrices + kept_cov_mask = np.isfinite(delta_hat) + + # Replace NaN coefficients with 0 for adjustment + # (rank-deficient covariates are dropped) + delta_hat_clean = np.where(np.isfinite(delta_hat), delta_hat, 0.0) + + # Step C: Recover FE from covariate-adjusted outcome using iterative FE + y_adj = y - np.dot(X_raw, delta_hat_clean) + unit_fe, time_fe = self._iterative_fe(y_adj, units, times, df_0.index, weights=w_0) + + # grand_mean = 0: iterative FE absorb the intercept + return unit_fe, time_fe, 0.0, delta_hat_clean, kept_cov_mask + + def _impute_treatment_effects( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + covariates: Optional[List[str]], + omega_1_mask: pd.Series, + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + ) -> Tuple[np.ndarray, np.ndarray]: + """ + Step 2: Impute Y(0) for treated observations and compute tau_hat. + + Returns + ------- + tau_hat : np.ndarray + Imputed treatment effects for each treated observation. + y_hat_0 : np.ndarray + Imputed counterfactual Y(0). + """ + df_1 = df.loc[omega_1_mask] + + # Look up unit and time FE + alpha_i = df_1[unit].map(unit_fe).values + beta_t = df_1[time].map(time_fe).values + + # Handle missing FE (set to NaN) + alpha_i = np.where(pd.isna(alpha_i), np.nan, alpha_i).astype(float) + beta_t = np.where(pd.isna(beta_t), np.nan, beta_t).astype(float) + + y_hat_0 = grand_mean + alpha_i + beta_t + + if delta_hat is not None and covariates: + X_1 = df_1[covariates].values + y_hat_0 = y_hat_0 + np.dot(X_1, delta_hat) + + tau_hat = df_1[outcome].values - y_hat_0 + + return tau_hat, y_hat_0 + + def _compute_cluster_psi_sums( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + omega_1_mask: pd.Series, + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + weights: np.ndarray, + cluster_var: str, + kept_cov_mask: Optional[np.ndarray] = None, + survey_weights_0: Optional[np.ndarray] = None, + proj_cache: Optional[Dict[Any, _UntreatedProjection]] = None, + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Compute cluster-level influence function sums (Theorem 3). + + psi_i = sum_t v_it * epsilon_tilde_it, summed within each cluster. + + Returns + ------- + cluster_psi_sums : np.ndarray + Array of cluster-level psi sums. + cluster_ids_unique : np.ndarray + Unique cluster identifiers (matching order of psi sums). + """ + df_0 = df.loc[omega_0_mask] + df_1 = df.loc[omega_1_mask] + + # ---- Compute v_it for treated observations ---- + v_treated = weights.copy() + + # ---- Compute v_it for untreated observations ---- + # Exact two-way-FE imputation projection + # v_untreated = -A_0 (A_0' [W] A_0)^{-1} A_1' w_treated (Theorem 3 / the + # implied weights of Supplementary Proposition A3), used for BOTH the + # FE-only and the covariate case. The earlier FE-only closed form + # -(w_i/n0_i + w_t/n0_t - w/N_0) is exact only for a *balanced* untreated + # panel; Omega_0 is generically unbalanced in staggered designs (treated + # observations are removed), which biased the analytical SE downward + # (~27% on the parity panel). The projection matches R `didimputation` + # exactly -- see tests/test_methodology_imputation.py::TestImputationDiDParityR. + # Build the target-invariant projection design + factorization once per + # fit() (cached in proj_cache), then solve only the target-specific RHS. + # survey_weights is DELIBERATELY excluded from the key: the cache is a + # fit-LOCAL dict, and within one fit() survey_weights is a single fixed + # object, so the masks deterministically map to one sw_0 = + # survey_weights[omega_0_mask]. The masks + covariates + kept_cov_mask + # therefore FULLY identify the design (sw_0 itself is a fresh-sliced array + # per call -- keying on its id() would miss every time and balloon the + # cache to 1+H+G full A_0/A_1/factorization entries). id()-keys are safe: + # the masks are fit() locals alive for the whole fit and the cache is a + # fit-local dict, so no cross-fit leak / id reuse. + cov_list = covariates if covariates is not None else [] + ctx: Optional[_UntreatedProjection] = None + if proj_cache is not None: + key = ( + id(omega_0_mask), + id(omega_1_mask), + tuple(cov_list), + kept_cov_mask.tobytes() if kept_cov_mask is not None else None, + ) + ctx = proj_cache.get(key) + if ctx is None: + ctx = self._build_untreated_projection( + df_0, + df_1, + unit, + time, + cov_list, + kept_cov_mask=kept_cov_mask, + survey_weights_0=survey_weights_0, + ) + if proj_cache is not None: + proj_cache[key] = ctx + v_untreated = self._solve_untreated_v(ctx, weights) + + # ---- Compute auxiliary model residuals (Equation 8) ---- + epsilon_treated = self._compute_auxiliary_residuals_treated( + df_1, + outcome, + unit, + time, + first_treat, + covariates, + unit_fe, + time_fe, + grand_mean, + delta_hat, + v_treated, + ) + epsilon_untreated = self._compute_residuals_untreated( + df_0, outcome, unit, time, covariates, unit_fe, time_fe, grand_mean, delta_hat + ) + + # ---- psi_it = v_it * epsilon_tilde_it ---- + v_all = np.empty(len(df)) + v_all[omega_1_mask.values] = v_treated + v_all[omega_0_mask.values] = v_untreated + + eps_all = np.empty(len(df)) + eps_all[omega_1_mask.values] = epsilon_treated + eps_all[omega_0_mask.values] = epsilon_untreated + + ve_product = v_all * eps_all + # NaN eps from missing FE (rank condition violation). Zero their variance + # contribution — matches R's did_imputation which drops unimputable obs. + np.nan_to_num(ve_product, copy=False, nan=0.0) + + # Sum within clusters + cluster_ids = df[cluster_var].values + ve_series = pd.Series(ve_product, index=df.index) + cluster_sums = ve_series.groupby(cluster_ids).sum() + + return cluster_sums.values, cluster_sums.index.values, ve_product + + def _compute_conservative_variance( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + omega_1_mask: pd.Series, + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + weights: np.ndarray, + cluster_var: str, + kept_cov_mask: Optional[np.ndarray] = None, + survey_weights: Optional[np.ndarray] = None, + resolved_survey=None, + proj_cache: Optional[Dict[Any, _UntreatedProjection]] = None, + ) -> float: + """ + Compute conservative clustered variance (Theorem 3, Equation 7). + + Parameters + ---------- + weights : np.ndarray + Aggregation weights w_it for treated observations. + Shape: (n_treated,), must sum to 1. + survey_weights : np.ndarray, optional + Full-panel survey weights. When provided, they enter the untreated + v_it WLS projection (weighted normal equations plus the left + per-observation weight factor) and the design-based variance path. + resolved_survey : ResolvedSurveyDesign, optional + When provided, uses design-based variance via + ``compute_survey_if_variance()`` (supports strata, PSU, FPC). + + Returns + ------- + float + Standard error. + """ + sw_0 = survey_weights[omega_0_mask.values] if survey_weights is not None else None + try: + cluster_psi_sums, _, ve_product = self._compute_cluster_psi_sums( + df=df, + outcome=outcome, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=unit_fe, + time_fe=time_fe, + grand_mean=grand_mean, + delta_hat=delta_hat, + weights=weights, + cluster_var=cluster_var, + kept_cov_mask=kept_cov_mask, + survey_weights_0=sw_0, + proj_cache=proj_cache, + ) + except _LSMRUnconvergedError: + # Solver failure is GLOBAL (the untreated projection is invalid), + # unlike per-observation missing-FE NaNs — fail the whole SE + # closed instead of letting nan_to_num launder it to zeros. + return np.nan + + if resolved_survey is not None: + # Design-based variance with strata/PSU/FPC support + from diff_diff.survey import compute_survey_if_variance + + variance = compute_survey_if_variance(ve_product, resolved_survey) + if np.isnan(variance): + return np.nan + return np.sqrt(max(variance, 0.0)) + + sigma_sq = float((cluster_psi_sums**2).sum()) + return np.sqrt(max(sigma_sq, 0.0)) + + def _build_untreated_projection( + self, + df_0: pd.DataFrame, + df_1: pd.DataFrame, + unit: str, + time: str, + covariates: List[str], + kept_cov_mask: Optional[np.ndarray] = None, + survey_weights_0: Optional[np.ndarray] = None, + ) -> _UntreatedProjection: + """ + Build the target-INVARIANT pieces of the exact imputation projection + ``v_untreated = -A_0 (A_0' [W] A_0)^{-1} A_1' w_treated`` and factorize the + normal-equations matrix once. The result is cached per ``fit()`` (see + ``_compute_cluster_psi_sums``) and reused across all estimand targets; + only the target-specific RHS ``A_1' w`` is solved per target in + ``_solve_untreated_v``. + + This is the GENERAL path -- used for both the FE-only and the covariate + cases (an empty ``covariates`` list builds a pure two-way-FE design; + ``n_cov == 0`` is the FE-only path). When survey_weights_0 is provided, + uses the weighted normal equations ``A_0' W A_0`` (the per-observation + survey weight is reapplied to the solved v in ``_solve_untreated_v``). + + Uses scipy.sparse for FE dummy columns to reduce memory from O(N*(U+T)) + to O(N) for the FE portion. An exactly singular ``A_0'[W]A_0`` makes + ``sparse_factorized`` raise ``RuntimeError``; we emit a UserWarning (once + per fit) and record ``singular=True`` so the solve routes to the sparse + LSMR least-squares fallback (no dense materialization; see + :func:`_lsmr_minnorm_normal_solve`). + """ + # Exclude rank-deficient covariates from design matrices + if kept_cov_mask is not None and not np.all(kept_cov_mask): + covariates = [c for c, k in zip(covariates, kept_cov_mask) if k] + + units_0 = df_0[unit].values + times_0 = df_0[time].values + units_1 = df_1[unit].values + times_1 = df_1[time].values + + all_units = np.unique(np.concatenate([units_0, units_1])) + all_times = np.unique(np.concatenate([times_0, times_1])) + unit_to_idx = {u: i for i, u in enumerate(all_units)} + time_to_idx = {t: i for i, t in enumerate(all_times)} + n_units = len(all_units) + n_times = len(all_times) + n_cov = len(covariates) + # Two-way FE design = all unit dummies (their sum spans the intercept) + + # time dummies dropping the first (identification). Dropping the first + # unit dummy too -- with no intercept column -- would omit the baseline + # level dimension and project onto a space one rank short of the true + # two-way-FE span, biasing the imputation weights (and hence the SE). + n_fe_cols = n_units + (n_times - 1) + + def _build_A_sparse(df_sub, unit_vals, time_vals): + n = len(df_sub) + + # Unit dummies — keep ALL (together they span the intercept). + u_indices = np.array([unit_to_idx[u] for u in unit_vals]) + u_rows = np.arange(n) + u_cols = u_indices + + # Time dummies (drop first) — vectorized + t_indices = np.array([time_to_idx[t] for t in time_vals]) + t_mask = t_indices > 0 + t_rows = np.arange(n)[t_mask] + t_cols = n_units + (t_indices[t_mask] - 1) + + rows = np.concatenate([u_rows, t_rows]) + cols = np.concatenate([u_cols, t_cols]) + data = np.ones(len(rows)) + + A_fe = sparse.csr_matrix((data, (rows, cols)), shape=(n, n_fe_cols)) + + # Covariates (dense, typically few columns) + if n_cov > 0: + A_cov = sparse.csr_matrix(df_sub[covariates].values) + A = sparse.hstack([A_fe, A_cov], format="csr") + else: + A = A_fe + + return A + + A_0 = _build_A_sparse(df_0, units_0, times_0) + A_1 = _build_A_sparse(df_1, units_1, times_1) + + # Form (A_0' [W] A_0). When survey weights present, use the weighted + # normal equations A_0' W A_0. + if survey_weights_0 is not None: + A0tA0_sparse = A_0.T @ A_0.multiply(survey_weights_0[:, None]) + else: + A0tA0_sparse = A_0.T @ A_0 # stays sparse + A0tA0_csc = A0tA0_sparse.tocsc() + + # Factorize once (factorize-once / solve-many). An exactly singular + # matrix makes sparse_factorized raise RuntimeError -- the same condition + # that previously surfaced as spsolve's MatrixRankWarning -> non-finite + # solution. Warn once and fall back to the sparse LSMR least-squares + # solve per target (no dense materialization). (The factorized path is + # bit-identical to the prior per-target spsolve for a single dense + # RHS -- both use the SuperLU simple driver with the same defaults.) + try: + solver: Optional[Callable[[np.ndarray], np.ndarray]] = sparse_factorized(A0tA0_csc) + singular = False + except RuntimeError as exc: + # Silent-failure audit axis C: emit a UserWarning on fallback instead + # of swallowing the error. Keep the "sparse LSMR" substring (asserted + # by tests). + warnings.warn( + "ImputationDiD variance: sparse factorization of (A_0' [W] A_0) " + f"failed ({type(exc).__name__}); falling back to a sparse LSMR " + "least-squares solve (no dense materialization). This may " + "indicate a rank-deficient or near-singular normal-equations " + "matrix and variance estimates may be less reliable.", + UserWarning, + stacklevel=2, + ) + solver = None + singular = True + + return _UntreatedProjection( + A_0=A_0, + A_1=A_1, + solver=solver, + A0tA0_csc=A0tA0_csc, + survey_weights_0=survey_weights_0, + singular=singular, + ) + + def _solve_untreated_v(self, ctx: _UntreatedProjection, weights: np.ndarray) -> np.ndarray: + """ + Solve the target-SPECIFIC RHS of the untreated imputation projection using + the cached design + factorization in ``ctx``: + ``v_untreated = -[W_0] A_0 (A_0'[W]A_0)^{-1} A_1' w_treated``. + """ + A1_w = ctx.A_1.T @ weights # (p,) + + if ctx.singular: + # Factorization was singular at build time (warned once already). + z = _lsmr_minnorm_normal_solve(ctx.A0tA0_csc, A1_w) + else: + assert ctx.solver is not None + z = ctx.solver(A1_w) + if not np.all(np.isfinite(z)): + # Defensive, target-specific: a non-finite solve on an otherwise + # factorizable matrix routes this RHS to the LSMR fallback. Warn per + # target (silent-failure audit axis C) -- distinct from the + # once-per-fit build-time singular warning. + warnings.warn( + "ImputationDiD variance: sparse solve of (A_0' [W] A_0) z = " + "A_1' w returned a non-finite solution; falling back to a " + "sparse LSMR least-squares solve for this target. Variance " + "estimates may be less reliable.", + UserWarning, + stacklevel=2, + ) + z = _lsmr_minnorm_normal_solve(ctx.A0tA0_csc, A1_w) + + # v_untreated = -[W_0] A_0 z (WLS projection requires per-obs weight) + v_untreated = -(ctx.A_0 @ z) + if ctx.survey_weights_0 is not None: + v_untreated = v_untreated * ctx.survey_weights_0 + return v_untreated + + def _compute_auxiliary_residuals_treated( + self, + df_1: pd.DataFrame, + outcome: str, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + v_treated: np.ndarray, + ) -> np.ndarray: + """ + Compute auxiliary residuals for treated obs (Theorem 3, Equation 8). + + Implements the paper's *unit-clustered* group aggregator (Borusyak, + Jaravel & Spiess 2024, eq. 8, p. 3272), which minimizes the excess + variance of the conservative estimator under a within-group + constant-effect auxiliary model (Supplementary Appendix A.8): + + tau_tilde_g = sum_i (sum_{t in G_g,i} v_it)(sum_{t in G_g,i} v_it * tau_hat_it) + ---------------------------------------------------------------- + sum_i (sum_{t in G_g,i} v_it)^2 + + i.e. for each unit i form the within-unit weight sum a_{i,g} and the + within-unit weighted-effect sum b_{i,g} over the unit's observations in + group g, then combine across units. At the default cohort x event-time + partition (<=1 obs/unit/group) this reduces to sum(v^2 * tau_hat) / + sum(v^2) -- the form the R `didimputation` package implements -- and + equals the naive observation-level mean sum(v * tau_hat) / sum(v) only + when within-group weights are uniform. Under coarser `cohort` / `horizon` + partitions (a unit contributes several observations to a group) or + non-uniform v_it (e.g. survey weights) the two genuinely differ. + + epsilon_tilde_it = Y_it - alpha_i - beta_t [- X'delta] - tau_tilde_g + """ + n_1 = len(df_1) + + # Compute base residuals (Y - Y_hat(0) = tau_hat) + # NaN for missing FE (consistent with _impute_treatment_effects) + alpha_i = df_1[unit].map(unit_fe).values.astype(float) # NaN for missing + beta_t = df_1[time].map(time_fe).values.astype(float) # NaN for missing + y_hat_0 = grand_mean + alpha_i + beta_t + + if delta_hat is not None and covariates: + y_hat_0 = y_hat_0 + np.dot(df_1[covariates].values, delta_hat) + + tau_hat = df_1[outcome].values - y_hat_0 + + # Partition Omega_1 into groups G_g + if self.aux_partition == "cohort_horizon": + group_keys = list(zip(df_1[first_treat].values, df_1["_rel_time"].values)) + elif self.aux_partition == "cohort": + group_keys = list(df_1[first_treat].values) + elif self.aux_partition == "horizon": + group_keys = list(df_1["_rel_time"].values) + else: + group_keys = list(range(n_1)) # each obs is its own group + + # Factorize group keys to integer codes (robust to tuple-valued keys). + group_codes = pd.factorize(pd.Series(group_keys), sort=False)[0] + gc_series = pd.Series(group_codes, index=df_1.index) + tau_series = pd.Series(tau_hat, index=df_1.index) + + # Unit-clustered Equation 8. Only v_it != 0 observations contribute: a + # zero-weight row adds exactly 0 to both a_{i,g} and b_{i,g}, so dropping + # it is exact for finite tau_hat AND avoids letting an unimputable row + # (NaN tau_hat, which always carries v_it == 0 by construction in + # _compute_target_weights) poison its whole group via 0 * NaN = NaN. The + # previous observation-level pandas sum relied on skipna to drop them. + contrib = (v_treated != 0.0) & np.isfinite(tau_hat) + loo_factor: Optional[pd.Series] = None + n_single_loo = 0 + if contrib.any(): + inner = pd.DataFrame( + { + "g": group_codes[contrib], + "u": df_1[unit].values[contrib], + "v": v_treated[contrib], + "vt": v_treated[contrib] * tau_hat[contrib], + } + ) + # Per (group, unit): a_{i,g} = sum v_it, b_{i,g} = sum v_it * tau_hat + per_unit = inner.groupby(["g", "u"], sort=False).agg(a=("v", "sum"), b=("vt", "sum")) + # Per group: numerator sum_i a*b, denominator sum_i a^2 + per_group = ( + per_unit.assign(ab=per_unit["a"] * per_unit["b"], a2=per_unit["a"] ** 2) + .groupby(level="g") + .agg(num=("ab", "sum"), den=("a2", "sum")) + ) + den_ok = per_group["den"].abs() >= 1e-15 + tau_tilde_map = (per_group["num"] / per_group["den"]).where(den_ok) + # BJS 2024 App. A.9 leave-one-out refinement: rescale each treated + # residual by 1/(1 - v_ig^2 / sum_j v_jg^2) (== the direct-LOO tau_tilde + # exactly, at the per-unit cluster sum). Reuses a_{i,g} = per_unit['a'] + # and sum_j v_jg^2 = per_group['den']; applied to epsilon_treated below. + if self.leave_one_out: + loo_factor, n_single_loo = self._leave_one_out_factor(per_unit, per_group) + else: + tau_tilde_map = pd.Series(dtype=float) + + tau_tilde_per_obs = gc_series.map(tau_tilde_map) + + # Groups with no contributing (v_it != 0, finite tau_hat) observations -- + # e.g. off-target horizons in an event-study SE -- are a variance no-op + # (psi_g = sum_t v_it * eps_tilde_it = 0 there regardless of tau_tilde_g), + # so fall back to the unweighted group mean of tau_hat for a finite value. + if tau_tilde_per_obs.isna().any(): + simple_means = tau_series.groupby(gc_series).mean() + tau_tilde_per_obs = tau_tilde_per_obs.fillna(gc_series.map(simple_means)) + + tau_tilde = tau_tilde_per_obs.values + + # Auxiliary residuals + epsilon_treated = tau_hat - tau_tilde + + # Leave-one-out rescale (BJS 2024 App. A.9): map each treated obs to its + # (group, unit) factor and inflate the residual. Non-contributing rows + # (v_it == 0, psi == 0 anyway) and single-positive-weight-unit groups + # (LOO undefined, fn. 51) keep factor 1.0. + if self.leave_one_out and loo_factor is not None: + obs_index = pd.MultiIndex.from_arrays( + [group_codes, df_1[unit].values], names=["g", "u"] + ) + factor_per_obs = loo_factor.reindex(obs_index).to_numpy(dtype=float) + factor_per_obs = np.where(np.isfinite(factor_per_obs), factor_per_obs, 1.0) + epsilon_treated = epsilon_treated * factor_per_obs + if n_single_loo > 0: + warnings.warn( + f"leave_one_out=True: {n_single_loo} auxiliary group(s) have a single " + f"positive-weight unit, where the leave-one-out variance is undefined " + f"(Borusyak, Jaravel & Spiess 2024, Supp. App. A.9 fn. 51); those groups " + f"keep the non-leave-out residual. A coarser aux_partition reduces " + f"singleton groups.", + UserWarning, + stacklevel=2, + ) + + return epsilon_treated + + def _compute_residuals_untreated( + self, + df_0: pd.DataFrame, + outcome: str, + unit: str, + time: str, + covariates: Optional[List[str]], + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + ) -> np.ndarray: + """Compute Step 1 residuals for untreated observations.""" + # Preserve NaN for any missing FE, symmetric with the treated path in + # _compute_auxiliary_residuals_treated. On valid data this is inert -- + # every untreated observation's unit and period appear in the Step 1 FE + # dicts (the dicts are estimated FROM Omega_0) -- but it stops a missing + # FE from silently becoming a 0 residual, which would mask a rank- + # condition logic error. Any NaN is zeroed downstream in the variance + # product (np.nan_to_num), exactly like the treated path. + alpha_i = df_0[unit].map(unit_fe).values.astype(float) + beta_t = df_0[time].map(time_fe).values.astype(float) + y_hat = grand_mean + alpha_i + beta_t + + if delta_hat is not None and covariates: + y_hat = y_hat + np.dot(df_0[covariates].values, delta_hat) + + return df_0[outcome].values - y_hat + + def _aggregate_event_study( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + omega_1_mask: pd.Series, + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + cluster_var: str, + treatment_groups: List[Any], + balance_e: Optional[int] = None, + kept_cov_mask: Optional[np.ndarray] = None, + survey_weights: Optional[np.ndarray] = None, + survey_df: Optional[int] = None, + resolved_survey=None, + proj_cache: Optional[Dict[Any, _UntreatedProjection]] = None, + ) -> Dict[int, Dict[str, Any]]: + """Aggregate treatment effects by event-study horizon.""" + df_1 = df.loc[omega_1_mask] + tau_hat = df["_tau_hat"].loc[omega_1_mask].values + rel_times = df_1["_rel_time"].values + + # Get all horizons + all_horizons = sorted(set(int(h) for h in rel_times if np.isfinite(h))) + + # Apply horizon_max filter + if self.horizon_max is not None: + all_horizons = [h for h in all_horizons if abs(h) <= self.horizon_max] + + # Apply balance_e filter + if balance_e is not None: + cohort_rel_times = self._build_cohort_rel_times(df, first_treat) + balanced_mask = pd.Series( + self._compute_balanced_cohort_mask( + df_1, first_treat, all_horizons, balance_e, cohort_rel_times + ), + index=df_1.index, + ) + else: + balanced_mask = pd.Series(True, index=df_1.index) + + # Check Proposition 5: no never-treated units + has_never_treated = df["_never_treated"].any() + h_bar = np.inf + if not has_never_treated and len(treatment_groups) > 1: + h_bar = max(treatment_groups) - min(treatment_groups) + + # Reference period + ref_period = -1 - self.anticipation + + event_study_effects: Dict[int, Dict[str, Any]] = {} + + # Add reference period marker + event_study_effects[ref_period] = { + "effect": 0.0, + "se": 0.0, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (0.0, 0.0), + "n_obs": 0, + } + + # Pre-period coefficients via BJS Test 1 lead regression + if self.pretrends: + df_0 = df.loc[omega_0_mask].copy() + + # Determine which cohorts' lead indicators to include. + # balance_e restricts which cohorts contribute lead dummies, + # but the full Omega_0 sample (including never-treated controls) + # is kept for the within-transformed OLS (BJS Test 1, Equation 9). + balanced_cohorts = None + skip_preperiods = False + if balance_e is not None: + cohort_rel_times_0 = self._build_cohort_rel_times(df, first_treat) + balanced_cohorts = set() + if all_horizons: + max_h = max(all_horizons) + required_range = set(range(-balance_e, max_h + 1)) + for g, horizons in cohort_rel_times_0.items(): + if required_range.issubset(horizons): + balanced_cohorts.add(g) + if not balanced_cohorts: + skip_preperiods = True # No cohorts qualify — skip entirely + + if not skip_preperiods: + rel_time_0 = np.where( + ~df_0["_never_treated"], + df_0[time] - df_0[first_treat], + np.nan, + ) + + # When balance_e is set, only include leads from balanced cohorts + if balanced_cohorts is not None: + is_balanced = df_0[first_treat].isin(balanced_cohorts).values + rel_time_for_leads = np.where(is_balanced, rel_time_0, np.nan) + else: + rel_time_for_leads = rel_time_0 + + pre_rel_times = sorted( + set( + int(h) + for h in rel_time_for_leads + if np.isfinite(h) and h < -self.anticipation + ) + ) + pre_rel_times = [h for h in pre_rel_times if h != ref_period] + if self.horizon_max is not None: + pre_rel_times = [h for h in pre_rel_times if abs(h) <= self.horizon_max] + if pre_rel_times: + # Survey pretrends: pass full design (subpopulation approach) + _sw_0_pre = None + _rs_full_pre = None + _n_full_pre = None + _o0_idx_pre = None + if survey_weights is not None and resolved_survey is not None: + _sw_0_pre = survey_weights[omega_0_mask.values] + _rs_full_pre = resolved_survey + _n_full_pre = len(df) + _o0_idx_pre = np.where(omega_0_mask.values)[0] + _survey_df_pre = ( + resolved_survey.df_survey if resolved_survey is not None else None + ) + pre_effects, _, _ = self._compute_lead_coefficients( + df_0, + outcome, + unit, + time, + first_treat, + covariates, + cluster_var, + pre_rel_times, + alpha=self.alpha, + balanced_cohorts=balanced_cohorts, + survey_weights_0=_sw_0_pre, + resolved_survey_full=_rs_full_pre, + n_obs_full=_n_full_pre, + omega_0_indices=_o0_idx_pre, + survey_df=_survey_df_pre, + ) + event_study_effects.update(pre_effects) + + # Collect horizons with Proposition 5 violations + prop5_horizons = [] + + for h in all_horizons: + if h == ref_period: + continue + + # Select treated obs at this horizon from balanced cohorts + h_mask = (rel_times == h) & balanced_mask.values + n_h = int(h_mask.sum()) + + if n_h == 0: + continue + + # Proposition 5 check + if not has_never_treated and h >= h_bar: + prop5_horizons.append(h) + event_study_effects[h] = { + "effect": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "n_obs": n_h, + } + continue + + tau_h = tau_hat[h_mask] + finite_h = np.isfinite(tau_h) + valid_tau = tau_h[finite_h] + + if len(valid_tau) == 0: + event_study_effects[h] = { + "effect": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "n_obs": n_h, + } + continue + + # Survey-weighted or simple mean for per-horizon effect + if survey_weights is not None: + treated_sw = survey_weights[omega_1_mask.values] + sw_h = treated_sw[h_mask] + sw_valid = sw_h[finite_h] + effect = float(np.average(valid_tau, weights=sw_valid)) + else: + effect = float(np.mean(valid_tau)) + + # Compute SE via conservative variance with horizon-specific weights + # When survey, aggregation weights are proportional to survey weights + if survey_weights is not None: + treated_sw = survey_weights[omega_1_mask.values] + n_1 = len(tau_hat) + weights_h = np.zeros(n_1) + sw_h = treated_sw[h_mask] + finite_in_h = np.isfinite(tau_h) + sw_finite = sw_h[finite_in_h] + # Set weights proportional to survey weights, summing to 1 + if sw_finite.sum() > 0: + h_indices = np.where(h_mask)[0] + finite_indices = h_indices[finite_in_h] + weights_h[finite_indices] = sw_finite / sw_finite.sum() + n_valid = int(finite_in_h.sum()) + else: + weights_h, n_valid = _compute_target_weights(tau_hat, h_mask) + + se = self._compute_conservative_variance( + df=df, + outcome=outcome, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=unit_fe, + time_fe=time_fe, + grand_mean=grand_mean, + delta_hat=delta_hat, + weights=weights_h, + cluster_var=cluster_var, + kept_cov_mask=kept_cov_mask, + survey_weights=survey_weights, + resolved_survey=resolved_survey, + proj_cache=proj_cache, + ) + + t_stat, p_value, conf_int = safe_inference(effect, se, alpha=self.alpha, df=survey_df) + + event_study_effects[h] = { + "effect": effect, + "se": se, + "t_stat": t_stat, + "p_value": p_value, + "conf_int": conf_int, + "n_obs": n_h, + } + + # Proposition 5 warning + if prop5_horizons: + warnings.warn( + f"Horizons {prop5_horizons} are not identified without " + f"never-treated units (Proposition 5). Set to NaN.", + UserWarning, + stacklevel=3, + ) + + # Check for empty result set after filtering + real_effects = [ + h for h, v in event_study_effects.items() if h != ref_period and v.get("n_obs", 0) > 0 + ] + if len(real_effects) == 0: + filter_info = [] + if balance_e is not None: + filter_info.append(f"balance_e={balance_e}") + if self.horizon_max is not None: + filter_info.append(f"horizon_max={self.horizon_max}") + filter_str = " and ".join(filter_info) if filter_info else "filters" + warnings.warn( + f"Event study aggregation produced no horizons with observations " + f"after applying {filter_str}. The result contains only the " + f"reference period marker. Consider relaxing filter parameters.", + UserWarning, + stacklevel=3, + ) + + return event_study_effects + + def _aggregate_group( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + omega_1_mask: pd.Series, + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + cluster_var: str, + treatment_groups: List[Any], + kept_cov_mask: Optional[np.ndarray] = None, + survey_weights: Optional[np.ndarray] = None, + survey_df: Optional[int] = None, + resolved_survey=None, + proj_cache: Optional[Dict[Any, _UntreatedProjection]] = None, + ) -> Dict[Any, Dict[str, Any]]: + """Aggregate treatment effects by cohort.""" + df_1 = df.loc[omega_1_mask] + tau_hat = df["_tau_hat"].loc[omega_1_mask].values + cohorts = df_1[first_treat].values + + group_effects: Dict[Any, Dict[str, Any]] = {} + + for g in treatment_groups: + g_mask = cohorts == g + n_g = int(g_mask.sum()) + + if n_g == 0: + continue + + tau_g = tau_hat[g_mask] + finite_g = np.isfinite(tau_g) + valid_tau = tau_g[finite_g] + + if len(valid_tau) == 0: + group_effects[g] = { + "effect": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "n_obs": n_g, + } + continue + + # Survey-weighted or simple mean for per-group effect + if survey_weights is not None: + treated_sw = survey_weights[omega_1_mask.values] + sw_g = treated_sw[g_mask] + sw_valid = sw_g[finite_g] + effect = float(np.average(valid_tau, weights=sw_valid)) + else: + effect = float(np.mean(valid_tau)) + + # Compute SE with group-specific weights + # When survey, aggregation weights proportional to survey weights + if survey_weights is not None: + treated_sw = survey_weights[omega_1_mask.values] + n_1 = len(tau_hat) + weights_g = np.zeros(n_1) + sw_g = treated_sw[g_mask] + sw_finite = sw_g[finite_g] + if sw_finite.sum() > 0: + g_indices = np.where(g_mask)[0] + finite_indices = g_indices[finite_g] + weights_g[finite_indices] = sw_finite / sw_finite.sum() + else: + weights_g, _ = _compute_target_weights(tau_hat, g_mask) + + se = self._compute_conservative_variance( + df=df, + outcome=outcome, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=unit_fe, + time_fe=time_fe, + grand_mean=grand_mean, + delta_hat=delta_hat, + weights=weights_g, + cluster_var=cluster_var, + kept_cov_mask=kept_cov_mask, + survey_weights=survey_weights, + resolved_survey=resolved_survey, + proj_cache=proj_cache, + ) + + t_stat, p_value, conf_int = safe_inference(effect, se, alpha=self.alpha, df=survey_df) + + group_effects[g] = { + "effect": effect, + "se": se, + "t_stat": t_stat, + "p_value": p_value, + "conf_int": conf_int, + "n_obs": n_g, + # Per-row df provenance (M-021): the exact df THIS row's + # safe_inference received. Capture-at-use — the replicate + # override rewrites it (with the row's se/t/p/CI) and the + # bootstrap override clears it; the all-NaN cohort branch + # above writes no key (consumers read via .get). + "df_used": survey_df, + } + + return group_effects + + def _compute_lead_coefficients( + self, + df_0: pd.DataFrame, + outcome: str, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + cluster_var: str, + pre_rel_times: List[int], + alpha: float = 0.05, + balanced_cohorts: Optional[set] = None, + survey_weights_0: Optional[np.ndarray] = None, + resolved_survey_full=None, + n_obs_full: Optional[int] = None, + omega_0_indices: Optional[np.ndarray] = None, + survey_df: Optional[int] = None, + ) -> Tuple[Dict[int, Dict[str, Any]], np.ndarray, np.ndarray]: + """ + Compute pre-period lead coefficients via within-transformed OLS (Test 1). + + Adds lead indicator dummies W_it(h) = 1[K_it = h] to the untreated + model and estimates their coefficients. Uses cluster-robust SEs by + default, or design-based survey VCV when ``resolved_survey_full`` + is provided (subpopulation approach: scores zero-padded to full + panel length to preserve PSU/strata structure). + + The full Omega_0 sample (including never-treated controls) is always + used for within-transformation. When balanced_cohorts is provided, + lead indicators are restricted to observations from those cohorts only. + + Returns + ------- + effects : dict + Per-horizon event_study_effects entries. + gamma : ndarray + Lead coefficient vector. + V_gamma : ndarray + Sub-VCV matrix for lead coefficients. + """ + rel_time_0 = np.where( + ~df_0["_never_treated"], + df_0[time] - df_0[first_treat], + np.nan, + ) + + # Build lead indicators — restrict to balanced cohorts if specified + if balanced_cohorts is not None: + is_balanced = df_0[first_treat].isin(balanced_cohorts).values + else: + is_balanced = None + + lead_cols = [] + for h in pre_rel_times: + col_name = f"_lead_{h}" + indicator = (rel_time_0 == h).astype(float) + if is_balanced is not None: + indicator = indicator * is_balanced # zero out non-balanced cohorts + df_0[col_name] = indicator + lead_cols.append(col_name) + + all_x_cols = lead_cols[:] + if covariates: + all_x_cols.extend(covariates) + + # Within-transform through the shared MAP engine (survey-weighted when + # present), one dispatch for outcome + leads + covariates. Demean into + # a narrow copy: df_0's raw lead indicators must survive for the + # per-horizon n_obs counts below. within_transform pins [unit, time]; + # [time, unit] here preserves the historical time-then-unit sweep order. + narrow = df_0[[outcome, *all_x_cols, time, unit]].copy() + _pre_norms = pre_demean_norms(narrow, all_x_cols, weights=survey_weights_0) + demeaned, _ = demean_by_groups( + narrow, + [outcome, *all_x_cols], + [time, unit], + inplace=True, + weights=survey_weights_0, + max_iter=10_000, + tol=1e-10, + ) + # FE-spanned regressors demean to numerical junk, not exact zero; + # snap them so rank handling drops them deterministically (NaN + # coefficient for that horizon) instead of the junk direction + # perturbing the identified lead coefficients. Lead indicators are + # the most plausible FE-spanned regressors here: with a single + # (balanced-restricted) cohort a lead h collapses to a calendar-time + # dummy on Omega_0, which lies exactly in the span of the absorbed + # time FE. + snap_absorbed_regressors( + demeaned, + all_x_cols, + _pre_norms, + absorbed_desc="unit and time fixed effects (pretrends lead model)", + group_vars=[time, unit], + rank_deficient_action=self.rank_deficient_action, + display_names={f"_lead_{h}": f"lead[{h}]" for h in pre_rel_times}, + weights=survey_weights_0, + ) + y_dm = demeaned[outcome].to_numpy(dtype=np.float64) + X_dm = demeaned[all_x_cols].to_numpy(dtype=np.float64) + + # OLS for point estimates + VCV. When survey VCV will replace the + # cluster-robust VCV, skip cluster_ids to avoid errors on domains + # with few PSUs (the cluster-robust VCV is discarded anyway). + cluster_ids = df_0[cluster_var].values + _ols_weights = survey_weights_0 + _ols_weight_type = "pweight" if survey_weights_0 is not None else None + _use_survey_vcov = resolved_survey_full is not None + # Clustered-CR1 K_reference increment (variance-conventions.md D2): + # the demeaned lead design carries NO intercept column, so the + # absorbed constant contributes the +1 term. Survey fits pass 0 + # (cluster_ids is None there and the survey vcov replaces CR1). + _cr1_k_adj_imp = 0 + if not _use_survey_vcov: + _cr1_k_adj_imp = absorbed_fe_cr1_k_increment( + df_0, + [time, unit], + cluster_ids, + has_intercept_col=False, + weights=survey_weights_0, + ) + try: + result = solve_ols( + X_dm, + y_dm, + weights=_ols_weights, + weight_type=_ols_weight_type, + cluster_ids=None if _use_survey_vcov else cluster_ids, + cluster_k_adjustment=_cr1_k_adj_imp, + return_vcov=True, + rank_deficient_action=self.rank_deficient_action, + column_names=all_x_cols, + ) + except (IndexError, np.linalg.LinAlgError): + # All lead columns dropped (rank deficient after demeaning) + effects: Dict[int, Dict[str, Any]] = {} + for h in pre_rel_times: + n_obs = int(df_0[f"_lead_{h}"].sum()) + effects[h] = { + "effect": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "n_obs": n_obs, + } + for col in lead_cols: + df_0.drop(columns=col, inplace=True) + return ( + effects, + np.full(len(pre_rel_times), np.nan), + np.full((len(pre_rel_times), len(pre_rel_times)), np.nan), + ) + + coefficients = result[0] + vcov = result[2] + assert vcov is not None + + # Replace cluster-robust VCV with survey design-based VCV. + # Use the FULL survey design (subpopulation approach): zero-pad + # the Omega_0 scores back to full-panel length so PSU/strata + # structure is preserved for variance estimation. + if resolved_survey_full is not None: + from diff_diff.survey import compute_survey_vcov + + # Use residuals from solve_ols (safe for rank-deficient fits). + residuals_0 = result[1] + + # Reduce to kept (finite-coefficient) columns for VCV + kept_mask = np.isfinite(coefficients) + if np.all(kept_mask): + X_for_vcov = X_dm + res_for_vcov = residuals_0 + else: + X_for_vcov = X_dm[:, kept_mask] + res_for_vcov = residuals_0 + + # Zero-pad to full panel length (subpopulation approach): + # observations outside Omega_0 contribute zero to the score, + # but preserve PSU/strata structure for design-based variance. + # The survey full-design path always supplies the full obs count. + assert n_obs_full is not None + n_full_obs = n_obs_full + k_vcov = X_for_vcov.shape[1] + X_full = np.zeros((n_full_obs, k_vcov), dtype=np.float64) + res_full = np.zeros(n_full_obs, dtype=np.float64) + X_full[omega_0_indices] = X_for_vcov + res_full[omega_0_indices] = res_for_vcov + + vcov_kept = compute_survey_vcov(X_full, res_full, resolved_survey_full) + + if not np.all(kept_mask): + # Expand back: NaN rows/cols for dropped columns + n_coef = len(coefficients) + vcov = np.full((n_coef, n_coef), np.nan) + kept_idx = np.where(kept_mask)[0] + vcov[np.ix_(kept_idx, kept_idx)] = vcov_kept + else: + vcov = vcov_kept + + n_leads = len(lead_cols) + gamma = coefficients[:n_leads] + V_gamma = vcov[:n_leads, :n_leads] + + # Resolve the per-lead tail df: the full-design survey df keeps + # precedence on survey fits; otherwise resolve through the + # df_convention knob (3.9 fix: previously df=None → silent + # normal-theory z on plain clustered fits). Residual df = n − + # k_kept − absorbed [time, unit] rank (the demeaned lead design + # carries no intercept column); the "cluster" G counts + # positive-weight clusters on the SAME raw cluster ids the CR1 + # vcov partitioned on. NOTE this resolution serves BOTH callers — + # the fit-time pretrends event-study path and the post-fit public + # ``pretrend_test()`` (which consumes only gamma/V_gamma; its + # joint Wald F denominator is knob-independent). + if _use_survey_vcov: + _df = survey_df + else: + _k_kept_imp = int(np.count_nonzero(np.isfinite(coefficients))) + _df_res_imp = float( + len(df_0) + - _k_kept_imp + - absorbed_fe_rank( + df_0, + [time, unit], + has_intercept_col=False, + weights=survey_weights_0, + ) + ) + from diff_diff.linalg import effective_cluster_count + + _df = resolve_tail_df( + self.df_convention, + residual_df=_df_res_imp, + n_clusters=effective_cluster_count(cluster_ids, survey_weights_0), + ) + + # Build per-horizon effects + effects = {} + for j, h in enumerate(pre_rel_times): + effect = float(gamma[j]) + se = float(np.sqrt(max(V_gamma[j, j], 0.0))) + # n_obs from the lead indicator (respects balanced_cohorts restriction) + n_obs = int(df_0[f"_lead_{h}"].sum()) + t_stat, p_value, conf_int = safe_inference(effect, se, alpha=alpha, df=_df) + effects[h] = { + "effect": effect, + "se": se, + "t_stat": t_stat, + "p_value": p_value, + "conf_int": conf_int, + "n_obs": n_obs, + } + + # Clean up temporary columns + for col in lead_cols: + df_0.drop(columns=col, inplace=True) + + return effects, gamma, V_gamma + + @staticmethod + def _leave_one_out_factor( + per_unit: pd.DataFrame, per_group: pd.DataFrame + ) -> Tuple[pd.Series, int]: + """Per-(group, unit) leave-one-out residual-rescale factor (BJS 2024 A.9). + + ``factor_{g,i} = 1 / (1 - v_ig**2 / sum_j v_jg**2)`` with + ``v_ig = per_unit['a']`` and ``sum_j v_jg**2 = per_group['den']``. This + rescale of ``epsilon_tilde_it`` reproduces the direct leave-one-out + aggregate ``tau_tilde_it^LO`` exactly at the per-unit cluster sum + ``psi_i = sum_t v_it * epsilon_tilde_it`` (App. A.9). A group with a + single positive-weight unit has ``v_ig**2 == sum_j v_jg**2`` so the + factor diverges (LOO undefined, App. A.9 fn. 51); those groups fall back + to ``1.0`` (non-LOO). A genuinely unit-dominated but >=2-unit group keeps + its large finite factor -- that is the paper's intended inflation. + + Returns + ------- + (factor : pd.Series indexed like ``per_unit`` (g, u), n_single_unit_groups : int) + """ + a = per_unit["a"].to_numpy(dtype=float) + sq = a**2 + g_level = per_unit.index.get_level_values("g") + u_level = per_unit.index.get_level_values("u") + den = per_group["den"].reindex(g_level).to_numpy(dtype=float) # D_g per (g,u) + + # A group is "singleton" for LOO (App. A.9 fn. 51) when fewer than two + # units carry positive squared weight -- covers a true 1-unit group AND + # the effective-singleton case (>=2 rows, only one with a_ig != 0). + pos_per_group = pd.Series(sq > 0.0, index=per_unit.index).groupby(level="g").sum() + single_groups = pos_per_group.index[pos_per_group < 2] + is_single = pd.Index(g_level).isin(single_groups) + + den_ok = np.abs(den) >= 1e-15 + # factor = D_g / (D_g - v_ig^2) = D_g / sum_{j!=i} v_jg^2. Compute the + # leave-one-out denominator as the sum of the OTHER units' squared + # weights -- NOT as D_g - v_ig^2 after forming the ratio: for a genuinely + # dominated (but >=2-unit) group the subtraction loses precision (and can + # cancel to 0/negative) in float64 -- a finite-but-wrong or silently + # non-LOO factor. The fast subtraction is accurate away from the + # near-cancellation boundary; wherever the leave-one-out mass is a tiny + # fraction of D (relative loss of >~1e-6), recompute it exactly as the + # drop-then-sum of the OTHER units' squared weights. At most one unit per + # group can be that dominant, so the recompute stays O(units). + other_mass = den - sq + suspect = (~is_single) & den_ok & (other_mass <= 1e-6 * den) + if suspect.any(): + sq_series = pd.Series(sq, index=per_unit.index) + for pos in np.nonzero(suspect)[0]: + grp = sq_series.xs(g_level[pos], level="g") + other_mass[pos] = float(grp.drop(u_level[pos]).sum()) + # Fall back to non-LOO (factor 1.0) only where LOO is genuinely undefined: + # a singleton group (fn. 51), a degenerate den, or no other positive mass. + fallback = is_single | ~den_ok | (other_mass <= 0.0) + factor = np.where(fallback, 1.0, den / np.where(fallback, 1.0, other_mass)) + return pd.Series(factor, index=per_unit.index), int(len(single_groups)) + + def _replicate_override_aggregates( + self, + *, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + omega_1_mask: pd.Series, + resolved_survey: Any, + overall_att: float, + event_study_effects: Optional[Dict[int, Dict[str, Any]]], + group_effects: Optional[Dict[Any, Dict[str, Any]]], + balance_e: Optional[int], + survey_df_seed: Optional[int], + ) -> Tuple[np.ndarray, int, Optional[int]]: + """Replicate-weight inference override for the aggregation surfaces. + + Extracted verbatim from ``fit()``'s replicate block so the post-fit + ``aggregate()`` path can replay it from the kit payload. The stacked + layout is ``[overall, es..., grp...]`` built from whichever family + dicts are non-None. ``compute_replicate_refit_variance`` validates + replicates JOINTLY (all-finite rows, ``survey.py``), so inference is + LEVEL-MATCHED: post-fit ``aggregate('event_study')`` replays + ``[overall, ES]`` and reproduces ``fit(aggregate='event_study')`` + exactly; a ``fit(aggregate='all')`` surface is NOT the equivalence + target when a replicate NaNs on exactly one family's targets (its + joint stack drops that replicate for every row). + + Overrides ``se``/``t_stat``/``p_value``/``conf_int`` (and, on group + rows, the ``df_used`` provenance key) IN PLACE on the passed dicts + and returns ``(vcov, n_valid, survey_df_final)``. The fit-only tail + -- the overall-row quintet and the ``survey_metadata.df_survey`` + write -- deliberately stays in ``fit()``: this method must never + mutate results-owned metadata (post-fit ``aggregate()`` calls it on + kit refs and a throwaway host). Cost: one full per-replicate refit + pass per call (R refits), the same work the fit-time path does. + """ + from diff_diff.survey import compute_replicate_refit_variance + + _rel_times_treated = df.loc[omega_1_mask, "_rel_time"].values + _cohorts_treated = df.loc[omega_1_mask, first_treat].values + + # Derive keys from actual outputs (excludes filtered/Prop5/ref) + _es_effects = event_study_effects or {} + _grp_effects = group_effects or {} + _sorted_rel_times = sorted( + e + for e in _es_effects.keys() + if np.isfinite(_es_effects[e]["effect"]) and _es_effects[e].get("n_obs", 1) > 0 + ) + _sorted_groups = sorted( + g for g in _grp_effects.keys() if np.isfinite(_grp_effects[g]["effect"]) + ) + _n_es = len(_sorted_rel_times) + + # Pre-compute balanced cohort mask for balance_e + _balanced_mask_treated = None + if balance_e is not None and _sorted_rel_times: + df_1 = df.loc[omega_1_mask] + rel_times_all = df_1["_rel_time"].values + all_horizons_full = sorted(set(int(h) for h in rel_times_all if np.isfinite(h))) + if self.horizon_max is not None: + all_horizons_full = [h for h in all_horizons_full if abs(h) <= self.horizon_max] + cohort_rel_times = self._build_cohort_rel_times(df, first_treat) + _balanced_mask_treated = self._compute_balanced_cohort_mask( + df_1, first_treat, all_horizons_full, balance_e, cohort_rel_times + ) + + # Single vectorized refit: [overall, es_e0..., grp_g0...] + def _refit_imp(w_r): + ufe_r, tfe_r, gm_r, delta_r, _ = self._fit_untreated_model( + df, + outcome, + unit, + time, + covariates, + omega_0_mask, + weights=w_r, + ) + tau_r, _ = self._impute_treatment_effects( + df, + outcome, + unit, + time, + covariates, + omega_1_mask, + ufe_r, + tfe_r, + gm_r, + delta_r, + ) + fin = np.isfinite(tau_r) + treated_w = w_r[omega_1_mask.values] + results = [] + # [0] Overall ATT + tw_fin = treated_w[fin] + tw_sum = np.sum(tw_fin) + results.append(float(np.sum(tau_r[fin] * tw_fin) / tw_sum) if tw_sum > 0 else np.nan) + # [1..n_es] Event-study (identified only) + for e in _sorted_rel_times: + mask_e = fin & (_rel_times_treated == e) + if _balanced_mask_treated is not None: + mask_e = mask_e & _balanced_mask_treated + tw_e = treated_w[mask_e] + s = np.sum(tw_e) + results.append(float(np.sum(tau_r[mask_e] * tw_e) / s) if s > 0 else np.nan) + # [n_es+1..] Group (identified only) + for g in _sorted_groups: + mask_g = fin & (_cohorts_treated == g) + tw_g = treated_w[mask_g] + s = np.sum(tw_g) + results.append(float(np.sum(tau_r[mask_g] * tw_g) / s) if s > 0 else np.nan) + return np.array(results) + + # Build full-sample estimate from actual effects + _full_est = [overall_att] + _full_est.extend([_es_effects[e]["effect"] for e in _sorted_rel_times]) + _full_est.extend([_grp_effects[g]["effect"] for g in _sorted_groups]) + + _vcov_rep_imp, _n_valid_rep_imp = compute_replicate_refit_variance( + _refit_imp, np.array(_full_est), resolved_survey + ) + + # Override df if replicates were dropped + # Replicate-refit path is only reached with a resolved design. + assert resolved_survey is not None + survey_df_final = survey_df_seed + if _n_valid_rep_imp < resolved_survey.n_replicates: + survey_df_final = _n_valid_rep_imp - 1 if _n_valid_rep_imp > 1 else 0 + + # Override event-study SEs from vcov diagonal + for i, e in enumerate(_sorted_rel_times): + if event_study_effects is not None and e in event_study_effects: + se_e = float(np.sqrt(max(_vcov_rep_imp[1 + i, 1 + i], 0.0))) + eff_e = event_study_effects[e]["effect"] + t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=survey_df_final) + event_study_effects[e]["se"] = se_e + event_study_effects[e]["t_stat"] = t_e + event_study_effects[e]["p_value"] = p_e + event_study_effects[e]["conf_int"] = ci_e + + # Override group SEs from vcov diagonal + for j, g in enumerate(_sorted_groups): + if group_effects is not None and g in group_effects: + se_g = float(np.sqrt(max(_vcov_rep_imp[1 + _n_es + j, 1 + _n_es + j], 0.0))) + eff_g = group_effects[g]["effect"] + t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=survey_df_final) + group_effects[g]["se"] = se_g + group_effects[g]["t_stat"] = t_g + group_effects[g]["p_value"] = p_g + group_effects[g]["conf_int"] = ci_g + # the override rewrote se/t/p/CI under the final df -- + # the per-row provenance key follows the rewrite + group_effects[g]["df_used"] = survey_df_final + + return _vcov_rep_imp, _n_valid_rep_imp, survey_df_final diff --git a/diff_diff/imputation_bootstrap.py b/diff_diff/imputation_bootstrap.py index ec7d3da4..a9d1540c 100644 --- a/diff_diff/imputation_bootstrap.py +++ b/diff_diff/imputation_bootstrap.py @@ -20,6 +20,7 @@ from diff_diff.bootstrap_utils import ( generate_survey_multiplier_weights_batch as _generate_survey_multiplier_weights_batch, ) +from diff_diff.imputation_aggregation import _compute_target_weights # noqa: F401 (shared helper) from diff_diff.imputation_results import ImputationBootstrapResults __all__ = [ @@ -27,39 +28,6 @@ ] -def _compute_target_weights( - tau_hat: np.ndarray, - target_mask: np.ndarray, -) -> "tuple[np.ndarray, int]": - """ - Equal weights for finite tau_hat observations within target_mask. - - Used by both aggregation and bootstrap paths to avoid weight logic - duplication. - - Parameters - ---------- - tau_hat : np.ndarray - Per-observation treatment effects (may contain NaN). - target_mask : np.ndarray - Boolean mask selecting the target subset within tau_hat. - - Returns - ------- - weights : np.ndarray - Weight array (same length as tau_hat). 1/n_valid for finite - observations in target_mask, 0 elsewhere. - n_valid : int - Number of finite observations in the target subset. - """ - finite_target = np.isfinite(tau_hat) & target_mask - n_valid = int(finite_target.sum()) - weights = np.zeros(len(tau_hat)) - if n_valid > 0: - weights[np.where(finite_target)[0]] = 1.0 / n_valid - return weights, n_valid - - class ImputationDiDBootstrapMixin: """Mixin providing bootstrap inference methods for ImputationDiD.""" diff --git a/diff_diff/imputation_results.py b/diff_diff/imputation_results.py index 161c36f1..f3273ea7 100644 --- a/diff_diff/imputation_results.py +++ b/diff_diff/imputation_results.py @@ -5,14 +5,51 @@ dataclasses. Extracted from imputation.py for module size management. """ +import dataclasses from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd +from diff_diff.aggregation import AggregationMixin, AggregationResult +from diff_diff.imputation_aggregation import _ImputationAggregationMixin 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 + + +class _ImputationKitAggregator(_ImputationAggregationMixin): + """Throwaway per-call host for the post-fit recompute (M-021/M-118). + + Hosts the moved aggregation methods with exactly the mixin's declared + host-attribute contract, populated from KIT SNAPSHOTS — never from the + live estimator (whose config may have been mutated since the fit) and + never from mutable public results fields. A fresh instance per + ``aggregate()`` call; the moved methods write nothing to ``self``, so + the kit stays immutable either way. + """ + + def __init__( + self, + *, + alpha: float, + anticipation: int, + horizon_max: Optional[int], + pretrends: bool, + aux_partition: str, + leave_one_out: bool, + rank_deficient_action: str, + df_convention: str, + ) -> None: + self.alpha = alpha + self.anticipation = anticipation + self.horizon_max = horizon_max + self.pretrends = pretrends + self.aux_partition = aux_partition + self.leave_one_out = leave_one_out + self.rank_deficient_action = rank_deficient_action + self.df_convention = df_convention + __all__ = [ "ImputationBootstrapResults", @@ -75,7 +112,7 @@ class ImputationBootstrapResults: @dataclass -class ImputationDiDResults(BaseResults): +class ImputationDiDResults(BaseResults, AggregationMixin): """ Results from Borusyak-Jaravel-Spiess (2024) imputation DiD estimation. @@ -157,9 +194,18 @@ class ImputationDiDResults(BaseResults): # The estimator's df_convention configuration echoed onto the results # ("residual" | "cluster" | "normal"; added 3.9). It governs only the # pretrends lead regression's per-lead t/p/CI - the BJS overall / - # post-treatment inference is knob-independent. Appended LAST (the - # generated __init__ positional indexes are public API). + # post-treatment inference is knob-independent. df_convention: Optional[str] = None + # Private panel-backed post-fit aggregation kit (rows M-021/M-118), + # attached by ImputationDiD.fit(). None on results unpickled from a + # pre-3.9 release (aggregate() then fails with the re-fit message). + # Appended LAST (the generated __init__ positional indexes are public + # API). + _aggregation_kit: Optional[Any] = field(default=None, repr=False, compare=False) + + # Post-fit aggregation vocabulary (M-021). 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 @@ -182,6 +228,223 @@ def p_value(self) -> float: def t_stat(self) -> float: return self.overall_t_stat + # --- Post-fit aggregation (M-021/M-118) ------------------------------ + + 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 ImputationDiDResults carries no aggregation kit - it is " + "attached by ImputationDiD.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-target 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) + if level == "simple": + return self._aggregate_simple_result(kit) + if level == "event_study" and bk["uses_replicate"] and bk["pretrends"]: + # The same unsupported combination fit(aggregate='event_study') + # rejects: the pre-period lead regression's per-replicate refits + # are not implemented, and the replicate replay re-estimates + # post-treatment targets only. + raise NotImplementedError( + "aggregate('event_study') is not available on this fit: it " + "used pretrends=True with a replicate-weight survey design, " + "and the pre-period lead regression's per-replicate refits " + "are not yet implemented (fit(aggregate='event_study') " + "rejects the same combination). Re-fit with pretrends=False, " + "or use an analytical (strata/PSU/FPC) survey design." + ) + # Fresh throwaway host per call, populated from KIT snapshots only + # (estimator/config mutation after fit() must not leak in), and a + # call-local projection cache (the id()-keyed cache keys are the + # kit's own mask objects, so reuse within the call is exact). + agg = _ImputationKitAggregator( + alpha=kit.alpha, + anticipation=kit.anticipation, + horizon_max=bk["horizon_max"], + pretrends=bk["pretrends"], + aux_partition=bk["aux_partition"], + leave_one_out=bk["leave_one_out"], + rank_deficient_action=bk["rank_deficient_action"], + df_convention=bk["df_convention"], + ) + proj_cache: Dict[Any, Any] = {} + common: Dict[str, Any] = dict( + df=bk["df"], + outcome=bk["outcome"], + unit=bk["unit"], + time=bk["time"], + first_treat=bk["first_treat"], + covariates=bk["covariates"], + omega_0_mask=bk["omega_0_mask"], + omega_1_mask=bk["omega_1_mask"], + unit_fe=bk["unit_fe"], + time_fe=bk["time_fe"], + grand_mean=bk["grand_mean"], + delta_hat=bk["delta_hat"], + cluster_var=bk["cluster_var"], + treatment_groups=bk["treatment_groups"], + kept_cov_mask=bk["kept_cov_mask"], + survey_weights=bk["survey_weights"], + survey_df=bk["survey_df_seed"], + resolved_survey=(None if bk["uses_replicate"] else bk["resolved_survey"]), + proj_cache=proj_cache, + ) + if level == "group": + effects = agg._aggregate_group(**common) + if bk["uses_replicate"]: + # LEVEL-MATCHED replay: [overall, groups] - reproduces + # fit(aggregate='group') exactly (see the replay docstring). + agg._replicate_override_aggregates( + df=bk["df"], + outcome=bk["outcome"], + unit=bk["unit"], + time=bk["time"], + first_treat=bk["first_treat"], + covariates=bk["covariates"], + omega_0_mask=bk["omega_0_mask"], + omega_1_mask=bk["omega_1_mask"], + resolved_survey=bk["resolved_survey"], + overall_att=bk["overall_att"], + event_study_effects=None, + group_effects=effects, + balance_e=None, + survey_df_seed=bk["survey_df_seed"], + ) + return self._group_effects_to_aggregation(effects, kit) + # level == "event_study" (the mixin validated the vocabulary) + es = agg._aggregate_event_study(**common, balance_e=balance_e) + replay_df: Optional[int] = None + if bk["uses_replicate"]: + _, _, replay_df = agg._replicate_override_aggregates( + df=bk["df"], + outcome=bk["outcome"], + unit=bk["unit"], + time=bk["time"], + first_treat=bk["first_treat"], + covariates=bk["covariates"], + omega_0_mask=bk["omega_0_mask"], + omega_1_mask=bk["omega_1_mask"], + resolved_survey=bk["resolved_survey"], + overall_att=bk["overall_att"], + event_study_effects=es, + group_effects=None, + balance_e=balance_e, + survey_df_seed=bk["survey_df_seed"], + ) + # Carrier + shared builder: ImputationDiD is a _from_relative_dict + # producer, so the recomputed dict rides the same route as the + # fit-time surface (zero-count-sentinel reference marking, + # n_kind="obs", all-NaN per-row df - identical to fit-time output). + # The carrier's metadata is a copy-on-use of the KIT's fit-final + # metadata copy (never the mutable public field); on a replicate + # replay its df_survey is the REPLAYED level-matched value, + # normalized by the same rule fit applies. + meta = bk["survey_metadata"] + if meta is not None: + if bk["uses_replicate"]: + meta = dataclasses.replace( + meta, df_survey=(replay_df if replay_df and replay_df > 0 else None) + ) + else: + meta = dataclasses.replace(meta) + carrier = dataclasses.replace( + self, + event_study_effects=es, + survey_metadata=meta, + 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_obs`` (|Omega_1|) with ``n_kind="obs"``: + ImputationDiD's ``n_treated_units``/``n_control_units`` unit sets + OVERLAP (a treated unit with pre-periods counts in both), so the + CS/EDiD disjoint-units convention cannot apply (the StackedDiD + carve-out class); the treated-observation count is the population + the overall ATT averages over and matches every other Imputation + row's n semantics. + + ``df`` is the kit's ``survey_df_final`` snapshot - the exact value + the STORED overall ``safe_inference`` received (on a replicate fit + that value came from the ``[overall]``-only joint stack, which is + precisely why it must be snapshotted rather than re-derived). + None → all-NaN df column; the replicate-undefined 0 sentinel NaNs + out via post_init. + """ + 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_treated_obs"]], dtype=float), + df=kit.bookkeeping["survey_df_final"], + alpha=kit.alpha, + n_kind="obs", + 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`` key each row's + ``safe_inference`` recorded (capture-at-use: the analytical writer + and the replicate override genuinely use different values on + replicate fits; the all-NaN cohort branch writes no key, read via + ``.get`` → NaN). ``weight=None``: cohort means over their own + observations carry no cross-cohort mass (the CS rationale). + ``n_kind="obs"``: ``n_obs`` counts the cohort's treated + observations, matching the ES surface's n semantics. + """ + 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_obs"] for g in labels], dtype=float), + df=df_arr, + alpha=kit.alpha, + n_kind="obs", + weight=None, + estimator=type(self).__name__.replace("Results", ""), + ) + def __repr__(self) -> str: """Concise string representation.""" sig = _get_significance_stars(self.overall_p_value) @@ -440,8 +703,12 @@ def to_dataframe(self, level: str = "observation") -> pd.DataFrame: elif level == "event_study": if self.event_study_effects is None: raise ValueError( - "Event study effects not computed. " - "Use aggregate='event_study' or aggregate='all'." + "Event study effects not computed. Aggregate post-fit " + "instead - results.aggregate('event_study') returns the " + "EventStudyResults container (on a bootstrapped fit, " + "re-fit with n_bootstrap=0 or use the deprecated " + "fit-time aggregate=; a result unpickled from a pre-3.9 " + "release carries no kit and must be re-fit)." ) rows = [] for h, data in sorted(self.event_study_effects.items()): @@ -462,7 +729,12 @@ def to_dataframe(self, level: str = "observation") -> pd.DataFrame: elif level == "group": if self.group_effects is None: raise ValueError( - "Group effects not computed. " "Use aggregate='group' or aggregate='all'." + "Group effects not computed. Aggregate post-fit instead " + "- results.aggregate('group') returns the " + "AggregationResult container (on a bootstrapped fit, " + "re-fit with n_bootstrap=0 or use the deprecated " + "fit-time aggregate=; a result unpickled from a pre-3.9 " + "release carries no kit and must be re-fit)." ) rows = [] for g, data in sorted(self.group_effects.items()): diff --git a/diff_diff/practitioner.py b/diff_diff/practitioner.py index 010d3085..1cbd4585 100644 --- a/diff_diff/practitioner.py +++ b/diff_diff/practitioner.py @@ -499,6 +499,35 @@ def _handle_sa(results: Any): def _handle_imputation(results: Any): steps = [ _parallel_trends_step(staggered=True), + _step( + baker_step=7, + label="Aggregate treatment-effect heterogeneity post-fit", + why=( + "ImputationDiD aggregates post-fit from its panel-backed kit " + "(M-021) - 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 = imp.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", + ), _step( baker_step=6, label="Specification-based falsification", @@ -527,6 +556,35 @@ def _handle_imputation(results: Any): def _handle_two_stage(results: Any): steps = [ + _step( + baker_step=7, + label="Aggregate treatment-effect heterogeneity post-fit", + why=( + "TwoStageDiD aggregates post-fit from its panel-backed kit " + "(M-022) - 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 = ts.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", + ), _parallel_trends_step(staggered=True), _step( baker_step=6, diff --git a/diff_diff/pretrends.py b/diff_diff/pretrends.py index ea5eca53..d24618c0 100644 --- a/diff_diff/pretrends.py +++ b/diff_diff/pretrends.py @@ -1445,11 +1445,16 @@ 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); 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." + "SunAbrahamResults); EfficientDiDResults and " + "ImputationDiDResults containers are rejected BY DESIGN " + "(their surfaces carry no joint event-study covariance - " + "per-horizon SEs only; see the REGISTRY EfficientDiD and " + "ImputationDiD Notes); TwoStageDiDResults container " + "admission is DEFERRED pending a normalization derivation " + "- analytical fits carry the joint Gardner-GMM covariance, " + "but the pre-period coefficients are stage-1 residual " + "means, not contrasts against a reference period (see the " + "REGISTRY TwoStageDiD Note and DEFERRED.md)." ) _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 4b0adc44..a07130a9 100644 --- a/diff_diff/results_base.py +++ b/diff_diff/results_base.py @@ -613,8 +613,14 @@ def summary(self, alpha: Optional[float] = None) -> str: _ABSENT_SURFACE_HINTS: Dict[str, str] = { # Migrated to the post-fit surface (row M-020): no refit needed. "CallawaySantAnnaResults": "call results.aggregate('event_study')", - "ImputationDiDResults": "refit with aggregate='event_study' (or 'all')", - "TwoStageDiDResults": "refit with aggregate='event_study' (or 'all')", + "ImputationDiDResults": ( + "call results.aggregate('event_study') (on a bootstrapped fit, " + "re-fit with n_bootstrap=0 or the deprecated fit-time aggregate=)" + ), + "TwoStageDiDResults": ( + "call results.aggregate('event_study') (on a bootstrapped fit, " + "re-fit with n_bootstrap=0 or the deprecated fit-time aggregate=)" + ), # Absence only possible on pre-3.9 pickles: 3.9+ fits always # materialize the surface (row M-024). "StackedDiDResults": "re-fit with diff-diff >= 3.9, which always computes the surface", diff --git a/diff_diff/two_stage.py b/diff_diff/two_stage.py index 19aa1f6a..49c7ed8a 100644 --- a/diff_diff/two_stage.py +++ b/diff_diff/two_stage.py @@ -21,6 +21,7 @@ Difference-in-Differences. R Journal, 14(1), 162-173. """ +import dataclasses import warnings from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple @@ -30,6 +31,8 @@ from scipy.sparse.linalg import factorized as sparse_factorized from diff_diff._base import BaseEstimator +from diff_diff._deprecation import NOT_SUPPLIED +from diff_diff.aggregation import AggregationKit from diff_diff.conley import ( ConleyMetric, _compute_conley_meat, @@ -37,13 +40,18 @@ _validate_conley_kwargs, _validate_meat_psd, ) -from diff_diff.linalg import _rank_guarded_inv, solve_ols +from diff_diff.two_stage_aggregation import ( # noqa: F401 (compat re-exports) + _SPARSE_DENSE_THRESHOLD, + _lsmr_certified_normal_solve, + _LSMRUnconvergedError, + _TwoStageAggregationMixin, +) from diff_diff.two_stage_bootstrap import TwoStageDiDBootstrapMixin from diff_diff.two_stage_results import ( TwoStageBootstrapResults, # noqa: F401 TwoStageDiDResults, ) # noqa: F401 (re-export) -from diff_diff.utils import _iterative_fe_solve, demean_by_groups, safe_inference +from diff_diff.utils import safe_inference if TYPE_CHECKING: # Forward reference for the Wave E.1 survey-design path. Imported under @@ -51,87 +59,6 @@ # circular-import surprises with diff_diff.survey. from diff_diff.survey import ResolvedSurveyDesign, SurveyDesign -# Maximum number of elements before falling back to per-column sparse aggregation. -# 10M float64 elements ≈ 80 MB peak allocation. Above this, per-column .getcol() -# trades throughput for bounded memory. Keep in sync with two_stage_bootstrap.py. -_SPARSE_DENSE_THRESHOLD = 10_000_000 - - -class _LSMRUnconvergedError(RuntimeError): - """LSMR could not certify the Stage-1 normal-equation solve; the - variance boundary converts this to NaN inference (fail-closed).""" - - -def _lsmr_certified_normal_solve( - gram_csc, rhs: np.ndarray, context: str = "TwoStageDiD GMM sandwich" -) -> np.ndarray: - """Least-squares solve of the (possibly singular) sparse Stage-1 Gram - system ``gram @ out = rhs`` via per-column LSMR — no dense - materialization of the ``(p_1, p_1)`` normal matrix (`O((U+T+K)^2)` - OOM risk on large panels; the pattern the ImputationDiD LSMR fix - closed, ported here after the consumer-invariance analysis). - - OUTPUT-PRESERVING despite the min-norm ambiguity on singular systems: - least-squares solutions differ only by a ``null(X'X) = null(X_10)`` - component (weighted: ``null(X'WX) = null(W^{1/2}X_10)`` — zero-weight - rows are inert in every weighted consumer because Psi/score/residual - contributions carry the same ``W`` factor), and EVERY ``gamma_hat`` - consumer is an ``X_10``-range functional. One ``theta_exact`` consumer - (the bootstrap exact-residual helper's ``X_1_sparse @ theta_exact``) - evaluates theta on TREATED rows where a ``null(X_10)`` component would - NOT annihilate — parity there holds for a second reason: both dense - ``lstsq`` (SVD) and LSMR return the MIN-NORM least-squares solution, so - the two solvers agree on the whole vector (to iterative tolerance), not - just on range functionals; the fit-level singular-design parity test - locks this at the V/SE level. The remaining consumers are the - ``X_10``-range functionals — ``Psi_stage1 = X_10 @ gamma_hat``, the GMM - score correction ``c_g' gamma_hat`` with ``c_g = X_{10,g}' eps_{10,g}`` - in ``rowspace(X_10)``, and Stage-1 residuals ``y - X_10 theta`` — so - every null component annihilates. Locked by the singular-system parity - test against a dense-lstsq oracle. - - CONVERGENCE IS VALIDATED (fail-closed): ``istop`` in ``{0, 1, 2, 4, 5}`` - certifies a solution / least-squares solution within tolerance (4/5 are - the machine-precision analogues of 1/2 per SciPy); anything else gets - ONE retry with an uncapped condition limit, then raises - :class:`_LSMRUnconvergedError` — converted to NaN inference at the - variance boundary rather than feeding an unverified solution into the - GMM sandwich. - """ - import scipy.sparse.linalg as spla - - _certified = (0, 1, 2, 4, 5) - rhs_2d = np.atleast_2d(np.asarray(rhs, dtype=np.float64)) - if rhs_2d.shape[0] == 1 and np.asarray(rhs).ndim == 1: - rhs_2d = rhs_2d.T - dim = gram_csc.shape[0] - out = np.empty((dim, rhs_2d.shape[1])) - for j in range(rhs_2d.shape[1]): - result = spla.lsmr(gram_csc, rhs_2d[:, j], atol=1e-14, btol=1e-14) - z, istop = result[0], int(result[1]) - if istop not in _certified or not np.all(np.isfinite(z)): - result = spla.lsmr( - gram_csc, - rhs_2d[:, j], - atol=1e-14, - btol=1e-14, - conlim=1e16, - maxiter=max(50 * dim, 10_000), - ) - z, istop = result[0], int(result[1]) - if istop not in _certified or not np.all(np.isfinite(z)): - warnings.warn( - f"{context}: the LSMR fallback solve of the " - f"Stage-1 normal equations did not converge (istop={istop}); " - "the affected variance is reported as NaN rather than from " - "an unverified solution.", - UserWarning, - stacklevel=3, - ) - raise _LSMRUnconvergedError(f"LSMR uncertified (istop={istop})") - out[:, j] = z - return out - # ============================================================================= # Wave D — Gardner GMM-corrected meat for SpilloverDiD @@ -1291,7 +1218,7 @@ def _compute_stratified_serial_bartlett_meat( # ============================================================================= -class TwoStageDiD(TwoStageDiDBootstrapMixin, BaseEstimator): +class TwoStageDiD(TwoStageDiDBootstrapMixin, _TwoStageAggregationMixin, BaseEstimator): """ Gardner (2022) two-stage Difference-in-Differences estimator. @@ -1364,14 +1291,14 @@ class TwoStageDiD(TwoStageDiDBootstrapMixin, BaseEstimator): ... time='period', first_treat='first_treat') >>> results.print_summary() - With event study: + With a post-fit event study (M-022): >>> est = TwoStageDiD() >>> results = est.fit(data, outcome='outcome', unit='unit', - ... time='period', first_treat='first_treat', - ... aggregate='event_study') + ... time='period', first_treat='first_treat') + >>> es = results.aggregate('event_study') >>> from diff_diff import plot_event_study - >>> plot_event_study(results) + >>> plot_event_study(es) Notes ----- @@ -1434,8 +1361,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["SurveyDesign"] = None, ) -> TwoStageDiDResults: """ @@ -1457,11 +1384,18 @@ def fit( covariates : list of str, optional List of covariate column names. aggregate : str, optional - Aggregation mode: None/"simple" (overall ATT only), - "event_study", "group", or "all". + DEPRECATED (3.9, removed in 4.0; row M-022): aggregate as a + post-fit step instead — ``results.aggregate('event_study')`` / + ``.aggregate('group')`` / ``.aggregate('simple')``. Supplying + ANY value (``None`` included) warns; the deprecated path still + works and returns exactly the numbers it always did + (fit-time mode: None/"simple" overall only, "event_study", + "group", or "all"). balance_e : int, optional - When computing event study, restrict to cohorts observed at all - relative times in [-balance_e, max_h]. + DEPRECATED (3.9, removed in 4.0; row M-119): moves onto + ``results.aggregate('event_study', balance_e=...)``. Restricts + the event study to cohorts observed at every relative time in + ``[-balance_e, max_h]`` (the balanced-window rule). survey_design : SurveyDesign, optional Survey design specification for design-based inference. Supports pweight only (aweight/fweight raise ValueError). Supports strata, @@ -1479,6 +1413,35 @@ def fit( ValueError If required columns are missing or data validation fails. """ + # M-022/M-119 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"TwoStageDiD.fit({_args}) is deprecated and will be " + "removed in 4.0. Fit once, then aggregate as a post-fit " + "step: results = TwoStageDiD().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 + # Re-validate vcov_type at fit-time: set_params validates eagerly # (BaseEstimator probe re-init), so this only catches DIRECT # attribute mutation (est.vcov_type = ...). @@ -1772,6 +1735,11 @@ def fit( if _uses_replicate_ts and _survey_df is None: _survey_df = 0 + # Kit df-provenance SEED (M-022): the exact value the stage-2 + # aggregators below receive, captured BEFORE the replicate override + # can rebind _survey_df — post-fit recompute must re-seed from it. + _survey_df_stage2 = _survey_df + # Wave E.3 parity (PR #482 SpilloverDiD precedent): under the survey # path, `score_pad_mask_arg` is the FULL-DOMAIN keep_mask identifying # rows present in the stage-1 / stage-2 OLS sample after the @@ -1881,179 +1849,38 @@ def fit( _n_valid_rep_ts = None _vcov_rep_ts = None if _uses_replicate_ts: - from diff_diff.survey import compute_replicate_refit_variance - - # Derive keys from actual outputs (excludes filtered/Prop5 horizons) - _es_effects_ts = event_study_effects or {} - _grp_effects_ts = group_effects or {} - _sorted_es_periods_ts = sorted( - e for e in _es_effects_ts.keys() if np.isfinite(_es_effects_ts[e]["effect"]) - ) - _sorted_groups_ts = sorted( - g for g in _grp_effects_ts.keys() if np.isfinite(_grp_effects_ts[g]["effect"]) - ) - _n_es_ts = len(_sorted_es_periods_ts) - _n_grp_ts = len(_sorted_groups_ts) - - # Build full-sample estimate from actual outputs - _full_est_ts = [overall_att] - _full_est_ts.extend([_es_effects_ts[e]["effect"] for e in _sorted_es_periods_ts]) - _full_est_ts.extend([_grp_effects_ts[g]["effect"] for g in _sorted_groups_ts]) - - def _refit_ts(w_r): - # Wave E.3 parity (PR #482 SpilloverDiD precedent): the main fit - # path keeps `resolved_survey` at full-domain length but subsets - # `survey_weights` for stage-1 / stage-2 OLS arithmetic via - # `keep_mask` (always-treated drop). The replicate refit - # callback receives a FULL-DOMAIN replicate weight `w_r` - # (sourced from `resolved_survey.replicate_weights` which is - # also full-domain) and must apply the SAME `keep_mask` - # subsetting before threading through stage-1 / stage-2, - # otherwise `solve_ols` rejects the length mismatch - # (full-domain w_r vs post-drop df) and the ValueError is - # swallowed by `compute_replicate_refit_variance` → - # NaN replicate inference. - w_r_fit = np.asarray(w_r)[keep_mask.values] - ufe_r, tfe_r, gm_r, delta_r, kcm_r = self._fit_untreated_model( - df, - outcome, - unit, - time, - covariates, - omega_0_mask, - weights=w_r_fit, - ) - y_tilde_r = self._residualize( - df, - outcome, - unit, - time, - covariates, - ufe_r, - tfe_r, - gm_r, - delta_r, - ) - df_tmp = df.copy() - df_tmp["_y_tilde"] = y_tilde_r - results = [] - - att_r, _ = self._stage2_static( - df=df_tmp, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=ufe_r, - time_fe=tfe_r, - grand_mean=gm_r, - delta_hat=delta_r, - cluster_var=cluster_var, - kept_cov_mask=kcm_r, - survey_weights=w_r_fit, - survey_weight_type="pweight", - warn_nan=False, - ) - results.append(att_r) - - if _sorted_es_periods_ts: - # Replicate refits only need the point effects; the - # per-replicate V is irrelevant to the refit variance. - es_r, _, _ = self._stage2_event_study( - df=df_tmp, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=ufe_r, - time_fe=tfe_r, - grand_mean=gm_r, - delta_hat=delta_r, - cluster_var=cluster_var, - treatment_groups=treatment_groups, - ref_period=ref_period, - balance_e=balance_e, - kept_cov_mask=kcm_r, - survey_weights=w_r_fit, - survey_weight_type="pweight", - survey_df=None, - warn_nan=False, - ) - for e in _sorted_es_periods_ts: - results.append(es_r[e]["effect"] if e in es_r else np.nan) - - if _sorted_groups_ts: - grp_r = self._stage2_group( - df=df_tmp, - unit=unit, - time=time, - first_treat=first_treat, - covariates=covariates, - omega_0_mask=omega_0_mask, - omega_1_mask=omega_1_mask, - unit_fe=ufe_r, - time_fe=tfe_r, - grand_mean=gm_r, - delta_hat=delta_r, - cluster_var=cluster_var, - treatment_groups=treatment_groups, - kept_cov_mask=kcm_r, - survey_weights=w_r_fit, - survey_weight_type="pweight", - survey_df=None, - warn_nan=False, - ) - for g in _sorted_groups_ts: - results.append(grp_r[g]["effect"] if g in grp_r else np.nan) - - return np.array(results) - - _vcov_rep_ts, _n_valid_rep_ts = compute_replicate_refit_variance( - _refit_ts, np.array(_full_est_ts), resolved_survey + ( + _vcov_rep_ts, + _n_valid_rep_ts, + _survey_df, + ) = self._replay_replicate_inference( + df=df, + outcome=outcome, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + cluster_var=cluster_var, + treatment_groups=treatment_groups, + ref_period=ref_period, + balance_e=balance_e, + keep_mask=keep_mask, + resolved_survey=resolved_survey, + overall_att=overall_att, + event_study_effects=event_study_effects, + group_effects=group_effects, + survey_df_seed=_survey_df, ) overall_se = float(np.sqrt(max(_vcov_rep_ts[0, 0], 0.0))) - - # Override df if replicates were dropped - # Replicate-refit path is only reached with a resolved design. - assert resolved_survey is not None - if _n_valid_rep_ts < resolved_survey.n_replicates: - _survey_df = _n_valid_rep_ts - 1 if _n_valid_rep_ts > 1 else 0 if survey_metadata is not None: survey_metadata.df_survey = _survey_df if _survey_df and _survey_df > 0 else None - # Recompute overall inference with replicate SE/df overall_t, overall_p, overall_ci = safe_inference( overall_att, overall_se, alpha=self.alpha, df=_survey_df ) - # Override event-study SEs (only for identified effects) - for i, e in enumerate(_sorted_es_periods_ts): - if event_study_effects is not None and e in event_study_effects: - se_e = float(np.sqrt(max(_vcov_rep_ts[1 + i, 1 + i], 0.0))) - eff_e = event_study_effects[e]["effect"] - t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=_survey_df) - event_study_effects[e]["se"] = se_e - event_study_effects[e]["t_stat"] = t_e - event_study_effects[e]["p_value"] = p_e - event_study_effects[e]["conf_int"] = ci_e - - # Override group SEs (only for identified effects) - for j, g in enumerate(_sorted_groups_ts): - if group_effects is not None and g in group_effects: - se_g = float( - np.sqrt(max(_vcov_rep_ts[1 + _n_es_ts + j, 1 + _n_es_ts + j], 0.0)) - ) - eff_g = group_effects[g]["effect"] - t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=_survey_df) - group_effects[g]["se"] = se_g - group_effects[g]["t_stat"] = t_g - group_effects[g]["p_value"] = p_g - group_effects[g]["conf_int"] = ci_g - # Build treatment effects DataFrame treated_df = df.loc[omega_1_mask, [unit, time, "_y_tilde", "_rel_time"]].copy() treated_df = treated_df.rename(columns={"_y_tilde": "tau_hat", "_rel_time": "rel_time"}) @@ -2240,6 +2067,46 @@ def _refit_ts(w_r): event_study_df=_es_df_final, ) + # Attach the post-fit aggregation kit (M-022/M-119). Unconditional — + # including bootstrap fits, whose gate lives in _aggregate_compute + # (a FAILED bootstrap leaves bootstrap_results=None and the fit + # aggregates normally). + self.results_._aggregation_kit = _build_twostage_aggregation_kit( + df=df, + outcome=outcome, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=unit_fe, + time_fe=time_fe, + grand_mean=grand_mean, + delta_hat=delta_hat, + kept_cov_mask=kept_cov_mask, + cluster_var=cluster_var, + treatment_groups=treatment_groups, + ref_period=ref_period, + survey_weights=survey_weights, + survey_weight_type=survey_weight_type, + resolved_survey=resolved_survey, + uses_replicate=_uses_replicate_ts, + keep_mask=keep_mask, + score_pad_mask=score_pad_mask_arg, + cluster_ids_full=cluster_ids_full_arg, + overall_att=overall_att, + n_treated_obs=n_omega_1, + survey_df_stage2=_survey_df_stage2, + survey_df_final=_survey_df, + survey_metadata=survey_metadata, + pretrends=self.pretrends, + horizon_max=self.horizon_max, + rank_deficient_action=self.rank_deficient_action, + alpha=self.alpha, + anticipation=self.anticipation, + ) + self.is_fitted_ = True return self.results_ @@ -2247,1160 +2114,22 @@ def _refit_ts(w_r): # Stage 1: OLS on untreated observations # ========================================================================= - def _iterative_fe( - self, - y: np.ndarray, - unit_vals: np.ndarray, - time_vals: np.ndarray, - idx: pd.Index, - max_iter: int = 10_000, - tol: float = 1e-10, - weights: Optional[np.ndarray] = None, - ) -> Tuple[Dict[Any, float], Dict[Any, float]]: - """ - Estimate unit and time FE via iterative alternating projection. - - Thin wrapper over the shared bincount solver - (``diff_diff.utils._iterative_fe_solve``): factorize unit/time once, - solve on integer codes, map the level arrays back to dicts. - - Parameters - ---------- - idx : pd.Index - Unused; retained for call-site stability. - weights : np.ndarray, optional - Survey weights (weighted group means ``sum(w*x)/sum(w)``). A - unit/period whose observations ALL carry zero weight has no - identifying contribution and gets ``NaN`` FE (its key is kept so - the rank-condition membership check still sees the group). - - Returns - ------- - unit_fe : dict - Mapping from unit -> unit fixed effect. - time_fe : dict - Mapping from time -> time fixed effect. - """ - unit_codes, unit_uniques = pd.factorize(unit_vals, sort=False) - time_codes, time_uniques = pd.factorize(time_vals, sort=False) - if (unit_codes < 0).any() or (time_codes < 0).any(): - raise ValueError( - "TwoStageDiD: unit or time column contains NaN. Drop or " - "impute missing group keys before fitting." - ) - unit_fe_arr, time_fe_arr = _iterative_fe_solve( - np.asarray(y, dtype=np.float64), - unit_codes.astype(np.intp, copy=False), - time_codes.astype(np.intp, copy=False), - len(unit_uniques), - len(time_uniques), - weights=weights, - max_iter=max_iter, - tol=tol, - method_name="TwoStageDiD iterative FE solver", - ) - unit_fe = dict(zip(unit_uniques, unit_fe_arr)) - time_fe = dict(zip(time_uniques, time_fe_arr)) - return unit_fe, time_fe - - def _fit_untreated_model( - self, - df: pd.DataFrame, - outcome: str, - unit: str, - time: str, - covariates: Optional[List[str]], - omega_0_mask: pd.Series, - weights: Optional[np.ndarray] = None, - ) -> Tuple[ - Dict[Any, float], Dict[Any, float], float, Optional[np.ndarray], Optional[np.ndarray] - ]: - """ - Stage 1: Estimate unit + time FE on untreated observations. - - Parameters - ---------- - weights : np.ndarray, optional - Full-panel survey weights (same length as df). The untreated subset - is extracted internally via omega_0_mask. When None, unweighted. - - Returns - ------- - unit_fe, time_fe, grand_mean, delta_hat, kept_cov_mask - """ - df_0 = df.loc[omega_0_mask] - w_0 = weights[omega_0_mask.values] if weights is not None else None - - if covariates is None or len(covariates) == 0: - y = df_0[outcome].values.copy() - unit_fe, time_fe = self._iterative_fe( - y, df_0[unit].values, df_0[time].values, df_0.index, weights=w_0 - ) - return unit_fe, time_fe, 0.0, None, None - - else: - y = df_0[outcome].values.copy() - X_raw = df_0[covariates].values.copy() - units = df_0[unit].values - times = df_0[time].values - - # Within-transform Y and all X columns through the shared MAP - # engine (factorize-once + bincount + optional Rust kernel), one - # dispatch for every column. within_transform pins [unit, time]; - # [time, unit] here preserves the historical time-then-unit sweep - # order of the per-estimator loops. - narrow = df_0[[outcome, *covariates, time, unit]].copy() - demeaned, _ = demean_by_groups( - narrow, - [outcome, *covariates], - [time, unit], - inplace=True, - weights=w_0, - max_iter=10_000, - tol=1e-10, - ) - y_dm = demeaned[outcome].to_numpy(dtype=np.float64) - X_dm = demeaned[covariates].to_numpy(dtype=np.float64) - - result = solve_ols( - X_dm, - y_dm, - return_vcov=False, - rank_deficient_action=self.rank_deficient_action, - column_names=covariates, - weights=w_0, - ) - delta_hat = result[0] - kept_cov_mask = np.isfinite(delta_hat) - delta_hat_clean = np.where(np.isfinite(delta_hat), delta_hat, 0.0) - - y_adj = y - np.dot(X_raw, delta_hat_clean) - unit_fe, time_fe = self._iterative_fe(y_adj, units, times, df_0.index, weights=w_0) - - return unit_fe, time_fe, 0.0, delta_hat_clean, kept_cov_mask - # ========================================================================= # Residualization # ========================================================================= - def _residualize( - self, - df: pd.DataFrame, - outcome: str, - unit: str, - time: str, - covariates: Optional[List[str]], - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - grand_mean: float, - delta_hat: Optional[np.ndarray], - ) -> np.ndarray: - """ - Compute residualized outcome y_tilde for ALL observations. - - y_tilde_i = y_i - mu_hat_i - eta_hat_t [- X_i @ delta_hat] - """ - alpha_i = df[unit].map(unit_fe).values - beta_t = df[time].map(time_fe).values - - # Handle missing FE (NaN for units/periods not in untreated sample) - alpha_i = np.where(pd.isna(alpha_i), np.nan, alpha_i).astype(float) - beta_t = np.where(pd.isna(beta_t), np.nan, beta_t).astype(float) - - y_hat = grand_mean + alpha_i + beta_t - - if delta_hat is not None and covariates: - y_hat = y_hat + np.dot(df[covariates].values, delta_hat) - - y_tilde = df[outcome].values - y_hat - return y_tilde - # ========================================================================= # Stage 2 specifications # ========================================================================= - @staticmethod - def _mask_nan_ytilde(y_tilde, warn: bool = True): - """Mask non-finite y_tilde values and warn if any found. - - Returns the boolean mask of non-finite values. Modifies y_tilde in-place - (sets NaN values to 0.0). ``warn=False`` suppresses the UserWarning - - used ONLY by the replicate-refit closures, where zero-weight replicate - designs (JK1/BRR) make NaN FE for zeroed-out PSUs expected mechanics - (the main-fit warning still fires once; per-replicate repeats would - emit up to ~3x n_replicates copies of the same message). - """ - nan_mask = ~np.isfinite(y_tilde) - if nan_mask.any(): - n_nan = int(nan_mask.sum()) - if warn: - warnings.warn( - f"{n_nan} observation(s) have non-finite imputed outcomes " - f"(y_tilde) from unidentified fixed effects. These " - f"observations are excluded from ATT estimation.", - UserWarning, - stacklevel=3, - ) - y_tilde[nan_mask] = 0.0 - return nan_mask - - def _stage2_static( - self, - df: pd.DataFrame, - unit: str, - time: str, - first_treat: str, - covariates: Optional[List[str]], - omega_0_mask: pd.Series, - omega_1_mask: pd.Series, - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - grand_mean: float, - delta_hat: Optional[np.ndarray], - cluster_var: str, - kept_cov_mask: Optional[np.ndarray], - survey_weights: Optional[np.ndarray] = None, - survey_weight_type: str = "pweight", - resolved_survey=None, - score_pad_mask: Optional[np.ndarray] = None, - cluster_ids_full: Optional[np.ndarray] = None, - warn_nan: bool = True, - ) -> Tuple[float, float]: - """ - Static (simple ATT) Stage 2: OLS of y_tilde on D_it. - - Returns (att, se). - """ - y_tilde = df["_y_tilde"].values.copy() - nan_mask = self._mask_nan_ytilde(y_tilde, warn=warn_nan) - - D = omega_1_mask.values.astype(float) - # Zero out treatment indicator for NaN y_tilde obs (don't count in ATT) - D[nan_mask] = 0.0 - - # X_2: treatment indicator (no intercept) - X_2 = D.reshape(-1, 1) - - # Avoid degenerate case where all treated obs have NaN y_tilde - if D.sum() == 0: - return np.nan, np.nan - - # Stage 2 OLS for point estimate (discard naive SE) - coef, residuals, _ = solve_ols( - X_2, - y_tilde, - return_vcov=False, - weights=survey_weights, - weight_type=survey_weight_type, - ) - att = float(coef[0]) - - # GMM sandwich variance - # An uncertified LSMR Stage-1 fallback solve fails closed: - # NaN vcov -> NaN SE/t/p/CI (the helper already warned). - try: - V = self._compute_gmm_variance( - df=df, - unit=unit, - time=time, - covariates=covariates, - omega_0_mask=omega_0_mask, - unit_fe=unit_fe, - time_fe=time_fe, - delta_hat=delta_hat, - kept_cov_mask=kept_cov_mask, - X_2=X_2, - cluster_ids=df[cluster_var].values, - survey_weights=survey_weights, - resolved_survey=resolved_survey, - score_pad_mask=score_pad_mask, - cluster_ids_full=cluster_ids_full, - ) - except _LSMRUnconvergedError: - V = np.full((X_2.shape[1], X_2.shape[1]), np.nan) - - se = float(np.sqrt(max(V[0, 0], 0.0))) - return att, se - - def _stage2_event_study( - self, - df: pd.DataFrame, - unit: str, - time: str, - first_treat: str, - covariates: Optional[List[str]], - omega_0_mask: pd.Series, - omega_1_mask: pd.Series, - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - grand_mean: float, - delta_hat: Optional[np.ndarray], - cluster_var: str, - treatment_groups: List[Any], - ref_period: int, - balance_e: Optional[int], - kept_cov_mask: Optional[np.ndarray], - survey_weights: Optional[np.ndarray] = None, - survey_weight_type: str = "pweight", - survey_df: Optional[int] = None, - resolved_survey=None, - score_pad_mask: Optional[np.ndarray] = None, - cluster_ids_full: Optional[np.ndarray] = None, - warn_nan: bool = True, - ) -> Tuple[Dict[int, Dict[str, Any]], Optional[np.ndarray], Optional[List[int]]]: - """Event study Stage 2: OLS of y_tilde on relative-time dummies. - - Returns ``(effects, vcov, vcov_index)``: the per-horizon effects - dict, the full GMM variance-covariance matrix over the ESTIMATED - horizon coefficients, and the horizon labels ordering its - rows/columns. The reference period and Proposition-5 horizons are - never regression columns, so they appear in ``effects`` but not in - ``vcov_index``; all-filtered horizons (n_obs == 0) ARE columns, - with NaN-filled rows/columns from the rank guard. ``(dict, None, - None)`` on the degenerate early returns that fit no Stage-2 - regression. - """ - y_tilde = df["_y_tilde"].values.copy() - nan_mask = self._mask_nan_ytilde(y_tilde, warn=warn_nan) - rel_times = df["_rel_time"].values - n = len(df) - - # Get all horizons — include pre-periods when pretrends=True - if self.pretrends: - evt_rel = rel_times[~df["_never_treated"].values] - else: - evt_rel = rel_times[omega_1_mask.values] - all_horizons = sorted(set(int(h) for h in evt_rel if np.isfinite(h))) - - # Apply horizon_max filter - if self.horizon_max is not None: - all_horizons = [h for h in all_horizons if abs(h) <= self.horizon_max] - - # Apply balance_e filter - if balance_e is not None: - cohort_rel_times = self._build_cohort_rel_times(df, first_treat) - balanced_cohorts = set() - if all_horizons: - max_h = max(all_horizons) - required_range = set(range(-balance_e, max_h + 1)) - for g, horizons in cohort_rel_times.items(): - if required_range.issubset(horizons): - balanced_cohorts.add(g) - if not balanced_cohorts: - warnings.warn( - f"No cohorts satisfy balance_e={balance_e} requirement. " - "Event study results will contain only the reference period. " - "Consider reducing balance_e.", - UserWarning, - stacklevel=2, - ) - return ( - { - ref_period: { - "effect": 0.0, - "se": 0.0, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (0.0, 0.0), - "n_obs": 0, - } - }, - None, - None, - ) - balance_mask = df[first_treat].isin(balanced_cohorts).values - else: - balance_mask = np.ones(n, dtype=bool) - - # Check Proposition 5: no never-treated units - has_never_treated = df["_never_treated"].any() - h_bar = np.inf - if not has_never_treated and len(treatment_groups) > 1: - h_bar = max(treatment_groups) - min(treatment_groups) - - # Identify Prop 5 horizons and compute their actual treated obs counts. - # Treated obs have NaN y_tilde at these horizons (counterfactual - # unidentified), but actual_n counts them to distinguish from truly - # empty horizons. rel_times is NaN for untreated/never-treated obs - # (line ~653), so (rel_times == h) is False for them. - prop5_horizons = [] - prop5_effects: Dict[int, Dict[str, Any]] = {} - if h_bar < np.inf: - for h in all_horizons: - if h == ref_period: - continue - if h >= h_bar: - actual_n = int(np.sum((rel_times == h) & omega_1_mask.values & balance_mask)) - if actual_n > 0: - prop5_horizons.append(h) - prop5_effects[h] = { - "effect": np.nan, - "se": np.nan, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (np.nan, np.nan), - "n_obs": actual_n, - } - - # Remove reference period AND Prop 5 horizons from estimation - prop5_set = set(prop5_horizons) - est_horizons = [h for h in all_horizons if h != ref_period and h not in prop5_set] - - if len(est_horizons) == 0: - # No horizons to estimate — return just reference period - return ( - { - ref_period: { - "effect": 0.0, - "se": 0.0, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (0.0, 0.0), - "n_obs": 0, - } - }, - None, - None, - ) - - # Build Stage 2 design: one column per horizon (no intercept) - # Never-treated obs get all-zero rows (undefined relative time -> NaN) - # With no intercept, they contribute zero to X'_2 X_2 and X'_2 y_tilde - horizon_to_col = {h: j for j, h in enumerate(est_horizons)} - k = len(est_horizons) - X_2 = np.zeros((n, k)) - - for i in range(n): - if not balance_mask[i]: - continue - if nan_mask[i]: - continue # NaN y_tilde -> don't include in event study - h = rel_times[i] - if np.isfinite(h): - h_int = int(h) - if h_int in horizon_to_col: - X_2[i, horizon_to_col[h_int]] = 1.0 - - # Stage 2 OLS - coef, residuals, _ = solve_ols( - X_2, - y_tilde, - return_vcov=False, - weights=survey_weights, - weight_type=survey_weight_type, - ) - - # GMM variance for full coefficient vector - # An uncertified LSMR Stage-1 fallback solve fails closed: - # NaN vcov -> NaN SE/t/p/CI (the helper already warned). - try: - V = self._compute_gmm_variance( - df=df, - unit=unit, - time=time, - covariates=covariates, - omega_0_mask=omega_0_mask, - unit_fe=unit_fe, - time_fe=time_fe, - delta_hat=delta_hat, - kept_cov_mask=kept_cov_mask, - X_2=X_2, - cluster_ids=df[cluster_var].values, - survey_weights=survey_weights, - resolved_survey=resolved_survey, - score_pad_mask=score_pad_mask, - cluster_ids_full=cluster_ids_full, - ) - except _LSMRUnconvergedError: - V = np.full((X_2.shape[1], X_2.shape[1]), np.nan) - - # Build results dict - event_study_effects: Dict[int, Dict[str, Any]] = {} - - # Reference period marker - event_study_effects[ref_period] = { - "effect": 0.0, - "se": 0.0, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (0.0, 0.0), - "n_obs": 0, - } - - for h in est_horizons: - j = horizon_to_col[h] - n_obs = int(np.sum(X_2[:, j])) - - if n_obs == 0: - event_study_effects[h] = { - "effect": np.nan, - "se": np.nan, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (np.nan, np.nan), - "n_obs": 0, - } - continue - - effect = float(coef[j]) - se = float(np.sqrt(max(V[j, j], 0.0))) - - t_stat, p_val, ci = safe_inference(effect, se, alpha=self.alpha, df=survey_df) - - event_study_effects[h] = { - "effect": effect, - "se": se, - "t_stat": t_stat, - "p_value": p_val, - "conf_int": ci, - "n_obs": n_obs, - } - - # Add Proposition 5 entries (unidentified horizons with n_obs > 0) - event_study_effects.update(prop5_effects) - - if prop5_horizons: - warnings.warn( - f"Horizons {prop5_horizons} are not identified without " - f"never-treated units (Proposition 5). Set to NaN.", - UserWarning, - stacklevel=2, - ) - - return event_study_effects, V, [int(h) for h in est_horizons] - - def _stage2_group( - self, - df: pd.DataFrame, - unit: str, - time: str, - first_treat: str, - covariates: Optional[List[str]], - omega_0_mask: pd.Series, - omega_1_mask: pd.Series, - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - grand_mean: float, - delta_hat: Optional[np.ndarray], - cluster_var: str, - treatment_groups: List[Any], - kept_cov_mask: Optional[np.ndarray], - survey_weights: Optional[np.ndarray] = None, - survey_weight_type: str = "pweight", - survey_df: Optional[int] = None, - resolved_survey=None, - score_pad_mask: Optional[np.ndarray] = None, - cluster_ids_full: Optional[np.ndarray] = None, - warn_nan: bool = True, - ) -> Dict[Any, Dict[str, Any]]: - """Group (cohort) Stage 2: OLS of y_tilde on cohort dummies.""" - y_tilde = df["_y_tilde"].values.copy() - nan_mask = self._mask_nan_ytilde(y_tilde, warn=warn_nan) - n = len(df) - - # Build Stage 2 design: one column per cohort (no intercept) - group_to_col = {g: j for j, g in enumerate(treatment_groups)} - k = len(treatment_groups) - X_2 = np.zeros((n, k)) - - ft_vals = df[first_treat].values - treated_mask = omega_1_mask.values - for i in range(n): - if treated_mask[i] and not nan_mask[i]: - g = ft_vals[i] - if g in group_to_col: - X_2[i, group_to_col[g]] = 1.0 - - # Stage 2 OLS - coef, residuals, _ = solve_ols( - X_2, - y_tilde, - return_vcov=False, - weights=survey_weights, - weight_type=survey_weight_type, - ) - - # GMM variance - # An uncertified LSMR Stage-1 fallback solve fails closed: - # NaN vcov -> NaN SE/t/p/CI (the helper already warned). - try: - V = self._compute_gmm_variance( - df=df, - unit=unit, - time=time, - covariates=covariates, - omega_0_mask=omega_0_mask, - unit_fe=unit_fe, - time_fe=time_fe, - delta_hat=delta_hat, - kept_cov_mask=kept_cov_mask, - X_2=X_2, - cluster_ids=df[cluster_var].values, - survey_weights=survey_weights, - resolved_survey=resolved_survey, - score_pad_mask=score_pad_mask, - cluster_ids_full=cluster_ids_full, - ) - except _LSMRUnconvergedError: - V = np.full((X_2.shape[1], X_2.shape[1]), np.nan) - - group_effects: Dict[Any, Dict[str, Any]] = {} - for g in treatment_groups: - j = group_to_col[g] - n_obs = int(np.sum(X_2[:, j])) - - if n_obs == 0: - group_effects[g] = { - "effect": np.nan, - "se": np.nan, - "t_stat": np.nan, - "p_value": np.nan, - "conf_int": (np.nan, np.nan), - "n_obs": 0, - } - continue - - effect = float(coef[j]) - se = float(np.sqrt(max(V[j, j], 0.0))) - - t_stat, p_val, ci = safe_inference(effect, se, alpha=self.alpha, df=survey_df) - - group_effects[g] = { - "effect": effect, - "se": se, - "t_stat": t_stat, - "p_value": p_val, - "conf_int": ci, - "n_obs": n_obs, - } - - return group_effects - # ========================================================================= # GMM score computation # ========================================================================= - @staticmethod - def _compute_gmm_scores( - c_by_cluster: np.ndarray, - gamma_hat: np.ndarray, - s2_by_cluster: np.ndarray, - ) -> np.ndarray: - """ - Compute per-cluster GMM scores S_g = gamma_hat' c_g - X'_{2g} eps_{2g}. - - Handles NaN/overflow from rank-deficient FE by wrapping in errstate - and replacing non-finite values with 0. - - Parameters - ---------- - c_by_cluster : np.ndarray, shape (G, p) - Per-cluster Stage 1 scores. - gamma_hat : np.ndarray, shape (p, k) - Cross-moment correction matrix. - s2_by_cluster : np.ndarray, shape (G, k) - Per-cluster Stage 2 scores. - - Returns - ------- - np.ndarray, shape (G, k) - Per-cluster influence scores. - """ - with np.errstate(invalid="ignore", divide="ignore", over="ignore"): - correction = np.dot(c_by_cluster, gamma_hat) - np.nan_to_num(correction, copy=False, nan=0.0, posinf=0.0, neginf=0.0) - return correction - s2_by_cluster - # ========================================================================= # GMM Sandwich Variance (Butts & Gardner 2022) # ========================================================================= - def _compute_gmm_variance( - self, - df: pd.DataFrame, - unit: str, - time: str, - covariates: Optional[List[str]], - omega_0_mask: pd.Series, - unit_fe: Dict[Any, float], - time_fe: Dict[Any, float], - delta_hat: Optional[np.ndarray], - kept_cov_mask: Optional[np.ndarray], - X_2: np.ndarray, - cluster_ids: np.ndarray, - survey_weights: Optional[np.ndarray] = None, - resolved_survey=None, - score_pad_mask: Optional[np.ndarray] = None, - cluster_ids_full: Optional[np.ndarray] = None, - ) -> np.ndarray: - """ - Compute GMM sandwich variance (Butts & Gardner 2022). - - Matches the R `did2s` source code implementation: uses the GLOBAL - Hessian inverse (not per-cluster) and NO finite-sample adjustments. - - The per-observation influence function is: - IF_i = (X'_2 X_2)^{-1} [gamma_hat' x_{10i} eps_{10i} - x_{2i} eps_{2i}] - - where gamma_hat = (X'_{10} X_{10})^{-1} (X'_1 X_2) uses the GLOBAL - cross-moment. - - The cluster-robust variance is: - V = (X'_2 X_2)^{-1} (sum_g S_g S'_g) (X'_2 X_2)^{-1} - S_g = gamma_hat' c_g - X'_{2g} eps_{2g} - c_g = X'_{10g} eps_{10g} - - With survey weights W (diagonal): - Bread: (X'_2 W X_2)^{-1} - gamma_hat: (X'_{10} W X_{10})^{-1} (X'_1 W X_2) - c_g = sum_{i in g} w_i * x_{10i} * eps_{10i} - s2_g = sum_{i in g} w_i * x_{2i} * eps_{2i} - - Parameters - ---------- - X_2 : np.ndarray, shape (n, k) - Stage 2 design matrix (treatment indicators). The Stage-2 residual - ``eps_2`` is re-solved internally from the *exact* Stage-1 residuals - (see the exact-residual note below), so it is not a parameter. - cluster_ids : np.ndarray, shape (n,) - Cluster identifiers, fit-sample length. Used for the per-cluster - stage-1 / stage-2 score aggregation (OLS path). - survey_weights : np.ndarray, optional - Survey weights of shape (n,). When None, unweighted (identical - to current code). - resolved_survey : ResolvedSurveyDesign, optional - Resolved survey design. Under Wave E.3 parity (PR #482 SpilloverDiD - precedent) the design retains full-domain `n_psu` / `n_strata` / - `df_survey` / `strata` / `fpc` / `psu` arrays even when the - always-treated drop removes rows from the OLS sample. The - zero-padded per-cluster scores expand onto the full-domain PSU - list before stratified-meat dispatch. R `survey::svyrecvar(subset())` - convention (Lumley 2010 §2.5); mirrors `imputation.py:2175-2183` - (PreTrendsImputation) and `prep.py:1401-1432` (DCDH cell variance). - score_pad_mask : np.ndarray of shape (n_full,), bool, optional - Wave E.3 parity zero-pad mask. When supplied, indicates which - FULL-DOMAIN rows are present in the fit sample (True = kept - for OLS). Requires `n == int(np.sum(score_pad_mask))`. Co-supplied - with `cluster_ids_full`. Per-cluster stage-1 / stage-2 score - aggregates computed at fit-length are expanded onto the - full-domain unique-PSU list; PSUs absent from the fit sample - (e.g. PSUs containing only always-treated rows) get zero score - rows but still count toward `G_full` for `n_psu` / `df_survey`. - None (default) → no padding, exact pre-PR behavior. - cluster_ids_full : np.ndarray of shape (n_full,), optional - Full-domain PSU labels. Co-supplied with `score_pad_mask`. Must - share the same length. Provides the full-domain unique-PSU list - used both for score zero-pad expansion and for downstream - strata/FPC `obs_idx` lookups against the full-domain - `resolved_survey.strata` / `.fpc` arrays. None (default) → no - padding, exact pre-PR behavior. - - Returns - ------- - np.ndarray, shape (k, k) - Variance-covariance matrix. - """ - n = len(df) - k = X_2.shape[1] - - # Exclude rank-deficient covariates - cov_list = covariates - if covariates and kept_cov_mask is not None and not np.all(kept_cov_mask): - cov_list = [c for c, k_ in zip(covariates, kept_cov_mask) if k_] - - # Build sparse FE design matrices X_1 (all obs) and X_10 (untreated only) - X_1_sparse, X_10_sparse, unit_to_idx, time_to_idx = self._build_fe_design( - df, unit, time, cov_list, omega_0_mask - ) - - p = X_1_sparse.shape[1] - - # eps_10 = Y - X_10 @ gamma_hat - # Untreated: stage 1 residual (Y - fitted). Treated: Y (X_10 rows = 0). - # Reconstruct Y from y_tilde: Y = y_tilde + fitted_stage1. Because - # y_tilde = Y - fitted_1, the iterative FE in fitted_1 cancel exactly, so - # y_vals == Y (independent of the iterative solver's tolerance). - alpha_i = df[unit].map(unit_fe).values - beta_t = df[time].map(time_fe).values - # Identification mask: obs whose unit AND time FE are both identified by the - # untreated Stage-1 fit. Rank-deficient / Proposition-5 obs (NaN FE) keep the - # iterative-residual behavior; only identified obs get the exact residuals. - identified = np.isfinite(np.asarray(alpha_i, dtype=float)) & np.isfinite( - np.asarray(beta_t, dtype=float) - ) - alpha_i = np.where(pd.isna(alpha_i), 0.0, alpha_i).astype(float) - beta_t = np.where(pd.isna(beta_t), 0.0, beta_t).astype(float) - fitted_1 = alpha_i + beta_t - if delta_hat is not None and cov_list: - if kept_cov_mask is not None and not np.all(kept_cov_mask): - fitted_1 = fitted_1 + np.dot(df[cov_list].values, delta_hat[kept_cov_mask]) - else: - fitted_1 = fitted_1 + np.dot(df[cov_list].values, delta_hat) - - y_tilde = df["_y_tilde"].values - y_vals = y_tilde + fitted_1 # reconstruct Y - y_vals_clean = np.nan_to_num(y_vals, nan=0.0) - - omega_0 = omega_0_mask.values - - # 1. gamma_hat = (X'_{10} W X_{10})^{-1} (X'_1 W X_2) [p x k] - # With survey weights, both cross-products need W. We reuse the SAME - # factorization of (X'_{10} W X_{10}) to also solve the exact Stage-1 FE - # coefficients theta_exact (see exact-residual note below). - if survey_weights is not None: - XtWX_10 = X_10_sparse.T @ X_10_sparse.multiply(survey_weights[:, None]) - Xt1_WX2 = X_1_sparse.T @ (X_2 * survey_weights[:, None]) - rhs_fe = X_10_sparse.T @ (survey_weights * y_vals_clean) - else: - XtWX_10 = X_10_sparse.T @ X_10_sparse # (p x p) sparse - Xt1_WX2 = X_1_sparse.T @ X_2 # (p x k) dense - rhs_fe = X_10_sparse.T @ y_vals_clean # (p,) X'_{10} W Y - - try: - solve_XtX = sparse_factorized(XtWX_10.tocsc()) - if Xt1_WX2.ndim == 1: - gamma_hat = solve_XtX(Xt1_WX2).reshape(-1, 1) - else: - gamma_hat = np.column_stack( - [solve_XtX(Xt1_WX2[:, j]) for j in range(Xt1_WX2.shape[1])] - ) - theta_exact = np.asarray(solve_XtX(np.asarray(rhs_fe).ravel())).ravel() - except RuntimeError as exc: - # Singular matrix — fall back to certified sparse LSMR. Silent-failure - # audit axis C: emit a UserWarning on fallback instead of swallowing. - warnings.warn( - "TwoStageDiD GMM sandwich: sparse factorization of " - f"(X'_{{10}} W X_{{10}}) failed ({type(exc).__name__}); falling " - "back to sparse LSMR. This may indicate a rank-deficient or " - "near-singular Stage 1 design matrix and SE estimates may be " - "less reliable.", - UserWarning, - stacklevel=2, - ) - XtWX_10_csc = XtWX_10.tocsc() - gamma_hat = _lsmr_certified_normal_solve(XtWX_10_csc, Xt1_WX2) - theta_exact = _lsmr_certified_normal_solve( - XtWX_10_csc, np.asarray(rhs_fe).ravel() - ).ravel() - - # Exact Stage-1 / Stage-2 residuals. The point-estimate path uses the - # iterative alternating-projection FE solver (`_iterative_fe`), which - # converges only to ~1e-7 on unbalanced untreated panels; that error is - # negligible for the ATT but perturbs the variance by ~1% relative to the - # analytical GMM sandwich. The variance therefore re-solves the Stage-1 FE - # EXACTLY using the sparse normal equations already factorized for gamma_hat - # (theta_exact), matching R `did2s` to ~1e-7 and mirroring ImputationDiD's - # exact-sparse variance path. The shared `_exact_gmm_residuals` helper is - # used by BOTH this analytical path and the multiplier bootstrap - # (`_compute_cluster_S_scores`) so the influence function is single-sourced. - eps_10, eps_2 = self._exact_gmm_residuals( - X_1_sparse, - theta_exact, - y_vals_clean, - identified, - omega_0, - y_tilde, - X_2, - survey_weights, - ) - - # 2. Per-cluster Stage 1 scores: c_g = sum_{i in g} w_i * x_{10i} * eps_{10i} - # Only untreated obs have non-zero X_10 rows - # With survey weights: multiply eps_10 by survey_weights before sparse multiply - if survey_weights is not None: - weighted_eps_10 = survey_weights * eps_10 - else: - weighted_eps_10 = eps_10 - weighted_X10 = X_10_sparse.multiply(weighted_eps_10[:, None]) # sparse element-wise - - unique_clusters, cluster_indices = np.unique(cluster_ids, return_inverse=True) - G = len(unique_clusters) - - n_elements = weighted_X10.shape[0] * weighted_X10.shape[1] - c_by_cluster = np.zeros((G, p)) - if n_elements > _SPARSE_DENSE_THRESHOLD: - # Per-column path: limits peak memory for large FE matrices - weighted_X10_csc = weighted_X10.tocsc() - for j_col in range(p): - col_data = weighted_X10_csc.getcol(j_col).toarray().ravel() - np.add.at(c_by_cluster[:, j_col], cluster_indices, col_data) - else: - # Dense path: faster for moderate-size matrices - weighted_X10_dense = weighted_X10.toarray() - for j_col in range(p): - np.add.at(c_by_cluster[:, j_col], cluster_indices, weighted_X10_dense[:, j_col]) - - # 3. Per-cluster Stage 2 scores: s2_g = sum_{i in g} w_i * x_{2i} * eps_{2i} - if survey_weights is not None: - weighted_eps_2 = survey_weights * eps_2 - else: - weighted_eps_2 = eps_2 - weighted_X2 = X_2 * weighted_eps_2[:, None] # (n x k) dense - s2_by_cluster = np.zeros((G, k)) - for j_col in range(k): - np.add.at(s2_by_cluster[:, j_col], cluster_indices, weighted_X2[:, j_col]) - - # Wave E.3 parity (PR #482 SpilloverDiD precedent): when the caller - # supplies `score_pad_mask` + `cluster_ids_full`, expand per-cluster - # stage-1 / stage-2 score aggregates onto the FULL-DOMAIN unique-PSU - # list. PSUs absent from the fit sample (those containing only - # always-treated rows) get zero score rows but still count toward - # `G_full` for `n_psu` / `df_survey` accounting. Mirrors R - # `survey::svyrecvar(subset())` (Lumley 2010 §2.5) and the in-library - # convention at `imputation.py:2175-2183` (PreTrendsImputation) and - # `prep.py:1401-1432` (DCDH cell variance). Downstream strata / FPC - # lookups use `cluster_ids_for_lookup` so the obs_idx applies to the - # full-domain `resolved_survey.strata` / `.fpc` arrays. - if score_pad_mask is not None: - if cluster_ids_full is None: - raise ValueError( - "_compute_gmm_variance: score_pad_mask requires " - "cluster_ids_full to be co-supplied (Wave E.3 parity " - "contract — score zero-pad expansion needs the " - "full-domain PSU labels to align with resolved_survey)." - ) - if resolved_survey is None: - raise ValueError( - "_compute_gmm_variance: score_pad_mask requires " - "resolved_survey to be co-supplied (Wave E.3 parity " - "contract — zero-pad only meaningful under a survey " - "design that retains full-domain dimensions)." - ) - n_full = int(len(score_pad_mask)) - if int(len(cluster_ids_full)) != n_full: - raise ValueError( - "_compute_gmm_variance: score_pad_mask and " - "cluster_ids_full must share the FULL-DOMAIN length; " - f"got len(score_pad_mask)={n_full}, " - f"len(cluster_ids_full)={int(len(cluster_ids_full))}." - ) - if int(np.sum(score_pad_mask)) != n: - raise ValueError( - "_compute_gmm_variance: int(np.sum(score_pad_mask)) " - f"({int(np.sum(score_pad_mask))}) must equal the " - f"fit-sample length n ({n}) so the score expansion " - "is well-defined." - ) - unique_clusters_full = np.unique(cluster_ids_full) - G_full = int(len(unique_clusters_full)) - # Map fit-sample unique_clusters into positions in - # unique_clusters_full via searchsorted (both arrays sorted by - # np.unique). Verify the mapping is exact — otherwise the fit - # sample contains PSU labels absent from the full domain (a - # contract violation that should never occur under the upstream - # `_inject_cluster_as_psu` invariant). - fit_to_full_idx = np.searchsorted(unique_clusters_full, unique_clusters) - if not np.array_equal( - unique_clusters_full[fit_to_full_idx], np.asarray(unique_clusters) - ): - raise ValueError( - "_compute_gmm_variance: fit-sample unique cluster " - "labels are not a subset of full-domain cluster labels " - "(Wave E.3 parity invariant violated). This should be " - "impossible under `_inject_cluster_as_psu` — please " - "file an issue with a minimal reproducer." - ) - c_by_cluster_full = np.zeros((G_full, p)) - s2_by_cluster_full = np.zeros((G_full, k)) - c_by_cluster_full[fit_to_full_idx] = c_by_cluster - s2_by_cluster_full[fit_to_full_idx] = s2_by_cluster - c_by_cluster = c_by_cluster_full - s2_by_cluster = s2_by_cluster_full - unique_clusters = unique_clusters_full - G = G_full - cluster_ids_for_lookup = np.asarray(cluster_ids_full) - else: - cluster_ids_for_lookup = cluster_ids - - # 4. S_g = gamma_hat' c_g - X'_{2g} eps_{2g} - S = self._compute_gmm_scores(c_by_cluster, gamma_hat, s2_by_cluster) - - # 5. Meat: sum_g S_g S'_g = S' S - _use_stratified_meat = resolved_survey is not None and ( - resolved_survey.strata is not None or resolved_survey.fpc is not None - ) - if _use_stratified_meat: - from diff_diff.survey import _compute_stratified_meat_from_psu_scores - - # Build PSU→stratum and PSU→FPC mappings from observation-level arrays. - # cluster_ids_for_lookup is full-domain length under Wave E.3 parity - # (score_pad_mask path) and fit-sample length otherwise; either way it - # aligns with `resolved_survey.strata` / `resolved_survey.fpc` so the - # obs_idx lookup resolves to the correct stratum / FPC value. - G_meat = len(unique_clusters) - - # Strata: synthesize single stratum when strata is None (unstratified FPC) - if resolved_survey.strata is not None: - psu_strata = np.empty(G_meat, dtype=resolved_survey.strata.dtype) - for idx, c in enumerate(unique_clusters): - obs_idx = np.where(cluster_ids_for_lookup == c)[0][0] - psu_strata[idx] = resolved_survey.strata[obs_idx] - else: - psu_strata = np.zeros(G_meat, dtype=int) - - # FPC: map observation-level FPC to PSU level - psu_fpc = None - if resolved_survey.fpc is not None: - psu_fpc = np.empty(G_meat, dtype=np.float64) - for idx, c in enumerate(unique_clusters): - obs_idx = np.where(cluster_ids_for_lookup == c)[0][0] - psu_fpc[idx] = resolved_survey.fpc[obs_idx] - - # Unstratified single-PSU: variance is unidentified (matches - # _compute_stratified_psu_meat at survey.py:1225 which returns - # zero meat with no variance_computed flag for n_psu < 2). - # Under Wave E.3 parity, G_meat = G_full (post zero-pad), so the - # gate fires on the full-domain PSU count, not the fit-sample. - if resolved_survey.strata is None and G_meat < 2: - return np.full((k, k), np.nan) - - # Reorder S rows to match unique_clusters ordering - # S is built using np.add.at with cluster_indices from pd.factorize, - # which uses the same order as unique_clusters from the data. - meat, _var_computed, _legit_zero = _compute_stratified_meat_from_psu_scores( - psu_scores=S, - psu_strata=psu_strata, - fpc_per_psu=psu_fpc, - lonely_psu=resolved_survey.lonely_psu, - ) - # If no variance was computed and no legitimate zeros, variance - # is unidentified — return NaN VCV so caller gets NaN SE. - if not _var_computed and _legit_zero == 0: - return np.full((k, k), np.nan) - else: - with np.errstate(invalid="ignore", over="ignore"): - meat = S.T @ S # (k x k) - - # 6. Bread: (X'_2 W X_2)^{-1} - with np.errstate(invalid="ignore", over="ignore", divide="ignore"): - if survey_weights is not None: - XtWX_2 = X_2.T @ (X_2 * survey_weights[:, None]) - else: - XtWX_2 = X_2.T @ X_2 - # np.linalg.solve only raises on an *exactly* singular Gram; a *near*- - # singular X_2'WX_2 would otherwise flow a garbage inverse (~1e13) - # straight into the SE. `_rank_guarded_inv` truncates redundant - # directions on the equilibrated Gram -> finite SE on the identified - # subspace (NaN only at rank 0), matching the covariate IF rank-guard. - # Sibling of finding #17 (axis A): the prior fallback fired only on an - # exactly-singular matrix. X_2 is the Stage-2 indicator design (not user - # covariates), so the diagnostic guidance points at that layer. - bread, n_dropped, _, dropped = _rank_guarded_inv(XtWX_2, return_dropped=True) - if n_dropped: - warnings.warn( - "Rank-deficient second-stage design matrix X_2'WX_2 in " - "TwoStageDiD TSL variance; rank-reducing to a finite SE on the " - f"identified subspace ({n_dropped} redundant direction(s) " - "dropped, NaN if rank 0). The Stage-2 design is built from " - "treatment, event-time, or group indicators, so this typically " - "indicates a zero-weight or all-zero indicator column " - "(e.g. an aggregation path with no qualifying observations).", - UserWarning, - stacklevel=2, - ) - - # 7. V = bread @ meat @ bread - V = bread @ meat @ bread - # A dropped (unidentified) Stage-2 coefficient is zero-filled in `bread`, - # which would report se=0 for that named coefficient; NaN its row/col in - # the FINAL vcov so per-coefficient SE extraction yields NaN (not 0). - if dropped.any(): - V[dropped, :] = np.nan - V[:, dropped] = np.nan - return V - - def _build_fe_design( - self, - df: pd.DataFrame, - unit: str, - time: str, - covariates: Optional[List[str]], - omega_0_mask: pd.Series, - ) -> Tuple[sparse.csr_matrix, sparse.csr_matrix, Dict[Any, int], Dict[Any, int]]: - """ - Build sparse FE design matrices X_1 (all obs) and X_10 (untreated rows only). - - Column layout: [intercept, unit_1, ..., unit_{U-1}, time_1, ..., time_{T-1}, - cov_1, ..., cov_C] (drop first unit and first time for identification, with an - intercept). The intercept makes the column space span the constant (the grand - mean); the prior intercept-free layout silently omitted the grand mean from the - FE span, which biased the GMM-sandwich residuals when re-solved exactly. With - the intercept this is the standard full-rank two-way FE (matches fixest / R - ``did2s``). - - X_10 is identical to X_1 except that rows for treated observations are zeroed out. - - Returns - ------- - X_1_sparse : sparse.csr_matrix, shape (n, p) - X_10_sparse : sparse.csr_matrix, shape (n, p) - unit_to_idx : dict - time_to_idx : dict - """ - n = len(df) - unit_vals = df[unit].values - time_vals = df[time].values - omega_0 = omega_0_mask.values - - all_units = np.unique(unit_vals) - all_times = np.unique(time_vals) - unit_to_idx = {u: i for i, u in enumerate(all_units)} - time_to_idx = {t: i for i, t in enumerate(all_times)} - n_units = len(all_units) - n_times = len(all_times) - n_cov = len(covariates) if covariates else 0 - # [intercept, unit_1..unit_{U-1}, time_1..time_{T-1}] — the intercept (col 0) - # makes the column space span the constant / grand mean (see docstring). - n_fe_cols = 1 + (n_units - 1) + (n_times - 1) - - def _build_rows(mask=None): - """Build sparse matrix for given observation mask.""" - all_rows = np.arange(n) - - # Intercept (col 0): 1 for every (masked) row. - i_rows = all_rows if mask is None else all_rows[mask] - i_cols = np.zeros(len(i_rows), dtype=int) - - # Unit dummies (drop first) at cols 1..n_units-1 - u_indices = np.array([unit_to_idx[u] for u in unit_vals]) - u_mask = u_indices > 0 - if mask is not None: - u_mask = u_mask & mask - - u_rows = all_rows[u_mask] - u_cols = u_indices[u_mask] # 1..n_units-1 (intercept occupies col 0) - - # Time dummies (drop first) at cols n_units..n_units+n_times-2 - t_indices = np.array([time_to_idx[t] for t in time_vals]) - t_mask = t_indices > 0 - if mask is not None: - t_mask = t_mask & mask - - t_rows = all_rows[t_mask] - t_cols = n_units + t_indices[t_mask] - 1 - - rows = np.concatenate([i_rows, u_rows, t_rows]) - cols = np.concatenate([i_cols, u_cols, t_cols]) - data = np.ones(len(rows)) - - A_fe = sparse.csr_matrix((data, (rows, cols)), shape=(n, n_fe_cols)) - - if n_cov > 0: - cov_data = df[covariates].values.copy() - if mask is not None: - cov_data[~mask] = 0.0 - A_cov = sparse.csr_matrix(cov_data) - A = sparse.hstack([A_fe, A_cov], format="csr") - else: - A = A_fe - - return A - - X_1 = _build_rows(mask=None) - X_10 = _build_rows(mask=omega_0) - - return X_1, X_10, unit_to_idx, time_to_idx - # ========================================================================= # sklearn-compatible interface # ========================================================================= @@ -3469,6 +2198,134 @@ def print_summary(self) -> None: print(self.summary()) +# ============================================================================= +# Post-fit aggregation kit (M-022/M-119) +# ============================================================================= + + +def _build_twostage_aggregation_kit( + *, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + omega_1_mask: pd.Series, + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + kept_cov_mask: Optional[np.ndarray], + cluster_var: str, + treatment_groups: List[Any], + ref_period: int, + survey_weights: Optional[np.ndarray], + survey_weight_type: str, + resolved_survey: Any, + uses_replicate: bool, + keep_mask: pd.Series, + score_pad_mask: Optional[np.ndarray], + cluster_ids_full: Optional[np.ndarray], + overall_att: float, + n_treated_obs: int, + survey_df_stage2: Optional[int], + survey_df_final: Optional[int], + survey_metadata: Optional[Any], + pretrends: bool, + horizon_max: Optional[int], + rank_deficient_action: str, + alpha: float, + anticipation: int, +) -> AggregationKit: + """Build the PANEL-BACKED post-fit aggregation kit (rows M-022/M-119). + + TwoStageDiD's event-study/group aggregation is a fresh Stage-2 OLS + + joint GMM sandwich per level — no compact influence payload can honor + a different ``balance_e`` post-fit — so the kit retains a + COLUMN-SUBSET COPY of the working frame (only the columns the moved + methods read by name, deduplicated: ``cluster=`` may legally name the + unit/time/first_treat column) plus the Stage-1 FE model and survey + objects. This is the FIRST new panel retention on TwoStageDiD results + (a documented memory-contract change, O(n_obs); replicate designs + additionally retain the (n_obs x R) replicate matrix via + ``resolved_survey`` — the identifier-minimization guarantee of the + CS/EDiD kits deliberately does NOT hold here). + + ``score_pad_mask``/``cluster_ids_full`` are the Wave-E.3-GATED values + fit actually passed (None unless the always-treated pad is active) — + storing the ungated arrays would fire the zero-pad branch on baseline + survey fits and break inertness. + + Value snapshots isolate recompute from public-field mutation + (``treatment_groups`` copied; ``survey_metadata`` a + ``dataclasses.replace`` copy the ES carrier builds from). df channels: + ``survey_df_stage2`` (what the stage-2 aggregators received — the + recompute seed), ``survey_df_final`` (what the stored overall + inference received — the 'simple' relay df), and the metadata copy's + own ``df_survey`` (the fit-final container channel). + + ``influence`` is EMPTY BY DESIGN: the recompute is panel-backed + (scores re-derived from the retained frame), not IF-payload-backed. + """ + subset_cols = list( + dict.fromkeys( + [unit, time, outcome, first_treat] + + list(covariates or []) + + [cluster_var] + + (["_survey_cluster"] if "_survey_cluster" in df.columns else []) + + ["_never_treated", "_rel_time", "_y_tilde"] + ) + ) + bookkeeping: Dict[str, Any] = { + # Panel-backed working objects (column-subset copy; index/row order + # preserved so numerics are identical) + "df": df[subset_cols].copy(), + "outcome": outcome, + "unit": unit, + "time": time, + "first_treat": first_treat, + "covariates": list(covariates) if covariates else None, + "omega_0_mask": omega_0_mask, + "omega_1_mask": omega_1_mask, + "unit_fe": unit_fe, + "time_fe": time_fe, + "grand_mean": grand_mean, + "delta_hat": delta_hat, + "kept_cov_mask": kept_cov_mask, + "cluster_var": cluster_var, + "treatment_groups": list(treatment_groups), + "ref_period": ref_period, + "survey_weights": survey_weights, + "survey_weight_type": survey_weight_type, + "resolved_survey": resolved_survey, + "uses_replicate": bool(uses_replicate), + "keep_mask": keep_mask, + "score_pad_mask": score_pad_mask, + "cluster_ids_full": cluster_ids_full, + # Value snapshots (isolation from public-field / estimator mutation) + "overall_att": float(overall_att), + "n_treated_obs": int(n_treated_obs), + "survey_df_stage2": survey_df_stage2, + "survey_df_final": survey_df_final, + "survey_metadata": ( + dataclasses.replace(survey_metadata) if survey_metadata is not None else None + ), + "pretrends": pretrends, + "horizon_max": horizon_max, + "rank_deficient_action": rank_deficient_action, + } + return AggregationKit( + bookkeeping=bookkeeping, + influence={}, + alpha=alpha, + anticipation=anticipation, + cband=False, # no simultaneous-band concept on this estimator + bootstrap=None, # replay not wired; results.aggregate() fails closed + ) + + # ============================================================================= # Convenience function # ============================================================================= @@ -3481,8 +2338,8 @@ def two_stage_did( 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["SurveyDesign"] = None, vcov_type: str = "hc1", **kwargs, @@ -3507,9 +2364,14 @@ def two_stage_did( covariates : list of str, optional Covariate column names. aggregate : str, optional - Aggregation mode: None, "simple", "event_study", "group", "all". + DEPRECATED (3.9, removed in 4.0; row M-022): forwarded to ``fit()``, + which warns — aggregate post-fit via + ``results.aggregate('event_study')`` instead. A plain wrapper call + (kwarg not supplied) never warns. balance_e : int, optional - Balance event study to cohorts observed at all relative times. + DEPRECATED (3.9, removed in 4.0; row M-119): forwarded to ``fit()``, + which warns — moves onto ``results.aggregate('event_study', + balance_e=...)``. survey_design : SurveyDesign, optional Survey design specification for design-based inference. Supports pweight only (aweight/fweight raise ValueError). Supports strata, @@ -3534,8 +2396,9 @@ def two_stage_did( >>> from diff_diff import two_stage_did, generate_staggered_data >>> data = generate_staggered_data(seed=42) >>> results = two_stage_did(data, 'outcome', 'unit', 'period', - ... 'first_treat', aggregate='event_study') + ... 'first_treat') >>> results.print_summary() + >>> results.aggregate('event_study').summary() # post-fit aggregation """ est = TwoStageDiD(vcov_type=vcov_type, **kwargs) return est.fit( diff --git a/diff_diff/two_stage_aggregation.py b/diff_diff/two_stage_aggregation.py new file mode 100644 index 00000000..37692131 --- /dev/null +++ b/diff_diff/two_stage_aggregation.py @@ -0,0 +1,1564 @@ +"""Stage-2 aggregation + GMM variance engine for the Gardner two-stage estimator. + +Extracted verbatim from ``diff_diff/two_stage.py`` (and, for the two shared +static helpers, ``two_stage_bootstrap.py``) for the M-022/M-119 post-fit +``aggregate()`` migration: ``two_stage.py`` imports ``two_stage_results.py`` +and ``two_stage_bootstrap.py`` (which imports ``two_stage_results.py`` too), +so the results module can import neither -- the shared machinery lives here, +an import-leaf module both sides can reach (the +``efficient_did_aggregation.py`` / ``imputation_aggregation.py`` precedent). + +Contents: + +- module helpers ``_SPARSE_DENSE_THRESHOLD``, ``_LSMRUnconvergedError`` and + ``_lsmr_certified_normal_solve`` (``two_stage.py`` re-imports all three -- + ``spillover.py`` and the bootstrap module's lazy imports keep working); +- :class:`_TwoStageAggregationMixin` -- the three Stage-2 aggregation levels + (static / event-study / group), the joint GMM sandwich they recompute + through, the Stage-1 helpers the replicate replay refits with, and the + replicate-weight inference override replay. Inherited by ``TwoStageDiD`` + (fit-time behavior byte-identical) and hosted post-fit by the throwaway + ``_TwoStageKitAggregator`` (``two_stage_results.py``). +""" + +import warnings +from typing import Any, Dict, List, Optional, Set, Tuple + +import numpy as np +import pandas as pd +from scipy import sparse +from scipy.sparse.linalg import factorized as sparse_factorized + +from diff_diff.linalg import _rank_guarded_inv, solve_ols +from diff_diff.utils import _iterative_fe_solve, demean_by_groups, safe_inference + +# Maximum number of elements before falling back to per-column sparse aggregation. +# 10M float64 elements ≈ 80 MB peak allocation. Above this, per-column .getcol() +# trades throughput for bounded memory. Keep in sync with two_stage_bootstrap.py. +_SPARSE_DENSE_THRESHOLD = 10_000_000 + + +class _LSMRUnconvergedError(RuntimeError): + """LSMR could not certify the Stage-1 normal-equation solve; the + variance boundary converts this to NaN inference (fail-closed).""" + + +def _lsmr_certified_normal_solve( + gram_csc, rhs: np.ndarray, context: str = "TwoStageDiD GMM sandwich" +) -> np.ndarray: + """Least-squares solve of the (possibly singular) sparse Stage-1 Gram + system ``gram @ out = rhs`` via per-column LSMR — no dense + materialization of the ``(p_1, p_1)`` normal matrix (`O((U+T+K)^2)` + OOM risk on large panels; the pattern the ImputationDiD LSMR fix + closed, ported here after the consumer-invariance analysis). + + OUTPUT-PRESERVING despite the min-norm ambiguity on singular systems: + least-squares solutions differ only by a ``null(X'X) = null(X_10)`` + component (weighted: ``null(X'WX) = null(W^{1/2}X_10)`` — zero-weight + rows are inert in every weighted consumer because Psi/score/residual + contributions carry the same ``W`` factor), and EVERY ``gamma_hat`` + consumer is an ``X_10``-range functional. One ``theta_exact`` consumer + (the bootstrap exact-residual helper's ``X_1_sparse @ theta_exact``) + evaluates theta on TREATED rows where a ``null(X_10)`` component would + NOT annihilate — parity there holds for a second reason: both dense + ``lstsq`` (SVD) and LSMR return the MIN-NORM least-squares solution, so + the two solvers agree on the whole vector (to iterative tolerance), not + just on range functionals; the fit-level singular-design parity test + locks this at the V/SE level. The remaining consumers are the + ``X_10``-range functionals — ``Psi_stage1 = X_10 @ gamma_hat``, the GMM + score correction ``c_g' gamma_hat`` with ``c_g = X_{10,g}' eps_{10,g}`` + in ``rowspace(X_10)``, and Stage-1 residuals ``y - X_10 theta`` — so + every null component annihilates. Locked by the singular-system parity + test against a dense-lstsq oracle. + + CONVERGENCE IS VALIDATED (fail-closed): ``istop`` in ``{0, 1, 2, 4, 5}`` + certifies a solution / least-squares solution within tolerance (4/5 are + the machine-precision analogues of 1/2 per SciPy); anything else gets + ONE retry with an uncapped condition limit, then raises + :class:`_LSMRUnconvergedError` — converted to NaN inference at the + variance boundary rather than feeding an unverified solution into the + GMM sandwich. + """ + import scipy.sparse.linalg as spla + + _certified = (0, 1, 2, 4, 5) + rhs_2d = np.atleast_2d(np.asarray(rhs, dtype=np.float64)) + if rhs_2d.shape[0] == 1 and np.asarray(rhs).ndim == 1: + rhs_2d = rhs_2d.T + dim = gram_csc.shape[0] + out = np.empty((dim, rhs_2d.shape[1])) + for j in range(rhs_2d.shape[1]): + result = spla.lsmr(gram_csc, rhs_2d[:, j], atol=1e-14, btol=1e-14) + z, istop = result[0], int(result[1]) + if istop not in _certified or not np.all(np.isfinite(z)): + result = spla.lsmr( + gram_csc, + rhs_2d[:, j], + atol=1e-14, + btol=1e-14, + conlim=1e16, + maxiter=max(50 * dim, 10_000), + ) + z, istop = result[0], int(result[1]) + if istop not in _certified or not np.all(np.isfinite(z)): + warnings.warn( + f"{context}: the LSMR fallback solve of the " + f"Stage-1 normal equations did not converge (istop={istop}); " + "the affected variance is reported as NaN rather than from " + "an unverified solution.", + UserWarning, + stacklevel=3, + ) + raise _LSMRUnconvergedError(f"LSMR uncertified (istop={istop})") + out[:, j] = z + return out + + +class _TwoStageAggregationMixin: + """Shared Stage-2/GMM methods (moved verbatim from ``TwoStageDiD``). + + HOST-ATTRIBUTE CONTRACT -- the complete ``self.`` surface the moved + methods read (typed class-level declarations for ``mypy diff_diff`` at + zero errors on both hosts). Zero methods WRITE to ``self`` -- the + post-fit throwaway host exists for estimator-mutation isolation only. + """ + + alpha: float + pretrends: bool + horizon_max: Optional[int] + rank_deficient_action: str + + def _iterative_fe( + self, + y: np.ndarray, + unit_vals: np.ndarray, + time_vals: np.ndarray, + idx: pd.Index, + max_iter: int = 10_000, + tol: float = 1e-10, + weights: Optional[np.ndarray] = None, + ) -> Tuple[Dict[Any, float], Dict[Any, float]]: + """ + Estimate unit and time FE via iterative alternating projection. + + Thin wrapper over the shared bincount solver + (``diff_diff.utils._iterative_fe_solve``): factorize unit/time once, + solve on integer codes, map the level arrays back to dicts. + + Parameters + ---------- + idx : pd.Index + Unused; retained for call-site stability. + weights : np.ndarray, optional + Survey weights (weighted group means ``sum(w*x)/sum(w)``). A + unit/period whose observations ALL carry zero weight has no + identifying contribution and gets ``NaN`` FE (its key is kept so + the rank-condition membership check still sees the group). + + Returns + ------- + unit_fe : dict + Mapping from unit -> unit fixed effect. + time_fe : dict + Mapping from time -> time fixed effect. + """ + unit_codes, unit_uniques = pd.factorize(unit_vals, sort=False) + time_codes, time_uniques = pd.factorize(time_vals, sort=False) + if (unit_codes < 0).any() or (time_codes < 0).any(): + raise ValueError( + "TwoStageDiD: unit or time column contains NaN. Drop or " + "impute missing group keys before fitting." + ) + unit_fe_arr, time_fe_arr = _iterative_fe_solve( + np.asarray(y, dtype=np.float64), + unit_codes.astype(np.intp, copy=False), + time_codes.astype(np.intp, copy=False), + len(unit_uniques), + len(time_uniques), + weights=weights, + max_iter=max_iter, + tol=tol, + method_name="TwoStageDiD iterative FE solver", + ) + unit_fe = dict(zip(unit_uniques, unit_fe_arr)) + time_fe = dict(zip(time_uniques, time_fe_arr)) + return unit_fe, time_fe + + def _fit_untreated_model( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + weights: Optional[np.ndarray] = None, + ) -> Tuple[ + Dict[Any, float], Dict[Any, float], float, Optional[np.ndarray], Optional[np.ndarray] + ]: + """ + Stage 1: Estimate unit + time FE on untreated observations. + + Parameters + ---------- + weights : np.ndarray, optional + Full-panel survey weights (same length as df). The untreated subset + is extracted internally via omega_0_mask. When None, unweighted. + + Returns + ------- + unit_fe, time_fe, grand_mean, delta_hat, kept_cov_mask + """ + df_0 = df.loc[omega_0_mask] + w_0 = weights[omega_0_mask.values] if weights is not None else None + + if covariates is None or len(covariates) == 0: + y = df_0[outcome].values.copy() + unit_fe, time_fe = self._iterative_fe( + y, df_0[unit].values, df_0[time].values, df_0.index, weights=w_0 + ) + return unit_fe, time_fe, 0.0, None, None + + else: + y = df_0[outcome].values.copy() + X_raw = df_0[covariates].values.copy() + units = df_0[unit].values + times = df_0[time].values + + # Within-transform Y and all X columns through the shared MAP + # engine (factorize-once + bincount + optional Rust kernel), one + # dispatch for every column. within_transform pins [unit, time]; + # [time, unit] here preserves the historical time-then-unit sweep + # order of the per-estimator loops. + narrow = df_0[[outcome, *covariates, time, unit]].copy() + demeaned, _ = demean_by_groups( + narrow, + [outcome, *covariates], + [time, unit], + inplace=True, + weights=w_0, + max_iter=10_000, + tol=1e-10, + ) + y_dm = demeaned[outcome].to_numpy(dtype=np.float64) + X_dm = demeaned[covariates].to_numpy(dtype=np.float64) + + result = solve_ols( + X_dm, + y_dm, + return_vcov=False, + rank_deficient_action=self.rank_deficient_action, + column_names=covariates, + weights=w_0, + ) + delta_hat = result[0] + kept_cov_mask = np.isfinite(delta_hat) + delta_hat_clean = np.where(np.isfinite(delta_hat), delta_hat, 0.0) + + y_adj = y - np.dot(X_raw, delta_hat_clean) + unit_fe, time_fe = self._iterative_fe(y_adj, units, times, df_0.index, weights=w_0) + + return unit_fe, time_fe, 0.0, delta_hat_clean, kept_cov_mask + + def _residualize( + self, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + covariates: Optional[List[str]], + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + ) -> np.ndarray: + """ + Compute residualized outcome y_tilde for ALL observations. + + y_tilde_i = y_i - mu_hat_i - eta_hat_t [- X_i @ delta_hat] + """ + alpha_i = df[unit].map(unit_fe).values + beta_t = df[time].map(time_fe).values + + # Handle missing FE (NaN for units/periods not in untreated sample) + alpha_i = np.where(pd.isna(alpha_i), np.nan, alpha_i).astype(float) + beta_t = np.where(pd.isna(beta_t), np.nan, beta_t).astype(float) + + y_hat = grand_mean + alpha_i + beta_t + + if delta_hat is not None and covariates: + y_hat = y_hat + np.dot(df[covariates].values, delta_hat) + + y_tilde = df[outcome].values - y_hat + return y_tilde + + @staticmethod + def _mask_nan_ytilde(y_tilde, warn: bool = True): + """Mask non-finite y_tilde values and warn if any found. + + Returns the boolean mask of non-finite values. Modifies y_tilde in-place + (sets NaN values to 0.0). ``warn=False`` suppresses the UserWarning - + used ONLY by the replicate-refit closures, where zero-weight replicate + designs (JK1/BRR) make NaN FE for zeroed-out PSUs expected mechanics + (the main-fit warning still fires once; per-replicate repeats would + emit up to ~3x n_replicates copies of the same message). + """ + nan_mask = ~np.isfinite(y_tilde) + if nan_mask.any(): + n_nan = int(nan_mask.sum()) + if warn: + warnings.warn( + f"{n_nan} observation(s) have non-finite imputed outcomes " + f"(y_tilde) from unidentified fixed effects. These " + f"observations are excluded from ATT estimation.", + UserWarning, + stacklevel=3, + ) + y_tilde[nan_mask] = 0.0 + return nan_mask + + def _stage2_static( + self, + df: pd.DataFrame, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + omega_1_mask: pd.Series, + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + cluster_var: str, + kept_cov_mask: Optional[np.ndarray], + survey_weights: Optional[np.ndarray] = None, + survey_weight_type: str = "pweight", + resolved_survey=None, + score_pad_mask: Optional[np.ndarray] = None, + cluster_ids_full: Optional[np.ndarray] = None, + warn_nan: bool = True, + ) -> Tuple[float, float]: + """ + Static (simple ATT) Stage 2: OLS of y_tilde on D_it. + + Returns (att, se). + """ + y_tilde = df["_y_tilde"].values.copy() + nan_mask = self._mask_nan_ytilde(y_tilde, warn=warn_nan) + + D = omega_1_mask.values.astype(float) + # Zero out treatment indicator for NaN y_tilde obs (don't count in ATT) + D[nan_mask] = 0.0 + + # X_2: treatment indicator (no intercept) + X_2 = D.reshape(-1, 1) + + # Avoid degenerate case where all treated obs have NaN y_tilde + if D.sum() == 0: + return np.nan, np.nan + + # Stage 2 OLS for point estimate (discard naive SE) + coef, residuals, _ = solve_ols( + X_2, + y_tilde, + return_vcov=False, + weights=survey_weights, + weight_type=survey_weight_type, + ) + att = float(coef[0]) + + # GMM sandwich variance + # An uncertified LSMR Stage-1 fallback solve fails closed: + # NaN vcov -> NaN SE/t/p/CI (the helper already warned). + try: + V = self._compute_gmm_variance( + df=df, + unit=unit, + time=time, + covariates=covariates, + omega_0_mask=omega_0_mask, + unit_fe=unit_fe, + time_fe=time_fe, + delta_hat=delta_hat, + kept_cov_mask=kept_cov_mask, + X_2=X_2, + cluster_ids=df[cluster_var].values, + survey_weights=survey_weights, + resolved_survey=resolved_survey, + score_pad_mask=score_pad_mask, + cluster_ids_full=cluster_ids_full, + ) + except _LSMRUnconvergedError: + V = np.full((X_2.shape[1], X_2.shape[1]), np.nan) + + se = float(np.sqrt(max(V[0, 0], 0.0))) + return att, se + + def _stage2_event_study( + self, + df: pd.DataFrame, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + omega_1_mask: pd.Series, + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + cluster_var: str, + treatment_groups: List[Any], + ref_period: int, + balance_e: Optional[int], + kept_cov_mask: Optional[np.ndarray], + survey_weights: Optional[np.ndarray] = None, + survey_weight_type: str = "pweight", + survey_df: Optional[int] = None, + resolved_survey=None, + score_pad_mask: Optional[np.ndarray] = None, + cluster_ids_full: Optional[np.ndarray] = None, + warn_nan: bool = True, + ) -> Tuple[Dict[int, Dict[str, Any]], Optional[np.ndarray], Optional[List[int]]]: + """Event study Stage 2: OLS of y_tilde on relative-time dummies. + + Returns ``(effects, vcov, vcov_index)``: the per-horizon effects + dict, the full GMM variance-covariance matrix over the ESTIMATED + horizon coefficients, and the horizon labels ordering its + rows/columns. The reference period and Proposition-5 horizons are + never regression columns, so they appear in ``effects`` but not in + ``vcov_index``; all-filtered horizons (n_obs == 0) ARE columns, + with NaN-filled rows/columns from the rank guard. ``(dict, None, + None)`` on the degenerate early returns that fit no Stage-2 + regression. + """ + y_tilde = df["_y_tilde"].values.copy() + nan_mask = self._mask_nan_ytilde(y_tilde, warn=warn_nan) + rel_times = df["_rel_time"].values + n = len(df) + + # Get all horizons — include pre-periods when pretrends=True + if self.pretrends: + evt_rel = rel_times[~df["_never_treated"].values] + else: + evt_rel = rel_times[omega_1_mask.values] + all_horizons = sorted(set(int(h) for h in evt_rel if np.isfinite(h))) + + # Apply horizon_max filter + if self.horizon_max is not None: + all_horizons = [h for h in all_horizons if abs(h) <= self.horizon_max] + + # Apply balance_e filter + if balance_e is not None: + cohort_rel_times = self._build_cohort_rel_times(df, first_treat) + balanced_cohorts = set() + if all_horizons: + max_h = max(all_horizons) + required_range = set(range(-balance_e, max_h + 1)) + for g, horizons in cohort_rel_times.items(): + if required_range.issubset(horizons): + balanced_cohorts.add(g) + if not balanced_cohorts: + warnings.warn( + f"No cohorts satisfy balance_e={balance_e} requirement. " + "Event study results will contain only the reference period. " + "Consider reducing balance_e.", + UserWarning, + stacklevel=2, + ) + return ( + { + ref_period: { + "effect": 0.0, + "se": 0.0, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (0.0, 0.0), + "n_obs": 0, + } + }, + None, + None, + ) + balance_mask = df[first_treat].isin(balanced_cohorts).values + else: + balance_mask = np.ones(n, dtype=bool) + + # Check Proposition 5: no never-treated units + has_never_treated = df["_never_treated"].any() + h_bar = np.inf + if not has_never_treated and len(treatment_groups) > 1: + h_bar = max(treatment_groups) - min(treatment_groups) + + # Identify Prop 5 horizons and compute their actual treated obs counts. + # Treated obs have NaN y_tilde at these horizons (counterfactual + # unidentified), but actual_n counts them to distinguish from truly + # empty horizons. rel_times is NaN for untreated/never-treated obs + # (line ~653), so (rel_times == h) is False for them. + prop5_horizons = [] + prop5_effects: Dict[int, Dict[str, Any]] = {} + if h_bar < np.inf: + for h in all_horizons: + if h == ref_period: + continue + if h >= h_bar: + actual_n = int(np.sum((rel_times == h) & omega_1_mask.values & balance_mask)) + if actual_n > 0: + prop5_horizons.append(h) + prop5_effects[h] = { + "effect": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "n_obs": actual_n, + } + + # Remove reference period AND Prop 5 horizons from estimation + prop5_set = set(prop5_horizons) + est_horizons = [h for h in all_horizons if h != ref_period and h not in prop5_set] + + if len(est_horizons) == 0: + # No horizons to estimate — return the reference row PLUS any + # Proposition-5 rows (local-review fix, 2(b) PR-3b): when EVERY + # non-reference horizon is Prop-5-unidentified, the rows must + # still surface as all-NaN with n_obs > 0 and the consolidated + # warning, exactly as on the normal path below — dropping them + # here reported real treated horizons as absent instead of + # unidentified (contra REGISTRY Prop-5 contract). + if prop5_horizons: + warnings.warn( + f"Horizons {prop5_horizons} are not identified without " + f"never-treated units (Proposition 5). Set to NaN.", + UserWarning, + stacklevel=2, + ) + ref_only: Dict[int, Dict[str, Any]] = { + ref_period: { + "effect": 0.0, + "se": 0.0, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (0.0, 0.0), + "n_obs": 0, + } + } + ref_only.update(prop5_effects) + return (ref_only, None, None) + + # Build Stage 2 design: one column per horizon (no intercept) + # Never-treated obs get all-zero rows (undefined relative time -> NaN) + # With no intercept, they contribute zero to X'_2 X_2 and X'_2 y_tilde + horizon_to_col = {h: j for j, h in enumerate(est_horizons)} + k = len(est_horizons) + X_2 = np.zeros((n, k)) + + for i in range(n): + if not balance_mask[i]: + continue + if nan_mask[i]: + continue # NaN y_tilde -> don't include in event study + h = rel_times[i] + if np.isfinite(h): + h_int = int(h) + if h_int in horizon_to_col: + X_2[i, horizon_to_col[h_int]] = 1.0 + + # Stage 2 OLS + coef, residuals, _ = solve_ols( + X_2, + y_tilde, + return_vcov=False, + weights=survey_weights, + weight_type=survey_weight_type, + ) + + # GMM variance for full coefficient vector + # An uncertified LSMR Stage-1 fallback solve fails closed: + # NaN vcov -> NaN SE/t/p/CI (the helper already warned). + try: + V = self._compute_gmm_variance( + df=df, + unit=unit, + time=time, + covariates=covariates, + omega_0_mask=omega_0_mask, + unit_fe=unit_fe, + time_fe=time_fe, + delta_hat=delta_hat, + kept_cov_mask=kept_cov_mask, + X_2=X_2, + cluster_ids=df[cluster_var].values, + survey_weights=survey_weights, + resolved_survey=resolved_survey, + score_pad_mask=score_pad_mask, + cluster_ids_full=cluster_ids_full, + ) + except _LSMRUnconvergedError: + V = np.full((X_2.shape[1], X_2.shape[1]), np.nan) + + # Build results dict + event_study_effects: Dict[int, Dict[str, Any]] = {} + + # Reference period marker + event_study_effects[ref_period] = { + "effect": 0.0, + "se": 0.0, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (0.0, 0.0), + "n_obs": 0, + } + + for h in est_horizons: + j = horizon_to_col[h] + n_obs = int(np.sum(X_2[:, j])) + + if n_obs == 0: + event_study_effects[h] = { + "effect": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "n_obs": 0, + } + continue + + effect = float(coef[j]) + se = float(np.sqrt(max(V[j, j], 0.0))) + + t_stat, p_val, ci = safe_inference(effect, se, alpha=self.alpha, df=survey_df) + + event_study_effects[h] = { + "effect": effect, + "se": se, + "t_stat": t_stat, + "p_value": p_val, + "conf_int": ci, + "n_obs": n_obs, + } + + # Add Proposition 5 entries (unidentified horizons with n_obs > 0) + event_study_effects.update(prop5_effects) + + if prop5_horizons: + warnings.warn( + f"Horizons {prop5_horizons} are not identified without " + f"never-treated units (Proposition 5). Set to NaN.", + UserWarning, + stacklevel=2, + ) + + return event_study_effects, V, [int(h) for h in est_horizons] + + def _stage2_group( + self, + df: pd.DataFrame, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + omega_1_mask: pd.Series, + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + grand_mean: float, + delta_hat: Optional[np.ndarray], + cluster_var: str, + treatment_groups: List[Any], + kept_cov_mask: Optional[np.ndarray], + survey_weights: Optional[np.ndarray] = None, + survey_weight_type: str = "pweight", + survey_df: Optional[int] = None, + resolved_survey=None, + score_pad_mask: Optional[np.ndarray] = None, + cluster_ids_full: Optional[np.ndarray] = None, + warn_nan: bool = True, + ) -> Dict[Any, Dict[str, Any]]: + """Group (cohort) Stage 2: OLS of y_tilde on cohort dummies.""" + y_tilde = df["_y_tilde"].values.copy() + nan_mask = self._mask_nan_ytilde(y_tilde, warn=warn_nan) + n = len(df) + + # Build Stage 2 design: one column per cohort (no intercept) + group_to_col = {g: j for j, g in enumerate(treatment_groups)} + k = len(treatment_groups) + X_2 = np.zeros((n, k)) + + ft_vals = df[first_treat].values + treated_mask = omega_1_mask.values + for i in range(n): + if treated_mask[i] and not nan_mask[i]: + g = ft_vals[i] + if g in group_to_col: + X_2[i, group_to_col[g]] = 1.0 + + # Stage 2 OLS + coef, residuals, _ = solve_ols( + X_2, + y_tilde, + return_vcov=False, + weights=survey_weights, + weight_type=survey_weight_type, + ) + + # GMM variance + # An uncertified LSMR Stage-1 fallback solve fails closed: + # NaN vcov -> NaN SE/t/p/CI (the helper already warned). + try: + V = self._compute_gmm_variance( + df=df, + unit=unit, + time=time, + covariates=covariates, + omega_0_mask=omega_0_mask, + unit_fe=unit_fe, + time_fe=time_fe, + delta_hat=delta_hat, + kept_cov_mask=kept_cov_mask, + X_2=X_2, + cluster_ids=df[cluster_var].values, + survey_weights=survey_weights, + resolved_survey=resolved_survey, + score_pad_mask=score_pad_mask, + cluster_ids_full=cluster_ids_full, + ) + except _LSMRUnconvergedError: + V = np.full((X_2.shape[1], X_2.shape[1]), np.nan) + + group_effects: Dict[Any, Dict[str, Any]] = {} + for g in treatment_groups: + j = group_to_col[g] + n_obs = int(np.sum(X_2[:, j])) + + if n_obs == 0: + group_effects[g] = { + "effect": np.nan, + "se": np.nan, + "t_stat": np.nan, + "p_value": np.nan, + "conf_int": (np.nan, np.nan), + "n_obs": 0, + } + continue + + effect = float(coef[j]) + se = float(np.sqrt(max(V[j, j], 0.0))) + + t_stat, p_val, ci = safe_inference(effect, se, alpha=self.alpha, df=survey_df) + + group_effects[g] = { + "effect": effect, + "se": se, + "t_stat": t_stat, + "p_value": p_val, + "conf_int": ci, + "n_obs": n_obs, + } + + return group_effects + + @staticmethod + def _compute_gmm_scores( + c_by_cluster: np.ndarray, + gamma_hat: np.ndarray, + s2_by_cluster: np.ndarray, + ) -> np.ndarray: + """ + Compute per-cluster GMM scores S_g = gamma_hat' c_g - X'_{2g} eps_{2g}. + + Handles NaN/overflow from rank-deficient FE by wrapping in errstate + and replacing non-finite values with 0. + + Parameters + ---------- + c_by_cluster : np.ndarray, shape (G, p) + Per-cluster Stage 1 scores. + gamma_hat : np.ndarray, shape (p, k) + Cross-moment correction matrix. + s2_by_cluster : np.ndarray, shape (G, k) + Per-cluster Stage 2 scores. + + Returns + ------- + np.ndarray, shape (G, k) + Per-cluster influence scores. + """ + with np.errstate(invalid="ignore", divide="ignore", over="ignore"): + correction = np.dot(c_by_cluster, gamma_hat) + np.nan_to_num(correction, copy=False, nan=0.0, posinf=0.0, neginf=0.0) + return correction - s2_by_cluster + + def _compute_gmm_variance( + self, + df: pd.DataFrame, + unit: str, + time: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + unit_fe: Dict[Any, float], + time_fe: Dict[Any, float], + delta_hat: Optional[np.ndarray], + kept_cov_mask: Optional[np.ndarray], + X_2: np.ndarray, + cluster_ids: np.ndarray, + survey_weights: Optional[np.ndarray] = None, + resolved_survey=None, + score_pad_mask: Optional[np.ndarray] = None, + cluster_ids_full: Optional[np.ndarray] = None, + ) -> np.ndarray: + """ + Compute GMM sandwich variance (Butts & Gardner 2022). + + Matches the R `did2s` source code implementation: uses the GLOBAL + Hessian inverse (not per-cluster) and NO finite-sample adjustments. + + The per-observation influence function is: + IF_i = (X'_2 X_2)^{-1} [gamma_hat' x_{10i} eps_{10i} - x_{2i} eps_{2i}] + + where gamma_hat = (X'_{10} X_{10})^{-1} (X'_1 X_2) uses the GLOBAL + cross-moment. + + The cluster-robust variance is: + V = (X'_2 X_2)^{-1} (sum_g S_g S'_g) (X'_2 X_2)^{-1} + S_g = gamma_hat' c_g - X'_{2g} eps_{2g} + c_g = X'_{10g} eps_{10g} + + With survey weights W (diagonal): + Bread: (X'_2 W X_2)^{-1} + gamma_hat: (X'_{10} W X_{10})^{-1} (X'_1 W X_2) + c_g = sum_{i in g} w_i * x_{10i} * eps_{10i} + s2_g = sum_{i in g} w_i * x_{2i} * eps_{2i} + + Parameters + ---------- + X_2 : np.ndarray, shape (n, k) + Stage 2 design matrix (treatment indicators). The Stage-2 residual + ``eps_2`` is re-solved internally from the *exact* Stage-1 residuals + (see the exact-residual note below), so it is not a parameter. + cluster_ids : np.ndarray, shape (n,) + Cluster identifiers, fit-sample length. Used for the per-cluster + stage-1 / stage-2 score aggregation (OLS path). + survey_weights : np.ndarray, optional + Survey weights of shape (n,). When None, unweighted (identical + to current code). + resolved_survey : ResolvedSurveyDesign, optional + Resolved survey design. Under Wave E.3 parity (PR #482 SpilloverDiD + precedent) the design retains full-domain `n_psu` / `n_strata` / + `df_survey` / `strata` / `fpc` / `psu` arrays even when the + always-treated drop removes rows from the OLS sample. The + zero-padded per-cluster scores expand onto the full-domain PSU + list before stratified-meat dispatch. R `survey::svyrecvar(subset())` + convention (Lumley 2010 §2.5); mirrors `imputation.py:2175-2183` + (PreTrendsImputation) and `prep.py:1401-1432` (DCDH cell variance). + score_pad_mask : np.ndarray of shape (n_full,), bool, optional + Wave E.3 parity zero-pad mask. When supplied, indicates which + FULL-DOMAIN rows are present in the fit sample (True = kept + for OLS). Requires `n == int(np.sum(score_pad_mask))`. Co-supplied + with `cluster_ids_full`. Per-cluster stage-1 / stage-2 score + aggregates computed at fit-length are expanded onto the + full-domain unique-PSU list; PSUs absent from the fit sample + (e.g. PSUs containing only always-treated rows) get zero score + rows but still count toward `G_full` for `n_psu` / `df_survey`. + None (default) → no padding, exact pre-PR behavior. + cluster_ids_full : np.ndarray of shape (n_full,), optional + Full-domain PSU labels. Co-supplied with `score_pad_mask`. Must + share the same length. Provides the full-domain unique-PSU list + used both for score zero-pad expansion and for downstream + strata/FPC `obs_idx` lookups against the full-domain + `resolved_survey.strata` / `.fpc` arrays. None (default) → no + padding, exact pre-PR behavior. + + Returns + ------- + np.ndarray, shape (k, k) + Variance-covariance matrix. + """ + n = len(df) + k = X_2.shape[1] + + # Exclude rank-deficient covariates + cov_list = covariates + if covariates and kept_cov_mask is not None and not np.all(kept_cov_mask): + cov_list = [c for c, k_ in zip(covariates, kept_cov_mask) if k_] + + # Build sparse FE design matrices X_1 (all obs) and X_10 (untreated only) + X_1_sparse, X_10_sparse, unit_to_idx, time_to_idx = self._build_fe_design( + df, unit, time, cov_list, omega_0_mask + ) + + p = X_1_sparse.shape[1] + + # eps_10 = Y - X_10 @ gamma_hat + # Untreated: stage 1 residual (Y - fitted). Treated: Y (X_10 rows = 0). + # Reconstruct Y from y_tilde: Y = y_tilde + fitted_stage1. Because + # y_tilde = Y - fitted_1, the iterative FE in fitted_1 cancel exactly, so + # y_vals == Y (independent of the iterative solver's tolerance). + alpha_i = df[unit].map(unit_fe).values + beta_t = df[time].map(time_fe).values + # Identification mask: obs whose unit AND time FE are both identified by the + # untreated Stage-1 fit. Rank-deficient / Proposition-5 obs (NaN FE) keep the + # iterative-residual behavior; only identified obs get the exact residuals. + identified = np.isfinite(np.asarray(alpha_i, dtype=float)) & np.isfinite( + np.asarray(beta_t, dtype=float) + ) + alpha_i = np.where(pd.isna(alpha_i), 0.0, alpha_i).astype(float) + beta_t = np.where(pd.isna(beta_t), 0.0, beta_t).astype(float) + fitted_1 = alpha_i + beta_t + if delta_hat is not None and cov_list: + if kept_cov_mask is not None and not np.all(kept_cov_mask): + fitted_1 = fitted_1 + np.dot(df[cov_list].values, delta_hat[kept_cov_mask]) + else: + fitted_1 = fitted_1 + np.dot(df[cov_list].values, delta_hat) + + y_tilde = df["_y_tilde"].values + y_vals = y_tilde + fitted_1 # reconstruct Y + y_vals_clean = np.nan_to_num(y_vals, nan=0.0) + + omega_0 = omega_0_mask.values + + # 1. gamma_hat = (X'_{10} W X_{10})^{-1} (X'_1 W X_2) [p x k] + # With survey weights, both cross-products need W. We reuse the SAME + # factorization of (X'_{10} W X_{10}) to also solve the exact Stage-1 FE + # coefficients theta_exact (see exact-residual note below). + if survey_weights is not None: + XtWX_10 = X_10_sparse.T @ X_10_sparse.multiply(survey_weights[:, None]) + Xt1_WX2 = X_1_sparse.T @ (X_2 * survey_weights[:, None]) + rhs_fe = X_10_sparse.T @ (survey_weights * y_vals_clean) + else: + XtWX_10 = X_10_sparse.T @ X_10_sparse # (p x p) sparse + Xt1_WX2 = X_1_sparse.T @ X_2 # (p x k) dense + rhs_fe = X_10_sparse.T @ y_vals_clean # (p,) X'_{10} W Y + + try: + solve_XtX = sparse_factorized(XtWX_10.tocsc()) + if Xt1_WX2.ndim == 1: + gamma_hat = solve_XtX(Xt1_WX2).reshape(-1, 1) + else: + gamma_hat = np.column_stack( + [solve_XtX(Xt1_WX2[:, j]) for j in range(Xt1_WX2.shape[1])] + ) + theta_exact = np.asarray(solve_XtX(np.asarray(rhs_fe).ravel())).ravel() + except RuntimeError as exc: + # Singular matrix — fall back to certified sparse LSMR. Silent-failure + # audit axis C: emit a UserWarning on fallback instead of swallowing. + warnings.warn( + "TwoStageDiD GMM sandwich: sparse factorization of " + f"(X'_{{10}} W X_{{10}}) failed ({type(exc).__name__}); falling " + "back to sparse LSMR. This may indicate a rank-deficient or " + "near-singular Stage 1 design matrix and SE estimates may be " + "less reliable.", + UserWarning, + stacklevel=2, + ) + XtWX_10_csc = XtWX_10.tocsc() + gamma_hat = _lsmr_certified_normal_solve(XtWX_10_csc, Xt1_WX2) + theta_exact = _lsmr_certified_normal_solve( + XtWX_10_csc, np.asarray(rhs_fe).ravel() + ).ravel() + + # Exact Stage-1 / Stage-2 residuals. The point-estimate path uses the + # iterative alternating-projection FE solver (`_iterative_fe`), which + # converges only to ~1e-7 on unbalanced untreated panels; that error is + # negligible for the ATT but perturbs the variance by ~1% relative to the + # analytical GMM sandwich. The variance therefore re-solves the Stage-1 FE + # EXACTLY using the sparse normal equations already factorized for gamma_hat + # (theta_exact), matching R `did2s` to ~1e-7 and mirroring ImputationDiD's + # exact-sparse variance path. The shared `_exact_gmm_residuals` helper is + # used by BOTH this analytical path and the multiplier bootstrap + # (`_compute_cluster_S_scores`) so the influence function is single-sourced. + eps_10, eps_2 = self._exact_gmm_residuals( + X_1_sparse, + theta_exact, + y_vals_clean, + identified, + omega_0, + y_tilde, + X_2, + survey_weights, + ) + + # 2. Per-cluster Stage 1 scores: c_g = sum_{i in g} w_i * x_{10i} * eps_{10i} + # Only untreated obs have non-zero X_10 rows + # With survey weights: multiply eps_10 by survey_weights before sparse multiply + if survey_weights is not None: + weighted_eps_10 = survey_weights * eps_10 + else: + weighted_eps_10 = eps_10 + weighted_X10 = X_10_sparse.multiply(weighted_eps_10[:, None]) # sparse element-wise + + unique_clusters, cluster_indices = np.unique(cluster_ids, return_inverse=True) + G = len(unique_clusters) + + n_elements = weighted_X10.shape[0] * weighted_X10.shape[1] + c_by_cluster = np.zeros((G, p)) + if n_elements > _SPARSE_DENSE_THRESHOLD: + # Per-column path: limits peak memory for large FE matrices + weighted_X10_csc = weighted_X10.tocsc() + for j_col in range(p): + col_data = weighted_X10_csc.getcol(j_col).toarray().ravel() + np.add.at(c_by_cluster[:, j_col], cluster_indices, col_data) + else: + # Dense path: faster for moderate-size matrices + weighted_X10_dense = weighted_X10.toarray() + for j_col in range(p): + np.add.at(c_by_cluster[:, j_col], cluster_indices, weighted_X10_dense[:, j_col]) + + # 3. Per-cluster Stage 2 scores: s2_g = sum_{i in g} w_i * x_{2i} * eps_{2i} + if survey_weights is not None: + weighted_eps_2 = survey_weights * eps_2 + else: + weighted_eps_2 = eps_2 + weighted_X2 = X_2 * weighted_eps_2[:, None] # (n x k) dense + s2_by_cluster = np.zeros((G, k)) + for j_col in range(k): + np.add.at(s2_by_cluster[:, j_col], cluster_indices, weighted_X2[:, j_col]) + + # Wave E.3 parity (PR #482 SpilloverDiD precedent): when the caller + # supplies `score_pad_mask` + `cluster_ids_full`, expand per-cluster + # stage-1 / stage-2 score aggregates onto the FULL-DOMAIN unique-PSU + # list. PSUs absent from the fit sample (those containing only + # always-treated rows) get zero score rows but still count toward + # `G_full` for `n_psu` / `df_survey` accounting. Mirrors R + # `survey::svyrecvar(subset())` (Lumley 2010 §2.5) and the in-library + # convention at `imputation.py:2175-2183` (PreTrendsImputation) and + # `prep.py:1401-1432` (DCDH cell variance). Downstream strata / FPC + # lookups use `cluster_ids_for_lookup` so the obs_idx applies to the + # full-domain `resolved_survey.strata` / `.fpc` arrays. + if score_pad_mask is not None: + if cluster_ids_full is None: + raise ValueError( + "_compute_gmm_variance: score_pad_mask requires " + "cluster_ids_full to be co-supplied (Wave E.3 parity " + "contract — score zero-pad expansion needs the " + "full-domain PSU labels to align with resolved_survey)." + ) + if resolved_survey is None: + raise ValueError( + "_compute_gmm_variance: score_pad_mask requires " + "resolved_survey to be co-supplied (Wave E.3 parity " + "contract — zero-pad only meaningful under a survey " + "design that retains full-domain dimensions)." + ) + n_full = int(len(score_pad_mask)) + if int(len(cluster_ids_full)) != n_full: + raise ValueError( + "_compute_gmm_variance: score_pad_mask and " + "cluster_ids_full must share the FULL-DOMAIN length; " + f"got len(score_pad_mask)={n_full}, " + f"len(cluster_ids_full)={int(len(cluster_ids_full))}." + ) + if int(np.sum(score_pad_mask)) != n: + raise ValueError( + "_compute_gmm_variance: int(np.sum(score_pad_mask)) " + f"({int(np.sum(score_pad_mask))}) must equal the " + f"fit-sample length n ({n}) so the score expansion " + "is well-defined." + ) + unique_clusters_full = np.unique(cluster_ids_full) + G_full = int(len(unique_clusters_full)) + # Map fit-sample unique_clusters into positions in + # unique_clusters_full via searchsorted (both arrays sorted by + # np.unique). Verify the mapping is exact — otherwise the fit + # sample contains PSU labels absent from the full domain (a + # contract violation that should never occur under the upstream + # `_inject_cluster_as_psu` invariant). + fit_to_full_idx = np.searchsorted(unique_clusters_full, unique_clusters) + if not np.array_equal( + unique_clusters_full[fit_to_full_idx], np.asarray(unique_clusters) + ): + raise ValueError( + "_compute_gmm_variance: fit-sample unique cluster " + "labels are not a subset of full-domain cluster labels " + "(Wave E.3 parity invariant violated). This should be " + "impossible under `_inject_cluster_as_psu` — please " + "file an issue with a minimal reproducer." + ) + c_by_cluster_full = np.zeros((G_full, p)) + s2_by_cluster_full = np.zeros((G_full, k)) + c_by_cluster_full[fit_to_full_idx] = c_by_cluster + s2_by_cluster_full[fit_to_full_idx] = s2_by_cluster + c_by_cluster = c_by_cluster_full + s2_by_cluster = s2_by_cluster_full + unique_clusters = unique_clusters_full + G = G_full + cluster_ids_for_lookup = np.asarray(cluster_ids_full) + else: + cluster_ids_for_lookup = cluster_ids + + # 4. S_g = gamma_hat' c_g - X'_{2g} eps_{2g} + S = self._compute_gmm_scores(c_by_cluster, gamma_hat, s2_by_cluster) + + # 5. Meat: sum_g S_g S'_g = S' S + _use_stratified_meat = resolved_survey is not None and ( + resolved_survey.strata is not None or resolved_survey.fpc is not None + ) + if _use_stratified_meat: + from diff_diff.survey import _compute_stratified_meat_from_psu_scores + + # Build PSU→stratum and PSU→FPC mappings from observation-level arrays. + # cluster_ids_for_lookup is full-domain length under Wave E.3 parity + # (score_pad_mask path) and fit-sample length otherwise; either way it + # aligns with `resolved_survey.strata` / `resolved_survey.fpc` so the + # obs_idx lookup resolves to the correct stratum / FPC value. + G_meat = len(unique_clusters) + + # Strata: synthesize single stratum when strata is None (unstratified FPC) + if resolved_survey.strata is not None: + psu_strata = np.empty(G_meat, dtype=resolved_survey.strata.dtype) + for idx, c in enumerate(unique_clusters): + obs_idx = np.where(cluster_ids_for_lookup == c)[0][0] + psu_strata[idx] = resolved_survey.strata[obs_idx] + else: + psu_strata = np.zeros(G_meat, dtype=int) + + # FPC: map observation-level FPC to PSU level + psu_fpc = None + if resolved_survey.fpc is not None: + psu_fpc = np.empty(G_meat, dtype=np.float64) + for idx, c in enumerate(unique_clusters): + obs_idx = np.where(cluster_ids_for_lookup == c)[0][0] + psu_fpc[idx] = resolved_survey.fpc[obs_idx] + + # Unstratified single-PSU: variance is unidentified (matches + # _compute_stratified_psu_meat at survey.py:1225 which returns + # zero meat with no variance_computed flag for n_psu < 2). + # Under Wave E.3 parity, G_meat = G_full (post zero-pad), so the + # gate fires on the full-domain PSU count, not the fit-sample. + if resolved_survey.strata is None and G_meat < 2: + return np.full((k, k), np.nan) + + # Reorder S rows to match unique_clusters ordering + # S is built using np.add.at with cluster_indices from pd.factorize, + # which uses the same order as unique_clusters from the data. + meat, _var_computed, _legit_zero = _compute_stratified_meat_from_psu_scores( + psu_scores=S, + psu_strata=psu_strata, + fpc_per_psu=psu_fpc, + lonely_psu=resolved_survey.lonely_psu, + ) + # If no variance was computed and no legitimate zeros, variance + # is unidentified — return NaN VCV so caller gets NaN SE. + if not _var_computed and _legit_zero == 0: + return np.full((k, k), np.nan) + else: + with np.errstate(invalid="ignore", over="ignore"): + meat = S.T @ S # (k x k) + + # 6. Bread: (X'_2 W X_2)^{-1} + with np.errstate(invalid="ignore", over="ignore", divide="ignore"): + if survey_weights is not None: + XtWX_2 = X_2.T @ (X_2 * survey_weights[:, None]) + else: + XtWX_2 = X_2.T @ X_2 + # np.linalg.solve only raises on an *exactly* singular Gram; a *near*- + # singular X_2'WX_2 would otherwise flow a garbage inverse (~1e13) + # straight into the SE. `_rank_guarded_inv` truncates redundant + # directions on the equilibrated Gram -> finite SE on the identified + # subspace (NaN only at rank 0), matching the covariate IF rank-guard. + # Sibling of finding #17 (axis A): the prior fallback fired only on an + # exactly-singular matrix. X_2 is the Stage-2 indicator design (not user + # covariates), so the diagnostic guidance points at that layer. + bread, n_dropped, _, dropped = _rank_guarded_inv(XtWX_2, return_dropped=True) + if n_dropped: + warnings.warn( + "Rank-deficient second-stage design matrix X_2'WX_2 in " + "TwoStageDiD TSL variance; rank-reducing to a finite SE on the " + f"identified subspace ({n_dropped} redundant direction(s) " + "dropped, NaN if rank 0). The Stage-2 design is built from " + "treatment, event-time, or group indicators, so this typically " + "indicates a zero-weight or all-zero indicator column " + "(e.g. an aggregation path with no qualifying observations).", + UserWarning, + stacklevel=2, + ) + + # 7. V = bread @ meat @ bread + V = bread @ meat @ bread + # A dropped (unidentified) Stage-2 coefficient is zero-filled in `bread`, + # which would report se=0 for that named coefficient; NaN its row/col in + # the FINAL vcov so per-coefficient SE extraction yields NaN (not 0). + if dropped.any(): + V[dropped, :] = np.nan + V[:, dropped] = np.nan + return V + + def _build_fe_design( + self, + df: pd.DataFrame, + unit: str, + time: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + ) -> Tuple[sparse.csr_matrix, sparse.csr_matrix, Dict[Any, int], Dict[Any, int]]: + """ + Build sparse FE design matrices X_1 (all obs) and X_10 (untreated rows only). + + Column layout: [intercept, unit_1, ..., unit_{U-1}, time_1, ..., time_{T-1}, + cov_1, ..., cov_C] (drop first unit and first time for identification, with an + intercept). The intercept makes the column space span the constant (the grand + mean); the prior intercept-free layout silently omitted the grand mean from the + FE span, which biased the GMM-sandwich residuals when re-solved exactly. With + the intercept this is the standard full-rank two-way FE (matches fixest / R + ``did2s``). + + X_10 is identical to X_1 except that rows for treated observations are zeroed out. + + Returns + ------- + X_1_sparse : sparse.csr_matrix, shape (n, p) + X_10_sparse : sparse.csr_matrix, shape (n, p) + unit_to_idx : dict + time_to_idx : dict + """ + n = len(df) + unit_vals = df[unit].values + time_vals = df[time].values + omega_0 = omega_0_mask.values + + all_units = np.unique(unit_vals) + all_times = np.unique(time_vals) + unit_to_idx = {u: i for i, u in enumerate(all_units)} + time_to_idx = {t: i for i, t in enumerate(all_times)} + n_units = len(all_units) + n_times = len(all_times) + n_cov = len(covariates) if covariates else 0 + # [intercept, unit_1..unit_{U-1}, time_1..time_{T-1}] — the intercept (col 0) + # makes the column space span the constant / grand mean (see docstring). + n_fe_cols = 1 + (n_units - 1) + (n_times - 1) + + def _build_rows(mask=None): + """Build sparse matrix for given observation mask.""" + all_rows = np.arange(n) + + # Intercept (col 0): 1 for every (masked) row. + i_rows = all_rows if mask is None else all_rows[mask] + i_cols = np.zeros(len(i_rows), dtype=int) + + # Unit dummies (drop first) at cols 1..n_units-1 + u_indices = np.array([unit_to_idx[u] for u in unit_vals]) + u_mask = u_indices > 0 + if mask is not None: + u_mask = u_mask & mask + + u_rows = all_rows[u_mask] + u_cols = u_indices[u_mask] # 1..n_units-1 (intercept occupies col 0) + + # Time dummies (drop first) at cols n_units..n_units+n_times-2 + t_indices = np.array([time_to_idx[t] for t in time_vals]) + t_mask = t_indices > 0 + if mask is not None: + t_mask = t_mask & mask + + t_rows = all_rows[t_mask] + t_cols = n_units + t_indices[t_mask] - 1 + + rows = np.concatenate([i_rows, u_rows, t_rows]) + cols = np.concatenate([i_cols, u_cols, t_cols]) + data = np.ones(len(rows)) + + A_fe = sparse.csr_matrix((data, (rows, cols)), shape=(n, n_fe_cols)) + + if n_cov > 0: + cov_data = df[covariates].values.copy() + if mask is not None: + cov_data[~mask] = 0.0 + A_cov = sparse.csr_matrix(cov_data) + A = sparse.hstack([A_fe, A_cov], format="csr") + else: + A = A_fe + + return A + + X_1 = _build_rows(mask=None) + X_10 = _build_rows(mask=omega_0) + + return X_1, X_10, unit_to_idx, time_to_idx + + @staticmethod + def _exact_gmm_residuals( + X_1_sparse, + theta_exact: np.ndarray, + y_vals_clean: np.ndarray, + identified: np.ndarray, + omega_0: np.ndarray, + y_tilde: np.ndarray, + X_2: np.ndarray, + survey_weights: Optional[np.ndarray], + ) -> Tuple[np.ndarray, np.ndarray]: + """Exact Stage-1 / Stage-2 residuals for the GMM influence function. + + Given the EXACT Stage-1 FE coefficients ``theta_exact`` (solved from the + same ``(X'_{10} W X_{10})`` factorization used for ``gamma_hat``), return the + exact Stage-1 residual ``eps_10`` (untreated rows) and the exact Stage-2 + residual ``eps_2``. **Shared** by the analytical GMM variance + (``TwoStageDiD._compute_gmm_variance``) and the multiplier bootstrap + (``_compute_cluster_S_scores``) so both build the per-cluster influence + score ``S_g = gamma_hat' c_g - X'_{2g} eps_{2g}`` from the same exact + residuals. The iterative alternating-projection FE used for the point + estimate is only ~1e-7-accurate on unbalanced untreated panels, which + perturbs the variance ~1% relative to the analytical sandwich; obs whose FE + are unidentified (rank-deficient / Proposition-5) fall back to the iterative + residual ``y_tilde`` so those edge cases are unchanged. + """ + n = X_1_sparse.shape[0] + fitted_exact = np.asarray(X_1_sparse @ theta_exact).ravel() + y_tilde_exact = y_vals_clean - fitted_exact + use_exact = identified & np.isfinite(y_tilde_exact) + y_tilde_use = np.where(use_exact, y_tilde_exact, y_tilde) + eps_10 = np.empty(n) + eps_10[omega_0] = y_tilde_use[omega_0] # exact Stage-1 residual (untreated) + eps_10[~omega_0] = y_vals_clean[~omega_0] # x_{10i} = 0, so value is inert + # Exact Stage-2 residual: re-solve delta on the exact residualized outcome + # (X_2 already has NaN-y_tilde rows zeroed by the caller, so masked obs + # contribute nothing to the normal equations). + y_tilde_s2 = np.where(np.isfinite(y_tilde_use), y_tilde_use, 0.0) + if survey_weights is not None: + XtWX2 = X_2.T @ (X_2 * survey_weights[:, None]) + XtWy2 = X_2.T @ (survey_weights * y_tilde_s2) + else: + XtWX2 = X_2.T @ X_2 + XtWy2 = X_2.T @ y_tilde_s2 + try: + delta_2 = np.linalg.solve(XtWX2, XtWy2) + except np.linalg.LinAlgError: + # Silent-failure audit convention: warn before the dense fallback. + warnings.warn( + "TwoStageDiD GMM sandwich: Stage-2 design (X'_2 W X_2) is " + "singular; falling back to dense lstsq for the exact-residual " + "re-solve. This may indicate collinear treatment/horizon " + "indicators and SE estimates may be less reliable.", + UserWarning, + stacklevel=2, + ) + delta_2 = np.linalg.lstsq(XtWX2, XtWy2, rcond=None)[0] + eps_2 = y_tilde_s2 - X_2 @ delta_2 + return eps_10, eps_2 + + @staticmethod + def _build_cohort_rel_times( + df: pd.DataFrame, + first_treat: str, + ) -> Dict[Any, Set[int]]: + """Build mapping of cohort -> set of observed relative times.""" + treated_mask = ~df["_never_treated"] + treated_df = df.loc[treated_mask] + result: Dict[Any, Set[int]] = {} + ft_vals = treated_df[first_treat].values + rt_vals = treated_df["_rel_time"].values + for i in range(len(treated_df)): + h = rt_vals[i] + if np.isfinite(h): + result.setdefault(ft_vals[i], set()).add(int(h)) + return result + + def _replay_replicate_inference( + self, + *, + df: pd.DataFrame, + outcome: str, + unit: str, + time: str, + first_treat: str, + covariates: Optional[List[str]], + omega_0_mask: pd.Series, + omega_1_mask: pd.Series, + cluster_var: str, + treatment_groups: List[Any], + ref_period: int, + balance_e: Optional[int], + keep_mask: pd.Series, + resolved_survey: Any, + overall_att: float, + event_study_effects: Optional[Dict[int, Dict[str, Any]]], + group_effects: Optional[Dict[Any, Dict[str, Any]]], + survey_df_seed: Optional[int], + ) -> Tuple[np.ndarray, int, Optional[int]]: + """Replicate-weight inference override for the aggregation surfaces. + + Extracted verbatim from ``fit()``'s replicate block so the post-fit + ``aggregate()`` path can replay it from the kit payload. The stacked + layout is ``[overall, es..., grp...]`` built from whichever family + dicts are non-None. ``compute_replicate_refit_variance`` validates + replicates JOINTLY (all-finite rows), so inference is LEVEL-MATCHED: + post-fit ``aggregate('event_study')`` replays ``[overall, ES]`` and + reproduces ``fit(aggregate='event_study')`` exactly; a + ``fit(aggregate='all')`` surface is NOT the equivalence target when + a replicate NaNs on exactly one family's targets. + + Overrides ``se``/``t_stat``/``p_value``/``conf_int`` IN PLACE on the + passed dicts and returns ``(vcov, n_valid, survey_df_final)``. The + fit-only tail -- the overall-row quintet and the + ``survey_metadata.df_survey`` write -- deliberately stays in + ``fit()``: this method must never mutate results-owned metadata + (post-fit ``aggregate()`` calls it on kit refs and a throwaway + host). Cost: one full per-replicate two-stage refit pass per call + (R refits + R frame copies), the same work the fit-time path does. + """ + from diff_diff.survey import compute_replicate_refit_variance + + # Derive keys from actual outputs (excludes filtered/Prop5 horizons) + _es_effects_ts = event_study_effects or {} + _grp_effects_ts = group_effects or {} + _sorted_es_periods_ts = sorted( + e for e in _es_effects_ts.keys() if np.isfinite(_es_effects_ts[e]["effect"]) + ) + _sorted_groups_ts = sorted( + g for g in _grp_effects_ts.keys() if np.isfinite(_grp_effects_ts[g]["effect"]) + ) + _n_es_ts = len(_sorted_es_periods_ts) + _n_grp_ts = len(_sorted_groups_ts) + + # Build full-sample estimate from actual outputs + _full_est_ts = [overall_att] + _full_est_ts.extend([_es_effects_ts[e]["effect"] for e in _sorted_es_periods_ts]) + _full_est_ts.extend([_grp_effects_ts[g]["effect"] for g in _sorted_groups_ts]) + + def _refit_ts(w_r): + # Wave E.3 parity (PR #482 SpilloverDiD precedent): the main fit + # path keeps `resolved_survey` at full-domain length but subsets + # `survey_weights` for stage-1 / stage-2 OLS arithmetic via + # `keep_mask` (always-treated drop). The replicate refit + # callback receives a FULL-DOMAIN replicate weight `w_r` + # (sourced from `resolved_survey.replicate_weights` which is + # also full-domain) and must apply the SAME `keep_mask` + # subsetting before threading through stage-1 / stage-2, + # otherwise `solve_ols` rejects the length mismatch + # (full-domain w_r vs post-drop df) and the ValueError is + # swallowed by `compute_replicate_refit_variance` → + # NaN replicate inference. + w_r_fit = np.asarray(w_r)[keep_mask.values] + ufe_r, tfe_r, gm_r, delta_r, kcm_r = self._fit_untreated_model( + df, + outcome, + unit, + time, + covariates, + omega_0_mask, + weights=w_r_fit, + ) + y_tilde_r = self._residualize( + df, + outcome, + unit, + time, + covariates, + ufe_r, + tfe_r, + gm_r, + delta_r, + ) + df_tmp = df.copy() + df_tmp["_y_tilde"] = y_tilde_r + results = [] + + att_r, _ = self._stage2_static( + df=df_tmp, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=ufe_r, + time_fe=tfe_r, + grand_mean=gm_r, + delta_hat=delta_r, + cluster_var=cluster_var, + kept_cov_mask=kcm_r, + survey_weights=w_r_fit, + survey_weight_type="pweight", + warn_nan=False, + ) + results.append(att_r) + + if _sorted_es_periods_ts: + # Replicate refits only need the point effects; the + # per-replicate V is irrelevant to the refit variance. + es_r, _, _ = self._stage2_event_study( + df=df_tmp, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=ufe_r, + time_fe=tfe_r, + grand_mean=gm_r, + delta_hat=delta_r, + cluster_var=cluster_var, + treatment_groups=treatment_groups, + ref_period=ref_period, + balance_e=balance_e, + kept_cov_mask=kcm_r, + survey_weights=w_r_fit, + survey_weight_type="pweight", + survey_df=None, + warn_nan=False, + ) + for e in _sorted_es_periods_ts: + results.append(es_r[e]["effect"] if e in es_r else np.nan) + + if _sorted_groups_ts: + grp_r = self._stage2_group( + df=df_tmp, + unit=unit, + time=time, + first_treat=first_treat, + covariates=covariates, + omega_0_mask=omega_0_mask, + omega_1_mask=omega_1_mask, + unit_fe=ufe_r, + time_fe=tfe_r, + grand_mean=gm_r, + delta_hat=delta_r, + cluster_var=cluster_var, + treatment_groups=treatment_groups, + kept_cov_mask=kcm_r, + survey_weights=w_r_fit, + survey_weight_type="pweight", + survey_df=None, + warn_nan=False, + ) + for g in _sorted_groups_ts: + results.append(grp_r[g]["effect"] if g in grp_r else np.nan) + + return np.array(results) + + _vcov_rep_ts, _n_valid_rep_ts = compute_replicate_refit_variance( + _refit_ts, np.array(_full_est_ts), resolved_survey + ) + + # Override df if replicates were dropped + # Replicate-refit path is only reached with a resolved design. + assert resolved_survey is not None + survey_df_final = survey_df_seed + if _n_valid_rep_ts < resolved_survey.n_replicates: + survey_df_final = _n_valid_rep_ts - 1 if _n_valid_rep_ts > 1 else 0 + + # Override event-study SEs (only for identified effects) + for i, e in enumerate(_sorted_es_periods_ts): + if event_study_effects is not None and e in event_study_effects: + se_e = float(np.sqrt(max(_vcov_rep_ts[1 + i, 1 + i], 0.0))) + eff_e = event_study_effects[e]["effect"] + t_e, p_e, ci_e = safe_inference(eff_e, se_e, alpha=self.alpha, df=survey_df_final) + event_study_effects[e]["se"] = se_e + event_study_effects[e]["t_stat"] = t_e + event_study_effects[e]["p_value"] = p_e + event_study_effects[e]["conf_int"] = ci_e + + # Override group SEs (only for identified effects) + for j, g in enumerate(_sorted_groups_ts): + if group_effects is not None and g in group_effects: + se_g = float(np.sqrt(max(_vcov_rep_ts[1 + _n_es_ts + j, 1 + _n_es_ts + j], 0.0))) + eff_g = group_effects[g]["effect"] + t_g, p_g, ci_g = safe_inference(eff_g, se_g, alpha=self.alpha, df=survey_df_final) + group_effects[g]["se"] = se_g + group_effects[g]["t_stat"] = t_g + group_effects[g]["p_value"] = p_g + group_effects[g]["conf_int"] = ci_g + + return _vcov_rep_ts, _n_valid_rep_ts, survey_df_final diff --git a/diff_diff/two_stage_bootstrap.py b/diff_diff/two_stage_bootstrap.py index b63d9ac4..055f71ac 100644 --- a/diff_diff/two_stage_bootstrap.py +++ b/diff_diff/two_stage_bootstrap.py @@ -64,65 +64,22 @@ def _compute_gmm_scores( s2_by_cluster: np.ndarray, ) -> np.ndarray: ... - @staticmethod - def _exact_gmm_residuals( - X_1_sparse, - theta_exact: np.ndarray, - y_vals_clean: np.ndarray, - identified: np.ndarray, - omega_0: np.ndarray, - y_tilde: np.ndarray, - X_2: np.ndarray, - survey_weights: Optional[np.ndarray], - ) -> Tuple[np.ndarray, np.ndarray]: - """Exact Stage-1 / Stage-2 residuals for the GMM influence function. - - Given the EXACT Stage-1 FE coefficients ``theta_exact`` (solved from the - same ``(X'_{10} W X_{10})`` factorization used for ``gamma_hat``), return the - exact Stage-1 residual ``eps_10`` (untreated rows) and the exact Stage-2 - residual ``eps_2``. **Shared** by the analytical GMM variance - (``TwoStageDiD._compute_gmm_variance``) and the multiplier bootstrap - (``_compute_cluster_S_scores``) so both build the per-cluster influence - score ``S_g = gamma_hat' c_g - X'_{2g} eps_{2g}`` from the same exact - residuals. The iterative alternating-projection FE used for the point - estimate is only ~1e-7-accurate on unbalanced untreated panels, which - perturbs the variance ~1% relative to the analytical sandwich; obs whose FE - are unidentified (rank-deficient / Proposition-5) fall back to the iterative - residual ``y_tilde`` so those edge cases are unchanged. - """ - n = X_1_sparse.shape[0] - fitted_exact = np.asarray(X_1_sparse @ theta_exact).ravel() - y_tilde_exact = y_vals_clean - fitted_exact - use_exact = identified & np.isfinite(y_tilde_exact) - y_tilde_use = np.where(use_exact, y_tilde_exact, y_tilde) - eps_10 = np.empty(n) - eps_10[omega_0] = y_tilde_use[omega_0] # exact Stage-1 residual (untreated) - eps_10[~omega_0] = y_vals_clean[~omega_0] # x_{10i} = 0, so value is inert - # Exact Stage-2 residual: re-solve delta on the exact residualized outcome - # (X_2 already has NaN-y_tilde rows zeroed by the caller, so masked obs - # contribute nothing to the normal equations). - y_tilde_s2 = np.where(np.isfinite(y_tilde_use), y_tilde_use, 0.0) - if survey_weights is not None: - XtWX2 = X_2.T @ (X_2 * survey_weights[:, None]) - XtWy2 = X_2.T @ (survey_weights * y_tilde_s2) - else: - XtWX2 = X_2.T @ X_2 - XtWy2 = X_2.T @ y_tilde_s2 - try: - delta_2 = np.linalg.solve(XtWX2, XtWy2) - except np.linalg.LinAlgError: - # Silent-failure audit convention: warn before the dense fallback. - warnings.warn( - "TwoStageDiD GMM sandwich: Stage-2 design (X'_2 W X_2) is " - "singular; falling back to dense lstsq for the exact-residual " - "re-solve. This may indicate collinear treatment/horizon " - "indicators and SE estimates may be less reliable.", - UserWarning, - stacklevel=2, - ) - delta_2 = np.linalg.lstsq(XtWX2, XtWy2, rcond=None)[0] - eps_2 = y_tilde_s2 - X_2 @ delta_2 - return eps_10, eps_2 + # Provided by _TwoStageAggregationMixin on the estimator MRO + # (moved there with the M-022 aggregate() migration). + @staticmethod + def _exact_gmm_residuals( + X_1_sparse: Any, + theta_exact: np.ndarray, + y_vals_clean: np.ndarray, + identified: np.ndarray, + omega_0: np.ndarray, + y_tilde: np.ndarray, + X_2: np.ndarray, + survey_weights: Optional[np.ndarray], + ) -> Tuple[np.ndarray, np.ndarray]: ... + + @staticmethod + def _build_cohort_rel_times(df: pd.DataFrame, first_treat: str) -> Dict[Any, Set[int]]: ... def _compute_cluster_S_scores( self, @@ -707,20 +664,3 @@ def _run_bootstrap( # ========================================================================= # Utility # ========================================================================= - - @staticmethod - def _build_cohort_rel_times( - df: pd.DataFrame, - first_treat: str, - ) -> Dict[Any, Set[int]]: - """Build mapping of cohort -> set of observed relative times.""" - treated_mask = ~df["_never_treated"] - treated_df = df.loc[treated_mask] - result: Dict[Any, Set[int]] = {} - ft_vals = treated_df[first_treat].values - rt_vals = treated_df["_rel_time"].values - for i in range(len(treated_df)): - h = rt_vals[i] - if np.isfinite(h): - result.setdefault(ft_vals[i], set()).add(int(h)) - return result diff --git a/diff_diff/two_stage_results.py b/diff_diff/two_stage_results.py index 271f732f..c37d2eaf 100644 --- a/diff_diff/two_stage_results.py +++ b/diff_diff/two_stage_results.py @@ -5,14 +5,43 @@ dataclasses. Extracted from two_stage.py for module size management. """ +import dataclasses from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd +from diff_diff.aggregation import AggregationMixin, AggregationResult from diff_diff.results 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 +from diff_diff.two_stage_aggregation import _TwoStageAggregationMixin + + +class _TwoStageKitAggregator(_TwoStageAggregationMixin): + """Throwaway per-call host for the post-fit recompute (M-022/M-119). + + Hosts the moved Stage-2/GMM methods with exactly the mixin's declared + host-attribute contract, populated from KIT SNAPSHOTS — never from the + live estimator (whose config may have been mutated since the fit) and + never from mutable public results fields. A fresh instance per + ``aggregate()`` call; the moved methods write nothing to ``self``, so + the kit stays immutable either way. + """ + + def __init__( + self, + *, + alpha: float, + pretrends: bool, + horizon_max: Optional[int], + rank_deficient_action: str, + ) -> None: + self.alpha = alpha + self.pretrends = pretrends + self.horizon_max = horizon_max + self.rank_deficient_action = rank_deficient_action + __all__ = [ "TwoStageBootstrapResults", @@ -77,7 +106,7 @@ class TwoStageBootstrapResults: @dataclass -class TwoStageDiDResults(BaseResults): +class TwoStageDiDResults(BaseResults, AggregationMixin): """ Results from Gardner (2022) two-stage DiD estimation. @@ -168,6 +197,16 @@ class TwoStageDiDResults(BaseResults): event_study_vcov: Optional[np.ndarray] = field(default=None, repr=False) event_study_vcov_index: Optional[List[int]] = field(default=None, repr=False) event_study_df: Optional[float] = field(default=None, repr=False) + # Private panel-backed post-fit aggregation kit (rows M-022/M-119), + # attached by TwoStageDiD.fit(). None on results unpickled from a + # pre-3.9 release (aggregate() then fails with the re-fit message). + # Appended LAST (the generated __init__ positional indexes are public + # API). + _aggregation_kit: Optional[Any] = field(default=None, repr=False, compare=False) + + # Post-fit aggregation vocabulary (M-022). 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 @@ -190,6 +229,223 @@ def p_value(self) -> float: def t_stat(self) -> float: return self.overall_t_stat + # --- Post-fit aggregation (M-022/M-119) ------------------------------ + + 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 TwoStageDiDResults carries no aggregation kit - it is " + "attached by TwoStageDiD.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-target 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) + if level == "simple": + return self._aggregate_simple_result(kit) + # Fresh throwaway host per call, populated from KIT snapshots only + # (estimator/config mutation after fit() must not leak in). + agg = _TwoStageKitAggregator( + alpha=kit.alpha, + pretrends=bk["pretrends"], + horizon_max=bk["horizon_max"], + rank_deficient_action=bk["rank_deficient_action"], + ) + common: Dict[str, Any] = dict( + df=bk["df"], + unit=bk["unit"], + time=bk["time"], + first_treat=bk["first_treat"], + covariates=bk["covariates"], + omega_0_mask=bk["omega_0_mask"], + omega_1_mask=bk["omega_1_mask"], + unit_fe=bk["unit_fe"], + time_fe=bk["time_fe"], + grand_mean=bk["grand_mean"], + delta_hat=bk["delta_hat"], + cluster_var=bk["cluster_var"], + treatment_groups=bk["treatment_groups"], + kept_cov_mask=bk["kept_cov_mask"], + survey_weights=bk["survey_weights"], + survey_weight_type=bk["survey_weight_type"], + survey_df=bk["survey_df_stage2"], + resolved_survey=(None if bk["uses_replicate"] else bk["resolved_survey"]), + score_pad_mask=bk["score_pad_mask"], + cluster_ids_full=bk["cluster_ids_full"], + ) + if level == "group": + effects = agg._stage2_group(**common) + replay_df_g: Optional[int] = bk["survey_df_stage2"] + if bk["uses_replicate"]: + # LEVEL-MATCHED replay: [overall, groups] - reproduces + # fit(aggregate='group') exactly (see the replay docstring). + _, _, replay_df_g = agg._replay_replicate_inference( + df=bk["df"], + outcome=bk["outcome"], + unit=bk["unit"], + time=bk["time"], + first_treat=bk["first_treat"], + covariates=bk["covariates"], + omega_0_mask=bk["omega_0_mask"], + omega_1_mask=bk["omega_1_mask"], + cluster_var=bk["cluster_var"], + treatment_groups=bk["treatment_groups"], + ref_period=bk["ref_period"], + balance_e=None, + keep_mask=bk["keep_mask"], + resolved_survey=bk["resolved_survey"], + overall_att=bk["overall_att"], + event_study_effects=None, + group_effects=effects, + survey_df_seed=bk["survey_df_stage2"], + ) + return self._group_effects_to_aggregation(effects, kit, group_df=replay_df_g) + # level == "event_study" (the mixin validated the vocabulary) + es, es_vcov, es_vcov_index = agg._stage2_event_study( + ref_period=bk["ref_period"], balance_e=balance_e, **common + ) + replay_df: Optional[int] = None + if bk["uses_replicate"]: + _, _, replay_df = agg._replay_replicate_inference( + df=bk["df"], + outcome=bk["outcome"], + unit=bk["unit"], + time=bk["time"], + first_treat=bk["first_treat"], + covariates=bk["covariates"], + omega_0_mask=bk["omega_0_mask"], + omega_1_mask=bk["omega_1_mask"], + cluster_var=bk["cluster_var"], + treatment_groups=bk["treatment_groups"], + ref_period=bk["ref_period"], + balance_e=balance_e, + keep_mask=bk["keep_mask"], + resolved_survey=bk["resolved_survey"], + overall_att=bk["overall_att"], + event_study_effects=es, + group_effects=None, + survey_df_seed=bk["survey_df_stage2"], + ) + # Carrier + shared builder, reproducing fit's M-092 mode gates + # exactly (two_stage.py): analytical -> recomputed V + index + the + # finite-and->0 df scalar; replicate -> vcov/index None with the + # REPLAYED level-matched df; bootstrap unreachable (failed closed + # above). The carrier's metadata is a copy-on-use of the KIT's + # fit-final metadata copy (never the mutable public field). + if bk["uses_replicate"]: + vcov_final = None + index_final = None + df_gate = replay_df + else: + vcov_final = es_vcov + index_final = es_vcov_index + df_gate = bk["survey_df_stage2"] + es_df_final: Optional[float] = ( + float(df_gate) if df_gate is not None and np.isfinite(df_gate) and df_gate > 0 else None + ) + meta = bk["survey_metadata"] + if meta is not None: + if bk["uses_replicate"]: + meta = dataclasses.replace( + meta, df_survey=(replay_df if replay_df and replay_df > 0 else None) + ) + else: + meta = dataclasses.replace(meta) + carrier = dataclasses.replace( + self, + event_study_effects=es, + event_study_vcov=vcov_final, + event_study_vcov_index=index_final, + event_study_df=es_df_final, + survey_metadata=meta, + 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_obs`` (|Omega_1|, the D-column support the ATT + averages over) with ``n_kind="obs"``: TwoStageDiD's + ``n_treated_units``/``n_control_units`` unit sets OVERLAP (an + eventually-treated unit contributes untreated observations), so + the CS/EDiD disjoint-units convention cannot apply (the StackedDiD + carve-out class). + + ``df`` is the kit's ``survey_df_final`` snapshot - the exact value + the STORED overall ``safe_inference`` received (on a replicate fit + that value came from the ``[overall]``-only joint stack, which is + precisely why it must be snapshotted rather than re-derived: a + post-fit level-matched replay produces a different n_valid). + None → all-NaN df column; the replicate-undefined 0 sentinel NaNs + out via post_init. + """ + 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_treated_obs"]], dtype=float), + df=kit.bookkeeping["survey_df_final"], + alpha=kit.alpha, + n_kind="obs", + 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, *, group_df: Optional[int] + ) -> AggregationResult: + """Per-cohort AggregationResult from the recomputed group dict. + + ``df`` is a SCALAR relay (a deliberate divergence from + ImputationDiD's per-row ``df_used`` capture, documented in the + REGISTRY note): ``_stage2_group`` passes ONE immutable + ``survey_df`` parameter to every row's ``safe_inference``, so a + scalar broadcast is provenance-exact by construction and keeps the + moved method verbatim. Analytical fits relay the stage-2 seed; + replicate fits relay the REPLAYED level-matched value (the replay + rewrote every row's inference under it). ``weight=None``: Stage-2 + cohort dummies carry no cross-cohort mass (the CS rationale). + ``n_kind="obs"``: ``n_obs`` counts the cohort's treated + observations backing its indicator column. + """ + labels = list(effects.keys()) + 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_obs"] for g in labels], dtype=float), + df=(float(group_df) if group_df is not None else None), + alpha=kit.alpha, + n_kind="obs", + weight=None, + estimator=type(self).__name__.replace("Results", ""), + ) + def __repr__(self) -> str: """Concise string representation.""" sig = _get_significance_stars(self.overall_p_value) @@ -426,8 +682,12 @@ def to_dataframe(self, level: str = "event_study") -> pd.DataFrame: elif level == "event_study": if self.event_study_effects is None: raise ValueError( - "Event study effects not computed. " - "Use aggregate='event_study' or aggregate='all'." + "Event study effects not computed. Aggregate post-fit " + "instead - results.aggregate('event_study') returns the " + "EventStudyResults container (on a bootstrapped fit, " + "re-fit with n_bootstrap=0 or use the deprecated " + "fit-time aggregate=; a result unpickled from a pre-3.9 " + "release carries no kit and must be re-fit)." ) rows = [] for h, data in sorted(self.event_study_effects.items()): @@ -448,7 +708,12 @@ def to_dataframe(self, level: str = "event_study") -> pd.DataFrame: elif level == "group": if self.group_effects is None: raise ValueError( - "Group effects not computed. " "Use aggregate='group' or aggregate='all'." + "Group effects not computed. Aggregate post-fit instead " + "- results.aggregate('group') returns the " + "AggregationResult container (on a bootstrapped fit, " + "re-fit with n_bootstrap=0 or use the deprecated " + "fit-time aggregate=; a result unpickled from a pre-3.9 " + "release carries no kit and must be re-fit)." ) rows = [] for g, data in sorted(self.group_effects.items()): diff --git a/docs/api/_autosummary/diff_diff.ImputationDiDResults.rst b/docs/api/_autosummary/diff_diff.ImputationDiDResults.rst index 0ab48283..12addb10 100644 --- a/docs/api/_autosummary/diff_diff.ImputationDiDResults.rst +++ b/docs/api/_autosummary/diff_diff.ImputationDiDResults.rst @@ -12,6 +12,7 @@ .. autosummary:: ~ImputationDiDResults.__init__ + ~ImputationDiDResults.aggregate ~ImputationDiDResults.pretrend_test ~ImputationDiDResults.print_summary ~ImputationDiDResults.summary diff --git a/docs/api/_autosummary/diff_diff.TwoStageDiDResults.rst b/docs/api/_autosummary/diff_diff.TwoStageDiDResults.rst index 35aadf13..95b60171 100644 --- a/docs/api/_autosummary/diff_diff.TwoStageDiDResults.rst +++ b/docs/api/_autosummary/diff_diff.TwoStageDiDResults.rst @@ -12,6 +12,7 @@ .. autosummary:: ~TwoStageDiDResults.__init__ + ~TwoStageDiDResults.aggregate ~TwoStageDiDResults.print_summary ~TwoStageDiDResults.summary ~TwoStageDiDResults.to_dataframe diff --git a/docs/api/imputation.rst b/docs/api/imputation.rst index 751a97ba..f7b797e3 100644 --- a/docs/api/imputation.rst +++ b/docs/api/imputation.rst @@ -65,6 +65,7 @@ Results container for imputation DiD estimation. .. autosummary:: + ~ImputationDiDResults.aggregate ~ImputationDiDResults.summary ~ImputationDiDResults.print_summary ~ImputationDiDResults.to_dataframe @@ -105,9 +106,9 @@ Event study with visualization:: est = ImputationDiD() results = est.fit(data, outcome='outcome', unit='unit', - time='period', first_treat='first_treat', - aggregate='event_study') - plot_event_study(results) + time='period', first_treat='first_treat') + es = results.aggregate('event_study') # post-fit (M-021) + plot_event_study(es) Pre-trend test:: diff --git a/docs/api/two_stage.rst b/docs/api/two_stage.rst index 4f7c01aa..501bd4e4 100644 --- a/docs/api/two_stage.rst +++ b/docs/api/two_stage.rst @@ -68,6 +68,7 @@ Results container for two-stage DiD estimation. .. autosummary:: + ~TwoStageDiDResults.aggregate ~TwoStageDiDResults.summary ~TwoStageDiDResults.print_summary ~TwoStageDiDResults.to_dataframe @@ -108,9 +109,9 @@ Event study with visualization:: est = TwoStageDiD() results = est.fit(data, outcome='outcome', unit='unit', - time='period', first_treat='first_treat', - aggregate='event_study') - plot_event_study(results) + time='period', first_treat='first_treat') + es = results.aggregate('event_study') # post-fit (M-022) + plot_event_study(es) Comparison with other estimators:: diff --git a/docs/choosing_estimator.rst b/docs/choosing_estimator.rst index 0fa6f0f8..cdd1a6d6 100644 --- a/docs/choosing_estimator.rst +++ b/docs/choosing_estimator.rst @@ -518,9 +518,9 @@ impute counterfactual Y(0) for treated observations, then aggregate. imp = ImputationDiD() results = imp.fit(data, outcome='y', unit='unit_id', - time='period', first_treat='first_treat', - aggregate='event_study') + time='period', first_treat='first_treat') results.print_summary() + results.aggregate('event_study').summary() # post-fit aggregation .. note:: @@ -545,9 +545,9 @@ then regresses residuals on treatment indicators (stage 2). ts = TwoStageDiD() results = ts.fit(data, outcome='y', unit='unit_id', - time='period', first_treat='first_treat', - aggregate='event_study') + time='period', first_treat='first_treat') results.print_summary() + results.aggregate('event_study').summary() # post-fit aggregation .. note:: diff --git a/docs/dev-status.md b/docs/dev-status.md index f94cddc6..7ec61d46 100644 --- a/docs/dev-status.md +++ b/docs/dev-status.md @@ -19,12 +19,14 @@ Target: ideally < 1000 lines per module; modules ≥3000 lines are candidates fo | `had_pretests.py` | 4664 | Consider splitting (Stute / Yatchew / QUG / joint pretests) | | `diagnostic_report.py` | 4135 | Consider splitting (per-method renderers + provenance) | | `spillover.py` | 3655 | Consider splitting | -| `two_stage.py` | 3512 | Consider splitting | +| `two_stage.py` | 2430 | Monitor — exited the splitting band when the M-022 aggregate() migration extracted the Stage-2/GMM engine into `two_stage_aggregation.py` | | `power.py` | 3488 | Consider splitting (power analysis + MDE + sample size) | | `utils.py` | 3483 | Consider splitting | | `synthetic_control_results.py` | 3294 | Consider splitting | | `honest_did.py` | 3068 | Consider splitting | -| `imputation.py` | 2898 | Monitor | +| `imputation.py` | 1491 | Acceptable — dropped below 2000 when the M-021 aggregate() migration extracted the Theorem-3 engine into `imputation_aggregation.py` | +| `imputation_aggregation.py` | 1858 | Acceptable — verbatim-moved Theorem-3 aggregation/variance engine (M-021/M-118) | +| `two_stage_aggregation.py` | 1555 | Acceptable — verbatim-moved Stage-2/GMM aggregation engine (M-022/M-119) | | `synthetic_did.py` | 2826 | Monitor — variance methods + survey paths | | `business_report.py` | 2728 | Monitor — per-method narrative renderers | | `survey.py` | 2681 | Monitor — grew with Phase 6 features | diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index 5ff3663d..e38a6342 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -38,10 +38,12 @@ groups: - diff_diff/imputation.py - diff_diff/imputation_bootstrap.py - diff_diff/imputation_results.py + - diff_diff/imputation_aggregation.py two_stage: - diff_diff/two_stage.py - diff_diff/two_stage_bootstrap.py - diff_diff/two_stage_results.py + - diff_diff/two_stage_aggregation.py efficient_did: - diff_diff/efficient_did.py - diff_diff/efficient_did_aggregation.py diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 54062223..5f524452 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -1811,9 +1811,18 @@ where `W_it(h) = 1[K_it = h]` are lead indicators, estimated on `Omega_0` only. joint-test convention; it consumes gamma/V_gamma, not the per-lead t's, coheres with `df_convention="cluster"`, and converges with the t's at the v4 flip) — the per-lead-t vs joint-F pairing is a stated convention, not - a silent split. A `df_convention` value that never surfaces the per-lead - inference (`pretrends=False`, or `aggregate` outside - `{"event_study","all"}`) warns at fit time when explicitly non-default. + a silent split. Since the M-021 post-fit migration the + inert-config warning predicate is REACHABILITY-based, no longer + aggregate-keyed: an explicitly non-default value warns iff the per-lead + inference is unreachable for the fit config — reachable iff + `pretrends=True` AND not a replicate-weight design AND (the deprecated + fit-time `aggregate='event_study'/'all'` was supplied OR + `n_bootstrap <= 0`, since post-fit `results.aggregate('event_study')` + reaches the leads on any analytical pretrends fit but fails closed on + bootstrapped ones). Reachability-BASED, not exact: a fit whose bootstrap + later FAILS (`bootstrap_results=None`) can still aggregate post-fit, so + that corner warns spuriously — the warning fires before the bootstrap + runs and cannot know. *Edge cases:* - **Unbalanced panels:** FE estimated via iterative alternating projection (Gauss-Seidel), equivalent to OLS with unit+time dummies. Converges in O(max_iter) passes; typically 5-20 iterations for unbalanced panels, 1-2 for balanced. One-pass demeaning is only exact for balanced panels. @@ -1858,6 +1867,8 @@ where `W_it(h) = 1[K_it = h]` are lead indicators, estimated on `Omega_0` only. - [x] Supports balanced and unbalanced panels (iterative Gauss-Seidel demeaning for exact FE) - [x] Event study and group aggregation +- **Note (post-fit aggregate() - rows M-021/M-118):** `fit(aggregate=, balance_e=)` is deprecated (3.9; removed 4.0; joint FutureWarning, warn-and-still-work; the `imputation_did` wrapper forwards the shared sentinel so plain wrapper calls never warn) in favor of post-fit `ImputationDiDResults.aggregate(type, balance_e=)` - a PANEL-BACKED lazy recompute kit (not an EIF-payload kit: ES/group aggregation is a target-specific Theorem-3 recompute - each `balance_e` re-masks which treated observations enter every horizon and re-solves the untreated projection - so no compact influence payload can replace the frame). (a) RETAINED BUFFERS (memory contract): the kit's bookkeeping holds REFERENCES to the SAME per-fit objects `_fit_data` already retains for `pretrend_test()` - the working panel copy (all user columns plus `_tau_hat`/`_rel_time`/`_never_treated`), the Omega masks, `unit_fe`/`time_fe`/`grand_mean`/`delta_hat`/`kept_cov_mask`, the resolved survey design, and `survey_weights` - ZERO marginal memory, and pickles are unchanged via memoization (`_estimator_ref` already ships these objects); plus value SNAPSHOTS for isolation (a `treatment_groups` copy, config scalars, a `dataclasses.replace` copy of `survey_metadata`, `overall_att`, `n_treated_obs`) and TWO df-provenance scalars (`survey_df_seed`, what the analytical aggregators received; `survey_df_final`, what the stored overall inference received). Each `aggregate()` call runs on a fresh throwaway host with a call-local projection cache (the fit-local factorizations are unpicklable and never retained). (b) `balance_e` uses the BALANCED-WINDOW rule: a cohort is retained iff its observed relative-time set - checked against the FULL panel via `_build_cohort_rel_times()` - covers the contiguous window `[-balance_e, max_h]`; the SAME rule TwoStageDiD uses, divergent from CS/EfficientDiD's anchor-horizon rule. A window no cohort satisfies warns and yields the reference-marker-only dict (a legal near-empty container). (c) BOOTSTRAP fits fail closed for ALL levels including 'simple' (uniform CS/EDiD parity; the per-target psi machinery makes exact replay tractable - a TODO row). (d) CONTAINER ADMISSION into `compute_honest_did`/`compute_pretrends_power` is REJECTED BY DESIGN (both terminal TypeErrors state it): the surface carries no joint event-study covariance - per-horizon conservative SEs only (container `vcov=None`; the scalar `df_survey` channel is its only df provenance, the per-row hole being the tracked M-092-completion TODO row). (e) RELAY CONVENTIONS: 'simple' relays the stored overall quintet bit-exact with `n = n_treated_obs`, `n_kind="obs"` (the treated/control UNIT sets overlap - a treated unit with pre-periods counts in both - so the CS/EDiD disjoint-units convention cannot apply; |Omega_1| is the population the ATT averages over) and `df = survey_df_final`; 'group' rows carry per-row `df_used` captured at each row's `safe_inference` (the replicate override rewrites it, the bootstrap override clears it, the all-NaN cohort branch writes no key - consumers read via `.get`); 'event_study' rides the shared `_from_relative_dict` builder via a carrier whose metadata is a copy-on-use of the KIT's fit-final metadata copy. REPLICATE-WEIGHT fits replay the extracted `_replicate_override_aggregates` with a LEVEL-MATCHED stack: `compute_replicate_refit_variance` validates replicates JOINTLY (all-finite rows), so `aggregate(L)` reproduces `fit(aggregate=L)` exactly, a `fit(aggregate='all')` surface is NOT the equivalence target when a replicate NaNs on exactly one family's targets, and - the documented migration delta - moving a replicate fit from `fit(aggregate=)` to plain fit + post-fit `aggregate()` can change the public OVERALL row's se/CI/df on such degenerate designs (each surface self-consistent; pinned in the contract tests). `pretrends=True` + replicate: post-fit `aggregate('event_study')` raises the same NotImplementedError the fit-time gate raises (per-replicate lead refits unimplemented); 'group'/'simple' still work. Recompute re-emits the fit-time warnings (LSMR, Prop-5, empty-window) with fit-tuned stacklevels - post-fit attribution lands on a library frame, an accepted verbatim-move trade-off. + --- ## TwoStageDiD @@ -1916,11 +1927,12 @@ Our implementation uses multiplier bootstrap on the GMM influence function: clus - **Event study aggregation:** Horizon-specific effects use the same two-stage procedure with horizon indicator dummies in Stage 2. Unidentified horizons (e.g., long-run effects without never-treated units, per Proposition 5 of Borusyak et al. 2024) produce NaN. - **Pre-period event study coefficients (`pretrends=True`):** When enabled, the Stage 2 design matrix `X_2` includes pre-period relative-time dummies. Pre-period observations have `y_tilde = Step 1 residual` by construction. The GMM sandwich variance accounts for Stage 1 estimation error (Gardner 2022 §3.3; Newey-McFadden 1994, Theorem 6.1 — the paper has no numbered theorems). Only affects event study aggregation; overall ATT unchanged. - **balance_e with no qualifying cohorts:** If no cohorts have sufficient pre/post coverage for the requested `balance_e`, a warning is emitted and event study results contain only the reference period. +- **All-Proposition-5 event studies (3.9 fix, 2(b) PR-3b):** when every non-reference horizon is Proposition-5-unidentified (`h >= h_bar` with treated observations), the previously reference-only early return now retains the Prop-5 rows — all-NaN inference with `n_obs > 0` — and emits the consolidated Proposition-5 warning, identically on the fit-time and post-fit routes (previously the rows were dropped from the dict in this corner; the accompanying unidentified-FE warnings still fired, but real treated horizons read as absent rather than unidentified). - **No never-treated units (Proposition 5):** When there are no never-treated units and multiple treatment cohorts, horizons h >= h_bar (where h_bar = max(groups) - min(groups)) are unidentified per Proposition 5 of Borusyak et al. (2024). These produce NaN inference with n_obs > 0 (treated observations exist but counterfactual is unidentified) and a warning listing affected horizons. Matches ImputationDiD behavior. Proposition 5 applies to event study horizons only, not cohort aggregation — a cohort whose treated obs all fall at Prop 5 horizons naturally gets n_obs=0 in group effects because all its y_tilde values are NaN. - **Zero-observation horizons after filtering:** When `balance_e` or NaN `y_tilde` filtering results in zero observations for some non-Prop-5 event study horizons, those horizons produce NaN for all inference fields (effect, SE, t-stat, p-value, CI) with n_obs=0. - **Zero-observation cohorts in group effects:** If all treated observations for a cohort have NaN `y_tilde` (excluded from estimation), that cohort's group effect is NaN with n_obs=0. - **Note:** Survey weights in TwoStageDiD GMM sandwich via weighted cross-products: bread uses (X'_2 W X_2)^{-1}, gamma_hat uses (X'_{10} W X_{10})^{-1}(X'_1 W X_2), per-cluster scores multiply by survey weights. PSU clustering, stratification, and FPC are fully supported in the meat matrix via `_compute_stratified_meat_from_psu_scores()`. When strata or FPC are present, the meat computation replaces `S' S` with the stratified formula `sum_h (1 - f_h) * (n_h/(n_h-1)) * centered_h' centered_h`. Strata also enters survey df (n_PSU - n_strata) for t-distribution inference. Bootstrap + survey supported (Phase 6) via PSU-level multiplier weights. -- **Note (documented synthesis — Wave E.3 parity, full-domain survey design under always-treated drop):** when the always-treated handler drops units that lack untreated observations, TwoStageDiD preserves the FULL-DOMAIN resolved survey design (`n_psu`, `n_strata`, `df_survey`, `strata`, `fpc`, `psu`) for variance estimation. Per-cluster stage-1 / stage-2 score aggregates are computed at the post-drop fit-sample length and then zero-padded onto the full-domain unique-PSU list via `score_pad_mask` + `cluster_ids_full` kwargs on `_compute_gmm_variance`; PSUs that contain only always-treated rows get zero score rows but still count toward `G_full` for `n_psu` / `df_survey` accounting. Stage-1 / stage-2 OLS solve continues to operate on the post-drop sample (`survey_weights` subsetted for OLS arithmetic; bread `(X'_2 W X_2)^{-1}` unchanged because dropped rows contribute zero score under zero-padded weights). Mirrors SpilloverDiD Wave E.3 (PR #482, merge 24de9062) and adopts the canonical "zero-pad scores to full panel + retain full-design resolved survey" convention from R `survey::svyrecvar(subset())` (Lumley 2010 §2.5 "Domains and subpopulations") and the in-library precedents at `imputation.py:2175-2183` (PreTrendsImputation) and `prep.py:1401-1432` (DCDH cell variance). Cluster-injection (`_inject_cluster_as_psu`) operates on the FULL-DOMAIN cluster column (sourced from `data` pre-drop, not the post-drop `df`) so `resolved_survey.strata` and the injected `psu` array stay length-aligned. Pre-PR, the always-treated drop physically subsetted `resolved_survey.weights / strata / psu / fpc / replicate_weights` via `replace(resolved_survey, ...)` and recomputed `n_psu` / `n_strata` on the post-drop sample, producing artificially-deflated `df_survey` when a PSU contained only always-treated rows; tests at `tests/test_two_stage.py::TestTwoStageDiDWaveE3ParityAlwaysTreated` lock the parity contract. +- **Note (documented synthesis — Wave E.3 parity, full-domain survey design under always-treated drop):** when the always-treated handler drops units that lack untreated observations, TwoStageDiD preserves the FULL-DOMAIN resolved survey design (`n_psu`, `n_strata`, `df_survey`, `strata`, `fpc`, `psu`) for variance estimation. Per-cluster stage-1 / stage-2 score aggregates are computed at the post-drop fit-sample length and then zero-padded onto the full-domain unique-PSU list via `score_pad_mask` + `cluster_ids_full` kwargs on `_compute_gmm_variance`; PSUs that contain only always-treated rows get zero score rows but still count toward `G_full` for `n_psu` / `df_survey` accounting. Stage-1 / stage-2 OLS solve continues to operate on the post-drop sample (`survey_weights` subsetted for OLS arithmetic; bread `(X'_2 W X_2)^{-1}` unchanged because dropped rows contribute zero score under zero-padded weights). Mirrors SpilloverDiD Wave E.3 (PR #482, merge 24de9062) and adopts the canonical "zero-pad scores to full panel + retain full-design resolved survey" convention from R `survey::svyrecvar(subset())` (Lumley 2010 §2.5 "Domains and subpopulations") and the in-library precedents at `imputation_aggregation.py` (`_compute_lead_coefficients`, PreTrendsImputation) and `prep.py:1401-1432` (DCDH cell variance). Cluster-injection (`_inject_cluster_as_psu`) operates on the FULL-DOMAIN cluster column (sourced from `data` pre-drop, not the post-drop `df`) so `resolved_survey.strata` and the injected `psu` array stay length-aligned. Pre-PR, the always-treated drop physically subsetted `resolved_survey.weights / strata / psu / fpc / replicate_weights` via `replace(resolved_survey, ...)` and recomputed `n_psu` / `n_strata` on the post-drop sample, producing artificially-deflated `df_survey` when a PSU contained only always-treated rows; tests at `tests/test_two_stage.py::TestTwoStageDiDWaveE3ParityAlwaysTreated` lock the parity contract. - **Note:** The Stage-1 iterative FE solver (`_iterative_fe`) routes through the shared bincount Gauss-Seidel helper `diff_diff.utils._iterative_fe_solve`, and the covariate within-transformation routes through the shared MAP engine `diff_diff.utils.demean_by_groups` (factorize-once + `np.bincount`, optional Rust kernel; group order `[time, unit]` preserving the historical time-then-unit sweep) — the same convergence contract, accumulation-order numerics (~1e-10 vs the pre-3.7 pandas loops, not bit-for-bit), and `max_iter=10_000` budget documented under "Absorbed Fixed Effects with Survey Weights". Both surfaces emit `UserWarning` via `diff_diff.utils.warn_if_not_converged` when `max_iter` exhausts without reaching `tol` (the demean warning now carries the shared-engine label naming the affected variables rather than the estimator name). Silent return of the current iterate was classified as a silent failure under the Phase 2 audit and replaced with an explicit signal to match the logistic/Poisson IRLS pattern in `linalg.py`. - **Note:** Zero-total-weight groups (e.g. whole PSUs zeroed by JK1/BRR replicate weights, which reach Stage 1 unmasked — `keep_mask` only drops always-treated units): a unit/period whose observations ALL carry zero weight surfaces as `NaN` FE (key retained for the rank-condition membership check; matches the SpilloverDiD `_iterative_fe_subset` REGISTRY contract — never a silent finite `0.0`), and the shared demean engine's inert-row guard leaves those rows un-demeaned instead of NaN-poisoning the column. Before 3.7 the pandas loops divided 0/0 there: the covariate replicate path NaN-poisoned `y_dm`/`X_dm`, failed EVERY replicate refit inside `solve_ols(check_finite=True)`, and returned NaN SEs after a non-convergence warning storm. It now produces finite replicate SEs. `_mask_nan_ytilde`'s "non-finite imputed outcomes" `UserWarning` is suppressed (via `warn_nan=False`) ONLY inside the replicate-refit closures, where NaN FE for zeroed PSUs is expected mechanics — the main-fit warning is unchanged. - **Note:** When the Stage-2 bread `X'_2 W X_2` is singular, both the analytical TSL variance (`two_stage.py`) and the multiplier-bootstrap bread (`two_stage_bootstrap.py`) now emit a `UserWarning` before falling back to `np.linalg.lstsq`. Previously this fallback was silent. Sibling of axis-A finding #17 in the Phase 2 silent-failures audit; surfaced by the repo-wide lstsq-fallback pattern grep that accompanied the StaggeredTripleDifference fix. @@ -1944,6 +1956,8 @@ Our implementation uses multiplier bootstrap on the GMM influence function: clus - [x] Multiplier bootstrap on GMM influence function - [x] Event study and overall ATT aggregation +- **Note (post-fit aggregate() - rows M-022/M-119):** `fit(aggregate=, balance_e=)` is deprecated (3.9; removed 4.0; joint FutureWarning, warn-and-still-work; the `two_stage_did` wrapper forwards the shared sentinel so plain wrapper calls never warn) in favor of post-fit `TwoStageDiDResults.aggregate(type, balance_e=)` - a PANEL-BACKED lazy recompute kit: each level is a fresh Stage-2 OLS + joint Gardner-GMM sandwich on a level-specific design, so no compact influence payload exists. (a) RETAINED BUFFERS (memory contract - the FIRST panel retention on TwoStageDiD results, a deliberate break from the CS/EDiD identifier-minimization guarantee, with a `store_kit` opt-out tracked in DEFERRED.md): a COLUMN-SUBSET COPY of the working frame - `unit`/`time`/`outcome`/`first_treat` + covariates + the cluster column (deduplicated: `cluster=` may legally name a core column) + `_never_treated`/`_rel_time`/`_y_tilde` - O(n_obs) on every results object and pickle; the Stage-1 FE model (`unit_fe`/`time_fe`/`grand_mean`/`delta_hat`/`kept_cov_mask`), the Omega masks, the full-domain `keep_mask`, the Wave-E.3-GATED `score_pad_mask`/`cluster_ids_full` values fit actually passed (None unless the always-treated pad was active), `survey_weights`, and the resolved survey design - on replicate designs that adds the O(n_obs x R) replicate matrix; plus value snapshots (`treatment_groups` copy, `ref_period`, `overall_att`, `n_treated_obs`, a `dataclasses.replace` copy of `survey_metadata`) and TWO df scalars (`survey_df_stage2`, the recompute seed; `survey_df_final`, what the stored overall inference received). (b) `balance_e` uses the BALANCED-WINDOW rule (`[-balance_e, max_h]` coverage against the full panel - the ImputationDiD rule, divergent from CS/EfficientDiD's anchor-horizon rule); zero qualifying cohorts warns and yields the reference-row-only dict with `vcov=None`. (c) BOOTSTRAP fits fail closed for ALL levels including 'simple' (per-level GMM scores are function-locals; replay is a TODO row); a fit whose bootstrap FAILED (`bootstrap_results=None`, analytical inference retained) aggregates normally. (d) CONTAINER ADMISSION into `compute_honest_did`/`compute_pretrends_power` is DEFERRED, not by-design (both terminal TypeErrors state it): analytical surfaces DO carry the real joint Gardner-GMM covariance (M-092), but the pre-period coefficients are stage-1 residual MEANS - the reference horizon is dropped from the no-intercept Stage-2 design and the zero anchor row is appended mechanically - not contrasts against the advertised reference, while HonestDiD's Delta^RM/Delta^SD arithmetic hard-codes the `delta_0 = 0` normalization into its boundary/bridge constraints; admission awaits a normalization derivation (either re-estimating Stage 2 with the reference horizon in the design or deriving the residual-to-reference mapping) - the DEFERRED.md paper-gated row. (e) RELAY CONVENTIONS: 'simple' relays the stored overall quintet bit-exact with `n = n_treated_obs`, `n_kind="obs"` (overlapping unit sets - the StackedDiD carve-out class) and `df = survey_df_final` (on replicate fits that value came from the `[overall]`-only joint stack - snapshotted, never re-derived); 'group' relays a SCALAR df broadcast (deliberate divergence from ImputationDiD's per-row `df_used`: `_stage2_group` passes one immutable `survey_df` to every row's `safe_inference`, so the scalar is provenance-exact by construction and the moved method stays verbatim); 'event_study' reproduces the M-092 container contract exactly - analytical fits thread the recomputed joint vcov + `vcov_index` + the finite-and->0 df scalar through the carrier, replicate fits thread `vcov=None`/`index=None` with the REPLAYED level-matched df, and the carrier's metadata is a copy-on-use of the KIT's fit-final metadata copy. REPLICATE-WEIGHT fits replay the extracted `_replay_replicate_inference` with a LEVEL-MATCHED stack (the ImputationDiD semantics: `aggregate(L)` reproduces `fit(aggregate=L)`; `fit(aggregate='all')` is not the equivalence target on degenerate designs; the OVERALL-row migration delta on such designs is documented and pinned). Recompute re-emits fit-time warnings with fit-tuned stacklevels - an accepted verbatim-move trade-off. + --- ## StackedDiD @@ -5005,7 +5019,7 @@ where `D_it` is the treatment indicator (1 if unit `i` is treated by time `t`) a ### Event-study mode (Wave C, `event_study=True`) -Replaces the aggregate spec with the per-event-time × ring decomposition from Butts Section 5 / Table 2. Direct effects `tau_k` and per-(ring, event-time) spillover effects `delta_jk` are emitted in `att_dynamic` and a MultiIndex `spillover_effects`. A TwoStageDiD-compatible `event_study_effects: Dict[int, Dict]` alias (mirroring `two_stage.py:1355-1389` schema with `conf_int = (low, high)` tuple) is also emitted for `plot_event_study` consumption — `_extract_plot_data` prefers the new `reference_period` attribute over the legacy `n_obs==0` heuristic. `DiagnosticReport` routing is now wired (see the DiagnosticReport applicability Note below). +Replaces the aggregate spec with the per-event-time × ring decomposition from Butts Section 5 / Table 2. Direct effects `tau_k` and per-(ring, event-time) spillover effects `delta_jk` are emitted in `att_dynamic` and a MultiIndex `spillover_effects`. A TwoStageDiD-compatible `event_study_effects: Dict[int, Dict]` alias (mirroring `two_stage_aggregation.py::_stage2_event_study` row schema with `conf_int = (low, high)` tuple) is also emitted for `plot_event_study` consumption — `_extract_plot_data` prefers the new `reference_period` attribute over the legacy `n_obs==0` heuristic. `DiagnosticReport` routing is now wired (see the DiagnosticReport applicability Note below). - **Note (DiagnosticReport applicability):** `SpilloverDiDResults` is registered in `DiagnosticReport`'s `_APPLICABILITY` / `_PT_METHOD` tables with applicable checks {`parallel_trends` (method `event_study`, on the per-event-time DIRECT-effect dynamics — populated when `event_study=True`; Bonferroni fallback across pre-period coefs since there is no `event_study_vcov`), `design_effect` (instance-gated on `survey_metadata`), `heterogeneity` (reads `event_study_effects`)}. This mirrors the `TwoStageDiD` set **minus `bacon`**. `bacon` is **deliberately excluded**: SpilloverDiD identifies the direct effect off FAR-AWAY control observations (`d_it > d_bar`, Assumption 5), not off the TWFE 2×2 comparisons a Goodman-Bacon decomposition enumerates; running `bacon_decompose` on the raw binary treatment would ignore the ring/distance structure and pool spillover-contaminated in-ring units into the control group — the exact SUTVA violation this estimator handles (same rationale that excludes `bacon` for SyntheticControl / TROP / ContinuousDiD). `pretrends_power` / `sensitivity` / `epv` / `estimator_native` are not applicable (no `compute_pretrends_power` adapter, no HonestDiD adapter, no `epv_diagnostics`, not an SDiD/TROP/SCM native path). @@ -5013,7 +5027,7 @@ Replaces the aggregate spec with the per-event-time × ring decomposition from B **Note (endpoint binning, divergence from TwoStageDiD):** `horizon_max` bins event-times outside `[-H, +H]` into endpoint pools (`k ≤ -H` aggregated into one pre-bin dummy; `k ≥ +H` into one post-bin dummy). No observations are dropped. This **diverges intentionally** from `TwoStageDiD`'s `horizon_max` semantic, which **filters** rows with `|K_it| > H` out of the stage-2 sample. SpilloverDiD's bin-into-endpoint behavior honors the no-silent-data-drop policy; the divergence is documented on both estimators' docstrings to prevent future unification. With `horizon_max=None`, the helper auto-detects the event-time bin set from the observed K values (no binning). -**Note (reference period `-1 - anticipation`):** The reference period (the event-time dummy dropped to anchor the level interpretation) is `ref_period = -1 - anticipation`. With `anticipation=0`, `ref_period = -1` (standard event-study convention; coefficients are relative to one period before treatment). With `anticipation > 0`, the reference period shifts to `-1 - anticipation` so the "pre-treatment" anchor sits BEFORE the anticipation window. Mirrors `TwoStageDiD`'s convention at `two_stage.py:486`. The reference row appears in `att_dynamic` and `event_study_effects` with `coef = 0.0`, `se = 0.0`, `n_obs = 0`, `conf_int = (0.0, 0.0)` (TwoStageDiD parity, `two_stage.py:1355-1362`). When `horizon_max` is set and `ref_period < -horizon_max` (i.e., `anticipation > horizon_max - 1`), the fit raises `ValueError` — silently floor-shifting the reference to `-horizon_max` would change identification (rejected per `feedback_no_silent_failures`). +**Note (reference period `-1 - anticipation`):** The reference period (the event-time dummy dropped to anchor the level interpretation) is `ref_period = -1 - anticipation`. With `anticipation=0`, `ref_period = -1` (standard event-study convention; coefficients are relative to one period before treatment). With `anticipation > 0`, the reference period shifts to `-1 - anticipation` so the "pre-treatment" anchor sits BEFORE the anticipation window. Mirrors `TwoStageDiD`'s convention at `two_stage_aggregation.py` (`ref_period = -1 - anticipation`). The reference row appears in `att_dynamic` and `event_study_effects` with `coef = 0.0`, `se = 0.0`, `n_obs = 0`, `conf_int = (0.0, 0.0)` (TwoStageDiD parity, `two_stage_aggregation.py::_stage2_event_study`). When `horizon_max` is set and `ref_period < -horizon_max` (i.e., `anticipation > horizon_max - 1`), the fit raises `ValueError` — silently floor-shifting the reference to `-horizon_max` would change identification (rejected per `feedback_no_silent_failures`). **Note (post-finite_mask sample):** `att_dynamic["n_obs"]`, `event_study_effects[k]["n_obs"]`, AND the scalar `att` share weights all reflect the POST-`finite_mask` stage-2 estimation sample — not the pre-mask design built by `_build_event_study_design`. On warn-and-drop fits (baseline-treated units without Omega_0 rows are excluded via `finite_mask`), counts and weights are recomputed from `X_2_fit` so the reported metadata matches the actual stage-2 sample that `solve_ols` sees. diff --git a/docs/methodology/papers/gardner-2022-review.md b/docs/methodology/papers/gardner-2022-review.md index 25e535e9..84fc74c7 100644 --- a/docs/methodology/papers/gardner-2022-review.md +++ b/docs/methodology/papers/gardner-2022-review.md @@ -104,7 +104,7 @@ where `W_gpit = [Y_gpit, (1g)_gpit, (1p)_gpit, D_gp]`. The first moment block is - ⚠️ **Paper-permitted but NOT exposed by `TwoStageDiD`** (recorded as scope, not a deliverable): the **`P̄`-average estimand eq. (5)** (duration-restricted Stage-2 sample) and the **full-sample first-stage variant (fn 8)** — both described in the paper as valid modifications, but the library has no public parameter for either (`get_params()` = `anticipation, alpha, cluster, n_bootstrap, bootstrap_weights, seed, rank_deficient_action, horizon_max, pretrends, vcov_type`; Stage 1 is always untreated-only). The fn-9 last-pre-treatment / doubly-robust covariate variants are likewise not exposed (only the in-both-stages approach is). - [x] **Dedicated `tests/test_methodology_two_stage.py`** with eq./section-numbered Verified Components (Stage-1 FE recovery; Stage-2 overall ATT eq. 4 + event-study eq. 6; GMM first-stage-correction behavior; always-treated drop). — *PR-B.* - [x] **R parity fixture vs `did2s`** (`benchmarks/R/generate_did2s_golden.R` + `benchmarks/data/did2s_golden.json` + `did2s_test_panel.csv`); overall ATT (abs 1e-6) + event-study ATT (1e-6) / SE (1e-7). — *PR-B.* -- [x] **`two_stage.py` docstring/formula faithful; variance made exact in PR-B.** `_compute_gmm_variance`'s docstring (`two_stage.py:2842-2857`) states the code uses the GLOBAL Hessian inverse matching `did2s` and gives the correct sandwich `V = (X'₂X₂)^{-1}(Σ_g S_g S_g')(X'₂X₂)^{-1}` — the *formula* needed no correction (the Eq.6 misattribution was confined to the REGISTRY note, fixed in PR-A). However, the PR-B `did2s` SE parity surfaced a ~1% **numerical** gap: the variance derived its residuals from the *iterative* alternating-projection first-stage FE (`_iterative_fe`, converged to ~1e-7 on unbalanced untreated panels) while computing `gamma_hat` exactly. PR-B corrects this by re-solving the Stage-1 FE **exactly** inside the variance (sparse OLS, reusing the `gamma_hat` factorization) and adding an intercept to `_build_fe_design` so its column space spans the grand mean — yielding ~1e-7 `did2s` SE parity, mirroring ImputationDiD's exact-sparse variance. The point estimate stays on the iterative FE (twin-equivalence preserved). See REGISTRY `## TwoStageDiD` Notes. +- [x] **`two_stage.py` docstring/formula faithful; variance made exact in PR-B.** `_compute_gmm_variance`'s docstring (`two_stage_aggregation.py`) states the code uses the GLOBAL Hessian inverse matching `did2s` and gives the correct sandwich `V = (X'₂X₂)^{-1}(Σ_g S_g S_g')(X'₂X₂)^{-1}` — the *formula* needed no correction (the Eq.6 misattribution was confined to the REGISTRY note, fixed in PR-A). However, the PR-B `did2s` SE parity surfaced a ~1% **numerical** gap: the variance derived its residuals from the *iterative* alternating-projection first-stage FE (`_iterative_fe`, converged to ~1e-7 on unbalanced untreated panels) while computing `gamma_hat` exactly. PR-B corrects this by re-solving the Stage-1 FE **exactly** inside the variance (sparse OLS, reusing the `gamma_hat` factorization) and adding an intercept to `_build_fe_design` so its column space spans the grand mean — yielding ~1e-7 `did2s` SE parity, mirroring ImputationDiD's exact-sparse variance. The point estimate stays on the iterative FE (twin-equivalence preserved). See REGISTRY `## TwoStageDiD` Notes. --- diff --git a/docs/methodology/variance-conventions.md b/docs/methodology/variance-conventions.md index 4b136013..0d28c953 100644 --- a/docs/methodology/variance-conventions.md +++ b/docs/methodology/variance-conventions.md @@ -131,9 +131,11 @@ output). - **L3 — CallawaySantAnna / TwoStageDiD / ImputationDiD (default)**: different variance theory (influence functions / two-stage / BJS imputation), never the shared CR1 sandwich. CS is anchored to Stata csdid outright. - **ImputationDiD is conditional**: its pretrends=True + - aggregate="event_study" lead regression DOES run the shared clustered CR1 - and carries the K_reference increment there (+6, the [time, unit] no-intercept + **ImputationDiD is conditional**: its pretrends=True lead regression DOES + run the shared clustered CR1 — surfaced fit-time via the deprecated + aggregate="event_study" or post-fit via results.aggregate('event_study') + since the M-021 migration — and carries the K_reference increment there + (+6, the [time, unit] no-intercept increment on df_0 — pinned via expected_adjustment on its matrix row). - **L4 — hc2/hc2_bm** (leverage / Satterthwaite DOF — no CR1 factor), **survey TSL** (n_PSU - n_strata over the full design), and **Wooldridge diff --git a/docs/references.rst b/docs/references.rst index 80082986..b67cd3c8 100644 --- a/docs/references.rst +++ b/docs/references.rst @@ -260,7 +260,7 @@ Multi-Period and Staggered Adoption - **Lumley, T. (2010).** *Complex Surveys: A Guide to Analysis Using R.* Hoboken, NJ: John Wiley & Sons. https://doi.org/10.1002/9780470580066 - Authoritative reference for the R ``survey`` package. §2.5 "Domains and subpopulations" documents the ``svyrecvar(subset(design, mask))`` zero-pad convention: subpopulation analyses preserve the full survey design (strata / PSU / FPC) and treat out-of-domain observations as zero-score padding rows rather than physically subsetting the design. SpilloverDiD's Wave E.3 adopts this convention at the ``_compute_gmm_corrected_meat`` boundary, matching the existing in-library precedents at ``diff_diff/imputation.py:2175-2183`` (PreTrendsImputation lead regression) and ``diff_diff/prep.py:1401-1432`` (DCDH cell variance); see REGISTRY section "Variance (Wave E.3)". + Authoritative reference for the R ``survey`` package. §2.5 "Domains and subpopulations" documents the ``svyrecvar(subset(design, mask))`` zero-pad convention: subpopulation analyses preserve the full survey design (strata / PSU / FPC) and treat out-of-domain observations as zero-score padding rows rather than physically subsetting the design. SpilloverDiD's Wave E.3 adopts this convention at the ``_compute_gmm_corrected_meat`` boundary, matching the existing in-library precedents at ``diff_diff/imputation_aggregation.py::_compute_lead_coefficients`` (PreTrendsImputation lead regression) and ``diff_diff/prep.py:1401-1432`` (DCDH cell variance); see REGISTRY section "Variance (Wave E.3)". - **Wing, C., Freedman, S. M., & Hollingsworth, A. (2024).** "Stacked Difference-in-Differences." *NBER Working Paper* 32054. https://www.nber.org/papers/w32054 diff --git a/docs/tutorials/11_imputation_did.ipynb b/docs/tutorials/11_imputation_did.ipynb index 6b993adc..abd336f5 100644 --- a/docs/tutorials/11_imputation_did.ipynb +++ b/docs/tutorials/11_imputation_did.ipynb @@ -76,7 +76,20 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": "# Fit with event study aggregation and pre-period coefficients\nest = ImputationDiD(pretrends=True)\nresults_es = est.fit(data, outcome='outcome', unit='unit', time='period',\n first_treat='first_treat', aggregate='event_study')\n\n# Plot event study — pre-period region is automatically shaded\nif HAS_MATPLOTLIB:\n plot_event_study(results_es, title='Imputation DiD Event Study (with Pre-Trends)')\nelse:\n print(\"Install matplotlib to see visualizations: pip install matplotlib\")" + "source": [ + "# Fit once, then aggregate the event study POST-FIT (3.9, row M-021)\n", + "# - pre-period coefficients ride the pretrends=True config.\n", + "est = ImputationDiD(pretrends=True)\n", + "results_es = est.fit(data, outcome='outcome', unit='unit', time='period',\n", + " first_treat='first_treat')\n", + "es_surface = results_es.aggregate('event_study')\n", + "\n", + "# Plot event study — pre-period region is automatically shaded\n", + "if HAS_MATPLOTLIB:\n", + " plot_event_study(es_surface, title='Imputation DiD Event Study (with Pre-Trends)')\n", + "else:\n", + " print(\"Install matplotlib to see visualizations: pip install matplotlib\")" + ] }, { "cell_type": "code", @@ -84,8 +97,9 @@ "metadata": {}, "outputs": [], "source": [ - "# View event study effects as a table\n", - "results_es.to_dataframe(level='event_study')" + "# View event study effects as a table (unified container schema:\n", + "# event_time / att / se / ... / is_reference)\n", + "es_surface.to_dataframe()" ] }, { @@ -155,11 +169,10 @@ "metadata": {}, "outputs": [], "source": [ - "# Fit with group aggregation\n", + "# Aggregate per-cohort effects POST-FIT - no refit needed\n", "results_grp = ImputationDiD().fit(data, outcome='outcome', unit='unit',\n", - " time='period', first_treat='first_treat',\n", - " aggregate='group')\n", - "results_grp.to_dataframe(level='group')" + " time='period', first_treat='first_treat')\n", + "results_grp.aggregate('group').to_dataframe()" ] }, { @@ -235,4 +248,4 @@ }, "nbformat": 4, "nbformat_minor": 4 -} \ No newline at end of file +} diff --git a/docs/tutorials/12_two_stage_did.ipynb b/docs/tutorials/12_two_stage_did.ipynb index 4ff1ab60..d465953b 100644 --- a/docs/tutorials/12_two_stage_did.ipynb +++ b/docs/tutorials/12_two_stage_did.ipynb @@ -83,14 +83,15 @@ "metadata": {}, "outputs": [], "source": [ - "# Fit with event study aggregation\n", + "# Fit once, then aggregate the event study POST-FIT (3.9, row M-022)\n", "est = TwoStageDiD()\n", "results_es = est.fit(data, outcome='outcome', unit='unit', time='period',\n", - " first_treat='first_treat', aggregate='event_study')\n", + " first_treat='first_treat')\n", + "es_surface = results_es.aggregate('event_study')\n", "\n", "# Plot event study\n", "if HAS_MATPLOTLIB:\n", - " plot_event_study(results_es, title='Two-Stage DiD Event Study')\n", + " plot_event_study(es_surface, title='Two-Stage DiD Event Study')\n", "else:\n", " print(\"Install matplotlib to see visualizations: pip install matplotlib\")" ] @@ -101,8 +102,8 @@ "metadata": {}, "outputs": [], "source": [ - "# View event study effects as a table\n", - "results_es.to_dataframe(level='event_study')" + "# View event study effects as a table (unified container schema)\n", + "es_surface.to_dataframe()" ] }, { @@ -186,11 +187,10 @@ "metadata": {}, "outputs": [], "source": [ - "# Fit with group aggregation\n", + "# Aggregate per-cohort effects POST-FIT - no refit needed\n", "results_grp = TwoStageDiD().fit(data, outcome='outcome', unit='unit',\n", - " time='period', first_treat='first_treat',\n", - " aggregate='group')\n", - "results_grp.to_dataframe(level='group')" + " time='period', first_treat='first_treat')\n", + "results_grp.aggregate('group').to_dataframe()" ] }, { @@ -238,28 +238,29 @@ "metadata": {}, "outputs": [], "source": [ - "# Horizon-by-horizon SE comparison\n", + "# Horizon-by-horizon SE comparison (post-fit containers, unified schema)\n", "ts_es = TwoStageDiD().fit(data, outcome='outcome', unit='unit',\n", - " time='period', first_treat='first_treat',\n", - " aggregate='event_study')\n", + " time='period', first_treat='first_treat'\n", + " ).aggregate('event_study')\n", "imp_es = ImputationDiD().fit(data, outcome='outcome', unit='unit',\n", - " time='period', first_treat='first_treat',\n", - " aggregate='event_study')\n", + " time='period', first_treat='first_treat'\n", + " ).aggregate('event_study')\n", + "ts_df = ts_es.to_dataframe().set_index('event_time')\n", + "imp_df = imp_es.to_dataframe().set_index('event_time')\n", "\n", "print(\"Horizon-by-Horizon Comparison: GMM Sandwich vs Conservative Variance\")\n", "print(\"=\" * 70)\n", "print(f\"{'Horizon':>8} {'Effect':>10} {'GMM SE':>10} {'Cons. SE':>10} {'Ratio':>8}\")\n", "print(\"-\" * 70)\n", "\n", - "for h in sorted(ts_es.event_study_effects.keys()):\n", - " ts_eff = ts_es.event_study_effects[h]\n", - " imp_eff = imp_es.event_study_effects[h]\n", - " if ts_eff.get('n_obs', 0) == 0:\n", + "for h in ts_df.index:\n", + " row = ts_df.loc[h]\n", + " if row['is_reference'] or not np.isfinite(row['att']):\n", " print(f\"{h:>8} {'[ref]':>10} {'---':>10} {'---':>10} {'---':>8}\")\n", " continue\n", - " effect = ts_eff['effect']\n", - " gmm_se = ts_eff['se']\n", - " cons_se = imp_eff['se']\n", + " effect = row['att']\n", + " gmm_se = row['se']\n", + " cons_se = imp_df.loc[h, 'se']\n", " ratio = gmm_se / cons_se if cons_se > 0 else np.nan\n", " print(f\"{h:>8} {effect:>10.4f} {gmm_se:>10.4f} {cons_se:>10.4f} {ratio:>8.3f}\")" ] diff --git a/docs/tutorials/15_efficient_did.ipynb b/docs/tutorials/15_efficient_did.ipynb index 02f0c2c0..cdf75010 100644 --- a/docs/tutorials/15_efficient_did.ipynb +++ b/docs/tutorials/15_efficient_did.ipynb @@ -503,8 +503,8 @@ "outputs": [], "source": [ "# 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", + "# EfficientDiD, CallawaySantAnna, and ImputationDiD all aggregate\n", + "# POST-FIT (3.9 - rows M-023/M-020/M-021).\n", "edid_es_r = EfficientDiD(pt_assumption=\"all\").fit(\n", " data, outcome='outcome', unit='unit', time='period',\n", " first_treat='first_treat')\n", @@ -513,15 +513,15 @@ " 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", + " first_treat='first_treat')\n", "\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", + "# rename to the legacy native column names so the plotting block below\n", + "# stays unchanged across all 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", + "imp_es_df = imp_es_r.aggregate('event_study').to_dataframe().rename(columns=_ren)\n", "\n", "if HAS_MATPLOTLIB:\n", " fig, ax = plt.subplots(figsize=(10, 6))\n", diff --git a/docs/v4-deprecations.yaml b/docs/v4-deprecations.yaml index c2faec89..e6921203 100644 --- a/docs/v4-deprecations.yaml +++ b/docs/v4-deprecations.yaml @@ -235,11 +235,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/imputation.py] - notes: "balance_e moves too." + test_ref: tests/test_aggregate_contract.py + code_refs: [diff_diff/imputation.py, diff_diff/imputation_aggregation.py, diff_diff/imputation_results.py, diff_diff/imputation_bootstrap.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 (plain fit() never warns; supplying ANY value, None included, warns - CS-style joint warning with balance_e [M-118], warn-and-still-work; the imputation_did wrapper forwards the sentinel so a plain wrapper call never warns). NO fit-time value validation existed and none is added (unknown strings silently act like None; the post-fit successor fails closed via the mixin vocabulary - a behavior improvement). The successor is a PANEL-BACKED lazy recompute kit: ES/group aggregation is a target-specific Theorem-3 recompute from the working panel + untreated FE model (no compact influence payload can honor a different balance_e), so the kit's bookkeeping holds REFERENCES to the SAME per-fit objects self._fit_data already retains for pretrend_test() - ZERO marginal memory and unchanged pickles via memoization (the _estimator_ref field already ships the panel; enumeration in the REGISTRY ImputationDiD Note). Value snapshots (treatment_groups copy, config scalars, a dataclasses.replace copy of survey_metadata, and the survey_df_seed/survey_df_final df channels) isolate recompute and the ES carrier from public-field mutation; aggregate() reads NOTHING mutable off the results object except the deliberate overall-quintet relay and the bootstrap_results gate. SUPPORTED SUBSET simple|event_study|group; calendar/'all' fail closed via the mixin; weights= rejected. Bootstrap fits fail closed for ALL levels including 'simple' (uniform CS/EDiD parity; the per-target psi machinery makes replay tractable - a TODO row). Replicate-weight fits REPLAY the extracted _replicate_override_aggregates with a LEVEL-MATCHED stack ([overall, ES] or [overall, groups]): compute_replicate_refit_variance validates replicates jointly, so aggregate(L) reproduces fit(aggregate=L) exactly and a fit(aggregate='all') surface is NOT the equivalence target when a replicate NaNs on one family's targets; the same joint-stack coupling means migrating a replicate fit from fit(aggregate=) to plain fit changes the OVERALL row's se/CI/df on degenerate designs (documented migration delta, CHANGELOG + REGISTRY note (e) + a contract-test pin). pretrends=True + replicate: post-fit aggregate('event_study') raises the same NotImplementedError the fit-time gate raises (per-replicate lead refits unimplemented); group/simple still work. Simple relay: n = n_treated_obs with n_kind='obs' (the treated/control unit sets OVERLAP, so the CS/EDiD disjoint-units convention cannot apply - the StackedDiD carve-out class); df = the survey_df_final snapshot (what the stored overall inference received). Group rows record per-row df_used at each safe_inference call (additive row-dict key; the replicate override rewrites it, the bootstrap override clears it, the all-NaN cohort branch writes no key - consumers read via .get). The M-127 df_convention inert-config warning predicate is REVISED to reachability (pretrends AND not-replicate AND (deprecated fit-time ES/all supplied OR n_bootstrap <= 0)) because post-fit aggregate() made the old aggregate-keyed claim false; reachability-BASED, not exact - a fit whose bootstrap later fails (bootstrap_results=None) can still aggregate post-fit, so that corner warns spuriously (recorded on M-127 too). Container admission NOT widened: ImputationDiD is rejected BY DESIGN (no joint ES covariance - per-horizon conservative SEs only; see M-093). balance_e moves as its own row [M-118]." - id: M-022 kind: param group: aggregate-postfit @@ -248,11 +249,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/two_stage.py] - notes: "balance_e moves too." + test_ref: tests/test_aggregate_contract.py + code_refs: [diff_diff/two_stage.py, diff_diff/two_stage_aggregation.py, diff_diff/two_stage_results.py, diff_diff/two_stage_bootstrap.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 (plain fit() never warns; supplying ANY value, None included, warns - CS-style joint warning with balance_e [M-119], warn-and-still-work; the two_stage_did wrapper forwards the sentinel so a plain wrapper call never warns). NO fit-time value validation existed and none is added (unknown strings silently act like None; the post-fit successor fails closed via the mixin vocabulary). The successor is a PANEL-BACKED lazy recompute kit: each level is a fresh Stage-2 OLS + joint Gardner-GMM sandwich, so the kit retains a COLUMN-SUBSET COPY of the working frame (only the columns the moved methods read by name, deduplicated - cluster= may legally name the unit/time/first_treat column) plus the Stage-1 FE model, masks, and survey objects. MEMORY CONTRACT: this is the FIRST panel retention on TwoStageDiD results - O(n_obs) incl. unit/time/cluster identifier columns on every results object and pickle, and replicate designs additionally retain the (n_obs x R) replicate matrix via resolved_survey; the CS/EDiD identifier-minimization guarantee deliberately does NOT hold (a store_kit opt-out is a DEFERRED row). score_pad_mask/cluster_ids_full are stored as the Wave-E.3-GATED values fit actually passed. Value snapshots (treatment_groups copy, overall_att, survey_df_stage2/survey_df_final, a dataclasses.replace copy of survey_metadata) isolate recompute and the ES carrier from public-field mutation. SUPPORTED SUBSET simple|event_study|group; calendar/'all' fail closed; weights= rejected. Bootstrap fits fail closed for ALL levels including 'simple' (per-level GMM scores are function-locals; replay is a TODO row); a fit whose bootstrap FAILED (bootstrap_results=None, analytical inference retained) aggregates normally. Replicate-weight fits REPLAY the extracted _replay_replicate_inference with a LEVEL-MATCHED stack - aggregate(L) reproduces fit(aggregate=L) exactly; fit(aggregate='all') is NOT the equivalence target on degenerate designs, and the same joint-stack coupling makes the migration to plain fit change the OVERALL row's se/CI/df there (documented migration delta). Post-fit aggregate('event_study') reproduces the M-092 container contract exactly: analytical fits thread the recomputed joint vcov + vcov_index + the finite-and->0 df scalar through the carrier; replicate fits thread vcov=None/index=None with the REPLAYED level-matched df. Simple relay: n = n_treated_obs with n_kind='obs' (overlapping unit sets - StackedDiD carve-out class); df = the survey_df_final snapshot (on replicate fits that value came from the [overall]-only stack - snapshotted, never re-derived). Group relay df is a SCALAR broadcast (deliberate divergence from ImputationDiD's per-row df_used: _stage2_group passes one immutable survey_df to every row, so the scalar is provenance-exact by construction and the moved method stays verbatim). Container admission NOT widened: DEFERRED pending a normalization derivation - analytical surfaces carry the real joint Gardner-GMM covariance (M-092), but pre-period coefficients are stage-1 residual means, not reference-normalized contrasts, while HonestDiD's Delta arithmetic hard-codes delta_0=0 (see M-093 + the DEFERRED.md paper-gated row). balance_e moves as its own row [M-119]." - id: M-023 kind: param group: aggregate-postfit @@ -806,7 +808,7 @@ rows: status: planned phase: 2 warning: FutureWarning - code_refs: [diff_diff/imputation.py, diff_diff/__init__.py, diff_diff/guides/llms-full.txt] + code_refs: [diff_diff/imputation.py, diff_diff/__init__.py, diff_diff/guides/llms-full.txt, docs/methodology/REGISTRY.md] notes: "Use ImputationDiD. Only 8 of 24 estimators ever had wrappers; the surface is retired rather than completed." - id: M-071 kind: function @@ -818,7 +820,7 @@ rows: status: planned phase: 2 warning: FutureWarning - code_refs: [diff_diff/two_stage.py, diff_diff/__init__.py, diff_diff/guides/llms-full.txt] + code_refs: [diff_diff/two_stage.py, diff_diff/__init__.py, diff_diff/guides/llms-full.txt, docs/methodology/REGISTRY.md] notes: "Use TwoStageDiD." - id: M-072 kind: function @@ -986,7 +988,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). 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." + 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. The qualifier EXTENDS to ImputationDiD (2(b) PR-3b, with M-021): its newly public container likewise exposes only the scalar df_survey channel (no per-row df source; all-NaN per-row column, identical to its fit-time surface) - the same TODO.md completion row names it." - id: M-093 kind: behavior group: results-contract @@ -997,7 +999,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. 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." + 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. FOURTH PRE-CUT AMENDMENT (2(b) PR-3b, with M-021/M-022): admission NOT widened by the Imputation/TwoStage migrations - ImputationDiD rejected BY DESIGN (no joint ES covariance; per-horizon Theorem-3 conservative SEs only); TwoStageDiD DEFERRED despite carrying the joint Gardner-GMM covariance on analytical fits (M-092), because its pre-period coefficients are stage-1 residual means (the reference horizon is dropped from the no-intercept Stage-2 design; the zero anchor row is appended mechanically), not contrasts against the advertised reference, while HonestDiD's Delta^RM/Delta^SD arithmetic hard-codes the delta_0=0 normalization into its boundary/bridge constraints - admission awaits a normalization derivation (DEFERRED.md paper-gated row); both terminal TypeErrors state the per-producer grounds. Source scoping continues to reject all other producers. 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 @@ -1360,10 +1362,11 @@ rows: # ---- Phase 2b PR 1: the post-fit aggregation surface --------------------- # 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. + # slot matching the M-021/M-022/M-023 ordering), and M-118/M-119 were + # claimed by the Imputation/TwoStage balance_e rows (2b PR-3b, the first + # and second slots). Ids M-116 and M-121 remain intentionally UNUSED + # (HAD rename, Wooldridge) - reserved rather than reassigned, since ids + # are never reused. - id: M-117 kind: param group: aggregate-postfit @@ -1392,6 +1395,34 @@ rows: 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-118 + kind: param + group: aggregate-postfit + old: "diff_diff:ImputationDiD.fit[balance_e]" + new: "diff_diff:ImputationDiDResults.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/imputation.py, diff_diff/imputation_aggregation.py, diff_diff/imputation_results.py, diff_diff/aggregation.py] + notes: "balance_e moves from fit() onto aggregate() with [M-021] (joint FutureWarning; the M-117/M-120 twin; the imputation_did wrapper forwards the sentinel). Applies to event-study aggregation only (the mixin default vocabulary), so aggregate(type='simple'|'group', balance_e=...) raises. ImputationDiD's balance rule is the BALANCED-WINDOW rule - a cohort is retained iff its observed relative-time set (checked against the FULL panel via _build_cohort_rel_times) covers the contiguous window [-balance_e, max_h] (imputation_aggregation.py _compute_balanced_cohort_mask) - divergent from CS/EfficientDiD's anchor-horizon rule ([M-117]/[M-120] document theirs; TwoStageDiD shares the balanced-window rule, [M-119]). A window no cohort satisfies warns (UserWarning) and yields the reference-marker-only dict -> a legal near-empty container." + - id: M-119 + kind: param + group: aggregate-postfit + old: "diff_diff:TwoStageDiD.fit[balance_e]" + new: "diff_diff:TwoStageDiDResults.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/two_stage.py, diff_diff/two_stage_aggregation.py, diff_diff/two_stage_results.py, diff_diff/aggregation.py] + notes: "balance_e moves from fit() onto aggregate() with [M-022] (joint FutureWarning; the M-117/M-120 twin; the two_stage_did wrapper forwards the sentinel). Applies to event-study aggregation only (the mixin default vocabulary), so aggregate(type='simple'|'group', balance_e=...) raises. TwoStageDiD's balance rule is the BALANCED-WINDOW rule - a cohort is retained iff its observed relative-time set (checked against the full panel via _build_cohort_rel_times) covers the contiguous window [-balance_e, max_h] (two_stage_aggregation.py _stage2_event_study) - the ImputationDiD rule ([M-118]), divergent from CS/EfficientDiD's anchor-horizon rule. Zero qualifying cohorts warns ('No cohorts satisfy balance_e=...') and yields the reference-row-only dict with vcov=None -> a legal near-empty container." - id: M-126 kind: behavior group: variance-cr1-convergence @@ -1469,7 +1500,7 @@ rows: phase: 2 test_ref: tests/test_variance_conventions.py code_refs: [diff_diff/utils.py, diff_diff/linalg.py, diff_diff/estimators.py, diff_diff/twfe.py, diff_diff/results.py, diff_diff/sun_abraham.py, diff_diff/wooldridge.py, diff_diff/wooldridge_results.py, diff_diff/stacked_did.py, diff_diff/stacked_did_results.py, diff_diff/imputation.py, diff_diff/imputation_results.py, diff_diff/lpdid.py, diff_diff/lpdid_results.py, docs/methodology/REGISTRY.md, docs/methodology/variance-conventions.md] - notes: "Tail-df consolidation (D4 + the normal-theory tail-df defect family of the 3.9 variance-consolidation program, PR C). df_convention becomes a THREE-VALUE knob {'residual','cluster','normal'} on DiD/MPD/TWFE/LinearRegression (previously two-value; existing values unchanged, M-004..M-006 old_default untouched; 'normal' is deliberate z at the fallback level on EVERY fit, unclustered and conley included) and is EXTENDED to SunAbraham / WooldridgeDiD(OLS arms) / StackedDiD / ImputationDiD(pretrends lead regression) / LPDiD via the new utils.resolve_tail_df fallback resolver (survey/replicate df and hc2_bm Bell-McCaffrey per-coefficient/contrast DOF keep precedence under every value; 'cluster' fails closed warn+NaN at G<=1; unclustered/conley fits are inert for 'cluster'; a non-positive residual df falls back to normal theory with a warning). DEFECT FIX shipped as the new 'residual' default: WooldridgeDiD default-hc1 (within AND the hc1+cohort_trends full-dummy arm), StackedDiD, SunAbraham aggregates (D4: one fit reported residual-t cells and normal-theory aggregates - aggregates now reuse the saturated regression's df_ so cells and aggregates share ONE df source), and ImputationDiD pretrends leads previously passed df=None (silent normal-theory z) on clustered analytical paths; at |t|=2 z understates the t(G-1) p-value by 24.2% at G=20. Under the residual default the numeric movement is z vs t(large residual df) - small but honest; the G-1 convergence with fixest/Stata is the 4.0 flip. 'normal' reproduces the pre-3.9 numbers on the fixed surfaces' affected lanes (on lanes already at t - unclustered DiD-family fits, Wooldridge classical/hc2, SA cells - it is a NEW deliberate z option, not a compatibility path). LPDiD keeps its Stata-anchored t(G-1) as default 'cluster' (bit-identical; already at the 4.0 target, no flip row; degenerate lanes - unclustered refit, RA G<=1, saturated early return - keep literal df=None under all values; RA-path 'residual' = n_total - k0_kept - 1, the pooled M-estimator parameter count, a library convention with no external anchor since Stata teffects ra reports z). Wooldridge logit/poisson arms (survey df or normal theory, knob-independent, explicit-non-default warns) and all IF-based estimators (CallawaySantAnna, TwoStageDiD, dCDH, EfficientDiD, ContinuousDiD, ChangesInChanges, SyntheticDiD, HAD - L3) plus OLS-based TripleDifference/TROP are out of the knob's scope. Wild-cluster-bootstrap p/CIs are test-inversion based and invariant. Provenance: df_convention echoed on the five standalone results containers; SunAbraham/StackedDiD gain inference_df (the overall-ATT df actually used; None under bootstrap); WooldridgeDiDResults._df_one_way renamed _df_analytic_fallback (now the resolved survey-first fallback on ALL OLS arms, live on bootstrap fits so post-fit aggregate() reproduces fit-time inference) with __setstate__ pickle migration. Defaults for the four new residual surfaces flip to 'cluster' at 4.0 (M-128..M-131, same policy as M-004..M-006). introduced_in gates the 3.9 cut; deprecated_in stays null so the early-flip guard does not fire against the PR that ships it." + notes: "Tail-df consolidation (D4 + the normal-theory tail-df defect family of the 3.9 variance-consolidation program, PR C). df_convention becomes a THREE-VALUE knob {'residual','cluster','normal'} on DiD/MPD/TWFE/LinearRegression (previously two-value; existing values unchanged, M-004..M-006 old_default untouched; 'normal' is deliberate z at the fallback level on EVERY fit, unclustered and conley included) and is EXTENDED to SunAbraham / WooldridgeDiD(OLS arms) / StackedDiD / ImputationDiD(pretrends lead regression) / LPDiD via the new utils.resolve_tail_df fallback resolver (survey/replicate df and hc2_bm Bell-McCaffrey per-coefficient/contrast DOF keep precedence under every value; 'cluster' fails closed warn+NaN at G<=1; unclustered/conley fits are inert for 'cluster'; a non-positive residual df falls back to normal theory with a warning). DEFECT FIX shipped as the new 'residual' default: WooldridgeDiD default-hc1 (within AND the hc1+cohort_trends full-dummy arm), StackedDiD, SunAbraham aggregates (D4: one fit reported residual-t cells and normal-theory aggregates - aggregates now reuse the saturated regression's df_ so cells and aggregates share ONE df source), and ImputationDiD pretrends leads previously passed df=None (silent normal-theory z) on clustered analytical paths; at |t|=2 z understates the t(G-1) p-value by 24.2% at G=20. Under the residual default the numeric movement is z vs t(large residual df) - small but honest; the G-1 convergence with fixest/Stata is the 4.0 flip. 'normal' reproduces the pre-3.9 numbers on the fixed surfaces' affected lanes (on lanes already at t - unclustered DiD-family fits, Wooldridge classical/hc2, SA cells - it is a NEW deliberate z option, not a compatibility path). LPDiD keeps its Stata-anchored t(G-1) as default 'cluster' (bit-identical; already at the 4.0 target, no flip row; degenerate lanes - unclustered refit, RA G<=1, saturated early return - keep literal df=None under all values; RA-path 'residual' = n_total - k0_kept - 1, the pooled M-estimator parameter count, a library convention with no external anchor since Stata teffects ra reports z). Wooldridge logit/poisson arms (survey df or normal theory, knob-independent, explicit-non-default warns) and all IF-based estimators (CallawaySantAnna, TwoStageDiD, dCDH, EfficientDiD, ContinuousDiD, ChangesInChanges, SyntheticDiD, HAD - L3) plus OLS-based TripleDifference/TROP are out of the knob's scope. Wild-cluster-bootstrap p/CIs are test-inversion based and invariant. Provenance: df_convention echoed on the five standalone results containers; SunAbraham/StackedDiD gain inference_df (the overall-ATT df actually used; None under bootstrap); WooldridgeDiDResults._df_one_way renamed _df_analytic_fallback (now the resolved survey-first fallback on ALL OLS arms, live on bootstrap fits so post-fit aggregate() reproduces fit-time inference) with __setstate__ pickle migration. Defaults for the four new residual surfaces flip to 'cluster' at 4.0 (M-128..M-131, same policy as M-004..M-006). introduced_in gates the 3.9 cut; deprecated_in stays null so the early-flip guard does not fire against the PR that ships it. AMENDED (2(b) PR-3b, with M-021): the ImputationDiD inert-config warning predicate is REVISED from aggregate-keyed to REACHABILITY-based - warn iff df_convention is explicitly non-default AND the per-lead inference is unreachable for this fit config (reachable iff pretrends AND not replicate-weight AND (the deprecated fit-time aggregate='event_study'/'all' was supplied OR n_bootstrap <= 0)), because post-fit results.aggregate('event_study') made the old 'aggregate outside {event_study,all} is inert' claim false on analytical pretrends fits. Reachability-BASED, not exact: a fit whose bootstrap later FAILS (bootstrap_results=None) can still aggregate post-fit, so that corner warns spuriously - the warning fires before the bootstrap runs and cannot know." - id: M-128 kind: default-flip group: df-convention-flip diff --git a/docs/v4-design.md b/docs/v4-design.md index cea81717..a67c874b 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, 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 | +| 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-118, M-119, 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` @@ -869,12 +869,14 @@ forever - a removed symbol resurrecting is a test failure. class/function rows and alias rows also assert `__all__` membership consistent with their status (stale `import *` entries fail). The shipped row ids are a - committed snapshot in the enforcement test (114 as of the alias-diet - family: Phase 1 + the diagnostic-family amendment + + committed snapshot in the enforcement test (120 as of 2(b) PR-3b's + Imputation/TwoStage balance_e rows: Phase 1 + the diagnostic-family + amendment + the M-092/M-093 results-contract rows + the M-094..M-096 amendment rows + - the M-097..M-115 completeness sweep + M-117/M-122 + the ETWFE + the M-097..M-115 completeness sweep + M-117..M-120/M-122 + the ETWFE reference-period pair M-123/M-124 + M-125 + M-126 + M-127..M-131 + - the alias-diet family M-132..M-135; + the alias-diet family M-132..M-135 + the 2(c)-ii amendments + M-136..M-138; the snapshot extends by a new id range in the same diff that appends rows): ids are never deleted or reused, and the test fails if any snapshot id disappears. diff --git a/tests/test_aggregate_contract.py b/tests/test_aggregate_contract.py index 06eb7d11..c4893717 100644 --- a/tests/test_aggregate_contract.py +++ b/tests/test_aggregate_contract.py @@ -1,7 +1,10 @@ """Behavioral contract for post-fit ``results.aggregate()`` (spec section 6). -The ``test_ref`` for ledger rows M-020 (``fit(aggregate=)`` shim), M-117 -(``balance_e`` moves onto ``aggregate()``) and M-122 (``AggregationResult``). +The ``test_ref`` for every aggregate-postfit ledger row this file pins: +M-020/M-023/M-021/M-022 (the CS / EfficientDiD / Imputation / TwoStage +``fit(aggregate=)`` shims), M-024/M-026 (the Stacked / dCDH shims + view +relays), M-117/M-120/M-118/M-119 (``balance_e`` moves onto ``aggregate()``) +and M-122 (``AggregationResult``). The headline gate is NUMERICAL INERTNESS: for every supported type, ``fit(aggregate=T)`` and ``fit(); .aggregate(T)`` must agree to 1e-14. The @@ -1111,6 +1114,8 @@ def test_mixin_hooks_are_not_dataclass_fields(self): assert "ChaisemartinDHaultfoeuilleResults" in checked assert "StackedDiDResults" in checked assert "EfficientDiDResults" in checked + assert "ImputationDiDResults" in checked + assert "TwoStageDiDResults" in checked # --------------------------------------------------------------------------- # @@ -1945,3 +1950,872 @@ def test_hausman_pretest_emits_no_future_warning(self, efficient_panel): first_treat="first_treat", ) assert [w for w in caught if issubclass(w.category, FutureWarning)] == [] + + +# --------------------------------------------------------------------------- # +# ImputationDiD (rows M-021/M-118): fit(aggregate=/balance_e=) shim + the +# PANEL-BACKED recompute aggregate() (kit refs = the _fit_data objects) +# --------------------------------------------------------------------------- # + +IMPUTATION_KW = dict(outcome="outcome", unit="unit", time="period", first_treat="first_treat") + + +def _imputation_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 _imputation_clustered_panel(seed=11): + d = _imputation_panel(seed=seed).copy() + d["cl"] = (d["unit"] // 3).astype(int) + return d + + +def _imputation_survey_panel(seed=5, replicate=False, degenerate=None): + """Unit-constant pweights; optionally a JK replicate design. + + degenerate: None | "dropped" (2 all-zero replicate columns) | + "undefined" (all but one column all-zero -> n_valid <= 1) | + "cohort_zero" (replicate rw0 zeroes EVERY row of one cohort -> that + replicate NaNs the cohort's GROUP target while overall stays finite, + so the joint [overall, groups] stack drops a replicate the + [overall]-only stack keeps - the deterministic overall-row + migration-delta shape). + """ + d = _imputation_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 + cohort4_units = set(d.loc[d["first_treat"] == 4, "unit"].unique()) + 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.1, 2.0) for u in d["unit"].unique()} + d[col] = d["unit"].map(jitter) * d["w"] + if degenerate == "cohort_zero" and r == 0: + d.loc[d["unit"].isin(cohort4_units), col] = 0.0 + return d, rep_cols + + +def _imputation_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_imputation(data, *, est_kw=None, **fit_kw): + from diff_diff import ImputationDiD + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return ImputationDiD(**(est_kw or {})).fit(data, **IMPUTATION_KW, **fit_kw) + + +def _assert_es_container_matches_fit_time(container, fit_time_results, name=""): + """Post-fit EventStudyResults must equal the fit-time-built surface + column-for-column at 1e-14 (incl. reference marking and df provenance).""" + from diff_diff.results_base import build_event_study_surface + + ref = build_event_study_surface(fit_time_results) + da, db = container.to_dataframe(), ref.to_dataframe() + assert list(da.columns) == list(db.columns) + for c in da.columns: + x, y = da[c].to_numpy(), db[c].to_numpy() + if x.dtype.kind in "fc": + np.testing.assert_allclose( + x, y, rtol=0, atol=1e-14, equal_nan=True, err_msg=f"{name}/{c}" + ) + else: + assert (x == y).all(), (name, c) + for attr in ("df_survey", "anticipation", "alpha"): + va, vb = getattr(container, attr), getattr(ref, attr) + same = va == vb or (va is None and vb is None) + try: + same = same or (np.isnan(va) and np.isnan(vb)) + except TypeError: + pass + assert same, (name, attr, va, vb) + + +def _assert_group_matches_fit_time(agg, fit_time_results, name=""): + for i, g in enumerate(agg.label): + row = fit_time_results.group_effects[g] + for field_, key in ( + ("att", "effect"), + ("se", "se"), + ("t_stat", "t_stat"), + ("p_value", "p_value"), + ): + np.testing.assert_allclose( + getattr(agg, field_)[i], + row[key], + rtol=0, + atol=1e-14, + equal_nan=True, + err_msg=f"{name}/{g}/{field_}", + ) + + +@pytest.fixture(scope="module") +def imputation_panel(): + return _imputation_panel() + + +@pytest.fixture(scope="module") +def imputation_fitted(imputation_panel): + """Plain fit - the kit (refs to _fit_data) powers aggregate().""" + return _fit_imputation(imputation_panel) + + +@pytest.fixture(scope="module") +def imputation_fit_time(imputation_panel): + """Deprecated fit-time aggregate="all" - the analytical inertness reference.""" + return _fit_imputation(imputation_panel, aggregate="all") + + +class TestImputationShim: + def test_plain_fit_does_not_warn(self, imputation_panel): + from diff_diff import ImputationDiD + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + ImputationDiD().fit(imputation_panel, **IMPUTATION_KW) + assert [w for w in caught if issubclass(w.category, FutureWarning)] == [] + + def test_aggregate_kwarg_warns_even_at_none(self, imputation_panel): + from diff_diff import ImputationDiD + + with pytest.warns(FutureWarning, match=r"ImputationDiD\.fit\(aggregate=\)"): + ImputationDiD().fit(imputation_panel, **IMPUTATION_KW, aggregate=None) + + def test_balance_e_kwarg_warns_alone(self, imputation_panel): + from diff_diff import ImputationDiD + + with pytest.warns(FutureWarning, match=r"ImputationDiD\.fit\(balance_e=\)"): + ImputationDiD().fit(imputation_panel, **IMPUTATION_KW, balance_e=None) + + def test_joint_supply_warns_once_naming_both(self, imputation_panel): + from diff_diff import ImputationDiD + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + ImputationDiD().fit( + imputation_panel, **IMPUTATION_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, imputation_panel): + res = _fit_imputation(imputation_panel, aggregate="nonsense") + plain = _fit_imputation(imputation_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, imputation_fitted, imputation_fit_time): + assert imputation_fit_time.event_study_effects is not None + assert imputation_fit_time.group_effects is not None + es = imputation_fitted.aggregate("event_study") + for e, row in imputation_fit_time.event_study_effects.items(): + i = list(es.event_time).index(e) + if np.isfinite(row["effect"]) and not es.is_reference[i]: + np.testing.assert_allclose(row["effect"], es.att[i], rtol=1e-14) + + def test_wrapper_forwarded_aggregate_warns(self, imputation_panel): + from diff_diff import imputation_did + + with pytest.warns(FutureWarning, match=r"ImputationDiD\.fit\(aggregate=\)"): + imputation_did( + imputation_panel, + "outcome", + "unit", + "period", + "first_treat", + aggregate="event_study", + ) + + def test_wrapper_forwarded_balance_e_warns(self, imputation_panel): + from diff_diff import imputation_did + + with pytest.warns(FutureWarning, match=r"ImputationDiD\.fit\(balance_e=\)"): + imputation_did( + imputation_panel, "outcome", "unit", "period", "first_treat", balance_e=1 + ) + + def test_plain_wrapper_call_does_not_warn(self, imputation_panel): + from diff_diff import imputation_did + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + imputation_did(imputation_panel, "outcome", "unit", "period", "first_treat") + assert [w for w in caught if issubclass(w.category, FutureWarning)] == [] + + +class TestImputationAggregate: + def test_event_study_inert(self, imputation_fitted, imputation_fit_time): + es = imputation_fitted.aggregate("event_study") + _assert_es_container_matches_fit_time(es, imputation_fit_time, "imp/es") + + @pytest.mark.parametrize("balance_e", [0, 1, 2]) + def test_balance_e_inert(self, imputation_panel, imputation_fitted, balance_e): + ref = _fit_imputation(imputation_panel, aggregate="event_study", balance_e=balance_e) + es = imputation_fitted.aggregate("event_study", balance_e=balance_e) + _assert_es_container_matches_fit_time(es, ref, f"imp/es/be{balance_e}") + + def test_group_inert(self, imputation_fitted, imputation_fit_time): + _assert_group_matches_fit_time( + imputation_fitted.aggregate("group"), imputation_fit_time, "imp/gr" + ) + + def test_simple_relay_bit_exact(self, imputation_fitted): + sm = imputation_fitted.aggregate("simple") + assert sm.att[0] == imputation_fitted.overall_att + assert sm.se[0] == imputation_fitted.overall_se + assert sm.t_stat[0] == imputation_fitted.overall_t_stat + assert sm.p_value[0] == imputation_fitted.overall_p_value + assert sm.n_kind == "obs" + assert sm.n[0] == imputation_fitted.n_treated_obs + assert np.isnan(sm.to_dataframe()["df"].to_numpy()).all() # non-survey fit + + def test_cluster_arm_inert(self): + d = _imputation_clustered_panel() + ref = _fit_imputation(d, est_kw={"cluster": "cl"}, aggregate="all") + plain = _fit_imputation(d, est_kw={"cluster": "cl"}) + _assert_es_container_matches_fit_time(plain.aggregate("event_study"), ref, "imp/cluster") + _assert_group_matches_fit_time(plain.aggregate("group"), ref, "imp/cluster") + + def test_covariate_arm_inert(self): + # Time-varying covariate: the kit must carry delta_hat / + # kept_cov_mask / the covariate columns; an omission escapes every + # covariate-free arm (local-review P2). + d = _imputation_panel(seed=23).copy() + rng = np.random.default_rng(23) + d["x1"] = rng.normal(size=len(d)) + 0.1 * d["period"] + ref = _fit_imputation(d, aggregate="all", covariates=["x1"]) + plain = _fit_imputation(d, covariates=["x1"]) + _assert_es_container_matches_fit_time(plain.aggregate("event_study"), ref, "imp/cov") + _assert_group_matches_fit_time(plain.aggregate("group"), ref, "imp/cov") + + def test_survey_tsl_arm_inert(self): + d, _ = _imputation_survey_panel() + sd = _imputation_survey_design() + ref = _fit_imputation(d, aggregate="all", survey_design=sd) + plain = _fit_imputation(d, survey_design=sd) + _assert_es_container_matches_fit_time(plain.aggregate("event_study"), ref, "imp/tsl") + _assert_group_matches_fit_time(plain.aggregate("group"), ref, "imp/tsl") + sm = plain.aggregate("simple") + df_col = sm.to_dataframe()["df"].to_numpy() + assert np.isfinite(df_col).all() # survey df threads the simple row + + @pytest.mark.parametrize("degenerate", [None, "dropped", "undefined"]) + def test_replicate_arms_level_matched(self, degenerate): + # LEVEL-MATCHED references: aggregate(L) reproduces fit(aggregate=L), + # not fit(aggregate='all') (the joint replicate stack couples rows). + d, rep_cols = _imputation_survey_panel(replicate=True, degenerate=degenerate) + sd = _imputation_survey_design(rep_cols) + ref_es = _fit_imputation(d, aggregate="event_study", survey_design=sd) + ref_gr = _fit_imputation(d, aggregate="group", survey_design=sd) + plain = _fit_imputation(d, survey_design=sd) + _assert_es_container_matches_fit_time( + plain.aggregate("event_study"), ref_es, f"imp/rep/{degenerate}" + ) + _assert_group_matches_fit_time(plain.aggregate("group"), ref_gr, f"imp/rep/{degenerate}") + sm = plain.aggregate("simple") + np.testing.assert_allclose(sm.se[0], plain.overall_se, rtol=0, atol=0, equal_nan=True) + + def test_replicate_undefined_df_sentinel(self): + # n_valid <= 1: the working df degenerates; the ES container's + # df_survey resolves to the 0.0 replicate-undefined sentinel on + # BOTH routes (the metadata-copy discriminator). + d, rep_cols = _imputation_survey_panel(replicate=True, degenerate="undefined") + sd = _imputation_survey_design(rep_cols) + plain = _fit_imputation(d, survey_design=sd) + es = plain.aggregate("event_study") + ref = _fit_imputation(d, aggregate="event_study", survey_design=sd) + from diff_diff.results_base import build_event_study_surface + + assert es.df_survey == build_event_study_surface(ref).df_survey == 0.0 + + def test_replicate_overall_row_migration_delta(self): + # The documented migration consequence: the deprecated fit(aggregate=) + # coupled the OVERALL row to the joint replicate stack. On the + # cohort-zero design (one replicate NaNs one cohort's group target + # while overall stays finite) the joint [overall, groups] stack drops + # a replicate the [overall]-only stack keeps -> plain-fit overall_se + # differs from fit(aggregate='group') overall_se; each surface is + # self-consistent, and post-fit aggregate('group') level-matches the + # deprecated fit(aggregate='group') rows exactly. + d, rep_cols = _imputation_survey_panel(replicate=True, degenerate="cohort_zero") + sd = _imputation_survey_design(rep_cols) + plain = _fit_imputation(d, survey_design=sd) + ref_gr = _fit_imputation(d, aggregate="group", survey_design=sd) + assert plain.overall_se != ref_gr.overall_se + _assert_group_matches_fit_time(plain.aggregate("group"), ref_gr, "imp/delta") + # healthy design: no coupling -> equality + dh, rep_h = _imputation_survey_panel(replicate=True) + sdh = _imputation_survey_design(rep_h) + # Healthy design: no replicate-drop coupling, so the two stacks + # agree - but NOT bit-identically: the [overall] vs + # [overall, groups] layouts route the replicate-variance matmul + # through different BLAS kernel shapes (~1 ULP reassociation; + # bit-identical on Accelerate, not on OpenBLAS-ARM/Windows). + # A REAL coupling delta is O(se) itself, far above this band. + np.testing.assert_allclose( + _fit_imputation(dh, survey_design=sdh).overall_se, + _fit_imputation(dh, aggregate="group", survey_design=sdh).overall_se, + rtol=1e-12, + ) + + def test_pretrends_arm_inert(self, imputation_panel): + ref = _fit_imputation(imputation_panel, est_kw={"pretrends": True}, aggregate="event_study") + plain = _fit_imputation(imputation_panel, est_kw={"pretrends": True}) + _assert_es_container_matches_fit_time(plain.aggregate("event_study"), ref, "imp/pretrends") + + def test_isolation_under_public_field_mutation(self, imputation_panel): + res = _fit_imputation(imputation_panel) + base_es = res.aggregate("event_study").to_dataframe() + base_gr = res.aggregate("group").to_dataframe() + base_sm = res.aggregate("simple").to_dataframe() + res.groups.pop() + res.time_periods.pop() + object.__setattr__(res, "alpha", 0.5) + object.__setattr__(res, "anticipation", 3) + # Imputation-only public config fields the kit snapshots: + object.__setattr__(res, "leave_one_out", True) + object.__setattr__(res, "df_convention", "cluster") + pd.testing.assert_frame_equal(res.aggregate("event_study").to_dataframe(), base_es) + pd.testing.assert_frame_equal(res.aggregate("group").to_dataframe(), base_gr) + pd.testing.assert_frame_equal(res.aggregate("simple").to_dataframe(), base_sm) + + def test_metadata_never_mutated_and_isolated(self): + d, rep_cols = _imputation_survey_panel(replicate=True) + sd = _imputation_survey_design(rep_cols) + res = _fit_imputation(d, survey_design=sd) + base = res.aggregate("event_study").to_dataframe() + base_df_survey = res.aggregate("event_study").df_survey + res.survey_metadata.df_survey = 999.0 + res.survey_metadata.replicate_method = None # the 0.0-sentinel discriminator + after = res.aggregate("event_study") + pd.testing.assert_frame_equal(after.to_dataframe(), base) + assert after.df_survey == base_df_survey + assert res.survey_metadata.df_survey == 999.0 # aggregate() never writes back + + def test_repeated_and_order_independent(self, imputation_fitted): + a = imputation_fitted.aggregate("group").to_dataframe() + imputation_fitted.aggregate("event_study") + imputation_fitted.aggregate("simple") + b = imputation_fitted.aggregate("group").to_dataframe() + pd.testing.assert_frame_equal(a, b) + + def test_bootstrap_fails_closed_all_levels(self, imputation_panel): + res = _fit_imputation(imputation_panel, est_kw={"n_bootstrap": 19, "seed": 1}) + for level in ("simple", "event_study", "group"): + with pytest.raises(NotImplementedError, match="bootstrap"): + res.aggregate(level) + + def test_pretrends_replicate_es_fails_closed(self): + d, rep_cols = _imputation_survey_panel(replicate=True) + sd = _imputation_survey_design(rep_cols) + res = _fit_imputation(d, est_kw={"pretrends": True}, survey_design=sd) + with pytest.raises(NotImplementedError, match="per-replicate"): + res.aggregate("event_study") + res.aggregate("group") + res.aggregate("simple") + + def test_fail_closed_vocabulary(self, imputation_fitted): + with pytest.raises(ValueError, match="Unsupported aggregation type"): + imputation_fitted.aggregate("calendar") + with pytest.raises(ValueError, match="Unsupported aggregation type"): + imputation_fitted.aggregate("all") + with pytest.raises(ValueError, match="Unsupported aggregation type"): + imputation_fitted.aggregate("nonsense") + with pytest.raises(ValueError, match="balance_e"): + imputation_fitted.aggregate("group", balance_e=1) + with pytest.raises(ValueError, match="balance_e"): + imputation_fitted.aggregate("simple", balance_e=1) + with pytest.raises(ValueError): + imputation_fitted.aggregate("group", weights="cell") + + def test_legacy_pickle_without_kit(self, imputation_fitted): + import copy + + legacy = copy.copy(imputation_fitted) + object.__setattr__(legacy, "_aggregation_kit", None) + with pytest.raises(ValueError, match="aggregation kit"): + legacy.aggregate("group") + + def test_pickle_roundtrip(self, imputation_panel): + import pickle + + res = _fit_imputation(imputation_panel) + clone = pickle.loads(pickle.dumps(res)) + pd.testing.assert_frame_equal( + clone.aggregate("group").to_dataframe(), res.aggregate("group").to_dataframe() + ) + + def test_group_df_used_relay(self): + # Survey fit: every group row's df equals the survey df its + # safe_inference received (per-row df_used capture). + d, _ = _imputation_survey_panel() + sd = _imputation_survey_design() + plain = _fit_imputation(d, survey_design=sd) + gr = plain.aggregate("group") + df_col = gr.to_dataframe()["df"].to_numpy() + finite_p = np.isfinite(gr.p_value) + assert np.isfinite(df_col[finite_p]).all() + # Plain fit: normal theory -> all-NaN df column. + plain2 = _fit_imputation(_imputation_panel()) + assert np.isnan(plain2.aggregate("group").to_dataframe()["df"].to_numpy()).all() + + def test_bootstrap_group_rows_clear_df_used(self, imputation_panel): + # Fit-time bootstrap override must never publish an analytical df + # beside percentile inference (public group_effects row dicts). + res = _fit_imputation( + imputation_panel, est_kw={"n_bootstrap": 19, "seed": 1}, aggregate="group" + ) + for row in res.group_effects.values(): + if np.isfinite(row["effect"]): + assert row.get("df_used") is None + + def test_empty_balance_window(self, imputation_fitted): + with pytest.warns(UserWarning, match="no horizons"): + es = imputation_fitted.aggregate("event_study", balance_e=100) + df_ = es.to_dataframe() + assert (df_["is_reference"] | ~np.isfinite(df_["att"])).all() + + +# --------------------------------------------------------------------------- # +# TwoStageDiD (rows M-022/M-119): fit(aggregate=/balance_e=) shim + the +# PANEL-BACKED recompute aggregate() (column-subset working-frame kit) +# --------------------------------------------------------------------------- # + +TWOSTAGE_KW = dict(outcome="outcome", unit="unit", time="period", first_treat="first_treat") + + +def _twostage_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 _twostage_prop5_panel(seed=13): + """No never-treated units, multiple cohorts -> Proposition-5 NaN rows.""" + d = _twostage_panel(seed=seed) + return d[d["first_treat"] > 0].copy() + + +def _twostage_survey_panel(seed=5, replicate=False, degenerate=None): + d = _twostage_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 + cohort4_units = set(d.loc[d["first_treat"] == 4, "unit"].unique()) + 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.1, 2.0) for u in d["unit"].unique()} + d[col] = d["unit"].map(jitter) * d["w"] + if degenerate == "cohort_zero" and r == 0: + d.loc[d["unit"].isin(cohort4_units), col] = 0.0 + return d, rep_cols + + +def _twostage_always_treated_panel(seed=7): + """Adds always-treated units so the survey Wave-E.3 pad activates + (the only shape where the kit's score_pad_mask/cluster_ids_full + snapshots are non-None).""" + d, _ = _twostage_survey_panel(seed=seed) + rng = np.random.default_rng(seed) + extra = [] + base_unit = int(d["unit"].max()) + 1 + for k in range(8): + u = base_unit + k + for t in sorted(d["period"].unique()): + extra.append((u, t, rng.normal() + 2.0, 1, 1.0)) + extra_df = pd.DataFrame(extra, columns=["unit", "period", "outcome", "first_treat", "w"]) + return pd.concat([d, extra_df], ignore_index=True) + + +def _twostage_survey_design(rep_cols=None, psu=False): + from diff_diff import SurveyDesign + + if rep_cols: + return SurveyDesign(weights="w", replicate_weights=rep_cols, replicate_method="JK1") + if psu: + return SurveyDesign(weights="w", psu="unit") + return SurveyDesign(weights="w") + + +def _fit_twostage(data, *, est_kw=None, **fit_kw): + from diff_diff import TwoStageDiD + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return TwoStageDiD(**(est_kw or {})).fit(data, **TWOSTAGE_KW, **fit_kw) + + +@pytest.fixture(scope="module") +def twostage_panel(): + return _twostage_panel() + + +@pytest.fixture(scope="module") +def twostage_fitted(twostage_panel): + return _fit_twostage(twostage_panel) + + +@pytest.fixture(scope="module") +def twostage_fit_time(twostage_panel): + return _fit_twostage(twostage_panel, aggregate="all") + + +class TestTwoStageShim: + def test_plain_fit_does_not_warn(self, twostage_panel): + from diff_diff import TwoStageDiD + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + TwoStageDiD().fit(twostage_panel, **TWOSTAGE_KW) + assert [w for w in caught if issubclass(w.category, FutureWarning)] == [] + + def test_aggregate_kwarg_warns_even_at_none(self, twostage_panel): + from diff_diff import TwoStageDiD + + with pytest.warns(FutureWarning, match=r"TwoStageDiD\.fit\(aggregate=\)"): + TwoStageDiD().fit(twostage_panel, **TWOSTAGE_KW, aggregate=None) + + def test_balance_e_kwarg_warns_alone(self, twostage_panel): + from diff_diff import TwoStageDiD + + with pytest.warns(FutureWarning, match=r"TwoStageDiD\.fit\(balance_e=\)"): + TwoStageDiD().fit(twostage_panel, **TWOSTAGE_KW, balance_e=None) + + def test_joint_supply_warns_once_naming_both(self, twostage_panel): + from diff_diff import TwoStageDiD + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + TwoStageDiD().fit(twostage_panel, **TWOSTAGE_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, twostage_panel): + res = _fit_twostage(twostage_panel, aggregate="nonsense") + plain = _fit_twostage(twostage_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, twostage_fitted, twostage_fit_time): + assert twostage_fit_time.event_study_effects is not None + assert twostage_fit_time.group_effects is not None + es = twostage_fitted.aggregate("event_study") + for e, row in twostage_fit_time.event_study_effects.items(): + i = list(es.event_time).index(e) + if np.isfinite(row["effect"]) and not es.is_reference[i]: + np.testing.assert_allclose(row["effect"], es.att[i], rtol=1e-14) + + def test_wrapper_forwarded_aggregate_warns(self, twostage_panel): + from diff_diff import two_stage_did + + with pytest.warns(FutureWarning, match=r"TwoStageDiD\.fit\(aggregate=\)"): + two_stage_did( + twostage_panel, + "outcome", + "unit", + "period", + "first_treat", + aggregate="event_study", + ) + + def test_wrapper_forwarded_balance_e_warns(self, twostage_panel): + from diff_diff import two_stage_did + + with pytest.warns(FutureWarning, match=r"TwoStageDiD\.fit\(balance_e=\)"): + two_stage_did(twostage_panel, "outcome", "unit", "period", "first_treat", balance_e=1) + + def test_plain_wrapper_call_does_not_warn(self, twostage_panel): + from diff_diff import two_stage_did + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + two_stage_did(twostage_panel, "outcome", "unit", "period", "first_treat") + assert [w for w in caught if issubclass(w.category, FutureWarning)] == [] + + +class TestTwoStageAggregate: + def test_event_study_inert(self, twostage_fitted, twostage_fit_time): + es = twostage_fitted.aggregate("event_study") + _assert_es_container_matches_fit_time(es, twostage_fit_time, "ts/es") + + @pytest.mark.parametrize("balance_e", [0, 1, 2]) + def test_balance_e_inert(self, twostage_panel, twostage_fitted, balance_e): + ref = _fit_twostage(twostage_panel, aggregate="event_study", balance_e=balance_e) + es = twostage_fitted.aggregate("event_study", balance_e=balance_e) + _assert_es_container_matches_fit_time(es, ref, f"ts/es/be{balance_e}") + + def test_group_inert(self, twostage_fitted, twostage_fit_time): + _assert_group_matches_fit_time( + twostage_fitted.aggregate("group"), twostage_fit_time, "ts/gr" + ) + + def test_simple_relay_bit_exact(self, twostage_fitted): + sm = twostage_fitted.aggregate("simple") + assert sm.att[0] == twostage_fitted.overall_att + assert sm.se[0] == twostage_fitted.overall_se + assert sm.n_kind == "obs" + assert sm.n[0] == twostage_fitted.n_treated_obs + assert np.isnan(sm.to_dataframe()["df"].to_numpy()).all() + + def test_m092_vcov_parity_analytical(self, twostage_fitted, twostage_panel): + # The post-fit container threads the recomputed joint GMM vcov + + # index + df exactly as the level-matched fit-time container does. + from diff_diff.results_base import build_event_study_surface + + ref = build_event_study_surface(_fit_twostage(twostage_panel, aggregate="event_study")) + es = twostage_fitted.aggregate("event_study") + assert es.vcov is not None and ref.vcov is not None + np.testing.assert_allclose(es.vcov, ref.vcov, rtol=0, atol=1e-14, equal_nan=True) + assert list(es.vcov_index) == list(ref.vcov_index) + + def test_m092_vcov_cleared_on_replicate(self): + d, rep_cols = _twostage_survey_panel(replicate=True) + sd = _twostage_survey_design(rep_cols) + plain = _fit_twostage(d, survey_design=sd) + es = plain.aggregate("event_study") + assert es.vcov is None and es.vcov_index is None + # df provenance still threads (level-matched replayed value) + ref = _fit_twostage(d, aggregate="event_study", survey_design=sd) + from diff_diff.results_base import build_event_study_surface + + ref_surface = build_event_study_surface(ref) + assert (es.df_survey == ref_surface.df_survey) or ( + es.df_survey is None and ref_surface.df_survey is None + ) + + def test_cluster_arm_inert(self): + d = _twostage_panel(seed=11).copy() + d["cl"] = (d["unit"] // 3).astype(int) + ref = _fit_twostage(d, est_kw={"cluster": "cl"}, aggregate="all") + plain = _fit_twostage(d, est_kw={"cluster": "cl"}) + _assert_es_container_matches_fit_time(plain.aggregate("event_study"), ref, "ts/cl") + _assert_group_matches_fit_time(plain.aggregate("group"), ref, "ts/cl") + + def test_cluster_naming_unit_column_inert(self, twostage_panel): + # cluster= legally names the unit column - the kit's column-subset + # dedup must keep the frame single-labeled (a duplicated column + # would break df[unit].map in the moved Stage-1 helpers). + ref = _fit_twostage(twostage_panel, est_kw={"cluster": "unit"}, aggregate="all") + plain = _fit_twostage(twostage_panel, est_kw={"cluster": "unit"}) + _assert_es_container_matches_fit_time(plain.aggregate("event_study"), ref, "ts/cl-unit") + _assert_group_matches_fit_time(plain.aggregate("group"), ref, "ts/cl-unit") + + def test_covariate_arm_inert(self): + d = _twostage_panel(seed=23).copy() + rng = np.random.default_rng(23) + d["x1"] = rng.normal(size=len(d)) + 0.1 * d["period"] + ref = _fit_twostage(d, aggregate="all", covariates=["x1"]) + plain = _fit_twostage(d, covariates=["x1"]) + _assert_es_container_matches_fit_time(plain.aggregate("event_study"), ref, "ts/cov") + _assert_group_matches_fit_time(plain.aggregate("group"), ref, "ts/cov") + + def test_prop5_arm_inert(self): + d = _twostage_prop5_panel() + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ref = _fit_twostage(d, aggregate="event_study") + plain = _fit_twostage(d) + es = plain.aggregate("event_study") + _assert_es_container_matches_fit_time(es, ref, "ts/prop5") + + def test_survey_tsl_arm_inert(self): + d, _ = _twostage_survey_panel() + sd = _twostage_survey_design(psu=True) + ref = _fit_twostage(d, aggregate="all", survey_design=sd) + plain = _fit_twostage(d, survey_design=sd) + _assert_es_container_matches_fit_time(plain.aggregate("event_study"), ref, "ts/tsl") + _assert_group_matches_fit_time(plain.aggregate("group"), ref, "ts/tsl") + gr = plain.aggregate("group") + df_col = gr.to_dataframe()["df"].to_numpy() + finite_p = np.isfinite(gr.p_value) + assert np.isfinite(df_col[finite_p]).all() # scalar survey-df broadcast + + def test_always_treated_pad_arm_inert(self): + # The only arm where score_pad_mask/cluster_ids_full are non-None: + # an implementation storing None unconditionally passes every other + # arm but breaks inertness here. + d = _twostage_always_treated_panel() + sd = _twostage_survey_design(psu=True) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ref = _fit_twostage(d, aggregate="all", survey_design=sd) + plain = _fit_twostage(d, survey_design=sd) + es = plain.aggregate("event_study") + gr = plain.aggregate("group") + _assert_es_container_matches_fit_time(es, ref, "ts/pad") + _assert_group_matches_fit_time(gr, ref, "ts/pad") + + @pytest.mark.parametrize("degenerate", [None, "dropped", "undefined"]) + def test_replicate_arms_level_matched(self, degenerate): + d, rep_cols = _twostage_survey_panel(replicate=True, degenerate=degenerate) + sd = _twostage_survey_design(rep_cols) + ref_es = _fit_twostage(d, aggregate="event_study", survey_design=sd) + ref_gr = _fit_twostage(d, aggregate="group", survey_design=sd) + plain = _fit_twostage(d, survey_design=sd) + _assert_es_container_matches_fit_time( + plain.aggregate("event_study"), ref_es, f"ts/rep/{degenerate}" + ) + _assert_group_matches_fit_time(plain.aggregate("group"), ref_gr, f"ts/rep/{degenerate}") + + def test_all_prop5_horizons_survive_early_return(self): + # Local-review P1 (pre-existing corner in the verbatim-moved code, + # fixed in this PR): when EVERY non-reference horizon is + # Proposition-5-unidentified, est_horizons empties and the early + # return used to DROP the built prop5 rows - real treated horizons + # reported as absent instead of unidentified. Both the fit-time + # and post-fit surfaces must keep them as all-NaN rows with + # n_obs > 0, plus the consolidated Prop-5 warning. + rng = np.random.default_rng(5) + rows = [] + for u in range(12): # cohort 2: pre at t=1, post only at h >= h_bar + for t in (1, 8, 9, 10): + rows.append((u, t, rng.normal() + (2.0 if t >= 2 else 0.0), 2)) + for u in range(20, 32): # cohort 8: pre-periods only + for t in range(1, 8): + rows.append((u, t, rng.normal(), 8)) + d = pd.DataFrame(rows, columns=["unit", "period", "outcome", "first_treat"]) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + ref = _fit_twostage(d.copy(), aggregate="event_study") + plain = _fit_twostage(d.copy()) + es = plain.aggregate("event_study") + assert any("Proposition 5" in str(w.message) for w in caught) + for surface_dict in (ref.event_study_effects,): + prop5 = {h: r for h, r in surface_dict.items() if h != -1} + assert prop5, "Prop-5 rows must survive the early return" + for h, r in prop5.items(): + assert h >= 6 and r["n_obs"] > 0 + assert np.isnan(r["effect"]) and np.isnan(r["se"]) + df_ = es.to_dataframe() + prop5_rows = df_[(df_["event_time"] >= 6)] + assert len(prop5_rows) > 0 + assert np.isnan(prop5_rows["att"]).all() + assert (prop5_rows["n"] > 0).all() + _assert_es_container_matches_fit_time(es, ref, "ts/all-prop5") + + def test_replicate_overall_row_migration_delta(self): + d, rep_cols = _twostage_survey_panel(replicate=True, degenerate="cohort_zero") + sd = _twostage_survey_design(rep_cols) + plain = _fit_twostage(d, survey_design=sd) + ref_gr = _fit_twostage(d, aggregate="group", survey_design=sd) + assert plain.overall_se != ref_gr.overall_se + _assert_group_matches_fit_time(plain.aggregate("group"), ref_gr, "ts/delta") + dh, rep_h = _twostage_survey_panel(replicate=True) + sdh = _twostage_survey_design(rep_h) + # Same BLAS-shape caveat as the Imputation twin above. + np.testing.assert_allclose( + _fit_twostage(dh, survey_design=sdh).overall_se, + _fit_twostage(dh, aggregate="group", survey_design=sdh).overall_se, + rtol=1e-12, + ) + + def test_pretrends_arm_inert(self, twostage_panel): + ref = _fit_twostage(twostage_panel, est_kw={"pretrends": True}, aggregate="event_study") + plain = _fit_twostage(twostage_panel, est_kw={"pretrends": True}) + _assert_es_container_matches_fit_time(plain.aggregate("event_study"), ref, "ts/pretrends") + + def test_isolation_under_public_field_mutation(self, twostage_panel): + res = _fit_twostage(twostage_panel) + base_es = res.aggregate("event_study").to_dataframe() + base_gr = res.aggregate("group").to_dataframe() + base_sm = res.aggregate("simple").to_dataframe() + res.groups.pop() + res.time_periods.pop() + object.__setattr__(res, "alpha", 0.5) + object.__setattr__(res, "anticipation", 3) + pd.testing.assert_frame_equal(res.aggregate("event_study").to_dataframe(), base_es) + pd.testing.assert_frame_equal(res.aggregate("group").to_dataframe(), base_gr) + pd.testing.assert_frame_equal(res.aggregate("simple").to_dataframe(), base_sm) + + def test_metadata_never_mutated_and_isolated(self): + d, rep_cols = _twostage_survey_panel(replicate=True) + sd = _twostage_survey_design(rep_cols) + res = _fit_twostage(d, survey_design=sd) + base = res.aggregate("event_study").to_dataframe() + base_df_survey = res.aggregate("event_study").df_survey + res.survey_metadata.df_survey = 999.0 + res.survey_metadata.replicate_method = None + after = res.aggregate("event_study") + pd.testing.assert_frame_equal(after.to_dataframe(), base) + assert after.df_survey == base_df_survey + assert res.survey_metadata.df_survey == 999.0 + + def test_repeated_and_order_independent(self, twostage_fitted): + a = twostage_fitted.aggregate("group").to_dataframe() + twostage_fitted.aggregate("event_study") + twostage_fitted.aggregate("simple") + b = twostage_fitted.aggregate("group").to_dataframe() + pd.testing.assert_frame_equal(a, b) + + def test_bootstrap_fails_closed_all_levels(self, twostage_panel): + res = _fit_twostage(twostage_panel, est_kw={"n_bootstrap": 19, "seed": 1}) + for level in ("simple", "event_study", "group"): + with pytest.raises(NotImplementedError, match="bootstrap"): + res.aggregate(level) + + def test_fail_closed_vocabulary(self, twostage_fitted): + for bad in ("calendar", "all", "nonsense"): + with pytest.raises(ValueError, match="Unsupported aggregation type"): + twostage_fitted.aggregate(bad) + with pytest.raises(ValueError, match="balance_e"): + twostage_fitted.aggregate("group", balance_e=1) + with pytest.raises(ValueError): + twostage_fitted.aggregate("group", weights="cell") + + def test_legacy_pickle_without_kit(self, twostage_fitted): + import copy + + legacy = copy.copy(twostage_fitted) + object.__setattr__(legacy, "_aggregation_kit", None) + with pytest.raises(ValueError, match="aggregation kit"): + legacy.aggregate("group") + + def test_pickle_roundtrip(self, twostage_panel): + import pickle + + res = _fit_twostage(twostage_panel) + clone = pickle.loads(pickle.dumps(res)) + pd.testing.assert_frame_equal( + clone.aggregate("group").to_dataframe(), res.aggregate("group").to_dataframe() + ) + + def test_empty_balance_window(self, twostage_fitted): + with pytest.warns(UserWarning, match="balance_e"): + es = twostage_fitted.aggregate("event_study", balance_e=100) + df_ = es.to_dataframe() + assert (df_["is_reference"] | ~np.isfinite(df_["att"])).all() diff --git a/tests/test_event_study_consumers.py b/tests/test_event_study_consumers.py index 6eb6acb9..1c16e68c 100644 --- a/tests/test_event_study_consumers.py +++ b/tests/test_event_study_consumers.py @@ -1951,3 +1951,67 @@ def test_real_efficient_container_rejected_by_design(self): compute_honest_did(surface, M=1.0) with pytest.raises(TypeError, match="rejected BY DESIGN"): compute_pretrends_power(surface, M=1.0) + + +class TestImputationContainerRejection: + """A REAL ImputationDiD post-fit container is rejected BY DESIGN. + + The match string pins the NEW clause text, not the class name (the + got source={...!r} interpolation would match "ImputationDiDResults" + even without the message edit). + """ + + def test_real_imputation_container_rejected_by_design(self): + from diff_diff import ImputationDiD + 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 = ImputationDiD().fit( + d, + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + ) + surface = res.aggregate("event_study") + assert surface.source == "ImputationDiDResults" + 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) + + +class TestTwoStageContainerRejection: + """A REAL TwoStageDiD post-fit container is rejected with the DEFERRED + clause (not the by-design clause: analytical fits DO carry the joint + Gardner-GMM covariance, but the pre-period coefficients are stage-1 + residual means, not reference-normalized contrasts - admission awaits a + normalization derivation; see the REGISTRY TwoStageDiD Note and the + DEFERRED.md paper-gated row). A pretrends=True fit - the strongest + admission candidate (estimated pre-periods + joint vcov) - is still + rejected. + """ + + def test_real_twostage_container_rejected_deferred(self): + import warnings as _warnings + + from diff_diff import TwoStageDiD + 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) + with _warnings.catch_warnings(): + _warnings.simplefilter("ignore") + res = TwoStageDiD(pretrends=True).fit( + d, + outcome="outcome", + unit="unit", + time="period", + first_treat="first_treat", + ) + surface = res.aggregate("event_study") + assert surface.source == "TwoStageDiDResults" + assert surface.vcov is not None # the joint GMM covariance IS present + with pytest.raises(TypeError, match="DEFERRED pending a normalization"): + compute_honest_did(surface, M=1.0) + with pytest.raises(TypeError, match="DEFERRED pending a normalization"): + compute_pretrends_power(surface, M=1.0) diff --git a/tests/test_event_study_surface.py b/tests/test_event_study_surface.py index b8793da0..cb8f5ce5 100644 --- a/tests/test_event_study_surface.py +++ b/tests/test_event_study_surface.py @@ -850,7 +850,7 @@ class _Malformed: def test_zero_count_nan_horizon_is_not_a_reference(): # TwoStageDiD emits effect=NaN, n_obs=0 for an estimated horizon whose - # observations are all filtered (two_stage.py:2669-2681) - distinct from + # observations are all filtered (two_stage_aggregation.py::_stage2_event_study) - distinct from # its effect=0.0, n_obs=0 reference. The count-sentinel adapter path must # mark ONLY the finite-effect (0.0) row as reference and preserve the # NaN-effect horizon as a non-reference NaN row (never normalize it to 0). diff --git a/tests/test_guides.py b/tests/test_guides.py index 9f362121..4fb13bbd 100644 --- a/tests/test_guides.py +++ b/tests/test_guides.py @@ -890,3 +890,32 @@ def test_llms_full_efficient_fit_aggregate_line_documents_shim(self): ) assert "NOT_SUPPLIED" in bal_line assert "DEPRECATED (M-120)" in bal_line + + def _section(self, header): + text = get_llm_guide("full") + start = text.index(header) + nxt = text.index("\n### ", start + 1) + return text[start:nxt] + + def _assert_shim_lines(self, section, fit_call, agg_row, bal_row): + fit_start = section.index(fit_call) + 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 f"DEPRECATED ({agg_row})" 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 f"DEPRECATED ({bal_row})" in bal_line + + def test_llms_full_imputation_fit_aggregate_line_documents_shim(self): + # The M-021/M-118 twin of the EfficientDiD pin above - no other + # test covers the ImputationDiD signature block's shim comments. + self._assert_shim_lines(self._section("### ImputationDiD"), "imp.fit(", "M-021", "M-118") + + def test_llms_full_two_stage_fit_aggregate_line_documents_shim(self): + self._assert_shim_lines(self._section("### TwoStageDiD"), ".fit(", "M-022", "M-119") diff --git a/tests/test_imputation.py b/tests/test_imputation.py index 24decca3..74357579 100644 --- a/tests/test_imputation.py +++ b/tests/test_imputation.py @@ -16,6 +16,21 @@ ) from diff_diff.survey import SurveyDesign +# --------------------------------------------------------------------------- +# Rows M-021/M-022 (+ M-118/M-119): ImputationDiD / TwoStageDiD +# ``fit(aggregate=, balance_e=)`` is deprecated (3.9, removed 4.0) and warns on +# ANY supplied value. The deprecated fit-time route is kept DELIBERATELY here: +# these tests pin FIT-TIME surface behaviour (bit-equality grids, bootstrap +# aggregation, R/Stata parity, replicate overrides, native effect dicts) that +# the post-fit ``results.aggregate(...)`` container route does not reproduce +# shape-for-shape. The shim warning is therefore filtered BY MESSAGE, scoped to +# these two estimators only - every other FutureWarning (including the other +# estimators' aggregate() shims) still surfaces. +# --------------------------------------------------------------------------- +pytestmark = pytest.mark.filterwarnings( + r"ignore:(ImputationDiD|TwoStageDiD)\.fit\((aggregate=|balance_e=|aggregate= / balance_e=)\):FutureWarning" +) + # ============================================================================= # Shared test data generation # ============================================================================= @@ -877,7 +892,8 @@ def test_sparse_solver_lsmr_fallback(self): # Monkey-patch the sparse factorization to force the LSMR fallback. with unittest.mock.patch( - "diff_diff.imputation.sparse_factorized", side_effect=RuntimeError("test failure") + "diff_diff.imputation_aggregation.sparse_factorized", + side_effect=RuntimeError("test failure"), ): results = est.fit( data, @@ -904,7 +920,8 @@ def test_sparse_solver_lsmr_fallback_emits_warning(self): est = ImputationDiD() with unittest.mock.patch( - "diff_diff.imputation.sparse_factorized", side_effect=RuntimeError("test failure") + "diff_diff.imputation_aggregation.sparse_factorized", + side_effect=RuntimeError("test failure"), ): with pytest.warns( UserWarning, match="sparse factorization.*falling back to a sparse LSMR" @@ -2513,7 +2530,7 @@ def test_bootstrap_cluster_hc1_bit_equal(self, ci_params, aggregate): == r_explicit.bootstrap_results.overall_att_se ) # Per-horizon / per-group bootstrap SE override branches at - # diff_diff/imputation.py:854-887 must also agree. + # imputation_aggregation.py::_replicate_override_aggregates must also agree. if aggregate == "event_study": assert r_default.bootstrap_results.event_study_ses is not None assert r_explicit.bootstrap_results.event_study_ses is not None @@ -3165,12 +3182,12 @@ def test_no_dense_materialization_on_fallback(self, monkeypatch): normal matrix (the O((U+T+K)^2) OOM risk this closes).""" import unittest.mock - import diff_diff.imputation as imp + import diff_diff.imputation_aggregation as imp data = generate_test_data(n_units=60, n_periods=6, seed=7) with unittest.mock.patch( - "diff_diff.imputation.sparse_factorized", side_effect=RuntimeError("forced") + "diff_diff.imputation_aggregation.sparse_factorized", side_effect=RuntimeError("forced") ): with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") @@ -3221,7 +3238,7 @@ def _fake_lsmr(A, b, **kwargs): data = generate_test_data(n_units=60, n_periods=6, seed=7) monkeypatch.setattr("scipy.sparse.linalg.lsmr", _fake_lsmr) with unittest.mock.patch( - "diff_diff.imputation.sparse_factorized", side_effect=RuntimeError("forced") + "diff_diff.imputation_aggregation.sparse_factorized", side_effect=RuntimeError("forced") ): with warnings.catch_warnings(): warnings.simplefilter("ignore") @@ -3363,18 +3380,46 @@ def test_survey_df_precedence_on_leads(self): assert e["conf_int"] == e0["conf_int"] def test_inert_config_warns_on_explicit_nondefault(self): + """M-127 REACHABILITY predicate (revised with the M-021 post-fit + migration): warn iff the per-lead inference is unreachable on every + route for this fit config — pretrends=False always warns; + pretrends=True analytical fits never warn (post-fit + results.aggregate('event_study') reaches the leads regardless of + the deprecated fit-time aggregate value); pretrends=True with + n_bootstrap>0 and aggregate unset warns (no ES surface is built + and post-fit aggregate() fails closed); pretrends=True with the + deprecated fit-time ES supplied never warns, bootstrap included + (fit-time leads use analytical inference).""" data = self._panel() base = dict(outcome="outcome", unit="unit", time="time", first_treat="first_treat") with pytest.warns(UserWarning, match="affects only the pretrends"): ImputationDiD(df_convention="cluster").fit(data, **base) - with pytest.warns(UserWarning, match="affects only the pretrends"): + # pretrends=True + deprecated aggregate='group': the knob is NO + # LONGER inert (post-fit ES reaches the leads) — no inert-config + # warning (the FutureWarning from aggregate= is separate). + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") ImputationDiD(pretrends=True, df_convention="cluster").fit( data, aggregate="group", **base ) + assert not any("affects only the pretrends" in str(w.message) for w in caught) with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") ImputationDiD(pretrends=True, df_convention="cluster").fit(data, **self._kw) assert not any("affects only the pretrends" in str(w.message) for w in caught) + # pretrends=True + bootstrap, aggregate unset: unreachable → warns. + with pytest.warns(UserWarning, match="affects only the pretrends"): + ImputationDiD(pretrends=True, df_convention="cluster", n_bootstrap=9, seed=1).fit( + data, **base + ) + # pretrends=True + bootstrap + deprecated fit-time ES: reachable + # (analytical lead inference rides the fit-time surface) → no warn. + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + ImputationDiD(pretrends=True, df_convention="cluster", n_bootstrap=9, seed=1).fit( + data, aggregate="event_study", **base + ) + assert not any("affects only the pretrends" in str(w.message) for w in caught) def test_validation_and_transactional_set_params(self): with pytest.raises(ValueError, match="df_convention"): diff --git a/tests/test_imputation_loo_stata_parity.py b/tests/test_imputation_loo_stata_parity.py index bf9a1f47..3ab16b0f 100644 --- a/tests/test_imputation_loo_stata_parity.py +++ b/tests/test_imputation_loo_stata_parity.py @@ -48,6 +48,21 @@ from diff_diff import ImputationDiD +# --------------------------------------------------------------------------- +# Rows M-021/M-022 (+ M-118/M-119): ImputationDiD / TwoStageDiD +# ``fit(aggregate=, balance_e=)`` is deprecated (3.9, removed 4.0) and warns on +# ANY supplied value. The deprecated fit-time route is kept DELIBERATELY here: +# these tests pin FIT-TIME surface behaviour (bit-equality grids, bootstrap +# aggregation, R/Stata parity, replicate overrides, native effect dicts) that +# the post-fit ``results.aggregate(...)`` container route does not reproduce +# shape-for-shape. The shim warning is therefore filtered BY MESSAGE, scoped to +# these two estimators only - every other FutureWarning (including the other +# estimators' aggregate() shims) still surfaces. +# --------------------------------------------------------------------------- +pytestmark = pytest.mark.filterwarnings( + r"ignore:(ImputationDiD|TwoStageDiD)\.fit\((aggregate=|balance_e=|aggregate= / balance_e=)\):FutureWarning" +) + _DATA = Path(__file__).parent.parent / "benchmarks" / "data" STATA_GOLDEN_PATH = _DATA / "didimputation_loo_stata_golden.json" R_GOLDEN_PATH = _DATA / "didimputation_golden.json" diff --git a/tests/test_lpdid.py b/tests/test_lpdid.py index 1f15b8d6..d69daa10 100644 --- a/tests/test_lpdid.py +++ b/tests/test_lpdid.py @@ -729,6 +729,12 @@ def test_ce2_reweighted_equals_callaway_santanna(self): cs.event_study_effects[h]["effect"], abs=1e-9 ) + # The BJS cross-check reads `bjs.event_study_effects` (the native fit-time + # dict) so it keeps the deprecated `fit(aggregate=)` route (rows M-021); + # scoped by message to the ImputationDiD shim only. + @pytest.mark.filterwarnings( + r"ignore:ImputationDiD\.fit\((aggregate=|balance_e=|aggregate= / balance_e=)\):FutureWarning" + ) def test_ce3_pmd_single_cohort_equals_bjs_imputation(self): # PMD LP-DiD (k=t-1, "max") with a SINGLE treated cohort == BJS # imputation [paper Section 3.4, footnotes 10-11; only single-cohort diff --git a/tests/test_methodology_imputation.py b/tests/test_methodology_imputation.py index 45e54a54..f74355d3 100644 --- a/tests/test_methodology_imputation.py +++ b/tests/test_methodology_imputation.py @@ -52,6 +52,21 @@ from diff_diff import ImputationDiD +# --------------------------------------------------------------------------- +# Rows M-021/M-022 (+ M-118/M-119): ImputationDiD / TwoStageDiD +# ``fit(aggregate=, balance_e=)`` is deprecated (3.9, removed 4.0) and warns on +# ANY supplied value. The deprecated fit-time route is kept DELIBERATELY here: +# these tests pin FIT-TIME surface behaviour (bit-equality grids, bootstrap +# aggregation, R/Stata parity, replicate overrides, native effect dicts) that +# the post-fit ``results.aggregate(...)`` container route does not reproduce +# shape-for-shape. The shim warning is therefore filtered BY MESSAGE, scoped to +# these two estimators only - every other FutureWarning (including the other +# estimators' aggregate() shims) still surfaces. +# --------------------------------------------------------------------------- +pytestmark = pytest.mark.filterwarnings( + r"ignore:(ImputationDiD|TwoStageDiD)\.fit\((aggregate=|balance_e=|aggregate= / balance_e=)\):FutureWarning" +) + # ============================================================================= # Module-level R-fixture availability + per-class seed decorrelation # ============================================================================= diff --git a/tests/test_methodology_two_stage.py b/tests/test_methodology_two_stage.py index 9fcaed3c..4c9c1a4f 100644 --- a/tests/test_methodology_two_stage.py +++ b/tests/test_methodology_two_stage.py @@ -56,6 +56,21 @@ from diff_diff import ImputationDiD, TwoStageDiD +# --------------------------------------------------------------------------- +# Rows M-021/M-022 (+ M-118/M-119): ImputationDiD / TwoStageDiD +# ``fit(aggregate=, balance_e=)`` is deprecated (3.9, removed 4.0) and warns on +# ANY supplied value. The deprecated fit-time route is kept DELIBERATELY here: +# these tests pin FIT-TIME surface behaviour (bit-equality grids, bootstrap +# aggregation, R/Stata parity, replicate overrides, native effect dicts) that +# the post-fit ``results.aggregate(...)`` container route does not reproduce +# shape-for-shape. The shim warning is therefore filtered BY MESSAGE, scoped to +# these two estimators only - every other FutureWarning (including the other +# estimators' aggregate() shims) still surfaces. +# --------------------------------------------------------------------------- +pytestmark = pytest.mark.filterwarnings( + r"ignore:(ImputationDiD|TwoStageDiD)\.fit\((aggregate=|balance_e=|aggregate= / balance_e=)\):FutureWarning" +) + # ============================================================================= # Module-level R-fixture availability + per-class seed decorrelation # ============================================================================= @@ -589,7 +604,7 @@ def test_no_never_treated_horizons_nan(self) -> None: """Proposition 5 (Borusyak et al. 2024): with no never-treated units and h_bar = max(groups) - min(groups), horizons h >= h_bar are not identified -> NaN effect with n_obs > 0 and a warning. Regression-pins the behavior - implemented at two_stage.py:2531-2674 (mirror of ImputationDiD).""" + implemented in two_stage_aggregation.py::_stage2_event_study (mirror of ImputationDiD).""" rng = np.random.default_rng(_BASE_SEED_IDENT + 2) # Cohorts 3 and 5, NO never-treated => h_bar = 5 - 3 = 2. panel = _make_staggered_panel( diff --git a/tests/test_naming_guard.py b/tests/test_naming_guard.py index 314d735f..ffcfcbad 100644 --- a/tests/test_naming_guard.py +++ b/tests/test_naming_guard.py @@ -48,9 +48,9 @@ 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, 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 + 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 version-aware lifecycle gate (``_NEXT_RELEASE``) and arm automatically at the version bump. Ambiguous-param readers that build kwargs INDIRECTLY - diff --git a/tests/test_practitioner.py b/tests/test_practitioner.py index 4d88d9f4..7f72f6c5 100644 --- a/tests/test_practitioner.py +++ b/tests/test_practitioner.py @@ -577,6 +577,71 @@ def test_aggregation_step_name_is_non_steps_key(self, mock_efficient_results): assert len(self._agg_step(output)) == 1 +class _AggregationStepMixin: + """Shared pins for the post-fit aggregation guidance (M-021/M-022, + mirroring the EfficientDiD M-023 pins above).""" + + row_id = "" + fit_var = "" + + @staticmethod + def _agg_step(output): + return [ + s + for s in output["next_steps"] + if "Aggregate treatment-effect heterogeneity" in s["label"] + ] + + def _results(self): # pragma: no cover - overridden + raise NotImplementedError + + def test_aggregation_step_post_fit_branch(self): + output = practitioner_next_steps(self._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"] + assert self.row_id in steps[0]["why"] + + def test_aggregation_step_bootstrap_branch(self): + r = self._results() + r.bootstrap_results = object() + output = practitioner_next_steps(r, 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): + from diff_diff.practitioner import STEPS + + assert "aggregation" not in STEPS + output = practitioner_next_steps( + self._results(), completed_steps=["heterogeneity"], verbose=False + ) + assert len(self._agg_step(output)) == 1 + + +class TestImputationAggregationStep(_AggregationStepMixin): + row_id = "M-021" + + def _results(self): + r = ImputationDiDResults.__new__(ImputationDiDResults) + r.overall_att = 0.6 + r.overall_se = 0.15 + return r + + +class TestTwoStageAggregationStep(_AggregationStepMixin): + row_id = "M-022" + + def _results(self): + r = TwoStageDiDResults.__new__(TwoStageDiDResults) + r.overall_att = 0.6 + r.overall_se = 0.15 + return r + + # --------------------------------------------------------------------------- # Tests: unknown result type fallback # --------------------------------------------------------------------------- diff --git a/tests/test_pretrends_event_study.py b/tests/test_pretrends_event_study.py index a0ab8990..4f90c18b 100644 --- a/tests/test_pretrends_event_study.py +++ b/tests/test_pretrends_event_study.py @@ -7,10 +7,26 @@ import numpy as np import pandas as pd +import pytest from diff_diff.imputation import ImputationDiD from diff_diff.two_stage import TwoStageDiD +# --------------------------------------------------------------------------- +# Rows M-021/M-022 (+ M-118/M-119): ImputationDiD / TwoStageDiD +# ``fit(aggregate=, balance_e=)`` is deprecated (3.9, removed 4.0) and warns on +# ANY supplied value. The deprecated fit-time route is kept DELIBERATELY here: +# these tests pin FIT-TIME surface behaviour (bit-equality grids, bootstrap +# aggregation, R/Stata parity, replicate overrides, native effect dicts) that +# the post-fit ``results.aggregate(...)`` container route does not reproduce +# shape-for-shape. The shim warning is therefore filtered BY MESSAGE, scoped to +# these two estimators only - every other FutureWarning (including the other +# estimators' aggregate() shims) still surfaces. +# --------------------------------------------------------------------------- +pytestmark = pytest.mark.filterwarnings( + r"ignore:(ImputationDiD|TwoStageDiD)\.fit\((aggregate=|balance_e=|aggregate= / balance_e=)\):FutureWarning" +) + def generate_test_data( n_units: int = 100, diff --git a/tests/test_replicate_weight_expansion.py b/tests/test_replicate_weight_expansion.py index d870749d..4880ba71 100644 --- a/tests/test_replicate_weight_expansion.py +++ b/tests/test_replicate_weight_expansion.py @@ -17,6 +17,21 @@ ) from diff_diff.survey import SurveyDesign +# --------------------------------------------------------------------------- +# Rows M-021/M-022 (+ M-118/M-119): ImputationDiD / TwoStageDiD +# ``fit(aggregate=, balance_e=)`` is deprecated (3.9, removed 4.0) and warns on +# ANY supplied value. The deprecated fit-time route is kept DELIBERATELY here: +# these tests pin FIT-TIME surface behaviour (bit-equality grids, bootstrap +# aggregation, R/Stata parity, replicate overrides, native effect dicts) that +# the post-fit ``results.aggregate(...)`` container route does not reproduce +# shape-for-shape. The shim warning is therefore filtered BY MESSAGE, scoped to +# these two estimators only - every other FutureWarning (including the other +# estimators' aggregate() shims) still surfaces. +# --------------------------------------------------------------------------- +pytestmark = pytest.mark.filterwarnings( + r"ignore:(ImputationDiD|TwoStageDiD)\.fit\((aggregate=|balance_e=|aggregate= / balance_e=)\):FutureWarning" +) + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- diff --git a/tests/test_spillover.py b/tests/test_spillover.py index b8e2cdcb..723f09b0 100644 --- a/tests/test_spillover.py +++ b/tests/test_spillover.py @@ -5547,7 +5547,7 @@ def test_p2_finite_mask_forces_drop_under_survey(self): Wave E.3 (shipped): warn-and-dropped rows are RETAINED in the resolved survey design as zero-score padding (matches R - `survey::svyrecvar(subset())` + `imputation.py:2175-2183` + `survey::svyrecvar(subset())` + `imputation_aggregation.py (_compute_lead_coefficients)` precedent). `n_psu` / `n_strata` / `df_survey` reflect the FULL domain (all 10 PSUs across 2 strata) rather than the post-drop fit sample (the prior Wave E.1 behavior of 8 PSUs / df_survey=6). @@ -7562,7 +7562,7 @@ class TestSpilloverDiDWaveE3SubpopulationFullDesign: Methodology anchor: R `survey::svyrecvar(subset(design, mask))` (Lumley 2010 §2.5) — zero-pad scores at the meat-helper boundary; resolved survey design retains full-panel `n_psu` / `n_strata` / `df_survey` / Binder centering. - Library precedent at `imputation.py:2175-2183` and `prep.py:1401-1432`. + Library precedent at `imputation_aggregation.py (_compute_lead_coefficients)` and `prep.py:1401-1432`. A2 invariant (locked in `_scratch/wave_e3_smoke.py`): warn-and-drop and `SurveyDesign.subpopulation()` apply the same zero-pad mechanism — both diff --git a/tests/test_survey_phase4.py b/tests/test_survey_phase4.py index 2e1fbbd5..715e0e59 100644 --- a/tests/test_survey_phase4.py +++ b/tests/test_survey_phase4.py @@ -18,6 +18,21 @@ ) from diff_diff.linalg import solve_logit +# --------------------------------------------------------------------------- +# Rows M-021/M-022 (+ M-118/M-119): ImputationDiD / TwoStageDiD +# ``fit(aggregate=, balance_e=)`` is deprecated (3.9, removed 4.0) and warns on +# ANY supplied value. The deprecated fit-time route is kept DELIBERATELY here: +# these tests pin FIT-TIME surface behaviour (bit-equality grids, bootstrap +# aggregation, R/Stata parity, replicate overrides, native effect dicts) that +# the post-fit ``results.aggregate(...)`` container route does not reproduce +# shape-for-shape. The shim warning is therefore filtered BY MESSAGE, scoped to +# these two estimators only - every other FutureWarning (including the other +# estimators' aggregate() shims) still surfaces. +# --------------------------------------------------------------------------- +pytestmark = pytest.mark.filterwarnings( + r"ignore:(ImputationDiD|TwoStageDiD)\.fit\((aggregate=|balance_e=|aggregate= / balance_e=)\):FutureWarning" +) + # ============================================================================= # Shared Fixtures # ============================================================================= diff --git a/tests/test_survey_phase8.py b/tests/test_survey_phase8.py index 0eeea543..a27b15f1 100644 --- a/tests/test_survey_phase8.py +++ b/tests/test_survey_phase8.py @@ -11,6 +11,21 @@ compute_replicate_vcov, ) +# --------------------------------------------------------------------------- +# Rows M-021/M-022 (+ M-118/M-119): ImputationDiD / TwoStageDiD +# ``fit(aggregate=, balance_e=)`` is deprecated (3.9, removed 4.0) and warns on +# ANY supplied value. The deprecated fit-time route is kept DELIBERATELY here: +# these tests pin FIT-TIME surface behaviour (bit-equality grids, bootstrap +# aggregation, R/Stata parity, replicate overrides, native effect dicts) that +# the post-fit ``results.aggregate(...)`` container route does not reproduce +# shape-for-shape. The shim warning is therefore filtered BY MESSAGE, scoped to +# these two estimators only - every other FutureWarning (including the other +# estimators' aggregate() shims) still surfaces. +# --------------------------------------------------------------------------- +pytestmark = pytest.mark.filterwarnings( + r"ignore:(ImputationDiD|TwoStageDiD)\.fit\((aggregate=|balance_e=|aggregate= / balance_e=)\):FutureWarning" +) + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- diff --git a/tests/test_two_stage.py b/tests/test_two_stage.py index 8b0333ae..4f674c18 100644 --- a/tests/test_two_stage.py +++ b/tests/test_two_stage.py @@ -15,6 +15,21 @@ two_stage_did, ) +# --------------------------------------------------------------------------- +# Rows M-021/M-022 (+ M-118/M-119): ImputationDiD / TwoStageDiD +# ``fit(aggregate=, balance_e=)`` is deprecated (3.9, removed 4.0) and warns on +# ANY supplied value. The deprecated fit-time route is kept DELIBERATELY here: +# these tests pin FIT-TIME surface behaviour (bit-equality grids, bootstrap +# aggregation, R/Stata parity, replicate overrides, native effect dicts) that +# the post-fit ``results.aggregate(...)`` container route does not reproduce +# shape-for-shape. The shim warning is therefore filtered BY MESSAGE, scoped to +# these two estimators only - every other FutureWarning (including the other +# estimators' aggregate() shims) still surfaces. +# --------------------------------------------------------------------------- +pytestmark = pytest.mark.filterwarnings( + r"ignore:(ImputationDiD|TwoStageDiD)\.fit\((aggregate=|balance_e=|aggregate= / balance_e=)\):FutureWarning" +) + # ============================================================================= # Shared test data generation # ============================================================================= @@ -504,7 +519,7 @@ def test_sparse_factorized_dense_fallback_emits_warning(self): data = generate_test_data() with unittest.mock.patch( - "diff_diff.two_stage.sparse_factorized", + "diff_diff.two_stage_aggregation.sparse_factorized", side_effect=RuntimeError("test failure"), ): with pytest.warns( @@ -565,7 +580,7 @@ def _no_lstsq(*a, **k): ) with unittest.mock.patch( - "diff_diff.two_stage.sparse_factorized", + "diff_diff.two_stage_aggregation.sparse_factorized", side_effect=RuntimeError("test failure"), ): with unittest.mock.patch("numpy.linalg.lstsq", _no_lstsq): @@ -594,7 +609,7 @@ def _fake_lsmr(A, b, **kwargs): return (x, 7, 0, 0.0, 0.0, 0.0, 0.0, 0.0) with unittest.mock.patch( - "diff_diff.two_stage.sparse_factorized", + "diff_diff.two_stage_aggregation.sparse_factorized", side_effect=RuntimeError("test failure"), ): with unittest.mock.patch("scipy.sparse.linalg.lsmr", _fake_lsmr): @@ -633,7 +648,7 @@ def _no_lstsq(*a, **k): ) with unittest.mock.patch( - "diff_diff.two_stage.sparse_factorized", + "diff_diff.two_stage_aggregation.sparse_factorized", side_effect=RuntimeError("test failure"), ): with unittest.mock.patch("numpy.linalg.lstsq", _no_lstsq): @@ -661,7 +676,7 @@ def test_forced_fallback_fit_level_lsmr_vs_dense_oracle(self): singular-Omega_0 regression.""" import unittest.mock - import diff_diff.two_stage as ts + import diff_diff.two_stage_aggregation as ts data = generate_test_data() @@ -675,7 +690,7 @@ def _fit(): ) with unittest.mock.patch( - "diff_diff.two_stage.sparse_factorized", + "diff_diff.two_stage_aggregation.sparse_factorized", side_effect=RuntimeError("test failure"), ): with warnings.catch_warnings(): @@ -1439,7 +1454,7 @@ def test_print_summary(self, capsys): def test_sparse_fallback_path(self): """Size guard falls back to per-column path and produces same results.""" - import diff_diff.two_stage as ts_mod + import diff_diff.two_stage_aggregation as ts_mod data = generate_test_data(n_units=50, n_periods=6, seed=42) @@ -1648,7 +1663,7 @@ def test_analytical_bread_rank_reduces_and_warns(self): variance.""" from unittest.mock import patch - import diff_diff.two_stage as ts_mod + import diff_diff.two_stage_aggregation as ts_mod data = generate_test_data(n_units=80, n_periods=6, seed=77) est = TwoStageDiD() @@ -1756,7 +1771,7 @@ def test_dropped_coefficient_propagates_nan_inference(self, aggregate, attr): Covers the analytical event-study AND group surfaces.""" from unittest.mock import patch - import diff_diff.two_stage as ts_mod + import diff_diff.two_stage_aggregation as ts_mod data = generate_test_data(n_units=80, n_periods=6, seed=77) with patch.object( @@ -1894,7 +1909,7 @@ class TestTwoStageDiDWaveE3ParityAlwaysTreated: Mirrors PR #482 SpilloverDiD Wave E.3 (merge 24de9062) which established the same invariant for SpilloverDiD's finite_mask / subpopulation drops. Adopts the R `survey::svyrecvar(subset())` convention (Lumley 2010 §2.5) - and the in-library precedents at `imputation.py:2175-2183` + and the in-library precedents at `imputation_aggregation.py (_compute_lead_coefficients)` (PreTrendsImputation) and `prep.py:1401-1432` (DCDH cell variance). Scope: this PR tests only `vcov_type` paths reachable from TwoStageDiD's diff --git a/tests/test_v4_matrix.py b/tests/test_v4_matrix.py index 7132705d..959aa61e 100644 --- a/tests/test_v4_matrix.py +++ b/tests/test_v4_matrix.py @@ -123,11 +123,12 @@ # 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, 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 = 118 +# reserved pool) = 118, plus 2b PR-3b's Imputation/TwoStage balance_e rows +# (M-118, M-119, claimed from the reserved pool) = 120. +# 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 = 120 # 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 @@ -152,9 +153,11 @@ # 3.9 warning, since the surviving target classes have no shim; EDiD was # initially slated but retained - the CSX paper's own estimator label). # (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. +# 2b PR-3a's EfficientDiD balance_e move (claimed from the reserved pool); +# (118,119) = 2b PR-3b's Imputation/TwoStage balance_e moves (the first and +# second reserved slots, claimed with M-021/M-022). +# M-116 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), @@ -175,6 +178,7 @@ (127, 131), (132, 135), (136, 138), + (118, 119), ] EXPECTED_INITIAL_IDS = frozenset( f"M-{n:03d}" for lo, hi in _INITIAL_ID_RANGES for n in range(lo, hi + 1) @@ -573,13 +577,14 @@ 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 (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-120/M-122 + M-123/M-124 + - M-125 + M-126 + M-127..M-131 + M-132..M-135 + M-136..M-138).""" + Extends as rows ship (120 as of 2b PR-3b's Imputation/TwoStage balance_e + rows: Phase 1 + diagnostic-family + M-092/M-093 + M-094..M-096 + the + M-097..M-115 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) == 118 + assert len(EXPECTED_INITIAL_IDS) == 120 def test_version_tuple_pads_to_three_components():