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

wrfrun/cli.py

@@ -7,7 +7,7 @@ from shutil import copyfile
 import tomli
 import tomli_w
 
-from .core import WRFRUNConfig
+from .core import WRFRunConfig
 from .res import CONFIG_MAIN_TOML_TEMPLATE, CONFIG_WRF_TOML_TEMPLATE
 from .utils import logger
 
@@ -16,6 +16,9 @@ MODEL_MAP = {
 }
 
 
+wrfrun_config = WRFRunConfig("./.wrfrun")
+
+
 def _entry_init(args: argparse.Namespace):
     """
     Initialize a wrfrun project.
@@ -35,11 +38,11 @@ def _entry_init(args: argparse.Namespace):
     makedirs(f"{project_name}/configs")
     makedirs(f"{project_name}/data")
 
-    copyfile(WRFRUNConfig.parse_resource_uri(CONFIG_MAIN_TOML_TEMPLATE), f"{project_name}/config.toml")
+    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 = WRFRUNConfig.parse_resource_uri(MODEL_MAP[_model])
+            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}.")
@@ -95,7 +98,7 @@ def _entry_model(args: argparse.Namespace):
                 }
 
     for _new_model in new_models:
-        copyfile(WRFRUNConfig.parse_resource_uri(model_config_map[_new_model]), f"{config_dir_path}/{_new_model}.toml")
+        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)

wrfrun/core/base.py

@@ -13,7 +13,8 @@ Defines what :class:`ExecutableBase <Executable>` is, how it works and how ``wrf
     FileConfigDict
     ExecutableClassConfig
     ExecutableConfig
-    _ExecutableConfigRecord
+    ExecutableConfigRecord
+    create_recorder
     ExecutableBase
 
 Executable
@@ -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
 
 
@@ -276,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.
@@ -305,7 +306,6 @@ class _ExecutableConfigRecord:
 
         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 = {}
@@ -318,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):
         """
@@ -417,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:
@@ -588,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
@@ -845,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"]

wrfrun/core/config.py

@@ -11,6 +11,9 @@ All classes in this module is used to manage various configurations of ``wrfrun`
     _WRFRunConstants
     _WRFRunNamelist
     _WRFRunResources
+    init_wrfrun_config
+    get_wrfrun_config
+    set_register_func
 
 WRFRunConfig
 ************
@@ -20,18 +23,24 @@ It inherits from three classes: :class:`_WRFRunResources`, :class:`_WRFRunConsta
 Users can use the global variable ``WRFRUNConfig``, which is the instance of this class being created when users import ``wrfrun``.
 """
 
+# TODO:
+#   1. NEW FEATURE: Allow reading work directory from config file.
+#   2. The first one will be a break change, fix the following errors.
+#   3. The structure of wrfrun may need to be changed again.
+
+import threading
 from copy import deepcopy
 from os import environ, makedirs
 from os.path import abspath, dirname, exists
 from shutil import copyfile
 from sys import platform
-from typing import Optional, Tuple, Union
+from typing import Callable, Optional, Tuple, Union
 
 import f90nml
 import tomli
 import tomli_w
 
-from .error import ModelNameError, NamelistError, NamelistIDError, ResourceURIError, WRFRunContextError
+from .error import ModelNameError, NamelistError, NamelistIDError, ResourceURIError, WRFRunContextError, ConfigError
 from ..utils import logger
 
 
@@ -140,38 +149,42 @@ class _WRFRunConstants:
     Define all variables that will be used by other components.
     """
 
-    def __init__(self):
+    def __init__(self, work_dir: str):
         """
         Define all variables that will be used by other components.
 
         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.
+
+        :param work_dir: ``wrfrun`` work directory path.
+        :type work_dir: str
         """
         # check system
         if platform != "linux":
             logger.debug(f"Not Linux system!")
 
+        if work_dir != "" or platform != "linux":
             # set temporary dir path
-            self._WRFRUN_TEMP_PATH = "./tmp/wrfrun"
-            user_home_path = "./tmp/wrfrun"
+            self._WRFRUN_TEMP_PATH = abspath(f"{work_dir}/tmp")
+            self._WRFRUN_HOME_PATH = abspath(work_dir)
 
         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
-        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"
+            # 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
+            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_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"
@@ -534,21 +547,25 @@ class _WRFRunNamelist:
         :return: ``True`` if it is registered and loaded, else ``False``.
         :rtype: bool
         """
-        if namelist_id in self._namelist_id_list and self._namelist_dict:
+        if namelist_id in self._namelist_id_list and namelist_id in self._namelist_dict:
             return True
 
         else:
             return False
 
 
+_URI_REGISTER_FUNC_LIST: list[Callable[["WRFRunConfig"], None]] = []
+
+
 class WRFRunConfig(_WRFRunConstants, _WRFRunNamelist, _WRFRunResources):
     """
     Class to manage wrfrun config, runtime constants, namelists and resource files.
     """
     _instance = None
     _initialized = False
+    _lock = threading.Lock()
 
-    def __init__(self):
+    def __init__(self, work_dir: str):
         """
         This class provides various interfaces to access ``wrfrun``'s config, namelist values of NWP models,
         runtime constants and resource files by inheriting from:
@@ -556,23 +573,32 @@ class WRFRunConfig(_WRFRunConstants, _WRFRunNamelist, _WRFRunResources):
 
         An instance of this class called ``WRFRUNConfig`` will be created after the user import ``wrfrun``,
         and you should use the instance to access configs or other things instead of creating another instance.
+
+        :param work_dir: ``wrfrun`` work directory path.
+        :type work_dir: str
         """
         if self._initialized:
             return
 
-        _WRFRunConstants.__init__(self)
-        _WRFRunNamelist.__init__(self)
-        _WRFRunResources.__init__(self)
+        with self._lock:
+            global _URI_REGISTER_FUNC_LIST
 
-        self._config = {}
+            self._initialized = True
 
-        # register uri for wrfrun constants
-        for key, value in self._get_uri_map().items():
-            self.register_resource_uri(key, value)
+            _WRFRunConstants.__init__(self, work_dir)
+            _WRFRunNamelist.__init__(self)
+            _WRFRunResources.__init__(self)
+
+            self._config = {}
+
+            self._config_template_file_path = None
+
+            self._register_wrfrun_uris()
 
-        self._config_template_file_path = None
+            for _fun in _URI_REGISTER_FUNC_LIST:
+                _fun(self)
 
-        self._initialized = True
+            _URI_REGISTER_FUNC_LIST = []
 
     def __new__(cls, *args, **kwargs):
         if cls._instance is None:
@@ -580,6 +606,39 @@ class WRFRunConfig(_WRFRunConstants, _WRFRunNamelist, _WRFRunResources):
 
         return cls._instance
 
+    def __getattribute__(self, item):
+        if item not in ("_initialized", "_instance", "_lock", "__class__"):
+            if not object.__getattribute__(self, "_initialized"):
+                logger.error(f"`WRFRUNConfig` hasn't been initialized.")
+                logger.error(f"Use `WRFRun` to load config automatically, or use function `init_wrfrun_config` to load config manually.")
+                raise ConfigError(f"`WRFRUNConfig` hasn't been initialized.")
+
+        return object.__getattribute__(self, item)
+
+    @classmethod
+    def from_config_file(cls, config_file: str) -> "WRFRunConfig":
+        """
+        Read the config file and reinitialize.
+
+        :param config_file: Config file path.
+        :type config_file: str
+        :return: New instance
+        :rtype: WRFRunConfig
+        """
+        cls._initialized = False
+
+        with open(config_file, "rb") as f:
+            config = tomli.load(f)
+
+        instance = cls(work_dir=config["work_dir"])
+        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.
@@ -652,6 +711,11 @@ class WRFRunConfig(_WRFRunConstants, _WRFRunNamelist, _WRFRunResources):
         output_path = abspath(self["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``'s config to a file.
@@ -685,6 +749,18 @@ class WRFRunConfig(_WRFRunConstants, _WRFRunNamelist, _WRFRunResources):
 
         return deepcopy(self._config[item])
 
+    def __setitem__(self, key: str, value):
+        if key == "model":
+            logger.error(f"Use `update_model_config` to change model configurations.")
+            raise KeyError(f"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.
@@ -711,6 +787,23 @@ class WRFRunConfig(_WRFRunConstants, _WRFRunNamelist, _WRFRunResources):
 
         return deepcopy(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` 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["model"][model_name] = self["model"][model_name] | value
+
     def get_log_path(self) -> str:
         """
         Get the directory path to save logs.
@@ -762,6 +855,69 @@ class WRFRunConfig(_WRFRunConstants, _WRFRunNamelist, _WRFRunResources):
         super().write_namelist(save_path, namelist_id, overwrite)
 
 
-WRFRUNConfig = WRFRunConfig()
+WRFRUNConfig = WRFRunConfig.__new__(WRFRunConfig)
+
+
+def set_register_func(func: Callable[["WRFRunConfig"], None]):
+    """
+    Set the function to register URIs.
+
+    These functions should accept ``WRFRUNConfig`` instance,
+    and will be called when initializing ``WRFRUNConfig``.
+
+    If ``WRFRUNConfig`` has been initialized, ``func`` will be called immediately.
+
+    Normal users should use :meth:`WRFRunConfig.register_resource_uri`,
+    because ``WRFRUNConfig`` should (and must) be initialized.
+
+    :param func: Functions to register URIs.
+    :type func: Callable
+    """
+    global _URI_REGISTER_FUNC_LIST
+
+    if object.__getattribute__(WRFRUNConfig, "_initialized"):
+        func(WRFRUNConfig)
+
+    else:
+        if func not in _URI_REGISTER_FUNC_LIST:
+            _URI_REGISTER_FUNC_LIST.append(func)
+
+
+def init_wrfrun_config(config_file: str) -> WRFRunConfig:
+    """
+    Initialize ``WRFRUNConfig`` with the given config file.
+
+    :param config_file: Config file path.
+    :type config_file: str
+    :return: ``WRFRUNConfig`` instance.
+    :rtype: WRFRunConfig
+    """
+    global WRFRUNConfig
+
+    logger.info(f"Initialize `WRFRUNConfig` with config: {config_file}")
+
+    WRFRUNConfig = WRFRunConfig.from_config_file(config_file)
+
+    return WRFRUNConfig
+
+
+def get_wrfrun_config() -> WRFRunConfig:
+    """
+    Get ``WRFRUNConfig`` instance.
+
+    An exception :class:`ConfigError` will be raised if you haven't initialized it.
+
+    :return: ``WRFRUNConfig`` instance.
+    :rtype: WRFRunConfig
+    """
+    global WRFRUNConfig
+
+    if WRFRUNConfig is None:
+        logger.error(f"`WRFRUNConfig` hasn't been initialized.")
+        logger.error(f"Use `WRFRun` to load config automatically, or use function `init_wrfrun_config` to load config manually.")
+        raise ConfigError(f"`WRFRUNConfig` hasn't been initialized.")
+
+    return WRFRUNConfig
+
 
-__all__ = ["WRFRUNConfig"]
+__all__ = ["WRFRunConfig", "WRFRUNConfig", "init_wrfrun_config", "get_wrfrun_config", "set_register_func"]

wrfrun/core/error.py

@@ -106,5 +106,12 @@ class ModelNameError(WRFRunBasicError):
     pass
 
 
+class RecordError(WRFRunBasicError):
+    """
+    Exception indicates ``wrfrun`` can't record simulations.
+    """
+    pass
+
+
 __all__ = ["WRFRunBasicError", "ConfigError", "WRFRunContextError", "CommandError", "OutputFileError", "ResourceURIError", "InputFileError",
-           "NamelistError", "ExecRegisterError", "GetExecClassError", "ModelNameError", "NamelistIDError"]
+           "NamelistError", "ExecRegisterError", "GetExecClassError", "ModelNameError", "NamelistIDError", "RecordError"]

wrfrun/core/server.py

@@ -8,6 +8,7 @@ In order to report the progress to user, ``wrfrun`` provides :class:`WRFRunServe
 .. autosummary::
     :toctree: generated/
 
+    set_log_parse_func
     WRFRunServer
     WRFRunServerHandler
     stop_server
@@ -15,6 +16,7 @@ In order to report the progress to user, ``wrfrun`` provides :class:`WRFRunServe
 
 import socket
 import socketserver
+import threading
 from collections.abc import Callable
 from datetime import datetime
 from json import dumps
@@ -28,6 +30,24 @@ WRFRUN_SERVER_INSTANCE = None
 WRFRUN_SERVER_THREAD = None
 
 
+SET_LOG_PARSER_LOCK = threading.Lock()
+LOG_PARSER: Callable[[datetime], int] | None = None
+
+
+def set_log_parse_func(func: Callable[[datetime], int]):
+    """
+    Set log parse function used by socket server.
+
+    :param func: Function used to get simulated seconds from model's log file.
+                 If the function can't parse the simulated seconds, it should return ``-1``.
+    :type func: Callable[[datetime], int]
+    """
+    global SET_LOG_PARSER_LOCK, LOG_PARSER
+
+    with SET_LOG_PARSER_LOCK:
+        LOG_PARSER = func
+
+
 class WRFRunServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
     """
     A socket server to report time usage.
@@ -155,7 +175,7 @@ class WRFRunServerHandler(socketserver.StreamRequestHandler):
     ``status`` represents work status,
     ``progress`` represents simulation progress of the status in percentage.
     """
-    def __init__(self, request, client_address, server: WRFRunServer, log_parse_func: Callable[[datetime], int] | None) -> None:
+    def __init__(self, request, client_address, server: WRFRunServer) -> None:
         """
         :class:`WRFRunServer` handler.
 
@@ -165,15 +185,11 @@ class WRFRunServerHandler(socketserver.StreamRequestHandler):
         :type client_address:
         :param server: :class:`WRFRunServer` instance.
         :type server: WRFRunServer
-        :param log_parse_func: Function used to get simulated seconds from model's log file.
-                               If the function can't parse the simulated seconds, it should return ``-1``.
-        :type log_parse_func: Callable[[datetime], int]
         """
         super().__init__(request, client_address, server)
 
         # get server
         self.server: WRFRunServer = server
-        self.log_parse_func = log_parse_func
 
     def calculate_time_usage(self) -> int:
         """
@@ -202,10 +218,11 @@ class WRFRunServerHandler(socketserver.StreamRequestHandler):
         """
         start_date, simulate_seconds = self.server.get_model_simulate_settings()
 
-        if self.log_parse_func is None:
-            simulated_seconds = -1
-        else:
-            simulated_seconds = self.log_parse_func(start_date)
+        with SET_LOG_PARSER_LOCK:
+            if LOG_PARSER is None:
+                simulated_seconds = -1
+            else:
+                simulated_seconds = LOG_PARSER(start_date)
 
         if simulated_seconds > 0:
             progress = simulated_seconds * 100 // simulate_seconds
@@ -269,4 +286,4 @@ def stop_server(socket_ip: str, socket_port: int):
         logger.warning("Fail to stop WRFRunServer, maybe it doesn't start at all.")
 
 
-__all__ = ["WRFRunServer", "WRFRunServerHandler", "stop_server"]
+__all__ = ["WRFRunServer", "WRFRunServerHandler", "stop_server", "set_log_parse_func"]