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

wrfrun/__init__.py

@@ -1,3 +1,8 @@
-from .core import *
-from .run import *
-from .res import *
+"""
+wrfrun
+######
+"""
+
+from . import core, data, extension, model, plot, res, run, scheduler, utils, workspace
+
+__all__ = ["core", "data", "extension", "model", "plot", "res", "run", "scheduler", "utils", "workspace"]

wrfrun/cli.py

@@ -1,6 +1,23 @@
+"""
+wrfrun.cli
+##########
+
+.. autosummary::
+    :toctree: generated/
+
+    main_entry
+
+``wrfrun`` command line interface (CLI).
+
+**This is an experimental tool, NOT BEING STABLE YET.**
+
+After installing ``wrfrun``, you can use command ``wrfrun`` to create a simulation project.
+Get more information by running ``wrfrun --help``.
+"""
+
 import argparse
 import sys
-from os import makedirs
+from os import listdir, makedirs
 from os.path import abspath, dirname, exists
 from shutil import copyfile
 
@@ -8,15 +25,17 @@ import tomli
 import tomli_w
 
 from .core import WRFRunConfig
-from .res import CONFIG_MAIN_TOML_TEMPLATE, CONFIG_WRF_TOML_TEMPLATE
-from .utils import logger
+from .log import logger
+from .res import CONFIG_MAIN_TOML_TEMPLATE, CONFIG_PALM_TOML_TEMPLATE, CONFIG_WRF_TOML_TEMPLATE, _register_res_uri
 
 MODEL_MAP = {
-    "wrf": CONFIG_WRF_TOML_TEMPLATE
+    "wrf": CONFIG_WRF_TOML_TEMPLATE,
+    "palm": CONFIG_PALM_TOML_TEMPLATE,
 }
 
-
+# need some mannual calls to make cli work without a config file.
 wrfrun_config = WRFRunConfig("./.wrfrun")
+_register_res_uri(wrfrun_config)
 
 
 def _entry_init(args: argparse.Namespace):
@@ -26,17 +45,29 @@ def _entry_init(args: argparse.Namespace):
     :param args: Arguments namespace.
     :type args: argparse.Namespace
     """
-    args = vars(args)
+    configs = vars(args)
+
+    project_name = configs["name"]
+    models = configs["models"]
 
-    project_name = args["name"]
-    models = args["models"]
+    if project_name is None:
+        project_name = "."
 
     if exists(project_name):
-        logger.error(f"{project_name} already exists.")
-        exit(1)
+        # we need to check if this directory isn't empty.
+        files = listdir()
+        # exclude:
+        exclude_target = [".venv", ".git", ".gitignore", ".gitattribute", "pyproject.toml", "uv.lock"]
+        files = [x for x in files if x not in exclude_target]
+
+        if len(files) != 0:
+            logger.error(f"{project_name} isn't empty, choose an empty directory, or backup and delete your files first.")
+            exit(1)
 
     makedirs(f"{project_name}/configs")
     makedirs(f"{project_name}/data")
+    namelist_path = f"{project_name}/namelists"
+    makedirs(namelist_path)
 
     copyfile(wrfrun_config.parse_resource_uri(CONFIG_MAIN_TOML_TEMPLATE), f"{project_name}/config.toml")
 
@@ -44,9 +75,15 @@ def _entry_init(args: argparse.Namespace):
         for _model in models:
             src_path = wrfrun_config.parse_resource_uri(MODEL_MAP[_model])
             copyfile(src_path, f"{project_name}/configs/{_model}.toml")
+            makedirs(f"{namelist_path}/{_model}")
 
-    logger.info(f"Created project {project_name}.")
-    logger.info(f"Use command `wrfrun add MODEL_NAME` to add a new model to project.")
+    logger.info(f"Created project {project_name}")
+    logger.info(f"All your configs should be placed in {project_name}/configs.")
+    logger.info(f"All your data should be placed in {project_name}/data.")
+    logger.info(f"All your namelist files should be placed in {namelist_path}")
+    logger.info(f"Make sure to set `use = True` to enable models in {project_name}/config.toml .")
+    logger.info("It is recommanded to use git to track your wrfrun and model configs.")
+    logger.info("Use command `wrfrun add MODEL_NAME` to add a new model to project.")
 
 
 def _entry_model(args: argparse.Namespace):
@@ -56,20 +93,16 @@ def _entry_model(args: argparse.Namespace):
     :param args: Arguments namespace.
     :type args: argparse.Namespace
     """
-    args = vars(args)
-    new_models = args["add"]
-    config_path = args["config"]
+    configs = vars(args)
+    new_models = configs["add"]
+    config_path = configs["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:
+        if _new_model not in MODEL_MAP:
             logger.error(f"Unknow model type: '{_new_model}'")
             exit(1)
 
@@ -84,21 +117,27 @@ def _entry_model(args: argparse.Namespace):
     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.",
+                "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"
+                "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.",
+                    "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"
+                    "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")
+        copyfile(wrfrun_config.parse_resource_uri(MODEL_MAP[_new_model]), f"{config_dir_path}/{_new_model}.toml")
 
     with open(config_path, "wb") as f:
         tomli_w.dump(main_config, f)
@@ -110,18 +149,19 @@ 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.add_argument("-n", "--name", type=str, help="Name of the wrfrun project.")
+    init_parser.add_argument("--models", nargs="*", type=str, help="List of models to use.", choices=["wrf", "palm"])
     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.add_argument(
+        "-a", "--add", nargs="+", required=True, type=str, help="Add models to the project.", choices=["wrf", "palm"]
+    )
     model_parser.set_defaults(func=_entry_model)
 
     args = args_parser.parse_args(args=None if sys.argv[1:] else ["--help"])

wrfrun/core/__init__.py

@@ -12,27 +12,44 @@ The functionalities are split into several submodules, which are listed in the t
 Submodules
 **********
 
-================================ ========================================================
-:doc:`base </api/core.base>`     Executable base class and related classes.
-:doc:`config </api/core.config>` ``wrfrun`` config classes.
-:doc:`error </api/core.error>`   ``wrfrun`` error classes.
-:doc:`replay </api/core.replay>` Functions and classes to record and replay simulations.
-:doc:`server </api/core.server>` Functions and classes to start socket server.
-================================ ========================================================
+====================================== ========================================================
+:doc:`_config </api/core._config>`     Definition of ``WRFRunConfig``.
+:doc:`_constant </api/core._constant>` Definition of ``ConstantMixIn``.
+:doc:`_exec_db </api/core._exec_db>`   Definition of ``ExecutableDB``.
+:doc:`_namelist </api/core._namelist>` Definition of ``NamelistMixIn``.
+:doc:`_record </api/core._record>`     Definition of ``ExecutableRecorder``.
+:doc:`_resource </api/core._resource>` Definition of ``ResourceMixIn``.
+:doc:`base </api/core.base>`           Definition of ``Executable`` base class.
+:doc:`core </api/core.core>`           Definition of proxy class ``WRFRUNProxy``.
+:doc:`error </api/core.error>`         Definition of ``wrfrun`` error exceptions.
+:doc:`replay </api/core.replay>`       Functions and classes to replay simulations.
+:doc:`server </api/core.server>`       Functions and classes to start socket server.
+:doc:`type </api/core.type>`           Definition of various types used in ``wrfrun``.
+====================================== ========================================================
 
-.. toctree:: 
+.. toctree::
     :maxdepth: 1
     :hidden:
 
+    _config <core._config>
+    _constant <core._constant>
+    _exec_db <core._exec_db>
+    _namelist <core._namelist>
+    _record <core._record>
+    _resource <core._resource>
     base <core.base>
-    config <core.config>
+    core <core.core>
     error <core.error>
     replay <core.replay>
     server <core.server>
+    type <core.type>
 """
 
+from ._config import *
+from ._exec_db import *
 from .base import *
-from .config import *
+from .core import *
 from .error import *
 from .replay import *
 from .server import *
+from .type import *

wrfrun/core/_config.py

@@ -0,0 +1,308 @@
+"""
+wrfrun.core._config
+###################
+
+.. autosummary::
+    :toctree: generated/
+
+    WRFRunConfig
+
+WRFRunConfig
+************
+
+This class inherits :class:`ConstantMixIn <wrfrun.core._constant.ConstantMixIn>`,
+:class:`NamelistMixIn <wrfrun.core._namelist.NamelistMixIn>`,
+and :class:`ResourceMixIn <wrfrun.core._resource.ResourceMixIn>`.
+
+Besides the methods from its parents, :class:`WRFRunConfig` provides methods to read and access user config files.
+"""
+
+from copy import deepcopy
+from os import makedirs
+from os.path import abspath, dirname, exists
+from shutil import copyfile
+from typing import Callable, Tuple
+
+import tomli
+import tomli_w
+
+from ..log import logger
+from ._constant import ConstantMixIn
+from ._namelist import NamelistMixIn
+from ._resource import ResourceMixIn
+from .error import ModelNameError
+
+
+class WRFRunConfig(ConstantMixIn, NamelistMixIn, ResourceMixIn):
+    """
+    Comprehensive class to manage wrfrun config, runtime constants, namelists and resource files.
+    """
+
+    def __init__(self, work_dir: str):
+        """
+
+        :param work_dir: ``wrfrun`` work directory path.
+        :type work_dir: str
+        """
+        super().__init__(work_dir=work_dir)
+
+        self._config = {}
+
+        self._config_template_file_path = None
+
+        self._register_wrfrun_uris()
+
+    def apply_register_func(self, func_list: list[Callable[["WRFRunConfig"], None]]):
+        """
+        Call register functions provided by other submodules.
+
+        :param func_list: A list contains register functions.
+        :type func_list: list[Callable[["WRFRunConfig"], None]]
+        """
+        while len(func_list) != 0:
+            _func = func_list.pop(0)
+            _func(self)
+
+    @classmethod
+    def from_config_file(
+        cls, config_file: str, register_funcs: list[Callable[["WRFRunConfig"], None]]
+    ) -> "WRFRunConfig":
+        """
+        Read the config file and create a instance.
+
+        :param config_file: Config file path.
+        :type config_file: str
+        :param register_funcs: Register function list.
+        :type register_funcs: list[Callable[["WRFRunConfig"], None]]
+        :return: New instance
+        :rtype: WRFRunConfig
+        """
+        with open(config_file, "rb") as f:
+            config = tomli.load(f)
+
+        instance = cls(work_dir=config["work_dir"])
+        instance.apply_register_func(register_funcs)
+        instance.load_wrfrun_config(config_file)
+
+        return instance
+
+    def _register_wrfrun_uris(self):
+        for key, value in self._get_uri_map().items():
+            self.register_resource_uri(key, value)
+
+    def set_config_template_path(self, file_path: str):
+        """
+        Set file path of the config template file.
+
+        ``WRFRunConfig`` needs to know
+
+        :param file_path: Template file path.
+        :type file_path: str
+        """
+        self._config_template_file_path = file_path
+
+    def load_wrfrun_config(self, config_path: str):
+        """
+        Load configs from a config file.
+
+        If the config path is invalid, ``WRFRunConfig`` will create a new config file at the same place,
+        and raise :class:`FileNotFoundError`.
+
+        :param config_path: TOML config file.
+        :type config_path: str
+        """
+        config_template_path = self.parse_resource_uri(self._config_template_file_path)
+
+        if not exists(config_path):
+            logger.error(f"Config file doesn't exist, copy template config to {config_path}")
+            logger.error("Please modify it.")
+
+            if not exists(dirname(config_path)):
+                makedirs(dirname(config_path))
+
+            copyfile(config_template_path, config_path)
+            raise FileNotFoundError(config_path)
+
+        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 "include" not in self._config["model"][model_key]:
+                continue
+
+            # use = True, and have "include" key
+            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:
+                    # keep "use" key, as other components may use this key
+                    _mode_config = tomli.load(f)
+                    _mode_config.update({"use": True})
+                    self._config["model"][model_key] = _mode_config
+
+            else:
+                self._config["model"].pop(model_key)
+
+        # register URI for output directory.
+        output_path = abspath(self["output_path"])
+        self.unregister_resource_uri(self.WRFRUN_OUTPUT_PATH)
+        self.register_resource_uri(self.WRFRUN_OUTPUT_PATH, output_path)
+
+        # some additional check
+        if self._config["input_data_path"] == "":
+            logger.warning("It seems you forget to set 'input_data_path', set it to 'data'.")
+            self._config["input_data_path"] = "data"
+
+    def save_wrfrun_config(self, save_path: str):
+        """
+        Save wrfrun config to a file.
+
+        :param save_path: Save path of the config.
+        :type save_path: str
+        """
+        save_path = self.parse_resource_uri(save_path)
+
+        path_dir = dirname(save_path)
+        if not exists(path_dir):
+            makedirs(path_dir)
+
+        with open(save_path, "wb") as f:
+            tomli_w.dump(self._config, f)
+
+    def __getitem__(self, item: str):
+        """
+        You can access wrfrun config like the way to access values in a dictionary.
+
+        For example:
+
+        >>> from wrfrun.core import WRFRUN
+        >>> model_config = WRFRUN.config["model"]    # get all model configs.
+
+        :param item: Keys.
+        :type item: str
+        """
+        if len(self._config) == 0:
+            logger.error("Attempt to read value before load config")
+            raise RuntimeError("Attempt to read value before load config")
+
+        return deepcopy(self._config[item])
+
+    def __setitem__(self, key: str, value):
+        if key == "model":
+            logger.error("Use `update_model_config` to change model configurations.")
+            raise KeyError("Use `update_model_config` to change model configurations.")
+
+        if key in self._config:
+            self._config[key] = value
+
+        else:
+            logger.error(f"Can't find key '{key}' in your config.")
+            raise KeyError(f"Can't find key '{key}' in your config.")
+
+    def get_input_data_path(self) -> str:
+        """
+        Get the path of directory in which stores the input data.
+
+        :return: Directory path.
+        :rtype: str
+        """
+        return deepcopy(self["input_data_path"])
+
+    def get_model_config(self, model_name: str) -> dict:
+        """
+        Get the config of a NWP model.
+
+        An exception :class:`ModelNameError <wrfrun.core.error.ModelNameError>` will be raised
+        if the config can't be found.
+
+        :param model_name: Name of the model. For example, ``wrf``.
+        :type model_name: str
+        :return: A dictionary.
+        :rtype: dict
+        """
+        if model_name not in self["model"]:
+            logger.error(f"Config of model '{model_name}' isn't found in your config file.")
+            raise ModelNameError(f"Config of model '{model_name}' isn't found in your config file.")
+
+        return self["model"][model_name]
+
+    def update_model_config(self, model_name: str, value: dict):
+        """
+        Update the config of a NWP model.
+
+        An exception :class:`ModelNameError <wrfrun.core.error.ModelNameError>` will be raised
+        if the config can't be found.
+
+        :param model_name: Name of the model. For example, ``wrf``.
+        :type model_name: str
+        :param value: Dictionary contains new values.
+        :type value: dict
+        """
+        if model_name not in self["model"]:
+            logger.error(f"Config of model '{model_name}' isn't found in your config file.")
+            raise ModelNameError(f"Config of model '{model_name}' isn't found in your config file.")
+
+        self._config["model"][model_name].update(value)
+
+    def get_log_path(self) -> str:
+        """
+        Get the directory path to save logs.
+
+        :return: A directory path.
+        :rtype: str
+        """
+        return self["log_path"]
+
+    def get_socket_server_config(self) -> Tuple[str, int]:
+        """
+        Get settings of the socket server.
+
+        :return: ("host", port)
+        :rtype: tuple
+        """
+        return self["server_host"], self["server_port"]
+
+    def get_job_scheduler_config(self) -> dict:
+        """
+        Get configs of job scheduler.
+
+        :return: A dict object.
+        :rtype: dict
+        """
+        return deepcopy(self["job_scheduler"])
+
+    def get_core_num(self) -> int:
+        """
+        Get the number of CPU cores to be used.
+        :return: Core numbers
+        :rtype: int
+        """
+        return self["core_num"]
+
+    def write_namelist(self, save_path: str, namelist_id: str, overwrite=True):
+        """
+        Write namelist values of a ``namelist_id`` to a file.
+        This method is overwritten to convert URIs in ``save_path`` first.
+
+        :param save_path: Target file path.
+        :type save_path: str
+        :param namelist_id: Registered ``namelist_id``.
+        :type namelist_id: str
+        :param overwrite: If overwrite the existed file.
+        :type overwrite: bool
+        """
+        save_path = self.parse_resource_uri(save_path)
+        super().write_namelist(save_path, namelist_id, overwrite)
+
+
+__all__ = ["WRFRunConfig"]