Climate Stripes plotting capability (#867)

* ADD: New example showing how to use ACT to work with ARM and Satellite data

* ENH: Updating tests

* ENH: Adding in averaging to the TSI data.

* ADD: Adding new test

* ADD: New plot

* ENH: Adding climate stripes plot

* ENH: Updates for stripe plot

* ENH: Adding example

* ENH: Updating documentation

* ENH: New image with new cmap

act/plotting/timeseriesdisplay.py

@@ -9,11 +9,14 @@
 from copy import deepcopy
 from re import search, search as re_search
 
+import numpy as np
+import pandas as pd
+from scipy import stats
 import matplotlib as mpl
 import matplotlib.dates as mdates
 import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
+from matplotlib.patches import Rectangle
+from matplotlib.collections import PatchCollection
 from matplotlib import colors as mplcolors
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 from scipy.interpolate import NearestNDInterpolator
@@ -1847,3 +1850,130 @@ def fill_between(
         ax.set_title(set_title)
         self.axes[subplot_index] = ax
         return self.axes[subplot_index]
+
+    def plot_stripes(
+        self,
+        field,
+        dsname=None,
+        subplot_index=(0,),
+        set_title=None,
+        reference_period=None,
+        cmap='bwr',
+        cbar_label=None,
+        colorbar=True,
+        **kwargs,
+    ):
+        """
+        Makes a climate stripe plot with or without a baseline period specified
+
+        Parameters
+        ----------
+        field : str
+            The name of the field to plot.
+        dsname : None or str
+            If there is more than one datastream in the display object the
+            name of the datastream needs to be specified. If set to None and
+            there is only one datastream ACT will use the sole datastream
+            in the object.
+        subplot_index : 1 or 2D tuple, list, or array
+            The index of the subplot to set the x range of.
+        set_title : str
+            The title for the plot.
+        reference_period : list
+            List of a start and end date for a reference period ['2020-01-01', '2020-04-01']
+            If this is set, the plot will subtract the mean of the reference period from the
+            field to create an anomaly calculation.
+        cmap : string
+            Colormap to use for plotting.  Defaults to bwr
+        cbar_label : str
+            Option to overwrite default colorbar label.
+        colorbar : boolean
+            Option to not plot the colorbar. Default is to plot it
+        **kwargs : keyword arguments
+            The keyword arguments for :func:`plt.plot` (1D timeseries) or
+            :func:`plt.pcolormesh` (2D timeseries).
+
+        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 TimeSeriesDisplay '
+                'object.'
+            )
+        elif dsname is None:
+            dsname = list(self._ds.keys())[0]
+
+        # Get data and dimensions
+        data = self._ds[dsname][field]
+        dim = list(self._ds[dsname][field].dims)
+        xdata = self._ds[dsname][dim[0]]
+
+        start = int(mdates.date2num(xdata.values[0]))
+        end = int(mdates.date2num(xdata.values[-1]))
+        delta = stats.mode(xdata.diff('time').values)[0] / np.timedelta64(1, 'D')
+
+        # Calculate mean for reference period and subtract from the data
+        if reference_period is not None:
+            reference = data.sel(time=slice(reference_period[0], reference_period[1])).mean('time')
+            data.values = data.values - reference.values
+
+        # Get the current plotting axis, add day/night background and plot data
+        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 ax to appropriate axis
+        ax = self.axes[subplot_index]
+
+        # Plot up data using rectangles
+        col = PatchCollection(
+            [Rectangle((y, 0), delta, 1) for y in np.arange(start, end + 1, delta)]
+        )
+        col.set_array(data)
+        col.set_cmap(cmap)
+        col.set_clim(np.nanmin(data), np.nanmax(data))
+        ax.add_collection(col)
+
+        locator = mdates.AutoDateLocator(minticks=3)
+        formatter = mdates.AutoDateFormatter(locator)
+        ax.xaxis.set_major_locator(locator)
+        ax.xaxis.set_major_formatter(formatter)
+
+        ax.set_ylim(0, 1)
+        ax.set_yticks([])
+        ax.set_xlim(start, end + 1)
+
+        # Set Title
+        if set_title is None:
+            set_title = ' '.join(
+                [
+                    dsname,
+                    field,
+                    'Stripes on',
+                    dt_utils.numpy_to_arm_date(self._ds[dsname].time.values[0]),
+                ]
+            )
+        ax.set_title(set_title)
+
+        # Set Colorbar
+        if colorbar:
+            if 'units' in data.attrs:
+                ytitle = ''.join(['(', data.attrs['units'], ')'])
+            else:
+                ytitle = field
+            if cbar_label is None:
+                cbar_title = ytitle
+            else:
+                cbar_title = ''.join(['(', cbar_label, ')'])
+            self.add_colorbar(col, title=cbar_title, subplot_index=subplot_index)
+
+        self.axes[subplot_index] = ax
+        return self.axes[subplot_index]

examples/plotting/plot_stripes.py

@@ -0,0 +1,28 @@
+"""
+Example plot using stripes
+--------------------------
+
+Plot up climate stripes plots from already
+existing climatologies from ARM data.
+Author: Adam Theisen
+
+"""
+
+import act
+import matplotlib.pyplot as plt
+
+# SGP E13 MET data has already been processed to yearly averages,
+# removing data flagged by embedded qc and DQRs
+url = 'https://raw.githubusercontent.com/AdamTheisen/ARM-Climatologies/refs/heads/main/results/sgpmetE13.b1_temp_mean_Y.csv'
+col_names = ['time', 'temperature', 'count']
+ds = act.io.read_csv(url, column_names=col_names, index_col=0, parse_dates=True)
+
+# Drop years with less than 500000 samples
+ds = ds.where(ds['count'] > 500000)
+
+# Create plot display
+display = act.plotting.TimeSeriesDisplay(ds, figsize=(10, 2))
+reference = ['2003-01-01', '2013-01-01']
+display.plot_stripes('temperature', reference_period=reference)
+
+plt.show()

tests/plotting/test_timeseriesdisplay.py

@@ -685,3 +685,16 @@ def test_xlim_correction_plot():
     ds.close()
 
     return display.fig
+
+
+@pytest.mark.mpl_image_compare(tolerance=10)
+def test_plot_stripes():
+    ds = act.io.read_arm_netcdf(sample_files.EXAMPLE_MET_WILDCARD)
+    ds = ds.resample(time='1H').mean()
+    print(ds)
+    reference_period = ['2019-01-01', '2019-10-02']
+
+    display = act.plotting.TimeSeriesDisplay(ds, figsize=(10, 2))
+    display.plot_stripes('temp_mean', reference_period=reference_period)
+
+    return display.fig