"""
Module that contains various functions dedicated to plotting using |pypago|
"""
import numpy as np
import pylab as plt
from matplotlib.colors import ListedColormap
import pypago.misc
[docs]def plot_dom_mask(grid, gridsec=None, mask=None, ax=None):
"""
Draws the final plot.
It draws the model mask and the model sections,
and it fills the mask of the area
:param pago_obj grid: |pypago| object that contains the model grid
:param pago_obj gridsec: |pypago| object that contains the model
sections
:param numpy.array mask: mask of the area
"""
if ax is None:
ax = plt.gca()
cmap2 = ListedColormap(['none', 'white'])
cmap = ListedColormap(['black', 'white', 'gray'])
plt.imshow(grid.mask, interpolation='none', cmap=cmap)
if mask is not None:
plt.imshow(mask, interpolation='none', cmap=cmap2)
if gridsec is not None:
for sec in gridsec:
sec.plotsecfaces(ax)
plt.title('sections in grid')
plt.xlabel('grid points')
plt.ylabel('grid points')
plt.xlim(0, grid.mask.shape[1])
plt.ylim(0, grid.mask.shape[0])
[docs]def _data_prep(secstruct, vartoplot, istracer=None, itime=None):
""" Data extraction, preparation and masking prior to plotting """
if not(hasattr(secstruct, vartoplot)):
message = 'The "%s" variables is not an attribute '
message += 'of the input section.'
raise PypagoErrors(message)
secdata = getattr(secstruct, vartoplot)
sectionname = secstruct.name
if secdata.ndim == 3:
if itime is None:
message = 'The time average of the %s ' %vartoplot
message += 'is plotted.'
print(message)
secdata = np.mean(secdata, axis=0)
else:
try:
secdata = secdata[itime, :, :]
except:
message = 'The data could not be extracted. '
message += 'Check that the time index is ok. '
message += 'This program will be stopped'
raise PypagoErrors(message)
atracer, zvect, lvect, \
atracer_pcol, zvect_pcol, lvect_pcol = preplot(secstruct, secdata, istracer)
atracer_pcol = np.ma.masked_where(pypago.misc.extract_mask(atracer_pcol), atracer_pcol)
atracer = np.ma.masked_where(pypago.misc.extract_mask(atracer), atracer)
lvect_pcol = np.ma.masked_where(pypago.misc.extract_mask(lvect_pcol), lvect_pcol)
lvect = np.ma.masked_where(pypago.misc.extract_mask(lvect), lvect)
zvect_pcol = np.ma.masked_where(zvect_pcol==0, zvect_pcol)
lvect_pcol = np.ma.masked_where(np.ma.getmaskarray(zvect_pcol), lvect_pcol)
return atracer, zvect, lvect, \
atracer_pcol, zvect_pcol, lvect_pcol
[docs]def pcolplot(secstruct, vartoplot, istracer, itime=None, ax=None):
""" Pcolor of a variable contained in a gridded section.
:param pypago.sections.GridSection secstruct: Gridded section
:param str vartoplot: Variable to plot
:param bool istracer: True if tracer field (i.e. temperature),
False if velocity field (or transport field).
:param int itime: Time index to extract. If None,
temporal mean is computed
:param ax ax: Axis. If None, gca()
:return: cs, cb
"""
atracer, zvect, lvect, \
atracer_pcol, zvect_pcol, lvect_pcol = _data_prep(secstruct, vartoplot, istracer, itime)
zmin, zmax = zvect.min(), zvect.max()
lmin, lmax = lvect.min(), lvect.max()
if ax is None:
ax = plt.gca()
cs = plt.pcolormesh(lvect_pcol, zvect_pcol, atracer_pcol)
cb = plt.colorbar(cs)
cmin, cmax = pypago.misc.make_percentile_cmap(atracer_pcol, 5)
cs.set_clim(cmin, cmax)
if not istracer:
cmin, cmax = cs.get_clim()
cmax = np.max(np.abs([cmin, cmax]))
cs.set_clim(-cmax, cmax)
ax.set_facecolor("gray")
ax.set_xlim(lmin, lmax)
ax.set_ylim(zmin, zmax)
ax.set_xticks(np.linspace(lmin, lmax, 5))
return cs, cb
[docs]def contourplot(secstruct, vartoplot, istracer, itime=None, ax=None, **kwargs):
""" Contourplot of a variable contained in a gridded section.
:param pypago.sections.GridSection secstruct: Gridded section
:param str vartoplot: Variable to plot
:param bool istracer: True if tracer field (i.e. temperature),
False if velocity field (or transport field).
:param int itime: Time index to extract. If None,
temporal mean is computed
:param ax ax: Axis. If None, gca()
:return: cs, cb
"""
atracer, zvect, lvect, \
atracer_pcol, zvect_pcol, lvect_pcol = _data_prep(secstruct, vartoplot, istracer, itime)
zmin, zmax = zvect.min(), zvect.max()
lmin, lmax = lvect.min(), lvect.max()
if ax is None:
ax = plt.gca()
cl = plt.contour(lvect, zvect, atracer, **kwargs)
ax.set_facecolor("gray")
ax.set_xlim(lmin, lmax)
ax.set_ylim(zmin, zmax)
ax.set_xticks(np.linspace(lmin, lmax, 5))
return cl
[docs]def preplot(secstruct, secdata, istracer):
"""
Function that prepares a section tracer (T or S array)
variable to do both a contourf and a pcolor plot.
:param pago_obj secstruct: The |pypago| object that contains the section
variables (`veci`, `vecj`, etc.)
:param numpy.array secdata: Array that contains the value to
plot (vect or vecs, could be a mean, a snapshot, etc.)
:param str istracer: True if a tracer field (vect, vecs) is plotted
"""
nz = len(secstruct.depthvect[:, 0])
nl = len(secstruct.lvect)
atracer = np.zeros((nz, nl))
zvect = np.zeros((nz, nl))
for indz in xrange(0, nz):
atracer[indz, :] = _nodouble_tsr(secstruct.veci,
secstruct.vecj, secdata[indz, :])
if istracer:
zvect[indz, :] = _nodouble_tsr(secstruct.veci, secstruct.vecj,
secstruct.depthvect[indz, :])
else:
zvect[indz, :] = _nodouble_v(secstruct.veci, secstruct.vecj,
secstruct.depthvect[indz, :])
# barrier.n: removed mask from zvect since it causes problem
# when plotting. we set to 0 so that no effect on the cumsum
zvect = np.array(zvect)
zvect = np.ma.masked_where(np.ma.getmaskarray(zvect), zvect)
zvect = np.ma.masked_where(np.isnan(zvect), zvect)
zvect = -np.cumsum(zvect, axis=0)
lvect = np.tile(secstruct.lvect, (nz, 1))
lvect_pcol = _make_lvect_pcol(lvect)
zvect_pcol = _make_zvect_pcol(zvect)
atracer_pcol = _make_atracer_pcol(atracer)
return atracer, zvect, lvect, \
atracer_pcol, zvect_pcol, lvect_pcol
[docs]def _nodouble_v(veci, vecj, vecv):
"""
Returns the velocity field after removing
the duplicated values of (`veci`, `vecj`)
.. warning::
Relies on the assumption that all sections go
from West to East
:param numpy.array veci: i indexes of the section
:param numpy.array vecj: j indexes of the section
:param numpy.array vecv: velocity field of the section
:return: the new velocity
:rtype: numpy.array
"""
veci = np.squeeze(veci)
vecj = np.squeeze(vecj)
vecv = np.squeeze(vecv)
newvecv = np.ma.masked_array([])
newvecv = np.ma.append(newvecv, vecv[0])
for indl in xrange(1, veci.shape[0]):
if not ((veci[indl] == veci[indl-1]) & (vecj[indl] == vecj[indl-1])):
newvecv = np.ma.append(newvecv, vecv[indl])
else:
temp = np.arctan(vecv[indl]/(-vecv[indl-1]))
if -vecv[indl-1] < 0:
if vecv[indl] >= 0:
temp += np.pi
else:
temp -= np.pi
if (temp > np.pi/4.) | (temp < -3*np.pi/4.):
signe = 1
else:
signe = -1
# barrier.n: correction according to Matlab's version
#if vecv[indl-1]*vecv[indl-1] + vecv[indl]*vecv[indl] >= 0:
# newvecv[-1] = signe*np.sqrt(vecv[indl-1]*vecv[indl-1] +
# vecv[indl]*vecv[indl])
#else:
# newvecv[-1] = np.nan
newvecv[-1] = signe * np.sqrt(vecv[l-1]**2 + vecv[l]**2)
return newvecv
[docs]def _nodouble_tsr(veci, vecj, vectsr):
"""
Returns the tracer field after removing
the duplicated values of (`veci`, `vecj`)
:param numpy.array veci: i indexes of the section
:param numpy.array vecj: j indexes of the section
:param numpy.array vectsr: tracer field of the section
:return: the tracer field without the duplications
:rtype: list
"""
veci = np.squeeze(veci)
vecj = np.squeeze(vecj)
vectsr = np.squeeze(vectsr)
newvectsr = []
newvectsr.append(vectsr[0])
for indl in xrange(1, veci.shape[0]):
if not ((veci[indl] == veci[indl-1]) & (vecj[indl] == vecj[indl-1])):
newvectsr.append(vectsr[indl])
newvectstr = np.array(newvectsr)
return newvectsr
[docs]def _make_lvect_pcol(lvect):
"""
Returns the length vector for a pcolor plot
:param numpy.array lvect: length vector
:return: array that contains the `lvect` array formatted for a pcolor plot
:rtype: numpy.array
"""
nl = lvect.shape[1]
nz = lvect.shape[0]
lvect = np.ma.masked_where(pypago.misc.extract_mask(lvect), lvect)
lvect_pcol = 1.5 * lvect[:, 0:1] - 0.5 * lvect[:, 1:2]
lvect_pcol = np.concatenate((lvect_pcol, .5 * (lvect[:, :-1] + lvect[:, 1:])), axis=1)
lvect_pcol = np.concatenate((lvect_pcol, 1.5 * lvect[:, -1:] - 0.5 * lvect[:, -2:-1]), axis=1)
lvect_pcol = np.concatenate((lvect_pcol[0:1, :], lvect_pcol), axis=0)
return lvect_pcol
[docs]def _make_zvect_pcol(zvect):
"""
Returns the depth vector for a pcolor plot.
:param numpy.array zvect: vector that contains the
`zvect` array without the duplications (obtained with the
:py:func:`pypago.plot._nodouble_tsr` module)
:return: array that contains the `zvect` array formatted for a pcolor plot
:rtype: numpy.array
"""
nz, nl = zvect.shape
zvect = np.ma.masked_where(pypago.misc.extract_mask(zvect), zvect)
# initialisation of a temporary zvect_pcol array
zvect_pcol_temp = []
for indl in xrange(0, nl):
zl = len(np.nonzero(np.ma.getmaskarray(zvect[:, indl]) == 0)[0])
if zl > 1:
temp = np.concatenate(([0],
.5*(zvect[:zl-1, indl] + zvect[1:zl, indl]),
1.5*zvect[zl-1:zl, indl] - 0.5*zvect[zl-2:zl-1, indl],
zvect[zl:nz+1, indl]))
else:
temp = np.concatenate(([0], zvect[0:1, indl],
zvect[1:nz, indl]), axis=0)
zvect_pcol_temp.append(temp)
zvect_pcol_temp = np.array(zvect_pcol_temp)
zvect_pcol_temp = np.ma.masked_where(zvect_pcol_temp != zvect_pcol_temp,
zvect_pcol_temp).T
# creation of the final zvect_pcol array
zvect_pcol = np.zeros((zvect_pcol_temp.shape[0],
zvect_pcol_temp.shape[1]+1))
zvect_pcol[:, 0] = zvect_pcol_temp[:, 0]
for indl in xrange(1, len(zvect_pcol_temp[0, :])):
ind1 = np.nonzero(np.ma.getmaskarray(zvect_pcol_temp[:, indl]) == 0)[0]
ind2 = np.nonzero(np.ma.getmaskarray(zvect_pcol_temp[:, indl-1]) == 0)[0]
zmin = np.min([len(ind1), len(ind2)])
zmax = np.max([len(ind1), len(ind2)])
zvect_pcol[:zmin, indl] = 0.5*(zvect_pcol_temp[:zmin, indl-1] + zvect_pcol_temp[:zmin, indl])
for indz in xrange(zmin, zmax):
if np.ma.getmaskarray(zvect_pcol_temp[indz, indl]):
zvect_pcol[indz, indl] = zvect_pcol_temp[indz, indl-1]
else:
zvect_pcol[indz, indl] = zvect_pcol_temp[indz, indl]
if zmax <= nz:
zvect_pcol[zmax:, indl] = np.nan
zvect_pcol[:, -1] = zvect_pcol[:, -2]
zvect_pcol = np.ma.array(zvect_pcol, mask=np.isnan(zvect_pcol))
return zvect_pcol
[docs]def _make_atracer_pcol(atracer):
"""
Returns the field vector
for a pcolor plot
:param numpy.array atracer: vector that contains the
field array without the duplications (obtained with the
:py:func:`pypago.plot._nodouble_tsr` or
:py:func:`pypago.plot._nodouble_v`)
:return: array that contains the field array formatted for a pcolor plot
:rtype: numpy.array
"""
nz, nl = atracer.shape
atracer_pcol = np.concatenate((atracer, np.nan*np.ones((nz, 1))), axis=1)
atracer_pcol = np.concatenate((atracer_pcol, np.nan*np.ones((1, nl+1))), axis=0)
return atracer_pcol