Merge pull request #328 from PNNL-CompBio/327-fix-balancing-of-datasets-in-train_test_validate

added balancing flag to all splitting methods

Author: ymahlich

coderdata/dataset/dataset.py

@@ -343,6 +343,7 @@ def split_train_other(
             ]='mixed-set',
         ratio: tuple[int, int, int]=(8,2),
         stratify_by: Optional[str]=None,
+        balance: bool=False,
         random_state: Optional[Union[int,RandomState]]=None,
         **kwargs: dict, 
         ) -> TwoWaySplit:
@@ -352,6 +353,7 @@ def split_train_other(
             split_type=split_type,
             ration=ratio,
             stratify_by=stratify_by,
+            balance=balance,
             random_state=random_state,
             **kwargs
         )
@@ -366,15 +368,16 @@ def split_train_test_validate(
             ]='mixed-set',
         ratio: tuple[int, int, int]=(8,1,1),
         stratify_by: Optional[str]=None,
+        balance: bool=False,
         random_state: Optional[Union[int,RandomState]]=None,
         **kwargs: dict,
         ) -> Split:
-
         split = split_train_test_validate(
             data=self,
             split_type=split_type,
             ratio=ratio,
             stratify_by=stratify_by,
+            balance=balance,
             random_state=random_state,
             **kwargs
         )
@@ -389,6 +392,7 @@ def train_test_validate(
             ]='mixed-set',
         ratio: tuple[int, int, int]=(8,1,1),
         stratify_by: Optional[str]=None,
+        balance: bool=False,
         random_state: Optional[Union[int,RandomState]]=None,
         **kwargs: dict,
         ) -> Split:
@@ -398,6 +402,7 @@ def train_test_validate(
             split_type=split_type,
             ratio=ratio,
             stratify_by=stratify_by,
+            balance=balance,
             random_state=random_state,
             **kwargs
         )
@@ -725,6 +730,7 @@ def split_train_other(
             ]='mixed-set',
         ratio: tuple[int, int, int]=(8,2),
         stratify_by: Optional[str]=None,
+        balance: bool=False,
         random_state: Optional[Union[int,RandomState]]=None,
         **kwargs: dict, 
     ):
@@ -733,8 +739,9 @@ def split_train_other(
         split_type,
         ratio,
         stratify_by,
+        balance,
         random_state,
-        kwargs=kwargs
+        **kwargs
     )
     if stratify_by is not None:
         train.experiments = train.experiments[train.experiments['dose_response_metric'] != 'split_class']
@@ -749,10 +756,10 @@ def split_train_test_validate(
             ]='mixed-set',
         ratio: tuple[int, int, int]=(8,1,1),
         stratify_by: Optional[str]=None,
+        balance: bool=False,
         random_state: Optional[Union[int,RandomState]]=None,
         **kwargs: dict,
         ) -> Split:
-
     # Type checking split_type
     if split_type not in [
         'mixed-set', 'drug-blind', 'cancer-blind'
@@ -766,17 +773,19 @@ def split_train_test_validate(
         split_type=split_type,
         ratio=[ratio[0], ratio[1] + ratio[2]],
         stratify_by=stratify_by,
+        balance=balance,
         random_state=random_state,
-        kwargs=kwargs,
+        **kwargs,
         )
 
     test, val = _split_two_way(
         data=other,
         split_type=split_type,
         ratio=[ratio[1], ratio[2]],
         stratify_by=stratify_by,
+        balance=balance,
         random_state=random_state,
-        kwargs=kwargs,
+        **kwargs,
     )
 
     if stratify_by is not None:
@@ -794,6 +803,7 @@ def train_test_validate(
             ]='mixed-set',
         ratio: tuple[int, int, int]=(8,1,1),
         stratify_by: Optional[str]=None,
+        balance: bool=False,
         random_state: Optional[Union[int,RandomState]]=None,
         **kwargs: dict,
         ) -> Split:
@@ -871,8 +881,9 @@ def train_test_validate(
         split_type=split_type,
         ratio=ratio,
         stratify_by=stratify_by,
+        balance=balance,
         random_state=random_state,
-        kwargs=kwargs,
+        **kwargs,
         )
 
 
@@ -976,6 +987,20 @@ def _filter(data: Dataset, split: pd.DataFrame) -> Dataset:
 
     return data_ret
 
+def _balance_data(
+        data: pd.Dataframe,
+        random_state: Optional[Union[int,RandomState]]=None,
+        # oversample: bool=False,
+        ) -> pd.Dataframe:
+    tmp = deepcopy(data)
+    counts = tmp.value_counts('split_class')
+    ret_df = (
+        tmp
+        .groupby('split_class')
+        .sample(n=min(counts), random_state=random_state)
+        )
+    return ret_df
+
 
 def _create_classes(
         data: pd.DataFrame,
@@ -1072,6 +1097,7 @@ def _split_two_way(
             'mixed-set', 'drug-blind', 'cancer-blind'
             ]='mixed-set',
         ratio: tuple[int, int, int]=(8,2),
+        balance: bool=False,
         stratify_by: Optional[str]=None,
         random_state: Optional[Union[int,RandomState]]=None,
         **kwargs: dict, 
@@ -1150,7 +1176,6 @@ def _split_two_way(
     thresh = kwargs.get('thresh', None)
     num_classes = kwargs.get('num_classes', 2)
     quantiles = kwargs.get('quantiles', True)
-
     # Type checking split_type
     if split_type not in [
         'mixed-set', 'drug-blind', 'cancer-blind'
@@ -1290,6 +1315,11 @@ def _split_two_way(
                 thresh=thresh,
                 quantiles=quantiles,
                 )
+            if balance:
+                df_full = _balance_data(
+                    data=df_full,
+                    random_state=random_state
+                )
         if split_type == 'mixed-set':
             # Using StratifiedShuffleSplit to generate randomized train 
             # and 'other' set, since there is no need for grouping.