Climate Stripes plotting capability

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]

act/plotting/xsectiondisplay.py

@@ -146,13 +146,14 @@ def set_yrng(self, yrng, subplot_index=(0,)):
 
     def plot_xsection(
         self,
-        dsname,
-        varname,
+        field,
+        dsname=None,
         x=None,
         y=None,
         subplot_index=(0,),
         sel_kwargs=None,
         isel_kwargs=None,
+        set_title=None,
         **kwargs,
     ):
         """
@@ -162,11 +163,10 @@ def plot_xsection(
 
         Parameters
         ----------
-        dsname : str or None
-            The name of the datastream to plot from. Set to None to have
-            ACT attempt to automatically detect this.
-        varname : str
+        field : str
             The name of the variable to plot.
+        dsname : str or None
+            The name of the datastream to plot from.
         x : str or None
             The name of the x coordinate variable.
         y : str or None
@@ -181,6 +181,8 @@ def plot_xsection(
             to slice datasets, see the documentation on :func:`xarray.DataArray.sel`.
         isel_kwargs : dict
             The keyword arguments to pass into :py:func:`xarray.DataArray.sel`
+        set_title : str
+            Title for the plot
         **kwargs : keyword arguments
             Additional keyword arguments will be passed into
             :func:`xarray.DataArray.plot`.
@@ -220,9 +222,9 @@ def plot_xsection(
             y_coord_dim = temp_ds[y].dims[0]
             coord_list[y] = y_coord_dim
             new_ds = data_utils.assign_coordinates(temp_ds, coord_list)
-            my_dataarray = new_ds[varname]
+            my_dataarray = new_ds[field]
         else:
-            my_dataarray = temp_ds[varname]
+            my_dataarray = temp_ds[field]
 
         coord_keys = [key for key in my_dataarray.coords.keys()]
         # X-array will sometimes shorten latitude and longitude variables
@@ -255,28 +257,42 @@ def plot_xsection(
         self.set_xrng(xrng, subplot_index)
         yrng = self.axes[subplot_index].get_ylim()
         self.set_yrng(yrng, subplot_index)
+
+        if set_title is None:
+            if 'long_name' in self._ds[dsname][field].attrs:
+                set_title = self._ds[dsname][field].attrs['long_name']
+        plt.title(set_title)
+
         del temp_ds
         return self.axes[subplot_index]
 
     def plot_xsection_map(
-        self, dsname, varname, subplot_index=(0,), coastlines=True, background=False, **kwargs
+        self,
+        field,
+        dsname=None,
+        subplot_index=(0,),
+        coastlines=True,
+        background=False,
+        set_title=None,
+        **kwargs,
     ):
         """
         Plots a cross section of 2D data on a geographical map.
 
         Parameters
         ----------
-        dsname : str or None
-            The name of the datastream to plot from. Set to None
-            to have ACT attempt to automatically detect this.
-        varname : str
+        field : str
             The name of the variable to plot.
+        dsname : str or None
+            The name of the datastream to plot from.
         subplot_index : tuple
             The index of the subplot to plot inside.
         coastlines : bool
             Set to True to plot the coastlines.
         background : bool
             Set to True to plot a stock image background.
+        set_title : str
+            Title for the plot
         **kwargs : keyword arguments
             Additional keyword arguments will be passed into
             :func:`act.plotting.XSectionDisplay.plot_xsection`
@@ -292,7 +308,9 @@ def plot_xsection_map(
                 'Cartopy needs to be installed in order to plot ' + 'cross sections on maps!'
             )
         self.set_subplot_to_map(subplot_index)
-        self.plot_xsection(dsname, varname, subplot_index=subplot_index, **kwargs)
+        self.plot_xsection(
+            field, dsname=dsname, subplot_index=subplot_index, set_title=set_title, **kwargs
+        )
         xlims = self.xrng[subplot_index].flatten()
         ylims = self.yrng[subplot_index].flatten()
         self.axes[subplot_index].set_xticks(np.linspace(round(xlims[0], 0), round(xlims[1], 0), 10))

examples/plotting/plot_satellite.py

@@ -0,0 +1,53 @@
+"""
+Using ACT for Satellite data
+----------------------------
+
+Simple example for working with satellite data.
+It is recommended that users explore other
+satellite specific libraries suchas SatPy.
+
+Author: Adam Theisen
+"""
+
+import act
+import matplotlib.pyplot as plt
+import numpy as np
+from arm_test_data import DATASETS
+
+# Read in VISST Data
+files = DATASETS.fetch('enavisstgridm11minnisX1.c1.20230307.000000.cdf')
+ds = act.io.read_arm_netcdf(files)
+
+# Plot up the VISST cloud percentage using XSectionDisplay
+display = act.plotting.XSectionDisplay(ds, figsize=(10, 8))
+display.plot_xsection_map(
+    'cloud_percentage', x='longitude', y='latitude', isel_kwargs={'time': 0, 'cld_type': 0}
+)
+plt.show()
+
+# Download ARM TSI Data
+files = DATASETS.fetch('enatsiskycoverC1.b1.20230307.082100.cdf')
+ds_tsi = act.io.read_arm_netcdf(files)
+ds_tsi = ds_tsi.where(ds_tsi.percent_opaque > 0)
+ds_tsi = ds_tsi.resample(time='30min').mean()
+
+# Set coordinates to extra data for ENA
+ena_lat = 39.091600
+ena_lon = 28.025700
+
+lat = ds['lat'].values
+lon = ds['lon'].values
+
+# Find the nearest pixel for the satellite data and extract it
+lat_ind = np.argmin(np.abs(lat - ena_lat))
+lon_ind = np.argmin(np.abs(lon - ena_lon))
+
+ds_new = ds.isel(lat=lat_ind, lon=lon_ind, cld_type=0)
+
+# Plot the comparison using TimeSeriesDisplay
+display = act.plotting.TimeSeriesDisplay({'Satellite': ds_new, 'ARM': ds_tsi}, figsize=(15, 8))
+display.plot('cloud_percentage', dsname='Satellite', label='VISST Cloud Percentage')
+display.plot('percent_opaque', dsname='ARM', label='ARM TSI Percent Opaque')
+display.day_night_background(dsname='ARM')
+plt.legend()
+plt.show()

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

tests/plotting/test_xsectiondisplay.py

@@ -28,7 +28,7 @@ def test_xsection_errors():
 
     display = XSectionDisplay(ds, figsize=(10, 8), subplot_shape=(1,))
     with np.testing.assert_raises(RuntimeError):
-        display.plot_xsection(None, 'backscatter', x='time', cmap='HomeyerRainbow')
+        display.plot_xsection('backscatter', x='time', cmap='HomeyerRainbow')
 
     ds.close()
     matplotlib.pyplot.close(fig=display.fig)
@@ -39,9 +39,7 @@ def test_xsection_plot():
     visst_ds = act.io.arm.read_arm_netcdf(sample_files.EXAMPLE_CEIL1)
 
     xsection = XSectionDisplay(visst_ds, figsize=(10, 8))
-    xsection.plot_xsection(
-        None, 'backscatter', x='time', y='range', cmap='coolwarm', vmin=0, vmax=320
-    )
+    xsection.plot_xsection('backscatter', x='time', y='range', cmap='coolwarm', vmin=0, vmax=320)
     visst_ds.close()
 
     try:
@@ -63,7 +61,6 @@ def test_xsection_plot_map():
         ):
             xsection = XSectionDisplay(radar_ds, figsize=(15, 8))
         xsection.plot_xsection_map(
-            None,
             'ir_temperature',
             vmin=220,
             vmax=300,