Move group_data and reset_groups from survey

niimpy/preprocessing/survey.py

@@ -2,6 +2,7 @@
 
 import pandas as pd
 import numpy as np
+from niimpy.preprocessing import util
 
 
 # Below, we provide some mappings between standardized survey raw questions and their respective codes
@@ -116,17 +117,6 @@
 
 group_by_columns = set(["user", "device"])
 
-def group_data(df):
-    """ Group the dataframe by a standard set of columns listed in
-    group_by_columns."""
-    columns = list(group_by_columns & set(df.columns))
-    return df.groupby(columns)
-
-def reset_groups(df):
-    """ Reset the grouping, keeping only the original index columns. """
-    columns = list(group_by_columns & set(df.index.names))
-    return df.reset_index(columns)
-
 
 def clean_survey_column_names(df):
     """
@@ -265,8 +255,8 @@ def calculate_statistic(df):
             result[answer_col+"_std"] = df[answer_col].std()
         return pd.Series(result)
 
-    res = group_data(df).resample(**resample_args).apply(calculate_statistic)
-    res = reset_groups(res)
+    res = util.group_data(df).resample(**resample_args).apply(calculate_statistic)
+    res = util.reset_groups(res)
     return res
 
 
@@ -353,5 +343,5 @@ def extract_features_survey(df, features=None):
     if 'group' in df:
         computed_features['group'] = df.groupby('user')['group'].first()
 
-    computed_features = reset_groups(computed_features)
+    computed_features = util.reset_groups(computed_features)
     return computed_features