update args for average & shade

Author: TomDonoghue

fooof/plts/spectra.py

@@ -8,6 +8,7 @@
 from itertools import repeat, cycle
 
 import numpy as np
+from scipy.stats import sem
 
 from fooof.core.modutils import safe_import, check_dependency
 from fooof.plts.settings import PLT_FIGSIZES
@@ -177,8 +178,9 @@ def plot_spectra_shading(freqs, power_spectra, shades, shade_colors='r',
 @savefig
 @style_plot
 @check_dependency(plt, 'matplotlib')
-def plot_spectra_yshade(freqs, power_spectra, shade=None, scale=1, log_freqs=False,
-                        log_powers=False, ax=None, **plot_kwargs):
+def plot_spectra_yshade(freqs, power_spectra, shade='std', average='mean', scale=1,
+                        log_freqs=False, log_powers=False, color=None, label=None,
+                        ax=None, **plot_kwargs):
     """Plot standard deviation or error as a shaded region around the mean spectrum.
 
     Parameters
@@ -187,23 +189,27 @@ def plot_spectra_yshade(freqs, power_spectra, shade=None, scale=1, log_freqs=Fal
         Frequency values, to be plotted on the x-axis.
     power_spectra : 1d or 2d array
         Power values, to be plotted on the y-axis. ``shade`` must be provided if 1d.
-    shade : 1d array, optional, default: None
-        Powers to shade above/below the mean spectrum. None defaults to one standard deviation.
+    shade : 'std', 'sem', 1d array or callable, optional, default: 'std'
+        Approach for shading above/below the mean spectrum.
+    average : 'mean', 'median' or callable, optional, default: 'mean'
+        Averaging approach for the average spectrum to plot. Only used if power_spectra is 2d.
     scale : int, optional, default: 1
-        Factor to multiply the the standard deviation, or ``shade``, by.
+        Factor to multiply the plotted shade by.
     log_freqs : bool, optional, default: False
         Whether to plot the frequency axis in log spacing.
     log_powers : bool, optional, default: False
         Whether to plot the power axis in log spacing.
+    color : str, optional, default: None
+        Line color of the spectrum.
+    label : str, optional, default: None
+        Legend label for the spectrum.
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
-    plot_style : callable, optional, default: style_spectrum_plot
-        A function to call to apply styling & aesthetics to the plot.
     **plot_kwargs
         Keyword arguments to be passed to `plot_spectra` or to the plot call.
     """
 
-    if shade is None and power_spectra.ndim != 2:
+    if isinstance(shade, str) and power_spectra.ndim != 2:
         raise ValueError('Power spectra must be 2d if shade is not given.')
 
     ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['spectral']))
@@ -212,16 +218,25 @@ def plot_spectra_yshade(freqs, power_spectra, shade=None, scale=1, log_freqs=Fal
     plt_freqs = np.log10(freqs) if log_freqs else freqs
     plt_powers = np.log10(power_spectra) if log_powers else power_spectra
 
-    # Plot mean
-    powers_mean = np.mean(plt_powers, axis=0) if plt_powers.ndim == 2 else plt_powers
-    ax.plot(plt_freqs, powers_mean)
+    # Organize mean spectrum to plot
+    avg_funcs = {'mean' : np.mean, 'median' : np.median}
+    avg_func = avg_funcs[average] if isinstance(average, str) else average
+    avg_powers = avg_func(plt_powers, axis=0) if plt_powers.ndim == 2 else plt_powers
+
+    # Plot average power spectrum
+    ax.plot(plt_freqs, avg_powers, linewidth=2.0, color=color, label=label)
 
-    # Shade +/- scale * (standard deviation or shade)
-    shade = scale * np.std(plt_powers, axis=0) if shade is None else scale * shade
-    upper_shade = powers_mean + shade
-    lower_shade = powers_mean - shade
+    # Organize shading to plot
+    shade_funcs = {'std' : np.std, 'sem' : sem}
+    shade_func = shade_funcs[shade] if isinstance(shade, str) else shade
+    shade_vals = scale * shade_func(plt_powers, axis=0) \
+        if isinstance(shade, str) else scale * shade
+    upper_shade = avg_powers + shade_vals
+    lower_shade = avg_powers - shade_vals
 
+    # Plot +/- yshading around spectrum
     alpha = plot_kwargs.pop('alpha', 0.25)
-    ax.fill_between(plt_freqs, lower_shade, upper_shade, alpha=alpha, **plot_kwargs)
+    ax.fill_between(plt_freqs, lower_shade, upper_shade,
+                    alpha=alpha, color=color, **plot_kwargs)
 
     style_spectrum_plot(ax, log_freqs, log_powers)

fooof/tests/plts/test_spectra.py

@@ -72,11 +72,18 @@ def test_plot_spectra_yshade(skip_if_no_mpl, tfg):
     with raises(ValueError):
         plot_spectra_yshade(freqs, powers[0])
 
-    # Valid 1d array with shade
-    plot_spectra_yshade(freqs, np.mean(powers, axis=0), shade=np.std(powers, axis=0),
+    # Plot with 2d array
+    plot_spectra_yshade(freqs, powers, shade='std',
                         save_fig=True, file_path=TEST_PLOTS_PATH,
                         file_name='test_plot_spectra_yshade1.png')
 
-    # 2d array
-    plot_spectra_yshade(freqs, powers, save_fig=True, file_path=TEST_PLOTS_PATH,
+    # Plot shade with given 1d array
+    plot_spectra_yshade(freqs, np.mean(powers, axis=0),
+                        shade=np.std(powers, axis=0),
+                        save_fig=True, file_path=TEST_PLOTS_PATH,
                         file_name='test_plot_spectra_yshade2.png')
+
+    # Plot shade with different average and shade approaches
+    plot_spectra_yshade(freqs, powers, shade='sem', average='median',
+                        save_fig=True, file_path=TEST_PLOTS_PATH,
+                        file_name='test_plot_spectra_yshade3.png')