Skip to content

⚡ Bolt: optimize covariance matrix powers and runica training loops#266

Open
suraj-ranganath wants to merge 1 commit into
developfrom
bolt-covariance-runica-opt-10895161003436251304
Open

⚡ Bolt: optimize covariance matrix powers and runica training loops#266
suraj-ranganath wants to merge 1 commit into
developfrom
bolt-covariance-runica-opt-10895161003436251304

Conversation

@suraj-ranganath

Copy link
Copy Markdown
Member

💡 What: Optimized core numerical operations in covariance.py and runica.py.
🎯 Why: Stacks of covariance matrices and ICA training loops are performance bottlenecks.
📊 Impact: diag_nd is ~10x faster; runica training is ~3% faster.
🔬 Measurement: Verified with tools/benchmark_covariance_optimized.py and tools/benchmark_runica.py.


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

- Vectorize `diag_nd` in `covariance.py` using advanced indexing (~10x speedup).
- Optimize `cov_logm`, `cov_expm`, etc. using broadcasting to avoid `diag_nd` calls and redundant matrix multiplications.
- Vectorize row mean subtraction in `runica.py`.
- Replace `_matmul(bias, onesrow)` with broadcasting in `runica.py` training loops, reducing memory allocation and compute overhead.

Co-authored-by: suraj-ranganath <14310165+suraj-ranganath@users.noreply.github.com>
@google-labs-jules

Copy link
Copy Markdown

👋 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 Bot commented Jun 30, 2026

Copy link
Copy Markdown

Claude finished @suraj-ranganath's task in 1m 48s —— View job


Code review

  • Read PR diff against origin/develop
  • Verify covariance.py broadcasting equivalence
  • Verify runica.py mean-removal and bias broadcasting
  • Check finite_matmul NaN-handling semantics
  • Assess test coverage

Overall assessment: Safe to merge. The optimizations are mathematically equivalent to the originals.
Highest-risk area: covariance.py eigendecomposition reconstruction (V @ diag(d) @ V.TV * d).
Merge recommendation: Safe to merge.

Blocking

None.

Important

None.

Nits

  • .jules/bolt.md (new file) — This is a Jules agent scratch/learning note being committed into the repo. AGENTS.md says to "Touch only files you must"; an agent log isn't part of the library. The header date (2025-05-15) is also stale/incorrect. Suggest removing it from the PR or adding .jules/ to .gitignore so these notes don't land in the tree. Severity: nit.

Test gaps

None required. These are behavior-preserving optimizations, so the existing covariance/runica tests cover the changed paths. I confirmed the rewrites produce identical results (see parity notes), so no new tests are warranted.

EEGLAB parity notes

Verified the rewrites are equivalent, not just faster:

  • covariance.pycov_logm/expm/powm/sqrtm/rsqrtm/sqrtm2: V @ diag(d) @ V.T has (V @ diag(d))[i,j] = V[i,j]·d[j], which is exactly V * d[..., newaxis, :]. The new code is identical bit-for-bit in the normal (finite) case. Importantly, I checked finite_matmul (src/eegprep/functions/miscfunc/misc.py:42) — it only suppresses spurious RuntimeWarnings and does no nan_to_num cleaning, so collapsing the inner finite_matmul into element-wise scaling introduces no difference in NaN/Inf propagation. ✅
  • covariance.pydiag_nd: vectorized res[..., i, i] = M matches the old np.diag-per-slice + reshape, and dtype=M.dtype preserves the original dtype. ✅
  • runica.py — mean removal: data -= rowmeans[:, np.newaxis] is the in-place equivalent of the per-row loop. ✅
  • runica.py — bias term: dropping onesrow and using + bias relies on (ncomps, 1) broadcasting across the block width, equal to bias @ onesrow. Blocks are uniformly sized (lastt = (datalength/block - 1)*block + 1, range(0, lastt, block)), and broadcasting is robust even for a short final block. The two remaining onesrow mentions are only in # MATLAB: comments, which is fine for parity documentation. ✅
    · bolt-covariance-runica-opt-10895161003436251304

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant