File size: 5,214 Bytes
e69d4e4 | 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 | """Splits the model C cohort and patient days into stratified train and test sets.
The train set retains these characteristics of the full data set:
Exac days to non-exac days ratio (within 5%). Individual patients can only appear in
either train or test.
Sex ratio (within 0.05)
Age distribution (minimum p-value for Kolmogorov-Smirnov test=0.9)
This script also splits the train data into balanced folds for cross-validation. Patient
IDs for train, test and all data folds are stored for use in subsequent scripts.
All data sets are divided into train and test and stored in separate folders.
"""
import numpy as np
import os
import pandas as pd
import pickle
from lenusml import splits
data_dir = '<YOUR_DATA_PATH>/copd-dataset/'
output_train_data_dir = '<YOUR_DATA_PATH>/train_data'
output_test_data_dir = '<YOUR_DATA_PATH>/test_data'
cohort_info_dir = '../data/cohort_info/'
save_cohort_info = True
data = pd.read_pickle(os.path.join(data_dir, 'exac_data.pkl'))
##########################################
# Prepare demographic info for splitting
##########################################
# Calculate decimal age on DateOfEvent
data['DateOfBirth'] = pd.to_datetime(data['DateOfBirth'], utc=True)
def calculate_age_decimal(dob, date):
age = date - dob
decimal_age = (age.days + age.seconds / 86400.0) / 365.2425
return decimal_age
data['Age'] = data.apply(lambda x: calculate_age_decimal(
x['DateOfBirth'], x['DateOfEvent']), axis=1)
data = data.drop(columns=['DateOfBirth'])
##########################################
# Merge with COPD status and inhaler data
##########################################
patient_details = pd.read_csv(os.path.join(data_dir, 'CopdDatasetPatientDetails.txt'),
usecols=['StudyId', 'CopdStatusDetailsId'],
delimiter="|")
copd_status = pd.read_csv(os.path.join(data_dir, 'CopdDatasetCopdStatusDetails.txt'),
usecols=['Id', 'SmokingStatus', 'RequiredAcuteNIV',
'RequiredICUAdmission',
'LungFunction_FEV1PercentPredicted',
'LabsHighestEosinophilCount'],
delimiter="|")
# Strip out % signs from spirometry and convert to float
copd_status['LungFunction_FEV1PercentPredicted'] = copd_status[
'LungFunction_FEV1PercentPredicted'].str.strip('%').astype('float')
patient_details = patient_details.merge(
copd_status, left_on='CopdStatusDetailsId', right_on='Id',
how='left').drop(columns=['CopdStatusDetailsId', 'Id'])
data = data.merge(patient_details, on='StudyId', how='left')
#################################
# Define train and test cohorts
#################################
print('Split data into train and test')
# Set the class ratio tolerance to 5% of the data class ratio
class_ratio_tolerance = 0.05 * data.IsExac.value_counts(normalize=True)[0] /\
data.IsExac.value_counts(normalize=True)[1]
print("Class ratio tolerance: ", class_ratio_tolerance)
# Set the sex ratio tolerance to 5% of the data class ratio
sex_ratio_tolerance = 0.05 * data.Sex.value_counts(normalize=True)['M'] /\
data.Sex.value_counts(normalize=True)['F']
print("Sex ratio tolerance: ", sex_ratio_tolerance)
train_data, test_data, train_ids, test_ids = splits.train_test_stratified_class_sex(
data=data, id_column='StudyId', class_column='IsExac', sex_column='Sex',
train_proportion=0.85,
proportion_tolerance=0.05, class_ratio_tolerance=class_ratio_tolerance,
sex_ratio_tolerance=sex_ratio_tolerance, random_seed=42)
#################################
# Create cross validation folds
#################################
fold_proportions, fold_class_ratios, fold_patients = splits.group_kfold_class_balanced(
data=train_data, id_column='StudyId', class_column='IsExac', K=5,
fold_proportion_tolerance=0.05,
fold_class_ratio_tolerance=class_ratio_tolerance, random_seed=42)
if save_cohort_info:
os.makedirs(cohort_info_dir, exist_ok=True)
with open(os.path.join(cohort_info_dir, "test_ids.pkl"), 'wb') as f:
pickle.dump(list(test_ids), f)
with open(os.path.join(cohort_info_dir, "train_ids.pkl"), 'wb') as f:
pickle.dump(list(train_ids), f)
print('Train and test patient IDs saved')
with open(os.path.join(cohort_info_dir, "fold_proportions.pkl"), 'wb') as f:
pickle.dump(list(fold_proportions), f)
with open(os.path.join(cohort_info_dir, "fold_class_ratios.pkl"), 'wb') as f:
pickle.dump(list(fold_class_ratios), f)
np.save(os.path.join(cohort_info_dir, 'fold_patients.npy'), fold_patients,
allow_pickle=True)
print('Cross validation fold information saved')
###############################
# Save train and test sets
###############################
# Create the output directories
os.makedirs(output_train_data_dir, exist_ok=True)
os.makedirs(output_test_data_dir, exist_ok=True)
# Save exac and patient details info
train_data.to_pickle(os.path.join(output_train_data_dir, 'train_data.pkl'))
test_data.to_pickle(os.path.join(output_test_data_dir, 'test_data.pkl'))
print('Patient details/exac data saved')
|