t5g/network.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.nn import Sequential, Linear, ReLU
from torch_geometric.nn import GraphConv, GINConv, GATConv, SAGEConv, GPSConv, GINEConv, GATv2Conv
from torch_geometric.nn import global_mean_pool, global_add_pool, global_max_pool, GlobalAttention, Set2Set, MulAggregation
from torch_geometric.nn import aggr
from torch_geometric.nn import GraphNorm


import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleSelfAttention(nn.Module):
    def __init__(self, embedding_dim, num_heads=1):
        super(SimpleSelfAttention, self).__init__()
        self.embedding_dim = embedding_dim
        self.num_heads = num_heads

        # Assuming key_channels = value_channels = embedding_dim
        self.key_channels = self.embedding_dim
        self.value_channels = self.embedding_dim

        # Linear layers for queries, keys, and values
        self.query = nn.Linear(embedding_dim, self.key_channels * num_heads)
        self.key = nn.Linear(embedding_dim, self.key_channels * num_heads)
        self.value = nn.Linear(embedding_dim, self.value_channels * num_heads)

        # Output projection layer
        self.proj = nn.Linear(self.value_channels * num_heads, embedding_dim)

        # Scaling for dot-product attention
        self.scale = nn.Parameter(torch.sqrt(torch.FloatTensor([self.key_channels // num_heads])))

    def forward(self, x1, x2, x3):
        # x1, x2, x3 shapes: [Batch_size, Embedding_dim]
        # Stack the inputs along a new dimension (sequence dimension)
        batch_size = x1.shape[0]
        x = torch.stack((x1, x2, x3), dim=1)  # [Batch_size, 3, Embedding_dim]

        # Compute queries, keys, values for all three inputs
        Q = self.query(x)  # [Batch_size, 3, num_heads * embedding_dim]
        K = self.key(x)    # [Batch_size, 3, num_heads * embedding_dim]
        V = self.value(x)  # [Batch_size, 3, num_heads * embedding_dim]

        # Reshape for multi-head attention
        Q = Q.view(batch_size, -1, self.num_heads, self.embedding_dim).transpose(1, 2)  # [Batch_size, num_heads, 3, embedding_dim]
        K = K.view(batch_size, -1, self.num_heads, self.embedding_dim).transpose(1, 2)  # [Batch_size, num_heads, 3, embedding_dim]
        V = V.view(batch_size, -1, self.num_heads, self.embedding_dim).transpose(1, 2)  # [Batch_size, num_heads, 3, embedding_dim]

        # Calculate dot product attention
        attention_scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
        attention = F.softmax(attention_scores, dim=-1)

        # Apply attention to V
        x = torch.matmul(attention, V)  # [Batch_size, num_heads, 3, embedding_dim]

        # Concatenate heads and put through final linear layer
        x = x.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.embedding_dim)
        x = self.proj(x)  # [Batch_size, 3, embedding_dim]

        # Sum the outputs from the three inputs
        out = x.sum(dim=1)  # [Batch_size, embedding_dim]
        return out

def cosine_similarity(x,y):
    num = x.dot(y.T)
    denom = np.linalg.norm(x) * np.linalg.norm(y)
    return num / denom

class LogCoshLoss(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, y_t, y_prime_t):
        ey_t = y_t - y_prime_t
        return torch.mean(torch.log(torch.cosh(ey_t + 1e-12)))


class WeightedMSELoss(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, y, y_t, weights=None):
        loss = (y - y_t) ** 2
        if weights is not None:
            loss *= weights.expand_as(loss)
        return torch.mean(loss)


class GNN(nn.Module):
    def __init__(self, num_layer, input_dim, emb_dim, JK="last", drop_ratio=0, gnn_type="gin"):
        super(GNN, self).__init__()
        self.num_layer = num_layer
        self.drop_ratio = drop_ratio
        self.JK = JK
        # self.fc2 = nn.Linear(200, 200)
        self.gnns = torch.nn.ModuleList()
        for layer in range(num_layer):
            in_dim = input_dim if layer == 0 else emb_dim
            if gnn_type == "gin":
                # self.gnns.append(GINConv(nn.Sequential(nn.Linear(in_dim, emb_dim), nn.BatchNorm1d(emb_dim), nn.ReLU(),
                #                                        nn.Linear(emb_dim, emb_dim))))
                self.gnns.append(GINConv(nn.Sequential(nn.Linear(in_dim, emb_dim), GraphNorm(emb_dim), nn.ReLU(),
                                                       nn.Linear(emb_dim, emb_dim), nn.ReLU())))
            elif gnn_type == "gps":
                nn_ = Sequential(
                    Linear(in_dim, emb_dim),
                    ReLU(),
                    Linear(emb_dim, emb_dim),
                )
                conv = GPSConv(emb_dim, GINEConv(nn_), heads=4)
                self.gnns.append(conv)
            elif gnn_type == "gcn":
                self.gnns.append(GraphConv(in_dim, emb_dim))
            elif gnn_type == "gat":
                self.gnns.append(GATConv(in_dim, emb_dim))
            elif gnn_type == "gatv2":
                self.gnns.append(GATv2Conv(in_dim, emb_dim))
            elif gnn_type == "graphsage":
                self.gnns.append(SAGEConv(in_dim, emb_dim))
            else:
                raise ValueError("Invalid GNN type.")

    def forward(self, x, edge_index, edge_attr=None):
        h_list = [x]
        mut_site = []
        for layer in range(self.num_layer):

            h = self.gnns[layer](h_list[layer], edge_index, edge_attr)
            # if layer == self.num_layer - 1:
            #     # remove relu from the last layer
            #     h = F.dropout(h, self.drop_ratio, training=self.training)
            # else:
            #     h = F.dropout(F.relu(h), self.drop_ratio, training=self.training) # F.relu()
            h_list.append(h)
            # if len(h_list) == 2:
            #     previous_mut_site_feature = h_list[-2][mut_res_idx]
            #     current_mut_site_feature = h_list[-1][mut_res_idx]
            #     # print(previous_mut_site_feature.shape, current_mut_site_feature.shape)
            #     h_feature = self.global_encoder(previous_mut_site_feature)
            #     h_list[-1][mut_res_idx] = h_feature + current_mut_site_feature
            # if len(h_list) == 3:
            #     previous_mut_site_feature = h_list[-2][mut_res_idx].squeeze(0)
            #     current_mut_site_feature = h_list[-1][mut_res_idx].squeeze(0)

            #     h_feature = self.fc2(previous_mut_site_feature) + current_mut_site_feature
            #     h_list[-1][mut_res_idx] = h_feature.unsqueeze(0)
            # mut_site.append()
        # print(len(h_list))
        if self.JK == "last":
            node_representation = h_list[-1]
        elif self.JK == "sum":
            h_list = [h.unsqueeze_(0) for h in h_list]
            node_representation = torch.sum(torch.cat(h_list[1:], dim=0), dim=0)
        # print('node_rep', node_representation.shape)
        return h_list[-1]


# orthogonal initialization
def init_gru_orth(model, gain=1):
    model.reset_parameters()
    # orthogonal initialization of gru weights
    for _, hh, _, _ in model.all_weights:
        for i in range(0, hh.size(0), model.hidden_size):
            torch.nn.init.orthogonal_(hh[i:i + model.hidden_size], gain=gain)


def init_lstm_orth(model, gain=1):
    init_gru_orth(model, gain)

    # positive forget gate bias (Jozefowicz es at. 2015)
    for _, _, ih_b, hh_b in model.all_weights:
        l = len(ih_b)
        ih_b[l // 4: l // 2].data.fill_(1.0)
        hh_b[l // 4: l // 2].data.fill_(1.0)


class MLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers, dropout_rate):
        super(MLP, self).__init__()
        
        layers = []
        layers.append(nn.Linear(input_dim, hidden_dim))
        layers.append(nn.ReLU())
        layers.append(nn.Dropout(dropout_rate))
        
        for _ in range(num_layers - 1):
            layers.append(nn.Linear(hidden_dim, hidden_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout_rate))
        
        layers.append(nn.Linear(hidden_dim, output_dim))
        
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)


class FusionGraph(nn.Module):
    def __init__(self, num_layer, input_dim, emb_dim, out_dim, JK="last", drop_ratio=0.5, graph_pooling="attention",
                 gnn_type="gat", concat_type=None, fds=False, feature_level='both', contrast_curri=False, aux_mode='11') -> object:
        super().__init__()
        self.num_layer = num_layer
        self.drop_ratio = drop_ratio
        self.JK = JK
        self.input_dim = input_dim
        self.emb_dim = emb_dim
        self.out_dim = out_dim
        self.concat_type = concat_type
        self.feature_level = feature_level
        self.contrast_curri = contrast_curri
        self.mode = [False, False]
        final_dim = emb_dim

        if aux_mode[0] == '1':
            final_dim += 128
            self.mode[0] = True
            self.q_encoder = nn.LSTM(
                input_size=21,
                hidden_size=128,
                num_layers=2,
                batch_first=True,  # input & output will take batch size as 1 dim (batch, time_step, input_size)
                bidirectional=True
            )
            self.q_fc = nn.Linear(256, 128)
        
        if aux_mode[1] == '1':
            final_dim += 128
            self.mode[1] = True
            self.g_encoder = MLP(10, 128, 128, 3, 0.3)

        self.fc = nn.Sequential(
            nn.Linear(final_dim, self.emb_dim), nn.LeakyReLU(0.1), nn.Dropout(p=self.drop_ratio), 
            nn.Linear(self.emb_dim, self.emb_dim // 2), nn.LeakyReLU(0.1), nn.Dropout(p=self.drop_ratio),
            nn.Linear(self.emb_dim // 2, self.out_dim))

        self.gnn = GNN(num_layer, input_dim, emb_dim, JK, drop_ratio, gnn_type=gnn_type)

        if graph_pooling == "sum":
            self.pool = global_add_pool
        elif graph_pooling == "mean":
            self.pool = global_mean_pool
        elif graph_pooling == "max":
            self.pool = global_max_pool
        elif graph_pooling == "mul":
            self.pool = MulAggregation()
        elif graph_pooling == "attention":
            self.pool = GlobalAttention(gate_nn=torch.nn.Linear(emb_dim, 1))
        elif graph_pooling == "set2set":
            self.pool = Set2Set(emb_dim, processing_steps=2)
        elif graph_pooling == "lstm":
            self.pool = aggr.LSTMAggregation(emb_dim, emb_dim)
        else:
            raise ValueError("Invalid graph pooling type.")

    def forward_once(self, x, edge_index, batch):
        node_representation = self.gnn(x, edge_index)
        graph_rep = self.pool(node_representation, batch)
        return graph_rep

    def forward(self, data):
        fusion = [self.forward_once(data.x_s, data.edge_index_s, data.x_s_batch)]
        seq, globf = data.seq, data.global_f
        device = self.fc[0].bias.device
        if self.mode[0]:
            seq = torch.tensor(np.asarray(seq, dtype=np.float32), device=device)
            fusion.append(self.q_fc(self.q_encoder(seq)[0][:, -1, :]))
        if self.mode[1]:
            globf = torch.tensor(np.asarray(globf, dtype=np.float32), device=device)
            fusion.append(self.g_encoder(globf))
        fusion = torch.cat(fusion, dim=-1)
        x = self.fc(fusion)
        return x



class MMGraph(nn.Module):
    def __init__(self, num_layer, input_dim, emb_dim, out_dim, JK="last", drop_ratio=0.5, graph_pooling="attention",
                 gnn_type="gat", concat_type=None, fds=False, feature_level='both', contrast_curri=False, max_length=50) -> object:
        super(MMGraph, self).__init__()
        self.num_layer = num_layer
        self.drop_ratio = drop_ratio
        self.JK = JK
        self.input_dim = input_dim
        self.emb_dim = emb_dim
        self.out_dim = out_dim
        self.concat_type = concat_type
        self.feature_level = feature_level
        self.contrast_curri = contrast_curri

        self.graph_pool = nn.Linear(self.emb_dim, 1)

        self.fc = nn.Sequential(
            nn.Linear(self.emb_dim, self.emb_dim), nn.LeakyReLU(0.1), nn.Dropout(p=self.drop_ratio),
            nn.Linear(self.emb_dim, self.out_dim))

        if fds:
            self.dir = True
        else:
            self.dir = False
        self.global_encoder = nn.Sequential(nn.Linear(10, self.emb_dim), nn.LeakyReLU(0.1), nn.Dropout(p=self.drop_ratio),)
        self.seq_encoder = nn.Sequential(
            nn.Linear(max_length, self.emb_dim), nn.LeakyReLU(0.1), nn.Dropout(p=self.drop_ratio),
            )
        self.gnn = GNN(num_layer, input_dim, emb_dim, JK, drop_ratio, gnn_type=gnn_type)

        if graph_pooling == "sum":
            self.pool = global_add_pool
        elif graph_pooling == "mean":
            self.pool = global_mean_pool
        elif graph_pooling == "max":
            self.pool = global_max_pool
        elif graph_pooling == "mul":
            self.pool = MulAggregation()
        elif graph_pooling == "attention":
            self.pool = GlobalAttention(gate_nn=torch.nn.Linear(emb_dim, 1))
        elif graph_pooling == "set2set":
            self.pool = Set2Set(emb_dim, processing_steps=2)
        elif graph_pooling == "lstm":
            self.pool = aggr.LSTMAggregation(emb_dim, emb_dim)
        else:
            raise ValueError("Invalid graph pooling type.")
        self.att = SimpleSelfAttention(emb_dim, num_heads=4)

    def forward_once(self, x, edge_index, batch):
        node_representation = self.gnn(x, edge_index)
        graph_rep = self.pool(node_representation, batch)
        return graph_rep

    def forward(self, data):
        seq1, global_1 = data.seq, data.global_f
        device = self.graph_pool.bias.device
        seq1 = torch.tensor(seq1, dtype=torch.float, device=device)
        global_1 = torch.tensor(global_1, dtype=torch.float, device=device)

        graph_rep_be = self.forward_once(data.x_s, data.edge_index_s, data.x_s_batch)
        seq1_rep_be = self.seq_encoder(seq1)
        global1 = self.global_encoder(global_1)

        a1 = self.att(graph_rep_be, seq1_rep_be, global1)
        return self.fc(a1)



class PMMGraph(nn.Module):
    def __init__(self, num_layer, input_dim, emb_dim, out_dim, JK="last", drop_ratio=0.5, graph_pooling="attention",
                 gnn_type="gat", concat_type=None, fds=False, feature_level='both', contrast_curri=False):
        super(PMMGraph, self).__init__()
        self.num_layer = num_layer
        self.drop_ratio = drop_ratio
        self.JK = JK
        self.input_dim = input_dim
        self.emb_dim = emb_dim
        self.out_dim = out_dim
        self.concat_type = concat_type
        self.feature_level = feature_level
        self.contrast_curri = contrast_curri

        # Define the learnable prompt token
        self.prompt_token = nn.Parameter(torch.randn(1, 10))

        self.graph_pool = nn.Linear(self.emb_dim, 1)

        self.fc = nn.Sequential(
            nn.Linear(self.emb_dim + 10, self.emb_dim),  # Adjust input size to include the prompt token
            nn.LeakyReLU(0.1), nn.Dropout(p=self.drop_ratio),
            nn.Linear(self.emb_dim, self.out_dim))

        self.dir = fds
        self.global_encoder = nn.Sequential(nn.Linear(10, self.emb_dim), nn.LeakyReLU(0.1), nn.Dropout(p=self.drop_ratio))
        self.seq_encoder = nn.Sequential(
            nn.Linear(30, self.emb_dim), nn.LeakyReLU(0.1), nn.Dropout(p=self.drop_ratio),
        )
        self.gnn = GNN(num_layer, input_dim, emb_dim, JK, drop_ratio, gnn_type=gnn_type)

        # Initialize pooling based on the specified type
        if graph_pooling in ["sum", "mean", "max", "mul", "attention", "set2set", "lstm"]:
            pooling_classes = {
                "sum": global_add_pool,
                "mean": global_mean_pool,
                "max": global_max_pool,
                "mul": aggr.MulAggregation(),
                "attention": GlobalAttention(gate_nn=torch.nn.Linear(emb_dim, 1)),
                "set2set": Set2Set(emb_dim, processing_steps=2),
                "lstm": aggr.LSTMAggregation(emb_dim, emb_dim)
            }
            self.pool = pooling_classes[graph_pooling]
        else:
            raise ValueError("Invalid graph pooling type.")

        self.att = SimpleSelfAttention(emb_dim + 10, num_heads=4)  # Adjust for prompt dimension

    def forward_once(self, x, edge_index, batch):
        node_representation = self.gnn(x, edge_index)
        graph_rep = self.pool(node_representation, batch)
        # Concatenate the prompt token
        graph_rep = torch.cat([graph_rep, self.prompt_token.expand(graph_rep.size(0), -1)], dim=1)
        return graph_rep

    def forward(self, data):
        seq1 = torch.tensor(np.array(data.seq, dtype=np.float32)).to(device='cuda')
        global_1 = torch.tensor(np.array(data.global_f, dtype=np.float32)).to(device='cuda')

        graph_rep_be = self.forward_once(data.x_s, data.edge_index_s, data.x_s_batch)
        seq1_rep_be = self.seq_encoder(seq1)
        global1_rep = self.global_encoder(global_1)
        # Concatenate the prompt token to other representations as well
        seq1_rep_be = torch.cat([seq1_rep_be, self.prompt_token.expand(seq1_rep_be.size(0), -1)], dim=1)
        global1_rep = torch.cat([global1_rep, self.prompt_token.expand(global1_rep.size(0), -1)], dim=1)

        # Process combined representations
        a1 = self.att(graph_rep_be, seq1_rep_be, global1_rep)
        return self.fc(a1)