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

wrfrun/core/server.py

@@ -8,7 +8,7 @@ In order to report the progress to user, ``wrfrun`` provides :class:`WRFRunServe
 .. autosummary::
     :toctree: generated/
 
-    get_wrf_simulated_seconds
+    set_log_parse_func
     WRFRunServer
     WRFRunServerHandler
     stop_server
@@ -16,11 +16,12 @@ In order to report the progress to user, ``wrfrun`` provides :class:`WRFRunServe
 
 import socket
 import socketserver
-import subprocess
+import threading
+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
@@ -29,38 +30,22 @@ WRFRUN_SERVER_INSTANCE = None
 WRFRUN_SERVER_THREAD = None
 
 
-def get_wrf_simulated_seconds(start_datetime: datetime, log_file_path: Optional[str] = None) -> int:
-    """
-    Read the latest line of WRF's log file and calculate 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: Integrated seconds. If this method fails to calculate the time, the returned value is ``-1``.
-    :rtype: int
-    """
-    # 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()
+SET_LOG_PARSER_LOCK = threading.Lock()
+LOG_PARSER: Callable[[datetime], int] | None = None
 
-    if not (log_text.startswith("d01") or log_text.startswith("d02")):
-        return -1
 
-    time_string = log_text.split()[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
+def set_log_parse_func(func: Callable[[datetime], int]):
+    """
+    Set log parse function used by socket server.
 
-    except ValueError:
-        seconds = -1
+    :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
 
-    return seconds
+    with SET_LOG_PARSER_LOCK:
+        LOG_PARSER = func
 
 
 class WRFRunServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
@@ -98,13 +83,13 @@ class WRFRunServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
         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 wrf_simulate_seconds: The total seconds the simulation will integrate.
-        :type wrf_simulate_seconds: int
+        :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.
@@ -119,12 +104,7 @@ 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):
         """
@@ -144,33 +124,80 @@ 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:
+        """
+        :class:`WRFRunServer` handler.
+
+        :param request:
+        :type request:
+        :param client_address:
+        :type client_address:
+        :param server: :class:`WRFRunServer` instance.
+        :type server: WRFRunServer
+        """
         super().__init__(request, client_address, server)
 
         # get server
         self.server: WRFRunServer = server
 
-    def calculate_time_usage(self) -> str:
+    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.
 
-        :return: A time string in ``%H:%M:%S`` format.
-        :rtype: str
+        :return: Seconds.
+        :rtype: int
         """
         # get current timestamp
         current_timestamp = datetime.fromtimestamp(time())
@@ -179,42 +206,35 @@ 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
+        return seconds_diff
 
-        time_usage = ":".join([
-            str(hours).rjust(2, '0'),
-            str(minutes).rjust(2, '0'),
-            str(seconds).rjust(2, '0')
-        ])
-
-        return time_usage
-
-    def calculate_progress(self) -> str:
+    def calculate_progress(self) -> tuple[str, int]:
         """
         Read the log file and calculate the simulation progress.
 
-        :return: A JSON string with two keys: ``status`` and ``progress``.
-        :rtype: str
+        :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)
+        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
-
         else:
-            progress = 0
+            progress = -1
 
         status = WRFRUNConfig.WRFRUN_WORK_STATUS
 
         if status == "":
             status = "*"
 
-        return dumps({"status": status, "progress": progress})
+        return status, progress
 
     def handle(self) -> None:
         """
@@ -229,17 +249,17 @@ 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):
@@ -266,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", "get_wrf_simulated_seconds", "stop_server"]
+__all__ = ["WRFRunServer", "WRFRunServerHandler", "stop_server", "set_log_parse_func"]

wrfrun/data.py

@@ -1,14 +1,15 @@
 from datetime import datetime
 from os import makedirs
-from os.path import exists, dirname
-from typing import Union, List, Tuple
+from os.path import dirname, exists
+from typing import List, Tuple, Union
 
 from pandas import date_range
-# from seafog import goos_sst_find_data
 
-from .core.config import WRFRUNConfig
+from .core import get_wrfrun_config
 from .utils import logger
 
+# from seafog import goos_sst_find_data
+
 # lazy initialize
 CDS_CLIENT = None
 
@@ -264,8 +265,7 @@ def find_era5_data(date: Union[List[str], List[datetime]], area: Tuple[int, int,
     # check if we need to add pressure_level to params dict
     if dataset == ERA5CONFIG.DATASET_ERA5_PRESSURE_LEVEL:
         if pressure_level is None:
-            logger.error(
-                f"You need to provide pressure levels to download data")
+            logger.error("You need to provide pressure levels to download data")
             exit(1)
         # convert value to str
         if not isinstance(pressure_level, list):
@@ -277,12 +277,11 @@ def find_era5_data(date: Union[List[str], List[datetime]], area: Tuple[int, int,
         if _check_pressure_level(pressure_level):   # type: ignore
             params_dict["pressure_level"] = pressure_level
         else:
-            logger.error(
-                f"You have passed wrong pressure level to download data, check it")
+            logger.error("You have passed wrong pressure level to download data, check it")
             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
@@ -299,29 +298,28 @@ def prepare_wps_input_data(area: Tuple[int, int, int, int]):
     Args:
         area (Tuple[int, int, int, int]): Range of longitude and latitude, `[lon1, lon2, lat1, lat2]`.
     """
-    wrf_config = WRFRUNConfig.get_model_config("wrf")
+    wrf_config = get_wrfrun_config().get_model_config("wrf")
     # get start and end date from config
     start_date = wrf_config["time"]["start_date"]
     end_date = wrf_config["time"]["end_date"]
 
     # remove second part
-    start_date = start_date[:-3]
-    end_date = end_date[:-3]
+    start_date = start_date.strftime("%Y-%m-%d %H:%M")
+    end_date = end_date.strftime("%Y-%m-%d %H:%M")
 
     # get hour step
     hour_step = wrf_config["time"]["input_data_interval"] // 3600
 
     # get data save path
-    bg_save_path = wrf_config["wps_input_data_folder"]
-    sst_save_path = wrf_config["near_goos_data_folder"]
+    data_save_path = get_wrfrun_config().get_input_data_path()
 
     # download data
     logger.info(f"Download background data of surface level...")
-    download_data(start_date, end_date, hour_step, area, f"{bg_save_path}/surface.grib",
+    download_data(start_date, end_date, hour_step, area, f"{data_save_path}/surface.grib",
                   data_format="grib", data_type="surface", overwrite=True)
 
     logger.info(f"Download background data of pressure level...")
-    download_data(start_date, end_date, hour_step, area, f"{bg_save_path}/pressure.grib",
+    download_data(start_date, end_date, hour_step, area, f"{data_save_path}/pressure.grib",
                   data_format="grib", data_type="pressure", overwrite=True)
 
     # logger.info(f"Download NearGOOS data...")

wrfrun/extension/goos_sst/__init__.py

@@ -4,11 +4,11 @@ wrfrun.extension.goos_sst
 
 This extension can help you create a GRIB file from ERA5 skin temperature (SKT) data and NEAR-GOOS sea surface temperature (SST) data.
 
-================================================== =============================================
-:doc:`goos_sst </api/extension.goos_sst.goos_sst>` Core functionality submodule.
-:doc:`res </api/extension.goos_sst.res>`           Resource files provided by this extension.
-:doc:`utils </api/extension.goos_sst.utils>`       Utility submodule used by the core submodule.
-================================================== =============================================
+============================================= =============================================
+:doc:`core </api/extension.goos_sst.core>`    Core functionality submodule.
+:doc:`res </api/extension.goos_sst.res>`      Resource files provided by this extension.
+:doc:`utils </api/extension.goos_sst.utils>`  Utility submodule used by the core submodule.
+============================================= =============================================
 
 Important Note
 **************

wrfrun/extension/goos_sst/core.py

@@ -16,7 +16,6 @@ from os.path import dirname
 import cfgrib as cf
 import numpy as np
 from pandas import to_datetime
-from seafog import goos_sst_find_data, goos_sst_parser
 from xarray import DataArray
 
 from .utils import create_sst_grib
@@ -40,6 +39,10 @@ def merge_era5_goos_sst_grib(surface_grib_path: str, save_path: str, sst_data_sa
                        Please check ``seafog.goos_sst_find_data`` for more information.
     :type resolution: str
     """
+    # lazy import seafog to fix libcurl error in readthedocs
+    # T^T
+    from seafog import goos_sst_find_data, goos_sst_parser
+
     dataset_list = cf.open_datasets(surface_grib_path)
 
     dataset = None

wrfrun/extension/goos_sst/res/Vtable.ERA_GOOS_SST

@@ -4,4 +4,4 @@ Code | Code |   1  |   2  | Name     |  Units   | Description
   34 |  1   |   0  |      | SST      | K        | Sea-Surface Temperature                  |
 -----+------+------+------+----------+----------+------------------------------------------+
 
-# Vtable setting for GRIB file created by wrfrun.extension.sst
+# Vtable setting for GRIB file created by wrfrun.extension.goos_sst

wrfrun/extension/goos_sst/res/__init__.py

@@ -1,3 +1,20 @@
+"""
+wrfrun.extension.goos_sst.res
+#############################
+
+Resource files provided by :doc:`/api/extension.goos_sst`.
+
+VTABLE_ERA_GOOS_SST
+*******************
+
+.. py:data:: VTABLE_ERA_GOOS_SST
+    :type: str
+    :value: Absolute file path.
+
+    Vtable file used to input the GRIB data created by :doc:`/api/extension.goos_sst` to WRF.
+
+"""
+
 from os.path import abspath, dirname
 
 

wrfrun/extension/goos_sst/utils.py

@@ -1,3 +1,15 @@
+"""
+wrfrun.extension.goos_sst.utils
+###############################
+
+Functions that are used by :doc:`/api/extension.goos_sst`.
+
+.. autosummary::
+    :toctree: generated/
+    
+    create_sst_grib
+"""
+
 from cfgrib.xarray_to_grib import to_grib
 from numpy.dtypes import DateTime64DType
 from pandas import to_datetime
@@ -8,12 +20,16 @@ from wrfrun.utils import logger
 
 def create_sst_grib(data: DataArray, save_path: str):
     """
-    Create a dataset and save it to a GRIB file.
-
-    Args:
-        data: DataArray data, it should at least contain three dimensions: ``["time", "latitude", "longitude"]``.
-        save_path: GRIB file path.
+    Write SST data to a GRIB file.
+    
+    This function creates GRIB file using ``cfgrib`` package.
+    While GRIB write support is experimental in ``cfgrib``,
+    this function may **FAIL TO CREATE GRIB FILE**.
 
+    :param data: ``xarray.DataArray``, which at least has three dimensions: ``["time", "latitude", "longitude"]``.
+    :type data: DataArray
+    :param save_path: Output GRIB file path.
+    :type save_path: str
     """
     # check the data's dimensions.
     for _dim in ["time", "longitude", "latitude"]:

wrfrun/extension/littler/__init__.py

@@ -1 +1,57 @@
-from .utils import *
+"""
+wrfrun.extension.littler
+########################
+
+This extension can help you manage observation data, and create ``LITTLE_R`` file for data assimilation.
+
+========================================= =============================
+:doc:`core </api/extension.littler.core>` Core functionality submodule.
+========================================= =============================
+
+What Can This Extension Do?
+***************************
+
+According to the `WRFDA Online Tutorial <https://www2.mmm.ucar.edu/wrf/users/wrfda/OnlineTutorial/Help/littler.html>`_,
+``LITTLE_R`` is an ASCII-based observation file format that is designed to be an intermediate format
+so that WRFDA might be able to assimilate as many observation types as possible in a universal manner.
+
+However, ``LITTLE_R`` is really hard to process elegantly from the point of view of Python.
+To help users create ``LITTLE_R`` file easily, this extension introduces :class:`LittleR <core.LittleR>`,
+and **Zipped Little R** file.
+
+:class:`LittleR <core.LittleR>` accepts observation datas, and can generate observation reports in proper format.
+Besides, it can save an observation report to a Zipped Little R file,
+so you can read the report later or process the observation data with other program.
+Please check :class:`LittleR <core.LittleR>` for more information.
+
+How To Use This Extension?
+**************************
+
+The code snap below shows you how to use this extension.
+
+.. code-block:: Python
+    :caption: main.py
+    
+    from wrfrun.extension.littler import LittleR
+
+
+    if __name__ == '__main__':
+        littler = LittleR()
+        littler.set_header(
+            longitude=120, latitude=60, fm="FM-19", elevation=0,
+            is_bogus=True, date="20250902070000"
+        )
+        # write to LITTLE_R file
+        with open("data/test", "w") as f:
+            f.write(str(littler))
+        # write to zlr file
+        littler.to_zlr("data/test.zlr")
+
+.. toctree::
+    :maxdepth: 1
+    :hidden:
+
+    core <extension.littler.core>
+"""
+
+from .core import *