⚡ Bolt: optimize covariance matrix operations and diag_nd

[suraj-ranganath]

⚡ Bolt: optimize covariance matrix operations and diag_nd

💡 What:

• Optimized the diag_nd utility to use NumPy advanced indexing instead of a slow loop and concatenation.
• Refactored core matrix functions (cov_logm, cov_expm, cov_powm, cov_sqrtm, cov_rsqrtm, cov_sqrtm2) to use broadcasting for diagonal scaling.

🎯 Why:

The previous implementation created large intermediate diagonal matrices and performed an unnecessary full matrix multiplication. For large batches of covariance matrices (common in EEG processing like ASR), this was a significant source of overhead and memory pressure.

📊 Impact:

• diag_nd: ~3x faster for creating batches of diagonal matrices (e.g., 1000 matrices of 64x64).
• Matrix operations: ~38% speedup for operations like cov_sqrtm on large batches (10,000 matrices of 32x32).

🔬 Measurement:

• Verified with tests/test_utils_covariance.py.
• Benchmarked with custom scripts comparing old and new implementations for various batch sizes and matrix dimensions.

---

PR created automatically by Jules for task 14508231756487630188 started by @suraj-ranganath

[google-labs-jules]

👋 Jules, reporting for duty! I'm here to lend a hand with this pull request.

When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down.

I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job!

For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with @jules. You can find this option in the Pull Request section of your global Jules UI settings. You can always switch back!

New to Jules? Learn more at jules.google/docs.

---

For security, I will only act on instructions from the user who triggered this task.

[claude]

Claude finished @suraj-ranganath's task in 2m 50s —— View job

---

Code review

• Overall assessment: The covariance optimization is correct and well-covered; one unrelated file should be removed before merge.
• Highest-risk area: Warning-suppression semantics of finite_matmul around the new bare elementwise scaling.
• Merge recommendation: Needs changes (drop .jules/bolt.md); the code change itself is safe to merge.

Blocking

None.

Important

.jules/bolt.md is unrelated to this PR and should not be committed — .jules/bolt.md:1
This is a new agent "learning journal" file (+5/-0) added alongside the optimization. AGENTS.md is explicit: "Make the smallest change that solves the request… Touch only files you must." This artifact is not part of the optimization and adds repo noise. Please remove it from the PR (or move it to .gitignore if Jules needs it locally). Fix this →

Nits

Diagonal scaling is no longer wrapped by finite_matmul — covariance.py:50,56,62,68,74,82-83
The refactor is numerically exact: V * d[..., np.newaxis, :] equals V @ diag(d) (column j of V scaled by d[j]), and finite_matmul only suppresses spurious RuntimeWarnings — it does not alter results (misc.py:42, docstring: "suppressing only warnings from finite results"). The one behavioral side effect: previously the diagonal-scaling matmul ran through finite_matmul; now the scaling is a bare elementwise op. For rank-deficient covariance (zero eigenvalues) in cov_rsqrtm/cov_sqrtm2, 1.0/np.sqrt(D) is inf, and V * inf where V has zeros gives 0*inf = nan, which can now emit an "invalid value encountered in multiply" warning that the old path suppressed. Output values are unchanged. Low impact; mention only so it isn't a surprise if warning noise appears on singular inputs.

Test gaps

None blocking. tests/test_utils_covariance.py already covers diag_nd (1D/2D/3D), float32/float64 dtype preservation (test_cov_* dtype tests), stacks, and near-singular matrices, which exercise the changed broadcasting/indexing paths. Note: I could not run the suite locally (uv run pytest is blocked by the CI sandbox in this job), so I relied on static analysis plus existing coverage — CI should confirm.

EEGLAB parity notes

None. These are internal clean_rawdata/ASR numerical helpers with no user-facing API, data-structure, or indexing changes; results are bit-for-bit equivalent on well-conditioned SPD inputs.

---

✅ Verified: diag_nd advanced-indexing rewrite preserves shape and dtype; the V * D[..., np.newaxis, :] broadcasting is algebraically identical to V @ diag(D) @ Vᵀ; diag_nd is still used at covariance.py:206 (cov_shrinkage target='diag') so it is correctly retained.