Improved types and socd

Author: lucasimi

src/tdamapper/_common.py

@@ -98,7 +98,6 @@ def get_params(self, deep: bool = True) -> Dict[str, Any]:
         Get all public parameters of the object as a dictionary.
 
         :param deep: A flag for returning also nested parameters.
-        :type deep: bool, optional.
         """
         params = {}
         for k, v in self.__dict__.items():

src/tdamapper/core.py

@@ -85,20 +85,13 @@ def mapper_labels(
     than those at position j.
 
     :param X: A dataset of n points.
-    :type X: array-like of shape (n, m) or list-like of length n
     :param y: Lens values for the n points of the dataset.
-    :type y: array-like of shape (n, k) or list-like of length n
     :param cover: A cover algorithm.
-    :type cover: A class compatible with :class:`tdamapper.core.Cover`
     :param clustering: A clustering algorithm.
-    :type clustering: An estimator compatible with scikit-learn's clustering
-        interface, typically from :mod:`sklearn.cluster`.
     :param n_jobs: The maximum number of parallel clustering jobs. This
         parameter is passed to the constructor of :class:`joblib.Parallel`.
         Defaults to 1.
-    :type n_jobs: int
     :return: A list of node labels for each point in the dataset.
-    :rtype: list[list[int]]
     """
 
     def _run_clustering(local_ids, X_local, clust):
@@ -146,22 +139,15 @@ def mapper_connected_components(
     :func:`networkx.connected_components` on it.
 
     :param X: A dataset of n points.
-    :type X: array-like of shape (n, m) or list-like of length n
     :param y: Lens values for the n points of the dataset.
-    :type y: array-like of shape (n, k) or list-like of length n
     :param cover: A cover algorithm.
-    :type cover: A class compatible with :class:`tdamapper.core.Cover`
     :param clustering: The clustering algorithm to apply to each subset of the
         dataset.
-    :type clustering: An estimator compatible with scikit-learn's clustering
-        interface, typically from :mod:`sklearn.cluster`.
     :param n_jobs: The maximum number of parallel clustering jobs. This
         parameter is passed to the constructor of :class:`joblib.Parallel`.
         Defaults to 1.
-    :type n_jobs: int
     :return: A list of labels. The label at position i identifies the connected
         component of the point at position i in the dataset.
-    :rtype: list[int]
     """
     itm_lbls = mapper_labels(X, y, cover, clustering, n_jobs=n_jobs)
     label_values = set()
@@ -202,21 +188,14 @@ def mapper_graph(
     contained in the cluster.
 
     :param X: A dataset of n points.
-    :type X: array-like of shape (n, m) or list-like of length n
     :param y: Lens values for the n points of the dataset.
-    :type y: array-like of shape (n, k) or list-like of length n
     :param cover: A cover algorithm.
-    :type cover: A class compatible with :class:`tdamapper.core.Cover`
     :param clustering: The clustering algorithm to apply to each subset of the
         dataset.
-    :type clustering: An estimator compatible with scikit-learn's clustering
-        interface, typically from :mod:`sklearn.cluster`.
     :param n_jobs: The maximum number of parallel clustering jobs. This
         parameter is passed to the constructor of :class:`joblib.Parallel`.
         Defaults to 1.
-    :type n_jobs: int
     :return: The Mapper graph.
-    :rtype: :class:`networkx.Graph`
     """
     itm_lbls = mapper_labels(X, y, cover, clustering, n_jobs=n_jobs)
     graph = nx.Graph()
@@ -253,13 +232,9 @@ def aggregate_graph(X: ArrayLike, graph: nx.Graph, agg: Callable) -> Dict:
     and the values are the aggregation values.
 
     :param X: A dataset of n points.
-    :type X: array-like of shape (n, m) or list-like of length n
     :param graph: The graph to apply the aggregation function to.
-    :type graph: :class:`networkx.Graph`.
     :param agg: The aggregation function to use.
-    :type agg: Callable.
     :return: A dictionary of node-aggregation pairs.
-    :rtype: dict
     """
     agg_values = {}
     nodes = graph.nodes()
@@ -286,7 +261,6 @@ def fit(self, X: ArrayLike) -> SpatialSearch:
         Fit the spatial search algorithm to the data.
 
         :param X: A dataset of n points.
-        :type X: array-like of shape (n, m) or list-like of length n
         :return: self
         """
 
@@ -295,9 +269,7 @@ def search(self, x: PointLike) -> List[int]:
         Search for the nearest neighbors of a point.
 
         :param x: A point to search for.
-        :type x: A point-like object, such as a list or a numpy array.
         :return: A list of indices of the nearest neighbors of the point.
-        :rtype: list[int]
         """
 
 
@@ -317,7 +289,6 @@ def fit(self, X: ArrayLike) -> Cover:
         Fit the cover algorithm to the data.
 
         :param X: A dataset of n points.
-        :type X: array-like of shape (n, m) or list-like of length n
         :return: self
         """
 
@@ -329,7 +300,6 @@ def fit_transform(self, X: ArrayLike) -> Generator[List[int], None, None]:
         the indices of the points in the dataset that belong to the open set.
 
         :param X: A dataset of n points.
-        :type X: array-like of shape (n, m) or list-like of length n
         :yield: A generator of lists of indices.
         """
 
@@ -341,7 +311,6 @@ def transform(self, X: ArrayLike) -> Generator[List[int], None, None]:
         the indices of the points in the dataset that belong to the open set.
 
         :param X: A dataset of n points.
-        :type X: array-like of shape (n, m) or list-like of length n
         :yield: A generator of lists of indices.
         """
 
@@ -374,7 +343,6 @@ def fit(self, X: ArrayLike, y: Any = None) -> Clustering:
         Fit the clustering algorithm to the data.
 
         :param X: A dataset of n points.
-        :type X: array-like of shape (n, m) or list-like of length n
         :param y: Ignored.
         :return: self
         """
@@ -402,6 +370,13 @@ def fit_transform(self, X: ArrayLike) -> Generator[List[int], None, None]:
 
 class _MapperAlgorithm(EstimatorMixin, ParamsMixin):
 
+    _cover: Cover
+    _clustering: Clustering
+    _verbose: bool
+    _failsafe: bool
+    _n_jobs: int
+    graph_: nx.Graph
+
     def __init__(
         self,
         cover: Optional[Cover] = None,
@@ -472,13 +447,14 @@ class FailSafeClustering(ParamsMixin):
     returned. This can be useful for robustness and debugging purposes.
 
     :param clustering: A clustering algorithm to delegate to.
-    :type clustering: An estimator compatible with scikit-learn's clustering
-        interface, typically from :mod:`sklearn.cluster`.
     :param verbose: A flag to log clustering exceptions. Set to True to
         enable logging, or False to suppress it. Defaults to True.
-    :type verbose: bool, optional.
     """
 
+    _clustering: Clustering
+    _verbose: bool
+    labels_: List[int]
+
     def __init__(self, clustering: Optional[Clustering] = None, verbose: bool = True):
         self.clustering = clustering
         self.verbose = verbose
@@ -488,7 +464,7 @@ def fit(self, X: ArrayLike, y: Optional[ArrayLike] = None):
             TrivialClustering() if self.clustering is None else self.clustering
         )
         self._verbose = self.verbose
-        self.labels_ = None
+        self.labels_ = []
         try:
             self._clustering.fit(X, y)
             self.labels_ = self._clustering.labels_
@@ -509,15 +485,16 @@ class TrivialClustering(ParamsMixin):
     construction of the Mapper graph.
     """
 
+    labels_: List[int]
+
     def __init__(self):
         pass
 
-    def fit(self, X: ArrayLike, y: Optional[ArrayLike] = None) -> TrivialClustering:
+    def fit(self, X: ArrayLike, _: Optional[ArrayLike] = None) -> TrivialClustering:
         """
         Fit the clustering algorithm to the data.
 
         :param X: A dataset of n points.
-        :type X: array-like of shape (n, m) or list-like of length n
         :param y: Ignored.
         :return: self
         """

src/tdamapper/cover.py

@@ -58,9 +58,7 @@ def _cover(self, X: ArrayLike) -> Generator[Tuple[int, List[int]], None, None]:
         landmark, and the list of ids contains the indices of its neighbors.
 
         :param X: A dataset of n points.
-        :type X: array-like of shape (n, m) or list-like of length n
         :return: A generator of pairs of indices and lists of ids.
-        :rtype: generator of tuples (int, List[int])
         :yield: A tuple containing the index of the point and a list of ids of
                 its covered neighbors. If the index is -1, it indicates a
                 landmark, and the list contains the indices of its neighbors.

src/tdamapper/heap.py

@@ -4,6 +4,10 @@
 and iterate over the elements in the heap.
 """
 
+from __future__ import annotations
+
+from typing import Generic, Optional, Tuple, TypeVar
+
 
 def _left(i: int) -> int:
     return 2 * i + 1
@@ -17,7 +21,11 @@ def _parent(i: int) -> int:
     return max(0, (i - 1) // 2)
 
 
-class _HeapNode:
+K = TypeVar("K")
+T = TypeVar("T")
+
+
+class _HeapNode(Generic[K, T]):
     """
     A node in the max-heap, storing a key and a value.
 
@@ -28,36 +36,34 @@ class _HeapNode:
 
     :param key: The key used for ordering in the heap.
     :param value: The value associated with the key.
-    :type key: Any
-    :type value: Any
     """
 
-    def __init__(self, key, value):
+    def __init__(self, key: K, value: T):
         self._key = key
         self._value = value
 
-    def get(self):
+    def get(self) -> Tuple[K, T]:
         """
         Returns the key and value of the node.
         :return: A tuple containing the key and value.
         :rtype: tuple
         """
         return self._key, self._value
 
-    def __lt__(self, other):
+    def __lt__(self, other: _HeapNode[K, T]) -> bool:
         return self._key < other._key
 
-    def __le__(self, other):
+    def __le__(self, other: _HeapNode[K, T]) -> bool:
         return self._key <= other._key
 
-    def __gt__(self, other):
+    def __gt__(self, other: _HeapNode[K, T]) -> bool:
         return self._key > other._key
 
-    def __ge__(self, other):
+    def __ge__(self, other: _HeapNode[K, T]) -> bool:
         return self._key >= other._key
 
 
-class MaxHeap:
+class MaxHeap(Generic[K, T]):
     """
     A max-heap data structure that allows for efficient retrieval of the maximum element.
 
@@ -69,58 +75,54 @@ def __init__(self):
         self._heap = []
         self._iter = None
 
-    def __iter__(self):
+    def __iter__(self) -> MaxHeap[K, T]:
         self._iter = iter(self._heap)
         return self
 
-    def __next__(self):
+    def __next__(self) -> Tuple[K, T]:
         node = next(self._iter)
         return node.get()
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self._heap)
 
-    def top(self):
+    def top(self) -> Tuple[Optional[K], Optional[T]]:
         """
         Returns the maximum element in the heap without removing it.
 
         :return: A tuple containing the key and value of the maximum element, or (None, None) if the heap is empty.
-        :rtype: tuple
         """
         if not self._heap:
             return (None, None)
         return self._heap[0].get()
 
-    def pop(self):
+    def pop(self) -> Optional[Tuple[K, T]]:
         """
         Removes and returns the maximum element from the heap.
 
         :return: A tuple containing the key and value of the maximum element, or None if the heap is empty.
         :rtype: tuple
         """
         if not self._heap:
-            return
+            return None
         max_val = self._heap[0]
         self._heap[0] = self._heap[-1]
         self._heap.pop()
         self._bubble_down()
         return max_val.get()
 
-    def add(self, key, val):
+    def add(self, key: K, val: T) -> None:
         """
         Adds a new element to the heap with the specified key and value.
 
         :param key: The key used for ordering in the heap.
         :param val: The value associated with the key.
-        :type key: Any
-        :type val: Any
         :return: None
-        :rtype: None
         """
         self._heap.append(_HeapNode(key, val))
         self._bubble_up()
 
-    def _get_local_max(self, i):
+    def _get_local_max(self, i: int) -> int:
         heap_len = len(self._heap)
         left = _left(i)
         right = _right(i)
@@ -137,7 +139,7 @@ def _get_local_max(self, i):
             return max_child
         return i
 
-    def _fix_down(self, i):
+    def _fix_down(self, i: int) -> int:
         local_max = self._get_local_max(i)
         if i < local_max:
             self._heap[i], self._heap[local_max] = (
@@ -147,22 +149,22 @@ def _fix_down(self, i):
             return local_max
         return i
 
-    def _fix_up(self, i):
+    def _fix_up(self, i: int) -> int:
         parent = _parent(i)
         if self._heap[parent] < self._heap[i]:
             self._heap[i], self._heap[parent] = self._heap[parent], self._heap[i]
             return parent
         return i
 
-    def _bubble_down(self):
+    def _bubble_down(self) -> None:
         current = 0
         done = False
         while not done:
             local_max = self._fix_down(current)
             done = current == local_max
             current = local_max
 
-    def _bubble_up(self):
+    def _bubble_up(self) -> None:
         current = len(self._heap) - 1
         done = False
         while not done: