Add ocean turbulence test

compass/ocean/__init__.py

@@ -1,4 +1,5 @@
 from compass.mpas_core import MpasCore
+from compass.ocean.tests.turbulence_closure import TurbulenceClosure
 from compass.ocean.tests.baroclinic_channel import BaroclinicChannel
 from compass.ocean.tests.dam_break import DamBreak
 from compass.ocean.tests.drying_slope import DryingSlope
@@ -32,6 +33,7 @@ def __init__(self):
         """
         super().__init__(name='ocean')
 
+        self.add_test_group(TurbulenceClosure(mpas_core=self))
         self.add_test_group(BaroclinicChannel(mpas_core=self))
         self.add_test_group(DamBreak(mpas_core=self))
         self.add_test_group(DryingSlope(mpas_core=self))

compass/ocean/tests/turbulence_closure/__init__.py

@@ -0,0 +1,85 @@
+from compass.ocean.tests.turbulence_closure.decomp_test import DecompTest
+from compass.ocean.tests.turbulence_closure.default import Default
+from compass.ocean.tests.turbulence_closure.restart_test import RestartTest
+from compass.testgroup import TestGroup
+
+
+class TurbulenceClosure(TestGroup):
+    """
+    A test group for turbulence closure test cases
+    """
+    def __init__(self, mpas_core):
+        """
+        mpas_core : compass.MpasCore
+            the MPAS core that this test group belongs to
+        """
+        super().__init__(mpas_core=mpas_core, name='turbulence_closure')
+
+        for resolution in [1e4]:
+            self.add_test_case(
+                DecompTest(test_group=self, resolution=resolution))
+            self.add_test_case(
+                RestartTest(test_group=self, resolution=resolution))
+        for resolution in [1, 2, 1e4]:
+            for forcing in ['cooling', 'evaporation']:
+                self.add_test_case(
+                    Default(test_group=self, resolution=resolution,
+                            forcing=forcing))
+
+
+def configure(resolution, forcing, config):
+    """
+    Modify the configuration options for one of the turbulence closure test
+    cases
+
+    Parameters
+    ----------
+    resolution : float
+        The resolution of the test case in meters
+
+    config : configparser.ConfigParser
+        Configuration options for this test case
+    """
+    # The resolution parameters are different for different resolutions
+    # to match existing simulations
+    if resolution > 1e3:
+        nx = 16
+        ny = 50
+        vert_levels = 20
+        bottom_depth = 1e3
+    elif resolution <= 1e3 and resolution > 5:
+        nx = 50
+        ny = 50
+        vert_levels = 50
+        bottom_depth = 100.0
+    elif resolution <= 5 and resolution > 1:
+        nx = 150
+        ny = 150
+        vert_levels = 50
+        bottom_depth = 100.0
+    elif resolution <= 1:
+        nx = 128
+        ny = 128
+        vert_levels = 128
+        bottom_depth = 128.0
+
+    config.set('turbulence_closure', 'nx', f'{nx}')
+    config.set('turbulence_closure', 'ny', f'{ny}')
+    config.set('turbulence_closure', 'dc', f'{resolution}')
+    config.set('vertical_grid', 'vert_levels', f'{vert_levels}')
+    config.set('vertical_grid', 'bottom_depth', f'{bottom_depth}')
+
+    if forcing == 'cooling':
+        config.set('turbulence_closure', 'surface_heat_flux', '-100')
+        config.set('turbulence_closure', 'surface_freshwater_flux', '0')
+        config.set('turbulence_closure', 'interior_temperature_gradient',
+                   '0.1')
+        config.set('turbulence_closure', 'interior_salinity_gradient', '0')
+        config.set('turbulence_closure', 'wind_stress_zonal', '0')
+    if forcing == 'evaporation':
+        config.set('turbulence_closure', 'surface_heat_flux', '0')
+        config.set('turbulence_closure', 'surface_freshwater_flux', '0.429')
+        config.set('turbulence_closure', 'interior_temperature_gradient', '0')
+        config.set('turbulence_closure', 'interior_salinity_gradient',
+                   '-0.025')
+        config.set('turbulence_closure', 'wind_stress_zonal', '0')

compass/ocean/tests/turbulence_closure/boundary_values.py

@@ -0,0 +1,210 @@
+import numpy
+
+
+def add_boundary_arrays(mesh, x_nonperiodic, y_nonperiodic):  # noqa: C901
+    """
+    Computes boundary arrays necessary for LES test cases
+    Values are appended to a MPAS mesh dataset
+
+    if the domain is doubly periodic, returns array of all zeros
+    as arrays are meaningless in that case.
+
+    Parameters
+    ----------
+    mesh : xarray dataset containing an MPAS mesh
+
+    x_nonperiodic : logical if x direction is non periodic
+
+    y_nonperiodic : logical if y direction is non periodic
+    """
+
+    nCells = mesh.dims['nCells']
+    neonc = mesh.nEdgesOnCell.values
+    eonc = mesh.edgesOnCell.values - 1
+    conc = mesh.cellsOnCell.values - 1
+    nVertLevels = mesh.dims['nVertLevels']
+    mlc = mesh.maxLevelCell.values - 1
+    cone = mesh.cellsOnEdge.values - 1
+
+    if (not x_nonperiodic) and (not y_nonperiodic):
+        mesh['distanceToBoundary'] = (
+            ['nCells', 'nVertLevels'],
+            numpy.zeros((nCells, nVertLevels)))
+        mesh['boundaryNormalAngle'] = (
+            ['nCells', 'nVertLevels'],
+            numpy.zeros((nCells, nVertLevels)))
+        mesh['boundaryCellMask'] = (
+            ['nCells', 'nVertLevels'],
+            numpy.zeros((nCells, nVertLevels)).astype(int))
+        return mesh
+
+    # surface first and save until later
+    indVals = []
+    angle = mesh.angleEdge.values
+    xc = mesh.xEdge.values
+    yc = mesh.yEdge.values
+    xCell = mesh.xCell.values
+    yCell = mesh.yCell.values
+    aCell = numpy.zeros(nCells)
+    bCell = numpy.zeros(nCells)
+    xEdge = []
+    yEdge = []
+    for i in range(nCells):
+        if conc[i, :].min() < 0:
+            aEdge = 0
+            count = 0
+            for j in range(neonc[i]):
+                if conc[i, j] < 0:
+                    indVals.append(i)
+                    xEdge.append(xc[eonc[i, j]])
+                    yEdge.append(yc[eonc[i, j]])
+                    # compute dx, dy
+                    # check angle edge
+                    # add pi if needed
+                    # check if over 2pi correct if needed
+                    dx = xc[eonc[i, j]] - xCell[i]
+                    dy = yc[eonc[i, j]] - yCell[i]
+                    if dx <= 0 and dy <= 0:  # first quadrant
+                        if angle[eonc[i, j]] >= numpy.pi / 2.:
+                            aEdge += angle[eonc[i, j]] - numpy.pi
+                            count += 1
+                        else:
+                            aEdge += angle[eonc[i, j]]
+                            count += 1
+
+                    # fourth quadrant, but reference to 0 not 2pi
+                    elif dx <= 0 and dy >= 0:
+                        if angle[eonc[i, j]] < 3.0 * numpy.pi / 2.0:
+                            aEdge += angle[eonc[i, j]] + numpy.pi
+                            count += 1
+                        else:
+                            aEdge += angle[eonc[i, j]]
+                            count += 1
+
+                    elif dx >= 0 and dy >= 0:  # third quadrant
+                        if angle[eonc[i, j]] < numpy.pi / 2.0:
+                            # not in correct place
+                            aEdge += angle[eonc[i, j]] + numpy.pi
+                            count += 1
+                        else:
+                            aEdge += angle[eonc[i, j]]
+                            count += 1
+
+                    else:  # quadrant 2
+                        if angle[eonc[i, j]] > 3.0 * numpy.pi / 2.0:
+                            # wrong place
+                            aEdge += angle[eonc[i, j]] - numpy.pi
+                            count += 1
+                        else:
+                            aEdge += angle[eonc[i, j]]
+                            count += 1
+            if count > 0:
+                aCellTemp = aEdge / count
+                if aCellTemp > numpy.pi:
+                    aCellTemp = 2.0 * numpy.pi - aCellTemp
+                    aCell[i] = aCellTemp
+                    bCell[i] = 1
+
+    # with angle and index arrays, now can do distance
+    dist = numpy.zeros((nCells, nVertLevels))
+    angleCells = numpy.zeros((nCells, nVertLevels))
+    boundaryCells = numpy.zeros((nCells, nVertLevels))
+    for i in range(nCells):
+        d = numpy.sqrt((xCell[i] - xEdge)**2 + (yCell[i] - yEdge)**2)
+        ind = d.argmin()
+        angleCells[i, 0] = aCell[indVals[ind]]
+        boundaryCells[i, 0] = bCell[i]
+        dist[i, 0] = d.min()
+
+    # now to the rest
+    for k in range(1, nVertLevels):
+        inds = numpy.where(k == mlc)[0]
+        edgeList = []
+        cellList = []
+        if len(inds) > 0:
+            # First step is forming list of edges bordering maxLC
+            for i in range(len(inds)):
+                for j in range(neonc[inds[i]]):
+                    if conc[inds[i], j] not in inds:
+                        edgeList.append(eonc[inds[i], j])
+            for i in range(len(edgeList)):
+                c1 = cone[edgeList[i], 0]
+                c2 = cone[edgeList[i], 1]
+                if c1 in inds:
+                    if c2 not in cellList:
+                        cellList.append(c2)
+                if c2 in inds:
+                    if c1 not in cellList:
+                        cellList.append(c1)
+
+            for i in range(len(cellList)):
+                aEdge = 0
+                count = 0
+                for j in range(neonc[cellList[i]]):
+                    if eonc[cellList[i], j] in edgeList:
+                        indVals.append(cellList[i])
+                        xEdge.append(xc[eonc[cellList[i], j]])
+                        yEdge.append(yc[eonc[cellList[i], j]])
+                        # compute dx, dy
+                        # check angle edge
+                        # add pi if needed
+                        # check if over 2pi correct if needed
+                        dx = xc[eonc[cellList[i], j]] - xCell[cellList[i]]
+                        dy = yc[eonc[cellList[i], j]] - yCell[cellList[i]]
+                        if dx <= 0 and dy <= 0:  # first quadrant
+                            if angle[eonc[cellList[i], j]] >= numpy.pi / 2.:
+                                aEdge += angle[eonc[cellList[i], j]] - numpy.pi
+                                count += 1
+                            else:
+                                aEdge += angle[eonc[cellList[i], j]]
+                                count += 1
+                        # fourth quadrant, but reference to 0 not 2pi
+                        elif dx <= 0 and dy >= 0:
+                            if {(angle[eonc[cellList[i], j]] <
+                                 3.0 * numpy.pi / 2.0)}:
+                                aEdge += angle[eonc[cellList[i], j]] + numpy.pi
+                                count += 1
+                            else:
+                                aEdge += angle[eonc[cellList[i], j]]
+                                count += 1
+                        elif dx >= 0 and dy >= 0:  # third quadrant
+                            # not in correct place
+                            if angle[eonc[cellList[i], j]] < numpy.pi / 2.0:
+                                aEdge += angle[eonc[cellList[i], j]] + numpy.pi
+                                count += 1
+                            else:
+                                aEdge += angle[eonc[cellList[i], j]]
+                                count += 1
+                        else:  # quadrant 2
+                            # wrong place
+                            if {(angle[eonc[cellList[i], j]] >
+                                 3.0 * numpy.pi / 2.0)}:
+                                aEdge += angle[eonc[cellList[i], j]] - numpy.pi
+                                count += 1
+                            else:
+                                aEdge += angle[eonc[cellList[i], j]]
+                                count += 1
+                if count > 0:
+                    aCellTemp = aEdge / count
+                    if aCellTemp > numpy.pi:
+                        aCellTemp = 2.0 * numpy.pi - aCellTemp
+                        aCell[cellList[i]] = aCellTemp
+                        bCell[cellList[i]] = 1
+
+                for ii in range(nCells):
+                    d = numpy.sqrt((xCell[ii] - xEdge)**2 +
+                                   (yCell[ii] - yEdge)**2)
+                    ind = d.argmin()
+                    angleCells[ii, k] = aCell[indVals[ind]]
+                    boundaryCells[ii, k] = bCell[ii]
+                    dist[ii, k] = d.min()
+        else:
+            dist[:, k] = dist[:, 0]
+            angleCells[:, k] = angleCells[:, 0]
+            boundaryCells[:, k] = boundaryCells[:, 0]
+    mesh['distanceToBoundary'] = (['nCells', 'nVertLevels'], dist)
+    mesh['boundaryNormalAngle'] = (['nCells', 'nVertLevels'], angleCells)
+    mesh['boundaryCellMask'] = (['nCells', 'nVertLevels'],
+                                boundaryCells.astype(int))
+
+    return mesh

compass/ocean/tests/turbulence_closure/decomp_test/__init__.py

@@ -0,0 +1,67 @@
+from compass.testcase import TestCase
+from compass.ocean.tests.turbulence_closure.initial_state import InitialState
+from compass.ocean.tests.turbulence_closure.forward import Forward
+from compass.ocean.tests import turbulence_closure
+from compass.validate import compare_variables
+
+
+class DecompTest(TestCase):
+    """
+    A decomposition test case for the turbulence closure test group, which
+    makes sure the model produces identical results on 1 and 4 cores.
+
+    Attributes
+    ----------
+    resolution : float
+        The resolution of the test case in meters
+    """
+
+    def __init__(self, test_group, resolution, forcing='cooling'):
+        """
+        Create the test case
+
+        Parameters
+        ----------
+        test_group : compass.ocean.tests.turbulence_closure.TurbulenceClosure
+            The test group that this test case belongs to
+
+        resolution : float
+            The resolution of the test case in meters
+        """
+        name = 'decomp_test'
+        self.resolution = resolution
+        if resolution >= 1e3:
+            res_name = f'{int(resolution/1e3)}km'
+        else:
+            res_name = f'{int(resolution)}m'
+        subdir = f'{res_name}/{forcing}/{name}'
+        super().__init__(test_group=test_group, name=name,
+                         subdir=subdir)
+
+        self.add_step(
+            InitialState(test_case=self, resolution=resolution))
+
+        for procs in [4, 8]:
+            name = '{}proc'.format(procs)
+            self.add_step(
+                Forward(test_case=self, name=name, subdir=name, ntasks=procs,
+                        openmp_threads=1, resolution=resolution))
+
+    def configure(self):
+        """
+        Modify the configuration options for this test case.
+        """
+        turbulence_closure.configure(self.resolution, self.config)
+
+    # no run() method is needed
+
+    def validate(self):
+        """
+        Test cases can override this method to perform validation of variables
+        and timers
+        """
+        variables = ['temperature', 'salinity', 'layerThickness',
+                     'normalVelocity']
+        compare_variables(test_case=self, variables=variables,
+                          filename1='4proc/output.nc',
+                          filename2='8proc/output.nc')