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

wrfrun/core/error.py

@@ -1,80 +1,110 @@
+"""
+wrfrun.core.error
+#################
+
+This module defines all exceptions used in ``wrfrun``.
+
+.. autosummary::
+    :toctree: generated/
+
+    WRFRunBasicError
+    ConfigError
+    WRFRunContextError
+    CommandError
+    OutputFileError
+    ResourceURIError
+    InputFileError
+    NamelistError
+    NamelistIDError
+    ExecRegisterError
+    GetExecClassError
+    ModelNameError
+"""
+
+
 class WRFRunBasicError(Exception):
     """
-    Basic error type.
+    Basic exception class of ``wrfrun``. New exception **MUST** inherit this class.
     """
     pass
 
 
 class ConfigError(WRFRunBasicError):
     """
-    Config cannot be used error.
+    Exception indicates the config of ``wrfrun`` or NWP model can't be used.
     """
     pass
 
 
 class WRFRunContextError(WRFRunBasicError):
     """
-    Need to enter wrfrun context.
+    Exception indicates ``wrfrun`` is running out of the ``wrfrun`` context.
     """
     pass
 
 
 class CommandError(WRFRunBasicError):
     """
-    The command is wrong.
+    Exception indicates the command of ``Executable`` can't be executed successfully.
     """
     pass
 
 
-class LoadConfigError(WRFRunBasicError):
+class OutputFileError(WRFRunBasicError):
     """
-    Failed to load config.
+    Exception indicates ``wrfrun`` can't find any output files with the given rules.
     """
     pass
 
 
-class OutputFileError(WRFRunBasicError):
+class ResourceURIError(WRFRunBasicError):
     """
-    No output found.
+    Exception indicates ``wrfrun`` can't parse the URI.
     """
     pass
 
 
-class ResourceURIError(WRFRunBasicError):
+class InputFileError(WRFRunBasicError):
     """
-    Error about resource namespace.
+    Exception indicates ``wrfrun`` can't find specified input files.
     """
     pass
 
 
-class InputFileError(WRFRunBasicError):
+class NamelistError(WRFRunBasicError):
     """
-    Input file error.
+    Exception indicates ``wrfrun`` can't find the namelist user want to use.
     """
     pass
 
 
-class NamelistError(WRFRunBasicError):
+class NamelistIDError(WRFRunBasicError):
     """
-    Error about namelist.
+    Exception indicates ``wrfrun`` can't register the specified namelist id.
     """
     pass
 
 
 class ExecRegisterError(WRFRunBasicError):
+    """
+    Exception indicates ``wrfrun`` can't register the specified ``Executable``.
+    """
     pass
 
 
 class GetExecClassError(WRFRunBasicError):
+    """
+    Exception indicates ``wrfrun`` can't find the specified ``Executable``.
+    """
     pass
 
 
 class ModelNameError(WRFRunBasicError):
     """
-    Name of the model isn't found in the config file.
+    Exception indicates ``wrfrun`` can't find config of the specified NWP model in the config file.
     """
     pass
 
 
-__all__ = ["WRFRunBasicError", "ConfigError", "WRFRunContextError", "CommandError", "LoadConfigError", "OutputFileError", "ResourceURIError", "InputFileError",
-           "NamelistError", "ExecRegisterError", "GetExecClassError", "ModelNameError"]
+__all__ = ["WRFRunBasicError", "ConfigError", "WRFRunContextError", "CommandError", "OutputFileError", "ResourceURIError", "InputFileError",
+           "NamelistError", "ExecRegisterError", "GetExecClassError", "ModelNameError", "NamelistIDError"]

wrfrun/core/replay.py

@@ -1,6 +1,29 @@
+"""
+wrfrun.core.replay
+##################
+
+This module provides methods to read config from ``.replay`` file and reproduce the simulation.
+
+.. autosummary::
+    :toctree: generated/
+
+    WRFRunExecutableRegisterCenter
+    replay_config_generator
+
+WRFRUNExecDB
+************
+
+In order to load ``Executable`` correctly based on the stored ``name`` in ``.replay`` files,
+``wrfrun`` uses ``WRFRUNExecDB``, which is the instance of :class:`WRFRunExecutableRegisterCenter`,
+to records all ``Executable`` classes and corresponding ``name``.
+When ``wrfrun`` replays the simulation, it gets the right ``Executable`` from ``WRFRUNExecDB`` and executes it.
+"""
+
+from collections.abc import Generator
 from json import loads
 from os.path import exists
 from shutil import unpack_archive
+from typing import Any
 
 from .base import ExecutableBase, ExecutableConfig
 from .config import WRFRUNConfig
@@ -12,7 +35,8 @@ WRFRUN_REPLAY_URI = ":WRFRUN_REPLAY:"
 
 class WRFRunExecutableRegisterCenter:
     """
-    Record all executable class.
+    This class provides the method to records ``Executable``'s class with a unique ``name``.
+    Later you can get the class with the ``name``.
     """
     _instance = None
     _initialized = False
@@ -31,67 +55,76 @@ class WRFRunExecutableRegisterCenter:
 
         return cls._instance
 
-    def register_exec(self, unique_iq: str, cls: type):
+    def register_exec(self, name: str, cls: type):
         """
-        Register an executable class.
+        Register an ``Executable``'s class with a unique ``name``.
+
+        If the ``name`` has been used, :class:`ExecRegisterError` will be raised.
 
-        :param unique_iq: Unique ID.
-        :type unique_iq: str
-        :param cls: Class.
+        :param name: ``Executable``'s unique name.
+        :type name: str
+        :param cls: ``Executable``'s class.
         :type cls: type
-        :return:
-        :rtype:
         """
-        if unique_iq in self._exec_db:
-            logger.error(f"'{unique_iq}' has been registered.")
-            raise ExecRegisterError(f"'{unique_iq}' has been registered.")
+        if name in self._exec_db:
+            logger.error(f"'{name}' has been registered.")
+            raise ExecRegisterError(f"'{name}' has been registered.")
 
-        self._exec_db[unique_iq] = cls
+        self._exec_db[name] = cls
 
-    def is_registered(self, unique_id: str) -> bool:
+    def is_registered(self, name: str) -> bool:
         """
-        Check if an executable class has been registered.
+        Check if an ``Executable``'s class has been registered.
 
-        :param unique_id: Unique ID.
-        :type unique_id: str
+        :param name: ``Executable``'s unique name.
+        :type name: str
         :return: True or False.
         :rtype: bool
         """
-        if unique_id in self._exec_db:
+        if name in self._exec_db:
             return True
         else:
             return False
 
-    def get_cls(self, unique_id: str) -> type:
+    def get_cls(self, name: str) -> type:
         """
-        Get a registered executable class.
+        Get an ``Executable``'s class with the ``name``.
+
+        If the ``name`` can't be found, :class:`GetExecClassError` will be raised.
 
-        :param unique_id: Unique ID.
-        :type unique_id: str
-        :return: Executable class.
+        :param name: ``Executable``'s unique name.
+        :type name: str
+        :return: ``Executable``'s class.
         :rtype: type
         """
-        if unique_id not in self._exec_db:
-            logger.error(f"Executable class '{unique_id}' not found.")
-            raise GetExecClassError(f"Executable class '{unique_id}' not found.")
+        if name not in self._exec_db:
+            logger.error(f"Executable class '{name}' not found.")
+            raise GetExecClassError(f"Executable class '{name}' not found.")
 
-        return self._exec_db[unique_id]
+        return self._exec_db[name]
 
 
 WRFRUNExecDB = WRFRunExecutableRegisterCenter()
 
 
-def replay_config_generator(replay_config_file: str):
+def replay_config_generator(replay_config_file: str) -> Generator[tuple[str, ExecutableBase], Any, None]:
     """
-    Replay the simulations.
+    This method can read the ``.replay`` file and returns a generator which yields ``Executable`` and their names.
+    If this method doesn't find ``config.json`` in the ``.replay`` file, ``FileNotFoundError`` will be raised.
+
+    The ``Executable`` you get from the generator has been initialized so you can execute it directly.
+
+    >>> replay_file = "./example.replay"
+    >>> for name, _exec in replay_config_generator(replay_file):
+    >>>     _exec.replay()
 
-    :param replay_config_file: Replay file path.
+    :param replay_config_file: Path of the ``.replay`` file.
     :type replay_config_file: str
-    :return:
-    :rtype:
+    :return: A generator that yields: ``(name, Executable)``
+    :rtype: Generator
     """
     logger.info(f"Loading replay resources from: {replay_config_file}")
-    work_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_REPLAY_WORK_PATH)
+    work_path = WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_WORKSPACE_REPLAY)
 
     unpack_archive(replay_config_file, work_path, "zip")
 

wrfrun/core/server.py

@@ -1,9 +1,25 @@
+"""
+wrfrun.core.server
+##################
+
+Currently, ``wrfrun`` provides the method to reads log file of WRF and calculates simulation progress.
+In order to report the progress to user, ``wrfrun`` provides :class:`WRFRunServer` to set up a socket server.
+
+.. autosummary::
+    :toctree: generated/
+
+    WRFRunServer
+    WRFRunServerHandler
+    stop_server
+"""
+
 import socket
 import socketserver
-import subprocess
+from collections.abc import Callable
 from datetime import datetime
+from json import dumps
 from time import time
-from typing import Tuple, Optional
+from typing import Tuple
 
 from .config import WRFRUNConfig
 from ..utils import logger
@@ -12,48 +28,53 @@ WRFRUN_SERVER_INSTANCE = None
 WRFRUN_SERVER_THREAD = None
 
 
-def get_wrf_simulated_seconds(start_datetime: datetime, log_file_path: Optional[str] = None) -> int:
-    """
-    Get how many seconds wrf has integrated.
-
-    :param start_datetime: WRF start datetime.
-    :type start_datetime: datetime
-    :param log_file_path: Absolute path of the log file to be parsed.
-    :type log_file_path: str
-    :return: Seconds.
-    :rtype: int
+class WRFRunServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
     """
-    # use linux cmd to get the latest line of wrf log files
-    if log_file_path is None:
-        log_file_path = WRFRUNConfig.parse_resource_uri(f"{WRFRUNConfig.WRF_WORK_PATH}/rsl.out.0000")
-    res = subprocess.run(["tail", "-n", "1", log_file_path], capture_output=True)
-    log_text = res.stdout.decode()
+    A socket server to report time usage.
 
-    if not (log_text.startswith("d01") or log_text.startswith("d02")):
-        return -1
+    If you want to use the socket server, you need to give four arguments: ``start_date``,
+    ``wrf_simulate_seconds``, ``socket_address_port`` and ``WRFRunServerHandler``.
 
-    time_string = log_text.split()[1]
+    >>> start_date = datetime(2021, 3, 25, 8)
+    >>> simulate_seconds = 60 * 60 * 72
+    >>> socket_address_port = ("0.0.0.0", 54321)
+    >>> _wrfrun_server = WRFRunServer(start_date, simulate_seconds, socket_address_port, WRFRunServerHandler)
 
-    seconds = -1
-    try:
-        current_datetime = datetime.strptime(time_string, "%Y-%m-%d_%H:%M:%S")
-        # remove timezone info so we can calculate.
-        date_delta = current_datetime - start_datetime.replace(tzinfo=None)
-        seconds = date_delta.days * 24 * 60 * 60 + date_delta.seconds
+    Usually you don't need to create socket server manually, because :class:`WRFRun` will handle this.
 
-    except ValueError:
-        seconds = -1
+    .. py:attribute:: start_timestamp
+        :type: datetime
 
-    finally:
-        return seconds
+        The time the socket server start.
 
+    .. py:attribute:: start_date
+        :type: datetime
 
-class WRFRunServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
-    """
-    A socket server to report time usage.
+        The simulation's start date.
+
+    .. py:attribute:: wrf_simulate_seconds
+        :type: int
+
+        The total seconds the simulation will integrate.
+
+    .. py:attribute:: wrf_log_path
+        :type: str
+
+        Path of the log file the server will read.
     """
 
-    def __init__(self, start_date: datetime, wrf_simulate_seconds: int, *args, **kwargs) -> None:
+    def __init__(self, start_date: datetime, total_simulate_seconds: int, *args, **kwargs) -> None:
+        """
+
+        :param start_date: The simulation's start date.
+        :type start_date: datetime
+        :param total_simulate_seconds: The total seconds the simulation will integrate.
+        :type total_simulate_seconds: int
+        :param args: Other positional arguments passed to parent class.
+        :type args:
+        :param kwargs: Other keyword arguments passed to parent class.
+        :type kwargs:
+        """
         super().__init__(*args, **kwargs)
 
         # record the time the server starts
@@ -63,16 +84,11 @@ class WRFRunServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
         self.start_date = start_date
 
         # record how many seconds the wrf will integral
-        self.wrf_simulate_seconds = wrf_simulate_seconds
-
-        # we need to parse the log file to track the simulation progress.
-        self.wrf_log_path = WRFRUNConfig.parse_resource_uri(f"{WRFRUNConfig.WRF_WORK_PATH}/rsl.out.0000")
-        logger.debug("WRFRun Server will try to track simulation progress with following log files:")
-        logger.debug(f"WRF: {self.wrf_log_path}")
+        self.total_simulate_seconds = total_simulate_seconds
 
     def server_bind(self):
         """
-        Bind address and port.
+        Bind and listen on the address and port.
         """
         # reuse address and port to prevent the error `Address already in use`
         self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
@@ -88,32 +104,84 @@ class WRFRunServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
         """
         return self.start_timestamp
 
-    def get_wrf_simulate_settings(self) -> Tuple[datetime, int]:
+    def get_model_simulate_settings(self) -> Tuple[datetime, int]:
         """
         Get the start date of the case the NWP simulates and the total seconds of the simulation.
 
         :return: (start date, simulation seconds)
         :rtype: tuple
         """
-        return self.start_date, self.wrf_simulate_seconds
+        return self.start_date, self.total_simulate_seconds
 
 
 class WRFRunServerHandler(socketserver.StreamRequestHandler):
     """
-    Socket server handler.
+    :class:`WRFRunServer` handler.
+
+    This handler can report time usage and simulation progress of the running model, simulation settings, and stop the server:
 
+    1. If this handler receives ``"stop"``, it stops the server.
+    2. If this handler receives ``"debug"``, it returns the simulation settings in JSON.
+    3. It returns time usage and simulation progress in JSON when receiving any other messages.
+
+    **On receiving "stop"**
+
+    This handler will stop the server, and return a plain message ``Server stop``.
+
+    **On receiving "debug"**
+
+    This handler will return simulation settings in a JSON string like:
+
+    .. code-block:: json
+
+        {
+            "start_date": "2021-03-25 00:00",
+            "total_simulate_seconds": 360000,
+        }
+
+    **On receiving any other messages**
+
+    This handler will return time usage and simulation progress in a JSON string like:
+
+    .. code-block:: json
+
+        {
+            "usage": 3600,
+            "status": "geogrid",
+            "progress": 35,
+        }
+
+    where ``usage`` represents the seconds ``wrfrun`` has spent running the NWP model,
+    ``status`` represents work status,
+    ``progress`` represents simulation progress of the status in percentage.
     """
-    def __init__(self, request, client_address, server: WRFRunServer) -> None:
+    def __init__(self, request, client_address, server: WRFRunServer, log_parse_func: Callable[[datetime], int] | None) -> None:
+        """
+        :class:`WRFRunServer` handler.
+
+        :param request:
+        :type request:
+        :param client_address:
+        :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) -> str:
-        """Calculate time usage from server start (usually the time to run wrfrun)
+    def calculate_time_usage(self) -> int:
+        """
+        Calculate the duration from the server's start time to the present,
+        which represents the time ``wrfrun`` has spent running the NWP model.
 
-        Returns:
-            str: Time usage in `"%H:%M:%S"`
+        :return: Seconds.
+        :rtype: int
         """
         # get current timestamp
         current_timestamp = datetime.fromtimestamp(time())
@@ -122,42 +190,34 @@ class WRFRunServerHandler(socketserver.StreamRequestHandler):
         seconds_diff = current_timestamp - self.server.get_start_time()
         seconds_diff = seconds_diff.seconds
 
-        # calculate hours, minutes and seconds
-        seconds = seconds_diff % 60
-        minutes = (seconds_diff % 3600) // 60
-        hours = seconds_diff // 3600
-
-        time_usage = ":".join([
-            str(hours).rjust(2, '0'),
-            str(minutes).rjust(2, '0'),
-            str(seconds).rjust(2, '0')
-        ])
+        return seconds_diff
 
-        return time_usage
-
-    def calculate_progress(self) -> str:
+    def calculate_progress(self) -> tuple[str, int]:
         """
-        Calculate the simulation progress.
+        Read the log file and calculate the simulation progress.
 
-        :return:
-        :rtype:
+        :return: ``(status, progress)``. ``status`` represents work status,
+                 ``progress`` represents simulation progress of the status in percentage.
+        :rtype: tuple[str, int]
         """
-        start_date, simulate_seconds = self.server.get_wrf_simulate_settings()
+        start_date, simulate_seconds = self.server.get_model_simulate_settings()
 
-        simulated_seconds = get_wrf_simulated_seconds(start_date, self.server.wrf_log_path)
+        if self.log_parse_func is None:
+            simulated_seconds = -1
+        else:
+            simulated_seconds = self.log_parse_func(start_date)
 
         if simulated_seconds > 0:
             progress = simulated_seconds * 100 // simulate_seconds
-
         else:
-            progress = 0
+            progress = -1
 
         status = WRFRUNConfig.WRFRUN_WORK_STATUS
 
         if status == "":
             status = "*"
 
-        return f"{status}: {progress}%"
+        return status, progress
 
     def handle(self) -> None:
         """
@@ -172,25 +232,25 @@ class WRFRunServerHandler(socketserver.StreamRequestHandler):
             self.wfile.write(f"Server stop\n".encode())
         
         elif msg == "debug":
-            start_date, simulate_seconds = self.server.get_wrf_simulate_settings()
+            start_date, simulate_seconds = self.server.get_model_simulate_settings()
             start_date = start_date.strftime("%Y-%m-%d %H:%M")
             
-            self.wfile.write(f"{start_date}\n{simulate_seconds}\n".encode())
+            self.wfile.write(dumps({"start_date": start_date, "total_seconds": simulate_seconds}).encode())
             
         else:
-            progress = self.calculate_progress()
+            status, progress = self.calculate_progress()
             time_usage = self.calculate_time_usage()
 
             # send the message
-            self.wfile.write(f"{progress}\n{time_usage}".encode())
+            self.wfile.write(dumps({"usage": time_usage, "status": status, "progress": progress}).encode())
 
 
 def stop_server(socket_ip: str, socket_port: int):
-    """Try to stop server.
+    """
+    Stop the socket server.
 
-    Args:
-        socket_ip (str): Server IP.
-        socket_port (int): Server Port.
+    :param socket_ip: Server address.
+    :param socket_port: Server port.
     """
     try:
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -209,4 +269,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", "get_wrf_simulated_seconds", "stop_server"]
+__all__ = ["WRFRunServer", "WRFRunServerHandler", "stop_server"]

wrfrun/data.py

@@ -6,12 +6,11 @@ from typing import Union, List, Tuple
 from pandas import date_range
 # from seafog import goos_sst_find_data
 
-import cdsapi
-
 from .core.config import WRFRUNConfig
 from .utils import logger
 
-CDS_CLIENT = cdsapi.Client()
+# lazy initialize
+CDS_CLIENT = None
 
 
 class ERA5CONFIG:
@@ -214,6 +213,8 @@ def find_era5_data(date: Union[List[str], List[datetime]], area: Tuple[int, int,
     Returns: data path
 
     """
+    global CDS_CLIENT
+
     # check variables and datasets
     if not _check_variables_and_datasets(variables, dataset):
         logger.error(
@@ -281,8 +282,12 @@ def find_era5_data(date: Union[List[str], List[datetime]], area: Tuple[int, int,
             exit(1)
 
     # download data
-    logger.info(
-        f"Downloading data to {save_path}, it may take several tens of minutes, please wait...")
+    logger.info(f"Downloading data to {save_path}, it may take several tens of minutes, please wait...")
+
+    if CDS_CLIENT is None:
+        import cdsapi
+        CDS_CLIENT = cdsapi.Client()
+
     CDS_CLIENT.retrieve(dataset, params_dict, save_path)
 
     return save_path

wrfrun/extension/__init__.py

@@ -1 +1,29 @@
+"""
+wrfrun.extension
+################
+
+Through extensions, developers can add more functionality to ``wrfrun``.
+``wrfrun`` now has the following extensions:
+
+========================================= ==========================================
+:doc:`goos_sst </api/extension.goos_sst>` Process NEAR-GOOS SST data.
+:doc:`littler </api/extension.littler>`   Process LITTLE_R observation data.
+========================================= ==========================================
+
+Other Submodules
+****************
+
+=================================== ==========================================
+:doc:`utils </api/extension.utils>` Utility methods can be used by extensions.
+=================================== ==========================================
+
+.. toctree::
+    :maxdepth: 1
+    :hidden:
+
+    goos_sst <extension.goos_sst>
+    littler <extension.littler>
+    utils <extension.utils>
+"""
+
 from .utils import *