perf(callaway): fancy-+= per-(g,t) IF scatters; drop dead _compute_aggregated_se

CHANGELOG.md

@@ -398,6 +398,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   (GMM-variance-dominated). Accumulation-order numerics documented in
   `docs/methodology/REGISTRY.md` (both estimator sections + "Absorbed Fixed Effects").
 
+### Performance
+- **`CallawaySantAnna` per-(g,t) IF scatters converted from `np.add.at` to fancy `+=`**
+  (`staggered.py::_cluster_robust_se_from_per_gt_if` — runs once per (g,t) cell when
+  `cluster=` is set — and the general combined-IF assembly path in
+  `staggered_aggregation.py`). The index arrays are duplicate-free by construction at
+  every producer (`np.where` on disjoint masks — the invariant the aggregation fast
+  path already relies on), so the fancy scatter is bit-identical to the unbuffered
+  `np.add.at` while avoiding its 5-20x per-element overhead. Also removed the dead
+  `_compute_aggregated_se` (zero callers; superseded by
+  `_compute_aggregated_se_with_wif` since the WIF adjustment landed).
+
 ## [3.6.2] - 2026-07-03
 
 ### Added

TODO.md

@@ -58,9 +58,7 @@ 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
 

diff_diff/staggered.py

@@ -162,11 +162,16 @@ def _cluster_robust_se_from_per_gt_if(
         and (control_idx.max(initial=-1) >= n or control_idx.min(initial=0) < 0)
     ):
         return None
+    # Index arrays are unique within each cell by construction at every
+    # producer (np.where on disjoint masks), so fancy += is exact — same
+    # scatter contract as staggered_aggregation._combined_if_fast, without
+    # np.add.at's unbuffered-ufunc overhead (runs once per (g,t) cell when
+    # cluster= is set).
     psi_per_index = np.zeros(n)
     if treated_idx.size:
-        np.add.at(psi_per_index, treated_idx, treated_inf)
+        psi_per_index[treated_idx] += treated_inf
     if control_idx.size:
-        np.add.at(psi_per_index, control_idx, control_inf)
+        psi_per_index[control_idx] += control_inf
     # Route through the shared survey helper so the per-cell variance
     # gets the same G/(G-1) finite-sample correction, PSU centering,
     # FPC handling, and lonely-PSU/G<2→NaN behavior as overall +

diff_diff/staggered_aggregation.py

@@ -169,72 +169,6 @@ def _aggregate_simple(
 
         return overall_att, overall_se, effective_df
 
-    def _compute_aggregated_se(
-        self,
-        gt_pairs: List[Tuple[Any, Any]],
-        weights: np.ndarray,
-        influence_func_info: Dict,
-        n_units: Optional[int] = None,
-    ) -> float:
-        """
-        Compute standard error using influence function aggregation.
-
-        This properly accounts for covariances across (g,t) pairs by
-        aggregating unit-level influence functions:
-
-            ψ_i(overall) = Σ_{(g,t)} w_(g,t) × ψ_i(g,t)
-            Var(overall) = (1/n) Σ_i [ψ_i]²
-
-        This matches R's `did` package analytical SE formula.
-
-        Parameters
-        ----------
-        n_units : int, optional
-            Size of the canonical index space (len(precomputed['all_units'])).
-            When provided, influence function indices (treated_idx, control_idx)
-            index directly into this space, eliminating dict lookups.
-        """
-        if not influence_func_info:
-            return 0.0
-
-        if n_units is None:
-            # Fallback: infer size from influence function info
-            max_idx = 0
-            for g, t in gt_pairs:
-                if (g, t) in influence_func_info:
-                    info = influence_func_info[(g, t)]
-                    if len(info["treated_idx"]) > 0:
-                        max_idx = max(max_idx, info["treated_idx"].max())
-                    if len(info["control_idx"]) > 0:
-                        max_idx = max(max_idx, info["control_idx"].max())
-            n_units = max_idx + 1
-
-        if n_units == 0:
-            return 0.0
-
-        # Aggregate influence functions across (g,t) pairs
-        psi_overall = np.zeros(n_units)
-
-        for j, (g, t) in enumerate(gt_pairs):
-            if (g, t) not in influence_func_info:
-                continue
-
-            info = influence_func_info[(g, t)]
-            w = weights[j]
-
-            # Vectorized influence function aggregation using index arrays
-            treated_idx = info["treated_idx"]
-            if len(treated_idx) > 0:
-                np.add.at(psi_overall, treated_idx, w * info["treated_inf"])
-
-            control_idx = info["control_idx"]
-            if len(control_idx) > 0:
-                np.add.at(psi_overall, control_idx, w * info["control_inf"])
-
-        # Compute variance: Var(θ̄) = (1/n) Σᵢ ψᵢ²
-        variance = np.sum(psi_overall**2)
-        return np.sqrt(variance)
-
     @staticmethod
     def _get_agg_cache(precomputed: "PrecomputedData") -> Dict[str, Any]:
         """
@@ -474,8 +408,10 @@ def _compute_combined_influence_function(
         # global_unit_to_idx and precomputed["unit_to_idx"] are None, so the
         # identity guard alone would spuriously pass). Anything not exactly
         # matched (direct callers with foreign index maps, size mismatches,
-        # non-numeric cohorts) falls through to the general path below,
-        # which is preserved unchanged.
+        # non-numeric cohorts) falls through to the general path below
+        # (same mathematical contract; its IF scatter uses the fancy-+=
+        # form, bit-identical to np.add.at on the producers' duplicate-free
+        # index arrays).
         if precomputed is not None and n_global_units is not None:
             _fast_ok = False
             if _is_rcs:
@@ -571,14 +507,17 @@ def _compute_combined_influence_function(
             info = influence_func_info[(g, t)]
             w = weights[j]
 
-            # Vectorized influence function aggregation using precomputed index arrays
+            # Vectorized IF aggregation using precomputed index arrays. Index
+            # arrays are unique within each cell by construction at every
+            # producer (np.where on disjoint masks), so fancy += is exact —
+            # same scatter contract as _combined_if_fast.
             treated_idx = info["treated_idx"]
             if len(treated_idx) > 0:
-                np.add.at(psi_standard, treated_idx, w * info["treated_inf"])
+                psi_standard[treated_idx] += w * info["treated_inf"]
 
             control_idx = info["control_idx"]
             if len(control_idx) > 0:
-                np.add.at(psi_standard, control_idx, w * info["control_inf"])
+                psi_standard[control_idx] += w * info["control_inf"]
 
         # Build unit-group array: normalize iterator to (idx, uid) pairs
         unit_groups_array = np.full(n_units, -1, dtype=np.float64)

docs/methodology/REGISTRY.md

@@ -610,8 +610,10 @@ Aggregations:
   (see "Absorbed Fixed Effects with Survey Weights"). Units whose cohort is not
   among the keeper (g,t) groups get an exact 0 WIF contribution (the dense form
   realized the same value through cancelling terms). The pre-rewrite general path
-  is preserved verbatim as the fallback for direct callers with foreign index maps,
-  non-numeric cohort dtypes, or absent precomputed structures.
+  remains the fallback for direct callers with foreign index maps, non-numeric
+  cohort dtypes, or absent precomputed structures; its per-cell IF scatter uses
+  fancy `+=` (bit-identical to the prior `np.add.at` — index arrays are
+  duplicate-free by construction at every producer), same mathematical contract.
 - Block structure preserves within-unit correlation
 - Simultaneous confidence bands (`cband=True`, default): Uses sup-t bootstrap to compute
   a uniform critical value across event times, controlling family-wise error rate. Matches

tests/test_se_accuracy.py

@@ -364,7 +364,8 @@ class TestSEFormulaComparison:
         V <- Matrix::t(inffunc) %*% inffunc / (n)
         se <- sqrt(Matrix::diag(V) / n)
 
-    Python formula (from staggered.py _compute_aggregated_se):
+    Python formula (staggered_aggregation.py _se_from_psi, consumed by
+    _compute_aggregated_se_with_wif):
         variance = np.sum(psi_overall ** 2)
         return np.sqrt(variance)