Add Omega support to Cosine Bell tests

Author: xylar

polaris/ocean/tasks/cosine_bell/analysis.py

@@ -79,7 +79,7 @@ def exact_solution(self, refinement_factor, field_name, time, zidx=None):
         ds_mesh = xr.open_dataset(f'mesh_r{refinement_factor:02g}.nc')
         sphere_radius = ds_mesh.sphere_radius
 
-        ds_init = xr.open_dataset(f'init_r{refinement_factor:02g}.nc')
+        ds_init = self.open_model_dataset(f'init_r{refinement_factor:02g}.nc')
         latCell = ds_init.latCell.values
         lonCell = ds_init.lonCell.values
 

polaris/ocean/tasks/cosine_bell/cosine_bell.cfg

@@ -36,7 +36,10 @@ error_type = l2
 # config options for convergence forward steps
 [convergence_forward]
 
-# time integrator: {'split_explicit', 'RK4'}
+# time integrator
+#  either: {'RK4'}
+#  mpas-ocean: {'split_explicit'}
+#  omega: {'Forward-Backward', 'RungeKutta2'}
 time_integrator = RK4
 
 # RK4 time step per resolution (s/km), since dt is proportional to resolution

polaris/ocean/tasks/cosine_bell/forward.py

@@ -46,3 +46,13 @@ def __init__(self, component, name, subdir, mesh, init,
                          graph_target=f'{mesh.path}/graph.info',
                          refinement_factor=refinement_factor,
                          refinement=refinement)
+
+    def setup(self):
+        """
+        TEMP: symlink initial condition to name hard-coded in Omega
+        """
+        super().setup()
+        model = self.config.get('ocean', 'model')
+        # TODO: remove as soon as Omega no longer hard-codes this file
+        if model == 'omega':
+            self.add_input_file(filename='OmegaMesh.nc', target='init.nc')

polaris/ocean/tasks/cosine_bell/forward.yaml

@@ -1,14 +1,16 @@
-mpas-ocean:
-  run_modes:
-    config_ocean_run_mode: forward
+ocean:
   time_management:
     config_stop_time: none
     config_run_duration: {{ run_duration }}
-  decomposition:
-    config_block_decomp_file_prefix: graph.info.part.
   time_integration:
     config_dt: {{ dt }}
     config_time_integrator: {{ time_integrator }}
+
+mpas-ocean:
+  run_modes:
+    config_ocean_run_mode: forward
+  decomposition:
+    config_block_decomp_file_prefix: graph.info.part.
   debug:
     config_disable_thick_all_tend: true
     config_disable_thick_hadv: true
@@ -65,3 +67,32 @@ mpas-ocean:
       - refLayerThickness
       - kineticEnergyCell
       - relativeVorticityCell
+
+Omega:
+  Tendencies:
+    ThicknessFluxTendencyEnable : false
+    PVTendencyEnable : false
+    KETendencyEnable : false
+    SSHTendencyEnable : false
+    VelDiffTendencyEnable: false
+    VelHyperDiffTendencyEnable: false
+    TracerHorzAdvTendencyEnable : true
+    TracerDiffTendencyEnable : false
+    TracerHyperDiffTendencyEnable : false
+  Dimension:
+    NVertLevels: 1
+  IOStreams:
+    InitialState:
+      Filename: init.nc
+      Contents:
+        - State
+        - Tracers
+    History:
+      Filename: output.nc
+      Freq: {{ output_freq }}
+      FreqUnits: Seconds
+      Contents:
+        - Tracers
+        - LayerThickness
+        - NormalVelocity
+    RestartRead: {}

polaris/ocean/tasks/cosine_bell/init.py

@@ -1,14 +1,13 @@
 import numpy as np
 import xarray as xr
-from mpas_tools.io import write_netcdf
 from mpas_tools.transects import lon_lat_to_cartesian
 from mpas_tools.vector import Vector
 
-from polaris import Step
+from polaris.ocean.model import OceanIOStep
 from polaris.ocean.vertical import init_vertical_coord
 
 
-class Init(Step):
+class Init(OceanIOStep):
     """
     A step for an initial condition for for the cosine bell test case
     """
@@ -36,10 +35,6 @@ def __init__(self, component, name, subdir, base_mesh):
             filename='mesh.nc',
             work_dir_target=f'{base_mesh.path}/base_mesh.nc')
 
-        self.add_input_file(
-            filename='graph.info',
-            work_dir_target=f'{base_mesh.path}/graph.info')
-
         self.add_output_file(filename='initial_state.nc')
 
     def run(self):
@@ -107,7 +102,7 @@ def run(self):
         ds['fEdge'] = xr.zeros_like(ds_mesh.xEdge)
         ds['fVertex'] = xr.zeros_like(ds_mesh.xVertex)
 
-        write_netcdf(ds, 'initial_state.nc')
+        self.write_model_dataset(ds, 'initial_state.nc')
 
 
 def cosine_bell(max_value, ri, r):

polaris/ocean/tasks/cosine_bell/viz.py

@@ -1,11 +1,10 @@
 import cmocean  # noqa: F401
-import xarray as xr
 
-from polaris import Step
+from polaris.ocean.model import OceanIOStep
 from polaris.viz import plot_global_mpas_field
 
 
-class Viz(Step):
+class Viz(OceanIOStep):
     """
     A step for plotting fields from the cosine bell output
 
@@ -64,7 +63,7 @@ def run(self):
         mesh_name = self.mesh_name
         run_duration = config.getfloat('convergence_forward', 'run_duration')
 
-        ds_init = xr.open_dataset('initial_state.nc')
+        ds_init = self.open_model_dataset('initial_state.nc')
         da = ds_init['tracer1'].isel(Time=0, nVertLevels=0)
 
         plot_global_mpas_field(
@@ -73,11 +72,11 @@ def run(self):
             title=f'{mesh_name} tracer at init', plot_land=False,
             central_longitude=180.)
 
-        ds_out = xr.open_dataset('output.nc')
+        ds_out = self.open_model_dataset('output.nc')
         da = ds_out['tracer1'].isel(Time=-1, nVertLevels=0)
 
         plot_global_mpas_field(
             mesh_filename='mesh.nc', da=da, out_filename='final.png',
             config=config, colormap_section='cosine_bell_viz',
-            title=f'{mesh_name} tracer after {run_duration/24.:g} days',
+            title=f'{mesh_name} tracer after {run_duration / 24.:g} days',
             plot_land=False, central_longitude=180.)