Merge pull request #300 from fooof-tools/psdgrid

[MNT] - Make grid an updateable argument in PSD plots

Author: TomDonoghue

doc/api.rst

@@ -326,6 +326,27 @@ Annotated plots that describe the model and fitting process.
     plot_annotated_model
     plot_annotated_peak_search
 
+Plot Utilities & Styling
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+Plot related utilies for styling and managing plots.
+
+.. currentmodule:: fooof.plts.style
+
+.. autosummary::
+    :toctree: generated/
+
+    check_style_options
+
+.. currentmodule:: fooof.plts.utils
+
+.. autosummary::
+    :toctree: generated/
+
+    check_ax
+    recursive_plot
+    save_figure
+
 Utilities
 ---------
 

fooof/objs/fit.py

@@ -407,7 +407,7 @@ def report(self, freqs=None, power_spectrum=None, freq_range=None,
             Only relevant / effective if `freqs` and `power_spectrum` passed in in this call.
         **plot_kwargs
             Keyword arguments to pass into the plot method.
-            Plot options with a name conflict be passed by pre-pending 'plot_'.
+            Plot options with a name conflict be passed by pre-pending `plot_`.
             e.g. `freqs`, `power_spectrum` and `freq_range`.
 
         Notes

fooof/plts/aperiodic.py

@@ -33,7 +33,7 @@ def plot_aperiodic_params(aps, colors=None, labels=None, ax=None, **plot_kwargs)
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Keyword arguments to pass into the ``style_plot``.
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
     """
 
     ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['params']))
@@ -83,7 +83,7 @@ def plot_aperiodic_fits(aps, freq_range, control_offset=False,
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Keyword arguments to pass into the ``style_plot``.
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
     """
 
     ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['params']))

fooof/plts/error.py

@@ -33,7 +33,7 @@ def plot_spectral_error(freqs, error, shade=None, log_freqs=False, ax=None, **pl
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Keyword arguments to pass into the ``style_plot``.
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
     """
 
     ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['spectral']))

fooof/plts/fg.py

@@ -33,6 +33,8 @@ def plot_fg(fg, save_fig=False, file_name=None, file_path=None, **plot_kwargs):
         Name to give the saved out file.
     file_path : Path or str, optional
         Path to directory to save to. If None, saves to current directory.
+    **plot_kwargs
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
 
     Raises
     ------
@@ -76,7 +78,7 @@ def plot_fg_ap(fg, ax=None, **plot_kwargs):
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Keyword arguments to pass into the ``style_plot``.
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
     """
 
     if fg.aperiodic_mode == 'knee':
@@ -101,7 +103,7 @@ def plot_fg_gf(fg, ax=None, **plot_kwargs):
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Keyword arguments to pass into the ``style_plot``.
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
     """
 
     plot_scatter_2(fg.get_params('error'), 'Error',
@@ -121,7 +123,7 @@ def plot_fg_peak_cens(fg, ax=None, **plot_kwargs):
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Keyword arguments to pass into the ``style_plot``.
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
     """
 
     plot_hist(fg.get_params('peak_params', 0)[:, 0], 'Center Frequency',

fooof/plts/fm.py

@@ -63,7 +63,7 @@ def plot_fm(fm, plot_peaks=None, plot_aperiodic=True, freqs=None, power_spectrum
     data_kwargs, model_kwargs, aperiodic_kwargs, peak_kwargs : None or dict, optional
         Keyword arguments to pass into the plot call for each plot element.
     **plot_kwargs
-        Keyword arguments to pass into the ``style_plot``.
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
 
     Notes
     -----
@@ -169,7 +169,7 @@ def _add_peaks_shade(fm, plt_log, ax, **plot_kwargs):
     ax : matplotlib.Axes
         Figure axes upon which to plot.
     **plot_kwargs
-        Keyword arguments to pass into the ``fill_between``.
+        Keyword arguments to pass into ``fill_between``.
     """
 
     defaults = {'color' : PLT_COLORS['periodic'], 'alpha' : 0.25}

fooof/plts/periodic.py

@@ -35,7 +35,7 @@ def plot_peak_params(peaks, freq_range=None, colors=None, labels=None, ax=None,
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Keyword arguments to pass into the ``style_plot``.
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
     """
 
     ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['params']))
@@ -86,7 +86,7 @@ def plot_peak_fits(peaks, freq_range=None, colors=None, labels=None, ax=None, **
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Keyword arguments to pass into the plot call.
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
     """
 
     ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['params']))

fooof/plts/settings.py

@@ -28,7 +28,8 @@
                'linestyle' : ['ls', 'linestyle']}
 
 # Plot style arguments are those that can be defined on an axis object
-AXIS_STYLE_ARGS = ['title', 'xlabel', 'ylabel', 'xlim', 'ylim']
+AXIS_STYLE_ARGS = ['title', 'xlabel', 'ylabel', 'xlim', 'ylim',
+                   'xticks', 'yticks', 'xticklabels', 'yticklabels']
 
 # Line style arguments are those that can be defined on a line object
 LINE_STYLE_ARGS = ['alpha', 'lw', 'linewidth', 'ls', 'linestyle',
@@ -40,8 +41,13 @@
 # Custom style arguments are those that are custom-handled by the plot style function
 CUSTOM_STYLE_ARGS = ['title_fontsize', 'label_size', 'tick_labelsize',
                      'legend_size', 'legend_loc']
-STYLERS = ['axis_styler', 'line_styler', 'custom_styler']
-STYLE_ARGS = AXIS_STYLE_ARGS + LINE_STYLE_ARGS + CUSTOM_STYLE_ARGS + STYLERS
+
+# Define list of available style functions - these can also be replaced by arguments
+STYLERS = ['axis_styler', 'line_styler', 'collection_styler', 'custom_styler']
+
+# Collect the full set of possible style related input keyword arguments
+STYLE_ARGS = \
+    AXIS_STYLE_ARGS + LINE_STYLE_ARGS + COLLECTION_STYLE_ARGS + CUSTOM_STYLE_ARGS + STYLERS
 
 ## Define default values for plot aesthetics
 # These are all custom style arguments

fooof/plts/spectra.py

@@ -47,13 +47,14 @@ def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=False, freq_r
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Additional plot related keyword arguments.
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
+        For spectra plots, boolean input `grid` can be used to control if the figure has a grid.
     """
 
+    # Create the plot & collect plot kwargs of interest
     ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['spectral']))
-
-    # Create the plot
     plot_kwargs = check_plot_kwargs(plot_kwargs, {'linewidth' : 2.0})
+    grid = plot_kwargs.pop('grid', True)
 
     # Check for frequency range input, and log if x-axis is in log space
     if freq_range is not None:
@@ -82,7 +83,7 @@ def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=False, freq_r
 
     ax.set_xlim(freq_range)
 
-    style_spectrum_plot(ax, log_freqs, log_powers)
+    style_spectrum_plot(ax, log_freqs, log_powers, grid)
 
 
 # Alias `plot_spectrum` to `plot_spectra` for backwards compatibility
@@ -110,8 +111,9 @@ def plot_spectra_shading(freqs, power_spectra, shades, shade_colors='r',
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Additional plot related keyword arguments.
-        This can include additional inputs into :func:`~.plot_spectra`.
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
+        For spectra plots, boolean input `grid` can be used to control if the figure has a grid.
+        This can also include additional inputs into :func:`~.plot_spectra`.
 
     Notes
     -----
@@ -127,7 +129,8 @@ def plot_spectra_shading(freqs, power_spectra, shades, shade_colors='r',
     add_shades(ax, shades, shade_colors, add_center, plot_kwargs.get('log_freqs', False))
 
     style_spectrum_plot(ax, plot_kwargs.get('log_freqs', False),
-                        plot_kwargs.get('log_powers', False))
+                        plot_kwargs.get('log_powers', False),
+                        plot_kwargs.get('grid', True))
 
 
 # Alias `plot_spectrum_shading` to `plot_spectra_shading` for backwards compatibility
@@ -165,13 +168,16 @@ def plot_spectra_yshade(freqs, power_spectra, shade='std', average='mean', scale
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Additional plot related keyword arguments.
+        Additional plot related keyword arguments, with styling options managed by ``style_plot``.
+        For spectra plots, boolean input `grid` can be used to control if the figure has a grid.
+        This can also include additional inputs into :func:`~.plot_spectra`.
     """
 
     if (isinstance(shade, str) or isfunction(shade)) 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']))
+    grid = plot_kwargs.pop('grid', True)
 
     # Set plot data & labels, logging if requested
     plt_freqs = np.log10(freqs) if log_freqs else freqs
@@ -208,4 +214,4 @@ def plot_spectra_yshade(freqs, power_spectra, shade='std', average='mean', scale
     ax.fill_between(plt_freqs, lower_shade, upper_shade,
                     alpha=alpha, color=color, **plot_kwargs)
 
-    style_spectrum_plot(ax, log_freqs, log_powers)
+    style_spectrum_plot(ax, log_freqs, log_powers, grid)

fooof/plts/style.py

@@ -6,13 +6,23 @@
 import matplotlib.pyplot as plt
 
 from fooof.plts.settings import (AXIS_STYLE_ARGS, LINE_STYLE_ARGS, COLLECTION_STYLE_ARGS,
-                                 STYLE_ARGS, LABEL_SIZE, LEGEND_SIZE, LEGEND_LOC,
-                                 TICK_LABELSIZE, TITLE_FONTSIZE)
+                                 CUSTOM_STYLE_ARGS, STYLE_ARGS, TICK_LABELSIZE, TITLE_FONTSIZE,
+                                 LABEL_SIZE, LEGEND_SIZE, LEGEND_LOC)
 
 ###################################################################################################
 ###################################################################################################
 
-def style_spectrum_plot(ax, log_freqs, log_powers):
+def check_style_options():
+    """Check the list of valid style arguments that can be passed into plot functions."""
+
+    print('Valid style arguments:')
+    for label, options in zip(['Axis', 'Line', 'Collection', 'Custom'],
+                              [AXIS_STYLE_ARGS, LINE_STYLE_ARGS,
+                               COLLECTION_STYLE_ARGS, CUSTOM_STYLE_ARGS]):
+        print('    {:10s}    {}'.format(label, ', '.join(options)))
+
+
+def style_spectrum_plot(ax, log_freqs, log_powers, grid=True):
     """Apply style and aesthetics to a power spectrum plot.
 
     Parameters
@@ -23,6 +33,8 @@ def style_spectrum_plot(ax, log_freqs, log_powers):
         Whether the frequency axis is plotted in log space.
     log_powers : bool
         Whether the power axis is plotted in log space.
+    grid : bool, optional, default: True
+        Whether to add grid lines to the plot.
     """
 
     # Get labels, based on log status
@@ -33,7 +45,7 @@ def style_spectrum_plot(ax, log_freqs, log_powers):
     ax.set_xlabel(xlabel, fontsize=20)
     ax.set_ylabel(ylabel, fontsize=20)
     ax.tick_params(axis='both', which='major', labelsize=16)
-    ax.grid(True)
+    ax.grid(grid)
 
     # If labels were provided, add a legend
     if ax.get_legend_handles_labels()[0]:
@@ -227,9 +239,24 @@ def style_plot(func, *args, **kwargs):
     By default, this function applies styling with the `apply_style` function. Custom
     functions for applying style can be passed in using `apply_style` as a keyword argument.
 
-    The `apply_style` function calls sub-functions for applying style different plot elements,
-    and these sub-functions can be overridden by passing in alternatives for `axis_styler`,
-    `line_styler`, and `custom_styler`.
+    The `apply_style` function calls sub-functions for applying different plot elements, including:
+
+    - `axis_styler`: apply style options to an axis
+    - `line_styler`: applies style options to lines objects in a plot
+    - `collection_styler`: applies style options to collections objects in a plot
+    - `custom_style`: applies custom style options
+
+    Each of these sub-functions can be overridden by passing in alternatives.
+
+    To see the full set of style arguments that are supported, run the following code:
+
+    >>> from fooof.plts.style import check_style_options
+    >>> check_style_options()
+    Valid style arguments:
+        Axis          title, xlabel, ylabel, xlim, ylim, xticks, yticks, xticklabels, yticklabels
+        Line          alpha, lw, linewidth, ls, linestyle, marker, ms, markersize
+        Collection    alpha, edgecolor
+        Custom        title_fontsize, label_size, tick_labelsize, legend_size, legend_loc
     """
 
     @wraps(func)