Add barotropic jet test group

polaris/ocean/__init__.py

@@ -1,5 +1,6 @@
 from polaris import Component
 from polaris.ocean.tests.baroclinic_channel import BaroclinicChannel
+from polaris.ocean.tests.galewsky_jet import GalewskyJet
 from polaris.ocean.tests.global_convergence import GlobalConvergence
 from polaris.ocean.tests.single_column import SingleColumn
 
@@ -17,6 +18,7 @@ def __init__(self):
 
         # please keep these in alphabetical order
         self.add_test_group(BaroclinicChannel(component=self))
+        self.add_test_group(GalewskyJet(component=self))
         self.add_test_group(GlobalConvergence(component=self))
         self.add_test_group(SingleColumn(component=self))
 

polaris/ocean/tests/galewsky_jet/__init__.py

@@ -0,0 +1,24 @@
+from polaris import TestGroup
+from polaris.ocean.tests.galewsky_jet.test_balance import TestBalance
+from polaris.ocean.tests.galewsky_jet.short_adjustment import ShortAdjustment
+from polaris.ocean.tests.galewsky_jet.barotropic_instability import BarotropicInstability
+
+
+class GalewskyJet(TestGroup):
+    """
+    A test group for "galewsky jet" test cases
+    """
+    def __init__(self, component):
+        """
+        component : polaris.ocean.Ocean
+            the ocean component that this test group belongs to
+        """
+        super().__init__(component=component, name='galewsky_jet')
+
+        for resolution in [120.]:
+            self.add_test_case(
+                TestBalance(test_group=self, resolution=resolution))
+            self.add_test_case(
+                ShortAdjustment(test_group=self, resolution=resolution))
+            self.add_test_case(
+                BarotropicInstability(test_group=self, resolution=resolution))

polaris/ocean/tests/galewsky_jet/barotropic_instability/__init__.py

@@ -0,0 +1,65 @@
+from polaris.mesh.spherical import IcosahedralMeshStep
+from polaris.ocean.tests.galewsky_jet.barotropic_instability.forward import (
+    Forward,
+)
+from polaris.ocean.tests.galewsky_jet.barotropic_instability.initial_state import (
+    InitialState,
+)
+from polaris.ocean.tests.galewsky_jet.barotropic_instability.viz import Viz
+from polaris.testcase import TestCase
+
+
+class BarotropicInstability(TestCase):
+    """
+    A class to define the Galewsky jet barotropic instability
+    test case (6 days for instbaility development)
+
+    Attributes
+    ----------
+    resolution : str
+        The resolution of the test case
+    """
+
+    def __init__(self, test_group, resolution):
+        """
+        Create the test case
+
+        Parameters
+        ----------
+        test_group :
+        polaris.ocean.tests.galewsky_jet.GalewskyJet
+            The test group that this test case belongs to
+
+        resolution : float
+            The resolution of the test case (in km)
+        """
+        name = 'barotropic_instability'
+        self.resolution = resolution
+        res = int(resolution)
+        subdir = f'{res}km/{name}'
+        super().__init__(test_group=test_group, name=name,
+                         subdir=subdir)
+
+        self.add_step(IcosahedralMeshStep(
+            test_case=self, cell_width=resolution))
+        self.add_step(
+            InitialState(test_case=self, resolution=resolution))
+        self.add_step(
+            Forward(test_case=self, resolution=resolution))
+
+        mesh_name = f'Icos{res}'
+        name = f'{mesh_name}_viz'
+        subdir = 'viz'
+        self.add_step(Viz(test_case=self, name=name, subdir=subdir,
+                          mesh_name=mesh_name))
+
+    def configure(self):
+        """
+        Set config options for the test case
+        """
+        super().configure()
+        config = self.config
+        config.add_from_package('polaris.mesh', 'mesh.cfg')
+
+        config.set('spherical_mesh', 'mpas_mesh_filename',
+                   'mesh.nc')

polaris/ocean/tests/galewsky_jet/barotropic_instability/forward.py

@@ -0,0 +1,99 @@
+import time
+
+import xarray as xr
+
+from polaris.ocean.model import OceanModelStep
+
+
+class Forward(OceanModelStep):
+    """
+    A step for performing forward ocean component runs as part of
+    the galewsky jet - barotropic instability test case
+
+    Attributes
+    ----------
+    resolution : int
+        The resolution of the (uniform) mesh in km
+    """
+
+    def __init__(self, test_case, resolution):
+        """
+        Create a new step
+
+        Parameters
+        ----------
+        test_case : polaris.ocean.tests.galewsky_jet.barotropic_instability.BarotropicInstability  # noqa: E501
+            The test case this step belongs to
+
+        resolution : int
+            The resolution of the (uniform) mesh in km
+        """
+        super().__init__(test_case=test_case,
+                         name='forward',
+                         subdir='forward',
+                         openmp_threads=1)
+
+        self.resolution = resolution
+
+        # make sure output is double precision
+        self.add_yaml_file('polaris.ocean.config', 'output.yaml')
+
+        self.add_yaml_file(
+            'polaris.ocean.tests.galewsky_jet.barotropic_instability',
+            'forward.yaml')
+
+        self.add_input_file(filename='initial_state.nc',
+                            target='../initial_state/initial_state.nc')
+        self.add_input_file(filename='graph.info',
+                            target='../base_mesh/graph.info')
+
+        self.add_output_file(filename='output.nc')
+
+    def compute_cell_count(self, at_setup):
+        """
+        Compute the approximate number of cells in the mesh, used to constrain
+        resources
+
+        Parameters
+        ----------
+        at_setup : bool
+            Whether this method is being run during setup of the step, as
+            opposed to at runtime
+
+        Returns
+        -------
+        cell_count : int or None
+            The approximate number of cells in the mesh
+        """
+        if at_setup:
+            # use a heuristic based on QU30 (65275 cells) and QU240 (10383
+            # cells)
+            cell_count = 6e8 / self.resolution**2
+        else:
+            # get nCells from the input file
+            with xr.open_dataset('initial_state.nc') as ds:
+                cell_count = ds.sizes['nCells']
+        return cell_count
+
+    def dynamic_model_config(self, at_setup):
+        """
+        Set the model time step from config options at setup and runtime
+
+        Parameters
+        ----------
+        at_setup : bool
+            Whether this method is being run during setup of the step, as
+            opposed to at runtime
+        """
+        super().dynamic_model_config(at_setup=at_setup)
+
+        config = self.config
+        # dt is proportional to resolution: default 30 seconds per km
+        dt_per_km = config.getfloat('galewsky_jet', 'dt_per_km')
+
+        dt = dt_per_km * self.resolution
+        # https://stackoverflow.com/a/1384565/7728169
+        dt_str = time.strftime('%H:%M:%S', time.gmtime(dt))
+
+        options = dict(config_dt=dt_str)
+        self.add_model_config_options(options)

polaris/ocean/tests/galewsky_jet/barotropic_instability/forward.yaml

@@ -0,0 +1,75 @@
+omega:
+  run_modes:
+    config_ocean_run_mode: forward
+  time_management:
+    config_run_duration: 0006_00:00:00
+  ALE_vertical_grid:
+    config_vert_coord_movement: impermeable_interfaces
+    config_ALE_thickness_proportionality: weights_only
+  decomposition:
+    config_block_decomp_file_prefix: graph.info.part.
+  time_integration:
+    config_time_integrator: RK4
+  debug:
+    config_disable_thick_all_tend: false
+    config_disable_thick_hadv: false
+    config_disable_thick_vadv: true
+    config_disable_thick_sflux: true
+    config_disable_vel_all_tend: false
+    config_disable_vel_coriolis: false
+    config_disable_vel_pgrad: false
+    config_disable_vel_hmix: false
+    config_disable_vel_surface_stress: true
+    config_disable_vel_explicit_bottom_drag: true
+    config_disable_vel_vmix: true
+    config_disable_vel_vadv: true
+    config_disable_tr_all_tend: true
+    config_disable_tr_hmix: true
+    config_disable_tr_vmix: true
+    config_disable_tr_sflux: true
+    config_disable_tr_nonlocalflux: true
+    config_check_ssh_consistency: false
+  cvmix:
+    config_use_cvmix: false
+  eos:
+    config_eos_type: linear
+  forcing:
+    config_use_bulk_wind_stress: false
+    config_use_bulk_thickness_flux: false
+  pressure_gradient:
+    config_pressure_gradient_type: ssh_gradient
+  tracer_forcing_activeTracers:
+    config_use_activeTracers_surface_restoring: false
+    config_use_activeTracers_interior_restoring: false
+    config_use_activeTracers: true
+    config_use_activeTracers_surface_bulk_forcing: false
+  tracer_forcing_debugTracers:
+    config_use_debugTracers: false
+  AM_mixedLayerDepths:
+    config_AM_mixedLayerDepths_enable: false
+  streams:
+    mesh:
+      filename_template: initial_state.nc
+    input:
+      filename_template: initial_state.nc
+    restart:
+      output_interval: 0000_01:00:00
+    output:
+      type: output
+      filename_template: output.nc
+      output_interval: 0001_00:00:00
+      clobber_mode: truncate
+      reference_time: 0001-01-01_00:00:00
+      contents:
+      - mesh
+      - xtime
+      - normalVelocity
+      - layerThickness
+      - refZMid
+      - refLayerThickness
+      - ssh
+      - kineticEnergyCell
+      - relativeVorticityCell
+      - relativeVorticity
+      - normalizedPlanetaryVorticityEdge
+      - normalizedRelativeVorticityEdge

polaris/ocean/tests/galewsky_jet/barotropic_instability/initial_state.py

@@ -0,0 +1,99 @@
+import cmocean  # noqa: F401
+import xarray as xr
+from mpas_tools.io import write_netcdf
+from mpas_tools.mesh.jet import init as jet_init
+
+from polaris import Step
+from polaris.ocean.vertical import init_vertical_coord
+
+
+class InitialState(Step):
+    """
+    A step for creating a mesh and initial condition
+    for galewsky jet barotropic instability
+    test case
+
+    Attributes
+    ----------
+    resolution : float
+        The resolution of the test case in km
+    """
+    def __init__(self, test_case, resolution):
+        """
+        Create the step
+
+        Parameters
+        ----------
+        test_case : polaris.TestCase
+            The test case this step belongs to
+
+        resolution : float
+            The resolution of the test case in km
+        """
+        super().__init__(test_case=test_case, name='initial_state')
+        self.resolution = resolution
+
+        self.add_input_file(
+            filename='mesh.nc',
+            target='../base_mesh/mesh.nc')
+        self.add_input_file(
+            filename='graph.info',
+            target='../base_mesh/graph.info')
+        # self.add_input_file(
+        #     filename='init.nc')
+        self.add_output_file('initial_state.nc')
+
+    def run(self):
+        """
+        Run this step of the test case
+        """
+        config = self.config
+# update rsph from cime.constants
+        # jet_init(name='mesh.nc', save='velocity_ic.nc',
+        #         rsph=6371220.0, pert=False)
+        # ds2 = xr.open_dataset('velocity_ic.nc')
+
+        dsMesh = xr.open_dataset('mesh.nc')
+
+        ds = dsMesh.copy()
+        x_cell = ds.xCell
+
+        bottom_depth = config.getfloat('vertical_grid', 'bottom_depth')
+
+        ds['bottomDepth'] = bottom_depth * xr.ones_like(x_cell)
+        ds['ssh'] = xr.zeros_like(x_cell)
+        # ds['ssh'] = ds2.h - ds['bottomDepth']
+
+        init_vertical_coord(config, ds)
+
+        # resolution = self.resolution
+        section = config['galewsky_jet']
+        temperature = section.getfloat('temperature')
+        salinity = section.getfloat('salinity')
+        # coriolis_parameter = section.getfloat('coriolis_parameter')
+
+        temperature_array = temperature * xr.ones_like(x_cell)
+        temperature_array, _ = xr.broadcast(temperature_array, ds.refZMid)
+        ds['temperature'] = temperature_array.expand_dims(dim='Time', axis=0)
+        ds['salinity'] = salinity * xr.ones_like(ds.temperature)
+
+        # ds2 = xr.open_dataset('init.nc')
+        jet_init(name='mesh.nc', save='velocity_ic.nc',
+                 rsph=6371220.0, pert=True)
+        ds2 = xr.open_dataset('velocity_ic.nc')
+
+        unrm_array, _ = xr.broadcast(ds2.u, ds.refZMid)
+        ds['normalVelocity'] = unrm_array
+        # ds['normalVelocity'] = ds2.u
+        h_array, _ = xr.broadcast(ds2.h, ds.refZMid)
+        ds['layerThickness'] = h_array
+        ds['fCell'] = ds2.fCell
+        ds['fEdge'] = ds2.fEdge
+        ds['fVertex'] = ds2.fVertex
+        ds['ssh'] = ds2.h - ds['bottomDepth']
+
+        # if (config.getfloat('vertical_grid', 'grid_type') == 'uniform'):
+        # nlev = config.getfloat('vertical_grid', 'vert_levels')
+        # ds['layerThickness'], _ = xr.broadcast(ds2.h / nlev, ds.refZMid)
+
+        write_netcdf(ds, 'initial_state.nc')