diff --git a/src/eegprep/functions/sigprocfunc/runica.py b/src/eegprep/functions/sigprocfunc/runica.py index 5c23534b..0279abe8 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 @@ -606,7 +604,7 @@ def runica(data, **kwargs): signs_str = ' '.join([str(int(signs[k])) for k in range(ncomps)]) logger.info(f'Fixed extended-ICA sign assignments: {signs_str}') - signs = np.diag(signs) # make diagonal matrix + # Keep signs as a vector so applying them is a row-wise scaling operation. oldsigns = np.zeros_like(signs) signcount = 0 signcounts = [] @@ -679,86 +677,85 @@ def runica(data, **kwargs): timeperm = rand_permutation(datalength, rng) # Process data in blocks (MATLAB line 831) - 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) - - # 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, - ) - - # Bias update for tanh (MATLAB line 850) - # bias = bias + lrate*sum((-2*y)')'; - bias = bias + lrate * np.sum(-2 * y, axis=1, keepdims=True) - - # Add momentum if enabled (MATLAB lines 852-856) - if momentum > 0: - weights = weights + momentum * prevwtchange - prevwtchange = weights - prevweights - prevweights = weights.copy() - - # Check for weight blowup (MATLAB lines 858-861) - if np.max(np.abs(weights)) > MAX_WEIGHT: - wts_blowup = 1 - change = nochange - - # Extended-ICA kurtosis estimation (MATLAB lines 862-900) - if not wts_blowup: - # Recompute signs vector using kurtosis (MATLAB line 866) - if extblocks > 0 and blockno % extblocks == 0: - # Random subset selection or whole data (MATLAB lines 868-879) - if kurtsize < frames: - # Pick random subset (MATLAB lines 869-876) - # Use randint to avoid index overflow (rand() * datalength could equal datalength) - rp = rng.randint(1, datalength, size=kurtsize) - partact = _matmul(weights, data[:, rp[:kurtsize]]) - else: - # For small data sets, use whole data (MATLAB lines 877-878) - partact = _matmul(weights, data) - - # Compute kurtosis (MATLAB lines 880-882) - m2 = np.mean(partact**2, axis=1) ** 2 - m4 = np.mean(partact**4, axis=1) - # Add epsilon to prevent division by zero for near-zero variance components - kk = (m4 / (m2 + 1e-10)) - 3.0 # kurtosis estimates - - # Apply momentum to kurtosis (MATLAB lines 883-886) - if extmomentum: - kk = extmomentum * old_kk + (1.0 - extmomentum) * kk - old_kk = kk - - # Update signs based on kurtosis (MATLAB line 887) - signs = np.diag(np.sign(kk + signsbias)) - - # Track sign changes (MATLAB lines 888-898) - if np.array_equal(signs, oldsigns): - signcount = signcount + 1 - else: - signcount = 0 - - oldsigns = signs.copy() - signcounts.append(signcount) - - # Make kurtosis estimation less frequent if signs stable (MATLAB lines 895-898) - if signcount >= SIGNCOUNT_THRESHOLD: - extblocks = int(extblocks * SIGNCOUNT_STEP) - signcount = 0 - - # Increment block counter (MATLAB line 901) - blockno = blockno + 1 - - # Break if weights blew up (MATLAB lines 902-904) - if wts_blowup: - break + with np.errstate(divide="ignore", over="ignore", invalid="ignore"): + 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 = 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 + signed_y = signs[:, np.newaxis] * y + weights = weights + lrate * ((BI - (signed_y + u) @ u.T) @ weights) + + # Bias update for tanh (MATLAB line 850) + # bias = bias + lrate*sum((-2*y)')'; + bias = bias + lrate * np.sum(-2.0 * y, axis=1, keepdims=True) + + # Add momentum if enabled (MATLAB lines 852-856) + if momentum > 0: + weights = weights + momentum * prevwtchange + prevwtchange = weights - prevweights + prevweights = weights.copy() + + # Check for weight blowup (MATLAB lines 858-861) + if np.max(np.abs(weights)) > MAX_WEIGHT: + wts_blowup = 1 + change = nochange + + # Extended-ICA kurtosis estimation (MATLAB lines 862-900) + if not wts_blowup: + # Recompute signs vector using kurtosis (MATLAB line 866) + if extblocks > 0 and blockno % extblocks == 0: + # Random subset selection or whole data (MATLAB lines 868-879) + if kurtsize < frames: + # Pick random subset (MATLAB lines 869-876) + # Use randint to avoid index overflow (rand() * datalength could equal datalength) + rp = rng.randint(1, datalength, size=kurtsize) + partact = weights @ data[:, rp[:kurtsize]] + else: + # For small data sets, use whole data (MATLAB lines 877-878) + partact = weights @ data + + # Compute kurtosis (MATLAB lines 880-882) + m2 = np.mean(partact**2, axis=1) ** 2 + m4 = np.mean(partact**4, axis=1) + # Add epsilon to prevent division by zero for near-zero variance components + kk = (m4 / (m2 + 1e-10)) - 3.0 # kurtosis estimates + + # Apply momentum to kurtosis (MATLAB lines 883-886) + if extmomentum: + kk = extmomentum * old_kk + (1.0 - extmomentum) * kk + old_kk = kk + + # Update signs based on kurtosis (MATLAB line 887) + signs = np.sign(kk + signsbias) + + # Track sign changes (MATLAB lines 888-898) + if np.array_equal(signs, oldsigns): + signcount = signcount + 1 + else: + signcount = 0 + + oldsigns = signs.copy() + signcounts.append(signcount) + + # Make kurtosis estimation less frequent if signs stable (MATLAB lines 895-898) + if signcount >= SIGNCOUNT_THRESHOLD: + extblocks = int(extblocks * SIGNCOUNT_STEP) + signcount = 0 + + # Increment block counter (MATLAB line 901) + blockno = blockno + 1 + + # Break if weights blew up (MATLAB lines 902-904) + if wts_blowup: + break - # End of block loop (MATLAB line 905) + # End of block loop (MATLAB line 905) # Compute weight changes if no blowup (MATLAB lines 907-917) if not wts_blowup: @@ -800,7 +797,7 @@ def runica(data, **kwargs): signs_vec = np.ones(ncomps) for k in range(nsub): signs_vec[k] = -1 - signs = np.diag(signs_vec) + signs = signs_vec oldsigns = np.zeros_like(signs) # Check if we can continue (MATLAB lines 947-960) @@ -866,45 +863,47 @@ def runica(data, **kwargs): timeperm = rand_permutation(datalength, rng) # Process data in blocks (MATLAB line 1007) - 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) - - # Apply logistic nonlinearity (MATLAB line 1022) - # Clip u to prevent overflow in exp - u = np.maximum(u, -MAX_WEIGHT) - u = np.minimum(u, MAX_WEIGHT) - y = 1.0 / (1.0 + np.exp(-u)) - - # 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) - - # Bias update (MATLAB line 1024) - # bias = bias + lrate*sum((1-2*y)')'; - bias = bias + lrate * np.sum(1 - 2 * y, axis=1, keepdims=True) - - # Add momentum if enabled (MATLAB lines 1026-1030) - if momentum > 0: - weights = weights + momentum * prevwtchange - prevwtchange = weights - prevweights - prevweights = weights.copy() - - # Check for weight blowup (MATLAB lines 1032-1035) - if np.max(np.abs(weights)) > MAX_WEIGHT: - wts_blowup = 1 - change = nochange - - # Increment block counter (MATLAB line 1036) - blockno = blockno + 1 - - # Break if weights blew up (MATLAB lines 1037-1039) - if wts_blowup: - break + with np.errstate(divide="ignore", over="ignore", invalid="ignore"): + 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 = weights @ data[:, timeperm[t : t + block]] + bias + + # Apply logistic nonlinearity (MATLAB line 1022) + # Clip u to prevent overflow in exp + u = np.maximum(u, -MAX_WEIGHT) + u = np.minimum(u, MAX_WEIGHT) + y = 1.0 / (1.0 + np.exp(-u)) + + # Natural gradient weight update (MATLAB line 1023) + # weights = weights + lrate*(BI+(1-2*y)*u')*weights + y_update = 1.0 - 2.0 * y + weights = weights + lrate * ((BI + y_update @ u.T) @ weights) + + # Bias update (MATLAB line 1024) + # bias = bias + lrate*sum((1-2*y)')'; + bias = bias + lrate * np.sum(y_update, axis=1, keepdims=True) + + # Add momentum if enabled (MATLAB lines 1026-1030) + if momentum > 0: + weights = weights + momentum * prevwtchange + prevwtchange = weights - prevweights + prevweights = weights.copy() + + # Check for weight blowup (MATLAB lines 1032-1035) + if np.max(np.abs(weights)) > MAX_WEIGHT: + wts_blowup = 1 + change = nochange + + # Increment block counter (MATLAB line 1036) + blockno = blockno + 1 + + # Break if weights blew up (MATLAB lines 1037-1039) + if wts_blowup: + break - # End of block loop (MATLAB line 1040) + # End of block loop (MATLAB line 1040) # Compute weight changes if no blowup (MATLAB lines 1042-1052) if not wts_blowup: @@ -1009,69 +1008,68 @@ def runica(data, **kwargs): timeperm = rand_permutation(datalength, rng) # Process data in blocks (MATLAB line 1131) - for t in range(0, lastt, block): - # Extract and process block - NO BIAS (MATLAB line 1145) - u = _matmul(weights, data[:, timeperm[t : t + block]]) - - # Apply tanh nonlinearity (MATLAB line 1146) - 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, - ) - - # NO BIAS UPDATE for no-bias variant - - # Add momentum if enabled (MATLAB lines 1149-1153) - if momentum > 0: - weights = weights + momentum * prevwtchange - prevwtchange = weights - prevweights - prevweights = weights.copy() - - # Check for weight blowup (MATLAB lines 1155-1158) - if np.max(np.abs(weights)) > MAX_WEIGHT: - wts_blowup = 1 - change = nochange - - # Extended-ICA kurtosis estimation (MATLAB lines 1159-1197) - if not wts_blowup: - if extblocks > 0 and blockno % extblocks == 0: - if kurtsize < frames: - # Use randint to avoid index overflow (rand() * datalength could equal datalength) - rp = rng.randint(1, datalength, size=kurtsize) - partact = _matmul(weights, data[:, rp[:kurtsize]]) - else: - partact = _matmul(weights, data) - - m2 = np.mean(partact**2, axis=1) ** 2 - m4 = np.mean(partact**4, axis=1) - # Add epsilon to prevent division by zero for near-zero variance components - kk = (m4 / (m2 + 1e-10)) - 3.0 - - if extmomentum: - kk = extmomentum * old_kk + (1.0 - extmomentum) * kk - old_kk = kk - - signs = np.diag(np.sign(kk + signsbias)) - - if np.array_equal(signs, oldsigns): - signcount = signcount + 1 - else: - signcount = 0 - - oldsigns = signs.copy() - signcounts.append(signcount) - - if signcount >= SIGNCOUNT_THRESHOLD: - extblocks = int(extblocks * SIGNCOUNT_STEP) - signcount = 0 - - blockno = blockno + 1 - - if wts_blowup: - break + with np.errstate(divide="ignore", over="ignore", invalid="ignore"): + for t in range(0, lastt, block): + # Extract and process block - NO BIAS (MATLAB line 1145) + u = weights @ data[:, timeperm[t : t + block]] + + # Apply tanh nonlinearity (MATLAB line 1146) + y = np.tanh(u) + + # Extended-ICA natural gradient weight update (MATLAB line 1147) + signed_y = signs[:, np.newaxis] * y + weights = weights + lrate * ((BI - (signed_y + u) @ u.T) @ weights) + + # NO BIAS UPDATE for no-bias variant + + # Add momentum if enabled (MATLAB lines 1149-1153) + if momentum > 0: + weights = weights + momentum * prevwtchange + prevwtchange = weights - prevweights + prevweights = weights.copy() + + # Check for weight blowup (MATLAB lines 1155-1158) + if np.max(np.abs(weights)) > MAX_WEIGHT: + wts_blowup = 1 + change = nochange + + # Extended-ICA kurtosis estimation (MATLAB lines 1159-1197) + if not wts_blowup: + if extblocks > 0 and blockno % extblocks == 0: + if kurtsize < frames: + # Use randint to avoid index overflow (rand() * datalength could equal datalength) + rp = rng.randint(1, datalength, size=kurtsize) + partact = weights @ data[:, rp[:kurtsize]] + else: + partact = weights @ data + + m2 = np.mean(partact**2, axis=1) ** 2 + m4 = np.mean(partact**4, axis=1) + # Add epsilon to prevent division by zero for near-zero variance components + kk = (m4 / (m2 + 1e-10)) - 3.0 + + if extmomentum: + kk = extmomentum * old_kk + (1.0 - extmomentum) * kk + old_kk = kk + + signs = np.sign(kk + signsbias) + + if np.array_equal(signs, oldsigns): + signcount = signcount + 1 + else: + signcount = 0 + + oldsigns = signs.copy() + signcounts.append(signcount) + + if signcount >= SIGNCOUNT_THRESHOLD: + extblocks = int(extblocks * SIGNCOUNT_STEP) + signcount = 0 + + blockno = blockno + 1 + + if wts_blowup: + break # Compute weight changes if no blowup (MATLAB lines 1204-1214) if not wts_blowup: @@ -1107,7 +1105,7 @@ def runica(data, **kwargs): signs_vec = np.ones(ncomps) for k in range(nsub): signs_vec[k] = -1 - signs = np.diag(signs_vec) + signs = signs_vec oldsigns = np.zeros_like(signs) if lrate > MIN_LRATE: @@ -1171,35 +1169,37 @@ def runica(data, **kwargs): timeperm = rand_permutation(datalength, rng) # Process data in blocks (MATLAB line 1302) - for t in range(0, lastt, block): - # Extract and process block - NO BIAS (MATLAB line 1315) - u = _matmul(weights, data[:, timeperm[t : t + block]]) + with np.errstate(divide="ignore", over="ignore", invalid="ignore"): + for t in range(0, lastt, block): + # Extract and process block - NO BIAS (MATLAB line 1315) + u = weights @ data[:, timeperm[t : t + block]] - # Apply logistic nonlinearity (MATLAB line 1316) - u = np.maximum(u, -MAX_WEIGHT) - u = np.minimum(u, MAX_WEIGHT) - y = 1.0 / (1.0 + np.exp(-u)) + # Apply logistic nonlinearity (MATLAB line 1316) + u = np.maximum(u, -MAX_WEIGHT) + u = np.minimum(u, MAX_WEIGHT) + 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) + # Natural gradient weight update (MATLAB line 1317) + y_update = 1.0 - 2.0 * y + weights = weights + lrate * ((BI + y_update @ u.T) @ weights) - # NO BIAS UPDATE for no-bias variant + # NO BIAS UPDATE for no-bias variant - # Add momentum if enabled (MATLAB lines 1319-1323) - if momentum > 0: - weights = weights + momentum * prevwtchange - prevwtchange = weights - prevweights - prevweights = weights.copy() + # Add momentum if enabled (MATLAB lines 1319-1323) + if momentum > 0: + weights = weights + momentum * prevwtchange + prevwtchange = weights - prevweights + prevweights = weights.copy() - # Check for weight blowup (MATLAB lines 1325-1328) - if np.max(np.abs(weights)) > MAX_WEIGHT: - wts_blowup = 1 - change = nochange + # Check for weight blowup (MATLAB lines 1325-1328) + if np.max(np.abs(weights)) > MAX_WEIGHT: + wts_blowup = 1 + change = nochange - blockno = blockno + 1 + blockno = blockno + 1 - if wts_blowup: - break + if wts_blowup: + break # Compute weight changes if no blowup (MATLAB lines 1336-1346) if not wts_blowup: @@ -1304,14 +1304,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 @@ -1361,8 +1359,8 @@ def runica(data, **kwargs): weights = weights[windex, :] # reorder the weight matrix (MATLAB line 1527) bias = bias[windex] # reorder bias (MATLAB line 1528) - # Convert signs diagonal matrix to vector and reorder (MATLAB lines 1529-1530) - signs_vec = np.diag(signs) # vectorize the signs matrix + # Convert signs vector to reordered vector (MATLAB lines 1529-1530) + signs_vec = signs signs_vec = signs_vec[windex] # reorder them # Prepare final outputs diff --git a/tests/test_runica.py b/tests/test_runica.py index d4b0ee60..1dd74b5d 100644 --- a/tests/test_runica.py +++ b/tests/test_runica.py @@ -90,6 +90,40 @@ def test_sample_data_extended_ica_does_not_surface_finite_matmul_warnings(self): self.assertTrue(np.isfinite(signs).all()) self.assertFalse([warning for warning in captured if "matmul" in str(warning.message)]) + def test_strict_numpy_error_policy_does_not_escape_matmul_products(self): + """Backend floating-point status from ICA products stays internal.""" + + class StrictMatmulArray(np.ndarray): + def __new__(cls, values): + return np.asarray(values, dtype=float).view(cls) + + def __matmul__(self, other): + # Reliably emulate a BLAS backend surfacing a finite-product + # floating-point status flag under a strict NumPy policy. + np.multiply(np.finfo(float).max, 2.0) + return np.asarray(self) @ np.asarray(other) + + data = np.random.RandomState(0).standard_normal((2, 20)) + initial_weights = StrictMatmulArray(np.eye(2)) + + previous_policy = np.seterr(over="raise") + try: + weights, sphere, _compvars, bias, signs, _lrates = runica( + data, + weights=initial_weights, + sphering="none", + maxsteps=1, + verbose=False, + rndreset="off", + ) + finally: + np.seterr(**previous_policy) + + self.assertTrue(np.isfinite(weights).all()) + self.assertTrue(np.isfinite(sphere).all()) + self.assertTrue(np.isfinite(bias).all()) + self.assertTrue(np.isfinite(signs).all()) + def test_pca_reduction(self): """Test PCA dimension reduction.""" np.random.seed(42) 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..5fee715c --- /dev/null +++ b/tools/benchmark_runica.py @@ -0,0 +1,186 @@ +"""Repeatable microbenchmarks for the allocation kernels optimized in RunICA. + +Run from the repository root with ``python tools/benchmark_runica.py``. +This reports timings for human inspection; it is not a CI performance test. +""" + +from __future__ import annotations + +import argparse +from collections.abc import Callable +from time import perf_counter + +import numpy as np + + +Benchmark = Callable[[], object] + + +def _elapsed(operation: Benchmark, iterations: int) -> float: + start = perf_counter() + for _ in range(iterations): + operation() + 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_block_projection(weights: np.ndarray, blocks: np.ndarray) -> float: + checksum = 0.0 + for block_data in blocks: + with np.errstate(divide="ignore", over="ignore", invalid="ignore"): + checksum += float(np.sum(weights @ block_data)) + return checksum + + +def _optimized_block_projection(weights: np.ndarray, blocks: np.ndarray) -> float: + checksum = 0.0 + with np.errstate(divide="ignore", over="ignore", invalid="ignore"): + for block_data in blocks: + checksum += float(np.sum(weights @ block_data)) + return checksum + + +def _guarded_matmul(left: np.ndarray, right: np.ndarray) -> np.ndarray: + with np.errstate(divide="ignore", over="ignore", invalid="ignore"): + return left @ right + + +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 + - _guarded_matmul(_guarded_matmul(signs, activations), projected.T) + - _guarded_matmul(projected, projected.T) + ) + return weights + learning_rate * _guarded_matmul(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: + with np.errstate(divide="ignore", over="ignore", invalid="ignore"): + signed_activations = 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 _guarded_matmul(weights, block_data) + _guarded_matmul(bias, ones) + + +def _optimized_bias_projection( + weights: np.ndarray, + block_data: np.ndarray, + bias: np.ndarray, +) -> np.ndarray: + with np.errstate(divide="ignore", over="ignore", invalid="ignore"): + return weights @ block_data + bias + + +def _report(label: str, legacy: Benchmark, optimized: Benchmark, 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 " + f"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("--projection-iterations", type=int, default=10) + 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 = rng.choice((-1.0, 1.0), size=args.channels) + signs_matrix = np.diag(signs) + identity = np.eye(args.channels) * args.block + bias = rng.standard_normal((args.channels, 1)) * 0.01 + ones = np.ones((1, args.block)) + training_frames = args.frames // args.block * args.block + blocks = data[:, :training_frames].reshape(args.channels, -1, args.block).transpose(1, 0, 2) + learning_rate = 0.0001 + + _report( + "channel centering", + lambda: _legacy_center(data), + lambda: _optimized_center(data), + args.center_iterations, + ) + _report( + "training block projection", + lambda: _legacy_block_projection(weights, blocks), + lambda: _optimized_block_projection(weights, blocks), + args.projection_iterations, + ) + _report( + "extended weight update", + lambda: _legacy_extended_update( + weights, + identity, + signs_matrix, + 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()