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

wrfrun/core/__init__.py

@@ -1,3 +1,36 @@
+"""
+wrfrun.core
+###########
+
+The core functionalities of ``wrfrun`` are all implemented in the submodule ``wrfrun.core``,
+such as processing namelist files, managing resource files required during model's running (e.g., VTable files),
+calling the numerical model and process its output and log files, monitoring simulation process via model's log,
+recording and replaying the simulation.
+
+The functionalities are split into several submodules, which are listed in the table below.
+
+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.
+================================ ========================================================
+
+.. toctree:: 
+    :maxdepth: 1
+    :hidden:
+
+    base <core.base>
+    config <core.config>
+    error <core.error>
+    replay <core.replay>
+    server <core.server>
+"""
+
 from .base import *
 from .config import *
 from .error import *

wrfrun/core/base.py

@@ -1,4 +1,49 @@
+"""
+wrfrun.core.base
+################
+
+Defines what :class:`ExecutableBase <Executable>` is, how it works and how ``wrfrun`` records simulations.
+
+.. autosummary::
+    :toctree: generated/
+
+    check_subprocess_status
+    call_subprocess
+    InputFileType
+    FileConfigDict
+    ExecutableClassConfig
+    ExecutableConfig
+    _ExecutableConfigRecord
+    ExecutableBase
+
+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``.
+``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:
+
+* Store all the information about the program (e.g., its inputs and outputs, its configuration).
+* Provide the interface to import and export ``Executable``'s config.
+* Support the ``replay`` functionality of ``wrfrun``.
+
+If you want to use all the ``wrfrun``'s features, you **HAVE TO** implement your code with :class:`ExecutableBase`.
+
+ExecConfigRecorder
+******************
+
+To be able to record the whole simulation, ``wrfrun`` introduces the global variable ``ExecConfigRecorder``, 
+which is an instance of class :class:`_ExecutableConfigRecord`.
+``ExecConfigRecorder`` can store the config of any ``Executable`` in the correct order and export them to a JSON file,
+enabling the precise replay of the simulation. Depending on the internal implementation of the ``Executable``,
+the recorded config may even include the complete namelist values, Vtable config, and more.
+Please check :meth:`ExecutableBase.generate_custom_config` for more details.
+"""
+
 import subprocess
+from copy import deepcopy
 from enum import Enum
 from json import dumps
 from os import chdir, getcwd, listdir, makedirs, remove, symlink
@@ -15,7 +60,8 @@ from ..utils import check_path, logger
 
 def check_subprocess_status(status: subprocess.CompletedProcess):
     """
-    Check subprocess return code and print log if ``return_code != 0``.
+    Check subprocess return code.
+    An ``RuntimeError`` exception will be raised if ``return_code != 0``, and the ``stdout`` and ``stderr`` of the subprocess will be logged.
 
     :param status: Status from subprocess.
     :type status: CompletedProcess
@@ -39,7 +85,7 @@ def check_subprocess_status(status: subprocess.CompletedProcess):
 
 def call_subprocess(command: list[str], work_path: Optional[str] = None, print_output=False):
     """
-    Execute the given command in system shell.
+    Execute the given command in the system shell.
 
     :param command: A list contains the command and parameters to be executed.
     :type command: list
@@ -48,8 +94,6 @@ def call_subprocess(command: list[str], work_path: Optional[str] = None, print_o
     :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()
@@ -88,9 +132,21 @@ def _json_default(obj):
 
 class InputFileType(Enum):
     """
-    Input file type.
-    ``WRFRUN_RES`` means the input file is from the NWP or wrfrun package.
-    ``CUSTOM_RES`` means the input file is from the user, which may be a customized file.
+    This class is an ``Enum`` class, providing the following values:
+
+    .. py:attribute:: WRFRUN_RES
+        :type: int
+        :value: 1
+
+        Indicating resource files are from the model or ``wrfrun``.
+        ``wrfrun`` won't save these files when recording the simulation.
+
+    .. py:attribute:: CUSTOM_RES
+        :type: int
+        :value: 2
+
+        Indicating resource files are provided by the user.
+        ``wrfrun`` will also save these files to ``.replay`` file when recording the simulation to ensure the simulation is replayable.
     """
     WRFRUN_RES = 1
     CUSTOM_RES = 2
@@ -98,7 +154,32 @@ class InputFileType(Enum):
 
 class FileConfigDict(TypedDict):
     """
-    Dict class to give information to process files.
+    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
+
+        A real file path or a valid URI which can be converted to a file path.
+
+    .. py:attribute:: save_path
+        :type: str
+
+        Save path of the file.
+
+    .. py:attribute:: save_name
+        :type: str
+
+        Save name of the file.
+
+    .. 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.
+
+    .. py:attribute:: is_output
+        :type: bool
+
+        If the file is model's output. Output file will never be saved to ``.replay`` file.
     """
     file_path: str
     save_path: str
@@ -109,7 +190,17 @@ class FileConfigDict(TypedDict):
 
 class ExecutableClassConfig(TypedDict):
     """
-    Executable class initialization config template.
+    This dict is used to store arguments of ``Executable``'s ``__init__`` function.
+
+    .. py:attribute:: class_args
+        :type: tuple
+
+        Positional arguments of the class.
+
+    .. py:attribute:: class_kwargs
+        :type: dict
+
+        Keyword arguments of the class.
     """
     # only list essential config
     class_args: tuple
@@ -118,7 +209,57 @@ class ExecutableClassConfig(TypedDict):
 
 class ExecutableConfig(TypedDict):
     """
-    Executable config template.
+    This dict is used to store all configs of a :class:`ExecutableBase`.
+
+    .. py:attribute:: name
+        :type: str
+
+        Name of the executable. Each type of executable has a unique name.
+
+    .. py:attribute:: cmd
+        :type: str | list[str]
+
+        Command of the executable.
+
+    .. py:attribute:: work_path
+        :type: str | None
+
+        Work path of the executable.
+
+    .. py:attribute:: mpi_use
+        :type: bool
+
+        If the executable will use MPI.
+
+    .. py:attribute:: mpi_cmd
+        :type: str | None
+
+        Command name of the MPI.
+
+    .. py:attribute:: mpi_core_num
+        :type: int | None
+
+        Number of the CPU core to use with MPI.
+
+    .. py:attribute:: class_config
+        :type: ExecutableClassConfig | None
+
+        A dict stores arguments of ``Executable``'s ``__init__`` function.
+
+    .. py:attribute:: input_file_config
+        :type: list[FileConfigDict] | None
+
+        A list stores information about input files of the executable.
+
+    .. py:attribute:: output_file_config
+        :type: list[FileConfigDict] | None
+
+        A list stores information about output files of the executable.
+
+    .. py:attribute:: custom_config
+        :type: dict | None
+
+        A dict that can be used by subclass to store other configs.
     """
     name: str
     cmd: Union[str, list[str]]
@@ -134,17 +275,16 @@ class ExecutableConfig(TypedDict):
 
 class _ExecutableConfigRecord:
     """
-    Record executable configs and export them.
+    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):
         """
-        Record executable configs and export them.
 
         :param save_path: Save path of the exported config file.
-        :type save_path: str
+        :type save_path: str | None
         :param include_data: If includes input data.
         :type include_data: bool
         """
@@ -179,8 +319,8 @@ class _ExecutableConfigRecord:
         """
         Reinitialize this instance.
 
-        :return:
-        :rtype:
+        :return: New instance.
+        :rtype: _ExecutableConfigRecord
         """
         self._initialized = False
         return _ExecutableConfigRecord(save_path, include_data)
@@ -191,8 +331,6 @@ class _ExecutableConfigRecord:
 
         :param exported_config: Executable config.
         :type exported_config: ExecutableConfig
-        :return:
-        :rtype:
         """
         if not self.include_data:
             self._recorded_config.append(exported_config)
@@ -236,19 +374,13 @@ class _ExecutableConfigRecord:
 
     def clear_records(self):
         """
-        Clean old configs.
-
-        :return:
-        :rtype:
+        Clean recorded configs.
         """
         self._recorded_config = []
 
     def export_replay_file(self):
         """
-        Save replay file to the specific save path.
-
-        :return:
-        :rtype:
+        Save replay file to the save path.
         """
         if len(self._recorded_config) == 0:
             logger.warning("No replay config has been recorded.")
@@ -288,12 +420,37 @@ ExecConfigRecorder = _ExecutableConfigRecord()
 class ExecutableBase:
     """
     Base class for all executables.
+    
+    .. py:attribute:: class_config
+        :type: ExecutableClassConfig
+        :value: {"class_args": (), "class_kwargs": {}}
+
+        A dict stores arguments of ``Executable``'s ``__init__`` function.
+
+    .. py:attribute:: custom_config
+        :type: dict
+        :value: {}
+
+        A dict that can be used by subclass to store custom configs.
+
+    .. py:attribute:: input_file_config
+        :type: list[FileConfigDict]
+        :value: []
+
+        A list stores information about input files of the executable.
+
+    .. py:attribute:: output_file_config
+        :type: list[FileConfigDict]
+        :value: []
+
+        A list stores information about output files of the executable.
+
     """
     _instance = None
 
-    def __init__(self, name: str, cmd: Union[str, list[str]], work_path: str, mpi_use=False, mpi_cmd: Optional[str] = None, mpi_core_num: Optional[int] = None):
+    def __init__(self, name: str, cmd: Union[str, list[str]], work_path: str,
+                 mpi_use=False, mpi_cmd: Optional[str] = None, mpi_core_num: Optional[int] = None):
         """
-        Base class for all executables.
 
         :param name: Unique name to identify different executables.
         :type name: str
@@ -338,19 +495,23 @@ class ExecutableBase:
 
     def generate_custom_config(self):
         """
-        Generate custom configs.
+        Generate custom configs. This method should be overwritten in the child class,
+        and **MUST STORE THE CUSTOM CONFIG IN THE ATTRIBUTE** :attr:`ExecutableBase.custom_config`,
+        or it will do nothing except print a debug log.
 
-        :return:
-        :rtype:
+        You can export various configs in this method, like the complete namelist values of a NWP model binary,
+        or the path of Vtable file this executable will use.
+
+        If you overwrite this method to generate custom configs,
+        you also have to overwrite :meth:`ExecutableBase.load_custom_config` to load your custom configs.
         """
         logger.debug(f"Method 'generate_custom_config' not implemented in '{self.name}'")
 
     def load_custom_config(self):
         """
         Load custom configs.
-
-        :return:
-        :rtype:
+        This method should be overwritten in the child class to process the custom config stored in :attr:`ExecutableBase.custom_config`,
+        or it will do nothing except print a debug log.
         """
         logger.debug(f"Method 'load_custom_config' not implemented in '{self.name}'")
 
@@ -370,20 +531,18 @@ class ExecutableBase:
             "mpi_use": self.mpi_use,
             "mpi_cmd": self.mpi_cmd,
             "mpi_core_num": self.mpi_core_num,
-            "class_config": self.class_config,
-            "custom_config": self.custom_config,
-            "input_file_config": self.input_file_config,
-            "output_file_config": self.output_file_config
+            "class_config": deepcopy(self.class_config),
+            "custom_config": deepcopy(self.custom_config),
+            "input_file_config": deepcopy(self.input_file_config),
+            "output_file_config": deepcopy(self.output_file_config)
         }
 
     def load_config(self, config: ExecutableConfig):
         """
-        Load config from a dict.
+        Load executable config from a dict.
 
         :param config: Config dict. It must contain some essential keys. Check ``ExecutableConfig`` for details.
         :type config: ExecutableConfig
-        :return:
-        :rtype:
         """
         if "name" not in config:
             logger.error("A valid config is required. Please check ``ExecutableConfig``.")
@@ -398,10 +557,10 @@ class ExecutableBase:
         self.mpi_use = config["mpi_use"]
         self.mpi_cmd = config["mpi_cmd"]
         self.mpi_core_num = config["mpi_core_num"]
-        self.class_config = config["class_config"]
-        self.custom_config = config["custom_config"]
-        self.input_file_config = config["input_file_config"]
-        self.output_file_config = config["output_file_config"]
+        self.class_config = deepcopy(config["class_config"])
+        self.custom_config = deepcopy(config["custom_config"])
+        self.input_file_config = deepcopy(config["input_file_config"])
+        self.output_file_config = deepcopy(config["output_file_config"])
 
         self.load_custom_config()
 
@@ -409,23 +568,22 @@ class ExecutableBase:
         """
         This method will be called when replay the simulation.
         This method should take care every job that will be done when replaying the simulation.
-
-        :return:
-        :rtype:
+        By default, this method will call ``__call__`` method of the instance.
         """
         logger.debug(f"Method 'replay' not implemented in '{self.name}', fall back to default action.")
         self()
 
-    def add_input_files(self, input_files: Union[str, list[str], FileConfigDict, list[FileConfigDict]], is_data=True, is_output=True):
+    def add_input_files(self, input_files: Union[str, list[str], FileConfigDict, list[FileConfigDict]],
+                        is_data=True, is_output=True):
         """
-        Add input files the extension will use.
+        Add input files the executable will use.
 
         You can give a single file path or a list contains files' path.
 
         >>> self.add_input_files("data/custom_file")
         >>> self.add_input_files(["data/custom_file_1", "data/custom_file_2"])
 
-        You can give more information with a ``FileConfigDict``, like the path and the name to store, and if it is data.
+        You can give more information with a ``FileConfigDict``.
 
         >>> file_dict: FileConfigDict = {
         ...     "file_path": "data/custom_file.nc",
@@ -452,12 +610,14 @@ class ExecutableBase:
         ... }
         >>> self.add_input_files([file_dict_1, file_dict_2])
 
-        :param input_files: Custom files' path.
+        Please check :class:`FileConfigDict` for more details.
+
+        :param input_files: Custom files.
         :type input_files: str | list | dict
-        :param is_data: If its data. This parameter will be overwritten by the value in ``input_files``.
+        :param is_data: If it is a data file. This parameter will be overwritten by the value in ``input_files``.
         :type is_data: bool
-        :return:
-        :rtype:
+        :param is_output: If it is an output from another executable. This parameter will be overwritten by the value in ``input_files``.
+        :type is_output: bool
         """
         if isinstance(input_files, str):
             self.input_file_config.append(
@@ -496,22 +656,39 @@ class ExecutableBase:
         endswith: Union[None, str, tuple[str, ...]] = None, outputs: Union[None, str, list[str]] = None, no_file_error=True
     ):
         """
-        Add save file rules.
+        Find and save model's outputs to the output save path.
+        An ``OutputFileError`` exception will be raised if no file can be found and ``no_file_error==True``.
+
+        You can give the specific path of a file or multiple files.
+
+        >>> self.add_output_files(outputs="wrfout.d01")
+        >>> self.add_output_files(outputs=["wrfout.d01", "wrfout.d02"])
+
+        If you have too many outputs, but they have the same prefix or postfix, you can use ``startswith`` or ``endswith``.
+
+        >>> self.add_output_files(startswith="rsl.out.")
+        >>> self.add_output_files(endswith="log")
+        >>> self.add_output_files(startswith=("rsl", "wrfout"), endswith="log")
+
+        ``startswith``, ``endswith`` and ``outputs`` can be used together.
 
-        :param output_dir: Output dir paths.
+        ``output_dir`` specify the search path of outputs, by default it is the work path of the executable.
+        You can change its value if the output path of the executable isn't its work path.
+
+        >>> self.add_output_files(output_dir=f"/absolute/dir/path", outputs=...)
+
+        :param output_dir: Search path of outputs.
         :type output_dir: str
-        :param save_path: Save path.
+        :param save_path: New save path of outputs. By default, it is ``f"{WRFRUNConfig.WRFRUN_OUTPUT_PATH}/{self.name}"``.
         :type save_path: str
         :param startswith: Prefix string or prefix list of output files.
         :type startswith: str | list
-        :param endswith: Postfix string or Postfix list of output files.
+        :param endswith: Postfix string or postfix list of output files.
         :type endswith: str | list
         :param outputs: Files name list. All files in the list will be saved.
         :type outputs: str | list
-        :param no_file_error: If True, an error will be raised with the ``error_message``. Defaults to True.
+        :param no_file_error: If True, an OutputFileError will be raised if no output file can be found. Defaults to True.
         :type no_file_error: bool
-        :return:
-        :rtype:
         """
         if WRFRUNConfig.FAKE_SIMULATION_MODE:
             return
@@ -567,10 +744,7 @@ class ExecutableBase:
 
     def before_exec(self):
         """
-        Prepare input files.
-
-        :return:
-        :rtype:
+        Prepare input files before executing the external program.
         """
         if WRFRUNConfig.FAKE_SIMULATION_MODE:
             logger.info(f"We are in fake simulation mode, skip preparing input files for '{self.name}'")
@@ -603,10 +777,7 @@ class ExecutableBase:
 
     def after_exec(self):
         """
-        Save outputs and logs.
-        
-        :return: 
-        :rtype: 
+        Save outputs and logs after executing the external program.
         """
         if WRFRUNConfig.FAKE_SIMULATION_MODE:
             logger.info(f"We are in fake simulation mode, skip saving outputs for '{self.name}'")
@@ -639,9 +810,6 @@ class ExecutableBase:
     def exec(self):
         """
         Execute the given command.
-
-        :return:
-        :rtype:
         """
         work_path = WRFRUNConfig.parse_resource_uri(self.work_path)
 
@@ -664,7 +832,7 @@ class ExecutableBase:
 
     def __call__(self):
         """
-        Execute the given command.
+        Execute the given command by calling ``before_exec``, ``exec`` and ``after_exec``.
 
         :return:
         :rtype: