feat(v4): PR 3 of the 4.0 cut - remove the fit-time aggregate=/balance_e= keywords

TODO.md

@@ -37,8 +37,8 @@ Related tracking surfaces:
 | Opt-in singleton-group pruning for TwoWayFixedEffects (static + event-study mode; reghdfe parity): singleton units/periods are currently RETAINED class-wide - the within-demeaned row is zero so points are unchanged, but N/G/residual-df count it and CR1/finite-sample SEs shift (~0.41019 -> 0.40962 measured; REGISTRY "Deviation from R" Note, R5 review) - reghdfe iteratively drops singletons by default while fixest retains them (diff-diff matches fixest); an opt-in knob needs iterative unit+period pruning with consistent cluster/survey/replicate/Conley array subsetting and a default-flip decision protocol (moves published SEs) | `diff_diff/twfe.py`, `diff_diff/estimators.py`, `diff_diff/utils.py` | 3(a) R5 | Mid | Low |
 | Cohort-timing validation input for the simultaneous-adoption event-study family (TWFE `event_study=True` + MultiPeriodDiD through 3.9): an optional `first_treat=`/`cohort=` column so simultaneous adoption becomes checkable under the contract-valid time-invariant `D_i` indicator - today the staggered-adoption advisory derives timing from within-unit 0->1 transitions, so it can only fire on off-contract time-varying `D_it` input, and with valid `D_i` adoption timing is not observable in the inputs at all (REGISTRY "staggered-adoption detection limit" Notes, both sections); design questions: validate-only vs steering error, and interplay with the M-011 removal | `diff_diff/twfe.py`, `diff_diff/estimators.py` | 3(a) R2 | Mid | Medium |
 | CS parity with DMLDiD's replicate conventions: per-cell replicate SEs for `CallawaySantAnna` (CS leaves per-cell SEs on the weighted sqrt-sum under replicate designs; DMLDiD now computes them via `compute_replicate_if_variance`) AND flip the replace-style effective-df relays to min-cap (`min(df_survey, n_valid - 1)`) at all five sites: CS overall (`staggered.py` overall relay), the shared event-study/group aggregation sites (`staggered_aggregation.py` `df_survey_val = min(non_none_dfs)` x2), and the staggered-DDD engine twins (`_staggered_triple_diff_engine.py` overall + event-study) — the replace convention can RAISE df above the QR-rank design df (anti-conservative); latent, fires only when `n_valid < R` (all-zero/non-finite replicate column). DMLDiD's min-cap + the pinned inherited ES/group behavior are the templates (REGISTRY DMLDiD replicate Note) | `diff_diff/staggered.py`, `diff_diff/staggered_aggregation.py`, `diff_diff/_staggered_triple_diff_engine.py` | DML replicate PR | Mid | Low |
-| ImputationDiD/TwoStageDiD `aggregate()` recompute levels on bootstrapped fits fail closed ('simple' relays since the M-027 per-level convergence; 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 |
-| ContinuousDiD `aggregate('event_study')` on bootstrapped fits fails closed (M-025); a seeded post-fit bootstrap-ES replay is tractable - the multiplier draws are seeded (`np.random.default_rng(self.seed)`) - but needs the FULL per-cell `_bootstrap_info` (bread/ee_treated/Psi_eval/dPsi_*/beta_pred) the pruned kit deliberately drops, so shipping it means a kit-payload change with its own memory contract | `diff_diff/continuous_did_aggregation.py`, `diff_diff/continuous_did_results.py` | 2(b) PR-3c | Mid | Low |
+| `balance_e` replay on bootstrapped ImputationDiD/TwoStageDiD fits: the unrestricted ES/group surface is relayed since 4.0 PR 3 (computed at fit, `balance_e=None`; M-021/M-022), but `aggregate('event_study', balance_e=k)` still raises there - a balanced-window replay needs retained per-target psi (ImputationDiD: the panel-backed kit retains everything the psi precompute reads, so a seeded replay is tractable) / retained per-level GMM scores (TwoStageDiD: function-locals today) | `diff_diff/imputation_results.py`, `diff_diff/two_stage_results.py`, `diff_diff/aggregation.py` | 2(b) PR-3b | Mid | Low |
+| ContinuousDiD post-fit bootstrap-ES REPLAY: since 4.0 PR 3 a bootstrapped fit relays the event-study surface computed at fit (M-025), so the gap is a genuine post-fit replay (fresh draws, no fit-side cost, a route for pre-4.0 pickles); a seeded replay is tractable - the multiplier draws are seeded (`np.random.default_rng(self.seed)`) - but needs the FULL per-cell `_bootstrap_info` (bread/ee_treated/Psi_eval/dPsi_*/beta_pred) the pruned kit deliberately drops, so shipping it means a kit-payload change with its own memory contract | `diff_diff/continuous_did_aggregation.py`, `diff_diff/continuous_did_results.py` | 2(b) PR-3c | Mid | 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 |

benchmarks/python/benchmark_callaway.py

@@ -105,7 +105,6 @@ def main():
             time="time",
             unit="unit",
             first_treat="first_treat",
-            aggregate="all",
         )
 
     cs = CallawaySantAnna(
@@ -122,8 +121,12 @@ def main():
             time="time",
             unit="unit",
             first_treat="first_treat",
-            aggregate="all",  # Get event study and group aggregations
         )
+        # Post-fit aggregations stay inside the timed region: the R `aggte`
+        # yardstick is fit + simple/dynamic/group aggregation (the 3.x
+        # fit(aggregate="all") route, removed at 4.0).
+        es = results.aggregate("event_study")
+        grp = results.aggregate("group")
 
     estimation_time = estimation_timer.elapsed
     total_time = estimation_time
@@ -145,25 +148,25 @@ def main():
             }
         )
 
-    # Format event study effects (if available)
+    # Format event study effects from the post-fit container: one record per
+    # row INCLUDING the reference row (the legacy dict carried it too);
+    # run_refresh.py joins these records against R on `event_time`.
     es_effects = []
-    if results.event_study_effects:
-        for rel_t, effect_data in sorted(results.event_study_effects.items()):
-            es_effects.append({
-                "event_time": int(rel_t),
-                "att": float(effect_data["effect"]),
-                "se": float(effect_data["se"]),
-            })
-
-    # Format group effects (if available)
+    for row in es.to_dataframe().sort_values("event_time").itertuples(index=False):
+        es_effects.append({
+            "event_time": int(row.event_time),
+            "att": float(row.att),
+            "se": float(row.se),
+        })
+
+    # Format group effects from the post-fit container (joined on `group`).
     grp_effects = []
-    if results.group_effects:
-        for g, effect_data in sorted(results.group_effects.items()):
-            grp_effects.append({
-                "group": int(g),
-                "att": float(effect_data["effect"]),
-                "se": float(effect_data["se"]),
-            })
+    for row in grp.to_dataframe().sort_values("label").itertuples(index=False):
+        grp_effects.append({
+            "group": int(row.label),
+            "att": float(row.att),
+            "se": float(row.se),
+        })
 
     # Build output
     output = {

benchmarks/python/benchmark_imputation.py

@@ -79,8 +79,10 @@ def main():
             time="time",
             unit="unit",
             first_treat="first_treat",
-            aggregate="event_study",
         )
+        # Post-fit event study stays inside the timed region (the 3.x
+        # fit(aggregate="event_study") route, removed at 4.0).
+        es = results.aggregate("event_study")
 
     estimation_time = estimation_timer.elapsed
     total_time = estimation_time
@@ -90,18 +92,17 @@ def main():
     n_periods = len(df["time"].unique())
     n_obs = len(df)
 
-    # Format event study effects (if available)
+    # Format event study effects from the post-fit container, skipping the
+    # reference-period marker row exactly as the legacy reader did.
     es_effects = []
-    if results.event_study_effects:
-        for rel_t, effect_data in sorted(results.event_study_effects.items()):
-            # Skip reference period marker (n_obs == 0)
-            if effect_data.get("n_obs", 1) == 0:
-                continue
-            es_effects.append({
-                "event_time": int(rel_t),
-                "att": float(effect_data["effect"]),
-                "se": float(effect_data["se"]),
-            })
+    for row in es.to_dataframe().sort_values("event_time").itertuples(index=False):
+        if bool(row.is_reference):
+            continue
+        es_effects.append({
+            "event_time": int(row.event_time),
+            "att": float(row.att),
+            "se": float(row.se),
+        })
 
     # Build output
     output = {

benchmarks/python/benchmark_stacked_did.py

@@ -96,8 +96,7 @@ def main():
     )
 
     # The event-study surface is always computed at fit (3.9, row M-024);
-    # the deprecated aggregate= kwarg would emit a FutureWarning inside
-    # the timed block and contaminate the measurement.
+    # fit() takes no aggregate= kwarg since 4.0.
     with Timer() as estimation_timer:
         results = est.fit(
             df,

benchmarks/speed_review/bench_brfss_panel.py

@@ -93,20 +93,28 @@ def cs_fit():
         results["cs"] = cs.fit(
             results["panel"], outcome="y_mean",
             unit="state", time="year", first_treat="first_treat",
-            survey_design=results["stage2"], aggregate="all",
+            survey_design=results["stage2"],
         )
+        # Materialize BOTH post-fit containers ONCE inside the timed fit
+        # phase (the 3.x fit(aggregate="all") workload = fit + event-study
+        # + group aggregation): the fit is bootstrapped, so every
+        # aggregate() call replays the bootstrap.
+        results["cs_es"] = results["cs"].aggregate("event_study")
+        results["cs_grp"] = results["cs"].aggregate("group")
 
     def inspect_pretrends():
-        es = results["cs"].event_study_effects or {}
+        es = results["cs_es"].to_dataframe()
+        pre = es[es["event_time"] < 0]
         results["pretrends"] = {
-            rel_t: eff for rel_t, eff in es.items() if rel_t < 0
+            int(r.event_time): {"effect": float(r.att), "se": float(r.se)}
+            for r in pre.itertuples(index=False)
         }
 
     def honest_grid():
         out = {}
         for M in (0.5, 1.0, 1.5):
             out[M] = compute_honest_did(
-                results["cs"], method="relative_magnitude", M=M,
+                results["cs_es"], method="relative_magnitude", M=M,
             )
         results["honest"] = out
 

benchmarks/speed_review/bench_callaway.py

@@ -57,15 +57,18 @@ def bench_fit(n_units, n_bootstrap=0, covariates=None, n_cohorts=5, n_runs=3,
     times = []
     for _ in range(n_runs):
         start = time.perf_counter()
-        cs.fit(
+        res = cs.fit(
             df,
             outcome="outcome",
             unit="unit",
             time="time",
             first_treat="first_treat",
             covariates=covariates,
-            aggregate="all",
         )
+        # Keep "fit + aggregations" as the measured work (the 3.x
+        # fit(aggregate="all") route, removed at 4.0).
+        res.aggregate("event_study")
+        res.aggregate("group")
         elapsed = time.perf_counter() - start
         times.append(elapsed)
 

benchmarks/speed_review/bench_campaign_staggered.py

@@ -1,7 +1,7 @@
 """
 Scenario 1: Staggered marketing campaign.
 
-CallawaySantAnna with covariates + bootstrap + aggregate='all', wrapped in
+CallawaySantAnna with covariates + bootstrap + post-fit aggregate(), wrapped in
 the 8-step Baker workflow: Bacon -> CS fit -> event-study pre-trend
 inspection -> HonestDiD M-grid -> SunAbraham + ImputationDiD robustness
 -> with/without-covariates refit -> practitioner_next_steps.
@@ -64,21 +64,27 @@ def cs_fit():
             control_group="never_treated", estimation_method="dr",
             cluster="unit", n_bootstrap=999, seed=123,
         )
-        results["cs"] = cs.fit(
-            **fit_kwargs, covariates=covars, aggregate="all",
-        )
+        results["cs"] = cs.fit(**fit_kwargs, covariates=covars)
+        # Materialize BOTH post-fit containers ONCE inside the timed fit
+        # phase (the 3.x fit(aggregate="all") workload = fit + event-study
+        # + group aggregation) and reuse them below: the fit has
+        # n_bootstrap=999, so every aggregate() call replays the bootstrap.
+        results["cs_es"] = results["cs"].aggregate("event_study")
+        results["cs_grp"] = results["cs"].aggregate("group")
 
     def inspect_pretrends():
-        es = results["cs"].event_study_effects or {}
+        es = results["cs_es"].to_dataframe()
+        pre = es[es["event_time"] < 0]
         results["pretrends"] = {
-            rel_t: eff for rel_t, eff in es.items() if rel_t < 0
+            int(r.event_time): {"effect": float(r.att), "se": float(r.se)}
+            for r in pre.itertuples(index=False)
         }
 
     def honest_did_grid():
         out = {}
         for M in (0.5, 1.0, 1.5, 2.0):
             out[M] = compute_honest_did(
-                results["cs"], method="relative_magnitude", M=M,
+                results["cs_es"], method="relative_magnitude", M=M,
             )
         results["honest"] = out
 
@@ -88,7 +94,8 @@ def sun_abraham():
 
     def imputation():
         bjs = ImputationDiD(cluster="unit")
-        results["bjs"] = bjs.fit(**fit_kwargs, aggregate="event_study")
+        results["bjs"] = bjs.fit(**fit_kwargs)
+        results["bjs_es"] = results["bjs"].aggregate("event_study")
 
     def cs_no_covariates():
         # Match phase 2's estimator config exactly; the only axis that
@@ -99,7 +106,10 @@ def cs_no_covariates():
             control_group="never_treated", estimation_method="dr",
             cluster="unit", n_bootstrap=999, seed=123,
         )
-        results["cs_nocov"] = cs.fit(**fit_kwargs, aggregate="all")
+        results["cs_nocov"] = cs.fit(**fit_kwargs)
+        # Same fit + event-study + group workload as the covariate arm.
+        results["cs_nocov_es"] = results["cs_nocov"].aggregate("event_study")
+        results["cs_nocov_grp"] = results["cs_nocov"].aggregate("group")
 
     def next_steps():
         results["guidance"] = practitioner_next_steps(results["cs"])
@@ -137,7 +147,7 @@ def run_scale(scale, config):
             "n_obs": int(len(data)),
             "covariates": covars,
             "n_bootstrap": 999,
-            "aggregate": "all",
+            "aggregate": "fit + post-fit aggregate",
             "estimation_method": "dr",
         },
     )

benchmarks/speed_review/bench_memory_scaling.py

@@ -128,7 +128,6 @@ def _run_one(cfg):
             "dose",
             "period",
             "unit",
-            aggregate="event_study",
             survey_design=SurveyDesign(weights="w"),
         )
     elif est == "twfe":

benchmarks/speed_review/bench_solve_ols_fastpath.py

@@ -124,8 +124,11 @@ def fit(d):
                 time="time",
                 first_treat="first_treat",
                 covariates=covs,
-                aggregate="all",
             )
+            # fit + aggregations is the measured work (the 3.x
+            # fit(aggregate="all") route, removed at 4.0).
+            res.aggregate("event_study")
+            res.aggregate("group")
             return float(res.overall_att), float(res.overall_se)
 
         return df, fit

benchmarks/speed_review/validate_results.py

@@ -64,8 +64,11 @@ def run_estimator(df, estimation_method="reg", covariates=None, control_group="n
         time="time",
         first_treat="first_treat",
         covariates=covariates,
-        aggregate="all",
     )
+    # Post-fit containers (the 3.x fit(aggregate="all") route, removed at 4.0);
+    # the fit is bootstrapped, so each call replays the fit-time bootstrap.
+    es_container = results.aggregate("event_study")
+    grp_container = results.aggregate("group")
 
     out = {
         "overall_att": float(results.overall_att),
@@ -87,31 +90,30 @@ def run_estimator(df, estimation_method="reg", covariates=None, control_group="n
         }
     out["group_time_effects"] = gt_effects
 
-    # Event study
-    if results.event_study_effects:
-        es = {}
-        for e, data in sorted(results.event_study_effects.items()):
-            es[str(e)] = {
-                "effect": float(data["effect"]),
-                "se": float(data["se"]),
-                "t_stat": float(data["t_stat"]),
-                "p_value": float(data["p_value"]),
-                "conf_int": [float(data["conf_int"][0]), float(data["conf_int"][1])],
-            }
-        out["event_study"] = es
+    # Event study (same keyed shape as the legacy dict, one record per
+    # container row including the reference row)
+    es = {}
+    for row in es_container.to_dataframe().sort_values("event_time").itertuples(index=False):
+        es[str(int(row.event_time))] = {
+            "effect": float(row.att),
+            "se": float(row.se),
+            "t_stat": float(row.t_stat),
+            "p_value": float(row.p_value),
+            "conf_int": [float(row.conf_int_lower), float(row.conf_int_upper)],
+        }
+    out["event_study"] = es
 
     # Group effects
-    if results.group_effects:
-        ge = {}
-        for g_key, data in sorted(results.group_effects.items()):
-            ge[str(g_key)] = {
-                "effect": float(data["effect"]),
-                "se": float(data["se"]),
-                "t_stat": float(data["t_stat"]),
-                "p_value": float(data["p_value"]),
-                "conf_int": [float(data["conf_int"][0]), float(data["conf_int"][1])],
-            }
-        out["group_effects"] = ge
+    ge = {}
+    for row in grp_container.to_dataframe().sort_values("label").itertuples(index=False):
+        ge[str(row.label)] = {
+            "effect": float(row.att),
+            "se": float(row.se),
+            "t_stat": float(row.t_stat),
+            "p_value": float(row.p_value),
+            "conf_int": [float(row.conf_int_lower), float(row.conf_int_upper)],
+        }
+    out["group_effects"] = ge
 
     return out