forked from 3rdparty/wrf-python
Bill Ladwig
10 years ago
37 changed files with 1764 additions and 989 deletions
@ -0,0 +1,182 @@
@@ -0,0 +1,182 @@
|
||||
from __future__ import (absolute_import, division, print_function, |
||||
unicode_literals) |
||||
|
||||
from math import floor, ceil |
||||
|
||||
import numpy as np |
||||
|
||||
from .extension import interp2dxy |
||||
|
||||
__all__ = ["to_positive_idxs", "calc_xy", "get_xy_z_params", "get_xy"] |
||||
|
||||
def to_positive_idxs(shape, coord): |
||||
if (coord[-2] >= 0 and coord[-1] >= 0): |
||||
return coord |
||||
|
||||
return [x if (x >= 0) else shape[i]+x for (i,x) in enumerate(coord) ] |
||||
|
||||
def calc_xy(xdim, ydim, pivot_point=None, angle=None, |
||||
start_point=None, end_point=None): |
||||
"""Returns the x,y points for the horizontal cross section line. |
||||
|
||||
xdim - maximum x-dimension |
||||
ydim - maximum y-dimension |
||||
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 sequence of [south_north1, west_east1] |
||||
end_point - an end point sequence of [south_north2, west_east2] |
||||
|
||||
""" |
||||
|
||||
# Have a pivot point with an angle to find cross section |
||||
if pivot_point is not None and angle is not None: |
||||
xp = pivot_point[-1] |
||||
yp = pivot_point[-2] |
||||
|
||||
if (angle > 315.0 or angle < 45.0 |
||||
or ((angle > 135.0) and (angle < 225.0))): |
||||
|
||||
#x = y*slope + intercept |
||||
slope = -(360.-angle)/45. |
||||
if( angle < 45. ): |
||||
slope = angle/45. |
||||
if( angle > 135.): |
||||
slope = (angle-180.)/45. |
||||
|
||||
intercept = xp - yp*slope |
||||
|
||||
# find intersections with domain boundaries |
||||
y0 = 0. |
||||
x0 = y0*slope + intercept |
||||
|
||||
if( x0 < 0.): # intersect outside of left boundary |
||||
x0 = 0. |
||||
y0 = (x0 - intercept)/slope |
||||
if( x0 > xdim-1): #intersect outside of right boundary |
||||
x0 = xdim-1 |
||||
y0 = (x0 - intercept)/slope |
||||
y1 = ydim-1. #need to make sure this will be a float? |
||||
x1 = y1*slope + intercept |
||||
|
||||
if( x1 < 0.): # intersect outside of left boundary |
||||
x1 = 0. |
||||
y1 = (x1 - intercept)/slope |
||||
|
||||
if( x1 > xdim-1): # intersect outside of right boundary |
||||
x1 = xdim-1 |
||||
y1 = (x1 - intercept)/slope |
||||
else: |
||||
# y = x*slope + intercept |
||||
slope = (90.-angle)/45. |
||||
if( angle > 225. ): |
||||
slope = (270.-angle)/45. |
||||
intercept = yp - xp*slope |
||||
|
||||
#find intersections with domain boundaries |
||||
x0 = 0. |
||||
y0 = x0*slope + intercept |
||||
|
||||
if( y0 < 0.): # intersect outside of bottom boundary |
||||
y0 = 0. |
||||
x0 = (y0 - intercept)/slope |
||||
|
||||
if( y0 > ydim-1): # intersect outside of top boundary |
||||
y0 = ydim-1 |
||||
x0 = (y0 - intercept)/slope |
||||
|
||||
x1 = xdim-1. # need to make sure this will be a float? |
||||
y1 = x1*slope + intercept |
||||
|
||||
if( y1 < 0.): # intersect outside of bottom boundary |
||||
y1 = 0. |
||||
x1 = (y1 - intercept)/slope |
||||
|
||||
if( y1 > ydim-1):# intersect outside of top boundary |
||||
y1 = ydim-1 |
||||
x1 = (y1 - intercept)/slope |
||||
elif start_point is not None and end_point is not None: |
||||
x0 = start_point[-1] |
||||
y0 = start_point[-2] |
||||
x1 = end_point[-1] |
||||
y1 = end_point[-2] |
||||
if ( x1 > xdim-1 ): |
||||
x1 = xdim |
||||
if ( y1 > ydim-1): |
||||
y1 = ydim |
||||
else: |
||||
raise ValueError("invalid start/end or pivot/angle arguments") |
||||
|
||||
dx = x1 - x0 |
||||
dy = y1 - y0 |
||||
distance = (dx*dx + dy*dy)**0.5 |
||||
npts = int(distance) |
||||
dxy = distance/npts |
||||
|
||||
xy = np.zeros((npts,2), "float") |
||||
|
||||
dx = dx/npts |
||||
dy = dy/npts |
||||
|
||||
for i in xrange(npts): |
||||
xy[i,0] = x0 + i*dx |
||||
xy[i,1] = y0 + i*dy |
||||
|
||||
return xy |
||||
|
||||
def get_xy_z_params(z, pivot_point=None, angle=None, |
||||
start_point=None, end_point=None): |
||||
|
||||
xy = get_xy(z, pivot_point, angle, start_point, end_point) |
||||
|
||||
# Interp z |
||||
var2dz = interp2dxy(z, xy) |
||||
|
||||
extra_dim_num = z.ndim - 3 |
||||
idx1 = tuple([0]*extra_dim_num + [0,0]) |
||||
idx2 = tuple([0]*extra_dim_num + [-1,0]) |
||||
|
||||
# interp to constant z grid |
||||
if(var2dz[idx1] > var2dz[idx2]): # monotonically decreasing coordinate |
||||
z_max = floor(np.amax(z) / 10) * 10 # bottom value |
||||
z_min = ceil(np.amin(z) / 10) * 10 # top value |
||||
dz = 10 |
||||
nlevels = int((z_max-z_min) / dz) |
||||
z_var2d = np.zeros((nlevels), dtype=z.dtype) |
||||
z_var2d[0] = z_max |
||||
dz = -dz |
||||
else: |
||||
z_max = np.amax(z) |
||||
z_min = 0. |
||||
dz = 0.01 * z_max |
||||
nlevels = int(z_max / dz) |
||||
z_var2d = np.zeros((nlevels), dtype=z.dtype) |
||||
z_var2d[0] = z_min |
||||
|
||||
for i in xrange(1,nlevels): |
||||
z_var2d[i] = z_var2d[0] + i*dz |
||||
|
||||
return xy, var2dz, z_var2d |
||||
|
||||
def get_xy(var, pivot_point=None, angle=None, start_point=None, end_point=None): |
||||
if pivot_point is not None: |
||||
pos_pivot = to_positive_idxs(var.shape[-2:], pivot_point) |
||||
else: |
||||
pos_pivot = pivot_point |
||||
|
||||
if start_point is not None: |
||||
pos_start = to_positive_idxs(var.shape[-2:], start_point) |
||||
else: |
||||
pos_start = start_point |
||||
|
||||
if end_point is not None: |
||||
pos_end = to_positive_idxs(var.shape[-2:], end_point) |
||||
else: |
||||
pos_end = start_point |
||||
|
||||
xdim = var.shape[-1] |
||||
ydim = var.shape[-2] |
||||
|
||||
xy = calc_xy(xdim, ydim, pos_pivot, angle, pos_start, pos_end) |
||||
|
||||
return xy |
||||
@ -0,0 +1,695 @@
@@ -0,0 +1,695 @@
|
||||
from __future__ import (absolute_import, division, print_function, |
||||
unicode_literals) |
||||
import wrapt |
||||
from collections import OrderedDict |
||||
|
||||
import numpy as np |
||||
import numpy.ma as ma |
||||
|
||||
|
||||
from .util import (viewkeys, viewitems, extract_vars, |
||||
combine_with, either, from_args, arg_location, |
||||
_is_coord_var, XYCoord, npvalues) |
||||
from .interputils import get_xy_z_params, get_xy |
||||
from .config import xarray_enabled |
||||
|
||||
if xarray_enabled(): |
||||
from xarray import DataArray |
||||
|
||||
__all__ = ["copy_and_set_metadata", "set_wind_metadata", |
||||
"set_latlon_metadata", "set_height_metadata", |
||||
"set_interp_metadata"] |
||||
|
||||
def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, |
||||
remove_dims=None, dimnames=None, |
||||
coords=None, **fixed_attrs): |
||||
"""Decorator to set the metadata for a WRF method. |
||||
|
||||
A cache is inserted/updated to include the extracted variable that will |
||||
have its metadata copied. This prevents the variable being extracted more |
||||
than once. This extraction can be slow with sequences of large files. |
||||
|
||||
""" |
||||
@wrapt.decorator |
||||
def func_wrapper(wrapped, instance, args, kwargs): |
||||
if not xarray_enabled(): |
||||
return wrapped(*args, **kwargs) |
||||
|
||||
argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", |
||||
"squeeze", "cache", "units"), |
||||
*args, **kwargs) |
||||
wrfnc = argvars["wrfnc"] |
||||
timeidx = argvars["timeidx"] |
||||
units = argvars["units"] |
||||
method = argvars["method"] |
||||
squeeze = argvars["squeeze"] |
||||
cache = argvars["cache"] |
||||
if cache is None: |
||||
cache = {} |
||||
|
||||
# Note: can't modify nonlocal var |
||||
if (callable(copy_varname)): |
||||
_copy_varname = copy_varname(wrfnc) |
||||
else: |
||||
_copy_varname = copy_varname |
||||
|
||||
# Extract the copy_from argument |
||||
var_to_copy = None if cache is None else cache.get(_copy_varname, |
||||
None) |
||||
|
||||
|
||||
if var_to_copy is None: |
||||
var_to_copy = extract_vars(wrfnc, timeidx, (_copy_varname,), |
||||
method, squeeze, cache, |
||||
nometa=False)[_copy_varname] |
||||
|
||||
# Make a copy so we don't modify a user supplied cache |
||||
new_cache = dict(cache) |
||||
new_cache[_copy_varname] = var_to_copy |
||||
|
||||
# Don't modify the original args/kargs. The args need to be a list |
||||
# so it can be modified. |
||||
new_args, cache_argloc = arg_location(wrapped, "cache", args, kwargs) |
||||
new_args[cache_argloc] = new_cache |
||||
|
||||
result = wrapped(*new_args) |
||||
|
||||
outname = "" |
||||
outdimnames = list() |
||||
outcoords = OrderedDict() |
||||
outattrs = OrderedDict() |
||||
|
||||
if copy_varname is not None: |
||||
outname = unicode(var_to_copy.name) |
||||
|
||||
if dimnames is not None: |
||||
if isinstance(dimnames, combine_with): |
||||
outdimnames, outcoords = dimnames(var_to_copy) |
||||
else: |
||||
outdimnames = dimnames |
||||
outcoords = coords |
||||
else: |
||||
outdimnames += var_to_copy.dims |
||||
outcoords.update(var_to_copy.coords) |
||||
|
||||
outattrs.update(var_to_copy.attrs) |
||||
|
||||
if remove_dims is not None: |
||||
for dimname in remove_dims: |
||||
outdimnames.remove(dimname) |
||||
|
||||
try: |
||||
del outcoords[dimname] |
||||
except KeyError: |
||||
pass |
||||
|
||||
|
||||
if name is not None: |
||||
outname = name |
||||
|
||||
if units is not None: |
||||
outattrs["units"] = units |
||||
|
||||
for argname, val in viewitems(fixed_attrs): |
||||
outattrs[argname] = val |
||||
|
||||
if delete_attrs is not None: |
||||
for attr in delete_attrs: |
||||
try: |
||||
del outattrs[attr] |
||||
except KeyError: |
||||
pass |
||||
|
||||
if isinstance(result, ma.MaskedArray): |
||||
outattrs["_FillValue"] = result.fill_value |
||||
outattrs["missing_value"] = result.fill_value |
||||
|
||||
return DataArray(result, name=outname, coords=outcoords, |
||||
dims=outdimnames, attrs=outattrs) |
||||
|
||||
return func_wrapper |
||||
|
||||
|
||||
def set_wind_metadata(copy_varname, name, description, |
||||
wind_ncvar=False, |
||||
two_d=False, wspd_wdir=False): |
||||
@wrapt.decorator |
||||
def func_wrapper(wrapped, instance, args, kwargs): |
||||
if not xarray_enabled(): |
||||
return wrapped(*args, **kwargs) |
||||
|
||||
argvars = from_args(wrapped, ("wrfnc", "timeidx", "units", |
||||
"method", "squeeze", "ten_m", "cache"), |
||||
*args, **kwargs) |
||||
wrfnc = argvars["wrfnc"] |
||||
timeidx = argvars["timeidx"] |
||||
units = argvars["units"] |
||||
method = argvars["method"] |
||||
squeeze = argvars["squeeze"] |
||||
ten_m = argvars["ten_m"] |
||||
cache = argvars["cache"] |
||||
if cache is None: |
||||
cache = {} |
||||
|
||||
if isinstance(copy_varname, either): |
||||
_copy_varname = copy_varname(wrfnc) |
||||
else: |
||||
_copy_varname = copy_varname |
||||
|
||||
copy_var = extract_vars(wrfnc, timeidx, _copy_varname, |
||||
method, squeeze, cache, |
||||
nometa=False)[_copy_varname] |
||||
|
||||
# Make a copy so we don't modify a user supplied cache |
||||
new_cache = dict(cache) |
||||
new_cache[_copy_varname] = copy_var |
||||
|
||||
# Don't modify the original args/kargs. The args need to be a list |
||||
# so it can be modified. |
||||
new_args, cache_argloc = arg_location(wrapped, "cache", args, kwargs) |
||||
new_args[cache_argloc] = new_cache |
||||
|
||||
result = wrapped(*new_args) |
||||
|
||||
outcoords = OrderedDict() |
||||
outattrs = OrderedDict() |
||||
|
||||
outdimnames = list(copy_var.dims) |
||||
outcoords.update(copy_var.coords) |
||||
outattrs.update(copy_var.attrs) |
||||
|
||||
if wind_ncvar: |
||||
pass |
||||
|
||||
elif not wspd_wdir: |
||||
if not two_d: |
||||
outdimnames.insert(-3, "u_v") |
||||
else: |
||||
outdimnames.insert(-2, "u_v") |
||||
outattrs["MemoryOrder"] = "XY" |
||||
outcoords["u_v"] = ["u", "v"] |
||||
else: |
||||
if not two_d: |
||||
outdimnames.insert(-3, "wspd_wdir") |
||||
else: |
||||
outdimnames.insert(-2, "wspd_wdir") |
||||
outattrs["MemoryOrder"] = "XY" |
||||
|
||||
outcoords["wspd_wdir"] = ["wspd", "wdir"] |
||||
|
||||
if units is not None: |
||||
outattrs["units"] = units |
||||
|
||||
outname = name |
||||
outattrs["description"] = description |
||||
|
||||
return DataArray(result, name=outname, coords=outcoords, |
||||
dims=outdimnames, attrs=outattrs) |
||||
|
||||
return func_wrapper |
||||
|
||||
def set_latlon_metadata(ij=False): |
||||
@wrapt.decorator |
||||
def func_wrapper(wrapped, instance, args, kwargs): |
||||
if not xarray_enabled(): |
||||
return wrapped(*args, **kwargs) |
||||
|
||||
res = wrapped(*args, **kwargs) |
||||
|
||||
argnames = ("latitude", "longitude") if not ij else ("i", "j") |
||||
outname = "latlon" if not ij else "ij" |
||||
|
||||
if res.ndim <= 2: |
||||
dimnames = (["lat_lon", "i_j"] if not ij |
||||
else ["i_j", "lat_lon"]) |
||||
else: |
||||
dimnames = (["lat_lon", "domain", "i_j"] if not ij |
||||
else ["i_j", "domain", "lat_lon"]) |
||||
|
||||
|
||||
argvars = from_args(wrapped, argnames, *args, **kwargs) |
||||
|
||||
var1 = argvars[argnames[0]] |
||||
var2 = argvars[argnames[1]] |
||||
|
||||
arr1 = np.asarray(var1).ravel() |
||||
arr2 = np.asarray(var2).ravel() |
||||
|
||||
coords = {} |
||||
coords[dimnames[0]] = [x for x in zip(arr1, arr2)] |
||||
|
||||
return DataArray(res, name=outname, dims=dimnames, coords=coords) |
||||
|
||||
return func_wrapper |
||||
|
||||
def set_height_metadata(geopt=False): |
||||
@wrapt.decorator |
||||
def func_wrapper(wrapped, instance, args, kwargs): |
||||
if not xarray_enabled(): |
||||
return wrapped(*args, **kwargs) |
||||
|
||||
argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", |
||||
"squeeze", "units", "msl", "cache"), |
||||
*args, **kwargs) |
||||
wrfnc = argvars["wrfnc"] |
||||
timeidx = argvars["timeidx"] |
||||
units = argvars["units"] |
||||
method = argvars["method"] |
||||
squeeze = argvars["squeeze"] |
||||
msl = argvars["msl"] |
||||
cache = argvars["cache"] |
||||
|
||||
if cache is None: |
||||
cache = {} |
||||
|
||||
# For height, either copy the met_em GHT variable or copy and modify |
||||
# pressure (which has the same dims as destaggered height) |
||||
ht_metadata_varname = either("P", "GHT")(wrfnc) |
||||
ht_var = extract_vars(wrfnc, timeidx, ht_metadata_varname, |
||||
method, squeeze, cache, nometa=False) |
||||
ht_metadata_var = ht_var[ht_metadata_varname] |
||||
|
||||
# Make a copy so we don't modify a user supplied cache |
||||
new_cache = dict(cache) |
||||
new_cache[ht_metadata_varname] = ht_metadata_var |
||||
|
||||
# Don't modify the original args/kargs. The args need to be a list |
||||
# so it can be modified. |
||||
new_args, cache_argloc = arg_location(wrapped, "cache", args, kwargs) |
||||
new_args[cache_argloc] = new_cache |
||||
|
||||
result = wrapped(*new_args) |
||||
|
||||
outcoords = OrderedDict() |
||||
outattrs = OrderedDict() |
||||
outdimnames = list(ht_metadata_var.dims) |
||||
outcoords.update(ht_metadata_var.coords) |
||||
outattrs.update(ht_metadata_var.attrs) |
||||
|
||||
if geopt: |
||||
outname = "geopt" |
||||
outattrs["units"] = "m2 s-2" |
||||
outattrs["description"] = "full model geopotential" |
||||
else: |
||||
outname = "height" if msl else "height_agl" |
||||
outattrs["units"] = units |
||||
height_type = "MSL" if msl else "AGL" |
||||
outattrs["description"] = "model height ({})".format(height_type) |
||||
|
||||
|
||||
return DataArray(result, name=outname, |
||||
dims=outdimnames, coords=outcoords, attrs=outattrs) |
||||
return func_wrapper |
||||
|
||||
def _set_horiz_meta(wrapped, instance, args, kwargs): |
||||
argvars = from_args(wrapped, ("field3d", "z", "desiredloc", |
||||
"missingval"), |
||||
*args, **kwargs) |
||||
|
||||
field3d = argvars["field3d"] |
||||
z = argvars["z"] |
||||
desiredloc = argvars["desiredloc"] |
||||
missingval = argvars["missingval"] |
||||
|
||||
result = wrapped(*args, **kwargs) |
||||
|
||||
# Defaults, in case the data isn't a DataArray |
||||
outname = None |
||||
outdimnames = None |
||||
outcoords = None |
||||
outattrs = None |
||||
|
||||
# Get the vertical level units |
||||
vert_units = None |
||||
if isinstance(z, DataArray): |
||||
vert_units = z.attrs.get("units", None) |
||||
|
||||
# If we have no metadata to start with, only set the level |
||||
levelstr = ("{0} {1}".format(desiredloc, vert_units) |
||||
if vert_units is not None |
||||
else "{0}".format(desiredloc)) |
||||
|
||||
if isinstance(field3d, DataArray): |
||||
outcoords = OrderedDict() |
||||
outattrs = OrderedDict() |
||||
outdimnames = list(field3d.dims) |
||||
outcoords.update(field3d.coords) |
||||
outdimnames.remove(field3d.dims[-3]) |
||||
del outcoords[field3d.dims[-3]] |
||||
outattrs.update(field3d.attrs) |
||||
outname = "{0}_{1}".format(field3d.name, levelstr) |
||||
|
||||
else: |
||||
outname = "field3d_{0}".format(levelstr) |
||||
outattrs = OrderedDict() |
||||
|
||||
outattrs["PlotLevelID"] = levelstr |
||||
outattrs["missing_value"] = missingval |
||||
outattrs["_FillValue"] = missingval |
||||
|
||||
for key in ("MemoryOrder", "description"): |
||||
try: |
||||
del outattrs[key] |
||||
except KeyError: |
||||
pass |
||||
|
||||
return DataArray(result, name=outname, dims=outdimnames, |
||||
coords=outcoords, attrs=outattrs) |
||||
|
||||
def _set_cross_meta(wrapped, instance, args, kwargs): |
||||
argvars = from_args(wrapped, ("field3d", "z", "missingval", |
||||
"pivot_point", "angle", |
||||
"start_point", "end_point", |
||||
"cache"), |
||||
*args, **kwargs) |
||||
|
||||
field3d = argvars["field3d"] |
||||
z = argvars["z"] |
||||
missingval = argvars["missingval"] |
||||
pivot_point = argvars["pivot_point"] |
||||
angle = argvars["angle"] |
||||
start_point = argvars["start_point"] |
||||
end_point = argvars["end_point"] |
||||
cache = argvars["cache"] |
||||
|
||||
xy, var2dz, z_var2d = get_xy_z_params(npvalues(z), pivot_point, angle, |
||||
start_point, end_point) |
||||
|
||||
# Make a copy so we don't modify a user supplied cache |
||||
if cache is not None: |
||||
new_cache = dict(cache) |
||||
else: |
||||
new_cache = {} |
||||
new_cache["xy"] = xy |
||||
new_cache["var2dz"] = var2dz |
||||
new_cache["z_var2d"] = z_var2d |
||||
|
||||
# Don't modify the original args/kargs. The args need to be a list |
||||
# so it can be modified. |
||||
new_args, cache_argloc = arg_location(wrapped, "cache", args, kwargs) |
||||
new_args[cache_argloc] = new_cache |
||||
|
||||
result = wrapped(*new_args) |
||||
|
||||
# Defaults, in case the data isn't a DataArray |
||||
outname = None |
||||
outdimnames = None |
||||
outcoords = None |
||||
outattrs = None |
||||
|
||||
# Use XY to set the cross-section metadata |
||||
st_x = xy[0,0] |
||||
st_y = xy[0,1] |
||||
ed_x = xy[-1,0] |
||||
ed_y = xy[-1,1] |
||||
|
||||
cross_str = "cross-section: ({0}, {1}) to ({2}, {3})".format(st_x, st_y, |
||||
ed_x, ed_y) |
||||
if angle is not None: |
||||
cross_str += " ; center={0} ; angle={1}".format(pivot_point, |
||||
angle) |
||||
|
||||
if isinstance(field3d, DataArray): |
||||
outcoords = OrderedDict() |
||||
outattrs = OrderedDict() |
||||
outdimnames = list(field3d.dims) |
||||
outcoords.update(field3d.coords) |
||||
for i in xrange(-3,0,1): |
||||
outdimnames.remove(field3d.dims[i]) |
||||
del outcoords[field3d.dims[i]] |
||||
|
||||
|
||||
# Delete any lat,lon coords |
||||
for key in viewkeys(outcoords): |
||||
if _is_coord_var(key): |
||||
del outcoords[key] |
||||
|
||||
outdimnames.append("vertical") |
||||
outdimnames.append("xy") |
||||
outattrs.update(field3d.attrs) |
||||
|
||||
outname = "{0}_cross".format(field3d.name) |
||||
|
||||
for key in ("MemoryOrder",): |
||||
try: |
||||
del outattrs[key] |
||||
except KeyError: |
||||
pass |
||||
|
||||
outcoords["xy"] = [XYCoord(xy[i,0], xy[i,1]) |
||||
for i in xrange(xy.shape[-2])] |
||||
|
||||
outcoords["vertical"] = z_var2d[:] |
||||
|
||||
else: |
||||
outname = "field3d_cross" |
||||
outattrs = OrderedDict() |
||||
|
||||
outattrs["orientation"] = cross_str |
||||
outattrs["missing_value"] = missingval |
||||
outattrs["_FillValue"] = missingval |
||||
|
||||
return DataArray(result, name=outname, dims=outdimnames, |
||||
coords=outcoords, attrs=outattrs) |
||||
|
||||
|
||||
|
||||
def _set_line_meta(wrapped, instance, args, kwargs): |
||||
argvars = from_args(wrapped, ("field2d", "pivot_point", "angle", |
||||
"start_point", "end_point", "cache"), |
||||
*args, **kwargs) |
||||
|
||||
field2d = argvars["field2d"] |
||||
pivot_point = argvars["pivot_point"] |
||||
angle = argvars["angle"] |
||||
start_point = argvars["start_point"] |
||||
end_point = argvars["end_point"] |
||||
cache = argvars["cache"] |
||||
|
||||
if cache is None: |
||||
cache = {} |
||||
|
||||
xy = get_xy(field2d, pivot_point, angle, start_point, end_point) |
||||
|
||||
# Make a copy so we don't modify a user supplied cache |
||||
new_cache = dict(cache) |
||||
new_cache["xy"] = xy |
||||
|
||||
# Don't modify the original args/kargs. The args need to be a list |
||||
# so it can be modified. |
||||
new_args, cache_argloc = arg_location(wrapped, "cache", args, kwargs) |
||||
new_args[cache_argloc] = new_cache |
||||
|
||||
result = wrapped(*new_args) |
||||
|
||||
# Defaults, in case the data isn't a DataArray |
||||
outname = None |
||||
outdimnames = None |
||||
outcoords = None |
||||
outattrs = None |
||||
|
||||
# Use XY to set the cross-section metadata |
||||
st_x = xy[0,0] |
||||
st_y = xy[0,1] |
||||
ed_x = xy[-1,0] |
||||
ed_y = xy[-1,1] |
||||
|
||||
cross_str = "({0}, {1}) to ({2}, {3})".format(st_x, st_y, |
||||
ed_x, ed_y) |
||||
if angle is not None: |
||||
cross_str += " ; center={0} ; angle={1}".format(pivot_point, |
||||
angle) |
||||
|
||||
if isinstance(field2d, DataArray): |
||||
outcoords = OrderedDict() |
||||
outattrs = OrderedDict() |
||||
outdimnames = list(field2d.dims) |
||||
outcoords.update(field2d.coords) |
||||
for i in xrange(-2,0,1): |
||||
outdimnames.remove(field2d.dims[i]) |
||||
del outcoords[field2d.dims[i]] |
||||
|
||||
# Delete any lat,lon coords |
||||
for key in viewkeys(outcoords): |
||||
if _is_coord_var(key): |
||||
del outcoords[key] |
||||
|
||||
outdimnames.append("xy") |
||||
outattrs.update(field2d.attrs) |
||||
|
||||
outname = "{0}_line".format(field2d.name) |
||||
|
||||
for key in ("MemoryOrder",): |
||||
try: |
||||
del outattrs[key] |
||||
except KeyError: |
||||
pass |
||||
|
||||
outcoords["xy"] = [XYCoord(xy[i,0], xy[i,1]) |
||||
for i in xrange(xy.shape[-2])] |
||||
|
||||
else: |
||||
outname = "field2d_line" |
||||
outattrs = OrderedDict() |
||||
|
||||
outattrs["orientation"] = cross_str |
||||
|
||||
return DataArray(result, name=outname, dims=outdimnames, |
||||
coords=outcoords, attrs=outattrs) |
||||
|
||||
|
||||
def _set_vinterp_meta(wrapped, instance, args, kwargs): |
||||
argvars = from_args(wrapped, ("wrfnc", "field", "vert_coord", |
||||
"interp_levels", "extrapolate", |
||||
"field_type", "log_p", |
||||
"timeidx", "method", "squeeze", |
||||
"cache"), |
||||
*args, **kwargs) |
||||
|
||||
field = argvars["field"] |
||||
vert_coord = argvars["vert_coord"] |
||||
interp_levels = argvars["interp_levels"] |
||||
field_type = argvars["field_type"] |
||||
|
||||
result = wrapped(*args, **kwargs) |
||||
|
||||
# Defaults, in case the data isn't a DataArray |
||||
outname = None |
||||
outdimnames = None |
||||
outcoords = None |
||||
outattrs = None |
||||
|
||||
|
||||
if isinstance(field, DataArray): |
||||
outcoords = OrderedDict() |
||||
outattrs = OrderedDict() |
||||
outdimnames = list(field.dims) |
||||
outcoords.update(field.coords) |
||||
|
||||
outdimnames.remove(field.dims[-3]) |
||||
del outcoords[field.dims[-3]] |
||||
|
||||
outdimnames.insert(-2, "interp_level") |
||||
outcoords["interp_level"] = interp_levels |
||||
outattrs.update(field.attrs) |
||||
outattrs["vert_interp_type"] = vert_coord |
||||
|
||||
outname = field.name |
||||
|
||||
else: |
||||
outname = field_type |
||||
|
||||
|
||||
return DataArray(result, name=outname, dims=outdimnames, |
||||
coords=outcoords, attrs=outattrs) |
||||
|
||||
|
||||
def _set_2dxy_meta(wrapped, instance, args, kwargs): |
||||
argvars = from_args(wrapped, ("field3d", "xy"), |
||||
*args, **kwargs) |
||||
|
||||
field3d = argvars["field3d"] |
||||
xy = argvars["xy"] |
||||
|
||||
result = wrapped(*args, **kwargs) |
||||
|
||||
# Use XY to set the cross-section metadata |
||||
st_x = xy[0,0] |
||||
st_y = xy[0,1] |
||||
ed_x = xy[-1,0] |
||||
ed_y = xy[-1,1] |
||||
|
||||
cross_str = "({0},{1}) to ({2},{3})".format(st_x, st_y, |
||||
ed_x, ed_y) |
||||
|
||||
# Dims are (...,xy,z) |
||||
if isinstance(field3d, DataArray): |
||||
outcoords = OrderedDict() |
||||
outattrs = OrderedDict() |
||||
outdimnames = list(field3d.dims) |
||||
outcoords.update(field3d.coords) |
||||
for i in xrange(-2,0,1): |
||||
outdimnames.remove(field3d.dims[i]) |
||||
del outcoords[field3d.dims[i]] |
||||
|
||||
outdimnames[-2] = "xy" |
||||
outattrs.update(field3d.attrs) |
||||
|
||||
outname = "{0}_xy".format(field3d.name) |
||||
|
||||
outcoords["xy"] = [XYCoord(xy[i,0], xy[i,1]) |
||||
for i in xrange(xy.shape[-2])] |
||||
|
||||
for key in ("MemoryOrder",): |
||||
try: |
||||
del outattrs[key] |
||||
except KeyError: |
||||
pass |
||||
|
||||
else: |
||||
outname = "field3d_xy" |
||||
|
||||
outattrs["Orientation"] = cross_str |
||||
|
||||
return DataArray(result, name=outname, dims=outdimnames, |
||||
coords=outcoords, attrs=outattrs) |
||||
|
||||
|
||||
def _set_1d_meta(wrapped, instance, args, kwargs): |
||||
argvars = from_args(wrapped, ("v_in", "z_in", "z_out", "missingval"), |
||||
*args, **kwargs) |
||||
|
||||
v_in = argvars["v_in"] |
||||
z_in = argvars["z_in"] |
||||
z_out = argvars["z_out"] |
||||
missingval = argvars["missingval"] |
||||
|
||||
result = wrapped(*args, **kwargs) |
||||
|
||||
# Dims are (...,xy,z) |
||||
if isinstance(v_in, DataArray): |
||||
outcoords = OrderedDict() |
||||
outattrs = OrderedDict() |
||||
outdimnames = list(v_in.dims) |
||||
outcoords.update(v_in.coords) |
||||
|
||||
|
||||
outdimnames.remove(v_in.dims[-1]) |
||||
del outcoords[v_in.dims[-1]] |
||||
outdimnames.append("z") |
||||
outname = "{0}_z".format(v_in.name) |
||||
outcoords["z"] = z_out |
||||
|
||||
outattrs.update(v_in.attrs) |
||||
outattrs["_FillValue"] = missingval |
||||
outattrs["missing_value"] = missingval |
||||
|
||||
else: |
||||
outname = "v_in_z" |
||||
|
||||
|
||||
return DataArray(result, name=outname, dims=outdimnames, |
||||
coords=outcoords, attrs=outattrs) |
||||
|
||||
|
||||
|
||||
|
||||
def set_interp_metadata(interp_type): |
||||
@wrapt.decorator |
||||
def func_wrapper(wrapped, instance, args, kwargs): |
||||
if not xarray_enabled(): |
||||
return wrapped(*args, **kwargs) |
||||
|
||||
if interp_type == "horiz": |
||||
return _set_horiz_meta(wrapped, instance, args, kwargs) |
||||
elif interp_type == "cross": |
||||
return _set_cross_meta(wrapped, instance, args, kwargs) |
||||
elif interp_type == "line": |
||||
return _set_line_meta(wrapped, instance, args, kwargs) |
||||
elif interp_type == "vinterp": |
||||
return _set_vinterp_meta(wrapped, instance, args, kwargs) |
||||
elif interp_type == "2dxy": |
||||
return _set_2dxy_meta(wrapped, instance, args, kwargs) |
||||
elif interp_type == "1d": |
||||
return _set_1d_meta(wrapped, instance, args, kwargs) |
||||
return func_wrapper |
||||
@ -0,0 +1,90 @@
@@ -0,0 +1,90 @@
|
||||
from __future__ import (absolute_import, division, print_function, |
||||
unicode_literals) |
||||
|
||||
import numpy as np |
||||
|
||||
import wrapt |
||||
|
||||
from .destag import destagger |
||||
from .util import iter_left_indexes |
||||
|
||||
__all__ = ["uvmet_left_iter"] |
||||
|
||||
# Placed in separate module to resolve a circular dependency with destagger |
||||
# module |
||||
|
||||
def uvmet_left_iter(): |
||||
"""Decorator to handle iterating over leftmost dimensions when using |
||||
multiple files and/or multiple times with the uvmet product. |
||||
|
||||
""" |
||||
@wrapt.decorator |
||||
def func_wrapper(wrapped, instance, args, kwargs): |
||||
u = args[0] |
||||
v = args[1] |
||||
lat = args[2] |
||||
lon = args[3] |
||||
cen_long = args[4] |
||||
cone = args[5] |
||||
|
||||
if u.ndim == lat.ndim: |
||||
num_right_dims = 2 |
||||
is_3d = False |
||||
else: |
||||
num_right_dims = 3 |
||||
is_3d = True |
||||
|
||||
is_stag = False |
||||
if ((u.shape[-1] != lat.shape[-1]) or |
||||
(u.shape[-2] != lat.shape[-2])): |
||||
is_stag = True |
||||
|
||||
if is_3d: |
||||
extra_dim_num = u.ndim - 3 |
||||
else: |
||||
extra_dim_num = u.ndim - 2 |
||||
|
||||
if is_stag: |
||||
u = destagger(u,-1) |
||||
v = destagger(v,-2) |
||||
|
||||
# No special left side iteration, return the function result |
||||
if (extra_dim_num == 0): |
||||
return wrapped(u, v, lat, lon, cen_long, cone) |
||||
|
||||
# Start by getting the left-most 'extra' dims |
||||
outdims = [u.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 |
||||
outdims += [2] # For u/v components |
||||
|
||||
outdims += [u.shape[x] for x in xrange(-num_right_dims,0,1)] |
||||
|
||||
output = np.empty(outdims, u.dtype) |
||||
|
||||
for left_idxs in iter_left_indexes(extra_dims): |
||||
# Make the left indexes plus a single slice object |
||||
# The single slice will handle all the dimensions to |
||||
# the right (e.g. [1,1,:]) |
||||
left_and_slice_idxs = tuple([x for x in left_idxs] + [slice(None)]) |
||||
|
||||
new_u = u[left_and_slice_idxs] |
||||
new_v = v[left_and_slice_idxs] |
||||
new_lat = lat[left_and_slice_idxs] |
||||
new_lon = lon[left_and_slice_idxs] |
||||
|
||||
# Call the numerical routine |
||||
res = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone) |
||||
|
||||
# Note: The 2D version will return a 3D array with a 1 length |
||||
# dimension. Numpy is unable to broadcast this without |
||||
# sqeezing first. |
||||
res = np.squeeze(res) |
||||
|
||||
output[left_and_slice_idxs] = res[:] |
||||
|
||||
return output |
||||
|
||||
return func_wrapper |
||||
|
||||
Loading…
Reference in new issue