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

wrfrun/cli.py

@@ -0,0 +1,128 @@
+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
+}
+
+
+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(WRFRUNConfig.parse_resource_uri(CONFIG_MAIN_TOML_TEMPLATE), f"{project_name}/config.toml")
+
+    if models is not None:
+        for _model in models:
+            src_path = WRFRUNConfig.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(WRFRUNConfig.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

@@ -20,7 +20,7 @@ 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:
@@ -154,7 +154,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 +176,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
@@ -300,7 +303,7 @@ 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)
 
@@ -348,7 +351,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)
 

wrfrun/core/config.py

@@ -22,15 +22,16 @@ Users can use the global variable ``WRFRUNConfig``, which is the instance of thi
 
 from copy import deepcopy
 from os import environ, makedirs
-from os.path import abspath, basename, dirname, exists
+from os.path import abspath, dirname, exists
 from shutil import copyfile
+from sys import platform
 from typing import Optional, Tuple, Union
 
 import f90nml
 import tomli
 import tomli_w
 
-from .error import ResourceURIError, WRFRunContextError, ModelNameError, NamelistIDError, NamelistError
+from .error import ModelNameError, NamelistError, NamelistIDError, ResourceURIError, WRFRunContextError
 from ..utils import logger
 
 
@@ -110,7 +111,7 @@ class _WRFRunResources:
 
         For example, you can get the real path of ``wrfrun`` workspace with this method:
 
-        >>> workspace_path = f"{WRFRUNConfig.WRFRUN_WORKSPACE_PATH}/WPS"    # ":WRFRUN_WORKSPACE_PATH:/WPS"
+        >>> workspace_path = f"{WRFRUNConfig.WRFRUN_WORKSPACE_ROOT}/WPS"    # ":WRFRUN_WORKSPACE_PATH:/WPS"
         >>> real_path = WRFRUNConfig.parse_resource_uri(workspace_path)     # should be a valid path like: "/home/syize/.config/wrfrun/workspace/WPS"
 
         """
@@ -146,41 +147,42 @@ class _WRFRunConstants:
         These variables are related to ``wrfrun`` installation environments, configuration files and more.
         They are defined either directly or mapped using URIs to ensure consistent access across all components.
         """
-        # the path we may need to store temp files,
-        # don't worry, it will be deleted once the system reboots
-        self._WRFRUN_TEMP_PATH = "/tmp/wrfrun"
+        # check system
+        if platform != "linux":
+            logger.debug(f"Not Linux system!")
+
+            # set temporary dir path
+            self._WRFRUN_TEMP_PATH = "./tmp/wrfrun"
+            user_home_path = "./tmp/wrfrun"
+
+        else:
+
+            # the path we may need to store temp files,
+            # don't worry, it will be deleted once the system reboots
+            self._WRFRUN_TEMP_PATH = "/tmp/wrfrun"
+            user_home_path = f"{environ['HOME']}"
 
         # WRF may need a large disk space to store output, we can't run wrf in /tmp,
         # so we will create a folder in $HOME/.config to run wrf.
         # we need to check if we're running as a root user
-        USER_HOME_PATH = f"{environ['HOME']}"
-        if USER_HOME_PATH in ["/", "/root", ""]:
-            logger.warning(f"User's home path is '{USER_HOME_PATH}', which means you are running this program as a root user")
+        if user_home_path in ["/", "/root", ""]:
+            logger.warning(f"User's home path is '{user_home_path}', which means you are running this program as a root user")
             logger.warning("It's not recommended to use wrfrun as a root user")
-            logger.warning("Set USER_HOME_PATH as /root")
-            USER_HOME_PATH = "/root"
-
-        self._WRFRUN_HOME_PATH = f"{USER_HOME_PATH}/.config/wrfrun"
-        self._WRFRUN_REPLAY_WORK_PATH = f"{self._WRFRUN_HOME_PATH}/replay"
+            logger.warning("Set user_home_path as /root")
+            user_home_path = "/root"
 
-        # work path to run WPS, WRF and WRFDA
-        self._WORK_PATH = f"{self._WRFRUN_HOME_PATH}/workspace"
-        self._WPS_WORK_PATH = f"{self._WORK_PATH}/WPS"
-        self._WRF_WORK_PATH = f"{self._WORK_PATH}/WRF"
-        self._WRFDA_WORK_PATH = f"{self._WORK_PATH}/WRFDA"
+        self._WRFRUN_HOME_PATH = f"{user_home_path}/.config/wrfrun"
+        # workspace root path
+        self._WRFRUN_WORKSPACE_ROOT = f"{self._WRFRUN_HOME_PATH}/workspace"
+        self._WRFRUN_WORKSPACE_MODEL = f"{self._WRFRUN_WORKSPACE_ROOT}/model"
+        self._WRFRUN_WORKSPACE_REPLAY = f"{self._WRFRUN_WORKSPACE_ROOT}/replay"
 
         # record WRF progress status
-        self._WORK_STATUS = ""
+        self._WRFRUN_WORK_STATUS = ""
 
         # record context status
         self._WRFRUN_CONTEXT_STATUS = False
 
-        # WRFDA is not necessary
-        self.USE_WRFDA: bool = False
-
-        # output directory of ungrib
-        self._UNGRIB_OUT_DIR = "./outputs"
-
         self._WRFRUN_OUTPUT_PATH = ":WRFRUN_OUTPUT_PATH:"
         self._WRFRUN_RESOURCE_PATH = ":WRFRUN_RESOURCE_PATH:"
 
@@ -203,22 +205,20 @@ class _WRFRunConstants:
         return {
             self.WRFRUN_TEMP_PATH: self._WRFRUN_TEMP_PATH,
             self.WRFRUN_HOME_PATH: self._WRFRUN_HOME_PATH,
-            self.WRFRUN_WORKSPACE_PATH: self._WORK_PATH,
-            self.WPS_WORK_PATH: self._WPS_WORK_PATH,
-            self.WRF_WORK_PATH: self._WRF_WORK_PATH,
-            self.WRFDA_WORK_PATH: self._WRFDA_WORK_PATH,
-            self.WRFRUN_REPLAY_WORK_PATH: self._WRFRUN_REPLAY_WORK_PATH,
+            self.WRFRUN_WORKSPACE_ROOT: self._WRFRUN_WORKSPACE_ROOT,
+            self.WRFRUN_WORKSPACE_MODEL: self._WRFRUN_WORKSPACE_MODEL,
+            self.WRFRUN_WORKSPACE_REPLAY: self._WRFRUN_WORKSPACE_REPLAY,
         }
 
     @property
-    def WRFRUN_REPLAY_WORK_PATH(self) -> str:
+    def WRFRUN_WORKSPACE_REPLAY(self) -> str:
         """
-        Path (URI) to store related files of ``wrfrun``'s replay functionality.
+        Path (URI) to store related files of ``wrfrun`` replay functionality.
 
         :return: URI.
         :rtype: str
         """
-        return ":WRFRUN_REPLAY_WORK_PATH:"
+        return ":WRFRUN_WORKSPACE_REPLAY:"
 
     @property
     def WRFRUN_TEMP_PATH(self) -> str:
@@ -241,44 +241,24 @@ class _WRFRunConstants:
         return ":WRFRUN_HOME_PATH:"
 
     @property
-    def WRFRUN_WORKSPACE_PATH(self) -> str:
-        """
-        Path of the workspace, in which ``wrfrun`` runs NWP models.
-
-        :return: URI
-        :rtype: str
-        """
-        return ":WRFRUN_WORKSPACE_PATH:"
-
-    @property
-    def WPS_WORK_PATH(self) -> str:
+    def WRFRUN_WORKSPACE_ROOT(self) -> str:
         """
-        Workspace in which ``wrfrun`` runs WPS.
+        Path of the root workspace.
 
         :return: URI
         :rtype: str
         """
-        return ":WRFRUN_WPS_WORK_PATH:"
+        return ":WRFRUN_WORKSPACE_ROOT:"
 
     @property
-    def WRF_WORK_PATH(self) -> str:
+    def WRFRUN_WORKSPACE_MODEL(self) -> str:
         """
-        Workspace in which ``wrfrun`` runs WRF.
+        Path of the model workspace, in which ``wrfrun`` runs numerical models.
 
         :return: URI
         :rtype: str
         """
-        return ":WRFRUN_WRF_WORK_PATH:"
-
-    @property
-    def WRFDA_WORK_PATH(self) -> str:
-        """
-        Workspace in which ``wrfrun`` runs WRFDA.
-
-        :return: URI
-        :rtype: str
-        """
-        return ":WRFRUN_WRFDA_WORK_PATH:"
+        return ":WRFRUN_WORKSPACE_MODEL:"
 
     @property
     def WRFRUN_WORK_STATUS(self) -> str:
@@ -291,7 +271,7 @@ class _WRFRunConstants:
         :return: A string reflect the current work progress.
         :rtype: str
         """
-        return self._WORK_STATUS
+        return self._WRFRUN_WORK_STATUS
 
     @WRFRUN_WORK_STATUS.setter
     def WRFRUN_WORK_STATUS(self, value: str):
@@ -306,29 +286,7 @@ class _WRFRunConstants:
         """
         if not isinstance(value, str):
             value = str(value)
-        self._WORK_STATUS = value
-
-    @property
-    def UNGRIB_OUT_DIR(self) -> str:
-        """
-        Output directory path of ``ungrib.exe``.
-
-        :return: URI
-        :rtype: str
-        """
-        return self._UNGRIB_OUT_DIR
-
-    @UNGRIB_OUT_DIR.setter
-    def UNGRIB_OUT_DIR(self, value: str):
-        """
-        Set the output directory path of ``ungrib.exe``.
-
-        :param value: A real path or a URI represents the directory path of ``ungrib.exe``'s output.
-        :type value: str
-        """
-        if not isinstance(value, str):
-            value = str(value)
-        self._UNGRIB_OUT_DIR = value
+        self._WRFRUN_WORK_STATUS = value
 
     @property
     def WRFRUN_OUTPUT_PATH(self) -> str:
@@ -394,9 +352,6 @@ class _WRFRunNamelist:
         If you want to use a new ``namelist_id`` other than the defaults to store namelist,
         you can register a new ``namelist_id`` with :meth:`_WRFRunNamelist.register_custom_namelist_id`.
         """
-        self._wps_namelist = {}
-        self._wrf_namelist = {}
-        self._wrfda_namelist = {}
         self._namelist_dict = {}
         self._namelist_id_list = ("param", "geog_static_data", "wps", "wrf", "wrfda")
 
@@ -570,6 +525,21 @@ class _WRFRunNamelist:
 
         self._namelist_dict.pop(namelist_id)
 
+    def check_namelist(self, namelist_id: str) -> bool:
+        """
+        Check if a namelist has been registered and loaded.
+
+        :param namelist_id: Registered ``namelist_id``.
+        :type namelist_id: str
+        :return: ``True`` if it is registered and loaded, else ``False``.
+        :rtype: bool
+        """
+        if namelist_id in self._namelist_id_list and self._namelist_dict:
+            return True
+
+        else:
+            return False
+
 
 class WRFRunConfig(_WRFRunConstants, _WRFRunNamelist, _WRFRunResources):
     """
@@ -654,6 +624,30 @@ class WRFRunConfig(_WRFRunConstants, _WRFRunNamelist, _WRFRunResources):
         with open(config_path, "rb") as f:
             self._config = tomli.load(f)
 
+        config_dir_path = abspath(dirname(config_path))
+
+        # merge model config.
+        keys_list = list(self._config["model"].keys())
+        for model_key in keys_list:
+
+            # skip the key that isn't model.
+            if model_key == "debug_level":
+                continue
+
+            if "use" not in self._config["model"][model_key]:
+                continue
+
+            if self._config["model"][model_key]["use"]:
+                include_file = self._config["model"][model_key]["include"]
+                if include_file[0] != "/":
+                    include_file = f"{config_dir_path}/{include_file}"
+
+                with open(include_file, "rb") as f:
+                    self._config["model"][model_key] = tomli.load(f)
+
+            else:
+                self._config["model"].pop(model_key)
+
         # register URI for output directory.
         output_path = abspath(self["output_path"])
         self.register_resource_uri(self.WRFRUN_OUTPUT_PATH, output_path)
@@ -752,69 +746,6 @@ class WRFRunConfig(_WRFRunConstants, _WRFRunNamelist, _WRFRunResources):
         """
         return self["core_num"]
 
-    def get_ungrib_out_dir_path(self) -> str:
-        """
-        Get the output directory of ungrib output (WRF intermediate file).
-
-        :return: URI path.
-        :rtype: str
-        """
-        wif_prefix = self.get_namelist("wps")["ungrib"]["prefix"]
-        wif_path = f"{self.WPS_WORK_PATH}/{dirname(wif_prefix)}"
-
-        return wif_path
-
-    def get_ungrib_out_prefix(self) -> str:
-        """
-        Get the prefix string of ungrib output (WRF intermediate file).
-
-        :return: Prefix string of ungrib output (WRF intermediate file).
-        :rtype: str
-        """
-        wif_prefix = self.get_namelist("wps")["ungrib"]["prefix"]
-        wif_prefix = basename(wif_prefix)
-        return wif_prefix
-
-    def set_ungrib_out_prefix(self, prefix: str):
-        """
-        Set the prefix string of ungrib output (WRF intermediate file).
-
-        :param prefix: Prefix string of ungrib output (WRF intermediate file).
-        :type prefix: str
-        """
-        self.update_namelist(
-            {
-                "ungrib": {"prefix": f"{self.UNGRIB_OUT_DIR}/{prefix}"}
-            }, "wps"
-        )
-
-    def get_metgrid_fg_names(self) -> list[str]:
-        """
-        Get prefix strings from "fg_name" in namelist "metgrid" section.
-
-        :return: Prefix strings list.
-        :rtype: list
-        """
-        fg_names = self.get_namelist("wps")["metgrid"]["fg_name"]
-        fg_names = [basename(x) for x in fg_names]
-        return fg_names
-
-    def set_metgrid_fg_names(self, prefix: Union[str, list[str]]):
-        """
-        Set prefix strings of "fg_name" in namelist "metgrid" section.
-
-        :param prefix: Prefix strings list.
-        :type prefix: str | list
-        """
-        if isinstance(prefix, str):
-            prefix = [prefix, ]
-        fg_names = [f"{self.UNGRIB_OUT_DIR}/{x}" for x in prefix]
-        self.update_namelist(
-            {
-                "metgrid": {"fg_name": fg_names}
-            }, "wps"
-        )
-
     def write_namelist(self, save_path: str, namelist_id: str, overwrite=True):
         """
         Write namelist values of a ``namelist_id`` to a file.

wrfrun/core/replay.py

@@ -124,7 +124,7 @@ def replay_config_generator(replay_config_file: str) -> Generator[tuple[str, Exe
     :rtype: Generator
     """
     logger.info(f"Loading replay resources from: {replay_config_file}")
-    work_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_REPLAY_WORK_PATH)
+    work_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_WORKSPACE_REPLAY)
 
     unpack_archive(replay_config_file, work_path, "zip")