def test_stratified_fold_split():
df = _setup_data()
splitted = dataset.stratified_fold_split(dataframe=df,
class_column="class")
assert int == splitted["fold"].dtype
assert set(range(5)) == set(splitted["fold"].unique())
ants_folds = set(splitted[splitted["tag"] == "ants"]["fold"])
bees_folds = set(splitted[splitted["tag"] == "bees"]["fold"])
assert ants_folds == bees_folds
def test_stratified_fold_split_num_folds():
df = _setup_data()
splitted = dataset.stratified_fold_split(df, "class", n_folds=2)
assert set(range(2)) == set(splitted["fold"].unique())
def split_dataframe(
dataframe: pd.DataFrame,
train_folds: List[int],
valid_folds: Optional[List[int]] = None,
infer_folds: Optional[List[int]] = None,
tag2class: Optional[Dict[str, int]] = None,
tag_column: str = None,
class_column: str = None,
seed: int = 42,
n_folds: int = 5
) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
"""
Split a Pandas DataFrame into folds.
Args:
dataframe (pd.DataFrame): input dataframe
train_folds (List[int]): train folds
valid_folds (List[int], optional): valid folds.
If none takes all folds not included in ``train_folds``
infer_folds (List[int], optional): infer folds.
If none takes all folds not included in ``train_folds``
and ``valid_folds``
tag2class (Dict[str, int], optional): mapping from label names into int
tag_column (str, optional): column with label names
class_column (str, optional): column to use for split
seed (int): seed for split
n_folds (int): number of folds
Returns:
(tuple): tuple with 4 dataframes
whole dataframe, train part, valid part and infer part
"""
if args_are_not_none(tag2class, tag_column, class_column):
dataframe = map_dataframe(dataframe, tag_column, class_column,
tag2class)
if class_column is not None:
result_dataframe = stratified_fold_split(dataframe,
class_column=class_column,
random_state=seed,
n_folds=n_folds)
else:
result_dataframe = default_fold_split(dataframe,
random_state=seed,
n_folds=n_folds)
fold_series = result_dataframe["fold"]
train_folds = folds_to_list(train_folds)
df_train = result_dataframe[fold_series.isin(train_folds)]
if valid_folds is None:
mask = ~fold_series.isin(train_folds)
valid_folds = result_dataframe[mask]["fold"]
valid_folds = folds_to_list(valid_folds)
df_valid = result_dataframe[fold_series.isin(valid_folds)]
infer_folds = folds_to_list(infer_folds or [])
df_infer = result_dataframe[fold_series.isin(infer_folds)]
return result_dataframe, df_train, df_valid, df_infer