PPI_Graph_Functions.py

import Data_Processing as dp
import numpy as np
import networkx as nx
import random
import pandas as pd

from torch_geometric.nn import global_mean_pool
from torch_geometric.nn.pool import MemPooling
from torch_geometric.nn.norm import LayerNorm
from torch_geometric.loader import DataLoader
import torch_geometric.utils.convert as pyg_convert
import torch

from torch_geometric.nn import PDNConv, GATv2Conv
from torch_geometric.nn import GAE



class PPIGraphCollection:
    def __init__(self, pred_dict=None, zs_cutoff=None, relv_zs_cutoff=None, percentile=70, score_saver_percentile=99,
                 gene_dict_percentile=80, min_interactors=20, max_degree=300, esm2_dict=None,
                 del_raw_pred_dict=True, name_dict=None, id_len=15, gen_attrs=False, gs_edges=None, gs_nodes=None,
                 load_address=None, save_address=None, graph_load_mode=False, verbose=0, del_esm2_dict=True,
                 prebuild_neighbors=True, num_workers=0, convert_to_pyg_graph=True, strict_nodes=False,
                 rescue_prots=True):

        self.merged = False
        self.homo = True
        self.hetero = False
        self.verbose = verbose
        self.num_workers = num_workers
        self.convert_to_pyg_graph = convert_to_pyg_graph
        self.prebuild_neighbors = prebuild_neighbors
        self.min_interactors = min_interactors
        self.max_degree = max_degree
        self.strict_nodes = strict_nodes
        self.rescue_prots = rescue_prots

        if load_address == None:
            if pred_dict == None or esm2_dict == None:
                print('Error: Please provide pred_dict and esm2_dict')
                return

            self.raw_pred_dict = pred_dict
            self.zs_cutoff = zs_cutoff
            self.relv_zs_cutoff = relv_zs_cutoff
            self.percentile = percentile
            self.score_saver_percentile = score_saver_percentile
            self.gene_dict_percentile = gene_dict_percentile
            self.esm2_dict = esm2_dict
            self.name_dict = name_dict
            self.id_len = id_len
            self.gen_attrs = gen_attrs
            self.gs_edges = gs_edges
            self.gs_nodes = gs_nodes


            self.num_node_features = None
            self.num_edge_features = None

            self.build_master_pred_dict()
            self.build_master_graph()
            self.reset_node_order()
            if del_esm2_dict:
                del self.esm2_dict

            if del_raw_pred_dict:
                del self.raw_pred_dict

            if save_address != None:
                self.save_pred_dict(save_address)

        else:
            self.esm2_dict = esm2_dict
            self.name_dict = name_dict
            self.max_degree = max_degree
            self.id_len = id_len
            self.gen_attrs = gen_attrs
            self.gs_edges = gs_edges
            self.gs_nodes = gs_nodes

            self.num_node_features = None
            self.num_edge_features = None

            if graph_load_mode:
                self.load_graph(load_address)
                self.add_node_features()
                self.reset_node_order()
                if del_esm2_dict:
                    del self.esm2_dict
            else:
                self.load_pred_dict(load_address)
                self.build_master_graph()
                self.reset_node_order()

        if self.prebuild_neighbors:
            self.build_neighbors_dict()

        if convert_to_pyg_graph:
            self.convert_master_graph()

    def build_zs_pred_dict(self, pred_dict):
        zs_pred_dict = {}

        mean = np.nanmean(list(pred_dict.values()))
        std = np.nanstd(list(pred_dict.values()))

        for pair, val in pred_dict.items():
            zs_pred_dict[pair] = (val - mean) / std

        return zs_pred_dict

    def zs_filter_pred_dict(self, pred_dict):
        filtered_pred_dict = {}

        mean = np.nanmean(list(pred_dict.values()))
        std = np.nanstd(list(pred_dict.values()))

        for pair, val in pred_dict.items():
            zs = (val - mean) / std
            if zs >= self.zs_cutoff:
                filtered_pred_dict[pair] = val

        return filtered_pred_dict

    def relv_zs_pred_dict_filter(self, pred_dict):
        filtered_pred_dict = {}

        gene_dict = dp.create_gene_interaction_dict(pred_dict)
        for prot, sub_dict in gene_dict.items():
            genes = sub_dict['genes']
            values = sub_dict['values']
            zs_values = sub_dict['z_scores']
            for i in range(len(genes)):
                if zs_values[i] >= self.relv_zs_cutoff:
                    gene = genes[i]
                    if (prot, gene) not in filtered_pred_dict and (gene, prot) not in filtered_pred_dict:
                        filtered_pred_dict[prot, gene] = values[i]

        return filtered_pred_dict

    def build_master_pred_dict(self):
        ori_prot_set = set()
        for pair in self.raw_pred_dict.keys():
            ori_prot_set.add(pair[0])
            ori_prot_set.add(pair[1])
       # Filter Preds
        cut_pred_dict = {}
        if self.percentile == None:
            zs_filter_pred_dict = self.zs_filter_pred_dict(self.raw_pred_dict)
            relv_zs_filter_pred_dict = self.relv_zs_pred_dict_filter(self.raw_pred_dict)

            # prot_set = set()
            for pair, val in zs_filter_pred_dict.items():
                if pair in relv_zs_filter_pred_dict:
                    cut_pred_dict[pair] = val
                elif (pair[1], pair[0]) in relv_zs_filter_pred_dict:
                    cut_pred_dict[pair[1], pair[0]] = val

            del zs_filter_pred_dict, relv_zs_filter_pred_dict

        else:
            score_cutoff = np.percentile(list(self.raw_pred_dict.values()), self.percentile)
            score_saver_cutoff = np.percentile(list(self.raw_pred_dict.values()), self.score_saver_percentile)
            for pair, val in self.raw_pred_dict.items():
                if val >= score_cutoff:
                    cut_pred_dict[pair] = val


            gene_cut_pred_dict = {}
            gene_dict = dp.create_gene_interaction_dict(cut_pred_dict)
            for prot, sub_dict in gene_dict.items():
                genes = sub_dict['genes']
                values = sub_dict['values']

                relv_score_cutoff = np.percentile(values, self.gene_dict_percentile)

                gene_vals = list(zip(genes, values))
                gene_vals.sort(reverse=True, key=lambda x: x[1])

                cut_gene_vals = [tup for tup in gene_vals if
                                 tup[1] >= relv_score_cutoff or tup[1] >= score_saver_cutoff]
                if len(cut_gene_vals) < self.min_interactors:
                    if len(cut_gene_vals) < self.min_interactors:
                        cut_gene_vals = gene_vals
                    else:
                        cut_gene_vals = gene_vals[:self.min_interactors]

                for tup in cut_gene_vals:
                    gene = tup[0]
                    val = tup[1]
                    if (prot, gene) not in gene_cut_pred_dict and (gene, prot) not in gene_cut_pred_dict:
                        gene_cut_pred_dict[prot, gene] = val
            cut_pred_dict = gene_cut_pred_dict

        # Build all values pred dict (S, ZS, RZS1, RZS2)
        zs_pred_dict = self.build_zs_pred_dict(self.raw_pred_dict)
        pred_dict = {}
        prot_degree_dict = {}
        for pair, val in cut_pred_dict.items():
            if pair in zs_pred_dict:
                pred_dict[pair] = [val, zs_pred_dict[pair]]
            else:
                pred_dict[pair[1], pair[0]] = [val, zs_pred_dict[pair[1], pair[0]]]

            if pair[0] not in prot_degree_dict:
                prot_degree_dict[pair[0]] = 0
            if pair[1] not in prot_degree_dict:
                prot_degree_dict[pair[1]] = 0
            prot_degree_dict[pair[0]] += 1
            prot_degree_dict[pair[1]] += 1

        if self.rescue_prots:
            print('RESCUE')
            rescue_prots_set = set()
            for prot in ori_prot_set:
                if prot not in prot_degree_dict:
                    rescue_prots_set.add(prot)
                else:
                    if prot_degree_dict[prot] < self.min_interactors:
                        rescue_prots_set.add(prot)

            print('len resuce set',len(rescue_prots_set))
            if len(rescue_prots_set) > 0:
                rescue_prot_dict = {}
                for pair,val in self.raw_pred_dict.items():
                    if pair[0] in rescue_prots_set:
                        if pair[0] not in rescue_prot_dict:
                            rescue_prot_dict[pair[0]] = []
                        rescue_prot_dict[pair[0]].append((pair,val))

                    if pair[1] in rescue_prots_set:
                        if pair[1] not in rescue_prot_dict:
                            rescue_prot_dict[pair[1]] = []
                        rescue_prot_dict[pair[1]].append((pair, val))

                for prot, pair_val_tup_list in rescue_prot_dict.items():
                    pair_val_tup_list.sort(reverse=True, key=lambda x:x[1])
                    relv_tups = pair_val_tup_list[0:self.min_interactors]
                    for tup in relv_tups:
                        pair = tup[0]
                        val = tup[1]

                        if pair in pred_dict or (pair[1],pair[0]) in pred_dict:
                            continue

                        if pair in zs_pred_dict:
                            pred_dict[pair] = [val, zs_pred_dict[pair]]
                            cut_pred_dict[pair] = val
                        else:
                            pred_dict[pair[1], pair[0]] = [val, zs_pred_dict[pair[1], pair[0]]]
                            cut_pred_dict[pair[1], pair[0]] = val

        del zs_pred_dict

        gene_dict = dp.create_gene_interaction_dict(cut_pred_dict)
        for prot, sub_dict in gene_dict.items():
            genes = sub_dict['genes']
            zs_values = sub_dict['z_scores']
            for i in range(len(genes)):
                gene = genes[i]
                if (prot, gene) in pred_dict:
                    pred_dict[prot, gene].append(zs_values[i])
                else:
                    pred_dict[gene, prot].append(zs_values[i])

        self.pred_dict = pred_dict

    def build_master_graph(self):

        self.master_graph = nx.from_edgelist(list(self.pred_dict.keys()))

        if self.strict_nodes:
            del_nodes = []
            for node in self.master_graph.nodes:
                if type(node) == str:
                    prot = node.split('-')[0].split(';')[0].split('.0')[0]
                else:
                    prot = node
                if prot not in self.esm2_dict:
                    del_nodes.append(node)
            self.master_graph.remove_nodes_from(del_nodes)

        # Assign features to edges
        for pair, val_list in self.pred_dict.items():

            if self.strict_nodes:
                if pair[0] not in self.master_graph.nodes or pair[1] not in self.master_graph.nodes:
                    continue
            self.master_graph[pair[0]][pair[1]]['edge_x'] = np.array(val_list[1:]).astype('float64')
            self.master_graph[pair[0]][pair[1]]['edge_y'] = 1

            if self.num_edge_features == None:
                self.num_edge_features = len(val_list)


        self.add_node_features()

        if self.max_degree != None:
            # Drop Edges for nodes that have too many edges
            nx_degree_dict = self.master_graph.degree
            degree_dict = {}
            for prot, degree in nx_degree_dict:
                degree_dict[prot] = degree

            for prot, degree in degree_dict.items():
                if degree > self.max_degree:
                    current_degree = degree
                    edges = self.master_graph.edges(prot)
                    edge_list = []
                    for edge in edges:
                        rel_zs_vals = self.master_graph.edges[edge]['edge_x'][1:]
                        edge_list.append((edge, rel_zs_vals))

                    edge_list.sort(reverse=True, key=lambda x: np.nanmax(x[1]))

                    del_edge_list = []
                    while current_degree > self.max_degree:
                        edge = edge_list.pop()[0]
                        del_edge_list.append(edge)
                        current_degree -= 1
                        if edge[0] == prot:
                            degree_dict[edge[1]] -= 1
                        else:
                            degree_dict[edge[0]] -= 1

                    self.master_graph.remove_edges_from(del_edge_list)

        self.master_nodes = list(self.master_graph.nodes)



    def build_neighbors_dict(self):
        if self.merged:
            graph = self.merged_master_graph
        else:
            graph = self.master_graph

        self.neighbor_dict = {}
        prots = set(graph.nodes())

        for prot in prots:
            self.neighbor_dict[prot] = list(nx.all_neighbors(graph, prot))

        # self.master_nodes = list(self.neighbor_dict.keys())

    def convert_master_graph(self):
        self.convert_to_pyg_graph = True
        if self.merged:
            self.node_dict = {}
            for i,node in enumerate(self.merged_master_graph.nodes()):
                self.node_dict[node] = i

            self.merged_master_graph = pyg_convert.from_networkx(self.merged_master_graph)
        else:
            self.node_dict = {}
            for i, node in enumerate(self.master_graph.nodes()):
                self.node_dict[node] = i

            self.master_graph = pyg_convert.from_networkx(self.master_graph)

    def add_node_features(self):
        if self.merged:
            graph = self.merged_master_graph
        else:
            graph = self.master_graph

        for node in graph.nodes():
            if type(node) == str:
                prot = node.split('-')[0].split(';')[0].split('.0')[0]
            else:
                prot = node
            if prot in self.esm2_dict:
                esm_matrix = self.esm2_dict[prot]
            else:
                if self.verbose > 0:
                    print('Missing Protein:', node, prot)
                esm_matrix = np.zeros(np.shape(self.esm2_dict[list(self.esm2_dict.keys())[0]]))

            if self.num_node_features == None:
                self.num_node_features = len(esm_matrix)

            graph.nodes[node]['node_x'] = esm_matrix.astype('float64')
            graph.nodes[node]['node_y'] = 1

        if self.merged:
            self.merged_master_graph = graph
        else:
            self.master_graph = graph

    def reset_node_order(self, randomize=True):

        if self.merged:
            self.node_order = []
            if self.homo:
                self.node_order += list(self.homo_merge_set)
            if self.hetero:
                self.node_order += list(self.hetero_merge_set)

            # print('node_order_list initial', len(self.node_order))
        else:
            self.node_order = self.master_nodes

        if self.merged:
            if self.merged_gs_dict:
                # print('nodes', self.node_order)
                # print('gs_dict', self.merged_gs_dict)
                self.node_order = [node_pair for node_pair in self.node_order if node_pair in self.merged_gs_dict
                                   or (node_pair[1], node_pair[0]) in self.merged_gs_dict]
                # print('node_order_list after gs_dict', len(self.node_order))

        else:
            if self.gs_nodes != None:
                self.node_order = [node for node in self.node_order if node in self.gs_nodes]

        if randomize:
            random.shuffle(self.node_order)

    def gen_batch(self, batch_size, random=True, convert_to_tensor=True, generate_neg_edges=True,
                  prot_batch_list=None, pred_mode=False, include_center_prot=True,
                  adaptive_batching=False, low_batch_size_cutoff=300, medium_batch_size_cutoff=600):
        return self.gen_converted_batch_helper(batch_size, random=random, convert_to_tensor=convert_to_tensor,
                                            generate_neg_edges=generate_neg_edges, prot_batch_list=prot_batch_list,
                                            pred_mode=pred_mode, include_center_prot=include_center_prot,
                                            adaptive_batching=adaptive_batching,
                                            low_batch_size_cutoff=low_batch_size_cutoff,
                                            medium_batch_size_cutoff=medium_batch_size_cutoff
                                                )


    def gen_converted_batch_helper(self, batch_size, random=True, convert_to_tensor=True, generate_neg_edges=False,
                                prot_batch_list=None, pred_mode=False, include_center_prot=True,
                                adaptive_batching=False, low_batch_size_cutoff=300, medium_batch_size_cutoff=600):

        x_batch = []
        y_batch = []
        node_pair_list = []
        prot_list = []
        if prot_batch_list == None:
            num_prots = batch_size
        else:
            num_prots = len(prot_batch_list)

        if adaptive_batching:
            small_batch_size = batch_size
            medium_batch_size = int(batch_size/2)
            large_batch_size = 1
            small_batch_data_dict = {'x_batch':[],'prot_list':[]}
            medium_batch_data_dict = {'x_batch':[],'prot_list':[]}
            large_batch_data_dict = {'x_batch':[],'prot_list':[]}

        for i in range(num_prots):
            if self.merged:
                node_pair = self.node_order.pop()
                node_pair_list.append(node_pair)


                if self.merged_gs_dict != None:
                    if node_pair in self.merged_gs_dict:
                        y_batch.append(self.merged_gs_dict[node_pair])
                    else:
                        y_batch.append(self.merged_gs_dict[node_pair[1], node_pair[0]])

                prot = node_pair[0]
                other_prot = node_pair[1]
                neighbors = self.neighbor_dict[prot] + [prot]
                if prot != other_prot:
                    neighbors = self.neighbor_dict[prot] + [other_prot]
                neighbors = set(neighbors)

                if pred_mode:
                    if prot == other_prot:
                        prot_list.append(prot)
                    else:
                        prot_list.append((prot,other_prot))

                node_subset = torch.tensor([self.node_dict[node] for node in neighbors])
                prot_graph = self.merged_master_graph.subgraph(node_subset)
            else:
                if prot_batch_list == None:
                    prot = self.node_order.pop()
                else:
                    prot = prot_batch_list.pop()

                if include_center_prot:
                    neighbors = self.neighbor_dict[prot] + [prot]
                else:
                    neighbors = self.neighbor_dict[prot]

                if adaptive_batching:
                    if len(neighbors) <= low_batch_size_cutoff:
                        batch_mode = 's'
                    elif len(neighbors) <= medium_batch_size_cutoff:
                        batch_mode = 'm'
                    else:
                        batch_mode = 'l'

                if adaptive_batching:
                    if batch_mode == 's':
                        small_batch_data_dict['prot_list'].append(prot)
                    elif batch_mode == 'm':
                        medium_batch_data_dict['prot_list'].append(prot)
                    elif batch_mode == 'l':
                        large_batch_data_dict['prot_list'].append(prot)
                else:
                    prot_list.append(prot)

                node_subset = torch.tensor([self.node_dict[node] for node in neighbors])
                prot_graph = self.master_graph.subgraph(node_subset)

            if generate_neg_edges:
                prot_graph = self.generate_neg_edges(prot_graph)

            if adaptive_batching:
                if batch_mode == 's':
                    small_batch_data_dict['x_batch'].append(prot_graph)
                elif batch_mode == 'm':
                    medium_batch_data_dict['x_batch'].append(prot_graph)
                elif batch_mode == 'l':
                    large_batch_data_dict['x_batch'].append(prot_graph)
            else:
                x_batch.append(prot_graph)



            if len(self.node_order) == 0:
                self.reset_node_order(randomize=random)

        if convert_to_tensor:
            if adaptive_batching:
                if batch_mode == 's':
                    small_batch_data_dict['x_batch'] = DataLoader(x_batch, batch_size=small_batch_size,
                                                                  num_workers=self.num_workers, pin_memory=True)
                elif batch_mode == 'm':
                    medium_batch_data_dict['x_batch'] = DataLoader(x_batch, batch_size=medium_batch_size,
                                                                  num_workers=self.num_workers, pin_memory=True)
                elif batch_mode == 'l':
                    large_batch_data_dict['x_batch'] = DataLoader(x_batch, batch_size=large_batch_size,
                                                                  num_workers=self.num_workers, pin_memory=True)
            else:
                x_batch = DataLoader(x_batch, batch_size=batch_size, num_workers=self.num_workers, pin_memory=True)


        if self.merged:
            if pred_mode:
                return node_pair_list, x_batch
            else:
                if convert_to_tensor:
                    y_batch = torch.from_numpy(np.reshape(y_batch, (len(y_batch), 1))).float()
                return x_batch, y_batch
        else:
            if pred_mode:
                if adaptive_batching:
                    prot_list_tup = (small_batch_data_dict['prot_list'],medium_batch_data_dict['prot_list'],
                                   large_batch_data_dict['prot_list'])
                    x_batch_tup = (small_batch_data_dict['x_batch'],medium_batch_data_dict['x_batch'],
                                   large_batch_data_dict['x_batch'])
                    return prot_list_tup, x_batch_tup
                else:
                    return prot_list, x_batch
            else:
                if adaptive_batching:
                    return (small_batch_data_dict['x_batch'],medium_batch_data_dict['x_batch'],
                            large_batch_data_dict['x_batch'])
                else:
                    return x_batch



    def save_pred_dict(self, save_address):
        prot_1_list = []
        prot_2_list = []
        val_dict_list = {}

        for i in range(len(list(self.pred_dict.values())[0])):
            val_dict_list[i] = []

        for pair, val_list in self.pred_dict.items():
            prot_1_list.append(pair[0])
            prot_2_list.append(pair[1])
            for i, val in enumerate(val_list):
                val_dict_list[i].append(val)

        table_dict = {'Protein 1': prot_1_list, 'Protein 2': prot_2_list}
        for i, val_list in val_dict_list.items():
            table_dict[i] = val_list

        table = pd.DataFrame(table_dict)
        table.to_csv(save_address, index=None)

    def load_pred_dict(self, load_address):
        self.pred_dict = {}
        table = pd.read_csv(load_address)
        val_cols = [col for col in table.columns if 'Protein' not in col]
        for row in table.iloc:
            prot_1 = str(row['Protein 1'])
            prot_2 = str(row['Protein 2'])
            val_list = []
            for col in val_cols:
                val_list.append(row[col])

            self.pred_dict[prot_1, prot_2] = val_list

    def strip_graph_node_features(self,graph):
        for node in graph.nodes:
            try:
                del graph.nodes[node]['node_x']
            except:
                print(node,graph.nodes[node]['node_x'])
        return graph

    def save_graph(self, save_address):
        if self.merged:
            self.strip_graph_node_features(self.merged_master_graph)
            save_package = (
                self.merged_master_graph,
                self.merged,
                self.merged_gs_dict,
                self.homo_merge_set,
                self.hetero_merge_set,
            )
            dp.save_object(save_package, save_address)
        else:
            self.strip_graph_node_features(self.master_graph)
            dp.save_object(self.master_graph, save_address)

    def load_graph(self, load_address):
        save_package = dp.load_object(load_address)
        if type(save_package) == tuple:
            self.merged_master_graph = save_package[0]
            self.merged = save_package[1]
            self.merged_gs_dict = save_package[2]
            self.homo_merge_set = save_package[3]
            self.hetero_merge_set = save_package[4]
        else:
            self.master_graph = save_package


class NetworkEmbedder(torch.nn.Module):
    def __init__(self, in_channels=5120, hidden_channels=4096, edge_dim=3, num_layers=3, num_blocks=3,
                 conv_type='pdn', num_heads=2, hidden_layer_channel=64,
                 dropout=0.3, activation=torch.nn.LeakyReLU()):
        super().__init__()

        self.num_blocks = num_blocks
        self.activation = activation

        self.block1 = Conv_Block(in_channels, hidden_channels, edge_dim, num_layers, hidden_channels,
                                 conv_type=conv_type,
                                 num_heads=num_heads, hidden_layer_channel=hidden_layer_channel, dropout=dropout,
                                 activation=activation)
        self.block2 = Conv_Block(hidden_channels, hidden_channels, edge_dim, num_layers, hidden_channels,
                                 conv_type=conv_type,
                                 num_heads=num_heads, hidden_layer_channel=hidden_layer_channel, dropout=dropout,
                                 activation=activation)
        self.block3 = Conv_Block(hidden_channels, int(np.ceil(hidden_channels / 2)), edge_dim, num_layers,
                                 int(np.ceil(hidden_channels / 2)), conv_type=conv_type,
                                 num_heads=num_heads, hidden_layer_channel=hidden_layer_channel,
                                 dropout=dropout * 1.25,
                                 activation=activation)

        self.linear1 = torch.nn.Linear(hidden_channels, hidden_channels // 2)

    def forward(self, x, edge_index, edge_attr):
        x1 = self.block1(x, edge_index, edge_attr)
        x2 = self.block2(x1, edge_index, edge_attr)
        x2 = x2 + x1

        x3 = self.block3(x2, edge_index, edge_attr)
        x3 = x3 + self.linear1(x2)
        x = x3

        return x



class Protea(torch.nn.Module):
    def __init__(self, in_channels=5120, hidden_channels=4096, edge_dim=3, num_layers=3, num_blocks=3,
                 conv_type='pdn', num_heads=2, hidden_layer_channel=64,
                 dropout=0.3, linear_dropout=0.5, activation=torch.nn.LeakyReLU(),
                 num_clusters=4, num_mem_heads=4,
                 networkembedder_address=None, freeze_embedder_weights=True, device='cpu',
                 mempool=True, pred_mode=False):
        super().__init__()

        self.pred_mode = pred_mode
        self.num_blocks = num_blocks
        self.activation = activation
        self.mempool = mempool

        self.embedder = GAE(NetworkEmbedder(in_channels, hidden_channels, edge_dim, num_layers, num_blocks,
                                            conv_type=conv_type, num_heads=num_heads,
                                            hidden_layer_channel=hidden_layer_channel,
                                            dropout=dropout, activation=activation))

        if networkembedder_address != None:
            self.embedder.load_state_dict(torch.load(networkembedder_address, weights_only=True,
                                                     map_location=torch.device(device)))

        self.embedder = self.embedder.encoder

        if networkembedder_address != None:
            if freeze_embedder_weights:
                for param in self.embedder.parameters():
                    param.requires_grad = False
            else:
                for param in self.embedder.parameters():
                    param.requires_grad = True

        hidden_division = 2

        self.mempooling_1 = MemPooling(int(np.ceil(hidden_channels / hidden_division)),
                                       int(np.ceil(hidden_channels / hidden_division)),
                                       heads=num_mem_heads,
                                       num_clusters=num_clusters,
                                       tau=1.0)
        self.lnorm1 = LayerNorm(int(np.ceil(hidden_channels / hidden_division)))

        self.mempooling_2 = MemPooling(int(np.ceil(hidden_channels / hidden_division)),
                                       int(np.ceil(hidden_channels / hidden_division)),
                                       heads=num_mem_heads,
                                       num_clusters=1,
                                       tau=1.0)
        self.lnorm2 = LayerNorm(int(np.ceil(hidden_channels / hidden_division)))

        self.dropout_mem = torch.nn.Dropout(dropout)

        self.dropout1 = torch.nn.Dropout(linear_dropout)
        self.head_1 = torch.nn.Linear(int(np.ceil(hidden_channels / hidden_division)),
                                      int(np.ceil(hidden_channels / hidden_division)))
        self.dropout2 = torch.nn.Dropout(linear_dropout)
        self.head_2 = torch.nn.Linear(int(np.ceil(hidden_channels / hidden_division)), 1)

    def forward(self, x1, edge_index_1, edge_attr_1, batch_idx_1, x2, edge_index_2, edge_attr_2, batch_idx_2):

        # Get Embeddings for the Graphs
        x1 = self.embedder(x1, edge_index_1, edge_attr_1)
        x2 = self.embedder(x2, edge_index_2, edge_attr_2)

        # Global Pooling on Graphs and Combining the Embeddings
        x1, s1 = self.mempooling_1(x1, batch_idx_1)  # global pooling
        x1 = self.activation(x1)
        x1 = self.dropout_mem(x1)
        x1 = self.lnorm1(x1)

        x2, s2 = self.mempooling_1(x2, batch_idx_2)  # global pooling
        x2 = self.activation(x2)
        x2 = self.dropout_mem(x2)
        x2 = self.lnorm1(x2)

        x = torch.cat((x1, x2), dim=1)

        x = self.activation(x)
        x = self.dropout_mem(x)

        x, s = self.mempooling_2(x)
        x = self.lnorm2(x)

        kl_loss = MemPooling.kl_loss(s1) + MemPooling.kl_loss(s2) + MemPooling.kl_loss(s)

        x = x.squeeze()

        x = self.dropout1(x)
        x = self.head_1(x)
        x = self.activation(x)
        x = self.dropout2(x)
        x = self.head_2(x)

        if self.pred_mode:
            return x

        return x, kl_loss



class Conv_Block(torch.nn.Module):
    def __init__(self, in_channels, conv_hidden_channels, edge_dim, num_layers, out_channels, conv_type='pdn',
                 num_heads=2, hidden_layer_channel=64, dropout=0.3, activation=torch.nn.LeakyReLU()):
        super().__init__()

        self.activation = activation
        self.num_layers = num_layers

        self.conv1 = PDNConv(in_channels, conv_hidden_channels, edge_dim=edge_dim,
                             hidden_channels=hidden_layer_channel)
        self.conv2 = PDNConv(conv_hidden_channels, conv_hidden_channels, edge_dim=edge_dim,
                             hidden_channels=hidden_layer_channel)
        self.conv3 = PDNConv(conv_hidden_channels, out_channels, edge_dim=edge_dim,
                             hidden_channels=hidden_layer_channel)

        self.dropout = torch.nn.Dropout(dropout)

    def forward(self, x, edge_index, edge_attr):
        x = self.conv1(x, edge_index, edge_attr=edge_attr)
        x = self.activation(x)
        x = self.dropout(x)

        x = self.conv2(x, edge_index, edge_attr=edge_attr)
        x = self.activation(x)
        x = self.dropout(x)

        x = self.conv3(x, edge_index, edge_attr=edge_attr)
        x = self.activation(x)
        x = self.dropout(x)
        return x



def protea_predictor(model,device,ppi_graph_1,ppi_graph_2,include_center_prot,
                     save_address, hetero_prots={}, batch_size=1):
    model.eval()

    nodes_1 = ppi_graph_1.master_nodes
    nodes_2 = ppi_graph_2.master_nodes

    if len(hetero_prots) != 0:
        relv_nodes_1 = [node for node in nodes_1 if node in hetero_prots]
        relv_nodes_2 = [node for node in nodes_2 if node in hetero_prots]

        num_batchs_1 = int(np.ceil(len(relv_nodes_1) / batch_size))
        num_batchs_2 = int(np.ceil(len(relv_nodes_2) / batch_size))

        prot_list_1, loader_1 = ppi_graph_1.gen_batch(batch_size, random=False, convert_to_tensor=True,
                                                      generate_neg_edges=False,
                                                      prot_batch_list=relv_nodes_1.copy(),
                                                      pred_mode=True,
                                                      include_center_prot=include_center_prot)
        loader_1_iter = iter(loader_1)
        prot_list_2, loader_2 = ppi_graph_2.gen_batch(batch_size, random=False, convert_to_tensor=True,
                                                      generate_neg_edges=False,
                                                      prot_batch_list=relv_nodes_2.copy(),
                                                      pred_mode=True,
                                                      include_center_prot=include_center_prot)

        pred_dict = {}
        for i in range(num_batchs_1):
            loader_2_iter = iter(loader_2)
            try:
                ori_batch_1 = next(loader_1_iter)

            except:
                break
            for j in range(num_batchs_2):
                try:
                    batch_2 = next(loader_2_iter)
                except:
                    break
                batch_1 = ori_batch_1.detach().clone()
                print('Num batch',j,len(loader_2_iter))
                print('In Batch')
                print(batch_1, batch_2)
                batch_1 = batch_1.to(device, non_blocking=False)
                batch_2 = batch_2.to(device, non_blocking=False)
                relv_prot_list_1 = prot_list_1[i * batch_size:(i + 1) * batch_size]
                relv_prot_list_2 = prot_list_2[j * batch_size:(j + 1) * batch_size]

                preds = prediction_helper(batch_1, batch_2, model)
                del batch_1
                for k in range(len(preds)):
                    if len(preds) > 1:
                        pred_dict[(relv_prot_list_1[k],relv_prot_list_2[k])] = preds[k][0]
                    else:
                        pred_dict[(relv_prot_list_1[k],relv_prot_list_2[k])] = preds[k]

    else:
        relv_nodes = [node for node in nodes_1 if node in nodes_2]
        num_batchs = int(np.ceil(len(relv_nodes) / batch_size))

        prot_list_1, loader_1 = ppi_graph_1.gen_batch(batch_size, random=False, convert_to_tensor=True,
                                                                    generate_neg_edges=False,
                                                                    prot_batch_list=relv_nodes.copy(),
                                                                    pred_mode=True,
                                                                    include_center_prot=include_center_prot)
        loader_1_iter = iter(loader_1)
        prot_list_2, loader_2 = ppi_graph_2.gen_batch(batch_size, random=False, convert_to_tensor=True,
                                         generate_neg_edges=False,
                                         prot_batch_list=relv_nodes.copy(),
                                         pred_mode=True,
                                         include_center_prot=include_center_prot)
        loader_2_iter = iter(loader_2)

        pred_dict = {}
        for i in range(num_batchs):
            try:
                batch_1 = next(loader_1_iter)
                batch_2 = next(loader_2_iter)
            except:
                break

            batch_1 = batch_1.to(device, non_blocking=True)
            batch_2 = batch_2.to(device, non_blocking=True)
            prot_list = prot_list_1[i*batch_size:(i+1)*batch_size]

            preds = prediction_helper(batch_1, batch_2, model)
            for j in range(len(preds)):
                if len(preds) > 1:
                    pred_dict[prot_list[j]] = preds[j][0]
                else:
                    pred_dict[prot_list[j]] = preds[j]

    save_table_dict = {
        'Protein 1': [],
        'Protein 2': [],
        'Score':[]
    }
    for key, val in pred_dict.items():
        if type(key) == tuple:
            save_table_dict['Protein 1'].append(key[0])
            save_table_dict['Protein 2'].append(key[1])
        else:
            save_table_dict['Protein 1'].append(key)
            save_table_dict['Protein 2'].append('')

        save_table_dict['Score'].append(val)

    save_table = pd.DataFrame(save_table_dict)
    save_table.to_csv(save_address, index=False)


@torch.no_grad()
def prediction_helper(x_data_1, x_data_2, model):

    x_1 = x_data_1.node_x.float()
    edge_x_1 = x_data_1.edge_x.float()
    edge_index_1 = x_data_1.edge_index
    batch_index_1 = x_data_1.batch

    x_2 = x_data_2.node_x.float()
    edge_x_2 = x_data_2.edge_x.float()
    edge_index_2 = x_data_2.edge_index
    batch_index_2 = x_data_2.batch

    z = model(x_1, edge_index_1, edge_x_1, batch_index_1,
                       x_2, edge_index_2, edge_x_2, batch_index_2)

    z = torch.nn.Sigmoid()(z)
    z = z.detach().cpu().numpy()

    return z