Adding vertex plotting option to plot_network

wntr/graphics/network.py

@@ -43,7 +43,7 @@ def _format_link_attribute(link_attribute, wn):
 
     return link_attribute
 
-def plot_network(wn, node_attribute=None, link_attribute=None, title=None,
+def plot_network(wn, node_attribute=None, link_attribute=None, title=None, plot_vertices=False,
                node_size=20, node_range=[None,None], node_alpha=1, node_cmap=None, node_labels=False,
                link_width=1, link_range=[None,None], link_alpha=1, link_cmap=None, link_labels=False,
                add_colorbar=True, node_colorbar_label='Node', link_colorbar_label='Link', 
@@ -79,7 +79,10 @@ def plot_network(wn, node_attribute=None, link_attribute=None, title=None,
           {linkid: x} where linkid is a string and x is a float.
 
     title: str, optional
-        Plot title 
+        Plot title
+
+    plot_vertices: bool, optional
+        Draw the plot using vertices. May be slow on large models.
 
     node_size: int, optional
         Node size 
@@ -140,7 +143,10 @@ def plot_network(wn, node_attribute=None, link_attribute=None, title=None,
         ax = plt.gca()
 
     # Graph
-    G = wn.to_graph()
+    if plot_vertices:
+        G = wn.to_graph(include_vertices=True)
+    else: 
+        G = wn.to_graph()
     if not directed:
         G = G.to_undirected()
 
@@ -169,9 +175,13 @@ def plot_network(wn, node_attribute=None, link_attribute=None, title=None,
         nodelist,nodecolor = zip(*node_attribute.items())
 
     else:
-        nodelist = None
+        nodelist = list(G.nodes())
         nodecolor = 'k'
-
+        if plot_vertices:
+            node_types = nx.get_node_attributes(G, "type")
+            vertex_nodes = [key for key in node_types if node_types[key] == "Vertex"]
+            nodelist = list(set(nodelist) - set(vertex_nodes))
+
     add_link_colorbar = add_colorbar
     if link_attribute is not None:
 
@@ -194,27 +204,109 @@ def plot_network(wn, node_attribute=None, link_attribute=None, title=None,
         attr = {}
         for link_name, value in link_attribute.items():
             link = wn.get_link(link_name)
-            attr[(link.start_node_name, link.end_node_name, link_name)] = value
+            attr[(link.start_node_name, link.end_node_name, link.name)] = value
+        if plot_vertices:
+            vertex_attr = {}
+            edge_base_links = nx.get_edge_attributes(G, "base_link")
+            for vertex_edge in edge_base_links:
+                try:
+                    vertex_attr[vertex_edge] = link_attribute[edge_base_links[vertex_edge]]
+                except KeyError:
+                    pass
+            attr = vertex_attr
+
         link_attribute = attr
-
-        linklist,linkcolor = zip(*link_attribute.items())
+        linklist, linkcolor = zip(*link_attribute.items())
+        # if plot_vertices:
+        #     link_names = nx.get_edge_attributes(G, "name")
+
+        #     for edge in G.edges:
+        #     node_types = nx.get_node_attributes(G, "type")
+
+        #         u, v, _ = edge
     else:
         linklist = None
         linkcolor = 'k'
 
     if title is not None:
         ax.set_title(title)
+
+    # create graph with new edges generated from vertices for vertex plotting
+    # if plot_vertices:
+    #     if linklist:
+    #         vertex_linklist = []
+    #         vertex_linkcolor = []
+    #     else:
+    #         vertex_linklist = None
+    #         vertex_linkcolor = 'k'
+    #     G_vertices = nx.Graph()
+
+    #     for link_name in wn.link_name_list:
+    #         link = wn.get_link(link_name)
+    #         edge = (link.start_node_name, link.end_node_name, link_name)
+    #         edge_in_linklist = linklist and edge in linklist
+    #         num_vertices = len(link.vertices)
+    #         start_node = link.start_node.name
+    #         end_node = link.end_node.name
+
+    #         if num_vertices > 0:
+    #             base_name = link_name+"_v"
+
+    #             for vertex_ind in range(num_vertices):
+    #                 val = vertex_ind + 1
+    #                 name = base_name+str(val)
+    #                 # G_vertices.add_node(name, size=0)
+    #                 G_vertices.add_edge(start_node, name)
+
+    #                 if edge_in_linklist:
+    #                     vertex_linklist.append((start_node, name))
+    #                     vertex_linkcolor.append(link_attribute[edge])
+
+    #                 if val == num_vertices:
+    #                     G_vertices.add_edge(name, end_node)
+    #                     if edge_in_linklist:
+    #                         vertex_linklist.append((name, end_node))
+    #                         vertex_linkcolor.append(link_attribute[edge])
+
+    #                 pos[name] = link.vertices[vertex_ind]
+    #                 start_node = name
+    #         else:
+    #             G_vertices.add_edge(start_node, end_node)
+
+    #             if edge_in_linklist:
+    #                 vertex_linklist.append((start_node, end_node))
+    #                 vertex_linkcolor.append(link_attribute[edge])
+
+    #     # Create vertex_nodelist
+    #     if nodelist is None:
+    #         vertex_nodelist = wn.node_name_list
+    #     else:
+    #         vertex_nodelist = nodelist
+
+    #     edge_background = nx.draw_networkx_edges(G_vertices, pos, edge_color='grey', 
+    #                                     width=0.5, ax=ax)
+
+    #     nodes = nx.draw_networkx_nodes(G_vertices, pos, node_color=nodecolor, nodelist=vertex_nodelist, node_size=node_size,
+    #                                    alpha=node_alpha, cmap=node_cmap, vmin=node_range[0], 
+    #                                    vmax=node_range[1], linewidths=0, ax=ax)
+
+    #     edges = nx.draw_networkx_edges(G_vertices, pos, edgelist=vertex_linklist, 
+    #             edge_color=vertex_linkcolor, width=link_width, alpha=link_alpha, edge_cmap=link_cmap, 
+    #             edge_vmin=link_range[0], edge_vmax=link_range[1], ax=ax)
+
+
 
     edge_background = nx.draw_networkx_edges(G, pos, edge_color='grey', 
-                                             width=0.5, ax=ax)
+                                            width=0.5, ax=ax)
 
     nodes = nx.draw_networkx_nodes(G, pos, 
             nodelist=nodelist, node_color=nodecolor, node_size=node_size, 
-            alpha=node_alpha, cmap=node_cmap, vmin=node_range[0], vmax = node_range[1], 
+            alpha=node_alpha, cmap=node_cmap, vmin=node_range[0], vmax=node_range[1], 
             linewidths=0, ax=ax)
     edges = nx.draw_networkx_edges(G, pos, edgelist=linklist, 
             edge_color=linkcolor, width=link_width, alpha=link_alpha, edge_cmap=link_cmap, 
             edge_vmin=link_range[0], edge_vmax=link_range[1], ax=ax)
+
     if node_labels:
         labels = dict(zip(wn.node_name_list, wn.node_name_list))
         nx.draw_networkx_labels(G, pos, labels, font_size=7, ax=ax)
@@ -663,7 +755,7 @@ def plot_leaflet_network(wn, node_attribute=None, link_attribute=None,
 
     m.save(filename)
 
-def network_animation(wn, node_attribute=None, link_attribute=None, title=None,
+def network_animation(wn, node_attribute=None, link_attribute=None, title=None, plot_vertices=False,
                node_size=20, node_range=[None,None], node_alpha=1, node_cmap=None, node_labels=False,
                link_width=1, link_range=[None,None], link_alpha=1, link_cmap=None, link_labels=False,
                add_colorbar=True, directed=False, ax=None, repeat=True):
@@ -766,7 +858,7 @@ def network_animation(wn, node_attribute=None, link_attribute=None, title=None,
     else:
         title_name = '0'
 
-    ax = plot_network(wn, node_attribute=initial_node_values, link_attribute=initial_link_values, title=title_name,
+    ax = plot_network(wn, node_attribute=initial_node_values, link_attribute=initial_link_values, title=title_name, plot_vertices=False,
                node_size=node_size, node_range=node_range, node_alpha=node_alpha, node_cmap=node_cmap, node_labels=node_labels,
                link_width=link_width, link_range=link_range, link_alpha=link_alpha, link_cmap=link_cmap, link_labels=link_labels,
                add_colorbar=add_colorbar, directed=directed, ax=ax)
@@ -790,7 +882,7 @@ def update(n):
         fig.clf()  
         ax = plt.gca()
 
-        ax = plot_network(wn, node_attribute=node_values, link_attribute=link_values, title=title_name,
+        ax = plot_network(wn, node_attribute=node_values, link_attribute=link_values, title=title_name, plot_vertices=False,
                node_size=node_size, node_range=node_range, node_alpha=node_alpha, node_cmap=node_cmap, node_labels=node_labels,
                link_width=link_width, link_range=link_range, link_alpha=link_alpha, link_cmap=link_cmap, link_labels=link_labels,
                add_colorbar=add_colorbar, directed=directed, ax=ax)

wntr/network/io.py

@@ -373,7 +373,7 @@ def from_gis(gis_data, append=None):
 
     return wn
 
-def to_graph(wn, node_weight=None, link_weight=None, modify_direction=False):
+def to_graph(wn, node_weight=None, link_weight=None, modify_direction=False, include_vertices=False):
     """
     Convert a WaterNetworkModel into a networkx MultiDiGraph
 
@@ -410,8 +410,31 @@ def to_graph(wn, node_weight=None, link_weight=None, modify_direction=False):
     for name, link in wn.links():
         start_node = link.start_node_name
         end_node = link.end_node_name
-        G.add_edge(start_node, end_node, key=name)
-        nx.set_edge_attributes(G, name="type", values={(start_node, end_node, name): link.link_type})
+        vertices = link.vertices
+        if include_vertices and vertices:
+            node_base_name = name + "_vertex_"
+            link_base_name = name + "_vertexlink_"
+            for vertex_ind in range(len(vertices)):
+                end_node = node_base_name+str(vertex_ind)
+                link_name = link_base_name+str(vertex_ind)
+                G.add_edge(start_node, end_node, link_name)
+                nx.set_edge_attributes(G, name="type", values={(start_node, end_node, link_name): link.link_type})
+                nx.set_edge_attributes(G, name="base_link", values={(start_node, end_node, link_name): link.name})
+                nx.set_node_attributes(G, name="pos", values={end_node: vertices[vertex_ind]})
+                nx.set_node_attributes(G, name="type", values={end_node: "Vertex"})
+                nx.set_node_attributes(G, name="base_link", values={end_node: link.name})
+
+                start_node = end_node
+            end_node = link.end_node_name
+            link_name = link_base_name+str(vertex_ind+1)
+            G.add_edge(start_node, end_node, link_name)
+            nx.set_edge_attributes(G, name="type", values={(start_node, end_node, link_name): link.link_type})
+            nx.set_edge_attributes(G, name="base_link", values={(start_node, end_node, link_name): link.name})
+
+        else:
+            G.add_edge(start_node, end_node, key=name)
+            nx.set_edge_attributes(G, name="type", values={(start_node, end_node, name): link.link_type})
+            nx.set_edge_attributes(G, name="base_link", values={(start_node, end_node, name): link.name})
 
         if link_weight is not None:
             try:  # weight links

wntr/network/model.py

@@ -1237,7 +1237,7 @@ def from_gis(self, gis_data):
         return wntr.network.io.from_gis(gis_data, append=self)
 
     def to_graph(self, node_weight=None, link_weight=None, 
-                 modify_direction=False):
+                 modify_direction=False, include_vertices=False):
         """
         Convert a WaterNetworkModel into a networkx MultiDiGraph
 
@@ -1258,7 +1258,7 @@ def to_graph(self, node_weight=None, link_weight=None,
         networkx MultiDiGraph
         """
         return wntr.network.io.to_graph(self, node_weight, link_weight, 
-                                        modify_direction)
+                                        modify_direction, include_vertices)
 
 
     def get_graph(self, node_weight=None, link_weight=None, modify_direction=False):

wntr/tests/test_graphics.py

@@ -103,6 +103,23 @@ def test_plot_network5(self):
         plt.close()
 
         self.assertTrue(isfile(filename))
+
+    def test_plot_network_vertices(self):
+        filename = abspath(join(testdir, "plot_network_vertices.png"))
+        if isfile(filename):
+            os.remove(filename)
+
+        inp_file = join(test_datadir, "io.inp")
+        wn = wntr.network.WaterNetworkModel(inp_file)
+
+        wntr.graphics.plot_network(
+            wn, link_attribute="diameter", node_attribute="elevation",
+            plot_vertices=True
+        )
+        plt.savefig(filename, format="png")
+        plt.close()
+
+        self.assertTrue(isfile(filename))
 
     def test_plot_interactive_network1(self):