UW-532: Ungrib driver (ufs-community#433)

* copy paste all

* remove diag

* fix

* remove fieldtable

* adding fieldtable back...

* edit

* ungrib files

* mpas_init files

* add vtable and gfs

* fixes

* WIP w/ PM

* added namelist dict

* WIP

* rm wpsfiles

* WIP

* remove local

* schema edits

* remove mpas

* gribfile test

* it works

* WIP

* add ungrib.yaml

* fix

* doc strings

* remove files

* doc strings

* Update src/uwtools/drivers/ungrib.py

Co-authored-by: Paul Madden <[email protected]>

* Update docs/shared/ungrib.yaml

Co-authored-by: Paul Madden <[email protected]>

* Update docs/shared/ungrib.yaml

Co-authored-by: Paul Madden <[email protected]>

* Update src/uwtools/drivers/ungrib.py

Co-authored-by: Paul Madden <[email protected]>

* misc

* remove GFS

* ungribbing again

* rm _runscript_path()

* namelist_wps to namelist_file

* fix time

* namelist checking

---------

Co-authored-by: Paul Madden <[email protected]>

docs/shared/ungrib.yaml

@@ -0,0 +1,10 @@
+ungrib:
+  execution:
+    mpicmd: srun
+    batchargs:
+      cores: 1
+      walltime: 00:01:00
+    executable: "/path/to/ungrib.exe"
+  gfs_file: "/path/to/dir/gfs.t12z.pgrb2.0p25.f000"
+  run_dir: "/path/to/dir/run"
+  vtable: "/path/to/Vtable.GFS"

src/uwtools/api/ungrib.py

@@ -0,0 +1,53 @@
+"""
+API access to the ``uwtools`` ``ungrib`` driver.
+"""
+import datetime as dt
+from pathlib import Path
+from typing import Dict, Optional
+
+import uwtools.drivers.support as _support
+from uwtools.drivers.ungrib import Ungrib as _Ungrib
+
+
+def execute(
+    task: str,
+    config_file: Path,
+    cycle: dt.datetime,
+    batch: bool = False,
+    dry_run: bool = False,
+    graph_file: Optional[Path] = None,
+) -> bool:
+    """
+    Execute an ``ungrib`` 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_file: Path to YAML config file
+    :param cycle: The cycle to run
+    :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 = _Ungrib(config_file=config_file, 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(_Ungrib)

src/uwtools/cli.py

@@ -18,6 +18,7 @@
 import uwtools.api.rocoto
 import uwtools.api.sfc_climo_gen
 import uwtools.api.template
+import uwtools.api.ungrib
 import uwtools.config.jinja2
 import uwtools.rocoto
 from uwtools.exceptions import UWConfigRealizeError, UWError, UWTemplateRenderError
@@ -61,6 +62,7 @@ def main() -> None:
             STR.rocoto: _dispatch_rocoto,
             STR.sfcclimogen: _dispatch_sfc_climo_gen,
             STR.template: _dispatch_template,
+            STR.ungrib: _dispatch_ungrib,
         }
         sys.exit(0 if modes[args[STR.mode]](args) else 1)
     except UWError as e:
@@ -658,6 +660,60 @@ def _dispatch_template_translate(args: Args) -> bool:
     )
 
 
+# Mode ungrib
+
+
+def _add_subparser_ungrib(subparsers: Subparsers) -> ModeChecks:
+    """
+    Subparser for mode: ungrib
+
+    :param subparsers: Parent parser's subparsers, to add this subparser to.
+    """
+    parser = _add_subparser(subparsers, STR.ungrib, "Execute Ungrib tasks")
+    _basic_setup(parser)
+    subparsers = _add_subparsers(parser, STR.action, STR.task.upper())
+    return {
+        task: _add_subparser_ungrib_task(subparsers, task, helpmsg)
+        for task, helpmsg in uwtools.api.ungrib.tasks().items()
+    }
+
+
+def _add_subparser_ungrib_task(subparsers: Subparsers, task: str, helpmsg: str) -> ActionChecks:
+    """
+    Subparser for mode: ungrib <task>
+
+    :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_config_file(group=required, required=True)
+    _add_arg_cycle(required)
+    optional = _basic_setup(parser)
+    _add_arg_batch(optional)
+    _add_arg_dry_run(optional)
+    _add_arg_graph_file(optional)
+    checks = _add_args_verbosity(optional)
+    return checks
+
+
+def _dispatch_ungrib(args: Args) -> bool:
+    """
+    Dispatch logic for ungrib mode.
+
+    :param args: Parsed command-line args.
+    """
+    return uwtools.api.ungrib.execute(
+        task=args[STR.action],
+        config_file=args[STR.cfgfile],
+        cycle=args[STR.cycle],
+        batch=args[STR.batch],
+        dry_run=args[STR.dryrun],
+        graph_file=args[STR.graphfile],
+    )
+
+
 # Arguments
 
 # pylint: disable=missing-function-docstring
@@ -1031,6 +1087,7 @@ def _parse_args(raw_args: List[str]) -> Tuple[Args, Checks]:
         STR.rocoto: _add_subparser_rocoto(subparsers),
         STR.sfcclimogen: _add_subparser_sfc_climo_gen(subparsers),
         STR.template: _add_subparser_template(subparsers),
+        STR.ungrib: _add_subparser_ungrib(subparsers),
     }
     return vars(parser.parse_args(raw_args)), checks
 

src/uwtools/drivers/ungrib.py

@@ -0,0 +1,148 @@
+"""
+A driver for the ungrib component.
+"""
+
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict
+
+from iotaa import asset, dryrun, task, tasks
+
+from uwtools.config.formats.nml import NMLConfig
+from uwtools.drivers.driver import Driver
+from uwtools.strings import STR
+from uwtools.utils.tasks import file
+
+
+class Ungrib(Driver):
+    """
+    A driver for ungrib.
+    """
+
+    def __init__(
+        self, config_file: Path, cycle: datetime, dry_run: bool = False, batch: bool = False
+    ):
+        """
+        The driver.
+
+        :param config_file: Path to config file.
+        :param cycle: The forecast cycle.
+        :param dry_run: Run in dry-run mode?
+        :param batch: Run component via the batch system?
+        """
+        super().__init__(config_file=config_file, dry_run=dry_run, batch=batch)
+        self._config.dereference(context={"cycle": cycle})
+        if self._dry_run:
+            dryrun()
+        self._cycle = cycle
+
+    # Workflow tasks
+
+    @task
+    def gribfile_aaa(self):
+        """
+        The gribfile.
+        """
+        path = self._rundir / "GRIBFILE.AAA"
+        yield self._taskname(str(path))
+        yield asset(path, path.is_symlink)
+        infile = Path(self._driver_config["gfs_file"])
+        yield file(path=infile)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.symlink_to(infile)
+
+    @task
+    def namelist_file(self):
+        """
+        The namelist file.
+        """
+        d = {
+            "update_values": {
+                "share": {
+                    "end_date": self._cycle.strftime("%Y-%m-%d_%H:00:00"),
+                    "interval_seconds": 1,
+                    "max_dom": 1,
+                    "start_date": self._cycle.strftime("%Y-%m-%d_%H:00:00"),
+                    "wrf_core": "ARW",
+                },
+                "ungrib": {
+                    "out_format": "WPS",
+                    "prefix": "FILE",
+                },
+            }
+        }
+        path = self._rundir / "namelist.wps"
+        yield self._taskname(str(path))
+        yield asset(path, path.is_file)
+        yield None
+        path.parent.mkdir(parents=True, exist_ok=True)
+        self._create_user_updated_config(
+            config_class=NMLConfig,
+            config_values=d,
+            path=path,
+        )
+
+    @tasks
+    def provisioned_run_directory(self):
+        """
+        Run directory provisioned with all required content.
+        """
+        yield self._taskname("provisioned run directory")
+        yield [
+            self.gribfile_aaa(),
+            self.namelist_file(),
+            self.runscript(),
+            self.vtable(),
+        ]
+
+    @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 vtable(self):
+        """
+        The Vtable.
+        """
+        path = self._rundir / "Vtable"
+        yield self._taskname(str(path))
+        yield asset(path, path.is_symlink)
+        yield None
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.symlink_to(Path(self._driver_config["vtable"]))
+
+    # Private helper methods
+
+    @property
+    def _driver_name(self) -> str:
+        """
+        Returns the name of this driver.
+        """
+        return STR.ungrib
+
+    @property
+    def _resources(self) -> Dict[str, Any]:
+        """
+        Returns configuration data for the runscript.
+        """
+        return {
+            "account": self._config["platform"]["account"],
+            "rundir": self._rundir,
+            "scheduler": self._config["platform"]["scheduler"],
+            **self._driver_config.get("execution", {}).get("batchargs", {}),
+        }
+
+    def _taskname(self, suffix: str) -> str:
+        """
+        Returns a common tag for graph-task log messages.
+
+        :param suffix: Log-string suffix.
+        """
+        return "%s ungrib %s" % (self._cycle.strftime("%Y%m%d %HZ"), suffix)

src/uwtools/resources/jsonschema/ungrib.jsonschema

@@ -0,0 +1,29 @@
+{
+  "properties": {
+    "ungrib": {
+      "additionalProperties": false,
+      "properties": {
+        "execution": {
+          "$ref": "urn:uwtools:execution"
+        },
+        "gfs_file": {
+          "type": "string"
+        },
+        "run_dir": {
+          "type": "string"
+        },
+        "vtable": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "execution",
+        "gfs_file",
+        "run_dir",
+        "vtable"
+      ],
+      "type": "object"
+    }
+  },
+  "type": "object"
+}

src/uwtools/strings.py

@@ -104,6 +104,7 @@ class STR:
     template: str = "template"
     total: str = "total"
     translate: str = "translate"
+    ungrib: str = "ungrib"
     validate: str = "validate"
     valsfile: str = "values_file"
     valsfmt: str = "values_format"

src/uwtools/tests/api/test_ungrib.py

@@ -0,0 +1,34 @@
+# pylint: disable=missing-function-docstring,protected-access
+
+import datetime as dt
+from unittest.mock import patch
+
+from uwtools.api import ungrib
+
+
+def test_execute(tmp_path):
+    dot = tmp_path / "graph.dot"
+    args: dict = {
+        "config_file": "config.yaml",
+        "cycle": dt.datetime.utcnow(),
+        "batch": False,
+        "dry_run": True,
+        "graph_file": dot,
+    }
+    with patch.object(ungrib, "_Ungrib") as Ungrib:
+        assert ungrib.execute(**args, task="foo") is True
+    del args["graph_file"]
+    Ungrib.assert_called_once_with(**args)
+    Ungrib().foo.assert_called_once_with()
+
+
+def test_graph():
+    with patch.object(ungrib._support, "graph") as graph:
+        ungrib.graph()
+    graph.assert_called_once_with()
+
+
+def test_tasks():
+    with patch.object(ungrib._support, "tasks") as _tasks:
+        ungrib.tasks()
+    _tasks.assert_called_once_with(ungrib._Ungrib)