From b8eacbfb24d52b1315c47b687ee1de663586ffb0 Mon Sep 17 00:00:00 2001 From: Emily Carpenter <137525341+elcarpenterNOAA@users.noreply.github.com> Date: Mon, 1 Apr 2024 10:09:55 -0400 Subject: [PATCH] UW-533 & UW-534 MPAS driver for init_atmosphere and atmosphere forecast (#447) * init * WIP * mpas files * Adding the initial bits of an MPAS driver. * WIP * WIP * WIP * WIP * WIP * WIP * WIP * Driver WIP * WIP * edit namelist name * WIP * WIP * WIP * WIP * WIP * Fixing namelist and streams creation. * WIP * WIP * clean up * renaming * WIP * mpas_init.yaml * new line * punctuation * A tested MPAS driver. * Adding schema with updated tests. * WIP * Works for running mpas. * PR feedback * feedback * config optional path * Formatting. * WIP * WIP * Updating schemas for consistency. * WIP * Default dict for mpas boundary files algo. * update CLI * DRY boundary files * doc string * update namelist_file * mpas updates * added try block * remove comment * simplify tests * updated runscript tests and configs --------- Co-authored-by: Christina Holt --- docs/shared/mpas_init.yaml | 62 +++++ src/uwtools/api/mpas.py | 54 +++++ src/uwtools/api/mpas_init.py | 54 +++++ src/uwtools/cli.py | 114 +++++++++ src/uwtools/drivers/mpas.py | 173 ++++++++++++++ src/uwtools/drivers/mpas_init.py | 176 ++++++++++++++ src/uwtools/drivers/ungrib.py | 2 +- .../resources/jsonschema/mpas-init.jsonschema | 94 ++++++++ .../resources/jsonschema/mpas.jsonschema | 93 ++++++++ src/uwtools/strings.py | 2 + src/uwtools/tests/api/test_mpas.py | 35 +++ src/uwtools/tests/api/test_mpas_init.py | 35 +++ src/uwtools/tests/drivers/test_fv3.py | 2 +- src/uwtools/tests/drivers/test_mpas.py | 209 +++++++++++++++++ src/uwtools/tests/drivers/test_mpas_init.py | 218 ++++++++++++++++++ src/uwtools/tests/test_cli.py | 64 +++++ 16 files changed, 1385 insertions(+), 2 deletions(-) create mode 100644 docs/shared/mpas_init.yaml create mode 100644 src/uwtools/api/mpas.py create mode 100644 src/uwtools/api/mpas_init.py create mode 100644 src/uwtools/drivers/mpas.py create mode 100644 src/uwtools/drivers/mpas_init.py create mode 100644 src/uwtools/resources/jsonschema/mpas-init.jsonschema create mode 100644 src/uwtools/resources/jsonschema/mpas.jsonschema create mode 100644 src/uwtools/tests/api/test_mpas.py create mode 100644 src/uwtools/tests/api/test_mpas_init.py create mode 100644 src/uwtools/tests/drivers/test_mpas.py create mode 100644 src/uwtools/tests/drivers/test_mpas_init.py diff --git a/docs/shared/mpas_init.yaml b/docs/shared/mpas_init.yaml new file mode 100644 index 000000000..c8402b925 --- /dev/null +++ b/docs/shared/mpas_init.yaml @@ -0,0 +1,62 @@ +user: + mpas_app: /path/to/mpas_app +platform: + account: me + scheduler: slurm +mpas_init: + boundary_conditions: + interval_hours: 6 + length: 6 + execution: + executable: "{{ user.mpas_app }}/exec/init_atmosphere_model" + batchargs: + walltime: 01:30:00 + cores: 4 + mpiargs: + - "--ntasks=4" + mpicmd: srun + envcmds: + - module use {{ user.mpas_app }}/modulefiles + - module load build_jet_intel + files_to_copy: + conus.static.nc: /path/to/conus.static.nc + conus.graph.info.part.{{mpas_init.execution["batchargs"]["cores"]}}: /path/to/conus.graph.info.part.{{mpas_init.execution["batchargs"]["cores"]}} + conus.init.nc: /path/to/conus.init.nc + stream_list.atmosphere.diagnostics: "{{ user.mpas_app }}/src/MPAS-Model/stream_list.atmosphere.diagnostics" + stream_list.atmosphere.output: "{{ user.mpas_app }}/src/MPAS-Model/stream_list.atmosphere.output" + stream_list.atmosphere.surface: "{{ user.mpas_app }}/src/MPAS-Model/stream_list.atmosphere.surface" + files_to_link: + CAM_ABS_DATA.DBL: "{{ user.mpas_app }}/src/MPAS-Model/CAM_ABS_DATA.DBL" + CAM_AEROPT_DATA.DBL: "{{ user.mpas_app }}/src/MPAS-Model/CAM_AEROPT_DATA.DBL" + GENPARM.TBL: "{{ user.mpas_app }}/src/MPAS-Model/GENPARM.TBL" + LANDUSE.TBL: "{{ user.mpas_app }}/src/MPAS-Model/LANDUSE.TBL" + OZONE_DAT.TBL: "{{ user.mpas_app }}/src/MPAS-Model/OZONE_DAT.TBL" + OZONE_LAT.TBL: "{{ user.mpas_app }}/src/MPAS-Model/OZONE_LAT.TBL" + OZONE_PLEV.TBL: "{{ user.mpas_app }}/src/MPAS-Model/OZONE_PLEV.TBL" + RRTMG_LW_DATA: "{{ user.mpas_app }}/src/MPAS-Model/RRTMG_LW_DATA" + RRTMG_LW_DATA.DBL: "{{ user.mpas_app }}/src/MPAS-Model/RRTMG_LW_DATA.DBL" + RRTMG_SW_DATA: "{{ user.mpas_app }}/src/MPAS-Model/RRTMG_SW_DATA" + RRTMG_SW_DATA.DBL: "{{ user.mpas_app }}/src/MPAS-Model/RRTMG_SW_DATA.DBL" + SOILPARM.TBL: "{{ user.mpas_app }}/src/MPAS-Model/SOILPARM.TBL" + VEGPARM.TBL: "{{ user.mpas_app }}/src/MPAS-Model/VEGPARM.TBL" + namelist: + base_file: "{{ user.mpas_app }}/src/MPAS-Model/namelist.init_atmosphere" + update_values: + nhyd_model: + config_init_case: 9 + data_sources: + config_met_prefix: FILE + config_fg_interval: !int "{{ mpas_init.boundary_conditions['interval_hours'] * 3600 }}" + vertical_grid: + config_blend_bdy_terrain: true + preproc_stages: + config_static_interp: false + config_native_gwd_static: false + decomposition: + config_block_decomp_file_prefix: conus.graph.info.part. + run_dir: /path/to/rundir + streams: + path: /path/to/streams.init_atmosphere.IN + values: + input_filename: conus.init.nc + output_filename: foo.nc diff --git a/src/uwtools/api/mpas.py b/src/uwtools/api/mpas.py new file mode 100644 index 000000000..dfcee9fc8 --- /dev/null +++ b/src/uwtools/api/mpas.py @@ -0,0 +1,54 @@ +""" +API access to the ``uwtools`` ``mpas`` driver. +""" + +import datetime as dt +from pathlib import Path +from typing import Dict, Optional + +import uwtools.drivers.support as _support +from uwtools.drivers.mpas import MPAS as _MPAS + + +def execute( + task: str, + cycle: dt.datetime, + config: Optional[Path] = None, + batch: bool = False, + dry_run: bool = False, + graph_file: Optional[Path] = None, +) -> bool: + """ + Execute an ``mpas`` task. + + If ``batch`` is specified, a runscript will be written and submitted to the batch system. + Otherwise, the executable will be run directly on the current system. + + :param task: The task to execute. + :param cycle: The cycle. + :param config: Path to config file (read stdin if missing or None). + :param batch: Submit run to the batch system. + :param dry_run: Do not run the executable, just report what would have been done. + :param graph_file: Write Graphviz DOT output here. + :return: ``True`` if task completes without raising an exception. + """ + obj = _MPAS(config=config, cycle=cycle, batch=batch, dry_run=dry_run) + getattr(obj, task)() + if graph_file: + with open(graph_file, "w", encoding="utf-8") as f: + print(graph(), file=f) + return True + + +def graph() -> str: + """ + Returns Graphviz DOT code for the most recently executed task. + """ + return _support.graph() + + +def tasks() -> Dict[str, str]: + """ + Returns a mapping from task names to their one-line descriptions. + """ + return _support.tasks(_MPAS) diff --git a/src/uwtools/api/mpas_init.py b/src/uwtools/api/mpas_init.py new file mode 100644 index 000000000..a698e5705 --- /dev/null +++ b/src/uwtools/api/mpas_init.py @@ -0,0 +1,54 @@ +""" +API access to the ``uwtools`` ``mpas-init`` driver. +""" + +import datetime as dt +from pathlib import Path +from typing import Dict, Optional + +import uwtools.drivers.support as _support +from uwtools.drivers.mpas_init import MPASInit as _MPASInit + + +def execute( + task: str, + cycle: dt.datetime, + config: Optional[Path] = None, + batch: bool = False, + dry_run: bool = False, + graph_file: Optional[Path] = None, +) -> bool: + """ + Execute an MPAS ``init-atmosphere`` task. + + If ``batch`` is specified, a runscript will be written and submitted to the batch system. + Otherwise, the executable will be run directly on the current system. + + :param task: The task to execute. + :param config: Path to config file (read stdin if missing or None). + :param cycle: The cycle. + :param batch: Submit run to the batch system? + :param dry_run: Do not run the executable, just report what would have been done. + :param graph_file: Write Graphviz DOT output here. + :return: ``True`` if task completes without raising an exception. + """ + obj = _MPASInit(config=config, cycle=cycle, batch=batch, dry_run=dry_run) + getattr(obj, task)() + if graph_file: + with open(graph_file, "w", encoding="utf-8") as f: + print(graph(), file=f) + return True + + +def graph() -> str: + """ + Returns Graphviz DOT code for the most recently executed task. + """ + return _support.graph() + + +def tasks() -> Dict[str, str]: + """ + Returns a mapping from task names to their one-line descriptions. + """ + return _support.tasks(_MPASInit) diff --git a/src/uwtools/cli.py b/src/uwtools/cli.py index 30adb7114..e46fc063f 100644 --- a/src/uwtools/cli.py +++ b/src/uwtools/cli.py @@ -17,6 +17,8 @@ import uwtools.api.config import uwtools.api.file import uwtools.api.fv3 +import uwtools.api.mpas +import uwtools.api.mpas_init import uwtools.api.rocoto import uwtools.api.sfc_climo_gen import uwtools.api.template @@ -61,6 +63,8 @@ def main() -> None: STR.config: _dispatch_config, STR.file: _dispatch_file, STR.fv3: _dispatch_fv3, + STR.mpas: _dispatch_mpas, + STR.mpasinit: _dispatch_mpas_init, STR.rocoto: _dispatch_rocoto, STR.sfcclimogen: _dispatch_sfc_climo_gen, STR.template: _dispatch_template, @@ -425,6 +429,114 @@ def _dispatch_fv3(args: Args) -> bool: ) +# Mode mpas + + +def _add_subparser_mpas(subparsers: Subparsers) -> ModeChecks: + """ + Subparser for mode: mpas + + :param subparsers: Parent parser's subparsers, to add this subparser to. + """ + parser = _add_subparser(subparsers, STR.mpas, "Execute MPAS tasks") + _basic_setup(parser) + subparsers = _add_subparsers(parser, STR.action, STR.task.upper()) + return { + task: _add_subparser_mpas_task(subparsers, task, helpmsg) + for task, helpmsg in uwtools.api.mpas.tasks().items() + } + + +def _add_subparser_mpas_task(subparsers: Subparsers, task: str, helpmsg: str) -> ActionChecks: + """ + Subparser for mode: mpas + + :param subparsers: Parent parser's subparsers, to add this subparser to. + :param task: The task to add a subparser for. + :param helpmsg: Help message for task. + """ + parser = _add_subparser(subparsers, task, helpmsg.rstrip(".")) + required = parser.add_argument_group(TITLE_REQ_ARG) + _add_arg_cycle(required) + optional = _basic_setup(parser) + _add_arg_config_file(group=optional, required=False) + _add_arg_batch(optional) + _add_arg_dry_run(optional) + _add_arg_graph_file(optional) + checks = _add_args_verbosity(optional) + return checks + + +def _dispatch_mpas(args: Args) -> bool: + """ + Dispatch logic for mpas mode. + + :param args: Parsed command-line args. + """ + return uwtools.api.mpas.execute( + task=args[STR.action], + config=args[STR.cfgfile], + cycle=args[STR.cycle], + batch=args[STR.batch], + dry_run=args[STR.dryrun], + graph_file=args[STR.graphfile], + ) + + +# Mode mpas_init + + +def _add_subparser_mpas_init(subparsers: Subparsers) -> ModeChecks: + """ + Subparser for mode: mpas_init + + :param subparsers: Parent parser's subparsers, to add this subparser to. + """ + parser = _add_subparser(subparsers, STR.mpasinit, "Execute MPASInit tasks") + _basic_setup(parser) + subparsers = _add_subparsers(parser, STR.action, STR.task.upper()) + return { + task: _add_subparser_mpas_init_task(subparsers, task, helpmsg) + for task, helpmsg in uwtools.api.mpas_init.tasks().items() + } + + +def _add_subparser_mpas_init_task(subparsers: Subparsers, task: str, helpmsg: str) -> ActionChecks: + """ + Subparser for mode: mpas_init + + :param subparsers: Parent parser's subparsers, to add this subparser to. + :param task: The task to add a subparser for. + :param helpmsg: Help message for task. + """ + parser = _add_subparser(subparsers, task, helpmsg.rstrip(".")) + required = parser.add_argument_group(TITLE_REQ_ARG) + _add_arg_cycle(required) + optional = _basic_setup(parser) + _add_arg_config_file(group=optional, required=False) + _add_arg_batch(optional) + _add_arg_dry_run(optional) + _add_arg_graph_file(optional) + checks = _add_args_verbosity(optional) + return checks + + +def _dispatch_mpas_init(args: Args) -> bool: + """ + Dispatch logic for mpas_init mode. + + :param args: Parsed command-line args. + """ + return uwtools.api.mpas_init.execute( + task=args[STR.action], + config=args[STR.cfgfile], + cycle=args[STR.cycle], + batch=args[STR.batch], + dry_run=args[STR.dryrun], + graph_file=args[STR.graphfile], + ) + + # Mode rocoto @@ -1092,6 +1204,8 @@ def _parse_args(raw_args: List[str]) -> Tuple[Args, Checks]: STR.config: _add_subparser_config(subparsers), STR.file: _add_subparser_file(subparsers), STR.fv3: _add_subparser_fv3(subparsers), + STR.mpas: _add_subparser_mpas(subparsers), + STR.mpasinit: _add_subparser_mpas_init(subparsers), STR.rocoto: _add_subparser_rocoto(subparsers), STR.sfcclimogen: _add_subparser_sfc_climo_gen(subparsers), STR.template: _add_subparser_template(subparsers), diff --git a/src/uwtools/drivers/mpas.py b/src/uwtools/drivers/mpas.py new file mode 100644 index 000000000..65714e754 --- /dev/null +++ b/src/uwtools/drivers/mpas.py @@ -0,0 +1,173 @@ +""" +A driver for the MPAS component. +""" + +from datetime import datetime, timedelta +from pathlib import Path +from typing import Optional + +from iotaa import asset, dryrun, task, tasks + +from uwtools.api.template import render +from uwtools.config.formats.nml import NMLConfig +from uwtools.drivers.driver import Driver +from uwtools.exceptions import UWConfigError +from uwtools.strings import STR +from uwtools.utils.tasks import file, filecopy, symlink + + +class MPAS(Driver): + """ + A driver for MPAS. + """ + + def __init__( + self, + cycle: datetime, + config: Optional[Path] = None, + dry_run: bool = False, + batch: bool = False, + ): + """ + The driver. + + :param cycle: The cycle. + :param config: Path to config file (read stdin if missing or None). + :param dry_run: Run in dry-run mode? + :param batch: Run component via the batch system? + """ + super().__init__(config=config, dry_run=dry_run, batch=batch, cycle=cycle) + if self._dry_run: + dryrun() + self._cycle = cycle + + # Workflow tasks + + @task + def boundary_files(self): + """ + Boundary files. + """ + yield self._taskname("boundary files") + lbcs = self._driver_config["lateral_boundary_conditions"] + endhour = self._driver_config["length"] + interval = lbcs["interval_hours"] + symlinks = {} + for boundary_hour in range(0, endhour + 1, interval): + file_date = self._cycle + timedelta(hours=boundary_hour) + fn = f"lbc.{file_date.strftime('%Y-%m-%d_%H.%M.%S')}.nc" + linkname = self._rundir / fn + symlinks[linkname] = Path(lbcs["path"]) / fn + yield [symlink(target=t, linkname=l) for l, t in symlinks.items()] + + @tasks + def files_copied(self): + """ + Files copied for run. + """ + yield self._taskname("files copied") + yield [ + filecopy(src=Path(src), dst=self._rundir / dst) + for dst, src in self._driver_config.get("files_to_copy", {}).items() + ] + + @tasks + def files_linked(self): + """ + Files linked for run. + """ + yield self._taskname("files linked") + yield [ + symlink(target=Path(target), linkname=self._rundir / linkname) + for linkname, target in self._driver_config.get("files_to_link", {}).items() + ] + + @task + def namelist_file(self): + """ + The namelist file. + """ + path = self._rundir / "namelist.atmosphere" + yield self._taskname(str(path)) + yield asset(path, path.is_file) + yield None + duration = timedelta(hours=self._driver_config["length"]) + str_duration = str(duration).replace(" days, ", "") + try: + namelist = self._driver_config["namelist"] + except KeyError as e: + raise UWConfigError( + "Provide either a 'namelist' YAML block or the %s file" % path + ) from e + update_values = namelist.get("update_values", {}) + update_values.setdefault("nhyd_model", {}).update( + { + "config_start_time": self._cycle.strftime("%Y-%m-%d_%H:00:00"), + "config_run_duration": str_duration, + } + ) + namelist["update_values"] = update_values + self._create_user_updated_config( + config_class=NMLConfig, + config_values=namelist, + path=path, + ) + + @tasks + def provisioned_run_directory(self): + """ + Run directory provisioned with all required content. + """ + yield self._taskname("provisioned run directory") + yield [ + self.boundary_files(), + self.files_copied(), + self.files_linked(), + self.namelist_file(), + self.runscript(), + self.streams_file(), + ] + + @task + def runscript(self): + """ + The runscript. + """ + path = self._runscript_path + yield self._taskname(path.name) + yield asset(path, path.is_file) + yield None + self._write_runscript(path=path, envvars={}) + + @task + def streams_file(self): + """ + The streams file. + """ + fn = "streams.atmosphere" + yield self._taskname(fn) + path = self._rundir / fn + yield asset(path, path.is_file) + yield file(path=Path(self._driver_config["streams"]["path"])) + render( + input_file=Path(self._driver_config["streams"]["path"]), + output_file=path, + values_src=self._driver_config["streams"]["values"], + ) + + # Private helper methods + + @property + def _driver_name(self) -> str: + """ + Returns the name of this driver. + """ + return STR.mpas + + def _taskname(self, suffix: str) -> str: + """ + Returns a common tag for graph-task log messages. + + :param suffix: Log-string suffix. + """ + return "%s %s %s" % (self._cycle.strftime("%Y%m%d %HZ"), self._driver_name, suffix) diff --git a/src/uwtools/drivers/mpas_init.py b/src/uwtools/drivers/mpas_init.py new file mode 100644 index 000000000..ab6981a55 --- /dev/null +++ b/src/uwtools/drivers/mpas_init.py @@ -0,0 +1,176 @@ +""" +A driver for the mpas-init component. +""" + +from datetime import datetime, timedelta +from pathlib import Path +from typing import Optional + +from iotaa import asset, dryrun, task, tasks + +from uwtools.api.template import render +from uwtools.config.formats.nml import NMLConfig +from uwtools.drivers.driver import Driver +from uwtools.exceptions import UWConfigError +from uwtools.strings import STR +from uwtools.utils.tasks import file, filecopy, symlink + + +class MPASInit(Driver): + """ + A driver for mpas-init. + """ + + def __init__( + self, + cycle: datetime, + config: Optional[Path] = None, + dry_run: bool = False, + batch: bool = False, + ): + """ + The driver. + + :param config_file: Path to config file (read stdin if missing or None). + :param cycle: The cycle. + :param dry_run: Run in dry-run mode? + :param batch: Run component via the batch system? + """ + super().__init__(config=config, cycle=cycle, dry_run=dry_run, batch=batch) + if self._dry_run: + dryrun() + self._cycle = cycle + + # Workflow tasks + + @tasks + def boundary_files(self): + """ + Boundary files. + """ + yield self._taskname("boundary files") + lbcs = self._driver_config["boundary_conditions"] + endhour = lbcs["length"] + interval = lbcs["interval_hours"] + symlinks = {} + boundary_filepath = lbcs["path"] + for boundary_hour in range(0, endhour + 1, interval): + file_date = self._cycle + timedelta(hours=boundary_hour) + fn = f"FILE:{file_date.strftime('%Y-%m-%d_%H')}" + target = Path(boundary_filepath) / fn + linkname = self._rundir / fn + symlinks[target] = linkname + yield [symlink(target=t, linkname=l) for t, l in symlinks.items()] + + @tasks + def files_copied(self): + """ + Files copied for run. + """ + yield self._taskname("files copied") + yield [ + filecopy(src=Path(src), dst=self._rundir / dst) + for dst, src in self._driver_config.get("files_to_copy", {}).items() + ] + + @tasks + def files_linked(self): + """ + Files linked for run. + """ + yield self._taskname("files linked") + yield [ + symlink(target=Path(target), linkname=self._rundir / linkname) + for linkname, target in self._driver_config.get("files_to_link", {}).items() + ] + + @task + def namelist_file(self): + """ + The namelist file. + """ + fn = "namelist.init_atmosphere" + yield self._taskname(fn) + path = self._rundir / fn + yield asset(path, path.is_file) + yield None + duration = timedelta(hours=self._driver_config["boundary_conditions"]["length"]) + str_duration = str(duration).replace(" days, ", "") + try: + namelist = self._driver_config["namelist"] + except KeyError as e: + raise UWConfigError( + "Provide either a 'namelist' YAML block or the %s file" % path + ) from e + update_values = namelist.get("update_values", {}) + update_values.setdefault("nhyd_model", {}).update( + { + "config_start_time": self._cycle.strftime("%Y-%m-%d_%H:00:00"), + "config_run_duration": str_duration, + } + ) + namelist["update_values"] = update_values + self._create_user_updated_config( + config_class=NMLConfig, + config_values=namelist, + path=path, + ) + + @tasks + def provisioned_run_directory(self): + """ + Run directory provisioned with all required content. + """ + yield self._taskname("provisioned run directory") + yield [ + self.boundary_files(), + self.files_copied(), + self.files_linked(), + self.namelist_file(), + self.runscript(), + self.streams_file(), + ] + + @task + def runscript(self): + """ + The runscript. + """ + path = self._runscript_path + yield self._taskname(path.name) + yield asset(path, path.is_file) + yield None + self._write_runscript(path=path, envvars={}) + + @task + def streams_file(self): + """ + The streams file. + """ + fn = "streams.init_atmosphere" + yield self._taskname(fn) + path = self._rundir / fn + yield asset(path, path.is_file) + yield file(path=Path(self._driver_config["streams"]["path"])) + render( + input_file=Path(self._driver_config["streams"]["path"]), + output_file=path, + values_src=self._driver_config["streams"]["values"], + ) + + # Private helper methods + + @property + def _driver_name(self) -> str: + """ + Returns the name of this driver. + """ + return STR.mpasinit + + def _taskname(self, suffix: str) -> str: + """ + Returns a common tag for graph-task log messages. + + :param suffix: Log-string suffix. + """ + return "%s %s %s" % (self._cycle.strftime("%Y%m%d %HZ"), self._driver_name, suffix) diff --git a/src/uwtools/drivers/ungrib.py b/src/uwtools/drivers/ungrib.py index 905342b7b..3da7d6a8d 100644 --- a/src/uwtools/drivers/ungrib.py +++ b/src/uwtools/drivers/ungrib.py @@ -163,7 +163,7 @@ def _taskname(self, suffix: str) -> str: :param suffix: Log-string suffix. """ - return "%s ungrib %s" % (self._cycle.strftime("%Y%m%d %HZ"), suffix) + return "%s %s %s" % (self._cycle.strftime("%Y%m%d %HZ"), self._driver_name, suffix) def _ext(n): diff --git a/src/uwtools/resources/jsonschema/mpas-init.jsonschema b/src/uwtools/resources/jsonschema/mpas-init.jsonschema new file mode 100644 index 000000000..7c3fd2ffc --- /dev/null +++ b/src/uwtools/resources/jsonschema/mpas-init.jsonschema @@ -0,0 +1,94 @@ +{ + "properties": { + "mpas_init": { + "additionalProperties": false, + "properties": { + "boundary_conditions": { + "additionalProperties": false, + "properties": { + "interval_hours": { + "minimum": 1, + "type": "integer" + }, + "length": { + "minimum": 1, + "type": "integer" + }, + "offset": { + "minimum": 0, + "type": "integer" + }, + "path": { + "type": "string" + } + }, + "required": [ + "length", + "interval_hours", + "offset", + "path" + ], + "type": "object" + }, + "execution": { + "$ref": "urn:uwtools:execution" + }, + "files_to_copy": { + "$ref": "urn:uwtools:files-to-stage" + }, + "files_to_link": { + "$ref": "urn:uwtools:files-to-stage" + }, + "namelist": { + "additionalProperties": false, + "anyOf": [ + { + "required": [ + "base_file" + ] + }, + { + "required": [ + "update_values" + ] + } + ], + "properties": { + "base_file": { + "type": "string" + }, + "update_values": { + "$ref": "urn:uwtools:namelist" + } + }, + "type": "object" + }, + "run_dir": { + "type": "string" + }, + "streams": { + "properties": { + "path": { + "type": "string" + }, + "values": { + "type": "object" + } + }, + "required": [ + "path", + "values" + ], + "type": "object" + } + }, + "required": [ + "execution", + "run_dir", + "streams" + ], + "type": "object" + } + }, + "type": "object" +} diff --git a/src/uwtools/resources/jsonschema/mpas.jsonschema b/src/uwtools/resources/jsonschema/mpas.jsonschema new file mode 100644 index 000000000..b95904f38 --- /dev/null +++ b/src/uwtools/resources/jsonschema/mpas.jsonschema @@ -0,0 +1,93 @@ +{ + "properties": { + "mpas": { + "additionalProperties": false, + "properties": { + "execution": { + "$ref": "urn:uwtools:execution" + }, + "files_to_copy": { + "$ref": "urn:uwtools:files-to-stage" + }, + "files_to_link": { + "$ref": "urn:uwtools:files-to-stage" + }, + "lateral_boundary_conditions": { + "additionalProperties": false, + "properties": { + "interval_hours": { + "minimum": 1, + "type": "integer" + }, + "offset": { + "minimum": 0, + "type": "integer" + }, + "path": { + "type": "string" + } + }, + "required": [ + "interval_hours", + "offset", + "path" + ], + "type": "object" + }, + "length": { + "minimum": 1, + "type": "integer" + }, + "namelist": { + "additionalProperties": false, + "anyOf": [ + { + "required": [ + "base_file" + ] + }, + { + "required": [ + "update_values" + ] + } + ], + "properties": { + "base_file": { + "type": "string" + }, + "update_values": { + "$ref": "urn:uwtools:namelist" + } + }, + "type": "object" + }, + "run_dir": { + "type": "string" + }, + "streams": { + "properties": { + "path": { + "type": "string" + }, + "values": { + "type": "object" + } + }, + "required": [ + "path", + "values" + ], + "type": "object" + } + }, + "required": [ + "execution", + "run_dir", + "streams" + ], + "type": "object" + } + }, + "type": "object" +} diff --git a/src/uwtools/strings.py b/src/uwtools/strings.py index 614b30c28..2122069fe 100644 --- a/src/uwtools/strings.py +++ b/src/uwtools/strings.py @@ -86,6 +86,8 @@ class STR: link: str = "link" mode: str = "mode" model: str = "model" + mpas: str = "mpas" + mpasinit: str = "mpas_init" outfile: str = "output_file" outfmt: str = "output_format" partial: str = "partial" diff --git a/src/uwtools/tests/api/test_mpas.py b/src/uwtools/tests/api/test_mpas.py new file mode 100644 index 000000000..a406edc2b --- /dev/null +++ b/src/uwtools/tests/api/test_mpas.py @@ -0,0 +1,35 @@ +# pylint: disable=missing-function-docstring,protected-access + +import datetime as dt +from unittest.mock import patch + +from uwtools.api import mpas + + +def test_execute(tmp_path): + cycle = dt.datetime.now() + dot = tmp_path / "graph.dot" + args: dict = { + "batch": False, + "config": "config.yaml", + "cycle": cycle, + "dry_run": True, + "graph_file": dot, + } + with patch.object(mpas, "_MPAS") as MPAS: + assert mpas.execute(**args, task="foo") is True + del args["graph_file"] + MPAS.assert_called_once_with(**args) + MPAS().foo.assert_called_once_with() + + +def test_graph(): + with patch.object(mpas._support, "graph") as graph: + mpas.graph() + graph.assert_called_once_with() + + +def test_tasks(): + with patch.object(mpas._support, "tasks") as _tasks: + mpas.tasks() + _tasks.assert_called_once_with(mpas._MPAS) diff --git a/src/uwtools/tests/api/test_mpas_init.py b/src/uwtools/tests/api/test_mpas_init.py new file mode 100644 index 000000000..36568807d --- /dev/null +++ b/src/uwtools/tests/api/test_mpas_init.py @@ -0,0 +1,35 @@ +# pylint: disable=missing-function-docstring,protected-access + +import datetime as dt +from unittest.mock import patch + +from uwtools.api import mpas_init + + +def test_execute(tmp_path): + cycle = dt.datetime.now() + dot = tmp_path / "graph.dot" + args: dict = { + "batch": False, + "config": "config.yaml", + "cycle": cycle, + "dry_run": True, + "graph_file": dot, + } + with patch.object(mpas_init, "_MPASInit") as MPASInit: + assert mpas_init.execute(**args, task="foo") is True + del args["graph_file"] + MPASInit.assert_called_once_with(**args) + MPASInit().foo.assert_called_once_with() + + +def test_graph(): + with patch.object(mpas_init._support, "graph") as graph: + mpas_init.graph() + graph.assert_called_once_with() + + +def test_tasks(): + with patch.object(mpas_init._support, "tasks") as _tasks: + mpas_init.tasks() + _tasks.assert_called_once_with(mpas_init._MPASInit) diff --git a/src/uwtools/tests/drivers/test_fv3.py b/src/uwtools/tests/drivers/test_fv3.py index f5d1860b4..697f9157b 100644 --- a/src/uwtools/tests/drivers/test_fv3.py +++ b/src/uwtools/tests/drivers/test_fv3.py @@ -120,7 +120,7 @@ def test_FV3_field_table(driverobj): "key,task,test", [("files_to_copy", "files_copied", "is_file"), ("files_to_link", "files_linked", "is_symlink")], ) -def test_FV3_files_copied(config, cycle, key, task, test, tmp_path): +def test_FV3_files_copied_and_linked(config, cycle, key, task, test, tmp_path): atm, sfc = "gfs.t%sz.atmanl.nc", "gfs.t%sz.sfcanl.nc" atm_cfg_dst, sfc_cfg_dst = [x % "{{ cycle.strftime('%H') }}" for x in [atm, sfc]] atm_cfg_src, sfc_cfg_src = [str(tmp_path / (x + ".in")) for x in [atm_cfg_dst, sfc_cfg_dst]] diff --git a/src/uwtools/tests/drivers/test_mpas.py b/src/uwtools/tests/drivers/test_mpas.py new file mode 100644 index 000000000..4b44f32a3 --- /dev/null +++ b/src/uwtools/tests/drivers/test_mpas.py @@ -0,0 +1,209 @@ +# pylint: disable=missing-function-docstring,protected-access,redefined-outer-name +""" +MPAS driver tests. +""" +import datetime as dt +from pathlib import Path +from unittest.mock import DEFAULT as D +from unittest.mock import patch + +import f90nml # type: ignore +import pytest +import yaml +from pytest import fixture, raises + +from uwtools.drivers import mpas +from uwtools.exceptions import UWConfigError +from uwtools.scheduler import Slurm +from uwtools.tests.support import fixture_path + +# Fixtures + + +@fixture +def cycle(): + return dt.datetime(2024, 3, 22, 6) + + +# Driver fixtures + + +@fixture +def config(tmp_path): + return { + "mpas": { + "execution": { + "executable": "atmosphere_model", + "batchargs": { + "walltime": "01:30:00", + }, + }, + "lateral_boundary_conditions": { + "interval_hours": 1, + "offset": 0, + "path": str(tmp_path / "input_files"), + }, + "length": 1, + "namelist": { + "base_file": str(fixture_path("simple.nml")), + "update_values": { + "nhyd_model": {"config_start_time": "12", "config_stop_time": "12"}, + }, + }, + "run_dir": str(tmp_path), + "streams": { + "path": str(tmp_path / "streams.atmosphere.in"), + "values": { + "world": "user", + }, + }, + "files_to_link": { + "CAM_ABS_DATA.DBL": "src/MPAS-Model/CAM_ABS_DATA.DBL", + "CAM_AEROPT_DATA.DBL": "src/MPAS-Model/CAM_AEROPT_DATA.DBL", + }, + }, + "platform": { + "account": "me", + "scheduler": "slurm", + }, + } + + +@fixture +def config_file(config, tmp_path): + path = tmp_path / "config.yaml" + with open(path, "w", encoding="utf-8") as f: + yaml.dump(config, f) + return path + + +@fixture +def driverobj(config_file, cycle): + return mpas.MPAS(config=config_file, cycle=cycle, batch=True) + + +# Driver tests + + +def test_MPAS(driverobj): + assert isinstance(driverobj, mpas.MPAS) + + +def test_MPAS_boundary_files(driverobj, cycle): + ns = (0, 1) + links = [ + driverobj._rundir / f"lbc.{(cycle+dt.timedelta(hours=n)).strftime('%Y-%m-%d_%H.%M.%S')}.nc" + for n in ns + ] + assert not any(link.is_file() for link in links) + infile_path = Path(driverobj._driver_config["lateral_boundary_conditions"]["path"]) + infile_path.mkdir() + for n in ns: + ( + infile_path / f"lbc.{(cycle+dt.timedelta(hours=n)).strftime('%Y-%m-%d_%H.%M.%S')}.nc" + ).touch() + driverobj.boundary_files() + assert all(link.is_symlink() for link in links) + + +def test_MPAS_dry_run(config_file, cycle): + with patch.object(mpas, "dryrun") as dryrun: + driverobj = mpas.MPAS(config=config_file, cycle=cycle, batch=True, dry_run=True) + assert driverobj._dry_run is True + dryrun.assert_called_once_with() + + +@pytest.mark.parametrize( + "key,task,test", + [("files_to_copy", "files_copied", "is_file"), ("files_to_link", "files_linked", "is_symlink")], +) +def test_MPAS_files_copied_and_linked(config, cycle, key, task, test, tmp_path): + atm, sfc = "gfs.t%sz.atmanl.nc", "gfs.t%sz.sfcanl.nc" + atm_cfg_dst, sfc_cfg_dst = [x % "{{ cycle.strftime('%H') }}" for x in [atm, sfc]] + atm_cfg_src, sfc_cfg_src = [str(tmp_path / (x + ".in")) for x in [atm_cfg_dst, sfc_cfg_dst]] + config["mpas"].update({key: {atm_cfg_dst: atm_cfg_src, sfc_cfg_dst: sfc_cfg_src}}) + path = tmp_path / "config.yaml" + with open(path, "w", encoding="utf-8") as f: + yaml.dump(config, f) + driverobj = mpas.MPAS(config=path, cycle=cycle, batch=True) + atm_dst, sfc_dst = [tmp_path / (x % cycle.strftime("%H")) for x in [atm, sfc]] + assert not any(dst.is_file() for dst in [atm_dst, sfc_dst]) + atm_src, sfc_src = [Path(str(x) + ".in") for x in [atm_dst, sfc_dst]] + for src in (atm_src, sfc_src): + src.touch() + getattr(driverobj, task)() + assert all(getattr(dst, test)() for dst in [atm_dst, sfc_dst]) + + +def test_MPAS_namelist_file(driverobj): + dst = driverobj._rundir / "namelist.atmosphere" + assert not dst.is_file() + driverobj.namelist_file() + assert dst.is_file() + assert isinstance(f90nml.read(dst), f90nml.Namelist) + + +def test_MPAS_namelist_missing(driverobj): + path = driverobj._rundir / "namelist.atmosphere" + del driverobj._driver_config["namelist"] + with raises(UWConfigError) as e: + assert driverobj.namelist_file() + assert str(e.value) == ("Provide either a 'namelist' YAML block or the %s file" % path) + + +def test_MPAS_provisioned_run_directory(driverobj): + with patch.multiple( + driverobj, + boundary_files=D, + files_copied=D, + files_linked=D, + namelist_file=D, + runscript=D, + streams_file=D, + ) as mocks: + driverobj.provisioned_run_directory() + for m in mocks: + mocks[m].assert_called_once_with() + + +def test_MPAS_run_batch(driverobj): + with patch.object(driverobj, "_run_via_batch_submission") as func: + driverobj.run() + func.assert_called_once_with() + + +def test_MPAS_run_local(driverobj): + driverobj._batch = False + with patch.object(driverobj, "_run_via_local_execution") as func: + driverobj.run() + func.assert_called_once_with() + + +def test_MPAS_runscript(driverobj): + with patch.object(driverobj, "_runscript") as runscript: + driverobj.runscript() + runscript.assert_called_once() + args = ("envcmds", "envvars", "execution", "scheduler") + types = [list, dict, list, Slurm] + assert [type(runscript.call_args.kwargs[x]) for x in args] == types + + +def test_MPAS_streams(driverobj): + src = driverobj._driver_config["streams"]["path"] + with open(src, "w", encoding="utf-8") as f: + f.write("Hello, {{ world }}") + assert not (driverobj._rundir / "streams.atmosphere").is_file() + driverobj.streams_file() + assert (driverobj._rundir / "streams.atmosphere").is_file() + + +def test_MPAS__runscript_path(driverobj): + assert driverobj._runscript_path == driverobj._rundir / "runscript.mpas" + + +def test_MPAS__taskanme(driverobj): + assert driverobj._taskname("foo") == "20240322 06Z mpas foo" + + +def test_MPAS__validate(driverobj): + driverobj._validate() diff --git a/src/uwtools/tests/drivers/test_mpas_init.py b/src/uwtools/tests/drivers/test_mpas_init.py new file mode 100644 index 000000000..9734f1822 --- /dev/null +++ b/src/uwtools/tests/drivers/test_mpas_init.py @@ -0,0 +1,218 @@ +# pylint: disable=missing-function-docstring,protected-access,redefined-outer-name +""" +MPASInit driver tests. +""" +import datetime as dt +from pathlib import Path +from unittest.mock import DEFAULT as D +from unittest.mock import patch + +import f90nml # type: ignore +import pytest +import yaml +from pytest import fixture, raises + +from uwtools.drivers import mpas_init +from uwtools.exceptions import UWConfigError +from uwtools.scheduler import Slurm +from uwtools.tests.support import fixture_path + +# Fixtures + + +@fixture +def cycle(): + return dt.datetime(2024, 2, 1, 18) + + +# Driver fixtures + + +@fixture +def config(tmp_path): + return { + "mpas_init": { + "execution": { + "executable": "mpas_init", + "batchargs": { + "walltime": "01:30:00", + }, + }, + "boundary_conditions": { + "interval_hours": 1, + "length": 1, + "offset": 0, + "path": str(tmp_path / "input_path"), + }, + "namelist": { + "base_file": str(fixture_path("simple.nml")), + "update_values": { + "nhyd_model": {"config_start_time": "12", "config_stop_time": "12"}, + }, + }, + "run_dir": str(tmp_path), + "streams": { + "path": str(tmp_path / "streams.init_atmosphere.in"), + "values": { + "world": "user", + }, + }, + "files_to_link": { + "CAM_ABS_DATA.DBL": "src/MPAS-Model/CAM_ABS_DATA.DBL", + "CAM_AEROPT_DATA.DBL": "src/MPAS-Model/CAM_AEROPT_DATA.DBL", + "GENPARM.TBL": "src/MPAS-Model/GENPARM.TBL", + "LANDUSE.TBL": "src/MPAS-Model/LANDUSE.TBL", + "OZONE_DAT.TBL": "src/MPAS-Model/OZONE_DAT.TBL", + "OZONE_LAT.TBL": "src/MPAS-Model/OZONE_LAT.TBL", + "OZONE_PLEV.TBL": "src/MPAS-Model/OZONE_PLEV.TBL", + "RRTMG_LW_DATA": "src/MPAS-Model/RRTMG_LW_DATA", + "RRTMG_LW_DATA.DBL": "src/MPAS-Model/RRTMG_LW_DATA.DBL", + "RRTMG_SW_DATA": "src/MPAS-Model/RRTMG_SW_DATA", + "RRTMG_SW_DATA.DBL": "src/MPAS-Model/RRTMG_SW_DATA.DBL", + "SOILPARM.TBL": "src/MPAS-Model/SOILPARM.TBL", + "VEGPARM.TBL": "src/MPAS-Model/VEGPARM.TBL", + }, + }, + "platform": { + "account": "me", + "scheduler": "slurm", + }, + } + + +@fixture +def config_file(config, tmp_path): + path = tmp_path / "config.yaml" + with open(path, "w", encoding="utf-8") as f: + yaml.dump(config, f) + return path + + +@fixture +def driverobj(config_file, cycle): + return mpas_init.MPASInit(config=config_file, cycle=cycle, batch=True) + + +# Driver tests + + +def test_MPASInit(driverobj): + assert isinstance(driverobj, mpas_init.MPASInit) + + +def test_MPASInit_boundary_files(cycle, driverobj): + ns = (0, 1) + links = [ + driverobj._rundir / f"FILE:{(cycle+dt.timedelta(hours=n)).strftime('%Y-%m-%d_%H')}" + for n in ns + ] + assert not any(link.is_file() for link in links) + input_path = Path(driverobj._driver_config["boundary_conditions"]["path"]) + input_path.mkdir() + for n in ns: + (input_path / f"FILE:{(cycle+dt.timedelta(hours=n)).strftime('%Y-%m-%d_%H')}").touch() + driverobj.boundary_files() + assert all(link.is_symlink() for link in links) + + +def test_MPASInit_dry_run(config_file, cycle): + with patch.object(mpas_init, "dryrun") as dryrun: + driverobj = mpas_init.MPASInit(config=config_file, cycle=cycle, batch=True, dry_run=True) + assert driverobj._dry_run is True + dryrun.assert_called_once_with() + + +@pytest.mark.parametrize( + "key,task,test", + [("files_to_copy", "files_copied", "is_file"), ("files_to_link", "files_linked", "is_symlink")], +) +def test_MPASInit_files_copied_and_linked(config, cycle, key, task, test, tmp_path): + atm, sfc = "gfs.t%sz.atmanl.nc", "gfs.t%sz.sfcanl.nc" + atm_cfg_dst, sfc_cfg_dst = [x % "{{ cycle.strftime('%H') }}" for x in [atm, sfc]] + atm_cfg_src, sfc_cfg_src = [str(tmp_path / (x + ".in")) for x in [atm_cfg_dst, sfc_cfg_dst]] + config["mpas_init"].update({key: {atm_cfg_dst: atm_cfg_src, sfc_cfg_dst: sfc_cfg_src}}) + path = tmp_path / "config.yaml" + with open(path, "w", encoding="utf-8") as f: + yaml.dump(config, f) + driverobj = mpas_init.MPASInit(config=path, cycle=cycle, batch=True) + atm_dst, sfc_dst = [tmp_path / (x % cycle.strftime("%H")) for x in [atm, sfc]] + assert not any(dst.is_file() for dst in [atm_dst, sfc_dst]) + atm_src, sfc_src = [Path(str(x) + ".in") for x in [atm_dst, sfc_dst]] + for src in (atm_src, sfc_src): + src.touch() + getattr(driverobj, task)() + assert all(getattr(dst, test)() for dst in [atm_dst, sfc_dst]) + + +def test_MPASInit_namelist_file(driverobj): + dst = driverobj._rundir / "namelist.init_atmosphere" + assert not dst.is_file() + driverobj.namelist_file() + assert dst.is_file() + assert isinstance(f90nml.read(dst), f90nml.Namelist) + + +def test_MPASInit_namelist_missing(driverobj): + path = driverobj._rundir / "namelist.init_atmosphere" + del driverobj._driver_config["namelist"] + with raises(UWConfigError) as e: + assert driverobj.namelist_file() + assert str(e.value) == ("Provide either a 'namelist' YAML block or the %s file" % path) + + +def test_MPASInit_provisioned_run_directory(driverobj): + with patch.multiple( + driverobj, + boundary_files=D, + files_copied=D, + files_linked=D, + namelist_file=D, + runscript=D, + streams_file=D, + ) as mocks: + driverobj.provisioned_run_directory() + for m in mocks: + mocks[m].assert_called_once_with() + + +def test_MPASInit_run_batch(driverobj): + with patch.object(driverobj, "_run_via_batch_submission") as func: + driverobj.run() + func.assert_called_once_with() + + +def test_MPASInit_run_local(driverobj): + driverobj._batch = False + with patch.object(driverobj, "_run_via_local_execution") as func: + driverobj.run() + func.assert_called_once_with() + + +def test_MPASInit_runscript(driverobj): + with patch.object(driverobj, "_runscript") as runscript: + driverobj.runscript() + runscript.assert_called_once() + args = ("envcmds", "envvars", "execution", "scheduler") + types = [list, dict, list, Slurm] + assert [type(runscript.call_args.kwargs[x]) for x in args] == types + + +def test_MPASInit_streams_file(driverobj): + src = driverobj._driver_config["streams"]["path"] + with open(src, "w", encoding="utf-8") as f: + f.write("Hello, {{ world }}") + assert not (driverobj._rundir / "streams.init_atmosphere").is_file() + driverobj.streams_file() + assert (driverobj._rundir / "streams.init_atmosphere").is_file() + + +def test_MPASInit__runscript_path(driverobj): + assert driverobj._runscript_path == driverobj._rundir / "runscript.mpas_init" + + +def test_MPASInit__taskanme(driverobj): + assert driverobj._taskname("foo") == "20240201 18Z mpas_init foo" + + +def test_MPASInit__validate(driverobj): + driverobj._validate() diff --git a/src/uwtools/tests/test_cli.py b/src/uwtools/tests/test_cli.py index 381d1d313..5472fb2d5 100644 --- a/src/uwtools/tests/test_cli.py +++ b/src/uwtools/tests/test_cli.py @@ -14,11 +14,15 @@ import uwtools.api.config import uwtools.api.fv3 +import uwtools.api.mpas +import uwtools.api.mpas_init import uwtools.api.rocoto import uwtools.api.sfc_climo_gen import uwtools.api.template import uwtools.api.ungrib import uwtools.drivers.fv3 +import uwtools.drivers.mpas +import uwtools.drivers.mpas_init import uwtools.drivers.sfc_climo_gen import uwtools.drivers.ungrib from uwtools import cli @@ -130,6 +134,34 @@ def test__add_subparser_fv3(subparsers): ] +def test__add_subparser_mpas_init(subparsers): + cli._add_subparser_mpas_init(subparsers) + assert actions(subparsers.choices[STR.mpasinit]) == [ + "boundary_files", + "files_copied", + "files_linked", + "namelist_file", + "provisioned_run_directory", + "run", + "runscript", + "streams_file", + ] + + +def test__add_subparser_mpas(subparsers): + cli._add_subparser_mpas(subparsers) + assert actions(subparsers.choices[STR.mpas]) == [ + "boundary_files", + "files_copied", + "files_linked", + "namelist_file", + "provisioned_run_directory", + "run", + "runscript", + "streams_file", + ] + + def test__add_subparser_rocoto(subparsers): cli._add_subparser_rocoto(subparsers) assert subparsers.choices[STR.rocoto] @@ -448,6 +480,38 @@ def test__dispatch_fv3(): ) +def test__dispatch_mpas(): + cycle = dt.datetime.now() + args: dict = { + "batch": True, + "config_file": "config.yaml", + "cycle": cycle, + "dry_run": False, + "graph_file": None, + } + with patch.object(uwtools.api.mpas, "execute") as execute: + cli._dispatch_mpas({**args, "action": "foo"}) + execute.assert_called_once_with( + batch=True, config="config.yaml", cycle=cycle, dry_run=False, graph_file=None, task="foo" + ) + + +def test__dispatch_mpas_init(): + cycle = dt.datetime.now() + args: dict = { + "batch": True, + "config_file": "config.yaml", + "cycle": cycle, + "dry_run": False, + "graph_file": None, + } + with patch.object(uwtools.api.mpas_init, "execute") as execute: + cli._dispatch_mpas_init({**args, "action": "foo"}) + execute.assert_called_once_with( + batch=True, config="config.yaml", cycle=cycle, dry_run=False, graph_file=None, task="foo" + ) + + @pytest.mark.parametrize( "params", [