Fix poisson disc sampling mask generation

Before fixing, the function segfaults or stalls for some rectangular shapes.
The main reason is that poisson disc mask requires a radius parameter
such that each sample is away from each other.
The problem was for rectangular shapes, the radius parameter is
adjusted such that it was not stable.

sigpy/mri/samp.py

@@ -8,19 +8,30 @@
 __all__ = ['poisson', 'radial', 'spiral']
 
 
-def poisson(img_shape, accel, K=30, calib=(0, 0), dtype=np.complex,
-            crop_corner=True, return_density=False, seed=0):
-    """Generate Poisson-disc sampling pattern
+def poisson(img_shape, accel, calib=(0, 0), dtype=np.complex,
+            crop_corner=True, return_density=False, seed=0,
+            max_attempts=30, tol=0.1):
+    """Generate variable-density Poisson-disc sampling pattern.
+
+    The function generates a variable density Poisson-disc sampling
+    mask with density proportional to :math:`1 / (1 + s * |r|)`,
+    where :math:`r` represents the k-space radius, and :math:`s`
+    represents the slope. A binary search is performed on the slope :math:`s`
+    such that the resulting acceleration factor is close to the
+    prescribed acceleration factor `accel`. The parameter `tol`
+    determines how much they can deviate.
 
     Args:
         img_shape (tuple of ints): length-2 image shape.
-        accel (float): Target acceleration factor. Greater than 1.
-        K (float): maximum number of samples to reject.
+        accel (float): Target acceleration factor. Must be greater than 1.
         calib (tuple of ints): length-2 calibration shape.
         dtype (Dtype): data type.
         crop_corner (bool): Toggle whether to crop sampling corners.
-        return_density (bool): Toggle whether to return sampling density.
         seed (int): Random seed.
+        max_attempts (float): maximum number of samples to reject in Poisson
+           disc calculation.
+        tol (float): Tolerance for how much the resulting acceleration can
+            deviate form `accel`.
 
     Returns:
         array: Poisson-disc sampling mask.
@@ -33,40 +44,44 @@ def poisson(img_shape, accel, K=30, calib=(0, 0), dtype=np.complex,
     if seed is not None:
         rand_state = np.random.get_state()
 
-    y, x = np.mgrid[:img_shape[-2], :img_shape[-1]]
+    ny, nx = img_shape
+    y, x = np.mgrid[:ny, :nx]
     x = np.maximum(abs(x - img_shape[-1] / 2) - calib[-1] / 2, 0)
     x /= x.max()
     y = np.maximum(abs(y - img_shape[-2] / 2) - calib[-2] / 2, 0)
     y /= y.max()
-    r = np.sqrt(x ** 2 + y ** 2)
+    r = np.sqrt(x**2 + y**2)
 
-    slope_max = 40
+    slope_max = max(nx, ny)
     slope_min = 0
     while slope_min < slope_max:
-        slope = (slope_max + slope_min) / 2.0
-        R = (1.0 + r * slope)
-        mask = _poisson(img_shape[-1], img_shape[-2], K, R, calib, seed)
+        slope = (slope_max + slope_min) / 2
+        radius_x = np.clip((1 + r * slope) * nx / max(nx, ny), 1, None)
+        radius_y = np.clip((1 + r * slope) * ny / max(nx, ny), 1, None)
+        mask = _poisson(
+            img_shape[-1], img_shape[-2], max_attempts,
+            radius_x, radius_y, calib, seed)
         if crop_corner:
             mask *= r < 1
 
-        est_accel = img_shape[-1] * img_shape[-2] / np.sum(mask[:])
+        actual_accel = img_shape[-1] * img_shape[-2] / np.sum(mask)
 
-        if abs(est_accel - accel) < 0.1:
+        if abs(actual_accel - accel) < tol:
             break
-        if est_accel < accel:
+        if actual_accel < accel:
             slope_min = slope
         else:
             slope_max = slope
 
+    if abs(actual_accel - accel) >= tol:
+        raise ValueError(f'Cannot generate mask to satisfy accel={accel}.')
+
     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:
-        return mask
+    return mask
 
 
 def radial(coord_shape, img_shape, golden=True, dtype=np.float):
@@ -137,69 +152,63 @@ 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=None):
-
+def _poisson(nx, ny, max_attempts, radius_x, radius_y, calib, seed=None):
     mask = np.zeros((ny, nx))
-    f = ny / nx
 
     if seed is not None:
         np.random.seed(int(seed))
 
+    # initialize active list
     pxs = np.empty(nx * ny, np.int32)
     pys = np.empty(nx * ny, np.int32)
     pxs[0] = np.random.randint(0, nx)
     pys[0] = np.random.randint(0, ny)
-    m = 1
-    while (m > 0):
-
-        i = np.random.randint(0, m)
+    num_actives = 1
+    while num_actives > 0:
+        i = np.random.randint(0, num_actives)
         px = pxs[i]
         py = pys[i]
-        rad = R[py, px]
+        rx = radius_x[py, px]
+        ry = radius_y[py, px]
 
         # Attempt to generate point
         done = False
         k = 0
-        while not done and k < K:
-
-            # Generate point randomly from R and 2R
-            rd = rad * (np.random.random() * 3 + 1)**0.5
+        while not done and k < max_attempts:
+            # Generate point randomly from r and 2 * r
+            v = (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)
+            qx = px + v * rx * np.cos(t)
+            qy = py + v * ry * np.sin(t)
 
             # Reject if outside grid or close to other points
             if qx >= 0 and qx < nx and qy >= 0 and qy < ny:
-
-                startx = max(int(qx - rad), 0)
-                endx = min(int(qx + rad + 1), nx)
-                starty = max(int(qy - rad * f), 0)
-                endy = min(int(qy + rad * f + 1), ny)
+                startx = max(int(qx - rx), 0)
+                endx = min(int(qx + rx + 1), nx)
+                starty = max(int(qy - ry), 0)
+                endy = min(int(qy + ry + 1), ny)
 
                 done = True
                 for x in range(startx, endx):
                     for y in range(starty, endy):
                         if (mask[y, x] == 1
-                            and (((qx - x) / R[y, x]) ** 2 +
-                                 ((qy - y) / (R[y, x] * f)) ** 2 < 1)):
+                            and (((qx - x) / radius_x[y, x])**2 +
+                                 ((qy - y) / (radius_y[y, x]))**2 < 1)):
                             done = False
                             break
 
-                    if not done:
-                        break
-
             k += 1
 
-        # Add point if done else remove active
+        # Add point if done else remove from active list
         if done:
-            pxs[m] = qx
-            pys[m] = qy
+            pxs[num_actives] = qx
+            pys[num_actives] = qy
             mask[int(qy), int(qx)] = 1
-            m += 1
+            num_actives += 1
         else:
-            pxs[i] = pxs[m - 1]
-            pys[i] = pys[m - 1]
-            m -= 1
+            pxs[i] = pxs[num_actives - 1]
+            pys[i] = pys[num_actives - 1]
+            num_actives -= 1
 
     # Add calibration region
     mask[int(ny / 2 - calib[-2] / 2):int(ny / 2 + calib[-2] / 2),

tests/mri/test_samp.py

@@ -16,18 +16,28 @@ def test_numpy_random_state(self):
         np.random.seed(0)
         expected_state = np.random.get_state()
 
-        _ = samp.poisson((320, 320), accel=6, seed=80)
+        _ = samp.poisson((120, 120), 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)
+        mask1 = samp.poisson((120, 120), 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)
+        mask2 = samp.poisson((120, 120), accel=6, seed=80)
 
         npt.assert_allclose(mask2, mask1)
+
+    def test_poisson_accel(self):
+        """Verify that poisson generates the correct acceleration."""
+        for x in [60, 120]:
+            for y in [60, 120]:
+                for tol in [0.1, 0.2]:
+                    for accel in [4, 5, 6, 7, 8]:
+                        mask = samp.poisson(
+                            (x, y), accel=accel, seed=80, tol=tol)
+                        assert abs(mask.size / np.sum(mask) - accel) < tol