diff --git a/wrf_open/var/src/python/wrf/var/decorators.py b/wrf_open/var/src/python/wrf/var/decorators.py index d0a1c80..3e46d3a 100644 --- a/wrf_open/var/src/python/wrf/var/decorators.py +++ b/wrf_open/var/src/python/wrf/var/decorators.py @@ -7,7 +7,8 @@ import numpy as n from wrf.var.units import do_conversion, check_units from wrf.var.destag import destagger -__all__ = ["convert_units", "handle_left_iter"] +__all__ = ["convert_units", "handle_left_iter", "uvmet_left_iter", + "handle_casting"] def convert_units(unit_type, alg_unit): """A decorator that applies unit conversion to a function's output array. @@ -248,5 +249,38 @@ 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. + + """ + def cast_decorator(func): + @wraps(func) + def func_wrapper(*args, **kargs): + orig_type = args[ref_idx].dtype + + new_args = [arg.astype(dtype) + if i in arg_idxs else arg + for i,arg in enumerate(args)] + + new_kargs = {key:(val.astype(dtype) + if key in karg_names else val) + for key,val in kargs.iteritems()} + + + res = func(*new_args, **new_kargs) + + if isinstance(res, n.ndarray): + if res.dtype == orig_type: + return res + return res.astype(orig_type) + else: # got back a sequence of arrays + return tuple(arr.astype(orig_type) + if arr.dtype != orig_type else arr + for arr in res) + + return func_wrapper + + return cast_decorator diff --git a/wrf_open/var/src/python/wrf/var/extension.py b/wrf_open/var/src/python/wrf/var/extension.py index 842cf68..99a93b7 100755 --- a/wrf_open/var/src/python/wrf/var/extension.py +++ b/wrf_open/var/src/python/wrf/var/extension.py @@ -1,5 +1,4 @@ import numpy as n -import numpy.ma as ma from wrf.var.constants import Constants from wrf.var.psadlookup import get_lookup_tables @@ -11,7 +10,8 @@ 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, uvmet_left_iter +from wrf.var.decorators import (handle_left_iter, uvmet_left_iter, + handle_casting) __all__ = ["FortranException", "computeslp", "computetk", "computetd", "computerh", "computeavo", "computepvo", "computeeth", @@ -23,108 +23,119 @@ class FortranException(Exception): def __call__(self, message): raise self.__class__(message) +@handle_casting(arg_idxs=(0,1)) def interpz3d(data3d,zdata,desiredloc,missingval): - res = f_interpz3d(data3d.astype("float64").T, - zdata.astype("float64").T, + res = f_interpz3d(data3d.T, + zdata.T, desiredloc, missingval) - return res.astype("float32").T + return res.T +@handle_casting(arg_idxs=(0,1)) def interpz2d(data3d,xy): - res = f_interp2dxy(data3d.astype("float64").T, - xy.astype("float64").T) - # Note: Fortran routine does not support missing values, so no masking - return res.astype("float32").T + res = f_interp2dxy(data3d.T, + xy.T) + return res.T +@handle_casting(arg_idxs=(0,1,2)) def interp1d(v_in,z_in,z_out,missingval): - res = f_interp1d(v_in.astype("float64"), - z_in.astype("float64"), - z_out.astype("float64"), + res = f_interp1d(v_in, + z_in, + z_out, missingval) - return res.astype("float32") + return res @handle_left_iter(2,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") 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, - p.astype("float64").T, - q.astype("float64").T, + res = f_computeslp(z.T, + t.T, + p.T, + q.T, t_sea_level.T, t_surf.T, level.T, FortranException()) - return res.astype("float32").T + return res.T @handle_left_iter(3,3,0) +@handle_casting(arg_idxs=(0,1)) def computetk(pres, theta): # No need to transpose here since operations on 1D array shape = pres.shape - res = f_computetk(pres.astype("float64").flatten("A"), - theta.astype("float64").flatten("A")) + res = f_computetk(pres.flatten("A"), + theta.flatten("A")) res = n.reshape(res, shape, "A") - return res.astype("float32") + return res -# Note: No decorator needed with 1D arrays +# Note: No left iteration decorator needed with 1D arrays +@handle_casting(arg_idxs=(0,1)) def computetd(pressure,qv_in): shape = pressure.shape - res = f_computetd(pressure.astype("float64").flatten("A"), qv_in.astype("float64").flatten("A")) + res = f_computetd(pressure.flatten("A"), + qv_in.flatten("A")) res = n.reshape(res, shape, "A") - return res.astype("float32") + return res # Note: No decorator needed with 1D arrays +@handle_casting(arg_idxs=(0,1,2)) def computerh(qv,q,t): shape = qv.shape - res = f_computerh(qv.astype("float64").flatten("A"), - q.astype("float64").flatten("A"), - t.astype("float64").flatten("A")) + res = f_computerh(qv.flatten("A"), + q.flatten("A"), + t.flatten("A")) res = n.reshape(res, shape, "A") - return res.astype("float32") + return res @handle_left_iter(3,3,0, ignore_args=(6,7), ref_stag_dim=-1) +@handle_casting(arg_idxs=(0,1,2,3,4,5)) def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy): - res = f_computeabsvort(u.astype("float64").T, - v.astype("float64").T, - msfu.astype("float64").T, - msfv.astype("float64").T, - msfm.astype("float64").T, - cor.astype("float64").T, + res = f_computeabsvort(u.T, + v.T, + msfu.T, + msfv.T, + msfm.T, + cor.T, dx, dy) - return res.astype("float32").T + return res.T @handle_left_iter(3,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): - res = f_computepvo(u.astype("float64").T, - v.astype("float64").T, - theta.astype("float64").T, - prs.astype("float64").T, - msfu.astype("float64").T, - msfv.astype("float64").T, - msfm.astype("float64").T, - cor.astype("float64").T, + res = f_computepvo(u.T, + v.T, + theta.T, + prs.T, + msfu.T, + msfv.T, + msfm.T, + cor.T, dx, dy) - return res.astype("float32").T + return res.T @handle_left_iter(3,3,0) +@handle_casting(arg_idxs=(0,1,2)) def computeeth(qv, tk, p): - res = f_computeeth(qv.astype("float64").T, - tk.astype("float64").T, - p.astype("float64").T) + res = f_computeeth(qv.T, + tk.T, + p.T) - return res.astype("float32").T + return res.T @uvmet_left_iter() +@handle_casting(arg_idxs=(0,1,2,3)) def computeuvmet(u,v,lat,lon,cen_long,cone): longca = n.zeros((lat.shape[-2], lat.shape[-1]), "float64") longcb = n.zeros((lon.shape[-2], lon.shape[-1]), "float64") @@ -138,12 +149,12 @@ def computeuvmet(u,v,lat,lon,cen_long,cone): # istag will always be false since winds are destaggered already # Missing values don't appear to be used, so setting to 0 - res = f_computeuvmet(u.astype("float64").T, - v.astype("float64").T, + res = f_computeuvmet(u.T, + v.T, longca.T, longcb.T, - lon.astype("float64").T, - lat.astype("float64").T, + lon.T, + lat.T, cen_long, cone, rpd, @@ -154,102 +165,112 @@ def computeuvmet(u,v,lat,lon,cen_long,cone): 0) - return n.squeeze(res.astype("float32").T) + return n.squeeze(res.T) @handle_left_iter(3,3,0) +@handle_casting(arg_idxs=(0,1,2,3)) def computeomega(qv, tk, w, p): - res = f_computeomega(qv.astype("float64").T, - tk.astype("float64").T, - w.astype("float64").T, - p.astype("float64").T) + res = f_computeomega(qv.T, + tk.T, + w.T, + p.T) #return res.T - return res.astype("float32").T + return res.T @handle_left_iter(3,3,0) +@handle_casting(arg_idxs=(0,1)) def computetv(tk,qv): - res = f_computetv(tk.astype("float64").T, - qv.astype("float64").T) + res = f_computetv(tk.T, + qv.T) - return res.astype("float32").T + return res.T @handle_left_iter(3,3,0) +@handle_casting(arg_idxs=(0,1,2)) def computewetbulb(p,tk,qv): PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables() - res = f_computewetbulb(p.astype("float64").T, - tk.astype("float64").T, - qv.astype("float64").T, + res = f_computewetbulb(p.T, + tk.T, + qv.T, PSADITHTE, PSADIPRS, PSADITMK.T, FortranException()) - return res.astype("float32").T + return res.T @handle_left_iter(2,3,0, ignore_args=(4,)) +@handle_casting(arg_idxs=(0,1,2,3)) def computesrh(u, v, z, ter, top): - res = f_computesrh(u.astype("float64").T, - v.astype("float64").T, - z.astype("float64").T, - ter.astype("float64").T, + res = f_computesrh(u.T, + v.T, + z.T, + ter.T, top) - return res.astype("float32").T + return res.T @handle_left_iter(2,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): 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, + res = f_computeuh(zp.T, + mapfct.T, dx, dy, bottom, top, - u.astype("float64").T, - v.astype("float64").T, - wstag.astype("float64").T, + u.T, + v.T, + wstag.T, tem1.T, tem2.T) - return res.astype("float32").T + return res.T @handle_left_iter(2,3,0) +@handle_casting(arg_idxs=(0,1,2,3)) def computepw(p,tv,qv,ht): - # Note, dim 0 is height, we only want y and x + # Note, dim -3 is height, we only want y and x 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, - ht.astype("float64").T, + res = f_computepw(p.T, + tv.T, + qv.T, + ht.T, zdiff.T) - return res.astype("float32").T + return res.T @handle_left_iter(3,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): - res = f_computedbz(p.astype("float64").T, - tk.astype("float64").T, - qv.astype("float64").T, - qr.astype("float64").T, - qs.astype("float64").T, - qg.astype("float64").T, + res = f_computedbz(p.T, + tk.T, + qv.T, + qr.T, + qs.T, + qg.T, sn0, ivarint, iliqskin) - return res.astype("float32").T + return res.T @handle_left_iter(3,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]), "float64") - flip_cin = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), "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, @@ -259,12 +280,12 @@ def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow): flip_qv = qv[::-1,:,:] flip_ht = ht[::-1,:,:] - f_computecape(flip_p.astype("float64").T, - flip_tk.astype("float64").T, - flip_qv.astype("float64").T, - flip_ht.astype("float64").T, - ter.astype("float64").T, - sfp.astype("float64").T, + f_computecape(flip_p.T, + flip_tk.T, + flip_qv.T, + flip_ht.T, + ter.T, + sfp.T, flip_cape.T, flip_cin.T, PSADITHTE, @@ -276,10 +297,11 @@ def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow): FortranException()) # Don't need to transpose since we only passed a view to fortran - cape = flip_cape.astype("float32")[::-1,:,:] - cin = flip_cin.astype("float32")[::-1,:,:] - # Remember to flip cape and cin back to descending p coordinates + cape = flip_cape[::-1,:,:] + cin = flip_cin[::-1,:,:] + + return (cape, cin) # TODO: This should handle lists of coords @@ -307,33 +329,35 @@ def computell(map_proj,truelat1,truelat2,stdlon,lat1,lon1, return res[1],res[0] @handle_left_iter(3,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") mean_t = n.zeros(full_t.shape, "float64") temp_t = n.zeros(full_t.shape, "float64") - res = f_converteta(full_t.astype("float64").T, - znu.astype("float64"), - psfc.astype("float64").T, + res = f_converteta(full_t.T, + znu.T, + psfc.T, ptop, pcalc.T, mean_t.T, temp_t.T) - return res.astype("float32").T + return res.T @handle_left_iter(2,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.astype("float64").T, - tk.astype("float64").T, - qice.astype("float64").T, - qcld.astype("float64").T, - qv.astype("float64").T, - ght.astype("float64").T, - ter.astype("float64").T, + res = f_computectt(p_hpa.T, + tk.T, + qice.T, + qcld.T, + qv.T, + ght.T, + ter.T, haveqci) - return res.astype("float32").T + return res.T diff --git a/wrf_open/var/src/python/wrf/var/wrfext.f90 b/wrf_open/var/src/python/wrf/var/wrfext.f90 index 098db5a..7d5580c 100755 --- a/wrf_open/var/src/python/wrf/var/wrfext.f90 +++ b/wrf_open/var/src/python/wrf/var/wrfext.f90 @@ -710,7 +710,7 @@ SUBROUTINE f_computewetbulb(prs,tmk,qvp,PSADITHTE,PSADIPRS,PSADITMK,throw_except jt=-1 DO jtch=1,150-1 - IF (eth.GE.psadithte(jtch).AND.eth.LT.psadithte(jtch+1)) then + IF (eth.GE.psadithte(jtch).AND.eth.LT.psadithte(jtch+1)) THEN jt=jtch EXIT ENDIF @@ -719,7 +719,7 @@ SUBROUTINE f_computewetbulb(prs,tmk,qvp,PSADITHTE,PSADIPRS,PSADITMK,throw_except ! 213 CONTINUE ip=-1 DO ipch=1,150-1 - IF (p.LE.psadiprs(ipch).AND.p.GT.psadiprs(ipch+1)) then + IF (p.LE.psadiprs(ipch).AND.p.GT.psadiprs(ipch+1)) THEN ip=ipch EXIT ENDIF @@ -734,7 +734,7 @@ SUBROUTINE f_computewetbulb(prs,tmk,qvp,PSADITHTE,PSADIPRS,PSADITMK,throw_except fracjt2=1.-fracjt fracip=(psadiprs(ip)-p)/(psadiprs(ip)-psadiprs(ip+1)) fracip2=1.-fracip - IF (psaditmk(ip,jt).GT.1e9.OR.psaditmk(ip+1,jt).GT.1e9.OR.psaditmk(ip,jt+1).GT.1e9.OR.psaditmk(ip+1,jt+1).GT.1e9) then + IF (psaditmk(ip,jt).GT.1e9.OR.psaditmk(ip+1,jt).GT.1e9.OR.psaditmk(ip,jt+1).GT.1e9.OR.psaditmk(ip+1,jt+1).GT.1e9) THEN !PRINT*,'Tried to access missing tmperature in lookup table.' CALL throw_exception('Prs and Thte probably unreasonable. prs,thte=',p,eth) !STOP ! TODO: Need to make python throw an exception here