diff --git a/wrf_open/var/src/python/wrf/var/cape.py b/wrf_open/var/src/python/wrf/var/cape.py index 4f14c82..e21ed9f 100755 --- a/wrf_open/var/src/python/wrf/var/cape.py +++ b/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) \ No newline at end of file diff --git a/wrf_open/var/src/python/wrf/var/decorators.py b/wrf_open/var/src/python/wrf/var/decorators.py index c0ee771..334bac8 100644 --- a/wrf_open/var/src/python/wrf/var/decorators.py +++ b/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 + diff --git a/wrf_open/var/src/python/wrf/var/extension.py b/wrf_open/var/src/python/wrf/var/extension.py index 461b425..16c1768 100755 --- a/wrf_open/var/src/python/wrf/var/extension.py +++ b/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, diff --git a/wrf_open/var/src/python/wrf/var/uvmet.py b/wrf_open/var/src/python/wrf/var/uvmet.py index fb037bc..0f6e66b 100755 --- a/wrf_open/var/src/python/wrf/var/uvmet.py +++ b/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: