from copy import deepcopy
from tqdm import tqdm
import os
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
import pandas as pd
import numpy as np
import pickle
import subprocess
import torch
import torch.nn.functional as F
import torch.optim as optim
import torch.multiprocessing
torch.multiprocessing.set_sharing_strategy('file_system')
from torch_geometric.loader import NeighborLoader
from .utils import print_sys, compute_metrics, save_dict, \
load_dict, load_pretrained, save_model, \
evaluate_minibatch_clean, process_data, \
get_network_weight, generate_viz
from .eval_utils import storey_ribshirani_integrate, get_clumps_gold_label, get_meta_clumps, \
get_mega_clump_query, get_curve, find_closest_x
from .model import HeteroGNN
class KGWAS:
def __init__(self,
data,
weight_bias_track = False,
device = 'cuda',
proj_name = 'KGWAS',
exp_name = 'KGWAS',
seed = 42):
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
self.seed = seed
torch.backends.cudnn.enabled = False
use_cuda = torch.cuda.is_available()
self.device = device if use_cuda else "cpu"
self.data = data
self.data_path = data.data_path
if weight_bias_track:
import wandb
wandb.init(project=proj_name, name=exp_name)
self.wandb = wandb
else:
self.wandb = False
self.exp_name = exp_name
def initialize_model(self, gnn_num_layers = 2, gnn_hidden_dim = 128, gnn_backbone = 'GAT', gnn_aggr = 'sum', gat_num_head = 1, no_relu = False):
self.config = {
'gnn_num_layers': gnn_num_layers,
'gnn_hidden_dim': gnn_hidden_dim,
'gnn_backbone': gnn_backbone,
'gnn_aggr': gnn_aggr,
'gat_num_head': gat_num_head
}
self.gnn_num_layers = gnn_num_layers
self.model = HeteroGNN(self.data.data, gnn_hidden_dim, 1,
gnn_num_layers, gnn_backbone, gnn_aggr,
self.data.snp_init_dim_size,
self.data.gene_init_dim_size,
self.data.go_init_dim_size,
gat_num_head,
no_relu = no_relu,
).to(self.device)
def load_pretrained(self, path):
with open(os.path.join(path, 'config.pkl'), 'rb') as f:
config = pickle.load(f)
self.initialize_model(**config)
self.config = config
self.model = load_pretrained(path, self.model)
self.best_model = self.model
self.kgwas_res = pd.read_csv(os.path.join(path, 'pred.csv'), sep = None, engine = 'python')
self.save_name = path.split('/')[-1]
def train(self, batch_size = 512, num_workers = 0, lr = 1e-4,
weight_decay = 5e-4, epoch = 10, save_best_model = True,
save_name = None, data_to_cuda = False):
total_epoch = epoch
if save_name is None:
save_name = self.exp_name
self.save_name = save_name
print_sys('Creating data loader...')
kwargs = {'batch_size': batch_size, 'num_workers': num_workers, 'drop_last': True}
eval_kwargs = {'batch_size': 512, 'num_workers': num_workers, 'drop_last': False}
if data_to_cuda:
self.data.data = self.data.data.to(self.device)
self.train_loader = NeighborLoader(self.data.data, num_neighbors=[-1] * self.gnn_num_layers,
sampler = None,
input_nodes=self.data.train_input_nodes, **kwargs)
self.val_loader = NeighborLoader(self.data.data, num_neighbors=[-1] * self.gnn_num_layers,
input_nodes=self.data.val_input_nodes, **kwargs)
self.test_loader = NeighborLoader(self.data.data, num_neighbors=[-1] * self.gnn_num_layers,
input_nodes=self.data.test_input_nodes, **eval_kwargs)
X_infer = self.data.lr_uni.ID.values
#print_sys('# of to-infer SNPs: ' + str(len(X_infer)))
infer_idx = np.array([self.data.id2idx['SNP'][i] for i in X_infer])
infer_input_nodes = ('SNP', infer_idx)
self.infer_loader = NeighborLoader(self.data.data, num_neighbors=[-1] * self.gnn_num_layers,
input_nodes=infer_input_nodes, **eval_kwargs)
## model training
optimizer = optim.Adam(self.model.parameters(), lr=lr, weight_decay = weight_decay)
loss_fct = F.mse_loss
earlystop_validation_metric = 'pearsonr'
binary_output = False
earlystop_direction = 'ascend'
min_val = -1000
self.best_model = deepcopy(self.model).to(self.device)
print_sys('Start Training...')
for epoch in range(total_epoch):
self.model.train()
for step, batch in enumerate(tqdm(self.train_loader, desc=f"Training Progress Epoch {epoch+1}/{total_epoch}", total=len(self.train_loader))):
optimizer.zero_grad()
if data_to_cuda:
pass
#batch = batch.to(self.device, 'edge_index')
else:
batch = batch.to(self.device)
bs_batch = batch['SNP'].batch_size
out = self.model(batch.x_dict, batch.edge_index_dict, bs_batch)
pred = out.reshape(-1)
y_batch = batch['SNP'].y[:bs_batch]
rs_id = [self.data.idx2id['SNP'][i.item()] for i in batch['SNP']['n_id'][:bs_batch]]
ld_weight = torch.tensor([self.data.rs_id_to_ldsc_weight[i] for i in rs_id]).to(self.device)
loss = torch.mean(ld_weight * (pred - y_batch)**2)
if self.wandb:
self.wandb.log({'training_loss': loss.item()})
loss.backward()
optimizer.step()
if (step % 500 == 0) and (step >= 500):
log = "Epoch {} Step {} Train Loss: {:.4f}"
print_sys(log.format(epoch + 1, step + 1, loss.item()))
val_res = evaluate_minibatch_clean(self.val_loader, self.model, self.device)
val_metrics = compute_metrics(val_res, binary_output, -1, -1, loss_fct)
log = "Epoch {}: Validation MSE: {:.4f} " \
"Validation Pearson: {:.4f}. "
print_sys(log.format(epoch + 1, val_metrics['mse'],
val_metrics['pearsonr']))
if self.wandb:
for i,j in val_metrics.items():
self.wandb.log({'val_' + i: j})
if val_metrics[earlystop_validation_metric] > min_val:
min_val = val_metrics[earlystop_validation_metric]
self.best_model = deepcopy(self.model)
best_epoch = epoch
if save_best_model:
save_model_path = self.data_path + '/model/'
print_sys('Saving models to ' + os.path.join(save_model_path, save_name))
save_model(self.best_model, self.config, os.path.join(save_model_path, save_name))
test_res = evaluate_minibatch_clean(self.test_loader, self.best_model, self.device)
test_metric = compute_metrics(test_res, binary_output, -1, -1, loss_fct)
if self.wandb:
for i,j in test_metric.items():
self.wandb.log({'test_' + i: j})
infer_res = evaluate_minibatch_clean(self.infer_loader, self.best_model, self.device)
self.data.lr_uni['pred'] = infer_res['pred']
lr_uni_to_save = deepcopy(self.data.lr_uni)
self.data.lr_uni['abs_pred'] = np.abs(self.data.lr_uni['pred'])
self.data.lr_uni['SR_P_val'] = storey_ribshirani_integrate(self.data.lr_uni, column = 'abs_pred', num_bins = 500)
self.data.lr_uni['SR'] = -(np.log10(self.data.lr_uni['SR_P_val'].astype(float).values))
lr_uni_to_save['P_weighted'] = self.data.lr_uni['SR_P_val']
## calibration
scale_factor = find_closest_x(lr_uni_to_save)
lr_uni_to_save['KGWAS_P'] = scale_factor * lr_uni_to_save['P_weighted']
lr_uni_to_save['KGWAS_P'] = lr_uni_to_save['KGWAS_P'].clip(lower=0, upper=1)
if not os.path.exists(self.data_path + '/model_pred/'):
os.makedirs(self.data_path + '/model_pred/')
os.makedirs(self.data_path + '/model_pred/new_experiments/')
lr_uni_to_save.to_csv(self.data_path + '/model_pred/new_experiments/' + save_name + '_pred.csv', index = False, sep = '\t')
print('KGWAS prediction and p-values saved to ' + self.data_path + '/model_pred/new_experiments/' + save_name + '_pred.csv')
if save_best_model:
lr_uni_to_save.to_csv(self.data_path + '/model/' + save_name + '/pred.csv', index = False, sep = '\t')
self.kgwas_res = lr_uni_to_save
def run_magma(self, path_to_magma, bfile):
if 'N' in self.kgwas_res.columns:
n_value = self.kgwas_res['N'].values[0]
else:
n_value = input("Please provide the sample size for the GWAS analysis.")
url = "https://dataverse.harvard.edu/api/access/datafile/10731670"
annot_file_path = os.path.join(self.data_path, 'gene_annotation.genes.annot')
# Check if the example file is already downloaded
if not os.path.exists(annot_file_path):
print('Annotation file not found locally. Downloading...')
self.data._download_with_progress(url, annot_file_path)
print('Annotation file downloaded successfully.')
else:
print('Annotation file already exists locally.')
gene_annot = annot_file_path
magma_path = self.data_path + '/model_pred/new_experiments/' + self.save_name + '_magma_format.csv'
self.kgwas_res[['ID', 'KGWAS_P']].rename(columns = {'ID': 'SNP', 'KGWAS_P': 'P'}).to_csv(magma_path, index = False, sep = '\t')
# Construct the MAGMA command
command = [
path_to_magma,
"--bfile", bfile,
"--gene-annot", gene_annot,
"--pval", magma_path, f"N={n_value}",
"--out", self.data_path + '/model_pred/new_experiments/' + self.save_name + '_magma_out'
]
try:
# Run the command with real-time output
process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
print("Running MAGMA...")
# Stream stdout line by line
for line in process.stdout:
print(line, end="") # Print each line as it's received
# Wait for the process to complete and capture stderr
stderr = process.communicate()[1]
if process.returncode == 0:
print("MAGMA command executed successfully.")
else:
print("MAGMA encountered an error.")
print("Error message:", stderr)
except FileNotFoundError:
print("MAGMA executable not found. Ensure it is in the specified path.")
except Exception as e:
print(f"An unexpected error occurred: {e}")
def get_disease_critical_network(self, variant_threshold = 5e-8,
magma_path = None, magma_threshold = 0.05, program_threshold = 0.05,
K_neighbors = 3, num_cpus = 1):
df_network_weight = get_network_weight(self, self.data)
df_variant_interpretation, disease_critical_network = generate_viz(self, df_network_weight, self.data_path, variant_threshold, magma_path, magma_threshold, program_threshold, K_neighbors, num_cpus)
return df_network_weight, df_variant_interpretation, disease_critical_network
Datasets
Models
