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
Split 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 @@ -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(): @@ -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

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

@ -1,5 +1,4 @@ @@ -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, @@ -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): @@ -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): @@ -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): @@ -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): @@ -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): @@ -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, @@ -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

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 @@ -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 @@ -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 @@ -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

Write Preview
Loading…
Cancel
Save