add option to pass custom freqs & power spectrum to plot_fm

Author: TomDonoghue

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/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(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):