Skip to content

Commit

Permalink
ACCESS-OM3: use per-variable checksums from mom restart files
Browse files Browse the repository at this point in the history
  • Loading branch information
dougiesquire committed Feb 23, 2025
1 parent fe4f2e7 commit 399e8d0
Show file tree
Hide file tree
Showing 11 changed files with 109 additions and 60 deletions.
2 changes: 2 additions & 0 deletions src/model_config_tests/exp_test_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion src/model_config_tests/models/accessesm1p5.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/model_config_tests/models/accessom2.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
57 changes: 22 additions & 35 deletions src/model_config_tests/models/accessom3.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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)
Expand All @@ -61,36 +53,31 @@ 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()

Check warning on line 56 in src/model_config_tests/models/accessom3.py

View check run for this annotation

Codecov / codecov/patch

src/model_config_tests/models/accessom3.py#L56

Added line #L56 was not covered by tests

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"

Check warning on line 63 in src/model_config_tests/models/accessom3.py

View check run for this annotation

Codecov / codecov/patch

src/model_config_tests/models/accessom3.py#L63

Added line #L63 was not covered by tests
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 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
Expand Down
3 changes: 3 additions & 0 deletions src/model_config_tests/models/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/model_config_tests/test_bit_reproducibility.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Check warning on line 182 in src/model_config_tests/test_bit_reproducibility.py

View check run for this annotation

Codecov / codecov/patch

src/model_config_tests/test_bit_reproducibility.py#L182

Added line #L182 was not covered by tests

checksums_2d = exp_2day.extract_checksums()

Expand Down
95 changes: 76 additions & 19 deletions tests/resources/access-om3/checksums/1-0-0.json
Original file line number Diff line number Diff line change
@@ -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"
]
}
}
}
3 changes: 0 additions & 3 deletions tests/resources/access-om3/output000/ocean.stats

This file was deleted.

Binary file not shown.
1 change: 1 addition & 0 deletions tests/resources/access-om3/restart000/rpointer.ocn
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
access-om3.mom6.r.1900-01-01-10800.nc
2 changes: 2 additions & 0 deletions tests/test_model_extract_checksums.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 399e8d0

Please sign in to comment.