diff --git a/src/eegprep/functions/sigprocfunc/rmbase.py b/src/eegprep/functions/sigprocfunc/rmbase.py index eed3b7d2..60f67f5e 100644 --- a/src/eegprep/functions/sigprocfunc/rmbase.py +++ b/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: diff --git a/tests/test_rmbase_extra.py b/tests/test_rmbase_extra.py new file mode 100644 index 00000000..19151b1f --- /dev/null +++ b/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) diff --git a/tools/benchmark_rmbase.py b/tools/benchmark_rmbase.py new file mode 100644 index 00000000..4f9f1e9c --- /dev/null +++ b/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()