File size: 11,097 Bytes
000de75 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 | """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
|