from __future__ import (absolute_import, division, print_function) from .util import (get_iterable, is_standard_wrf_var, extract_vars, viewkeys, get_id) from .g_cape import (get_2dcape, get_3dcape, get_cape2d_only, get_cin2d_only, get_lcl, get_lfc, get_3dcape_only, get_3dcin_only) from .g_ctt import get_ctt from .g_dbz import get_dbz, get_max_dbz from .g_dewpoint import get_dp, get_dp_2m from .g_geoht import get_geopt, get_height, get_stag_geopt, get_stag_height from .g_helicity import get_srh, get_uh from .g_latlon import get_lat, get_lon from .g_omega import get_omega from .g_pressure import get_pressure, get_pressure_hpa from .g_pw import get_pw from .g_rh import get_rh, get_rh_2m from .g_slp import get_slp from .g_temp import get_tc, get_eth, get_temp, get_theta, get_tk, get_tv, get_tw from .g_terrain import get_terrain from .g_uvmet import (get_uvmet, get_uvmet10, get_uvmet10_wspd_wdir, get_uvmet_wspd_wdir, get_uvmet_wspd, get_uvmet_wdir, get_uvmet10_wspd, get_uvmet10_wdir) from .g_vorticity import get_avo, get_pvo from .g_wind import (get_destag_wspd_wdir, get_destag_wspd_wdir10, get_u_destag, get_v_destag, get_w_destag, get_destag_wspd, get_destag_wdir, get_destag_wspd10, get_destag_wdir10) from .g_times import get_times, get_xtimes from .g_cloudfrac import (get_cloudfrac, get_low_cloudfrac, get_mid_cloudfrac, get_high_cloudfrac) # func is the function to call. kargs are required arguments that should # not be altered by the user _FUNC_MAP = {"cape2d" : get_2dcape, "cape3d" : get_3dcape, "dbz" : get_dbz, "maxdbz" : get_max_dbz, "dp" : get_dp, "dp2m" : get_dp_2m, "height" : get_height, "geopt" : get_geopt, "srh" : get_srh, "uhel" : get_uh, "omega" : get_omega, "pw" : get_pw, "rh" : get_rh, "rh2m" : get_rh_2m, "slp" : get_slp, "theta" : get_theta, "temp" : get_temp, "tk" : get_tk, "tc" : get_tc, "theta_e" : get_eth, "tv" : get_tv, "twb" : get_tw, "terrain" : get_terrain, "times" : get_times, "xtimes" : get_xtimes, "uvmet" : get_uvmet, "uvmet10" : get_uvmet10, "avo" : get_avo, "pvo" : get_pvo, "ua" : get_u_destag, "va" : get_v_destag, "wa" : get_w_destag, "lat" : get_lat, "lon" : get_lon, "pressure" : get_pressure_hpa, "pres" : get_pressure, "wspd_wdir" : get_destag_wspd_wdir, "wspd_wdir10" : get_destag_wspd_wdir10, "uvmet_wspd_wdir" : get_uvmet_wspd_wdir, "uvmet10_wspd_wdir" : get_uvmet10_wspd_wdir, "ctt" : get_ctt, "cloudfrac" : get_cloudfrac, "geopt_stag" : get_stag_geopt, "zstag" : get_stag_height, # Diagnostics below are extracted from multi-product diagnostics "cape2d_only" : get_cape2d_only, "cin2d_only" : get_cin2d_only, "lcl" : get_lcl, "lfc" : get_lfc, "cape3d_only" : get_3dcape_only, "cin3d_only": get_3dcin_only, "uvmet_wspd" : get_uvmet_wspd, "uvmet_wdir" : get_uvmet_wdir, "uvmet10_wspd" : get_uvmet10_wspd, "uvmet10_wdir" : get_uvmet10_wdir, "wspd" : get_destag_wspd, "wdir" : get_destag_wdir, "wspd10" : get_destag_wspd10, "wdir10" : get_destag_wdir10, "low_cloudfrac" : get_low_cloudfrac, "mid_cloudfrac" : get_mid_cloudfrac, "high_cloudfrac" : get_high_cloudfrac } _VALID_KARGS = {"cape2d" : ["missing"], "cape3d" : ["missing"], "dbz" : ["do_variant", "do_liqskin"], "maxdbz" : ["do_variant", "do_liqskin"], "dp" : ["units"], "dp2m" : ["units"], "height" : ["msl", "units"], "geopt" : [], "srh" : ["top"], "uhel" : ["bottom", "top"], "omega" : [], "pw" : [], "rh" : [], "rh2m" : [], "slp" : ["units"], "temp" : ["units"], "tk" : [], "tc" : [], "theta" : ["units"], "theta_e" : ["units"], "tv" : ["units"], "twb" : ["units"], "terrain" : ["units"], "times" : [], "xtimes" : [], "uvmet" : ["units"], "uvmet10" : ["units"], "avo" : [], "pvo" : [], "ua" : ["units"], "va" : ["units"], "wa" : ["units"], "lat" : [], "lon" : [], "pres" : ["units"], "pressure" : ["units"], "wspd_wdir" : ["units"], "wspd_wdir10" : ["units"], "uvmet_wspd_wdir" : ["units"], "uvmet10_wspd_wdir" : ["units"], "ctt" : ["fill_nocloud", "missing", "opt_thresh", "units"], "cloudfrac" : ["vert_type", "low_thresh", "mid_thresh", "high_thresh"], "geopt_stag" : [], "zstag" : ["msl", "units"], "cape2d_only" : ["missing"], "cin2d_only" : ["missing"], "lcl" : ["missing"], "lfc" : ["missing"], "cape3d_only" : ["missing"], "cin3d_only": ["missing"], "uvmet_wspd" : ["units"], "uvmet_wdir" : ["units"], "uvmet10_wspd" : ["units"], "uvmet10_wdir" : ["units"], "wspd" : ["units"], "wdir" : ["units"], "wspd10" : ["units"], "wdir10" : ["units"], "low_cloudfrac" : ["vert_type", "low_thresh", "mid_thresh", "high_thresh"], "mid_cloudfrac" : ["vert_type", "low_thresh", "mid_thresh", "high_thresh"], "high_cloudfrac" : ["vert_type", "low_thresh", "mid_thresh", "high_thresh"], "default" : [] } _ALIASES = {"cape_2d" : "cape2d", "cape_3d" : "cape3d", "eth" : "theta_e", "mdbz" : "maxdbz", "geopotential" : "geopt", "helicity" : "srh", "latitude" : "lat", "longitude" : "lon", "omg" : "omega", "p" : "pres", "rh2" : "rh2m", "z": "height", "ter" : "terrain", "updraft_helicity" : "uhel", "td" : "dp", "td2" : "dp2m", "cfrac" : "cloudfrac", "wspd_wdir_uvmet" : "uvmet_wspd_wdir", "wspd_wdir_uvmet10" : "uvmet10_wspd_wdir", "th" : "theta", "low_cfrac" : "low_cloudfrac", "mid_cfrac" : "mid_cloudfrac", "high_cfrac" : "high_cloudfrac", "wspd_uvmet" : "uvmet_wspd" , "wdir_uvmet" : "uvmet_wdir" , "wspd_uvmet10" : "uvmet10_wspd" , "wdir_uvmet10" : "uvmet10_wdir" , } class ArgumentError(Exception): def __init__(self, msg): self.msg = msg def __str__(self): return self.msg def _undo_alias(alias): actual = _ALIASES.get(alias, None) if actual is None: return alias else: return actual def _check_kargs(var, kargs): for arg in viewkeys(kargs): if arg not in _VALID_KARGS[var]: if var != "default": raise ValueError("'{}' is an invalid keyword " "argument for '{}'".format(arg, var)) else: raise ValueError("'{}' is an invalid keyword " "argument".format(arg)) def getvar(wrfin, varname, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, **kwargs): """Returns basic diagnostics from the WRF ARW model output. A table of all available diagnostics is below. .. include:: ../../_templates/product_table.txt Args: wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ iterable): WRF-ARW NetCDF data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` or an iterable sequence of the aforementioned types. varname (:obj:`str`) : The variable name. timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The desired time index. This value can be a positive integer, negative integer, or :data:`wrf.ALL_TIMES` (an alias for None) to return all times in the file or sequence. The default is 0. method (:obj:`str`, optional): The aggregation method to use for sequences. Must be either 'cat' or 'join'. 'cat' combines the data along the Time dimension. 'join' creates a new dimension for the file index. The default is 'cat'. squeeze (:obj:`bool`, optional): Set to False to prevent dimensions with a size of 1 from being automatically removed from the shape of the output. Default is True. cache (:obj:`dict`, optional): A dictionary of (varname, ndarray) that can be used to supply pre-extracted NetCDF variables to the computational routines. It is primarily used for internal purposes, but can also be used to improve performance by eliminating the need to repeatedly extract the same variables used in multiple diagnostics calculations, particularly when using large sequences of files. Default is None. meta (:obj:`bool`, optional): Set to False to disable metadata and return :class:`numpy.ndarray` instead of :class:`xarray.DataArray`. Default is True. **kwargs: Optional keyword arguments for certain diagnostics. See table above. Returns: :class:`xarray.DataArray` or :class:`numpy.ndarray`: If xarray is enabled and the *meta* parameter is True, then the result will be a :class:`xarray.DataArray` object. Otherwise, the result will be a :class:`numpy.ndarray` object with no metadata. Raises: :class:`ValueError`: Raised when an invalid diagnostic type or keyword argument is passed to the routine. :class:`FortranError`: Raised when a problem occurs during a Fortran calculation. See Also: :class:`numpy.ndarray`, :class:`xarray.DataArray` Examples: Using netCDF4 .. code-block:: python from netCDF4 import Dataset from wrf import getvar wrfnc = Dataset("wrfout_d02_2010-06-13_21:00:00") slp = getvar(wrfnc, "slp") Using PyNIO .. code-block:: python from Nio import open_file from wrf import getvar wrfnc = open_file("wrfout_d02_2010-06-13_21:00:00"+".nc", "r") slp = getvar(wrfnc, "slp") Using Iterables: .. code-block:: python import os from netCDF4 import Dataset from wrf import getvar filedir = "/path/to/wrf/files" wrfin = [Dataset(f) for f in os.listdir(filedir) if f.startswith("wrfout_d02_")] uvmet = getvar(wrfin, "uvmet", timeidx=3, units="kt") """ _key = get_id(wrfin) wrfin = get_iterable(wrfin) if is_standard_wrf_var(wrfin, varname) and varname != "Times": _check_kargs("default", kwargs) return extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta, _key)[varname] elif varname == "Times": varname = "times" # Diverting to the get_times routine actual_var = _undo_alias(varname) if actual_var not in _VALID_KARGS: raise ValueError("'%s' is not a valid variable name" % (varname)) _check_kargs(actual_var, kwargs) return _FUNC_MAP[actual_var](wrfin, timeidx, method, squeeze, cache, meta, _key, **kwargs)