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.
2131 lines
76 KiB
2131 lines
76 KiB
from __future__ import (absolute_import, division, print_function, |
|
unicode_literals) |
|
import wrapt |
|
from collections import OrderedDict |
|
|
|
import numpy as np |
|
import numpy.ma as ma |
|
|
|
from .extension import _interpline |
|
from .util import (extract_vars, either, from_args, arg_location, |
|
is_coordvar, latlon_coordvars, to_np, |
|
from_var, iter_left_indexes) |
|
from .coordpair import CoordPair |
|
from .py3compat import viewkeys, viewitems, py3range, ucode |
|
from .interputils import get_xy_z_params, get_xy, to_xy_coords |
|
from .config import xarray_enabled |
|
|
|
if xarray_enabled(): |
|
from xarray import DataArray |
|
|
|
|
|
def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, |
|
remove_dims=None, dimnames=None, |
|
coords=None, **fixed_attrs): |
|
"""A decorator that sets the metadata for a wrapped function's output. |
|
|
|
Generally, the metadata is copied from the variable specified by |
|
*copy_varname*, with other fixed fields set by the other decorator |
|
arguments. |
|
|
|
The *cache* argument used by most diagnostic routines is supplied by |
|
this decorator when the *copy_varname* variable is extracted, in order |
|
to prevent the variable from being extracted again by the wrapped function. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
copy_varname (:obj:`str`, optional): The NetCDF variable name to copy. |
|
Default is None. |
|
|
|
delete_attrs (sequence of :obj:`str`, optional): A sequence of key |
|
names to remove from the :attr:`xarray.DataArray.attrs` attribute |
|
in the wrapped function output (after being copied from |
|
*copy_varname*). Default is None. |
|
|
|
name (:obj:`str`): The name to use for the |
|
:attr:`xarray.DataArray.name` attribute. |
|
|
|
remove_dims (sequence of :obj:`int`, optional): A sequence of dimension |
|
indexes to be removed from the wrapped function output (after being |
|
copied from *copy_varname*). This is useful when the copy |
|
variable is three dimensional but the wrapped function output |
|
is two dimensional, and you still want to keep the names of |
|
the rightmost two dimensions. Default is None. |
|
|
|
dimnames (sequence of :obj:`str`, optional): A sequence of dimension |
|
names in order to manually set the |
|
:attr:`xarray.DataArray.dims` attribute. Default is None. |
|
|
|
coords (:obj:`dict`): A mapping of coordinate name to coordinate |
|
value to manually specify the :attr:`xarray.DataArray.coords` |
|
attribute. Default is None. |
|
|
|
**fixed_attrs: These keyword arguments are added to the |
|
:attr:`xarray.DataArray.attrs` attribute. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
function output with or without metadata. 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. |
|
|
|
|
|
""" |
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
argvars = from_args(wrapped, ("wrfin", "timeidx", "method", |
|
"squeeze", "cache", "units", "meta", |
|
"_key"), |
|
*args, **kwargs) |
|
|
|
wrfin = argvars["wrfin"] |
|
timeidx = argvars["timeidx"] |
|
units = argvars["units"] |
|
method = argvars["method"] |
|
squeeze = argvars["squeeze"] |
|
cache = argvars["cache"] |
|
_key = argvars["_key"] |
|
if cache is None: |
|
cache = {} |
|
|
|
# Note: can't modify nonlocal var |
|
if (callable(copy_varname)): |
|
_copy_varname = copy_varname(wrfin) |
|
else: |
|
_copy_varname = copy_varname |
|
|
|
# Extract the copy_from argument |
|
var_to_copy = None if cache is None else cache.get(_copy_varname, |
|
None) |
|
|
|
|
|
if var_to_copy is None: |
|
var_to_copy = extract_vars(wrfin, timeidx, (_copy_varname,), |
|
method, squeeze, cache, |
|
meta=True, _key=_key)[_copy_varname] |
|
|
|
# Make a copy so we don't modify a user supplied cache |
|
new_cache = dict(cache) |
|
new_cache[_copy_varname] = var_to_copy |
|
|
|
# Don't modify the original args/kargs. The args need to be a list |
|
# so it can be modified. |
|
new_args, cache_argloc = arg_location(wrapped, "cache", args, kwargs) |
|
new_args[cache_argloc] = new_cache |
|
|
|
result = wrapped(*new_args) |
|
|
|
outname = "" |
|
outdimnames = list() |
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
|
|
if copy_varname is not None: |
|
outname = ucode(var_to_copy.name) |
|
|
|
if dimnames is not None: |
|
outdimnames = dimnames |
|
outcoords = coords |
|
else: |
|
outdimnames += var_to_copy.dims |
|
outcoords.update(var_to_copy.coords) |
|
|
|
outattrs.update(var_to_copy.attrs) |
|
|
|
if remove_dims is not None: |
|
for dimname in remove_dims: |
|
outdimnames.remove(dimname) |
|
|
|
try: |
|
del outcoords[dimname] |
|
except KeyError: |
|
pass |
|
|
|
|
|
if name is not None: |
|
outname = name |
|
|
|
if units is not None: |
|
outattrs["units"] = units |
|
|
|
for argname, val in viewitems(fixed_attrs): |
|
outattrs[argname] = val |
|
|
|
if delete_attrs is not None: |
|
for attr in delete_attrs: |
|
try: |
|
del outattrs[attr] |
|
except KeyError: |
|
pass |
|
|
|
if isinstance(result, ma.MaskedArray): |
|
outattrs["_FillValue"] = result.fill_value |
|
outattrs["missing_value"] = result.fill_value |
|
|
|
return DataArray(result, name=outname, coords=outcoords, |
|
dims=outdimnames, attrs=outattrs) |
|
|
|
return func_wrapper |
|
|
|
|
|
def set_wind_metadata(copy_varname, name, description, |
|
wind_ncvar=False, |
|
two_d=False, wspd_wdir=False): |
|
"""A decorator that sets the metadata for a wrapped wind function's output. |
|
|
|
This is a special metadata decorator for working with wind functions, which |
|
include wind extraction routines, uvmet routines, and |
|
wind speed / wind direction routines. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
copy_varname (:obj:`str`, optional): The NetCDF variable name to copy. |
|
Default is None. |
|
|
|
name (:obj:`str`): The name to use for the |
|
:attr:`xarray.DataArray.name` attribute. |
|
|
|
description (:obj:`str`): The description to use for the 'description' |
|
key in the :attr:`xarray.DataArray.attrs` attribute. |
|
|
|
wind_ncvar (:obj:`bool`, optional): Set to True when the wrapped |
|
function is simply extracting a wind variable (U, V, W) from the |
|
NetCDF file. Set to False for other types of wind algorithms |
|
(uvmet, wspd_wdir, etc). Default is False. |
|
|
|
two_d (:obj:`bool`, optional): Set to True if the wind field is |
|
two-dimensional. Set to False for a three-dimensional wind field. |
|
Default is False. |
|
|
|
wspd_wdir (:obj:`bool`): Set to True if the wrapped function is a |
|
wind speed / wind direction algorithm. Otherwise, set to False. |
|
Default is False. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
wind function output with or without metadata. 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. |
|
|
|
|
|
""" |
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
argvars = from_args(wrapped, ("wrfin", "timeidx", "units", |
|
"method", "squeeze", "ten_m", "cache", |
|
"_key"), |
|
*args, **kwargs) |
|
wrfin = argvars["wrfin"] |
|
timeidx = argvars["timeidx"] |
|
units = argvars["units"] |
|
method = argvars["method"] |
|
squeeze = argvars["squeeze"] |
|
ten_m = argvars["ten_m"] |
|
cache = argvars["cache"] |
|
_key = argvars["_key"] |
|
if cache is None: |
|
cache = {} |
|
|
|
if isinstance(copy_varname, either): |
|
_copy_varname = copy_varname(wrfin) |
|
else: |
|
_copy_varname = copy_varname |
|
|
|
copy_var = extract_vars(wrfin, timeidx, _copy_varname, |
|
method, squeeze, cache, |
|
meta=True, _key=_key)[_copy_varname] |
|
|
|
# Make a copy so we don't modify a user supplied cache |
|
new_cache = dict(cache) |
|
new_cache[_copy_varname] = copy_var |
|
|
|
# Don't modify the original args/kargs. The args need to be a list |
|
# so it can be modified. |
|
new_args, cache_argloc = arg_location(wrapped, "cache", args, kwargs) |
|
new_args[cache_argloc] = new_cache |
|
|
|
result = wrapped(*new_args) |
|
|
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
|
|
outdimnames = list(copy_var.dims) |
|
#outcoords.update(copy_var.coords) |
|
outattrs.update(copy_var.attrs) |
|
|
|
|
|
if wind_ncvar: |
|
outcoords.update(copy_var.coords) |
|
elif not wspd_wdir: |
|
if not two_d: |
|
outdimnames.insert(0, "u_v") |
|
else: |
|
outdimnames.insert(0, "u_v") |
|
outattrs["MemoryOrder"] = "XY" |
|
outcoords["u_v"] = ["u", "v"] |
|
else: |
|
if not two_d: |
|
outdimnames.insert(0, "wspd_wdir") |
|
else: |
|
outdimnames.insert(0, "wspd_wdir") |
|
outattrs["MemoryOrder"] = "XY" |
|
|
|
outcoords["wspd_wdir"] = ["wspd", "wdir"] |
|
|
|
if units is not None: |
|
outattrs["units"] = units |
|
|
|
# xarray doesn't line up coordinate dimensions based on |
|
# names, it just remembers the index it originally mapped to. |
|
# So, need to rebuild the XLAT, XLONG, coordinates again since the |
|
# leftmost index changed. |
|
if not wind_ncvar: |
|
for key,dataarray in viewitems(copy_var.coords): |
|
if is_coordvar(key): |
|
outcoords[key] = dataarray.dims, to_np(dataarray) |
|
elif key == "XTIME": |
|
outcoords[key] = dataarray.dims, to_np(dataarray) |
|
elif key == "Time": |
|
outcoords[key] = to_np(dataarray) |
|
|
|
outname = name |
|
outattrs["description"] = description |
|
|
|
return DataArray(result, name=outname, coords=outcoords, |
|
dims=outdimnames, attrs=outattrs) |
|
|
|
return func_wrapper |
|
|
|
def set_cape_metadata(is2d): |
|
"""A decorator that sets the metadata for a wrapped CAPE function's output. |
|
|
|
This is a special metadata decorator for working with CAPE functions. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
is2d (:obj:`bool`): Set to True if the wrapped function is for a |
|
two-dimensional CAPE routine. Set to False for a |
|
three-dimensional CAPE routine. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
CAPE function output with or without metadata. 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. |
|
|
|
""" |
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
argvars = from_args(wrapped, ("wrfin", "timeidx", "method", "squeeze", |
|
"cache", "_key", "missing"), |
|
*args, **kwargs) |
|
wrfin = argvars["wrfin"] |
|
timeidx = argvars["timeidx"] |
|
method = argvars["method"] |
|
squeeze = argvars["squeeze"] |
|
cache = argvars["cache"] |
|
missing = argvars["missing"] |
|
_key = argvars["_key"] |
|
if cache is None: |
|
cache = {} |
|
|
|
_copy_varname = "P" |
|
copy_var = extract_vars(wrfin, timeidx, _copy_varname, method, squeeze, |
|
cache, meta=True, _key=_key)[_copy_varname] |
|
|
|
# Make a copy so we don't modify a user supplied cache |
|
new_cache = dict(cache) |
|
new_cache[_copy_varname] = copy_var |
|
|
|
# Don't modify the original args/kargs. The args need to be a list |
|
# so it can be modified. |
|
new_args, cache_argloc = arg_location(wrapped, "cache", args, kwargs) |
|
new_args[cache_argloc] = new_cache |
|
|
|
result = wrapped(*new_args) |
|
|
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
outattrs.update(copy_var.attrs) |
|
outdimnames = [None] * result.ndim |
|
|
|
if is2d: |
|
# Right dims |
|
outdimnames[-2:] = copy_var.dims[-2:] |
|
# Left dims |
|
outdimnames[1:-2] = copy_var.dims[0:-3] |
|
outdimnames[0] = "mcape_mcin_lcl_lfc" |
|
outattrs["description"] = "mcape ; mcin ; lcl ; lfc" |
|
outattrs["MemoryOrder"] = "XY" |
|
outattrs["units"] = "J kg-1 ; J kg-1 ; m ; m" |
|
outname = "cape_2d" |
|
else: |
|
# Right dims |
|
outdimnames[-3:] = copy_var.dims[-3:] |
|
# Left dims |
|
outdimnames[1:-3] = copy_var.dims[0:-3] |
|
outdimnames[0] = "cape_cin" |
|
outattrs["description"] = "cape; cin" |
|
outattrs["units"] = "J kg-1 ; J kg-1" |
|
outattrs["MemoryOrder"] = "XYZ" |
|
outname = "cape_3d" |
|
|
|
outattrs["_FillValue"] = missing |
|
outattrs["missing_value"] = missing |
|
|
|
|
|
# xarray doesn't line up coordinate dimensions based on |
|
# names, it just remembers the index it originally mapped to. |
|
# So, need to rebuild the XLAT, XLONG, coordinates again since the |
|
# leftmost index changed. |
|
|
|
for key,dataarray in viewitems(copy_var.coords): |
|
if is_coordvar(key): |
|
outcoords[key] = dataarray.dims, to_np(dataarray) |
|
elif key == "XTIME": |
|
outcoords[key] = dataarray.dims, to_np(dataarray) |
|
elif key == "Time": |
|
outcoords[key] = to_np(dataarray) |
|
|
|
if is2d: |
|
outcoords["mcape_mcin_lcl_lfc"] = ["mcape", "mcin", "lcl", "lfc"] |
|
else: |
|
outcoords["cape_cin"] = ["cape", "cin"] |
|
|
|
|
|
return DataArray(result, name=outname, coords=outcoords, |
|
dims=outdimnames, attrs=outattrs) |
|
|
|
return func_wrapper |
|
|
|
|
|
def set_cloudfrac_metadata(): |
|
"""A decorator that sets the metadata for a wrapped cloud fraction |
|
function's output. |
|
|
|
This is a special metadata decorator for working with cloud fraction |
|
functions. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
cloud fraction function output with or without metadata. 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. |
|
|
|
""" |
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
argvars = from_args(wrapped, ("wrfin", "timeidx", "method", "squeeze", |
|
"cache", "_key"), |
|
*args, **kwargs) |
|
wrfin = argvars["wrfin"] |
|
timeidx = argvars["timeidx"] |
|
method = argvars["method"] |
|
squeeze = argvars["squeeze"] |
|
cache = argvars["cache"] |
|
_key = argvars["_key"] |
|
if cache is None: |
|
cache = {} |
|
|
|
_copy_varname = "P" |
|
copy_var = extract_vars(wrfin, timeidx, _copy_varname, method, squeeze, |
|
cache, meta=True, _key=_key)[_copy_varname] |
|
|
|
# Make a copy so we don't modify a user supplied cache |
|
new_cache = dict(cache) |
|
new_cache[_copy_varname] = copy_var |
|
|
|
# Don't modify the original args/kargs. The args need to be a list |
|
# so it can be modified. |
|
new_args, cache_argloc = arg_location(wrapped, "cache", args, kwargs) |
|
new_args[cache_argloc] = new_cache |
|
|
|
result = wrapped(*new_args) |
|
|
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
outattrs.update(copy_var.attrs) |
|
outdimnames = [None] * result.ndim |
|
|
|
# Right dims |
|
outdimnames[-2:] = copy_var.dims[-2:] |
|
# Left dims |
|
outdimnames[1:-2] = copy_var.dims[0:-3] |
|
outdimnames[0] = "low_mid_high" |
|
outattrs["description"] = "low, mid, high clouds" |
|
outattrs["MemoryOrder"] = "XY" |
|
outattrs["units"] = "%" |
|
outname = "cloudfrac" |
|
|
|
# xarray doesn't line up coordinate dimensions based on |
|
# names, it just remembers the index it originally mapped to. |
|
# So, need to rebuild the XLAT, XLONG, coordinates again since the |
|
# leftmost index changed. |
|
|
|
for key,dataarray in viewitems(copy_var.coords): |
|
if is_coordvar(key): |
|
outcoords[key] = dataarray.dims, to_np(dataarray) |
|
elif key == "XTIME": |
|
outcoords[key] = dataarray.dims, to_np(dataarray) |
|
elif key == "Time": |
|
outcoords[key] = to_np(dataarray) |
|
|
|
outcoords["low_mid_high"] = ["low", "mid", "high"] |
|
|
|
return DataArray(result, name=outname, coords=outcoords, |
|
dims=outdimnames, attrs=outattrs) |
|
|
|
return func_wrapper |
|
|
|
def set_latlon_metadata(xy=False): |
|
"""A decorator that sets the metadata for a wrapped latlon function's |
|
output. |
|
|
|
This is a special metadata decorator for working with latlon functions. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
xy (:obj:`bool`, optional): Set to True if the wrapped function returns |
|
xy values (ll_to_xy). Set to False if the wrapped function returns |
|
latlon values (xy_to_ll). Default is False. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
latlon function output with or without metadata. 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. |
|
|
|
""" |
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
|
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
# Set squeeze to False. Squeeze will be handled in here |
|
new_args, squeeze_argloc = arg_location(wrapped, "squeeze", args, |
|
kwargs) |
|
new_args[squeeze_argloc] = False |
|
|
|
result = wrapped(*new_args) |
|
|
|
# Want to preserve the input coordinate pair in metadata |
|
if result.ndim == 1: |
|
result = result[:, np.newaxis] |
|
|
|
argnames = ["x", "y"] if not xy else ["latitude", "longitude"] |
|
argnames.append("squeeze") |
|
outname = "latlon" if not xy else "xy" |
|
|
|
if result.ndim == 2: |
|
dimnames = (["lat_lon", "idx"] if not xy else ["x_y", "idx"]) |
|
else: |
|
dimnames = (["lat_lon", "domain_idx", "idx"] if not xy |
|
else ["x_y", "domain_idx", "idx"]) |
|
|
|
argvars = from_args(wrapped, argnames, *args, **kwargs) |
|
|
|
var1 = argvars[argnames[0]] |
|
var2 = argvars[argnames[1]] |
|
squeeze = argvars["squeeze"] |
|
|
|
arr1 = np.asarray(var1).ravel() |
|
arr2 = np.asarray(var2).ravel() |
|
|
|
coords = {} |
|
if not xy: |
|
coords["xy_coord"] = (dimnames[-1], [CoordPair(x=x[0], y=x[1]) |
|
for x in zip(arr1, arr2)]) |
|
coords[dimnames[0]] = ["lat", "lon"] |
|
else: |
|
coords["latlon_coord"] = (dimnames[-1], [CoordPair(lat=x[0], |
|
lon=x[1]) |
|
for x in zip(arr1, arr2)]) |
|
coords[dimnames[0]] = ["x", "y"] |
|
|
|
da = DataArray(result, name=outname, dims=dimnames, coords=coords) |
|
|
|
if squeeze: |
|
da = da.squeeze() |
|
|
|
return da |
|
|
|
return func_wrapper |
|
|
|
def set_height_metadata(geopt=False): |
|
"""A decorator that sets the metadata for a wrapped height function's |
|
output. |
|
|
|
This is a special metadata decorator for working with height functions. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
geopt (:obj:`bool`, optional): Set to True if the wrapped function |
|
returns geopotential. Set to True if the wrapped function |
|
returns geopotential height. Default is False. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
height function output with or without metadata. 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. |
|
|
|
""" |
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
argvars = from_args(wrapped, ("wrfin", "timeidx", "method", |
|
"squeeze", "units", "msl", "cache", |
|
"_key"), |
|
*args, **kwargs) |
|
wrfin = argvars["wrfin"] |
|
timeidx = argvars["timeidx"] |
|
units = argvars["units"] |
|
method = argvars["method"] |
|
squeeze = argvars["squeeze"] |
|
msl = argvars["msl"] |
|
cache = argvars["cache"] |
|
_key = argvars["_key"] |
|
|
|
if cache is None: |
|
cache = {} |
|
|
|
# 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) |
|
ht_var = extract_vars(wrfin, timeidx, ht_metadata_varname, |
|
method, squeeze, cache, meta=True, |
|
_key=_key) |
|
ht_metadata_var = ht_var[ht_metadata_varname] |
|
|
|
# Make a copy so we don't modify a user supplied cache |
|
new_cache = dict(cache) |
|
new_cache[ht_metadata_varname] = ht_metadata_var |
|
|
|
# Don't modify the original args/kargs. The args need to be a list |
|
# so it can be modified. |
|
new_args, cache_argloc = arg_location(wrapped, "cache", args, kwargs) |
|
new_args[cache_argloc] = new_cache |
|
|
|
result = wrapped(*new_args) |
|
|
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
outdimnames = list(ht_metadata_var.dims) |
|
outcoords.update(ht_metadata_var.coords) |
|
outattrs.update(ht_metadata_var.attrs) |
|
|
|
if geopt: |
|
outname = "geopt" |
|
outattrs["units"] = "m2 s-2" |
|
outattrs["description"] = "full model geopotential" |
|
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) |
|
|
|
|
|
return DataArray(result, name=outname, |
|
dims=outdimnames, coords=outcoords, attrs=outattrs) |
|
return func_wrapper |
|
|
|
def _set_horiz_meta(wrapped, instance, args, kwargs): |
|
"""A decorator implementation that sets the metadata for a wrapped |
|
horizontal interpolation function. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
wrapped: The wrapped function which in turns needs to be called by your |
|
wrapper function. |
|
|
|
instance: The object to which the wrapped function was bound when it |
|
was called. |
|
|
|
args: The list of positional arguments supplied when the decorated |
|
function was called. |
|
|
|
kwargs: The dictionary of keyword arguments supplied when the decorated |
|
function was called. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
horiztontal interpolation function output with or without metadata. |
|
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. |
|
|
|
See Also: |
|
|
|
:mod:`wrapt` |
|
|
|
""" |
|
argvars = from_args(wrapped, ("field3d", "vert", "desiredlev", |
|
"missing"), |
|
*args, **kwargs) |
|
|
|
field3d = argvars["field3d"] |
|
z = argvars["vert"] |
|
desiredloc = argvars["desiredlev"] |
|
missingval = argvars["missing"] |
|
|
|
result = wrapped(*args, **kwargs) |
|
|
|
# Defaults, in case the data isn't a DataArray |
|
outname = None |
|
outdimnames = None |
|
outcoords = None |
|
outattrs = None |
|
|
|
# Get the vertical level units |
|
vert_units = None |
|
if isinstance(z, DataArray): |
|
vert_units = z.attrs.get("units", None) |
|
|
|
# If we have no metadata to start with, only set the level |
|
levelstr = ("{0} {1}".format(desiredloc, vert_units) |
|
if vert_units is not None |
|
else "{0}".format(desiredloc)) |
|
|
|
name_levelstr = ("{0}_{1}".format(desiredloc, vert_units) |
|
if vert_units is not None |
|
else "{0}".format(desiredloc)) |
|
|
|
if isinstance(field3d, DataArray): |
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
outdimnames = list(field3d.dims) |
|
outcoords.update(field3d.coords) |
|
outdimnames.remove(field3d.dims[-3]) |
|
try: |
|
del outcoords[field3d.dims[-3]] |
|
except KeyError: |
|
pass # xarray 0.9 |
|
outattrs.update(field3d.attrs) |
|
outname = "{0}_{1}".format(field3d.name, name_levelstr) |
|
|
|
else: |
|
outname = "field3d_{0}".format(name_levelstr) |
|
outattrs = OrderedDict() |
|
|
|
outattrs["level"] = levelstr |
|
outattrs["missing_value"] = missingval |
|
outattrs["_FillValue"] = missingval |
|
|
|
for key in ("MemoryOrder", "description"): |
|
try: |
|
del outattrs[key] |
|
except KeyError: |
|
pass |
|
|
|
return DataArray(result, name=outname, dims=outdimnames, |
|
coords=outcoords, attrs=outattrs) |
|
|
|
def _set_cross_meta(wrapped, instance, args, kwargs): |
|
"""A decorator implementation that sets the metadata for a wrapped cross \ |
|
section interpolation function. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
wrapped: The wrapped function which in turns needs to be called by your |
|
wrapper function. |
|
|
|
instance: The object to which the wrapped function was bound when it |
|
was called. |
|
|
|
args: The list of positional arguments supplied when the decorated |
|
function was called. |
|
|
|
kwargs: The dictionary of keyword arguments supplied when the decorated |
|
function was called. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
cross section interpolation function output with or without metadata. |
|
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. |
|
|
|
See Also: |
|
|
|
:mod:`wrapt` |
|
|
|
""" |
|
argvars = from_args(wrapped, ("field3d", "vert", "levels", |
|
"latlon", "missing", |
|
"wrfin", "timeidx", "stagger", "projection", |
|
"pivot_point", "angle", |
|
"start_point", "end_point", |
|
"cache"), |
|
*args, **kwargs) |
|
|
|
field3d = argvars["field3d"] |
|
z = argvars["vert"] |
|
levels = argvars["levels"] |
|
inc_latlon = argvars["latlon"] |
|
missingval = argvars["missing"] |
|
wrfin = argvars["wrfin"] |
|
timeidx = argvars["timeidx"] |
|
stagger = argvars["stagger"] |
|
projection = argvars["projection"] |
|
pivot_point = argvars["pivot_point"] |
|
angle = argvars["angle"] |
|
start_point = argvars["start_point"] |
|
end_point = argvars["end_point"] |
|
cache = argvars["cache"] |
|
|
|
start_point_xy = None |
|
end_point_xy = None |
|
pivot_point_xy = None |
|
|
|
if pivot_point is not None: |
|
if pivot_point.lat is not None and pivot_point.lon is not None: |
|
xy_coords = to_xy_coords(pivot_point, wrfin, timeidx, |
|
stagger, projection) |
|
pivot_point_xy = (xy_coords.x, xy_coords.y) |
|
else: |
|
pivot_point_xy = (pivot_point.x, pivot_point.y) |
|
|
|
if start_point is not None and end_point is not None: |
|
if start_point.lat is not None and start_point.lon is not None: |
|
xy_coords = to_xy_coords(start_point, wrfin, timeidx, |
|
stagger, projection) |
|
start_point_xy = (xy_coords.x, xy_coords.y) |
|
else: |
|
start_point_xy = (start_point.x, start_point.y) |
|
|
|
if end_point.lat is not None and end_point.lon is not None: |
|
xy_coords = to_xy_coords(end_point, wrfin, timeidx, |
|
stagger, projection) |
|
end_point_xy = (xy_coords.x, xy_coords.y) |
|
else: |
|
end_point_xy = (end_point.x, end_point.y) |
|
|
|
xy, var2dz, z_var2d = get_xy_z_params(to_np(z), pivot_point_xy, angle, |
|
start_point_xy, end_point_xy, levels) |
|
|
|
# Make a copy so we don't modify a user supplied cache |
|
if cache is not None: |
|
new_cache = dict(cache) |
|
else: |
|
new_cache = {} |
|
new_cache["xy"] = xy |
|
new_cache["var2dz"] = var2dz |
|
new_cache["z_var2d"] = z_var2d |
|
|
|
# Don't modify the original args/kargs. The args need to be a list |
|
# so it can be modified. |
|
new_args, cache_argloc = arg_location(wrapped, "cache", args, kwargs) |
|
new_args[cache_argloc] = new_cache |
|
|
|
result = wrapped(*new_args) |
|
|
|
# Defaults, in case the data isn't a DataArray |
|
outname = None |
|
outdimnames = None |
|
outcoords = None |
|
outattrs = None |
|
|
|
# Use XY to set the cross-section metadata |
|
st_x = xy[0,0] |
|
st_y = xy[0,1] |
|
ed_x = xy[-1,0] |
|
ed_y = xy[-1,1] |
|
|
|
cross_str = "({0}, {1}) to ({2}, {3})".format(st_x, st_y, ed_x, ed_y) |
|
if angle is not None: |
|
cross_str += " ; center={0} ; angle={1}".format(pivot_point, |
|
angle) |
|
|
|
if isinstance(field3d, DataArray): |
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
outdimnames = list(field3d.dims) |
|
outcoords.update(field3d.coords) |
|
for i in py3range(-3,0,1): |
|
outdimnames.remove(field3d.dims[i]) |
|
try: |
|
del outcoords[field3d.dims[i]] |
|
except KeyError: |
|
pass # Xarray 0.9 |
|
|
|
|
|
# Delete any lat,lon coords |
|
delkeys = [key for key in viewkeys(outcoords) if is_coordvar(key)] |
|
for key in delkeys: |
|
del outcoords[key] |
|
|
|
outdimnames.append("vertical") |
|
outdimnames.append("cross_line_idx") |
|
outattrs.update(field3d.attrs) |
|
|
|
outname = "{0}_cross".format(field3d.name) |
|
|
|
for key in ("MemoryOrder",): |
|
try: |
|
del outattrs[key] |
|
except KeyError: |
|
pass |
|
|
|
# Interpolate to get the lat/lon coords, if desired |
|
if inc_latlon: |
|
latcoordname, loncoordname = latlon_coordvars(field3d.coords) |
|
|
|
if latcoordname is not None and loncoordname is not None: |
|
latcoord = field3d.coords[latcoordname] |
|
loncoord = field3d.coords[loncoordname] |
|
|
|
if latcoord.ndim == 2: |
|
lats = _interpline(latcoord, xy) |
|
lons = _interpline(loncoord, xy) |
|
|
|
outcoords["xy_loc"] = ("cross_line_idx", |
|
np.asarray(tuple( |
|
CoordPair(x=xy[i,0], y=xy[i,1], |
|
lat=lats[i], lon=lons[i]) |
|
for i in py3range(xy.shape[-2]))) |
|
) |
|
# Moving domain |
|
else: |
|
extra_dims = latcoord.shape[0:-2] |
|
outdims = extra_dims + xy.shape[-2:-1] |
|
|
|
latlon_loc = np.empty(outdims, np.object_) |
|
for left_dims in iter_left_indexes(extra_dims): |
|
idxs = left_dims + (slice(None),) |
|
lats = _interpline(latcoord[idxs], xy) |
|
lons = _interpline(loncoord[idxs], xy) |
|
|
|
latlon_loc[idxs] = np.asarray(tuple( |
|
CoordPair(x=xy[i,0], y=xy[i,1], |
|
lat=lats[i], lon=lons[i]) |
|
for i in py3range(xy.shape[-2])) |
|
)[:] |
|
|
|
|
|
extra_dimnames = latcoord.dims[0:-2] |
|
loc_dimnames = extra_dimnames + ("cross_line_idx",) |
|
outcoords["xy_loc"] = (loc_dimnames, latlon_loc) |
|
|
|
else: |
|
outcoords["xy_loc"] = ("cross_line_idx", np.asarray(tuple( |
|
CoordPair(xy[i,0], xy[i,1]) |
|
for i in py3range(xy.shape[-2])))) |
|
|
|
else: |
|
outcoords["xy_loc"] = ("cross_line_idx", np.asarray(tuple( |
|
CoordPair(xy[i,0], xy[i,1]) |
|
for i in py3range(xy.shape[-2])))) |
|
|
|
outcoords["vertical"] = z_var2d[:] |
|
|
|
else: |
|
outname = "field3d_cross" |
|
outattrs = OrderedDict() |
|
|
|
outattrs["orientation"] = cross_str |
|
outattrs["missing_value"] = missingval |
|
outattrs["_FillValue"] = missingval |
|
|
|
return DataArray(result, name=outname, dims=outdimnames, |
|
coords=outcoords, attrs=outattrs) |
|
|
|
|
|
|
|
def _set_line_meta(wrapped, instance, args, kwargs): |
|
"""A decorator implementation that sets the metadata for a wrapped line |
|
interpolation function. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
wrapped: The wrapped function which in turns needs to be called by your |
|
wrapper function. |
|
|
|
instance: The object to which the wrapped function was bound when it |
|
was called. |
|
|
|
args: The list of positional arguments supplied when the decorated |
|
function was called. |
|
|
|
kwargs: The dictionary of keyword arguments supplied when the decorated |
|
function was called. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
line interpolation function output with or without metadata. |
|
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. |
|
|
|
See Also: |
|
|
|
:mod:`wrapt` |
|
|
|
""" |
|
argvars = from_args(wrapped, ("field2d", |
|
"wrfin", "timeidx", "stagger", "projection", |
|
"pivot_point", "angle", |
|
"start_point", "end_point", "latlon", |
|
"cache"), |
|
*args, **kwargs) |
|
|
|
field2d = argvars["field2d"] |
|
wrfin = argvars["wrfin"] |
|
timeidx = argvars["timeidx"] |
|
stagger = argvars["stagger"] |
|
projection = argvars["projection"] |
|
pivot_point = argvars["pivot_point"] |
|
angle = argvars["angle"] |
|
start_point = argvars["start_point"] |
|
end_point = argvars["end_point"] |
|
inc_latlon = argvars["latlon"] |
|
cache = argvars["cache"] |
|
|
|
if cache is None: |
|
cache = {} |
|
|
|
start_point_xy = None |
|
end_point_xy = None |
|
pivot_point_xy = None |
|
|
|
if pivot_point is not None: |
|
if pivot_point.lat is not None and pivot_point.lon is not None: |
|
xy_coords = to_xy_coords(pivot_point, wrfin, timeidx, |
|
stagger, projection) |
|
pivot_point_xy = (xy_coords.x, xy_coords.y) |
|
else: |
|
pivot_point_xy = (pivot_point.x, pivot_point.y) |
|
|
|
|
|
if start_point is not None and end_point is not None: |
|
if start_point.lat is not None and start_point.lon is not None: |
|
xy_coords = to_xy_coords(start_point, wrfin, timeidx, |
|
stagger, projection) |
|
start_point_xy = (xy_coords.x, xy_coords.y) |
|
else: |
|
start_point_xy = (start_point.x, start_point.y) |
|
|
|
if end_point.lat is not None and end_point.lon is not None: |
|
xy_coords = to_xy_coords(end_point, wrfin, timeidx, |
|
stagger, projection) |
|
end_point_xy = (xy_coords.x, xy_coords.y) |
|
else: |
|
end_point_xy = (end_point.x, end_point.y) |
|
|
|
xy = get_xy(field2d, pivot_point_xy, angle, start_point_xy, end_point_xy) |
|
|
|
# Make a copy so we don't modify a user supplied cache |
|
new_cache = dict(cache) |
|
new_cache["xy"] = xy |
|
|
|
# Don't modify the original args/kargs. The args need to be a list |
|
# so it can be modified. |
|
new_args, cache_argloc = arg_location(wrapped, "cache", args, kwargs) |
|
new_args[cache_argloc] = new_cache |
|
|
|
result = wrapped(*new_args) |
|
|
|
# Defaults, in case the data isn't a DataArray |
|
outname = None |
|
outdimnames = None |
|
outcoords = None |
|
outattrs = None |
|
|
|
# Use XY to set the cross-section metadata |
|
st_x = xy[0,0] |
|
st_y = xy[0,1] |
|
ed_x = xy[-1,0] |
|
ed_y = xy[-1,1] |
|
|
|
cross_str = "({0}, {1}) to ({2}, {3})".format(st_x, st_y, ed_x, ed_y) |
|
if angle is not None: |
|
cross_str += " ; center={0} ; angle={1}".format(pivot_point, |
|
angle) |
|
|
|
if isinstance(field2d, DataArray): |
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
outdimnames = list(field2d.dims) |
|
outcoords.update(field2d.coords) |
|
for i in py3range(-2,0,1): |
|
outdimnames.remove(field2d.dims[i]) |
|
try: |
|
del outcoords[field2d.dims[i]] |
|
except KeyError: |
|
pass # xarray 0.9 |
|
|
|
# Delete any lat,lon coords |
|
delkeys = [key for key in viewkeys(outcoords) if is_coordvar(key)] |
|
for key in delkeys: |
|
del outcoords[key] |
|
|
|
outdimnames.append("line_idx") |
|
outattrs.update(field2d.attrs) |
|
|
|
outname = "{0}_line".format(field2d.name) |
|
|
|
for key in ("MemoryOrder",): |
|
try: |
|
del outattrs[key] |
|
except KeyError: |
|
pass |
|
|
|
# Interpolate to get the lat/lon coords, if desired |
|
if inc_latlon: |
|
latcoordname, loncoordname = latlon_coordvars(field2d.coords) |
|
|
|
if latcoordname is not None and loncoordname is not None: |
|
latcoord = field2d.coords[latcoordname] |
|
loncoord = field2d.coords[loncoordname] |
|
|
|
if latcoord.ndim == 2: |
|
lats = _interpline(latcoord, xy) |
|
lons = _interpline(loncoord, xy) |
|
|
|
outcoords["xy_loc"] = ("line_idx", |
|
np.asarray(tuple( |
|
CoordPair(x=xy[i,0], y=xy[i,1], |
|
lat=lats[i], lon=lons[i]) |
|
for i in py3range(xy.shape[-2]))) |
|
) |
|
|
|
# Moving domain |
|
else: |
|
extra_dims = latcoord.shape[0:-2] |
|
outdims = extra_dims + xy.shape[-2:-1] |
|
|
|
latlon_loc = np.empty(outdims, np.object_) |
|
for left_dims in iter_left_indexes(extra_dims): |
|
idxs = left_dims + (slice(None),) |
|
lats = _interpline(latcoord[idxs], xy) |
|
lons = _interpline(loncoord[idxs], xy) |
|
|
|
latlon_loc[idxs] = np.asarray(tuple( |
|
CoordPair(x=xy[i,0], y=xy[i,1], |
|
lat=lats[i], lon=lons[i]) |
|
for i in py3range(xy.shape[-2])) |
|
)[:] |
|
|
|
|
|
extra_dimnames = latcoord.dims[0:-2] |
|
loc_dimnames = extra_dimnames + ("line_idx",) |
|
outcoords["xy_loc"] = (loc_dimnames, latlon_loc) |
|
|
|
else: |
|
outcoords["xy_loc"] = ("line_idx", np.asarray(tuple( |
|
CoordPair(xy[i,0], xy[i,1]) |
|
for i in py3range(xy.shape[-2])))) |
|
|
|
else: |
|
outcoords["xy_loc"] = ("line_idx", np.asarray(tuple( |
|
CoordPair(xy[i,0], xy[i,1]) |
|
for i in py3range(xy.shape[-2])))) |
|
|
|
else: |
|
outname = "field2d_line" |
|
outattrs = OrderedDict() |
|
|
|
outattrs["orientation"] = cross_str |
|
|
|
return DataArray(result, name=outname, dims=outdimnames, |
|
coords=outcoords, attrs=outattrs) |
|
|
|
|
|
def _set_vinterp_meta(wrapped, instance, args, kwargs): |
|
"""A decorator implementation that sets the metadata for a wrapped |
|
vertical coordinate interpolation function. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
wrapped: The wrapped function which in turns needs to be called by your |
|
wrapper function. |
|
|
|
instance: The object to which the wrapped function was bound when it |
|
was called. |
|
|
|
args: The list of positional arguments supplied when the decorated |
|
function was called. |
|
|
|
kwargs: The dictionary of keyword arguments supplied when the decorated |
|
function was called. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
vertical coordinate interpolation function output with or without |
|
metadata. |
|
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. |
|
|
|
See Also: |
|
|
|
:mod:`wrapt` |
|
|
|
""" |
|
argvars = from_args(wrapped, ("wrfin", "field", "vert_coord", |
|
"interp_levels", "extrapolate", |
|
"field_type", "log_p", |
|
"timeidx", "method", "squeeze", |
|
"cache"), |
|
*args, **kwargs) |
|
|
|
field = argvars["field"] |
|
vert_coord = argvars["vert_coord"] |
|
interp_levels = argvars["interp_levels"] |
|
field_type = argvars["field_type"] |
|
|
|
result = wrapped(*args, **kwargs) |
|
|
|
# Defaults, in case the data isn't a DataArray |
|
outname = None |
|
outdimnames = None |
|
outcoords = None |
|
outattrs = None |
|
|
|
|
|
if isinstance(field, DataArray): |
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
outdimnames = list(field.dims) |
|
outcoords.update(field.coords) |
|
|
|
outdimnames.remove(field.dims[-3]) |
|
try: |
|
del outcoords[field.dims[-3]] |
|
except KeyError: |
|
pass # xarray 0.9 |
|
|
|
outdimnames.insert(-2, "interp_level") |
|
outcoords["interp_level"] = interp_levels |
|
outattrs.update(field.attrs) |
|
outattrs["vert_interp_type"] = vert_coord |
|
|
|
outname = field.name |
|
|
|
else: |
|
outname = field_type |
|
|
|
|
|
return DataArray(result, name=outname, dims=outdimnames, |
|
coords=outcoords, attrs=outattrs) |
|
|
|
|
|
def _set_2dxy_meta(wrapped, instance, args, kwargs): |
|
"""A decorator implementation that sets the metadata for a wrapped line |
|
cross section interpolation function. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
wrapped: The wrapped function which in turns needs to be called by your |
|
wrapper function. |
|
|
|
instance: The object to which the wrapped function was bound when it |
|
was called. |
|
|
|
args: The list of positional arguments supplied when the decorated |
|
function was called. |
|
|
|
kwargs: The dictionary of keyword arguments supplied when the decorated |
|
function was called. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
line cross section interpolation function output with or without |
|
metadata. |
|
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. |
|
|
|
See Also: |
|
|
|
:mod:`wrapt` |
|
|
|
""" |
|
argvars = from_args(wrapped, ("field3d", "xy"), *args, **kwargs) |
|
|
|
field3d = argvars["field3d"] |
|
xy = argvars["xy"] |
|
xy = to_np(xy) |
|
|
|
result = wrapped(*args, **kwargs) |
|
|
|
# Use XY to set the cross-section metadata |
|
st_x = xy[0,0] |
|
st_y = xy[0,1] |
|
ed_x = xy[-1,0] |
|
ed_y = xy[-1,1] |
|
|
|
cross_str = "({0},{1}) to ({2},{3})".format(st_x, st_y, |
|
ed_x, ed_y) |
|
|
|
outname = None |
|
outdimnames = None |
|
outcoords = None |
|
outattrs = None |
|
|
|
# Dims are (...,xy,z) |
|
if isinstance(field3d, DataArray): |
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
outdimnames = list(field3d.dims) |
|
outcoords.update(field3d.coords) |
|
|
|
for i in py3range(-2,0,1): |
|
try: |
|
del outcoords[field3d.dims[i]] |
|
except KeyError: |
|
pass # xarray 0.9 |
|
outdimnames.remove(field3d.dims[i]) |
|
|
|
# Need to remove XLAT, XLONG... |
|
delkeys = (key for key,arr in viewitems(field3d.coords) |
|
if arr.ndim > 1) |
|
|
|
for key in delkeys: |
|
del outcoords[key] |
|
|
|
outdimnames.append("line_idx") |
|
#outattrs.update(field3d.attrs) |
|
|
|
desc = field3d.attrs.get("description", None) |
|
if desc is not None: |
|
outattrs["description"] = desc |
|
|
|
units = field3d.attrs.get("units", None) |
|
if units is not None: |
|
outattrs["units"] = units |
|
|
|
outname = "{0}_2dxy".format(field3d.name) |
|
|
|
outcoords["xy_loc"] = ("line_idx", [CoordPair(xy[i,0], xy[i,1]) |
|
for i in py3range(xy.shape[-2])]) |
|
|
|
for key in ("MemoryOrder",): |
|
try: |
|
del outattrs[key] |
|
except KeyError: |
|
pass |
|
|
|
else: |
|
outname = "field3d_2dxy" |
|
|
|
outattrs["orientation"] = cross_str |
|
|
|
return DataArray(result, name=outname, dims=outdimnames, |
|
coords=outcoords, attrs=outattrs) |
|
|
|
|
|
def _set_1d_meta(wrapped, instance, args, kwargs): |
|
"""A decorator implementation that sets the metadata for a wrapped 1D |
|
interpolation function. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
wrapped: The wrapped function which in turns needs to be called by your |
|
wrapper function. |
|
|
|
instance: The object to which the wrapped function was bound when it |
|
was called. |
|
|
|
args: The list of positional arguments supplied when the decorated |
|
function was called. |
|
|
|
kwargs: The dictionary of keyword arguments supplied when the decorated |
|
function was called. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
1D interpolation function output with or without metadata. |
|
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. |
|
|
|
See Also: |
|
|
|
:mod:`wrapt` |
|
|
|
""" |
|
argvars = from_args(wrapped, ("field", "z_in", "z_out", "missingval"), |
|
*args, **kwargs) |
|
|
|
field = argvars["field"] |
|
z_in = argvars["z_in"] |
|
z_out = argvars["z_out"] |
|
missingval = argvars["missingval"] |
|
|
|
result = wrapped(*args, **kwargs) |
|
|
|
outname = None |
|
outdimnames = None |
|
outcoords = None |
|
outattrs = None |
|
|
|
# Dims are (...,xy,z) |
|
if isinstance(field, DataArray): |
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
outdimnames = list(field.dims) |
|
|
|
outdimnames.pop(-1) |
|
|
|
for name in outdimnames: |
|
try: |
|
outcoords[name] = field.coords[name] |
|
except KeyError: |
|
continue |
|
|
|
outdimnames.append("z") |
|
outname = "{0}_z".format(field.name) |
|
outcoords["z"] = z_out |
|
|
|
outattrs["_FillValue"] = missingval |
|
outattrs["missing_value"] = missingval |
|
|
|
desc = field.attrs.get("description", None) |
|
if desc is not None: |
|
outattrs["description"] = desc |
|
|
|
units = field.attrs.get("units", None) |
|
if units is not None: |
|
outattrs["units"] = units |
|
|
|
else: |
|
outname = "field_z" |
|
|
|
|
|
return DataArray(result, name=outname, dims=outdimnames, |
|
coords=outcoords, attrs=outattrs) |
|
|
|
|
|
def _set_xy_meta(wrapped, instance, args, kwargs): |
|
"""A decorator implementation that sets the metadata for a wrapped xy line |
|
interpolation function. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
wrapped: The wrapped function which in turns needs to be called by your |
|
wrapper function. |
|
|
|
instance: The object to which the wrapped function was bound when it |
|
was called. |
|
|
|
args: The list of positional arguments supplied when the decorated |
|
function was called. |
|
|
|
kwargs: The dictionary of keyword arguments supplied when the decorated |
|
function was called. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
xy line interpolation function output with or without metadata. |
|
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. |
|
|
|
See Also: |
|
|
|
:mod:`wrapt` |
|
|
|
""" |
|
argvars = from_args(wrapped, ("field", "pivot_point", "angle", |
|
"start_point", "end_point"), |
|
*args, **kwargs) |
|
|
|
field = argvars["field"] |
|
pivot_point = argvars["pivot_point"] |
|
angle = argvars["angle"] |
|
start_point = argvars["start_point"] |
|
end_point = argvars["end_point"] |
|
|
|
result = wrapped(*args, **kwargs) |
|
|
|
if isinstance(field, DataArray): |
|
outname = "{0}_xy".format(field.name) |
|
else: |
|
outname = "xy" |
|
|
|
outdimnames = ["line_idx", "x_y"] |
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
|
|
outcoords["x_y"] = ["x", "y"] |
|
|
|
if pivot_point is not None and angle is not None: |
|
outattrs["pivot_point"] = pivot_point |
|
outattrs["angle"] = angle |
|
|
|
if start_point is not None and end_point is not None: |
|
outattrs["start_point"] = start_point |
|
outattrs["end_point"] = end_point |
|
|
|
return DataArray(result, name=outname, dims=outdimnames, |
|
coords=outcoords, attrs=outattrs) |
|
|
|
|
|
def set_interp_metadata(interp_type): |
|
"""A decorator that sets the metadata for a wrapped interpolation |
|
function. |
|
|
|
If the wrapped function's *meta* argument is False, then this decorator |
|
returns the wrapped function output without applying the metadata. |
|
|
|
Args: |
|
|
|
interp_type (:obj:`str`): The type of interpolation routine. Choices |
|
are: 'horiz', 'cross', 'line', 'vinterp', '2dxy', '1d', |
|
'xy'. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
interpolation function output with or without metadata. |
|
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. |
|
|
|
""" |
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
if interp_type == "horiz": |
|
return _set_horiz_meta(wrapped, instance, args, kwargs) |
|
elif interp_type == "cross": |
|
return _set_cross_meta(wrapped, instance, args, kwargs) |
|
elif interp_type == "line": |
|
return _set_line_meta(wrapped, instance, args, kwargs) |
|
elif interp_type == "vinterp": |
|
return _set_vinterp_meta(wrapped, instance, args, kwargs) |
|
elif interp_type == "2dxy": |
|
return _set_2dxy_meta(wrapped, instance, args, kwargs) |
|
elif interp_type == "1d": |
|
return _set_1d_meta(wrapped, instance, args, kwargs) |
|
elif interp_type == "xy": |
|
return _set_xy_meta(wrapped, instance, args, kwargs) |
|
|
|
return func_wrapper |
|
|
|
|
|
def set_alg_metadata(alg_ndims, refvarname, |
|
refvarndims=None, missingarg=None, |
|
stagdim=None, stagsubvar=None, |
|
units=None, description=None): |
|
"""A decorator that sets the metadata for a wrapped raw diagnostic |
|
function. |
|
|
|
Args: |
|
|
|
alg_ndims (:obj:`int`): The number of dimensions returned by the |
|
wrapped function. |
|
|
|
refvarname (:obj:`str`): The wrapped function argument name for the |
|
reference variable. |
|
|
|
refvarndims (:obj:`int`, optional): The number of right dimensions for |
|
the reference variable. This paramter is required when the |
|
wrapped function result has less dimensions than reference. |
|
Default is None. |
|
|
|
missingarg (:obj:`str`, optional): The wrapped function argument name |
|
for the missing value variable. Default is None. |
|
|
|
stagdim (:obj`int`, optional): The staggered dimension for the |
|
reference. This is only needed if the reference variable is |
|
on a staggered grid. Default is None. |
|
|
|
stagsubvar (:obj:`str`, optional): The wrapped function argument name |
|
to use to supply the unstaggered dimension name for the staggered |
|
dimension in the reference. This is needed if *stagdim* is not |
|
None. It is primarily used for absolute vorticity. Default is None. |
|
|
|
units (:obj:`str`, optional): The units to use if if there is no |
|
'units' argument for the wrapped function. Default is None. |
|
|
|
description (:obj:`str`, optional): A description for the wrapped |
|
algorithm, which is stored in the :attr:`xarray.DataArray.attrs` |
|
attribute under the 'description' key. Default is None. |
|
|
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
numerical function output with or without metadata. |
|
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. |
|
|
|
""" |
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
|
|
result = wrapped(*args, **kwargs) |
|
|
|
outname = wrapped.__name__ |
|
outattrs = OrderedDict() |
|
|
|
# Default dimension names |
|
outdims = ["dim_{}".format(i) for i in py3range(result.ndim)] |
|
|
|
if missingarg is not None: |
|
missingval = from_args(wrapped, (missingarg,), |
|
*args, **kwargs)[missingarg] |
|
else: |
|
missingval = None |
|
|
|
if missingval is not None: |
|
outattrs["_FillValue"] = missingval |
|
outattrs["missing_value"] = missingval |
|
result = np.ma.masked_values(result, missingval) |
|
|
|
if units is not None: |
|
if isinstance(units, from_var): |
|
_units = units(wrapped, *args, **kwargs) |
|
if _units is not None: |
|
outattrs["units"] = _units |
|
else: |
|
outattrs["units"] = units |
|
|
|
else: |
|
# Check for a units argument, if not, just ignore |
|
_units = from_args(wrapped, ("units",), *args, **kwargs)["units"] |
|
if _units is not None: |
|
outattrs["units"] = _units |
|
|
|
|
|
if description is not None: |
|
if isinstance(description, from_var): |
|
desc = description(wrapped, *args, **kwargs) |
|
if desc is not None: |
|
outattrs["description"] = desc |
|
else: |
|
outattrs["description"] = description |
|
|
|
|
|
# Copy the dimnames from the reference variable, otherwise, use |
|
# the supplied dimnames |
|
if refvarname is not None: |
|
refvar = from_args(wrapped, (refvarname,), |
|
*args, **kwargs)[refvarname] |
|
else: |
|
refvar = None |
|
|
|
if stagsubvar is not None: |
|
stagvar = from_args(wrapped, (stagsubvar,), |
|
*args, **kwargs)[stagsubvar] |
|
else: |
|
stagvar = None |
|
|
|
if isinstance(refvar, DataArray): |
|
|
|
# Copy the right dims |
|
outdims[-alg_ndims:] = refvar.dims[-alg_ndims:] |
|
|
|
# Use the stagsubvar if applicable |
|
if stagvar is not None and stagdim is not None: |
|
outdims[stagdim] = stagvar.dims[stagdim] |
|
|
|
# Left dims |
|
if refvarndims is None: |
|
# Used when result and reference are aligned on right |
|
if result.ndim > alg_ndims: |
|
result_extra = result.ndim - alg_ndims |
|
|
|
for i in py3range(1, result_extra + 1): |
|
idx = -alg_ndims - i |
|
if -idx <= refvar.ndim: |
|
outdims[idx] = refvar.dims[idx] |
|
else: |
|
continue |
|
# When reference and result aren't exactly aligned (slp) |
|
# (reference is 3D, result is 2D) |
|
else: |
|
ref_extra = refvar.ndim - refvarndims |
|
ref_left_dimnames = refvar.dims[0:ref_extra] |
|
|
|
for i,dimname in enumerate(ref_left_dimnames[::-1], 1): |
|
if i <= result.ndim: |
|
outdims[-alg_ndims - i] = dimname |
|
else: |
|
continue |
|
|
|
out = DataArray(result, name=outname, dims=outdims, attrs=outattrs) |
|
|
|
return out |
|
|
|
return func_wrapper |
|
|
|
def set_smooth_metdata(): |
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
argvars = from_args(wrapped, ("field", "passes"), |
|
*args, **kwargs) |
|
|
|
field = argvars["field"] |
|
passes = argvars["passes"] |
|
|
|
result = wrapped(*args, **kwargs) |
|
|
|
outname = "smooth2d" |
|
outdimnames = ["dim_{}".format(i) for i in py3range(result.ndim)] |
|
outcoords = OrderedDict() |
|
outattrs = OrderedDict() |
|
|
|
if isinstance(field, DataArray): |
|
outname = "smooth_" + ucode(field.name) |
|
outdimnames = list(field.dims) |
|
outcoords.update(field.coords) |
|
outattrs.update(field.attrs) |
|
outattrs["passes"] = passes |
|
|
|
if isinstance(result, ma.MaskedArray): |
|
outattrs["_FillValue"] = result.fill_value |
|
outattrs["missing_value"] = result.fill_value |
|
|
|
return DataArray(result, name=outname, coords=outcoords, |
|
dims=outdimnames, attrs=outattrs) |
|
|
|
return func_wrapper |
|
|
|
|
|
def set_uvmet_alg_metadata(units=None, description="earth rotated u,v", |
|
latarg="lat", windarg="u"): |
|
"""A decorator that sets the metadata for the wrapped raw UVMET diagnostic |
|
function. |
|
|
|
Args: |
|
|
|
units (:obj:`str`, optional): The units to use if if there is no |
|
'units' argument for the wrapped function. Default is None. |
|
|
|
description (:obj:`str`, optional): A description for the wrapped |
|
algorithm, which is stored in the :attr:`xarray.DataArray.attrs` |
|
attribute under the 'description' key. Default is None. |
|
|
|
latarg (:obj:'str`, optional): The wrapped function argument name for |
|
latitude. Default is 'lat'. |
|
|
|
windarg (:obj:`str`, optional): The wrapped function argument name for |
|
the u wind component. Default is 'u'. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
UVMET function output with or without metadata. |
|
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. |
|
|
|
""" |
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
result = wrapped(*args, **kwargs) |
|
|
|
# Default dimension names |
|
outdims = ["dim_{}".format(i) for i in py3range(result.ndim)] |
|
|
|
outname = "uvmet" |
|
outattrs = OrderedDict() |
|
|
|
if units is not None: |
|
outattrs["units"] = units |
|
else: |
|
_units = from_args(wrapped, ("units",), *args, **kwargs)["units"] |
|
if _units is not None: |
|
outattrs["units"] = _units |
|
|
|
if description is not None: |
|
outattrs["description"] = description |
|
|
|
latvar = from_args(wrapped, latarg, *args, **kwargs)[latarg] |
|
uvar = from_args(wrapped, windarg, *args, **kwargs)[windarg] |
|
|
|
if isinstance(uvar, DataArray) and isinstance(latvar, DataArray): |
|
# Right dims come from latvar |
|
outdims[-2:] = latvar.dims[-2:] |
|
|
|
# Left dims come from u-var |
|
outdims[1:-2] = uvar.dims[0:-2] |
|
|
|
# Left-most is always u_v |
|
outdims[0] = "u_v" |
|
outcoords = {} |
|
outcoords["u_v"] = ["u", "v"] |
|
|
|
out = DataArray(result, name=outname, dims=outdims, coords=outcoords, |
|
attrs=outattrs) |
|
|
|
return out |
|
|
|
return func_wrapper |
|
|
|
|
|
def set_cape_alg_metadata(is2d, copyarg="pres_hpa"): |
|
"""A decorator that sets the metadata for the wrapped raw CAPE diagnostic |
|
function. |
|
|
|
Args: |
|
|
|
is2d (:obj:`bool`): Set to True for the two-dimensional CAPE |
|
calculation, False for three-dimensional CAPE. |
|
|
|
copyarg (:obj:`str`): The wrapped function argument to use for |
|
copying dimension names. Default is 'pres_hpa'. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
CAPE function output with or without metadata. |
|
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. |
|
|
|
""" |
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
result = wrapped(*args, **kwargs) |
|
|
|
argvals = from_args(wrapped, (copyarg,"missing"), *args, **kwargs) |
|
p = argvals[copyarg] |
|
missing = argvals["missing"] |
|
|
|
# Note: cape_3d supports using only a single column of data |
|
is1d = p.ndim == 1 |
|
|
|
# Need to squeeze the right dimensions for 1D cape |
|
if is1d: |
|
result = np.squeeze(result) |
|
|
|
# Default dimension names |
|
outdims = ["dim_{}".format(i) for i in py3range(result.ndim)] |
|
|
|
outattrs = OrderedDict() |
|
|
|
if is2d: |
|
outname = "cape_2d" |
|
outattrs["description"] = "mcape ; mcin ; lcl ; lfc" |
|
outattrs["units"] = "J kg-1 ; J kg-1 ; m ; m" |
|
outattrs["MemoryOrder"] = "XY" |
|
else: |
|
if not is1d: |
|
outname = "cape_3d" |
|
outattrs["description"] = "cape; cin" |
|
outattrs["units"] = "J kg-1 ; J kg-1" |
|
outattrs["MemoryOrder"] = "XYZ" |
|
else: |
|
outname = "cape_column" |
|
outattrs["description"] = "cape; cin" |
|
outattrs["units"] = "J kg-1 ; J kg-1" |
|
outattrs["MemoryOrder"] = "Z" |
|
|
|
|
|
if isinstance(p, DataArray): |
|
if is2d: |
|
# Right dims |
|
outdims[-2:] = p.dims[-2:] |
|
# Left dims |
|
outdims[1:-2] = p.dims[0:-3] |
|
|
|
else: |
|
if not is1d: |
|
# Right dims |
|
outdims[-3:] = p.dims[-3:] |
|
# Left dims |
|
outdims[1:-3] = p.dims[0:-3] |
|
else: |
|
outdims[1] = p.dims[0] |
|
|
|
outcoords = {} |
|
# Left-most is always cape_cin or cape_cin_lcl_lfc |
|
if is2d: |
|
outdims[0] = "cape_cin_lcl_lfc" |
|
outcoords["cape_cin_lcl_lfc"] = ["cape", "cin", "lcl", "lfc"] |
|
else: |
|
outdims[0] = "cape_cin" |
|
outcoords["cape_cin"] = ["cape", "cin"] |
|
|
|
outattrs["_FillValue"] = missing |
|
outattrs["missing_value"] = missing |
|
|
|
out = DataArray(result, name=outname, dims=outdims, coords=outcoords, |
|
attrs=outattrs) |
|
|
|
return out |
|
|
|
return func_wrapper |
|
|
|
|
|
def set_cloudfrac_alg_metadata(copyarg="pres"): |
|
"""A decorator that sets the metadata for the wrapped raw cloud fraction |
|
diagnostic function. |
|
|
|
Args: |
|
|
|
copyarg (:obj:`str`): The wrapped function argument to use for |
|
copying dimension names. Default is 'pres'. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
cloud fraction function output with or without metadata. |
|
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. |
|
|
|
""" |
|
|
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
result = wrapped(*args, **kwargs) |
|
|
|
# Default dimension names |
|
outdims = ["dim_{}".format(i) for i in py3range(result.ndim)] |
|
|
|
outattrs = OrderedDict() |
|
|
|
outname = "cloudfrac" |
|
outattrs["description"] = "low, mid, high clouds" |
|
outattrs["units"] = "%" |
|
outattrs["MemoryOrder"] = "XY" |
|
|
|
|
|
p = from_args(wrapped, copyarg, *args, **kwargs)[copyarg] |
|
|
|
if isinstance(p, DataArray): |
|
# Right dims |
|
outdims[-2:] = p.dims[-2:] |
|
# Left dims |
|
outdims[1:-2] = p.dims[0:-3] |
|
|
|
|
|
outcoords = {} |
|
# Left-most is always low_mid_high |
|
outdims[0] = "low_mid_high" |
|
outcoords["low_mid_high"] = ["low", "mid", "high"] |
|
|
|
out = DataArray(result, name=outname, dims=outdims, coords=outcoords, |
|
attrs=outattrs) |
|
|
|
return out |
|
|
|
return func_wrapper |
|
|
|
|
|
def set_destag_metadata(): |
|
"""A decorator that sets the metadata for the wrapped raw destaggering |
|
function. |
|
|
|
Returns: |
|
|
|
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped |
|
destaggering function output with or without metadata. |
|
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. |
|
|
|
""" |
|
|
|
@wrapt.decorator |
|
def func_wrapper(wrapped, instance, args, kwargs): |
|
do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] |
|
|
|
if do_meta is None: |
|
do_meta = True |
|
|
|
if not xarray_enabled() or not do_meta: |
|
return wrapped(*args, **kwargs) |
|
|
|
result = wrapped(*args, **kwargs) |
|
|
|
# Default dimension names |
|
outdims = ["dim{}".format(i) for i in py3range(result.ndim)] |
|
|
|
destag_args = from_args(wrapped, ("var", "stagger_dim"), |
|
*args, **kwargs) |
|
var = destag_args["var"] |
|
destag_dim = destag_args["stagger_dim"] |
|
|
|
if isinstance(var, DataArray): |
|
|
|
if var.name is not None: |
|
outname = "destag_{}".format(var.name) |
|
else: |
|
outname = "destag_var" |
|
|
|
outattrs = OrderedDict() |
|
outattrs.update(var.attrs) |
|
|
|
outattrs["destag_dim"] = destag_dim |
|
|
|
outdims = [] |
|
outdims += var.dims |
|
|
|
destag_dim_name = outdims[destag_dim] |
|
if destag_dim_name.find("_stag") >= 0: |
|
new_dim_name = destag_dim_name.replace("_stag", "") |
|
else: |
|
if destag_dim >= 0: |
|
new_dim_name = "dim_{}".format(destag_dim) |
|
else: |
|
dim_num = result.ndim + destag_dim |
|
new_dim_name = "dim_{}".format(dim_num) |
|
|
|
outdims[destag_dim] = new_dim_name |
|
|
|
out = DataArray(result, name=outname, dims=outdims, attrs=outattrs) |
|
|
|
return out |
|
|
|
return func_wrapper |
|
|
|
|
|
|
|
|
|
|
|
|