From 3441467ecb26d83185f8f467dd5867b27042d25b Mon Sep 17 00:00:00 2001 From: Bill Ladwig Date: Fri, 22 Dec 2017 14:13:21 -0700 Subject: [PATCH] Added a product for the vertical staggered grid. Fixes #20. --- doc/source/_templates/product_table.txt | 16 ++- src/wrf/g_geoht.py | 158 +++++++++++++++++++++++- src/wrf/metadecorators.py | 28 ++++- src/wrf/routines.py | 8 +- 4 files changed, 197 insertions(+), 13 deletions(-) diff --git a/doc/source/_templates/product_table.txt b/doc/source/_templates/product_table.txt index 8dfb388..c49c3ac 100644 --- a/doc/source/_templates/product_table.txt +++ b/doc/source/_templates/product_table.txt @@ -35,7 +35,9 @@ | | | | | | | | | **do_liqskin** (boolean): Set to True to enable liquid skin calculation. Default is *False*. | +--------------------+---------------------------------------------------------------+-----------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------+ -| geopt/geopotential | Full Model Geopotential | m2 s-2 | | +| geopt/geopotential | Geopotential for the Mass Grid | m2 s-2 | | ++--------------------+---------------------------------------------------------------+-----------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------+ +| geopt_stag | Geopotential for the Vertically Staggered Grid | m2 s-2 | | +--------------------+---------------------------------------------------------------+-----------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------+ | helicity | Storm Relative Helicity | m2 s-2 | **top** (float): The top level for the calculation in meters. Default is *3000.0*. | +--------------------+---------------------------------------------------------------+-----------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -231,7 +233,17 @@ | | | | | | | | ft s-1 | | +--------------------+---------------------------------------------------------------+-----------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------+ -| z/height | Full Model Height on Mass Levels | m | **msl** (boolean): Set to False to return AGL values. True is for MSL. Default is *True*. | +| z/height | Model Height for Mass Grid | m | **msl** (boolean): Set to False to return AGL values. True is for MSL. Default is *True*. | +| | | | | +| | | km | **units** (str) : Set to desired units. Default is *'m'*. | +| | | | | +| | | dm | | +| | | | | +| | | ft | | +| | | | | +| | | mi | | ++--------------------+---------------------------------------------------------------+-----------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------+ +| zstag | Model Height for Vertically Staggered Grid | m | **msl** (boolean): Set to False to return AGL values. True is for MSL. Default is *True*. | | | | | | | | | km | **units** (str) : Set to desired units. Default is *'m'*. | | | | | | diff --git a/src/wrf/g_geoht.py b/src/wrf/g_geoht.py index 98d9ceb..d9c7624 100755 --- a/src/wrf/g_geoht.py +++ b/src/wrf/g_geoht.py @@ -1,6 +1,8 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) +import warnings + from .constants import Constants from .destag import destagger from .decorators import convert_units @@ -9,7 +11,7 @@ from .util import extract_vars, either def _get_geoht(wrfin, timeidx, method="cat", squeeze=True, cache=None, meta=True, _key=None, - height=True, msl=True): + height=True, msl=True, stag=False): """Return the geopotential or geopotential height. If *height* is False, then geopotential is returned in units of @@ -67,6 +69,9 @@ def _get_geoht(wrfin, timeidx, method="cat", squeeze=True, as Mean Sea Level (MSL). Set to False to return the geopotential height as Above Ground Level (AGL) by subtracting the terrain height. Default is True. + + stag (:obj:`bool`, optional): Set to True to use the vertical + staggered grid, rather than the mass grid. Default is False. Returns: :class:`xarray.DataArray` or :class:`numpy.ndarray`: The @@ -86,14 +91,21 @@ def _get_geoht(wrfin, timeidx, method="cat", squeeze=True, phb = ph_vars["PHB"] hgt = ph_vars["HGT"] geopt = ph + phb - geopt_unstag = destagger(geopt, -3) + if not stag: + geopt_unstag = destagger(geopt, -3) + else: + geopt_unstag = geopt else: ght_vars = extract_vars(wrfin, timeidx, ("GHT", "HGT_M"), method, squeeze, cache, meta=False, _key=_key) geopt_unstag = ght_vars["GHT"] * Constants.G hgt = ght_vars["HGT_M"] - + + if stag: + warnings.warn("file contains no vertically staggered geopotential " + "height variable, returning unstaggered result " + "instead" ) if height: if msl: return geopt_unstag / Constants.G @@ -110,7 +122,7 @@ def _get_geoht(wrfin, timeidx, method="cat", squeeze=True, return geopt_unstag -@set_height_metadata(geopt=True) +@set_height_metadata(geopt=True, stag=False) def get_geopt(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): """Return the geopotential. @@ -171,7 +183,7 @@ def get_geopt(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, False, True) -@set_height_metadata(geopt=False) +@set_height_metadata(geopt=False, stag=False) @convert_units("height", "m") def get_height(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, @@ -245,3 +257,139 @@ def get_height(wrfin, timeidx=0, method="cat", squeeze=True, return _get_geoht(wrfin, timeidx, method, squeeze, cache, meta, _key, True, msl) + +@set_height_metadata(geopt=True, stag=True) +def get_stag_geopt(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, + meta=True, _key=None): + """Return the geopotential for the vertically staggered grid. + + The geopotential is returned in units of [m2 s-2]. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + 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. + + 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. + + _key (:obj:`int`, optional): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + geopotential. + 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. + + """ + return _get_geoht(wrfin, timeidx, method, squeeze, cache, meta, _key, + False, True, stag=True) + + +@set_height_metadata(geopt=False, stag=True) +@convert_units("height", "m") +def get_stag_height(wrfin, timeidx=0, method="cat", squeeze=True, + cache=None, meta=True, _key=None, + msl=True, units="m"): + """Return the geopotential height for the vertically staggered grid. + + If *msl* is True, then geopotential height is returned as Mean Sea Level + (MSL). If *msl* is False, then geopotential height is returned as + Above Ground Level (AGL) by subtracting the terrain height. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + 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. + + 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. + + _key (:obj:`int`, optional): A caching key. This is used for internal + purposes only. Default is None. + + msl (:obj:`bool`, optional): Set to True to return geopotential height + as Mean Sea Level (MSL). Set to False to return the + geopotential height as Above Ground Level (AGL) by subtracting + the terrain height. Default is True. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'z'. Default + is 'm'. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + geopotential height. + 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. + + """ + + return _get_geoht(wrfin, timeidx, method, squeeze, cache, meta, _key, + True, msl, stag=True) + \ No newline at end of file diff --git a/src/wrf/metadecorators.py b/src/wrf/metadecorators.py index c397ba1..0549220 100644 --- a/src/wrf/metadecorators.py +++ b/src/wrf/metadecorators.py @@ -645,7 +645,7 @@ def set_latlon_metadata(xy=False): return func_wrapper -def set_height_metadata(geopt=False): +def set_height_metadata(geopt=False, stag=False): """A decorator that sets the metadata for a wrapped height function's output. @@ -660,6 +660,9 @@ def set_height_metadata(geopt=False): returns geopotential. Set to True if the wrapped function returns geopotential height. Default is False. + stag (:obj:`bool`, optional): Set to True to use the vertical + staggered grid, rather than the mass grid. Default is False. + Returns: :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped @@ -695,9 +698,17 @@ def set_height_metadata(geopt=False): if cache is None: cache = {} + is_met_em = False # For height, either copy the met_em GHT variable or copy and modify # pressure (which has the same dims as destaggered height) - ht_metadata_varname = either("P", "GHT")(wrfin) + if not stag: + ht_metadata_varname = either("P", "GHT")(wrfin) + else: + ht_metadata_varname = either("PH", "GHT")(wrfin) + + if ht_metadata_varname == "GHT": + is_met_em = True + ht_var = extract_vars(wrfin, timeidx, ht_metadata_varname, method, squeeze, cache, meta=True, _key=_key) @@ -723,12 +734,21 @@ def set_height_metadata(geopt=False): if geopt: outname = "geopt" outattrs["units"] = "m2 s-2" - outattrs["description"] = "full model geopotential" + if not stag or is_met_em: + outattrs["description"] = "geopotential (mass grid)" + else: + outattrs["description"] = ("geopotential (vertically " + "staggered grid)") else: outname = "height" if msl else "height_agl" outattrs["units"] = units height_type = "MSL" if msl else "AGL" - outattrs["description"] = "model height ({})".format(height_type) + if not stag or is_met_em: + outattrs["description"] = ("model height - [{}] " + "(mass grid)".format(height_type)) + else: + outattrs["description"] = ("model height - [{}] (vertically " + "staggered grid)".format(height_type)) return DataArray(result, name=outname, diff --git a/src/wrf/routines.py b/src/wrf/routines.py index 3267598..5442956 100644 --- a/src/wrf/routines.py +++ b/src/wrf/routines.py @@ -7,7 +7,7 @@ from .g_cape import get_2dcape, get_3dcape 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 +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 @@ -68,7 +68,9 @@ _FUNC_MAP = {"cape2d" : get_2dcape, "uvmet_wspd_wdir" : get_uvmet_wspd_wdir, "uvmet10_wspd_wdir" : get_uvmet10_wspd_wdir, "ctt" : get_ctt, - "cloudfrac" : get_cloudfrac + "cloudfrac" : get_cloudfrac, + "geopt_stag" : get_stag_geopt, + "zstag" : get_stag_height } _VALID_KARGS = {"cape2d" : ["missing"], @@ -113,6 +115,8 @@ _VALID_KARGS = {"cape2d" : ["missing"], "ctt" : [], "cloudfrac" : ["vert_type", "low_thresh", "mid_thresh", "high_thresh"], + "geopt_stag" : [], + "zstag" : ["msl", "units"], "default" : [] }