Added types and docs

Author: lucasimi

src/tdamapper/plot_backends/plot_plotly.py

@@ -11,6 +11,7 @@
 import numpy as np
 import plotly.colors as pc
 import plotly.graph_objects as go
+from numpy.typing import NDArray
 
 from tdamapper.core import ATTR_SIZE, aggregate_graph
 
@@ -304,13 +305,15 @@ def plot(
         )
         return self.fig
 
-    def _node_pos_array(self):
+    def _node_pos_array(self) -> tuple[list[float], ...]:
         return tuple(
             [self.positions[n][i] for n in self.graph.nodes()] for i in range(self.dim)
         )
 
-    def _edge_pos_array(self):
-        edges_arr = tuple([] for i in range(self.dim))
+    def _edge_pos_array(self) -> tuple[list[Optional[float]], ...]:
+        edges_arr: tuple[list[Optional[float]], ...] = tuple(
+            [] for i in range(self.dim)
+        )
         for edge in self.graph.edges():
             pos0, pos1 = self.positions[edge[0]], self.positions[edge[1]]
             for i in range(self.dim):
@@ -530,7 +533,7 @@ def update_figure(
         if cmaps is not None:
             self.set_cmap(cmaps[0])
 
-    def _nodes_trace(self, node_pos_arr: ) -> Union[go.Scatter, go.Scatter3d]:
+    def _nodes_trace(self, node_pos_arr) -> Union[go.Scatter, go.Scatter3d]:
         scatter = dict(
             name=_NODES_TRACE,
             x=node_pos_arr[0],
@@ -699,7 +702,8 @@ def set_ui(
 
         if cmaps is not None:
             cmaps_plotly = [PLOTLY_CMAPS.get(c.lower()) for c in cmaps]
-            ui_menu_cmap = self._ui_menu_cmap(cmaps_plotly)
+            cmaps_plotly_ok = [c for c in cmaps_plotly if c is not None]
+            ui_menu_cmap = self._ui_menu_cmap(cmaps_plotly_ok)
 
         if colors is not None and agg is not None and titles is not None:
             ui_menu_color = self._ui_menu_color(colors, titles, agg)
@@ -798,7 +802,9 @@ def _update_cmap(cmap: str) -> dict:
             yanchor="top",
         )
 
-    def _ui_menu_color(self, colors: NDArray[np.float64], titles: list[str], agg: Callable) -> dict:
+    def _ui_menu_color(
+        self, colors: NDArray[np.float64], titles: list[str], agg: Callable
+    ) -> dict:
         colors_arr = np.array(colors)
         colors_num = colors_arr.shape[1] if colors_arr.ndim == 2 else 1
 

src/tdamapper/plot_backends/plot_pyvis.py

@@ -4,10 +4,13 @@
 """
 
 import math
+from typing import Any, Callable
 
+import numpy as np
 import plotly.colors as pc
 import plotly.graph_objects as go
 import plotly.io as pio
+from numpy.typing import NDArray
 from pyvis.network import Network
 
 from tdamapper.core import aggregate_graph
@@ -19,12 +22,18 @@
 _TICKS_NUM = 10
 
 
-def _fmt(x, max_len=3):
+def _fmt(x: Any, max_len: int = 3) -> str:
     fmt = f".{max_len}g"
     return f"{x:{fmt}}"
 
 
-def _colorbar(height, cmap, cmin, cmax, title):
+def _colorbar(
+    height: int,
+    cmap: str,
+    cmin: float,
+    cmax: float,
+    title: str,
+) -> go.Figure:
     colorbar_fig = go.Figure()
     colorbar_fig.add_trace(
         go.Scatter(
@@ -68,7 +77,7 @@ def _colorbar(height, cmap, cmin, cmax, title):
     return colorbar_fig
 
 
-def _combine(network, colorbar):
+def _combine(network: Network, colorbar: go.Figure) -> str:
     network_html = network.generate_html()
     colorbar_html = pio.to_html(
         colorbar,
@@ -130,15 +139,28 @@ def _combine(network, colorbar):
 
 def plot_pyvis(
     mapper_plot,
-    output_file,
-    colors,
-    node_size,
-    agg,
-    title,
-    width,
-    height,
-    cmap,
-):
+    output_file: str,
+    colors: NDArray[np.float64],
+    node_size: float,
+    agg: Callable,
+    title: str,
+    width: int,
+    height: int,
+    cmap: str,
+) -> None:
+    """
+    Generates a pyvis network visualization of the Mapper graph and saves it to an HTML file.
+
+    :param mapper_plot: The Mapper plot object containing the graph and positions.
+    :param output_file: The path to the output HTML file.
+    :param colors: A 2D array of colors for the nodes.
+    :param node_size: The size of the nodes in the graph.
+    :param agg: A callable function to aggregate the graph.
+    :param title: The title for the colorbar.
+    :param width: The width of the network visualization.
+    :param height: The height of the network visualization.
+    :param cmap: The colormap to use for the nodes.
+    """
     net, cmin, cmax = _compute_net(
         mapper_plot=mapper_plot,
         width=width,
@@ -150,19 +172,19 @@ def plot_pyvis(
     )
     colorbar = _colorbar(height=height, cmap=cmap, cmin=cmin, cmax=cmax, title=title)
     combined_html = _combine(net, colorbar)
-    with open(output_file, "w") as file:
+    with open(output_file, "w", encoding="utf-8") as file:
         file.write(combined_html)
 
 
 def _compute_net(
     mapper_plot,
-    colors,
-    node_size,
-    agg,
-    width,
-    height,
-    cmap,
-):
+    colors: NDArray[np.float64],
+    node_size: float,
+    agg: Callable,
+    width: int,
+    height: int,
+    cmap: str,
+) -> tuple[Network, float, float]:
     net = Network(
         height=f"{height}px",
         width=f"{width}px",