Added metadata decorators (untested and buggy)

wrf_open/var/src/python/wrf/var/__init__.py

@@ -1,63 +1,61 @@
 import warnings
 
-from config import *
+from . import config
-import config
+from .config import *
-from extension import *
+from . import extension
-import extension
+from .extension import *
-from util import *
+from . import util
-import util
+from .util import *
-from cape import *
+from . import cape
-import cape
+from .cape import *
-from constants import *
+from . import constants
-import constants
+from .constants import *
-from ctt import *
+from . import ctt
-import ctt
+from .ctt import *
-from dbz import *
+from . import dbz
-import dbz
+from .dbz import *
-from destag import *
+from . import destag
-import destag
+from .destag import *
-from dewpoint import *
+from . import dewpoint
-import dewpoint
+from .dewpoint import *
-from etaconv import *
+from . import geoht
-import etaconv
+from .geoht import *
-from geoht import *
+from . import helicity
-import geoht
+from .helicity import *
-from helicity import *
+from . import interp
-import helicity
+from .interp import *
-from interp import *
+from . import latlon
-import interp
+from .latlon import *
-from latlon import *
+from . import omega
-import latlon
+from .omega import *
-from omega import *
+from . import precip
-import omega
+from .precip import *
-from precip import *
+from . import pressure
-import precip
+from .pressure import *
-from pressure import *
+from . import psadlookup
-import pressure
+from .psadlookup import *
-from psadlookup import *
+from . import pw
-import psadlookup
+from .pw import *
-from pw import *
+from . import rh
-import pw
+from .rh import *
-from rh import *
+from . import slp
-import rh
+from .slp import *
-from slp import *
+from . import temp
-import slp
+from .temp import *
-from temp import *
+from . import terrain
-import temp
+from .terrain import *
-from terrain import *
+from . import uvmet
-import terrain
+from .uvmet import *
-from uvmet import *
+from . import vorticity
-import uvmet
+from .vorticity import *
-from vorticity import *
+from . import wind
-import vorticity
+from .wind import *
-from wind import *
+from . import times
-import wind
+from .times import *
-from times import *
+from . import units
-import times
+from .units import *
-from units import *
+from . import projection
-import units
+from .projection import *
-from projection import *
-import projection
 
 __all__ = ["getvar"]
 __all__.extend(config.__all__)
@@ -69,7 +67,6 @@ __all__.extend(ctt.__all__)
 __all__.extend(dbz.__all__)
 __all__.extend(destag.__all__)
 __all__.extend(dewpoint.__all__)
-__all__.extend(etaconv.__all__)
 __all__.extend(geoht.__all__)
 __all__.extend(helicity.__all__)
 __all__.extend(interp.__all__)
@@ -127,9 +124,10 @@ _FUNC_MAP = {"cape2d" : get_2dcape,
              "lon" : get_lon,
              "pressure" : get_pressure_hpa,
              "pres" : get_pressure,
-             "wspddir" : get_destag_wspd_wdir,
+             "wspd_wdir" : get_destag_wspd_wdir,
-             "wspddir_uvmet" : get_uvmet_wspd_wdir,
+             "wspd_wdir10" : get_destag_wspd_wdir10,
-             "wspddir_uvmet10" : get_uvmet10_wspd_wdir,
+             "wspd_wdir_uvmet" : get_uvmet_wspd_wdir,
+             "wspd_wdir_uvmet10" : get_uvmet10_wspd_wdir,
              "ctt" : get_ctt
              }
 
@@ -214,9 +212,11 @@ def _check_kargs(var, kargs):
                           "argument for '%s'" % (arg, var))
             
 
-def getvar(wrfnc, var, timeidx=0, **kargs):
+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)[var]
+        return extract_vars(wrfnc, timeidx, var, method, squeeze, cache)[var]
     
     actual_var = _undo_alias(var)
     if actual_var not in _VALID_KARGS:

wrf_open/var/src/python/wrf/var/cape.py

@@ -1,17 +1,26 @@
 import numpy as n
 import numpy.ma as ma
 
-from wrf.var.extension import computetk,computecape
+from .extension import computetk,computecape
-from wrf.var.destag import destagger
+from .destag import destagger
-from wrf.var.constants import Constants, ConversionFactors
+from .constants import Constants, ConversionFactors
-from wrf.var.util import extract_vars
+from .util import extract_vars, combine_with
+from .decorators import copy_and_set_metadata
 
 __all__ = ["get_2dcape", "get_3dcape"]
-
+@copy_and_set_metadata(copy_varname="T", 
-def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL):
+                       name="cape_2d",
+                       dimnames=combine_with("T", remove_dims=("bottom_top",),
+                                        insert_before="south_north",
+                                        new_dimnames=("mcape_mcin_lcl_lfc",)), 
+                       description="mcape ; mcin ; lcl ; lfc",
+                       units="J/kg ; J/kg ; m ; m",
+                       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"""
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR", "PH",
+    varnames = ("T", "P", "PB", "QVAPOR", "PH","PHB", "HGT", "PSFC")
-                                              "PHB", "HGT", "PSFC"))
+    ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
     
     t = ncvars["T"]
     p = ncvars["P"]
@@ -51,7 +60,7 @@ def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL):
     resdim = left_dims + [4] + right_dims
     
     # Make a new output array for the result
-    res = n.zeros((resdim), cape.dtype)
+    res = n.zeros(resdim, cape.dtype)
     
     res[...,0,:,:] = cape[:]
     res[...,1,:,:] = cin[:]
@@ -60,11 +69,18 @@ def get_2dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL):
     
     return ma.masked_values(res, missing)
 
-def get_3dcape(wrfnc, timeidx=0, missing=Constants.DEFAULT_FILL):
+@copy_and_set_metadata(copy_varname="T", name="cape_3d",
+                       dimnames=combine_with("T",
+                                             insert_before="bottom_top",
+                                             new_dimnames=("cape_cin",)),
+                       description="cape ; cin",
+                       units="J kg-1 ; J kg-1",
+                       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=("T", "P", "PB", "QVAPOR", 
+    ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
-                                                "PH", "PHB", "HGT", "PSFC"))
     t = ncvars["T"]
     p = ncvars["P"]
     pb = ncvars["PB"]

wrf_open/var/src/python/wrf/var/ctt.py

@@ -1,21 +1,24 @@
 
 import numpy as n
 
-from wrf.var.extension import computectt, computetk
+from .extension import computectt, computetk
-from wrf.var.constants import Constants, ConversionFactors
+from .constants import Constants, ConversionFactors
-from wrf.var.destag import destagger
+from .destag import destagger
-from wrf.var.decorators import convert_units
+from .decorators import convert_units, copy_and_set_metadata
-from wrf.var.util import extract_vars
+from .util import extract_vars
 
 __all__ = ["get_ctt"]
 
+@copy_and_set_metadata(copy_varname="T", name="ctt", 
+                       description="cloud top temperature")
 @convert_units("temp", "c")
-def get_ctt(wrfnc, timeidx=0, units="c"):
+def get_ctt(wrfnc, timeidx=0, units="c", method="cat", squeeze=True, 
+              cache=None):
     """Return the cloud top temperature.
     
     """
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "PH" ,
+    varnames = ("T", "P", "PB", "PH", "PHB", "HGT", "QVAPOR")
-                                            "PHB", "HGT", "QVAPOR"))
+    ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
     t = ncvars["T"]
     p = ncvars["P"]
     pb = ncvars["PB"]
@@ -26,7 +29,7 @@ def get_ctt(wrfnc, timeidx=0, units="c"):
     
     haveqci = 1
     try:
-        icevars = extract_vars(wrfnc, timeidx, vars="QICE")
+        icevars = extract_vars(wrfnc, timeidx, "QICE", method, squeeze, cache)
     except KeyError:
         qice = n.zeros(qv.shape, qv.dtype)
         haveqci = 0
@@ -34,7 +37,7 @@ def get_ctt(wrfnc, timeidx=0, units="c"):
         qice = icevars["QICE"] * 1000.0 #g/kg
     
     try:
-        cldvars = extract_vars(wrfnc, timeidx, vars="QCLOUD")
+        cldvars = extract_vars(wrfnc, timeidx, "QCLOUD", method, squeeze, cache)
     except KeyError:
         raise RuntimeError("'QCLOUD' not found in NetCDF file")
     else:

wrf_open/var/src/python/wrf/var/dbz.py

@@ -1,12 +1,17 @@
 import numpy as n
 
-from wrf.var.extension import computedbz,computetk
+from .extension import computedbz,computetk
-from wrf.var.constants import Constants
+from .constants import Constants
-from wrf.var.util import extract_vars
+from .util import extract_vars, combine_with
+from .decorators import copy_and_set_metadata
 
 __all__ = ["get_dbz", "get_max_dbz"]
 
-def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False):
+@copy_and_set_metadata(copy_varname="T", name="dbz", 
+                       description="radar reflectivity",
+                       units="dBz")
+def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False,
+            method="cat", squeeze=True, cache=None):
     """ Return the dbz
     
     do_varint - do variable intercept (if False, constants are used.  Otherwise, 
@@ -17,8 +22,8 @@ def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False):
     as liquid)
     
     """
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR", 
+    varnames = ("T", "P", "PB", "QVAPOR", "QRAIN")
-                                              "QRAIN"))
+    ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
     t = ncvars["T"]
     p = ncvars["P"]
     pb = ncvars["PB"]
@@ -26,14 +31,16 @@ def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False):
     qr = ncvars["QRAIN"]
     
     try:
-        snowvars = extract_vars(wrfnc, timeidx, vars="QSNOW")
+        snowvars = extract_vars(wrfnc, timeidx, "QSNOW", 
+                                method, squeeze, cache)
     except KeyError:
         qs = n.zeros(qv.shape, "float")
     else:
         qs = snowvars["QSNOW"]
     
     try:
-        graupvars = extract_vars(wrfnc, timeidx, vars="QGRAUP")
+        graupvars = extract_vars(wrfnc, timeidx, "QGRAUP", 
+                                 method, squeeze, cache)
     except KeyError:
         qg = n.zeros(qv.shape, "float")
     else:
@@ -58,7 +65,13 @@ def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False):
     
     return computedbz(full_p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin)
 
-def get_max_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False):
+@copy_and_set_metadata(copy_varname="T", name="dbz", 
-    return n.amax(get_dbz(wrfnc, do_varint, do_liqskin, timeidx), 
+                       dimnames=combine_with("T", remove_dims=("bottom_top",)),
+                       description="maximum radar reflectivity",
+                       units="dBz")
+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, 
+                          method, squeeze, cache), 
                   axis=-3)
 

wrf_open/var/src/python/wrf/var/decorators.py

@@ -1,13 +1,16 @@
-from functools import wraps
+#from functools import wraps
+import wrapt 
 from inspect import getargspec
+from collections import OrderedDict
 
-import numpy as n
+import numpy as np
 import numpy.ma as ma
 
-from wrf.var.units import do_conversion, check_units
+from .units import do_conversion, check_units
-from wrf.var.destag import destagger
+from .destag import destagger
-from wrf.var.util import iter_left_indexes
+from .util import (iter_left_indexes, viewitems, extract_vars,
-from wrf.var.config import xarray_enabled
+                          combine_with, either)
+from .config import xarray_enabled
 
 if xarray_enabled():
     from xarray import DataArray
@@ -15,40 +18,60 @@ if xarray_enabled():
 
 
 __all__ = ["convert_units", "handle_left_iter", "uvmet_left_iter", 
-           "handle_casting" "set_metadata"]
+           "handle_casting" "copy_and_set_metadata", "set_wind_metadata",
+           "set_latlon_metadata", "set_height_metadata"]
 
 # TODO:  In python 3.x, getargspec is deprecated
-class from_args(object):
+def from_args(func, argnames, *args, **kwargs):
-    def __init__(self, argname):
-        self.argname
-        
-    def __call__(self, func, *args, **kargs):
         """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 units are provided to the method call, use them.  
+        if isinstance(argnames, str):
-        # Otherwise, need to parse the argspec to find what the default 
+            arglist = [argnames]
-        # value is since wraps does not preserve this.
+        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)
     
-        if self.argname not in argspec.args and self.argname not in kargs:
+    print argspec
+    
+    if argname not in argspec.args and argname not in kwargs:
         return None
         
     try:
-            result_idx = argspec.args.index(self.argname)
+        result_idx = argspec.args.index(argname)
     except ValueError:
         result_idx = None
     
-        if (self.argname in kargs):
+    result = None
-            result = kargs[self.argname]
+    if (argname in kwargs):
+        result = kwargs, argname
     elif (len(args) > result_idx and result_idx is not None):
-            result = args[result_idx]
+        result = args, result_idx
     else:    
         arg_idx_from_right = (len(argspec.args) - 
-                                  argspec.args.index(self.argname))
+                              argspec.args.index(argname))
-            result = argspec.defaults[-arg_idx_from_right]
+        result = list(argspec.defaults), -arg_idx_from_right
         
     return result
   
@@ -61,19 +84,19 @@ def convert_units(unit_type, alg_unit):
         - alg_unit - the units that the function returns by default
     
     """
-    def convert_decorator(func):
+#    def convert_decorator(func):
-        @wraps(func)
+    @wrapt.decorator
-        def func_wrapper(*args, **kargs):
+    def func_wrapper(wrapped, instance, args, kwargs):
         
-            desired_units = from_args("units")(func, *args, **kargs)
+        desired_units = from_args(wrapped, "units", *args, **kwargs)["units"]
         check_units(desired_units, unit_type)
         
         # Unit conversion done here
-            return do_conversion(func(*args, **kargs), unit_type, 
+        return do_conversion(wrapped(*args, **kwargs), unit_type, 
                              alg_unit, desired_units)
     return func_wrapper
     
-    return convert_decorator
+#    return convert_decorator
 
 
 def _calc_out_dims(outvar, left_dims):
@@ -105,23 +128,23 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
         the output dimensions (e.g. avo)
     
     """
-    def indexing_decorator(func):
+#    def indexing_decorator(func):
-        @wraps(func)
+    @wrapt.decorator
-        def func_wrapper(*args, **kargs):
+    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 ()
         
         if ref_var_idx >= 0:
             ref_var = args[ref_var_idx]
         else:
-                ref_var = kargs[ref_var_name]
+            ref_var = kwargs[ref_var_name]
         
         ref_var_shape = ref_var.shape
         extra_dim_num = ref_var.ndim - ref_var_expected_dims
         
         # No special left side iteration, return the function result
         if (extra_dim_num == 0):
-                return func(*args, **kargs)
+            return wrapped(*args, **kwargs)
         
         # Start by getting the left-most 'extra' dims
         extra_dims = [ref_var_shape[x] for x in xrange(extra_dim_num)]            
@@ -141,22 +164,22 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
             # Slice the kargs if applicable
             new_kargs = {key:(val[left_and_slice_idxs] 
                          if key not in ignore_kargs else val)
-                             for key,val in kargs.iteritems()}
+                         for key,val in viewitems(kwargs)}
             
             # Call the numerical routine
-                res = func(*new_args, **new_kargs)
+            res = wrapped(*new_args, **new_kargs)
             
-                if isinstance(res, n.ndarray):
+            if isinstance(res, np.ndarray):
                 # Output array
                 if not out_inited:
                     outdims = _calc_out_dims(res, extra_dims)
                     if not isinstance(res, ma.MaskedArray):
-                            output = n.zeros(outdims, ref_var.dtype)
+                        output = np.empty(outdims, ref_var.dtype)
                         masked = False
                     else:
                         output = ma.MaskedArray(
-                                            n.zeros(outdims, ref_var.dtype),
+                                        np.zeros(outdims, ref_var.dtype),
-                                            mask=n.zeros(outdims, n.bool_),
+                                        mask=np.zeros(outdims, np.bool_),
                                         fill_value=res.fill_value)
                         masked = True
                     
@@ -172,13 +195,13 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
                 if not out_inited:
                     outdims = _calc_out_dims(res[0], extra_dims)
                     if not isinstance(res[0], ma.MaskedArray):
-                            output = [n.zeros(outdims, ref_var.dtype) 
+                        output = [np.empty(outdims, ref_var.dtype) 
                                   for i in xrange(len(res))]
                         masked = False
                     else:
                         output = [ma.MaskedArray(
-                                        n.zeros(outdims, ref_var.dtype),
+                                    np.zeros(outdims, ref_var.dtype),
-                                        mask=n.zeros(outdims, n.bool_),
+                                    mask=np.zeros(outdims, np.bool_),
                                     fill_value=res[0].fill_value) 
                                   for i in xrange(len(res))]
                         masked = True
@@ -197,16 +220,16 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
     
     return func_wrapper
     
-    return indexing_decorator
+#    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):
+    #def indexing_decorator(func):
-        @wraps(func)
+    @wrapt.decorator
-        def func_wrapper(*args):
+    def func_wrapper(wrapped, instance, args, kwargs):
         u = args[0]
         v = args[1]
         lat = args[2]
@@ -237,7 +260,7 @@ def uvmet_left_iter():
         
         # No special left side iteration, return the function result
         if (extra_dim_num == 0):
-                return func(u,v,lat,lon,cen_long,cone)
+            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)]
@@ -248,7 +271,7 @@ def uvmet_left_iter():
         
         outdims += [u.shape[x] for x in xrange(-num_right_dims,0,1)]
         
-            output = n.zeros(outdims, u.dtype)
+        output = np.empty(outdims, u.dtype)
         
         for left_idxs in iter_left_indexes(extra_dims):
             # Make the left indexes plus a single slice object
@@ -264,12 +287,12 @@ def uvmet_left_iter():
             new_lon = lon[left_and_slice_idxs]
             
             # Call the numerical routine
-                res = func(new_u, new_v, new_lat, new_lon, cen_long, cone)
+            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 = n.squeeze(res) 
+            res = np.squeeze(res) 
             
             output[left_and_slice_idxs] = res[:]
             
@@ -278,16 +301,16 @@ def uvmet_left_iter():
     
     return func_wrapper
     
-    return indexing_decorator
+#    return indexing_decorator
 
-def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=n.float64):
+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):
+#    def cast_decorator(func):
-        @wraps(func)
+    @wrapt.decorator
-        def func_wrapper(*args, **kargs):
+    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 ()
         
@@ -299,11 +322,11 @@ def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=n.float64):
         
         new_kargs = {key:(val.astype(dtype)
                     if key in karg_names else val)
-                        for key,val in kargs.iteritems()}
+                    for key,val in viewitems()}
         
-            res = func(*new_args, **new_kargs) 
+        res = wrapped(*new_args, **new_kargs) 
         
-            if isinstance(res, n.ndarray):
+        if isinstance(res, np.ndarray):
             if res.dtype == orig_type:
                 return res
             return res.astype(orig_type)
@@ -314,14 +337,14 @@ def handle_casting(ref_idx=0, arg_idxs=None, karg_names=None, dtype=n.float64):
     
     return func_wrapper
     
-    return cast_decorator
+#    return cast_decorator
 
 def _extract_and_transpose(arg):
     if xarray_enabled():
         if isinstance(arg, DataArray):
             arg = arg.values
     
-    if isinstance(arg, n.ndarray):
+    if isinstance(arg, np.ndarray):
         if not arg.flags.F_CONTIGUOUS:
             return arg.T
         
@@ -332,18 +355,18 @@ def handle_extract_transpose():
     transposes if the data is not fortran contiguous.
     
     """
-    def extract_transpose_decorator(func):
+#    def extract_transpose_decorator(func):
-        @wraps(func)
+    @wrapt.decorator
-        def func_wrapper(*args, **kargs):
+    def func_wrapper(wrapped, instance, args, kwargs):
         
         new_args = [_extract_and_transpose(arg) for arg in args]
         
         new_kargs = {key:_extract_and_transpose(val) 
-                         for key,val in kargs.iteritems()}
+                     for key,val in viewitems(kwargs)}
             
-            res = func(*new_args, **new_kargs) 
+        res = wrapped(*new_args, **new_kargs) 
         
-            if isinstance(res, n.ndarray):
+        if isinstance(res, np.ndarray):
             if res.flags.F_CONTIGUOUS:
                 return res.T
         else:
@@ -353,32 +376,284 @@ def handle_extract_transpose():
     
     return func_wrapper
     
-    return extract_transpose_decorator
+#    return extract_transpose_decorator
 
-def set_metadata(copy_from=None, ignore=None, **fixed_kargs):
+
-    """Decorator to set the attributes for a WRF method.
+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.
     
     """
-    def attr_decorator(func):
+    @wrapt.decorator
-        @wraps(func)
+    def func_wrapper(wrapped, instance, args, kwargs):
-        def func_wrapper(*args, **kargs):
         if not xarray_enabled():
-                return func(*args, **kargs)
+            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 = {}
+        
         
-            result = func(*args, **kargs)
+        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
         
-            units = from_args("units")(func, *args, **kargs)
         if units is not None:
-                result.attrs["units"] = units
+            outattrs["units"] = units
             
-            for argname, val in fixed_kargs.iteritems():
+        for argname, val in viewitems(fixed_attrs):
-                result[argname] = val
+            outattrs[argname] = val
         
-            return result
+        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
     
-    return attr_decorator
+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)
+        
+        
     
     
 

wrf_open/var/src/python/wrf/var/destag.py

@@ -1,6 +1,6 @@
 
 
-from wrf.var.util import extract_vars
+from .util import extract_vars
 
 __all__ = ["destagger", "destagger_windcomp", "destagger_winds"]
 
@@ -40,24 +40,25 @@ def destagger(var, stagger_dim):
     
     return result
 
-def destagger_windcomp(wrfnc, comp, timeidx=0):
+# def destagger_windcomp(wrfnc, comp, timeidx=0, 
-    if comp.lower() == "u":
+#                        method="cat", squeeze=True, cache=None):
-        wrfvar = "U"
+#     if comp.lower() == "u":
-        stagdim = -1
+#         wrfvar = "U"
-    elif comp.lower() == "v":
+#         stagdim = -1
-        wrfvar = "V"
+#     elif comp.lower() == "v":
-        stagdim = -2
+#         wrfvar = "V"
-    elif comp.lower() == "w":
+#         stagdim = -2
-        wrfvar = "W"
+#     elif comp.lower() == "w":
-        stagdim = -3
+#         wrfvar = "W"
-        
+#         stagdim = -3
-    #wind_data = wrfnc.variables[wrfvar][timeidx,:,:,:]
+#         
-    ncvars = extract_vars(wrfnc, timeidx, vars=wrfvar)
+#     ncvars = extract_vars(wrfnc, timeidx, wrfvar, method, squeeze, cache)
-    wind_data = ncvars[wrfvar]
+#     wind_data = ncvars[wrfvar]
-    return destagger(wind_data, stagdim)
+#     
-
+#     return destagger(wind_data, stagdim)
-def destagger_winds(wrfnc, timeidx=0):
+# 
-    return (destagger_windcomp(wrfnc, "u", timeidx),
+# def destagger_winds(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None):
-            destagger_windcomp(wrfnc, "v", timeidx),
+#     return (destagger_windcomp(wrfnc, "u", timeidx, method, squeeze, cache),
-            destagger_windcomp(wrfnc, "w", timeidx))
+#             destagger_windcomp(wrfnc, "v", timeidx, method, squeeze, cache),
+#             destagger_windcomp(wrfnc, "w", timeidx, method, squeeze, cache))
     

wrf_open/var/src/python/wrf/var/dewpoint.py

@@ -1,13 +1,17 @@
-from wrf.var.extension import computetd
+from .extension import computetd
-from wrf.var.decorators import convert_units
+from .decorators import convert_units, copy_and_set_metadata
-from wrf.var.util import extract_vars
+from .util import extract_vars
 
 __all__ = ["get_dp", "get_dp_2m"]
 
+@copy_and_set_metadata(copy_varname="QVAPOR", name="td", 
+                       description="dew point temperature")
 @convert_units("temp", "c")
-def get_dp(wrfnc, timeidx=0, units="c"):
+def get_dp(wrfnc, timeidx=0, units="c",
+           method="cat", squeeze=True, cache=None):
     
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("P", "PB", "QVAPOR"))
+    varnames=("P", "PB", "QVAPOR")
+    ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
     
     p = ncvars["P"]
     pb = ncvars["PB"]
@@ -20,9 +24,13 @@ def get_dp(wrfnc, timeidx=0, units="c"):
     td = computetd(full_p, qvapor)
     return td
 
+@copy_and_set_metadata(copy_varname="Q2", name="td2", 
+                       description="2m dew point temperature")
 @convert_units("temp", "c")
-def get_dp_2m(wrfnc, timeidx=0, units="c"):
+def get_dp_2m(wrfnc, timeidx=0, units="c",
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("PSFC", "Q2"))
+              method="cat", squeeze=True, cache=None):
+    varnames=("PSFC", "Q2")
+    ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
 
     # Algorithm requires hPa
     psfc = .01*(ncvars["PSFC"])

wrf_open/var/src/python/wrf/var/extension.py

@@ -1,8 +1,8 @@
-import numpy as n
+import numpy as np
 
-from wrf.var.constants import Constants
+from .constants import Constants
-from wrf.var.psadlookup import get_lookup_tables
+from .psadlookup import get_lookup_tables
-from wrf.var._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d,
+from ._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d,
                      f_computeslp, f_computetk, f_computetd, f_computerh, 
                      f_computeabsvort,f_computepvo, f_computeeth, 
                      f_computeuvmet, 
@@ -10,8 +10,8 @@ from wrf.var._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d,
                      f_computesrh, f_computeuh, f_computepw, f_computedbz,
                      f_lltoij, f_ijtoll, f_converteta, f_computectt,
                      f_monotonic, f_filter2d, f_vintrp)
-from wrf.var._wrfcape import f_computecape
+from ._wrfcape import f_computecape
-from wrf.var.decorators import (handle_left_iter, uvmet_left_iter, 
+from .decorators import (handle_left_iter, uvmet_left_iter, 
                                 handle_casting, handle_extract_transpose)
 
 __all__ = ["FortranException", "computeslp", "computetk", "computetd", 
@@ -56,9 +56,9 @@ def interp1d(v_in,z_in,z_out,missingval):
 @handle_casting(arg_idxs=(0,1,2,3))
 @handle_extract_transpose()
 def computeslp(z, t, p, q):
-    t_surf = n.zeros((z.shape[-2], z.shape[-1]), "float64", order="F")
+    t_surf = np.zeros((z.shape[-2], z.shape[-1]), np.float64, order="F")
-    t_sea_level = n.zeros((z.shape[-2], z.shape[-1]), "float64", order="F")
+    t_sea_level = np.zeros((z.shape[-2], z.shape[-1]), np.float64, order="F")
-    level = n.zeros((z.shape[-2], z.shape[-1]), "int32", order="F")
+    level = np.zeros((z.shape[-2], z.shape[-1]), np.int32, order="F")
     
     res = f_computeslp(z, 
                        t, 
@@ -79,7 +79,7 @@ def computetk(pres, theta):
     shape = pres.shape
     res = f_computetk(pres.ravel(order="A"), 
                       theta.ravel(order="A"))
-    res = n.reshape(res, shape)
+    res = np.reshape(res, shape)
     return res
 
 # Note: No left iteration decorator needed with 1D arrays
@@ -89,7 +89,7 @@ def computetd(pressure,qv_in):
     shape = pressure.shape
     res = f_computetd(pressure.ravel(order="A"), 
                       qv_in.ravel(order="A"))
-    res = n.reshape(res, shape)
+    res = np.reshape(res, shape)
     return res
 
 # Note:  No decorator needed with 1D arrays
@@ -100,7 +100,7 @@ def computerh(qv,q,t):
     res = f_computerh(qv.ravel(order="A"),
                       q.ravel(order="A"),
                       t.ravel(order="A"))
-    res = n.reshape(res, shape)
+    res = np.reshape(res, shape)
     return res
 
 @handle_left_iter(3,0, ignore_args=(6,7))
@@ -151,8 +151,8 @@ 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 = n.zeros((lat.shape[-2], lat.shape[-1]), "float64", order="F")
+    longca = np.zeros((lat.shape[-2], lat.shape[-1]), np.float64, order="F")
-    longcb = n.zeros((lon.shape[-2], lon.shape[-1]), "float64", order="F")
+    longcb = np.zeros((lon.shape[-2], lon.shape[-1]), np.float64, order="F")
     rpd = Constants.PI/180.
     
     
@@ -179,7 +179,7 @@ def computeuvmet(u,v,lat,lon,cen_long,cone):
                0)
 
     
-    return n.squeeze(res)
+    return np.squeeze(res)
 
 @handle_left_iter(3,0)
 @handle_casting(arg_idxs=(0,1,2,3))
@@ -238,8 +238,10 @@ def computesrh(u, v, z, ter, top):
 @handle_extract_transpose()
 def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top):
     
-    tem1 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "float64", order="F")
+    tem1 = np.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), np.float64, 
-    tem2 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "float64", order="F")
+                    order="F")
+    tem2 = np.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), np.float64, 
+                    order="F")
     
     res = f_computeuh(zp,
                       mapfct,
@@ -260,7 +262,7 @@ def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top):
 @handle_extract_transpose()
 def computepw(p, tv, qv, ht):
     # Note, dim -3 is height, we only want y and x
-    zdiff = n.zeros((p.shape[-2], p.shape[-1]), "float64", order="F")
+    zdiff = np.zeros((p.shape[-2], p.shape[-1]), np.float64, order="F")
     res = f_computepw(p,
                       tv,
                       qv,
@@ -290,10 +292,10 @@ 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 = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), 
+    flip_cape = np.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), 
-                        "float64", order="F")
+                        np.float64, order="F")
-    flip_cin = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), 
+    flip_cin = np.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), 
-                       "float64", order="F")
+                       np.float64, order="F")
     PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables()
     PSADITMK = PSADITMK.T
     
@@ -327,7 +329,6 @@ def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow):
     
     return (cape, cin)
 
-# TODO:  This should handle lists of coords
 def computeij(map_proj, truelat1, truelat2, stdlon,
                lat1, lon1, pole_lat, pole_lon,
                knowni, knownj, dx, latinc, loninc, lat, lon):
@@ -337,9 +338,8 @@ def computeij(map_proj,truelat1,truelat2,stdlon,
                    knowni,knownj,dx,latinc,loninc,lat,lon,
                    FortranException())
     
-    return res[0],res[1]
+    return res
 
-# TODO:  This should handle lists of coords
 def computell(map_proj, truelat1, truelat2, stdlon, lat1, lon1,
              pole_lat, pole_lon, knowni, knownj, dx, latinc,
              loninc, i, j):
@@ -348,16 +348,15 @@ def computell(map_proj,truelat1,truelat2,stdlon,lat1,lon1,
                    pole_lat,pole_lon,knowni,knownj,dx,latinc,
                    loninc,i,j,FortranException())
     
-    # Want lon,lat
+    return res
-    return res[1],res[0]
 
 @handle_left_iter(3,0, ignore_args=(3,))
 @handle_casting(arg_idxs=(0,1,2))
 @handle_extract_transpose()
 def computeeta(full_t, znu, psfc, ptop):
-    pcalc = n.zeros(full_t.shape, "float64", order="F")
+    pcalc = np.zeros(full_t.shape, np.float64, order="F")
-    mean_t = n.zeros(full_t.shape, "float64", order="F")
+    mean_t = np.zeros(full_t.shape, np.float64, order="F")
-    temp_t = n.zeros(full_t.shape, "float64", order="F")
+    temp_t = np.zeros(full_t.shape, np.float64, order="F")
     
     res = f_converteta(full_t, 
                        znu, 
@@ -392,13 +391,13 @@ def smooth2d(field, passes):
     # copies the original data before modifying it.  This allows the decorator
     # to work properly and also behaves like the other methods.
     
-    if isinstance(field, n.ma.MaskedArray):
+    if isinstance(field, np.ma.MaskedArray):
         missing = field.fill_value
     else:
         missing = Constants.DEFAULT_FILL
     
     field_copy = field.copy()
-    field_tmp = n.zeros(field_copy.shape, field_copy.dtype, order="F")  
+    field_tmp = np.zeros(field_copy.shape, field_copy.dtype, order="F")  
     
     f_filter2d(field_copy, 
                field_tmp, 

wrf_open/var/src/python/wrf/var/geoht.py

@@ -1,11 +1,12 @@
-from wrf.var.constants import Constants
+from .constants import Constants
-from wrf.var.destag import destagger
+from .destag import destagger
-from wrf.var.decorators import convert_units
+from .decorators import convert_units, set_height_metadata
-from wrf.var.util import extract_vars
+from .util import extract_vars, either
 
 __all__ = ["get_geopt", "get_height"]
 
-def _get_geoht(wrfnc, timeidx, height=True, msl=True):
+def _get_geoht(wrfnc, timeidx, height=True, msl=True,
+               method="cat", squeeze=True, cache=None):
     """Return the geopotential in units of m2 s-2 if height is False,
     otherwise return the geopotential height in meters.  If height is True,
     then if msl is True the result will be in MSL, otherwise AGL (the terrain
@@ -13,23 +14,18 @@ def _get_geoht(wrfnc, timeidx, height=True, msl=True):
     
     """
     
-    try:
+    varname = either("PH", "GHT")(wrfnc)
+    if varname == "PH":
         ph_vars = extract_vars(wrfnc, timeidx, varnames=("PH", "PHB", "HGT"))
-    except KeyError:
-        try:
-            ght_vars = extract_vars(wrfnc, timeidx, varnames=("GHT", "HGT_U"))
-        except KeyError:
-            raise RuntimeError("Cannot calculate height with variables in "
-                               "NetCDF file")
-        else:
-            geopt_unstag = ght_vars["GHT"] * Constants.G
-            hgt = destagger(ght_vars["HGT_U"], -1)
-    else:
         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"))
+        geopt_unstag = ght_vars["GHT"] * Constants.G
+        hgt = ght_vars["HGT_M"]
     
     if height:
         if msl:
@@ -46,11 +42,14 @@ def _get_geoht(wrfnc, timeidx, height=True, msl=True):
     else:
         return geopt_unstag
 
-def get_geopt(wrfnc, timeidx=0):
+@set_height_metadata(geopt=True)
-    return _get_geoht(wrfnc, timeidx, False)
+def get_geopt(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None):
+    return _get_geoht(wrfnc, timeidx, False, True, method, squeeze, cache)
 
+@set_height_metadata(geopt=False)
 @convert_units("height", "m")
-def get_height(wrfnc, timeidx=0, msl=True, units="m"):
+def get_height(wrfnc, timeidx=0, msl=True, units="m",
-    z = _get_geoht(wrfnc, timeidx, True, msl)
+               method="cat", squeeze=True, cache=None):
-    return z
+    
+    return _get_geoht(wrfnc, timeidx, True, msl, method, squeeze, cache)
 

wrf_open/var/src/python/wrf/var/helicity.py

@@ -1,43 +1,34 @@
-from wrf.var.constants import Constants
+from .constants import Constants
 
-from wrf.var.extension import computesrh, computeuh
+from .extension import computesrh, computeuh
-from wrf.var.destag import destagger
+from .destag import destagger
-from wrf.var.util import extract_vars, extract_global_attrs
+from .util import extract_vars, extract_global_attrs, either
+from .decorators import copy_and_set_metadata
 
 __all__ = ["get_srh", "get_uh"]
 
-def get_srh(wrfnc, timeidx=0, top=3000.0):
+@copy_and_set_metadata(copy_varname="HGT", name="srh", 
+                       description="storm relative helicity",
+                       units="m-2/s-2")
+def get_srh(wrfnc, timeidx=0, top=3000.0,
+            method="cat", squeeze=True, cache=None):
     # Top can either be 3000 or 1000 (for 0-1 srh or 0-3 srh)
     
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("HGT", "PH", "PHB"))
+    ncvars = extract_vars(wrfnc, timeidx, ("HGT", "PH", "PHB"),
+                          method, squeeze, cache)
     
     ter = ncvars["HGT"]
     ph = ncvars["PH"]
     phb = ncvars["PHB"]
     
-    try:
+    # As coded in NCL, but not sure this is possible
-        u_vars = extract_vars(wrfnc, timeidx, varnames="U")
+    varname = either("U", "UU")(wrfnc)
-    except KeyError:
+    u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
-        try:
+    u = destagger(u_vars[varname], -1) 
-            uu_vars = extract_vars(wrfnc, timeidx, varnames="UU")
+    
-        except KeyError:
+    varname = either("V", "VV")(wrfnc)
-            raise RuntimeError("No valid wind data found in NetCDF file")
+    v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
-        else:
+    v = destagger(v_vars[varname], -2)
-            u = destagger(uu_vars["UU"], -1) # support met_em files
-    else:
-        u = destagger(u_vars["U"], -1)   
-        
-    try:
-        v_vars = extract_vars(wrfnc, timeidx, varnames="V")
-    except KeyError:
-        try:
-            vv_vars = extract_vars(wrfnc, timeidx, varnames="VV")
-        except KeyError:
-            raise RuntimeError("No valid wind data found in NetCDF file")
-        else:
-            v = destagger(vv_vars["VV"], -2) # support met_em files
-    else:
-        v = destagger(v_vars["V"], -2) 
 
     geopt = ph + phb
     geopt_unstag = destagger(geopt, -3)
@@ -53,9 +44,14 @@ def get_srh(wrfnc, timeidx=0, top=3000.0):
     
     return srh
 
-def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0):
+@copy_and_set_metadata(copy_varname="MAPFAC_M", name="updraft_helicity", 
+                       description="updraft helicity",
+                       units="m-2/s-2")
+def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0,
+           method="cat", squeeze=True, cache=None):
     
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("W", "PH", "PHB", "MAPFAC_M"))
+    ncvars = extract_vars(wrfnc, timeidx, ("W", "PH", "PHB", "MAPFAC_M"),
+                          method, squeeze, cache)
     
     wstag = ncvars["W"]
     ph = ncvars["PH"]
@@ -66,29 +62,14 @@ def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0):
     dx = attrs["DX"]
     dy = attrs["DY"]
     
-    try:
+    # As coded in NCL, but not sure this is possible
-        u_vars = extract_vars(wrfnc, timeidx, varnames="U")
+    varname = either("U", "UU")(wrfnc)
-    except KeyError:
+    u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
-        try:
+    u = destagger(u_vars[varname], -1) 
-            uu_vars = extract_vars(wrfnc, timeidx, varnames="UU")
+    
-        except KeyError:
+    varname = either("V", "VV")(wrfnc)
-            raise RuntimeError("No valid wind data found in NetCDF file")
+    v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
-        else:
+    v = destagger(v_vars[varname], -2) 
-            u = destagger(uu_vars["UU"], -1) # support met_em files
-    else:
-        u = destagger(u_vars["U"], -1)   
-        
-    try:
-        v_vars = extract_vars(wrfnc, timeidx, varnames="V")
-    except KeyError:
-        try:
-            vv_vars = extract_vars(wrfnc, timeidx, varnames="VV")
-        except KeyError:
-            raise RuntimeError("No valid wind data found in NetCDF file")
-        else:
-            v = destagger(vv_vars["VV"], -2) # support met_em files
-    else:
-        v = destagger(v_vars["V"], -2)  
     
     zp = ph + phb
     

wrf_open/var/src/python/wrf/var/interp.py

@@ -3,15 +3,15 @@ from math import floor, ceil
 import numpy as n
 import numpy.ma as ma
 
-from wrf.var.extension import (interpz3d,interp2dxy,interp1d,
+from .extension import (interpz3d,interp2dxy,interp1d,
                                smooth2d,monotonic,vintrp)
-from wrf.var.decorators import handle_left_iter, handle_casting
+from .decorators import handle_left_iter, handle_casting
-from wrf.var.util import extract_vars, is_staggered
+from .util import extract_vars, is_staggered
-from wrf.var.constants import Constants, ConversionFactors
+from .constants import Constants, ConversionFactors
-from wrf.var.terrain import get_terrain
+from .terrain import get_terrain
-from wrf.var.geoht import get_height
+from .geoht import get_height
-from wrf.var.temp import get_theta, get_temp, get_eth
+from .temp import get_theta, get_temp, get_eth
-from wrf.var.pressure import get_pressure
+from .pressure import get_pressure
 
 __all__ = ["interplevel", "vertcross", "interpline", "vinterp"]
 
@@ -28,11 +28,13 @@ def interplevel(data3d,zdata,desiredloc,missingval=Constants.DEFAULT_FILL):
     """
     r1 = interpz3d(data3d, zdata, 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, 
@@ -41,7 +43,8 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None,
     
     xdim - maximum x-dimension
     ydim - maximum y-dimension
-    pivot_point - a pivot point of (south_north, west_east) (must be used with angle)
+    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]
@@ -156,7 +159,8 @@ def vertcross(data3d, z, missingval=Constants.DEFAULT_FILL,
     Arguments:
         data3d - a 3D data field
         z - 3D height field
-        pivot_point - a pivot point of (south_north,west_east) (must be used with angle)
+        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 tuple of (south_north1,west_east1)
         end_point - an end point tuple of (south_north2,west_east2)

wrf_open/var/src/python/wrf/var/latlon.py

@@ -1,43 +1,57 @@
 from collections import Iterable
 
-from wrf.var.constants import Constants
+import numpy as np
-from wrf.var.extension import computeij, computell
-from wrf.var.util import extract_vars, extract_global_attrs
 
-__all__ = ["get_lat", "get_lon", "get_ij", "get_ll"]
+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,
+                   iter_left_indexes)
+from .decorators import set_latlon_metadata
 
-def get_lat(wrfnc, timeidx=0):
+if xarray_enabled():
+    from xarray import DataArray
 
-    try:
+__all__ = ["get_lat", "get_lon", "get_ij", "get_ll"]
-        lat_vars = extract_vars(wrfnc, timeidx, varnames="XLAT")
+
-    except KeyError:
+def _lat_varname(stagger):
-        try:
+    if stagger is None or stagger.lower() == "m":
-            latm_vars = extract_vars(wrfnc, timeidx, varnames="XLAT_M")
+        varname = either("XLAT", "XLAT_M")
-        except:
+    elif stagger.lower() == "u" or stagger.lower() == "v":
-            raise RuntimeError("Latitude variable not found in NetCDF file")
+        varname = "XLAT_{}".format(stagger.upper())
-        else:
-            xlat = latm_vars["XLAT_M"]
-    else:
-        xlat = lat_vars["XLAT"]
-    
-    return xlat
-        
-def get_lon(wrfnc, timeidx=0):
-    try:
-        lon_vars = extract_vars(wrfnc, timeidx, varnames="XLONG")
-    except KeyError:
-        try:
-            lonm_vars = extract_vars(wrfnc, timeidx, varnames="XLONG_M")
-        except:
-            raise RuntimeError("Latitude variable not found in NetCDF file")
     else:
-            xlon = lonm_vars["XLONG_M"]
+        raise ValueError("invalid 'stagger' value")
+    
+    return varname
+    
+def _lon_varname(stagger):
+    if stagger is None or stagger.lower() == "m":
+        varname = either("XLONG", "XLONG_M")
+    elif stagger.lower() == "u" or stagger.lower() == "v":
+        varname = "XLONG_{}".format(stagger.upper())
     else:
-        xlon = lon_vars["XLONG"]
+        raise ValueError("invalid 'stagger' value")
+    
+    return varname
+
+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)
+    
+    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)
     
-    return xlon
+    return lon_var[varname]
 
-def _get_proj_params(wrfnc, timeidx):
+def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache):
     if timeidx < 0:
         raise ValueError("'timeidx' must be greater than 0")
     
@@ -64,12 +78,22 @@ def _get_proj_params(wrfnc, timeidx):
         latinc = 0.0
         loninc = 0.0
     
-    xlat = get_lat(wrfnc, timeidx)
+    latvar = _lat_varname(stagger)
-    xlon = get_lon(wrfnc, timeidx)
+    lonvar = _lon_varname(stagger)
     
-    ref_lat = xlat[0,0]
+    lat_timeidx = timeidx
-    ref_lon = xlon[0,0]
+    # Only need all the lats/lons if it's a moving domain file/files
+    if _is_multi_time(timeidx):
+        if not _is_moving_domain(wrfnc, latvar=latvar, lonvar=lonvar):
+            lat_timeidx = 0
         
+    xlat = get_lat(wrfnc, lat_timeidx, stagger, method, squeeze, cache)
+    xlon = get_lon(wrfnc, lat_timeidx, stagger, method, squeeze, cache)
+    
+    ref_lat = np.ravel(xlat[...,0,0])
+    ref_lon = np.ravel(xlon[...,0,0])
+    
+    # Note: fortran index
     known_i = 1.0
     known_j = 1.0
     
@@ -78,20 +102,70 @@ def _get_proj_params(wrfnc, timeidx):
              loninc)
     
 
-# TODO:  longitude and latitude can also be lists
+@set_latlon_metadata(ij=True) 
-def get_ij(wrfnc, longitude, latitude, timeidx=0):
+def get_ij(wrfnc, latitude, longitude, timeidx=0,
-    if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str):
+           stagger=None, method="cat", squeeze=True, cache=None):
-        raise TypeError("'get_ij' is only applicabe for single files")
     
     (map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon,
     pole_lat,pole_lon,known_i,known_j,dx,latinc,
     loninc) = _get_proj_params(wrfnc, timeidx)
     
-    return computeij(map_proj,truelat1,truelat2,stdlon,
+    if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str):
+        lats = np.asarray(latitude)
+        lons = np.asarray(longitude)
+        
+        if lats.ndim > 1:
+            lats = lats.ravel()
+        
+        if lons.ndim > 1:
+            lons = lons.ravel()
+        
+        if (lats.size != lons.size):
+            raise ValueError("'latitude' and 'longitude' "
+                             "must be the same length")
+        
+        
+        if ref_lat.size == 1:
+            outdim = [lats.size, 2]
+            extra_dims = outdim[0]
+        else:
+            # Moving domain will have moving ref_lats/ref_lons
+            outdim = [lats.size, ref_lat.size, 2]
+            extra_dims = outdim[0:2]
+
+        res = np.empty(outdim, np.float64)
+        
+        for left_idxs in iter_left_indexes(extra_dims):
+            left_and_slice_idxs = left_idxs + [slice(None, None, None)]
+            
+            if ref_lat.size == 1:
+                ref_lat_val = ref_lat[0]
+                ref_lon_val = ref_lon[0]
+            else:
+                ref_lat_val = ref_lat[left_idxs[-1]]
+                ref_lon_val = ref_lon[left_idxs[-1]]
+            
+            lat = lats[left_idxs[0]]
+            lon = lons[left_idxs[0]]
+            
+            ij = computeij(map_proj, truelat1, truelat2, stdlon,
+               ref_lat_val, ref_lon_val, pole_lat, pole_lon,
+               known_i, known_j, dx, latinc, loninc,
+               lat, lon)
+            
+            res[left_and_slice_idxs] = ij[:]
+            
+    else:
+        
+        res = computeij(map_proj, truelat1, truelat2, stdlon,
                ref_lat, ref_lon, pole_lat, pole_lon,
-               known_i,known_j,dx,latinc,loninc,latitude,longitude)
+               known_i, known_j, dx, latinc, loninc,
+               latitude, longitude)
     
-# TODO:  i and j can also be lists
+    return res
+
+
+@set_latlon_metadata(ij=False)
 def get_ll(wrfnc, i, j, timeidx=0):
     if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str):
         raise TypeError("'get_ll' is only applicabe for single files")
@@ -100,9 +174,59 @@ def get_ll(wrfnc, i, j, timeidx=0):
     pole_lat,pole_lon,known_i,known_j,dx,latinc,
     loninc) = _get_proj_params(wrfnc, timeidx)
     
-    return computell(map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon,
+    if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str):
-             pole_lat,pole_lon,known_i,known_j,dx,latinc,
+        i_arr = np.asarray(i)
-             loninc,i,j)
+        j_arr = np.asarray(j)
+        
+        if i_arr.ndim > 1:
+            i_arr = i_arr.ravel()
+        
+        if j_arr.ndim > 1:
+            j_arr = j_arr.ravel()
+        
+        if (i_arr.size != j_arr.size):
+            raise ValueError("'i' and 'j' "
+                             "must be the same length")
+            
+        if ref_lat.size == 1:
+            outdim = [i_arr.size, 2]
+            extra_dims = outdim[0]
+        else:
+            # Moving domain will have moving ref_lats/ref_lons
+            outdim = [i_arr.size, ref_lat.size, 2]
+            extra_dims = outdim[0:2]
+
+        res = np.empty(outdim, np.float64)
+        
+        for left_idxs in iter_left_indexes(extra_dims):
+            left_and_slice_idxs = left_idxs + [slice(None, None, None)]
+            
+            if ref_lat.size == 1:
+                ref_lat_val = ref_lat[0]
+                ref_lon_val = ref_lon[0]
+            else:
+                ref_lat_val = ref_lat[left_idxs[-1]]
+                ref_lon_val = ref_lon[left_idxs[-1]]
+            
+            i_val = i_arr[left_idxs[0]]
+            j_val = j_arr[left_idxs[0]]
+            
+            ll = computell(map_proj, truelat1, truelat2,
+                           stdlon, ref_lat_val, ref_lon_val,
+                           pole_lat, pole_lon, known_i, known_j,
+                           dx, latinc, loninc,
+                           i_val, j_val)
+            
+            res[left_and_slice_idxs] = ll[:]
+            
+    else:
+        res = computell(map_proj, truelat1, truelat2,
+                        stdlon, ref_lat, ref_lon,
+                        pole_lat, pole_lon, known_i, known_j,
+                        dx, latinc, loninc,
+                        i_val, j_val)
+        
+    return res
     
     
     

wrf_open/var/src/python/wrf/var/omega.py

@@ -1,14 +1,18 @@
 
-from wrf.var.constants import Constants
+from .constants import Constants
-from wrf.var.destag import destagger
+from .destag import destagger
-from wrf.var.extension import computeomega,computetk
+from .extension import computeomega,computetk
-from wrf.var.util import extract_vars
+from .util import extract_vars
+from .decorators import copy_and_set_metadata
 
 __all__ = ["get_omega"]
 
-def get_omega(wrfnc, timeidx=0):
+@copy_and_set_metadata(copy_varname="T", name="omega", 
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "W", 
+                       description="omega",
-                                                    "PB", "QVAPOR"))
+                       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)
     t = ncvars["T"]
     p = ncvars["P"]
     w = ncvars["W"]

wrf_open/var/src/python/wrf/var/precip.py

@@ -1,6 +1,6 @@
 
 
-from wrf.var.util import extract_vars
+from .util import extract_vars
 
 __all__ = ["get_accum_precip", "get_precip_diff"]
 

wrf_open/var/src/python/wrf/var/pressure.py

@@ -1,30 +1,31 @@
 
-from wrf.var.decorators import convert_units
+from .decorators import convert_units, copy_and_set_metadata
-from wrf.var.util import extract_vars
+from .util import extract_vars, either
 
 __all__ = ["get_pressure", "get_pressure_hpa"]
 
+@copy_and_set_metadata(copy_varname=either("P", "PRES"), name="pressure", 
+                       description="pressure")
 @convert_units("pressure", "pa")
-def get_pressure(wrfnc, timeidx=0, units="pa"):
+def get_pressure(wrfnc, timeidx=0, units="pa",
-
+                 method="cat", squeeze=True, cache=None):
-    try:
+    
-        p_vars = extract_vars(wrfnc, timeidx, varnames=("P", "PB"))
+    varname = either("P", "PRES")(wrfnc)
-    except KeyError:
+    if varname == "P":
-        try:
+        p_vars = extract_vars(wrfnc, timeidx, ("P", "PB"), 
-            pres_vars = extract_vars(wrfnc, timeidx, varnames="PRES")
+                              method, squeeze, cache)
-        except:
-            raise RuntimeError("pressure variable not found in NetCDF file")
-        else:
-            pres = pres_vars["PRES"]
-    else:
         p = p_vars["P"]
         pb = p_vars["PB"]
         pres = p + pb
+    else:
+        pres = extract_vars(wrfnc, timeidx, "PRES", 
+                            method, squeeze, cache)["PRES"]
     
     return pres
 
-def get_pressure_hpa(wrfnc, timeidx=0, units="hpa"):
+def get_pressure_hpa(wrfnc, timeidx=0, units="hpa", 
-    return get_pressure(wrfnc, timeidx, units=units)
+                     method="cat", squeeze=True, cache=None):
+    return get_pressure(wrfnc, timeidx, units, method, squeeze, cache)
 
 
     

wrf_open/var/src/python/wrf/var/projection.py

@@ -1,8 +1,8 @@
 import numpy as np
 import math
 
-from wrf.var.config import basemap_enabled, cartopy_enabled, pyngl_enabled
+from .config import basemap_enabled, cartopy_enabled, pyngl_enabled
-from wrf.var.constants import Constants
+from .constants import Constants
 
 if cartopy_enabled():
     from cartopy import crs

wrf_open/var/src/python/wrf/var/pw.py

@@ -1,13 +1,17 @@
 
-from wrf.var.extension import computepw,computetv,computetk
+from .extension import computepw,computetv,computetk
-from wrf.var.constants import Constants
+from .constants import Constants
-from wrf.var.util import extract_vars
+from .util import extract_vars
+from .decorators import copy_and_set_metadata
 
 __all__ = ["get_pw"]
 
-def get_pw(wrfnc, timeidx=0):
+@copy_and_set_metadata(copy_varname="T", name="pw", 
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "PH", 
+                       description="precipitable water",
-                                                    "PHB", "QVAPOR"))
+                       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)
     
     t = ncvars["T"]
     p = ncvars["P"]

wrf_open/var/src/python/wrf/var/rh.py

@@ -1,12 +1,17 @@
 
-from wrf.var.constants import Constants   
+from .constants import Constants   
-from wrf.var.extension import computerh, computetk
+from .extension import computerh, computetk
-from wrf.var.util import extract_vars
+from .util import extract_vars
+from .decorators import copy_and_set_metadata
 
 __all__ = ["get_rh", "get_rh_2m"]
 
-def get_rh(wrfnc, timeidx=0):
+@copy_and_set_metadata(copy_varname="T", name="rh", 
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR"))
+                       description="relative humidity",
+                       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)
     t = ncvars["T"]
     p = ncvars["P"]
     pb = ncvars["PB"]
@@ -20,8 +25,12 @@ def get_rh(wrfnc, timeidx=0):
     
     return rh
 
-def get_rh_2m(wrfnc, timeidx=0):
+@copy_and_set_metadata(copy_varname="T2", name="rh2", 
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("T2", "PSFC", "Q2"))
+                       description="2m relative humidity",
+                       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)
     t2 = ncvars["T2"]
     psfc = ncvars["PSFC"]
     q2 = ncvars["Q2"]

wrf_open/var/src/python/wrf/var/slp.py

@@ -1,15 +1,19 @@
-from wrf.var.extension import computeslp, computetk
+from .extension import computeslp, computetk
-from wrf.var.constants import Constants
+from .constants import Constants
-from wrf.var.destag import destagger
+from .destag import destagger
-from wrf.var.decorators import convert_units
+from .decorators import convert_units, copy_and_set_metadata
-from wrf.var.util import extract_vars
+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")
 @convert_units("pressure", "hpa")
-def get_slp(wrfnc, timeidx=0, units="hpa"):
+def get_slp(wrfnc, timeidx=0, units="hpa", 
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR",
+            method="cat", squeeze=True, cache=None):
-                                              "PH", "PHB"))
+    varnames=("T", "P", "PB", "QVAPOR", "PH", "PHB")
+    ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
 
     t = ncvars["T"]
     p = ncvars["P"]

wrf_open/var/src/python/wrf/var/temp.py

@@ -1,26 +1,34 @@
 
-from wrf.var.constants import Constants
+from .constants import Constants
-from wrf.var.extension import computetk, computeeth, computetv, computewetbulb
+from .extension import computetk, computeeth, computetv, computewetbulb
-from wrf.var.decorators import convert_units, set_metadata
+from .decorators import convert_units, copy_and_set_metadata
-from wrf.var.util import extract_vars
+from .util import extract_vars
 
 __all__ = ["get_theta", "get_temp", "get_eth", "get_tv", "get_tw",
            "get_tk", "get_tc"]
 
-@set_metadata()
+@copy_and_set_metadata(copy_varname="T", name="theta", 
+                       description="potential temperature")
 @convert_units("temp", "k")
-def get_theta(wrfnc, timeidx=0, units="k"):
+def get_theta(wrfnc, timeidx=0, units="k", 
-    ncvars = extract_vars(wrfnc, timeidx, varnames="T")
+              method="cat", squeeze=True, cache=None):
+    varnames = ("T",)
+    
+    ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
     t = ncvars["T"]
     full_t = t + Constants.T_BASE
 
     return full_t
 
+@copy_and_set_metadata(copy_varname="T", name="temp", 
+                       description="temperature")
 @convert_units("temp", "k")
-def get_temp(wrfnc, timeidx=0, units="k"):
+def get_temp(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, 
+             cache=None):
     """Return the temperature in Kelvin or Celsius"""
     
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB"))
+    varnames=("T", "P", "PB")
+    ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
     t = ncvars["T"]
     p = ncvars["P"]
     pb = ncvars["PB"]
@@ -31,11 +39,15 @@ def get_temp(wrfnc, timeidx=0, units="k"):
     
     return tk
 
+@copy_and_set_metadata(copy_varname="T", name="theta_e", 
+                       description="equivalent potential temperature")
 @convert_units("temp", "k")
-def get_eth(wrfnc, timeidx=0, units="k"):
+def get_eth(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, 
+            cache=None):
     "Return equivalent potential temperature (Theta-e) in Kelvin"
     
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR"))
+    varnames=("T", "P", "PB", "QVAPOR")
+    ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
     t = ncvars["T"]
     p = ncvars["P"]
     pb = ncvars["PB"]
@@ -49,11 +61,15 @@ def get_eth(wrfnc, timeidx=0, units="k"):
     
     return eth
 
+@copy_and_set_metadata(copy_varname="T", name="tv", 
+                       description="virtual temperature")
 @convert_units("temp", "k")
-def get_tv(wrfnc, timeidx=0, units="k"):
+def get_tv(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, 
+           cache=None):
     "Return the virtual temperature (tv) in Kelvin or Celsius"
     
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR"))
+    varnames=("T", "P", "PB", "QVAPOR")
+    ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
     
     t = ncvars["T"]
     p = ncvars["P"]
@@ -68,12 +84,15 @@ def get_tv(wrfnc, timeidx=0, units="k"):
     
     return tv
     
-
+@copy_and_set_metadata(copy_varname="T", name="twb", 
+                       description="wetbulb temperature")
 @convert_units("temp", "k")
-def get_tw(wrfnc, timeidx=0, units="k"):
+def get_tw(wrfnc, timeidx=0, units="k", method="cat", squeeze=True, 
+           cache=None):
     "Return the wetbulb temperature (tw)"
     
-    ncvars = extract_vars(wrfnc, timeidx, varnames=("T", "P", "PB", "QVAPOR"))
+    varnames=("T", "P", "PB", "QVAPOR")
+    ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache)
     t = ncvars["T"]
     p = ncvars["P"]
     pb = ncvars["PB"]
@@ -87,11 +106,13 @@ def get_tw(wrfnc, timeidx=0, units="k"):
     
     return tw
 
-def get_tk(wrfnc, timeidx=0):
+def get_tk(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None):
-    return get_temp(wrfnc, timeidx, units="k")
+    return get_temp(wrfnc, timeidx, units="k", 
+                    method=method, squeeze=squeeze, cache=cache)
 
-def get_tc(wrfnc, timeidx=0):
+def get_tc(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None):
-    return get_temp(wrfnc, timeidx, units="c")
+    return get_temp(wrfnc, timeidx, units="c",
+                    method=method, squeeze=squeeze, cache=cache)
     
     
 

wrf_open/var/src/python/wrf/var/terrain.py

@@ -1,24 +1,18 @@
 
-from wrf.var.decorators import convert_units
+from .decorators import convert_units, copy_and_set_metadata
-from wrf.var.util import extract_vars
+from .util import extract_vars, either
 
 __all__ = ["get_terrain"]
 
+# Need to handle either
+@copy_and_set_metadata(copy_varname=either("HGT", "HGT_M"), name="terrain", 
+                       description="terrain height")
 @convert_units("height", "m")
-def get_terrain(wrfnc, timeidx=0, units="m"):
+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]
 
-    try:
-        hgt_vars = extract_vars(wrfnc, timeidx, varnames="HGT")
-    except KeyError:
-        try:
-            hgt_m_vars = extract_vars(wrfnc, timeidx, varnames="HGT_M")
-        except KeyError:
-            raise RuntimeError("height variable not found in NetCDF file")
-        else:
-            hgt = hgt_m_vars["HGT_M"]
-    else:
-        hgt = hgt_vars["HGT"]
-    
-    return hgt
     
         

wrf_open/var/src/python/wrf/var/times.py

@@ -1,5 +1,5 @@
 
-from wrf.var.util import extract_times
+from .util import extract_times
 
 __all__ = ["get_times"]
 

wrf_open/var/src/python/wrf/var/units.py

@@ -1,5 +1,5 @@
 
-from wrf.var.constants import Constants, ConversionFactors
+from .constants import Constants, ConversionFactors
 
 __all__ = ["check_units", "do_conversion"]
 

wrf_open/var/src/python/wrf/var/util.py

@@ -1,12 +1,11 @@
 from collections import Iterable, Mapping, OrderedDict
 from itertools import product
 import datetime as dt
-import warnings
 
-import numpy as n
+import numpy as np
 
-from wrf.var.config import xarray_enabled
+from .config import xarray_enabled
-from wrf.var.projection import getproj
+from .projection import getproj
 
 if xarray_enabled():
     from xarray import DataArray
@@ -14,7 +13,8 @@ if xarray_enabled():
 __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"]
+           "is_staggered", "get_proj_params", "viewitems", "viewkeys",
+           "viewvalues"]
 
 def _is_multi_time(timeidx):
     if timeidx == -1:
@@ -31,23 +31,77 @@ def _is_mapping(wrfnc):
         return True
     return False
 
-def _is_moving_domain(wrfnc, latvar="XLAT", lonvar="XLONG"):
+def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar):
-    lat1 = wrfnc.variables[latvar][0,:]
+    lats = wrfnc.variables[latvar]
-    lat2 = wrfnc.variables[latvar][-1,:]
+    lons = wrfnc.variables[lonvar]
-    lon1 = wrfnc.variables[lonvar][0,:]
+    # Need to check all times
-    lon2 = wrfnc.variables[lonvar][-1,:]
+    for i in xrange(lats.shape[-3]):
+        start_idxs = [0]*lats.ndim
+        start_idxs[-3] = i
+        
+        end_idxs = [-1]*lats.ndim
+        end_idxs[-3] = i
+        
+        if (first_ll_corner[0] != lats[start_idxs] or 
+            first_ll_corner[1] != lons[start_idxs] or 
+            first_ur_corner[0] != lats[end_idxs] or 
+            first_ur_corner[1] != lons[end_idxs]):
+            return True
+    
+    return False
+
+def _is_moving_domain(wrfseq, varname=None, latvar="XLAT", lonvar="XLONG"):
+    # In case it's just a single file
+    if not _is_multi_file(wrfseq):
+        wrfseq = [wrfseq]
+        
+    # Slow, but safe. Compare the corner points to the first item and see
+    # any move.  User iterator protocol in case wrfseq is not a list/tuple.
+    wrf_iter = iter(wrfseq)
+    
+    first_wrfnc = next(wrf_iter)
+    
+    if varname is not None:
+        coord_names = getattr(first_wrfnc.variables[varname], 
+                              "coordinates").split()
+        lon_coord = coord_names[0]
+        lat_coord = coord_names[1]
+    else:
+        lon_coord = lonvar
+        lat_coord = latvar
+    
+    lats = first_wrfnc.variables[lat_coord]
+    lons = first_wrfnc.variables[lon_coord]
+    
+    zero_idxs = [0] * first_wrfnc.variables[lat_coord].ndim
+    last_idxs = list(zero_idxs)
+    last_idxs[-2:] = [-1]*2
+    
+    lat0 = lats[zero_idxs]
+    lat1 = lats[last_idxs]
+    lon0 = lons[zero_idxs]
+    lon1 = lons[last_idxs]
     
-    if (lat1[0,0] != lat2[0,0] or lat1[-1,-1] != lat2[-1,-1] or 
+    ll_corner = (lat0, lon0)
-        lon1[0,0] != lon2[0,0] or lon1[-1,-1] != lon2[-1,-1]):
+    ur_corner = (lat1, lon1)
+    
+    while True:
+        try:
+            wrfnc = next(wrf_iter)
+        except StopIteration:
+            break
+        else:
+            if _corners_moved(wrfnc, ll_corner, ur_corner, 
+                              lat_coord, lon_coord):
                 return True
     
     return False
 
-def _get_attr(wrfnc, attr):
+def _get_global_attr(wrfnc, attr):
     val = getattr(wrfnc, attr, None)
     
     # PyNIO puts single values in to an array
-    if isinstance(val,n.ndarray):
+    if isinstance(val,np.ndarray):
         if len(val) == 1:
             return val[0] 
     return val
@@ -66,7 +120,7 @@ def extract_global_attrs(wrfnc, attrs):
         else:
             wrfnc = wrfnc[next(wrfnc.iterkeys())]
         
-    return {attr:_get_attr(wrfnc, attr) for attr in attrlist}
+    return {attr:_get_global_attr(wrfnc, attr) for attr in attrlist}
 
 def extract_dim(wrfnc, dim):
     if _is_multi_file(wrfnc):
@@ -80,71 +134,81 @@ def extract_dim(wrfnc, dim):
         return len(d) #netCDF4
     return d # PyNIO
         
-
+def _combine_dict(wrfdict, varname, timeidx, method):
-# TODO
-def _combine_dict(wrfseq, var, timeidx, method):
     """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 = []
+    numkeys = len(wrfdict)  
     
-    multitime = _is_multi_time(timeidx)
+    key_iter = iter(viewkeys(wrfdict))
-    numfiles = len(wrfseq)  
+    first_key = next(key_iter)
+    keynames.append(first_key)
     
-    if not multitime:
+    first_array = _extract_var(wrfdict[first_key], varname, 
-        time_idx_or_slice = timeidx
+                              timeidx, method, squeeze=False)
-    else:
-        time_idx_or_slice = slice(None, None, None)
     
-    keys = (list(x for x in xrange(numfiles)) if not _is_mapping(wrfseq) else 
-            list(key for key in wrfseq.iterkeys()))
-    
-    # Check if 
-    
-    if not xarray_enabled():
-        first_var = wrfseq[keys[0]].variables[var][time_idx_or_slice, :]
-    else:
-        first_var = _build_data_array(wrfseq[keys[0]], var, timeidx)
     
     # Create the output data numpy array based on the first array
-    outdims = [numfiles]
+    outdims = [numkeys]
-    outdims += first_var.shape
+    outdims += first_array.shape
-    outdata = n.zeros(outdims, first_var.dtype)
+    outdata = np.empty(outdims, first_array.dtype)
-    outdata[0,:] = first_var[:]
+    outdata[0,:] = first_array[:]
+    
+    idx = 1
+    while True:
+        try:
+            key = next(key_iter)
+        except StopIteration:
+            break
+        else:
+            keynames.append(key)
+            vardata = _extract_var(wrfdict[key], varname, timeidx, 
+                                   method, squeeze=False)
             
-    for idx, key in enumerate(keys[1:], start=1):
+            if outdata.shape[1:] != vardata.shape:
-        outdata[idx,:] = wrfseq[key].variables[var][time_idx_or_slice, :]
+                raise ValueError("data sequences must have the "
+                                   "same size for all dictionary keys")
+            outdata[idx,:] = vardata.values[:]
+            idx += 1
     
     if not xarray_enabled():
         outarr = outdata
     else:
-        outname = str(first_var.name)
+        outname = str(first_array.name)
-        outcoords = dict(first_var.coords)
+        # Note: assumes that all entries in dict have same coords
-        outdims = ["sequence"] + list(first_var.dims)
+        outcoords = OrderedDict(first_array.coords)
-        outcoords["sequence"] = keys
+        outdims = ["key"] + list(first_array.dims)
-        outattrs = dict(first_var.attrs)
+        outcoords["key"] = keynames
+        outattrs = OrderedDict(first_array.attrs)
         
         outarr = DataArray(outdata, name=outname, coords=outcoords, 
                            dims=outdims, attrs=outattrs)
         
     return outarr
 
+# TODO:  implement in C
+def _cat_files(wrfseq, varname, timeidx):
+    is_moving = _is_moving_domain(wrfseq, varname)
     
-def _cat_files(wrfseq, var, timeidx):
     file_times = extract_times(wrfseq, timeidx)
     
     multitime = _is_multi_time(timeidx) 
 
     time_idx_or_slice = timeidx if not multitime else slice(None, None, None)
     
-    first_var = (_build_data_array(wrfseq[0], var, timeidx)
+    # 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[var][time_idx_or_slice, :])
+                 wrfseq[0].variables[varname][time_idx_or_slice, :])
     
     outdims = [len(file_times)]
     
     # Making a new time dim, so ignore this one
     outdims += first_var.shape[1:]
         
-    outdata = n.zeros(outdims, first_var.dtype)
+    outdata = np.empty(outdims, first_var.dtype)
     
     numtimes = first_var.shape[0]
     startidx = 0
@@ -153,8 +217,14 @@ def _cat_files(wrfseq, var, timeidx):
     outdata[startidx:endidx, :] = first_var[:]
     
     startidx = endidx
-    for wrfnc in wrfseq[1:]:
+    while True:
-        vardata = wrfnc.variables[var][time_idx_or_slice, :]
+        try:
+            wrfnc = next(wrf_iter)
+        except StopIteration:
+            break
+        else:
+            vardata = wrfnc.variables[varname][time_idx_or_slice, :]
+            
             if multitime:
                 numtimes = vardata.shape[0]
             else:
@@ -170,8 +240,8 @@ def _cat_files(wrfseq, var, timeidx):
         # 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)
-        outattrs = dict(first_var.attrs)
+        outattrs = OrderedDict(first_var.attrs)
-        outcoords = dict(first_var.coords)
+        outcoords = OrderedDict(first_var.coords)
         outdimnames = list(first_var.dims)
         outcoords[outdimnames[0]] = file_times # New time dimension values
         
@@ -183,8 +253,9 @@ def _cat_files(wrfseq, var, timeidx):
         
     return outarr
 
-def _join_files(wrfseq, var, timeidx):
+# TODO:  implement in C
-        
+def _join_files(wrfseq, varname, timeidx):
+    is_moving = _is_moving_domain(wrfseq, varname)
     multitime = _is_multi_time(timeidx)
     numfiles = len(wrfseq)  
     
@@ -193,26 +264,36 @@ def _join_files(wrfseq, var, timeidx):
     else:
         time_idx_or_slice = slice(None, None, None)
     
+    # wrfseq might be a generator
+    wrf_iter = iter(wrfseq)
+    
     if xarray_enabled():
-        first_var = _build_data_array(wrfseq[0], var, timeidx)
+        first_var = _build_data_array(next(wrf_iter), varname, 
+                                      timeidx, is_moving)
     else:
-        first_var = wrfseq[0].variables[var][time_idx_or_slice, :]
+        first_var = (next(wrf_iter)).variables[varname][time_idx_or_slice, :]
     
     # Create the output data numpy array based on the first array
     outdims = [numfiles]
     outdims += first_var.shape
-    outdata = n.zeros(outdims, first_var.dtype)
+    outdata = np.empty(outdims, first_var.dtype)
     
     outdata[0,:] = first_var[:]
     
-    for idx, wrfnc in enumerate(wrfseq[1:], 1):
+    idx=1
-        print idx
+    while True:
-        outdata[idx,:] = wrfnc.variables[var][time_idx_or_slice, :]
+        try:
+            wrfnc = next(wrf_iter)
+        except StopIteration:
+            break
+        else:
+            outdata[idx,:] = wrfnc.variables[varname][time_idx_or_slice, :]
+            idx += 1
         
     if xarray_enabled():
         outname = str(first_var.name)
-        outcoords = dict(first_var.coords)
+        outcoords = OrderedDict(first_var.coords)
-        outattrs = dict(first_var.attrs)
+        outattrs = OrderedDict(first_var.attrs)
         # New dimensions
         outdimnames = ["file_idx"] + list(first_var.dims)
         outcoords["file_idx"] = [i for i in xrange(numfiles)]
@@ -225,42 +306,57 @@ def _join_files(wrfseq, var, timeidx):
         
     return outarr
 
-def combine_files(wrfseq, var, timeidx, method="cat", squeeze=True):
+def combine_files(wrfseq, varname, timeidx, method="cat", squeeze=True):
     # Dictionary is unique
     if _is_mapping(wrfseq):
-        outarr = _combine_dict(wrfseq, var, timeidx, method)
+        outarr = _combine_dict(wrfseq, varname, timeidx, method)
-    
+    elif method.lower() == "cat":
-    if method.lower() == "cat":
+        outarr = _cat_files(wrfseq, varname, timeidx)
-        outarr = _cat_files(wrfseq, var, timeidx)
     elif method.lower() == "join":
-        outarr = _join_files(wrfseq, var, timeidx)
+        outarr = _join_files(wrfseq, varname, timeidx)
     else:
         raise ValueError("method must be 'cat' or 'join'")
     
-    if squeeze:
+    return outarr.squeeze() if squeeze else outarr
-        return outarr.squeeze()
-    
-    return outarr
 
 # Note, always returns the full data set with the time dimension included
-def _build_data_array(wrfnc, varname, timeidx):
+def _build_data_array(wrfnc, varname, timeidx, is_moving_domain):
     multitime = _is_multi_time(timeidx)
     var = wrfnc.variables[varname]
     data = var[:]
     attrs = OrderedDict(var.__dict__)
     dimnames = var.dimensions
     
-    # Add the coordinate variables here.
+    # WRF variables will have a coordinates attribute.  MET_EM files have 
-    coord_names = getattr(var, "coordinates").split()
+    # 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
+        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]
         lat_coord = coord_names[1]
     
-    coords = {}
+    coords = OrderedDict()
+    
+    # Handle lat/lon coordinates and projection information if available
+    if lon_coord is not None and lat_coord is not None:
         lon_coord_var = wrfnc.variables[lon_coord]
         lat_coord_var = wrfnc.variables[lat_coord]
     
         if multitime:
-        if _is_moving_domain(wrfnc, lat_coord, lon_coord):
+            if is_moving_domain:
                 # Special case with a moving domain in a multi-time file,
                 # otherwise the projection parameters don't change
                 coords[lon_coord] = lon_coord_var.dimensions, lon_coord_var[:]
@@ -291,38 +387,53 @@ def _build_data_array(wrfnc, varname, timeidx):
             lats, lons, proj_params = get_proj_params(wrfnc, 0, varname)
             proj = getproj(lats=lats, lons=lons, **proj_params)
         
+        attrs["projection"] = proj
         
-    coords[dimnames[0]] = extract_times(wrfnc, timeidx)
     
-    attrs["projection"] = proj
+    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
 
-def _extract_var(wrfnc, varname, timeidx, method, squeeze):
+# Cache is a dictionary of already extracted variables
+def _extract_var(wrfnc, varname, timeidx, method, squeeze, cache):
+    # 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]
+        except KeyError:
+            pass
+    
+    is_moving = _is_moving_domain(wrfnc, varname)
     multitime = _is_multi_time(timeidx)
     multifile = _is_multi_file(wrfnc)
     
     if not multifile:
         if xarray_enabled():
-            return _build_data_array(wrfnc, varname, timeidx)
+            result = _build_data_array(wrfnc, varname, timeidx, is_moving)
         else:
             if not multitime:
-                return wrfnc.variables[varname][timeidx,:]
+                result = wrfnc.variables[varname][timeidx,:]
             else:
-                return wrfnc.variables[varname][:]
+                result = wrfnc.variables[varname][:]
     else:
-        return combine_files(wrfnc, varname, timeidx, method)
+        # Squeeze handled in this routine, so just return it
+        return combine_files(wrfnc, varname, timeidx, method, squeeze)
         
-def extract_vars(wrfnc, timeidx, varnames, method="cat", squeeze=True):
+    return result.squeeze() if squeeze else result
+
+def extract_vars(wrfnc, timeidx, varnames, method="cat", squeeze=True, 
+                 cache=None):
     if isinstance(varnames, str):
         varlist = [varnames]
     else:
         varlist = varnames
     
-    return {var:_extract_var(wrfnc, var, timeidx, method, squeeze)
+    return {var:_extract_var(wrfnc, var, timeidx, method, squeeze, cache)
             for var in varlist}
 
 def _make_time(timearr):
@@ -430,6 +541,78 @@ def get_proj_params(wrfnc, timeidx=0, varname=None):
             wrfnc.variables[lon_coord][time_idx_or_slice,:],
             proj_params)
     
+# Dictionary python 2-3 compatibility stuff
+def viewitems(d):
+    func = getattr(d, "viewitems", None)
+    if func is None:
+        func = d.items
+    return func()
+
+
+def viewkeys(d):
+    func = getattr(d, "viewkeys", None)
+    if func is None:
+        func = d.keys
+    return func()
+
+
+def viewvalues(d):
+    func = getattr(d, "viewvalues", None)
+    if func is None:
+        func = d.values
+    return func()
+
+# Helper utilities for metadata
+class either(object):
+    def __init__(self, *varnames):
+        self.varnames = varnames
+    
+    def __call__(self, wrfnc):
+        if _is_multi_file(wrfnc):
+            wrfnc = next(iter(wrfnc))
+            
+        for varname in self.varnames:
+            if varname in wrfnc:
+                return varname
+        
+        raise ValueError("{} are not valid variable names".format(
+                                                            self.varnames))
+
+class combine_with:
+    # Remove remove_idx first, then insert_idx is applied to removed set
+    def __init__(self, varname, remove_dims=None, insert_before=None, 
+                 new_dimnames=None, new_coords=None):
+        self.varname = varname
+        self.remove_dims = remove_dims
+        self.insert_before = insert_before
+        self.new_dimnames = new_dimnames
+        self.new_coords = new_coords
+    
+    def __call__(self, var):
+        new_dims = list(var.dims)
+        new_coords = OrderedDict(var.coords)
+        
+        if self.remove_dims is not None:
+            for dim in self.remove_dims:
+                new_dims.remove(dim)
+                del new_coords[dim]
+        
+        if self.insert_before is not None:     
+            insert_idx = new_dims.index(self.insert_before)
+            new_dims = (new_dims[0:insert_idx] + self.new_dimnames + 
+                    new_dims[insert_idx:])
+        elif self.new_dimnames is not None:
+            new_dims = self.new_dimnames
+        
+        if self.new_coords is not None:
+            new_coords.update(self.new_coords)
+        
+        return new_dims, new_coords
+        
+        
+
+
+        
     
     
     

wrf_open/var/src/python/wrf/var/uvmet.py

@@ -1,72 +1,39 @@
 from math import fabs, log, tan, sin, cos
 
-from wrf.var.extension import computeuvmet
+from .extension import computeuvmet
-from wrf.var.destag import destagger
+from .destag import destagger
-from wrf.var.constants import Constants
+from .constants import Constants
-from wrf.var.wind import _calc_wspd_wdir
+from .wind import _calc_wspd_wdir
-from wrf.var.decorators import convert_units
+from .decorators import convert_units, set_wind_metadata
-from wrf.var.util import extract_vars, extract_global_attrs
+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:
-        try:
+        varname = either("U", "UU")(wrfnc)
-            u_vars = extract_vars(wrfnc, timeidx, varnames="U")
+        u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
-        except KeyError:
+        u = destagger(u_vars[varname], -1)
-            try:
-                uu_vars = extract_vars(wrfnc, timeidx, varnames="UU")
-            except KeyError:
-                raise RuntimeError("No valid wind data found in NetCDF file")
-            else:
-                u = destagger(uu_vars["UU"], -1) # support met_em files
-        else:
-            u = destagger(u_vars["U"], -1)   
-        
-        try:
-            v_vars = extract_vars(wrfnc, timeidx, varnames="V")
-        except KeyError:
-            try:
-                vv_vars = extract_vars(wrfnc, timeidx, varnames="VV")
-            except KeyError:
-                raise RuntimeError("No valid wind data found in NetCDF file")
-            else:
-                v = destagger(vv_vars["VV"], -2) # support met_em files
-        else:
-            v = destagger(v_vars["V"], -2) 
         
+        varname = either("V", "VV")(wrfnc)
+        v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
+        v = destagger(v_vars[varname], -2)
     else:
-        try:
+        varname = either("U10", "UU")
-            u_vars = extract_vars(wrfnc, timeidx, varnames="U10")
+        u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
-        except KeyError:
+        u = (u_vars[varname] if varname == "U10" else 
-            try:
+             destagger(u_vars[varname][...,0,:,:], -1)) 
-                uu_vars = extract_vars(wrfnc, timeidx, varnames="UU")
-            except KeyError:
-                raise RuntimeError("No valid wind data found in NetCDF file")
-            else:
-                # For met_em files, this just takes the lowest level winds
-                # (3rd dimension from right is bottom_top)
-                u = destagger(uu_vars["UU"][...,0,:,:], -1) # support met_em files
-        else:
-            u = u_vars["U10"] 
         
-        try:
+        varname = either("V10", "VV")
-            v_vars = extract_vars(wrfnc, timeidx, varnames="V10")
+        v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
-        except KeyError:
+        v = (v_vars[varname] if varname == "V10" else 
-            try:
+             destagger(v_vars[varname][...,0,:,:], -2))
-                vv_vars = extract_vars(wrfnc, timeidx, varnames="VV")
-            except KeyError:
-                raise RuntimeError("No valid wind data found in NetCDF file")
-            else:
-                # For met_em files, this just takes the lowest level winds
-                # (3rd dimension from right is bottom_top)
-                v = destagger(vv_vars["VV"][...,0,:,:], -2) # support met_em files
-        else:
-            v = v_vars["V10"]
     
     map_proj_attrs = extract_global_attrs(wrfnc, attrs="MAP_PROJ")
     map_proj = map_proj_attrs["MAP_PROJ"]
@@ -101,37 +68,25 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"):
         else:
             cen_lon = lon_attrs["STAND_LON"]
         
-        try:
-            xlat_m_vars = extract_vars(wrfnc, timeidx, varnames="XLAT_M")
-        except KeyError:
-            try:
-                xlat_vars = extract_vars(wrfnc, timeidx, varnames="XLAT")
-            except KeyError:
-                raise RuntimeError("xlat not found in NetCDF file")
-            else:
-                lat = xlat_vars["XLAT"]
-        else:
-            lat = xlat_m_vars["XLAT_M"]
         
-        try:
+        varname = either("XLAT_M", "XLAT")(wrfnc)
-            xlon_m_vars = extract_vars(wrfnc, timeidx, varnames="XLONG_M")
+        xlat_var = extract_vars(wrfnc, timeidx, varname, 
-        except KeyError:
+                                method, squeeze, cache)
-            try:
+        lat = xlat_var[varname]
-                xlon_vars = extract_vars(wrfnc, timeidx, varnames="XLONG")
+        
-            except KeyError:
+        varname = either("XLONG_M", "XLONG")
-                raise RuntimeError("xlong not found in NetCDF file")
+        xlon_var = extract_vars(wrfnc, timeidx, varname, 
-            else:
+                                method, squeeze, cache)
-                lon = xlon_vars["XLONG"]
+        lon = xlon_var[varname]
-        else:
-            lon = xlon_m_vars["XLONG_M"]
         
         if map_proj == 1:
             if((fabs(true_lat1 - true_lat2) > 0.1) and
                     (fabs(true_lat2 - 90.) > 0.1)): 
                 cone = (log(cos(true_lat1*radians_per_degree)) 
                     - log(cos(true_lat2*radians_per_degree)))
-                cone = cone / (log(tan((45.-fabs(true_lat1/2.))*radians_per_degree)) 
+                cone = (cone / 
-                        - log(tan((45.-fabs(true_lat2/2.))*radians_per_degree))) 
+                        (log(tan((45.-fabs(true_lat1/2.))*radians_per_degree)) 
+                    - log(tan((45.-fabs(true_lat2/2.))*radians_per_degree)))) 
             else:
                 cone = sin(fabs(true_lat1)*radians_per_degree)
         else:
@@ -142,16 +97,27 @@ def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"):
         return res
             
 
-def get_uvmet10(wrfnc, timeidx=0, units="mps"):
+def get_uvmet10(wrfnc, timeidx=0, units="mps",
+                method="cat", squeeze=True, cache=None):
     return get_uvmet(wrfnc, timeidx, True, units)
 
-def get_uvmet_wspd_wdir(wrfnc, timeidx=0, units="mps"):
+@set_wind_metadata(wspd_wdir=True)
-    u,v = get_uvmet(wrfnc, timeidx, False, units)
+def get_uvmet_wspd_wdir(wrfnc, timeidx=0, units="mps",
-    return _calc_wspd_wdir(u, v, units)
+                        method="cat", squeeze=True, cache=None):
+    
+    uvmet = get_uvmet(wrfnc, timeidx, False, units, method, squeeze, cache)
+    
+    return _calc_wspd_wdir(uvmet[0,:], uvmet[1,:], False, units)
+    
+
+@set_wind_metadata(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)
+    
+    return _calc_wspd_wdir(uvmet10[0,:], uvmet10[1,:], True, units)
     
-def get_uvmet10_wspd_wdir(wrfnc, timeidx=0, units="mps"):
-    u,v = get_uvmet10(wrfnc, timeidx, units="mps")
-    return _calc_wspd_wdir(u, v, units)
 
             
         

wrf_open/var/src/python/wrf/var/vorticity.py

@@ -1,12 +1,17 @@
-from wrf.var.extension import computeavo, computepvo
+from .extension import computeavo, computepvo
-from wrf.var.util import extract_vars, extract_global_attrs
+from .util import extract_vars, extract_global_attrs
+from .decorators import copy_and_set_metadata
 
 __all__ = ["get_avo", "get_pvo"]
 
-def get_avo(wrfnc, timeidx=0):
+@copy_and_set_metadata(copy_varname="F", 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"))
+                                              "F"),
+                          method, squeeze, cache)
     
     attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
     u = ncvars["U"]
@@ -22,11 +27,15 @@ def get_avo(wrfnc, timeidx=0):
     return computeavo(u,v,msfu,msfv,msfm,cor,dx,dy)
 
 
-def get_pvo(wrfnc, timeidx=0):
+@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"))
+                                              "F"),
+                          method, squeeze, cache)
     attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
     
     u = ncvars["U"]

wrf_open/var/src/python/wrf/var/wind.py

@@ -1,42 +1,88 @@
 
-import numpy as n
+import numpy as np
 
-from wrf.var.constants import Constants
+from .constants import Constants
-from wrf.var.destag import destagger_windcomp
+from .destag import destagger
-from wrf.var.decorators import convert_units
+from .util import extract_vars, either
+from .decorators import convert_units, set_wind_metadata
 
 __all__ = ["get_u_destag", "get_v_destag", "get_w_destag",
-           "get_destag_wspd_wdir"]
+           "get_destag_wspd_wdir", "get_destag_wspd_wdir10"]
 
 @convert_units("wind", "mps")
 def _calc_wspd(u, v, units="mps"):
-    return n.sqrt(u**2 + v**2)
+    return np.sqrt(u**2 + v**2)
 
 def _calc_wdir(u, v):
-    wdir = 270.0 - n.arctan2(v,u) * (180.0/Constants.PI)
+    wdir = 270.0 - np.arctan2(v,u) * (180.0/Constants.PI)
-    return n.remainder(wdir, 360.0)
+    return np.remainder(wdir, 360.0)
 
-def _calc_wspd_wdir(u, v, units):
+def _calc_wspd_wdir(u, v, two_d, units):
-    return (_calc_wspd(u,v, units), _calc_wdir(u,v))
+    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[:]
+    
+    return res
+
+@set_wind_metadata(wspd_wdir=False)
 @convert_units("wind", "mps")
-def get_u_destag(wrfnc, timeidx=0, units="mps"):
+def get_u_destag(wrfnc, timeidx=0, units="mps",
-    u = destagger_windcomp(wrfnc,"u", timeidx)
+                 method="cat", squeeze=True, cache=None):
+    varname = either("U", "UU")(wrfnc)
+    u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
+    u = destagger(u_vars[varname], -1)
+    
     return u
 
+@set_wind_metadata(wspd_wdir=False)
 @convert_units("wind", "mps")
-def get_v_destag(wrfnc, timeidx=0, units="mps"):
+def get_v_destag(wrfnc, timeidx=0, units="mps",
-    v = destagger_windcomp(wrfnc,"v", timeidx)
+                 method="cat", squeeze=True, cache=None):
+    varname = either("V", "VV")(wrfnc)
+    v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
+    v = destagger(v_vars[varname], -2)
     return v
 
+@set_wind_metadata(wspd_wdir=False)
 @convert_units("wind", "mps")
-def get_w_destag(wrfnc, timeidx=0, units="mps"):
+def get_w_destag(wrfnc, timeidx=0, units="mps", 
-    w = destagger_windcomp(wrfnc,"w", timeidx)
+                 method="cat", squeeze=True, cache=None):
+    w_vars = extract_vars(wrfnc, timeidx, "W", method, squeeze, cache)
+    w = destagger(w_vars["W"], -3)
     return w
 
-def get_destag_wspd_wdir(wrfnc, timeidx=0, units="mps"):
+@set_wind_metadata(wspd_wdir=True)
-    u = destagger_windcomp(wrfnc,"u", timeidx)
+def get_destag_wspd_wdir(wrfnc, timeidx=0, units="mps",
-    v = destagger_windcomp(wrfnc,"v", timeidx)
+                         method="cat", squeeze=True, cache=None):
+    varname = either("U", "UU")(wrfnc)
+    u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
+    u = destagger(u_vars[varname], -1)
+    
+    varname = either("V", "VV")(wrfnc)
+    v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache)
+    v = destagger(v_vars[varname], -2)
+    
+    return _calc_wspd_wdir(u, v, False, units)
+
+@set_wind_metadata(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 = (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 = (v_vars[varname] if varname == "V10" else 
+         destagger(v_vars[varname][...,0,:,:], -2))
     
-    return _calc_wspd_wdir(u,v,units)
+    return _calc_wspd_wdir(u,v,True,units)