From ea7adb994f0e3cc570c630246d33d0ee4e599b97 Mon Sep 17 00:00:00 2001 From: Bill Ladwig Date: Fri, 20 May 2016 16:23:07 -0600 Subject: [PATCH] Fixed problems with the moving domain. Added performance tweaks to the cat operation and fixed the behavior of the cat operation to properly extract the correct time index. Fixed the join operation with multiple files/times. Now supports negative indexes and uses None to get all times. Added an ALL_TIMES constant for readability. --- wrf_open/var/src/python/wrf/var/constants.py | 4 +- wrf_open/var/src/python/wrf/var/interp.py | 8 +- wrf_open/var/src/python/wrf/var/projection.py | 11 +- wrf_open/var/src/python/wrf/var/util.py | 705 ++++++++++++++---- wrf_open/var/test/snippet.py | 30 + wrf_open/var/test/utests.py | 48 +- 6 files changed, 641 insertions(+), 165 deletions(-) create mode 100644 wrf_open/var/test/snippet.py 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)