Improved types

Author: lucasimi

src/tdamapper/utils/vptree_flat/ball_search.py

@@ -1,24 +1,37 @@
-from tdamapper.utils.vptree_flat.common import _mid
+from typing import Generic, Iterable, TypeVar
 
+from tdamapper.utils.metrics import Metric
+from tdamapper.utils.vptree_flat.common import VPArray, VPTreeType, _mid
 
-class BallSearch:
+T = TypeVar("T")
 
-    def __init__(self, vpt, point, eps, inclusive=True):
+
+class BallSearch(Generic[T]):
+
+    _arr: VPArray[T]
+    _distance: Metric[T]
+    _point: T
+    _eps: float
+    _inclusive: bool
+
+    def __init__(
+        self, vpt: VPTreeType[T], point: T, eps: float, inclusive: bool = True
+    ):
         self._arr = vpt._get_arr()
         self._distance = vpt._get_distance()
         self._point = point
         self._eps = eps
         self._inclusive = inclusive
 
-    def search(self):
+    def search(self) -> Iterable[T]:
         return self._search_iter()
 
-    def _inside(self, dist):
+    def _inside(self, dist: float) -> bool:
         if self._inclusive:
             return dist <= self._eps
         return dist < self._eps
 
-    def _search_iter(self):
+    def _search_iter(self) -> Iterable[T]:
         stack = [(0, self._arr.size())]
         result = []
         while stack:

src/tdamapper/utils/vptree_flat/builder.py

@@ -1,17 +1,20 @@
 from random import randrange
+from typing import Generic, Iterable, TypeVar
 
 import numpy as np
 
-from tdamapper.utils.vptree_flat.common import VPArray
+from tdamapper.utils.vptree_flat.common import VPArray, VPTreeType
+
+T = TypeVar("T")
 
 
 def _mid(start, end):
     return (start + end) // 2
 
 
-class Builder:
+class Builder(Generic[T]):
 
-    def __init__(self, vpt, X):
+    def __init__(self, vpt: VPTreeType[T], X: Iterable[T]) -> None:
         self._distance = vpt._get_distance()
 
         dataset = [x for x in X]
@@ -29,17 +32,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)
@@ -51,7 +54,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)
@@ -60,7 +63,7 @@ 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)
         is_terminal = self._arr.is_terminal(start)
@@ -69,11 +72,11 @@ def _update(self, start, end):
             self._arr.set_distance(i, self._distance(v_point, point))
             self._arr.set_terminal(i, is_terminal)
 
-    def build(self):
+    def build(self) -> VPArray[T]:
         self._build_iter()
         return self._arr
 
-    def _build_iter(self):
+    def _build_iter(self) -> None:
         stack = [(0, self._arr.size())]
         while stack:
             start, end = stack.pop()

src/tdamapper/utils/vptree_flat/common.py

@@ -1,42 +1,72 @@
+from __future__ import annotations
+
+from typing import Generic, Iterable, Optional, Protocol, TypeVar
+
+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):
     return (start + end) // 2
 
 
-class VPArray:
+class VPTreeType(Protocol[T]):
+
+    def _get_arr(self) -> VPArray[T]: ...
+
+    def _get_distance(self) -> Metric[T]: ...
+
+    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, distances, indices, is_terminal):
+    def __init__(
+        self,
+        dataset: ArrayLike[T],
+        distances: NDArray[np.float64],
+        indices: NDArray[np.int64],
+        is_terminal: NDArray[np.bool_],
+    ):
         self._dataset = dataset
         self._distances = distances
         self._indices = indices
         self._is_terminal = is_terminal
 
-    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 set_terminal(self, i, terminal):
+    def set_terminal(self, i: int, terminal: bool) -> None:
         self._is_terminal[i] = terminal
 
-    def is_terminal(self, i):
+    def is_terminal(self, i: int) -> bool:
         return self._is_terminal[i]
 
-    def swap(self, i, j):
+    def swap(self, i: int, j: int) -> None:
         swap_all(self._distances, i, j, self._indices, self._is_terminal)
 
-    def partition(self, s, e, k):
+    def partition(self, s: int, e: int, k: int) -> None:
         quickselect(self._distances, s, e, k, self._indices, self._is_terminal)

src/tdamapper/utils/vptree_flat/knn_search.py

@@ -1,30 +1,41 @@
+from typing import Generic, TypeVar
+
 from tdamapper.utils.heap import MaxHeap
-from tdamapper.utils.vptree_flat.common import _mid
+from tdamapper.utils.metrics import Metric
+from tdamapper.utils.vptree_flat.common import VPArray, VPTreeType, _mid
 
 _PRE = 0
 _POST = 1
 
+T = TypeVar("T")
+
+
+class KnnSearch(Generic[T]):
 
-class KnnSearch:
+    _arr: VPArray[T]
+    _distance: Metric[T]
+    _point: T
+    _neighbors: int
+    _result: MaxHeap[float, T]
 
-    def __init__(self, vpt, point, neighbors):
+    def __init__(self, vpt: VPTreeType[T], point: T, neighbors: int) -> None:
         self._arr = vpt._get_arr()
         self._distance = vpt._get_distance()
         self._point = point
         self._neighbors = neighbors
         self._radius = float("inf")
         self._result = MaxHeap()
 
-    def _get_items(self):
+    def _get_items(self) -> list[T]:
         while len(self._result) > self._neighbors:
             self._result.pop()
         return [x for (_, x) in self._result]
 
-    def search(self):
+    def search(self) -> list[T]:
         self._search_iter()
         return self._get_items()
 
-    def _process(self, x):
+    def _process(self, x: T) -> float:
         dist = self._distance(self._point, x)
         if dist >= self._radius:
             return dist
@@ -35,7 +46,7 @@ def _process(self, x):
             self._radius, _ = self._result.top()
         return dist
 
-    def _search_iter(self):
+    def _search_iter(self) -> list[T]:
         self._result = MaxHeap()
         stack = [(0, self._arr.size(), 0.0, _PRE)]
         while stack:

src/tdamapper/utils/vptree_flat/vptree.py

@@ -1,51 +1,58 @@
-from tdamapper.utils.metrics import get_metric
+from typing import Any, Generic, Iterable, Optional, TypeVar, Union
+
+import numpy as np
+from numpy.typing import NDArray
+
+from tdamapper.utils.metrics import Metric, get_metric
 from tdamapper.utils.vptree_flat.ball_search import BallSearch
 from tdamapper.utils.vptree_flat.builder import Builder
 from tdamapper.utils.vptree_flat.knn_search import KnnSearch
 
+T = TypeVar("T")
+
 
-class VPTree:
+class VPTree(Generic[T]):
 
     def __init__(
         self,
-        X,
-        metric="euclidean",
-        metric_params=None,
-        leaf_capacity=1,
-        leaf_radius=0.0,
-        pivoting=None,
-    ):
+        X: Iterable[T],
+        metric: Union[str, Metric[T]] = "euclidean",
+        metric_params: Optional[dict[str, Any]] = None,
+        leaf_capacity: int = 1,
+        leaf_radius: float = 0.0,
+        pivoting: Optional[str] = None,
+    ) -> None:
         self._metric = metric
         self._metric_params = metric_params
         self._leaf_capacity = leaf_capacity
         self._leaf_radius = leaf_radius
         self._pivoting = pivoting
         self._arr = Builder(self, X).build()
 
-    def get_metric(self):
+    def get_metric(self) -> Union[str, Metric[T]]:
         return self._metric
 
-    def get_metric_params(self):
+    def get_metric_params(self) -> Optional[dict[str, Any]]:
         return self._metric_params
 
-    def get_leaf_capacity(self):
+    def get_leaf_capacity(self) -> int:
         return self._leaf_capacity
 
-    def get_leaf_radius(self):
+    def get_leaf_radius(self) -> float:
         return self._leaf_radius
 
-    def get_pivoting(self):
+    def get_pivoting(self) -> Optional[str]:
         return self._pivoting
 
     def _get_arr(self):
         return self._arr
 
-    def _get_distance(self):
+    def _get_distance(self) -> Union[Metric[NDArray[np.float64]], Metric[T]]:
         metric_params = self._metric_params or {}
         return get_metric(self._metric, **metric_params)
 
-    def ball_search(self, point, eps, inclusive=True):
+    def ball_search(self, point: T, eps: float, inclusive: bool = True) -> Iterable[T]:
         return BallSearch(self, point, eps, inclusive).search()
 
-    def knn_search(self, point, k):
+    def knn_search(self, point: T, k: int) -> Iterable[T]:
         return KnnSearch(self, point, k).search()