import xarray
import numpy
import argparse
import sys
from datetime import datetime
import netCDF4
from mpas_tools.io import write_netcdf
[docs]
def compute_depth(refBottomDepth):
"""
Computes depth and depth bounds given refBottomDepth
Parameters
----------
refBottomDepth : xarray.DataArray
the depth of the bottom of each vertical layer in the initial state
(perfect z-level coordinate)
Returns
-------
depth : numpy.ndarray
the vertical coordinate defining the middle of each layer
depth_bnds : numpy.ndarray
the vertical coordinate defining the top and bottom of each layer
"""
# Authors
# -------
# Xylar Asay-Davis
refBottomDepth = refBottomDepth.values
depth_bnds = numpy.zeros((len(refBottomDepth), 2))
depth_bnds[0, 0] = 0.
depth_bnds[1:, 0] = refBottomDepth[0:-1]
depth_bnds[:, 1] = refBottomDepth
depth = 0.5 * (depth_bnds[:, 0] + depth_bnds[:, 1])
return depth, depth_bnds
[docs]
def compute_zmid(bottomDepth, maxLevelCell, layerThickness,
depth_dim='nVertLevels'):
"""
Computes zMid given data arrays for bottomDepth, maxLevelCell and
layerThickness
Parameters
----------
bottomDepth : xarray.DataArray
the depth of the ocean bottom (positive)
maxLevelCell : xarray.DataArray
the 1-based vertical index of the bottom of the ocean
layerThickness : xarray.DataArray
the thickness of MPAS-Ocean layers (possibly as a function of time)
depth_dim : str, optional
the name of the vertical dimension
Returns
-------
zMid : xarray.DataArray
the vertical coordinate defining the middle of each layer, masked below
the bathymetry
"""
# Authors
# -------
# Xylar Asay-Davis
nDepth = layerThickness.sizes[depth_dim]
vertIndex = \
xarray.DataArray.from_dict({'dims': (depth_dim,),
'data': numpy.arange(nDepth)})
layerThickness = layerThickness.where(vertIndex < maxLevelCell)
thicknessSum = layerThickness.sum(dim=depth_dim)
thicknessCumSum = layerThickness.cumsum(dim=depth_dim)
zSurface = -bottomDepth + thicknessSum
zLayerBot = zSurface - thicknessCumSum
zMid = zLayerBot + 0.5 * layerThickness
zMid = zMid.where(vertIndex < maxLevelCell)
if 'Time' in zMid.dims:
zMid = zMid.transpose('Time', 'nCells', depth_dim)
else:
zMid = zMid.transpose('nCells', depth_dim)
return zMid
[docs]
def add_depth(inFileName, outFileName, coordFileName=None):
"""
Add a 1D depth coordinate to an MPAS-Ocean file.
Parameters
----------
inFileName : str
An input MPAS-Ocean file that depth should be added to, used for coords
if another file is not provided via ``coordFileName``.
outFileName : str
An output MPAS-Ocean file with depth added
coordFileName : str, optional
An MPAS-Ocean file with ``refBottomDepth``
"""
if coordFileName is None:
coordFileName = inFileName
ds = xarray.open_dataset(inFileName, mask_and_scale=False)
if 'nVertLevels' in ds.dims:
ds = ds.rename({'nVertLevels': 'depth'})
dsCoord = xarray.open_dataset(coordFileName, mask_and_scale=False)
dsCoord = dsCoord.rename({'nVertLevels': 'depth'})
depth, depth_bnds = compute_depth(dsCoord.refBottomDepth)
ds.coords['depth'] = ('depth', depth)
ds.depth.attrs['long_name'] = 'reference depth of the center of ' \
'each vertical level'
ds.depth.attrs['standard_name'] = 'depth'
ds.depth.attrs['units'] = 'meters'
ds.depth.attrs['axis'] = 'Z'
ds.depth.attrs['positive'] = 'down'
ds.depth.attrs['valid_min'] = depth_bnds[0, 0]
ds.depth.attrs['valid_max'] = depth_bnds[-1, 1]
ds.depth.attrs['bounds'] = 'depth_bnds'
ds.coords['depth_bnds'] = (('depth', 'nbnd'), depth_bnds)
ds.depth_bnds.attrs['long_name'] = 'Gridcell depth interfaces'
for varName in ds.data_vars:
var = ds[varName]
if 'depth' in var.dims:
var = var.assign_coords(depth=ds.depth)
ds[varName] = var
time = datetime.now().strftime('%c')
history = '{}: {}'.format(time, ' '.join(sys.argv))
if 'history' in ds.attrs:
ds.attrs['history'] = f'{history}\n{ds.attrs["history"]}'
else:
ds.attrs['history'] = history
write_netcdf(ds, outFileName)
def main_add_depth():
parser = argparse.ArgumentParser(
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("-c", "--coordFileName", dest="coordFileName",
type=str, required=False,
help="An MPAS-Ocean file with refBottomDepth")
parser.add_argument("-i", "--inFileName", dest="inFileName", type=str,
required=True,
help="An input MPAS-Ocean file that depth should be"
"added to, used for coords if another file is"
"not provided via -c.")
parser.add_argument("-o", "--outFileName", dest="outFileName", type=str,
required=True,
help="An output MPAS-Ocean file with depth added")
args = parser.parse_args()
add_depth(args.inFileName, args.outFileName,
coordFileName=args.coordFileName)
[docs]
def add_zmid(inFileName, outFileName, coordFileName=None):
"""
Add a 3D, time-independent depth coordinate to an MPAS-Ocean file.
Parameters
----------
inFileName : str
An input MPAS-Ocean file that ``zMid`` should be added to, used for
coords if another file is not provided via ``coordFileName``.
outFileName : str
An output MPAS-Ocean file with ``zMid`` added
coordFileName : str, optional
An MPAS-Ocean file with ``bottomDepth``, ``maxLevelCell`` and
``layerThickness`` but not ``zMid``
"""
if coordFileName is None:
coordFileName = inFileName
ds = xarray.open_dataset(inFileName, mask_and_scale=False)
if 'nVertLevels' in ds.dims:
ds = ds.rename({'nVertLevels': 'depth'})
# dsCoord doesn't have masking disabled because we want it for zMid
dsCoord = xarray.open_dataset(coordFileName)
dsCoord = dsCoord.rename({'nVertLevels': 'depth'})
if 'Time' in dsCoord.dims:
dsCoord = dsCoord.isel(Time=0)
ds.coords['zMid'] = compute_zmid(dsCoord.bottomDepth,
dsCoord.maxLevelCell,
dsCoord.layerThickness,
depth_dim='depth')
fillValue = netCDF4.default_fillvals['f8']
ds.coords['zMid'] = ds.zMid.where(ds.zMid.notnull(), other=fillValue)
ds.zMid.attrs['units'] = 'meters'
ds.zMid.attrs['positive'] = 'up'
ds.zMid.attrs['_FillValue'] = fillValue
for varName in ds.data_vars:
var = ds[varName]
if 'nCells' in var.dims and 'depth' in var.dims:
var = var.assign_coords(zMid=ds.zMid)
ds[varName] = var
time = datetime.now().strftime('%c')
history = '{}: {}'.format(time, ' '.join(sys.argv))
if 'history' in ds.attrs:
ds.attrs['history'] = f'{history}\n{ds.attrs["history"]}'
else:
ds.attrs['history'] = history
write_netcdf(ds, outFileName)
def main_add_zmid():
parser = argparse.ArgumentParser(
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("-c", "--coordFileName", dest="coordFileName",
type=str, required=False,
help="An MPAS-Ocean file with bottomDepth, maxLevelCell"
"and layerThickness but not zMid")
parser.add_argument("-i", "--inFileName", dest="inFileName", type=str,
required=True,
help="An input MPAS-Ocean file that zMid should be"
"added to, used for coords if another file is"
"not provided via -c.")
parser.add_argument("-o", "--outFileName", dest="outFileName", type=str,
required=True,
help="An output MPAS-Ocean file with zMid added")
args = parser.parse_args()
add_zmid(args.inFileName, args.outFileName,
coordFileName=args.coordFileName)
[docs]
def write_time_varying_zmid(inFileName, outFileName, coordFileName=None,
prefix=''):
"""
Add a 3D, time-independent depth coordinate to an MPAS-Ocean file.
Parameters
----------
inFileName : str
An input MPAS-Ocean file with some form of ``layerThickness``, and also
``bottomDepth`` and ``maxLevelCell`` if no ``coordFileName``
is provided.
outFileName : str
An output MPAS-Ocean file with ``zMid`` for each time in the input file
coordFileName : str, optional
An MPAS-Ocean file with ``bottomDepth`` and ``maxLevelCell``
prefix : str, optional
A prefix on ``layerThickness`` (in) and ``zMid`` (out), such as
``timeMonthly_avg_``
"""
if coordFileName is None:
coordFileName = inFileName
dsCoord = xarray.open_dataset(coordFileName)
dsCoord = dsCoord.rename({'nVertLevels': 'depth'})
dsIn = xarray.open_dataset(inFileName)
dsIn = dsIn.rename({'nVertLevels': 'depth'})
inVarName = f'{prefix}layerThickness'
outVarName = f'{prefix}zMid'
layerThickness = dsIn[inVarName]
zMid = compute_zmid(dsCoord.bottomDepth, dsCoord.maxLevelCell,
layerThickness, depth_dim='depth')
dsOut = xarray.Dataset()
dsOut[outVarName] = zMid
fillValue = netCDF4.default_fillvals['f8']
dsOut[outVarName] = dsOut[outVarName].where(dsOut[outVarName].notnull(),
other=fillValue)
dsOut[outVarName].attrs['units'] = 'meters'
dsOut[outVarName].attrs['positive'] = 'up'
dsOut[outVarName].attrs['_FillValue'] = fillValue
time = datetime.now().strftime('%c')
history = '{}: {}'.format(time, ' '.join(sys.argv))
dsOut.attrs['history'] = history
write_netcdf(dsOut, outFileName)
def main_write_time_varying_zmid():
parser = argparse.ArgumentParser(
formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("-c", "--coordFileName", dest="coordFileName",
type=str, required=False,
help="An MPAS-Ocean file with bottomDepth and "
"maxLevelCell")
parser.add_argument("-i", "--inFileName", dest="inFileName", type=str,
required=True,
help="An input MPAS-Ocean file with some form of"
"layerThickness, and also bottomDepth and"
"maxLevelCell if no coordinate file is provided.")
parser.add_argument("-o", "--outFileName", dest="outFileName", type=str,
required=True,
help="An output MPAS-Ocean file with zMid for each"
"time in the input file")
parser.add_argument("-p", "--prefix", dest="prefix", type=str,
required=False, default="",
help="A prefix on layerThickness (in) and zMid (out),"
"such as 'timeMonthly_avg_'")
args = parser.parse_args()
write_time_varying_zmid(
args.inFileName, args.outFileName, coordFileName=args.coordFileName,
prefix=args.prefix)