nilyin
/
wrf-python
forked from 3rdparty/wrf-python
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
Browse Source

Added type casting decorator

main
Bill Ladwig 10 years ago
parent
8b2843263a
commit
5135d89896
3 changed files with 174 additions and 116 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 36
      wrf_open/var/src/python/wrf/var/decorators.py
  2. 248
      wrf_open/var/src/python/wrf/var/extension.py
  3. 6
      wrf_open/var/src/python/wrf/var/wrfext.f90

36
wrf_open/var/src/python/wrf/var/decorators.py
Unescape View File

@ -7,7 +7,8 @@ import numpy as n
from wrf.var.units import do_conversion, check_units from wrf.var.units import do_conversion, check_units
from wrf.var.destag import destagger 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): def convert_units(unit_type, alg_unit):
"""A decorator that applies unit conversion to a function's output array. """A decorator that applies unit conversion to a function's output array.
@ -248,5 +249,38 @@ def uvmet_left_iter():
return indexing_decorator 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

248
wrf_open/var/src/python/wrf/var/extension.py
Unescape View File

@ -1,5 +1,4 @@
import numpy as n import numpy as n
import numpy.ma as ma
from wrf.var.constants import Constants from wrf.var.constants import Constants
from wrf.var.psadlookup import get_lookup_tables 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_computesrh, f_computeuh, f_computepw, f_computedbz,
f_lltoij, f_ijtoll, f_converteta, f_computectt) f_lltoij, f_ijtoll, f_converteta, f_computectt)
from wrf.var._wrfcape import f_computecape 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", __all__ = ["FortranException", "computeslp", "computetk", "computetd",
"computerh", "computeavo", "computepvo", "computeeth", "computerh", "computeavo", "computepvo", "computeeth",
@ -23,108 +23,119 @@ class FortranException(Exception):
def __call__(self, message): def __call__(self, message):
raise self.__class__(message) raise self.__class__(message)
@handle_casting(arg_idxs=(0,1))
def interpz3d(data3d,zdata,desiredloc,missingval): def interpz3d(data3d,zdata,desiredloc,missingval):
res = f_interpz3d(data3d.astype("float64").T, res = f_interpz3d(data3d.T,
zdata.astype("float64").T, zdata.T,
desiredloc, desiredloc,
missingval) missingval)
return res.astype("float32").T return res.T
@handle_casting(arg_idxs=(0,1))
def interpz2d(data3d,xy): def interpz2d(data3d,xy):
res = f_interp2dxy(data3d.astype("float64").T, res = f_interp2dxy(data3d.T,
xy.astype("float64").T) xy.T)
# Note: Fortran routine does not support missing values, so no masking return res.T
return res.astype("float32").T
@handle_casting(arg_idxs=(0,1,2))
def interp1d(v_in,z_in,z_out,missingval): def interp1d(v_in,z_in,z_out,missingval):
res = f_interp1d(v_in.astype("float64"), res = f_interp1d(v_in,
z_in.astype("float64"), z_in,
z_out.astype("float64"), z_out,
missingval) missingval)
return res.astype("float32") return res
@handle_left_iter(2,3,0) @handle_left_iter(2,3,0)
@handle_casting(arg_idxs=(0,1,2,3))
def computeslp(z,t,p,q): def computeslp(z,t,p,q):
t_surf = n.zeros((z.shape[-2], z.shape[-1]), "float64") t_surf = n.zeros((z.shape[-2], z.shape[-1]), "float64")
t_sea_level = 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") level = n.zeros((z.shape[-2], z.shape[-1]), "int32")
res = f_computeslp(z.astype("float64").T, res = f_computeslp(z.T,
t.astype("float64").T, t.T,
p.astype("float64").T, p.T,
q.astype("float64").T, q.T,
t_sea_level.T, t_sea_level.T,
t_surf.T, t_surf.T,
level.T, level.T,
FortranException()) FortranException())
return res.astype("float32").T return res.T
@handle_left_iter(3,3,0) @handle_left_iter(3,3,0)
@handle_casting(arg_idxs=(0,1))
def computetk(pres, theta): def computetk(pres, theta):
# No need to transpose here since operations on 1D array # No need to transpose here since operations on 1D array
shape = pres.shape shape = pres.shape
res = f_computetk(pres.astype("float64").flatten("A"), res = f_computetk(pres.flatten("A"),
theta.astype("float64").flatten("A")) theta.flatten("A"))
res = n.reshape(res, shape, "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): def computetd(pressure,qv_in):
shape = pressure.shape 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") res = n.reshape(res, shape, "A")
return res.astype("float32") return res
# Note: No decorator needed with 1D arrays # Note: No decorator needed with 1D arrays
@handle_casting(arg_idxs=(0,1,2))
def computerh(qv,q,t): def computerh(qv,q,t):
shape = qv.shape shape = qv.shape
res = f_computerh(qv.astype("float64").flatten("A"), res = f_computerh(qv.flatten("A"),
q.astype("float64").flatten("A"), q.flatten("A"),
t.astype("float64").flatten("A")) t.flatten("A"))
res = n.reshape(res, shape, "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_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): def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy):
res = f_computeabsvort(u.astype("float64").T, res = f_computeabsvort(u.T,
v.astype("float64").T, v.T,
msfu.astype("float64").T, msfu.T,
msfv.astype("float64").T, msfv.T,
msfm.astype("float64").T, msfm.T,
cor.astype("float64").T, cor.T,
dx, dx,
dy) dy)
return res.astype("float32").T return res.T
@handle_left_iter(3,3,2, ignore_args=(8,9)) @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): def computepvo(u,v,theta,prs,msfu,msfv,msfm,cor,dx,dy):
res = f_computepvo(u.astype("float64").T, res = f_computepvo(u.T,
v.astype("float64").T, v.T,
theta.astype("float64").T, theta.T,
prs.astype("float64").T, prs.T,
msfu.astype("float64").T, msfu.T,
msfv.astype("float64").T, msfv.T,
msfm.astype("float64").T, msfm.T,
cor.astype("float64").T, cor.T,
dx, dx,
dy) dy)
return res.astype("float32").T return res.T
@handle_left_iter(3,3,0) @handle_left_iter(3,3,0)
@handle_casting(arg_idxs=(0,1,2))
def computeeth(qv, tk, p): def computeeth(qv, tk, p):
res = f_computeeth(qv.astype("float64").T, res = f_computeeth(qv.T,
tk.astype("float64").T, tk.T,
p.astype("float64").T) p.T)
return res.astype("float32").T return res.T
@uvmet_left_iter() @uvmet_left_iter()
@handle_casting(arg_idxs=(0,1,2,3))
def computeuvmet(u,v,lat,lon,cen_long,cone): def computeuvmet(u,v,lat,lon,cen_long,cone):
longca = n.zeros((lat.shape[-2], lat.shape[-1]), "float64") longca = n.zeros((lat.shape[-2], lat.shape[-1]), "float64")
longcb = n.zeros((lon.shape[-2], lon.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 # istag will always be false since winds are destaggered already
# Missing values don't appear to be used, so setting to 0 # Missing values don't appear to be used, so setting to 0
res = f_computeuvmet(u.astype("float64").T, res = f_computeuvmet(u.T,
v.astype("float64").T, v.T,
longca.T, longca.T,
longcb.T, longcb.T,
lon.astype("float64").T, lon.T,
lat.astype("float64").T, lat.T,
cen_long, cen_long,
cone, cone,
rpd, rpd,
@ -154,102 +165,112 @@ def computeuvmet(u,v,lat,lon,cen_long,cone):
0) 0)
return n.squeeze(res.astype("float32").T) return n.squeeze(res.T)
@handle_left_iter(3,3,0) @handle_left_iter(3,3,0)
@handle_casting(arg_idxs=(0,1,2,3))
def computeomega(qv, tk, w, p): def computeomega(qv, tk, w, p):
res = f_computeomega(qv.astype("float64").T, res = f_computeomega(qv.T,
tk.astype("float64").T, tk.T,
w.astype("float64").T, w.T,
p.astype("float64").T) p.T)
#return res.T #return res.T
return res.astype("float32").T return res.T
@handle_left_iter(3,3,0) @handle_left_iter(3,3,0)
@handle_casting(arg_idxs=(0,1))
def computetv(tk,qv): def computetv(tk,qv):
res = f_computetv(tk.astype("float64").T, res = f_computetv(tk.T,
qv.astype("float64").T) qv.T)
return res.astype("float32").T return res.T
@handle_left_iter(3,3,0) @handle_left_iter(3,3,0)
@handle_casting(arg_idxs=(0,1,2))
def computewetbulb(p,tk,qv): def computewetbulb(p,tk,qv):
PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables() PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables()
res = f_computewetbulb(p.astype("float64").T, res = f_computewetbulb(p.T,
tk.astype("float64").T, tk.T,
qv.astype("float64").T, qv.T,
PSADITHTE, PSADITHTE,
PSADIPRS, PSADIPRS,
PSADITMK.T, PSADITMK.T,
FortranException()) FortranException())
return res.astype("float32").T return res.T
@handle_left_iter(2,3,0, ignore_args=(4,)) @handle_left_iter(2,3,0, ignore_args=(4,))
@handle_casting(arg_idxs=(0,1,2,3))
def computesrh(u, v, z, ter, top): def computesrh(u, v, z, ter, top):
res = f_computesrh(u.astype("float64").T, res = f_computesrh(u.T,
v.astype("float64").T, v.T,
z.astype("float64").T, z.T,
ter.astype("float64").T, ter.T,
top) top)
return res.astype("float32").T return res.T
@handle_left_iter(2,3,2, ignore_args=(5,6,7,8)) @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): def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top):
tem1 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "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") tem2 = n.zeros((u.shape[-3],u.shape[-2],u.shape[-1]), "float64")
res = f_computeuh(zp.astype("float64").T, res = f_computeuh(zp.T,
mapfct.astype("float64").T, mapfct.T,
dx, dx,
dy, dy,
bottom, bottom,
top, top,
u.astype("float64").T, u.T,
v.astype("float64").T, v.T,
wstag.astype("float64").T, wstag.T,
tem1.T, tem1.T,
tem2.T) tem2.T)
return res.astype("float32").T return res.T
@handle_left_iter(2,3,0) @handle_left_iter(2,3,0)
@handle_casting(arg_idxs=(0,1,2,3))
def computepw(p,tv,qv,ht): 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") zdiff = n.zeros((p.shape[-2], p.shape[-1]), "float64")
res = f_computepw(p.astype("float64").T, res = f_computepw(p.T,
tv.astype("float64").T, tv.T,
qv.astype("float64").T, qv.T,
ht.astype("float64").T, ht.T,
zdiff.T) zdiff.T)
return res.astype("float32").T return res.T
@handle_left_iter(3,3,0, ignore_args=(6,7,8)) @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): def computedbz(p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin):
res = f_computedbz(p.astype("float64").T, res = f_computedbz(p.T,
tk.astype("float64").T, tk.T,
qv.astype("float64").T, qv.T,
qr.astype("float64").T, qr.T,
qs.astype("float64").T, qs.T,
qg.astype("float64").T, qg.T,
sn0, sn0,
ivarint, ivarint,
iliqskin) iliqskin)
return res.astype("float32").T return res.T
@handle_left_iter(3,3,0,ignore_args=(6,7,8)) @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): 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_cape = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]),
flip_cin = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]), "float64") "float64")
flip_cin = n.zeros((p_hpa.shape[-3],p_hpa.shape[-2],p_hpa.shape[-1]),
"float64")
PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables() PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables()
# The fortran routine needs pressure to be ascending in z-direction, # 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_qv = qv[::-1,:,:]
flip_ht = ht[::-1,:,:] flip_ht = ht[::-1,:,:]
f_computecape(flip_p.astype("float64").T, f_computecape(flip_p.T,
flip_tk.astype("float64").T, flip_tk.T,
flip_qv.astype("float64").T, flip_qv.T,
flip_ht.astype("float64").T, flip_ht.T,
ter.astype("float64").T, ter.T,
sfp.astype("float64").T, sfp.T,
flip_cape.T, flip_cape.T,
flip_cin.T, flip_cin.T,
PSADITHTE, PSADITHTE,
@ -276,10 +297,11 @@ def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow):
FortranException()) FortranException())
# Don't need to transpose since we only passed a view to fortran # 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 # Remember to flip cape and cin back to descending p coordinates
cape = flip_cape[::-1,:,:]
cin = flip_cin[::-1,:,:]
return (cape, cin) return (cape, cin)
# TODO: This should handle lists of coords # 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] return res[1],res[0]
@handle_left_iter(3,3,0, ignore_args=(3,)) @handle_left_iter(3,3,0, ignore_args=(3,))
@handle_casting(arg_idxs=(0,1,2))
def computeeta(full_t, znu, psfc, ptop): def computeeta(full_t, znu, psfc, ptop):
pcalc = n.zeros(full_t.shape, "float64") pcalc = n.zeros(full_t.shape, "float64")
mean_t = n.zeros(full_t.shape, "float64") mean_t = n.zeros(full_t.shape, "float64")
temp_t = n.zeros(full_t.shape, "float64") temp_t = n.zeros(full_t.shape, "float64")
res = f_converteta(full_t.astype("float64").T, res = f_converteta(full_t.T,
znu.astype("float64"), znu.T,
psfc.astype("float64").T, psfc.T,
ptop, ptop,
pcalc.T, pcalc.T,
mean_t.T, mean_t.T,
temp_t.T) temp_t.T)
return res.astype("float32").T return res.T
@handle_left_iter(2,3,0,ignore_args=(7,)) @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): def computectt(p_hpa,tk,qice,qcld,qv,ght,ter,haveqci):
res = f_computectt(p_hpa.astype("float64").T, res = f_computectt(p_hpa.T,
tk.astype("float64").T, tk.T,
qice.astype("float64").T, qice.T,
qcld.astype("float64").T, qcld.T,
qv.astype("float64").T, qv.T,
ght.astype("float64").T, ght.T,
ter.astype("float64").T, ter.T,
haveqci) haveqci)
return res.astype("float32").T return res.T

6
wrf_open/var/src/python/wrf/var/wrfext.f90
Unescape View File

@ -710,7 +710,7 @@ SUBROUTINE f_computewetbulb(prs,tmk,qvp,PSADITHTE,PSADIPRS,PSADITMK,throw_except
jt=-1 jt=-1
DO jtch=1,150-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 jt=jtch
EXIT EXIT
ENDIF ENDIF
@ -719,7 +719,7 @@ SUBROUTINE f_computewetbulb(prs,tmk,qvp,PSADITHTE,PSADIPRS,PSADITMK,throw_except
! 213 CONTINUE ! 213 CONTINUE
ip=-1 ip=-1
DO ipch=1,150-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 ip=ipch
EXIT EXIT
ENDIF ENDIF
@ -734,7 +734,7 @@ SUBROUTINE f_computewetbulb(prs,tmk,qvp,PSADITHTE,PSADIPRS,PSADITMK,throw_except
fracjt2=1.-fracjt fracjt2=1.-fracjt
fracip=(psadiprs(ip)-p)/(psadiprs(ip)-psadiprs(ip+1)) fracip=(psadiprs(ip)-p)/(psadiprs(ip)-psadiprs(ip+1))
fracip2=1.-fracip 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.' !PRINT*,'Tried to access missing tmperature in lookup table.'
CALL throw_exception('Prs and Thte probably unreasonable. prs,thte=',p,eth) CALL throw_exception('Prs and Thte probably unreasonable. prs,thte=',p,eth)
!STOP ! TODO: Need to make python throw an exception here !STOP ! TODO: Need to make python throw an exception here

Write Preview
Loading…
Cancel
Save