FIX sort radius neighbors results when sort_results=True and algorithm="brute" (scikit-learn#18612)

Co-authored-by: Nicolas Hug <[email protected]>

Author: TomDLT

doc/whats_new/v0.24.rst

@@ -542,6 +542,14 @@ Changelog
   will raise a `ValueError` when fitting on data with all constant features.
   :pr:`18370` by :user:`Trevor Waite <trewaite>`.
 
+- |Fix| In  methods `radius_neighbors` and
+  `radius_neighbors_graph` of :class:`neighbors.NearestNeighbors`,
+  :class:`neighbors.RadiusNeighborsClassifier`,
+  :class:`neighbors.RadiusNeighborsRegressor`, and
+  :class:`neighbors.RadiusNeighborsTransformer`, using `sort_results=True` now
+  correctly sorts the results even when fitting with the "brute" algorithm.
+  :pr:`18612` by `Tom Dupre la Tour`_.
+
 :mod:`sklearn.neural_network`
 .............................
 

sklearn/neighbors/_base.py

@@ -925,10 +925,10 @@ def radius_neighbors(self, X=None, radius=None, return_distance=True,
             Whether or not to return the distances.
 
         sort_results : bool, default=False
-            If True, the distances and indices will be sorted before being
-            returned. If `False`, the results will not be sorted. If
-            `return_distance=False`, setting `sort_results=True` will
-            result in an error.
+            If True, the distances and indices will be sorted by increasing
+            distances before being returned. If False, the results may not
+            be sorted. If `return_distance=False`, setting `sort_results=True`
+            will result in an error.
 
             .. versionadded:: 0.22
 
@@ -1021,6 +1021,16 @@ class from an array representing our data set and ask who's
                 neigh_ind_list = sum(chunked_results, [])
                 results = _to_object_array(neigh_ind_list)
 
+            if sort_results:
+                if not return_distance:
+                    raise ValueError("return_distance must be True "
+                                     "if sort_results is True.")
+                for ii in range(len(neigh_dist)):
+                    order = np.argsort(neigh_dist[ii], kind='mergesort')
+                    neigh_ind[ii] = neigh_ind[ii][order]
+                    neigh_dist[ii] = neigh_dist[ii][order]
+                results = neigh_dist, neigh_ind
+
         elif self._fit_method in ['ball_tree', 'kd_tree']:
             if issparse(X):
                 raise ValueError(
@@ -1097,9 +1107,9 @@ def radius_neighbors_graph(self, X=None, radius=None, mode='connectivity',
             edges are Euclidean distance between points.
 
         sort_results : bool, default=False
-            If True, the distances and indices will be sorted before being
-            returned. If False, the results will not be sorted.
-            Only used with mode='distance'.
+            If True, in each row of the result, the non-zero entries will be
+            sorted by increasing distances. If False, the non-zero entries may
+            not be sorted. Only used with mode='distance'.
 
             .. versionadded:: 0.22
 

sklearn/neighbors/tests/test_neighbors.py

@@ -678,6 +678,39 @@ def test_radius_neighbors_returns_array_of_objects():
     assert_array_equal(neigh_ind, expected_ind)
 
 
+@pytest.mark.parametrize(["algorithm", "metric"], [("ball_tree", "euclidean"),
+                                                   ("kd_tree", "euclidean"),
+                                                   ("brute", "euclidean"),
+                                                   ("brute", "precomputed")])
+def test_radius_neighbors_sort_results(algorithm, metric):
+    # Test radius_neighbors[_graph] output when sort_result is True
+    n_samples = 10
+    rng = np.random.RandomState(42)
+    X = rng.random_sample((n_samples, 4))
+
+    if metric == "precomputed":
+        X = neighbors.radius_neighbors_graph(X, radius=np.inf, mode="distance")
+    model = neighbors.NearestNeighbors(algorithm=algorithm, metric=metric)
+    model.fit(X)
+
+    # self.radius_neighbors
+    distances, indices = model.radius_neighbors(X=X, radius=np.inf,
+                                                sort_results=True)
+    for ii in range(n_samples):
+        assert_array_equal(distances[ii], np.sort(distances[ii]))
+
+    # sort_results=True and return_distance=False
+    if metric != "precomputed":  # no need to raise with precomputed graph
+        with pytest.raises(ValueError, match="return_distance must be True"):
+            model.radius_neighbors(X=X, radius=np.inf, sort_results=True,
+                                   return_distance=False)
+
+    # self.radius_neighbors_graph
+    graph = model.radius_neighbors_graph(X=X, radius=np.inf, mode="distance",
+                                         sort_results=True)
+    assert _is_sorted_by_data(graph)
+
+
 def test_RadiusNeighborsClassifier_multioutput():
     # Test k-NN classifier on multioutput data
     rng = check_random_state(0)