from __future__ import (absolute_import, division, print_function)

from .constants import Constants

from .extension import _tk, _eth, _tv, _wetbulb

from .decorators import convert_units

from .metadecorators import copy_and_set_metadata

from .util import extract_vars

@copy_and_set_metadata(copy_varname="T", name="theta",

description="potential temperature")

@convert_units("temp", "k")

def get_theta(wrfin, timeidx=0, method="cat", squeeze=True,

cache=None, meta=True, _key=None,

units="K"):

"""Return the potential temperature.

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.

units (:obj:`str`): The desired units. Refer to the :meth:`getvar`

product table for a list of available units for 'theta'. Default

is 'K'.

Returns:

:class:`xarray.DataArray` or :class:`numpy.ndarray`: The

potential temperature.

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.

"""

varnames = ("T", )

ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,

meta=False, _key=_key)

t = ncvars["T"]

full_t = t + Constants.T_BASE

return full_t

@copy_and_set_metadata(copy_varname="T", name="temp",

description="temperature")

@convert_units("temp", "k")

def get_temp(wrfin, timeidx=0, method="cat", squeeze=True,

cache=None, meta=True, _key=None,

units="K"):

"""Return the temperature in the specified units.

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.

units (:obj:`str`): The desired units. Refer to the :meth:`getvar`

product table for a list of available units for 'temp'. Default

is 'K'.

Returns:

:class:`xarray.DataArray` or :class:`numpy.ndarray`: The

temperature in the specified units.

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.

"""

varnames = ("T", "P", "PB")

ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,

meta=False, _key=_key)

t = ncvars["T"]

p = ncvars["P"]

pb = ncvars["PB"]

full_t = t + Constants.T_BASE

full_p = p + pb

tk = _tk(full_p, full_t)

return tk

@copy_and_set_metadata(copy_varname="T", name="theta_e",

description="equivalent potential temperature")

@convert_units("temp", "K")

def get_eth(wrfin, timeidx=0, method="cat", squeeze=True,

cache=None, meta=True, _key=None,

units="K"):

"""Return the equivalent potential temperature.

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.

units (:obj:`str`): The desired units. Refer to the :meth:`getvar`

product table for a list of available units for 'eth'. Default

is 'K'.

Returns:

:class:`xarray.DataArray` or :class:`numpy.ndarray`: The

equivalent potential temperature.

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.

"""

varnames = ("T", "P", "PB", "QVAPOR")

ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,

meta=False, _key=_key)

t = ncvars["T"]

p = ncvars["P"]

pb = ncvars["PB"]

qv = ncvars["QVAPOR"]

full_t = t + Constants.T_BASE

full_p = p + pb

tk = _tk(full_p, full_t)

eth = _eth(qv, tk, full_p)

return eth

@copy_and_set_metadata(copy_varname="T", name="tv",

description="virtual temperature")

@convert_units("temp", "k")

def get_tv(wrfin, timeidx=0, method="cat", squeeze=True,

cache=None, meta=True, _key=None,

units="K"):

"""Return the virtual temperature.

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.

units (:obj:`str`): The desired units. Refer to the :meth:`getvar`

product table for a list of available units for 'tv'. Default

is 'K'.

Returns:

:class:`xarray.DataArray` or :class:`numpy.ndarray`: The

virtual temperature.

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.

"""

varnames = ("T", "P", "PB", "QVAPOR")

ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,

meta=False, _key=_key)

t = ncvars["T"]

p = ncvars["P"]

pb = ncvars["PB"]

qv = ncvars["QVAPOR"]

full_t = t + Constants.T_BASE

full_p = p + pb

tk = _tk(full_p, full_t)

tv = _tv(tk, qv)

return tv

@copy_and_set_metadata(copy_varname="T", name="twb",

description="wetbulb temperature")

@convert_units("temp", "k")

def get_tw(wrfin, timeidx=0, method="cat", squeeze=True,

cache=None, meta=True, _key=None,

units="K"):

"""Return the wetbulb temperature.

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.

units (:obj:`str`): The desired units. Refer to the :meth:`getvar`

product table for a list of available units for 'twb'. Default

is 'K'.

Returns:

:class:`xarray.DataArray` or :class:`numpy.ndarray`: The

wetbulb temperature.

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.

"""

varnames = ("T", "P", "PB", "QVAPOR")

ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,

meta=False, _key=_key)

t = ncvars["T"]

p = ncvars["P"]

pb = ncvars["PB"]

qv = ncvars["QVAPOR"]

full_t = t + Constants.T_BASE

full_p = p + pb

tk = _tk(full_p, full_t)

tw = _wetbulb(full_p, tk, qv)

return tw

def get_tk(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,

meta=True, _key=None):

"""Return the temperature in [K].

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

temperature in [K].

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_temp(wrfin, timeidx, method, squeeze, cache, meta, _key,

units="K")

def get_tc(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,

meta=True, _key=None):

"""Return the temperature in [degC].

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

temperature in [degC].

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_temp(wrfin, timeidx, method, squeeze, cache, meta, _key,

units="degC")