From f45ac9756e790571cf5c051116555da2bab677b3 Mon Sep 17 00:00:00 2001 From: Bill Ladwig Date: Sat, 19 Dec 2015 16:42:35 -0700 Subject: [PATCH] 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 | 42 +++------ wrf_open/var/src/python/wrf/var/extension.py | 36 ++++---- wrf_open/var/src/python/wrf/var/interp.py | 86 +++++++++++++++---- 3 files changed, 100 insertions(+), 64 deletions(-) diff --git a/wrf_open/var/src/python/wrf/var/decorators.py b/wrf_open/var/src/python/wrf/var/decorators.py index 3e46d3a..124f538 100644 --- a/wrf_open/var/src/python/wrf/var/decorators.py +++ b/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): diff --git a/wrf_open/var/src/python/wrf/var/extension.py b/wrf_open/var/src/python/wrf/var/extension.py index a22fc4a..d11cda1 100755 --- a/wrf_open/var/src/python/wrf/var/extension.py +++ b/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, diff --git a/wrf_open/var/src/python/wrf/var/interp.py b/wrf_open/var/src/python/wrf/var/interp.py index cd508e4..a963000 100755 --- a/wrf_open/var/src/python/wrf/var/interp.py +++ b/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,:] +