update data & models for Band with n_bands

Author: TomDonoghue

specparam/data/conversions.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from specparam import Bands
+from specparam.bands.bands import Bands, check_bands
 from specparam.modutils.dependencies import safe_import, check_dependency
 from specparam.data.periodic import get_band_peak_arr
 from specparam.data.utils import flatten_results_dict
@@ -21,15 +21,18 @@ def model_to_dict(fit_results, modes, bands):
         Results of a model fit.
     modes : Modes
         Model modes definition.
-    bands : Bands or int
-        How to organize peaks, based on band definitions (Bands) or number of peaks (int).
+    bands : Bands or dict or int
+        How to organize peaks, based on band definitions.
+        Can be Bands object or object that can be converted into a Bands object.
 
     Returns
     -------
     dict
         Model results organized into a dictionary.
     """
 
+    bands = check_bands(bands)
+
     fr_dict = {}
 
     # aperiodic parameters
@@ -38,18 +41,18 @@ def model_to_dict(fit_results, modes, bands):
 
     # periodic parameters
     peaks = fit_results.peak_params
+    if not bands.bands and bands.n_bands:
 
-    if isinstance(bands, int):
-
-        if len(peaks) < bands:
-            nans = [np.array([np.nan] * 3) for ind in range(bands-len(peaks))]
+        # If bands if defined in terms of number of peaks
+        if len(peaks) < bands.n_bands:
+            nans = [np.array([np.nan] * 3) for ind in range(bands.n_bands-len(peaks))]
             peaks = np.vstack((peaks, nans))
 
-        for ind, peak in enumerate(peaks[:bands, :]):
+        for ind, peak in enumerate(peaks[:bands.n_bands, :]):
             for pe_label, pe_param in zip(modes.periodic.params.indices, peak):
                 fr_dict[pe_label + '_' + str(ind)] = pe_param
 
-    elif isinstance(bands, Bands):
+    elif bands.bands:
         for band, f_range in bands:
             for label, param in zip(modes.periodic.params.indices,
                                     get_band_peak_arr(peaks, f_range)):
@@ -72,16 +75,17 @@ def model_to_dataframe(fit_results, modes, bands):
         Results of a model fit.
     modes : Modes
         Model modes definition.
-    bands : Bands or int
-        How to organize peaks, based on band definitions (Bands) or number of peaks (int).
+    bands : Bands or dict or int
+        How to organize peaks, based on band definitions.
+        Can be Bands object or object that can be converted into a Bands object.
 
     Returns
     -------
     pd.Series
         Model results organized into a dataframe.
     """
 
-    return pd.Series(model_to_dict(fit_results, modes, bands))
+    return pd.Series(model_to_dict(fit_results, modes, check_bands(bands)))
 
 
 def group_to_dict(group_results, modes, bands):
@@ -93,15 +97,18 @@ def group_to_dict(group_results, modes, bands):
         List of FitResults objects, reflecting model results across a group of power spectra.
     modes : Modes
         Model modes definition.
-    bands : Bands or int
-        How to organize peaks, based on band definitions (Bands) or number of peaks (int).
+    bands : Bands or dict or int
+        How to organize peaks, based on band definitions.
+        Can be Bands object or object that can be converted into a Bands object.
 
     Returns
     -------
     dict
         Model results organized into a dictionary.
     """
 
+    bands = check_bands(bands)
+
     nres = len(group_results)
     fr_dict = {ke : np.zeros(nres) for ke in model_to_dict(group_results[0], modes, bands)}
     for ind, f_res in enumerate(group_results):
@@ -121,16 +128,17 @@ def group_to_dataframe(group_results, modes, bands):
         List of FitResults objects.
     modes : Modes
         Model modes definition.
-    bands : Bands or int
-        How to organize peaks, based on band definitions (Bands) or number of peaks (int).
+    bands : Bands or dict or int
+        How to organize peaks, based on band definitions.
+        Can be Bands object or object that can be converted into a Bands object.
 
     Returns
     -------
     pd.DataFrame
         Model results organized into a dataframe.
     """
 
-    return pd.DataFrame(group_to_dict(group_results, modes, bands))
+    return pd.DataFrame(group_to_dict(group_results, modes, check_bands(bands)))
 
 
 def event_group_to_dict(event_group_results, modes, bands):
@@ -142,8 +150,9 @@ def event_group_to_dict(event_group_results, modes, bands):
         Model fit results from across a set of events.
     modes : Modes
         Model modes definition.
-    bands : Bands or int
-        How to organize peaks, based on band definitions (Bands) or number of peaks (int).
+    bands : Bands or dict or int
+        How to organize peaks, based on band definitions.
+        Can be Bands object or object that can be converted into a Bands object.
 
     Returns
     -------
@@ -152,6 +161,7 @@ def event_group_to_dict(event_group_results, modes, bands):
     """
 
     event_time_results = {}
+    bands = check_bands(bands)
 
     for key in group_to_dict(event_group_results[0], modes, bands):
         event_time_results[key] = []
@@ -177,8 +187,9 @@ def event_group_to_dataframe(event_group_results, modes, bands):
         List of FitResults objects.
     modes : Modes
         Model modes definition.
-    bands : Bands or int
-        How to organize peaks, based on band definitions (Bands) or number of peaks (int).
+    bands : Bands or dict or int
+        How to organize peaks, based on band definitions.
+        Can be Bands object or object that can be converted into a Bands object.
 
     Returns
     -------
@@ -187,7 +198,7 @@ def event_group_to_dataframe(event_group_results, modes, bands):
     """
 
     return pd.DataFrame(flatten_results_dict(\
-        event_group_to_dict(event_group_results, modes, bands)))
+        event_group_to_dict(event_group_results, modes, check_bands(bands))))
 
 
 @check_dependency(pd, 'pandas')

specparam/models/event.py

@@ -109,10 +109,9 @@ def fit(self, freqs=None, spectrograms=None, freq_range=None, bands=None,
             If 3d array, should have shape [n_events, n_freqs, n_time_windows].
         freq_range : list of [float, float], optional
             Frequency range to fit the model to. If not provided, fits the entire given range.
-        bands : Bands or int, optional
+        bands : Bands or dict or int, optional
             How to organize peaks into bands.
-            If Bands, extracts peaks based on band definitions.
-            If int, extracts the first n peaks.
+            If Bands or dict, uses band definitions. If int, extracts the first 'n' peaks.
         n_jobs : int, optional, default: 1
             Number of jobs to run in parallel.
             1 is no parallelization. -1 uses all available cores.
@@ -172,10 +171,9 @@ def report(self, freqs=None, spectrograms=None, freq_range=None,
             If a 3d array, should have shape [n_events, n_freqs, n_time_windows].
         freq_range : list of [float, float], optional
             Frequency range to fit the model to. If not provided, fits the entire given range.
-        bands : Bands or int, optional
+        bands : Bands or dict or int, optional
             How to organize peaks into bands.
-            If Bands, extracts peaks based on band definitions.
-            If int, extracts the first n peaks.
+            If Bands or dict, uses band definitions. If int, extracts the first 'n' peaks.
         n_jobs : int, optional, default: 1
             Number of jobs to run in parallel.
             1 is no parallelization. -1 uses all available cores.
@@ -385,11 +383,10 @@ def to_df(self, bands=None):
 
         Parameters
         ----------
-        bands : Bands or int, optional
+        bands : Bands or dict or int, optional
             How to organize peaks into bands.
-            If Bands, extracts peaks based on band definitions.
-            If int, extracts the first n peaks.
-            If provided, re-extracts peak features; if not provided, converts from `time_results`.
+            If Bands or dict, uses band definitions. If int, extracts the first 'n' peaks.
+            If provided, re-extracts peak features; if not, converts from `event_group_results`.
 
         Returns
         -------
@@ -410,10 +407,9 @@ def convert_results(self, bands=None):
 
         Parameters
         ----------
-        bands : Bands or int, optional
+        bands : Bands or dict or int, optional
             How to organize peaks into bands.
-            If Bands, extracts peaks based on band definitions.
-            If int, extracts the first 'n' peaks.
+            If Bands or dict, uses band definitions. If int, extracts the first 'n' peaks.
             If not provided, uses band definition available in object.
         """
 

specparam/models/time.py

@@ -78,10 +78,9 @@ def fit(self, freqs=None, spectrogram=None, freq_range=None, bands=None,
             Spectrogram of power spectrum values, in linear space.
         freq_range : list of [float, float], optional
             Frequency range to fit the model to. If not provided, fits the entire given range.
-        bands : Bands or int, optional
+        bands : Bands or dict or int, optional
             How to organize peaks into bands.
-            If Bands, extracts peaks based on band definitions.
-            If int, extracts the first n peaks.
+            If Bands or dict, uses band definitions. If int, extracts the first 'n' peaks.
         n_jobs : int, optional, default: 1
             Number of jobs to run in parallel.
             1 is no parallelization. -1 uses all available cores.
@@ -121,10 +120,9 @@ def report(self, freqs=None, spectrogram=None, freq_range=None,
             Spectrogram of power spectrum values, in linear space.
         freq_range : list of [float, float], optional
             Frequency range to fit the model to. If not provided, fits the entire given range.
-        bands : Bands or int, optional
+        bands : Bands or dict or int, optional
             How to organize peaks into bands.
-            If Bands, extracts peaks based on band definitions.
-            If int, extracts the first n peaks.
+            If Bands or dict, uses band definitions. If int, extracts the first 'n' peaks.
         n_jobs : int, optional, default: 1
             Number of jobs to run in parallel.
             1 is no parallelization. -1 uses all available cores.
@@ -251,11 +249,10 @@ def to_df(self, bands=None):
 
         Parameters
         ----------
-        bands : Bands or int, optional
+        bands : Bands or dict or int, optional
             How to organize peaks into bands.
-            If Bands, extracts peaks based on band definitions.
-            If int, extracts the first n peaks.
-            If provided, re-extracts peak features; if not provided, converts from `time_results`.
+            If Bands or dict, uses band definitions. If int, extracts the first 'n' peaks.
+            If provided, re-extracts peak features; if not, converts from `time_results`.
 
         Returns
         -------
@@ -276,10 +273,9 @@ def convert_results(self, bands):
 
         Parameters
         ----------
-        bands : Bands or int, optional
+        bands : Bands or dict or int, optional
             How to organize peaks into bands.
-            If Bands, extracts peaks based on band definitions.
-            If int, extracts the first 'n' peaks.
+            If Bands or dict, uses band definitions. If int, extracts the first 'n' peaks.
             If not provided, uses band definition available in object.
         """
 

specparam/tests/data/test_conversions.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 
+from specparam.bands import Bands
 from specparam.modutils.dependencies import safe_import
 pd = safe_import('pandas')
 
@@ -14,18 +15,18 @@
 
 def test_model_to_dict(tresults, tmodes, tbands):
 
-    out = model_to_dict(tresults, tmodes, 1)
+    out = model_to_dict(tresults, tmodes, Bands(n_bands=1))
     assert isinstance(out, dict)
     assert 'cf_0' in out
     assert out['cf_0'] == tresults.peak_params[0, 0]
     assert 'cf_1' not in out
 
-    out = model_to_dict(tresults, tmodes, 2)
+    out = model_to_dict(tresults, tmodes, Bands(n_bands=2))
     assert 'cf_0' in out
     assert 'cf_1' in out
     assert out['cf_1'] == tresults.peak_params[1, 0]
 
-    out = model_to_dict(tresults, tmodes, 3)
+    out = model_to_dict(tresults, tmodes, Bands(n_bands=3))
     assert 'cf_2' in out
     assert np.isnan(out['cf_2'])
 
@@ -35,7 +36,7 @@ def test_model_to_dict(tresults, tmodes, tbands):
 def test_model_to_dataframe(tresults, tmodes, tbands, skip_if_no_pandas):
 
     for nbands in [1, 2, 3]:
-        out = model_to_dataframe(tresults, tmodes, nbands)
+        out = model_to_dataframe(tresults, tmodes, Bands(n_bands=nbands))
         assert isinstance(out, pd.Series)
 
     out = model_to_dataframe(tresults, tmodes, tbands)
@@ -46,7 +47,7 @@ def test_group_to_dict(tresults, tmodes, tbands):
     fit_results = [deepcopy(tresults), deepcopy(tresults), deepcopy(tresults)]
 
     for nbands in [1, 2, 3]:
-        out = group_to_dict(fit_results, tmodes, nbands)
+        out = group_to_dict(fit_results, tmodes, Bands(n_bands=nbands))
         assert isinstance(out, dict)
 
     out = group_to_dict(fit_results, tmodes, tbands)
@@ -57,7 +58,7 @@ def test_group_to_dataframe(tresults, tmodes, tbands, skip_if_no_pandas):
     fit_results = [deepcopy(tresults), deepcopy(tresults), deepcopy(tresults)]
 
     for nbands in [1, 2, 3]:
-        out = group_to_dataframe(fit_results, tmodes, nbands)
+        out = group_to_dataframe(fit_results, tmodes, Bands(n_bands=nbands))
         assert isinstance(out, pd.DataFrame)
 
     out = group_to_dataframe(fit_results, tmodes, tbands)
@@ -69,7 +70,7 @@ def test_event_group_to_dict(tresults, tmodes, tbands):
                    [deepcopy(tresults), deepcopy(tresults), deepcopy(tresults)]]
 
     for nbands in [1, 2, 3]:
-        out = event_group_to_dict(fit_results, tmodes, nbands)
+        out = event_group_to_dict(fit_results, tmodes, Bands(n_bands=nbands))
         assert isinstance(out, dict)
 
     out = event_group_to_dict(fit_results, tmodes, tbands)
@@ -81,7 +82,7 @@ def test_event_group_to_dataframe(tresults, tmodes, tbands, skip_if_no_pandas):
                    [deepcopy(tresults), deepcopy(tresults), deepcopy(tresults)]]
 
     for nbands in [1, 2, 3]:
-        out = event_group_to_dataframe(fit_results, tmodes, nbands)
+        out = event_group_to_dataframe(fit_results, tmodes, Bands(n_bands=nbands))
         assert isinstance(out, pd.DataFrame)
 
     out = event_group_to_dataframe(fit_results, tmodes, tbands)