diff --git a/src/wrf/computation.py b/src/wrf/computation.py index dcca9b8..d01e1ab 100644 --- a/src/wrf/computation.py +++ b/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 - top_bottom x south_north x west_east or bottom_top x south_north x - west_east. + least three dimensions when operating on a grid of values. The + rightmost dimensions can be top_bottom x south_north x 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 height in [m]. This is at least a two-dimensional array + terrain (:class:`xarray.DataArray`, :class:`numpy.ndarray`, \ + 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`): - The surface pressure in [hPa]. This is at least a two-dimensional - array with the same dimensionality as *pres_hpa*, excluding the - vertical (bottom_top/top_bottom) dimension. + psfc_hpa (:class:`xarray.DataArray`, :class:`numpy.ndarray`, \ + or a scalar): Surface pressure in [hPa]. 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. When operating on a single + vertical column, this argument must be a scalar (:obj:`float`, + :class:`numpy.float32`, or :class:`numpy.float64`). Note: diff --git a/src/wrf/extension.py b/src/wrf/extension.py index d1d43c5..2b86650 100755 --- a/src/wrf/extension.py +++ b/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")) diff --git a/src/wrf/metadecorators.py b/src/wrf/metadecorators.py index 1227ee1..e34b71d 100644 --- a/src/wrf/metadecorators.py +++ b/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" - outattrs["description"] = "cape; cin" - outattrs["units"] = "J kg-1 ; J kg-1" - outattrs["MemoryOrder"] = "XYZ" + if not is1d: + outname = "cape_3d" + outattrs["description"] = "cape; cin" + 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 - outdims[-3:] = p.dims[-3:] - # Left dims - outdims[1:-3] = p.dims[0:-3] + if not is1d: + # Right dims + outdims[-3:] = p.dims[-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 diff --git a/src/wrf/specialdec.py b/src/wrf/specialdec.py index 77f6161..db8e5ae 100644 --- a/src/wrf/specialdec.py +++ b/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 - - # 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 + is2d = i3dflag != 0 + is1d = np.isscalar(sfp) + 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 + diff --git a/test/comp_utest.py b/test/comp_utest.py index 4cb4e5a..846466f 100644 --- a/test/comp_utest.py +++ b/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() \ No newline at end of file