def draw(self, func, scaling=1, segfunc=False, clim=None, cmap=None):
""" Draw cell in matplotlib line plot collection
"""
from numpy import array, linspace
from matplotlib.collections import LineCollection
from matplotlib import pyplot
self.build_tree(func, segfunc)
pts = self.xyzdv[:, :2]
edges = self.connections
diam = self.xyzdv[:, 3]
data = self.xyzdv[:, 4]
print("DATA RANGE: ", data.min(), data.max())
# Define colors
if not cmap:
from matplotlib.cm import jet as cmap
if not clim:
clim = [data.min(), data.max()]
a = (data - clim[0]) / (clim[1] - clim[0])
# Define line segments
segments = []
for edge in edges:
segments.append([pts[edge[0], :], pts[edge[1], :]])
# Build Line Collection
collection = LineCollection(segments=array(segments),
linewidths=diam * scaling,
colors=cmap(a))
collection.set_array(data)
collection.set_clim(clim[0], clim[1])
pyplot.gca().add_collection(collection, autolim=True)
pyplot.axis('equal')
return collection
def draw_edges(tree,
steiner_pos,
diams=None,
fig=None,
ax=None,
lim=None,
colorbar=True):
'''Draw edges with given positions.'''
if fig is None:
fig, ax = new_fig()
if lim is None:
lim = diam_min, diam_max
pos = merge_pos(tree, steiner_pos)
nodes = tree.get_nodes()
arcs = tree.get_arcs()
x = np.array([pos[n][0] for n in nodes])
y = np.array([pos[n][1] for n in nodes])
segments = [(pos[u], pos[v]) for (u, v) in arcs]
if diams is None:
lines = LineCollection(segments, colors='k', zorder=1)
else:
diams = np.array([diams[a] for a in arcs])
lw = 7 * diams + 1
lines = LineCollection(segments, linewidths=lw, zorder=1)
# set colors
lines.set_array(diams)
lines.set_cmap(_diam_cmap)
lines.set_clim(*lim)
if colorbar:
plt.colorbar(lines, orientation='horizontal')
ax.add_collection(lines)
def plot_multicolor_line(splot,
x,
y,
z,
cmap=None,
vmin=None,
vmax=None,
lw=None,
label=None,
off_x=0,
off_y=0):
"""
Adds a line to a plot that changes color based on the z-value and a cmap (default=viridis)
"""
# Based on the matplotlib example: https://matplotlib.org/gallery/lines_bars_and_markers/multicolored_line.html
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
xo = np.ones(len(segments)) * off_x
yo = np.ones(len(segments)) * off_y
# xyo = list(zip(xo, yo))
xyo = np.array([xo, yo]).T
lc = LineCollection(segments,
cmap=cmap,
offsets=xyo,
transOffset=matplotlib.transforms.IdentityTransform())
lc.set_array(z)
lc.set_clim(vmin=vmin, vmax=vmax)
if lw is not None:
lc.set_linewidth(lw)
if label is not None:
lc.set_label(label)
return splot.add_collection(lc)
def draw_edges(tree, steiner_pos, diams=None, fig=None, ax=None, lim=None,
colorbar=True):
'''Draw edges with given positions.'''
if fig is None:
fig, ax = new_fig()
if lim is None:
lim = diam_min, diam_max
pos = merge_pos(tree, steiner_pos)
nodes = tree.get_nodes()
arcs = tree.get_arcs()
x = np.array([pos[n][0] for n in nodes])
y = np.array([pos[n][1] for n in nodes])
segments = [(pos[u], pos[v]) for (u,v) in arcs]
if diams is None:
lines = LineCollection(segments, colors='k', zorder=1)
else:
diams = np.array([diams[a] for a in arcs])
lw = 7*diams + 1
lines = LineCollection(segments, linewidths=lw, zorder=1)
# set colors
lines.set_array(diams)
lines.set_cmap(_diam_cmap)
lines.set_clim(*lim)
if colorbar:
plt.colorbar(lines, orientation='horizontal')
ax.add_collection(lines)
def _plot_linestring_collection(ax,
geoms,
values=None,
color=None,
cmap=None,
vmin=None,
vmax=None,
**kwargs):
"""
Plots a collection of LineString and MultiLineString geometries to `ax`
Parameters
----------
ax : matplotlib.axes.Axes
where shapes will be plotted
geoms : a sequence of `N` LineStrings and/or MultiLineStrings (can be
mixed)
values : a sequence of `N` values, optional
Values will be mapped to colors using vmin/vmax/cmap. They should
have 1:1 correspondence with the geometries (not their components).
color : single color or sequence of `N` colors
Cannot be used together with `values`.
Returns
-------
collection : matplotlib.collections.Collection that was plotted
"""
from matplotlib.collections import LineCollection
geoms, multiindex = _flatten_multi_geoms(geoms)
if values is not None:
values = np.take(values, multiindex, axis=0)
# LineCollection does not accept some kwargs.
kwargs = {
att: value
for att, value in kwargs.items()
if att not in ["markersize", "marker"]
}
# Add to kwargs for easier checking below.
if color is not None:
kwargs["color"] = color
_expand_kwargs(kwargs, multiindex)
segments = [np.array(linestring.coords)[:, :2] for linestring in geoms]
collection = LineCollection(segments, **kwargs)
if values is not None:
collection.set_array(np.asarray(values))
collection.set_cmap(cmap)
if "norm" not in kwargs:
collection.set_clim(vmin, vmax)
ax.add_collection(collection, autolim=True)
ax.autoscale_view()
return collection
def plot_linestring_collection(ax,
geoms,
values=None,
color=None,
cmap=None,
vmin=None,
vmax=None,
**kwargs):
"""
Plots a collection of LineString and MultiLineString geometries to `ax`
Parameters
----------
ax : matplotlib.axes.Axes
where shapes will be plotted
geoms : a sequence of `N` LineStrings and/or MultiLineStrings (can be
mixed)
values : a sequence of `N` values, optional
Values will be mapped to colors using vmin/vmax/cmap. They should
have 1:1 correspondence with the geometries (not their components).
color : single color or sequence of `N` colors
Cannot be used together with `values`.
Returns
-------
collection : matplotlib.collections.Collection that was plotted
"""
from matplotlib.collections import LineCollection
geoms, values = _flatten_multi_geoms(geoms, values)
if None in values:
values = None
# LineCollection does not accept some kwargs.
if 'markersize' in kwargs:
del kwargs['markersize']
# color=None gives black instead of default color cycle
if color is not None:
kwargs['color'] = color
segments = [np.array(linestring)[:, :2] for linestring in geoms]
collection = LineCollection(segments, **kwargs)
if values is not None:
collection.set_array(np.asarray(values))
collection.set_cmap(cmap)
collection.set_clim(vmin, vmax)
ax.add_collection(collection, autolim=True)
ax.autoscale_view()
return collection
def plot_linestring_collection(ax, geoms, values=None, color=None,
cmap=None, vmin=None, vmax=None, **kwargs):
"""
Plots a collection of LineString and MultiLineString geometries to `ax`
Parameters
----------
ax : matplotlib.axes.Axes
where shapes will be plotted
geoms : a sequence of `N` LineStrings and/or MultiLineStrings (can be
mixed)
values : a sequence of `N` values, optional
Values will be mapped to colors using vmin/vmax/cmap. They should
have 1:1 correspondence with the geometries (not their components).
color : single color or sequence of `N` colors
Cannot be used together with `values`.
Returns
-------
collection : matplotlib.collections.Collection that was plotted
"""
from matplotlib.collections import LineCollection
geoms, values = _flatten_multi_geoms(geoms, values)
if None in values:
values = None
# LineCollection does not accept some kwargs.
if 'markersize' in kwargs:
del kwargs['markersize']
# color=None gives black instead of default color cycle
if color is not None:
kwargs['color'] = color
segments = [np.array(linestring)[:, :2] for linestring in geoms]
collection = LineCollection(segments, **kwargs)
if values is not None:
collection.set_array(np.asarray(values))
collection.set_cmap(cmap)
collection.set_clim(vmin, vmax)
ax.add_collection(collection, autolim=True)
ax.autoscale_view()
return collection
def dg_plot(self):
xylines = [self.xp[self.edges], self.yp[self.edges]]
fig = plt.figure(figsize=(10, 15))
ax = fig.add_subplot(111)
clim = [0, 500]
xlim = np.asarray([616832, 709049]) - self.originx
ylim = np.asarray([2854748, 3050710]) - self.originy
# Find the colorbar limits if unspecified
if clim == None:
clim = []
clim.append(np.min(self.dg))
clim.append(np.max(self.dg))
# Create the inputs needed by line collection
Ne = xylines[0].shape[0]
# Put this into the format needed by LineCollection
linesc = [
zip(xylines[0][ii, :], xylines[1][ii, :]) for ii in range(Ne)
]
collection = LineCollection(linesc,
array=self.dg,
cmap='bone',
linewidth=0.2)
collection.set_clim(vmin=clim[0], vmax=clim[1])
ax.add_collection(collection)
## extremely small dg
ind = self.badedges(threshold=50)
print "Bad edge indexes are ...\n"
print ind
#ind = np.squeeze(np.argwhere(self.dg<50))
#pdb.set_trace()
linesc2 = [zip(xylines[0][ii, :], xylines[1][ii, :]) for ii in ind]
collection2 = LineCollection(linesc2, colors='r', linewidths=2.5)
ax.add_collection(collection2)
## plot outline
#lc = self.outline() #line collection of outline
#ax.add_collection(lc)
ax.set_aspect('equal')
ax.set_xlim(xlim)
ax.set_ylim(ylim)
cb = fig.colorbar(collection, orientation='vertical')
plt.show()
def create_trafo_connection_collection(net, trafos=None, bus_geodata=None, infofunc=None,
cmap=None, clim=None, norm=None, z=None,
cbar_title="Transformer Loading", **kwargs):
"""
Creates a matplotlib line collection of pandapower transformers.
Input:
**net** (pandapowerNet) - The pandapower network
OPTIONAL:
**trafos** (list, None) - The transformers for which the collections are created.
If None, all transformers in the network are considered.
**bus_geodata** (DataFrame, None) - coordinates to use for plotting
If None, net["bus_geodata"] is used
**infofunc** (function, None) - infofunction for the patch element
**kwargs - key word arguments are passed to the patch function
OUTPUT:
**lc** - line collection
"""
trafos = net.trafo if trafos is None else net.trafo.loc[trafos]
if bus_geodata is None:
bus_geodata = net["bus_geodata"]
hv_geo = list(zip(bus_geodata.loc[trafos["hv_bus"], "x"].values,
bus_geodata.loc[trafos["hv_bus"], "y"].values))
lv_geo = list(zip(bus_geodata.loc[trafos["lv_bus"], "x"].values,
bus_geodata.loc[trafos["lv_bus"], "y"].values))
tg = list(zip(hv_geo, lv_geo))
info = [infofunc(tr) if infofunc is not None else [] for tr in trafos.index.values]
lc = LineCollection([(tgd[0], tgd[1]) for tgd in tg], **kwargs)
lc.info = info
if cmap is not None:
if z is None:
z = net.res_trafo.loading_percent.loc[trafos.index]
lc.set_cmap(cmap)
lc.set_norm(norm)
if clim is not None:
lc.set_clim(clim)
lc.set_array(np.ma.masked_invalid(z))
lc.has_colormap = True
lc.cbar_title = cbar_title
return lc
def plot_dft_angles(snap_times,angles,amps,sim_time,logx=False,logy=False,plot_wheel=False,lw=1.5,cmap='hsv',wheel_axis=None,label_size=10):
"""
Plot 2D-DFT amplitudes at a specific frequency as a function of time and angle
"""
if logy : pl.gca().set_yscale('log')
if logx : pl.gca().set_xscale('log')
ax=pl.gca()
custom_axes()
ax.set_xlim(snap_times[1] if pl.gca().get_xscale()=='log' else 0,sim_time)
det_lines = LineCollection([list(zip(snap_times[1:], amp[1:])) for amp in amps],
linewidths=lw,
linestyles='-')
ax.set_ylim((np.amin(amps), np.amax(amps)))
det_lines.set_array(angles)
det_lines.set_clim(0,180)
det_lines.set_cmap(cmap)
ax.add_collection(det_lines)
pl.sci(det_lines)
if plot_wheel is True:
color_steps = 2056
pos = ax.get_position()
wheel_pos= [pos.x0+pos.width/20., pos.y0+pos.height*0.7, pos.width / 4.0, pos.height / 4.0]
if wheel_axis is None:
wheel_axis = pl.gcf().add_axes(wheel_pos, projection='polar')
else:
wheel_axis.projection='polar'
wheel_axis._direction = 2*np.pi
norm = matplotlib.colors.Normalize(0.0, 180.)
cb = matplotlib.colorbar.ColorbarBase(wheel_axis, cmap=cm.get_cmap(cmap,color_steps),
norm=norm,
orientation='horizontal',ticks=[0, 30, 60, 90, 120,150])
cb.ax.tick_params(labelsize=label_size)
cb.outline.set_visible(False)
return ax
def lines(xy, c='b', vmin=None, vmax=None, **kwargs):
"""
xy : sequence of array
Coordinates of points in lines.
`xy` is a sequence of array (line0, line1, ..., lineN) where
line = [(x0, y0), (x1, y1), ... (xm, ym)]
c : color or sequence of color, optional, default : 'b'
`c` can be a single color format string, or a sequence of color
specifications of length `N`, or a sequence of `N` numbers to be
mapped to colors using the `cmap` and `norm` specified via kwargs.
Note that `c` should not be a single numeric RGB or RGBA sequence
because that is indistinguishable from an array of values
to be colormapped. (If you insist, use `color` instead.)
`c` can be a 2-D array in which the rows are RGB or RGBA, however.
vmin, vmax : scalar, optional, default: None
`vmin` and `vmax` are used in conjunction with `norm` to normalize
luminance data. If either are `None`, the min and max of the
color array is used.
kwargs : `~matplotlib.collections.Collection` properties
Eg. alpha, linewidth(lw), linestyle(ls), norm, cmap, transform, etc.
Returns
-------
collection : `~matplotlib.collections.LineCollection`
"""
if np.isscalar(c):
kwargs.setdefault('color', c)
c = None
if 'ls' in kwargs:
kwargs.setdefault('linestyle', kwargs.pop('ls'))
if 'lw' in kwargs:
kwargs.setdefault('linewidth', kwargs.pop('lw'))
collection = LineCollection(xy, **kwargs)
if c is not None:
collection.set_array(np.asarray(c))
collection.set_clim(vmin, vmax)
ax = plt.gca()
ax.add_collection(collection)
ax.autoscale_view()
plt.draw_if_interactive()
if c is not None:
plt.sci(collection)
return collection
def lines(xy, c='b', vmin=None, vmax=None, **kwargs):
"""
xy : sequence of array
Coordinates of points in lines.
`xy` is a sequence of array (line0, line1, ..., lineN) where
line = [(x0, y0), (x1, y1), ... (xm, ym)]
c : color or sequence of color, optional, default : 'b'
`c` can be a single color format string, or a sequence of color
specifications of length `N`, or a sequence of `N` numbers to be
mapped to colors using the `cmap` and `norm` specified via kwargs.
Note that `c` should not be a single numeric RGB or RGBA sequence
because that is indistinguishable from an array of values
to be colormapped. (If you insist, use `color` instead.)
`c` can be a 2-D array in which the rows are RGB or RGBA, however.
vmin, vmax : scalar, optional, default: None
`vmin` and `vmax` are used in conjunction with `norm` to normalize
luminance data. If either are `None`, the min and max of the
color array is used.
kwargs : `~matplotlib.collections.Collection` properties
Eg. alpha, linewidth(lw), linestyle(ls), norm, cmap, transform, etc.
Returns
-------
collection : `~matplotlib.collections.LineCollection`
"""
if np.isscalar(c):
kwargs.setdefault('color', c)
c = None
if 'ls' in kwargs:
kwargs.setdefault('linestyle', kwargs.pop('ls'))
if 'lw' in kwargs:
kwargs.setdefault('linewidth', kwargs.pop('lw'))
collection = LineCollection(xy, **kwargs)
if c is not None:
collection.set_array(np.asarray(c))
collection.set_clim(vmin, vmax)
ax = plt.gca()
ax.add_collection(collection)
ax.autoscale_view()
plt.draw_if_interactive()
if c is not None:
plt.sci(collection)
return collection
def plotedgedata(self, z, xlims=None, ylims=None, colorbar=True, **kwargs):
"""
Plot the unstructured grid edge data
"""
ax = plt.gca()
fig = plt.gcf()
assert(z.shape[0] == self.Ne),\
' length of z scalar vector not equal to number of edges, Ne.'
# Find the colorbar limits if unspecified
if self.clim is None:
self.clim = [z.min(), z.max()]
# Set the xy limits
if xlims is None or ylims is None:
xlims = self.xlims()
ylims = self.ylims()
xylines = [self.xp[self.edges], self.yp[self.edges]]
# Create the inputs needed by line collection
Ne = xylines[0].shape[0]
# Put this into the format needed by LineCollection
linesc = [
list(zip(xylines[0][ii, :], xylines[1][ii, :])) for ii in range(Ne)
]
collection = LineCollection(linesc, array=z, **kwargs)
collection.set_clim(vmin=self.clim[0], vmax=self.clim[1])
ax.add_collection(collection)
ax.set_aspect('equal')
ax.set_xlim(xlims)
ax.set_ylim(ylims)
axcb = None
if colorbar:
axcb = fig.colorbar(collection)
return fig, ax, collection, axcb
def plotedgedata(self, z, xlims=None,ylims=None, colorbar=True, **kwargs):
"""
Plot the unstructured grid edge data
"""
ax=plt.gca()
fig = plt.gcf()
assert(z.shape[0] == self.Ne),\
' length of z scalar vector not equal to number of edges, Ne.'
# Find the colorbar limits if unspecified
if self.clim is None:
self.clim = [z.min(),z.max()]
# Set the xy limits
if xlims is None or ylims is None:
xlims=self.xlims()
ylims=self.ylims()
xylines = [self.xp[self.edges],self.yp[self.edges]]
# Create the inputs needed by line collection
Ne = xylines[0].shape[0]
# Put this into the format needed by LineCollection
linesc = [zip(xylines[0][ii,:],xylines[1][ii,:]) for ii in range(Ne)]
collection = LineCollection(linesc,array=z,**kwargs)
collection.set_clim(vmin=self.clim[0],vmax=self.clim[1])
ax.add_collection(collection)
ax.set_aspect('equal')
ax.set_xlim(xlims)
ax.set_ylim(ylims)
axcb=None
if colorbar:
axcb = fig.colorbar(collection)
return fig, ax, collection, axcb
def construct_eigvect_DOS_plot(xy,
fig,
DOS_ax,
eig_ax,
eigval,
eigvect,
en,
sim_type,
Ni,
Nk,
marker_num=0,
color_scheme='default',
sub_lattice=-1,
cmap_lines='BlueBlackRed',
line_climv=None,
cmap_patches='isolum_rainbow',
draw_strain=False,
lw=3,
bondval_matrix=None,
dotsz=.04,
normalization=0.9,
xycolor=lecmaps.green(),
wzcolor=lecmaps.yellow()):
"""puts together lattice and DOS plots and draws normal mode ellipsoids on top
Parameters
----------
xy: array 2N x 3
Equilibrium position of the gyroscopes
fig :
figure with lattice and DOS drawn
DOS_ax:
axis for the DOS plot
eig_ax
axis for the eigenvalue plot
eigval : array of dimension 2nx1
Eigenvalues of matrix for system
eigvect : array of dimension 2nx2n
Eigenvectors of matrix for system.
Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like:
x0, y0, x1, y1, ... xNP, yNP.
en: int
Number of the eigenvalue you are plotting
marker_num : int
the index of the 'timestep' (or phase of the eigvect) at which we place the t=0 line/dot for each particle
color_scheme : str (default='default')
sub_lattice : int
cmap_lines : str
line_climv : tuple of floats or None
cmap_patches : str
draw_strain : bool
lw : float
bondval_matrix : (NP x max#NN) float array or None
a color specification for the bonds in the network
Returns
----------
fig :
completed figure for normal mode
[scat_fg, p, f_mark] :
things to be cleared before next normal mode is drawn
"""
# ppu = leplt.get_points_per_unit()
s = leplt.absolute_sizer()
plt.sca(DOS_ax)
ev = eigval[en]
ev1 = ev
# Show where current eigenvalue is in DOS plot
(f_mark, ) = plt.plot([np.real(ev), np.real(ev)], plt.ylim(), '-r')
NP = len(xy[:, 0])
im1 = np.imag(ev)
re1 = np.real(ev)
plt.sca(eig_ax)
plt.title('Mode %d; $\Omega=( %0.6f + %0.6f i)$' % (en, re1, im1))
# Preallocate ellipsoid plot vars
angles_arr = np.zeros(NP)
tangles_arr = np.zeros(NP)
# patch = []
polygons, tiltgons = [], []
colors = np.zeros(2 * NP + 2)
# x_mag = np.zeros(NP)
# y_mag = np.zeros(NP)
x0s, y0s = np.zeros(NP), np.zeros(NP)
w0s, z0s = np.zeros(NP), np.zeros(NP)
mag1 = eigvect[en]
# Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like:
# x0, y0, x1, y1, ... xNP, yNP ... w0, z0, ... wNP, zNP.
mag1x = np.array([mag1[2 * i] for i in range(NP)])
mag1y = np.array([mag1[2 * i + 1] for i in range(NP)])
mag1w = np.array([mag1[2 * i] for i in np.arange(NP, 2 * NP)])
mag1z = np.array([mag1[2 * i + 1] for i in np.arange(NP, 2 * NP)])
# Pick a series of times to draw out the ellipsoids
if abs(ev1) > 0:
time_arr = np.arange(21) * 2 * np.pi / (np.abs(ev1) * 20)
exp1 = np.exp(1j * ev1 * time_arr)
else:
time_arr = np.arange(21) * 2 * np.pi / 20
exp1 = np.exp(1j * time_arr)
# Normalization for the ellipsoids
lim_mag1 = np.max(
np.array([
np.sqrt(
np.abs(exp1 * mag1x[i])**2 + np.abs(exp1 * mag1y[i])**2 +
np.abs(exp1 * mag1w[i])**2 + np.abs(exp1 * mag1z[i])**2)
for i in range(len(mag1x))
]).flatten())
if np.isnan(lim_mag1):
print 'found nan for limiting magnitude, replacing lim_mag1 with 1.0'
lim_mag1 = 1.
mag1x *= normalization / lim_mag1
mag1y *= normalization / lim_mag1
mag1w *= normalization / lim_mag1
mag1z *= normalization / lim_mag1
# sys.exit()
cw = []
ccw = []
lines_stretch = []
lines_twist = []
for i in range(NP):
# Draw COM movement of each node as an ellipse
x_disps = 0.5 * (exp1 * mag1x[i]).real
y_disps = 0.5 * (exp1 * mag1y[i]).real
x_vals = xy[i, 0] + x_disps
y_vals = xy[i, 1] + y_disps
# Draw movement of the orientation of each node as an ellipse
w_disps = 0.5 * (exp1 * mag1w[i]).real
z_disps = 0.5 * (exp1 * mag1z[i]).real
w_vals = xy[i, 0] + w_disps
z_vals = xy[i, 1] + z_disps
poly_points = np.array([x_vals, y_vals]).T
tilt_points = np.array([w_vals, z_vals]).T
# polygon = Polygon(poly_points, True, lw=lw, ec='g')
# tiltgon = Polygon(tilt_points, True, lw=lw, ec='r')
# polygon = plt.plot(poly_points[:, 0], poly_points[:, 1], 'g-', lw=lw)
# tiltgon = plt.plot(tilt_points[:, 0], tilt_points[:, 1], 'r-', lw=lw)
npolypts = len(poly_points[:, 0])
for ii in range(npolypts):
lines_s = [[
poly_points[ii, 0], poly_points[(ii + 1) % npolypts, 0]
], [poly_points[ii, 1], poly_points[(ii + 1) % npolypts, 1]]]
lines_t = [[
tilt_points[ii, 0], tilt_points[(ii + 1) % npolypts, 0]
], [tilt_points[ii, 1], tilt_points[(ii + 1) % npolypts, 1]]]
lines_stretch.append(lines_s)
lines_twist.append(lines_t)
# x0 is the marker_num^th element of x_disps
x0 = x_disps[marker_num]
y0 = y_disps[marker_num]
w0 = w_disps[marker_num]
z0 = z_disps[marker_num]
# x0s is the position (global pos, not relative) of each gyro at time = marker_num(out of 81)
x0s[i] = x_vals[marker_num]
y0s[i] = y_vals[marker_num]
w0s[i] = w_vals[marker_num]
z0s[i] = z_vals[marker_num]
# Get angle for xy
mag = np.sqrt(x0**2 + y0**2)
if mag > 0:
anglez = np.arccos(x0 / mag)
else:
anglez = 0
if y0 < 0:
anglez = 2 * np.pi - anglez
# Get angle for wz
tmag = np.sqrt(w0**2 + z0**2)
if tmag > 0:
tanglez = np.arccos(w0 / tmag)
else:
tanglez = 0
if y0 < 0:
tanglez = 2 * np.pi - tanglez
angles_arr[i] = anglez
tangles_arr[i] = tanglez
# polygons.append(polygon)
# tiltgons.append(tiltgon)
# Do Fast Fourier Transform (FFT)
# ff = abs(fft.fft(x_disps + 1j*y_disps))**2
# ff_freq = fft.fftfreq(len(x_vals), 1)
# mm_f = ff_freq[ff == max(ff)][0]
if color_scheme == 'default':
colors[2 * i] = anglez
colors[2 * i + 1] = anglez
# add two more colors to ensure clims go from 0 to 2pi, I think...
colors[2 * NP] = 0
colors[2 * NP + 1] = 2 * np.pi
plt.yticks([])
plt.xticks([])
# this is the part that puts a dot a t=0 point
scat_fg = eig_ax.scatter(x0s,
y0s,
s=s(dotsz),
edgecolors='k',
facecolors='none')
scat_fg2 = eig_ax.scatter(w0s,
z0s,
s=s(dotsz),
edgecolors='gray',
facecolors='none')
try:
NN = np.shape(Ni)[1]
except IndexError:
NN = 0
Rnorm = np.array([x0s, y0s]).T
# Bond Stretches
if draw_strain:
inc = 0
stretches = np.zeros(4 * len(xy))
for i in range(len(xy)):
if NN > 0:
for j, k in zip(Ni[i], Nk[i]):
if i < j and abs(k) > 0:
n1 = float(np.linalg.norm(Rnorm[i] - Rnorm[j]))
n2 = np.linalg.norm(xy[i] - xy[j])
stretches[inc] = (n1 - n2)
inc += 1
stretch = np.array(stretches[0:inc])
else:
# simply get length of BL (len(BL) = inc) by iterating over all bondsssa
inc = 0
for i in range(len(xy)):
if NN > 0:
for j, k in zip(Ni[i], Nk[i]):
if i < j and abs(k) > 0:
inc += 1
# For particles with neighbors, get list of bonds to draw.
# If bondval_matrix is not None, color by the elements of that matrix
if bondval_matrix is not None or draw_strain:
test = list(np.zeros([inc, 1]))
bondvals = list(np.ones([inc, 1]))
inc = 0
xy = np.array([x0s, y0s]).T
for i in range(len(xy)):
if NN > 0:
for j, k in zip(Ni[i], Nk[i]):
if i < j and abs(k) > 0:
test[inc] = [xy[(i, j), 0], xy[(i, j), 1]]
if bondval_matrix is not None:
bondvals[inc] = bondval_matrix[i, j]
inc += 1
# lines connect sites (bonds), while lines_12 draw the black lines from the pinning to location sites
lines = [zip(x, y) for x, y in test]
# Check that we have all the cmaps
if cmap_lines not in plt.colormaps() or cmap_patches not in plt.colormaps(
):
lecmaps.register_cmaps()
# Add lines colored by strain here
if bondval_matrix is not None:
lines_st = LineCollection(lines,
array=bondvals,
cmap=cmap_lines,
linewidth=0.8)
if line_climv is None:
maxk = np.max(np.abs(bondvals))
mink = np.min(np.abs(bondvals))
if (bondvals - bondvals[0] < 1e-8).all():
lines_st.set_clim([mink - 1., maxk + 1.])
else:
lines_st.set_clim([mink, maxk])
lines_st.set_zorder(2)
eig_ax.add_collection(lines_st)
else:
if draw_strain:
lines_st = LineCollection(lines,
array=stretch,
cmap=cmap_lines,
linewidth=0.8)
if line_climv is None:
maxstretch = np.max(np.abs(stretch))
if maxstretch < 1e-8:
line_climv = 1.0
else:
line_climv = maxstretch
lines_st.set_clim([-line_climv, line_climv])
lines_st.set_zorder(2)
eig_ax.add_collection(lines_st)
# Draw lines for movement
lines_stretch = [zip(x, y) for x, y in lines_stretch]
polygons = LineCollection(lines_stretch,
linewidth=lw,
linestyles='solid',
color=xycolor)
polygons.set_zorder(1)
eig_ax.add_collection(polygons)
lines_twist = [zip(x, y) for x, y in lines_twist]
tiltgons = LineCollection(lines_twist,
linewidth=lw,
linestyles='solid',
color=wzcolor)
tiltgons.set_zorder(1)
eig_ax.add_collection(tiltgons)
# Draw polygons
# polygons = PatchCollection(polygons, cmap=cmap_patches, alpha=0.6)
# # p.set_array(np.array(colors))
# polygons.set_clim([0, 2 * np.pi])
# polygons.set_zorder(1)
# eig_ax.add_collection(polygons)
# tiltgons = PatchCollection(tiltgons, cmap=cmap_patches, alpha=0.6)
# # p.set_array(np.array(colors))
# tiltgons.set_clim([0, 2 * np.pi])
# tiltgons.set_zorder(1000)
# eig_ax.add_collection(tiltgons)
eig_ax.set_aspect('equal')
# erased ev/(2*pi) here npm 2016
cw_ccw = [cw, ccw, ev]
# print cw_ccw[1]
# If on a virtualenv, check it here
# if not hasattr(sys, 'real_prefix'):
# plt.show()
# eig_ax.set_facecolor('#000000')
# print 'leplt: construct_eigvect_DOS_plot() exiting'
return fig, [scat_fg, scat_fg2, f_mark, polygons, tiltgons], cw_ccw
def plot_linestring_collection(ax,
geoms,
values=None,
color=None,
cmap=None,
vmin=None,
vmax=None,
**kwargs):
"""
Plots a collection of LineString and MultiLineString geometries to `ax`
Parameters
----------
ax : matplotlib.axes.Axes
where shapes will be plotted
geoms : a sequence of `N` LineStrings and/or MultiLineStrings (can be
mixed)
values : a sequence of `N` values, optional
Values will be mapped to colors using vmin/vmax/cmap. They should
have 1:1 correspondence with the geometries (not their components).
color : single color or sequence of `N` colors
Cannot be used together with `values`.
Returns
-------
collection : matplotlib.collections.Collection that was plotted
"""
from matplotlib.collections import LineCollection
from matplotlib.colors import is_color_like
geoms, multiindex = _flatten_multi_geoms(geoms)
if values is not None:
values = np.take(values, multiindex, axis=0)
# LineCollection does not accept some kwargs.
if "markersize" in kwargs:
del kwargs["markersize"]
# color=None gives black instead of default color cycle
if color is not None:
if is_color_like(color):
kwargs["color"] = color
elif pd.api.types.is_list_like(color):
kwargs["color"] = np.take(color, multiindex, axis=0)
else:
raise TypeError(
"Color attribute has to be a single color or sequence of colors."
)
segments = [np.array(linestring)[:, :2] for linestring in geoms]
collection = LineCollection(segments, **kwargs)
if values is not None:
collection.set_array(np.asarray(values))
collection.set_cmap(cmap)
if "norm" not in kwargs:
collection.set_clim(vmin, vmax)
ax.add_collection(collection, autolim=True)
ax.autoscale_view()
return collection
def construct_haldane_eigvect_DOS_plot(xy,
fig,
DOS_ax,
eig_ax,
eigval,
eigvect,
en,
NL,
KL,
marker_num=0,
color_scheme='default',
sub_lattice=-1,
normalization=None):
"""puts together lattice and DOS plots and draws normal mode ellipsoids on top
Parameters
----------
xy: array 2N x 3
Equilibrium position of the gyroscopes
fig :
figure with lattice and DOS drawn
DOS_ax:
axis for the DOS plot
eig_ax
axis for the eigenvalue plot
eigval : array of dimension 2nx1
Eigenvalues of matrix for system
eigvect : array of dimension 2nx2n
Eigenvectors of matrix for system.
Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like:
x0, y0, x1, y1, ... xNP, yNP.
en: int
Number of the eigenvalue you are plotting
Returns
----------
fig :
completed figure for normal mode
[scat_fg, p, f_mark] :
things to be cleared before next normal mode is drawn
"""
s = leplt.absolute_sizer()
plt.sca(DOS_ax)
ev = eigval[en]
ev1 = ev
# Show where current eigenvalue is in DOS plot
(f_mark, ) = plt.plot([ev, ev], plt.ylim(), '-r')
NP = len(xy)
im1 = np.imag(ev)
re1 = np.real(ev)
plt.sca(eig_ax)
plt.title('Mode %d; $\Omega=( %0.6f + %0.6f i)$' % (en, re1, im1))
# Preallocate ellipsoid plot vars
angles_arr = np.zeros(NP)
patch = []
colors = np.zeros(NP + 2)
x0s = np.zeros(NP)
y0s = np.zeros(NP)
mag1 = eigvect[en]
if normalization is None:
mag1 /= np.max(np.abs(mag1))
else:
mag1 *= normalization * float(len(xy))
# Pick a series of times to draw out the ellipsoid
time_arr = np.arange(81.0) * 2. * np.pi / float(abs(ev1) * 80)
exp1 = np.exp(1j * ev1 * time_arr)
cw = []
ccw = []
lines_1 = []
for i in range(NP):
x_disps = 0.5 * (exp1 * mag1[i]).real
y_disps = 0.5 * (exp1 * mag1[i]).imag
x_vals = xy[i, 0] + x_disps
y_vals = xy[i, 1] + y_disps
poly_points = np.array([x_vals, y_vals]).T
polygon = Polygon(poly_points, True)
# x0 is the marker_num^th element of x_disps
x0 = x_disps[marker_num]
y0 = y_disps[marker_num]
x0s[i] = x_vals[marker_num]
y0s[i] = y_vals[marker_num]
# These are the black lines protruding from pivot point to current position
lines_1.append([[xy[i, 0], x_vals[marker_num]],
[xy[i, 1], y_vals[marker_num]]])
mag = np.sqrt(x0**2 + y0**2)
if mag > 0:
anglez = np.arccos(x0 / mag)
else:
anglez = 0
if y0 < 0:
anglez = 2 * np.pi - anglez
angles_arr[i] = anglez
patch.append(polygon)
if color_scheme == 'default':
colors[i] = anglez
else:
if sub_lattice[i] == 0:
colors[i] = 0
else:
colors[i] = np.pi
ccw.append(i)
colors[NP] = 0
colors[NP + 1] = 2 * np.pi
plt.yticks([])
plt.xticks([])
# this is the part that puts a dot a t=0 point
scat_fg = eig_ax.scatter(x0s[cw], y0s[cw], s=s(.02), c='DodgerBlue')
scat_fg2 = eig_ax.scatter(x0s[ccw], y0s[ccw], s=s(.02), c='Red', zorder=3)
NP = len(xy)
try:
NN = np.shape(NL)[1]
except IndexError:
NN = 0
z = np.zeros(NP)
Rnorm = np.array([x0s, y0s, z]).T
# Bond Stretches
inc = 0
stretches = np.zeros(4 * len(xy))
for i in range(len(xy)):
if NN > 0:
for j, k in zip(NL[i], KL[i]):
if i < j and abs(k) > 0:
n1 = float(linalg.norm(Rnorm[i] - Rnorm[j]))
n2 = linalg.norm(xy[i] - xy[j])
stretches[inc] = (n1 - n2)
inc += 1
# For particles with neighbors, get list of bonds to draw by stretches
test = list(np.zeros([inc, 1]))
inc = 0
xy = np.array([x0s, y0s, z]).T
for i in range(len(xy)):
if NN > 0:
for j, k in zip(NL[i], KL[i]):
if i < j and abs(k) > 0:
test[inc] = [xy[(i, j), 0], xy[(i, j), 1]]
inc += 1
stretch = np.array(stretches[0:inc])
# lines connect sites (bonds), while lines_12 draw the black lines from the pinning to location sites
lines = [zip(x, y) for x, y in test]
lines_12 = [zip(x, y) for x, y in lines_1]
lines_st = LineCollection(lines,
array=stretch,
cmap='seismic',
linewidth=8)
lines_st.set_clim([-1. * 0.25, 1 * 0.25])
lines_st.set_zorder(2)
lines_12_st = LineCollection(lines_12, linewidth=0.8)
lines_12_st.set_color('k')
p = PatchCollection(patch, cmap='hsv', alpha=0.6)
p.set_array(np.array(colors))
p.set_clim([0, 2 * np.pi])
p.set_zorder(1)
# eig_ax.add_collection(lines_st)
eig_ax.add_collection(lines_12_st)
eig_ax.add_collection(p)
eig_ax.set_aspect('equal')
# erased ev/(2*pi) here npm 2016
cw_ccw = [cw, ccw, ev]
# print cw_ccw[1]
return fig, [scat_fg, scat_fg2, p, f_mark, lines_12_st], cw_ccw
def draw_animation_edges(G,
pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
**kwds):
try:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cbook as cb
from matplotlib.colors import colorConverter, Colormap
from matplotlib.collections import LineCollection
import numpy
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
# set edge positions
box_pos = numpy.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
p = 0.25
edge_pos = []
for edge in edgelist:
src, dst = np.array(pos[edge[0]]), np.array(pos[edge[1]])
s = dst - src
# src = src + p * s # Box at beginning
# dst = src + (1-p) * s # Box at the end
dst = src # No edge at all
edge_pos.append((src, dst))
edge_pos = numpy.asarray(edge_pos)
if not cb.iterable(width):
lw = (width, )
else:
lw = width
if not cb.is_scalar_or_string(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color) == len(edge_pos):
if numpy.alltrue([cb.is_scalar_or_string(c) for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple(
[colorConverter.to_rgba(c, alpha) for c in edge_color])
elif numpy.alltrue(
[not cb.is_scalar_or_string(c) for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if numpy.alltrue(
[cb.iterable(c) and len(c) in (3, 4) for c in edge_color]):
edge_colors = tuple(edge_color)
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError(
'edge_color must consist of either color names or numbers')
else:
if cb.is_scalar_or_string(edge_color) or len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError(
'edge_color must be a single color or list of exactly m colors where m is the number or edges'
)
'''
modEdgeColors = list(edge_colors)
modEdgeColors = tuple(modEdgeColors + [colorConverter.to_rgba('w', alpha)
for c in edge_color])
#print(modEdgeColors)
edge_collection = LineCollection(np.asarray(list(edge_pos)*2),
colors=modEdgeColors,
linewidths=[6]*len(list(edge_colors))+[4]*len(list(edge_colors)),
antialiaseds=(1,),
linestyle=style,
transOffset=ax.transData,
)
'''
edge_collection = LineCollection(
edge_pos,
colors=edge_colors,
linewidths=6,
antialiaseds=(1, ),
linestyle=style,
transOffset=ax.transData,
)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
tube_collection = LineCollection(
edge_pos,
colors=tuple([
colorConverter.to_rgba('lightgrey', alpha) for c in edge_color
]),
linewidths=4,
antialiaseds=(1, ),
linestyle=style,
transOffset=ax.transData,
)
tube_collection.set_zorder(1) # edges go behind nodes
tube_collection.set_label(label)
ax.add_collection(tube_collection)
# Note: there was a bug in mpl regarding the handling of alpha values for
# each line in a LineCollection. It was fixed in matplotlib in r7184 and
# r7189 (June 6 2009). We should then not set the alpha value globally,
# since the user can instead provide per-edge alphas now. Only set it
# globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
if edge_colors is None:
if edge_cmap is not None:
assert (isinstance(edge_cmap, Colormap))
edge_collection.set_array(numpy.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
box_collection = Utilities.get_boxes(edge_colors=edge_colors,
edge_pos=box_pos)
box_collection.set_zorder(1) # edges go behind nodes
box_collection.set_label(label)
ax.add_collection(box_collection)
arrow_collection = Utilities.get_arrows_on_edges(
edge_colors=edge_colors, edge_pos=box_pos)
arrow_collection.set_zorder(0)
if arrows:
# Visualize them only if wanted
ax.add_collection(arrow_collection)
return edge_collection, box_collection, tube_collection, arrow_collection
def colorplotMD(self,
xvals,
yvals,
cvals,
title="",
xlabel="",
ylabel="",
clabel="Mass Defect",
cmap="hsv",
config=None,
color="black",
max=1,
marker=None,
zoom="box",
nopaint=False,
test_kda=False,
integerticks=False,
**kwargs):
self._axes = [0.11, 0.1, 0.64, 0.8]
self.clear_plot("nopaint")
self.xlabel = xlabel
self.ylabel = ylabel
self.zoomtype = zoom
if "nticks" in kwargs:
nticks = kwargs["nticks"]
del kwargs['nticks']
else:
nticks = None
if test_kda:
self.kda_test(xvals)
xvals = xvals / self.kdnorm
pubflag = 0
if config is not None:
if config.publicationmode != 0:
pubflag = 1
self.subplot1 = self.figure.add_axes(self._axes)
points = np.array([xvals, yvals]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
# print segments
t = cvals # np.linspace(0, 1, len(xvals), endpoint=True)
lc = LineCollection(segments, cmap=cmap)
lc.set_array(t)
lc.set_clim(0, max)
# print(max, np.amax(t))
self.subplot1.add_collection(lc)
if pubflag == 0:
self.subplot1.set_ylabel(self.ylabel)
self.subplot1.set_title(title)
else:
self.subplot1.spines['top'].set_visible(False)
self.subplot1.spines['right'].set_visible(False)
self.subplot1.get_xaxis().tick_bottom()
self.subplot1.get_yaxis().tick_left()
self.subplot1.get_yaxis().set_tick_params(direction='out')
self.subplot1.get_xaxis().set_tick_params(direction='out')
if config.peaknorm != 2:
self.subplot1.get_yaxis().set_ticks(
[0, np.amax(yvals) / 2,
np.amax(yvals)])
self.subplot1.get_yaxis().set_ticklabels(["0", '%', "100"])
else:
self.subplot1.set_ylabel("Relative Intensity")
if nticks is not None:
self.subplot1.xaxis.set_major_locator(MaxNLocator(nbins=nticks))
if integerticks:
self.subplot1.xaxis.set_major_locator(MaxNLocator(integer=True))
self.subplot1.set_xlabel(self.xlabel)
cax = self.figure.add_axes([0.77, 0.1, 0.04, 0.8])
ticks = np.linspace(0., 1., 11, endpoint=True)
cmap2 = cm.get_cmap(cmap)
self.cbar = colorbar.ColorbarBase(cax,
cmap=cmap2,
orientation="vertical",
ticks=ticks)
ticks = ticks * max
if max > 10:
ticks = np.round(ticks).astype(np.int)
else:
ticks = np.round(ticks, 1)
ticks = ticks.astype(str)
self.cbar.set_ticklabels(ticks)
self.cbar.set_label(clabel)
self.setup_zoom([self.subplot1],
self.zoomtype,
data_lims=[
np.amin(xvals),
np.amin(yvals),
np.max(xvals),
np.amax(yvals)
])
if not nopaint:
self.repaint()
self.flag = True
self.mlist = []
self.x1, self.x2 = None, None
self.colors = []
def create_trafo3w_collection(net, trafo3ws=None, picker=False, infofunc=None, cmap=None, norm=None,
z=None, clim=None, cbar_title="3W-Transformer Loading",
plot_colormap=True, **kwargs):
"""
Creates a matplotlib line collection of pandapower transformers.
Input:
**net** (pandapowerNet) - The pandapower network
OPTIONAL:
**trafo3ws** (list, None) - The three winding transformers for which the collections are
created. If None, all three winding transformers in the network are considered.
**picker** (bool, False) - picker argument passed to the patch collection
**infofunc** (function, None) - infofunction for the patch element
**kwargs - key word arguments are passed to the patch function
OUTPUT:
**lc** - line collection
**pc** - patch collection
"""
trafo3ws = get_index_array(trafo3ws, net.trafo3w.index)
trafo3w_table = net.trafo3w.loc[trafo3ws]
lines = []
circles = []
infos = []
color = kwargs.pop("color", "k")
linewidth = kwargs.pop("linewidths", 2.)
if cmap is not None and z is None:
z = net.res_trafo3w.loading_percent
for i, idx in enumerate(trafo3w_table.index):
# get bus geodata
p1 = net.bus_geodata[["x", "y"]].loc[net.trafo3w.at[idx, "hv_bus"]].values
p2 = net.bus_geodata[["x", "y"]].loc[net.trafo3w.at[idx, "mv_bus"]].values
p3 = net.bus_geodata[["x", "y"]].loc[net.trafo3w.at[idx, "lv_bus"]].values
if np.all(p1 == p2) and np.all(p1 == p3):
continue
p = np.array([p1, p2, p3])
# determine center of buses and minimum distance center-buses
center = sum(p) / 3
d = np.linalg.norm(p - center, axis=1)
r = d.min() / 3
# determine closest bus to center and vector from center to circle midpoint in closest
# direction
closest = d.argmin()
to_closest = (p[closest] - center) / d[closest] * 2 * r / 3
# determine vectors from center to circle midpoint
order = list(range(closest, 3)) + list(range(closest))
cm = np.empty((3, 2))
cm[order.pop(0)] = to_closest
ang = 2 * np.pi / 3 # 120 degree
cm[order.pop(0)] = _rotate_dim2(to_closest, ang)
cm[order.pop(0)] = _rotate_dim2(to_closest, -ang)
# determine midpoints of circles
m = center + cm
# determine endpoints of circles
e = (center - p) * (1 - 5 * r / 3 / d).reshape(3, 1) + p
# save circle and line collection data
ec = color if cmap is None else cmap(norm(z.at[idx]))
for j in range(3):
circles.append(Circle(m[j], r, fc=(1, 0, 0, 0), ec=ec))
lines.append([p[j], e[j]])
if infofunc is not None:
infos.append(infofunc(i))
infos.append(infofunc(i))
if len(circles) == 0:
return None, None
lc = LineCollection(lines, color=color, picker=picker, linewidths=linewidth, **kwargs)
lc.info = infos
pc = PatchCollection(circles, match_original=True, picker=picker, linewidth=linewidth, **kwargs)
pc.info = infos
if cmap is not None:
z_duplicated = np.repeat(z.values, 3)
lc.set_cmap(cmap)
lc.set_norm(norm)
if clim is not None:
lc.set_clim(clim)
lc.set_array(np.ma.masked_invalid(z_duplicated))
lc.has_colormap = plot_colormap
lc.cbar_title = cbar_title
return lc, pc
def var_trace_plot(var_trace, var_type):
assert var_trace.ndim == 3
""" initialize the plotting params. """
num_of_xcoords = var_trace.shape[0]
num_of_line = var_trace.shape[1]
num_of_col = var_trace.shape[2]
x_coord_1st = 0 if var_type != 'gradients' else 1
x_coords = range(x_coord_1st, num_of_xcoords + x_coord_1st)
fig, axes = plt.subplots(1,
num_of_col,
figsize=(8, 4),
sharex=True,
sharey=True)
fig.subplots_adjust(bottom=0.13)
# else:
# fig, axes = plt.subplots(
# 2, 1, figsize=(11,22), sharex=True, sharey=True)
axes = np.array(axes).flatten()
""" determine the var type and plot. """
for case in switch(var_type):
if case('weights') \
or case('gradients'):
legend = [
r'$w(%s)$' % term for term in \
get_terms_expression(self.nn.fp_transformation)
] + [r'$bias$'] if case('weights') \
else [
r'$\partial{J}\ / \partial{w_{%sj}}$' % i
for i in range(num_of_line - 1)
] + [r'$\partial{J}\ / \partial{bias_j}$']
""" stack `W` and `b` together can avoid redundant \
plotting code in the following. """
# var_trace =np.vstack((trace_dict['W'], trace_dict['b']))
for dim_j in range(num_of_col):
for i in range(num_of_line):
axes[dim_j].plot(x_coords,
var_trace[:, i, dim_j],
linestyle='-',
marker='')
""" `legend` should be calculated finally here when \
var type is `weights`. """
if case('gradients'):
axes[dim_j].legend(
[_.replace('j', str(dim_j)) for _ in legend])
else:
axes[dim_j].legend(legend)
axes[dim_j].set_title(r'dimension %d of new space' %
dim_j,
y=1.05)
axes[dim_j].set_ylabel(var_type)
axes[dim_j].set_xlabel('step')
axes[dim_j].set_xticks(
[_ for _ in x_coords if _ % 2 == 0])
plt.show()
break
if case('loss'):
""" TODO: loss-element-wise & loss-dim-wise """
lc_s, annot_s = [], []
for dim_j in range(num_of_col):
"""
fisrt add x coordinates (`stack`), e.g:
[0.2 [0.3 ____\ [(0,0.2) [(0,0.3)
0.6], -0.4] / (1,0.6)], -(1,0.4)]
then permute dims to adapt to the input format of \
`LineCollection` (`transpose`).
"""
lines = np.transpose(
np.stack([
np.broadcast_to(
x_coords, (num_of_line, num_of_xcoords)).T,
var_trace[:, :, dim_j]
],
axis=2), (1, 0, 2))
""" split according to the type of class. """
lc_s_dim_j = []
cmaps = self.plt_cfg.vt_kwargs['cmaps']
assert len(cmaps) == \
self.nn.config.num_of_class
for cls_j in range(self.nn.config.num_of_class):
"""
`ls_cls_j`: lines of class j
`sample_size_j`: sample size of class j
"""
ls_cls_j = lines[self._indices(cls_j)]
sample_size_j = len(ls_cls_j)
""" attention the two meanings of `j` """
lc = LineCollection(ls_cls_j, cmap=cmaps[cls_j])
""" set the line colors """
lc.set_array(
np.linspace(int(sample_size_j * 0.2),
int(sample_size_j * 0.8),
sample_size_j))
lc.set_clim(0, sample_size_j)
axes[dim_j].add_collection(lc)
lc_s_dim_j.append(lc)
axes[dim_j].plot(
x_coords,
np.mean(ls_cls_j[:, :, 1], 0),
linestyle='-',
color=self.plt_cfg.colors[cls_j], \
linewidth=1.5, marker=''
)
axes[dim_j].autoscale()
axes[dim_j].set_title(r'dimension %d of new space' %
dim_j,
y=1.05)
axes[dim_j].set_ylabel('loss per sample')
axes[dim_j].set_xlabel('step')
axes[dim_j].set_xticks(
[_ for _ in x_coords if _ % 2 == 0])
axes[dim_j].legend(handles=self.plt_cfg.legend)
lc_s.append(lc_s_dim_j)
""" ref temporarily: https://stackoverflow.com/questions/7908636/possible-to-make-labels-appear-when-hovering-over-a-point-in-matplotlib """
annot = axes[dim_j].annotate(
"",
xy=(0, 0),
xytext=(4, 4),
textcoords="offset points",
bbox=dict(boxstyle="round", fc="c"),
)
annot.set_visible(True)
annot_s.append(annot)
def hover(event):
if event.inaxes:
""" idx of target ax. """
idx = list(axes).index(event.inaxes)
""" retrieve lcs of all classes in the target ax.\
"""
for j, lc in enumerate(lc_s[idx]):
is_contain, idx_dict = lc.contains(event)
if is_contain:
annot_s[idx].xy = event.xdata, event.ydata
""" ref temporarily: https://stackoverflow.com/questions/9647202/ordinal-numbers-replacement """
""" find the sample size of `j-1` class. \
"""
line_quantity_of_pre_classes = \
np.sum(np.argmax(self.y_s, 1) < j) \
if j > 0 else 0
annot_s[idx].set_text(
'%s sample' % \
ordinal(idx_dict['ind'][0] +
line_quantity_of_pre_classes)
)
annot_s[idx].get_bbox_patch(). \
set_alpha(0.4)
annot_s[idx].set_visible(True)
fig.canvas.draw_idle()
return
for _ in annot_s:
_.set_visible(False)
fig.canvas.draw_idle()
fig.canvas.mpl_connect("motion_notify_event", hover)
plt.show()
break
elif case('default'):
raise AssertionError('invalid var type!')
def draw_edges(G,
pos,
ax,
edgelist=None,
width=1.0,
width_adjuster=50,
edge_color='k',
style='solid',
alpha=None,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
traversal_weight=1.0,
edge_delengthify=0.15,
arrows=True,
label=None,
zorder=1,
**kwds):
"""
Code cleaned-up version of networkx.draw_networkx_edges
New args:
width_adjuster - the line width is generated from the weight if present, use this adjuster to thicken the lines (multiply)
"""
if edgelist is None:
edgelist = G.edges()
if not edgelist or len(edgelist) == 0: # no edges!
return None
# set edge positions
edge_pos = [(pos[e[0]], pos[e[1]]) for e in edgelist]
new_ep = []
for e in edge_pos:
x, y = e[0]
dx, dy = e[1]
# Get edge length
elx = (dx - x) * edge_delengthify
ely = (dy - y) * edge_delengthify
x += elx
y += ely
dx -= elx
dy -= ely
new_ep.append(((x, y), (dx, dy)))
edge_pos = numpy.asarray(new_ep)
if not cb.iterable(width):
#print [G.get_edge_data(n[0], n[1])['weight'] for n in edgelist]
# see if I can find an edge attribute:
if 'weight' in G.get_edge_data(edgelist[0][0],
edgelist[0][1]): # Test an edge
lw = [
0.5 +
((G.get_edge_data(n[0], n[1])['weight'] - traversal_weight) *
width_adjuster) for n in edgelist
]
else:
lw = (width, )
else:
lw = width
if not is_string_like(edge_color) and cb.iterable(edge_color) and len(
edge_color) == len(edge_pos):
if numpy.alltrue([cb.is_string_like(c) for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple(
[colorConverter.to_rgba(c, alpha) for c in edge_color])
elif numpy.alltrue([not cb.is_string_like(c) for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if numpy.alltrue(
[cb.iterable(c) and len(c) in (3, 4) for c in edge_color]):
edge_colors = tuple(edge_color)
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError(
'edge_color must consist of either color names or numbers')
else:
if is_string_like(edge_color) or len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError(
'edge_color must be a single color or list of exactly m colors where m is the number or edges'
)
edge_collection = LineCollection(edge_pos,
colors=edge_colors,
linewidths=lw,
antialiaseds=(1, ),
linestyle=style,
transOffset=ax.transData,
zorder=zorder)
edge_collection.set_label(label)
ax.add_collection(edge_collection)
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
if edge_colors is None:
if edge_cmap is not None:
assert (isinstance(edge_cmap, Colormap))
edge_collection.set_array(numpy.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
# update view
'''
minx = numpy.amin(numpy.ravel(edge_pos[:,:,0]))
maxx = numpy.amax(numpy.ravel(edge_pos[:,:,0]))
miny = numpy.amin(numpy.ravel(edge_pos[:,:,1]))
maxy = numpy.amax(numpy.ravel(edge_pos[:,:,1]))
w = maxx-minx
h = maxy-miny
padx, pady = 0.05*w, 0.05*h
corners = (minx-padx, miny-pady), (maxx+padx, maxy+pady)
ax.update_datalim(corners)
ax.autoscale_view()
'''
return (edge_collection)
def create_line_collection(net,
lines=None,
line_geodata=None,
bus_geodata=None,
use_bus_geodata=False,
infofunc=None,
cmap=None,
norm=None,
picker=False,
z=None,
cbar_title="Line Loading [%]",
clim=None,
**kwargs):
"""
Creates a matplotlib line collection of pandapower lines.
Input:
**net** (pandapowerNet) - The pandapower network
OPTIONAL:
**lines** (list, None) - The lines for which the collections are created. If None, all lines
in the network are considered.
**line_geodata** (DataFrame, None) - coordinates to use for plotting If None,
net["line_geodata"] is used
**infofunc** (function, None) - infofunction for the patch element
**kwargs - key word arguments are passed to the patch function
OUTPUT:
**lc** - line collection
"""
lines = net.line.index.tolist() if lines is None else list(lines)
if len(lines) == 0:
return None
if line_geodata is None:
line_geodata = net["line_geodata"]
if bus_geodata is None:
bus_geodata = net["bus_geodata"]
if len(lines) == 0:
return None
if use_bus_geodata:
data = [
([(bus_geodata.at[a, "x"], bus_geodata.at[a, "y"]),
(bus_geodata.at[b, "x"], bus_geodata.at[b, "y"])],
infofunc(line) if infofunc else [])
for line, (a,
b) in net.line.loc[lines,
["from_bus", "to_bus"]].iterrows()
if a in bus_geodata.index.values and b in bus_geodata.index.values
]
else:
data = [(line_geodata.loc[line, "coords"],
infofunc(line) if infofunc else []) for line in lines
if line in line_geodata.index.values]
if len(data) == 0:
return None
data, info = list(zip(*data))
# This would be done anyways by matplotlib - doing it explicitly makes it a) clear and
# b) prevents unexpected behavior when observing colors being "none"
lc = LineCollection(data, picker=picker, **kwargs)
lc.line_indices = np.array(lines)
if cmap:
if z is None:
z = net.res_line.loading_percent.loc[lines]
lc.set_cmap(cmap)
lc.set_norm(norm)
if clim is not None:
lc.set_clim(clim)
lc.set_array(np.array(z))
lc.has_colormap = True
lc.cbar_title = cbar_title
lc.info = info
return lc
def create_line_collection(net, lines=None, line_geodata=None, bus_geodata=None,
use_bus_geodata=False, infofunc=None,
cmap=None, norm=None, picker=False, z=None,
cbar_title="Line Loading [%]", clim=None, **kwargs):
"""
Creates a matplotlib line collection of pandapower lines.
Input:
**net** (pandapowerNet) - The pandapower network
OPTIONAL:
**lines** (list, None) - The lines for which the collections are created. If None, all lines
in the network are considered.
**line_geodata** (DataFrame, None) - coordinates to use for plotting. If None,
net["line_geodata"] is used
**bus_geodata** (DataFrame, None) - coordinates to use for plotting
If None, net["bus_geodata"] is used
**use_bus_geodata** (bool, False) - Defines whether bus or line geodata are used.
**infofunc** (function, None) - infofunction for the patch element
**cmap** - colormap for the patch colors
**norm** (matplotlib norm object, None) - matplotlib norm object
**picker** (bool, False) - picker argument passed to the patch collection
**z** (array, None) - array of bus voltage magnitudes for colormap. Used in case of given
cmap. If None net.res_bus.vm_pu is used.
**cbar_title** (str, "Bus Voltage [pu]") - colormap bar title in case of given cmap
**clim** (tuple of floats, None) - setting the norm limits for image scaling
**kwargs - key word arguments are passed to the patch function
OUTPUT:
**lc** - line collection
"""
if use_bus_geodata:
linetab = net.line if lines is None else net.line.loc[lines]
lines = net.line.index.tolist() if lines is None else list(lines)
if len(lines) == 0:
return None
if line_geodata is None:
line_geodata = net["line_geodata"]
if bus_geodata is None:
bus_geodata = net["bus_geodata"]
if len(lines) == 0:
return None
lines_with_geo = []
if use_bus_geodata:
data = []
buses_with_geodata = bus_geodata.index.values
bg_dict = bus_geodata.to_dict() #transforming to dict to make lookup faster
for line, fb, tb in zip(linetab.index, linetab.from_bus.values, linetab.to_bus.values):
if fb in buses_with_geodata and tb in buses_with_geodata:
lines_with_geo.append(line)
data.append(([(bg_dict["x"][fb], bg_dict["y"][fb]),
(bg_dict["x"][tb], bg_dict["y"][tb])],
infofunc(line) if infofunc else[]))
lines_without_geo = set(lines)-set(lines_with_geo)
if lines_without_geo:
logger.warning("Could not plot lines %s. Bus geodata is missing for those lines!"
% lines_without_geo)
else:
data = []
for line in lines:
if line in line_geodata.index.values:
lines_with_geo.append(line)
data.append((line_geodata.loc[line, "coords"], infofunc(line) if infofunc else []))
lines_without_geo = set(lines)-set(lines_with_geo)
if len(lines_without_geo) > 0:
logger.warning("Could not plot lines %s. Line geodata is missing for those lines!"
% lines_without_geo)
if len(data) == 0:
return None
data, info = list(zip(*data))
# This would be done anyways by matplotlib - doing it explicitly makes it a) clear and
# b) prevents unexpected behavior when observing colors being "none"
lc = LineCollection(data, picker=picker, **kwargs)
lc.line_indices = np.array(lines_with_geo)
if cmap is not None:
if z is None:
z = net.res_line.loading_percent.loc[lines_with_geo]
lc.set_cmap(cmap)
lc.set_norm(norm)
if clim is not None:
lc.set_clim(clim)
lc.set_array(np.array(z))
lc.has_colormap = True
lc.cbar_title = cbar_title
lc.info = info
return lc
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=None,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
**kwds):
"""Draw the edges of the graph G
This draws only the edges of the graph G.
pos is a dictionary keyed by vertex with a two-tuple
of x-y positions as the value.
See networkx.layout for functions that compute node positions.
edgelist is an optional list of the edges in G to be drawn.
If provided, only the edges in edgelist will be drawn.
edgecolor can be a list of matplotlib color letters such as 'k' or
'b' that lists the color of each edge; the list must be ordered in
the same way as the edge list. Alternatively, this list can contain
numbers and those number are mapped to a color scale using the color
map edge_cmap. Finally, it can also be a list of (r,g,b) or (r,g,b,a)
tuples, in which case these will be used directly to color the edges. If
the latter mode is used, you should not provide a value for alpha, as it
would be applied globally to all lines.
For directed graphs, "arrows" (actually just thicker stubs) are drawn
at the head end. Arrows can be turned off with keyword arrows=False.
See draw_networkx for the list of other optional parameters.
"""
try:
import matplotlib
import matplotlib.pylab as pylab
import matplotlib.cbook as cb
from matplotlib.colors import colorConverter,Colormap
from matplotlib.collections import LineCollection
import numpy
except ImportError:
raise ImportError, "Matplotlib required for draw()"
except RuntimeError:
pass # unable to open display
if ax is None:
ax=pylab.gca()
if edgelist is None:
edgelist=G.edges()
if not edgelist or len(edgelist)==0: # no edges!
return None
# set edge positions
edge_pos=numpy.asarray([(pos[e[0]],pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width,)
else:
lw = width
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color)==len(edge_pos):
if numpy.alltrue([cb.is_string_like(c)
for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple([colorConverter.to_rgba(c,alpha)
for c in edge_color])
elif numpy.alltrue([not cb.is_string_like(c)
for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if numpy.alltrue([cb.iterable(c) and len(c) in (3,4)
for c in edge_color]):
edge_colors = tuple(edge_color)
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must consist of either color names or numbers')
else:
if len(edge_color)==1:
edge_colors = ( colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError('edge_color must be a single color or list of exactly m colors where m is the number or edges')
edge_collection = LineCollection(edge_pos,
colors = edge_colors,
linewidths = lw,
antialiaseds = (1,),
linestyle = style,
transOffset = ax.transData,
)
# Note: there was a bug in mpl regarding the handling of alpha values for
# each line in a LineCollection. It was fixed in matplotlib in r7184 and
# r7189 (June 6 2009). We should then not set the alpha value globally,
# since the user can instead provide per-edge alphas now. Only set it
# globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
# need 0.87.7 or greater for edge colormaps
mpl_version=matplotlib.__version__
if mpl_version.endswith('svn'):
mpl_version=matplotlib.__version__[0:-3]
if mpl_version.endswith('pre'):
mpl_version=matplotlib.__version__[0:-3]
if map(int,mpl_version.split('.'))>=[0,87,7]:
if edge_colors is None:
if edge_cmap is not None: assert(isinstance(edge_cmap, Colormap))
edge_collection.set_array(numpy.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
pylab.sci(edge_collection)
# else:
# sys.stderr.write(\
# """matplotlib version >= 0.87.7 required for colormapped edges.
# (version %s detected)."""%matplotlib.__version__)
# raise UserWarning(\
# """matplotlib version >= 0.87.7 required for colormapped edges.
# (version %s detected)."""%matplotlib.__version__)
arrow_collection=None
if G.is_directed() and arrows:
# a directed graph hack
# draw thick line segments at head end of edge
# waiting for someone else to implement arrows that will work
arrow_colors = ( colorConverter.to_rgba('k', alpha), )
a_pos=[]
p=1.0-0.25 # make head segment 25 percent of edge length
for src,dst in edge_pos:
x1,y1=src
x2,y2=dst
dx=x2-x1 # x offset
dy=y2-y1 # y offset
d=numpy.sqrt(float(dx**2+dy**2)) # length of edge
if d==0: # source and target at same position
continue
if dx==0: # vertical edge
xa=x2
ya=dy*p+y1
if dy==0: # horizontal edge
ya=y2
xa=dx*p+x1
else:
theta=numpy.arctan2(dy,dx)
xa=p*d*numpy.cos(theta)+x1
ya=p*d*numpy.sin(theta)+y1
a_pos.append(((xa,ya),(x2,y2)))
arrow_collection = LineCollection(a_pos,
colors = arrow_colors,
linewidths = [4*ww for ww in lw],
antialiaseds = (1,),
transOffset = ax.transData,
)
# update view
minx = numpy.amin(numpy.ravel(edge_pos[:,:,0]))
maxx = numpy.amax(numpy.ravel(edge_pos[:,:,0]))
miny = numpy.amin(numpy.ravel(edge_pos[:,:,1]))
maxy = numpy.amax(numpy.ravel(edge_pos[:,:,1]))
w = maxx-minx
h = maxy-miny
padx, pady = 0.05*w, 0.05*h
corners = (minx-padx, miny-pady), (maxx+padx, maxy+pady)
ax.update_datalim( corners)
ax.autoscale_view()
edge_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(edge_collection)
if arrow_collection:
arrow_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(arrow_collection)
return edge_collection
class SunPlotPy(wx.Frame, Spatial, Grid ):
"""
The main frame of the application
"""
title = 'sunplot(py)'
# Plotting options
autoclim=True
showedges=False
bgcolor='k'
textcolor='w'
cmap='RdBu'
particlesize = 1.8
particlecolor = 'm'
# other flags
collectiontype='cells'
oldcollectiontype='cells'
#
tindex=0
depthlevs = [0., 10., 100., 200., 300., 400., 500.,\
1000.,2000.,3000.,4000.,5000]
_FillValue=999999
def __init__(self):
wx.Frame.__init__(self, None, -1, self.title)
self.create_menu()
self.create_status_bar()
self.create_main_panel()
#self.draw_figure()
def create_menu(self):
self.menubar = wx.MenuBar()
###
# File Menu
###
menu_file = wx.Menu()
# Load a hydro output file
m_expt = menu_file.Append(-1, "&Open file\tCtrl-O", "Open netcdf file")
self.Bind(wx.EVT_MENU, self.on_open_file, m_expt)
# Load a grid file
m_grid = menu_file.Append(-1, "&Load grid\tCtrl-G", "Load SUNTANS grid from folder")
self.Bind(wx.EVT_MENU, self.on_load_grid, m_grid)
# Load a particle file
m_part = menu_file.Append(-1, "&Load PTM file\tCtrl-Shift-P", "Load a PTM file")
self.Bind(wx.EVT_MENU, self.on_load_ptm, m_part)
# Save current scene as an animation
m_anim = menu_file.Append(-1,"&Save animation of current scene\tCtrl-S","Save animation")
self.Bind(wx.EVT_MENU, self.on_save_anim, m_anim)
# Save the current figure
m_prin = menu_file.Append(-1,"&Print current scene\tCtrl-P","Save figure")
self.Bind(wx.EVT_MENU, self.on_save_fig, m_prin)
menu_file.AppendSeparator()
# Exit
m_exit = menu_file.Append(-1, "E&xit\tCtrl-X", "Exit")
self.Bind(wx.EVT_MENU, self.on_exit, m_exit)
###
# Tools menu
###
menu_tools = wx.Menu()
m_gridstat = menu_tools.Append(-1, "&Plot grid size statistics", "SUNTANS grid size")
self.Bind(wx.EVT_MENU, self.on_plot_gridstat, m_gridstat)
m_countcells = menu_tools.Append(-1, "&Count # grid cells", "Grid cell count")
self.Bind(wx.EVT_MENU, self.on_count_cells, m_countcells)
m_overlaybathy = menu_tools.Append(-1, "&Overlay depth contours", "Depth overlay")
self.Bind(wx.EVT_MENU, self.on_overlay_bathy, m_overlaybathy)
###
# Help Menu
###
menu_help = wx.Menu()
m_about = menu_help.Append(-1, "&About\tF1", "About the demo")
self.Bind(wx.EVT_MENU, self.on_about, m_about)
# Add all of the menu bars
self.menubar.Append(menu_file, "&File")
self.menubar.Append(menu_tools, "&Tools")
self.menubar.Append(menu_help, "&Help")
self.SetMenuBar(self.menubar)
def create_main_panel(self):
""" Creates the main panel with all the controls on it:
* mpl canvas
* mpl navigation toolbar
* Control panel for interaction
"""
self.panel = wx.Panel(self)
# Create the mpl Figure and FigCanvas objects.
# 5x4 inches, 100 dots-per-inch
#
self.dpi = 100
#self.fig = Figure((7.0, 6.0), dpi=self.dpi,facecolor=self.bgcolor)
self.fig = Figure((7.0, 6.0), dpi=self.dpi)
self.canvas = FigCanvas(self.panel, -1, self.fig)
# Since we have only one plot, we can use add_axes
# instead of add_subplot, but then the subplot
# configuration tool in the navigation toolbar wouldn't
# work.
#
self.axes = self.fig.add_subplot(111)
#SetAxColor(self.axes,self.textcolor,self.bgcolor)
# Bind the 'pick' event for clicking on one of the bars
#
#self.canvas.mpl_connect('pick_event', self.on_pick)
########
# Create widgets
########
self.variable_list = wx.ComboBox(
self.panel,
size=(200,-1),
choices=['Select a variable...'],
style=wx.CB_READONLY)
self.variable_list.Bind(wx.EVT_COMBOBOX, self.on_select_variable)
self.time_list = wx.ComboBox(
self.panel,
size=(200,-1),
choices=['Select a time step...'],
style=wx.CB_READONLY)
self.time_list.Bind(wx.EVT_COMBOBOX, self.on_select_time)
self.depthlayer_list = wx.ComboBox(
self.panel,
size=(200,-1),
choices=['Select a vertical layer...'],
style=wx.CB_READONLY)
self.depthlayer_list.Bind(wx.EVT_COMBOBOX, self.on_select_depth)
self.show_edge_check = wx.CheckBox(self.panel, -1,
"Show Edges",
style=wx.ALIGN_RIGHT)
self.show_edge_check.Bind(wx.EVT_CHECKBOX, self.on_show_edges)
if USECMOCEAN:
cmaps=[]
for cmap in cm.cmapnames:
cmaps.append(cmap)
cmaps.append(cmap+'_r') # Add all reverse map options
else:
# Use matplotlib standard
cmaps = matplotlib.cm.datad.keys()
cmaps.sort()
self.colormap_list = wx.ComboBox(
self.panel,
size=(100,-1),
choices=cmaps,
style=wx.CB_READONLY)
self.colormap_list.Bind(wx.EVT_COMBOBOX, self.on_select_cmap)
self.colormap_label = wx.StaticText(self.panel, -1,"Colormap:")
self.clim_check = wx.CheckBox(self.panel, -1,
"Manual color limits ",
style=wx.ALIGN_RIGHT)
self.clim_check.Bind(wx.EVT_CHECKBOX, self.on_clim_check)
self.climlow = wx.TextCtrl(
self.panel,
size=(100,-1),
style=wx.TE_PROCESS_ENTER)
self.climlow.Bind(wx.EVT_TEXT_ENTER, self.on_climlow)
self.climhigh = wx.TextCtrl(
self.panel,
size=(100,-1),
style=wx.TE_PROCESS_ENTER)
self.climhigh.Bind(wx.EVT_TEXT_ENTER, self.on_climhigh)
# Labels
self.variable_label = wx.StaticText(self.panel, -1,"Variable:",size=(200,-1))
self.time_label = wx.StaticText(self.panel, -1,"Time step:",size=(200,-1))
self.depth_label = wx.StaticText(self.panel, -1,"Vertical level:",size=(200,-1))
# Create the navigation toolbar, tied to the canvas
#
self.toolbar = NavigationToolbar(self.canvas)
#self.toolbar.toolitems[8][3]='my_save_fig'
#def my_save_fig(self,*args):
# print 'saving figure'
# return "break"
#########
# Layout with box sizers
#########
self.vbox = wx.BoxSizer(wx.VERTICAL)
self.vbox.Add(self.canvas, 1, wx.LEFT | wx.TOP | wx.GROW)
self.vbox.Add(self.toolbar, 0, wx.EXPAND)
self.vbox.AddSpacer(10)
#self.vbox.Add((-1,25))
flags = wx.ALIGN_LEFT | wx.ALL | wx.ALIGN_CENTER_VERTICAL
self.hbox0 = wx.BoxSizer(wx.HORIZONTAL)
self.hbox0.Add(self.show_edge_check, 0, border=10, flag=flags)
self.hbox0.Add(self.colormap_label, 0, border=10, flag=flags)
self.hbox0.Add(self.colormap_list, 0, border=10, flag=flags)
self.hbox0.Add(self.clim_check, 0, border=10, flag=flags)
self.hbox0.Add(self.climlow, 0, border=10, flag=flags)
self.hbox0.Add(self.climhigh, 0, border=10, flag=flags)
self.vbox.AddSpacer(5)
self.hbox1 = wx.BoxSizer(wx.HORIZONTAL)
self.hbox1.Add(self.variable_label, 0, border=10, flag=flags)
self.hbox1.Add(self.time_label, 0, border=10, flag=flags)
self.hbox1.Add(self.depth_label, 0, border=10, flag=flags)
self.vbox.AddSpacer(5)
self.hbox2 = wx.BoxSizer(wx.HORIZONTAL)
self.hbox2.Add(self.variable_list, 0, border=10, flag=flags)
self.hbox2.Add(self.time_list, 0, border=10, flag=flags)
self.hbox2.Add(self.depthlayer_list, 0, border=10, flag=flags)
self.vbox.Add(self.hbox1, 0, flag = wx.ALIGN_LEFT | wx.TOP)
self.vbox.Add(self.hbox2, 0, flag = wx.ALIGN_LEFT | wx.TOP)
self.vbox.Add(self.hbox0, 0, flag = wx.ALIGN_LEFT | wx.TOP)
self.panel.SetSizer(self.vbox)
self.vbox.Fit(self)
##########
# Event functions
##########
def create_figure(self):
"""
Creates the figure
"""
# Find the colorbar limits if unspecified
if self.autoclim:
self.clim = [self.data.min(),self.data.max()]
self.climlow.SetValue('%3.1f'%self.clim[0])
self.climhigh.SetValue('%3.1f'%self.clim[1])
if self.__dict__.has_key('collection'):
#self.collection.remove()
self.axes.collections.remove(self.collection)
else:
# First call - set the axes limits
self.axes.set_aspect('equal')
self.axes.set_xlim(self.xlims)
self.axes.set_ylim(self.ylims)
if self.collectiontype=='cells':
self.collection = PolyCollection(self.xy,cmap=self.cmap)
self.collection.set_array(np.array(self.data[:]))
if not self.showedges:
self.collection.set_edgecolors(self.collection.to_rgba(np.array((self.data[:]))))
elif self.collectiontype=='edges':
xylines = [self.xp[self.edges],self.yp[self.edges]]
linesc = [zip(xylines[0][ii,:],xylines[1][ii,:]) for ii in range(self.Ne)]
self.collection = LineCollection(linesc,array=np.array(self.data[:]),cmap=self.cmap)
self.collection.set_clim(vmin=self.clim[0],vmax=self.clim[1])
self.axes.add_collection(self.collection)
self.title=self.axes.set_title(self.genTitle(),color=self.textcolor)
self.axes.set_xlabel('Easting [m]')
self.axes.set_ylabel('Northing [m]')
# create a colorbar
if not self.__dict__.has_key('cbar'):
self.cbar = self.fig.colorbar(self.collection)
#SetAxColor(self.cbar.ax.axes,self.textcolor,self.bgcolor)
else:
#pass
print 'Updating colorbar...'
#self.cbar.check_update(self.collection)
self.cbar.on_mappable_changed(self.collection)
self.canvas.draw()
def update_figure(self):
if self.autoclim:
self.clim = [self.data.min(),self.data.max()]
self.climlow.SetValue('%3.1f'%self.clim[0])
self.climhigh.SetValue('%3.1f'%self.clim[1])
else:
self.clim = [float(self.climlow.GetValue()),\
float(self.climhigh.GetValue())]
# check whether it is cell or edge type
if self.hasDim(self.variable,self.griddims['Ne']):
self.collectiontype='edges'
elif self.hasDim(self.variable,self.griddims['Nc']):
self.collectiontype='cells'
# Create a new figure if the variable has gone from cell to edge of vice
# versa
if not self.collectiontype==self.oldcollectiontype:
self.create_figure()
self.oldcollectiontype=self.collectiontype
self.collection.set_array(np.array(self.data[:]))
self.collection.set_clim(vmin=self.clim[0],vmax=self.clim[1])
# Cells only
if self.collectiontype=='cells':
if not self.showedges:
self.collection.set_edgecolors(self.collection.to_rgba(np.array((self.data[:]))))
else:
self.collection.set_edgecolors('k')
self.collection.set_linewidths(0.2)
# Update the title
self.title=self.axes.set_title(self.genTitle(),color=self.textcolor)
#Update the colorbar
self.cbar.update_normal(self.collection)
# redraw the figure
self.canvas.draw()
def on_pick(self, event):
# The event received here is of the type
# matplotlib.backend_bases.PickEvent
#
# It carries lots of information, of which we're using
# only a small amount here.
#
box_points = event.artist.get_bbox().get_points()
msg = "You've clicked on a bar with coords:\n %s" % box_points
dlg = wx.MessageDialog(
self,
msg,
"Click!",
wx.OK | wx.ICON_INFORMATION)
dlg.ShowModal()
dlg.Destroy()
def on_select_variable(self, event):
vname = event.GetString()
self.flash_status_message("Selecting variable: %s"%vname)
# update the spatial object and load the data
self.variable = vname
self.loadData(variable=self.variable)
# Check if the variable has a depth coordinate
depthstr = ['']
# If so populate the vertical layer box
if self.hasDim(self.variable,self.griddims['Nk']):
depthstr = ['%3.1f'%self.z_r[k] for k in range(self.Nkmax)]
depthstr += ['surface','seabed']
elif self.hasDim(self.variable,'Nkw'):
depthstr = ['%3.1f'%self.z_w[k] for k in range(self.Nkmax+1)]
self.depthlayer_list.SetItems(depthstr)
# Update the plot
self.update_figure()
def on_select_time(self, event):
self.tindex = event.GetSelection()
# Update the object time index and reload the data
if self.plot_type=='hydro':
if not self.tstep==self.tindex:
self.tstep=self.tindex
self.loadData()
self.flash_status_message("Selecting variable: %s..."%event.GetString())
# Update the plot
self.update_figure()
elif self.plot_type=='particles':
self.PTM.plot(self.tindex,ax=self.axes,\
xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
self.canvas.draw()
def on_select_depth(self, event):
kindex = event.GetSelection()
if not self.klayer[0]==kindex:
# Check if its the seabed or surface value
if kindex>=self.Nkmax:
kindex=event.GetString()
self.klayer = [kindex]
self.loadData()
self.flash_status_message("Selecting depth: %s..."%event.GetString())
# Update the plot
self.update_figure()
def on_open_file(self, event):
file_choices = "SUNTANS NetCDF (*.nc)|*.nc*|UnTRIM NetCDF (*.nc)|*.nc*|All Files (*.*)|*.*"
dlg = wx.FileDialog(
self,
message="Open SUNTANS file...",
defaultDir=os.getcwd(),
defaultFile="",
wildcard=file_choices,
style= wx.FD_MULTIPLE)
if dlg.ShowModal() == wx.ID_OK:
self.plot_type='hydro'
path = dlg.GetPaths()
# Initialise the class
if dlg.GetFilterIndex() == 0 or dlg.GetFilterIndex() > 1: #SUNTANS
self.flash_status_message("Opening SUNTANS file: %s" % path)
try:
Spatial.__init__(self, path, _FillValue=self._FillValue)
except:
Spatial.__init__(self, path, _FillValue=-999999)
startvar='dv'
if dlg.GetFilterIndex()==1: #UnTRIM
self.flash_status_message("Opening UnTRIMS file: %s" % path)
#Spatial.__init__(self,path,gridvars=untrim_gridvars,griddims=untrim_griddims)
UNTRIMSpatial.__init__(self,path)
startvar='Mesh2_face_depth'
# Populate the drop down menus
vnames = self.listCoordVars()
self.variable_list.SetItems(vnames)
# Update the time drop down list
if self.__dict__.has_key('time'):
self.timestr = [datetime.strftime(tt,'%d-%b-%Y %H:%M:%S') for tt in self.time]
else:
# Assume that it is a harmonic-type file
self.timestr = self.nc.Constituent_Names.split()
self.time_list.SetItems(self.timestr)
# Draw the depth
if startvar in vnames:
self.variable=startvar
self.loadData()
self.create_figure()
def on_load_grid(self, event):
dlg = wx.DirDialog(
self,
message="Open SUNTANS grid from folder...",
defaultPath=os.getcwd(),
style= wx.DD_DEFAULT_STYLE)
if dlg.ShowModal() == wx.ID_OK:
path = dlg.GetPath()
# Initialise the class
self.flash_status_message("Opening SUNTANS grid from folder: %s" % path)
Grid.__init__(self,path)
# Plot the Grid
if self.__dict__.has_key('collection'):
self.axes.collections.remove(self.collection)
self.axes,self.collection = self.plotmesh(ax=self.axes,edgecolors='y')
# redraw the figure
self.canvas.draw()
def on_load_ptm(self, event):
file_choices = "PTM NetCDF (*.nc)|*.nc|PTM Binary (*_bin.out)|*_bin.out|All Files (*.*)|*.*"
dlg = wx.FileDialog(
self,
message="Open PTM file...",
defaultDir=os.getcwd(),
defaultFile="",
wildcard=file_choices,
style= wx.FD_MULTIPLE)
if dlg.ShowModal() == wx.ID_OK:
self.plot_type = 'particles'
path = dlg.GetPath()
# Initialise the class
if dlg.GetFilterIndex() == 0: #SUNTANS
self.flash_status_message("Opening PTM netcdf file: %s" % path)
self.PTM = PtmNC(path)
elif dlg.GetFilterIndex() == 1: #PTM
self.flash_status_message("Opening PTM binary file: %s" % path)
self.PTM = PtmBin(path)
self.Nt = self.PTM.nt
# Update the time drop down list
self.timestr = [datetime.strftime(tt,'%d-%b-%Y %H:%M:%S') for tt in self.PTM.time]
self.time_list.SetItems(self.timestr)
# Plot the first time step
if self.__dict__.has_key('xlims'):
self.PTM.plot(self.PTM.nt-1,ax=self.axes,xlims=self.xlims,\
ylims=self.ylims,color=self.particlecolor,\
fontcolor='w',markersize=self.particlesize)
else:
self.PTM.plot(self.PTM.nt-1,ax=self.axes,fontcolor='w',\
color=self.particlecolor,markersize=self.particlesize)
# redraw the figure
self.canvas.draw()
def on_show_edges(self,event):
sender=event.GetEventObject()
self.showedges = sender.GetValue()
# Update the figure
self.update_figure()
def on_clim_check(self,event):
sender=event.GetEventObject()
if sender.GetValue() == True:
self.autoclim=False
self.update_figure()
else:
self.autoclim=True
def on_climlow(self,event):
self.clim[0] = event.GetString()
#self.update_figure()
def on_climhigh(self,event):
self.clim[1] = event.GetString()
#self.update_figure()
def on_select_cmap(self,event):
self.cmap=event.GetString()
if USECMOCEAN:
self.collection.set_cmap(getattr(cm,self.cmap))
else:
self.collection.set_cmap(self.cmap)
# Update the figure
self.update_figure()
def on_save_fig(self,event):
"""
Save a figure of the current scene to a file
"""
file_choices = " (*.png)|*.png| (*.pdf)|*.pdf |(*.jpg)|*.jpg |(*.eps)|*eps "
filters=['.png','.pdf','.png','.png']
dlg = wx.FileDialog(
self,
message="Save figure to file...",
defaultDir=os.getcwd(),
defaultFile="",
wildcard=file_choices,
style= wx.FD_SAVE | wx.FD_OVERWRITE_PROMPT)
if dlg.ShowModal() == wx.ID_OK:
path = dlg.GetPath()
ext = filters[dlg.GetFilterIndex()]
if ext in path:
outfile=path
else:
outfile = path+ext
self.fig.savefig(outfile)
def on_save_anim(self,event):
"""
Save an animation of the current scene to a file
"""
file_choices = "Quicktime (*.mov)|*.mov| (*.gif)|*.gif| (*.avi)|*.avi |(*.mp4)|*.mp4 "
filters=['.mov','.gif','.avi','.mp4']
dlg = wx.FileDialog(
self,
message="Output animation file...",
defaultDir=os.getcwd(),
defaultFile="",
wildcard=file_choices,
style= wx.FD_SAVE | wx.FD_OVERWRITE_PROMPT)
if dlg.ShowModal() == wx.ID_OK:
path = dlg.GetPath()
ext = filters[dlg.GetFilterIndex()]
if ext in path:
outfile=path
else:
outfile = path+ext
self.flash_status_message("Saving figure to file: %s" %outfile)
self.flash_status_message("Saving animation to file: %s" %outfile)
# Create the animation
#self.tstep = range(self.Nt) # Use all time steps for animation
#self.animate(cbar=self.cbar,cmap=self.cmap,\
# xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
def initanim():
if not self.plot_type=='particles':
return (self.title, self.collection)
else:
return (self.PTM.title,self.PTM.p_handle)
def updateScalar(i):
if not self.plot_type=='particles':
self.tstep=[i]
self.loadData()
self.update_figure()
return (self.title,self.collection)
elif self.plot_type=='particles':
self.PTM.plot(i,ax=self.axes,\
xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
return (self.PTM.title,self.PTM.p_handle)
self.anim = animation.FuncAnimation(self.fig, \
updateScalar, init_func = initanim, frames=self.Nt, interval=50, blit=True)
if ext=='.gif':
self.anim.save(outfile,writer='imagemagick',fps=6)
elif ext=='.mp4':
print 'Saving html5 video...'
# Ensures html5 compatibility
self.anim.save(outfile,writer='mencoder',fps=6,\
bitrate=3600,extra_args=['-ovc','x264']) # mencoder options
#bitrate=3600,extra_args=['-vcodec','libx264'])
else:
self.anim.save(outfile,writer='mencoder',fps=6,bitrate=3600)
# Return the figure back to its status
del self.anim
self.tstep=self.tindex
if not self.plot_type=='particles':
self.loadData()
self.update_figure()
# Bring up a dialog box
dlg2= wx.MessageDialog(self, 'Animation complete.', "Done", wx.OK)
dlg2.ShowModal()
dlg2.Destroy()
def on_exit(self, event):
self.Destroy()
def on_about(self, event):
msg = """ SUNTANS NetCDF visualization tool
*Author: Matt Rayson
*Institution: Stanford University
*Created: October 2013
"""
dlg = wx.MessageDialog(self, msg, "About", wx.OK)
dlg.ShowModal()
dlg.Destroy()
def on_count_cells(self,eveny):
msg = "Total 3-D grid cells = %d"%(self.count_cells())
dlg = wx.MessageDialog(self, msg, "No. cells", wx.OK)
dlg.ShowModal()
dlg.Destroy()
def on_overlay_bathy(self,event):
# Plot depth contours
print 'Plotting contours...'
self.contourf(z=self.dv, clevs=self.depthlevs,\
ax=self.axes,\
filled=False, colors='0.5', linewidths=0.5, zorder=1e6)
print 'Done'
def on_plot_gridstat(self, event):
"""
Plot the grid size histogram in a new figure
"""
matplotlib.pyplot.figure()
self.plothist()
matplotlib.pyplot.show()
def create_status_bar(self):
self.statusbar = self.CreateStatusBar()
def flash_status_message(self, msg, flash_len_ms=1500):
self.statusbar.SetStatusText(msg)
self.timeroff = wx.Timer(self)
self.Bind(
wx.EVT_TIMER,
self.on_flash_status_off,
self.timeroff)
self.timeroff.Start(flash_len_ms, oneShot=True)
def on_flash_status_off(self, event):
self.statusbar.SetStatusText('')
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
**kwds):
"""Draw the edges of the graph G.
This draws only the edges of the graph G.
Parameters
----------
G : graph
A networkx graph
pos : dictionary
A dictionary with nodes as keys and positions as values.
Positions should be sequences of length 2.
edgelist : collection of edge tuples
Draw only specified edges(default=G.edges())
width : float, or array of floats
Line width of edges (default=1.0)
edge_color : color string, or array of floats
Edge color. Can be a single color format string (default='r'),
or a sequence of colors with the same length as edgelist.
If numeric values are specified they will be mapped to
colors using the edge_cmap and edge_vmin,edge_vmax parameters.
style : string
Edge line style (default='solid') (solid|dashed|dotted,dashdot)
alpha : float
The edge transparency (default=1.0)
edge_ cmap : Matplotlib colormap
Colormap for mapping intensities of edges (default=None)
edge_vmin,edge_vmax : floats
Minimum and maximum for edge colormap scaling (default=None)
ax : Matplotlib Axes object, optional
Draw the graph in the specified Matplotlib axes.
arrows : bool, optional (default=True)
For directed graphs, if True draw arrowheads.
label : [None| string]
Label for legend
Returns
-------
matplotlib.collection.LineCollection
`LineCollection` of the edges
Notes
-----
For directed graphs, "arrows" (actually just thicker stubs) are drawn
at the head end. Arrows can be turned off with keyword arrows=False.
Yes, it is ugly but drawing proper arrows with Matplotlib this
way is tricky.
Examples
--------
>>> G=nx.dodecahedral_graph()
>>> edges=nx.draw_networkx_edges(G,pos=nx.spring_layout(G))
Also see the NetworkX drawing examples at
http://networkx.github.io/documentation/latest/gallery.html
See Also
--------
draw()
draw_networkx()
draw_networkx_nodes()
draw_networkx_labels()
draw_networkx_edge_labels()
"""
try:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cbook as cb
from matplotlib.colors import colorConverter, Colormap
from matplotlib.collections import LineCollection
import numpy
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
# set edge positions
edge_pos = numpy.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width,)
else:
lw = width
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color) == len(edge_pos):
if numpy.alltrue([cb.is_string_like(c)
for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple([colorConverter.to_rgba(c, alpha)
for c in edge_color])
elif numpy.alltrue([not cb.is_string_like(c)
for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if numpy.alltrue([cb.iterable(c) and len(c) in (3, 4)
for c in edge_color]):
edge_colors = tuple(edge_color)
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must consist of either color names or numbers')
else:
if cb.is_string_like(edge_color) or len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError('edge_color must be a single color or list of exactly m colors where m is the number or edges')
edge_collection = LineCollection(edge_pos,
colors=edge_colors,
linewidths=lw,
antialiaseds=(1,),
linestyle=style,
transOffset = ax.transData,
)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
# Note: there was a bug in mpl regarding the handling of alpha values for
# each line in a LineCollection. It was fixed in matplotlib in r7184 and
# r7189 (June 6 2009). We should then not set the alpha value globally,
# since the user can instead provide per-edge alphas now. Only set it
# globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
if edge_colors is None:
if edge_cmap is not None:
assert(isinstance(edge_cmap, Colormap))
edge_collection.set_array(numpy.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
arrow_collection = None
if G.is_directed() and arrows:
# a directed graph hack
# draw thick line segments at head end of edge
# waiting for someone else to implement arrows that will work
arrow_colors = edge_colors
a_pos = []
p = 1.0-0.25 # make head segment 25 percent of edge length
for src, dst in edge_pos:
x1, y1 = src
x2, y2 = dst
dx = x2-x1 # x offset
dy = y2-y1 # y offset
d = numpy.sqrt(float(dx**2 + dy**2)) # length of edge
if d == 0: # source and target at same position
continue
if dx == 0: # vertical edge
xa = x2
ya = dy*p+y1
if dy == 0: # horizontal edge
ya = y2
xa = dx*p+x1
else:
theta = numpy.arctan2(dy, dx)
xa = p*d*numpy.cos(theta)+x1
ya = p*d*numpy.sin(theta)+y1
a_pos.append(((xa, ya), (x2, y2)))
arrow_collection = LineCollection(a_pos,
colors=arrow_colors,
linewidths=[4*ww for ww in lw],
antialiaseds=(1,),
transOffset = ax.transData,
)
arrow_collection.set_zorder(1) # edges go behind nodes
arrow_collection.set_label(label)
ax.add_collection(arrow_collection)
# update view
minx = numpy.amin(numpy.ravel(edge_pos[:, :, 0]))
maxx = numpy.amax(numpy.ravel(edge_pos[:, :, 0]))
miny = numpy.amin(numpy.ravel(edge_pos[:, :, 1]))
maxy = numpy.amax(numpy.ravel(edge_pos[:, :, 1]))
w = maxx-minx
h = maxy-miny
padx, pady = 0.05*w, 0.05*h
corners = (minx-padx, miny-pady), (maxx+padx, maxy+pady)
ax.update_datalim(corners)
ax.autoscale_view()
# if arrow_collection:
return edge_collection
def draw_networkx_edges(
G,
pos,
edgelist=None,
width=1.0,
edge_color="k",
style="solid",
alpha=None,
arrowstyle="-|>",
arrowsize=10,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
node_size=300,
nodelist=None,
node_shape="o",
connectionstyle=None,
min_source_margin=0,
min_target_margin=0,
):
"""Draw the edges of the graph G.
This draws only the edges of the graph G.
Parameters
----------
G : graph
A networkx graph
pos : dictionary
A dictionary with nodes as keys and positions as values.
Positions should be sequences of length 2.
edgelist : collection of edge tuples
Draw only specified edges(default=G.edges())
width : float, or array of floats
Line width of edges (default=1.0)
edge_color : color or array of colors (default='k')
Edge color. Can be a single color or a sequence of colors with the same
length as edgelist. Color can be string, or rgb (or rgba) tuple of
floats from 0-1. If numeric values are specified they will be
mapped to colors using the edge_cmap and edge_vmin,edge_vmax parameters.
style : string
Edge line style (default='solid') (solid|dashed|dotted,dashdot)
alpha : float
The edge transparency (default=None)
edge_ cmap : Matplotlib colormap
Colormap for mapping intensities of edges (default=None)
edge_vmin,edge_vmax : floats
Minimum and maximum for edge colormap scaling (default=None)
ax : Matplotlib Axes object, optional
Draw the graph in the specified Matplotlib axes.
arrows : bool, optional (default=True)
For directed graphs, if True draw arrowheads.
Note: Arrows will be the same color as edges.
arrowstyle : str, optional (default='-|>')
For directed graphs, choose the style of the arrow heads.
See :py:class: `matplotlib.patches.ArrowStyle` for more
options.
arrowsize : int, optional (default=10)
For directed graphs, choose the size of the arrow head head's length and
width. See :py:class: `matplotlib.patches.FancyArrowPatch` for attribute
`mutation_scale` for more info.
connectionstyle : str, optional (default=None)
Pass the connectionstyle parameter to create curved arc of rounding
radius rad. For example, connectionstyle='arc3,rad=0.2'.
See :py:class: `matplotlib.patches.ConnectionStyle` and
:py:class: `matplotlib.patches.FancyArrowPatch` for more info.
label : [None| string]
Label for legend
min_source_margin : int, optional (default=0)
The minimum margin (gap) at the begining of the edge at the source.
min_target_margin : int, optional (default=0)
The minimum margin (gap) at the end of the edge at the target.
Returns
-------
matplotlib.collection.LineCollection
`LineCollection` of the edges
list of matplotlib.patches.FancyArrowPatch
`FancyArrowPatch` instances of the directed edges
Depending whether the drawing includes arrows or not.
Notes
-----
For directed graphs, arrows are drawn at the head end. Arrows can be
turned off with keyword arrows=False. Be sure to include `node_size` as a
keyword argument; arrows are drawn considering the size of nodes.
Examples
--------
>>> G = nx.dodecahedral_graph()
>>> edges = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> G = nx.DiGraph()
>>> G.add_edges_from([(1, 2), (1, 3), (2, 3)])
>>> arcs = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> alphas = [0.3, 0.4, 0.5]
>>> for i, arc in enumerate(arcs): # change alpha values of arcs
... arc.set_alpha(alphas[i])
Also see the NetworkX drawing examples at
https://networkx.github.io/documentation/latest/auto_examples/index.html
See Also
--------
draw()
draw_networkx()
draw_networkx_nodes()
draw_networkx_labels()
draw_networkx_edge_labels()
"""
try:
import matplotlib.pyplot as plt
from matplotlib.colors import colorConverter, Colormap, Normalize
from matplotlib.collections import LineCollection
from matplotlib.patches import FancyArrowPatch
import numpy as np
except ImportError as e:
raise ImportError("Matplotlib required for draw()") from e
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if len(edgelist) == 0: # no edges!
if not G.is_directed() or not arrows:
return LineCollection(None)
else:
return []
if nodelist is None:
nodelist = list(G.nodes())
# FancyArrowPatch handles color=None different from LineCollection
if edge_color is None:
edge_color = "k"
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
# Check if edge_color is an array of floats and map to edge_cmap.
# This is the only case handled differently from matplotlib
if (np.iterable(edge_color) and (len(edge_color) == len(edge_pos))
and np.alltrue([isinstance(c, Number) for c in edge_color])):
if edge_cmap is not None:
assert isinstance(edge_cmap, Colormap)
else:
edge_cmap = plt.get_cmap()
if edge_vmin is None:
edge_vmin = min(edge_color)
if edge_vmax is None:
edge_vmax = max(edge_color)
color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)
edge_color = [edge_cmap(color_normal(e)) for e in edge_color]
if not G.is_directed() or not arrows:
edge_collection = LineCollection(
edge_pos,
colors=edge_color,
linewidths=width,
antialiaseds=(1, ),
linestyle=style,
transOffset=ax.transData,
alpha=alpha,
)
edge_collection.set_cmap(edge_cmap)
edge_collection.set_clim(edge_vmin, edge_vmax)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
return edge_collection
arrow_collection = None
if G.is_directed() and arrows:
# Note: Waiting for someone to implement arrow to intersection with
# marker. Meanwhile, this works well for polygons with more than 4
# sides and circle.
def to_marker_edge(marker_size, marker):
if marker in "s^>v<d": # `large` markers need extra space
return np.sqrt(2 * marker_size) / 2
else:
return np.sqrt(marker_size) / 2
# Draw arrows with `matplotlib.patches.FancyarrowPatch`
arrow_collection = []
mutation_scale = arrowsize # scale factor of arrow head
# FancyArrowPatch doesn't handle color strings
arrow_colors = colorConverter.to_rgba_array(edge_color, alpha)
for i, (src, dst) in enumerate(edge_pos):
x1, y1 = src
x2, y2 = dst
shrink_source = 0 # space from source to tail
shrink_target = 0 # space from head to target
if np.iterable(node_size): # many node sizes
source, target = edgelist[i][:2]
source_node_size = node_size[nodelist.index(source)]
target_node_size = node_size[nodelist.index(target)]
shrink_source = to_marker_edge(source_node_size, node_shape)
shrink_target = to_marker_edge(target_node_size, node_shape)
else:
shrink_source = shrink_target = to_marker_edge(
node_size, node_shape)
if shrink_source < min_source_margin:
shrink_source = min_source_margin
if shrink_target < min_target_margin:
shrink_target = min_target_margin
if len(arrow_colors) == len(edge_pos):
arrow_color = arrow_colors[i]
elif len(arrow_colors) == 1:
arrow_color = arrow_colors[0]
else: # Cycle through colors
arrow_color = arrow_colors[i % len(arrow_colors)]
if np.iterable(width):
if len(width) == len(edge_pos):
line_width = width[i]
else:
line_width = width[i % len(width)]
else:
line_width = width
arrow = FancyArrowPatch(
(x1, y1),
(x2, y2),
arrowstyle=arrowstyle,
shrinkA=shrink_source,
shrinkB=shrink_target,
mutation_scale=mutation_scale,
color=arrow_color,
linewidth=line_width,
connectionstyle=connectionstyle,
linestyle=style,
zorder=1,
) # arrows go behind nodes
# There seems to be a bug in matplotlib to make collections of
# FancyArrowPatch instances. Until fixed, the patches are added
# individually to the axes instance.
arrow_collection.append(arrow)
ax.add_patch(arrow)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
ax.tick_params(
axis="both",
which="both",
bottom=False,
left=False,
labelbottom=False,
labelleft=False,
)
return arrow_collection
def construct_eigenvalue_DOS_plot_haldane(xy,
fig,
dos_ax,
eig_ax,
eigval,
eigvect,
en,
Ni,
Nk,
marker_num=0,
color_scheme='default',
sub_lattice=-1,
PVx=[],
PVy=[],
black_t0lines=False,
mark_t0=True,
title='auto',
normalization=1.,
alpha=0.6,
lw=1,
zorder=10):
"""puts together lattice and DOS plots and draws normal mode magitudes as circles on top
Parameters
----------
xy: array 2nx3
Equilibrium position of the gyroscopes
fig :
figure with lattice and DOS drawn
dos_ax:
axis for the DOS plot
eig_ax
axis for the eigenvalue plot
eigval : array of dimension 2nx1
Eigenvalues of matrix for system
eigvect : array of dimension 2nx2n
Eigenvectors of matrix for system.
Eigvect is stored as NModes x NP array, like: mode0_psi0, mode0_psi1, ... / mode1_psi0, ...
en: int
Number of the eigenvalue you are plotting
Returns
----------
fig :
completed figure for normal mode
[scat_fg, p, f_mark] :
things to be cleared before next normal mode is drawn
"""
ppu = leplt.get_points_per_unit()
s = leplt.absolute_sizer()
# re_eigvals = sum(abs(real(eigval)))
# im_eigvals = sum(abs(imag(eigval)))
ev = eigval[en]
ev1 = ev
# Show where current eigenvalue is in DOS plot (red line ticking current eigval)
if dos_ax is not None:
(f_mark, ) = dos_ax.plot([ev.real, ev.real], P.ylim(), '-r')
plt.sca(dos_ax)
NP = len(xy)
im1 = np.imag(ev)
re1 = np.real(ev)
P.sca(eig_ax)
if title == 'auto':
eig_ax.set_title('Mode %d: $\omega=( %0.6f + %0.6f i)$' %
(en, re1, im1))
elif title is not None and title not in ['', 'none']:
eig_ax.set_title(title)
# Preallocate ellipsoid plot vars
shap = eigvect.shape
angles_arr = np.zeros(NP)
major_Ax = np.zeros(NP)
patch = []
polygon = []
colors = np.zeros(NP + 2)
# x_mag = np.zeros(NP)
# y_mag = np.zeros(NP)
x0s = np.zeros(NP)
y0s = np.zeros(NP)
mag1 = eigvect[en]
# Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like:
# x0, y0, x1, y1, ... xNP, yNP.
mag1x = np.array([mag1[i] for i in range(NP)])
mag1y = np.array([mag1[i] for i in range(NP)])
# Pick a series of times to draw out the ellipsoid
time_arr = np.arange(81) * 2 * pi / (abs(ev1) * 80)
exp1 = np.exp(1j * ev1 * time_arr)
# Normalization for the ellipsoids
lim_mag1 = max(
np.array([
np.sqrt(2 * abs(exp1 * mag1x[i])**2) for i in range(len(mag1x))
]).flatten())
mag1x /= lim_mag1
mag1y /= lim_mag1
mag1x *= normalization
mag1y *= normalization
cw = []
ccw = []
lines_1 = []
for i in range(NP):
unit = mag1x[i]
x_disps = 0.5 * (exp1 * unit).real
y_disps = 0.5 * (exp1 * unit).imag
x_vals = xy[i, 0] + x_disps
y_vals = xy[i, 1] + y_disps
# x_mag[i] = max(x_vals-xy[i,0]).real
# y_mag[i] = max(y_vals-xy[i,1]).real
poly_points = array([x_vals, y_vals]).T
polygon = Polygon(poly_points, True)
# x0 is the marker_num^th element of x_disps
x0 = x_disps[marker_num]
y0 = y_disps[marker_num]
x0s[i] = x_vals[marker_num]
y0s[i] = y_vals[marker_num]
# These are the black lines protruding from pivot point to current position
lines_1.append([[xy[i, 0], x_vals[marker_num]],
[xy[i, 1], y_vals[marker_num]]])
mag = sqrt(x0**2 + y0**2)
if mag > 0:
anglez = np.arccos(x0 / mag)
else:
anglez = 0
if y0 < 0:
anglez = 2 * np.pi - anglez
# testangle = arctan2(y0,x0)
# print ' x0 - x_disps[0] =', x0-x_disps[marker_num]
angles_arr[i] = anglez
# print 'polygon = ', poly_points
patch.append(polygon)
# Do Fast Fourier Transform (FFT)
# ff = abs(fft.fft(x_disps + 1j*y_disps))**2
# ff_freq = fft.fftfreq(len(x_vals), 1)
# mm_f = ff_freq[ff == max(ff)][0]
if color_scheme == 'default':
colors[i] = anglez
else:
if sub_lattice[i] == 0:
colors[i] = 0
else:
colors[i] = pi
# if mm_f > 0:
# colors[i] = 0
# else:
# colors[i] = pi
colors[NP] = 0
colors[NP + 1] = 2 * pi
plt.yticks([])
plt.xticks([])
# this is the part that puts a dot a t=0 point
if mark_t0:
scat_fg = eig_ax.scatter(x0s[cw], y0s[cw], s=s(.02), c='k')
scat_fg2 = eig_ax.scatter(x0s[cw], y0s[cw], s=s(.02), c='r')
scat_fg = [scat_fg, scat_fg2]
else:
scat_fg = []
NP = len(xy)
try:
NN = shape(Ni)[1]
except IndexError:
NN = 0
z = np.zeros(NP)
Rnorm = np.array([x0s, y0s, z]).T
# Bond Stretches
inc = 0
stretches = zeros(3 * len(xy))
if PVx == [] and PVy == []:
'''There are no periodic boundaries supplied'''
for i in range(len(xy)):
if NN > 0:
for j, k in zip(Ni[i], Nk[i]):
if i < j and abs(k) > 0:
n1 = float(linalg.norm(Rnorm[i] - Rnorm[j]))
n2 = linalg.norm(xy[i] - xy[j])
stretches[inc] = (n1 - n2)
inc += 1
else:
'''There are periodic boundaries supplied'''
# get boundary particle indices
KLabs = np.zeros_like(Nk, dtype='int')
KLabs[Nk > 0] = 1
boundary = extract_boundary_from_NL(xy, Ni, KLabs)
for i in range(len(xy)):
if NN > 0:
for j, k in zip(Ni[i], Nk[i]):
# if i in boundary and j in boundary:
# col = np.where( Ni[i] == j)[0][0]
# print 'col = ', col
# n1 = float( np.linalg.norm(Rnorm[i]-Rnorm[j]) )
# n2 = np.linalg.norm(R[i] - R[j] )
# stretches[inc] = (n1 - n2)
# inc += 1
#
# #test[inc] = [R[i], np.array([R[j,0]+PVx[i,col], R[j,1] + PVy[i,col], 0])]
# else:
if i < j and abs(k) > 0:
n1 = float(np.linalg.norm(Rnorm[i] - Rnorm[j]))
n2 = np.linalg.norm(xy[i] - xy[j])
stretches[inc] = (n1 - n2)
inc += 1
stretch = np.array(stretches[0:inc])
# For particles with neighbors, get list of bonds to draw by stretches
test = list(np.zeros([inc, 1]))
inc = 0
xy = np.array([x0s, y0s, z]).T
if PVx == [] and PVy == []:
'''There are no periodic boundaries supplied'''
for i in range(len(xy)):
if NN > 0:
for j, k in zip(Ni[i], Nk[i]):
if i < j and abs(k) > 0:
test[inc] = [xy[(i, j), 0], xy[(i, j), 1]]
inc += 1
else:
'''There are periodic boundaries supplied'''
# get boundary particle indices
KLabs = np.zeros_like(Nk, dtype='int')
KLabs[Nk > 0] = 1
boundary = extract_boundary_from_NL(xy, Ni, KLabs)
for i in range(len(xy)):
if NN > 0:
for j, k in zip(Ni[i], Nk[i]):
# if i in boundary and j in boundary:
# col = np.where( Ni[i] == j)[0][0]
# print 'i,j = (', i,j, ')'
# print 'PVx[i,col] = ', PVx[i,col]
# print 'PVy[i,col] = ', PVy[i,col]
# test[inc] = [xy[i], np.array([xy[j,0]+PVx[i,col], xy[j,1] + PVy[i,col], 0])]
# #plt.plot([ xy[i,0], xy[j,0]+PVx[i,col]], [xy[i,1], xy[j,1] + PVy[i,col] ], 'k-')
# #plt.plot(xy[:,0], xy[:,1],'b.')
# #plt.show()
# print 'test = ', test
# inc += 1
# else:
if i < j and abs(k) > 0:
test[inc] = [xy[(i, j), 0], xy[(i, j), 1]]
inc += 1
lines = [zip(x, y) for x, y in test]
# angles[-1] = 0
# angles[-2] = 2*pi
lines_st = LineCollection(lines,
array=stretch,
cmap='seismic',
linewidth=8)
lines_st.set_clim([-1. * 0.25, 1 * 0.25])
lines_st.set_zorder(2)
if black_t0lines:
lines_12 = [zip(x, y) for x, y in lines_1]
lines_12_st = LineCollection(lines_12, linewidth=0.8)
lines_12_st.set_color('k')
eig_ax.add_collection(lines_12_st)
else:
lines_12_st = []
p = PatchCollection(patch, cmap='hsv', lw=lw, alpha=alpha, zorder=zorder)
p.set_array(array(colors))
p.set_clim([0, 2 * pi])
p.set_zorder(1)
# eig_ax.add_collection(lines_st)
eig_ax.add_collection(p)
eig_ax.set_aspect('equal')
s = leplt.absolute_sizer()
# erased ev/(2*pi) here npm 2016
cw_ccw = [cw, ccw, ev]
# print cw_ccw[1]
return fig, [scat_fg, p, f_mark, lines_12_st], cw_ccw
def plot_linestring_collection(ax,
geoms,
colors_or_values,
plot_values,
vmin=None,
vmax=None,
cmap=None,
linewidth=1.0,
**kwargs):
"""
Plots a collection of LineString and MultiLineString geometries to `ax`
Parameters
----------
ax : matplotlib.axes.Axes
where shapes will be plotted
geoms : a sequence of `N` LineStrings and/or MultiLineStrings (can be mixed)
colors_or_values : a sequence of `N` values or RGBA tuples
It should have 1:1 correspondence with the geometries (not their components).
plot_values : bool
If True, `colors_or_values` is interpreted as a list of values, and will
be mapped to colors using vmin/vmax/cmap (which become required).
Otherwise `colors_or_values` is interpreted as a list of colors.
Returns
-------
collection : matplotlib.collections.Collection that was plotted
"""
from matplotlib.collections import LineCollection
components, component_colors_or_values = _flatten_multi_geoms(
geoms, colors_or_values)
# LineCollection does not accept some kwargs.
if 'markersize' in kwargs:
del kwargs['markersize']
segments = [np.array(linestring)[:, :2] for linestring in components]
collection = LineCollection(segments, linewidth=linewidth, **kwargs)
if plot_values:
collection.set_array(np.array(component_colors_or_values))
collection.set_cmap(cmap)
collection.set_clim(vmin, vmax)
else:
# set_color magically sets the correct combination of facecolor and
# edgecolor, based on collection type.
collection.set_color(component_colors_or_values)
# If the user set facecolor and/or edgecolor explicitly, the previous
# call to set_color might have overridden it (remember, the 'color' may
# have come from plot_series, not from the user). The user should be
# able to override matplotlib's default behavior, by setting them again
# after set_color.
if 'facecolor' in kwargs:
collection.set_facecolor(kwargs['facecolor'])
if 'edgecolor' in kwargs:
collection.set_edgecolor(kwargs['edgecolor'])
ax.add_collection(collection, autolim=True)
ax.autoscale_view()
return collection
def overlay_ssies_data(self, mobj, ax,
interval=60, velscl=500.e3,
vec_cmap=None,
vel_scale=[0, 1000.],
plot_path=True,
plot_vecs_on_full_path=False,
quality_flag=False):
"""Overlays data onto a map"""
import matplotlib
from matplotlib.colors import ListedColormap as lcm
import sys
import datetime as dt
import math
import numpy as np
from matplotlib.collections import LineCollection
# create a colormap for the quality flags of Vx and Vy
cmj = matplotlib.cm.jet
cmap_flag = lcm(['k', 'y', 'r', cmj(.27)])
bounds_flag = np.round(np.linspace(0.5, 4.5, 5))
norm_flag = matplotlib.colors.BoundaryNorm(bounds_flag, cmap_flag.N)
# drop NaN values for 'Vy', 'GLAT', and 'GLONG'
df = self.data[['Vy', 'GLAT', 'GLONG', 'I']].dropna()
df_vys_tmp = df['Vy']
df_pos = df[['GLAT', 'GLONG']]
df_Is = df['I']
df_vy = df_vys_tmp.loc[self.datetime.strftime("%Y%m%d%H%M%S"):\
(self.datetime+dt.timedelta(seconds=interval-1)).\
strftime("%Y%m%d%H%M%S")]
df_lat = df_pos.loc[self.datetime.strftime("%Y%m%d%H%M%S"):\
(self.datetime+dt.timedelta(seconds=interval-1)).\
strftime("%Y%m%d%H%M%S"), 'GLAT']
df_lon = df_pos.loc[self.datetime.strftime("%Y%m%d%H%M%S"):\
(self.datetime+dt.timedelta(seconds=interval-1)).\
strftime("%Y%m%d%H%M%S"), 'GLONG']
df_I = df_Is.loc[self.datetime.strftime("%Y%m%d%H%M%S"):\
(self.datetime+dt.timedelta(seconds=interval-1)).\
strftime("%Y%m%d%H%M%S")]
try:
xxs, yys = mobj(df_lon[0], df_lat[0], coords='geo')
xxe, yye = mobj(df_lon[-1], df_lat[-1], coords='geo')
except IndexError:
return
the_x = math.atan2(yye-yys, xxe-xxs)
the_vel = the_x - np.deg2rad(90)
# Get the locs along the full path over polar region
df_lats = df_pos.loc[(self.datetime-dt.\
timedelta(seconds=20*interval)).strftime("%Y%m%d%H%M%S"):\
(self.datetime+dt.timedelta(seconds=20*interval-1)).\
strftime("%Y%m%d%H%M%S"), 'GLAT']
df_lons = df_pos.loc[(self.datetime-dt.\
timedelta(seconds=20*interval)).strftime("%Y%m%d%H%M%S"):\
(self.datetime+dt.timedelta(seconds=20*interval-1)).\
strftime("%Y%m%d%H%M%S"), 'GLONG']
df_vys = df_vys_tmp.loc[(self.datetime-dt.\
timedelta(seconds=20*interval)).strftime("%Y%m%d%H%M%S"):\
(self.datetime+dt.timedelta(seconds=20*interval-1)).\
strftime("%Y%m%d%H%M%S")]
# plot the DMSP path
if plot_path:
for k in range(len(df_lats)):
x1s, y1s = mobj(df_lons[k], df_lats[k], coords='geo')
mobj.scatter(np.array(x1s), np.array(y1s),
s=1.0, zorder=5, marker='o', color='gray',
edgecolors='face', linewidths=.5)
if plot_vecs_on_full_path:
verts = [[],[]]
tails_dmsp = []
# Plot vectors along the full path
for k in range(len(df_lats)):
x1, y1 = mobj(df_lons[k], df_lats[k], coords='geo')
verts[0].append(x1)
verts[1].append(y1)
x2 = x1+df_vys[k]*velscl*(-1.0)*math.cos(the_vel)
y2 = y1+df_vys[k]*velscl*(-1.0)*math.sin(the_vel)
tails_dmsp.append(((x1,y1),(x2,y2)))
lcoll = LineCollection(np.array(tails_dmsp),
linewidths=.6,
zorder=12, alpha=0.6, color="b")
ax.add_collection(lcoll)
verts = [[],[]]
tails_dmsp = []
# plot the measurement points and velocity vectors at a specified time
for k in range(len(df_lat)):
x1, y1 = mobj(df_lon[k], df_lat[k], coords='geo')
verts[0].append(x1)
verts[1].append(y1)
x2 = x1+df_vy[k]*velscl*(-1.0)*math.cos(the_vel)
y2 = y1+df_vy[k]*velscl*(-1.0)*math.sin(the_vel)
tails_dmsp.append(((x1,y1),(x2,y2)))
xx = mobj.scatter(np.array(verts[0]),np.array(verts[1]),
s=2.5,zorder=5,marker='o',
c=np.abs(df_vy.as_matrix()),
vmin=vel_scale[0], vmax=vel_scale[1],
edgecolors='face', linewidths=.5,
cmap=vec_cmap)
if quality_flag:
lcoll = LineCollection(np.array(tails_dmsp), linewidths=.6,
zorder=12, cmap=cmap_flag,
norm=norm_flag, alpha=1)
lcoll.set_array(df_I)
else:
lcoll = LineCollection(np.array(tails_dmsp),
linewidths=.6,
zorder=12, alpha=1,
cmap=vec_cmap)
lcoll.set_array(np.abs(df_vy.as_matrix()))
lcoll.set_clim(vmin=vel_scale[0], vmax=vel_scale[1])
ax.add_collection(lcoll)
return
def draw_networkx_edges(
G,
pos,
edgelist=None,
width=1.0,
edge_color="k",
style="solid",
alpha=None,
arrowstyle="-|>",
arrowsize=3,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
node_size=300,
nodelist=None,
node_shape="o",
connectionstyle=None,
min_source_margin=0,
min_target_margin=0,
):
"""Draw the edges of the graph G. Adjusted from networkx."""
try:
import matplotlib.pyplot as plt
from matplotlib.colors import colorConverter, Colormap, Normalize
from matplotlib.collections import LineCollection
from matplotlib.patches import FancyArrowPatch
from numbers import Number
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
if nodelist is None:
nodelist = list(G.nodes())
# FancyArrowPatch handles color=None different from LineCollection
if edge_color is None:
edge_color = "k"
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
# Check if edge_color is an array of floats and map to edge_cmap.
# This is the only case handled differently from matplotlib
if (
np.iterable(edge_color)
and (len(edge_color) == len(edge_pos))
and np.alltrue([isinstance(c, Number) for c in edge_color])
):
if edge_cmap is not None:
assert isinstance(edge_cmap, Colormap)
else:
edge_cmap = plt.get_cmap()
if edge_vmin is None:
edge_vmin = min(edge_color)
if edge_vmax is None:
edge_vmax = max(edge_color)
color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)
edge_color = [edge_cmap(color_normal(e)) for e in edge_color]
if not G.is_directed() or not arrows:
edge_collection = LineCollection(
edge_pos,
colors=edge_color,
linewidths=width,
antialiaseds=(1,),
linestyle=style,
transOffset=ax.transData,
alpha=alpha,
)
edge_collection.set_cmap(edge_cmap)
edge_collection.set_clim(edge_vmin, edge_vmax)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
return edge_collection
arrow_collection = None
if G.is_directed() and arrows:
# Note: Waiting for someone to implement arrow to intersection with
# marker. Meanwhile, this works well for polygons with more than 4
# sides and circle.
def to_marker_edge(marker_size, marker):
if marker in "s^>v<d": # `large` markers need extra space
return np.sqrt(2 * marker_size) / 2
else:
return np.sqrt(marker_size) / 2
# Draw arrows with `matplotlib.patches.FancyarrowPatch`
arrow_collection = []
mutation_scale = arrowsize # scale factor of arrow head
# FancyArrowPatch doesn't handle color strings
arrow_colors = colorConverter.to_rgba_array(edge_color, alpha)
for i, (src, dst) in enumerate(edge_pos):
x1, y1 = src
x2, y2 = dst
shrink_source = 0 # space from source to tail
shrink_target = 0 # space from head to target
if np.iterable(node_size): # many node sizes
source, target = edgelist[i][:2]
source_node_size = node_size[nodelist.index(source)]
target_node_size = node_size[nodelist.index(target)]
shrink_source = to_marker_edge(source_node_size, node_shape)
shrink_target = to_marker_edge(target_node_size, node_shape)
else:
shrink_source = shrink_target = to_marker_edge(node_size, node_shape)
if shrink_source < min_source_margin:
shrink_source = min_source_margin
if shrink_target < min_target_margin:
shrink_target = min_target_margin
if len(arrow_colors) == len(edge_pos):
arrow_color = arrow_colors[i]
elif len(arrow_colors) == 1:
arrow_color = arrow_colors[0]
else: # Cycle through colors
arrow_color = arrow_colors[i % len(arrow_colors)]
if np.iterable(width):
if len(width) == len(edge_pos):
line_width = width[i]
else:
line_width = width[i % len(width)]
else:
line_width = width
arrow = FancyArrowPatch(
(x1, y1),
(x2, y2),
arrowstyle=arrowstyle,
shrinkA=shrink_source,
shrinkB=shrink_target,
mutation_scale=mutation_scale,
color=arrow_color,
linewidth=line_width,
connectionstyle=connectionstyle,
linestyle=style,
zorder=1,
) # arrows go behind nodes
# There seems to be a bug in matplotlib to make collections of
# FancyArrowPatch instances. Until fixed, the patches are added
# individually to the axes instance.
arrow_collection.append(arrow)
ax.add_patch(arrow)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
ax.tick_params(
axis="both",
which="both",
bottom=False,
left=False,
labelbottom=False,
labelleft=False,
)
return arrow_collection
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
**kwds):
"""Draw the edges of the graph G
This draws only the edges of the graph G.
pos is a dictionary keyed by vertex with a two-tuple
of x-y positions as the value.
See networkx_v099.layout for functions that compute node positions.
edgelist is an optional list of the edges in G to be drawn.
If provided, only the edges in edgelist will be drawn.
edgecolor can be a list of matplotlib color letters such as 'k' or
'b' that lists the color of each edge; the list must be ordered in
the same way as the edge list. Alternatively, this list can contain
numbers and those number are mapped to a color scale using the color
map edge_cmap.
For directed graphs, "arrows" (actually just thicker stubs) are drawn
at the head end. Arrows can be turned off with keyword arrows=False.
See draw_networkx_v099 for the list of other optional parameters.
"""
if ax is None:
ax=matplotlib.pylab.gca()
if edgelist is None:
edgelist=G.edges()
if not edgelist or len(edgelist)==0: # no edges!
return None
# set edge positions
edge_pos=asarray([(pos[e[0]],pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width,)
else:
lw = width
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color)==len(edge_pos):
if matplotlib.numerix.alltrue([cb.is_string_like(c)
for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple([colorConverter.to_rgba(c,alpha)
for c in edge_color])
elif matplotlib.numerix.alltrue([not cb.is_string_like(c)
for c in edge_color]):
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must consist of either color names or numbers')
else:
if len(edge_color)==1:
edge_colors = ( colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError('edge_color must be a single color or list of exactly m colors where m is the number or edges')
edge_collection = LineCollection(edge_pos,
colors = edge_colors,
linewidths = lw,
antialiaseds = (1,),
linestyle = style,
transOffset = ax.transData,
)
edge_collection.set_alpha(alpha)
# need 0.87.7 or greater for edge colormaps
mpl_version=matplotlib.__version__
if mpl_version.endswith('svn'):
mpl_version=matplotlib.__version__[0:-3]
if mpl_version.endswith('pre'):
mpl_version=matplotlib.__version__[0:-3]
if map(int,mpl_version.split('.'))>=[0,87,7]:
if edge_colors is None:
if edge_cmap is not None: assert(isinstance(edge_cmap, Colormap))
edge_collection.set_array(asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
matplotlib.pylab.sci(edge_collection)
# else:
# sys.stderr.write(\
# """matplotlib version >= 0.87.7 required for colormapped edges.
# (version %s detected)."""%matplotlib.__version__)
# raise UserWarning(\
# """matplotlib version >= 0.87.7 required for colormapped edges.
# (version %s detected)."""%matplotlib.__version__)
arrow_collection=None
if G.directed and arrows:
# a directed graph hack
# draw thick line segments at head end of edge
# waiting for someone else to implement arrows that will work
arrow_colors = ( colorConverter.to_rgba('k', alpha), )
a_pos=[]
p=1.0-0.25 # make head segment 25 percent of edge length
for src,dst in edge_pos:
x1,y1=src
x2,y2=dst
dx=x2-x1 # x offset
dy=y2-y1 # y offset
d=sqrt(float(dx**2+dy**2)) # length of edge
if d==0: # source and target at same position
continue
if dx==0: # vertical edge
xa=x2
ya=dy*p+y1
if dy==0: # horizontal edge
ya=y2
xa=dx*p+x1
else:
theta=arctan2(dy,dx)
xa=p*d*cos(theta)+x1
ya=p*d*sin(theta)+y1
a_pos.append(((xa,ya),(x2,y2)))
arrow_collection = LineCollection(a_pos,
colors = arrow_colors,
linewidths = [4*ww for ww in lw],
antialiaseds = (1,),
transOffset = ax.transData,
)
# update view
minx = amin(ravel(edge_pos[:,:,0]))
maxx = amax(ravel(edge_pos[:,:,0]))
miny = amin(ravel(edge_pos[:,:,1]))
maxy = amax(ravel(edge_pos[:,:,1]))
w = maxx-minx
h = maxy-miny
padx, pady = 0.05*w, 0.05*h
corners = (minx-padx, miny-pady), (maxx+padx, maxy+pady)
ax.update_datalim( corners)
ax.autoscale_view()
edge_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(edge_collection)
if arrow_collection:
arrow_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(arrow_collection)
return edge_collection
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=None,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
**kwds):
"""Draw the edges of the graph G
This draws only the edges of the graph G.
pos is a dictionary keyed by vertex with a two-tuple
of x-y positions as the value.
See networkx.layout for functions that compute node positions.
edgelist is an optional list of the edges in G to be drawn.
If provided, only the edges in edgelist will be drawn.
edgecolor can be a list of matplotlib color letters such as 'k' or
'b' that lists the color of each edge; the list must be ordered in
the same way as the edge list. Alternatively, this list can contain
numbers and those number are mapped to a color scale using the color
map edge_cmap. Finally, it can also be a list of (r,g,b) or (r,g,b,a)
tuples, in which case these will be used directly to color the edges. If
the latter mode is used, you should not provide a value for alpha, as it
would be applied globally to all lines.
For directed graphs, "arrows" (actually just thicker stubs) are drawn
at the head end. Arrows can be turned off with keyword arrows=False.
See draw_networkx for the list of other optional parameters.
"""
try:
import matplotlib
import matplotlib.pylab as pylab
import numpy as np
from matplotlib.colors import colorConverter,Colormap
from matplotlib.collections import LineCollection
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
pass # unable to open display
if ax is None:
ax=pylab.gca()
if edgelist is None:
edgelist=G.edges()
if not edgelist or len(edgelist)==0: # no edges!
return None
# set edge positions
edge_pos=np.asarray([(pos[e[0]],pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width,)
else:
lw = width
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color)==len(edge_pos):
if np.alltrue([cb.is_string_like(c)
for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple([colorConverter.to_rgba(c,alpha)
for c in edge_color])
elif np.alltrue([not cb.is_string_like(c)
for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if np.alltrue([cb.iterable(c) and len(c) in (3,4)
for c in edge_color]):
edge_colors = tuple(edge_color)
alpha=None
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must consist of either color names or numbers')
else:
if len(edge_color)==1:
edge_colors = ( colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError('edge_color must be a single color or list of exactly m colors where m is the number or edges')
edge_collection = LineCollection(edge_pos,
colors = edge_colors,
linewidths = lw,
antialiaseds = (1,),
linestyle = style,
transOffset = ax.transData,
)
# Note: there was a bug in mpl regarding the handling of alpha values for
# each line in a LineCollection. It was fixed in matplotlib in r7184 and
# r7189 (June 6 2009). We should then not set the alpha value globally,
# since the user can instead provide per-edge alphas now. Only set it
# globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
# need 0.87.7 or greater for edge colormaps. No checks done, this will
# just not work with an older mpl
if edge_colors is None:
if edge_cmap is not None: assert(isinstance(edge_cmap, Colormap))
edge_collection.set_array(np.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
pylab.sci(edge_collection)
arrow_collection=None
if G.is_directed() and arrows:
# a directed graph hack
# draw thick line segments at head end of edge
# waiting for someone else to implement arrows that will work
arrow_colors = ( colorConverter.to_rgba('k', alpha), )
a_pos=[]
p=1.0-0.25 # make head segment 25 percent of edge length
for src,dst in edge_pos:
x1,y1=src
x2,y2=dst
dx=x2-x1 # x offset
dy=y2-y1 # y offset
d=np.sqrt(float(dx**2+dy**2)) # length of edge
if d==0: # source and target at same position
continue
if dx==0: # vertical edge
xa=x2
ya=dy*p+y1
if dy==0: # horizontal edge
ya=y2
xa=dx*p+x1
else:
theta=np.arctan2(dy,dx)
xa=p*d*np.cos(theta)+x1
ya=p*d*np.sin(theta)+y1
a_pos.append(((xa,ya),(x2,y2)))
arrow_collection = LineCollection(a_pos,
colors = arrow_colors,
linewidths = [4*ww for ww in lw],
antialiaseds = (1,),
transOffset = ax.transData,
)
# update view
minx = np.amin(np.ravel(edge_pos[:,:,0]))
maxx = np.amax(np.ravel(edge_pos[:,:,0]))
miny = np.amin(np.ravel(edge_pos[:,:,1]))
maxy = np.amax(np.ravel(edge_pos[:,:,1]))
w = maxx-minx
h = maxy-miny
padx, pady = 0.05*w, 0.05*h
corners = (minx-padx, miny-pady), (maxx+padx, maxy+pady)
ax.update_datalim( corners)
ax.autoscale_view()
edge_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(edge_collection)
if arrow_collection:
arrow_collection.set_zorder(1) # edges go behind nodes
ax.add_collection(arrow_collection)
return edge_collection
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
arrowstyle='-|>',
arrowsize=10,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
node_size=300,
nodelist=None,
node_shape="o",
**kwds):
"""Draw the edges of the graph G.
This draws only the edges of the graph G.
Parameters
----------
G : graph
A networkx graph
pos : dictionary
A dictionary with nodes as keys and positions as values.
Positions should be sequences of length 2.
edgelist : collection of edge tuples
Draw only specified edges(default=G.edges())
width : float, or array of floats
Line width of edges (default=1.0)
edge_color : color string, or array of floats
Edge color. Can be a single color format string (default='r'),
or a sequence of colors with the same length as edgelist.
If numeric values are specified they will be mapped to
colors using the edge_cmap and edge_vmin,edge_vmax parameters.
style : string
Edge line style (default='solid') (solid|dashed|dotted,dashdot)
alpha : float
The edge transparency (default=1.0)
edge_ cmap : Matplotlib colormap
Colormap for mapping intensities of edges (default=None)
edge_vmin,edge_vmax : floats
Minimum and maximum for edge colormap scaling (default=None)
ax : Matplotlib Axes object, optional
Draw the graph in the specified Matplotlib axes.
arrows : bool, optional (default=True)
For directed graphs, if True draw arrowheads.
Note: Arrows will be the same color as edges.
arrowstyle : str, optional (default='-|>')
For directed graphs, choose the style of the arrow heads.
See :py:class: `matplotlib.patches.ArrowStyle` for more
options.
arrowsize : int, optional (default=10)
For directed graphs, choose the size of the arrow head head's length and
width. See :py:class: `matplotlib.patches.FancyArrowPatch` for attribute
`mutation_scale` for more info.
label : [None| string]
Label for legend
Returns
-------
matplotlib.collection.LineCollection
`LineCollection` of the edges
list of matplotlib.patches.FancyArrowPatch
`FancyArrowPatch` instances of the directed edges
Depending whether the drawing includes arrows or not.
Notes
-----
For directed graphs, arrows are drawn at the head end. Arrows can be
turned off with keyword arrows=False. Be sure to include `node_size' as a
keyword argument; arrows are drawn considering the size of nodes.
Examples
--------
>>> G = nx.dodecahedral_graph()
>>> edges = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> G = nx.DiGraph()
>>> G.add_edges_from([(1, 2), (1, 3), (2, 3)])
>>> arcs = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> alphas = [0.3, 0.4, 0.5]
>>> for i, arc in enumerate(arcs): # change alpha values of arcs
... arc.set_alpha(alphas[i])
Also see the NetworkX drawing examples at
https://networkx.github.io/documentation/latest/auto_examples/index.html
See Also
--------
draw()
draw_networkx()
draw_networkx_nodes()
draw_networkx_labels()
draw_networkx_edge_labels()
"""
try:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cbook as cb
from matplotlib.colors import colorConverter, Colormap, Normalize
from matplotlib.collections import LineCollection
from matplotlib.patches import FancyArrowPatch
import numpy as np
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
if nodelist is None:
nodelist = list(G.nodes())
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width,)
else:
lw = width
if not is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color) == len(edge_pos):
if np.alltrue([is_string_like(c) for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple([colorConverter.to_rgba(c, alpha)
for c in edge_color])
elif np.alltrue([not is_string_like(c) for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if np.alltrue([cb.iterable(c) and len(c) in (3, 4)
for c in edge_color]):
edge_colors = tuple(edge_color)
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must contain color names or numbers')
else:
if is_string_like(edge_color) or len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
msg = 'edge_color must be a color or list of one color per edge'
raise ValueError(msg)
if (not G.is_directed() or not arrows):
edge_collection = LineCollection(edge_pos,
colors=edge_colors,
linewidths=lw,
antialiaseds=(1,),
linestyle=style,
transOffset=ax.transData,
)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
# Note: there was a bug in mpl regarding the handling of alpha values
# for each line in a LineCollection. It was fixed in matplotlib by
# r7184 and r7189 (June 6 2009). We should then not set the alpha
# value globally, since the user can instead provide per-edge alphas
# now. Only set it globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
if edge_colors is None:
if edge_cmap is not None:
assert(isinstance(edge_cmap, Colormap))
edge_collection.set_array(np.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
return edge_collection
arrow_collection = None
if G.is_directed() and arrows:
# Note: Waiting for someone to implement arrow to intersection with
# marker. Meanwhile, this works well for polygons with more than 4
# sides and circle.
def to_marker_edge(marker_size, marker):
if marker in "s^>v<d": # `large` markers need extra space
return np.sqrt(2 * marker_size) / 2
else:
return np.sqrt(marker_size) / 2
# Draw arrows with `matplotlib.patches.FancyarrowPatch`
arrow_collection = []
mutation_scale = arrowsize # scale factor of arrow head
arrow_colors = edge_colors
if arrow_colors is None:
if edge_cmap is not None:
assert(isinstance(edge_cmap, Colormap))
else:
edge_cmap = plt.get_cmap() # default matplotlib colormap
if edge_vmin is None:
edge_vmin = min(edge_color)
if edge_vmax is None:
edge_vmax = max(edge_color)
color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)
for i, (src, dst) in enumerate(edge_pos):
x1, y1 = src
x2, y2 = dst
arrow_color = None
line_width = None
shrink_source = 0 # space from source to tail
shrink_target = 0 # space from head to target
if cb.iterable(node_size): # many node sizes
src_node, dst_node = edgelist[i]
index_node = nodelist.index(dst_node)
marker_size = node_size[index_node]
shrink_target = to_marker_edge(marker_size, node_shape)
else:
shrink_target = to_marker_edge(node_size, node_shape)
if arrow_colors is None:
arrow_color = edge_cmap(color_normal(edge_color[i]))
elif len(arrow_colors) > 1:
arrow_color = arrow_colors[i]
else:
arrow_color = arrow_colors[0]
if len(lw) > 1:
line_width = lw[i]
else:
line_width = lw[0]
arrow = FancyArrowPatch((x1, y1), (x2, y2),
arrowstyle=arrowstyle,
shrinkA=shrink_source,
shrinkB=shrink_target,
mutation_scale=mutation_scale,
color=arrow_color,
linewidth=line_width,
zorder=1) # arrows go behind nodes
# There seems to be a bug in matplotlib to make collections of
# FancyArrowPatch instances. Until fixed, the patches are added
# individually to the axes instance.
arrow_collection.append(arrow)
ax.add_patch(arrow)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
return arrow_collection
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
**kwds):
"""Draw the edges of the graph G.
This draws only the edges of the graph G.
Parameters
----------
G : graph
A networkx graph
pos : dictionary
A dictionary with nodes as keys and positions as values.
Positions should be sequences of length 2.
edgelist : collection of edge tuples
Draw only specified edges(default=G.edges())
width : float, or array of floats
Line width of edges (default=1.0)
edge_color : color string, or array of floats
Edge color. Can be a single color format string (default='r'),
or a sequence of colors with the same length as edgelist.
If numeric values are specified they will be mapped to
colors using the edge_cmap and edge_vmin,edge_vmax parameters.
style : string
Edge line style (default='solid') (solid|dashed|dotted,dashdot)
alpha : float
The edge transparency (default=1.0)
edge_ cmap : Matplotlib colormap
Colormap for mapping intensities of edges (default=None)
edge_vmin,edge_vmax : floats
Minimum and maximum for edge colormap scaling (default=None)
ax : Matplotlib Axes object, optional
Draw the graph in the specified Matplotlib axes.
arrows : bool, optional (default=True)
For directed graphs, if True draw arrowheads.
label : [None| string]
Label for legend
Returns
-------
matplotlib.collection.LineCollection
`LineCollection` of the edges
Notes
-----
For directed graphs, "arrows" (actually just thicker stubs) are drawn
at the head end. Arrows can be turned off with keyword arrows=False.
Yes, it is ugly but drawing proper arrows with Matplotlib this
way is tricky.
Examples
--------
>>> G = nx.dodecahedral_graph()
>>> edges = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
Also see the NetworkX drawing examples at
https://networkx.github.io/documentation/latest/auto_examples/index.html
See Also
--------
draw()
draw_networkx()
draw_networkx_nodes()
draw_networkx_labels()
draw_networkx_edge_labels()
"""
try:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cbook as cb
from matplotlib.colors import colorConverter, Colormap
from matplotlib.collections import LineCollection
import numpy as np
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width,)
else:
lw = width
if not cb.is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color) == len(edge_pos):
if np.alltrue([cb.is_string_like(c)
for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple([colorConverter.to_rgba(c, alpha)
for c in edge_color])
elif np.alltrue([not cb.is_string_like(c)
for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if np.alltrue([cb.iterable(c) and len(c) in (3, 4)
for c in edge_color]):
edge_colors = tuple(edge_color)
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must consist of either color names or numbers')
else:
if cb.is_string_like(edge_color) or len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
raise ValueError(
'edge_color must be a single color or list of exactly m colors where m is the number or edges')
edge_collection = LineCollection(edge_pos,
colors=edge_colors,
linewidths=lw,
antialiaseds=(1,),
linestyle=style,
transOffset=ax.transData,
)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
# Note: there was a bug in mpl regarding the handling of alpha values for
# each line in a LineCollection. It was fixed in matplotlib in r7184 and
# r7189 (June 6 2009). We should then not set the alpha value globally,
# since the user can instead provide per-edge alphas now. Only set it
# globally if provided as a scalar.
if cb.is_numlike(alpha):
edge_collection.set_alpha(alpha)
if edge_colors is None:
if edge_cmap is not None:
assert(isinstance(edge_cmap, Colormap))
edge_collection.set_array(np.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
arrow_collection = None
if G.is_directed() and arrows:
# a directed graph hack
# draw thick line segments at head end of edge
# waiting for someone else to implement arrows that will work
arrow_colors = edge_colors
a_pos = []
p = 1.0 - 0.25 # make head segment 25 percent of edge length
for src, dst in edge_pos:
x1, y1 = src
x2, y2 = dst
dx = x2 - x1 # x offset
dy = y2 - y1 # y offset
d = np.sqrt(float(dx**2 + dy**2)) # length of edge
if d == 0: # source and target at same position
continue
if dx == 0: # vertical edge
xa = x2
ya = dy * p + y1
if dy == 0: # horizontal edge
ya = y2
xa = dx * p + x1
else:
theta = np.arctan2(dy, dx)
xa = p * d * np.cos(theta) + x1
ya = p * d * np.sin(theta) + y1
a_pos.append(((xa, ya), (x2, y2)))
arrow_collection = LineCollection(a_pos,
colors=arrow_colors,
linewidths=[4 * ww for ww in lw],
antialiaseds=(1,),
transOffset=ax.transData,
)
arrow_collection.set_zorder(1) # edges go behind nodes
arrow_collection.set_label(label)
ax.add_collection(arrow_collection)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
# if arrow_collection:
return edge_collection
class SunPlotPy(wx.Frame, Spatial, Grid):
"""
The main frame of the application
"""
title = 'sunplot(py)'
# Plotting options
autoclim = True
showedges = False
bgcolor = 'k'
textcolor = 'w'
cmap = 'RdBu'
particlesize = 1.8
particlecolor = 'm'
# other flags
collectiontype = 'cells'
oldcollectiontype = 'cells'
#
tindex = 0
depthlevs = [0., 10., 100., 200., 300., 400., 500.,\
1000.,2000.,3000.,4000.,5000]
_FillValue = 999999
def __init__(self):
wx.Frame.__init__(self, None, -1, self.title)
self.create_menu()
self.create_status_bar()
self.create_main_panel()
#self.draw_figure()
def create_menu(self):
self.menubar = wx.MenuBar()
###
# File Menu
###
menu_file = wx.Menu()
# Load a hydro output file
m_expt = menu_file.Append(-1, "&Open file\tCtrl-O", "Open netcdf file")
self.Bind(wx.EVT_MENU, self.on_open_file, m_expt)
# Load a grid file
m_grid = menu_file.Append(-1, "&Load grid\tCtrl-G",
"Load SUNTANS grid from folder")
self.Bind(wx.EVT_MENU, self.on_load_grid, m_grid)
# Load a particle file
m_part = menu_file.Append(-1, "&Load PTM file\tCtrl-Shift-P",
"Load a PTM file")
self.Bind(wx.EVT_MENU, self.on_load_ptm, m_part)
# Save current scene as an animation
m_anim = menu_file.Append(-1,
"&Save animation of current scene\tCtrl-S",
"Save animation")
self.Bind(wx.EVT_MENU, self.on_save_anim, m_anim)
# Save the current figure
m_prin = menu_file.Append(-1, "&Print current scene\tCtrl-P",
"Save figure")
self.Bind(wx.EVT_MENU, self.on_save_fig, m_prin)
menu_file.AppendSeparator()
# Exit
m_exit = menu_file.Append(-1, "E&xit\tCtrl-X", "Exit")
self.Bind(wx.EVT_MENU, self.on_exit, m_exit)
###
# Tools menu
###
menu_tools = wx.Menu()
m_gridstat = menu_tools.Append(-1, "&Plot grid size statistics",
"SUNTANS grid size")
self.Bind(wx.EVT_MENU, self.on_plot_gridstat, m_gridstat)
m_countcells = menu_tools.Append(-1, "&Count # grid cells",
"Grid cell count")
self.Bind(wx.EVT_MENU, self.on_count_cells, m_countcells)
m_overlaybathy = menu_tools.Append(-1, "&Overlay depth contours",
"Depth overlay")
self.Bind(wx.EVT_MENU, self.on_overlay_bathy, m_overlaybathy)
###
# Help Menu
###
menu_help = wx.Menu()
m_about = menu_help.Append(-1, "&About\tF1", "About the demo")
self.Bind(wx.EVT_MENU, self.on_about, m_about)
# Add all of the menu bars
self.menubar.Append(menu_file, "&File")
self.menubar.Append(menu_tools, "&Tools")
self.menubar.Append(menu_help, "&Help")
self.SetMenuBar(self.menubar)
def create_main_panel(self):
""" Creates the main panel with all the controls on it:
* mpl canvas
* mpl navigation toolbar
* Control panel for interaction
"""
self.panel = wx.Panel(self)
# Create the mpl Figure and FigCanvas objects.
# 5x4 inches, 100 dots-per-inch
#
self.dpi = 100
#self.fig = Figure((7.0, 6.0), dpi=self.dpi,facecolor=self.bgcolor)
self.fig = Figure((7.0, 6.0), dpi=self.dpi)
self.canvas = FigCanvas(self.panel, -1, self.fig)
# Since we have only one plot, we can use add_axes
# instead of add_subplot, but then the subplot
# configuration tool in the navigation toolbar wouldn't
# work.
#
self.axes = self.fig.add_subplot(111)
#SetAxColor(self.axes,self.textcolor,self.bgcolor)
# Bind the 'pick' event for clicking on one of the bars
#
#self.canvas.mpl_connect('pick_event', self.on_pick)
########
# Create widgets
########
self.variable_list = wx.ComboBox(self.panel,
size=(200, -1),
choices=['Select a variable...'],
style=wx.CB_READONLY)
self.variable_list.Bind(wx.EVT_COMBOBOX, self.on_select_variable)
self.time_list = wx.ComboBox(self.panel,
size=(200, -1),
choices=['Select a time step...'],
style=wx.CB_READONLY)
self.time_list.Bind(wx.EVT_COMBOBOX, self.on_select_time)
self.depthlayer_list = wx.ComboBox(
self.panel,
size=(200, -1),
choices=['Select a vertical layer...'],
style=wx.CB_READONLY)
self.depthlayer_list.Bind(wx.EVT_COMBOBOX, self.on_select_depth)
self.show_edge_check = wx.CheckBox(self.panel,
-1,
"Show Edges",
style=wx.ALIGN_RIGHT)
self.show_edge_check.Bind(wx.EVT_CHECKBOX, self.on_show_edges)
if USECMOCEAN:
cmaps = []
for cmap in cm.cmapnames:
cmaps.append(cmap)
cmaps.append(cmap + '_r') # Add all reverse map options
else:
# Use matplotlib standard
cmaps = list(matplotlib.cm.datad.keys())
cmaps.sort()
self.colormap_list = wx.ComboBox(self.panel,
size=(100, -1),
choices=cmaps,
style=wx.CB_READONLY)
self.colormap_list.Bind(wx.EVT_COMBOBOX, self.on_select_cmap)
self.colormap_label = wx.StaticText(self.panel, -1, "Colormap:")
self.clim_check = wx.CheckBox(self.panel,
-1,
"Manual color limits ",
style=wx.ALIGN_RIGHT)
self.clim_check.Bind(wx.EVT_CHECKBOX, self.on_clim_check)
self.climlow = wx.TextCtrl(self.panel,
size=(100, -1),
style=wx.TE_PROCESS_ENTER)
self.climlow.Bind(wx.EVT_TEXT_ENTER, self.on_climlow)
self.climhigh = wx.TextCtrl(self.panel,
size=(100, -1),
style=wx.TE_PROCESS_ENTER)
self.climhigh.Bind(wx.EVT_TEXT_ENTER, self.on_climhigh)
# Labels
self.variable_label = wx.StaticText(self.panel,
-1,
"Variable:",
size=(200, -1))
self.time_label = wx.StaticText(self.panel,
-1,
"Time step:",
size=(200, -1))
self.depth_label = wx.StaticText(self.panel,
-1,
"Vertical level:",
size=(200, -1))
# Create the navigation toolbar, tied to the canvas
#
self.toolbar = NavigationToolbar(self.canvas)
#self.toolbar.toolitems[8][3]='my_save_fig'
#def my_save_fig(self,*args):
# print 'saving figure'
# return "break"
#########
# Layout with box sizers
#########
self.vbox = wx.BoxSizer(wx.VERTICAL)
self.vbox.Add(self.canvas, 1, wx.LEFT | wx.TOP | wx.GROW)
self.vbox.Add(self.toolbar, 0, wx.EXPAND)
self.vbox.AddSpacer(10)
#self.vbox.Add((-1,25))
flags = wx.ALIGN_LEFT | wx.ALL | wx.ALIGN_CENTER_VERTICAL
self.hbox0 = wx.BoxSizer(wx.HORIZONTAL)
self.hbox0.Add(self.show_edge_check, 0, border=10, flag=flags)
self.hbox0.Add(self.colormap_label, 0, border=10, flag=flags)
self.hbox0.Add(self.colormap_list, 0, border=10, flag=flags)
self.hbox0.Add(self.clim_check, 0, border=10, flag=flags)
self.hbox0.Add(self.climlow, 0, border=10, flag=flags)
self.hbox0.Add(self.climhigh, 0, border=10, flag=flags)
self.vbox.AddSpacer(5)
self.hbox1 = wx.BoxSizer(wx.HORIZONTAL)
self.hbox1.Add(self.variable_label, 0, border=10, flag=flags)
self.hbox1.Add(self.time_label, 0, border=10, flag=flags)
self.hbox1.Add(self.depth_label, 0, border=10, flag=flags)
self.vbox.AddSpacer(5)
self.hbox2 = wx.BoxSizer(wx.HORIZONTAL)
self.hbox2.Add(self.variable_list, 0, border=10, flag=flags)
self.hbox2.Add(self.time_list, 0, border=10, flag=flags)
self.hbox2.Add(self.depthlayer_list, 0, border=10, flag=flags)
self.vbox.Add(self.hbox1, 0, flag=wx.ALIGN_LEFT | wx.TOP)
self.vbox.Add(self.hbox2, 0, flag=wx.ALIGN_LEFT | wx.TOP)
self.vbox.Add(self.hbox0, 0, flag=wx.ALIGN_LEFT | wx.TOP)
self.panel.SetSizer(self.vbox)
self.vbox.Fit(self)
##########
# Event functions
##########
def create_figure(self):
"""
Creates the figure
"""
# Find the colorbar limits if unspecified
if self.autoclim:
self.clim = [self.data.min(), self.data.max()]
self.climlow.SetValue('%3.1f' % self.clim[0])
self.climhigh.SetValue('%3.1f' % self.clim[1])
if 'collection' in self.__dict__:
#self.collection.remove()
self.axes.collections.remove(self.collection)
else:
# First call - set the axes limits
self.axes.set_aspect('equal')
self.axes.set_xlim(self.xlims)
self.axes.set_ylim(self.ylims)
if self.collectiontype == 'cells':
self.collection = PolyCollection(self.xy, cmap=self.cmap)
self.collection.set_array(np.array(self.data[:]))
if not self.showedges:
self.collection.set_edgecolors(
self.collection.to_rgba(np.array((self.data[:]))))
elif self.collectiontype == 'edges':
xylines = [self.xp[self.edges], self.yp[self.edges]]
linesc = [
list(zip(xylines[0][ii, :], xylines[1][ii, :]))
for ii in range(self.Ne)
]
self.collection = LineCollection(linesc,
array=np.array(self.data[:]),
cmap=self.cmap)
self.collection.set_clim(vmin=self.clim[0], vmax=self.clim[1])
self.axes.add_collection(self.collection)
self.title = self.axes.set_title(self.genTitle(), color=self.textcolor)
self.axes.set_xlabel('Easting [m]')
self.axes.set_ylabel('Northing [m]')
# create a colorbar
if 'cbar' not in self.__dict__:
self.cbar = self.fig.colorbar(self.collection)
#SetAxColor(self.cbar.ax.axes,self.textcolor,self.bgcolor)
else:
#pass
print('Updating colorbar...')
#self.cbar.check_update(self.collection)
self.cbar.on_mappable_changed(self.collection)
self.canvas.draw()
def update_figure(self):
if self.autoclim:
self.clim = [self.data.min(), self.data.max()]
self.climlow.SetValue('%3.1f' % self.clim[0])
self.climhigh.SetValue('%3.1f' % self.clim[1])
else:
self.clim = [float(self.climlow.GetValue()),\
float(self.climhigh.GetValue())]
# check whether it is cell or edge type
if self.hasDim(self.variable, self.griddims['Ne']):
self.collectiontype = 'edges'
elif self.hasDim(self.variable, self.griddims['Nc']):
self.collectiontype = 'cells'
# Create a new figure if the variable has gone from cell to edge of vice
# versa
if not self.collectiontype == self.oldcollectiontype:
self.create_figure()
self.oldcollectiontype = self.collectiontype
self.collection.set_array(np.array(self.data[:]))
self.collection.set_clim(vmin=self.clim[0], vmax=self.clim[1])
# Cells only
if self.collectiontype == 'cells':
if not self.showedges:
self.collection.set_edgecolors(
self.collection.to_rgba(np.array((self.data[:]))))
else:
self.collection.set_edgecolors('k')
self.collection.set_linewidths(0.2)
# Update the title
self.title = self.axes.set_title(self.genTitle(), color=self.textcolor)
#Update the colorbar
self.cbar.update_normal(self.collection)
# redraw the figure
self.canvas.draw()
def on_pick(self, event):
# The event received here is of the type
# matplotlib.backend_bases.PickEvent
#
# It carries lots of information, of which we're using
# only a small amount here.
#
box_points = event.artist.get_bbox().get_points()
msg = "You've clicked on a bar with coords:\n %s" % box_points
dlg = wx.MessageDialog(self, msg, "Click!",
wx.OK | wx.ICON_INFORMATION)
dlg.ShowModal()
dlg.Destroy()
def on_select_variable(self, event):
vname = event.GetString()
self.flash_status_message("Selecting variable: %s" % vname)
# update the spatial object and load the data
self.variable = vname
self.loadData(variable=self.variable)
# Check if the variable has a depth coordinate
depthstr = ['']
# If so populate the vertical layer box
if self.hasDim(self.variable, self.griddims['Nk']):
depthstr = ['%3.1f' % self.z_r[k] for k in range(self.Nkmax)]
depthstr += ['surface', 'seabed']
elif self.hasDim(self.variable, 'Nkw'):
depthstr = ['%3.1f' % self.z_w[k] for k in range(self.Nkmax + 1)]
self.depthlayer_list.SetItems(depthstr)
# Update the plot
self.update_figure()
def on_select_time(self, event):
self.tindex = event.GetSelection()
# Update the object time index and reload the data
if self.plot_type == 'hydro':
if not self.tstep == self.tindex:
self.tstep = self.tindex
self.loadData()
self.flash_status_message("Selecting variable: %s..." %
event.GetString())
# Update the plot
self.update_figure()
elif self.plot_type == 'particles':
self.PTM.plot(self.tindex,ax=self.axes,\
xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
self.canvas.draw()
def on_select_depth(self, event):
kindex = event.GetSelection()
if not self.klayer[0] == kindex:
# Check if its the seabed or surface value
if kindex >= self.Nkmax:
kindex = event.GetString()
self.klayer = [kindex]
self.loadData()
self.flash_status_message("Selecting depth: %s..." %
event.GetString())
# Update the plot
self.update_figure()
def on_open_file(self, event):
file_choices = "SUNTANS NetCDF (*.nc)|*.nc*|UnTRIM NetCDF (*.nc)|*.nc*|All Files (*.*)|*.*"
dlg = wx.FileDialog(self,
message="Open SUNTANS file...",
defaultDir=os.getcwd(),
defaultFile="",
wildcard=file_choices,
style=wx.FD_MULTIPLE)
if dlg.ShowModal() == wx.ID_OK:
self.plot_type = 'hydro'
path = dlg.GetPaths()
# Initialise the class
if dlg.GetFilterIndex() == 0 or dlg.GetFilterIndex() > 1: #SUNTANS
self.flash_status_message("Opening SUNTANS file: %s" % path)
try:
Spatial.__init__(self, path, _FillValue=self._FillValue)
except:
Spatial.__init__(self, path, _FillValue=-999999)
startvar = 'dv'
if dlg.GetFilterIndex() == 1: #UnTRIM
self.flash_status_message("Opening UnTRIMS file: %s" % path)
#Spatial.__init__(self,path,gridvars=untrim_gridvars,griddims=untrim_griddims)
UNTRIMSpatial.__init__(self, path)
startvar = 'Mesh2_face_depth'
# Populate the drop down menus
vnames = self.listCoordVars()
self.variable_list.SetItems(vnames)
# Update the time drop down list
if 'time' in self.__dict__:
self.timestr = [
datetime.strftime(tt, '%d-%b-%Y %H:%M:%S')
for tt in self.time
]
else:
# Assume that it is a harmonic-type file
self.timestr = self.nc.Constituent_Names.split()
self.time_list.SetItems(self.timestr)
# Draw the depth
if startvar in vnames:
self.variable = startvar
self.loadData()
self.create_figure()
def on_load_grid(self, event):
dlg = wx.DirDialog(self,
message="Open SUNTANS grid from folder...",
defaultPath=os.getcwd(),
style=wx.DD_DEFAULT_STYLE)
if dlg.ShowModal() == wx.ID_OK:
path = dlg.GetPath()
# Initialise the class
self.flash_status_message("Opening SUNTANS grid from folder: %s" %
path)
Grid.__init__(self, path)
# Plot the Grid
if 'collection' in self.__dict__:
self.axes.collections.remove(self.collection)
self.axes, self.collection = self.plotmesh(ax=self.axes,
edgecolors='y')
# redraw the figure
self.canvas.draw()
def on_load_ptm(self, event):
file_choices = "PTM NetCDF (*.nc)|*.nc|PTM Binary (*_bin.out)|*_bin.out|All Files (*.*)|*.*"
dlg = wx.FileDialog(self,
message="Open PTM file...",
defaultDir=os.getcwd(),
defaultFile="",
wildcard=file_choices,
style=wx.FD_MULTIPLE)
if dlg.ShowModal() == wx.ID_OK:
self.plot_type = 'particles'
path = dlg.GetPath()
# Initialise the class
if dlg.GetFilterIndex() == 0: #SUNTANS
self.flash_status_message("Opening PTM netcdf file: %s" % path)
self.PTM = PtmNC(path)
elif dlg.GetFilterIndex() == 1: #PTM
self.flash_status_message("Opening PTM binary file: %s" % path)
self.PTM = PtmBin(path)
self.Nt = self.PTM.nt
# Update the time drop down list
self.timestr = [
datetime.strftime(tt, '%d-%b-%Y %H:%M:%S')
for tt in self.PTM.time
]
self.time_list.SetItems(self.timestr)
# Plot the first time step
if 'xlims' in self.__dict__:
self.PTM.plot(self.PTM.nt-1,ax=self.axes,xlims=self.xlims,\
ylims=self.ylims,color=self.particlecolor,\
fontcolor='w',markersize=self.particlesize)
else:
self.PTM.plot(self.PTM.nt-1,ax=self.axes,fontcolor='w',\
color=self.particlecolor,markersize=self.particlesize)
# redraw the figure
self.canvas.draw()
def on_show_edges(self, event):
sender = event.GetEventObject()
self.showedges = sender.GetValue()
# Update the figure
self.update_figure()
def on_clim_check(self, event):
sender = event.GetEventObject()
if sender.GetValue() == True:
self.autoclim = False
self.update_figure()
else:
self.autoclim = True
def on_climlow(self, event):
self.clim[0] = event.GetString()
#self.update_figure()
def on_climhigh(self, event):
self.clim[1] = event.GetString()
#self.update_figure()
def on_select_cmap(self, event):
self.cmap = event.GetString()
if USECMOCEAN:
self.collection.set_cmap(getattr(cm, self.cmap))
else:
self.collection.set_cmap(self.cmap)
# Update the figure
self.update_figure()
def on_save_fig(self, event):
"""
Save a figure of the current scene to a file
"""
file_choices = " (*.png)|*.png| (*.pdf)|*.pdf |(*.jpg)|*.jpg |(*.eps)|*eps "
filters = ['.png', '.pdf', '.png', '.png']
dlg = wx.FileDialog(self,
message="Save figure to file...",
defaultDir=os.getcwd(),
defaultFile="",
wildcard=file_choices,
style=wx.FD_SAVE | wx.FD_OVERWRITE_PROMPT)
if dlg.ShowModal() == wx.ID_OK:
path = dlg.GetPath()
ext = filters[dlg.GetFilterIndex()]
if ext in path:
outfile = path
else:
outfile = path + ext
self.fig.savefig(outfile)
def on_save_anim(self, event):
"""
Save an animation of the current scene to a file
"""
file_choices = "Quicktime (*.mov)|*.mov| (*.gif)|*.gif| (*.avi)|*.avi |(*.mp4)|*.mp4 "
filters = ['.mov', '.gif', '.avi', '.mp4']
dlg = wx.FileDialog(self,
message="Output animation file...",
defaultDir=os.getcwd(),
defaultFile="",
wildcard=file_choices,
style=wx.FD_SAVE | wx.FD_OVERWRITE_PROMPT)
if dlg.ShowModal() == wx.ID_OK:
path = dlg.GetPath()
ext = filters[dlg.GetFilterIndex()]
if ext in path:
outfile = path
else:
outfile = path + ext
self.flash_status_message("Saving figure to file: %s" % outfile)
self.flash_status_message("Saving animation to file: %s" % outfile)
# Create the animation
#self.tstep = range(self.Nt) # Use all time steps for animation
#self.animate(cbar=self.cbar,cmap=self.cmap,\
# xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
def initanim():
if not self.plot_type == 'particles':
return (self.title, self.collection)
else:
return (self.PTM.title, self.PTM.p_handle)
def updateScalar(i):
if not self.plot_type == 'particles':
self.tstep = [i]
self.loadData()
self.update_figure()
return (self.title, self.collection)
elif self.plot_type == 'particles':
self.PTM.plot(i,ax=self.axes,\
xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
return (self.PTM.title, self.PTM.p_handle)
self.anim = animation.FuncAnimation(self.fig, \
updateScalar, init_func = initanim, frames=self.Nt, interval=50, blit=True)
if ext == '.gif':
self.anim.save(outfile, writer='imagemagick', fps=6)
elif ext == '.mp4':
print('Saving html5 video...')
# Ensures html5 compatibility
self.anim.save(outfile,writer='mencoder',fps=6,\
bitrate=3600,extra_args=['-ovc','x264']) # mencoder options
#bitrate=3600,extra_args=['-vcodec','libx264'])
else:
self.anim.save(outfile, writer='mencoder', fps=6, bitrate=3600)
# Return the figure back to its status
del self.anim
self.tstep = self.tindex
if not self.plot_type == 'particles':
self.loadData()
self.update_figure()
# Bring up a dialog box
dlg2 = wx.MessageDialog(self, 'Animation complete.', "Done", wx.OK)
dlg2.ShowModal()
dlg2.Destroy()
def on_exit(self, event):
self.Destroy()
def on_about(self, event):
msg = """ SUNTANS NetCDF visualization tool
*Author: Matt Rayson
*Institution: Stanford University
*Created: October 2013
"""
dlg = wx.MessageDialog(self, msg, "About", wx.OK)
dlg.ShowModal()
dlg.Destroy()
def on_count_cells(self, eveny):
msg = "Total 3-D grid cells = %d" % (self.count_cells())
dlg = wx.MessageDialog(self, msg, "No. cells", wx.OK)
dlg.ShowModal()
dlg.Destroy()
def on_overlay_bathy(self, event):
# Plot depth contours
print('Plotting contours...')
self.contourf(z=self.dv, clevs=self.depthlevs,\
ax=self.axes,\
filled=False, colors='0.5', linewidths=0.5, zorder=1e6)
print('Done')
def on_plot_gridstat(self, event):
"""
Plot the grid size histogram in a new figure
"""
matplotlib.pyplot.figure()
self.plothist()
matplotlib.pyplot.show()
def create_status_bar(self):
self.statusbar = self.CreateStatusBar()
def flash_status_message(self, msg, flash_len_ms=1500):
self.statusbar.SetStatusText(msg)
self.timeroff = wx.Timer(self)
self.Bind(wx.EVT_TIMER, self.on_flash_status_off, self.timeroff)
self.timeroff.Start(flash_len_ms, oneShot=True)
def on_flash_status_off(self, event):
self.statusbar.SetStatusText('')
def draw_networkx_edges(G, pos,
edgelist=None,
width=1.0,
edge_color='k',
style='solid',
alpha=1.0,
arrowstyle='-|>',
arrowsize=10,
edge_cmap=None,
edge_vmin=None,
edge_vmax=None,
ax=None,
arrows=True,
label=None,
node_size=300,
nodelist=None,
node_shape="o",
connectionstyle=None,
**kwds):
"""Draw the edges of the graph G.
This draws only the edges of the graph G.
Parameters
----------
G : graph
A networkx graph
pos : dictionary
A dictionary with nodes as keys and positions as values.
Positions should be sequences of length 2.
edgelist : collection of edge tuples
Draw only specified edges(default=G.edges())
width : float, or array of floats
Line width of edges (default=1.0)
edge_color : color string, or array of floats
Edge color. Can be a single color format string (default='r'),
or a sequence of colors with the same length as edgelist.
If numeric values are specified they will be mapped to
colors using the edge_cmap and edge_vmin,edge_vmax parameters.
style : string
Edge line style (default='solid') (solid|dashed|dotted,dashdot)
alpha : float
The edge transparency (default=1.0)
edge_ cmap : Matplotlib colormap
Colormap for mapping intensities of edges (default=None)
edge_vmin,edge_vmax : floats
Minimum and maximum for edge colormap scaling (default=None)
ax : Matplotlib Axes object, optional
Draw the graph in the specified Matplotlib axes.
arrows : bool, optional (default=True)
For directed graphs, if True draw arrowheads.
Note: Arrows will be the same color as edges.
arrowstyle : str, optional (default='-|>')
For directed graphs, choose the style of the arrow heads.
See :py:class: `matplotlib.patches.ArrowStyle` for more
options.
arrowsize : int, optional (default=10)
For directed graphs, choose the size of the arrow head head's length and
width. See :py:class: `matplotlib.patches.FancyArrowPatch` for attribute
`mutation_scale` for more info.
connectionstyle : str, optional (default=None)
Pass the connectionstyle parameter to create curved arc of rounding
radius rad. For example, connectionstyle='arc3,rad=0.2'.
See :py:class: `matplotlib.patches.ConnectionStyle` and
:py:class: `matplotlib.patches.FancyArrowPatch` for more info.
label : [None| string]
Label for legend
Returns
-------
matplotlib.collection.LineCollection
`LineCollection` of the edges
list of matplotlib.patches.FancyArrowPatch
`FancyArrowPatch` instances of the directed edges
Depending whether the drawing includes arrows or not.
Notes
-----
For directed graphs, arrows are drawn at the head end. Arrows can be
turned off with keyword arrows=False. Be sure to include `node_size` as a
keyword argument; arrows are drawn considering the size of nodes.
Examples
--------
>>> G = nx.dodecahedral_graph()
>>> edges = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> G = nx.DiGraph()
>>> G.add_edges_from([(1, 2), (1, 3), (2, 3)])
>>> arcs = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
>>> alphas = [0.3, 0.4, 0.5]
>>> for i, arc in enumerate(arcs): # change alpha values of arcs
... arc.set_alpha(alphas[i])
Also see the NetworkX drawing examples at
https://networkx.github.io/documentation/latest/auto_examples/index.html
See Also
--------
draw()
draw_networkx()
draw_networkx_nodes()
draw_networkx_labels()
draw_networkx_edge_labels()
"""
try:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.cbook as cb
from matplotlib.colors import colorConverter, Colormap, Normalize
from matplotlib.collections import LineCollection
from matplotlib.patches import FancyArrowPatch
import numpy as np
except ImportError:
raise ImportError("Matplotlib required for draw()")
except RuntimeError:
print("Matplotlib unable to open display")
raise
if ax is None:
ax = plt.gca()
if edgelist is None:
edgelist = list(G.edges())
if not edgelist or len(edgelist) == 0: # no edges!
return None
if nodelist is None:
nodelist = list(G.nodes())
# set edge positions
edge_pos = np.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])
if not cb.iterable(width):
lw = (width,)
else:
lw = width
if not is_string_like(edge_color) \
and cb.iterable(edge_color) \
and len(edge_color) == len(edge_pos):
if np.alltrue([is_string_like(c) for c in edge_color]):
# (should check ALL elements)
# list of color letters such as ['k','r','k',...]
edge_colors = tuple([colorConverter.to_rgba(c, alpha)
for c in edge_color])
elif np.alltrue([not is_string_like(c) for c in edge_color]):
# If color specs are given as (rgb) or (rgba) tuples, we're OK
if np.alltrue([cb.iterable(c) and len(c) in (3, 4)
for c in edge_color]):
edge_colors = tuple(edge_color)
else:
# numbers (which are going to be mapped with a colormap)
edge_colors = None
else:
raise ValueError('edge_color must contain color names or numbers')
else:
if is_string_like(edge_color) or len(edge_color) == 1:
edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
else:
msg = 'edge_color must be a color or list of one color per edge'
raise ValueError(msg)
if (not G.is_directed() or not arrows):
edge_collection = LineCollection(edge_pos,
colors=edge_colors,
linewidths=lw,
antialiaseds=(1,),
linestyle=style,
transOffset=ax.transData,
)
edge_collection.set_zorder(1) # edges go behind nodes
edge_collection.set_label(label)
ax.add_collection(edge_collection)
# Note: there was a bug in mpl regarding the handling of alpha values
# for each line in a LineCollection. It was fixed in matplotlib by
# r7184 and r7189 (June 6 2009). We should then not set the alpha
# value globally, since the user can instead provide per-edge alphas
# now. Only set it globally if provided as a scalar.
if isinstance(alpha, Number):
edge_collection.set_alpha(alpha)
if edge_colors is None:
if edge_cmap is not None:
assert(isinstance(edge_cmap, Colormap))
edge_collection.set_array(np.asarray(edge_color))
edge_collection.set_cmap(edge_cmap)
if edge_vmin is not None or edge_vmax is not None:
edge_collection.set_clim(edge_vmin, edge_vmax)
else:
edge_collection.autoscale()
return edge_collection
arrow_collection = None
if G.is_directed() and arrows:
# Note: Waiting for someone to implement arrow to intersection with
# marker. Meanwhile, this works well for polygons with more than 4
# sides and circle.
def to_marker_edge(marker_size, marker):
if marker in "s^>v<d": # `large` markers need extra space
return np.sqrt(2 * marker_size) / 2
else:
return np.sqrt(marker_size) / 2
# Draw arrows with `matplotlib.patches.FancyarrowPatch`
arrow_collection = []
mutation_scale = arrowsize # scale factor of arrow head
arrow_colors = edge_colors
if arrow_colors is None:
if edge_cmap is not None:
assert(isinstance(edge_cmap, Colormap))
else:
edge_cmap = plt.get_cmap() # default matplotlib colormap
if edge_vmin is None:
edge_vmin = min(edge_color)
if edge_vmax is None:
edge_vmax = max(edge_color)
color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)
for i, (src, dst) in enumerate(edge_pos):
x1, y1 = src
x2, y2 = dst
arrow_color = None
line_width = None
shrink_source = 0 # space from source to tail
shrink_target = 0 # space from head to target
if cb.iterable(node_size): # many node sizes
src_node, dst_node = edgelist[i][:2]
index_node = nodelist.index(dst_node)
marker_size = node_size[index_node]
shrink_target = to_marker_edge(marker_size, node_shape)
else:
shrink_target = to_marker_edge(node_size, node_shape)
if arrow_colors is None:
arrow_color = edge_cmap(color_normal(edge_color[i]))
elif len(arrow_colors) > 1:
arrow_color = arrow_colors[i]
else:
arrow_color = arrow_colors[0]
if len(lw) > 1:
line_width = lw[i]
else:
line_width = lw[0]
arrow = FancyArrowPatch((x1, y1), (x2, y2),
arrowstyle=arrowstyle,
shrinkA=shrink_source,
shrinkB=shrink_target,
mutation_scale=mutation_scale,
color=arrow_color,
linewidth=line_width,
connectionstyle=connectionstyle,
zorder=1) # arrows go behind nodes
# There seems to be a bug in matplotlib to make collections of
# FancyArrowPatch instances. Until fixed, the patches are added
# individually to the axes instance.
arrow_collection.append(arrow)
ax.add_patch(arrow)
# update view
minx = np.amin(np.ravel(edge_pos[:, :, 0]))
maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
miny = np.amin(np.ravel(edge_pos[:, :, 1]))
maxy = np.amax(np.ravel(edge_pos[:, :, 1]))
w = maxx - minx
h = maxy - miny
padx, pady = 0.05 * w, 0.05 * h
corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
ax.update_datalim(corners)
ax.autoscale_view()
plt.tick_params(
axis='both',
which='both',
bottom=False,
left=False,
labelbottom=False,
labelleft=False)
return arrow_collection
def plot_linestring_collection(ax, geoms, colors_or_values, plot_values,
vmin=None, vmax=None, cmap=None,
linewidth=1.0, **kwargs):
"""
Plots a collection of LineString and MultiLineString geometries to `ax`
Parameters
----------
ax : matplotlib.axes.Axes
where shapes will be plotted
geoms : a sequence of `N` LineStrings and/or MultiLineStrings (can be mixed)
colors_or_values : a sequence of `N` values or RGBA tuples
It should have 1:1 correspondence with the geometries (not their components).
plot_values : bool
If True, `colors_or_values` is interpreted as a list of values, and will
be mapped to colors using vmin/vmax/cmap (which become required).
Otherwise `colors_or_values` is interpreted as a list of colors.
Returns
-------
collection : matplotlib.collections.Collection that was plotted
"""
from matplotlib.collections import LineCollection
components, component_colors_or_values = _flatten_multi_geoms(
geoms, colors_or_values)
# LineCollection does not accept some kwargs.
if 'markersize' in kwargs:
del kwargs['markersize']
segments = [np.array(linestring)[:, :2] for linestring in components]
collection = LineCollection(segments,
linewidth=linewidth, **kwargs)
if plot_values:
collection.set_array(np.array(component_colors_or_values))
collection.set_cmap(cmap)
collection.set_clim(vmin, vmax)
else:
# set_color magically sets the correct combination of facecolor and
# edgecolor, based on collection type.
collection.set_color(component_colors_or_values)
# If the user set facecolor and/or edgecolor explicitly, the previous
# call to set_color might have overridden it (remember, the 'color' may
# have come from plot_series, not from the user). The user should be
# able to override matplotlib's default behavior, by setting them again
# after set_color.
if 'facecolor' in kwargs:
collection.set_facecolor(kwargs['facecolor'])
if 'edgecolor' in kwargs:
collection.set_edgecolor(kwargs['edgecolor'])
ax.add_collection(collection, autolim=True)
ax.autoscale_view()
return collection
