import xarray
import numpy
from compass.ocean.vertical.grid_1d import add_1d_grid
from compass.ocean.vertical.partial_cells import alter_bottom_depth
from compass.ocean.vertical.zlevel import compute_z_level_layer_thickness, \
compute_min_max_level_cell
[docs]
def init_z_star_vertical_coord(config, ds):
"""
Create a z-star vertical coordinate based on the config options in the
``vertical_grid`` section and the ``bottomDepth`` and ``ssh`` variables of
the mesh data set.
The following new variables will be added to the data set:
* ``minLevelCell`` - the index of the top valid layer
* ``maxLevelCell`` - the index of the bottom valid layer
* ``cellMask`` - a mask of where cells are valid
* ``layerThickness`` - the thickness of each layer
* ``restingThickness`` - the thickness of each layer stretched as if
``ssh = 0``
* ``zMid`` - the elevation of the midpoint of each layer
So far, all supported coordinates make use of a 1D reference vertical grid.
The following variables associated with that field are also added to the
mesh:
* ``refTopDepth`` - the positive-down depth of the top of each ref. level
* ``refZMid`` - the positive-down depth of the middle of each ref. level
* ``refBottomDepth`` - the positive-down depth of the bottom of each ref.
level
* ``refInterfaces`` - the positive-down depth of the interfaces between
ref. levels (with ``nVertLevels`` + 1 elements).
* ``vertCoordMovementWeights`` - the weights (all ones) for coordinate
movement
There is considerable redundancy between these variables but each is
sometimes convenient.
Parameters
----------
config : compass.config.CompassConfigParser
Configuration options with parameters used to construct the vertical
grid
ds : xarray.Dataset
A data set containing ``bottomDepth`` and ``ssh`` variables used to
construct the vertical coordinate
"""
add_1d_grid(config, ds)
ds['vertCoordMovementWeights'] = xarray.ones_like(ds.refBottomDepth)
restingSSH = xarray.zeros_like(ds.bottomDepth)
ds['minLevelCell'], ds['maxLevelCell'], ds['cellMask'] = \
compute_min_max_level_cell(ds.refTopDepth, ds.refBottomDepth,
restingSSH, ds.bottomDepth)
ds['bottomDepth'], ds['maxLevelCell'] = alter_bottom_depth(
config, ds.bottomDepth, ds.refBottomDepth, ds.maxLevelCell)
ds['restingThickness'] = compute_z_level_layer_thickness(
ds.refTopDepth, ds.refBottomDepth, restingSSH, ds.bottomDepth,
ds.minLevelCell, ds.maxLevelCell)
ds['layerThickness'] = _compute_z_star_layer_thickness(
ds.restingThickness, ds.ssh, ds.bottomDepth, ds.minLevelCell,
ds.maxLevelCell)
def _compute_z_star_layer_thickness(restingThickness, ssh, bottomDepth,
minLevelCell, maxLevelCell):
"""
Compute z-star layer thickness by stretching restingThickness based on ssh
and bottomDepth
Parameters
----------
restingThickness : xarray.DataArray
The thickness of z-star layers when ssh = 0
ssh : xarray.DataArray
The sea surface height
bottomDepth : xarray.DataArray
The positive-down depth of the seafloor
minLevelCell : xarray.DataArray
The zero-based index of the top valid level
maxLevelCell : xarray.DataArray
The zero-based index of the bottom valid level
Returns
-------
layerThickness : xarray.DataArray
The thickness of each layer (level)
"""
nVertLevels = restingThickness.sizes['nVertLevels']
layerThickness = []
layerStretch = (ssh + bottomDepth) / bottomDepth
for zIndex in range(nVertLevels):
mask = numpy.logical_and(zIndex >= minLevelCell,
zIndex <= maxLevelCell)
thickness = layerStretch*restingThickness.isel(nVertLevels=zIndex)
thickness = thickness.where(mask, 0.)
layerThickness.append(thickness)
layerThickness = xarray.concat(layerThickness, dim='nVertLevels')
layerThickness = layerThickness.transpose('nCells', 'nVertLevels')
return layerThickness