wrfrun 0.1.8__py3-none-any.whl → 0.2.0__py3-none-any.whl

wrfrun/cli.py

@@ -0,0 +1,131 @@
+import argparse
+import sys
+from os import makedirs
+from os.path import abspath, dirname, exists
+from shutil import copyfile
+
+import tomli
+import tomli_w
+
+from .core import WRFRunConfig
+from .res import CONFIG_MAIN_TOML_TEMPLATE, CONFIG_WRF_TOML_TEMPLATE
+from .utils import logger
+
+MODEL_MAP = {
+    "wrf": CONFIG_WRF_TOML_TEMPLATE
+}
+
+
+wrfrun_config = WRFRunConfig("./.wrfrun")
+
+
+def _entry_init(args: argparse.Namespace):
+    """
+    Initialize a wrfrun project.
+
+    :param args: Arguments namespace.
+    :type args: argparse.Namespace
+    """
+    args = vars(args)
+
+    project_name = args["name"]
+    models = args["models"]
+
+    if exists(project_name):
+        logger.error(f"{project_name} already exists.")
+        exit(1)
+
+    makedirs(f"{project_name}/configs")
+    makedirs(f"{project_name}/data")
+
+    copyfile(wrfrun_config.parse_resource_uri(CONFIG_MAIN_TOML_TEMPLATE), f"{project_name}/config.toml")
+
+    if models is not None:
+        for _model in models:
+            src_path = wrfrun_config.parse_resource_uri(MODEL_MAP[_model])
+            copyfile(src_path, f"{project_name}/configs/{_model}.toml")
+
+    logger.info(f"Created project {project_name}.")
+    logger.info(f"Use command `wrfrun add MODEL_NAME` to add a new model to project.")
+
+
+def _entry_model(args: argparse.Namespace):
+    """
+    Manage models used by wrfrun project.
+
+    :param args: Arguments namespace.
+    :type args: argparse.Namespace
+    """
+    args = vars(args)
+    new_models = args["add"]
+    config_path = args["config"]
+
+    if not exists(config_path):
+        logger.error(f"Can't find '{config_path}', initialize this project first.")
+        exit(1)
+
+    model_config_map = {
+        "wrf": CONFIG_WRF_TOML_TEMPLATE
+    }
+
+    for _new_model in new_models:
+        if _new_model not in model_config_map:
+            logger.error(f"Unknow model type: '{_new_model}'")
+            exit(1)
+
+    config_dir_path = f"{abspath(dirname(config_path))}/configs"
+
+    if not exists(config_dir_path):
+        makedirs(config_dir_path)
+
+    with open(config_path, "rb") as f:
+        main_config = tomli.load(f)
+
+    for _new_model in new_models:
+        if _new_model not in main_config["model"]:
+            main_config["model"][_new_model] = {
+                "note": "Config of this model is generated by wrfrun cli command. DO NOT ADD CUSTOM CONFIG IN THIS SECTION because they may be overwrite by wrfrun cli tools.",
+                "use": True,
+                "include": f"./configs/{_new_model}.toml"
+            }
+
+        else:
+            if not ("use" in main_config["model"] and main_config["model"]["use"]):
+                main_config["model"][_new_model] = {
+                    "note": "Config of this model is generated by wrfrun cli command. DO NOT ADD CUSTOM CONFIG IN THIS SECTION because they may be overwrite by wrfrun cli tools.",
+                    "use": True,
+                    "include": f"./configs/{_new_model}.toml"
+                }
+
+    for _new_model in new_models:
+        copyfile(wrfrun_config.parse_resource_uri(model_config_map[_new_model]), f"{config_dir_path}/{_new_model}.toml")
+
+    with open(config_path, "wb") as f:
+        tomli_w.dump(main_config, f)
+
+    logger.info(f"Added models: {new_models}")
+
+
+def main_entry():
+    """
+    CLI entry point.
+    """
+
+    args_parser = argparse.ArgumentParser()
+    subparsers = args_parser.add_subparsers(title="Subcommands", description="Valid Subcommands", help="Subcommands")
+
+    init_parser = subparsers.add_parser("init", help="Initialize a wrfrun project.", add_help=True)
+    init_parser.add_argument("-n", "--name", type=str, required=True, help="Name of the wrfrun project.")
+    init_parser.add_argument("--models", nargs="*", type=str, help="List of models to use.", choices=["wrf"])
+    init_parser.set_defaults(func=_entry_init)
+
+    model_parser = subparsers.add_parser("model", help="Manage models used by wrfrun project.", add_help=True)
+    model_parser.add_argument("-c", "--config", type=str, default="config.toml", help="Path of the main config file.")
+    model_parser.add_argument("-a", "--add", nargs="+", required=True, type=str, help="Add models to the project.")
+    model_parser.set_defaults(func=_entry_model)
+
+    args = args_parser.parse_args(args=None if sys.argv[1:] else ["--help"])
+    args.func(args)
+
+
+__all__ = ["main_entry"]

wrfrun/core/base.py

@@ -13,14 +13,15 @@ Defines what :class:`ExecutableBase <Executable>` is, how it works and how ``wrf
     FileConfigDict
     ExecutableClassConfig
     ExecutableConfig
-    _ExecutableConfigRecord
+    ExecutableConfigRecord
+    create_recorder
     ExecutableBase
 
 Executable
 **********
 
 While ``wrfrun`` aims to provide Python interfaces to various Numerical Weather Prediction model,
-it is important to provide a clear standard about how should a external executable file be implemented in ``wrfrun``.
+it is important to provide a clear standard about how should an external executable file be implemented in ``wrfrun``.
 ``wrfrun`` provides a class called :class:`ExecutableBase`, which is the parent class for all ``Executable`` classes.
 It not only provide the method to execute external programs,
 but also:
@@ -54,7 +55,7 @@ from typing import Optional, TypedDict, Union
 import numpy as np
 
 from .config import WRFRUNConfig
-from .error import CommandError, ConfigError, OutputFileError
+from .error import CommandError, ConfigError, OutputFileError, RecordError
 from ..utils import check_path, logger
 
 
@@ -154,7 +155,9 @@ class InputFileType(Enum):
 
 class FileConfigDict(TypedDict):
     """
-    This dict is used to store information about the file, including its path, the path it will be copied or moved to, its new name, etc. This dict contains following keys:
+    This dict is used to store information about the file, including its path,
+    the path it will be copied or moved to, its new name, etc.
+    This dict contains following keys:
 
     .. py:attribute:: file_path
         :type: str
@@ -174,7 +177,8 @@ class FileConfigDict(TypedDict):
     .. py:attribute:: is_data
         :type: bool
 
-        If the file is data. If not, ``wrfrun`` will treat it as a config file, and always save it to ``.replay`` file when recording the simulation.
+        If the file is data. If not, ``wrfrun`` will treat it as a config file,
+        and always save it to ``.replay`` file when recording the simulation.
 
     .. py:attribute:: is_output
         :type: bool
@@ -273,14 +277,14 @@ class ExecutableConfig(TypedDict):
     custom_config: Optional[dict]
 
 
-class _ExecutableConfigRecord:
+class ExecutableConfigRecord:
     """
     A class to helps store configs of various executables and exports them to a file.
     """
     _instance = None
     _initialized = False
 
-    def __init__(self, save_path: Optional[str] = None, include_data=False):
+    def __init__(self, save_path: str, include_data=False):
         """
 
         :param save_path: Save path of the exported config file.
@@ -300,9 +304,8 @@ class _ExecutableConfigRecord:
         self.save_path = save_path
         self.include_data = include_data
 
-        self.work_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_REPLAY_WORK_PATH)
+        self.work_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_WORKSPACE_REPLAY)
         self.content_path = f"{self.work_path}/config_and_data"
-        check_path(self.content_path)
 
         self._recorded_config = []
         self._name_count = {}
@@ -315,15 +318,20 @@ class _ExecutableConfigRecord:
 
         return cls._instance
 
-    def reinit(self, save_path: Optional[str] = None, include_data=False):
+    @classmethod
+    def reinit(cls, save_path: str, include_data=False):
         """
         Reinitialize this instance.
 
+        :param save_path: Save path of the exported config file.
+        :type save_path: str
+        :param include_data: If includes input data.
+        :type include_data: bool
         :return: New instance.
-        :rtype: _ExecutableConfigRecord
+        :rtype: ExecutableConfigRecord
         """
-        self._initialized = False
-        return _ExecutableConfigRecord(save_path, include_data)
+        cls._initialized = False
+        return cls(save_path, include_data)
 
     def record(self, exported_config: ExecutableConfig):
         """
@@ -348,7 +356,7 @@ class _ExecutableConfigRecord:
             self._name_count[name] = 1
             index = 1
 
-        data_save_uri = f"{WRFRUNConfig.WRFRUN_REPLAY_WORK_PATH}/{name}/{index}"
+        data_save_uri = f"{WRFRUNConfig.WRFRUN_WORKSPACE_REPLAY}/{name}/{index}"
         data_save_path = f"{self.content_path}/{name}/{index}"
         makedirs(data_save_path)
 
@@ -414,7 +422,25 @@ class _ExecutableConfigRecord:
         logger.info(f"Replay config exported to {self.save_path}")
 
 
-ExecConfigRecorder = _ExecutableConfigRecord()
+ExecConfigRecorder: Optional[ExecutableConfigRecord] = None
+
+
+def create_recorder(save_path: str, include_data=False) -> ExecutableConfigRecord:
+    """
+    Create a recorder to record simulations.
+
+    :param save_path: Save path of the exported config file.
+    :type save_path: str
+    :param include_data: If includes input data.
+    :type include_data: bool
+    :return: Recorder instance.
+    :rtype: ExecutableConfigRecord
+    """
+    global ExecConfigRecorder
+
+    ExecConfigRecorder = ExecutableConfigRecord.reinit(save_path, include_data)
+
+    return ExecConfigRecorder
 
 
 class ExecutableBase:
@@ -585,25 +611,27 @@ class ExecutableBase:
 
         You can give more information with a ``FileConfigDict``.
 
+        >>> from wrfrun.workspace.wrf import get_wrf_workspace_path
         >>> file_dict: FileConfigDict = {
         ...     "file_path": "data/custom_file.nc",
-        ...     "save_path": f"{WRFRUNConfig.WPS_WORK_PATH}",
+        ...     "save_path": get_wrf_workspace_path("wps"),
         ...     "save_name": "custom_file.nc",
         ...     "is_data": True,
         ...     "is_output": False
         ... }
         >>> self.add_input_files(file_dict)
 
+        >>> from wrfrun.workspace.wrf import get_wrf_workspace_path
         >>> file_dict_1: FileConfigDict = {
         ...     "file_path": "data/custom_file",
-        ...     "save_path": f"{WRFRUNConfig.WPS_WORK_PATH}/geogrid",
+        ...     "save_path": f"{get_wrf_workspace_path('wps')}/geogrid",
         ...     "save_name": "GEOGRID.TBL",
         ...     "is_data": False,
         ...     "is_output": False
         ... }
         >>> file_dict_2: FileConfigDict = {
         ...     "file_path": "data/custom_file",
-        ...     "save_path": f"{WRFRUNConfig.WPS_WORK_PATH}/outputs",
+        ...     "save_path": f"{get_wrf_workspace_path('wps')}/outputs",
         ...     "save_name": "test_file",
         ...     "is_data": True,
         ...     "is_output": True
@@ -842,7 +870,12 @@ class ExecutableBase:
         self.after_exec()
 
         if not WRFRUNConfig.IS_IN_REPLAY and WRFRUNConfig.IS_RECORDING:
+            global ExecConfigRecorder
+            if ExecConfigRecorder is None:
+                logger.error(f"Trying to record simulation before create the recorder.")
+                raise RecordError(f"Trying to record simulation before create the recorder.")
+
             ExecConfigRecorder.record(self.export_config())
 
 
-__all__ = ["ExecutableBase", "FileConfigDict", "InputFileType", "ExecutableConfig", "ExecutableClassConfig", "ExecConfigRecorder"]
+__all__ = ["ExecutableBase", "FileConfigDict", "InputFileType", "ExecutableConfig", "ExecutableClassConfig", "ExecConfigRecorder", "create_recorder"]