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

@@ -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

wrf_open/var/src/python/wrf/var/interp.py

@@ -227,6 +227,12 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False,
     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)

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"]
 
@@ -181,6 +181,15 @@ class WrfProj(object):
         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, 
                  lats=None, lons=None, **proj_params):

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",
+           "viewvalues", "py2round", "combine_with", "either", 
-           "args_to_list", "npvalues", "CoordPair"]
+           "from_args", "arg_location", "args_to_list", "npvalues", 
+           "CoordPair"]
 
 
 _COORD_PAIR_MAP = {"XLAT" : ("XLAT", "XLONG"),
@@ -41,21 +52,29 @@ _COORD_PAIR_MAP = {"XLAT" : ("XLAT", "XLONG"),
 _COORD_VARS = ("XLAT", "XLONG", "XLAT_M", "XLONG_M", "XLAT_U", "XLONG_U",
                "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)
@@ -158,6 +177,18 @@ def isstr(s):
     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):
     if not isinstance(da, DataArray):
@@ -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")):
     
@@ -270,6 +302,20 @@ def _is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"),
         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) 
     
-    time_idx_or_slice = timeidx if not multitime else slice(None)
+    # 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)
 
-    # wrfseq might be a generator
+    #time_idx_or_slice = timeidx if not multitime else slice(None)
+    
+    # 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, :]
+        first_var = (next(wrf_iter)).variables[varname][:]
-        if not multitime:
-            first_var = first_var[np.newaxis,:]
     
     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
                 
             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")
         
-        outcoords[outdimnames[0]] = file_times # New time dimension values
+        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:]
+            
+            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, 
+        first_var = _build_data_array(wrfnc, varname, ALL_TIMES, is_moving)
-                                      timeidx, 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, :]
+        first_var = wrfnc.variables[varname][:]
-        if not multitime:
-            first_var = first_var[np.newaxis, :]
     
-    # 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
+    outdims += [maxtimes]
-    outdata = np.empty(outdims, first_var.dtype)
+    outdims += first_var.shape[1:]
     
-    outdata[0,:] = first_var[:]
+    # 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[:]
     
-    idx=1
+    # 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)[:]
+    
+    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)
+        outdimnames = ["file"] + list(first_var.dims)
-        outcoords["file_idx"] = [i for i in xrange(numfiles)]
+        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):
@@ -876,8 +1257,7 @@ def from_args(func, argnames, *args, **kwargs):
         
         return result
 
-def args_to_list(func, args, kwargs):
+def _args_to_list2(func, args, kwargs):
-    """Converts the mixed args/kwargs to a single list of args"""
     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):
+def _args_to_list3(func, args, kwargs):
-    """Parses the function args, kargs and signature looking for the desired 
+    sig = signature(func)
-    argument location (either in args, kargs, or argspec.defaults), 
+    bound = sig.bind(*args, **kwargs)
-    and returns a list containing representing all arguments in the 
+    bound.apply_defaults()
-    correct order with defaults filled in.
     
-    """
+    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:
@@ -914,6 +1307,30 @@ def arg_location(func, argname, args, kwargs):
     
     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)
+    
+    
     
         
 

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()

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")
+            hts = getvar(in_wrfnc, "z", timeidx=timeidx)
-            p = getvar(in_wrfnc, "pressure")
+            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")
+            hts = getvar(in_wrfnc, "z", timeidx=timeidx)
-            p = getvar(in_wrfnc, "pressure")
+            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]
             
@@ -265,6 +269,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.
@@ -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.
@@ -305,6 +311,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)
@@ -321,6 +328,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)
@@ -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)