Bill Ladwig 6 years ago
parent
commit
f3068c2b69
  1. 71
      src/wrf/geobnds.py
  2. 940
      src/wrf/interp.py
  3. 454
      src/wrf/interputils.py
  4. 547
      src/wrf/latlonutils.py
  5. 1841
      src/wrf/metadecorators.py
  6. 988
      src/wrf/projection.py
  7. 13
      src/wrf/projutils.py
  8. 118
      src/wrf/py3compat.py
  9. 489
      src/wrf/routines.py
  10. 463
      src/wrf/specialdec.py
  11. 440
      src/wrf/units.py
  12. 3451
      src/wrf/util.py
  13. 1
      src/wrf/version.py

71
src/wrf/geobnds.py

@ -2,46 +2,47 @@ from __future__ import (absolute_import, division, print_function)
from .coordpair import CoordPair from .coordpair import CoordPair
class GeoBounds(object): class GeoBounds(object):
"""A class that stores the geographic boundaries. """A class that stores the geographic boundaries.
Currently, only corner points are used, specified as the bottom left and Currently, only corner points are used, specified as the bottom left and
top right corners. Users can specify the corner points directly, or top right corners. Users can specify the corner points directly, or
specify two-dimensional latitude and longitude arrays and the corner points specify two-dimensional latitude and longitude arrays and the corner points
will be extracted from them. will be extracted from them.
Attributes: Attributes:
bottom_left (:class:`wrf.CoordPair`): The bottom left coordinate. bottom_left (:class:`wrf.CoordPair`): The bottom left coordinate.
top_right (:class:`wrf.CoordPair`): The top right coordinate. top_right (:class:`wrf.CoordPair`): The top right coordinate.
""" """
def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None): def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None):
""" Initialize a :class:`wrf.GeoBounds` object. """ Initialize a :class:`wrf.GeoBounds` object.
Args: Args:
bottom_left (:class:`wrf.CoordPair`, optional): The lower left bottom_left (:class:`wrf.CoordPair`, optional): The lower left
corner. Must also specify *top_right* if used. corner. Must also specify *top_right* if used.
Default is None. Default is None.
top_right (:class:`wrf.CoordPair`, optional): The upper right top_right (:class:`wrf.CoordPair`, optional): The upper right
corner. Must also specify *bottom_left* if used. corner. Must also specify *bottom_left* if used.
Default is None. Default is None.
lats (:class:`numpy.ndarray`, optional): An array of at least lats (:class:`numpy.ndarray`, optional): An array of at least
two dimensions containing all of the latitude values. Must two dimensions containing all of the latitude values. Must
also specify *lons* if used. Default is None. also specify *lons* if used. Default is None.
lons (:class:`numpy.ndarray`, optional): An array of at least lons (:class:`numpy.ndarray`, optional): An array of at least
two dimensions containing all of the longitude values. Must two dimensions containing all of the longitude values. Must
also specify *lats* if used. Default is None. also specify *lats* if used. Default is None.
""" """
if bottom_left is not None and top_right is not None: if bottom_left is not None and top_right is not None:
self.bottom_left = bottom_left self.bottom_left = bottom_left
self.top_right = top_right self.top_right = top_right
# Make sure the users set lat/lon coordinates # Make sure the users set lat/lon coordinates
if self.bottom_left.lat is None: if self.bottom_left.lat is None:
raise ValueError("'bottom_left' parameter does not contain a " raise ValueError("'bottom_left' parameter does not contain a "
@ -56,35 +57,31 @@ class GeoBounds(object):
raise ValueError("'top_right' parameter does not contain a" raise ValueError("'top_right' parameter does not contain a"
"'lon' attribute") "'lon' attribute")
elif lats is not None and lons is not None: elif lats is not None and lons is not None:
self.bottom_left = CoordPair(lat=lats[0,0], lon=lons[0,0]) self.bottom_left = CoordPair(lat=lats[0, 0], lon=lons[0, 0])
self.top_right = CoordPair(lat=lats[-1,-1], lon=lons[-1,-1]) self.top_right = CoordPair(lat=lats[-1, -1], lon=lons[-1, -1])
else: else:
raise ValueError("must specify either 'bottom_top' and " raise ValueError("must specify either 'bottom_top' and "
"'top_right' parameters " "'top_right' parameters "
"or 'lats' and 'lons' parameters") "or 'lats' and 'lons' parameters")
def __repr__(self): def __repr__(self):
argstr = "{}, {}".format(repr(self.bottom_left), argstr = "{}, {}".format(repr(self.bottom_left),
repr(self.top_right)) repr(self.top_right))
return "{}({})".format(self.__class__.__name__, argstr) return "{}({})".format(self.__class__.__name__, argstr)
class NullGeoBounds(GeoBounds): class NullGeoBounds(GeoBounds):
"""An emtpy :class:`wrf.GeoBounds` subclass. """An emtpy :class:`wrf.GeoBounds` subclass.
This is used for initializing arrays of :class:`wrf.GeoBounds`, in This is used for initializing arrays of :class:`wrf.GeoBounds`, in
particular when working with moving domains and variables combined with the particular when working with moving domains and variables combined with the
'join' method. 'join' method.
""" """
def __init__(self): def __init__(self):
""" Initialize a :class:`wrf.NullGeoBounds` object.""" """ Initialize a :class:`wrf.NullGeoBounds` object."""
pass pass
def __repr__(self): def __repr__(self):
return "{}()".format(self.__class__.__name__) return "{}()".format(self.__class__.__name__)

940
src/wrf/interp.py

File diff suppressed because it is too large Load Diff

454
src/wrf/interputils.py

@ -14,165 +14,164 @@ from .util import pairs_to_latlon
def to_positive_idxs(shape, coord): def to_positive_idxs(shape, coord):
"""Return the positive index values. """Return the positive index values.
This function converts negative index values to positive index values. This function converts negative index values to positive index values.
Args: Args:
shape (indexable sequence): The array shape. shape (indexable sequence): The array shape.
coord (indexable sequence): The coordinate pair for x and y. coord (indexable sequence): The coordinate pair for x and y.
Returns: Returns:
:obj:`list`: The coordinate values with all positive indexes. :obj:`list`: The coordinate values with all positive indexes.
""" """
if (coord[-2] >= 0 and coord[-1] >= 0): if (coord[-2] >= 0 and coord[-1] >= 0):
return coord return coord
return [x if (x >= 0) else shape[-i-1]+x for (i,x) in enumerate(coord)]
return [x if (x >= 0) else shape[-i-1]+x for (i, x) in enumerate(coord)]
def _calc_xy(xdim, ydim, pivot_point=None, angle=None,
start_point=None, end_point=None): def _calc_xy(xdim, ydim, pivot_point=None, angle=None,
start_point=None, end_point=None):
"""Return the x,y points for the horizontal cross section line. """Return the x,y points for the horizontal cross section line.
Args: Args:
xdim (:obj:`int`): The x-dimension size. xdim (:obj:`int`): The x-dimension size.
ydim (:obj:`int`): The y-dimension size. ydim (:obj:`int`): The y-dimension size.
pivot_point (:obj:`tuple` or :obj:`list`, optional): A pivot_point (:obj:`tuple` or :obj:`list`, optional): A
:obj:`tuple` or :obj:`list` with two entries, :obj:`tuple` or :obj:`list` with two entries,
in the form of [x, y] (or [west_east, south_north]), which in the form of [x, y] (or [west_east, south_north]), which
indicates the x,y location through which the plane will pass. indicates the x,y location through which the plane will pass.
Must also specify `angle`. Must also specify `angle`.
angle (:obj:`float`, optional): Only valid for cross sections where angle (:obj:`float`, optional): Only valid for cross sections where
a plane will be plotted through a plane will be plotted through
a given point on the model domain. 0.0 represents a S-N cross a given point on the model domain. 0.0 represents a S-N cross
section. 90.0 is a W-E cross section. section. 90.0 is a W-E cross section.
start_point (:obj:`tuple` or :obj:`list`, optional): A start_point (:obj:`tuple` or :obj:`list`, optional): A
:obj:`tuple` or :obj:`list` with two entries, in the form of :obj:`tuple` or :obj:`list` with two entries, in the form of
[x, y] (or [west_east, south_north]), which indicates the start [x, y] (or [west_east, south_north]), which indicates the start
x,y location through which the plane will pass. x,y location through which the plane will pass.
end_point (:obj:`tuple` or :obj:`list`, optional): A end_point (:obj:`tuple` or :obj:`list`, optional): A
:obj:`tuple` or :obj:`list` with two entries, in the form of :obj:`tuple` or :obj:`list` with two entries, in the form of
[x, y] (or [west_east, south_north]), which indicates the end x,y [x, y] (or [west_east, south_north]), which indicates the end x,y
location through which the plane will pass. location through which the plane will pass.
Returns: Returns:
:class:`np.ndarray`: A two-dimensional array with the left index :class:`np.ndarray`: A two-dimensional array with the left index
representing each point along the line, and the rightmost dimension representing each point along the line, and the rightmost dimension
having two values for the x and y coordinates [0=X, 1=Y]. having two values for the x and y coordinates [0=X, 1=Y].
""" """
# 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[-2] xp = pivot_point[-2]
yp = pivot_point[-1] yp = pivot_point[-1]
if xp >= xdim or yp >= ydim: if xp >= xdim or yp >= ydim:
raise ValueError("pivot point {} is outside of domain " raise ValueError("pivot point {} is outside of domain "
"with shape {}".format(pivot_point, "with shape {}".format(pivot_point,
(xdim, ydim))) (xdim, ydim)))
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))):
#x = y*slope + intercept
slope = -(360.-angle)/45. slope = -(360.-angle)/45.
if( angle < 45. ): if(angle < 45.):
slope = angle/45. slope = angle/45.
if( angle > 135.): if(angle > 135.):
slope = (angle-180.)/45. slope = (angle-180.)/45.
intercept = xp - yp*slope intercept = xp - yp*slope
# find intersections with domain boundaries # find intersections with domain boundaries
y0 = 0. y0 = 0.
x0 = y0*slope + intercept x0 = y0*slope + intercept
if( x0 < 0.): # intersect outside of left boundary if(x0 < 0.): # intersect outside of left boundary
x0 = 0. x0 = 0.
y0 = (x0 - intercept)/slope y0 = (x0 - intercept)/slope
if( x0 > xdim-1): #intersect outside of right boundary if(x0 > xdim-1): # intersect outside of right boundary
x0 = xdim-1 x0 = xdim-1
y0 = (x0 - intercept)/slope y0 = (x0 - intercept)/slope
y1 = ydim-1. #need to make sure this will be a float? y1 = ydim-1. # need to make sure this will be a float?
x1 = y1*slope + intercept x1 = y1*slope + intercept
if( x1 < 0.): # intersect outside of left boundary if(x1 < 0.): # intersect outside of left boundary
x1 = 0. x1 = 0.
y1 = (x1 - intercept)/slope y1 = (x1 - intercept)/slope
if( x1 > xdim-1): # intersect outside of right boundary if(x1 > xdim-1): # intersect outside of right boundary
x1 = xdim-1 x1 = xdim-1
y1 = (x1 - intercept)/slope y1 = (x1 - intercept)/slope
else: else:
# y = x*slope + intercept # y = x*slope + intercept
slope = (90.-angle)/45. slope = (90.-angle)/45.
if( angle > 225. ): if (angle > 225.):
slope = (270.-angle)/45. slope = (270.-angle)/45.
intercept = yp - xp*slope intercept = yp - xp*slope
#find intersections with domain boundaries # Find intersections with domain boundaries
x0 = 0. x0 = 0.
y0 = x0*slope + intercept y0 = x0*slope + intercept
if( y0 < 0.): # intersect outside of bottom boundary if (y0 < 0.): # intersect outside of bottom boundary
y0 = 0. y0 = 0.
x0 = (y0 - intercept)/slope x0 = (y0 - intercept)/slope
if( y0 > ydim-1): # intersect outside of top boundary if (y0 > ydim-1): # intersect outside of top boundary
y0 = ydim-1 y0 = ydim-1
x0 = (y0 - intercept)/slope x0 = (y0 - intercept)/slope
x1 = xdim-1. # need to make sure this will be a float? x1 = xdim-1. # need to make sure this will be a float?
y1 = x1*slope + intercept y1 = x1*slope + intercept
if( y1 < 0.): # intersect outside of bottom boundary if (y1 < 0.): # intersect outside of bottom boundary
y1 = 0. y1 = 0.
x1 = (y1 - intercept)/slope x1 = (y1 - intercept)/slope
if( y1 > ydim-1):# intersect outside of top boundary if (y1 > ydim-1): # intersect outside of top boundary
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[-2] x0 = start_point[-2]
y0 = start_point[-1] y0 = start_point[-1]
x1 = end_point[-2] x1 = end_point[-2]
y1 = end_point[-1] y1 = end_point[-1]
if x0 >= xdim or y0 >= ydim: if x0 >= xdim or y0 >= ydim:
raise ValueError("start_point {} is outside of domain " raise ValueError("start_point {} is outside of domain "
"with shape {}".format(start_point, (xdim, ydim))) "with shape {}".format(start_point, (xdim, ydim)))
if x1 >= xdim or y1 >= ydim: if x1 >= xdim or y1 >= ydim:
raise ValueError("end_point {} is outside of domain " raise ValueError("end_point {} is outside of domain "
"with shape {}".format(end_point, (xdim, ydim))) "with shape {}".format(end_point, (xdim, ydim)))
else: else:
raise ValueError("invalid start/end or pivot/angle arguments") raise ValueError("invalid start/end or pivot/angle arguments")
dx = x1 - x0 dx = x1 - x0
dy = y1 - y0 dy = y1 - y0
distance = (dx*dx + dy*dy)**0.5 distance = (dx*dx + dy*dy)**0.5
npts = int(distance) + 1 npts = int(distance) + 1
xy = np.zeros((npts,2), "float") xy = np.zeros((npts, 2), "float")
dx = dx/(npts-1) dx = dx/(npts-1)
dy = dy/(npts-1) dy = dy/(npts-1)
for i in py3range(npts): for i in py3range(npts):
xy[i,0] = x0 + i*dx xy[i, 0] = x0 + i*dx
xy[i,1] = y0 + i*dy xy[i, 1] = y0 + i*dy
return xy return xy
@ -180,66 +179,66 @@ def get_xy_z_params(z, pivot_point=None, angle=None,
start_point=None, end_point=None, start_point=None, end_point=None,
levels=None, autolevels=100): levels=None, autolevels=100):
"""Return the cross section parameters. """Return the cross section parameters.
This function returns the xy horizontal cross section line coordinates, This function returns the xy horizontal cross section line coordinates,
the xy x z vertical values interpolated along the xy cross section the xy x z vertical values interpolated along the xy cross section
line, and the fixed vertical levels to be used by the cross section line, and the fixed vertical levels to be used by the cross section
algorithm (at ~1% increments for the minimum to maximum vertical algorithm (at ~1% increments for the minimum to maximum vertical
span). span).
Args: Args:
z (:class:`numpy.ndarray`): The vertical coordinate, whose rightmost z (:class:`numpy.ndarray`): The vertical coordinate, whose rightmost
dimensions are bottom_top x south_north x west_east. dimensions are bottom_top x south_north x west_east.
pivot_point (:obj:`tuple` or :obj:`list`, optional): A pivot_point (:obj:`tuple` or :obj:`list`, optional): A
:obj:`tuple` or :obj:`list` with two entries, :obj:`tuple` or :obj:`list` with two entries,
in the form of [x, y] (or [west_east, south_north]), which in the form of [x, y] (or [west_east, south_north]), which
indicates the x,y location through which the plane will pass. indicates the x,y location through which the plane will pass.
Must also specify `angle`. Must also specify `angle`.
angle (:obj:`float`, optional): Only valid for cross sections where angle (:obj:`float`, optional): Only valid for cross sections where
a plane will be plotted through a plane will be plotted through
a given point on the model domain. 0.0 represents a S-N cross a given point on the model domain. 0.0 represents a S-N cross
section. 90.0 is a W-E cross section. section. 90.0 is a W-E cross section.
start_point (:obj:`tuple` or :obj:`list`, optional): A start_point (:obj:`tuple` or :obj:`list`, optional): A
:obj:`tuple` or :obj:`list` with two entries, in the form of :obj:`tuple` or :obj:`list` with two entries, in the form of
[x, y] (or [west_east, south_north]), which indicates the start [x, y] (or [west_east, south_north]), which indicates the start
x,y location through which the plane will pass. x,y location through which the plane will pass.
end_point (:obj:`tuple` or :obj:`list`, optional): A end_point (:obj:`tuple` or :obj:`list`, optional): A
:obj:`tuple` or :obj:`list` with two entries, in the form of :obj:`tuple` or :obj:`list` with two entries, in the form of
[x, y] (or [west_east, south_north]), which indicates the end x,y [x, y] (or [west_east, south_north]), which indicates the end x,y
location through which the plane will pass. location through which the plane will pass.
levels (sequence): A sequence of :obj:`float` for the desired levels (sequence): A sequence of :obj:`float` for the desired
vertical levels in the output array. If None, a fixed set of vertical levels in the output array. If None, a fixed set of
vertical levels is provided. Default is None. vertical levels is provided. Default is None.
autolevels(:obj:`int`, optional): The number of evenly spaced autolevels(:obj:`int`, optional): The number of evenly spaced
automatically chosen vertical levels to use when *levels* automatically chosen vertical levels to use when *levels*
is None. Default is 100. is None. Default is 100.
Returns: Returns:
:obj:`tuple`: A tuple containing the xy horizontal cross section :obj:`tuple`: A tuple containing the xy horizontal cross section
coordinates, the vertical values interpolated along the xy cross coordinates, the vertical values interpolated along the xy cross
section line, and the fixed vertical levels used by the section line, and the fixed vertical levels used by the
cross section algorithm at ~1% increments of minimum to maximum cross section algorithm at ~1% increments of minimum to maximum
vertical span. vertical span.
""" """
xy = get_xy(z, pivot_point, angle, start_point, end_point) xy = get_xy(z, pivot_point, angle, start_point, end_point)
# Interp z # Interp z
var2dz = _interp2dxy(z, xy) var2dz = _interp2dxy(z, xy)
extra_dim_num = z.ndim - 3 extra_dim_num = z.ndim - 3
idx1 = tuple([0]*extra_dim_num + [0,0]) idx1 = tuple([0]*extra_dim_num + [0, 0])
idx2 = tuple([0]*extra_dim_num + [-1,0]) idx2 = tuple([0]*extra_dim_num + [-1, 0])
if levels is None: if levels is None:
# interp to constant z grid # interp to constant z grid
if(var2dz[idx1] > var2dz[idx2]): # monotonically decreasing coordinate if(var2dz[idx1] > var2dz[idx2]): # monotonically decreasing coordinate
@ -255,181 +254,182 @@ def get_xy_z_params(z, pivot_point=None, angle=None,
dz = (1.0/autolevels)*z_max dz = (1.0/autolevels)*z_max
z_var2d = np.zeros((autolevels), dtype=z.dtype) z_var2d = np.zeros((autolevels), dtype=z.dtype)
z_var2d[0] = z_min z_var2d[0] = z_min
for i in py3range(1,autolevels): for i in py3range(1, autolevels):
z_var2d[i] = z_var2d[0] + i*dz z_var2d[i] = z_var2d[0] + i*dz
else: else:
z_var2d = np.asarray(levels, z.dtype) z_var2d = np.asarray(levels, z.dtype)
return xy, var2dz, z_var2d return xy, var2dz, z_var2d
def get_xy(var, pivot_point=None, angle=None, def get_xy(var, pivot_point=None, angle=None,
start_point=None, end_point=None): start_point=None, end_point=None):
"""Return the x,y points for the horizontal cross section line. """Return the x,y points for the horizontal cross section line.
Args: Args:
var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A variable var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A variable
that contains a :attr:`shape` attribute. that contains a :attr:`shape` attribute.
pivot_point (:obj:`tuple` or :obj:`list`, optional): A pivot_point (:obj:`tuple` or :obj:`list`, optional): A
:obj:`tuple` or :obj:`list` with two entries, :obj:`tuple` or :obj:`list` with two entries,
in the form of [x, y] (or [west_east, south_north]), which in the form of [x, y] (or [west_east, south_north]), which
indicates the x,y location through which the plane will pass. indicates the x,y location through which the plane will pass.
Must also specify `angle`. Must also specify `angle`.
angle (:obj:`float`, optional): Only valid for cross sections where angle (:obj:`float`, optional): Only valid for cross sections where
a plane will be plotted through a plane will be plotted through
a given point on the model domain. 0.0 represents a S-N cross a given point on the model domain. 0.0 represents a S-N cross
section. 90.0 is a W-E cross section. section. 90.0 is a W-E cross section.
start_point (:obj:`tuple` or :obj:`list`, optional): A start_point (:obj:`tuple` or :obj:`list`, optional): A
:obj:`tuple` or :obj:`list` with two entries, in the form of :obj:`tuple` or :obj:`list` with two entries, in the form of
[x, y] (or [west_east, south_north]), which indicates the start [x, y] (or [west_east, south_north]), which indicates the start
x,y location through which the plane will pass. x,y location through which the plane will pass.
end_point (:obj:`tuple` or :obj:`list`, optional): A end_point (:obj:`tuple` or :obj:`list`, optional): A
:obj:`tuple` or :obj:`list` with two entries, in the form of :obj:`tuple` or :obj:`list` with two entries, in the form of
[x, y] (or [west_east, south_north]), which indicates the end x,y [x, y] (or [west_east, south_north]), which indicates the end x,y
location through which the plane will pass. location through which the plane will pass.
Returns: Returns:
:class:`np.ndarray`: A two-dimensional array with the left index :class:`np.ndarray`: A two-dimensional array with the left index
representing each point along the line, and the rightmost dimension representing each point along the line, and the rightmost dimension
having two values for the x and y coordinates [0=X, 1=Y]. having two values for the x and y coordinates [0=X, 1=Y].
""" """
if pivot_point is not None: if pivot_point is not None:
pos_pivot = to_positive_idxs(var.shape[-2:], pivot_point) pos_pivot = to_positive_idxs(var.shape[-2:], pivot_point)
else: else:
pos_pivot = pivot_point pos_pivot = pivot_point
if start_point is not None: if start_point is not None:
pos_start = to_positive_idxs(var.shape[-2:], start_point) pos_start = to_positive_idxs(var.shape[-2:], start_point)
else: else:
pos_start = start_point pos_start = start_point
if end_point is not None: if end_point is not None:
pos_end = to_positive_idxs(var.shape[-2:], end_point) pos_end = to_positive_idxs(var.shape[-2:], end_point)
else: else:
pos_end = start_point pos_end = start_point
xdim = var.shape[-1] xdim = var.shape[-1]
ydim = var.shape[-2] ydim = var.shape[-2]
xy = _calc_xy(xdim, ydim, pos_pivot, angle, pos_start, pos_end) xy = _calc_xy(xdim, ydim, pos_pivot, angle, pos_start, pos_end)
return xy return xy
def to_xy_coords(pairs, wrfin=None, timeidx=0, stagger=None, projection=None, def to_xy_coords(pairs, wrfin=None, timeidx=0, stagger=None, projection=None,
ll_point=None): ll_point=None):
"""Return the coordinate pairs in grid space. """Return the coordinate pairs in grid space.
This function converts latitude,longitude coordinate pairs to This function converts latitude,longitude coordinate pairs to
x,y coordinate pairs. x,y coordinate pairs.
Args: Args:
pairs (:class:`CoordPair` or sequence): A single coordinate pair or pairs (:class:`CoordPair` or sequence): A single coordinate pair or
a sequence of coordinate pairs to be converted. a sequence of coordinate pairs to be converted.
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable, optional): WRF-ARW NetCDF iterable, optional): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. This is used or an iterable sequence of the aforementioned types. This is used
to obtain the map projection when using latitude,longitude to obtain the map projection when using latitude,longitude
coordinates. Should not be used when working with x,y coordinates. Should not be used when working with x,y
coordinates. Default is None. coordinates. Default is None.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index when obtaining map boundary information desired time index when obtaining map boundary information
from moving nests. This value can be a positive integer, from moving nests. This value can be a positive integer,
negative integer, or negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return :data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. Only required when all times in the file or sequence. Only required when
*wrfin* is specified and the nest is moving. Default is 0. *wrfin* is specified and the nest is moving. Default is 0.
stagger (:obj:`str`): If using latitude, longitude coordinate pairs stagger (:obj:`str`): If using latitude, longitude coordinate pairs
for *start_point*, *end_point*, or *pivot_point*, for *start_point*, *end_point*, or *pivot_point*,
set the appropriate grid staggering type for *field2d*. By default, set the appropriate grid staggering type for *field2d*. By default,
the mass grid is used. The options are: the mass grid is used. The options are:
- 'm': Use the mass grid (default). - 'm': Use the mass grid (default).
- 'u': Use the same staggered grid as the u wind component, - 'u': Use the same staggered grid as the u wind component,
which has a staggered west_east (x) dimension. which has a staggered west_east (x) dimension.
- 'v': Use the same staggered grid as the v wind component, - 'v': Use the same staggered grid as the v wind component,
which has a staggered south_north (y) dimension. which has a staggered south_north (y) dimension.
projection (:class:`wrf.WrfProj`, optional): The map projection (:class:`wrf.WrfProj`, optional): The map
projection object to use when working with latitude, longitude projection object to use when working with latitude, longitude
coordinates, and must be specified if *wrfin* is None. Default coordinates, and must be specified if *wrfin* is None. Default
is None. is None.
ll_point (:class:`wrf.CoordPair`, sequence of :class:`wrf.CoordPair`, \ ll_point (:class:`wrf.CoordPair`, sequence of :class:`wrf.CoordPair`, \
optional): The lower left latitude, longitude point for your domain, optional): The lower left latitude, longitude point for your domain,
and must be specified and must be specified
if *wrfin* is None. If the domain is a moving nest, this should be if *wrfin* is None. If the domain is a moving nest, this should be
a sequence of :class:`wrf.CoordPair`. Default is None. a sequence of :class:`wrf.CoordPair`. Default is None.
Returns: Returns:
:class:`wrf.CoordPair` or sequence: The coordinate pair(s) in :class:`wrf.CoordPair` or sequence: The coordinate pair(s) in
x,y grid coordinates. x,y grid coordinates.
""" """
if (wrfin is None and (projection is None or ll_point is None)): if (wrfin is None and (projection is None or ll_point is None)):
raise ValueError ("'wrfin' parameter or " raise ValueError("'wrfin' parameter or "
"'projection' and 'll_point' parameters " "'projection' and 'll_point' parameters "
"are required") "are required")
lat, lon = pairs_to_latlon(pairs) lat, lon = pairs_to_latlon(pairs)
if wrfin is not None: if wrfin is not None:
xy_vals = _ll_to_xy(lat, lon, wrfin=wrfin, timeidx=timeidx, xy_vals = _ll_to_xy(lat, lon, wrfin=wrfin, timeidx=timeidx,
squeeze=True, meta=False, stagger=stagger, as_int=True) squeeze=True, meta=False, stagger=stagger,
as_int=True)
else: else:
map_proj = projection.map_proj map_proj = projection.map_proj
if map_proj == ProjectionTypes.LAT_LON: if map_proj == ProjectionTypes.LAT_LON:
pole_lat = projection.pole_lat pole_lat = projection.pole_lat
pole_lon = projection.pole_lon pole_lon = projection.pole_lon
latinc = ((projection.dy*360.0)/2.0 / latinc = ((projection.dy*360.0)/2.0 /
Constants.PI/Constants.WRF_EARTH_RADIUS) Constants.PI/Constants.WRF_EARTH_RADIUS)
loninc = ((projection.dx*360.0)/2.0 / loninc = ((projection.dx*360.0)/2.0 /
Constants.PI/Constants.WRF_EARTH_RADIUS) Constants.PI/Constants.WRF_EARTH_RADIUS)
else: else:
pole_lat = 90.0 pole_lat = 90.0
pole_lon = 0.0 pole_lon = 0.0
latinc = 0.0 latinc = 0.0
loninc = 0.0 loninc = 0.0
ll_lat, ll_lon = pairs_to_latlon(ll_point) ll_lat, ll_lon = pairs_to_latlon(ll_point)
xy_vals = _ll_to_xy(lat, lon, meta=False, squeeze=True, xy_vals = _ll_to_xy(lat, lon, meta=False, squeeze=True,
as_int=True, as_int=True,
map_proj=projection.map_proj, map_proj=projection.map_proj,
truelat1=projection.truelat1, truelat1=projection.truelat1,
truelat2=projection.truelat2, truelat2=projection.truelat2,
stand_lon=projection.stand_lon, stand_lon=projection.stand_lon,
ref_lat=ll_lat, ref_lat=ll_lat,
ref_lon=ll_lon, ref_lon=ll_lon,
pole_lat=pole_lat, pole_lat=pole_lat,
pole_lon=pole_lon, pole_lon=pole_lon,
known_x=0, known_x=0,
known_y=0, known_y=0,
dx=projection.dx, dx=projection.dx,
dy=projection.dy, dy=projection.dy,
latinc=latinc, latinc=latinc,
loninc=loninc) loninc=loninc)
xy_vals = xy_vals.squeeze() xy_vals = xy_vals.squeeze()
if xy_vals.ndim == 1: if xy_vals.ndim == 1:
return CoordPair(x=xy_vals[0], y=xy_vals[1]) return CoordPair(x=xy_vals[0], y=xy_vals[1])
else: else:
return [CoordPair(x=xy_vals[0,i], y=xy_vals[1,i]) return [CoordPair(x=xy_vals[0, i], y=xy_vals[1, i])
for i in py3range(xy_vals.shape[1])] for i in py3range(xy_vals.shape[1])]

547
src/wrf/latlonutils.py

@ -6,36 +6,36 @@ import numpy as np
from .constants import Constants, ProjectionTypes from .constants import Constants, ProjectionTypes
from .extension import _lltoxy, _xytoll from .extension import _lltoxy, _xytoll
from .util import (extract_vars, extract_global_attrs, from .util import (extract_vars, extract_global_attrs,
either, is_moving_domain, is_multi_time_req, either, is_moving_domain, iter_left_indexes,
iter_left_indexes, is_mapping, is_multi_file) is_mapping, is_multi_file)
from .py3compat import viewkeys, viewitems from .py3compat import viewkeys, viewitems
from .projutils import dict_keys_to_upper from .projutils import dict_keys_to_upper
def _lat_varname(wrfin, stagger): def _lat_varname(wrfin, stagger):
"""Return the latitude variable name for the specified stagger type. """Return the latitude variable name for the specified stagger type.
Args: Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. or an iterable sequence of the aforementioned types.
stagger (:obj:`str`): The staggered grid type which is one of the stagger (:obj:`str`): The staggered grid type which is one of the
following: following:
- 'm': Use the mass grid (default). - 'm': Use the mass grid (default).
- 'u': Use the same staggered grid as the u wind component, - 'u': Use the same staggered grid as the u wind component,
which has a staggered west_east (x) dimension. which has a staggered west_east (x) dimension.
- 'v': Use the same staggered grid as the v wind component, - 'v': Use the same staggered grid as the v wind component,
which has a staggered south_north (y) dimension. which has a staggered south_north (y) dimension.
Returns: Returns:
:obj:`str`: The latitude variable name. :obj:`str`: The latitude variable name.
""" """
if stagger is None or stagger.lower() == "m": if stagger is None or stagger.lower() == "m":
varname = either("XLAT", "XLAT_M")(wrfin) varname = either("XLAT", "XLAT_M")(wrfin)
@ -43,32 +43,33 @@ def _lat_varname(wrfin, stagger):
varname = "XLAT_{}".format(stagger.upper()) varname = "XLAT_{}".format(stagger.upper())
else: else:
raise ValueError("invalid 'stagger' value") raise ValueError("invalid 'stagger' value")
return varname return varname
def _lon_varname(wrfin, stagger): def _lon_varname(wrfin, stagger):
"""Return the longitude variable name for the specified stagger type. """Return the longitude variable name for the specified stagger type.
Args: Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. or an iterable sequence of the aforementioned types.
stagger (:obj:`str`): The staggered grid type, which is one of the stagger (:obj:`str`): The staggered grid type, which is one of the
following: following:
- 'm': Use the mass grid (default). - 'm': Use the mass grid (default).
- 'u': Use the same staggered grid as the u wind component, - 'u': Use the same staggered grid as the u wind component,
which has a staggered west_east (x) dimension. which has a staggered west_east (x) dimension.
- 'v': Use the same staggered grid as the v wind component, - 'v': Use the same staggered grid as the v wind component,
which has a staggered south_north (y) dimension. which has a staggered south_north (y) dimension.
Returns: Returns:
:obj:`str`: The latitude variable name. :obj:`str`: The latitude variable name.
""" """
if stagger is None or stagger.lower() == "m": if stagger is None or stagger.lower() == "m":
varname = either("XLONG", "XLONG_M")(wrfin) varname = either("XLONG", "XLONG_M")(wrfin)
@ -76,64 +77,65 @@ def _lon_varname(wrfin, stagger):
varname = "XLONG_{}".format(stagger.upper()) varname = "XLONG_{}".format(stagger.upper())
else: else:
raise ValueError("invalid 'stagger' value") raise ValueError("invalid 'stagger' value")
return varname return varname
def _get_proj_params(wrfin, timeidx, stagger, method, squeeze, cache, _key): def _get_proj_params(wrfin, timeidx, stagger, method, squeeze, cache, _key):
"""Return the map projection parameters. """Return the map projection parameters.
Args: Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. or an iterable sequence of the aforementioned types.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer, desired time index. This value can be a positive integer,
negative integer, or negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return :data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0. all times in the file or sequence. The default is 0.
stagger (:obj:`str`): The staggered grid type, which is one of the stagger (:obj:`str`): The staggered grid type, which is one of the
following: following:
- 'm': Use the mass grid (default). - 'm': Use the mass grid (default).
- 'u': Use the same staggered grid as the u wind component, - 'u': Use the same staggered grid as the u wind component,
which has a staggered west_east (x) dimension. which has a staggered west_east (x) dimension.
- 'v': Use the same staggered grid as the v wind component, - 'v': Use the same staggered grid as the v wind component,
which has a staggered south_north (y) dimension. which has a staggered south_north (y) dimension.
method (:obj:`str`, optional): The aggregation method to use for method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'. sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension. 'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index. 'join' creates a new dimension for the file index.
The default is 'cat'. The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape with a size of 1 from being automatically removed from the shape
of the output. Default is True. of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray) cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using used in multiple diagnostics calculations, particularly when using
large sequences of files. large sequences of files.
Default is None. Default is None.
_key (:obj:`int`, optional): A caching key. This is used for internal _key (:obj:`int`, optional): A caching key. This is used for internal
purposes only. Default is None. purposes only. Default is None.
Returns: Returns:
""" """
if timeidx is not None: if timeidx is not None:
if timeidx < 0: if timeidx < 0:
raise ValueError("'timeidx' must be greater than 0") raise ValueError("'timeidx' must be greater than 0")
attrs = extract_global_attrs(wrfin, attrs=("MAP_PROJ", "TRUELAT1", attrs = extract_global_attrs(wrfin, attrs=("MAP_PROJ", "TRUELAT1",
"TRUELAT2", "STAND_LON", "TRUELAT2", "STAND_LON",
"DX", "DY")) "DX", "DY"))
@ -143,9 +145,9 @@ def _get_proj_params(wrfin, timeidx, stagger, method, squeeze, cache, _key):
stdlon = attrs["STAND_LON"] stdlon = attrs["STAND_LON"]
dx = attrs["DX"] dx = attrs["DX"]
dy = attrs["DY"] dy = attrs["DY"]
if map_proj == ProjectionTypes.LAT_LON: if map_proj == ProjectionTypes.LAT_LON:
pole_attrs = extract_global_attrs(wrfin, attrs=("POLE_LAT", pole_attrs = extract_global_attrs(wrfin, attrs=("POLE_LAT",
"POLE_LON")) "POLE_LON"))
pole_lat = pole_attrs["POLE_LAT"] pole_lat = pole_attrs["POLE_LAT"]
pole_lon = pole_attrs["POLE_LON"] pole_lon = pole_attrs["POLE_LON"]
@ -156,20 +158,20 @@ def _get_proj_params(wrfin, timeidx, stagger, method, squeeze, cache, _key):
pole_lon = 0.0 pole_lon = 0.0
latinc = 0.0 latinc = 0.0
loninc = 0.0 loninc = 0.0
latvar = _lat_varname(wrfin, stagger) latvar = _lat_varname(wrfin, stagger)
lonvar = _lon_varname(wrfin, stagger) lonvar = _lon_varname(wrfin, stagger)
lat_timeidx = timeidx lat_timeidx = timeidx
is_moving = is_moving_domain(wrfin, latvar=latvar, lonvar=lonvar, is_moving = is_moving_domain(wrfin, latvar=latvar, lonvar=lonvar,
_key=_key) _key=_key)
# Only need one file and one time if the domain is not moving # Only need one file and one time if the domain is not moving
if not is_moving: if not is_moving:
# Always use the 0th time for non-moving domains to avoid problems # Always use the 0th time for non-moving domains to avoid problems
lat_timeidx = 0 lat_timeidx = 0
if is_multi_file(wrfin): if is_multi_file(wrfin):
if not is_mapping(wrfin): if not is_mapping(wrfin):
wrfin = next(iter(wrfin)) # only need one file wrfin = next(iter(wrfin)) # only need one file
@ -177,43 +179,43 @@ def _get_proj_params(wrfin, timeidx, stagger, method, squeeze, cache, _key):
first_entry = next(iter(viewkeys(wrfin))) first_entry = next(iter(viewkeys(wrfin)))
wrfin = wrfin[first_entry] wrfin = wrfin[first_entry]
key = _key[first_entry] key = _key[first_entry]
return _get_proj_params(wrfin, timeidx, stagger, return _get_proj_params(wrfin, timeidx, stagger,
method, squeeze, cache, key) method, squeeze, cache, key)
xlat = extract_vars(wrfin, lat_timeidx, (latvar,), method, squeeze, cache, xlat = extract_vars(wrfin, lat_timeidx, (latvar,), method, squeeze, cache,
meta=False, _key=_key)[latvar] meta=False, _key=_key)[latvar]
xlon = extract_vars(wrfin, lat_timeidx, (lonvar,), method, squeeze, cache, xlon = extract_vars(wrfin, lat_timeidx, (lonvar,), method, squeeze, cache,
meta=False, _key=_key)[lonvar] meta=False, _key=_key)[lonvar]
ref_lat = np.ravel(xlat[..., 0, 0]) ref_lat = np.ravel(xlat[..., 0, 0])
ref_lon = np.ravel(xlon[..., 0, 0]) ref_lon = np.ravel(xlon[..., 0, 0])
# Note: fortran index # Note: fortran index
known_x = 1.0 known_x = 1.0
known_y = 1.0 known_y = 1.0
return (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, return (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
pole_lat, pole_lon, known_x, known_y, dx, dy, latinc, loninc) pole_lat, pole_lon, known_x, known_y, dx, dy, latinc, loninc)
# known_x and known_y are 0-based # known_x and known_y are 0-based
def _kwarg_proj_params(**projparams): def _kwarg_proj_params(**projparams):
"""Return the map projection parameters. """Return the map projection parameters.
This function aggregates the projection parameter keyword This function aggregates the projection parameter keyword
arguments and also performs sanity checking on them. arguments and also performs sanity checking on them.
Args: Args:
**projparams: Projection parameter keyword arguments. **projparams: Projection parameter keyword arguments.
Returns: Returns:
:obj:`tuple`: The map projection parameters. :obj:`tuple`: The map projection parameters.
""" """
projparams = dict_keys_to_upper(projparams) projparams = dict_keys_to_upper(projparams)
map_proj = projparams.get("MAP_PROJ") map_proj = projparams.get("MAP_PROJ")
truelat1 = projparams.get("TRUELAT1") truelat1 = projparams.get("TRUELAT1")
truelat2 = projparams.get("TRUELAT2") truelat2 = projparams.get("TRUELAT2")
@ -222,164 +224,164 @@ def _kwarg_proj_params(**projparams):
ref_lon = projparams.get("REF_LON") ref_lon = projparams.get("REF_LON")
pole_lat = projparams.get("POLE_LAT", 90.0) pole_lat = projparams.get("POLE_LAT", 90.0)
pole_lon = projparams.get("POLE_LON", 0.0) pole_lon = projparams.get("POLE_LON", 0.0)
known_x = projparams.get("KNOWN_X") # Use 0-based known_x = projparams.get("KNOWN_X") # Use 0-based
known_y = projparams.get("KNOWN_Y") # Use 0-based known_y = projparams.get("KNOWN_Y") # Use 0-based
dx = projparams.get("DX") dx = projparams.get("DX")
dy = projparams.get("DY") dy = projparams.get("DY")
latinc = projparams.get("LATINC") latinc = projparams.get("LATINC")
loninc = projparams.get("LONINC") loninc = projparams.get("LONINC")
# Sanity checks # Sanity checks
# Required args for all projections # Required args for all projections
for name, var in viewitems({"MAP_PROJ" : map_proj, for name, var in viewitems({"MAP_PROJ": map_proj,
"REF_LAT" : ref_lat, "REF_LAT": ref_lat,
"REF_LON" : ref_lon, "REF_LON": ref_lon,
"KNOWN_X" : known_x, "KNOWN_X": known_x,
"KNOWN_Y" : known_y, "KNOWN_Y": known_y,
"DX" : dx}): "DX": dx}):
if var is None: if var is None:
raise ValueError("'{}' argument required".format(name)) raise ValueError("'{}' argument required".format(name))
# ref_lat and ref_lon are expected to be lists # ref_lat and ref_lon are expected to be lists
ref_lat = np.ravel(np.asarray([ref_lat])) ref_lat = np.ravel(np.asarray([ref_lat]))
ref_lon = np.ravel(np.asarray([ref_lon])) ref_lon = np.ravel(np.asarray([ref_lon]))
# Fortran wants 1-based indexing # Fortran wants 1-based indexing
known_x = known_x + 1 known_x = known_x + 1
known_y = known_y + 1 known_y = known_y + 1
if map_proj in (ProjectionTypes.LAMBERT_CONFORMAL, if map_proj in (ProjectionTypes.LAMBERT_CONFORMAL,
ProjectionTypes.POLAR_STEREOGRAPHIC, ProjectionTypes.POLAR_STEREOGRAPHIC,
ProjectionTypes.MERCATOR): ProjectionTypes.MERCATOR):
if truelat1 is None: if truelat1 is None:
raise ValueError("'TRUELAT1' argument required") raise ValueError("'TRUELAT1' argument required")
else: else:
if truelat1 is None: if truelat1 is None:
truelat1 = 0.0 truelat1 = 0.0
# Map projection 6 (lat lon) required latinc, loninc, and dy # Map projection 6 (lat lon) required latinc, loninc, and dy
if map_proj == ProjectionTypes.LAT_LON: if map_proj == ProjectionTypes.LAT_LON:
if latinc is None: if latinc is None:
raise ValueError("'LATINC' argument required") raise ValueError("'LATINC' argument required")
if loninc is None: if loninc is None:
raise ValueError("'LONINC' argument required") raise ValueError("'LONINC' argument required")
if dy is None: if dy is None:
raise ValueError("'DY' argument required") raise ValueError("'DY' argument required")
else: else:
latinc = 0.0 latinc = 0.0
loninc = 0.0 loninc = 0.0
dy = 0.0 dy = 0.0
return (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, pole_lat, return (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, pole_lat,
pole_lon, known_x, known_y, dx, dy, latinc, loninc) pole_lon, known_x, known_y, dx, dy, latinc, loninc)
# Will return 0-based indexes # Will return 0-based indexes
def _ll_to_xy(latitude, longitude, wrfin=None, timeidx=0, def _ll_to_xy(latitude, longitude, wrfin=None, timeidx=0,
stagger=None, method="cat", squeeze=True, cache=None, stagger=None, method="cat", squeeze=True, cache=None,
_key=None, as_int=True, **projparams): _key=None, as_int=True, **projparams):
"""Return the x,y coordinates for a specified latitude and longitude. """Return the x,y coordinates for a specified latitude and longitude.
The *latitude* and *longitude* arguments can be a single value or a The *latitude* and *longitude* arguments can be a single value or a
sequence of values. sequence of values.
The leftmost dimension of the returned array represents two different The leftmost dimension of the returned array represents two different
quantities: quantities:
- return_val[0,...] will contain the X (west_east) values. - return_val[0,...] will contain the X (west_east) values.
- return_val[1,...] will contain the Y (south_north) values. - return_val[1,...] will contain the Y (south_north) values.
Args: Args:
latitude (:obj:`float` or sequence): A single latitude or a sequence latitude (:obj:`float` or sequence): A single latitude or a sequence
of latitude values to be converted. of latitude values to be converted.
longitude (:obj:`float` or sequence): A single longitude or a sequence longitude (:obj:`float` or sequence): A single longitude or a sequence
of latitude values to be converted. of latitude values to be converted.
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. or an iterable sequence of the aforementioned types.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer, desired time index. This value can be a positive integer,
negative integer, or negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return :data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0. all times in the file or sequence. The default is 0.
stagger (:obj:`str`): By default, the latitude and longitude are stagger (:obj:`str`): By default, the latitude and longitude are
returned on the mass grid, but a staggered grid can be chosen returned on the mass grid, but a staggered grid can be chosen
with the following options: with the following options:
- 'm': Use the mass grid (default). - 'm': Use the mass grid (default).
- 'u': Use the same staggered grid as the u wind component, - 'u': Use the same staggered grid as the u wind component,
which has a staggered west_east (x) dimension. which has a staggered west_east (x) dimension.
- 'v': Use the same staggered grid as the v wind component, - 'v': Use the same staggered grid as the v wind component,
which has a staggered south_north (y) dimension. which has a staggered south_north (y) dimension.
method (:obj:`str`, optional): The aggregation method to use for method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'. sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension. 'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index. 'join' creates a new dimension for the file index.
The default is 'cat'. The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape with a size of 1 from being automatically removed from the shape
of the output. Default is True. of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray) cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using used in multiple diagnostics calculations, particularly when using
large sequences of files. large sequences of files.
Default is None. Default is None.
_key (:obj:`int`, optional): A caching key. This is used for internal _key (:obj:`int`, optional): A caching key. This is used for internal
purposes only. Default is None. purposes only. Default is None.
as_int (:obj:`bool`): Set to True to return the x,y values as as_int (:obj:`bool`): Set to True to return the x,y values as
:obj:`int`, otherwise they will be returned as :obj:`float`. :obj:`int`, otherwise they will be returned as :obj:`float`.
**projparams: Map projection keyword arguments to set manually. **projparams: Map projection keyword arguments to set manually.
Returns: Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The :class:`xarray.DataArray` or :class:`numpy.ndarray`: The
x,y coordinate value(s) whose leftmost dimension is 2 (0=X, 1=Y). x,y coordinate value(s) whose leftmost dimension is 2 (0=X, 1=Y).
If xarray is enabled and the *meta* parameter is True, then the result If xarray is enabled and the *meta* parameter is True, then the result
will be a :class:`xarray.DataArray` object. Otherwise, the result will will be a :class:`xarray.DataArray` object. Otherwise, the result will
be a :class:`numpy.ndarray` object with no metadata. be a :class:`numpy.ndarray` object with no metadata.
""" """
if wrfin is not None: if wrfin is not None:
(map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
pole_lat, pole_lon, known_x, known_y, dx, dy, latinc, pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
loninc) = _get_proj_params(wrfin, timeidx, stagger, method, squeeze, loninc) = _get_proj_params(wrfin, timeidx, stagger, method, squeeze,
cache, _key) cache, _key)
else: else:
(map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
pole_lat, pole_lon, known_x, known_y, dx, dy, latinc, pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
loninc) = _kwarg_proj_params(**projparams) loninc) = _kwarg_proj_params(**projparams)
if isinstance(latitude, Iterable): if isinstance(latitude, Iterable):
lats = np.asarray(latitude) lats = np.asarray(latitude)
lons = np.asarray(longitude) lons = np.asarray(longitude)
# Note: For scalars, this will make a single element array # Note: For scalars, this will make a single element array
lats = lats.ravel() lats = lats.ravel()
lons = lons.ravel() lons = lons.ravel()
if (lats.size != lons.size): if (lats.size != lons.size):
raise ValueError("'latitude' and 'longitude' " raise ValueError("'latitude' and 'longitude' "
"must be the same length") "must be the same length")
if ref_lat.size == 1: if ref_lat.size == 1:
outdim = [2, lats.size] outdim = [2, lats.size]
extra_dims = [outdim[1]] extra_dims = [outdim[1]]
@ -387,11 +389,10 @@ def _ll_to_xy(latitude, longitude, wrfin=None, timeidx=0,
# Moving domain will have moving ref_lats/ref_lons # Moving domain will have moving ref_lats/ref_lons
outdim = [2, ref_lat.size, lats.size] outdim = [2, ref_lat.size, lats.size]
extra_dims = outdim[1:] extra_dims = outdim[1:]
result = np.empty(outdim, np.float64) result = np.empty(outdim, np.float64)
for left_idxs in iter_left_indexes(extra_dims): for left_idxs in iter_left_indexes(extra_dims):
#left_and_slice_idxs = left_idxs + (slice(None), )
# Left indexes is a misnomer, since these will be on the right # Left indexes is a misnomer, since these will be on the right
x_idxs = (0,) + left_idxs x_idxs = (0,) + left_idxs
y_idxs = (1,) + left_idxs y_idxs = (1,) + left_idxs
@ -401,149 +402,148 @@ def _ll_to_xy(latitude, longitude, wrfin=None, timeidx=0,
else: else:
ref_lat_val = ref_lat[left_idxs[-2]] ref_lat_val = ref_lat[left_idxs[-2]]
ref_lon_val = ref_lon[left_idxs[-2]] ref_lon_val = ref_lon[left_idxs[-2]]
lat = lats[left_idxs[-1]] lat = lats[left_idxs[-1]]
lon = lons[left_idxs[-1]] lon = lons[left_idxs[-1]]
xy = _lltoxy(map_proj, truelat1, truelat2, stdlon, xy = _lltoxy(map_proj, truelat1, truelat2, stdlon,
ref_lat_val, ref_lon_val, pole_lat, pole_lon, ref_lat_val, ref_lon_val, pole_lat, pole_lon,
known_x, known_y, dx, dy, latinc, loninc, known_x, known_y, dx, dy, latinc, loninc,
lat, lon) lat, lon)
# Note: comes back from fortran as y,x # Note: comes back from fortran as y,x
result[x_idxs] = xy[1] result[x_idxs] = xy[1]
result[y_idxs] = xy[0] result[y_idxs] = xy[0]
else: else:
result = np.empty((2,), np.float64) result = np.empty((2,), np.float64)
fort_out = _lltoxy(map_proj, truelat1, truelat2, stdlon, fort_out = _lltoxy(map_proj, truelat1, truelat2, stdlon,
ref_lat, ref_lon, pole_lat, pole_lon, ref_lat, ref_lon, pole_lat, pole_lon,
known_x, known_y, dx, dy, latinc, loninc, known_x, known_y, dx, dy, latinc, loninc,
latitude, longitude) latitude, longitude)
# Note, comes back from fortran as y,x. So, need to swap them. # Note, comes back from fortran as y,x. So, need to swap them.
result[0] = fort_out[1] result[0] = fort_out[1]
result[1] = fort_out[0] result[1] = fort_out[0]
# Make indexes 0-based # Make indexes 0-based
result = result - 1 result = result - 1
if as_int: if as_int:
result = np.rint(result).astype(int) result = np.rint(result).astype(int)
return result return result
# X and Y should be 0-based # X and Y should be 0-based
def _xy_to_ll(x, y, wrfin=None, timeidx=0, stagger=None, def _xy_to_ll(x, y, wrfin=None, timeidx=0, stagger=None,
method="cat", squeeze=True, cache=None, _key=None, method="cat", squeeze=True, cache=None, _key=None,
**projparams): **projparams):
"""Return the latitude and longitude for specified x,y coordinates. """Return the latitude and longitude for specified x,y coordinates.
The *x* and *y* arguments can be a single value or a sequence of values. The *x* and *y* arguments can be a single value or a sequence of values.
The leftmost dimension of the returned array represents two different The leftmost dimension of the returned array represents two different
quantities: quantities:
- return_val[0,...] will contain the latitude values. - return_val[0,...] will contain the latitude values.
- return_val[1,...] will contain the longitude values. - return_val[1,...] will contain the longitude values.
Args: Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. or an iterable sequence of the aforementioned types.
x (:obj:`float` or sequence): A single x-coordinate or a sequence x (:obj:`float` or sequence): A single x-coordinate or a sequence
of x-coordinate values to be converted. of x-coordinate values to be converted.
y (:obj:`float` or sequence): A single y-coordinate or a sequence y (:obj:`float` or sequence): A single y-coordinate or a sequence
of y-coordinate values to be converted. of y-coordinate values to be converted.
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. or an iterable sequence of the aforementioned types.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer, desired time index. This value can be a positive integer,
negative integer, or negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return :data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0. all times in the file or sequence. The default is 0.
stagger (:obj:`str`): By default, the latitude and longitude are stagger (:obj:`str`): By default, the latitude and longitude are
returned on the mass grid, but a staggered grid can be chosen returned on the mass grid, but a staggered grid can be chosen
with the following options: with the following options:
- 'm': Use the mass grid (default). - 'm': Use the mass grid (default).
- 'u': Use the same staggered grid as the u wind component, - 'u': Use the same staggered grid as the u wind component,
which has a staggered west_east (x) dimension. which has a staggered west_east (x) dimension.
- 'v': Use the same staggered grid as the v wind component, - 'v': Use the same staggered grid as the v wind component,
which has a staggered south_north (y) dimension. which has a staggered south_north (y) dimension.
method (:obj:`str`, optional): The aggregation method to use for method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'. sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension. 'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index. 'join' creates a new dimension for the file index.
The default is 'cat'. The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape with a size of 1 from being automatically removed from the shape
of the output. Default is True. of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray) cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using used in multiple diagnostics calculations, particularly when using
large sequences of files. large sequences of files.
Default is None. Default is None.
_key (:obj:`int`, optional): A caching key. This is used for internal _key (:obj:`int`, optional): A caching key. This is used for internal
purposes only. Default is None. purposes only. Default is None.
**projparams: Map projection keyword arguments to set manually. **projparams: Map projection keyword arguments to set manually.
Returns: Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The :class:`xarray.DataArray` or :class:`numpy.ndarray`: The
latitude and longitude values whose leftmost dimension is 2 latitude and longitude values whose leftmost dimension is 2
(0=latitude, 1=longitude). (0=latitude, 1=longitude).
If xarray is enabled and the *meta* parameter is True, then the result If xarray is enabled and the *meta* parameter is True, then the result
will be a :class:`xarray.DataArray` object. Otherwise, the result will will be a :class:`xarray.DataArray` object. Otherwise, the result will
be a :class:`numpy.ndarray` object with no metadata. be a :class:`numpy.ndarray` object with no metadata.
""" """
if wrfin is not None: if wrfin is not None:
(map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
pole_lat, pole_lon, known_x, known_y, dx, dy, latinc, pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
loninc) = _get_proj_params(wrfin, timeidx, stagger, method, squeeze, loninc) = _get_proj_params(wrfin, timeidx, stagger, method, squeeze,
cache, _key) cache, _key)
else: else:
(map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
pole_lat, pole_lon, known_x, known_y, dx, dy, latinc, pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
loninc) = _kwarg_proj_params(**projparams) loninc) = _kwarg_proj_params(**projparams)
if isinstance(x, Iterable): if isinstance(x, Iterable):
x_arr = np.asarray(x) x_arr = np.asarray(x)
y_arr = np.asarray(y) y_arr = np.asarray(y)
# Convert 0-based x, y to 1-based # Convert 0-based x, y to 1-based
x_arr = x_arr + 1 x_arr = x_arr + 1
y_arr = y_arr + 1 y_arr = y_arr + 1
x_arr = x_arr.ravel() x_arr = x_arr.ravel()
y_arr = y_arr.ravel() y_arr = y_arr.ravel()
if (x_arr.size != y_arr.size): if (x_arr.size != y_arr.size):
raise ValueError("'x' and 'y' must be the same length") raise ValueError("'x' and 'y' must be the same length")
if ref_lat.size == 1: if ref_lat.size == 1:
outdim = [2, x_arr.size] outdim = [2, x_arr.size]
extra_dims = [outdim[1]] extra_dims = [outdim[1]]
@ -551,44 +551,37 @@ def _xy_to_ll(x, y, wrfin=None, timeidx=0, stagger=None,
# Moving domain will have moving ref_lats/ref_lons # Moving domain will have moving ref_lats/ref_lons
outdim = [2, ref_lat.size, x_arr.size] outdim = [2, ref_lat.size, x_arr.size]
extra_dims = outdim[1:] extra_dims = outdim[1:]
result = np.empty(outdim, np.float64) result = np.empty(outdim, np.float64)
for left_idxs in iter_left_indexes(extra_dims): for left_idxs in iter_left_indexes(extra_dims):
#left_and_slice_idxs = left_idxs + (slice(None), )
lat_idxs = (0,) + left_idxs lat_idxs = (0,) + left_idxs
lon_idxs = (1,) + left_idxs lon_idxs = (1,) + left_idxs
if ref_lat.size == 1: if ref_lat.size == 1:
ref_lat_val = ref_lat[0] ref_lat_val = ref_lat[0]
ref_lon_val = ref_lon[0] ref_lon_val = ref_lon[0]
else: else:
ref_lat_val = ref_lat[left_idxs[-2]] ref_lat_val = ref_lat[left_idxs[-2]]
ref_lon_val = ref_lon[left_idxs[-2]] ref_lon_val = ref_lon[left_idxs[-2]]
x_val = x_arr[left_idxs[-1]] x_val = x_arr[left_idxs[-1]]
y_val = y_arr[left_idxs[-1]] y_val = y_arr[left_idxs[-1]]
ll = _xytoll(map_proj, truelat1, truelat2, stdlon, ref_lat_val, ll = _xytoll(map_proj, truelat1, truelat2, stdlon, ref_lat_val,
ref_lon_val, pole_lat, pole_lon, known_x, known_y, ref_lon_val, pole_lat, pole_lon, known_x, known_y,
dx, dy, latinc, loninc, x_val, y_val) dx, dy, latinc, loninc, x_val, y_val)
#result[left_and_slice_idxs] = ll[:]
result[lat_idxs] = ll[0] result[lat_idxs] = ll[0]
result[lon_idxs] = ll[1] result[lon_idxs] = ll[1]
else: else:
# Convert 0-based to 1-based for Fortran # Convert 0-based to 1-based for Fortran
x_val = x + 1 x_val = x + 1
y_val = y + 1 y_val = y + 1
result = _xytoll(map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,
pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
loninc, x_val, y_val)
return result
result = _xytoll(map_proj, truelat1, truelat2, stdlon, ref_lat,
ref_lon, pole_lat, pole_lon, known_x, known_y,
dx, dy, latinc, loninc, x_val, y_val)
return result

1841
src/wrf/metadecorators.py

File diff suppressed because it is too large Load Diff

988
src/wrf/projection.py

File diff suppressed because it is too large Load Diff

13
src/wrf/projutils.py

@ -2,16 +2,17 @@ from __future__ import (absolute_import, division, print_function)
from .py3compat import viewitems from .py3compat import viewitems
def dict_keys_to_upper(d): def dict_keys_to_upper(d):
"""Return a dictionary with the keys changed to uppercase. """Return a dictionary with the keys changed to uppercase.
Args: Args:
d (:obj:`dict`): A dictionary. d (:obj:`dict`): A dictionary.
Returns: Returns:
:obj:`dict`: A dictionary with uppercase keys. :obj:`dict`: A dictionary with uppercase keys.
""" """
return {key.upper() : val for key, val in viewitems(d)} return {key.upper(): val for key, val in viewitems(d)}

118
src/wrf/py3compat.py

@ -3,18 +3,19 @@ from __future__ import (absolute_import, division, print_function)
from sys import version_info from sys import version_info
from math import floor, copysign from math import floor, copysign
# Dictionary python 2-3 compatibility stuff # Dictionary python 2-3 compatibility stuff
def viewitems(d): def viewitems(d):
"""Return either the items or viewitems method for a dictionary. """Return either the items or viewitems method for a dictionary.
Args: Args:
d (:obj:`dict`): A dictionary. d (:obj:`dict`): A dictionary.
Returns: Returns:
view method: Either the items or viewitems method. view method: Either the items or viewitems method.
""" """
func = getattr(d, "viewitems", None) func = getattr(d, "viewitems", None)
if func is None: if func is None:
@ -24,15 +25,15 @@ def viewitems(d):
def viewkeys(d): def viewkeys(d):
"""Return either the keys or viewkeys method for a dictionary. """Return either the keys or viewkeys method for a dictionary.
Args: Args:
d (:obj:`dict`): A dictionary. d (:obj:`dict`): A dictionary.
Returns: Returns:
view method: Either the keys or viewkeys method. view method: Either the keys or viewkeys method.
""" """
func = getattr(d, "viewkeys", None) func = getattr(d, "viewkeys", None)
if func is None: if func is None:
@ -42,32 +43,33 @@ def viewkeys(d):
def viewvalues(d): def viewvalues(d):
"""Return either the values or viewvalues method for a dictionary. """Return either the values or viewvalues method for a dictionary.
Args: Args:
d (:obj:`dict`): A dictionary. d (:obj:`dict`): A dictionary.
Returns: Returns:
view method: Either the values or viewvalues method. view method: Either the values or viewvalues method.
""" """
func = getattr(d, "viewvalues", None) func = getattr(d, "viewvalues", None)
if func is None: if func is None:
func = d.values func = d.values
return func() return func()
def isstr(s): def isstr(s):
"""Return True if the object is a string type. """Return True if the object is a string type.
Args: Args:
s (string): A string (str, unicode, bytes). s (string): A string (str, unicode, bytes).
Returns: Returns:
:obj:`bool`: True if the object is a type of string. Otherwise, False. :obj:`bool`: True if the object is a type of string. Otherwise, False.
""" """
try: try:
return isinstance(s, basestring) return isinstance(s, basestring)
@ -75,25 +77,25 @@ def isstr(s):
return isinstance(s, str) return isinstance(s, str)
# Python 2 rounding behavior # Python 2 rounding behavior
def _round2(x, d=None): def _round2(x, d=None):
"""Return the result of Python 2.x rounding, which is to round the number """Return the result of Python 2.x rounding, which is to round the number
to the nearest integer. to the nearest integer.
Python 3.x uses banker's rounding, which is not applicable for nearest Python 3.x uses banker's rounding, which is not applicable for nearest
neighbor approaches with grid boxes. neighbor approaches with grid boxes.
Args: Args:
x (:obj:`float`): A number, usually a float. x (:obj:`float`): A number, usually a float.
d (:obj:`int`, optional): The number of digits. Default is None, d (:obj:`int`, optional): The number of digits. Default is None,
which indicates the nearest integer. which indicates the nearest integer.
Returns: Returns:
:obj:`float`: The rounded number. :obj:`float`: The rounded number.
""" """
d = 0 if d is None else d d = 0 if d is None else d
p = 10 ** d p = 10 ** d
@ -101,69 +103,67 @@ def _round2(x, d=None):
def py2round(x, d=None): def py2round(x, d=None):
"""Return the result of Python 2.x rounding, which is to round the number """Return the result of Python 2.x rounding, which is to round the number
to the nearest integer. to the nearest integer.
Python 3.x uses banker's rounding, which is not applicable for nearest Python 3.x uses banker's rounding, which is not applicable for nearest
neighbor approaches with grid boxes. neighbor approaches with grid boxes.
Args: Args:
x (:obj:`float`): A number, usually a float. x (:obj:`float`): A number, usually a float.
d (:obj:`int`, optional): The number of digits. Default is None, d (:obj:`int`, optional): The number of digits. Default is None,
which indicates the nearest integer. which indicates the nearest integer.
Returns: Returns:
:obj:`float`: The rounded number. :obj:`float`: The rounded number.
""" """
if version_info >= (3,): if version_info >= (3,):
return _round2(x, d) return _round2(x, d)
return round(x, d) return round(x, d)
def py3range(*args): def py3range(*args):
"""Return the equivalent of the range function in Python 3.x. """Return the equivalent of the range function in Python 3.x.
For Python 2.x, this is the same as the xrange function. For Python 2.x, this is the same as the xrange function.
Args: Args:
*args: The function arguments for range or xrange. *args: The function arguments for range or xrange.
Returns: Returns:
iterable: An iterable sequence. iterable: An iterable sequence.
""" """
if version_info >= (3,): if version_info >= (3,):
return range(*args) return range(*args)
return xrange(*args) return xrange(*args)
def ucode(*args, **kwargs): def ucode(*args, **kwargs):
"""Return a Python 3.x unicode string. """Return a Python 3.x unicode string.
For Python 2.x, this is accomplished by using the unicode function. For Python 2.x, this is accomplished by using the unicode function.
Args: Args:
*args: The function positional arguments for str or unicode. *args: The function positional arguments for str or unicode.
**kwargs: The function keyword arguments for str or unicode. **kwargs: The function keyword arguments for str or unicode.
Returns: Returns:
string: A unicode string. string: A unicode string.
""" """
if version_info >= (3, ): if version_info >= (3, ):
return str(*args, **kwargs) return str(*args, **kwargs)
return unicode(*args, **kwargs)
return unicode(*args, **kwargs)

489
src/wrf/routines.py

@ -2,8 +2,8 @@ from __future__ import (absolute_import, division, print_function)
from .util import (get_iterable, is_standard_wrf_var, extract_vars, viewkeys, from .util import (get_iterable, is_standard_wrf_var, extract_vars, viewkeys,
get_id) get_id)
from .g_cape import (get_2dcape, get_3dcape, get_cape2d_only, from .g_cape import (get_2dcape, get_3dcape, get_cape2d_only,
get_cin2d_only, get_lcl, get_lfc, get_3dcape_only, get_cin2d_only, get_lcl, get_lfc, get_3dcape_only,
get_3dcin_only) get_3dcin_only)
from .g_ctt import get_ctt from .g_ctt import get_ctt
from .g_dbz import get_dbz, get_max_dbz from .g_dbz import get_dbz, get_max_dbz
@ -16,334 +16,337 @@ from .g_pressure import get_pressure, get_pressure_hpa
from .g_pw import get_pw from .g_pw import get_pw
from .g_rh import get_rh, get_rh_2m from .g_rh import get_rh, get_rh_2m
from .g_slp import get_slp from .g_slp import get_slp
from .g_temp import get_tc, get_eth, get_temp, get_theta, get_tk, get_tv, get_tw from .g_temp import (get_tc, get_eth, get_temp, get_theta, get_tk, get_tv,
get_tw)
from .g_terrain import get_terrain from .g_terrain import get_terrain
from .g_uvmet import (get_uvmet, get_uvmet10, get_uvmet10_wspd_wdir, from .g_uvmet import (get_uvmet, get_uvmet10, get_uvmet10_wspd_wdir,
get_uvmet_wspd_wdir, get_uvmet_wspd, get_uvmet_wdir, get_uvmet_wspd_wdir, get_uvmet_wspd, get_uvmet_wdir,
get_uvmet10_wspd, get_uvmet10_wdir) get_uvmet10_wspd, get_uvmet10_wdir)
from .g_vorticity import get_avo, get_pvo from .g_vorticity import get_avo, get_pvo
from .g_wind import (get_destag_wspd_wdir, get_destag_wspd_wdir10, from .g_wind import (get_destag_wspd_wdir, get_destag_wspd_wdir10,
get_u_destag, get_v_destag, get_w_destag, get_u_destag, get_v_destag, get_w_destag,
get_destag_wspd, get_destag_wdir, get_destag_wspd10, get_destag_wspd, get_destag_wdir, get_destag_wspd10,
get_destag_wdir10) get_destag_wdir10)
from .g_times import get_times, get_xtimes from .g_times import get_times, get_xtimes
from .g_cloudfrac import (get_cloudfrac, get_low_cloudfrac, get_mid_cloudfrac, from .g_cloudfrac import (get_cloudfrac, get_low_cloudfrac, get_mid_cloudfrac,
get_high_cloudfrac) get_high_cloudfrac)
# func is the function to call. kargs are required arguments that should # func is the function to call. kargs are required arguments that should
# not be altered by the user # not be altered by the user
_FUNC_MAP = {"cape2d" : get_2dcape, _FUNC_MAP = {"cape2d": get_2dcape,
"cape3d" : get_3dcape, "cape3d": get_3dcape,
"dbz" : get_dbz, "dbz": get_dbz,
"maxdbz" : get_max_dbz, "maxdbz": get_max_dbz,
"dp" : get_dp, "dp": get_dp,
"dp2m" : get_dp_2m, "dp2m": get_dp_2m,
"height" : get_height, "height": get_height,
"geopt" : get_geopt, "geopt": get_geopt,
"srh" : get_srh, "srh": get_srh,
"uhel" : get_uh, "uhel": get_uh,
"omega" : get_omega, "omega": get_omega,
"pw" : get_pw, "pw": get_pw,
"rh" : get_rh, "rh": get_rh,
"rh2m" : get_rh_2m, "rh2m": get_rh_2m,
"slp" : get_slp, "slp": get_slp,
"theta" : get_theta, "theta": get_theta,
"temp" : get_temp, "temp": get_temp,
"tk" : get_tk, "tk": get_tk,
"tc" : get_tc, "tc": get_tc,
"theta_e" : get_eth, "theta_e": get_eth,
"tv" : get_tv, "tv": get_tv,
"twb" : get_tw, "twb": get_tw,
"terrain" : get_terrain, "terrain": get_terrain,
"times" : get_times, "times": get_times,
"xtimes" : get_xtimes, "xtimes": get_xtimes,
"uvmet" : get_uvmet, "uvmet": get_uvmet,
"uvmet10" : get_uvmet10, "uvmet10": get_uvmet10,
"avo" : get_avo, "avo": get_avo,
"pvo" : get_pvo, "pvo": get_pvo,
"ua" : get_u_destag, "ua": get_u_destag,
"va" : get_v_destag, "va": get_v_destag,
"wa" : get_w_destag, "wa": get_w_destag,
"lat" : get_lat, "lat": get_lat,
"lon" : get_lon, "lon": get_lon,
"pressure" : get_pressure_hpa, "pressure": get_pressure_hpa,
"pres" : get_pressure, "pres": get_pressure,
"wspd_wdir" : get_destag_wspd_wdir, "wspd_wdir": get_destag_wspd_wdir,
"wspd_wdir10" : get_destag_wspd_wdir10, "wspd_wdir10": get_destag_wspd_wdir10,
"uvmet_wspd_wdir" : get_uvmet_wspd_wdir, "uvmet_wspd_wdir": get_uvmet_wspd_wdir,
"uvmet10_wspd_wdir" : get_uvmet10_wspd_wdir, "uvmet10_wspd_wdir": get_uvmet10_wspd_wdir,
"ctt" : get_ctt, "ctt": get_ctt,
"cloudfrac" : get_cloudfrac, "cloudfrac": get_cloudfrac,
"geopt_stag" : get_stag_geopt, "geopt_stag": get_stag_geopt,
"zstag" : get_stag_height, "zstag": get_stag_height,
# Diagnostics below are extracted from multi-product diagnostics # Diagnostics below are extracted from multi-product diagnostics
"cape2d_only" : get_cape2d_only, "cape2d_only": get_cape2d_only,
"cin2d_only" : get_cin2d_only, "cin2d_only": get_cin2d_only,
"lcl" : get_lcl, "lcl": get_lcl,
"lfc" : get_lfc, "lfc": get_lfc,
"cape3d_only" : get_3dcape_only, "cape3d_only": get_3dcape_only,
"cin3d_only": get_3dcin_only, "cin3d_only": get_3dcin_only,
"uvmet_wspd" : get_uvmet_wspd, "uvmet_wspd": get_uvmet_wspd,
"uvmet_wdir" : get_uvmet_wdir, "uvmet_wdir": get_uvmet_wdir,
"uvmet10_wspd" : get_uvmet10_wspd, "uvmet10_wspd": get_uvmet10_wspd,
"uvmet10_wdir" : get_uvmet10_wdir, "uvmet10_wdir": get_uvmet10_wdir,
"wspd" : get_destag_wspd, "wspd": get_destag_wspd,
"wdir" : get_destag_wdir, "wdir": get_destag_wdir,
"wspd10" : get_destag_wspd10, "wspd10": get_destag_wspd10,
"wdir10" : get_destag_wdir10, "wdir10": get_destag_wdir10,
"low_cloudfrac" : get_low_cloudfrac, "low_cloudfrac": get_low_cloudfrac,
"mid_cloudfrac" : get_mid_cloudfrac, "mid_cloudfrac": get_mid_cloudfrac,
"high_cloudfrac" : get_high_cloudfrac "high_cloudfrac": get_high_cloudfrac
} }
_VALID_KARGS = {"cape2d" : ["missing"], _VALID_KARGS = {"cape2d": ["missing"],
"cape3d" : ["missing"], "cape3d": ["missing"],
"dbz" : ["do_variant", "do_liqskin"], "dbz": ["do_variant", "do_liqskin"],
"maxdbz" : ["do_variant", "do_liqskin"], "maxdbz": ["do_variant", "do_liqskin"],
"dp" : ["units"], "dp": ["units"],
"dp2m" : ["units"], "dp2m": ["units"],
"height" : ["msl", "units"], "height": ["msl", "units"],
"geopt" : [], "geopt": [],
"srh" : ["top"], "srh": ["top"],
"uhel" : ["bottom", "top"], "uhel": ["bottom", "top"],
"omega" : [], "omega": [],
"pw" : [], "pw": [],
"rh" : [], "rh": [],
"rh2m" : [], "rh2m": [],
"slp" : ["units"], "slp": ["units"],
"temp" : ["units"], "temp": ["units"],
"tk" : [], "tk": [],
"tc" : [], "tc": [],
"theta" : ["units"], "theta": ["units"],
"theta_e" : ["units"], "theta_e": ["units"],
"tv" : ["units"], "tv": ["units"],
"twb" : ["units"], "twb": ["units"],
"terrain" : ["units"], "terrain": ["units"],
"times" : [], "times": [],
"xtimes" : [], "xtimes": [],
"uvmet" : ["units"], "uvmet": ["units"],
"uvmet10" : ["units"], "uvmet10": ["units"],
"avo" : [], "avo": [],
"pvo" : [], "pvo": [],
"ua" : ["units"], "ua": ["units"],
"va" : ["units"], "va": ["units"],
"wa" : ["units"], "wa": ["units"],
"lat" : [], "lat": [],
"lon" : [], "lon": [],
"pres" : ["units"], "pres": ["units"],
"pressure" : ["units"], "pressure": ["units"],
"wspd_wdir" : ["units"], "wspd_wdir": ["units"],
"wspd_wdir10" : ["units"], "wspd_wdir10": ["units"],
"uvmet_wspd_wdir" : ["units"], "uvmet_wspd_wdir": ["units"],
"uvmet10_wspd_wdir" : ["units"], "uvmet10_wspd_wdir": ["units"],
"ctt" : ["fill_nocloud", "missing", "opt_thresh", "units"], "ctt": ["fill_nocloud", "missing", "opt_thresh", "units"],
"cloudfrac" : ["vert_type", "low_thresh", "cloudfrac": ["vert_type", "low_thresh",
"mid_thresh", "high_thresh"], "mid_thresh", "high_thresh"],
"geopt_stag" : [], "geopt_stag": [],
"zstag" : ["msl", "units"], "zstag": ["msl", "units"],
"cape2d_only" : ["missing"], "cape2d_only": ["missing"],
"cin2d_only" : ["missing"], "cin2d_only": ["missing"],
"lcl" : ["missing"], "lcl": ["missing"],
"lfc" : ["missing"], "lfc": ["missing"],
"cape3d_only" : ["missing"], "cape3d_only": ["missing"],
"cin3d_only": ["missing"], "cin3d_only": ["missing"],
"uvmet_wspd" : ["units"], "uvmet_wspd": ["units"],
"uvmet_wdir" : ["units"], "uvmet_wdir": ["units"],
"uvmet10_wspd" : ["units"], "uvmet10_wspd": ["units"],
"uvmet10_wdir" : ["units"], "uvmet10_wdir": ["units"],
"wspd" : ["units"], "wspd": ["units"],
"wdir" : ["units"], "wdir": ["units"],
"wspd10" : ["units"], "wspd10": ["units"],
"wdir10" : ["units"], "wdir10": ["units"],
"low_cloudfrac" : ["vert_type", "low_thresh", "low_cloudfrac": ["vert_type", "low_thresh",
"mid_thresh", "high_thresh"], "mid_thresh", "high_thresh"],
"mid_cloudfrac" : ["vert_type", "low_thresh", "mid_cloudfrac": ["vert_type", "low_thresh",
"mid_thresh", "high_thresh"], "mid_thresh", "high_thresh"],
"high_cloudfrac" : ["vert_type", "low_thresh", "high_cloudfrac": ["vert_type", "low_thresh",
"mid_thresh", "high_thresh"], "mid_thresh", "high_thresh"],
"default" : [] "default": []
} }
_ALIASES = {"cape_2d" : "cape2d", _ALIASES = {"cape_2d": "cape2d",
"cape_3d" : "cape3d", "cape_3d": "cape3d",
"eth" : "theta_e", "eth": "theta_e",
"mdbz" : "maxdbz", "mdbz": "maxdbz",
"geopotential" : "geopt", "geopotential": "geopt",
"helicity" : "srh", "helicity": "srh",
"latitude" : "lat", "latitude": "lat",
"longitude" : "lon", "longitude": "lon",
"omg" : "omega", "omg": "omega",
"p" : "pres", "p": "pres",
"rh2" : "rh2m", "rh2": "rh2m",
"z": "height", "z": "height",
"ter" : "terrain", "ter": "terrain",
"updraft_helicity" : "uhel", "updraft_helicity": "uhel",
"td" : "dp", "td": "dp",
"td2" : "dp2m", "td2": "dp2m",
"cfrac" : "cloudfrac", "cfrac": "cloudfrac",
"wspd_wdir_uvmet" : "uvmet_wspd_wdir", "wspd_wdir_uvmet": "uvmet_wspd_wdir",
"wspd_wdir_uvmet10" : "uvmet10_wspd_wdir", "wspd_wdir_uvmet10": "uvmet10_wspd_wdir",
"th" : "theta", "th": "theta",
"low_cfrac" : "low_cloudfrac", "low_cfrac": "low_cloudfrac",
"mid_cfrac" : "mid_cloudfrac", "mid_cfrac": "mid_cloudfrac",
"high_cfrac" : "high_cloudfrac", "high_cfrac": "high_cloudfrac",
"wspd_uvmet" : "uvmet_wspd" , "wspd_uvmet": "uvmet_wspd",
"wdir_uvmet" : "uvmet_wdir" , "wdir_uvmet": "uvmet_wdir",
"wspd_uvmet10" : "uvmet10_wspd" , "wspd_uvmet10": "uvmet10_wspd",
"wdir_uvmet10" : "uvmet10_wdir" , "wdir_uvmet10": "uvmet10_wdir",
} }
class ArgumentError(Exception): class ArgumentError(Exception):
def __init__(self, msg): def __init__(self, msg):
self.msg = msg self.msg = msg
def __str__(self): def __str__(self):
return self.msg return self.msg
def _undo_alias(alias): def _undo_alias(alias):
actual = _ALIASES.get(alias, None) actual = _ALIASES.get(alias, None)
if actual is None: if actual is None:
return alias return alias
else: else:
return actual return actual
def _check_kargs(var, kargs): def _check_kargs(var, kargs):
for arg in viewkeys(kargs): for arg in viewkeys(kargs):
if arg not in _VALID_KARGS[var]: if arg not in _VALID_KARGS[var]:
if var != "default": if var != "default":
raise ValueError("'{}' is an invalid keyword " raise ValueError("'{}' is an invalid keyword "
"argument for '{}'".format(arg, var)) "argument for '{}'".format(arg, var))
else: else:
raise ValueError("'{}' is an invalid keyword " raise ValueError("'{}' is an invalid keyword "
"argument".format(arg)) "argument".format(arg))
def getvar(wrfin, varname, timeidx=0, def getvar(wrfin, varname, timeidx=0,
method="cat", squeeze=True, cache=None, meta=True, method="cat", squeeze=True, cache=None, meta=True,
**kwargs): **kwargs):
"""Returns basic diagnostics from the WRF ARW model output. """Returns basic diagnostics from the WRF ARW model output.
A table of all available diagnostics is below. A table of all available diagnostics is below.
.. include:: ../../_templates/product_table.txt .. include:: ../../_templates/product_table.txt
Args: Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. or an iterable sequence of the aforementioned types.
varname (:obj:`str`) : The variable name. varname (:obj:`str`) : The variable name.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer, desired time index. This value can be a positive integer,
negative integer, or negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return :data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0. all times in the file or sequence. The default is 0.
method (:obj:`str`, optional): The aggregation method to use for method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'. sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension. 'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index. 'join' creates a new dimension for the file index.
The default is 'cat'. The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape with a size of 1 from being automatically removed from the shape
of the output. Default is True. of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray) cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using used in multiple diagnostics calculations, particularly when using
large sequences of files. large sequences of files.
Default is None. Default is None.
meta (:obj:`bool`, optional): Set to False to disable metadata and meta (:obj:`bool`, optional): Set to False to disable metadata and
return :class:`numpy.ndarray` instead of return :class:`numpy.ndarray` instead of
:class:`xarray.DataArray`. Default is True. :class:`xarray.DataArray`. Default is True.
**kwargs: Optional keyword arguments for certain diagnostics. **kwargs: Optional keyword arguments for certain diagnostics.
See table above. See table above.
Returns: Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: If xarray is :class:`xarray.DataArray` or :class:`numpy.ndarray`: If xarray is
enabled and the *meta* parameter is True, then the result will be a enabled and the *meta* parameter is True, then the result will be a
:class:`xarray.DataArray` object. Otherwise, the result will be a :class:`xarray.DataArray` object. Otherwise, the result will be a
:class:`numpy.ndarray` object with no metadata. :class:`numpy.ndarray` object with no metadata.
Raises: Raises:
:class:`ValueError`: Raised when an invalid diagnostic type or :class:`ValueError`: Raised when an invalid diagnostic type or
keyword argument is passed to the routine. keyword argument is passed to the routine.
:class:`FortranError`: Raised when a problem occurs during a Fortran :class:`FortranError`: Raised when a problem occurs during a Fortran
calculation. calculation.
See Also: See Also:
:class:`numpy.ndarray`, :class:`xarray.DataArray` :class:`numpy.ndarray`, :class:`xarray.DataArray`
Examples: Examples:
Using netCDF4 Using netCDF4
.. code-block:: python .. code-block:: python
from netCDF4 import Dataset from netCDF4 import Dataset
from wrf import getvar from wrf import getvar
wrfnc = Dataset("wrfout_d02_2010-06-13_21:00:00") wrfnc = Dataset("wrfout_d02_2010-06-13_21:00:00")
slp = getvar(wrfnc, "slp") slp = getvar(wrfnc, "slp")
Using PyNIO Using PyNIO
.. code-block:: python .. code-block:: python
from Nio import open_file from Nio import open_file
from wrf import getvar from wrf import getvar
wrfnc = open_file("wrfout_d02_2010-06-13_21:00:00"+".nc", "r") wrfnc = open_file("wrfout_d02_2010-06-13_21:00:00"+".nc", "r")
slp = getvar(wrfnc, "slp") slp = getvar(wrfnc, "slp")
Using Iterables: Using Iterables:
.. code-block:: python .. code-block:: python
import os import os
from netCDF4 import Dataset from netCDF4 import Dataset
from wrf import getvar from wrf import getvar
filedir = "/path/to/wrf/files" filedir = "/path/to/wrf/files"
wrfin = [Dataset(f) for f in os.listdir(filedir) wrfin = [Dataset(f) for f in os.listdir(filedir)
if f.startswith("wrfout_d02_")] if f.startswith("wrfout_d02_")]
uvmet = getvar(wrfin, "uvmet", timeidx=3, units="kt") uvmet = getvar(wrfin, "uvmet", timeidx=3, units="kt")
""" """
_key = get_id(wrfin) _key = get_id(wrfin)
wrfin = get_iterable(wrfin) wrfin = get_iterable(wrfin)
if is_standard_wrf_var(wrfin, varname) and varname != "Times": if is_standard_wrf_var(wrfin, varname) and varname != "Times":
_check_kargs("default", kwargs) _check_kargs("default", kwargs)
return extract_vars(wrfin, timeidx, varname, return extract_vars(wrfin, timeidx, varname,
method, squeeze, cache, meta, _key)[varname] method, squeeze, cache, meta, _key)[varname]
elif varname == "Times": elif varname == "Times":
varname = "times" # Diverting to the get_times routine varname = "times" # Diverting to the get_times routine
actual_var = _undo_alias(varname) actual_var = _undo_alias(varname)
if actual_var not in _VALID_KARGS: if actual_var not in _VALID_KARGS:
raise ValueError("'%s' is not a valid variable name" % (varname)) raise ValueError("'{}' is not a valid variable name".format(varname))
_check_kargs(actual_var, kwargs) _check_kargs(actual_var, kwargs)
return _FUNC_MAP[actual_var](wrfin, timeidx, method, squeeze, cache, return _FUNC_MAP[actual_var](wrfin, timeidx, method, squeeze, cache,
meta, _key, **kwargs) meta, _key, **kwargs)

463
src/wrf/specialdec.py

@ -2,7 +2,7 @@ from __future__ import (absolute_import, division, print_function)
import numpy as np import numpy as np
import wrapt import wrapt
from .util import iter_left_indexes, to_np from .util import iter_left_indexes, to_np
from .py3compat import py3range from .py3compat import py3range
@ -14,28 +14,28 @@ if xarray_enabled():
def uvmet_left_iter(alg_dtype=np.float64): def uvmet_left_iter(alg_dtype=np.float64):
"""A decorator to handle iterating over the leftmost dimensions for the """A decorator to handle iterating over the leftmost dimensions for the
uvmet diagnostic. uvmet diagnostic.
For example, if a wrapped function works with three-dimensional arrays, but For example, if a wrapped function works with three-dimensional arrays, but
the variables include a 4th leftmost dimension for 'Time', this decorator the variables include a 4th leftmost dimension for 'Time', this decorator
will iterate over all times, call the 3D Fortran routine, and aggregate the will iterate over all times, call the 3D Fortran routine, and aggregate the
results in to a 4D output array. results in to a 4D output array.
It is also important to note that the final output array is allocated It is also important to note that the final output array is allocated
first, and then views are passed to the wrapped function so that values first, and then views are passed to the wrapped function so that values
do not need to get copied in to the final output array. do not need to get copied in to the final output array.
Args: Args:
alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used
in the wrapped function. in the wrapped function.
Returns: Returns:
:class:`numpy.ndarray`: The aggregated uvmet output array that includes :class:`numpy.ndarray`: The aggregated uvmet output array that includes
all extra leftmost dimensions. all extra leftmost dimensions.
""" """
@wrapt.decorator @wrapt.decorator
def func_wrapper(wrapped, instance, args, kwargs): def func_wrapper(wrapped, instance, args, kwargs):
@ -43,64 +43,62 @@ def uvmet_left_iter(alg_dtype=np.float64):
v = args[1] v = args[1]
lat = args[2] lat = args[2]
lon = args[3] lon = args[3]
cen_long = args[4] cen_long = args[4]
cone = args[5] cone = args[5]
orig_dtype = u.dtype orig_dtype = u.dtype
lat_lon_fixed = False lat_lon_fixed = False
if lat.ndim == 2: if lat.ndim == 2:
lat_lon_fixed = True lat_lon_fixed = True
if lon.ndim == 2 and not lat_lon_fixed: if lon.ndim == 2 and not lat_lon_fixed:
raise ValueError("'lat' and 'lon' shape mismatch") raise ValueError("'lat' and 'lon' shape mismatch")
num_left_dims_u = u.ndim - 2 num_left_dims_u = u.ndim - 2
num_left_dims_lat = lat.ndim - 2 num_left_dims_lat = lat.ndim - 2
if (num_left_dims_lat > num_left_dims_u): if (num_left_dims_lat > num_left_dims_u):
raise ValueError("number of 'lat' dimensions is greater than 'u'") raise ValueError("number of 'lat' dimensions is greater than 'u'")
if lat_lon_fixed: if lat_lon_fixed:
mode = 0 # fixed lat/lon mode = 0 # fixed lat/lon
else: else:
if num_left_dims_u == num_left_dims_lat: if num_left_dims_u == num_left_dims_lat:
mode = 1 # lat/lon same as u mode = 1 # lat/lon same as u
else: else:
mode = 2 # probably 3D with 2D lat/lon plus Time mode = 2 # probably 3D with 2D lat/lon plus Time
has_missing = False has_missing = False
u_arr = to_np(u) u_arr = to_np(u)
v_arr = to_np(v) v_arr = to_np(v)
umissing = default_fill(np.float64) umissing = default_fill(np.float64)
if isinstance(u_arr, np.ma.MaskedArray): if isinstance(u_arr, np.ma.MaskedArray):
has_missing = True has_missing = True
umissing = u_arr.fill_value umissing = u_arr.fill_value
vmissing = default_fill(np.float64) vmissing = default_fill(np.float64)
if isinstance(v_arr, np.ma.MaskedArray): if isinstance(v_arr, np.ma.MaskedArray):
has_missing = True has_missing = True
vmissing = v_arr.fill_value vmissing = v_arr.fill_value
uvmetmissing = umissing uvmetmissing = umissing
is_stag = 0 is_stag = 0
if (u.shape[-1] != lat.shape[-1] or u.shape[-2] != lat.shape[-2]): if (u.shape[-1] != lat.shape[-1] or u.shape[-2] != lat.shape[-2]):
is_stag = 1 is_stag = 1
# Sanity check # Sanity check
if (v.shape[-1] == lat.shape[-1] or v.shape[-2] == lat.shape[-2]): if (v.shape[-1] == lat.shape[-1] or v.shape[-2] == lat.shape[-2]):
raise ValueError("u is staggered but v is not") raise ValueError("u is staggered but v is not")
if (v.shape[-1] != lat.shape[-1] or v.shape[-2] != lat.shape[-2]): if (v.shape[-1] != lat.shape[-1] or v.shape[-2] != lat.shape[-2]):
is_stag = 1 is_stag = 1
# Sanity check # Sanity check
if (u.shape[-1] == lat.shape[-1] or u.shape[-2] == lat.shape[-2]): if (u.shape[-1] == lat.shape[-1] or u.shape[-2] == lat.shape[-2]):
raise ValueError("v is staggered but u is not") raise ValueError("v is staggered but u is not")
# No special left side iteration, return the function result # No special left side iteration, return the function result
if (num_left_dims_u == 0): if (num_left_dims_u == 0):
return wrapped(u, v, lat, lon, cen_long, cone, isstag=is_stag, return wrapped(u, v, lat, lon, cen_long, cone, isstag=is_stag,
@ -109,44 +107,43 @@ def uvmet_left_iter(alg_dtype=np.float64):
# Initial output is time,nz,2,ny,nx to create contiguous views # Initial output is time,nz,2,ny,nx to create contiguous views
outdims = u.shape[0:num_left_dims_u] outdims = u.shape[0:num_left_dims_u]
extra_dims = tuple(outdims) # Copy the left-most dims for iteration extra_dims = tuple(outdims) # Copy the left-most dims for iteration
outdims += (2,) outdims += (2,)
outdims += lat.shape[-2:] outdims += lat.shape[-2:]
outview_array = np.empty(outdims, alg_dtype) outview_array = np.empty(outdims, alg_dtype)
# Final Output moves the u_v dimension to left side # Final Output moves the u_v dimension to left side
output_dims = (2,) output_dims = (2,)
output_dims += extra_dims output_dims += extra_dims
output_dims += lat.shape[-2:] output_dims += lat.shape[-2:]
output = np.empty(output_dims, orig_dtype) output = np.empty(output_dims, orig_dtype)
for left_idxs in iter_left_indexes(extra_dims): for left_idxs in iter_left_indexes(extra_dims):
left_and_slice_idxs = left_idxs + (slice(None),) left_and_slice_idxs = left_idxs + (slice(None),)
if mode == 0: if mode == 0:
lat_left_and_slice = (slice(None),) lat_left_and_slice = (slice(None),)
elif mode == 1: elif mode == 1:
lat_left_and_slice = left_and_slice_idxs lat_left_and_slice = left_and_slice_idxs
elif mode == 2: elif mode == 2:
# Only need the left-most # Only need the left-most
lat_left_and_slice = tuple(left_idx lat_left_and_slice = tuple(
for left_idx in left_idxs[0:num_left_dims_lat]) left_idx for left_idx in left_idxs[0:num_left_dims_lat])
u_output_idxs = (0,) + left_idxs + (slice(None),) u_output_idxs = (0,) + left_idxs + (slice(None),)
v_output_idxs = (1,) + left_idxs + (slice(None),) v_output_idxs = (1,) + left_idxs + (slice(None),)
u_view_idxs = left_idxs + (0, slice(None)) u_view_idxs = left_idxs + (0, slice(None))
v_view_idxs = left_idxs + (1, slice(None)) v_view_idxs = left_idxs + (1, slice(None))
new_u = u[left_and_slice_idxs] new_u = u[left_and_slice_idxs]
new_v = v[left_and_slice_idxs] new_v = v[left_and_slice_idxs]
new_lat = lat[lat_left_and_slice] new_lat = lat[lat_left_and_slice]
new_lon = lon[lat_left_and_slice] new_lon = lon[lat_left_and_slice]
outview = outview_array[left_and_slice_idxs] outview = outview_array[left_and_slice_idxs]
# Skip the possible empty/missing arrays for the join method # Skip the possible empty/missing arrays for the join method
skip_missing = False skip_missing = False
for arg in (new_u, new_v, new_lat, new_lon): for arg in (new_u, new_v, new_lat, new_lon):
@ -154,70 +151,69 @@ def uvmet_left_iter(alg_dtype=np.float64):
if arg.mask.all(): if arg.mask.all():
output[u_output_idxs] = uvmetmissing output[u_output_idxs] = uvmetmissing
output[v_output_idxs] = uvmetmissing output[v_output_idxs] = uvmetmissing
skip_missing = True skip_missing = True
has_missing = True has_missing = True
if skip_missing: if skip_missing:
continue continue
# Call the numerical routine # Call the numerical routine
result = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone, result = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone,
isstag=is_stag, has_missing=has_missing, isstag=is_stag, has_missing=has_missing,
umissing=umissing, vmissing=vmissing, umissing=umissing, vmissing=vmissing,
uvmetmissing=uvmetmissing, outview=outview) uvmetmissing=uvmetmissing, outview=outview)
# Make sure the result is the same data as what got passed in # Make sure the result is the same data as what got passed in
# Can delete this once everything works # Can delete this once everything works
if (result.__array_interface__["data"][0] != if (result.__array_interface__["data"][0] !=
outview.__array_interface__["data"][0]): outview.__array_interface__["data"][0]):
raise RuntimeError("output array was copied") raise RuntimeError("output array was copied")
output[u_output_idxs] = ( output[u_output_idxs] = (
outview_array[u_view_idxs].astype(orig_dtype)) outview_array[u_view_idxs].astype(orig_dtype))
output[v_output_idxs] = ( output[v_output_idxs] = (
outview_array[v_view_idxs].astype(orig_dtype)) outview_array[v_view_idxs].astype(orig_dtype))
if has_missing: if has_missing:
output = np.ma.masked_values(output, uvmetmissing) output = np.ma.masked_values(output, uvmetmissing)
return output return output
return func_wrapper return func_wrapper
def cape_left_iter(alg_dtype=np.float64): def cape_left_iter(alg_dtype=np.float64):
"""A decorator to handle iterating over the leftmost dimensions for the """A decorator to handle iterating over the leftmost dimensions for the
cape diagnostic. cape diagnostic.
For example, if a wrapped function works with three-dimensional arrays, but For example, if a wrapped function works with three-dimensional arrays, but
the variables include a 4th leftmost dimension for 'Time', this decorator the variables include a 4th leftmost dimension for 'Time', this decorator
will iterate over all times, call the 3D Fortran routine, and aggregate the will iterate over all times, call the 3D Fortran routine, and aggregate the
results in to a 4D output array. results in to a 4D output array.
It is also important to note that the final output array is allocated It is also important to note that the final output array is allocated
first, and then views are passed to the wrapped function so that values first, and then views are passed to the wrapped function so that values
do not need to get copied in to the final output array. do not need to get copied in to the final output array.
Args: Args:
alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used
in the wrapped function. in the wrapped function.
Returns: Returns:
:class:`numpy.ndarray`: The aggregated cape output array that includes :class:`numpy.ndarray`: The aggregated cape output array that includes
all extra leftmost dimensions. all extra leftmost dimensions.
""" """
@wrapt.decorator @wrapt.decorator
def func_wrapper(wrapped, instance, args, kwargs): def func_wrapper(wrapped, instance, args, kwargs):
# The cape calculations use an ascending vertical pressure coordinate # The cape calculations use an ascending vertical pressure coordinate
new_args = list(args) new_args = list(args)
new_kwargs = dict(kwargs) new_kwargs = dict(kwargs)
p_hpa = args[0] p_hpa = args[0]
tk = args[1] tk = args[1]
qv = args[2] qv = args[2]
@ -227,19 +223,19 @@ def cape_left_iter(alg_dtype=np.float64):
missing = args[6] missing = args[6]
i3dflag = args[7] i3dflag = args[7]
ter_follow = args[8] ter_follow = args[8]
is2d = i3dflag == 0 is2d = i3dflag == 0
# Note: This should still work with DataArrays # Note: This should still work with DataArrays
is1d = np.isscalar(sfp) or np.size(sfp) == 1 is1d = np.isscalar(sfp) or np.size(sfp) == 1
# Make sure sfp and terrain are regular floats for 1D case # Make sure sfp and terrain are regular floats for 1D case
# This should also work with DataArrays # This should also work with DataArrays
if is1d: if is1d:
ter = float(ter) ter = float(ter)
sfp = float(sfp) sfp = float(sfp)
orig_dtype = p_hpa.dtype orig_dtype = p_hpa.dtype
if not is1d: if not is1d:
# Need to order in ascending pressure order # Need to order in ascending pressure order
flip = False flip = False
@ -247,48 +243,48 @@ def cape_left_iter(alg_dtype=np.float64):
top_idxs = list(bot_idxs) top_idxs = list(bot_idxs)
top_idxs[-3] = -1 top_idxs[-3] = -1
top_idxs = tuple(top_idxs) top_idxs = tuple(top_idxs)
if p_hpa[bot_idxs] > p_hpa[top_idxs]: if p_hpa[bot_idxs] > p_hpa[top_idxs]:
flip = True flip = True
p_hpa = np.ascontiguousarray(p_hpa[...,::-1,:,:]) p_hpa = np.ascontiguousarray(p_hpa[..., ::-1, :, :])
tk = np.ascontiguousarray(tk[...,::-1,:,:]) tk = np.ascontiguousarray(tk[..., ::-1, :, :])
qv = np.ascontiguousarray(qv[...,::-1,:,:]) qv = np.ascontiguousarray(qv[..., ::-1, :, :])
ht = np.ascontiguousarray(ht[...,::-1,:,:]) ht = np.ascontiguousarray(ht[..., ::-1, :, :])
new_args[0] = p_hpa new_args[0] = p_hpa
new_args[1] = tk new_args[1] = tk
new_args[2] = qv new_args[2] = qv
new_args[3] = ht new_args[3] = ht
num_left_dims = p_hpa.ndim - 3 num_left_dims = p_hpa.ndim - 3
else: else:
# Need to order in ascending pressure order # Need to order in ascending pressure order
flip = False flip = False
if p_hpa[0] > p_hpa[-1]: if p_hpa[0] > p_hpa[-1]:
flip = True flip = True
p_hpa = np.ascontiguousarray(p_hpa[::-1]) p_hpa = np.ascontiguousarray(p_hpa[::-1])
tk = np.ascontiguousarray(tk[::-1]) tk = np.ascontiguousarray(tk[::-1])
qv = np.ascontiguousarray(qv[::-1]) qv = np.ascontiguousarray(qv[::-1])
ht = np.ascontiguousarray(ht[::-1]) ht = np.ascontiguousarray(ht[::-1])
# Need to make 3D views for the fortran code. # Need to make 3D views for the fortran code.
# Going to make these fortran ordered, since the f_contiguous and # Going to make these fortran ordered, since the f_contiguous and
# c_contiguous flags are broken in numpy 1.11 (always false). This # c_contiguous flags are broken in numpy 1.11 (always false). This
# should work across all numpy versions. # should work across all numpy versions.
new_args[0] = p_hpa.reshape((1, 1, p_hpa.shape[0]), order='F') new_args[0] = p_hpa.reshape((1, 1, p_hpa.shape[0]), order='F')
new_args[1] = tk.reshape((1, 1, tk.shape[0]), order='F') new_args[1] = tk.reshape((1, 1, tk.shape[0]), order='F')
new_args[2] = qv.reshape((1, 1, qv.shape[0]), order='F') new_args[2] = qv.reshape((1, 1, qv.shape[0]), order='F')
new_args[3] = ht.reshape((1, 1, ht.shape[0]), order='F') new_args[3] = ht.reshape((1, 1, ht.shape[0]), order='F')
new_args[4] = np.full((1,1), ter, orig_dtype) new_args[4] = np.full((1, 1), ter, orig_dtype)
new_args[5] = np.full((1,1), sfp, orig_dtype) new_args[5] = np.full((1, 1), sfp, orig_dtype)
num_left_dims = 0 num_left_dims = 0
# No special left side iteration, build the output from the cape,cin # No special left side iteration, build the output from the cape,cin
# result # result
if (num_left_dims == 0): if (num_left_dims == 0):
cape, cin = wrapped(*new_args, **new_kwargs) cape, cin = wrapped(*new_args, **new_kwargs)
output_dims = (2,) output_dims = (2,)
if not is1d: if not is1d:
output_dims += p_hpa.shape[-3:] output_dims += p_hpa.shape[-3:]
@ -296,26 +292,26 @@ def cape_left_iter(alg_dtype=np.float64):
output_dims += (p_hpa.shape[0], 1, 1) output_dims += (p_hpa.shape[0], 1, 1)
output = np.empty(output_dims, orig_dtype) output = np.empty(output_dims, orig_dtype)
if flip and not is2d: if flip and not is2d:
output[0,:] = cape[::-1,:,:] output[0, :] = cape[::-1, :, :]
output[1,:] = cin[::-1,:,:] output[1, :] = cin[::-1, :, :]
else: else:
output[0,:] = cape[:] output[0, :] = cape[:]
output[1,:] = cin[:] output[1, :] = cin[:]
return output return output
# Initial output is ...,cape_cin,nz,ny,nx to create contiguous views # Initial output is ...,cape_cin,nz,ny,nx to create contiguous views
outdims = p_hpa.shape[0:num_left_dims] outdims = p_hpa.shape[0:num_left_dims]
extra_dims = tuple(outdims) # Copy the left-most dims for iteration extra_dims = tuple(outdims) # Copy the left-most dims for iteration
outdims += (2,) # cape_cin outdims += (2,) # cape_cin
outdims += p_hpa.shape[-3:] outdims += p_hpa.shape[-3:]
outview_array = np.empty(outdims, alg_dtype) outview_array = np.empty(outdims, alg_dtype)
# Create the output array where the leftmost dim is the product type # Create the output array where the leftmost dim is the product type
output_dims = (2,) output_dims = (2,)
output_dims += extra_dims output_dims += extra_dims
@ -326,14 +322,14 @@ def cape_left_iter(alg_dtype=np.float64):
left_and_slice_idxs = left_idxs + (slice(None),) left_and_slice_idxs = left_idxs + (slice(None),)
cape_idxs = left_idxs + (0, slice(None)) cape_idxs = left_idxs + (0, slice(None))
cin_idxs = left_idxs + (1, slice(None)) cin_idxs = left_idxs + (1, slice(None))
cape_output_idxs = (0,) + left_idxs + (slice(None),) cape_output_idxs = (0,) + left_idxs + (slice(None),)
cin_output_idxs = (1,) + left_idxs + (slice(None),) cin_output_idxs = (1,) + left_idxs + (slice(None),)
view_cape_reverse_idxs = left_idxs + (0, slice(None,None,-1), view_cape_reverse_idxs = left_idxs + (0, slice(None, None, -1),
slice(None)) slice(None))
view_cin_reverse_idxs = left_idxs + (1, slice(None,None,-1), view_cin_reverse_idxs = left_idxs + (1, slice(None, None, -1),
slice(None)) slice(None))
new_args[0] = p_hpa[left_and_slice_idxs] new_args[0] = p_hpa[left_and_slice_idxs]
new_args[1] = tk[left_and_slice_idxs] new_args[1] = tk[left_and_slice_idxs]
new_args[2] = qv[left_and_slice_idxs] new_args[2] = qv[left_and_slice_idxs]
@ -342,9 +338,9 @@ def cape_left_iter(alg_dtype=np.float64):
new_args[5] = sfp[left_and_slice_idxs] new_args[5] = sfp[left_and_slice_idxs]
capeview = outview_array[cape_idxs] capeview = outview_array[cape_idxs]
cinview = outview_array[cin_idxs] cinview = outview_array[cin_idxs]
# Skip the possible empty/missing arrays for the join method # Skip the possible empty/missing arrays for the join method
# Note: Masking handled by cape.py or computation.py, so only # Note: Masking handled by cape.py or computation.py, so only
# supply the fill values here. # supply the fill values here.
skip_missing = False skip_missing = False
for arg in (new_args[0:6]): for arg in (new_args[0:6]):
@ -356,25 +352,24 @@ def cape_left_iter(alg_dtype=np.float64):
else: else:
output[cape_output_idxs] = missing output[cape_output_idxs] = missing
output[cin_output_idxs] = missing output[cin_output_idxs] = missing
skip_missing = True skip_missing = True
if skip_missing: if skip_missing:
continue continue
# Call the numerical routine # Call the numerical routine
new_kwargs["capeview"] = capeview new_kwargs["capeview"] = capeview
new_kwargs["cinview"] = cinview new_kwargs["cinview"] = cinview
cape, cin = wrapped(*new_args, **new_kwargs) cape, cin = wrapped(*new_args, **new_kwargs)
# Make sure the result is the same data as what got passed in # Make sure the result is the same data as what got passed in
# Can delete this once everything works # Can delete this once everything works
if (cape.__array_interface__["data"][0] != if (cape.__array_interface__["data"][0] !=
capeview.__array_interface__["data"][0]): capeview.__array_interface__["data"][0]):
raise RuntimeError("output array was copied") raise RuntimeError("output array was copied")
if flip and not is2d: if flip and not is2d:
output[cape_output_idxs] = ( output[cape_output_idxs] = (
outview_array[view_cape_reverse_idxs].astype(orig_dtype)) outview_array[view_cape_reverse_idxs].astype(orig_dtype))
@ -382,82 +377,82 @@ def cape_left_iter(alg_dtype=np.float64):
outview_array[view_cin_reverse_idxs].astype(orig_dtype)) outview_array[view_cin_reverse_idxs].astype(orig_dtype))
else: else:
output[cape_output_idxs] = ( output[cape_output_idxs] = (
outview_array[cape_idxs].astype(orig_dtype)) outview_array[cape_idxs].astype(orig_dtype))
output[cin_output_idxs] = ( output[cin_output_idxs] = (
outview_array[cin_idxs].astype(orig_dtype)) outview_array[cin_idxs].astype(orig_dtype))
return output return output
return func_wrapper return func_wrapper
def cloudfrac_left_iter(alg_dtype=np.float64): def cloudfrac_left_iter(alg_dtype=np.float64):
"""A decorator to handle iterating over the leftmost dimensions for the """A decorator to handle iterating over the leftmost dimensions for the
cloud fraction diagnostic. cloud fraction diagnostic.
For example, if a wrapped function works with three-dimensional arrays, but For example, if a wrapped function works with three-dimensional arrays, but
the variables include a 4th leftmost dimension for 'Time', this decorator the variables include a 4th leftmost dimension for 'Time', this decorator
will iterate over all times, call the 3D Fortran routine, and aggregate the will iterate over all times, call the 3D Fortran routine, and aggregate the
results in to a 4D output array. results in to a 4D output array.
It is also important to note that the final output array is allocated It is also important to note that the final output array is allocated
first, and then views are passed to the wrapped function so that values first, and then views are passed to the wrapped function so that values
do not need to get copied in to the final output array. do not need to get copied in to the final output array.
Args: Args:
alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used
in the wrapped function. in the wrapped function.
Returns: Returns:
:class:`numpy.ndarray`: The aggregated cloud fraction output array :class:`numpy.ndarray`: The aggregated cloud fraction output array
that includes all extra leftmost dimensions. that includes all extra leftmost dimensions.
""" """
@wrapt.decorator @wrapt.decorator
def func_wrapper(wrapped, instance, args, kwargs): def func_wrapper(wrapped, instance, args, kwargs):
new_args = list(args) new_args = list(args)
new_kwargs = dict(kwargs) new_kwargs = dict(kwargs)
vert = args[0] vert = args[0]
rh = args[1] rh = args[1]
num_left_dims = vert.ndim - 3 num_left_dims = vert.ndim - 3
orig_dtype = vert.dtype orig_dtype = vert.dtype
# No special left side iteration, build the output from the # No special left side iteration, build the output from the
# low, mid, high results. # low, mid, high results.
if (num_left_dims == 0): if (num_left_dims == 0):
low, mid, high = wrapped(*new_args, **new_kwargs) low, mid, high = wrapped(*new_args, **new_kwargs)
output_dims = (3,) output_dims = (3,)
output_dims += vert.shape[-2:] output_dims += vert.shape[-2:]
output = np.empty(output_dims, orig_dtype) output = np.empty(output_dims, orig_dtype)
output[0,:] = low[:] output[0, :] = low[:]
output[1,:] = mid[:] output[1, :] = mid[:]
output[2,:] = high[:] output[2, :] = high[:]
return output return output
# Initial output is ...,low_mid_high,nz,ny,nx to create contiguous views # Initial output is ...,low_mid_high,nz,ny,nx to create contiguous
# views
outdims = vert.shape[0:num_left_dims] outdims = vert.shape[0:num_left_dims]
extra_dims = tuple(outdims) # Copy the left-most dims for iteration extra_dims = tuple(outdims) # Copy the left-most dims for iteration
outdims += (3,) # low_mid_high outdims += (3,) # low_mid_high
outdims += vert.shape[-2:] outdims += vert.shape[-2:]
outview_array = np.empty(outdims, alg_dtype) outview_array = np.empty(outdims, alg_dtype)
# Create the output array where the leftmost dim is the cloud type # Create the output array where the leftmost dim is the cloud type
output_dims = (3,) output_dims = (3,)
output_dims += extra_dims output_dims += extra_dims
output_dims += vert.shape[-2:] output_dims += vert.shape[-2:]
output = np.empty(output_dims, orig_dtype) output = np.empty(output_dims, orig_dtype)
has_missing = False has_missing = False
missing = default_fill(np.float64) missing = default_fill(np.float64)
for left_idxs in iter_left_indexes(extra_dims): for left_idxs in iter_left_indexes(extra_dims):
@ -465,16 +460,16 @@ def cloudfrac_left_iter(alg_dtype=np.float64):
low_idxs = left_idxs + (0, slice(None)) low_idxs = left_idxs + (0, slice(None))
mid_idxs = left_idxs + (1, slice(None)) mid_idxs = left_idxs + (1, slice(None))
high_idxs = left_idxs + (2, slice(None)) high_idxs = left_idxs + (2, slice(None))
low_output_idxs = (0,) + left_idxs + (slice(None),) low_output_idxs = (0,) + left_idxs + (slice(None),)
mid_output_idxs = (1,) + left_idxs + (slice(None),) mid_output_idxs = (1,) + left_idxs + (slice(None),)
high_output_idxs = (2,) + left_idxs + (slice(None),) high_output_idxs = (2,) + left_idxs + (slice(None),)
new_args[0] = vert[left_and_slice_idxs] new_args[0] = vert[left_and_slice_idxs]
new_args[1] = rh[left_and_slice_idxs] new_args[1] = rh[left_and_slice_idxs]
# Skip the possible empty/missing arrays for the join method # Skip the possible empty/missing arrays for the join method
# Note: Masking handled by cloudfrac.py or computation.py, so only # Note: Masking handled by cloudfrac.py or computation.py, so only
# supply the fill values here. # supply the fill values here.
skip_missing = False skip_missing = False
for arg in (new_args[0:2]): for arg in (new_args[0:2]):
@ -483,41 +478,41 @@ def cloudfrac_left_iter(alg_dtype=np.float64):
output[low_output_idxs] = missing output[low_output_idxs] = missing
output[mid_output_idxs] = missing output[mid_output_idxs] = missing
output[high_output_idxs] = missing output[high_output_idxs] = missing
skip_missing = True skip_missing = True
has_missing = True has_missing = True
if skip_missing: if skip_missing:
continue continue
lowview = outview_array[low_idxs] lowview = outview_array[low_idxs]
midview = outview_array[mid_idxs] midview = outview_array[mid_idxs]
highview = outview_array[high_idxs] highview = outview_array[high_idxs]
new_kwargs["lowview"] = lowview new_kwargs["lowview"] = lowview
new_kwargs["midview"] = midview new_kwargs["midview"] = midview
new_kwargs["highview"] = highview new_kwargs["highview"] = highview
low, mid, high = wrapped(*new_args, **new_kwargs) low, mid, high = wrapped(*new_args, **new_kwargs)
# Make sure the result is the same data as what got passed in # Make sure the result is the same data as what got passed in
# Can delete this once everything works # Can delete this once everything works
if (low.__array_interface__["data"][0] != if (low.__array_interface__["data"][0] !=
lowview.__array_interface__["data"][0]): lowview.__array_interface__["data"][0]):
raise RuntimeError("output array was copied") raise RuntimeError("output array was copied")
output[low_output_idxs] = ( output[low_output_idxs] = (
outview_array[low_idxs].astype(orig_dtype)) outview_array[low_idxs].astype(orig_dtype))
output[mid_output_idxs] = ( output[mid_output_idxs] = (
outview_array[mid_idxs].astype(orig_dtype)) outview_array[mid_idxs].astype(orig_dtype))
output[high_output_idxs] = ( output[high_output_idxs] = (
outview_array[high_idxs].astype(orig_dtype)) outview_array[high_idxs].astype(orig_dtype))
if has_missing: if has_missing:
output = np.ma.masked_values(output, missing) output = np.ma.masked_values(output, missing)
return output return output
return func_wrapper return func_wrapper
@ -526,95 +521,94 @@ def interplevel_left_iter(is2dlev, alg_dtype=np.float64):
def func_wrapper(wrapped, instance, args, kwargs): def func_wrapper(wrapped, instance, args, kwargs):
new_args = list(args) new_args = list(args)
new_kwargs = dict(kwargs) new_kwargs = dict(kwargs)
field3d = args[0] field3d = args[0]
z = args[1] z = args[1]
levels = args[2] levels = args[2]
num_left_dims = z.ndim - 3 num_left_dims = z.ndim - 3
orig_dtype = field3d.dtype orig_dtype = field3d.dtype
left_dims = z.shape[0:num_left_dims] left_dims = z.shape[0:num_left_dims]
multiproduct = True if field3d.ndim - z.ndim == 1 else False multiproduct = True if field3d.ndim - z.ndim == 1 else False
# No special left side iteration, build the output from the # No special left side iteration, build the output from the
# low, mid, high results. # low, mid, high results.
if (num_left_dims == 0): if (num_left_dims == 0):
if multiproduct: if multiproduct:
if not is2dlev: if not is2dlev:
outshape = (field3d.shape[0:-3] + levels.shape + outshape = (field3d.shape[0:-3] + levels.shape +
field3d.shape[-2:]) field3d.shape[-2:])
else: else:
outshape = (field3d.shape[0:-3] + field3d.shape[-2:]) outshape = (field3d.shape[0:-3] + field3d.shape[-2:])
output = np.empty(outshape, dtype=alg_dtype) output = np.empty(outshape, dtype=alg_dtype)
for i in py3range(field3d.shape[0]): for i in py3range(field3d.shape[0]):
new_args[0] = field3d[i,:] new_args[0] = field3d[i, :]
new_kwargs["outview"] = output[i,:] new_kwargs["outview"] = output[i, :]
_ = wrapped(*new_args, **new_kwargs) _ = wrapped(*new_args, **new_kwargs)
else: else:
output = wrapped(*args, **kwargs) output = wrapped(*args, **kwargs)
return output return output
if multiproduct: if multiproduct:
outdims = field3d.shape[0:1] + left_dims outdims = field3d.shape[0:1] + left_dims
else: else:
outdims = left_dims outdims = left_dims
extra_dims = tuple(outdims) extra_dims = tuple(outdims)
if not is2dlev: if not is2dlev:
outdims += levels.shape outdims += levels.shape
outdims += z.shape[-2:] outdims += z.shape[-2:]
outview_array = np.empty(outdims, alg_dtype) outview_array = np.empty(outdims, alg_dtype)
for left_idxs in iter_left_indexes(extra_dims): for left_idxs in iter_left_indexes(extra_dims):
field_out_slice_idxs = left_idxs + (slice(None),) field_out_slice_idxs = left_idxs + (slice(None),)
if multiproduct: if multiproduct:
z_slice_idxs = left_idxs[1:] + (slice(None),) z_slice_idxs = left_idxs[1:] + (slice(None),)
else: else:
z_slice_idxs = left_idxs + (slice(None),) z_slice_idxs = left_idxs + (slice(None),)
new_args[0] = field3d[field_out_slice_idxs] new_args[0] = field3d[field_out_slice_idxs]
new_args[1] = z[z_slice_idxs] new_args[1] = z[z_slice_idxs]
if is2dlev: if is2dlev:
if levels.ndim > 2: if levels.ndim > 2:
new_args[2] = levels[z_slice_idxs] new_args[2] = levels[z_slice_idxs]
new_kwargs["outview"] = outview_array[field_out_slice_idxs] new_kwargs["outview"] = outview_array[field_out_slice_idxs]
_ = wrapped(*new_args, **new_kwargs) _ = wrapped(*new_args, **new_kwargs)
output = outview_array.astype(orig_dtype) output = outview_array.astype(orig_dtype)
return output return output
return func_wrapper return func_wrapper
def check_cape_args(): def check_cape_args():
"""A decorator to check that the cape_3d arguments are valid. """A decorator to check that the cape_3d arguments are valid.
An exception is raised when an invalid argument is found. An exception is raised when an invalid argument is found.
Returns: Returns:
None None
Raises: Raises:
:class:`ValueError`: Raised when an invalid argument is detected. :class:`ValueError`: Raised when an invalid argument is detected.
""" """
@wrapt.decorator @wrapt.decorator
def func_wrapper(wrapped, instance, args, kwargs): def func_wrapper(wrapped, instance, args, kwargs):
p_hpa = args[0] p_hpa = args[0]
tk = args[1] tk = args[1]
qv = args[2] qv = args[2]
@ -624,11 +618,11 @@ def check_cape_args():
missing = args[6] missing = args[6]
i3dflag = args[7] i3dflag = args[7]
ter_follow = args[8] ter_follow = args[8]
is2d = False if i3dflag != 0 else True is2d = False if i3dflag != 0 else True
is1d = ((np.isscalar(sfp) or np.size(sfp) == 1) or is1d = ((np.isscalar(sfp) or np.size(sfp) == 1) or
(np.isscalar(ter) or np.size(ter) == 1)) (np.isscalar(ter) or np.size(ter) == 1))
if not (p_hpa.shape == tk.shape == qv.shape == ht.shape): if not (p_hpa.shape == tk.shape == qv.shape == ht.shape):
raise ValueError("arguments 0, 1, 2, 3 must be the same shape") raise ValueError("arguments 0, 1, 2, 3 must be the same shape")
@ -641,41 +635,41 @@ def check_cape_args():
if np.size(ter) != np.size(sfp): if np.size(ter) != np.size(sfp):
raise ValueError("arguments 4 and 5 must both be scalars or " raise ValueError("arguments 4 and 5 must both be scalars or "
"both be arrays") "both be arrays")
# Only need to test p_hpa since we assured args 0-3 have same ndim # Only need to test p_hpa since we assured args 0-3 have same ndim
if p_hpa.ndim != 1: if p_hpa.ndim != 1:
raise ValueError("arguments 0-3 " raise ValueError("arguments 0-3 "
"must be 1-dimensional when " "must be 1-dimensional when "
"arguments 4 and 5 are scalars") "arguments 4 and 5 are scalars")
return wrapped(*args, **kwargs) return wrapped(*args, **kwargs)
return func_wrapper return func_wrapper
def check_interplevel_args(is2dlev): def check_interplevel_args(is2dlev):
"""A decorator to check that the interplevel arguments are valid. """A decorator to check that the interplevel arguments are valid.
An exception is raised when an invalid argument is found. An exception is raised when an invalid argument is found.
Returns: Returns:
None None
Raises: Raises:
:class:`ValueError`: Raised when an invalid argument is detected. :class:`ValueError`: Raised when an invalid argument is detected.
""" """
@wrapt.decorator @wrapt.decorator
def func_wrapper(wrapped, instance, args, kwargs): def func_wrapper(wrapped, instance, args, kwargs):
field3d = args[0] field3d = args[0]
z = args[1] z = args[1]
levels = args[2] levels = args[2]
multiproduct = True if (field3d.ndim - z.ndim) == 1 else False multiproduct = True if (field3d.ndim - z.ndim) == 1 else False
if not multiproduct: if not multiproduct:
if field3d.shape != z.shape: if field3d.shape != z.shape:
raise ValueError("arguments 0 and 1 must have the same shape") raise ValueError("arguments 0 and 1 must have the same shape")
@ -683,16 +677,15 @@ def check_interplevel_args(is2dlev):
if field3d.shape[1:] != z.shape: if field3d.shape[1:] != z.shape:
raise ValueError("argument 0 and 1 must have same rightmost " raise ValueError("argument 0 and 1 must have same rightmost "
"dimensions") "dimensions")
if is2dlev: if is2dlev:
if levels.ndim != 2: if levels.ndim != 2:
if (levels.shape[0:-2] != z.shape[0:-3] or if (levels.shape[0:-2] != z.shape[0:-3] or
levels.shape[-2:] != z.shape[-2:]): levels.shape[-2:] != z.shape[-2:]):
raise ValueError("argument 1 and 2 must have " raise ValueError("argument 1 and 2 must have "
"the same leftmost and rightmost " "the same leftmost and rightmost "
"dimensions") "dimensions")
return wrapped(*args, **kwargs) return wrapped(*args, **kwargs)
return func_wrapper
return func_wrapper

440
src/wrf/units.py

@ -4,31 +4,31 @@ from .constants import Constants, ConversionFactors
def _apply_conv_fact(var, vartype, var_unit, dest_unit): def _apply_conv_fact(var, vartype, var_unit, dest_unit):
"""Return the variable converted to different units using a conversion """Return the variable converted to different units using a conversion
factor. factor.
Args: Args:
var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A
variable. variable.
vartype (:obj:`str`): The type of variable. Choices are: 'wind', vartype (:obj:`str`): The type of variable. Choices are: 'wind',
'pressure', 'temp', or 'height'. 'pressure', 'temp', or 'height'.
var_unit (:obj:`str`): The variable's current units. var_unit (:obj:`str`): The variable's current units.
dest_unit (:obj:`str`): The desired units. dest_unit (:obj:`str`): The desired units.
Returns: Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in :class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in
the desired units. the desired units.
""" """
if var_unit == dest_unit: if var_unit == dest_unit:
return var return var
# Note, case where var_unit and dest_unit are base unit, should be # Note, case where var_unit and dest_unit are base unit, should be
# handled above # handled above
if var_unit == _BASE_UNITS[vartype]: if var_unit == _BASE_UNITS[vartype]:
return var*(_CONV_FACTORS[vartype]["to_dest"][dest_unit]) return var*(_CONV_FACTORS[vartype]["to_dest"][dest_unit])
@ -36,62 +36,62 @@ def _apply_conv_fact(var, vartype, var_unit, dest_unit):
if dest_unit == _BASE_UNITS[vartype]: if dest_unit == _BASE_UNITS[vartype]:
return var*(_CONV_FACTORS[vartype]["to_base"][var_unit]) return var*(_CONV_FACTORS[vartype]["to_base"][var_unit])
else: else:
return var*(_CONV_FACTORS[vartype]["to_base"][var_unit] * return var*(_CONV_FACTORS[vartype]["to_base"][var_unit] *
_CONV_FACTORS[vartype]["to_dest"][dest_unit]) _CONV_FACTORS[vartype]["to_dest"][dest_unit])
def _to_kelvin(var, var_unit): def _to_kelvin(var, var_unit):
"""Return the variable in Kelvin. """Return the variable in Kelvin.
Args: Args:
var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A
variable. variable.
var_unit (:obj:`str`): The variable's current units. var_unit (:obj:`str`): The variable's current units.
Returns: Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in :class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in
Kelvin. Kelvin.
""" """
if var_unit == "c": if var_unit == "c":
return var + Constants.CELKEL return var + Constants.CELKEL
elif var_unit == "f": elif var_unit == "f":
return (var - 32.0) * (5.0/9.0) + Constants.CELKEL return (var - 32.0) * (5.0/9.0) + Constants.CELKEL
def _k_to_c(var): def _k_to_c(var):
"""Return the variable in Celsius. """Return the variable in Celsius.
Args: Args:
var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A
variable in units of Kelvin. variable in units of Kelvin.
Returns: Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in :class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in
Celsius. Celsius.
""" """
return var - Constants.CELKEL return var - Constants.CELKEL
def _k_to_f(var): def _k_to_f(var):
"""Return the variable in Fahrenheit. """Return the variable in Fahrenheit.
Args: Args:
var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A
variable in units of Kelvin. variable in units of Kelvin.
Returns: Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in :class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in
Fahrenheit. Fahrenheit.
""" """
return 1.8 * _k_to_c(var) + 32.0 return 1.8 * _k_to_c(var) + 32.0
@ -99,25 +99,25 @@ def _k_to_f(var):
def _apply_temp_conv(var, var_unit, dest_unit): def _apply_temp_conv(var, var_unit, dest_unit):
"""Return the variable converted to different units using a temperature """Return the variable converted to different units using a temperature
conversion algorithm. conversion algorithm.
Args: Args:
var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A
variable. variable.
var_unit (:obj:`str`): The variable's current units. var_unit (:obj:`str`): The variable's current units.
dest_unit (:obj:`str`): The desired units. dest_unit (:obj:`str`): The desired units.
Returns: Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in :class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in
the desired units. the desired units.
""" """
if dest_unit == var_unit: if dest_unit == var_unit:
return var return var
if var_unit != _BASE_UNITS["temp"]: if var_unit != _BASE_UNITS["temp"]:
tk = _to_kelvin(var, var_unit) tk = _to_kelvin(var, var_unit)
if dest_unit == _BASE_UNITS["temp"]: if dest_unit == _BASE_UNITS["temp"]:
@ -126,247 +126,237 @@ def _apply_temp_conv(var, var_unit, dest_unit):
return (_TEMP_CONV_METHODS[dest_unit])(tk) return (_TEMP_CONV_METHODS[dest_unit])(tk)
else: else:
return (_TEMP_CONV_METHODS[dest_unit])(var) return (_TEMP_CONV_METHODS[dest_unit])(var)
# A mapping of unit names to their dictionary key names # A mapping of unit names to their dictionary key names
_UNIT_ALIASES = {"mps" : "m s-1", _UNIT_ALIASES = {"mps": "m s-1",
"m/s" : "m s-1", "m/s": "m s-1",
"ms-1" : "m s-1", "ms-1": "m s-1",
"meters_per_second" : "m s-1", "meters_per_second": "m s-1",
"metres_per_second" : "m s-1", "metres_per_second": "m s-1",
"knots" : "kt", "knots": "kt",
"knot" : "kt", "knot": "kt",
"kts" : "kt", "kts": "kt",
"kn" : "kt", "kn": "kt",
"miles_per_hour" : "mi h-1", "miles_per_hour": "mi h-1",
"mih-1" : "mi h-1", "mih-1": "mi h-1",
"mph" : "mi h-1", "mph": "mi h-1",
"mi/h" : "mi h-1", "mi/h": "mi h-1",
"kmph" : "km h-1", "kmph": "km h-1",
"kmh-1" : "km h-1", "kmh-1": "km h-1",
"km/h" : "km h-1", "km/h": "km h-1",
"kilometers_per_hour" : "km h-1", "kilometers_per_hour": "km h-1",
"kilometres_per_hour" : "km h-1", "kilometres_per_hour": "km h-1",
"ft/s" : "ft s-1", "ft/s": "ft s-1",
"ft/sec" : "ft s-1", "ft/sec": "ft s-1",
"fps" : "ft s-1", "fps": "ft s-1",
"fs-1" : "ft s-1", "fs-1": "ft s-1",
"feet_per_second" : "ft s-1", "feet_per_second": "ft s-1",
"pascal": "pa",
"pascal" : "pa", "pascals": "pa",
"pascals" : "pa", "hecto_pascal": "hpa",
"hecto_pascal" : "hpa", "hecto_pascals": "hpa",
"hecto_pascals" : "hpa", "millibar": "mb",
"millibar" : "mb", "millibars": "mb",
"millibars" : "mb", "mbar": "mb",
"mbar" : "mb", "kelvin": "k",
"degree_kelvin": "k",
"kelvin" : "k", "degrees_kelvin": "k",
"degree_kelvin" : "k", "degree_k": "k",
"degrees_kelvin" : "k", "degrees_k": "k",
"degree_k" : "k", "degreek": "k",
"degrees_k" : "k", "degreesk": "k",
"degreek" : "k", "degk": "k",
"degreesk" : "k", "degsk": "k",
"degk" : "k", "deg_k": "k",
"degsk" : "k", "degs_k": "k",
"deg_k" : "k", "deg k": "k",
"degs_k" : "k", "degs k": "k",
"deg k" : "k", "celsius": "c",
"degs k" : "k", "degree_celsius": "c",
"degrees_celsius": "c",
"celsius" : "c", "degree_c": "c",
"degree_celsius" : "c", "degrees_c": "c",
"degrees_celsius" : "c", "degreec": "c",
"degree_c" : "c", "degreesc": "c",
"degrees_c" : "c", "degc": "c",
"degreec" : "c", "degsc": "c",
"degreesc" : "c", "deg_c": "c",
"degc" : "c", "degs_c": "c",
"degsc" : "c", "deg c": "c",
"deg_c" : "c", "degs c": "c",
"degs_c" : "c", "fahrenheit": "f",
"deg c" : "c", "degree_fahrenheit": "f",
"degs c" : "c", "degrees_fahrenheit": "f",
"degree_f": "f",
"fahrenheit" : "f", "degrees_f": "f",
"degree_fahrenheit" : "f", "degreef": "f",
"degrees_fahrenheit" : "f", "degreesf": "f",
"degree_f" : "f", "degf": "f",
"degrees_f" : "f", "degsf": "f",
"degreef" : "f", "deg_f": "f",
"degreesf" : "f", "degs_f": "f",
"degf" : "f", "deg f": "f",
"degsf" : "f", "degs f": "f",
"deg_f" : "f", "meter": "m",
"degs_f" : "f", "meters": "m",
"deg f" : "f", "metre": "m",
"degs f" : "f", "metres": "m",
"kilometer": "km",
"meter" : "m", "kilometers": "km",
"meters" : "m", "dekameter": "dm",
"metre" : "m", "dekameters": "dm",
"metres" : "m", "decameter": "dm",
"kilometer" : "km", "decameters": "dm",
"kilometers" : "km", "dekametre": "dm",
"dekameter" : "dm", "dekametres": "dm",
"dekameters" : "dm", "decametre": "dm",
"decameter" : "dm", "decametres": "dm",
"decameters" : "dm", "dam": "dm",
"dekametre" : "dm", "dkm": "dm",
"dekametres" : "dm", "feet": "ft",
"decametre" : "dm", "foot": "ft",
"decametres" : "dm", "mile": "mi",
"dam" : "dm", "miles": "mi"
"dkm" : "dm", }
"feet" : "ft",
"foot" : "ft",
"mile" : "mi",
"miles" : "mi"
}
# A mapping of unit types to the avaible units # A mapping of unit types to the avaible units
_VALID_UNITS = {"wind" : ["m s-1", "kt", "mi h-1", "km h-1", "ft s-1"], _VALID_UNITS = {"wind": ["m s-1", "kt", "mi h-1", "km h-1", "ft s-1"],
"pressure" : ["pa", "hpa", "mb", "torr", "mmhg", "atm"], "pressure": ["pa", "hpa", "mb", "torr", "mmhg", "atm"],
"temp" : ["k", "f", "c"], "temp": ["k", "f", "c"],
"height" : ["m", "km", "dm", "ft", "mi"] "height": ["m", "km", "dm", "ft", "mi"]
} }
# Conversion factor map for wind from base units # Conversion factor map for wind from base units
_WIND_BASE_FACTORS = {"kt" : ConversionFactors.MPS_TO_KTS, _WIND_BASE_FACTORS = {"kt": ConversionFactors.MPS_TO_KTS,
"km h-1" : ConversionFactors.MPS_TO_KMPH, "km h-1": ConversionFactors.MPS_TO_KMPH,
"mi h-1" : ConversionFactors.MPS_TO_MPH, "mi h-1": ConversionFactors.MPS_TO_MPH,
"ft s-1" : ConversionFactors.MPS_TO_FPS "ft s-1": ConversionFactors.MPS_TO_FPS
} }
# Conversion factor map to base units # Conversion factor map to base units
_WIND_TOBASE_FACTORS = {"kt" : 1.0/ConversionFactors.MPS_TO_KTS, _WIND_TOBASE_FACTORS = {"kt": 1.0/ConversionFactors.MPS_TO_KTS,
"km h-1" : 1.0/ConversionFactors.MPS_TO_KMPH, "km h-1": 1.0/ConversionFactors.MPS_TO_KMPH,
"mi h-1" : 1.0/ConversionFactors.MPS_TO_MPH, "mi h-1": 1.0/ConversionFactors.MPS_TO_MPH,
"ft s-1" : 1.0/ConversionFactors.MPS_TO_FPS "ft s-1": 1.0/ConversionFactors.MPS_TO_FPS
} }
# Conversion factor map for pressure from base units # Conversion factor map for pressure from base units
_PRES_BASE_FACTORS = {"hpa" : ConversionFactors.PA_TO_HPA, _PRES_BASE_FACTORS = {"hpa": ConversionFactors.PA_TO_HPA,
"mb" : ConversionFactors.PA_TO_HPA, "mb": ConversionFactors.PA_TO_HPA,
"torr" : ConversionFactors.PA_TO_TORR, "torr": ConversionFactors.PA_TO_TORR,
"mmhg" : ConversionFactors.PA_TO_MMHG, "mmhg": ConversionFactors.PA_TO_MMHG,
"atm" : ConversionFactors.PA_TO_ATM "atm": ConversionFactors.PA_TO_ATM
} }
# Conversion factor map for pressure to base units # Conversion factor map for pressure to base units
_PRES_TOBASE_FACTORS = {"hpa" : 1.0/ConversionFactors.PA_TO_HPA, _PRES_TOBASE_FACTORS = {"hpa": 1.0/ConversionFactors.PA_TO_HPA,
"mb" : 1.0/ConversionFactors.PA_TO_HPA, "mb": 1.0/ConversionFactors.PA_TO_HPA,
"torr" : 1.0/ConversionFactors.PA_TO_TORR, "torr": 1.0/ConversionFactors.PA_TO_TORR,
"mmhg" : 1.0/ConversionFactors.PA_TO_MMHG, "mmhg": 1.0/ConversionFactors.PA_TO_MMHG,
"atm" : 1.0/ConversionFactors.PA_TO_ATM "atm": 1.0/ConversionFactors.PA_TO_ATM
} }
# Conversion factor map for height from base units # Conversion factor map for height from base units
_HEIGHT_BASE_FACTORS = {"km" : ConversionFactors.M_TO_KM, _HEIGHT_BASE_FACTORS = {"km": ConversionFactors.M_TO_KM,
"dm" : ConversionFactors.M_TO_DM, "dm": ConversionFactors.M_TO_DM,
"ft" : ConversionFactors.M_TO_FT, "ft": ConversionFactors.M_TO_FT,
"mi" : ConversionFactors.M_TO_MILES "mi": ConversionFactors.M_TO_MILES
} }
# Conversion factor map for height to base units # Conversion factor map for height to base units
_HEIGHT_TOBASE_FACTORS = {"km" : 1.0/ConversionFactors.M_TO_KM, _HEIGHT_TOBASE_FACTORS = {"km": 1.0/ConversionFactors.M_TO_KM,
"dm" : 1.0/ConversionFactors.M_TO_DM, "dm": 1.0/ConversionFactors.M_TO_DM,
"ft" : 1.0/ConversionFactors.M_TO_FT, "ft": 1.0/ConversionFactors.M_TO_FT,
"mi" : 1.0/ConversionFactors.M_TO_MILES "mi": 1.0/ConversionFactors.M_TO_MILES
} }
# Mapping of unit type to base unit type # Mapping of unit type to base unit type
_BASE_UNITS = {"wind" : "m s-1", _BASE_UNITS = {"wind": "m s-1",
"pressure" : "pa", "pressure": "pa",
"temp" : "k", "temp": "k",
"height" : "m" "height": "m"
} }
# A mapping of unit type to a mapping of to/from base conversion factors # A mapping of unit type to a mapping of to/from base conversion factors
_CONV_FACTORS = {"wind" : {"to_dest" : _WIND_BASE_FACTORS, _CONV_FACTORS = {"wind": {"to_dest": _WIND_BASE_FACTORS,
"to_base" : _WIND_TOBASE_FACTORS}, "to_base": _WIND_TOBASE_FACTORS},
"pressure" : {"to_dest" : _PRES_BASE_FACTORS, "pressure": {"to_dest": _PRES_BASE_FACTORS,
"to_base" : _PRES_TOBASE_FACTORS}, "to_base": _PRES_TOBASE_FACTORS},
"height" : {"to_dest" : _HEIGHT_BASE_FACTORS, "height": {"to_dest": _HEIGHT_BASE_FACTORS,
"to_base" : _HEIGHT_TOBASE_FACTORS} "to_base": _HEIGHT_TOBASE_FACTORS}
} }
# A mapping of temperature type to the conversion function # A mapping of temperature type to the conversion function
_TEMP_CONV_METHODS = {"c" : _k_to_c, _TEMP_CONV_METHODS = {"c": _k_to_c,
"f" : _k_to_f "f": _k_to_f
} }
def dealias_and_clean_unit(unit): def dealias_and_clean_unit(unit):
"""Return the properly cleaned and dealiased unit name. """Return the properly cleaned and dealiased unit name.
Args: Args:
unit (:obj:`str`): The unit name. unit (:obj:`str`): The unit name.
Returns: Returns:
:obj:`str`: A unit name suitable for dictionary key lookups. :obj:`str`: A unit name suitable for dictionary key lookups.
""" """
cleaned_unit = " ".join(unit.lower().split()) cleaned_unit = " ".join(unit.lower().split())
dealiased = _UNIT_ALIASES.get(cleaned_unit, None) dealiased = _UNIT_ALIASES.get(cleaned_unit, None)
return cleaned_unit if dealiased is None else dealiased return cleaned_unit if dealiased is None else dealiased
def check_units(unit, unit_type): def check_units(unit, unit_type):
"""Raise an exception if the unit name is invalid. """Raise an exception if the unit name is invalid.
Args: Args:
unit (:obj:`str`): The unit name. unit (:obj:`str`): The unit name.
unit_type (:obj:`str`): The type of unit. unit_type (:obj:`str`): The type of unit.
Returns: Returns:
None None
Raises: Raises:
:class:`ValueError`: Raised when the unit name is invalid. :class:`ValueError`: Raised when the unit name is invalid.
""" """
u_cleaned = dealias_and_clean_unit(unit) u_cleaned = dealias_and_clean_unit(unit)
if u_cleaned not in _VALID_UNITS[unit_type]: if u_cleaned not in _VALID_UNITS[unit_type]:
raise ValueError("invalid unit type '%s'" % unit) raise ValueError("invalid unit type '{}'".format(unit))
def do_conversion(var, vartype, var_unit, dest_unit): def do_conversion(var, vartype, var_unit, dest_unit):
"""Return the variable converted to different units. """Return the variable converted to different units.
Args: Args:
var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A
variable. variable.
vartype (:obj:`str`): The type of variable. Choices are: 'wind', vartype (:obj:`str`): The type of variable. Choices are: 'wind',
'pressure', 'temp', or 'height'. 'pressure', 'temp', or 'height'.
var_unit (:obj:`str`): The variable's current units. var_unit (:obj:`str`): The variable's current units.
dest_unit (:obj:`str`): The desired units. dest_unit (:obj:`str`): The desired units.
Returns: Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in :class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in
the desired units. the desired units.
""" """
u_cleaned = dealias_and_clean_unit(dest_unit) u_cleaned = dealias_and_clean_unit(dest_unit)
if vartype != "temp": if vartype != "temp":
return _apply_conv_fact(var, vartype, var_unit.lower(), u_cleaned) return _apply_conv_fact(var, vartype, var_unit.lower(), u_cleaned)
else: else:
return _apply_temp_conv(var, var_unit.lower(), u_cleaned) return _apply_temp_conv(var, var_unit.lower(), u_cleaned)

3451
src/wrf/util.py

File diff suppressed because it is too large Load Diff

1
src/wrf/version.py

@ -1,2 +1 @@
__version__ = "1.3.1" __version__ = "1.3.1"

Loading…
Cancel
Save