define_vertical_grid_64_layers_xsad.py

#!/usr/bin/env python

from make_vertical_grid_xsad import create_vertical_grid

# Standard 64 layer vertical grid
create_vertical_grid(bottom_depth=5500, nz=64, dz1_in=2., dz2_in=200.,
                     plot_vertical_grid=True, outFile= 'vertical_grid_xsad.nc')

make_vertical_grid.py

#!/usr/bin/env python
"""
This script creates the vertical grid for MPAS-Ocean and writes it to a netcdf file.
"""
# import modules
# {{{
from netCDF4 import Dataset
import numpy as np
import argparse
from scipy.optimize import root_scalar
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
# }}}


def main():
    # parser
    # {{{
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        '-o', '--output_file_name', dest='output_filename_name',
        default='MPAS-Ocean_vertical_grid.nc',
        help='MPAS file name for output of vertical grid.',
        metavar='NAME')
    parser.add_argument(
        '-p', '--plot_vertical_grid', dest='plot_vertical_grid',
        action='store_true')
    parser.add_argument(
        '-bd', '--bottom_depth', dest='bottom_depth',
        default=5000.0,
        help='bottom depth for the chosen vertical coordinate [m]',
        type=float)
    parser.add_argument(
        '-nz', '--num_vert_levels', dest='nz',
        default=64,
        help='Number of vertical levels for the grid',
        type=int)
    parser.add_argument(
        '-dz1', '--layer1_thickness', dest='dz1',
        default=2.0,
        help='Target thickness of the first layer [m]',
        type=float)
    parser.add_argument(
        '-dz2', '--maxLayer_thickness', dest='dz2',
        default=250.0,
        help='Target maximum thickness in column [m]',
        type=float)
    args = parser.parse_args()

    create_vertical_grid(args.bottom_depth, args.nz, args.dz1,
                         args.dz2, args.plot_vertical_grid,
                         outFile=args.output_filename_name)
# }}}


def create_vertical_grid(
        bottom_depth=5000.0,
        nz=64,
        dz1_in=2.0,
        dz2_in=250.0,
        plot_vertical_grid=False,
        outFile='MPAS-Ocean_vertical_grid.nc'):
    # {{{
    """
    This function creates the vertical grid for MPAS-Ocean and writes it to a
    NetCDF file.

    Parameters
    ----------

    bottom_depth : float, optional
        bottom depth for the chosen vertical coordinate [m]

    nz : int, optional
        Number of vertical levels for the grid

    dz1_in : float, optional
        Target thickness of the first layer [m]

    dz2_in : float, optional
        Target maximum thickness in column [m]

    plot_vertical_grid : bool, optional
        Whether to plot the vertical grid

    outFile : str, optional
        MPAS file name for output of vertical grid
    """

    print('Creating mesh with ', nz, ' layers...')
    dz1 = dz1_in
    dz2 = dz2_in
    # open a new netCDF file for writing.
    ncfile = Dataset(outFile, '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',))

    sol = root_scalar(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)
    nVertLevels = nz
    botDepth = -z[1:]
    midDepth = -0.5*(z[0:-1] + z[1:])

    refBottomDepth[:] = botDepth
    refMidDepth[:] = midDepth
    refLayerThickness[:] = layerThickness[:nVertLevels]
    ncfile.close()

    if plot_vertical_grid:
        fig = plt.figure()
        fig.set_size_inches(16.0, 8.0)
        zInd = np.arange(1, nVertLevels + 1)
        plt.clf()

        plt.subplot(2, 2, 1)
        plt.plot(zInd, midDepth, '.')
        plt.gca().invert_yaxis()
        plt.xlabel('vertical index (one-based)')
        plt.ylabel('layer mid-depth [m]')
        plt.grid()

        plt.subplot(2, 2, 2)
        plt.plot(layerThickness, midDepth, '.')
        plt.gca().invert_yaxis()
        plt.xlabel('layer thickness [m]')
        plt.ylabel('layer mid-depth [m]')
        plt.grid()

        plt.subplot(2, 2, 3)
        plt.plot(zInd, layerThickness, '.')
        plt.xlabel('vertical index (one-based)')
        plt.ylabel('layer thickness [m]')
        plt.grid()

        txt = \
            'number layers: {}\n'.format(nz) + \
            'bottom depth requested:  {:8.2f}\n'.format(bottom_depth) +  \
            'bottom depth actual:     {:8.2f}\n'.format(np.amax(botDepth)) +  \
            'min thickness reqeusted: {:8.2f}\n'.format(dz1_in) + \
            'min thickness actual:    {:8.2f}\n'.format(np.amin(layerThickness[:])) + \
            'max thickness reqeusted: {:8.2f}\n'.format(dz2_in) + \
            'max thickness actual:    {:8.2f}'.format(np.amax(layerThickness[:]))
        print(txt)
        plt.subplot(2, 2, 4)
        plt.text(0, 0, txt, fontsize=12)
        plt.axis('off')
        plt.savefig('vertical_grid_xsad.png')

# }}}


def match_bottom(delta, nz, dz1, dz2, bottom_depth):
    _, z = cumsum_z(delta, nz, dz1, dz2)
    diff = -bottom_depth - z[-1]
    return diff


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)
        z[zindex+1] = z[zindex] - dz[zindex]
    return dz, z


def dz_z(z, dz1, dz2, delta):
    return (dz2 - dz1) * np.tanh(-z * np.pi / delta) + dz1


if __name__ == '__main__':
    # If called as a primary module, run main
    main()

# vim: foldmethod=marker ai ts=4 sts=4 et sw=4 ft=python

root_finding_algorithm.py

import numpy as np
from scipy.optimize import root_scalar

def create_vertical_grid(
        bottom_depth=5500.0,
        nz=64,
        dz1=2.0,
        dz2=200.0):
    sol = root_scalar(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)
    nVertLevels = nz
    botDepth = -z[1:]
    midDepth = -0.5*(z[0:-1] + z[1:])


def match_bottom(delta, nz, dz1, dz2, bottom_depth):
    _, z = cumsum_z(delta, nz, dz1, dz2)
    diff = -bottom_depth - z[-1]
    return diff


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)
        z[zindex+1] = z[zindex] - dz[zindex]
    return dz, z


def dz_z(z, dz1, dz2, delta):
    return (dz2 - dz1) * np.tanh(-z * np.pi / delta) + dz1