From f48e8b2d09174a6cb50828a8d7983ebd509943dc Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Fri, 26 Jun 2026 12:36:03 +0200 Subject: [PATCH 1/5] Satisfy pre-commit on files touched by Omega changes ruff-format was added late in development and not all files were updated. Reformat depth.py, barotropic.py and plot.py to satisfy the ruff-format, flynt and import-sorting hooks, and add an explicit strict=True to the existing zip() calls flagged by ruff (B905), so the subsequent Omega changes land on a pre-commit-clean base. Co-Authored-By: Claude Opus 4.8 --- conda_package/mpas_tools/ocean/depth.py | 213 ++++++++++++------ .../ocean/streamfunction/barotropic.py | 6 +- .../mpas_tools/ocean/viz/transect/plot.py | 2 +- 3 files changed, 146 insertions(+), 75 deletions(-) diff --git a/conda_package/mpas_tools/ocean/depth.py b/conda_package/mpas_tools/ocean/depth.py index 7d00b7b68..12e2634d8 100644 --- a/conda_package/mpas_tools/ocean/depth.py +++ b/conda_package/mpas_tools/ocean/depth.py @@ -1,9 +1,10 @@ -import xarray -import numpy import argparse import sys from datetime import datetime + import netCDF4 +import numpy +import xarray from mpas_tools.io import write_netcdf @@ -34,7 +35,7 @@ def compute_depth(refBottomDepth): depth_bnds = numpy.zeros((len(refBottomDepth), 2)) - depth_bnds[0, 0] = 0. + depth_bnds[0, 0] = 0.0 depth_bnds[1:, 0] = refBottomDepth[0:-1] depth_bnds[:, 1] = refBottomDepth depth = 0.5 * (depth_bnds[:, 0] + depth_bnds[:, 1]) @@ -42,8 +43,9 @@ def compute_depth(refBottomDepth): return depth, depth_bnds -def compute_zmid(bottomDepth, maxLevelCell, layerThickness, - depth_dim='nVertLevels'): +def compute_zmid( + bottomDepth, maxLevelCell, layerThickness, depth_dim='nVertLevels' +): """ Computes zMid given data arrays for bottomDepth, maxLevelCell and layerThickness @@ -74,9 +76,9 @@ def compute_zmid(bottomDepth, maxLevelCell, layerThickness, nDepth = layerThickness.sizes[depth_dim] - vertIndex = \ - xarray.DataArray.from_dict({'dims': (depth_dim,), - 'data': numpy.arange(nDepth)}) + vertIndex = xarray.DataArray.from_dict( + {'dims': (depth_dim,), 'data': numpy.arange(nDepth)} + ) layerThickness = layerThickness.where(vertIndex < maxLevelCell) @@ -126,8 +128,9 @@ def add_depth(inFileName, outFileName, coordFileName=None): 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['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' @@ -147,7 +150,7 @@ def add_depth(inFileName, outFileName, coordFileName=None): time = datetime.now().strftime('%c') - history = '{}: {}'.format(time, ' '.join(sys.argv)) + history = f'{time}: {" ".join(sys.argv)}' if 'history' in ds.attrs: ds.attrs['history'] = f'{history}\n{ds.attrs["history"]}' @@ -159,22 +162,39 @@ def add_depth(inFileName, outFileName, coordFileName=None): 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") + 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) + add_depth( + args.inFileName, args.outFileName, coordFileName=args.coordFileName + ) def add_zmid(inFileName, outFileName, coordFileName=None): @@ -207,10 +227,12 @@ def add_zmid(inFileName, outFileName, coordFileName=None): if 'Time' in dsCoord.dims: dsCoord = dsCoord.isel(Time=0) - ds.coords['zMid'] = compute_zmid(dsCoord.bottomDepth, - dsCoord.maxLevelCell, - dsCoord.layerThickness, - depth_dim='depth') + 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' @@ -225,7 +247,7 @@ def add_zmid(inFileName, outFileName, coordFileName=None): time = datetime.now().strftime('%c') - history = '{}: {}'.format(time, ' '.join(sys.argv)) + history = f'{time}: {" ".join(sys.argv)}' if 'history' in ds.attrs: ds.attrs['history'] = f'{history}\n{ds.attrs["history"]}' @@ -237,27 +259,45 @@ def add_zmid(inFileName, outFileName, coordFileName=None): 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") + 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) + add_zmid( + args.inFileName, args.outFileName, coordFileName=args.coordFileName + ) -def write_time_varying_zmid(inFileName, outFileName, coordFileName=None, - prefix=''): +def write_time_varying_zmid( + inFileName, outFileName, coordFileName=None, prefix='' +): """ Add a 3D, time-independent depth coordinate to an MPAS-Ocean file. @@ -290,21 +330,26 @@ def write_time_varying_zmid(inFileName, outFileName, coordFileName=None, outVarName = f'{prefix}zMid' layerThickness = dsIn[inVarName] - zMid = compute_zmid(dsCoord.bottomDepth, dsCoord.maxLevelCell, - layerThickness, depth_dim='depth') + 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] = 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)) + history = f'{time}: {" ".join(sys.argv)}' dsOut.attrs['history'] = history write_netcdf(dsOut, outFileName) @@ -312,26 +357,50 @@ def write_time_varying_zmid(inFileName, outFileName, coordFileName=None, 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_'") + 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) + args.inFileName, + args.outFileName, + coordFileName=args.coordFileName, + prefix=args.prefix, + ) diff --git a/conda_package/mpas_tools/ocean/streamfunction/barotropic.py b/conda_package/mpas_tools/ocean/streamfunction/barotropic.py index a0e0c722b..09be36bc0 100644 --- a/conda_package/mpas_tools/ocean/streamfunction/barotropic.py +++ b/conda_package/mpas_tools/ocean/streamfunction/barotropic.py @@ -202,7 +202,9 @@ def _build_minimal_boundary_constraints( """ # Create a graph from the boundary edges graph = nx.Graph() - edges = list(zip(boundary_vertex0.values, boundary_vertex1.values)) + edges = list( + zip(boundary_vertex0.values, boundary_vertex1.values, strict=True) + ) graph.add_edges_from(edges) minimal_constraints = [] @@ -264,7 +266,7 @@ def _identify_boundary_vertices(ds_mesh, logger, all_vertices): ) # Unpack minimal_constraints into boundary_vertex0 and boundary_vertex1 - boundary_vertex0, boundary_vertex1 = zip(*minimal_constraints) + boundary_vertex0, boundary_vertex1 = zip(*minimal_constraints, strict=True) boundary_vertex0 = xr.DataArray( np.array(boundary_vertex0), dims=('nVertices',) ) diff --git a/conda_package/mpas_tools/ocean/viz/transect/plot.py b/conda_package/mpas_tools/ocean/viz/transect/plot.py index 296beb1ee..61de48677 100644 --- a/conda_package/mpas_tools/ocean/viz/transect/plot.py +++ b/conda_package/mpas_tools/ocean/viz/transect/plot.py @@ -511,7 +511,7 @@ def _compute_feature_transects(fc, ds_mesh, flip): assert transect['geometry']['type'] == 'LineString' coordinates = transect['geometry']['coordinates'] - transect_lon, transect_lat = zip(*coordinates) + transect_lon, transect_lat = zip(*coordinates, strict=True) transect_lon = np.array(transect_lon) transect_lat = np.array(transect_lat) if flip: From a89db9358a916c7e12c98095cdfafc6882059cf3 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Fri, 26 Jun 2026 12:36:37 +0200 Subject: [PATCH 2/5] Support Omega split mesh and vertical coordinate datasets Omega stores the vertical coordinate variables (minLevelCell, maxLevelCell and bottomDepth) in a dataset separate from the horizontal mesh. Add an optional ds_vert_coord argument (the same as ds_mesh by default) to compute_vertically_integrated_velocity, compute_barotropic_streamfunction and plot_feature_transects, and a corresponding --vert_coord option to the plot_feature_transects command, so these variables can be read from a separate dataset. layerThickness is a dynamic field, so it is now always read from the dynamic dataset rather than the mesh, and add_zmid reads it from its input file rather than the coordinate file. This is not backwards compatible. add_depth and compute_depth rely on the 1D refBottomDepth, which Omega and other general vertical coordinates do not provide. They now raise a clear error when refBottomDepth is missing and are documented as not supported for Omega, pointing users to the 3D zMid coordinate instead. Co-Authored-By: Claude Opus 4.8 --- conda_package/mpas_tools/ocean/depth.py | 44 ++++++++++-- .../ocean/streamfunction/barotropic.py | 12 +++- .../ocean/streamfunction/velocity.py | 24 +++++-- .../mpas_tools/ocean/viz/transect/plot.py | 69 ++++++++++++++----- 4 files changed, 120 insertions(+), 29 deletions(-) diff --git a/conda_package/mpas_tools/ocean/depth.py b/conda_package/mpas_tools/ocean/depth.py index 12e2634d8..0932b2e73 100644 --- a/conda_package/mpas_tools/ocean/depth.py +++ b/conda_package/mpas_tools/ocean/depth.py @@ -13,6 +13,14 @@ def compute_depth(refBottomDepth): """ Computes depth and depth bounds given refBottomDepth + .. note:: + This produces a 1D depth coordinate from ``refBottomDepth``, a single + reference depth profile that is independent of horizontal location. + This assumes a z-level/z-star vertical coordinate and is **not + supported for Omega or other general vertical coordinates**, which do + not provide ``refBottomDepth``. Use :py:func:`add_zmid` for a 3D + ``zMid`` coordinate instead. + Parameters ---------- refBottomDepth : xarray.DataArray @@ -103,6 +111,13 @@ def add_depth(inFileName, outFileName, coordFileName=None): """ Add a 1D depth coordinate to an MPAS-Ocean file. + .. note:: + This adds a 1D ``depth`` coordinate derived from ``refBottomDepth``, + which assumes a z-level/z-star vertical coordinate. It is **not + supported for Omega or other general vertical coordinates**, which do + not provide ``refBottomDepth``. Use :py:func:`add_zmid` for a 3D + ``zMid`` coordinate instead. + Parameters ---------- inFileName : str @@ -126,6 +141,15 @@ def add_depth(inFileName, outFileName, coordFileName=None): dsCoord = xarray.open_dataset(coordFileName, mask_and_scale=False) dsCoord = dsCoord.rename({'nVertLevels': 'depth'}) + if 'refBottomDepth' not in dsCoord: + raise ValueError( + 'add_depth() requires a 1D refBottomDepth variable, which ' + 'defines a z-level/z-star reference depth for each layer. ' + 'This is not available in Omega or for general vertical ' + 'coordinates. Use add_zmid() to add a 3D zMid coordinate ' + 'instead.' + ) + depth, depth_bnds = compute_depth(dsCoord.refBottomDepth) ds.coords['depth'] = ('depth', depth) ds.depth.attrs['long_name'] = ( @@ -211,8 +235,9 @@ def add_zmid(inFileName, outFileName, coordFileName=None): An output MPAS-Ocean file with ``zMid`` added coordFileName : str, optional - An MPAS-Ocean file with ``bottomDepth``, ``maxLevelCell`` and - ``layerThickness`` but not ``zMid`` + An MPAS-Ocean file with ``bottomDepth`` and ``maxLevelCell`` but not + ``zMid``. ``layerThickness`` is a dynamic field and is always read + from ``inFileName``. """ if coordFileName is None: coordFileName = inFileName @@ -221,6 +246,14 @@ def add_zmid(inFileName, outFileName, coordFileName=None): if 'nVertLevels' in ds.dims: ds = ds.rename({'nVertLevels': 'depth'}) + # dsIn doesn't have masking disabled because we want NaNs below the + # bathymetry for zMid; layerThickness is a dynamic field so it comes + # from the input file, not the (static) coordinate file + dsIn = xarray.open_dataset(inFileName) + dsIn = dsIn.rename({'nVertLevels': 'depth'}) + if 'Time' in dsIn.dims: + dsIn = dsIn.isel(Time=0) + # dsCoord doesn't have masking disabled because we want it for zMid dsCoord = xarray.open_dataset(coordFileName) dsCoord = dsCoord.rename({'nVertLevels': 'depth'}) @@ -230,7 +263,7 @@ def add_zmid(inFileName, outFileName, coordFileName=None): ds.coords['zMid'] = compute_zmid( dsCoord.bottomDepth, dsCoord.maxLevelCell, - dsCoord.layerThickness, + dsIn.layerThickness, depth_dim='depth', ) fillValue = netCDF4.default_fillvals['f8'] @@ -267,8 +300,9 @@ def main_add_zmid(): dest='coordFileName', type=str, required=False, - help='An MPAS-Ocean file with bottomDepth, ' - 'maxLevelCell and layerThickness but not zMid', + help='An MPAS-Ocean file with bottomDepth and ' + 'maxLevelCell but not zMid (layerThickness is ' + 'always read from the input file)', ) parser.add_argument( '-i', diff --git a/conda_package/mpas_tools/ocean/streamfunction/barotropic.py b/conda_package/mpas_tools/ocean/streamfunction/barotropic.py index 09be36bc0..2e18f7d7a 100644 --- a/conda_package/mpas_tools/ocean/streamfunction/barotropic.py +++ b/conda_package/mpas_tools/ocean/streamfunction/barotropic.py @@ -20,6 +20,7 @@ def compute_barotropic_streamfunction( ds_mesh, ds, + ds_vert_coord=None, logger=None, min_depth=None, max_depth=None, @@ -37,12 +38,18 @@ def compute_barotropic_streamfunction( Parameters ---------- ds_mesh : xarray.Dataset - A dataset containing MPAS mesh variables + A dataset containing MPAS horizontal mesh variables ds : xarray.Dataset A dataset containing MPAS output variables ``normalVelocity`` and ``layerThickness`` (possibly with a ``prefix``) + ds_vert_coord : xarray.Dataset, optional + A dataset with the vertical coordinate variables ``minLevelCell``, + ``maxLevelCell`` and ``bottomDepth``. The same as ``ds_mesh`` by + default, but in Omega these are stored separately from the horizontal + mesh. + logger : logging.Logger, optional A logger for the output if not stdout @@ -106,6 +113,7 @@ def compute_barotropic_streamfunction( bsf_vertex = _compute_barotropic_streamfunction_vertex( ds_mesh, ds, + ds_vert_coord, prefix, include_bolus, include_submesoscale, @@ -420,6 +428,7 @@ def _assemble_matrix( def _compute_barotropic_streamfunction_vertex( ds_mesh, ds, + ds_vert_coord, prefix, include_bolus, include_submesoscale, @@ -442,6 +451,7 @@ def _compute_barotropic_streamfunction_vertex( vert_integ_velocity = compute_vertically_integrated_velocity( ds_mesh=ds_mesh, ds=ds, + ds_vert_coord=ds_vert_coord, logger=logger, min_depth=min_depth, max_depth=max_depth, diff --git a/conda_package/mpas_tools/ocean/streamfunction/velocity.py b/conda_package/mpas_tools/ocean/streamfunction/velocity.py index 114d59db9..e61d340cb 100644 --- a/conda_package/mpas_tools/ocean/streamfunction/velocity.py +++ b/conda_package/mpas_tools/ocean/streamfunction/velocity.py @@ -7,6 +7,7 @@ def compute_vertically_integrated_velocity( ds_mesh, ds, + ds_vert_coord=None, logger=None, min_depth=None, max_depth=None, @@ -21,12 +22,18 @@ def compute_vertically_integrated_velocity( Parameters ---------- ds_mesh : xarray.Dataset - A dataset containing MPAS mesh variables + A dataset containing MPAS horizontal mesh variables ds : xarray.Dataset A dataset containing MPAS output variables ``normalVelocity`` and ``layerThickness`` among others, possibly with a ``prefix`` + ds_vert_coord : xarray.Dataset, optional + A dataset with the vertical coordinate variables ``minLevelCell``, + ``maxLevelCell`` and ``bottomDepth``. The same as ``ds_mesh`` by + default, but in Omega these are stored separately from the horizontal + mesh. + logger : logging.Logger, optional A logger for the output if not stdout @@ -55,6 +62,9 @@ def compute_vertically_integrated_velocity( vert_integ_velocity : xarray.DataArray The vertically integrated velocity on the mesh edges """ + if ds_vert_coord is None: + ds_vert_coord = ds_mesh + if nedges_chunk is not None: ds = ds.chunk({'nEdges': nedges_chunk}) @@ -73,9 +83,9 @@ def compute_vertically_integrated_velocity( n_vert_levels = ds.sizes['nVertLevels'] layer_thickness = ds[f'{prefix}layerThickness'] - max_level_cell = ds_mesh.maxLevelCell - 1 - if 'minLevelCell' in ds_mesh: - min_level_cell = ds_mesh.minLevelCell - 1 + max_level_cell = ds_vert_coord.maxLevelCell - 1 + if 'minLevelCell' in ds_vert_coord: + min_level_cell = ds_vert_coord.minLevelCell - 1 else: min_level_cell = xr.zeros_like(max_level_cell) @@ -85,7 +95,7 @@ def compute_vertically_integrated_velocity( if min_depth is not None or max_depth is not None: z_mid_edge = _compute_zmid_edge( - ds, ds_mesh, prefix, logger, cell0, cell1 + ds, ds_vert_coord, prefix, logger, cell0, cell1 ) else: z_mid_edge = None @@ -139,13 +149,13 @@ def compute_vertically_integrated_velocity( return vert_integ_velocity -def _compute_zmid_edge(ds, ds_mesh, prefix, logger, cell0, cell1): +def _compute_zmid_edge(ds, ds_vert_coord, prefix, logger, cell0, cell1): layer_thickness = ds[f'{prefix}layerThickness'] if logger: logger.info(' Computing z_mid and z_mid_edge.') z_mid = compute_zmid( - ds_mesh.bottomDepth, ds_mesh.maxLevelCell, layer_thickness + ds_vert_coord.bottomDepth, ds_vert_coord.maxLevelCell, layer_thickness ) z_mid_edge = 0.5 * (z_mid.isel(nCells=cell0) + z_mid.isel(nCells=cell1)) return z_mid_edge diff --git a/conda_package/mpas_tools/ocean/viz/transect/plot.py b/conda_package/mpas_tools/ocean/viz/transect/plot.py index 61de48677..07e1f1652 100644 --- a/conda_package/mpas_tools/ocean/viz/transect/plot.py +++ b/conda_package/mpas_tools/ocean/viz/transect/plot.py @@ -192,6 +192,7 @@ def plot_feature_transects( fc, ds, ds_mesh=None, + ds_vert_coord=None, variable_list=None, cmap=None, flip=False, @@ -210,10 +211,19 @@ def plot_feature_transects( The transects to plot ds : xarray.Dataset - The MPAS-Ocean dataset to plot + The MPAS-Ocean dataset to plot. It must contain ``layerThickness``, + which is a time-evolving (dynamic) field and so is no longer read from + the mesh dataset. ds_mesh : xarray.Dataset, optional - The MPAS-Ocean mesh to use for plotting, the same as ``ds`` by default + The MPAS-Ocean horizontal mesh to use for plotting, the same as ``ds`` + by default + + ds_vert_coord : xarray.Dataset, optional + A dataset with the vertical coordinate variables ``minLevelCell``, + ``maxLevelCell`` and ``bottomDepth``. The same as ``ds_mesh`` by + default, but in Omega these are stored separately from the horizontal + mesh. variable_list : list of str, optional The variables to plot @@ -237,17 +247,27 @@ def plot_feature_transects( add_z : bool, optional Whether to add zMid and zInterface to the mesh dataset """ + if ds_mesh is None: + ds_mesh = ds + if ds_vert_coord is None: + ds_vert_coord = ds_mesh + if 'Time' in ds.dims: ds = ds.isel(Time=0) if 'Time' in ds_mesh.dims: ds_mesh = ds_mesh.isel(Time=0) + if 'Time' in ds_vert_coord.dims: + ds_vert_coord = ds_vert_coord.isel(Time=0) + if add_z: - _add_z(ds_mesh) + _add_z(ds, ds_mesh, ds_vert_coord) print('\nBuilding transect geometry...') - transects = _compute_feature_transects(fc, ds_mesh, flip) + transects = _compute_feature_transects( + fc, ds, ds_mesh, ds_vert_coord, flip + ) fc_transects = dict() for transect in fc.features: @@ -311,6 +331,14 @@ def plot_feature_transects_main(): help='An MPAS-Ocean mesh file. If not specified, the ' 'MPAS-Ocean data file must contain the mesh.', ) + parser.add_argument( + '-vc', + '--vert_coord', + dest='vert_coord_filename', + help='An MPAS-Ocean/Omega file with minLevelCell, maxLevelCell and ' + 'bottomDepth. If not specified, the mesh file (or data file) must ' + 'contain these.', + ) parser.add_argument( '-f', '--file', @@ -369,6 +397,11 @@ def plot_feature_transects_main(): else: ds_mesh = ds + if args.vert_coord_filename is not None: + ds_vert_coord = xr.open_dataset(args.vert_coord_filename) + else: + ds_vert_coord = ds_mesh + variable_list = args.variable_list if 'Time' in ds.dims: @@ -377,10 +410,14 @@ def plot_feature_transects_main(): if 'Time' in ds_mesh.dims: ds_mesh = ds_mesh.isel(Time=0) + if 'Time' in ds_vert_coord.dims: + ds_vert_coord = ds_vert_coord.isel(Time=0) + plot_feature_transects( fc=fc, ds=ds, ds_mesh=ds_mesh, + ds_vert_coord=ds_vert_coord, variable_list=variable_list, cmap=args.colormap, flip=args.flip, @@ -488,18 +525,18 @@ def _plot_outline( ) -def _compute_feature_transects(fc, ds_mesh, flip): +def _compute_feature_transects(fc, ds, ds_mesh, ds_vert_coord, flip): """ build a sequence of triangles showing the transect intersecting mpas cells """ transects = dict() - layer_thickness = ds_mesh.layerThickness - bottom_depth = ds_mesh.bottomDepth - max_level_cell = ds_mesh.maxLevelCell - 1 - if 'minLevelCell' in ds_mesh: - min_level_cell = ds_mesh.minLevelCell - 1 + layer_thickness = ds.layerThickness + bottom_depth = ds_vert_coord.bottomDepth + max_level_cell = ds_vert_coord.maxLevelCell - 1 + if 'minLevelCell' in ds_vert_coord: + min_level_cell = ds_vert_coord.minLevelCell - 1 else: min_level_cell = xr.zeros_like(max_level_cell) @@ -592,16 +629,16 @@ def _plot_feature_transect( ds_transect.to_netcdf(f'{transect_prefix}_{var_name}.nc') -def _add_z(ds_mesh): +def _add_z(ds, ds_mesh, ds_vert_coord): """ Add zMid and zInterface to ``ds_mesh``, useful for debugging """ - layer_thickness = ds_mesh.layerThickness - bottom_depth = ds_mesh.bottomDepth - max_level_cell = ds_mesh.maxLevelCell - 1 - if 'minLevelCell' in ds_mesh: - min_level_cell = ds_mesh.minLevelCell - 1 + layer_thickness = ds.layerThickness + bottom_depth = ds_vert_coord.bottomDepth + max_level_cell = ds_vert_coord.maxLevelCell - 1 + if 'minLevelCell' in ds_vert_coord: + min_level_cell = ds_vert_coord.minLevelCell - 1 else: min_level_cell = xr.zeros_like(max_level_cell) From f04d1895aa070b916c92710511fd7217701208a2 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Fri, 26 Jun 2026 12:36:46 +0200 Subject: [PATCH 3/5] Update depth tests for the three-dataset model add_zmid now reads layerThickness from the input file, so include it there in test_add_zmid. Add a test that add_depth raises a clear error when refBottomDepth is missing, as is the case for Omega. Co-Authored-By: Claude Opus 4.8 --- conda_package/tests/test_depth.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/conda_package/tests/test_depth.py b/conda_package/tests/test_depth.py index 4e8fc8098..ef6af7b40 100644 --- a/conda_package/tests/test_depth.py +++ b/conda_package/tests/test_depth.py @@ -100,13 +100,37 @@ def test_add_depth(): ) +def test_add_depth_without_ref_bottom_depth(): + # add_depth relies on refBottomDepth (a 1D z-level/z-star assumption) and + # should raise a clear error when it is missing, as is the case for Omega + dsMesh = create_3d_mesh() + out_filename = 'test_depth_out.nc' + + dsIn = xarray.Dataset() + dsIn['temperature'] = xarray.ones_like(dsMesh.layerThickness) + dsIn['maxLevelCell'] = dsMesh.maxLevelCell + dsIn['bottomDepth'] = dsMesh.bottomDepth + write_netcdf(dsIn, 'test_depth_no_ref.nc') + + try: + add_depth('test_depth_no_ref.nc', out_filename) + except ValueError as exc: + assert 'refBottomDepth' in str(exc) + else: + raise AssertionError('add_depth did not raise without refBottomDepth') + + def test_add_zmid(): dsMesh = create_3d_mesh() mesh_filename = 'test_depth_mesh.nc' out_filename = 'test_depth_out.nc' + # layerThickness is a dynamic field and is always read from the input file, + # so it must be present there (the coordinate file only needs bottomDepth + # and maxLevelCell) dsIn = xarray.Dataset() dsIn['temperature'] = xarray.ones_like(dsMesh.layerThickness) + dsIn['layerThickness'] = dsMesh.layerThickness write_netcdf(dsIn, 'test_depth_in.nc') # test adding zMid once to the mesh and once to the input file, with the @@ -187,5 +211,6 @@ def test_write_time_varying_zmid(): test_compute_depth() test_compute_zmid() test_add_depth() + test_add_depth_without_ref_bottom_depth() test_add_zmid() test_write_time_varying_zmid() From 7276613cd656ff9bf137e5deb81ad11a9b0afb7b Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Fri, 26 Jun 2026 12:36:54 +0200 Subject: [PATCH 4/5] Document that the 1D depth coordinate is not Omega compatible Add a warning to the "Adding a Depth Coordinate" docs that add_depth and compute_depth rely on refBottomDepth and so assume a z-level/z-star vertical coordinate, which Omega does not provide. Point users to the 3D zMid coordinate instead. Co-Authored-By: Claude Opus 4.8 --- conda_package/docs/ocean/depth.rst | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/conda_package/docs/ocean/depth.rst b/conda_package/docs/ocean/depth.rst index e0c0b05d1..17705b77a 100644 --- a/conda_package/docs/ocean/depth.rst +++ b/conda_package/docs/ocean/depth.rst @@ -17,6 +17,14 @@ directly if one has a suitable ``refBottomDepth`` data array indicating the reference depth of the bottom of each layer in a 1D coordinate that is independent of both time and horizontal coordinate. +.. warning:: + + This 1D ``depth`` coordinate relies on ``refBottomDepth`` and therefore + assumes a z-level/z-star vertical coordinate. It is **not supported for + Omega or other general vertical coordinates**, which do not provide + ``refBottomDepth``. In those cases, use the 3D ``zMid`` coordinate + (:py:func:`mpas_tools.ocean.depth.add_zmid()`) described below. + Adding a 3D zMid coordinate --------------------------- From 2830ba857f9ba588652a1072b69d622c5526e9b2 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Fri, 26 Jun 2026 12:38:17 +0200 Subject: [PATCH 5/5] Update version to 2.0.0 Bump the major version to reflect the non-backwards-compatible Omega changes: layerThickness must now come from the dynamic dataset rather than the mesh, and the refBottomDepth-based 1D depth coordinate is no longer supported for Omega. Co-Authored-By: Claude Opus 4.8 --- CITATION.cff | 4 ++-- conda_package/mpas_tools/__init__.py | 2 +- conda_package/recipe/recipe.yaml | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/CITATION.cff b/CITATION.cff index 337adc3ec..1720d7876 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -79,5 +79,5 @@ authors: family-names: Smith repository-code: 'https://github.com/MPAS-Dev/MPAS-Tools' url: 'https://mpas-dev.github.io/MPAS-Tools/master/' -version: 1.6.0 -date-released: '2026-05-09' +version: 2.0.0 +date-released: '2026-06-26' diff --git a/conda_package/mpas_tools/__init__.py b/conda_package/mpas_tools/__init__.py index f0c86403d..09b6d0b04 100644 --- a/conda_package/mpas_tools/__init__.py +++ b/conda_package/mpas_tools/__init__.py @@ -1,2 +1,2 @@ -__version_info__ = (1, 6, 0) +__version_info__ = (2, 0, 0) __version__ = '.'.join(str(vi) for vi in __version_info__) diff --git a/conda_package/recipe/recipe.yaml b/conda_package/recipe/recipe.yaml index e8754f81f..552e9f426 100644 --- a/conda_package/recipe/recipe.yaml +++ b/conda_package/recipe/recipe.yaml @@ -2,7 +2,7 @@ schema_version: 1 context: name: mpas_tools - version: 1.6.0 + version: 2.0.0 package: name: ${{ name | lower }}