Browse Source

Added tests for the ll_to_xy and xy_to_ll routines. Fixed bugs found during testing. Fixed numerous typos and removed obsolete code.

main
Bill Ladwig 9 years ago
parent
commit
00f7170e58
  1. 4
      src/wrf/api.py
  2. 3
      src/wrf/cache.py
  3. 2
      src/wrf/cape.py
  4. 10
      src/wrf/cloudfrac.py
  5. 4
      src/wrf/computation.py
  6. 6
      src/wrf/decorators.py
  7. 14
      src/wrf/interp.py
  8. 1
      src/wrf/interputils.py
  9. 24
      src/wrf/latlon.py
  10. 37
      src/wrf/latlonutils.py
  11. 16
      src/wrf/metadecorators.py
  12. 1
      src/wrf/omega.py
  13. 2
      src/wrf/projection.py
  14. 1
      src/wrf/py3compat.py
  15. 1
      src/wrf/rh.py
  16. 109
      src/wrf/specialdec.py
  17. 7
      src/wrf/temp.py
  18. 6
      src/wrf/units.py
  19. 30
      src/wrf/util.py
  20. 46
      test/utests.py

4
src/wrf/api.py

@ -10,7 +10,7 @@ from .computation import (xy, interp1d, interp2dxy, interpz3d, slp, tk, td, rh,
dbz, srhel, udhel, avo, pvo, eth, wetbulb, tvirtual, dbz, srhel, udhel, avo, pvo, eth, wetbulb, tvirtual,
omega, pw) omega, pw)
from .interp import (interplevel, vertcross, interpline, vinterp) from .interp import (interplevel, vertcross, interpline, vinterp)
from .latlon import (xy_to_ll, ll_to_xy) from .latlon import (xy_to_ll, ll_to_xy, xy_to_ll_proj, ll_to_xy_proj)
from .py3compat import (viewitems, viewkeys, viewvalues, isstr, py2round, from .py3compat import (viewitems, viewkeys, viewvalues, isstr, py2round,
py3range, ucode) py3range, ucode)
from .util import (npvalues, extract_global_attrs, from .util import (npvalues, extract_global_attrs,
@ -32,7 +32,7 @@ __all__ += ["xy", "interp1d", "interp2dxy", "interpz3d", "slp", "tk", "td",
"ctt", "dbz", "srhel", "udhel", "avo", "pvo", "eth", "wetbulb", "ctt", "dbz", "srhel", "udhel", "avo", "pvo", "eth", "wetbulb",
"tvirtual", "omega", "pw"] "tvirtual", "omega", "pw"]
__all__ += ["interplevel", "vertcross", "interpline", "vinterp"] __all__ += ["interplevel", "vertcross", "interpline", "vinterp"]
__all__ += ["xy_to_ll", "ll_to_xy"] __all__ += ["xy_to_ll", "ll_to_xy", "xy_to_ll_proj", "ll_to_xy_proj"]
__all__ += ["viewitems", "viewkeys", "viewvalues", "isstr", "py2round", __all__ += ["viewitems", "viewkeys", "viewvalues", "isstr", "py2round",
"py3range", "ucode"] "py3range", "ucode"]
__all__ += ["npvalues", "extract_global_attrs", __all__ += ["npvalues", "extract_global_attrs",

3
src/wrf/cache.py

@ -21,13 +21,12 @@ def cache_item(key, product, value):
cache = _local_storage.cache cache = _local_storage.cache
try: try:
prod_dict = 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()
prod_dict = cache[key]
cache[key][product] = value cache[key][product] = value

2
src/wrf/cape.py

@ -8,7 +8,7 @@ import numpy.ma as ma
from .extension import _tk, _cape from .extension import _tk, _cape
from .destag import destagger from .destag import destagger
from .constants import Constants, ConversionFactors from .constants import Constants, ConversionFactors
from .util import extract_vars, combine_with from .util import extract_vars
from .metadecorators import set_cape_metadata from .metadecorators import set_cape_metadata

10
src/wrf/cloudfrac.py

@ -10,14 +10,14 @@ from .util import extract_vars
def get_cloudfrac(wrfnc, timeidx=0, method="cat", squeeze=True, def get_cloudfrac(wrfnc, timeidx=0, method="cat", squeeze=True,
cache=None, meta=True, _key=None): cache=None, meta=True, _key=None):
vars = extract_vars(wrfnc, timeidx, ("P", "PB", "QVAPOR", "T"), ncvars = extract_vars(wrfnc, timeidx, ("P", "PB", "QVAPOR", "T"),
method, squeeze, cache, meta=False, method, squeeze, cache, meta=False,
_key=_key) _key=_key)
p = vars["P"] p = ncvars["P"]
pb = vars["PB"] pb = ncvars["PB"]
qv = vars["QVAPOR"] qv = ncvars["QVAPOR"]
t = vars["T"] t = ncvars["T"]
full_p = p + pb full_p = p + pb
full_t = t + Constants.T_BASE full_t = t + Constants.T_BASE

4
src/wrf/computation.py

@ -75,7 +75,7 @@ def smooth2d(field, passes, meta=True):
@set_cape_alg_metadata(is2d=True, copyarg="pres_hpa") @set_cape_alg_metadata(is2d=True, copyarg="pres_hpa")
def cape_2d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow, def cape_2d(pres_hpa, tkel, qvapor, height, terrain, psfc_hpa, ter_follow,
missing=Constants.DEFAULT_FILL, meta=True): missing=Constants.DEFAULT_FILL, meta=True):
if isinstance(ter_follow, bool): if isinstance(ter_follow, bool):
@ -128,7 +128,7 @@ def ctt(pres_hpa, tkel, qv, qcld, height, terrain, qice=None, meta=True):
# Qice and QCLD need to be in g/kg # Qice and QCLD need to be in g/kg
if qice is None: if qice is None:
qice = n.zeros(qv.shape, qv.dtype) qice = np.zeros(qv.shape, qv.dtype)
haveqci = 0 haveqci = 0
else: else:
haveqci = 1 if qice.any() else 0 haveqci = 1 if qice.any() else 0

6
src/wrf/decorators.py

@ -9,7 +9,7 @@ import numpy.ma as ma
from .units import do_conversion, check_units from .units import do_conversion, check_units
from .util import iter_left_indexes, from_args, npvalues, combine_dims from .util import iter_left_indexes, from_args, npvalues, combine_dims
from .py3compat import viewitems, viewvalues, py3range, isstr from .py3compat import viewitems, viewvalues, isstr
from .config import xarray_enabled from .config import xarray_enabled
from .constants import Constants from .constants import Constants
@ -192,9 +192,9 @@ def left_iteration(ref_var_expected_dims,
# Mostly when used with join # Mostly when used with join
if mask_output: if mask_output:
if isinstance(output, np.ndarray): if isinstance(output, np.ndarray):
output = np.ma.masked_values(output, Constants.DEFAULT_FILL) output = ma.masked_values(output, Constants.DEFAULT_FILL)
else: else:
output = tuple(masked_values(arr, Constants.DEFAULT_FILL) output = tuple(ma.masked_values(arr, Constants.DEFAULT_FILL)
for arr in output) for arr in output)
return output return output

14
src/wrf/interp.py

@ -4,12 +4,8 @@ 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 (interpz3d, interp2dxy, interp1d, from .extension import (_interpz3d, _vertcross, _interpline, _smooth2d,
# smooth2d, monotonic, vintrp, computevertcross, _monotonic, _vintrp)
# computeinterpline)
from .extension import (_interpz3d, _interp2dxy, _interp1d, _vertcross,
_interpline, _smooth2d, _monotonic, _vintrp)
from .metadecorators import set_interp_metadata from .metadecorators import set_interp_metadata
from .util import extract_vars, is_staggered, get_id from .util import extract_vars, is_staggered, get_id
@ -312,9 +308,9 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False,
"eth" : 6} "eth" : 6}
# These constants match what's in the fortran code. # These constants match what's in the fortran code.
rgas = Constants.RD#287.04 #J/K/kg rgas = Constants.RD
ussalr = Constants.USSALR#.0065 # deg C per m, avg lapse rate ussalr = Constants.USSALR
sclht = Constants.SCLHT #rgas*256./9.81 sclht = Constants.SCLHT
# interp_levels might be a list or tuple, make a numpy array # interp_levels might be a list or tuple, make a numpy array
if not isinstance(interp_levels, np.ndarray): if not isinstance(interp_levels, np.ndarray):

1
src/wrf/interputils.py

@ -15,6 +15,7 @@ def to_positive_idxs(shape, 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, def calc_xy(xdim, ydim, pivot_point=None, angle=None,
start_point=None, end_point=None): start_point=None, end_point=None):
"""Returns the x,y points for the horizontal cross section line. """Returns the x,y points for the horizontal cross section line.

24
src/wrf/latlon.py

@ -41,7 +41,18 @@ def ll_to_xy(wrfnc, latitude, longitude, timeidx=0, stagger=None, method="cat",
@set_latlon_metadata(xy=True) @set_latlon_metadata(xy=True)
def ll_to_xy_proj(latitude, longitude, meta=True, squeeze=True, as_int=True, def ll_to_xy_proj(latitude, longitude, meta=True, squeeze=True, as_int=True,
**projparams): map_proj=None, truelat1=None, truelat2=None, stand_lon=None,
ref_lat=None, ref_lon=None, pole_lat=None, pole_lon=None,
known_x=None, known_y=None, dx=None, dy=None,
latinc=None, loninc=None):
loc = locals()
projparams = {name : loc[name] for name in ("map_proj", "truelat1",
"truelat2", "stand_lon", "ref_lat",
"ref_lon", "pole_lat", "pole_lon",
"known_x", "known_y", "dx", "dy",
"latinc", "loninc")}
return _ll_to_xy(latitude, longitude, None, 0, squeeze, "cat", True, None, return _ll_to_xy(latitude, longitude, None, 0, squeeze, "cat", True, None,
None, as_int, **projparams) None, as_int, **projparams)
@ -55,7 +66,16 @@ def xy_to_ll(wrfnc, x, y, timeidx=0, stagger=None, method="cat", squeeze=True,
@set_latlon_metadata(xy=False) @set_latlon_metadata(xy=False)
def xy_to_ll_proj(x, y, meta=True, squeeze=True, **projparams): def xy_to_ll_proj(x, y, meta=True, squeeze=True, map_proj=None, truelat1=None,
truelat2=None, stand_lon=None, ref_lat=None, ref_lon=None,
pole_lat=None, pole_lon=None, known_x=None, known_y=None,
dx=None, dy=None, latinc=None, loninc=None):
loc = locals()
projparams = {name : loc[name] for name in ("map_proj", "truelat1",
"truelat2", "stand_lon", "ref_lat",
"ref_lon", "pole_lat", "pole_lon",
"known_x", "known_y", "dx", "dy",
"latinc", "loninc")}
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)

37
src/wrf/latlonutils.py

@ -101,28 +101,29 @@ def _dict_keys_to_upper(d):
return {key.upper(): val for key, val in viewitems(d)} return {key.upper(): val for key, val in viewitems(d)}
# 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):
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")
stdlon = projparams.get("STDLON"), stdlon = projparams.get("STAND_LON")
ref_lat = projparams.get("REF_LAT"), ref_lat = projparams.get("REF_LAT")
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"),
dy = projparams.get("DY"), dx = projparams.get("DX")
latinc = projparams.get("LATINC"), dy = projparams.get("DY")
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,
"STDLON" : stdlon, "STAND_LON" : stdlon,
"REF_LAT" : ref_lat, "REF_LAT" : ref_lat,
"REF_LON" : ref_lon, "REF_LON" : ref_lon,
"KNOWN_X" : known_x, "KNOWN_X" : known_x,
@ -131,6 +132,10 @@ def _kwarg_proj_params(projparams):
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 expectd to be lists
ref_lat = np.asarray([ref_lat])
ref_lon = 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
@ -166,7 +171,7 @@ def _kwarg_proj_params(projparams):
# Will return 0-based indexes # Will return 0-based indexes
def _ll_to_xy(latitude, longitude, wrfnc=None, timeidx=0, def _ll_to_xy(latitude, longitude, wrfnc=None, timeidx=0,
stagger=None, method="cat", squeeze=True, cache=None, stagger=None, method="cat", squeeze=True, cache=None,
_key=None, as_int=True, **projparms): _key=None, as_int=True, **projparams):
if wrfnc is not None: if wrfnc is not None:
(map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon,

16
src/wrf/metadecorators.py

@ -1127,6 +1127,9 @@ def set_alg_metadata(alg_ndims, refvarname,
result = wrapped(*args, **kwargs) result = wrapped(*args, **kwargs)
outname = wrapped.__name__
outattrs = OrderedDict()
# Default dimension names # Default dimension names
outdims = ["dim_{}".format(i) for i in py3range(result.ndim)] outdims = ["dim_{}".format(i) for i in py3range(result.ndim)]
@ -1141,13 +1144,10 @@ def set_alg_metadata(alg_ndims, refvarname,
outattrs["missing_value"] = missingval outattrs["missing_value"] = missingval
result = np.ma.masked_values(result, missingval) result = np.ma.masked_values(result, missingval)
outname = wrapped.__name__
outattrs = OrderedDict()
if units is not None: if units is not None:
if isinstance(description, from_var): if isinstance(units, from_var):
_units = units(wrapped, *args, **kwargs) _units = units(wrapped, *args, **kwargs)
if uts is not None: if _units is not None:
outattrs["units"] = _units outattrs["units"] = _units
else: else:
outattrs["units"] = units outattrs["units"] = units
@ -1214,6 +1214,7 @@ def set_alg_metadata(alg_ndims, refvarname,
return func_wrapper return func_wrapper
def set_uvmet_alg_metadata(units="mps", description="earth rotated u,v", def set_uvmet_alg_metadata(units="mps", description="earth rotated u,v",
latarg="lat", windarg="u"): latarg="lat", windarg="u"):
@ -1263,6 +1264,7 @@ def set_uvmet_alg_metadata(units="mps", description="earth rotated u,v",
return func_wrapper return func_wrapper
def set_cape_alg_metadata(is2d, copyarg="pres_hpa"): def set_cape_alg_metadata(is2d, copyarg="pres_hpa"):
@wrapt.decorator @wrapt.decorator
@ -1293,7 +1295,6 @@ def set_cape_alg_metadata(is2d, copyarg="pres_hpa"):
outattrs["units"] = "J kg-1 ; J kg-1" outattrs["units"] = "J kg-1 ; J kg-1"
outattrs["MemoryOrder"] = "XYZ" outattrs["MemoryOrder"] = "XYZ"
argvals = from_args(wrapped, (copyarg,"missing"), *args, **kwargs) argvals = from_args(wrapped, (copyarg,"missing"), *args, **kwargs)
p = argvals[copyarg] p = argvals[copyarg]
missing = argvals["missing"] missing = argvals["missing"]
@ -1377,6 +1378,7 @@ def set_cloudfrac_alg_metadata(copyarg="pres"):
return func_wrapper return func_wrapper
def set_destag_metadata(): def set_destag_metadata():
@wrapt.decorator @wrapt.decorator
@ -1404,7 +1406,7 @@ def set_destag_metadata():
if var.name is not None: if var.name is not None:
outname = "destag_{}".format(var.name) outname = "destag_{}".format(var.name)
else: else:
outnames = "destag_var" outname = "destag_var"
outattrs = OrderedDict() outattrs = OrderedDict()
outattrs.update(var.attrs) outattrs.update(var.attrs)

1
src/wrf/omega.py

@ -3,7 +3,6 @@ from __future__ import (absolute_import, division, print_function,
from .constants import Constants from .constants import Constants
from .destag import destagger from .destag import destagger
#from .extension import computeomega,computetk
from .extension import _omega, _tk from .extension import _omega, _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

2
src/wrf/projection.py

@ -257,7 +257,7 @@ class LambertConformal(WrfProj):
def _cart_extents(self): def _cart_extents(self):
# Need to modify the extents for the new projection # Need to modify the extents for the new projection
pc = crs.PlateCarree() pc = crs.PlateCarree()
xs, ys, zs = self._cartopy().transform_points(pc, xs, ys, _ = self._cartopy().transform_points(pc,
np.array([self.ll_lon, self.ur_lon]), np.array([self.ll_lon, self.ur_lon]),
np.array([self.ll_lat, self.ur_lat])).T np.array([self.ll_lat, self.ur_lat])).T

1
src/wrf/py3compat.py

@ -1,4 +1,5 @@
from sys import version_info from sys import version_info
from math import floor, copysign
# Dictionary python 2-3 compatibility stuff # Dictionary python 2-3 compatibility stuff
def viewitems(d): def viewitems(d):

1
src/wrf/rh.py

@ -29,6 +29,7 @@ def get_rh(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
return rh return rh
@copy_and_set_metadata(copy_varname="T2", name="rh2", @copy_and_set_metadata(copy_varname="T2", name="rh2",
description="2m relative humidity", description="2m relative humidity",
delete_attrs=("units",)) delete_attrs=("units",))

109
src/wrf/specialdec.py

@ -5,9 +5,7 @@ import numpy as np
import wrapt import wrapt
#from .destag import destagger
from .util import iter_left_indexes, npvalues from .util import iter_left_indexes, npvalues
from .py3compat import py3range
from .config import xarray_enabled from .config import xarray_enabled
from .constants import Constants from .constants import Constants
@ -15,113 +13,6 @@ if xarray_enabled():
from xarray import DataArray from xarray import DataArray
# def uvmet_left_iter():
# """Decorator to handle iterating over leftmost dimensions when using
# multiple files and/or multiple times with the uvmet product.
#
# """
# @wrapt.decorator
# def func_wrapper(wrapped, instance, args, kwargs):
# u = args[0]
# v = args[1]
# lat = args[2]
# lon = args[3]
# cen_long = args[4]
# cone = args[5]
#
# if u.ndim == lat.ndim:
# num_right_dims = 2
# is_3d = False
# else:
# num_right_dims = 3
# is_3d = True
#
# is_stag = False
# if ((u.shape[-1] != lat.shape[-1]) or
# (u.shape[-2] != lat.shape[-2])):
# is_stag = True
#
# if is_3d:
# extra_dim_num = u.ndim - 3
# else:
# extra_dim_num = u.ndim - 2
#
# if is_stag:
# u = destagger(u,-1)
# v = destagger(v,-2)
#
# # No special left side iteration, return the function result
# if (extra_dim_num == 0):
# return wrapped(u, v, lat, lon, cen_long, cone)
#
# # Start by getting the left-most 'extra' dims
# outdims = u.shape[0:extra_dim_num]
# extra_dims = list(outdims) # Copy the left-most dims for iteration
#
# # Append the right-most dimensions
# outdims += [2] # For u/v components
#
# #outdims += [u.shape[x] for x in py3range(-num_right_dims,0,1)]
# outdims += list(u.shape[-num_right_dims:])
#
# output = np.empty(outdims, u.dtype)
#
# for left_idxs in iter_left_indexes(extra_dims):
# # Make the left indexes plus a single slice object
# # The single slice will handle all the dimensions to
# # the right (e.g. [1,1,:])
# left_and_slice_idxs = tuple([x for x in left_idxs] + [slice(None)])
#
# new_u = u[left_and_slice_idxs]
# new_v = v[left_and_slice_idxs]
# new_lat = lat[left_and_slice_idxs]
# new_lon = lon[left_and_slice_idxs]
#
# # Call the numerical routine
# result = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone)
#
# # Note: The 2D version will return a 3D array with a 1 length
# # dimension. Numpy is unable to broadcast this without
# # sqeezing first.
# result = np.squeeze(result)
#
# output[left_and_slice_idxs] = result[:]
#
# return output
#
# return func_wrapper
def _move_dim_to_left(arr, dimidx):
if isinstance(arr, np.ma.MaskedArray):
has_missing = True
missing = result.fill_value
shape = list(arr.shape)
move_dim_size = shape.pop(dimidx)
output_dims = [move_dim_size] + shape
output = np.empty(output_dims, arr.dtype)
if dimidx < 0:
right_ndim = -dimidx
else:
right_ndim = arr.ndim - dimidx
rightdims = [slice(None)] * right_ndim
for i in py3range(move_dim_size):
rightdims[0] = i
outidxs = (Ellipsis,) + tuple(rightdims)
output[i,:] = outview_array[outidxs]
if has_missing:
output = np.ma.masked_values(output, missing)
return output
def uvmet_left_iter_nocopy(alg_dtype=np.float64): def uvmet_left_iter_nocopy(alg_dtype=np.float64):
"""Decorator to handle iterating over leftmost dimensions when using """Decorator to handle iterating over leftmost dimensions when using
multiple files and/or multiple times with the uvmet product. multiple files and/or multiple times with the uvmet product.

7
src/wrf/temp.py

@ -2,7 +2,6 @@ from __future__ import (absolute_import, division, print_function,
unicode_literals) unicode_literals)
from .constants import Constants from .constants import Constants
#from .extension import computetk, computeeth, computetv, computewetbulb
from .extension import _tk, _eth, _tv, _wetbulb from .extension import _tk, _eth, _tv, _wetbulb
from .decorators import convert_units from .decorators import convert_units
from .metadecorators import copy_and_set_metadata from .metadecorators import copy_and_set_metadata
@ -24,6 +23,7 @@ def get_theta(wrfnc, timeidx=0, method="cat", squeeze=True,
return full_t return full_t
@copy_and_set_metadata(copy_varname="T", name="temp", @copy_and_set_metadata(copy_varname="T", name="temp",
description="temperature") description="temperature")
@convert_units("temp", "k") @convert_units("temp", "k")
@ -45,6 +45,7 @@ def get_temp(wrfnc, timeidx=0, method="cat", squeeze=True,
return tk return tk
@copy_and_set_metadata(copy_varname="T", name="theta_e", @copy_and_set_metadata(copy_varname="T", name="theta_e",
description="equivalent potential temperature") description="equivalent potential temperature")
@convert_units("temp", "k") @convert_units("temp", "k")
@ -69,6 +70,7 @@ def get_eth(wrfnc, timeidx=0, method="cat", squeeze=True,
return eth return eth
@copy_and_set_metadata(copy_varname="T", name="tv", @copy_and_set_metadata(copy_varname="T", name="tv",
description="virtual temperature") description="virtual temperature")
@convert_units("temp", "k") @convert_units("temp", "k")
@ -94,6 +96,7 @@ def get_tv(wrfnc, timeidx=0, method="cat", squeeze=True,
return tv return tv
@copy_and_set_metadata(copy_varname="T", name="twb", @copy_and_set_metadata(copy_varname="T", name="twb",
description="wetbulb temperature") description="wetbulb temperature")
@convert_units("temp", "k") @convert_units("temp", "k")
@ -118,11 +121,13 @@ def get_tw(wrfnc, timeidx=0, method="cat", squeeze=True,
return tw return tw
def get_tk(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, def get_tk(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True, _key=None): meta=True, _key=None):
return get_temp(wrfnc, timeidx, method, squeeze, cache, meta, _key, return get_temp(wrfnc, timeidx, method, squeeze, cache, meta, _key,
units="k") units="k")
def get_tc(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, def get_tc(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None,
meta=True, _key=None): meta=True, _key=None):
return get_temp(wrfnc, timeidx, method, squeeze, cache, meta, _key, return get_temp(wrfnc, timeidx, method, squeeze, cache, meta, _key,

6
src/wrf/units.py

@ -20,18 +20,22 @@ def _apply_conv_fact(var, vartype, var_unit, dest_unit):
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_celsius(var, var_unit): def _to_celsius(var, var_unit):
if var_unit == "k": if var_unit == "k":
return var - Constants.CELKEL return var - Constants.CELKEL
elif var_unit == "f": elif var_unit == "f":
return (var - 32.0) * (5.0/9.0) return (var - 32.0) * (5.0/9.0)
def _c_to_k(var): def _c_to_k(var):
return var + Constants.TCK0 return var + Constants.TCK0
def _c_to_f(var): def _c_to_f(var):
return 1.8 * var + 32.0 return 1.8 * var + 32.0
# Temperature is a more complicated operation so requres functions # Temperature is a more complicated operation so requres functions
def _apply_temp_conv(var, var_unit, dest_unit): def _apply_temp_conv(var, var_unit, dest_unit):
if dest_unit == var_unit: if dest_unit == var_unit:
@ -109,11 +113,13 @@ _TEMP_CONV_METHODS = {"k" : _c_to_k,
"f" : _c_to_f "f" : _c_to_f
} }
def check_units(unit, unit_type): def check_units(unit, unit_type):
unitl = unit.lower() unitl = unit.lower()
if unitl not in _VALID_UNITS[unit_type]: if unitl not in _VALID_UNITS[unit_type]:
raise ValueError("invalid unit type '%s'" % unit) raise ValueError("invalid unit type '%s'" % unit)
def do_conversion(var, vartype, var_unit, dest_unit): def do_conversion(var, vartype, var_unit, dest_unit):
if vartype != "temp": if vartype != "temp":
return _apply_conv_fact(var, vartype, return _apply_conv_fact(var, vartype,

30
src/wrf/util.py

@ -8,7 +8,6 @@ from collections import Iterable, Mapping, OrderedDict
from itertools import product from itertools import product
from types import GeneratorType from types import GeneratorType
import datetime as dt import datetime as dt
from math import floor, copysign
from inspect import getmodule from inspect import getmodule
try: try:
@ -27,8 +26,7 @@ import numpy.ma as ma
from .config import xarray_enabled from .config import xarray_enabled
from .projection import getproj, NullProjection from .projection import getproj, NullProjection
from .constants import Constants, ALL_TIMES from .constants import Constants, ALL_TIMES
from .py3compat import (viewitems, viewkeys, viewvalues, isstr, py2round, from .py3compat import (viewitems, viewkeys, isstr, py3range, ucode)
py3range, ucode)
from .cache import cache_item, get_cached_item from .cache import cache_item, get_cached_item
if xarray_enabled(): if xarray_enabled():
@ -73,12 +71,15 @@ def is_multi_time_req(timeidx):
def is_multi_file(wrfnc): def is_multi_file(wrfnc):
return (isinstance(wrfnc, Iterable) and not isstr(wrfnc)) return (isinstance(wrfnc, Iterable) and not isstr(wrfnc))
def has_time_coord(wrfnc): def has_time_coord(wrfnc):
return "XTIME" in wrfnc.variables return "XTIME" in wrfnc.variables
def is_mapping(wrfnc): def is_mapping(wrfnc):
return isinstance(wrfnc, Mapping) return isinstance(wrfnc, Mapping)
def _generator_copy(gen): def _generator_copy(gen):
funcname = gen.__name__ funcname = gen.__name__
try: try:
@ -96,11 +97,13 @@ def _generator_copy(gen):
return res return res
def test(): def test():
q = [1,2,3] q = [1,2,3]
for i in q: for i in q:
yield i yield i
class TestGen(object): class TestGen(object):
def __init__(self, count=3): def __init__(self, count=3):
self._total = count self._total = count
@ -121,6 +124,7 @@ class TestGen(object):
def __next__(self): def __next__(self):
return self.next() return self.next()
def latlon_coordvars(d): def latlon_coordvars(d):
lat_coord = None lat_coord = None
lon_coord = None lon_coord = None
@ -137,9 +141,11 @@ def latlon_coordvars(d):
return lat_coord, lon_coord return lat_coord, lon_coord
def is_coordvar(varname): def is_coordvar(varname):
return varname in _COORD_VARS return varname in _COORD_VARS
class IterWrapper(Iterable): class IterWrapper(Iterable):
"""A wrapper class for generators and custom iterable classes which returns """A wrapper class for generators and custom iterable classes which returns
a new iterator from the start of the sequence when __iter__ is called""" a new iterator from the start of the sequence when __iter__ is called"""
@ -172,6 +178,7 @@ def get_iterable(wrfseq):
else: else:
return dict(wrfseq) # generator/custom iterable class return dict(wrfseq) # generator/custom iterable class
# Helper to extract masked arrays from DataArrays that convert to NaN # Helper to extract masked arrays from DataArrays that convert to NaN
def npvalues(array_type): def npvalues(array_type):
if not isinstance(array_type, DataArray): if not isinstance(array_type, DataArray):
@ -187,6 +194,7 @@ def npvalues(array_type):
return result return result
# Helper utilities for metadata # Helper utilities for metadata
class either(object): class either(object):
def __init__(self, *varnames): def __init__(self, *varnames):
@ -207,6 +215,7 @@ class either(object):
raise ValueError("{} are not valid variable names".format( raise ValueError("{} are not valid variable names".format(
self.varnames)) self.varnames))
class combine_with(object): class combine_with(object):
# Remove remove_idx first, then insert_idx is applied to removed set # Remove remove_idx first, then insert_idx is applied to removed set
def __init__(self, varname, remove_dims=None, insert_before=None, def __init__(self, varname, remove_dims=None, insert_before=None,
@ -239,6 +248,7 @@ class combine_with(object):
return new_dims, new_coords return new_dims, new_coords
# This should look like: # This should look like:
# [(0, (-3,-2)), # variable 1 # [(0, (-3,-2)), # variable 1
# (1, -1)] # variable 2 # (1, -1)] # variable 2
@ -387,6 +397,7 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"),
return False return False
def _get_global_attr(wrfnc, attr): def _get_global_attr(wrfnc, attr):
val = getattr(wrfnc, attr, None) val = getattr(wrfnc, attr, None)
@ -396,6 +407,7 @@ def _get_global_attr(wrfnc, attr):
return val[0] return val[0]
return val return val
def extract_global_attrs(wrfnc, attrs): def extract_global_attrs(wrfnc, attrs):
if isstr(attrs): if isstr(attrs):
attrlist = [attrs] attrlist = [attrs]
@ -413,6 +425,7 @@ def extract_global_attrs(wrfnc, attrs):
return {attr:_get_global_attr(wrfnc, attr) for attr in attrlist} return {attr:_get_global_attr(wrfnc, attr) for attr in attrlist}
def extract_dim(wrfnc, dim): def extract_dim(wrfnc, dim):
if is_multi_file(wrfnc): if is_multi_file(wrfnc):
if not is_mapping(wrfnc): if not is_mapping(wrfnc):
@ -426,6 +439,7 @@ def extract_dim(wrfnc, dim):
return len(d) #netCDF4 return len(d) #netCDF4
return d # PyNIO return d # PyNIO
def _combine_dict(wrfdict, varname, timeidx, method, meta, _key): def _combine_dict(wrfdict, varname, timeidx, method, meta, _key):
"""Dictionary combination creates a new left index for each key, then """Dictionary combination creates a new left index for each key, then
does a cat or join for the list of files for that key""" does a cat or join for the list of files for that key"""
@ -798,6 +812,7 @@ def _find_arr_for_time(wrfseq, varname, timeidx, is_moving, meta, _key):
else: else:
return _find_reverse(wrfseq, varname, timeidx, is_moving, meta, _key) return _find_reverse(wrfseq, varname, timeidx, is_moving, meta, _key)
# TODO: implement in C # TODO: implement in C
def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta, _key): def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta, _key):
if is_moving is None: if is_moving is None:
@ -1234,6 +1249,7 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key):
return outarr return outarr
def combine_files(wrfseq, varname, timeidx, is_moving=None, def combine_files(wrfseq, varname, timeidx, is_moving=None,
method="cat", squeeze=True, meta=True, method="cat", squeeze=True, meta=True,
_key=None): _key=None):
@ -1416,7 +1432,6 @@ def is_staggered(var, wrfnc):
return False return False
def get_left_indexes(ref_var, expected_dims): def get_left_indexes(ref_var, expected_dims):
"""Returns the extra left side dimensions for a variable with an expected """Returns the extra left side dimensions for a variable with an expected
shape. shape.
@ -1432,6 +1447,7 @@ def get_left_indexes(ref_var, expected_dims):
return tuple([ref_var.shape[x] for x in py3range(extra_dim_num)]) return tuple([ref_var.shape[x] for x in py3range(extra_dim_num)])
def iter_left_indexes(dims): def iter_left_indexes(dims):
"""A generator which yields the iteration tuples for a sequence of """A generator which yields the iteration tuples for a sequence of
dimensions sizes. dimensions sizes.
@ -1449,6 +1465,7 @@ def iter_left_indexes(dims):
for idxs in product(*arg): for idxs in product(*arg):
yield idxs yield idxs
def get_right_slices(var, right_ndims, fixed_val=0): def get_right_slices(var, right_ndims, fixed_val=0):
extra_dim_num = var.ndim - right_ndims extra_dim_num = var.ndim - right_ndims
if extra_dim_num == 0: if extra_dim_num == 0:
@ -1457,6 +1474,7 @@ def get_right_slices(var, right_ndims, fixed_val=0):
return tuple([fixed_val]*extra_dim_num + return tuple([fixed_val]*extra_dim_num +
[slice(None)]*right_ndims) [slice(None)]*right_ndims)
def get_proj_params(wrfnc, timeidx=0, varname=None): def get_proj_params(wrfnc, timeidx=0, varname=None):
proj_params = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", proj_params = extract_global_attrs(wrfnc, attrs=("MAP_PROJ",
"CEN_LAT", "CEN_LON", "CEN_LAT", "CEN_LON",
@ -1556,6 +1574,7 @@ def from_args(func, argnames, *args, **kwargs):
return result return result
def _args_to_list2(func, args, kwargs): def _args_to_list2(func, args, kwargs):
argspec = getargspec(func) argspec = getargspec(func)
@ -1576,6 +1595,7 @@ def _args_to_list2(func, args, kwargs):
return outargs return outargs
def _args_to_list3(func, args, kwargs): def _args_to_list3(func, args, kwargs):
sig = signature(func) sig = signature(func)
bound = sig.bind(*args, **kwargs) bound = sig.bind(*args, **kwargs)
@ -1584,6 +1604,7 @@ def _args_to_list3(func, args, kwargs):
return [x for x in bound.arguments.values()] return [x for x in bound.arguments.values()]
# Note: Doesn't allow for **kwargs or *args
def args_to_list(func, args, kwargs): def args_to_list(func, args, kwargs):
"""Converts the mixed args/kwargs to a single list of args""" """Converts the mixed args/kwargs to a single list of args"""
if version_info > (3,): if version_info > (3,):
@ -1607,6 +1628,7 @@ def _arg_location2(func, argname, args, kwargs):
return list_args, result_idx return list_args, result_idx
def _arg_location3(func, argname, args, kwargs): def _arg_location3(func, argname, args, kwargs):
sig = signature(func) sig = signature(func)
params = list(sig.parameters.keys()) params = list(sig.parameters.keys())

46
test/utests.py

@ -7,7 +7,9 @@ import subprocess
from wrf import (getvar, interplevel, interpline, vertcross, vinterp, from wrf import (getvar, interplevel, interpline, vertcross, vinterp,
disable_xarray, xarray_enabled, npvalues, disable_xarray, xarray_enabled, npvalues,
xy_to_ll, ll_to_xy ) xy_to_ll, ll_to_xy, xy_to_ll_proj, ll_to_xy_proj,
extract_global_attrs, viewitems)
from wrf.util import is_multi_file
NCL_EXE = "/Users/ladwig/nclbuild/6.3.0/bin/ncl" NCL_EXE = "/Users/ladwig/nclbuild/6.3.0/bin/ncl"
TEST_FILE = "/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00" TEST_FILE = "/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00"
@ -458,6 +460,23 @@ def make_interp_test(varname, wrf_in, referent, multi=False,
return test return test
def extract_proj_params(wrfnc):
attrs = extract_global_attrs(wrfnc, ("MAP_PROJ", "TRUELAT1", "TRUELAT2",
"STAND_LON", "POLE_LAT", "POLE_LON",
"DX", "DY"))
result = {key.lower(): val for key,val in viewitems(attrs)}
if is_multi_file(wrfnc):
wrfnc = wrfnc[0]
result["known_x"] = 0
result["known_y"] = 0
result["ref_lat"] = wrfnc.variables["XLAT"][0,0,0]
result["ref_lon"] = wrfnc.variables["XLONG"][0,0,0]
return result
def make_latlon_test(testid, wrf_in, referent, single, multi=False, repeat=3, def make_latlon_test(testid, wrf_in, referent, single, multi=False, repeat=3,
pynio=False): pynio=False):
def test(self): def test(self):
@ -514,6 +533,13 @@ def make_latlon_test(testid, wrf_in, referent, single, multi=False, repeat=3,
nt.assert_allclose(npvalues(xy), ref) nt.assert_allclose(npvalues(xy), ref)
# Next make sure the 'proj' version works
projparams = extract_proj_params(in_wrfnc)
xy_proj = ll_to_xy_proj(lats[0], lons[0], **projparams)
nt.assert_allclose(npvalues(xy_proj), npvalues(xy-1))
else: else:
xy = ll_to_xy(in_wrfnc, lats, lons) xy = ll_to_xy(in_wrfnc, lats, lons)
xy = xy + 1 # NCL uses fortran indexing xy = xy + 1 # NCL uses fortran indexing
@ -521,6 +547,12 @@ def make_latlon_test(testid, wrf_in, referent, single, multi=False, repeat=3,
nt.assert_allclose(npvalues(xy), ref) nt.assert_allclose(npvalues(xy), ref)
# Next make sure the 'proj' version works
projparams = extract_proj_params(in_wrfnc)
xy_proj = ll_to_xy_proj(lats, lons, **projparams)
nt.assert_allclose(npvalues(xy_proj), npvalues(xy-1))
else: else:
# Since this domain is not moving, the reference values are the # Since this domain is not moving, the reference values are the
# same whether there are multiple or single files # same whether there are multiple or single files
@ -536,12 +568,24 @@ def make_latlon_test(testid, wrf_in, referent, single, multi=False, repeat=3,
ref = ref_vals[::-1,0] ref = ref_vals[::-1,0]
nt.assert_allclose(npvalues(ll), ref) nt.assert_allclose(npvalues(ll), ref)
# Next make sure the 'proj' version works
projparams = extract_proj_params(in_wrfnc)
ll_proj = xy_to_ll_proj(i_s[0], j_s[0], **projparams)
nt.assert_allclose(npvalues(ll_proj), npvalues(ll))
else: else:
ll = xy_to_ll(in_wrfnc, i_s, j_s) ll = xy_to_ll(in_wrfnc, i_s, j_s)
ref = ref_vals[::-1,:] ref = ref_vals[::-1,:]
nt.assert_allclose(npvalues(ll), ref) nt.assert_allclose(npvalues(ll), ref)
# Next make sure the 'proj' version works
projparams = extract_proj_params(in_wrfnc)
ll_proj = xy_to_ll_proj(i_s, j_s, **projparams)
nt.assert_allclose(npvalues(ll_proj), npvalues(ll))
return test return test
class WRFVarsTest(ut.TestCase): class WRFVarsTest(ut.TestCase):

Loading…
Cancel
Save