From 2fe12ed9603be746974f7097f9d1a18884f9f251 Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 11 Jun 2024 12:57:24 -0400 Subject: [PATCH 01/21] first draft of plot_network function with a GIS backend --- wntr/graphics/__init__.py | 2 +- wntr/graphics/network.py | 145 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 146 insertions(+), 1 deletion(-) diff --git a/wntr/graphics/__init__.py b/wntr/graphics/__init__.py index b0d086674..5ea000ee9 100644 --- a/wntr/graphics/__init__.py +++ b/wntr/graphics/__init__.py @@ -1,7 +1,7 @@ """ The wntr.graphics package contains graphic functions """ -from wntr.graphics.network import plot_network, plot_interactive_network, plot_leaflet_network, network_animation +from wntr.graphics.network import plot_network, plot_network_gis, plot_interactive_network, plot_leaflet_network, network_animation from wntr.graphics.layer import plot_valve_layer from wntr.graphics.curve import plot_fragility_curve, plot_pump_curve, plot_tank_volume_curve from wntr.graphics.color import custom_colormap, random_colormap diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index 45b420b13..66b2ff82b 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -3,6 +3,7 @@ water network model. """ import logging +import math import networkx as nx import pandas as pd import matplotlib.pyplot as plt @@ -21,6 +22,15 @@ logger = logging.getLogger(__name__) +def _get_angle(line, loc=0.5): + # calculate orientation angle + p1 = line.interpolate(loc-0.01, normalized=True) + p2 = line.interpolate(loc+0.01, normalized=True) + angle = math.atan2(p2.y-p1.y, p2.x - p1.x) + return angle + +# def _create_oriented_arrow(line, length=0.01) + def _format_node_attribute(node_attribute, wn): if isinstance(node_attribute, str): @@ -42,6 +52,141 @@ def _format_link_attribute(link_attribute, wn): link_attribute = dict(link_attribute) return link_attribute + +def plot_network_gis( + wn, node_attribute=None, link_attribute=None, title=None, + 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', + directed=False, ax=None, filename=None): + + + # # Define node properties + # add_node_colorbar = add_colorbar + # if node_attribute is not None: + + # if isinstance(node_attribute, list): + # if node_cmap is None: + # node_cmap = ['red', 'red'] + # add_node_colorbar = False + + # if node_cmap is None: + # node_cmap = plt.get_cmap('Spectral_r') + # elif isinstance(node_cmap, list): + # if len(node_cmap) == 1: + # node_cmap = node_cmap*2 + # node_cmap = custom_colormap(len(node_cmap), node_cmap) + + # node_attribute = _format_node_attribute(node_attribute, wn) + # nodelist,nodecolor = zip(*node_attribute.items()) + + # else: + # nodelist = None + # nodecolor = 'k' + + # # Define link properties + # add_link_colorbar = add_colorbar + # if link_attribute is not None: + + # if isinstance(link_attribute, list): + # if link_cmap is None: + # link_cmap = ['red', 'red'] + # add_link_colorbar = False + + # if link_cmap is None: + # link_cmap = plt.get_cmap('Spectral_r') + # elif isinstance(link_cmap, list): + # if len(link_cmap) == 1: + # link_cmap = link_cmap*2 + # link_cmap = custom_colormap(len(link_cmap), link_cmap) + + # link_attribute = _format_link_attribute(link_attribute, wn) + + # # Replace link_attribute dictionary defined as + # # {link_name: attr} with {(start_node, end_node, link_name): attr} + # 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 + # link_attribute = attr + + # linklist,linkcolor = zip(*link_attribute.items()) + # else: + # linklist = None + # linkcolor = 'k' + + + if ax is None: # create a new figure + plt.figure(facecolor='w', edgecolor='k') + ax = plt.gca() + + if title is not None: + ax.set_title(title) + + # set aspect setting + aspect = None + + wn_gis = wn.to_gis() + + + # link preprocessing + pipes_kwds = {} + if link_attribute is not None: + pipes_kwds["column"] = link_attribute + pipes_kwds["cmap"] = "Spectral_r" + pipes_kwds["legend"] = True + else: + pipes_kwds["color"] = "black" + + + # node preprocessing + if node_attribute is not None: + node_color = node_attribute + else: + node_color = "black" + + # plot pipes + wn_gis.pipes.plot(ax=ax, aspect=aspect, zorder=0, linewidth=link_width, **pipes_kwds) + + # plot pumps + if len(wn_gis.pumps) >0: + wn_gis.pumps.plot(ax=ax, color="purple", aspect=aspect) + wn_gis.pumps["midpoint"] = wn_gis.pumps.geometry.interpolate(0.5, normalized=True) + wn_gis.pumps["angle"] = wn_gis.pumps.apply(lambda row: _get_angle(row.geometry), axis=1) + # valve_midpoints.plot(ax=ax, marker=">", aspect=aspect) + for idx , row in wn_gis.pumps.iterrows(): + x,y = row["midpoint"].x, row["midpoint"].y + # dx = math.cos(math.radians(row["angle"])) + # dy = math.sin(math.radians(row["angle"])) + angle = row["angle"] + ax.scatter(x,y, color="purple", s=100, marker=(3,0, angle-90)) + # ax.arrow(x,y, dx*0.1, dy*0.1, head_width=0.05, head_length=0.1, fc="blue", ec="blue") + + # plot valves + if len(wn_gis.valves) >0: + wn_gis.valves.plot(ax=ax, color="green", aspect=aspect) + wn_gis.valves["midpoint"] = wn_gis.valves.geometry.interpolate(0.5, normalized=True) + wn_gis.valves["angle"] = wn_gis.valves.apply(lambda row: _get_angle(row.geometry), axis=1) + # valve_midpoints.plot(ax=ax, marker=">", aspect=aspect) + for idx , row in wn_gis.valves.iterrows(): + x,y = row["midpoint"].x, row["midpoint"].y + # dx = math.cos(math.radians(row["angle"])) + # dy = math.sin(math.radians(row["angle"])) + angle = row["angle"] + ax.scatter(x,y, color="green", s=100, marker=(3,0, angle-90)) + # ax.arrow(x,y, dx*0.1, dy*0.1, head_width=0.05, head_length=0.1, fc="blue", ec="blue") + + # plot junctions + wn_gis.junctions.plot(ax=ax, aspect=aspect, color=node_color, markersize=node_size, zorder=1) + + # plot tanks + wn_gis.tanks.plot(ax=ax, marker="P", aspect=aspect, zorder=1) + + # plot reservoirs + wn_gis.reservoirs.plot(ax=ax, marker="s", aspect=aspect, zorder=1) + + return ax + def plot_network(wn, node_attribute=None, link_attribute=None, title=None, node_size=20, node_range=[None,None], node_alpha=1, node_cmap=None, node_labels=False, From 48ac43e2c75b201dfdeddda38b3117b194beb96c Mon Sep 17 00:00:00 2001 From: kbonney Date: Mon, 30 Sep 2024 15:57:42 -0400 Subject: [PATCH 02/21] incorporating the rest of the keywords from plot_network in plot_network_gis --- wntr/graphics/network.py | 213 +++++++++++++++++++++++++++------------ 1 file changed, 147 insertions(+), 66 deletions(-) diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index 66b2ff82b..b29794bed 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -55,67 +55,96 @@ def _format_link_attribute(link_attribute, wn): def plot_network_gis( wn, node_attribute=None, link_attribute=None, title=None, - 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, + node_size=20, node_range=None, node_alpha=1, node_cmap=None, node_labels=False, + link_width=1, link_range=None, link_alpha=1, link_cmap=None, link_labels=False, add_colorbar=True, node_colorbar_label='Node', link_colorbar_label='Link', directed=False, ax=None, filename=None): + """ + Plot network graphic + + Parameters + ---------- + wn : wntr WaterNetworkModel + A WaterNetworkModel object + + node_attribute : None, str, list, pd.Series, or dict, optional + + - If node_attribute is a string, then a node attribute dictionary is + created using node_attribute = wn.query_node_attribute(str) + - If node_attribute is a list, then each node in the list is given a + value of 1. + - If node_attribute is a pd.Series, then it should be in the format + {nodeid: x} where nodeid is a string and x is a float. + - If node_attribute is a dict, then it should be in the format + {nodeid: x} where nodeid is a string and x is a float + link_attribute : None, str, list, pd.Series, or dict, optional + + - If link_attribute is a string, then a link attribute dictionary is + created using edge_attribute = wn.query_link_attribute(str) + - If link_attribute is a list, then each link in the list is given a + value of 1. + - If link_attribute is a pd.Series, then it should be in the format + {linkid: x} where linkid is a string and x is a float. + - If link_attribute is a dict, then it should be in the format + {linkid: x} where linkid is a string and x is a float. + + title: str, optional + Plot title + + node_size: int, optional + Node size + + node_range: list, optional + Node color range ([None,None] indicates autoscale) + + node_alpha: int, optional + Node transparency + + node_cmap: matplotlib.pyplot.cm colormap or list of named colors, optional + Node colormap + + node_labels: bool, optional + If True, the graph will include each node labelled with its name. + + link_width: int, optional + Link width + + link_range : list, optional + Link color range ([None,None] indicates autoscale) + + link_alpha : int, optional + Link transparency - # # Define node properties - # add_node_colorbar = add_colorbar - # if node_attribute is not None: + link_cmap: matplotlib.pyplot.cm colormap or list of named colors, optional + Link colormap - # if isinstance(node_attribute, list): - # if node_cmap is None: - # node_cmap = ['red', 'red'] - # add_node_colorbar = False + link_labels: bool, optional + If True, the graph will include each link labelled with its name. - # if node_cmap is None: - # node_cmap = plt.get_cmap('Spectral_r') - # elif isinstance(node_cmap, list): - # if len(node_cmap) == 1: - # node_cmap = node_cmap*2 - # node_cmap = custom_colormap(len(node_cmap), node_cmap) - - # node_attribute = _format_node_attribute(node_attribute, wn) - # nodelist,nodecolor = zip(*node_attribute.items()) - - # else: - # nodelist = None - # nodecolor = 'k' - - # # Define link properties - # add_link_colorbar = add_colorbar - # if link_attribute is not None: - - # if isinstance(link_attribute, list): - # if link_cmap is None: - # link_cmap = ['red', 'red'] - # add_link_colorbar = False + add_colorbar: bool, optional + Add colorbar - # if link_cmap is None: - # link_cmap = plt.get_cmap('Spectral_r') - # elif isinstance(link_cmap, list): - # if len(link_cmap) == 1: - # link_cmap = link_cmap*2 - # link_cmap = custom_colormap(len(link_cmap), link_cmap) - - # link_attribute = _format_link_attribute(link_attribute, wn) + node_colorbar_label: str, optional + Node colorbar label - # # Replace link_attribute dictionary defined as - # # {link_name: attr} with {(start_node, end_node, link_name): attr} - # 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 - # link_attribute = attr + link_colorbar_label: str, optional + Link colorbar label - # linklist,linkcolor = zip(*link_attribute.items()) - # else: - # linklist = None - # linkcolor = 'k' + directed: bool, optional + If True, plot the directed graph + ax: matplotlib axes object, optional + Axes for plotting (None indicates that a new figure with a single + axes will be used) + filename : str, optional + Filename used to save the figure + + Returns + ------- + ax : matplotlib axes object + """ if ax is None: # create a new figure plt.figure(facecolor='w', edgecolor='k') ax = plt.gca() @@ -125,28 +154,74 @@ def plot_network_gis( # set aspect setting aspect = None - + wn_gis = wn.to_gis() + # colormap + if link_cmap is None: + link_cmap = plt.get_cmap('Spectral_r') + if node_cmap is None: + node_cmap = plt.get_cmap('Spectral_r') + + # ranges + if link_range is None: + link_range = (None, None) + if node_range is None: + node_range = (None, None) + - # link preprocessing + # prepare pipe plotting keywords pipes_kwds = {} if link_attribute is not None: pipes_kwds["column"] = link_attribute - pipes_kwds["cmap"] = "Spectral_r" - pipes_kwds["legend"] = True + pipes_kwds["cmap"] = link_cmap + if add_colorbar: + pipes_kwds["legend"] = True else: pipes_kwds["color"] = "black" + pipes_kwds["alpha"] = link_alpha + + pipes_cbar_kwds = {} + pipes_cbar_kwds["shrink"] = 0.5 + pipes_cbar_kwds["pad"] = 0.0 + pipes_cbar_kwds["label"] = link_colorbar_label - - # node preprocessing + # prepare junctin plotting keywords + junction_kwds = {} if node_attribute is not None: - node_color = node_attribute + junction_kwds["column"] = node_attribute + junction_kwds["cmap"] = node_cmap + if add_colorbar: + junction_kwds["legend"] = True else: - node_color = "black" + junction_kwds["color"] = "black" + junction_kwds["alpha"] = node_alpha + + junction_cbar_kwds = {} + junction_cbar_kwds["shrink"] = 0.5 + junction_cbar_kwds["pad"] = 0.0 + junction_cbar_kwds["label"] = node_colorbar_label + + # TODO handle node/link labels + + # colorbar kwds + + + # plot junctions + wn_gis.junctions.plot( + ax=ax, aspect=aspect, markersize=node_size, zorder=1, + vmax=node_range[0], vmin=node_range[1],legend_kwds=junction_cbar_kwds, **junction_kwds) + + # plot tanks + wn_gis.tanks.plot(ax=ax, marker="P", aspect=aspect, zorder=1) + + # plot reservoirs + wn_gis.reservoirs.plot(ax=ax, marker="s", aspect=aspect, zorder=1) # plot pipes - wn_gis.pipes.plot(ax=ax, aspect=aspect, zorder=0, linewidth=link_width, **pipes_kwds) + wn_gis.pipes.plot( + ax=ax, aspect=aspect, zorder=0, linewidth=link_width, + vmax=link_range[0], vmin=link_range[1], legend_kwds=pipes_cbar_kwds, **pipes_kwds) # plot pumps if len(wn_gis.pumps) >0: @@ -175,15 +250,21 @@ def plot_network_gis( angle = row["angle"] ax.scatter(x,y, color="green", s=100, marker=(3,0, angle-90)) # ax.arrow(x,y, dx*0.1, dy*0.1, head_width=0.05, head_length=0.1, fc="blue", ec="blue") - - # plot junctions - wn_gis.junctions.plot(ax=ax, aspect=aspect, color=node_color, markersize=node_size, zorder=1) - # plot tanks - wn_gis.tanks.plot(ax=ax, marker="P", aspect=aspect, zorder=1) + # annotation + if node_labels: + for x, y, label in zip(wn_gis.junctions.geometry.x, wn_gis.junctions.geometry.y, wn_gis.junctions.index): + ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") + if link_labels: + # compute midpoints + midpoints = wn_gis.pipes.geometry.apply(lambda x: x.interpolate(0.5, normalized=True)) + for x, y, label in zip(midpoints.geometry.x, midpoints.geometry.y, wn_gis.pipes.index): + ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") - # plot reservoirs - wn_gis.reservoirs.plot(ax=ax, marker="s", aspect=aspect, zorder=1) + ax.axis('off') + + if filename: + plt.savefig(filename) return ax From 3baa93a44dd678459bee5eaa6d90e24a531995d8 Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 1 Oct 2024 12:43:27 -0400 Subject: [PATCH 03/21] combining link-like and node-like gis files --- wntr/graphics/network.py | 125 +++++++++++++++++++++------------------ 1 file changed, 67 insertions(+), 58 deletions(-) diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index 10a966185..42282d10c 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -58,7 +58,7 @@ def plot_network_gis( node_size=20, node_range=None, node_alpha=1, node_cmap=None, node_labels=False, link_width=1, link_range=None, link_alpha=1, link_cmap=None, link_labels=False, add_colorbar=True, node_colorbar_label='Node', link_colorbar_label='Link', - directed=False, ax=None, filename=None): + directed=False, ax=None, show_plot=True, filename=None): """ Plot network graphic @@ -137,6 +137,9 @@ def plot_network_gis( ax: matplotlib axes object, optional Axes for plotting (None indicates that a new figure with a single axes will be used) + + show_plot: bool, optional + If True, show plot with plt.show() filename : str, optional Filename used to save the figure @@ -153,10 +156,22 @@ def plot_network_gis( ax.set_title(title) # set aspect setting - aspect = None + # aspect = None + aspect = "auto" + # aspect = "equal" + # initialize gis objects wn_gis = wn.to_gis() + link_gdf = pd.concat((wn_gis.pipes, wn_gis.pumps, wn_gis.valves)) + + node_gdf = pd.concat((wn_gis.junctions, wn_gis.tanks, wn_gis.reservoirs)) + + # missing keyword args + # these are used for elements that do not have a value for the link_attribute + missing_kwds = {"color": "black"} + + # colormap if link_cmap is None: link_cmap = plt.get_cmap('Spectral_r') @@ -171,36 +186,36 @@ def plot_network_gis( # prepare pipe plotting keywords - pipes_kwds = {} + link_kwds = {} if link_attribute is not None: - pipes_kwds["column"] = link_attribute - pipes_kwds["cmap"] = link_cmap + link_kwds["column"] = link_attribute + link_kwds["cmap"] = link_cmap if add_colorbar: - pipes_kwds["legend"] = True + link_kwds["legend"] = True else: - pipes_kwds["color"] = "black" - pipes_kwds["alpha"] = link_alpha + link_kwds["color"] = "black" + link_kwds["alpha"] = link_alpha - pipes_cbar_kwds = {} - pipes_cbar_kwds["shrink"] = 0.5 - pipes_cbar_kwds["pad"] = 0.0 - pipes_cbar_kwds["label"] = link_colorbar_label + link_cbar_kwds = {} + link_cbar_kwds["shrink"] = 0.5 + link_cbar_kwds["pad"] = 0.0 + link_cbar_kwds["label"] = link_colorbar_label # prepare junctin plotting keywords - junction_kwds = {} + node_kwds = {} if node_attribute is not None: - junction_kwds["column"] = node_attribute - junction_kwds["cmap"] = node_cmap + node_kwds["column"] = node_attribute + node_kwds["cmap"] = node_cmap if add_colorbar: - junction_kwds["legend"] = True + node_kwds["legend"] = True else: - junction_kwds["color"] = "black" - junction_kwds["alpha"] = node_alpha + node_kwds["color"] = "black" + node_kwds["alpha"] = node_alpha - junction_cbar_kwds = {} - junction_cbar_kwds["shrink"] = 0.5 - junction_cbar_kwds["pad"] = 0.0 - junction_cbar_kwds["label"] = node_colorbar_label + node_cbar_kwds = {} + node_cbar_kwds["shrink"] = 0.5 + node_cbar_kwds["pad"] = 0.0 + node_cbar_kwds["label"] = node_colorbar_label # TODO handle node/link labels @@ -208,48 +223,42 @@ def plot_network_gis( # plot junctions - wn_gis.junctions.plot( - ax=ax, aspect=aspect, markersize=node_size, zorder=1, - vmax=node_range[0], vmin=node_range[1],legend_kwds=junction_cbar_kwds, **junction_kwds) + # node_gdf.plot( + # ax=ax, aspect=aspect, markersize=node_size, zorder=1, + # vmax=node_range[0], vmin=node_range[1], legend_kwds=node_cbar_kwds, **node_kwds) - # plot tanks - wn_gis.tanks.plot(ax=ax, marker="P", aspect=aspect, zorder=1) + # # plot tanks + # wn_gis.tanks.plot(ax=ax, marker="P", aspect=aspect, zorder=1) - # plot reservoirs - wn_gis.reservoirs.plot(ax=ax, marker="s", aspect=aspect, zorder=1) + # # plot reservoirs + # wn_gis.reservoirs.plot(ax=ax, marker="s", aspect=aspect, zorder=1) # plot pipes - wn_gis.pipes.plot( + minx, miny, maxx, maxy = link_gdf.total_bounds + link_gdf.plot( ax=ax, aspect=aspect, zorder=0, linewidth=link_width, - vmax=link_range[0], vmin=link_range[1], legend_kwds=pipes_cbar_kwds, **pipes_kwds) - - # plot pumps - if len(wn_gis.pumps) >0: - wn_gis.pumps.plot(ax=ax, color="purple", aspect=aspect) - wn_gis.pumps["midpoint"] = wn_gis.pumps.geometry.interpolate(0.5, normalized=True) - wn_gis.pumps["angle"] = wn_gis.pumps.apply(lambda row: _get_angle(row.geometry), axis=1) - # valve_midpoints.plot(ax=ax, marker=">", aspect=aspect) - for idx , row in wn_gis.pumps.iterrows(): - x,y = row["midpoint"].x, row["midpoint"].y - # dx = math.cos(math.radians(row["angle"])) - # dy = math.sin(math.radians(row["angle"])) - angle = row["angle"] - ax.scatter(x,y, color="purple", s=100, marker=(3,0, angle-90)) - # ax.arrow(x,y, dx*0.1, dy*0.1, head_width=0.05, head_length=0.1, fc="blue", ec="blue") + vmax=link_range[0], vmin=link_range[1], missing_kwds=missing_kwds, legend_kwds=link_cbar_kwds, **link_kwds) + ax.set_xlim([minx, maxx]) + ax.set_ylim([miny, maxy]) + # # plot pumps + # if len(wn_gis.pumps) >0: + # wn_gis.pumps.plot(ax=ax, color="purple", aspect=aspect) + # wn_gis.pumps["midpoint"] = wn_gis.pumps.geometry.interpolate(0.5, normalized=True) + # wn_gis.pumps["angle"] = wn_gis.pumps.apply(lambda row: _get_angle(row.geometry), axis=1) + # for idx , row in wn_gis.pumps.iterrows(): + # x,y = row["midpoint"].x, row["midpoint"].y + # angle = row["angle"] + # ax.scatter(x,y, color="purple", s=100, marker=(3,0, angle-90)) # plot valves - if len(wn_gis.valves) >0: - wn_gis.valves.plot(ax=ax, color="green", aspect=aspect) - wn_gis.valves["midpoint"] = wn_gis.valves.geometry.interpolate(0.5, normalized=True) - wn_gis.valves["angle"] = wn_gis.valves.apply(lambda row: _get_angle(row.geometry), axis=1) - # valve_midpoints.plot(ax=ax, marker=">", aspect=aspect) - for idx , row in wn_gis.valves.iterrows(): - x,y = row["midpoint"].x, row["midpoint"].y - # dx = math.cos(math.radians(row["angle"])) - # dy = math.sin(math.radians(row["angle"])) - angle = row["angle"] - ax.scatter(x,y, color="green", s=100, marker=(3,0, angle-90)) - # ax.arrow(x,y, dx*0.1, dy*0.1, head_width=0.05, head_length=0.1, fc="blue", ec="blue") + # if len(wn_gis.valves) >0: + # # wn_gis.valves.plot(ax=ax, color="green", aspect=aspect) + # wn_gis.valves["midpoint"] = wn_gis.valves.geometry.interpolate(0.5, normalized=True) + # wn_gis.valves["angle"] = wn_gis.valves.apply(lambda row: _get_angle(row.geometry), axis=1) + # for idx , row in wn_gis.valves.iterrows(): + # x,y = row["midpoint"].x, row["midpoint"].y + # angle = row["angle"] + # ax.scatter(x,y, color="green", s=100, marker=(3,0, angle-90)) # annotation if node_labels: @@ -261,7 +270,7 @@ def plot_network_gis( for x, y, label in zip(midpoints.geometry.x, midpoints.geometry.y, wn_gis.pipes.index): ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") - ax.axis('off') + # ax.axis('off') if filename: plt.savefig(filename) From 487511afa21f237cdcd62f8886bda8662175589e Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 1 Oct 2024 18:30:28 -0400 Subject: [PATCH 04/21] adding tank/reservoir marker shapes and implementating other type options for attributes --- wntr/graphics/network.py | 104 +++++++++++++++++++++++++++++---------- 1 file changed, 79 insertions(+), 25 deletions(-) diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index 42282d10c..37529d3a7 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -57,7 +57,7 @@ def plot_network_gis( wn, node_attribute=None, link_attribute=None, title=None, node_size=20, node_range=None, node_alpha=1, node_cmap=None, node_labels=False, link_width=1, link_range=None, link_alpha=1, link_cmap=None, link_labels=False, - add_colorbar=True, node_colorbar_label='Node', link_colorbar_label='Link', + add_colorbar=True, node_colorbar_label="", link_colorbar_label="", directed=False, ax=None, show_plot=True, filename=None): """ Plot network graphic @@ -156,8 +156,8 @@ def plot_network_gis( ax.set_title(title) # set aspect setting - # aspect = None - aspect = "auto" + aspect = None + # aspect = "auto" # aspect = "equal" # initialize gis objects @@ -169,8 +169,11 @@ def plot_network_gis( # missing keyword args # these are used for elements that do not have a value for the link_attribute - missing_kwds = {"color": "black"} + # missing_kwds = {"color": "black"} + # set tank and reservoir marker + tank_marker = "P" + reservoir_marker = "s" # colormap if link_cmap is None: @@ -188,14 +191,41 @@ def plot_network_gis( # prepare pipe plotting keywords link_kwds = {} if link_attribute is not None: - link_kwds["column"] = link_attribute + # if dict convert to a series + if isinstance(link_attribute, dict): + link_attribute = pd.Series(link_attribute) + # if series add as a column to link gdf + if isinstance(link_attribute, pd.Series): + link_gdf["_link_attribute"] = link_attribute + link_kwds["column"] = "_link_attribute" + # if list, create new boolean column that captures which indices are in the list + # TODO need to check this with original behavior + elif isinstance(link_attribute, list): + link_gdf["_link_attribute"] = link_gdf.index.isin(link_attribute).astype(int) + link_kwds["column"] = "_link_attribute" + # if str, assert that column name exists + elif isinstance(link_attribute, str): + if link_attribute not in link_gdf.columns: + raise KeyError(f"link_attribute {link_attribute} does not exist.") + link_kwds["column"] = link_attribute + else: + raise TypeError("link_attribute must be dict, Series, list, or str") link_kwds["cmap"] = link_cmap if add_colorbar: link_kwds["legend"] = True + link_kwds["vmin"] = link_range[0] + link_kwds["vmax"] = link_range[1] else: link_kwds["color"] = "black" + + link_kwds["linewidth"] = link_width link_kwds["alpha"] = link_alpha + background_link_kwds = {} + background_link_kwds["color"] = "grey" + background_link_kwds["linewidth"] = link_width + background_link_kwds["alpha"] = link_alpha + link_cbar_kwds = {} link_cbar_kwds["shrink"] = 0.5 link_cbar_kwds["pad"] = 0.0 @@ -204,43 +234,67 @@ def plot_network_gis( # prepare junctin plotting keywords node_kwds = {} if node_attribute is not None: - node_kwds["column"] = node_attribute + # if dict convert to a series + if isinstance(node_attribute, dict): + node_attribute = pd.Series(node_attribute) + # if series add as a column to node gdf + if isinstance(node_attribute, pd.Series): + node_gdf["_node_attribute"] = node_attribute + node_kwds["column"] = "_node_attribute" + # if list, create new boolean column that captures which indices are in the list + # TODO need to check this with original behavior + elif isinstance(node_attribute, list): + node_gdf["_node_attribute"] = node_gdf.index.isin(node_attribute).astype(int) + node_kwds["column"] = "_node_attribute" + # if str, assert that column name exists + elif isinstance(node_attribute, str): + if node_attribute not in node_gdf.columns: + raise KeyError(f"node_attribute {node_attribute} does not exist.") + node_kwds["column"] = node_attribute + else: + raise TypeError("node_attribute must be dict, Series, list, or str") node_kwds["cmap"] = node_cmap if add_colorbar: node_kwds["legend"] = True + node_kwds["vmin"] = node_range[0] + node_kwds["vmax"] = node_range[1] else: node_kwds["color"] = "black" node_kwds["alpha"] = node_alpha + node_kwds["markersize"] = node_size node_cbar_kwds = {} node_cbar_kwds["shrink"] = 0.5 node_cbar_kwds["pad"] = 0.0 node_cbar_kwds["label"] = node_colorbar_label - # TODO handle node/link labels - - # colorbar kwds - - # plot junctions - # node_gdf.plot( - # ax=ax, aspect=aspect, markersize=node_size, zorder=1, - # vmax=node_range[0], vmin=node_range[1], legend_kwds=node_cbar_kwds, **node_kwds) + # junction_mask + node_gdf[node_gdf.node_type == "Junction"].plot( + ax=ax, aspect=aspect, zorder=3, legend_kwds=node_cbar_kwds, **node_kwds) + + # turn off legend for subsequent node plots + node_kwds["legend"] = False - # # plot tanks - # wn_gis.tanks.plot(ax=ax, marker="P", aspect=aspect, zorder=1) + # plot tanks + node_kwds["markersize"] = node_size * 1.5 + node_gdf[node_gdf.node_type == "Tank"].plot( + ax=ax, aspect=aspect, zorder=4, marker=tank_marker, **node_kwds) - # # plot reservoirs - # wn_gis.reservoirs.plot(ax=ax, marker="s", aspect=aspect, zorder=1) + # plot reservoirs + node_kwds["markersize"] = node_size * 2.0 + node_gdf[node_gdf.node_type == "Reservoir"].plot( + ax=ax, aspect=aspect, zorder=5, marker=reservoir_marker, **node_kwds) # plot pipes - minx, miny, maxx, maxy = link_gdf.total_bounds + # background + link_gdf.plot( + ax=ax, aspect=aspect, zorder=1, **background_link_kwds) + link_gdf.plot( - ax=ax, aspect=aspect, zorder=0, linewidth=link_width, - vmax=link_range[0], vmin=link_range[1], missing_kwds=missing_kwds, legend_kwds=link_cbar_kwds, **link_kwds) - ax.set_xlim([minx, maxx]) - ax.set_ylim([miny, maxy]) - # # plot pumps + ax=ax, aspect=aspect, zorder=2, legend_kwds=link_cbar_kwds, **link_kwds) + + # plot pumps # if len(wn_gis.pumps) >0: # wn_gis.pumps.plot(ax=ax, color="purple", aspect=aspect) # wn_gis.pumps["midpoint"] = wn_gis.pumps.geometry.interpolate(0.5, normalized=True) @@ -270,7 +324,7 @@ def plot_network_gis( for x, y, label in zip(midpoints.geometry.x, midpoints.geometry.y, wn_gis.pipes.index): ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") - # ax.axis('off') + ax.axis('off') if filename: plt.savefig(filename) From 9bc07cfa255d25f0cdead32d874bb69c4539e56e Mon Sep 17 00:00:00 2001 From: kbonney Date: Mon, 21 Oct 2024 09:52:42 -0400 Subject: [PATCH 05/21] set geopandas backend to the plot_network function name. update the geopandas backend to handle lists for attributes and add shapes for valves and pumps. --- wntr/graphics/__init__.py | 2 +- wntr/graphics/network.py | 182 +++++++++++++++++++------------------- 2 files changed, 94 insertions(+), 90 deletions(-) diff --git a/wntr/graphics/__init__.py b/wntr/graphics/__init__.py index 5ea000ee9..7238bd7ea 100644 --- a/wntr/graphics/__init__.py +++ b/wntr/graphics/__init__.py @@ -1,7 +1,7 @@ """ The wntr.graphics package contains graphic functions """ -from wntr.graphics.network import plot_network, plot_network_gis, plot_interactive_network, plot_leaflet_network, network_animation +from wntr.graphics.network import plot_network, plot_network_nx, plot_interactive_network, plot_leaflet_network, network_animation from wntr.graphics.layer import plot_valve_layer from wntr.graphics.curve import plot_fragility_curve, plot_pump_curve, plot_tank_volume_curve from wntr.graphics.color import custom_colormap, random_colormap diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index 37529d3a7..c8d12ce78 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -7,7 +7,9 @@ import networkx as nx import pandas as pd import matplotlib.pyplot as plt +import matplotlib.path as mpath from matplotlib import animation +import numpy as np try: import plotly @@ -22,13 +24,56 @@ logger = logging.getLogger(__name__) + +arrow_verts = [ + (0.0, 0.0), + (0.5, 0.5), + (0.5, -0.5), + (0.0, 0.0), +] + +arrow_marker = mpath.Path(arrow_verts) + def _get_angle(line, loc=0.5): # calculate orientation angle p1 = line.interpolate(loc-0.01, normalized=True) p2 = line.interpolate(loc+0.01, normalized=True) - angle = math.atan2(p2.y-p1.y, p2.x - p1.x) + angle = math.atan2(p2.y-p1.y, p2.x - p1.x) # radians + angle = math.degrees(angle) return angle + +def _prepare_attribute(attribute, gdf): + kwds = {} + if attribute is not None: + # if dict convert to a series + if isinstance(attribute, dict): + attribute = pd.Series(attribute) + # if series add as a column to link gdf + if isinstance(attribute, pd.Series): + gdf["_attribute"] = attribute + kwds["column"] = "_attribute" + # if list, create new boolean column that captures which indices are in the list + # TODO need to check this with original behavior + elif isinstance(attribute, list): + gdf["_attribute"] = np.nan + gdf.loc[gdf.index.isin(attribute), "_attribute"] = 1 + kwds["column"] = "_attribute" + # if str, assert that column name exists + elif isinstance(attribute, str): + if attribute not in gdf.columns: + raise KeyError(f"attribute {attribute} does not exist.") + kwds["column"] = attribute + else: + raise TypeError("attribute must be dict, Series, list, or str") + else: + kwds["color"] = "black" + return kwds + + + + + # def _create_oriented_arrow(line, length=0.01) def _format_node_attribute(node_attribute, wn): @@ -53,7 +98,7 @@ def _format_link_attribute(link_attribute, wn): return link_attribute -def plot_network_gis( +def plot_network( wn, node_attribute=None, link_attribute=None, title=None, node_size=20, node_range=None, node_alpha=1, node_cmap=None, node_labels=False, link_width=1, link_range=None, link_alpha=1, link_cmap=None, link_labels=False, @@ -155,66 +200,37 @@ def plot_network_gis( if title is not None: ax.set_title(title) - # set aspect setting aspect = None - # aspect = "auto" - # aspect = "equal" - - # initialize gis objects - wn_gis = wn.to_gis() - - link_gdf = pd.concat((wn_gis.pipes, wn_gis.pumps, wn_gis.valves)) - node_gdf = pd.concat((wn_gis.junctions, wn_gis.tanks, wn_gis.reservoirs)) - - # missing keyword args - # these are used for elements that do not have a value for the link_attribute - # missing_kwds = {"color": "black"} - - # set tank and reservoir marker - tank_marker = "P" + tank_marker = "D" reservoir_marker = "s" - # colormap if link_cmap is None: link_cmap = plt.get_cmap('Spectral_r') if node_cmap is None: node_cmap = plt.get_cmap('Spectral_r') - # ranges if link_range is None: link_range = (None, None) if node_range is None: node_range = (None, None) + + wn_gis = wn.to_gis() + link_gdf = pd.concat((wn_gis.pipes, wn_gis.pumps, wn_gis.valves)) + node_gdf = pd.concat((wn_gis.junctions, wn_gis.tanks, wn_gis.reservoirs)) + # process link attribute + link_kwds = _prepare_attribute(link_attribute, link_gdf) - # prepare pipe plotting keywords - link_kwds = {} - if link_attribute is not None: - # if dict convert to a series - if isinstance(link_attribute, dict): - link_attribute = pd.Series(link_attribute) - # if series add as a column to link gdf - if isinstance(link_attribute, pd.Series): - link_gdf["_link_attribute"] = link_attribute - link_kwds["column"] = "_link_attribute" - # if list, create new boolean column that captures which indices are in the list - # TODO need to check this with original behavior - elif isinstance(link_attribute, list): - link_gdf["_link_attribute"] = link_gdf.index.isin(link_attribute).astype(int) - link_kwds["column"] = "_link_attribute" - # if str, assert that column name exists - elif isinstance(link_attribute, str): - if link_attribute not in link_gdf.columns: - raise KeyError(f"link_attribute {link_attribute} does not exist.") - link_kwds["column"] = link_attribute - else: - raise TypeError("link_attribute must be dict, Series, list, or str") + if isinstance(link_attribute, list): + link_kwds["column"] = "_attribute" + link_kwds["cmap"] = custom_colormap(2,("red", "red")) + link_kwds["legend"] = False + elif isinstance(link_attribute, (dict, pd.Series, str)): link_kwds["cmap"] = link_cmap - if add_colorbar: - link_kwds["legend"] = True link_kwds["vmin"] = link_range[0] link_kwds["vmax"] = link_range[1] + link_kwds["legend"] = add_colorbar else: link_kwds["color"] = "black" @@ -230,36 +246,21 @@ def plot_network_gis( link_cbar_kwds["shrink"] = 0.5 link_cbar_kwds["pad"] = 0.0 link_cbar_kwds["label"] = link_colorbar_label - - # prepare junctin plotting keywords - node_kwds = {} - if node_attribute is not None: - # if dict convert to a series - if isinstance(node_attribute, dict): - node_attribute = pd.Series(node_attribute) - # if series add as a column to node gdf - if isinstance(node_attribute, pd.Series): - node_gdf["_node_attribute"] = node_attribute - node_kwds["column"] = "_node_attribute" - # if list, create new boolean column that captures which indices are in the list - # TODO need to check this with original behavior - elif isinstance(node_attribute, list): - node_gdf["_node_attribute"] = node_gdf.index.isin(node_attribute).astype(int) - node_kwds["column"] = "_node_attribute" - # if str, assert that column name exists - elif isinstance(node_attribute, str): - if node_attribute not in node_gdf.columns: - raise KeyError(f"node_attribute {node_attribute} does not exist.") - node_kwds["column"] = node_attribute - else: - raise TypeError("node_attribute must be dict, Series, list, or str") + + # process node attribute + node_kwds = _prepare_attribute(node_attribute, node_gdf) + + if isinstance(node_attribute, list): + node_kwds["cmap"] = custom_colormap(2,("red", "red")) + node_kwds["legend"] = False + elif isinstance(node_attribute, (dict, pd.Series, str)): node_kwds["cmap"] = node_cmap - if add_colorbar: - node_kwds["legend"] = True node_kwds["vmin"] = node_range[0] node_kwds["vmax"] = node_range[1] + node_kwds["legend"] = add_colorbar else: node_kwds["color"] = "black" + node_kwds["alpha"] = node_alpha node_kwds["markersize"] = node_size @@ -268,8 +269,8 @@ def plot_network_gis( node_cbar_kwds["pad"] = 0.0 node_cbar_kwds["label"] = node_colorbar_label + # plot nodes - each type is plotted separately to allow for different marker types # plot junctions - # junction_mask node_gdf[node_gdf.node_type == "Junction"].plot( ax=ax, aspect=aspect, zorder=3, legend_kwds=node_cbar_kwds, **node_kwds) @@ -277,12 +278,12 @@ def plot_network_gis( node_kwds["legend"] = False # plot tanks - node_kwds["markersize"] = node_size * 1.5 + node_kwds["markersize"] = node_size * 2.0 node_gdf[node_gdf.node_type == "Tank"].plot( ax=ax, aspect=aspect, zorder=4, marker=tank_marker, **node_kwds) # plot reservoirs - node_kwds["markersize"] = node_size * 2.0 + node_kwds["markersize"] = node_size * 3.0 node_gdf[node_gdf.node_type == "Reservoir"].plot( ax=ax, aspect=aspect, zorder=5, marker=reservoir_marker, **node_kwds) @@ -295,24 +296,24 @@ def plot_network_gis( ax=ax, aspect=aspect, zorder=2, legend_kwds=link_cbar_kwds, **link_kwds) # plot pumps - # if len(wn_gis.pumps) >0: - # wn_gis.pumps.plot(ax=ax, color="purple", aspect=aspect) - # wn_gis.pumps["midpoint"] = wn_gis.pumps.geometry.interpolate(0.5, normalized=True) - # wn_gis.pumps["angle"] = wn_gis.pumps.apply(lambda row: _get_angle(row.geometry), axis=1) - # for idx , row in wn_gis.pumps.iterrows(): - # x,y = row["midpoint"].x, row["midpoint"].y - # angle = row["angle"] - # ax.scatter(x,y, color="purple", s=100, marker=(3,0, angle-90)) + if len(wn_gis.pumps) >0: + # wn_gis.pumps.plot(ax=ax, color="purple", aspect=aspect) + wn_gis.pumps["midpoint"] = wn_gis.pumps.geometry.interpolate(0.5, normalized=True) + wn_gis.pumps["angle"] = wn_gis.pumps.apply(lambda row: _get_angle(row.geometry), axis=1) + for idx , row in wn_gis.pumps.iterrows(): + x,y = row["midpoint"].x, row["midpoint"].y + angle = row["angle"] + ax.scatter(x,y, color="black", s=100, marker=(3, 0, angle-90)) + # ax.scatter(x,y, color="purple", s=100, marker=arrow_marker) # plot valves - # if len(wn_gis.valves) >0: - # # wn_gis.valves.plot(ax=ax, color="green", aspect=aspect) - # wn_gis.valves["midpoint"] = wn_gis.valves.geometry.interpolate(0.5, normalized=True) - # wn_gis.valves["angle"] = wn_gis.valves.apply(lambda row: _get_angle(row.geometry), axis=1) - # for idx , row in wn_gis.valves.iterrows(): - # x,y = row["midpoint"].x, row["midpoint"].y - # angle = row["angle"] - # ax.scatter(x,y, color="green", s=100, marker=(3,0, angle-90)) + if len(wn_gis.valves) >0: + wn_gis.valves["midpoint"] = wn_gis.valves.geometry.interpolate(0.5, normalized=True) + wn_gis.valves["angle"] = wn_gis.valves.apply(lambda row: _get_angle(row.geometry), axis=1) + for idx , row in wn_gis.valves.iterrows(): + x,y = row["midpoint"].x, row["midpoint"].y + angle = row["angle"] + ax.scatter(x,y, color="black", s=200, marker=(2,0, angle)) # annotation if node_labels: @@ -329,10 +330,13 @@ def plot_network_gis( if filename: plt.savefig(filename) + if show_plot is True: + plt.show(block=False) + return ax -def plot_network(wn, node_attribute=None, link_attribute=None, title=None, +def plot_network_nx(wn, node_attribute=None, link_attribute=None, title=None, 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', From 3bb3c4b206512a275e402de4ff0809fd68591a6b Mon Sep 17 00:00:00 2001 From: kbonney Date: Mon, 21 Oct 2024 09:56:02 -0400 Subject: [PATCH 06/21] archive networkx plotting function in test suite --- wntr/tests/test_graphics.py | 238 ++++++++++++++++++++++++++++++++++++ 1 file changed, 238 insertions(+) diff --git a/wntr/tests/test_graphics.py b/wntr/tests/test_graphics.py index 0374ebef3..20ce5c055 100644 --- a/wntr/tests/test_graphics.py +++ b/wntr/tests/test_graphics.py @@ -5,7 +5,10 @@ import warnings from os.path import abspath, dirname, isfile, join +import networkx as nx import matplotlib.pylab as plt +from wntr.graphics.color import custom_colormap +import pandas as pd import wntr testdir = dirname(abspath(str(__file__))) @@ -235,6 +238,241 @@ def test_custom_colormap(self): ) self.assertEqual(cmp.N, 3) self.assertEqual(cmp.name, "custom") + + +# old plotting function using networkx backend to compare with geopandas +def plot_network_nx(wn, node_attribute=None, link_attribute=None, title=None, + 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', + directed=False, ax=None, show_plot=True, filename=None): + """ + Plot network graphic + + Parameters + ---------- + wn : wntr WaterNetworkModel + A WaterNetworkModel object + + node_attribute : None, str, list, pd.Series, or dict, optional + + - If node_attribute is a string, then a node attribute dictionary is + created using node_attribute = wn.query_node_attribute(str) + - If node_attribute is a list, then each node in the list is given a + value of 1. + - If node_attribute is a pd.Series, then it should be in the format + {nodeid: x} where nodeid is a string and x is a float. + - If node_attribute is a dict, then it should be in the format + {nodeid: x} where nodeid is a string and x is a float + + link_attribute : None, str, list, pd.Series, or dict, optional + + - If link_attribute is a string, then a link attribute dictionary is + created using edge_attribute = wn.query_link_attribute(str) + - If link_attribute is a list, then each link in the list is given a + value of 1. + - If link_attribute is a pd.Series, then it should be in the format + {linkid: x} where linkid is a string and x is a float. + - If link_attribute is a dict, then it should be in the format + {linkid: x} where linkid is a string and x is a float. + + title: str, optional + Plot title + + node_size: int, optional + Node size + + node_range: list, optional + Node color range ([None,None] indicates autoscale) + + node_alpha: int, optional + Node transparency + + node_cmap: matplotlib.pyplot.cm colormap or list of named colors, optional + Node colormap + + node_labels: bool, optional + If True, the graph will include each node labelled with its name. + + link_width: int, optional + Link width + + link_range : list, optional + Link color range ([None,None] indicates autoscale) + + link_alpha : int, optional + Link transparency + + link_cmap: matplotlib.pyplot.cm colormap or list of named colors, optional + Link colormap + + link_labels: bool, optional + If True, the graph will include each link labelled with its name. + + add_colorbar: bool, optional + Add colorbar + + node_colorbar_label: str, optional + Node colorbar label + + link_colorbar_label: str, optional + Link colorbar label + + directed: bool, optional + If True, plot the directed graph + + ax: matplotlib axes object, optional + Axes for plotting (None indicates that a new figure with a single + axes will be used) + + show_plot: bool, optional + If True, show plot with plt.show() + + filename : str, optional + Filename used to save the figure + + Returns + ------- + ax : matplotlib axes object + """ + + def _format_node_attribute(node_attribute, wn): + + if isinstance(node_attribute, str): + node_attribute = wn.query_node_attribute(node_attribute) + if isinstance(node_attribute, list): + node_attribute = dict(zip(node_attribute,[1]*len(node_attribute))) + if isinstance(node_attribute, pd.Series): + node_attribute = dict(node_attribute) + + return node_attribute + + def _format_link_attribute(link_attribute, wn): + + if isinstance(link_attribute, str): + link_attribute = wn.query_link_attribute(link_attribute) + if isinstance(link_attribute, list): + link_attribute = dict(zip(link_attribute,[1]*len(link_attribute))) + if isinstance(link_attribute, pd.Series): + link_attribute = dict(link_attribute) + + return link_attribute + + if ax is None: # create a new figure + plt.figure(facecolor='w', edgecolor='k') + ax = plt.gca() + + # Graph + G = wn.to_graph() + if not directed: + G = G.to_undirected() + + # Position + pos = nx.get_node_attributes(G,'pos') + if len(pos) == 0: + pos = None + + # Define node properties + add_node_colorbar = add_colorbar + if node_attribute is not None: + + if isinstance(node_attribute, list): + if node_cmap is None: + node_cmap = ['red', 'red'] + add_node_colorbar = False + + if node_cmap is None: + node_cmap = plt.get_cmap('Spectral_r') + elif isinstance(node_cmap, list): + if len(node_cmap) == 1: + node_cmap = node_cmap*2 + node_cmap = custom_colormap(len(node_cmap), node_cmap) + + node_attribute = _format_node_attribute(node_attribute, wn) + nodelist,nodecolor = zip(*node_attribute.items()) + + else: + nodelist = None + nodecolor = 'k' + + add_link_colorbar = add_colorbar + if link_attribute is not None: + + if isinstance(link_attribute, list): + if link_cmap is None: + link_cmap = ['red', 'red'] + add_link_colorbar = False + + if link_cmap is None: + link_cmap = plt.get_cmap('Spectral_r') + elif isinstance(link_cmap, list): + if len(link_cmap) == 1: + link_cmap = link_cmap*2 + link_cmap = custom_colormap(len(link_cmap), link_cmap) + + link_attribute = _format_link_attribute(link_attribute, wn) + + # Replace link_attribute dictionary defined as + # {link_name: attr} with {(start_node, end_node, link_name): attr} + 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 + link_attribute = attr + + linklist,linkcolor = zip(*link_attribute.items()) + else: + linklist = None + linkcolor = 'k' + + if title is not None: + ax.set_title(title) + + edge_background = nx.draw_networkx_edges(G, pos, edge_color='grey', + 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], + linewidths=0, ax=ax) + edges = nx.draw_networkx_edges(G, pos, edgelist=linklist, arrows=directed, + 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) + if link_labels: + labels = {} + for link_name in wn.link_name_list: + link = wn.get_link(link_name) + labels[(link.start_node_name, link.end_node_name)] = link_name + nx.draw_networkx_edge_labels(G, pos, labels, font_size=7, ax=ax) + if add_node_colorbar and node_attribute: + clb = plt.colorbar(nodes, shrink=0.5, pad=0, ax=ax) + clb.ax.set_title(node_colorbar_label, fontsize=10) + if add_link_colorbar and link_attribute: + if link_range[0] is None: + vmin = min(link_attribute.values()) + else: + vmin = link_range[0] + if link_range[1] is None: + vmax = max(link_attribute.values()) + else: + vmax = link_range[1] + sm = plt.cm.ScalarMappable(cmap=link_cmap, norm=plt.Normalize(vmin=vmin, vmax=vmax)) + sm.set_array([]) + clb = plt.colorbar(sm, shrink=0.5, pad=0.05, ax=ax) + clb.ax.set_title(link_colorbar_label, fontsize=10) + + ax.axis('off') + + if filename: + plt.savefig(filename) + + if show_plot is True: + plt.show(block=False) + + return ax if __name__ == "__main__": From 1a7e0a593ca9e56cbb15c2348891e2983479ba86 Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 5 Nov 2024 10:05:22 -0500 Subject: [PATCH 07/21] clean up function and add directed functionality --- wntr/graphics/network.py | 43 +++++++++++++++++----------------------- 1 file changed, 18 insertions(+), 25 deletions(-) diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index c8d12ce78..8ef9b396c 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -54,7 +54,6 @@ def _prepare_attribute(attribute, gdf): gdf["_attribute"] = attribute kwds["column"] = "_attribute" # if list, create new boolean column that captures which indices are in the list - # TODO need to check this with original behavior elif isinstance(attribute, list): gdf["_attribute"] = np.nan gdf.loc[gdf.index.isin(attribute), "_attribute"] = 1 @@ -70,12 +69,6 @@ def _prepare_attribute(attribute, gdf): kwds["color"] = "black" return kwds - - - - -# def _create_oriented_arrow(line, length=0.01) - def _format_node_attribute(node_attribute, wn): if isinstance(node_attribute, str): @@ -270,60 +263,60 @@ def plot_network( node_cbar_kwds["label"] = node_colorbar_label # plot nodes - each type is plotted separately to allow for different marker types - # plot junctions node_gdf[node_gdf.node_type == "Junction"].plot( ax=ax, aspect=aspect, zorder=3, legend_kwds=node_cbar_kwds, **node_kwds) # turn off legend for subsequent node plots node_kwds["legend"] = False - # plot tanks node_kwds["markersize"] = node_size * 2.0 node_gdf[node_gdf.node_type == "Tank"].plot( ax=ax, aspect=aspect, zorder=4, marker=tank_marker, **node_kwds) - # plot reservoirs node_kwds["markersize"] = node_size * 3.0 node_gdf[node_gdf.node_type == "Reservoir"].plot( ax=ax, aspect=aspect, zorder=5, marker=reservoir_marker, **node_kwds) - # plot pipes - # background + # plot links link_gdf.plot( ax=ax, aspect=aspect, zorder=1, **background_link_kwds) link_gdf.plot( ax=ax, aspect=aspect, zorder=2, legend_kwds=link_cbar_kwds, **link_kwds) - # plot pumps if len(wn_gis.pumps) >0: # wn_gis.pumps.plot(ax=ax, color="purple", aspect=aspect) - wn_gis.pumps["midpoint"] = wn_gis.pumps.geometry.interpolate(0.5, normalized=True) - wn_gis.pumps["angle"] = wn_gis.pumps.apply(lambda row: _get_angle(row.geometry), axis=1) + wn_gis.pumps["_midpoint"] = wn_gis.pumps.geometry.interpolate(0.5, normalized=True) + wn_gis.pumps["_angle"] = wn_gis.pumps.apply(lambda row: _get_angle(row.geometry), axis=1) for idx , row in wn_gis.pumps.iterrows(): - x,y = row["midpoint"].x, row["midpoint"].y - angle = row["angle"] + x,y = row["_midpoint"].x, row["_midpoint"].y + angle = row["_angle"] ax.scatter(x,y, color="black", s=100, marker=(3, 0, angle-90)) - # ax.scatter(x,y, color="purple", s=100, marker=arrow_marker) - # plot valves if len(wn_gis.valves) >0: - wn_gis.valves["midpoint"] = wn_gis.valves.geometry.interpolate(0.5, normalized=True) - wn_gis.valves["angle"] = wn_gis.valves.apply(lambda row: _get_angle(row.geometry), axis=1) + wn_gis.valves["_midpoint"] = wn_gis.valves.geometry.interpolate(0.5, normalized=True) + wn_gis.valves["_angle"] = wn_gis.valves.apply(lambda row: _get_angle(row.geometry), axis=1) for idx , row in wn_gis.valves.iterrows(): - x,y = row["midpoint"].x, row["midpoint"].y - angle = row["angle"] + x,y = row["_midpoint"].x, row["_midpoint"].y + angle = row["_angle"] ax.scatter(x,y, color="black", s=200, marker=(2,0, angle)) - # annotation if node_labels: for x, y, label in zip(wn_gis.junctions.geometry.x, wn_gis.junctions.geometry.y, wn_gis.junctions.index): ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") + if link_labels: - # compute midpoints midpoints = wn_gis.pipes.geometry.apply(lambda x: x.interpolate(0.5, normalized=True)) for x, y, label in zip(midpoints.geometry.x, midpoints.geometry.y, wn_gis.pipes.index): ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") + + if directed: + link_gdf["_midpoint"] = link_gdf.geometry.interpolate(0.5, normalized=True) + link_gdf["_angle"] = link_gdf.apply(lambda row: _get_angle(row.geometry), axis=1) + for idx , row in link_gdf.iterrows(): + x,y = row["_midpoint"].x, row["_midpoint"].y + angle = row["_angle"] + ax.scatter(x,y, color="black", s=200, marker=(3,0, angle-90)) ax.axis('off') From 08f9ee9b824d4029a525b67272aaa1da5707c730 Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 5 Nov 2024 10:05:54 -0500 Subject: [PATCH 08/21] add comparison test for plotting --- wntr/tests/test_graphics.py | 62 +++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/wntr/tests/test_graphics.py b/wntr/tests/test_graphics.py index 20ce5c055..a049c15ad 100644 --- a/wntr/tests/test_graphics.py +++ b/wntr/tests/test_graphics.py @@ -9,6 +9,7 @@ import matplotlib.pylab as plt from wntr.graphics.color import custom_colormap import pandas as pd +import numpy as np import wntr testdir = dirname(abspath(str(__file__))) @@ -119,6 +120,67 @@ def test_plot_network5(self): plt.close() self.assertTrue(isfile(filename)) + + def test_plot_network_options(self): + # NOTE:to compare with the old plot_network set compare=True. + # this should be set to false for regular testing + compare = False + + inp_file = join(ex_datadir, "Net3.inp") + wn = wntr.network.WaterNetworkModel(inp_file) + + random_node_values = pd.Series( + np.random.rand(len(wn.node_name_list)), index=wn.node_name_list) + random_link_values = pd.Series( + np.random.rand(len(wn.link_name_list)), index=wn.link_name_list) + random_node_dict_subset = dict(random_node_values.iloc[:10]) + random_link_dict_subset = dict(random_link_values.iloc[:10]) + node_list = list(wn.node_name_list[:10]) + link_list = list(wn.link_name_list[:10]) + + kwarg_list = [ + {"node_attribute": "elevation", + "node_range": [0,1], + "node_alpha": 0.5, + "node_colorbar_label": "test_label"}, + {"link_attribute": "diameter", + "link_range": [0,1], + "link_alpha": 0.5, + "link_colorbar_label": "test_label"}, + {"node_labels": True, + "link_labels": True}, + {"node_attribute": "elevation", + "add_colorbar": False}, + {"link_attribute": "diameter", + "add_colorbar": False}, + {"node_attribute": node_list}, + {"node_attribute": random_node_values}, + {"node_attribute": random_node_dict_subset}, + {"link_attribute": link_list}, + {"link_attribute": random_link_values}, + {"link_attribute": random_link_dict_subset}, + {"directed": True} + ] + + for kwargs in kwarg_list: + filename = abspath(join(testdir, "plot_network_options.png")) + if isfile(filename): + os.remove(filename) + if compare: + fig, ax = plt.subplots(1,2) + wntr.graphics.plot_network(wn, ax=ax[0], title="GIS plot_network", **kwargs) + wntr.graphics.plot_network_nx(wn, ax=ax[1], title="NX plot_network", **kwargs) + fig.savefig(filename, format="png") + plt.close(fig) + else: + plt.figure() + wntr.graphics.plot_network(wn, **kwargs) + plt.savefig(filename, format="png") + plt.close() + + self.assertTrue(isfile(filename)) + os.remove(filename) + def test_plot_interactive_network1(self): From 1eaaade86bc29c000d51a518c4a0d143ccdd7f5b Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 5 Nov 2024 11:56:18 -0500 Subject: [PATCH 09/21] handle cbar manually to avoid error in earthquake demo --- wntr/graphics/network.py | 39 +++++++++++++++++++++++++++------------ 1 file changed, 27 insertions(+), 12 deletions(-) diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index 8ef9b396c..a974b7c2f 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -9,6 +9,7 @@ import matplotlib.pyplot as plt import matplotlib.path as mpath from matplotlib import animation +import matplotlib as mpl import numpy as np try: @@ -218,14 +219,15 @@ def plot_network( if isinstance(link_attribute, list): link_kwds["column"] = "_attribute" link_kwds["cmap"] = custom_colormap(2,("red", "red")) - link_kwds["legend"] = False + link_cbar = False elif isinstance(link_attribute, (dict, pd.Series, str)): link_kwds["cmap"] = link_cmap link_kwds["vmin"] = link_range[0] link_kwds["vmax"] = link_range[1] - link_kwds["legend"] = add_colorbar + link_cbar = add_colorbar else: link_kwds["color"] = "black" + link_cbar = False link_kwds["linewidth"] = link_width link_kwds["alpha"] = link_alpha @@ -245,15 +247,16 @@ def plot_network( if isinstance(node_attribute, list): node_kwds["cmap"] = custom_colormap(2,("red", "red")) - node_kwds["legend"] = False + node_cbar = False elif isinstance(node_attribute, (dict, pd.Series, str)): node_kwds["cmap"] = node_cmap node_kwds["vmin"] = node_range[0] node_kwds["vmax"] = node_range[1] - node_kwds["legend"] = add_colorbar + node_cbar = add_colorbar else: node_kwds["color"] = "black" - + node_cbar = False + node_kwds["alpha"] = node_alpha node_kwds["markersize"] = node_size @@ -264,25 +267,37 @@ def plot_network( # plot nodes - each type is plotted separately to allow for different marker types node_gdf[node_gdf.node_type == "Junction"].plot( - ax=ax, aspect=aspect, zorder=3, legend_kwds=node_cbar_kwds, **node_kwds) + ax=ax, aspect=aspect, zorder=3, legend=False, **node_kwds) - # turn off legend for subsequent node plots - node_kwds["legend"] = False + if node_cbar: + norm = plt.Normalize(vmin=node_kwds["vmin"], vmax=node_kwds["vmax"],) + sm = mpl.cm.ScalarMappable(cmap=node_kwds["cmap"], norm=norm) + sm.set_array([]) + + node_cbar = ax.figure.colorbar(sm, ax=ax, **node_cbar_kwds) node_kwds["markersize"] = node_size * 2.0 node_gdf[node_gdf.node_type == "Tank"].plot( - ax=ax, aspect=aspect, zorder=4, marker=tank_marker, **node_kwds) + ax=ax, aspect=aspect, zorder=4, marker=tank_marker, legend=False, **node_kwds) node_kwds["markersize"] = node_size * 3.0 node_gdf[node_gdf.node_type == "Reservoir"].plot( - ax=ax, aspect=aspect, zorder=5, marker=reservoir_marker, **node_kwds) + ax=ax, aspect=aspect, zorder=5, marker=reservoir_marker, legend=False, **node_kwds) # plot links link_gdf.plot( - ax=ax, aspect=aspect, zorder=1, **background_link_kwds) + ax=ax, aspect=aspect, zorder=1, legend=False, **background_link_kwds) link_gdf.plot( - ax=ax, aspect=aspect, zorder=2, legend_kwds=link_cbar_kwds, **link_kwds) + ax=ax, aspect=aspect, zorder=2, legend=False, **link_kwds) + + # Create a ScalarMappable for the colorbar + if link_cbar: + norm = plt.Normalize(vmin=link_kwds["vmin"], vmax=link_kwds["vmax"]) + sm = mpl.cm.ScalarMappable(cmap=link_kwds["cmap"], norm=norm) + sm.set_array([]) # Needed to create an empty array for the colorbar + + ax.figure.colorbar(sm, ax=ax, **link_cbar_kwds) # Adjusts size and position of colorbar if len(wn_gis.pumps) >0: # wn_gis.pumps.plot(ax=ax, color="purple", aspect=aspect) From 80eaf964af563204b434f935b56b5c0d608a5c03 Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 5 Nov 2024 13:45:34 -0500 Subject: [PATCH 10/21] extend test cases --- wntr/tests/test_graphics.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/wntr/tests/test_graphics.py b/wntr/tests/test_graphics.py index a049c15ad..112797fa6 100644 --- a/wntr/tests/test_graphics.py +++ b/wntr/tests/test_graphics.py @@ -7,6 +7,7 @@ import networkx as nx import matplotlib.pylab as plt +import matplotlib from wntr.graphics.color import custom_colormap import pandas as pd import numpy as np @@ -126,6 +127,9 @@ def test_plot_network_options(self): # this should be set to false for regular testing compare = False + cmap = matplotlib.colormaps['viridis'] + + inp_file = join(ex_datadir, "Net3.inp") wn = wntr.network.WaterNetworkModel(inp_file) @@ -133,6 +137,8 @@ def test_plot_network_options(self): np.random.rand(len(wn.node_name_list)), index=wn.node_name_list) random_link_values = pd.Series( np.random.rand(len(wn.link_name_list)), index=wn.link_name_list) + random_pipe_values = pd.Series( + np.random.rand(len(wn.pipe_name_list)), index=wn.pipe_name_list) random_node_dict_subset = dict(random_node_values.iloc[:10]) random_link_dict_subset = dict(random_link_values.iloc[:10]) node_list = list(wn.node_name_list[:10]) @@ -159,7 +165,13 @@ def test_plot_network_options(self): {"link_attribute": link_list}, {"link_attribute": random_link_values}, {"link_attribute": random_link_dict_subset}, - {"directed": True} + {"directed": True}, + {"link_attribute": random_pipe_values, + "node_size": 0, + "link_cmap": cmap, + "link_range": [0,1], + "link_width": 1.5} + ] for kwargs in kwarg_list: From d0d59f46ca7ff3d27079787f6e068d049f1fb11a Mon Sep 17 00:00:00 2001 From: kbonney Date: Wed, 20 Nov 2024 14:02:50 -0500 Subject: [PATCH 11/21] fix bug caused by node_type no longer provided by GIS geodataframes, fix bugs pointed out by meghna: colorbar scale and spacing, reduce directed arrow size, add option for pumps/valves (default to not plot), --- wntr/graphics/network.py | 84 +++++++++++++++++++++++++++------------- 1 file changed, 57 insertions(+), 27 deletions(-) diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index a974b7c2f..6634cc226 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -6,6 +6,7 @@ import math import networkx as nx import pandas as pd +import matplotlib import matplotlib.pyplot as plt import matplotlib.path as mpath from matplotlib import animation @@ -97,7 +98,7 @@ def plot_network( node_size=20, node_range=None, node_alpha=1, node_cmap=None, node_labels=False, link_width=1, link_range=None, link_alpha=1, link_cmap=None, link_labels=False, add_colorbar=True, node_colorbar_label="", link_colorbar_label="", - directed=False, ax=None, show_plot=True, filename=None): + directed=False, plot_valves=False, plot_pumps=False, ax=None, show_plot=True, filename=None): """ Plot network graphic @@ -210,20 +211,33 @@ def plot_network( node_range = (None, None) wn_gis = wn.to_gis() + # add node_type so that node assets can be plotted separately + wn_gis.junctions["node_type"] = "Junction" + wn_gis.tanks["node_type"] = "Tank" + wn_gis.reservoirs["node_type"] = "Reservoir" link_gdf = pd.concat((wn_gis.pipes, wn_gis.pumps, wn_gis.valves)) node_gdf = pd.concat((wn_gis.junctions, wn_gis.tanks, wn_gis.reservoirs)) # process link attribute link_kwds = _prepare_attribute(link_attribute, link_gdf) + # handle cbar/cmap if isinstance(link_attribute, list): - link_kwds["column"] = "_attribute" link_kwds["cmap"] = custom_colormap(2,("red", "red")) link_cbar = False elif isinstance(link_attribute, (dict, pd.Series, str)): link_kwds["cmap"] = link_cmap - link_kwds["vmin"] = link_range[0] - link_kwds["vmax"] = link_range[1] + + link_attribute_values = link_gdf[link_kwds["column"]] + if link_range[0] is None: + link_kwds["vmin"] = np.nanmin(link_attribute_values) + else: + link_kwds["vmin"] = link_range[0] + if link_range[1] is None: + link_kwds["vmax"] = np.nanmax(link_attribute_values) + else: + link_kwds["vmax"] = link_range[1] + link_cbar = add_colorbar else: link_kwds["color"] = "black" @@ -239,7 +253,7 @@ def plot_network( link_cbar_kwds = {} link_cbar_kwds["shrink"] = 0.5 - link_cbar_kwds["pad"] = 0.0 + link_cbar_kwds["pad"] = 0.05 link_cbar_kwds["label"] = link_colorbar_label # process node attribute @@ -250,8 +264,17 @@ def plot_network( node_cbar = False elif isinstance(node_attribute, (dict, pd.Series, str)): node_kwds["cmap"] = node_cmap - node_kwds["vmin"] = node_range[0] - node_kwds["vmax"] = node_range[1] + + node_attribute_values = node_gdf[node_kwds["column"]] + if node_range[0] is None: + node_kwds["vmin"] = np.nanmin(node_attribute_values) + else: + node_kwds["vmin"] = node_range[0] + if node_range[1] is None: + node_kwds["vmax"] = np.nanmax(node_attribute_values) + else: + node_kwds["vmax"] = node_range[1] + node_cbar = add_colorbar else: node_kwds["color"] = "black" @@ -267,39 +290,39 @@ def plot_network( # plot nodes - each type is plotted separately to allow for different marker types node_gdf[node_gdf.node_type == "Junction"].plot( - ax=ax, aspect=aspect, zorder=3, legend=False, **node_kwds) + ax=ax, aspect=aspect, zorder=3, label="Junctions", legend=False, **node_kwds) - if node_cbar: - norm = plt.Normalize(vmin=node_kwds["vmin"], vmax=node_kwds["vmax"],) - sm = mpl.cm.ScalarMappable(cmap=node_kwds["cmap"], norm=norm) - sm.set_array([]) - - node_cbar = ax.figure.colorbar(sm, ax=ax, **node_cbar_kwds) node_kwds["markersize"] = node_size * 2.0 node_gdf[node_gdf.node_type == "Tank"].plot( - ax=ax, aspect=aspect, zorder=4, marker=tank_marker, legend=False, **node_kwds) + ax=ax, aspect=aspect, zorder=4, marker=tank_marker, label="Tanks", legend=False, **node_kwds) node_kwds["markersize"] = node_size * 3.0 node_gdf[node_gdf.node_type == "Reservoir"].plot( - ax=ax, aspect=aspect, zorder=5, marker=reservoir_marker, legend=False, **node_kwds) + ax=ax, aspect=aspect, zorder=5, marker=reservoir_marker, label="Reservoirs", legend=False, **node_kwds) + + if node_cbar: + sm = mpl.cm.ScalarMappable(cmap=node_kwds["cmap"]) + sm.set_clim(node_kwds["vmin"], node_kwds["vmax"]) + + node_cbar = ax.figure.colorbar(sm, ax=ax, **node_cbar_kwds) # plot links + # background link_gdf.plot( ax=ax, aspect=aspect, zorder=1, legend=False, **background_link_kwds) + # main plot link_gdf.plot( ax=ax, aspect=aspect, zorder=2, legend=False, **link_kwds) - # Create a ScalarMappable for the colorbar if link_cbar: - norm = plt.Normalize(vmin=link_kwds["vmin"], vmax=link_kwds["vmax"]) - sm = mpl.cm.ScalarMappable(cmap=link_kwds["cmap"], norm=norm) - sm.set_array([]) # Needed to create an empty array for the colorbar + sm = mpl.cm.ScalarMappable(cmap=link_kwds["cmap"]) + sm.set_clim(link_kwds["vmin"], link_kwds["vmax"]) - ax.figure.colorbar(sm, ax=ax, **link_cbar_kwds) # Adjusts size and position of colorbar + link_cbar = ax.figure.colorbar(sm, ax=ax, **link_cbar_kwds) - if len(wn_gis.pumps) >0: + if plot_pumps & (len(wn_gis.pumps) > 0): # wn_gis.pumps.plot(ax=ax, color="purple", aspect=aspect) wn_gis.pumps["_midpoint"] = wn_gis.pumps.geometry.interpolate(0.5, normalized=True) wn_gis.pumps["_angle"] = wn_gis.pumps.apply(lambda row: _get_angle(row.geometry), axis=1) @@ -308,7 +331,7 @@ def plot_network( angle = row["_angle"] ax.scatter(x,y, color="black", s=100, marker=(3, 0, angle-90)) - if len(wn_gis.valves) >0: + if plot_valves & (len(wn_gis.valves) > 0): wn_gis.valves["_midpoint"] = wn_gis.valves.geometry.interpolate(0.5, normalized=True) wn_gis.valves["_angle"] = wn_gis.valves.apply(lambda row: _get_angle(row.geometry), axis=1) for idx , row in wn_gis.valves.iterrows(): @@ -317,12 +340,19 @@ def plot_network( ax.scatter(x,y, color="black", s=200, marker=(2,0, angle)) if node_labels: - for x, y, label in zip(wn_gis.junctions.geometry.x, wn_gis.junctions.geometry.y, wn_gis.junctions.index): + for x, y, label in zip(node_gdf.geometry.x, node_gdf.geometry.y, node_gdf.index): ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") + # for x, y, label in zip(wn_gis.tanks.geometry.x, wn_gis.tanks.geometry.y, wn_gis.tanks.index): + # ax.annotate(label, xy=(x, y)) + # for x, y, label in zip(wn_gis.junctions.geometry.x, wn_gis.junctions.geometry.y, wn_gis.junctions.index): + # ax.annotate(label, xy=(x, y)) if link_labels: - midpoints = wn_gis.pipes.geometry.apply(lambda x: x.interpolate(0.5, normalized=True)) - for x, y, label in zip(midpoints.geometry.x, midpoints.geometry.y, wn_gis.pipes.index): + # midpoints = wn_gis.pipes.geometry.apply(lambda x: x.interpolate(0.5, normalized=True)) + # for x, y, label in zip(midpoints.geometry.x, midpoints.geometry.y, wn_gis.pipes.index): + # ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") + midpoints = link_gdf.geometry.apply(lambda x: x.interpolate(0.5, normalized=True)) + for x, y, label in zip(midpoints.geometry.x, midpoints.geometry.y, link_gdf.index): ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") if directed: @@ -331,7 +361,7 @@ def plot_network( for idx , row in link_gdf.iterrows(): x,y = row["_midpoint"].x, row["_midpoint"].y angle = row["_angle"] - ax.scatter(x,y, color="black", s=200, marker=(3,0, angle-90)) + ax.scatter(x,y, color="black", s=50, marker=(3,0, angle-90)) ax.axis('off') From 6e3fcf6903839de37aba669ad4c5ba81ed07ac4f Mon Sep 17 00:00:00 2001 From: kbonney Date: Wed, 20 Nov 2024 14:03:56 -0500 Subject: [PATCH 12/21] add extra test cases, remove unneccesary calls to plt.figure --- wntr/tests/test_graphics.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/wntr/tests/test_graphics.py b/wntr/tests/test_graphics.py index 112797fa6..8cf28e708 100644 --- a/wntr/tests/test_graphics.py +++ b/wntr/tests/test_graphics.py @@ -27,7 +27,6 @@ def test_plot_network1(self): inp_file = join(ex_datadir, "Net6.inp") wn = wntr.network.WaterNetworkModel(inp_file) - plt.figure() wntr.graphics.plot_network(wn) plt.savefig(filename, format="png") plt.close() @@ -42,7 +41,7 @@ def test_plot_network2(self): filename = abspath(join(testdir, "plot_network2_undirected.png")) if isfile(filename): os.remove(filename) - plt.figure() + wntr.graphics.plot_network( wn, node_attribute="elevation", link_attribute="length" ) @@ -55,7 +54,7 @@ def test_plot_network2(self): filename = abspath(join(testdir, "plot_network2_directed.png")) if isfile(filename): os.remove(filename) - plt.figure() + wntr.graphics.plot_network( wn, node_attribute="elevation", link_attribute="length", directed=True ) @@ -72,7 +71,6 @@ def test_plot_network3(self): inp_file = join(ex_datadir, "Net1.inp") wn = wntr.network.WaterNetworkModel(inp_file) - plt.figure() wntr.graphics.plot_network( wn, node_attribute=["11", "21"], @@ -92,7 +90,6 @@ def test_plot_network4(self): inp_file = join(ex_datadir, "Net1.inp") wn = wntr.network.WaterNetworkModel(inp_file) - plt.figure() wntr.graphics.plot_network( wn, node_attribute={"11": 5, "21": 10}, @@ -113,7 +110,6 @@ def test_plot_network5(self): wn = wntr.network.WaterNetworkModel(inp_file) pop = wntr.metrics.population(wn) - plt.figure() wntr.graphics.plot_network( wn, node_attribute=pop, node_range=[0, 500], title="Population" ) @@ -125,7 +121,7 @@ def test_plot_network5(self): def test_plot_network_options(self): # NOTE:to compare with the old plot_network set compare=True. # this should be set to false for regular testing - compare = False + compare = True cmap = matplotlib.colormaps['viridis'] @@ -145,6 +141,10 @@ def test_plot_network_options(self): link_list = list(wn.link_name_list[:10]) kwarg_list = [ + {"node_attribute": "elevation", + "node_range": [0,20], + "node_alpha": 0.5, + "node_colorbar_label": "test_label"}, {"node_attribute": "elevation", "node_range": [0,1], "node_alpha": 0.5, @@ -170,8 +170,9 @@ def test_plot_network_options(self): "node_size": 0, "link_cmap": cmap, "link_range": [0,1], - "link_width": 1.5} - + "link_width": 1.5}, + {"plot_pumps": True, + "plot_valves": False} ] for kwargs in kwarg_list: @@ -185,7 +186,6 @@ def test_plot_network_options(self): fig.savefig(filename, format="png") plt.close(fig) else: - plt.figure() wntr.graphics.plot_network(wn, **kwargs) plt.savefig(filename, format="png") plt.close() From a3b25761a918cb441a2fbc404cd8f1d822c3fd54 Mon Sep 17 00:00:00 2001 From: kbonney Date: Wed, 20 Nov 2024 14:16:19 -0500 Subject: [PATCH 13/21] setting compare to false --- wntr/tests/test_graphics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wntr/tests/test_graphics.py b/wntr/tests/test_graphics.py index 8cf28e708..301ed77f9 100644 --- a/wntr/tests/test_graphics.py +++ b/wntr/tests/test_graphics.py @@ -121,7 +121,7 @@ def test_plot_network5(self): def test_plot_network_options(self): # NOTE:to compare with the old plot_network set compare=True. # this should be set to false for regular testing - compare = True + compare = False cmap = matplotlib.colormaps['viridis'] From c267ac194620bd9d89cf0a9c86026f350f8da16d Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 26 Nov 2024 10:26:34 -0500 Subject: [PATCH 14/21] add alpha to colorbars --- wntr/graphics/network.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index 6634cc226..e61724bb5 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -255,6 +255,7 @@ def plot_network( link_cbar_kwds["shrink"] = 0.5 link_cbar_kwds["pad"] = 0.05 link_cbar_kwds["label"] = link_colorbar_label + link_cbar_kwds["alpha"] = link_alpha # process node attribute node_kwds = _prepare_attribute(node_attribute, node_gdf) @@ -286,6 +287,7 @@ def plot_network( node_cbar_kwds = {} node_cbar_kwds["shrink"] = 0.5 node_cbar_kwds["pad"] = 0.0 + node_cbar_kwds["alpha"] = node_alpha node_cbar_kwds["label"] = node_colorbar_label # plot nodes - each type is plotted separately to allow for different marker types From 664e4a56804cd5fcd7b06d33c28570901046b757 Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 26 Nov 2024 14:17:50 -0800 Subject: [PATCH 15/21] fix color bug on mpl 3.8, change plot_valve/pumps kwarg name, set aspect to "equal" --- wntr/graphics/network.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index e61724bb5..df9100c3c 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -98,7 +98,7 @@ def plot_network( node_size=20, node_range=None, node_alpha=1, node_cmap=None, node_labels=False, link_width=1, link_range=None, link_alpha=1, link_cmap=None, link_labels=False, add_colorbar=True, node_colorbar_label="", link_colorbar_label="", - directed=False, plot_valves=False, plot_pumps=False, ax=None, show_plot=True, filename=None): + directed=False, show_valve_direction=False, show_pump_direction=False, ax=None, show_plot=True, filename=None): """ Plot network graphic @@ -195,7 +195,7 @@ def plot_network( if title is not None: ax.set_title(title) - aspect = None + aspect = "equal" tank_marker = "D" reservoir_marker = "s" @@ -223,7 +223,7 @@ def plot_network( # handle cbar/cmap if isinstance(link_attribute, list): - link_kwds["cmap"] = custom_colormap(2,("red", "red")) + link_kwds["cmap"] = custom_colormap(2,["red", "red"]) link_cbar = False elif isinstance(link_attribute, (dict, pd.Series, str)): link_kwds["cmap"] = link_cmap @@ -261,7 +261,7 @@ def plot_network( node_kwds = _prepare_attribute(node_attribute, node_gdf) if isinstance(node_attribute, list): - node_kwds["cmap"] = custom_colormap(2,("red", "red")) + node_kwds["cmap"] = custom_colormap(2,["red", "red"]) node_cbar = False elif isinstance(node_attribute, (dict, pd.Series, str)): node_kwds["cmap"] = node_cmap @@ -324,23 +324,23 @@ def plot_network( link_cbar = ax.figure.colorbar(sm, ax=ax, **link_cbar_kwds) - if plot_pumps & (len(wn_gis.pumps) > 0): + if show_pump_direction & (len(wn_gis.pumps) > 0): # wn_gis.pumps.plot(ax=ax, color="purple", aspect=aspect) wn_gis.pumps["_midpoint"] = wn_gis.pumps.geometry.interpolate(0.5, normalized=True) wn_gis.pumps["_angle"] = wn_gis.pumps.apply(lambda row: _get_angle(row.geometry), axis=1) for idx , row in wn_gis.pumps.iterrows(): x,y = row["_midpoint"].x, row["_midpoint"].y angle = row["_angle"] - ax.scatter(x,y, color="black", s=100, marker=(3, 0, angle-90)) + ax.scatter(x,y, color="black", s=50, marker=(3, 0, angle-90)) - if plot_valves & (len(wn_gis.valves) > 0): + if show_valve_direction & (len(wn_gis.valves) > 0): wn_gis.valves["_midpoint"] = wn_gis.valves.geometry.interpolate(0.5, normalized=True) wn_gis.valves["_angle"] = wn_gis.valves.apply(lambda row: _get_angle(row.geometry), axis=1) for idx , row in wn_gis.valves.iterrows(): x,y = row["_midpoint"].x, row["_midpoint"].y angle = row["_angle"] - ax.scatter(x,y, color="black", s=200, marker=(2,0, angle)) - + ax.scatter(x,y, color="black", s=50, marker=(3, 0, angle-90)) + if node_labels: for x, y, label in zip(node_gdf.geometry.x, node_gdf.geometry.y, node_gdf.index): ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") From 64017163a83bb81f195b367abc4ac824d2f5aac3 Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 26 Nov 2024 14:26:34 -0800 Subject: [PATCH 16/21] fix test case to adjust for new kwarg names --- wntr/tests/test_graphics.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/wntr/tests/test_graphics.py b/wntr/tests/test_graphics.py index 301ed77f9..9151918a6 100644 --- a/wntr/tests/test_graphics.py +++ b/wntr/tests/test_graphics.py @@ -171,8 +171,8 @@ def test_plot_network_options(self): "link_cmap": cmap, "link_range": [0,1], "link_width": 1.5}, - {"plot_pumps": True, - "plot_valves": False} + {"show_pump_direction": True, + "show_pump_direction": True} ] for kwargs in kwarg_list: From 105d243a53c0a3e692ab11777ccdc75dda0047dc Mon Sep 17 00:00:00 2001 From: kbonney Date: Mon, 2 Dec 2024 06:06:00 -0800 Subject: [PATCH 17/21] add legend and plot all nodes regardless of node_attribute --- wntr/graphics/network.py | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index df9100c3c..0dbc179a4 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -98,7 +98,7 @@ def plot_network( node_size=20, node_range=None, node_alpha=1, node_cmap=None, node_labels=False, link_width=1, link_range=None, link_alpha=1, link_cmap=None, link_labels=False, add_colorbar=True, node_colorbar_label="", link_colorbar_label="", - directed=False, show_valve_direction=False, show_pump_direction=False, ax=None, show_plot=True, filename=None): + directed=False, legend=False, show_valve_direction=False, show_pump_direction=False, ax=None, show_plot=True, filename=None): """ Plot network graphic @@ -289,19 +289,24 @@ def plot_network( node_cbar_kwds["pad"] = 0.0 node_cbar_kwds["alpha"] = node_alpha node_cbar_kwds["label"] = node_colorbar_label - + + # # prepare legend item list + # legend_items = [] + missing_node_kwds={"color": "black", "markersize": node_size / 2} # TODO: customize per element type + missing_link_kwds={"color": "black", "linewidth": link_width / 2} + # plot nodes - each type is plotted separately to allow for different marker types node_gdf[node_gdf.node_type == "Junction"].plot( - ax=ax, aspect=aspect, zorder=3, label="Junctions", legend=False, **node_kwds) - + ax=ax, aspect=aspect, zorder=3, label="Junctions", legend=False, missing_kwds=missing_node_kwds, **node_kwds) + # legend_items.append(plt.Line2D([0], [0], marker='o', color='w', label='Junctions', markerfacecolor='blue', markersize=6)) node_kwds["markersize"] = node_size * 2.0 node_gdf[node_gdf.node_type == "Tank"].plot( - ax=ax, aspect=aspect, zorder=4, marker=tank_marker, label="Tanks", legend=False, **node_kwds) + ax=ax, aspect=aspect, zorder=4, marker=tank_marker, label="Tanks", legend=False, missing_kwds=missing_node_kwds, **node_kwds) node_kwds["markersize"] = node_size * 3.0 node_gdf[node_gdf.node_type == "Reservoir"].plot( - ax=ax, aspect=aspect, zorder=5, marker=reservoir_marker, label="Reservoirs", legend=False, **node_kwds) + ax=ax, aspect=aspect, zorder=5, marker=reservoir_marker, label="Reservoirs", legend=False, missing_kwds=missing_node_kwds,**node_kwds) if node_cbar: sm = mpl.cm.ScalarMappable(cmap=node_kwds["cmap"]) @@ -316,7 +321,7 @@ def plot_network( # main plot link_gdf.plot( - ax=ax, aspect=aspect, zorder=2, legend=False, **link_kwds) + ax=ax, aspect=aspect, zorder=2, legend=False, missing_kwds=missing_link_kwds, **link_kwds) if link_cbar: sm = mpl.cm.ScalarMappable(cmap=link_kwds["cmap"]) @@ -364,6 +369,14 @@ def plot_network( x,y = row["_midpoint"].x, row["_midpoint"].y angle = row["_angle"] ax.scatter(x,y, color="black", s=50, marker=(3,0, angle-90)) + + # NOTE: The coloring on the symbols will change based on the colors of the underlying object. + # If this isn't desired behavior, handles and labels can be build manually using: + # handle = plt.Line2D([0], [0], marker='o', color='w', label='Junctions', markerfacecolor='black', markersize=6) + if legend: + handles, labels = ax.get_legend_handles_labels() + ax.legend(handles, labels, loc='upper right', title="Legend") + ax.axis('off') From c8f9fd5b9007a3d9c6b0d3cf8d6bce1b0c084555 Mon Sep 17 00:00:00 2001 From: kbonney Date: Mon, 2 Dec 2024 06:06:24 -0800 Subject: [PATCH 18/21] add additional tests --- wntr/tests/test_graphics.py | 52 ++++++++++++++++++++++++++++++++++--- 1 file changed, 49 insertions(+), 3 deletions(-) diff --git a/wntr/tests/test_graphics.py b/wntr/tests/test_graphics.py index 9151918a6..1ed01c659 100644 --- a/wntr/tests/test_graphics.py +++ b/wntr/tests/test_graphics.py @@ -117,6 +117,52 @@ def test_plot_network5(self): plt.close() self.assertTrue(isfile(filename)) + + def test_plot_network6(self): + # pumps/valves + filename = abspath(join(testdir, "plot_network6.png")) + if isfile(filename): + os.remove(filename) + + inp_file = join(ex_datadir, "Net6.inp") + wn = wntr.network.WaterNetworkModel(inp_file) + + # verify that direction points away from start node + start_nodes = [] + for link_name in wn.pump_name_list+wn.valve_name_list: + link = wn.get_link(link_name) + start_nodes.append(link.start_node_name) + + # pump=0, valve=1 + link_type = pd.Series(0, index=wn.pump_name_list+wn.valve_name_list) + link_type[wn.valve_name_list] = 1 + + wntr.graphics.plot_network( + wn, node_attribute=start_nodes, link_attribute=link_type, + show_pump_direction=True, show_valve_direction=True + ) + plt.savefig(filename, format="png") + plt.close() + + self.assertTrue(isfile(filename)) + + def test_plot_network7(self): + # legend + filename = abspath(join(testdir, "plot_network7.png")) + if isfile(filename): + os.remove(filename) + + inp_file = join(ex_datadir, "Net6.inp") + wn = wntr.network.WaterNetworkModel(inp_file) + + wntr.graphics.plot_network( + wn, node_attribute="elevation", link_attribute="diameter", + add_colorbar=True, legend=True + ) + plt.savefig(filename, format="png") + plt.close() + + self.assertTrue(isfile(filename)) def test_plot_network_options(self): # NOTE:to compare with the old plot_network set compare=True. @@ -126,7 +172,7 @@ def test_plot_network_options(self): cmap = matplotlib.colormaps['viridis'] - inp_file = join(ex_datadir, "Net3.inp") + inp_file = join(ex_datadir, "Net6.inp") wn = wntr.network.WaterNetworkModel(inp_file) random_node_values = pd.Series( @@ -171,8 +217,8 @@ def test_plot_network_options(self): "link_cmap": cmap, "link_range": [0,1], "link_width": 1.5}, - {"show_pump_direction": True, - "show_pump_direction": True} + # {"show_pump_direction": True, + # "show_pump_direction": True} ] for kwargs in kwarg_list: From 0558ff8824b96500e6fa770abf47a3a6284b1bdb Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 3 Dec 2024 09:52:34 -0500 Subject: [PATCH 19/21] remove old plot network code, clean up plot network, add default colorbar label when passing a str for attribute, fix missingkwds markersize --- wntr/graphics/__init__.py | 2 +- wntr/graphics/network.py | 271 ++++---------------------------------- 2 files changed, 26 insertions(+), 247 deletions(-) diff --git a/wntr/graphics/__init__.py b/wntr/graphics/__init__.py index 7238bd7ea..b0d086674 100644 --- a/wntr/graphics/__init__.py +++ b/wntr/graphics/__init__.py @@ -1,7 +1,7 @@ """ The wntr.graphics package contains graphic functions """ -from wntr.graphics.network import plot_network, plot_network_nx, plot_interactive_network, plot_leaflet_network, network_animation +from wntr.graphics.network import plot_network, plot_interactive_network, plot_leaflet_network, network_animation from wntr.graphics.layer import plot_valve_layer from wntr.graphics.curve import plot_fragility_curve, plot_pump_curve, plot_tank_volume_curve from wntr.graphics.color import custom_colormap, random_colormap diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index 0dbc179a4..2433b0003 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -6,7 +6,6 @@ import math import networkx as nx import pandas as pd -import matplotlib import matplotlib.pyplot as plt import matplotlib.path as mpath from matplotlib import animation @@ -97,8 +96,8 @@ def plot_network( wn, node_attribute=None, link_attribute=None, title=None, node_size=20, node_range=None, node_alpha=1, node_cmap=None, node_labels=False, link_width=1, link_range=None, link_alpha=1, link_cmap=None, link_labels=False, - add_colorbar=True, node_colorbar_label="", link_colorbar_label="", - directed=False, legend=False, show_valve_direction=False, show_pump_direction=False, ax=None, show_plot=True, filename=None): + add_colorbar=True, node_colorbar_label=None, link_colorbar_label=None, + directed=False, legend=False, ax=None, show_plot=True, filename=None): """ Plot network graphic @@ -194,7 +193,7 @@ def plot_network( if title is not None: ax.set_title(title) - + aspect = "equal" tank_marker = "D" @@ -209,6 +208,13 @@ def plot_network( link_range = (None, None) if node_range is None: node_range = (None, None) + + # use attribute name if no other label is provided + if node_colorbar_label is None and isinstance(node_attribute, str): + node_colorbar_label = node_attribute + if link_colorbar_label is None and isinstance(link_attribute, str): + link_colorbar_label = link_attribute + wn_gis = wn.to_gis() # add node_type so that node assets can be plotted separately @@ -219,7 +225,17 @@ def plot_network( node_gdf = pd.concat((wn_gis.junctions, wn_gis.tanks, wn_gis.reservoirs)) # process link attribute + # link_kwds = {} link_kwds = _prepare_attribute(link_attribute, link_gdf) + # Node attribute + # if node_attribute is not None: + # if isinstance(node_attribute, list): + # node_cmap = 'Reds' + # add_colorbar = False + # node_attribute = _format_node_attribute(node_attribute, wn) + # else: + # add_colorbar = False + # link_ # handle cbar/cmap if isinstance(link_attribute, list): @@ -292,21 +308,21 @@ def plot_network( # # prepare legend item list # legend_items = [] - missing_node_kwds={"color": "black", "markersize": node_size / 2} # TODO: customize per element type + missing_node_kwds={"color": "black"} missing_link_kwds={"color": "black", "linewidth": link_width / 2} # plot nodes - each type is plotted separately to allow for different marker types node_gdf[node_gdf.node_type == "Junction"].plot( - ax=ax, aspect=aspect, zorder=3, label="Junctions", legend=False, missing_kwds=missing_node_kwds, **node_kwds) + ax=ax, aspect=aspect, zorder=3, legend=False, label="Junction", missing_kwds=missing_node_kwds, **node_kwds) # legend_items.append(plt.Line2D([0], [0], marker='o', color='w', label='Junctions', markerfacecolor='blue', markersize=6)) node_kwds["markersize"] = node_size * 2.0 node_gdf[node_gdf.node_type == "Tank"].plot( - ax=ax, aspect=aspect, zorder=4, marker=tank_marker, label="Tanks", legend=False, missing_kwds=missing_node_kwds, **node_kwds) + ax=ax, aspect=aspect, zorder=4, marker=tank_marker, legend=False, label="Tank", missing_kwds=missing_node_kwds, **node_kwds) node_kwds["markersize"] = node_size * 3.0 node_gdf[node_gdf.node_type == "Reservoir"].plot( - ax=ax, aspect=aspect, zorder=5, marker=reservoir_marker, label="Reservoirs", legend=False, missing_kwds=missing_node_kwds,**node_kwds) + ax=ax, aspect=aspect, zorder=5, marker=reservoir_marker, legend=False, label="Reservoir", missing_kwds=missing_node_kwds,**node_kwds) if node_cbar: sm = mpl.cm.ScalarMappable(cmap=node_kwds["cmap"]) @@ -329,35 +345,11 @@ def plot_network( link_cbar = ax.figure.colorbar(sm, ax=ax, **link_cbar_kwds) - if show_pump_direction & (len(wn_gis.pumps) > 0): - # wn_gis.pumps.plot(ax=ax, color="purple", aspect=aspect) - wn_gis.pumps["_midpoint"] = wn_gis.pumps.geometry.interpolate(0.5, normalized=True) - wn_gis.pumps["_angle"] = wn_gis.pumps.apply(lambda row: _get_angle(row.geometry), axis=1) - for idx , row in wn_gis.pumps.iterrows(): - x,y = row["_midpoint"].x, row["_midpoint"].y - angle = row["_angle"] - ax.scatter(x,y, color="black", s=50, marker=(3, 0, angle-90)) - - if show_valve_direction & (len(wn_gis.valves) > 0): - wn_gis.valves["_midpoint"] = wn_gis.valves.geometry.interpolate(0.5, normalized=True) - wn_gis.valves["_angle"] = wn_gis.valves.apply(lambda row: _get_angle(row.geometry), axis=1) - for idx , row in wn_gis.valves.iterrows(): - x,y = row["_midpoint"].x, row["_midpoint"].y - angle = row["_angle"] - ax.scatter(x,y, color="black", s=50, marker=(3, 0, angle-90)) - if node_labels: for x, y, label in zip(node_gdf.geometry.x, node_gdf.geometry.y, node_gdf.index): ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") - # for x, y, label in zip(wn_gis.tanks.geometry.x, wn_gis.tanks.geometry.y, wn_gis.tanks.index): - # ax.annotate(label, xy=(x, y)) - # for x, y, label in zip(wn_gis.junctions.geometry.x, wn_gis.junctions.geometry.y, wn_gis.junctions.index): - # ax.annotate(label, xy=(x, y)) if link_labels: - # midpoints = wn_gis.pipes.geometry.apply(lambda x: x.interpolate(0.5, normalized=True)) - # for x, y, label in zip(midpoints.geometry.x, midpoints.geometry.y, wn_gis.pipes.index): - # ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") midpoints = link_gdf.geometry.apply(lambda x: x.interpolate(0.5, normalized=True)) for x, y, label in zip(midpoints.geometry.x, midpoints.geometry.y, link_gdf.index): ax.annotate(label, xy=(x, y))#, xytext=(3, 3),)# textcoords="offset points") @@ -375,221 +367,8 @@ def plot_network( # handle = plt.Line2D([0], [0], marker='o', color='w', label='Junctions', markerfacecolor='black', markersize=6) if legend: handles, labels = ax.get_legend_handles_labels() - ax.legend(handles, labels, loc='upper right', title="Legend") - - - ax.axis('off') - - if filename: - plt.savefig(filename) - - if show_plot is True: - plt.show(block=False) + leg = ax.legend(handles, labels, loc='upper right', title="Legend") - return ax - - -def plot_network_nx(wn, node_attribute=None, link_attribute=None, title=None, - 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', - directed=False, ax=None, show_plot=True, filename=None): - """ - Plot network graphic - - Parameters - ---------- - wn : wntr WaterNetworkModel - A WaterNetworkModel object - - node_attribute : None, str, list, pd.Series, or dict, optional - - - If node_attribute is a string, then a node attribute dictionary is - created using node_attribute = wn.query_node_attribute(str) - - If node_attribute is a list, then each node in the list is given a - value of 1. - - If node_attribute is a pd.Series, then it should be in the format - {nodeid: x} where nodeid is a string and x is a float. - - If node_attribute is a dict, then it should be in the format - {nodeid: x} where nodeid is a string and x is a float - - link_attribute : None, str, list, pd.Series, or dict, optional - - - If link_attribute is a string, then a link attribute dictionary is - created using edge_attribute = wn.query_link_attribute(str) - - If link_attribute is a list, then each link in the list is given a - value of 1. - - If link_attribute is a pd.Series, then it should be in the format - {linkid: x} where linkid is a string and x is a float. - - If link_attribute is a dict, then it should be in the format - {linkid: x} where linkid is a string and x is a float. - - title: str, optional - Plot title - - node_size: int, optional - Node size - - node_range: list, optional - Node color range ([None,None] indicates autoscale) - - node_alpha: int, optional - Node transparency - - node_cmap: matplotlib.pyplot.cm colormap or list of named colors, optional - Node colormap - - node_labels: bool, optional - If True, the graph will include each node labelled with its name. - - link_width: int, optional - Link width - - link_range : list, optional - Link color range ([None,None] indicates autoscale) - - link_alpha : int, optional - Link transparency - - link_cmap: matplotlib.pyplot.cm colormap or list of named colors, optional - Link colormap - - link_labels: bool, optional - If True, the graph will include each link labelled with its name. - - add_colorbar: bool, optional - Add colorbar - - node_colorbar_label: str, optional - Node colorbar label - - link_colorbar_label: str, optional - Link colorbar label - - directed: bool, optional - If True, plot the directed graph - - ax: matplotlib axes object, optional - Axes for plotting (None indicates that a new figure with a single - axes will be used) - - show_plot: bool, optional - If True, show plot with plt.show() - - filename : str, optional - Filename used to save the figure - - Returns - ------- - ax : matplotlib axes object - """ - - if ax is None: # create a new figure - plt.figure(facecolor='w', edgecolor='k') - ax = plt.gca() - - # Graph - G = wn.to_graph() - if not directed: - G = G.to_undirected() - - # Position - pos = nx.get_node_attributes(G,'pos') - if len(pos) == 0: - pos = None - - # Define node properties - add_node_colorbar = add_colorbar - if node_attribute is not None: - - if isinstance(node_attribute, list): - if node_cmap is None: - node_cmap = ['red', 'red'] - add_node_colorbar = False - - if node_cmap is None: - node_cmap = plt.get_cmap('Spectral_r') - elif isinstance(node_cmap, list): - if len(node_cmap) == 1: - node_cmap = node_cmap*2 - node_cmap = custom_colormap(len(node_cmap), node_cmap) - - node_attribute = _format_node_attribute(node_attribute, wn) - nodelist,nodecolor = zip(*node_attribute.items()) - - else: - nodelist = None - nodecolor = 'k' - - add_link_colorbar = add_colorbar - if link_attribute is not None: - - if isinstance(link_attribute, list): - if link_cmap is None: - link_cmap = ['red', 'red'] - add_link_colorbar = False - - if link_cmap is None: - link_cmap = plt.get_cmap('Spectral_r') - elif isinstance(link_cmap, list): - if len(link_cmap) == 1: - link_cmap = link_cmap*2 - link_cmap = custom_colormap(len(link_cmap), link_cmap) - - link_attribute = _format_link_attribute(link_attribute, wn) - - # Replace link_attribute dictionary defined as - # {link_name: attr} with {(start_node, end_node, link_name): attr} - 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 - link_attribute = attr - - linklist,linkcolor = zip(*link_attribute.items()) - else: - linklist = None - linkcolor = 'k' - - if title is not None: - ax.set_title(title) - - edge_background = nx.draw_networkx_edges(G, pos, edge_color='grey', - 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], - linewidths=0, ax=ax) - edges = nx.draw_networkx_edges(G, pos, edgelist=linklist, arrows=directed, - 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) - if link_labels: - labels = {} - for link_name in wn.link_name_list: - link = wn.get_link(link_name) - labels[(link.start_node_name, link.end_node_name)] = link_name - nx.draw_networkx_edge_labels(G, pos, labels, font_size=7, ax=ax) - if add_node_colorbar and node_attribute: - clb = plt.colorbar(nodes, shrink=0.5, pad=0, ax=ax) - clb.ax.set_title(node_colorbar_label, fontsize=10) - if add_link_colorbar and link_attribute: - if link_range[0] is None: - vmin = min(link_attribute.values()) - else: - vmin = link_range[0] - if link_range[1] is None: - vmax = max(link_attribute.values()) - else: - vmax = link_range[1] - sm = plt.cm.ScalarMappable(cmap=link_cmap, norm=plt.Normalize(vmin=vmin, vmax=vmax)) - sm.set_array([]) - clb = plt.colorbar(sm, shrink=0.5, pad=0.05, ax=ax) - clb.ax.set_title(link_colorbar_label, fontsize=10) - ax.axis('off') if filename: From 4a8f5a484589ae0e449369e65b1ea6cc69e1fdb9 Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 3 Dec 2024 09:53:04 -0500 Subject: [PATCH 20/21] remove pump/valve direction plotting tests --- wntr/tests/test_graphics.py | 32 +------------------------------- 1 file changed, 1 insertion(+), 31 deletions(-) diff --git a/wntr/tests/test_graphics.py b/wntr/tests/test_graphics.py index 1ed01c659..f5c5c6c7c 100644 --- a/wntr/tests/test_graphics.py +++ b/wntr/tests/test_graphics.py @@ -119,36 +119,8 @@ def test_plot_network5(self): self.assertTrue(isfile(filename)) def test_plot_network6(self): - # pumps/valves - filename = abspath(join(testdir, "plot_network6.png")) - if isfile(filename): - os.remove(filename) - - inp_file = join(ex_datadir, "Net6.inp") - wn = wntr.network.WaterNetworkModel(inp_file) - - # verify that direction points away from start node - start_nodes = [] - for link_name in wn.pump_name_list+wn.valve_name_list: - link = wn.get_link(link_name) - start_nodes.append(link.start_node_name) - - # pump=0, valve=1 - link_type = pd.Series(0, index=wn.pump_name_list+wn.valve_name_list) - link_type[wn.valve_name_list] = 1 - - wntr.graphics.plot_network( - wn, node_attribute=start_nodes, link_attribute=link_type, - show_pump_direction=True, show_valve_direction=True - ) - plt.savefig(filename, format="png") - plt.close() - - self.assertTrue(isfile(filename)) - - def test_plot_network7(self): # legend - filename = abspath(join(testdir, "plot_network7.png")) + filename = abspath(join(testdir, "plot_network6.png")) if isfile(filename): os.remove(filename) @@ -217,8 +189,6 @@ def test_plot_network_options(self): "link_cmap": cmap, "link_range": [0,1], "link_width": 1.5}, - # {"show_pump_direction": True, - # "show_pump_direction": True} ] for kwargs in kwarg_list: From 046006a2b848ea169fd0621415f7a0255c7a3ffd Mon Sep 17 00:00:00 2001 From: kbonney Date: Tue, 3 Dec 2024 10:20:45 -0500 Subject: [PATCH 21/21] update plot_network to use _format functions instead of _prepare function --- wntr/graphics/network.py | 161 ++++++++++++++---------------------- wntr/tests/test_graphics.py | 11 +-- 2 files changed, 68 insertions(+), 104 deletions(-) diff --git a/wntr/graphics/network.py b/wntr/graphics/network.py index 2433b0003..7b49704fd 100644 --- a/wntr/graphics/network.py +++ b/wntr/graphics/network.py @@ -43,33 +43,6 @@ def _get_angle(line, loc=0.5): angle = math.degrees(angle) return angle - -def _prepare_attribute(attribute, gdf): - kwds = {} - if attribute is not None: - # if dict convert to a series - if isinstance(attribute, dict): - attribute = pd.Series(attribute) - # if series add as a column to link gdf - if isinstance(attribute, pd.Series): - gdf["_attribute"] = attribute - kwds["column"] = "_attribute" - # if list, create new boolean column that captures which indices are in the list - elif isinstance(attribute, list): - gdf["_attribute"] = np.nan - gdf.loc[gdf.index.isin(attribute), "_attribute"] = 1 - kwds["column"] = "_attribute" - # if str, assert that column name exists - elif isinstance(attribute, str): - if attribute not in gdf.columns: - raise KeyError(f"attribute {attribute} does not exist.") - kwds["column"] = attribute - else: - raise TypeError("attribute must be dict, Series, list, or str") - else: - kwds["color"] = "black" - return kwds - def _format_node_attribute(node_attribute, wn): if isinstance(node_attribute, str): @@ -213,8 +186,7 @@ def plot_network( if node_colorbar_label is None and isinstance(node_attribute, str): node_colorbar_label = node_attribute if link_colorbar_label is None and isinstance(link_attribute, str): - link_colorbar_label = link_attribute - + link_colorbar_label = link_attribute wn_gis = wn.to_gis() # add node_type so that node assets can be plotted separately @@ -224,37 +196,71 @@ def plot_network( link_gdf = pd.concat((wn_gis.pipes, wn_gis.pumps, wn_gis.valves)) node_gdf = pd.concat((wn_gis.junctions, wn_gis.tanks, wn_gis.reservoirs)) - # process link attribute - # link_kwds = {} - link_kwds = _prepare_attribute(link_attribute, link_gdf) # Node attribute - # if node_attribute is not None: - # if isinstance(node_attribute, list): - # node_cmap = 'Reds' - # add_colorbar = False - # node_attribute = _format_node_attribute(node_attribute, wn) - # else: - # add_colorbar = False - # link_ - - # handle cbar/cmap - if isinstance(link_attribute, list): - link_kwds["cmap"] = custom_colormap(2,["red", "red"]) - link_cbar = False - elif isinstance(link_attribute, (dict, pd.Series, str)): - link_kwds["cmap"] = link_cmap + node_kwds = {} + node_cbar = add_colorbar + if node_attribute is not None: + node_gdf["_attribute"] = _format_node_attribute(node_attribute, wn) + node_kwds["column"] = "_attribute" - link_attribute_values = link_gdf[link_kwds["column"]] - if link_range[0] is None: - link_kwds["vmin"] = np.nanmin(link_attribute_values) + # handle cbar/cmap + if isinstance(node_attribute, list): + node_kwds["cmap"] = custom_colormap(2,["red", "red"]) + node_cbar = False + elif isinstance(node_attribute, (dict, pd.Series, str)): + node_kwds["cmap"] = node_cmap + + # manually extract min/max if no range is given + node_attribute_values = node_gdf[node_kwds["column"]] + if node_range[0] is None: + node_kwds["vmin"] = np.nanmin(node_attribute_values) + else: + node_kwds["vmin"] = node_range[0] + if node_range[1] is None: + node_kwds["vmax"] = np.nanmax(node_attribute_values) + else: + node_kwds["vmax"] = node_range[1] else: - link_kwds["vmin"] = link_range[0] - if link_range[1] is None: - link_kwds["vmax"] = np.nanmax(link_attribute_values) + raise TypeError("attribute must be dict, Series, list, or str") + else: + node_kwds["color"] = "black" + node_cbar = False + + node_kwds["alpha"] = node_alpha + node_kwds["markersize"] = node_size + + node_cbar_kwds = {} + node_cbar_kwds["shrink"] = 0.5 + node_cbar_kwds["pad"] = 0.0 + node_cbar_kwds["alpha"] = node_alpha + node_cbar_kwds["label"] = node_colorbar_label + + # Link attribute + link_kwds = {} + link_cbar = add_colorbar + if link_attribute is not None: + link_gdf["_attribute"] = pd.Series(_format_link_attribute(link_attribute, wn)) + link_kwds["column"] = "_attribute" + + # handle cbar/cmap + if isinstance(link_attribute, list): + link_kwds["cmap"] = custom_colormap(2,["red", "red"]) + link_cbar = False + elif isinstance(link_attribute, (dict, pd.Series, str)): + link_kwds["cmap"] = link_cmap + + # manually extract min/max if no range is given + link_attribute_values = link_gdf[link_kwds["column"]] + if link_range[0] is None: + link_kwds["vmin"] = np.nanmin(link_attribute_values) + else: + link_kwds["vmin"] = link_range[0] + if link_range[1] is None: + link_kwds["vmax"] = np.nanmax(link_attribute_values) + else: + link_kwds["vmax"] = link_range[1] else: - link_kwds["vmax"] = link_range[1] - - link_cbar = add_colorbar + raise TypeError("attribute must be dict, Series, list, or str") else: link_kwds["color"] = "black" link_cbar = False @@ -264,7 +270,7 @@ def plot_network( background_link_kwds = {} background_link_kwds["color"] = "grey" - background_link_kwds["linewidth"] = link_width + background_link_kwds["linewidth"] = link_width / 2 background_link_kwds["alpha"] = link_alpha link_cbar_kwds = {} @@ -273,48 +279,12 @@ def plot_network( link_cbar_kwds["label"] = link_colorbar_label link_cbar_kwds["alpha"] = link_alpha - # process node attribute - node_kwds = _prepare_attribute(node_attribute, node_gdf) - - if isinstance(node_attribute, list): - node_kwds["cmap"] = custom_colormap(2,["red", "red"]) - node_cbar = False - elif isinstance(node_attribute, (dict, pd.Series, str)): - node_kwds["cmap"] = node_cmap - - node_attribute_values = node_gdf[node_kwds["column"]] - if node_range[0] is None: - node_kwds["vmin"] = np.nanmin(node_attribute_values) - else: - node_kwds["vmin"] = node_range[0] - if node_range[1] is None: - node_kwds["vmax"] = np.nanmax(node_attribute_values) - else: - node_kwds["vmax"] = node_range[1] - - node_cbar = add_colorbar - else: - node_kwds["color"] = "black" - node_cbar = False - - node_kwds["alpha"] = node_alpha - node_kwds["markersize"] = node_size - - node_cbar_kwds = {} - node_cbar_kwds["shrink"] = 0.5 - node_cbar_kwds["pad"] = 0.0 - node_cbar_kwds["alpha"] = node_alpha - node_cbar_kwds["label"] = node_colorbar_label - - # # prepare legend item list - # legend_items = [] missing_node_kwds={"color": "black"} - missing_link_kwds={"color": "black", "linewidth": link_width / 2} + missing_link_kwds={"color": "black"} # plot nodes - each type is plotted separately to allow for different marker types node_gdf[node_gdf.node_type == "Junction"].plot( ax=ax, aspect=aspect, zorder=3, legend=False, label="Junction", missing_kwds=missing_node_kwds, **node_kwds) - # legend_items.append(plt.Line2D([0], [0], marker='o', color='w', label='Junctions', markerfacecolor='blue', markersize=6)) node_kwds["markersize"] = node_size * 2.0 node_gdf[node_gdf.node_type == "Tank"].plot( @@ -362,9 +332,6 @@ def plot_network( angle = row["_angle"] ax.scatter(x,y, color="black", s=50, marker=(3,0, angle-90)) - # NOTE: The coloring on the symbols will change based on the colors of the underlying object. - # If this isn't desired behavior, handles and labels can be build manually using: - # handle = plt.Line2D([0], [0], marker='o', color='w', label='Junctions', markerfacecolor='black', markersize=6) if legend: handles, labels = ax.get_legend_handles_labels() leg = ax.legend(handles, labels, loc='upper right', title="Legend") diff --git a/wntr/tests/test_graphics.py b/wntr/tests/test_graphics.py index f5c5c6c7c..8caed7e45 100644 --- a/wntr/tests/test_graphics.py +++ b/wntr/tests/test_graphics.py @@ -143,7 +143,6 @@ def test_plot_network_options(self): cmap = matplotlib.colormaps['viridis'] - inp_file = join(ex_datadir, "Net6.inp") wn = wntr.network.WaterNetworkModel(inp_file) @@ -163,14 +162,12 @@ def test_plot_network_options(self): "node_range": [0,20], "node_alpha": 0.5, "node_colorbar_label": "test_label"}, - {"node_attribute": "elevation", - "node_range": [0,1], - "node_alpha": 0.5, - "node_colorbar_label": "test_label"}, {"link_attribute": "diameter", - "link_range": [0,1], + "link_range": [0,None], "link_alpha": 0.5, "link_colorbar_label": "test_label"}, + {"link_attribute": "diameter", + "node_attribute": "elevation"}, {"node_labels": True, "link_labels": True}, {"node_attribute": "elevation", @@ -198,7 +195,7 @@ def test_plot_network_options(self): if compare: fig, ax = plt.subplots(1,2) wntr.graphics.plot_network(wn, ax=ax[0], title="GIS plot_network", **kwargs) - wntr.graphics.plot_network_nx(wn, ax=ax[1], title="NX plot_network", **kwargs) + plot_network_nx(wn, ax=ax[1], title="NX plot_network", **kwargs) fig.savefig(filename, format="png") plt.close(fig) else: