Merge pull request #577 from jrleeman/SRH_Bounds

Storm Relative Helicity Improvements

metpy/calc/kinematics.py

@@ -9,7 +9,7 @@
 
 import numpy as np
 
-from ..calc.tools import get_layer
+from ..calc.tools import get_layer_heights
 from ..cbook import is_string_like, iterable
 from ..constants import Cp_d, g
 from ..package_tools import Exporter
@@ -539,18 +539,14 @@ def montgomery_streamfunction(height, temperature):
 
 
 @exporter.export
-@check_units('[speed]', '[speed]', '[pressure]', '[length]',
-             '[length]', '[length]', '[speed]', '[speed]')
-def storm_relative_helicity(u, v, p, hgt, top, bottom=0 * units('meter'),
+@check_units('[speed]', '[speed]', '[length]', '[length]', '[length]',
+             '[speed]', '[speed]')
+def storm_relative_helicity(u, v, heights, depth, bottom=0 * units.m,
                             storm_u=0 * units('m/s'), storm_v=0 * units('m/s')):
     # Partially adapted from similar SharpPy code
     r"""Calculate storm relative helicity.
 
-    Needs u and v wind components, heights and pressures,
-    and top and bottom of SRH layer. An optional storm
-    motion vector can be specified. SRH is calculated using the
-    equation specified on p. 230-231 in the Markowski and Richardson
-    meso textbook [Markowski2010].
+    Calculates storm relatively helicity following [Markowski2010] 230-231.
 
     .. math:: \int\limits_0^d (\bar v - c) \cdot \bar\omega_{h} \,dz
 
@@ -562,48 +558,37 @@ def storm_relative_helicity(u, v, p, hgt, top, bottom=0 * units('meter'),
     Parameters
     ----------
     u : array-like
-        The u components of winds, same length as hgts
+        u component winds
     v : array-like
-        The u components of winds, same length as hgts
-    p : array-like
-        Pressure in hPa, same length as hgts
-    hgt : array-like
-        The heights associatd with the data, provided in meters above mean
-        sea level and converted into meters AGL.
-    top : number
-        The height of the top of the desired layer for SRH.
+        v component winds
+    heights : array-like
+        atmospheric heights, will be converted to AGL
+    depth : number
+        depth of the layer
     bottom : number
-        The height at the bottom of the SRH layer. Default is sfc (None).
+        height of layer bottom AGL (default is surface)
     storm_u : number
-        u component of storm motion
+        u component of storm motion (default is 0 m/s)
     storm_v : number
-        v component of storm motion
+        v component of storm motion (default is 0 m/s)
 
     Returns
     -------
-    number
-        p_srh : positive storm-relative helicity
-    number
-        n_srh : negative storm-relative helicity
-    number
-        t_srh : total storm-relative helicity
+    `pint.Quantity, pint.Quantity, pint.Quantity`
+        positive, negative, total storm-relative helicity
 
     """
-    # Converting to m/s to make sure output is in m^2/s^2
-    u = u.to('meters/second')
-    v = v.to('meters/second')
-    storm_u = storm_u.to('meters/second')
-    storm_v = storm_v.to('meters/second')
+    _, u, v = get_layer_heights(heights, depth, u, v, with_agl=True, bottom=bottom)
 
-    w_int = get_layer(p, u, v, heights=hgt, bottom=bottom, depth=top - bottom)
+    storm_relative_u = u - storm_u
+    storm_relative_v = v - storm_v
 
-    sru = w_int[1] - storm_u
-    srv = w_int[2] - storm_v
+    int_layers = (storm_relative_u[1:] * storm_relative_v[:-1] -
+                  storm_relative_u[:-1] * storm_relative_v[1:])
 
-    int_layers = sru[1:] * srv[:-1] - sru[:-1] * srv[1:]
+    positive_srh = int_layers[int_layers.magnitude > 0.].sum()
+    negative_srh = int_layers[int_layers.magnitude < 0.].sum()
 
-    p_srh = int_layers[int_layers.magnitude > 0.].sum()
-    n_srh = int_layers[int_layers.magnitude < 0.].sum()
-    t_srh = p_srh + n_srh
-
-    return p_srh, n_srh, t_srh
+    return (positive_srh.to('meter ** 2 / second ** 2'),
+            negative_srh.to('meter ** 2 / second ** 2'),
+            (positive_srh + negative_srh).to('meter ** 2 / second ** 2'))

metpy/calc/tests/test_kinematics.py

@@ -7,10 +7,10 @@
 import pytest
 
 from metpy.calc import (advection, convergence_vorticity, frontogenesis, geostrophic_wind,
-                        get_wind_components, h_convergence,
-                        montgomery_streamfunction, shearing_deformation,
-                        shearing_stretching_deformation, storm_relative_helicity,
-                        stretching_deformation, total_deformation, v_vorticity)
+                        get_wind_components, h_convergence, montgomery_streamfunction,
+                        shearing_deformation, shearing_stretching_deformation,
+                        storm_relative_helicity, stretching_deformation, total_deformation,
+                        v_vorticity)
 from metpy.constants import g, omega, Re
 from metpy.testing import assert_almost_equal, assert_array_equal
 from metpy.units import concatenate, units
@@ -394,12 +394,58 @@ def test_streamfunc():
     assert_almost_equal(msf, 337468.2500 * units('m^2 s^-2'), 4)
 
 
-def test_helicity():
-    """Test function for SRH calculations on an eigth-circle hodograph."""
-    pres_test = np.asarray([1013.25, 954.57955706, 898.690770743,
-                            845.481604002, 794.85264282]) * units('hPa')
-    hgt_test = hgt_test = np.asarray([0, 500, 1000, 1500, 2000])
+def test_storm_relative_helicity_no_storm_motion():
+    """Test storm relative helicity with no storm motion and differing input units."""
+    u = np.array([0, 20, 10, 0]) * units('m/s')
+    v = np.array([20, 0, 0, 10]) * units('m/s')
+    u = u.to('knots')
+    heights = np.array([0, 250, 500, 750]) * units.m
+
+    positive_srh, negative_srh, total_srh = storm_relative_helicity(u, v, heights,
+                                                                    depth=750 * units.meters)
+
+    assert_almost_equal(positive_srh, 400. * units('meter ** 2 / second ** 2 '), 6)
+    assert_almost_equal(negative_srh, -100. * units('meter ** 2 / second ** 2 '), 6)
+    assert_almost_equal(total_srh, 300. * units('meter ** 2 / second ** 2 '), 6)
+
+
+def test_storm_relative_helicity_storm_motion():
+    """Test storm relative helicity with storm motion and differing input units."""
+    u = np.array([5, 25, 15, 5]) * units('m/s')
+    v = np.array([30, 10, 10, 20]) * units('m/s')
+    u = u.to('knots')
+    heights = np.array([0, 250, 500, 750]) * units.m
+
+    pos_srh, neg_srh, total_srh = storm_relative_helicity(u, v, heights,
+                                                          depth=750 * units.meters,
+                                                          storm_u=5 * units('m/s'),
+                                                          storm_v=10 * units('m/s'))
+
+    assert_almost_equal(pos_srh, 400. * units('meter ** 2 / second ** 2 '), 6)
+    assert_almost_equal(neg_srh, -100. * units('meter ** 2 / second ** 2 '), 6)
+    assert_almost_equal(total_srh, 300. * units('meter ** 2 / second ** 2 '), 6)
 
+
+def test_storm_relative_helicity_with_interpolation():
+    """Test storm relative helicity with interpolation."""
+    u = np.array([-5, 15, 25, 15, -5]) * units('m/s')
+    v = np.array([40, 20, 10, 10, 30]) * units('m/s')
+    u = u.to('knots')
+    heights = np.array([0, 100, 200, 300, 400]) * units.m
+
+    pos_srh, neg_srh, total_srh = storm_relative_helicity(u, v, heights,
+                                                          bottom=50 * units.meters,
+                                                          depth=300 * units.meters,
+                                                          storm_u=5 * units('m/s'),
+                                                          storm_v=10 * units('m/s'))
+
+    assert_almost_equal(pos_srh, 400. * units('meter ** 2 / second ** 2 '), 6)
+    assert_almost_equal(neg_srh, -100. * units('meter ** 2 / second ** 2 '), 6)
+    assert_almost_equal(total_srh, 300. * units('meter ** 2 / second ** 2 '), 6)
+
+
+def test_storm_relative_helicity():
+    """Test function for SRH calculations on an eigth-circle hodograph."""
     # Create larger arrays for everything except pressure to make a smoother graph
     hgt_int = np.arange(0, 2050, 50)
     hgt_int = hgt_int * units('meter')
@@ -408,10 +454,6 @@ def test_helicity():
     spd_int[:] = 2.
     u_int, v_int = get_wind_components(spd_int * units('m/s'), dir_int * units.degree)
 
-    # Interpolate pressure to that graph
-    pres_int = np.interp(hgt_int, hgt_test, np.log(pres_test.magnitude))
-    pres_int = np.exp(pres_int) * units('hPa')
-
     # Put in the correct value of SRH for a eighth-circle, 2 m/s hodograph
     # (SRH = 2 * area under hodo, in this case...)
     srh_true_p = (.25 * np.pi * (2 ** 2)) * units('m^2/s^2')
@@ -420,7 +462,7 @@ def test_helicity():
     # negative SRH will be zero.
     srh_true_t = srh_true_p
     srh_true_n = 0 * units('m^2/s^2')
-    p_srh, n_srh, T_srh = storm_relative_helicity(u_int, v_int, pres_int,
+    p_srh, n_srh, T_srh = storm_relative_helicity(u_int, v_int,
                                                   hgt_int, 1000 * units('meter'),
                                                   bottom=0 * units('meter'),
                                                   storm_u=0 * units.knot,

metpy/calc/tests/test_tools.py

@@ -7,7 +7,8 @@
 import numpy.ma as ma
 import pytest
 
-from metpy.calc import (find_intersections, get_layer, interp, interpolate_nans, log_interp,
+from metpy.calc import (find_intersections, get_layer, get_layer_heights,
+                        interp, interpolate_nans, log_interp,
                         nearest_intersection_idx, pressure_to_height_std,
                         reduce_point_density, resample_nn_1d)
 from metpy.calc.tools import (_get_bound_pressure_height, _greater_or_close, _less_or_close,
@@ -475,3 +476,49 @@ def test_less_or_close():
     truth = np.array([True, True, True, True, True, False])
     res = _less_or_close(x, comparison_value)
     assert_array_equal(res, truth)
+
+
+def test_get_layer_heights_interpolation():
+    """Test get_layer_heights with interpolation."""
+    heights = np.arange(10) * units.km
+    data = heights.m * 2 * units.degC
+    heights, data = get_layer_heights(heights, 5000 * units.m, data, bottom=1500 * units.m)
+    heights_true = np.array([1.5, 2, 3, 4, 5, 6, 6.5]) * units.km
+    data_true = heights_true.m * 2 * units.degC
+    assert_array_almost_equal(heights_true, heights, 6)
+    assert_array_almost_equal(data_true, data, 6)
+
+
+def test_get_layer_heights_no_interpolation():
+    """Test get_layer_heights without interpolation."""
+    heights = np.arange(10) * units.km
+    data = heights.m * 2 * units.degC
+    heights, data = get_layer_heights(heights, 5000 * units.m, data,
+                                      bottom=1500 * units.m, interpolate=False)
+    heights_true = np.array([2, 3, 4, 5, 6]) * units.km
+    data_true = heights_true.m * 2 * units.degC
+    assert_array_almost_equal(heights_true, heights, 6)
+    assert_array_almost_equal(data_true, data, 6)
+
+
+def test_get_layer_heights_agl():
+    """Test get_layer_heights with interpolation."""
+    heights = np.arange(300, 1200, 100) * units.m
+    data = heights.m * 0.1 * units.degC
+    heights, data = get_layer_heights(heights, 500 * units.m, data, with_agl=True)
+    heights_true = np.array([0, 0.1, 0.2, 0.3, 0.4, 0.5]) * units.km
+    data_true = np.array([30, 40, 50, 60, 70, 80]) * units.degC
+    assert_array_almost_equal(heights_true, heights, 6)
+    assert_array_almost_equal(data_true, data, 6)
+
+
+def test_get_layer_heights_agl_bottom_no_interp():
+    """Test get_layer_heights with no interpolation and a bottom."""
+    heights = np.arange(300, 1200, 100) * units.m
+    data = heights.m * 0.1 * units.degC
+    heights, data = get_layer_heights(heights, 500 * units.m, data, with_agl=True,
+                                      interpolation=False, bottom=200 * units.m)
+    heights_true = np.array([0.2, 0.3, 0.4, 0.5, 0.6, 0.7]) * units.km
+    data_true = np.array([50, 60, 70, 80, 90, 100]) * units.degC
+    assert_array_almost_equal(heights_true, heights, 6)
+    assert_array_almost_equal(data_true, data, 6)

metpy/calc/tools.py

@@ -394,6 +394,97 @@ def _get_bound_pressure_height(pressure, bound, heights=None, interpolate=True):
     return bound_pressure, bound_height
 
 
+@exporter.export
+@check_units('[length]')
+def get_layer_heights(heights, depth, *args, **kwargs):
+    """Return an atmospheric layer from upper air data with the requested bottom and depth.
+
+    This function will subset an upper air dataset to contain only the specified layer using
+    the heights only.
+
+    Parameters
+    ----------
+    heights : array-like
+        Atmospheric heights
+    depth : `pint.Quantity`
+        The thickness of the layer
+    *args : array-like
+        Atmospheric variable(s) measured at the given pressures
+    bottom : `pint.Quantity`, optional
+        The bottom of the layer
+    interpolate : bool, optional
+        Interpolate the top and bottom points if they are not in the given data. Defaults
+        to True.
+    with_agl : bool, optional
+        Returns the heights as above ground level by subtracting the minimum height in the
+        provided heights. Defaults to False.
+
+    Returns
+    -------
+    `pint.Quantity, pint.Quantity`
+        The height and data variables of the layer
+
+    """
+    bottom = kwargs.pop('bottom', None)
+    interpolate = kwargs.pop('interpolate', True)
+    with_agl = kwargs.pop('with_agl', False)
+
+    # Make sure pressure and datavars are the same length
+    for datavar in args:
+        if len(heights) != len(datavar):
+            raise ValueError('Height and data variables must have the same length.')
+
+    # If we want things in AGL, subtract the minimum height from all height values
+    if with_agl:
+        sfc_height = np.min(heights)
+        heights -= sfc_height
+
+    # If the bottom is not specified, make it the surface
+    if bottom is None:
+        bottom = heights[0]
+
+    # Make heights and arguments base units
+    heights = heights.to_base_units()
+    bottom = bottom.to_base_units()
+
+    # Calculate the top of the layer
+    top = bottom + depth
+
+    ret = []  # returned data variables in layer
+
+    # Ensure heights are sorted in ascending order
+    sort_inds = np.argsort(heights)
+    heights = heights[sort_inds]
+
+    # Mask based on top and bottom
+    inds = (heights >= bottom) & (heights <= top)
+    heights_interp = heights[inds]
+
+    # Interpolate heights at bounds if necessary and sort
+    if interpolate:
+        # If we don't have the bottom or top requested, append them
+        if top not in heights_interp:
+            heights_interp = np.sort(np.append(heights_interp, top)) * heights.units
+        if bottom not in heights_interp:
+            heights_interp = np.sort(np.append(heights_interp, bottom)) * heights.units
+
+    ret.append(heights_interp)
+
+    for datavar in args:
+        # Ensure that things are sorted in ascending order
+        datavar = datavar[sort_inds]
+
+        if interpolate:
+            # Interpolate for the possibly missing bottom/top values
+            datavar_interp = interp(heights_interp, heights, datavar)
+            datavar = datavar_interp
+        else:
+            datavar = datavar[inds]
+
+        ret.append(datavar)
+    return ret
+
+
 @exporter.export
 @check_units('[pressure]')
 def get_layer(pressure, *args, **kwargs):