# -*- coding: utf-8 -*-
"""
The viswaternet.utils.base module contains plotting functions that are
frequently utilized by other plotting functions. This includes base element
drawing, legend drawing, color map, and label drawing functions.
"""
import numpy as np
import pandas as pd
import networkx.drawing.nx_pylab as nxp
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.lines import Line2D
from mpl_toolkits.axes_grid1 import make_axes_locatable
from viswaternet.utils import save_fig, normalize_parameter
[docs]def draw_nodes(
self,
ax,
node_list,
parameter_results=None,
node_size=None,
node_shape=None,
node_border_width=None,
node_border_color=None,
node_color=None,
vmin=None,
vmax=None,
label=None,
style=None):
"""Draws continuous nodal data onto the figure.
Arguments
---------
ax : axes._subplots.AxesSubplot
Matplotlib axes object.
node_list : string, array-like
List of draw_nodes to be drawn.
parameter_results : array-like
The data associated with each node.
node_size : integer, array-like
The size of the nodes being drawn. Can either be a single string value
or an array of values for each node being drawn.
Should only be used if calling draw_nodes() manually.
node_shape : string, array-like
The shape of the nodes being drawn. Can either be a single string value
or an array of values for each node being drawn.
Should only be used if calling draw_nodes() manually.
Refer to the matplotlib documentation for available marker types.
https://matplotlib.org/stable/api/markers_api.html
node_border_color : string, array-like
The color of the node borders for the nodes being drawn. Can either be
a single string value or an array of values for each node being drawn.
Should only be used if calling draw_nodes() manually.
node_border_width : integer, array-like
The width of the node borders. Can either be a single
string value or an array of values for each node being drawn.
Should only be used if calling draw_nodes() manually.
vmin : integer
The minimum value of the color bar.
vmax : integer
The maximum value of the color bar.
label : string, array-like
Labels to be assigned to each of the nodes drawn. Can either be a
single string value or an array of values for each node being drawn.
(Not Implemented Yet)
Should only be used if calling draw_nodes() manually.
style : VisWaterNet Style Object
The style object to be used.
"""
# Initalize parameters
model = self.model
if style is None:
style = self.default_style
args = style.args
if parameter_results is None:
parameter_results = pd.DataFrame([])
args = style.args
if node_size is None:
node_size = args['node_size']
draw_tanks = args['draw_tanks']
draw_reservoirs = args['draw_reservoirs']
cmap = args['cmap']
if node_shape is None:
node_shape = args['node_shape']
if node_border_width is None:
node_border_width = args['node_border_width']
if node_border_color is None:
node_border_color = args['node_border_color']
if node_color is None:
node_color = args['node_color']
if node_size is None:
node_size = (np.ones(len(node_list)) * 100).tolist()
# Checks if some data values are given
if parameter_results.values.tolist():
# If values is less than this value, we treat it as a negative.
node_list = [node_list[node_list.index(name)]
for name in node_list
if ((name not in model["tank_names"]
or draw_tanks is False)
and (name not in model["reservoir_names"]
or draw_reservoirs is False))]
parameter_results = parameter_results.loc[node_list]
parameter_results = parameter_results.values.tolist()
if isinstance(node_size, tuple):
min_size = node_size[0]
max_size = node_size[1]
if min_size is not None and max_size is not None:
node_size = normalize_parameter(
parameter_results, min_size, max_size)
if np.min(parameter_results) < -1e-5:
# Gets the cmap object from matplotlib
try:
cmap = mpl.colormaps[cmap]
except Exception:
if isinstance(cmap, mpl.colors.LinearSegmentedColormap) \
or isinstance(cmap, mpl.colors.ListedColormap):
pass
else:
raise Exception('Invalid cmap!')
# If both vmin and vmax are None, set vmax to the max data
# value and vmin to the negative of the max data value. This
# ensures that the colorbar is centered at 0.
if vmin is None and vmax is None:
g = nxp.draw_networkx_nodes(
model["G"],
model["pos_dict"],
ax=ax,
nodelist=node_list,
node_size=node_size,
node_color=parameter_results,
cmap=cmap,
vmax=np.max(parameter_results),
vmin=-np.max(parameter_results),
node_shape=node_shape,
linewidths=node_border_width,
edgecolors=node_border_color)
# Otherwise, just pass the user-given parameters
else:
g = nxp.draw_networkx_nodes(
model["G"],
model["pos_dict"],
ax=ax,
nodelist=node_list,
node_size=node_size,
node_color=parameter_results,
vmax=vmax,
vmin=vmin,
cmap=cmap,
node_shape=node_shape,
linewidths=node_border_width,
edgecolors=node_border_color)
# Return networkx object
return g
else:
# Gets the cmap object from matplotlib
try:
cmap = mpl.colormaps[cmap]
except Exception:
if isinstance(cmap, mpl.colors.LinearSegmentedColormap) \
or isinstance(cmap, mpl.colors.ListedColormap):
pass
else:
raise Exception('Invalid cmap!')
# If both vmin and vmax are None, don't pass vmin and vmax,
# as networkx will handle the limits of the colorbar
# itself.
if vmin is None and vmax is None:
g = nxp.draw_networkx_nodes(
model["G"],
model["pos_dict"],
ax=ax,
nodelist=node_list,
node_size=node_size,
node_color=parameter_results,
cmap=cmap,
node_shape=node_shape,
linewidths=node_border_width,
edgecolors=node_border_color)
# Otherwise, just pass the user-given parameters
else:
g = nxp.draw_networkx_nodes(
model["G"],
model["pos_dict"],
ax=ax,
nodelist=node_list,
node_size=node_size,
node_color=parameter_results,
cmap=cmap,
node_shape=node_shape,
linewidths=node_border_width,
edgecolors=node_border_color,
vmin=vmin,
vmax=vmax)
# Return networkx object
return g
# Draw without any data associated with draw_nodes
else:
nxp.draw_networkx_nodes(
model["G"],
model["pos_dict"],
ax=ax,
nodelist=node_list,
node_size=node_size,
node_color=node_color,
node_shape=node_shape,
edgecolors=node_border_color,
linewidths=node_border_width)
[docs]def draw_links(
self,
ax,
link_list,
parameter_results=None,
link_width=None,
link_style=None,
link_arrows=None,
link_color=None,
vmin=None,
vmax=None,
label=None,
style=None):
"""Draws continuous link data onto the figure.
Arguments
---------
ax : axes._subplots.AxesSubplot
Matplotlib axes object.
link_list : string, array-like
List of draw_links to be drawn.
parameter_results : array-like
The data associated with each node.
link_width : string, array-like
The width of the link being drawn. Can either be a single string value
or an array of values for each link being drawn.
Should only be used if calling draw_links() manually.
link_style : string, array-like
The style of the link being drawn. Can either be a single string value
or an array of values for each link being drawn.
Should only be used if calling draw_links() manually.
Refer to the matplotlib documentation for available link styles.
https://matplotlib.org/stable/gallery/lines_bars_and_markers/linestyles.html
link_arrows : string, array-like
Whether arrows should be drawn for each link. Can either be a single
string value or an array of values for each link being drawn.
Should only be used if calling draw_links() manually.
link_color : string, array-like
The color the link being drawn. Can either be a single string value
or an array of values for each link being drawn.
Should only be used if calling draw_links() manually.
vmin : integer
The minimum value of the color bar.
vmax : integer
The maximum value of the color bar.
label : string, array-like
Labels to be assigned to each of the nodes drawn. Can either be a
single string value or an array of values for each node being drawn.
(Not Implemented Yet)
Should only be used if calling draw_links() manually.
style : VisWaterNet Style Object
The style object to be used.
"""
# Initalize parameters
model = self.model
if style is None:
style = self.default_style
args = style.args
if link_width is None:
link_width = args['link_width']
pump_element = args['pump_element']
draw_pumps = args['draw_pumps']
valve_element = args['valve_element']
draw_valves = args['draw_valves']
cmap = args['cmap']
if link_style is None:
link_style = args['link_style']
if link_arrows is None:
link_arrows = args['link_arrows']
if link_color is None:
link_color = args['link_color']
if isinstance(link_list, np.ndarray):
link_list = link_list.tolist()
if parameter_results is None:
parameter_results = pd.DataFrame([])
# Creates default list of link widths
if link_width is None:
link_width = (np.ones(len(link_list)) * 1).tolist()
# Checks if some data values are given
if parameter_results.values.tolist():
link_list = [link_list[link_list.index(name)]
for name in link_list
if ((name not in model["pump_names"]
or pump_element == 'node'
or draw_pumps is False)
and (name not in model["valve_names"]
or valve_element == 'node'
or draw_valves is False))]
edges = [model["pipe_list"][model['G_pipe_name_list'].index(name)]
for name in link_list]
parameter_results = parameter_results.loc[link_list]
parameter_results = parameter_results.values.tolist()
if isinstance(link_width, tuple):
min_size = link_width[0]
max_size = link_width[1]
if min_size is not None and max_size is not None:
link_width = normalize_parameter(
parameter_results, min_size, max_size)
if np.min(parameter_results) < -1e-5:
# Gets the cmap object from matplotlib
try:
cmap = mpl.colormaps[cmap]
except Exception:
if isinstance(cmap, mpl.colors.LinearSegmentedColormap) \
or isinstance(cmap, mpl.colors.ListedColormap):
pass
else:
raise Exception('Invalid cmap!')
# If both vmin and vmax are None, set vmax to the max data
# value and vmin to the negative of the max data value. This
# ensures that the colorbar is centered at 0.
if vmin is None and vmax is None:
g = nxp.draw_networkx_edges(
model["G"],
model["pos_dict"],
ax=ax,
edgelist=edges,
edge_color=parameter_results,
edge_vmax=np.max(parameter_results),
edge_vmin=-np.max(parameter_results),
edge_cmap=cmap,
style=link_style,
arrows=link_arrows,
width=link_width,
node_size=0)
# Otherwise, just pass the user-given parameters
else:
g = nxp.draw_networkx_edges(
model["G"],
model["pos_dict"],
ax=ax,
edgelist=edges,
edge_color=parameter_results,
edge_vmax=vmax,
edge_vmin=vmin,
edge_cmap=cmap,
style=link_style,
arrows=link_arrows,
width=link_width,
node_size=0)
# Return networkx object
return g
else:
# Gets the cmap object from matplotlib
try:
cmap = mpl.colormaps[cmap]
except Exception:
if isinstance(cmap, mpl.colors.LinearSegmentedColormap) \
or isinstance(cmap, mpl.colors.ListedColormap):
pass
else:
raise Exception('Invalid cmap!')
# If both vmin and vmax are None, don't pass vmin and vmax,
# as networkx will handle the limits of the colorbar
# itself.
if vmin is None and vmax is None:
g = nxp.draw_networkx_edges(
model["G"],
model["pos_dict"],
ax=ax,
edgelist=edges,
edge_color=parameter_results,
edge_cmap=cmap,
style=link_style,
arrows=link_arrows,
width=link_width,
node_size=0)
# Otherwise, just pass the user-given parameters
else:
g = nxp.draw_networkx_edges(
model["G"],
model["pos_dict"],
ax=ax,
edgelist=edges,
edge_color=parameter_results,
edge_cmap=cmap,
style=link_style,
arrows=link_arrows,
width=link_width,
edge_vmin=vmin,
edge_vmax=vmax,
node_size=0)
# Return networkx object
return g
# Draw without any data associated with draw_links
else:
edges = [model["pipe_list"][model['G_pipe_name_list'].index(name)]
for name in link_list]
nxp.draw_networkx_edges(
model["G"],
model["pos_dict"],
ax=ax,
edgelist=edges,
edge_color=link_color,
style=link_style,
arrows=link_arrows,
width=link_width,
node_size=0)
[docs]def draw_base_elements(
self,
ax,
draw_nodes=True,
element_list=None,
draw_originator=None,
style=None):
"""
Draws base elements (draw_nodes, draw_links, draw_reservoirs, draw_tanks,
draw_pumps, and draw_valves) without any data associated with the elements.
Arguments
---------
ax : axes._subplots.AxesSubplot
Matplotlib axes object.
draw_nodes : boolean
Determines if base draw_nodes with no data associated with them are
drawn. Set to False for all functions excep plot_basic_elements by
default.
element_list : array-like
The list of elements being drawn. Can either be links or nodes, but
not a combination. When this is None, all links/nodes are drawn, but
if not None only the ones listed are drawn.
draw_originator : string
Used with element_list and stores whether the elements being drawn
are nodes or links.
style : VisWaterNet Style Object
The style object to be used.
"""
model = self.model
if style is None:
style = self.default_style
args = style.args
draw_tanks = args['draw_tanks']
draw_reservoirs = args['draw_reservoirs']
base_node_size = args['base_node_size']
base_node_color = args['base_node_color']
reservoir_size = args['reservoir_size']
reservoir_color = args['reservoir_color']
reservoir_border_color = args['reservoir_border_color']
reservoir_border_width = args['reservoir_border_width']
reservoir_shape = args['reservoir_shape']
tank_size = args['tank_size']
tank_color = args['tank_color']
tank_border_color = args['tank_border_color']
tank_border_width = args['tank_border_width']
tank_shape = args['tank_shape']
pump_element = args['pump_element']
draw_pumps = args['draw_pumps']
valve_element = args['valve_element']
draw_valves = args['draw_valves']
base_link_color = args['base_link_color']
base_link_width = args['base_link_width']
base_link_line_style = args['base_link_line_style']
base_link_arrows = args['base_link_arrows']
valve_size = args['valve_size']
valve_color = args['valve_color']
valve_border_color = args['valve_border_color']
valve_border_width = args['valve_border_width']
valve_shape = args['valve_shape']
pump_size = args['pump_size']
pump_color = args['pump_color']
pump_border_color = args['pump_border_color']
pump_border_width = args['pump_border_width']
pump_shape = args['pump_shape']
valve_width = args['valve_width']
valve_line_style = args['valve_line_style']
valve_arrows = args['valve_arrows']
pump_width = args['pump_width']
pump_line_style = args['pump_line_style']
pump_arrows = args['pump_arrows']
# If draw_nodes is True, then draw draw_nodes
if draw_nodes:
node_list = model['node_names']
if element_list is None or draw_originator == 'link':
node_list = [node_list[node_list.index(name)]
for name in node_list
if ((name not in model["tank_names"]
or draw_tanks is False)
and (name not in model["reservoir_names"]
or draw_reservoirs is False))]
else:
node_list = [node_list[node_list.index(name)]
for name in node_list
if ((name not in model["tank_names"]
or draw_tanks is False)
and (name not in model["reservoir_names"]
or draw_reservoirs is False)
and (name not in element_list))]
nxp.draw_networkx_nodes(
model["G"],
model["pos_dict"],
node_size=base_node_size,
nodelist=node_list,
node_color=base_node_color,
ax=ax)
# If draw_reservoirs is True, then draw draw_reservoirs
if draw_reservoirs:
nxp.draw_networkx_nodes(
model["G"],
model["pos_dict"],
ax=ax,
nodelist=model["reservoir_names"],
node_size=reservoir_size,
node_color=reservoir_color,
edgecolors=reservoir_border_color,
linewidths=reservoir_border_width,
node_shape=reservoir_shape,
label="Reservoirs")
# If draw_tanks is True, then draw draw_tanks
if draw_tanks:
nxp.draw_networkx_nodes(
model["G"],
model["pos_dict"],
ax=ax,
nodelist=model["tank_names"],
node_size=tank_size,
node_color=tank_color,
edgecolors=tank_border_color,
linewidths=tank_border_width,
node_shape=tank_shape,
label="Tanks")
# If draw_links is True, then draw draw_links
if draw_links:
pipe_name_list = model['G_pipe_name_list']
if element_list is None or draw_originator == 'node':
edgelist = [model['pipe_list'][pipe_name_list.index(name)]
for name in pipe_name_list
if ((name not in model["pump_names"]
or pump_element == 'node'
or draw_pumps is False)
and (name not in model["valve_names"]
or valve_element == 'node'
or draw_valves is False))]
else:
edgelist = [model['pipe_list'][pipe_name_list.index(name)]
for name in pipe_name_list
if ((name not in model["pump_names"]
or pump_element == 'node'
or draw_pumps is False)
and (name not in model["valve_names"]
or valve_element == 'node'
or draw_valves is False)
and (name not in element_list))]
nxp.draw_networkx_edges(
model["G"],
model["pos_dict"],
edgelist=edgelist,
ax=ax,
edge_color=base_link_color,
width=base_link_width,
style=base_link_line_style,
arrows=base_link_arrows)
# If draw_valves is True, then draw draw_valves
if draw_valves:
if valve_element == 'node':
valve_coordinates = {}
# For each valve, calculate midpoint along link it is located at
# then store the coordinates of where valve should be drawn
for i, (point1, point2) in enumerate(model["G_list_valves_only"]):
midpoint = [(model["wn"].get_node(point1).coordinates[0]
+ model["wn"].get_node(point2).coordinates[0])/2,
(model["wn"].get_node(point1).coordinates[1]
+ model["wn"].get_node(point2).coordinates[1])/2]
valve_coordinates[model["valve_names"][i]] = midpoint
# Draw draw_valves after midpoint calculations
nxp.draw_networkx_nodes(
model["G"],
valve_coordinates,
ax=ax,
nodelist=model["valve_names"],
node_size=valve_size,
node_color=valve_color,
edgecolors=valve_border_color,
linewidths=valve_border_width,
node_shape=valve_shape,
label="Valves")
elif valve_element == 'link':
nxp.draw_networkx_edges(
model["G"],
model["pos_dict"],
ax=ax,
edgelist=model["G_list_valves_only"],
edge_color=valve_color,
width=valve_width,
style=valve_line_style,
arrows=valve_arrows)
# If draw_pumps is True, then draw draw_pumps
if draw_pumps:
if pump_element == 'node':
pump_coordinates = {}
# For each valve, calculate midpoint along link it is located at
# then store the coordinates of where pump should be drawn
for i, (point1, point2) in enumerate(model["G_list_pumps_only"]):
midpoint = [(model["wn"].get_node(point1).coordinates[0]
+ model["wn"].get_node(point2).coordinates[0])/2,
(model["wn"].get_node(point1).coordinates[1]
+ model["wn"].get_node(point2).coordinates[1])/2]
pump_coordinates[model["pump_names"][i]] = midpoint
# Draw draw_valves after midpoint calculations
nxp.draw_networkx_nodes(
model["G"],
pump_coordinates,
ax=ax,
nodelist=model["pump_names"],
node_size=pump_size,
node_color=pump_color,
edgecolors=pump_border_color,
linewidths=pump_border_width,
node_shape=pump_shape,
label="Pumps")
elif pump_element == 'link':
nxp.draw_networkx_edges(
model["G"],
model["pos_dict"],
ax=ax,
edgelist=model["G_list_pumps_only"],
edge_color=pump_color,
width=pump_width,
style=pump_line_style,
arrows=pump_arrows)
[docs]def plot_basic_elements(
self,
ax=None,
draw_nodes=True,
savefig=False,
save_name=None,
style=None):
"""User-level function that draws base elements with no data assocaited
with them, draws a legend, and saves the figure.
Arguments
---------
ax : axes._subplots.AxesSubplot
Matplotlib axes object.
draw_nodes : boolean
Determines if base draw_nodes with no data associated with them are
drawn. Set to False for all functions excep plot_basic_elements by
default.
savefig : boolean
Determines if the figure is saved.
save_name : string
The inputted string will be appended to the name of the network.
Example
-------
>>>import viswaternet as vis
>>>model = vis.VisWNModel(r'Networks/Net3.inp')
...
>>>model.save_fig(save_name='_example')
<Net3_example.png>
style : VisWaterNet Style Object
The style object to be used.
"""
if style is None:
style = self.default_style
# Checks if an axis as been specified
if ax is None:
if ax is None:
fig, ax = plt.subplots(figsize=self.figsize)
ax.set_frame_on(self.axis_frame)
# Draw all base elements w/o data associated with them
draw_base_elements(
self,
ax,
draw_nodes=draw_nodes,
style=style)
# Draw legend if legend is True. Only draws base elements legend
draw_legend(
self,
ax,
style=style)
# Save figure if savefig is set to True
if savefig:
save_fig(self, save_name=save_name, style=style)
[docs]def draw_legend(
self,
ax,
intervals=None,
title=None,
element_size_intervals=None,
element_size_legend_title=None,
element_size_legend_loc=None,
element_size_legend_labels=None,
style=None):
"""Draws the legends for all other plotting functions. There are two
legends that might be drawn. One is the base elements legend with displays
what markers are associated with each element type (draw_nodes, draw_links,
etc.) The other legend is the intervals legend which is the legend for
discrete drawing. Under normal use, draw_legends is not normally called by
the user directly, even with more advanced applications. However, some
specialized plots may require draw_legend to be called directly.
Arguments
---------
ax : axes._subplots.AxesSubplot
Matplotlib axes object.
intervals : array-like, string
If set to 'automatic' then intervals are created automatically on an
equal interval basis. Otherwise, it is the edges of the intervals to be
created. intervals array length should be num_intervals + 1.
color_list : string, array-like
The list of node colors for each interval. Both cmap and color_list can
not be used at the same time to color draw_nodes. If both are, then
color_list takes priority.
cmap : string
The matplotlib color map to be used for plotting. Refer to matplotlib
documentation for possible inputs.
https://matplotlib.org/stable/users/explain/colors/colormaps.html
title : string
The title text of the legend.
element_size_intervals : integer
The number of intervals to be used if an element size legend is used.
element_size_legend_title : string
The title of the element size legend.
element_size_legend_loc : string
The location of the element size legend on the figure.
element_size_legend_labels : array-like
The labels of each interval of the element size legend.
style : VisWaterNet Style Object
The style object to be used.
"""
# If no intervals for data legend are specified, then create empty array
if intervals is None:
intervals = []
if style is None:
style = self.default_style
args = style.args
pump_color = args['pump_color']
pump_element = args['pump_element']
draw_pumps = args['draw_pumps']
valve_color = args['valve_color']
valve_element = args['valve_element']
draw_valves = args['draw_valves']
valve_line_style = args['valve_line_style']
valve_arrows = args['valve_arrows']
pump_line_style = args['pump_line_style']
pump_arrows = args['pump_arrows']
base_link_color = args['base_link_color']
base_link_line_style = args['base_link_line_style']
base_link_arrows = args['base_link_arrows']
draw_base_legend = args['draw_base_legend']
base_legend_loc = args['base_legend_loc']
base_legend_label_font_size = args['base_legend_label_font_size']
base_legend_label_color = args['base_legend_label_color']
draw_legend_frame = args['draw_legend_frame']
draw_discrete_legend = args['draw_discrete_legend']
discrete_legend_loc = args['discrete_legend_loc']
discrete_legend_label_font_size = args['discrete_legend_label_font_size']
discrete_legend_label_color = args['discrete_legend_label_color']
draw_legend_frame = args['draw_legend_frame']
discrete_legend_title_color = args['discrete_legend_title_color']
discrete_legend_title_font_size = args['discrete_legend_title_font_size']
cmap = args['cmap']
color_list = args['color_list']
node_size = args['node_size']
node_border_width = args['node_border_width']
node_border_color = args['node_border_color']
link_width = args['link_width']
# Get handles, labels
handles, labels = ax.get_legend_handles_labels()
# Where new handles will be stored
extensions = []
# If draw_pumps is True, then add legend element. Note that right now
# pump_arrows does not affect legend entry, but that it may in the future,
# hence the if statement
if draw_pumps and pump_element == 'link':
if pump_arrows:
extensions.append(Line2D([0], [0], color=pump_color,
linestyle=pump_line_style, lw=4, label='Pumps'))
else:
extensions.append(Line2D([0], [0], color=pump_color,
linestyle=pump_line_style, lw=4, label='Pumps'))
if draw_valves and valve_element == 'link':
if valve_arrows:
extensions.append(Line2D([0], [0], color=valve_color,
linestyle=valve_line_style, lw=4,
label='Valves'))
else:
extensions.append(Line2D([0], [0], color=valve_color,
linestyle=valve_line_style, lw=4,
label='Valves'))
# If draw_base_links is True, then add legend element. Note that right now
# base_link_arrows does not affect legend entry, but that it may in the
# future, hence the if statement
if draw_links:
if base_link_arrows:
extensions.append(Line2D([0], [0], color=base_link_color,
linestyle=base_link_line_style, lw=4,
label='Pipes'))
else:
extensions.append(Line2D([0], [0], color=base_link_color,
linestyle=base_link_line_style, lw=4,
label='Pipes'))
# Extend handles list
handles.extend(extensions)
# If discrete intervals are given
if intervals:
# Draws base legend, which includes the legend for draw_reservoirs,
# draw_tanks, and so on
if draw_base_legend is True:
legend = ax.legend(handles=handles[len(intervals):],
loc=base_legend_loc,
fontsize=base_legend_label_font_size,
labelcolor=base_legend_label_color,
frameon=draw_legend_frame)
# Align legend text to the left, add legend to ax
legend._legend_box.align = "left"
ax.add_artist(legend)
# Draws intervals, or data, legend to the ax
if draw_discrete_legend is True:
if isinstance(discrete_legend_label_color, str) \
and discrete_legend_label_color != 'interval_color':
legend2 = ax.legend(
title=title,
handles=handles[: len(intervals)],
loc=discrete_legend_loc,
fontsize=discrete_legend_label_font_size,
labelcolor=discrete_legend_label_color,
title_fontsize=discrete_legend_title_font_size,
frameon=draw_legend_frame)
if discrete_legend_label_color == 'interval_color':
legend2 = ax.legend(
title=title, handles=handles[: len(intervals)],
loc=discrete_legend_loc,
fontsize=discrete_legend_label_font_size,
title_fontsize=discrete_legend_title_font_size,
frameon=draw_legend_frame)
if color_list:
for i, text in enumerate(legend2.get_texts()):
text.set_color(color_list[i])
elif cmap:
try:
cmap = mpl.colormaps[cmap]
except Exception:
if isinstance(cmap, mpl.colors.LinearSegmentedColormap) \
or isinstance(cmap, mpl.colors.ListedColormap):
pass
else:
raise Exception('Invalid cmap!')
cmap_value = 1 / len(intervals)
for i, text in enumerate(legend2.get_texts()):
text.set_color(cmap(float(cmap_value)))
cmap_value += 1 / len(intervals)
if isinstance(discrete_legend_label_color, list):
legend2 = ax.legend(
title=title,
handles=handles[: len(intervals)],
loc=discrete_legend_loc,
fontsize=discrete_legend_label_font_size,
title_fontsize=discrete_legend_title_font_size,
frameon=draw_legend_frame)
for i, text in enumerate(legend2.get_texts()):
text.set_color(discrete_legend_label_color[i])
# Align legend text to the left, adds title, and adds to ax
legend2._legend_box.align = "left"
legend2.get_title().set_color(discrete_legend_title_color)
ax.add_artist(legend2)
# If there are no intervals, just draw base legend
else:
# Draws base legend, which includes the legend for draw_reservoirs,
# draw_tanks, and so on
if draw_base_legend is True:
legend = ax.legend(handles=handles,
loc=base_legend_loc,
fontsize=base_legend_label_font_size,
labelcolor=base_legend_label_color,
frameon=draw_legend_frame)
# Align legend text to the left, add legend to ax
legend._legend_box.align = "left"
ax.add_artist(legend)
# The following code is for a node/link legend. This adds a 2nd dimension
# to the data that can be plotted, by allowing for changes in size of
# a node/link to represent some parameter. For now it is limited to
# discrete node sizes, which works the same as discrete data for the color
# of draw_nodes. To use it requires a much more involved process than all
# other functions that viswaternet performs, and improvements to ease of
# use will be made in the future.
if node_size is not None and element_size_intervals is not None:
if isinstance(node_size, list):
handles_2 = []
min_size = np.min(node_size)
max_size = np.max(node_size)
marker_sizes = np.linspace(
min_size, max_size, element_size_intervals)
for size, label in zip(marker_sizes, element_size_legend_labels):
handles_2.append(Line2D([], [], marker='.', color='w',
markeredgecolor=node_border_color,
markeredgewidth=node_border_width,
label=label, markerfacecolor='k',
markersize=np.sqrt(size)))
legend3 = ax.legend(
handles=handles_2,
title=element_size_legend_title,
loc=element_size_legend_loc,
fontsize=discrete_legend_label_font_size, # Change later!
title_fontsize=discrete_legend_title_font_size,
labelcolor=discrete_legend_label_color,
frameon=draw_legend_frame)
legend3._legend_box.align = "left"
ax.add_artist(legend3)
if link_width is not None and element_size_intervals is not None:
if isinstance(link_width, list):
handles_2 = []
min_size = np.min(link_width)
max_size = np.max(link_width)
marker_sizes = np.linspace(
min_size, max_size, element_size_intervals)
for size, label in zip(marker_sizes, element_size_legend_labels):
handles_2.append(Line2D([], [], marker=None, color='k',
linewidth=size, label=label))
legend3 = ax.legend(
handles=handles_2,
title=element_size_legend_title,
loc=element_size_legend_loc,
fontsize=discrete_legend_label_font_size,
title_fontsize=discrete_legend_title_font_size,
labelcolor=discrete_legend_label_color,
frameon=draw_legend_frame)
legend3._legend_box.align = "left"
ax.add_artist(legend3)
[docs]def draw_color_bar(
self,
ax,
g,
color_bar_title=None,
style=None):
"""Draws the color bar for all continuous plotting functions. Like
draw_legends, under normal use, draw_color_bar is not normally called by
the user directly, even with more advanced applications. However, some
specialized plots may require draw_color_bar to be called directly.
Arguments
---------
ax : axes._subplots.AxesSubplot
Matplotlib axes object.
g : NetworkX path collection
The list of elements drawn by NetworkX function.
color_bar_title : string
The title of the color bar.
style : VisWaterNet Style Object
The style object to be used.
"""
# Unruly code to make colorbar location nice and symmetrical when dealing
# with subplots especially.
if style is None:
style = self.default_style
args = style.args
color_bar_height = args['color_bar_height']
color_bar_width = args['color_bar_width']
color_bar_loc = args['color_bar_loc']
color_bar_label_loc = args['color_bar_label_loc']
color_bar_label_font_size = args['color_bar_label_font_size']
color_bar_label_font_color = args['color_bar_label_font_color']
divider = make_axes_locatable(ax)
fig = plt.gcf()
if color_bar_loc == 'right':
cax = fig.add_axes([divider.get_position()[0]+divider.get_position()[2]
+ 0.02, (divider.get_position()[1])
+ ((divider.get_position()[3]
* (1-color_bar_height)))/2,
color_bar_width,
divider.get_position()[3]*color_bar_height])
cbar = fig.colorbar(g, cax=cax)
if color_bar_label_loc == 'left':
cbar.ax.yaxis.set_label_position('left')
else:
pass
if color_bar_loc == 'left':
cax = fig.add_axes([divider.get_position()[0] - 0.02 - color_bar_width,
(divider.get_position()[1])
+ ((divider.get_position()[3]
* (1-color_bar_height)))/2,
color_bar_width,
divider.get_position()[3]*color_bar_height])
cbar = fig.colorbar(g, cax=cax)
if color_bar_label_loc == 'left':
cbar.ax.yaxis.set_label_position('left')
else:
pass
if color_bar_loc == 'top':
cax = fig.add_axes([(divider.get_position()[0])
+ (divider.get_position()[2]
* (1-color_bar_width))/2,
divider.get_position()[3] + 0.15,
divider.get_position()[2]*color_bar_width,
color_bar_height])
cbar = fig.colorbar(g, cax=cax, orientation='horizontal')
if color_bar_label_loc == 'top':
cbar.ax.xaxis.set_label_position('top')
else:
pass
if color_bar_loc == 'bottom':
cax = fig.add_axes([(divider.get_position()[0])
+ (divider.get_position()[2]
* (1-color_bar_width))/2,
divider.get_position()[1],
divider.get_position()[2]*color_bar_width,
color_bar_height])
cbar = fig.colorbar(g, cax=cax, orientation='horizontal')
if color_bar_label_loc == 'top':
cbar.ax.xaxis.set_label_position('top')
else:
pass
cbar.set_label(color_bar_title, fontsize=10)
cbar.ax.yaxis.label.set_fontsize(color_bar_label_font_size)
cbar.ax.yaxis.label.set_color(color_bar_label_font_color)
cbar.ax.xaxis.label.set_fontsize(color_bar_label_font_size)
cbar.ax.xaxis.label.set_color(color_bar_label_font_color)
[docs]def draw_label(
self,
labels,
x_coords,
y_coords,
ax=None,
draw_nodes=None,
draw_arrow=True,
label_font_size=11,
label_text_color='k',
label_face_color='white',
label_edge_color='k',
label_alpha=0.9,
label_font_style=None,
label_edge_width=None
):
"""Draws customizable labels on the figure.
There are two modes of coordinate input: If the 'draw_nodes' argument is
not specified, then the label coordinates are processed as absolute
coordinates with possible values from 0 to 1. For instance, (0,0) would
place the label in the bottom left of the figure, while (1,1) would place
the label in the top right of the figure. If the 'draw_nodes' argument is
specified, then the coordinates are processed as coordinates relative to
it's associated node. The scale of the coordinates scaling differs between
networks. For instance, (50,100) would place the label 50 units to the
right, and 100 units above the associated node.
Arguments
---------
ax : axes._subplots.AxesSubplot
Matplotlib axes object.
labels : string, array-like
The label(s) textual content.
x_coords : integer, array-like
The x coordinate(s) of the labels.
y_coords : integer, array-like
The y coordinate(s) of the labels.
draw_nodes : string, array-like
A list of the draw_nodes the labels are to be associated with.
draw_arrow : boolean
Determine if an arrow is drawn from the associated draw_nodes to
labels.
label_font_size : integer
The font size of the labels.
label_text_color : string
The color of the text of the labels.
label_face_color : string
The face color of the label.
label_edge_color : string
The edge color of the label
label_alpha : integer
The transparency of the label. Takes a value between 0.0 and 1.0.
label_font_style : string
The font style of the label. Takes 'normal', 'italic', or 'oblique'.
label_edge_width : integer
The width of the label edge.
"""
model = self.model
if ax is None:
ax = self.ax
if draw_nodes is not None:
for label, node, xCoord, yCoord in \
zip(labels, draw_nodes, x_coords, y_coords):
if draw_arrow:
edge_list = []
if label == node:
pass
else:
model["G"].add_node(label, pos=(xCoord, yCoord))
model["pos_dict"][label] = (
model["wn"].get_node(node).coordinates[0] + xCoord,
model["wn"].get_node(node).coordinates[1] + yCoord)
edge_list.append((node, label))
nxp.draw_networkx_edges(
model["G"], model["pos_dict"], edgelist=edge_list,
edge_color="g", width=0.8, arrows=False)
model["G"].remove_node(label)
model["pos_dict"].pop(label, None)
edge_list.append((node, label))
if draw_arrow is True:
if xCoord < 0:
ax.text(
model["wn"].get_node(node).coordinates[0] + xCoord,
model["wn"].get_node(node).coordinates[1] + yCoord,
s=label,
color=label_text_color,
style=label_font_style,
bbox=dict(facecolor=label_face_color,
alpha=label_alpha,
edgecolor=label_edge_color,
lw=label_edge_width),
horizontalalignment="right",
verticalalignment="center",
fontsize=label_font_size)
if xCoord >= 0:
ax.text(
model["wn"].get_node(node).coordinates[0] + xCoord,
model["wn"].get_node(node).coordinates[1] + yCoord,
s=label,
color=label_text_color,
style=label_font_style,
bbox=dict(facecolor=label_face_color,
alpha=label_alpha,
edgecolor=label_edge_color,
lw=label_edge_width),
horizontalalignment="left",
verticalalignment="center",
fontsize=label_font_size)
else:
ax.text(
model["wn"].get_node(node).coordinates[0] + xCoord,
model["wn"].get_node(node).coordinates[1] + yCoord,
s=label, color=label_text_color, style=label_font_style,
bbox=dict(facecolor=label_face_color,
alpha=label_alpha, edgecolor=label_edge_color,
lw=label_edge_width),
horizontalalignment="center",
verticalalignment="center", fontsize=label_font_size)
elif draw_nodes is None:
for label, xCoord, yCoord in zip(labels, x_coords, y_coords):
ax.text(
xCoord,
yCoord,
s=label,
color=label_text_color,
style=label_font_style,
bbox=dict(facecolor=label_face_color,
alpha=label_alpha,
edgecolor=label_edge_color,
lw=label_edge_width),
horizontalalignment="center",
fontsize=label_font_size,
transform=ax.transAxes)