Optimize runica training kernels

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

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)

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()