def set_zpts(self, vs):
"""
Get an array of z elevations based on minimum of cell elevation
(self.elev) or passed vs numpy.ndarray
Parameters
----------
vs : numpy.ndarray
Three-dimensional array to plot.
Returns
-------
zpts : numpy.ndarray
"""
zpts = []
for k in xrange(self.layer0, self.layer1):
e = self.elev[k, :, :]
if k < self.dis.nlay:
v = vs[k, :, :]
idx = v < e
e[idx] = v[idx]
zpts.append(
plotutil.cell_value_points(self.xpts, self.xedge, self.yedge,
e))
return np.array(zpts)
def set_zpts(self, vs):
"""
Get an array of z elevations based on minimum of cell elevation
(self.elev) or passed vs numpy.ndarray
Parameters
----------
vs : numpy.ndarray
Three-dimensional array to plot.
Returns
-------
zpts : numpy.ndarray
"""
zpts = []
for k in xrange(self.layer0, self.layer1):
e = self.elev[k, :, :]
if k < self.dis.nlay:
v = vs[k, :, :]
idx = v < e
e[idx] = v[idx]
zpts.append(plotutil.cell_value_points(self.xpts, self.xedge,
self.yedge, e))
return np.array(zpts)
def set_zcentergrid(self, vs):
"""
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.
Returns
-------
zcentergrid : numpy.ndarray
"""
vpts = []
for k in xrange(self.layer0, self.layer1):
if k < self.dis.nlay:
e = vs[k, :, :]
else:
e = self.elev[k, :, :]
vpts.append(
plotutil.cell_value_points(self.xpts, self.xedge, self.yedge,
e))
vpts = np.array(vpts)
zcentergrid = []
nz = 0
if self.dis.nlay == 1:
for k in xrange(0, self.zpts.shape[0]):
nz += 1
nx = 0
for i in xrange(0, self.xpts.shape[0], 2):
nx += 1
vp = vpts[k, i]
zp = self.zpts[k, i]
if k == 0:
if vp < zp:
zp = vp
zcentergrid.append(zp)
else:
for k in xrange(0, self.zpts.shape[0], 2):
nz += 1
nx = 0
for i in xrange(0, self.xpts.shape[0], 2):
nx += 1
vp = vpts[k, i]
ep = self.zpts[k, i]
if vp < ep:
ep = vp
zp = 0.5 * (ep + self.zpts[k + 1, i + 1])
zcentergrid.append(zp)
return np.array(zcentergrid).reshape((nz, nx))
def set_zcentergrid(self, vs):
"""
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.
Returns
-------
zcentergrid : numpy.ndarray
"""
vpts = []
for k in xrange(self.layer0, self.layer1):
if k < self.dis.nlay:
e = vs[k, :, :]
else:
e = self.elev[k, :, :]
vpts.append(plotutil.cell_value_points(self.xpts, self.xedge,
self.yedge, e))
vpts = np.array(vpts)
zcentergrid = []
nz = 0
if self.dis.nlay == 1:
for k in xrange(0, self.zpts.shape[0]):
nz += 1
nx = 0
for i in xrange(0, self.xpts.shape[0], 2):
nx += 1
vp = vpts[k, i]
zp = self.zpts[k, i]
if k == 0:
if vp < zp:
zp = vp
zcentergrid.append(zp)
else:
for k in xrange(0, self.zpts.shape[0], 2):
nz += 1
nx = 0
for i in xrange(0, self.xpts.shape[0], 2):
nx += 1
vp = vpts[k, i]
ep = self.zpts[k, i]
if vp < ep:
ep = vp
zp = 0.5 * (ep + self.zpts[k+1, i+1])
zcentergrid.append(zp)
return np.array(zcentergrid).reshape((nz, nx))
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
"""
if 'ax' in kwargs:
ax = kwargs.pop('ax')
else:
ax = self.ax
plotarray = a
if masked_values is not None:
for mval in masked_values:
plotarray = np.ma.masked_equal(plotarray, mval)
vpts = []
for k in xrange(self.dis.nlay):
vpts.append(plotutil.cell_value_points(self.xpts, self.xedge,
self.yedge,
plotarray[k, :, :]))
vpts = np.array(vpts)
if isinstance(head, np.ndarray):
zpts = self.set_zpts(head)
else:
zpts = self.zpts
if masked_values is not None:
for mval in masked_values:
vpts = np.ma.masked_equal(vpts, mval)
pc = self.get_grid_patch_collection(zpts, vpts, **kwargs)
if pc != None:
ax.add_collection(pc)
return pc
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
"""
if 'ax' in kwargs:
ax = kwargs.pop('ax')
else:
ax = self.ax
plotarray = a
if masked_values is not None:
for mval in masked_values:
plotarray = np.ma.masked_equal(plotarray, mval)
vpts = []
for k in xrange(self.dis.nlay):
vpts.append(
plotutil.cell_value_points(self.xpts, self.xedge, self.yedge,
plotarray[k, :, :]))
vpts = np.array(vpts)
if isinstance(head, np.ndarray):
zpts = self.set_zpts(head)
else:
zpts = self.zpts
if masked_values is not None:
for mval in masked_values:
vpts = np.ma.masked_equal(vpts, mval)
pc = self.get_grid_patch_collection(zpts, vpts, **kwargs)
if pc != None:
ax.add_collection(pc)
return pc
def contour_array(self, a, masked_values=None, head=None, **kwargs):
"""
Contour a three-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
"""
plotarray = a
vpts = []
for k in xrange(self.dis.nlay):
vpts.append(plotutil.cell_value_points(self.xpts, self.xedge,
self.yedge,
plotarray[k, :, :]))
vpts = np.array(vpts)
vpts = vpts[:, ::2]
if self.dis.nlay == 1:
vpts = np.vstack((vpts, vpts))
if masked_values is not None:
for mval in masked_values:
vpts = np.ma.masked_equal(vpts, mval)
if isinstance(head, np.ndarray):
zcentergrid = self.set_zcentergrid(head)
else:
zcentergrid = self.zcentergrid
contour_set = self.ax.contour(self.xcentergrid, zcentergrid,
vpts, **kwargs)
return contour_set
def contour_array(self, a, masked_values=None, head=None, **kwargs):
"""
Contour a three-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
"""
plotarray = a
vpts = []
for k in xrange(self.dis.nlay):
vpts.append(
plotutil.cell_value_points(self.xpts, self.xedge, self.yedge,
plotarray[k, :, :]))
vpts = np.array(vpts)
vpts = vpts[:, ::2]
if self.dis.nlay == 1:
vpts = np.vstack((vpts, vpts))
if masked_values is not None:
for mval in masked_values:
vpts = np.ma.masked_equal(vpts, mval)
if isinstance(head, np.ndarray):
zcentergrid = self.set_zcentergrid(head)
else:
zcentergrid = self.zcentergrid
contour_set = self.ax.contour(self.xcentergrid, zcentergrid, vpts,
**kwargs)
return contour_set
def plot_surface(self, a, masked_values=None, **kwargs):
"""
Plot a three-dimensional array as lines.
Parameters
----------
a : numpy.ndarray
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
"""
if 'ax' in kwargs:
ax = kwargs.pop('ax')
else:
ax = self.ax
plotarray = a
vpts = []
for k in xrange(self.dis.nlay):
vpts.append(plotutil.cell_value_points(self.xpts, self.xedge,
self.yedge,
plotarray[k, :, :]))
vpts = np.array(vpts)
if masked_values is not None:
for mval in masked_values:
vpts = np.ma.masked_equal(vpts, mval)
plot = []
for k in xrange(vpts.shape[0]):
plot.append(ax.plot(self.d, vpts[k, :], **kwargs))
return plot
def plot_surface(self, a, masked_values=None, **kwargs):
"""
Plot a three-dimensional array as lines.
Parameters
----------
a : numpy.ndarray
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
"""
if 'ax' in kwargs:
ax = kwargs.pop('ax')
else:
ax = self.ax
plotarray = a
vpts = []
for k in xrange(self.dis.nlay):
vpts.append(
plotutil.cell_value_points(self.xpts, self.xedge, self.yedge,
plotarray[k, :, :]))
vpts = np.array(vpts)
if masked_values is not None:
for mval in masked_values:
vpts = np.ma.masked_equal(vpts, mval)
plot = []
for k in xrange(vpts.shape[0]):
plot.append(ax.plot(self.d, vpts[k, :], **kwargs))
return plot
def plot_discharge(self, frf, fff, flf=None, head=None,
kstep=1, hstep=1,
**kwargs):
"""
Use quiver to plot vectors.
Parameters
----------
frf : numpy.ndarray
MODFLOW's 'flow right face'
fff : numpy.ndarray
MODFLOW's 'flow front face'
flf : numpy.ndarray
MODFLOW's 'flow lower face' (Default is None.)
head : numpy.ndarray
MODFLOW's head array. If not provided, then will assume confined
conditions in order to calculated saturated thickness.
kstep : int
layer frequency to plot. (Default is 1.)
hstep : int
horizontal frequency to plot. (Default is 1.)
kwargs : dictionary
Keyword arguments passed to plt.quiver()
Returns
-------
quiver : matplotlib.pyplot.quiver
Vectors
"""
# Calculate specific discharge
delr = self.dis.delr.array
delc = self.dis.delc.array
top = self.dis.top.array
botm = self.dis.botm.array
nlay, nrow, ncol = botm.shape
laytyp = None
hnoflo = 999.
hdry = 999.
if self.model is not None:
lpf = self.model.get_package('LPF')
if lpf is not None:
laytyp = lpf.laytyp.array
hdry = lpf.hdry
bas = self.model.get_package('BAS6')
if bas is not None:
hnoflo = bas.hnoflo
# If no access to head or laytyp, then calculate confined saturated
# thickness by setting laytyp to zeros
if head is None or laytyp is None:
head = np.zeros(botm.shape, np.float32)
laytyp = np.zeros((nlay), dtype=np.int)
sat_thk = plotutil.saturated_thickness(head, top, botm, laytyp,
[hnoflo, hdry])
# Calculate specific discharge
qx, qy, qz = plotutil.centered_specific_discharge(frf, fff, flf, delr,
delc, sat_thk)
if qz == None:
qz = np.zeros((qx.shape), dtype=np.float)
# Select correct specific discharge direction
if self.direction == 'x':
u = -qx[:, :, :]
u2 = -qy[:, :, :]
v = qz[:, :, :]
elif self.direction == 'y':
u = -qy[:, :, :]
u2 = -qx[:, :, :]
v = qz[:, :, :]
elif self.direction == 'xy':
print 'csplot_discharge does not support arbitrary cross-sections'
return None
if isinstance(head, np.ndarray):
zcentergrid = self.set_zcentergrid(head)
else:
zcentergrid = self.zcentergrid
if nlay == 1:
x = []
z = []
for k in xrange(1):
for i in xrange(self.xcentergrid.shape[1]):
x.append(self.xcentergrid[k, i])
z.append(0.5 * (zcentergrid[k, i] + zcentergrid[k+1, i]))
x = np.array(x).reshape((1,self.xcentergrid.shape[1]))
z = np.array(z).reshape((1,self.xcentergrid.shape[1]))
else:
x = self.xcentergrid
z = zcentergrid
upts = []
u2pts = []
vpts = []
for k in xrange(self.dis.nlay):
upts.append(plotutil.cell_value_points(self.xpts, self.xedge,
self.yedge, u[k, :, :]))
u2pts.append(plotutil.cell_value_points(self.xpts, self.xedge,
self.yedge, u2[k, :, :]))
vpts.append(plotutil.cell_value_points(self.xpts, self.xedge,
self.yedge, v[k, :, :]))
upts = np.array(upts)
u2pts = np.array(u2pts)
vpts = np.array(vpts)
x = x[::kstep, ::hstep]
z = z[::kstep, ::hstep]
upts = upts[::kstep, ::hstep]
u2pts = u2pts[::kstep, ::hstep]
vpts = vpts[::kstep, ::hstep]
N = np.sqrt(upts**2. + u2pts**2. + vpts**2.)
idx = N > 0.
upts[idx] /= N[idx]
u2pts[idx] /= N[idx]
vpts[idx] /= N[idx]
# plot the vectors
quiver = self.ax.quiver(x, z, upts, vpts, **kwargs)
return quiver
def __init__(self, ax=None, model=None, dis=None, line=None,
xul=None, yul=None, rotation=0., extent=None):
self.model = model
if dis is None:
if model is None:
raise Exception('Cannot find discretization package')
else:
self.dis = model.get_package('DIS')
else:
self.dis = dis
if line == None:
s = 'line must be specified.'
raise Exception(s)
linekeys = [linekeys.lower() for linekeys in line.keys()]
if len(linekeys) < 1:
s = 'only row, column, or line can be specified in line dictionary.\n'
s += 'keys specified: '
for k in linekeys:
s += '{} '.format(k)
raise Exception(s)
if 'row' in linekeys and 'column' in linekeys:
s = 'row and column cannot both be specified in line dictionary.'
raise Exception(s)
if 'row' in linekeys and 'line' in linekeys:
s = 'row and line cannot both be specified in line dictionary.'
raise Exception(s)
if 'column' in linekeys and 'line' in linekeys:
s = 'column and line cannot both be specified in line dictionary.'
raise Exception(s)
if ax is None:
self.ax = plt.gca()
else:
self.ax = ax
# Set origin and rotation
if xul is None:
self.xul = 0.
else:
self.xul = xul
if yul is None:
self.yul = 0
else:
self.yul = yul
self.rotation = -rotation * np.pi / 180.
# Create edge arrays and meshgrid for pcolormesh
self.xedge = self.get_xedge_array()
self.yedge = self.get_yedge_array()
# Create x and y center arrays and meshgrid of centers
self.xcenter = self.get_xcenter_array()
self.ycenter = self.get_ycenter_array()
onkey = line.keys()[0]
if 'row' in linekeys:
self.direction = 'x'
pts = [(self.xedge[0]+0.1, self.ycenter[int(line[onkey])]-0.1),
(self.xedge[-1]-0.1, self.ycenter[int(line[onkey])]+0.1)]
elif 'column' in linekeys:
self.direction = 'y'
pts = [(self.xcenter[int(line[onkey])]+0.1, self.yedge[0]-0.1),
(self.xcenter[int(line[onkey])]-0.1, self.yedge[-1]+0.1)]
else:
self.direction = 'xy'
verts = line[onkey]
xp = []
yp = []
for [v1, v2] in verts:
xp.append(v1)
yp.append(v2)
xp, yp = np.array(xp), np.array(yp)
# remove offset and rotation from line
xp -= self.xul
yp -= self.yul
xp, yp = rotate(xp, yp, -self.rotation, 0, self.yedge[0])
pts = []
for xt, yt in zip(xp, yp):
pts.append((xt, yt))
# convert pts list to numpy array
self.pts = np.array(pts)
# get points along the line
self.xpts = plotutil.line_intersect_grid(self.pts, self.xedge,
self.yedge)
if len(self.xpts) < 2:
s = 'cross-section cannot be created\n.'
s += ' less than 2 points intersect the model grid\n'
s += ' {} points intersect the grid.'.format(len(self.xpts))
raise Exception(s)
# set horizontal distance
d = []
for v in self.xpts:
d.append(v[2])
self.d = np.array(d)
top = self.dis.top.array
botm = self.dis.botm.array
elev = [top.copy()]
for k in xrange(self.dis.nlay):
elev.append(botm[k, :, :])
self.elev = np.array(elev)
self.layer0 = 0
self.layer1 = self.dis.nlay + 1
zpts = []
for k in xrange(self.layer0, self.layer1):
zpts.append(plotutil.cell_value_points(self.xpts, self.xedge,
self.yedge,
self.elev[k, :, :]))
self.zpts = np.array(zpts)
xcentergrid = []
zcentergrid = []
nz = 0
if self.dis.nlay == 1:
for k in xrange(0, self.zpts.shape[0]):
nz += 1
nx = 0
for i in xrange(0, self.xpts.shape[0], 2):
try:
xp = 0.5 * (self.xpts[i][2] + self.xpts[i+1][2])
zp = self.zpts[k, i]
xcentergrid.append(xp)
zcentergrid.append(zp)
nx += 1
except:
break
else:
for k in xrange(0, self.zpts.shape[0]-1):
nz += 1
nx = 0
for i in xrange(0, self.xpts.shape[0], 2):
try:
xp = 0.5 * (self.xpts[i][2] + self.xpts[i+1][2])
zp = 0.5 * (self.zpts[k, i] + self.zpts[k+1, i+1])
xcentergrid.append(xp)
zcentergrid.append(zp)
nx += 1
except:
break
self.xcentergrid = np.array(xcentergrid).reshape((nz, nx))
self.zcentergrid = np.array(zcentergrid).reshape((nz, nx))
# Create cross-section extent
if extent is None:
self.extent = self.get_extent()
else:
self.extent = extent
# Set axis limits
self.ax.set_xlim(self.extent[0], self.extent[1])
self.ax.set_ylim(self.extent[2], self.extent[3])
return
def plot_fill_between(self, a, colors=['blue', 'red'],
masked_values=None, **kwargs):
"""
Plot a three-dimensional array as lines.
Parameters
----------
a : numpy.ndarray
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.fillbetween objects
"""
if 'ax' in kwargs:
ax = kwargs.pop('ax')
else:
ax = self.ax
plotarray = a
vpts = []
for k in xrange(self.dis.nlay):
vpts.append(plotutil.cell_value_points(self.xpts, self.xedge,
self.yedge,
plotarray[k, :, :]))
vpts = np.ma.array(vpts, mask=False)
if masked_values is not None:
for mval in masked_values:
vpts = np.ma.masked_equal(vpts, mval)
idxm = np.ma.getmask(vpts)
plot = []
for k in xrange(self.dis.nlay):
idxmk = idxm[k, :]
v = vpts[k, :]
y1 = self.zpts[k, :]
y2 = self.zpts[k+1, :]
#--make sure y1 is not below y2
idx = y1 < y2
y1[idx] = y2[idx]
#--make sure v is not below y2
idx = v < y2
v[idx] = y2[idx]
#--make sure v is not above y1
idx = v > y1
v[idx] = y1[idx]
#--set y2 to v
y2 = v
#--mask cells
y1[idxmk] = np.nan
y2[idxmk] = np.nan
plot.append(ax.fill_between(self.d, y1=y1, y2=y2, color=colors[0],
**kwargs))
y1 = y2
y2 = self.zpts[k+1, :]
y2[idxmk] = np.nan
plot.append(ax.fill_between(self.d, y1=y1, y2=y2, color=colors[1],
**kwargs))
return plot
def plot_discharge(self,
frf,
fff,
flf=None,
head=None,
kstep=1,
hstep=1,
**kwargs):
"""
Use quiver to plot vectors.
Parameters
----------
frf : numpy.ndarray
MODFLOW's 'flow right face'
fff : numpy.ndarray
MODFLOW's 'flow front face'
flf : numpy.ndarray
MODFLOW's 'flow lower face' (Default is None.)
head : numpy.ndarray
MODFLOW's head array. If not provided, then will assume confined
conditions in order to calculated saturated thickness.
kstep : int
layer frequency to plot. (Default is 1.)
hstep : int
horizontal frequency to plot. (Default is 1.)
kwargs : dictionary
Keyword arguments passed to plt.quiver()
Returns
-------
quiver : matplotlib.pyplot.quiver
Vectors
"""
# Calculate specific discharge
delr = self.dis.delr.array
delc = self.dis.delc.array
top = self.dis.top.array
botm = self.dis.botm.array
nlay, nrow, ncol = botm.shape
laytyp = None
hnoflo = 999.
hdry = 999.
if self.model is not None:
lpf = self.model.get_package('LPF')
if lpf is not None:
laytyp = lpf.laytyp.array
hdry = lpf.hdry
bas = self.model.get_package('BAS6')
if bas is not None:
hnoflo = bas.hnoflo
# If no access to head or laytyp, then calculate confined saturated
# thickness by setting laytyp to zeros
if head is None or laytyp is None:
head = np.zeros(botm.shape, np.float32)
laytyp = np.zeros((nlay), dtype=np.int)
sat_thk = plotutil.saturated_thickness(head, top, botm, laytyp,
[hnoflo, hdry])
# Calculate specific discharge
qx, qy, qz = plotutil.centered_specific_discharge(
frf, fff, flf, delr, delc, sat_thk)
if qz == None:
qz = np.zeros((qx.shape), dtype=np.float)
# Select correct specific discharge direction
if self.direction == 'x':
u = -qx[:, :, :]
u2 = -qy[:, :, :]
v = qz[:, :, :]
elif self.direction == 'y':
u = -qy[:, :, :]
u2 = -qx[:, :, :]
v = qz[:, :, :]
elif self.direction == 'xy':
print 'csplot_discharge does not support arbitrary cross-sections'
return None
if isinstance(head, np.ndarray):
zcentergrid = self.set_zcentergrid(head)
else:
zcentergrid = self.zcentergrid
if nlay == 1:
x = []
z = []
for k in xrange(1):
for i in xrange(self.xcentergrid.shape[1]):
x.append(self.xcentergrid[k, i])
z.append(0.5 * (zcentergrid[k, i] + zcentergrid[k + 1, i]))
x = np.array(x).reshape((1, self.xcentergrid.shape[1]))
z = np.array(z).reshape((1, self.xcentergrid.shape[1]))
else:
x = self.xcentergrid
z = zcentergrid
upts = []
u2pts = []
vpts = []
for k in xrange(self.dis.nlay):
upts.append(
plotutil.cell_value_points(self.xpts, self.xedge, self.yedge,
u[k, :, :]))
u2pts.append(
plotutil.cell_value_points(self.xpts, self.xedge, self.yedge,
u2[k, :, :]))
vpts.append(
plotutil.cell_value_points(self.xpts, self.xedge, self.yedge,
v[k, :, :]))
upts = np.array(upts)
u2pts = np.array(u2pts)
vpts = np.array(vpts)
x = x[::kstep, ::hstep]
z = z[::kstep, ::hstep]
upts = upts[::kstep, ::hstep]
u2pts = u2pts[::kstep, ::hstep]
vpts = vpts[::kstep, ::hstep]
N = np.sqrt(upts**2. + u2pts**2. + vpts**2.)
idx = N > 0.
upts[idx] /= N[idx]
u2pts[idx] /= N[idx]
vpts[idx] /= N[idx]
# plot the vectors
quiver = self.ax.quiver(x, z, upts, vpts, **kwargs)
return quiver
def __init__(self,
ax=None,
model=None,
dis=None,
line=None,
xul=None,
yul=None,
rotation=0.,
extent=None):
self.model = model
if dis is None:
if model is None:
raise Exception('Cannot find discretization package')
else:
self.dis = model.get_package('DIS')
else:
self.dis = dis
if line == None:
s = 'line must be specified.'
raise Exception(s)
linekeys = [linekeys.lower() for linekeys in line.keys()]
if len(linekeys) < 1:
s = 'only row, column, or line can be specified in line dictionary.\n'
s += 'keys specified: '
for k in linekeys:
s += '{} '.format(k)
raise Exception(s)
if 'row' in linekeys and 'column' in linekeys:
s = 'row and column cannot both be specified in line dictionary.'
raise Exception(s)
if 'row' in linekeys and 'line' in linekeys:
s = 'row and line cannot both be specified in line dictionary.'
raise Exception(s)
if 'column' in linekeys and 'line' in linekeys:
s = 'column and line cannot both be specified in line dictionary.'
raise Exception(s)
if ax is None:
self.ax = plt.gca()
else:
self.ax = ax
# Set origin and rotation
if xul is None:
self.xul = 0.
else:
self.xul = xul
if yul is None:
self.yul = 0
else:
self.yul = yul
self.rotation = -rotation * np.pi / 180.
# Create edge arrays and meshgrid for pcolormesh
self.xedge = self.get_xedge_array()
self.yedge = self.get_yedge_array()
# Create x and y center arrays and meshgrid of centers
self.xcenter = self.get_xcenter_array()
self.ycenter = self.get_ycenter_array()
onkey = line.keys()[0]
if 'row' in linekeys:
self.direction = 'x'
pts = [(self.xedge[0] + 0.1, self.ycenter[int(line[onkey])] - 0.1),
(self.xedge[-1] - 0.1, self.ycenter[int(line[onkey])] + 0.1)
]
elif 'column' in linekeys:
self.direction = 'y'
pts = [(self.xcenter[int(line[onkey])] + 0.1, self.yedge[0] - 0.1),
(self.xcenter[int(line[onkey])] - 0.1, self.yedge[-1] + 0.1)
]
else:
self.direction = 'xy'
verts = line[onkey]
xp = []
yp = []
for [v1, v2] in verts:
xp.append(v1)
yp.append(v2)
xp, yp = np.array(xp), np.array(yp)
# remove offset and rotation from line
xp -= self.xul
yp -= self.yul
xp, yp = rotate(xp, yp, -self.rotation, 0, self.yedge[0])
pts = []
for xt, yt in zip(xp, yp):
pts.append((xt, yt))
# convert pts list to numpy array
self.pts = np.array(pts)
# get points along the line
self.xpts = plotutil.line_intersect_grid(self.pts, self.xedge,
self.yedge)
if len(self.xpts) < 2:
s = 'cross-section cannot be created\n.'
s += ' less than 2 points intersect the model grid\n'
s += ' {} points intersect the grid.'.format(len(self.xpts))
raise Exception(s)
# set horizontal distance
d = []
for v in self.xpts:
d.append(v[2])
self.d = np.array(d)
top = self.dis.top.array
botm = self.dis.botm.array
elev = [top.copy()]
for k in xrange(self.dis.nlay):
elev.append(botm[k, :, :])
self.elev = np.array(elev)
self.layer0 = 0
self.layer1 = self.dis.nlay + 1
zpts = []
for k in xrange(self.layer0, self.layer1):
zpts.append(
plotutil.cell_value_points(self.xpts, self.xedge, self.yedge,
self.elev[k, :, :]))
self.zpts = np.array(zpts)
xcentergrid = []
zcentergrid = []
nz = 0
if self.dis.nlay == 1:
for k in xrange(0, self.zpts.shape[0]):
nz += 1
nx = 0
for i in xrange(0, self.xpts.shape[0], 2):
try:
xp = 0.5 * (self.xpts[i][2] + self.xpts[i + 1][2])
zp = self.zpts[k, i]
xcentergrid.append(xp)
zcentergrid.append(zp)
nx += 1
except:
break
else:
for k in xrange(0, self.zpts.shape[0] - 1):
nz += 1
nx = 0
for i in xrange(0, self.xpts.shape[0], 2):
try:
xp = 0.5 * (self.xpts[i][2] + self.xpts[i + 1][2])
zp = 0.5 * (self.zpts[k, i] + self.zpts[k + 1, i + 1])
xcentergrid.append(xp)
zcentergrid.append(zp)
nx += 1
except:
break
self.xcentergrid = np.array(xcentergrid).reshape((nz, nx))
self.zcentergrid = np.array(zcentergrid).reshape((nz, nx))
# Create cross-section extent
if extent is None:
self.extent = self.get_extent()
else:
self.extent = extent
# Set axis limits
self.ax.set_xlim(self.extent[0], self.extent[1])
self.ax.set_ylim(self.extent[2], self.extent[3])
return
def plot_fill_between(self,
a,
colors=['blue', 'red'],
masked_values=None,
**kwargs):
"""
Plot a three-dimensional array as lines.
Parameters
----------
a : numpy.ndarray
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.fillbetween objects
"""
if 'ax' in kwargs:
ax = kwargs.pop('ax')
else:
ax = self.ax
plotarray = a
vpts = []
for k in xrange(self.dis.nlay):
vpts.append(
plotutil.cell_value_points(self.xpts, self.xedge, self.yedge,
plotarray[k, :, :]))
vpts = np.ma.array(vpts, mask=False)
if masked_values is not None:
for mval in masked_values:
vpts = np.ma.masked_equal(vpts, mval)
idxm = np.ma.getmask(vpts)
plot = []
for k in xrange(self.dis.nlay):
idxmk = idxm[k, :]
v = vpts[k, :]
y1 = self.zpts[k, :]
y2 = self.zpts[k + 1, :]
#--make sure y1 is not below y2
idx = y1 < y2
y1[idx] = y2[idx]
#--make sure v is not below y2
idx = v < y2
v[idx] = y2[idx]
#--make sure v is not above y1
idx = v > y1
v[idx] = y1[idx]
#--set y2 to v
y2 = v
#--mask cells
y1[idxmk] = np.nan
y2[idxmk] = np.nan
plot.append(
ax.fill_between(self.d,
y1=y1,
y2=y2,
color=colors[0],
**kwargs))
y1 = y2
y2 = self.zpts[k + 1, :]
y2[idxmk] = np.nan
plot.append(
ax.fill_between(self.d,
y1=y1,
y2=y2,
color=colors[1],
**kwargs))
return plot
