import os, sys
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import connected_components
from sklearn.preprocessing import OneHotEncoder
import pandas as pd
import numpy as np
import pickle
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.sampler import WeightedRandomSampler
import torch
from .params import main_data_path, cohort_data_path, kinship_path, withdraw_path, fam_path
from .utils import get_fields, get_row_last_values, remove_kinships, save_dict, load_dict, print_sys
class ukbb_cohort:
def __init__(self, main_data_path, cohort_data_path, withdraw_path, keep_relatives = False):
self.keep_relatives = keep_relatives
self.cohort_data_path = cohort_data_path
self.main_data_path = main_data_path
if keep_relatives:
cohort_path = os.path.join(cohort_data_path, 'cohort_with_relatives.pkl')
else:
cohort_path = os.path.join(cohort_data_path, 'cohort_no_relatives.pkl')
if not os.path.exists(cohort_path):
print_sys('construct from scratch...')
'''
exclusions:
the original uk biobank nature paper supplementary S 3.4:
22006: Genetic ethnic grouping -> to retain only white british ancestry
https://www.frontiersin.org/articles/10.3389/fgene.2022.866042/full
22018: genetic relatedness exclusions
22019: sex chromosome aneuploidy
31 <-> 22001: mismatch between self-reported sex and genetically determined sex
22010: recommended genomic analysis exclusions, signs of insufficient data quality
(optional) further remove relatives based on KING relative scores, choose the first one in relative group
remove the list of eids who do not want to be in the study anymore
'''
all_field_ids = [22006, 22018, 22019, 22001, 22010, 31]
df_main = get_fields(all_field_ids, main_data_path)
cur_size = len(df_main)
print_sys('Total sample size: ' + str(cur_size))
df_main = df_main[df_main['22006-0.0'] == 1]
print_sys('Keeping only white british ancestry (ID: 22006), cutting from ' + str(cur_size) + ' to ' + str(len(df_main)))
cur_size = len(df_main)
df_main = df_main[df_main['22018-0.0'].isnull()]
print_sys('Removing genetics related samples (ID: 22018), cutting from ' + str(cur_size) + ' to ' + str(len(df_main)))
cur_size = len(df_main)
df_main = df_main[df_main['22019-0.0'].isnull()]
print_sys('Removing sex chromosome aneuploidy (ID: 22019), cutting from ' + str(cur_size) + ' to ' + str(len(df_main)))
cur_size = len(df_main)
df_main = df_main[df_main['31-0.0'] == df_main['22001-0.0']]
print_sys('Removing samples with mismatched self-reported sex and genetic determined sex (ID: 31 <-> 22001), cutting from ' + str(cur_size) + ' to ' + str(len(df_main)))
cur_size = len(df_main)
df_main = df_main[df_main['22010-0.0'].isnull()]
print_sys('Removing samples with genomic data quality (ID: 22010), cutting from ' + str(cur_size) + ' to ' + str(len(df_main)))
cur_size = len(df_main)
save_dict(os.path.join(cohort_data_path, 'cohort_with_relatives.pkl'), df_main.eid.values)
kinship_mask = remove_kinships(df_main.eid)
df_main = df_main[kinship_mask]
save_dict(os.path.join(cohort_data_path, 'cohort_no_relatives.pkl'), df_main.eid.values)
else:
print_sys('Found local copy...')
self.cohort = load_dict(cohort_path)
print_sys('There are ' + str(len(self.cohort)) + ' samples!')
if keep_relatives:
self.no_rel_eid = load_dict(os.path.join(cohort_data_path, 'cohort_no_relatives.pkl'))
if os.path.exists(withdraw_path):
## todo: when there is a withdraw file, implement this...
pass
def get_covariates(self, to_plink = False, plink_num_pca = 15, return_full = False, plink_filter = False):
'''
covariates:
31: sex
21003: Age when attended assessment centre
22009: pca
54: assessment center
batch from params.fam_path file
'''
covar_path = os.path.join(self.cohort_data_path, 'covariates_all.pkl')
if os.path.exists(covar_path):
print_sys('Found local copy...')
self.covar = load_dict(covar_path)
else:
print_sys('construct co-variates from scratch...')
df_covar = get_fields([31, 54, 21003, 22009], self.main_data_path)
column_name_map = {'22009-0.' + str(i): 'pca ' + str(i) for i in range(1, 41)}
column_name_map['31-0.0'] = 'sex'
column_name_map['21003-0.0'] = 'age'
column_name_map['54-0.0'] = 'assessment_center'
self.covar = df_covar.rename(columns = column_name_map)
enc = OneHotEncoder(handle_unknown='ignore')
enc.fit(self.covar['assessment_center'].unique().reshape(-1,1))
center_array = enc.transform(self.covar['assessment_center'].values.reshape(-1,1)).toarray()
center_one_hot = pd.DataFrame(center_array).astype('int').rename(columns = dict(zip(range(22), ['center_' + str(i) for i in range(22)])))
self.covar = self.covar.drop(['21003-1.0', '21003-2.0', '21003-3.0', 'assessment_center', '54-1.0', '54-2.0', '54-3.0'], axis = 1)
self.covar = self.covar.join(center_one_hot)
df_fam = pd.read_csv(fam_path)
enc = OneHotEncoder(handle_unknown='ignore')
enc.fit(df_fam.trait.unique().reshape(-1,1))
batch_one_hot = enc.transform(df_fam['trait'].values.reshape(-1,1)).toarray()
batch_num = batch_one_hot.shape[1]
id2batch = dict(zip(df_fam.fid.values, batch_one_hot.astype(int)))
df_batch = pd.DataFrame(np.stack(self.covar['eid'].apply(lambda x: id2batch[x] if x in id2batch else np.zeros(batch_one_hot.shape[1]).astype(int)).values)).rename(columns = dict(zip(range(batch_num), ['batch_' + str(i) for i in range(batch_num)])))
self.covar = self.covar.join(df_batch)
save_dict(covar_path, self.covar)
print_sys('Done! Saving...')
if not to_plink:
if return_full:
return self.covar.reset_index(drop = True)
else:
return self.covar[self.covar.eid.isin(self.cohort)].reset_index(drop = True)
else:
plink_path = os.path.join(self.cohort_data_path, 'covar_pca' + str(plink_num_pca) + '_all_real_value')
if plink_filter:
plink_path += '_null_removed'
plink_path += '.txt'
if not os.path.exists(plink_path):
pca_columns = [i for i in self.covar.columns.values if (i[:3]=='pca') and int(i.split()[-1]) <= plink_num_pca]
#center_one_hot_columns = ['center_' + str(i) for i in range(22)]
#batch_columns = ['batch_' + str(i) for i in range(batch_num)]
#self.covar[['eid', 'eid', 'age', 'sex'] + pca_columns + center_one_hot_columns + batch_columns].to_csv(plink_path, header=None, index=None, sep=' ')
center = np.argmax(self.covar.loc[:, self.covar.columns.str.contains('center')].values, axis = 1)
batch = np.argmax(self.covar.loc[:, self.covar.columns.str.contains('batch')].values, axis = 1)
self.covar = self.covar.iloc[:, :43]
self.covar['assessment_center'] = center
self.covar['batch'] = batch
if plink_filter:
self.covar = self.covar[self.covar.eid.isin(self.cohort)].reset_index(drop = True)
self.covar[['eid', 'eid', 'age', 'sex', 'assessment_center', 'batch'] + pca_columns].to_csv(plink_path, header=None, index=None, sep=' ')
self.covar_plink = pd.read_csv(plink_path, header = None, sep = ' ')
return self.covar_plink
def get_external_traits(self, trait_name, to_plink = False, to_str = True, random_seed = 42, sep_cohort = False, randomize = False, use_sample_size = False, sample_size = -1, randomize_seed = 42):
'''
example:
standing heights: 50
'''
if trait_name in ['body_BALDING1', 'cancer_BREAST', 'disease_ALLERGY_ECZEMA_DIAGNOSED', 'disease_HYPOTHYROIDISM_SELF_REP', 'other_MORNINGPERSON', 'pigment_SUNBURN']:
trait_type = 'binary'
else:
trait_type = 'continuous'
pheno_path = os.path.join(self.cohort_data_path, str(trait_name) + '_pheno.pkl')
if os.path.exists(pheno_path):
print_sys('Found local copy...')
self.pheno = load_dict(pheno_path)
else:
print_sys('construct phenotype from scratch...')
self.pheno = pd.read_csv(os.path.join(data_path, 'full_gwas', trait_name+'_'+trait_type+'.csv'))
self.pheno['eid'] = self.pheno.eid.astype('int')
self.pheno = self.pheno[self.pheno['pheno'].notnull()]
if trait_type == 'binary':
self.pheno['pheno'] += 1
self.pheno['pheno'] = self.pheno['pheno'].astype(int)
save_dict(pheno_path, self.pheno)
print_sys('Done! Saving...')
# filtering to cohorts incl. with/without relatives
self.pheno = self.pheno[self.pheno.eid.isin(self.cohort)].reset_index(drop = True)
if to_str:
self.pheno['eid'] = self.pheno['eid'].astype('str')
if not to_plink:
return self.pheno
else:
plink_path = os.path.join(self.cohort_data_path, str(trait_name) + '_plink')
if self.keep_relatives:
plink_path = plink_path + '_with_relatives'
else:
plink_path = plink_path + '_no_relatives'
if use_sample_size:
plink_path = plink_path + '_' + str(sample_size) + '_' + str(random_seed)
if sep_cohort:
plink_path += '_sep_cohort'
if randomize:
plink_path += '_randomize' + str(randomize_seed)
plink_path = plink_path + '.txt'
if randomize:
self.pheno['pheno'] = self.pheno['pheno'].sample(frac = 1, random_state = randomize_seed).values
if use_sample_size:
from sklearn.model_selection import train_test_split
print('random_seed:', random_seed)
pheno_shuffle = self.pheno.sample(frac = 1, random_state = random_seed)
all_ids, y = pheno_shuffle.eid.values, pheno_shuffle['pheno'].values
train_val_ids, test_ids, y_train_val, y_test = all_ids[:sample_size], all_ids[sample_size:], y[:sample_size], y[sample_size:]
if sep_cohort:
self.pheno = self.pheno[self.pheno.eid.isin(test_ids)]
else:
self.pheno = self.pheno[self.pheno.eid.isin(train_val_ids)]
if not os.path.exists(plink_path):
print_sys('Saving...')
self.pheno[['eid', 'eid', self.pheno.columns.values[-1]]].to_csv(plink_path, header=None, index=None, sep=' ')
else:
print_sys('Already existed! Loading...')
self.pheno_plink = pd.read_csv(plink_path, header = None, sep = ' ')
return self.pheno_plink
def get_phenotype(self, field_id, aggregate = 'last_value', to_plink = False, to_str = True, normalize = 'None', frac = 1, random_seed = 42, fastgwa_match = False, icd10 = False, icd10_level = 2, sep_cohort = False, randomize = False, use_sample_size = False, sample_size = -1, randomize_seed = 42):
'''
example:
standing heights: 50
'''
pheno_path = os.path.join(self.cohort_data_path, str(field_id) + '_pheno.pkl')
if os.path.exists(pheno_path):
print_sys('Found local copy...')
self.pheno = load_dict(pheno_path)
else:
print_sys('construct phenotype from scratch...')
if icd10:
## field_id is icd10 level
icd10_df = self.get_icd10(to_plink = True, level = icd10_level, get_all = True)
self.pheno = icd10_df[['FID', field_id]].rename(columns = {'FID': 'eid'})
self.pheno['eid'] = self.pheno.eid.astype('int')
else:
## from raw data field id
self.pheno = get_fields([field_id], self.main_data_path)
save_dict(pheno_path, self.pheno)
print_sys('Done! Saving...')
if len(self.pheno.columns.values) > 2:
print_sys('There are multiple index for this phenotype... aggregate...')
if aggregate == 'last_value':
print_sys('Getting the latest measure...')
tmp = pd.DataFrame()
tmp['eid'] = self.pheno.loc[:, 'eid']
tmp[str(field_id)] = get_row_last_values(self.pheno.iloc[:, 1:])
self.pheno = tmp
print_sys('There are ' + str(len(self.pheno[self.pheno[str(field_id)].isnull()])) + ' samples with NaN values. Removing them ...')
self.pheno = self.pheno[self.pheno[str(field_id)].notnull()]
if fastgwa_match:
# get the number of without relatives:
if not self.keep_relatives:
raise ValueError('If you turned fastgwa_match = True, then keep_relatives = True!')
self.rel_ratio = len(self.pheno[self.pheno.eid.isin(self.no_rel_eid)])/len(self.pheno[self.pheno.eid.isin(self.cohort)])
# filtering to cohorts incl. with/without relatives
self.pheno = self.pheno[self.pheno.eid.isin(self.cohort)].reset_index(drop = True)
if normalize != 'None':
y = self.pheno[str(field_id)].values
if normalize == 'log':
y = np.log(y)
elif normalize == 'std':
y = (y - np.mean(y))/np.std(y)
elif normalize == 'quantile_normalization':
from sklearn.preprocessing import quantile_transform
y = quantile_transform(y.reshape(-1,1), output_distribution = 'normal', random_state = 42).reshape(-1)
self.pheno[str(field_id)] = y
if to_str:
self.pheno['eid'] = self.pheno['eid'].astype('str')
if not to_plink:
return self.pheno
else:
plink_path = os.path.join(self.cohort_data_path, str(field_id) + '_plink')
if self.keep_relatives:
plink_path = plink_path + '_with_relatives'
else:
plink_path = plink_path + '_no_relatives'
if normalize != 'None':
plink_path = plink_path + '_' + str(normalize)
if use_sample_size:
plink_path = plink_path + '_' + str(sample_size) + '_' + str(random_seed)
else:
if frac != 1:
plink_path = plink_path + '_' + str(frac) + '_' + str(random_seed)
if fastgwa_match:
plink_path += '_match'
if sep_cohort:
plink_path += '_sep_cohort'
if randomize:
plink_path += '_randomize' + str(randomize_seed)
plink_path = plink_path + '.txt'
if randomize:
self.pheno[str(field_id)] = self.pheno[str(field_id)].sample(frac = 1, random_state = randomize_seed).values
if use_sample_size:
from sklearn.model_selection import train_test_split
if icd10:
df_cases = self.pheno[self.pheno[str(field_id)] == 2]
df_cases_shuffle = df_cases.sample(frac = 1, random_state = random_seed)
all_ids, y = df_cases_shuffle.eid.values, df_cases_shuffle[str(field_id)].values
train_val_ids, test_ids, y_train_val, y_test = all_ids[:sample_size], all_ids[sample_size:], y[:sample_size], y[sample_size:]
train_val_ids = np.concatenate((train_val_ids, self.pheno[self.pheno[str(field_id)] == 1].eid.values))
self.pheno = self.pheno[self.pheno.eid.isin(train_val_ids)]
if sep_cohort:
raise NotImplementedError
else:
print('random_seed', random_seed)
pheno_shuffle = self.pheno.sample(frac = 1, random_state = random_seed)
all_ids, y = pheno_shuffle.eid.values, pheno_shuffle[str(field_id)].values
train_val_ids, test_ids, y_train_val, y_test = all_ids[:sample_size], all_ids[sample_size:], y[:sample_size], y[sample_size:]
if fastgwa_match:
raise ValueError('Not used anymore...')
if sep_cohort:
self.pheno = self.pheno[self.pheno.eid.isin(test_ids)]
else:
self.pheno = self.pheno[self.pheno.eid.isin(train_val_ids)]
else:
if frac!=1:
from sklearn.model_selection import train_test_split
all_ids, y = self.pheno.eid.values, self.pheno[str(field_id)].values
train_val_ids, test_ids, y_train_val, y_test = train_test_split(all_ids, y, test_size=frac, random_state=random_seed)
if fastgwa_match:
train_val_ids, test_ids, y_train_val, y_test = train_test_split(train_val_ids, y_train_val, test_size=1-self.rel_ratio, random_state=42)
if sep_cohort:
self.pheno = self.pheno[self.pheno.eid.isin(test_ids)]
else:
self.pheno = self.pheno[self.pheno.eid.isin(train_val_ids)]
if not os.path.exists(plink_path):
self.pheno[['eid', 'eid', self.pheno.columns.values[-1]]].to_csv(plink_path, header=None, index=None, sep=' ')
else:
print_sys('Already existed! Loading...')
self.pheno_plink = pd.read_csv(plink_path, header = None, sep = ' ')
return self.pheno_plink
def get_icd10(self, to_plink = False, level = 2, get_all = False):
'''
icd10: 41270
'''
pheno_path = os.path.join(self.cohort_data_path, 'icd10.pkl')
level_str = 'level' + str(level)
if os.path.exists(pheno_path):
print_sys('Found local copy...')
self.icd10 = load_dict(pheno_path)
else:
print_sys('construct from scratch...')
icd10_raw_concat = get_fields([41270], self.main_data_path)
icd10_columns = icd10_raw_concat.columns.values[1:]
icd10_tuple = icd10_raw_concat.apply(lambda x: (x.eid, x[icd10_columns][x[icd10_columns].notnull()].values), axis = 1)
icd10 = pd.DataFrame(list(icd10_tuple.values)).rename(columns = {0: 'eid', 1: 'level3'})
icd10['level2'] = icd10['level3'].apply(lambda x: np.unique([i[:3] for i in x]))
save_dict(pheno_path, icd10)
print_sys('Done! Saving...')
self.icd10 = icd10
if get_all:
self.pheno = self.icd10.reset_index(drop = True)
else:
self.pheno = self.icd10[self.icd10.eid.isin(self.cohort)].reset_index(drop = True)
if not to_plink:
return self.pheno
else:
if self.keep_relatives or get_all:
plink_path = os.path.join(self.cohort_data_path, 'icd10_plink_with_relatives_' + level_str + '.txt')
else:
plink_path = os.path.join(self.cohort_data_path, 'icd10_plink_no_relatives_' + level_str + '.txt')
if os.path.exists(plink_path):
print_sys("Found local copy...")
self.icd10_plink = pd.read_csv(plink_path, sep=' ')
else:
print_sys('transforming to plink files... takes around 1 min...')
unique_icd10 = np.unique([item for sublist in self.pheno[level_str].values for item in sublist])
icd10_2_idx = dict(zip(unique_icd10, range(len(unique_icd10))))
idx_2_icd10 = dict(zip(range(len(unique_icd10)), unique_icd10))
self.pheno[level_str + '_idx'] = self.pheno[level_str].apply(lambda x: [icd10_2_idx[i] for i in x])
tmp = np.zeros((len(self.pheno), len(unique_icd10)), dtype=np.int8)
for idx, i in enumerate(self.pheno[level_str + '_idx'].values):
tmp[idx, i] = 1
icd10_plink = pd.DataFrame(tmp).rename(columns = idx_2_icd10)
icd102sample_size = dict(icd10_plink.sum(axis = 0))
icd_100 = [i for i,j in icd102sample_size.items() if j > 100]
icd10_plink = icd10_plink + 1
icd10_plink['IID'] = self.pheno.eid.values
icd10_plink['FID'] = self.pheno.eid.values
icd10_plink = icd10_plink.loc[:, ['FID', 'IID'] + icd_100]
print_sys('Only using ICD10 codes with at least 100 cases...')
print_sys('There are ' + str(len(icd_100)) + ' ICD10 codes with at least 100 cases.')
icd10_plink.to_csv(plink_path, index=None, sep=' ')
self.icd10_plink = icd10_plink
return self.icd10_plink
Datasets
Models
