Refactored quickselect using numba

Author: lucasimi

src/tdamapper/utils/quickselect.py

@@ -1,54 +1,41 @@
-def __swap(arr, i, j):
-    arr[i], arr[j] = arr[j], arr[i]
+from numba import njit
 
 
-def partition(data, start, end, p_ord):
-    higher = start
-    for j in range(start, end):
-        j_ord, _ = data[j]
-        if j_ord < p_ord:
-            __swap(data, higher, j)
-            higher += 1
-    return higher
+@njit
+def swap(arr, i, j):
+    arr[i], arr[j] = arr[j], arr[i]
 
 
-def quickselect(data, start, end, k):
-    if (k < start) or (k >= end):
-        return
-    start_, end_, higher = start, end, None
-    while higher != k + 1:
-        p, _ = data[k]
-        __swap(data, start_, k)
-        higher = partition(data, start_ + 1, end_, p)
-        __swap(data, start_, higher - 1)
-        if k <= higher - 1:
-            end_ = higher
-        else:
-            start_ = higher
+@njit
+def swap_all(arr, i, j, extra1=None, extra2=None):
+    swap(arr, i, j)
+    if extra1 is not None:
+        swap(extra1, i, j)
+    if extra2 is not None:
+        swap(extra2, i, j)
 
 
-def partition_tuple(data_ord, data_arr, start, end, p_ord):
+@njit
+def partition(data, start, end, p_ord, *extra):
     higher = start
     for j in range(start, end):
-        j_ord = data_ord[j]
+        j_ord = data[j]
         if j_ord < p_ord:
-            __swap(data_arr, higher, j)
-            __swap(data_ord, higher, j)
+            swap_all(data, higher, j, *extra)
             higher += 1
     return higher
 
 
-def quickselect_tuple(data_ord, data_arr, start, end, k):
+@njit
+def quickselect(data, start, end, k, *extra):
     if (k < start) or (k >= end):
         return
     start_, end_, higher = start, end, None
     while higher != k + 1:
-        p_ord = data_ord[k]
-        __swap(data_arr, start_, k)
-        __swap(data_ord, start_, k)
-        higher = partition_tuple(data_ord, data_arr, start_ + 1, end_, p_ord)
-        __swap(data_arr, start_, higher - 1)
-        __swap(data_ord, start_, higher - 1)
+        p = data[k]
+        swap_all(data, start_, k, *extra)
+        higher = partition(data, start_ + 1, end_, p, *extra)
+        swap_all(data, start_, higher - 1, *extra)
         if k <= higher - 1:
             end_ = higher
         else:

src/tdamapper/utils/vptree_flat.py

@@ -1,60 +1,16 @@
 from random import randrange
 
 import numpy as np
-from numba import njit
 
 from tdamapper.utils.heap import MaxHeap
 from tdamapper.utils.metrics import get_metric
-
-
-@njit
-def _swap(arr, i, j):
-    arr[i], arr[j] = arr[j], arr[i]
+from tdamapper.utils.quickselect import quickselect, swap_all
 
 
 def _mid(start, end):
     return (start + end) // 2
 
 
-@njit
-def _partition(distances, indices, is_terminal, start, end, p_ord):
-    higher = start
-    for j in range(start, end):
-        j_ord = distances[j]
-        if j_ord < p_ord:
-            _swap(distances, higher, j)
-            _swap(indices, higher, j)
-            _swap(is_terminal, higher, j)
-            higher += 1
-    return higher
-
-
-def _quickselect(distances, indices, is_terminal, start, end, k):
-    if (k < start) or (k >= end):
-        return
-    start_, end_, higher = start, end, None
-    while higher != k + 1:
-        # TODO: pivot_index = randrange(start_, end_)
-        pivot_index = k
-
-        p = distances[pivot_index]
-
-        _swap(distances, start_, pivot_index)
-        _swap(indices, start_, pivot_index)
-        _swap(is_terminal, start_, pivot_index)
-
-        higher = _partition(distances, indices, is_terminal, start_ + 1, end_, p)
-
-        _swap(distances, start_, higher - 1)
-        _swap(indices, start_, higher - 1)
-        _swap(is_terminal, start_, higher - 1)
-
-        if k <= higher - 1:
-            end_ = higher
-        else:
-            start_ = higher
-
-
 class VPTree:
 
     def __init__(
@@ -71,9 +27,12 @@ def __init__(
         self.__leaf_capacity = leaf_capacity
         self.__leaf_radius = leaf_radius
         self.__pivoting = pivoting
-        self.__dataset, self.__distances, self.__indices, self.__is_terminal = (
-            self._Build(self, X).build()
-        )
+        (
+            self.__dataset,
+            self.__arr_distances,
+            self.__arr_indices,
+            self.__arr_is_terminal,
+        ) = self._Build(self, X).build()
 
     def get_metric(self):
         return self.__metric
@@ -98,23 +57,23 @@ def _get_distance(self):
         return get_metric(self.__metric, **metric_params)
 
     def _get_distances(self):
-        return self.__distances
+        return self.__arr_distances
 
     def _get_indices(self):
-        return self.__indices
+        return self.__arr_indices
 
     def _get_is_terminal(self):
-        return self.__is_terminal
+        return self.__arr_is_terminal
 
     class _Build:
 
         def __init__(self, vpt, X):
             self.__distance = vpt._get_distance()
 
             self.__dataset = [x for x in X]
-            self.__indices = np.array([i for i in range(len(self.__dataset))])
-            self.__distances = np.array([0.0 for _ in X])
-            self.__is_terminal = np.array([False for _ in X])
+            self.__arr_indices = np.array([i for i in range(len(self.__dataset))])
+            self.__arr_distances = np.array([0.0 for _ in X])
+            self.__arr_is_terminal = np.array([False for _ in X])
 
             self.__leaf_capacity = vpt.get_leaf_capacity()
             self.__leaf_radius = vpt.get_leaf_radius()
@@ -133,20 +92,25 @@ def _pivoting_random(self, start, end):
                 return
             pivot = randrange(start, end)
             if pivot > start:
-                _swap(self.__distances, start, pivot)
-                _swap(self.__indices, start, pivot)
-                _swap(self.__is_terminal, start, pivot)
+                swap_all(
+                    self.__arr_distances,
+                    start,
+                    pivot,
+                    self.__arr_indices,
+                    self.__arr_is_terminal,
+                )
+
+        def _get_point(self, i):
+            return self.__dataset[self.__arr_indices[i]]
 
         def _furthest(self, start, end, i):
             furthest_dist = 0.0
             furthest = start
 
-            i_point_index = self.__indices[i]
-            i_point = self.__dataset[i_point_index]
+            i_point = self._get_point(i)
 
             for j in range(start, end):
-                j_point_index = self.__indices[j]
-                j_point = self.__dataset[j_point_index]
+                j_point = self._get_point(j)
 
                 j_dist = self.__distance(i_point, j_point)
                 if j_dist > furthest_dist:
@@ -161,27 +125,36 @@ def _pivoting_furthest(self, start, end):
             furthest_rnd = self._furthest(start, end, rnd)
             furthest = self._furthest(start, end, furthest_rnd)
             if furthest > start:
-                _swap(self.__distances, start, furthest)
-                _swap(self.__indices, start, furthest)
-                _swap(self.__is_terminal, start, furthest)
+                swap_all(
+                    self.__arr_distances,
+                    start,
+                    furthest,
+                    self.__arr_indices,
+                    self.__arr_is_terminal,
+                )
 
         def _update(self, start, end):
             self.__pivoting(start, end)
 
-            v_point_index = self.__indices[start]
+            v_point_index = self.__arr_indices[start]
             v_point = self.__dataset[v_point_index]
-            is_terminal = self.__is_terminal[start]
+            is_terminal = self.__arr_is_terminal[start]
 
             for i in range(start + 1, end):
-                point_index = self.__indices[i]
+                point_index = self.__arr_indices[i]
                 point = self.__dataset[point_index]
 
-                self.__distances[i] = self.__distance(v_point, point)
-                self.__is_terminal[i] = is_terminal
+                self.__arr_distances[i] = self.__distance(v_point, point)
+                self.__arr_is_terminal[i] = is_terminal
 
         def build(self):
             self._build_iter()
-            return self.__dataset, self.__distances, self.__indices, self.__is_terminal
+            return (
+                self.__dataset,
+                self.__arr_distances,
+                self.__arr_indices,
+                self.__arr_is_terminal,
+            )
 
         def _build_iter(self):
             stack = [(0, len(self.__dataset))]
@@ -190,32 +163,32 @@ def _build_iter(self):
                 mid = _mid(start, end)
                 self._update(start, end)
 
-                v_point_index = self.__indices[start]
+                # v_point_index = self.__indices[start]
 
-                _quickselect(
-                    self.__distances,
-                    self.__indices,
-                    self.__is_terminal,
+                quickselect(
+                    self.__arr_distances,
                     start + 1,
                     end,
                     mid,
+                    self.__arr_indices,
+                    self.__arr_is_terminal,
                 )
 
-                v_radius = self.__distances[mid]
+                v_radius = self.__arr_distances[mid]
 
                 if (end - start > 2 * self.__leaf_capacity) and (
                     v_radius > self.__leaf_radius
                 ):
-                    self.__distances[start] = v_radius
-                    self.__indices[start] = v_point_index
-                    self.__is_terminal[start] = False
+                    self.__arr_distances[start] = v_radius
+                    # self.__indices[start] = v_point_index
+                    self.__arr_is_terminal[start] = False
 
                     stack.append((mid, end))
                     stack.append((start + 1, mid))
                 else:
-                    self.__distances[start] = v_radius
-                    self.__indices[start] = v_point_index
-                    self.__is_terminal[start] = True
+                    self.__arr_distances[start] = v_radius
+                    # self.__indices[start] = v_point_index
+                    self.__arr_is_terminal[start] = True
 
     def ball_search(self, point, eps, inclusive=True):
         return self._BallSearch(self, point, eps, inclusive).search()