Merge pull request #547 from alicebarthel/add-mitgcm-baroclinic-gyre

Add baroclinic gyre tests

compass/ocean/__init__.py

@@ -1,5 +1,6 @@
 from compass.mpas_core import MpasCore
 from compass.ocean.tests.baroclinic_channel import BaroclinicChannel
+from compass.ocean.tests.baroclinic_gyre import BaroclinicGyre
 from compass.ocean.tests.buttermilk_bay import ButtermilkBay
 from compass.ocean.tests.dam_break import DamBreak
 from compass.ocean.tests.drying_slope import DryingSlope
@@ -38,6 +39,7 @@ def __init__(self):
         super().__init__(name='ocean')
 
         self.add_test_group(BaroclinicChannel(mpas_core=self))
+        self.add_test_group(BaroclinicGyre(mpas_core=self))
         self.add_test_group(ButtermilkBay(mpas_core=self))
         self.add_test_group(DamBreak(mpas_core=self))
         self.add_test_group(DryingSlope(mpas_core=self))

compass/ocean/tests/baroclinic_gyre/__init__.py

@@ -0,0 +1,22 @@
+from compass.ocean.tests.baroclinic_gyre.gyre_test_case import GyreTestCase
+from compass.testgroup import TestGroup
+
+
+class BaroclinicGyre(TestGroup):
+    """
+    A test group for baroclinic gyre 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='baroclinic_gyre')
+
+        for resolution in [20000., 80000.]:
+            self.add_test_case(
+                GyreTestCase(test_group=self, resolution=resolution,
+                             long=False))
+            self.add_test_case(
+                GyreTestCase(test_group=self, resolution=resolution,
+                             long=True))

compass/ocean/tests/baroclinic_gyre/baroclinic_gyre.cfg

@@ -0,0 +1,57 @@
+
+# Options related to the vertical grid
+[vertical_grid]
+
+# the type of vertical grid
+grid_type = linear_dz
+
+# the linear rate of thickness (m) increase for linear_dz
+linear_dz_rate = 10.
+
+# Number of vertical levels
+vert_levels = 15
+
+# Total water column depth in m
+bottom_depth = 1800.
+
+# The type of vertical coordinate (e.g. z-level, z-star)
+coord_type = z-star
+
+# Whether to use "partial" or "full", or "None" to not alter the topography
+partial_cell_type = None
+
+# The minimum fraction of a layer for partial cells
+min_pc_fraction = 0.1
+
+
+# config options for the baroclinic gyre
+[baroclinic_gyre]
+# Basin dimensions
+lat_min = 15
+lat_max = 75
+lon_min = 0
+lon_max = 60
+
+# Initial vertical temperature profile (C)
+initial_temp_top = 33.
+initial_temp_bot = 1.
+
+# Constant salinity value (also used in restoring)
+initial_salinity = 34.
+
+# Maximum zonal wind stress value (N m-2)
+wind_stress_max = 0.1
+
+# Surface temperature restoring profile
+restoring_temp_min = 0.
+restoring_temp_max = 30.
+
+# Restoring timescale for surface temperature (in days)
+restoring_temp_timescale = 30.
+
+# config options for the post processing (moc and viz)
+[baroclinic_gyre_post]
+# latitude bin increment for the moc calculation
+dlat = 0.25
+# number of years to average over for the mean state plots
+time_averaging_length = 1

compass/ocean/tests/baroclinic_gyre/cull_mesh.py

@@ -0,0 +1,61 @@
+import numpy as np
+import xarray as xr
+from mpas_tools.io import write_netcdf
+from mpas_tools.mesh.conversion import convert, cull
+
+from compass.step import Step
+
+
+class CullMesh(Step):
+    """
+    Cull a global mesh to only a single basin
+    Attributes
+    ----------
+    """
+    def __init__(self, test_case):
+        """
+        Create the step
+        Parameters
+        ----------
+        test_case : compass.TestCase
+            The test case this step belongs to
+        """
+        super().__init__(test_case=test_case, name='cull_mesh')
+
+        self.add_input_file(
+            filename='base_mesh.nc',
+            target='../base_mesh/base_mesh.nc')
+
+        for file in ['culled_mesh.nc', 'culled_graph.info']:
+            self.add_output_file(file)
+
+    def run(self):
+        """
+        Run this step of the test case
+        """
+        config = self.config
+        section = config['baroclinic_gyre']
+        logger = self.logger
+        ds_mesh = xr.open_dataset('base_mesh.nc')
+        ds_mask = xr.Dataset()
+
+        lon = np.rad2deg(ds_mesh.lonCell.values)
+        lat = np.rad2deg(ds_mesh.latCell.values)
+        lon_min = section.getfloat('lon_min')
+        lon_max = section.getfloat('lon_max')
+        lat_min = section.getfloat('lat_min')
+        lat_max = section.getfloat('lat_max')
+
+        mask = np.logical_and(
+            np.logical_and(lon >= lon_min, lon <= lon_max),
+            np.logical_and(lat >= lat_min, lat <= lat_max))
+
+        n_cells = ds_mesh.sizes['nCells']
+        ds_mask['regionCellMasks'] = (('nCells', 'nRegions'),
+                                      mask.astype(int).reshape(n_cells, 1))
+
+        ds_mesh = cull(ds_mesh, dsInverse=ds_mask, logger=logger)
+        ds_mesh = convert(ds_mesh, graphInfoFileName='culled_graph.info',
+                          logger=logger)
+
+        write_netcdf(ds_mesh, 'culled_mesh.nc')

compass/ocean/tests/baroclinic_gyre/forward.py

@@ -0,0 +1,120 @@
+from compass.model import run_model
+from compass.step import Step
+
+
+class Forward(Step):
+    """
+    A step for performing forward MPAS-Ocean runs as part of
+    the baroclinic gyre test cases.
+
+    Attributes
+    ----------
+    resolution : float
+        The resolution of the test case (m)
+
+    """
+    def __init__(self, test_case, resolution, name='forward', subdir=None,
+                 long=False):
+        """
+        Create a new test case
+
+        Parameters
+        ----------
+        test_case : compass.TestCase
+            The test case this step belongs to
+
+        resolution : float
+            The resolution of the test case (m)
+
+        name : str
+            the name of the test case
+
+        subdir : str, optional
+            the subdirectory for the step.  The default is ``name``
+
+        long : bool, optional
+            Whether to run a long (3-year) simulation to quasi-equilibrium
+        """
+        self.long = long
+        self.resolution = resolution
+        if resolution >= 1e3:
+            res_name = f'{int(resolution/1e3)}km'
+        else:
+            res_name = f'{int(resolution)}m'
+
+        res_params = {'80km': {'ntasks': 32,
+                               'min_tasks': 3,
+                               'dt': "'01:00:00'",
+                               'btr_dt': "'00:03:00'",
+                               'mom_del4': "5.0e11",
+                               'run_duration': "'0000_03:00:00'"},
+                      '20km': {'ntasks': 384,
+                               'min_tasks': 32,
+                               'dt': "'00:20:00'",
+                               'btr_dt': "'0000_00:00:20'",
+                               'mom_del4': "2.0e10 ",
+                               'run_duration': "'0000_01:00:00'"}}
+
+        if res_name not in res_params:
+            raise ValueError(
+                f'Unsupported resolution {res_name}. Supported values are: '
+                f'{list(res_params)}')
+
+        res_params = res_params[res_name]
+
+        ntasks = res_params['ntasks']
+        min_tasks = res_params['min_tasks']
+
+        super().__init__(test_case=test_case, name=name, subdir=subdir,
+                         ntasks=ntasks, min_tasks=min_tasks, openmp_threads=1)
+
+        # make sure output is double precision
+        self.add_streams_file('compass.ocean.streams', 'streams.output')
+
+        self.add_namelist_file('compass.ocean.tests.baroclinic_gyre',
+                               'namelist.forward')
+        if long:
+            output_interval = "0000-01-00_00:00:00"
+            restart_interval = "0001-00-00_00:00:00"
+        else:
+            output_interval = res_params['run_duration'].replace("'", "")
+            restart_interval = "0030_00:00:00"
+        replacements = dict(
+            output_interval=output_interval, restart_interval=restart_interval)
+        self.add_streams_file(package='compass.ocean.tests.baroclinic_gyre',
+                              streams='streams.forward',
+                              template_replacements=replacements)
+        options = dict()
+        for option in ['dt', 'btr_dt', 'mom_del4', 'run_duration']:
+            options[f'config_{option}'] = res_params[option]
+        if long:
+            # run for 3 years instead of 3 time steps
+            options['config_run_duration'] = "'0003-00-00_00:00:00'"
+
+        self.add_namelist_options(options=options)
+        self.add_input_file(filename='init.nc',
+                            target='../initial_state/initial_state.nc')
+        self.add_input_file(filename='forcing.nc',
+                            target='../initial_state/forcing.nc')
+        self.add_input_file(filename='graph.info',
+                            target='../initial_state/culled_graph.info')
+
+        self.add_model_as_input()
+
+        self.add_output_file(filename='output/output.0001-01-01_00.00.00.nc')
+        if long:
+            self.add_output_file(filename='output/'
+                                 'timeSeriesStatsMonthly.0001-01-01.nc')
+
+    # no setup() is needed
+
+    def run(self):
+        """
+        Run this step of the test case
+        """
+        run_model(self, partition_graph=True)
+
+        if self.long:
+            replacements = {'config_do_restart': '.true.',
+                            'config_start_time': "'file'"}
+            self.update_namelist_at_runtime(replacements)

compass/ocean/tests/baroclinic_gyre/gyre_test_case.py

@@ -0,0 +1,81 @@
+from compass.mesh import QuasiUniformSphericalMeshStep
+from compass.ocean.tests.baroclinic_gyre.cull_mesh import CullMesh
+from compass.ocean.tests.baroclinic_gyre.forward import Forward
+from compass.ocean.tests.baroclinic_gyre.initial_state import InitialState
+from compass.ocean.tests.baroclinic_gyre.moc import Moc
+from compass.ocean.tests.baroclinic_gyre.viz import Viz
+from compass.testcase import TestCase
+from compass.validate import compare_variables
+
+
+class GyreTestCase(TestCase):
+    """
+    A class to define the baroclinic gyre test cases
+
+    Attributes
+    ----------
+    resolution : float
+        The resolution of the test case (m)
+    """
+
+    def __init__(self, test_group, resolution, long):
+        """
+        Create the test case
+
+        Parameters
+        ----------
+        test_group :
+        compass.ocean.tests.baroclinic_gyre.BaroclinicGyre
+            The test group that this test case belongs to
+
+        resolution : float
+            The resolution of the test case (m)
+
+        long : bool
+            Whether to run a long (3-year) simulation to quasi-equilibrium
+        """
+        name = 'performance_test'
+        self.resolution = resolution
+        self.long = long
+
+        if long:
+            name = '3_year_test'
+
+        if resolution >= 1e3:
+            res_name = f'{int(resolution/1e3)}km'
+        else:
+            res_name = f'{int(resolution)}m'
+        subdir = f'{res_name}/{name}'
+        super().__init__(test_group=test_group, name=name,
+                         subdir=subdir)
+
+        self.add_step(QuasiUniformSphericalMeshStep(
+            test_case=self, cell_width=int(resolution / 1e3)))
+        self.add_step(CullMesh(test_case=self))
+        self.add_step(
+            InitialState(test_case=self, resolution=resolution))
+        self.add_step(
+            Forward(test_case=self, resolution=resolution,
+                    long=long))
+        if long:
+            self.add_step(
+                Moc(test_case=self, resolution=resolution))
+            self.add_step(
+                Viz(test_case=self, resolution=resolution))
+
+    def configure(self):
+        """
+        Set config options for the test case
+        """
+        config = self.config
+        config.add_from_package('compass.mesh', 'mesh.cfg')
+
+    def validate(self):
+        """
+        Validate variables against a baseline
+        """
+        compare_variables(test_case=self,
+                          variables=['layerThickness', 'temperature',
+                                     'ssh'],
+                          filename1='forward/output/'
+                                    'output.0001-01-01_00.00.00.nc')