Height Scale for SkewT

src/metpy/plots/skewt.py

@@ -17,13 +17,15 @@
 from matplotlib.patches import Circle
 from matplotlib.projections import register_projection
 import matplotlib.spines as mspines
-from matplotlib.ticker import MultipleLocator, NullFormatter, ScalarFormatter
+from matplotlib.ticker import (MultipleLocator, FixedLocator, NullFormatter,
+                               ScalarFormatter, NullLocator)
 import matplotlib.transforms as transforms
 import numpy as np
 
 from ._util import colored_line
-from ..calc import dewpoint, dry_lapse, el, lcl, moist_lapse, vapor_pressure
+from ..calc import (dewpoint, dry_lapse, el, lcl, moist_lapse, vapor_pressure,
+                    pressure_to_height_std, height_to_pressure_std)
 from ..calc.tools import _delete_masked_points
 from ..interpolate import interpolate_1d
 from ..package_tools import Exporter
 from ..units import concatenate, is_quantity, units
@@ -261,7 +263,8 @@
 
     """
 
-    def __init__(self, fig=None, rotation=30, subplot=None, rect=None, aspect=80.5):
+    def __init__(self, fig=None, rotation=30, subplot=None, rect=None,
+                 aspect=80.5, show_heights=False):
         r"""Create SkewT - logP plots.
 
         Parameters
@@ -286,6 +289,10 @@
             Aspect ratio (i.e. ratio of y-scale to x-scale) to maintain in the plot.
             Defaults to 80.5. Passing the string ``'auto'`` tells matplotlib to handle
             the aspect ratio automatically (this is not recommended for SkewT).
+        show_heights : boolean, optional
+            Flag for showing heights as a secondary y axis.
+            Calculated from pressure_to_height_std.
+            (defaults to false)
 
         """
         if fig is None:
@@ -335,6 +342,28 @@
         self.dry_adiabats = None
         self.moist_adiabats = None
 
+        if (show_heights):
+            # Set a secondary axis with height from pressure_to_height_standard
+            # Requires direct and inverse fctns - pressure axis and height axis
+            def pressure_axis(p):
+                return pressure_to_height_std(units.Quantity(p, 'hPa')).m_as('km')
+
+            def height_axis(h):
+                return height_to_pressure_std(units.Quantity(h, 'km')).m
+            # Positions the axis .12 normalized units to the left of the pressure axis
+            self.heightax = self.ax.secondary_yaxis(-0.12,
+                                                    functions=(pressure_axis, height_axis))
+            # Set ylim based on pressure limits
+            self.heightax.set_ylim(pressure_to_height_std(units.Quantity
+                                                          (self.ax.get_ylim(), 'hPa')))
+            self.heightax.yaxis.set_units(units.km)
+            self.heightax.yaxis.set_minor_locator(NullLocator())
+            self.heightax.yaxis.set_major_formatter(ScalarFormatter())
+            # Create ticks on the height axis counting by 1 from min to max
+            ymin, ymax = self.heightax.get_ylim()
+            yticks = np.arange(ymin, ymax + 1, 1)
+            self.heightax.yaxis.set_major_locator(FixedLocator(yticks))
+
         # Maintain a reasonable ratio of data limits.
         self.ax.set_aspect(aspect, adjustable='box')
 

tests/plots/test_skewt.py

@@ -46,6 +46,49 @@ def test_skewt_api():
     return fig
 
 
+@pytest.mark.mpl_image_compare(style='default', tolerance=0.069)
+def test_skewt_api_with_heights():
+    """Test the SkewT API with height axis."""
+    with matplotlib.rc_context({'axes.autolimit_mode': 'data'}):
+        fig = plt.figure(figsize=(9, 9))
+        skew = SkewT(fig, aspect='auto', show_heights=True)
+
+        # Plot the data using normal plotting functions, in this case using
+        # log scaling in Y, as dictated by the typical meteorological plot
+        p = np.linspace(1000, 100, 10)
+        t = np.linspace(20, -20, 10)
+        u = np.linspace(-10, 10, 10)
+        skew.plot(p, t, 'r')
+        skew.plot_barbs(p, u, u)
+
+        skew.ax.set_xlim(-20, 30)
+        skew.ax.set_ylim(1000, 100)
+
+        # Add the relevant special lines
+        skew.plot_dry_adiabats()
+        skew.plot_moist_adiabats()
+        skew.plot_mixing_lines()
+
+        # Call again to hit removal statements
+        skew.plot_dry_adiabats()
+        skew.plot_moist_adiabats()
+        skew.plot_mixing_lines()
+
+        # You can't remove text from a secax with remove_text so do it manually
+        skew.heightax.set_ylabel('')
+        skew.heightax.set_yticklabels([])
+        skew.ax.set_title('')
+        skew.ax.set_xlabel('')
+        skew.ax.set_ylabel('')
+        skew.ax.set_xticklabels([])
+        skew.ax.set_yticklabels([])
+
+        # prevents label from being cut off by savefig
+        plt.tight_layout()
+
+    return fig
+
+
 @pytest.mark.mpl_image_compare(remove_text=True, style='default', tolerance=0.32)
 def test_skewt_api_units():
     """Test the SkewT API when units are provided."""
@@ -70,6 +113,42 @@ def test_skewt_api_units():
     return fig
 
 
+@pytest.mark.mpl_image_compare(style='default', tolerance=.32)
+def test_skewt_api_units_heights():
+    """Test the SkewT API when units are provided and show_heights is true."""
+    with matplotlib.rc_context({'axes.autolimit_mode': 'data'}):
+        fig = plt.figure(figsize=(9, 9))
+        skew = SkewT(fig, show_heights=True)
+        p = (np.linspace(950, 100, 10) * units.hPa).to(units.Pa)
+        t = (np.linspace(18, -20, 10) * units.degC).to(units.kelvin)
+        u = np.linspace(-20, 20, 10) * units.knots
+
+        skew.plot(p, t, 'r')
+        skew.plot_barbs(p, u, u)
+
+        # Add the relevant special lines
+        skew.plot_dry_adiabats()
+        skew.plot_moist_adiabats()
+        skew.plot_mixing_lines()
+
+        # This works around the fact that newer pint versions default to degrees_Celsius
+        skew.ax.set_xlabel('degC')
+
+        # You can't remove text from a secax with remove_text so do it manually
+        skew.heightax.set_ylabel('')
+        skew.heightax.set_yticklabels([])
+        skew.ax.set_title('')
+        skew.ax.set_xlabel('')
+        skew.ax.set_ylabel('')
+        skew.ax.set_xticklabels([])
+        skew.ax.set_yticklabels([])
+
+        # Prevents labels from being cut off by savefig
+        plt.tight_layout()
+
+    return fig
+
+
 @pytest.mark.mpl_image_compare(tolerance=0., remove_text=True, style='default')
 def test_skewt_default_aspect_empty():
     """Test SkewT with default aspect and no plots, only special lines."""
@@ -113,6 +192,24 @@ def test_skewt_subplot():
     return fig
 
 
+@pytest.mark.mpl_image_compare(tolerance=.811, style='default')
+def test_skewt_subplot_heights():
+    """Test using skewT on a sub-plot with height axis."""
+    fig = plt.figure(figsize=(9, 9))
+    skew = SkewT(fig, subplot=(2, 2, 1), aspect='auto', show_heights=True)
+
+    # You can't remove text from a secax with remove_text so do it manually
+    skew.heightax.set_ylabel('')
+    skew.heightax.set_yticklabels([])
+    skew.ax.set_title('')
+    skew.ax.set_xlabel('')
+    skew.ax.set_ylabel('')
+    skew.ax.set_xticklabels([])
+    skew.ax.set_yticklabels([])
+
+    return fig
+
+
 @pytest.mark.mpl_image_compare(tolerance=0, remove_text=True, style='default')
 def test_skewt_gridspec():
     """Test using SkewT on a GridSpec sub-plot."""
@@ -130,6 +227,24 @@ def test_skewt_with_grid_enabled():
         plt.close(s.ax.figure)
 
 
+@pytest.mark.mpl_image_compare(tolerance=0, style='default')
+def test_skewt_gridspec_heights():
+    """Test using SkewT on a GridSpec sub-plot."""
+    fig = plt.figure(figsize=(9, 9))
+    gs = GridSpec(1, 2)
+    skew = SkewT(fig, subplot=gs[0, 1], aspect='auto', show_heights=True)
+
+    # You can't remove text from a secax with remove_text so do it manually
+    skew.heightax.set_ylabel('')
+    skew.heightax.set_yticklabels([])
+    skew.ax.set_title('')
+    skew.ax.set_xlabel('')
+    skew.ax.set_ylabel('')
+    skew.ax.set_xticklabels([])
+    skew.ax.set_yticklabels([])
+    return fig
+
+
 @pytest.mark.mpl_image_compare(tolerance=0., remove_text=True, style='default')
 def test_skewt_arbitrary_rect():
     """Test placing the SkewT in an arbitrary rectangle."""
@@ -144,6 +259,23 @@ def test_skewt_subplot_rect_conflict():
         SkewT(fig, rect=(0.15, 0.35, 0.8, 0.3), subplot=(1, 1, 1))
 
 
+@pytest.mark.mpl_image_compare(tolerance=0., style='default')
+def test_skewt_arbitrary_rect_heights():
+    """Test placing the SkewT in an arbitrary rectangle with height axis."""
+    fig = plt.figure(figsize=(9, 9))
+    skew = SkewT(fig, rect=(0.15, 0.35, 0.8, 0.3), aspect='auto', show_heights=True)
+
+    # You can't remove text from a secax with remove_text so do it manually
+    skew.heightax.set_ylabel('')
+    skew.heightax.set_yticklabels([])
+    skew.ax.set_title('')
+    skew.ax.set_xlabel('')
+    skew.ax.set_ylabel('')
+    skew.ax.set_xticklabels([])
+    skew.ax.set_yticklabels([])
+    return fig
+
+
 @pytest.mark.mpl_image_compare(tolerance=0.0198, remove_text=True, style='default')
 def test_skewt_units():
     """Test that plotting with SkewT works with units properly."""
@@ -308,6 +440,31 @@ def test_skewt_wide_aspect_ratio(test_profile):
     return fig
 
 
+@pytest.mark.mpl_image_compare(tolerance=0.039, style='default')
+def test_skewt_wide_aspect_ratio_heights(test_profile):
+    """Test plotting a skewT with a wide aspect ratio with height axis."""
+    p, t, _, tp = test_profile
+
+    fig = plt.figure(figsize=(12.5, 3))
+    skew = SkewT(fig, aspect='auto', show_heights=True)
+    skew.plot(p, t, 'r')
+    skew.plot(p, tp, 'k')
+    skew.ax.set_xlim(-30, 50)
+    skew.ax.set_ylim(1050, 700)
+
+    # You can't remove text from a secax with remove_text so do it manually
+    skew.heightax.set_ylabel('')
+    skew.heightax.set_yticklabels([])
+    skew.ax.set_title('')
+    skew.ax.set_xlabel('')
+    skew.ax.set_xticklabels([])
+    skew.ax.set_yticklabels([])
+
+    # This works around the fact that newer pint versions default to degrees_Celsius
+    skew.ax.set_xlabel('degC')
+    return fig
+
+
 @pytest.mark.mpl_image_compare(tolerance=0, remove_text=True)
 def test_hodograph_api():
     """Basic test of Hodograph API."""