diff --git a/payu/models/__init__.py b/payu/models/__init__.py index 5c5d2859..c3b4e971 100644 --- a/payu/models/__init__.py +++ b/payu/models/__init__.py @@ -1,6 +1,7 @@ from payu.models.access import Access from payu.models.accessom2 import AccessOm2 from payu.models.cesm_cmeps import AccessOm3 +from payu.models.cable import Cable from payu.models.cice import Cice from payu.models.cice5 import Cice5 from payu.models.gold import Gold @@ -36,6 +37,7 @@ 'ww3': WW3, 'mom6': Mom6, 'qgcm': Qgcm, + 'cable': Cable, # Default 'default': Model, diff --git a/payu/models/cable.py b/payu/models/cable.py new file mode 100644 index 00000000..91bd75d3 --- /dev/null +++ b/payu/models/cable.py @@ -0,0 +1,182 @@ +"""payu.models.cable + ================ + + Driver interface to CABLE + + :copyright: Copyright 2021 Marshall Ward, see AUTHORS for details + :license: Apache License, Version 2.0, see LICENSE for details +""" + +# Standard Library +import glob +import os +import shutil + +# Extensions +import f90nml +import yaml + +# Local +from payu.fsops import mkdir_p +from payu.models.model import Model + + +def _get_forcing_path(variable, year, input_dir, offset=None, repeat=None): + """Return the met forcing file path for a given variable and year. + + Parameters + ---------- + variable : str + Variable name. + year : int + Year value. + input_dir : str + Path to work input directory. + offset : list of int, optional + Offset the current simulation year from `offset[0]` to `offset[1]` + before inferring the met forcing path. + repeat : list of int, optional + Constrain the current simulation year between `repeat[0]` and + `repeat[1]` (inclusive) before inferring the met forcing path. If the + year is outside the interval, the constrained year repeats over the + interval. + + Returns + ------- + path : str + Path (relative to control directory) to the inferred met forcing file. + + Raises + ------ + FileNotFoundError + If unable to infer met forcing path. + """ + if offset: + year += offset[1] - offset[0] + if repeat: + year = repeat[0] + ((year - repeat[0]) % (repeat[1] - repeat[0] + 1)) + pattern = os.path.join(input_dir, f"*{variable}*{year}*.nc") + for path in glob.glob(pattern): + return path + msg = f"Unable to infer met forcing path for variable {variable} for year {year}." + raise FileNotFoundError(msg) + + +class Cable(Model): + + def __init__(self, expt, name, config): + super(Cable, self).__init__(expt, name, config) + + self.model_type = 'cable' + self.default_exec = 'cable' + + self.cable_nml_fname = 'cable.nml' + + self.config_files = [ + self.cable_nml_fname, + 'cable_soilparm.nml', + 'pft_params.nml', + ] + + self.forcing_year_config = 'cable.forcing_year.yaml' + self.optional_config_files = [self.forcing_year_config] + + self.met_forcing_vars = [ + "Rainf", + "Snowf", + "LWdown", + "SWdown", + "PSurf", + "Qair", + "Tair", + "Wind", + ] + + def set_model_pathnames(self): + super(Cable, self).set_model_pathnames() + + # TODO: Check for path in filename%type + self.work_input_path = os.path.join(self.work_path, 'INPUT') + self.work_init_path = self.work_input_path + # TODO: Check for path in filename%restart_out + self.work_restart_path = os.path.join(self.work_path, 'RESTART') + + self.restart_calendar_file = self.model_type + '.res.yaml' + self.restart_calendar_path = os.path.join(self.work_init_path, + self.restart_calendar_file) + + self.cable_nml_path = os.path.join(self.work_path, + self.cable_nml_fname) + + def setup(self): + super(Cable, self).setup() + + self.cable_nml = f90nml.read(self.cable_nml_path) + if self.prior_restart_path: + with open(self.restart_calendar_path, 'r') as restart_file: + self.restart_info = yaml.safe_load(restart_file) + else: + self.restart_info = {'year': self.cable_nml['cable']['ncciy']} + + year = self.cable_nml['cable']['ncciy'] = self.restart_info['year'] + + self.cable_nml['cable']['filename']['restart_in'] = ( + os.path.join('INPUT', 'restart.nc') + ) + self.cable_nml['cable']['filename']['restart_out'] = ( + os.path.join('RESTART', 'restart.nc') + ) + self.cable_nml['cable']['output']['restart'] = True + + forcing_year_config_path = os.path.join(self.work_path, self.forcing_year_config) + if os.path.exists(forcing_year_config_path): + with open(forcing_year_config_path, 'r') as file: + conf = yaml.safe_load(file) + forcing_year_config = conf if conf else {} + for var in self.met_forcing_vars: + path = _get_forcing_path( + var, year, self.work_input_path, **forcing_year_config + ) + self.cable_nml["cable"]["gswpfile"][var] = ( + os.path.relpath(path, start=self.work_path) + ) + + # Write modified namelist file to work dir + self.cable_nml.write( + os.path.join(self.work_path, self.cable_nml_fname), + force=True + ) + + def archive(self, **kwargs): + + # Save model time to restart next run + with open(os.path.join(self.restart_path, + self.restart_calendar_file), 'w') as restart_file: + restart = {'year': self.restart_info['year'] + 1} + restart_file.write(yaml.dump(restart, default_flow_style=False)) + + super(Cable, self).archive() + + # Archive the restart files + mkdir_p(self.restart_path) + + restart_files = [f for f in os.listdir(self.work_restart_path) + if f.endswith('restart.nc')] + + for f in restart_files: + f_src = os.path.join(self.work_restart_path, f) + shutil.move(f_src, self.restart_path) + + os.rmdir(self.work_restart_path) + + # Move all logs into a logs subdir + log_path = os.path.join(self.work_path, 'logs') + mkdir_p(log_path) + log_files = [f for f in os.listdir(self.work_path) + if f.startswith('cable_log')] + for f in log_files: + f_src = os.path.join(self.work_path, f) + shutil.move(f_src, log_path) + + def collate(self): + pass diff --git a/test/models/test_cable.py b/test/models/test_cable.py new file mode 100644 index 00000000..7e62b10d --- /dev/null +++ b/test/models/test_cable.py @@ -0,0 +1,52 @@ +import os +import shutil +import tempfile + +import pytest + +import payu.models.cable as cable + +from test.common import make_random_file + + +class TestGetForcingPath: + """Tests for `payu.models.cable._get_forcing_path()`.""" + + @pytest.fixture() + def input_dir(self): + """Create a temporary input directory and return its path.""" + _input_dir = tempfile.mkdtemp(prefix="payu_test_get_forcing_path") + yield _input_dir + shutil.rmtree(_input_dir) + + @pytest.fixture(autouse=True) + def _make_forcing_inputs(self, input_dir): + """Create forcing inputs from 1900 to 1903.""" + for year in [1900, 1901, 1903]: + make_random_file(os.path.join(input_dir, f"crujra_LWdown_{year}.nc")) + + def test_get_forcing_path(self, input_dir): + """Success case: test correct path can be inferred.""" + assert cable._get_forcing_path("LWdown", 1900, input_dir) == os.path.join( + input_dir, "crujra_LWdown_1900.nc" + ) + + def test_year_offset(self, input_dir): + """Success case: test correct path can be inferred with offset.""" + assert cable._get_forcing_path( + "LWdown", 2000, input_dir, offset=[2000, 1900] + ) == os.path.join(input_dir, "crujra_LWdown_1900.nc") + + def test_year_repeat(self, input_dir): + """Success case: test correct path can be inferred with repeat.""" + assert cable._get_forcing_path( + "LWdown", 1904, input_dir, repeat=[1900, 1903] + ) == os.path.join(input_dir, "crujra_LWdown_1900.nc") + + def test_file_not_found_exception(self, input_dir): + """Failure case: test exception is raised if path cannot be inferred.""" + with pytest.raises( + FileNotFoundError, + match="Unable to infer met forcing path for variable LWdown for year 1904.", + ): + _ = cable._get_forcing_path("LWdown", 1904, input_dir)