diff --git a/src/eegprep/functions/sigprocfunc/runica.py b/src/eegprep/functions/sigprocfunc/runica.py index 5c23534b..f1f6c23b 100644 --- a/src/eegprep/functions/sigprocfunc/runica.py +++ b/src/eegprep/functions/sigprocfunc/runica.py @@ -488,9 +488,8 @@ def runica(data, **kwargs): if verbose: logger.info('Removing mean of each channel ...') - rowmeans = np.mean(data, axis=1) # shape: (chans,) - for i in range(data.shape[0]): - data[i, :] = data[i, :] - rowmeans[i] + rowmeans = np.mean(data, axis=1, keepdims=True) + data -= rowmeans if verbose: logger.info(f'Final training data range: {np.min(data):g} to {np.max(data):g}') @@ -593,7 +592,6 @@ def runica(data, **kwargs): prevwtchange = np.zeros((chans, ncomps)) oldwtchange = np.zeros((chans, ncomps)) lrates = np.zeros(maxsteps) - onesrow = np.ones((1, block)) bias = np.zeros((ncomps, 1)) # Initialize signs for extended-ICA @@ -682,17 +680,15 @@ def runica(data, **kwargs): for t in range(0, lastt, block): # Extract and process block (MATLAB line 846) # MATLAB: u = weights*double(data(:,timeperm(t:t+block-1))) + bias*onesrow - u = _matmul(weights, data[:, timeperm[t : t + block]]) + _matmul(bias, onesrow) + u = _matmul(weights, data[:, timeperm[t : t + block]]) + bias # Apply tanh nonlinearity (MATLAB line 848) y = np.tanh(u) # Extended-ICA natural gradient weight update (MATLAB line 849) # weights = weights + lrate*(BI-signs*y*u'-u*u')*weights - weights = weights + lrate * _matmul( - BI - _matmul(_matmul(signs, y), u.T) - _matmul(u, u.T), - weights, - ) + signs_y = np.diag(signs)[:, np.newaxis] * y + weights = weights + lrate * _matmul(BI - _matmul(signs_y + u, u.T), weights) # Bias update for tanh (MATLAB line 850) # bias = bias + lrate*sum((-2*y)')'; @@ -869,8 +865,7 @@ def runica(data, **kwargs): for t in range(0, lastt, block): # Extract and process block (MATLAB line 1021) # MATLAB: u = weights*double(data(:,timeperm(t:t+block-1))) + bias*onesrow - # Note: MATLAB uses 1-based indexing, so t:t+block-1 means t to t+block - u = _matmul(weights, data[:, timeperm[t : t + block]]) + _matmul(bias, onesrow) + u = _matmul(weights, data[:, timeperm[t : t + block]]) + bias # Apply logistic nonlinearity (MATLAB line 1022) # Clip u to prevent overflow in exp @@ -880,11 +875,12 @@ def runica(data, **kwargs): # Natural gradient weight update (MATLAB line 1023) # weights = weights + lrate*(BI+(1-2*y)*u')*weights - weights = weights + lrate * _matmul(BI + _matmul(1 - 2 * y, u.T), weights) + y_update = 1.0 - 2.0 * y + weights = weights + lrate * _matmul(BI + _matmul(y_update, u.T), weights) # Bias update (MATLAB line 1024) # bias = bias + lrate*sum((1-2*y)')'; - bias = bias + lrate * np.sum(1 - 2 * y, axis=1, keepdims=True) + bias = bias + lrate * np.sum(y_update, axis=1, keepdims=True) # Add momentum if enabled (MATLAB lines 1026-1030) if momentum > 0: @@ -1017,10 +1013,8 @@ def runica(data, **kwargs): y = np.tanh(u) # Extended-ICA natural gradient weight update (MATLAB line 1147) - weights = weights + lrate * _matmul( - BI - _matmul(_matmul(signs, y), u.T) - _matmul(u, u.T), - weights, - ) + signs_y = np.diag(signs)[:, np.newaxis] * y + weights = weights + lrate * _matmul(BI - _matmul(signs_y + u, u.T), weights) # NO BIAS UPDATE for no-bias variant @@ -1181,7 +1175,8 @@ def runica(data, **kwargs): y = 1.0 / (1.0 + np.exp(-u)) # Natural gradient weight update (MATLAB line 1317) - weights = weights + lrate * _matmul(BI + _matmul(1 - 2 * y, u.T), weights) + y_update = 1.0 - 2.0 * y + weights = weights + lrate * _matmul(BI + _matmul(y_update, u.T), weights) # NO BIAS UPDATE for no-bias variant @@ -1304,14 +1299,12 @@ def runica(data, **kwargs): # Add back the row means removed from data before sphering (MATLAB lines 1442-1447) if pcaflag == 'off': sr = _matmul(sphere, rowmeans) - for r in range(ncomps): - data[r, :] = data[r, :] + sr[r] + data += sr activations_unsorted = _matmul(weights, data) # MATLAB line 1447 else: # For PCA case (MATLAB lines 1449-1453) ser = _matmul(_matmul(sphere, eigenvectors[:, :ncomps].T), rowmeans) - for r in range(ncomps): - data[r, :] = data[r, :] + ser[r] + data += ser activations_unsorted = _matmul(weights, data) # Now 'activations_unsorted' are the component activations = weights*sphere*raw_data diff --git a/tests/test_runica_optimization.py b/tests/test_runica_optimization.py new file mode 100644 index 00000000..5c7a4677 --- /dev/null +++ b/tests/test_runica_optimization.py @@ -0,0 +1,75 @@ +import numpy as np +import pytest + +from eegprep.functions.sigprocfunc.runica import runica + + +_REFERENCE_CASES = [ + pytest.param( + 0, + "off", + [ + [0.006138258426678086, 0.920038988944051, -0.007614426473801105], + [1.245639957427586, 0.01054020803403042, -0.0025770443524238138], + [-0.00156237277009202, -0.009895611292063627, 1.424320270769803], + ], + [[0.0], [0.0], [0.0]], + id="logistic-no-bias", + ), + pytest.param( + 0, + "on", + [ + [0.006133360652087168, 0.9200026096663453, -0.0076195912896603], + [1.2455962713117488, 0.01054249400139514, -0.002576216189956501], + [-0.0015617542160055618, -0.009908514977658586, 1.4242983431233893], + ], + [[0.0027986461069580363], [-0.0013498284793111981], [0.0006208222287476582]], + id="logistic-bias", + ), + pytest.param( + -1, + "off", + [ + [0.016542806300490903, 0.4837462583100016, -0.009894727973969961], + [1.2466400017060382, 0.009666327764259823, 0.002293221790050116], + [-0.004086194330539649, -0.011541930062800668, 1.2120278671034428], + ], + [[0.0], [0.0], [0.0]], + id="extended-no-bias", + ), + pytest.param( + -1, + "on", + [ + [0.01657556669489547, 0.48362814000133625, -0.009935969624553129], + [1.24643159652647, 0.009643847125636948, 0.002320749638652659], + [-0.0040687687430069844, -0.011580978864895756, 1.211819845044342], + ], + [[0.005632991536067746], [-0.008324510388924372], [0.0033542275952815784]], + id="extended-bias", + ), +] + + +@pytest.mark.parametrize(("extended", "bias", "expected_weights", "expected_bias"), _REFERENCE_CASES) +def test_optimized_training_modes_match_preoptimization_reference(extended, bias, expected_weights, expected_bias): + """Lock down all four optimized training branches against the old implementation.""" + data = np.random.RandomState(20260716).standard_normal((3, 240)) * np.array([[1.0], [2.0], [0.5]]) + + weights, sphere, _compvars, actual_bias, signs, lrates = runica( + data, + extended=extended, + bias=bias, + sphering="none", + block=24, + maxsteps=3, + verbose=False, + rndreset="off", + ) + + np.testing.assert_allclose(weights, expected_weights, rtol=5e-13, atol=5e-13) + np.testing.assert_allclose(actual_bias, expected_bias, rtol=5e-13, atol=5e-13) + np.testing.assert_array_equal(sphere, np.eye(3)) + np.testing.assert_array_equal(signs, [1.0, -1.0, 1.0]) + np.testing.assert_allclose(lrates, np.full(3, 0.0005916554973074442), rtol=0.0, atol=1e-18) diff --git a/tools/benchmark_runica.py b/tools/benchmark_runica.py new file mode 100644 index 00000000..48a66d8e --- /dev/null +++ b/tools/benchmark_runica.py @@ -0,0 +1,127 @@ +"""Benchmark the allocation and matrix kernels optimized in ``runica``. + +This is a repeatable microbenchmark, not a CI performance assertion. Run it +from the repository root with ``python tools/benchmark_runica.py``. +""" + +from __future__ import annotations + +import argparse +from collections.abc import Callable +from time import perf_counter + +import numpy as np + + +ArrayFactory = Callable[[], np.ndarray] + + +def _elapsed(factory: ArrayFactory, iterations: int) -> float: + start = perf_counter() + for _ in range(iterations): + factory() + return (perf_counter() - start) / iterations + + +def _legacy_center(data: np.ndarray) -> np.ndarray: + output = data.copy() + rowmeans = np.mean(output, axis=1) + for row in range(output.shape[0]): + output[row, :] -= rowmeans[row] + return output + + +def _optimized_center(data: np.ndarray) -> np.ndarray: + output = data.copy() + output -= np.mean(output, axis=1, keepdims=True) + return output + + +def _legacy_extended_update( + weights: np.ndarray, + identity: np.ndarray, + signs: np.ndarray, + activations: np.ndarray, + projected: np.ndarray, + learning_rate: float, +) -> np.ndarray: + gradient = identity - signs @ activations @ projected.T - projected @ projected.T + return weights + learning_rate * gradient @ weights + + +def _optimized_extended_update( + weights: np.ndarray, + identity: np.ndarray, + signs: np.ndarray, + activations: np.ndarray, + projected: np.ndarray, + learning_rate: float, +) -> np.ndarray: + signed_activations = np.diag(signs)[:, np.newaxis] * activations + gradient = identity - (signed_activations + projected) @ projected.T + return weights + learning_rate * gradient @ weights + + +def _legacy_bias_projection( + weights: np.ndarray, block_data: np.ndarray, bias: np.ndarray, ones: np.ndarray +) -> np.ndarray: + return weights @ block_data + bias @ ones + + +def _optimized_bias_projection(weights: np.ndarray, block_data: np.ndarray, bias: np.ndarray) -> np.ndarray: + return weights @ block_data + bias + + +def _report(label: str, legacy: ArrayFactory, optimized: ArrayFactory, iterations: int) -> None: + np.testing.assert_allclose(legacy(), optimized(), rtol=1e-12, atol=1e-12) + legacy_seconds = _elapsed(legacy, iterations) + optimized_seconds = _elapsed(optimized, iterations) + print( + f"{label}: legacy={legacy_seconds * 1e3:.3f} ms " + f"optimized={optimized_seconds * 1e3:.3f} ms speedup={legacy_seconds / optimized_seconds:.2f}x" + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--channels", type=int, default=64) + parser.add_argument("--frames", type=int, default=30_720) + parser.add_argument("--block", type=int, default=52) + parser.add_argument("--center-iterations", type=int, default=20) + parser.add_argument("--training-iterations", type=int, default=2_000) + args = parser.parse_args() + + rng = np.random.default_rng(42) + data = rng.standard_normal((args.channels, args.frames)) + weights = rng.standard_normal((args.channels, args.channels)) * 0.01 + block_data = rng.standard_normal((args.channels, args.block)) + projected = weights @ block_data + activations = np.tanh(projected) + signs = np.diag(rng.choice((-1.0, 1.0), size=args.channels)) + identity = np.eye(args.channels) * args.block + bias = rng.standard_normal((args.channels, 1)) * 0.01 + ones = np.ones((1, args.block)) + learning_rate = 0.0001 + + _report( + "channel centering", + lambda: _legacy_center(data), + lambda: _optimized_center(data), + args.center_iterations, + ) + _report( + "extended weight update", + lambda: _legacy_extended_update(weights, identity, signs, activations, projected, learning_rate), + lambda: _optimized_extended_update(weights, identity, signs, activations, projected, learning_rate), + args.training_iterations, + ) + _report( + "bias projection", + lambda: _legacy_bias_projection(weights, block_data, bias, ones), + lambda: _optimized_bias_projection(weights, block_data, bias), + args.training_iterations, + ) + + +if __name__ == "__main__": + main()