⚡ Bolt: vectorized rmbase

[suraj-ranganath]

⚡ Bolt: vectorized rmbase

💡 What

Vectorized the baseline removal logic in src/eegprep/functions/sigprocfunc/rmbase.py by replacing an explicit Python loop over epochs with a 3D NumPy reshape and broadcasting.

🎯 Why

The original implementation iterated through every epoch, calling np.nanmean and performing subtraction individually. For datasets with thousands of epochs (common in long recordings or high-density EEG), this Python-level overhead became a measurable bottleneck.

📊 Impact

• ~5-6% speedup for typical datasets (e.g., 64 channels, 500 frames, 200 epochs).
• ~45% speedup for datasets with many small epochs (e.g., 2000 epochs, 50 frames).
• Maintains bit-perfect parity with the original implementation's numerical results and dtype handling.

🔬 Measurement

Verified using a temporary benchmark script (tools/benchmark_rmbase.py) comparing the original loop-based logic against the vectorized implementation across multiple scenarios. Correctness confirmed via tests/test_pop_rmbase.py.

---

PR created automatically by Jules for task 3581584258411826506 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 57s —— View job

---

Code review — PR #268: vectorized rmbase

• Gather context (diff, files)
• Analyze correctness & EEGLAB parity
• Check test coverage
• Post review

---

Code review

• Overall assessment: Safe to merge. The vectorization is correct and preserves bit-for-bit parity with the original loop.
• Highest-risk area: dtype handling and peak memory on large datasets.
• Merge recommendation: Safe to merge (one optional memory improvement below).

What I checked

• Epoch layout parity — matrix.reshape(chans, epochs, frames) is C-order, so reshaped[c, e, f] == matrix[c, e*frames + f], which matches the original loop's output[:, e*frames : e*frames+frames] exactly. ✅
• Baseline selection — reshaped[:, :, baseline] with 0-based baseline matches matrix[:, start + baseline]. ✅
• NaN handling — still np.nanmean(..., dtype=np.float64) over the same samples. ✅
• dtype behavior — int → float64, floating → preserved. The end-of-function cast is numerically identical to the original per-epoch cast, because each element is rounded to the target dtype the same way. ✅
• No input mutation — output is a fresh array from the subtraction, so the input is never modified (the original relied on .copy(); the new code is safe by construction). ✅
• return_mean — means has the same (chans, epochs) shape and float64 dtype. ✅

Blocking

None.

Important

Peak memory roughly doubles for float32 data — the exact "high-density / long recording" case the PR targets. src/eegprep/functions/sigprocfunc/rmbase.py:59

(reshaped - means[:, :, np.newaxis]) allocates a full-dataset float64 temporary before casting back down to the input dtype at line 65. For a float32 input (e.g. 256 ch × 1e6 samples ≈ 1 GB), that temporary is ~2 GB, whereas the original subtracted into a single float32 buffer per epoch. Correctness is unaffected, but it partly works against the stated motivation.

A single-allocation, still bit-parity-preserving rewrite (in-place float32 -= float64 keeps the float32 dtype and rounds per element exactly like the original):

out_dtype = matrix.dtype if np.issubdtype(matrix.dtype, np.floating) else np.float64
output = reshaped.astype(out_dtype, copy=True)
output -= means[:, :, np.newaxis]
output = output.reshape(chans, total_frames)

This drops the whole-array float64 temporary while keeping the same numerical result. Fix this →

Nits

• .jules/bolt.md is a bot learnings artifact rather than product code; it's harmless but note it's now committed to the repo.

Test gaps

None strictly required — tests/test_pop_rmbase.py::test_rmbase_removes_epoch_baseline_and_returns_means exercises the epoched + baseline + return_mean path this change touches. If you want to lock in the dtype-preservation contract that the vectorized cast now depends on, a small parametrized case over float32 / int16 inputs asserting out.dtype would be a cheap regression guard.

EEGLAB parity notes

None. Behavior, dtype semantics, and (chans, epochs) mean output are unchanged from the pre-existing EEGLAB-style implementation.

---

Note: I was unable to run uv run pytest in this environment (sandbox denied the command), so parity was verified by static analysis rather than execution.