From cbe40cadf21ad4e1c81f03e70b4286735be11677 Mon Sep 17 00:00:00 2001 From: igerber Date: Tue, 7 Jul 2026 15:08:04 -0400 Subject: [PATCH] perf(callaway): fancy-+= per-(g,t) IF scatters; drop dead _compute_aggregated_se MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-(g,t) cluster-SE helper (_cluster_robust_se_from_per_gt_if) and the general combined-IF assembly scattered per-cell influence functions with np.add.at (unbuffered ufunc, 5-20x slower than fancy +=). Index arrays are duplicate-free by construction at every producer (np.where on disjoint masks — the invariant _combined_if_fast already relies on), so the fancy scatter is bit-identical. Also removes the dead _compute_aggregated_se (zero in-package callers; superseded by _compute_aggregated_se_with_wif), updates the stale test_se_accuracy docstring pointer, and refreshes the REGISTRY/dispatch-comment wording about the general fallback path. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 11 ++++ TODO.md | 2 - diff_diff/staggered.py | 9 +++- diff_diff/staggered_aggregation.py | 81 ++++-------------------------- docs/methodology/REGISTRY.md | 6 ++- tests/test_se_accuracy.py | 3 +- 6 files changed, 34 insertions(+), 78 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 60e5c940..eb50ecbf 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/TODO.md b/TODO.md index c35b548f..f37cb984 100644 --- a/TODO.md +++ b/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 --git a/diff_diff/staggered.py b/diff_diff/staggered.py index fee3878f..4d9742bc 100644 --- a/diff_diff/staggered.py +++ b/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 --git a/diff_diff/staggered_aggregation.py b/diff_diff/staggered_aggregation.py index bc36bf56..ff519688 100644 --- a/diff_diff/staggered_aggregation.py +++ b/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) diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 98fb462d..34351d0e 100644 --- a/docs/methodology/REGISTRY.md +++ b/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 diff --git a/tests/test_se_accuracy.py b/tests/test_se_accuracy.py index e5687658..144695a1 100644 --- a/tests/test_se_accuracy.py +++ b/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)