Completed first draft of numpy migration

Author: seanlaw

stumpy/core.py

@@ -3717,7 +3717,8 @@ def check_ignore_trivial(T_A, T_B, ignore_trivial):
     import numpy as np
     import warnings
 
-    T = np.random.rand(10_000)
+    rng = np.random.default_rng()
+    T = rng.random(10_000)
     m = 50
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore", message="Arrays T_A, T_B are equal")
@@ -4493,7 +4494,8 @@ def check_self_join(ignore_trivial):
     import numpy as np
     import warnings
 
-    T = np.random.rand(10_000)
+    rng = np.random.default_rng()
+    T = rng.random(10_000)
     m = 50
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore", message="`ignore_trivial` cannot be `False`")

stumpy/floss.py

@@ -84,13 +84,13 @@ def _iac(
     IAC : numpy.ndarray
         Idealized arc curve (IAC)
     """
-    np.random.seed(seed)
+    local_rng = np.random.default_rng(seed)
 
-    I = np.random.randint(0, width, size=width, dtype=np.int64)
+    I = local_rng.integers(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)
+            I[i] = local_rng.integers(i + 1, width, dtype=np.int64)
 
     target_AC = _nnmark(I)
 
@@ -99,7 +99,7 @@ def _iac(
         hist_dist = scipy.stats.rv_histogram(
             (target_AC, np.append(np.arange(width), width))
         )
-        data = hist_dist.rvs(size=n_samples)
+        data = hist_dist.rvs(size=n_samples, random_state=local_rng)
         a, b, c, d = scipy.stats.beta.fit(data, floc=0, fscale=width)
 
         params[i, 0] = a

stumpy/rng.py

@@ -1,3 +1,5 @@
+from contextlib import contextmanager
+
 import numpy as np
 
 bit_gen = np.random.PCG64()
@@ -85,7 +87,28 @@ def _fix_state():
     _set_state(FIXED_STATE)
 
 
+@contextmanager
+def fix_state(*args):
+    """
+    A context manager for setting the RNG state to a fixed, hardcoded, safe state
+
+    This is exceptionally rare and you probably want to use `state = get_state()`
+    and `set_state(state)` instead.
+
+    Parameters
+    ----------
+    None
+
+    Returns
+    -------
+    None
+    """
+    _fix_state()
+    try:
+        yield
+    finally:
+        _reset_state()
+
+
 get_state = _get_state
 set_state = _set_state
-fix_state = _fix_state
-unfix_state = _reset_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,
@@ -378,6 +378,7 @@ def _compute_PI(
                         P_squared[thread_idx, j], idx, squared_distance_profile[j]
                     )
                     core._shift_insert_at_index(I[thread_idx, j], idx, i)
+                    # core._shift_insert_at_index(I[thread_idx, j, :idx], idx, i)
 
 
 @njit(
@@ -997,7 +998,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 +1008,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 = [
@@ -51,32 +51,30 @@
 
 
 def test_aamp_mmotifs_default_parameters():
-    motif_distances_ref = np.array(
-        [[0.0, 0.06315749, 0.25275899, 0.34087884, 0.3452315]]
-    )
-    motif_indices_ref = np.array([[19, 77, 63, 52, 71]])
-    motif_subspaces_ref = [np.array([2])]
+    motif_distances_ref = np.array([[0.0, 0.10660435, 0.27927693]])
+    motif_indices_ref = np.array([[19, 41, 34]])
+    motif_subspaces_ref = [np.array([1])]
     motif_mdls_ref = [
-        np.array([411.60964047, 423.69925001, 449.11032383, 476.95855027, 506.62406252])
+        np.array([411.60964047, 433.21928095, 429.77443751, 457.74925683, 477.80793557])
     ]
 
-    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_state():
+        T = rng.RNG.random(size=(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)