From 27f720184164844c3490a9b8cd9fff51feaa8619 Mon Sep 17 00:00:00 2001 From: Bill Ladwig Date: Wed, 20 Apr 2016 16:49:27 -0600 Subject: [PATCH] Bug fixes for demo --- wrf_open/var/src/python/wrf/var/__init__.py | 220 +----------- wrf_open/var/src/python/wrf/var/decorators.py | 2 +- wrf_open/var/src/python/wrf/var/interp.py | 3 +- wrf_open/var/src/python/wrf/var/latlon.py | 211 +---------- .../var/src/python/wrf/var/latlonutils.py | 227 ++++++++++++ .../var/src/python/wrf/var/metadecorators.py | 57 ++- wrf_open/var/src/python/wrf/var/projection.py | 38 +- wrf_open/var/src/python/wrf/var/routines.py | 166 +++++++++ wrf_open/var/src/python/wrf/var/util.py | 332 +++++++++++------- wrf_open/var/src/python/wrf/var/uvmet.py | 22 +- wrf_open/var/test/utests.py | 2 +- 11 files changed, 698 insertions(+), 582 deletions(-) create mode 100644 wrf_open/var/src/python/wrf/var/latlonutils.py create mode 100644 wrf_open/var/src/python/wrf/var/routines.py diff --git a/wrf_open/var/src/python/wrf/var/__init__.py b/wrf_open/var/src/python/wrf/var/__init__.py index 4737bc8..2b5e45f 100755 --- a/wrf_open/var/src/python/wrf/var/__init__.py +++ b/wrf_open/var/src/python/wrf/var/__init__.py @@ -5,228 +5,24 @@ import warnings from . import config from .config import * -from . import extension -from .extension import * +from . import routines +from .routines import * from . import util from .util import * -from . import cape -from .cape import * -from . import constants -from .constants import * -from . import ctt -from .ctt import * -from . import dbz -from .dbz import * -from . import destag -from .destag import * -from . import dewpoint -from .dewpoint import * -from . import geoht -from .geoht import * -from . import helicity -from .helicity import * from . import interp from .interp import * -from . import latlon -from .latlon import * -from . import omega -from .omega import * -from . import precip -from .precip import * -from . import pressure -from .pressure import * -from . import psadlookup -from .psadlookup import * -from . import pw -from .pw import * -from . import rh -from .rh import * -from . import slp -from .slp import * -from . import temp -from .temp import * -from . import terrain -from .terrain import * -from . import uvmet -from .uvmet import * -from . import vorticity -from .vorticity import * -from . import wind -from .wind import * -from . import times -from .times import * -from . import units -from .units import * from . import projection from .projection import * - -__all__ = ["getvar"] + +__all__ = [] +__all__.extend(routines.__all__) +__all__.extend(interp.__all__) __all__.extend(config.__all__) -__all__.extend( extension.__all__) __all__.extend(util.__all__) -__all__.extend(cape.__all__) -__all__.extend(constants.__all__) -__all__.extend(ctt.__all__) -__all__.extend(dbz.__all__) -__all__.extend(destag.__all__) -__all__.extend(dewpoint.__all__) -__all__.extend(geoht.__all__) -__all__.extend(helicity.__all__) -__all__.extend(interp.__all__) -__all__.extend(latlon.__all__) -__all__.extend(omega.__all__) -__all__.extend(precip.__all__) -__all__.extend(psadlookup.__all__) -__all__.extend(pw.__all__) -__all__.extend(rh.__all__) -__all__.extend(slp.__all__) -__all__.extend(temp.__all__) -__all__.extend(terrain.__all__) -__all__.extend(uvmet.__all__) -__all__.extend(vorticity.__all__) -__all__.extend(wind.__all__) -__all__.extend(times.__all__) -__all__.extend(pressure.__all__) -__all__.extend(units.__all__) __all__.extend(projection.__all__) + -# func is the function to call. kargs are required arguments that should -# not be altered by the user -_FUNC_MAP = {"cape2d" : get_2dcape, - "cape3d" : get_3dcape, - "dbz" : get_dbz, - "maxdbz" : get_max_dbz, - "dp" : get_dp, - "dp2m" : get_dp_2m, - "height" : get_height, - "geopt" : get_geopt, - "srh" : get_srh, - "uhel" : get_uh, - "omega" : get_omega, - "pw" : get_pw, - "rh" : get_rh, - "rh2m" : get_rh_2m, - "slp" : get_slp, - "theta" : get_theta, - "temp" : get_temp, - "tk" : get_tk, - "tc" : get_tc, - "theta_e" : get_eth, - "tv" : get_tv, - "twb" : get_tw, - "terrain" : get_terrain, - "times" : get_times, - "uvmet" : get_uvmet, - "uvmet10" : get_uvmet10, - "avo" : get_avo, - "pvo" : get_pvo, - "ua" : get_u_destag, - "va" : get_v_destag, - "wa" : get_w_destag, - "lat" : get_lat, - "lon" : get_lon, - "pressure" : get_pressure_hpa, - "pres" : get_pressure, - "wspd_wdir" : get_destag_wspd_wdir, - "wspd_wdir10" : get_destag_wspd_wdir10, - "wspd_wdir_uvmet" : get_uvmet_wspd_wdir, - "wspd_wdir_uvmet10" : get_uvmet10_wspd_wdir, - "ctt" : get_ctt - } - -_VALID_KARGS = {"cape2d" : ["missing"], - "cape3d" : ["missing"], - "dbz" : ["do_variant", "do_liqskin"], - "maxdbz" : ["do_variant", "do_liqskin"], - "dp" : ["units"], - "dp2m" : ["units"], - "height" : ["msl", "units"], - "geopt" : [], - "srh" : ["top"], - "uhel" : ["bottom", "top"], - "omega" : [], - "pw" : [], - "rh" : [], - "rh2m" : [], - "slp" : ["units"], - "temp" : ["units"], - "tk" : [], - "tc" : [], - "theta" : ["units"], - "theta_e" : ["units"], - "tv" : ["units"], - "twb" : ["units"], - "terrain" : ["units"], - "times" : [], - "uvmet" : ["units"], - "uvmet10" : ["units"], - "avo" : [], - "pvo" : [], - "ua" : ["units"], - "va" : ["units"], - "wa" : ["units"], - "lat" : [], - "lon" : [], - "pres" : ["units"], - "pressure" : ["units"], - "wspddir" : ["units"], - "wspddir_uvmet" : ["units"], - "wspddir_uvmet10" : ["units"], - "ctt" : [], - "default" : [] - } - -_ALIASES = {"cape_2d" : "cape2d", - "cape_3d" : "cape3d", - "eth" : "theta_e", - "mdbz" : "maxdbz", - "geopotential" : "geopt", - "helicity" : "srh", - "latitude" : "lat", - "longitude" : "lon", - "omg" : "omega", - "p" : "pres", - "rh2" : "rh2m", - "z": "height", - "ter" : "terrain", - "updraft_helicity" : "uhel", - "td" : "dp", - "td2" : "dp2m" - } - -class ArgumentError(Exception): - def __init__(self, msg): - self.msg = msg - - def __str__(self): - return self.msg - -def _undo_alias(alias): - actual = _ALIASES.get(alias, None) - if actual is None: - return alias - else: - return actual - -def _check_kargs(var, kargs): - for arg in kargs.iterkeys(): - if arg not in _VALID_KARGS[var]: - raise ArgumentError("'%s' is an invalid keyword " - "argument for '%s'" % (arg, var)) - - -def getvar(wrfnc, var, timeidx=0, - method="cat", squeeze=True, cache=None, - **kargs): - if is_standard_wrf_var(wrfnc, var): - return extract_vars(wrfnc, timeidx, var, method, squeeze, cache)[var] - - actual_var = _undo_alias(var) - if actual_var not in _VALID_KARGS: - raise ArgumentError("'%s' is not a valid variable name" % (var)) - - _check_kargs(actual_var, kargs) - return _FUNC_MAP[actual_var](wrfnc,timeidx,**kargs) + diff --git a/wrf_open/var/src/python/wrf/var/decorators.py b/wrf_open/var/src/python/wrf/var/decorators.py index 2a57a1a..7aaa6a6 100644 --- a/wrf_open/var/src/python/wrf/var/decorators.py +++ b/wrf_open/var/src/python/wrf/var/decorators.py @@ -93,7 +93,7 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, # 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)]) + left_and_slice_idxs = left_idxs + (slice(None), ) # Slice the args if applicable new_args = [arg[left_and_slice_idxs] diff --git a/wrf_open/var/src/python/wrf/var/interp.py b/wrf_open/var/src/python/wrf/var/interp.py index 89559f8..ebeb7d0 100755 --- a/wrf_open/var/src/python/wrf/var/interp.py +++ b/wrf_open/var/src/python/wrf/var/interp.py @@ -7,8 +7,7 @@ import numpy.ma as ma from .extension import (interpz3d, interp2dxy, interp1d, smooth2d, monotonic, vintrp, computevertcross, computeinterpline) -from .decorators import (handle_left_iter, handle_casting, - handle_extract_transpose) + from .metadecorators import set_interp_metadata from .util import extract_vars, is_staggered from .interputils import get_xy, get_xy_z_params diff --git a/wrf_open/var/src/python/wrf/var/latlon.py b/wrf_open/var/src/python/wrf/var/latlon.py index c709458..1072f9c 100755 --- a/wrf_open/var/src/python/wrf/var/latlon.py +++ b/wrf_open/var/src/python/wrf/var/latlon.py @@ -1,42 +1,13 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) -from collections import Iterable - -import numpy as np - -from .config import xarray_enabled -from .constants import Constants -from .extension import computeij, computell -from .util import (extract_vars, extract_global_attrs, - either, _is_moving_domain, _is_multi_time_req, - iter_left_indexes) +from .util import (extract_vars) +from .latlonutils import (_lat_varname, _lon_varname, ll_to_ij, ij_to_ll) from .metadecorators import set_latlon_metadata -if xarray_enabled(): - from xarray import DataArray __all__ = ["get_lat", "get_lon", "get_ij", "get_ll"] -def _lat_varname(wrfnc, stagger): - if stagger is None or stagger.lower() == "m": - varname = either("XLAT", "XLAT_M")(wrfnc) - elif stagger.lower() == "u" or stagger.lower() == "v": - varname = "XLAT_{}".format(stagger.upper()) - else: - raise ValueError("invalid 'stagger' value") - - return varname - -def _lon_varname(wrfnc, stagger): - if stagger is None or stagger.lower() == "m": - varname = either("XLONG", "XLONG_M")(wrfnc) - elif stagger.lower() == "u" or stagger.lower() == "v": - varname = "XLONG_{}".format(stagger.upper()) - else: - raise ValueError("invalid 'stagger' value") - - return varname def get_lat(wrfnc, timeidx=0, stagger=None, method="cat", squeeze=True, cache=None): @@ -46,6 +17,7 @@ def get_lat(wrfnc, timeidx=0, stagger=None, nometa=False) return lat_var[varname] + def get_lon(wrfnc, timeidx=0, stagger=None, method="cat", squeeze=True, cache=None): @@ -56,182 +28,19 @@ def get_lon(wrfnc, timeidx=0, stagger=None, return lon_var[varname] -def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache): - if timeidx < 0: - raise ValueError("'timeidx' must be greater than 0") - - attrs = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", "TRUELAT1", - "TRUELAT2", "STAND_LON", - "DX", "DY")) - map_proj = attrs["MAP_PROJ"] - truelat1 = attrs["TRUELAT1"] - truelat2 = attrs["TRUELAT2"] - stdlon = attrs["STAND_LON"] - dx = attrs["DX"] - dy = attrs["DY"] - - if map_proj == 6: - pole_attrs = extract_global_attrs(wrfnc, attrs=("POLE_LAT", - "POLE_LON")) - pole_lat = pole_attrs["POLE_LAT"] - pole_lon = pole_attrs["POLE_LON"] - latinc = (dy*360.0)/2.0 / Constants.PI/Constants.WRF_EARTH_RADIUS - loninc = (dx*360.0)/2.0 / Constants.PI/Constants.WRF_EARTH_RADIUS - else: - pole_lat = 90.0 - pole_lon = 0.0 - latinc = 0.0 - loninc = 0.0 - - latvar = _lat_varname(stagger) - lonvar = _lon_varname(stagger) - - lat_timeidx = timeidx - # Only need all the lats/lons if it's a moving domain file/files - if _is_multi_time_req(timeidx): - if not _is_moving_domain(wrfnc, latvar=latvar, lonvar=lonvar): - lat_timeidx = 0 - - xlat = get_lat(wrfnc, lat_timeidx, stagger, method, squeeze, cache) - xlon = get_lon(wrfnc, lat_timeidx, stagger, method, squeeze, cache) - - ref_lat = np.ravel(xlat[...,0,0]) - ref_lon = np.ravel(xlon[...,0,0]) - - # Note: fortran index - known_i = 1.0 - known_j = 1.0 - - return (map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, - pole_lat,pole_lon,known_i,known_j,dx,latinc, - loninc) - @set_latlon_metadata(ij=True) def get_ij(wrfnc, latitude, longitude, timeidx=0, stagger=None, method="cat", squeeze=True, cache=None): + return ll_to_ij(wrfnc, latitude, longitude, timeidx, stagger, + method, squeeze, cache) - (map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, - pole_lat,pole_lon,known_i,known_j,dx,latinc, - loninc) = _get_proj_params(wrfnc, timeidx) - if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str): - lats = np.asarray(latitude) - lons = np.asarray(longitude) - - if lats.ndim > 1: - lats = lats.ravel() - - if lons.ndim > 1: - lons = lons.ravel() - - if (lats.size != lons.size): - raise ValueError("'latitude' and 'longitude' " - "must be the same length") - - - if ref_lat.size == 1: - outdim = [lats.size, 2] - extra_dims = outdim[0] - else: - # Moving domain will have moving ref_lats/ref_lons - outdim = [lats.size, ref_lat.size, 2] - extra_dims = outdim[0:2] - - res = np.empty(outdim, np.float64) - - for left_idxs in iter_left_indexes(extra_dims): - left_and_slice_idxs = left_idxs + [slice(None, None, None)] - - if ref_lat.size == 1: - ref_lat_val = ref_lat[0] - ref_lon_val = ref_lon[0] - else: - ref_lat_val = ref_lat[left_idxs[-1]] - ref_lon_val = ref_lon[left_idxs[-1]] - - lat = lats[left_idxs[0]] - lon = lons[left_idxs[0]] - - ij = computeij(map_proj, truelat1, truelat2, stdlon, - ref_lat_val, ref_lon_val, pole_lat, pole_lon, - known_i, known_j, dx, latinc, loninc, - lat, lon) - - res[left_and_slice_idxs] = ij[:] - - else: - - res = computeij(map_proj, truelat1, truelat2, stdlon, - ref_lat, ref_lon, pole_lat, pole_lon, - known_i, known_j, dx, latinc, loninc, - latitude, longitude) - - return res - - -@set_latlon_metadata(ij=False) -def get_ll(wrfnc, i, j, timeidx=0): - if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str): - raise TypeError("'get_ll' is only applicabe for single files") - - (map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, - pole_lat,pole_lon,known_i,known_j,dx,latinc, - loninc) = _get_proj_params(wrfnc, timeidx) - - if isinstance(wrfnc, Iterable) and not isinstance(wrfnc, str): - i_arr = np.asarray(i) - j_arr = np.asarray(j) - - if i_arr.ndim > 1: - i_arr = i_arr.ravel() - - if j_arr.ndim > 1: - j_arr = j_arr.ravel() - - if (i_arr.size != j_arr.size): - raise ValueError("'i' and 'j' " - "must be the same length") - - if ref_lat.size == 1: - outdim = [i_arr.size, 2] - extra_dims = outdim[0] - else: - # Moving domain will have moving ref_lats/ref_lons - outdim = [i_arr.size, ref_lat.size, 2] - extra_dims = outdim[0:2] - - res = np.empty(outdim, np.float64) - - for left_idxs in iter_left_indexes(extra_dims): - left_and_slice_idxs = left_idxs + [slice(None, None, None)] - - if ref_lat.size == 1: - ref_lat_val = ref_lat[0] - ref_lon_val = ref_lon[0] - else: - ref_lat_val = ref_lat[left_idxs[-1]] - ref_lon_val = ref_lon[left_idxs[-1]] - - i_val = i_arr[left_idxs[0]] - j_val = j_arr[left_idxs[0]] - - ll = computell(map_proj, truelat1, truelat2, - stdlon, ref_lat_val, ref_lon_val, - pole_lat, pole_lon, known_i, known_j, - dx, latinc, loninc, - i_val, j_val) - - res[left_and_slice_idxs] = ll[:] - - else: - res = computell(map_proj, truelat1, truelat2, - stdlon, ref_lat, ref_lon, - pole_lat, pole_lon, known_i, known_j, - dx, latinc, loninc, - i_val, j_val) - - return res +@set_latlon_metadata(ij=False) +def get_ll(wrfnc, i, j, timeidx=0, + stagger=None, method="cat", squeeze=True, cache=None): + return ij_to_ll(wrfnc, i, j, timeidx, stagger, + method, squeeze, cache) \ No newline at end of file diff --git a/wrf_open/var/src/python/wrf/var/latlonutils.py b/wrf_open/var/src/python/wrf/var/latlonutils.py new file mode 100644 index 0000000..281d9b5 --- /dev/null +++ b/wrf_open/var/src/python/wrf/var/latlonutils.py @@ -0,0 +1,227 @@ +from __future__ import (absolute_import, division, print_function, + unicode_literals) + +from collections import Iterable + +import numpy as np + +from .constants import Constants +from .extension import computeij, computell +from .util import (extract_vars, extract_global_attrs, + either, _is_moving_domain, _is_multi_time_req, + iter_left_indexes, _is_mapping, _is_multi_file, + viewkeys) + +def _lat_varname(wrfnc, stagger): + if stagger is None or stagger.lower() == "m": + varname = either("XLAT", "XLAT_M")(wrfnc) + elif stagger.lower() == "u" or stagger.lower() == "v": + varname = "XLAT_{}".format(stagger.upper()) + else: + raise ValueError("invalid 'stagger' value") + + return varname + +def _lon_varname(wrfnc, stagger): + if stagger is None or stagger.lower() == "m": + varname = either("XLONG", "XLONG_M")(wrfnc) + elif stagger.lower() == "u" or stagger.lower() == "v": + varname = "XLONG_{}".format(stagger.upper()) + else: + raise ValueError("invalid 'stagger' value") + + return varname + +def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache): + if timeidx < 0: + raise ValueError("'timeidx' must be greater than 0") + + attrs = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", "TRUELAT1", + "TRUELAT2", "STAND_LON", + "DX", "DY")) + map_proj = attrs["MAP_PROJ"] + truelat1 = attrs["TRUELAT1"] + truelat2 = attrs["TRUELAT2"] + stdlon = attrs["STAND_LON"] + dx = attrs["DX"] + dy = attrs["DY"] + + if map_proj == 6: + pole_attrs = extract_global_attrs(wrfnc, attrs=("POLE_LAT", + "POLE_LON")) + pole_lat = pole_attrs["POLE_LAT"] + pole_lon = pole_attrs["POLE_LON"] + latinc = (dy*360.0)/2.0 / Constants.PI/Constants.WRF_EARTH_RADIUS + loninc = (dx*360.0)/2.0 / Constants.PI/Constants.WRF_EARTH_RADIUS + else: + pole_lat = 90.0 + pole_lon = 0.0 + latinc = 0.0 + loninc = 0.0 + + latvar = _lat_varname(wrfnc, stagger) + lonvar = _lon_varname(wrfnc, stagger) + + lat_timeidx = timeidx + + is_moving = _is_moving_domain(wrfnc, latvar=latvar, lonvar=lonvar) + + # Only need one file and one time if the domain is not moving + if not is_moving: + if _is_multi_time_req(timeidx): + lat_timeidx = 0 + + if _is_multi_file(wrfnc): + if not _is_mapping(wrfnc): + wrfnc = next(iter(wrfnc)) # only need one file + else: + wrfnc = wrfnc[next(iter(viewkeys(wrfnc)))] + return _get_proj_params(wrfnc, timeidx, stagger, + method, squeeze, cache) + + xlat = extract_vars(wrfnc, lat_timeidx, (latvar,), method, squeeze, cache, + nometa=True)[latvar] + xlon = extract_vars(wrfnc, lat_timeidx, (lonvar,), method, squeeze, cache, + nometa=True)[lonvar] + + ref_lat = np.ravel(xlat[...,0,0]) + ref_lon = np.ravel(xlon[...,0,0]) + + # Note: fortran index + known_i = 1.0 + known_j = 1.0 + + return (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, + pole_lat, pole_lon, known_i, known_j, dx, latinc, loninc) + + +def ll_to_ij(wrfnc, latitude, longitude, timeidx=0, + stagger=None, method="cat", squeeze=True, cache=None): + + (map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, + pole_lat,pole_lon,known_i,known_j,dx,latinc, + loninc) = _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache) + + if isinstance(latitude, Iterable): + lats = np.asarray(latitude) + lons = np.asarray(longitude) + + if lats.ndim > 1: + lats = lats.ravel() + + if lons.ndim > 1: + lons = lons.ravel() + + if (lats.size != lons.size): + raise ValueError("'latitude' and 'longitude' " + "must be the same length") + + + if ref_lat.size == 1: + outdim = [lats.size, 2] + extra_dims = [outdim[0]] + else: + # Moving domain will have moving ref_lats/ref_lons + outdim = [lats.size, ref_lat.size, 2] + extra_dims = outdim[0:2] + + res = np.empty(outdim, np.float64) + + for left_idxs in iter_left_indexes(extra_dims): + left_and_slice_idxs = left_idxs + (slice(None), ) + + if ref_lat.size == 1: + ref_lat_val = ref_lat[0] + ref_lon_val = ref_lon[0] + else: + ref_lat_val = ref_lat[left_idxs[-1]] + ref_lon_val = ref_lon[left_idxs[-1]] + + lat = lats[left_idxs[0]] + lon = lons[left_idxs[0]] + + ij = computeij(map_proj, truelat1, truelat2, stdlon, + ref_lat_val, ref_lon_val, pole_lat, pole_lon, + known_i, known_j, dx, latinc, loninc, + lat, lon) + + res[left_and_slice_idxs] = ij[:] + + else: + + res = computeij(map_proj, truelat1, truelat2, stdlon, + ref_lat, ref_lon, pole_lat, pole_lon, + known_i, known_j, dx, latinc, loninc, + latitude, longitude) + + if squeeze: + res = res.squeeze() + + return res + +def ij_to_ll(wrfnc, i, j, timeidx=0, + stagger=None, method="cat", squeeze=True, cache=None): + + (map_proj,truelat1,truelat2,stdlon,ref_lat,ref_lon, + pole_lat,pole_lon,known_i,known_j,dx,latinc, + loninc) = _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache) + + if isinstance(i, Iterable): + i_arr = np.asarray(i) + j_arr = np.asarray(j) + + if i_arr.ndim > 1: + i_arr = i_arr.ravel() + + if j_arr.ndim > 1: + j_arr = j_arr.ravel() + + if (i_arr.size != j_arr.size): + raise ValueError("'i' and 'j' " + "must be the same length") + + if ref_lat.size == 1: + outdim = [i_arr.size, 2] + extra_dims = [outdim[0]] + else: + # Moving domain will have moving ref_lats/ref_lons + outdim = [i_arr.size, ref_lat.size, 2] + extra_dims = outdim[0:2] + + res = np.empty(outdim, np.float64) + + for left_idxs in iter_left_indexes(extra_dims): + left_and_slice_idxs = left_idxs + (slice(None), ) + + if ref_lat.size == 1: + ref_lat_val = ref_lat[0] + ref_lon_val = ref_lon[0] + else: + ref_lat_val = ref_lat[left_idxs[-1]] + ref_lon_val = ref_lon[left_idxs[-1]] + + i_val = i_arr[left_idxs[0]] + j_val = j_arr[left_idxs[0]] + + ll = computell(map_proj, truelat1, truelat2, + stdlon, ref_lat_val, ref_lon_val, + pole_lat, pole_lon, known_i, known_j, + dx, latinc, loninc, + i_val, j_val) + + res[left_and_slice_idxs] = ll[:] + + else: + i_val = i + j_val = j + + res = computell(map_proj, truelat1, truelat2, + stdlon, ref_lat, ref_lon, + pole_lat, pole_lon, known_i, known_j, + dx, latinc, loninc, + i_val, j_val) + + if squeeze: + res = res.squeeze() + + return res \ No newline at end of file diff --git a/wrf_open/var/src/python/wrf/var/metadecorators.py b/wrf_open/var/src/python/wrf/var/metadecorators.py index 1d0061f..678a5c2 100644 --- a/wrf_open/var/src/python/wrf/var/metadecorators.py +++ b/wrf_open/var/src/python/wrf/var/metadecorators.py @@ -9,8 +9,9 @@ import numpy.ma as ma from .util import (viewkeys, viewitems, extract_vars, combine_with, either, from_args, arg_location, - _is_coord_var, XYCoord, npvalues) + _is_coord_var, CoordPair, npvalues) from .interputils import get_xy_z_params, get_xy +from .latlonutils import ij_to_ll, ll_to_ij from .config import xarray_enabled if xarray_enabled(): @@ -211,34 +212,52 @@ def set_wind_metadata(copy_varname, name, description, def set_latlon_metadata(ij=False): @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): + if not xarray_enabled(): return wrapped(*args, **kwargs) res = wrapped(*args, **kwargs) - argnames = ("latitude", "longitude") if not ij else ("i", "j") + # Want to preserve the input coordinate pair in metadata + if res.ndim == 1: + res = res[np.newaxis, :] + + argnames = ["i", "j"] if not ij else ["latitude", "longitude"] + argnames.append("squeeze") outname = "latlon" if not ij else "ij" - if res.ndim <= 2: - dimnames = (["lat_lon", "i_j"] if not ij - else ["i_j", "lat_lon"]) + if res.ndim == 2: + dimnames = (["ij", "lat_lon"] if not ij + else ["latlon", "i_j"]) else: - dimnames = (["lat_lon", "domain", "i_j"] if not ij - else ["i_j", "domain", "lat_lon"]) - + dimnames = (["ij", "domain", "lat_lon"] if not ij + else ["latlon", "domain", "i_j"]) argvars = from_args(wrapped, argnames, *args, **kwargs) var1 = argvars[argnames[0]] var2 = argvars[argnames[1]] + squeeze = argvars["squeeze"] arr1 = np.asarray(var1).ravel() arr2 = np.asarray(var2).ravel() coords = {} - coords[dimnames[0]] = [x for x in zip(arr1, arr2)] + if not ij: + coords["coord_pair"] = (dimnames[0], [CoordPair(i=x[0], j=x[1]) + for x in zip(arr1, arr2)]) + coords[dimnames[-1]] = ["lat", "lon"] + else: + coords["coord_pair"] = (dimnames[0], [CoordPair(lat=x[0], lon=x[1]) + for x in zip(arr1, arr2)]) + coords[dimnames[-1]] = ["i", "j"] - return DataArray(res, name=outname, dims=dimnames, coords=coords) + da = DataArray(res, name=outname, dims=dimnames, coords=coords) + + if squeeze: + da = da.squeeze() + + return da return func_wrapper @@ -329,6 +348,10 @@ def _set_horiz_meta(wrapped, instance, args, kwargs): if vert_units is not None else "{0}".format(desiredloc)) + name_levelstr = ("{0}_{1}".format(desiredloc, vert_units) + if vert_units is not None + else "{0}".format(desiredloc)) + if isinstance(field3d, DataArray): outcoords = OrderedDict() outattrs = OrderedDict() @@ -337,7 +360,7 @@ def _set_horiz_meta(wrapped, instance, args, kwargs): outdimnames.remove(field3d.dims[-3]) del outcoords[field3d.dims[-3]] outattrs.update(field3d.attrs) - outname = "{0}_{1}".format(field3d.name, levelstr) + outname = "{0}_{1}".format(field3d.name, name_levelstr) else: outname = "field3d_{0}".format(levelstr) @@ -436,8 +459,8 @@ def _set_cross_meta(wrapped, instance, args, kwargs): except KeyError: pass - outcoords["xy"] = [XYCoord(xy[i,0], xy[i,1]) - for i in xrange(xy.shape[-2])] + outcoords["xy_loc"] = ("xy", [CoordPair(xy[i,0], xy[i,1]) + for i in xrange(xy.shape[-2])]) outcoords["vertical"] = z_var2d[:] @@ -525,8 +548,8 @@ def _set_line_meta(wrapped, instance, args, kwargs): except KeyError: pass - outcoords["xy"] = [XYCoord(xy[i,0], xy[i,1]) - for i in xrange(xy.shape[-2])] + outcoords["xy_loc"] = ("xy", [CoordPair(xy[i,0], xy[i,1]) + for i in xrange(xy.shape[-2])]) else: outname = "field2d_line" @@ -617,8 +640,8 @@ def _set_2dxy_meta(wrapped, instance, args, kwargs): outname = "{0}_xy".format(field3d.name) - outcoords["xy"] = [XYCoord(xy[i,0], xy[i,1]) - for i in xrange(xy.shape[-2])] + outcoords["xy_loc"] = ("xy", [CoordPair(xy[i,0], xy[i,1]) + for i in xrange(xy.shape[-2])]) for key in ("MemoryOrder",): try: diff --git a/wrf_open/var/src/python/wrf/var/projection.py b/wrf_open/var/src/python/wrf/var/projection.py index 972323c..7e15413 100644 --- a/wrf_open/var/src/python/wrf/var/projection.py +++ b/wrf_open/var/src/python/wrf/var/projection.py @@ -15,13 +15,13 @@ if basemap_enabled(): if pyngl_enabled(): from Ngl import Resources -__all__ = ["WrfProj", "LambertConformalProj", "MercatorProj", - "PolarStereographicProj", "LatLonProj", "RotLatLonProj", +__all__ = ["WrfProj", "LambertConformal", "Mercator", + "PolarStereographic", "LatLon", "RotatedLatLon", "getproj"] if cartopy_enabled(): - class MercatorWithLatTsProj(crs.Mercator): + class MercatorWithLatTS(crs.Mercator): def __init__(self, central_longitude=0.0, latitude_true_scale=0.0, min_latitude=-80.0, @@ -181,10 +181,10 @@ class WrfProj(object): return self._cf_params -class LambertConformalProj(WrfProj): +class LambertConformal(WrfProj): def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None, **proj_params): - super(LambertConformalProj, self).__init__(bottom_left, + super(LambertConformal, self).__init__(bottom_left, top_right, lats, lons, **proj_params) self._std_parallels = [self.truelat1] @@ -287,10 +287,10 @@ class LambertConformalProj(WrfProj): self.stand_lon)) return _proj4 -class MercatorProj(WrfProj): +class Mercator(WrfProj): def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None, **proj_params): - super(MercatorProj, self).__init__(bottom_left, top_right, + super(Mercator, self).__init__(bottom_left, top_right, lats, lons, **proj_params) self._lat_ts = (None @@ -354,7 +354,7 @@ class MercatorProj(WrfProj): globe = self._globe) else: - _cartopy = MercatorWithLatTsProj( + _cartopy = MercatorWithLatTS( central_longitude = self.stand_lon, latitude_true_scale = self._lat_ts, globe = self._globe) @@ -387,10 +387,10 @@ class MercatorProj(WrfProj): return _proj4 -class PolarStereographicProj(WrfProj): +class PolarStereographic(WrfProj): def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None, **proj_params): - super(PolarStereographicProj, self).__init__(bottom_left, + super(PolarStereographic, self).__init__(bottom_left, top_right, lats, lons, **proj_params) self._hemi = -90. if self.truelat1 < 0 else 90. self._lat_ts = (None @@ -486,10 +486,10 @@ class PolarStereographicProj(WrfProj): -class LatLonProj(WrfProj): +class LatLon(WrfProj): def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None, **proj_params): - super(LatLonProj, self).__init__(bottom_left, top_right, + super(LatLon, self).__init__(bottom_left, top_right, lats, lons, **proj_params) @property @@ -588,10 +588,10 @@ class LatLonProj(WrfProj): # this reason, the proj4 string for this class will use cartopy's values # to keep things in the -180 to 180, -90 to 90 range. # 12) This projection makes me sad. -class RotLatLonProj(WrfProj): +class RotatedLatLon(WrfProj): def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None, **proj_params): - super(RotLatLonProj, self).__init__(bottom_left, top_right, + super(RotatedLatLon, self).__init__(bottom_left, top_right, lats, lons, **proj_params) # Need to determine hemisphere, typically pole_lon is 0 for southern @@ -726,21 +726,21 @@ def getproj(bottom_left=None, top_right=None, proj_type = proj_params.get("MAP_PROJ", 0) if proj_type == 1: - return LambertConformalProj(bottom_left, top_right, + return LambertConformal(bottom_left, top_right, lats, lons, **proj_params) elif proj_type == 2: - return PolarStereographicProj(bottom_left, top_right, + return PolarStereographic(bottom_left, top_right, lats, lons, **proj_params) elif proj_type == 3: - return MercatorProj(bottom_left, top_right, + return Mercator(bottom_left, top_right, lats, lons, **proj_params) elif proj_type == 0 or proj_type == 6: if (proj_params.get("POLE_LAT", None) == 90. and proj_params.get("POLE_LON", None) == 0.): - return LatLonProj(bottom_left, top_right, + return LatLon(bottom_left, top_right, lats, lons, **proj_params) else: - return RotLatLonProj(bottom_left, top_right, + return RotatedLatLon(bottom_left, top_right, lats, lons, **proj_params) else: # Unknown projection diff --git a/wrf_open/var/src/python/wrf/var/routines.py b/wrf_open/var/src/python/wrf/var/routines.py new file mode 100644 index 0000000..858bf12 --- /dev/null +++ b/wrf_open/var/src/python/wrf/var/routines.py @@ -0,0 +1,166 @@ + + +from .util import _unpack_sequence, is_standard_wrf_var, extract_vars +from .cape import get_2dcape, get_3dcape +from .ctt import get_ctt +from .dbz import get_dbz, get_max_dbz +from .dewpoint import get_dp, get_dp_2m +from .geoht import get_geopt, get_height +from .helicity import get_srh, get_uh +from .latlon import get_lat, get_lon +from .omega import get_omega +from .pressure import get_pressure, get_pressure_hpa +from .pw import get_pw +from .rh import get_rh, get_rh_2m +from .slp import get_slp +from .temp import get_tc, get_eth, get_temp, get_theta, get_tk, get_tv, get_tw +from .terrain import get_terrain +from .uvmet import (get_uvmet, get_uvmet10, get_uvmet10_wspd_wdir, + get_uvmet_wspd_wdir) +from .vorticity import get_avo, get_pvo +from .wind import (get_destag_wspd_wdir, get_destag_wspd_wdir10, + get_u_destag, get_v_destag, get_w_destag) +from .times import get_times + +__all__ = ["getvar"] + +# func is the function to call. kargs are required arguments that should +# not be altered by the user +_FUNC_MAP = {"cape2d" : get_2dcape, + "cape3d" : get_3dcape, + "dbz" : get_dbz, + "maxdbz" : get_max_dbz, + "dp" : get_dp, + "dp2m" : get_dp_2m, + "height" : get_height, + "geopt" : get_geopt, + "srh" : get_srh, + "uhel" : get_uh, + "omega" : get_omega, + "pw" : get_pw, + "rh" : get_rh, + "rh2m" : get_rh_2m, + "slp" : get_slp, + "theta" : get_theta, + "temp" : get_temp, + "tk" : get_tk, + "tc" : get_tc, + "theta_e" : get_eth, + "tv" : get_tv, + "twb" : get_tw, + "terrain" : get_terrain, + "times" : get_times, + "uvmet" : get_uvmet, + "uvmet10" : get_uvmet10, + "avo" : get_avo, + "pvo" : get_pvo, + "ua" : get_u_destag, + "va" : get_v_destag, + "wa" : get_w_destag, + "lat" : get_lat, + "lon" : get_lon, + "pressure" : get_pressure_hpa, + "pres" : get_pressure, + "wspd_wdir" : get_destag_wspd_wdir, + "wspd_wdir10" : get_destag_wspd_wdir10, + "wspd_wdir_uvmet" : get_uvmet_wspd_wdir, + "wspd_wdir_uvmet10" : get_uvmet10_wspd_wdir, + "ctt" : get_ctt + } + +_VALID_KARGS = {"cape2d" : ["missing"], + "cape3d" : ["missing"], + "dbz" : ["do_variant", "do_liqskin"], + "maxdbz" : ["do_variant", "do_liqskin"], + "dp" : ["units"], + "dp2m" : ["units"], + "height" : ["msl", "units"], + "geopt" : [], + "srh" : ["top"], + "uhel" : ["bottom", "top"], + "omega" : [], + "pw" : [], + "rh" : [], + "rh2m" : [], + "slp" : ["units"], + "temp" : ["units"], + "tk" : [], + "tc" : [], + "theta" : ["units"], + "theta_e" : ["units"], + "tv" : ["units"], + "twb" : ["units"], + "terrain" : ["units"], + "times" : [], + "uvmet" : ["units"], + "uvmet10" : ["units"], + "avo" : [], + "pvo" : [], + "ua" : ["units"], + "va" : ["units"], + "wa" : ["units"], + "lat" : [], + "lon" : [], + "pres" : ["units"], + "pressure" : ["units"], + "wspddir" : ["units"], + "wspddir_uvmet" : ["units"], + "wspddir_uvmet10" : ["units"], + "ctt" : [], + "default" : [] + } + +_ALIASES = {"cape_2d" : "cape2d", + "cape_3d" : "cape3d", + "eth" : "theta_e", + "mdbz" : "maxdbz", + "geopotential" : "geopt", + "helicity" : "srh", + "latitude" : "lat", + "longitude" : "lon", + "omg" : "omega", + "p" : "pres", + "rh2" : "rh2m", + "z": "height", + "ter" : "terrain", + "updraft_helicity" : "uhel", + "td" : "dp", + "td2" : "dp2m" + } + +class ArgumentError(Exception): + def __init__(self, msg): + self.msg = msg + + def __str__(self): + return self.msg + +def _undo_alias(alias): + actual = _ALIASES.get(alias, None) + if actual is None: + return alias + else: + return actual + +def _check_kargs(var, kargs): + for arg in kargs.iterkeys(): + if arg not in _VALID_KARGS[var]: + raise ArgumentError("'%s' is an invalid keyword " + "argument for '%s'" % (arg, var)) + + +def getvar(wrfnc, var, timeidx=0, + method="cat", squeeze=True, cache=None, + **kargs): + + wrfnc = _unpack_sequence(wrfnc) + + if is_standard_wrf_var(wrfnc, var): + return extract_vars(wrfnc, timeidx, var, method, squeeze, cache)[var] + + actual_var = _undo_alias(var) + if actual_var not in _VALID_KARGS: + raise ArgumentError("'%s' is not a valid variable name" % (var)) + + _check_kargs(actual_var, kargs) + return _FUNC_MAP[actual_var](wrfnc,timeidx,**kargs) diff --git a/wrf_open/var/src/python/wrf/var/util.py b/wrf_open/var/src/python/wrf/var/util.py index 9f037ad..c03bb7c 100644 --- a/wrf_open/var/src/python/wrf/var/util.py +++ b/wrf_open/var/src/python/wrf/var/util.py @@ -20,8 +20,8 @@ __all__ = ["extract_vars", "extract_global_attrs", "extract_dim", "combine_files", "is_standard_wrf_var", "extract_times", "iter_left_indexes", "get_left_indexes", "get_right_slices", "is_staggered", "get_proj_params", "viewitems", "viewkeys", - "viewvalues", "combine_with", "from_args", "arg_location", - "args_to_list", "npvalues", "XYCoord"] + "viewvalues", "combine_with", "either", "from_args", "arg_location", + "args_to_list", "npvalues", "CoordPair"] _COORD_PAIR_MAP = {"XLAT" : ("XLAT", "XLONG"), @@ -39,9 +39,6 @@ _COORD_PAIR_MAP = {"XLAT" : ("XLAT", "XLONG"), _COORD_VARS = ("XLAT", "XLONG", "XLAT_M", "XLONG_M", "XLAT_U", "XLONG_U", "XLAT_V", "XLONG_V", "CLAT", "CLONG") - - - def _is_coord_var(varname): return varname in _COORD_VARS @@ -61,6 +58,123 @@ def _is_multi_file(wrfnc): def _is_mapping(wrfnc): return isinstance(wrfnc, Mapping) +def _unpack_sequence(wrfseq): + """Unpacks generators in to lists or dictionaries if applicable, otherwise + returns the original object. + + This is apparently the easiest and + fastest way of being able to re-iterate through generators when used + more than once. + + """ + if not _is_multi_file(wrfseq): + return wrfseq + else: + if not _is_mapping(wrfseq): + if isinstance(wrfseq, (list, tuple)): + return wrfseq + else: + return list(wrfseq) # generator/custom iterable class + else: + if isinstance(wrfseq, dict): + return wrfseq + else: + return dict(wrfseq) # generator/custom iterable class + +# Dictionary python 2-3 compatibility stuff +def viewitems(d): + func = getattr(d, "viewitems", None) + if func is None: + func = d.items + return func() + + +def viewkeys(d): + func = getattr(d, "viewkeys", None) + if func is None: + func = d.keys + return func() + + +def viewvalues(d): + func = getattr(d, "viewvalues", None) + if func is None: + func = d.values + return func() + +def isstr(s): + try: + return isinstance(s, basestring) + except NameError: + return isinstance(s, str) + +# Helper to extract masked arrays from DataArrays that convert to NaN +def npvalues(da): + if not isinstance(da, DataArray): + result = da + else: + try: + fill_value = da.attrs["_FillValue"] + except KeyError: + result = da.values + else: + result = ma.masked_invalid(da.values, copy=False) + result.set_fill_value(fill_value) + + return result + +# Helper utilities for metadata +class either(object): + def __init__(self, *varnames): + self.varnames = varnames + + def __call__(self, wrfnc): + if _is_multi_file(wrfnc): + if not _is_mapping(wrfnc): + wrfnc = next(iter(wrfnc)) + else: + entry = wrfnc[next(iter(viewkeys(wrfnc)))] + return self(entry) + + for varname in self.varnames: + if varname in wrfnc.variables: + return varname + + raise ValueError("{} are not valid variable names".format( + self.varnames)) + +class combine_with: + # Remove remove_idx first, then insert_idx is applied to removed set + def __init__(self, varname, remove_dims=None, insert_before=None, + new_dimnames=None, new_coords=None): + self.varname = varname + self.remove_dims = remove_dims + self.insert_before = insert_before + self.new_dimnames = new_dimnames if new_dimnames is not None else [] + self.new_coords = (new_coords if new_coords is not None + else OrderedDict()) + + def __call__(self, var): + new_dims = list(var.dims) + new_coords = OrderedDict(var.coords) + + if self.remove_dims is not None: + for dim in self.remove_dims: + new_dims.remove(dim) + del new_coords[dim] + + if self.insert_before is not None: + insert_idx = new_dims.index(self.insert_before) + new_dims = (new_dims[0:insert_idx] + self.new_dimnames + + new_dims[insert_idx:]) + elif self.new_dimnames is not None: + new_dims = self.new_dimnames + + if self.new_coords is not None: + new_coords.update(self.new_coords) + + return new_dims, new_coords + def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar): lats = wrfnc.variables[latvar] @@ -84,17 +198,28 @@ def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar): return False -def _is_moving_domain(wrfseq, varname=None, latvar="XLAT", lonvar="XLONG"): +def _is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), + lonvar=either("XLONG", "XLONG_M")): + + if isinstance(latvar, either): + latvar = latvar(wrfseq) + + if isinstance(lonvar, either): + lonvar = lonvar(wrfseq) + # In case it's just a single file if not _is_multi_file(wrfseq): wrfseq = [wrfseq] - + # Slow, but safe. Compare the corner points to the first item and see # any move. User iterator protocol in case wrfseq is not a list/tuple. - wrf_iter = iter(wrfseq) - - first_wrfnc = next(wrf_iter) - + if not _is_mapping(wrfseq): + wrf_iter = iter(wrfseq) + first_wrfnc = next(wrf_iter) + else: + entry = wrfseq[next(iter(viewkeys(wrfseq)))] + return _is_moving_domain(entry, varname, latvar, lonvar) + if varname is not None: try: coord_names = getattr(first_wrfnc.variables[varname], @@ -146,7 +271,7 @@ def _get_global_attr(wrfnc, attr): val = getattr(wrfnc, attr, None) # PyNIO puts single values in to an array - if isinstance(val,np.ndarray): + if isinstance(val, np.ndarray): if len(val) == 1: return val[0] return val @@ -161,18 +286,20 @@ def extract_global_attrs(wrfnc, attrs): if multifile: if not _is_mapping(wrfnc): - wrfnc = wrfnc[0] + wrfnc = next(iter(wrfnc)) else: - wrfnc = wrfnc[next(wrfnc.iterkeys())] + entry = wrfnc[next(iter(viewkeys(wrfnc)))] + return extract_global_attrs(entry, 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): - wrfnc = wrfnc[0] + wrfnc = next(iter(wrfnc)) else: - wrfnc = wrfnc[next(wrfnc.iterkeys())] + entry = wrfnc[next(iter(viewkeys(wrfnc)))] + return extract_dim(entry, dim) d = wrfnc.dimensions[dim] if not isinstance(d, int): @@ -189,9 +316,11 @@ def _combine_dict(wrfdict, varname, timeidx, method, nometa): first_key = next(key_iter) keynames.append(first_key) + is_moving = _is_moving_domain(wrfdict, varname) + first_array = _extract_var(wrfdict[first_key], varname, - timeidx, method, squeeze=False, - nometa=nometa) + timeidx, is_moving=is_moving, method=method, + squeeze=False, cache=None, nometa=nometa) # Create the output data numpy array based on the first array @@ -209,12 +338,13 @@ def _combine_dict(wrfdict, varname, timeidx, method, nometa): else: keynames.append(key) vardata = _extract_var(wrfdict[key], varname, timeidx, - method, squeeze=False) + is_moving=is_moving, method=method, + squeeze=False, cache=None, nometa=nometa) if outdata.shape[1:] != vardata.shape: raise ValueError("data sequences must have the " "same size for all dictionary keys") - outdata[idx,:] = vardata.values[:] + outdata[idx,:] = npvalues(vardata)[:] idx += 1 if xarray_enabled() and not nometa: @@ -233,8 +363,10 @@ def _combine_dict(wrfdict, varname, timeidx, method, nometa): return outarr # TODO: implement in C -def _cat_files(wrfseq, varname, timeidx, squeeze, nometa): - is_moving = _is_moving_domain(wrfseq, varname) +# Note: is moving argument needed for dictionary combination +def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, nometa): + if is_moving is None: + is_moving = _is_moving_domain(wrfseq, varname) file_times = extract_times(wrfseq, timeidx) @@ -307,8 +439,9 @@ def _cat_files(wrfseq, varname, timeidx, squeeze, nometa): return outarr # TODO: implement in C -def _join_files(wrfseq, varname, timeidx, nometa): - is_moving = _is_moving_domain(wrfseq, varname) +def _join_files(wrfseq, varname, timeidx, is_moving, nometa): + if is_moving is None: + is_moving = _is_moving_domain(wrfseq, varname) multitime = _is_multi_time_req(timeidx) numfiles = len(wrfseq) @@ -361,15 +494,20 @@ def _join_files(wrfseq, varname, timeidx, nometa): return outarr -def combine_files(wrfseq, varname, timeidx, method="cat", squeeze=True, - nometa=False): +def combine_files(wrfseq, varname, timeidx, is_moving=None, + method="cat", squeeze=True, nometa=False): + + # Handles generators, single files, lists, tuples, custom classes + wrfseq = _unpack_sequence(wrfseq) + # Dictionary is unique if _is_mapping(wrfseq): outarr = _combine_dict(wrfseq, varname, timeidx, method, nometa) elif method.lower() == "cat": - outarr = _cat_files(wrfseq, varname, timeidx, squeeze, nometa) + outarr = _cat_files(wrfseq, varname, timeidx, is_moving, + squeeze, nometa) elif method.lower() == "join": - outarr = _join_files(wrfseq, varname, timeidx, nometa) + outarr = _join_files(wrfseq, varname, timeidx, is_moving, nometa) else: raise ValueError("method must be 'cat' or 'join'") @@ -467,7 +605,8 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): return data_array # Cache is a dictionary of already extracted variables -def _extract_var(wrfnc, varname, timeidx, method, squeeze, cache, nometa): +def _extract_var(wrfnc, varname, timeidx, is_moving, + method, squeeze, cache, nometa): # Mainly used internally so variables don't get extracted multiple times, # particularly to copy metadata. This can be slow. if cache is not None: @@ -481,28 +620,29 @@ def _extract_var(wrfnc, varname, timeidx, method, squeeze, cache, nometa): return cache_var.values return cache_var - - is_moving = _is_moving_domain(wrfnc, varname) multitime = _is_multi_time_req(timeidx) multifile = _is_multi_file(wrfnc) if not multifile: if xarray_enabled() and not nometa: + if is_moving is None: + is_moving = _is_moving_domain(wrfnc, varname) result = _build_data_array(wrfnc, varname, timeidx, is_moving) else: if not multitime: result = wrfnc.variables[varname][timeidx,:] result = result[np.newaxis, :] # So that no squeeze works - else: result = wrfnc.variables[varname][:] else: # Squeeze handled in this routine, so just return it - return combine_files(wrfnc, varname, timeidx, method, squeeze, nometa) + return combine_files(wrfnc, varname, timeidx, is_moving, + method, squeeze, nometa) return result.squeeze() if squeeze else result + def extract_vars(wrfnc, timeidx, varnames, method="cat", squeeze=True, cache=None, nometa=False): if isstr(varnames): @@ -510,7 +650,7 @@ def extract_vars(wrfnc, timeidx, varnames, method="cat", squeeze=True, else: varlist = varnames - return {var:_extract_var(wrfnc, var, timeidx, + return {var:_extract_var(wrfnc, var, timeidx, None, method, squeeze, cache, nometa) for var in varlist} @@ -543,7 +683,12 @@ def extract_times(wrfnc, timeidx): def is_standard_wrf_var(wrfnc, var): multifile = _is_multi_file(wrfnc) if multifile: - wrfnc = wrfnc[0] + if not _is_mapping(wrfnc): + wrfnc = next(iter(wrfnc)) + else: + entry = wrfnc[next(iter(viewkeys(wrfnc)))] + return is_standard_wrf_var(entry, var) + return var in wrfnc.variables @@ -627,104 +772,37 @@ def get_proj_params(wrfnc, timeidx=0, varname=None): wrfnc.variables[lon_coord][time_idx_or_slice,:], proj_params) -# Dictionary python 2-3 compatibility stuff -def viewitems(d): - func = getattr(d, "viewitems", None) - if func is None: - func = d.items - return func() - - -def viewkeys(d): - func = getattr(d, "viewkeys", None) - if func is None: - func = d.keys - return func() - -def viewvalues(d): - func = getattr(d, "viewvalues", None) - if func is None: - func = d.values - return func() - -def isstr(s): - try: - return isinstance(s, basestring) - except NameError: - return isinstance(s, str) - -# Helper to extract masked arrays from DataArrays that convert to NaN -def npvalues(da): - if not isinstance(da, DataArray): - result = da - else: - try: - fill_value = da.attrs["_FillValue"] - except KeyError: - result = da.values - else: - result = ma.masked_invalid(da.values, copy=False) - result.set_fill_value(fill_value) - - return result +class CoordPair(object): + def __init__(self, x=None, y=None, i=None, j=None, lat=None, lon=None): + self.x = x + self.y = y + self.i = i + self.j = j + self.lat = lat + self.lon = lon + + def __repr__(self): + args = [] + if self.x is not None: + args.append("x={}".format(self.x)) + args.append("y={}".format(self.y)) - -# Helper utilities for metadata -class either(object): - def __init__(self, *varnames): - self.varnames = varnames - - def __call__(self, wrfnc): - if _is_multi_file(wrfnc): - wrfnc = next(iter(wrfnc)) + if self.i is not None: + args.append("i={}".format(self.i)) + args.append("j={}".format(self.j)) + + if self.lat is not None: + args.append("lat={}".format(self.lat)) + args.append("lon={}".format(self.lon)) - for varname in self.varnames: - if varname in wrfnc.variables: - return varname + argstr = ", ".join(args) - raise ValueError("{} are not valid variable names".format( - self.varnames)) - -class combine_with: - # Remove remove_idx first, then insert_idx is applied to removed set - def __init__(self, varname, remove_dims=None, insert_before=None, - new_dimnames=None, new_coords=None): - self.varname = varname - self.remove_dims = remove_dims - self.insert_before = insert_before - self.new_dimnames = new_dimnames if new_dimnames is not None else [] - self.new_coords = (new_coords if new_coords is not None - else OrderedDict()) + return "{}({})".format(self.__class__.__name__, argstr) - def __call__(self, var): - new_dims = list(var.dims) - new_coords = OrderedDict(var.coords) - - if self.remove_dims is not None: - for dim in self.remove_dims: - new_dims.remove(dim) - del new_coords[dim] - - if self.insert_before is not None: - insert_idx = new_dims.index(self.insert_before) - new_dims = (new_dims[0:insert_idx] + self.new_dimnames + - new_dims[insert_idx:]) - elif self.new_dimnames is not None: - new_dims = self.new_dimnames - - if self.new_coords is not None: - new_coords.update(self.new_coords) - - return new_dims, new_coords + def __str__(self): + return self.__repr__() -class XYCoord(object): - def __init__(self, x, y): - self.x = x - self.y = y - - def __repr__(self): - return "{}({}, {})".format(self.__class__.__name__, self.x, self.y) def from_args(func, argnames, *args, **kwargs): """Parses the function args and kargs looking for the desired argument diff --git a/wrf_open/var/src/python/wrf/var/uvmet.py b/wrf_open/var/src/python/wrf/var/uvmet.py index a3cf7fa..3c16a6e 100755 --- a/wrf_open/var/src/python/wrf/var/uvmet.py +++ b/wrf_open/var/src/python/wrf/var/uvmet.py @@ -3,6 +3,8 @@ from __future__ import (absolute_import, division, print_function, from math import fabs, log, tan, sin, cos +import numpy as np + from .extension import computeuvmet from .destag import destagger from .constants import Constants @@ -55,8 +57,24 @@ def _get_uvmet(wrfnc, timeidx=0, ten_m=False, units ="mps", # don't appear to be supported any longer if map_proj in (0,3,6): - # No rotation needed for Mercator and Lat/Lon - return u,v + # No rotation needed for Mercator and Lat/Lon, but still need + # u,v aggregated in to one array + + end_idx = -3 if not ten_m else -2 + resdim = list(u.shape[0:end_idx]) + [2] + list(u.shape[end_idx:]) + + # Make a new output array for the result + res = np.empty(resdim, u.dtype) + + # For 2D array, this makes (...,0,:,:) and (...,1,:,:) + # For 3D array, this makes (...,0,:,:,:) and (...,1,:,:,:) + idx0 = tuple([Ellipsis] + [0] + [slice(None)]*(-end_idx)) + idx1 = tuple([Ellipsis] + [1] + [slice(None)]*(-end_idx)) + + res[idx0] = u[:] + res[idx1] = v[:] + + return res elif map_proj in (1,2): lat_attrs = extract_global_attrs(wrfnc, attrs=("TRUELAT1", "TRUELAT2")) diff --git a/wrf_open/var/test/utests.py b/wrf_open/var/test/utests.py index 765706a..e1b8f6d 100644 --- a/wrf_open/var/test/utests.py +++ b/wrf_open/var/test/utests.py @@ -408,7 +408,7 @@ if __name__ == "__main__": "geopt", "helicity", "lat", "lon", "omg", "p", "pressure", "pvo", "pw", "rh2", "rh", "slp", "ter", "td2", "td", "tc", "theta", "tk", "tv", "twb", "updraft_helicity", "ua", "va", - "wa", "uvmet10", "uvmet", "z", "ctt", "cape_2d", "cape_3d"] + "wa", "uvmet10", "uvmet", "z", "ctt"] interp_methods = ["interplevel", "vertcross", "interpline", "vinterp"] try: