Add omega support to barotropic gyre

Author: cbegeman

polaris/tasks/ocean/barotropic_gyre/analysis.py

@@ -4,7 +4,6 @@
 import matplotlib.pyplot as plt
 import mosaic
 import numpy as np
-import xarray as xr
 from matplotlib import colors as mcolors
 from mpas_tools.ocean import compute_barotropic_streamfunction
 
@@ -50,9 +49,9 @@ def __init__(
 
     def run(self):
         logger = self.logger
-        ds_mesh = xr.open_dataset('mesh.nc')
-        ds_init = xr.open_dataset('init.nc')
-        ds = xr.open_dataset('output.nc')
+        ds_mesh = self.open_model_dataset('mesh.nc')
+        ds_init = self.open_model_dataset('init.nc')
+        ds = self.open_model_dataset('output.nc')
 
         field_mpas = compute_barotropic_streamfunction(
             ds_init.isel(Time=0), ds, prefix='', time_index=-1
@@ -84,7 +83,6 @@ def run(self):
 
         use_mplstyle()
         pad = 20
-        fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(12, 2))
         x0 = ds_mesh.xEdge.min().values
         y0 = ds_mesh.yEdge.min().values
 
@@ -94,6 +92,8 @@ def run(self):
         # convert to km
         descriptor.vertex_patches *= 1.0e-3
 
+        fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(12, 2))
+
         eta0 = max(
             np.max(np.abs(field_exact.values)),
             np.max(np.abs(field_mpas.values)),

polaris/tasks/ocean/barotropic_gyre/barotropic_gyre.cfg

@@ -1,5 +1,8 @@
 [barotropic_gyre]
 
+# time integrator
+time_integrator = RK4
+
 # distance in kilometers between cell centers
 resolution = 10
 

polaris/tasks/ocean/barotropic_gyre/forward.py

@@ -93,7 +93,6 @@ def __init__(
         self.add_yaml_file('polaris.ocean.config', 'output.yaml')
 
         self.add_input_file(filename='init.nc', target='../init/init.nc')
-        self.add_input_file(filename='forcing.nc', target='../init/forcing.nc')
 
         self.add_output_file(
             filename='output.nc',
@@ -195,32 +194,58 @@ def dynamic_model_config(self, at_setup):
                 f'maximum value is {nu_max} or decrease the time step'
             )
 
+        model = config.get('ocean', 'model')
         options = {'config_dt': dt_str, 'config_density0': rho_0}
         self.add_model_config_options(
             options=options, config_model='mpas-ocean'
         )
 
         if self.run_time_steps is not None:
+            output_interval_units = 'Seconds'
             run_duration = ceil(self.run_time_steps * dt)
             stop_time_str = time.strftime(
                 '0001-01-01_%H:%M:%S', time.gmtime(run_duration)
             )
-            output_interval_str = time.strftime(
-                '0000_%H:%M:%S', time.gmtime(run_duration)
-            )
+            if model == 'omega':
+                output_interval_str = str(run_duration)
+            else:
+                output_interval_str = get_time_interval_string(
+                    seconds=run_duration
+                )
         else:
+            output_interval = 1
+            output_interval_units = 'Months'
             run_duration = config.getfloat('barotropic_gyre', 'run_duration')
             stop_time_str = time.strftime(
                 f'{run_duration + 1.0:04g}-01-01_00:00:00'
             )
-            output_interval_str = time.strftime('0000-01-00_00:00:00')
+            if model == 'omega':
+                output_interval_str = str(output_interval)
+            else:
+                output_interval_str = get_time_interval_string(
+                    days=output_interval * 30.0
+                )
 
         slip_factor_dict = {'no-slip': 0.0, 'free-slip': 1.0}
+        time_integrator = config.get('barotropic_gyre', 'time_integrator')
+        time_integrator_map = dict([('RK4', 'RungeKutta4')])
+        if model == 'omega':
+            if time_integrator in time_integrator_map.keys():
+                time_integrator = time_integrator_map[time_integrator]
+            else:
+                print(
+                    'Warning: mapping from time integrator '
+                    f'{time_integrator} to omega not found, '
+                    'retaining name given in config'
+                )
+
         replacements = dict(
             dt=dt_str,
             dt_btr=dt_btr_str,
             stop_time=stop_time_str,
             output_interval=output_interval_str,
+            output_interval_units=output_interval_units,
+            time_integrator=time_integrator,
             nu=f'{nu:02g}',
             slip_factor=f'{slip_factor_dict[self.boundary_condition]:02g}',
         )
@@ -232,6 +257,16 @@ def dynamic_model_config(self, at_setup):
             template_replacements=replacements,
         )
 
+    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')
+
     def compute_max_time_step(self, config):
         """
         Compute the approximate maximum time step for stability

polaris/tasks/ocean/barotropic_gyre/forward.yaml

@@ -4,17 +4,39 @@ ocean:
     config_run_duration: none
   time_integration:
     config_dt: {{ dt }}
-    config_time_integrator: split_explicit_ab2
+    config_time_integrator: {{ time_integrator }}
   hmix_del2:
     config_use_mom_del2: true
     config_mom_del2: {{ nu }}
+  hmix_del4:
+    config_use_mom_del4: false
+  forcing:
+    config_use_bulk_wind_stress: true
+  debug:
+    config_disable_vel_explicit_bottom_drag: true
+
+mpas-ocean:
+  run_modes:
+    config_ocean_run_mode: forward
+  decomposition:
+    config_block_decomp_file_prefix: graph.info.part.
+  advection:
+    config_thickness_flux_type: constant
+  lateral_walls:
+    config_wall_slip_factor: {{ slip_factor }}
+  debug:
+    config_disable_thick_vadv: true
+    config_disable_vel_hadv: true
+    config_disable_vel_hmix: false
+    config_disable_tr_all_tend: true
+
   streams:
     mesh:
       filename_template: init.nc
     input:
       filename_template: init.nc
     forcing:
-      filename_template: forcing.nc
+      filename_template: init.nc
       input_interval: initial_only
       type: input
       contents:
@@ -33,24 +55,39 @@ ocean:
       - ssh
       - relativeVorticity
 
-mpas-ocean:
-  advection:
-    config_thickness_flux_type: constant
-  bottom_drag:
-    config_implicit_constant_bottom_drag_coeff: 0.
-  forcing:
-    config_use_bulk_wind_stress: true
-  lateral_walls:
-    config_wall_slip_factor: {{ slip_factor }}
-  debug:
-    config_disable_thick_vadv: true
-    config_disable_vel_hadv: true
-    config_disable_vel_hmix: false
-    config_disable_tr_all_tend: true
-
 Omega:
   Tendencies:
-    VelDiffTendencyEnable: true
-    VelHyperDiffTendencyEnable: false
+    TracerHorzAdvTendencyEnable : false
+    TracerDiffTendencyEnable : false
+    TracerHyperDiffTendencyEnable : false
+  Advection:
+    FluxThicknessType: Center
   Dimension:
     NVertLevels: 1
+  IOStreams:
+    InitialState:
+      UsePointerFile: false
+      Filename: init.nc
+      Mode: read
+      Precision: double
+      Freq: 1
+      FreqUnits: OnStartup
+      UseStartEnd: false
+      Contents:
+      - State
+      - WindStressZonal
+      - WindStressMeridional
+    History:
+      UsePointerFile: false
+      Filename: output.nc
+      Mode: write
+      IfExists: append
+      Precision: double
+      Freq: {{ output_interval }}
+      FreqUnits: {{ output_interval_units }}
+      UseStartEnd: false
+      Contents:
+      - State
+      - SshCell
+      FileFreq: 9999
+      FileFreqUnits: years

polaris/tasks/ocean/barotropic_gyre/init.py

@@ -1,16 +1,15 @@
 import numpy as np
 import xarray as xr
-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 polaris import Step
 from polaris.mesh.planar import compute_planar_hex_nx_ny
+from polaris.ocean.model import OceanIOStep
 from polaris.ocean.vertical import init_vertical_coord
 from polaris.viz import plot_horiz_field
 
 
-class Init(Step):
+class Init(OceanIOStep):
     """
     A step for creating a mesh and initial condition for baroclinic channel
     tasks
@@ -43,7 +42,6 @@ def __init__(
             'base_mesh.nc',
             'culled_mesh.nc',
             'culled_graph.info',
-            'forcing.nc',
         ]:
             self.add_output_file(file)
         self.name = name
@@ -72,13 +70,13 @@ def run(self):
         ds_mesh = make_planar_hex_mesh(
             nx=nx, ny=ny, dc=dc, nonperiodic_x=True, nonperiodic_y=True
         )
-        write_netcdf(ds_mesh, 'base_mesh.nc')
+        self.write_model_dataset(ds_mesh, 'base_mesh.nc')
 
         ds_mesh = cull(ds_mesh, logger=logger)
         ds_mesh = convert(
             ds_mesh, graphInfoFileName='culled_graph.info', logger=logger
         )
-        write_netcdf(ds_mesh, 'culled_mesh.nc')
+        self.write_model_dataset(ds_mesh, 'culled_mesh.nc')
 
         # vertical coordinate parameters
         bottom_depth = config.getfloat('vertical_grid', 'bottom_depth')
@@ -116,30 +114,28 @@ def run(self):
         normal_velocity = normal_velocity.expand_dims(dim='Time', axis=0)
         ds['normalVelocity'] = normal_velocity
 
-        # write the initial condition file
-        write_netcdf(ds, 'init.nc')
-
         # set the wind stress forcing
-        ds_forcing = xr.Dataset()
         # Convert from km to m
         ly = ly * 1e3
         wind_stress_zonal = -tau_0 * np.cos(
             np.pi * (ds.yCell - ds.yCell.min()) / ly
         )
         wind_stress_meridional = xr.zeros_like(ds.xCell)
-        ds_forcing['windStressZonal'] = wind_stress_zonal.expand_dims(
+        ds['windStressZonal'] = wind_stress_zonal.expand_dims(
             dim='Time', axis=0
         )
-        ds_forcing['windStressMeridional'] = (
-            wind_stress_meridional.expand_dims(dim='Time', axis=0)
+        ds['windStressMeridional'] = wind_stress_meridional.expand_dims(
+            dim='Time', axis=0
         )
-        write_netcdf(ds_forcing, 'forcing.nc')
+
+        # write the initial condition file
+        self.write_model_dataset(ds, 'init.nc')
 
         cell_mask = ds.maxLevelCell >= 1
 
         plot_horiz_field(
             ds_mesh,
-            ds_forcing['windStressZonal'],
+            ds['windStressZonal'],
             'forcing_wind_stress_zonal.png',
             cmap='cmo.balance',
             show_patch_edges=True,