test(v4): PR 3a of the 4.0 cut - migrate the suites to the post-fit aggregation containers

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, deprecated since 3.9 and removed at the 4.0 cut).
+        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, deprecated since 3.9 and removed at the 4.0 cut).
+        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/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, deprecated since 3.9 and removed at the 4.0 cut).
+        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, deprecated since 3.9 and removed at the 4.0 cut).
+            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 is deprecated and goes at the 4.0 cut);
+    # 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
 

changelog.d/20261020-v4-pr3a-aggregate-test-migration.md

@@ -0,0 +1,25 @@
+### Internal
+- **4.0 cut, PR 3a - tests-first migration to the post-fit aggregation containers.** The
+  test and benchmark suites now read event-study / group surfaces from
+  `results.aggregate("event_study")` / `results.aggregate("group")` (the container API that
+  has existed since 3.9) instead of the deprecated fit-time `fit(aggregate=...)` keyword and
+  the legacy `event_study_effects` / `group_effects` fields, so the keyword removals in PRs 3b
+  and 3c land against already-migrated suites. No library source changes and no ledger row
+  flips in this PR; the fit-time keywords still work and still warn. Bootstrapped
+  ImputationDiD / TwoStageDiD / ContinuousDiD sites stay on the fit-time route until PR 3c
+  adds the post-fit relay for them.
+- **Committed per-backend aggregation oracles.** `tests/test_aggregate_contract.py` no longer
+  uses the fit-time route as its numerical reference: every analytic surface is pinned at
+  `rtol=1e-12` (confidence-interval bounds at `1e-9`: scipy's `t.ppf` moved by up to 3.4e-11
+  relative between 1.13 and 1.14+) and every bootstrapped CS / EfficientDiD replay at the
+  replay-parity tolerances against `tests/fixtures/v4_aggregate_oracles.json`, captured on the unmodified
+  tree at `8724fecd` by `tests/_capture_v4_aggregate_oracles.py` under both weight backends
+  (see `tests/fixtures/README.md`). New `tests/helpers/legacy_surface.py` (`es_dict`,
+  `group_dict`, `legacy_standin`) gives consumer tests a faithful stand-in for the 3.x
+  fit-time surface where their subject is the native (legacy-pickle) route.
+- **Benchmarks** (`benchmarks/python/benchmark_callaway.py`, `benchmark_imputation.py`, the
+  `speed_review` scripts) call `aggregate()` post-fit inside the timed region and emit the
+  same JSON row shape the R refresh joins on.
+- `docs/v4-design.md` section 9 records the three-way split of cut-sequence item 3 (3a
+  tests-first, 3b mechanical removals, 3c the Imputation/TwoStage/ContinuousDiD removals +
+  relay).

docs/v4-design.md

@@ -1002,7 +1002,8 @@ re-enumerate the cells' M-id lists:
 cut's PR order; order and rationale only - the M-id rosters live in the
 phase-5 cell above and the lifecycle facts in the ledger. Every PR lands on
 main under the section 2 cut window with `__version__` unchanged; each flips
-its rows terminal and names them in its changelog fragment.
+its rows terminal and names them in its changelog fragment (3a excepted: a
+tests-first PR that flips no rows).
 
 0. Spec/ledger refresh (the PR that wrote this subsection): the cut-window
    rule, the roster/M-080/M-008 amendments, the [M-150] locator, the id
@@ -1017,7 +1018,16 @@ its rows terminal and names them in its changelog fragment.
    [M-045]/[M-115]; `_deprecation.py` stays (section 8 rule 11 corollary).
 3. Fit-time `aggregate=`/`balance_e=` removals, plus the runtime messages that
    still steer users to the deprecated route (readers recorded in the rows'
-   `code_refs`) and the benchmark/guide sweep.
+   `code_refs`) and the benchmark/guide sweep. Split three ways (2026-10-20)
+   so each piece fits the AI reviewer's input cap:
+   - 3a - tests-first migration of the test and benchmark suites to the
+     post-fit containers; flips no rows, by design (it carries the test bulk
+     so 3b/3c fit the reviewer) (SHIPPED, PR 3a).
+   - 3b - the mechanical removals on the estimators whose post-fit route
+     already exists (CallawaySantAnna, EfficientDiD, StackedDiD, dCDH, HAD,
+     the HAD pretest workflow).
+   - 3c - the ImputationDiD/TwoStageDiD/ContinuousDiD removals plus the
+     bootstrapped-fit relay (section 6 amendment).
 4. The three merges: the MultiPeriodDiD class, `MultiPeriodDiDResults`,
    `PeriodEffect` and `EventStudy`; the QDiD class and `QDiDResults`; the
    `StaggeredTripleDifference` CLASS and `SDDD` - but NOT