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

wrfrun/log.py

@@ -0,0 +1,117 @@
+"""
+wrfrun.log
+##########
+
+.. autosummary::
+    :toctree: generated/
+
+    get_wrfrun_rich_console
+    logger_add_file_handler
+    set_logger
+    unify_logger_format
+
+``wrfrun`` submodule to manage logs. This submodule only take care of ``wrfrun``'s log.
+
+For functions and classes which parse numerical models' log, please check :doc:`model </api/model>`.
+"""
+
+import logging
+from datetime import datetime
+from os import environ
+from time import time
+from typing import Dict, List, Optional
+
+from rich.console import Console
+from rich.logging import RichHandler
+
+from .utils import check_path
+
+WRFRUN_RICH_CONSOLE = Console()
+
+
+def set_logger(logger_list: List[str], logger_level: Optional[Dict] = None):
+    """
+    This function will replace all handlers of each logger in ``logger_list`` with RichHandler.
+
+    If there are some custom handlers in logger, they will be replaced too.
+
+    :param logger_list: A list contains loggers.
+    :type logger_list: list
+    :param logger_level: You can specify the log level in ``logger_level``, with the name of logger is the key,
+                         and the level of logger is the value.
+                         Default if None, with which all loggers' level will be set to ``logging.WARNING``.
+    :type logger_level: list | None
+    """
+    formatter = logging.Formatter("%(name)s :: %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
+    # use rich handler
+    handler = RichHandler(console=WRFRUN_RICH_CONSOLE)
+    handler.setFormatter(formatter)
+
+    for logger_name in logger_list:
+        if logger_name in logging.root.manager.loggerDict:
+            _logger = logging.getLogger(logger_name)
+            for _handler in _logger.handlers:
+                _logger.removeHandler(_handler)
+            _logger.addHandler(handler)
+
+            if logger_level is not None and logger_name in logger_level:
+                _logger.setLevel(logger_level[logger_name])
+            else:
+                _logger.setLevel(logging.WARNING)
+
+
+def unify_logger_format():
+    """
+    This function is only supposed to be used internally.
+
+    This function will replace all handlers of each logger with ``rich.logging.RichHandler``.
+    Use this carefully.
+    """
+    set_logger(
+        ["cdsapi", "cfgrib", "datapi"],
+        {"cdsapi": logging.INFO, "cfgrib": logging.ERROR, "datapi": logging.INFO},
+    )
+
+
+# init wrfrun logger
+logger = logging.getLogger("wrfrun")
+# check environment variables and set logger level
+if "WRFRUN_DEBUG_MODE" in environ and environ["WRFRUN_DEBUG_MODE"]:
+    _logger_level = logging.DEBUG
+    _debug_mode = True
+else:
+    _logger_level = logging.INFO
+    _debug_mode = False
+set_logger(["wrfrun"], {"wrfrun": _logger_level})
+logger.debug(f"DEBUG MODE: {_debug_mode}")
+
+
+def logger_add_file_handler(log_path: str):
+    """
+    Set log file of logger.
+
+    :param log_path: Log file path.
+    :type log_path: str
+    """
+    # check log save path
+    check_path(log_path)
+
+    # add file handler
+    file_handler = logging.FileHandler(f"{log_path}/{datetime.fromtimestamp(time()).strftime('%Y-%m-%d %H:%M:%S')}.log")
+    file_handler.setFormatter(
+        logging.Formatter("%(asctime)s - %(name)s - %(levelname)s :: %(message)s", datefmt="%m-%d %H:%M:%S")
+    )
+    logger.addHandler(file_handler)
+
+
+def get_wrfrun_rich_console() -> Console:
+    """
+    Get ``rich.console.Console`` instance used in wrfrun.
+
+    :return: Console instance.
+    :rtype: Console
+    """
+    return WRFRUN_RICH_CONSOLE
+
+
+__all__ = ["set_logger", "unify_logger_format", "logger_add_file_handler", "logger", "get_wrfrun_rich_console"]

wrfrun/model/__init__.py

@@ -3,12 +3,14 @@ wrfrun.model
 ############
 
 To be able to handle different numerical models,
-``wrfrun`` implements various ``Executable`` for model's binary executable based on :class:`ExecutableBase <wrfrun.core.base.ExecutableBase>`,
+``wrfrun`` implements various ``Executable`` for model's binary executable
+based on :class:`ExecutableBase <wrfrun.core.base.ExecutableBase>`,
 which are all placed in this module.
 
 ``wrfrun.model`` currently supports following numerical models:
 
 ========================================= ==========================================
+:doc:`palm </api/model.palm>`             Support for PALM
 :doc:`wrf </api/model.wrf>`               Support for WRF
 ========================================= ==========================================
 
@@ -16,15 +18,17 @@ which are all placed in this module.
     :maxdepth: 1
     :hidden:
 
-    wrf <model.wrf>
-    base <model.base>
+    constants <model.constants>
+    palm <model.palm>
     plot <model.plot>
+    type <model.type>
     utils <model.utils>
+    wrf <model.wrf>
 """
 
-from .base import *
+# just to register executables
+from . import palm as _
+from . import wrf as _
+from .constants import *
 from .plot import *
 from .utils import *
-
-# just to register executables
-from . import wrf

wrfrun/model/constants.py

@@ -0,0 +1,52 @@
+"""
+wrfrun.model.constants
+######################
+
+Definition of constants used by models.
+
+.. autosummary::
+    :toctree: generated/
+
+    NamelistName
+"""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class NamelistName:
+    """
+    Namelist file names.
+
+    .. py:attribute:: WPS
+        :type: str
+        :value: namelist.wps
+
+        WPS namelist file name.
+
+    .. py:attribute:: WRF
+        :type: str
+        :value: namelist.input
+
+        WRF namelist file name.
+
+    .. py:attribute:: WRFDA
+        :type: str
+        :value: namelist.input
+
+        WRFDA namelist file name.
+
+    .. py:attribute:: PALM
+        :type: str
+        :value: wrfrun_p3d
+
+        PALM namelist file name.
+    """
+
+    WPS = "namelist.wps"
+    WRF = "namelist.input"
+    WRFDA = "namelist.input"
+    PALM = "wrfrun_p3d"
+
+
+__all__ = ["NamelistName"]

wrfrun/model/palm/__init__.py

@@ -0,0 +1,30 @@
+"""
+wrfrun.model.palm
+#################
+
+Implementation of PALM model.
+
+Submodules
+**********
+
+============================================ ==================================================================================
+:doc:`core </api/model.palm.core>`           Core implementation of PALM model.
+:doc:`namelist </api/model.palm.namelist>`   Functions to process PALM namelist files.
+============================================ ==================================================================================
+
+About the implementation of PALM
+********************************
+
+Since ``PALM`` provides a bash script ``palmrun`` to take care of the running of ``palm``,
+``wrfrun`` just simply wraps the bash script.
+
+.. toctree::
+    :maxdepth: 1
+    :hidden:
+
+    core <model.palm.core>
+    namelist <model.palm.namelist>
+"""
+
+from .core import *
+from .namelist import *

wrfrun/model/palm/core.py

@@ -0,0 +1,145 @@
+"""
+wrfrun.model.palm.core
+######################
+
+Core implementation of PALM model. All ``Executable`` and function interface of PALM model are defined here.
+
+.. autosummary::
+    :toctree: generated/
+
+    _check_and_prepare_namelist
+    PALMRun
+    palmrun
+"""
+
+from os.path import exists
+from typing import Optional
+
+from wrfrun.core import WRFRUN, ExecutableBase, ExecutableDB
+from wrfrun.log import logger
+from wrfrun.workspace.palm import get_palm_workspace_path
+
+from ..constants import NamelistName
+from .namelist import prepare_palm_namelist
+
+
+def _check_and_prepare_namelist():
+    """
+    Check if namelist of ``PALM`` has been loaded.
+    If not, call :func:`prepare_palm_namelist <wrfrun.model.palm.namelist.prepare_palm_namelist>` to load it.
+    """
+    if not WRFRUN.config.check_namelist("palm"):
+        prepare_palm_namelist()
+
+
+class PALMRun(ExecutableBase):
+    """
+    ``Executable`` for bash script "palmrun".
+    """
+
+    def __init__(self, config_id: str = "default", core_num: Optional[int] = None):
+        """
+        ``Executable`` for bash script "palmrun".
+
+        :param config_id: Configuration identifier of ``PALM``, defaults to "default"
+        :type config_id: str, optional
+        :param core_num: CPU cores to use, defaults to None
+        :type core_num: Optional[int], optional
+        """
+        if isinstance(core_num, int) and core_num <= 0:
+            logger.warning("`core_num` should be greater than 0")
+            core_num = None
+
+        mpi_use = False
+        mpi_cmd = None
+        mpi_core_num = None
+
+        cmd = f"./palmrun -r wrfrun -c {config_id} -a d3# -X {core_num} -v"
+
+        super().__init__("palmrun", cmd, get_palm_workspace_path(), mpi_use, mpi_cmd, mpi_core_num)
+
+        _check_and_prepare_namelist()
+
+    def generate_custom_config(self):
+        """
+        Store custom configs, including:
+
+        * Namelist settings.
+        """
+        self.custom_config.update({"namelist": WRFRUN.config.get_namelist("palm")})
+
+    def load_custom_config(self):
+        """
+        Load custom configs, including:
+
+        * Namelist settings.
+        """
+        WRFRUN.config.update_namelist(self.custom_config["namelist"], "palm")
+
+    def before_exec(self):
+        WRFRUN.config.check_wrfrun_context(True)
+        WRFRUN.config.WRFRUN_WORK_STATUS = "palm"
+
+        config = WRFRUN.config.get_model_config("palm")
+
+        if not WRFRUN.config.IS_IN_REPLAY:
+            # check if user provides topography files
+            if exists(config["topography_file"]):
+                self.add_input_files(
+                    {
+                        "file_path": config["topography_file"],
+                        "save_path": get_palm_workspace_path("input"),
+                        "save_name": "palmrun.topo",
+                        "is_data": True,
+                        "is_output": False,
+                    }
+                )
+
+        super().before_exec()
+
+        WRFRUN.config.write_namelist(f"{get_palm_workspace_path('input')}/{NamelistName.PALM}", "palm")
+
+        # print debug logs
+        logger.debug("Namelist settings of 'palmrun':")
+        logger.debug(WRFRUN.config.get_namelist("palm"))
+
+    def after_exec(self):
+        if not WRFRUN.config.IS_IN_REPLAY:
+            self.add_output_files(
+                output_dir=get_palm_workspace_path("output"), save_path=self._output_save_path, startswith="wrfrun_"
+            )
+
+        super().after_exec()
+
+        logger.info(f"All PALM output files have been copied to {WRFRUN.config.parse_resource_uri(self._output_save_path)}")
+
+
+def palmrun():
+    """
+    Function interface for :class:`PALMRun`.
+
+    Parameters needed to initialize :class:`PALMRun` is read from global variable :doc:`WRFRUN </api/core.core>`.
+    """
+    config = WRFRUN.config.get_model_config("palm")
+    PALMRun(config["config_identifier"], WRFRUN.config.get_core_num())()
+
+
+def _exec_register_func(exec_db: ExecutableDB):
+    """
+    Function to register ``Executable``.
+
+    :param exec_db: ``ExecutableDB`` instance.
+    :type exec_db: ExecutableDB
+    """
+    class_list = [PALMRun]
+    class_id_list = ["palmrun"]
+
+    for _class, _id in zip(class_list, class_id_list):
+        if not exec_db.is_registered(_id):
+            exec_db.register_exec(_id, _class)
+
+
+WRFRUN.set_exec_db_register_func(_exec_register_func)
+
+
+__all__ = ["PALMRun", "palmrun"]

wrfrun/model/palm/namelist.py

@@ -0,0 +1,33 @@
+"""
+wrfrun.model.palm.namelist
+##########################
+
+Functions to process namelist for ``PALM``.
+
+.. autosummary::
+    :toctree: generated/
+
+    prepare_palm_namelist
+"""
+
+from os.path import exists
+
+from wrfrun.core import WRFRUN
+from wrfrun.log import logger
+
+
+def prepare_palm_namelist():
+    """
+    This function loads user PALM namelist file and save it in :doc:`WRFRUN </api/core.core>`.
+    """
+    palm_config = WRFRUN.config.get_model_config("palm")
+    namelist_file = palm_config["user_namelist"]
+
+    if not exists(namelist_file):
+        logger.error(f"Can't find PALM namelist: {namelist_file}")
+        raise FileNotFoundError(f"Can't find PALM namelist: {namelist_file}")
+
+    WRFRUN.config.read_namelist(namelist_file, "palm")
+
+
+__all__ = ["prepare_palm_namelist"]