Add restart test to cosine bell

polaris/ocean/tasks/cosine_bell/__init__.py

@@ -11,6 +11,7 @@
 from polaris.ocean.tasks.cosine_bell.analysis import Analysis
 from polaris.ocean.tasks.cosine_bell.forward import Forward
 from polaris.ocean.tasks.cosine_bell.init import Init
+from polaris.ocean.tasks.cosine_bell.restart import Restart
 from polaris.ocean.tasks.cosine_bell.viz import Viz
 
 
@@ -22,7 +23,8 @@ def add_cosine_bell_tasks(component):
         the ocean component that the tasks will be added to
     """
 
-    for icosahedral, prefix in [(True, 'icos'), (False, 'qu')]:
+    for icosahedral, prefix, restart_refinement in [(True, 'icos', 8.0),
+                                                    (False, 'qu', 2.0)]:
 
         filepath = f'spherical/{prefix}/cosine_bell/cosine_bell.cfg'
         config = PolarisConfigParser(filepath=filepath)
@@ -42,6 +44,12 @@ def add_cosine_bell_tasks(component):
                                               include_viz=include_viz,
                                               refinement=refinement))
 
+        component.add_task(Restart(component=component,
+                                   config=config,
+                                   icosahedral=icosahedral,
+                                   refinement_factor=restart_refinement,
+                                   refinement='both'))
+
 
 class CosineBell(Task):
     """

polaris/ocean/tasks/cosine_bell/restart/__init__.py

@@ -0,0 +1,86 @@
+from polaris import Task
+from polaris.ocean.convergence import get_resolution_for_task
+from polaris.ocean.mesh.spherical import add_spherical_base_mesh_step
+from polaris.ocean.tasks.cosine_bell.init import Init
+from polaris.ocean.tasks.cosine_bell.restart.restart_step import RestartStep
+from polaris.ocean.tasks.cosine_bell.restart.validate import Validate
+
+
+class Restart(Task):
+    """
+    A cosine bell restart test case, which makes sure the model produces
+    identical results with one longer run and two shorter runs with a restart
+    in between.
+    """
+
+    def __init__(self, component, config, icosahedral, refinement_factor,
+                 refinement):
+        """
+        Create the convergence test
+
+        Parameters
+        ----------
+        component : polaris.ocean.Ocean
+            The ocean component that this task belongs to
+
+        config : polaris.config.PolarisConfigParser
+            A shared config parser
+
+        icosahedral : bool
+            Whether to use icosahedral, as opposed to less regular, JIGSAW
+            meshes
+
+        refinement_factor : float
+            The factor by which to scale space, time or both
+
+        refinement : str
+            Refinement type. One of 'space', 'time' or 'both' indicating both
+            space and time
+        """
+
+        if icosahedral:
+            prefix = 'icos'
+        else:
+            prefix = 'qu'
+
+        task_subdir = f'spherical/{prefix}/cosine_bell/restart'
+        name = f'{prefix}_cosine_bell_restart'
+        config_filename = 'cosine_bell.cfg'
+
+        super().__init__(component=component, name=name, subdir=task_subdir)
+
+        self.set_shared_config(config, link=config_filename)
+
+        resolution = get_resolution_for_task(
+            config, refinement_factor, refinement=refinement)
+
+        base_mesh_step, mesh_name = add_spherical_base_mesh_step(
+            component, resolution, icosahedral)
+
+        name = f'{prefix}_init_{mesh_name}'
+        init_subdir = f'spherical/{prefix}/cosine_bell/init/{mesh_name}'
+        if init_subdir in component.steps:
+            init_step = component.steps[init_subdir]
+        else:
+            init_step = Init(component=component, name=name,
+                             subdir=init_subdir, base_mesh=base_mesh_step)
+            init_step.set_shared_config(config, link=config_filename)
+
+        self.add_step(base_mesh_step, symlink=f'base_mesh/{mesh_name}')
+        self.add_step(init_step, symlink=f'init/{mesh_name}')
+
+        step_names = ['full_run', 'restart_run']
+        for name in step_names:
+            subdir = f'{task_subdir}/{name}'
+            step = RestartStep(
+                component=component, name=name, subdir=subdir,
+                mesh=base_mesh_step, init=init_step,
+                refinement_factor=refinement_factor,
+                refinement=refinement)
+            step.set_shared_config(
+                config, link=config_filename)
+            self.add_step(step)
+
+        self.add_step(Validate(component=component,
+                               step_subdirs=step_names,
+                               indir=task_subdir))

polaris/ocean/tasks/cosine_bell/restart/forward.yaml

@@ -0,0 +1,35 @@
+ocean:
+  time_management:
+    config_start_time: {{ start_time }}
+    config_stop_time: none
+    config_run_duration: {{ run_duration }}
+
+mpas-ocean:
+  time_management:
+    config_do_restart: {{ do_restart }}
+  io:
+    config_write_output_on_startup: false
+  streams:
+    restart:
+      filename_template: ../restarts/rst.$Y-$M-$D_$h.$m.$s.nc
+      filename_interval: output_interval
+      output_interval: 0000-00-00_00:00:01
+    output:
+      output_interval: {{ output_interval }}
+
+Omega:
+  IOStreams:
+    InitialState:
+      FreqUnits: {{ init_freq_units }}
+    RestartRead:
+      UsePointerFile: false
+      Filename: ../restarts/rst.$Y-$M-$D_$h.$m.$s
+      StartTime: {{ restart_time }}
+    RestartWrite:
+      UsePointerFile: false
+      Filename: ../restarts/rst.$Y-$M-$D_$h.$m.$s
+      Freq: 1
+      FreqUnits: seconds
+    History:
+      Freq: {{ output_freq }}
+      FreqUnits: seconds

polaris/ocean/tasks/cosine_bell/restart/restart_step.py

@@ -0,0 +1,85 @@
+import os
+import time
+
+import numpy as np
+
+from polaris.ocean.convergence import get_timestep_for_task
+from polaris.ocean.model import get_time_interval_string
+from polaris.ocean.tasks.cosine_bell.forward import Forward
+
+
+class RestartStep(Forward):
+    """
+    A forward model step in the restart test case
+    """
+    def dynamic_model_config(self, at_setup):
+        """
+        Add model config options, namelist, streams and yaml files using config
+        options or template replacements that need to be set both during step
+        setup and at 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)
+
+        dt, _ = get_timestep_for_task(
+            self.config, self.refinement_factor, refinement=self.refinement)
+        dt = np.ceil(dt)
+
+        if self.name == 'full_run':
+            # 2 time steps without a restart
+            do_restart = False
+            start_time = 0.
+            run_duration = 2. * dt
+            output_interval = 2. * dt
+        elif self.name == 'restart_run':
+            # 1 time step from the restart at 1 time step
+            do_restart = True
+            start_time = dt
+            run_duration = dt
+            output_interval = dt
+        else:
+            raise ValueError(f'Unexpected step name {self.name}')
+
+        # to keep the time formatting from getting too complicated, we'll
+        # assume 2 time steps is never more than a day
+        start_time_str = time.strftime('0001-01-01_%H:%M:%S',
+                                       time.gmtime(start_time))
+
+        run_duration_str = get_time_interval_string(seconds=run_duration)
+
+        output_interval_str = get_time_interval_string(seconds=output_interval)
+
+        # For Omega, we want the output interval as a number of seconds
+        output_freq = int(output_interval)
+
+        if do_restart:
+            restart_time_str = start_time_str
+            init_freq_units = 'never'
+        else:
+            # Effectively never
+            restart_time_str = '99999-12-31_00:00:00'
+            init_freq_units = 'OnStartup'
+
+        package = 'polaris.ocean.tasks.cosine_bell.restart'
+        replacements = dict(
+            do_restart=do_restart,
+            start_time=start_time_str,
+            run_duration=run_duration_str,
+            output_interval=output_interval_str,
+            restart_time=restart_time_str,
+            init_freq_units=init_freq_units,
+            output_freq=f'{output_freq}'
+        )
+
+        self.add_yaml_file(package=package,
+                           yaml='forward.yaml',
+                           template_replacements=replacements)
+
+        restart_dir = os.path.abspath(os.path.join(
+            self.work_dir, '..', 'restarts'))
+        os.makedirs(restart_dir, exist_ok=True)

polaris/ocean/tasks/cosine_bell/restart/validate.py

@@ -0,0 +1,68 @@
+import os
+
+from polaris.ocean.model import OceanIOStep
+from polaris.validate import compare_variables
+
+
+class Validate(OceanIOStep):
+    """
+    A step for comparing outputs between steps in a cosine bell run
+
+    Attributes
+    ----------
+    step_subdirs : list of str
+        The number of processors used in each run
+    """
+    def __init__(self, component, step_subdirs, indir):
+        """
+        Create the step
+
+        Parameters
+        ----------
+        component : polaris.Component
+            The component the step belongs to
+
+        step_subdirs : list of str
+            The number of processors used in each run
+
+        indir : str
+            the directory the step is in, to which ``name`` will be appended
+       """
+        super().__init__(component=component, name='validate', indir=indir)
+
+        self.step_subdirs = step_subdirs
+
+        for subdir in step_subdirs:
+            self.add_input_file(filename=f'output_{subdir}.nc',
+                                target=f'../{subdir}/output.nc')
+
+    def run(self):
+        """
+        Compare ``tracer1``, ``layerThickness`` and ``normalVelocity`` in the
+        outputs of two previous steps with each other
+        """
+        super().run()
+        step_subdirs = self.step_subdirs
+        logger = self.logger
+        variables = ['tracer1', 'layerThickness', 'normalVelocity']
+
+        filename1 = self.inputs[0]
+        filename2 = self.inputs[1]
+
+        all_pass = True
+        for filename in [filename1, filename2]:
+            if not os.path.exists(filename):
+                logger.error(f'File {filename} does not exist.')
+                all_pass = False
+
+        if all_pass:
+            ds1 = self.open_model_dataset(filename1)
+            ds2 = self.open_model_dataset(filename2)
+
+            all_pass = compare_variables(variables=variables,
+                                         filename1=filename1,
+                                         filename2=filename2, logger=logger,
+                                         ds1=ds1, ds2=ds2)
+        if not all_pass:
+            raise ValueError(f'Validation failed comparing outputs between '
+                             f'{step_subdirs[0]} and {step_subdirs[1]}.')