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

wrfrun/utils.py

@@ -1,82 +1,33 @@
-import logging
-import subprocess
-from datetime import datetime
-from os import chdir, getcwd, makedirs, environ
-from os.path import exists
-from shutil import rmtree
-from time import time
-from typing import Optional, List, Dict
-
-from rich.logging import RichHandler
-
-
-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
-    :return:
-    :rtype:
-    """
-    formatter = logging.Formatter(
-        "%(name)s :: %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
-    )
-    # use rich handler
-    handler = RichHandler()
-    handler.setFormatter(formatter)
+"""
+wrfrun.utils
+############
 
-    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)
+.. autosummary::
+    :toctree: generated/
 
-            if logger_level is not None and logger_name in logger_level:
-                _logger.setLevel(logger_level[logger_name])
-            else:
-                _logger.setLevel(logging.WARNING)
-                
-                
-# 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
-else:
-    _logger_level = logging.INFO
-set_logger(["wrfrun", ], {"wrfrun": _logger_level})
+    check_path
+    rectify_domain_size
+    _calculate_domain_shape
+    calculate_domain_shape
+    _check_domain_shape
+    check_domain_shape
 
+Utility submodule.
+"""
 
-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.
-
-    :return:
-    :rtype:
-    """
-    set_logger(
-        ["cdsapi", "cfgrib", "datapi"],
-        {
-            "cdsapi": logging.INFO,
-            "cfgrib": logging.ERROR,
-            "datapi": logging.INFO
-        }
-    )
+from os import makedirs
+from os.path import exists
+from shutil import rmtree
 
 
 def check_path(*args, force=False):
-    """Check and create all the path in *args
-
-    Returns:
+    """
+    Check and create all the path in args.
 
+    :param args: Path list.
+    :type args: list[str]
+    :param force: If ``True``, delete existed directory and create a new.
+    :type force: bool
     """
     for _path in args:
         if exists(_path) and force:
@@ -85,32 +36,16 @@ def check_path(*args, force=False):
             makedirs(_path)
 
 
-def logger_add_file_handler(log_path: str):
-    """Add file handler to logger
-
-    Args:
-        log_path (str): Folder to place log file
-    """
-    # 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 rectify_domain_size(point_num: int, nest_ratio: int) -> int:
-    """Rectify domain size.
-
-    Args:
-        point_num (int): Point number of one side.
-        nest_ratio (int): The nesting ratio relative to the domain’s parent.
-
-    Returns:
-        int: New size.
+    """
+    Rectify domain size.
+
+    :param point_num: Point number of one side.
+    :type point_num: int
+    :param nest_ratio: The nesting ratio relative to the domain's parent.
+    :type nest_ratio: int
+    :return: New size.
+    :rtype: int
     """
     # calculate remainder
     point_num_mod = (point_num - 1) % nest_ratio
@@ -123,22 +58,23 @@ def rectify_domain_size(point_num: int, nest_ratio: int) -> int:
         point_num -= point_num_mod
     else:
         # # point_num_mod is closer to nest_ratio than (nest_ratio - 1)
-        point_num += (nest_ratio - point_num_mod)
+        point_num += nest_ratio - point_num_mod
 
     return point_num
 
 
 def _calculate_domain_shape(step: float, resolution: int, grid_resolution=110, nest_ratio=1) -> int:
-    """Calculate domain shape based on its area
-
-    Args:
-        step (float): Length of the side. Unit: degree.
-        resolution (int): Resolution of domain. Unit: km.
-        grid_resolution (int, optional): Resolution of grid. Unit: km. Defaults to 110.
-        nest_ratio (int, optional): The nesting ratio relative to the domain’s parent.
-
-    Returns:
-        tuple[int, int]: (length of x, length of y)
+    """
+    Calculate domain shape based on its area.
+
+    :param step: Length of the side. Unit: degree.
+    :type step: float
+    :param resolution: Resolution of domain. Unit: km.
+    :type resolution: int
+    :param grid_resolution: Resolution of grid. Unit: km. Defaults to 110.
+    :param nest_ratio: The nesting ratio relative to the domain’s parent.
+    :return: ``(length of x, length of y)``.
+    :rtype: int
     """
     # calculate res based on doamin area and resolution
     res = int(step * grid_resolution) // resolution
@@ -148,37 +84,39 @@ def _calculate_domain_shape(step: float, resolution: int, grid_resolution=110, n
     return rectify_domain_size(res, nest_ratio=nest_ratio)
 
 
-def calculate_domain_shape(lon_step: float, lat_step: float, resolution: int, grid_resolution=110, nest_ratio=1) -> tuple[int, int]:
-    """Calculate domain shape based on its area
-
-    Args:
-        lon_step (float): Length in X direction. Unit: degree.
-        lat_step (float): Length in Y direction. Unit: degree.
-        resolution (int): Resolution of domain. Unit: km.
-        grid_resolution (int, optional): Resolution of grid. Unit: km. Defaults to 110.
-        nest_ratio (int, optional): The nesting ratio relative to the domain’s parent.
-
-    Returns:
-        tuple[int, int]: (length of x, length of y)
+def calculate_domain_shape(
+    lon_step: float, lat_step: float, resolution: int, grid_resolution=110, nest_ratio=1
+) -> tuple[int, int]:
+    """
+    Calculate domain shape based on its area.
+
+    :param lon_step: Length in X direction. Unit: degree.
+    :type lon_step: float
+    :param lat_step: Length in Y direction. Unit: degree.
+    :type lat_step: float
+    :param resolution: Resolution of domain. Unit: km.
+    :type resolution: int
+    :param grid_resolution: Resolution of grid. Unit: km. Defaults to 110.
+    :param nest_ratio: The nesting ratio relative to the domain’s parent.
+    :return: ``(length of x, length of y)``.
+    :rtype: tuple[int, int]
     """
-
     return (
-        _calculate_domain_shape(
-            lon_step, resolution, grid_resolution=grid_resolution, nest_ratio=nest_ratio),
-        _calculate_domain_shape(
-            lat_step, resolution, grid_resolution=grid_resolution, nest_ratio=nest_ratio)
+        _calculate_domain_shape(lon_step, resolution, grid_resolution=grid_resolution, nest_ratio=nest_ratio),
+        _calculate_domain_shape(lat_step, resolution, grid_resolution=grid_resolution, nest_ratio=nest_ratio),
     )
 
 
 def _check_domain_shape(point_num: int, nest_ratio: int) -> bool:
-    """Check domain shape.
-
-    Args:
-        point_num (int): Point number of one side.
-        nest_ratio (int): The nesting ratio relative to the domain’s parent.
-
-    Returns:
-        bool: True if pass check else False.
+    """
+    Check domain shape.
+
+    :param point_num: Point number of one side.
+    :type point_num: int
+    :param nest_ratio: The nesting ratio relative to the domain’s parent.
+    :type nest_ratio: int
+    :return: True if pass check else False.
+    :rtype: bool
     """
     if (point_num - 1) % nest_ratio != 0:
         return False
@@ -187,76 +125,27 @@ def _check_domain_shape(point_num: int, nest_ratio: int) -> bool:
 
 
 def check_domain_shape(x_point_num: int, y_point_num: int, nest_ratio: int) -> tuple[bool, bool]:
-    """Check domain shape.
-
-    Args:
-        x_point_num (int): Point number of X side.
-        y_point_num (int): Point number of Y side.
-        nest_ratio (int): The nesting ratio relative to the domain’s parent.
-
-    Returns:
-        tuple[bool, bool]: Tuple of bool. True if pass check else False
+    """
+    Check domain shape.
+
+    :param x_point_num: Point number of X side.
+    :type x_point_num: int
+    :param y_point_num: Point number of Y side.
+    :type y_point_num: int
+    :param nest_ratio: The nesting ratio relative to the domain's parent.
+    :type nest_ratio: int
+    :return: Tuple of bool. True if pass check else False
+    :rtype: tuple[bool, bool]
     """
     return (
         _check_domain_shape(x_point_num, nest_ratio=nest_ratio),
-        _check_domain_shape(y_point_num, nest_ratio=nest_ratio)
+        _check_domain_shape(y_point_num, nest_ratio=nest_ratio),
     )
 
 
-def check_subprocess_status(status: subprocess.CompletedProcess):
-    """Check subprocess return code and print log if their return code != 0
-
-    Args:
-        status (CompletedProcess): Status from subprocess.
-    """
-    if status.returncode != 0:
-        # print command
-        command = status.args
-        logger.error(f"Failed to exec command: {command}")
-
-        # print log
-        logger.error(f"====== stdout ======")
-        logger.error(status.stdout.decode())
-        logger.error(f"====== ====== ======")
-        logger.error(f"====== stderr ======")
-        logger.error(status.stderr.decode())
-        logger.error(f"====== ====== ======")
-
-        # raise error
-        raise RuntimeError(f"Failed to exec command: '{command}'. Please check the log above.")
-
-
-def call_subprocess(command: list[str], work_path: Optional[str] = None, print_output=False):
-    """
-    Execute the given command in system shell.
-
-    :param command: A list contains the command and parameters to be executed.
-    :type command: list
-    :param work_path: The work path of the command.
-                      If None, works in current directory.
-    :type work_path: str | None
-    :param print_output: If print standard output and error in the logger.
-    :type print_output: bool
-    :return:
-    :rtype:
-    """
-    if work_path is not None:
-        origin_path = getcwd()
-        chdir(work_path)
-    else:
-        origin_path = None
-
-    status = subprocess.run(' '.join(command), shell=True, capture_output=True)
-
-    if origin_path is not None:
-        chdir(origin_path)
-
-    check_subprocess_status(status)
-
-    if print_output:
-        logger.info(status.stdout.decode())
-        logger.warning(status.stderr.decode())
-
-
-__all__ = ["logger", "check_path", "logger_add_file_handler", "calculate_domain_shape", "rectify_domain_size", "check_domain_shape",
-           "check_subprocess_status", "call_subprocess", "set_logger", "unify_logger_format"]
+__all__ = [
+    "check_path",
+    "calculate_domain_shape",
+    "rectify_domain_size",
+    "check_domain_shape",
+]

wrfrun/workspace/__init__.py

@@ -9,6 +9,7 @@ Submodules
 
 =================================    ===========================================================
 :doc:`core </api/workspace.core>`    Core functions of this submodule.
+:doc:`palm </api/workspace.palm>`    Functions to prepare workspace for PALM model.
 :doc:`wrf </api/workspace.wrf>`      Functions to prepare workspace for WPS/WRF model.
 =================================    ===========================================================
 
@@ -16,15 +17,16 @@ Workspace
 *********
 
 ``workspace`` is a collection of several directories where ``wrfrun``, extensions and numerical model works.
-These directories and their purpose are listed below.
+These directories and their purpose are listed below. ``$ROOT`` is the root work path set in config file,
+or ``$HOME/.config/wrfrun``.
 
 ===================================         ===========================================================
 Director Path                               Purpose
 ===================================         ===========================================================
-``/tmp/wrfrun``                             Store temporary files.
-``$HOME/.config/wrfrun``                    Main work directory.
-``$HOME/.config/wrfrun/replay``             Work directory for :doc:`replay <wrfrun.core.replay>`.
-``$HOME/.config/wrfrun/model``              Work directory for numerical models.
+``$ROOT``                                   Main work directory.
+``$ROOT/tmp``                               Store temporary files.
+``$ROOT/replay``                            Work directory for :doc:`replay </api/core.replay>`.
+``$ROOT/workspace``                         Work directory for numerical models.
 ===================================         ===========================================================
 
 .. toctree::
@@ -32,6 +34,7 @@ Director Path                               Purpose
     :hidden:
 
     core <workspace.core>
+    palm <workspace.palm>
     wrf <workspace.wrf>
 """
 

wrfrun/workspace/core.py

@@ -8,13 +8,16 @@ Core functions to prepare ``wrfrun`` workspace.
     :toctree: generated/
 
     prepare_workspace
+    check_workspace
 """
 
 from os.path import exists
 from shutil import rmtree
 
-from wrfrun.core import get_wrfrun_config
-from wrfrun.utils import check_path, logger
+from wrfrun.core import WRFRUN
+from wrfrun.log import check_path, logger
+
+from .palm import prepare_palm_workspace
 from .wrf import check_wrf_workspace, prepare_wrf_workspace
 
 
@@ -34,16 +37,17 @@ def prepare_workspace():
 
     1. :doc:`WPS/WRF model </api/workspace.wrf>`
     """
-    WRFRUNConfig = get_wrfrun_config()
-    logger.info(f"Initialize main workspace.")
+    WRFRUNConfig = WRFRUN.config
 
     wrfrun_temp_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_TEMP_PATH)
     workspace_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_WORKSPACE_ROOT)
     replay_work_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_WORKSPACE_REPLAY)
     output_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_OUTPUT_PATH)
 
+    logger.info(f"Initialize main workspace at: {workspace_path}")
+
     if exists(workspace_path):
-        logger.info(f"Remove old files in workspace.")
+        logger.info("Remove old files in workspace.")
         rmtree(workspace_path)
 
     # check folder
@@ -51,12 +55,14 @@ def prepare_workspace():
     check_path(replay_work_path)
     check_path(output_path)
 
-    func_map = {
-        "wrf": prepare_wrf_workspace
-    }
+    func_map = {"wrf": prepare_wrf_workspace, "palm": prepare_palm_workspace}
     model_configs = WRFRUNConfig["model"]
 
     for model_name in model_configs:
+        if model_name not in func_map:
+            logger.warning(f"Function to prepare '{model_name}' workspace not found, workspace may be incomplete")
+            continue
+
         func_map[model_name](model_configs[model_name])
 
 
@@ -67,7 +73,7 @@ def check_workspace() -> bool:
     :return: ``True`` if workspace exists, ``False`` otherwise.
     :rtype: bool
     """
-    WRFRUNConfig = get_wrfrun_config()
+    WRFRUNConfig = WRFRUN.config
 
     wrfrun_temp_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_TEMP_PATH)
     workspace_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_WORKSPACE_ROOT)
@@ -77,15 +83,17 @@ def check_workspace() -> bool:
     flag = True
     flag = flag & exists(wrfrun_temp_path) & exists(replay_work_path) & exists(output_path) & exists(workspace_path)
 
-    func_map = {
-        "wrf": check_wrf_workspace
-    }
+    func_map = {"wrf": check_wrf_workspace}
     model_configs = WRFRUNConfig["model"]
 
     for model_name in model_configs:
         if model_name == "debug_level":
             continue
 
+        if model_name not in func_map:
+            logger.info(f"Function to check '{model_name}' workspace not found, skip")
+            continue
+
         flag = flag & func_map[model_name](model_configs[model_name])
 
     return True

wrfrun/workspace/palm.py

@@ -0,0 +1,137 @@
+"""
+wrfrun.workspace.palm
+#####################
+
+Functions to prepare workspace for PALM model.
+
+.. autosummary::
+    :toctree: generated/
+
+    get_palm_workspace_path
+    prepare_palm_workspace
+"""
+
+from os import remove, symlink
+from os.path import abspath, exists, islink
+from shutil import copyfile, move, rmtree
+from typing import Literal
+
+from wrfrun.core import WRFRUN, WRFRunConfig
+from wrfrun.log import logger
+from wrfrun.utils import check_path
+
+WORKSPACE_PALM = ""
+
+
+def _register_palm_workspace_uri(wrfrun_config: WRFRunConfig):
+    """
+    This function doesn't register any URI.
+
+    This is a hook to initializes some global strings.
+
+    :param wrfrun_config: ``WRFRunConfig`` instance.
+    :type wrfrun_config: WRFRunConfig
+    """
+    global WORKSPACE_PALM
+
+    WORKSPACE_PALM = f"{wrfrun_config.WRFRUN_WORKSPACE_MODEL}/PALM"
+
+
+WRFRUN.set_config_register_func(_register_palm_workspace_uri)
+
+
+def get_palm_workspace_path(node: Literal["root", "job", "input", "output"] = "root") -> str:
+    """
+    Get workspace of PALM model.
+
+    :param node: Which dir.
+    :type node: str
+    :return: Workspace path.
+    :rtype: str
+    """
+    match node:
+        case "root":
+            return WORKSPACE_PALM
+
+        case "job":
+            return f"{WORKSPACE_PALM}/job"
+
+        case "input":
+            return f"{WORKSPACE_PALM}/job/wrfrun/INPUT"
+
+        case "output":
+            return f"{WORKSPACE_PALM}/job/wrfrun/OUTPUT"
+
+
+def prepare_palm_workspace(model_config: dict):
+    """
+    Initialize workspace for PALM model.
+
+    This function will check following paths,
+    and create them or delete old files inside:
+
+    1. ``$HOME/.config/wrfrun/model/PALM``
+
+    :param model_config: Model config.
+    :type model_config: dict
+    """
+    logger.info("Initialize workspace for PALM.")
+
+    WRFRUNConfig = WRFRUN.config
+
+    palm_path = model_config["palm_path"]
+    config_id = model_config["config_identifier"]
+    config_file = model_config["config_file_path"]
+
+    palm_path = abspath(palm_path)
+
+    if not (exists(palm_path) and exists(f"{palm_path}/bin")):
+        logger.error("Your PALM path is wrong.")
+        raise FileNotFoundError("Your PALM path is wrong.")
+
+    palm_work_path = WRFRUNConfig.parse_resource_uri(WORKSPACE_PALM)
+    workspace_job_path = f"{palm_work_path}/job"
+    workspace_input_path = f"{workspace_job_path}/wrfrun/INPUT"
+    check_path(palm_work_path, workspace_job_path, workspace_input_path, force=True)
+
+    if not exists(f"{palm_path}/bin/palmrun"):
+        logger.error("Script 'palmrun' not found in your PALM dir.")
+        raise FileNotFoundError("Script 'palmrun' not found in your PALM dir.")
+
+    symlink(f"{palm_path}/bin/palmrun", f"{palm_work_path}/palmrun")
+
+    # we need some tricks to hack palm runtime directory.
+    job_path = f"{palm_path}/JOBS"
+    if exists(job_path):
+        if islink(job_path):
+            remove(job_path)
+
+        else:
+            new_job_path = f"{job_path}_wrfrun_bak"
+            if exists(new_job_path):
+                rmtree(new_job_path)
+
+            logger.warning(f"You have an existed PALM JOBS dir: {job_path}")
+            logger.warning(f"wrfrun has renamed it to: {new_job_path}")
+            logger.warning("If you have important files in it, please backup it to other positions,")
+            logger.warning("cause wrfrun may delete the renamed depository in the future.")
+
+            move(job_path, new_job_path)
+
+    symlink(workspace_job_path, job_path)
+
+    # PALM config file.
+    if config_file == "":
+        if not exists(f"{palm_path}/.palm.config.default"):
+            logger.error(f"Can't find the default config file: '{palm_path}/.palm.config.default', please provide one.")
+            raise FileNotFoundError(
+                f"Can't find the default config file: '{palm_path}/.palm.config.default', please provide one."
+            )
+
+        copyfile(f"{palm_path}/.palm.config.default", f"{palm_work_path}/.palm.config.{config_id}")
+
+    else:
+        symlink(abspath(config_file), f"{palm_work_path}/.palm.config.{config_id}")
+
+
+__all__ = ["get_palm_workspace_path", "prepare_palm_workspace"]