ACCESS-OM3: use per-variable checksums from mom restart files

Author: dougiesquire

src/model_config_tests/exp_test_helper.py

@@ -27,6 +27,8 @@ def __init__(self, control_path: Path, lab_path: Path, disable_payu_run=False):
         self.work_path = lab_path / "work" / self.exp_name
         self.output000 = self.archive_path / "output000"
         self.output001 = self.archive_path / "output001"
+        self.restart000 = self.archive_path / "restart000"
+        self.restart001 = self.archive_path / "restart001"
 
         with open(self.config_path) as f:
             self.config = yaml.safe_load(f)

src/model_config_tests/models/accessesm1p5.py

@@ -19,7 +19,7 @@ def __init__(self, experiment):
         # Override model default runtime
         self.default_runtime_seconds = DEFAULT_RUNTIME_SECONDS
 
-        self.output_file = self.experiment.output000 / "access.out"
+        self.output_file = self.output_0 / "access.out"
 
     def set_model_runtime(
         self, years: int = 0, months: int = 0, seconds: int = DEFAULT_RUNTIME_SECONDS

src/model_config_tests/models/accessom2.py

@@ -16,7 +16,7 @@
 class AccessOm2(Model):
     def __init__(self, experiment):
         super().__init__(experiment)
-        self.output_file = self.experiment.output000 / "access-om2.out"
+        self.output_file = self.output_0 / "access-om2.out"
 
         self.accessom2_config = experiment.control_path / "accessom2.nml"
         self.ocean_config = experiment.control_path / "ocean" / "input.nml"

src/model_config_tests/models/accessom3.py

@@ -1,28 +1,29 @@
 """Specific Access-OM3 Model setup and post-processing"""
 
-import re
 from collections import defaultdict
 from pathlib import Path
 from typing import Any
 
+from netCDF4 import Dataset
 from payu.models.cesm_cmeps import Runconfig
 
 from model_config_tests.models.model import (
     DEFAULT_RUNTIME_SECONDS,
     SCHEMA_VERSION_1_0_0,
     Model,
 )
-from model_config_tests.util import DAY_IN_SECONDS
 
 
 class AccessOm3(Model):
     def __init__(self, experiment):
         super().__init__(experiment)
-        self.output_file = self.experiment.output000 / "ocean.stats"
 
+        # ACCESS-OM3 uses restarts for repro testing
+        self.output_0 = self.experiment.restart000
+        self.output_1 = self.experiment.restart001
+
+        self.mom_restart_pointer = self.output_0 / "rpointer.ocn"
         self.runconfig = experiment.control_path / "nuopc.runconfig"
-        self.mom_override = experiment.control_path / "MOM_override"
-        self.ocean_config = experiment.control_path / "input.nml"
 
     def set_model_runtime(
         self, years: int = 0, months: int = 0, seconds: int = DEFAULT_RUNTIME_SECONDS
@@ -35,15 +36,6 @@ def set_model_runtime(
             freq = "nseconds"
             n = str(seconds)
 
-            # Ensure that ocean.stats are written at the end of the run
-            if seconds < DAY_IN_SECONDS:
-                with open(self.mom_override, "a") as f:
-                    f.writelines(
-                        [
-                            f"\n#override TIMEUNIT = {n}",
-                            "\n#override ENERGYSAVEDAYS = 1.0",
-                        ]
-                    )
         elif seconds == 0:
             freq = "nmonths"
             n = str(12 * years + months)
@@ -61,36 +53,32 @@ def set_model_runtime(
 
     def output_exists(self) -> bool:
         """Check for existing output file"""
-        return self.output_file.exists()
+        return self.mom_restart_pointer.exists()
 
     def extract_checksums(
         self, output_directory: Path = None, schema_version: str = None
     ) -> dict[str, Any]:
         """Parse output file and create checksum using defined schema"""
         if output_directory:
-            output_filename = output_directory / "ocean.stats"
+            mom_restart_pointer = output_directory / "rpointer.ocn"
         else:
-            output_filename = self.output_file
-
-        # ocean.stats is used for regression testing in MOM6's own test suite
-        # See https://github.com/mom-ocean/MOM6/blob/2ab885eddfc47fc0c8c0bae46bc61531104428d5/.testing/Makefile#L495-L501
-        # Rows in ocean.stats look like:
-        #      0,  693135.000,     0, En 3.0745627134675957E-23, CFL  0.00000, ...
-        # where the first three columns are Step, Day, Truncs and the remaining
-        # columns include a label for what they are (e.g. En = Energy/Mass)
-        # Header info is only included for new runs so can't be relied on
+            mom_restart_pointer = self.mom_restart_pointer
+
+        # MOM6 saves checksums for each variable in its restart files. Extract these
+        # attributes for each restart
         output_checksums: dict[str, list[any]] = defaultdict(list)
 
-        with open(output_filename) as f:
-            lines = f.readlines()
-            # Skip header if it exists (for new runs)
-            istart = 2 if "Step" in lines[0] else 0
-            for line in lines[istart:]:
-                for col in line.split(","):
-                    # Only keep columns with labels (ie not Step, Day, Truncs)
-                    col = re.split(" +", col.strip().rstrip("\n"))
-                    if len(col) > 1:
-                        output_checksums[col[0]].append(col[-1])
+        mom_restarts = []
+        with open(mom_restart_pointer) as f:
+            for restart in f.readlines():
+                mom_restarts.append(mom_restart_pointer.parent / restart.rstrip())
+
+        for mom_restart in mom_restarts:
+            rootgrp = Dataset(mom_restart, "r")
+            for v in rootgrp.variables:
+                var = rootgrp[v]
+                if "checksum" in var.ncattrs():
+                    output_checksums[var.long_name.strip()].append(var.checksum.strip())
 
         if schema_version is None:
             schema_version = self.default_schema_version

src/model_config_tests/models/model.py

@@ -23,6 +23,9 @@ def __init__(self, experiment):
 
         self.default_runtime_seconds = DEFAULT_RUNTIME_SECONDS
 
+        self.output_0 = self.experiment.output000
+        self.output_1 = self.experiment.output001
+
     def extract_checksums(self, output_directory: Path, schema_version: str):
         """Extract checksums from output directory"""
         raise NotImplementedError

src/model_config_tests/test_bit_reproducibility.py

@@ -179,7 +179,7 @@ def test_restart_repro(self, output_path: Path, control_path: Path):
 
         # Now compare the output between our two short and one long run.
         checksums_1d_0 = exp_2x1day.extract_checksums()
-        checksums_1d_1 = exp_2x1day.extract_checksums(exp_2x1day.output001)
+        checksums_1d_1 = exp_2x1day.extract_checksums(exp_2x1day.model.output_1)
 
         checksums_2d = exp_2day.extract_checksums()
 

tests/resources/access-om3/checksums/1-0-0.json

@@ -1,32 +1,89 @@
 {
   "schema_version": "1-0-0",
   "output": {
-    "En": [
-      "3.0745627134675957E-23"
+    "Potential Temperature": [
+      "E441AC22DEDE8930"
     ],
-    "CFL": [
-      "0.00000"
+    "Salinity": [
+      "8E3D02158AE4EAED"
     ],
-    "SL": [
-      "1.5112E-10"
+    "Layer Thickness": [
+      "9138701970F7E8A"
     ],
-    "M": [
-      "1.36404E+21"
+    "Zonal velocity": [
+      "803BF5B7E9C239C1"
     ],
-    "S": [
-      "34.7263"
+    "Meridional velocity": [
+      "7D9693C22E4D52B5"
     ],
-    "T": [
-      "3.6362"
+    "Frazil heat flux into ocean": [
+      "5FE9C9F905239FC4"
     ],
-    "Me": [
-      "0.00E+00"
+    "Ocean surface pressure used in EoS": [
+      "161DADB852A006F1"
     ],
-    "Se": [
-      "0.00E+00"
+    "Time average sea surface height": [
+      "40E5A21D6F0D1F87"
     ],
-    "Te": [
-      "0.00E+00"
+    "Indicator of the first direction in split calculations.": [
+      "0"
+    ],
+    "Free surface Height": [
+      "68FCD5E9DE45D1DF"
+    ],
+    "Auxiliary Zonal velocity": [
+      "32A58A90C668C32D"
+    ],
+    "Auxiliary Meridional velocity": [
+      "DA212F15199DC960"
+    ],
+    "Zonal Coriolis and advactive acceleration": [
+      "A4C51EC2A34D8EF1"
+    ],
+    "Meridional Coriolis and advactive  acceleration": [
+      "DA7F07794C0B7279"
+    ],
+    "Zonal horizontal viscous acceleration": [
+      "749F723BA63A013C"
+    ],
+    "Meridional horizontal viscous acceleration": [
+      "CD7614674E679DA8"
+    ],
+    "Time mean barotropic zonal velocity": [
+      "2E16365E105FFB9"
+    ],
+    "Time mean barotropic meridional velocity": [
+      "A6344B9A6D8CB3A9"
+    ],
+    "Barotropic timestep": [
+      "4055968059DD63AF"
+    ],
+    "Ideal Age Tracer": [
+      "B64AD181B58D3F4B"
+    ],
+    "Mesoscale Eddy Kinetic Energy": [
+      "A205FF94D0736FD0"
+    ],
+    "Lateral diffusivity from Mesoscale Eddy Kinetic Energy": [
+      "C5E23A24466765D4"
+    ],
+    "Lateral viscosity from Mesoscale Eddy Kinetic Energy": [
+      "A9959752C4D3DE41"
+    ],
+    "Copy of thickness diffusivity for diffusing MEKE": [
+      "C62261CF4678ACEE"
+    ],
+    "Shear-driven turbulent diffusivity at interfaces": [
+      "8A1CB3488BF7FAC7"
+    ],
+    "Shear-driven turbulent viscosity at interfaces": [
+      "55036CBF887185D4"
+    ],
+    "Instantaneous active mixing layer thickness": [
+      "A815BF54039C2D82"
+    ],
+    "Time-filtered MLD for use in MLE": [
+      "7463C5D823501487"
     ]
   }
-}
+}

tests/resources/access-om3/output000/ocean.stats

@@ -0,0 +1 @@
+access-om3.mom6.r.1900-01-01-10800.nc

tests/resources/access-om3/restart000/access-om3.mom6.r.1900-01-01-10800.nc

@@ -19,6 +19,7 @@ def test_extract_checksums(model_name):
     # Mock ExpTestHelper
     mock_experiment = Mock()
     mock_experiment.output000 = resources_dir / "output000"
+    mock_experiment.restart000 = resources_dir / "restart000"
     mock_experiment.control_path = Path("test/tmp")
 
     # Create Model instance
@@ -53,6 +54,7 @@ def test_extract_checksums_unsupported_version(model_name):
     # Mock ExpTestHelper
     mock_experiment = Mock()
     mock_experiment.output000 = resources_dir / "output000"
+    mock_experiment.restart000 = resources_dir / "restart000"
     mock_experiment.control_path = Path("test/tmp")
 
     # Create Model instance