import pandas as pd
import numpy as np
def get_splits(age_at_exam, patient_ids, exam_ids, splits, min_age_valid=16, max_age_valid=85, seed=0):
rng = np.random.RandomState(seed)
if sum(splits) > 1.0:
raise ValueError('splits should be sum to a number smaller than one.')
n_exams = len(exam_ids)
# Get patients
patients = np.unique(patient_ids)
n_patients = len(patients)
# Create correspondence
hash_exams = dict(zip(exam_ids, range(n_exams)))
hash_patients = dict(zip(patients, range(n_patients)))
inverse_hash_patients = dict(zip(range(n_patients), patients))
# Get all exams for each patient
patient_exams = [[] for _ in range(n_patients)]
for exam_idx in range(n_exams):
exam_id = exam_ids[exam_idx]
patient_id = patient_ids[exam_idx]
patient_idx = hash_patients[patient_id]
patient_exams[patient_idx].append(exam_id)
# Get the age at one of the exams for each patient
ages, _, _ = np.unique(age_at_exam, return_inverse=True, return_counts=True)
patient_idx_per_age = {a: [] for a in ages}
patient_single_exam = np.zeros(n_patients, dtype=int)
for patient_idx in range(n_patients):
# Pick random exam id for the given patient
# OBS:Another formulation that could make sense could be to always pick the first exam....
id_exam = rng.choice(patient_exams[patient_idx])
patient_single_exam[patient_idx] = id_exam
exam_idx = hash_exams[id_exam]
a = age_at_exam[exam_idx]
patient_idx_per_age[a].append(patient_idx)
# Get number of patient in each split
n_splits = [int(np.floor(s * n_patients)) for s in splits]
n_splits += [n_patients - sum(n_splits)]
# Shuffle
rng.shuffle(ages) # Shuffle ages
for a, patient_idx in patient_idx_per_age.items(): # Shuffle within the same age
rng.shuffle(patient_idx)
# Pick one id per age and build a list from that
all_patient_idx = []
stop = False
# Pick ids within the given range first (which will probably be used for training, validation and test)
while not stop:
stop = True
for a in ages[(ages >= min_age_valid) & (ages <= max_age_valid)]:
if patient_idx_per_age[a]:
patient_idx = patient_idx_per_age[a].pop()
all_patient_idx.append(patient_idx)
stop = False
# Pick remaining ids last (which will probably be used only for training)
stop = False
while not stop:
stop = True
for a in ages[(ages < min_age_valid) | (ages > max_age_valid)]:
if patient_idx_per_age[a]:
patient_idx = patient_idx_per_age[a].pop()
all_patient_idx.append(patient_idx)
stop = False
# Save ids
patients_in_splits = [[] for n in n_splits]
single_exam_in_split = [[] for n in n_splits]
exams_in_splits = [[] for n in n_splits]
for i, patient_idx in enumerate(all_patient_idx):
last_n = 0
for s, n in enumerate(np.cumsum(n_splits)):
if last_n <= i < n:
patients_in_splits[s].append(inverse_hash_patients[patient_idx])
single_exam_in_split[s].append(patient_single_exam[patient_idx])
exams_in_splits[s] += patient_exams[patient_idx]
last_n = n
return patients_in_splits, single_exam_in_split, exams_in_splits
if __name__ == "__main__":
import argparse
import warnings
parser = argparse.ArgumentParser(description='Generate data summary for the age prediction problem')
parser.add_argument('file',
help='csv file to read data from.')
parser.add_argument('--exam_id_col', default='N_exame',
help='column in csv containing exam id')
parser.add_argument('--age_col', default='Idade',
help='column in csv containing age')
parser.add_argument('--patient_id_col', default='N_paciente_univoco',
help='column in csv containing patient id')
parser.add_argument('--splits', default=[0.15, 0.05], nargs='*', type=float,
help='percentage of data in each split')
parser.add_argument('--splits_names', default=['test', 'val', 'train'], nargs='*', type=str,
help='split names')
parser.add_argument('--no_plot', action='store_true',
help='dont show plots')
args, unk = parser.parse_known_args()
if unk:
warnings.warn("Unknown arguments:" + str(unk) + ".")
# Open csv file
df = pd.read_csv(args.file, low_memory=False)
# Remove duplicated rows
df.drop_duplicates(args.exam_id_col, inplace=True)
# Get ids from csv file
exam_ids = np.array(df[args.exam_id_col], dtype=int)
age_at_exam = np.array(df[args.age_col])
patient_ids = np.array(df[args.patient_id_col], dtype=int)
# define splits
patients_in_splits, single_exam_in_split, exams_in_splits = get_splits(age_at_exam, patient_ids, exam_ids, args.splits)
if not args.no_plot:
import seaborn as sns
import matplotlib.pyplot as plt
n = len(args.splits) + 1
fig, ax = plt.subplots(nrows=n)
for i in range(n):
age_single_exam = age_at_exam[np.isin(exam_ids, single_exam_in_split[i])]
age = age_at_exam[np.isin(exam_ids, exams_in_splits[i])]
sns.histplot(age, ax=ax[i], kde=False, bins=range(0, 130, 1))
sns.histplot(age_single_exam, ax=ax[i], kde=False, bins=range(0, 130, 1))
plt.show()