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

[igerber]

Summary

• _ridge_solve_weights (the batched Omega* ridge solve behind the conditional-path efficient weights): when no row of the stack is zero-masked (rest.size == n, the common case), consume omega_stack directly instead of paying the O(n·H²) omega_stack[rest] fancy-index copy, and return the solved weights directly instead of scattering through a preallocated uniform matrix. The zero_mask abs scan is retained (needed for correctness); mixed/all-zero stacks keep the prior arms verbatim.
• Safe by construction: om is never mutated on either backend — the Rust kernel takes a read-only borrow (PyReadonlyArray3) and adds the ridge in-kernel; the numpy chain builds its own om_ridged.
• Numbers (honest read): function-level 5.6ms → 4.5ms per call at (10k, H=24) with byte-identical outputs on common and zero-masked stacks; end-to-end cond-10k (20 periods, 5 cohorts, 5 covariates, 3-rep medians) 20.09s → 19.99s with bit-identical overall ATT — a memory-traffic cleanup (the TODO row's ~28 GB/10k-fit copy traffic), not a headline speedup.
• New TestRidgeSolveRustDispatch::test_zero_mask_arms_all_zero_mixed_none locks all three zero-mask arms (all-zero → uniform 1/H; mixed → zero rows uniform + solved rows byte-equal to a standalone sub-stack solve; no-zero fast path already locked byte-identical by test_symbol_none_degrades_to_legacy_exactly).
• TODO row narrowed to the remaining conditional-path lever (sieve/nuisance stage); CHANGELOG under Unreleased → Performance.

Methodology references (required if estimator / math changes)

• Method name(s): Roth & Sant'Anna (2023) efficient DiD — Omega* efficient-weight solve (implementation routing only; REGISTRY documents the ridge solve (Omega* + λ·max(trace/H,0)·I)x = 1 and backend dispatch as implementation choices)
• Paper / source link(s): REGISTRY.md § EfficientDiD (Omega* ridge Note)
• Any intentional deviations from the source (and why): None — byte-identical outputs on both backends.

Validation

• Tests added/updated: tests/test_efficient_did.py (new three-arm zero-mask test).
• Evidence: tests/test_efficient_did.py full file passes on the Rust backend (202) and the ridge/tile subset under DIFF_DIFF_BACKEND=python; function-level byte-identity BEFORE vs AFTER on saved outputs; cond-10k end-to-end ATT bit-identical.

Security / privacy

• Confirm no secrets/PII in this PR: Yes

---

Generated with Claude Code

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: 15981b1174016ef8c38b885dbbd8da7add406400

---

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings.

Executive Summary

• The affected method is EfficientDiD, specifically the conditional-path Omega* ridge weight solve.
• The PR preserves the documented ridge equation and only changes allocation/routing when rest.size == n.
• Registry notes already document the Omega* ridge and backend solve choices; this is an implementation optimization, not a methodology deviation.
• Zero-mask edge cases are handled for all-zero, mixed, and no-zero stacks.
• No inference, SE, default parameter, control-group, or public API behavior is changed.

Methodology

• Severity: None
  Impact: No undocumented mismatch found. _ridge_solve_weights still solves (Omega* + lambda * max(trace/H, 0) * I)x = 1 and normalizes by 1'x, matching the documented EfficientDiD Omega* ridge note in docs/methodology/REGISTRY.md:L1287.
  Concrete fix: None required.

• Severity: P3 informational
  Impact: The PR changes backend preparation only: no-zero stacks now pass omega_stack directly and return solved directly at diff_diff/efficient_did_covariates.py:L1630-L1661. This is covered by the registry’s documented “pure implementation change” posture for ridge backend dispatch.
  Concrete fix: None required.

Code Quality

• Severity: None
  Impact: The branch logic preserves the old all-zero, mixed-zero, and no-zero semantics. In the no-zero case, rest is the identity index, so returning solved is equivalent to scattering into a uniform matrix at diff_diff/efficient_did_covariates.py:L1658-L1663.
  Concrete fix: None required.

Performance

• Severity: None
  Impact: The change removes an avoidable omega_stack[rest] copy and final scatter in the common no-zero path, matching the stated performance goal at CHANGELOG.md:L415-L426.
  Concrete fix: None required.

Maintainability

• Severity: None
  Impact: The added comment explains the no-copy safety condition and is consistent with the Rust wrapper, which reads through ndarray views and copies each matrix into scratch before solving.
  Concrete fix: None required.

Tech Debt

• Severity: P3 informational
  Impact: The remaining conditional-path performance work is tracked in TODO.md:L48, so it is not a blocker under the deferred-work rules.
  Concrete fix: None required.

Security

• Severity: None
  Impact: No secrets, file/network access changes, or security-sensitive surfaces are introduced.
  Concrete fix: None required.

Documentation/Tests

• Severity: None
  Impact: The PR adds targeted coverage for all-zero and mixed zero-mask branches, and indirectly exercises the no-zero direct solve via the standalone sub-stack comparison at tests/test_efficient_did.py:L1064-L1086. Changelog/TODO updates are present.
  Concrete fix: None required.

[github-actions]

🔁 AI review rerun (requested by @igerber)

Head SHA: d2389b1ca372dddd2093ded7b8fe7c1c47d55a3c

---

Overall Assessment

✅ Looks good — no unmitigated P0/P1 findings.

Executive Summary

• Affected method: EfficientDiD conditional-path Omega* ridge weight solve.
• The PR changes allocation/routing only; the documented ridge equation and normalization are unchanged.
• All-zero, mixed zero-mask, and no-zero paths preserve the prior semantics.
• No inference, SE, public parameter, default behavior, or control-group logic changed.
• Remaining performance work is tracked in TODO.md.
• Targeted test execution was not possible locally because pytest is unavailable.

Methodology

• Severity: P3 informational
  Impact: _ridge_solve_weights still solves (Omega* + omega_ridge * max(trace/H, 0) * I)x = 1 and normalizes by 1'x, matching the registry’s EfficientDiD Omega* ridge note. The changed no-zero branch at diff_diff/efficient_did_covariates.py:L1630-L1664 is an implementation optimization, not a methodology deviation. Registry basis: docs/methodology/REGISTRY.md:L1215-L1228, docs/methodology/REGISTRY.md:L1280-L1289.
  Concrete fix: None required.

Code Quality

• Severity: None
  Impact: Branch semantics are preserved: all-zero stacks return uniform weights, mixed stacks scatter solved nonzero rows back into a uniform matrix, and no-zero stacks return the solved matrix directly. See diff_diff/efficient_did_covariates.py:L1624-L1664.
  Concrete fix: None required.

Performance

• Severity: None
  Impact: The common no-zero path avoids omega_stack[rest] and the final scatter, matching the documented performance goal in CHANGELOG.md:L435-L445.
  Concrete fix: None required.

Maintainability

• Severity: None
  Impact: The safety comment is consistent with the Rust kernel contract: the kernel takes PyReadonlyArray3 and copies each matrix into scratch before adding the ridge. Context: rust/src/batched_solve.rs:L43-L47, rust/src/batched_solve.rs:L103-L110.
  Concrete fix: None required.

Tech Debt

• Severity: P3 informational
  Impact: Remaining conditional-path performance work is tracked in TODO.md:L48, so it is not a blocker under the deferred-work policy.
  Concrete fix: None required.

Security

• Severity: None
  Impact: No secrets, file/network access, or security-sensitive behavior introduced.
  Concrete fix: None required.

Documentation/Tests

• Severity: None
  Impact: The added test covers all-zero and mixed zero-mask behavior and compares mixed solved rows against a standalone nonzero sub-stack solve at tests/test_efficient_did.py:L1064-L1086. Changelog and TODO updates are present. I attempted the targeted test, but pytest was not installed in the runner.
  Concrete fix: None required.