From a8ca356583f851b187aecccf1c2fedbdcb601b65 Mon Sep 17 00:00:00 2001 From: igerber Date: Tue, 7 Jul 2026 16:15:13 -0400 Subject: [PATCH] perf(callaway): stop materializing per-cell unit-label arrays in influence_func_info MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every (g,t) cell's internal IF record carried treated_units / control_units label arrays (all_units[positions], two O(n_control) allocations per cell) that no in-package code path ever read — the only consumers were the precomputed-less fallbacks of the combined-IF assembly and the multiplier bootstrap, which every in-package caller bypasses by threading precomputed structures. Labels stay recoverable as all_units[treated_idx]. The two fallbacks now raise an actionable ValueError when reached without label arrays instead of silently KeyErroring. Applies to CallawaySantAnna (panel, vectorized, RCS, reference cells) and StaggeredTripleDifference producers. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 11 +++++++++++ TODO.md | 1 - diff_diff/staggered.py | 20 ++++++-------------- diff_diff/staggered_aggregation.py | 22 ++++++++++++++++++++-- diff_diff/staggered_bootstrap.py | 22 +++++++++++++++++++--- diff_diff/staggered_triple_diff.py | 2 -- 6 files changed, 56 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f7c8dab..db0b1b27 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -454,6 +454,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 memory-traffic cleanup, not a headline speedup). The `zero_mask` abs scan is retained (correctness); the remaining conditional-path lever is the sieve/nuisance stage (see TODO). +- **`CallawaySantAnna` / `StaggeredTripleDifference` per-cell unit-label arrays no longer + materialized.** Every (g,t) cell's internal influence-function record carried + `treated_units` / `control_units` label arrays (`all_units[positions]`, two O(n_control) + allocations per cell) that no in-package code path ever read: the only consumers were the + `precomputed`-less fallbacks of the combined-IF assembly and the multiplier bootstrap, + which every in-package caller bypasses by threading the precomputed structures. Labels + remain recoverable as `all_units[treated_idx]`. The two fallbacks now raise an actionable + `ValueError` if a direct caller reaches them without label arrays (previously they would + have been reached only by external callers hand-building the internal dict). No public + API or numerical change — the dict is internal (`_influence_func_info` whitebox surface + keeps `treated_idx`/`control_idx`/`treated_inf`/`control_inf`). ## [3.6.2] - 2026-07-03 diff --git a/TODO.md b/TODO.md index f6b1024a..e1b66978 100644 --- a/TODO.md +++ b/TODO.md @@ -56,7 +56,6 @@ 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 | -| 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 | ### Testing / docs diff --git a/diff_diff/staggered.py b/diff_diff/staggered.py index 4d9742bc..2e1c367d 100644 --- a/diff_diff/staggered.py +++ b/diff_diff/staggered.py @@ -1100,20 +1100,22 @@ 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. + # O(n) Python dict lookups. Unit-LABEL arrays are deliberately NOT + # materialized (labels are always recoverable as all_units[idx]); + # only the precomputed-None fallbacks of the combined-IF assembly + # and the multiplier bootstrap ever consumed them — no in-package + # caller reaches those — so dropping them cuts two O(n_control) + # allocations per (g,t) cell. # 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] control_positions = np.where(control_valid)[0] inf_func_info = { "treated_idx": treated_positions, "control_idx": control_positions, - "treated_units": all_units[treated_positions], - "control_units": all_units[control_positions], "treated_inf": inf_func[:n_t], "control_inf": inf_func[n_t:], } @@ -1292,14 +1294,11 @@ def _compute_all_att_gt_vectorized( gte_entry["survey_weight_sum"] = sw_sum group_time_effects[(g, t)] = gte_entry - all_units = precomputed["all_units"] treated_positions = np.where(treated_valid)[0] control_positions = np.where(control_valid)[0] inf_info_gt = { "treated_idx": treated_positions, "control_idx": control_positions, - "treated_units": all_units[treated_positions], - "control_units": all_units[control_positions], "treated_inf": inf_treated, "control_inf": inf_control, } @@ -1724,14 +1723,11 @@ def _compute_all_att_gt_covariate_reg( # 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] inf_info_gt = { "treated_idx": treated_positions, "control_idx": control_positions, - "treated_units": all_units[treated_positions], - "control_units": all_units[control_positions], "treated_inf": inf_treated, "control_inf": inf_control, } @@ -2609,8 +2605,6 @@ def fit( "control_idx": np.array([], dtype=int), "treated_inf": np.array([]), "control_inf": np.array([]), - "treated_units": np.array([]), - "control_units": np.array([]), } # Compute overall ATT (simple aggregation) @@ -4002,8 +3996,6 @@ def _compute_att_gt_rc( inf_func_info = { "treated_idx": treated_idx, "control_idx": control_idx, - "treated_units": treated_idx, # For RCS, obs indices = "units" - "control_units": control_idx, "treated_inf": inf_func_all[:n_treated_combined], "control_inf": inf_func_all[n_treated_combined:], } diff --git a/diff_diff/staggered_aggregation.py b/diff_diff/staggered_aggregation.py index ff519688..96338abc 100644 --- a/diff_diff/staggered_aggregation.py +++ b/diff_diff/staggered_aggregation.py @@ -442,12 +442,30 @@ def _compute_combined_influence_function( n_units = n_global_units all_units = None # caller already has the unit list else: + # Local-units fallback for direct callers without precomputed / + # global unit ids. It needs per-cell unit-LABEL arrays, which + # in-package fits stopped materializing (v3.8 per-cell + # allocation shave — labels are always all_units[idx], and every + # in-package caller threads global_unit_to_idx + n_global_units, + # so this branch is unreachable from a fit). Direct callers must + # either thread the global ids or supply label arrays. all_units_set: Set[Any] = set() for g, t in gt_pairs: if (g, t) in influence_func_info: info = influence_func_info[(g, t)] - all_units_set.update(info["treated_units"]) - all_units_set.update(info["control_units"]) + t_units = info.get("treated_units") + c_units = info.get("control_units") + if t_units is None or c_units is None: + raise ValueError( + "Combined-IF assembly without precomputed/global unit " + "ids requires per-cell 'treated_units'/'control_units' " + "label arrays in influence_func_info; in-package fits " + "no longer materialize them. Pass global_unit_to_idx " + "and n_global_units (as all in-package callers do), " + "or add the label arrays to your influence_func_info." + ) + all_units_set.update(t_units) + all_units_set.update(c_units) if not all_units_set: return np.zeros(0), [] diff --git a/diff_diff/staggered_bootstrap.py b/diff_diff/staggered_bootstrap.py index 9dc66d56..1b9acf8c 100644 --- a/diff_diff/staggered_bootstrap.py +++ b/diff_diff/staggered_bootstrap.py @@ -210,11 +210,27 @@ def _run_multiplier_bootstrap( n_units = precomputed.get("canonical_size", len(all_units)) unit_to_idx = precomputed["unit_to_idx"] # None for RCS else: - # Fallback: collect units from influence functions + # Fallback: collect units from influence functions. Needs the + # per-cell unit-LABEL arrays, which in-package fits stopped + # materializing (v3.8 per-cell allocation shave) — every + # in-package caller threads `precomputed`, so this branch is + # unreachable from a fit. Direct callers must thread + # `precomputed` or supply label arrays. all_units_set = set() for (g, t), info in influence_func_info.items(): - all_units_set.update(info["treated_units"]) - all_units_set.update(info["control_units"]) + t_units = info.get("treated_units") + c_units = info.get("control_units") + if t_units is None or c_units is None: + raise ValueError( + "Multiplier bootstrap without `precomputed` requires " + "per-cell 'treated_units'/'control_units' label arrays " + "in influence_func_info; in-package fits no longer " + "materialize them. Thread `precomputed` (as all " + "in-package callers do), or add the label arrays to " + "your influence_func_info." + ) + all_units_set.update(t_units) + all_units_set.update(c_units) all_units = sorted(all_units_set) # Use global N from dataframe when available n_units = ( diff --git a/diff_diff/staggered_triple_diff.py b/diff_diff/staggered_triple_diff.py index 18e57ab9..f3492745 100644 --- a/diff_diff/staggered_triple_diff.py +++ b/diff_diff/staggered_triple_diff.py @@ -504,8 +504,6 @@ def fit( influence_func_info[(g, t)] = { "treated_idx": treated_idx, "control_idx": control_idx, - "treated_units": all_units[treated_idx], - "control_units": all_units[control_idx], "treated_inf": treated_inf, "control_inf": control_inf, }