feat(v4): ImputationDiD + TwoStageDiD post-fit aggregate() via panel-backed kits (2(b) PR-3b, M-021/M-022/M-118/M-119)

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]

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 |

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** |

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")`.

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_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",