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
commit
f45ac9756e
  1. 40
      wrf_open/var/src/python/wrf/var/decorators.py
  2. 36
      wrf_open/var/src/python/wrf/var/extension.py
  3. 84
      wrf_open/var/src/python/wrf/var/interp.py

40
wrf_open/var/src/python/wrf/var/decorators.py

@ -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
@ -102,26 +105,7 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1,
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

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

84
wrf_open/var/src/python/wrf/var/interp.py

@ -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
# is returned. In this case, the result of the xy calculation determines
# output dimensions on right.
@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)
# TODO: Add flag to use lat/lon points by doing conversion """
def get_vertcross(data3d, z, missingval=-99999,
pivot_point=None,angle=None,start_point=None,end_point=None):
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,:]

Loading…
Cancel
Save