Improved type-hints

Author: lucasimi

Makefile

@@ -8,6 +8,10 @@ all: install
 install:
 	$(PIP) install -e .[dev]
 
+.PHONY: typecheck
+typecheck:
+	mypy src --ignore-missing-imports --install-types
+
 .PHONY: test
 test:
 	coverage run --source=src -m pytest tests/test_unit_*.py

pyproject.toml

@@ -55,6 +55,7 @@ dev = [
     "nbformat>=4.2.0",
     "umap-learn<0.6.0",
     "nicegui>=2.18.0,<3.0.0",
+    "mypy",
 ]
 app = [
     "pandas<3.0.0",
@@ -113,3 +114,6 @@ filterwarnings = [
 markers = [
     "module_under_test",
 ]
+
+[tool.mypy]
+ignore_missing_imports = true

src/tdamapper/app.py

@@ -176,7 +176,7 @@ def _umap(X: NDArray[np.float64]) -> NDArray[np.float64]:
 
 def run_mapper(
     df: pd.DataFrame, **kwargs: dict[str, Any]
-) -> Optional[tuple[MapperAlgorithm, pd.DataFrame]]:
+) -> Optional[tuple[nx.Graph, pd.DataFrame]]:
     """
     Run the Mapper algorithm on the provided DataFrame.
 

src/tdamapper/core.py

@@ -213,14 +213,13 @@ def mapper_graph(
     :rtype: :class:`networkx.Graph`
     """
     itm_lbls = mapper_labels(X, y, cover, clustering, n_jobs=n_jobs)
-    graph = nx.Graph()
+    graph: nx.Graph = nx.Graph()
     for n, lbls in enumerate(itm_lbls):
         for lbl in lbls:
             if not graph.has_node(lbl):
                 graph.add_node(lbl, **{ATTR_SIZE: 0, ATTR_IDS: []})
-            nodes = graph.nodes()
-            nodes[lbl][ATTR_SIZE] += 1
-            nodes[lbl][ATTR_IDS].append(n)
+            graph.nodes[lbl][ATTR_SIZE] += 1
+            graph.nodes[lbl][ATTR_IDS].append(n)
     for lbls in itm_lbls:
         lbls_len = len(lbls)
         for i in range(lbls_len):
@@ -256,9 +255,8 @@ def aggregate_graph(X: ArrayLike, graph: nx.Graph, agg: Callable) -> dict[Any, A
     :rtype: dict
     """
     agg_values = {}
-    nodes = graph.nodes()
-    for node_id in nodes:
-        node_values = [X[i] for i in nodes[node_id][ATTR_IDS]]
+    for node_id in graph.nodes:
+        node_values = [X[i] for i in graph.nodes[node_id][ATTR_IDS]]
         agg_value = agg(node_values)
         agg_values[node_id] = agg_value
     return agg_values

src/tdamapper/cover.py

@@ -261,9 +261,9 @@ class BaseCubicalCover:
 
     _n_intervals: int
     _overlap_frac: float
-    _min: NDArray
-    _max: NDArray
-    _delta: NDArray
+    _min: NDArray[np.float64]
+    _max: NDArray[np.float64]
+    _delta: NDArray[np.float64]
     _cover: BallCover
 
     def __init__(
@@ -282,7 +282,7 @@ def __init__(
         self.leaf_radius = leaf_radius
         self.pivoting = pivoting
 
-    def _get_center(self, x: NDArray) -> tuple[tuple, NDArray]:
+    def _get_center(self, x: NDArray[np.float64]) -> tuple[tuple, NDArray]:
         offset = self._offset(x)
         center = self._phi(x)
         return tuple(offset), center
@@ -291,20 +291,22 @@ def _get_overlap_frac(self, dim: int, overlap_vol_frac: float) -> float:
         beta = math.pow(1.0 - overlap_vol_frac, 1.0 / dim)
         return 1.0 - 1.0 / (2.0 - beta)
 
-    def _offset(self, x: NDArray) -> NDArray:
+    def _offset(self, x: NDArray[np.float64]) -> NDArray[np.float64]:
         return np.minimum(self._n_intervals - 1, np.floor(self._gamma_n(x)))
 
-    def _phi(self, x: NDArray) -> NDArray:
+    def _phi(self, x: NDArray[np.float64]) -> NDArray[np.float64]:
         offset = self._offset(x)
         return self._gamma_n_inv(0.5 + offset)
 
-    def _gamma_n(self, x: NDArray) -> NDArray:
+    def _gamma_n(self, x: NDArray[np.float64]) -> NDArray[np.float64]:
         return self._n_intervals * (x - self._min) / self._delta
 
-    def _gamma_n_inv(self, x: NDArray) -> NDArray:
+    def _gamma_n_inv(self, x: NDArray[np.float64]) -> NDArray[np.float64]:
         return self._min + self._delta * x / self._n_intervals
 
-    def _get_bounds(self, X: ArrayLike) -> tuple[NDArray, NDArray, NDArray]:
+    def _get_bounds(
+        self, X: ArrayLike
+    ) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
         if (X is None) or len(X) == 0:
             raise ValueError("The dataset is empty or None.")
         _min, _max = X[0], X[0]
@@ -408,12 +410,12 @@ class ProximityCubicalCover(BaseCubicalCover, ParamsMixin, Proximity):
 
     def __init__(
         self,
-        n_intervals=1,
-        overlap_frac=None,
-        kind="flat",
-        leaf_capacity=1,
-        leaf_radius=None,
-        pivoting=None,
+        n_intervals: int = 1,
+        overlap_frac: Optional[float] = None,
+        kind: str = "flat",
+        leaf_capacity: int = 1,
+        leaf_radius: Optional[float] = None,
+        pivoting: Optional[str] = None,
     ):
         super().__init__(
             n_intervals=n_intervals,
@@ -479,15 +481,15 @@ def __init__(
             pivoting=pivoting,
         )
 
-    def _landmarks(self, X: ArrayLike) -> dict[tuple, NDArray]:
+    def _landmarks(self, X: NDArray[np.float64]) -> dict[tuple, NDArray[np.float64]]:
         lmrks = {}
         for x in X:
             lmrk, _ = self._get_center(x)
             if lmrk not in lmrks:
                 lmrks[lmrk] = x
         return lmrks
 
-    def apply(self, X: ArrayLike) -> Generator[list[int]]:
+    def apply(self, X: NDArray[np.float64]) -> Generator[list[int]]:
         """
         Covers the dataset using landmarks.
 
@@ -575,7 +577,7 @@ def __init__(
         self.pivoting = pivoting
 
     def _get_cubical_cover(self) -> Union[ProximityCubicalCover, StandardCubicalCover]:
-        params = dict(
+        params: dict[str, Any] = dict(
             n_intervals=self.n_intervals,
             overlap_frac=self.overlap_frac,
             kind=self.kind,
@@ -593,7 +595,7 @@ def _get_cubical_cover(self) -> Union[ProximityCubicalCover, StandardCubicalCove
                 "'proximity'."
             )
 
-    def fit(self, X: ArrayLike) -> CubicalCover:
+    def fit(self, X: NDArray[np.float64]) -> CubicalCover:
         """
         Train internal parameters.
 
@@ -623,7 +625,7 @@ def search(self, x: Any) -> list[int]:
         """
         return self._cubical_cover.search(x)
 
-    def apply(self, X: ArrayLike) -> Generator[list[int]]:
+    def apply(self, X: NDArray[np.float64]) -> Generator[list[int]]:
         """
         Covers the dataset using hypercubes.
 

src/tdamapper/learn.py

@@ -9,8 +9,16 @@
 scikit-learn's conventions for estimators.
 """
 
-from tdamapper._common import EstimatorMixin, ParamsMixin, clone
+from __future__ import annotations
+
+from typing import Optional
+
+import networkx as nx
+
+from tdamapper._common import ArrayLike, EstimatorMixin, ParamsMixin, clone
 from tdamapper.core import (
+    Clustering,
+    Cover,
     FailSafeClustering,
     TrivialClustering,
     TrivialCover,
@@ -45,17 +53,19 @@ class MapperClustering(EstimatorMixin, ParamsMixin):
     :type n_jobs: int
     """
 
+    labels_: list[int]
+
     def __init__(
         self,
-        cover=None,
-        clustering=None,
-        n_jobs=1,
+        cover: Optional[Cover] = None,
+        clustering: Optional[Clustering] = None,
+        n_jobs: int = 1,
     ):
         self.cover = cover
         self.clustering = clustering
         self.n_jobs = n_jobs
 
-    def fit(self, X, y=None):
+    def fit(self, X: ArrayLike, y: Optional[ArrayLike] = None) -> MapperClustering:
         """
         Fit the clustering algorithm to the data.
 
@@ -123,21 +133,28 @@ class MapperAlgorithm(EstimatorMixin, ParamsMixin):
     :type n_jobs: int
     """
 
+    _cover: Cover
+    _clustering: Clustering
+    _verbose: bool
+    _failsafe: bool
+    _n_jobs: int
+    graph_: nx.Graph
+
     def __init__(
         self,
-        cover=None,
-        clustering=None,
-        failsafe=True,
-        verbose=True,
-        n_jobs=1,
+        cover: Optional[Cover] = None,
+        clustering: Optional[Clustering] = None,
+        failsafe: bool = True,
+        verbose: bool = True,
+        n_jobs: int = 1,
     ):
         self.cover = cover
         self.clustering = clustering
         self.failsafe = failsafe
         self.verbose = verbose
         self.n_jobs = n_jobs
 
-    def fit(self, X, y=None):
+    def fit(self, X: ArrayLike, y: Optional[ArrayLike] = None) -> MapperAlgorithm:
         """
         Create the Mapper graph and store it for later use.
 
@@ -150,6 +167,8 @@ def fit(self, X, y=None):
         :type y: array-like of shape (n, k) or list-like of length n
         :return: The object itself.
         """
+        if y is None:
+            y = X
         X, y = self._validate_X_y(X, y)
         self._cover = TrivialCover() if self.cover is None else self.cover
         self._clustering = (
@@ -176,7 +195,7 @@ def fit(self, X, y=None):
         self._set_n_features_in(X)
         return self
 
-    def fit_transform(self, X, y):
+    def fit_transform(self, X: ArrayLike, y: ArrayLike) -> nx.Graph:
         """
         Create the Mapper graph.
 

src/tdamapper/plot.py

@@ -7,9 +7,9 @@
 import igraph as ig
 import networkx as nx
 import numpy as np
+from numpy.typing import NDArray
 from plotly import graph_objects as go
 
-from tdamapper._common import ArrayLike
 from tdamapper.plot_backends.plot_matplotlib import plot_matplotlib
 from tdamapper.plot_backends.plot_plotly import plot_plotly, plot_plotly_update
 from tdamapper.plot_backends.plot_pyvis import plot_pyvis
@@ -155,7 +155,7 @@ def plot_matplotlib(
 
     def plot_plotly(
         self,
-        colors: Union[np.ndarray, list[float]],
+        colors: Union[NDArray[np.float64], list[float]],
         node_size: Union[int, float, list[Union[int, float]]] = 1,
         agg: Callable = np.nanmean,
         title: Optional[Union[str, list[str]]] = None,
@@ -210,7 +210,7 @@ def plot_plotly(
     def plot_plotly_update(
         self,
         fig: go.Figure,
-        colors: Optional[Union[np.ndarray, list[float]]] = None,
+        colors: Optional[Union[NDArray[np.float64], list[float]]] = None,
         node_size: Optional[Union[int, float, list[Union[int, float]]]] = None,
         agg: Optional[Callable] = None,
         title: Optional[Union[str, list[str]]] = None,
@@ -271,7 +271,7 @@ def plot_plotly_update(
     def plot_pyvis(
         self,
         output_file: str,
-        colors: ArrayLike,
+        colors: NDArray[np.float64],
         node_size: int = 1,
         agg: Callable = np.nanmean,
         title: Optional[str] = None,

src/tdamapper/plot_backends/plot_matplotlib.py

@@ -4,10 +4,11 @@
 """
 
 import math
-from typing import Any, Callable
+from typing import Any, Callable, Optional
 
 import matplotlib.pyplot as plt
 import networkx as nx
+import numpy as np
 from matplotlib.collections import LineCollection
 from numpy.typing import NDArray
 
@@ -26,8 +27,8 @@ def plot_matplotlib(
     mapper_plot,
     width: int,
     height: int,
-    title: str,
-    colors: NDArray,
+    title: Optional[str],
+    colors: NDArray[np.float64],
     node_size: float,
     agg: Callable,
     cmap: str,
@@ -56,8 +57,8 @@ def plot_matplotlib(
 
 def _plot_nodes(
     mapper_plot,
-    title: str,
-    colors: NDArray,
+    title: Optional[str],
+    colors: NDArray[np.float64],
     node_size: float,
     agg: Callable,
     cmap: str,
@@ -95,9 +96,10 @@ def _plot_nodes(
         ax=ax,
         format="%.2g",
     )
-    colorbar.set_label(title, color=_NODE_OUTER_COLOR)
+    if title is not None:
+        colorbar.set_label(title, color=_NODE_OUTER_COLOR)
     colorbar.set_alpha(1.0)
-    colorbar.outline.set_color(_NODE_OUTER_COLOR)
+    # colorbar.outline.set_color(_NODE_OUTER_COLOR)
     colorbar.ax.yaxis.set_tick_params(
         color=_NODE_OUTER_COLOR, labelcolor=_NODE_OUTER_COLOR
     )

src/tdamapper/plot_backends/plot_plotly.py

@@ -94,7 +94,7 @@ def _to_cmaps(cmap: Optional[Union[str, list[str]]]) -> list[str]:
         raise ValueError(f"Invalid cmap type: {type(cmap)}. Expected str or list[str].")
 
 
-def _to_colors(colors: Union[np.ndarray, list[float]]) -> np.ndarray:
+def _to_colors(colors: Union[NDArray[np.float64], list[float]]) -> np.ndarray:
     """Convert colors to a numpy array."""
     colors_arr = np.array(colors)
     if colors_arr.ndim == 1:
@@ -146,7 +146,7 @@ def _get_cmap_rgb(cmap: str):
 
 def plot_plotly(
     mapper_plot,
-    colors: Union[np.ndarray, list[float]],
+    colors: Union[NDArray[np.float64], list[float]],
     node_size: Optional[Union[int, float, list[Union[int, float]]]] = None,
     title: Optional[Union[str, list[str]]] = None,
     agg: Callable = np.nanmean,
@@ -192,7 +192,7 @@ def plot_plotly_update(
     width: Optional[int] = None,
     height: Optional[int] = None,
     node_size: Optional[Union[int, float, list[Union[int, float]]]] = None,
-    colors: Optional[Union[np.ndarray, list[float]]] = None,
+    colors: Optional[Union[NDArray[np.float64], list[float]]] = None,
     title: Optional[Union[str, list[str]]] = None,
     agg: Optional[Callable] = None,
     cmap: Optional[Union[str, list[str]]] = None,
@@ -266,7 +266,7 @@ def __init__(self, mapper_plot, fig: Optional[go.Figure] = None):
 
     def plot(
         self,
-        colors: np.ndarray,
+        colors: NDArray[np.float64],
         node_sizes: list[float],
         titles: list[str],
         agg: Callable,