diff --git a/act/plotting/distributiondisplay.py b/act/plotting/distributiondisplay.py index dd8cee55e5..5cb967b265 100644 --- a/act/plotting/distributiondisplay.py +++ b/act/plotting/distributiondisplay.py @@ -5,7 +5,7 @@ import xarray as xr import pandas as pd -from ..utils import datetime_utils as dt_utils +from ..utils import datetime_utils as dt_utils, calculate_percentages from .plot import Display @@ -863,3 +863,98 @@ def plot_violin( self.axes[subplot_index].set_yticks([]) return self.axes[subplot_index] + + def plot_pie_chart( + self, + fields, + time=None, + time_slice=None, + threshold=None, + fill_value=0.0, + dsname=None, + subplot_index=(0,), + set_title=None, + autopct='%1.1f%%', + **kwargs, + ): + """ + This procedure will produce a pie chart for the selected fields. + + Parameters + ---------- + fields : list + The list of fields to calculate percentages on for the pie chart. + time : datetime + A single datetime to be passed into the act.utils.calculate percentages function + if desired. Default is None and all data will be included. + time_slice : tuple + A tuple of two datetimes to grab all data between those two datetimes for + act.utils.calculate_percentages. Default is None and all data will be included. + threshold : float + Threshold in which anything below will be considered invalid. + Default is None. + fill_value : float + Fill value for invalid data. Only used if a threshold is provided. + dsname : str or None + The name of the datastream the field is contained in. Set + to None to let ACT automatically determine this. + subplot_index : tuple + The subplot index to place the plot in + set_title : str + The title of the plot. + autopct : str + Format string for the percentages. Default is float with one + decimal place. If this parameter is set to None, no percentage + string values are displayed. + **kwargs : keywords + Keywords to pass through to :func:`matplotlib.pyplot.pie`. + + Returns + ------- + ax : matplotlib axis handle + The matplotlib axis handle of the plot + + """ + if dsname is None and len(self._ds.keys()) > 1: + raise ValueError( + 'You must choose a datastream when there are 2 ' + + 'or more datasets in the DistributionDisplay ' + + 'object.' + ) + elif dsname is None: + dsname = list(self._ds.keys())[0] + + # Get the current plotting axis + if self.fig is None: + self.fig = plt.figure() + if self.axes is None: + self.axes = np.array([plt.axes()]) + self.fig.add_axes(self.axes[0]) + + # Set Title + if set_title is None: + set_title = ' '.join( + [ + dsname, + 'on', + dt_utils.numpy_to_arm_date(self._ds[dsname].time.values[0]), + ] + ) + self.axes[subplot_index].set_title(set_title) + + percentages = calculate_percentages( + self._ds[dsname], + fields, + time=time, + time_slice=time_slice, + threshold=threshold, + fill_value=fill_value, + ) + + self.axes[subplot_index].pie( + [percentages[field] for field in percentages.keys()], + labels=percentages.keys(), + autopct=autopct, + **kwargs, + ) + return self.axes[subplot_index] diff --git a/examples/plotting/plot_pie_chart.py b/examples/plotting/plot_pie_chart.py new file mode 100644 index 0000000000..c2f0a3a6c4 --- /dev/null +++ b/examples/plotting/plot_pie_chart.py @@ -0,0 +1,52 @@ +""" +Calculate and View Aerosol Percentages +-------------------------------------- + +Calculate the percentages of different aerosols in a Aerosol +Chemical Speciation (AOS) monitor dataset and view the percentages +in a pie chart. + +Written: Zach Sherman + +""" + +from arm_test_data import DATASETS +import matplotlib.pyplot as plt + +import act +from act.io.arm import read_arm_netcdf + +# Read an ARM AOS dataset +filename = DATASETS.fetch('sgpaosacsmE13.b2.20230420.000109.nc') +ds = read_arm_netcdf(filename) + +# Let us print out the fields in the dataset and see what it contains. +print(ds.data_vars.keys()) + +# Knowing what fields the dataset contains, let's create a list of fields +# to use in the plot. + +fields = ['sulfate', 'ammonium', 'nitrate', 'chloride'] + +# We also want to provide some keyword arguments to avoid invalid data such +# as negative values. +threshold = 0.0 +fill_value = 0.0 + +# Create a DistributionDisplay object to compare fields +display = act.plotting.DistributionDisplay(ds) + +# We can set one of the slices to explode and give it a nice shadow. +explode = (0, 0.1, 0, 0) +shadow = True + +# Create a pie chart using the fields list. The percentages of the +# fields will be calculated using act.utils.calculate_percentages. +display.plot_pie_chart( + fields, + threshold=threshold, + fill_value=fill_value, + explode=explode, + shadow=True, +) +plt.show() diff --git a/tests/plotting/baseline/test_plot_pie_chart.png b/tests/plotting/baseline/test_plot_pie_chart.png new file mode 100644 index 0000000000..79eeee2b48 Binary files /dev/null and b/tests/plotting/baseline/test_plot_pie_chart.png differ diff --git a/tests/plotting/baseline/test_plot_pie_chart_kwargs.png b/tests/plotting/baseline/test_plot_pie_chart_kwargs.png new file mode 100644 index 0000000000..2d3d8ae441 Binary files /dev/null and b/tests/plotting/baseline/test_plot_pie_chart_kwargs.png differ diff --git a/tests/plotting/test_distributiondisplay.py b/tests/plotting/test_distributiondisplay.py index 517a229a29..ad906ab443 100644 --- a/tests/plotting/test_distributiondisplay.py +++ b/tests/plotting/test_distributiondisplay.py @@ -412,3 +412,38 @@ def test_scatter2(): return display.fig finally: matplotlib.pyplot.close(display.fig) + + +@pytest.mark.mpl_image_compare(tolerance=10) +def test_plot_pie_chart(): + ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_AOSACSM) + fields = ['sulfate', 'ammonium', 'nitrate', 'chloride'] + display = DistributionDisplay(ds) + display.plot_pie_chart(fields) + ds.close() + + try: + return display.fig + finally: + matplotlib.pyplot.close(display.fig) + + +@pytest.mark.mpl_image_compare(tolerance=10) +def test_plot_pie_chart_kwargs(): + ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_AOSACSM) + fields = ['sulfate', 'ammonium', 'nitrate', 'chloride'] + threshold = 0.0 + fill_value = 0.0 + display = DistributionDisplay(ds) + display.plot_pie_chart( + fields, + threshold=threshold, + fill_value=fill_value, + colors=['olivedrab', 'rosybrown', 'gray', 'saddlebrown'], + ) + ds.close() + + try: + return display.fig + finally: + matplotlib.pyplot.close(display.fig)