perf(efficient_did): skip _ridge_solve_weights fancy-index copy when no row is zero-masked

CHANGELOG.md

@@ -432,6 +432,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `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).
+- **`EfficientDiD` `_ridge_solve_weights` Python-prep shave.** When no row of the Omega*
+  stack is zero-masked (the common case), the batched ridge solve now consumes the stack
+  directly instead of paying the `omega_stack[rest]` fancy-index copy, and returns the
+  solved weights without the tail scatter through a preallocated uniform matrix — cutting
+  ~O(n·H²) memory traffic per tile call on both backends (the Rust kernel takes a
+  read-only borrow; the numpy chain builds its own ridged copy). Outputs are
+  **byte-identical** (function-level identity checked on common and zero-masked stacks;
+  cond-10k end-to-end ATT bit-identical, fit time 20.09s → 19.99s median — a
+  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).
 
 ## [3.6.2] - 2026-07-03
 

TODO.md

@@ -45,7 +45,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m
 
 | Issue | Location | Origin | Effort | Priority |
 |-------|----------|--------|--------|----------|
-| `EfficientDiD` conditional path: post-Cholesky-dispatch (2026-07, 10k fit 22.7s) the ridge-solve stage is 3.7s, of which the Rust kernel is only ~1.6s — the remainder is shared Python prep inside `_ridge_solve_weights`: the `zero_mask` full-stack abs scan plus the `omega_stack[rest]` fancy-index copy (~28 GB of extra memory traffic per 10k fit). The copy is avoidable when no rows are zero-masked (`rest.size == n`, the common case) — value-identical on both backends since `om` is never mutated. After that, the largest O(n) stage is the sieve/nuisance construction outside the tiled pass (~9s at 10k). | `efficient_did_covariates.py::_ridge_solve_weights` | CS-scaling | Quick | Low |
+| `EfficientDiD` conditional path: the largest remaining O(n) stage is the sieve/nuisance construction outside the tiled pass (~9s at 10k). (The `_ridge_solve_weights` Python-prep shave landed 2026-07-07 — the `omega_stack[rest]` fancy-index copy and tail scatter are skipped when no row is zero-masked, byte-identical outputs; the `zero_mask` abs scan itself remains, needed for correctness.) | `efficient_did_covariates.py` | CS-scaling | Mid | Low |
 | `CallawaySantAnna` repeated-cross-section ipw/dr paths (`_ipw_estimation_rc` / `_doubly_robust_rc`) re-solve the propensity logit for EVERY (g,t) cell with no cache — the panel paths dedup via `pscore_cache` keyed `(g, base_period[, t])`. The Phase 3 IRLS fast path already speeds each RC solve; the missing (g,t)-dedup cache is the remaining lever. Unprofiled surface — measure an RC covariate scenario first. | `staggered.py::_ipw_estimation_rc`, `_doubly_robust_rc` | CS-scaling | Mid | Low |
 | Rust `solve_ols` always runs a full thin SVD (equilibrated, gelsd-parity); at 40 covariates it is the top per-cell solver item on a CS dr fit (1.04s of 3.1s at 2M rows, 95 calls). A Cholesky/QR fast path with SVD fallback — mirroring the Phase 3 Python-side pattern (certify well-conditioned, fall back verbatim) — is the natural lever, but `solve_ols` is the universal OLS entry point (every estimator), so the blast radius needs the full backend-parity treatment (`TestSolveOLSSkipRankCheckParity` posture: fitted-values parity, not beta). | `rust/src/linalg.rs::solve_ols` | CS-scaling | Heavy | Low |
 | `ImputationDiD` dense `(A0'A0).toarray()` scales `O((U+T+K)^2)` — OOM risk on large panels (only triggers when the sparse solver fails). Needs an alternative dense fallback or richer sparse strategy. | `imputation.py` | #141 | Heavy | Medium |

diff_diff/efficient_did_covariates.py

@@ -1620,15 +1620,19 @@ def _ridge_solve_weights(omega_stack: np.ndarray, omega_ridge: float) -> np.ndar
     """
     n, H, _ = omega_stack.shape
     ones = np.ones(H)
-    weights = np.full((n, H), 1.0 / H)
 
     # np.allclose(omega_i, 0.0) == elementwise |x| <= atol (1e-8), rtol inert
     zero_mask = np.all(np.abs(omega_stack) <= 1e-8, axis=(1, 2))
     rest = np.flatnonzero(~zero_mask)
     if rest.size == 0:
-        return weights
-
-    om = omega_stack[rest]
+        return np.full((n, H), 1.0 / H)
+
+    # Common case: no row is zero-masked — solve on the stack directly
+    # instead of paying the O(n*H^2) fancy-index copy (`om` is never
+    # mutated: the Rust kernel takes a read-only borrow and adds the ridge
+    # in-kernel; the numpy chain builds its own om_ridged). Value-identical
+    # on both backends.
+    om = omega_stack if rest.size == n else omega_stack[rest]
     trace = np.trace(om, axis1=1, axis2=2)
     ridge = omega_ridge * np.maximum(trace / H, 0.0)
     if (
@@ -1653,6 +1657,9 @@ def _ridge_solve_weights(omega_stack: np.ndarray, omega_ridge: float) -> np.ndar
     ok = np.abs(den) >= 1e-15
     solved = np.full((rest.size, H), 1.0 / H)
     solved[ok] = num[ok] / den[ok, None]
+    if rest.size == n:
+        return solved
+    weights = np.full((n, H), 1.0 / H)
     weights[rest] = solved
     return weights
 

tests/test_efficient_did.py

@@ -1061,6 +1061,30 @@ def test_f32_stack_declines_dispatch(self):
         expected = self._legacy_weights(om32, 1e-6)
         np.testing.assert_array_equal(got, expected)
 
+    def test_zero_mask_arms_all_zero_mixed_none(self):
+        """Zero-mask handling across all three arms after the no-copy fast
+        path (v3.7 prep shave): an all-zero stack returns uniform 1/H; a
+        mixed stack gives zero rows uniform 1/H while solved rows match a
+        standalone solve of the nonzero sub-stack byte-for-byte (each row's
+        solve is independent, so batch composition must not matter); a
+        no-zero stack takes the direct fast path, already locked
+        byte-identical to the legacy chain by
+        test_symbol_none_degrades_to_legacy_exactly."""
+        import diff_diff.efficient_did_covariates as cov_mod
+
+        H = 4
+        w0 = cov_mod._ridge_solve_weights(np.zeros((3, H, H)), 1e-6)
+        np.testing.assert_array_equal(w0, np.full((3, H), 1.0 / H))
+
+        mixed = self._spd_stack(m=5, H=H, seed=13)
+        mixed[[1, 3]] = 0.0
+        w_mixed = cov_mod._ridge_solve_weights(mixed, 1e-6)
+        np.testing.assert_array_equal(w_mixed[1], np.full(H, 1.0 / H))
+        np.testing.assert_array_equal(w_mixed[3], np.full(H, 1.0 / H))
+        keep = [0, 2, 4]
+        w_sub = cov_mod._ridge_solve_weights(mixed[keep], 1e-6)
+        np.testing.assert_array_equal(w_mixed[keep], w_sub)
+
 
 class TestValidTriples:
     """Test enumerate_valid_triples with hand-worked examples."""