Add pytest for Stuve and Emagram

New test files tests/plots/test_emagram.py tests/plots/test_stuve.py replicate the SkewT tests done in tests/plots/test_skewt.py Whitespace changes for linter. Signed-off-by: Christopher Holmes <cdholmes@fsu.edu>
Unidata · Jan 9, 2025 · 24d925f · 24d925f
1 parent e2a0d93
commit 24d925f
Show file tree

Hide file tree

Showing 21 changed files with 311 additions and 1 deletion.
diff --git a/tests/plots/baseline/test_emagram_api.png b/tests/plots/baseline/test_emagram_api.png
diff --git a/tests/plots/baseline/test_emagram_arbitrary_rect.png b/tests/plots/baseline/test_emagram_arbitrary_rect.png
diff --git a/tests/plots/baseline/test_emagram_default_aspect_empty.png b/tests/plots/baseline/test_emagram_default_aspect_empty.png
diff --git a/tests/plots/baseline/test_emagram_gridspec.png b/tests/plots/baseline/test_emagram_gridspec.png
diff --git a/tests/plots/baseline/test_emagram_mixing_line_args.png b/tests/plots/baseline/test_emagram_mixing_line_args.png
diff --git a/tests/plots/baseline/test_emagram_shade_area.png b/tests/plots/baseline/test_emagram_shade_area.png
diff --git a/tests/plots/baseline/test_emagram_shade_cape_cin.png b/tests/plots/baseline/test_emagram_shade_cape_cin.png
diff --git a/tests/plots/baseline/test_emagram_subplot.png b/tests/plots/baseline/test_emagram_subplot.png
diff --git a/tests/plots/baseline/test_emagram_tight_bbox.png b/tests/plots/baseline/test_emagram_tight_bbox.png
diff --git a/tests/plots/baseline/test_emagram_wide_aspect_ratio.png b/tests/plots/baseline/test_emagram_wide_aspect_ratio.png
diff --git a/tests/plots/baseline/test_stuve_api.png b/tests/plots/baseline/test_stuve_api.png
diff --git a/tests/plots/baseline/test_stuve_arbitrary_rect.png b/tests/plots/baseline/test_stuve_arbitrary_rect.png
diff --git a/tests/plots/baseline/test_stuve_default_aspect_empty.png b/tests/plots/baseline/test_stuve_default_aspect_empty.png
diff --git a/tests/plots/baseline/test_stuve_gridspec.png b/tests/plots/baseline/test_stuve_gridspec.png
diff --git a/tests/plots/baseline/test_stuve_mixing_line_args.png b/tests/plots/baseline/test_stuve_mixing_line_args.png
diff --git a/tests/plots/baseline/test_stuve_shade_area.png b/tests/plots/baseline/test_stuve_shade_area.png
diff --git a/tests/plots/baseline/test_stuve_shade_cape_cin.png b/tests/plots/baseline/test_stuve_shade_cape_cin.png
diff --git a/tests/plots/baseline/test_stuve_subplot.png b/tests/plots/baseline/test_stuve_subplot.png
diff --git a/tests/plots/baseline/test_stuve_tight_bbox.png b/tests/plots/baseline/test_stuve_tight_bbox.png
diff --git a/tests/plots/baseline/test_stuve_wide_aspect_ratio.png b/tests/plots/baseline/test_stuve_wide_aspect_ratio.png
diff --git a/tests/plots/test_skewt.py b/tests/plots/test_skewt.py
@@ -10,7 +10,7 @@
 import numpy as np
 import pytest
 
-from metpy.plots import Hodograph, SkewT
+from metpy.plots import Emagram, Hodograph, SkewT, Stuve
 from metpy.testing import autoclose_figure, version_check
 from metpy.units import units
 
@@ -576,3 +576,313 @@ def test_hodograph_range_with_units():
     with autoclose_figure(figsize=(6, 6)) as fig:
         ax = fig.add_subplot(1, 1, 1)
         Hodograph(ax, component_range=60. * units.knots)
+
+
+@pytest.mark.mpl_image_compare(remove_text=True, style='default', tolerance=0.069)
+def test_stuve_api():
+    """Test the Stuve API."""
+    with matplotlib.rc_context({'axes.autolimit_mode': 'data'}):
+        fig = plt.figure(figsize=(9, 9))
+        stuve = Stuve(fig, aspect='auto')
+
+        # 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)
+        stuve.plot(p, t, 'r')
+        stuve.plot_barbs(p, u, u)
+
+        stuve.ax.set_xlim(-20, 30)
+        stuve.ax.set_ylim(1000, 100)
+
+        # Add the relevant special lines
+        stuve.plot_dry_adiabats()
+        stuve.plot_moist_adiabats()
+        stuve.plot_mixing_lines()
+
+        # Call again to hit removal statements
+        stuve.plot_dry_adiabats()
+        stuve.plot_moist_adiabats()
+        stuve.plot_mixing_lines()
+
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0., remove_text=True, style='default')
+def test_stuve_default_aspect_empty():
+    """Test Stuve with default aspect and no plots, only special lines."""
+    fig = plt.figure(figsize=(12, 9))
+    stuve = Stuve(fig)
+    stuve.plot_dry_adiabats()
+    stuve.plot_moist_adiabats()
+    stuve.plot_mixing_lines()
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0., remove_text=True, style='default')
+def test_stuve_mixing_line_args():
+    """Test plot_mixing_lines accepting kwargs for mixing ratio and pressure levels."""
+    fig = plt.figure(figsize=(12, 9))
+    stuve = Stuve(fig)
+    mlines = np.array([0.0004, 0.001, 0.002, 0.004, 0.007, 0.01, 0.016, 0.024, 0.032])
+    press = units.Quantity(np.linspace(600, max(stuve.ax.get_ylim())), 'mbar')
+    stuve.plot_dry_adiabats()
+    stuve.plot_moist_adiabats()
+    stuve.plot_mixing_lines(mixing_ratio=mlines, pressure=press)
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0., remove_text=False, style='default',
+                               savefig_kwargs={'bbox_inches': 'tight'})
+def test_stuve_tight_bbox():
+    """Test Stuve when saved with `savefig(..., bbox_inches='tight')`."""
+    fig = plt.figure(figsize=(12, 9))
+    Stuve(fig)
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0.811, remove_text=True, style='default')
+def test_stuve_subplot():
+    """Test using Stuve on a sub-plot."""
+    fig = plt.figure(figsize=(9, 9))
+    Stuve(fig, subplot=(2, 2, 1), aspect='auto')
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0, remove_text=True, style='default')
+def test_stuve_gridspec():
+    """Test using Stuve on a GridSpec sub-plot."""
+    fig = plt.figure(figsize=(9, 9))
+    gs = GridSpec(1, 2)
+    Stuve(fig, subplot=gs[0, 1], aspect='auto')
+    return fig
+
+
+def test_stuve_with_grid_enabled():
+    """Test using Stuve when gridlines are already enabled (#271)."""
+    with plt.rc_context(rc={'axes.grid': True}):
+        # Also tests when we don't pass in Figure
+        s = Stuve(aspect='auto')
+        plt.close(s.ax.figure)
+
+
+@pytest.mark.mpl_image_compare(tolerance=0., remove_text=True, style='default')
+def test_stuve_arbitrary_rect():
+    """Test placing the Stuve in an arbitrary rectangle."""
+    fig = plt.figure(figsize=(9, 9))
+    Stuve(fig, rect=(0.15, 0.35, 0.8, 0.3), aspect='auto')
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=.033, remove_text=True, style='default')
+def test_stuve_shade_cape_cin(test_profile):
+    """Test shading CAPE and CIN on a Stuve plot."""
+    p, t, td, tp = test_profile
+
+    with matplotlib.rc_context({'axes.autolimit_mode': 'data'}):
+        fig = plt.figure(figsize=(9, 9))
+        stuve = Stuve(fig, aspect='auto')
+        stuve.plot(p, t, 'r')
+        stuve.plot(p, tp, 'k')
+        stuve.shade_cape(p, t, tp)
+        stuve.shade_cin(p, t, tp, td)
+        stuve.ax.set_xlim(-50, 50)
+        stuve.ax.set_ylim(1000, 100)
+
+        # This works around the fact that newer pint versions default to degrees_Celsius
+        stuve.ax.set_xlabel('degC')
+
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0.033, remove_text=True, style='default')
+def test_stuve_shade_area(test_profile):
+    """Test shading areas on a Stuve plot."""
+    p, t, _, tp = test_profile
+
+    with matplotlib.rc_context({'axes.autolimit_mode': 'data'}):
+        fig = plt.figure(figsize=(9, 9))
+        stuve = Stuve(fig, aspect='auto')
+        stuve.plot(p, t, 'r')
+        stuve.plot(p, tp, 'k')
+        stuve.shade_area(p, t, tp)
+        stuve.ax.set_xlim(-50, 50)
+        stuve.ax.set_ylim(1000, 100)
+
+        # This works around the fact that newer pint versions default to degrees_Celsius
+        stuve.ax.set_xlabel('degC')
+
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0.039, remove_text=True, style='default')
+def test_stuve_wide_aspect_ratio(test_profile):
+    """Test plotting a stuveT with a wide aspect ratio."""
+    p, t, _, tp = test_profile
+
+    fig = plt.figure(figsize=(12.5, 3))
+    stuve = Stuve(fig, aspect='auto')
+    stuve.plot(p, t, 'r')
+    stuve.plot(p, tp, 'k')
+    stuve.ax.set_xlim(-30, 50)
+    stuve.ax.set_ylim(1050, 700)
+
+    # This works around the fact that newer pint versions default to degrees_Celsius
+    stuve.ax.set_xlabel('degC')
+    return fig
+
+
+@pytest.mark.mpl_image_compare(remove_text=True, style='default', tolerance=0.069)
+def test_emagram_api():
+    """Test the Emagram API."""
+    with matplotlib.rc_context({'axes.autolimit_mode': 'data'}):
+        fig = plt.figure(figsize=(9, 9))
+        emagram = Emagram(fig, aspect='auto')
+
+        # 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)
+        emagram.plot(p, t, 'r')
+        emagram.plot_barbs(p, u, u)
+
+        emagram.ax.set_xlim(-20, 30)
+        emagram.ax.set_ylim(1000, 100)
+
+        # Add the relevant special lines
+        emagram.plot_dry_adiabats()
+        emagram.plot_moist_adiabats()
+        emagram.plot_mixing_lines()
+
+        # Call again to hit removal statements
+        emagram.plot_dry_adiabats()
+        emagram.plot_moist_adiabats()
+        emagram.plot_mixing_lines()
+
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0.001, remove_text=True, style='default')
+def test_emagram_default_aspect_empty():
+    """Test Emagram with default aspect and no plots, only special lines."""
+    fig = plt.figure(figsize=(12, 9))
+    emagram = Emagram(fig)
+    emagram.plot_dry_adiabats()
+    emagram.plot_moist_adiabats()
+    emagram.plot_mixing_lines()
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0.001, remove_text=True, style='default')
+def test_emagram_mixing_line_args():
+    """Test plot_mixing_lines accepting kwargs for mixing ratio and pressure levels."""
+    fig = plt.figure(figsize=(12, 9))
+    emagram = Emagram(fig)
+    mlines = np.array([0.0004, 0.001, 0.002, 0.004, 0.007, 0.01, 0.016, 0.024, 0.032])
+    press = units.Quantity(np.linspace(600, max(emagram.ax.get_ylim())), 'mbar')
+    emagram.plot_dry_adiabats()
+    emagram.plot_moist_adiabats()
+    emagram.plot_mixing_lines(mixing_ratio=mlines, pressure=press)
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0., remove_text=False, style='default',
+                               savefig_kwargs={'bbox_inches': 'tight'})
+def test_emagram_tight_bbox():
+    """Test Emagram when saved with `savefig(..., bbox_inches='tight')`."""
+    fig = plt.figure(figsize=(12, 9))
+    Emagram(fig)
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0.811, remove_text=True, style='default')
+def test_emagram_subplot():
+    """Test using Emagram on a sub-plot."""
+    fig = plt.figure(figsize=(9, 9))
+    Emagram(fig, subplot=(2, 2, 1), aspect='auto')
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0, remove_text=True, style='default')
+def test_emagram_gridspec():
+    """Test using Emagram on a GridSpec sub-plot."""
+    fig = plt.figure(figsize=(9, 9))
+    gs = GridSpec(1, 2)
+    Emagram(fig, subplot=gs[0, 1], aspect='auto')
+    return fig
+
+
+def test_emagram_with_grid_enabled():
+    """Test using Emagram when gridlines are already enabled (#271)."""
+    with plt.rc_context(rc={'axes.grid': True}):
+        # Also tests when we don't pass in Figure
+        s = Emagram(aspect='auto')
+        plt.close(s.ax.figure)
+
+
+@pytest.mark.mpl_image_compare(tolerance=0., remove_text=True, style='default')
+def test_emagram_arbitrary_rect():
+    """Test placing the Emagram in an arbitrary rectangle."""
+    fig = plt.figure(figsize=(9, 9))
+    Emagram(fig, rect=(0.15, 0.35, 0.8, 0.3), aspect='auto')
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=.033, remove_text=True, style='default')
+def test_emagram_shade_cape_cin(test_profile):
+    """Test shading CAPE and CIN on an Emagram plot."""
+    p, t, td, tp = test_profile
+
+    with matplotlib.rc_context({'axes.autolimit_mode': 'data'}):
+        fig = plt.figure(figsize=(9, 9))
+        emagram = Emagram(fig, aspect='auto')
+        emagram.plot(p, t, 'r')
+        emagram.plot(p, tp, 'k')
+        emagram.shade_cape(p, t, tp)
+        emagram.shade_cin(p, t, tp, td)
+        emagram.ax.set_xlim(-50, 50)
+        emagram.ax.set_ylim(1000, 100)
+
+        # This works around the fact that newer pint versions default to degrees_Celsius
+        emagram.ax.set_xlabel('degC')
+
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0.033, remove_text=True, style='default')
+def test_emagram_shade_area(test_profile):
+    """Test shading areas on an Emagram plot."""
+    p, t, _, tp = test_profile
+
+    with matplotlib.rc_context({'axes.autolimit_mode': 'data'}):
+        fig = plt.figure(figsize=(9, 9))
+        emagram = Emagram(fig, aspect='auto')
+        emagram.plot(p, t, 'r')
+        emagram.plot(p, tp, 'k')
+        emagram.shade_area(p, t, tp)
+        emagram.ax.set_xlim(-50, 50)
+        emagram.ax.set_ylim(1000, 100)
+
+        # This works around the fact that newer pint versions default to degrees_Celsius
+        emagram.ax.set_xlabel('degC')
+
+    return fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=0.039, remove_text=True, style='default')
+def test_emagram_wide_aspect_ratio(test_profile):
+    """Test plotting an Emagram with a wide aspect ratio."""
+    p, t, _, tp = test_profile
+
+    fig = plt.figure(figsize=(12.5, 3))
+    emagram = Emagram(fig, aspect='auto')
+    emagram.plot(p, t, 'r')
+    emagram.plot(p, tp, 'k')
+    emagram.ax.set_xlim(-30, 50)
+    emagram.ax.set_ylim(1050, 700)
+
+    # This works around the fact that newer pint versions default to degrees_Celsius
+    emagram.ax.set_xlabel('degC')
+    return fig