Vectorize rmbase with bounded memory

src/eegprep/functions/sigprocfunc/rmbase.py

@@ -6,6 +6,8 @@
 
 import numpy as np
 
+_NANMEAN_CHUNK_BYTES = 4 * 1024 * 1024
+
 
 def rmbase(data: Any, frames: int | None = 0, basevector: Any = 0, *, return_mean: bool = False):
     """Subtract per-channel baseline means from continuous or epoched data.
@@ -29,8 +31,8 @@ def rmbase(data: Any, frames: int | None = 0, basevector: Any = 0, *, return_mea
         raise ValueError("rmbase(): data must be 2D or 3D")
 
     original_shape = array.shape
-    matrix = array.transpose(0, 2, 1).reshape(array.shape[0], -1) if array.ndim == 3 else array
-    chans, total_frames = matrix.shape
+    channels = array.shape[0]
+    total_frames = array.shape[1] * array.shape[2] if array.ndim == 3 else array.shape[1]
     frames = int(frames or 0)
     if frames == 0:
         frames = total_frames
@@ -44,19 +46,34 @@ def rmbase(data: Any, frames: int | None = 0, basevector: Any = 0, *, return_mea
         raise ValueError("rmbase(): total sample count must be an integer multiple of frames")
 
     baseline = _baseline_indices(basevector, frames)
-    output = matrix.astype(np.float64, copy=True) if not np.issubdtype(matrix.dtype, np.floating) else matrix.copy()
-    means = np.zeros((chans, epochs), dtype=np.result_type(matrix.dtype, np.float64))
-    for epoch in range(epochs):
-        start = epoch * frames
-        stop = start + frames
-        epoch_slice = output[:, start:stop]
-        if baseline is None:
-            mean = np.nanmean(matrix[:, start:stop], axis=1, keepdims=True, dtype=np.float64)
-        else:
-            mean = np.nanmean(matrix[:, start + baseline], axis=1, keepdims=True, dtype=np.float64)
-        means[:, epoch : epoch + 1] = mean
-        output[:, start:stop] = epoch_slice - mean
 
+    # Keep frames contiguous, as in the original epoch loop, while using this
+    # copy as the output buffer. Integer input retains the legacy float64 output.
+    output_dtype = np.float64 if not np.issubdtype(array.dtype, np.floating) else array.dtype
+    epoch_order = array.transpose(0, 2, 1) if array.ndim == 3 else array
+    output = np.array(epoch_order, dtype=output_dtype, order="C", copy=True)
+    output_reshaped = output.reshape(channels, epochs, frames)
+    means = np.empty((channels, epochs), dtype=np.result_type(array.dtype, np.float64))
+
+    # np.nanmean makes a data copy and a validity mask. Process several epochs
+    # at a time so those temporaries stay bounded without returning to a Python
+    # loop per epoch.
+    baseline_frames = frames if baseline is None else baseline.size
+    mean_bytes_per_epoch = channels * baseline_frames * output_reshaped.dtype.itemsize
+    chunk_epochs = max(1, min(epochs, _NANMEAN_CHUNK_BYTES // max(1, mean_bytes_per_epoch)))
+    for start in range(0, epochs, chunk_epochs):
+        stop = min(epochs, start + chunk_epochs)
+        output_chunk = output_reshaped[:, start:stop, :]
+        baseline_chunk = output_chunk if baseline is None else output_chunk[:, :, baseline]
+        chunk_means = np.nanmean(baseline_chunk, axis=2, dtype=np.float64)
+        means[:, start:stop] = chunk_means
+
+        # Compute with the float64 means but write directly into the intended
+        # output dtype. This preserves legacy float32 rounding without a
+        # recording-sized float64 subtraction result.
+        np.subtract(output_chunk, chunk_means[:, :, np.newaxis], out=output_chunk, casting="unsafe")
+
+    output = output_reshaped.reshape(channels, total_frames)
     if array.ndim == 3:
         output = output.reshape(original_shape[0], original_shape[2], original_shape[1]).transpose(0, 2, 1)
     else:

tests/test_rmbase_extra.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from eegprep.functions.sigprocfunc.rmbase import rmbase
+
+
+def _legacy_rmbase(
+    data: np.ndarray,
+    frames: int,
+    basevector: list[int] | int = 0,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Run the pre-vectorization loop as a numerical reference."""
+    array = np.asarray(data)
+    original_shape = array.shape
+    matrix = array.transpose(0, 2, 1).reshape(array.shape[0], -1) if array.ndim == 3 else array
+    channels, total_frames = matrix.shape
+    epochs = total_frames // frames
+    baseline = None if basevector == 0 else np.asarray(basevector, dtype=int) - 1
+
+    output = matrix.astype(np.float64, copy=True) if not np.issubdtype(matrix.dtype, np.floating) else matrix.copy()
+    means = np.zeros((channels, epochs), dtype=np.result_type(matrix.dtype, np.float64))
+    for epoch in range(epochs):
+        start = epoch * frames
+        stop = start + frames
+        if baseline is None:
+            mean = np.nanmean(matrix[:, start:stop], axis=1, keepdims=True, dtype=np.float64)
+        else:
+            mean = np.nanmean(matrix[:, start + baseline], axis=1, keepdims=True, dtype=np.float64)
+        means[:, epoch : epoch + 1] = mean
+        output[:, start:stop] = output[:, start:stop] - mean
+
+    if array.ndim == 3:
+        output = output.reshape(original_shape[0], original_shape[2], original_shape[1]).transpose(0, 2, 1)
+    return output.reshape(original_shape), means
+
+
+def test_rmbase_3d_default_frames_matches_legacy_grand_mean():
+    data = np.arange(24, dtype=np.float32).reshape(2, 4, 3)
+    total_frames = data.shape[1] * data.shape[2]
+
+    expected, expected_means = _legacy_rmbase(data, frames=total_frames)
+    actual, actual_means = rmbase(data, return_mean=True)
+
+    np.testing.assert_array_equal(actual, expected)
+    np.testing.assert_array_equal(actual_means, expected_means)
+    assert actual_means.shape == (data.shape[0], 1)
+
+
+@pytest.mark.parametrize("shape", [(3, 85), (3, 17, 5)])
+@pytest.mark.parametrize("basevector", [0, [1, 4, 7, 11]])
+def test_rmbase_float32_matches_legacy_rounding(shape: tuple[int, ...], basevector: list[int] | int):
+    rng = np.random.default_rng(268)
+    data = (rng.standard_normal(shape) * 1_000).astype(np.float32)
+    original = data.copy()
+
+    expected, expected_means = _legacy_rmbase(data, frames=17, basevector=basevector)
+    actual, actual_means = rmbase(data, frames=17, basevector=basevector, return_mean=True)
+
+    np.testing.assert_array_equal(actual, expected)
+    np.testing.assert_array_equal(actual_means, expected_means)
+    np.testing.assert_array_equal(data, original)
+    assert actual.dtype == np.float32
+    assert actual_means.dtype == np.float64
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.int16])
+def test_rmbase_matches_legacy_for_other_dtypes(dtype: type[np.generic]):
+    rng = np.random.default_rng(269)
+    if np.issubdtype(dtype, np.integer):
+        data = rng.integers(-2_000, 2_000, size=(4, 13, 3), dtype=dtype)
+    else:
+        data = (rng.standard_normal((4, 13, 3)) * 1_000).astype(dtype)
+
+    expected, expected_means = _legacy_rmbase(data, frames=13, basevector=[2, 5, 9])
+    actual, actual_means = rmbase(data, frames=13, basevector=[2, 5, 9], return_mean=True)
+
+    np.testing.assert_array_equal(actual, expected)
+    np.testing.assert_array_equal(actual_means, expected_means)
+    assert actual.dtype == (np.dtype(np.float64) if np.issubdtype(dtype, np.integer) else dtype)
+
+
+def test_rmbase_preserves_nan_results_and_warning_behavior():
+    data = np.arange(16, dtype=np.float32).reshape(2, 8)
+    data[0, :2] = np.nan
+    data[1, 4:6] = np.nan
+
+    with pytest.warns(RuntimeWarning, match="Mean of empty slice"):
+        actual, actual_means = rmbase(data, frames=4, basevector=[1, 2], return_mean=True)
+    with pytest.warns(RuntimeWarning, match="Mean of empty slice"):
+        expected, expected_means = _legacy_rmbase(data, frames=4, basevector=[1, 2])
+
+    np.testing.assert_array_equal(actual, expected)
+    np.testing.assert_array_equal(actual_means, expected_means)

tools/benchmark_rmbase.py

@@ -0,0 +1,94 @@
+"""Compare vectorized ``rmbase`` with its pre-vectorization loop."""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import statistics
+import time
+import tracemalloc
+from collections.abc import Callable
+
+import numpy as np
+
+from eegprep.functions.sigprocfunc.rmbase import rmbase
+
+
+def _legacy_rmbase(data: np.ndarray, frames: int) -> np.ndarray:
+    """Reproduce the implementation replaced by the optimization."""
+    original_shape = data.shape
+    matrix = data.transpose(0, 2, 1).reshape(data.shape[0], -1) if data.ndim == 3 else data
+    channels, total_frames = matrix.shape
+    epochs = total_frames // frames
+    output = matrix.copy()
+
+    for epoch in range(epochs):
+        start = epoch * frames
+        stop = start + frames
+        mean = np.nanmean(matrix[:, start:stop], axis=1, keepdims=True, dtype=np.float64)
+        output[:, start:stop] = output[:, start:stop] - mean
+
+    if data.ndim == 3:
+        return output.reshape(channels, original_shape[2], original_shape[1]).transpose(0, 2, 1)
+    return output
+
+
+def _median_seconds(operation: Callable[[], np.ndarray], repeats: int) -> float:
+    timings: list[float] = []
+    for _ in range(repeats):
+        gc.collect()
+        start = time.perf_counter()
+        operation()
+        timings.append(time.perf_counter() - start)
+    return statistics.median(timings)
+
+
+def _peak_bytes(operation: Callable[[], np.ndarray]) -> int:
+    gc.collect()
+    tracemalloc.start()
+    operation()
+    _, peak = tracemalloc.get_traced_memory()
+    tracemalloc.stop()
+    return peak
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--channels", type=int, default=64)
+    parser.add_argument("--frames", type=int, default=500)
+    parser.add_argument("--epochs", type=int, default=200)
+    parser.add_argument("--repeats", type=int, default=5)
+    parser.add_argument("--seed", type=int, default=268)
+    args = parser.parse_args()
+    if min(args.channels, args.frames, args.epochs, args.repeats) < 1:
+        parser.error("channels, frames, epochs, and repeats must be positive")
+
+    rng = np.random.default_rng(args.seed)
+    data = rng.standard_normal((args.channels, args.frames, args.epochs), dtype=np.float32)
+
+    def legacy() -> np.ndarray:
+        return _legacy_rmbase(data, args.frames)
+
+    def vectorized() -> np.ndarray:
+        return rmbase(data, frames=args.frames)
+
+    expected = legacy()
+    actual = vectorized()
+    np.testing.assert_array_equal(actual, expected)
+
+    legacy_seconds = _median_seconds(legacy, args.repeats)
+    vectorized_seconds = _median_seconds(vectorized, args.repeats)
+    legacy_peak = _peak_bytes(legacy)
+    vectorized_peak = _peak_bytes(vectorized)
+
+    print(
+        f"Input: {args.channels} channels x {args.frames} frames x {args.epochs} epochs "
+        f"({data.nbytes / 2**20:.1f} MiB, {data.dtype})"
+    )
+    print(f"Median of {args.repeats} runs: legacy={legacy_seconds:.4f}s, vectorized={vectorized_seconds:.4f}s")
+    print(f"Observed speed ratio: {legacy_seconds / vectorized_seconds:.2f}x")
+    print(f"Tracemalloc peak: legacy={legacy_peak / 2**20:.1f} MiB, vectorized={vectorized_peak / 2**20:.1f} MiB")
+
+
+if __name__ == "__main__":
+    main()