Check CICE4 restart file dates (#539)

blimlim · web-flow · commit 0c2a837bb9fd · 2024-12-13T15:52:56.000+11:00
* Add CICE4 restart date checks in access.py driver
diff --git a/payu/models/access.py b/payu/models/access.py
@@ -218,6 +218,13 @@ def setup(self):
                 f90nml.write(cpl_nml, nml_work_path + '~')
                 shutil.move(nml_work_path + '~', nml_work_path)
 
+            if model.model_type == 'cice':
+                if model.prior_restart_path and not self.expt.repeat_run:
+                    # Set up and check the cice restart files.
+                    model.overwrite_restart_ptr(run_start_date,
+                                                previous_runtime,
+                                                start_date_fpath)
+
         # Now change the oasis runtime. This needs to be done after the others.
         for model in self.expt.models:
             if model.model_type == 'oasis':
diff --git a/payu/models/cice.py b/payu/models/cice.py
@@ -17,6 +17,7 @@
 import sys
 import shutil
 import datetime
+import struct
 import re
 import tarfile
 
@@ -173,27 +174,7 @@ def setup(self):
         setup_nml = self.ice_in['setup_nml']
 
         if self.prior_restart_path:
-            # Generate ice.restart_file
-            # TODO: better check of restart filename
-            iced_restart_file = None
-            iced_restart_files = [f for f in self.get_prior_restart_files()
-                                  if f.startswith('iced.')]
-
-            if len(iced_restart_files) > 0:
-                iced_restart_file = sorted(iced_restart_files)[-1]
-
-            if iced_restart_file is None:
-                raise FileNotFoundError(
-                    f'No iced restart file found in {self.prior_restart_path}')
-
-            res_ptr_path = os.path.join(self.work_init_path,
-                                        'ice.restart_file')
-            if os.path.islink(res_ptr_path):
-                # If we've linked in a previous pointer it should be deleted
-                os.remove(res_ptr_path)
-            with open(res_ptr_path, 'w') as res_ptr:
-                res_dir = self.get_ptr_restart_dir()
-                print(os.path.join(res_dir, iced_restart_file), file=res_ptr)
+            self._make_restart_ptr()
 
             # Update input namelist
             setup_nml['runtype'] = 'continue'
@@ -394,3 +375,130 @@ def link_restart(self, fpath):
             )
 
         make_symlink(input_path, input_work_path)
+
+    def _make_restart_ptr(self):
+        """
+        CICE4 restart pointers are created in the access driver, where
+        the correct run start dates are available.
+        """
+        pass
+
+    def overwrite_restart_ptr(self,
+                              run_start_date,
+                              previous_runtime,
+                              calendar_file):
+        """
+        Generate restart pointer file 'ice.restart_file' pointing to
+        'iced.YYYYMMDD' with the correct start date.
+        Additionally check that the `iced.YYYYMNDD` restart file's header
+        has the correct previous runtime.
+        Typically called from the access driver, which provides the
+        the correct date and runtime.
+
+        Parameters
+        ----------
+        run_start_date: datetime.date
+            Start date of the new simulation
+        previous_runtime:  int
+            Seconds between experiment initialisation date and start date
+        calendar_file:  str
+            Calendar restart file used to calculate timing information
+        """
+        # Expected iced restart file name
+        iced_restart_file = self.find_matching_iced(self.prior_restart_path,
+                                                    run_start_date)
+
+        res_ptr_path = os.path.join(self.work_init_path,
+                                    'ice.restart_file')
+        if os.path.islink(res_ptr_path):
+            # If we've linked in a previous pointer it should be deleted
+            os.remove(res_ptr_path)
+
+        iced_path = os.path.join(self.prior_restart_path,
+                                 iced_restart_file)
+
+        # Check binary restart has correct time
+        self._cice4_check_date_consistency(iced_path,
+                                           previous_runtime,
+                                           calendar_file)
+
+        with open(res_ptr_path, 'w') as res_ptr:
+            res_dir = self.get_ptr_restart_dir()
+            res_ptr.write(os.path.join(res_dir, iced_restart_file))
+
+    def _cice4_check_date_consistency(self,
+                                      iced_path,
+                                      previous_runtime,
+                                      calendar_file):
+        """
+        Check that the previous runtime in iced restart file header
+        matches the runtime calculated from the calendar restart file.
+
+        Parameters
+        ----------
+        iced_path: str or Path
+            Path to iced restart file
+        previous_runtime:  int
+            Seconds between experiment initialisation date and start date
+        calendar_file:  str or Path
+            Calendar restart file used to calculate timing information
+        """
+        _, _, cice_iced_runtime, _ = read_binary_iced_header(iced_path)
+        if previous_runtime != cice_iced_runtime:
+            msg = (f"Previous runtime from calendar file "
+                   f"{calendar_file}: {previous_runtime} "
+                   "does not match previous runtime in restart"
+                   f"file {iced_path}: {cice_iced_runtime}.")
+            raise RuntimeError(msg)
+
+    def find_matching_iced(self, dir_path, date):
+        """
+        Check a directory for an iced.YYYYMMDD restart file matching a
+        specified date.
+        Raises an error if the expected file is not found.
+
+        Parameters
+        ----------
+        dir_path: str or Path
+            Path to directory containing iced restart files
+        date: datetime.date
+            Date for matching iced file names
+
+        Returns
+        -------
+        iced_file_name: str
+            Name of iced restart file found in dir_path matching
+            the specified date
+        """
+        # Expected iced restart file name
+        date_int = cal.date_to_int(date)
+        iced_restart_file = f"iced.{date_int:08d}"
+
+        dir_files = [f for f in os.listdir(dir_path)
+                     if os.path.isfile(os.path.join(dir_path, f))]
+
+        if iced_restart_file not in dir_files:
+            msg = (f"CICE restart file not found in {dir_path}. Expected "
+                   f"{iced_restart_file} to exist. Is 'dumpfreq' set "
+                   f"in {self.ice_nml_fname} consistently with the run-length?"
+                   )
+            raise FileNotFoundError(msg)
+
+        return iced_restart_file
+
+
+CICE4_RESTART_HEADER_SIZE = 24
+CICE4_RESTART_HEADER_FORMAT = '>iidd'
+
+
+def read_binary_iced_header(iced_path):
+    """
+    Read header information from a CICE4 binary restart file.
+    """
+    with open(iced_path, 'rb') as iced_file:
+        header = iced_file.read(CICE4_RESTART_HEADER_SIZE)
+        bint, istep0, time, time_forc = struct.unpack(
+                                            CICE4_RESTART_HEADER_FORMAT,
+                                            header)
+
+    return (bint, istep0, time, time_forc)
diff --git a/payu/models/cice5.py b/payu/models/cice5.py
@@ -81,3 +81,28 @@ def _calc_runtime(self):
         the timing information in the cice_in.nml namelist.
         """
         pass
+
+    def _make_restart_ptr(self):
+        """
+        Generate restart pointer which points to the latest iced.YYYYMMDD
+        restart file.
+        """
+        iced_restart_file = None
+        iced_restart_files = [f for f in self.get_prior_restart_files()
+                              if f.startswith('iced.')]
+
+        if len(iced_restart_files) > 0:
+            iced_restart_file = sorted(iced_restart_files)[-1]
+
+        if iced_restart_file is None:
+            raise FileNotFoundError(
+                f'No iced restart file found in {self.prior_restart_path}')
+
+        res_ptr_path = os.path.join(self.work_init_path,
+                                    'ice.restart_file')
+        if os.path.islink(res_ptr_path):
+            # If we've linked in a previous pointer it should be deleted
+            os.remove(res_ptr_path)
+        with open(res_ptr_path, 'w') as res_ptr:
+            res_dir = self.get_ptr_restart_dir()
+            res_ptr.write(os.path.join(res_dir, iced_restart_file))
diff --git a/test/models/test_access.py b/test/models/test_access.py
@@ -4,10 +4,12 @@
 
 import pytest
 import cftime
+import f90nml
+from unittest.mock import patch
 
 import payu
 
-from test.common import cd
+from test.common import cd, expt_workdir
 from test.common import tmpdir, ctrldir, labdir, workdir, archive_dir
 from test.common import config as config_orig
 from test.common import write_config
@@ -16,7 +18,6 @@
 from test.common import make_expt_archive_dir, remove_expt_archive_dirs
 from test.common import config_path
 from payu.calendar import GREGORIAN, NOLEAP
-import f90nml
 
 
 verbose = True
@@ -44,7 +45,6 @@ def setup_module(module):
         tmpdir.mkdir()
         labdir.mkdir()
         ctrldir.mkdir()
-        workdir.mkdir()
         archive_dir.mkdir()
         make_all_files()
     except Exception as e:
@@ -65,6 +65,23 @@ def teardown_module(module):
         print(e)
 
 
+@pytest.fixture(autouse=True)
+def empty_workdir():
+    """
+    Model setup tests require a clean work directory and symlink from
+    the control directory.
+    """
+    expt_workdir.mkdir(parents=True)
+    # Symlink must exist for setup to use correct locations
+    workdir.symlink_to(expt_workdir)
+
+    yield expt_workdir
+    try:
+        shutil.rmtree(expt_workdir)
+    except FileNotFoundError:
+        pass
+    workdir.unlink()
+
 @pytest.fixture
 def access_1year_config():
     # Write an access model config file with 1 year runtime
@@ -229,7 +246,15 @@ def test_access_cice_calendar_cycling_500(
             # which we are trying to bypass.
             shutil.copy(default_input_ice, cice_model.work_path)
 
-            access_model.setup()
+            # Skip writing restart pointer as it requires iced file
+            # with valid header. Restart pointer functionality is tested
+            # in test_cice.py.
+            with patch(
+                'payu.models.cice.Cice.overwrite_restart_ptr',
+                return_value=None
+            ):
+                access_model.setup()
+
             access_model.archive()
 
         end_date_fpath = os.path.join(
@@ -269,7 +294,7 @@ def test_access_cice_1year_runtimes(
     expected_runtime
 ):
     """
-    The large setup/archive cycling test won't pick up situations 
+    The large setup/archive cycling test won't pick up situations
     where the calculations during setup and archive are simultaneously
     wrong, e.g. if they both used the wrong calendar.
     Hence test seperately that the correct runtimes for cice are
@@ -331,7 +356,14 @@ def test_access_cice_1year_runtimes(
         # which we are trying to bypass.
         shutil.copy(ctrl_input_ice_path, cice_model.work_path)
 
-        access_model.setup()
+        # Skip writing restart pointer as it requires iced file
+        # with valid header. Restart pointer functionality is tested
+        # in test_cice.py
+        with patch(
+            'payu.models.cice.Cice.overwrite_restart_ptr',
+            return_value=None
+        ):
+            access_model.setup()
 
         # Check that the correct runtime is written to the work directory's
         # input ice namelist.
diff --git a/test/models/test_cice.py b/test/models/test_cice.py
