Merge pull request #134 from fooof-tools/plts

Update plots

Author: TomDonoghue

fooof/plts/spectra.py

@@ -5,18 +5,24 @@
 This file contains functions for plotting power spectra, that take in data directly.
 """
 
+from itertools import repeat
+
 import numpy as np
 
-from fooof.plts.utils import check_ax, add_shades
 from fooof.core.modutils import safe_import, check_dependency
 
+from fooof.plts.settings import DEFAULT_FIGSIZE
+from fooof.plts.utils import check_ax, add_shades
+from fooof.plts.style import check_n_style, style_spectrum_plot
+
 plt = safe_import('.pyplot', 'matplotlib')
 
 ###################################################################################################
 ###################################################################################################
 
 @check_dependency(plt, 'matplotlib')
-def plot_spectrum(freqs, power_spectrum, log_freqs=False, log_powers=False, ax=None, **kwargs):
+def plot_spectrum(freqs, power_spectrum, log_freqs=False, log_powers=False,
+                  ax=None, plot_style=style_spectrum_plot, **kwargs):
     """Plot a power spectrum.
 
     Parameters
@@ -31,102 +37,121 @@ def plot_spectrum(freqs, power_spectrum, log_freqs=False, log_powers=False, ax=N
         Whether or not to take the log of the power axis before plotting.
     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.
     **kwargs
         Keyword arguments to be passed to the plot call.
     """
 
     # Create plot axes, if not provided
     if not ax:
-        _, ax = plt.subplots(figsize=(12, 10))
+        _, ax = plt.subplots(figsize=DEFAULT_FIGSIZE)
 
-    # Set plot data, logging if requested
+    # Set plot data & labels, logging if requested
     plt_freqs = np.log10(freqs) if log_freqs else freqs
     plt_powers = np.log10(power_spectrum) if log_powers else power_spectrum
 
-    # Create the plot
-    ax.plot(plt_freqs, plt_powers, **kwargs)
-
-    # Aesthetics and axis labels
-    ax.set_xlabel('Frequency', fontsize=20)
-    ax.set_ylabel('Power', fontsize=20)
-    ax.tick_params(axis='both', which='major', labelsize=16)
-    ax.grid(True)
+    # Set default plot settings, that only apply if not over-written in kwargs
+    if 'linewidth' not in kwargs:
+        kwargs['linewidth'] = 2.0
 
-    # If labels were provided, add a legend
-    if ax.get_legend_handles_labels()[0]:
-        ax.legend(prop={'size': 16})
+    # Create the plot & style
+    ax.plot(plt_freqs, plt_powers, **kwargs)
+    check_n_style(plot_style, ax, log_freqs, log_powers)
 
 
 @check_dependency(plt, 'matplotlib')
-def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=False, ax=None, **kwargs):
+def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=False, labels=None,
+                 ax=None, plot_style=style_spectrum_plot, **kwargs):
     """Plot multiple power spectra on the same plot.
 
     Parameters
     ----------
-    freqs : 1d array
+    freqs : 2d array or 1d array or list of 1d array
         X-axis data, frequency values.
-    power_spectra : list of 1d array
+    power_spectra : 2d array or list of 1d array
         Y-axis data, power values for spectra to plot.
     log_freqs : boolean, optional, default: False
         Whether or not to take the log of the power axis before plotting.
     log_powers : boolean, optional, default: False
         Whether or not to take the log of the power axis before plotting.
+    labels " "
     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.
     **kwargs
         Keyword arguments to be passed to the plot call.
     """
 
+    freqs = repeat(freqs) if isinstance(freqs, np.ndarray) and freqs.ndim == 1 else freqs
+    labels = repeat(labels) if not isinstance(labels, list) else labels
+
     ax = check_ax(ax)
-    for power_spectrum in power_spectra:
-        plot_spectrum(freqs, power_spectrum, log_freqs, log_powers, ax=ax, **kwargs)
+    for freq, power_spectrum, label in zip(freqs, power_spectra, labels):
+        plot_spectrum(freq, power_spectrum, log_freqs, log_powers, label=label,
+                      plot_style=None, ax=ax, **kwargs)
+    check_n_style(plot_style, ax, log_freqs, log_powers)
 
 
 @check_dependency(plt, 'matplotlib')
-def plot_spectrum_shading(freqs, power_spectrum, shades, add_center=False, ax=None, **kwargs):
+def plot_spectrum_shading(freqs, power_spectrum, shades, add_center=False,
+                          ax=None, plot_style=style_spectrum_plot, **kwargs):
     """Plot a power spectrum with a shaded frequency region (or regions).
 
     Parameters
     ----------
     freqs : 1d array
         X-axis data, frequency values.
-    power_spectrum : list of 1d array
+    power_spectrum : 1d array
         Y-axis data, power values for spectrum to plot.
     shades : list of [float, float] or list of list of [float, float]
         Shaded region(s) to add to plot, defined as [lower_bound, upper_bound].
     add_center : boolean, optional, default: False
         Whether to add a line at the center point of the shaded regions.
     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.
     **kwargs
         Keyword arguments to be passed to the plot call.
     """
 
     ax = check_ax(ax)
-    plot_spectrum(freqs, power_spectrum, ax=ax, **kwargs)
+    plot_spectrum(freqs, power_spectrum, plot_style=None, ax=ax, **kwargs)
     add_shades(ax, shades, add_center, kwargs.get('log_freqs', False))
+    check_n_style(plot_style, ax, kwargs.get('log_freqs', False), kwargs.get('log_powers', False))
 
 
 @check_dependency(plt, 'matplotlib')
-def plot_spectra_shading(freqs, power_spectra, shades, add_center=False, ax=None, **kwargs):
+def plot_spectra_shading(freqs, power_spectra, shades, add_center=False,
+                         ax=None, plot_style=style_spectrum_plot, **kwargs):
     """Plot a group of power spectra with a shaded frequency region (or regions).
 
     Parameters
     ----------
-    freqs : 1d array
+    freqs : 2d array or 1d array or list of 1d array
         X-axis data, frequency values.
-    power_spectra : list of 1d array
+    power_spectra : 2d array or list of 1d array
         Y-axis data, power values for spectra to plot.
     shades : list of [float, float] or list of list of [float, float]
         Shaded region(s) to add to plot, defined as [lower_bound, upper_bound].
     add_center : boolean, optional, default: False
         Whether to add a line at the center point of the shaded regions.
     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.
     **kwargs
-        Keyword arguments to be passed to the plot call.
+        Keyword arguments to be passed to plot_spectra or the plot call.
+
+    Notes
+    -----
+    Parameters for `plot_spectra` can also be passed into this function as **kwargs.
+    This includes `log_freqs`, `log_powers` & `labels`. See `plot_spectra for usage details.
     """
 
     ax = check_ax(ax)
-    plot_spectra(freqs, power_spectra, ax=ax, **kwargs)
+    plot_spectra(freqs, power_spectra, ax=ax, plot_style=None, **kwargs)
     add_shades(ax, shades, add_center, kwargs.get('log_freqs', False))
+    check_n_style(plot_style, ax, kwargs.get('log_freqs', False), kwargs.get('log_powers', False))

fooof/plts/style.py

@@ -0,0 +1,28 @@
+"""Helper functions that apply style and decoration to plots."""
+
+###################################################################################################
+###################################################################################################
+
+def check_n_style(style_func, *args):
+    """"Check is a style function has been passed, and apply if so."""
+
+    if style_func:
+        style_func(*args)
+
+
+def style_spectrum_plot(ax, log_freqs, log_powers):
+    """Define to styling for a power spectrum plot."""
+
+    # Get labels, based on log status
+    xlabel = 'Frequency' if not log_freqs else 'log(Frequency)'
+    ylabel = 'Power' if not log_powers else 'log(Power)'
+
+    # Aesthetics and axis labels
+    ax.set_xlabel(xlabel, fontsize=20)
+    ax.set_ylabel(ylabel, fontsize=20)
+    ax.tick_params(axis='both', which='major', labelsize=16)
+    ax.grid(True)
+
+    # If labels were provided, add a legend
+    if ax.get_legend_handles_labels()[0]:
+        ax.legend(prop={'size': 16})

fooof/tests/test_plts_spectra.py

@@ -9,13 +9,27 @@
 @plot_test
 def test_plot_spectrum(tfm, skip_if_no_mpl):
 
-    plot_spectrum(tfm.freqs, tfm.power_spectrum, True)
+    plot_spectrum(tfm.freqs, tfm.power_spectrum)
+
+    # Test with logging both axes
+    plot_spectrum(tfm.freqs, tfm.power_spectrum, True, True)
 
 @plot_test
 def test_plot_spectra(tfg, skip_if_no_mpl):
 
+    # Test with 1d inputs - 1d freq array and list of 1d power spectra
     plot_spectra(tfg.freqs, [tfg.power_spectra[0, :], tfg.power_spectra[1, :]])
 
+    # Test with multiple freq inputs - list of 1d freq array and list of 1d power spectra
+    plot_spectra([tfg.freqs, tfg.freqs], [tfg.power_spectra[0, :], tfg.power_spectra[1, :]])
+
+    # Test with 2d array inputs
+    plot_spectra(np.vstack([tfg.freqs, tfg.freqs]),
+                 np.vstack([tfg.power_spectra[0, :], tfg.power_spectra[1, :]]))
+
+    # Test with labels
+    plot_spectra(tfg.freqs, [tfg.power_spectra[0, :], tfg.power_spectra[1, :]], labels=['A', 'B'])
+
 @plot_test
 def test_plot_spectrum_shading(tfm, skip_if_no_mpl):
 
@@ -26,3 +40,8 @@ def test_plot_spectra_shading(tfg, skip_if_no_mpl):
 
     plot_spectra_shading(tfg.freqs, [tfg.power_spectra[0, :], tfg.power_spectra[1, :]],
                          shades=[8, 12], add_center=True)
+
+    # Test with **kwargs that pass into plot_spectra
+    plot_spectra_shading(tfg.freqs, [tfg.power_spectra[0, :], tfg.power_spectra[1, :]],
+                         shades=[8, 12], add_center=True, log_freqs=True, log_powers=True,
+                         labels=['A', 'B'])

fooof/tests/test_plts_styles.py

@@ -0,0 +1,29 @@
+"""Tests for fooof.plts.styles."""
+
+from fooof.plts.style import *
+
+###################################################################################################
+###################################################################################################
+
+def test_check_n_style(skip_if_no_mpl):
+
+    # Check can pass None and do nothing
+    check_n_style(None)
+    assert True
+
+    # Check can pass a callable
+    def checker(*args):
+        return True
+    check_n_style(checker)
+
+def test_style_spectrum_plot(skip_if_no_mpl):
+
+    # Create a dummy plot and style it
+    from fooof.core.modutils import safe_import
+    plt = safe_import('.pyplot', 'matplotlib')
+    _, ax = plt.subplots()
+    style_spectrum_plot(ax, False, False)
+
+    # Check that axis labels are added - use as proxy that it ran correctly
+    assert ax.xaxis.get_label().get_text()
+    assert ax.yaxis.get_label().get_text()