From b6dcacb13539142b6bbc86e4d339c3b6c199cc10 Mon Sep 17 00:00:00 2001
From: Xylar Asay-Davis <xylarstorm@gmail.com>
Date: Mon, 23 Dec 2024 07:48:58 -0600
Subject: [PATCH] Address reviewer comments

Add a `do_restart` attribute to keep track of which runs are
restart runs, rather than using the step name.
---
 polaris/ocean/tasks/cosine_bell/forward.py            | 11 ++++++++++-
 polaris/ocean/tasks/cosine_bell/restart/__init__.py   |  3 ++-
 .../ocean/tasks/cosine_bell/restart/restart_step.py   | 10 ++++------
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/polaris/ocean/tasks/cosine_bell/forward.py b/polaris/ocean/tasks/cosine_bell/forward.py
index 4bdf26479..5e4da4b84 100644
--- a/polaris/ocean/tasks/cosine_bell/forward.py
+++ b/polaris/ocean/tasks/cosine_bell/forward.py
@@ -5,10 +5,15 @@ class Forward(SphericalConvergenceForward):
     """
     A step for performing forward ocean component runs as part of the cosine
     bell test case
+
+    Attributes
+    ----------
+    do_restart : bool
+        Whether this is a restart run
     """
 
     def __init__(self, component, name, subdir, mesh, init,
-                 refinement_factor, refinement):
+                 refinement_factor, refinement, do_restart=False):
         """
         Create a new step
 
@@ -35,6 +40,9 @@ def __init__(self, component, name, subdir, mesh, init,
         refinement : str
             Refinement type. One of 'space', 'time' or 'both' indicating both
             space and time
+
+        do_restart : bool, optional
+            Whether this is a restart run
         """
         package = 'polaris.ocean.tasks.cosine_bell'
         validate_vars = ['normalVelocity', 'tracer1']
@@ -46,6 +54,7 @@ def __init__(self, component, name, subdir, mesh, init,
                          graph_target=f'{mesh.path}/graph.info',
                          refinement_factor=refinement_factor,
                          refinement=refinement)
+        self.do_restart = do_restart
 
     def setup(self):
         """
diff --git a/polaris/ocean/tasks/cosine_bell/restart/__init__.py b/polaris/ocean/tasks/cosine_bell/restart/__init__.py
index e33f1f8ca..c35a1252e 100644
--- a/polaris/ocean/tasks/cosine_bell/restart/__init__.py
+++ b/polaris/ocean/tasks/cosine_bell/restart/__init__.py
@@ -72,11 +72,12 @@ def __init__(self, component, config, icosahedral, refinement_factor,
         step_names = ['full_run', 'restart_run']
         for name in step_names:
             subdir = f'{task_subdir}/{name}'
+            do_restart = (name == 'restart_run')
             step = RestartStep(
                 component=component, name=name, subdir=subdir,
                 mesh=base_mesh_step, init=init_step,
                 refinement_factor=refinement_factor,
-                refinement=refinement)
+                refinement=refinement, do_restart=do_restart)
             step.set_shared_config(
                 config, link=config_filename)
             self.add_step(step)
diff --git a/polaris/ocean/tasks/cosine_bell/restart/restart_step.py b/polaris/ocean/tasks/cosine_bell/restart/restart_step.py
index 28934c0ac..7f8771e35 100644
--- a/polaris/ocean/tasks/cosine_bell/restart/restart_step.py
+++ b/polaris/ocean/tasks/cosine_bell/restart/restart_step.py
@@ -26,24 +26,22 @@ def dynamic_model_config(self, at_setup):
         """
         super().dynamic_model_config(at_setup)
 
+        do_restart = self.do_restart
+
         dt, _ = get_timestep_for_task(
             self.config, self.refinement_factor, refinement=self.refinement)
         dt = np.ceil(dt)
 
-        if self.name == 'full_run':
+        if not do_restart:
             # 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':
+        else:
             # 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