Fixed issue with None colors argument

Author: lucasimi

src/tdamapper/plot.py

@@ -42,6 +42,13 @@ def _node_pos(graph, dim, seed, iterations):
         iterations=iterations)
 
 
+def _node_col(graph, colors, agg, default=0.5):
+    if colors is not None:
+        return aggregate_graph(colors, graph, agg)
+    else:
+        return [default for _ in graph.nodes()]
+
+
 def _node_pos_array(graph, dim, node_pos):
     return tuple([node_pos[n][i] for n in graph.nodes()] for i in range(dim))
 
@@ -116,7 +123,7 @@ def __init__(self, graph, dim,
         self.__width = width
         self.__height = height
         self.__cmap = cmap
-        node_col = aggregate_graph(self.__colors, self.__graph, self.__agg)
+        node_col = _node_col(self.__graph, self.__colors, self.__agg)
         self.__fig = self._figure(node_col)
         #self._update_traces_col()
         #self._update_layout()
@@ -311,10 +318,10 @@ def _update_traces_pos(self):
 
     def _update_traces_col(self):
         if (self.__colors is not None) and (self.__agg is not None):
-            colors_agg = aggregate_graph(self.__colors, self.__graph, self.__agg)
-            colors_list = list(colors_agg.values())
-            self._update_node_trace_col(colors_agg, colors_list)
-            self._update_edge_trace_col(colors_agg, colors_list)
+            nodes_col = _node_col(self.__graph, self.__colors, self.__agg)
+            colors_list = list(nodes_col.values())
+            self._update_node_trace_col(nodes_col, colors_list)
+            self._update_edge_trace_col(nodes_col, colors_list)
 
     def _update_edge_trace_col(self, colors_agg, colors_list):
         colors_avg = []
@@ -538,7 +545,7 @@ def _plot_nodes(self, ax, nodes_pos):
         nodes_arr = _node_pos_array(self.__graph, self.__dim, nodes_pos)
         attr_size = nx.get_node_attributes(self.__graph, ATTR_SIZE)
         max_size = max(attr_size.values()) if attr_size else 1.0
-        colors_agg = aggregate_graph(self.__colors, self.__graph, self.__agg)
+        colors_agg = _node_col(self.__graph, self.__colors, self.__agg)
         marker_color = [colors_agg[n] for n in self.__graph.nodes()]
         marker_size = [200.0 * math.sqrt(attr_size[n] / max_size) for n in self.__graph.nodes()]
         verts = ax.scatter(