wrfrun 0.1.7__py3-none-any.whl

wrfrun/model/wrf/vtable.py

@@ -0,0 +1,65 @@
+from dataclasses import dataclass
+
+from wrfrun.core import WRFRUNConfig
+
+
+VTABLE_URI = ":WRFRUN_VTABLE:"
+
+
+@dataclass
+class VtableFiles:
+    """
+    Represent WPS Vtable files (v4.5).
+    With the ":/" we can tell from user custom Vtable files.
+    """
+    AFWAICE = f"{VTABLE_URI}/Vtable.AFWAICE"
+    AGRMETSNOW = f"{VTABLE_URI}/Vtable.AGRMETSNOW"
+    AGRMETSOIL = f"{VTABLE_URI}/Vtable.AGRMETSOIL"
+    AGRMETSOIL2 = f"{VTABLE_URI}/Vtable.AGRMETSOIL2"
+    AGRWRF = f"{VTABLE_URI}/Vtable.AGRWRF"
+    ARW = f"{VTABLE_URI}/Vtable.ARW.UPP"
+    ARWP = f"{VTABLE_URI}/Vtable.ARWp.UPP"
+    AVN0P5WRF = f"{VTABLE_URI}/Vtable.AVN0P5WRF"
+    AWIP = f"{VTABLE_URI}/Vtable.AWIP"
+    CFSR = f"{VTABLE_URI}/Vtable.CFSR"
+    CFSR_MEAN = f"{VTABLE_URI}/Vtable.CFSR_mean"
+    ECMWF = f"{VTABLE_URI}/Vtable.ECMWF"
+    ECMWF_SIGMA = f"{VTABLE_URI}/Vtable.ECMWF_sigma"
+    ERA_ML = f"{VTABLE_URI}/Vtable.ERA-interim.ml"
+    ERA_PL = f"{VTABLE_URI}/Vtable.ERA-interim.pl"
+    GFDL = f"{VTABLE_URI}/Vtable.GFDL"
+    GFS = f"{VTABLE_URI}/Vtable.GFS"
+    GFS_OZONE = f"{VTABLE_URI}/Vtable.GFS_OZONE"
+    GFSENS = f"{VTABLE_URI}/Vtable.GFSENS"
+    GODAS = f"{VTABLE_URI}/Vtable.GODAS"
+    GSM = f"{VTABLE_URI}/Vtable.GSM"
+    ICONM = f"{VTABLE_URI}/Vtable.ICONm"
+    ICONP = f"{VTABLE_URI}/Vtable.ICONp"
+    JMAGSM = f"{VTABLE_URI}/Vtable.JMAGSM"
+    NAM = f"{VTABLE_URI}/Vtable.NAM"
+    NARR = f"{VTABLE_URI}/Vtable.NARR"
+    NAVY_SST = f"{VTABLE_URI}/Vtable.NavySST"
+    NCEP2 = f"{VTABLE_URI}/Vtable.NCEP2"
+    NNRP = f"{VTABLE_URI}/Vtable.NNRP"
+    NOGAPS = f"{VTABLE_URI}/Vtable.NOGAPS"
+    NOGAPS_GFS_SOIL = f"{VTABLE_URI}/Vtable.NOGAPS_needs_GFS_soil"
+    RAP_HYBRID_NCEP = f"{VTABLE_URI}/Vtable.RAP.hybrid.ncep"
+    RAP_PRESSURE_NCEP = f"{VTABLE_URI}/Vtable.RAP.pressure.ncep"
+    RAP_SIGMA_GSD = f"{VTABLE_URI}/Vtable.RAP.sigma.gsd"
+    RAPHRRR = f"{VTABLE_URI}/Vtable.raphrrr"
+    RUCB = f"{VTABLE_URI}/Vtable.RUCb"
+    RUCP = f"{VTABLE_URI}/Vtable.RUCp"
+    SREF = f"{VTABLE_URI}/Vtable.SREF"
+    SST = f"{VTABLE_URI}/Vtable.SST"
+    TCRP = f"{VTABLE_URI}/Vtable.TCRP"
+    UKMO_END_GAME = f"{VTABLE_URI}/Vtable.UKMO_ENDGame"
+    UKMO_LANDSEA = f"{VTABLE_URI}/Vtable.UKMO_LANDSEA"
+    UKMO_NO_HEIGHTS = f"{VTABLE_URI}/Vtable.UKMO_no_heights"
+
+
+# register uri
+if not WRFRUNConfig.check_resource_uri(VTABLE_URI):
+    WRFRUNConfig.register_resource_uri(VTABLE_URI, f"{WRFRUNConfig.WPS_WORK_PATH}/ungrib/Variable_Tables")
+
+
+__all__ = ["VtableFiles"]

wrfrun/pbs.py

@@ -0,0 +1,86 @@
+from os import environ
+from os.path import exists, dirname, abspath
+
+from .core import WRFRUNConfig
+from .res import TASKSYS_PBS_TEMPLATE
+from .utils import logger
+
+
+def prepare_pbs_script(main_file_path: str):
+    """Prepare the bash script to be submitted to PBS work system.
+
+    Args:
+        main_file_path (str): The path of main Python file.
+    """
+    # check main file
+    if not exists(main_file_path):
+        logger.error(f"Wrong path of main Python file: {main_file_path}")
+        raise FileNotFoundError
+
+    # get absolute path of parent directory
+    dir_path = abspath(dirname(main_file_path))
+
+    # read log path and PBS setting from config
+    log_path = WRFRUNConfig.get_log_path()
+    PBS_setting = WRFRUNConfig.get_job_scheduler_config()
+
+    # set PBS log path
+    stdout_log_path = f"{log_path}/PBS.log"
+    stderr_log_path = f"{log_path}/PBS.log"
+
+    # set environment parameter
+    env_settings = ''
+    if len(PBS_setting["env_settings"]) != 0:
+        for key in PBS_setting["env_settings"]:
+            env_settings += f"{key}={PBS_setting['env_settings'][key]}\n"
+
+    # set command
+    exec_cmd = f"{PBS_setting['python_interpreter']} {main_file_path}"
+
+    # read template and write to file
+    pbs_template_path = WRFRUNConfig.parse_resource_uri(TASKSYS_PBS_TEMPLATE)
+    with open(f"{dir_path}/run.sh", "w") as f:
+
+        with open(pbs_template_path, "r") as f_template:
+            template = f_template.read()
+
+        template = template.format(
+            STDOUT_LOG_PATH=stdout_log_path,
+            STDERR_LOG_PATH=stderr_log_path,
+            NODE_NUM=PBS_setting["node_num"],
+            CORE_NUM=PBS_setting["core_num"],
+            ENV_SETTINGS=env_settings,
+            WORK_PATH=dir_path,
+            WORK_COMMAND=exec_cmd
+        )
+
+        f.write(template)
+
+    logger.info(
+        f"PBS script has been generated and write to file {dir_path}/run.sh. Check it and submit it to PBS system to run wrfrun.")
+
+
+def get_core_num() -> int:
+    """Read core num from config.
+
+    Returns:
+        int: Core number.
+    """
+    return WRFRUNConfig["core_num"]
+
+
+def in_pbs() -> bool:
+    """Check if we're in a PBS work.
+
+    Returns:
+        bool: True if we're in, False is we aren't.
+    """
+    # check if we're a PBS task
+    if "PBS_ENVIRONMENT" in environ:
+        # we're in a PBS task
+        return True
+    else:
+        return False
+
+
+__all__ = ["prepare_pbs_script", "get_core_num", "in_pbs"]

wrfrun/plot/__init__.py

@@ -0,0 +1 @@
+from .wps import *

wrfrun/plot/wps.py

@@ -0,0 +1,188 @@
+from typing import Tuple
+
+from cartopy import crs
+from cartopy.mpl.geoaxes import GeoAxes
+from matplotlib.collections import QuadMesh
+from matplotlib.colors import ListedColormap
+from matplotlib.figure import Figure
+from netCDF4 import Dataset  # type: ignore
+from wrf import getvar, latlon_coords, get_cartopy, to_np
+from xarray import DataArray
+
+from wrfrun.utils import logger
+
+# Here defines the colormap for landuse and soiltype
+LANDUSE_CMAP = ListedColormap([
+    [0, .4, 0],         # 1 Evergreen Needleleaf Forest
+    [0, .4, .2],        # 2 Evergreen Broadleaf Forest
+    [.2, .8, .2],       # 3 Deciduous Needleleaf Forest
+    [.2, .8, .4],       # 4 Deciduous Broadleaf Forest
+    [.2, .6, .2],       # 5 Mixed Forests
+    [.3, .7, 0],        # 6 Closed Shrublands
+    [.82, .41, .12],    # 7 Open Shurblands
+    [.74, .71, .41],    # 8 Woody Savannas
+    [1, .84, .0],       # 9 Savannas
+    [0, 1, 0],          # 10 Grasslands
+    [0, 1, 1],          # 11 Permanant Wetlands
+    [1, 1, 0],          # 12 Croplands
+    [1, 0, 0],          # 13 Urban and Built-up
+    [.7, .9, .3],       # 14 Cropland/Natual Vegation Mosaic
+    [1, 1, 1],          # 15 Snow and Ice
+    [.914, .914, .7],   # 16 Barren or Sparsely Vegetated
+    [.5, .7, 1],        # 17 Water (like oceans)
+    [.86, .08, .23],    # 18 Wooded Tundra
+    [.97, .5, .31],     # 19 Mixed Tundra
+    [.91, .59, .48],    # 20 Barren Tundra
+    [0, 0, .88],        # 21 Lake
+])
+
+LANDUSE_LABELS = [
+    'Evergreen Needleleaf Forest',
+    'Evergreen Broadleaf Forest',
+    'Deciduous Needleleaf Forest',
+    'Deciduous Broadleaf Forest',
+    'Mixed Forests',
+    'Closed Shrublands',
+    'Open Shrublands',
+    'Woody Savannas',
+    'Savannas',
+    'Grasslands',
+    'Permanent Wetlands',
+    'Croplands',
+    'Urban and Built-Up',
+    'Cropland/Natural Vegetation Mosaic',
+    'Snow and Ice',
+    'Barren or Sparsely Vegetated',
+    'Water',
+    'Wooded Tundra',
+    'Mixed Tundra',
+    'Barren Tundra',
+    'Lake',
+]
+
+
+# Here defines the colormap and labels for soil type
+SOILTYPE_CMAP = ListedColormap([
+    [0, .4, 0],         # 1 Sand
+    [0, .4, .2],        # 2 Loamy Sand
+    [.2, .8, .2],       # 3 Sandy Loam
+    [.2, .8, .4],       # 4 Silt Loam
+    [.2, .6, .2],       # 5 Silt
+    [.3, .7, 0],        # 6 Loam
+    [.82, .41, .12],    # 7 Sandy Clay Loam
+    [.74, .71, .41],    # 8 Silty Clay Loam
+    [1, .84, .0],       # 9 Clay Loam
+    [0, 1, 0],          # 10 Sandy Clay
+    [0, 1, 1],          # 11 Silty Clay
+    [1, 1, 0],          # 12 Clay
+    [1, 0, 0],          # 13 Organic Material
+    [.7, .9, .3],       # 14 Water
+    [1, 1, 1],          # 15 Bedrock
+    [.914, .914, .7]    # 16 Other
+])
+
+SOILTYPE_LABELS = [
+    "Sand",
+    "Loamy Sand",
+    "Sandy Loam",
+    "Silt Loam",
+    "Silt",
+    "Loam",
+    "Sandy Clay Loam",
+    "Silty Clay Loam",
+    "Clay Loam",
+    "Sandy Clay",
+    "Silty Clay",
+    "Clay",
+    "Organic Material",
+    "Water",
+    "Bedrock",
+    "Other"
+]
+
+
+def get_cmap_ticks(name: str) -> tuple:
+    """Get corresponding colormap, labels and ticks.
+
+    Args:
+        name (str): Field name.
+
+    Returns:
+        tuple: (colormap, labels, ticks)
+    """
+    # name map
+    cmap_ticks_map = {
+        "LU_INDEX": {
+            "cmap": LANDUSE_CMAP,
+            "labels": LANDUSE_LABELS,
+        },
+        "SOILCTOP": {
+            "cmap": SOILTYPE_CMAP,
+            "labels": SOILTYPE_LABELS,
+        },
+        "SOILCBOT": {
+            "cmap": SOILTYPE_CMAP,
+            "labels": SOILTYPE_LABELS,
+        },
+    }
+
+    # take out colormap and labels
+    if name not in cmap_ticks_map:
+        logger.error(f"Can't found colormap and labels for {name}")
+        raise ValueError
+
+    cmap = cmap_ticks_map[name]["cmap"]
+    labels = cmap_ticks_map[name]["labels"]
+
+    # generate ticks
+    ticks = [x + 1.5 for x in range(len(labels))]
+
+    return cmap, labels, ticks
+
+
+def draw_geogrid(data_path: str, field: str, fig: Figure, nrow: int, ncol: int, index: int) -> Tuple[GeoAxes, QuadMesh]:
+    """Draw specific data in geo_em*
+
+    Args:
+        data_path (str): data path.
+        field (str): Which field you want to draw.
+        fig (Figure): Matplotlib figure object.
+        nrow (int): Row number of axes.
+        ncol (int): Column number of axes.
+        index (int): Index of axes.
+    """
+    # take out data
+    data: DataArray = getvar(Dataset(data_path), field)     # type: ignore
+
+    # soiltype data is different from landuse data
+    if field in ["SOILCTOP", "SOILCBOT"]:
+        data = data.argmax(dim="soil_cat", keep_attrs=True)     # type: ignore
+        # get data attrs
+        data_attrs = data.attrs
+        data = data + 1
+        data.attrs = data_attrs
+
+    # get latitude and longitude
+    lats, lons = latlon_coords(data)
+
+    # get projection
+    proj = get_cartopy(data)
+
+    # add axes
+    ax = fig.add_subplot(nrow, ncol, index, projection=proj)
+
+    # get colormap, labels and ticks
+    cmap, labels, ticks = get_cmap_ticks(field)
+
+    # convert data to numpy
+    data, lats, lons = to_np(data), to_np(lats), to_np(lons)    # type: ignore
+
+    # plot
+    im = ax.pcolormesh(
+        lons, lats, data, vmin=1, vmax=len(labels) + 1, cmap=cmap, transform=crs.PlateCarree()
+    )
+
+    return ax, im   # type: ignore
+
+
+__all__ = ["draw_geogrid", "get_cmap_ticks"]

wrfrun/res/__init__.py

@@ -0,0 +1,22 @@
+from os.path import abspath, dirname
+
+from wrfrun.core import WRFRUNConfig
+
+
+RES_PATH = abspath(dirname(__file__))
+WRFRUNConfig.register_resource_uri(WRFRUNConfig.WRFRUN_RESOURCE_PATH, RES_PATH)
+
+CONFIG_TOML_TEMPLATE = ":WRFRUN_RESOURCE_PATH:/config.toml.template"
+EXT_NCL_PLOT_SCRIPT = ":WRFRUN_RESOURCE_PATH:/extension/plotgrids.ncl"
+TASKSYS_PBS_TEMPLATE = ":WRFRUN_RESOURCE_PATH:/job_scheduler/pbs.template"
+TASKSYS_SLURM_TEMPLATE = ":WRFRUN_RESOURCE_PATH:/job_scheduler/slurm.template"
+NAMELIST_WRFDA = ":WRFRUN_RESOURCE_PATH:/namelist/namelist.input.da_wrfvar.template"
+NAMELIST_DFI = ":WRFRUN_RESOURCE_PATH:/namelist/namelist.input.dfi.template"
+NAMELIST_REAL = ":WRFRUN_RESOURCE_PATH:/namelist/namelist.input.real.template"
+NAMELIST_WRF = ":WRFRUN_RESOURCE_PATH:/namelist/namelist.input.wrf.template"
+NAMELIST_WPS = ":WRFRUN_RESOURCE_PATH:/namelist/namelist.wps.template"
+
+WRFRUNConfig.set_config_template_path(CONFIG_TOML_TEMPLATE)
+
+
+__all__ = ["CONFIG_TOML_TEMPLATE", "EXT_NCL_PLOT_SCRIPT", "TASKSYS_PBS_TEMPLATE", "TASKSYS_SLURM_TEMPLATE", "NAMELIST_WRFDA", "NAMELIST_DFI", "NAMELIST_REAL", "NAMELIST_WRF", "NAMELIST_WPS"]

wrfrun/res/config.toml.template

@@ -0,0 +1,136 @@
+# Config template of wrfrun.
+# Path of the directory which contains input data.
+input_data_path = ""
+
+# Path of the directory to store all outputs.
+output_path = "./outputs"
+log_path = "./logs"
+
+# wrfrun can launch a socket server during NWP execution to report simulation progress.
+# To enable the server, you need to configure the IP address and port on which it will listen to.
+server_host = "localhost"
+server_port = 54321
+
+# How many cores you will use.
+# Note that if you use a job scheduler (like PBS), this value means the number of cores each node you use.
+core_num = 36
+
+
+[job_scheduler]
+# Job scheduler settings.
+# How many nodes you will use.
+node_num = 1
+
+# Custom environment settings
+env_settings = {}
+
+# Path of the python interpreter that will be used to run wrfrun.
+# You can also give its name only.
+python_interpreter = "/usr/bin/python3"     # or just "python3"
+
+
+[model]
+# Model debug level
+debug_level = 100
+
+# ################################################### Only settings above is necessary ###########################################
+
+
+# ####################################### You can give more settings about the NWP you will use ##################################
+
+[model.wrf]
+# Config for WRF
+# WRF model path
+wps_path = '/path/to/your/WPS/folder'
+wrf_path = '/path/to/your/WRF/folder'
+# WRFDA is optional.
+wrfda_path = ''
+
+# static geographic data path
+geog_data_path = '/path/to/your/geog/data'
+
+# Your can give your custom namelist files here.
+# The value in it will overwrite the default value in the namelist template file.
+user_wps_namelist = ''
+user_real_namelist = ''
+user_wrf_namelist = ''
+user_wrfda_namelist = ''
+
+# If you make a restart run?
+restart_mode = false
+
+
+[model.wrf.time]
+# Advance time config for WRF
+# Set the start and end date. It will be used for all domains.
+# You can also provide all the dates as a list, with each date for the corresponding domain.
+start_date = 2021-03-24T12:00:00Z   # or [2021-03-24T12:00:00Z, 2021-03-24T12:00:00Z]
+end_date = 2021-03-26T00:00:00Z     # or [2021-03-26T00:00:00Z, 2021-03-24T12:00:00Z]
+
+# Set input data time interval. Unit: seconds
+input_data_interval = 10800
+
+# Set output data time interval. Unit: minutes
+output_data_interval = 180
+
+# Note that there are various reasons which could crash wrf,
+# and in most cases you can deal with them by decrease time step.
+# Unit: seconds
+time_step = 120
+
+# Time ratio to the first domain for each domain
+parent_time_step_ratio = [1, 3, 4]
+
+# Time interval to write restart files.
+# This help you can restart WRF after it stops.
+# By default, it equals to output_data_interval. Unit: minutes.
+restart_interval = -1
+
+
+[model.wrf.domain]
+# Advance domain config for WRF.
+# Set domain number.
+domain_num = 3
+
+# It's very hard to process wrf domain settings because it's related to various settings, so I keep it.
+# The resolution ratio to the first domain.
+parent_grid_ratio = [1, 3, 9]
+
+# Index of the start point.
+i_parent_start = [1, 17, 72]
+j_parent_start = [1, 17, 36]
+
+# Number of points.
+e_we = [120, 250, 1198]
+e_sn = [120, 220, 1297]
+
+# Resolution of the first domain.
+dx = 9000
+dy = 9000
+
+# Projection.
+map_proj = { name = 'lambert', truelat1 = 34.0, truelat2 = 40.0}
+
+# Central point of the first area.
+ref_lat = 37.0
+ref_lon = 120.5
+stand_lon = 120.5
+
+
+[model.wrf.scheme]
+# Advance physics scheme config for WRF.
+# To loop up the nickname for all physics schemes, please see: https://wrfrun.syize.cn.
+# Option contains many other settings related to the scheme.
+# Sometimes some option can only be used for a specific scheme.
+# You can check it in online namelist variables: https://www2.mmm.ucar.edu/wrf/users/wrf_users_guide/build/html/namelist_variables.html
+# You can set option with its `wrf name` and its `wrf value`.
+# For example, `ghg_input=1` works with rrtm scheme. If you want to set `ghg_input=1` when using rrtm, set option: {"ghg_input": 1}
+# However, sometimes some options work with various schemes, and some options themselves are schem.
+# Use this carefully.
+# You can set multiple keys in option.
+long_wave_scheme = { name = "rrtm", option = {} }
+short_wave_scheme = { name = "rrtmg", option = {} }
+cumulus_scheme = { name = "kf", option = {} }
+pbl_scheme = { name = "ysu", option = { ysu_topdown_pblmix = 1} }
+land_surface_scheme = { name = "noah", option = {} }
+surface_layer_scheme = { name = "mm5", option = {} }