--- a
+++ b/kgwas/utils.py
@@ -0,0 +1,725 @@
+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
\ No newline at end of file