Added support for using a single column of data for the cape_3d routine. This was an NCL feature that should have been included. Fixes #9.

src/wrf/computation.py

@@ -859,6 +859,11 @@ def cape_3d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow,
         - return_val[0,...] will contain CAPE [J kg-1]
         - return_val[1,...] will contain CIN [J kg-1]
     
+    This function also supports computing CAPE along a single vertical 
+    column.  In this mode, the *pres_hpa*, *tkel*, *qv* and *height* arguments
+    must be one-dimensional vertical columns, and the *terrain* and 
+    *psfc_hpa* arguments must be scalar values 
+    (:obj:`float`, :class:`numpy.float32` or :class:`numpy.float64`).
     
     This is the raw computational algorithm and does not extract any variables 
     from WRF output files.  Use :meth:`wrf.getvar` to both extract and compute
@@ -868,9 +873,12 @@ def cape_3d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow,
             
         pres_hpa (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full 
             pressure (perturbation + base state pressure) in [hPa] with at 
-            least three dimensions. The rightmost dimensions can be 
+            least three dimensions when operating on a grid of values. The 
-            top_bottom x south_north x west_east or bottom_top x south_north x
+            rightmost dimensions can be top_bottom x south_north x west_east 
-            west_east.
+            or bottom_top x south_north x west_east.  
+            When operating on only a single column of values, the vertical 
+            column can be bottom_top or top_bottom.  In this case, *terrain* 
+            and *psfc_hpa* must be scalars.
             
             Note:
             
@@ -893,15 +901,21 @@ def cape_3d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow,
             Geopotential height in [m] with the same dimensionality as 
             *pres_hpa*.
             
-        terrain (:class:`xarray.DataArray` or :class:`numpy.ndarray`): 
+        terrain (:class:`xarray.DataArray`, :class:`numpy.ndarray`, \
-            Terrain height in [m].  This is at least a two-dimensional array 
+            or a scalar): Terrain height in [m].  When operating on a grid of 
+            values, this argument is at least a two-dimensional array 
             with the same dimensionality as *pres_hpa*, excluding the vertical 
-            (bottom_top/top_bottom) dimension.
+            (bottom_top/top_bottom) dimension.  When operating on a single 
+            vertical column, this argument must be a scalar (:obj:`float`, 
+            :class:`numpy.float32`, or :class:`numpy.float64`).
             
-        psfc_hpa (:class:`xarray.DataArray` or :class:`numpy.ndarray`): 
+        psfc_hpa (:class:`xarray.DataArray`, :class:`numpy.ndarray`, \
-            The surface pressure in [hPa].  This is at least a two-dimensional 
+            or a scalar): Surface pressure in [hPa].  When operating on a 
-            array with the same dimensionality as *pres_hpa*, excluding the 
+            grid of values, this argument is at least a two-dimensional array 
-            vertical (bottom_top/top_bottom) dimension.
+            with the same dimensionality as *pres_hpa*, excluding the vertical 
+            (bottom_top/top_bottom) dimension.  When operating on a single 
+            vertical column, this argument must be a scalar (:obj:`float`, 
+            :class:`numpy.float32`, or :class:`numpy.float64`).
             
             Note:
             

src/wrf/extension.py

@@ -18,7 +18,7 @@ from .decorators import (left_iteration, cast_type,
 from .util import combine_dims, npbytes_to_str, psafilepath
 from .py3compat import py3range
 from .specialdec import (uvmet_left_iter, cape_left_iter, 
-                         cloudfrac_left_iter)
+                         cloudfrac_left_iter, check_cape_args)
 
 class DiagnosticError(Exception):
     """Raised when an error occurs in a diagnostic routine."""
@@ -576,7 +576,7 @@ def _dbz(p, tk, qv, qr, qs, qg, sn0, ivarint, iliqskin, outview=None):
     return result
 
 
-@check_args(0, 3, (3,3,3,3,2,2))
+@check_cape_args()
 @cape_left_iter()
 @cast_type(arg_idxs=(0,1,2,3,4,5), outviews=("capeview", "cinview"))
 @extract_and_transpose(outviews=("capeview", "cinview"))

src/wrf/metadecorators.py

@@ -1921,9 +1921,20 @@ def set_cape_alg_metadata(is2d, copyarg="pres_hpa"):
             
         if not xarray_enabled() or not do_meta:
             return wrapped(*args, **kwargs)
-            
+    
         result = wrapped(*args, **kwargs)
         
+        argvals = from_args(wrapped, (copyarg,"missing"), *args, **kwargs)
+        p = argvals[copyarg]
+        missing = argvals["missing"]
+        
+        # Note: cape_3d supports using only a single column of data
+        is1d = p.ndim == 1
+        
+        # Need to squeeze the right dimensions for 1D cape
+        if is1d:
+            result = np.squeeze(result)
+        
         # Default dimension names
         outdims = ["dim_{}".format(i) for i in py3range(result.ndim)]
         
@@ -1935,15 +1946,18 @@ def set_cape_alg_metadata(is2d, copyarg="pres_hpa"):
             outattrs["units"] = "J kg-1 ; J kg-1 ; m ; m"
             outattrs["MemoryOrder"] = "XY"
         else:
-            outname = "cape_3d"
+            if not is1d:
-            outattrs["description"] = "cape; cin"
+                outname = "cape_3d"
-            outattrs["units"] = "J kg-1 ; J kg-1"
+                outattrs["description"] = "cape; cin"
-            outattrs["MemoryOrder"] = "XYZ"
+                outattrs["units"] = "J kg-1 ; J kg-1"
+                outattrs["MemoryOrder"] = "XYZ"
+            else:
+                outname = "cape_column"
+                outattrs["description"] = "cape; cin"
+                outattrs["units"] = "J kg-1 ; J kg-1"
+                outattrs["MemoryOrder"] = "Z"
+                
         
-        argvals = from_args(wrapped, (copyarg,"missing"), *args, **kwargs)
-        p = argvals[copyarg]
-        missing = argvals["missing"]
-            
         if isinstance(p, DataArray):
             if is2d:
                 # Right dims
@@ -1952,10 +1966,13 @@ def set_cape_alg_metadata(is2d, copyarg="pres_hpa"):
                 outdims[1:-2] = p.dims[0:-3]
                 
             else:
-                # Right dims
+                if not is1d:
-                outdims[-3:] = p.dims[-3:]
+                    # Right dims
-                # Left dims
+                    outdims[-3:] = p.dims[-3:]
-                outdims[1:-3] = p.dims[0:-3]
+                    # Left dims
+                    outdims[1:-3] = p.dims[0:-3]
+                else:
+                    outdims[1] = p.dims[0]
         
         outcoords = {}     
         # Left-most is always cape_cin or cape_cin_lcl_lfc

src/wrf/specialdec.py

@@ -218,36 +218,64 @@ def cape_left_iter(alg_dtype=np.float64):
         i3dflag = args[7]
         ter_follow = args[8]
         
-        is2d = False if i3dflag != 0 else True
+        is2d = i3dflag != 0
-            
+        is1d = np.isscalar(sfp)
-        # Need to order in ascending pressure order
+        
-        flip = False
-        bot_idxs = (0,) * p_hpa.ndim
-        top_idxs = list(bot_idxs)
-        top_idxs[-3] = -1
-        top_idxs = tuple(top_idxs)
-        
-        if p_hpa[bot_idxs] > p_hpa[top_idxs]:
-            flip = True
-            p_hpa = np.ascontiguousarray(p_hpa[...,::-1,:,:])
-            tk = np.ascontiguousarray(tk[...,::-1,:,:])
-            qv = np.ascontiguousarray(qv[...,::-1,:,:])
-            ht = np.ascontiguousarray(ht[...,::-1,:,:])
-            new_args[0] = p_hpa
-            new_args[1] = tk
-            new_args[2] = qv
-            new_args[3] = ht
-            
-        num_left_dims = p_hpa.ndim - 3
         orig_dtype = p_hpa.dtype
         
+        if not is1d:
+            # Need to order in ascending pressure order
+            flip = False
+            bot_idxs = (0,) * p_hpa.ndim
+            top_idxs = list(bot_idxs)
+            top_idxs[-3] = -1
+            top_idxs = tuple(top_idxs)
+        
+            if p_hpa[bot_idxs] > p_hpa[top_idxs]:
+                flip = True
+                p_hpa = np.ascontiguousarray(p_hpa[...,::-1,:,:])
+                tk = np.ascontiguousarray(tk[...,::-1,:,:])
+                qv = np.ascontiguousarray(qv[...,::-1,:,:])
+                ht = np.ascontiguousarray(ht[...,::-1,:,:])
+                new_args[0] = p_hpa
+                new_args[1] = tk
+                new_args[2] = qv
+                new_args[3] = ht
+            
+            num_left_dims = p_hpa.ndim - 3
+        else:
+            # Need to order in ascending pressure order
+            flip = False
+        
+            if p_hpa[0] > p_hpa[-1]:
+                
+                flip = True
+                p_hpa = np.ascontiguousarray(p_hpa[::-1])
+                tk = np.ascontiguousarray(tk[::-1])
+                qv = np.ascontiguousarray(qv[::-1])
+                ht = np.ascontiguousarray(ht[::-1])
+            
+            # Need to make 3D views for the fortran code.
+            new_args[0] = p_hpa[:, np.newaxis, np.newaxis]
+            new_args[1] = tk[:, np.newaxis, np.newaxis]
+            new_args[2] = qv[:, np.newaxis, np.newaxis]
+            new_args[3] = ht[:, np.newaxis, np.newaxis]
+            new_args[4] = np.full((1,1), ter, orig_dtype)
+            new_args[5] = np.full((1,1), sfp, orig_dtype)
+            
+            num_left_dims = 0
+        
         # No special left side iteration, build the output from the cape,cin
         # result
         if (num_left_dims == 0):
             cape, cin = wrapped(*new_args, **new_kwargs)
             
             output_dims = (2,)
-            output_dims += p_hpa.shape[-3:]
+            if not is1d:
+                output_dims += p_hpa.shape[-3:]
+            else:
+                output_dims += (p_hpa.shape[0], 1, 1)
+
             output = np.empty(output_dims, orig_dtype)
             
             if flip and not is2d:
@@ -258,7 +286,6 @@ def cape_left_iter(alg_dtype=np.float64):
                 output[1,:] = cin[:]
             
             return output
-                
 
         # Initial output is ...,cape_cin,nz,ny,nx to create contiguous views
         outdims = p_hpa.shape[0:num_left_dims]
@@ -433,3 +460,70 @@ def cloudfrac_left_iter(alg_dtype=np.float64):
     
     return func_wrapper
 
+
+def check_cape_args():
+    """A decorator to check that the cape_3d arguments are valid.
+    
+    An exception is raised when an invalid argument is found.
+            
+    Returns:
+    
+        None
+        
+    Raises:
+    
+        :class:`ValueError`: Raised when an invalid argument is detected.
+        
+    """
+    @wrapt.decorator
+    def func_wrapper(wrapped, instance, args, kwargs):
+        
+        p_hpa = args[0]
+        tk = args[1]
+        qv = args[2]
+        ht = args[3]
+        ter = args[4]
+        sfp = args[5]
+        missing = args[6]
+        i3dflag = args[7]
+        ter_follow = args[8]
+        
+        is2d = False if i3dflag != 0 else True
+        
+        if not (p_hpa.shape == tk.shape == qv.shape == ht.shape):
+            raise ValueError("arguments 0, 1, 2, 3 must be the same shape")
+
+        # 2D CAPE does not allow for scalars
+        if is2d:
+            if np.isscalar(ter) or np.isscalar(sfp):
+                raise ValueError("arguments 4 and 5 cannot be scalars with "
+                                 "cape_2d routine")
+                
+            if ter.ndim != p_hpa.ndim-1 or sfp.ndim != p_hpa.ndim-1:
+                raise ValueError("arguments 4 and 5 must have "
+                                 "{} dimensions".format(p_hpa.ndim-1))
+                
+        # 3D cape is allowed to be just a vertical column with scalars
+        # for terrain and psfc_hpa.
+        else:
+            if np.isscalar(ter) and np.isscalar(sfp):
+                if p_hpa.ndim != 1:
+                    raise ValueError("arguments 0-3 "
+                                     "must be 1-dimensional when "
+                                     "arguments 4 and 5 are scalars")
+                if is2d:
+                    raise ValueError("argument 7 must be 0 when using 1D data")
+            else:
+                if ((np.isscalar(ter) and not np.isscalar(sfp)) or 
+                    (not np.isscalar(ter) and np.isscalar(sfp))):
+                    raise ValueError("arguments 4 and 5 must both be scalars")
+                else:
+                    if ter.ndim != p_hpa.ndim-1 or sfp.ndim != p_hpa.ndim-1:
+                        raise ValueError("arguments 4 and 5 "
+                                         "must have {} dimensions".format(
+                                             p_hpa.ndim-1))
+      
+        return wrapped(*args, **kwargs)
+    
+    return func_wrapper
+

test/comp_utest.py

@@ -515,7 +515,6 @@ def get_args(varname, wrfnc, timeidx, method, squeeze):
         return (full_p, relh)
         
         
-                
 class WRFVarsTest(ut.TestCase):
     longMessage = True
     
@@ -541,6 +540,60 @@ def make_func(varname, wrfnc, timeidx, method, squeeze, meta):
             self.assertEqual(result.dims, ref.dims)
     
     return func
+
+
+def test_cape3d_1d(wrfnc):
+    
+    def func(self):
+        varnames = ("T", "P", "PB", "QVAPOR", "PH", "PHB", "HGT", "PSFC")
+        ncvars = extract_vars(wrfnc, 0, varnames, method="cat", squeeze=True, 
+                              cache=None, meta=True)
+        
+        t = ncvars["T"]
+        p = ncvars["P"]
+        pb = ncvars["PB"]
+        qv = ncvars["QVAPOR"]
+        ph = ncvars["PH"]
+        phb = ncvars["PHB"]
+        ter = ncvars["HGT"]
+        psfc = ncvars["PSFC"]
+        
+        col_idxs = (slice(None), t.shape[-2]//2, t.shape[-1]//2)
+        
+        t = t[col_idxs]
+        p = p[col_idxs]
+        pb = pb[col_idxs]
+        qv = qv[col_idxs]
+        ph = ph[col_idxs]
+        phb = phb[col_idxs]
+        ter = float(ter[col_idxs[1:]])
+        psfc = float(psfc[col_idxs[1:]])
+        
+        full_t = t + Constants.T_BASE
+        full_p = p + pb
+        tkel = tk(full_p, full_t, meta=False)
+        
+        geopt = ph + phb
+        geopt_unstag = destagger(geopt, -1)
+        z = geopt_unstag/Constants.G
+        
+        # Convert pressure to hPa
+        p_hpa = ConversionFactors.PA_TO_HPA * full_p
+        psfc_hpa = ConversionFactors.PA_TO_HPA * psfc 
+        
+        i3dflag = 1
+        ter_follow = 1
+        
+        result = cape_3d(p_hpa, tkel, qv, z, ter, psfc_hpa, ter_follow)
+    
+        ref = getvar(wrfnc, "cape_3d")
+
+        ref = ref[(slice(None),) + col_idxs]
+        
+        nt.assert_allclose(to_np(result), to_np(ref))
+    
+    return func
+        
         
 if __name__ == "__main__":
     varnames = ["avo", "eth", "cape_2d", "cape_3d", "ctt", "dbz", "mdbz", 
@@ -575,6 +628,10 @@ if __name__ == "__main__":
     
                             setattr(WRFVarsTest, test_name, func)
     
+    func = test_cape3d_1d(wrfnc)
+    setattr(WRFVarsTest, "test_cape3d_1d", func)
+    
+    
     ut.main()