Now squeezes single element multidim arrays. Timeidx made a positional argument and removed from kargs check

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

@@ -124,43 +124,43 @@ _FUNC_MAP = {"cape2d" : get_2dcape,
              "ctt" : get_ctt
              }
 
-_VALID_ARGS = {"cape2d" : ["missing", "timeidx"],
-             "cape3d" : ["missing", "timeidx"],
-             "dbz" : ["do_variant", "do_liqskin", "timeidx"],
-             "maxdbz" : ["do_variant", "do_liqskin", "timeidx"],
-             "dp" : ["timeidx", "units"],
-             "dp2m" : ["timeidx", "units"], 
-             "height" : ["msl", "units", "timeidx"],
-             "geopt" : ["timeidx"],
-             "srh" : ["top", "timeidx"],
-             "uhel" : ["bottom", "top", "timeidx"],
-             "omega" : ["timeidx"],
-             "pw" : ["timeidx"],
-             "rh" : ["timeidx"], 
-             "rh2m" : ["timeidx"], 
-             "slp" : ["units", "timeidx"], 
-             "temp" : ["units", "timeidx"], 
-             "theta" : ["units", "timeidx"],
-             "theta_e" : ["timeidx", "units"], 
-             "tv" : ["units", "timeidx"], 
-             "twb" : ["units", "timeidx"], 
-             "terrain" : ["units", "timeidx"],
-             "times" : ["timeidx"], 
-             "uvmet" : ["units", "timeidx"], 
-             "uvmet10" : ["units", "timeidx"],
-             "avo" : ["timeidx"],
-             "pvo" : ["timeidx"],
-             "ua" : ["units", "timeidx"], 
-             "va" : ["units", "timeidx"],
-             "wa" : ["units", "timeidx"],
-             "lat" : ["timeidx"],
-             "lon" : ["timeidx"],
-             "pressure" : ["units", "timeidx"],
-             "wspddir" : ["units", "timeidx"],
-             "wspddir_uvmet" : ["units", "timeidx"],
-             "wspddir_uvmet10" : ["units", "timeidx"],  
-             "ctt" : ["timeidx"],
-             "default" : ["timeidx"]
+_VALID_KARGS = {"cape2d" : ["missing"],
+             "cape3d" : ["missing"],
+             "dbz" : ["do_variant", "do_liqskin"],
+             "maxdbz" : ["do_variant", "do_liqskin"],
+             "dp" : ["units"],
+             "dp2m" : ["units"], 
+             "height" : ["msl", "units"],
+             "geopt" : [],
+             "srh" : ["top"],
+             "uhel" : ["bottom", "top"],
+             "omega" : [],
+             "pw" : [],
+             "rh" : [], 
+             "rh2m" : [], 
+             "slp" : ["units"], 
+             "temp" : ["units"], 
+             "theta" : ["units"],
+             "theta_e" : ["units"], 
+             "tv" : ["units"], 
+             "twb" : ["units"], 
+             "terrain" : ["units"],
+             "times" : [], 
+             "uvmet" : ["units"], 
+             "uvmet10" : ["units"],
+             "avo" : [],
+             "pvo" : [],
+             "ua" : ["units"], 
+             "va" : ["units"],
+             "wa" : ["units"],
+             "lat" : [],
+             "lon" : [],
+             "pressure" : ["units"],
+             "wspddir" : ["units"],
+             "wspddir_uvmet" : ["units"],
+             "wspddir_uvmet10" : ["units"],  
+             "ctt" : [],
+             "default" : []
             }
 
 _ALIASES = {"cape_2d" : "cape2d",
@@ -198,25 +198,21 @@ def _undo_alias(alias):
 
 def _check_kargs(var, kargs):
     for arg in kargs.iterkeys():
-        if arg not in _VALID_ARGS[var]:
+        if arg not in _VALID_KARGS[var]:
             raise ArgumentError("'%s' is an invalid keyword "
-                          "argument for '%s" % (arg, var))
+                          "argument for '%s'" % (arg, var))
             
 
-def getvar(wrfnc, var, **kargs):
+def getvar(wrfnc, var, timeidx=0, **kargs):
     if is_standard_wrf_var(wrfnc, var):
-        if "timeidx" in kargs:
-            timeidx = kargs["timeidx"]
-        else:
-            timeidx = 0 
         return extract_vars(wrfnc, timeidx, var)[var]
     
     actual_var = _undo_alias(var)
-    if actual_var not in _VALID_ARGS:
+    if actual_var not in _VALID_KARGS:
         raise ArgumentError("'%s' is not a valid variable name" % (var))
     
     _check_kargs(actual_var, kargs)
-    return _FUNC_MAP[actual_var](wrfnc,**kargs)
+    return _FUNC_MAP[actual_var](wrfnc,timeidx,**kargs)
     
 
     

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

@@ -8,7 +8,7 @@ from wrf.var.util import extract_vars
 
 __all__ = ["get_2dcape", "get_3dcape"]
 
-def get_2dcape(wrfnc, missing=-999999.0, timeidx=0):
+def get_2dcape(wrfnc, timeidx=0, missing=-999999.0):
     """Return the 2d fields of cape, cin, lcl, and lfc"""
     ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", "PH",
                                               "PHB", "HGT", "PSFC"))
@@ -60,7 +60,7 @@ def get_2dcape(wrfnc, missing=-999999.0, timeidx=0):
     
     return ma.masked_values(res,missing)
 
-def get_3dcape(wrfnc, missing=-999999.0, timeidx=0):
+def get_3dcape(wrfnc, timeidx=0, missing=-999999.0):
     """Return the 3d fields of cape and cin"""
     
     ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", 

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

@@ -10,7 +10,7 @@ from wrf.var.util import extract_vars
 __all__ = ["get_ctt"]
 
 @convert_units("temp", "c")
-def get_ctt(wrfnc, units="c", timeidx=0):
+def get_ctt(wrfnc, timeidx=0, units="c"):
     """Return the cloud top temperature.
     
     """

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

@@ -6,7 +6,7 @@ from wrf.var.util import extract_vars
 
 __all__ = ["get_dbz", "get_max_dbz"]
 
-def get_dbz(wrfnc, do_varint=False, do_liqskin=False, timeidx=0):
+def get_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False):
     """ Return the dbz
     
     do_varint - do variable intercept (if False, constants are used.  Otherwise, 
@@ -58,7 +58,7 @@ def get_dbz(wrfnc, do_varint=False, do_liqskin=False, timeidx=0):
     
     return computedbz(full_p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin)
 
-def get_max_dbz(wrfnc, do_varint=False, do_liqskin=False, timeidx=0):
+def get_max_dbz(wrfnc, timeidx=0, do_varint=False, do_liqskin=False):
     return n.amax(get_dbz(wrfnc, do_varint, do_liqskin, timeidx), 
                   axis=-3)
 

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

@@ -15,7 +15,7 @@ def destagger(var, stagger_dim):
     
     """
     var_shape = var.shape
-    num_dims = len(var_shape)
+    num_dims = var.ndim
     stagger_dim_size = var_shape[stagger_dim]
     
     # Dynamically building the range slices to create the appropriate 
@@ -29,8 +29,8 @@ def destagger(var, stagger_dim):
     slice2 = slice(1, stagger_dim_size, 1)
     
     # default to full slices
-    dim_ranges_1 = [full_slice for x in xrange(num_dims)]
-    dim_ranges_2 = [full_slice for x in xrange(num_dims)]
+    dim_ranges_1 = [full_slice] * num_dims
+    dim_ranges_2 = [full_slice] * num_dims
     
     # for the stagger dim, insert the appropriate slice range
     dim_ranges_1[stagger_dim] = slice1

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

@@ -5,7 +5,7 @@ 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):
+def get_dp(wrfnc, timeidx=0, units="c"):
     
     ncvars = extract_vars(wrfnc, timeidx, vars=("P", "PB", "QVAPOR"))
     
@@ -21,7 +21,7 @@ def get_dp(wrfnc, units="c", timeidx=0):
     return td
     
 @convert_units("temp", "c")
-def get_dp_2m(wrfnc, units="c", timeidx=0):
+def get_dp_2m(wrfnc, timeidx=0, units="c"):
     ncvars = extract_vars(wrfnc, timeidx, vars=("PSFC", "Q2"))
 
     # Algorithm requires hPa

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

@@ -23,6 +23,7 @@ class FortranException(Exception):
     def __call__(self, message):
         raise self.__class__(message)
 
+@handle_left_iter(2,3,0)
 @handle_casting(arg_idxs=(0,1))
 def interpz3d(data3d,zdata,desiredloc,missingval):
     res = f_interpz3d(data3d.T, 
@@ -32,6 +33,7 @@ def interpz3d(data3d,zdata,desiredloc,missingval):
     return res.T
 
 @handle_casting(arg_idxs=(0,1))
+@handle_left_iter(2,2,0)
 def interpz2d(data3d,xy):
     res = f_interp2dxy(data3d.T,
                        xy.T)

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

@@ -5,7 +5,7 @@ from wrf.var.util import extract_vars
 
 __all__ = ["get_geopt", "get_height"]
 
-def _get_geoht(wrfnc, height=True, msl=True, timeidx=0):
+def _get_geoht(wrfnc, timeidx, height=True, msl=True):
     """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
@@ -40,10 +40,10 @@ def _get_geoht(wrfnc, height=True, msl=True, timeidx=0):
         return geopt_unstag
 
 def get_geopt(wrfnc, timeidx=0):
-    return _get_geoht(wrfnc, False, timeidx=timeidx)
+    return _get_geoht(wrfnc, timeidx, False)
 
 @convert_units("height", "m")
-def get_height(wrfnc, msl=True, units="m", timeidx=0):
-    z = _get_geoht(wrfnc, True, msl, timeidx)
+def get_height(wrfnc, timeidx=0, msl=True, units="m"):
+    z = _get_geoht(wrfnc, timeidx, True, msl)
     return z
 

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

@@ -6,7 +6,7 @@ from wrf.var.util import extract_vars, extract_global_attrs
 
 __all__ = ["get_srh", "get_uh"]
 
-def get_srh(wrfnc, top=3000.0, timeidx=0):
+def get_srh(wrfnc, timeidx=0, top=3000.0):
     # Top can either be 3000 or 1000 (for 0-1 srh or 0-3 srh)
     
     ncvars = extract_vars(wrfnc, timeidx, vars=("HGT", "PH", "PHB"))
@@ -53,7 +53,7 @@ def get_srh(wrfnc, top=3000.0, timeidx=0):
     
     return srh
 
-def get_uh(wrfnc, bottom=2000.0, top=5000.0, timeidx=0):
+def get_uh(wrfnc, timeidx=0, bottom=2000.0, top=5000.0):
     
     ncvars = extract_vars(wrfnc, timeidx, vars=("W", "PH", "PHB", "MAPFAC_M"))
     

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

@@ -98,7 +98,7 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None,
                 x1 =  (y1 - intercept)/slope
             
             if( y1 > ydim-1):# intersect outside of top boundary
-                y1 = ydim1
+                y1 = ydim-1
                 x1 =  (y1 - intercept)/slope
     elif start_point is not None and end_point is not None:
         x0 = start_point[0]
@@ -118,7 +118,7 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None,
     npts = int(distance)
     dxy = distance/npts
     
-    xz = n.zeros((npts,2), "float")
+    xy = n.zeros((npts,2), "float")
 
     dx = dx/npts
     dy = dy/npts
@@ -134,8 +134,8 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None,
 def get_vertcross(data3d, z, missingval=-99999, 
                   pivot_point=None,angle=None,start_point=None,end_point=None):
     
-    xdim = z.shape[2]
-    ydim = z.shape[1]
+    xdim = z.shape[-1]
+    ydim = z.shape[-2]
     
     xy = _get_xy(xdim, ydim, pivot_point, angle, start_point, end_point)
     

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

@@ -6,7 +6,7 @@ from wrf.var.util import extract_vars
 __all__ = ["get_pressure"]
 
 @convert_units("pressure", "pa")
-def get_pressure(wrfnc, units="hpa", timeidx=0):
+def get_pressure(wrfnc, timeidx=0, units="hpa"):
 
     try:
         p_vars = extract_vars(wrfnc, timeidx, vars=("P", "PB"))

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

@@ -7,7 +7,7 @@ from wrf.var.util import extract_vars
 __all__ = ["get_slp"]
 
 @convert_units("pressure", "hpa")
-def get_slp(wrfnc, units="hpa", timeidx=0):
+def get_slp(wrfnc, timeidx=0, units="hpa"):
     ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR",
                                               "PH", "PHB"))
 

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

@@ -7,7 +7,7 @@ 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):
+def get_theta(wrfnc, timeidx=0, units="k"):
     ncvars = extract_vars(wrfnc, timeidx, vars="T")
     t = ncvars["T"]
     full_t = t + Constants.T_BASE
@@ -15,7 +15,7 @@ def get_theta(wrfnc, units="k", timeidx=0):
     return full_t
 
 @convert_units("temp", "k")
-def get_temp(wrfnc, units="k", timeidx=0):
+def get_temp(wrfnc, timeidx=0, units="k"):
     """Return the temperature in Kelvin or Celsius"""
     
     ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB"))
@@ -30,7 +30,7 @@ def get_temp(wrfnc, units="k", timeidx=0):
     return tk
 
 @convert_units("temp", "k")
-def get_eth(wrfnc, units="k", timeidx=0):
+def get_eth(wrfnc, timeidx=0, units="k"):
     "Return equivalent potential temperature (Theta-e) in Kelvin"
     
     ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR"))
@@ -48,7 +48,7 @@ def get_eth(wrfnc, units="k", timeidx=0):
     return eth
     
 @convert_units("temp", "k")
-def get_tv(wrfnc, units="k", timeidx=0):
+def get_tv(wrfnc, timeidx=0, units="k"):
     "Return the virtual temperature (tv) in Kelvin or Celsius"
     
     ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR"))
@@ -68,7 +68,7 @@ def get_tv(wrfnc, units="k", timeidx=0):
     
 
 @convert_units("temp", "k")
-def get_tw(wrfnc, units="k", timeidx=0):
+def get_tw(wrfnc, timeidx=0, units="k"):
     "Return the wetbulb temperature (tw)"
     
     ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR"))

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

@@ -5,7 +5,7 @@ from wrf.var.util import extract_vars
 __all__ = ["get_terrain"]
 
 @convert_units("height", "m")
-def get_terrain(wrfnc, units="m", timeidx=0):
+def get_terrain(wrfnc, timeidx=0, units="m"):
     
     try:
         hgt_vars = extract_vars(wrfnc, timeidx, vars="HGT")

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

@@ -1,12 +1,7 @@
 
-import datetime as dt
+from wrf.var.util import extract_times
 
 __all__ = ["get_times"]
 
-def _make_time(timearr):
-    return dt.datetime.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])]
+def get_times(wrfnc,timeidx=0):
+    return extract_times(wrfnc,timeidx)

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

@@ -108,9 +108,9 @@ _TEMP_CONV_METHODS = {"k" : _c_to_k,
                       "f" : _c_to_f
                       }
 
-def check_units(unit, type):
+def check_units(unit, unit_type):
     unitl = unit.lower()
-    if unitl not in _VALID_UNITS[type]:
+    if unitl not in _VALID_UNITS[unit_type]:
         raise ValueError("invalid unit type '%s'" % unit)
 
 def do_conversion(var, vartype, var_unit, dest_unit):

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

@@ -1,9 +1,10 @@
 from collections import Iterable
+import datetime as dt
 
 import numpy as n
 
 __all__ = ["extract_vars", "extract_global_attrs", "extract_dim",
-           "combine_files", "is_standard_wrf_var"]
+           "combine_files", "is_standard_wrf_var", "extract_times"]
 
 def _is_multi_time(timeidx):
     if timeidx == -1:
@@ -53,14 +54,14 @@ def combine_files(wrflist, var, timeidx):
     else:
         time_idx_or_slice = slice(None, None, None)
         
-    first_var = wrflist[0].variables[var][time_idx_or_slice, :]
+    first_var = n.squeeze(wrflist[0].variables[var][time_idx_or_slice, :])
     outdims = [numfiles]
     outdims += first_var.shape
     outarr = n.zeros(outdims, first_var.dtype)
     outarr[0,:] = first_var[:]
     
     for idx, wrfnc in enumerate(wrflist[1:], start=1):
-        outarr[idx,:] = wrfnc.variables[var][time_idx_or_slice, :]
+        outarr[idx,:] = n.squeeze(wrfnc.variables[var][time_idx_or_slice, :])
             
     return outarr
 
@@ -77,9 +78,35 @@ def extract_vars(wrfnc, timeidx, vars):
     if not multitime and not multifile:
         return {var:wrfnc.variables[var][timeidx,:] for var in varlist}
     elif multitime and not multifile:
-        return {var:wrfnc.variables[var][:] for var in varlist}
+        return {var:n.squeeze(wrfnc.variables[var][:]) for var in varlist}
     else:
         return {var:combine_files(wrfnc, var, timeidx) for var in varlist}
+
+def _make_time(timearr):
+    return dt.datetime.strptime("".join(timearr[:]), "%Y-%m-%d_%H:%M:%S")
+
+def _file_times(wrfnc,timeidx):
+    multitime = _is_multi_time(timeidx)
+    if multitime:
+        times = wrfnc.variables["Times"][:,:]
+        for i in xrange(times.shape[0]):
+            yield _make_time(times[i,:])
+    else:
+        times = wrfnc.variables["Times"][timeidx,:]
+        yield _make_time(times)
+        
+
+def extract_times(wrfnc,timeidx):
+    multi_file = _is_multi_file(wrfnc)
+    if not multi_file:
+        wrf_list = [wrfnc]
+    else:
+        wrf_list = wrfnc
+    
+    return [file_time 
+            for wrf_file in wrf_list 
+            for file_time in _file_times(wrf_file,timeidx) ]        
+        
     
 def is_standard_wrf_var(wrfnc, var):
     multifile = _is_multi_file(wrfnc)

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

@@ -11,7 +11,7 @@ __all__=["get_uvmet", "get_uvmet10", "get_uvmet_wspd_wdir",
          "get_uvmet10_wspd_wdir"]
 
 @convert_units("wind", "mps")
-def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0):
+def get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps"):
     """ Return a tuple of u,v with the winds rotated in to earth space"""
     
     if not ten_m:
@@ -144,15 +144,15 @@ def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0):
         return res
             
     
-def get_uvmet10(wrfnc, units="mps", timeidx=0):
-    return get_uvmet(wrfnc, True, units, timeidx)
+def get_uvmet10(wrfnc, timeidx=0, units="mps"):
+    return get_uvmet(wrfnc, timeidx, True, units)
 
-def get_uvmet_wspd_wdir(wrfnc, units="mps", timeidx=0):
-    u,v = get_uvmet(wrfnc, False, units, timeidx)
+def get_uvmet_wspd_wdir(wrfnc, timeidx=0, units="mps"):
+    u,v = get_uvmet(wrfnc, timeidx, False, units)
     return _calc_wspd_wdir(u, v, units)
 
-def get_uvmet10_wspd_wdir(wrfnc, units="mps", timeidx=0):
-    u,v = get_uvmet10(wrfnc, units="mps", timeidx=0)
+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/wind.py

@@ -16,25 +16,25 @@ def _calc_wdir(u, v):
     wdir = 270.0 - n.arctan2(v,u) * (180.0/Constants.PI)
     return n.remainder(wdir, 360.0)
 
-def _calc_wspd_wdir(u, v, units="mps"):
+def _calc_wspd_wdir(u, v, units):
     return (_calc_wspd(u,v, units), _calc_wdir(u,v))
 
 @convert_units("wind", "mps")
-def get_u_destag(wrfnc, units="mps", timeidx=0):
+def get_u_destag(wrfnc, timeidx=0, units="mps"):
     u = destagger_windcomp(wrfnc,"u", timeidx)
     return u
 
 @convert_units("wind", "mps")
-def get_v_destag(wrfnc, units="mps", timeidx=0):
+def get_v_destag(wrfnc, timeidx=0, units="mps"):
     v = destagger_windcomp(wrfnc,"v", timeidx)
     return v
 
 @convert_units("wind", "mps")
-def get_w_destag(wrfnc, units="mps", timeidx=0):
+def get_w_destag(wrfnc, timeidx=0, units="mps"):
     w = destagger_windcomp(wrfnc,"w", timeidx)
     return w
 
-def get_destag_wspd_wdir(wrfnc, units="mps", timeidx=0):
+def get_destag_wspd_wdir(wrfnc, timeidx=0, units="mps"):
     u = destagger_windcomp(wrfnc,"u", timeidx)
     v = destagger_windcomp(wrfnc,"v", timeidx)