Merge pull request #279 from fooof-tools/dataf

[ENH] - Add and use `group_to_dict` function

Author: TomDonoghue

fooof/data/conversions.py

@@ -84,6 +84,32 @@ def model_to_dataframe(fit_results, peak_org):
     return pd.Series(model_to_dict(fit_results, peak_org))
 
 
+def group_to_dict(fit_results, peak_org):
+    """Convert a group of model fit results into a dictionary.
+
+    Parameters
+    ----------
+    fit_results : list of FOOOFResults
+        List of FOOOFResults objects.
+    peak_org : int or Bands
+        How to organize peaks.
+        If int, extracts the first n peaks.
+        If Bands, extracts peaks based on band definitions.
+
+    Returns
+    -------
+    dict
+        Model results organized into a dictionary.
+    """
+
+    fr_dict = {ke : [] for ke in model_to_dict(fit_results[0], peak_org).keys()}
+    for f_res in fit_results:
+        for key, val in model_to_dict(f_res, peak_org).items():
+            fr_dict[key].append(val)
+
+    return fr_dict
+
+
 @check_dependency(pd, 'pandas')
 def group_to_dataframe(fit_results, peak_org):
     """Convert a group of model fit results into a dataframe.
@@ -103,4 +129,4 @@ def group_to_dataframe(fit_results, peak_org):
         Model results organized into a dataframe.
     """
 
-    return pd.DataFrame([model_to_dataframe(f_res, peak_org) for f_res in fit_results])
+    return pd.DataFrame(group_to_dict(fit_results, peak_org))

fooof/tests/data/test_conversions.py

@@ -41,6 +41,17 @@ def test_model_to_dataframe(tresults, tbands, skip_if_no_pandas):
     out = model_to_dataframe(tresults, peak_org=tbands)
     assert isinstance(out, pd.Series)
 
+def test_group_to_dict(tresults, tbands):
+
+    fit_results = [deepcopy(tresults), deepcopy(tresults), deepcopy(tresults)]
+
+    for peak_org in [1, 2, 3]:
+        out = group_to_dict(fit_results, peak_org=peak_org)
+        assert isinstance(out, dict)
+
+    out = group_to_dict(fit_results, peak_org=tbands)
+    assert isinstance(out, dict)
+
 def test_group_to_dataframe(tresults,  tbands, skip_if_no_pandas):
 
     fit_results = [deepcopy(tresults), deepcopy(tresults), deepcopy(tresults)]