feat(datasets): add aggregate, basic and reshape classes, rename tran…

…sform to combine, process to plot, filter to slice

src/hyfi/utils/datasets/aggregate.py

@@ -0,0 +1,11 @@
+"""
+This module contains the class for aggregating datasets.
+"""
+from hyfi.utils.logging import LOGGING
+
+logger = LOGGING.getLogger(__name__)
+
+
+class DSAggregate:
+    def __init__(self):
+        pass

src/hyfi/utils/datasets/basic.py

@@ -0,0 +1,11 @@
+"""
+This file contains the basic dataset functions.
+"""
+from hyfi.utils.logging import LOGGING
+
+logger = LOGGING.getLogger(__name__)
+
+
+class DSBasic:
+    def __init__(self):
+        pass

src/hyfi/utils/datasets/transform.py → src/hyfi/utils/datasets/combine.py

@@ -1,3 +1,6 @@
+"""
+Dataset transformation functions. Concatenate, merge, join, etc.
+"""
 from typing import Dict, List, Optional, Sequence, Union
 
 import datasets as hfds
@@ -13,7 +16,7 @@
 logger = LOGGING.getLogger(__name__)
 
 
-class DSTransform:
+class DSCombine:
     @staticmethod
     def concatenate_data(
         data: Union[Dict[str, pd.DataFrame], Sequence[pd.DataFrame], List[DatasetType]],
@@ -28,14 +31,14 @@ def concatenate_data(
     ) -> Union[pd.DataFrame, DatasetType]:
         # if data is a list of datasets, concatenate them
         if isinstance(data, list) and isinstance(data[0], Dataset):
-            return DSTransform.concatenate_datasets(
+            return DSCombine.concatenate_datasets(
                 data,
                 axis=axis,
                 split=split,
                 **kwargs,
             )
         else:
-            return DSTransform.concatenate_dataframes(
+            return DSCombine.concatenate_dataframes(
                 data,
                 columns=columns,
                 add_split_key_column=add_split_key_column,

src/hyfi/utils/datasets/process.py → src/hyfi/utils/datasets/plot.py

@@ -1,8 +1,11 @@
+"""
+Plotting functions for datasets.
+"""
 from hyfi.utils.logging import LOGGING
 
 logger = LOGGING.getLogger(__name__)
 
 
-class DSProcess:
+class DSPlot:
     def __init__(self):
         pass

src/hyfi/utils/datasets/reshape.py

@@ -0,0 +1,11 @@
+"""
+This module contains the class for reshaping the dataset. Pivot, melt, etc.
+"""
+from hyfi.utils.logging import LOGGING
+
+logger = LOGGING.getLogger(__name__)
+
+
+class DSReshape:
+    def __init__(self):
+        pass

src/hyfi/utils/datasets/filter.py → src/hyfi/utils/datasets/slice.py

@@ -1,6 +1,10 @@
+"""
+Filter datasets. Slice, sample, and filter datasets.
+"""
 import random
-from typing import List, Optional, Union
+from typing import List, Optional, Sequence, Union
 
+import numpy as np
 import pandas as pd
 from datasets.arrow_dataset import Dataset
 from datasets.dataset_dict import DatasetDict
@@ -14,7 +18,7 @@
 logger = LOGGING.getLogger(__name__)
 
 
-class DSFilter:
+class DSSlice:
     @staticmethod
     def sample_dataset(
         data: DatasetLikeType,
@@ -71,15 +75,15 @@ def filter_and_sample_data(
         """
         # Filter by queries
         if queries:
-            data_ = DSFilter.filter_data_by_queries(data, queries, verbose=verbose)
+            data_ = DSSlice.filter_data_by_queries(data, queries, verbose=verbose)
         else:
             logger.warning("No query specified")
             data_ = data.copy()
 
         # Create a sample for analysis
         sample = None
         if sample_size:
-            sample = DSFilter.sample_data(
+            sample = DSSlice.sample_data(
                 data_,
                 sample_size_per_group=sample_size,
                 sample_seed=sample_seed,
@@ -237,3 +241,24 @@ def sample_data(
                 print(grp_dists)
 
         return _sample
+
+    @staticmethod
+    def split_dataframe(
+        data,
+        indices_or_sections: Union[int, Sequence[int]],
+        verbose: bool = False,
+    ) -> List[pd.DataFrame]:
+        """
+        Split a dataframe into multiple dataframes
+
+        Args:
+            data (pd.DataFrame): dataframe to split
+            indices_or_sections (int or sequence of ints): if int, number of chunks to split the dataframe into
+
+        Returns:
+            List[pd.DataFrame]: list of dataframes
+        """
+
+        if verbose:
+            logger.info("Splitting dataframe into %s", indices_or_sections)
+        return np.array_split(data, indices_or_sections)