Merge pull request #110 from ACCESS-NRI/86-better-om3-checksums

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,31 @@
 """Specific Access-OM3 Model setup and post-processing"""
 
-import re
 from collections import defaultdict
 from pathlib import Path
 from typing import Any
 
+import f90nml
+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"
+        self.wav_in = experiment.control_path / "wav_in"
 
     def set_model_runtime(
         self, years: int = 0, months: int = 0, seconds: int = DEFAULT_RUNTIME_SECONDS
@@ -31,19 +34,20 @@ def set_model_runtime(
         Default is 3 hours"""
         runconfig = Runconfig(self.runconfig)
 
+        # Check that ocean model component is MOM since checksums are obtained from
+        # MOM6 restarts. Fail early if not
+        ocn_model = runconfig.get("ALLCOMP_attributes", "OCN_model")
+        if ocn_model != "mom":
+            raise ValueError(
+                "ACCESS-OM3 reproducibility checks utilize checksums written in MOM6 "
+                "restarts and hence can only be used with ACCESS-OM3 configurations that "
+                f"use MOM6. This configuration uses OCN_model = {ocn_model}."
+            )
+
         if years == months == 0:
             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)
@@ -59,38 +63,42 @@ def set_model_runtime(
 
         runconfig.write()
 
+        # Unfortunately WW3 doesn't (yet) obey the nuopc.runconfig. This should change in a
+        # future release, but for now we have to set WW3 runtime in wav_in. See
+        # https://github.com/COSIMA/access-om3/issues/239
+        if self.wav_in.exists():
+            with open(self.wav_in) as f:
+                nml = f90nml.read(f)
+
+            nml["output_date_nml"]["date"]["restart"]["stride"] = int(n)
+            nml.write(self.wav_in, force=True)
+
     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])
+        with open(mom_restart_pointer) as f:
+            for restart_file in f.readlines():
+                restart = mom_restart_pointer.parent / restart_file.rstrip()
+                rootgrp = Dataset(restart, "r")
+                for variable in sorted(rootgrp.variables):
+                    var = rootgrp[variable]
+                    if "checksum" in var.ncattrs():
+                        output_checksums[variable.strip()].append(var.checksum.strip())
+                rootgrp.close()
 
         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/qa/test_test_access_esm1p5_config.py

@@ -9,7 +9,8 @@ def test_test_access_esm1p5_config_release_release_preindustrial():
     access_esm1p5_configs = RESOURCES_DIR / "access" / "configurations"
     test_config = access_esm1p5_configs / "release-preindustrial+concentrations"
 
-    assert test_config.exists()
+    if not test_config.exists():
+        raise FileNotFoundError(f"The test configuration {test_config} does not exist.")
 
     test_cmd = (
         "model-config-tests -s "

tests/qa/test_test_access_esm1p6_config.py

@@ -11,7 +11,8 @@ def test_test_access_esm1p6_config_release_release_preindustrial():
     # access_esm1p6_configs = RESOURCES_DIR / "access" / "configurations"
     # test_config = access_esm1p6_configs / "release-preindustrial+concentrations"
 
-    # assert test_config.exists()
+    # if not test_config.exists():
+    #     raise FileNotFoundError(f"The test configuration {test_config} does not exist.")
 
     # test_cmd = (
     #     "model-config-tests -s "

tests/qa/test_test_access_om2_config.py

@@ -12,7 +12,8 @@ def test_test_access_om2_config_release_1deg_jra55_ryf():
     access_om2_configs = RESOURCES_DIR / "access-om2" / "configurations"
     test_config = access_om2_configs / "release-1deg_jra55_ryf"
 
-    assert test_config.exists()
+    if not test_config.exists():
+        raise FileNotFoundError(f"The test configuration {test_config} does not exist.")
 
     test_cmd = (
         "model-config-tests -s "

tests/qa/test_test_access_om3_config.py

@@ -0,0 +1,73 @@
+import shlex
+import shutil
+import subprocess
+
+import yaml
+
+from tests.common import RESOURCES_DIR
+
+
+def test_test_access_om3_config_release_1deg_jra55_ryf():
+    """Test ACCESS-OM3 specific config tests"""
+    access_om3_configs = RESOURCES_DIR / "access-om3" / "configurations"
+    test_config = access_om3_configs / "om3-dev-1deg_jra55do_ryf"
+
+    if not test_config.exists():
+        raise FileNotFoundError(f"The test configuration {test_config} does not exist.")
+
+    test_cmd = (
+        "model-config-tests -s "
+        # Run all access_om3 specific tests
+        "-m access_om3 "
+        f"--control-path {test_config} "
+        # Use target branch as can't mock get_git_branch function in utils
+        f"--target-branch om3-dev-1deg_jra55do_ryf"
+    )
+
+    result = subprocess.run(shlex.split(test_cmd), capture_output=True, text=True)
+
+    # Expect the tests to have passed
+    if result.returncode:
+        # Print out test logs if there are errors
+        print(f"Test stdout: {result.stdout}\nTest stderr: {result.stderr}")
+
+    assert result.returncode == 0
+
+
+def test_test_access_om3_config_modified_module_version(tmp_path):
+    """Test changing model module version in config.yaml,
+    will cause tests to fail if paths in exe manifests don't
+    match released spack.location file"""
+    access_om3_configs = RESOURCES_DIR / "access-om3" / "configurations"
+
+    # Copy test configuration
+    test_config = access_om3_configs / "om3-dev-1deg_jra55do_ryf"
+    mock_control_path = tmp_path / "mock_control_path"
+    shutil.copytree(test_config, mock_control_path)
+
+    mock_config = mock_control_path / "config.yaml"
+
+    with open(mock_config) as f:
+        config = yaml.safe_load(f)
+
+    # Use a different released version of access-om3 module
+    config["modules"]["load"] = ["access-om3/2024.09.0"]
+
+    with open(mock_config, "w") as f:
+        yaml.dump(config, f)
+
+    test_cmd = (
+        "model-config-tests -s "
+        # Only test the manifest exe in release spack location test
+        "-k test_access_om3_manifest_exe_in_release_spack_location "
+        f"--control-path {mock_control_path} "
+        # Use target branch as can't mock get_git_branch function in utils
+        f"--target-branch om3-dev-1deg_jra55do_ryf"
+    )
+
+    result = subprocess.run(shlex.split(test_cmd), capture_output=True, text=True)
+
+    # Expect test to have failed
+    assert result.returncode == 1
+    error_msg = "Expected exe path in exe manifest to match an install path in released spack.location"
+    assert error_msg in result.stdout

tests/qa/test_test_config.py

@@ -10,7 +10,8 @@ def test_test_config_access_om2():
     access_om2_configs = RESOURCES_DIR / "access-om2" / "configurations"
     test_config = access_om2_configs / branch_name
 
-    assert test_config.exists()
+    if not test_config.exists():
+        raise FileNotFoundError(f"The test configuration {test_config} does not exist.")
 
     test_cmd = (
         "model-config-tests -s "