Added left iteration decorator

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

@@ -1,9 +1,10 @@
+import numpy as n
 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
+from wrf.var.util import extract_vars, hold_dim_fixed
 
 __all__ = ["get_2dcape", "get_3dcape"]
 
@@ -38,22 +39,30 @@ def get_2dcape(wrfnc, missing=-999999.0, timeidx=0):
     
     cape_res,cin_res = computecape(p_hpa,tk,qv,z,ter,psfc_hpa,
                                    missing,i3dflag,ter_follow)
+     
+    cape = hold_dim_fixed(cape_res, -3, 0)
+    cin = hold_dim_fixed(cin_res, -3, 0)
+    lcl = hold_dim_fixed(cin_res, -3, 1)
+    lfc = hold_dim_fixed(cin_res, -3, 2)
     
-    cape = cape_res[0,:,:]
-    cin = cin_res[0,:,:]
-    lcl = cin_res[1,:,:]
-    lfc = cin_res[2,:,:]
+    resdim = [4]
+    resdim += cape.shape
     
-    return (ma.masked_values(cape,missing), 
-            ma.masked_values(cin,missing), 
-            ma.masked_values(lcl,missing), 
-            ma.masked_values(lfc,missing))
+    # Make a new output array for the result
+    res = n.zeros((resdim), cape.dtype)
+    
+    res[0,:] = cape[:]
+    res[1,:] = cin[:]
+    res[2,:] = lcl[:]
+    res[3,:] = lfc[:]
+    
+    return ma.masked_values(res,missing)
 
 def get_3dcape(wrfnc, missing=-999999.0, timeidx=0):
     """Return the 3d fields of cape and cin"""
     
-    ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", "PH",
-                                              "PHB", "HGT", "PSFC"))
+    ncvars = extract_vars(wrfnc, timeidx, vars=("T", "P", "PB", "QVAPOR", 
+                                                "PH", "PHB", "HGT", "PSFC"))
     t = ncvars["T"]
     p = ncvars["P"]
     pb = ncvars["PB"]
@@ -80,8 +89,16 @@ def get_3dcape(wrfnc, missing=-999999.0, timeidx=0):
     
     cape,cin = computecape(p_hpa,tk,qv,z,ter,psfc_hpa,
                            missing,i3dflag,ter_follow)
-    return (ma.masked_values(cape, missing), 
-            ma.masked_values(cin, missing))
+    
+    # Make a new output array for the result
+    resdim = [2]
+    resdim += cape.shape
+    res = n.zeros((resdim), cape.dtype)
+    
+    res[0,:] = cape[:]
+    res[1,:] = cin[:]
+    
+    return ma.masked_values(res, missing)
     
     
     

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

@@ -1,30 +1,36 @@
 from functools import wraps
 from inspect import getargspec
-from collections import Iterable
+from itertools import product
+
+import numpy as n
 
 from wrf.var.units import do_conversion, check_units
 
-__all__ = ["convert_units"]
+__all__ = ["convert_units", "handle_left_iter"]
 
 def convert_units(unit_type, alg_unit):
+    """A decorator that applies unit conversion to a function's output array.
+    
+    Arguments:
+    
+        - unit_type - the unit type category (wind, pressure, etc)
+        - alg_unit - the units that the function returns by default
+    
+    """
     def convert_decorator(func):
         @wraps(func)
         def func_wrapper(*args, **kargs):
             # If units are provided to the method call, use them.  
             # Otherwise, need to parse the argspec to find what the default 
-            # value is.
+            # value is since wraps does not preserve this.
             if ("units" in kargs):
                 desired_units = kargs["units"]
             else:
                 argspec = getargspec(func)
-                print argspec
-                arg_idx_from_right = len(argspec.args) - argspec.args.index("units")
+                arg_idx_from_right = (len(argspec.args) 
+                                      - argspec.args.index("units"))
                 desired_units = argspec.defaults[-arg_idx_from_right]
                 
-                #print desired_idx
-                #desired_units = argspec.defaults[desired_idx]
-                print desired_units
-                
             check_units(desired_units, unit_type)
             
             # Unit conversion done here
@@ -34,4 +40,124 @@ def convert_units(unit_type, alg_unit):
     
     return convert_decorator
 
+def _left_indexes(dims):
+    """A generator which yields the iteration tuples for a sequence of 
+    dimensions sizes.
+    
+    For example, if an array shape is (3,3), then this will yield:
+    
+    (0,0), (0,1), (1,0), (1,1)
+    
+    Arguments:
+    
+        - dims - a sequence of dimensions sizes (e.g. ndarry.shape)
+    
+    """
+    arg = [xrange(dim) for dim in dims]
+    for idxs in product(*arg):
+        yield idxs
+
+def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1,
+                   ref_var_name=None, alg_out_fixed_dims=(), 
+                   ignore_args=(), ignore_kargs={}):
+    """Decorator to handle iterating over leftmost dimensions when using 
+    multiple files and/or multiple times.
+    
+    Arguments:
+        - alg_out_num_dims - the return dimension count from the numerical 
+        algorithm
+        - ref_var_expected_dims - the number of dimensions that the Fortran 
+        algorithm is expecting for the reference variable
+        - ref_var_idx - the index in args used as the reference variable for 
+        calculating leftmost dimensions
+        - ref_var_name - the keyword argument name for kargs used as the 
+        reference varible for calculating leftmost dimensions
+        - alg_out_fixed_dims - additional fixed dimension sizes for the 
+        numerical algorithm (e.g. uvmet has a fixed left dimsize of 2)
+        - ignore_args - indexes of any arguments which should be passed 
+        directly without any slicing
+        - ignore_kargs - keys of any keyword arguments which should be passed
+        directly without slicing
+    
+    """
+    def indexing_decorator(func):
+        @wraps(func)
+        def func_wrapper(*args, **kargs):
+            if ref_var_idx >= 0:
+                ref_var = args[ref_var_idx]
+            else:
+                ref_var = kargs[ref_var_name]
+            
+            ref_var_shape = ref_var.shape
+            extra_dim_num = len(ref_var_shape) - ref_var_expected_dims
+            
+            # No special left side iteration, return the function result
+            if (extra_dim_num == 0):
+                return func(*args, **kargs)
+            
+            # Start by getting the left-most 'extra' dims
+            outdims = [ref_var_shape[x] for x in xrange(extra_dim_num)]
+            extra_dims = list(outdims) # Copy the left-most dims for iteration
+            
+            # Append the right-most dimensions
+            # Note: numerical algorithm output dim count might be less than 
+            # or greater than reference variable
+            if alg_out_num_dims <= ref_var_expected_dims:
+                outdims += [ref_var_shape[x] 
+                        for x in xrange(-alg_out_num_dims,0,1)]
+            else: 
+                # numerical algorithm has greater dim count than reference, 
+                # Note: user must provide this information to decorator
+                right_dims = [dim for dim in alg_out_fixed_dims]
+                right_dims += [ref_var_shape[x] 
+                        for x in xrange(-alg_out_num_dims,0,1)]
+                outdims += right_dims
+                
+            
+            out_inited = False
+            for left_idxs in _left_indexes(extra_dims):
+                # Make the left indexes plus a single slice object
+                # The single slice will handle all the dimensions to
+                # the right (e.g. [1,1,:])
+                left_and_slice_idxs = ([x for x in left_idxs] + 
+                                       [slice(None, None, None)])
+                
+                # Slice the args if applicable
+                new_args = [arg[left_and_slice_idxs] 
+                            if i not in ignore_args else arg 
+                            for i,arg in enumerate(args)]
+                          
+                
+                # 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()}
+                
+                
+                # Call the numerical routine
+                res = func(*new_args, **new_kargs)
+                
+                if isinstance(res, n.ndarray):
+                    # Output array
+                    if not out_inited:
+                        output = n.zeros(outdims, ref_var.dtype)
+                        out_inited = True
+                    
+                    output[left_and_slice_idxs] = res[:]
+                    
+                else:   # This should be a list or a tuple (cape)
+                    if not out_inited:
+                        output = [n.zeros(outdims, ref_var.dtype)]
+                        out_inited = True
+                    
+                    for outarr in output:
+                        outarr[left_and_slice_idxs] = res[:]
+                
+            
+            return output
+        
+        return func_wrapper
+    
+    return indexing_decorator
+
 

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

@@ -11,6 +11,7 @@ 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)
 from wrf.var._wrfcape import f_computecape
+from wrf.var.decorators import handle_left_iter
 
 __all__ = ["FortranException", "computeslp", "computetk", "computetd", 
            "computerh", "computeavo", "computepvo", "computeeth", 
@@ -43,10 +44,11 @@ def interp1d(v_in,z_in,z_out,missingval):
     
     return res.astype("float32")
 
+@handle_left_iter(2,3,0)
 def computeslp(z,t,p,q):
-    t_surf = n.zeros((z.shape[1], z.shape[2]), "float64")
-    t_sea_level = n.zeros((z.shape[1], z.shape[2]), "float64")
-    level = n.zeros((z.shape[1], z.shape[2]), "int32")
+    t_surf = n.zeros((z.shape[-2], z.shape[-1]), "float64")
+    t_sea_level = n.zeros((z.shape[-2], z.shape[-1]), "float64")
+    level = n.zeros((z.shape[-2], z.shape[-1]), "int32")
     
     res = f_computeslp(z.astype("float64").T, 
                        t.astype("float64").T, 
@@ -59,6 +61,7 @@ def computeslp(z,t,p,q):
     
     return res.astype("float32").T
 
+@handle_left_iter(3,3,0)
 def computetk(pres, theta):
     # No need to transpose here since operations on 1D array
     shape = pres.shape
@@ -67,12 +70,14 @@ def computetk(pres, theta):
     res = n.reshape(res, shape, "A")
     return res.astype("float32")
 
+# Note: No decorator needed with 1D arrays
 def computetd(pressure,qv_in):
     shape = pressure.shape
     res = f_computetd(pressure.astype("float64").flatten("A"), qv_in.astype("float64").flatten("A"))
     res = n.reshape(res, shape, "A")
     return res.astype("float32")
 
+# Note:  No decorator needed with 1D arrays
 def computerh(qv,q,t):
     shape = qv.shape
     res = f_computerh(qv.astype("float64").flatten("A"),
@@ -81,6 +86,7 @@ def computerh(qv,q,t):
     res = n.reshape(res, shape, "A")
     return res.astype("float32")
 
+@handle_left_iter(3,3,0, ignore_args=(6,7))
 def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy):
     res = f_computeabsvort(u.astype("float64").T,
                            v.astype("float64").T,
@@ -93,6 +99,7 @@ def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy):
     
     return res.astype("float32").T
 
+@handle_left_iter(3,3,0, ignore_args=(8,9))
 def computepvo(u,v,theta,prs,msfu,msfv,msfm,cor,dx,dy):
     
     res = f_computepvo(u.astype("float64").T,
@@ -107,7 +114,8 @@ def computepvo(u,v,theta,prs,msfu,msfv,msfm,cor,dx,dy):
                        dy)
     
     return res.astype("float32").T
-    
+
+@handle_left_iter(3,3,0)
 def computeeth(qv, tk, p):
     
     res = f_computeeth(qv.astype("float64").T,
@@ -116,16 +124,17 @@ def computeeth(qv, tk, p):
     
     return res.astype("float32").T
 
+@handle_left_iter(4,2,2, ignore_args=(4,5), alg_out_fixed_dims=(2,))
 def computeuvmet(u,v,lat,lon,cen_long,cone):
-    longca = n.zeros((lat.shape[0], lat.shape[1]), "float64")
-    longcb = n.zeros((lon.shape[0], lon.shape[1]), "float64")
+    longca = n.zeros((lat.shape[-2], lat.shape[-1]), "float64")
+    longcb = n.zeros((lon.shape[-2], lon.shape[-1]), "float64")
     rpd = Constants.PI/180.
     
     
     # Make the 2D array a 3D array with 1 dimension
     if u.ndim != 3:
-        u = u.reshape((1,u.shape[0], u.shape[1]))
-        v = v.reshape((1,v.shape[0], v.shape[1]))
+        u = u.reshape((1,u.shape[-2], u.shape[-1]))
+        v = v.reshape((1,v.shape[-2], v.shape[-1]))
 
     # istag will always be false since winds are destaggered already
     # Missing values don't appear to be used, so setting to 0
@@ -147,6 +156,7 @@ def computeuvmet(u,v,lat,lon,cen_long,cone):
     
     return res.astype("float32").T
 
+@handle_left_iter(3,3,0)
 def computeomega(qv, tk, w, p):
     
     res = f_computeomega(qv.astype("float64").T,
@@ -157,12 +167,14 @@ def computeomega(qv, tk, w, p):
     #return res.T
     return res.astype("float32").T
 
+@handle_left_iter(3,3,0)
 def computetv(tk,qv):
     res = f_computetv(tk.astype("float64").T,
                       qv.astype("float64").T)
     
     return res.astype("float32").T
 
+@handle_left_iter(3,3,0)
 def computewetbulb(p,tk,qv):
     PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables()
     
@@ -176,6 +188,7 @@ def computewetbulb(p,tk,qv):
     
     return res.astype("float32").T
 
+@handle_left_iter(2,3,0, ignore_args=(4,))
 def computesrh(u, v, z, ter, top):
     
     res = f_computesrh(u.astype("float64").T, 
@@ -186,10 +199,11 @@ def computesrh(u, v, z, ter, top):
     
     return res.astype("float32").T
 
+@handle_left_iter(2,3,2, ignore_args=(5,6,7,8))
 def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top):
     
-    tem1 = n.zeros((u.shape[0],u.shape[1],u.shape[2]), "float64")
-    tem2 = n.zeros((u.shape[0],u.shape[1],u.shape[2]), "float64")
+    tem1 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "float64")
+    tem2 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "float64")
     
     res = f_computeuh(zp.astype("float64").T,
                       mapfct.astype("float64").T,
@@ -205,9 +219,10 @@ def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top):
     
     return res.astype("float32").T
 
+@handle_left_iter(2,3,0)
 def computepw(p,tv,qv,ht):
     # Note, dim 0 is height, we only want y and x
-    zdiff = n.zeros((p.shape[1], p.shape[2]), "float64")
+    zdiff = n.zeros((p.shape[-2], p.shape[-1]), "float64")
     res = f_computepw(p.astype("float64").T,
                       tv.astype("float64").T,
                       qv.astype("float64").T,
@@ -216,6 +231,7 @@ def computepw(p,tv,qv,ht):
     
     return res.astype("float32").T
 
+@handle_left_iter(3,3,0, ignore_args=(6,7,8))
 def computedbz(p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin):
     
     res = f_computedbz(p.astype("float64").T,
@@ -230,9 +246,10 @@ def computedbz(p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin):
     
     return res.astype("float32").T
 
+@handle_left_iter(3,3,0,ignore_args=(6,7,8))
 def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow):
-    flip_cape = n.zeros((p_hpa.shape[0],p_hpa.shape[1],p_hpa.shape[2]), "float64")
-    flip_cin = n.zeros((p_hpa.shape[0],p_hpa.shape[1],p_hpa.shape[2]), "float64")
+    flip_cape = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), "float64")
+    flip_cin = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), "float64")
     PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables()
     
     # The fortran routine needs pressure to be ascending in z-direction, 
@@ -264,7 +281,7 @@ def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow):
     # Remember to flip cape and cin back to descending p coordinates
     return (cape[::-1,:,:],cin[::-1,:,:])
 
-
+# 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):
@@ -276,6 +293,7 @@ def computeij(map_proj,truelat1,truelat2,stdlon,
     
     return res[0],res[1]
 
+# 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):
@@ -287,6 +305,7 @@ def computell(map_proj,truelat1,truelat2,stdlon,lat1,lon1,
     # Want lon,lat
     return res[1],res[0]
 
+@handle_left_iter(3,3,0, ignore_args=(3,))
 def computeeta(full_t, znu, psfc, ptop):
     pcalc = n.zeros(full_t.shape, "float64")
     mean_t = n.zeros(full_t.shape, "float64")
@@ -302,6 +321,7 @@ def computeeta(full_t, znu, psfc, ptop):
     
     return res.astype("float32").T
 
+@handle_left_iter(2,3,0,ignore_args=(7,))
 def computectt(p_hpa,tk,qice,qcld,qv,ght,ter,haveqci):
     res = f_computectt(p_hpa.astype("float64").T,
                     tk.astype("float64").T,

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

@@ -40,16 +40,6 @@ def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0):
             v = destagger(v_vars["V"], -2) 
             
     else:
-#         # 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: