Bill Ladwig 6 years ago
parent
commit
9d35d87ae6
  1. 1
      src/wrf/api.py
  2. 7
      src/wrf/cache.py
  3. 31
      src/wrf/computation.py
  4. 4
      src/wrf/config.py
  5. 37
      src/wrf/constants.py
  6. 43
      src/wrf/coordpair.py
  7. 82
      src/wrf/decorators.py
  8. 2
      src/wrf/destag.py
  9. 229
      src/wrf/extension.py
  10. 43
      src/wrf/g_cape.py
  11. 48
      src/wrf/g_cloudfrac.py
  12. 7
      src/wrf/g_ctt.py
  13. 6
      src/wrf/g_dbz.py
  14. 7
      src/wrf/g_dewpoint.py
  15. 11
      src/wrf/g_geoht.py
  16. 5
      src/wrf/g_helicity.py
  17. 38
      src/wrf/g_latlon.py
  18. 3
      src/wrf/g_omega.py
  19. 2
      src/wrf/g_precip.py
  20. 5
      src/wrf/g_pressure.py
  21. 7
      src/wrf/g_pw.py
  22. 6
      src/wrf/g_rh.py
  23. 4
      src/wrf/g_slp.py
  24. 13
      src/wrf/g_temp.py
  25. 5
      src/wrf/g_terrain.py
  26. 1
      src/wrf/g_times.py
  27. 41
      src/wrf/g_uvmet.py
  28. 1
      src/wrf/g_vorticity.py
  29. 21
      src/wrf/g_wind.py

1
src/wrf/api.py

@ -110,4 +110,3 @@ __all__ += ["CoordPair"]
__all__ += ["to_xy_coords"] __all__ += ["to_xy_coords"]
__all__ += ["cache_item", "get_cached_item"] __all__ += ["cache_item", "get_cached_item"]
__all__ += ["__version__"] __all__ += ["__version__"]

7
src/wrf/cache.py

@ -84,7 +84,7 @@ def cache_item(key, product, value):
_ = cache[key] _ = cache[key]
except KeyError: except KeyError:
if len(cache) >= get_cache_size(): if len(cache) >= get_cache_size():
cache.popitem(last=False) # Remove the oldest dataset cache.popitem(last=False) # Remove the oldest dataset
cache[key] = OrderedDict() cache[key] = OrderedDict()
@ -161,8 +161,3 @@ def _get_cache():
_shrink_cache() _shrink_cache()
return getattr(_local_storage, "cache", None) return getattr(_local_storage, "cache", None)

31
src/wrf/computation.py

@ -15,6 +15,7 @@ from .metadecorators import (set_alg_metadata, set_uvmet_alg_metadata,
set_smooth_metdata) set_smooth_metdata)
from .interputils import get_xy from .interputils import get_xy
@set_interp_metadata("xy") @set_interp_metadata("xy")
def xy(field, pivot_point=None, angle=None, start_point=None, end_point=None, def xy(field, pivot_point=None, angle=None, start_point=None, end_point=None,
meta=True): meta=True):
@ -601,8 +602,8 @@ def uvmet(u, v, lat, lon, cen_long, cone, meta=True, units="m s-1"):
and *v*, but with rightmost dimensions south_north x and *v*, but with rightmost dimensions south_north x
west_east and the same leftmost dimensions as *u* and *v* west_east and the same leftmost dimensions as *u* and *v*
- multi-dimensional with one fewer dimensions as *u* and *v*, - multi-dimensional with one fewer dimensions as *u* and *v*,
with rightmost dimensions south_north x west_east and the same with rightmost dimensions south_north x west_east and the
leftmost dimensions as *u* and *v*, minus the same leftmost dimensions as *u* and *v*, minus the
third-from-the-right dimension of *u* and *v*. third-from-the-right dimension of *u* and *v*.
Note: Note:
@ -622,8 +623,8 @@ def uvmet(u, v, lat, lon, cen_long, cone, meta=True, units="m s-1"):
and *v*, but with rightmost dimensions south_north x and *v*, but with rightmost dimensions south_north x
west_east and the same leftmost dimensions as *u* and *v* west_east and the same leftmost dimensions as *u* and *v*
- multi-dimensional with one fewer dimensions as *u* and *v*, - multi-dimensional with one fewer dimensions as *u* and *v*,
with rightmost dimensions south_north x west_east and the same with rightmost dimensions south_north x west_east and the
leftmost dimensions as *u* and *v*, minus the same leftmost dimensions as *u* and *v*, minus the
third-from-the-right dimension of *u* and *v*. third-from-the-right dimension of *u* and *v*.
@ -870,10 +871,10 @@ def cape_2d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow,
# Cape 2D output is not flipped in the vertical, so index from the # Cape 2D output is not flipped in the vertical, so index from the
# end # end
result[0,...,:,:] = cape_cin[0,...,-1,:,:] result[0, ..., :, :] = cape_cin[0, ..., -1, :, :]
result[1,...,:,:] = cape_cin[1,...,-1,:,:] result[1, ..., :, :] = cape_cin[1, ..., -1, :, :]
result[2,...,:,:] = cape_cin[1,...,-2,:,:] result[2, ..., :, :] = cape_cin[1, ..., -2, :, :]
result[3,...,:,:] = cape_cin[1,...,-3,:,:] result[3, ..., :, :] = cape_cin[1, ..., -3, :, :]
return ma.masked_values(result, missing) return ma.masked_values(result, missing)
@ -1081,7 +1082,7 @@ def cloudfrac(vert, relh, vert_inc_w_height, low_thresh, mid_thresh,
""" """
cfrac = _cloudfrac(vert, relh, vert_inc_w_height, low_thresh, mid_thresh, cfrac = _cloudfrac(vert, relh, vert_inc_w_height, low_thresh, mid_thresh,
high_thresh, missing) high_thresh, missing)
return ma.masked_values(cfrac, missing) return ma.masked_values(cfrac, missing)
@ -1196,12 +1197,11 @@ def ctt(pres_hpa, tkel, qv, qcld, height, terrain, qice=None,
_fill_nocloud = 1 if fill_nocloud else 0 _fill_nocloud = 1 if fill_nocloud else 0
ctt = _ctt(pres_hpa, tkel, qice, qcld, qv, height, terrain, haveqci, ctt = _ctt(pres_hpa, tkel, qice, qcld, qv, height, terrain, haveqci,
_fill_nocloud, missing, opt_thresh) _fill_nocloud, missing, opt_thresh)
return ma.masked_values(ctt, missing) return ma.masked_values(ctt, missing)
@set_alg_metadata(3, "pres", units="dBZ", @set_alg_metadata(3, "pres", units="dBZ",
description="radar reflectivity") description="radar reflectivity")
def dbz(pres, tkel, qv, qr, qs=None, qg=None, use_varint=False, def dbz(pres, tkel, qv, qr, qs=None, qg=None, use_varint=False,
@ -1374,9 +1374,9 @@ def srhel(u, v, height, terrain, top=3000.0, lats=None, meta=True):
""" """
# u, v get swapped in vertical # u, v get swapped in vertical
_u = np.ascontiguousarray(u[...,::-1,:,:]) _u = np.ascontiguousarray(u[..., ::-1, :, :])
_v = np.ascontiguousarray(v[...,::-1,:,:]) _v = np.ascontiguousarray(v[..., ::-1, :, :])
_height = np.ascontiguousarray(height[...,::-1,:,:]) _height = np.ascontiguousarray(height[..., ::-1, :, :])
if lats is None: if lats is None:
_lats = np.ones_like(terrain) _lats = np.ones_like(terrain)
@ -1955,6 +1955,3 @@ def pw(pres, tkel, qv, height, meta=True):
tv = _tv(tkel, qv) tv = _tv(tkel, qv)
return _pw(pres, tv, qv, height) return _pw(pres, tv, qv, height)

4
src/wrf/config.py

@ -9,6 +9,7 @@ from ._wrffortran import (fomp_enabled, fomp_set_num_threads,
_local_config = local() _local_config = local()
def _init_local(): def _init_local():
global _local_config global _local_config
@ -191,7 +192,8 @@ def set_cache_size(size):
@init_local() @init_local()
def get_cache_size(): def get_cache_size():
"""Return the maximum number of items that the threadlocal cache can retain. """Return the maximum number of items that the threadlocal cache can
retain.
Returns: Returns:

37
src/wrf/constants.py

@ -10,10 +10,12 @@ from ._wrffortran import wrf_constants, omp_constants
#: Indicates that all times should be used in a diagnostic routine. #: Indicates that all times should be used in a diagnostic routine.
ALL_TIMES = None ALL_TIMES = None
class Constants(object): class Constants(object):
pass pass
for key,val in viewitems(wrf_constants.__dict__):
for key, val in viewitems(wrf_constants.__dict__):
setattr(Constants, key.upper(), np.asscalar(val)) setattr(Constants, key.upper(), np.asscalar(val))
OMP_SCHED_STATIC = omp_constants.fomp_sched_static OMP_SCHED_STATIC = omp_constants.fomp_sched_static
@ -44,21 +46,22 @@ class ProjectionTypes(object):
MERCATOR = 3 MERCATOR = 3
LAT_LON = 6 LAT_LON = 6
# Create the default fill mapping based on type. # Create the default fill mapping based on type.
_DEFAULT_FILL_MAP = {None: Constants.DEFAULT_FILL, _DEFAULT_FILL_MAP = {None: Constants.DEFAULT_FILL,
np.dtype(np.bool_) : False, np.dtype(np.bool_): False,
np.dtype(np.intc) : Constants.DEFAULT_FILL_INT32, # Usually true np.dtype(np.intc): Constants.DEFAULT_FILL_INT32,
np.dtype(np.int8) : Constants.DEFAULT_FILL_INT8, np.dtype(np.int8): Constants.DEFAULT_FILL_INT8,
np.dtype(np.uint8) : 255, np.dtype(np.uint8): 255,
np.dtype(np.int16) : Constants.DEFAULT_FILL_INT16, np.dtype(np.int16): Constants.DEFAULT_FILL_INT16,
np.dtype(np.uint16) : 65535, np.dtype(np.uint16): 65535,
np.dtype(np.int32) : Constants.DEFAULT_FILL_INT32, np.dtype(np.int32): Constants.DEFAULT_FILL_INT32,
np.dtype(np.uint32) : 4294967295, np.dtype(np.uint32): 4294967295,
np.dtype(np.int64) : Constants.DEFAULT_FILL_INT64, np.dtype(np.int64): Constants.DEFAULT_FILL_INT64,
np.dtype(np.uint64) : 18446744073709551614, np.dtype(np.uint64): 18446744073709551614,
np.dtype(np.float_) : Constants.DEFAULT_FILL_DOUBLE, np.dtype(np.float_): Constants.DEFAULT_FILL_DOUBLE,
np.dtype(np.float32) : Constants.DEFAULT_FILL_FLOAT, np.dtype(np.float32): Constants.DEFAULT_FILL_FLOAT,
np.dtype(np.float64) : Constants.DEFAULT_FILL_DOUBLE np.dtype(np.float64): Constants.DEFAULT_FILL_DOUBLE
} }
if version_info >= (3, ): if version_info >= (3, ):
@ -76,9 +79,3 @@ else:
def default_fill(dtype=None): def default_fill(dtype=None):
dt = np.dtype(dtype) if dtype is not None else None dt = np.dtype(dtype) if dtype is not None else None
return _DEFAULT_FILL_MAP.get(dt, Constants.DEFAULT_FILL) return _DEFAULT_FILL_MAP.get(dt, Constants.DEFAULT_FILL)

43
src/wrf/coordpair.py

@ -35,15 +35,15 @@ def _binary_operator(operator):
""" """
if isinstance(other, CoordPair): if isinstance(other, CoordPair):
args = [ args = [None if getattr(self, attr) is None or
None if getattr(self, attr) is None or getattr(other, attr) is None getattr(other, attr) is None else
else getattr(getattr(self, attr), operator)(getattr(other, attr)) getattr(getattr(self, attr), operator)(getattr(other,
for attr in ("x", "y", "lat", "lon")] attr))
for attr in ("x", "y", "lat", "lon")]
else: else:
args = [ args = [None if getattr(self, attr) is None
None if getattr(self, attr) is None else getattr(getattr(self, attr), operator)(other)
else getattr(getattr(self, attr), operator)(other) for attr in ("x", "y", "lat", "lon")]
for attr in ("x", "y", "lat", "lon")]
return CoordPair(*args) return CoordPair(*args)
@ -151,7 +151,6 @@ class CoordPair(object):
self.lat = lat self.lat = lat
self.lon = lon self.lon = lon
def __repr__(self): def __repr__(self):
args = [] args = []
if self.x is not None: if self.x is not None:
@ -166,11 +165,9 @@ class CoordPair(object):
return "{}({})".format(self.__class__.__name__, argstr) return "{}({})".format(self.__class__.__name__, argstr)
def __str__(self): def __str__(self):
return self.__repr__() return self.__repr__()
def xy_str(self, fmt="{:.4f}, {:.4f}"): def xy_str(self, fmt="{:.4f}, {:.4f}"):
"""Return a :obj:`str` for the (x,y) coordinate pair. """Return a :obj:`str` for the (x,y) coordinate pair.
@ -188,7 +185,6 @@ class CoordPair(object):
return fmt.format(self.x, self.y) return fmt.format(self.x, self.y)
def latlon_str(self, fmt="{:.4f}, {:.4f}"): def latlon_str(self, fmt="{:.4f}, {:.4f}"):
"""Return a :obj:`str` for the (latitude, longitude) coordinate pair. """Return a :obj:`str` for the (latitude, longitude) coordinate pair.
@ -206,7 +202,6 @@ class CoordPair(object):
return fmt.format(self.lat, self.lon) return fmt.format(self.lat, self.lon)
def __round__(self, ndigits=None): def __round__(self, ndigits=None):
"""Return a new :class:`CoordPair` object with all coordinate values """Return a new :class:`CoordPair` object with all coordinate values
rounded to the nearest integer. rounded to the nearest integer.
@ -226,23 +221,20 @@ class CoordPair(object):
return CoordPair(*args) return CoordPair(*args)
def __pow__(self, other, modulo=None): def __pow__(self, other, modulo=None):
if isinstance(other, CoordPair): if isinstance(other, CoordPair):
args = [ args = [None if getattr(self, attr) is None or
None if getattr(self, attr) is None or getattr(other, attr) is None getattr(other, attr) is None
else getattr(getattr(self, attr), "__pow__")(getattr(other, attr), else getattr(getattr(self, attr), "__pow__")(
modulo) getattr(other, attr), modulo)
for attr in ("x", "y", "lat", "lon")] for attr in ("x", "y", "lat", "lon")]
else: else:
args = [ args = [None if getattr(self, attr) is None
None if getattr(self, attr) is None else getattr(getattr(self, attr), "__pow__")(other, modulo)
else getattr(getattr(self, attr), "__pow__")(other, modulo) for attr in ("x", "y", "lat", "lon")]
for attr in ("x", "y", "lat", "lon")]
return CoordPair(*args) return CoordPair(*args)
def __rpow__(self, other): def __rpow__(self, other):
return self.__pow__(other) return self.__pow__(other)
@ -260,6 +252,3 @@ for operator in ("__neg__", "__pos__", "__abs__", "__invert__"):
for operator in ("__lt__", "__le__", "__eq__", "__ne__", "__gt__", "__ge__"): for operator in ("__lt__", "__le__", "__eq__", "__ne__", "__gt__", "__ge__"):
setattr(CoordPair, operator, _cmp_operator(operator)) setattr(CoordPair, operator, _cmp_operator(operator))

82
src/wrf/decorators.py

@ -15,6 +15,7 @@ from .constants import default_fill
if xarray_enabled(): if xarray_enabled():
from xarray import DataArray from xarray import DataArray
def convert_units(unit_type, alg_unit): def convert_units(unit_type, alg_unit):
"""A decorator that converts the units from the wrapped function's output. """A decorator that converts the units from the wrapped function's output.
@ -47,17 +48,6 @@ def convert_units(unit_type, alg_unit):
return func_wrapper return func_wrapper
#def _calc_out_dims(outvar, left_dims):
# """
#
# """
# #left_dims = [x for x in left_dims]
# #right_dims = [x for x in outvar.shape]
# #return left_dims + right_dims
#
# return left_dims + outvar.shape
def left_iteration(ref_var_expected_dims, def left_iteration(ref_var_expected_dims,
ref_var_right_ndims, ref_var_right_ndims,
insert_dims=None, insert_dims=None,
@ -162,7 +152,7 @@ def left_iteration(ref_var_expected_dims,
if "outview" not in kwargs: if "outview" not in kwargs:
outd = OrderedDict((outkey, np.empty(outdims, alg_dtype)) outd = OrderedDict((outkey, np.empty(outdims, alg_dtype))
for outkey in _outkeys) for outkey in _outkeys)
mask_output = False mask_output = False
for left_idxs in iter_left_indexes(extra_dims): for left_idxs in iter_left_indexes(extra_dims):
@ -171,16 +161,15 @@ def left_iteration(ref_var_expected_dims,
# the right (e.g. [1,1,:]) # the right (e.g. [1,1,:])
left_and_slice_idxs = left_idxs + (slice(None), ) left_and_slice_idxs = left_idxs + (slice(None), )
# Slice the args if applicable # Slice the args if applicable
new_args = [arg[left_and_slice_idxs] new_args = [arg[left_and_slice_idxs]
if i not in _ignore_args else arg if i not in _ignore_args else arg
for i,arg in enumerate(args)] for i, arg in enumerate(args)]
# Slice the kwargs if applicable # Slice the kwargs if applicable
new_kargs = {key:(val[left_and_slice_idxs] new_kargs = {key: (val[left_and_slice_idxs]
if key not in _ignore_kargs else val) if key not in _ignore_kargs else val)
for key,val in viewitems(kwargs)} for key, val in viewitems(kwargs)}
# 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
@ -188,14 +177,14 @@ def left_iteration(ref_var_expected_dims,
try: try:
_ = arg.ndim _ = arg.ndim
except AttributeError: except AttributeError:
continue # Not an array object continue # Not an array object
else: else:
arr = to_np(arg) arr = to_np(arg)
try: try:
all_masked = arr.mask.all() all_masked = arr.mask.all()
except AttributeError: except AttributeError:
pass # Not a masked array pass # Not a masked array
else: else:
if all_masked: if all_masked:
for output in viewvalues(outd): for output in viewvalues(outd):
@ -210,7 +199,7 @@ def left_iteration(ref_var_expected_dims,
# Insert the output views if one hasn't been provided # Insert the output views if one hasn't been provided
if "outview" not in new_kargs: if "outview" not in new_kargs:
for outkey,output in viewitems(outd): for outkey, output in viewitems(outd):
outview = output[left_and_slice_idxs] outview = output[left_and_slice_idxs]
new_kargs[outkey] = outview new_kargs[outkey] = outview
@ -219,17 +208,14 @@ def left_iteration(ref_var_expected_dims,
# 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")
if len(outd) == 1: if len(outd) == 1:
output = next(iter(viewvalues(outd))) output = next(iter(viewvalues(outd)))
else: else:
output = tuple(arr for arr in viewvalues(outd)) output = tuple(arr for arr in viewvalues(outd))
if cast_output: if cast_output:
if isinstance(output, np.ndarray): if isinstance(output, np.ndarray):
output = output.astype(ref_var_dtype) output = output.astype(ref_var_dtype)
@ -262,8 +248,8 @@ def cast_type(ref_idx=0, arg_idxs=None, karg_names=None,
positional arguments to be used as the reference variable for positional arguments to be used as the reference variable for
determining the :class:`numpy.dtype` to return. Default is 0. determining the :class:`numpy.dtype` to return. Default is 0.
arg_idxs (sequence of :obj:`int`, optional): A sequence of indexes in the arg_idxs (sequence of :obj:`int`, optional): A sequence of indexes in
wrapped function's positional arguments that indicate which the wrapped function's positional arguments that indicate which
arguments to cast. Must be specified if *karg_names* is None. arguments to cast. Must be specified if *karg_names* is None.
Default is None. Default is None.
@ -272,8 +258,8 @@ def cast_type(ref_idx=0, arg_idxs=None, karg_names=None,
arguments to cast. Must be specified if *arg_idxs* is None. arguments to cast. Must be specified if *arg_idxs* is None.
Default is None. Default is None.
alg_dtype (:class:`numpy.dtype` or :obj:`str`): The numpy data type used alg_dtype (:class:`numpy.dtype` or :obj:`str`): The numpy data type
in the wrapped function. used in the wrapped function.
outviews (:obj:`str` or a sequence): A single key or sequence of keys outviews (:obj:`str` or a sequence): A single key or sequence of keys
that indicate the wrapped function's keyword argument to use that indicate the wrapped function's keyword argument to use
@ -300,16 +286,15 @@ def cast_type(ref_idx=0, arg_idxs=None, karg_names=None,
if _outview is not None: if _outview is not None:
has_outview = True has_outview = True
orig_type = args[ref_idx].dtype orig_type = args[ref_idx].dtype
new_args = [arg.astype(alg_dtype) new_args = [arg.astype(alg_dtype)
if i in _arg_idxs else arg if i in _arg_idxs else arg
for i,arg in enumerate(args)] for i, arg in enumerate(args)]
new_kargs = {key:(val.astype(alg_dtype) new_kargs = {key: (val.astype(alg_dtype)
if key in _karg_names else val) if key in _karg_names else val)
for key,val in viewitems(kwargs)} for key, val in viewitems(kwargs)}
result = wrapped(*new_args, **new_kargs) result = wrapped(*new_args, **new_kargs)
@ -319,7 +304,7 @@ def cast_type(ref_idx=0, arg_idxs=None, karg_names=None,
if result.dtype == orig_type: if result.dtype == orig_type:
return result return result
return result.astype(orig_type) return result.astype(orig_type)
elif isinstance(result, Iterable): # got back a sequence of arrays elif isinstance(result, Iterable): # got back a sequence of arrays
return tuple(arr.astype(orig_type) return tuple(arr.astype(orig_type)
if arr.dtype != orig_type else arr if arr.dtype != orig_type else arr
for arr in result) for arr in result)
@ -401,8 +386,8 @@ def extract_and_transpose(do_transpose=True, outviews="outview"):
new_args = [_extract_and_transpose(arg, do_transpose) for arg in args] new_args = [_extract_and_transpose(arg, do_transpose) for arg in args]
new_kargs = {key:_extract_and_transpose(val, do_transpose) new_kargs = {key: _extract_and_transpose(val, do_transpose)
for key,val in viewitems(kwargs)} for key, val in viewitems(kwargs)}
result = wrapped(*new_args, **new_kargs) result = wrapped(*new_args, **new_kargs)
@ -481,7 +466,7 @@ def check_args(refvaridx, refvarndim, rightdims, stagger=None,
else: else:
_stagger = stagger _stagger = stagger
for i,ndim in enumerate(rightdims): for i, ndim in enumerate(rightdims):
if ndim is None: if ndim is None:
continue continue
@ -498,9 +483,10 @@ def check_args(refvaridx, refvarndim, rightdims, stagger=None,
# Check that the number of dims is correct # Check that the number of dims is correct
if (var.ndim - extra_dims != right_var_ndims): if (var.ndim - extra_dims != right_var_ndims):
raise ValueError("invalid number of dimensions for argument " raise ValueError("invalid number of dimensions for argument "
"{} (got {}, expected {}).".format(i, "{} (got {}, expected {}).".format(
var.ndim, i,
right_var_ndims + extra_dims)) var.ndim,
right_var_ndims + extra_dims))
# Add 1 to the reference staggered dim index before doing the check # Add 1 to the reference staggered dim index before doing the check
if _stagger[i] is not None: if _stagger[i] is not None:
@ -514,22 +500,14 @@ def check_args(refvaridx, refvarndim, rightdims, stagger=None,
# Check that right dimensions are lined up # Check that right dimensions are lined up
if (var.shape[-right_var_ndims:] != if (var.shape[-right_var_ndims:] !=
ref_right_sizes[-right_var_ndims:]): ref_right_sizes[-right_var_ndims:]):
raise ValueError("invalid shape for argument " raise ValueError("invalid shape for argument "
"{} (got {}, expected {})".format(i, "{} (got {}, expected {})".format(
var.shape[-right_var_ndims:], i,
ref_right_sizes[-right_var_ndims:])) var.shape[-right_var_ndims:],
ref_right_sizes[-right_var_ndims:]))
return wrapped(*args, **kwargs) return wrapped(*args, **kwargs)
return func_wrapper return func_wrapper

2
src/wrf/destag.py

@ -60,5 +60,3 @@ def destagger(var, stagger_dim, meta=False):
result = .5*(var[tuple(dim_ranges_1)] + var[tuple(dim_ranges_2)]) result = .5*(var[tuple(dim_ranges_1)] + var[tuple(dim_ranges_2)])
return result return result

229
src/wrf/extension.py

@ -5,29 +5,30 @@ import numpy as np
from .constants import Constants, default_fill from .constants import Constants, default_fill
from wrf._wrffortran import (dcomputetk, dinterp3dz, dinterp2dxy, dinterp1d, from wrf._wrffortran import (dcomputetk, dinterp3dz, dinterp2dxy, dinterp1d,
dcomputeseaprs, dfilter2d, dcomputerh, dcomputeuvmet, dcomputeseaprs, dfilter2d, dcomputerh,
dcomputetd, dcapecalc2d, dcapecalc3d, dcloudfrac2, dcomputeuvmet, dcomputetd, dcapecalc2d,
wrfcttcalc, calcdbz, dcalrelhl, dcalcuh, dcomputepv, dcapecalc3d, dcloudfrac2, wrfcttcalc, calcdbz,
dcomputeabsvort, dlltoij, dijtoll, deqthecalc, dcalrelhl, dcalcuh, dcomputepv, dcomputeabsvort,
omgcalc, virtual_temp, wetbulbcalc, dcomputepw, dlltoij, dijtoll, deqthecalc, omgcalc,
wrf_monotonic, wrf_vintrp, dcomputewspd, virtual_temp, wetbulbcalc, dcomputepw,
dcomputewdir, dinterp3dz_2dlev, wrf_monotonic, wrf_vintrp, dcomputewspd,
fomp_set_num_threads, fomp_get_num_threads, dcomputewdir, dinterp3dz_2dlev,
fomp_get_max_threads, fomp_get_thread_num, fomp_set_num_threads, fomp_get_num_threads,
fomp_get_num_procs, fomp_in_parallel, fomp_get_max_threads, fomp_get_thread_num,
fomp_set_dynamic, fomp_get_dynamic, fomp_set_nested, fomp_get_num_procs, fomp_in_parallel,
fomp_get_nested, fomp_set_schedule, fomp_set_dynamic, fomp_get_dynamic,
fomp_get_schedule, fomp_get_thread_limit, fomp_set_nested, fomp_get_nested,
fomp_set_max_active_levels, fomp_set_schedule, fomp_get_schedule,
fomp_get_max_active_levels, fomp_get_level, fomp_get_thread_limit, fomp_set_max_active_levels,
fomp_get_ancestor_thread_num, fomp_get_team_size, fomp_get_max_active_levels, fomp_get_level,
fomp_get_active_level, fomp_in_final, fomp_get_ancestor_thread_num, fomp_get_team_size,
fomp_init_lock, fomp_init_nest_lock, fomp_get_active_level, fomp_in_final,
fomp_destroy_lock, fomp_destroy_nest_lock, fomp_init_lock, fomp_init_nest_lock,
fomp_set_lock, fomp_set_nest_lock, fomp_destroy_lock, fomp_destroy_nest_lock,
fomp_unset_lock, fomp_unset_nest_lock, fomp_set_lock, fomp_set_nest_lock,
fomp_test_lock, fomp_test_nest_lock, fomp_unset_lock, fomp_unset_nest_lock,
fomp_get_wtime, fomp_get_wtick, fomp_enabled) fomp_test_lock, fomp_test_nest_lock,
fomp_get_wtime, fomp_get_wtick, fomp_enabled)
from .decorators import (left_iteration, cast_type, from .decorators import (left_iteration, cast_type,
extract_and_transpose, check_args) extract_and_transpose, check_args)
@ -37,6 +38,7 @@ from .specialdec import (uvmet_left_iter, cape_left_iter,
cloudfrac_left_iter, check_cape_args, cloudfrac_left_iter, check_cape_args,
interplevel_left_iter, check_interplevel_args) interplevel_left_iter, check_interplevel_args)
class DiagnosticError(Exception): class DiagnosticError(Exception):
"""Raised when an error occurs in a diagnostic routine.""" """Raised when an error occurs in a diagnostic routine."""
def __init__(self, message=None): def __init__(self, message=None):
@ -66,6 +68,7 @@ class DiagnosticError(Exception):
""" """
raise self.__class__(message) raise self.__class__(message)
# The routines below are thin wrappers around the Fortran functions. These # The routines below are thin wrappers around the Fortran functions. These
# are not meant to be called by end users. Use the public API instead for # are not meant to be called by end users. Use the public API instead for
# that purpose. # that purpose.
@ -73,10 +76,9 @@ class DiagnosticError(Exception):
# IMPORTANT! Unless otherwise noted, all variables used in the routines # IMPORTANT! Unless otherwise noted, all variables used in the routines
# below assume that Fortran-ordered views are being used. This allows # below assume that Fortran-ordered views are being used. This allows
# f2py to pass the array pointers directly to the Fortran routine. # f2py to pass the array pointers directly to the Fortran routine.
@check_interplevel_args(is2dlev=False) @check_interplevel_args(is2dlev=False)
@interplevel_left_iter(is2dlev=False) @interplevel_left_iter(is2dlev=False)
@cast_type(arg_idxs=(0,1,2)) @cast_type(arg_idxs=(0, 1, 2))
@extract_and_transpose() @extract_and_transpose()
def _interpz3d(field3d, z, desiredloc, missingval, outview=None): def _interpz3d(field3d, z, desiredloc, missingval, outview=None):
"""Wrapper for dinterp3dz. """Wrapper for dinterp3dz.
@ -98,7 +100,7 @@ def _interpz3d(field3d, z, desiredloc, missingval, outview=None):
@check_interplevel_args(is2dlev=True) @check_interplevel_args(is2dlev=True)
@interplevel_left_iter(is2dlev=True) @interplevel_left_iter(is2dlev=True)
@cast_type(arg_idxs=(0,1,2)) @cast_type(arg_idxs=(0, 1, 2))
@extract_and_transpose() @extract_and_transpose()
def _interpz3d_lev2d(field3d, z, lev2d, missingval, outview=None): def _interpz3d_lev2d(field3d, z, lev2d, missingval, outview=None):
"""Wrapper for dinterp3dz. """Wrapper for dinterp3dz.
@ -117,10 +119,10 @@ def _interpz3d_lev2d(field3d, z, lev2d, missingval, outview=None):
return result return result
@check_args(0, 3, (3,)) @check_args(0, 3, (3, ))
@left_iteration(3, combine_dims([(0,-3),(1,-2)]), ref_var_idx=0, @left_iteration(3, combine_dims([(0, -3), (1, -2)]), ref_var_idx=0,
ignore_args=(1,)) ignore_args=(1, ))
@cast_type(arg_idxs=(0,1)) @cast_type(arg_idxs=(0, 1))
@extract_and_transpose() @extract_and_transpose()
def _interp2dxy(field3d, xy, outview=None): def _interp2dxy(field3d, xy, outview=None):
"""Wrapper for dinterp2dxy. """Wrapper for dinterp2dxy.
@ -138,9 +140,9 @@ def _interp2dxy(field3d, xy, outview=None):
return result return result
@check_args(0, 1, (1,1,None,None)) @check_args(0, 1, (1, 1, None, None))
@left_iteration(1, combine_dims([(2,0)]), ref_var_idx=0, ignore_args=(2,3)) @left_iteration(1, combine_dims([(2, 0)]), ref_var_idx=0, ignore_args=(2, 3))
@cast_type(arg_idxs=(0,1,2)) @cast_type(arg_idxs=(0, 1, 2))
@extract_and_transpose() @extract_and_transpose()
def _interp1d(v_in, z_in, z_out, missingval, outview=None): def _interp1d(v_in, z_in, z_out, missingval, outview=None):
"""Wrapper for dinterp1d. """Wrapper for dinterp1d.
@ -160,9 +162,9 @@ def _interp1d(v_in, z_in, z_out, missingval, outview=None):
return result return result
@left_iteration(3, combine_dims([(3,0), (1,0)]), @left_iteration(3, combine_dims([(3, 0), (1, 0)]),
ref_var_idx=0, ignore_args=(1,3,4)) ref_var_idx=0, ignore_args=(1, 3, 4))
@cast_type(arg_idxs=(0,)) @cast_type(arg_idxs=(0, ))
@extract_and_transpose(do_transpose=False) @extract_and_transpose(do_transpose=False)
def _vertcross(field3d, xy, var2dz, z_var2d, missingval, outview=None): def _vertcross(field3d, xy, var2dz, z_var2d, missingval, outview=None):
"""Return the vertical cross section. """Return the vertical cross section.
@ -180,14 +182,14 @@ def _vertcross(field3d, xy, var2dz, z_var2d, missingval, outview=None):
var2dtmp = _interp2dxy(field3d, xy) var2dtmp = _interp2dxy(field3d, xy)
for i in py3range(xy.shape[0]): for i in py3range(xy.shape[0]):
outview[:,i] = _interp1d(var2dtmp[:,i], var2dz[:,i], z_var2d, outview[:, i] = _interp1d(var2dtmp[:, i], var2dz[:, i], z_var2d,
missingval) missingval)
return outview return outview
@left_iteration(2, combine_dims([(1,0)]), ref_var_idx=0, ignore_args=(1,)) @left_iteration(2, combine_dims([(1, 0)]), ref_var_idx=0, ignore_args=(1, ))
@cast_type(arg_idxs=(0,)) @cast_type(arg_idxs=(0, ))
@extract_and_transpose(do_transpose=False) @extract_and_transpose(do_transpose=False)
def _interpline(field2d, xy, outview=None): def _interpline(field2d, xy, outview=None):
"""Return the two-dimensional field interpolated to a line. """Return the two-dimensional field interpolated to a line.
@ -204,7 +206,7 @@ def _interpline(field2d, xy, outview=None):
tmp_shape = (1,) + field2d.shape tmp_shape = (1,) + field2d.shape
var2dtmp = np.empty(tmp_shape, field2d.dtype) var2dtmp = np.empty(tmp_shape, field2d.dtype)
var2dtmp[0,:,:] = field2d[:,:] var2dtmp[0, :, :] = field2d[:, :]
var1dtmp = _interp2dxy(var2dtmp, xy) var1dtmp = _interp2dxy(var2dtmp, xy)
@ -213,9 +215,9 @@ def _interpline(field2d, xy, outview=None):
return outview return outview
@check_args(0, 3, (3,3,3,3)) @check_args(0, 3, (3, 3, 3, 3))
@left_iteration(3, 2, ref_var_idx=0) @left_iteration(3, 2, ref_var_idx=0)
@cast_type(arg_idxs=(0,1,2,3)) @cast_type(arg_idxs=(0, 1, 2, 3))
@extract_and_transpose() @extract_and_transpose()
def _slp(z, t, p, q, outview=None): def _slp(z, t, p, q, outview=None):
"""Wrapper for dcomputeseaprs. """Wrapper for dcomputeseaprs.
@ -250,9 +252,9 @@ def _slp(z, t, p, q, outview=None):
return result return result
@check_args(0, 3, (3,3)) @check_args(0, 3, (3, 3))
@left_iteration(3, 3, ref_var_idx=0) @left_iteration(3, 3, ref_var_idx=0)
@cast_type(arg_idxs=(0,1)) @cast_type(arg_idxs=(0, 1))
@extract_and_transpose() @extract_and_transpose()
def _tk(pressure, theta, outview=None): def _tk(pressure, theta, outview=None):
"""Wrapper for dcomputetk. """Wrapper for dcomputetk.
@ -265,16 +267,16 @@ def _tk(pressure, theta, outview=None):
if outview is None: if outview is None:
outview = np.empty_like(pressure) outview = np.empty_like(pressure)
result = dcomputetk(outview.ravel(order="A"), result = dcomputetk(outview.ravel(order="A"),
pressure.ravel(order="A"), pressure.ravel(order="A"),
theta.ravel(order="A")) theta.ravel(order="A"))
result = np.reshape(result, shape, order="F") result = np.reshape(result, shape, order="F")
return result return result
@check_args(0, 2, (2,2)) @check_args(0, 2, (2, 2))
@left_iteration(2, 2, ref_var_idx=0) @left_iteration(2, 2, ref_var_idx=0)
@cast_type(arg_idxs=(0,1)) @cast_type(arg_idxs=(0, 1))
@extract_and_transpose() @extract_and_transpose()
def _td(pressure, qv_in, outview=None): def _td(pressure, qv_in, outview=None):
"""Wrapper for dcomputetd. """Wrapper for dcomputetd.
@ -294,9 +296,9 @@ def _td(pressure, qv_in, outview=None):
return result return result
@check_args(0, 2, (2,2,2)) @check_args(0, 2, (2, 2, 2))
@left_iteration(2, 2, ref_var_idx=0) @left_iteration(2, 2, ref_var_idx=0)
@cast_type(arg_idxs=(0,1,2)) @cast_type(arg_idxs=(0, 1, 2))
@extract_and_transpose() @extract_and_transpose()
def _rh(qv, q, t, outview=None): def _rh(qv, q, t, outview=None):
"""Wrapper for dcomputerh. """Wrapper for dcomputerh.
@ -308,9 +310,9 @@ def _rh(qv, q, t, outview=None):
if outview is None: if outview is None:
outview = np.empty_like(qv) outview = np.empty_like(qv)
result = dcomputerh(qv.ravel(order="A"), result = dcomputerh(qv.ravel(order="A"),
q.ravel(order="A"), q.ravel(order="A"),
t.ravel(order="A"), t.ravel(order="A"),
outview.ravel(order="A")) outview.ravel(order="A"))
result = np.reshape(result, shape, order="F") result = np.reshape(result, shape, order="F")
return result return result
@ -318,12 +320,12 @@ def _rh(qv, q, t, outview=None):
# Note: combining the -3 and -2 dimensions from u, then the -1 dimension # Note: combining the -3 and -2 dimensions from u, then the -1 dimension
# from v # from v
@check_args(0, 3, (3,3,2,2,2,2), stagger=(-1,-2,-1,-2,None,None), @check_args(0, 3, (3, 3, 2, 2, 2, 2), stagger=(-1, -2, -1, -2, None, None),
refstagdim=-1) refstagdim=-1)
@left_iteration(3, combine_dims([(0, (-3,-2)), @left_iteration(3, combine_dims([(0, (-3, -2)),
(1, (-1,))]), (1, (-1, ))]),
ref_var_idx=0, ignore_args=(6,7)) ref_var_idx=0, ignore_args=(6, 7))
@cast_type(arg_idxs=(0,1,2,3,4,5)) @cast_type(arg_idxs=(0, 1, 2, 3, 4, 5))
@extract_and_transpose() @extract_and_transpose()
def _avo(u, v, msfu, msfv, msfm, cor, dx, dy, outview=None): def _avo(u, v, msfu, msfv, msfm, cor, dx, dy, outview=None):
"""Wrapper for dcomputeabsvort. """Wrapper for dcomputeabsvort.
@ -332,7 +334,7 @@ def _avo(u, v, msfu, msfv, msfm, cor, dx, dy, outview=None):
""" """
if outview is None: if outview is None:
outshape = (v.shape[0],) + u.shape[1:] outshape = (v.shape[0], ) + u.shape[1:]
outview = np.empty(outshape, np.float64, order="F") outview = np.empty(outshape, np.float64, order="F")
result = dcomputeabsvort(outview, result = dcomputeabsvort(outview,
@ -348,11 +350,10 @@ def _avo(u, v, msfu, msfv, msfm, cor, dx, dy, outview=None):
return result return result
@check_args(0, 3, (3,3,3,3,2,2,2,2), stagger=(-1,-2,None,None,-1,-2,None, @check_args(0, 3, (3, 3, 3, 3, 2, 2, 2, 2),
None), stagger=(-1, -2, None, None, -1, -2, None, None), refstagdim=-1)
refstagdim=-1) @left_iteration(3, 3, ref_var_idx=2, ignore_args=(8, 9))
@left_iteration(3, 3, ref_var_idx=2, ignore_args=(8,9)) @cast_type(arg_idxs=(0, 1, 2, 3, 4, 5, 6, 7))
@cast_type(arg_idxs=(0,1,2,3,4,5,6,7))
@extract_and_transpose() @extract_and_transpose()
def _pvo(u, v, theta, prs, msfu, msfv, msfm, cor, dx, dy, outview=None): def _pvo(u, v, theta, prs, msfu, msfv, msfm, cor, dx, dy, outview=None):
"""Wrapper for dcomputepv. """Wrapper for dcomputepv.
@ -378,9 +379,9 @@ def _pvo(u, v, theta, prs, msfu, msfv, msfm, cor, dx, dy, outview=None):
return result return result
@check_args(0, 3, (3,3,3)) @check_args(0, 3, (3, 3, 3))
@left_iteration(3, 3, ref_var_idx=0) @left_iteration(3, 3, ref_var_idx=0)
@cast_type(arg_idxs=(0,1,2)) @cast_type(arg_idxs=(0, 1, 2))
@extract_and_transpose() @extract_and_transpose()
def _eth(qv, tk, p, outview=None): def _eth(qv, tk, p, outview=None):
"""Wrapper for deqthecalc. """Wrapper for deqthecalc.
@ -400,11 +401,13 @@ def _eth(qv, tk, p, outview=None):
@uvmet_left_iter() @uvmet_left_iter()
@cast_type(arg_idxs=(0,1,2,3)) @cast_type(arg_idxs=(0, 1, 2, 3))
@extract_and_transpose() @extract_and_transpose()
def _uvmet(u, v, lat, lon, cen_long, cone, isstag=0, has_missing=False, def _uvmet(u, v, lat, lon, cen_long, cone, isstag=0, has_missing=False,
umissing=default_fill(np.float64), vmissing=default_fill(np.float64), umissing=default_fill(np.float64),
uvmetmissing=default_fill(np.float64), outview=None): vmissing=default_fill(np.float64),
uvmetmissing=default_fill(np.float64),
outview=None):
"""Wrapper for dcomputeuvmet. """Wrapper for dcomputeuvmet.
Located in wrf_user.f90. Located in wrf_user.f90.
@ -437,9 +440,9 @@ def _uvmet(u, v, lat, lon, cen_long, cone, isstag=0, has_missing=False,
return result return result
@check_args(0, 3, (3,3,3,3)) @check_args(0, 3, (3, 3, 3, 3))
@left_iteration(3, 3, ref_var_idx=0) @left_iteration(3, 3, ref_var_idx=0)
@cast_type(arg_idxs=(0,1,2,3)) @cast_type(arg_idxs=(0, 1, 2, 3))
@extract_and_transpose() @extract_and_transpose()
def _omega(qv, tk, w, p, outview=None): def _omega(qv, tk, w, p, outview=None):
"""Wrapper for omgcalc. """Wrapper for omgcalc.
@ -459,9 +462,9 @@ def _omega(qv, tk, w, p, outview=None):
return result return result
@check_args(0, 3, (3,3)) @check_args(0, 3, (3, 3))
@left_iteration(3, 3, ref_var_idx=0) @left_iteration(3, 3, ref_var_idx=0)
@cast_type(arg_idxs=(0,1)) @cast_type(arg_idxs=(0, 1))
@extract_and_transpose() @extract_and_transpose()
def _tv(tk, qv, outview=None): def _tv(tk, qv, outview=None):
"""Wrapper for virtual_temp. """Wrapper for virtual_temp.
@ -479,9 +482,9 @@ def _tv(tk, qv, outview=None):
return result return result
@check_args(0, 3, (3,3,3)) @check_args(0, 3, (3, 3, 3))
@left_iteration(3, 3, ref_var_idx=0, ignore_args=(3,)) @left_iteration(3, 3, ref_var_idx=0, ignore_args=(3,))
@cast_type(arg_idxs=(0,1,2)) @cast_type(arg_idxs=(0, 1, 2))
@extract_and_transpose() @extract_and_transpose()
def _wetbulb(p, tk, qv, psafile=psafilepath(), outview=None): def _wetbulb(p, tk, qv, psafile=psafilepath(), outview=None):
"""Wrapper for wetbulbcalc. """Wrapper for wetbulbcalc.
@ -509,9 +512,9 @@ def _wetbulb(p, tk, qv, psafile=psafilepath(), outview=None):
return result return result
@check_args(0, 3, (3,3,3,2,2)) @check_args(0, 3, (3, 3, 3, 2, 2))
@left_iteration(3, 2, ref_var_idx=0, ignore_args=(5,)) @left_iteration(3, 2, ref_var_idx=0, ignore_args=(5, ))
@cast_type(arg_idxs=(0,1,2,3,4)) @cast_type(arg_idxs=(0, 1, 2, 3, 4))
@extract_and_transpose() @extract_and_transpose()
def _srhel(u, v, z, ter, lats, top, outview=None): def _srhel(u, v, z, ter, lats, top, outview=None):
"""Wrapper for dcalrelhl. """Wrapper for dcalrelhl.
@ -533,9 +536,9 @@ def _srhel(u, v, z, ter, lats, top, outview=None):
return result return result
@check_args(2, 3, (3,2,3,3,3), stagger=(-3,None,None,None,-3)) @check_args(2, 3, (3, 2, 3, 3, 3), stagger=(-3, None, None, None, -3))
@left_iteration(3, 2, ref_var_idx=2, ignore_args=(5,6,7,8)) @left_iteration(3, 2, ref_var_idx=2, ignore_args=(5, 6, 7, 8))
@cast_type(arg_idxs=(0,1,2,3,4)) @cast_type(arg_idxs=(0, 1, 2, 3, 4))
@extract_and_transpose() @extract_and_transpose()
def _udhel(zstag, mapfct, u, v, wstag, dx, dy, bottom, top, outview=None): def _udhel(zstag, mapfct, u, v, wstag, dx, dy, bottom, top, outview=None):
"""Wrapper for dcalcuh. """Wrapper for dcalcuh.
@ -567,9 +570,9 @@ def _udhel(zstag, mapfct, u, v, wstag, dx, dy, bottom, top, outview=None):
return result return result
@check_args(0, 3, (3,3,3,3), stagger=(None, None, None, -3)) @check_args(0, 3, (3, 3, 3, 3), stagger=(None, None, None, -3))
@left_iteration(3, 2, ref_var_idx=0) @left_iteration(3, 2, ref_var_idx=0)
@cast_type(arg_idxs=(0,1,2,3)) @cast_type(arg_idxs=(0, 1, 2, 3))
@extract_and_transpose() @extract_and_transpose()
def _pw(p, tv, qv, ht, outview=None): def _pw(p, tv, qv, ht, outview=None):
"""Wrapper for dcomputepw. """Wrapper for dcomputepw.
@ -589,9 +592,9 @@ def _pw(p, tv, qv, ht, outview=None):
return result return result
@check_args(0, 3, (3,3,3,3,3,3)) @check_args(0, 3, (3, 3, 3, 3, 3, 3))
@left_iteration(3, 3, ref_var_idx=0, ignore_args=(6,7,8)) @left_iteration(3, 3, ref_var_idx=0, ignore_args=(6, 7, 8))
@cast_type(arg_idxs=(0,1,2,3,4,5)) @cast_type(arg_idxs=(0, 1, 2, 3, 4, 5))
@extract_and_transpose() @extract_and_transpose()
def _dbz(p, tk, qv, qr, qs, qg, sn0, ivarint, iliqskin, outview=None): def _dbz(p, tk, qv, qr, qs, qg, sn0, ivarint, iliqskin, outview=None):
"""Wrapper for calcdbz. """Wrapper for calcdbz.
@ -618,7 +621,7 @@ def _dbz(p, tk, qv, qr, qs, qg, sn0, ivarint, iliqskin, outview=None):
@check_cape_args() @check_cape_args()
@cape_left_iter() @cape_left_iter()
@cast_type(arg_idxs=(0,1,2,3,4,5), outviews=("capeview", "cinview")) @cast_type(arg_idxs=(0, 1, 2, 3, 4, 5), outviews=("capeview", "cinview"))
@extract_and_transpose(outviews=("capeview", "cinview")) @extract_and_transpose(outviews=("capeview", "cinview"))
def _cape(p_hpa, tk, qv, ht, ter, sfp, missing, i3dflag, ter_follow, def _cape(p_hpa, tk, qv, ht, ter, sfp, missing, i3dflag, ter_follow,
psafile=psafilepath(), capeview=None, cinview=None): psafile=psafilepath(), capeview=None, cinview=None):
@ -674,12 +677,14 @@ def _cape(p_hpa, tk, qv, ht, ter, sfp, missing, i3dflag, ter_follow,
return result return result
@check_args(0, 3, (3,3))
@check_args(0, 3, (3, 3))
@cloudfrac_left_iter() @cloudfrac_left_iter()
@cast_type(arg_idxs=(0, 1), outviews=("lowview", "midview", "highview")) @cast_type(arg_idxs=(0, 1), outviews=("lowview", "midview", "highview"))
@extract_and_transpose(outviews=("lowview", "midview", "highview")) @extract_and_transpose(outviews=("lowview", "midview", "highview"))
def _cloudfrac(vert, rh, vert_inc_w_height, low_thresh, mid_thresh, def _cloudfrac(vert, rh, vert_inc_w_height, low_thresh, mid_thresh,
high_thresh, missing, lowview=None, midview=None, highview=None): high_thresh, missing, lowview=None, midview=None,
highview=None):
"""Wrapper for dcloudfrac2. """Wrapper for dcloudfrac2.
Located in wrf_cloud_fracf.f90. Located in wrf_cloud_fracf.f90.
@ -750,8 +755,8 @@ def _lltoxy(map_proj, truelat1, truelat2, stdlon,
def _xytoll(map_proj, truelat1, truelat2, stdlon, lat1, lon1, def _xytoll(map_proj, truelat1, truelat2, stdlon, lat1, lon1,
pole_lat, pole_lon, known_x, known_y, dx, dy, latinc, pole_lat, pole_lon, known_x, known_y, dx, dy, latinc,
loninc, x, y, outview=None): loninc, x, y, outview=None):
"""Wrapper for dijtoll. """Wrapper for dijtoll.
Located in wrf_user_latlon_routines.f90. Located in wrf_user_latlon_routines.f90.
@ -789,9 +794,9 @@ def _xytoll(map_proj, truelat1, truelat2, stdlon, lat1, lon1,
return result return result
@check_args(0, 3, (3,3,3,3,3,3,2)) @check_args(0, 3, (3, 3, 3, 3, 3, 3, 2))
@left_iteration(3, 2, ref_var_idx=0, ignore_args=(7,8,9,10)) @left_iteration(3, 2, ref_var_idx=0, ignore_args=(7, 8, 9, 10))
@cast_type(arg_idxs=(0,1,2,3,4,5,6)) @cast_type(arg_idxs=(0, 1, 2, 3, 4, 5, 6))
@extract_and_transpose() @extract_and_transpose()
def _ctt(p_hpa, tk, qice, qcld, qv, ght, ter, haveqci, fill_nocloud, def _ctt(p_hpa, tk, qice, qcld, qv, ght, ter, haveqci, fill_nocloud,
missing, opt_thresh, outview=None): missing, opt_thresh, outview=None):
@ -822,9 +827,9 @@ def _ctt(p_hpa, tk, qice, qcld, qv, ght, ter, haveqci, fill_nocloud,
return result return result
@check_args(0, 2, (2,)) @check_args(0, 2, (2, ))
@left_iteration(2, 2, ref_var_idx=0, ignore_args=(1,2)) @left_iteration(2, 2, ref_var_idx=0, ignore_args=(1, 2))
@cast_type(arg_idxs=(0,)) @cast_type(arg_idxs=(0, ))
@extract_and_transpose() @extract_and_transpose()
def _smooth2d(field, passes, cenweight, outview=None): def _smooth2d(field, passes, cenweight, outview=None):
"""Wrapper for dfilter2d. """Wrapper for dfilter2d.
@ -856,9 +861,9 @@ def _smooth2d(field, passes, cenweight, outview=None):
return outview return outview
@check_args(0, 3, (3,3,2)) @check_args(0, 3, (3, 3, 2))
@left_iteration(3, 3, ref_var_idx=0, ignore_args=(3,4,5)) @left_iteration(3, 3, ref_var_idx=0, ignore_args=(3, 4, 5))
@cast_type(arg_idxs=(0,1,2)) @cast_type(arg_idxs=(0, 1, 2))
@extract_and_transpose() @extract_and_transpose()
def _monotonic(var, lvprs, coriolis, idir, delta, icorsw, outview=None): def _monotonic(var, lvprs, coriolis, idir, delta, icorsw, outview=None):
"""Wrapper for wrf_monotonic. """Wrapper for wrf_monotonic.
@ -885,11 +890,11 @@ def _monotonic(var, lvprs, coriolis, idir, delta, icorsw, outview=None):
# Output shape is interp_levels.shape + field.shape[-2:] # Output shape is interp_levels.shape + field.shape[-2:]
@check_args(0, 3, (3,3,3,3,3,2,2,2,3)) @check_args(0, 3, (3, 3, 3, 3, 3, 2, 2, 2, 3))
@left_iteration(3, combine_dims([(9, (-1,)), @left_iteration(3, combine_dims([(9, (-1, )),
(0, (-2,-1))]), (0, (-2, -1))]),
ref_var_idx=0, ignore_args=(9,10,11,12,13,14)) ref_var_idx=0, ignore_args=(9, 10, 11, 12, 13, 14))
@cast_type(arg_idxs=(0,1,2,3,4,5,6,7,8,9)) @cast_type(arg_idxs=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9))
@extract_and_transpose() @extract_and_transpose()
def _vintrp(field, pres, tk, qvp, ght, terrain, sfp, smsfp, def _vintrp(field, pres, tk, qvp, ght, terrain, sfp, smsfp,
vcarray, interp_levels, icase, extrap, vcor, logp, vcarray, interp_levels, icase, extrap, vcor, logp,
@ -933,9 +938,10 @@ def _vintrp(field, pres, tk, qvp, ght, terrain, sfp, smsfp,
return result return result
@check_args(0, 2, (2,2))
@check_args(0, 2, (2, 2))
@left_iteration(2, 2, ref_var_idx=0) @left_iteration(2, 2, ref_var_idx=0)
@cast_type(arg_idxs=(0,1)) @cast_type(arg_idxs=(0, 1))
@extract_and_transpose() @extract_and_transpose()
def _wspd(u, v, outview=None): def _wspd(u, v, outview=None):
"""Wrapper for dcomputewspd. """Wrapper for dcomputewspd.
@ -953,9 +959,9 @@ def _wspd(u, v, outview=None):
return result return result
@check_args(0, 2, (2,2)) @check_args(0, 2, (2, 2))
@left_iteration(2, 2, ref_var_idx=0) @left_iteration(2, 2, ref_var_idx=0)
@cast_type(arg_idxs=(0,1)) @cast_type(arg_idxs=(0, 1))
@extract_and_transpose() @extract_and_transpose()
def _wdir(u, v, outview=None): def _wdir(u, v, outview=None):
"""Wrapper for dcomputewdir. """Wrapper for dcomputewdir.
@ -1651,6 +1657,3 @@ def omp_get_wtick():
""" """
return fomp_get_wtick() return fomp_get_wtick()

43
src/wrf/g_cape.py

@ -9,6 +9,7 @@ from .constants import default_fill, Constants, ConversionFactors
from .util import extract_vars from .util import extract_vars
from .metadecorators import set_cape_metadata from .metadecorators import set_cape_metadata
@set_cape_metadata(is2d=True) @set_cape_metadata(is2d=True)
def get_2dcape(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, def get_2dcape(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True, _key=None, missing=default_fill(np.float64)): meta=True, _key=None, missing=default_fill(np.float64)):
@ -77,7 +78,7 @@ def get_2dcape(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
""" """
varnames = ("T", "P", "PB", "QVAPOR", "PH","PHB", "HGT", "PSFC") varnames = ("T", "P", "PB", "QVAPOR", "PH", "PHB", "HGT", "PSFC")
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
@ -118,10 +119,10 @@ def get_2dcape(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
# Cape 2D output is not flipped in the vertical, so index from the # Cape 2D output is not flipped in the vertical, so index from the
# end # end
result[0,...,:,:] = cape_cin[0,...,-1,:,:] result[0, ..., :, :] = cape_cin[0, ..., -1, :, :]
result[1,...,:,:] = cape_cin[1,...,-1,:,:] result[1, ..., :, :] = cape_cin[1, ..., -1, :, :]
result[2,...,:,:] = cape_cin[1,...,-2,:,:] result[2, ..., :, :] = cape_cin[1, ..., -2, :, :]
result[3,...,:,:] = cape_cin[1,...,-3,:,:] result[3, ..., :, :] = cape_cin[1, ..., -3, :, :]
return ma.masked_values(result, missing) return ma.masked_values(result, missing)
@ -221,12 +222,11 @@ def get_3dcape(wrfin, timeidx=0, method="cat",
cape_cin = _cape(p_hpa, tk, qv, z, ter, psfc_hpa, missing, i3dflag, cape_cin = _cape(p_hpa, tk, qv, z, ter, psfc_hpa, missing, i3dflag,
ter_follow) ter_follow)
return ma.masked_values(cape_cin, missing) return ma.masked_values(cape_cin, missing)
def get_cape2d_only(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, def get_cape2d_only(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True, _key=None, missing=default_fill(np.float64)): meta=True, _key=None, missing=default_fill(np.float64)):
"""Return the two-dimensional field of MCAPE (Max Convective Available """Return the two-dimensional field of MCAPE (Max Convective Available
Potential Energy). Potential Energy).
@ -284,7 +284,7 @@ def get_cape2d_only(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
""" """
result = get_2dcape(wrfin, timeidx, method, squeeze, cache, result = get_2dcape(wrfin, timeidx, method, squeeze, cache,
meta, _key, missing)[0,:] meta, _key, missing)[0, :]
if meta: if meta:
result.attrs["description"] = "mcape" result.attrs["description"] = "mcape"
@ -294,7 +294,7 @@ def get_cape2d_only(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
def get_cin2d_only(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, def get_cin2d_only(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True, _key=None, missing=default_fill(np.float64)): meta=True, _key=None, missing=default_fill(np.float64)):
"""Return the two-dimensional field of MCIN (Max Convective Inhibition). """Return the two-dimensional field of MCIN (Max Convective Inhibition).
This functions extracts the necessary variables from the NetCDF file This functions extracts the necessary variables from the NetCDF file
@ -351,7 +351,7 @@ def get_cin2d_only(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
""" """
result = get_2dcape(wrfin, timeidx, method, squeeze, cache, result = get_2dcape(wrfin, timeidx, method, squeeze, cache,
meta, _key, missing)[1,:] meta, _key, missing)[1, :]
if meta: if meta:
result.attrs["description"] = "mcin" result.attrs["description"] = "mcin"
@ -361,7 +361,7 @@ def get_cin2d_only(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
def get_lcl(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, def get_lcl(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True, _key=None, missing=default_fill(np.float64)): meta=True, _key=None, missing=default_fill(np.float64)):
"""Return the two-dimensional field of LCL (Lifted Condensation Level). """Return the two-dimensional field of LCL (Lifted Condensation Level).
This functions extracts the necessary variables from the NetCDF file This functions extracts the necessary variables from the NetCDF file
@ -418,7 +418,7 @@ def get_lcl(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
""" """
result = get_2dcape(wrfin, timeidx, method, squeeze, cache, result = get_2dcape(wrfin, timeidx, method, squeeze, cache,
meta, _key, missing)[2,:] meta, _key, missing)[2, :]
if meta: if meta:
result.attrs["description"] = "lcl" result.attrs["description"] = "lcl"
@ -428,7 +428,7 @@ def get_lcl(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
def get_lfc(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, def get_lfc(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True, _key=None, missing=default_fill(np.float64)): meta=True, _key=None, missing=default_fill(np.float64)):
"""Return the two-dimensional field of LFC (Level of Free Convection). """Return the two-dimensional field of LFC (Level of Free Convection).
This functions extracts the necessary variables from the NetCDF file This functions extracts the necessary variables from the NetCDF file
@ -485,7 +485,7 @@ def get_lfc(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
""" """
result = get_2dcape(wrfin, timeidx, method, squeeze, cache, result = get_2dcape(wrfin, timeidx, method, squeeze, cache,
meta, _key, missing)[3,:] meta, _key, missing)[3, :]
if meta: if meta:
result.attrs["description"] = "lfc" result.attrs["description"] = "lfc"
@ -495,8 +495,8 @@ def get_lfc(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
def get_3dcape_only(wrfin, timeidx=0, method="cat", def get_3dcape_only(wrfin, timeidx=0, method="cat",
squeeze=True, cache=None, meta=True, squeeze=True, cache=None, meta=True,
_key=None, missing=default_fill(np.float64)): _key=None, missing=default_fill(np.float64)):
"""Return the three-dimensional field of CAPE (Convective Available """Return the three-dimensional field of CAPE (Convective Available
Potential Energy). Potential Energy).
@ -554,7 +554,7 @@ def get_3dcape_only(wrfin, timeidx=0, method="cat",
""" """
result = get_3dcape(wrfin, timeidx, method, squeeze, cache, meta, result = get_3dcape(wrfin, timeidx, method, squeeze, cache, meta,
_key, missing)[0,:] _key, missing)[0, :]
if meta: if meta:
result.attrs["description"] = "cape" result.attrs["description"] = "cape"
@ -564,8 +564,8 @@ def get_3dcape_only(wrfin, timeidx=0, method="cat",
def get_3dcin_only(wrfin, timeidx=0, method="cat", def get_3dcin_only(wrfin, timeidx=0, method="cat",
squeeze=True, cache=None, meta=True, squeeze=True, cache=None, meta=True,
_key=None, missing=default_fill(np.float64)): _key=None, missing=default_fill(np.float64)):
"""Return the three-dimensional field of CIN (Convective Inhibition). """Return the three-dimensional field of CIN (Convective Inhibition).
This functions extracts the necessary variables from the NetCDF file This functions extracts the necessary variables from the NetCDF file
@ -622,13 +622,10 @@ def get_3dcin_only(wrfin, timeidx=0, method="cat",
""" """
result = get_3dcape(wrfin, timeidx, method, squeeze, cache, meta, result = get_3dcape(wrfin, timeidx, method, squeeze, cache, meta,
_key, missing)[1,:] _key, missing)[1, :]
if meta: if meta:
result.attrs["description"] = "cin" result.attrs["description"] = "cin"
result.attrs["units"] = "J kg-1" result.attrs["units"] = "J kg-1"
return result return result

48
src/wrf/g_cloudfrac.py

@ -12,9 +12,9 @@ from .g_geoht import _get_geoht
@set_cloudfrac_metadata() @set_cloudfrac_metadata()
def get_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True, def get_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, _key=None, cache=None, meta=True, _key=None,
vert_type="height_agl", low_thresh=None, mid_thresh=None, vert_type="height_agl", low_thresh=None, mid_thresh=None,
high_thresh=None, missing=default_fill(np.float64)): high_thresh=None, missing=default_fill(np.float64)):
"""Return the cloud fraction for low, mid, and high level clouds. """Return the cloud fraction for low, mid, and high level clouds.
The leftmost dimension of the returned array represents three different The leftmost dimension of the returned array represents three different
@ -157,15 +157,16 @@ def get_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True,
"or 'height_agl'") "or 'height_agl'")
cfrac = _cloudfrac(v_coord, rh, vert_inc_w_height, cfrac = _cloudfrac(v_coord, rh, vert_inc_w_height,
_low_thresh, _mid_thresh, _high_thresh, missing) _low_thresh, _mid_thresh, _high_thresh, missing)
return ma.masked_values(cfrac, missing) return ma.masked_values(cfrac, missing)
def get_low_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True, def get_low_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, _key=None, cache=None, meta=True, _key=None,
vert_type="height_agl", low_thresh=None, mid_thresh=None, vert_type="height_agl", low_thresh=None,
high_thresh=None, missing=default_fill(np.float64)): mid_thresh=None, high_thresh=None,
missing=default_fill(np.float64)):
"""Return the cloud fraction for the low level clouds. """Return the cloud fraction for the low level clouds.
If the vertical coordinate type is 'height_agl' or 'height_msl', the If the vertical coordinate type is 'height_agl' or 'height_msl', the
@ -263,10 +264,9 @@ def get_low_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True,
be a :class:`numpy.ndarray` object with no metadata. be a :class:`numpy.ndarray` object with no metadata.
""" """
result = get_cloudfrac(wrfin, timeidx, method, squeeze, result = get_cloudfrac(wrfin, timeidx, method, squeeze,
cache, meta, _key, cache, meta, _key, vert_type, low_thresh,
vert_type, low_thresh, mid_thresh, mid_thresh, high_thresh, missing)[0, :]
high_thresh, missing)[0,:]
if meta: if meta:
result.attrs["description"] = "low clouds" result.attrs["description"] = "low clouds"
@ -275,9 +275,10 @@ def get_low_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True,
def get_mid_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True, def get_mid_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, _key=None, cache=None, meta=True, _key=None,
vert_type="height_agl", low_thresh=None, mid_thresh=None, vert_type="height_agl", low_thresh=None,
high_thresh=None, missing=default_fill(np.float64)): mid_thresh=None, high_thresh=None,
missing=default_fill(np.float64)):
"""Return the cloud fraction for the mid level clouds. """Return the cloud fraction for the mid level clouds.
If the vertical coordinate type is 'height_agl' or 'height_msl', the If the vertical coordinate type is 'height_agl' or 'height_msl', the
@ -376,9 +377,9 @@ def get_mid_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True,
""" """
result = get_cloudfrac(wrfin, timeidx, method, squeeze, result = get_cloudfrac(wrfin, timeidx, method, squeeze,
cache, meta, _key, cache, meta, _key,
vert_type, low_thresh, mid_thresh, vert_type, low_thresh, mid_thresh,
high_thresh, missing)[1,:] high_thresh, missing)[1, :]
if meta: if meta:
result.attrs["description"] = "mid clouds" result.attrs["description"] = "mid clouds"
@ -387,9 +388,10 @@ def get_mid_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True,
def get_high_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True, def get_high_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, _key=None, cache=None, meta=True, _key=None,
vert_type="height_agl", low_thresh=None, mid_thresh=None, vert_type="height_agl", low_thresh=None,
high_thresh=None, missing=default_fill(np.float64)): mid_thresh=None, high_thresh=None,
missing=default_fill(np.float64)):
"""Return the cloud fraction for the high level clouds. """Return the cloud fraction for the high level clouds.
If the vertical coordinate type is 'height_agl' or 'height_msl', the If the vertical coordinate type is 'height_agl' or 'height_msl', the
@ -488,9 +490,9 @@ def get_high_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True,
""" """
result = get_cloudfrac(wrfin, timeidx, method, squeeze, result = get_cloudfrac(wrfin, timeidx, method, squeeze,
cache, meta, _key, cache, meta, _key,
vert_type, low_thresh, mid_thresh, vert_type, low_thresh, mid_thresh,
high_thresh, missing)[2,:] high_thresh, missing)[2, :]
if meta: if meta:
result.attrs["description"] = "high clouds" result.attrs["description"] = "high clouds"

7
src/wrf/g_ctt.py

@ -3,7 +3,6 @@ from __future__ import (absolute_import, division, print_function)
import numpy as np import numpy as np
import numpy.ma as ma import numpy.ma as ma
#from .extension import computectt, computetk
from .extension import _ctt, _tk from .extension import _ctt, _tk
from .constants import Constants, ConversionFactors, default_fill from .constants import Constants, ConversionFactors, default_fill
from .destag import destagger from .destag import destagger
@ -107,7 +106,7 @@ def get_ctt(wrfin, timeidx=0, method="cat",
ph = ncvars["PH"] ph = ncvars["PH"]
phb = ncvars["PHB"] phb = ncvars["PHB"]
ter = ncvars["HGT"] ter = ncvars["HGT"]
qv = ncvars["QVAPOR"] * 1000.0 # g/kg qv = ncvars["QVAPOR"] * 1000.0 # g/kg
haveqci = 1 haveqci = 1
try: try:
@ -118,7 +117,7 @@ def get_ctt(wrfin, timeidx=0, method="cat",
qice = np.zeros(qv.shape, qv.dtype) qice = np.zeros(qv.shape, qv.dtype)
haveqci = 0 haveqci = 0
else: else:
qice = icevars["QICE"] * 1000.0 #g/kg qice = icevars["QICE"] * 1000.0 # g/kg
try: try:
cldvars = extract_vars(wrfin, timeidx, "QCLOUD", cldvars = extract_vars(wrfin, timeidx, "QCLOUD",
@ -127,7 +126,7 @@ def get_ctt(wrfin, timeidx=0, method="cat",
except KeyError: except KeyError:
raise RuntimeError("'QCLOUD' not found in NetCDF file") raise RuntimeError("'QCLOUD' not found in NetCDF file")
else: else:
qcld = cldvars["QCLOUD"] * 1000.0 #g/kg qcld = cldvars["QCLOUD"] * 1000.0 # g/kg
full_p = p + pb full_p = p + pb
p_hpa = full_p * ConversionFactors.PA_TO_HPA p_hpa = full_p * ConversionFactors.PA_TO_HPA

6
src/wrf/g_dbz.py

@ -2,7 +2,6 @@ from __future__ import (absolute_import, division, print_function)
import numpy as np import numpy as np
#from .extension import computedbz,computetk
from .extension import _dbz, _tk from .extension import _dbz, _tk
from .constants import Constants from .constants import Constants
from .util import extract_vars, to_np from .util import extract_vars, to_np
@ -194,6 +193,5 @@ def get_max_dbz(wrfin, timeidx=0, method="cat",
""" """
return np.amax(to_np(get_dbz(wrfin, timeidx, method, squeeze, cache, meta, return np.amax(to_np(get_dbz(wrfin, timeidx, method, squeeze, cache, meta,
_key, use_varint, use_liqskin)), _key, use_varint, use_liqskin)),
axis=-3) axis=-3)

7
src/wrf/g_dewpoint.py

@ -1,6 +1,5 @@
from __future__ import (absolute_import, division, print_function) from __future__ import (absolute_import, division, print_function)
#from .extension import computetd
from .extension import _td from .extension import _td
from .decorators import convert_units from .decorators import convert_units
from .metadecorators import copy_and_set_metadata from .metadecorators import copy_and_set_metadata
@ -69,7 +68,7 @@ def get_dp(wrfin, timeidx=0, method="cat", squeeze=True,
""" """
varnames=("P", "PB", "QVAPOR") varnames = ("P", "PB", "QVAPOR")
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
@ -86,6 +85,7 @@ def get_dp(wrfin, timeidx=0, method="cat", squeeze=True,
td = _td(full_p, qvapor) td = _td(full_p, qvapor)
return td return td
@copy_and_set_metadata(copy_varname="Q2", name="td2", @copy_and_set_metadata(copy_varname="Q2", name="td2",
description="2m dew point temperature") description="2m dew point temperature")
@convert_units("temp", "c") @convert_units("temp", "c")
@ -147,7 +147,7 @@ def get_dp_2m(wrfin, timeidx=0, method="cat", squeeze=True,
be a :class:`numpy.ndarray` object with no metadata. be a :class:`numpy.ndarray` object with no metadata.
""" """
varnames=("PSFC", "Q2") varnames = ("PSFC", "Q2")
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
@ -161,4 +161,3 @@ def get_dp_2m(wrfin, timeidx=0, method="cat", squeeze=True,
td = _td(psfc, q2) td = _td(psfc, q2)
return td return td

11
src/wrf/g_geoht.py

@ -8,6 +8,7 @@ from .decorators import convert_units
from .metadecorators import set_height_metadata from .metadecorators import set_height_metadata
from .util import extract_vars, either from .util import extract_vars, either
def _get_geoht(wrfin, timeidx, method="cat", squeeze=True, def _get_geoht(wrfin, timeidx, method="cat", squeeze=True,
cache=None, meta=True, _key=None, cache=None, meta=True, _key=None,
height=True, msl=True, stag=False): height=True, msl=True, stag=False):
@ -104,7 +105,7 @@ def _get_geoht(wrfin, timeidx, method="cat", squeeze=True,
if stag: if stag:
warnings.warn("file contains no vertically staggered geopotential " warnings.warn("file contains no vertically staggered geopotential "
"height variable, returning unstaggered result " "height variable, returning unstaggered result "
"instead" ) "instead")
if height: if height:
if msl: if msl:
return geopt_unstag / Constants.G return geopt_unstag / Constants.G
@ -113,7 +114,7 @@ def _get_geoht(wrfin, timeidx, method="cat", squeeze=True,
# array needs to be reshaped to a 3D array so the right dims # array needs to be reshaped to a 3D array so the right dims
# line up # line up
new_dims = list(hgt.shape) new_dims = list(hgt.shape)
new_dims.insert(-2,1) new_dims.insert(-2, 1)
hgt = hgt.reshape(new_dims) hgt = hgt.reshape(new_dims)
return (geopt_unstag / Constants.G) - hgt return (geopt_unstag / Constants.G) - hgt
@ -259,7 +260,7 @@ def get_height(wrfin, timeidx=0, method="cat", squeeze=True,
@set_height_metadata(geopt=True, stag=True) @set_height_metadata(geopt=True, stag=True)
def get_stag_geopt(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, def get_stag_geopt(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True, _key=None): meta=True, _key=None):
"""Return the geopotential for the vertically staggered grid. """Return the geopotential for the vertically staggered grid.
The geopotential is returned in units of [m2 s-2]. The geopotential is returned in units of [m2 s-2].
@ -321,8 +322,7 @@ def get_stag_geopt(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
@set_height_metadata(geopt=False, stag=True) @set_height_metadata(geopt=False, stag=True)
@convert_units("height", "m") @convert_units("height", "m")
def get_stag_height(wrfin, timeidx=0, method="cat", squeeze=True, def get_stag_height(wrfin, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, _key=None, cache=None, meta=True, _key=None, msl=True, units="m"):
msl=True, units="m"):
"""Return the geopotential height for the vertically staggered grid. """Return the geopotential height for the vertically staggered grid.
If *msl* is True, then geopotential height is returned as Mean Sea Level If *msl* is True, then geopotential height is returned as Mean Sea Level
@ -391,4 +391,3 @@ def get_stag_height(wrfin, timeidx=0, method="cat", squeeze=True,
return _get_geoht(wrfin, timeidx, method, squeeze, cache, meta, _key, return _get_geoht(wrfin, timeidx, method, squeeze, cache, meta, _key,
True, msl, stag=True) True, msl, stag=True)

5
src/wrf/g_helicity.py

@ -9,6 +9,7 @@ from .util import extract_vars, extract_global_attrs, either
from .metadecorators import copy_and_set_metadata from .metadecorators import copy_and_set_metadata
from .g_latlon import get_lat from .g_latlon import get_lat
@copy_and_set_metadata(copy_varname="HGT", name="srh", @copy_and_set_metadata(copy_varname="HGT", name="srh",
description="storm relative helicity", description="storm relative helicity",
units="m2 s-2") units="m2 s-2")
@ -207,7 +208,3 @@ def get_uh(wrfin, timeidx=0, method="cat", squeeze=True,
return uh return uh

38
src/wrf/g_latlon.py

@ -244,7 +244,7 @@ def _llxy_mapping(wrfin, x_or_lat, y_or_lon, func, timeidx, stagger,
outdims = [numkeys] outdims = [numkeys]
outdims += first_array.shape outdims += first_array.shape
outdata = np.empty(outdims, first_array.dtype) outdata = np.empty(outdims, first_array.dtype)
outdata[0,:] = first_array[:] outdata[0, :] = first_array[:]
idx = 1 idx = 1
while True: while True:
@ -265,7 +265,7 @@ def _llxy_mapping(wrfin, x_or_lat, y_or_lon, func, timeidx, stagger,
if outdata.shape[1:] != result_array.shape: if outdata.shape[1:] != result_array.shape:
raise ValueError("data sequences must have the " raise ValueError("data sequences must have the "
"same size for all dictionary keys") "same size for all dictionary keys")
outdata[idx,:] = to_np(result_array)[:] outdata[idx, :] = to_np(result_array)[:]
idx += 1 idx += 1
if xarray_enabled() and meta: if xarray_enabled() and meta:
@ -301,12 +301,10 @@ def _llxy_mapping(wrfin, x_or_lat, y_or_lon, func, timeidx, stagger,
# make it so that key_0 is leftmost # make it so that key_0 is leftmost
outdims = key_coordnames + list(first_array.dims[existing_cnt:]) outdims = key_coordnames + list(first_array.dims[existing_cnt:])
# Create the new 'key_n', value pairs # Create the new 'key_n', value pairs
for coordname, coordval in zip(key_coordnames, coord_vals): for coordname, coordval in zip(key_coordnames, coord_vals):
outcoords[coordname] = coordval outcoords[coordname] = coordval
outattrs = OrderedDict(first_array.attrs) outattrs = OrderedDict(first_array.attrs)
outarr = DataArray(outdata, name=outname, coords=outcoords, outarr = DataArray(outdata, name=outname, coords=outcoords,
@ -473,7 +471,7 @@ def _set_defaults(projparams):
if _params.get("latinc") is None: if _params.get("latinc") is None:
_params["latinc"] = ((_params["dy"]*360.0)/2.0 / _params["latinc"] = ((_params["dy"]*360.0)/2.0 /
Constants.PI/Constants.WRF_EARTH_RADIUS) Constants.PI/Constants.WRF_EARTH_RADIUS)
if _params.get("loninc") is None: if _params.get("loninc") is None:
_params["loninc"] = ((_params["dx"]*360.0)/2.0 / _params["loninc"] = ((_params["dx"]*360.0)/2.0 /
@ -604,11 +602,12 @@ def ll_to_xy_proj(latitude, longitude, meta=True, squeeze=True, as_int=True,
""" """
loc = locals() loc = locals()
_projparams = {name : loc[name] for name in ("map_proj", "truelat1", _projparams = {name: loc[name] for name in ("map_proj", "truelat1",
"truelat2", "stand_lon", "ref_lat", "truelat2", "stand_lon",
"ref_lon", "pole_lat", "pole_lon", "ref_lat", "ref_lon",
"known_x", "known_y", "dx", "dy", "pole_lat", "pole_lon",
"latinc", "loninc")} "known_x", "known_y", "dx",
"dy", "latinc", "loninc")}
projparams = _set_defaults(_projparams) projparams = _set_defaults(_projparams)
@ -772,20 +771,13 @@ def xy_to_ll_proj(x, y, meta=True, squeeze=True, map_proj=None, truelat1=None,
be a :class:`numpy.ndarray` object with no metadata. be a :class:`numpy.ndarray` object with no metadata.
""" """
loc = locals() loc = locals()
_projparams = {name : loc[name] for name in ("map_proj", "truelat1", _projparams = {name: loc[name] for name in ("map_proj", "truelat1",
"truelat2", "stand_lon", "ref_lat", "truelat2", "stand_lon",
"ref_lon", "pole_lat", "pole_lon", "ref_lat", "ref_lon",
"known_x", "known_y", "dx", "dy", "pole_lat", "pole_lon",
"latinc", "loninc")} "known_x", "known_y", "dx",
"dy", "latinc", "loninc")}
projparams = _set_defaults(_projparams) projparams = _set_defaults(_projparams)
return _xy_to_ll(x, y, None, 0, None, "cat", squeeze, None, None, return _xy_to_ll(x, y, None, 0, None, "cat", squeeze, None, None,
**projparams) **projparams)

3
src/wrf/g_omega.py

@ -64,7 +64,7 @@ def get_omega(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
:class:`numpy.ndarray` object with no metadata. :class:`numpy.ndarray` object with no metadata.
""" """
varnames=("T", "P", "W", "PB", "QVAPOR") varnames = ("T", "P", "W", "PB", "QVAPOR")
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
t = ncvars["T"] t = ncvars["T"]
@ -81,4 +81,3 @@ def get_omega(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
omega = _omega(qv, tk, wa, full_p) omega = _omega(qv, tk, wa, full_p)
return omega return omega

2
src/wrf/g_precip.py

@ -2,7 +2,6 @@ from __future__ import (absolute_import, division, print_function)
from .util import extract_vars from .util import extract_vars
__all__ = ["get_accum_precip", "get_precip_diff"]
def get_accum_precip(wrfin, timeidx=0): def get_accum_precip(wrfin, timeidx=0):
ncvars = extract_vars(wrfin, timeidx, varnames=("RAINC", "RAINNC")) ncvars = extract_vars(wrfin, timeidx, varnames=("RAINC", "RAINNC"))
@ -13,6 +12,7 @@ def get_accum_precip(wrfin, timeidx=0):
return rainsum return rainsum
def get_precip_diff(wrfin1, wrfin2, timeidx=0): def get_precip_diff(wrfin1, wrfin2, timeidx=0):
vars1 = extract_vars(wrfin1, timeidx, varnames=("RAINC", "RAINNC")) vars1 = extract_vars(wrfin1, timeidx, varnames=("RAINC", "RAINNC"))
vars2 = extract_vars(wrfin2, timeidx, varnames=("RAINC", "RAINNC")) vars2 = extract_vars(wrfin2, timeidx, varnames=("RAINC", "RAINNC"))

5
src/wrf/g_pressure.py

@ -85,6 +85,7 @@ def get_pressure(wrfin, timeidx=0, method="cat", squeeze=True,
return pres return pres
def get_pressure_hpa(wrfin, timeidx=0, method="cat", squeeze=True, def get_pressure_hpa(wrfin, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, _key=None): cache=None, meta=True, _key=None):
"""Return the pressure in [hPa]. """Return the pressure in [hPa].
@ -142,7 +143,3 @@ def get_pressure_hpa(wrfin, timeidx=0, method="cat", squeeze=True,
""" """
return get_pressure(wrfin, timeidx, method, squeeze, cache, meta, _key, return get_pressure(wrfin, timeidx, method, squeeze, cache, meta, _key,
units="hPa") units="hPa")

7
src/wrf/g_pw.py

@ -1,6 +1,5 @@
from __future__ import (absolute_import, division, print_function) from __future__ import (absolute_import, division, print_function)
#from .extension import computepw,computetv,computetk
from .extension import _pw, _tv, _tk from .extension import _pw, _tv, _tk
from .constants import Constants from .constants import Constants
from .util import extract_vars from .util import extract_vars
@ -66,7 +65,7 @@ def get_pw(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
:class:`numpy.ndarray` object with no metadata. :class:`numpy.ndarray` object with no metadata.
""" """
varnames=("T", "P", "PB", "PH", "PHB", "QVAPOR") varnames = ("T", "P", "PB", "PH", "PHB", "QVAPOR")
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
@ -85,7 +84,3 @@ def get_pw(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
tv = _tv(tk, qv) tv = _tv(tk, qv)
return _pw(full_p, tv, qv, ht) return _pw(full_p, tv, qv, ht)

6
src/wrf/g_rh.py

@ -1,7 +1,6 @@
from __future__ import (absolute_import, division, print_function) from __future__ import (absolute_import, division, print_function)
from .constants import Constants from .constants import Constants
#from .extension import computerh, computetk
from .extension import _rh, _tk from .extension import _rh, _tk
from .util import extract_vars from .util import extract_vars
from .metadecorators import copy_and_set_metadata from .metadecorators import copy_and_set_metadata
@ -64,7 +63,7 @@ def get_rh(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
:class:`numpy.ndarray` object with no metadata. :class:`numpy.ndarray` object with no metadata.
""" """
varnames=("T", "P", "PB", "QVAPOR") varnames = ("T", "P", "PB", "QVAPOR")
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
t = ncvars["T"] t = ncvars["T"]
@ -140,7 +139,7 @@ def get_rh_2m(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
:class:`numpy.ndarray` object with no metadata. :class:`numpy.ndarray` object with no metadata.
""" """
varnames=("T2", "PSFC", "Q2") varnames = ("T2", "PSFC", "Q2")
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
t2 = ncvars["T2"] t2 = ncvars["T2"]
@ -153,4 +152,3 @@ def get_rh_2m(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
rh = _rh(q2, psfc, t2) rh = _rh(q2, psfc, t2)
return rh return rh

4
src/wrf/g_slp.py

@ -1,6 +1,5 @@
from __future__ import (absolute_import, division, print_function) from __future__ import (absolute_import, division, print_function)
#from .extension import computeslp, computetk
from .extension import _slp, _tk from .extension import _slp, _tk
from .constants import Constants from .constants import Constants
from .destag import destagger from .destag import destagger
@ -75,7 +74,7 @@ def get_slp(wrfin, timeidx=0, method="cat", squeeze=True,
:class:`numpy.ndarray` object with no metadata. :class:`numpy.ndarray` object with no metadata.
""" """
varnames=("T", "P", "PB", "QVAPOR", "PH", "PHB") varnames = ("T", "P", "PB", "QVAPOR", "PH", "PHB")
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
@ -100,4 +99,3 @@ def get_slp(wrfin, timeidx=0, method="cat", squeeze=True,
slp = _slp(destag_ph, tk, full_p, qvapor) slp = _slp(destag_ph, tk, full_p, qvapor)
return slp return slp

13
src/wrf/g_temp.py

@ -69,7 +69,7 @@ def get_theta(wrfin, timeidx=0, method="cat", squeeze=True,
be a :class:`numpy.ndarray` object with no metadata. be a :class:`numpy.ndarray` object with no metadata.
""" """
varnames = ("T",) varnames = ("T", )
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
@ -142,7 +142,7 @@ def get_temp(wrfin, timeidx=0, method="cat", squeeze=True,
""" """
varnames=("T", "P", "PB") varnames = ("T", "P", "PB")
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
t = ncvars["T"] t = ncvars["T"]
@ -219,7 +219,7 @@ def get_eth(wrfin, timeidx=0, method="cat", squeeze=True,
""" """
varnames=("T", "P", "PB", "QVAPOR") varnames = ("T", "P", "PB", "QVAPOR")
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
t = ncvars["T"] t = ncvars["T"]
@ -299,7 +299,7 @@ def get_tv(wrfin, timeidx=0, method="cat", squeeze=True,
""" """
varnames=("T", "P", "PB", "QVAPOR") varnames = ("T", "P", "PB", "QVAPOR")
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
@ -380,7 +380,7 @@ def get_tw(wrfin, timeidx=0, method="cat", squeeze=True,
""" """
varnames=("T", "P", "PB", "QVAPOR") varnames = ("T", "P", "PB", "QVAPOR")
ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
t = ncvars["T"] t = ncvars["T"]
@ -511,6 +511,3 @@ def get_tc(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
""" """
return get_temp(wrfin, timeidx, method, squeeze, cache, meta, _key, return get_temp(wrfin, timeidx, method, squeeze, cache, meta, _key,
units="degC") units="degC")

5
src/wrf/g_terrain.py

@ -10,7 +10,7 @@ from .util import extract_vars, either
description="terrain height") description="terrain height")
@convert_units("height", "m") @convert_units("height", "m")
def get_terrain(wrfin, timeidx=0, method="cat", squeeze=True, def get_terrain(wrfin, timeidx=0, method="cat", squeeze=True,
cache=None, meta=False, _key=None, units="m"): cache=None, meta=False, _key=None, units="m"):
"""Return the terrain height in the specified units. """Return the terrain height in the specified units.
This functions extracts the necessary variables from the NetCDF file This functions extracts the necessary variables from the NetCDF file
@ -73,6 +73,3 @@ def get_terrain(wrfin, timeidx=0, method="cat", squeeze=True,
return extract_vars(wrfin, timeidx, varname, return extract_vars(wrfin, timeidx, varname,
method, squeeze, cache, meta=False, method, squeeze, cache, meta=False,
_key=_key)[varname] _key=_key)[varname]

1
src/wrf/g_times.py

@ -118,4 +118,3 @@ def get_xtimes(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
""" """
return extract_times(wrfin, timeidx, method, squeeze, cache, return extract_times(wrfin, timeidx, method, squeeze, cache,
meta=meta, do_xtime=True) meta=meta, do_xtime=True)

41
src/wrf/g_uvmet.py

@ -4,7 +4,6 @@ from math import fabs, log, tan, sin, cos
import numpy as np import numpy as np
#from .extension import computeuvmet
from .extension import _uvmet from .extension import _uvmet
from .destag import destagger from .destag import destagger
from .constants import Constants from .constants import Constants
@ -101,13 +100,13 @@ def _get_uvmet(wrfin, timeidx=0, method="cat", squeeze=True,
u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
u = (u_vars[varname] if varname == "U10" else u = (u_vars[varname] if varname == "U10" else
destagger(u_vars[varname][...,0,:,:], -1)) destagger(u_vars[varname][..., 0, :, :], -1))
varname = either("V10", "VV")(wrfin) varname = either("V10", "VV")(wrfin)
v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
v = (v_vars[varname] if varname == "V10" else v = (v_vars[varname] if varname == "V10" else
destagger(v_vars[varname][...,0,:,:], -2)) destagger(v_vars[varname][..., 0, :, :], -2))
map_proj_attrs = extract_global_attrs(wrfin, attrs="MAP_PROJ") map_proj_attrs = extract_global_attrs(wrfin, attrs="MAP_PROJ")
map_proj = map_proj_attrs["MAP_PROJ"] map_proj = map_proj_attrs["MAP_PROJ"]
@ -119,7 +118,7 @@ def _get_uvmet(wrfin, timeidx=0, method="cat", squeeze=True,
# Note: NCL has no code to handle other projections (0,4,5) as they # Note: NCL has no code to handle other projections (0,4,5) as they
# don't appear to be supported any longer # don't appear to be supported any longer
if map_proj in (0,3,6): if map_proj in (0, 3, 6):
# No rotation needed for Mercator and Lat/Lon, but still need # No rotation needed for Mercator and Lat/Lon, but still need
# u,v aggregated in to one array # u,v aggregated in to one array
@ -145,7 +144,7 @@ def _get_uvmet(wrfin, timeidx=0, method="cat", squeeze=True,
result = np.ma.masked_values(result, fill) result = np.ma.masked_values(result, fill)
return result return result
elif map_proj in (1,2): elif map_proj in (1, 2):
lat_attrs = extract_global_attrs(wrfin, attrs=("TRUELAT1", lat_attrs = extract_global_attrs(wrfin, attrs=("TRUELAT1",
"TRUELAT2")) "TRUELAT2"))
radians_per_degree = Constants.PI/180.0 radians_per_degree = Constants.PI/180.0
@ -165,7 +164,6 @@ def _get_uvmet(wrfin, timeidx=0, method="cat", squeeze=True,
else: else:
cen_lon = lon_attrs["STAND_LON"] cen_lon = lon_attrs["STAND_LON"]
varname = either("XLAT_M", "XLAT")(wrfin) varname = either("XLAT_M", "XLAT")(wrfin)
xlat_var = extract_vars(wrfin, timeidx, varname, xlat_var = extract_vars(wrfin, timeidx, varname,
method, squeeze, cache, meta=False, method, squeeze, cache, meta=False,
@ -181,11 +179,12 @@ def _get_uvmet(wrfin, timeidx=0, method="cat", squeeze=True,
if map_proj == 1: if map_proj == 1:
if((fabs(true_lat1 - true_lat2) > 0.1) and if((fabs(true_lat1 - true_lat2) > 0.1) and
(fabs(true_lat2 - 90.) > 0.1)): (fabs(true_lat2 - 90.) > 0.1)):
cone = (log(cos(true_lat1*radians_per_degree)) cone = (log(cos(true_lat1*radians_per_degree)) -
- log(cos(true_lat2*radians_per_degree))) log(cos(true_lat2*radians_per_degree)))
cone = (cone / cone = (cone /
(log(tan((45.-fabs(true_lat1/2.))*radians_per_degree)) (log(tan((45.-fabs(true_lat1/2.))*radians_per_degree))
- log(tan((45.-fabs(true_lat2/2.))*radians_per_degree)))) - log(tan((45.-fabs(true_lat2/2.)) *
radians_per_degree))))
else: else:
cone = sin(fabs(true_lat1)*radians_per_degree) cone = sin(fabs(true_lat1)*radians_per_degree)
else: else:
@ -434,7 +433,7 @@ def get_uvmet_wspd_wdir(wrfin, timeidx=0, method="cat", squeeze=True,
uvmet = _get_uvmet(wrfin, timeidx, method, squeeze, uvmet = _get_uvmet(wrfin, timeidx, method, squeeze,
cache, meta, _key, False, units="m s-1") cache, meta, _key, False, units="m s-1")
return _calc_wspd_wdir(uvmet[0,...,:,:,:], uvmet[1,...,:,:,:], return _calc_wspd_wdir(uvmet[0, ..., :, :, :], uvmet[1, ..., :, :, :],
False, units) False, units)
@ -517,7 +516,8 @@ def get_uvmet10_wspd_wdir(wrfin, timeidx=0, method="cat", squeeze=True,
uvmet10 = _get_uvmet(wrfin, timeidx, method, squeeze, cache, meta, _key, uvmet10 = _get_uvmet(wrfin, timeidx, method, squeeze, cache, meta, _key,
True, units="m s-1") True, units="m s-1")
return _calc_wspd_wdir(uvmet10[0,...,:,:], uvmet10[1,...,:,:], True, units) return _calc_wspd_wdir(uvmet10[0, ..., :, :], uvmet10[1, ..., :, :],
True, units)
def get_uvmet_wspd(wrfin, timeidx=0, method="cat", squeeze=True, def get_uvmet_wspd(wrfin, timeidx=0, method="cat", squeeze=True,
@ -583,7 +583,7 @@ def get_uvmet_wspd(wrfin, timeidx=0, method="cat", squeeze=True,
""" """
result = get_uvmet_wspd_wdir(wrfin, timeidx, method, squeeze, result = get_uvmet_wspd_wdir(wrfin, timeidx, method, squeeze,
cache, meta, _key, units)[0,:] cache, meta, _key, units)[0, :]
if meta: if meta:
result.attrs["description"] = "earth rotated wspd" result.attrs["description"] = "earth rotated wspd"
@ -654,7 +654,7 @@ def get_uvmet_wdir(wrfin, timeidx=0, method="cat", squeeze=True,
""" """
result = get_uvmet_wspd_wdir(wrfin, timeidx, method, squeeze, result = get_uvmet_wspd_wdir(wrfin, timeidx, method, squeeze,
cache, meta, _key, units)[1,:] cache, meta, _key, units)[1, :]
if meta: if meta:
result.attrs["description"] = "earth rotated wdir" result.attrs["description"] = "earth rotated wdir"
@ -663,8 +663,8 @@ def get_uvmet_wdir(wrfin, timeidx=0, method="cat", squeeze=True,
def get_uvmet10_wspd(wrfin, timeidx=0, method="cat", squeeze=True, def get_uvmet10_wspd(wrfin, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, _key=None, cache=None, meta=True, _key=None,
units="m s-1"): units="m s-1"):
"""Return the wind speed for the 10m wind rotated to earth coordinates. """Return the wind speed for the 10m wind rotated to earth coordinates.
This function should be used when comparing with observations. This function should be used when comparing with observations.
@ -725,7 +725,7 @@ def get_uvmet10_wspd(wrfin, timeidx=0, method="cat", squeeze=True,
""" """
result = get_uvmet10_wspd_wdir(wrfin, timeidx, method, squeeze, result = get_uvmet10_wspd_wdir(wrfin, timeidx, method, squeeze,
cache, meta, _key, units)[0,:] cache, meta, _key, units)[0, :]
if meta: if meta:
result.attrs["description"] = "10m earth rotated wspd" result.attrs["description"] = "10m earth rotated wspd"
@ -733,8 +733,8 @@ def get_uvmet10_wspd(wrfin, timeidx=0, method="cat", squeeze=True,
def get_uvmet10_wdir(wrfin, timeidx=0, method="cat", squeeze=True, def get_uvmet10_wdir(wrfin, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, _key=None, cache=None, meta=True, _key=None,
units="m s-1"): units="m s-1"):
"""Return the wind direction for the 10m wind rotated to earth coordinates. """Return the wind direction for the 10m wind rotated to earth coordinates.
This function should be used when comparing with observations. This function should be used when comparing with observations.
@ -795,12 +795,9 @@ def get_uvmet10_wdir(wrfin, timeidx=0, method="cat", squeeze=True,
""" """
result = get_uvmet10_wspd_wdir(wrfin, timeidx, method, squeeze, result = get_uvmet10_wspd_wdir(wrfin, timeidx, method, squeeze,
cache, meta, _key, units)[1,:] cache, meta, _key, units)[1, :]
if meta: if meta:
result.attrs["description"] = "10m earth rotated wdir" result.attrs["description"] = "10m earth rotated wdir"
return result return result

1
src/wrf/g_vorticity.py

@ -164,4 +164,3 @@ def get_pvo(wrfin, timeidx=0, method="cat", squeeze=True, cache=None,
full_p = p + pb full_p = p + pb
return _pvo(u, v, full_t, full_p, msfu, msfv, msfm, cor, dx, dy) return _pvo(u, v, full_t, full_p, msfu, msfv, msfm, cor, dx, dy)

21
src/wrf/g_wind.py

@ -92,13 +92,11 @@ def _calc_wspd_wdir(u, v, two_d, units):
result = np.zeros(outdims, wspd.dtype) result = np.zeros(outdims, wspd.dtype)
idxs0 = ((0,Ellipsis, slice(None), slice(None), slice(None)) idxs0 = ((0, Ellipsis, slice(None), slice(None), slice(None))
if not two_d else if not two_d else (0, Ellipsis, slice(None), slice(None)))
(0, Ellipsis, slice(None), slice(None)))
idxs1 = ((1, Ellipsis, slice(None), slice(None), slice(None)) idxs1 = ((1, Ellipsis, slice(None), slice(None), slice(None))
if not two_d else if not two_d else (1, Ellipsis, slice(None), slice(None)))
(1, Ellipsis, slice(None), slice(None)))
result[idxs0] = wspd[:] result[idxs0] = wspd[:]
result[idxs1] = wdir[:] result[idxs1] = wdir[:]
@ -497,13 +495,13 @@ def get_destag_wspd_wdir10(wrfin, timeidx=0, method="cat",
u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
u = (u_vars[varname] if varname == "U10" else u = (u_vars[varname] if varname == "U10" else
destagger(u_vars[varname][...,0,:,:], -1)) destagger(u_vars[varname][..., 0, :, :], -1))
varname = either("V10", "VV")(wrfin) varname = either("V10", "VV")(wrfin)
v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache,
meta=False, _key=_key) meta=False, _key=_key)
v = (v_vars[varname] if varname == "V10" else v = (v_vars[varname] if varname == "V10" else
destagger(v_vars[varname][...,0,:,:], -2)) destagger(v_vars[varname][..., 0, :, :], -2))
return _calc_wspd_wdir(u, v, True, units) return _calc_wspd_wdir(u, v, True, units)
@ -569,7 +567,7 @@ def get_destag_wspd(wrfin, timeidx=0, method="cat",
""" """
result = get_destag_wspd_wdir(wrfin, timeidx, method, squeeze, cache, result = get_destag_wspd_wdir(wrfin, timeidx, method, squeeze, cache,
meta, _key, units)[0,:] meta, _key, units)[0, :]
if meta: if meta:
result.attrs["description"] = "wspd in projection space" result.attrs["description"] = "wspd in projection space"
@ -638,7 +636,7 @@ def get_destag_wdir(wrfin, timeidx=0, method="cat",
""" """
result = get_destag_wspd_wdir(wrfin, timeidx, method, squeeze, cache, result = get_destag_wspd_wdir(wrfin, timeidx, method, squeeze, cache,
meta, _key, units)[1,:] meta, _key, units)[1, :]
if meta: if meta:
result.attrs["description"] = "wdir in projection space" result.attrs["description"] = "wdir in projection space"
@ -708,7 +706,7 @@ def get_destag_wspd10(wrfin, timeidx=0, method="cat",
""" """
result = get_destag_wspd_wdir10(wrfin, timeidx, method, result = get_destag_wspd_wdir10(wrfin, timeidx, method,
squeeze, cache, meta, _key, units)[0,:] squeeze, cache, meta, _key, units)[0, :]
if meta: if meta:
result.attrs["description"] = "10m wspd in projection space" result.attrs["description"] = "10m wspd in projection space"
@ -778,10 +776,9 @@ def get_destag_wdir10(wrfin, timeidx=0, method="cat",
""" """
result = get_destag_wspd_wdir10(wrfin, timeidx, method, result = get_destag_wspd_wdir10(wrfin, timeidx, method,
squeeze, cache, meta, _key, units)[1,:] squeeze, cache, meta, _key, units)[1, :]
if meta: if meta:
result.attrs["description"] = "10m wdir in projection space" result.attrs["description"] = "10m wdir in projection space"
return result return result

Loading…
Cancel
Save