import numpy
import xarray
import os
import copy
import progressbar
import subprocess
import glob
import matplotlib.pyplot as plt
import cmocean
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
[docs]class TimeSeriesPlotter(object):
"""
A plotter object to hold on to some info needed for plotting time series
from ISOMIP+ simulation results
Attributes
----------
inFolder : str
The folder with simulation results
outFolder : str
The folder where images will be written
expt : {'Ocean0', 'Ocean1', 'Ocean2'}
The name of the experiment
"""
[docs] def __init__(self, inFolder='.', outFolder='plots', expt='Ocean0',
dsMesh=None, ds=None):
"""
Create a plotter object to hold on to some info needed for plotting
time series from ISOMIP+ simulation results
Parameters
----------
inFolder : str, optional
The folder with simulation results
outFolder : str, optional
The folder where images will be written
expt : {'Ocean0', 'Ocean1', 'Ocean2'}, optional
The name of the experiment
dsMesh : xarray.Dataset, optional
The MPAS mesh
ds : xarray.Dataset, optional
The time series output
"""
self.inFolder = inFolder
self.outFolder = outFolder
self.expt = expt
if dsMesh is not None:
self.dsMesh = dsMesh
else:
self.dsMesh = xarray.open_dataset(
'{}/init.nc'.format(self.inFolder))
if ds is not None:
self.ds = ds
else:
self.ds = xarray.open_mfdataset(
'{}/timeSeriesStatsMonthly*.nc'.format(self.inFolder),
concat_dim='Time', combine='nested')
try:
os.makedirs(self.outFolder)
except OSError:
pass
plt.switch_backend('Agg')
[docs] def plot_melt_time_series(self, sshMax=None):
"""
Plot a series of image for each of several variables related to melt
at the ice shelf-ocean interface: mean melt rate, total melt flux,
mean thermal driving, mean friction velocity
"""
rho_fw = 1000.
secPerYear = 365*24*60*60
areaCell = self.dsMesh.areaCell
iceMask = self.ds.timeMonthly_avg_landIceFraction
meltFlux = self.ds.timeMonthly_avg_landIceFreshwaterFlux
if sshMax is not None:
ssh = self.ds.timeMonthly_avg_ssh
iceMask = iceMask.where(ssh < sshMax)
totalMeltFlux = (meltFlux*areaCell*iceMask).sum(dim='nCells')
totalArea = (areaCell*iceMask).sum(dim='nCells')
meanMeltRate = totalMeltFlux/totalArea/rho_fw*secPerYear
self.plot_time_series(meanMeltRate, 'mean melt rate', 'meanMeltRate',
'm/yr')
self.plot_time_series(1e-6*totalMeltFlux, 'total melt flux',
'totalMeltFlux', 'kT/yr')
da = (self.ds.timeMonthly_avg_landIceBoundaryLayerTracers_landIceBoundaryLayerTemperature -
self.ds.timeMonthly_avg_landIceInterfaceTracers_landIceInterfaceTemperature)
da = (da*areaCell*iceMask).sum(dim='nCells')/totalArea
self.plot_time_series(da, 'mean thermal driving',
'meanThermalDriving', 'deg C')
da = self.ds.timeMonthly_avg_landIceFrictionVelocity
da = (da*areaCell*iceMask).sum(dim='nCells')/totalArea
self.plot_time_series(da, 'mean friction velocity',
'meanFrictionVelocity', 'm/s')
[docs] def plot_time_series(self, da, nameInTitle, prefix, units=None,
figsize=(12, 6), color=None, overwrite=True):
fileName = '{}/{}.png'.format(self.outFolder, prefix)
if(not overwrite and os.path.exists(fileName)):
return
nTime = da.sizes['Time']
time = numpy.arange(nTime)/12.
plt.figure(1, figsize=figsize)
plt.plot(time, da.values, color=color)
if units is None:
ylabel = nameInTitle
else:
ylabel = '{} ({})'.format(nameInTitle, units)
plt.ylabel(ylabel)
plt.xlabel('time (yrs)')
plt.savefig(fileName)
plt.close()
[docs]class MoviePlotter(object):
"""
A plotter object to hold on to some info needed for plotting images from
ISOMIP+ simulation results
Attributes
----------
inFolder : str
The folder with simulation results
streamfunctionFolder : str
The folder where streamfunction input files were computed
outFolder : str
The folder where images will be written
expt : {'Ocean0', 'Ocean1', 'Ocean2'}
The name of the experiment
sectionY : float
The location along the y axis of a transect in the x-z plane to plot
dsMesh : ``xarray.Dataset``
A data set with mesh data
ds : ``xarray.Dataset``
A data set with the montly-mean simulation results
oceanMask : ``numpy.ndarray``
A mask of cells that are in the ocean domain (probably all ones)
cavityMask : ``numpy.ndarray``
A mask of cells that are in the sub-ice-shelf cavity
oceanPatches : ``PatchCollection``
A set of polygons covering ocean cells
cavityPatches : ``PatchCollection``
A set of polygons covering only cells in the cavity
X, Z : ``numpy.ndarray``
The horiz. and vert. coordinates of the x-z cross section
sectionMask : ``numpy.ndarray``
A mask for the cross section indicating where values are valid (i.e.
above the bathymetry)
showProgress : bool
Whether to show a progressbar
"""
[docs] def __init__(self, inFolder, streamfunctionFolder, outFolder, expt,
sectionY, dsMesh, ds, showProgress):
"""
Create a plotter object to hold on to some info needed for plotting
images from ISOMIP+ simulation results
Parameters
----------
inFolder : str
The folder with simulation results
streamfunctionFolder : str
The folder where streamfunction input files were computed
outFolder : str
The folder where images will be written
expt : {'Ocean0', 'Ocean1', 'Ocean2'}
The name of the experiment
sectionY : float
The location along the y axis of a transect in the x-z plane to
plot
dsMesh : xarray.Dataset
The MPAS mesh
ds : xarray.Dataset
The time series output
showProgress : bool
Whether to show a progressbar
"""
plt.switch_backend('Agg')
self.inFolder = inFolder
self.streamfunctionFolder = streamfunctionFolder
self.outFolder = outFolder
self.expt = expt
self.sectionY = sectionY
self.showProgress = showProgress
self.dsMesh = dsMesh
self.ds = ds
landIceMask = self.dsMesh.landIceFraction.isel(Time=0) > 0.01
self.oceanMask = self.dsMesh.maxLevelCell-1 >= 0
self.cavityMask = numpy.logical_and(self.oceanMask, landIceMask)
self.oceanPatches = _compute_cell_patches(
self.dsMesh, self.oceanMask)
self.cavityPatches = _compute_cell_patches(
self.dsMesh, self.cavityMask)
self.sectionCellIndices = _compute_section_cell_indices(self.sectionY,
self.dsMesh)
self._compute_section_x_z()
[docs] def plot_barotropic_streamfunction(self, vmin=None, vmax=None):
"""
Plot a series of image of the barotropic streamfunction
Parameters
----------
vmin, vmax : float, optional
The minimum and maximum values for the colorbar, defaults are
chosen depending on ``expt``
"""
ds = xarray.open_dataset('{}/barotropicStreamfunction.nc'.format(
self.streamfunctionFolder))
if vmin is None or vmax is None:
if self.expt in ['Ocean0', 'Ocean1']:
vmin = -1
vmax = 1
else:
vmin = -0.5
vmax = 0.5
nTime = ds.sizes['Time']
if self.showProgress:
widgets = ['plotting barotropic streamfunction: ',
progressbar.Percentage(), ' ',
progressbar.Bar(), ' ', progressbar.ETA()]
bar = progressbar.ProgressBar(widgets=widgets,
maxval=nTime).start()
else:
bar = None
for tIndex in range(nTime):
self.update_date(tIndex)
bsf = ds.bsfCell.isel(Time=tIndex)
outFileName = '{}/bsf/bsf_{:04d}.png'.format(
self.outFolder, tIndex+1)
self._plot_horiz_field(bsf, title='barotropic streamfunction (Sv)',
outFileName=outFileName, oceanDomain=True,
vmin=vmin, vmax=vmax, cmap='cmo.curl')
if self.showProgress:
bar.update(tIndex+1)
if self.showProgress:
bar.finish()
[docs] def plot_overturning_streamfunction(self, vmin=-0.3, vmax=0.3):
"""
Plot a series of image of the overturning streamfunction
Parameters
----------
vmin, vmax : float, optional
The minimum and maximum values for the colorbar
"""
ds = xarray.open_dataset('{}/overturningStreamfunction.nc'.format(
self.streamfunctionFolder))
nTime = ds.sizes['Time']
if self.showProgress:
widgets = ['plotting overturning streamfunction: ',
progressbar.Percentage(), ' ',
progressbar.Bar(), ' ', progressbar.ETA()]
bar = progressbar.ProgressBar(widgets=widgets,
maxval=nTime).start()
else:
bar = None
for tIndex in range(nTime):
self.update_date(tIndex)
osf = ds.osf.isel(Time=tIndex)
outFileName = '{}/osf/osf_{:04d}.png'.format(self.outFolder,
tIndex+1)
x = _interp_extrap_corner(ds.x.values)
z = _interp_extrap_corner(ds.z.values)
self._plot_vert_field(
x, z, osf, title='overturning streamfunction (Sv)',
outFileName=outFileName, vmin=vmin, vmax=vmax, cmap='cmo.curl')
if self.showProgress:
bar.update(tIndex+1)
if self.showProgress:
bar.finish()
[docs] def plot_melt_rates(self, vmin=-100., vmax=100.):
"""
Plot a series of image of the melt rate
Parameters
----------
vmin, vmax : float, optional
The minimum and maximum values for the colorbar
"""
rho_fw = 1000.
secPerYear = 365*24*60*60
da = secPerYear/rho_fw*self.ds.timeMonthly_avg_landIceFreshwaterFlux
self.plot_horiz_series(da, 'melt rate', prefix='meltRate',
oceanDomain=False, units='m/yr', vmin=vmin,
vmax=vmax, cmap='cmo.curl')
[docs] def plot_ice_shelf_boundary_variables(self):
"""
Plot a series of image for each of several variables related to the
ice shelf-ocean interface: heat flux from the ocean, heat flux into the
ice, thermal driving, haline driving, and the friction velocity under
ice
"""
self.plot_horiz_series(self.ds.timeMonthly_avg_landIceHeatFlux,
'heat flux from ocean to ice-ocean interface',
prefix='oceanHeatFlux',
oceanDomain=False, units='W/s',
vmin=-1e3, vmax=1e3, cmap='cmo.curl')
self.plot_horiz_series(self.ds.timeMonthly_avg_heatFluxToLandIce,
'heat flux into ice at ice-ocean interface',
prefix='iceHeatFlux',
oceanDomain=False, units='W/s',
vmin=-1e1, vmax=1e1, cmap='cmo.curl')
da = (self.ds.timeMonthly_avg_landIceBoundaryLayerTracers_landIceBoundaryLayerTemperature -
self.ds.timeMonthly_avg_landIceInterfaceTracers_landIceInterfaceTemperature)
self.plot_horiz_series(da, 'thermal driving',
prefix='thermalDriving',
oceanDomain=False, units='deg C',
vmin=-2, vmax=2, cmap='cmo.thermal')
da = (self.ds.timeMonthly_avg_landIceBoundaryLayerTracers_landIceBoundaryLayerSalinity -
self.ds.timeMonthly_avg_landIceInterfaceTracers_landIceInterfaceSalinity)
self.plot_horiz_series(da, 'haline driving',
prefix='halineDriving',
oceanDomain=False, units='PSU',
vmin=-10, vmax=10, cmap='cmo.haline')
self.plot_horiz_series(self.ds.timeMonthly_avg_landIceFrictionVelocity,
'friction velocity',
prefix='frictionVelocity',
oceanDomain=True, units='m/s',
vmin=0, vmax=0.05, cmap='cmo.speed')
[docs] def plot_temperature(self):
"""
Plot a series of images of temperature at the sea surface or
ice-ocean interface, sea floor and in an x-z section
"""
da = self.ds.timeMonthly_avg_activeTracers_temperature
self.plot_3d_field_top_bot_section(da,
nameInTitle='temperature',
prefix='Temp', units='deg C',
vmin=-2.5, vmax=1.0,
cmap='cmo.thermal')
[docs] def plot_salinity(self):
"""
Plot a series of images of salinity at the sea surface or
ice-ocean interface, sea floor and in an x-z section
"""
da = self.ds.timeMonthly_avg_activeTracers_salinity
self.plot_3d_field_top_bot_section(da,
nameInTitle='salinity',
prefix='Salinity', units='PSU',
vmin=33.8, vmax=34.7,
cmap='cmo.haline')
[docs] def plot_potential_density(self):
"""
Plot a series of images of salinity at the sea surface or
ice-ocean interface, sea floor and in an x-z section
"""
da = self.ds.timeMonthly_avg_potentialDensity
self.plot_3d_field_top_bot_section(da,
nameInTitle='potential density',
prefix='PotRho', units='kg/m^3',
vmin=1027., vmax=1028.,
cmap='cmo.dense')
def plot_haney_number(self, haneyFolder=None):
"""
Plot a series of images of the Haney number rx1 at the sea surface or
ice-ocean interface, sea floor and in an x-z section
Parameters
----------
haneyFolder : str, optional
The location of the haney number output, default is ``inFolder``
"""
if haneyFolder is None:
haneyFolder = self.inFolder
ds = xarray.open_dataset('{}/haney.nc'.format(haneyFolder))
self.plot_3d_field_top_bot_section(ds.haneyCell,
nameInTitle='Haney Number (rx1)',
prefix='Haney', units=None,
vmin=0., vmax=8., cmap='cmo.matter')
[docs] def plot_horiz_series(self, da, nameInTitle, prefix, oceanDomain,
units=None, vmin=None, vmax=None, cmap=None):
"""
Plot a series of image of a given variable
Parameters
----------
da : ``xarray.DataArray``
The data array of time series to plot
nameInTitle : str
The name of the variable to use in the title and the progress bar
prefix : str
The nae of the variable to use in the subfolder and file prefix
oceanDomain : bool
True if the variable is for the full ocean, False if only for the
cavity
units : str, optional
The units of the variable to be included in the title
vmin, vmax : float, optional
The minimum and maximum values for the colorbar
cmap : Colormap or str
A color map to plot
"""
nTime = self.ds.sizes['Time']
if self.showProgress:
widgets = ['plotting {}: '.format(nameInTitle),
progressbar.Percentage(), ' ',
progressbar.Bar(), ' ', progressbar.ETA()]
bar = progressbar.ProgressBar(widgets=widgets,
maxval=nTime).start()
else:
bar = None
for tIndex in range(nTime):
self.update_date(tIndex)
field = da.isel(Time=tIndex).values
outFileName = '{}/{}/{}_{:04d}.png'.format(
self.outFolder, prefix, prefix, tIndex+1)
if units is None:
title = nameInTitle
else:
title = '{} ({})'.format(nameInTitle, units)
self._plot_horiz_field(field, title=title, outFileName=outFileName,
oceanDomain=oceanDomain, vmin=vmin,
vmax=vmax, cmap=cmap)
if self.showProgress:
bar.update(tIndex+1)
if self.showProgress:
bar.finish()
[docs] def plot_3d_field_top_bot_section(self, da, nameInTitle, prefix,
units=None, vmin=None, vmax=None,
cmap=None):
"""
Plot a series of images of a given 3D variable showing the value
at the top (sea surface or ice-ocean interface), sea floor and in an
x-z section
Parameters
----------
da : ``xarray.DataArray``
The data array of time series to plot
nameInTitle : str
The name of the variable to use in the title and the progress bar
prefix : str
The nae of the variable to use in the subfolder and file prefix
units : str, optional
The units of the variable to be included in the title
vmin, vmax : float, optional
The minimum and maximum values for the colorbar
cmap : Colormap or str, optional
A color map to plot
"""
if vmin is None:
vmin = da.min()
if vmax is None:
vmax = da.max()
minLevelCell = self.dsMesh.minLevelCell-1
daTop = xarray.DataArray(da)
daTop.coords['verticalIndex'] = \
('nVertLevels',
numpy.arange(daTop.sizes['nVertLevels']))
# mask only the values with the right vertical index
daTop = daTop.where(daTop.verticalIndex == minLevelCell)
# Each vertical layer has at most one non-NaN value so the "sum"
# over the vertical is used to collapse the array in the vertical
# dimension
daTop = daTop.sum(dim='nVertLevels').where(minLevelCell >= 0)
self.plot_horiz_series(daTop,
'top {}'.format(nameInTitle),
'top{}'.format(prefix), oceanDomain=True,
vmin=vmin, vmax=vmax, cmap=cmap)
maxLevelCell = self.dsMesh.maxLevelCell-1
daBot = xarray.DataArray(da)
daBot.coords['verticalIndex'] = \
('nVertLevels',
numpy.arange(daBot.sizes['nVertLevels']))
# mask only the values with the right vertical index
daBot = daBot.where(daBot.verticalIndex == maxLevelCell)
# Each vertical layer has at most one non-NaN value so the "sum"
# over the vertical is used to collapse the array in the vertical
# dimension
daBot = daBot.sum(dim='nVertLevels').where(maxLevelCell >= 0)
self.plot_horiz_series(daBot,
'bot {}'.format(nameInTitle),
'bot{}'.format(prefix), oceanDomain=True,
vmin=vmin, vmax=vmax, cmap=cmap)
daSection = da.isel(nCells=self.sectionCellIndices)
nTime = self.ds.sizes['Time']
if self.showProgress:
widgets = ['plotting {} section: '.format(nameInTitle),
progressbar.Percentage(), ' ',
progressbar.Bar(), ' ', progressbar.ETA()]
bar = progressbar.ProgressBar(widgets=widgets,
maxval=nTime).start()
else:
bar = None
for tIndex in range(nTime):
self.update_date(tIndex)
mask = numpy.logical_not(self.sectionMask)
field = numpy.ma.masked_array(daSection.isel(Time=tIndex).values.T,
mask=mask)
outFileName = '{}/section{}/section{}_{:04d}.png'.format(
self.outFolder, prefix, prefix, tIndex+1)
if units is None:
title = nameInTitle
else:
title = '{} ({}) along section at y={:g} km'.format(
nameInTitle, units, 1e-3*self.sectionY)
self._plot_vert_field(self.X, self.Z[tIndex, :, :],
field, title=title,
outFileName=outFileName,
vmin=vmin, vmax=vmax, cmap=cmap)
if self.showProgress:
bar.update(tIndex+1)
if self.showProgress:
bar.finish()
[docs] def plot_layer_interfaces(self, figsize=(9, 5)):
"""
Plot layer interfaces, the sea surface height and the bottom topography
of the cross section at fixed y.
Parameters
----------
figsize : tuple, optional
The size of the figure
"""
nTime = self.Z.shape[0]
if self.showProgress:
widgets = ['plotting section of layer interfaces: ',
progressbar.Percentage(), ' ',
progressbar.Bar(), ' ', progressbar.ETA()]
bar = progressbar.ProgressBar(widgets=widgets,
maxval=nTime).start()
else:
bar = None
z_mask = numpy.ones(self.X.shape)
z_mask[0:-1, 0:-1] *= numpy.where(self.sectionMask, 1., numpy.nan)
z_mask[1:, 0:-1] *= numpy.where(self.sectionMask, 1., numpy.nan)
z_mask[0:-1, 1:] *= numpy.where(self.sectionMask, 1., numpy.nan)
z_mask[1:, 1:] *= numpy.where(self.sectionMask, 1., numpy.nan)
for tIndex in range(nTime):
Z = numpy.array(self.Z[tIndex, :, :])
ylim = [numpy.amin(Z), 20]
Z *= z_mask
X = self.X
self.update_date(tIndex)
outFileName = '{}/layers/layers_{:04d}.png'.format(self.outFolder,
tIndex+1)
if os.path.exists(outFileName):
continue
try:
os.makedirs(os.path.dirname(outFileName))
except OSError:
pass
plt.figure(figsize=figsize)
ax = plt.subplot(111)
for z_index in range(1, X.shape[0]):
plt.plot(1e-3 * X[z_index, :], Z[z_index, :], 'k')
plt.plot(1e-3 * X[0, :], Z[0, :], 'g')
plt.plot(1e-3 * X[0, :], self.zBotSection, 'g')
ax.autoscale(tight=True)
plt.ylim(ylim)
plt.title('{} {}'.format('layer interfaces', self.date))
plt.tight_layout(pad=0.5)
plt.savefig(outFileName)
plt.close()
if self.showProgress:
bar.update(tIndex+1)
if self.showProgress:
bar.finish()
[docs] def images_to_movies(self, outFolder, framesPerSecond=30, extension='mp4',
overwrite=True):
"""
Convert all the image sequences into movies with ffmpeg
"""
try:
os.makedirs('{}/logs'.format(outFolder))
except OSError:
pass
framesPerSecond = '{}'.format(framesPerSecond)
for fileName in sorted(glob.glob(
'{}/*/*0001.png'.format(self.outFolder))):
prefix = os.path.basename(fileName)[:-9]
outFileName = '{}/{}.{}'.format(outFolder, prefix, extension)
if not overwrite and os.path.exists(outFileName):
continue
imageFileTemplate = '{}/{}/{}_%04d.png'.format(self.outFolder,
prefix, prefix)
logFileName = '{}/logs/{}.log'.format(outFolder, prefix)
with open(logFileName, 'w') as logFile:
args = ['ffmpeg', '-y', '-r', framesPerSecond,
'-i', imageFileTemplate, '-b:v', '32000k',
'-r', framesPerSecond, '-pix_fmt', 'yuv420p',
outFileName]
print('running {}'.format(' '.join(args)))
subprocess.check_call(args, stdout=logFile, stderr=logFile)
def update_date(self, tIndex):
if 'xtime_startMonthly' in self.ds:
var = 'xtime_startMonthly'
elif 'xtime' in self.ds:
var = 'xtime'
else:
self.date = ''
return
xtime = self.ds[var].isel(Time=tIndex).values
xtime = ''.join(str(xtime.astype('U'))).strip()
year = xtime[0:4]
month = xtime[5:7]
self.date = '{}-{}'.format(year, month)
def _plot_horiz_field(self, field, title, outFileName, oceanDomain=True,
vmin=None, vmax=None, figsize=(9, 3), cmap=None):
try:
os.makedirs(os.path.dirname(outFileName))
except OSError:
pass
if os.path.exists(outFileName):
return
if oceanDomain:
localPatches = copy.copy(self.oceanPatches)
localPatches.set_array(field[self.oceanMask])
else:
localPatches = copy.copy(self.cavityPatches)
localPatches.set_array(field[self.cavityMask])
localPatches.set_edgecolor('face')
localPatches.set_clim(vmin=vmin, vmax=vmax)
if cmap is not None:
localPatches.set_cmap(cmap)
plt.figure(figsize=figsize)
ax = plt.subplot(111)
ax.add_collection(localPatches)
plt.colorbar(localPatches)
plt.axis([0, 500, 0, 1000])
ax.set_aspect('equal')
ax.autoscale(tight=True)
plt.title('{} {}'.format(title, self.date))
plt.tight_layout(pad=0.5)
plt.savefig(outFileName)
plt.close()
def _plot_vert_field(self, inX, inZ, field, title, outFileName, vmin=None,
vmax=None, figsize=(9, 5), cmap=None):
try:
os.makedirs(os.path.dirname(outFileName))
except OSError:
pass
if os.path.exists(outFileName):
return
plt.figure(figsize=figsize)
ax = plt.subplot(111)
plt.pcolormesh(1e-3*inX, inZ, field, vmin=vmin, vmax=vmax, cmap=cmap)
plt.colorbar()
ax.autoscale(tight=True)
plt.ylim([numpy.amin(inZ), 20])
plt.title('{} {}'.format(title, self.date))
plt.tight_layout(pad=0.5)
plt.savefig(outFileName)
plt.close()
def _compute_section_x_z(self):
x = _interp_extrap_corner(self.dsMesh.xCell[self.sectionCellIndices])
nx = len(x)
nVertLevels = self.dsMesh.sizes['nVertLevels']
nTime = self.ds.sizes['Time']
self.X = numpy.zeros((nVertLevels+1, nx))
for zIndex in range(nVertLevels+1):
self.X[zIndex, :] = x
self.sectionMask = numpy.zeros((nVertLevels, nx-1), dtype=bool)
for zIndex in range(nVertLevels):
minLevelCell = self.dsMesh.minLevelCell.isel(
nCells=self.sectionCellIndices) - 1
maxLevelCell = self.dsMesh.maxLevelCell.isel(
nCells=self.sectionCellIndices) - 1
self.sectionMask[zIndex, :] = numpy.logical_and(
zIndex >= minLevelCell, zIndex <= maxLevelCell)
self.Z = numpy.zeros((nTime, nVertLevels+1, nx))
self.zBotSection = -_interp_extrap_corner(
self.dsMesh.bottomDepth.isel(
nCells=self.sectionCellIndices).values)
for tIndex in range(nTime):
if 'timeMonthly_avg_layerThickness' in self.ds:
var = 'timeMonthly_avg_layerThickness'
else:
var = 'layerThickness'
layerThickness = self.ds[var].isel(
Time=tIndex, nCells=self.sectionCellIndices)
layerThickness = layerThickness.values*self.sectionMask.T
layerThickness = numpy.nan_to_num(layerThickness)
self.Z[tIndex, -1, :] = self.zBotSection
for zIndex in range(nVertLevels-1, -1, -1):
layerThicknessSection = _interp_extrap_corner(
layerThickness[:, zIndex])
self.Z[tIndex, zIndex, :] = self.Z[tIndex, zIndex+1, :] + \
layerThicknessSection
def _compute_cell_patches(dsMesh, mask):
patches = []
nVerticesOnCell = dsMesh.nEdgesOnCell.values
verticesOnCell = dsMesh.verticesOnCell.values - 1
xVertex = dsMesh.xVertex.values
yVertex = dsMesh.yVertex.values
for iCell in range(dsMesh.sizes['nCells']):
if(not mask[iCell]):
continue
nVert = nVerticesOnCell[iCell]
vertexIndices = verticesOnCell[iCell, :nVert]
vertices = numpy.zeros((nVert, 2))
vertices[:, 0] = 1e-3*xVertex[vertexIndices]
vertices[:, 1] = 1e-3*yVertex[vertexIndices]
polygon = Polygon(vertices, True)
patches.append(polygon)
p = PatchCollection(patches, alpha=1.)
return p
def _compute_section_cell_indices(y, dsMesh):
xCell = dsMesh.xCell.values
yCell = dsMesh.yCell.values
xMin = numpy.amin(xCell)
xMax = numpy.amax(xCell)
xs = numpy.linspace(xMin, xMax, 10000)
cellIndices = []
for x in xs:
distanceSquared = (x - xCell)**2 + (y-yCell)**2
index = numpy.argmin(distanceSquared)
if len(cellIndices) == 0 or cellIndices[-1] != index:
cellIndices.append(index)
return numpy.array(cellIndices)
def _interp_extrap_corner(inField):
"""Interpolate/extrapolate a 1D field from grid centers to grid corners"""
outField = numpy.zeros(len(inField) + 1)
outField[1:-1] = 0.5 * (inField[0:-1] + inField[1:])
# extrapolate the ends
outField[0] = 1.5 * inField[0] - 0.5 * inField[1]
outField[-1] = 1.5 * inField[-1] - 0.5 * inField[-2]
return outField