Improved variance of types

Author: lucasimi

src/tdamapper/_common.py

@@ -13,7 +13,7 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from tdamapper.protocols import Array
+from tdamapper.protocols import Array, ArrayRead
 
 warnings.filterwarnings("default", category=DeprecationWarning, module=r"^tdamapper\.")
 
@@ -35,48 +35,52 @@ def warn_user(msg: str) -> None:
 
 class EstimatorMixin:
 
-    def _is_sparse(self, X: Array[Any]) -> bool:
+    def _is_sparse(self, X: ArrayRead[Any]) -> bool:
         # simple alternative use scipy.sparse.issparse
         return hasattr(X, "toarray")
 
     def _validate_X_y(
-        self, X: Array[Any], y: Array[Any]
+        self, X: ArrayRead[Any], y: ArrayRead[Any]
     ) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
         if self._is_sparse(X):
             raise ValueError("Sparse data not supported.")
 
-        X = np.asarray(X)
-        y = np.asarray(y)
+        X_ = np.asarray(X)
+        y_ = np.asarray(y)
 
-        if X.size == 0:
-            msg = f"0 feature(s) (shape={X.shape}) while a minimum of 1 is " "required."
+        if X_.size == 0:
+            msg = (
+                f"0 feature(s) (shape={X_.shape}) while a minimum of 1 is " "required."
+            )
             raise ValueError(msg)
 
-        if y.size == 0:
-            msg = f"0 feature(s) (shape={y.shape}) while a minimum of 1 is " "required."
+        if y_.size == 0:
+            msg = (
+                f"0 feature(s) (shape={y_.shape}) while a minimum of 1 is " "required."
+            )
             raise ValueError(msg)
 
-        if X.ndim == 1:
+        if X_.ndim == 1:
             raise ValueError("1d-arrays not supported.")
 
-        if np.iscomplexobj(X) or np.iscomplexobj(y):
+        if np.iscomplexobj(X_) or np.iscomplexobj(y_):
             raise ValueError("Complex data not supported.")
 
-        if X.dtype == np.object_:
-            X = np.array(X, dtype=float)
+        if X_.dtype == np.object_:
+            X_ = np.array(X_, dtype=float)
 
-        if y.dtype == np.object_:
-            y = np.array(y, dtype=float)
+        if y_.dtype == np.object_:
+            y_ = np.array(y_, dtype=float)
 
         if (
-            np.isnan(X).any()
-            or np.isinf(X).any()
-            or np.isnan(y).any()
-            or np.isinf(y).any()
+            np.isnan(X_).any()
+            or np.isinf(X_).any()
+            or np.isnan(y_).any()
+            or np.isinf(y_).any()
         ):
             raise ValueError("NaNs or infinite values not supported.")
 
-        return X, y
+        return X_, y_
 
     def _set_n_features_in(self, X: Array[Any]) -> None:
         if hasattr(X, "shape"):

src/tdamapper/core.py

@@ -37,7 +37,7 @@
 from joblib import Parallel, delayed
 
 from tdamapper._common import ParamsMixin, clone
-from tdamapper.protocols import Array, Clustering, Cover, SpatialSearch
+from tdamapper.protocols import ArrayRead, Clustering, Cover, SpatialSearch
 from tdamapper.utils.unionfind import UnionFind
 
 ATTR_IDS = "ids"
@@ -60,8 +60,8 @@
 
 
 def mapper_labels(
-    X: Array[S],
-    y: Array[T],
+    X: ArrayRead[S],
+    y: ArrayRead[T],
     cover: Cover[T],
     clustering: Clustering[S],
     n_jobs: int = 1,
@@ -94,7 +94,7 @@ def mapper_labels(
     """
 
     def _run_clustering(
-        local_ids: list[int], X_local: Array[S], clust: Clustering[S]
+        local_ids: list[int], X_local: ArrayRead[S], clust: Clustering[S]
     ) -> tuple[list[int], list[int]]:
         local_lbls = clust.fit(X_local).labels_
         return local_ids, local_lbls
@@ -119,8 +119,8 @@ def _run_clustering(
 
 
 def mapper_connected_components(
-    X: Array[S],
-    y: Array[T],
+    X: ArrayRead[S],
+    y: ArrayRead[T],
     cover: Cover[T],
     clustering: Clustering[S],
     n_jobs: int = 1,
@@ -168,8 +168,8 @@ def mapper_connected_components(
 
 
 def mapper_graph(
-    X: Array[S],
-    y: Array[T],
+    X: ArrayRead[S],
+    y: ArrayRead[T],
     cover: Cover[T],
     clustering: Clustering[S],
     n_jobs: int = 1,
@@ -218,7 +218,7 @@ def mapper_graph(
 
 
 def aggregate_graph(
-    X: Array[S], graph: nx.Graph, agg: Callable[..., Any]
+    X: ArrayRead[S], graph: nx.Graph, agg: Callable[..., Any]
 ) -> dict[int, Any]:
     """
     Apply an aggregation function to the nodes of a graph.
@@ -246,7 +246,7 @@ def aggregate_graph(
     return agg_values
 
 
-def proximity_net(search: SpatialSearch[S], X: Array[S]) -> Iterator[list[int]]:
+def proximity_net(search: SpatialSearch[S], X: ArrayRead[S]) -> Iterator[list[int]]:
     """
     Covers the dataset using proximity-net.
 
@@ -283,7 +283,7 @@ class TrivialCover(ParamsMixin, Generic[T]):
     dataset.
     """
 
-    def apply(self, X: Array[T]) -> Iterator[list[int]]:
+    def apply(self, X: ArrayRead[T]) -> Iterator[list[int]]:
         """
         Covers the dataset with a single open set.
 
@@ -317,7 +317,9 @@ def __init__(
         self.clustering = clustering
         self.verbose = verbose
 
-    def fit(self, X: Array[T], y: Optional[Array[T]] = None) -> FailSafeClustering[T]:
+    def fit(
+        self, X: ArrayRead[T], y: Optional[ArrayRead[T]] = None
+    ) -> FailSafeClustering[T]:
         self._clustering = (
             TrivialClustering() if self.clustering is None else self.clustering
         )
@@ -347,7 +349,9 @@ class TrivialClustering(ParamsMixin, Generic[T]):
     def __init__(self) -> None:
         pass
 
-    def fit(self, X: Array[T], _y: Optional[Array[T]] = None) -> TrivialClustering[T]:
+    def fit(
+        self, X: ArrayRead[T], _y: Optional[ArrayRead[T]] = None
+    ) -> TrivialClustering[T]:
         """
         Fit the clustering algorithm to the data.
 

src/tdamapper/cover.py

@@ -16,29 +16,33 @@
 
 from tdamapper._common import ParamsMixin, warn_user
 from tdamapper.core import proximity_net
-from tdamapper.protocols import Array, Metric
+from tdamapper.protocols import ArrayRead, Metric
 from tdamapper.utils.metrics import MetricLiteral, chebyshev, get_metric
 from tdamapper.utils.vptree import PivotingStrategy, VPTree, VPTreeKind
 
 T = TypeVar("T")
+T_contra = TypeVar("T_contra", contravariant=True)
 S = TypeVar("S")
+S_contra = TypeVar("S_contra", contravariant=True)
 
 
-class _Pullback(Generic[S, T]):
+class _Pullback(Generic[S_contra, T_contra]):
 
-    def __init__(self, fun: Callable[[S], T], dist: Metric[T]):
+    def __init__(
+        self, fun: Callable[[S_contra], T_contra], dist: Metric[T_contra]
+    ) -> None:
         self.fun = fun
         self.dist = dist
 
-    def __call__(self, x: S, y: S) -> float:
+    def __call__(self, x: S_contra, y: S_contra) -> float:
         return self.dist(self.fun(x), self.fun(y))
 
 
-def _snd(x: tuple[Any, ...]) -> Any:
+def _snd(x: tuple[T, ...]) -> T:
     return x[1]
 
 
-class BallCover(ParamsMixin, Generic[T]):
+class BallCover(ParamsMixin, Generic[T_contra]):
     """
     Cover algorithm based on `ball proximity function`, which covers data with
     open balls of fixed radius.
@@ -67,8 +71,8 @@ class BallCover(ParamsMixin, Generic[T]):
     """
 
     _radius: float
-    _data: list[tuple[int, T]]
-    _vptree: VPTree[tuple[int, T]]
+    _data: list[tuple[int, T_contra]]
+    _vptree: VPTree[tuple[int, T_contra]]
 
     def __init__(
         self,
@@ -88,7 +92,7 @@ def __init__(
         self.leaf_radius = leaf_radius
         self.pivoting = pivoting
 
-    def fit(self, X: Array[T]) -> BallCover[T]:
+    def fit(self, X: ArrayRead[T_contra]) -> BallCover[T_contra]:
         """
         Train internal parameters.
 
@@ -112,7 +116,7 @@ def fit(self, X: Array[T]) -> BallCover[T]:
         )
         return self
 
-    def search(self, x: T) -> list[int]:
+    def search(self, x: T_contra) -> list[int]:
         """
         Return a list of neighbors for the query point.
 
@@ -130,7 +134,7 @@ def search(self, x: T) -> list[int]:
         )
         return [x for (x, _) in neighs]
 
-    def apply(self, X: Array[T]) -> Iterator[list[int]]:
+    def apply(self, X: ArrayRead[T_contra]) -> Iterator[list[int]]:
         """
         Covers the dataset using proximity-net.
 
@@ -144,7 +148,7 @@ def apply(self, X: Array[T]) -> Iterator[list[int]]:
         return proximity_net(self, X)
 
 
-class KNNCover(ParamsMixin, Generic[T]):
+class KNNCover(ParamsMixin, Generic[T_contra]):
     """
     Cover algorithm based on `KNN proximity function`, which covers data using
     k-nearest neighbors (KNN).
@@ -173,8 +177,8 @@ class KNNCover(ParamsMixin, Generic[T]):
     """
 
     _neighbors: int
-    _data: list[tuple[int, T]]
-    _vptree: VPTree[tuple[int, T]]
+    _data: list[tuple[int, T_contra]]
+    _vptree: VPTree[tuple[int, T_contra]]
 
     def __init__(
         self,
@@ -194,7 +198,7 @@ def __init__(
         self.leaf_radius = leaf_radius
         self.pivoting = pivoting
 
-    def fit(self, X: Array[T]) -> KNNCover[T]:
+    def fit(self, X: ArrayRead[T_contra]) -> KNNCover[T_contra]:
         """
         Train internal parameters.
 
@@ -218,7 +222,7 @@ def fit(self, X: Array[T]) -> KNNCover[T]:
         )
         return self
 
-    def search(self, x: T) -> list[int]:
+    def search(self, x: T_contra) -> list[int]:
         """
         Return a list of neighbors for the query point.
 
@@ -233,7 +237,7 @@ def search(self, x: T) -> list[int]:
         neighs = self._vptree.knn_search((-1, x), self._neighbors)
         return [x for (x, _) in neighs]
 
-    def apply(self, X: Array[T]) -> Iterator[list[int]]:
+    def apply(self, X: ArrayRead[T_contra]) -> Iterator[list[int]]:
         """
         Covers the dataset using proximity-net.
 
@@ -309,7 +313,7 @@ def _get_bounds(
         _delta[(_delta >= -eps) & (_delta <= eps)] = self._n_intervals
         return _min, _max, _delta
 
-    def fit(self, X: Array[NDArray[np.float_]]) -> BaseCubicalCover:
+    def fit(self, X: ArrayRead[NDArray[np.float_]]) -> BaseCubicalCover:
         """
         Train internal parameters.
 
@@ -408,7 +412,7 @@ def __init__(
             pivoting=pivoting,
         )
 
-    def apply(self, X: Array[NDArray[np.float_]]) -> Iterator[list[int]]:
+    def apply(self, X: ArrayRead[NDArray[np.float_]]) -> Iterator[list[int]]:
         """
         Covers the dataset using proximity-net.
 
@@ -471,7 +475,7 @@ def __init__(
         )
 
     def _landmarks(
-        self, X: Array[NDArray[np.float_]]
+        self, X: ArrayRead[NDArray[np.float_]]
     ) -> dict[tuple[float], NDArray[np.float_]]:
         lmrks = {}
         for x in X:
@@ -480,7 +484,7 @@ def _landmarks(
                 lmrks[lmrk] = x
         return lmrks
 
-    def apply(self, X: Array[NDArray[np.float_]]) -> Iterator[list[int]]:
+    def apply(self, X: ArrayRead[NDArray[np.float_]]) -> Iterator[list[int]]:
         """
         Covers the dataset using landmarks.
 
@@ -578,7 +582,7 @@ def _get_cubical_cover(self) -> Union[ProximityCubicalCover, StandardCubicalCove
             "The only possible values for algorithm are 'standard' and 'proximity'."
         )
 
-    def fit(self, X: Array[NDArray[np.float_]]) -> CubicalCover:
+    def fit(self, X: ArrayRead[NDArray[np.float_]]) -> CubicalCover:
         """
         Train internal parameters.
 
@@ -604,7 +608,7 @@ def search(self, x: NDArray[np.float_]) -> list[int]:
         """
         return self._cubical_cover.search(x)
 
-    def apply(self, X: Array[NDArray[np.float_]]) -> Iterator[list[int]]:
+    def apply(self, X: ArrayRead[NDArray[np.float_]]) -> Iterator[list[int]]:
         """
         Covers the dataset using hypercubes.