training/splitting.py

"""Module to perform splitting of data into train/test or K-folds."""

import pandas as pd
import numpy as np
from sklearn.model_selection import StratifiedGroupKFold


def subject_wise_train_test_split(
    data,
    target_col,
    id_col,
    test_size,
    stratify_by_sex=False,
    sex_col=None,
    stratify_by_age=False,
    age_bin_col=None,
    shuffle=False,
    random_state=None,
):
    """Subject wise splitting data into train and test sets.

    Splits data into train and test sets ensuring that the same patient can only appear in
    either train or test set. Stratifies data according to class label, with additional
    options to stratify by sex and age.

    Parameters
    ----------
    data : pd.DataFrame
        dataframe containing features and target.
    target_col : str
        name of target column.
    id_col : str
        name of patient id column.
    test_size : float
        represents the proportion of the dataset to include in the test split. Float should
        be between 0 and 1.
    stratify_by_sex : bool, optional
        option to stratify data by sex, by default False.
    sex_col : str, optional
        name of sex column, by default None.
    stratify_by_age : bool, optional
        option to stratify data by age, by default False.
    age_bin_col : str, optional
        name of age column, by default None. Age column must be provided in binned format.
    shuffle : bool, optional
        whether to shuffle each class's samples before splitting into batches, by default
        False.
    random_state : int, optional
        when shuffle is True, random_state affects the ordering of the indices, by default
        None.

    Returns
    -------
    train_data : pd.DataFrame
        train data stratified by class. Also stratified by age/sex as specified in input
        parameters.
    test_data : pd.DataFrame
        test data stratified by class. Also stratified by age/sex as specified in input
        parameters.

    Raises
    -------
    ValueError : error raised when boolean stratify_by_age or stratify_by_sex is True but
        the respective columns are not provided for stratifying.

    """
    # Raise error if stratify_by_sex/stratify_by_age is True but the respective columns are
    # not provided
    if (stratify_by_age is True) & (age_bin_col is None):
        raise ValueError(
            "Parameter stratify_by_age is True but age_bin_col not provided."
        )
    if (stratify_by_sex is True) & (sex_col is None):
        raise ValueError("Parameter stratify_by_sex is True but sex_col not provided.")

    # Adapt target column to contain all variables to split by to allow stratified splitting
    # by StratifiedGroupKFold.
    if (stratify_by_sex is True) and (stratify_by_age is True):
        data["TempTarget"] = (
            data[target_col].astype(str) + data[sex_col].astype(str) + data[age_bin_col]
        )
    elif (stratify_by_sex is True) and (stratify_by_age is False):
        data["TempTarget"] = data[target_col].astype(str) + data[sex_col].astype(str)
    elif (stratify_by_sex is False) and (stratify_by_age is True):
        data["TempTarget"] = data[target_col].astype(str) + data[sex_col].astype(str)
    else:
        data["TempTarget"] = data[target_col]
    temp_target_col = "TempTarget"

    # Calculate the number of folds to split data to using the size of the test data.
    num_folds = round(1 / test_size)
    sgkf = StratifiedGroupKFold(
        n_splits=num_folds, shuffle=shuffle, random_state=random_state
    )

    # Retrieve the first split and create train and test dfs
    train_test_splits = next(
        sgkf.split(data, data[temp_target_col], groups=data[id_col])
    )
    train_ids = train_test_splits[0].tolist()
    test_ids = train_test_splits[1].tolist()
    train_data = data.iloc[train_ids]
    test_data = data.iloc[test_ids]

    # Drop temporary target column
    train_data = train_data.drop(columns=temp_target_col)
    test_data = test_data.drop(columns=temp_target_col)

    return train_data, test_data


def subject_wise_kfold_split(
    train_data,
    target_col,
    id_col,
    num_folds,
    sex_col=None,
    age_col=None,
    stratify_by_sex=False,
    stratify_by_age=False,
    age_bin_col=None,
    shuffle=False,
    random_state=None,
    print_log=False,
):
    """Subject wise splitting data into balanced K-folds.

    Splits data into K-folds ensuring that the same patient can only appear in
    either train or validation set. Stratifies data according to class label, with additional
    options to stratify by sex and age.

    Parameters
    ----------
    train_data : pd.DataFrame
        dataframe containing features and target.
    target_col : str
        name of target column.
    id_col : str
        name of patient id column.
    num_folds : int
        number of folds.
    sex_col : str, optional
        name of sex column, by default None. Required if stratify_by_sex is True. Can be
        included when stratify_by_sex is False to get info on sex ratio across folds.
    age_col : str, optional
        name of age column, by default None. Column must be a continous variable. Can be
        included to get info on mean age across folds.
    stratify_by_sex : bool, optional
        option to stratify data by sex, by default False.
    stratify_by_age : bool, optional
        option to stratify data by age, by default False. The binned age (age_bin_col) is
        used for stratifying by age rather than age_col.
    age_bin_col : str, optional
        name of age column, by default None. Age column must be provided in binned format.
    shuffle : bool, optional
        whether to shuffle each class's samples before splitting into batches, by default
        False.
    random_state : int, optional
        when shuffle is True, random_state affects the ordering of the indices, by default
        None.
    print_log : bool, optional
        flag to print distributions across folds, by default False.

    Returns
    -------
    validation_fold_ids : list of arrays
        each array contains the validation patient IDs for each fold.

    Raises
    -------
    ValueError : error raised when boolean stratify_by_age or stratify_by_sex is True but
        the respective columns are not provided for stratifying.

    """
    # Raise error if stratify_by_sex/stratify_by_age is True but the respective columns are
    # not provided
    if (stratify_by_age is True) & (age_bin_col is None):
        raise ValueError(
            "Parameter stratify_by_age is True but age_bin_col not provided."
        )
    if (stratify_by_sex is True) & (sex_col is None):
        raise ValueError("Parameter stratify_by_sex is True but sex_col not provided.")
    
    # Adapt target column to contain all variables to split by to allow stratified splitting
    # by StratifiedGroupKFold.
    if (stratify_by_sex is True) and (stratify_by_age is True):
        train_data["TempTarget"] = (
            train_data[target_col].astype(str)
            + train_data[sex_col].astype(str)
            + train_data[age_bin_col]
        )
    elif (stratify_by_sex is True) and (stratify_by_age is False):
        train_data["TempTarget"] = train_data[target_col].astype(str) + train_data[
            sex_col
        ].astype(str)
    elif (stratify_by_sex is False) and (stratify_by_age is True):
        train_data["TempTarget"] = train_data[target_col].astype(str) + train_data[
            sex_col
        ].astype(str)
    else:
        train_data["TempTarget"] = train_data[target_col]
    temp_target_col = "TempTarget"

    sgkf_train = StratifiedGroupKFold(
        n_splits=num_folds, shuffle=shuffle, random_state=random_state
    )

    validation_fold_ids = []
    class_fold_ratios = []
    sex_fold_ratios = []
    fold_mean_ages = []
    for i, (train_index, validation_index) in enumerate(
        sgkf_train.split(
            train_data, train_data[temp_target_col], groups=train_data[id_col]
        )
    ):
        # Get patient ID's for each validation fold
        validation_ids = train_data[id_col].iloc[validation_index].unique()
        validation_fold_ids.append(validation_ids)

        # Get class, sex and age distributions for the training data for each fold
        train_fold_data = train_data[~train_data[id_col].isin(validation_ids)]
        class_ratio = train_fold_data[target_col].value_counts()[1] / len(
            train_fold_data
        )
        class_fold_ratios.append(class_ratio)
        if not sex_col is None:
            sex_ratio = train_fold_data[sex_col].value_counts()[1] / len(
                train_fold_data
            )
            sex_fold_ratios.append(sex_ratio)
        if not age_col is None:
            mean_age = train_fold_data[age_col].mean()
            fold_mean_ages.append(mean_age)

    if print_log is True:
        print("Fold proportions:")
        print("Train class ratio:", class_fold_ratios)
        if not sex_col is None:
            print("Sex class ratio:", sex_fold_ratios)
        if not age_col is None:
            print("Mean age:", fold_mean_ages)

    # Allows inhomogenous array to be saved with np.save
    validation_fold_ids = np.asarray(validation_fold_ids, dtype="object")

    # Delete temporary target column
    del train_data[temp_target_col]

    return validation_fold_ids


def get_cv_fold_indices(validation_fold_ids, train_data, id_col):
    """
    Find train/val dataframe indices for each fold and format for cross validation.

    Creates a tuple with training and validation indices for each K-fold using the
    validation_fold_ids. These patients are assigned to the validation portion of the data
    and all other patients are assigned to the train portion for that fold.
    Checks that all patient IDs listed for the K folds are contained in the train data.
    For each fold, extracts the dataframe indices for patient data belonging that fold
    and assigns all other indices to the 'train' portion. The list returned contains tuples
    required to be passed to sklearn's cross_validate function (through the cv argument).

    Parameters
    ----------
    fold_patients : array
        lists of patient IDs for each of the K folds.
    train_data : pd.DataFrame
        train data (must contain id_col).
    id_col : str
        name of column containing patient ID.

    Returns
    -------
    cross_validation_fold_indices : list of tuples
        K lists of val/train dataframe indices.

    """
    # Create a tuple with training and validation indices for each fold.
    cross_val_fold_indices = []
    for fold in validation_fold_ids:
        fold_val_ids = train_data[train_data[id_col].isin(fold)]
        fold_train_ids = train_data[~train_data[id_col].isin(fold)]

        # Get index of rows in val and train
        fold_val_index = fold_val_ids.index
        fold_train_index = fold_train_ids.index

        # Append tuple of training and val indices
        cross_val_fold_indices.append((fold_train_index, fold_val_index))
    return cross_val_fold_indices