now uses extract methods to pull variables and attributes out of netcdf files

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

@@ -2,6 +2,8 @@ import warnings
 
 from extension import *
 import extension
+from util import *
+import util
 from cape import *
 import cape
 from constants import *
@@ -55,6 +57,7 @@ import units
 
 __all__ = ["getvar"]
 __all__ += extension.__all__
+__all__ += util.__all__
 __all__ += cape.__all__
 __all__ += constants.__all__
 __all__ += ctt.__all__

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

@@ -3,19 +3,23 @@ import numpy.ma as ma
 from wrf.var.extension import computetk,computecape
 from wrf.var.destagger import destagger
 from wrf.var.constants import Constants, ConversionFactors
+from wrf.var.util import extract_vars
 
 __all__ = ["get_2dcape", "get_3dcape"]
 
 def get_2dcape(wrfnc, missing=-999999.0, timeidx=0):
     """Return the 2d fields of cape, cin, lcl, and lfc"""
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
-    qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:]
-    ph = wrfnc.variables["PH"][timeidx,:,:,:]
-    phb = wrfnc.variables["PHB"][timeidx,:,:,:]
-    ter = wrfnc.variables["HGT"][timeidx,:,:]
-    psfc = wrfnc.variables["PSFC"][timeidx,:,:]
+    vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", "PH",
+                                              "PHB", "HGT", "PSFC"))
+    
+    t = vars["T"]
+    p = vars["P"]
+    pb = vars["PB"]
+    qv = vars["QVAPOR"]
+    ph = vars["PH"]
+    phb = vars["PHB"]
+    ter = vars["HGT"]
+    psfc = vars["PSFC"]
     
     full_t = t + Constants.T_BASE
     full_p = p + pb
@@ -49,14 +53,17 @@ def get_2dcape(wrfnc, missing=-999999.0, timeidx=0):
 
 def get_3dcape(wrfnc, missing=-999999.0, timeidx=0):
     """Return the 3d fields of cape and cin"""
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
-    qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:]
-    ph = wrfnc.variables["PH"][timeidx,:,:,:]
-    phb = wrfnc.variables["PHB"][timeidx,:,:,:]
-    ter = wrfnc.variables["HGT"][timeidx,:,:]
-    psfc = wrfnc.variables["PSFC"][timeidx,:,:]
+    
+    vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", "PH",
+                                              "PHB", "HGT", "PSFC"))
+    t = vars["T"]
+    p = vars["P"]
+    pb = vars["PB"]
+    qv = vars["QVAPOR"]
+    ph = vars["PH"]
+    phb = vars["PHB"]
+    ter = vars["HGT"]
+    psfc = vars["PSFC"]
     
     full_t = t + Constants.T_BASE
     full_p = p + pb

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

@@ -5,6 +5,7 @@ from wrf.var.extension import computectt, computetk
 from wrf.var.constants import Constants, ConversionFactors
 from wrf.var.destagger import destagger
 from wrf.var.decorators import convert_units
+from wrf.var.util import extract_vars
 
 __all__ = ["get_ctt"]
 
@@ -13,25 +14,31 @@ def get_ctt(wrfnc, units="c", timeidx=0):
     """Return the cloud top temperature.
     
     """
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
-    ph = wrfnc.variables["PH"][timeidx,:,:,:]
-    phb = wrfnc.variables["PHB"][timeidx,:,:,:]
-    ter = wrfnc.variables["HGT"][timeidx,:,:]
-    qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:] * 1000.0 # g/kg
+    vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "PH" ,"PHB",
+                                              "HGT", "QVAPOR"))
+    t = vars["T"]
+    p = vars["P"]
+    pb = vars["PB"]
+    ph = vars["PH"]
+    phb = vars["PHB"]
+    ter = vars["HGT"]
+    qv = vars["QVAPOR"] * 1000.0 # g/kg
     
     haveqci = 1
-    if "QICE" in wrfnc.variables:
-        qice = wrfnc.variables["QICE"][timeidx,:,:,:] * 1000.0 #g/kg
-    else:
+    try:
+        icevars = extract_vars(wrfnc, timeidx, vars="QICE")
+    except KeyError:
         qice = n.zeros(qv.shape, qv.dtype)
         haveqci = 0
-        
-    if "QCLOUD" in wrfnc.variables:
-        qcld = wrfnc.variables["QCLOUD"][timeidx,:,:,:] * 1000.0 #g/kg
     else:
+        qice = icevars["QICE"] * 1000.0 #g/kg
+    
+    try:
+        cldvars = extract_vars(wrfnc, timeidx, vars="QCLOUD")
+    except KeyError:
         raise RuntimeError("'QCLOUD' not found in NetCDF file")
+    else:
+        qcld = cldvars["QCLOUD"] * 1000.0 #g/kg
     
     full_p = p + pb
     p_hpa = full_p * ConversionFactors.PA_TO_HPA
@@ -44,4 +51,4 @@ def get_ctt(wrfnc, units="c", timeidx=0):
     
     ctt = computectt(p_hpa,tk,qice,qcld,qv,ght,ter,haveqci)
     
-    return ctt
+    return ctt

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

@@ -2,6 +2,7 @@ import numpy as n
 
 from wrf.var.extension import computedbz,computetk
 from wrf.var.constants import Constants
+from wrf.var.util import extract_vars
 
 __all__ = ["get_dbz", "get_max_dbz"]
 
@@ -16,22 +17,27 @@ def get_dbz(wrfnc, do_varint=False, do_liqskin=False, timeidx=0):
     as liquid)
     
     """
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
+    vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", 
+                                              "QRAIN"))
+    t = vars["T"]
+    p = vars["P"]
+    pb = vars["PB"]
+    qv = vars["QVAPOR"]
+    qr = vars["QRAIN"]
     
-    qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:]
-    qr = wrfnc.variables["QRAIN"][timeidx,:,:,:]
-    
-    if "QSNOW" in wrfnc.variables:
-        qs = wrfnc.variables["QSNOW"][timeidx,:,:,:]
+    try:
+        snowvars = extract_vars(wrfnc, timeidx, vars="QSNOW")
+    except KeyError:
+        qs = n.zeros(qv.shape, "float")
     else:
-        qs = n.zeros((qv.shape[0], qv.shape[1], qv.shape[2]), "float")
-        
-    if "QGRAUP" in wrfnc.variables:
-        qg = wrfnc.variables["QGRAUP"][timeidx,:,:,:]
+        qs = snowvars["QSNOW"]
+    
+    try:
+        graupvars = extract_vars(wrfnc, timeidx, vars="QGRAUP")
+    except KeyError:
+        qg = n.zeros(qv.shape, "float")
     else:
-        qg = n.zeros((qv.shape[0], qv.shape[1], qv.shape[2]), "float")
+        qg = graupvars["QGRAUP"]
     
     # If qsnow is all 0, set sn0 to 1
     sn0 = 0

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

@@ -1,5 +1,6 @@
 from functools import wraps
 from inspect import getargspec
+from collections import Iterable
 
 from wrf.var.units import do_conversion, check_units
 
@@ -37,7 +38,48 @@ def combine_list_and_times(alg_out_dim):
     def combine_decorator(func):
         @wraps(func)
         def func_wrapper(*args, **kargs):
-            argspec = getargspec(func)
+            # Multiple times?
+            multitime = False
+            if "timeidx" in kargs:
+                if kargs["timeidx"] == 0:
+                    multitime = True
+                    
+            # Multiple files?
+            multifile = False
+            if isinstance(args[0], Iterable) and not isinstance(args[0], str):
+                multifile = True
+            
+            # Single file, single time
+            if not multitime and not multifile:
+                return func(*args, **kargs)
+            
+            # Get the dimensions
+            if multifile:
+                wrffile = args[0][0]
+            else:
+                wrffile = args[0]
+            
+            # TODO:  Add PyNIO support
+            dims = wrffile.dimensions
+            we_size = len(dims["west_east"])
+            sn_size = len(dims["south_north"])
+            bt_size = len(dims["bottom_top"])
+            time_size = len(dims["Time"])
+            
+            
+            if alg_out_dim == 2:
+                pass
+            elif algout_dim == 3:
+                pass
+            elif algout_dim == 4:
+                pass
+            else:
+                raise RuntimeError("invalid algorithm output dimsize")
+            
+            
+                
+             
+            return func(*args, **kargs)
             
         return func_wrapper
     

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

@@ -9,7 +9,8 @@ def destagger(var, stagger_dim):
     Arguments:
         - var is a numpy array for the variable 
         - stagger_dim is the dimension of the numpy array to de-stagger 
-        (e.g. 0, 1, 2)
+        (e.g. 0, 1, 2).  Note:  negative values are acceptable to choose
+        a dimensions from the right hand side (e.g. -1, -2, -3)
     
     """
     var_shape = var.shape
@@ -41,13 +42,13 @@ def destagger(var, stagger_dim):
 def destagger_windcomp(wrfnc, comp, timeidx=0):
     if comp.lower() == "u":
         wrfvar = "U"
-        stagdim = 2
+        stagdim = -1
     elif comp.lower() == "v":
         wrfvar = "V"
-        stagdim = 1
+        stagdim = -2
     elif comp.lower() == "w":
         wrfvar = "W"
-        stagdim = 0
+        stagdim = -3
         
     wind_data = wrfnc.variables[wrfvar][timeidx,:,:,:]
     return destagger(wind_data, stagdim)

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

@@ -1,14 +1,17 @@
 from wrf.var.extension import computetd
 from wrf.var.decorators import convert_units
+from wrf.var.util import extract_vars
 
 __all__ = ["get_dp", "get_dp_2m"]
 
 @convert_units("temp", "c")
 def get_dp(wrfnc, units="c", timeidx=0):
     
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
-    qvapor = wrfnc.variables["QVAPOR"][timeidx,:,:,:]
+    vars = extract_vars(wrfnc, timeidx, vars=("P", "PB", "QVAPOR"))
+    
+    p = vars["P"]
+    pb = vars["PB"]
+    qvapor = vars["QVAPOR"]
     
     # Algorithm requires hPa
     full_p = .01*(p + pb)
@@ -19,10 +22,11 @@ def get_dp(wrfnc, units="c", timeidx=0):
     
 @convert_units("temp", "c")
 def get_dp_2m(wrfnc, units="c", timeidx=0):
+    vars = extract_vars(wrfnc, timeidx, vars=("PSFC", "Q2"))
 
     # Algorithm requires hPa
-    psfc = .01*(wrfnc.variables["PSFC"][timeidx,:,:])
-    q2 = wrfnc.variables["Q2"][timeidx,:,:]
+    psfc = .01*(vars["PSFC"])
+    q2 = vars["Q2"]
     q2[q2 < 0] = 0
     
     td = computetd(psfc, q2)

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

@@ -3,6 +3,7 @@ import numpy as n
 from wrf.var.extension import computeeta
 from wrf.var.constants import Constants
 from wrf.var.decorators import convert_units
+from wrf.var.util import extract_vars
 
 #__all__ = ["convert_eta"]
 __all__ = []

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

@@ -1,6 +1,7 @@
 from wrf.var.constants import Constants
 from wrf.var.destagger import destagger
 from wrf.var.decorators import convert_units
+from wrf.var.util import extract_vars
 
 __all__ = ["get_geopt", "get_height"]
 
@@ -11,16 +12,24 @@ def _get_geoht(wrfnc, height=True, msl=True, timeidx=0):
     height is subtracted).
     
     """
-
-    if "PH" in wrfnc.variables:
-        ph = wrfnc.variables["PH"][timeidx,:,:,:]
-        phb = wrfnc.variables["PHB"][timeidx,:,:,:]
-        hgt = wrfnc.variables["HGT"][timeidx,:,:]
+    
+    try:
+        ph_vars = extract_vars(wrfnc, timeidx, ("PH", "PHB", "HGT"))
+    except KeyError:
+        try:
+            ght_vars = extract_vars(wrfnc, timeidx, ("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, 0)
-    elif "GHT" in wrfnc.variables: # met_em files
-        geopt_unstag = wrfnc.variables["GHT"][timeidx,:,:,:] * Constants.G
-        hgt = destagger(wrfnc.variables["HGT_U"][timidx,:,:], 1)
     
     if height:
         if msl:

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

@@ -2,26 +2,43 @@ from wrf.var.constants import Constants
 
 from wrf.var.extension import computesrh, computeuh
 from wrf.var.destagger import destagger
+from wrf.var.util import extract_vars, extract_global_attrs
 
 __all__ = ["get_srh", "get_uh"]
 
 def get_srh(wrfnc, top=3000.0, timeidx=0):
     # Top can either be 3000 or 1000 (for 0-1 srh or 0-3 srh)
     
-    if "U" in wrfnc.variables:
-        u = destagger(wrfnc.variables["U"][timeidx,:,:,:], 2)
-    elif "UU" in wrfnc.variables:
-        u = destagger(wrfnc.variables["UU"][timeidx,:,:,:], 2) # support met_em files
-        
-    if "V" in wrfnc.variables:
-        v = destagger(wrfnc.variables["V"][timeidx,:,:,:], 1)
-    elif "VV" in wrfnc.variables:
-        v = destagger(wrfnc.variables["VV"][timeidx,:,:,:], 1) 
+    vars = extract_vars(wrfnc, timeidx, vars=("HGT", "PH", "PHB"))
     
-    ter = wrfnc.variables["HGT"][timeidx,:,:]
-    ph = wrfnc.variables["PH"][timeidx,:,:,:]
-    phb = wrfnc.variables["PHB"][timeidx,:,:,:]
+    ter = vars["HGT"]
+    ph = vars["PH"]
+    phb = vars["PHB"]
     
+    try:
+        u_vars = extract_vars(wrfnc, timeidx, vars="U")
+    except KeyError:
+        try:
+            uu_vars = extract_vars(wrfnc, timeidx, vars="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, vars="V")
+    except KeyError:
+        try:
+            vv_vars = extract_vars(wrfnc, timeidx, vars="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, 0)
     
@@ -38,26 +55,43 @@ def get_srh(wrfnc, top=3000.0, timeidx=0):
 
 def get_uh(wrfnc, bottom=2000.0, top=5000.0, timeidx=0):
     
-    if "U" in wrfnc.variables:
-            u = destagger(wrfnc.variables["U"][timeidx,:,:,:], 2)
-    elif "UU" in wrfnc.variables:
-        u = destagger(wrfnc.variables["UU"][timeidx,:,:,:], 2) # support met_em files
+    vars = extract_vars(wrfnc, timeidx, vars=("W", "PH", "PHB", "MAPFAC_M"))
+    
+    wstag = vars["W"]
+    ph = vars["PH"]
+    phb = vars["PHB"]
+    mapfct = vars["MAPFAC_M"]
+    
+    attrs  = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
+    dx = attrs["DX"]
+    dy = attrs["DY"]
+    
+    try:
+        u_vars = extract_vars(wrfnc, timeidx, vars="U")
+    except KeyError:
+        try:
+            uu_vars = extract_vars(wrfnc, timeidx, vars="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)   
         
-    if "V" in wrfnc.variables:
-        v = destagger(wrfnc.variables["V"][timeidx,:,:,:], 1)
-    elif "VV" in wrfnc.variables:
-        v = destagger(wrfnc.variables["VV"][timeidx,:,:,:], 1) 
-    
-    wstag = wrfnc.variables["W"][timeidx,:,:,:]
-    ph = wrfnc.variables["PH"][timeidx,:,:,:]
-    phb = wrfnc.variables["PHB"][timeidx,:,:,:]
+    try:
+        v_vars = extract_vars(wrfnc, timeidx, vars="V")
+    except KeyError:
+        try:
+            vv_vars = extract_vars(wrfnc, timeidx, vars="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
     
-    mapfct = wrfnc.variables["MAPFAC_M"][timeidx,:,:]
-    dx = wrfnc.getncattr("DX")
-    dy = wrfnc.getncattr("DY")
-       
-       
     uh = computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top)
     
     return uh

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

@@ -1,35 +1,59 @@
 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"]
 
 def get_lat(wrfnc, timeidx=0):
-    if "XLAT" in wrfnc.variables:
-        xlat = wrfnc.variables["XLAT"][timeidx,:,:]
-    elif "XLAT_M" in wrfnc.variables:
-        xlat = wrfnc.variables["XLAT_M"][timeidx,:,:]
+    
+    try:
+        lat_vars = extract_vars(wrfnc, timeidx, vars="XLAT")
+    except KeyError:
+        try:
+            latm_vars = extract_vars(wrfnc, timeidx, vars="XLAT_M")
+        except:
+            raise RuntimeError("Latitude variable not found in NetCDF file")
+        else:
+            xlat = latm_vars["XLAT_M"]
+    else:
+        xlat = lat_vars["XLAT"]
     
     return xlat
         
 def get_lon(wrfnc, timeidx=0):
-    if "XLONG" in wrfnc.variables:
-        xlon = wrfnc.variables["XLONG"][timeidx,:,:]
-    elif "XLONG_M" in wrfnc.variables:
-        xlon = wrfnc.variables["XLONG_M"][timeidx,:,:]
+    try:
+        lon_vars = extract_vars(wrfnc, timeidx, vars="XLONG")
+    except KeyError:
+        try:
+            lonm_vars = extract_vars(wrfnc, timeidx, vars="XLONG_M")
+        except:
+            raise RuntimeError("Latitude variable not found in NetCDF file")
+        else:
+            xlon = lonm_vars["XLONG_M"]
+    else:
+        xlon = lon_vars["XLONG"]
         
     return xlon
 
-def get_ij(wrfnc, longitude, latitude, timeidx=0):
-    map_proj = wrfnc.getncattr("MAP_PROJ")
-    truelat1 = wrfnc.getncattr("TRUELAT1")
-    truelat2 = wrfnc.getncattr("TRUELAT2")
-    stdlon = wrfnc.getncattr("STAND_LON")
-    dx = wrfnc.getncattr("DX")
-    dy = wrfnc.getncattr("DY")
-    stdlon = wrfnc.getncattr("STAND_LON")
+def _get_proj_params(wrfnc, timeidx):
+    if timeidx < 0:
+        raise ValueError("'timeidx' must be greater than 0")
+    
+    attrs = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", "TRUELAT1",
+                                               "TRUELAT2", "STAND_LON",
+                                               "DX", "DY", "STAND_LON"))
+    map_proj = attrs["MAP_PROJ"]
+    truelat1 = attrs["TRUELAT1"]
+    truelat2 = attrs["TRUELAT2"]
+    stdlon = attrs["STAND_LON"]
+    dx = attrs["DX"]
+    dy = attrs["DY"]
+    stdlon = attrs["STAND_LON"]
     
     if map_proj == 6:
-        pole_lat = wrfnc.getncattr("POLE_LAT")
-        pole_lon = wrfnc.getncattr("POLE_LON")
+        pole_attrs = extract_global_attrs(wrfnc, attrs=("POLE_LAT", 
+                                                        "POLE_LON"))
+        pole_lat = pole_attrs["POLE_LAT"]
+        pole_lon = pole_attrs["POLE_LON"]
         latinc = (dy*360.0)/2.0/3.141592653589793/6370000.
         loninc = (dx*360.0)/2.0/3.141592653589793/6370000.
     else:
@@ -38,15 +62,8 @@ def get_ij(wrfnc, longitude, latitude, timeidx=0):
         latinc = 0.0
         loninc = 0.0
     
-    if "XLAT" in wrfnc.variables:
-        xlat = wrfnc.variables["XLAT"][timeidx,:,:]
-    elif "XLAT_M" in wrfnc.variables:
-        xlat = wrfnc.variables["XLAT_M"][timeidx,:,:]
-        
-    if "XLONG" in wrfnc.variables:
-        xlon = wrfnc.variables["XLONG"][timeidx,:,:]
-    elif "XLONG_M" in wrfnc.variables:
-        xlon = wrfnc.variables["XLONG_M"][timeidx,:,:]
+    xlat = get_lat(wrfnc, timeidx)
+    xlon = get_lon(wrfnc, timeidx)
     
     ref_lat = xlat[0,0]
     ref_lon = xlon[0,0]
@@ -54,47 +71,32 @@ def get_ij(wrfnc, longitude, latitude, timeidx=0):
     known_i = 1.0
     known_j = 1.0
     
+    return (map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon,
+             pole_lat,pole_lon,known_i,known_j,dx,latinc,
+             loninc)
+    
+
+# TODO:  longitude and latitude can also be lists
+def get_ij(wrfnc, longitude, latitude, timeidx=0):
+    if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str):
+        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,
                ref_lat,ref_lon,pole_lat,pole_lon,
                known_i,known_j,dx,latinc,loninc,latitude,longitude)
     
-
+# TODO:  i and j can also be lists
 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")
     
-    map_proj = wrfnc.getncattr("MAP_PROJ")
-    truelat1 = wrfnc.getncattr("TRUELAT1")
-    truelat2 = wrfnc.getncattr("TRUELAT2")
-    stdlon = wrfnc.getncattr("STAND_LON")
-    dx = wrfnc.getncattr("DX")
-    dy = wrfnc.getncattr("DY")
-    stdlon = wrfnc.getncattr("STAND_LON")
-    
-    if map_proj == 6:
-        pole_lat = wrfnc.getncattr("POLE_LAT")
-        pole_lon = wrfnc.getncattr("POLE_LON")
-        latinc = (dy*360.0)/2.0/3.141592653589793/6370000.
-        loninc = (dx*360.0)/2.0/3.141592653589793/6370000.
-    else:
-        pole_lat = 90.0
-        pole_lon = 0.0
-        latinc = 0.0
-        loninc = 0.0
-    
-    if "XLAT" in wrfnc.variables:
-        xlat = wrfnc.variables["XLAT"][timeidx,:,:]
-    elif "XLAT_M" in wrfnc.variables:
-        xlat = wrfnc.variables["XLAT_M"][timeidx,:,:]
-        
-    if "XLONG" in wrfnc.variables:
-        xlon = wrfnc.variables["XLONG"][timeidx,:,:]
-    elif "XLONG_M" in wrfnc.variables:
-        xlon = wrfnc.variables["XLONG_M"][timeidx,:,:]
-    
-    ref_lat = xlat[0,0]
-    ref_lon = xlon[0,0]
-    
-    known_i = 1.0
-    known_j = 1.0
+    (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 computell(map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon,
              pole_lat,pole_lon,known_i,known_j,dx,latinc,

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

@@ -2,15 +2,17 @@
 from wrf.var.constants import Constants
 from wrf.var.destagger import destagger
 from wrf.var.extension import computeomega,computetk
+from wrf.var.util import extract_vars
 
 __all__ = ["get_omega"]
 
 def get_omega(wrfnc, timeidx=0):
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    w = wrfnc.variables["W"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
-    qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:]
+    vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "W", "PB", "QVAPOR"))
+    t = vars["T"]
+    p = vars["P"]
+    w = vars["W"]
+    pb = vars["PB"]
+    qv = vars["QVAPOR"]
     
     wa = destagger(w, 0)
     full_t = t + Constants.T_BASE

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

@@ -1,22 +1,27 @@
 
 import numpy as n
 
+from wrf.var.util import extract_vars
+
 __all__ = ["get_accum_precip", "get_precip_diff"]
 
 def get_accum_precip(wrfnc, timeidx=0):
-    rainc = wrfnc.variables["RAINC"][timeidx,:,:]
-    rainnc = wrfnc.variables["RAINNC"][timeidx,:,:]
+    vars = extract_vars(wrfnc, timeidx, vars=("RAINC", "RAINNC"))
+    rainc = vars["RAINC"]
+    rainnc = vars["RAINNC"]
     
     rainsum = rainc + rainnc
     
     return rainsum
 
 def get_precip_diff(wrfnc1, wrfnc2, timeidx=0):
-    rainc1 = wrfnc1.variables["RAINC"][timeidx,:,:]
-    rainnc1 = wrfnc1.variables["RAINNC"][timeidx,:,:]
+    vars1 = extract_vars(wrfnc, timeidx, vars=("RAINC", "RAINNC"))
+    vars2 = extract_vars(wrfnc, timeidx, vars=("RAINC", "RAINNC"))
+    rainc1 = vars1["RAINC"]
+    rainnc1 = vars1["RAINNC"]
     
-    rainc2 = wrfnc2.variables["RAINC"][timeidx,:,:]
-    rainnc2 = wrfnc2.variables["RAINNC"][timeidx,:,:]
+    rainc2 = vars2["RAINC"]
+    rainnc2 = vars2["RAINNC"]
     
     rainsum1 = rainc1 + rainnc1
     rainsum2 = rainc2 + rainnc2

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

@@ -1,18 +1,26 @@
 
 from wrf.var.constants import Constants
 from wrf.var.decorators import convert_units
+from wrf.var.util import extract_vars
 
 __all__ = ["get_pressure"]
 
 @convert_units("pressure", "pa")
 def get_pressure(wrfnc, units="hpa", timeidx=0):
 
-    if "P" in wrfnc.variables:
-        p = wrfnc.variables["P"][timeidx,:,:,:]
-        pb = wrfnc.variables["PB"][timeidx,:,:,:]
+    try:
+        p_vars = extract_vars(wrfnc, timeidx, vars=("P", "PB"))
+    except KeyError:
+        try:
+            pres_vars = extract_vars(wrfnc, timeidx, vars="PRES")
+        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
-    elif "PRES" in wrfnc.variables:
-        pres = wrfnc.variables["PRES"][timeidx,:,:,:]
     
     return pres
 

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

@@ -1,17 +1,20 @@
 
 from wrf.var.extension import computepw,computetv,computetk
 from wrf.var.constants import Constants
+from wrf.var.util import extract_vars
 
 __all__ = ["get_pw"]
 
 def get_pw(wrfnc, timeidx=0):
+    vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "PH", "PHB", 
+                                              "QVAPOR"))
     
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
-    ph = wrfnc.variables["PH"][timeidx,:,:,:]
-    phb = wrfnc.variables["PHB"][timeidx,:,:,:]
-    qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:]
+    t = vars["T"]
+    p = vars["P"]
+    pb = vars["PB"]
+    ph = vars["PH"]
+    phb = vars["PHB"]
+    qv = vars["QVAPOR"]
     
     # Change this to use real virtual temperature!
     full_p   =  p + pb

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

@@ -1,15 +1,16 @@
 
 from wrf.var.constants import Constants   
 from wrf.var.extension import computerh, computetk
+from wrf.var.util import extract_vars
 
 __all__ = ["get_rh", "get_rh_2m"]
 
 def get_rh(wrfnc, timeidx=0):
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    #t00 = wrfnc.variables["T00"][timeidx]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
-    qvapor = wrfnc.variables["QVAPOR"][timeidx,:,:,:]
+    vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR"))
+    t = vars["T"]
+    p = vars["P"]
+    pb = vars["PB"]
+    qvapor = vars["QVAPOR"]
     
     full_t = t + Constants.T_BASE
     full_p = p + pb
@@ -20,9 +21,10 @@ def get_rh(wrfnc, timeidx=0):
     return rh
 
 def get_rh_2m(wrfnc, timeidx=0):
-    t2 = wrfnc.variables["T2"][timeidx,:,:]
-    psfc = wrfnc.variables["PSFC"][timeidx,:,:]
-    q2 = wrfnc.variables["Q2"][timeidx,:,:]
+    vars = extract_vars(wrfnc, timeidx, vars=("T2", "PSFC", "Q2"))
+    t2 = vars["T2"]
+    psfc = vars["PSFC"]
+    q2 = vars["Q2"]
     
     q2[q2 < 0] = 0
     rh = computerh(q2, psfc, t2)

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

@@ -2,18 +2,21 @@ from wrf.var.extension import computeslp, computetk
 from wrf.var.constants import Constants
 from wrf.var.destagger import destagger
 from wrf.var.decorators import convert_units
+from wrf.var.util import extract_vars
 
 __all__ = ["get_slp"]
 
 @convert_units("pressure", "hpa")
 def get_slp(wrfnc, units="hpa", timeidx=0):
+    vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR",
+                                              "PH", "PHB"))
 
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
-    qvapor = wrfnc.variables["QVAPOR"][timeidx,:,:,:]
-    ph = wrfnc.variables["PH"][timeidx,:,:,:]
-    phb = wrfnc.variables["PHB"][timeidx,:,:,:]
+    t = vars["T"]
+    p = vars["P"]
+    pb = vars["PB"]
+    qvapor = vars["QVAPOR"]
+    ph = vars["PH"]
+    phb = vars["PHB"]
     
     full_t = t + Constants.T_BASE
     full_p = p + pb

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

@@ -2,12 +2,14 @@
 from wrf.var.constants import Constants
 from wrf.var.extension import computetk, computeeth, computetv, computewetbulb
 from wrf.var.decorators import convert_units
+from wrf.var.util import extract_vars
 
 __all__ = ["get_theta", "get_temp", "get_eth", "get_tv", "get_tw"]
 
 @convert_units("temp", "k")
 def get_theta(wrfnc, units="k", timeidx=0):
-    t = wrfnc.variables["T"][timeidx,:,:,:]
+    vars = extract_vars(wrfnc, timeidx, vars="T")
+    t = vars["T"]
     full_t = t + Constants.T_BASE
 
     return full_t
@@ -16,9 +18,10 @@ def get_theta(wrfnc, units="k", timeidx=0):
 def get_temp(wrfnc, units="k", timeidx=0):
     """Return the temperature in Kelvin or Celsius"""
     
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
+    vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB"))
+    t = vars["T"]
+    p = vars["P"]
+    pb = vars["PB"]
     
     full_t = t + Constants.T_BASE
     full_p = p + pb
@@ -30,10 +33,11 @@ def get_temp(wrfnc, units="k", timeidx=0):
 def get_eth(wrfnc, units="k", timeidx=0):
     "Return equivalent potential temperature (Theta-e) in Kelvin"
     
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
-    qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:]
+    vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR"))
+    t = vars["T"]
+    p = vars["P"]
+    pb = vars["PB"]
+    qv = vars["QVAPOR"]
     
     full_t = t + Constants.T_BASE
     full_p = p + pb
@@ -47,10 +51,12 @@ def get_eth(wrfnc, units="k", timeidx=0):
 def get_tv(wrfnc, units="k", timeidx=0):
     "Return the virtual temperature (tv) in Kelvin or Celsius"
     
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
-    qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:]
+    vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR"))
+    
+    t = vars["T"]
+    p = vars["P"]
+    pb = vars["PB"]
+    qv = vars["QVAPOR"]
     
     full_t = t + Constants.T_BASE
     full_p = p + pb
@@ -64,11 +70,12 @@ def get_tv(wrfnc, units="k", timeidx=0):
 @convert_units("temp", "k")
 def get_tw(wrfnc, units="k", timeidx=0):
     "Return the wetbulb temperature (tw)"
-
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
-    qv = wrfnc.variables["QVAPOR"][timeidx,:,:,:]
+    
+    vars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR"))
+    t = vars["T"]
+    p = vars["P"]
+    pb = vars["PB"]
+    qv = vars["QVAPOR"]
     
     full_t = t + Constants.T_BASE
     full_p = p + pb

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

@@ -1,15 +1,23 @@
 
 from wrf.var.decorators import convert_units
+from wrf.var.util import extract_vars
 
 __all__ = ["get_terrain"]
 
 @convert_units("height", "m")
 def get_terrain(wrfnc, units="m", timeidx=0):
-
-    if "HGT" in wrfnc.variables:
-        hgt = wrfnc.variables["HGT"][timeidx,:,:]
-    elif "HGT_M":
-        hgt = wrfnc.variables["HGT_M"][timeidx,:,:] 
+    
+    try:
+        hgt_vars = extract_vars(wrfnc, timeidx, vars="HGT")
+    except KeyError:
+        try:
+            hgt_m_vars = extract_vars(wrfnc, timeidx, vars="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

@@ -6,6 +6,7 @@ __all__ = ["get_times"]
 def _make_time(timearr):
     return dt.strptime("".join(timearr[:]), "%Y-%m-%d_%H:%M:%S")
 
+# TODO:  handle list of files and multiple times
 def get_times(wrfnc):
     times = wrfnc.variables["Times"][:,:]
     return [_make_time(times[i,:]) for i in xrange(times.shape[0])]    

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

@@ -1,7 +1,7 @@
 
 from wrf.var.constants import Constants, ConversionFactors
 
-__all__ = ["check_units", "do_conversion", "convert_units"]
+__all__ = ["check_units", "do_conversion"]
 
 # Handles unit conversions that only differ by multiplication factors
 def _apply_conv_fact(var, vartype, var_unit, dest_unit):
@@ -119,7 +119,6 @@ def do_conversion(var, vartype, var_unit, dest_unit):
     else:
         return _apply_temp_conv(var, var_unit, dest_unit)
     
-convert_units = do_conversion
 
 
     

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

@@ -0,0 +1,77 @@
+from collections import Iterable
+
+import numpy as n
+
+__all__ = ["extract_vars", "extract_global_attrs", "hold_dim_fixed"]
+
+def _is_multi_time(timeidx):
+    if timeidx == -1:
+        return True
+    return False
+
+def _is_multi_file(wrfnc):
+    if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str):
+        return True
+    return False
+
+def _get_attr(wrfnc, attr):
+    val = getattr(wrfnc, attr)
+    
+    # PyNIO puts single values in to an array
+    if isinstance(val,n.ndarray):
+        if len(val) == 1:
+            return val[0] 
+    return val
+        
+def extract_global_attrs(wrfnc, attrs):
+    if isinstance(attrs, str):
+        attrlist = [attrs]
+    else:
+        attrlist = attrs
+        
+    multifile = _is_multi_file(wrfnc)
+    
+    if multifile:
+        wrfnc = wrfnc[0]
+    
+    return {attr:_get_attr(wrfnc, attr) for attr in attrlist}
+        
+
+def extract_vars(wrfnc, timeidx, vars):
+    if isinstance(vars, str):
+        varlist = [vars]
+    else:
+        varlist = vars
+        
+    multitime = _is_multi_time(timeidx)
+    multifile = _is_multi_file(wrfnc)
+    
+    # Single file, single time
+    if not multitime and not multifile:
+        return {var:wrfnc.variables[var][timeidx,:] for var in varlist}
+    
+def hold_dim_fixed(var, dim, idx):
+    """Generic method to hold a single dimension to a fixed index when 
+    the array shape is unknown.  The values for 'dim' and 'idx' can 
+    be negative.
+    
+    For example, on a 4D array with 'dim' set to 
+    -1 and 'idx' set to 3, this is simply going to do this operation:
+    
+    var[:,:,:,3]
+    
+    """
+    
+    var_shape = var.shape
+    num_dims = len(var_shape) 
+    full_slice = slice(None, None, None)
+    
+    # Make the sequence of slices
+    dim_slices = [full_slice for x in xrange(num_dims)]
+    
+    # Replace the dimension with the fixed index
+    dim_slices[dim] = idx
+
+    return var[dim_slices]
+    
+    

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

@@ -5,6 +5,7 @@ from wrf.var.destagger import destagger
 from wrf.var.constants import Constants
 from wrf.var.wind import _calc_wspd_wdir
 from wrf.var.decorators import convert_units
+from wrf.var.util import extract_vars, extract_global_attrs, hold_dim_fixed
 
 __all__=["get_uvmet", "get_uvmet10", "get_uvmet_wspd_wdir", 
          "get_uvmet10_wspd_wdir"]
@@ -14,27 +15,71 @@ def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0):
     """ Return a tuple of u,v with the winds rotated in to earth space"""
     
     if not ten_m:
-        if "U" in wrfnc.variables:
-            u = destagger(wrfnc.variables["U"][timeidx,:,:,:],  2)
-        elif "UU" in wrfnc.variables:
-            u = destagger(wrfnc.variables["UU"][timeidx,:,:,:], 2) # support met_em files
+        try:
+            u_vars = extract_vars(wrfnc, timeidx, vars="U")
+        except KeyError:
+            try:
+                uu_vars = extract_vars(wrfnc, timeidx, vars="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, vars="V")
+        except KeyError:
+            try:
+                vv_vars = extract_vars(wrfnc, timeidx, vars="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) 
             
-        if "V" in wrfnc.variables:
-            v = destagger(wrfnc.variables["V"][timeidx,:,:,:], 1)
-        elif "VV" in wrfnc.variables:
-            v = destagger(wrfnc.variables["VV"][timeidx,:,:,:], 1) 
     else:
-        if "U10" in wrfnc.variables:
-            u = wrfnc.variables["U10"][timeidx,:,:]
-        elif "UU" in wrfnc.variables:
-            u = destagger(wrfnc.variables["UU"][timeidx,0,:,:], 1) # support met_em files
-        
-        if "V10" in wrfnc.variables:
-            v = wrfnc.variables["V10"][timeidx,:,:]
-        elif "VV" in wrfnc.variables:
-            v = destagger(wrfnc.variables["VV"][timeidx,0,:,:], 0) # support met_em files
+#         # Original code, keeping for reference
+#         if "U10" in wrfnc.variables:
+#             u = wrfnc.variables["U10"][timeidx,:,:]
+#         elif "UU" in wrfnc.variables:
+#             u = destagger(wrfnc.variables["UU"][timeidx,0,:,:], 1) # support met_em files
+#         
+#         if "V10" in wrfnc.variables:
+#             v = wrfnc.variables["V10"][timeidx,:,:]
+#         elif "VV" in wrfnc.variables:
+#             v = destagger(wrfnc.variables["VV"][timeidx,0,:,:], 0) # support met_em files
+        try:
+            u_vars = extract_vars(wrfnc, timeidx, vars="U10")
+        except KeyError:
+            try:
+                uu_vars = extract_vars(wrfnc, timeidx, vars="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(hold_dim_fixed(uu_vars["UU"], -3, 0), -1) # support met_em files
+        else:
+            u = u_vars["U10"] 
         
-    map_proj = wrfnc.getncattr("MAP_PROJ")
+        try:
+            v_vars = extract_vars(wrfnc, timeidx, vars="V10")
+        except KeyError:
+            try:
+                vv_vars = extract_vars(wrfnc, timeidx, vars="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(hold_dim_fixed(vv_vars["VV"], -3, 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"]
     
     # 1 - Lambert
     # 2 - Polar Stereographic
@@ -47,26 +92,50 @@ def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0):
         # No rotation needed for Mercator and Lat/Lon
         return u,v
     elif map_proj in (1,2):
+        lat_attrs = extract_global_attrs(wrfnc, attrs=("CEN_LAT",
+                                                        "TRUELAT1",
+                                                        "TRUELAT2"))
         radians_per_degree = Constants.PI/180.0
         # Rotation needed for Lambert and Polar Stereographic
-        cen_lat  = wrfnc.getncattr("CEN_LAT")
-        if "STAND_LON" in wrfnc.ncattrs():
-            cen_lon = wrfnc.getncattr("STAND_LON")
-        else:
-            cen_lon = wrfnc.getncattr("CEN_LON")
-            
-        true_lat1 = wrfnc.getncattr("TRUELAT1")
-        true_lat2 = wrfnc.getncattr("TRUELAT2")
+        cen_lat  = lat_attrs["CEN_LAT"]
+        true_lat1 = lat_attrs["TRUELAT1"]
+        true_lat2 = lat_attrs["TRUELAT2"]
         
-        if "XLAT_M" in wrfnc.variables:
-            lat = wrfnc.variables["XLAT_M"][timeidx,:,:]
+        try:
+            lon_attrs = extract_global_attrs(wrfnc, attrs="STAND_LON")
+        except AttributeError:
+            try:
+                cen_lon_attrs = extract_global_attrs(wrfnc, attrs="CEN_LON")
+            except AttributeError:
+                raise RuntimeError("longitude attributes not found in NetCDF")
+            else:
+                cen_lon = cen_lon_attrs["CEN_LON"]
+        else:
+            cen_lon = lon_attrs["STAND_LON"]
+              
+        try:
+            xlat_m_vars = extract_vars(wrfnc, timeidx, vars="XLAT_M")
+        except KeyError:
+            try:
+                xlat_vars = extract_vars(wrfnc, timeidx, vars="XLAT")
+            except KeyError:
+                raise RuntimeError("xlat not found in NetCDF file")
+            else:
+                lat = xlat_vars["XLAT"]
         else:
-            lat = wrfnc.variables["XLAT"][timeidx,:,:]
+            lat = xlat_m_vars["XLAT_M"]
             
-        if "XLONG_M" in wrfnc.variables:
-            lon = wrfnc.variables["XLONG_M"][timeidx,:,:]
+        try:
+            xlon_m_vars = extract_vars(wrfnc, timeidx, vars="XLONG_M")
+        except KeyError:
+            try:
+                xlon_vars = extract_vars(wrfnc, timeidx, vars="XLONG")
+            except KeyError:
+                raise RuntimeError("xlong not found in NetCDF file")
+            else:
+                lon = xlon_vars["XLONG"]
         else:
-            lon = wrfnc.variables["XLONG"][timeidx,:,:]
+            lon = xlon_m_vars["XLONG_M"]
             
         if map_proj == 1:
             if((fabs(true_lat1 - true_lat2) > 0.1) and

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

@@ -1,32 +1,46 @@
 from wrf.var.extension import computeavo, computepvo
+from wrf.var.util import extract_vars, extract_global_attrs
 
 __all__ = ["get_avo", "get_pvo"]
 
 def get_avo(wrfnc, timeidx=0):
-    u = wrfnc.variables["U"][timeidx,:,:,:]
-    v = wrfnc.variables["V"][timeidx,:,:,:]
-    msfu = wrfnc.variables["MAPFAC_U"][timeidx,:,:]
-    msfv = wrfnc.variables["MAPFAC_V"][timeidx,:,:]
-    msfm = wrfnc.variables["MAPFAC_M"][timeidx,:,:]
-    cor = wrfnc.variables["F"][timeidx,:,:]
-    dx = wrfnc.getncattr("DX")
-    dy = wrfnc.getncattr("DY")
+    vars = extract_vars(wrfnc, timeidx, vars=("U", "V", "MAPFAC_U",
+                                              "MAPFAC_V", "MAPFAC_M",
+                                              "F"))
+    
+    attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
+    u = vars["U"]
+    v = vars["V"]
+    msfu = vars["MAPFAC_U"]
+    msfv = vars["MAPFAC_V"]
+    msfm = vars["MAPFAC_M"]
+    cor = vars["F"]
+    
+    dx = attrs["DX"]
+    dy = attrs["DY"]
     
     return computeavo(u,v,msfu,msfv,msfm,cor,dx,dy)
 
 
 def get_pvo(wrfnc, timeidx=0):
-    u = wrfnc.variables["U"][timeidx,:,:,:]
-    v = wrfnc.variables["V"][timeidx,:,:,:]
-    t = wrfnc.variables["T"][timeidx,:,:,:]
-    p = wrfnc.variables["P"][timeidx,:,:,:]
-    pb = wrfnc.variables["PB"][timeidx,:,:,:]
-    msfu = wrfnc.variables["MAPFAC_U"][timeidx,:,:]
-    msfv = wrfnc.variables["MAPFAC_V"][timeidx,:,:]
-    msfm = wrfnc.variables["MAPFAC_M"][timeidx,:,:]
-    cor = wrfnc.variables["F"][timeidx,:,:]
-    dx = wrfnc.getncattr("DX")
-    dy = wrfnc.getncattr("DY")
+    vars = extract_vars(wrfnc, timeidx, vars=("U", "V", "T", "P",
+                                              "PB", "MAPFAC_U",
+                                              "MAPFAC_V", "MAPFAC_M",
+                                              "F"))
+    attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY"))
+    
+    u = vars["U"]
+    v = vars["V"]
+    t = vars["T"]
+    p = vars["P"]
+    pb = vars["PB"]
+    msfu = vars["MAPFAC_U"]
+    msfv = vars["MAPFAC_V"]
+    msfm = vars["MAPFAC_M"]
+    cor = vars["F"]
+    
+    dx = attrs["DX"]
+    dy = attrs["DY"]
     
     full_t = t + 300
     full_p = p + pb

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

@@ -4,21 +4,22 @@ import numpy as n
 from wrf.var.constants import Constants
 from wrf.var.destagger import destagger_windcomp
 from wrf.var.decorators import convert_units
+from wrf.var.util import extract_vars
 
 __all__ = ["get_u_destag", "get_v_destag", "get_w_destag",
            "get_destag_wspd_wdir"]
 
-def _calc_wspd(u, v):
+@convert_units("wind", "mps")
+def _calc_wspd(u, v, units="mps"):
     return n.sqrt(u**2 + v**2)
 
 def _calc_wdir(u, v):
     wdir = 270.0 - n.arctan2(v,u) * (180.0/Constants.PI)
     return n.remainder(wdir, 360.0)
 
-@convert_units("wind", "mps")
 def _calc_wspd_wdir(u, v, units="mps"):
     check_units(units, "wind")
-    return (_calc_wspd(u,v), _calc_wdir(u,v))
+    return (_calc_wspd(u,v, units), _calc_wdir(u,v))
 
 @convert_units("wind", "mps")
 def get_u_destag(wrfnc, units="mps", timeidx=0):