Merge branch 'main' into plts

Author: TomDonoghue

fooof/core/utils.py

@@ -13,13 +13,13 @@ def group_three(vec):
 
     Parameters
     ----------
-    vec : 1d array
-        Array of items to group by 3. Length of array must be divisible by three.
+    vec : list or 1d array
+        List or array of items to group by 3. Length of array must be divisible by three.
 
     Returns
     -------
-    list of list
-        List of lists, each with three items.
+    array or list of list
+        Array or list of lists, each with three items. Output type will match input type.
 
     Raises
     ------
@@ -30,7 +30,11 @@ def group_three(vec):
     if len(vec) % 3 != 0:
         raise ValueError("Wrong size array to group by three.")
 
-    return [list(vec[ii:ii+3]) for ii in range(0, len(vec), 3)]
+    # Reshape, if an array, as it's faster, otherwise asssume lise
+    if isinstance(vec, np.ndarray):
+        return np.reshape(vec, (-1, 3))
+    else:
+        return [list(vec[ii:ii+3]) for ii in range(0, len(vec), 3)]
 
 
 def nearest_ind(array, value):

fooof/objs/fit.py

@@ -74,7 +74,6 @@
                                 gen_issue_str, gen_width_warning_str)
 
 from fooof.plts.fm import plot_fm
-from fooof.plts.style import style_spectrum_plot
 from fooof.utils.data import trim_spectrum
 from fooof.utils.params import compute_gauss_std
 from fooof.data import FOOOFResults, FOOOFSettings, FOOOFMetaData
@@ -633,12 +632,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, plot_style=style_spectrum_plot,
-             data_kwargs=None, model_kwargs=None, aperiodic_kwargs=None, peak_kwargs=None):
+             ax=None, data_kwargs=None, model_kwargs=None,
+             aperiodic_kwargs=None, peak_kwargs=None, **plot_kwargs):
 
-        plot_fm(self, plot_peaks, plot_aperiodic, plt_log, add_legend,
-                save_fig, file_name, file_path, ax, plot_style,
-                data_kwargs, model_kwargs, aperiodic_kwargs, peak_kwargs)
+        plot_fm(self, plot_peaks=plot_peaks, plot_aperiodic=plot_aperiodic, 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)
 
 
     @copy_doc_func_to_method(save_report_fm)
@@ -1005,18 +1005,16 @@ def _create_peak_params(self, gaus_params):
         with `freqs`, `fooofed_spectrum_` and `_ap_fit` all required to be available.
         """
 
-        peak_params = np.empty([0, 3])
+        peak_params = np.empty((len(gaus_params), 3))
 
         for ii, peak in enumerate(gaus_params):
 
             # Gets the index of the power_spectrum at the frequency closest to the CF of the peak
-            ind = min(range(len(self.freqs)), key=lambda ii: abs(self.freqs[ii] - peak[0]))
+            ind = np.argmin(np.abs(self.freqs - peak[0]))
 
             # Collect peak parameter data
-            peak_params = np.vstack((peak_params,
-                                     [peak[0],
-                                      self.fooofed_spectrum_[ind] - self._ap_fit[ind],
-                                      peak[2] * 2]))
+            peak_params[ii] = [peak[0], self.fooofed_spectrum_[ind] - self._ap_fit[ind],
+                               peak[2] * 2]
 
         return peak_params
 
@@ -1035,8 +1033,8 @@ def _drop_peak_cf(self, guess):
             Guess parameters for gaussian peak fits. Shape: [n_peaks, 3].
         """
 
-        cf_params = [item[0] for item in guess]
-        bw_params = [item[2] * self._bw_std_edge for item in guess]
+        cf_params = guess[:, 0]
+        bw_params = guess[:, 2] * self._bw_std_edge
 
         # Check if peaks within drop threshold from the edge of the frequency range
         keep_peak = \

fooof/objs/group.py

@@ -396,9 +396,9 @@ def get_params(self, name, col=None):
 
 
     @copy_doc_func_to_method(plot_fg)
-    def plot(self, save_fig=False, file_name=None, file_path=None):
+    def plot(self, save_fig=False, file_name=None, file_path=None, **plot_kwargs):
 
-        plot_fg(self, save_fig, file_name, file_path)
+        plot_fg(self, save_fig=save_fig, file_name=file_name, file_path=file_path, **plot_kwargs)
 
 
     @copy_doc_func_to_method(save_report_fg)

fooof/objs/utils.py

@@ -219,9 +219,14 @@ def fit_fooof_3d(fg, freqs, power_spectra, freq_range=None, n_jobs=1):
     >>> fgs = fit_fooof_3d(fg, freqs, power_spectra, freq_range=[3, 30])  # doctest:+SKIP
     """
 
-    fgs = []
-    for cond_spectra in power_spectra:
-        fg.fit(freqs, cond_spectra, freq_range, n_jobs)
-        fgs.append(fg.copy())
+    # Reshape 3d data to 2d and fit, in order to fit with a single group model object
+    shape = np.shape(power_spectra)
+    powers_2d = np.reshape(power_spectra, (shape[0] * shape[1], shape[2]))
+
+    fg.fit(freqs, powers_2d, freq_range, n_jobs)
+
+    # Reorganize 2d results into a list of model group objects, to reflect original shape
+    fgs = [fg.get_group(range(dim_a * shape[1], (dim_a + 1) * shape[1])) \
+        for dim_a in range(shape[0])]
 
     return fgs

fooof/plts/annotate.py

@@ -6,30 +6,31 @@
 from fooof.core.errors import NoModelError
 from fooof.core.funcs import gaussian_function
 from fooof.core.modutils import safe_import, check_dependency
+
 from fooof.sim.gen import gen_aperiodic
-from fooof.plts.utils import check_ax
-from fooof.plts.spectra import plot_spectra
-from fooof.plts.settings import PLT_FIGSIZES, PLT_COLORS
-from fooof.plts.style import check_n_style, style_spectrum_plot
 from fooof.analysis.periodic import get_band_peak_fm
 from fooof.utils.params import compute_knee_frequency, compute_fwhm
 
+from fooof.plts.spectra import plot_spectra
+from fooof.plts.utils import check_ax, savefig
+from fooof.plts.settings import PLT_FIGSIZES, PLT_COLORS
+from fooof.plts.style import style_spectrum_plot
+
 plt = safe_import('.pyplot', 'matplotlib')
 mpatches = safe_import('.patches', 'matplotlib')
 
 ###################################################################################################
 ###################################################################################################
 
+@savefig
 @check_dependency(plt, 'matplotlib')
-def plot_annotated_peak_search(fm, plot_style=style_spectrum_plot):
+def plot_annotated_peak_search(fm):
     """Plot a series of plots illustrating the peak search from a flattened spectrum.
 
     Parameters
     ----------
     fm : FOOOF
         FOOOF object, with model fit, data and settings available.
-    plot_style : callable, optional, default: style_spectrum_plot
-        A function to call to apply styling & aesthetics to the plots.
     """
 
     # Recalculate the initial aperiodic fit and flattened spectrum that
@@ -46,14 +47,12 @@ def plot_annotated_peak_search(fm, plot_style=style_spectrum_plot):
         # This forces the creation of a new plotting axes per iteration
         ax = check_ax(None, PLT_FIGSIZES['spectral'])
 
-        plot_spectra(fm.freqs, flatspec, ax=ax, plot_style=None,
-                     label='Flattened Spectrum', color=PLT_COLORS['data'], linewidth=2.5)
-        plot_spectra(fm.freqs, [fm.peak_threshold * np.std(flatspec)]*len(fm.freqs),
-                     ax=ax, plot_style=None, label='Relative Threshold',
-                     color='orange', linewidth=2.5, linestyle='dashed')
-        plot_spectra(fm.freqs, [fm.min_peak_height]*len(fm.freqs),
-                     ax=ax, plot_style=None, label='Absolute Threshold',
-                     color='red', linewidth=2.5, linestyle='dashed')
+        plot_spectra(fm.freqs, flatspec, ax=ax, linewidth=2.5,
+                     label='Flattened Spectrum', color=PLT_COLORS['data'])
+        plot_spectra(fm.freqs, [fm.peak_threshold * np.std(flatspec)]*len(fm.freqs), ax=ax,
+                     label='Relative Threshold', color='orange', linewidth=2.5, linestyle='dashed')
+        plot_spectra(fm.freqs, [fm.min_peak_height]*len(fm.freqs), ax=ax,
+                     label='Absolute Threshold', color='red', linewidth=2.5, linestyle='dashed')
 
         maxi = np.argmax(flatspec)
         ax.plot(fm.freqs[maxi], flatspec[maxi], '.',
@@ -65,18 +64,18 @@ def plot_annotated_peak_search(fm, plot_style=style_spectrum_plot):
         if ind < fm.n_peaks_:
 
             gauss = gaussian_function(fm.freqs, *fm.gaussian_params_[ind, :])
-            plot_spectra(fm.freqs, gauss, ax=ax, plot_style=None,
-                         label='Gaussian Fit', color=PLT_COLORS['periodic'],
-                         linestyle=':', linewidth=3.0)
+            plot_spectra(fm.freqs, gauss, ax=ax, label='Gaussian Fit',
+                         color=PLT_COLORS['periodic'], linestyle=':', linewidth=3.0)
 
             flatspec = flatspec - gauss
 
-        check_n_style(plot_style, ax, False, True)
+        style_spectrum_plot(ax, False, True)
 
 
+@savefig
 @check_dependency(plt, 'matplotlib')
-def plot_annotated_model(fm, plt_log=False, annotate_peaks=True, annotate_aperiodic=True,
-                         ax=None, plot_style=style_spectrum_plot):
+def plot_annotated_model(fm, plt_log=False, annotate_peaks=True,
+                         annotate_aperiodic=True, ax=None):
     """Plot a an annotated power spectrum and model, from a FOOOF object.
 
     Parameters
@@ -91,8 +90,6 @@ def plot_annotated_model(fm, plt_log=False, annotate_peaks=True, annotate_aperio
         Whether to annotate the aperiodic components of the model fit.
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
-    plot_style : callable, optional, default: style_spectrum_plot
-        A function to call to apply styling & aesthetics to the plots.
 
     Raises
     ------
@@ -112,7 +109,7 @@ def plot_annotated_model(fm, plt_log=False, annotate_peaks=True, annotate_aperio
 
     # Create the baseline figure
     ax = check_ax(ax, PLT_FIGSIZES['spectral'])
-    fm.plot(plot_peaks='dot-shade-width', plt_log=plt_log, ax=ax, plot_style=None,
+    fm.plot(plot_peaks='dot-shade-width', plt_log=plt_log, ax=ax,
             data_kwargs={'lw' : lw1, 'alpha' : 0.6},
             aperiodic_kwargs={'lw' : lw1, 'zorder' : 10},
             model_kwargs={'lw' : lw1, 'alpha' : 0.5},
@@ -219,7 +216,7 @@ def plot_annotated_model(fm, plt_log=False, annotate_peaks=True, annotate_aperio
                     color=PLT_COLORS['aperiodic'], fontsize=fontsize)
 
     # Apply style to plot & tune grid styling
-    check_n_style(plot_style, ax, plt_log, True)
+    style_spectrum_plot(ax, plt_log, True)
     ax.grid(True, alpha=0.5)
 
     # Add labels to plot in the legend

fooof/plts/aperiodic.py

@@ -3,21 +3,23 @@
 from itertools import cycle
 
 import numpy as np
+import matplotlib.pyplot as plt
 
 from fooof.sim.gen import gen_freqs, gen_aperiodic
 from fooof.core.modutils import safe_import, check_dependency
 from fooof.plts.settings import PLT_FIGSIZES
-from fooof.plts.style import check_n_style, style_param_plot
-from fooof.plts.utils import check_ax, recursive_plot, check_plot_kwargs
+from fooof.plts.style import style_param_plot, style_plot
+from fooof.plts.utils import check_ax, recursive_plot, savefig, check_plot_kwargs
 
 plt = safe_import('.pyplot', 'matplotlib')
 
 ###################################################################################################
 ###################################################################################################
 
+@savefig
+@style_plot
 @check_dependency(plt, 'matplotlib')
-def plot_aperiodic_params(aps, colors=None, labels=None,
-                          ax=None, plot_style=style_param_plot, **plot_kwargs):
+def plot_aperiodic_params(aps, colors=None, labels=None, ax=None, **plot_kwargs):
     """Plot aperiodic parameters as dots representing offset and exponent value.
 
     Parameters
@@ -30,38 +32,38 @@ def plot_aperiodic_params(aps, colors=None, labels=None,
         Label(s) for plotted data, to be added in a legend.
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
-    plot_style : callable, optional, default: style_param_plot
-        A function to call to apply styling & aesthetics to the plot.
     **plot_kwargs
-        Keyword arguments to pass into the plot call.
+        Keyword arguments to pass into the ``style_plot``.
     """
 
     ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['params']))
 
     if isinstance(aps, list):
-        recursive_plot(aps, plot_aperiodic_params, ax, colors=colors, labels=labels,
-                       plot_style=plot_style, **plot_kwargs)
+        recursive_plot(aps, plot_aperiodic_params, ax, colors=colors, labels=labels)
 
     else:
 
         # Unpack data: offset as x; exponent as y
         xs, ys = aps[:, 0], aps[:, -1]
         sizes = plot_kwargs.pop('s', 150)
 
+        # Create the plot
         plot_kwargs = check_plot_kwargs(plot_kwargs, {'alpha' : 0.7})
         ax.scatter(xs, ys, sizes, c=colors, label=labels, **plot_kwargs)
 
     # Add axis labels
     ax.set_xlabel('Offset')
     ax.set_ylabel('Exponent')
 
-    check_n_style(plot_style, ax)
+    style_param_plot(ax)
 
 
+@savefig
+@style_plot
 @check_dependency(plt, 'matplotlib')
 def plot_aperiodic_fits(aps, freq_range, control_offset=False,
                         log_freqs=False, colors=None, labels=None,
-                        ax=None, plot_style=style_param_plot, **plot_kwargs):
+                        ax=None, **plot_kwargs):
     """Plot reconstructions of model aperiodic fits.
 
     Parameters
@@ -80,10 +82,8 @@ def plot_aperiodic_fits(aps, freq_range, control_offset=False,
         Label(s) for plotted data, to be added in a legend.
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
-    plot_style : callable, optional, default: style_param_plot
-        A function to call to apply styling & aesthetics to the plot.
     **plot_kwargs
-        Keyword arguments to pass into the plot call.
+        Keyword arguments to pass into the ``style_plot``.
     """
 
     ax = check_ax(ax, plot_kwargs.pop('figsize', PLT_FIGSIZES['params']))
@@ -93,11 +93,9 @@ def plot_aperiodic_fits(aps, freq_range, control_offset=False,
         if not colors:
             colors = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
 
-        recursive_plot(aps, plot_function=plot_aperiodic_fits, ax=ax,
-                       freq_range=tuple(freq_range),
-                       control_offset=control_offset,
-                       log_freqs=log_freqs, colors=colors, labels=labels,
-                       plot_style=plot_style, **plot_kwargs)
+        recursive_plot(aps, plot_aperiodic_fits, ax=ax, freq_range=tuple(freq_range),
+                       control_offset=control_offset, log_freqs=log_freqs, colors=colors,
+                       labels=labels, **plot_kwargs)
     else:
 
         freqs = gen_freqs(freq_range, 0.1)
@@ -118,17 +116,15 @@ def plot_aperiodic_fits(aps, freq_range, control_offset=False,
             # Recreate & plot the aperiodic component from parameters
             ap_vals = gen_aperiodic(freqs, ap_params)
 
-            plot_kwargs = check_plot_kwargs(plot_kwargs, {'alpha' : 0.35, 'linewidth' : 1.25})
-            ax.plot(plt_freqs, ap_vals, color=colors, **plot_kwargs)
+            ax.plot(plt_freqs, ap_vals, color=colors, alpha=0.35, linewidth=1.25)
 
             # Collect a running average across components
             avg_vals = np.nansum(np.vstack([avg_vals, ap_vals]), axis=0)
 
         # Plot the average component
         avg = avg_vals / aps.shape[0]
         avg_color = 'black' if not colors else colors
-        ax.plot(plt_freqs, avg, linewidth=plot_kwargs.get('linewidth')*3,
-                color=avg_color, label=labels)
+        ax.plot(plt_freqs, avg, linewidth=3.75, color=avg_color, label=labels)
 
     # Add axis labels
     ax.set_xlabel('log(Frequency)' if log_freqs else 'Frequency')
@@ -137,5 +133,4 @@ def plot_aperiodic_fits(aps, freq_range, control_offset=False,
     # Set plot limit
     ax.set_xlim(np.log10(freq_range) if log_freqs else freq_range)
 
-    # Apply plot style
-    check_n_style(plot_style, ax)
+    style_param_plot(ax)