From 249ad58d60c6488c7311332512ed6bcbe2cf18d1 Mon Sep 17 00:00:00 2001
From: Alice Barthel <abarthel@lanl.gov>
Date: Mon, 6 Feb 2023 13:14:42 -0600
Subject: [PATCH 1/5] modified names of functions specific to tanh thickness
 for clarity

---
 polaris/ocean/vertical/grid_1d.py | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/polaris/ocean/vertical/grid_1d.py b/polaris/ocean/vertical/grid_1d.py
index 36df88ff3..f0b2dee78 100644
--- a/polaris/ocean/vertical/grid_1d.py
+++ b/polaris/ocean/vertical/grid_1d.py
@@ -141,7 +141,7 @@ def _read_json(grid_type):
 def _create_tanh_dz_grid(num_vert_levels, bottom_depth, min_layer_thickness,
                          max_layer_thickness):
     """
-    Creates the vertical grid for MPAS-Ocean and writes it to a NetCDF file
+    Creates the tanh vertical grid for MPAS-Ocean and writes it to a NetCDF file
 
     Parameters
     ----------
@@ -173,20 +173,21 @@ def _create_tanh_dz_grid(num_vert_levels, bottom_depth, min_layer_thickness,
     # and the root finder will determine a value of delta (sol.root) such that
     # match_bottom is within a tolerance of zero, meaning the bottom of the
     # coordinate computed by cumsum_z hits bottom_depth almost exactly
-    sol = root_scalar(_match_bottom, method='brentq',
+    sol = root_scalar(_tanh_match_bottom, method='brentq',
                       bracket=[dz1, 10 * bottom_depth],
                       args=(nz, dz1, dz2, bottom_depth))
 
     delta = sol.root
-    layerThickness, z = _cumsum_z(delta, nz, dz1, dz2)
+    layerThickness, z = _tanh_cumsum_z(delta, nz, dz1, dz2)
     interfaces = -z
 
     return interfaces
 
 
-def _match_bottom(delta, nz, dz1, dz2, bottom_depth):
+def _tanh_match_bottom(delta, nz, dz1, dz2, bottom_depth):
     """
-    Compute the difference between the bottom depth computed with the given
+    For tanh layer thickness, compute the difference between the 
+    bottom depth computed with the given
     parameters and the target ``bottom_depth``, used in the root finding
     algorithm to determine which value of ``delta`` to use.
 
@@ -216,14 +217,14 @@ def _match_bottom(delta, nz, dz1, dz2, bottom_depth):
         The computed bottom depth minus the target ``bottom_depth``.  ``diff``
         should be zero when we have found the desired ``delta``.
     """
-    _, z = _cumsum_z(delta, nz, dz1, dz2)
+    _, z = _tanh_cumsum_z(delta, nz, dz1, dz2)
     diff = -bottom_depth - z[-1]
     return diff
 
 
-def _cumsum_z(delta, nz, dz1, dz2):
+def _tanh_cumsum_z(delta, nz, dz1, dz2):
     """
-    Compute layer interface depths and layer thicknesses over ``nz`` layers
+    Compute tanh layer interface depths and layer thicknesses over ``nz`` layers
 
     Parameters
     ----------
@@ -252,14 +253,14 @@ def _cumsum_z(delta, nz, dz1, dz2):
     dz = np.zeros(nz)
     z = np.zeros(nz + 1)
     for zindex in range(nz):
-        dz[zindex] = _dz_z(z[zindex], dz1, dz2, delta)
+        dz[zindex] = _tanh_dz_z(z[zindex], dz1, dz2, delta)
         z[zindex + 1] = z[zindex] - dz[zindex]
     return dz, z
 
 
-def _dz_z(z, dz1, dz2, delta):
+def _tanh_dz_z(z, dz1, dz2, delta):
     """
-    layer thickness as a function of depth
+    Tanh layer thickness as a function of depth
 
     Parameters
     ----------

From f0317c43e92d14e07163bab2c8b3ab7f021a08cb Mon Sep 17 00:00:00 2001
From: Xylar Asay-Davis <xylarstorm@gmail.com>
Date: Sun, 2 Apr 2023 17:28:15 +0200
Subject: [PATCH 2/5] Make grid_1d into a module

This merge also breaks out functions related to `tanh_dz` into
their own module.
---
 polaris/ocean/vertical/grid_1d.py          | 285 ---------------------
 polaris/ocean/vertical/grid_1d/__init__.py | 139 ++++++++++
 polaris/ocean/vertical/grid_1d/tanh_dz.py  | 153 +++++++++++
 3 files changed, 292 insertions(+), 285 deletions(-)
 delete mode 100644 polaris/ocean/vertical/grid_1d.py
 create mode 100644 polaris/ocean/vertical/grid_1d/__init__.py
 create mode 100644 polaris/ocean/vertical/grid_1d/tanh_dz.py

diff --git a/polaris/ocean/vertical/grid_1d.py b/polaris/ocean/vertical/grid_1d.py
deleted file mode 100644
index f0b2dee78..000000000
--- a/polaris/ocean/vertical/grid_1d.py
+++ /dev/null
@@ -1,285 +0,0 @@
-import importlib.resources as imp_res
-import json
-
-import numpy
-import numpy as np
-from netCDF4 import Dataset
-from scipy.optimize import root_scalar
-
-
-def generate_1d_grid(config):
-    """
-    Generate a vertical grid for a test case, using the config options in the
-    ``vertical_grid`` section
-
-    Parameters
-    ----------
-    config : polaris.config.PolarisConfigParser
-        Configuration options with parameters used to construct the vertical
-        grid
-
-    Returns
-    -------
-    interfaces : numpy.ndarray
-        A 1D array of positive depths for layer interfaces in meters
-    """
-    section = config['vertical_grid']
-    grid_type = section.get('grid_type')
-    if grid_type == 'uniform':
-        vert_levels = section.getint('vert_levels')
-        interfaces = _generate_uniform(vert_levels)
-    elif grid_type == 'tanh_dz':
-        vert_levels = section.getint('vert_levels')
-        min_layer_thickness = section.getfloat('min_layer_thickness')
-        max_layer_thickness = section.getfloat('max_layer_thickness')
-        bottom_depth = section.getfloat('bottom_depth')
-        interfaces = _create_tanh_dz_grid(vert_levels, bottom_depth,
-                                          min_layer_thickness,
-                                          max_layer_thickness)
-
-    elif grid_type in ['60layerPHC', '80layerE3SMv1', '100layerE3SMv1']:
-        interfaces = _read_json(grid_type)
-    else:
-        raise ValueError(f'Unexpected grid type: {grid_type}')
-
-    if config.has_option('vertical_grid', 'bottom_depth') and \
-            grid_type != 'tanh_dz':
-        bottom_depth = section.getfloat('bottom_depth')
-        # renormalize to the requested range
-        interfaces = (bottom_depth / interfaces[-1]) * interfaces
-
-    return interfaces
-
-
-def write_1d_grid(interfaces, out_filename):
-    """
-    write the vertical grid to a file
-
-    Parameters
-    ----------
-    interfaces : numpy.ndarray
-        A 1D array of positive depths for layer interfaces in meters
-
-    out_filename : str
-        MPAS file name for output of vertical grid
-    """
-
-    nz = len(interfaces) - 1
-
-    # open a new netCDF file for writing.
-    ncfile = Dataset(out_filename, 'w')
-    # create the depth_t dimension.
-    ncfile.createDimension('nVertLevels', nz)
-
-    refBottomDepth = ncfile.createVariable(
-        'refBottomDepth', np.dtype('float64').char, ('nVertLevels',))
-    refMidDepth = ncfile.createVariable(
-        'refMidDepth', np.dtype('float64').char, ('nVertLevels',))
-    refLayerThickness = ncfile.createVariable(
-        'refLayerThickness', np.dtype('float64').char, ('nVertLevels',))
-
-    botDepth = interfaces[1:]
-    midDepth = 0.5 * (interfaces[0:-1] + interfaces[1:])
-
-    refBottomDepth[:] = botDepth
-    refMidDepth[:] = midDepth
-    refLayerThickness[:] = interfaces[1:] - interfaces[0:-1]
-    ncfile.close()
-
-
-def add_1d_grid(config, ds):
-    """
-    Add a 1D vertical grid based on the config options in the ``vertical_grid``
-    section to a mesh data set
-
-    The following variables are 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).
-    There is considerable redundancy between these variables but each is
-    sometimes convenient.
-
-    Parameters
-    ----------
-    config : polaris.config.PolarisConfigParser
-        Configuration options with parameters used to construct the vertical
-        grid
-
-    ds : xarray.Dataset
-        A data set to add the grid variables to
-    """
-
-    interfaces = generate_1d_grid(config=config)
-
-    ds['refTopDepth'] = ('nVertLevels', interfaces[0:-1])
-    ds['refZMid'] = ('nVertLevels', -0.5 * (interfaces[1:] + interfaces[0:-1]))
-    ds['refBottomDepth'] = ('nVertLevels', interfaces[1:])
-    ds['refInterfaces'] = ('nVertLevelsP1', interfaces)
-
-
-def _generate_uniform(vert_levels):
-    """ Generate uniform layer interfaces between 0 and 1 """
-    interfaces = numpy.linspace(0., 1., vert_levels + 1)
-    return interfaces
-
-
-def _read_json(grid_type):
-    """ Read the grid interfaces from a json file """
-
-    package = 'polaris.ocean.vertical'
-    filename = f'{grid_type}.json'
-    with imp_res.files(package).joinpath(filename).open('r') as data_file:
-        data = json.load(data_file)
-        interfaces = numpy.array(data)
-
-    return interfaces
-
-
-def _create_tanh_dz_grid(num_vert_levels, bottom_depth, min_layer_thickness,
-                         max_layer_thickness):
-    """
-    Creates the tanh vertical grid for MPAS-Ocean and writes it to a NetCDF file
-
-    Parameters
-    ----------
-    num_vert_levels : int
-        Number of vertical levels for the grid
-
-    bottom_depth : float
-        bottom depth for the chosen vertical coordinate [m]
-
-    min_layer_thickness : float
-        Target thickness of the first layer [m]
-
-    max_layer_thickness : float
-        Target maximum thickness in column [m]
-
-    Returns
-    -------
-    interfaces : numpy.ndarray
-        A 1D array of positive depths for layer interfaces in meters
-    """
-
-    nz = num_vert_levels
-    dz1 = min_layer_thickness
-    dz2 = max_layer_thickness
-
-    # the bracket here is large enough that it should hopefully encompass any
-    # reasonable value of delta, the characteristic length scale over which
-    # dz varies.  The args are passed on to the match_bottom function below,
-    # and the root finder will determine a value of delta (sol.root) such that
-    # match_bottom is within a tolerance of zero, meaning the bottom of the
-    # coordinate computed by cumsum_z hits bottom_depth almost exactly
-    sol = root_scalar(_tanh_match_bottom, method='brentq',
-                      bracket=[dz1, 10 * bottom_depth],
-                      args=(nz, dz1, dz2, bottom_depth))
-
-    delta = sol.root
-    layerThickness, z = _tanh_cumsum_z(delta, nz, dz1, dz2)
-    interfaces = -z
-
-    return interfaces
-
-
-def _tanh_match_bottom(delta, nz, dz1, dz2, bottom_depth):
-    """
-    For tanh layer thickness, compute the difference between the 
-    bottom depth computed with the given
-    parameters and the target ``bottom_depth``, used in the root finding
-    algorithm to determine which value of ``delta`` to use.
-
-    Parameters
-    ----------
-    delta : float
-        The characteristic length scale over which dz varies (this parameter
-        will be optimized to hit a target depth in a target number of layers)
-
-    nz : int
-        The number of layers
-
-    dz1 : float
-        The layer thickness at the top of the ocean (z = 0)
-
-    dz2 : float
-        The layer thickness at z --> -infinity
-
-    bottom_depth: float
-        depth of the bottom of the ocean that should match the bottom layer
-        interface.  Note: the bottom_depth is positive, whereas the layer
-        interfaces are negative.
-
-    Returns
-    -------
-    diff : float
-        The computed bottom depth minus the target ``bottom_depth``.  ``diff``
-        should be zero when we have found the desired ``delta``.
-    """
-    _, z = _tanh_cumsum_z(delta, nz, dz1, dz2)
-    diff = -bottom_depth - z[-1]
-    return diff
-
-
-def _tanh_cumsum_z(delta, nz, dz1, dz2):
-    """
-    Compute tanh layer interface depths and layer thicknesses over ``nz`` layers
-
-    Parameters
-    ----------
-    delta : float
-        The characteristic length scale over which dz varies (this parameter
-        will be optimized to hit a target depth in a target number of layers)
-
-    nz : int
-        The number of layers
-
-    dz1 : float
-        The layer thickness at the top of the ocean (z = 0)
-
-    dz2 : float
-        The layer thickness at z --> -infinity
-
-    Returns
-    -------
-    dz : numpy.ndarray
-        The layer thicknesses for each layer
-
-    z : numpy.ndarray
-        The depth (positive up) of each layer interface (``nz + 1`` total
-        elements)
-    """
-    dz = np.zeros(nz)
-    z = np.zeros(nz + 1)
-    for zindex in range(nz):
-        dz[zindex] = _tanh_dz_z(z[zindex], dz1, dz2, delta)
-        z[zindex + 1] = z[zindex] - dz[zindex]
-    return dz, z
-
-
-def _tanh_dz_z(z, dz1, dz2, delta):
-    """
-    Tanh layer thickness as a function of depth
-
-    Parameters
-    ----------
-    z : float
-        Depth coordinate (positive up) at which to find the layer thickness
-
-    dz1 : float
-        The layer thickness at the top of the ocean (z = 0)
-
-    dz2 : float
-        The layer thickness at z --> -infinity
-
-    delta : float
-        The characteristic length scale over which dz varies (this parameter
-        will be optimized to hit a target depth in a target numer of layers)
-
-    Returns
-    -------
-    dz : float
-        The layer thickness
-    """
-    return (dz2 - dz1) * np.tanh(-z * np.pi / delta) + dz1
diff --git a/polaris/ocean/vertical/grid_1d/__init__.py b/polaris/ocean/vertical/grid_1d/__init__.py
new file mode 100644
index 000000000..392c1b61c
--- /dev/null
+++ b/polaris/ocean/vertical/grid_1d/__init__.py
@@ -0,0 +1,139 @@
+import json
+from importlib import resources
+
+import numpy as np
+from netCDF4 import Dataset
+
+from polaris.ocean.vertical.grid_1d.tanh_dz import create_tanh_dz_grid
+
+
+def generate_1d_grid(config):
+    """
+    Generate a vertical grid for a test case, using the config options in the
+    ``vertical_grid`` section
+
+    Parameters
+    ----------
+    config : polaris.config.PolarisConfigParser
+        Configuration options with parameters used to construct the vertical
+        grid
+
+    Returns
+    -------
+    interfaces : numpy.ndarray
+        A 1D array of positive depths for layer interfaces in meters
+    """
+    section = config['vertical_grid']
+    grid_type = section.get('grid_type')
+    if grid_type == 'uniform':
+        vert_levels = section.getint('vert_levels')
+        interfaces = _generate_uniform(vert_levels)
+    elif grid_type == 'tanh_dz':
+        vert_levels = section.getint('vert_levels')
+        min_layer_thickness = section.getfloat('min_layer_thickness')
+        max_layer_thickness = section.getfloat('max_layer_thickness')
+        bottom_depth = section.getfloat('bottom_depth')
+        interfaces = create_tanh_dz_grid(vert_levels,
+                                         bottom_depth,
+                                         min_layer_thickness,
+                                         max_layer_thickness)
+
+    elif grid_type in ['60layerPHC', '80layerE3SMv1', '100layerE3SMv1']:
+        interfaces = _read_json(grid_type)
+    else:
+        raise ValueError(f'Unexpected grid type: {grid_type}')
+
+    if config.has_option('vertical_grid', 'bottom_depth') and \
+            grid_type != 'tanh_dz':
+        bottom_depth = section.getfloat('bottom_depth')
+        # renormalize to the requested range
+        interfaces = (bottom_depth / interfaces[-1]) * interfaces
+
+    return interfaces
+
+
+def write_1d_grid(interfaces, out_filename):
+    """
+    write the vertical grid to a file
+
+    Parameters
+    ----------
+    interfaces : numpy.ndarray
+        A 1D array of positive depths for layer interfaces in meters
+
+    out_filename : str
+        MPAS file name for output of vertical grid
+    """
+
+    nz = len(interfaces) - 1
+
+    # open a new netCDF file for writing.
+    ncfile = Dataset(out_filename, 'w')
+    # create the depth_t dimension.
+    ncfile.createDimension('nVertLevels', nz)
+
+    refBottomDepth = ncfile.createVariable(
+        'refBottomDepth', np.dtype('float64').char, ('nVertLevels',))
+    refMidDepth = ncfile.createVariable(
+        'refMidDepth', np.dtype('float64').char, ('nVertLevels',))
+    refLayerThickness = ncfile.createVariable(
+        'refLayerThickness', np.dtype('float64').char, ('nVertLevels',))
+
+    botDepth = interfaces[1:]
+    midDepth = 0.5 * (interfaces[0:-1] + interfaces[1:])
+
+    refBottomDepth[:] = botDepth
+    refMidDepth[:] = midDepth
+    refLayerThickness[:] = interfaces[1:] - interfaces[0:-1]
+    ncfile.close()
+
+
+def add_1d_grid(config, ds):
+    """
+    Add a 1D vertical grid based on the config options in the ``vertical_grid``
+    section to a mesh data set
+
+    The following variables are 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).
+    There is considerable redundancy between these variables but each is
+    sometimes convenient.
+
+    Parameters
+    ----------
+    config : polaris.config.PolarisConfigParser
+        Configuration options with parameters used to construct the vertical
+        grid
+
+    ds : xarray.Dataset
+        A data set to add the grid variables to
+    """
+
+    interfaces = generate_1d_grid(config=config)
+
+    ds['refTopDepth'] = ('nVertLevels', interfaces[0:-1])
+    ds['refZMid'] = ('nVertLevels', -0.5 * (interfaces[1:] + interfaces[0:-1]))
+    ds['refBottomDepth'] = ('nVertLevels', interfaces[1:])
+    ds['refInterfaces'] = ('nVertLevelsP1', interfaces)
+
+
+def _generate_uniform(vert_levels):
+    """ Generate uniform layer interfaces between 0 and 1 """
+    interfaces = np.linspace(0., 1., vert_levels + 1)
+    return interfaces
+
+
+def _read_json(grid_type):
+    """ Read the grid interfaces from a json file """
+
+    filename = f'{grid_type}.json'
+    with resources.open_text("polaris.ocean.vertical.grid_1d", filename) as \
+            data_file:
+        data = json.load(data_file)
+        interfaces = np.array(data)
+
+    return interfaces
diff --git a/polaris/ocean/vertical/grid_1d/tanh_dz.py b/polaris/ocean/vertical/grid_1d/tanh_dz.py
new file mode 100644
index 000000000..ab43829fd
--- /dev/null
+++ b/polaris/ocean/vertical/grid_1d/tanh_dz.py
@@ -0,0 +1,153 @@
+import numpy as np
+from scipy.optimize import root_scalar
+
+
+def create_tanh_dz_grid(num_vert_levels, bottom_depth, min_layer_thickness,
+                        max_layer_thickness):
+    """
+    Creates the vertical grid for MPAS-Ocean in which the layer thicknesses
+    vary as a tanh function in depth with origin at z=0.  The tanh function
+    is scaled in depth such that the bottom layer interface is exactly at
+    ``bottom_depth``.
+
+    Parameters
+    ----------
+    num_vert_levels : int
+        Number of vertical levels for the grid
+
+    bottom_depth : float
+        bottom depth for the chosen vertical coordinate [m]
+
+    min_layer_thickness : float
+        Target thickness of the first layer [m]
+
+    max_layer_thickness : float
+        Target maximum thickness in column [m]
+
+    Returns
+    -------
+    interfaces : numpy.ndarray
+        A 1D array of positive depths for layer interfaces in meters
+    """
+
+    nz = num_vert_levels
+    dz1 = min_layer_thickness
+    dz2 = max_layer_thickness
+
+    # the bracket here is large enough that it should hopefully encompass any
+    # reasonable value of delta, the characteristic length scale over which
+    # dz varies.  The args are passed on to the match_bottom function below,
+    # and the root finder will determine a value of delta (sol.root) such that
+    # match_bottom is within a tolerance of zero, meaning the bottom of the
+    # coordinate computed by cumsum_z hits bottom_depth almost exactly
+    sol = root_scalar(_tanh_match_bottom, method='brentq',
+                      bracket=[dz1, 10 * bottom_depth],
+                      args=(nz, dz1, dz2, bottom_depth))
+
+    delta = sol.root
+    layerThickness, z = _tanh_cumsum_z(delta, nz, dz1, dz2)
+    interfaces = -z
+
+    return interfaces
+
+
+def _tanh_match_bottom(delta, nz, dz1, dz2, bottom_depth):
+    """
+    For tanh layer thickness, compute the difference between the
+    bottom depth computed with the given
+    parameters and the target ``bottom_depth``, used in the root finding
+    algorithm to determine which value of ``delta`` to use.
+
+    Parameters
+    ----------
+    delta : float
+        The characteristic length scale over which dz varies (this parameter
+        will be optimized to hit a target depth in a target number of layers)
+
+    nz : int
+        The number of layers
+
+    dz1 : float
+        The layer thickness at the top of the ocean (z = 0)
+
+    dz2 : float
+        The layer thickness at z --> -infinity
+
+    bottom_depth: float
+        depth of the bottom of the ocean that should match the bottom layer
+        interface.  Note: the bottom_depth is positive, whereas the layer
+        interfaces are negative.
+
+    Returns
+    -------
+    diff : float
+        The computed bottom depth minus the target ``bottom_depth``.  ``diff``
+        should be zero when we have found the desired ``delta``.
+    """
+    _, z = _tanh_cumsum_z(delta, nz, dz1, dz2)
+    diff = -bottom_depth - z[-1]
+    return diff
+
+
+def _tanh_cumsum_z(delta, nz, dz1, dz2):
+    """
+    Compute tanh layer interface depths and layer thicknesses over ``nz``
+    layers
+
+    Parameters
+    ----------
+    delta : float
+        The characteristic length scale over which dz varies (this parameter
+        will be optimized to hit a target depth in a target number of layers)
+
+    nz : int
+        The number of layers
+
+    dz1 : float
+        The layer thickness at the top of the ocean (z = 0)
+
+    dz2 : float
+        The layer thickness at z --> -infinity
+
+    Returns
+    -------
+    dz : numpy.ndarray
+        The layer thicknesses for each layer
+
+    z : numpy.ndarray
+        The depth (positive up) of each layer interface (``nz + 1`` total
+        elements)
+    """
+    dz = np.zeros(nz)
+    z = np.zeros(nz + 1)
+    for zindex in range(nz):
+        dz[zindex] = _tanh_dz_z(z[zindex], dz1, dz2, delta)
+        z[zindex + 1] = z[zindex] - dz[zindex]
+    return dz, z
+
+
+def _tanh_dz_z(z, dz1, dz2, delta):
+    """
+    Tanh layer thickness as a function of depth
+
+    Parameters
+    ----------
+    z : float
+        Depth coordinate (positive up) at which to find the layer thickness
+
+    dz1 : float
+        The layer thickness at the top of the ocean (z = 0)
+
+    dz2 : float
+        The layer thickness at z --> -infinity
+
+    delta : float
+        The characteristic length scale over which dz varies (this parameter
+        will be optimized to hit a target depth in a target numer of layers)
+
+    Returns
+    -------
+    dz : float
+        The layer thickness
+    """
+    return (dz2 - dz1) * np.tanh(-z * np.pi / delta) + dz1

From 8e5e2b700691ad2993866a4dc84b148f20693401 Mon Sep 17 00:00:00 2001
From: Xylar Asay-Davis <xylarstorm@gmail.com>
Date: Sun, 2 Apr 2023 17:30:30 +0200
Subject: [PATCH 3/5] Add index_tanh_dz vertical grid

---
 polaris/ocean/vertical/grid_1d/__init__.py    |  16 ++
 .../ocean/vertical/grid_1d/index_tanh_dz.py   | 176 ++++++++++++++++++
 2 files changed, 192 insertions(+)
 create mode 100644 polaris/ocean/vertical/grid_1d/index_tanh_dz.py

diff --git a/polaris/ocean/vertical/grid_1d/__init__.py b/polaris/ocean/vertical/grid_1d/__init__.py
index 392c1b61c..716910c63 100644
--- a/polaris/ocean/vertical/grid_1d/__init__.py
+++ b/polaris/ocean/vertical/grid_1d/__init__.py
@@ -4,6 +4,9 @@
 import numpy as np
 from netCDF4 import Dataset
 
+from polaris.ocean.vertical.grid_1d.index_tanh_dz import (
+    create_index_tanh_dz_grid,
+)
 from polaris.ocean.vertical.grid_1d.tanh_dz import create_tanh_dz_grid
 
 
@@ -38,6 +41,19 @@ def generate_1d_grid(config):
                                          min_layer_thickness,
                                          max_layer_thickness)
 
+    elif grid_type == 'index_tanh_dz':
+        vert_levels = section.getint('vert_levels')
+        min_layer_thickness = section.getfloat('min_layer_thickness')
+        max_layer_thickness = section.getfloat('max_layer_thickness')
+        bottom_depth = section.getfloat('bottom_depth')
+        transition_levels = section.getfloat('transition_levels')
+        interfaces = create_index_tanh_dz_grid(
+            vert_levels,
+            bottom_depth,
+            min_layer_thickness,
+            max_layer_thickness,
+            transition_levels)
+
     elif grid_type in ['60layerPHC', '80layerE3SMv1', '100layerE3SMv1']:
         interfaces = _read_json(grid_type)
     else:
diff --git a/polaris/ocean/vertical/grid_1d/index_tanh_dz.py b/polaris/ocean/vertical/grid_1d/index_tanh_dz.py
new file mode 100644
index 000000000..2023225d4
--- /dev/null
+++ b/polaris/ocean/vertical/grid_1d/index_tanh_dz.py
@@ -0,0 +1,176 @@
+import numpy as np
+from scipy.optimize import root_scalar
+
+
+def create_index_tanh_dz_grid(num_vert_levels, bottom_depth,
+                              min_layer_thickness, max_layer_thickness,
+                              transition_levels):
+    """
+    Creates layer thicknesses that vary as a tanh function in vertical index
+    (as opposed to z for the ``tanh_dz`` profile) for MPAS-Ocean
+
+    Parameters
+    ----------
+    num_vert_levels : int
+        Number of vertical levels for the grid
+
+    bottom_depth : float
+        bottom depth for the chosen vertical coordinate [m]
+
+    min_layer_thickness : float
+        Target thickness of the first layer [m]
+
+    max_layer_thickness : float
+        Target maximum thickness in column [m]
+
+    transition_levels : float
+        The width of the transition in resolution in layers
+
+    Returns
+    -------
+    interfaces : numpy.ndarray
+        A 1D array of positive depths for layer interfaces in meters
+    """
+
+    nz = num_vert_levels
+    dz1 = min_layer_thickness
+    dz2 = max_layer_thickness
+    delta = transition_levels
+
+    bracket = [0, num_vert_levels]
+
+    # the bracket here is large enough that it should hopefully encompass any
+    # reasonable value of delta, the characteristic length scale over which
+    # dz varies.  The args are passed on to the match_bottom function below,
+    # and the root finder will determine a value of delta (sol.root) such that
+    # match_bottom is within a tolerance of zero, meaning the bottom of the
+    # coordinate computed by cumsum_z hits bottom_depth almost exactly
+    sol = root_scalar(_index_tanh_match_bottom, method='brentq',
+                      bracket=bracket,
+                      args=(nz, dz1, dz2, bottom_depth, delta))
+
+    origin = sol.root
+    layerThickness, z = _index_tanh_cumsum_z(delta, nz, dz1, dz2, origin)
+    interfaces = -z
+
+    return interfaces
+
+
+def _index_tanh_match_bottom(origin, nz, dz1, dz2, bottom_depth, delta):
+    """
+    For tanh layer thickness, compute the difference between the
+    bottom depth computed with the given
+    parameters and the target ``bottom_depth``, used in the root finding
+    algorithm to determine which value of ``delta`` to use.
+
+    Parameters
+    ----------
+    origin : float
+        The layer index around which resolution transitions from the min to the
+        max layer thickness (this parameter will be optimized to hit a target
+        depth in a target numer of layers)
+
+    nz : int
+        The number of layers
+
+    dz1 : float
+        The layer thickness at the top of the ocean (z = 0)
+
+    dz2 : float
+        The layer thickness at z --> -infinity
+
+    bottom_depth: float
+        depth of the bottom of the ocean that should match the bottom layer
+        interface.  Note: the bottom_depth is positive, whereas the layer
+        interfaces are negative.
+
+    delta : float
+        The characteristic number of layers over which dz varies
+
+    Returns
+    -------
+    diff : float
+        The computed bottom depth minus the target ``bottom_depth``.  ``diff``
+        should be zero when we have found the desired ``delta``.
+    """
+    _, z = _index_tanh_cumsum_z(delta, nz, dz1, dz2, origin)
+    diff = -bottom_depth - z[-1]
+    return diff
+
+
+def _index_tanh_cumsum_z(delta, nz, dz1, dz2, origin):
+    """
+    Compute tanh layer interface depths and layer thicknesses over ``nz``
+    layers
+
+    Parameters
+    ----------
+    delta : float
+        The characteristic number of layers over which dz varies
+
+    nz : int
+        The number of layers
+
+    dz1 : float
+        The layer thickness at the top of the ocean (z = 0)
+
+    dz2 : float
+        The layer thickness at z --> -infinity
+
+    origin : float
+        The layer index around which resolution transitions from the min to the
+        max layer thickness (this parameter will be optimized to hit a target
+        depth in a target numer of layers)
+
+
+    Returns
+    -------
+    dz : numpy.ndarray
+        The layer thicknesses for each layer
+
+    z : numpy.ndarray
+        The depth (positive up) of each layer interface (``nz + 1`` total
+        elements)
+    """
+    dz = np.zeros(nz)
+    z = np.zeros(nz + 1)
+    for zindex in range(nz):
+        dz[zindex] = _index_tanh_dz_z(zindex, dz1, dz2, delta, origin)
+        z[zindex + 1] = z[zindex] - dz[zindex]
+    return dz, z
+
+
+def _index_tanh_dz_z(zindex, dz1, dz2, delta, origin):
+    """
+    Tanh layer thickness as a function of depth
+
+    Parameters
+    ----------
+    zindex : float
+        The layer index at which to find the layer thickness
+
+    dz1 : float
+        The layer thickness at the top of the ocean (z = 0)
+
+    dz2 : float
+        The layer thickness at z --> -infinity
+
+    delta : float
+        The characteristic number of layers over which dz varies
+
+    origin : float
+        The layer index around which resolution transitions from the min to the
+        max layer thickness (this parameter will be optimized to hit a target
+        depth in a target numer of layers)
+
+
+    Returns
+    -------
+    dz : float
+        The layer thickness
+    """
+    tanh = np.tanh((zindex - origin) * np.pi / delta)
+    # rescale such that tanh hits zero at index 0
+    tanh0 = np.tanh(-origin * np.pi / delta)
+    tanh = (tanh - tanh0) / (1. - tanh0)
+    return (dz2 - dz1) * tanh + dz1

From 0fe7f24c9fd01b92d016f5d25c5f08045474acd0 Mon Sep 17 00:00:00 2001
From: Xylar Asay-Davis <xylarstorm@gmail.com>
Date: Tue, 31 Oct 2023 11:56:49 +0100
Subject: [PATCH 4/5] Add missing json files to grid_1d

---
 .../vertical/grid_1d/100layerE3SMv1.json      | 101 ++++++++++++++++++
 .../ocean/vertical/grid_1d/60layerPHC.json    |  63 +++++++++++
 .../ocean/vertical/grid_1d/80layerE3SMv1.json |  81 ++++++++++++++
 3 files changed, 245 insertions(+)
 create mode 100644 polaris/ocean/vertical/grid_1d/100layerE3SMv1.json
 create mode 100644 polaris/ocean/vertical/grid_1d/60layerPHC.json
 create mode 100644 polaris/ocean/vertical/grid_1d/80layerE3SMv1.json

diff --git a/polaris/ocean/vertical/grid_1d/100layerE3SMv1.json b/polaris/ocean/vertical/grid_1d/100layerE3SMv1.json
new file mode 100644
index 000000000..a5c98f2d8
--- /dev/null
+++ b/polaris/ocean/vertical/grid_1d/100layerE3SMv1.json
@@ -0,0 +1,101 @@
+[0.0000e0,
+ 0.1510e1,
+ 0.3135e1,
+ 0.4882e1,
+ 0.6761e1,
+ 0.8779e1,
+ 0.1095e2,
+ 0.1327e2,
+ 0.1577e2,
+ 0.1845e2,
+ 0.2132e2,
+ 0.2440e2,
+ 0.2769e2,
+ 0.3122e2,
+ 0.3500e2,
+ 0.3904e2,
+ 0.4335e2,
+ 0.4797e2,
+ 0.5289e2,
+ 0.5815e2,
+ 0.6377e2,
+ 0.6975e2,
+ 0.7614e2,
+ 0.8294e2,
+ 0.9018e2,
+ 0.9790e2,
+ 0.1061e3,
+ 0.1148e3,
+ 0.1241e3,
+ 0.1340e3,
+ 0.1445e3,
+ 0.1556e3,
+ 0.1674e3,
+ 0.1799e3,
+ 0.1932e3,
+ 0.2072e3,
+ 0.2221e3,
+ 0.2379e3,
+ 0.2546e3,
+ 0.2722e3,
+ 0.2909e3,
+ 0.3106e3,
+ 0.3314e3,
+ 0.3534e3,
+ 0.3766e3,
+ 0.4011e3,
+ 0.4269e3,
+ 0.4541e3,
+ 0.4827e3,
+ 0.5128e3,
+ 0.5445e3,
+ 0.5779e3,
+ 0.6130e3,
+ 0.6498e3,
+ 0.6885e3,
+ 0.7291e3,
+ 0.7717e3,
+ 0.8164e3,
+ 0.8633e3,
+ 0.9124e3,
+ 0.9638e3,
+ 0.1018e4,
+ 0.1074e4,
+ 0.1133e4,
+ 0.1194e4,
+ 0.1259e4,
+ 0.1326e4,
+ 0.1396e4,
+ 0.1469e4,
+ 0.1546e4,
+ 0.1625e4,
+ 0.1708e4,
+ 0.1794e4,
+ 0.1884e4,
+ 0.1978e4,
+ 0.2075e4,
+ 0.2176e4,
+ 0.2281e4,
+ 0.2390e4,
+ 0.2503e4,
+ 0.2620e4,
+ 0.2742e4,
+ 0.2868e4,
+ 0.2998e4,
+ 0.3134e4,
+ 0.3274e4,
+ 0.3418e4,
+ 0.3568e4,
+ 0.3723e4,
+ 0.3882e4,
+ 0.4047e4,
+ 0.4218e4,
+ 0.4393e4,
+ 0.4574e4,
+ 0.4761e4,
+ 0.4953e4,
+ 0.5151e4,
+ 0.5354e4,
+ 0.5564e4,
+ 0.5779e4,
+ 0.6000e4]
diff --git a/polaris/ocean/vertical/grid_1d/60layerPHC.json b/polaris/ocean/vertical/grid_1d/60layerPHC.json
new file mode 100644
index 000000000..4e3c53184
--- /dev/null
+++ b/polaris/ocean/vertical/grid_1d/60layerPHC.json
@@ -0,0 +1,63 @@
+[
+ 0.0,
+ 10.0,
+ 20.0,
+ 30.0,
+ 40.0,
+ 50.0,
+ 60.0,
+ 70.0,
+ 80.0,
+ 90.0,
+ 100.0,
+ 110.0,
+ 120.0,
+ 130.0,
+ 140.0,
+ 150.0,
+ 160.0,
+ 170.19677734375,
+ 180.76129150390625,
+ 191.82119750976562,
+ 203.49929809570312,
+ 215.92340087890625,
+ 229.23312377929688,
+ 243.58447265625,
+ 259.1557922363281,
+ 276.1524963378906,
+ 294.8147277832031,
+ 315.4236145019531,
+ 338.3122863769531,
+ 363.8746032714844,
+ 392.5804748535156,
+ 424.9887390136719,
+ 461.7665710449219,
+ 503.7067565917969,
+ 551.7491760253906,
+ 606.9965515136719,
+ 670.7285461425781,
+ 744.3980407714844,
+ 829.6069641113281,
+ 928.0434265136719,
+ 1041.3681945800781,
+ 1171.0402526855469,
+ 1318.0935363769531,
+ 1482.9008483886719,
+ 1664.9919738769531,
+ 1863.0146179199219,
+ 2074.873809814453,
+ 2298.039276123047,
+ 2529.903594970703,
+ 2768.098846435547,
+ 3010.670196533203,
+ 3256.138885498047,
+ 3503.448516845703,
+ 3751.892791748047,
+ 4001.011505126953,
+ 4250.525604248047,
+ 4500.259552001953,
+ 4750.121307373047,
+ 5000.045684814453,
+ 5250.010955810547,
+ 5499.989044189453
+]
diff --git a/polaris/ocean/vertical/grid_1d/80layerE3SMv1.json b/polaris/ocean/vertical/grid_1d/80layerE3SMv1.json
new file mode 100644
index 000000000..5b43c4e8c
--- /dev/null
+++ b/polaris/ocean/vertical/grid_1d/80layerE3SMv1.json
@@ -0,0 +1,81 @@
+[0.0,
+ 1.986786469355805,
+ 4.154119395541748,
+ 6.518473015615179,
+ 9.097824574809877,
+ 11.911737056573685,
+ 14.981608934309552,
+ 18.330839881040465,
+ 21.984828733039574,
+ 25.971389018482437,
+ 30.320829117011147,
+ 35.066198637149725,
+ 40.24353341642255,
+ 45.892098753893194,
+ 52.054713914204456,
+ 58.778073608483226,
+ 66.11298212263647,
+ 74.11483297053151,
+ 82.84399843035216,
+ 92.36612734933736,
+ 102.75259907640329,
+ 114.08104666296263,
+ 126.43577878360325,
+ 139.90834494484784,
+ 154.5979878835862,
+ 170.61222402211808,
+ 188.06721008633497,
+ 207.08846155392177,
+ 227.81108029827277,
+ 250.38029644815975,
+ 274.95164090923384,
+ 301.69113003347167,
+ 330.7752620963125,
+ 362.3906244929844,
+ 396.73330941558294,
+ 434.00792602635505,
+ 474.42608083399733,
+ 518.2042060869917,
+ 565.5607828175309,
+ 616.7129401252857,
+ 671.8720767795905,
+ 731.2388094472287,
+ 794.9973708678924,
+ 863.3093165271582,
+ 936.3071845957666,
+ 1014.0882623825473,
+ 1096.7088376639208,
+ 1184.1797979766206,
+ 1276.4637610648406,
+ 1373.474191183026,
+ 1475.0766739591154,
+ 1581.0924977739978,
+ 1691.3042518457269,
+ 1805.4629509798497,
+ 1923.2961967654483,
+ 2044.5166633553195,
+ 2168.83040784201,
+ 2295.9444441559367,
+ 2425.5731465435715,
+ 2557.443466443468,
+ 2691.2987229942432,
+ 2826.9011458061473,
+ 2964.033287582243,
+ 3102.4984729509856,
+ 3242.1205209147015,
+ 3382.742909836675,
+ 3524.227580019397,
+ 3666.4535034604037,
+ 3809.3151441356854,
+ 3952.7208868168414,
+ 4096.591510555103,
+ 4240.85873802501,
+ 4385.46388801116,
+ 4530.356647311017,
+ 4675.493968648996,
+ 4820.839083119443,
+ 4966.360630500305,
+ 5112.031898664707,
+ 5257.830157844073,
+ 5403.736084165608,
+ 5549.733263212151]

From c78288d68f2242120bde253fb9ea2d45d0b5ee9e Mon Sep 17 00:00:00 2001
From: Xylar Asay-Davis <xylarstorm@gmail.com>
Date: Tue, 31 Oct 2023 15:12:20 +0100
Subject: [PATCH 5/5] Update user's guide

---
 docs/users_guide/ocean/framework/vertical.md | 107 ++++++++++++++++---
 1 file changed, 95 insertions(+), 12 deletions(-)

diff --git a/docs/users_guide/ocean/framework/vertical.md b/docs/users_guide/ocean/framework/vertical.md
index 1cff8422b..2d87b98a6 100644
--- a/docs/users_guide/ocean/framework/vertical.md
+++ b/docs/users_guide/ocean/framework/vertical.md
@@ -24,6 +24,10 @@ min_layer_thickness = 3.0
 # The maximum layer thickness
 max_layer_thickness = 500.0
 
+# The characteristic number of levels over which the index_tanh_dz
+# transition between the min and max occurs
+transition_levels = 28
+
 # The type of vertical coordinate (e.g. z-level, z-star)
 coord_type = z-star
 
@@ -35,12 +39,12 @@ min_pc_fraction = 0.1
 ```
 
 The vertical coordinate is typically defined based on a 1D reference grid.
-Possible 1D grid types are: `uniform`, `tanh_dz`, `60layerPHC`, and
-`100layerE3SMv1`.
+Possible 1D grid types are: `uniform`, `tanh_dz`, `index_tanh_dz`, 
+`60layerPHC`, `80layerE3SMv1`, and `100layerE3SMv1`.
 
 The meaning of the config options `vert_levels`, `bottom_depth`,
-`min_layer_thickness` and `max_layer_thickness` depends grid type and is
-described below.
+`min_layer_thickness`, `max_layer_thickness`, and `transition_levels` depends 
+grid type and is described below.
 
 The options `coord_type`, `partial_cell_type` and `min_pc_fraction`
 relate to {ref}`ocean-vert-3d`, described below.
@@ -117,6 +121,61 @@ min_layer_thickness = 2.0
 max_layer_thickness = 210.0
 ```
 
+### index_tanh_dz
+
+This is similar to `tanh_dz` but the hyperbolic tangent function is defined
+in layer index space rather than physical depth. Layer thickness is defined by:
+
+$$
+\Delta z\left(k\right) = (\Delta z_2 - \Delta z_1)
+           \mathrm{tanh}\left[\frac{\pi \left(k - k_0\right)}{\Delta}\right]+ \Delta z_1,
+$$
+
+where $\Delta z_1$ (`min_layer_thickness`) is the value of the layer
+thickness $\Delta z$ at $z = 0$ and $\Delta z_2$
+(`max_layer_thickness`) is the same as $z \rightarrow \infty$.  The
+vertical layer index is $k$, $\Delta$ (`transition_levels`) is the number of
+vertical levels over which the `tanh` transitions from the finer to the
+coarser resolution, and $k_0$ is the origin in vertical index space of the
+transition. Interface locations $z_k$ are defined by:
+
+$$
+z_0 = 0, \\
+z_{k+1} = z_k - \Delta z\left(k\right).
+$$
+
+We use a root finder to solve for $k_0$, such that
+$z_{n_z+1} = -H_\mathrm{bot}$, where $n_z$ is `vert_levels`, the
+number of vertical levels (one less than the number of layer interfaces) and
+$H_\mathrm{bot}$ is `bottom_depth`, the depth of the seafloor.
+
+The following config options are all required.  This is an example of a
+64-layer vertical grid that has been explored in E3SM v3:
+
+```cfg
+# Options related to the vertical grid
+[vertical_grid]
+
+# the type of vertical grid
+grid_type = index_tanh_dz
+
+# Number of vertical levels
+vert_levels = 64
+
+# Depth of the bottom of the ocean
+bottom_depth = 5500.0
+
+# The minimum layer thickness
+min_layer_thickness = 10.0
+
+# The maximum layer thickness
+max_layer_thickness = 250.0
+
+# The characteristic number of levels over which the index_tanh_dz
+# transition between the min and max occurs
+transition_levels = 28
+```
+
 ### 60layerPHC
 
 This is the vertical grid used by the Polar science center Hydrographic Climatology
@@ -135,6 +194,22 @@ grid_type = 60layerPHC
 If the `bottom_depth` option is also defined, the depths will be renormalized
 so that bottom of the deepest layer is at `z = -bottom_depth`
 
+### 80layerE3SMv1
+
+This is the vertical grid was used in some E3SM v1 and v2 meshes, such as the
+ARRM10to60 mesh. Layer thicknesses vary over 80 layers from 2 m at the surface
+to 146 m at the seafloor, which is at 5550 m depth.  To get the default grid,
+use:
+```cfg
+# Options related to the vertical grid
+[vertical_grid]
+
+# the type of vertical grid
+grid_type = 80layerE3SMv1
+```
+If the `bottom_depth` option is also defined, the depths will be renormalized
+so that bottom of the deepest layer is at `z = -bottom_depth`.
+
 ### 100layerE3SMv1
 
 This is the vertical grid was used in some E3SM v1 experiments. Layer
@@ -171,13 +246,14 @@ with the sea floor at 2500 m.
 
 ## 3D vertical coordinates
 
-Currently, `z-star` and `z-level` vertical coordinates are supported
-(`coord_type`).  Each supports 3 options for `partial_cell_type`: `full`
-meaning the topography (bottom depth and sea-surface height) are expanded so
-that all layers have their full thickness; `partial` meaning that cells
-adjacent to the topography are allowed to be a small fraction of a full layer
-thickness; or `None` to indicate that no alteration is needed for full or
-partial cells (typically only in cases where the topography is already flat).
+Currently, `z-star`, `z-level` and `sigma` vertical coordinates are supported
+(`coord_type`).  The `z-star` and `z-level` types support 3 options for 
+`partial_cell_type`: `full` meaning the topography (bottom depth and 
+sea-surface height) are expanded so that all layers have their full thickness; 
+`partial` meaning that cells adjacent to the topography are allowed to be a 
+small fraction of a full layer thickness; or `None` to indicate that no 
+alteration is needed for full or partial cells (typically only in cases where 
+the topography is already flat).
 
 If `partial_cell_type = partial`, the option `min_pc_fraction` indicates
 the smallest fraction of a layer that a partial cell is allowed to have before
@@ -193,7 +269,7 @@ Most ocean tasks currently use the z\* vertical coordinate
 by default.  Typically (in the absence of ice-shelf cavities), the initial
 "resting" grid uses a {ref}`ocean-z-level` coordinate.  As the sea-surface
 height evolves, the coordinate stretches and squashes in proportion to changes
-in the local watercolumn thickness.
+in the local water-column thickness.
 
 In configurations with {ref}`ocean-ice-shelf-cavities`, the ice draft
 (elevation of the ice shelf-ocean interface) also acts the sea-surface height
@@ -214,3 +290,10 @@ parts of the mesh as "land" in the same way that cells below the batymetry are
 masked as land.  The topography at the top of the ocean is represented by
 "stair steps", using either "full" or "partial" cells to represent these steps
 in exactly the same way as at the seafloor.
+
+(ocean-sigma)=
+
+### sigma
+
+The `sigma` coordinate is a terrain-following coordinate that stretches a 1D
+reference coordinate between `z = -bottomDepth` and `z = ssh`.