perf(callaway): stop materializing per-cell unit-label arrays in influence_func_info

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
 

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_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_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_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_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,
                 }