Fixes #1153, Fixes 707 Track Random Seed (#1154)

* Initial commit

* Changed dtype

* Added scraamp-related files

* Completed up to test_aampi.py

* Completed seed tracking for all modules

* Force histogram distribution to use custom RNG

* Added no cover

* Addressed comments

* Added STUMPY_SEED support

Author: seanlaw

conftest.py

@@ -4,3 +4,13 @@
 # to fix eventual module import errors that can arise, for example when
 # running tests from inside VS code.
 # See https://stackoverflow.com/a/34520971
+
+from stumpy import rng
+
+
+def pytest_configure(config):
+    """
+    Called after command line options have been parsed
+    and all plugins and initial conftest files been loaded.
+    """
+    print(f"STUMPY_SEED={rng.SEED} pixi run tests custom {config.args[0]}")

stumpy/floss.py

@@ -7,7 +7,7 @@
 import numpy as np
 import scipy.stats
 
-from . import config, core
+from . import config, core, rng
 
 
 def _nnmark(I):
@@ -84,34 +84,36 @@ def _iac(
     IAC : numpy.ndarray
         Idealized arc curve (IAC)
     """
-    np.random.seed(seed)
+    with rng.fix_seed(seed):
+        I = rng.RNG.randint(0, width, size=width, dtype=np.int64)
+        if bidirectional is False:  # Idealized 1-dimensional matrix profile index
+            I[:-1] = width
+            for i in range(width - 1):
+                I[i] = rng.RNG.randint(i + 1, width, dtype=np.int64)
+
+        target_AC = _nnmark(I)
+
+        params = np.empty((n_iter, 2), dtype=np.float64)
+        for i in range(n_iter):
+            hist_dist = scipy.stats.rv_histogram(
+                (target_AC, np.append(np.arange(width), width))
+            )
+            hist_dist.random_state = rng.RNG
+            data = hist_dist.rvs(size=n_samples)
+            a, b, c, d = scipy.stats.beta.fit(data, floc=0, fscale=width)
 
-    I = np.random.randint(0, width, size=width, dtype=np.int64)
-    if bidirectional is False:  # Idealized 1-dimensional matrix profile index
-        I[:-1] = width
-        for i in range(width - 1):
-            I[i] = np.random.randint(i + 1, width, dtype=np.int64)
+            params[i, 0] = a
+            params[i, 1] = b
 
-    target_AC = _nnmark(I)
+        a_mean = np.round(np.mean(params[:, 0]), 2)
+        b_mean = np.round(np.mean(params[:, 1]), 2)
 
-    params = np.empty((n_iter, 2), dtype=np.float64)
-    for i in range(n_iter):
-        hist_dist = scipy.stats.rv_histogram(
-            (target_AC, np.append(np.arange(width), width))
+        IAC = scipy.stats.beta.pdf(np.arange(width), a_mean, b_mean, loc=0, scale=width)
+        slope, _, _, _ = np.linalg.lstsq(
+            np.expand_dims(IAC, axis=1), target_AC, rcond=None
         )
-        data = hist_dist.rvs(size=n_samples)
-        a, b, c, d = scipy.stats.beta.fit(data, floc=0, fscale=width)
-
-        params[i, 0] = a
-        params[i, 1] = b
-
-    a_mean = np.round(np.mean(params[:, 0]), 2)
-    b_mean = np.round(np.mean(params[:, 1]), 2)
-
-    IAC = scipy.stats.beta.pdf(np.arange(width), a_mean, b_mean, loc=0, scale=width)
-    slope, _, _, _ = np.linalg.lstsq(np.expand_dims(IAC, axis=1), target_AC, rcond=None)
 
-    IAC *= slope
+        IAC *= slope
 
     return IAC
 

stumpy/rng.py

@@ -0,0 +1,83 @@
+import os
+from contextlib import contextmanager
+
+import numpy as np
+
+# Note that an initial SEED = 0 is disallowed
+# in order to account for unit testing
+if os.getenv("STUMPY_SEED") is not None:  # pragma: no cover
+    SEED = int(os.getenv("STUMPY_SEED"))
+else:
+    SEED = np.random.randint(1, 4_294_967_296, dtype=np.uint32)
+RNG = np.random.RandomState(seed=SEED)
+
+
+def set_seed(seed):
+    """
+    Permanently set the RNG seed to a different value
+
+    Parameters
+    ----------
+    seed : int
+        The random seed for (permanently) setting the random number generator to
+
+    Returns
+    -------
+    None
+    """
+    global SEED
+    global RNG
+    SEED = seed
+    RNG = np.random.RandomState(seed=SEED)
+
+
+@contextmanager
+def fix_seed(seed):
+    """
+    A context manager for setting the RNG seed to a fixed, hardcoded, safe seed
+    and then returning the RNG back to its previous state prior to the seed change
+
+    This is typically used when you want to generate a specific random sequence once.
+    To repeat the same random sequence, use `fix_state` instead. If you are picking
+    a random seed directly before calling `fix_seed` then you probably want to use
+    `fix_state` instead!
+
+    Parameters
+    ----------
+    seed : int
+        The random seed for (temporarily) setting the random number generator to
+
+    Returns
+    -------
+    None
+    """
+    curr_state = RNG.get_state()
+    RNG.seed(seed)
+    try:
+        yield
+    finally:
+        RNG.set_state(curr_state)
+
+
+@contextmanager
+def fix_state():
+    """
+    A context manager for setting the RNG state to a fixed, hardcoded, safe state
+    and then returning the RNG back to its previous state prior to the state change
+
+    This is typically used when you want to repeat the same random sequence more than
+    once.
+
+    Parameters
+    ----------
+    None
+
+    Returns
+    -------
+    None
+    """
+    curr_state = RNG.get_state()
+    try:
+        yield
+    finally:
+        RNG.set_state(curr_state)

stumpy/scraamp.py

@@ -6,7 +6,7 @@
 import numpy as np
 from numba import njit, prange
 
-from . import config, core
+from . import config, core, rng
 from .aamp import _aamp
 
 
@@ -78,7 +78,7 @@ def _preprocess_prescraamp(T_A, m, T_B=None, s=None):
         else:  # AB-join
             s = int(np.ceil(m / config.STUMPY_EXCL_ZONE_DENOM))
 
-    indices = np.random.permutation(range(0, l, s)).astype(np.int64)
+    indices = rng.RNG.permutation(range(0, l, s)).astype(np.int64)
 
     return (T_A, T_B, T_A_subseq_isfinite, T_B_subseq_isfinite, indices, s, excl_zone)
 
@@ -718,7 +718,7 @@ def __init__(
             core._merge_topk_PI(self._P, P, self._I, I)
 
         if self._ignore_trivial:
-            self._diags = np.random.permutation(
+            self._diags = rng.RNG.permutation(
                 range(self._excl_zone + 1, self._n_A - self._m + 1)
             ).astype(np.int64)
             if self._diags.shape[0] == 0:  # pragma: no cover
@@ -728,7 +728,7 @@ def __init__(
                     f"Please try a value of `m <= {max_m}`"
                 )
         else:
-            self._diags = np.random.permutation(
+            self._diags = rng.RNG.permutation(
                 range(-(self._n_A - self._m + 1) + 1, self._n_B - self._m + 1)
             ).astype(np.int64)
 

stumpy/scrump.py

@@ -6,7 +6,7 @@
 import numpy as np
 from numba import njit, prange
 
-from . import config, core, sdp
+from . import config, core, rng, sdp
 from .scraamp import prescraamp, scraamp
 from .stump import _stump
 
@@ -116,7 +116,7 @@ def _preprocess_prescrump(
         else:  # AB-join
             s = int(np.ceil(m / config.STUMPY_EXCL_ZONE_DENOM))
 
-    indices = np.random.permutation(range(0, l, s)).astype(np.int64)
+    indices = rng.RNG.permutation(range(0, l, s)).astype(np.int64)
 
     return (
         T_A,
@@ -997,7 +997,7 @@ def __init__(
             core._merge_topk_PI(self._P, P, self._I, I)
 
         if self._ignore_trivial:
-            self._diags = np.random.permutation(
+            self._diags = rng.RNG.permutation(
                 range(self._excl_zone + 1, self._n_A - self._m + 1)
             ).astype(np.int64)
             if self._diags.shape[0] == 0:  # pragma: no cover
@@ -1007,7 +1007,7 @@ def __init__(
                     f"Please try a value of `m <= {max_m}`"
                 )
         else:
-            self._diags = np.random.permutation(
+            self._diags = rng.RNG.permutation(
                 range(-(self._n_A - self._m + 1) + 1, self._n_B - self._m + 1)
             ).astype(np.int64)
 

tests/naive.py

@@ -4,7 +4,7 @@
 from scipy.spatial.distance import cdist
 from scipy.stats import norm
 
-from stumpy import config, core
+from stumpy import config, core, rng
 
 
 def is_ptp_zero_1d(a, w):  # `a` is 1-D
@@ -1818,7 +1818,7 @@ def prescrump(
     P = np.full((l, k), np.inf, dtype=np.float64)
     I = np.full((l, k), -1, dtype=np.int64)
 
-    for i in np.random.permutation(range(0, l, s)):
+    for i in rng.RNG.permutation(range(0, l, s)):
         distance_profile = dist_matrix[i]
         if exclusion_zone is not None:
             apply_exclusion_zone(distance_profile, i, exclusion_zone, np.inf)
@@ -1914,11 +1914,11 @@ def scrump(
                     pass
 
     if exclusion_zone is not None:
-        diags = np.random.permutation(range(exclusion_zone + 1, n_A - m + 1)).astype(
+        diags = rng.RNG.permutation(range(exclusion_zone + 1, n_A - m + 1)).astype(
             np.int64
         )
     else:
-        diags = np.random.permutation(range(-(n_A - m + 1) + 1, n_B - m + 1)).astype(
+        diags = rng.RNG.permutation(range(-(n_A - m + 1) + 1, n_B - m + 1)).astype(
             np.int64
         )
 
@@ -1981,7 +1981,7 @@ def prescraamp(T_A, m, T_B, s, exclusion_zone=None, p=2.0, k=1):
     P = np.full((l, k), np.inf, dtype=np.float64)
     I = np.full((l, k), -1, dtype=np.int64)
 
-    for i in np.random.permutation(range(0, l, s)):
+    for i in rng.RNG.permutation(range(0, l, s)):
         distance_profile = distance_matrix[i]
         if exclusion_zone is not None:
             apply_exclusion_zone(distance_profile, i, exclusion_zone, np.inf)
@@ -2054,11 +2054,11 @@ def scraamp(T_A, m, T_B, percentage, exclusion_zone, pre_scraamp, s, p=2.0, k=1)
     l = n_A - m + 1
 
     if exclusion_zone is not None:
-        diags = np.random.permutation(range(exclusion_zone + 1, n_A - m + 1)).astype(
+        diags = rng.RNG.permutation(range(exclusion_zone + 1, n_A - m + 1)).astype(
             np.int64
         )
     else:
-        diags = np.random.permutation(range(-(n_A - m + 1) + 1, n_B - m + 1)).astype(
+        diags = rng.RNG.permutation(range(-(n_A - m + 1) + 1, n_B - m + 1)).astype(
             np.int64
         )
 

tests/test_aamp.py

@@ -4,7 +4,7 @@
 import pandas as pd
 import pytest
 
-from stumpy import config
+from stumpy import config, rng
 from stumpy.aamp import aamp
 
 test_data = [
@@ -13,8 +13,8 @@
         np.array([584, -11, 23, 79, 1001, 0, -19], dtype=np.float64),
     ),
     (
-        np.random.uniform(-1000, 1000, [8]).astype(np.float64),
-        np.random.uniform(-1000, 1000, [64]).astype(np.float64),
+        rng.RNG.uniform(-1000, 1000, size=8).astype(np.float64),
+        rng.RNG.uniform(-1000, 1000, size=64).astype(np.float64),
     ),
 ]
 
@@ -72,7 +72,7 @@ def test_aamp_constant_subsequence_self_join():
 
 
 def test_aamp_one_constant_subsequence_A_B_join():
-    T_A = np.random.rand(20)
+    T_A = rng.RNG.random(20)
     T_B = np.concatenate((np.zeros(20, dtype=np.float64), np.ones(5, dtype=np.float64)))
     m = 3
     ref_mp = naive.aamp(T_A, m, T_B=T_B)
@@ -122,8 +122,8 @@ def test_aamp_two_constant_subsequences_A_B_join():
 
 
 def test_aamp_identical_subsequence_self_join():
-    identical = np.random.rand(8)
-    T_A = np.random.rand(20)
+    identical = rng.RNG.random(8)
+    T_A = rng.RNG.random(20)
     T_A[1 : 1 + identical.shape[0]] = identical
     T_A[11 : 11 + identical.shape[0]] = identical
     m = 3
@@ -143,9 +143,9 @@ def test_aamp_identical_subsequence_self_join():
 
 
 def test_aamp_identical_subsequence_A_B_join():
-    identical = np.random.rand(8)
-    T_A = np.random.rand(20)
-    T_B = np.random.rand(20)
+    identical = rng.RNG.random(8)
+    T_A = rng.RNG.random(20)
+    T_B = rng.RNG.random(20)
     T_A[1 : 1 + identical.shape[0]] = identical
     T_B[11 : 11 + identical.shape[0]] = identical
     m = 3

tests/test_aamp_mmotifs.py

@@ -3,7 +3,7 @@
 import numpy.testing as npt
 import pytest
 
-from stumpy import config
+from stumpy import config, rng
 from stumpy.aamp_mmotifs import aamp_mmotifs
 
 test_data = [
@@ -60,23 +60,23 @@ def test_aamp_mmotifs_default_parameters():
         np.array([411.60964047, 423.69925001, 449.11032383, 476.95855027, 506.62406252])
     ]
 
-    np.random.seed(0)
-    T = np.random.rand(500).reshape(5, 100)
-
-    m = 5
-    excl_zone = int(np.ceil(m / config.STUMPY_EXCL_ZONE_DENOM))
-    P, I = naive.maamp(T, m, excl_zone)
-    (
-        motif_distances_cmp,
-        motif_indices_cmp,
-        motif_subspaces_cmp,
-        motif_mdls_cmp,
-    ) = aamp_mmotifs(T, P, I)
-
-    npt.assert_array_almost_equal(motif_distances_ref, motif_distances_cmp)
-    npt.assert_array_almost_equal(motif_indices_ref, motif_indices_cmp)
-    npt.assert_array_almost_equal(motif_subspaces_ref, motif_subspaces_cmp)
-    npt.assert_array_almost_equal(motif_mdls_ref, motif_mdls_cmp)
+    with rng.fix_seed(0):
+        T = rng.RNG.rand(500).reshape(5, 100)
+
+        m = 5
+        excl_zone = int(np.ceil(m / config.STUMPY_EXCL_ZONE_DENOM))
+        P, I = naive.maamp(T, m, excl_zone)
+        (
+            motif_distances_cmp,
+            motif_indices_cmp,
+            motif_subspaces_cmp,
+            motif_mdls_cmp,
+        ) = aamp_mmotifs(T, P, I)
+
+        npt.assert_array_almost_equal(motif_distances_ref, motif_distances_cmp)
+        npt.assert_array_almost_equal(motif_indices_ref, motif_indices_cmp)
+        npt.assert_array_almost_equal(motif_subspaces_ref, motif_subspaces_cmp)
+        npt.assert_array_almost_equal(motif_mdls_ref, motif_mdls_cmp)
 
 
 @pytest.mark.parametrize("T", test_data)