forked from 3rdparty/wrf-python
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
394 lines
18 KiB
394 lines
18 KiB
from __future__ import (absolute_import, division, print_function) |
|
|
|
import warnings |
|
|
|
from .constants import Constants |
|
from .destag import destagger |
|
from .decorators import convert_units |
|
from .metadecorators import set_height_metadata |
|
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, stag=False): |
|
"""Return the geopotential or geopotential height. |
|
|
|
If *height* is False, then geopotential is returned in units of |
|
[m2 s-2]. If *height* is True, then geopotential height is |
|
returned in units of [m]. If *msl* is True, then geopotential height |
|
is return as Mean Sea Level (MSL). If *msl* is False, then geopotential |
|
height is returned as Above Ground Level (AGL). |
|
|
|
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. |
|
|
|
height (:obj:`bool`, optional): Set to True to return geopotential |
|
height instead of geopotential. Default is True. |
|
|
|
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. |
|
|
|
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 |
|
geopotential or 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. |
|
|
|
""" |
|
|
|
varname = either("PH", "GHT")(wrfin) |
|
if varname == "PH": |
|
ph_vars = extract_vars(wrfin, timeidx, ("PH", "PHB", "HGT"), |
|
method, squeeze, cache, meta=False, |
|
_key=_key) |
|
ph = ph_vars["PH"] |
|
phb = ph_vars["PHB"] |
|
hgt = ph_vars["HGT"] |
|
geopt = ph + phb |
|
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 |
|
else: |
|
# Due to broadcasting with multifile/multitime, the 2D terrain |
|
# array needs to be reshaped to a 3D array so the right dims |
|
# line up |
|
new_dims = list(hgt.shape) |
|
new_dims.insert(-2,1) |
|
hgt = hgt.reshape(new_dims) |
|
|
|
return (geopt_unstag / Constants.G) - hgt |
|
else: |
|
return geopt_unstag |
|
|
|
|
|
@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. |
|
|
|
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) |
|
|
|
|
|
@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, |
|
msl=True, units="m"): |
|
"""Return the geopotential height. |
|
|
|
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) |
|
|
|
|
|
@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) |
|
|