diff --git a/src/eegprep/plugins/clean_rawdata/private/ransac.py b/src/eegprep/plugins/clean_rawdata/private/ransac.py index 1a80bee1..7dfa0ebe 100644 --- a/src/eegprep/plugins/clean_rawdata/private/ransac.py +++ b/src/eegprep/plugins/clean_rawdata/private/ransac.py @@ -5,7 +5,6 @@ import numpy as np from ....functions.adminfunc.eeglabcompat import get_eeglab -from ....functions.miscfunc.misc import round_mat from .sphericalSplineInterpolate import sphericalSplineInterpolate @@ -26,7 +25,7 @@ def rand_sample(n: int, m: int, stream: np.random.RandomState) -> np.ndarray: Performance: O(n) time complexity (was O(n²) in previous implementation) - For n=1M: ~3s (was ~80s) - 25x faster + Vectorized RNG calls. Note: This implementation uses Fisher-Yates shuffle for efficiency. @@ -36,11 +35,18 @@ def rand_sample(n: int, m: int, stream: np.random.RandomState) -> np.ndarray: # Start with identity permutation pool = np.arange(n) + if m <= 0: + return np.array([], dtype=int) + + # Pre-generate random numbers to avoid scalar RNG overhead in the loop. + rands = stream.rand(m) + # Fisher-Yates shuffle: only shuffle first m elements for k in range(m): # Choose from remaining elements (k to n-1) remaining = n - k - choice = int(round_mat((remaining - 1) * stream.rand())) + # Vectorized equivalent of int(round_mat((remaining - 1) * stream.rand())) + choice = int(np.floor((remaining - 1) * rands[k] + 0.5)) # Swap pool[k] with pool[k + choice] idx = k + choice @@ -69,7 +75,7 @@ def rand_permutation(n: int, stream: np.random.RandomState) -> np.ndarray: Performance: O(n) time complexity (was O(n²)) - For n=1M: ~3s (was ~80s) - 25x faster + Vectorized RNG calls. Example: >>> rng = np.random.RandomState(5489) @@ -86,10 +92,18 @@ def rand_permutation(n: int, stream: np.random.RandomState) -> np.ndarray: # Start with identity permutation [0, 1, 2, ..., n-1] result = np.arange(n) + if n <= 1: + return result + + # Pre-generate random numbers to avoid scalar RNG overhead in the loop. + rands = stream.rand(n - 1) + # Fisher-Yates shuffle: iterate backward from n-1 to 1 - for k in range(n - 1, 0, -1): + # We use rands in the same order as stream.rand() would be called + for i, k in enumerate(range(n - 1, 0, -1)): # Pick random index from 0 to k (inclusive) - j = int(round_mat(k * stream.rand())) + # Vectorized equivalent of int(round_mat(k * stream.rand())) + j = int(np.floor(k * rands[i] + 0.5)) # Swap elements k and j result[k], result[j] = result[j], result[k] diff --git a/tests/test_utils_ransac.py b/tests/test_utils_ransac.py index 3e386ece..5c85f442 100644 --- a/tests/test_utils_ransac.py +++ b/tests/test_utils_ransac.py @@ -2,10 +2,28 @@ import numpy as np from unittest.mock import patch, MagicMock -from eegprep.plugins.clean_rawdata.private.ransac import rand_sample, calc_projector +from eegprep.plugins.clean_rawdata.private.ransac import rand_permutation, rand_sample, calc_projector from eegprep.plugins.clean_rawdata.private.sphericalSplineInterpolate import sphericalSplineInterpolate +def _scalar_rand_sample(n, m, stream): + pool = np.arange(n) + for k in range(m): + remaining = n - k + choice = int(np.floor((remaining - 1) * stream.rand() + 0.5)) + idx = k + choice + pool[k], pool[idx] = pool[idx], pool[k] + return pool[:m].copy() + + +def _scalar_rand_permutation(n, stream): + result = np.arange(n) + for k in range(n - 1, 0, -1): + j = int(np.floor(k * stream.rand() + 0.5)) + result[k], result[j] = result[j], result[k] + return result + + class TestRandSample(unittest.TestCase): """Test the rand_sample function for random sampling without replacement.""" @@ -94,6 +112,36 @@ def test_sampling_algorithm_coverage(self): unique_indices = set(results) self.assertEqual(len(unique_indices), n) # All indices should appear + def test_vectorized_rng_matches_scalar_reference(self): + """Test that batched random draws preserve the scalar RNG sequence.""" + for seed, n, m in ((42, 10, 5), (5489, 20, 20), (123, 8, 0)): + with self.subTest(seed=seed, n=n, m=m): + vectorized_rng = np.random.RandomState(seed) + scalar_rng = np.random.RandomState(seed) + + result = rand_sample(n, m, vectorized_rng) + expected = _scalar_rand_sample(n, m, scalar_rng) + + np.testing.assert_array_equal(result, expected) + np.testing.assert_allclose(vectorized_rng.rand(3), scalar_rng.rand(3)) + + +class TestRandPermutation(unittest.TestCase): + """Test the rand_permutation function for MATLAB-parity shuffling.""" + + def test_vectorized_rng_matches_scalar_reference(self): + """Test that batched random draws preserve the scalar RNG sequence.""" + for seed, n in ((42, 10), (5489, 20), (123, 1), (321, 0)): + with self.subTest(seed=seed, n=n): + vectorized_rng = np.random.RandomState(seed) + scalar_rng = np.random.RandomState(seed) + + result = rand_permutation(n, vectorized_rng) + expected = _scalar_rand_permutation(n, scalar_rng) + + np.testing.assert_array_equal(result, expected) + np.testing.assert_allclose(vectorized_rng.rand(3), scalar_rng.rand(3)) + class TestCalcProjector(unittest.TestCase): """Test the calc_projector function for RANSAC reconstruction matrices."""