ENH: Updating tests for plotting to increase coverage (#777)

act/plotting/distributiondisplay.py

@@ -161,7 +161,11 @@ def plot_stacked_bar(
 
         if sortby_bins is None and sortby_field is not None:
             # We will defaut the y direction to have the same # of bins as x
-            sortby_bins = np.linspace(ydata.values.min(), ydata.values.max(), len(bins))
+            if isinstance(bins, int):
+                n_bins = bins
+            else:
+                n_bins = len(bins)
+            sortby_bins = np.linspace(ydata.values.min(), ydata.values.max(), n_bins)
 
         # Get the current plotting axis
         if self.fig is None:
@@ -175,18 +179,12 @@ def plot_stacked_bar(
                 ytitle = ''.join(['(', ydata.attrs['units'], ')'])
             else:
                 ytitle = field
-            if sortby_bins is None:
-                my_hist, x_bins, y_bins = np.histogram2d(
-                    xdata.values.flatten(), ydata.values.flatten(), density=density,
-                    bins=bins,
-                    **hist_kwargs)
-            else:
-                my_hist, x_bins, y_bins = np.histogram2d(
-                    xdata.values.flatten(),
-                    ydata.values.flatten(),
-                    density=density,
-                    bins=[bins, sortby_bins],
-                    **hist_kwargs)
+            my_hist, x_bins, y_bins = np.histogram2d(
+                xdata.values.flatten(),
+                ydata.values.flatten(),
+                density=density,
+                bins=[bins, sortby_bins],
+                **hist_kwargs)
             x_inds = (x_bins[:-1] + x_bins[1:]) / 2.0
             self.axes[subplot_index].bar(
                 x_inds,
@@ -402,8 +400,12 @@ def plot_stairstep(
             ydata = self._ds[dsname][sortby_field]
 
         if sortby_bins is None and sortby_field is not None:
+            if isinstance(bins, int):
+                n_bins = bins
+            else:
+                n_bins = len(bins)
             # We will defaut the y direction to have the same # of bins as x
-            sortby_bins = np.linspace(ydata.values.min(), ydata.values.max(), len(bins))
+            sortby_bins = np.linspace(ydata.values.min(), ydata.values.max(), n_bins)
 
         # Get the current plotting axis, add day/night background and plot data
         if self.fig is None:
@@ -418,18 +420,13 @@ def plot_stairstep(
                 ytitle = ''.join(['(', ydata.attrs['units'], ')'])
             else:
                 ytitle = field
-            if sortby_bins is None:
-                my_hist, x_bins, y_bins = np.histogram2d(
-                    xdata.values.flatten(), ydata.values.flatten(), bins=bins,
-                    density=density, **hist_kwargs)
-            else:
-                my_hist, x_bins, y_bins = np.histogram2d(
-                    xdata.values.flatten(),
-                    ydata.values.flatten(),
-                    density=density,
-                    bins=[bins, sortby_bins],
-                    **hist_kwargs
-                )
+            my_hist, x_bins, y_bins = np.histogram2d(
+                xdata.values.flatten(),
+                ydata.values.flatten(),
+                density=density,
+                bins=[bins, sortby_bins],
+                **hist_kwargs
+            )
             x_inds = (x_bins[:-1] + x_bins[1:]) / 2.0
             self.axes[subplot_index].step(
                 x_inds,

act/tests/plotting/test_distributiondisplay.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+import pandas as pd
 import xarray as xr
 from numpy.testing import assert_allclose
 
@@ -66,6 +67,17 @@ def test_distribution_errors():
     assert histdisplay.fig is not None
     assert histdisplay.axes is not None
 
+    histdisplay = DistributionDisplay({'thing1': sonde_ds, 'thing2': sonde_ds})
+    with np.testing.assert_raises(ValueError):
+        histdisplay.plot_stacked_bar('tdry')
+    with np.testing.assert_raises(ValueError):
+        histdisplay.plot_size_distribution('tdry', 'time')
+    with np.testing.assert_raises(ValueError):
+        histdisplay.plot_stairstep('tdry')
+    with np.testing.assert_raises(ValueError):
+        histdisplay.plot_heatmap('tdry', 'alt')
+    with np.testing.assert_raises(ValueError):
+        histdisplay.plot_violin('tdry')
     matplotlib.pyplot.close(fig=histdisplay.fig)
 
 
@@ -83,6 +95,21 @@ def test_stair_graph():
         matplotlib.pyplot.close(histdisplay.fig)
 
 
+@pytest.mark.mpl_image_compare(tolerance=30)
+def test_stair_graph2():
+    sonde_ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_SONDE1)
+    del sonde_ds['tdry'].attrs['units']
+
+    histdisplay = DistributionDisplay({'sgpsondewnpnC1.b1': sonde_ds})
+    histdisplay.plot_stairstep('tdry', sortby_field='alt')
+    sonde_ds.close()
+
+    try:
+        return histdisplay.fig
+    finally:
+        matplotlib.pyplot.close(histdisplay.fig)
+
+
 @pytest.mark.mpl_image_compare(tolerance=30)
 def test_stair_graph_sorted():
     sonde_ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_SONDE1)
@@ -132,6 +159,21 @@ def test_stacked_bar_graph2():
         matplotlib.pyplot.close(histdisplay.fig)
 
 
+@pytest.mark.mpl_image_compare(tolerance=30)
+def test_stacked_bar_graph3():
+    sonde_ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_SONDE1)
+    del sonde_ds['tdry'].attrs['units']
+
+    histdisplay = DistributionDisplay({'sgpsondewnpnC1.b1': sonde_ds})
+    histdisplay.plot_stacked_bar('tdry', sortby_field='alt')
+    sonde_ds.close()
+
+    try:
+        return histdisplay.fig
+    finally:
+        matplotlib.pyplot.close(histdisplay.fig)
+
+
 @pytest.mark.mpl_image_compare(tolerance=30)
 def test_stacked_bar_graph_sorted():
     sonde_ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_SONDE1)
@@ -171,6 +213,47 @@ def test_heatmap():
         matplotlib.pyplot.close(histdisplay.fig)
 
 
+@pytest.mark.mpl_image_compare(tolerance=30)
+def test_heatmap2():
+    sonde_ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_SONDE1)
+    del sonde_ds['tdry'].attrs['units']
+
+    histdisplay = DistributionDisplay({'sgpsondewnpnC1.b1': sonde_ds})
+    histdisplay.plot_heatmap(
+        'tdry',
+        'alt',
+        x_bins=10,
+        y_bins=10,
+        cmap='coolwarm',
+    )
+    sonde_ds.close()
+
+    try:
+        return histdisplay.fig
+    finally:
+        matplotlib.pyplot.close(histdisplay.fig)
+
+
+@pytest.mark.mpl_image_compare(tolerance=30)
+def test_heatmap3():
+    sonde_ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_SONDE1)
+    del sonde_ds['tdry'].attrs['units']
+
+    histdisplay = DistributionDisplay({'sgpsondewnpnC1.b1': sonde_ds})
+    histdisplay.plot_heatmap(
+        'tdry',
+        'alt',
+        threshold=1,
+        cmap='coolwarm',
+    )
+    sonde_ds.close()
+
+    try:
+        return histdisplay.fig
+    finally:
+        matplotlib.pyplot.close(histdisplay.fig)
+
+
 @pytest.mark.mpl_image_compare(tolerance=30)
 def test_size_distribution():
     sigma = 10
@@ -188,6 +271,24 @@ def test_size_distribution():
         matplotlib.pyplot.close(histdisplay.fig)
 
 
+@pytest.mark.mpl_image_compare(tolerance=30)
+def test_size_distribution2():
+    sigma = 10
+    mu = 50
+    bins = pd.date_range('2023-01-01', '2023-01-02', periods=mu)
+    ydata = 1 / (sigma * np.sqrt(2 * np.pi)) * np.exp(-((np.array(range(len(bins))) - mu) ** 2) / (2 * sigma**2))
+    y_array = xr.DataArray(ydata, dims={'time': bins})
+    bins = xr.DataArray(bins, dims={'time': bins})
+    my_fake_ds = xr.Dataset({'time': bins, 'ydata': y_array})
+    my_fake_ds['ydata'].attrs['units'] = 'units'
+    histdisplay = DistributionDisplay(my_fake_ds)
+    histdisplay.plot_size_distribution('ydata', bins, time=bins.values[10])
+    try:
+        return histdisplay.fig
+    finally:
+        matplotlib.pyplot.close(histdisplay.fig)
+
+
 def test_histogram_kwargs():
     files = sample_files.EXAMPLE_MET1
     ds = act.io.arm.read_arm_netcdf(files)
@@ -242,7 +343,29 @@ def test_violin():
 
     ds.close()
 
-    return display.fig
+    try:
+        return display.fig
+    finally:
+        matplotlib.pyplot.close(display.fig)
+
+
+@pytest.mark.mpl_image_compare(tolerance=30)
+def test_violin2():
+    ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_MET1)
+    del ds['temp_mean'].attrs['units']
+
+    # Create a DistributionDisplay object to compare fields
+    display = DistributionDisplay(ds)
+
+    # Create violin display of mean temperature
+    display.plot_violin('temp_mean', vert=False)
+
+    ds.close()
+
+    try:
+        return display.fig
+    finally:
+        matplotlib.pyplot.close(display.fig)
 
 
 @pytest.mark.mpl_image_compare(tolerance=30)
@@ -262,4 +385,26 @@ def test_scatter():
 
     ds.close()
 
-    return display.fig
+    try:
+        return display.fig
+    finally:
+        matplotlib.pyplot.close(display.fig)
+
+
+@pytest.mark.mpl_image_compare(tolerance=30)
+def test_scatter2():
+    ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_MET1)
+    del ds['wspd_arith_mean'].attrs['units']
+    del ds['wspd_vec_mean'].attrs['units']
+    # Create a DistributionDisplay object to compare fields
+    display = DistributionDisplay(ds)
+    display.plot_scatter(
+        'wspd_arith_mean', 'wspd_vec_mean',
+    )
+    display.set_ratio_line()
+    ds.close()
+
+    try:
+        return display.fig
+    finally:
+        matplotlib.pyplot.close(display.fig)

act/tests/plotting/test_geodisplay.py

@@ -1,7 +1,6 @@
 import matplotlib
-
 import pytest
-
+import numpy as np
 import act
 from act.tests import sample_files
 from act.plotting import GeographicPlotDisplay
@@ -73,3 +72,24 @@ def test_geoplot_tile():
     except Exception:
         pass
     sonde_ds.close()
+
+
+def test_geo_errors():
+    sonde_ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_SONDE1)
+    display = GeographicPlotDisplay({'thing1': sonde_ds, 'thing2': sonde_ds})
+    with np.testing.assert_raises(ValueError):
+        display.geoplot('tdry')
+    with np.testing.assert_raises(ValueError):
+        display.geoplot()
+
+    del sonde_ds['lat']
+    display = GeographicPlotDisplay({'thing1': sonde_ds, 'thing2': sonde_ds})
+    with np.testing.assert_raises(ValueError):
+        display.geoplot('tdry')
+
+    sonde_ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_SONDE1)
+    del sonde_ds['lon']
+    del sonde_ds['tdry'].attrs['units']
+    display = GeographicPlotDisplay({'thing1': sonde_ds, 'thing2': sonde_ds})
+    with np.testing.assert_raises(ValueError):
+        display.geoplot('tdry')