Added ability to set state like np.random.seed(seed)

Author: seanlaw

conftest.py

@@ -14,6 +14,18 @@
 #     Set RNG state here
 # }
 
+pytest.LAST_RNG_STATE = None
+pytest.FIXED_RNG_STATE = {
+    # DO NOT CHANGE/ALTER!!
+    "bit_generator": "PCG64",
+    "state": {
+        "state": 195349167630453735115769518810051464980,
+        "inc": 247589055400886363559049235690497450585,
+    },
+    "has_uint32": 0,
+    "uinteger": 0,
+}
+
 
 def pytest_configure(config):
     """
@@ -22,3 +34,44 @@ def pytest_configure(config):
     """
     state = json.dumps(pytest.RNG.bit_generator.state, indent=4)
     print(f"conftest.py: pytest.RNG.bit_generator.state = {state}")
+
+
+def _get_rng_state():
+    """
+    Get a copy of the current RNG state
+    """
+    return pytest.RNG.bit_generator.state.copy()
+
+
+def _set_rng_state(state):
+    """
+    Store existing RNG state and set RNG state
+
+    If no state is provided, a hardcoded, safe state will be used
+    """
+    if state is None:
+        state = pytest.SEED_ZERO.copy()
+
+    pytest.LAST_RNG_STATE = pytest.RNG.bit_generator.state.copy()
+    pytest.RNG.bit_generator.state = state
+
+
+def _reset_rng_state():
+    """
+    Restore the RNG state to the last recorded RNG state
+    """
+    pytest.RNG.bit_generator.state = pytest.LAST_RNG_STATE.copy()
+    pytest.LAST_RNG_STATE = None
+
+
+def _fix_rng_state():
+    """
+    Set the RNG state to a fixed, hardcoded, safe state
+    """
+    _set_rng_state(pytest.FIXED_RNG_STATE)
+
+
+pytest.get_rng_state = _get_rng_state
+pytest.set_rng_state = _set_rng_state
+pytest.fix_rng_state = _fix_rng_state
+pytest.unfix_rng_state = _reset_rng_state

tests/test_core.py

@@ -1906,7 +1906,7 @@ def test_update_incremental_PI_egressTrue_MemoryCheck():
     # a new data point is appended. However, the updated matrix profile index for the
     # middle subsequence `s` should still refer to  the first subsequence in
     # the historical data.
-
+    pytest.fix_rng_state()
     T = pytest.RNG.random(64)
     m = 3
     excl_zone = int(np.ceil(m / config.STUMPY_EXCL_ZONE_DENOM))
@@ -1965,6 +1965,8 @@ def test_update_incremental_PI_egressTrue_MemoryCheck():
         npt.assert_almost_equal(P_ref, P_comp)
         npt.assert_almost_equal(I_ref, I_comp)
 
+    pytest.unfix_rng_state()
+
 
 def test_check_self_join():
     with pytest.warns(UserWarning):