refact use of avg_func in objs utils

Author: TomDonoghue

fooof/objs/utils.py

@@ -65,18 +65,15 @@ def average_fg(fg, bands, avg_method='mean', regenerate=True):
         If there are no model fit results available to average across.
     """
 
-    if avg_method not in ['mean', 'median']:
-        raise ValueError("Requested average method not understood.")
     if not fg.has_model:
         raise NoModelError("No model fit results are available, can not proceed.")
 
-    if avg_method == 'mean':
-        avg_func = np.nanmean
-    elif avg_method == 'median':
-        avg_func = np.nanmedian
+    avg_funcs = {'mean' : np.nanmean, 'median' : np.nanmedian}
+    if avg_method not in avg_funcs.keys():
+        raise ValueError("Requested average method not understood.")
 
     # Aperiodic parameters: extract & average
-    ap_params = avg_func(fg.get_params('aperiodic_params'), 0)
+    ap_params = avg_funcs[avg_method](fg.get_params('aperiodic_params'), 0)
 
     # Periodic parameters: extract & average
     peak_params = []
@@ -90,15 +87,15 @@ def average_fg(fg, bands, avg_method='mean', regenerate=True):
         # Check if there are any extracted peaks - if not, don't add
         #   Note that we only check peaks, but gauss should be the same
         if not np.all(np.isnan(peaks)):
-            peak_params.append(avg_func(peaks, 0))
-            gauss_params.append(avg_func(gauss, 0))
+            peak_params.append(avg_funcs[avg_method](peaks, 0))
+            gauss_params.append(avg_funcs[avg_method](gauss, 0))
 
     peak_params = np.array(peak_params)
     gauss_params = np.array(gauss_params)
 
     # Goodness of fit measures: extract & average
-    r2 = avg_func(fg.get_params('r_squared'))
-    error = avg_func(fg.get_params('error'))
+    r2 = avg_funcs[avg_method](fg.get_params('r_squared'))
+    error = avg_funcs[avg_method](fg.get_params('error'))
 
     # Collect all results together, to be added to FOOOF object
     results = FOOOFResults(ap_params, peak_params, r2, error, gauss_params)
@@ -135,21 +132,18 @@ def average_reconstructions(fg, avg_method='mean'):
         Note that power values are in log10 space.
     """
 
-    if avg_method not in ['mean', 'median']:
-        raise ValueError("Requested average method not understood.")
     if not fg.has_model:
         raise NoModelError("No model fit results are available, can not proceed.")
 
-    if avg_method == 'mean':
-        avg_func = np.nanmean
-    elif avg_method == 'median':
-        avg_func = np.nanmedian
+    avg_funcs = {'mean' : np.nanmean, 'median' : np.nanmedian}
+    if avg_method not in avg_funcs.keys():
+        raise ValueError("Requested average method not understood.")
 
     models = np.zeros(shape=fg.power_spectra.shape)
     for ind in range(len(fg)):
         models[ind, :] = fg.get_fooof(ind, regenerate=True).fooofed_spectrum_
 
-    avg_model = avg_func(models, 0)
+    avg_model = avg_funcs[avg_method](models, 0)
 
     return fg.freqs, avg_model