from torch_geometric.data import HeteroData
import torch_geometric.transforms as T
from sklearn.model_selection import train_test_split
import pandas as pd
import torch
import numpy as np
import pickle
import os
import tarfile
import urllib.request
import shutil
from tqdm import tqdm
import subprocess
from .utils import ldsc_regression_weights, load_dict
from .params import scdrs_traits
class KGWAS_Data:
def __init__(self, data_path='./data/'):
self.data_path = data_path
# Ensure the data path exists
if not os.path.exists(data_path):
os.makedirs(data_path)
# Check if relevant data exists in the data_path
required_files = [
'cell_kg/network/node_idx2id.pkl',
'cell_kg/network/edge_index.pkl',
'cell_kg/network/node_id2idx.pkl',
'cell_kg/node_emb/variant_emb/enformer_feat.pkl',
'cell_kg/node_emb/gene_emb/esm_feat.pkl',
'ld_score/filter_genotyped_ldscores.csv',
'ld_score/ldscores_from_data.csv',
'ld_score/ukb_white_ld_10MB_no_hla.pkl',
'ld_score/ukb_white_ld_10MB.pkl',
'misc_data/ukb_white_with_cm.bim',
]
missing_files = [f for f in required_files if not os.path.exists(os.path.join(data_path, f))]
if missing_files:
print("Relevant data not found in the data_path. Downloading and extracting data...")
url = "https://dataverse.harvard.edu/api/access/datafile/10731230"
file_name = 'kgwas_core_data'
self._download_and_extract_data(url, file_name)
else:
print("All required data files are present.")
def download_all_data(self):
url = "https://dataverse.harvard.edu/api/access/datafile/XXXX"
file_name = 'kgwas_data'
self._download_and_extract_data(url, file_name)
def _merge_with_rsync(self, src, dst):
"""Merge directories using rsync."""
try:
subprocess.run(
["rsync", "-a", "--ignore-existing", src + "/", dst + "/"],
check=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
except subprocess.CalledProcessError as e:
print(f"Error during rsync: {e.stderr.decode()}")
def _download_and_extract_data(self, url, file_name):
"""Download, extract, and merge directories using rsync."""
tar_file_path = os.path.join(self.data_path, f"{file_name}.tar.gz")
# Download the file
print(f"Downloading {file_name}.tar.gz...")
self._download_with_progress(url, tar_file_path)
print("Download complete.")
# Extract the tar.gz file
print("Extracting files...")
with tarfile.open(tar_file_path, 'r:gz') as tar:
tar.extractall(self.data_path)
print("Extraction complete.")
# Clean up the tar.gz file
os.remove(tar_file_path)
# Merge extracted contents into the data_path directory
extracted_dir = os.path.join(self.data_path, file_name)
if os.path.exists(extracted_dir):
print(f"Merging extracted directory '{extracted_dir}' into '{self.data_path}'...")
self._merge_with_rsync(extracted_dir, self.data_path)
# Remove the now-empty extracted directory
shutil.rmtree(extracted_dir)
def _download_with_progress(self, url, file_path):
"""Download a file with a progress bar."""
request = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
response = urllib.request.urlopen(request)
total_size = int(response.getheader('Content-Length').strip())
block_size = 1024 # 1 KB
with open(file_path, 'wb') as file, tqdm(
total=total_size, unit='B', unit_scale=True, desc="Downloading"
) as pbar:
while True:
buffer = response.read(block_size)
if not buffer:
break
file.write(buffer)
pbar.update(len(buffer))
def load_kg(self, snp_init_emb = 'enformer',
go_init_emb = 'random',
gene_init_emb = 'esm',
sample_edges = False,
sample_ratio = 1):
data_path = self.data_path
## Load KG
print('--loading KG---')
idx2id = load_dict(os.path.join(data_path, 'cell_kg/network/node_idx2id.pkl'))
edge_index_all = load_dict(os.path.join(data_path, 'cell_kg/network/edge_index.pkl'))
id2idx = load_dict(os.path.join(data_path, 'cell_kg/network/node_id2idx.pkl'))
self.id2idx = id2idx
self.idx2id = idx2id
data = HeteroData()
## Load initialized embeddings
if snp_init_emb == 'random':
print('--using random SNP embedding--')
data['SNP'].x = torch.rand((len(idx2id['SNP']), 128), requires_grad = False)
snp_init_dim_size = 128
elif snp_init_emb == 'kg':
print('--using KG SNP embedding--')
id2idx_kg = load_dict(os.path.join(data_path, 'cell_kg/node_emb/transe_emb/transe_emb_id2idx_kg.pkl'))
kg_emb = load_dict(os.path.join(data_path, 'cell_kg/node_emb/transe_emb/transe_emb_inverse_triplets.pkl'))
node_map = idx2id['SNP']
data['SNP'].x = torch.vstack([torch.tensor(kg_emb[id2idx_kg[node_map[i]]]) if node_map[i] in id2idx_kg \
else torch.rand(50, requires_grad = False) for i in range(len(node_map))])
snp_init_dim_size = 50
elif snp_init_emb == 'cadd':
print('--using CADD SNP embedding--')
df_variant = pd.read_csv(os.path.join(data_path, 'cell_kg/node_emb/variant_emb/cadd_feat.csv'))
df_variant = df_variant.set_index('Unnamed: 0')
variant_feat = df_variant.values
node_map = idx2id['SNP']
rs2idx_feat = dict(zip(df_variant.index.values, range(len(df_variant.index.values))))
data['SNP'].x = torch.vstack([torch.tensor(variant_feat[rs2idx_feat[node_map[i]]]) if node_map[i] in rs2idx_feat \
else torch.rand(64, requires_grad = False) for i in range(len(node_map))]).float()
snp_init_dim_size = 64
elif snp_init_emb == 'baselineLD':
print('--using baselineLD SNP embedding--')
node_map = idx2id['SNP']
rs2idx_feat = load_dict(os.path.join(data_path, 'cell_kg/node_emb/variant_emb/baselineld_feat.pkl'))
data['SNP'].x = torch.vstack([torch.tensor(rs2idx_feat[node_map[i]]) if node_map[i] in rs2idx_feat \
else torch.rand(70, requires_grad = False) for i in range(len(node_map))]).float()
snp_init_dim_size = 70
elif snp_init_emb == 'SLDSC':
print('--using SLDSC SNP embedding--')
node_map = idx2id['SNP']
rs2idx_feat = load_dict(os.path.join(data_path, 'cell_kg/node_emb/variant_emb/sldsc_feat.pkl'))
data['SNP'].x = torch.vstack([torch.tensor(rs2idx_feat[node_map[i]]) if node_map[i] in rs2idx_feat \
else torch.rand(165, requires_grad = False) for i in range(len(node_map))]).float()
snp_init_dim_size = 165
elif snp_init_emb == 'enformer':
print('--using enformer SNP embedding--')
node_map = idx2id['SNP']
rs2idx_feat = load_dict(os.path.join(data_path, 'cell_kg/node_emb/variant_emb/enformer_feat.pkl'))
data['SNP'].x = torch.vstack([torch.tensor(rs2idx_feat[node_map[i]]) if node_map[i] in rs2idx_feat \
else torch.rand(20, requires_grad = False) for i in range(len(node_map))]).float()
snp_init_dim_size = 20
if go_init_emb == 'random':
print('--using random go embedding--')
for rel in ['CellularComponent', 'BiologicalProcess', 'MolecularFunction']:
data[rel].x = torch.rand((len(idx2id[rel]), 128), requires_grad = False)
go_init_dim_size = 128
elif go_init_emb == 'kg':
print('--using KG go embedding--')
id2idx_kg = load_dict(os.path.join(data_path, 'cell_kg/node_emb/transe_emb/transe_emb_id2idx_kg.pkl'))
kg_emb = load_dict(os.path.join(data_path, 'cell_kg/node_emb/transe_emb/transe_emb_inverse_triplets.pkl'))
for rel in ['CellularComponent', 'BiologicalProcess', 'MolecularFunction']:
node_map = idx2id[rel]
data[rel].x = torch.vstack([torch.tensor(kg_emb[id2idx_kg[node_map[i]]]) if node_map[i] in id2idx_kg \
else torch.rand(50, requires_grad = False) for i in range(len(node_map))])
go_init_dim_size = 50
elif go_init_emb == 'biogpt':
print('--using biogpt go embedding--')
go2idx_feat = load_dict(os.path.join(data_path, 'cell_kg/node_emb/program_emb/biogpt_feat.pkl'))
for rel in ['CellularComponent', 'BiologicalProcess', 'MolecularFunction']:
node_map = idx2id[rel]
data[rel].x = torch.vstack([torch.tensor(go2idx_feat[node_map[i]]) if node_map[i] in go2idx_feat \
else torch.rand(1600, requires_grad = False) for i in range(len(node_map))]).float()
go_init_dim_size = 1600
if gene_init_emb == 'random':
print('--using random gene embedding--')
data['Gene'].x = torch.rand((len(idx2id['Gene']), 128), requires_grad = False)
gene_init_dim_size = 128
elif gene_init_emb == 'kg':
print('--using KG gene embedding--')
id2idx_kg = load_dict(os.path.join(data_path, 'cell_kg/node_emb/transe_emb/transe_emb_id2idx_kg.pkl'))
kg_emb = load_dict(os.path.join(data_path, 'cell_kg/node_emb/transe_emb/transe_emb_inverse_triplets.pkl'))
node_map = idx2id['Gene']
data['Gene'].x = torch.vstack([torch.tensor(kg_emb[id2idx_kg[node_map[i]]]) if node_map[i] in id2idx_kg \
else torch.rand(50, requires_grad = False) for i in range(len(node_map))])
gene_init_dim_size = 50
elif gene_init_emb == 'esm':
print('--using ESM gene embedding--')
gene2idx_feat = load_dict(os.path.join(data_path, 'cell_kg/node_emb/gene_emb/esm_feat.pkl'))
node_map = idx2id['Gene']
data['Gene'].x = torch.vstack([torch.tensor(gene2idx_feat[node_map[i]]) if node_map[i] in gene2idx_feat \
else torch.rand(5120, requires_grad = False) for i in range(len(node_map))]).float()
gene_init_dim_size = 5120
elif gene_init_emb == 'pops':
print('--using PoPs expression+PPI+pathways gene embedding--')
gene2idx_feat = load_dict(os.path.join(data_path, 'cell_kg/node_emb/gene_emb/pops_feat.pkl'))
node_map = idx2id['Gene']
data['Gene'].x = torch.vstack([torch.tensor(gene2idx_feat[node_map[i]]) if node_map[i] in gene2idx_feat \
else torch.rand(57742, requires_grad = False) for i in range(len(node_map))]).float()
gene_init_dim_size = 57742
elif gene_init_emb == 'pops_expression':
print('--using PoPs expression only gene embedding--')
gene2idx_feat = load_dict(os.path.join(data_path, 'cell_kg/node_emb/gene_emb/pops_expression_feat.pkl'))
node_map = idx2id['Gene']
data['Gene'].x = torch.vstack([torch.tensor(gene2idx_feat[node_map[i]]) if node_map[i] in gene2idx_feat \
else torch.rand(40546, requires_grad = False) for i in range(len(node_map))]).float()
gene_init_dim_size = 40546
self.gene_init_dim_size = gene_init_dim_size
self.go_init_dim_size = go_init_dim_size
self.snp_init_dim_size = snp_init_dim_size
for i,j in edge_index_all.items():
if sample_edges:
edge_index = torch.tensor(j)
num_edges = edge_index.size(1)
num_samples = int(num_edges * sample_ratio)
indices = torch.randperm(num_edges)[:num_samples]
sampled_edge_index = edge_index[:, indices]
print(i, ' sampling ratio ', sample_ratio, ' from ', edge_index.shape[1], ' to ', sampled_edge_index.shape[1])
data[i].edge_index = sampled_edge_index
else:
data[i].edge_index = torch.tensor(j)
data = T.ToUndirected()(data)
data = T.AddSelfLoops()(data)
self.data = data
def load_simulation_gwas(self, simulation_type, seed):
data_path = self.data_path
print('Using simulation data....')
small_cohort = 5000
num_causal_hits = 20000
heritability = 0.3
self.sample_size = small_cohort
if simulation_type == 'causal_link':
lr_uni = pd.read_csv(os.path.join(data_path, 'simulation_gwas/causal_link_simulation/' + str(num_causal_hits) + '_' + str(seed) + '_' + str(heritability) + '_graph_funct_v2_ggi.fastGWA'), sep = '\t')
elif simulation_type == 'causal':
lr_uni = pd.read_csv(os.path.join(data_path, 'simulation_gwas/causal_simulation/' + str(num_causal_hits) + '_' + str(seed) + '_' + str(heritability) + '_' + str(small_cohort) + '_graph_funct_v2.fastGWA'), sep = '\t')
elif simulation_type == 'null':
lr_uni = pd.read_csv(os.path.join(data_path, 'simulation_gwas/null_simulation/' + str(num_causal_hits) + '_' + str(seed) + '_' + str(heritability) + '_' + str(small_cohort) + '.fastGWA'), sep = '\t')
if ('SNP' in lr_uni.columns.values) and ('ID' in lr_uni.columns.values):
self.lr_uni = lr_uni.rename(columns = {'CHR': '#CHROM'})
else:
self.lr_uni = lr_uni.rename(columns = {'CHR': '#CHROM', 'SNP': 'ID'})
self.seed = seed
self.pheno = 'simulation'
def load_external_gwas(self, path = None, seed = 42, example_file = False):
if example_file:
print('Loading example GWAS file...')
url = "https://dataverse.harvard.edu/api/access/datafile/10730346"
example_file_path = os.path.join(self.data_path, 'biochemistry_Creatinine_fastgwa_full_10000_1.fastGWA')
# Check if the example file is already downloaded
if not os.path.exists(example_file_path):
print('Example file not found locally. Downloading...')
self._download_with_progress(url, example_file_path)
print('Example file downloaded successfully.')
else:
print('Example file already exists locally.')
path = example_file_path
if path is None:
raise ValueError("A valid path must be provided or example_file must be set to True.")
print(f'Loading GWAS file from {path}...')
lr_uni = pd.read_csv(path, sep=None, engine='python')
if 'CHR' not in lr_uni.columns.values:
raise ValueError('CHR chromosome not in the file!')
if 'SNP' not in lr_uni.columns.values:
raise ValueError('SNP column not in the file!')
if 'P' not in lr_uni.columns.values:
raise ValueError('P column not in the file!')
if 'N' not in lr_uni.columns.values:
raise ValueError('N column number of sample size not in the file!')
lr_uni = lr_uni.rename(columns = {'CHR': '#CHROM', 'SNP': 'ID'})
## filtering to the current KG variant set
old_variant_set_len = len(lr_uni)
lr_uni = lr_uni[lr_uni.ID.isin(list(self.idx2id['SNP'].values()))]
print('Number of SNPs in the KG:', len(self.idx2id['SNP']))
print('Number of SNPs in the GWAS:', old_variant_set_len)
print('Number of SNPs in the KG variant set:', len(lr_uni))
self.lr_uni = lr_uni
self.sample_size = lr_uni.N.values[0]
self.pheno = 'EXTERNAL'
self.seed = seed
def load_full_gwas(self, pheno, seed=42):
data_path = self.data_path
if pheno in scdrs_traits:
print('Using scdrs traits...')
self.pheno = pheno
lr_uni = pd.read_csv(os.path.join(data_path, 'scDRS_Data/sumstats_ukb_snps.csv'))
lr_uni = lr_uni[['CHR', 'SNP', 'POS', 'A1', 'A2', 'N', 'AF1', pheno]]
lr_uni = lr_uni[lr_uni[pheno].notnull()].reset_index(drop = True)
lr_uni = lr_uni.rename(columns = {'CHR': '#CHROM', 'SNP': 'ID', pheno: 'chi'})
print('number of SNPs:', len(lr_uni))
self.lr_uni = lr_uni
self.seed = seed
trait2size = pickle.load(open(os.path.join(data_path, 'scDRS_data/trait2size.pkl'), 'rb'))
self.sample_size = trait2size[pheno]
else:
## load GWAS files
self.pheno = pheno
lr_uni = pd.read_csv(os.path.join(data_path, 'full_gwas/' + str(self.pheno) + '_with_rel_fastgwa.fastGWA'), sep = '\t')
lr_uni = lr_uni.rename(columns = {'CHR': '#CHROM', 'SNP': 'ID'})
self.lr_uni = lr_uni
self.seed = seed
self.sample_size = 387113
def load_gwas_subsample(self, pheno, sample_size, seed):
data_path = self.data_path
if pheno in ['body_BALDING1', 'cancer_BREAST', 'disease_ALLERGY_ECZEMA_DIAGNOSED', 'disease_HYPOTHYROIDISM_SELF_REP', 'other_MORNINGPERSON', 'pigment_SUNBURN']:
binary = True
else:
binary = False
## load GWAS files
self.sample_size = sample_size
self.pheno = pheno
if (sample_size > 3000):
lr_uni = pd.read_csv(os.path.join(data_path, 'subsample_gwas/' + str(self.pheno) + \
'_fastgwa_full_'+ str(sample_size) + '_' + str(seed) + '.fastGWA'), sep = '\t')
lr_uni = lr_uni.rename(columns = {'CHR': '#CHROM', 'SNP': 'ID'})
else:
## use PLINK if sample size <3000
if binary:
lr_uni = pd.read_csv(os.path.join(data_path, 'subsample_gwas/' + str(self.pheno) + \
'_plink_'+ str(sample_size) + '_' + str(seed) + '.PHENO1.glm.logistic.hybrid'), sep = '\t')
else:
lr_uni = pd.read_csv(os.path.join(data_path, 'subsample_gwas/' + + str(self.pheno) + \
'_plink_'+ str(sample_size) + '_' + str(seed) + '.PHENO1.glm.linear'), sep = '\t')
self.lr_uni = lr_uni
self.seed = seed
def process_gwas_file(self, label = 'chi'):
data_path = self.data_path
lr_uni = self.lr_uni
## LD scores
ld_scores = pd.read_csv(os.path.join(data_path, 'ld_score/filter_genotyped_ldscores.csv'))
w_ld_scores = pd.read_csv(os.path.join(data_path, 'ld_score/ldscores_from_data.csv'))
m = 15000000
if 'N' not in lr_uni.columns.values:
n = self.sample_size
else:
n = np.mean(lr_uni.N)
h_g_2 = 0.5
rs_id_2_ld_scores = dict(ld_scores.values)
rs_id_2_ld_scores = dict(ld_scores.values)
rs_id_2_w_ld = dict(w_ld_scores.values)
## use min ld score for snps with no ld score
min_ld = min(rs_id_2_ld_scores.values())
lr_uni['ld_score'] = lr_uni.ID.apply(lambda x: rs_id_2_ld_scores[x] if x in rs_id_2_ld_scores else min_ld)
rs_id_2_ld_scores = dict(lr_uni[['ID', 'ld_score']].values)
min_ld = min(rs_id_2_w_ld.values())
## the data LD is without the query SNP itself. so here add 1
lr_uni['w_ld_score'] = 1 + lr_uni.ID.apply(lambda x: rs_id_2_w_ld[x] if x in rs_id_2_w_ld else min_ld)
rs_id_2_w_ld = dict(lr_uni[['ID', 'w_ld_score']].values)
print('Using ldsc weight...')
ld = np.array([rs_id_2_ld_scores[rs_id] for rs_id in lr_uni.ID.values])
w_ld = np.array([rs_id_2_w_ld[rs_id] for rs_id in lr_uni.ID.values])
ldsc_weight = ldsc_regression_weights(ld, w_ld, n, m, h_g_2)
ldsc_weight = ldsc_weight/np.mean(ldsc_weight)
print('ldsc_weight mean: ', np.mean(ldsc_weight))
self.rs_id_to_ldsc_weight = dict(zip(lr_uni.ID.values, ldsc_weight))
## chi-square label
if label == 'chi':
if 'chi' in lr_uni.columns.values:
print('chi pre-computed...')
lr_uni['y'] = lr_uni['chi'].values
else:
if self.pheno in (['body_BALDING1', 'cancer_BREAST', 'disease_ALLERGY_ECZEMA_DIAGNOSED', 'disease_HYPOTHYROIDISM_SELF_REP', 'other_MORNINGPERSON', 'pigment_SUNBURN']) and (self.sample_size <= 3000):
lr_uni['y'] = lr_uni['Z_STAT'].values**2
lr_uni['y'] = lr_uni.y.fillna(0)
else:
if ('BETA' in lr_uni.columns.values) and ('SE' in lr_uni.columns.values):
lr_uni['y'] = (lr_uni['BETA']/lr_uni['SE']).values**2
lr_uni['y'] = lr_uni.y.fillna(0)
else:
from scipy.stats import chi2
## convert from p-values
lr_uni['y'] = chi2.ppf(1 - lr_uni['P'].values, 1)
lr_uni['y'] = lr_uni.y.fillna(0)
elif label == 'residual-w-ld':
lr_uni['y'] = (lr_uni['BETA']/lr_uni['SE']).values**2
lr_uni['y'] = lr_uni.y.fillna(0)
lr_uni['ld_weight'] = lr_uni.ID.apply(lambda x: self.rs_id_to_ldsc_weight[x])
import statsmodels.api as sm
X = lr_uni.w_ld_score.values
y = lr_uni.y.values
weights = lr_uni.ld_weight.values
X = sm.add_constant(X)
model = sm.WLS(y, X, weights=weights)
results = model.fit()
y_pred = results.params[0] + results.params[1] * lr_uni.w_ld_score.values
lr_uni['y'] = y - y_pred
elif label == 'residual-ld':
lr_uni['y'] = (lr_uni['BETA']/lr_uni['SE']).values**2
lr_uni['y'] = lr_uni.y.fillna(0)
lr_uni['ld_weight'] = lr_uni.ID.apply(lambda x: self.rs_id_to_ldsc_weight[x])
import statsmodels.api as sm
X = lr_uni.ld_score.values
y = lr_uni.y.values
weights = lr_uni.ld_weight.values
X = sm.add_constant(X)
model = sm.WLS(y, X, weights=weights)
results = model.fit()
y_pred = results.params[0] + results.params[1] * lr_uni.w_ld_score.values
lr_uni['y'] = y - y_pred
elif label == 'residual-ld-ols':
lr_uni['y'] = (lr_uni['BETA']/lr_uni['SE']).values**2
lr_uni['y'] = lr_uni.y.fillna(0)
import statsmodels.api as sm
X = lr_uni.ld_score.values
y = lr_uni.y.values
X = sm.add_constant(X)
model = sm.OLS(y, X)
results = model.fit()
y_pred = results.params[0] + results.params[1] * lr_uni.w_ld_score.values
lr_uni['y'] = y - y_pred
elif label == 'residual-ld-ols-abs':
lr_uni['y'] = (lr_uni['BETA']/lr_uni['SE']).values**2
lr_uni['y'] = lr_uni.y.fillna(0)
import statsmodels.api as sm
X = lr_uni.ld_score.values
y = lr_uni.y.values
X = sm.add_constant(X)
model = sm.OLS(y, X)
results = model.fit()
y_pred = results.params[0] + results.params[1] * lr_uni.w_ld_score.values
lr_uni['y'] = np.abs(y - y_pred)
elif label == 'residual-w-ld-ols':
lr_uni['y'] = (lr_uni['BETA']/lr_uni['SE']).values**2
lr_uni['y'] = lr_uni.y.fillna(0)
import statsmodels.api as sm
X = lr_uni.w_ld_score.values
y = lr_uni.y.values
X = sm.add_constant(X)
model = sm.OLS(y, X)
results = model.fit()
y_pred = results.params[0] + results.params[1] * lr_uni.w_ld_score.values
lr_uni['y'] = y - y_pred
id2y = dict(lr_uni[['ID', 'y']].values)
all_ids = lr_uni.ID.values
self.all_ids = np.array([self.id2idx['SNP'][i] for i in all_ids])
self.y = lr_uni.y.values
#idx2y = dict(zip(self.all_ids, y))
self.lr_uni = lr_uni
def prepare_split(self, test_set_fraction_data = 0.05):
## split SNPs to train/test/valid
train_val_ids, test_ids, y_train_val, y_test = train_test_split(self.all_ids, self.y, test_size=test_set_fraction_data, random_state=self.seed)
train_ids, val_ids, y_train, y_val = train_test_split(train_val_ids, y_train_val, test_size=0.05, random_state=self.seed)
self.train_input_nodes = ('SNP', train_ids)
self.val_input_nodes = ('SNP', val_ids)
self.test_input_nodes = ('SNP', test_ids)
y_snp = torch.zeros(self.data['SNP'].x.shape[0]) - 1
y_snp[train_ids] = torch.tensor(y_train).float()
y_snp[val_ids] = torch.tensor(y_val).float()
y_snp[test_ids] = torch.tensor(y_test).float()
self.data['SNP'].y = y_snp
for i in self.data.node_types:
self.data[i].n_id = torch.arange(self.data[i].x.shape[0])
self.data.train_mask = train_ids
self.data.val_mask = val_ids
self.data.test_mask = test_ids
self.data.all_mask = self.all_ids
#data = data.to(args.device)
def get_pheno_list(self):
return {"large_cohort": scdrs_traits,
"21_indep_traits": ['body_BALDING1',
'disease_ALLERGY_ECZEMA_DIAGNOSED',
'disease_HYPOTHYROIDISM_SELF_REP', 'pigment_SUNBURN',
'21001', '50', '30080', '30070', '30010', '30000',
'biochemistry_AlkalinePhosphatase',
'biochemistry_AspartateAminotransferase',
'biochemistry_Cholesterol', 'biochemistry_Creatinine',
'biochemistry_IGF1', 'biochemistry_Phosphate',
'biochemistry_Testosterone_Male', 'biochemistry_TotalBilirubin',
'biochemistry_TotalProtein', 'biochemistry_VitaminD',
'bmd_HEEL_TSCOREz']}
Datasets
Models
