iteration decorator now uses first output to determine output dimensions. Added 2D interpolation to a line

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

@@ -59,15 +59,18 @@ def _left_indexes(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={}, ref_stag_dim=None):
+def _calc_out_dims(outvar, left_dims):
+    left_dims = [x for x in left_dims]
+    right_dims = [x for x in outvar.shape]
+    return left_dims + right_dims 
+
+def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1,
+                   ref_var_name=None,
+                   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 
@@ -95,33 +98,14 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1,
                 ref_var = kargs[ref_var_name]
             
             ref_var_shape = ref_var.shape
-            extra_dim_num = ref_var.ndim- ref_var_expected_dims
+            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)
             
             # 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]
-                neg_start_idx = -(alg_out_num_dims - len(alg_out_fixed_dims))
-                right_dims += [ref_var_shape[x] 
-                        for x in xrange(neg_start_idx,0,1)]
-                outdims += right_dims
-            
-            if ref_stag_dim is not None:
-                outdims[ref_stag_dim] -= 1  
+            extra_dims = [ref_var_shape[x] for x in xrange(extra_dim_num)]            
             
             out_inited = False
             for left_idxs in _left_indexes(extra_dims):
@@ -147,6 +131,7 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1,
                 if isinstance(res, n.ndarray):
                     # Output array
                     if not out_inited:
+                        outdims = _calc_out_dims(res, extra_dims)
                         output = n.zeros(outdims, ref_var.dtype)
                         out_inited = True
                     
@@ -154,6 +139,7 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1,
                     
                 else:   # This should be a list or a tuple (cape)
                     if not out_inited:
+                        outdims = _calc_out_dims(res[0], extra_dims)
                         output = [n.zeros(outdims, ref_var.dtype) 
                                   for i in xrange(len(res))]
                         out_inited = True
@@ -250,8 +236,8 @@ def uvmet_left_iter():
     return indexing_decorator
 
 def handle_casting(ref_idx=0, arg_idxs=(), karg_names=None,dtype=n.float64):
-    """Decorator to handle iterating over leftmost dimensions when using 
-    multiple files and/or multiple times with the uvmet product.
+    """Decorator to handle casting to/from required dtype used in 
+    numerical routine.
     
     """
     def cast_decorator(func):

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

@@ -23,7 +23,7 @@ class FortranException(Exception):
     def __call__(self, message):
         raise self.__class__(message)
 
-@handle_left_iter(2,3,0)
+@handle_left_iter(3,0, ignore_args=(2,3))
 @handle_casting(arg_idxs=(0,1))
 def interpz3d(data3d,zdata,desiredloc,missingval):
     res = f_interpz3d(data3d.T, 
@@ -32,9 +32,9 @@ def interpz3d(data3d,zdata,desiredloc,missingval):
                       missingval)
     return res.T
 
+@handle_left_iter(2,0)
 @handle_casting(arg_idxs=(0,1))
-@handle_left_iter(2,2,0)
-def interpz2d(data3d,xy):
+def interp2dxy(data3d,xy):
     res = f_interp2dxy(data3d.T,
                        xy.T)
     return res.T
@@ -48,7 +48,7 @@ def interp1d(v_in,z_in,z_out,missingval):
     
     return res
 
-@handle_left_iter(2,3,0)
+@handle_left_iter(3,0)
 @handle_casting(arg_idxs=(0,1,2,3))
 def computeslp(z,t,p,q):
     t_surf = n.zeros((z.shape[-2], z.shape[-1]), "float64")
@@ -66,7 +66,7 @@ def computeslp(z,t,p,q):
     
     return res.T
 
-@handle_left_iter(3,3,0)
+@handle_left_iter(3,0)
 @handle_casting(arg_idxs=(0,1))
 def computetk(pres, theta):
     # No need to transpose here since operations on 1D array
@@ -95,7 +95,7 @@ def computerh(qv,q,t):
     res = n.reshape(res, shape, "A")
     return res
 
-@handle_left_iter(3,3,0, ignore_args=(6,7), ref_stag_dim=-1)
+@handle_left_iter(3,0, ignore_args=(6,7))
 @handle_casting(arg_idxs=(0,1,2,3,4,5))
 def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy):
     res = f_computeabsvort(u.T,
@@ -109,7 +109,7 @@ def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy):
     
     return res.T
 
-@handle_left_iter(3,3,2, ignore_args=(8,9))
+@handle_left_iter(3,2, ignore_args=(8,9))
 @handle_casting(arg_idxs=(0,1,2,3,4,5,6,7))
 def computepvo(u,v,theta,prs,msfu,msfv,msfm,cor,dx,dy):
     
@@ -126,7 +126,7 @@ def computepvo(u,v,theta,prs,msfu,msfv,msfm,cor,dx,dy):
     
     return res.T
 
-@handle_left_iter(3,3,0)
+@handle_left_iter(3,0)
 @handle_casting(arg_idxs=(0,1,2))
 def computeeth(qv, tk, p):
     
@@ -169,7 +169,7 @@ def computeuvmet(u,v,lat,lon,cen_long,cone):
     
     return n.squeeze(res.T)
 
-@handle_left_iter(3,3,0)
+@handle_left_iter(3,0)
 @handle_casting(arg_idxs=(0,1,2,3))
 def computeomega(qv, tk, w, p):
     
@@ -181,7 +181,7 @@ def computeomega(qv, tk, w, p):
     #return res.T
     return res.T
 
-@handle_left_iter(3,3,0)
+@handle_left_iter(3,0)
 @handle_casting(arg_idxs=(0,1))
 def computetv(tk,qv):
     res = f_computetv(tk.T,
@@ -189,7 +189,7 @@ def computetv(tk,qv):
     
     return res.T
 
-@handle_left_iter(3,3,0)
+@handle_left_iter(3,0)
 @handle_casting(arg_idxs=(0,1,2))
 def computewetbulb(p,tk,qv):
     PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables()
@@ -204,7 +204,7 @@ def computewetbulb(p,tk,qv):
     
     return res.T
 
-@handle_left_iter(2,3,0, ignore_args=(4,))
+@handle_left_iter(3,0, ignore_args=(4,))
 @handle_casting(arg_idxs=(0,1,2,3))
 def computesrh(u, v, z, ter, top):
 
@@ -216,7 +216,7 @@ def computesrh(u, v, z, ter, top):
     
     return res.T
 
-@handle_left_iter(2,3,2, ignore_args=(5,6,7,8))
+@handle_left_iter(3,2, ignore_args=(5,6,7,8))
 @handle_casting(arg_idxs=(0,1,2,3,4))
 def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top):
     
@@ -237,7 +237,7 @@ def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top):
     
     return res.T
 
-@handle_left_iter(2,3,0)
+@handle_left_iter(3,0)
 @handle_casting(arg_idxs=(0,1,2,3))
 def computepw(p,tv,qv,ht):
     # Note, dim -3 is height, we only want y and x
@@ -250,7 +250,7 @@ def computepw(p,tv,qv,ht):
     
     return res.T
 
-@handle_left_iter(3,3,0, ignore_args=(6,7,8))
+@handle_left_iter(3,0, ignore_args=(6,7,8))
 @handle_casting(arg_idxs=(0,1,2,3,4,5))
 def computedbz(p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin):
     
@@ -266,7 +266,7 @@ def computedbz(p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin):
     
     return res.T
 
-@handle_left_iter(3,3,0,ignore_args=(6,7,8))
+@handle_left_iter(3,0,ignore_args=(6,7,8))
 @handle_casting(arg_idxs=(0,1,2,3,4,5))
 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]), 
@@ -330,7 +330,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,))
+@handle_left_iter(3,0, ignore_args=(3,))
 @handle_casting(arg_idxs=(0,1,2))
 def computeeta(full_t, znu, psfc, ptop):
     pcalc = n.zeros(full_t.shape, "float64")
@@ -347,7 +347,7 @@ def computeeta(full_t, znu, psfc, ptop):
     
     return res.T
 
-@handle_left_iter(2,3,0,ignore_args=(7,))
+@handle_left_iter(3,0,ignore_args=(7,))
 @handle_casting(arg_idxs=(0,1,2,3,4,5,6))
 def computectt(p_hpa,tk,qice,qcld,qv,ght,ter,haveqci):
     res = f_computectt(p_hpa.T,

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

@@ -3,11 +3,13 @@ from math import floor, ceil
 import numpy as n
 import numpy.ma as ma
 
-from wrf.var.extension import interpz3d,interpz2d,interp1d
+from wrf.var.extension import interpz3d,interp2dxy,interp1d
+from wrf.var.decorators import handle_left_iter, handle_casting
 
-__all__ = ["get_interplevel", "get_vertcross"]
+__all__ = ["interplevel", "vertcross", "interpline"]
 
-def get_interplevel(data3d,zdata,desiredloc,missingval=-99999):
+#  Note:  Extension decorator is good enough to handle left dims
+def interplevel(data3d,zdata,desiredloc,missingval=-99999):
     """Return the horizontally interpolated data at the provided level
     
     data3d - the 3D field to interpolate
@@ -27,17 +29,17 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None,
     
     xdim - maximum x-dimension
     ydim - maximum y-dimension
-    pivot_point - a pivot point of (x,y) (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 (x,y)
-    end_point - an end point tuple of (x,y)
+    start_point - a start_point tuple of (south_north1, west_east1)
+    end_point - an end point tuple of (south_north2, west_east2)
     
     """ 
     
     # Have a pivot point with an angle to find cross section
     if pivot_point is not None and angle is not None:
-        xp = pivot_point[0]
-        yp = pivot_point[1]
+        xp = pivot_point[-1]
+        yp = pivot_point[-2]
         
         if (angle > 315.0 or angle < 45.0 
             or ((angle > 135.0) and (angle < 225.0))):
@@ -101,10 +103,10 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None,
                 y1 = ydim-1
                 x1 =  (y1 - intercept)/slope
     elif start_point is not None and end_point is not None:
-        x0 = start_point[0]
-        y0 = start_point[1]
-        x1 = end_point[0]
-        y1 = end_point[1]
+        x0 = start_point[-1]
+        y0 = start_point[-2]
+        x1 = end_point[-1]
+        y1 = end_point[-2]
         if ( x1 > xdim-1 ): 
             x1 = xdim
         if ( y1 > ydim-1): 
@@ -129,10 +131,25 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None,
         
     return xy
 
-
-# TODO:  Add flag to use lat/lon points by doing conversion
-def get_vertcross(data3d, z, missingval=-99999, 
-                  pivot_point=None,angle=None,start_point=None,end_point=None):
+# TODO:  Need a decorator that can handle right dimensions based on what
+# is returned.  In this case, the result of the xy calculation determines
+# output dimensions on right.
+@handle_casting(arg_idxs=(0,1))
+def vertcross(data3d, z, missingval=-99999, 
+              pivot_point=None,angle=None,
+              start_point=None,end_point=None):
+    """Return the vertical cross section for a 3D field, interpolated 
+    to a verical plane defined by a horizontal line.
+    
+    Arguments:
+        data3d - a 3D data field
+        z - 3D height field
+        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)
+        
+    """
     
     xdim = z.shape[-1]
     ydim = z.shape[-2]
@@ -140,7 +157,7 @@ def get_vertcross(data3d, z, missingval=-99999,
     xy = _get_xy(xdim, ydim, pivot_point, angle, start_point, end_point)
     
     # Interp z
-    var2dz   = interpz2d(z, xy)
+    var2dz   = interp2dxy(z, xy)
     
     #  interp to constant z grid
     if(var2dz[0,0] > var2dz[1,0]):  # monotonically decreasing coordinate
@@ -165,13 +182,46 @@ def get_vertcross(data3d, z, missingval=-99999,
     #interp the variable
     
     var2d = n.zeros((nlevels, xy.shape[0]),dtype=var2dz.dtype)
-    var2dtmp = interpz2d(data3d, xy)
+    var2dtmp = interp2dxy(data3d, xy)
     
     for i in xrange(xy.shape[0]):
         var2d[:,i] = interp1d(var2dtmp[:,i], var2dz[:,i], z_var2d, missingval)
         
     return ma.masked_values(var2d, missingval)
+
+# TODO:  Need a decorator that can handle right dimensions based on what
+# is returned.  In this case, the result of the xy calculation determines
+# output dimensions on right.
+@handle_casting(arg_idxs=(0,))
+def interpline(data2d, pivot_point=None, 
+                 angle=None, start_point=None,
+                 end_point=None):
+    """Return the 2D field interpolated along a line.
+    
+    Arguments:
+        var2d - a 2D data field
+        pivot_point - a pivot point of (south_north,west_east)
+        angle - the angle through the pivot point in degrees
+        start_point - a start_point tuple of (south_north1,west_east1)
+        end_point - an end point tuple of (south_north2,west_east2)
         
+    """
+    
+    tmp_shape = [0] + [x for x in data2d.shape]
+    
+    var2dtmp = n.zeros(tmp_shape, data2d.dtype)
+    
+    xdim = data2d.shape[-1]
+    ydim = data2d.shape[-2]
+    
+    xy = _get_xy(xdim, ydim, pivot_point, angle, start_point, end_point)
+    
+    var2dtmp[0,:,:] = data2d[:,:]
+    
+    var1dtmp = interp2dxy(var2dtmp, xy)
+    
+    return var1dtmp[0,:]
+