Repository: scikit-image/scikit-image

Issue:
ensure_spacing function may be slow
## Description

The `skimage._shared.coord.ensure_spacing` function introduced in #4760 may be slow when the number of input points is large.

## Way to reproduce
```python
In [1]: import numpy as np

In [2]: from skimage._shared.coord import ensure_spacing

In [3]: for n in range(500, 10001, 500):
    ...:     print(f"n = {n}")
    ...:     x = np.random.rand(n, 2)
    ...:     %timeit ensure_spacing(x, 0.2)
    ...: 
n = 500
27.6 ms ± 857 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
n = 1000
64.7 ms ± 1.93 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
n = 1500
105 ms ± 2.27 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
n = 2000
158 ms ± 1.18 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
n = 2500
213 ms ± 7.32 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 3000
299 ms ± 11.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 3500
378 ms ± 3.05 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 4000
428 ms ± 11.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 4500
554 ms ± 14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 5000
647 ms ± 9.42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 5500
777 ms ± 7.05 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 6000
870 ms ± 13.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 6500
1.03 s ± 9.56 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 7000
1.17 s ± 12.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 7500
1.33 s ± 28 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 8000
1.45 s ± 18.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 8500
1.63 s ± 36.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 9000
1.87 s ± 50.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 9500
2.07 s ± 23.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
n = 10000
2.29 s ± 51.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```



Discussion:

Resolving patch:
diff --git a/skimage/_shared/coord.py b/skimage/_shared/coord.py
--- a/skimage/_shared/coord.py
+++ b/skimage/_shared/coord.py
@@ -2,7 +2,7 @@
 from scipy.spatial import cKDTree, distance
 
 
-def ensure_spacing(coord, spacing=1, p_norm=np.inf):
+def _ensure_spacing(coord, spacing, p_norm):
     """Returns a subset of coord where a minimum spacing is guaranteed.
 
     Parameters
@@ -24,28 +24,72 @@ def ensure_spacing(coord, spacing=1, p_norm=np.inf):
 
     """
 
-    output = coord
-    if len(coord):
-        # Use KDtree to find the peaks that are too close to each other
-        tree = cKDTree(coord)
-
-        indices = tree.query_ball_point(coord, r=spacing, p=p_norm)
-        rejected_peaks_indices = set()
-        for idx, candidates in enumerate(indices):
-            if idx not in rejected_peaks_indices:
-                # keep current point and the points at exactly spacing from it
-                candidates.remove(idx)
-                dist = distance.cdist([coord[idx]],
-                                      coord[candidates],
-                                      distance.minkowski,
-                                      p=p_norm).reshape(-1)
-                candidates = [c for c, d in zip(candidates, dist)
-                              if d < spacing]
-
-                # candidates.remove(keep)
-                rejected_peaks_indices.update(candidates)
-
-        # Remove the peaks that are too close to each other
-        output = np.delete(coord, tuple(rejected_peaks_indices), axis=0)
+    # Use KDtree to find the peaks that are too close to each other
+    tree = cKDTree(coord)
+
+    indices = tree.query_ball_point(coord, r=spacing, p=p_norm)
+    rejected_peaks_indices = set()
+    for idx, candidates in enumerate(indices):
+        if idx not in rejected_peaks_indices:
+            # keep current point and the points at exactly spacing from it
+            candidates.remove(idx)
+            dist = distance.cdist([coord[idx]],
+                                  coord[candidates],
+                                  distance.minkowski,
+                                  p=p_norm).reshape(-1)
+            candidates = [c for c, d in zip(candidates, dist)
+                          if d < spacing]
+
+            # candidates.remove(keep)
+            rejected_peaks_indices.update(candidates)
+
+    # Remove the peaks that are too close to each other
+    output = np.delete(coord, tuple(rejected_peaks_indices), axis=0)
+
+    return output
+
+
+def ensure_spacing(coords, spacing=1, p_norm=np.inf, min_split_size=50):
+    """Returns a subset of coord where a minimum spacing is guaranteed.
+
+    Parameters
+    ----------
+    coord : array_like
+        The coordinates of the considered points.
+    spacing : float
+        the maximum allowed spacing between the points.
+    p_norm : float
+        Which Minkowski p-norm to use. Should be in the range [1, inf].
+        A finite large p may cause a ValueError if overflow can occur.
+        ``inf`` corresponds to the Chebyshev distance and 2 to the
+        Euclidean distance.
+    min_split_size : int
+        Minimum split size used to process ``coord`` by batch to save
+        memory. If None, the memory saving strategy is not applied.
+
+    Returns
+    -------
+    output : array_like
+        A subset of coord where a minimum spacing is guaranteed.
+
+    """
+
+    output = coords
+    if len(coords):
+
+        coords = np.atleast_2d(coords)
+        if min_split_size is None:
+            batch_list = [coords]
+        else:
+            coord_count = len(coords)
+            split_count = int(np.log2(coord_count / min_split_size)) + 1
+            split_idx = np.cumsum(
+                [coord_count // (2 ** i) for i in range(1, split_count)])
+            batch_list = np.array_split(coords, split_idx)
+
+        output = np.zeros((0, coords.shape[1]), dtype=coords.dtype)
+        for batch in batch_list:
+            output = _ensure_spacing(np.vstack([output, batch]),
+                                     spacing, p_norm)
 
     return output


Test patch:
diff --git a/skimage/_shared/tests/test_coord.py b/skimage/_shared/tests/test_coord.py
--- a/skimage/_shared/tests/test_coord.py
+++ b/skimage/_shared/tests/test_coord.py
@@ -5,39 +5,65 @@
 
 
 @pytest.mark.parametrize("p", [1, 2, np.inf])
-def test_ensure_spacing_trivial(p):
+@pytest.mark.parametrize("size", [30, 50, None])
+def test_ensure_spacing_trivial(p, size):
     # --- Empty input
     assert ensure_spacing([], p_norm=p) == []
 
+    # --- A unique point
+    coord = np.random.randn(1, 2)
+    assert np.array_equal(coord, ensure_spacing(coord, p_norm=p,
+                                                min_split_size=size))
+
     # --- Verified spacing
     coord = np.random.randn(100, 2)
 
     # --- 0 spacing
-    assert np.array_equal(coord, ensure_spacing(coord, spacing=0, p_norm=p))
+    assert np.array_equal(coord, ensure_spacing(coord, spacing=0, p_norm=p,
+                                                min_split_size=size))
 
     # Spacing is chosen to be half the minimum distance
-    spacing = pdist(coord, minkowski, p=p).min() * 0.5
+    spacing = pdist(coord, metric=minkowski, p=p).min() * 0.5
 
-    out = ensure_spacing(coord, spacing=spacing, p_norm=p)
+    out = ensure_spacing(coord, spacing=spacing, p_norm=p,
+                         min_split_size=size)
 
     assert np.array_equal(coord, out)
 
 
 @pytest.mark.parametrize("ndim", [1, 2, 3, 4, 5])
-def test_ensure_spacing_nD(ndim):
+@pytest.mark.parametrize("size", [2, 10, None])
+def test_ensure_spacing_nD(ndim, size):
     coord = np.ones((5, ndim))
 
     expected = np.ones((1, ndim))
 
-    assert np.array_equal(ensure_spacing(coord), expected)
+    assert np.array_equal(ensure_spacing(coord, min_split_size=size), expected)
+
+
+@pytest.mark.parametrize("p", [1, 2, np.inf])
+@pytest.mark.parametrize("size", [50, 100, None])
+def test_ensure_spacing_batch_processing(p, size):
+    coord = np.random.randn(100, 2)
+
+    # --- Consider the average distance btween the point as spacing
+    spacing = np.median(pdist(coord, metric=minkowski, p=p))
+
+    expected = ensure_spacing(coord, spacing=spacing, p_norm=p)
+
+    assert np.array_equal(ensure_spacing(coord, spacing=spacing, p_norm=p,
+                                         min_split_size=size),
+                          expected)
+
 
 
 @pytest.mark.parametrize("p", [1, 2, np.inf])
-def test_ensure_spacing_p_norm(p):
+@pytest.mark.parametrize("size", [30, 50, None])
+def test_ensure_spacing_p_norm(p, size):
     coord = np.random.randn(100, 2)
 
     # --- Consider the average distance btween the point as spacing
-    spacing = np.median(pdist(coord, minkowski, p))
-    out = ensure_spacing(coord, spacing=spacing, p_norm=p)
+    spacing = np.median(pdist(coord, metric=minkowski, p=p))
+    out = ensure_spacing(coord, spacing=spacing, p_norm=p, min_split_size=size)
 
-    assert pdist(out, minkowski, p).min() > spacing
+    assert pdist(out, metric=minkowski, p=p).min() > spacing