From 633f8e67f0e7c62e1cbbfbf472715bcc01ef3d00 Mon Sep 17 00:00:00 2001
From: Gabriel Cuendet <gabriel.cuendet@pix4d.com>
Date: Fri, 18 Oct 2019 16:28:25 +0200
Subject: [PATCH 1/2] Fix scalability issue with cKDTree.query

---
 pygsp/_nearest_neighbor.py | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/pygsp/_nearest_neighbor.py b/pygsp/_nearest_neighbor.py
index 15dd97ce..95f54764 100644
--- a/pygsp/_nearest_neighbor.py
+++ b/pygsp/_nearest_neighbor.py
@@ -54,20 +54,17 @@ def _scipy_ckdtree(features, _, order, kind, k, radius, params):
     params = dict(p=order, eps=eps, n_jobs=-1)
     if kind == 'knn':
         params['k'] = k + 1
-    elif kind == 'radius':
-        params['k'] = features.shape[0]  # number of vertices
-        params['distance_upper_bound'] = radius
-    distances, neighbors = tree.query(features, **params)
-    if kind == 'knn':
+        distances, neighbors = tree.query(features, **params)
         return neighbors, distances
     elif kind == 'radius':
+        neighbors = tree.query_ball_point(features, 
+                                          radius * np.ones((features.shape[0],)),
+                                          p=order)
         dist = []
-        neigh = []
-        for distance, neighbor in zip(distances, neighbors):
-            mask = (distance != np.inf)
-            dist.append(distance[mask])
-            neigh.append(neighbor[mask])
-        return neigh, dist
+        for i, neighbor in enumerate(neighbors):
+            dist.append(spatial.distance.cdist([features[i]],
+                                                features[neighbor]).flatten())
+        return neighbors, dist
 
 
 def _flann(features, metric, order, kind, k, radius, params):

From 38e8f94df3dd28d93a416b6a7e2db022d8386496 Mon Sep 17 00:00:00 2001
From: Gabriel Cuendet <gabriel.cuendet@pix4d.com>
Date: Fri, 18 Oct 2019 17:20:09 +0200
Subject: [PATCH 2/2] Add metrics

---
 pygsp/_nearest_neighbor.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/pygsp/_nearest_neighbor.py b/pygsp/_nearest_neighbor.py
index 95f54764..c3cb1ae5 100644
--- a/pygsp/_nearest_neighbor.py
+++ b/pygsp/_nearest_neighbor.py
@@ -46,7 +46,7 @@ def _scipy_kdtree(features, _, order, kind, k, radius, params):
     return neighbors, distances
 
 
-def _scipy_ckdtree(features, _, order, kind, k, radius, params):
+def _scipy_ckdtree(features, metric, order, kind, k, radius, params):
     if order is None:
         raise ValueError('invalid metric for scipy-kdtree')
     eps = params.pop('eps', 0)
@@ -61,9 +61,15 @@ def _scipy_ckdtree(features, _, order, kind, k, radius, params):
                                           radius * np.ones((features.shape[0],)),
                                           p=order)
         dist = []
+        metric = 'cityblock' if metric == 'manhattan' else metric
+        metric = 'chebyshev' if metric == 'max_dist' else metric
+        params = dict(metric=metric)
+        if metric == 'minkowski':
+            params['p'] = order
         for i, neighbor in enumerate(neighbors):
             dist.append(spatial.distance.cdist([features[i]],
-                                                features[neighbor]).flatten())
+                                                features[neighbor],
+                                                **params).flatten())
         return neighbors, dist