Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add restart pruning to OM3 driver #562

Merged
merged 11 commits into from
Feb 13, 2025
6 changes: 3 additions & 3 deletions docs/source/config.rst
Original file line number Diff line number Diff line change
Expand Up @@ -217,9 +217,9 @@ configuration.
For example, ``restart_freq: 10YS`` would save earliest restart of the year,
10 years from the last permanently archived restart's datetime.

Please note that currently, only ACCESS-OM2, MOM5 and MOM6 models support
date-based restart frequency, as it depends on the payu model driver being
able to parse restarts files for a datetime.
Please note that currently, only ACCESS-ESM1.5, ACCESS-OM2, ACCESS-OM3, MOM5
and MOM6 models support date-based restart frequency, as it depends on the payu
model driver being able to parse restarts files for a datetime.

``restart_history``
Specifies how many of the most recent restart files to retain regardless of
Expand Down
47 changes: 47 additions & 0 deletions payu/models/cesm_cmeps.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import os
import re
import shutil
import cftime
from warnings import warn

from payu.fsops import mkdir_p, make_symlink
Expand All @@ -23,6 +24,12 @@
NUOPC_CONFIG = "nuopc.runconfig"
NUOPC_RUNSEQ = "nuopc.runseq"

#mapping of runconfig to cftime calendars:
CFTIME_CALENDARS = {
"NO_LEAP" : "noleap" ,
"GREGORIAN" : "proleptic_gregorian"
}

# Add as needed
component_info = {
"mom": {
Expand Down Expand Up @@ -330,6 +337,46 @@ def collate(self):
else:
super().collate()

def get_restart_datetime(self, restart_path):
"""Given a restart path, parse the restart files and
return a cftime datetime (for date-based restart pruning)
Supports noleap and proleptic gregorian calendars"""

# Check for rpointer.cpl file
rpointer_path = os.path.join(restart_path, 'rpointer.cpl')
if not os.path.exists(rpointer_path):
raise FileNotFoundError(
'Cannot find rpointer.cpl file, which is required for '
'date-based restart pruning')

with open(rpointer_path, 'r') as ocean_solo:
lines = ocean_solo.readlines()
# example lines would be access-om3.cpl.r.1900-01-02-00000.nc

date_str = lines[0].split('.')[3]
year, month, day, seconds = date_str.split('-')

seconds = int(seconds)
hour = seconds // 3600 ; min = (seconds % 3600) // 60 ; sec = seconds % 60

self.get_runconfig(self.expt.control_path)
run_calendar = self.runconfig.get("CLOCK_attributes", "calendar")

try:
cf_cal = CFTIME_CALENDARS[run_calendar]
except KeyError:
warn(
f"Unsupported calendar for restart pruning: {run_calendar},"
f" try {' or '.join(CFTIME_CALENDARS.keys())}"
)
return False

return cftime.datetime(
int(year), int(month), int(day),
hour, min, sec,
calendar = cf_cal
)


class AccessOm3(CesmCmeps):

Expand Down
110 changes: 110 additions & 0 deletions test/models/access-om3/test_access_om3.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,12 @@
import pytest

import payu
import cftime

from test.common import cd, tmpdir, ctrldir, labdir, workdir, write_config, config_path
from test.common import config as config_orig
from test.common import make_inputs, make_exe
from test.common import list_expt_archive_dirs, make_expt_archive_dir, remove_expt_archive_dirs

MODEL = 'access-om3'

Expand Down Expand Up @@ -312,3 +314,111 @@ def test__setup_checks_bad_io(pio_numiotasks, pio_stride):
model._setup_checks()

teardown_cmeps_config()


# test restart datetime pruning

def make_restart_dir(start_dt):
"""Create restart directory with rpointer.cpl file"""
# Create restart directory
restart_path = make_expt_archive_dir(type='restart')

rpath = os.path.join(restart_path, "rpointer.cpl")
with open(rpath, "w") as rpointer_file:
rpointer_file.write(
f"access-om3.cpl.r.{start_dt}.nc"
)

@pytest.mark.parametrize(
"start_dt, calendar, cmeps_calendar, expected_cftime",
[
(
"0001-01-01-00000",
"proleptic_gregorian",
"GREGORIAN",
cftime.datetime(1, 1, 1, calendar="proleptic_gregorian")
),
(
"9999-12-31-86399",
"proleptic_gregorian",
"GREGORIAN",
cftime.datetime(9999, 12, 31,23,59,59, calendar="proleptic_gregorian")
),
(
"1900-02-01-00000",
"noleap",
"NO_LEAP",
cftime.datetime(1900, 2, 1, calendar="noleap")
),
])
@pytest.mark.filterwarnings("error")
def test_get_restart_datetime(start_dt, calendar, cmeps_calendar, expected_cftime):

cmeps_config(1)

make_restart_dir(start_dt)

test_runconf = {
"CLOCK_attributes": {
"calendar":cmeps_calendar
}
}

with cd(ctrldir):
lab = payu.laboratory.Laboratory(lab_path=str(labdir))
expt = payu.experiment.Experiment(lab, reproduce=False)

model = expt.model
model.get_runconfig = lambda a : (True) # mock reading runconf from file
model.runconfig = MockRunConfig(test_runconf)

print(model.runconfig.get("CLOCK_attributes","calendar"))

restart_path = list_expt_archive_dirs()[0]
parsed_run_dt = expt.model.get_restart_datetime(restart_path)
assert parsed_run_dt == expected_cftime

teardown_cmeps_config()
remove_expt_archive_dirs(type='restart')

@pytest.mark.parametrize(
"start_dt, calendar, cmeps_calendar, expected_cftime",
[
(
"1900-02-01-00000",
"julian",
"JULIAN",
cftime.datetime(1900, 2, 1, calendar="julian")
),
])
def test_get_restart_datetime_badcal(start_dt, calendar, cmeps_calendar, expected_cftime):

cmeps_config(1)

make_restart_dir(start_dt)

test_runconf = {
"CLOCK_attributes": {
"calendar":cmeps_calendar
}
}

with cd(ctrldir):
lab = payu.laboratory.Laboratory(lab_path=str(labdir))
expt = payu.experiment.Experiment(lab, reproduce=False)

model = expt.model
model.get_runconfig = lambda a : (True) # mock reading runconf from file
model.runconfig = MockRunConfig(test_runconf)

print(model.runconfig.get("CLOCK_attributes","calendar"))

restart_path = list_expt_archive_dirs()[0]
with pytest.warns(
Warning, match="Unsupported calendar"
):
expt.model.get_restart_datetime(restart_path)

teardown_cmeps_config()
remove_expt_archive_dirs(type='restart')

Loading