Merge pull request #572 from jrleeman/Hodograph_Coloring_Units

Hodograph coloring units

metpy/plots/skewt.py

@@ -22,7 +22,7 @@
 from ..calc import dewpoint, dry_lapse, moist_lapse, vapor_pressure
 from ..calc.tools import delete_masked_points, interp
 from ..package_tools import Exporter
-from ..units import units
+from ..units import concatenate, units
 
 exporter = Exporter(globals())
 
@@ -812,15 +812,18 @@ def plot_colormapped(self, u, v, c, bounds=None, colors=None, **kwargs):
         """
         u, v, c = delete_masked_points(u, v, c)
 
+        # Plotting a color segmented hodograph
         if colors:
             cmap = mcolors.ListedColormap(colors)
+            # If we are segmenting by height (a length), interpolate the bounds
             if bounds.dimensionality == {'[length]': 1.0}:
-                bounds = np.asarray(bounds + c[0]) * bounds.units
-                interp_vert = interp(bounds, c, c, u, v)
-                inds = np.searchsorted(c.magnitude, bounds.magnitude)
-                u = np.insert(u.magnitude, inds, interp_vert[1].magnitude)
-                v = np.insert(v.magnitude, inds, interp_vert[2].magnitude)
-                c = np.insert(c.magnitude, inds, interp_vert[0].magnitude)
+                bounds = bounds + c[0]  # Make all heights AGL
+                interpolation_heights = concatenate((bounds, c))
+                interpolation_heights = (np.sort(interpolation_heights) *
+                                         interpolation_heights.units)
+                c, u, v = interp(interpolation_heights, c, c, u, v)
+                c = c.to_base_units()  # TODO: This shouldn't be required!
+            # If segmenting by anything else, do not interpolate, just use the data
             else:
                 bounds = np.asarray(bounds) * bounds.units
 
@@ -830,8 +833,12 @@ def plot_colormapped(self, u, v, c, bounds=None, colors=None, **kwargs):
             kwargs['cmap'] = cmap
             kwargs['norm'] = norm
             line_args = self._form_line_args(kwargs)
+
+        # Plotting a continuously colored line
         else:
             line_args = self._form_line_args(kwargs)
+
+        # Do the plotting
         lc = colored_line(u, v, c, **line_args)
         self.ax.add_collection(lc)
         return lc

metpy/plots/tests/test_skewt.py

@@ -185,18 +185,37 @@ def test_skewt_barb_color():
 @pytest.mark.mpl_image_compare(tolerance=0, remove_text=True)
 def test_hodograph_plot_layers():
     """Test hodograph colored height layers with interpolation."""
-    u = np.arange(5, 65, 5) * units('knot')
-    v = np.arange(-5, -65, -5) * units('knot')
-    h = [178, 213, 610, 656, 721, 914, 1060,
-         1219, 1372, 1412, 1512, 1524] * units('meter')
-    colors = ['red', 'green']
-    levels = [0, 500, 1000] * units('meter')
-    fig = plt.figure(figsize=(9, 9))
-    ax = fig.add_subplot(1, 1, 1)
-    hodo = Hodograph(ax, component_range=80)
-    hodo.add_grid(increment=20, color='k')
-    hodo.plot_colormapped(u, v, h, bounds=levels, colors=colors)
+    u = np.zeros((6)) * units.knots
+    v = np.array([0, 10, 20, 30, 40, 50]) * units.knots
+    heights = np.array([0, 1000, 2000, 3000, 4000, 5000]) * units.m
+    bounds = np.array([500, 1500, 2500, 3500, 4500]) * units.m
+    colors = ['r', 'g', 'b', 'r']
+    fig = plt.figure(figsize=(7, 7))
+    ax1 = fig.add_subplot(1, 1, 1)
+    h = Hodograph(ax1)
+    h.add_grid(increment=10)
+    h.plot_colormapped(u, v, heights, colors=colors, bounds=bounds)
+    ax1.set_xlim(-50, 50)
+    ax1.set_ylim(-5, 50)
+
+    return fig
 
+
+@pytest.mark.mpl_image_compare(tolerance=0, remove_text=True)
+def test_hodograph_plot_layers_different_units():
+    """Test hodograph colored height layers with interpolation and different units."""
+    u = np.zeros((6)) * units.knots
+    v = np.array([0, 10, 20, 30, 40, 50]) * units.knots
+    heights = np.array([0, 1, 2, 3, 4, 5]) * units.km
+    bounds = np.array([500, 1500, 2500, 3500, 4500]) * units.m
+    colors = ['r', 'g', 'b', 'r']
+    fig = plt.figure(figsize=(7, 7))
+    ax1 = fig.add_subplot(1, 1, 1)
+    h = Hodograph(ax1)
+    h.add_grid(increment=10)
+    h.plot_colormapped(u, v, heights, colors=colors, bounds=bounds)
+    ax1.set_xlim(-50, 50)
+    ax1.set_ylim(-5, 50)
     return fig