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

iteration decorator now uses first output to determine output dimensions. Added 2D interpolation to a line

main
Bill Ladwig 10 years ago
parent
5b89a89c2e
commit
f45ac9756e
3 changed files with 100 additions and 64 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. 42
      wrf_open/var/src/python/wrf/var/decorators.py
  2. 36
      wrf_open/var/src/python/wrf/var/extension.py
  3. 86
      wrf_open/var/src/python/wrf/var/interp.py

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

@ -59,15 +59,18 @@ def _left_indexes(dims):
for idxs in product(*arg): for idxs in product(*arg):
yield idxs yield idxs
def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1, def _calc_out_dims(outvar, left_dims):
ref_var_name=None, alg_out_fixed_dims=(), left_dims = [x for x in left_dims]
ignore_args=(), ignore_kargs={}, ref_stag_dim=None): 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 """Decorator to handle iterating over leftmost dimensions when using
multiple files and/or multiple times. multiple files and/or multiple times.
Arguments: Arguments:
- alg_out_num_dims - the return dimension count from the numerical
algorithm
- ref_var_expected_dims - the number of dimensions that the Fortran - ref_var_expected_dims - the number of dimensions that the Fortran
algorithm is expecting for the reference variable algorithm is expecting for the reference variable
- ref_var_idx - the index in args used as the reference variable for - 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 = kargs[ref_var_name]
ref_var_shape = ref_var.shape 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 # No special left side iteration, return the function result
if (extra_dim_num == 0): if (extra_dim_num == 0):
return func(*args, **kargs) return func(*args, **kargs)
# Start by getting the left-most 'extra' dims # Start by getting the left-most 'extra' dims
outdims = [ref_var_shape[x] for x in xrange(extra_dim_num)] extra_dims = [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
out_inited = False out_inited = False
for left_idxs in _left_indexes(extra_dims): 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): if isinstance(res, n.ndarray):
# Output array # Output array
if not out_inited: if not out_inited:
outdims = _calc_out_dims(res, extra_dims)
output = n.zeros(outdims, ref_var.dtype) output = n.zeros(outdims, ref_var.dtype)
out_inited = True 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) else: # This should be a list or a tuple (cape)
if not out_inited: if not out_inited:
outdims = _calc_out_dims(res[0], extra_dims)
output = [n.zeros(outdims, ref_var.dtype) output = [n.zeros(outdims, ref_var.dtype)
for i in xrange(len(res))] for i in xrange(len(res))]
out_inited = True out_inited = True
@ -250,8 +236,8 @@ def uvmet_left_iter():
return indexing_decorator return indexing_decorator
def handle_casting(ref_idx=0, arg_idxs=(), karg_names=None,dtype=n.float64): def handle_casting(ref_idx=0, arg_idxs=(), karg_names=None,dtype=n.float64):
"""Decorator to handle iterating over leftmost dimensions when using """Decorator to handle casting to/from required dtype used in
multiple files and/or multiple times with the uvmet product. numerical routine.
""" """
def cast_decorator(func): def cast_decorator(func):

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

@ -23,7 +23,7 @@ class FortranException(Exception):
def __call__(self, message): def __call__(self, message):
raise self.__class__(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)) @handle_casting(arg_idxs=(0,1))
def interpz3d(data3d,zdata,desiredloc,missingval): def interpz3d(data3d,zdata,desiredloc,missingval):
res = f_interpz3d(data3d.T, res = f_interpz3d(data3d.T,
@ -32,9 +32,9 @@ def interpz3d(data3d,zdata,desiredloc,missingval):
missingval) missingval)
return res.T return res.T
@handle_left_iter(2,0)
@handle_casting(arg_idxs=(0,1)) @handle_casting(arg_idxs=(0,1))
@handle_left_iter(2,2,0) def interp2dxy(data3d,xy):
def interpz2d(data3d,xy):
res = f_interp2dxy(data3d.T, res = f_interp2dxy(data3d.T,
xy.T) xy.T)
return res.T return res.T
@ -48,7 +48,7 @@ def interp1d(v_in,z_in,z_out,missingval):
return res return res
@handle_left_iter(2,3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2,3)) @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")
@ -66,7 +66,7 @@ def computeslp(z,t,p,q):
return res.T return res.T
@handle_left_iter(3,3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1)) @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
@ -95,7 +95,7 @@ def computerh(qv,q,t):
res = n.reshape(res, shape, "A") res = n.reshape(res, shape, "A")
return res 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)) @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.T, res = f_computeabsvort(u.T,
@ -109,7 +109,7 @@ def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy):
return res.T 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)) @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):
@ -126,7 +126,7 @@ def computepvo(u,v,theta,prs,msfu,msfv,msfm,cor,dx,dy):
return res.T return res.T
@handle_left_iter(3,3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2)) @handle_casting(arg_idxs=(0,1,2))
def computeeth(qv, tk, p): def computeeth(qv, tk, p):
@ -169,7 +169,7 @@ def computeuvmet(u,v,lat,lon,cen_long,cone):
return n.squeeze(res.T) return n.squeeze(res.T)
@handle_left_iter(3,3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2,3)) @handle_casting(arg_idxs=(0,1,2,3))
def computeomega(qv, tk, w, p): def computeomega(qv, tk, w, p):
@ -181,7 +181,7 @@ def computeomega(qv, tk, w, p):
#return res.T #return res.T
return res.T return res.T
@handle_left_iter(3,3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1)) @handle_casting(arg_idxs=(0,1))
def computetv(tk,qv): def computetv(tk,qv):
res = f_computetv(tk.T, res = f_computetv(tk.T,
@ -189,7 +189,7 @@ def computetv(tk,qv):
return res.T return res.T
@handle_left_iter(3,3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2)) @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()
@ -204,7 +204,7 @@ def computewetbulb(p,tk,qv):
return res.T 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)) @handle_casting(arg_idxs=(0,1,2,3))
def computesrh(u, v, z, ter, top): def computesrh(u, v, z, ter, top):
@ -216,7 +216,7 @@ def computesrh(u, v, z, ter, top):
return res.T 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)) @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):
@ -237,7 +237,7 @@ def computeuh(zp, mapfct, u, v, wstag, dx, dy, bottom, top):
return res.T return res.T
@handle_left_iter(2,3,0) @handle_left_iter(3,0)
@handle_casting(arg_idxs=(0,1,2,3)) @handle_casting(arg_idxs=(0,1,2,3))
def computepw(p,tv,qv,ht): def computepw(p,tv,qv,ht):
# Note, dim -3 is height, we only want y and x # Note, dim -3 is height, we only want y and x
@ -250,7 +250,7 @@ def computepw(p,tv,qv,ht):
return res.T 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)) @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):
@ -266,7 +266,7 @@ def computedbz(p,tk,qv,qr,qs,qg,sn0,ivarint,iliqskin):
return res.T 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)) @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]), 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 # Want lon,lat
return res[1],res[0] 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)) @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")
@ -347,7 +347,7 @@ def computeeta(full_t, znu, psfc, ptop):
return res.T 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)) @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.T, res = f_computectt(p_hpa.T,

86
wrf_open/var/src/python/wrf/var/interp.py
Unescape View File

@ -3,11 +3,13 @@ from math import floor, ceil
import numpy as n import numpy as n
import numpy.ma as ma 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 """Return the horizontally interpolated data at the provided level
data3d - the 3D field to interpolate data3d - the 3D field to interpolate
@ -27,17 +29,17 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None,
xdim - maximum x-dimension xdim - maximum x-dimension
ydim - maximum y-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 angle - the angle through the pivot point in degrees
start_point - a start_point tuple of (x,y) start_point - a start_point tuple of (south_north1, west_east1)
end_point - an end point tuple of (x,y) end_point - an end point tuple of (south_north2, west_east2)
""" """
# Have a pivot point with an angle to find cross section # Have a pivot point with an angle to find cross section
if pivot_point is not None and angle is not None: if pivot_point is not None and angle is not None:
xp = pivot_point[0] xp = pivot_point[-1]
yp = pivot_point[1] yp = pivot_point[-2]
if (angle > 315.0 or angle < 45.0 if (angle > 315.0 or angle < 45.0
or ((angle > 135.0) and (angle < 225.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 y1 = ydim-1
x1 = (y1 - intercept)/slope x1 = (y1 - intercept)/slope
elif start_point is not None and end_point is not None: elif start_point is not None and end_point is not None:
x0 = start_point[0] x0 = start_point[-1]
y0 = start_point[1] y0 = start_point[-2]
x1 = end_point[0] x1 = end_point[-1]
y1 = end_point[1] y1 = end_point[-2]
if ( x1 > xdim-1 ): if ( x1 > xdim-1 ):
x1 = xdim x1 = xdim
if ( y1 > ydim-1): if ( y1 > ydim-1):
@ -129,10 +131,25 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None,
return xy return xy
# TODO: Need a decorator that can handle right dimensions based on what
# TODO: Add flag to use lat/lon points by doing conversion # is returned. In this case, the result of the xy calculation determines
def get_vertcross(data3d, z, missingval=-99999, # output dimensions on right.
pivot_point=None,angle=None,start_point=None,end_point=None): @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] xdim = z.shape[-1]
ydim = z.shape[-2] 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) xy = _get_xy(xdim, ydim, pivot_point, angle, start_point, end_point)
# Interp z # Interp z
var2dz = interpz2d(z, xy) var2dz = interp2dxy(z, xy)
# interp to constant z grid # interp to constant z grid
if(var2dz[0,0] > var2dz[1,0]): # monotonically decreasing coordinate if(var2dz[0,0] > var2dz[1,0]): # monotonically decreasing coordinate
@ -165,13 +182,46 @@ def get_vertcross(data3d, z, missingval=-99999,
#interp the variable #interp the variable
var2d = n.zeros((nlevels, xy.shape[0]),dtype=var2dz.dtype) 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]): for i in xrange(xy.shape[0]):
var2d[:,i] = interp1d(var2dtmp[:,i], var2dz[:,i], z_var2d, missingval) var2d[:,i] = interp1d(var2dtmp[:,i], var2dz[:,i], z_var2d, missingval)
return ma.masked_values(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,:]

Write Preview
Loading…
Cancel
Save