def draw_inner_divertor(self, mesh, iy_region):
inner_divertor_grid_IxIy = numpy.array(
(*mesh.GridPoint_IxIy[:mesh.ixcut[0]],
(*mesh.GridPoint_IxIy[mesh.ixcut[1]][:mesh.iycut],
*mesh.GridPoint_IxIy[mesh.ixcut[0]][mesh.iycut:]))
) # has to concatinate two sections to make the last radial surface.
if iy_region is "both":
coord = inner_divertor_grid_IxIy
self.stainedglass["inner_divertor"] = QuadMesh(
mesh.ny + 1, mesh.ixcut[0], coord)
elif iy_region is "SOL":
coord = inner_divertor_grid_IxIy[:, mesh.iycut:, :]
self.stainedglass["inner_divertor"] = QuadMesh(
mesh.ny + 1 - mesh.iycut, mesh.ixcut[0], coord)
else: # "PFR"
coord = inner_divertor_grid_IxIy[:, :mesh.iycut + 1, :]
self.stainedglass["inner_divertor"] = QuadMesh(
mesh.iycut, mesh.ixcut[0], coord)
self.axes.add_collection(self.stainedglass["inner_divertor"])
def draw_core(self, mesh, iy_region):
core_grid_IxIy = numpy.array(
(*mesh.GridPoint_IxIy[mesh.ixcut[0]:mesh.ixcut[1]],
(*mesh.GridPoint_IxIy[mesh.ixcut[0]][:mesh.iycut],
*mesh.GridPoint_IxIy[mesh.ixcut[1]][mesh.iycut:])))
if iy_region is "both":
coord = core_grid_IxIy
self.stainedglass["core"] = QuadMesh(mesh.ny + 1,
mesh.ixcut[1] - mesh.ixcut[0],
coord)
elif iy_region is "SOL":
coord = core_grid_IxIy[:, mesh.iycut:, :]
self.stainedglass["core"] = QuadMesh(mesh.ny + 1 - mesh.iycut,
mesh.ixcut[1] - mesh.ixcut[0],
coord)
else: # "PFR"
coord = core_grid_IxIy[:, :mesh.iycut + 1, :]
self.stainedglass["core"] = QuadMesh(mesh.iycut,
mesh.ixcut[1] - mesh.ixcut[0],
coord)
self.axes.add_collection(self.stainedglass["core"])
def draw_outer_divertor(self, mesh, iy_region):
# be careful, here quadmesh is drawn using IxIy order!
outer_divertor_grid_IxIy = numpy.array(
mesh.GridPoint_IxIy[mesh.ixcut[1]:])
# get the mesh collection
if iy_region is "both":
coord = outer_divertor_grid_IxIy
self.stainedglass["outer_divertor"] = QuadMesh(
mesh.ny + 1, mesh.nx + 1 - mesh.ixcut[1], coord)
elif iy_region is "SOL":
coord = outer_divertor_grid_IxIy[:, mesh.
iycut:, :] # you can not do multi-dimensional slicing in Python, have to slice on numpy array
self.stainedglass["outer_divertor"] = QuadMesh(
mesh.ny + 1 - mesh.iycut, mesh.nx + 1 - mesh.ixcut[1], coord)
else: # "PFR"
coord = outer_divertor_grid_IxIy[:, :mesh.iycut + 1, :]
self.stainedglass["outer_divertor"] = QuadMesh(
mesh.iycut, mesh.nx + 1 - mesh.ixcut[1], coord)
# add the collection to the axes
self.axes.add_collection(self.stainedglass["outer_divertor"])
def plot( self, title = None, scalelabel = None, xlim=None, ylim=None, xlog=False, ylog=False ):
"""
:param title:
:param scalelabel:
:param xlim:
:param ylim:
:param xlog:
:param ylog:
:return:
"""
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.collections import QuadMesh
# determine where to get the settings for the potentially mirrored second axis
if self.matrix.axes[1].style == 'link': axis1index = 2
else: axis1index = 1
x = self.matrix.axes[2].values
y = self.matrix.axes[axis1index].values
X, Y = np.meshgrid( x, y )
XY = np.hstack((X.ravel()[:,np.newaxis], Y.ravel()[:,np.newaxis]))
Z = (self.matrix.array.constructArray()).ravel()
ax = plt.subplot(1,1,1)
if title is None: title = str( self.toXLink() )
plt.suptitle(title)
qc = QuadMesh(
meshWidth=len(x)-1,
meshHeight=len(y)-1,
coordinates=XY,
# showedges=True,
antialiased=True,
shading='flat',
transOffset=ax.transData)
qc.set_array(Z)
ax.add_collection(qc,autolim=True)
if xlim is None: ax.set_xlim( x[0], x[-1] )
else: ax.set_xlim( xlim[0], xlim[1] )
if ylim is None: ax.set_ylim( y[0], y[-1] )
else: ax.set_ylim( ylim[0], ylim[1] )
if xlog: ax.set_xscale( 'log' )
if ylog: ax.set_yscale( 'log' )
xlabel = self.matrix.axes[2].label + ' (' + self.matrix.axes[2].unit +')'
ylabel = self.matrix.axes[axis1index].label + ' (' + self.matrix.axes[axis1index].unit +')'
ax.set_xlabel( xlabel )
ax.set_ylabel( ylabel )
cbar = plt.colorbar(qc)
if scalelabel is not None: cbar.set_label(scalelabel)
else: cbar.set_label(str(self.type)+' covariance ('+str(self.matrix.axes[0].unit)+')')
plt.show()
def test_quadmesh_deprecated_positional(fig_test, fig_ref):
# test that positional parameters are still accepted with the old signature
# and work correctly
# remove when the old QuadMesh.__init__ signature expires (v3.5+2)
from matplotlib.collections import QuadMesh
x = [0, 1, 2, 3.]
y = [1, 2, 3.]
X, Y = np.meshgrid(x, y)
X += 0.2 * Y
coords = np.stack([X, Y], axis=-1)
assert coords.shape == (3, 4, 2)
C = np.linspace(0, 2, 12).reshape(3, 4)
ax = fig_test.add_subplot()
ax.set(xlim=(0, 5), ylim=(0, 4))
qmesh = QuadMesh(coords, antialiased=False, shading='gouraud')
qmesh.set_array(C)
ax.add_collection(qmesh)
ax = fig_ref.add_subplot()
ax.set(xlim=(0, 5), ylim=(0, 4))
with pytest.warns(MatplotlibDeprecationWarning):
qmesh = QuadMesh(4 - 1, 3 - 1,
coords.copy().reshape(-1, 2), False, 'gouraud')
qmesh.set_array(C)
ax.add_collection(qmesh)
def test_quadmesh_deprecated_signature(fig_test, fig_ref, flat_ref, kwargs):
# test that the new and old quadmesh signature produce the same results
# remove when the old QuadMesh.__init__ signature expires (v3.5+2)
from matplotlib.collections import QuadMesh
x = [0, 1, 2, 3.]
y = [1, 2, 3.]
X, Y = np.meshgrid(x, y)
X += 0.2 * Y
coords = np.stack([X, Y], axis=-1)
assert coords.shape == (3, 4, 2)
C = np.linspace(0, 2, 6).reshape(2, 3)
ax = fig_test.add_subplot()
ax.set(xlim=(0, 5), ylim=(0, 4))
if 'transform' in kwargs:
kwargs['transform'] = mtransforms.Affine2D().scale(1.2) + ax.transData
qmesh = QuadMesh(coords, **kwargs)
qmesh.set_array(C)
ax.add_collection(qmesh)
assert qmesh._shading == 'flat'
ax = fig_ref.add_subplot()
ax.set(xlim=(0, 5), ylim=(0, 4))
if 'transform' in kwargs:
kwargs['transform'] = mtransforms.Affine2D().scale(1.2) + ax.transData
with pytest.warns(MatplotlibDeprecationWarning):
qmesh = QuadMesh(4 - 1, 3 - 1,
coords.copy().reshape(-1, 2) if flat_ref else coords,
**kwargs)
qmesh.set_array(C.flatten() if flat_ref else C)
ax.add_collection(qmesh)
assert qmesh._shading == 'flat'
def add_cam_shade(data, cam_min, cam_max, vmin, vmax, ax, t0=32):
l = data.shape[0]
cam_coords = np.stack(
(np.concatenate(
(np.arange(l + 1) + t0, np.arange(l + 1) + t0), axis=0),
np.concatenate(
(np.ones(l + 1) * cam_min, np.ones(l + 1) * cam_max), axis=0)),
axis=1)
cam_shade = QuadMesh(data.shape[0],
1,
cam_coords,
alpha=.25,
array=data,
cmap='coolwarm',
edgecolor=(1.0, 1.0, 1.0, 0.3),
linewidth=0.0016)
cam_shade.set_clim(vmin=vmin, vmax=vmax)
ax.add_collection(cam_shade)
def get_grid_map_plotter(self):
"""
Create a QuadMesh plotting object for this grid
Returns
-------
quadmesh : matplotlib.collections.QuadMesh
"""
from matplotlib.collections import QuadMesh
verts = np.vstack((self.xgrid.flatten(), self.ygrid.flatten())).T
qm = QuadMesh(self.ncol, self.nrow, verts)
return qm
def mesh(self):
if self._mesh is not None:
return self._mesh
px, py, dx, dy = self._corners
x, y, c = self.axes._pcolorargs('pcolormesh', dx, dy,
self.im[:, :, 0],
allmatch=False)
ny, nx = x.shape
coords = np.column_stack((x.ravel(), y.ravel()))
collection = QuadMesh(nx - 1, ny - 1, coords,
shading='flat', antialiased=False,
edgecolors='None',
cmap='gray')
collection.set_array(c.ravel())
collection.set_clip_path(self.axes.patch)
collection.set_transform(self.axes.transData)
self._mesh = collection
return self._mesh
def polar_imshow(ax,
Z,
extents=[0, 1, 0, 2 * np.pi],
vmin=None,
vmax=None,
cmap="viridis"):
Z = np.atleast_3d(Z)
nr, nt, d = Z.shape
rmin, rmax, tmin, tmax = extents
if d == 1:
cmap = plt.get_cmap(cmap)
vmin = vmin or Z.min()
vmax = vmax or Z.max()
norm = colors.Normalize(vmin=vmin, vmax=vmax)
facecolors = cmap(norm(Z))
else:
facecolors = Z.reshape(nr, nt, 3).reshape(-1, 3)
R = np.linspace(rmin, rmax, nr + 1)
T = np.linspace(tmin, tmax, nt + 1)
T, R = np.meshgrid(T, R)
nr, nt = R.shape
R, T = R.ravel(), T.ravel()
coords = np.column_stack((T, R))
collection = QuadMesh(
nt - 1,
nr - 1,
coords,
rasterized=True,
facecolors=facecolors,
edgecolors="None",
linewidth=0,
antialiased=False,
)
ax.add_collection(collection)
return collection
def clear_color_mesh(qmesh: QuadMesh, default_val=0):
qmesh.set_array(qmesh.get_array() * default_val)
def update_color_mesh(self, qmesh: QuadMesh, time_step, node_id,
value_name):
df = self.update_area(time_step, node_id, value_name)
data = np.array(df)
qmesh.set_array(data.ravel())
return qmesh
def plot(self, ax=None, mode="local", **kwargs):
"""
plot a frame. this can be useful for plotting as background for a velocimetry result. Plotting can be done in
three modes:
- "local": a simple planar view plot, with a local coordinate system in meters, with the top-left coordinate
being the 0, 0 point, and ascending coordinates towards the right and bottom.
- "geographical": a geographical plot, requiring the package `cartopy`, the results are plotted on a geographical
axes, so that combinations with tile layers such as OpenStreetMap, or shapefiles can be made.
- "camera": i.e. seen from the camera perspective. This is the most intuitive view for end users.
:param ax: pre-defined axes object. If not set, a new axes will be prepared. In case `mode=="geographical"`, a
cartopy GeoAxes needs to be provided, or will be made in case ax is not set.
:param mode: can be "local", "geographical", or "camera". For "geographical" a frames set from frames.project
should be used that contains "lon" and "lat" coordinates. For "camera", a non-projected frames set should
be used.
:param kwargs: dict, plotting parameters to be passed to matplotlib.pyplot.pcolormesh, for plotting the
background frame.
"""
# prepare axes
if "time" in self._obj.coords:
if self._obj.time.size > 1:
raise AttributeError(
f'Object contains dimension "time" with length {len(self._obj.time)}. Reduce dataset by selecting one time step or taking a median, mean or other statistic.'
)
ax = plot_orc.prepare_axes(ax=ax, mode=mode)
f = ax.figure # handle to figure
if mode == "local":
x = "x"
y = "y"
elif mode == "geographical":
# import some additional packages
import cartopy.crs as ccrs
# add transform for GeoAxes
kwargs["transform"] = ccrs.PlateCarree()
x = "lon"
y = "lat"
else:
# mode is camera
x = "xp"
y = "yp"
assert all(
v in self._obj.coords for v in [x, y]
), f'required coordinates "{x}" and/or "{y}" are not available'
if (len(self._obj.shape) == 3 and self._obj.shape[-1] == 3):
# looking at an rgb image
facecolors = self._obj.values.reshape(
self._obj.shape[0] * self._obj.shape[1], 3) / 255
facecolors = np.hstack([facecolors, np.ones((len(facecolors), 1))])
quad = ax.pcolormesh(self._obj[x],
self._obj[y],
self._obj.mean(dim="rgb"),
shading="nearest",
facecolors=facecolors,
**kwargs)
# remove array values, override .set_array, needed in case GeoAxes is provided, because GeoAxes asserts if array has dims
QuadMesh.set_array(quad, None)
else:
ax.pcolormesh(self._obj[x], self._obj[y], self._obj, **kwargs)
# fix axis limits to min and max of extent of frames
ax.set_xlim([self._obj[x].min(), self._obj[x].max()])
ax.set_ylim([self._obj[y].min(), self._obj[y].max()])
return ax
