From 00f7170e588e216ccc439711ec82035fd2d5a67b Mon Sep 17 00:00:00 2001 From: Bill Ladwig Date: Tue, 13 Sep 2016 16:40:42 -0600 Subject: [PATCH] Added tests for the ll_to_xy and xy_to_ll routines. Fixed bugs found during testing. Fixed numerous typos and removed obsolete code. --- src/wrf/api.py | 4 +- src/wrf/cache.py | 3 +- src/wrf/cape.py | 2 +- src/wrf/cloudfrac.py | 10 ++-- src/wrf/computation.py | 4 +- src/wrf/decorators.py | 6 +-- src/wrf/interp.py | 14 ++--- src/wrf/interputils.py | 1 + src/wrf/latlon.py | 24 ++++++++- src/wrf/latlonutils.py | 37 +++++++------ src/wrf/metadecorators.py | 22 ++++---- src/wrf/omega.py | 1 - src/wrf/projection.py | 2 +- src/wrf/py3compat.py | 1 + src/wrf/rh.py | 1 + src/wrf/specialdec.py | 109 -------------------------------------- src/wrf/temp.py | 7 ++- src/wrf/units.py | 6 +++ src/wrf/util.py | 36 ++++++++++--- test/utests.py | 48 ++++++++++++++++- 20 files changed, 165 insertions(+), 173 deletions(-) diff --git a/src/wrf/api.py b/src/wrf/api.py index f0e22f2..ae8d6bd 100644 --- a/src/wrf/api.py +++ b/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, omega, pw) 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, py3range, ucode) 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", "tvirtual", "omega", "pw"] __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", "py3range", "ucode"] __all__ += ["npvalues", "extract_global_attrs", diff --git a/src/wrf/cache.py b/src/wrf/cache.py index 867adf8..98672a8 100644 --- a/src/wrf/cache.py +++ b/src/wrf/cache.py @@ -21,13 +21,12 @@ def cache_item(key, product, value): cache = _local_storage.cache try: - prod_dict = cache[key] + _ = cache[key] except KeyError: if len(cache) >= get_cache_size(): cache.popitem(last=False) # Remove the oldest dataset cache[key] = OrderedDict() - prod_dict = cache[key] cache[key][product] = value diff --git a/src/wrf/cape.py b/src/wrf/cape.py index 7525f55..e86872d 100755 --- a/src/wrf/cape.py +++ b/src/wrf/cape.py @@ -8,7 +8,7 @@ import numpy.ma as ma from .extension import _tk, _cape from .destag import destagger from .constants import Constants, ConversionFactors -from .util import extract_vars, combine_with +from .util import extract_vars from .metadecorators import set_cape_metadata diff --git a/src/wrf/cloudfrac.py b/src/wrf/cloudfrac.py index 476f4ea..62848c6 100644 --- a/src/wrf/cloudfrac.py +++ b/src/wrf/cloudfrac.py @@ -10,14 +10,14 @@ from .util import extract_vars def get_cloudfrac(wrfnc, timeidx=0, method="cat", squeeze=True, 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, _key=_key) - p = vars["P"] - pb = vars["PB"] - qv = vars["QVAPOR"] - t = vars["T"] + p = ncvars["P"] + pb = ncvars["PB"] + qv = ncvars["QVAPOR"] + t = ncvars["T"] full_p = p + pb full_t = t + Constants.T_BASE diff --git a/src/wrf/computation.py b/src/wrf/computation.py index a76a014..7cd1408 100644 --- a/src/wrf/computation.py +++ b/src/wrf/computation.py @@ -75,7 +75,7 @@ def smooth2d(field, passes, meta=True): @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): 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 if qice is None: - qice = n.zeros(qv.shape, qv.dtype) + qice = np.zeros(qv.shape, qv.dtype) haveqci = 0 else: haveqci = 1 if qice.any() else 0 diff --git a/src/wrf/decorators.py b/src/wrf/decorators.py index d9dc645..2ab3979 100644 --- a/src/wrf/decorators.py +++ b/src/wrf/decorators.py @@ -9,7 +9,7 @@ import numpy.ma as ma from .units import do_conversion, check_units 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 .constants import Constants @@ -192,9 +192,9 @@ def left_iteration(ref_var_expected_dims, # Mostly when used with join if mask_output: if isinstance(output, np.ndarray): - output = np.ma.masked_values(output, Constants.DEFAULT_FILL) + output = ma.masked_values(output, Constants.DEFAULT_FILL) else: - output = tuple(masked_values(arr, Constants.DEFAULT_FILL) + output = tuple(ma.masked_values(arr, Constants.DEFAULT_FILL) for arr in output) return output diff --git a/src/wrf/interp.py b/src/wrf/interp.py index 6d384a0..771d353 100755 --- a/src/wrf/interp.py +++ b/src/wrf/interp.py @@ -4,12 +4,8 @@ from __future__ import (absolute_import, division, print_function, import numpy as np import numpy.ma as ma -# from .extension import (interpz3d, interp2dxy, interp1d, -# smooth2d, monotonic, vintrp, computevertcross, -# computeinterpline) - -from .extension import (_interpz3d, _interp2dxy, _interp1d, _vertcross, - _interpline, _smooth2d, _monotonic, _vintrp) +from .extension import (_interpz3d, _vertcross, _interpline, _smooth2d, + _monotonic, _vintrp) from .metadecorators import set_interp_metadata 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} # These constants match what's in the fortran code. - rgas = Constants.RD#287.04 #J/K/kg - ussalr = Constants.USSALR#.0065 # deg C per m, avg lapse rate - sclht = Constants.SCLHT #rgas*256./9.81 + rgas = Constants.RD + ussalr = Constants.USSALR + sclht = Constants.SCLHT # interp_levels might be a list or tuple, make a numpy array if not isinstance(interp_levels, np.ndarray): diff --git a/src/wrf/interputils.py b/src/wrf/interputils.py index 1ed89cd..1af701f 100644 --- a/src/wrf/interputils.py +++ b/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) ] + def calc_xy(xdim, ydim, pivot_point=None, angle=None, start_point=None, end_point=None): """Returns the x,y points for the horizontal cross section line. diff --git a/src/wrf/latlon.py b/src/wrf/latlon.py index fb3c4cd..4522c45 100755 --- a/src/wrf/latlon.py +++ b/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) 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, 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) -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, **projparams) diff --git a/src/wrf/latlonutils.py b/src/wrf/latlonutils.py index 2d070a3..6a0aabe 100644 --- a/src/wrf/latlonutils.py +++ b/src/wrf/latlonutils.py @@ -101,28 +101,29 @@ def _dict_keys_to_upper(d): return {key.upper(): val for key, val in viewitems(d)} # known_x and known_y are 0-based -def _kwarg_proj_params(projparams): +def _kwarg_proj_params(**projparams): projparams = _dict_keys_to_upper(projparams) - map_proj = projparams.get("MAP_PROJ"), - truelat1 = projparams.get("TRUELAT1"), - truelat2 = projparams.get("TRUELAT2"), - stdlon = projparams.get("STDLON"), - ref_lat = projparams.get("REF_LAT"), - ref_lon = projparams.get("REF_LON"), - pole_lat = projparams.get("POLE_LAT", 90.0), - pole_lon = projparams.get("POLE_LON", 0.0), - known_x = projparams.get("KNOWN_X"), # Use 0-based - known_y = projparams.get("KNOWN_Y"), # Use 0-based - dx = projparams.get("DX"), - dy = projparams.get("DY"), - latinc = projparams.get("LATINC"), + map_proj = projparams.get("MAP_PROJ") + truelat1 = projparams.get("TRUELAT1") + truelat2 = projparams.get("TRUELAT2") + stdlon = projparams.get("STAND_LON") + ref_lat = projparams.get("REF_LAT") + ref_lon = projparams.get("REF_LON") + pole_lat = projparams.get("POLE_LAT", 90.0) + pole_lon = projparams.get("POLE_LON", 0.0) + known_x = projparams.get("KNOWN_X") # Use 0-based + known_y = projparams.get("KNOWN_Y") # Use 0-based + + dx = projparams.get("DX") + dy = projparams.get("DY") + latinc = projparams.get("LATINC") loninc = projparams.get("LONINC") # Sanity checks # Required args for all projections for name, var in viewitems({"MAP_PROJ" : map_proj, - "STDLON" : stdlon, + "STAND_LON" : stdlon, "REF_LAT" : ref_lat, "REF_LON" : ref_lon, "KNOWN_X" : known_x, @@ -131,6 +132,10 @@ def _kwarg_proj_params(projparams): if var is None: 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 known_x = known_x + 1 known_y = known_y + 1 @@ -166,7 +171,7 @@ def _kwarg_proj_params(projparams): # Will return 0-based indexes def _ll_to_xy(latitude, longitude, wrfnc=None, timeidx=0, 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: (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, diff --git a/src/wrf/metadecorators.py b/src/wrf/metadecorators.py index 9eed8a0..c4fbba7 100644 --- a/src/wrf/metadecorators.py +++ b/src/wrf/metadecorators.py @@ -1127,6 +1127,9 @@ def set_alg_metadata(alg_ndims, refvarname, result = wrapped(*args, **kwargs) + outname = wrapped.__name__ + outattrs = OrderedDict() + # Default dimension names outdims = ["dim_{}".format(i) for i in py3range(result.ndim)] @@ -1137,17 +1140,14 @@ def set_alg_metadata(alg_ndims, refvarname, missingval = None if missingval is not None: - outattrs["_FillValue"] = missingval - outattrs["missing_value"] = missingval - result = np.ma.masked_values(result, missingval) - - outname = wrapped.__name__ - outattrs = OrderedDict() + outattrs["_FillValue"] = missingval + outattrs["missing_value"] = missingval + result = np.ma.masked_values(result, missingval) if units is not None: - if isinstance(description, from_var): + if isinstance(units, from_var): _units = units(wrapped, *args, **kwargs) - if uts is not None: + if _units is not None: outattrs["units"] = _units else: outattrs["units"] = units @@ -1214,6 +1214,7 @@ def set_alg_metadata(alg_ndims, refvarname, return func_wrapper + def set_uvmet_alg_metadata(units="mps", description="earth rotated u,v", latarg="lat", windarg="u"): @@ -1263,6 +1264,7 @@ def set_uvmet_alg_metadata(units="mps", description="earth rotated u,v", return func_wrapper + def set_cape_alg_metadata(is2d, copyarg="pres_hpa"): @wrapt.decorator @@ -1292,7 +1294,6 @@ def set_cape_alg_metadata(is2d, copyarg="pres_hpa"): outattrs["description"] = "cape; cin" outattrs["units"] = "J kg-1 ; J kg-1" outattrs["MemoryOrder"] = "XYZ" - argvals = from_args(wrapped, (copyarg,"missing"), *args, **kwargs) p = argvals[copyarg] @@ -1377,6 +1378,7 @@ def set_cloudfrac_alg_metadata(copyarg="pres"): return func_wrapper + def set_destag_metadata(): @wrapt.decorator @@ -1404,7 +1406,7 @@ def set_destag_metadata(): if var.name is not None: outname = "destag_{}".format(var.name) else: - outnames = "destag_var" + outname = "destag_var" outattrs = OrderedDict() outattrs.update(var.attrs) diff --git a/src/wrf/omega.py b/src/wrf/omega.py index 2d2a0ea..61dccb7 100755 --- a/src/wrf/omega.py +++ b/src/wrf/omega.py @@ -3,7 +3,6 @@ from __future__ import (absolute_import, division, print_function, from .constants import Constants from .destag import destagger -#from .extension import computeomega,computetk from .extension import _omega, _tk from .util import extract_vars from .metadecorators import copy_and_set_metadata diff --git a/src/wrf/projection.py b/src/wrf/projection.py index 726d362..b2a72c2 100644 --- a/src/wrf/projection.py +++ b/src/wrf/projection.py @@ -257,7 +257,7 @@ class LambertConformal(WrfProj): def _cart_extents(self): # Need to modify the extents for the new projection 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_lat, self.ur_lat])).T diff --git a/src/wrf/py3compat.py b/src/wrf/py3compat.py index f14af6a..737ade3 100644 --- a/src/wrf/py3compat.py +++ b/src/wrf/py3compat.py @@ -1,4 +1,5 @@ from sys import version_info +from math import floor, copysign # Dictionary python 2-3 compatibility stuff def viewitems(d): diff --git a/src/wrf/rh.py b/src/wrf/rh.py index 78428d6..e2d8c5f 100755 --- a/src/wrf/rh.py +++ b/src/wrf/rh.py @@ -29,6 +29,7 @@ def get_rh(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, return rh + @copy_and_set_metadata(copy_varname="T2", name="rh2", description="2m relative humidity", delete_attrs=("units",)) diff --git a/src/wrf/specialdec.py b/src/wrf/specialdec.py index 474d20d..ef911a0 100644 --- a/src/wrf/specialdec.py +++ b/src/wrf/specialdec.py @@ -5,9 +5,7 @@ import numpy as np import wrapt -#from .destag import destagger from .util import iter_left_indexes, npvalues -from .py3compat import py3range from .config import xarray_enabled from .constants import Constants @@ -15,113 +13,6 @@ if xarray_enabled(): 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): """Decorator to handle iterating over leftmost dimensions when using multiple files and/or multiple times with the uvmet product. diff --git a/src/wrf/temp.py b/src/wrf/temp.py index 3374a86..fa0776f 100755 --- a/src/wrf/temp.py +++ b/src/wrf/temp.py @@ -2,7 +2,6 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) from .constants import Constants -#from .extension import computetk, computeeth, computetv, computewetbulb from .extension import _tk, _eth, _tv, _wetbulb from .decorators import convert_units from .metadecorators import copy_and_set_metadata @@ -24,6 +23,7 @@ def get_theta(wrfnc, timeidx=0, method="cat", squeeze=True, return full_t + @copy_and_set_metadata(copy_varname="T", name="temp", description="temperature") @convert_units("temp", "k") @@ -45,6 +45,7 @@ def get_temp(wrfnc, timeidx=0, method="cat", squeeze=True, return tk + @copy_and_set_metadata(copy_varname="T", name="theta_e", description="equivalent potential temperature") @convert_units("temp", "k") @@ -69,6 +70,7 @@ def get_eth(wrfnc, timeidx=0, method="cat", squeeze=True, return eth + @copy_and_set_metadata(copy_varname="T", name="tv", description="virtual temperature") @convert_units("temp", "k") @@ -94,6 +96,7 @@ def get_tv(wrfnc, timeidx=0, method="cat", squeeze=True, return tv + @copy_and_set_metadata(copy_varname="T", name="twb", description="wetbulb temperature") @convert_units("temp", "k") @@ -118,11 +121,13 @@ def get_tw(wrfnc, timeidx=0, method="cat", squeeze=True, return tw + def get_tk(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): return get_temp(wrfnc, timeidx, method, squeeze, cache, meta, _key, units="k") + def get_tc(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): return get_temp(wrfnc, timeidx, method, squeeze, cache, meta, _key, diff --git a/src/wrf/units.py b/src/wrf/units.py index f21628a..153a4bd 100755 --- a/src/wrf/units.py +++ b/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] * _CONV_FACTORS[vartype]["to_dest"][dest_unit]) + def _to_celsius(var, var_unit): if var_unit == "k": return var - Constants.CELKEL elif var_unit == "f": return (var - 32.0) * (5.0/9.0) + def _c_to_k(var): return var + Constants.TCK0 + def _c_to_f(var): return 1.8 * var + 32.0 + # Temperature is a more complicated operation so requres functions def _apply_temp_conv(var, var_unit, dest_unit): if dest_unit == var_unit: @@ -109,11 +113,13 @@ _TEMP_CONV_METHODS = {"k" : _c_to_k, "f" : _c_to_f } + def check_units(unit, unit_type): unitl = unit.lower() if unitl not in _VALID_UNITS[unit_type]: raise ValueError("invalid unit type '%s'" % unit) + def do_conversion(var, vartype, var_unit, dest_unit): if vartype != "temp": return _apply_conv_fact(var, vartype, diff --git a/src/wrf/util.py b/src/wrf/util.py index a8943d3..3ab74e2 100644 --- a/src/wrf/util.py +++ b/src/wrf/util.py @@ -8,7 +8,6 @@ from collections import Iterable, Mapping, OrderedDict from itertools import product from types import GeneratorType import datetime as dt -from math import floor, copysign from inspect import getmodule try: @@ -27,8 +26,7 @@ import numpy.ma as ma from .config import xarray_enabled from .projection import getproj, NullProjection from .constants import Constants, ALL_TIMES -from .py3compat import (viewitems, viewkeys, viewvalues, isstr, py2round, - py3range, ucode) +from .py3compat import (viewitems, viewkeys, isstr, py3range, ucode) from .cache import cache_item, get_cached_item if xarray_enabled(): @@ -73,12 +71,15 @@ def is_multi_time_req(timeidx): def is_multi_file(wrfnc): return (isinstance(wrfnc, Iterable) and not isstr(wrfnc)) + def has_time_coord(wrfnc): return "XTIME" in wrfnc.variables + def is_mapping(wrfnc): return isinstance(wrfnc, Mapping) + def _generator_copy(gen): funcname = gen.__name__ try: @@ -96,10 +97,12 @@ def _generator_copy(gen): return res + def test(): q = [1,2,3] for i in q: yield i + class TestGen(object): def __init__(self, count=3): @@ -121,6 +124,7 @@ class TestGen(object): def __next__(self): return self.next() + def latlon_coordvars(d): lat_coord = None lon_coord = None @@ -137,9 +141,11 @@ def latlon_coordvars(d): return lat_coord, lon_coord + def is_coordvar(varname): return varname in _COORD_VARS + class IterWrapper(Iterable): """A wrapper class for generators and custom iterable classes which returns a new iterator from the start of the sequence when __iter__ is called""" @@ -172,6 +178,7 @@ def get_iterable(wrfseq): else: return dict(wrfseq) # generator/custom iterable class + # Helper to extract masked arrays from DataArrays that convert to NaN def npvalues(array_type): if not isinstance(array_type, DataArray): @@ -187,6 +194,7 @@ def npvalues(array_type): return result + # Helper utilities for metadata class either(object): def __init__(self, *varnames): @@ -207,6 +215,7 @@ class either(object): raise ValueError("{} are not valid variable names".format( self.varnames)) + class combine_with(object): # Remove remove_idx first, then insert_idx is applied to removed set def __init__(self, varname, remove_dims=None, insert_before=None, @@ -239,6 +248,7 @@ class combine_with(object): return new_dims, new_coords + # This should look like: # [(0, (-3,-2)), # variable 1 # (1, -1)] # variable 2 @@ -387,6 +397,7 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), return False + def _get_global_attr(wrfnc, attr): val = getattr(wrfnc, attr, None) @@ -395,6 +406,7 @@ def _get_global_attr(wrfnc, attr): if len(val) == 1: return val[0] return val + def extract_global_attrs(wrfnc, attrs): if isstr(attrs): @@ -413,6 +425,7 @@ def extract_global_attrs(wrfnc, attrs): return {attr:_get_global_attr(wrfnc, attr) for attr in attrlist} + def extract_dim(wrfnc, dim): if is_multi_file(wrfnc): if not is_mapping(wrfnc): @@ -426,6 +439,7 @@ def extract_dim(wrfnc, dim): return len(d) #netCDF4 return d # PyNIO + def _combine_dict(wrfdict, varname, timeidx, method, meta, _key): """Dictionary combination creates a new left index for each key, then does a cat or join for the list of files for that key""" @@ -656,7 +670,7 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain, is_multifile, time_coord_vals = None if time_coord is not None: # If not from a multifile sequence, then cache the time - # coordinate. Otherwise, handled in cat/join/ + # coordinate. Otherwise, handled in cat/join/ if not is_multifile: time_coord_vals = get_cached_item(_key, time_coord) @@ -798,6 +812,7 @@ def _find_arr_for_time(wrfseq, varname, timeidx, is_moving, meta, _key): else: return _find_reverse(wrfseq, varname, timeidx, is_moving, meta, _key) + # TODO: implement in C def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta, _key): if is_moving is None: @@ -1234,6 +1249,7 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key): return outarr + def combine_files(wrfseq, varname, timeidx, is_moving=None, method="cat", squeeze=True, meta=True, _key=None): @@ -1416,7 +1432,6 @@ def is_staggered(var, wrfnc): return False - def get_left_indexes(ref_var, expected_dims): """Returns the extra left side dimensions for a variable with an expected 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)]) + def iter_left_indexes(dims): """A generator which yields the iteration tuples for a sequence of dimensions sizes. @@ -1448,6 +1464,7 @@ def iter_left_indexes(dims): arg = [py3range(dim) for dim in dims] for idxs in product(*arg): yield idxs + def get_right_slices(var, right_ndims, fixed_val=0): extra_dim_num = var.ndim - right_ndims @@ -1457,6 +1474,7 @@ def get_right_slices(var, right_ndims, fixed_val=0): return tuple([fixed_val]*extra_dim_num + [slice(None)]*right_ndims) + def get_proj_params(wrfnc, timeidx=0, varname=None): proj_params = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", "CEN_LAT", "CEN_LON", @@ -1556,8 +1574,9 @@ def from_args(func, argnames, *args, **kwargs): return result + def _args_to_list2(func, args, kwargs): - argspec = getargspec(func) + argspec = getargspec(func) # Build the full tuple with defaults filled in outargs = [None]*len(argspec.args) @@ -1576,14 +1595,16 @@ def _args_to_list2(func, args, kwargs): return outargs + def _args_to_list3(func, args, kwargs): sig = signature(func) bound = sig.bind(*args, **kwargs) bound.apply_defaults() return [x for x in bound.arguments.values()] + - +# Note: Doesn't allow for **kwargs or *args def args_to_list(func, args, kwargs): """Converts the mixed args/kwargs to a single list of args""" if version_info > (3,): @@ -1607,6 +1628,7 @@ def _arg_location2(func, argname, args, kwargs): return list_args, result_idx + def _arg_location3(func, argname, args, kwargs): sig = signature(func) params = list(sig.parameters.keys()) diff --git a/test/utests.py b/test/utests.py index 4fcc267..a78940e 100644 --- a/test/utests.py +++ b/test/utests.py @@ -6,8 +6,10 @@ import os, sys import subprocess from wrf import (getvar, interplevel, interpline, vertcross, vinterp, - disable_xarray, xarray_enabled, npvalues, - xy_to_ll, ll_to_xy ) + disable_xarray, xarray_enabled, npvalues, + 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" 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 +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, pynio=False): def test(self): @@ -513,6 +532,13 @@ def make_latlon_test(testid, wrf_in, referent, single, multi=False, repeat=3, ref = ref_vals[:,0] 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: xy = ll_to_xy(in_wrfnc, lats, lons) @@ -521,6 +547,12 @@ def make_latlon_test(testid, wrf_in, referent, single, multi=False, repeat=3, 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: # Since this domain is not moving, the reference values are the # same whether there are multiple or single files @@ -536,11 +568,23 @@ def make_latlon_test(testid, wrf_in, referent, single, multi=False, repeat=3, ref = ref_vals[::-1,0] 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: ll = xy_to_ll(in_wrfnc, i_s, j_s) ref = ref_vals[::-1,:] 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