import copy
import warnings
import matplotlib.colors
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Polygon
from ..utils import geometry
from . import plotutil
warnings.simplefilter("always", PendingDeprecationWarning)
[docs]class PlotCrossSection:
"""
Class to create a cross sectional plot of a model.
Parameters
----------
ax : matplotlib.pyplot axis
The plot axis. If not provided it, plt.gca() will be used.
model : flopy.modflow object
flopy model object. (Default is None)
modelgrid : flopy.discretization.Grid object
can be a StructuredGrid, VertexGrid, or UnstructuredGrid object
line : dict
Dictionary with either "row", "column", or "line" key. If key
is "row" or "column" key value should be the zero-based row or
column index for cross-section. If key is "line" value should
be an array of (x, y) tuples with vertices of cross-section.
Vertices should be in map coordinates consistent with xul,
yul, and rotation.
extent : tuple of floats
(xmin, xmax, ymin, ymax) will be used to specify axes limits. If None
then these will be calculated based on grid, coordinates, and rotation.
geographic_coords : bool
boolean flag to allow the user to plot cross section lines in
geographic coordinates. If False (default), cross section is plotted
as the distance along the cross section line.
"""
def __init__(
self,
model=None,
modelgrid=None,
ax=None,
line=None,
extent=None,
geographic_coords=False,
):
self.ax = ax
self.geographic_coords = geographic_coords
self.model = model
if modelgrid is not None:
self.mg = modelgrid
elif model is not None:
self.mg = model.modelgrid
else:
raise Exception("Cannot find model grid")
if self.mg.top is None or self.mg.botm is None:
raise AssertionError("modelgrid top and botm must be defined")
if not isinstance(line, dict):
raise AssertionError("A line dictionary must be provided")
line = {k.lower(): v for k, v in line.items()}
if len(line) != 1:
s = (
"only row, column, or line can be specified in line "
"dictionary keys specified: "
)
for k in line.keys():
s += f"{k} "
raise AssertionError(s)
if ax is None:
self.ax = plt.gca()
else:
self.ax = ax
onkey = list(line.keys())[0]
self.__geographic_xpts = None
# un-translate model grid into model coordinates
xcellcenters, ycellcenters = geometry.transform(
self.mg.xcellcenters,
self.mg.ycellcenters,
self.mg.xoffset,
self.mg.yoffset,
self.mg.angrot_radians,
inverse=True,
)
xverts, yverts = self.mg.cross_section_vertices
(
xverts,
yverts,
) = plotutil.UnstructuredPlotUtilities.irregular_shape_patch(
xverts, yverts
)
self.xvertices, self.yvertices = geometry.transform(
xverts,
yverts,
self.mg.xoffset,
self.mg.yoffset,
self.mg.angrot_radians,
inverse=True,
)
if onkey in ("row", "column"):
eps = 1.0e-4
xedge, yedge = self.mg.xyedges
if onkey == "row":
self.direction = "x"
ycenter = ycellcenters.T[0]
pts = [
(xedge[0] - eps, ycenter[int(line[onkey])]),
(xedge[-1] + eps, ycenter[int(line[onkey])]),
]
else:
self.direction = "y"
xcenter = xcellcenters[0, :]
pts = [
(xcenter[int(line[onkey])], yedge[0] + eps),
(xcenter[int(line[onkey])], yedge[-1] - eps),
]
else:
verts = line[onkey]
xp = []
yp = []
for [v1, v2] in verts:
xp.append(v1)
yp.append(v2)
xp, yp = self.mg.get_local_coords(xp, yp)
if np.max(xp) - np.min(xp) > np.max(yp) - np.min(yp):
# this is x-projection and we should buffer x by small amount
idx0 = list(xp).index(np.max(xp))
idx1 = list(xp).index(np.min(xp))
xp[idx0] += 1e-04
xp[idx1] -= 1e-04
self.direction = "x"
else:
# this is y-projection and we should buffer y by small amount
idx0 = list(yp).index(np.max(yp))
idx1 = list(yp).index(np.min(yp))
yp[idx0] += 1e-04
yp[idx1] -= 1e-04
self.direction = "y"
pts = [(xt, yt) for xt, yt in zip(xp, yp)]
self.pts = np.array(pts)
self.xypts = plotutil.UnstructuredPlotUtilities.line_intersect_grid(
self.pts, self.xvertices, self.yvertices
)
if len(self.xypts) < 2:
s = "cross-section cannot be created\n."
s += " less than 2 points intersect the model grid\n"
s += f" {len(self.xypts)} points intersect the grid."
raise Exception(s)
if self.geographic_coords:
# transform back to geographic coordinates
xypts = {}
for nn, pt in self.xypts.items():
xp = [t[0] for t in pt]
yp = [t[1] for t in pt]
xp, yp = geometry.transform(
xp,
yp,
self.mg.xoffset,
self.mg.yoffset,
self.mg.angrot_radians,
)
xypts[nn] = [(xt, yt) for xt, yt in zip(xp, yp)]
self.xypts = xypts
laycbd = []
self.ncb = 0
if self.model is not None:
if self.model.laycbd is not None:
laycbd = list(self.model.laycbd)
self.ncb = np.count_nonzero(laycbd)
if laycbd:
self.active = []
for k in range(self.mg.nlay):
self.active.append(1)
if laycbd[k] > 0:
self.active.append(0)
self.active = np.array(self.active, dtype=int)
else:
self.active = np.ones(self.mg.nlay, dtype=int)
self._nlay, self._ncpl, self.ncb = self.mg.cross_section_lay_ncpl_ncb(
self.ncb
)
top = self.mg.top.reshape(1, self._ncpl)
botm = self.mg.botm.reshape(self._nlay + self.ncb, self._ncpl)
self.elev = np.concatenate((top, botm), axis=0)
self.idomain = self.mg.idomain
if self.mg.idomain is None:
self.idomain = np.ones(botm.shape, dtype=int)
self.projpts = self.set_zpts(None)
# Create cross-section extent
if extent is None:
self.extent = self.get_extent()
else:
self.extent = extent
# this is actually x or y based on projection
self.xcenters = [
np.mean(np.array(v).T[0]) for i, v in sorted(self.projpts.items())
]
self.mean_dx = np.mean(
np.max(self.xvertices, axis=1) - np.min(self.xvertices, axis=1)
)
self.mean_dy = np.mean(
np.max(self.yvertices, axis=1) - np.min(self.yvertices, axis=1)
)
self._polygons = {}
# Set axis limits
self.ax.set_xlim(self.extent[0], self.extent[1])
self.ax.set_ylim(self.extent[2], self.extent[3])
@property
def polygons(self):
"""
Method to return cached matplotlib polygons for a cross
section
Returns
-------
dict : [matplotlib.patches.Polygon]
"""
if not self._polygons:
for cell, poly in self.projpts.items():
if len(poly) > 4:
# this is the rare multipolygon instance...
n = 0
p = []
polys = []
for vn, v in enumerate(poly):
if vn == 3 + 4 * n:
n += 1
p.append(v)
polys.append(p)
p = []
else:
p.append(v)
else:
polys = [poly]
for polygon in polys:
verts = plotutil.UnstructuredPlotUtilities.arctan2(
np.array(polygon)
)
if cell not in self._polygons:
self._polygons[cell] = [Polygon(verts, closed=True)]
else:
self._polygons[cell].append(
Polygon(verts, closed=True)
)
return copy.copy(self._polygons)
[docs] def get_extent(self):
"""
Get the extent of the rotated and offset grid
Returns
-------
tuple : (xmin, xmax, ymin, ymax)
"""
xpts = []
for _, verts in self.projpts.items():
for v in verts:
xpts.append(v[0])
xmin = np.min(xpts)
xmax = np.max(xpts)
ymin = np.min(self.elev)
ymax = np.max(self.elev)
return xmin, xmax, ymin, ymax
[docs] def plot_array(self, a, masked_values=None, head=None, **kwargs):
"""
Plot a three-dimensional array as a patch collection.
Parameters
----------
a : numpy.ndarray
Three-dimensional array to plot.
masked_values : iterable of floats, ints
Values to mask.
head : numpy.ndarray
Three-dimensional array to set top of patches to the minimum
of the top of a layer or the head value. Used to create
patches that conform to water-level elevations.
**kwargs : dictionary
keyword arguments passed to matplotlib.collections.PatchCollection
Returns
-------
patches : matplotlib.collections.PatchCollection
"""
ax = kwargs.pop("ax", self.ax)
if not isinstance(a, np.ndarray):
a = np.array(a)
if a.ndim > 1:
a = np.ravel(a)
if masked_values is not None:
for mval in masked_values:
a = np.ma.masked_values(a, mval)
if isinstance(head, np.ndarray):
projpts = self.set_zpts(np.ravel(head))
else:
projpts = None
pc = self.get_grid_patch_collection(a, projpts, **kwargs)
if pc is not None:
ax.add_collection(pc)
ax.set_xlim(self.extent[0], self.extent[1])
ax.set_ylim(self.extent[2], self.extent[3])
return pc
[docs] def plot_surface(self, a, masked_values=None, **kwargs):
"""
Plot a two- or three-dimensional array as line(s).
Parameters
----------
a : numpy.ndarray
Two- or three-dimensional array to plot.
masked_values : iterable of floats, ints
Values to mask.
**kwargs : dictionary
keyword arguments passed to matplotlib.pyplot.plot
Returns
-------
plot : list containing matplotlib.plot objects
"""
ax = kwargs.pop("ax", self.ax)
color = kwargs.pop("color", "b")
color = kwargs.pop("c", color)
if not isinstance(a, np.ndarray):
a = np.array(a)
if a.ndim > 1:
a = np.ravel(a)
if a.size % self._ncpl != 0:
raise AssertionError("Array size must be a multiple of ncpl")
if masked_values is not None:
for mval in masked_values:
a = np.ma.masked_values(a, mval)
d = {
i: (np.min(np.array(v).T[0]), np.max(np.array(v).T[0]))
for i, v in sorted(self.projpts.items())
}
surface = []
for cell, val in d.items():
if cell >= a.size:
continue
elif np.isnan(a[cell]):
continue
elif a[cell] is np.ma.masked:
continue
else:
line = ax.plot(
d[cell], [a[cell], a[cell]], color=color, **kwargs
)
surface.append(line)
ax.set_xlim(self.extent[0], self.extent[1])
ax.set_ylim(self.extent[2], self.extent[3])
return surface
[docs] def plot_fill_between(
self,
a,
colors=("blue", "red"),
masked_values=None,
head=None,
**kwargs,
):
"""
Plot a three-dimensional array as lines.
Parameters
----------
a : numpy.ndarray
Three-dimensional array to plot.
colors : list
matplotlib fill colors, two required
masked_values : iterable of floats, ints
Values to mask.
head : numpy.ndarray
Three-dimensional array to set top of patches to the minimum
of the top of a layer or the head value. Used to create
patches that conform to water-level elevations.
**kwargs : dictionary
keyword arguments passed to matplotlib.pyplot.plot
Returns
-------
plot : list containing matplotlib.fillbetween objects
"""
ax = kwargs.pop("ax", self.ax)
kwargs["colors"] = colors
if not isinstance(a, np.ndarray):
a = np.array(a)
a = np.ravel(a)
if masked_values is not None:
for mval in masked_values:
a = np.ma.masked_values(a, mval)
if isinstance(head, np.ndarray):
projpts = self.set_zpts(head)
else:
projpts = self.projpts
pc = self.get_grid_patch_collection(
a, projpts, fill_between=True, **kwargs
)
if pc is not None:
ax.add_collection(pc)
ax.set_xlim(self.extent[0], self.extent[1])
ax.set_ylim(self.extent[2], self.extent[3])
return pc
[docs] def contour_array(self, a, masked_values=None, head=None, **kwargs):
"""
Contour a two-dimensional array.
Parameters
----------
a : numpy.ndarray
Three-dimensional array to plot.
masked_values : iterable of floats, ints
Values to mask.
head : numpy.ndarray
Three-dimensional array to set top of patches to the minimum
of the top of a layer or the head value. Used to create
patches that conform to water-level elevations.
**kwargs : dictionary
keyword arguments passed to matplotlib.pyplot.contour
Returns
-------
contour_set : matplotlib.pyplot.contour
"""
import matplotlib.tri as tri
if not isinstance(a, np.ndarray):
a = np.array(a)
if a.ndim > 1:
a = np.ravel(a)
ax = kwargs.pop("ax", self.ax)
xcenters = self.xcenters
plotarray = np.array([a[cell] for cell in sorted(self.projpts)])
(
plotarray,
xcenters,
zcenters,
mplcontour,
) = self.mg.cross_section_set_contour_arrays(
plotarray, xcenters, head, self.elev, self.projpts
)
if not mplcontour:
if isinstance(head, np.ndarray):
zcenters = self.set_zcentergrid(np.ravel(head))
else:
zcenters = np.array(
[
np.mean(np.array(v).T[1])
for i, v in sorted(self.projpts.items())
]
)
# work around for tri-contour ignore vmin & vmax
# necessary for the tri-contour NaN issue fix
if "levels" not in kwargs:
vmin = kwargs.pop("vmin", np.nanmin(plotarray))
vmax = kwargs.pop("vmax", np.nanmax(plotarray))
levels = np.linspace(vmin, vmax, 7)
kwargs["levels"] = levels
# workaround for tri-contour nan issue
plotarray[np.isnan(plotarray)] = -(2**31)
if masked_values is None:
masked_values = [-(2**31)]
else:
masked_values = list(masked_values)
if -(2**31) not in masked_values:
masked_values.append(-(2**31))
ismasked = None
if masked_values is not None:
for mval in masked_values:
if ismasked is None:
ismasked = np.isclose(plotarray, mval)
else:
t = np.isclose(plotarray, mval)
ismasked += t
filled = kwargs.pop("filled", False)
plot_triplot = kwargs.pop("plot_triplot", False)
if "extent" in kwargs:
extent = kwargs.pop("extent")
idx = (
(xcenters >= extent[0])
& (xcenters <= extent[1])
& (zcenters >= extent[2])
& (zcenters <= extent[3])
)
plotarray = plotarray[idx].flatten()
xcenters = xcenters[idx].flatten()
zcenters = zcenters[idx].flatten()
if mplcontour:
plotarray = np.ma.masked_array(plotarray, ismasked)
if filled:
contour_set = ax.contourf(
xcenters, zcenters, plotarray, **kwargs
)
else:
contour_set = ax.contour(
xcenters, zcenters, plotarray, **kwargs
)
else:
triang = tri.Triangulation(xcenters, zcenters)
analyze = tri.TriAnalyzer(triang)
mask = analyze.get_flat_tri_mask(rescale=False)
if ismasked is not None:
ismasked = ismasked.flatten()
mask2 = np.any(
np.where(ismasked[triang.triangles], True, False), axis=1
)
mask[mask2] = True
triang.set_mask(mask)
if filled:
contour_set = ax.tricontourf(triang, plotarray, **kwargs)
else:
contour_set = ax.tricontour(triang, plotarray, **kwargs)
if plot_triplot:
ax.triplot(triang, color="black", marker="o", lw=0.75)
ax.set_xlim(self.extent[0], self.extent[1])
ax.set_ylim(self.extent[2], self.extent[3])
return contour_set
[docs] def plot_inactive(self, ibound=None, color_noflow="black", **kwargs):
"""
Make a plot of inactive cells. If not specified, then pull ibound
from the self.ml
Parameters
----------
ibound : numpy.ndarray
ibound array to plot. (Default is ibound in 'BAS6' package.)
color_noflow : string
(Default is 'black')
Returns
-------
quadmesh : matplotlib.collections.QuadMesh
"""
if ibound is None:
if self.mg.idomain is None:
raise AssertionError("An idomain array must be provided")
else:
ibound = self.mg.idomain
plotarray = np.zeros(ibound.shape, dtype=int)
idx1 = ibound == 0
plotarray[idx1] = 1
plotarray = np.ma.masked_equal(plotarray, 0)
cmap = matplotlib.colors.ListedColormap(["0", color_noflow])
bounds = [0, 1, 2]
norm = matplotlib.colors.BoundaryNorm(bounds, cmap.N)
patches = self.plot_array(plotarray, cmap=cmap, norm=norm, **kwargs)
return patches
[docs] def plot_ibound(
self,
ibound=None,
color_noflow="black",
color_ch="blue",
color_vpt="red",
head=None,
**kwargs,
):
"""
Make a plot of ibound. If not specified, then pull ibound from the
self.model
Parameters
----------
ibound : numpy.ndarray
ibound array to plot. (Default is ibound in 'BAS6' package.)
color_noflow : string
(Default is 'black')
color_ch : string
Color for constant heads (Default is 'blue'.)
head : numpy.ndarray
Three-dimensional array to set top of patches to the minimum
of the top of a layer or the head value. Used to create
patches that conform to water-level elevations.
**kwargs : dictionary
keyword arguments passed to matplotlib.collections.PatchCollection
Returns
-------
patches : matplotlib.collections.PatchCollection
"""
if ibound is None:
if self.model is not None:
if self.model.version == "mf6":
color_ch = color_vpt
if self.mg.idomain is None:
raise AssertionError("Ibound/Idomain array must be provided")
ibound = self.mg.idomain
plotarray = np.zeros(ibound.shape, dtype=int)
idx1 = ibound == 0
idx2 = ibound < 0
plotarray[idx1] = 1
plotarray[idx2] = 2
plotarray = np.ma.masked_equal(plotarray, 0)
cmap = matplotlib.colors.ListedColormap(
["none", color_noflow, color_ch]
)
bounds = [0, 1, 2, 3]
norm = matplotlib.colors.BoundaryNorm(bounds, cmap.N)
# mask active cells
patches = self.plot_array(
plotarray,
masked_values=[0],
head=head,
cmap=cmap,
norm=norm,
**kwargs,
)
return patches
[docs] def plot_grid(self, **kwargs):
"""
Plot the grid lines.
Parameters
----------
kwargs : ax, colors. The remaining kwargs are passed into the
the LineCollection constructor.
Returns
-------
lc : matplotlib.collections.LineCollection
"""
ax = kwargs.pop("ax", self.ax)
col = self.get_grid_line_collection(**kwargs)
if col is not None:
ax.add_collection(col)
# ax.set_xlim(self.extent[0], self.extent[1])
# ax.set_ylim(self.extent[2], self.extent[3])
return col
[docs] def plot_bc(
self, name=None, package=None, kper=0, color=None, head=None, **kwargs
):
"""
Plot boundary conditions locations for a specific boundary
type from a flopy model
Parameters
----------
name : string
Package name string ('WEL', 'GHB', etc.). (Default is None)
package : flopy.modflow.Modflow package class instance
flopy package class instance. (Default is None)
kper : int
Stress period to plot
color : string
matplotlib color string. (Default is None)
head : numpy.ndarray
Three-dimensional array (structured grid) or
Two-dimensional array (vertex grid)
to set top of patches to the minimum of the top of a\
layer or the head value. Used to create
patches that conform to water-level elevations.
**kwargs : dictionary
keyword arguments passed to matplotlib.collections.PatchCollection
Returns
-------
patches : matplotlib.collections.PatchCollection
"""
if "ftype" in kwargs and name is None:
name = kwargs.pop("ftype")
# Find package to plot
if package is not None:
p = package
elif self.model is not None:
if name is None:
raise Exception("ftype not specified")
name = name.upper()
p = self.model.get_package(name)
else:
raise Exception("Cannot find package to plot")
# trap for mf6 'cellid' vs mf2005 'k', 'i', 'j' convention
if isinstance(p, list) or p.parent.version == "mf6":
if not isinstance(p, list):
p = [p]
idx = np.array([])
for pp in p:
if pp.package_type in ("lak", "sfr", "maw", "uzf"):
t = plotutil.advanced_package_bc_helper(pp, self.mg, kper)
else:
try:
mflist = pp.stress_period_data.array[kper]
except Exception as e:
raise Exception(
f"Not a list-style boundary package: {e!s}"
)
if mflist is None:
return
t = np.array(
[list(i) for i in mflist["cellid"]], dtype=int
).T
if len(idx) == 0:
idx = np.copy(t)
else:
idx = np.append(idx, t, axis=1)
else:
# modflow-2005 structured and unstructured grid
if p.package_type in ("uzf", "lak"):
idx = plotutil.advanced_package_bc_helper(p, self.mg, kper)
else:
try:
mflist = p.stress_period_data[kper]
except Exception as e:
raise Exception(
f"Not a list-style boundary package: {e!s}"
)
if mflist is None:
return
if len(self.mg.shape) == 3:
idx = [mflist["k"], mflist["i"], mflist["j"]]
else:
idx = mflist["node"]
if len(self.mg.shape) != 3:
plotarray = np.zeros((self._nlay, self._ncpl), dtype=int)
plotarray[tuple(idx)] = 1
else:
plotarray = np.zeros(
(self.mg.nlay, self.mg.nrow, self.mg.ncol), dtype=int
)
plotarray[idx[0], idx[1], idx[2]] = 1
plotarray = np.ma.masked_equal(plotarray, 0)
if color is None:
key = name[:3].upper()
if key in plotutil.bc_color_dict:
c = plotutil.bc_color_dict[key]
else:
c = plotutil.bc_color_dict["default"]
else:
c = color
cmap = matplotlib.colors.ListedColormap(["none", c])
bounds = [0, 1, 2]
norm = matplotlib.colors.BoundaryNorm(bounds, cmap.N)
patches = self.plot_array(
plotarray,
masked_values=[0],
head=head,
cmap=cmap,
norm=norm,
**kwargs,
)
return patches
[docs] def plot_vector(
self,
vx,
vy,
vz,
head=None,
kstep=1,
hstep=1,
normalize=False,
masked_values=None,
**kwargs,
):
"""
Plot a vector.
Parameters
----------
vx : np.ndarray
x component of the vector to be plotted (non-rotated)
array shape must be (nlay, nrow, ncol) for a structured grid
array shape must be (nlay, ncpl) for a unstructured grid
vy : np.ndarray
y component of the vector to be plotted (non-rotated)
array shape must be (nlay, nrow, ncol) for a structured grid
array shape must be (nlay, ncpl) for a unstructured grid
vz : np.ndarray
y component of the vector to be plotted (non-rotated)
array shape must be (nlay, nrow, ncol) for a structured grid
array shape must be (nlay, ncpl) for a unstructured grid
head : numpy.ndarray
MODFLOW's head array. If not provided, then the quivers will be
plotted in the cell center.
kstep : int
layer frequency to plot (default is 1)
hstep : int
horizontal frequency to plot (default is 1)
normalize : bool
boolean flag used to determine if vectors should be normalized
using the vector magnitude in each cell (default is False)
masked_values : iterable of floats
values to mask
kwargs : matplotlib.pyplot keyword arguments for the
plt.quiver method
Returns
-------
quiver : matplotlib.pyplot.quiver
result of the quiver function
"""
ax = kwargs.pop("ax", self.ax)
pivot = kwargs.pop("pivot", "middle")
# Check that the cross section is not arbitrary with a tolerance
# of the mean cell size in each direction
arbitrary = False
pts = self.pts
xuniform = [
True if abs(pts.T[0, 0] - i) < self.mean_dy else False
for i in pts.T[0]
]
yuniform = [
True if abs(pts.T[1, 0] - i) < self.mean_dx else False
for i in pts.T[1]
]
if not np.all(xuniform) and not np.all(yuniform):
arbitrary = True
if arbitrary:
err_msg = (
"plot_specific_discharge() does not "
"support arbitrary cross-sections"
)
raise AssertionError(err_msg)
# get ibound array to mask inactive cells
ib = np.ones((self.mg.nnodes,), dtype=int)
if self.mg.idomain is not None:
ib = self.mg.idomain.ravel()
# get the actual values to plot and set xcenters
if self.direction == "x":
u_tmp = vx
else:
u_tmp = vy * -1.0
# kstep implementation for vertex grid
projpts = {
key: value
for key, value in self.projpts.items()
if (key // self._ncpl) % kstep == 0
}
# set x and z centers
if isinstance(head, np.ndarray):
# pipe kstep to set_zcentergrid to assure consistent array size
zcenters = self.set_zcentergrid(np.ravel(head), kstep=kstep)
else:
zcenters = [
np.mean(np.array(v).T[1]) for i, v in sorted(projpts.items())
]
xcenters = np.array(
[np.mean(np.array(v).T[0]) for i, v in sorted(projpts.items())]
)
x = np.ravel(xcenters)
z = np.ravel(zcenters)
u = np.array([u_tmp.ravel()[cell] for cell in sorted(projpts)])
v = np.array([vz.ravel()[cell] for cell in sorted(projpts)])
ib = np.array([ib[cell] for cell in sorted(projpts)])
x = x[::hstep]
z = z[::hstep]
u = u[::hstep]
v = v[::hstep]
ib = ib[::hstep]
# mask values
if masked_values is not None:
for mval in masked_values:
to_mask = np.logical_or(u == mval, v == mval)
u[to_mask] = np.nan
v[to_mask] = np.nan
# normalize
if normalize:
vmag = np.sqrt(u**2.0 + v**2.0)
idx = vmag > 0.0
u[idx] /= vmag[idx]
v[idx] /= vmag[idx]
# mask with an ibound array
u[ib == 0] = np.nan
v[ib == 0] = np.nan
# plot with quiver
quiver = ax.quiver(x, z, u, v, pivot=pivot, **kwargs)
return quiver
[docs] def plot_pathline(
self, pl, travel_time=None, method="cell", head=None, **kwargs
):
"""
Plot the MODPATH pathlines
Parameters
----------
pl : list of rec arrays or a single rec array
rec array or list of rec arrays is data returned from
modpathfile PathlineFile get_data() or get_alldata()
methods. Data in rec array is 'x', 'y', 'z', 'time',
'k', and 'particleid'.
travel_time : float or str
travel_time is a travel time selection for the displayed
pathlines. If a float is passed then pathlines with times
less than or equal to the passed time are plotted. If a
string is passed a variety logical constraints can be added
in front of a time value to select pathlines for a select
period of time. Valid logical constraints are <=, <, >=, and
>. For example, to select all pathlines less than 10000 days
travel_time='< 10000' would be passed to plot_pathline.
(default is None)
method : str
"cell" shows only pathlines that intersect with a cell
"all" projects all pathlines onto the cross section regardless
of whether they intersect with a given cell
head : np.ndarray
optional adjustment to only show pathlines that are <= to
the top of the water table given a user supplied head array
kwargs : layer, ax, colors. The remaining kwargs are passed
into the LineCollection constructor.
Returns
-------
lc : matplotlib.collections.LineCollection
"""
from matplotlib.collections import LineCollection
# make sure pathlines is a list
if not isinstance(pl, list):
pl = [pl]
marker = kwargs.pop("marker", None)
markersize = kwargs.pop("markersize", None)
markersize = kwargs.pop("ms", markersize)
markercolor = kwargs.pop("markercolor", None)
markerevery = kwargs.pop("markerevery", 1)
ax = kwargs.pop("ax", self.ax)
if "colors" not in kwargs:
kwargs["colors"] = "0.5"
projpts = self.projpts
if head is not None:
projpts = self.set_zpts(head)
pl2 = []
for p in pl:
tp = plotutil.filter_modpath_by_travel_time(p, travel_time)
pl2.append(tp)
tp = plotutil.intersect_modpath_with_crosssection(
pl2,
projpts,
self.xvertices,
self.yvertices,
self.direction,
self._ncpl,
method=method,
)
plines = plotutil.reproject_modpath_to_crosssection(
tp,
projpts,
self.xypts,
self.direction,
self.mg,
self._ncpl,
self.geographic_coords,
)
# build linecollection and markers arrays
linecol = []
markers = []
for _, arr in plines.items():
arr = np.array(arr)
# sort by travel time
arr = arr[arr[:, -1].argsort()]
linecol.append(arr[:, :-1])
if marker is not None:
for xy in arr[::markerevery]:
markers.append(xy)
lc = None
if len(linecol) > 0:
lc = LineCollection(linecol, **kwargs)
ax.add_collection(lc)
if marker is not None:
markers = np.array(markers)
ax.plot(
markers[:, 0],
markers[:, 1],
lw=0,
marker=marker,
color=markercolor,
ms=markersize,
)
return lc
[docs] def plot_timeseries(
self, ts, travel_time=None, method="cell", head=None, **kwargs
):
"""
Plot the MODPATH timeseries.
Parameters
----------
ts : list of rec arrays or a single rec array
rec array or list of rec arrays is data returned from
modpathfile TimeseriesFile get_data() or get_alldata()
methods. Data in rec array is 'x', 'y', 'z', 'time',
'k', and 'particleid'.
travel_time : float or str
travel_time is a travel time selection for the displayed
pathlines. If a float is passed then pathlines with times
less than or equal to the passed time are plotted. If a
string is passed a variety logical constraints can be added
in front of a time value to select pathlines for a select
period of time. Valid logical constraints are <=, <, >=, and
>. For example, to select all pathlines less than 10000 days
travel_time='< 10000' would be passed to plot_pathline.
(default is None)
kwargs : layer, ax, colors. The remaining kwargs are passed
into the LineCollection constructor. If layer='all',
pathlines are output for all layers
Returns
-------
lo : list of Line2D objects
"""
if "color" in kwargs:
kwargs["markercolor"] = kwargs["color"]
return self.plot_pathline(
ts, travel_time=travel_time, method=method, head=head, **kwargs
)
[docs] def plot_endpoint(
self,
ep,
direction="ending",
selection=None,
selection_direction=None,
method="cell",
head=None,
**kwargs,
):
"""
Parameters
----------
Returns
-------
"""
ax = kwargs.pop("ax", self.ax)
# colorbar kwargs
createcb = kwargs.pop("colorbar", False)
colorbar_label = kwargs.pop("colorbar_label", "Endpoint Time")
shrink = float(kwargs.pop("shrink", 1.0))
# marker kwargs
s = kwargs.pop("s", np.sqrt(50))
s = float(kwargs.pop("size", s)) ** 2.0
cd = {}
if "c" not in kwargs:
vmin, vmax = 1e10, -1e10
for rec in ep:
tt = float(rec["time"] - rec["time0"])
if tt < vmin:
vmin = tt
if tt > vmax:
vmax = tt
cd[int(rec["particleid"])] = tt
kwargs["vmin"] = vmin
kwargs["vmax"] = vmax
else:
tc = kwargs.pop("c")
for rec in ep:
cd[int(rec["praticleid"])] = tc
tep, istart = plotutil.parse_modpath_selection_options(
ep, direction, selection, selection_direction
)[0:2]
projpts = self.projpts
if head is not None:
projpts = self.set_zpts(head)
tep = plotutil.intersect_modpath_with_crosssection(
tep,
projpts,
self.xvertices,
self.yvertices,
self.direction,
method=method,
starting=istart,
)
if not tep:
return
epdict = plotutil.reproject_modpath_to_crosssection(
tep,
projpts,
self.xypts,
self.direction,
self.mg,
self.geographic_coords,
starting=istart,
)
arr = []
c = []
for node, epl in sorted(epdict.items()):
c.append(cd[node])
for xy in epl:
arr.append(xy)
arr = np.array(arr)
sp = ax.scatter(arr[:, 0], arr[:, 1], c=c, s=s, **kwargs)
# add a colorbar for travel times
if createcb:
cb = plt.colorbar(sp, ax=ax, shrink=shrink)
cb.set_label(colorbar_label)
return sp
[docs] def get_grid_line_collection(self, **kwargs):
"""
Get a PatchCollection of the grid
Parameters
----------
**kwargs : dictionary
keyword arguments passed to matplotlib.collections.LineCollection
Returns
-------
PatchCollection : matplotlib.collections.LineCollection
"""
from matplotlib.collections import PatchCollection
edgecolor = kwargs.pop("colors", "grey")
edgecolor = kwargs.pop("color", edgecolor)
edgecolor = kwargs.pop("ec", edgecolor)
edgecolor = kwargs.pop("edgecolor", edgecolor)
facecolor = kwargs.pop("facecolor", "none")
facecolor = kwargs.pop("fc", facecolor)
polygons = [
p for _, polys in sorted(self.polygons.items()) for p in polys
]
if len(polygons) > 0:
patches = PatchCollection(
polygons, edgecolor=edgecolor, facecolor=facecolor, **kwargs
)
else:
patches = None
return patches
[docs] def set_zpts(self, vs):
"""
Get an array of projected vertices corrected with corrected
elevations based on minimum of cell elevation (self.elev) or
passed vs numpy.ndarray
Parameters
----------
vs : numpy.ndarray
Two-dimensional array to plot.
Returns
-------
zpts : dict
"""
# make vertex array based on projection direction
if vs is not None:
if not isinstance(vs, np.ndarray):
vs = np.array(vs)
if self.direction == "x":
xyix = 0
else:
xyix = -1
projpts = {}
nlay = self.mg.nlay + self.ncb
nodeskip = self.mg.cross_section_nodeskip(nlay, self.xypts)
cbcnt = 0
for k in range(1, nlay + 1):
if not self.active[k - 1]:
cbcnt += 1
continue
k, ns, ncbnn = self.mg.cross_section_adjust_indicies(k - 1, cbcnt)
top = self.elev[k - 1, :]
botm = self.elev[k, :]
d0 = 0
# trap to split multipolygons
xypts = []
for nn, verts in self.xypts.items():
if nn in nodeskip[ns - 1]:
continue
if len(verts) > 2:
i0 = 2
for ix in range(len(verts)):
if ix == i0 - 1:
xypts.append((nn, verts[i0 - 2 : i0]))
i0 += 2
else:
xypts.append((nn, verts))
xypts = sorted(xypts, key=lambda q: q[-1][xyix][xyix])
if self.direction == "y":
xypts = xypts[::-1]
for nn, verts in xypts:
if vs is None:
t = top[nn]
else:
t = vs[nn + ncbnn]
if np.isclose(t, -1e30):
t = botm[nn]
if t < botm[nn]:
t = botm[nn]
if top[nn] < t:
t = top[nn]
b = botm[nn]
if self.geographic_coords:
if self.direction == "x":
projt = [(v[0], t) for v in verts]
projb = [(v[0], b) for v in verts]
else:
projt = [(v[1], t) for v in verts]
projb = [(v[1], b) for v in verts]
else:
verts = np.array(verts).T
a2 = (np.max(verts[0]) - np.min(verts[0])) ** 2
b2 = (np.max(verts[1]) - np.min(verts[1])) ** 2
c = np.sqrt(a2 + b2)
d1 = d0 + c
projt = [(d0, t), (d1, t)]
projb = [(d0, b), (d1, b)]
d0 += c
projpt = projt + projb
node = nn + ncbnn
if node not in projpts:
projpts[node] = projpt
else:
projpts[node] += projpt
return projpts
[docs] def set_zcentergrid(self, vs, kstep=1):
"""
Get an array of z elevations at the center of a cell that is based
on minimum of cell top elevation (self.elev) or passed vs numpy.ndarray
Parameters
----------
vs : numpy.ndarray
Three-dimensional array to plot.
kstep : int
plotting layer interval
Returns
-------
zcentergrid : numpy.ndarray
"""
verts = self.set_zpts(vs)
zcenters = [
np.mean(np.array(v).T[1])
for i, v in sorted(verts.items())
if (i // self._ncpl) % kstep == 0
]
return zcenters
[docs] def get_grid_patch_collection(
self, plotarray, projpts=None, fill_between=False, **kwargs
):
"""
Get a PatchCollection of plotarray in unmasked cells
Parameters
----------
plotarray : numpy.ndarray
One-dimensional array to attach to the Patch Collection.
projpts : dict
dictionary defined by node number which contains model
patch vertices.
fill_between : bool
flag to create polygons that mimick the matplotlib fill between
method. Only used by the plot_fill_between method.
**kwargs : dictionary
keyword arguments passed to matplotlib.collections.PatchCollection
Returns
-------
patches : matplotlib.collections.PatchCollection
"""
from matplotlib.collections import PatchCollection
from matplotlib.patches import Polygon
use_cache = False
if projpts is None:
use_cache = True
projpts = self.polygons
vmin = kwargs.pop("vmin", None)
vmax = kwargs.pop("vmax", None)
match_original = False
if fill_between:
match_original = True
colors = kwargs.pop("colors")
rectcol = []
data = []
for cell, poly in sorted(projpts.items()):
if not use_cache:
if len(poly) > 4:
# multipolygon instance...
n = 0
p = []
polys = []
for vn, v in enumerate(poly):
if vn == 3 + 4 * n:
n += 1
p.append(v)
polys.append(p)
p = []
else:
p.append(v)
else:
polys = [poly]
else:
polys = poly
for polygon in polys:
if not use_cache:
polygon = plotutil.UnstructuredPlotUtilities.arctan2(
np.array(polygon)
)
if np.isnan(plotarray[cell]):
continue
elif plotarray[cell] is np.ma.masked:
continue
if use_cache:
rectcol.append(polygon)
elif fill_between:
x = list(set(np.array(polygon).T[0]))
y1 = np.max(np.array(polygon).T[1])
y = np.min(np.array(polygon).T[1])
v = plotarray[cell]
if v > y1:
v = y
if v < y:
v = y
p1 = [(x[0], y1), (x[1], y1), (x[1], v), (x[0], v)]
p2 = [(x[0], v), (x[1], v), (x[1], y), (x[0], y)]
rectcol.append(Polygon(p1, closed=True, color=colors[0]))
rectcol.append(Polygon(p2, closed=True, color=colors[1]))
else:
rectcol.append(Polygon(polygon, closed=True))
data.append(plotarray[cell])
if len(rectcol) > 0:
patches = PatchCollection(rectcol, match_original, **kwargs)
if not fill_between:
patches.set_array(np.array(data))
patches.set_clim(vmin, vmax)
else:
patches = None
return patches