Repository: mikgroup/sigpy

Issue:
RandomState not supported in numba
**Describe the bug**
An [earlier PR](https://github.com/mikgroup/sigpy/pull/55) attempted to resolve an issue where calling `sigpy.mri.samp.poisson` changed the numpy state (#54). However, `numba` only supports a select number of `numpy` operations and no longer supports the `RandomState` class in functions decorated with `nb.jit`. This is likely the source of the error.

**To Reproduce**
```python
from sigpy.mri.samp import poisson

# Not the issue occurs with any arbitrary image size and acceleration.
mask = poisson((320, 320), accel=6, seed=80)
```

**Expected behavior**
Expect to have no error in call while still ensuring that global numpy state does not change when seed is specified in `sigpy.mri.samp.poisson`.



Discussion:

Patch:
diff --git a/sigpy/mri/samp.py b/sigpy/mri/samp.py
index b2877ec..0bc79ef 100644
--- a/sigpy/mri/samp.py
+++ b/sigpy/mri/samp.py
@@ -30,6 +30,9 @@ def poisson(img_shape, accel, K=30, calib=[0, 0], dtype=np.complex,
         SIGGRAPH sketches. 2007.
 
     """
+    if seed is not None:
+        rand_state = np.random.get_state()
+
     y, x = np.mgrid[:img_shape[-2], :img_shape[-1]]
     x = np.maximum(abs(x - img_shape[-1] / 2) - calib[-1] / 2, 0)
     x /= x.max()
@@ -56,6 +59,10 @@ def poisson(img_shape, accel, K=30, calib=[0, 0], dtype=np.complex,
             slope_max = slope
 
     mask = mask.reshape(img_shape).astype(dtype)
+
+    if seed is not None:
+        np.random.set_state(rand_state)
+
     if return_density:
         return mask, R
     else:
@@ -130,24 +137,22 @@ def radial(coord_shape, img_shape, golden=True, dtype=np.float):
 
 
 @nb.jit(nopython=True, cache=True)  # pragma: no cover
-def _poisson(nx, ny, K, R, calib, seed):
+def _poisson(nx, ny, K, R, calib, seed=None):
 
     mask = np.zeros((ny, nx))
     f = ny / nx
 
     if seed is not None:
-        rand_state = np.random.RandomState(int(seed))
-    else:
-        rand_state = np.random
+        np.random.seed(int(seed))
 
     pxs = np.empty(nx * ny, np.int32)
     pys = np.empty(nx * ny, np.int32)
-    pxs[0] = rand_state.randint(0, nx)
-    pys[0] = rand_state.randint(0, ny)
+    pxs[0] = np.random.randint(0, nx)
+    pys[0] = np.random.randint(0, ny)
     m = 1
     while (m > 0):
 
-        i = rand_state.randint(0, m)
+        i = np.random.randint(0, m)
         px = pxs[i]
         py = pys[i]
         rad = R[py, px]
@@ -158,8 +163,8 @@ def _poisson(nx, ny, K, R, calib, seed):
         while not done and k < K:
 
             # Generate point randomly from R and 2R
-            rd = rad * (rand_state.random() * 3 + 1)**0.5
-            t = 2 * np.pi * rand_state.random()
+            rd = rad * (np.random.random() * 3 + 1)**0.5
+            t = 2 * np.pi * np.random.random()
             qx = px + rd * np.cos(t)
             qy = py + rd * f * np.sin(t)
 


Test patch:
diff --git a/tests/mri/test_samp.py b/tests/mri/test_samp.py
new file mode 100644
index 0000000..6e5d107
--- /dev/null
+++ b/tests/mri/test_samp.py
@@ -0,0 +1,33 @@
+import unittest
+import numpy as np
+import numpy.testing as npt
+
+from sigpy.mri import samp
+
+if __name__ == '__main__':
+    unittest.main()
+
+
+class TestPoisson(unittest.TestCase):
+    """Test poisson undersampling defined in `sigpy.mri.samp.poisson`."""
+
+    def test_numpy_random_state(self):
+        """Verify that random state is unchanged when seed is specified."""
+        np.random.seed(0)
+        expected_state = np.random.get_state()
+
+        _ = samp.poisson((320, 320), accel=6, seed=80)
+
+        state = np.random.get_state()
+        assert (expected_state[1] == state[1]).all()
+
+    def test_reproducibility(self):
+        """Verify that poisson is reproducible."""
+        np.random.seed(45)
+        mask1 = samp.poisson((320, 320), accel=6, seed=80)
+
+        # Changing internal numpy state should not affect mask.
+        np.random.seed(20)
+        mask2 = samp.poisson((320, 320), accel=6, seed=80)
+
+        npt.assert_allclose(mask2, mask1)