diff --git a/CHANGELOG.md b/CHANGELOG.md index 956cda27..5ebc400a 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/TODO.md b/TODO.md index 4f1ad39d..88a26e73 100644 --- a/TODO.md +++ b/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 --git a/diff_diff/efficient_did_covariates.py b/diff_diff/efficient_did_covariates.py index 7623f479..2611c0ac 100644 --- a/diff_diff/efficient_did_covariates.py +++ b/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 diff --git a/tests/test_efficient_did.py b/tests/test_efficient_did.py index d685f2a4..7c8612c7 100644 --- a/tests/test_efficient_did.py +++ b/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."""