⚡ Bolt: optimize covariance matrix operations

[suraj-ranganath]

💡 What: Optimized covariance matrix power functions and the diag_nd utility function in src/eegprep/plugins/clean_rawdata/private/covariance.py.

🎯 Why:

• diag_nd was using a loop-based np.concatenate which is slow for large batches.
• Covariance operations were creating intermediate diagonal matrices and performing full matrix multiplications ($O(N^3)$) when they could use NumPy broadcasting to scale eigenvectors directly ($O(N^2)$).

📊 Impact:

• diag_nd achieved a ~6x speedup in benchmarks.
• Covariance functions like cov_logm and cov_sqrtm achieved a ~2x speedup.
• Reduced memory overhead by avoiding intermediate diagonal matrices.

🔬 Measurement:

• Verified performance with a custom benchmark script (benchmarked 100 covariance matrices of size 32x32).
• Ensured numerical parity and no regressions by running tests/test_utils_covariance.py.

---

PR created automatically by Jules for task 14199838786116228957 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 1m 54s —— View job

---

Code review

• Overall assessment: Looks good — safe to merge.
• Highest-risk area: Numerical parity of the broadcast rewrite in cov_* functions.
• Merge recommendation: Safe to merge.

I verified the math, the diag_nd rewrite, and the finite_matmul semantics.

Blocking

None.

Important

None.

Nits

• .jules/bolt.md header date is 2025-05-14, but the PR/task is dated 2026-06-30. Cosmetic only.

Test gaps

None blocking. The existing tests/test_utils_covariance.py already covers diag_nd (1D/2D/3D) and the cov_* functions, which is exactly the changed surface. One optional addition worth considering: a test exercising the batched ...,N,N path through cov_logm/cov_sqrtm (the old code built diagonals via np.concatenate over the flattened batch; the new broadcast path is the one most prone to a silent axis mistake).

EEGLAB parity notes

None. These are internal clean_rawdata/private numerical helpers with no EEGLAB-facing API or data-structure change.

---

What I checked

1. cov_* broadcast rewrite is mathematically equivalent.
V * f(D)[..., np.newaxis, :] has shape (..., 1, N) and broadcasts across rows, scaling column j of V by f(D)[j]. That is exactly V @ diag(f(D)), so (V * f(D)[..., None, :]) @ V.T == V @ diag(f(D)) @ V.T. Correct for logm, expm, powm, sqrtm, rsqrtm, and cov_sqrtm2. The reused VT = V.swapaxes(-2, -1) in cov_sqrtm2 is a nice touch.

2. Removing the inner finite_matmul is safe.
finite_matmul (src/eegprep/functions/miscfunc/misc.py:42) only suppresses RuntimeWarnings when operands/result are finite — it returns a plain left @ right and never alters values. The inner call wrapped a true matmul (V @ diag); the replacement is an elementwise multiply that produces no matmul warning, so dropping the wrapper changes nothing numerically. (Minor bonus: for degenerate inputs where f(D) contains inf — e.g. log(0) or 1/sqrt(0) — the broadcast form avoids the 0*inf → nan off-diagonal contamination that the old V @ diag(inf) matmul could introduce. Strictly an improvement, not a regression.)

3. diag_nd advanced-indexing rewrite is equivalent.
res[..., idx, idx] = M with idx = np.arange(N) writes M onto the diagonal of a zero array, matching the old concatenate-of-np.diag output for the 1D, 2D, and 3D cases. dtype=M.dtype preserves the previous np.diag dtype behavior. diag_nd remains used at covariance.py:209 and is still exported/tested.

Clean, well-scoped optimization. 👍