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

wrfrun/data.py

@@ -1,12 +1,27 @@
+"""
+wrfrun.data
+###########
+
+.. autosummary::
+    :toctree: generated/
+
+    ERA5CONFIG
+    _check_variables_and_datasets
+    _check_pressure_level
+    find_era5_data
+    prepare_wps_input_data
+    download_data
+"""
+
 from datetime import datetime
 from os import makedirs
 from os.path import dirname, exists
-from typing import List, Tuple, Union
+from typing import List, Literal, Tuple, Union
 
 from pandas import date_range
 
-from .core import get_wrfrun_config
-from .utils import logger
+from .core import WRFRUN
+from .log import logger
 
 # from seafog import goos_sst_find_data
 
@@ -17,6 +32,68 @@ CDS_CLIENT = None
 class ERA5CONFIG:
     """
     A class to store parameters we will use to download ERA5 data from cdsapi.
+
+    .. py:attribute:: DATASET_ERA5_SINGLE_LEVEL
+        :type: str
+        :value: "reanalysis-era5-single-levels"
+
+        ERA5 dataset type.
+
+    .. py:attribute:: DATASET_ERA5_PRESSURE_LEVEL
+        :type: str
+        :value: "reanalysis-era5-pressure-levels"
+
+        ERA5 dataset type.
+
+    .. py:attribute:: TYPE_REANALYSIS
+        :type: str
+        :value: "reanalysis"
+
+        ERA5 product type.
+
+    .. py:attribute:: FORMAT_NETCDF
+        :type: str
+        :value: "netcdf"
+
+        ERA5 data format type.
+
+    .. py:attribute:: FORMAT_GRIB
+        :type: str
+        :value: "grib"
+
+        ERA5 data format type.
+
+    .. py:attribute:: DOWNLOAD_ZIP
+        :type: str
+        :value: "zip"
+
+        ERA5 download type.
+
+    .. py:attribute:: DOWNLOAD_UNZIP
+        :type: str
+        :value: "unarchived"
+
+        ERA5 download type.
+
+    .. py:attribute:: PRESSURE_LEVEL
+        :type: list
+
+        ERA5 available pressure levels.
+
+    .. py:attribute:: VARIABLE_*
+        :type: str
+
+        ERA5 variable names. Please check the code to see available values (too many to list here).
+
+    .. py:attribute:: NAME_MAP
+        :type: dict
+
+        Dictionary maps variable to its name in data.
+
+    .. py:attribute:: TYPE_MAP
+        :type: dict
+
+        Dictionary to distinguish between two types of data.
     """
 
     # dataset name
@@ -36,19 +113,43 @@ class ERA5CONFIG:
 
     # all level
     PRESSURE_LEVEL = [
-        '1', '2', '3',
-        '5', '7', '10',
-        '20', '30', '50',
-        '70', '100', '125',
-        '150', '175', '200',
-        '225', '250', '300',
-        '350', '400', '450',
-        '500', '550', '600',
-        '650', '700', '750',
-        '775', '800', '825',
-        '850', '875', '900',
-        '925', '950', '975',
-        '1000',
+        "1",
+        "2",
+        "3",
+        "5",
+        "7",
+        "10",
+        "20",
+        "30",
+        "50",
+        "70",
+        "100",
+        "125",
+        "150",
+        "175",
+        "200",
+        "225",
+        "250",
+        "300",
+        "350",
+        "400",
+        "450",
+        "500",
+        "550",
+        "600",
+        "650",
+        "700",
+        "750",
+        "775",
+        "800",
+        "825",
+        "850",
+        "875",
+        "900",
+        "925",
+        "950",
+        "975",
+        "1000",
     ]
 
     # variable name
@@ -102,7 +203,7 @@ class ERA5CONFIG:
         "volumetric_soil_water_layer_1": "swvl1",
         "volumetric_soil_water_layer_2": "swvl2",
         "volumetric_soil_water_layer_3": "swvl3",
-        "volumetric_soil_water_layer_4": "swvl4"
+        "volumetric_soil_water_layer_4": "swvl4",
     }
 
     # use a dict to distinguish between two types of data
@@ -125,8 +226,7 @@ class ERA5CONFIG:
             "volumetric_soil_water_layer_1",
             "volumetric_soil_water_layer_2",
             "volumetric_soil_water_layer_3",
-            "volumetric_soil_water_layer_4"
-
+            "volumetric_soil_water_layer_4",
         ),
         "reanalysis-era5-pressure-levels": (
             "specific_humidity",
@@ -134,20 +234,21 @@ class ERA5CONFIG:
             "v_component_of_wind",
             "geopotential",
             "relative_humidity",
-            "temperature"
-        )
+            "temperature",
+        ),
     }
 
 
 def _check_variables_and_datasets(variables: Union[str, Tuple[str, ...]], dataset: str) -> bool:
-    """Check if variables and datasets correspond
-
-    Args:
-        variables (str | tuple[str]): Variables type
-        dataset (str): Dataset type
-
-    Returns:
-        bool: If check passed, return True, else False
+    """
+    Check if variables and datasets correspond.
+
+    :param variables: Variables type.
+    :type variables: Union[str, Tuple[str, ...]]
+    :param dataset: Dataset type.
+    :type dataset: str
+    :return: If check passed, return ``True``, else ``False``.
+    :rtype: bool
     """
     if isinstance(variables, str):
         if variables in ERA5CONFIG.TYPE_MAP[dataset]:
@@ -163,28 +264,52 @@ def _check_variables_and_datasets(variables: Union[str, Tuple[str, ...]], datase
 
 
 def _check_pressure_level(pressure_level: Union[str, List[str]]) -> bool:
-    """Check pressure level
-
-    Args:
-        pressure_level (int | list[int]): A integer value or a list contains pressure values
+    """
+    Check pressure level.
 
-    Returns:
-        bool: If check passed, return True, else False
+    :param pressure_level: A integer value or a list contains pressure values.
+    :type pressure_level: Union[str, List[str]]
+    :return: If check passed, return ``True``, else ``False``.
+    :rtype: bool
     """
     valid_pressure_level = [
-        '1', '2', '3',
-        '5', '7', '10',
-        '20', '30', '50',
-        '70', '100', '125',
-        '150', '175', '200',
-        '225', '250', '300',
-        '350', '400', '450',
-        '500', '550', '600',
-        '650', '700', '750',
-        '775', '800', '825',
-        '850', '875', '900',
-        '925', '950', '975',
-        '1000',
+        "1",
+        "2",
+        "3",
+        "5",
+        "7",
+        "10",
+        "20",
+        "30",
+        "50",
+        "70",
+        "100",
+        "125",
+        "150",
+        "175",
+        "200",
+        "225",
+        "250",
+        "300",
+        "350",
+        "400",
+        "450",
+        "500",
+        "550",
+        "600",
+        "650",
+        "700",
+        "750",
+        "775",
+        "800",
+        "825",
+        "850",
+        "875",
+        "900",
+        "925",
+        "950",
+        "975",
+        "1000",
     ]
 
     for level in pressure_level:
@@ -194,32 +319,49 @@ def _check_pressure_level(pressure_level: Union[str, List[str]]) -> bool:
     return True
 
 
-def find_era5_data(date: Union[List[str], List[datetime]], area: Tuple[int, int, int, int], variables: Union[Tuple[str, ...], str], save_path: str,
-                   product_type=ERA5CONFIG.TYPE_REANALYSIS, data_format=ERA5CONFIG.FORMAT_NETCDF, dataset=ERA5CONFIG.DATASET_ERA5_SINGLE_LEVEL,
-                   download_format=ERA5CONFIG.DOWNLOAD_UNZIP, pressure_level: Union[int, List[int], str, List[str], None] = None, overwrite=False) -> str:
+def find_era5_data(
+    date: Union[List[str], List[datetime]],
+    area: Tuple[int, int, int, int],
+    variables: Union[Tuple[str, ...], str],
+    save_path: str,
+    product_type=ERA5CONFIG.TYPE_REANALYSIS,
+    data_format=ERA5CONFIG.FORMAT_NETCDF,
+    dataset=ERA5CONFIG.DATASET_ERA5_SINGLE_LEVEL,
+    download_format=ERA5CONFIG.DOWNLOAD_UNZIP,
+    pressure_level: Union[int, List[int], str, List[str], None] = None,
+    overwrite=False,
+) -> str:
     """
-    download era5 data
-    Args:
-        date: data date, string (for example, `2020-03-25 00:00`) or datetime object, UTC time
-        area: range of longitude and latitude, `[lon1, lon2, lat1, lat2]`
-        variables: variables, tuple of str or single string
-        save_path: save file path
-        product_type: product type, default is reanalysis
-        data_format: data format, default is netcdf
-        dataset: dataset type, default is reanalysis-era5-single-levels
-        download_format: zip or unarchived.
-        pressure_level: pressure levels.
-        overwrite (bool): If the data file exists, force to download it when `overwrite=True`
-
-    Returns: data path
-
+    Download data from ERA5.
+
+    :param date: List of date string or datetime object, UTC time.
+    :type date: Union[List[str], List[datetime]]
+    :param area: Range of longitude and latitude, ``[lon1, lon2, lat1, lat2]``.
+    :type area: Tuple[int, int, int, int]
+    :param variables: Variables to be download. Check :class:`ERA5CONFIG` for valid values.
+    :type variables: Union[Tuple[str, ...], str]
+    :param save_path: Data save path.
+    :type save_path: str
+    :param product_type: Product type. Check :class:`ERA5CONFIG` for valid values.
+    :type product_type: str
+    :param data_format: Data format. Check :class:`ERA5CONFIG` for valid values.
+    :type data_format: str
+    :param dataset: Dataset type. Check :class:`ERA5CONFIG` for valid values.
+    :type dataset: str
+    :param download_format: Download format. Check :class:`ERA5CONFIG` for valid values.
+    :type download_format: str
+    :param pressure_level: Pressure levels to be download. Check :class:`ERA5CONFIG` for valid values.
+    :type pressure_level: Union[int, List[int], str, List[str], None]
+    :param overwrite: If the data file exists, force to download it when ``overwrite==True``.
+    :type overwrite: bool
+    :return: Data path.
+    :rtype: str
     """
     global CDS_CLIENT
 
     # check variables and datasets
     if not _check_variables_and_datasets(variables, dataset):
-        logger.error(
-            f"Variables {variables} and dataset {dataset} doesn't correspond, check it")
+        logger.error(f"Variables {variables} and dataset {dataset} doesn't correspond, check it")
         exit(1)
 
     # check if we need to create directory
@@ -232,8 +374,10 @@ def find_era5_data(date: Union[List[str], List[datetime]], area: Tuple[int, int,
 
     # parse date
     if isinstance(date[0], str):
-        date = [datetime.strptime(_date, "%Y-%m-%d %H:%M")   # type: ignore
-                for _date in date]
+        date = [
+            datetime.strptime(_date, "%Y-%m-%d %H:%M")  # type: ignore
+            for _date in date
+        ]
     year = list(set(_date.strftime("%Y") for _date in date))  # type: ignore
     month = list(set(_date.strftime("%m") for _date in date))  # type: ignore
     day = list(set(_date.strftime("%d") for _date in date))  # type: ignore
@@ -251,15 +395,15 @@ def find_era5_data(date: Union[List[str], List[datetime]], area: Tuple[int, int,
 
     # create params dict
     params_dict = {
-        'product_type': product_type,
-        'data_format': data_format,
-        'download_format': download_format,
-        'variable': variables,
-        'year': year,
-        'month': month,
-        'day': day,
-        'time': time,
-        'area': area,
+        "product_type": product_type,
+        "data_format": data_format,
+        "download_format": download_format,
+        "variable": variables,
+        "year": year,
+        "month": month,
+        "day": day,
+        "time": time,
+        "area": area,
     }
 
     # check if we need to add pressure_level to params dict
@@ -269,12 +413,11 @@ def find_era5_data(date: Union[List[str], List[datetime]], area: Tuple[int, int,
             exit(1)
         # convert value to str
         if not isinstance(pressure_level, list):
-            pressure_level = [pressure_level]   # type: ignore
+            pressure_level = [pressure_level]  # type: ignore
         if not isinstance(pressure_level[0], str):  # type: ignore
-            pressure_level = [str(int(x))
-                              for x in pressure_level]  # type: ignore
+            pressure_level = [str(int(x)) for x in pressure_level]  # type: ignore
         # check
-        if _check_pressure_level(pressure_level):   # type: ignore
+        if _check_pressure_level(pressure_level):  # type: ignore
             params_dict["pressure_level"] = pressure_level
         else:
             logger.error("You have passed wrong pressure level to download data, check it")
@@ -285,6 +428,7 @@ def find_era5_data(date: Union[List[str], List[datetime]], area: Tuple[int, int,
 
     if CDS_CLIENT is None:
         import cdsapi
+
         CDS_CLIENT = cdsapi.Client()
 
     CDS_CLIENT.retrieve(dataset, params_dict, save_path)
@@ -293,12 +437,13 @@ def find_era5_data(date: Union[List[str], List[datetime]], area: Tuple[int, int,
 
 
 def prepare_wps_input_data(area: Tuple[int, int, int, int]):
-    """Download essential data for WPS.
+    """
+    Download essential data for WPS.
 
-    Args:
-        area (Tuple[int, int, int, int]): Range of longitude and latitude, `[lon1, lon2, lat1, lat2]`.
+    :param area: Range of longitude and latitude, ``[lon1, lon2, lat1, lat2]``.
+    :type area: Tuple[int, int, int, int]
     """
-    wrf_config = get_wrfrun_config().get_model_config("wrf")
+    wrf_config = 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"]
@@ -311,16 +456,32 @@ def prepare_wps_input_data(area: Tuple[int, int, int, int]):
     hour_step = wrf_config["time"]["input_data_interval"] // 3600
 
     # get data save path
-    data_save_path = get_wrfrun_config().get_input_data_path()
+    data_save_path = 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"{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"{data_save_path}/pressure.grib",
-                  data_format="grib", data_type="pressure", overwrite=True)
+    logger.info("Download background data of surface level...")
+    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("Download background data of pressure level...")
+    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...")
     # download_data(start_date, end_date, hour_step, area,
@@ -328,57 +489,61 @@ def prepare_wps_input_data(area: Tuple[int, int, int, int]):
 
 
 def download_data(
-        start_date: str,
-        end_date: str,
-        hour_step: int,
-        area: Tuple[int, int, int, int],
-        save_path: str,
-        data_format="nc",
-        data_type="pressure",
-        overwrite=False) -> str:
-    """Download essential data
-
-    Args:
-        start_date (str): Begin date, for example, "2022-05-19 12:00"
-        end_date (str): End date, for example, "2022-05-22 18:00"
-        hour_step (int): Hour step
-        area (tuple): Range of longitude and latitude, `[lon1, lon2, lat1, lat2]`
-        save_path (str): Data save path (era5 data) or data save folder path (goos sst)
-        data_format (str): Download data format, "nc" or "grib". Default is "nc"
-        data_type (str): Download data type, "pressure", "surface" or "goos". Default is "pressure".
-        overwrite (bool): If the data file exists, force to download it when `overwrite=True`
-
-    Returns:
-        str: Data save path
+    start_date: str,
+    end_date: str,
+    hour_step: int,
+    area: Tuple[int, int, int, int],
+    save_path: str,
+    data_format: Literal["nc", "grib"] = "nc",
+    data_type: Literal["pressure", "surface"] = "pressure",
+    overwrite=False,
+) -> str:
+    """
+    Download essential data from ERA5.
+
+    :param start_date: Begin date, for example, ``"2022-05-19 12:00"``.
+    :type start_date: str
+    :param end_date: End date, for example, ``"2022-05-22 18:00"``.
+    :type end_date: str
+    :param hour_step: Hour step.
+    :type hour_step: int
+    :param area: Range of longitude and latitude, ``[lon1, lon2, lat1, lat2]``.
+    :type area: Tuple[int, int, int, int]
+    :param save_path: Data save path.
+    :type save_path: str
+    :param data_format: Download data format. Default is ``"nc"``.
+    :type data_format: Literal["nc", "grib"]
+    :param data_type: Download data type. Default is ``"pressure"``.
+    :param overwrite: If the data file exists, force to download it when ``overwrite==True``.
+    :return: Data path.
+    :rtype: str
     """
     # generate date list
-    date_list = date_range(
-        start_date, end_date, freq=f"{hour_step}H"
-    ).strftime("%Y-%m-%d %H:%M").to_list()
+    date_list = date_range(start_date, end_date, freq=f"{hour_step}H").strftime("%Y-%m-%d %H:%M").to_list()
 
     # check format
     if data_format == "nc":
-        data_format = ERA5CONFIG.FORMAT_NETCDF
+        _data_format = ERA5CONFIG.FORMAT_NETCDF
     elif data_format == "grib":
-        data_format = ERA5CONFIG.FORMAT_GRIB
+        _data_format = ERA5CONFIG.FORMAT_GRIB
     else:
-        logger.error(f"Wrong data format: {data_format}")
+        logger.error(f"Wrong data format: {_data_format}")
         raise KeyError
 
     # check data type
     if data_type == "pressure":
-        data_type = ERA5CONFIG.DATASET_ERA5_PRESSURE_LEVEL
+        _data_type = ERA5CONFIG.DATASET_ERA5_PRESSURE_LEVEL
         variables = (
             ERA5CONFIG.VARIABLE_GEOPOTENTIAL,
             ERA5CONFIG.VARIABLE_RELATIVE_HUMIDITY,
             ERA5CONFIG.VARIABLE_SPECIFIC_HUMIDITY,
             ERA5CONFIG.VARIABLE_TEMPERATURE,
             ERA5CONFIG.VARIABLE_U_WIND,
-            ERA5CONFIG.VARIABLE_V_WIND
+            ERA5CONFIG.VARIABLE_V_WIND,
         )
         pressure_level = ERA5CONFIG.PRESSURE_LEVEL
     elif data_type == "surface":
-        data_type = ERA5CONFIG.DATASET_ERA5_SINGLE_LEVEL
+        _data_type = ERA5CONFIG.DATASET_ERA5_SINGLE_LEVEL
         variables = (
             ERA5CONFIG.VARIABLE_SURFACE_PRESSURE,
             ERA5CONFIG.VARIABLE_MEAN_SEA_LEVEL_PRESSURE,
@@ -397,25 +562,25 @@ def download_data(
             ERA5CONFIG.VARIABLE_VOLUMETRIC_SOIL_WATER_LAYER_3,
             ERA5CONFIG.VARIABLE_VOLUMETRIC_SOIL_WATER_LAYER_4,
             ERA5CONFIG.VARIABLE_SNOW_DEPTH,
-            ERA5CONFIG.VARIABLE_SNOW_DENSITY
+            ERA5CONFIG.VARIABLE_SNOW_DENSITY,
         )
         pressure_level = None
-    elif data_type == "goos":
-        logger.warning(f"NEAR-GOOS SST data hasn't been supported yet")
-        # download sst data
-        # for _date in date_list:
-        #     _ = goos_sst_find_data(_date, save_path=save_path)
-
-        return ""
     else:
-        logger.error(f"Wrong data type: {data_type}")
+        logger.error(f"Wrong data type: {_data_type}")
         raise KeyError
 
     # download data
-    return find_era5_data(date=date_list, area=area, variables=variables,   # type: ignore
-                          save_path=save_path, data_format=data_format,
-                          dataset=data_type, pressure_level=pressure_level,  # type: ignore
-                          download_format=ERA5CONFIG.DOWNLOAD_UNZIP, overwrite=overwrite)
+    return find_era5_data(
+        date=date_list,
+        area=area,
+        variables=variables,
+        save_path=save_path,
+        data_format=_data_format,
+        dataset=_data_type,
+        pressure_level=pressure_level,
+        download_format=ERA5CONFIG.DOWNLOAD_UNZIP,
+        overwrite=overwrite,
+    )
 
 
 __all__ = ["find_era5_data", "ERA5CONFIG", "download_data", "prepare_wps_input_data"]

wrfrun/extension/goos_sst/__init__.py

@@ -2,7 +2,7 @@
 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.
+Extension help user create a GRIB file from ERA5 skin temperature (SKT) data and NEAR-GOOS sea surface temperature (SST) data.
 
 ============================================= =============================================
 :doc:`core </api/extension.goos_sst.core>`    Core functionality submodule.
@@ -50,7 +50,7 @@ The code snap below shows you how to use this extension.
     if __name__ == '__main__':
         era5_data_path = "data/ear5-reanalysis-data.grib"
         merge_era5_goos_sst_grib(era5_data_path, "data/near-goos-data.grib")
-        
+
 Please check :func:`core.merge_era5_goos_sst_grib` for more information.
 
 .. toctree::

wrfrun/extension/goos_sst/core.py

@@ -7,19 +7,19 @@ Implementation of ``extension.goos_sst``'s core functionality.
 
 .. autosummary::
     :toctree: generated/
-    
+
     merge_era5_goos_sst_grib
 """
 
 from os.path import dirname
 
-import cfgrib as cf
 import numpy as np
 from pandas import to_datetime
 from xarray import DataArray
 
+from wrfrun.log import logger
+
 from .utils import create_sst_grib
-from wrfrun.utils import logger
 
 
 def merge_era5_goos_sst_grib(surface_grib_path: str, save_path: str, sst_data_save_path: str | None = None, resolution="low"):
@@ -41,6 +41,7 @@ def merge_era5_goos_sst_grib(surface_grib_path: str, save_path: str, sst_data_sa
     """
     # lazy import seafog to fix libcurl error in readthedocs
     # T^T
+    import cfgrib as cf
     from seafog import goos_sst_find_data, goos_sst_parser
 
     dataset_list = cf.open_datasets(surface_grib_path)
@@ -55,10 +56,10 @@ def merge_era5_goos_sst_grib(surface_grib_path: str, save_path: str, sst_data_sa
         logger.error(f"'skt' data not found in {surface_grib_path}")
         raise ValueError
 
-    skt: DataArray = dataset["skt"]     # type: ignore
+    skt: DataArray = dataset["skt"]  # type: ignore
 
-    longitude_start, longitude_end = skt["longitude"][0].data, skt["longitude"][-1].data    # type: ignore
-    latitude_start, latitude_end = skt["latitude"][-1].data, skt["latitude"][0].data    # type: ignore
+    longitude_start, longitude_end = skt["longitude"][0].data, skt["longitude"][-1].data  # type: ignore
+    latitude_start, latitude_end = skt["latitude"][-1].data, skt["latitude"][0].data  # type: ignore
 
     data_time = skt["time"].to_numpy()  # type: ignore
     data_time = to_datetime(data_time).strftime("%Y-%m-%d %H:%M")
@@ -95,14 +96,8 @@ def merge_era5_goos_sst_grib(surface_grib_path: str, save_path: str, sst_data_sa
         name="sst",
         data=sst,
         dims=["time", "latitude", "longitude"],
-        coords={
-            "time": skt["time"],
-            "latitude": skt["latitude"],
-            "longitude": skt["longitude"]
-        }, attrs={
-            "units": "K",
-            "long_name": "Sea surface temperature"
-        }
+        coords={"time": skt["time"], "latitude": skt["latitude"], "longitude": skt["longitude"]},
+        attrs={"units": "K", "long_name": "Sea surface temperature"},
     )
 
     create_sst_grib(sst, save_path=save_path)