diff --git a/wrf_open/var/src/python/wrf/var/constants.py b/wrf_open/var/src/python/wrf/var/constants.py index d463ed7..67bff0b 100755 --- a/wrf_open/var/src/python/wrf/var/constants.py +++ b/wrf_open/var/src/python/wrf/var/constants.py @@ -1,7 +1,9 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) -__all__ = ["Constants", "ConversionFactors"] +__all__ = ["ALL_TIMES", "Constants", "ConversionFactors"] + +ALL_TIMES = None class Constants(object): R = 287.06 diff --git a/wrf_open/var/src/python/wrf/var/interp.py b/wrf_open/var/src/python/wrf/var/interp.py index fc0d6b4..276bfce 100755 --- a/wrf_open/var/src/python/wrf/var/interp.py +++ b/wrf_open/var/src/python/wrf/var/interp.py @@ -226,7 +226,13 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, missing = field.fill_value else: missing = Constants.DEFAULT_FILL - + + if (field.shape != p.shape): + raise ValueError("'field' shape does not match other variable shapes. " + "Verify that the 'timeidx' parameter matches the " + "same value used when extracting the 'field' " + "variable.") + res = vintrp(field, p, tk, qv, ght, terht, sfp, smsfp, vcord_array, interp_levels, icase, extrap, vcor, log_p_int, missing) diff --git a/wrf_open/var/src/python/wrf/var/projection.py b/wrf_open/var/src/python/wrf/var/projection.py index 7e15413..b598015 100644 --- a/wrf_open/var/src/python/wrf/var/projection.py +++ b/wrf_open/var/src/python/wrf/var/projection.py @@ -15,7 +15,7 @@ if basemap_enabled(): if pyngl_enabled(): from Ngl import Resources -__all__ = ["WrfProj", "LambertConformal", "Mercator", +__all__ = ["WrfProj", "NullProjection", "LambertConformal", "Mercator", "PolarStereographic", "LatLon", "RotatedLatLon", "getproj"] @@ -180,6 +180,15 @@ class WrfProj(object): projection""" return self._cf_params + +# Used for 'missing' projection values during the 'join' method +class NullProjection(WrfProj): + def __init__(self): + pass + + def __repr__(self): + return "{}()".format(self.__class__.__name__) + class LambertConformal(WrfProj): def __init__(self, bottom_left=None, top_right=None, diff --git a/wrf_open/var/src/python/wrf/var/util.py b/wrf_open/var/src/python/wrf/var/util.py index 58e7dbb..3e2c3b3 100644 --- a/wrf_open/var/src/python/wrf/var/util.py +++ b/wrf_open/var/src/python/wrf/var/util.py @@ -1,29 +1,40 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) +from sys import version_info from copy import copy from collections import Iterable, Mapping, OrderedDict from itertools import product -from inspect import getargspec, getargvalues, getmodule from types import GeneratorType import datetime as dt +from math import floor, copysign +from inspect import getmodule +try: + from inspect import signature +except ImportError: + from inspect import getargspec, getargvalues + + import numpy as np import numpy.ma as ma from .config import xarray_enabled -from .projection import getproj +from .projection import getproj, NullProjection +from .constants import Constants, ALL_TIMES if xarray_enabled(): from xarray import DataArray + from pandas import NaT __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", "either", "from_args", "arg_location", - "args_to_list", "npvalues", "CoordPair"] + "viewvalues", "py2round", "combine_with", "either", + "from_args", "arg_location", "args_to_list", "npvalues", + "CoordPair"] _COORD_PAIR_MAP = {"XLAT" : ("XLAT", "XLONG"), @@ -39,23 +50,31 @@ _COORD_PAIR_MAP = {"XLAT" : ("XLAT", "XLONG"), _COORD_VARS = ("XLAT", "XLONG", "XLAT_M", "XLONG_M", "XLAT_U", "XLONG_U", - "XLAT_V", "XLONG_V", "CLAT", "CLONG") + "XLAT_V", "XLONG_V", "CLAT", "CLONG") + +_TIME_COORD_VARS = ("XTIME",) def _is_coord_var(varname): return varname in _COORD_VARS +def _is_time_coord_var(varname): + return varname in _TIME_COORD_VARS + + def _get_coord_pairs(varname): return _COORD_PAIR_MAP[varname] def _is_multi_time_req(timeidx): - return timeidx == -1 + return timeidx is None 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) @@ -157,6 +176,18 @@ def isstr(s): return isinstance(s, basestring) except NameError: return isinstance(s, str) + +# Python 2 rounding behavior +def _round2(x, d=0): + p = 10 ** d + return float(floor((x * p) + copysign(0.5, x)))/p + +def py2round(x, d=0): + if version_info >= (3,): + return _round2(x, d) + + return round(x, d) + # Helper to extract masked arrays from DataArrays that convert to NaN def npvalues(da): @@ -193,7 +224,7 @@ class either(object): raise ValueError("{} are not valid variable names".format( self.varnames)) -class combine_with: +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, new_dimnames=None, new_coords=None): @@ -248,6 +279,7 @@ def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar): return False + def _is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), lonvar=either("XLONG", "XLONG_M")): @@ -269,7 +301,21 @@ def _is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), else: entry = wrfseq[next(iter(viewkeys(wrfseq)))] return _is_moving_domain(entry, varname, latvar, lonvar) - + + # Quick check on pressure coordinates, bypassing the need to search the + # domain corner points + try: + coord_names = getattr(first_wrfnc.variables["P"], + "coordinates").split() + except KeyError: + pass + else: + if "XTIME" in coord_names: + return True + else: + return False + + # The long way of checking all lat/lon corner points if varname is not None: try: coord_names = getattr(first_wrfnc.variables[varname], @@ -280,16 +326,31 @@ def _is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), lon_coord = lonvar lat_coord = latvar else: + # If the XTIME variable is found to be a coordinate variable, + # then it's a moving domain file + try: + xtime_coord = coord_names[2] + except IndexError: + # XTIME is not a coordinate variable, if the variable is in the + # file, then this is not a moving domain file + if "XTIME" in first_wrfnc.variables: + return False + + else: + # XTIME is a coordinate, so this is a moving domain file + if xtime_coord == "XTIME": + return True + lon_coord = coord_names[0] lat_coord = coord_names[1] else: lon_coord = lonvar lat_coord = latvar + # Probably a met_em file, need to search all the files lats = first_wrfnc.variables[lat_coord] lons = first_wrfnc.variables[lon_coord] - zero_idxs = [0]*len(lats.shape) # PyNIO doesn't have ndim last_idxs = list(zero_idxs) last_idxs[-2:] = [-1]*2 @@ -412,28 +473,242 @@ def _combine_dict(wrfdict, varname, timeidx, method, meta): return outarr +def _find_coord_names(coords): + try: + lat_coord = [name for name in _COORD_VARS[0::2] if name in coords][0] + except IndexError: + lat_coord = None + + try: + lon_coord = [name for name in _COORD_VARS[1::2] if name in coords][0] + except IndexError: + lon_coord = None + + try: + xtime_coord = [name for name in _TIME_COORD_VARS if name in coords][0] + except IndexError: + xtime_coord = None + + return lat_coord, lon_coord, xtime_coord + + +def _find_max_time_size(wrfseq): + wrf_iter = iter(wrfseq) + + max_times = 0 + while True: + try: + wrfnc = next(wrf_iter) + except StopIteration: + break + else: + t = extract_dim(wrfnc, "Time") + max_times = t if t >= max_times else max_times + + return max_times + + +def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): + multitime = _is_multi_time_req(timeidx) + time_idx_or_slice = timeidx if not multitime else slice(None) + var = wrfnc.variables[varname] + data = var[time_idx_or_slice, :] + + # Want to preserve the time dimension + if not multitime: + data = data[np.newaxis, :] + + attrs = OrderedDict(var.__dict__) + dimnames = var.dimensions[-data.ndim:] + + # WRF variables will have a coordinates attribute. MET_EM files have + # a stagger attribute which indicates the coordinate variable. + try: + # WRF files + coord_attr = getattr(var, "coordinates") + except AttributeError: + if _is_coord_var(varname): + # Coordinate variable (most likely XLAT or XLONG) + lat_coord, lon_coord = _get_coord_pairs(varname) + time_coord = None + + if is_moving_domain and _has_time_coord(wrfnc): + time_coord = "XTIME" + + else: + try: + # met_em files + stag_attr = getattr(var, "stagger") + except AttributeError: + lon_coord = None + lat_coord = None + else: + # For met_em files, use the stagger name to get the lat/lon var + lat_coord = "XLAT_{}".format(stag_attr) + lon_coord = "XLONG_{}".format(stag_attr) + else: + coord_names = coord_attr.split() + lon_coord = coord_names[0] + lat_coord = coord_names[1] + + try: + time_coord = coord_names[2] + except IndexError: + time_coord = None + + coords = OrderedDict() + + # Handle lat/lon coordinates and projection information if available + if lon_coord is not None and lat_coord is not None: + lon_coord_var = wrfnc.variables[lon_coord] + lat_coord_var = wrfnc.variables[lat_coord] + time_coord_var = (wrfnc.variables[time_coord] + if time_coord is not None + else None) + + if multitime: + if is_moving_domain: + # Special case with a moving domain in a multi-time file, + # otherwise the projection parameters don't change + coords[lon_coord] = lon_coord_var.dimensions, lon_coord_var[:] + coords[lat_coord] = lat_coord_var.dimensions, lat_coord_var[:] + + # Returned lats/lons arrays will have a time dimension, so proj + # will need to be a list due to moving corner points + lats, lons, proj_params = get_proj_params(wrfnc, + timeidx, + varname) + proj = [getproj(lats=lats[i,:], + lons=lons[i,:], + **proj_params) for i in xrange(lats.shape[0])] + + if time_coord is not None: + coords[time_coord] = (lon_coord_var.dimensions[0], + time_coord_var[:]) + else: + coords[lon_coord] = (lon_coord_var.dimensions[1:], + lon_coord_var[0,:]) + coords[lat_coord] = (lat_coord_var.dimensions[1:], + lat_coord_var[0,:]) + + # Domain not moving, so just get the first time + lats, lons, proj_params = get_proj_params(wrfnc, 0, varname) + proj = getproj(lats=lats, lons=lons, **proj_params) + else: + coords[lon_coord] = (lon_coord_var.dimensions[1:], + lon_coord_var[timeidx,:]) + coords[lat_coord] = (lat_coord_var.dimensions[1:], + lat_coord_var[timeidx,:]) + lats, lons, proj_params = get_proj_params(wrfnc, 0, varname) + proj = getproj(lats=lats, lons=lons, **proj_params) + + attrs["projection"] = proj + + + if dimnames[0] == "Time": + coords[dimnames[0]] = extract_times(wrfnc, timeidx) + + data_array = DataArray(data, name=varname, dims=dimnames, coords=coords, + attrs=attrs) + + + return data_array + + +def _find_forward(wrfseq, varname, timeidx, is_moving, meta): + + wrf_iter = iter(wrfseq) + comboidx = 0 + + while True: + try: + wrfnc = next(wrf_iter) + except StopIteration: + break + else: + numtimes = extract_dim(wrfnc, "Time") + + if timeidx < comboidx + numtimes: + filetimeidx = timeidx - comboidx + + if meta: + return _build_data_array(wrfnc, varname, filetimeidx, + is_moving) + else: + return wrfnc.variables[varname][filetimeidx, :] + else: + comboidx += numtimes + + raise IndexError("invalid time index") + + +def _find_reverse(wrfseq, varname, timeidx, is_moving, meta): + try: + revwrfseq = reversed(wrfseq) + except TypeError: + revwrfseq = reversed(list(wrfseq)) + + wrf_iter = iter(revwrfseq) + revtimeidx = -timeidx - 1 + + comboidx = 0 + + while True: + try: + wrfnc = next(wrf_iter) + except StopIteration: + break + else: + numtimes = extract_dim(wrfnc, "Time") + + if revtimeidx < comboidx + numtimes: + # Finds the "forward" sequence index, then counts that + # number back from the back of the ncfile times, + # since the ncfile needs to be iterated backwards as well. + filetimeidx = numtimes - (revtimeidx - comboidx) - 1 + + if meta: + return _build_data_array(wrfnc, varname, filetimeidx, + is_moving) + else: + return wrfnc.variables[varname][filetimeidx, :] + else: + comboidx += numtimes + + raise IndexError("invalid time index") + + +def _find_arr_for_time(wrfseq, varname, timeidx, is_moving, meta): + if timeidx >= 0: + return _find_forward(wrfseq, varname, timeidx, is_moving, meta) + else: + return _find_reverse(wrfseq, varname, timeidx, is_moving, meta) + # TODO: implement in C -# Note: is moving argument needed for dictionary combination def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta): if is_moving is None: is_moving = _is_moving_domain(wrfseq, varname) - file_times = extract_times(wrfseq, timeidx) + file_times = extract_times(wrfseq, ALL_TIMES) multitime = _is_multi_time_req(timeidx) + + # For single times, just need to find the ncfile and appropriate + # time index, and return that array + if not multitime: + return _find_arr_for_time(wrfseq, varname, timeidx, is_moving, meta) - time_idx_or_slice = timeidx if not multitime else slice(None) + #time_idx_or_slice = timeidx if not multitime else slice(None) - # wrfseq might be a generator + # If all times are requested, need to build a new array and cat together + # all of the arrays in the sequence wrf_iter = iter(wrfseq) if xarray_enabled() and meta: first_var = _build_data_array(next(wrf_iter), varname, - timeidx, is_moving) + ALL_TIMES, is_moving) else: - first_var = (next(wrf_iter)).variables[varname][time_idx_or_slice, :] - if not multitime: - first_var = first_var[np.newaxis,:] + first_var = (next(wrf_iter)).variables[varname][:] outdims = [len(file_times)] @@ -448,6 +723,28 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta): outdata[startidx:endidx, :] = first_var[:] + if xarray_enabled() and meta and is_moving: + latname, lonname, timename = _find_coord_names(first_var.coords) + outcoorddims = outdims[0:1] + outdims[-2:] + + if latname is not None: + outlats = np.empty(outcoorddims, first_var.dtype) + outlats[startidx:endidx, :] = first_var.coords[latname][:] + + if lonname is not None: + outlons = np.empty(outcoorddims, first_var.dtype) + outlons[startidx:endidx, :] = first_var.coords[lonname][:] + + if timename is not None: + outxtimes = np.empty(outdims[0]) + outxtimes[startidx:endidx] = first_var.coords[timename][:] + + # Projections also need to be aggregated + outprojs = np.empty(outdims[0], np.object) + + outprojs[startidx:endidx] = np.asarray(first_var.attrs["projection"], + np.object)[:] + startidx = endidx while True: try: @@ -455,91 +752,265 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta): except StopIteration: break else: - vardata = wrfnc.variables[varname][time_idx_or_slice, :] + vardata = wrfnc.variables[varname][:] - if multitime: - numtimes = vardata.shape[0] - else: - numtimes = 1 + numtimes = vardata.shape[0] endidx = startidx + numtimes outdata[startidx:endidx, :] = vardata[:] + if xarray_enabled() and meta and is_moving: + if latname is not None: + latdata = wrfnc.variables[latname][:] + outlats[startidx:endidx, :] = latdata[:] + + if lonname is not None: + londata = wrfnc.variables[lonname][:] + outlons[startidx:endidx, :] = londata[:] + + if timename is not None: + xtimedata = wrfnc.variables[timename][:] + outxtimes[startidx:endidx] = xtimedata[:] + + lats, lons, proj_params = get_proj_params(wrfnc, + ALL_TIMES, + varname) + projs = [getproj(lats=lats[i,:], + lons=lons[i,:], + **proj_params) for i in xrange(lats.shape[0])] + + outprojs[startidx:endidx] = np.asarray(projs, np.object)[:] + startidx = endidx if xarray_enabled() and meta: - # FIXME: If it's a moving nest, then the coord arrays need to have same - # time indexes as the whole data set outname = unicode(first_var.name) outattrs = OrderedDict(first_var.attrs) outcoords = OrderedDict(first_var.coords) outdimnames = list(first_var.dims) + if "Time" not in outdimnames: outdimnames.insert(0, "Time") + + if not multitime: + file_times = [file_times[timeidx]] + + outcoords[outdimnames[0]] = file_times + + outcoords["datetime"] = outdimnames[0], file_times + + # If the domain is moving, need to create the lat/lon/xtime coords + # since they can't be copied + if is_moving: + outlatdims = [outdimnames[0]] + outdimnames[-2:] - outcoords[outdimnames[0]] = file_times # New time dimension values + if latname is not None: + outlats = outlats[:] + outcoords[latname] = outlatdims, outlats + if lonname is not None: + outlons = outlons[:] + outcoords[lonname] = outlatdims, outlons + if timename is not None: + outxtimes = outxtimes[:] + outcoords[timename] = outdimnames[0], outxtimes + + outattrs["projection"] = outprojs[:] + outdata = outdata[:] + outarr = DataArray(outdata, name=outname, coords=outcoords, dims=outdimnames, attrs=outattrs) else: + outdata = outdata[:] outarr = outdata return outarr +def _get_numfiles(wrfseq): + try: + return len(wrfseq) + except TypeError: + wrf_iter = iter(wrfseq) + return sum(1 for _ in wrf_iter) + # TODO: implement in C def _join_files(wrfseq, varname, timeidx, is_moving, meta): if is_moving is None: is_moving = _is_moving_domain(wrfseq, varname) multitime = _is_multi_time_req(timeidx) - numfiles = len(wrfseq) + numfiles = _get_numfiles(wrfseq) + maxtimes = _find_max_time_size(wrfseq) time_idx_or_slice = timeidx if not multitime else slice(None) + file_times_less_than_max = False + file_idx = 0 # wrfseq might be a generator wrf_iter = iter(wrfseq) - + wrfnc = next(wrf_iter) + numtimes = extract_dim(wrfnc, "Time") + if xarray_enabled() and meta: - first_var = _build_data_array(next(wrf_iter), varname, - timeidx, is_moving) + first_var = _build_data_array(wrfnc, varname, ALL_TIMES, is_moving) + time_coord = np.full((numfiles, maxtimes), int(NaT), "datetime64[ns]") + time_coord[file_idx, 0:numtimes] = first_var.coords["Time"][:] else: - first_var = (next(wrf_iter)).variables[varname][time_idx_or_slice, :] - if not multitime: - first_var = first_var[np.newaxis, :] + first_var = wrfnc.variables[varname][:] - # Create the output data numpy array based on the first array + if numtimes < maxtimes: + file_times_less_than_max = True + + # Out dimensions will be the number of files, maxtimes, then the + # non-time shapes from the first variable outdims = [numfiles] - outdims += first_var.shape - outdata = np.empty(outdims, first_var.dtype) + outdims += [maxtimes] + outdims += first_var.shape[1:] + + # For join, always need to start with full masked values + outdata = np.full(outdims, Constants.DEFAULT_FILL, first_var.dtype) + outdata[file_idx, 0:numtimes, :] = first_var[:] - outdata[0,:] = first_var[:] + # Create the secondary coordinate arrays + if xarray_enabled() and meta and is_moving: + latname, lonname, timename = _find_coord_names(first_var.coords) + outcoorddims = outdims[0:2] + outdims[-2:] + + if latname is not None: + outlats = np.full(outcoorddims, Constants.DEFAULT_FILL, + first_var.dtype) + outlats[file_idx, 0:numtimes, :] = first_var.coords[latname][:] + + if lonname is not None: + outlons = np.full(outcoorddims, Constants.DEFAULT_FILL, + first_var.dtype) + outlons[file_idx, 0:numtimes, :] = first_var.coords[lonname][:] + + if timename is not None: + outxtimes = np.full(outdims[0:2], Constants.DEFAULT_FILL, + first_var.dtype) + outxtimes[file_idx, 0:numtimes] = first_var.coords[timename][:] + + # Projections also need two dimensions + outprojs = np.full(outdims[0:2], NullProjection(), np.object) + + outprojs[file_idx, 0:numtimes] = np.asarray( + first_var.attrs["projection"], + np.object)[:] - idx=1 + file_idx=1 while True: try: wrfnc = next(wrf_iter) except StopIteration: break else: - outvar = wrfnc.variables[varname][time_idx_or_slice, :] + numtimes = extract_dim(wrfnc, "Time") + if numtimes < maxtimes: + file_times_less_than_max = True + outvar = wrfnc.variables[varname][:] + if not multitime: outvar = outvar[np.newaxis, :] - outdata[idx,:] = outvar[:] - idx += 1 - + + outdata[file_idx, 0:numtimes, :] = outvar[:] + + if xarray_enabled() and meta: + file_times = extract_times(wrfnc, ALL_TIMES) + time_coord[file_idx, 0:numtimes] = np.asarray(file_times, + "datetime64[ns]")[:] + + if xarray_enabled() and meta and is_moving: + if latname is not None: + latdata = wrfnc.variables[latname][:] + outlats[file_idx, 0:numtimes, :] = latdata[:] + + if lonname is not None: + londata = wrfnc.variables[lonname][:] + outlons[file_idx, 0:numtimes, :] = londata[:] + + if timename is not None: + xtimedata = wrfnc.variables[timename][:] + outxtimes[file_idx, 0:numtimes] = xtimedata[:] + + lats, lons, proj_params = get_proj_params(wrfnc, + ALL_TIMES, + varname) + projs = [getproj(lats=lats[i,:], + lons=lons[i,:], + **proj_params) for i in xrange(lats.shape[0])] + + outprojs[file_idx, 0:numtimes] = ( + np.asarray(projs, np.object)[:]) + + # Need to update coords here + file_idx += 1 + + # If any of the output files contain less than the max number of times, + # then a mask array is needed to flag all the missing arrays with + # missing values + if file_times_less_than_max: + outdata = np.ma.masked_values(outdata, Constants.DEFAULT_FILL) + if xarray_enabled() and meta: outname = unicode(first_var.name) outcoords = OrderedDict(first_var.coords) outattrs = OrderedDict(first_var.attrs) # New dimensions - outdimnames = ["file_idx"] + list(first_var.dims) - outcoords["file_idx"] = [i for i in xrange(numfiles)] + outdimnames = ["file"] + list(first_var.dims) + outcoords["file"] = [i for i in xrange(numfiles)] + + # Time needs to be multi dimensional, so use the default dimension + del outcoords["Time"] + time_coord = time_coord[:, time_idx_or_slice] + if not multitime: + time_coord = time_coord[:, np.newaxis] + outcoords["datetime"] = outdimnames[0:2], time_coord + + if isinstance(outdata, np.ma.MaskedArray): + outattrs["_FillValue"] = Constants.DEFAULT_FILL + outattrs["missing_value"] = Constants.DEFAULT_FILL + + # If the domain is moving, need to create the lat/lon/xtime coords + # since they can't be copied + if is_moving: + outlatdims = outdimnames[0:2] + outdimnames[-2:] + + if latname is not None: + outlats = outlats[:, time_idx_or_slice, :] + if not multitime: + outlats = outlats[:, np.newaxis, :] + outcoords[latname] = outlatdims, outlats + if lonname is not None: + outlons = outlons[:, time_idx_or_slice, :] + if not multitime: + outlons = outlons[:, np.newaxis, :] + outcoords[lonname] = outlatdims, outlons + if timename is not None: + outxtimes = outxtimes[:, time_idx_or_slice] + if not multitime: + outxtimes = outxtimes[:, np.newaxis] + outcoords[timename] = outdimnames[0:2], outxtimes[:] + + if not multitime: + outattrs["projection"] = outprojs[:, timeidx] + else: + outattrs["projection"] = outprojs + + if not multitime: + outdata = outdata[:, timeidx, :] + outdata = outdata[:, np.newaxis, :] + outarr = DataArray(outdata, name=outname, coords=outcoords, dims=outdimnames, attrs=outattrs) else: + if not multitime: + outdata = outdata[:, timeidx, :] + outdata = outdata[:, np.newaxis, :] + outarr = outdata return outarr @@ -564,96 +1035,6 @@ def combine_files(wrfseq, varname, timeidx, is_moving=None, return outarr.squeeze() if squeeze else outarr -def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): - multitime = _is_multi_time_req(timeidx) - time_idx_or_slice = timeidx if not multitime else slice(None) - var = wrfnc.variables[varname] - data = var[time_idx_or_slice, :] - - # Want to preserve the time dimension - if not multitime: - data = data[np.newaxis, :] - - attrs = OrderedDict(var.__dict__) - dimnames = var.dimensions[-data.ndim:] - - # WRF variables will have a coordinates attribute. MET_EM files have - # a stagger attribute which indicates the coordinate variable. - try: - # WRF files - coord_attr = getattr(var, "coordinates") - except AttributeError: - if _is_coord_var(varname): - # Coordinate variable (most likely XLAT or XLONG) - lat_coord, lon_coord = _get_coord_pairs(varname) - else: - try: - # met_em files - stag_attr = getattr(var, "stagger") - except AttributeError: - lon_coord = None - lat_coord = None - else: - # For met_em files, use the stagger name to get the lat/lon var - lat_coord = "XLAT_{}".format(stag_attr) - lon_coord = "XLONG_{}".format(stag_attr) - else: - coord_names = coord_attr.split() - lon_coord = coord_names[0] - lat_coord = coord_names[1] - - coords = OrderedDict() - - # Handle lat/lon coordinates and projection information if available - if lon_coord is not None and lat_coord is not None: - lon_coord_var = wrfnc.variables[lon_coord] - lat_coord_var = wrfnc.variables[lat_coord] - - if multitime: - if is_moving_domain: - # Special case with a moving domain in a multi-time file, - # otherwise the projection parameters don't change - coords[lon_coord] = lon_coord_var.dimensions, lon_coord_var[:] - coords[lat_coord] = lat_coord_var.dimensions, lat_coord_var[:] - - # Returned lats/lons arrays will have a time dimension, so proj - # will need to be a list due to moving corner points - lats, lons, proj_params = get_proj_params(wrfnc, - timeidx, - varname) - proj = [getproj(lats=lats[i,:], - lons=lons[i,:], - **proj_params) for i in xrange(lats.shape[0])] - else: - coords[lon_coord] = (lon_coord_var.dimensions[1:], - lon_coord_var[0,:]) - coords[lat_coord] = (lat_coord_var.dimensions[1:], - lat_coord_var[0,:]) - - # Domain not moving, so just get the first time - lats, lons, proj_params = get_proj_params(wrfnc, 0, varname) - proj = getproj(lats=lats, lons=lons, **proj_params) - else: - coords[lon_coord] = (lon_coord_var.dimensions[1:], - lon_coord_var[timeidx,:]) - coords[lat_coord] = (lat_coord_var.dimensions[1:], - lat_coord_var[timeidx,:]) - lats, lons, proj_params = get_proj_params(wrfnc, 0, varname) - proj = getproj(lats=lats, lons=lons, **proj_params) - - attrs["projection"] = proj - - - if dimnames[0] == "Time": - coords[dimnames[0]] = extract_times(wrfnc, timeidx) - - - data_array = DataArray(data, name=varname, dims=dimnames, coords=coords, - attrs=attrs) - - - return data_array - # Cache is a dictionary of already extracted variables def _extract_var(wrfnc, varname, timeidx, is_moving, method, squeeze, cache, meta): @@ -727,7 +1108,7 @@ def extract_times(wrfnc, timeidx): return [file_time for wrf_file in wrf_list - for file_time in _file_times(wrf_file,timeidx) ] + for file_time in _file_times(wrf_file, timeidx)] def is_standard_wrf_var(wrfnc, var): @@ -853,7 +1234,7 @@ class CoordPair(object): def __str__(self): return self.__repr__() - + def from_args(func, argnames, *args, **kwargs): """Parses the function args and kargs looking for the desired argument value. Otherwise, the value is taken from the default keyword argument @@ -876,8 +1257,7 @@ def from_args(func, argnames, *args, **kwargs): return result -def args_to_list(func, args, kwargs): - """Converts the mixed args/kwargs to a single list of args""" +def _args_to_list2(func, args, kwargs): argspec = getargspec(func) # Build the full tuple with defaults filled in @@ -896,15 +1276,28 @@ def args_to_list(func, args, kwargs): return outargs -def arg_location(func, argname, args, kwargs): - """Parses the function args, kargs and signature looking for the desired - argument location (either in args, kargs, or argspec.defaults), - and returns a list containing representing all arguments in the - correct order with defaults filled in. +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] + + +def args_to_list(func, args, kwargs): + """Converts the mixed args/kwargs to a single list of args""" + if version_info > (3,): + _args_to_list = _args_to_list3 + else: + _args_to_list = _args_to_list2 + + return _args_to_list(func, args, kwargs) + + +def _arg_location2(func, argname, args, kwargs): argspec = getargspec(func) - list_args = args_to_list(func, args, kwargs) + + list_args = _args_to_list2(func, args, kwargs) # Return the new sequence and location if argname not in argspec.args and argname not in kwargs: @@ -913,6 +1306,30 @@ def arg_location(func, argname, args, kwargs): result_idx = argspec.args.index(argname) return list_args, result_idx + +def _arg_location3(func, argname, args, kwargs): + sig = signature(func) + params = list(sig.params.keys()) + list_args = _args_to_list3(func, args, kwargs) + result_idx = params.index(argname) + + return list_args, result_idx + + +def arg_location(func, argname, args, kwargs): + """Parses the function args, kargs and signature looking for the desired + argument location (either in args, kargs, or argspec.defaults), + and returns a list containing representing all arguments in the + correct order with defaults filled in. + + """ + if version_info > (3,): + _arg_location = _arg_location3 + else: + _arg_location = _arg_location2 + + return _arg_location(func, argname, args, kwargs) + diff --git a/wrf_open/var/test/snippet.py b/wrf_open/var/test/snippet.py new file mode 100644 index 0000000..014810f --- /dev/null +++ b/wrf_open/var/test/snippet.py @@ -0,0 +1,30 @@ + +import matplotlib.pyplot as plt +from mpl_toolkits.basemap import Basemap + +def main(): + bm = Basemap(projection = "rotpole", + o_lat_p = 36.0, + o_lon_p = 180.0, + llcrnrlat = -10.590603, + urcrnrlat = 46.591976, + llcrnrlon = -139.08585, + urcrnrlon = 22.661009, + lon_0 = -106.0, + rsphere = 6370000, + resolution = 'l') + + fig = plt.figure(figsize=(8,8)) + ax = fig.add_axes([0.1,0.1,0.8,0.8]) + + bm.drawcoastlines(linewidth=.5) + + print bm.proj4string + + plt.savefig("basemap_map.png") + plt.close(fig) + + + +if __name__ == "__main__": + main() diff --git a/wrf_open/var/test/utests.py b/wrf_open/var/test/utests.py index e1b8f6d..28a6eac 100644 --- a/wrf_open/var/test/utests.py +++ b/wrf_open/var/test/utests.py @@ -12,7 +12,6 @@ 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" OUT_NC_FILE = "/tmp/wrftest.nc" - def setUpModule(): ncarg_root = os.environ.get("NCARG_ROOT", None) if ncarg_root is None: @@ -55,6 +54,7 @@ def make_test(varname, wrf_in, referent, multi=False, repeat=3, pynio=False): pass if not multi: + timeidx = 0 if not pynio: in_wrfnc = NetCDF(wrf_in) else: @@ -66,6 +66,7 @@ def make_test(varname, wrf_in, referent, multi=False, repeat=3, pynio=False): wrf_file = wrf_in in_wrfnc = Nio.open_file(wrf_file) else: + timeidx = None if not pynio: nc = NetCDF(wrf_in) in_wrfnc = [nc for i in xrange(repeat)] @@ -100,28 +101,29 @@ def make_test(varname, wrf_in, referent, multi=False, repeat=3, pynio=False): ref_vals.mask[i,:] = data.mask[:] if (varname == "tc"): - my_vals = getvar(in_wrfnc, "temp", units="c") + my_vals = getvar(in_wrfnc, "temp", timeidx=timeidx, units="c") tol = 0 atol = .1 # Note: NCL uses 273.16 as conversion for some reason + print my_vals.shape, ref_vals.shape nt.assert_allclose(npvalues(my_vals), ref_vals, tol, atol) elif (varname == "pw"): - my_vals = getvar(in_wrfnc, "pw") + my_vals = getvar(in_wrfnc, "pw", timeidx=timeidx) tol = .5/100.0 atol = 0 # NCL uses different constants and doesn't use same # handrolled virtual temp in method nt.assert_allclose(npvalues(my_vals), ref_vals, tol, atol) elif (varname == "cape_2d"): - cape_2d = getvar(in_wrfnc, varname) + cape_2d = getvar(in_wrfnc, varname, timeidx=timeidx) tol = 0/100. # Not sure why different, F77 vs F90? atol = .2 nt.assert_allclose(npvalues(cape_2d), ref_vals, tol, atol) elif (varname == "cape_3d"): - cape_3d = getvar(in_wrfnc, varname) + cape_3d = getvar(in_wrfnc, varname, timeidx=timeidx) tol = 0/100. # Not sure why different, F77 vs F90? atol = .01 nt.assert_allclose(npvalues(cape_3d), ref_vals, tol, atol) else: - my_vals = getvar(in_wrfnc, varname) + my_vals = getvar(in_wrfnc, varname, timeidx=timeidx) tol = 0/100. atol = 0.0001 nt.assert_allclose(npvalues(my_vals), ref_vals, tol, atol) @@ -173,6 +175,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, pass if not multi: + timeidx = 0 if not pynio: in_wrfnc = NetCDF(wrf_in) else: @@ -184,6 +187,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, wrf_file = wrf_in in_wrfnc = Nio.open_file(wrf_file) else: + timeidx = None if not pynio: nc = NetCDF(wrf_in) in_wrfnc = [nc for i in xrange(repeat)] @@ -197,8 +201,8 @@ def make_interp_test(varname, wrf_in, referent, multi=False, if (varname == "interplevel"): ref_ht_500 = _get_refvals(referent, "z_500", repeat, multi) - hts = getvar(in_wrfnc, "z") - p = getvar(in_wrfnc, "pressure") + hts = getvar(in_wrfnc, "z", timeidx=timeidx) + p = getvar(in_wrfnc, "pressure", timeidx=timeidx) hts_500 = interplevel(hts, p, 500) nt.assert_allclose(npvalues(hts_500), ref_ht_500) @@ -207,8 +211,8 @@ def make_interp_test(varname, wrf_in, referent, multi=False, ref_ht_cross = _get_refvals(referent, "ht_cross", repeat, multi) ref_p_cross = _get_refvals(referent, "p_cross", repeat, multi) - hts = getvar(in_wrfnc, "z") - p = getvar(in_wrfnc, "pressure") + hts = getvar(in_wrfnc, "z", timeidx=timeidx) + p = getvar(in_wrfnc, "pressure", timeidx=timeidx) pivot_point = (hts.shape[-2] / 2, hts.shape[-1] / 2) ht_cross = vertcross(hts, p, pivot_point=pivot_point, angle=90.) @@ -235,7 +239,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, ref_t2_line = _get_refvals(referent, "t2_line", repeat, multi) - t2 = getvar(in_wrfnc, "T2") + t2 = getvar(in_wrfnc, "T2", timeidx=timeidx) pivot_point = (t2.shape[-2] / 2, t2.shape[-1] / 2) t2_line1 = interpline(t2, pivot_point=pivot_point, angle=90.0) @@ -255,7 +259,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, fld_tk_theta = _get_refvals(referent, "fld_tk_theta", repeat, multi) fld_tk_theta = n.squeeze(fld_tk_theta) - tk = getvar(in_wrfnc, "temp", units="k") + tk = getvar(in_wrfnc, "temp", timeidx=timeidx, units="k") interp_levels = [200,300,500,1000] @@ -264,7 +268,8 @@ def make_interp_test(varname, wrf_in, referent, multi=False, vert_coord="theta", interp_levels=interp_levels, extrapolate=True, - field_type="tk", + field_type="tk", + timeidx=timeidx, log_p=True) tol = 0/100. @@ -285,6 +290,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, interp_levels=interp_levels, extrapolate=True, field_type="tk", + timeidx=timeidx, log_p=True) tol = 0/100. @@ -304,7 +310,8 @@ def make_interp_test(varname, wrf_in, referent, multi=False, vert_coord="pressure", interp_levels=interp_levels, extrapolate=True, - field_type="tk", + field_type="tk", + timeidx=timeidx, log_p=True) field = n.squeeze(field) @@ -320,7 +327,8 @@ def make_interp_test(varname, wrf_in, referent, multi=False, vert_coord="ght_msl", interp_levels=interp_levels, extrapolate=True, - field_type="tk", + field_type="tk", + timeidx=timeidx, log_p=True) field = n.squeeze(field) @@ -337,6 +345,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, interp_levels=interp_levels, extrapolate=True, field_type="tk", + timeidx=timeidx, log_p=True) field = n.squeeze(field) @@ -346,7 +355,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, fld_ht_pres = _get_refvals(referent, "fld_ht_pres", repeat, multi) fld_ht_pres = n.squeeze(fld_ht_pres) - z = getvar(in_wrfnc, "height", units="m") + z = getvar(in_wrfnc, "height", timeidx=timeidx, units="m") interp_levels = [500,50] field = vinterp(in_wrfnc, field=z, @@ -354,6 +363,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, interp_levels=interp_levels, extrapolate=True, field_type="ght", + timeidx=timeidx, log_p=True) field = n.squeeze(field) @@ -363,7 +373,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, fld_pres_theta = _get_refvals(referent, "fld_pres_theta", repeat, multi) fld_pres_theta = n.squeeze(fld_pres_theta) - p = getvar(in_wrfnc, "pressure") + p = getvar(in_wrfnc, "pressure", timeidx=timeidx) interp_levels = [200,300,500,1000] field = vinterp(in_wrfnc, field=p, @@ -371,6 +381,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, interp_levels=interp_levels, extrapolate=True, field_type="pressure", + timeidx=timeidx, log_p=True) field = n.squeeze(field) @@ -380,7 +391,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, fld_thetae_pres = _get_refvals(referent, "fld_thetae_pres", repeat, multi) fld_thetae_pres = n.squeeze(fld_thetae_pres) - eth = getvar(in_wrfnc, "eth") + eth = getvar(in_wrfnc, "eth", timeidx=timeidx) interp_levels = [850,500,5] field = vinterp(in_wrfnc, field=eth, @@ -388,6 +399,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, interp_levels=interp_levels, extrapolate=True, field_type="theta-e", + timeidx=timeidx, log_p=True) field = n.squeeze(field)