import numpy
import xarray
from compass.ocean.vertical.zlevel import init_z_level_vertical_coord
from compass.ocean.vertical.zstar import init_z_star_vertical_coord
[docs]def init_vertical_coord(config, ds):
"""
Create a 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
"""
for var in ['bottomDepth', 'ssh']:
if var not in ds:
raise ValueError('{} must be added to ds before this call.'.format(
var))
if 'Time' in ds.ssh.dims:
# drop it for now, we'll add it back at the end
ds['ssh'] = ds.ssh.isel(Time=0)
coord_type = config.get('vertical_grid', 'coord_type')
if coord_type == 'z-level':
init_z_level_vertical_coord(config, ds)
elif coord_type == 'z-star':
init_z_star_vertical_coord(config, ds)
elif coord_type == 'haney-number':
raise ValueError('Haney Number coordinate not yet supported.')
else:
raise ValueError('Unknown coordinate type {}'.format(coord_type))
# recompute the cell mask since min/max indices may have changed
ds['cellMask'] = _compute_cell_mask(ds.minLevelCell, ds.maxLevelCell,
ds.sizes['nVertLevels'])
# mask layerThickness and restingThickness
ds['layerThickness'] = ds.layerThickness.where(ds.cellMask)
ds['restingThickness'] = ds.restingThickness.where(ds.cellMask)
# add (back) Time dimension
ds['ssh'] = ds.ssh.expand_dims(dim='Time', axis=0)
ds['layerThickness'] = ds.layerThickness.expand_dims(dim='Time', axis=0)
ds['restingThickness'] = \
ds.restingThickness.expand_dims(dim='Time', axis=0)
ds['zMid'] = _compute_zmid_from_layer_thickness(
ds.layerThickness, ds.ssh, ds.cellMask)
# fortran 1-based indexing
ds['minLevelCell'] = ds.minLevelCell+1
ds['maxLevelCell'] = ds.maxLevelCell+1
def _compute_cell_mask(minLevelCell, maxLevelCell, nVertLevels):
cellMask = []
for zIndex in range(nVertLevels):
mask = numpy.logical_and(zIndex >= minLevelCell,
zIndex <= maxLevelCell)
cellMask.append(mask)
cellMask = xarray.concat(cellMask, dim='nVertLevels')
cellMask = cellMask.transpose('nCells', 'nVertLevels')
return cellMask
def _compute_zmid_from_layer_thickness(layerThickness, ssh, cellMask):
"""
Compute zMid from ssh and layerThickness for any vertical coordinate
Parameters
----------
layerThickness : xarray.DataArray
The thickness of each layer
ssh : xarray.DataArray
The sea surface height
cellMask : xarray.DataArray
A boolean mask of where there are valid cells
Returns
-------
zMid : xarray.DataArray
The elevation of layer centers
"""
zTop = ssh.copy()
nVertLevels = layerThickness.sizes['nVertLevels']
zMid = []
for zIndex in range(nVertLevels):
mask = cellMask.isel(nVertLevels=zIndex)
thickness = layerThickness.isel(nVertLevels=zIndex).where(mask, 0.)
z = (zTop - 0.5*thickness).where(mask)
zMid.append(z)
zTop -= thickness
zMid = xarray.concat(zMid, dim='nVertLevels').transpose('Time', 'nCells',
'nVertLevels')
return zMid