FIX: Fix geometry.nearest_neighbors when k is bigger than the numbe…

…r of candidates
sertit · Mar 7, 2024 · 538d369 · 538d369
1 parent 7dd558d
commit 538d369
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 1 deletion.
diff --git a/CHANGES.md b/CHANGES.md
@@ -2,6 +2,7 @@
 
 ## 1.36.1 (2024-xx-xx)
 
+- FIX: Fix `geometry.nearest_neighbors` when k is bigger than the number of candidates
 - DOC: Update some examples in documentation
 
 ## 1.36.0 (2024-02-27)

diff --git a/CI/SCRIPTS/test_geometry.py b/CI/SCRIPTS/test_geometry.py
@@ -188,3 +188,8 @@ def test_nearest_neighbors():
         assert (
             curr_dist[0] < radius
         ), f"distance superior to wanted distance: {curr_dist[0]} > {radius}"
+
+    # Ensure it works with k > nof candidates
+    closest, distances = nearest_neighbors(
+        src, candidates, method="k_neighbors", k_neighbors=len(candidates) + 10
+    )
diff --git a/sertit/geometry.py b/sertit/geometry.py
@@ -539,7 +539,9 @@ def _get_k_nearest(src_points: list, candidates: list, k_neighbors: int, **kwarg
     tree = BallTree(candidates, leaf_size=15)
 
     # Find the closest points and distances
-    closest_dist, closest = tree.query(src_points, k=k_neighbors, **kwargs)
+    closest_dist, closest = tree.query(
+        src_points, k=min(k_neighbors, len(candidates)), **kwargs
+    )
 
     # Return indices and distances
     return closest, closest_dist