Refactor of vptree code. Improved modularity and performance using numba

Author: lucasimi

src/tdamapper/utils/vptree.py

@@ -2,8 +2,8 @@
 A module for fast knn and range searches, depending only on a given metric
 """
 
-from tdamapper.utils.vptree_flat import VPTree as FVPT
-from tdamapper.utils.vptree_hier import VPTree as HVPT
+from tdamapper.utils.vptree_flat.vptree import VPTree as FVPT
+from tdamapper.utils.vptree_hier.vptree import VPTree as HVPT
 
 
 class VPTree:

src/tdamapper/utils/vptree_flat.py

@@ -0,0 +1,48 @@
+from tdamapper.utils.vptree_flat.common import _mid
+
+
+class BallSearch:
+
+    def __init__(self, vpt, point, eps, inclusive=True):
+        self._arr = vpt._get_arr()
+        self.__distance = vpt._get_distance()
+        self.__point = point
+        self.__eps = eps
+        self.__inclusive = inclusive
+
+    def search(self):
+        return self._search_iter()
+
+    def _inside(self, dist):
+        if self.__inclusive:
+            return dist <= self.__eps
+        return dist < self.__eps
+
+    def _search_iter(self):
+        stack = [(0, self._arr.size())]
+        result = []
+        while stack:
+            start, end = stack.pop()
+            v_radius = self._arr.get_distance(start)
+            v_point = self._arr.get_point(start)
+            is_terminal = self._arr.is_terminal(start)
+            if is_terminal:
+                for x in self._arr.get_points(start, end):
+                    dist = self.__distance(self.__point, x)
+                    if self._inside(dist):
+                        result.append(x)
+            else:
+                dist = self.__distance(self.__point, v_point)
+                mid = _mid(start, end)
+                if self._inside(dist):
+                    result.append(v_point)
+                if dist <= v_radius:
+                    fst = (start + 1, mid)
+                    snd = (mid, end)
+                else:
+                    fst = (mid, end)
+                    snd = (start + 1, mid)
+                if abs(dist - v_radius) <= self.__eps:
+                    stack.append(snd)
+                stack.append(fst)
+        return result