small doc & code notes tweaks

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,8 +30,8 @@ def group_three(vec):
     if len(vec) % 3 != 0:
         raise ValueError("Wrong size array to group by three.")
 
+    # Reshape, if an array, as it's faster, otherwise asssume lise
     if isinstance(vec, np.ndarray):
-        # Reshaping is faster if already an array
         return np.reshape(vec, (-1, 3))
     else:
         return [list(vec[ii:ii+3]) for ii in range(0, len(vec), 3)]

fooof/objs/utils.py

@@ -219,14 +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
     """
 
-    # Reshape 3d to 2d and fit
+    # 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
+    # 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]) ]
+        for dim_a in range(shape[0])]
 
     return fgs