perf(staggered): O(n_units) CS aggregation IF assembly + R did 2.5.1 yardstick

Rewrite CallawaySantAnna's combined influence-function assembly (also
inherited by StaggeredTripleDifference aggregation) - previously 56-85% of
analytical fit time at scale:

- Per-fit cohort tables (np.unique + np.bincount) cached on the precomputed
  structures with array-identity validation (safe against triple-diff's
  shallow-copy + zeroed-cohort aggregation pattern)
- Closed-form WIF (wif_i = w_i * (E(c_i)/S - K(c_i)*d/S^2)) replacing dense
  (n_units x n_gt) indicator/outer-product matrices - algebraically identical,
  documented in REGISTRY
- Fancy-index scatter replacing np.add.at (per-cell index arrays are
  duplicate-free by construction; invariant documented at all IF producers)
- Pre-rewrite general path preserved verbatim as fallback for direct callers

Point estimates bit-identical; aggregated SEs <=5e-16 relative (drift-bound
frozen-copy tests at rtol=0 atol=1e-9 across panel/survey/RCS/unbalanced/
triple-diff/inf-coded fixtures). Analytical fits 2.3-6.3x faster at 2-5M rows,
bootstrap fits 1.4-1.6x, RCS peak memory -55 to -68%. 3-11x faster than R did
2.5.1 at equal work (single- and multi-core R; benchmarks/R/benchmark_did.R
extended with bootstrap/covariate/parallelism flags for the yardstick).

Golden aggregated-SE assertions enabled for dr scenarios (match R fixtures at
1e-12..7e-6). Pre-existing reg-method aggregated-SE gap vs R fixtures (3-20%,
ATTs exact) documented in TODO.md for follow-up investigation.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01X4AzrFUMqJxcUumSH31mSr

Author: igerber

CHANGELOG.md

@@ -53,6 +53,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   supersedes it.
 
 ### Changed
+- **CallawaySantAnna aggregation SE assembly rewritten to O(n_units)** (also inherited by
+  `StaggeredTripleDifference` aggregation). The combined influence-function assembly behind
+  simple/event-study/group aggregated SEs — previously ~56-85% of analytical fit time at
+  scale — replaces its per-aggregation-target full-DataFrame cohort scans, per-unit Python
+  loops, and dense `(n_units × n_gt)` weight-influence-function matrices with per-fit cohort
+  tables and a closed-form WIF (algebraically identical; details in REGISTRY). Measured
+  (medians of 3, Apple M4 Max, Rust backend): analytical fits at 100k units × 20 periods
+  (2M rows) 1.32→0.21s no-covariate (6.3x), 2.49→1.08s 5-covariate DR (2.3x); long panels
+  40p×10 cohorts 9.09→3.91s and 60p×15 cohorts 9.86→4.34s (2.3x); bootstrap-999 fits
+  1.4-1.6x (the remaining draw-loop matmul is unchanged); repeated cross-sections
+  4.17→1.50s no-covariate (2.8x) with peak memory 6.4→2.1 GB (-67%; the dense WIF matrices
+  were observation-scale in RCS mode). Point estimates are bit-identical; aggregated SEs
+  move only at floating-point reassociation level (measured ≤5e-16 relative, drift-bound
+  tested at 1e-9 against a frozen copy of the prior implementation). For scale context,
+  a full `fit(aggregate="all")` at 2M rows now runs 3-11x faster than R `did` 2.5.1
+  (equal work, analytical or bootstrap-999, single- or multi-core R; R's `pl`/`cores`
+  parallelism is BLAS-bound on this path so both R arms time identically).
 - **ImputationDiD/TwoStageDiD demeaning modernized onto the shared MAP engine.** The
   private per-estimator pandas `_iterative_demean` loops (rebuilt a
   `pd.Series.groupby().transform()` hash table every alternating-projection iteration)

TODO.md

@@ -32,6 +32,7 @@ The `Origin` column (Actionable tables) and the `PR` column (Deferred tables) bo
 | `ContinuousDiD` CGBS-2024 extensions. (a) `covariates=` kwarg — **DONE (reg/dr)**; remaining: (b) discrete-treatment saturated regression (integer dose currently warned, not routed to per-level coefficients); (c) lowest-dose-as-control per Remark 3.1 when `P(D=0)=0`. Also deferred from the covariate work: `estimation_method="ipw"` on the dose curve (scalar-adjustment / degenerate — documented `NotImplementedError`), and `covariates=` × `survey_design=` (weighted OR + weighted nuisance IF). | `continuous_did.py` | CGBS-2024 | Heavy | Low |
 | `ImputationDiD` LOO conservative-variance refinement (BJS 2024 Supp. Appendix A.9) — a finite-sample improvement to the auxiliary-model residuals reducing overfit of `tau_tilde_g` to `epsilon`. Asymptotic Theorem-3 variance is implemented and matches R `didimputation` (which also omits LOO by default). | `imputation.py` | imputation-validation | Mid | Low |
 | `TwoWayFixedEffects(vcov_type in {hc2, hc2_bm})` with replicate-weight designs raises `NotImplementedError` (`twfe.py:~233`). The replicate path re-demeans per replicate, which doesn't compose with the full-dummy HC2/HC2-BM build — a correct impl needs per-replicate full-dummy refit. Workaround: `hc1` for replicate-weight CR1. | `twfe.py::fit` | follow-up | Heavy | Low |
+| `CallawaySantAnna` reg-method aggregated SEs sit 3-20% from the R golden fixtures (`two_period` simple/group/dynamic, `dynamic_effects` dynamic) while the ATTs match at 1e-11 and the dr-method aggregated SEs match at ≤7e-6. Pre-existing (byte-identical deltas on the pre-fast-path tree); golden aggregated-SE assertions were therefore enabled for dr scenarios only (`test_golden_simple_aggregation_se`, `test_golden_event_study_aggregation_se`). Investigate the reg-path per-cell IF vs R `DRDID::reg_did_panel`'s. | `staggered.py::_outcome_regression`, `tests/test_csdid_ported.py` | CS-scaling | Mid | Medium |
 | TWFE's HC2/HC2-BM inline full-dummy build (`twfe.py:280-315`) duplicates the dummy-construction logic in `DifferenceInDifferences(fixed_effects=...)` (`estimators.py:478-486`). Extract a shared helper, or delegate TWFE's HC2/HC2-BM path to DiD's `fixed_effects=` branch (with TWFE-specific cluster-default threading), to reduce drift risk on FE naming / survey behavior / result-surface conventions. Substantive refactor — touches both estimators. | `twfe.py::fit`, `estimators.py::DifferenceInDifferences.fit` | follow-up | Heavy | Low |
 
 ### Performance
@@ -53,6 +54,9 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m
 | Multiplier-bootstrap weight chunking (CallawaySantAnna, EfficientDiD, and HAD — all wired through `diff_diff/bootstrap_chunking.py`) covers the **unstratified** survey-PSU generation (the default unit-level bootstrap — `cluster=None`, equivalently `cluster="unit"` — the large-`n_units` OOM case). Remaining gap: the **stratified** survey-PSU generator (`generate_survey_multiplier_weights_batch`, per-stratum + lonely-PSU pooling + FPC) still materializes the full `(n_bootstrap × n_psu)` matrix (consumed via sliced blocks). Stratified designs have few PSUs so this rarely OOMs; tile per-stratum generation over draws (each stratum's draws are independent → contiguous draw-blocks reproduce the stream bit-identically) if a large-PSU stratified design hits memory. | `diff_diff/bootstrap_chunking.py::iter_survey_multiplier_weight_blocks` | follow-up | Mid | Low |
 | Migrate `spillover._iterative_fe_subset` onto the shared `diff_diff.utils._iterative_fe_solve` (same Gauss-Seidel-on-codes recursion, specialized to a masked Butts subsample; ImputationDiD/TwoStageDiD already route through the shared helper — this takes the FE-solver copy count from 2 to 1). Preserve the SpilloverDiD positive-weight/NaN-FE REGISTRY contract and `_FE_ITER_MAX=100` budget (or align it to 10k with a REGISTRY note). | `spillover.py` | demean-modernization | Mid | Low |
 | Adopt `snap_absorbed_regressors` (FE-spanned regressor two-stage snap + LSMR confirmation) on the ImputationDiD lead-indicator path — lead columns are the most plausible FE-spanned regressors, and the truncated-MAP-iterate exposure documented at `utils.py` applies; today only `solve_ols` rank detection guards it. Behavior change beyond the demean-modernization refactor, so deferred from that PR. | `imputation.py::_compute_lead_coefficients` | demean-modernization | Mid | Low |
+| `_cluster_robust_se_from_per_gt_if` scatters per-cell IFs with `np.add.at` (unbuffered ufunc, 5-20x slower than fancy `+=`; same slow-scatter class the aggregation fast path fixed). Index arrays come from `np.where` over boolean masks (duplicate-free), so per-cell `psi[idx] += vals` is exact. Runs once per (g,t) cell when `cluster=` is set — noticeable at many-cell fits. | `staggered.py::_cluster_robust_se_from_per_gt_if` | CS-scaling | Quick | Low |
+| Per-cell `treated_units`/`control_units` label arrays (`all_units[positions]`, ~O(n_control) alloc per (g,t) cell) are consumed only by the precomputed-None fallback of the combined-IF assembly, which no in-package caller reaches — build them lazily (or drop from the IF-info dict) to cut per-cell allocation at high cell counts. | `staggered.py::_compute_att_gt_fast` | CS-scaling | Mid | Low |
+| `_compute_aggregated_se` is dead code (zero in-package callers; superseded by `_compute_aggregated_se_with_wif`) — remove it, or fold its docstring into the WIF variant. Its `np.add.at` scatter also predates the fancy-`+=` convention. | `staggered_aggregation.py::_compute_aggregated_se` | CS-scaling | Quick | Low |
 
 ### Testing / docs
 

benchmarks/R/benchmark_did.R

@@ -2,7 +2,17 @@
 # Benchmark: Callaway-Sant'Anna Estimator (R `did` package)
 #
 # Usage:
-#   Rscript benchmark_did.R --data path/to/data.csv --output path/to/results.json
+#   Rscript benchmark_did.R --data path/to/data.csv --output path/to/results.json \
+#     [--method dr|ipw|reg] [--control-group nevertreated|notyettreated] \
+#     [--xformla "~ x1 + x2"] [--bstrap true|false] [--biters N] \
+#     [--cband true|false] [--pl true|false] [--cores N] \
+#     [--faster-mode true|false]
+#
+# Defaults reproduce the historical behavior (analytical SEs, no covariates,
+# single core), so existing callers (benchmarks/run_benchmarks.py) are
+# unaffected. The optional flags exist for the R-yardstick arms: bootstrap
+# inference (bstrap/biters/cband applied at BOTH att_gt and aggte), covariate
+# formulas, and did's parallel processing (pl/cores).
 
 library(did)
 library(jsonlite)
@@ -11,35 +21,73 @@ library(data.table)
 # Parse command line arguments
 args <- commandArgs(trailingOnly = TRUE)
 
+parse_bool <- function(x, flag) {
+  v <- tolower(x)
+  if (v %in% c("true", "t", "1", "yes")) return(TRUE)
+  if (v %in% c("false", "f", "0", "no")) return(FALSE)
+  stop(sprintf("Invalid boolean for %s: '%s' (use true/false)", flag, x))
+}
+
 parse_args <- function(args) {
   result <- list(
     data = NULL,
     output = NULL,
     method = "dr",
-    control_group = "nevertreated"
+    control_group = "nevertreated",
+    xformla = NULL,
+    bstrap = FALSE,
+    biters = 1000L,
+    cband = FALSE,
+    pl = FALSE,
+    cores = 1L,
+    faster_mode = NULL  # NULL -> use did's own default for this version
   )
 
   i <- 1
   while (i <= length(args)) {
-    if (args[i] == "--data") {
-      result$data <- args[i + 1]
-      i <- i + 2
-    } else if (args[i] == "--output") {
-      result$output <- args[i + 1]
-      i <- i + 2
-    } else if (args[i] == "--method") {
-      result$method <- args[i + 1]
-      i <- i + 2
-    } else if (args[i] == "--control-group") {
-      result$control_group <- args[i + 1]
-      i <- i + 2
+    flag <- args[i]
+    if (i + 1 > length(args) && flag != "") {
+      stop(sprintf("Missing value for flag: %s", flag))
+    }
+    val <- args[i + 1]
+    if (flag == "--data") {
+      result$data <- val
+    } else if (flag == "--output") {
+      result$output <- val
+    } else if (flag == "--method") {
+      result$method <- val
+    } else if (flag == "--control-group") {
+      result$control_group <- val
+    } else if (flag == "--xformla") {
+      result$xformla <- as.formula(val)
+    } else if (flag == "--bstrap") {
+      result$bstrap <- parse_bool(val, flag)
+    } else if (flag == "--biters") {
+      result$biters <- as.integer(val)
+    } else if (flag == "--cband") {
+      result$cband <- parse_bool(val, flag)
+    } else if (flag == "--pl") {
+      result$pl <- parse_bool(val, flag)
+    } else if (flag == "--cores") {
+      result$cores <- as.integer(val)
+    } else if (flag == "--faster-mode") {
+      result$faster_mode <- parse_bool(val, flag)
     } else {
-      i <- i + 1
+      # Unknown flags used to be silently skipped, which turned typos into
+      # silent default runs. Fail loudly instead.
+      stop(sprintf("Unknown flag: %s", flag))
     }
+    i <- i + 2
   }
 
   if (is.null(result$data) || is.null(result$output)) {
-    stop("Usage: Rscript benchmark_did.R --data <path> --output <path> [--method dr|ipw|reg] [--control-group nevertreated|notyettreated]")
+    stop("Usage: Rscript benchmark_did.R --data <path> --output <path> [--method dr|ipw|reg] [--control-group nevertreated|notyettreated] [--xformla '~ x1 + x2'] [--bstrap true|false] [--biters N] [--cband true|false] [--pl true|false] [--cores N] [--faster-mode true|false]")
+  }
+  if (is.na(result$biters) || result$biters <= 0) {
+    stop(sprintf("--biters must be a positive integer, got '%s'", result$biters))
+  }
+  if (is.na(result$cores) || result$cores <= 0) {
+    stop(sprintf("--cores must be a positive integer, got '%s'", result$cores))
   }
 
   return(result)
@@ -66,28 +114,45 @@ message(sprintf("Never-treated units (first_treat=Inf): %d", sum(is.infinite(dat
 message("Running Callaway-Sant'Anna estimation...")
 start_time <- Sys.time()
 
-out <- att_gt(
+att_gt_args <- list(
   yname = "outcome",
   tname = "time",
   idname = "unit",
   gname = "first_treat",
-  xformla = NULL,
+  xformla = config$xformla,
   data = data,
   est_method = config$method,
   control_group = config$control_group,
-  bstrap = FALSE,  # Use analytical SEs for speed
-  cband = FALSE
+  bstrap = config$bstrap,
+  biters = config$biters,
+  cband = config$cband,
+  pl = config$pl,
+  cores = config$cores
 )
+# faster_mode exists in recent did releases only; pass it only when both
+# requested and supported, so the script still runs on older installs.
+if (!is.null(config$faster_mode)) {
+  if ("faster_mode" %in% names(formals(att_gt))) {
+    att_gt_args$faster_mode <- config$faster_mode
+  } else {
+    message("faster_mode not supported by this did version; ignoring flag")
+  }
+}
+out <- do.call(att_gt, att_gt_args)
 
 estimation_time <- as.numeric(difftime(Sys.time(), start_time, units = "secs"))
 
-# Aggregate results
+# Aggregate results (same inference mode as att_gt so the timing arm does
+# equal work end-to-end)
 message("Aggregating results...")
 agg_start <- Sys.time()
 
-agg_simple <- aggte(out, type = "simple", bstrap = FALSE, cband = FALSE)
-agg_dynamic <- aggte(out, type = "dynamic", bstrap = FALSE, cband = FALSE)
-agg_group <- aggte(out, type = "group", bstrap = FALSE, cband = FALSE)
+agg_simple <- aggte(out, type = "simple", bstrap = config$bstrap,
+                    biters = config$biters, cband = FALSE)
+agg_dynamic <- aggte(out, type = "dynamic", bstrap = config$bstrap,
+                     biters = config$biters, cband = config$cband)
+agg_group <- aggte(out, type = "group", bstrap = config$bstrap,
+                   biters = config$biters, cband = FALSE)
 
 aggregation_time <- as.numeric(difftime(Sys.time(), agg_start, units = "secs"))
 total_time <- estimation_time + aggregation_time
@@ -131,13 +196,22 @@ results <- list(
     total_seconds = total_time
   ),
 
-  # Metadata
+  # Metadata (records the full inference/parallelism config so every
+  # yardstick number is reproducible from its own artifact)
   metadata = list(
     r_version = R.version.string,
     did_version = as.character(packageVersion("did")),
     n_units = length(unique(data$unit)),
     n_periods = length(unique(data$time)),
-    n_obs = nrow(data)
+    n_obs = nrow(data),
+    xformla = if (is.null(config$xformla)) NULL else deparse(config$xformla),
+    bstrap = config$bstrap,
+    biters = if (config$bstrap) config$biters else NULL,
+    cband = config$cband,
+    pl = config$pl,
+    cores = config$cores,
+    faster_mode = if (is.null(config$faster_mode)) "did-default" else config$faster_mode,
+    blas = tryCatch(sessionInfo()$BLAS, error = function(e) NULL)
   )
 )
 

diff_diff/staggered.py

@@ -1008,6 +1008,10 @@ def _compute_att_gt_fast(
         # Package influence function info with index arrays (positions into
         # precomputed['all_units']) for O(1) downstream lookups instead of
         # O(n) Python dict lookups.
+        # INVARIANT: treated_idx/control_idx are np.where over boolean masks,
+        # so each is duplicate-free - the aggregation fast path relies on this
+        # to scatter with fancy `psi[idx] += vals` (see
+        # staggered_aggregation._combined_if_fast).
         n_t = int(n_treated)
         all_units = precomputed["all_units"]
         treated_positions = np.where(treated_valid)[0]
@@ -1576,6 +1580,9 @@ def _compute_all_att_gt_covariate_reg(
                     "skip_reason": None,
                 }
 
+                # INVARIANT: np.where over boolean masks -> duplicate-free
+                # index arrays (fancy-+= scatter contract, see
+                # staggered_aggregation._combined_if_fast).
                 all_units = precomputed["all_units"]
                 treated_positions = np.where(treated_valid)[0]
                 control_positions = np.where(control_valid)[0]
@@ -3511,7 +3518,11 @@ def _compute_att_gt_rc(
             )
 
         # Build influence function info
-        # For RCS, treated_idx/control_idx combine obs from BOTH periods
+        # For RCS, treated_idx/control_idx combine obs from BOTH periods.
+        # INVARIANT: the two period masks are disjoint (obs_time == t vs
+        # == base period), so the concatenated index arrays stay
+        # duplicate-free (fancy-+= scatter contract, see
+        # staggered_aggregation._combined_if_fast).
         treated_idx = np.concatenate([np.where(treated_t)[0], np.where(treated_s)[0]])
         control_idx = np.concatenate([np.where(control_t)[0], np.where(control_s)[0]])