Merge pull request #738 from gcapodag/LTS_WD_PR

Add LTS option to dam_break and parabolic_bowl test cases

This PR enhances the dam_break and parabolic_bowl test cases with the option to run using the local time-stepping (LTS) time integrator. Note that LTS only works in a single layer setting. Please also see the companion PR E3SM-Project/E3SM#6074

compass/ocean/suites/lts.txt

@@ -0,0 +1,9 @@
+ocean/dam_break/40cm/default_lts
+ocean/dam_break/40cm/ramp_lts
+ocean/dam_break/120cm/default_lts
+ocean/dam_break/120cm/ramp_lts
+ocean/hurricane/DEQU120at30cr10rr2/mesh_lts
+ocean/hurricane/DEQU120at30cr10rr2/init_lts
+ocean/hurricane/DEQU120at30cr10rr2/sandy_lts
+ocean/parabolic_bowl/standard/ramp_lts
+ocean/parabolic_bowl/standard/noramp_lts

compass/ocean/suites/wetdry.txt

@@ -20,6 +20,14 @@ ocean/dam_break/40cm/default
 ocean/dam_break/40cm/ramp
 ocean/dam_break/120cm/default
 ocean/dam_break/120cm/ramp
+ocean/dam_break/40cm/default_lts
+ocean/dam_break/40cm/ramp_lts
+ocean/dam_break/120cm/default_lts
+ocean/dam_break/120cm/ramp_lts
+ocean/parabolic_bowl/standard/ramp
+ocean/parabolic_bowl/standard/noramp
+ocean/parabolic_bowl/standard/ramp_lts
+ocean/parabolic_bowl/standard/noramp_lts
 ocean/isomip_plus/planar/5km/sigma/thin_film_drying_Ocean0
 ocean/isomip_plus/planar/5km/sigma/thin_film_wetting_Ocean0
 ocean/isomip_plus/planar/5km/single_layer/thin_film_tidal_forcing_Ocean0

compass/ocean/tests/dam_break/__init__.py

@@ -16,7 +16,10 @@ def __init__(self, mpas_core):
         super().__init__(mpas_core=mpas_core, name='dam_break')
 
         for resolution in [0.04, 0.12]:
-            self.add_test_case(
-                Default(test_group=self, resolution=resolution))
-            self.add_test_case(
-                Ramp(test_group=self, resolution=resolution))
+            for use_lts in [True, False]:
+                self.add_test_case(
+                    Default(test_group=self, resolution=resolution,
+                            use_lts=use_lts))
+                self.add_test_case(
+                    Ramp(test_group=self, resolution=resolution,
+                         use_lts=use_lts))

compass/ocean/tests/dam_break/dam_break.cfg

@@ -1,5 +0,0 @@
-# Options related to the vertical grid
-[vertical_grid]
-
-# Number of vertical levels # redundant with config_dam_break_vert_levels
-vert_levels = 3

compass/ocean/tests/dam_break/default/__init__.py

@@ -1,8 +1,10 @@
 from math import floor
-from compass.testcase import TestCase
-from compass.ocean.tests.dam_break.initial_state import InitialState
+
 from compass.ocean.tests.dam_break.forward import Forward
+from compass.ocean.tests.dam_break.initial_state import InitialState
+from compass.ocean.tests.dam_break.lts.lts_regions import LTSRegions
 from compass.ocean.tests.dam_break.viz import Viz
+from compass.testcase import TestCase
 from compass.validate import compare_variables
 
 
@@ -17,7 +19,7 @@ class Default(TestCase):
 
     """
 
-    def __init__(self, test_group, resolution):
+    def __init__(self, test_group, resolution, use_lts):
         """
         Create the test case
 
@@ -29,22 +31,36 @@ def __init__(self, test_group, resolution):
         resolution : float
             The resolution of the test case in m
 
+        use_lts : bool
+            Whether local time-stepping is used
+
         """
-        name = 'default'
+        self.use_lts = use_lts
+
+        if use_lts:
+            name = 'default_lts'
+        else:
+            name = 'default'
 
         self.resolution = resolution
         if resolution < 1.:
             res_name = f'{int(resolution*1e3)}cm'
         else:
             res_name = f'{int(resolution)}m'
-        min_tasks = int(40/(resolution/0.04)**2)
-        ntasks = 10*min_tasks
+        min_tasks = int(40 / (resolution / 0.04) ** 2)
+        ntasks = 10 * min_tasks
         subdir = f'{res_name}/{name}'
         super().__init__(test_group=test_group, name=name,
                          subdir=subdir)
 
-        self.add_step(InitialState(test_case=self))
+        init_step = InitialState(test_case=self, use_lts=use_lts)
+        self.add_step(init_step)
+
+        if use_lts:
+            self.add_step(LTSRegions(test_case=self, init_step=init_step))
+
         self.add_step(Forward(test_case=self, resolution=resolution,
+                              use_lts=use_lts,
                               ntasks=ntasks, min_tasks=min_tasks,
                               openmp_threads=1))
         self.add_step(Viz(test_case=self))
@@ -56,11 +72,12 @@ def configure(self):
 
         resolution = self.resolution
         config = self.config
+        use_lts = self.use_lts
         dc = resolution  # cell width in m
         dx = 13          # width of the domain in m
         dy = 28          # length of the domain in m
-        nx = round(dx/dc)
-        ny = int(2*floor(dy/(2*dc)))  # guarantee that ny is even
+        nx = round(dx / dc)
+        ny = int(2 * floor(dy / (2 * dc)))  # guarantee that ny is even
 
         config.set('dam_break', 'nx', f'{nx}', comment='the number of '
                    'mesh cells in the x direction')
@@ -69,10 +86,19 @@ def configure(self):
         config.set('dam_break', 'dc', f'{dc}', comment='the distance '
                    'between adjacent cell centers')
 
+        if use_lts:
+            vert_levels = 1
+            config.set('vertical_grid', 'vert_levels', f'{vert_levels}',
+                       comment='number of vertical levels')
+        else:
+            vert_levels = 3
+            config.set('vertical_grid', 'vert_levels', f'{vert_levels}',
+                       comment='number of vertical levels')
+
     def validate(self):
         """
         Validate variables against a baseline
         """
         variables = ['layerThickness', 'normalVelocity', 'ssh']
         compare_variables(test_case=self, variables=variables,
-                          filename1=f'forward/output.nc')
+                          filename1=f'{"forward/output.nc"}')

compass/ocean/tests/dam_break/forward.py

@@ -7,7 +7,8 @@ class Forward(Step):
     A step for performing forward MPAS-Ocean runs as part of dam break
     test cases.
     """
-    def __init__(self, test_case, resolution, name='forward', subdir=None,
+    def __init__(self, test_case, resolution, use_lts,
+                 name='forward', subdir=None,
                  ntasks=1, min_tasks=None, openmp_threads=1):
         """
         Create a new test case
@@ -20,6 +21,9 @@ def __init__(self, test_case, resolution, name='forward', subdir=None,
         resolution : str
             The resolution of the test case
 
+        use_lts : bool
+            Whether local time-stepping is used
+
         name : str
             the name of the test case
 
@@ -50,18 +54,36 @@ def __init__(self, test_case, resolution, name='forward', subdir=None,
 
         self.add_namelist_file('compass.ocean.tests.dam_break',
                                'namelist.forward')
-        self.add_streams_file('compass.ocean.tests.dam_break',
-                              'streams.forward')
-
-        input_path = '../initial_state'
-        self.add_input_file(filename='mesh.nc',
-                            target=f'{input_path}/culled_mesh.nc')
-
-        self.add_input_file(filename='init.nc',
-                            target=f'{input_path}/ocean.nc')
 
-        self.add_input_file(filename='graph.info',
-                            target=f'{input_path}/culled_graph.info')
+        if use_lts:
+            self.add_namelist_options(
+                {'config_time_integrator': "'LTS'",
+                 'config_dt_scaling_LTS': "4",
+                 'config_number_of_time_levels': "4",
+                 'config_pressure_gradient_type': "'ssh_gradient'"})
+
+            self.add_streams_file('compass.ocean.tests.dam_break.lts',
+                                  'streams.forward')
+            input_path = '../lts_regions'
+            self.add_input_file(filename='mesh.nc',
+                                target=f'{input_path}/lts_mesh.nc')
+            self.add_input_file(filename='graph.info',
+                                target=f'{input_path}/lts_graph.info')
+            self.add_input_file(filename='init.nc',
+                                target=f'{input_path}/lts_ocean.nc')
+
+        else:
+            self.add_streams_file('compass.ocean.tests.dam_break',
+                                  'streams.forward')
+            input_path = '../initial_state'
+            self.add_input_file(filename='mesh.nc',
+                                target=f'{input_path}/culled_mesh.nc')
+
+            self.add_input_file(filename='init.nc',
+                                target=f'{input_path}/ocean.nc')
+
+            self.add_input_file(filename='graph.info',
+                                target=f'{input_path}/culled_graph.info')
 
         self.add_model_as_input()
 

compass/ocean/tests/dam_break/initial_state.py

@@ -1,21 +1,18 @@
 import xarray
-import numpy
-import matplotlib.pyplot as plt
-
-from mpas_tools.planar_hex import make_planar_hex_mesh
 from mpas_tools.io import write_netcdf
 from mpas_tools.mesh.conversion import convert, cull
+from mpas_tools.planar_hex import make_planar_hex_mesh
 
-from compass.step import Step
 from compass.model import run_model
+from compass.step import Step
 
 
 class InitialState(Step):
     """
     A step for creating a mesh and initial condition for dam break test
     cases
     """
-    def __init__(self, test_case):
+    def __init__(self, test_case, use_lts):
         """
         Create the step
 
@@ -27,8 +24,14 @@ def __init__(self, test_case):
         super().__init__(test_case=test_case, name='initial_state', ntasks=1,
                          min_tasks=1, openmp_threads=1)
 
-        self.add_namelist_file('compass.ocean.tests.dam_break',
-                               'namelist.init', mode='init')
+        self.use_lts = use_lts
+
+        if use_lts:
+            self.add_namelist_file('compass.ocean.tests.dam_break.lts',
+                                   'namelist.init', mode='init')
+        else:
+            self.add_namelist_file('compass.ocean.tests.dam_break',
+                                   'namelist.init', mode='init')
 
         self.add_streams_file('compass.ocean.tests.dam_break',
                               'streams.init', mode='init')
@@ -52,7 +55,7 @@ def run(self):
         dc = section.getfloat('dc')
 
         self.update_namelist_at_runtime(
-           {'config_dam_break_dc': f'{dc}'})
+            {'config_dam_break_dc': f'{dc}'})
 
         logger.info(' * Make planar hex mesh')
         dsMesh = make_planar_hex_mesh(nx=nx, ny=ny, dc=dc, nonperiodic_x=True,