@ -6,10 +6,11 @@ from collections import OrderedDict
import numpy as np import numpy as np
import numpy . ma as ma import numpy . ma as ma
from . extension import _interpline
from . util import ( viewkeys , viewitems , extract_vars , from . util import ( viewkeys , viewitems , extract_vars ,
combine_with , either , from_args , arg_location , combine_with , either , from_args , arg_location ,
_is_coord_var , CoordPair , npvalues , py3range , ucode ) is_coordvar , latlon_coordvars , CoordPair , npvalues ,
py3range , ucode , from_var , iter_left_indexes )
from . interputils import get_xy_z_params , get_xy from . interputils import get_xy_z_params , get_xy
from . latlonutils import ij_to_ll , ll_to_ij from . latlonutils import ij_to_ll , ll_to_ij
from . config import xarray_enabled from . config import xarray_enabled
@ -17,9 +18,6 @@ from .config import xarray_enabled
if xarray_enabled ( ) : if xarray_enabled ( ) :
from xarray import DataArray 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 , def copy_and_set_metadata ( copy_varname = None , delete_attrs = None , name = None ,
remove_dims = None , dimnames = None , remove_dims = None , dimnames = None ,
@ -187,24 +185,24 @@ def set_wind_metadata(copy_varname, name, description,
outattrs = OrderedDict ( ) outattrs = OrderedDict ( )
outdimnames = list ( copy_var . dims ) outdimnames = list ( copy_var . dims )
outcoords . update ( copy_var . coords ) #outcoords.update(copy_var.coords )
outattrs . update ( copy_var . attrs ) outattrs . update ( copy_var . attrs )
if wind_ncvar :
pass
if wind_ncvar :
outcoords . update ( copy_var . coords )
elif not wspd_wdir : elif not wspd_wdir :
if not two_d : if not two_d :
outdimnames . insert ( - 3 , " u_v " ) outdimnames . insert ( 0 , " u_v " )
else : else :
outdimnames . insert ( - 2 , " u_v " ) outdimnames . insert ( 0 , " u_v " )
outattrs [ " MemoryOrder " ] = " XY " outattrs [ " MemoryOrder " ] = " XY "
outcoords [ " u_v " ] = [ " u " , " v " ] outcoords [ " u_v " ] = [ " u " , " v " ]
else : else :
if not two_d : if not two_d :
outdimnames . insert ( - 3 , " wspd_wdir " ) outdimnames . insert ( 0 , " wspd_wdir " )
else : else :
outdimnames . insert ( - 2 , " wspd_wdir " ) outdimnames . insert ( 0 , " wspd_wdir " )
outattrs [ " MemoryOrder " ] = " XY " outattrs [ " MemoryOrder " ] = " XY "
outcoords [ " wspd_wdir " ] = [ " wspd " , " wdir " ] outcoords [ " wspd_wdir " ] = [ " wspd " , " wdir " ]
@ -212,6 +210,19 @@ def set_wind_metadata(copy_varname, name, description,
if units is not None : if units is not None :
outattrs [ " units " ] = units outattrs [ " units " ] = units
# xarray doesn't line up coordinate dimensions based on
# names, it just remembers the index it originally mapped to.
# So, need to rebuild the XLAT, XLONG, coordinates again since the
# leftmost index changed.
if not wind_ncvar :
for key , dataarray in viewitems ( copy_var . coords ) :
if is_coordvar ( key ) :
outcoords [ key ] = dataarray . dims , npvalues ( dataarray )
elif key == " XTIME " :
outcoords [ key ] = dataarray . dims , npvalues ( dataarray )
elif key == " Time " :
outcoords [ key ] = npvalues ( dataarray )
outname = name outname = name
outattrs [ " description " ] = description outattrs [ " description " ] = description
@ -342,14 +353,14 @@ def set_height_metadata(geopt=False):
return func_wrapper return func_wrapper
def _set_horiz_meta ( wrapped , instance , args , kwargs ) : def _set_horiz_meta ( wrapped , instance , args , kwargs ) :
argvars = from_args ( wrapped , ( " field3d " , " z " , " desiredloc " , argvars = from_args ( wrapped , ( " field3d " , " z " , " desiredlev " ,
" missingval " ) , " missing " ) ,
* args , * * kwargs ) * args , * * kwargs )
field3d = argvars [ " field3d " ] field3d = argvars [ " field3d " ]
z = argvars [ " z " ] z = argvars [ " z " ]
desiredloc = argvars [ " desiredloc " ] desiredloc = argvars [ " desiredlev " ]
missingval = argvars [ " missingval " ] missingval = argvars [ " missing " ]
result = wrapped ( * args , * * kwargs ) result = wrapped ( * args , * * kwargs )
@ -384,7 +395,7 @@ def _set_horiz_meta(wrapped, instance, args, kwargs):
outname = " {0} _ {1} " . format ( field3d . name , name_levelstr ) outname = " {0} _ {1} " . format ( field3d . name , name_levelstr )
else : else :
outname = " field3d_ {0} " . format ( levelstr ) outname = " field3d_ {0} " . format ( name_ levelstr)
outattrs = OrderedDict ( ) outattrs = OrderedDict ( )
outattrs [ " PlotLevelID " ] = levelstr outattrs [ " PlotLevelID " ] = levelstr
@ -401,7 +412,7 @@ def _set_horiz_meta(wrapped, instance, args, kwargs):
coords = outcoords , attrs = outattrs ) coords = outcoords , attrs = outattrs )
def _set_cross_meta ( wrapped , instance , args , kwargs ) : def _set_cross_meta ( wrapped , instance , args , kwargs ) :
argvars = from_args ( wrapped , ( " field3d " , " z " , " missingval " , argvars = from_args ( wrapped , ( " field3d " , " z " , " include_latlon " , " " ,
" pivot_point " , " angle " , " pivot_point " , " angle " ,
" start_point " , " end_point " , " start_point " , " end_point " ,
" cache " ) , " cache " ) ,
@ -409,7 +420,8 @@ def _set_cross_meta(wrapped, instance, args, kwargs):
field3d = argvars [ " field3d " ] field3d = argvars [ " field3d " ]
z = argvars [ " z " ] z = argvars [ " z " ]
missingval = argvars [ " missingval " ] inc_latlon = argvars [ " include_latlon " ]
missingval = argvars [ " missing " ]
pivot_point = argvars [ " pivot_point " ] pivot_point = argvars [ " pivot_point " ]
angle = argvars [ " angle " ] angle = argvars [ " angle " ]
start_point = argvars [ " start_point " ] start_point = argvars [ " start_point " ]
@ -447,7 +459,7 @@ def _set_cross_meta(wrapped, instance, args, kwargs):
ed_x = xy [ - 1 , 0 ] ed_x = xy [ - 1 , 0 ]
ed_y = xy [ - 1 , 1 ] ed_y = xy [ - 1 , 1 ]
cross_str = " cross-s ection: ({0} , {1} ) to ( {2} , {3} ) " . format ( st_x , st_y , cross_str = " Cross-S ection: ({0} , {1} ) to ( {2} , {3} ) " . format ( st_x , st_y ,
ed_x , ed_y ) ed_x , ed_y )
if angle is not None : if angle is not None :
cross_str + = " ; center= {0} ; angle= {1} " . format ( pivot_point , cross_str + = " ; center= {0} ; angle= {1} " . format ( pivot_point ,
@ -464,7 +476,7 @@ def _set_cross_meta(wrapped, instance, args, kwargs):
# Delete any lat,lon coords # Delete any lat,lon coords
delkeys = [ key for key in viewkeys ( outcoords ) if _ is_coord_ var( key ) ] delkeys = [ key for key in viewkeys ( outcoords ) if is_coordvar ( key ) ]
for key in delkeys : for key in delkeys :
del outcoords [ key ] del outcoords [ key ]
@ -480,8 +492,55 @@ def _set_cross_meta(wrapped, instance, args, kwargs):
except KeyError : except KeyError :
pass pass
outcoords [ " xy_loc " ] = ( " xy " , [ CoordPair ( xy [ i , 0 ] , xy [ i , 1 ] ) # Interpolate to get the lat/lon coords, if desired
for i in py3range ( xy . shape [ - 2 ] ) ] ) if inc_latlon :
latcoordname , loncoordname = latlon_coordvars ( field3d . coords )
if latcoordname is not None and loncoordname is not None :
latcoord = field3d . coords [ latcoordname ]
loncoord = field3d . coords [ loncoordname ]
if latcoord . ndim == 2 :
lats = _interpline ( latcoord , xy )
lons = _interpline ( loncoord , xy )
outcoords [ " xy_loc " ] = ( " xy " ,
np . asarray ( tuple (
CoordPair ( x = xy [ i , 0 ] , y = xy [ i , 1 ] ,
lat = lats [ i ] , lon = lons [ i ] )
for i in py3range ( xy . shape [ - 2 ] ) ) )
)
else :
extra_dims = latcoord . shape [ 0 : - 2 ]
outdims = extra_dims + xy . shape [ - 2 : - 1 ]
latlon_loc = np . empty ( outdims , np . object_ )
for left_dims in iter_left_indexes ( extra_dims ) :
idxs = left_dims + ( slice ( None ) , )
lats = _interpline ( latcoord [ idxs ] , xy )
lons = _interpline ( loncoord [ idxs ] , xy )
latlon_loc [ idxs ] = np . asarray ( tuple (
CoordPair ( x = xy [ i , 0 ] , y = xy [ i , 1 ] ,
lat = lats [ i ] , lon = lons [ i ] )
for i in py3range ( xy . shape [ - 2 ] ) )
) [ : ]
extra_dimnames = latcoord . dims [ 0 : - 2 ]
loc_dimnames = extra_dimnames + ( " xy " , )
outcoords [ " xy_loc " ] = ( loc_dimnames , latlon_loc )
else :
outcoords [ " xy_loc " ] = ( " xy " , np . asarray ( tuple (
CoordPair ( xy [ i , 0 ] , xy [ i , 1 ] )
for i in py3range ( xy . shape [ - 2 ] ) ) ) )
else :
outcoords [ " xy_loc " ] = ( " xy " , np . asarray ( tuple (
CoordPair ( xy [ i , 0 ] , xy [ i , 1 ] )
for i in py3range ( xy . shape [ - 2 ] ) ) ) )
outcoords [ " vertical " ] = z_var2d [ : ] outcoords [ " vertical " ] = z_var2d [ : ]
@ -489,7 +548,7 @@ def _set_cross_meta(wrapped, instance, args, kwargs):
outname = " field3d_cross " outname = " field3d_cross "
outattrs = OrderedDict ( ) outattrs = OrderedDict ( )
outattrs [ " o rientation" ] = cross_str outattrs [ " O rientation" ] = cross_str
outattrs [ " missing_value " ] = missingval outattrs [ " missing_value " ] = missingval
outattrs [ " _FillValue " ] = missingval outattrs [ " _FillValue " ] = missingval
@ -500,7 +559,8 @@ def _set_cross_meta(wrapped, instance, args, kwargs):
def _set_line_meta ( wrapped , instance , args , kwargs ) : def _set_line_meta ( wrapped , instance , args , kwargs ) :
argvars = from_args ( wrapped , ( " field2d " , " pivot_point " , " angle " , argvars = from_args ( wrapped , ( " field2d " , " pivot_point " , " angle " ,
" start_point " , " end_point " , " cache " ) , " start_point " , " end_point " , " include_latlon " ,
" cache " ) ,
* args , * * kwargs ) * args , * * kwargs )
field2d = argvars [ " field2d " ] field2d = argvars [ " field2d " ]
@ -508,6 +568,7 @@ def _set_line_meta(wrapped, instance, args, kwargs):
angle = argvars [ " angle " ] angle = argvars [ " angle " ]
start_point = argvars [ " start_point " ] start_point = argvars [ " start_point " ]
end_point = argvars [ " end_point " ] end_point = argvars [ " end_point " ]
inc_latlon = argvars [ " include_latlon " ]
cache = argvars [ " cache " ] cache = argvars [ " cache " ]
if cache is None : if cache is None :
@ -554,7 +615,7 @@ def _set_line_meta(wrapped, instance, args, kwargs):
del outcoords [ field2d . dims [ i ] ] del outcoords [ field2d . dims [ i ] ]
# Delete any lat,lon coords # Delete any lat,lon coords
delkeys = [ key for key in viewkeys ( outcoords ) if _ is_coord_ var( key ) ] delkeys = [ key for key in viewkeys ( outcoords ) if is_coordvar ( key ) ]
for key in delkeys : for key in delkeys :
del outcoords [ key ] del outcoords [ key ]
@ -569,14 +630,61 @@ def _set_line_meta(wrapped, instance, args, kwargs):
except KeyError : except KeyError :
pass pass
outcoords [ " xy_loc " ] = ( " xy " , [ CoordPair ( xy [ i , 0 ] , xy [ i , 1 ] ) # Interpolate to get the lat/lon coords, if desired
for i in py3range ( xy . shape [ - 2 ] ) ] ) if inc_latlon :
latcoordname , loncoordname = latlon_coordvars ( field2d . coords )
if latcoordname is not None and loncoordname is not None :
latcoord = field2d . coords [ latcoordname ]
loncoord = field2d . coords [ loncoordname ]
if latcoord . ndim == 2 :
lats = _interpline ( latcoord , xy )
lons = _interpline ( loncoord , xy )
outcoords [ " xy_loc " ] = ( " xy " ,
np . asarray ( tuple (
CoordPair ( x = xy [ i , 0 ] , y = xy [ i , 1 ] ,
lat = lats [ i ] , lon = lons [ i ] )
for i in py3range ( xy . shape [ - 2 ] ) ) )
)
else :
extra_dims = latcoord . shape [ 0 : - 2 ]
outdims = extra_dims + xy . shape [ - 2 : - 1 ]
latlon_loc = np . empty ( outdims , np . object_ )
for left_dims in iter_left_indexes ( extra_dims ) :
idxs = left_dims + ( slice ( None ) , )
lats = _interpline ( latcoord [ idxs ] , xy )
lons = _interpline ( loncoord [ idxs ] , xy )
latlon_loc [ idxs ] = np . asarray ( tuple (
CoordPair ( x = xy [ i , 0 ] , y = xy [ i , 1 ] ,
lat = lats [ i ] , lon = lons [ i ] )
for i in py3range ( xy . shape [ - 2 ] ) )
) [ : ]
extra_dimnames = latcoord . dims [ 0 : - 2 ]
loc_dimnames = extra_dimnames + ( " xy " , )
outcoords [ " xy_loc " ] = ( loc_dimnames , latlon_loc )
else :
outcoords [ " xy_loc " ] = ( " xy " , np . asarray ( tuple (
CoordPair ( xy [ i , 0 ] , xy [ i , 1 ] )
for i in py3range ( xy . shape [ - 2 ] ) ) ) )
else :
outcoords [ " xy_loc " ] = ( " xy " , np . asarray ( tuple (
CoordPair ( xy [ i , 0 ] , xy [ i , 1 ] )
for i in py3range ( xy . shape [ - 2 ] ) ) ) )
else : else :
outname = " field2d_line " outname = " field2d_line "
outattrs = OrderedDict ( ) outattrs = OrderedDict ( )
outattrs [ " orientation " ] = cross_str outattrs [ " O rientation" ] = cross_str
return DataArray ( result , name = outname , dims = outdimnames , return DataArray ( result , name = outname , dims = outdimnames ,
coords = outcoords , attrs = outattrs ) coords = outcoords , attrs = outattrs )
@ -634,6 +742,7 @@ def _set_2dxy_meta(wrapped, instance, args, kwargs):
field3d = argvars [ " field3d " ] field3d = argvars [ " field3d " ]
xy = argvars [ " xy " ] xy = argvars [ " xy " ]
xy = npvalues ( xy )
result = wrapped ( * args , * * kwargs ) result = wrapped ( * args , * * kwargs )
@ -652,14 +761,30 @@ def _set_2dxy_meta(wrapped, instance, args, kwargs):
outattrs = OrderedDict ( ) outattrs = OrderedDict ( )
outdimnames = list ( field3d . dims ) outdimnames = list ( field3d . dims )
outcoords . update ( field3d . coords ) outcoords . update ( field3d . coords )
for i in py3range ( - 2 , 0 , 1 ) : for i in py3range ( - 2 , 0 , 1 ) :
outdimnames . remove ( field3d . dims [ i ] )
del outcoords [ field3d . dims [ i ] ] del outcoords [ field3d . dims [ i ] ]
outdimnames . remove ( field3d . dims [ i ] )
outdimnames [ - 2 ] = " xy " # Need to remove XLAT, XLONG...
outattrs . update ( field3d . attrs ) delkeys = ( key for key , arr in viewitems ( field3d . coords )
if arr . ndim > 1 )
for key in delkeys :
del outcoords [ key ]
outdimnames . append ( " xy " )
#outattrs.update(field3d.attrs)
desc = field3d . attrs . get ( " description " , None )
if desc is not None :
outattrs [ " description " ] = desc
units = field3d . attrs . get ( " units " , None )
if units is not None :
outattrs [ " units " ] = units
outname = " {0} _xy " . format ( field3d . name ) outname = " {0} _2d xy " . format ( field3d . name )
outcoords [ " xy_loc " ] = ( " xy " , [ CoordPair ( xy [ i , 0 ] , xy [ i , 1 ] ) outcoords [ " xy_loc " ] = ( " xy " , [ CoordPair ( xy [ i , 0 ] , xy [ i , 1 ] )
for i in py3range ( xy . shape [ - 2 ] ) ] ) for i in py3range ( xy . shape [ - 2 ] ) ] )
@ -671,7 +796,7 @@ def _set_2dxy_meta(wrapped, instance, args, kwargs):
pass pass
else : else :
outname = " field3d_xy " outname = " field3d_2d xy "
outattrs [ " Orientation " ] = cross_str outattrs [ " Orientation " ] = cross_str
@ -695,19 +820,32 @@ def _set_1d_meta(wrapped, instance, args, kwargs):
outcoords = OrderedDict ( ) outcoords = OrderedDict ( )
outattrs = OrderedDict ( ) outattrs = OrderedDict ( )
outdimnames = list ( v_in . dims ) outdimnames = list ( v_in . dims )
outcoords . update ( v_in . coords ) #outcoords.update(v_in.coords)
outdimnames . pop ( - 1 )
for name in outdimnames :
try :
outcoords [ name ] = v_in . coords [ name ]
except KeyError :
continue
outdimnames . remove ( v_in . dims [ - 1 ] )
del outcoords [ v_in . dims [ - 1 ] ]
outdimnames . append ( " z " ) outdimnames . append ( " z " )
outname = " {0} _z " . format ( v_in . name ) outname = " {0} _z " . format ( v_in . name )
outcoords [ " z " ] = z_out outcoords [ " z " ] = z_out
outattrs . update ( v_in . attrs ) #outattrs.update(v_in.attrs )
outattrs [ " _FillValue " ] = missingval outattrs [ " _FillValue " ] = missingval
outattrs [ " missing_value " ] = missingval outattrs [ " missing_value " ] = missingval
desc = v_in . attrs . get ( " description " , None )
if desc is not None :
outattrs [ " description " ] = desc
units = v_in . attrs . get ( " units " , None )
if units is not None :
outattrs [ " units " ] = units
else : else :
outname = " v_in_z " outname = " v_in_z "
@ -716,6 +854,42 @@ def _set_1d_meta(wrapped, instance, args, kwargs):
coords = outcoords , attrs = outattrs ) coords = outcoords , attrs = outattrs )
def _set_xy_meta ( wrapped , instance , args , kwargs ) :
argvars = from_args ( wrapped , ( " field " , " pivot_point " , " angle " ,
" start_point " , " end_point " ) ,
* args , * * kwargs )
field = argvars [ " field " ]
pivot_point = argvars [ " pivot_point " ]
angle = argvars [ " angle " ]
start_point = argvars [ " start_point " ]
end_point = argvars [ " end_point " ]
result = wrapped ( * args , * * kwargs )
if isinstance ( field , DataArray ) :
outname = " {0} _xy " . format ( field . name )
else :
outname = " xy "
outdimnames = [ " idx " , " x_y " ]
outcoords = OrderedDict ( )
outattrs = OrderedDict ( )
outcoords [ " x_y " ] = [ " x " , " y " ]
if pivot_point is not None and angle is not None :
outattrs [ " pivot_point " ] = pivot_point
outattrs [ " angle " ] = angle
if start_point is not None and end_point is not None :
outattrs [ " start_point " ] = start_point
outattrs [ " end_point " ] = end_point
return DataArray ( result , name = outname , dims = outdimnames ,
coords = outcoords , attrs = outattrs )
def set_interp_metadata ( interp_type ) : def set_interp_metadata ( interp_type ) :
@wrapt . decorator @wrapt . decorator
def func_wrapper ( wrapped , instance , args , kwargs ) : def func_wrapper ( wrapped , instance , args , kwargs ) :
@ -739,5 +913,222 @@ def set_interp_metadata(interp_type):
return _set_2dxy_meta ( wrapped , instance , args , kwargs ) return _set_2dxy_meta ( wrapped , instance , args , kwargs )
elif interp_type == " 1d " : elif interp_type == " 1d " :
return _set_1d_meta ( wrapped , instance , args , kwargs ) return _set_1d_meta ( wrapped , instance , args , kwargs )
elif interp_type == " xy " :
return _set_xy_meta ( wrapped , instance , args , kwargs )
return func_wrapper
def set_alg_metadata ( alg_ndims , right_dimnames = None ,
refvarndims = None ,
refvarname = None , missingarg = None ,
insert_dimnames = None ,
units = None , description = None , squeeze = False ) :
@wrapt . decorator
def func_wrapper ( wrapped , instance , args , kwargs ) :
do_meta = from_args ( wrapped , ( " meta " , ) , * args , * * kwargs ) [ " meta " ]
if do_meta is None :
do_meta = True
if not xarray_enabled ( ) or not do_meta :
return wrapped ( * args , * * kwargs )
result = wrapped ( * args , * * kwargs )
# Default dimension names
outdims = [ " dim_ {} " . format ( i ) for i in py3range ( result . ndim ) ]
if missingarg is not None :
missingval = from_args ( wrapped , ( missingarg , ) ,
* args , * * kwargs ) [ missingarg ]
else :
missingval = None
if missingval is not None :
outattrs [ " _FillValue " ] = missingval
outattrs [ " missing_value " ] = missingval
result = np . ma . masked_values ( result , missingval )
outname = wrapped . __name__
outattrs = OrderedDict ( )
if units is not None :
if isinstance ( description , from_var ) :
_units = units ( wrapped , * args , * * kwargs )
if uts is not None :
outattrs [ " units " ] = _units
else :
outattrs [ " units " ] = units
if description is not None :
if isinstance ( description , from_var ) :
desc = description ( wrapped , * args , * * kwargs )
if desc is not None :
outattrs [ " description " ] = desc
else :
outattrs [ " description " ] = description
# Copy the dimnames from the reference variable, otherwise, use
# the supplied dimnames
if refvarname is not None :
refvar = from_args ( wrapped , ( refvarname , ) ,
* args , * * kwargs ) [ refvarname ]
else :
refvar = None
if isinstance ( refvar , DataArray ) :
# If right dims are provided, use them first
if right_dimnames is not None :
outdims [ - alg_ndims : ] = right_dimnames [ - alg_ndims : ]
else :
# Copy the right dims
outdims [ - alg_ndims : ] = refvar . dims [ - alg_ndims : ]
# Left dims
if refvarndims is None :
# Used when result and reference are aligned on right
if result . ndim > alg_ndims :
result_extra = result . ndim - alg_ndims
for i in py3range ( 1 , result_extra + 1 ) :
idx = - alg_ndims - i
if - idx < = refvar . ndim :
outdims [ idx ] = refvar . dims [ idx ]
else :
continue
# When reference and result aren't exactly aligned (slp,uvmet)
else :
ref_extra = refvar . ndim - refvarndims
ref_left_dimnames = refvar . dims [ 0 : ref_extra ]
for i , dimname in enumerate ( ref_left_dimnames [ : : - 1 ] , 1 ) :
idx = - i
if - idx < = result . shape :
outdims [ idx ] = dimname
else :
continute
if insert_dimnames is not None :
for pair in insert_dimnames :
outdims . insert ( pair [ 0 ] , pair [ 1 ] )
out = DataArray ( result , name = outname , dims = outdims ,
attrs = outattrs )
if squeeze :
return out . squeeze ( )
return out
return func_wrapper
def set_uvmet_alg_metadata ( units = " mps " , description = " earth rotated u,v " ) :
@wrapt . decorator
def func_wrapper ( wrapped , instance , args , kwargs ) :
do_meta = from_args ( wrapped , ( " meta " , ) , * args , * * kwargs ) [ " meta " ]
if do_meta is None :
do_meta = True
if not xarray_enabled ( ) or not do_meta :
return wrapped ( * args , * * kwargs )
result = wrapped ( * args , * * kwargs )
# Default dimension names
outdims = [ " dim_ {} " . format ( i ) for i in py3range ( result . ndim ) ]
outname = " uvmet "
outattrs = OrderedDict ( )
if units is not None :
outattrs [ " units " ] = units
if description is not None :
outattrs [ " description " ] = description
latvar = from_args ( wrapped , ( " lat " , ) , * args , * * kwargs ) [ " lat " ]
uvar = from_args ( wrapped , ( " u " , ) , * args , * * kwargs ) [ " u " ]
if isinstance ( uvar , DataArray ) :
# Right dims come from latvar
outdims [ - 2 : ] = latvar . dims [ - 2 : ]
# Left dims come from u-var
outdims [ 1 : - 2 ] = uvar . dims [ 0 : - 2 ]
# Left-most is always u_v
outdims [ 0 ] = " u_v "
out = DataArray ( result , name = outname , dims = outdims ,
attrs = outattrs )
return out
return func_wrapper return func_wrapper
def set_destag_metadata ( ) :
@wrapt . decorator
def func_wrapper ( wrapped , instance , args , kwargs ) :
do_meta = from_args ( wrapped , ( " meta " , ) , * args , * * kwargs ) [ " meta " ]
if do_meta is None :
do_meta = True
if not xarray_enabled ( ) or not do_meta :
return wrapped ( * args , * * kwargs )
result = wrapped ( * args , * * kwargs )
# Default dimension names
outdims = [ " dim {} " . format ( i ) for i in py3range ( result . ndim ) ]
destag_args = from_args ( wrapped , ( " var " , " stagger_dim " ) ,
* args , * * kwargs )
var = destag_args [ " var " ]
destag_dim = destag_args [ " stagger_dim " ]
if isinstance ( var , DataArray ) :
if var . name is not None :
outname = " destag_ {} " . format ( var . name )
else :
outnames = " destag_var "
outattrs = OrderedDict ( )
outattrs . update ( var . attrs )
outattrs [ " destag_dim " ] = destag_dim
outdims = [ ]
outdims + = var . dims
destag_dim_name = outdims [ destag_dim ]
if destag_dim_name . find ( " _stag " ) > = 0 :
new_dim_name = destag_dim_name . replace ( " _stag " , " " )
else :
if destag_dim > = 0 :
new_dim_name = " dim_ {} " . format ( destag_dim )
else :
dim_num = result . ndim + destag_dim
new_dim_name = " dim_ {} " . format ( dim_num )
outdims [ destag_dim ] = new_dim_name
out = DataArray ( result , name = outname , dims = outdims , attrs = outattrs )
return out
return func_wrapper
Bill Ladwig
9 years ago