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,
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):
"""
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 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 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)
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.get_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,
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 directed:
vmin = min(map(abs, link_attribute.values()))
vmax = max(map(abs, link_attribute.values()))
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)
else:
clb = plt.colorbar(edges, 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)
return ax
def plot_network(wn,
node_attribute=None,
link_attribute=None,
title=None,
node_size=20,
node_range=[None, None],
node_cmap=None,
node_labels=False,
link_width=1,
link_range=[None, None],
link_cmap=None,
link_labels=False,
add_colorbar=True,
directed=False,
ax=None):
"""
Plot network graphic using networkx.
Parameters
----------
wn : wntr WaterNetworkModel
A WaterNetworkModel object
node_attribute : str, list, pd.Series, or dict, optional
(default = None)
- 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,time): x} or {nodeid: x} where nodeid is a string and x is
a float. The time index is not used in the plot.
- 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 : str, list, pd.Series, or dict, optional
(default = None)
- 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,time): x} or {linkid: x} where linkid is a string and x is
a float. The time index is not used in the plot.
- 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 (default = None)
node_size : int, optional
Node size (default = 10)
node_range : list, optional
Node range (default = [None,None], autoscale)
node_cmap : matplotlib.pyplot.cm colormap, optional
Node colormap (default = jet)
node_labels: bool, optional
If True, the graph will include each node labelled with its name.
(default = False)
link_width : int, optional
Link width (default = 1)
link_range : list, optional
Link range (default = [None,None], autoscale)
link_cmap : matplotlib.pyplot.cm colormap, optional
Link colormap (default = jet)
link_labels: bool, optional
If True, the graph will include each link labelled with its name.
(default = False)
add_colorbar : bool, optional
Add colorbar (default = True)
directed : bool, optional
If True, plot the directed graph (default = False, converts the graph
to undirected)
ax : matplotlib axes object, optional
Axes for plotting (default = None, creates a new figure with a single
axes)
Returns
-------
nodes, edges
Notes
-----
For more network draw options, see nx.draw_networkx
"""
if plt is None:
raise ImportError('matplotlib is required')
if node_cmap is None:
node_cmap = plt.cm.Spectral_r
if link_cmap is None:
link_cmap = plt.cm.Spectral_r
if ax is None: # create a new figure
plt.figure(facecolor='w', edgecolor='k')
ax = plt.gca()
# Graph
G = wn.get_graph()
if not directed:
G = G.to_undirected()
# Position
pos = nx.get_node_attributes(G, 'pos')
if len(pos) == 0:
pos = None
# Node attribute
node_attr_from_list = False
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)))
node_attr_from_list = True
if isinstance(node_attribute, pd.Series):
if node_attribute.index.nlevels == 2: # (nodeid, time) index
# drop time
node_attribute.reset_index(level=1, drop=True, inplace=True)
node_attribute = dict(node_attribute)
# Define node list, color, and colormap
if node_attribute is None:
nodelist = None
nodecolor = 'k'
else:
nodelist, nodecolor = zip(*node_attribute.items())
if node_attr_from_list:
nodecolor = 'r'
add_colorbar = False
# Link attribute
link_attr_from_list = False
if isinstance(link_attribute, str):
link_attribute = wn.query_link_attribute(link_attribute)
if isinstance(link_attribute, list):
all_link_attribute = dict(
zip(wn.link_name_list, [0] * len(wn.link_name_list)))
for link in link_attribute:
all_link_attribute[link] = 1
link_attribute = all_link_attribute
link_attr_from_list = True
if isinstance(link_attribute, pd.Series):
if link_attribute.index.nlevels == 2: # (linkid, time) index
# drop time
link_attribute.reset_index(level=1, drop=True, inplace=True)
link_attribute = dict(link_attribute)
# Replace link_attribute dictionary defined as
# {link_name: attr} with {(start_node, end_node, link_name): attr}
if link_attribute is not None:
attr = {}
for link_name, value in link_attribute.items():
link = wn.get_link(link_name)
attr[(link.start_node_name, link.end_node_name, link_name)] = value
link_attribute = attr
if type(link_width) is dict:
attr = {}
for link_name, value in link_width.items():
link = wn.get_link(link_name)
attr[(link.start_node_name, link.end_node_name, link_name)] = value
link_width = attr
# Define link list, color, and colormap
if link_attribute is None:
linklist = None
linkcolor = 'k'
else:
linklist, linkcolor = zip(*link_attribute.items())
if link_attr_from_list:
link_cmap = custom_colormap(2, ['black', 'red'])
add_colorbar = False
if type(link_width) is dict:
linklist2, link_width = zip(*link_width.items())
if not linklist == linklist2:
logger.warning('Link color and width do not share the same \
indexes, link width changed to 1.')
link_width = 1
if title is not None:
ax.set_title(title)
nodes = nx.draw_networkx_nodes(G,
pos,
with_labels=False,
nodelist=nodelist,
node_color=nodecolor,
node_size=node_size,
cmap=node_cmap,
vmin=node_range[0],
vmax=node_range[1],
linewidths=0,
ax=ax)
edges = nx.draw_networkx_edges(G,
pos,
edgelist=linklist,
edge_color=linkcolor,
width=link_width,
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_colorbar and node_attribute:
plt.colorbar(nodes, shrink=0.5, pad=0, ax=ax)
if add_colorbar and link_attribute:
plt.colorbar(edges, shrink=0.5, pad=0.05, ax=ax)
ax.axis('off')
return nodes, edges
def plot_network(wn,
node_attribute=None,
link_attribute=None,
title=None,
node_size=20,
node_range=[None, None],
node_cmap=None,
node_labels=False,
link_width=1,
link_range=[None, None],
link_cmap=None,
link_labels=False,
add_colorbar=True,
directed=False,
ax=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 range ([None,None] indicates autoscale)
node_cmap : matplotlib.pyplot.cm colormap, 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 range ([None,None] indicates autoscale)
link_cmap : matplotlib.pyplot.cm colormap, 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
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)
Returns
-------
nodes, edges
Notes
-----
For more network draw options, see nx.draw_networkx
"""
if plt is None:
raise ImportError('matplotlib is required')
if node_cmap is None:
node_cmap = plt.cm.Spectral_r
if link_cmap is None:
link_cmap = plt.cm.Spectral_r
if ax is None: # create a new figure
plt.figure(facecolor='w', edgecolor='k')
ax = plt.gca()
# Graph
G = wn.get_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
if node_attribute is not None:
node_attribute_from_list = False
if isinstance(node_attribute, list):
node_attribute_from_list = True
add_colorbar = False
node_attribute = _format_node_attribute(node_attribute, wn)
nodelist, nodecolor = zip(*node_attribute.items())
if node_attribute_from_list:
nodecolor = 'r'
else:
nodelist = None
nodecolor = 'k'
if link_attribute is not None:
if isinstance(link_attribute, list):
link_cmap = custom_colormap(2, ['red', 'black'])
add_colorbar = False
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,
with_labels=False,
nodelist=nodelist,
node_color=nodecolor,
node_size=node_size,
cmap=node_cmap,
vmin=node_range[0],
vmax=node_range[1],
linewidths=0,
ax=ax)
edges = nx.draw_networkx_edges(G,
pos,
edgelist=linklist,
edge_color=linkcolor,
width=link_width,
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_colorbar and node_attribute:
plt.colorbar(nodes, shrink=0.5, pad=0, ax=ax)
if add_colorbar and link_attribute:
plt.colorbar(edges, shrink=0.5, pad=0.05, ax=ax)
ax.axis('off')
return nodes, edges
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,
link_width=1, link_range=[None,None], link_alpha=1, link_cmap=None, link_labels=False,
valve_layer=None, 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 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 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)
Returns
-------
nodes, edges : matplotlib objects for network nodes and edges
Notes
-----
For more network draw options, see nx.draw_networkx
"""
if plt is None:
raise ImportError('matplotlib is required')
if ax is None: # create a new figure
plt.figure(facecolor='w', edgecolor='k')
ax = plt.gca()
# Graph
G = wn.get_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.cm.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.cm.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, with_labels=False,
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,
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:
clb = plt.colorbar(edges, shrink=0.5, pad=0.05, ax=ax)
clb.ax.set_title(link_colorbar_label, fontsize=10)
ax.axis('off')
if valve_layer is not None:
for valve_name, (pipe_name, node_name) in valve_layer.iterrows():
pipe = wn.get_link(pipe_name)
if node_name == pipe.start_node_name:
start_node = pipe.start_node
end_node = pipe.end_node
elif node_name == pipe.end_node_name:
start_node = pipe.end_node
end_node = pipe.start_node
else:
print("Not valid")
continue
x0 = start_node.coordinates[0]
dx = end_node.coordinates[0] - x0
y0 = start_node.coordinates[1]
dy = end_node.coordinates[1] - y0
valve_coordinates = (x0 + dx * 0.1,
y0 + dy * 0.1)
ax.scatter(valve_coordinates[0], valve_coordinates[1], 15, 'r', 'v')
if filename:
plt.savefig(filename)
return nodes, edges
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,
link_width=1,
link_range=[None, None],
link_alpha=1,
link_cmap=None,
link_labels=False,
valve_layer=None,
add_colorbar=True,
node_colorbar_label='Node',
link_colorbar_label='Link',
directed=False,
ax=None,
filename=None,
plot_center=None,
plot_limits=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 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, list optional
If True, the graph will include each node labelled with its name.
If list, the graph will include only node labels in the list
link_width : int, optional
Link width
link_range : list, optional
Link 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, list optional
If True, the graph will include each link labelled with its name.
If list, the graph will only include link labels in the list.
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)
filename : str, optional
Path and file name to a grahpics output png file of the axes plotted
by this function (None indicates to not output a file).
plot_center : str or 2-list/tuple, optional
str = name of link/node in the wn water network to center the plot about
2-list/tuple = x-y coordinates
Either of these recenter the plot and change the viewed window.
If a link and node share the same name, then the node will be the center.
plot_limits : float, 2-list/tuple - window (horiz,vert) about plot_center
Returns
-------
nodes, edges : matplotlib objects for network nodes and edges
Notes
-----
For more network draw options, see nx.draw_networkx
"""
if ax is None: # create a new figure
plt.figure(facecolor='w', edgecolor='k')
ax = plt.gca()
# Graph
G = wn.get_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.cm.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.cm.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,
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:
if isinstance(node_labels, (list, tuple)):
labels = dict(zip(node_labels, node_labels))
else:
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:
if isinstance(link_labels, (list, tuple)):
links = [wn.get_link(name) for name in link_labels]
labels = dict(
zip([(link.start_node_name, link.end_node_name)
for link in links],
[linkname for linkname in link_labels]))
else:
links = [wn.get_link(name) for name in wn.links]
labels = dict(
zip([(link.start_node_name, link.end_node_name)
for link in links], [name for name in wn.links]))
# 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:
clb = plt.colorbar(edges, shrink=0.5, pad=0.05, ax=ax)
clb.ax.set_title(link_colorbar_label, fontsize=10)
ax.axis('off')
if valve_layer is not None:
for valve_name, (pipe_name, node_name) in valve_layer.iterrows():
pipe = wn.get_link(pipe_name)
if node_name == pipe.start_node_name:
start_node = pipe.start_node
end_node = pipe.end_node
elif node_name == pipe.end_node_name:
start_node = pipe.end_node
end_node = pipe.start_node
else:
print("Not valid")
continue
x0 = start_node.coordinates[0]
dx = end_node.coordinates[0] - x0
y0 = start_node.coordinates[1]
dy = end_node.coordinates[1] - y0
valve_coordinates = (x0 + dx * 0.1, y0 + dy * 0.1)
ax.scatter(valve_coordinates[0], valve_coordinates[1], 15, 'r',
'v')
# provide a way to focus in on a single node/link or coordinate and to
# zoom in
if plot_center:
# derive the length of x and y on the screen.
incorrect_input_string = ("The '{0}' input must be a list or" +
" tuple of length 2{1}")
plot_center_string1 = (", a node name or a link" +
" name in the input water network model!")
if not plot_limits:
xlim = ax.get_xlim()
len_x = xlim[1] - xlim[0]
ylim = ax.get_ylim()
len_y = ylim[1] - ylim[0]
else:
if isinstance(plot_limits, (list, tuple)):
if len(plot_limits) == 1:
len_x = plot_limits[0]
len_y = plot_limits[1]
elif len(plot_limits) == 2:
len_x = plot_limits[0]
len_y = plot_limits[1]
else:
raise ValueError(
incorrect_input_string.format("plot_limits",
", or a float."))
elif isinstance(plot_limits, float):
len_x = plot_limits
len_y = plot_limits
else:
raise ValueError(
incorrect_input_string.format("plot_limits",
", or a float."))
if isinstance(plot_center, (list, tuple)):
if len(plot_center) == 2:
center = plot_center
else:
raise ValueError(
incorrect_input_string.format("plot_center",
plot_center_string1))
else:
if plot_center in wn.node_name_list:
center_node = wn.get_node(plot_center)
center = center_node.coordinates
elif plot_center in wn.link_name_list:
center_link = wn.get_link(plot_center)
start_coord = center_link.start_node.coordinates
end_coord = center_link.end_node.coordinates
center = [(start_coord[0] + end_coord[0]) / 2.0,
(start_coord[1] + end_coord[1]) / 2.0]
else:
raise ValueError(
incorrect_input_string.format("plot_center",
plot_center_string1))
ax.set_xlim((center[0] - len_x / 2, center[0] + len_x / 2))
ax.set_ylim((center[1] - len_y / 2, center[1] + len_y / 2))
else:
if plot_limits:
raise ValueError("The plot_limits input must be used in" +
" coordination with the plot_center input!")
# end of plot_center code.
if filename:
plt.savefig(filename)
return nodes, edges
