Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jul 21, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 17 additions & 6 deletions fooof/objs/fit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)


Expand Down Expand Up @@ -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)
Expand Down
25 changes: 20 additions & 5 deletions fooof/plts/fm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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:
Expand Down
19 changes: 14 additions & 5 deletions fooof/plts/spectra.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)


Expand All @@ -98,7 +106,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
-----
Expand Down Expand Up @@ -148,7 +157,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:
Expand Down
16 changes: 16 additions & 0 deletions fooof/tests/plts/test_fm.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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):

Expand Down