Optimize RANSAC RNG sampling

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]

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."""