From 3c4fca12261ec0739c7bf6e255a3fa024807afc8 Mon Sep 17 00:00:00 2001 From: Tom Donoghue Date: Wed, 5 Jul 2023 22:49:09 -0400 Subject: [PATCH] add and use group_to_dict data conversion func --- fooof/data/conversions.py | 28 +++++++++++++++++++++++++++- fooof/tests/data/test_conversions.py | 11 +++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/fooof/data/conversions.py b/fooof/data/conversions.py index f2dbfeb7d..66cbb521a 100644 --- a/fooof/data/conversions.py +++ b/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)) diff --git a/fooof/tests/data/test_conversions.py b/fooof/tests/data/test_conversions.py index 08ab30bde..607a6f4ec 100644 --- a/fooof/tests/data/test_conversions.py +++ b/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)]