import os, sys
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import connected_components
import pandas as pd
import numpy as np
import pickle
from tqdm import tqdm
from scipy.stats import pearsonr
from sklearn.metrics import mean_squared_error, precision_score
import torch
from torch.nn import functional as F
from torch import nn
from multiprocessing import Pool
from tqdm import tqdm
from functools import partial
from .params import main_data_path, cohort_data_path, kinship_path, withdraw_path
def evaluate_minibatch_clean(loader, model, device):
model.eval()
pred_all = []
truth = []
results = {}
for step, batch in enumerate(tqdm(loader)):
batch = batch.to(device)
bs_batch = batch['SNP'].batch_size
out = model(batch.x_dict, batch.edge_index_dict, bs_batch)
pred = out.reshape(-1)
y_batch = batch['SNP'].y[:bs_batch]
pred_all.extend(pred.detach().cpu().numpy())
truth.extend(y_batch.detach().cpu().numpy())
del y_batch, pred, batch, out
results['pred'] = np.hstack(pred_all)
results['truth'] = np.hstack(truth)
return results
def compute_metrics(results, binary, coverage = None, uncertainty_reg = 1, loss_fct = None):
metrics = {}
metrics['mse'] = mean_squared_error(results['pred'], results['truth'])
metrics['pearsonr'] = pearsonr(results['pred'], results['truth'])[0]
return metrics
'''
requires to modify the pyg source code since it does not support heterogeneous graph attention
miniconda3/envs/a100_env/lib/python3.8/site-packages/torch_geometric/nn/conv/hgt_conv.py
def group(xs: List[Tensor], aggr: Optional[str]) -> Optional[Tensor]:
if len(xs) == 0:
return None
elif aggr is None:
return torch.stack(xs, dim=1)
elif len(xs) == 1:
return xs[0]
elif isinstance(xs, list) and isinstance(xs[0], tuple):
xs_old = [i[0] for i in xs]
out = torch.stack(xs_old, dim=0)
out = getattr(torch, aggr)(out, dim=0)
out = out[0] if isinstance(out, tuple) else out
att = [i[1] for i in xs]
return (out, att)
else:
out = torch.stack(xs, dim=0)
out = getattr(torch, aggr)(out, dim=0)
out = out[0] if isinstance(out, tuple) else out
return out
'''
def get_attention_weight(model, x_dict, edge_index_dict):
attention_all_layers = []
for conv in model.convs:
out = conv(x_dict, edge_index_dict, return_attention_weights_dict = dict(zip(list(data.edge_index_dict.keys()), [True] * len(list(data.edge_index_dict.keys())))))
x_dict = {i: j[0] for i,j in out.items()}
attention_layer = {i: j[1] for i,j in out.items()}
attention_all_layers.append(attention_layer)
x_dict = {key: x.relu() for key, x in x_dict.items()}
idx2n_id = {}
for i in batch.node_types:
idx2n_id[i] = dict(zip(range(len(batch[i].n_id)), batch[i].n_id.numpy()))
node_type = 'SNP'
edge2weight_l1 = {}
edge2weight_l2 = {}
edge_type_node = [i for i,j in batch.edge_index_dict.items() if i[2] == node_type]
edge_type_node_len = [j.shape[1] for i,j in batch.edge_index_dict.items() if i[2] == node_type]
for idx, edge_type in enumerate(edge_type_node):
edge2weight_l1[edge_type] = attention_all_layers[0][node_type][idx]
assert edge_type_node_len[idx] == edge2weight_l1[edge_type][0].shape[1]
edge2weight_l2[edge_type] = attention_all_layers[1][node_type][idx]
assert edge_type_node_len[idx] == edge2weight_l2[edge_type][0].shape[1]
edge2weight_l1[edge_type][0][0] = torch.LongTensor([idx2n_id[edge_type[0]][ent] for ent in edge2weight_l1[edge_type][0][0].detach().cpu().numpy()])
edge2weight_l1[edge_type][0][1] = torch.LongTensor([idx2n_id[edge_type[2]][ent] for ent in edge2weight_l1[edge_type][0][1].detach().cpu().numpy()])
return edge2weight_l1, edge2weight_l2
def get_fields(all_field_ids, main_data_path):
headers = pd.read_csv(main_data_path, nrows = 1).columns
relevant_headers = [i for i, header in enumerate(headers) if header == 'eid' or \
any([header.startswith('%d-' % field_id) for field_id in all_field_ids])]
return pd.read_csv(main_data_path, usecols = relevant_headers)
def get_row_last_values(df):
result = pd.Series(np.nan, index = df.index)
for column in df.columns[::-1]:
result = result.where(pd.notnull(result), df[column])
return result
def remove_kinships(eid, verbose = True):
'''
Determines which samples need to be removed such that the remaining samples will have no kinship connections whatsoever (according to the
kinship table provided by the UKBB). In order to determine that, kinship groups will first be determined (@see get_kinship_groups), and
only one sample will remain within each of the groups. For the sake of determinism, the sample with the lowest eid will be selected within
each kinship group, and the rest will be discarded.
@param eid (pd.Series): A series whose values are UKBB sample IDs, from which kinships should be removed.
@param verbose (bool): Whether to log details of the operation of this function.
@return: A mask of samples to keep (pd.Series with index corresponding to the eid input, and boolean values).
'''
all_eids = set(eid)
kinship_groups = get_kinship_groups()
relevant_kinship_groups = [kinship_group & all_eids for kinship_group in kinship_groups]
relevant_kinship_groups = [kinship_group for kinship_group in relevant_kinship_groups if len(kinship_group) >= 2]
unchosen_kinship_representatives = set.union(*[set(sorted(kinship_group)[1:]) for kinship_group in relevant_kinship_groups])
no_kinship_mask = ~eid.isin(unchosen_kinship_representatives)
if verbose:
print_sys(('Constructed %d kinship groups (%d samples), of which %d (%d samples) are relevant for the dataset (i.e. containing at least 2 ' + \
'samples in the dataset). Picking only one representative of each group and removing the %d other samples in those groups ' + \
'has reduced the dataset from %d to %d samples.') % (len(kinship_groups), len(set.union(*kinship_groups)), \
len(relevant_kinship_groups), len(set.union(*relevant_kinship_groups)), len(unchosen_kinship_representatives), len(no_kinship_mask), \
no_kinship_mask.sum()))
return no_kinship_mask
def get_kinship_groups():
'''
Uses the kinship table provided by the UKBB (as specified by the KINSHIP_TABLE_FILE_PATH configuration) in order to determine kinship groups.
Each kinship group is a connected component of samples in the graph of kinships (where each node is a UKBB sample, and an edge exists between
each pair of samples reported in the kinship table).
@return: A list of sets of strings (the strings are the sample IDs, i.e. eid). Each set of samples is a kinship group.
'''
kinship_table = pd.read_csv(kinship_path, sep = ' ')
kinship_ids = np.array(sorted(set(kinship_table['ID1']) | set(kinship_table['ID2'])))
n_kinship_ids = len(kinship_ids)
kinship_id_to_index = pd.Series(np.arange(n_kinship_ids), index = kinship_ids)
kinship_index1 = kinship_table['ID1'].map(kinship_id_to_index).values
kinship_index2 = kinship_table['ID2'].map(kinship_id_to_index).values
symmetric_kinship_index1 = np.concatenate([kinship_index1, kinship_index2])
symmetric_kinship_index2 = np.concatenate([kinship_index2, kinship_index1])
kinship_matrix = csr_matrix((np.ones(len(symmetric_kinship_index1), dtype = bool), (symmetric_kinship_index1, \
symmetric_kinship_index2)), shape = (n_kinship_ids, n_kinship_ids), dtype = bool)
_, kinship_labels = connected_components(kinship_matrix, directed = False)
kinship_labels = pd.Series(kinship_labels, index = kinship_ids)
return [set(group_kinship_labels.index) for _, group_kinship_labels in kinship_labels.groupby(kinship_labels)]
def save_dict(path, obj):
"""save an object to a pickle file
Args:
path (str): the path to save the pickle file
obj (object): any file
"""
with open(path, 'wb') as f:
pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)
def load_dict(path):
"""load an object from a path
Args:
path (str): the path where the pickle file locates
Returns:
object: loaded pickle file
"""
with open(path, 'rb') as f:
return pickle.load(f)
def save_model(model, config, path_dir):
if not os.path.exists(path_dir):
os.makedirs(path_dir)
torch.save(model.state_dict(), path_dir + '/model.pt')
save_dict(path_dir + '/config.pkl', config)
def load_pretrained(path, model):
state_dict = torch.load(os.path.join(path, 'model.pt'), map_location = torch.device('cpu'))
# to support training from multi-gpus data-parallel:
if next(iter(state_dict))[:7] == 'module.':
# the pretrained model is from data-parallel module
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
state_dict = new_state_dict
model.load_state_dict(state_dict)
return model
def get_args(path):
return load_dict(os.path.join(path, 'config.pkl'))
def print_sys(s):
"""system print
Args:
s (str): the string to print
"""
print(s, flush = True, file = sys.stderr)
def get_plink_QC_fam():
fam_path = '/dfs/project/datasets/20220524-ukbiobank/data/genetics/ukb_all.fam'
data = ukbb_cohort(main_data_path, cohort_data_path, withdraw_path, keep_relatives=True).cohort
df_fam = pd.read_csv(fam_path, sep = ' ', header = None)
df_fam[df_fam[0].isin(data)].reset_index(drop = True).to_csv('/dfs/project/datasets/20220524-ukbiobank/data/cohort/qc_cohort.txt', header = None, index = False, sep = ' ')
def get_plink_no_rel_fam():
fam_path = '/dfs/project/datasets/20220524-ukbiobank/data/genetics/ukb_all.fam'
data = ukbb_cohort(main_data_path, cohort_data_path, withdraw_path, keep_relatives=False).cohort
df_fam = pd.read_csv(fam_path, sep = ' ', header = None)
df_fam[df_fam[0].isin(data)].reset_index(drop = True).to_csv('/dfs/project/datasets/20220524-ukbiobank/data/cohort/no_rel.fam', header = None, index = False, sep = ' ')
def get_precision_recall_at_N(res, hits_all, input_dim, N, column_rsid = 'ID', thres = 5e-8):
eval_dict = {}
hits_sub = res[res.P < thres][column_rsid].values
p_sorted = res.sort_values('P')[column_rsid].values
for K in range(1, input_dim, 10000):
topK_true = np.intersect1d(hits_all, p_sorted[:K])
recall = len(topK_true)/len(hits_all)
if recall > N:
break
for K in range(K-10000, K, 1000):
topK_true = np.intersect1d(hits_all, p_sorted[:K])
recall = len(topK_true)/len(hits_all)
if recall > N:
break
for K in range(K-1000, K, 100):
topK_true = np.intersect1d(hits_all, p_sorted[:K])
recall = len(topK_true)/len(hits_all)
if recall > N:
break
for K in range(K-100, K, 10):
topK_true = np.intersect1d(hits_all, p_sorted[:K])
recall = len(topK_true)/len(hits_all)
if recall > N:
break
for K in range(K-10, K):
topK_true = np.intersect1d(hits_all, p_sorted[:K])
recall = len(topK_true)/len(hits_all)
if recall > N:
break
print_sys('PR@' + str(int(N * 100)) + ' is achieved when K = ' + str(K))
eval_dict['PR@' + str(int(N * 100)) + '_K'] = K
topK_true = [1 if i in hits_all else 0 for i in p_sorted[:K]]
precision = precision_score(topK_true, [1] * K)
eval_dict['PR@' + str(int(N * 100))] = precision
return eval_dict
def get_gwas_results(res, hits_all, input_dim, column_rsid = 'ID', thres = 5e-8):
eval_dict = {}
hits_sub = res[res.P < thres][column_rsid].values
eval_dict['overall_recall'] = len(np.intersect1d(hits_sub, hits_all))/len(hits_all)
if len(hits_sub) == 0:
eval_dict['overall_precision'] = 0
eval_dict['overall_f1'] = 0
else:
eval_dict['overall_precision'] = len(np.intersect1d(hits_sub, hits_all))/len(hits_sub)
eval_dict['overall_f1'] = 2 * eval_dict['overall_recall'] * eval_dict['overall_precision']/(eval_dict['overall_recall'] + eval_dict['overall_precision'])
for K in [100, 500, 1000, 5000]:
topK_true = [1 if i in hits_all else 0 for i in res.sort_values('P').iloc[:K][column_rsid].values]
eval_dict['precision_' + str(K)] = precision_score(topK_true, [1] * K)
eval_dict['recall_' + str(K)] = sum(topK_true)/len(hits_all)
eval_dict.update(get_precision_recall_at_N(res, hits_all, input_dim, 0.8, column_rsid, thres))
eval_dict.update(get_precision_recall_at_N(res, hits_all, input_dim, 0.9, column_rsid, thres))
eval_dict.update(get_precision_recall_at_N(res, hits_all, input_dim, 0.95, column_rsid, thres))
return eval_dict
def find_nearest(array, value):
array = np.asarray(array)
idx = (np.abs(array - value)).argmin()
return array[idx]
def get_preds(logits, multi_label):
if multi_label:
preds = (logits.sigmoid() > 0.5).float()
elif logits.shape[1] > 1: # multi-class
preds = logits.argmax(dim=1).float()
else: # binary
preds = (logits.sigmoid() > 0.5).float()
return preds
def process_data(data, use_edge_attr):
if not use_edge_attr:
data.edge_attr = None
if data.get('edge_label', None) is None:
data.edge_label = {i: torch.zeros(j.shape[1]) for i, j in data.edge_index_dict.items()}
return data
def load_checkpoint(model, model_dir, model_name, map_location=None):
checkpoint = torch.load(model_dir / (model_name + '.pt'), map_location=map_location)
model.load_state_dict(checkpoint['model_state_dict'])
def save_checkpoint(model, model_dir, model_name):
torch.save({'model_state_dict': model.state_dict()}, model_dir / (model_name + '.pt'))
def get_lr(optimizer):
for param_group in optimizer.param_groups:
return param_group['lr']
def flatten(list_of_lists):
return [item for sublist in list_of_lists for item in sublist]
def find_connected_components_details(edges):
graph = {}
for u, v in edges:
if u not in graph:
graph[u] = []
if v not in graph:
graph[v] = []
graph[u].append(v)
graph[v].append(u)
def dfs(vertex):
visited_nodes = set()
visited_edges = set()
stack = [vertex]
while stack:
current = stack.pop()
if current not in visited_nodes:
visited_nodes.add(current)
for neighbor in graph[current]:
stack.append(neighbor)
if (current, neighbor) not in visited_edges and (neighbor, current) not in visited_edges:
visited_edges.add((current, neighbor))
return list(visited_nodes), list(visited_edges)
visited = set()
components = []
for vertex in tqdm(graph):
if vertex not in visited:
nodes, edges = dfs(vertex)
components.append({
'nodes': nodes,
'edges': edges
})
visited.update(nodes)
return components
def flatten(lst):
return [item for sublist in lst for item in sublist]
def ldsc_regression_weights(ld, w_ld, N, M, hsq, intercept=None, ii=None):
'''
Regression weights.
Parameters
----------
ld : np.matrix with shape (n_snp, 1)
LD Scores (non-partitioned).
w_ld : np.matrix with shape (n_snp, 1)
LD Scores (non-partitioned) computed with sum r^2 taken over only those SNPs included
in the regression.
N : np.matrix of ints > 0 with shape (n_snp, 1)
Number of individuals sampled for each SNP.
M : float > 0
Number of SNPs used for estimating LD Score (need not equal number of SNPs included in
the regression).
hsq : float in [0,1]
Heritability estimate.
Returns
-------
w : np.matrix with shape (n_snp, 1)
Regression weights. Approx equal to reciprocal of conditional variance function.
'''
M = float(M)
if intercept is None:
intercept = 1
hsq = max(hsq, 0.0)
hsq = min(hsq, 1.0)
ld = np.fmax(ld, 1.0)
w_ld = np.fmax(w_ld, 1.0)
c = hsq * N / M
het_w = 1.0 / (2 * np.square(intercept + np.multiply(c, ld)))
oc_w = 1.0 / w_ld
w = np.multiply(het_w, oc_w)
return w
def get_network_weight(run, data):
model = run.best_model
model = model.to('cpu')
graph_data = data.data.to('cpu')
x_dict, edge_index_dict = graph_data.x_dict, graph_data.edge_index_dict
attention_all_layers = []
print('Retrieving weights...')
x_dict['SNP'] = model.snp_feat_mlp(x_dict['SNP'])
x_dict['Gene'] = model.gene_feat_mlp(x_dict['Gene'])
x_dict['CellularComponent'] = model.go_feat_mlp(x_dict['CellularComponent'])
x_dict['BiologicalProcess'] = model.go_feat_mlp(x_dict['BiologicalProcess'])
x_dict['MolecularFunction'] = model.go_feat_mlp(x_dict['MolecularFunction'])
for conv in model.convs:
x_dict = conv(x_dict, edge_index_dict,
return_attention_weights_dict = dict(zip(list(graph_data.edge_index_dict.keys()),
[True] * len(list(graph_data.edge_index_dict.keys())))),
return_raw_attention_weights_dict = dict(zip(list(graph_data.edge_index_dict.keys()),
[True] * len(list(graph_data.edge_index_dict.keys())))),
)
attention_layer = {i: j[1] for i,j in x_dict.items()}
attention_all_layers.append(attention_layer)
x_dict = {i: j[0] for i,j in x_dict.items()}
layer2rel2att = {
'l1': {},
'l2': {}
}
print('Aggregating across node types...')
for node_type in graph_data.x_dict.keys():
edge_type_node = [i for i,j in graph_data.edge_index_dict.items() if i[2] == node_type]
for idx, i in enumerate(attention_all_layers[0][node_type]):
layer2rel2att['l1'][edge_type_node[idx]] = np.vstack((i[0].detach().cpu().numpy(), i[1].detach().cpu().numpy().reshape(-1)))
for idx, i in enumerate(attention_all_layers[1][node_type]):
layer2rel2att['l2'][edge_type_node[idx]] = np.vstack((i[0].detach().cpu().numpy(), i[1].detach().cpu().numpy().reshape(-1)))
df_val_all = pd.DataFrame()
for rel, value in layer2rel2att['l1'].items():
df_val = pd.DataFrame(value).T.rename(columns = {0: 'h_idx', 1: 't_idx', 2: 'weight'})
df_val['h_type'] = rel[0]
df_val['rel_type'] = rel[1]
df_val['t_type'] = rel[2]
df_val['layer'] = 'l1'
df_val_all = df_val_all.append(df_val)
for rel, value in layer2rel2att['l2'].items():
df_val = pd.DataFrame(value).T.rename(columns = {0: 'h_idx', 1: 't_idx', 2: 'weight'})
df_val['h_type'] = rel[0]
df_val['rel_type'] = rel[1]
df_val['t_type'] = rel[2]
df_val['layer'] = 'l2'
df_val_all = df_val_all.append(df_val)
df_val_all = df_val_all.drop_duplicates(['h_idx', 't_idx', 'rel_type', 'layer'])
return df_val_all
def get_local_interpretation(query_snp, v2g, g2g, g2p, g2v, id2idx, K_neighbors):
try:
snp2gene_around_snp = v2g[v2g.t_idx == id2idx['SNP'][query_snp]]
snp2gene_around_snp = snp2gene_around_snp.sort_values('importance')[::-1]
gene_hit = snp2gene_around_snp.iloc[:K_neighbors]
gene_hit.loc[:, 'rel_type'] = gene_hit.rel_type.apply(lambda x: x[4:])
g2g_focal = pd.DataFrame()
for gene in gene_hit.h_id.values:
g2g_focal = g2g_focal.append(g2g[g2g.t_id == gene].sort_values('importance')[::-1].iloc[:K_neighbors])
g2g_focal.loc[:,'rel_type'] = g2g_focal.rel_type.apply(lambda x: x.split('-')[1])
g2p_focal = pd.DataFrame()
for gene in gene_hit.h_id.values:
g2p_focal = g2p_focal.append(g2p[g2p.t_id == gene].sort_values('importance')[::-1].iloc[:K_neighbors])
g2p_focal.loc[:,'rel_type'] = g2p_focal.rel_type.apply(lambda x: x.split('-')[1])
g2v_focal = pd.DataFrame()
for gene in gene_hit.h_id.values:
g2v_focal = g2v_focal.append(g2v[g2v.t_id == gene].sort_values('importance')[::-1].iloc[:K_neighbors])
local_neighborhood_around_snp = pd.concat((gene_hit, g2g_focal, g2p_focal, g2v_focal))
local_neighborhood_around_snp.loc[:,'QUERY_SNP'] = query_snp
return local_neighborhood_around_snp
except:
return None
def generate_viz(run, df_network, data_path, variant_threshold = 5e-8,
magma_path = None, magma_threshold = 0.05, program_threshold = 0.05,
K_neighbors = 3, num_cpus = 1):
gwas = run.kgwas_res
idx2id = run.data.idx2id
id2idx = run.data.id2idx
print('Start generating disease critical network...')
gene_sets = load_dict(os.path.join(data_path, 'misc_data/gene_set_bp.pkl'))
with open(os.path.join(data_path, 'misc_data/go2name.pkl'), 'rb') as f:
go2name = pickle.load(f)
df_network = df_network[~df_network.rel_type.isin(['TSS', 'rev_TSS'])]
snp2genes = df_network[(df_network.t_type == 'SNP')
& (df_network.h_type == 'Gene')]
gene2gene = df_network[(df_network.t_type == 'Gene')
& (df_network.h_type == 'Gene')]
gene2go = df_network[(df_network.t_type == 'Gene')
& (df_network.h_type.isin(['BiologicalProcess']))]
if 'SNP' not in gwas.columns.values:
gwas.loc[:, 'SNP'] = gwas['ID']
hit_snps = gwas[gwas.P < 5e-8].SNP.values
hit_snps_idx = [id2idx['SNP'][i] for i in hit_snps]
if magma_path is not None:
# use magma genes and GSEA programs
print('Using MAGMA genes to filter...')
gwas_gene = pd.read_csv(magma_path, sep = '\s+')
id2gene = dict(pd.read_csv(os.path.join(data_path, 'misc_data/NCBI37.3.gene.loc'), sep = '\t', header = None)[[0,5]].values)
gwas_gene.loc[:,'GENE'] = gwas_gene['GENE'].apply(lambda x: id2gene[x])
import statsmodels.api as sm
p_values = gwas_gene['P']
corrected_p_values = sm.stats.multipletests(p_values, alpha=magma_threshold, method='bonferroni')[1]
gwas_gene.loc[:,'corrected_p_value'] = corrected_p_values
df_gene_hits = gwas_gene[gwas_gene['corrected_p_value'] < magma_threshold]
rnk = df_gene_hits[['GENE', 'ZSTAT']].set_index('GENE')
gene_hit_idx = [id2idx['Gene'][i] for i in df_gene_hits.GENE.values if i in id2idx['Gene']]
try:
gsea_results_BP = gp.prerank(rnk=rnk, gene_sets=gene_sets,
outdir=None, permutation_num=100,
min_size=2, max_size=1000, seed = 42)
gsea_results_BP = gsea_results_BP.res2d
go_hits = gsea_results_BP[gsea_results_BP['NOM p-val'] < program_threshold].Term.values
if len(go_hits) <= 5:
go_hits = gsea_results_BP.sort_values('NOM p-val')[:5].Term.values
go_hits_idx = [id2idx['BiologicalProcess'][x] for x in go_hits]
print('Using GSEA gene programs to filter...')
except:
print('No significant gene programs found...')
go_hits_idx = []
else:
# use all genes and gene programs
print('No filters... Using all genes and gene programs...')
gene_hit_idx = list(id2idx['Gene'].values())
go_hits_idx = list(id2idx['BiologicalProcess'].values())
snp2genes_hit = snp2genes[snp2genes.t_idx.isin(hit_snps_idx) & snp2genes.h_idx.isin(gene_hit_idx)]
rel2mean = snp2genes_hit.groupby('rel_type').weight.mean().reset_index().rename(columns = {'weight': 'rel_type_mean'})
rel2std = snp2genes_hit.groupby('rel_type').weight.agg(np.std).reset_index().rename(columns = {'weight': 'rel_type_std'})
snp2genes_hit = snp2genes_hit.merge(rel2std)
snp2genes_hit = snp2genes_hit.merge(rel2mean)
snp2genes_hit.loc[:,'z_rel'] = (snp2genes_hit['weight'] - snp2genes_hit['rel_type_mean'])/snp2genes_hit['rel_type_std']
v2g_hit = snp2genes_hit.groupby(['h_idx', 't_idx']).z_rel.max().reset_index().rename(columns={'z_rel': 'importance'})
v2g_hit_with_rel_type = pd.merge(v2g_hit, snp2genes_hit, left_on=['h_idx', 't_idx', 'importance'], right_on=['h_idx', 't_idx', 'z_rel'], how='left')
v2g_hit = v2g_hit_with_rel_type[['h_idx', 't_idx', 'importance', 'h_type', 't_type', 'rel_type']]
v2g_hit.loc[:,'rel_type'] = v2g_hit.rel_type.apply(lambda x: x[4:])
v2g_hit.loc[:,'Category'] = 'V2G'
v2g_hit.loc[:,'h_id'] = v2g_hit['h_idx'].apply(lambda x: idx2id['Gene'][x])
v2g_hit.loc[:,'t_id'] = v2g_hit['t_idx'].apply(lambda x: idx2id['SNP'][x])
gene2gene_hit = gene2gene[gene2gene.h_idx.isin(gene_hit_idx) & gene2gene.t_idx.isin(gene_hit_idx)]
rel2mean = gene2gene_hit.groupby('rel_type').weight.mean().reset_index().rename(columns = {'weight': 'rel_type_mean'})
rel2std = gene2gene_hit.groupby('rel_type').weight.agg(np.std).reset_index().rename(columns = {'weight': 'rel_type_std'})
gene2gene_hit = gene2gene_hit.merge(rel2std)
gene2gene_hit = gene2gene_hit.merge(rel2mean)
gene2gene_hit.loc[:,'z_rel'] = (gene2gene_hit['weight'] - gene2gene_hit['rel_type_mean'])/gene2gene_hit['rel_type_std']
g2g_hit = gene2gene_hit.groupby(['h_idx', 't_idx']).z_rel.max().reset_index().rename(columns={'z_rel': 'importance'})
g2g_hit_with_rel_type = pd.merge(g2g_hit, gene2gene_hit, left_on=['h_idx', 't_idx', 'importance'], right_on=['h_idx', 't_idx', 'z_rel'], how='left')
g2g_hit = g2g_hit_with_rel_type[['h_idx', 't_idx', 'importance', 'h_type', 't_type', 'rel_type']]
g2g_hit.loc[:,'rel_type'] = g2g_hit.rel_type.apply(lambda x: x.split('-')[1])
g2g_hit.loc[:,'Category'] = 'G2G'
g2g_hit.loc[:,'h_id'] = g2g_hit['h_idx'].apply(lambda x: idx2id['Gene'][x])
g2g_hit.loc[:,'t_id'] = g2g_hit['t_idx'].apply(lambda x: idx2id['Gene'][x])
gene2program_hit = gene2go[gene2go.t_idx.isin(gene_hit_idx) & gene2go.h_idx.isin(go_hits_idx)]
rel2mean = gene2program_hit.groupby('rel_type').weight.mean().reset_index().rename(columns = {'weight': 'rel_type_mean'})
rel2std = gene2program_hit.groupby('rel_type').weight.agg(np.std).reset_index().rename(columns = {'weight': 'rel_type_std'})
gene2program_hit = gene2program_hit.merge(rel2std)
gene2program_hit = gene2program_hit.merge(rel2mean)
gene2program_hit.loc[:,'z_rel'] = (gene2program_hit['weight'] - gene2program_hit['rel_type_mean'])/gene2program_hit['rel_type_std']
g2p_hit = gene2program_hit.groupby(['h_idx', 't_idx']).z_rel.max().reset_index().rename(columns={'z_rel': 'importance'})
g2p_hit_with_rel_type = pd.merge(g2p_hit, gene2program_hit, left_on=['h_idx', 't_idx', 'importance'], right_on=['h_idx', 't_idx', 'z_rel'], how='left')
g2p_hit = g2p_hit_with_rel_type[['h_idx', 't_idx', 'importance', 'h_type', 't_type', 'rel_type']]
g2p_hit.loc[:,'rel_type'] = g2p_hit.rel_type.apply(lambda x: x.split('-')[1])
g2p_hit.loc[:,'Category'] = 'G2P'
g2p_hit.loc[:,'h_id'] = g2p_hit['h_idx'].apply(lambda x: idx2id['BiologicalProcess'][x])
g2p_hit.loc[:,'t_id'] = g2p_hit['t_idx'].apply(lambda x: idx2id['Gene'][x])
g2p_hit.loc[:,'h_id'] = g2p_hit.h_id.apply(lambda x: go2name[x].capitalize() if x in go2name else x)
disease_critical_network = pd.concat((v2g_hit, g2g_hit, g2p_hit)).reset_index(drop = True)
print('Disease critical network finished generating...')
print('Generating variant interpretation networks...')
#### get for variant interpretation -> since we are looking at top K neighbors, we don't filter
# V2G
rel2mean = snp2genes_hit.groupby('rel_type').weight.mean().reset_index().rename(columns = {'weight': 'rel_type_mean'})
rel2std = snp2genes_hit.groupby('rel_type').weight.agg(np.std).reset_index().rename(columns = {'weight': 'rel_type_std'})
snp2genes = snp2genes.merge(rel2std)
snp2genes = snp2genes.merge(rel2mean)
snp2genes.loc[:,'z_rel'] = (snp2genes['weight'] - snp2genes['rel_type_mean'])/snp2genes['rel_type_std']
snp2genes = snp2genes.rename(columns={'z_rel': 'importance'})
v2g = snp2genes.groupby(['h_idx', 't_idx']).importance.max().reset_index()
v2g_with_rel_type = pd.merge(v2g, snp2genes, left_on=['h_idx', 't_idx', 'importance'], right_on=['h_idx', 't_idx', 'importance'], how='left')
v2g = v2g_with_rel_type[['h_idx', 't_idx', 'importance', 'h_type', 't_type', 'rel_type']]
v2g.loc[:,'h_id'] = v2g['h_idx'].apply(lambda x: idx2id['Gene'][x])
v2g.loc[:,'t_id'] = v2g['t_idx'].apply(lambda x: idx2id['SNP'][x])
## G2G
rel2mean = gene2gene_hit.groupby('rel_type').weight.mean().reset_index().rename(columns = {'weight': 'rel_type_mean'})
rel2std = gene2gene_hit.groupby('rel_type').weight.agg(np.std).reset_index().rename(columns = {'weight': 'rel_type_std'})
gene2gene = gene2gene.merge(rel2std)
gene2gene = gene2gene.merge(rel2mean)
gene2gene.loc[:,'z_rel'] = (gene2gene['weight'] - gene2gene['rel_type_mean'])/gene2gene['rel_type_std']
gene2gene = gene2gene.rename(columns={'z_rel': 'importance'})
g2g = gene2gene.groupby(['h_idx', 't_idx']).importance.max().reset_index()
g2g_with_rel_type = pd.merge(g2g, gene2gene, left_on=['h_idx', 't_idx', 'importance'], right_on=['h_idx', 't_idx', 'importance'], how='left')
g2g = g2g_with_rel_type[['h_idx', 't_idx', 'importance', 'h_type', 't_type', 'rel_type']]
g2g.loc[:,'h_id'] = g2g['h_idx'].apply(lambda x: idx2id['Gene'][x])
g2g.loc[:,'t_id'] = g2g['t_idx'].apply(lambda x: idx2id['Gene'][x])
g2g = g2g[g2g.h_idx != g2g.t_idx]
## G2P
rel2mean = gene2program_hit.groupby('rel_type').weight.mean().reset_index().rename(columns = {'weight': 'rel_type_mean'})
rel2std = gene2program_hit.groupby('rel_type').weight.agg(np.std).reset_index().rename(columns = {'weight': 'rel_type_std'})
gene2go = gene2go.merge(rel2std)
gene2go = gene2go.merge(rel2mean)
gene2go.loc[:,'z_rel'] = (gene2go['weight'] - gene2go['rel_type_mean'])/gene2go['rel_type_std']
gene2go = gene2go.rename(columns={'z_rel': 'importance'})
g2p = gene2go.groupby(['h_idx', 't_idx']).importance.max().reset_index()
g2p_with_rel_type = pd.merge(g2p, gene2go, left_on=['h_idx', 't_idx', 'importance'], right_on=['h_idx', 't_idx', 'importance'], how='left')
g2p = g2p_with_rel_type[['h_idx', 't_idx', 'importance', 'h_type', 't_type', 'rel_type']]
g2p.loc[:,'h_id'] = g2p['h_idx'].apply(lambda x: go2name[idx2id['BiologicalProcess'][x]].capitalize() if idx2id['BiologicalProcess'][x] in go2name else idx2id['BiologicalProcess'][x])
g2p.loc[:,'t_id'] = g2p['t_idx'].apply(lambda x: idx2id['Gene'][x])
## G2V
gene2snp = df_network[(df_network.h_type == 'SNP')
& (df_network.t_type == 'Gene')]
gene2snp_hit = gene2snp[gene2snp.h_idx.isin(hit_snps_idx) & gene2snp.t_idx.isin(gene_hit_idx)]
rel2mean = gene2snp_hit.groupby('rel_type').weight.mean().reset_index().rename(columns = {'weight': 'rel_type_mean'})
rel2std = gene2snp_hit.groupby('rel_type').weight.agg(np.std).reset_index().rename(columns = {'weight': 'rel_type_std'})
gene2snp = gene2snp.merge(rel2std)
gene2snp = gene2snp.merge(rel2mean)
gene2snp.loc[:,'z_rel'] = (gene2snp['weight'] - gene2snp['rel_type_mean'])/gene2snp['rel_type_std']
gene2snp = gene2snp.rename(columns={'z_rel': 'importance'})
g2v = gene2snp.groupby(['h_idx', 't_idx']).importance.max().reset_index()
g2v_with_rel_type = pd.merge(g2v, gene2snp, left_on=['h_idx', 't_idx', 'importance'], right_on=['h_idx', 't_idx', 'importance'], how='left')
g2v = g2v_with_rel_type[['h_idx', 't_idx', 'importance', 'h_type', 't_type', 'rel_type']]
g2v.loc[:,'h_id'] = g2v['h_idx'].apply(lambda x: idx2id['SNP'][x])
g2v.loc[:,'t_id'] = g2v['t_idx'].apply(lambda x: idx2id['Gene'][x])
print('Number of hit snps: ', len(hit_snps))
process_func = partial(get_local_interpretation, v2g=v2g, g2g=g2g, g2p=g2p, g2v=g2v, id2idx=id2idx, K_neighbors=K_neighbors)
with Pool(num_cpus) as p:
res = list(tqdm(p.imap(process_func, hit_snps), total=len(hit_snps)))
try:
df_variant_interpretation = pd.concat([i for i in res if i is not None])
except:
df_variant_interpretation = pd.DataFrame()
return df_variant_interpretation, disease_critical_network
Datasets
Models
