Merge pull request #276 from fooof-tools/plotfull

[ENH] - Update options to plot a broader frequency range

Author: TomDonoghue

fooof/objs/fit.py

@@ -379,7 +379,7 @@ def add_results(self, fooof_result):
 
 
     def report(self, freqs=None, power_spectrum=None, freq_range=None,
-               plt_log=False, **plot_kwargs):
+               plt_log=False, plot_full_range=False, **plot_kwargs):
         """Run model fit, and display a report, which includes a plot, and printed results.
 
         Parameters
@@ -393,16 +393,26 @@ def report(self, freqs=None, power_spectrum=None, freq_range=None,
             If not provided, fits across the entire given range.
         plt_log : bool, optional, default: False
             Whether or not to plot the frequency axis in log space.
+        plot_full_range : bool, default: False
+            If True, plots the full range of the given power spectrum.
+            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_'.
+            e.g. `freqs`, `power_spectrum` and `freq_range`.
 
         Notes
         -----
         Data is optional, if data has already been added to the object.
         """
 
         self.fit(freqs, power_spectrum, freq_range)
-        self.plot(plt_log=plt_log, **plot_kwargs)
+        self.plot(plt_log=plt_log,
+                  freqs=freqs if plot_full_range else plot_kwargs.pop('plot_freqs', None),
+                  power_spectrum=power_spectrum if \
+                      plot_full_range else plot_kwargs.pop('plot_power_spectrum', None),
+                  freq_range=plot_kwargs.pop('plot_freq_range', None),
+                  **plot_kwargs)
         self.print_results(concise=False)
 
 
@@ -639,12 +649,13 @@ def get_results(self):
 
 
     @copy_doc_func_to_method(plot_fm)
-    def plot(self, plot_peaks=None, plot_aperiodic=True, plt_log=False,
-             add_legend=True, save_fig=False, file_name=None, file_path=None,
-             ax=None, data_kwargs=None, model_kwargs=None,
+    def plot(self, plot_peaks=None, plot_aperiodic=True, freqs=None, power_spectrum=None,
+             freq_range=None, plt_log=False, add_legend=True, save_fig=False, file_name=None,
+             file_path=None, ax=None, data_kwargs=None, model_kwargs=None,
              aperiodic_kwargs=None, peak_kwargs=None, **plot_kwargs):
 
-        plot_fm(self, plot_peaks=plot_peaks, plot_aperiodic=plot_aperiodic, plt_log=plt_log,
+        plot_fm(self, plot_peaks=plot_peaks, plot_aperiodic=plot_aperiodic, freqs=freqs,
+                power_spectrum=power_spectrum, freq_range=freq_range, plt_log=plt_log,
                 add_legend=add_legend, save_fig=save_fig, file_name=file_name,
                 file_path=file_path, ax=ax, data_kwargs=data_kwargs, model_kwargs=model_kwargs,
                 aperiodic_kwargs=aperiodic_kwargs, peak_kwargs=peak_kwargs, **plot_kwargs)

fooof/plts/fm.py

@@ -25,9 +25,10 @@
 @savefig
 @style_plot
 @check_dependency(plt, 'matplotlib')
-def plot_fm(fm, plot_peaks=None, plot_aperiodic=True, plt_log=False, add_legend=True,
-            save_fig=False, file_name=None, file_path=None, ax=None, data_kwargs=None,
-            model_kwargs=None, aperiodic_kwargs=None, peak_kwargs=None, **plot_kwargs):
+def plot_fm(fm, plot_peaks=None, plot_aperiodic=True, freqs=None, power_spectrum=None,
+            freq_range=None, plt_log=False, add_legend=True, save_fig=False, file_name=None,
+            file_path=None, ax=None, data_kwargs=None, model_kwargs=None, aperiodic_kwargs=None,
+            peak_kwargs=None, **plot_kwargs):
     """Plot the power spectrum and model fit results from a FOOOF object.
 
     Parameters
@@ -39,6 +40,14 @@ def plot_fm(fm, plot_peaks=None, plot_aperiodic=True, plt_log=False, add_legend=
         Can also be a combination of approaches, separated by '-', for example: 'shade-line'.
     plot_aperiodic : boolean, optional, default: True
         Whether to plot the aperiodic component of the model fit.
+    freqs : 1d array, optional
+        Frequency values of the power spectrum to plot, in linear space.
+        If provided, this overrides the values in the model object.
+    power_spectrum : 1d array, optional
+        Power values to plot, in linear space.
+        If provided, this overrides the values in the model object.
+    freq_range : list of [float, float], optional
+        Frequency range to plot, defined in linear space.
     plt_log : boolean, optional, default: False
         Whether to plot the frequency values in log10 spacing.
     add_legend : boolean, optional, default: False
@@ -64,16 +73,22 @@ def plot_fm(fm, plot_peaks=None, plot_aperiodic=True, plt_log=False, add_legend=
 
     ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['spectral']))
 
+    # Check inputs for what to plot
+    custom_spectrum = (np.any(freqs) and np.any(power_spectrum))
+
     # Log settings - note that power values in FOOOF objects are already logged
     log_freqs = plt_log
     log_powers = False
 
     # Plot the data, if available
-    if fm.has_data:
+    if fm.has_data or custom_spectrum:
         data_defaults = {'color' : PLT_COLORS['data'], 'linewidth' : 2.0,
                          'label' : 'Original Spectrum' if add_legend else None}
         data_kwargs = check_plot_kwargs(data_kwargs, data_defaults)
-        plot_spectra(fm.freqs, fm.power_spectrum, log_freqs, log_powers, ax=ax, **data_kwargs)
+        plot_spectra(freqs if custom_spectrum else fm.freqs,
+                     power_spectrum if custom_spectrum else fm.power_spectrum,
+                     log_freqs, log_powers if not custom_spectrum else True,
+                     freq_range, ax=ax, **data_kwargs)
 
     # Add the full model fit, and components (if requested)
     if fm.has_model:

fooof/plts/spectra.py

@@ -24,7 +24,7 @@
 @savefig
 @style_plot
 @check_dependency(plt, 'matplotlib')
-def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=False,
+def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=False, freq_range=None,
                  colors=None, labels=None, ax=None, **plot_kwargs):
     """Plot one or multiple power spectra.
 
@@ -38,21 +38,27 @@ def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=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.
+    freq_range : list of [float, float], optional
+        Frequency range to plot, defined in linear space.
     colors : list of str, optional, default: None
         Line colors of the spectra.
     labels : list of str, optional, default: None
         Legend labels for the spectra.
     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.
     """
 
     ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['spectral']))
 
     # Create the plot
     plot_kwargs = check_plot_kwargs(plot_kwargs, {'linewidth' : 2.0})
 
+    # Check for frequency range input, and log if x-axis is in log space
+    if freq_range is not None:
+        freq_range = np.log10(freq_range) if log_freqs else freq_range
+
     # Make inputs iterable if need to be passed multiple times to plot each spectrum
     plt_powers = np.reshape(power_spectra, (1, -1)) if np.ndim(power_spectra) == 1 else \
         power_spectra
@@ -63,7 +69,7 @@ def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=False,
     labels = repeat(labels) if not isinstance(labels, list) else cycle(labels)
     colors = repeat(colors) if not isinstance(colors, list) else cycle(colors)
 
-    # Plot
+    # Plot power spectra, looping across all spectra to plot
     for freqs, powers, color, label in zip(plt_freqs, plt_powers, colors, labels):
 
         # Set plot data, logging if requested, and collect color, if absent
@@ -74,6 +80,8 @@ def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=False,
 
         ax.plot(freqs, powers, label=label, **plot_kwargs)
 
+    ax.set_xlim(freq_range)
+
     style_spectrum_plot(ax, log_freqs, log_powers)
 
 
@@ -102,7 +110,8 @@ def plot_spectra_shading(freqs, power_spectra, shades, shade_colors='r',
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Keyword arguments to pass into :func:`~.plot_spectra`.
+        Additional plot related keyword arguments.
+        This can include additional inputs into :func:`~.plot_spectra`.
 
     Notes
     -----
@@ -152,7 +161,7 @@ def plot_spectra_yshade(freqs, power_spectra, shade='std', average='mean', scale
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **plot_kwargs
-        Keyword arguments to be passed to `plot_spectra` or to the plot call.
+        Additional plot related keyword arguments.
     """
 
     if (isinstance(shade, str) or isfunction(shade)) and power_spectra.ndim != 2:

fooof/tests/plts/test_fm.py

@@ -1,5 +1,7 @@
 """Tests for fooof.plts.fm."""
 
+import numpy as np
+
 from fooof.tests.tutils import plot_test
 from fooof.tests.settings import TEST_PLOTS_PATH
 
@@ -17,6 +19,20 @@ def test_plot_fm(tfm, skip_if_no_mpl):
     plot_fm(tfm, save_fig=True, file_path=TEST_PLOTS_PATH,
             file_name='test_plot_fm.png')
 
+@plot_test
+def test_plot_fm_custom(tfm, skip_if_no_mpl):
+
+    # Extract broader range of data available in the object
+    custom_freqs = tfm.freqs
+    custom_power_spectrum = np.power(10, tfm.power_spectrum)
+
+    # Make sure model has been fit - set custom frequency range
+    tfm.fit(custom_freqs, custom_power_spectrum, freq_range=[5, 35])
+
+    plot_fm(tfm, freqs=custom_freqs, power_spectrum=custom_power_spectrum,
+            freq_range=[1, 55], save_fig=True, file_path=TEST_PLOTS_PATH,
+            file_name='test_plot_fm_custom.png')
+
 @plot_test
 def test_plot_fm_add_peaks(tfm, skip_if_no_mpl):