From 07349af907c3ff3b37a7b6bfb9c4779a5deaf635 Mon Sep 17 00:00:00 2001 From: Bill Ladwig Date: Wed, 13 Apr 2016 15:23:58 -0600 Subject: [PATCH] Numerous changes to make DataArray implementation pass all unit tests. Properly uses setuptools to make namespace. --- .../wrf/var => ncl_reference}/etaconv.py | 0 wrf_open/var/setup.py | 1 + wrf_open/var/src/python/wrf/var/__init__.py | 29 +- wrf_open/var/src/python/wrf/var/cape.py | 24 +- wrf_open/var/src/python/wrf/var/config.py | 2 + wrf_open/var/src/python/wrf/var/constants.py | 2 + wrf_open/var/src/python/wrf/var/ctt.py | 22 +- wrf_open/var/src/python/wrf/var/dbz.py | 20 +- wrf_open/var/src/python/wrf/var/decorators.py | 493 +------------ wrf_open/var/src/python/wrf/var/destag.py | 34 +- wrf_open/var/src/python/wrf/var/dewpoint.py | 12 +- wrf_open/var/src/python/wrf/var/extension.py | 109 ++- wrf_open/var/src/python/wrf/var/geoht.py | 12 +- wrf_open/var/src/python/wrf/var/helicity.py | 19 +- wrf_open/var/src/python/wrf/var/interp.py | 310 ++------ .../var/src/python/wrf/var/interputils.py | 182 +++++ wrf_open/var/src/python/wrf/var/latlon.py | 27 +- .../var/src/python/wrf/var/metadecorators.py | 695 ++++++++++++++++++ wrf_open/var/src/python/wrf/var/omega.py | 7 +- wrf_open/var/src/python/wrf/var/precip.py | 3 +- wrf_open/var/src/python/wrf/var/pressure.py | 9 +- wrf_open/var/src/python/wrf/var/projection.py | 2 + wrf_open/var/src/python/wrf/var/psadlookup.py | 3 + wrf_open/var/src/python/wrf/var/pw.py | 9 +- wrf_open/var/src/python/wrf/var/rh.py | 10 +- wrf_open/var/src/python/wrf/var/slp.py | 13 +- wrf_open/var/src/python/wrf/var/temp.py | 20 +- wrf_open/var/src/python/wrf/var/terrain.py | 7 +- wrf_open/var/src/python/wrf/var/times.py | 2 + wrf_open/var/src/python/wrf/var/units.py | 2 + wrf_open/var/src/python/wrf/var/util.py | 344 ++++++--- .../var/src/python/wrf/var/uvdecorator.py | 90 +++ wrf_open/var/src/python/wrf/var/uvmet.py | 77 +- wrf_open/var/src/python/wrf/var/vorticity.py | 29 +- wrf_open/var/src/python/wrf/var/wind.py | 75 +- wrf_open/var/src/python/wrf/var/wrfext.f90 | 2 +- wrf_open/var/test/utests.py | 56 +- 37 files changed, 1764 insertions(+), 989 deletions(-) rename wrf_open/var/{src/python/wrf/var => ncl_reference}/etaconv.py (100%) create mode 100644 wrf_open/var/src/python/wrf/var/interputils.py create mode 100644 wrf_open/var/src/python/wrf/var/metadecorators.py create mode 100644 wrf_open/var/src/python/wrf/var/uvdecorator.py diff --git a/wrf_open/var/src/python/wrf/var/etaconv.py b/wrf_open/var/ncl_reference/etaconv.py similarity index 100% rename from wrf_open/var/src/python/wrf/var/etaconv.py rename to wrf_open/var/ncl_reference/etaconv.py diff --git a/wrf_open/var/setup.py b/wrf_open/var/setup.py index b81ac33..8619331 100755 --- a/wrf_open/var/setup.py +++ b/wrf_open/var/setup.py @@ -19,5 +19,6 @@ numpy.distutils.core.setup( packages = setuptools.find_packages("src/python"), ext_modules = [ext1,ext2], package_dir={"":"src/python"}, + namespace_packages=["wrf"], scripts=[], ) diff --git a/wrf_open/var/src/python/wrf/var/__init__.py b/wrf_open/var/src/python/wrf/var/__init__.py index 557771e..4737bc8 100755 --- a/wrf_open/var/src/python/wrf/var/__init__.py +++ b/wrf_open/var/src/python/wrf/var/__init__.py @@ -1,3 +1,6 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + import warnings from . import config @@ -56,7 +59,7 @@ from . import units from .units import * from . import projection from .projection import * - + __all__ = ["getvar"] __all__.extend(config.__all__) __all__.extend( extension.__all__) @@ -86,7 +89,7 @@ __all__.extend(times.__all__) __all__.extend(pressure.__all__) __all__.extend(units.__all__) __all__.extend(projection.__all__) - + # func is the function to call. kargs are required arguments that should # not be altered by the user _FUNC_MAP = {"cape2d" : get_2dcape, @@ -130,7 +133,7 @@ _FUNC_MAP = {"cape2d" : get_2dcape, "wspd_wdir_uvmet10" : get_uvmet10_wspd_wdir, "ctt" : get_ctt } - + _VALID_KARGS = {"cape2d" : ["missing"], "cape3d" : ["missing"], "dbz" : ["do_variant", "do_liqskin"], @@ -172,7 +175,7 @@ _VALID_KARGS = {"cape2d" : ["missing"], "ctt" : [], "default" : [] } - + _ALIASES = {"cape_2d" : "cape2d", "cape_3d" : "cape3d", "eth" : "theta_e", @@ -190,41 +193,41 @@ _ALIASES = {"cape_2d" : "cape2d", "td" : "dp", "td2" : "dp2m" } - + class ArgumentError(Exception): def __init__(self, msg): self.msg = msg - + def __str__(self): return self.msg - + def _undo_alias(alias): actual = _ALIASES.get(alias, None) if actual is None: return alias else: return actual - + def _check_kargs(var, kargs): for arg in kargs.iterkeys(): if arg not in _VALID_KARGS[var]: raise ArgumentError("'%s' is an invalid keyword " "argument for '%s'" % (arg, var)) - - + + def getvar(wrfnc, var, timeidx=0, method="cat", squeeze=True, cache=None, **kargs): if is_standard_wrf_var(wrfnc, var): return extract_vars(wrfnc, timeidx, var, method, squeeze, cache)[var] - + actual_var = _undo_alias(var) if actual_var not in _VALID_KARGS: raise ArgumentError("'%s' is not a valid variable name" % (var)) - + _check_kargs(actual_var, kargs) return _FUNC_MAP[actual_var](wrfnc,timeidx,**kargs) - + diff --git a/wrf_open/var/src/python/wrf/var/cape.py b/wrf_open/var/src/python/wrf/var/cape.py index 2cc134c..1330a00 100755 --- a/wrf_open/var/src/python/wrf/var/cape.py +++ b/wrf_open/var/src/python/wrf/var/cape.py @@ -1,3 +1,6 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + import numpy as n import numpy.ma as ma @@ -5,22 +8,24 @@ from .extension import computetk,computecape from .destag import destagger from .constants import Constants, ConversionFactors from .util import extract_vars, combine_with -from .decorators import copy_and_set_metadata +from .metadecorators import copy_and_set_metadata __all__ = ["get_2dcape", "get_3dcape"] + @copy_and_set_metadata(copy_varname="T", name="cape_2d", dimnames=combine_with("T", remove_dims=("bottom_top",), insert_before="south_north", - new_dimnames=("mcape_mcin_lcl_lfc",)), + new_dimnames=["mcape_mcin_lcl_lfc"]), description="mcape ; mcin ; lcl ; lfc", units="J/kg ; J/kg ; m ; m", - MemoryOrder = "XY") + MemoryOrder="XY") def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL, method="cat", squeeze=True, cache=None): """Return the 2d fields of cape, cin, lcl, and lfc""" varnames = ("T", "P", "PB", "QVAPOR", "PH","PHB", "HGT", "PSFC") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] p = ncvars["P"] @@ -48,7 +53,7 @@ def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL, cape_res,cin_res = computecape(p_hpa,tk,qv,z,ter,psfc_hpa, missing,i3dflag,ter_follow) - + cape = cape_res[...,0,:,:] cin = cin_res[...,0,:,:] lcl = cin_res[...,1,:,:] @@ -69,18 +74,20 @@ def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL, return ma.masked_values(res, missing) + @copy_and_set_metadata(copy_varname="T", name="cape_3d", dimnames=combine_with("T", insert_before="bottom_top", - new_dimnames=("cape_cin",)), + new_dimnames=["cape_cin"]), description="cape ; cin", units="J kg-1 ; J kg-1", - MemoryOrder = "XY") + MemoryOrder="XY") def get_3dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL, method="cat", squeeze=True, cache=None): """Return the 3d fields of cape and cin""" varnames = ("T", "P", "PB", "QVAPOR", "PH", "PHB", "HGT", "PSFC") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] @@ -119,7 +126,6 @@ def get_3dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL, res[...,0,:,:,:] = cape[:] res[...,1,:,:,:] = cin[:] - #return res return ma.masked_values(res, missing) diff --git a/wrf_open/var/src/python/wrf/var/config.py b/wrf_open/var/src/python/wrf/var/config.py index a7ec3de..27570f4 100644 --- a/wrf_open/var/src/python/wrf/var/config.py +++ b/wrf_open/var/src/python/wrf/var/config.py @@ -1,3 +1,5 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) try: from xarray import DataArray diff --git a/wrf_open/var/src/python/wrf/var/constants.py b/wrf_open/var/src/python/wrf/var/constants.py index 45167ce..d463ed7 100755 --- a/wrf_open/var/src/python/wrf/var/constants.py +++ b/wrf_open/var/src/python/wrf/var/constants.py @@ -1,3 +1,5 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) __all__ = ["Constants", "ConversionFactors"] diff --git a/wrf_open/var/src/python/wrf/var/ctt.py b/wrf_open/var/src/python/wrf/var/ctt.py index 0f06c82..0976d2a 100644 --- a/wrf_open/var/src/python/wrf/var/ctt.py +++ b/wrf_open/var/src/python/wrf/var/ctt.py @@ -1,16 +1,21 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) import numpy as n from .extension import computectt, computetk from .constants import Constants, ConversionFactors from .destag import destagger -from .decorators import convert_units, copy_and_set_metadata +from .decorators import convert_units +from .metadecorators import copy_and_set_metadata from .util import extract_vars __all__ = ["get_ctt"] -@copy_and_set_metadata(copy_varname="T", name="ctt", - description="cloud top temperature") +@copy_and_set_metadata(copy_varname="T", name="ctt", + remove_dims=("bottom_top",), + description="cloud top temperature", + MemoryOrder="XY") @convert_units("temp", "c") def get_ctt(wrfnc, timeidx=0, units="c", method="cat", squeeze=True, cache=None): @@ -18,7 +23,8 @@ def get_ctt(wrfnc, timeidx=0, units="c", method="cat", squeeze=True, """ varnames = ("T", "P", "PB", "PH", "PHB", "HGT", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] @@ -29,7 +35,8 @@ def get_ctt(wrfnc, timeidx=0, units="c", method="cat", squeeze=True, haveqci = 1 try: - icevars = extract_vars(wrfnc, timeidx, "QICE", method, squeeze, cache) + icevars = extract_vars(wrfnc, timeidx, "QICE", + method, squeeze, cache) except KeyError: qice = n.zeros(qv.shape, qv.dtype) haveqci = 0 @@ -37,7 +44,8 @@ def get_ctt(wrfnc, timeidx=0, units="c", method="cat", squeeze=True, qice = icevars["QICE"] * 1000.0 #g/kg try: - cldvars = extract_vars(wrfnc, timeidx, "QCLOUD", method, squeeze, cache) + cldvars = extract_vars(wrfnc, timeidx, "QCLOUD", + method, squeeze, cache) except KeyError: raise RuntimeError("'QCLOUD' not found in NetCDF file") else: @@ -52,6 +60,6 @@ def get_ctt(wrfnc, timeidx=0, units="c", method="cat", squeeze=True, geopt_unstag = destagger(geopt, -3) ght = geopt_unstag / Constants.G - ctt = computectt(p_hpa,tk,qice,qcld,qv,ght,ter,haveqci) + ctt = computectt(p_hpa, tk, qice, qcld, qv, ght, ter, haveqci) return ctt diff --git a/wrf_open/var/src/python/wrf/var/dbz.py b/wrf_open/var/src/python/wrf/var/dbz.py index ff0c67f..8f77955 100755 --- a/wrf_open/var/src/python/wrf/var/dbz.py +++ b/wrf_open/var/src/python/wrf/var/dbz.py @@ -1,9 +1,12 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + import numpy as n from .extension import computedbz,computetk from .constants import Constants -from .util import extract_vars, combine_with -from .decorators import copy_and_set_metadata +from .util import extract_vars +from .metadecorators import copy_and_set_metadata __all__ = ["get_dbz", "get_max_dbz"] @@ -23,7 +26,8 @@ def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False, """ varnames = ("T", "P", "PB", "QVAPOR", "QRAIN") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] @@ -65,13 +69,15 @@ def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False, return computedbz(full_p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin) -@copy_and_set_metadata(copy_varname="T", name="dbz", - dimnames=combine_with("T", remove_dims=("bottom_top",)), + +@copy_and_set_metadata(copy_varname="T", name="max_dbz", + remove_dims=("bottom_top",), description="maximum radar reflectivity", - units="dBz") + units="dBz", + MemoryOrder="XY") def get_max_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False, method="cat", squeeze=True, cache=None): - return n.amax(get_dbz(wrfnc, do_varint, do_liqskin, timeidx, method, + return n.amax(get_dbz(wrfnc, timeidx, do_varint, do_liqskin, method, squeeze, cache), axis=-3) diff --git a/wrf_open/var/src/python/wrf/var/decorators.py b/wrf_open/var/src/python/wrf/var/decorators.py index 8153ee3..2a57a1a 100644 --- a/wrf_open/var/src/python/wrf/var/decorators.py +++ b/wrf_open/var/src/python/wrf/var/decorators.py @@ -1,79 +1,21 @@ -#from functools import wraps -import wrapt -from inspect import getargspec -from collections import OrderedDict +from __future__ import (absolute_import, division, print_function, + unicode_literals) + +from collections import Iterable +import wrapt import numpy as np import numpy.ma as ma from .units import do_conversion, check_units -from .destag import destagger -from .util import (iter_left_indexes, viewitems, extract_vars, - combine_with, either) +from .util import (iter_left_indexes, viewitems, from_args, npvalues) from .config import xarray_enabled if xarray_enabled(): from xarray import DataArray - - __all__ = ["convert_units", "handle_left_iter", "uvmet_left_iter", - "handle_casting" "copy_and_set_metadata", "set_wind_metadata", - "set_latlon_metadata", "set_height_metadata"] - -# TODO: In python 3.x, getargspec is deprecated -def from_args(func, argnames, *args, **kwargs): - """Parses the function args and kargs looking for the desired argument - value. Otherwise, the value is taken from the default keyword argument - using the arg spec. - - """ - if isinstance(argnames, str): - arglist = [argnames] - else: - arglist = argnames - - result = {} - for argname in arglist: - arg_loc = arg_location(func, argname, args, kwargs) - - if arg_loc is not None: - result[argname] = arg_loc[0][arg_loc[1]] - else: - result[argname] = None - - return result - -def arg_location(func, argname, args, kwargs): - """Parses the function args, kargs and signature looking for the desired - argument location (either in args, kargs, or argspec.defaults), - and returns a tuple of argument_sequence, location. - - """ - - argspec = getargspec(func) - - print argspec - - if argname not in argspec.args and argname not in kwargs: - return None - - try: - result_idx = argspec.args.index(argname) - except ValueError: - result_idx = None - - result = None - if (argname in kwargs): - result = kwargs, argname - elif (len(args) > result_idx and result_idx is not None): - result = args, result_idx - else: - arg_idx_from_right = (len(argspec.args) - - argspec.args.index(argname)) - result = list(argspec.defaults), -arg_idx_from_right - - return result + "handle_casting", "handle_extract_transpose"] def convert_units(unit_type, alg_unit): """A decorator that applies unit conversion to a function's output array. @@ -84,7 +26,6 @@ def convert_units(unit_type, alg_unit): - alg_unit - the units that the function returns by default """ -# def convert_decorator(func): @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): @@ -94,9 +35,8 @@ def convert_units(unit_type, alg_unit): # Unit conversion done here return do_conversion(wrapped(*args, **kwargs), unit_type, alg_unit, desired_units) + return func_wrapper - -# return convert_decorator def _calc_out_dims(outvar, left_dims): @@ -115,7 +55,7 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, algorithm is expecting for the reference variable - ref_var_idx - the index in args used as the reference variable for calculating leftmost dimensions - - ref_var_name - the keyword argument name for kargs used as the + - ref_var_name - the keyword argument name for kwargs used as the reference varible for calculating leftmost dimensions - alg_out_fixed_dims - additional fixed dimension sizes for the numerical algorithm (e.g. uvmet has a fixed left dimsize of 2) @@ -128,11 +68,10 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, the output dimensions (e.g. avo) """ -# def indexing_decorator(func): @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): - ignore_args = ignore_args if ignore_args is not None else () - ignore_kargs = ignore_kargs if ignore_kargs is not None else () + _ignore_args = ignore_args if ignore_args is not None else () + _ignore_kargs = ignore_kargs if ignore_kargs is not None else () if ref_var_idx >= 0: ref_var = args[ref_var_idx] @@ -154,16 +93,16 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, # Make the left indexes plus a single slice object # The single slice will handle all the dimensions to # the right (e.g. [1,1,:]) - left_and_slice_idxs = ([x for x in left_idxs] + - [slice(None, None, None)]) + left_and_slice_idxs = tuple([x for x in left_idxs] + [slice(None)]) + # Slice the args if applicable new_args = [arg[left_and_slice_idxs] - if i not in ignore_args else arg + if i not in _ignore_args else arg for i,arg in enumerate(args)] - # Slice the kargs if applicable + # Slice the kwargs if applicable new_kargs = {key:(val[left_and_slice_idxs] - if key not in ignore_kargs else val) + if key not in _ignore_kargs else val) for key,val in viewitems(kwargs)} # Call the numerical routine @@ -219,110 +158,27 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, return output return func_wrapper - -# return indexing_decorator -def uvmet_left_iter(): - """Decorator to handle iterating over leftmost dimensions when using - multiple files and/or multiple times with the uvmet product. - - """ - #def indexing_decorator(func): - @wrapt.decorator - def func_wrapper(wrapped, instance, args, kwargs): - u = args[0] - v = args[1] - lat = args[2] - lon = args[3] - cen_long = args[4] - cone = args[5] - - if u.ndim == lat.ndim: - num_right_dims = 2 - is_3d = False - else: - num_right_dims = 3 - is_3d = True - - is_stag = False - if ((u.shape[-1] != lat.shape[-1]) or - (u.shape[-2] != lat.shape[-2])): - is_stag = True - - if is_3d: - extra_dim_num = u.ndim - 3 - else: - extra_dim_num = u.ndim - 2 - - if is_stag: - u = destagger(u,-1) - v = destagger(v,-2) - - # No special left side iteration, return the function result - if (extra_dim_num == 0): - return wrapped(u,v,lat,lon,cen_long,cone) - - # Start by getting the left-most 'extra' dims - outdims = [u.shape[x] for x in xrange(extra_dim_num)] - extra_dims = list(outdims) # Copy the left-most dims for iteration - - # Append the right-most dimensions - outdims += [2] # For u/v components - - outdims += [u.shape[x] for x in xrange(-num_right_dims,0,1)] - - output = np.empty(outdims, u.dtype) - - for left_idxs in iter_left_indexes(extra_dims): - # Make the left indexes plus a single slice object - # The single slice will handle all the dimensions to - # the right (e.g. [1,1,:]) - left_and_slice_idxs = ([x for x in left_idxs] + - [slice(None, None, None)]) - - - new_u = u[left_and_slice_idxs] - new_v = v[left_and_slice_idxs] - new_lat = lat[left_and_slice_idxs] - new_lon = lon[left_and_slice_idxs] - - # Call the numerical routine - res = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone) - - # Note: The 2D version will return a 3D array with a 1 length - # dimension. Numpy is unable to broadcast this without - # sqeezing first. - res = np.squeeze(res) - - output[left_and_slice_idxs] = res[:] - - - return output - - return func_wrapper - -# return indexing_decorator def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=np.float64): """Decorator to handle casting to/from required dtype used in numerical routine. """ -# def cast_decorator(func): @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): - arg_idxs = arg_idxs if arg_idxs is not None else () - karg_names = karg_names if karg_names is not None else () + _arg_idxs = arg_idxs if arg_idxs is not None else () + _karg_names = karg_names if karg_names is not None else () orig_type = args[ref_idx].dtype new_args = [arg.astype(dtype) - if i in arg_idxs else arg + if i in _arg_idxs else arg for i,arg in enumerate(args)] new_kargs = {key:(val.astype(dtype) - if key in karg_names else val) - for key,val in viewitems()} + if key in _karg_names else val) + for key,val in viewitems(kwargs)} res = wrapped(*new_args, **new_kargs) @@ -336,324 +192,49 @@ def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=np.float64): for arr in res) return func_wrapper - -# return cast_decorator -def _extract_and_transpose(arg): + +def _extract_and_transpose(arg, do_transpose): if xarray_enabled(): if isinstance(arg, DataArray): - arg = arg.values + arg = npvalues(arg) - if isinstance(arg, np.ndarray): - if not arg.flags.F_CONTIGUOUS: - return arg.T + if do_transpose: + if isinstance(arg, np.ndarray): + if not arg.flags.f_contiguous and arg.ndim > 1: + return arg.T return arg + -def handle_extract_transpose(): +def handle_extract_transpose(do_transpose=True): """Decorator to extract the data array from a DataArray and also - transposes if the data is not fortran contiguous. + transposes the view of the data if the data is not fortran contiguous. """ -# def extract_transpose_decorator(func): @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): - new_args = [_extract_and_transpose(arg) for arg in args] + new_args = [_extract_and_transpose(arg, do_transpose) for arg in args] new_kargs = {key:_extract_and_transpose(val) for key,val in viewitems(kwargs)} - + res = wrapped(*new_args, **new_kargs) if isinstance(res, np.ndarray): - if res.flags.F_CONTIGUOUS: + if res.flags.f_contiguous and res.ndim > 1: return res.T - else: - return tuple(x.T if x.flags.F_CONTIGUOUS else x for x in res) + elif isinstance(res, Iterable): + return tuple(x.T if x.flags.f_contiguous and x.ndim > 1 else x + for x in res) return res return func_wrapper -# return extract_transpose_decorator - - -def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, - dimnames=None, coords=None, **fixed_attrs): - """Decorator to set the metadata for a WRF method. - - A cache is inserted/updated to include the extracted variable that will - have its metadata copied. This prevents the variable being extracted more - than once. This extraction can be slow with sequences of large files. - - """ - @wrapt.decorator - def func_wrapper(wrapped, instance, args, kwargs): - if not xarray_enabled(): - return wrapped(*args, **kwargs) - - argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", - "squeeze", "cache", "units"), - *args, **kwargs) - wrfnc = argvars["wrfnc"] - timeidx = argvars["timeidx"] - units = argvars["units"] - method = argvars["method"] - squeeze = argvars["squeeze"] - cache = argvars["cache"] - if cache is None: - cache = {} - - - if (callable(copy_varname)): - copy_varname = copy_varname(wrfnc) - - # 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(wrfnc, timeidx, (copy_varname,), - method, squeeze, cache)[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 = list(args) - new_kargs = dict(kwargs) - cache_argseq, cache_argloc = arg_location(wrapped, "cache", - new_args, new_kargs) - - cache_argseq[cache_argloc] = new_cache - - result = wrapped(*new_args, **new_kargs) - - outname = "" - outdimnames = list() - outcoords = OrderedDict() - outattrs = OrderedDict() - - if copy_varname is not None: - outname = str(var_to_copy.name) - - if dimnames is not None: - if isinstance(combine_with): - outdimnames, outcoords = dimnames(copy_varname) - else: - outdimnames = dimnames - outcoords = coords - else: - outdimnames += var_to_copy.dims - outcoords.update(var_to_copy.coords) - - outattrs.update(var_to_copy.attrs) - - 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: - del outattrs[attr] - - if isinstance(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(wspd_wdir=False): - @wrapt.decorator - def func_wrapper(wrapped, instance, args, kwargs): - if not xarray_enabled(): - return wrapped(*args, **kwargs) - - argvars = from_args(wrapped, ("wrfnc", "timeidx", "units", - "method", "squeeze", "ten_m", "cache"), - *args, **kwargs) - wrfnc = argvars["wrfnc"] - timeidx = argvars["timeidx"] - units = argvars["units"] - method = argvars["method"] - squeeze = argvars["squeeze"] - ten_m = argvars["ten_m"] - cache = argvars["cache"] - if cache is None: - cache = {} - - lat_var = either("XLAT", "XLAT_M")(wrfnc) - xlat_var = extract_vars(wrfnc, timeidx, lat_var, - method, squeeze, cache) - lat = xlat_var[lat_var] - - lon_var = either("XLONG", "XLONG_M") - xlon_var = extract_vars(wrfnc, timeidx, lon_var, - method, squeeze, cache) - lon = xlon_var[lon_var] - - pres_var = either("P", "PRES") - p_var = extract_vars(wrfnc, timeidx, lon_var, method, squeeze, cache) - pres = p_var[pres_var] - - - # Make a copy so we don't modify a user supplied cache - new_cache = dict(cache) - new_cache[lat_var] = lat - new_cache[lon_var] = lon - new_cache[pres_var] = pres - - # Don't modify the original args/kargs. The args need to be a list - # so it can be modified. - new_args = list(args) - new_kargs = dict(kwargs) - cache_argseq, cache_argloc = arg_location(wrapped, "cache", - new_args, new_kargs) - - cache_argseq[cache_argloc] = new_cache - - result = wrapped(*new_args, **new_kargs) - - outcoords = OrderedDict() - outattrs = OrderedDict() - outname = "uvmet" if not ten_m else "uvmet10" - - outdimnames = list(pres.dims) - outcoords.update(pres.coords) - outattrs.update(pres.attrs) - - if not wspd_wdir: - outattrs["description"] = ("earth rotated u,v" if not ten_m else - "10m earth rotated u,v") - if not ten_m: - outdimnames.insert(-3, "u_v") - else: - outdimnames.insert(-2, "u_v") - else: - outattrs["description"] = ("earth rotated wspd,wdir" if not ten_m - else "10m earth rotated wspd,wdir") - if not ten_m: - outdimnames.insert(-3, "wspd_wdir") - else: - outdimnames.insert(-2, "wspd_wdir") - - if units is not None: - outattrs["units"] = units - - return DataArray(result, name=outname, coords=outcoords, - dims=outdimnames, attrs=outattrs) - - return func_wrapper - -def set_latlon_metadata(ij=False): - @wrapt.decorator - def func_wrapper(wrapped, instance, args, kwargs): - if not xarray_enabled(): - return wrapped(*args, **kwargs) - - res = wrapped(*args, **kwargs) - - argnames = ("latitude", "longitude") if not ij else ("i", "j") - outname = "latlon" if not ij else "ij" - - if res.ndim <= 2: - dimnames = (["lat_lon", "i_j"] if not ij - else ["i_j", "lat_lon"]) - else: - dimnames = (["lat_lon", "domain", "i_j"] if not ij - else ["i_j", "domain", "lat_lon"]) - - - argvars = from_args(wrapped, argnames, *args, **kwargs) - - var1 = argvars[argnames[0]] - var2 = argvars[argnames[1]] - - arr1 = np.asarray(var1).ravel() - arr2 = np.asarray(var2).ravel() - - coords = {} - coords[dimnames[0]] = [x for x in zip(arr1, arr2)] - - return DataArray(res, name=outname, dims=dimnames, coords=coords) - - return func_wrapper - -def set_height_metadata(geopt=False): - @wrapt.decorator - def func_wrapper(wrapped, instance, args, kwargs): - if not xarray_enabled(): - return wrapped(*args, **kwargs) - - argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", - "squeeze", "units", "msl", "cache"), - *args, **kwargs) - wrfnc = argvars["wrfnc"] - timeidx = argvars["timeidx"] - units = argvars["units"] - method = argvars["method"] - squeeze = argvars["squeeze"] - msl = argvars["msl"] - cache = argvars["cache"] - 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")(wrfnc) - ht_var = extract_vars(wrfnc, timeidx, ht_metadata_varname, - method, squeeze, cache) - 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_var] = ht_metadata_var - - # Don't modify the original args/kargs. The args need to be a list - # so it can be modified. - new_args = list(args) - new_kargs = dict(kwargs) - cache_argseq, cache_argloc = arg_location(wrapped, "cache", - new_args, new_kargs) - - cache_argseq[cache_argloc] = new_cache - - result = wrapped(*new_args, **new_kargs) - - 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) - - diff --git a/wrf_open/var/src/python/wrf/var/destag.py b/wrf_open/var/src/python/wrf/var/destag.py index 2ab2a6f..fd8e706 100755 --- a/wrf_open/var/src/python/wrf/var/destag.py +++ b/wrf_open/var/src/python/wrf/var/destag.py @@ -1,9 +1,11 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) +from .decorators import handle_extract_transpose -from .util import extract_vars - -__all__ = ["destagger", "destagger_windcomp", "destagger_winds"] +__all__ = ["destagger"] +@handle_extract_transpose(do_transpose=False) def destagger(var, stagger_dim): """ De-stagger the variable. @@ -21,10 +23,11 @@ def destagger(var, stagger_dim): # Dynamically building the range slices to create the appropriate # number of ':'s in the array accessor lists. # For example, for a 3D array, the calculation would be - # result = .5 * (var[:,:,0:stagger_dim_size-2] + var[:,:,1:stagger_dim_size-1]) + # result = .5 * (var[:,:,0:stagger_dim_size-2] + # + var[:,:,1:stagger_dim_size-1]) # for stagger_dim=2. So, full slices would be used for dims 0 and 1, but # dim 2 needs the special slice. - full_slice = slice(None, None, None) + full_slice = slice(None) slice1 = slice(0, stagger_dim_size - 1, 1) slice2 = slice(1, stagger_dim_size, 1) @@ -40,25 +43,4 @@ def destagger(var, stagger_dim): return result -# def destagger_windcomp(wrfnc, comp, timeidx=0, -# method="cat", squeeze=True, cache=None): -# if comp.lower() == "u": -# wrfvar = "U" -# stagdim = -1 -# elif comp.lower() == "v": -# wrfvar = "V" -# stagdim = -2 -# elif comp.lower() == "w": -# wrfvar = "W" -# stagdim = -3 -# -# ncvars = extract_vars(wrfnc, timeidx, wrfvar, method, squeeze, cache) -# wind_data = ncvars[wrfvar] -# -# return destagger(wind_data, stagdim) -# -# def destagger_winds(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): -# return (destagger_windcomp(wrfnc, "u", timeidx, method, squeeze, cache), -# destagger_windcomp(wrfnc, "v", timeidx, method, squeeze, cache), -# destagger_windcomp(wrfnc, "w", timeidx, method, squeeze, cache)) \ No newline at end of file diff --git a/wrf_open/var/src/python/wrf/var/dewpoint.py b/wrf_open/var/src/python/wrf/var/dewpoint.py index 4b55047..2d7f2ae 100755 --- a/wrf_open/var/src/python/wrf/var/dewpoint.py +++ b/wrf_open/var/src/python/wrf/var/dewpoint.py @@ -1,5 +1,9 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + from .extension import computetd -from .decorators import convert_units, copy_and_set_metadata +from .decorators import convert_units +from .metadecorators import copy_and_set_metadata from .util import extract_vars __all__ = ["get_dp", "get_dp_2m"] @@ -11,7 +15,8 @@ def get_dp(wrfnc, timeidx=0, units="c", method="cat", squeeze=True, cache=None): varnames=("P", "PB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) p = ncvars["P"] pb = ncvars["PB"] @@ -30,7 +35,8 @@ def get_dp(wrfnc, timeidx=0, units="c", def get_dp_2m(wrfnc, timeidx=0, units="c", method="cat", squeeze=True, cache=None): varnames=("PSFC", "Q2") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) # Algorithm requires hPa psfc = .01*(ncvars["PSFC"]) diff --git a/wrf_open/var/src/python/wrf/var/extension.py b/wrf_open/var/src/python/wrf/var/extension.py index 47af45f..765961a 100755 --- a/wrf_open/var/src/python/wrf/var/extension.py +++ b/wrf_open/var/src/python/wrf/var/extension.py @@ -1,3 +1,6 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + import numpy as np from .constants import Constants @@ -11,37 +14,42 @@ from ._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d, f_lltoij, f_ijtoll, f_converteta, f_computectt, f_monotonic, f_filter2d, f_vintrp) from ._wrfcape import f_computecape -from .decorators import (handle_left_iter, uvmet_left_iter, - handle_casting, handle_extract_transpose) +from .decorators import (handle_left_iter, handle_casting, + handle_extract_transpose) +from .uvdecorator import uvmet_left_iter __all__ = ["FortranException", "computeslp", "computetk", "computetd", "computerh", "computeavo", "computepvo", "computeeth", "computeuvmet","computeomega", "computetv", "computesrh", "computeuh", "computepw","computedbz","computecape", - "computeij", "computell", "computeeta", "computectt"] + "computeij", "computell", "computeeta", "computectt", + "interp2dxy", "interpz3d", "interp1d", "computeinterpline", + "computevertcross"] class FortranException(Exception): def __call__(self, message): raise self.__class__(message) + @handle_left_iter(3,0, ignore_args=(2,3)) @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() -def interpz3d(data3d, zdata, desiredloc, missingval): - res = f_interpz3d(data3d, - zdata, +def interpz3d(field3d, z, desiredloc, missingval): + res = f_interpz3d(field3d, + z, desiredloc, missingval) return res -@handle_left_iter(3,0) +@handle_left_iter(3,0, ignore_args=(1,)) @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() -def interp2dxy(data3d,xy): - res = f_interp2dxy(data3d, +def interp2dxy(field3d, xy): + res = f_interp2dxy(field3d, xy) return res +@handle_left_iter(1, 0, ignore_args=(2,3)) @handle_casting(arg_idxs=(0,1,2)) @handle_extract_transpose() def interp1d(v_in, z_in, z_out, missingval): @@ -52,19 +60,45 @@ def interp1d(v_in, z_in, z_out, missingval): return res +@handle_left_iter(3, 0, ignore_args=(1,4,3)) +@handle_casting(arg_idxs=(0,)) +@handle_extract_transpose(do_transpose=False) +def computevertcross(field3d, xy, var2dz, z_var2d, missingval): + var2d = np.empty((z_var2d.size, xy.shape[0]), dtype=var2dz.dtype) + var2dtmp = interp2dxy(field3d, xy) + + for i in xrange(xy.shape[0]): + var2d[:,i] = interp1d(var2dtmp[:,i], var2dz[:,i], z_var2d, missingval) + + return var2d + +@handle_left_iter(2, 0, ignore_args=(1,)) +@handle_casting(arg_idxs=(0,)) +@handle_extract_transpose(do_transpose=False) +def computeinterpline(field2d, xy): + + tmp_shape = [1] + [x for x in field2d.shape] + var2dtmp = np.empty(tmp_shape, field2d.dtype) + var2dtmp[0,:,:] = field2d[:,:] + + var1dtmp = interp2dxy(var2dtmp, xy) + + return var1dtmp[0,:] + @handle_left_iter(3,0) @handle_casting(arg_idxs=(0,1,2,3)) @handle_extract_transpose() def computeslp(z, t, p, q): - t_surf = np.zeros((z.shape[-2], z.shape[-1]), np.float64, order="F") - t_sea_level = np.zeros((z.shape[-2], z.shape[-1]), np.float64, order="F") - level = np.zeros((z.shape[-2], z.shape[-1]), np.int32, order="F") + + t_surf = np.zeros(z.shape[0:2], np.float64, order="F") + t_sea_level = np.zeros(z.shape[0:2], np.float64, order="F") + level = np.zeros(z.shape[0:2], np.int32, order="F") res = f_computeslp(z, t, p, q, - t_sea_level, # Should come in with fortran ordering + t_sea_level, t_surf, level, FortranException()) @@ -79,7 +113,8 @@ def computetk(pres, theta): shape = pres.shape res = f_computetk(pres.ravel(order="A"), theta.ravel(order="A")) - res = np.reshape(res, shape) + res = np.reshape(res, shape, order="F") + return res # Note: No left iteration decorator needed with 1D arrays @@ -89,7 +124,7 @@ def computetd(pressure, qv_in): shape = pressure.shape res = f_computetd(pressure.ravel(order="A"), qv_in.ravel(order="A")) - res = np.reshape(res, shape) + res = np.reshape(res, shape, order="F") return res # Note: No decorator needed with 1D arrays @@ -100,7 +135,7 @@ def computerh(qv, q, t): res = f_computerh(qv.ravel(order="A"), q.ravel(order="A"), t.ravel(order="A")) - res = np.reshape(res, shape) + res = np.reshape(res, shape, order="F") return res @handle_left_iter(3,0, ignore_args=(6,7)) @@ -151,15 +186,15 @@ def computeeth(qv, tk, p): @handle_casting(arg_idxs=(0,1,2,3)) @handle_extract_transpose() def computeuvmet(u, v, lat, lon, cen_long, cone): - longca = np.zeros((lat.shape[-2], lat.shape[-1]), np.float64, order="F") - longcb = np.zeros((lon.shape[-2], lon.shape[-1]), np.float64, order="F") + longca = np.zeros((lat.shape[0], lat.shape[1]), np.float64, order="F") + longcb = np.zeros((lon.shape[0], lon.shape[1]), np.float64, order="F") rpd = Constants.PI/180. # Make the 2D array a 3D array with 1 dimension if u.ndim < 3: - u = u.reshape((u.shape[-2], u.shape[-1], 1), order="F") - v = v.reshape((v.shape[-2], v.shape[-1], 1), order="F") + u = u.reshape((u.shape[0], u.shape[1], 1), order="F") + v = v.reshape((v.shape[0], v.shape[1], 1), order="F") # istag will always be false since winds are destaggered already # Missing values don't appear to be used, so setting to 0 @@ -238,9 +273,9 @@ def computesrh(u, v, z, ter, top): @handle_extract_transpose() def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top): - tem1 = np.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), np.float64, + tem1 = np.zeros((u.shape[0], u.shape[1], u.shape[2]), np.float64, order="F") - tem2 = np.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), np.float64, + tem2 = np.zeros((u.shape[0], u.shape[1], u.shape[2]), np.float64, order="F") res = f_computeuh(zp, @@ -261,8 +296,8 @@ def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top): @handle_casting(arg_idxs=(0,1,2,3)) @handle_extract_transpose() def computepw(p, tv, qv, ht): - # Note, dim -3 is height, we only want y and x - zdiff = np.zeros((p.shape[-2], p.shape[-1]), np.float64, order="F") + + zdiff = np.zeros((p.shape[0], p.shape[1]), np.float64, order="F") res = f_computepw(p, tv, qv, @@ -292,19 +327,22 @@ def computedbz(p, tk, qv, qr, qs, qg, sn0, ivarint, iliqskin): @handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_extract_transpose() def computecape(p_hpa, tk, qv, ht, ter, sfp, missing, i3dflag, ter_follow): - flip_cape = np.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), + flip_cape = np.zeros((p_hpa.shape[0], p_hpa.shape[1], p_hpa.shape[2]), np.float64, order="F") - flip_cin = np.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), + flip_cin = np.zeros((p_hpa.shape[0], p_hpa.shape[1], p_hpa.shape[2]), np.float64, order="F") PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables() PSADITMK = PSADITMK.T # The fortran routine needs pressure to be ascending in z-direction, # along with tk,qv,and ht. - flip_p = p_hpa[::-1,:,:] - flip_tk = tk[::-1,:,:] - flip_qv = qv[::-1,:,:] - flip_ht = ht[::-1,:,:] + # The extra mumbo-jumbo is so that the view created by numpy is fortran + # contiguous. 'ascontiguousarray' only works in C ordering, hence the extra + # transposes. + flip_p = np.ascontiguousarray(p_hpa[:,:,::-1].T).T + flip_tk = np.ascontiguousarray(tk[:,:,::-1].T).T + flip_qv = np.ascontiguousarray(qv[:,:,::-1].T).T + flip_ht = np.ascontiguousarray(ht[:,:,::-1].T).T f_computecape(flip_p, flip_tk, @@ -322,11 +360,10 @@ def computecape(p_hpa, tk, qv, ht, ter, sfp, missing, i3dflag, ter_follow): ter_follow, FortranException()) - # Don't need to transpose since we only passed a view to fortran - # Remember to flip cape and cin back to descending p coordinates - cape = flip_cape[::-1,:,:] - cin = flip_cin[::-1,:,:] - + # Need to flip the vertical back to decending pressure with height. + cape = np.ascontiguousarray(flip_cape[:,:,::-1].T).T + cin = np.ascontiguousarray(flip_cin[:,:,::-1].T).T + return (cape, cin) def computeij(map_proj, truelat1, truelat2, stdlon, @@ -396,7 +433,7 @@ def smooth2d(field, passes): else: missing = Constants.DEFAULT_FILL - field_copy = field.copy() + field_copy = field.copy(order="A") field_tmp = np.zeros(field_copy.shape, field_copy.dtype, order="F") f_filter2d(field_copy, diff --git a/wrf_open/var/src/python/wrf/var/geoht.py b/wrf_open/var/src/python/wrf/var/geoht.py index dc40b47..e2095b8 100755 --- a/wrf_open/var/src/python/wrf/var/geoht.py +++ b/wrf_open/var/src/python/wrf/var/geoht.py @@ -1,6 +1,10 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + from .constants import Constants from .destag import destagger -from .decorators import convert_units, set_height_metadata +from .decorators import convert_units +from .metadecorators import set_height_metadata from .util import extract_vars, either __all__ = ["get_geopt", "get_height"] @@ -16,14 +20,16 @@ def _get_geoht(wrfnc, timeidx, height=True, msl=True, varname = either("PH", "GHT")(wrfnc) if varname == "PH": - ph_vars = extract_vars(wrfnc, timeidx, varnames=("PH", "PHB", "HGT")) + ph_vars = extract_vars(wrfnc, timeidx, ("PH", "PHB", "HGT"), + method, squeeze, cache, nometa=True) ph = ph_vars["PH"] phb = ph_vars["PHB"] hgt = ph_vars["HGT"] geopt = ph + phb geopt_unstag = destagger(geopt, -3) else: - ght_vars = extract_vars(wrfnc, timeidx, varnames=("GHT", "HGT_M")) + ght_vars = extract_vars(wrfnc, timeidx, ("GHT", "HGT_M"), + method, squeeze, cache, nometa=True) geopt_unstag = ght_vars["GHT"] * Constants.G hgt = ght_vars["HGT_M"] diff --git a/wrf_open/var/src/python/wrf/var/helicity.py b/wrf_open/var/src/python/wrf/var/helicity.py index 47b86cc..ef6cced 100755 --- a/wrf_open/var/src/python/wrf/var/helicity.py +++ b/wrf_open/var/src/python/wrf/var/helicity.py @@ -1,9 +1,12 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + from .constants import Constants from .extension import computesrh, computeuh from .destag import destagger from .util import extract_vars, extract_global_attrs, either -from .decorators import copy_and_set_metadata +from .metadecorators import copy_and_set_metadata __all__ = ["get_srh", "get_uh"] @@ -15,7 +18,7 @@ def get_srh(wrfnc, timeidx=0, top=3000.0, # Top can either be 3000 or 1000 (for 0-1 srh or 0-3 srh) ncvars = extract_vars(wrfnc, timeidx, ("HGT", "PH", "PHB"), - method, squeeze, cache) + method, squeeze, cache, nometa=True) ter = ncvars["HGT"] ph = ncvars["PH"] @@ -23,11 +26,13 @@ def get_srh(wrfnc, timeidx=0, top=3000.0, # As coded in NCL, but not sure this is possible varname = either("U", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) u = destagger(u_vars[varname], -1) varname = either("V", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) v = destagger(v_vars[varname], -2) geopt = ph + phb @@ -64,11 +69,13 @@ def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0, # As coded in NCL, but not sure this is possible varname = either("U", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) u = destagger(u_vars[varname], -1) varname = either("V", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) v = destagger(v_vars[varname], -2) zp = ph + phb diff --git a/wrf_open/var/src/python/wrf/var/interp.py b/wrf_open/var/src/python/wrf/var/interp.py index 802c91e..89559f8 100755 --- a/wrf_open/var/src/python/wrf/var/interp.py +++ b/wrf_open/var/src/python/wrf/var/interp.py @@ -1,12 +1,17 @@ -from math import floor, ceil +from __future__ import (absolute_import, division, print_function, + unicode_literals) -import numpy as n +import numpy as np import numpy.ma as ma -from .extension import (interpz3d,interp2dxy,interp1d, - smooth2d,monotonic,vintrp) -from .decorators import handle_left_iter, handle_casting +from .extension import (interpz3d, interp2dxy, interp1d, + smooth2d, monotonic, vintrp, computevertcross, + computeinterpline) +from .decorators import (handle_left_iter, handle_casting, + handle_extract_transpose) +from .metadecorators import set_interp_metadata from .util import extract_vars, is_staggered +from .interputils import get_xy, get_xy_z_params from .constants import Constants, ConversionFactors from .terrain import get_terrain from .geoht import get_height @@ -16,148 +21,32 @@ from .pressure import get_pressure __all__ = ["interplevel", "vertcross", "interpline", "vinterp"] # Note: Extension decorator is good enough to handle left dims -def interplevel(data3d, zdata, desiredloc, missingval=Constants.DEFAULT_FILL): +@set_interp_metadata("horiz") +def interplevel(field3d, z, desiredloc, missingval=Constants.DEFAULT_FILL): """Return the horizontally interpolated data at the provided level - data3d - the 3D field to interpolate - zdata - the vertical values (height or pressure) + field3d - the 3D field to interpolate + z - the vertical values (height or pressure) desiredloc - the vertical level to interpolate at (must be same units as zdata) missingval - the missing data value (which will be masked on return) """ - r1 = interpz3d(data3d, zdata, desiredloc, missingval) + r1 = interpz3d(field3d, z, desiredloc, missingval) masked_r1 = ma.masked_values (r1, missingval) return masked_r1 -def _to_positive_idxs(shape, coord): - if (coord[-2] >= 0 and coord[-1] >= 0): - return coord - - return [x if (x >= 0) else shape[i]+x for (i,x) in enumerate(coord) ] - -def _get_xy(xdim, ydim, pivot_point=None, angle=None, - start_point=None, end_point=None): - """Returns the x,y points for the horizontal cross section line. - - xdim - maximum x-dimension - ydim - maximum y-dimension - pivot_point - a pivot point of (south_north, west_east) - (must be used with angle) - angle - the angle through the pivot point in degrees - start_point - a start_point sequence of [south_north1, west_east1] - end_point - an end point sequence of [south_north2, west_east2] - - """ - - # Have a pivot point with an angle to find cross section - if pivot_point is not None and angle is not None: - xp = pivot_point[-1] - yp = pivot_point[-2] - - if (angle > 315.0 or angle < 45.0 - or ((angle > 135.0) and (angle < 225.0))): - - #x = y*slope + intercept - slope = -(360.-angle)/45. - if( angle < 45. ): - slope = angle/45. - if( angle > 135.): - slope = (angle-180.)/45. - - intercept = xp - yp*slope - - # find intersections with domain boundaries - y0 = 0. - x0 = y0*slope + intercept - - if( x0 < 0.): # intersect outside of left boundary - x0 = 0. - y0 = (x0 - intercept)/slope - if( x0 > xdim-1): #intersect outside of right boundary - x0 = xdim-1 - y0 = (x0 - intercept)/slope - y1 = ydim-1. #need to make sure this will be a float? - x1 = y1*slope + intercept - - if( x1 < 0.): # intersect outside of left boundary - x1 = 0. - y1 = (x1 - intercept)/slope - - if( x1 > xdim-1): # intersect outside of right boundary - x1 = xdim-1 - y1 = (x1 - intercept)/slope - else: - # y = x*slope + intercept - slope = (90.-angle)/45. - if( angle > 225. ): - slope = (270.-angle)/45. - intercept = yp - xp*slope - - #find intersections with domain boundaries - x0 = 0. - y0 = x0*slope + intercept - - if( y0 < 0.): # intersect outside of bottom boundary - y0 = 0. - x0 = (y0 - intercept)/slope - - if( y0 > ydim-1): # intersect outside of top boundary - y0 = ydim-1 - x0 = (y0 - intercept)/slope - - x1 = xdim-1. # need to make sure this will be a float? - y1 = x1*slope + intercept - - if( y1 < 0.): # intersect outside of bottom boundary - y1 = 0. - x1 = (y1 - intercept)/slope - - if( y1 > ydim-1):# intersect outside of top boundary - y1 = ydim-1 - x1 = (y1 - intercept)/slope - elif start_point is not None and end_point is not None: - x0 = start_point[-1] - y0 = start_point[-2] - x1 = end_point[-1] - y1 = end_point[-2] - if ( x1 > xdim-1 ): - x1 = xdim - if ( y1 > ydim-1): - y1 = ydim - else: - raise ValueError("invalid combination of None arguments") - - dx = x1 - x0 - dy = y1 - y0 - distance = (dx*dx + dy*dy)**0.5 - npts = int(distance) - dxy = distance/npts - - xy = n.zeros((npts,2), "float") - - dx = dx/npts - dy = dy/npts - - for i in xrange(npts): - xy[i,0] = x0 + i*dx - xy[i,1] = y0 + i*dy - - return xy - -@handle_left_iter(3, 0, ignore_args=(2,3,4,5,6), - ignore_kargs=("missingval", "pivot_point", "angle", - "start_point", "end_point")) -@handle_casting(arg_idxs=(0,1)) -def vertcross(data3d, z, missingval=Constants.DEFAULT_FILL, +@set_interp_metadata("cross") +def vertcross(field3d, z, missingval=Constants.DEFAULT_FILL, pivot_point=None, angle=None, - start_point=None, end_point=None): + start_point=None, end_point=None, + cache=None): """Return the vertical cross section for a 3D field, interpolated to a verical plane defined by a horizontal line. Arguments: - data3d - a 3D data field + field3d - a 3D data field z - 3D height field pivot_point - a pivot point of (south_north,west_east) (must be used with angle) @@ -167,70 +56,27 @@ def vertcross(data3d, z, missingval=Constants.DEFAULT_FILL, """ - if pivot_point is not None: - pos_pivot = _to_positive_idxs(z.shape[-2:], pivot_point) - else: - pos_pivot = pivot_point + try: + xy = cache["xy"] + var2dz = cache["var2dz"] + z_var2d = cache["z_var2d"] + except (KeyError, TypeError): + xy, var2dz, z_var2d = get_xy_z_params(z, pivot_point, angle, + start_point, end_point) - if start_point is not None: - pos_start = _to_positive_idxs(z.shape[-2:], start_point) - else: - pos_start = start_point - - if end_point is not None: - pos_end = _to_positive_idxs(z.shape[-2:], end_point) - else: - pos_end = start_point - - xdim = z.shape[-1] - ydim = z.shape[-2] - - xy = _get_xy(xdim, ydim, pos_pivot, angle, pos_start, pos_end) - - # Interp z - var2dz = interp2dxy(z, xy) - - # interp to constant z grid - if(var2dz[0,0] > var2dz[-1,0]): # monotonically decreasing coordinate - z_max = floor(n.amax(z) / 10) * 10 # bottom value - z_min = ceil(n.amin(z) / 10) * 10 # top value - dz = 10 - nlevels = int((z_max-z_min) / dz) - z_var2d = n.zeros((nlevels), dtype=z.dtype) - z_var2d[0] = z_max - dz = -dz - else: - z_max = n.amax(z) - z_min = 0. - dz = 0.01 * z_max - nlevels = int(z_max / dz) - z_var2d = n.zeros((nlevels), dtype=z.dtype) - z_var2d[0] = z_min + res = computevertcross(field3d, xy, var2dz, z_var2d, missingval) - for i in xrange(1,nlevels): - z_var2d[i] = z_var2d[0] + i*dz - - #interp the variable - - var2d = n.zeros((nlevels, xy.shape[0]),dtype=var2dz.dtype) - var2dtmp = interp2dxy(data3d, xy) - - for i in xrange(xy.shape[0]): - var2d[:,i] = interp1d(var2dtmp[:,i], var2dz[:,i], z_var2d, missingval) - - return ma.masked_values(var2d, missingval) + return ma.masked_values(res, missingval) + -@handle_left_iter(2, 0, ignore_args=(1,2,3,4), - ignore_kargs=("pivot_point", "angle", - "start_point", "end_point")) -@handle_casting(arg_idxs=(0,)) -def interpline(data2d, pivot_point=None, +@set_interp_metadata("line") +def interpline(field2d, pivot_point=None, angle=None, start_point=None, - end_point=None): + end_point=None, cache=None): """Return the 2D field interpolated along a line. Arguments: - var2d - a 2D data field + field2d - a 2D data field pivot_point - a pivot point of (south_north,west_east) angle - the angle through the pivot point in degrees start_point - a start_point tuple of (south_north1,west_east1) @@ -238,38 +84,18 @@ def interpline(data2d, pivot_point=None, """ - if pivot_point is not None: - pos_pivot = _to_positive_idxs(data2d.shape[-2:], pivot_point) - else: - pos_pivot = pivot_point + try: + xy = cache["xy"] + except (KeyError, TypeError): + xy = get_xy(field2d, pivot_point, angle, start_point, end_point) - if start_point is not None: - pos_start = _to_positive_idxs(data2d.shape[-2:], start_point) - else: - pos_start = start_point - - if end_point is not None: - pos_end = _to_positive_idxs(data2d.shape[-2:], end_point) - else: - pos_end = start_point - - tmp_shape = [1] + [x for x in data2d.shape] - - var2dtmp = n.zeros(tmp_shape, data2d.dtype) - - xdim = data2d.shape[-1] - ydim = data2d.shape[-2] - - xy = _get_xy(xdim, ydim, pos_pivot, angle, pos_start, pos_end) - - var2dtmp[0,:,:] = data2d[:,:] - - var1dtmp = interp2dxy(var2dtmp, xy) - - return var1dtmp[0,:] + return computeinterpline(field2d, xy) + +@set_interp_metadata("vinterp") def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, - field_type=None, log_p=False): + field_type=None, log_p=False, timeidx=-1, method="cat", + squeeze=True, cache=None): # Remove case sensitivity field_type = field_type.lower() if field_type is not None else "none" vert_coord = vert_coord.lower() if vert_coord is not None else "none" @@ -301,8 +127,8 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, sclht = rgas*256./9.81 # interp_levels might be a list or tuple, make a numpy array - if not isinstance(interp_levels, n.ndarray): - interp_levels = n.array(interp_levels, "float64") + if not isinstance(interp_levels, np.ndarray): + interp_levels = np.asarray(interp_levels, np.float64) # TODO: Check if field is staggered if is_staggered(field, wrfnc): @@ -327,19 +153,26 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, icase = icase_lookup[field_type] # Extract vriables - timeidx = -1 # Should this be an argument? - ncvars = extract_vars(wrfnc, timeidx, varnames=("PSFC", "QVAPOR", "F")) + #timeidx = -1 # Should this be an argument? + ncvars = extract_vars(wrfnc, timeidx, ("PSFC", "QVAPOR", "F"), + method, squeeze, cache, nometa=True) sfp = ncvars["PSFC"] * ConversionFactors.PA_TO_HPA qv = ncvars["QVAPOR"] coriolis = ncvars["F"] - terht = get_terrain(wrfnc, timeidx) - t = get_theta(wrfnc, timeidx) - tk = get_temp(wrfnc, timeidx, units="k") - p = get_pressure(wrfnc, timeidx, units="pa") - ght = get_height(wrfnc, timeidx, msl=True) - ht_agl = get_height(wrfnc, timeidx, msl=False) + terht = get_terrain(wrfnc, timeidx, units="m", + method=method, squeeze=squeeze, cache=cache) + t = get_theta(wrfnc, timeidx, units="k", + method=method, squeeze=squeeze, cache=cache) + tk = get_temp(wrfnc, timeidx, units="k", + method=method, squeeze=squeeze, cache=cache) + p = get_pressure(wrfnc, timeidx, units="pa", + method=method, squeeze=squeeze, cache=cache) + ght = get_height(wrfnc, timeidx, msl=True, units="m", + method=method, squeeze=squeeze, cache=cache) + ht_agl = get_height(wrfnc, timeidx, msl=False, units="m", + method=method, squeeze=squeeze, cache=cache) smsfp = smooth2d(sfp, 3) @@ -352,11 +185,11 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, elif vert_coord == "ght_msl": vcor = 2 - vcord_array = n.exp(-ght/sclht) + vcord_array = np.exp(-ght/sclht) elif vert_coord == "ght_agl": vcor = 3 - vcord_array = n.exp(-ht_agl/sclht) + vcord_array = np.exp(-ht_agl/sclht) elif vert_coord in ("theta", "th"): vcor = 4 @@ -389,7 +222,7 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, icase = 0 # Set the missing value - if isinstance(field, n.ma.MaskedArray): + if isinstance(field, ma.MaskedArray): missing = field.fill_value else: missing = Constants.DEFAULT_FILL @@ -400,6 +233,23 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, return ma.masked_values(res, missing) +# Thin wrappers around fortran calls which allow for metadata +# Move to the new routines module +# TODO: Rename after the extensions are renamed +@set_interp_metadata("horiz") +def wrap_interpz3d(field3d, z, desiredloc, missingval): + return interpz3d(field3d, z, desiredloc, missingval) + + +@set_interp_metadata("2dxy") +def wrap_interp2dxy(field3d, xy): + return interp2dxy(field3d, xy) + + +@set_interp_metadata("1d") +def wrap_interp1d(v_in, z_in, z_out, missingval): + return interp1d(v_in, z_in, z_out, missingval) + diff --git a/wrf_open/var/src/python/wrf/var/interputils.py b/wrf_open/var/src/python/wrf/var/interputils.py new file mode 100644 index 0000000..67255de --- /dev/null +++ b/wrf_open/var/src/python/wrf/var/interputils.py @@ -0,0 +1,182 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + +from math import floor, ceil + +import numpy as np + +from .extension import interp2dxy + +__all__ = ["to_positive_idxs", "calc_xy", "get_xy_z_params", "get_xy"] + +def to_positive_idxs(shape, coord): + if (coord[-2] >= 0 and coord[-1] >= 0): + return coord + + return [x if (x >= 0) else shape[i]+x for (i,x) in enumerate(coord) ] + +def calc_xy(xdim, ydim, pivot_point=None, angle=None, + start_point=None, end_point=None): + """Returns the x,y points for the horizontal cross section line. + + xdim - maximum x-dimension + ydim - maximum y-dimension + pivot_point - a pivot point of (south_north, west_east) + (must be used with angle) + angle - the angle through the pivot point in degrees + start_point - a start_point sequence of [south_north1, west_east1] + end_point - an end point sequence of [south_north2, west_east2] + + """ + + # Have a pivot point with an angle to find cross section + if pivot_point is not None and angle is not None: + xp = pivot_point[-1] + yp = pivot_point[-2] + + if (angle > 315.0 or angle < 45.0 + or ((angle > 135.0) and (angle < 225.0))): + + #x = y*slope + intercept + slope = -(360.-angle)/45. + if( angle < 45. ): + slope = angle/45. + if( angle > 135.): + slope = (angle-180.)/45. + + intercept = xp - yp*slope + + # find intersections with domain boundaries + y0 = 0. + x0 = y0*slope + intercept + + if( x0 < 0.): # intersect outside of left boundary + x0 = 0. + y0 = (x0 - intercept)/slope + if( x0 > xdim-1): #intersect outside of right boundary + x0 = xdim-1 + y0 = (x0 - intercept)/slope + y1 = ydim-1. #need to make sure this will be a float? + x1 = y1*slope + intercept + + if( x1 < 0.): # intersect outside of left boundary + x1 = 0. + y1 = (x1 - intercept)/slope + + if( x1 > xdim-1): # intersect outside of right boundary + x1 = xdim-1 + y1 = (x1 - intercept)/slope + else: + # y = x*slope + intercept + slope = (90.-angle)/45. + if( angle > 225. ): + slope = (270.-angle)/45. + intercept = yp - xp*slope + + #find intersections with domain boundaries + x0 = 0. + y0 = x0*slope + intercept + + if( y0 < 0.): # intersect outside of bottom boundary + y0 = 0. + x0 = (y0 - intercept)/slope + + if( y0 > ydim-1): # intersect outside of top boundary + y0 = ydim-1 + x0 = (y0 - intercept)/slope + + x1 = xdim-1. # need to make sure this will be a float? + y1 = x1*slope + intercept + + if( y1 < 0.): # intersect outside of bottom boundary + y1 = 0. + x1 = (y1 - intercept)/slope + + if( y1 > ydim-1):# intersect outside of top boundary + y1 = ydim-1 + x1 = (y1 - intercept)/slope + elif start_point is not None and end_point is not None: + x0 = start_point[-1] + y0 = start_point[-2] + x1 = end_point[-1] + y1 = end_point[-2] + if ( x1 > xdim-1 ): + x1 = xdim + if ( y1 > ydim-1): + y1 = ydim + else: + raise ValueError("invalid start/end or pivot/angle arguments") + + dx = x1 - x0 + dy = y1 - y0 + distance = (dx*dx + dy*dy)**0.5 + npts = int(distance) + dxy = distance/npts + + xy = np.zeros((npts,2), "float") + + dx = dx/npts + dy = dy/npts + + for i in xrange(npts): + xy[i,0] = x0 + i*dx + xy[i,1] = y0 + i*dy + + return xy + +def get_xy_z_params(z, pivot_point=None, angle=None, + start_point=None, end_point=None): + + xy = get_xy(z, pivot_point, angle, start_point, end_point) + + # Interp z + var2dz = interp2dxy(z, xy) + + extra_dim_num = z.ndim - 3 + idx1 = tuple([0]*extra_dim_num + [0,0]) + idx2 = tuple([0]*extra_dim_num + [-1,0]) + + # interp to constant z grid + if(var2dz[idx1] > var2dz[idx2]): # monotonically decreasing coordinate + z_max = floor(np.amax(z) / 10) * 10 # bottom value + z_min = ceil(np.amin(z) / 10) * 10 # top value + dz = 10 + nlevels = int((z_max-z_min) / dz) + z_var2d = np.zeros((nlevels), dtype=z.dtype) + z_var2d[0] = z_max + dz = -dz + else: + z_max = np.amax(z) + z_min = 0. + dz = 0.01 * z_max + nlevels = int(z_max / dz) + z_var2d = np.zeros((nlevels), dtype=z.dtype) + z_var2d[0] = z_min + + for i in xrange(1,nlevels): + z_var2d[i] = z_var2d[0] + i*dz + + return xy, var2dz, z_var2d + +def get_xy(var, pivot_point=None, angle=None, start_point=None, end_point=None): + if pivot_point is not None: + pos_pivot = to_positive_idxs(var.shape[-2:], pivot_point) + else: + pos_pivot = pivot_point + + if start_point is not None: + pos_start = to_positive_idxs(var.shape[-2:], start_point) + else: + pos_start = start_point + + if end_point is not None: + pos_end = to_positive_idxs(var.shape[-2:], end_point) + else: + pos_end = start_point + + xdim = var.shape[-1] + ydim = var.shape[-2] + + xy = calc_xy(xdim, ydim, pos_pivot, angle, pos_start, pos_end) + + return xy diff --git a/wrf_open/var/src/python/wrf/var/latlon.py b/wrf_open/var/src/python/wrf/var/latlon.py index 912965e..c709458 100755 --- a/wrf_open/var/src/python/wrf/var/latlon.py +++ b/wrf_open/var/src/python/wrf/var/latlon.py @@ -1,3 +1,6 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + from collections import Iterable import numpy as np @@ -6,18 +9,18 @@ from .config import xarray_enabled from .constants import Constants from .extension import computeij, computell from .util import (extract_vars, extract_global_attrs, - either, _is_moving_domain, _is_multi_time, + either, _is_moving_domain, _is_multi_time_req, iter_left_indexes) -from .decorators import set_latlon_metadata +from .metadecorators import set_latlon_metadata if xarray_enabled(): from xarray import DataArray __all__ = ["get_lat", "get_lon", "get_ij", "get_ll"] -def _lat_varname(stagger): +def _lat_varname(wrfnc, stagger): if stagger is None or stagger.lower() == "m": - varname = either("XLAT", "XLAT_M") + varname = either("XLAT", "XLAT_M")(wrfnc) elif stagger.lower() == "u" or stagger.lower() == "v": varname = "XLAT_{}".format(stagger.upper()) else: @@ -25,9 +28,9 @@ def _lat_varname(stagger): return varname -def _lon_varname(stagger): +def _lon_varname(wrfnc, stagger): if stagger is None or stagger.lower() == "m": - varname = either("XLONG", "XLONG_M") + varname = either("XLONG", "XLONG_M")(wrfnc) elif stagger.lower() == "u" or stagger.lower() == "v": varname = "XLONG_{}".format(stagger.upper()) else: @@ -38,16 +41,18 @@ def _lon_varname(stagger): def get_lat(wrfnc, timeidx=0, stagger=None, method="cat", squeeze=True, cache=None): - varname = _lat_varname(stagger) - lat_var = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + varname = _lat_varname(wrfnc, stagger) + lat_var = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=False) return lat_var[varname] def get_lon(wrfnc, timeidx=0, stagger=None, method="cat", squeeze=True, cache=None): - varname = _lon_varname(stagger) - lon_var = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + varname = _lon_varname(wrfnc, stagger) + lon_var = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=False) return lon_var[varname] @@ -83,7 +88,7 @@ def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache): lat_timeidx = timeidx # Only need all the lats/lons if it's a moving domain file/files - if _is_multi_time(timeidx): + if _is_multi_time_req(timeidx): if not _is_moving_domain(wrfnc, latvar=latvar, lonvar=lonvar): lat_timeidx = 0 diff --git a/wrf_open/var/src/python/wrf/var/metadecorators.py b/wrf_open/var/src/python/wrf/var/metadecorators.py new file mode 100644 index 0000000..1d0061f --- /dev/null +++ b/wrf_open/var/src/python/wrf/var/metadecorators.py @@ -0,0 +1,695 @@ +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 .util import (viewkeys, viewitems, extract_vars, + combine_with, either, from_args, arg_location, + _is_coord_var, XYCoord, npvalues) +from .interputils import get_xy_z_params, get_xy +from .config import xarray_enabled + +if xarray_enabled(): + from xarray import DataArray + +__all__ = ["copy_and_set_metadata", "set_wind_metadata", + "set_latlon_metadata", "set_height_metadata", + "set_interp_metadata"] + +def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, + remove_dims=None, dimnames=None, + coords=None, **fixed_attrs): + """Decorator to set the metadata for a WRF method. + + A cache is inserted/updated to include the extracted variable that will + have its metadata copied. This prevents the variable being extracted more + than once. This extraction can be slow with sequences of large files. + + """ + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + if not xarray_enabled(): + return wrapped(*args, **kwargs) + + argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", + "squeeze", "cache", "units"), + *args, **kwargs) + wrfnc = argvars["wrfnc"] + timeidx = argvars["timeidx"] + units = argvars["units"] + method = argvars["method"] + squeeze = argvars["squeeze"] + cache = argvars["cache"] + if cache is None: + cache = {} + + # Note: can't modify nonlocal var + if (callable(copy_varname)): + _copy_varname = copy_varname(wrfnc) + 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(wrfnc, timeidx, (_copy_varname,), + method, squeeze, cache, + nometa=False)[_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 = unicode(var_to_copy.name) + + if dimnames is not None: + if isinstance(dimnames, combine_with): + outdimnames, outcoords = dimnames(var_to_copy) + else: + 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): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + if not xarray_enabled(): + return wrapped(*args, **kwargs) + + argvars = from_args(wrapped, ("wrfnc", "timeidx", "units", + "method", "squeeze", "ten_m", "cache"), + *args, **kwargs) + wrfnc = argvars["wrfnc"] + timeidx = argvars["timeidx"] + units = argvars["units"] + method = argvars["method"] + squeeze = argvars["squeeze"] + ten_m = argvars["ten_m"] + cache = argvars["cache"] + if cache is None: + cache = {} + + if isinstance(copy_varname, either): + _copy_varname = copy_varname(wrfnc) + else: + _copy_varname = copy_varname + + copy_var = extract_vars(wrfnc, timeidx, _copy_varname, + method, squeeze, cache, + nometa=False)[_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: + pass + + elif not wspd_wdir: + if not two_d: + outdimnames.insert(-3, "u_v") + else: + outdimnames.insert(-2, "u_v") + outattrs["MemoryOrder"] = "XY" + outcoords["u_v"] = ["u", "v"] + else: + if not two_d: + outdimnames.insert(-3, "wspd_wdir") + else: + outdimnames.insert(-2, "wspd_wdir") + outattrs["MemoryOrder"] = "XY" + + outcoords["wspd_wdir"] = ["wspd", "wdir"] + + if units is not None: + outattrs["units"] = units + + outname = name + outattrs["description"] = description + + return DataArray(result, name=outname, coords=outcoords, + dims=outdimnames, attrs=outattrs) + + return func_wrapper + +def set_latlon_metadata(ij=False): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + if not xarray_enabled(): + return wrapped(*args, **kwargs) + + res = wrapped(*args, **kwargs) + + argnames = ("latitude", "longitude") if not ij else ("i", "j") + outname = "latlon" if not ij else "ij" + + if res.ndim <= 2: + dimnames = (["lat_lon", "i_j"] if not ij + else ["i_j", "lat_lon"]) + else: + dimnames = (["lat_lon", "domain", "i_j"] if not ij + else ["i_j", "domain", "lat_lon"]) + + + argvars = from_args(wrapped, argnames, *args, **kwargs) + + var1 = argvars[argnames[0]] + var2 = argvars[argnames[1]] + + arr1 = np.asarray(var1).ravel() + arr2 = np.asarray(var2).ravel() + + coords = {} + coords[dimnames[0]] = [x for x in zip(arr1, arr2)] + + return DataArray(res, name=outname, dims=dimnames, coords=coords) + + return func_wrapper + +def set_height_metadata(geopt=False): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + if not xarray_enabled(): + return wrapped(*args, **kwargs) + + argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", + "squeeze", "units", "msl", "cache"), + *args, **kwargs) + wrfnc = argvars["wrfnc"] + timeidx = argvars["timeidx"] + units = argvars["units"] + method = argvars["method"] + squeeze = argvars["squeeze"] + msl = argvars["msl"] + cache = argvars["cache"] + + 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")(wrfnc) + ht_var = extract_vars(wrfnc, timeidx, ht_metadata_varname, + method, squeeze, cache, nometa=False) + 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): + argvars = from_args(wrapped, ("field3d", "z", "desiredloc", + "missingval"), + *args, **kwargs) + + field3d = argvars["field3d"] + z = argvars["z"] + desiredloc = argvars["desiredloc"] + missingval = argvars["missingval"] + + 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)) + + if isinstance(field3d, DataArray): + outcoords = OrderedDict() + outattrs = OrderedDict() + outdimnames = list(field3d.dims) + outcoords.update(field3d.coords) + outdimnames.remove(field3d.dims[-3]) + del outcoords[field3d.dims[-3]] + outattrs.update(field3d.attrs) + outname = "{0}_{1}".format(field3d.name, levelstr) + + else: + outname = "field3d_{0}".format(levelstr) + outattrs = OrderedDict() + + outattrs["PlotLevelID"] = 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): + argvars = from_args(wrapped, ("field3d", "z", "missingval", + "pivot_point", "angle", + "start_point", "end_point", + "cache"), + *args, **kwargs) + + field3d = argvars["field3d"] + z = argvars["z"] + missingval = argvars["missingval"] + pivot_point = argvars["pivot_point"] + angle = argvars["angle"] + start_point = argvars["start_point"] + end_point = argvars["end_point"] + cache = argvars["cache"] + + xy, var2dz, z_var2d = get_xy_z_params(npvalues(z), pivot_point, angle, + start_point, end_point) + + # 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 = "cross-section: ({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 xrange(-3,0,1): + outdimnames.remove(field3d.dims[i]) + del outcoords[field3d.dims[i]] + + + # Delete any lat,lon coords + for key in viewkeys(outcoords): + if _is_coord_var(key): + del outcoords[key] + + outdimnames.append("vertical") + outdimnames.append("xy") + outattrs.update(field3d.attrs) + + outname = "{0}_cross".format(field3d.name) + + for key in ("MemoryOrder",): + try: + del outattrs[key] + except KeyError: + pass + + outcoords["xy"] = [XYCoord(xy[i,0], xy[i,1]) + for i in xrange(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): + argvars = from_args(wrapped, ("field2d", "pivot_point", "angle", + "start_point", "end_point", "cache"), + *args, **kwargs) + + field2d = argvars["field2d"] + pivot_point = argvars["pivot_point"] + angle = argvars["angle"] + start_point = argvars["start_point"] + end_point = argvars["end_point"] + cache = argvars["cache"] + + if cache is None: + cache = {} + + xy = get_xy(field2d, pivot_point, angle, start_point, end_point) + + # 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 xrange(-2,0,1): + outdimnames.remove(field2d.dims[i]) + del outcoords[field2d.dims[i]] + + # Delete any lat,lon coords + for key in viewkeys(outcoords): + if _is_coord_var(key): + del outcoords[key] + + outdimnames.append("xy") + outattrs.update(field2d.attrs) + + outname = "{0}_line".format(field2d.name) + + for key in ("MemoryOrder",): + try: + del outattrs[key] + except KeyError: + pass + + outcoords["xy"] = [XYCoord(xy[i,0], xy[i,1]) + for i in xrange(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): + argvars = from_args(wrapped, ("wrfnc", "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]) + del outcoords[field.dims[-3]] + + 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): + argvars = from_args(wrapped, ("field3d", "xy"), + *args, **kwargs) + + field3d = argvars["field3d"] + xy = argvars["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) + + # Dims are (...,xy,z) + if isinstance(field3d, DataArray): + outcoords = OrderedDict() + outattrs = OrderedDict() + outdimnames = list(field3d.dims) + outcoords.update(field3d.coords) + for i in xrange(-2,0,1): + outdimnames.remove(field3d.dims[i]) + del outcoords[field3d.dims[i]] + + outdimnames[-2] = "xy" + outattrs.update(field3d.attrs) + + outname = "{0}_xy".format(field3d.name) + + outcoords["xy"] = [XYCoord(xy[i,0], xy[i,1]) + for i in xrange(xy.shape[-2])] + + for key in ("MemoryOrder",): + try: + del outattrs[key] + except KeyError: + pass + + else: + outname = "field3d_xy" + + outattrs["Orientation"] = cross_str + + return DataArray(result, name=outname, dims=outdimnames, + coords=outcoords, attrs=outattrs) + + +def _set_1d_meta(wrapped, instance, args, kwargs): + argvars = from_args(wrapped, ("v_in", "z_in", "z_out", "missingval"), + *args, **kwargs) + + v_in = argvars["v_in"] + z_in = argvars["z_in"] + z_out = argvars["z_out"] + missingval = argvars["missingval"] + + result = wrapped(*args, **kwargs) + + # Dims are (...,xy,z) + if isinstance(v_in, DataArray): + outcoords = OrderedDict() + outattrs = OrderedDict() + outdimnames = list(v_in.dims) + outcoords.update(v_in.coords) + + + outdimnames.remove(v_in.dims[-1]) + del outcoords[v_in.dims[-1]] + outdimnames.append("z") + outname = "{0}_z".format(v_in.name) + outcoords["z"] = z_out + + outattrs.update(v_in.attrs) + outattrs["_FillValue"] = missingval + outattrs["missing_value"] = missingval + + else: + outname = "v_in_z" + + + return DataArray(result, name=outname, dims=outdimnames, + coords=outcoords, attrs=outattrs) + + + + +def set_interp_metadata(interp_type): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + if not xarray_enabled(): + 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) + return func_wrapper diff --git a/wrf_open/var/src/python/wrf/var/omega.py b/wrf_open/var/src/python/wrf/var/omega.py index 6bf3235..7e545c0 100755 --- a/wrf_open/var/src/python/wrf/var/omega.py +++ b/wrf_open/var/src/python/wrf/var/omega.py @@ -1,9 +1,11 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) from .constants import Constants from .destag import destagger from .extension import computeomega,computetk from .util import extract_vars -from .decorators import copy_and_set_metadata +from .metadecorators import copy_and_set_metadata __all__ = ["get_omega"] @@ -12,7 +14,8 @@ __all__ = ["get_omega"] units="Pa/s") def get_omega(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): varnames=("T", "P", "W", "PB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] p = ncvars["P"] w = ncvars["W"] diff --git a/wrf_open/var/src/python/wrf/var/precip.py b/wrf_open/var/src/python/wrf/var/precip.py index 332ea2d..a7f4eea 100755 --- a/wrf_open/var/src/python/wrf/var/precip.py +++ b/wrf_open/var/src/python/wrf/var/precip.py @@ -1,4 +1,5 @@ - +from __future__ import (absolute_import, division, print_function, + unicode_literals) from .util import extract_vars diff --git a/wrf_open/var/src/python/wrf/var/pressure.py b/wrf_open/var/src/python/wrf/var/pressure.py index b528f77..5b0b204 100755 --- a/wrf_open/var/src/python/wrf/var/pressure.py +++ b/wrf_open/var/src/python/wrf/var/pressure.py @@ -1,5 +1,8 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) -from .decorators import convert_units, copy_and_set_metadata +from .decorators import convert_units +from .metadecorators import copy_and_set_metadata from .util import extract_vars, either __all__ = ["get_pressure", "get_pressure_hpa"] @@ -13,13 +16,13 @@ def get_pressure(wrfnc, timeidx=0, units="pa", varname = either("P", "PRES")(wrfnc) if varname == "P": p_vars = extract_vars(wrfnc, timeidx, ("P", "PB"), - method, squeeze, cache) + method, squeeze, cache, nometa=True) p = p_vars["P"] pb = p_vars["PB"] pres = p + pb else: pres = extract_vars(wrfnc, timeidx, "PRES", - method, squeeze, cache)["PRES"] + method, squeeze, cache, nometa=True)["PRES"] return pres diff --git a/wrf_open/var/src/python/wrf/var/projection.py b/wrf_open/var/src/python/wrf/var/projection.py index a16306d..972323c 100644 --- a/wrf_open/var/src/python/wrf/var/projection.py +++ b/wrf_open/var/src/python/wrf/var/projection.py @@ -1,3 +1,5 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) import numpy as np import math diff --git a/wrf_open/var/src/python/wrf/var/psadlookup.py b/wrf_open/var/src/python/wrf/var/psadlookup.py index 730144f..2565f9a 100755 --- a/wrf_open/var/src/python/wrf/var/psadlookup.py +++ b/wrf_open/var/src/python/wrf/var/psadlookup.py @@ -1,3 +1,6 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + import numpy as n __all__ = ["get_lookup_tables"] diff --git a/wrf_open/var/src/python/wrf/var/pw.py b/wrf_open/var/src/python/wrf/var/pw.py index 667bd53..79249b9 100755 --- a/wrf_open/var/src/python/wrf/var/pw.py +++ b/wrf_open/var/src/python/wrf/var/pw.py @@ -1,17 +1,22 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) from .extension import computepw,computetv,computetk from .constants import Constants from .util import extract_vars -from .decorators import copy_and_set_metadata +from .metadecorators import copy_and_set_metadata __all__ = ["get_pw"] @copy_and_set_metadata(copy_varname="T", name="pw", + remove_dims=("bottom_top",), description="precipitable water", + MemoryOrder="XY", units="kg m-2") def get_pw(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): varnames=("T", "P", "PB", "PH", "PHB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] p = ncvars["P"] diff --git a/wrf_open/var/src/python/wrf/var/rh.py b/wrf_open/var/src/python/wrf/var/rh.py index 65e245a..2514262 100755 --- a/wrf_open/var/src/python/wrf/var/rh.py +++ b/wrf_open/var/src/python/wrf/var/rh.py @@ -1,8 +1,10 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) from .constants import Constants from .extension import computerh, computetk from .util import extract_vars -from .decorators import copy_and_set_metadata +from .metadecorators import copy_and_set_metadata __all__ = ["get_rh", "get_rh_2m"] @@ -11,7 +13,8 @@ __all__ = ["get_rh", "get_rh_2m"] delete_attrs=("units",)) def get_rh(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): varnames=("T", "P", "PB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] @@ -30,7 +33,8 @@ def get_rh(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): delete_attrs=("units",)) def get_rh_2m(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): varnames=("T2", "PSFC", "Q2") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t2 = ncvars["T2"] psfc = ncvars["PSFC"] q2 = ncvars["Q2"] diff --git a/wrf_open/var/src/python/wrf/var/slp.py b/wrf_open/var/src/python/wrf/var/slp.py index 1171f73..cc45ef8 100755 --- a/wrf_open/var/src/python/wrf/var/slp.py +++ b/wrf_open/var/src/python/wrf/var/slp.py @@ -1,19 +1,25 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + from .extension import computeslp, computetk from .constants import Constants from .destag import destagger -from .decorators import convert_units, copy_and_set_metadata +from .decorators import convert_units +from .metadecorators import copy_and_set_metadata from .util import extract_vars __all__ = ["get_slp"] @copy_and_set_metadata(copy_varname="T", name="slp", remove_dims=("bottom_top",), - description="sea level pressure") + description="sea level pressure", + MemoryOrder="XY") @convert_units("pressure", "hpa") def get_slp(wrfnc, timeidx=0, units="hpa", method="cat", squeeze=True, cache=None): varnames=("T", "P", "PB", "QVAPOR", "PH", "PHB") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] p = ncvars["P"] @@ -25,6 +31,7 @@ def get_slp(wrfnc, timeidx=0, units="hpa", full_t = t + Constants.T_BASE full_p = p + pb qvapor[qvapor < 0] = 0. + full_ph = (ph + phb) / Constants.G destag_ph = destagger(full_ph, -3) diff --git a/wrf_open/var/src/python/wrf/var/temp.py b/wrf_open/var/src/python/wrf/var/temp.py index 390a3cf..5c4c8d1 100755 --- a/wrf_open/var/src/python/wrf/var/temp.py +++ b/wrf_open/var/src/python/wrf/var/temp.py @@ -1,7 +1,10 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) from .constants import Constants from .extension import computetk, computeeth, computetv, computewetbulb -from .decorators import convert_units, copy_and_set_metadata +from .decorators import convert_units +from .metadecorators import copy_and_set_metadata from .util import extract_vars __all__ = ["get_theta", "get_temp", "get_eth", "get_tv", "get_tw", @@ -14,7 +17,8 @@ def get_theta(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, cache=None): varnames = ("T",) - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] full_t = t + Constants.T_BASE @@ -28,7 +32,8 @@ def get_temp(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, """Return the temperature in Kelvin or Celsius""" varnames=("T", "P", "PB") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] @@ -47,7 +52,8 @@ def get_eth(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, "Return equivalent potential temperature (Theta-e) in Kelvin" varnames=("T", "P", "PB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] @@ -69,7 +75,8 @@ def get_tv(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, "Return the virtual temperature (tv) in Kelvin or Celsius" varnames=("T", "P", "PB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] p = ncvars["P"] @@ -92,7 +99,8 @@ def get_tw(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, "Return the wetbulb temperature (tw)" varnames=("T", "P", "PB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + nometa=True) t = ncvars["T"] p = ncvars["P"] pb = ncvars["PB"] diff --git a/wrf_open/var/src/python/wrf/var/terrain.py b/wrf_open/var/src/python/wrf/var/terrain.py index 5942d34..8fc4dca 100755 --- a/wrf_open/var/src/python/wrf/var/terrain.py +++ b/wrf_open/var/src/python/wrf/var/terrain.py @@ -1,5 +1,8 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) -from .decorators import convert_units, copy_and_set_metadata +from .decorators import convert_units +from .metadecorators import copy_and_set_metadata from .util import extract_vars, either __all__ = ["get_terrain"] @@ -12,7 +15,7 @@ def get_terrain(wrfnc, timeidx=0, units="m", method="cat", squeeze=True, cache=None): varname = either("HGT", "HGT_M")(wrfnc) return extract_vars(wrfnc, timeidx, varname, - method, squeeze, cache)[varname] + method, squeeze, cache, nometa=True)[varname] \ No newline at end of file diff --git a/wrf_open/var/src/python/wrf/var/times.py b/wrf_open/var/src/python/wrf/var/times.py index 7393845..61af303 100755 --- a/wrf_open/var/src/python/wrf/var/times.py +++ b/wrf_open/var/src/python/wrf/var/times.py @@ -1,3 +1,5 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) from .util import extract_times diff --git a/wrf_open/var/src/python/wrf/var/units.py b/wrf_open/var/src/python/wrf/var/units.py index e8fe358..be2795f 100755 --- a/wrf_open/var/src/python/wrf/var/units.py +++ b/wrf_open/var/src/python/wrf/var/units.py @@ -1,3 +1,5 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) from .constants import Constants, ConversionFactors diff --git a/wrf_open/var/src/python/wrf/var/util.py b/wrf_open/var/src/python/wrf/var/util.py index ef26359..9f037ad 100644 --- a/wrf_open/var/src/python/wrf/var/util.py +++ b/wrf_open/var/src/python/wrf/var/util.py @@ -1,12 +1,18 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + from collections import Iterable, Mapping, OrderedDict from itertools import product +from inspect import getargspec import datetime as dt import numpy as np +import numpy.ma as ma from .config import xarray_enabled from .projection import getproj + if xarray_enabled(): from xarray import DataArray @@ -14,33 +20,61 @@ __all__ = ["extract_vars", "extract_global_attrs", "extract_dim", "combine_files", "is_standard_wrf_var", "extract_times", "iter_left_indexes", "get_left_indexes", "get_right_slices", "is_staggered", "get_proj_params", "viewitems", "viewkeys", - "viewvalues"] + "viewvalues", "combine_with", "from_args", "arg_location", + "args_to_list", "npvalues", "XYCoord"] + + +_COORD_PAIR_MAP = {"XLAT" : ("XLAT", "XLONG"), + "XLONG" : ("XLAT", "XLONG"), + "XLAT_M" : ("XLAT_M", "XLONG_M"), + "XLONG_M" : ("XLAT_M", "XLONG_M"), + "XLAT_U" : ("XLAT_U", "XLONG_U"), + "XLONG_U" : ("XLAT_U", "XLONG_U"), + "XLAT_V" : ("XLAT_V", "XLONG_V"), + "XLONG_V" : ("XLAT_V", "XLONG_V"), + "CLAT" : ("CLAT", "CLONG"), + "CLONG" : ("CLAT", "CLONG")} + + +_COORD_VARS = ("XLAT", "XLONG", "XLAT_M", "XLONG_M", "XLAT_U", "XLONG_U", + "XLAT_V", "XLONG_V", "CLAT", "CLONG") + + + + +def _is_coord_var(varname): + return varname in _COORD_VARS + + +def _get_coord_pairs(varname): + return _COORD_PAIR_MAP[varname] + + +def _is_multi_time_req(timeidx): + return timeidx == -1 -def _is_multi_time(timeidx): - if timeidx == -1: - return True - return False def _is_multi_file(wrfnc): - if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str): - return True - return False + return (isinstance(wrfnc, Iterable) and not isstr(wrfnc)) + def _is_mapping(wrfnc): - if isinstance(wrfnc, Mapping): - return True - return False + return isinstance(wrfnc, Mapping) + def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar): lats = wrfnc.variables[latvar] lons = wrfnc.variables[lonvar] + # Need to check all times for i in xrange(lats.shape[-3]): - start_idxs = [0]*lats.ndim + start_idxs = [0]*len(lats.shape) # PyNIO does not support ndim start_idxs[-3] = i + start_idxs = tuple(start_idxs) - end_idxs = [-1]*lats.ndim + end_idxs = [-1]*len(lats.shape) end_idxs[-3] = i + end_idxs = tuple(end_idxs) if (first_ll_corner[0] != lats[start_idxs] or first_ll_corner[1] != lons[start_idxs] or @@ -62,10 +96,17 @@ def _is_moving_domain(wrfseq, varname=None, latvar="XLAT", lonvar="XLONG"): first_wrfnc = next(wrf_iter) if varname is not None: - coord_names = getattr(first_wrfnc.variables[varname], + try: + coord_names = getattr(first_wrfnc.variables[varname], "coordinates").split() - lon_coord = coord_names[0] - lat_coord = coord_names[1] + except AttributeError: + # Variable doesn't have a coordinates attribute, use the + # arguments + lon_coord = lonvar + lat_coord = latvar + else: + lon_coord = coord_names[0] + lat_coord = coord_names[1] else: lon_coord = lonvar lat_coord = latvar @@ -73,10 +114,14 @@ def _is_moving_domain(wrfseq, varname=None, latvar="XLAT", lonvar="XLONG"): lats = first_wrfnc.variables[lat_coord] lons = first_wrfnc.variables[lon_coord] - zero_idxs = [0] * first_wrfnc.variables[lat_coord].ndim + + zero_idxs = [0]*len(lats.shape) # PyNIO doesn't have ndim last_idxs = list(zero_idxs) last_idxs[-2:] = [-1]*2 + zero_idxs = tuple(zero_idxs) + last_idxs = tuple(last_idxs) + lat0 = lats[zero_idxs] lat1 = lats[last_idxs] lon0 = lons[zero_idxs] @@ -107,7 +152,7 @@ def _get_global_attr(wrfnc, attr): return val def extract_global_attrs(wrfnc, attrs): - if isinstance(attrs, str): + if isstr(attrs): attrlist = [attrs] else: attrlist = attrs @@ -134,7 +179,7 @@ def extract_dim(wrfnc, dim): return len(d) #netCDF4 return d # PyNIO -def _combine_dict(wrfdict, varname, timeidx, method): +def _combine_dict(wrfdict, varname, timeidx, method, nometa): """Dictionary combination creates a new left index for each key, then does a cat or join for the list of files for that key""" keynames = [] @@ -145,7 +190,8 @@ def _combine_dict(wrfdict, varname, timeidx, method): keynames.append(first_key) first_array = _extract_var(wrfdict[first_key], varname, - timeidx, method, squeeze=False) + timeidx, method, squeeze=False, + nometa=nometa) # Create the output data numpy array based on the first array @@ -170,11 +216,9 @@ def _combine_dict(wrfdict, varname, timeidx, method): "same size for all dictionary keys") outdata[idx,:] = vardata.values[:] idx += 1 - - if not xarray_enabled(): - outarr = outdata - else: - outname = str(first_array.name) + + if xarray_enabled() and not nometa: + outname = unicode(first_array.name) # Note: assumes that all entries in dict have same coords outcoords = OrderedDict(first_array.coords) outdims = ["key"] + list(first_array.dims) @@ -183,31 +227,37 @@ def _combine_dict(wrfdict, varname, timeidx, method): outarr = DataArray(outdata, name=outname, coords=outcoords, dims=outdims, attrs=outattrs) + else: + outarr = outdata return outarr # TODO: implement in C -def _cat_files(wrfseq, varname, timeidx): +def _cat_files(wrfseq, varname, timeidx, squeeze, nometa): is_moving = _is_moving_domain(wrfseq, varname) file_times = extract_times(wrfseq, timeidx) - multitime = _is_multi_time(timeidx) + multitime = _is_multi_time_req(timeidx) - time_idx_or_slice = timeidx if not multitime else slice(None, None, None) + time_idx_or_slice = timeidx if not multitime else slice(None) # wrfseq might be a generator wrf_iter = iter(wrfseq) - first_var = (_build_data_array(next(wrf_iter), varname, timeidx, is_moving) - if xarray_enabled() else - wrfseq[0].variables[varname][time_idx_or_slice, :]) + if xarray_enabled() and not nometa: + first_var = _build_data_array(next(wrf_iter), varname, + timeidx, is_moving) + else: + first_var = (next(wrf_iter)).variables[varname][time_idx_or_slice, :] + if not multitime: + first_var = first_var[np.newaxis,:] outdims = [len(file_times)] # Making a new time dim, so ignore this one outdims += first_var.shape[1:] - + outdata = np.empty(outdims, first_var.dtype) numtimes = first_var.shape[0] @@ -236,13 +286,16 @@ def _cat_files(wrfseq, varname, timeidx): startidx = endidx - if xarray_enabled(): + if xarray_enabled() and not nometa: # FIXME: If it's a moving nest, then the coord arrays need to have same # time indexes as the whole data set - outname = str(first_var.name) + outname = unicode(first_var.name) outattrs = OrderedDict(first_var.attrs) outcoords = OrderedDict(first_var.coords) outdimnames = list(first_var.dims) + if "Time" not in outdimnames: + outdimnames.insert(0, "Time") + outcoords[outdimnames[0]] = file_times # New time dimension values outarr = DataArray(outdata, name=outname, coords=outcoords, @@ -254,24 +307,23 @@ def _cat_files(wrfseq, varname, timeidx): return outarr # TODO: implement in C -def _join_files(wrfseq, varname, timeidx): +def _join_files(wrfseq, varname, timeidx, nometa): is_moving = _is_moving_domain(wrfseq, varname) - multitime = _is_multi_time(timeidx) + multitime = _is_multi_time_req(timeidx) numfiles = len(wrfseq) - if not multitime: - time_idx_or_slice = timeidx - else: - time_idx_or_slice = slice(None, None, None) - + time_idx_or_slice = timeidx if not multitime else slice(None) + # wrfseq might be a generator wrf_iter = iter(wrfseq) - if xarray_enabled(): + if xarray_enabled() and not nometa: first_var = _build_data_array(next(wrf_iter), varname, timeidx, is_moving) else: first_var = (next(wrf_iter)).variables[varname][time_idx_or_slice, :] + if not multitime: + first_var = first_var[np.newaxis, :] # Create the output data numpy array based on the first array outdims = [numfiles] @@ -287,11 +339,14 @@ def _join_files(wrfseq, varname, timeidx): except StopIteration: break else: - outdata[idx,:] = wrfnc.variables[varname][time_idx_or_slice, :] + outvar = wrfnc.variables[varname][time_idx_or_slice, :] + if not multitime: + outvar = outvar[np.newaxis, :] + outdata[idx,:] = outvar[:] idx += 1 - if xarray_enabled(): - outname = str(first_var.name) + if xarray_enabled() and not nometa: + outname = unicode(first_var.name) outcoords = OrderedDict(first_var.coords) outattrs = OrderedDict(first_var.attrs) # New dimensions @@ -306,43 +361,54 @@ def _join_files(wrfseq, varname, timeidx): return outarr -def combine_files(wrfseq, varname, timeidx, method="cat", squeeze=True): +def combine_files(wrfseq, varname, timeidx, method="cat", squeeze=True, + nometa=False): # Dictionary is unique if _is_mapping(wrfseq): - outarr = _combine_dict(wrfseq, varname, timeidx, method) + outarr = _combine_dict(wrfseq, varname, timeidx, method, nometa) elif method.lower() == "cat": - outarr = _cat_files(wrfseq, varname, timeidx) + outarr = _cat_files(wrfseq, varname, timeidx, squeeze, nometa) elif method.lower() == "join": - outarr = _join_files(wrfseq, varname, timeidx) + outarr = _join_files(wrfseq, varname, timeidx, nometa) else: raise ValueError("method must be 'cat' or 'join'") return outarr.squeeze() if squeeze else outarr -# Note, always returns the full data set with the time dimension included + def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): - multitime = _is_multi_time(timeidx) + multitime = _is_multi_time_req(timeidx) + time_idx_or_slice = timeidx if not multitime else slice(None) var = wrfnc.variables[varname] - data = var[:] + data = var[time_idx_or_slice, :] + + # Want to preserve the time dimension + if not multitime: + data = data[np.newaxis, :] + attrs = OrderedDict(var.__dict__) - dimnames = var.dimensions + dimnames = var.dimensions[-data.ndim:] # WRF variables will have a coordinates attribute. MET_EM files have # a stagger attribute which indicates the coordinate variable. try: # WRF files coord_attr = getattr(var, "coordinates") - except KeyError: - try: - # met_em files - stag_attr = getattr(var, "stagger") - except KeyError: - lon_coord = None - lat_coord = None + except AttributeError: + if _is_coord_var(varname): + # Coordinate variable (most likely XLAT or XLONG) + lat_coord, lon_coord = _get_coord_pairs(varname) else: - # For met_em files, use the stagger name to get the lat/lon var - lat_coord = "XLAT_{}".format(stag_attr) - lon_coord = "XLONG_{}".format(stag_attr) + try: + # met_em files + stag_attr = getattr(var, "stagger") + except AttributeError: + lon_coord = None + lat_coord = None + else: + # For met_em files, use the stagger name to get the lat/lon var + lat_coord = "XLAT_{}".format(stag_attr) + lon_coord = "XLONG_{}".format(stag_attr) else: coord_names = coord_attr.split() lon_coord = coord_names[0] @@ -393,54 +459,66 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): if dimnames[0] == "Time": coords[dimnames[0]] = extract_times(wrfnc, timeidx) + data_array = DataArray(data, name=varname, dims=dimnames, coords=coords, attrs=attrs) + return data_array # Cache is a dictionary of already extracted variables -def _extract_var(wrfnc, varname, timeidx, method, squeeze, cache): +def _extract_var(wrfnc, varname, timeidx, method, squeeze, cache, nometa): # Mainly used internally so variables don't get extracted multiple times, # particularly to copy metadata. This can be slow. if cache is not None: try: - return cache[varname] + cache_var = cache[varname] except KeyError: pass + else: + if nometa: + if isinstance(cache_var, DataArray): + return cache_var.values + + return cache_var + is_moving = _is_moving_domain(wrfnc, varname) - multitime = _is_multi_time(timeidx) + multitime = _is_multi_time_req(timeidx) multifile = _is_multi_file(wrfnc) if not multifile: - if xarray_enabled(): + if xarray_enabled() and not nometa: result = _build_data_array(wrfnc, varname, timeidx, is_moving) else: if not multitime: result = wrfnc.variables[varname][timeidx,:] + result = result[np.newaxis, :] # So that no squeeze works + else: result = wrfnc.variables[varname][:] else: # Squeeze handled in this routine, so just return it - return combine_files(wrfnc, varname, timeidx, method, squeeze) + return combine_files(wrfnc, varname, timeidx, method, squeeze, nometa) return result.squeeze() if squeeze else result def extract_vars(wrfnc, timeidx, varnames, method="cat", squeeze=True, - cache=None): - if isinstance(varnames, str): + cache=None, nometa=False): + if isstr(varnames): varlist = [varnames] else: varlist = varnames - return {var:_extract_var(wrfnc, var, timeidx, method, squeeze, cache) + return {var:_extract_var(wrfnc, var, timeidx, + method, squeeze, cache, nometa) for var in varlist} def _make_time(timearr): return dt.datetime.strptime("".join(timearr[:]), "%Y-%m-%d_%H:%M:%S") def _file_times(wrfnc, timeidx): - multitime = _is_multi_time(timeidx) + multitime = _is_multi_time_req(timeidx) if multitime: times = wrfnc.variables["Times"][:,:] for i in xrange(times.shape[0]): @@ -468,6 +546,7 @@ def is_standard_wrf_var(wrfnc, var): wrfnc = wrfnc[0] return var in wrfnc.variables + def is_staggered(var, wrfnc): we = extract_dim(wrfnc, "west_east") sn = extract_dim(wrfnc, "south_north") @@ -478,6 +557,8 @@ def is_staggered(var, wrfnc): return False + + def get_left_indexes(ref_var, expected_dims): """Returns the extra left side dimensions for a variable with an expected shape. @@ -491,7 +572,7 @@ def get_left_indexes(ref_var, expected_dims): if (extra_dim_num == 0): return [] - return [ref_var.shape[x] for x in xrange(extra_dim_num)] + return tuple([ref_var.shape[x] for x in xrange(extra_dim_num)]) def iter_left_indexes(dims): """A generator which yields the iteration tuples for a sequence of @@ -513,9 +594,10 @@ def iter_left_indexes(dims): def get_right_slices(var, right_ndims, fixed_val=0): extra_dim_num = var.ndim - right_ndims if extra_dim_num == 0: - return [slice(None,None,None)] * right_ndims + return [slice(None)] * right_ndims - return [fixed_val]*extra_dim_num + [slice(None,None,None)]*right_ndims + return tuple([fixed_val]*extra_dim_num + + [slice(None)]*right_ndims) def get_proj_params(wrfnc, timeidx=0, varname=None): proj_params = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", @@ -523,16 +605,20 @@ def get_proj_params(wrfnc, timeidx=0, varname=None): "TRUELAT1", "TRUELAT2", "MOAD_CEN_LAT", "STAND_LON", "POLE_LAT", "POLE_LON")) - multitime = _is_multi_time(timeidx) + multitime = _is_multi_time_req(timeidx) if not multitime: time_idx_or_slice = timeidx else: - time_idx_or_slice = slice(None, None, None) + time_idx_or_slice = slice(None) if varname is not None: - coord_names = getattr(wrfnc.variables[varname], "coordinates").split() - lon_coord = coord_names[0] - lat_coord = coord_names[1] + if not _is_coord_var(varname): + coord_names = getattr(wrfnc.variables[varname], + "coordinates").split() + lon_coord = coord_names[0] + lat_coord = coord_names[1] + else: + lat_coord, lon_coord = _get_coord_pairs(varname) else: lat_coord = "XLAT" lon_coord = "XLONG" @@ -562,6 +648,28 @@ def viewvalues(d): func = d.values return func() +def isstr(s): + try: + return isinstance(s, basestring) + except NameError: + return isinstance(s, str) + +# Helper to extract masked arrays from DataArrays that convert to NaN +def npvalues(da): + if not isinstance(da, DataArray): + result = da + else: + try: + fill_value = da.attrs["_FillValue"] + except KeyError: + result = da.values + else: + result = ma.masked_invalid(da.values, copy=False) + result.set_fill_value(fill_value) + + return result + + # Helper utilities for metadata class either(object): def __init__(self, *varnames): @@ -572,7 +680,7 @@ class either(object): wrfnc = next(iter(wrfnc)) for varname in self.varnames: - if varname in wrfnc: + if varname in wrfnc.variables: return varname raise ValueError("{} are not valid variable names".format( @@ -585,8 +693,9 @@ class combine_with: self.varname = varname self.remove_dims = remove_dims self.insert_before = insert_before - self.new_dimnames = new_dimnames - self.new_coords = new_coords + self.new_dimnames = new_dimnames if new_dimnames is not None else [] + self.new_coords = (new_coords if new_coords is not None + else OrderedDict()) def __call__(self, var): new_dims = list(var.dims) @@ -608,7 +717,76 @@ class combine_with: new_coords.update(self.new_coords) return new_dims, new_coords + +class XYCoord(object): + def __init__(self, x, y): + self.x = x + self.y = y + + def __repr__(self): + return "{}({}, {})".format(self.__class__.__name__, self.x, self.y) + +def from_args(func, argnames, *args, **kwargs): + """Parses the function args and kargs looking for the desired argument + value. Otherwise, the value is taken from the default keyword argument + using the arg spec. + + """ + if isstr(argnames): + arglist = [argnames] + else: + arglist = argnames + + result = {} + for argname in arglist: + arg_loc = arg_location(func, argname, args, kwargs) + + if arg_loc is not None: + result[argname] = arg_loc[0][arg_loc[1]] + else: + result[argname] = None + return result + +def args_to_list(func, args, kwargs): + """Converts the mixed args/kwargs to a single list of args""" + argspec = getargspec(func) + + # Build the full tuple with defaults filled in + outargs = [None]*len(argspec.args) + for i,default in enumerate(argspec.defaults[::-1], 1): + outargs[-i] = default + + # Add the supplied args + for i,arg in enumerate(args): + outargs[i] = arg + + # Fill in the supplied kargs + for argname,val in viewitems(kwargs): + argidx = argspec.args.index(argname) + outargs[argidx] = val + + return outargs + +def arg_location(func, argname, args, kwargs): + """Parses the function args, kargs and signature looking for the desired + argument location (either in args, kargs, or argspec.defaults), + and returns a list containing representing all arguments in the + correct order with defaults filled in. + + """ + argspec = getargspec(func) + list_args = args_to_list(func, args, kwargs) + + # Return the new sequence and location + if argname not in argspec.args and argname not in kwargs: + return None + + result_idx = argspec.args.index(argname) + + return list_args, result_idx + + diff --git a/wrf_open/var/src/python/wrf/var/uvdecorator.py b/wrf_open/var/src/python/wrf/var/uvdecorator.py new file mode 100644 index 0000000..12d9b26 --- /dev/null +++ b/wrf_open/var/src/python/wrf/var/uvdecorator.py @@ -0,0 +1,90 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + +import numpy as np + +import wrapt + +from .destag import destagger +from .util import iter_left_indexes + +__all__ = ["uvmet_left_iter"] + +# Placed in separate module to resolve a circular dependency with destagger +# module + +def uvmet_left_iter(): + """Decorator to handle iterating over leftmost dimensions when using + multiple files and/or multiple times with the uvmet product. + + """ + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + u = args[0] + v = args[1] + lat = args[2] + lon = args[3] + cen_long = args[4] + cone = args[5] + + if u.ndim == lat.ndim: + num_right_dims = 2 + is_3d = False + else: + num_right_dims = 3 + is_3d = True + + is_stag = False + if ((u.shape[-1] != lat.shape[-1]) or + (u.shape[-2] != lat.shape[-2])): + is_stag = True + + if is_3d: + extra_dim_num = u.ndim - 3 + else: + extra_dim_num = u.ndim - 2 + + if is_stag: + u = destagger(u,-1) + v = destagger(v,-2) + + # No special left side iteration, return the function result + if (extra_dim_num == 0): + return wrapped(u, v, lat, lon, cen_long, cone) + + # Start by getting the left-most 'extra' dims + outdims = [u.shape[x] for x in xrange(extra_dim_num)] + extra_dims = list(outdims) # Copy the left-most dims for iteration + + # Append the right-most dimensions + outdims += [2] # For u/v components + + outdims += [u.shape[x] for x in xrange(-num_right_dims,0,1)] + + output = np.empty(outdims, u.dtype) + + for left_idxs in iter_left_indexes(extra_dims): + # Make the left indexes plus a single slice object + # The single slice will handle all the dimensions to + # the right (e.g. [1,1,:]) + left_and_slice_idxs = tuple([x for x in left_idxs] + [slice(None)]) + + new_u = u[left_and_slice_idxs] + new_v = v[left_and_slice_idxs] + new_lat = lat[left_and_slice_idxs] + new_lon = lon[left_and_slice_idxs] + + # Call the numerical routine + res = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone) + + # Note: The 2D version will return a 3D array with a 1 length + # dimension. Numpy is unable to broadcast this without + # sqeezing first. + res = np.squeeze(res) + + output[left_and_slice_idxs] = res[:] + + return output + + return func_wrapper + diff --git a/wrf_open/var/src/python/wrf/var/uvmet.py b/wrf_open/var/src/python/wrf/var/uvmet.py index c650bbe..a3cf7fa 100755 --- a/wrf_open/var/src/python/wrf/var/uvmet.py +++ b/wrf_open/var/src/python/wrf/var/uvmet.py @@ -1,37 +1,46 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + from math import fabs, log, tan, sin, cos from .extension import computeuvmet from .destag import destagger from .constants import Constants from .wind import _calc_wspd_wdir -from .decorators import convert_units, set_wind_metadata +from .decorators import convert_units +from .metadecorators import set_wind_metadata from .util import extract_vars, extract_global_attrs, either __all__=["get_uvmet", "get_uvmet10", "get_uvmet_wspd_wdir", "get_uvmet10_wspd_wdir"] -@set_wind_metadata(wspd_wdir=False) + @convert_units("wind", "mps") -def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps", +def _get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps", method="cat", squeeze=True, cache=None): """ Return a tuple of u,v with the winds rotated in to earth space""" if not ten_m: varname = either("U", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) + u = destagger(u_vars[varname], -1) varname = either("V", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) v = destagger(v_vars[varname], -2) else: - varname = either("U10", "UU") - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + varname = either("U10", "UU")(wrfnc) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) u = (u_vars[varname] if varname == "U10" else destagger(u_vars[varname][...,0,:,:], -1)) - varname = either("V10", "VV") - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + varname = either("V10", "VV")(wrfnc) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) v = (v_vars[varname] if varname == "V10" else destagger(v_vars[varname][...,0,:,:], -2)) @@ -71,12 +80,12 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps", varname = either("XLAT_M", "XLAT")(wrfnc) xlat_var = extract_vars(wrfnc, timeidx, varname, - method, squeeze, cache) + method, squeeze, cache, nometa=True) lat = xlat_var[varname] - varname = either("XLONG_M", "XLONG") + varname = either("XLONG_M", "XLONG")(wrfnc) xlon_var = extract_vars(wrfnc, timeidx, varname, - method, squeeze, cache) + method, squeeze, cache, nometa=True) lon = xlon_var[varname] if map_proj == 1: @@ -92,31 +101,57 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps", else: cone = 1 - res = computeuvmet(u,v,lat,lon,cen_lon,cone) + res = computeuvmet(u, v ,lat, lon, cen_lon, cone) return res - + +@set_wind_metadata(copy_varname=either("P", "PRES"), + name="uvmet", + description="earth rotated u,v", + two_d=False, + wspd_wdir=False) +def get_uvmet(wrfnc, timeidx=0, units="mps", + method="cat", squeeze=True, cache=None): + + return _get_uvmet(wrfnc, timeidx, False, units, method, squeeze, cache) + + +@set_wind_metadata(copy_varname=either("PSFC", "F"), + name="uvmet10", + description="10m earth rotated u,v", + two_d=True, + wspd_wdir=False) def get_uvmet10(wrfnc, timeidx=0, units="mps", method="cat", squeeze=True, cache=None): - return get_uvmet(wrfnc, timeidx, True, units) + + return _get_uvmet(wrfnc, timeidx, True, units, method, squeeze, cache) + -@set_wind_metadata(wspd_wdir=True) +@set_wind_metadata(copy_varname=either("P", "PRES"), + name="uvmet_wspd_wdir", + description="earth rotated wspd,wdir", + two_d=False, + wspd_wdir=True) def get_uvmet_wspd_wdir(wrfnc, timeidx=0, units="mps", method="cat", squeeze=True, cache=None): - uvmet = get_uvmet(wrfnc, timeidx, False, units, method, squeeze, cache) + uvmet = _get_uvmet(wrfnc, timeidx, False, units, method, squeeze, cache) - return _calc_wspd_wdir(uvmet[0,:], uvmet[1,:], False, units) + return _calc_wspd_wdir(uvmet[...,0,:,:,:], uvmet[...,1,:,:,:], False, units) -@set_wind_metadata(wspd_wdir=True) +@set_wind_metadata(copy_varname=either("PSFC", "F"), + name="uvmet10_wspd_wdir", + description="10m earth rotated wspd,wdir", + two_d=True, + wspd_wdir=True) def get_uvmet10_wspd_wdir(wrfnc, timeidx=0, units="mps", method="cat", squeeze=True, cache=None): - uvmet10 = get_uvmet10(wrfnc, timeidx, units="mps", method, squeeze, cache) + uvmet10 = _get_uvmet(wrfnc, timeidx, True, units, method, squeeze, cache) - return _calc_wspd_wdir(uvmet10[0,:], uvmet10[1,:], True, units) + return _calc_wspd_wdir(uvmet10[...,0,:,:], uvmet10[...,1,:,:], True, units) diff --git a/wrf_open/var/src/python/wrf/var/vorticity.py b/wrf_open/var/src/python/wrf/var/vorticity.py index e19845d..6931951 100755 --- a/wrf_open/var/src/python/wrf/var/vorticity.py +++ b/wrf_open/var/src/python/wrf/var/vorticity.py @@ -1,17 +1,20 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + from .extension import computeavo, computepvo from .util import extract_vars, extract_global_attrs -from .decorators import copy_and_set_metadata +from .metadecorators import copy_and_set_metadata __all__ = ["get_avo", "get_pvo"] -@copy_and_set_metadata(copy_varname="F", name="avo", +@copy_and_set_metadata(copy_varname="T", name="avo", description="absolute vorticity", units="10-5 s-1") def get_avo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): - ncvars = extract_vars(wrfnc, timeidx, varnames=("U", "V", "MAPFAC_U", - "MAPFAC_V", "MAPFAC_M", - "F"), - method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, ("U", "V", "MAPFAC_U", + "MAPFAC_V", "MAPFAC_M", + "F"), + method, squeeze, cache, nometa=True) attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) u = ncvars["U"] @@ -24,18 +27,18 @@ def get_avo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): dx = attrs["DX"] dy = attrs["DY"] - return computeavo(u,v,msfu,msfv,msfm,cor,dx,dy) + return computeavo(u, v, msfu, msfv, msfm, cor, dx, dy) @copy_and_set_metadata(copy_varname="T", name="pvo", description="potential vorticity", units="PVU") def get_pvo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): - ncvars = extract_vars(wrfnc, timeidx, varnames=("U", "V", "T", "P", - "PB", "MAPFAC_U", - "MAPFAC_V", "MAPFAC_M", - "F"), - method, squeeze, cache) + ncvars = extract_vars(wrfnc, timeidx, ("U", "V", "T", "P", + "PB", "MAPFAC_U", + "MAPFAC_V", "MAPFAC_M", + "F"), + method, squeeze, cache, nometa=True) attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) u = ncvars["U"] @@ -54,5 +57,5 @@ def get_pvo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None): full_t = t + 300 full_p = p + pb - return computepvo(u,v,full_t,full_p,msfu,msfv,msfm,cor,dx,dy) + return computepvo(u, v, full_t, full_p, msfu, msfv, msfm, cor, dx, dy) \ No newline at end of file diff --git a/wrf_open/var/src/python/wrf/var/wind.py b/wrf_open/var/src/python/wrf/var/wind.py index 384a68e..071ce76 100755 --- a/wrf_open/var/src/python/wrf/var/wind.py +++ b/wrf_open/var/src/python/wrf/var/wind.py @@ -1,10 +1,13 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) import numpy as np from .constants import Constants from .destag import destagger from .util import extract_vars, either -from .decorators import convert_units, set_wind_metadata +from .decorators import convert_units +from .metadecorators import set_wind_metadata __all__ = ["get_u_destag", "get_v_destag", "get_w_destag", "get_destag_wspd_wdir", "get_destag_wspd_wdir10"] @@ -20,67 +23,107 @@ def _calc_wdir(u, v): def _calc_wspd_wdir(u, v, two_d, units): wspd = _calc_wspd(u, v, units) wdir = _calc_wdir(u, v) - + idx_end = -2 if two_d else -3 outdims = list(wspd.shape[0:idx_end]) + [2] + list(wspd.shape[idx_end:]) + res = np.zeros(outdims, wspd.dtype) - res[...,0,:] = wspd[:] - res[...,1,:] = wdir[:] + + idxs0 = ((Ellipsis, 0, slice(None), slice(None), slice(None)) + if not two_d else + (Ellipsis, 0, slice(None), slice(None))) + + idxs1 = ((Ellipsis, 1, slice(None), slice(None), slice(None)) + if not two_d else + (Ellipsis, 1, slice(None), slice(None))) + + res[idxs0] = wspd[:] + res[idxs1] = wdir[:] return res -@set_wind_metadata(wspd_wdir=False) +@set_wind_metadata(copy_varname=either("P", "PRES"), + name="ua", + description="destaggered u-wind component", + wind_ncvar=True, + two_d=False, + wspd_wdir=False) @convert_units("wind", "mps") def get_u_destag(wrfnc, timeidx=0, units="mps", method="cat", squeeze=True, cache=None): varname = either("U", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) u = destagger(u_vars[varname], -1) return u -@set_wind_metadata(wspd_wdir=False) +@set_wind_metadata(copy_varname=either("P", "PRES"), + name="va", + description="destaggered v-wind component", + two_d=False, + wind_ncvar=True, + wspd_wdir=False) @convert_units("wind", "mps") def get_v_destag(wrfnc, timeidx=0, units="mps", method="cat", squeeze=True, cache=None): varname = either("V", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) v = destagger(v_vars[varname], -2) return v -@set_wind_metadata(wspd_wdir=False) +@set_wind_metadata(copy_varname=either("P", "PRES"), + name="wa", + description="destaggered w-wind component", + two_d=False, + wind_ncvar=True, + wspd_wdir=False) @convert_units("wind", "mps") def get_w_destag(wrfnc, timeidx=0, units="mps", method="cat", squeeze=True, cache=None): - w_vars = extract_vars(wrfnc, timeidx, "W", method, squeeze, cache) + w_vars = extract_vars(wrfnc, timeidx, "W", method, squeeze, cache, + nometa=True) w = destagger(w_vars["W"], -3) return w -@set_wind_metadata(wspd_wdir=True) +@set_wind_metadata(copy_varname=either("P", "PRES"), + name="wspd_wdir", + description="wspd,wdir in projection space", + two_d=False, + wspd_wdir=True) def get_destag_wspd_wdir(wrfnc, timeidx=0, units="mps", method="cat", squeeze=True, cache=None): varname = either("U", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) u = destagger(u_vars[varname], -1) varname = either("V", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) v = destagger(v_vars[varname], -2) return _calc_wspd_wdir(u, v, False, units) -@set_wind_metadata(wspd_wdir=True) +@set_wind_metadata(copy_varname=either("PSFC", "F"), + name="wspd_wdir10", + description="10m wspd,wdir in projection space", + two_d=False, + wspd_wdir=True) def get_destag_wspd_wdir10(wrfnc, timeidx=0, units="mps", method="cat", squeeze=True, cache=None): varname = either("U10", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) u = (u_vars[varname] if varname == "U10" else destagger(u_vars[varname][...,0,:,:], -1)) varname = either("V10", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache) + v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + nometa=True) v = (v_vars[varname] if varname == "V10" else destagger(v_vars[varname][...,0,:,:], -2)) diff --git a/wrf_open/var/src/python/wrf/var/wrfext.f90 b/wrf_open/var/src/python/wrf/var/wrfext.f90 index 1964602..bed5088 100755 --- a/wrf_open/var/src/python/wrf/var/wrfext.f90 +++ b/wrf_open/var/src/python/wrf/var/wrfext.f90 @@ -85,7 +85,7 @@ SUBROUTINE f_interp1d(v_in,z_in,z_out,vmsg,v_out,nz_in,nz_out) IMPLICIT NONE - INTEGER,INTENT(IN) :: NZ_IN,NZ_OUT + INTEGER,INTENT(IN) :: nz_in, nz_out REAL(KIND=8),DIMENSION(nz_in),INTENT(IN) :: v_in,z_in REAL(KIND=8),DIMENSION(nz_out),INTENT(IN) :: z_out REAL(KIND=8),DIMENSION(nz_out),INTENT(OUT) :: v_out diff --git a/wrf_open/var/test/utests.py b/wrf_open/var/test/utests.py index aa0e369..765706a 100644 --- a/wrf_open/var/test/utests.py +++ b/wrf_open/var/test/utests.py @@ -5,12 +5,14 @@ import numpy.ma as ma import os, sys import subprocess -from wrf.var import getvar, interplevel, interpline, vertcross, vinterp +from wrf.var import (getvar, interplevel, interpline, vertcross, vinterp, + disable_xarray, xarray_enabled, npvalues) NCL_EXE = "/Users/ladwig/nclbuild/6.3.0/bin/ncl" TEST_FILE = "/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00" OUT_NC_FILE = "/tmp/wrftest.nc" + def setUpModule(): ncarg_root = os.environ.get("NCARG_ROOT", None) if ncarg_root is None: @@ -101,28 +103,28 @@ def make_test(varname, wrf_in, referent, multi=False, repeat=3, pynio=False): my_vals = getvar(in_wrfnc, "temp", units="c") tol = 0 atol = .1 # Note: NCL uses 273.16 as conversion for some reason - nt.assert_allclose(my_vals, ref_vals, tol, atol) + nt.assert_allclose(npvalues(my_vals), ref_vals, tol, atol) elif (varname == "pw"): my_vals = getvar(in_wrfnc, "pw") tol = .5/100.0 atol = 0 # NCL uses different constants and doesn't use same # handrolled virtual temp in method - nt.assert_allclose(my_vals, ref_vals, tol, atol) + nt.assert_allclose(npvalues(my_vals), ref_vals, tol, atol) elif (varname == "cape_2d"): cape_2d = getvar(in_wrfnc, varname) tol = 0/100. # Not sure why different, F77 vs F90? - atol = .1 - nt.assert_allclose(cape_2d, ref_vals, tol, atol) + atol = .2 + nt.assert_allclose(npvalues(cape_2d), ref_vals, tol, atol) elif (varname == "cape_3d"): cape_3d = getvar(in_wrfnc, varname) tol = 0/100. # Not sure why different, F77 vs F90? atol = .01 - nt.assert_allclose(cape_3d, ref_vals, tol, atol) + nt.assert_allclose(npvalues(cape_3d), ref_vals, tol, atol) else: my_vals = getvar(in_wrfnc, varname) tol = 0/100. atol = 0.0001 - nt.assert_allclose(my_vals, ref_vals, tol, atol) + nt.assert_allclose(npvalues(my_vals), ref_vals, tol, atol) return test @@ -199,7 +201,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, p = getvar(in_wrfnc, "pressure") hts_500 = interplevel(hts, p, 500) - nt.assert_allclose(hts_500, ref_ht_500) + nt.assert_allclose(npvalues(hts_500), ref_ht_500) elif (varname == "vertcross"): ref_ht_cross = _get_refvals(referent, "ht_cross", repeat, multi) @@ -209,14 +211,14 @@ def make_interp_test(varname, wrf_in, referent, multi=False, p = getvar(in_wrfnc, "pressure") pivot_point = (hts.shape[-2] / 2, hts.shape[-1] / 2) - ht_cross = vertcross(hts, p, pivot_point=pivot_point,angle=90.) - - nt.assert_allclose(ht_cross, ref_ht_cross, rtol=.01) + ht_cross = vertcross(hts, p, pivot_point=pivot_point, angle=90.) + + nt.assert_allclose(npvalues(ht_cross), ref_ht_cross, rtol=.01) # Test opposite p_cross1 = vertcross(p,hts,pivot_point=pivot_point, angle=90.0) - - nt.assert_allclose(p_cross1, + + nt.assert_allclose(npvalues(p_cross1), ref_p_cross, rtol=.01) # Test point to point @@ -225,10 +227,12 @@ def make_interp_test(varname, wrf_in, referent, multi=False, p_cross2 = vertcross(p,hts,start_point=start_point, end_point=end_point) - - nt.assert_allclose(p_cross1, p_cross2) - + + nt.assert_allclose(npvalues(p_cross1), + npvalues(p_cross2)) + elif (varname == "interpline"): + ref_t2_line = _get_refvals(referent, "t2_line", repeat, multi) t2 = getvar(in_wrfnc, "T2") @@ -236,7 +240,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, t2_line1 = interpline(t2, pivot_point=pivot_point, angle=90.0) - nt.assert_allclose(t2_line1, ref_t2_line) + nt.assert_allclose(npvalues(t2_line1), ref_t2_line) # Test point to point start_point = (t2.shape[-2]/2, 0) @@ -245,7 +249,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, t2_line2 = interpline(t2, start_point=start_point, end_point=end_point) - nt.assert_allclose(t2_line1, t2_line2) + nt.assert_allclose(npvalues(t2_line1), npvalues(t2_line2)) elif (varname == "vinterp"): # Tk to theta fld_tk_theta = _get_refvals(referent, "fld_tk_theta", repeat, multi) @@ -267,7 +271,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, atol = 0.0001 field = n.squeeze(field) - nt.assert_allclose(field, fld_tk_theta, tol, atol) + nt.assert_allclose(npvalues(field), fld_tk_theta, tol, atol) # Tk to theta-e fld_tk_theta_e = _get_refvals(referent, "fld_tk_theta_e", repeat, multi) @@ -287,7 +291,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, atol = 0.0001 field = n.squeeze(field) - nt.assert_allclose(field, fld_tk_theta_e, tol, atol) + nt.assert_allclose(npvalues(field), fld_tk_theta_e, tol, atol) # Tk to pressure fld_tk_pres = _get_refvals(referent, "fld_tk_pres", repeat, multi) @@ -304,7 +308,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, log_p=True) field = n.squeeze(field) - nt.assert_allclose(field, fld_tk_pres, tol, atol) + nt.assert_allclose(npvalues(field), fld_tk_pres, tol, atol) # Tk to geoht_msl fld_tk_ght_msl = _get_refvals(referent, "fld_tk_ght_msl", repeat, multi) @@ -320,7 +324,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, log_p=True) field = n.squeeze(field) - nt.assert_allclose(field, fld_tk_ght_msl, tol, atol) + nt.assert_allclose(npvalues(field), fld_tk_ght_msl, tol, atol) # Tk to geoht_agl fld_tk_ght_agl = _get_refvals(referent, "fld_tk_ght_agl", repeat, multi) @@ -336,7 +340,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, log_p=True) field = n.squeeze(field) - nt.assert_allclose(field, fld_tk_ght_agl, tol, atol) + nt.assert_allclose(npvalues(field), fld_tk_ght_agl, tol, atol) # Hgt to pressure fld_ht_pres = _get_refvals(referent, "fld_ht_pres", repeat, multi) @@ -353,7 +357,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, log_p=True) field = n.squeeze(field) - nt.assert_allclose(field, fld_ht_pres, tol, atol) + nt.assert_allclose(npvalues(field), fld_ht_pres, tol, atol) # Pressure to theta fld_pres_theta = _get_refvals(referent, "fld_pres_theta", repeat, multi) @@ -370,7 +374,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, log_p=True) field = n.squeeze(field) - nt.assert_allclose(field, fld_pres_theta, tol, atol) + nt.assert_allclose(npvalues(field), fld_pres_theta, tol, atol) # Theta-e to pres fld_thetae_pres = _get_refvals(referent, "fld_thetae_pres", repeat, multi) @@ -387,7 +391,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, log_p=True) field = n.squeeze(field) - nt.assert_allclose(field, fld_thetae_pres, tol, atol) + nt.assert_allclose(npvalues(field), fld_thetae_pres, tol, atol) return test