Height Scale for SkewT

src/metpy/plots/skewt.py

@@ -17,12 +17,13 @@
 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, ScalarFormatter, NullLocator, NullFormatter
 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
@@ -286,7 +287,6 @@
             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).
-
         """
         if fig is None:
             import matplotlib.pyplot as plt
@@ -336,7 +336,7 @@
         self.moist_adiabats = None
 
         # Maintain a reasonable ratio of data limits.
-        self.ax.set_aspect(aspect, adjustable='box')
+        self.ax.set_aspect(aspect, adjustable='datalim')
 
     def plot(self, pressure, t, *args, **kwargs):
         r"""Plot data.
@@ -743,6 +743,49 @@
         return self.shade_area(pressure[idx], t_parcel[idx], t[idx], which='negative',
                                **kwargs)
 
+    def update_heightax_limits(self, ax):
+        self.heightax.set_ylim(pressure_to_height_std(units.Quantity
+                                                      (self.ax.get_ylim(), 'hPa')))
+
+    def add_heightax(self):
+        r"""Add a secondary y axis with height values calculated from pressure_to_height_std.
+
+        Axis is created to .12 normalized units to the left of the pressure axis and can
+        be accessed with the name "heightax".
+
+        See Also
+        --------
+        :meth:`metpy.calc.pressure_to_height_std`
+
+        """
+        # Set a secondary axis with height from pressure_to_height_standard
+        # Requires direct and inverse fctns - pressure axis and height axis
+        def height_axis(p):
+            return pressure_to_height_std(units.Quantity(p, 'hPa')).m_as('km')
+
+        def pressure_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.twinx()
+        self.heightax.set_ylim(pressure_to_height_std(units.Quantity
+                                                      (self.ax.get_ylim(), 'hPa')))
+        self.heightax.spines["left"].set_position(("axes", -.12))
+        self.heightax.spines["left"].set_visible(True)
+        self.heightax.yaxis.set_label_position('left')
+        self.heightax.yaxis.set_ticks_position('left')
+        # Set height axis ylims based on pressure ylims
+        # This doesn't really seem to do anything except make it so that the unit is shown
+        # on the axis
+        self.heightax.set_yscale('linear')
+        self.heightax.yaxis.set_units(units.km)
+        self.heightax.yaxis.set_minor_locator(NullLocator())
+        self.heightax.yaxis.set_major_formatter(ScalarFormatter())
+        self.heightax.yaxis.set_major_locator(MultipleLocator(1))
+        self.ax.callbacks.connect('ylim_changed', self.update_heightax_limits)
+
+
+
+
 
 @exporter.export
 class Stuve(SkewT):

tests/plots/test_skewt.py

@@ -46,6 +46,50 @@ def test_skewt_api():
     return fig
 
 
+@pytest.mark.mpl_image_compare(style='default', tolerance=0.071)
+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')
+        skew.add_heightax()
+
+        # 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 +114,43 @@ 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 a secondary height axis is added."""
+    with matplotlib.rc_context({'axes.autolimit_mode': 'data'}):
+        fig = plt.figure(figsize=(9, 9))
+        skew = SkewT(fig)
+        skew.add_heightax()
+        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 +194,25 @@ 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')
+    skew.add_heightax()
+
+    # 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 +230,25 @@ 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 with a height axis."""
+    fig = plt.figure(figsize=(9, 9))
+    gs = GridSpec(1, 2)
+    skew = SkewT(fig, subplot=gs[0, 1], aspect='auto')
+    skew.add_heightax()
+
+    # 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 +263,24 @@ 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')
+    skew.add_heightax()
+
+    # 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."""
@@ -168,6 +305,38 @@ def test_skewt_units():
     return fig
 
 
+@pytest.mark.mpl_image_compare(tolerance=.069, style='default')
+def test_skewt_height_change_coords():
+    """Test plotting a skewt with a height axis then changing the pressure axis."""
+    fig = plt.figure(figsize=(9, 9))
+    skew = SkewT(fig, aspect='auto')
+    skew.add_heightax()
+
+    skew.ax.set_ylim(500, 100)
+
+    skew.ax.set_xticklabels([])
+    skew.ax.set_yticklabels([])
+
+    expected_ylim = np.array([15, 5])
+
+    # This updates the plot so get_ylim is accurate
+    fig.canvas.draw()
+
+    # Asserts that the ylims on height ax are as expected
+    assert np.array_equal(np.int64(skew.heightax.get_ylim()), expected_ylim)
+
+    # 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.fixture()
 def test_profile():
     """Return data for a test profile."""
@@ -308,6 +477,32 @@ 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')
+    skew.add_heightax()
+    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."""