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

wrfrun/model/wrf/exec_wrap.py

@@ -2,7 +2,7 @@
 wrfrun.model.wrf.exec_wrap
 ##########################
 
-Function wrapper of WPS / WRF :doc:`Executables </api/model.wrf.core>`.
+Function wrapper of WPS/WRF :doc:`Executables </api/model.wrf.core>`.
 
 .. autosummary::
     :toctree: generated/
@@ -18,8 +18,9 @@ Function wrapper of WPS / WRF :doc:`Executables </api/model.wrf.core>`.
 
 from typing import Optional, Union
 
-from wrfrun import WRFRUNConfig
-from .core import DFI, GeoGrid, MetGrid, NDown, Real, UnGrib, WRF
+from wrfrun.core import WRFRUN
+
+from .core import DFI, WRF, GeoGrid, MetGrid, NDown, Real, UnGrib
 
 
 def geogrid(geogrid_tbl_file: Union[str, None] = None):
@@ -28,7 +29,7 @@ def geogrid(geogrid_tbl_file: Union[str, None] = None):
 
     :param geogrid_tbl_file: Custom GEOGRID.TBL file path. Defaults to None.
     """
-    GeoGrid(geogrid_tbl_file, WRFRUNConfig.get_core_num())()
+    GeoGrid(geogrid_tbl_file, WRFRUN.config.get_core_num())()
 
 
 def ungrib(vtable_file: Union[str, None] = None, input_data_path: Optional[str] = None, prefix="FILE"):
@@ -47,7 +48,9 @@ def ungrib(vtable_file: Union[str, None] = None, input_data_path: Optional[str]
     UnGrib(vtable_file, input_data_path).set_ungrib_output_prefix(prefix)()
 
 
-def metgrid(geogrid_data_path: Optional[str] = None, ungrib_data_path: Optional[str] = None, fg_names: Union[str, list[str]] = "FILE"):
+def metgrid(
+    geogrid_data_path: Optional[str] = None, ungrib_data_path: Optional[str] = None, fg_names: Union[str, list[str]] = "FILE"
+):
     """
     Function interface for :class:`MetGrid <wrfrun.model.wrf.core.MetGrid>`.
 
@@ -60,9 +63,7 @@ def metgrid(geogrid_data_path: Optional[str] = None, ungrib_data_path: Optional[
     :param fg_names: ``fg_name`` of metgrid, a single prefix string or a string list.
     :type fg_names: str | list
     """
-    MetGrid(
-        geogrid_data_path, ungrib_data_path, WRFRUNConfig.get_core_num()
-    ).set_metgrid_fg_names(fg_names)()
+    MetGrid(geogrid_data_path, ungrib_data_path, WRFRUN.config.get_core_num()).set_metgrid_fg_names(fg_names)()
 
 
 def real(metgrid_data_path: Union[str, None] = None):
@@ -72,7 +73,7 @@ def real(metgrid_data_path: Union[str, None] = None):
     :param metgrid_data_path: Directory path of :class:`MetGrid <wrfrun.model.wrf.core.MetGrid>` outputs.
                               If is ``None``, try to use the workspace path or output path in the config file.
     """
-    Real(metgrid_data_path, WRFRUNConfig.get_core_num())()
+    Real(metgrid_data_path, WRFRUN.config.get_core_num())()
 
 
 def wrf(input_file_dir_path: Union[str, None] = None, restart_file_dir_path: Optional[str] = None, save_restarts=False):
@@ -83,7 +84,7 @@ def wrf(input_file_dir_path: Union[str, None] = None, restart_file_dir_path: Opt
     :param restart_file_dir_path: Directory path of restart files.
     :param save_restarts: If saving restart files. Defaults to False.
     """
-    WRF(input_file_dir_path, restart_file_dir_path, save_restarts, WRFRUNConfig.get_core_num())()
+    WRF(input_file_dir_path, restart_file_dir_path, save_restarts, WRFRUN.config.get_core_num())()
 
 
 def dfi(input_file_dir_path: Optional[str] = None, update_real_output=True):
@@ -95,7 +96,7 @@ def dfi(input_file_dir_path: Optional[str] = None, update_real_output=True):
     :param update_real_output: If update corresponding files in :class:`Real <wrfrun.model.wrf.core.Real>` outputs.
     :type update_real_output: bool
     """
-    DFI(input_file_dir_path, update_real_output, WRFRUNConfig.get_core_num())()
+    DFI(input_file_dir_path, update_real_output, WRFRUN.config.get_core_num())()
 
 
 def ndown(wrfout_file_path: str, real_output_dir_path: Optional[str] = None, update_namelist=True):
@@ -109,7 +110,7 @@ def ndown(wrfout_file_path: str, real_output_dir_path: Optional[str] = None, upd
     :param update_namelist: If update namelist settings for the final integral.
     :type update_namelist: bool
     """
-    NDown(wrfout_file_path, real_output_dir_path, update_namelist, WRFRUNConfig.get_core_num())()
+    NDown(wrfout_file_path, real_output_dir_path, update_namelist, WRFRUN.config.get_core_num())()
 
 
 __all__ = ["geogrid", "ungrib", "metgrid", "real", "wrf", "dfi", "ndown"]

wrfrun/model/wrf/geodata.py

@@ -1,3 +1,19 @@
+"""
+wrfrun.model.wrf.geodata
+########################
+
+Functions to read and write WPS geographical static data.
+
+.. autosummary::
+    :toctree: generated/
+
+    _get_data_type
+    _get_clip_area
+    parse_geographical_data_index
+    parse_geographical_data_file
+    read_geographical_static_data
+"""
+
 from os import listdir
 from os.path import exists
 from typing import OrderedDict, Union
@@ -5,44 +21,46 @@ from typing import OrderedDict, Union
 import numpy as np
 from xarray import DataArray
 
-from wrfrun.core import get_wrfrun_config
-from wrfrun.utils import logger
+from wrfrun.core import WRFRUN
+from wrfrun.log import logger
 
 # for example: 00001-00200.00201-00400
 DATA_NAME_TEMPLATE = "{}-{}.{}-{}"
 
 
 def _get_data_type(wordsize: int) -> type:
-    """Get data type based on wordsize value in index file.
-
-    Args:
-        wordsize (int): Wordsize in index file.
+    """
+    Get data type based on wordsize value in index file.
 
-    Returns:
-        type: NumPy value type.
+    :param wordsize: Wordsize in index file.
+    :type wordsize: int
+    :return: NumPy value type.
+    :rtype: type
     """
     # define map dict
-    map_dict = {
-        1: np.int8,
-        2: np.int16,
-        4: np.int32
-    }
+    map_dict = {1: np.int8, 2: np.int16, 4: np.int32}
 
     return map_dict[wordsize]
 
 
-def _get_clip_area(index_area: tuple[int, int, int, int], row_num: int, col_num: int, tile_x: int, tile_y: int) -> tuple[int, int, int, int]:
-    """Get clip area.
-
-    Args:
-        index_area (tuple[int, int, int, int]): Full area index.
-        row_num (int): Row number of the tile.
-        col_num (int): Column number of the file.
-        tile_x (int): X size of the tile.
-        tile_y (int): Y size of the tile.
-
-    Returns:
-        tuple[int, int, int, int]: Clip area.
+def _get_clip_area(
+    index_area: tuple[int, int, int, int], row_num: int, col_num: int, tile_x: int, tile_y: int
+) -> tuple[int, int, int, int]:
+    """
+    Get clip area.
+
+    :param index_area: Full area index.
+    :type index_area: tuple[int, int, int, int]
+    :param row_num: Row number of the tile.
+    :type row_num: int
+    :param col_num: Column number of the file.
+    :type col_num: int
+    :param tile_x: X size of the tile.
+    :type tile_x: int
+    :param tile_y: Y size of the tile.
+    :type tile_y: int
+    :return: Clip area.
+    :rtype: tuple[int, int, int, int]
     """
     # calculate tile area
     tile_area = (
@@ -64,25 +82,24 @@ def _get_clip_area(index_area: tuple[int, int, int, int], row_num: int, col_num:
 
 
 def parse_geographical_data_index(index_path: str) -> OrderedDict:
-    """Read geographical data index file.
-
-    Args:
-        index_path (str): Index file path.
+    """
+    Read geographical data index file.
 
-    Returns:
-        dict: Info stored in dict.
+    :param index_path: Index file path.
+    :type index_path: str
+    :return: Info stored in dict.
+    :rtype: OrderedDict[Any, Any]
     """
     # since the index file is very similar to fortran namelist file,
     # we can manually add "&index" and "/" and parse it as a namelist
     # temp file store path
-    WRFRUNConfig = get_wrfrun_config()
+    WRFRUNConfig = WRFRUN.config
     temp_file = f"{WRFRUNConfig.WRFRUN_TEMP_PATH}/geogrid_data.index"
     temp_file = WRFRUNConfig.parse_resource_uri(temp_file)
 
     # open file and add header and tail
     with open(index_path, "r") as _index_file:
         with open(temp_file, "w") as _temp_file:
-
             _temp_file.write("&index\n")
             _temp_file.write(_index_file.read())
             _temp_file.write("/")
@@ -93,20 +110,31 @@ def parse_geographical_data_index(index_path: str) -> OrderedDict:
     return WRFRUNConfig.get_namelist("geog_static_data")["index"]
 
 
-def parse_geographical_data_file(file_path: str, wordsize: int, endian: str, tile_shape: tuple[int, ...],
-                                 area: Union[tuple[int, ...], None] = None, miss_value: Union[int, float, None] = None) -> np.ndarray:
-    """Read geographical data file.
-
-    Args:
-        file_path (str): File path.
-        wordsize (int): How many bytes are used to store value in data file.
-        endian (str): "big" or "little".
-        tile_shape (tuple[int, ...]): The raw shape of the tile. Can be 2D or 3D.
-        area (Union[tuple[int, ...], None], optional): The range (x_start, x_stop, y_start, y_stop, ...) of data you want to read. Defaults to None.
-        miss_value (Union[int, float, None], optional): The value which represents NaN. Defaults to None.
-
-    Returns:
-        np.ndarray: NumPy array object.
+def parse_geographical_data_file(
+    file_path: str,
+    wordsize: int,
+    endian: str,
+    tile_shape: tuple[int, ...],
+    area: Union[tuple[int, ...], None] = None,
+    miss_value: Union[int, float, None] = None,
+) -> np.ndarray:
+    """
+    Read geographical data file.
+
+    :param file_path: File path.
+    :type file_path: str
+    :param wordsize: How many bytes are used to store value in data file.
+    :type wordsize: int
+    :param endian: "big" or "little".
+    :type endian: str
+    :param tile_shape: The raw shape of the tile. Can be 2D or 3D.
+    :type tile_shape: tuple[int, ...]
+    :param area: The range (x_start, x_stop, y_start, y_stop, ...) of data you want to read. Defaults to None.
+    :type area: Union[tuple[int, ...], None]
+    :param miss_value: The value which represents NaN. Defaults to None.
+    :type miss_value: Union[int, float, None]
+    :return: NumPy array object.
+    :rtype: ndarray[Any, Any]
     """
     # get data type
     data_type = _get_data_type(wordsize)
@@ -114,7 +142,7 @@ def parse_geographical_data_file(file_path: str, wordsize: int, endian: str, til
     # read data
     data = np.fromfile(file_path, dtype=data_type)
 
-    # swap byte if need
+    # swap byte if it needs
     if endian == "big":
         data = data.byteswap()
 
@@ -125,17 +153,16 @@ def parse_geographical_data_file(file_path: str, wordsize: int, endian: str, til
     if area:
         # check area
         if len(area) % 2 != 0:
-            logger.error(
-                f"The length of `area` must be even, but is {len(area)}")
+            logger.error(f"The length of `area` must be even, but is {len(area)}")
             exit(1)
 
         area_array = np.asarray(area).reshape(-1, 2)
         slice_index = tuple((slice(i[0], i[1]) for i in area_array))
 
         if len(slice_index) == 2:
-            slice_index += (slice(None), )
+            slice_index += (slice(None),)
 
-        data = data[slice_index[::-1]]    # type: ignore
+        data = data[slice_index[::-1]]  # type: ignore
 
     # fill nan
     if miss_value:
@@ -144,16 +171,20 @@ def parse_geographical_data_file(file_path: str, wordsize: int, endian: str, til
     return data
 
 
-def read_geographical_static_data(geog_data_folder_path: str, name: str, area: Union[tuple[float, float, float, float], None] = None) -> DataArray:
-    """Read WPS geographical static data
-
-    Args:
-        geog_data_folder_path (str): Data folder path.
-        name (str): Name that will be used to create DataArray.
-        area (Union[tuple[float, float, float, float], None]): Longitude and latitude area (lon_start, lon_stop, lat_start, lat_stop). Defaults to None.
-
-    Returns:
-        DataArray: DataArray object.
+def read_geographical_static_data(
+    geog_data_folder_path: str, name: str, area: Union[tuple[float, float, float, float], None] = None
+) -> DataArray:
+    """
+    Read WPS geographical static data
+
+    :param geog_data_folder_path: Data folder path.
+    :type geog_data_folder_path: str
+    :param name: Name that will be used to create DataArray.
+    :type name: str
+    :param area: Longitude and latitude area (lon_start, lon_stop, lat_start, lat_stop). Defaults to None.
+    :type area: Union[tuple[float, float, float, float], None]
+    :return: DataArray object.
+    :rtype: DataArray
     """
     # check if folder exists
     if not exists(geog_data_folder_path):
@@ -167,8 +198,7 @@ def read_geographical_static_data(geog_data_folder_path: str, name: str, area: U
     # extract info to read data
     # # check essential key
     if "wordsize" not in index_data:
-        logger.error(
-            f"Can't find key `wordsize` in index file, maybe it is corrupted.")
+        logger.error("Can't find key `wordsize` in index file, maybe it is corrupted.")
         exit(1)
     # # extract info
     wordsize = index_data["wordsize"]
@@ -194,15 +224,12 @@ def read_geographical_static_data(geog_data_folder_path: str, name: str, area: U
             int((area[3] - known_lat) // dy),
         )
         # check if negative value exists
-        if (
-            index_area[0] < 0 or
-            index_area[2] < 0
-        ):
-            logger.warning(f"Part of your area has exceeded data's area")
+        if index_area[0] < 0 or index_area[2] < 0:
+            logger.warning("Part of your area has exceeded data's area")
             # set negative value to 0
             index_area = tuple((i if i >= 0 else 0 for i in index_area))
     else:
-        logger.warning(f"You want to read all data, which may be very large")
+        logger.warning("You want to read all data, which may be very large")
         index_area = None
 
     # find the file we need to read
@@ -218,58 +245,49 @@ def read_geographical_static_data(geog_data_folder_path: str, name: str, area: U
         filenames = []
         # # generate filenames and clip area
         for row_num in range(tile_index_num[2], tile_index_num[3] + 1):
-
             _names = []
             for col_num in range(tile_index_num[0], tile_index_num[1] + 1):
-
                 _names.append(
                     [
                         DATA_NAME_TEMPLATE.format(
-                            str(col_num * tile_shape[-1] + 1).rjust(5, '0'),
-                            str((col_num + 1) * tile_shape[-1]).rjust(5, '0'),
-                            str(row_num * tile_shape[-2] + 1).rjust(5, '0'),
-                            str((row_num + 1) * tile_shape[-2]).rjust(5, '0'),
-                        ), _get_clip_area(index_area, row_num, col_num, tile_shape[-1], tile_shape[-2])   # type: ignore
+                            str(col_num * tile_shape[-1] + 1).rjust(5, "0"),
+                            str((col_num + 1) * tile_shape[-1]).rjust(5, "0"),
+                            str(row_num * tile_shape[-2] + 1).rjust(5, "0"),
+                            str((row_num + 1) * tile_shape[-2]).rjust(5, "0"),
+                        ),
+                        _get_clip_area(index_area, row_num, col_num, tile_shape[-1], tile_shape[-2]),  # type: ignore
                     ]
                 )
 
             filenames.append(_names)
     else:
-        raw_filenames = [x for x in listdir(
-            geog_data_folder_path) if x != "index"]
+        raw_filenames = [x for x in listdir(geog_data_folder_path) if x != "index"]
         raw_filenames.sort()
 
         # parse the last file to get row number and column number
         _last_filename = raw_filenames[-1]
-        total_col_num = int(_last_filename.split(
-            ".")[0].split("-")[1]) // tile_shape[-1]
-        total_row_num = int(_last_filename.split(
-            ".")[1].split("-")[1]) // tile_shape[-2]
+        total_col_num = int(_last_filename.split(".")[0].split("-")[1]) // tile_shape[-1]
+        total_row_num = int(_last_filename.split(".")[1].split("-")[1]) // tile_shape[-2]
 
         filenames = []
 
         for row_num in range(total_row_num):
-
             _names = []
             for col_num in range(total_col_num):
-
-                _names.append(
-                    [
-                        raw_filenames[row_num * total_col_num + col_num], None
-                    ]
-                )
+                _names.append([raw_filenames[row_num * total_col_num + col_num], None])
 
             filenames.append(_names)
 
     # read and concatenate
     array = []
     for _row in filenames:
-
         _array = []
         for _col in _row:
-
-            _array.append(parse_geographical_data_file(
-                f"{geog_data_folder_path}/{_col[0]}", wordsize, endian, tile_shape, _col[1], miss_value))
+            _array.append(
+                parse_geographical_data_file(
+                    f"{geog_data_folder_path}/{_col[0]}", wordsize, endian, tile_shape, _col[1], miss_value
+                )
+            )
 
         # concatenate _array
         array.append(np.concatenate(_array, axis=-1))
@@ -289,13 +307,11 @@ def read_geographical_static_data(geog_data_folder_path: str, name: str, area: U
     levels = np.arange(array.shape[-3])
 
     return DataArray(
-        name=name, data=array,
+        name=name,
+        data=array,
         dims=["levels", "latitude", "longitude"],
-        coords={
-            "longitude": longitude,
-            "latitude": latitude,
-            "levels": levels
-        }, attrs=index_data
+        coords={"longitude": longitude, "latitude": latitude, "levels": levels},
+        attrs=index_data,
     )
 
 

wrfrun/model/wrf/log.py

@@ -0,0 +1,103 @@
+"""
+wrfrun.model.wrf.log
+####################
+
+Functions to parse and clear WPS/WRF model logs.
+
+.. autosummary::
+    :toctree: generated/
+
+    get_wrf_simulated_seconds
+    clear_wrf_logs
+"""
+
+import subprocess
+from datetime import datetime
+from os import listdir
+from os.path import exists
+from shutil import move
+from typing import Optional
+
+from wrfrun.core import WRFRUN
+from wrfrun.log import logger
+from wrfrun.utils import check_path
+from wrfrun.workspace.wrf import get_wrf_workspace_path
+
+
+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 = WRFRUN.config.parse_resource_uri(f"{get_wrf_workspace_path('wrf')}/rsl.out.0000")
+    res = subprocess.run(["tail", "-n", "1", log_file_path], capture_output=True)
+    log_text = res.stdout.decode()
+
+    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
+
+    except ValueError:
+        seconds = -1
+
+    return seconds
+
+
+def clear_wrf_logs() -> None:
+    """
+    Collect unsaved WPS/WRF log files and save them to the corresponding
+    output directory of the ``Executable``.
+    """
+    WRFRUNConfig = WRFRUN.config
+
+    # wps
+    work_path = WRFRUNConfig.parse_resource_uri(get_wrf_workspace_path("wps"))
+
+    if exists(work_path):
+        log_files = [x for x in listdir(work_path) if x.endswith(".log")]
+
+        if len(log_files) > 0:
+            logger.warning("Found unprocessed log files of WPS model.")
+
+            log_save_path = f"{WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_OUTPUT_PATH)}/wps_unsaved_logs"
+            check_path(log_save_path)
+
+            for _file in log_files:
+                move(f"{work_path}/{_file}", f"{log_save_path}/{_file}")
+
+            logger.warning(f"Unprocessed log files of WPS model has been saved to {log_save_path}, check it")
+
+    # wrf
+    work_path = WRFRUNConfig.parse_resource_uri(get_wrf_workspace_path("wrf"))
+
+    if exists(work_path):
+        log_files = [x for x in listdir(work_path) if x.startswith("rsl.")]
+
+        if len(log_files) > 0:
+            logger.warning("Found unprocessed log files of WRF model.")
+
+            log_save_path = f"{WRFRUNConfig.parse_resource_uri(WRFRUNConfig.WRFRUN_OUTPUT_PATH)}/wrf_unsaved_logs"
+            check_path(log_save_path)
+
+            for _file in log_files:
+                move(f"{work_path}/{_file}", f"{log_save_path}/{_file}")
+
+            logger.warning(f"Unprocessed log files of WRF model has been saved to {log_save_path}, check it")
+
+
+__all__ = ["get_wrf_simulated_seconds", "clear_wrf_logs"]