diff --git a/act/plotting/timeseriesdisplay.py b/act/plotting/timeseriesdisplay.py index 45f321aa08..ffe573dbdf 100644 --- a/act/plotting/timeseriesdisplay.py +++ b/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] diff --git a/act/plotting/xsectiondisplay.py b/act/plotting/xsectiondisplay.py index 0538e8d08d..cc320e28cd 100644 --- a/act/plotting/xsectiondisplay.py +++ b/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)) diff --git a/examples/plotting/plot_satellite.py b/examples/plotting/plot_satellite.py new file mode 100644 index 0000000000..c66c780212 --- /dev/null +++ b/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() diff --git a/examples/plotting/plot_stripes.py b/examples/plotting/plot_stripes.py new file mode 100644 index 0000000000..c6def32e80 --- /dev/null +++ b/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() diff --git a/tests/plotting/baseline/test_plot_stripes.png b/tests/plotting/baseline/test_plot_stripes.png new file mode 100644 index 0000000000..da587f3a6c Binary files /dev/null and b/tests/plotting/baseline/test_plot_stripes.png differ diff --git a/tests/plotting/test_timeseriesdisplay.py b/tests/plotting/test_timeseriesdisplay.py index 508331f421..464dfb82f5 100644 --- a/tests/plotting/test_timeseriesdisplay.py +++ b/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 diff --git a/tests/plotting/test_xsectiondisplay.py b/tests/plotting/test_xsectiondisplay.py index 8dede9047a..1c45db5e68 100644 --- a/tests/plotting/test_xsectiondisplay.py +++ b/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,