Bug fixes for multi indexing after unit testing

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

@@ -4,7 +4,7 @@ import numpy.ma as ma
 from wrf.var.extension import computetk,computecape
 from wrf.var.destagger import destagger
 from wrf.var.constants import Constants, ConversionFactors
-from wrf.var.util import extract_vars, hold_dim_fixed
+from wrf.var.util import extract_vars
 
 __all__ = ["get_2dcape", "get_3dcape"]
 
@@ -27,7 +27,7 @@ def get_2dcape(wrfnc, missing=-999999.0, timeidx=0):
     tk = computetk(full_p, full_t)
     
     geopt = ph + phb
-    geopt_unstag = destagger(geopt, 0)
+    geopt_unstag = destagger(geopt, -3)
     z = geopt_unstag/Constants.G
     
     # Convert pressure to hPa
@@ -40,21 +40,23 @@ def get_2dcape(wrfnc, missing=-999999.0, timeidx=0):
     cape_res,cin_res = computecape(p_hpa,tk,qv,z,ter,psfc_hpa,
                                    missing,i3dflag,ter_follow)
      
-    cape = hold_dim_fixed(cape_res, -3, 0)
+    cape = cape_res[...,0,:,:]
-    cin = hold_dim_fixed(cin_res, -3, 0)
+    cin = cin_res[...,0,:,:]
-    lcl = hold_dim_fixed(cin_res, -3, 1)
+    lcl = cin_res[...,1,:,:]
-    lfc = hold_dim_fixed(cin_res, -3, 2)
+    lfc = cin_res[...,2,:,:]
     
-    resdim = [4]
+    left_dims = [x for x in cape_res.shape[0:-3]]
-    resdim += cape.shape
+    right_dims = [x for x in cape_res.shape[-2:]]
+    
+    resdim = left_dims + [4] + right_dims
     
     # Make a new output array for the result
     res = n.zeros((resdim), cape.dtype)
     
-    res[0,:] = cape[:]
+    res[...,0,:,:] = cape[:]
-    res[1,:] = cin[:]
+    res[...,1,:,:] = cin[:]
-    res[2,:] = lcl[:]
+    res[...,2,:,:] = lcl[:]
-    res[3,:] = lfc[:]
+    res[...,3,:,:] = lfc[:]
     
     return ma.masked_values(res,missing)
 
@@ -77,7 +79,7 @@ def get_3dcape(wrfnc, missing=-999999.0, timeidx=0):
     tk = computetk(full_p, full_t)
     
     geopt = ph + phb
-    geopt_unstag = destagger(geopt, 0)
+    geopt_unstag = destagger(geopt, -3)
     z = geopt_unstag/Constants.G
     
     # Convert pressure to hPa
@@ -91,13 +93,17 @@ def get_3dcape(wrfnc, missing=-999999.0, timeidx=0):
                            missing,i3dflag,ter_follow)
     
     # Make a new output array for the result
-    resdim = [2]
+    left_dims = [x for x in cape.shape[0:-3]]
-    resdim += cape.shape
+    right_dims = [x for x in cape.shape[-3:]]
-    res = n.zeros((resdim), cape.dtype)
+    
+    resdim = left_dims + [2] + right_dims
+    
+    res = n.zeros(resdim, cape.dtype)
     
-    res[0,:] = cape[:]
+    res[...,0,:,:,:] = cape[:]
-    res[1,:] = cin[:]
+    res[...,1,:,:,:] = cin[:]
     
+    #return res
     return ma.masked_values(res, missing)
     
     

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

@@ -46,7 +46,7 @@ def get_ctt(wrfnc, units="c", timeidx=0):
     tk = computetk(full_p, full_t)
     
     geopt = ph + phb
-    geopt_unstag = destagger(geopt, 0)
+    geopt_unstag = destagger(geopt, -3)
     ght = geopt_unstag / Constants.G
     
     ctt = computectt(p_hpa,tk,qice,qcld,qv,ght,ter,haveqci)

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

@@ -60,5 +60,5 @@ def get_dbz(wrfnc, do_varint=False, do_liqskin=False, timeidx=0):
 
 def get_max_dbz(wrfnc, do_varint=False, do_liqskin=False, timeidx=0):
     return n.amax(get_dbz(wrfnc, do_varint, do_liqskin, timeidx), 
-                  axis=0)
+                  axis=-3)
 

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

@@ -5,6 +5,7 @@ from itertools import product
 import numpy as n
 
 from wrf.var.units import do_conversion, check_units
+from wrf.var.destagger import destagger
 
 __all__ = ["convert_units", "handle_left_iter"]
 
@@ -59,7 +60,7 @@ def _left_indexes(dims):
 
 def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1,
                    ref_var_name=None, alg_out_fixed_dims=(), 
-                   ignore_args=(), ignore_kargs={}):
+                   ignore_args=(), ignore_kargs={}, ref_stag_dim=None):
     """Decorator to handle iterating over leftmost dimensions when using 
     multiple files and/or multiple times.
     
@@ -78,18 +79,22 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1,
         directly without any slicing
         - ignore_kargs - keys of any keyword arguments which should be passed
         directly without slicing
+        - ref_stag_dim - in some cases the reference variable dimensions
+        have a staggered dimension that needs to be corrected when calculating
+        the output dimensions (e.g. avo)
     
     """
     def indexing_decorator(func):
         @wraps(func)
         def func_wrapper(*args, **kargs):
+            
             if ref_var_idx >= 0:
                 ref_var = args[ref_var_idx]
             else:
                 ref_var = kargs[ref_var_name]
             
             ref_var_shape = ref_var.shape
-            extra_dim_num = len(ref_var_shape) - ref_var_expected_dims
+            extra_dim_num = ref_var.ndim- ref_var_expected_dims
             
             # No special left side iteration, return the function result
             if (extra_dim_num == 0):
@@ -109,10 +114,13 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1,
                 # numerical algorithm has greater dim count than reference, 
                 # Note: user must provide this information to decorator
                 right_dims = [dim for dim in alg_out_fixed_dims]
+                neg_start_idx = -(alg_out_num_dims - len(alg_out_fixed_dims))
                 right_dims += [ref_var_shape[x] 
-                        for x in xrange(-alg_out_num_dims,0,1)]
+                        for x in xrange(neg_start_idx,0,1)]
                 outdims += right_dims
             
+            if ref_stag_dim is not None:
+                outdims[ref_stag_dim] -= 1  
             
             out_inited = False
             for left_idxs in _left_indexes(extra_dims):
@@ -121,7 +129,6 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1,
                 # the right (e.g. [1,1,:])
                 left_and_slice_idxs = ([x for x in left_idxs] + 
                                        [slice(None, None, None)])
-                
                 # Slice the args if applicable
                 new_args = [arg[left_and_slice_idxs] 
                             if i not in ignore_args else arg 
@@ -133,7 +140,6 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1,
                              if key not in ignore_kargs else val)
                              for key,val in kargs.iteritems()}
                 
-                
                 # Call the numerical routine
                 res = func(*new_args, **new_kargs)
                 
@@ -147,11 +153,12 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1,
                     
                 else:   # This should be a list or a tuple (cape)
                     if not out_inited:
-                        output = [n.zeros(outdims, ref_var.dtype)]
+                        output = [n.zeros(outdims, ref_var.dtype) 
+                                  for i in xrange(len(res))]
                         out_inited = True
                     
-                    for outarr in output:
+                    for i,outarr in enumerate(res):
-                        outarr[left_and_slice_idxs] = res[:]
+                        (output[i])[left_and_slice_idxs] = outarr[:]
                 
             
             return output
@@ -160,4 +167,86 @@ def handle_left_iter(alg_out_num_dims, ref_var_expected_dims, ref_var_idx=-1,
     
     return indexing_decorator
 
+def uvmet_left_iter():
+    """Decorator to handle iterating over leftmost dimensions when using 
+    multiple files and/or multiple times with the uvmet product.
+    
+    """
+    def indexing_decorator(func):
+        @wraps(func)
+        def func_wrapper(*args):
+            u = args[0]
+            v = args[1]
+            lat = args[2]
+            lon = args[3]
+            cen_long  = args[4]
+            cone = args[5]
+            
+            if u.ndim == lat.ndim:
+                num_right_dims = 2
+                is_3d = False
+            else:
+                num_right_dims = 3
+                is_3d = True
+            
+            is_stag = False
+            if ((u.shape[-1] != lat.shape[-1]) or 
+                (u.shape[-2] != lat.shape[-2])):
+                is_stag = True
+            
+            if is_3d:
+                extra_dim_num = u.ndim - 3
+            else:
+                extra_dim_num = u.ndim - 2
+                
+            if is_stag:
+                u = destagger(u,-1)
+                v = destagger(v,-2)
+            
+            # No special left side iteration, return the function result
+            if (extra_dim_num == 0):
+                return func(u,v,lat,lon,cen_long,cone)
+            
+            # Start by getting the left-most 'extra' dims
+            outdims = [u.shape[x] for x in xrange(extra_dim_num)]
+            extra_dims = list(outdims) # Copy the left-most dims for iteration
+            
+            # Append the right-most dimensions
+            outdims += [2] # For u/v components
+            
+            outdims += [u.shape[x] for x in xrange(-num_right_dims,0,1)]
+            
+            output = n.zeros(outdims, u.dtype)
+            
+            for left_idxs in _left_indexes(extra_dims):
+                # Make the left indexes plus a single slice object
+                # The single slice will handle all the dimensions to
+                # the right (e.g. [1,1,:])
+                left_and_slice_idxs = ([x for x in left_idxs] + 
+                                       [slice(None, None, None)])
+                        
+                
+                new_u = u[left_and_slice_idxs]
+                new_v = v[left_and_slice_idxs]
+                new_lat = lat[left_and_slice_idxs]
+                new_lon = lon[left_and_slice_idxs]
+                
+                # Call the numerical routine
+                res = func(new_u, new_v, new_lat, new_lon, cen_long, cone)
+                
+                # Note:  The 2D version will return a 3D array with a 1 length
+                # dimension.  Numpy is unable to broadcast this without 
+                # sqeezing first.
+                res = n.squeeze(res) 
+                
+                output[left_and_slice_idxs] = res[:]
+                
+            
+            return output
+        
+        return func_wrapper
+    
+    return indexing_decorator
+
+
 

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

@@ -1,6 +1,8 @@
 
 import numpy as n
 
+from wrf.var.util import extract_vars
+
 __all__ = ["destagger", "destagger_windcomp", "destagger_winds"]
 
 def destagger(var, stagger_dim):
@@ -50,7 +52,9 @@ def destagger_windcomp(wrfnc, comp, timeidx=0):
         wrfvar = "W"
         stagdim = -3
         
-    wind_data = wrfnc.variables[wrfvar][timeidx,:,:,:]
+    #wind_data = wrfnc.variables[wrfvar][timeidx,:,:,:]
+    ncvars = extract_vars(wrfnc, timeidx, vars=wrfvar)
+    wind_data = ncvars[wrfvar]
     return destagger(wind_data, stagdim)
 
 def destagger_winds(wrfnc, timeidx=0):

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

@@ -11,7 +11,7 @@ from wrf.var._wrfext import (f_interpz3d, f_interp2dxy,f_interp1d,
                      f_computesrh, f_computeuh, f_computepw, f_computedbz,
                      f_lltoij, f_ijtoll, f_converteta, f_computectt)
 from wrf.var._wrfcape import f_computecape
-from wrf.var.decorators import handle_left_iter
+from wrf.var.decorators import handle_left_iter, uvmet_left_iter
 
 __all__ = ["FortranException", "computeslp", "computetk", "computetd", 
            "computerh", "computeavo", "computepvo", "computeeth", 
@@ -86,7 +86,7 @@ def computerh(qv,q,t):
     res = n.reshape(res, shape, "A")
     return res.astype("float32")
 
-@handle_left_iter(3,3,0, ignore_args=(6,7))
+@handle_left_iter(3,3,0, ignore_args=(6,7), ref_stag_dim=-1)
 def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy):
     res = f_computeabsvort(u.astype("float64").T,
                            v.astype("float64").T,
@@ -99,7 +99,7 @@ def computeavo(u,v,msfu,msfv,msfm,cor,dx,dy):
     
     return res.astype("float32").T
 
-@handle_left_iter(3,3,0, ignore_args=(8,9))
+@handle_left_iter(3,3,2, ignore_args=(8,9))
 def computepvo(u,v,theta,prs,msfu,msfv,msfm,cor,dx,dy):
     
     res = f_computepvo(u.astype("float64").T,
@@ -124,7 +124,7 @@ def computeeth(qv, tk, p):
     
     return res.astype("float32").T
 
-@handle_left_iter(4,2,2, ignore_args=(4,5), alg_out_fixed_dims=(2,))
+@uvmet_left_iter()
 def computeuvmet(u,v,lat,lon,cen_long,cone):
     longca = n.zeros((lat.shape[-2], lat.shape[-1]), "float64")
     longcb = n.zeros((lon.shape[-2], lon.shape[-1]), "float64")
@@ -132,7 +132,7 @@ def computeuvmet(u,v,lat,lon,cen_long,cone):
     
     
     # Make the 2D array a 3D array with 1 dimension
-    if u.ndim != 3:
+    if u.ndim < 3:
         u = u.reshape((1,u.shape[-2], u.shape[-1]))
         v = v.reshape((1,v.shape[-2], v.shape[-1]))
 
@@ -154,7 +154,7 @@ def computeuvmet(u,v,lat,lon,cen_long,cone):
                0)
 
     
-    return res.astype("float32").T
+    return n.squeeze(res.astype("float32").T)
 
 @handle_left_iter(3,3,0)
 def computeomega(qv, tk, w, p):
@@ -276,10 +276,11 @@ def computecape(p_hpa,tk,qv,ht,ter,sfp,missing,i3dflag,ter_follow):
                   FortranException())
     
     # Don't need to transpose since we only passed a view to fortran
-    cape = flip_cape.astype("float32")
+    cape = flip_cape.astype("float32")[::-1,:,:]
-    cin = flip_cin.astype("float32")
+    cin = flip_cin.astype("float32")[::-1,:,:]
+    
     # Remember to flip cape and cin back to descending p coordinates
-    return (cape[::-1,:,:],cin[::-1,:,:])
+    return (cape, cin)
 
 # TODO:  This should handle lists of coords
 def computeij(map_proj,truelat1,truelat2,stdlon,

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

@@ -23,13 +23,13 @@ def _get_geoht(wrfnc, height=True, msl=True, timeidx=0):
                                "NetCDF file")
         else:
             geopt_unstag = ght_vars["GHT"] * Constants.G
-            hgt = destagger(ght_vars["HGT_U"], 1)
+            hgt = destagger(ght_vars["HGT_U"], -1)
     else:
         ph = ph_vars["PH"]
         phb = ph_vars["PHB"]
         hgt = ph_vars["HGT"]
         geopt = ph + phb
-        geopt_unstag = destagger(geopt, 0)
+        geopt_unstag = destagger(geopt, -3)
     
     if height:
         if msl:

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

@@ -40,14 +40,14 @@ def get_srh(wrfnc, top=3000.0, timeidx=0):
         v = destagger(v_vars["V"], -2) 
 
     geopt = ph + phb
-    geopt_unstag = destagger(geopt, 0)
+    geopt_unstag = destagger(geopt, -3)
     
     z = geopt_unstag / Constants.G
     
     # Re-ordering from high to low
-    u1 = u[::-1,:,:]
+    u1 = u[...,::-1,:,:] 
-    v1 = v[::-1,:,:]
+    v1 = v[...,::-1,:,:]
-    z1 = z[::-1,:,:]
+    z1 = z[...,::-1,:,:]
     
     srh = computesrh(u1, v1, z1, ter, top)
     

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

@@ -35,7 +35,7 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None,
     """ 
     
     # Have a pivot point with an angle to find cross section
-    if pivot is not None and angle is not None:
+    if pivot_point is not None and angle is not None:
         xp = pivot_point[0]
         yp = pivot_point[1]
         
@@ -132,7 +132,7 @@ def _get_xy(xdim, ydim, pivot_point=None, angle=None,
 
 # TODO:  Add flag to use lat/lon points by doing conversion
 def get_vertcross(data3d, z, missingval=-99999, 
-                  pivot=None,angle=None,start_point=None,end_point=None):
+                  pivot_point=None,angle=None,start_point=None,end_point=None):
     
     xdim = z.shape[2]
     ydim = z.shape[1]

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

@@ -1,3 +1,5 @@
+from collections import Iterable
+
 from wrf.var.extension import computeij, computell
 from wrf.var.util import extract_vars, extract_global_attrs
 

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

@@ -14,7 +14,7 @@ def get_omega(wrfnc, timeidx=0):
     pb = ncvars["PB"]
     qv = ncvars["QVAPOR"]
     
-    wa = destagger(w, 0)
+    wa = destagger(w, -3)
     full_t = t + Constants.T_BASE
     full_p = p + pb
     tk = computetk(full_p, full_t)

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

@@ -23,7 +23,7 @@ def get_slp(wrfnc, units="hpa", timeidx=0):
     qvapor[qvapor < 0] = 0.
     full_ph = (ph + phb) / Constants.G
     
-    destag_ph = destagger(full_ph, 0)
+    destag_ph = destagger(full_ph, -3)
     
     tk = computetk(full_p, full_t)
     slp = computeslp(destag_ph, tk, full_p, qvapor)

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

@@ -4,11 +4,9 @@ import datetime as dt
 __all__ = ["get_times"]
 
 def _make_time(timearr):
-    return dt.strptime("".join(timearr[:]), "%Y-%m-%d_%H:%M:%S")
+    return dt.datetime.strptime("".join(timearr[:]), "%Y-%m-%d_%H:%M:%S")
 
 # TODO:  handle list of files and multiple times
 def get_times(wrfnc):
     times = wrfnc.variables["Times"][:,:]
     return [_make_time(times[i,:]) for i in xrange(times.shape[0])]
-
- 

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

@@ -3,7 +3,7 @@ from collections import Iterable
 import numpy as n
 
 __all__ = ["extract_vars", "extract_global_attrs", "extract_dim",
-           "hold_dim_fixed", "combine_files"]
+           "combine_files"]
 
 def _is_multi_time(timeidx):
     if timeidx == -1:
@@ -82,28 +82,5 @@ def extract_vars(wrfnc, timeidx, vars):
         return {var:combine_files(wrfnc, var, timeidx) for var in varlist}
         
     
-def hold_dim_fixed(var, dim, idx):
-    """Generic method to hold a single dimension to a fixed index when 
-    the array shape is unknown.  The values for 'dim' and 'idx' can 
-    be negative.
-    
-    For example, on a 4D array with 'dim' set to 
-    -1 and 'idx' set to 3, this is going to return this view:
-    
-    var[:,:,:,3]
-    
-    """
-    
-    var_shape = var.shape
-    num_dims = len(var_shape) 
-    full_slice = slice(None, None, None)
-    
-    # Make the sequence of slices
-    dim_slices = [full_slice for x in xrange(num_dims)]
-    
-    # Replace the dimension with the fixed index
-    dim_slices[dim] = idx
-
-    return var[dim_slices]
     
     

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

@@ -5,7 +5,7 @@ from wrf.var.destagger import destagger
 from wrf.var.constants import Constants
 from wrf.var.wind import _calc_wspd_wdir
 from wrf.var.decorators import convert_units
-from wrf.var.util import extract_vars, extract_global_attrs, hold_dim_fixed
+from wrf.var.util import extract_vars, extract_global_attrs
 
 __all__=["get_uvmet", "get_uvmet10", "get_uvmet_wspd_wdir", 
          "get_uvmet10_wspd_wdir"]
@@ -50,7 +50,7 @@ def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0):
             else:
                 # For met_em files, this just takes the lowest level winds
                 # (3rd dimension from right is bottom_top)
-                u = destagger(hold_dim_fixed(uu_vars["UU"], -3, 0), -1) # support met_em files
+                u = destagger(uu_vars["UU"][...,0,:,:], -1) # support met_em files
         else:
             u = u_vars["U10"] 
         
@@ -64,7 +64,7 @@ def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0):
             else:
                 # For met_em files, this just takes the lowest level winds
                 # (3rd dimension from right is bottom_top)
-                v = destagger(hold_dim_fixed(vv_vars["VV"], -3, 0), -2) # support met_em files
+                v = destagger(vv_vars["VV"][...,0,:,:], -2) # support met_em files
         else:
             v = v_vars["V10"]
     
@@ -141,10 +141,7 @@ def get_uvmet(wrfnc, ten_m=False, units ="mps", timeidx=0):
         
         res = computeuvmet(u,v,lat,lon,cen_lon,cone)
         
-        if u.ndim == 3:
         return res
-        else:
-            return res[:,0,:,:]
             
     
 def get_uvmet10(wrfnc, units="mps", timeidx=0):

wrf_open/var/test/utests.py

@@ -1,6 +1,7 @@
 import unittest as ut
 import numpy.testing as nt 
 import numpy as n
+import numpy.ma as ma
 import os, sys
 import subprocess
 
@@ -44,12 +45,36 @@ def setUpModule():
 
 # Using helpful information at: 
 # http://eli.thegreenplace.net/2014/04/02/dynamically-generating-python-test-cases
-def make_test(varname, wrf_in, referent):
+def make_test(varname, wrf_in, referent, multi=False, repeat=3):
     def test(self):
+        
+        if not multi:
             in_wrfnc = NetCDF(wrf_in)
+        else:
+            nc = NetCDF(wrf_in)
+            in_wrfnc = [nc for i in xrange(repeat)]
+        
         refnc = NetCDF(referent)
             
-        ref_vals = refnc.variables[varname][...]
+        if not multi:
+            ref_vals = refnc.variables[varname][:]
+        else:
+            data = refnc.variables[varname][:]
+            new_dims = [repeat] + [x for x in data.shape]
+            masked=False
+            if (isinstance(data, ma.core.MaskedArray)):
+                masked=True
+            
+            if not masked:
+                ref_vals = n.zeros(new_dims, data.dtype)
+            else:
+                ref_vals = ma.asarray(n.zeros(new_dims, data.dtype))
+                
+            for i in xrange(repeat):
+                ref_vals[i,:] = data[:]
+                
+                if masked:
+                    ref_vals.mask[i,:] = data.mask[:]
         
         if (varname == "tc"):
             my_vals = getvar(in_wrfnc, "temp", units="c")
@@ -104,20 +129,16 @@ def make_test(varname, wrf_in, referent):
             atol = 0
             nt.assert_allclose(my_vals, ref_vals, tol, atol)
         elif (varname == "cape_2d"):
-            mcape, mcin, lcl, lfc = getvar(in_wrfnc, varname)
+            cape_2d = getvar(in_wrfnc, varname)
             tol = 1/100.
             atol = 0
-            nt.assert_allclose(mcape, ref_vals[0,:,:], tol, atol)
+            nt.assert_allclose(cape_2d, ref_vals, tol, atol)
-            nt.assert_allclose(mcin, ref_vals[1,:,:], tol, atol)
-            nt.assert_allclose(lcl, ref_vals[2,:,:], tol, atol)
-            nt.assert_allclose(lfc, ref_vals[3,:,:], tol, atol)
         elif (varname == "cape_3d"):
-            cape, cin = getvar(in_wrfnc, varname)
+            cape_3d = getvar(in_wrfnc, varname)
             tol = 1/100.
             atol = 0
-            nt.assert_allclose(cape, ref_vals[0,:,:], tol, atol)
-            nt.assert_allclose(cin, ref_vals[1,:,:], tol, atol)
             
+            nt.assert_allclose(cape_3d, ref_vals, tol, atol)
             
         else:
             my_vals = getvar(in_wrfnc, varname)
@@ -144,8 +165,10 @@ if __name__ == "__main__":
         if var in ignore_vars:
             continue
         
-        test_func = make_test(var, TEST_FILE, OUT_NC_FILE)
+        test_func1 = make_test(var, TEST_FILE, OUT_NC_FILE)
-        setattr(WRFVarsTest, 'test_{0}'.format(var), test_func)
+        test_func2 = make_test(var, TEST_FILE, OUT_NC_FILE, multi=True)
+        setattr(WRFVarsTest, 'test_{0}'.format(var), test_func1)
+        setattr(WRFVarsTest, 'test_multi_{0}'.format(var), test_func2)
     
     
     ut.main()