def do_3d_projection(self, renderer):
xs, ys, zs = self._offsets3d
vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
fcs = (_zalpha(self._facecolor3d, vzs) if self._depthshade else
self._facecolor3d)
fcs = mcolors.to_rgba_array(fcs, self._alpha)
self.set_facecolors(fcs)
ecs = (_zalpha(self._edgecolor3d, vzs) if self._depthshade else
self._edgecolor3d)
# Sort the points based on z coordinates
# Performance optimization: Create a sorted index array and reorder
# points and point properties according to the index array
z_markers_idx = np.argsort(vzs)[::-1]
# Re-order items
vzs = vzs[z_markers_idx]
vxs = vxs[z_markers_idx]
vys = vys[z_markers_idx]
fcs = fcs[z_markers_idx]
ecs = ecs[z_markers_idx]
vps = np.column_stack((vxs, vys))
fcs = mcolors.to_rgba_array(fcs, self._alpha)
ecs = mcolors.to_rgba_array(ecs, self._alpha)
self.set_edgecolors(ecs)
self.set_facecolors(fcs)
PathCollection.set_offsets(self, vps)
return np.min(vzs) if vzs.size else np.nan
def _make_barb(self, temperature, theta, speed, angle):
"""Add the barb to the plot at the specified location."""
u, v = self._uv(speed, angle)
if 0 < speed < _BARB_BINS[0]:
# Plot the missing barbless 1-2 knots line.
length = self._kwargs['length']
pivot_points = dict(tip=0.0, middle=-length / 2.)
pivot = self._kwargs.get('pivot', 'tip')
offset = pivot_points[pivot]
verts = [(0.0, offset), (0.0, length + offset)]
rangle = math.radians(-angle)
verts = mtransforms.Affine2D().rotate(rangle).transform(verts)
codes = [Path.MOVETO, Path.LINETO]
path = Path(verts, codes)
size = length ** 2 / 4
xy = np.array([[temperature, theta]])
barb = PathCollection([path], (size,), offsets=xy,
transOffset=self._transform,
**self._custom_kwargs)
barb.set_transform(mtransforms.IdentityTransform())
self.axes.add_collection(barb)
else:
barb = plt.barbs(temperature, theta, u, v,
transform=self._transform, **self._kwargs)
return barb
def plot_map_native(lon_corners, lat_corners, data, **kwargs):
from matplotlib import pyplot, patches, path
import numpy
# Fix longitude coordinates; we need longitudes to run from
# -180 to 180, not from 0 to 360
lon_corners.values = numpy.where(lon_corners > 180, lon_corners - 360,
lon_corners)
# Loop over GLL nodes and create paths that describe the boundaries
# of each node; collect these into a list to plot all at once
path_list = []
for icol in range(lon_corners.shape[0]):
# Get corners for this node
x_corners = lon_corners[icol, :].values
y_corners = lat_corners[icol, :].values
# Repeat first vertex at end of array to close the path
x_corners = numpy.append(x_corners, x_corners[0])
y_corners = numpy.append(y_corners, y_corners[0])
# Create paths connecting the corners and append to list
vertices = numpy.column_stack([x_corners, y_corners])
path_list.append(path.Path(vertices, closed=True))
# Plot collection of patches
from matplotlib.collections import PathCollection
collection = PathCollection(path_list, transform=crs.Geodetic(), **kwargs)
collection.set_array(data)
ax = pyplot.gca()
pl = ax.add_collection(collection)
return pl
def __init__(self, simplices=None, **kwargs):
kwargs.setdefault('linewidths', 0.5)
kwargs.setdefault('edgecolors', 'k')
kwargs.setdefault('facecolors', (0.8, 0.9, 1.0))
PathCollection.__init__(self, [], **kwargs)
if simplices is not None:
self.set_simplices(simplices)
def __init__(self, simplices=None, **kwargs):
kwargs.setdefault('linewidths', 0.5)
kwargs.setdefault('edgecolors', 'k')
kwargs.setdefault('facecolors', (0.8, 0.9, 1.0))
PathCollection.__init__(self, [], **kwargs)
if simplices is not None:
self.set_simplices(simplices)
def do_3d_projection(self, renderer=None):
xs, ys, zs = self._offsets3d
vxs, vys, vzs, vis = proj3d.proj_transform_clip(
xs, ys, zs, self.axes.M)
# Sort the points based on z coordinates
# Performance optimization: Create a sorted index array and reorder
# points and point properties according to the index array
z_markers_idx = self._z_markers_idx = np.argsort(vzs)[::-1]
self._vzs = vzs
# we have to special case the sizes because of code in collections.py
# as the draw method does
# self.set_sizes(self._sizes, self.figure.dpi)
# so we can not rely on doing the sorting on the way out via get_*
if len(self._sizes3d) > 1:
self._sizes = self._sizes3d[z_markers_idx]
if len(self._linewidths3d) > 1:
self._linewidths = self._linewidths3d[z_markers_idx]
# Re-order items
vzs = vzs[z_markers_idx]
vxs = vxs[z_markers_idx]
vys = vys[z_markers_idx]
PathCollection.set_offsets(self, np.column_stack((vxs, vys)))
return np.min(vzs) if vzs.size else np.nan
def _make_barb(self, temperature, theta, speed, angle):
"""Add the barb to the plot at the specified location."""
u, v = self._uv(speed, angle)
if 0 < speed < _BARB_BINS[0]:
# Plot the missing barbless 1-2 knots line.
length = self._kwargs['length']
pivot_points = dict(tip=0.0, middle=-length / 2.)
pivot = self._kwargs.get('pivot', 'tip')
offset = pivot_points[pivot]
verts = [(0.0, offset), (0.0, length + offset)]
rangle = math.radians(-angle)
verts = mtransforms.Affine2D().rotate(rangle).transform(verts)
codes = [Path.MOVETO, Path.LINETO]
path = Path(verts, codes)
size = length**2 / 4
xy = np.array([[temperature, theta]])
barb = PathCollection([path], (size, ),
offsets=xy,
transOffset=self._transform,
**self._custom_kwargs)
barb.set_transform(mtransforms.IdentityTransform())
self.axes.add_collection(barb)
else:
barb = plt.barbs(temperature,
theta,
u,
v,
transform=self._transform,
**self._kwargs)
return barb
def makeplot(field2plot, lon_ts, lat_ts, vmin, vmax, m, findex, shp_info):
norm = colors.Normalize(vmin=vmin, vmax=vmax)
bounds = np.arange(vmin, vmax + .0001, dvar)
ax = plt.gca()
ax.cla()
paths = []
for line in shp_info[4]._paths:
paths.append(Path(line.vertices, codes=line.codes))
coll = PathCollection(paths, linewidths=0, facecolors='black', zorder=2)
ax.add_collection(coll)
m.drawcoastlines(ax=ax, linewidth=0.05)
m.drawcountries(ax=ax, linewidth=0.05)
pcm = plt.pcolormesh(lon_ts, lat_ts, field2plot, norm=norm, cmap=cmap)
# cbar = plt.colorbar(pcm, norm=norm, cmap=cmap, orientation='vertical', pad=0.05, aspect=15,
# shrink=0.7, extend='both')
# cbar.set_ticks(bounds)
figname = str(findex).zfill(5)
plt.savefig(figdir + figname, facecolor="black", dpi=150)
plt.close()
print findex
return findex
def as_patch(self, **kwds):
""" matplotlib pathpatch for hlines and vlines
**pathkwds : any valid PathCollection keyword (linestyle, edgecolor)
Notes
-----
First, finds coordinates of hlines, vlines. Then, it it finds
the x, y coords of each line; it extracts the points corresponding to
each isoline, and makes a path for each of these."""
hlines = zip(*np.where(self.hlines))
vlines = zip(*np.where(self.vlines))
xunique = set([x for x, y in hlines])
yunique = set([y for x, y in vlines])
hpaths = []
vpaths = []
# REVERSE PATHS FOR SAKE OF PLOT (y,x) !!
for xun in xunique:
hpaths.append(Path([(y, x) for x, y in hlines if x == xun]))
for yun in yunique:
vpaths.append(Path([(y, x) for x, y in vlines if y == yun]))
return PathCollection(hpaths + vpaths, **kwds)
def do_3d_projection(self, renderer):
xs, ys, zs = self._offsets3d
vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
fcs = (zalpha(self._facecolor3d, vzs) if self._depthshade else
self._facecolor3d)
fcs = mcolors.to_rgba_array(fcs, self._alpha)
self.set_facecolors(fcs)
ecs = (zalpha(self._edgecolor3d, vzs) if self._depthshade else
self._edgecolor3d)
ecs = mcolors.to_rgba_array(ecs, self._alpha)
self.set_edgecolors(ecs)
PathCollection.set_offsets(self, np.column_stack([vxs, vys]))
return np.min(vzs) if vzs.size else np.nan
def _add_bondlength_change_collection(self,
index,
threshold=0,
zorder=20,
**kwargs):
"""Comute and draw bondlength changes on the axes."""
codes = [Path.MOVETO, Path.LINETO]
col = []
amp = []
colors = []
for obbond in ob.OBMolBondIter(self.molecule):
atom1, atom2 = (self.molecule.GetAtom(obbond.GetBeginAtomIdx()),
self.molecule.GetAtom(obbond.GetEndAtomIdx()))
atom1nc, atom2nc = [
self._to_normal_coordinates(atom, index)
for atom in (atom1, atom2)
]
if obbond.GetLength() == 0.0:
logger.error(
"Bond between %i and %i with length %.1f ignored." %
(atom1.GetIdx(), atom2.GetIdx(), obbond.GetLength()))
continue
amplitude = atom1.GetDistance(atom2) - atom1nc.GetDistance(atom2nc)
if abs(amplitude * 100) <= threshold: continue
amp.append(abs(amplitude * 50))
colors.append(self.bond_colors[0 if amplitude < 0.0 else 1])
verts = (self._2Dcoords(atom1), self._2Dcoords(atom2))
segment = Path(verts, codes)
col.append(segment)
lw = 0.0 if not col else np.array(amp) * kwargs.pop("lw", 1.0)
kw = {'edgecolors': colors, 'linewidths': lw}
kwargs.update(kw)
self._vib_bonds = PathCollection(col, zorder=zorder, **kwargs)
self.axes.add_collection(self._vib_bonds)
def set_subplot_map(self, highlights, cgeorange):
###SUBPLOT: SET MAP
self.oldax = self.p.ax
self.p.ax = self.p.fig.add_axes([0.02, -0.39, 1.04, 0.44])
self.p.setmap(projection='cyl', georange=cgeorange, resolution='i')
self.p.draw('country coastline province city')
###SUBPLOT: HIGHLIGHT COASTLINES & PLOT LEGEND
colors = ['#AAFFAA', '#FFFF44', '#FF3333', '#BB0044']
descr = ['10~25%', '25~50%', '50~75%', '>75%']
handles = []
for i, clr in enumerate(colors, 0):
patch = PathCollection(geos_to_path(
MultiLineString(highlights[i]).buffer(self.buffersize)),
facecolor=clr)
self.p.ax.add_collection(patch)
handles.append(Patch(color=clr, label=descr[i]))
self.p.ax.text(0.98,
0.27,
'中心经过1经纬度\n范围内的几率',
transform=self.p.ax.transAxes,
va='bottom',
ha='right',
fontsize=6,
family='Source Han Sans CN')
self.p.legend(handles=handles, loc='lower right', framealpha=0.8)
def PolygonPath(shape):
"""Constructs a compound matplotlib path from a Shapely or GeoJSON-like
geometric object"""
if isinstance(shape, (sgeom.LineString, sgeom.Point)):
return Path(np.vstack(shape.xy).T)
elif isinstance(shape, sgeom.Polygon):
codes, vertices = PolygonCodes(shape)
return Path(vertices, codes)
elif isinstance(
shape,
(sgeom.MultiPolygon, sgeom.MultiLineString, sgeom.MultiPoint)):
codes, vertices = [], []
for poly in shape:
sub_codes, sub_vertices = PolygonCodes(poly)
codes.append(sub_codes)
vertices.append(sub_vertices)
codes = np.concatenate(codes)
vertices = np.concatenate(vertices)
return Path(vertices, codes)
elif isinstance(shape, sgeom.GeometryCollection):
return PathCollection([PolygonPath(geom) for geom in shape.geoms])
else:
raise ValueError('Unsupported shape type {}.'.format(type(shape)))
def do_3d_projection(self, renderer):
xs, ys, zs = self._offsets3d
vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
fcs = (_zalpha(self._facecolor3d, vzs)
if self._depthshade else self._facecolor3d)
fcs = mcolors.to_rgba_array(fcs, self._alpha)
self.set_facecolors(fcs)
ecs = (_zalpha(self._edgecolor3d, vzs)
if self._depthshade else self._edgecolor3d)
ecs = mcolors.to_rgba_array(ecs, self._alpha)
self.set_edgecolors(ecs)
PathCollection.set_offsets(self, np.column_stack([vxs, vys]))
return np.min(vzs) if vzs.size else np.nan
def _add_bond_collection(self, zorder=10, **kwargs):
"""Draw molecule skeleton on the axes."""
col = []
codes = [
Path.MOVETO,
Path.LINETO,
] # segment
for obbond in ob.OBMolBondIter(self.molecule):
atom1, atom2 = (self.molecule.GetAtom(obbond.GetBeginAtomIdx()),
self.molecule.GetAtom(obbond.GetEndAtomIdx()))
verts = self._2Dcoords(atom1), self._2Dcoords(atom2)
segment = Path(verts, codes)
col.append(segment)
kw = {'edgecolors': 'k'}
kwargs.update(kw)
self._mol_bonds = PathCollection(col, zorder=zorder, **kwargs)
self.axes.add_collection(self._mol_bonds)
def do_3d_projection(self, renderer):
xs, ys, zs = self._offsets3d
vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
fcs = zalpha(self._facecolor3d, vzs) if self._depthshade else self._facecolor3d
fcs = mcolors.colorConverter.to_rgba_array(fcs, self._alpha)
self.set_facecolors(fcs)
ecs = zalpha(self._edgecolor3d, vzs) if self._depthshade else self._edgecolor3d
ecs = mcolors.colorConverter.to_rgba_array(ecs, self._alpha)
self.set_edgecolors(ecs)
PathCollection.set_offsets(self, list(zip(vxs, vys)))
if vzs.size > 0:
return min(vzs)
else:
return np.nan
def _get_patch_collections(self, axes):
for key, collection in self._collection_map.items():
marker = self._markers[key]
coll = PathCollection([marker.get_path()],
offsets=collection,
transOffset=axes.transAxes,
facecolors=[self._colors[key]])
yield coll
def do_3d_projection(self, renderer=None):
# see _update_scalarmappable docstring for why this must be here
_update_scalarmappable(self)
xs, ys, zs = self._offsets3d
vxs, vys, vzs, vis = proj3d.proj_transform_clip(
xs, ys, zs, self.axes.M)
fcs = (_zalpha(self._facecolor3d, vzs)
if self._depthshade else self._facecolor3d)
ecs = (_zalpha(self._edgecolor3d, vzs)
if self._depthshade else self._edgecolor3d)
sizes = self._sizes3d
lws = self._linewidth3d
# Sort the points based on z coordinates
# Performance optimization: Create a sorted index array and reorder
# points and point properties according to the index array
z_markers_idx = np.argsort(vzs)[::-1]
# Re-order items
vzs = vzs[z_markers_idx]
vxs = vxs[z_markers_idx]
vys = vys[z_markers_idx]
if len(fcs) > 1:
fcs = fcs[z_markers_idx]
if len(ecs) > 1:
ecs = ecs[z_markers_idx]
if len(sizes) > 1:
sizes = sizes[z_markers_idx]
if len(lws) > 1:
lws = lws[z_markers_idx]
vps = np.column_stack((vxs, vys))
fcs = mcolors.to_rgba_array(fcs, self._alpha)
ecs = mcolors.to_rgba_array(ecs, self._alpha)
super().set_edgecolor(ecs)
super().set_facecolor(fcs)
super().set_sizes(sizes)
super().set_linewidth(lws)
PathCollection.set_offsets(self, vps)
return np.min(vzs) if vzs.size else np.nan
def on_ax(self, ax, pos, orientation):
trans = np.reshape(orientation, (2, 2)).T
for offset, normal, text, size, halign, valign, color in self.texts:
normal = np.matmul(normal, trans)
angle = math.atan2(-normal[1], normal[0]) * 180 / math.pi
x, y = (pos + np.matmul(offset, trans)) * Y_TRANS
if angle >= 180 or angle < 0:
if halign == 'right':
halign = 'left'
elif halign == 'left':
halign = 'right'
if valign == 'top':
valign = 'bottom'
elif valign == 'bottom':
valign = 'top'
angle -= 180
kw = {
'x': x,
'y': y,
'color': color,
# 'font': 'sans-serif',
'fontsize': size,
's': text,
'rotation': angle,
'rotation_mode': 'anchor',
'horizontalalignment': halign,
'verticalalignment': valign,
'clip_on': True,
}
ax.text(**kw)
if self.paths:
d = {}
d['paths'] = []
d['linewidths'] = []
d['edgecolors'] = []
d['facecolors'] = []
for path, linewidth, edgecolor, facecolor in self.paths:
verts = path.vertices
verts = np.matmul(verts, trans)
verts = (pos + verts) * Y_TRANS
path.vertices = verts
d['paths'].append(path)
d['linewidths'].append(linewidth)
d['edgecolors'].append(edgecolor)
d['facecolors'].append(facecolor)
pc = PathCollection(offset_position='data', **d)
ax.add_collection(pc)
def do_3d_projection(self, renderer):
xs, ys, zs = self._offsets3d
vxs, vys, vzs, vis = proj3d.proj_transform_clip(xs, ys, zs, renderer.M)
fcs = (zalpha(self._facecolor3d, vzs) if self._depthshade else
self._facecolor3d)
fcs = mcolors.colorConverter.to_rgba_array(fcs, self._alpha)
self.set_facecolors(fcs)
ecs = (zalpha(self._edgecolor3d, vzs) if self._depthshade else
self._edgecolor3d)
ecs = mcolors.colorConverter.to_rgba_array(ecs, self._alpha)
self.set_edgecolors(ecs)
PathCollection.set_offsets(self, list(zip(vxs, vys)))
if vzs.size > 0 :
return min(vzs)
else :
return np.nan
def _generate_node_artist(pos, styles, *, ax):
N = len(pos)
proto_node = next(iter(pos))
x = np.zeros(N) * np.nan
y = np.zeros(N) * np.nan
properties = {
k: [None] * N
for k in styles[proto_node] if k in _VALID_NODE_STYLE
}
for j, node in enumerate(pos):
x[j], y[j] = pos[node]
for key, values in properties.items():
values[j] = styles[node][key]
key_map = {
'size': 'sizes',
'color': 'facecolors',
'shape': 'marker',
'width': 'linewidths',
'edgecolor': 'edgecolors'
}
renamed_properties = {key_map[k]: v for k, v in properties.items()}
markers = renamed_properties.pop('marker', None)
if markers is None:
paths = (MarkerStyle('o'), )
else:
paths = [MarkerStyle(m) for m in markers]
paths = [p.get_path().transformed(p.get_transform()) for p in paths]
offsets = np.column_stack([x, y])
node_art = PathCollection(paths,
offsets=offsets,
transOffset=ax.transData,
**renamed_properties)
node_art.set_transform(mtransforms.IdentityTransform())
ax.add_collection(node_art)
ax.autoscale_view()
return node_art
def get_patches(self, ax):
ranges = LineCollection(self._cap_ranges, linestyle="solid")
links = LineCollection(self._oob_links,
linestyle="dotted",
colors=colorConverter.to_rgba_array("#808080"))
color = colorConverter.to_rgba_array("#DC143C")
scales = np.array((20, ))
marker_obj = MarkerStyle("o")
path = marker_obj.get_path().transformed(marker_obj.get_transform())
offsets = PathCollection((path, ),
scales,
facecolors=color,
offsets=self._oob_offsets,
transOffset=ax.transData)
offsets.set_transform(IdentityTransform())
return [ranges, links, offsets]
def plot(self,
*,
ax,
color,
lw=2,
n=30000,
fuzz=3,
rot=None,
rot_point=None,
**kwargs):
"""Add this plot to the axis.
Parameters
----------
ax: Axes
Matplotlib axis to add the art to
color: str
Color to use. Must be readable by `matplotlib.colors.to_rgba`. Can be
something like "black", "k", or "#ffffff".
lw: float > 0
Linewidth to use for chords
n: int > 0
Number of chords to use
fuzz: float > 0
How ragged edges should be
rot: float
How much to rotate the shape, in radians
rot_point: ndarray (optional)
What point to rotate the shape around (default is the center)
kwargs:
Passed to `PathCollection`
Returns
-------
None
"""
color = np.array(to_rgba(color)) * 255
colors = np.tile(color.reshape(-1, 1), n).T
colors[:, :3] += np.random.randint(-50, 50, (n, 3))
colors = np.minimum(255, np.maximum(0, colors)) / 255
alpha = kwargs.pop("alpha", 0.05)
colors[:, -1] = alpha
ax.add_collection(
PathCollection(
[
Path(p) for p in self.gen_points(
n=n, fuzz=fuzz, rot=rot, rot_point=rot_point)
],
linewidths=lw,
edgecolors=colors,
facecolors="none",
**kwargs,
))
def __init__(self, *args, **kwargs):
"""
Create a collection of flat 3D paths with its normal vector
pointed in *zdir* direction, and located at *zs* on the *zdir*
axis. 'zs' can be a scalar or an array-like of the same length as
the number of paths in the collection.
Constructor arguments are the same as for
:class:`~matplotlib.collections.PathCollection`. In addition,
keywords *zs=0* and *zdir='z'* are available.
Also, the keyword argument "depthshade" is available to
indicate whether or not to shade the patches in order to
give the appearance of depth (default is *True*).
This is typically desired in scatter plots.
"""
zs = kwargs.pop('zs', 0)
zdir = kwargs.pop('zdir', 'z')
self._depthshade = kwargs.pop('depthshade', True)
PathCollection.__init__(self, *args, **kwargs)
self.set_3d_properties(zs, zdir)
def _add_oop_angle_change_collection(self,
index,
threshold=0,
CURVE_TYPE=4,
zorder=50,
**kwargs):
"""Compute and draw torsion changes on the axes."""
CURVE_TYPE_3, CURVE_TYPE_4 = 3, 4
col = []
edgecolors = []
if CURVE_TYPE is CURVE_TYPE_3:
codes = [
Path.MOVETO,
Path.CURVE3,
Path.CURVE3,
]
elif CURVE_TYPE is CURVE_TYPE_4:
codes = [
Path.MOVETO,
Path.CURVE4,
Path.CURVE4,
Path.CURVE4,
]
for torsion in ob.OBMolTorsionIter(self.molecule):
atoms = [self.molecule.GetAtom(idx + 1) for idx in torsion]
atomsnc = [
self._to_normal_coordinates(atom, index) for atom in atoms
]
teq = self.molecule.GetTorsion(*atoms)
tnc = self.molecule.GetTorsion(*atomsnc)
amplitude = (tnc - teq + 360.0) % 360.0
if amplitude > 180.0:
amplitude -= 360.0
if abs(amplitude) <= threshold: continue
intensity = abs(amplitude / 40)
a, b, c, d = [self._2Dcoords(atom) for atom in atoms]
p2 = 0.5 * (b + c) # middle
p1 = intensity * (a - p2) + p2
p3 = intensity * (d - p2) + p2
color = self.oop_colors[0 if amplitude < 0.0 else 1]
if CURVE_TYPE is CURVE_TYPE_3:
verts = [p1, p2, p3]
elif CURVE_TYPE is CURVE_TYPE_4:
verts = [p1, b, c, p3]
curve = Path(verts, codes)
col.append(curve)
edgecolors.append(color)
kw = {'edgecolors': edgecolors, 'facecolors': 'none'}
kwargs.update(kw)
self._vib_oop = PathCollection(col, zorder=zorder, **kwargs)
self.axes.add_collection(self._vib_oop)
def _add_spectrum_collection(self, **kwargs):
"""Draw spectrum on the axes."""
codes = [
Path.MOVETO,
Path.LINETO,
]
col = []
for frequency, intensity in zip(self.frequencies, self.intensities):
verts = [(frequency, 0.0), (frequency, intensity)]
col.append(Path(verts, codes))
kw = {}
kwargs.update(kw)
self._spectrum = PathCollection(col, **kwargs)
self.axes.add_collection(self._spectrum)
def plot(self,
ax=None,
unfilled=(4, 5),
edgecolors='k',
facecolors=(0.75, 0.75, 0.75)):
from matplotlib.path import Path
from matplotlib.collections import PathCollection
ax = get_axes(ax, xlab="R (m)", ylab="z (m)")
ax.set_aspect('equal')
ax.set_xlim((0.5, 2.5))
ax.set_ylim((-1.5, 1.5))
paths = [Path(xy) for i, xy in enumerate(self.xy) if i not in unfilled]
pc = PathCollection(paths,
facecolors=facecolors,
edgecolors=edgecolors)
ax.add_collection(pc)
paths = [Path(self.xy[i]) for i in unfilled]
pc = PathCollection(paths, facecolors='none', edgecolors=edgecolors)
ax.add_collection(pc)
return ax
def my_scatter(self, ax, x, y):
from matplotlib.collections import PathCollection
from matplotlib.path import Path
import matplotlib.transforms as mtransforms
phi = np.linspace(0, 2 * np.pi, 100)
# Scale, in pixel coordinates
rad = 2
x_circle = np.cos(phi) * rad
y_circle = np.sin(phi) * rad
verts = np.vstack([x_circle, y_circle]).T
path = Path(verts, closed=False)
collection = PathCollection(
[path],
facecolor='blue',
edgecolor='black',
transOffset=ax.transData,
)
collection.set_transform(mtransforms.IdentityTransform())
ax.add_collection(collection, autolim=True)
ax.autoscale()
def __init__(self, n=100, marker=triangle):
v, c = marker
v, c = np.array(v), np.array(c)
self._marker = marker
self._base_vertices = np.tile(v.reshape(-1), n).reshape(n, len(v), 2)
self._vertices = np.tile(v.reshape(-1), n).reshape(n, len(v), 2)
self._codes = np.tile(c.reshape(-1), n)
self._scale = np.ones(n)
self._translate = np.zeros((n, 2))
self._rotate = np.zeros(n)
self._path = Path(vertices=self._vertices.reshape(n*len(v), 2),
codes=self._codes)
self._collection = PathCollection(
[self._path], facecolor="white", edgecolor="black")
def _add_bond_collection(self, zorder=10, **kwargs):
"""Draw molecule skeleton on the axes."""
col = []
codes = [Path.MOVETO, Path.LINETO,] # segment
for obbond in ob.OBMolBondIter(self.molecule):
atom1, atom2 = (self.molecule.GetAtom(obbond.GetBeginAtomIdx()),
self.molecule.GetAtom(obbond.GetEndAtomIdx()))
verts = self._2Dcoords(atom1), self._2Dcoords(atom2)
segment = Path(verts, codes)
col.append(segment)
kw = {'edgecolors': 'k'}
kwargs.update(kw)
self._mol_bonds = PathCollection(col, zorder=zorder, **kwargs)
self.axes.add_collection(self._mol_bonds)
def __init__(self, n=100):
v = np.array([(-0.25, -0.25), (+0.0, +0.5), (+0.25, -0.25), (0, 0)])
c = np.array([Path.MOVETO, Path.LINETO, Path.LINETO, Path.CLOSEPOLY])
self._base_vertices = np.tile(v.reshape(-1), n).reshape(n, len(v), 2)
self._vertices = np.tile(v.reshape(-1), n).reshape(n, len(v), 2)
self._codes = np.tile(c.reshape(-1), n)
self._scale = np.ones(n)
self._translate = np.zeros((n, 2))
self._rotate = np.zeros(n)
self._path = Path(vertices=self._vertices.reshape(n*len(v), 2),
codes=self._codes)
self._collection = PathCollection([self._path], linewidth=0.5,
facecolor="k", edgecolor="w")
def update_plot(num: int, args: MyProgramArgs, sa: SimulationArchive,
ed: ExtraData, dots: PathCollection, title: Text):
global mean_mass
total_frames = args.fps * args.duration
if args.log_time:
log_timestep = (log10(ed.meta.current_time) -
log10(50000)) / total_frames
time = 10**((num + log10(50000) / log_timestep) * log_timestep)
else:
timestep = ed.meta.current_time / total_frames
print(num / total_frames)
time = num * timestep
print(f"{num / total_frames:.2f}, {time:.0f}")
sim = sa.getSimulation(t=time)
if time < 1e3:
timestr = f"{time:.0f}"
elif time < 1e6:
timestr = f"{time / 1e3:.2f}K"
else:
timestr = f"{time / 1e6:.2f}M"
title.set_text(f"({len(sim.particles)}) {timestr} Years")
p: Particle
water_fractions = []
for p in sim.particles[1:]:
pd: ParticleData = ed.pd(p)
wf = pd.water_mass_fraction
water_fractions.append(wf)
# a, e, i, M, M_rat = data[line]
# title.set_text(f"({len(a)}) {ages[line]:.2f}K Years")
a = [p.a for p in sim.particles[1:]]
m = np.array([p.m for p in sim.particles[1:]])
m[:2] /= 1e2
if args.y_axis == "e":
bla = np.array([a, [p.e for p in sim.particles[1:]]])
elif args.y_axis == "i":
bla = np.array([a, [degrees(p.inc) for p in sim.particles[1:]]])
elif args.y_axis == "Omega":
bla = np.array([a, [degrees(p.Omega) for p in sim.particles[1:]]])
else:
raise ValueError("invalid y-axis")
dots.set_offsets(bla.T)
with np.errstate(divide='ignore'): # allow 0 water (becomes -inf)
color_val = (np.log10(water_fractions) + 5) / 5
colors = cmap(color_val)
if not mean_mass:
mean_mass = np.mean(m[3:])
dots.set_sizes(size_factor * m / mean_mass)
dots.set_color(colors)
if output_plots:
plt.savefig(f"tmp/{num}_{time / mega}.pdf", transparent=True)
return dots, title
def makeplot(field2plot, lon_ts, lat_ts, vmin, vmax, m, findex, shp_info):
norm = colors.Normalize(vmin=vmin, vmax=vmax)
bounds = np.arange(vmin, vmax + .0001, dvar)
#fig = plt.figure()
#ax = fig.add_subplot(111)
ax = plt.gca()
ax.cla()
paths = []
for line in shp_info[4]._paths:
paths.append(Path(line.vertices, codes=line.codes))
coll = PathCollection(paths, linewidths=0, facecolors='grey', zorder=2)
ax.add_collection(coll)
m.drawcoastlines(ax=ax, linewidth=0.05)
m.drawcountries(ax=ax, linewidth=0.05)
# m.fillcontinents(color = 'gray')
# m.drawmapboundary()
# m.drawparallels(np.arange(round(coordinates[2]), coordinates[3], dlat), linewidth=0.,
# labels=[1, 0, 0, 0], fontname='Times New Roman', fontsize=16, zorder=1)
# m.drawmeridians(np.arange(round(coordinates[0]), coordinates[1], dlon), linewidth=0.,
# labels=[0, 0, 0, 1], fontname='Times New Roman', fontsize=16, zorder=1)
# u2plot = u[tt, ::NN, ::NN].squeeze()
# v2plot = v[tt, ::NN, ::NN].squeeze()
# speed = np.sqrt(u2plot*u2plot + v2plot*v2plot)
# speed = np.ma.masked_greater(speed, 1.5)
pcm = plt.pcolormesh(lon_ts, lat_ts, field2plot, norm=norm, cmap=cmap)
# cbar = plt.colorbar(pcm, norm=norm, cmap=cmap, orientation='vertical', pad=0.05, aspect=15,
# shrink=0.7, extend='both')
# cbar.set_ticks(bounds)
figname = str(findex).zfill(5)
plt.savefig(figdir + figname, dpi=200)
plt.close()
print findex
return findex
'flag', 'prism', 'ocean', 'gist_earth', 'terrain', 'gist_stern',
'gnuplot', 'gnuplot2', 'CMRmap', 'cubehelix', 'brg',
'gist_rainbow', 'rainbow', 'jet', 'nipy_spectral', 'gist_ncar'])]
colormaplist = sum([x[1] for x in colormap_from_reference],[])
#
# Artist may not reture the same value as given by setter...
# Also, matplotlib has a routine to get a list of valid property values.
# With get_valid_values, it may be possible to implement more automatic
# property list generation.
#
# The following is an initial attempt. Needs to check more how well
# matplotlib is organized in this aspect.
obj = PathCollection(Path.unit_rectangle())
#plinestylelist, plinestyle_rlist = artist_property_checker(obj, 'linestyle')
pedgecolorlist, pedgecolor_rlist = artist_property_checker(
obj, 'edgecolor', values="'"+"','".join(collist)+"'")
plinestylelist = ['solid', 'dotted', 'dashed', 'dashdot']
plinestyle_rlist = ['solid', 'dotted', 'dashed', 'dashdot']
if isMPL2:
plinestylelist = plinestylelist[:4]
plinestyle_rlist = plinestyle_rlist[:4]
def colormap_list():
return colormaplist
def scratch_file():
class MoleculePlotter(Plotter):
"""Draw molecule and vibration.
Arguments:
axes (matplotlib.Axes): axes to draw on.
Attributes:
oop_curve_type: Use either 3- or 4-points bezier to represent bond
torsion.
bond_colors (tuple): Two matplotlib colors.
arc_colors (tuple): Two matplotlib colors.
oop_colors (tuple): Two matplotlib colors.
molecule (ob.OBMol): The molecule to draw.
"""
def __init__(self, axes):
super(MoleculePlotter, self).__init__(axes)
self.clear()
self.oop_curve_type = 4
self.bond_colors = self.arc_colors = ("b", "r")
self.oop_colors = ("g", "y")
self._molecule2D = ob.OBMol()
self._mol_bonds = None
self._mol_atoms = None
self._mol_labels = []
self._vib_bonds = None
self._vib_angles = None
self._vib_oop = None
def _2Dcoords(self, atom):
"""Returns:
The 2D coordinates for `atom`.
"""
atom2D = self._molecule2D.GetAtom(atom.GetIdx())
assert(atom2D.GetZ() == 0.0)
return np.array([atom2D.GetX(), atom2D.GetY()])
def _to_normal_coordinates(self, atom, index):
"""Returns:
The `atom`'s normal coordinates for normal mode at `index`.
"""
def ar(vec):
"""Returns a numpy array with the coordinates of the vector."""
return np.array((vec.GetX(), vec.GetY(), vec.GetZ()))
atomnc = ob.OBAtom()
nc = ar(atom) + self.lx[index][atom.GetIdx() - 1]
atomnc.SetVector(*nc)
return atomnc
@staticmethod
def _sdg(molecule):
"""Structure diagram generation."""
molecule2D = ob.OBMol(molecule)
gen2D = ob.OBOp.FindType("gen2D")
if not gen2D:
raise NameError("name 'gen2D' is not defined")
gen2D.Do(molecule2D)
assert(not molecule2D.Has3D())
return molecule2D
def _add_atom_labels(self, zorder=100, **kwargs):
"""Draw atom labels on the axes."""
box_props = dict(boxstyle='round', facecolor='white', edgecolor='none')
etab = ob.OBElementTable()
for atom in ob.OBMolAtomIter(self.molecule):
x, y = self._2Dcoords(atom)
kw = dict(horizontalalignment="center",
verticalalignment="center",
bbox=box_props)
kwargs.update(kw)
label = AtomText(x, y,
etab.GetSymbol(atom.GetAtomicNum()), atom.GetIdx(),
etab.GetRGB(atom.GetAtomicNum()),
zorder=zorder,
**kwargs)
self._mol_labels.append(label)
self.axes.add_artist(label)
def _add_atom_collection(self, zorder=100, **kwargs):
"""Draw atoms as colored circles on the axes."""
col = []
colors = []
etab = ob.OBElementTable()
for atom in ob.OBMolAtomIter(self.molecule):
colors.append(etab.GetRGB(atom.GetAtomicNum()))
radius = etab.GetCovalentRad(atom.GetAtomicNum())
circle = Circle(self._2Dcoords(atom), radius)
col.append(circle)
kw = {'facecolors': colors, 'edgecolors': colors}
kwargs.update(kw)
self._mol_atoms = PatchCollection(col, zorder=zorder, **kwargs)
self.axes.add_collection(self._mol_atoms)
def _add_bond_collection(self, zorder=10, **kwargs):
"""Draw molecule skeleton on the axes."""
col = []
codes = [Path.MOVETO, Path.LINETO,] # segment
for obbond in ob.OBMolBondIter(self.molecule):
atom1, atom2 = (self.molecule.GetAtom(obbond.GetBeginAtomIdx()),
self.molecule.GetAtom(obbond.GetEndAtomIdx()))
verts = self._2Dcoords(atom1), self._2Dcoords(atom2)
segment = Path(verts, codes)
col.append(segment)
kw = {'edgecolors': 'k'}
kwargs.update(kw)
self._mol_bonds = PathCollection(col, zorder=zorder, **kwargs)
self.axes.add_collection(self._mol_bonds)
def _add_bondlength_change_collection(
self, index, threshold=0, zorder=20, **kwargs):
"""Comute and draw bondlength changes on the axes."""
codes = [Path.MOVETO, Path.LINETO]
col = []
amp = []
colors = []
for obbond in ob.OBMolBondIter(self.molecule):
atom1, atom2 = (self.molecule.GetAtom(obbond.GetBeginAtomIdx()),
self.molecule.GetAtom(obbond.GetEndAtomIdx()))
atom1nc, atom2nc = [self._to_normal_coordinates(atom, index)
for atom in (atom1, atom2)]
if obbond.GetLength() == 0.0:
logger.error(
"Bond between %i and %i with length %.1f ignored."
% (atom1.GetIdx(), atom2.GetIdx(), obbond.GetLength()))
continue
amplitude = atom1.GetDistance(atom2) - atom1nc.GetDistance(atom2nc)
if abs(amplitude * 100) <= threshold: continue
amp.append(abs(amplitude * 50))
colors.append(self.bond_colors[0 if amplitude < 0.0 else 1])
verts = (self._2Dcoords(atom1), self._2Dcoords(atom2))
segment = Path(verts, codes)
col.append(segment)
lw = 0.0 if not col else np.array(amp) * kwargs.pop("lw", 1.0)
kw = {'edgecolors': colors, 'linewidths': lw}
kwargs.update(kw)
self._vib_bonds = PathCollection(col, zorder=zorder, **kwargs)
self.axes.add_collection(self._vib_bonds)
def _add_angle_change_collection(
self, index, threshold=0, zorder=25, **kwargs):
"""Compute and draw angle changes on the axes."""
col = []
colors = []
for angle in ob.OBMolAngleIter(self.molecule):
vertex, atom1, atom2 = [self.molecule.GetAtom(idx + 1)
for idx in angle]
vertexnc, atom1nc, atom2nc = [
self._to_normal_coordinates(atom, index)
for atom in (vertex, atom1, atom2)]
amplitude = (atom1nc.GetAngle(vertexnc, atom2nc) -
atom1.GetAngle(vertex, atom2))
if abs(amplitude) <= threshold: continue
width = height = abs(amplitude) / 20
d1, d2 = (self._2Dcoords(atom1) - self._2Dcoords(vertex),
self._2Dcoords(atom2) - self._2Dcoords(vertex))
theta1 = va.dangle2d(np.array([1.0, 0.0]), d1)
theta2 = va.dangle2d(np.array([1.0, 0.0]), d2)
# always plot smaller arc [ 0.0, 180.0 [
if (theta2 - theta1 + 360.0) % 360.0 > 180.0:
theta2, theta1 = theta1, theta2
color = self.arc_colors[0 if amplitude < 0.0 else 1]
colors.append(color)
arc = Arc(self._2Dcoords(vertex),
width, height, 0.0, theta1, theta2)
col.append(arc)
kw = {'edgecolors': colors, 'facecolors': 'none'}
kwargs.update(kw)
self._vib_angles = PatchCollection(col, zorder=zorder, **kwargs)
self.axes.add_collection(self._vib_angles)
def _add_oop_angle_change_collection(
self, index, threshold=0, CURVE_TYPE=4, zorder=50,
**kwargs):
"""Compute and draw torsion changes on the axes."""
CURVE_TYPE_3, CURVE_TYPE_4 = 3, 4
col = []
edgecolors = []
if CURVE_TYPE is CURVE_TYPE_3:
codes = [Path.MOVETO, Path.CURVE3, Path.CURVE3, ]
elif CURVE_TYPE is CURVE_TYPE_4:
codes = [Path.MOVETO, Path.CURVE4, Path.CURVE4, Path.CURVE4, ]
for torsion in ob.OBMolTorsionIter(self.molecule):
atoms = [self.molecule.GetAtom(idx + 1)
for idx in torsion]
atomsnc = [self._to_normal_coordinates(atom, index)
for atom in atoms]
teq = self.molecule.GetTorsion(*atoms)
tnc = self.molecule.GetTorsion(*atomsnc)
amplitude = (tnc - teq + 360.0) % 360.0
if amplitude > 180.0:
amplitude -= 360.0
if abs(amplitude) <= threshold: continue
intensity = abs(amplitude / 40)
a, b, c, d = [self._2Dcoords(atom) for atom in atoms]
p2 = 0.5 * (b + c) # middle
p1 = intensity * (a - p2) + p2
p3 = intensity * (d - p2) + p2
color = self.oop_colors[0 if amplitude < 0.0 else 1]
if CURVE_TYPE is CURVE_TYPE_3:
verts = [p1, p2, p3]
elif CURVE_TYPE is CURVE_TYPE_4:
verts = [p1, b, c, p3]
curve = Path(verts, codes)
col.append(curve)
edgecolors.append(color)
kw = {'edgecolors': edgecolors, 'facecolors': 'none'}
kwargs.update(kw)
self._vib_oop = PathCollection(col, zorder=zorder, **kwargs)
self.axes.add_collection(self._vib_oop)
def clear(self):
"""Clear the axes."""
super(MoleculePlotter, self).clear()
self.axes.set_xticks(())
self.axes.set_yticks(())
for __, spine in self.axes.spines.items():
spine.set_visible(False)
self.axes.figure.tight_layout()
self.axes.figure.subplots_adjust(left=0, bottom=0, right=1, top=1)
self._mol_labels = []
self._mol_atoms = self._mol_bonds = None
self._vib_bonds = self._vib_angles = self._vib_oop = None
def set_molecule(self, molecule):
"""Set molecule data for this plot."""
super(MoleculePlotter, self).set_molecule(molecule)
self._molecule2D = self._sdg(self.molecule)
def draw_molecule(self, padding=0.3, lw=1.0, fontsize=12.0):
"""Draw molecule on the axes."""
for artist in chain((self._mol_atoms, self._mol_bonds),
self._mol_labels):
if artist:
artist.remove()
self._add_bond_collection(lw=lw)
self._add_atom_labels(fontsize=fontsize)
self.axes.ignore_existing_data_limits = True
xmin, xmax, ymin, ymax = self.axes.axis("image")
self.axes.axis((xmin - padding, xmax + padding,
ymin - padding, ymax + padding))
self.draw()
def draw_vibration(self, row, bl_filter, angle_filter, torsion_filter):
"""Draw vibration on the axes."""
for artist in (self._vib_bonds, self._vib_angles, self._vib_oop):
if artist:
artist.remove()
if row is -1: return
self._add_bondlength_change_collection(row, bl_filter)
self._add_angle_change_collection(row, angle_filter)
self._add_oop_angle_change_collection(row, torsion_filter)
self.draw()
def show_atom_index(self, show=True):
"""Show or hide atom indexes."""
for artist in self._mol_labels:
artist.show_index(show)
self.draw()
def set_black_labels(self, black=True):
"""Colored or black atom labels"""
for artist in self._mol_labels:
artist.set_black_labels(black)
self.draw()
def set_fontsize(self, fontsize):
"""Set the font size for the atom labels."""
for artist in self._mol_labels:
artist.set_fontsize(fontsize)
self.draw()
def set_linewidth(self, linewidth):
"""Set the linewidth used for the skeleton."""
self._mol_bonds.set_linewidth(float(linewidth))
self.draw()
which = values > 1.
for shp_link in [shp_link]:
fig = plt.figure()
patchco = map_poly_shp(shp_link)
patchcoB = map_poly_shp(shp_link, which=which)
patchco.set_facecolor('none')
ax = setup_ax([patchco, patchcoB])
fig.add_axes(ax)
plt.show()
break
'''
xy = (((0, 0), (0, 0)), ((2, 1), (2, 1)), ((3, 1), (3, 1)), ((2, 5), (2,
5)))
xy = np.array([[10, 30], [20, 20]])
markerobj = mpl.markers.MarkerStyle('o')
path = markerobj.get_path().transformed(markerobj.get_transform())
scales = np.array([2, 2])
fig = plt.figure()
ax = fig.add_subplot(111)
pc = PathCollection((path,), scales, offsets=xy, \
facecolors='r', transOffset=mpl.transforms.IdentityTransform())
#pc.set_transform(mpl.transforms.IdentityTransform())
#_ = _add_axes2col(pc, [0, 0, 5, 5])
ax.add_collection(pc)
fig.add_axes(ax)
#ax = setup_ax([pc], ax)
plt.show()