diff --git a/fooof/objs/fit.py b/fooof/objs/fit.py index cef4cc7f..df8b002d 100644 --- a/fooof/objs/fit.py +++ b/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,8 +393,13 @@ 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 ----- @@ -402,7 +407,12 @@ def report(self, freqs=None, power_spectrum=None, freq_range=None, """ 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) diff --git a/fooof/plts/fm.py b/fooof/plts/fm.py index 91609887..016a6a7b 100644 --- a/fooof/plts/fm.py +++ b/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: diff --git a/fooof/plts/spectra.py b/fooof/plts/spectra.py index 8141b7cd..82ea646f 100644 --- a/fooof/plts/spectra.py +++ b/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,6 +38,8 @@ 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 @@ -45,7 +47,7 @@ def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=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. """ ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['spectral'])) @@ -53,6 +55,10 @@ def plot_spectra(freqs, power_spectra, log_freqs=False, log_powers=False, # 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: diff --git a/fooof/tests/plts/test_fm.py b/fooof/tests/plts/test_fm.py index 6d7a0f02..d650f84c 100644 --- a/fooof/tests/plts/test_fm.py +++ b/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):