Workaround for zstar min levels

E3SM-Project · Nov 9, 2023 · b007f81 · b007f81
1 parent 0527fc5
commit b007f81
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/polaris/ocean/vertical/zlevel.py b/polaris/ocean/vertical/zlevel.py
@@ -79,7 +79,8 @@ def init_z_level_vertical_coord(config, ds):
         ds.maxLevelCell)
 
 
-def compute_min_max_level_cell(refTopDepth, refBottomDepth, ssh, bottomDepth):
+def compute_min_max_level_cell(refTopDepth, refBottomDepth, ssh, bottomDepth,
+                               min_vert_levels=None, min_layer_thickness=None):
     """
     Compute ``minLevelCell`` and ``maxLevelCell`` indices as well as a cell
     mask for the given reference grid and top and bottom topography.
@@ -110,7 +111,10 @@ def compute_min_max_level_cell(refTopDepth, refBottomDepth, ssh, bottomDepth):
     cellMask : xarray.DataArray
         A boolean mask of where there are valid cells
     """
-    valid = bottomDepth > -ssh
+    if min_layer_thickness is not None:
+        valid = bottomDepth + min_layer_thickness * min_vert_levels >= -ssh
+    else:
+        valid = bottomDepth > -ssh
 
     aboveTopMask = (refBottomDepth <= -ssh).transpose('nCells', 'nVertLevels')
     aboveBottomMask = (refTopDepth < bottomDepth).transpose(
@@ -119,7 +123,9 @@ def compute_min_max_level_cell(refTopDepth, refBottomDepth, ssh, bottomDepth):
 
     minLevelCell = (aboveTopMask.sum(dim='nVertLevels')).where(valid, 0)
     maxLevelCell = (aboveBottomMask.sum(dim='nVertLevels') - 1).where(valid, 0)
-
+    if min_vert_levels is not None:
+        maxLevelCell = numpy.maximum(maxLevelCell,
+                                     minLevelCell + min_vert_levels - 1)
     cellMask = numpy.logical_and(numpy.logical_not(aboveTopMask),
                                  aboveBottomMask)
     cellMask = numpy.logical_and(cellMask, valid)

diff --git a/polaris/ocean/vertical/zstar.py b/polaris/ocean/vertical/zstar.py
@@ -65,9 +65,14 @@ def init_z_star_vertical_coord(config, ds):
     ds['vertCoordMovementWeights'] = xarray.ones_like(ds.refBottomDepth)
 
     restingSSH = xarray.zeros_like(ds.bottomDepth)
+    min_vert_levels = config.getint('vertical_grid', 'min_vert_levels')
+    min_layer_thickness = config.getfloat('vertical_grid',
+                                          'min_layer_thickness')
     ds['minLevelCell'], ds['maxLevelCell'], ds['cellMask'] = \
         compute_min_max_level_cell(ds.refTopDepth, ds.refBottomDepth,
-                                   restingSSH, ds.bottomDepth)
+                                   restingSSH, ds.bottomDepth,
+                                   min_vert_levels=min_vert_levels,
+                                   min_layer_thickness=min_layer_thickness)
 
     ds['bottomDepth'], ds['maxLevelCell'] = alter_bottom_depth(
         config, ds.bottomDepth, ds.refBottomDepth, ds.maxLevelCell)
@@ -113,7 +118,7 @@ def _compute_z_star_layer_thickness(restingThickness, ssh, bottomDepth,
     nVertLevels = restingThickness.sizes['nVertLevels']
     layerThickness = []
 
-    layerStretch = (ssh + bottomDepth) / bottomDepth
+    layerStretch = (ssh + bottomDepth) / numpy.sum(restingThickness, axis=1)
     for zIndex in range(nVertLevels):
         mask = numpy.logical_and(zIndex >= minLevelCell,
                                  zIndex <= maxLevelCell)