Added types

Author: lucasimi

src/tdamapper/utils/vptree_flat/ball_search.py

@@ -1,4 +1,4 @@
-from typing import Generic, Iterable, TypeVar
+from typing import Generic, TypeVar
 
 from tdamapper.utils.metrics import Metric
 from tdamapper.utils.vptree_flat.common import VPArray, VPTreeType, _mid
@@ -23,15 +23,15 @@ def __init__(
         self._eps = eps
         self._inclusive = inclusive
 
-    def search(self) -> Iterable[T]:
+    def search(self) -> list[T]:
         return self._search_iter()
 
     def _inside(self, dist: float) -> bool:
         if self._inclusive:
             return dist <= self._eps
         return dist < self._eps
 
-    def _search_iter(self) -> Iterable[T]:
+    def _search_iter(self) -> list[T]:
         stack = [(0, self._arr.size())]
         result = []
         while stack:

src/tdamapper/utils/vptree_flat/builder.py

@@ -1,8 +1,9 @@
 from random import randrange
-from typing import Generic, Iterable, TypeVar
+from typing import Callable, Generic, Iterable, TypeVar
 
 import numpy as np
 
+from tdamapper.utils.metrics import Metric
 from tdamapper.utils.vptree_flat.common import VPArray, VPTreeType
 
 T = TypeVar("T")
@@ -14,13 +15,19 @@ def _mid(start, end):
 
 class Builder(Generic[T]):
 
-    def __init__(self, vpt: VPTreeType[T], X: Iterable[T]) -> None:
+    _arr: VPArray[T]
+    _leaf_capacity: int
+    _leaf_radius: float
+    _distance: Metric[T]
+    _pivoting: Callable[[int, int], None]
+
+    def __init__(self, vpt: VPTreeType[T], items: Iterable[T]) -> None:
         self._distance = vpt._get_distance()
 
-        dataset = [x for x in X]
+        dataset = [x for x in items]
         indices = np.array([i for i in range(len(dataset))])
-        distances = np.array([0.0 for _ in X])
-        is_terminal = np.array([False for _ in X])
+        distances = np.array([0.0 for _ in items])
+        is_terminal = np.array([False for _ in items])
         self._arr = VPArray(dataset, distances, indices, is_terminal)
 
         self._leaf_capacity = vpt.get_leaf_capacity()

src/tdamapper/utils/vptree_flat/common.py

@@ -12,7 +12,7 @@
 T = TypeVar("T")
 
 
-def _mid(start, end):
+def _mid(start: int, end: int) -> int:
     return (start + end) // 2
 
 

src/tdamapper/utils/vptree_flat/knn_search.py

@@ -43,7 +43,9 @@ def _process(self, x: T) -> float:
         while len(self._result) > self._neighbors:
             self._result.pop()
         if len(self._result) == self._neighbors:
-            self._radius, _ = self._result.top()
+            radius, _ = self._result.top()
+            if radius is not None:
+                self._radius = radius
         return dist
 
     def _search_iter(self) -> list[T]:

src/tdamapper/utils/vptree_flat/vptree.py

@@ -13,9 +13,15 @@
 
 class VPTree(Generic[T]):
 
+    _metric: Union[str, Metric[T]]
+    _metric_params: Optional[dict[str, Any]]
+    _leaf_capacity: int
+    _leaf_radius: float
+    _pivoting: Optional[str]
+
     def __init__(
         self,
-        X: Iterable[T],
+        items: Iterable[T],
         metric: Union[str, Metric[T]] = "euclidean",
         metric_params: Optional[dict[str, Any]] = None,
         leaf_capacity: int = 1,
@@ -27,7 +33,7 @@ def __init__(
         self._leaf_capacity = leaf_capacity
         self._leaf_radius = leaf_radius
         self._pivoting = pivoting
-        self._arr = Builder(self, X).build()
+        self._arr = Builder(self, items).build()
 
     def get_metric(self) -> Union[str, Metric[T]]:
         return self._metric

src/tdamapper/utils/vptree_hier/ball_search.py

@@ -1,6 +1,24 @@
-class BallSearch:
+from typing import Callable, Generic, Iterable, Optional, TypeVar, Union
 
-    def __init__(self, vpt, point, eps, inclusive=True):
+from tdamapper.utils.metrics import Metric
+from tdamapper.utils.vptree_hier.common import Tree, VPArray, VPTreeType
+
+T = TypeVar("T")
+
+
+class BallSearch(Generic[T]):
+
+    _tree: Tree[T]
+    _arr: VPArray[T]
+    _distance: Metric[T]
+    _point: T
+    _eps: float
+    _inclusive: bool
+    _result: list[T]
+
+    def __init__(
+        self, vpt: VPTreeType[T], point: T, eps: float, inclusive: bool = True
+    ) -> None:
         self._tree = vpt._get_tree()
         self._arr = vpt._get_arr()
         self._distance = vpt._get_distance()
@@ -9,17 +27,17 @@ def __init__(self, vpt, point, eps, inclusive=True):
         self._inclusive = inclusive
         self._result = []
 
-    def search(self):
+    def search(self) -> list[T]:
         self._result.clear()
         self._search_rec(self._tree)
         return self._result
 
-    def _inside(self, dist):
+    def _inside(self, dist: float) -> bool:
         if self._inclusive:
             return dist <= self._eps
         return dist < self._eps
 
-    def _search_rec(self, tree):
+    def _search_rec(self, tree: Tree[T]) -> None:
         if tree.is_terminal():
             start, end = tree.get_bounds()
             for x in self._arr.get_points(start, end):

src/tdamapper/utils/vptree_hier/builder.py

@@ -1,18 +1,35 @@
 from random import randrange
+from typing import Callable, Generic, Iterable, TypeVar
 
 import numpy as np
 
-from tdamapper.utils.vptree_hier.common import Leaf, Node, VPArray, _mid
+from tdamapper.utils.metrics import Metric
+from tdamapper.utils.vptree_hier.common import (
+    Leaf,
+    Node,
+    Tree,
+    VPArray,
+    VPTreeType,
+    _mid,
+)
 
+T = TypeVar("T")
 
-class Builder:
 
-    def __init__(self, vpt, X):
+class Builder(Generic[T]):
+
+    _arr: VPArray[T]
+    _leaf_capacity: int
+    _leaf_radius: float
+    _distance: Metric[T]
+    _pivoting: Callable[[int, int], None]
+
+    def __init__(self, vpt: VPTreeType[T], items: Iterable[T]):
         self._distance = vpt._get_distance()
 
-        dataset = [x for x in X]
+        dataset = [x for x in items]
         indices = np.array([i for i in range(len(dataset))])
-        distances = np.array([0.0 for _ in X])
+        distances = np.array([0.0 for _ in items])
         self._arr = VPArray(dataset, distances, indices)
 
         self._leaf_capacity = vpt.get_leaf_capacity()
@@ -24,17 +41,17 @@ def __init__(self, vpt, X):
         elif pivoting == "furthest":
             self._pivoting = self._pivoting_furthest
 
-    def _pivoting_disabled(self, start, end):
+    def _pivoting_disabled(self, start: int, end: int) -> None:
         pass
 
-    def _pivoting_random(self, start, end):
+    def _pivoting_random(self, start: int, end: int) -> None:
         if end <= start:
             return
         pivot = randrange(start, end)
         if pivot > start:
             self._arr.swap(start, pivot)
 
-    def _furthest(self, start, end, i):
+    def _furthest(self, start: int, end: int, i: int) -> int:
         furthest_dist = 0.0
         furthest = start
         i_point = self._arr.get_point(i)
@@ -46,7 +63,7 @@ def _furthest(self, start, end, i):
                 furthest_dist = j_dist
         return furthest
 
-    def _pivoting_furthest(self, start, end):
+    def _pivoting_furthest(self, start: int, end: int) -> None:
         if end <= start:
             return
         rnd = randrange(start, end)
@@ -55,24 +72,26 @@ def _pivoting_furthest(self, start, end):
         if furthest > start:
             self._arr.swap(start, furthest)
 
-    def _update(self, start, end):
+    def _update(self, start: int, end: int) -> None:
         self._pivoting(start, end)
         v_point = self._arr.get_point(start)
         for i in range(start + 1, end):
             point = self._arr.get_point(i)
             self._arr.set_distance(i, self._distance(v_point, point))
 
-    def build(self):
+    def build(self) -> tuple[Tree[T], VPArray[T]]:
         tree = self._build_rec(0, self._arr.size())
         return tree, self._arr
 
-    def _build_rec(self, start, end):
+    def _build_rec(self, start: int, end: int) -> Tree[T]:
         mid = _mid(start, end)
         self._update(start, end)
         v_point = self._arr.get_point(start)
         self._arr.partition(start + 1, end, mid)
         v_radius = self._arr.get_distance(mid)
         self._arr.set_distance(start, v_radius)
+        left: Tree[T]
+        right: Tree[T]
         if (end - start <= 2 * self._leaf_capacity) or (v_radius <= self._leaf_radius):
             left = Leaf(start + 1, mid)
             right = Leaf(mid, end)

src/tdamapper/utils/vptree_hier/common.py

@@ -1,69 +1,103 @@
+from __future__ import annotations
+
+from typing import Generic, Iterable, Optional, Protocol, TypeVar, Union
+
+import numpy as np
+from numpy.typing import NDArray
+
+from tdamapper._common import ArrayLike
+from tdamapper.utils.metrics import Metric
 from tdamapper.utils.quickselect import quickselect, swap_all
 
+T = TypeVar("T")
+
 
-def _mid(start, end):
+def _mid(start: int, end: int) -> int:
     return (start + end) // 2
 
 
-class VPArray:
+class VPTreeType(Protocol[T]):
+
+    def _get_arr(self) -> VPArray[T]: ...
+
+    def _get_distance(self) -> Metric[T]: ...
+
+    def _get_tree(self) -> Tree[T]: ...
 
-    def __init__(self, dataset, distances, indices):
+    def get_leaf_capacity(self) -> int: ...
+
+    def get_leaf_radius(self) -> float: ...
+
+    def get_pivoting(self) -> Optional[str]: ...
+
+
+class VPArray(Generic[T]):
+
+    def __init__(
+        self,
+        dataset: ArrayLike[T],
+        distances: NDArray[np.float64],
+        indices: NDArray[np.bool_],
+    ):
         self._dataset = dataset
         self._distances = distances
         self._indices = indices
 
-    def size(self):
+    def size(self) -> int:
         return len(self._dataset)
 
-    def get_point(self, i):
+    def get_point(self, i: int) -> T:
         return self._dataset[self._indices[i]]
 
-    def get_points(self, s, e):
+    def get_points(self, s: int, e: int) -> Iterable[T]:
         for x_index in self._indices[s:e]:
             yield self._dataset[x_index]
 
-    def get_distance(self, i):
+    def get_distance(self, i: int) -> float:
         return self._distances[i]
 
-    def set_distance(self, i, dist):
+    def set_distance(self, i: int, dist: float) -> None:
         self._distances[i] = dist
 
-    def swap(self, i, j):
+    def swap(self, i: int, j: int) -> None:
         swap_all(self._distances, i, j, self._indices)
 
-    def partition(self, s, e, k):
+    def partition(self, s: int, e: int, k: int) -> None:
         quickselect(self._distances, s, e, k, self._indices)
 
 
-class Node:
+class Node(Generic[T]):
 
-    def __init__(self, radius, center, left, right):
+    def __init__(self, radius: float, center: T, left: Tree[T], right: Tree[T]):
         self._radius = radius
         self._center = center
         self._left = left
         self._right = right
 
-    def get_ball(self):
+    def get_ball(self) -> tuple[float, T]:
         return self._radius, self._center
 
-    def is_terminal(self):
+    def is_terminal(self) -> bool:
         return False
 
-    def get_left(self):
+    def get_left(self) -> Tree[T]:
         return self._left
 
-    def get_right(self):
+    def get_right(self) -> Tree[T]:
         return self._right
 
 
 class Leaf:
 
-    def __init__(self, start, end):
+    def __init__(self, start: int, end: int) -> None:
         self._start = start
         self._end = end
 
-    def get_bounds(self):
+    def get_bounds(self) -> tuple[int, int]:
         return self._start, self._end
 
-    def is_terminal(self):
+    def is_terminal(self) -> bool:
         return True
+
+
+Tree = Union[Node[T], Leaf]