layers.py

import torch
import torch.nn as nn
from torch.nn import functional as F
import math

class MLP(nn.Module):
    def __init__(self, in_feat, hid_feat=None, out_feat=None,
                 dropout=0.):
        super().__init__()
        if not hid_feat:
            hid_feat = in_feat
        if not out_feat:
            out_feat = in_feat
        self.fc1 = nn.Linear(in_feat, hid_feat)
        self.act = torch.nn.ReLU()
        self.fc2 = nn.Linear(hid_feat,out_feat)
        self.droprateout = nn.Dropout(dropout)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        return self.droprateout(x)

class Attention_new(nn.Module):
    def __init__(self, dim, heads, attention_dropout=0.):
        super().__init__()
        assert dim % heads == 0
        self.heads = heads
        self.scale = 1./dim**0.5

        self.q = nn.Linear(dim, dim)
        self.k = nn.Linear(dim, dim)
        self.v = nn.Linear(dim, dim)
        self.e = nn.Linear(dim, dim)
        #self.attention_dropout = nn.Dropout(attention_dropout)

        self.d_k = dim // heads  
        self.heads = heads
        self.out_e = nn.Linear(dim,dim)
        self.out_n = nn.Linear(dim, dim)
        
        
    def forward(self, node, edge):
        b, n, c = node.shape
        
        
        q_embed = self.q(node).view(-1, n, self.heads, c//self.heads)
        k_embed = self.k(node).view(-1, n, self.heads, c//self.heads)
        v_embed = self.v(node).view(-1, n, self.heads, c//self.heads)
   
        e_embed = self.e(edge).view(-1, n, n, self.heads, c//self.heads)
        
        q_embed = q_embed.unsqueeze(2)
        k_embed = k_embed.unsqueeze(1)
        
        attn = q_embed * k_embed
        
        attn = attn/ math.sqrt(self.d_k)
        
     
        attn = attn * (e_embed + 1) * e_embed

        edge = self.out_e(attn.flatten(3))  
      
        attn = F.softmax(attn, dim=2)
        
        v_embed = v_embed.unsqueeze(1)
  
        v_embed = attn * v_embed
        
        v_embed = v_embed.sum(dim=2).flatten(2)
        
        node  = self.out_n(v_embed)
           
        return node, edge

class Encoder_Block(nn.Module):
    def __init__(self, dim, heads,act, mlp_ratio=4, drop_rate=0.):
        super().__init__()
        self.ln1 = nn.LayerNorm(dim)
   
        self.attn = Attention_new(dim, heads, drop_rate)
        self.ln3 = nn.LayerNorm(dim)
        self.ln4 = nn.LayerNorm(dim)
        self.mlp = MLP(dim, dim*mlp_ratio, dim, dropout=drop_rate)
        self.mlp2 = MLP(dim, dim*mlp_ratio, dim, dropout=drop_rate)
        self.ln5 = nn.LayerNorm(dim)
        self.ln6 = nn.LayerNorm(dim)

    def forward(self, x,y):
        x1 = self.ln1(x)
        x2,y1 = self.attn(x1,y)
        x2 = x1 + x2
        y2 = y1 + y
        x2 = self.ln3(x2)   
        y2 = self.ln4(y2)   
        
        x = self.ln5(x2 + self.mlp(x2))
        y = self.ln6(y2 + self.mlp2(y2))
        return x, y


class TransformerEncoder(nn.Module):
    def __init__(self, dim, depth, heads, act, mlp_ratio=4, drop_rate=0.1):
        super().__init__()
        
        self.Encoder_Blocks = nn.ModuleList([
            Encoder_Block(dim, heads, act, mlp_ratio, drop_rate)
            for i in range(depth)])

    def forward(self, x,y):
        
        for Encoder_Block in self.Encoder_Blocks:
            x,  y = Encoder_Block(x,y)
            
        return x, y

class enc_dec_attention(nn.Module):
    def __init__(self, dim, heads, attention_dropout=0., proj_dropout=0.):
        super().__init__()
        self.dim = dim
        self.heads = heads
        self.scale = 1./dim**0.5
 
        
        "query is molecules"
        "key is prot"
        "values is again molecule"
        self.q_mx = nn.Linear(dim,dim)
        self.k_px = nn.Linear(dim,dim)
        self.v_mx = nn.Linear(dim,dim)
        
        
        self.k_pa = nn.Linear(dim,dim)
        self.v_ma = nn.Linear(dim,dim)
        
    

 
    
        #self.dropout_dec = nn.Dropout(proj_dropout)
        self.out_nd = nn.Linear(dim, dim)
        self.out_ed = nn.Linear(dim,dim)
        
    def forward(self, mol_annot, prot_annot, mol_adj, prot_adj):
        
        b, n, c = mol_annot.shape
        _, m, _ = prot_annot.shape
     
        
        query_mol_annot = self.q_mx(mol_annot).view(-1,m, self.heads, c//self.heads)
        key_prot_annot = self.k_px(prot_annot).view(-1,n, self.heads, c//self.heads)
        value_mol_annot = self.v_mx(mol_annot).view(-1,m, self.heads, c//self.heads)
        
        mol_e = self.v_ma(mol_adj).view(-1,m,m, self.heads, c//self.heads)
        prot_e = self.k_pa(prot_adj).view(-1,m,m, self.heads, c//self.heads)
        
        query_mol_annot = query_mol_annot.unsqueeze(2)
        key_prot_annot = key_prot_annot.unsqueeze(1)
        
        
        
        #attn = torch.einsum('bnchd,bmahd->bnahd', query_mol_annot, key_prot_annot)
        
        attn = query_mol_annot * key_prot_annot
        
        attn = attn/ math.sqrt(self.dim)

        
        attn = attn * (prot_e + 1) * mol_e      
      
        mol_e_new = attn.flatten(3)
        
        mol_adj = self.out_ed(mol_e_new)
        
        attn = F.softmax(attn, dim=2)
        
        value_mol_annot = value_mol_annot.unsqueeze(1)
        
        value_mol_annot = attn * value_mol_annot
        
        value_mol_annot = value_mol_annot.sum(dim=2).flatten(2)
        
        mol_annot  = self.out_nd(value_mol_annot)          

        return mol_annot, prot_annot, mol_adj, prot_adj

class Decoder_Block(nn.Module):
    def __init__(self, dim, heads, mlp_ratio=4, drop_rate=0.):
        super().__init__()

        
        self.ln1_ma = nn.LayerNorm(dim)
        self.ln1_pa = nn.LayerNorm(dim)
        self.ln1_mx = nn.LayerNorm(dim)
        self.ln1_px = nn.LayerNorm(dim)
        
        self.attn2 = Attention_new(dim, heads, drop_rate)
        
        self.ln2_pa = nn.LayerNorm(dim)
        self.ln2_px = nn.LayerNorm(dim)
        
        self.dec_attn = enc_dec_attention(dim, heads, drop_rate, drop_rate)
        
        self.ln3_ma = nn.LayerNorm(dim)
        self.ln3_mx = nn.LayerNorm(dim)

        self.mlp_ma = MLP(dim, dim, dropout=drop_rate)
        self.mlp_mx = MLP(dim, dim, dropout=drop_rate)
       
        self.ln4_ma = nn.LayerNorm(dim)
        self.ln4_mx = nn.LayerNorm(dim)
    
        
    def forward(self,mol_annot, prot_annot, mol_adj, prot_adj):

        mol_annot = self.ln1_mx(mol_annot)
        mol_adj = self.ln1_ma(mol_adj)
        
        prot_annot = self.ln1_px(prot_annot)
        prot_adj = self.ln1_pa(prot_adj)
        
        px1, pa1= self.attn2(prot_annot, prot_adj)
        
        prot_annot = prot_annot + px1
        prot_adj = prot_adj + pa1
        
        prot_annot = self.ln2_px(prot_annot)
        prot_adj = self.ln2_pa(prot_adj)
        
        mx1, prot_annot, ma1, prot_adj  = self.dec_attn(mol_annot,prot_annot,mol_adj,prot_adj)
        
        ma1 = mol_adj + ma1
        mx1 = mol_annot + mx1
        
        ma2 = self.ln3_ma(ma1)
        mx2 = self.ln3_mx(mx1)
        
        ma3 = self.mlp_ma(ma2)
        mx3 = self.mlp_mx(mx2)
        
        ma = ma3 + ma2
        mx = mx3 + mx2
        
        mol_adj = self.ln4_ma(ma)
        mol_annot = self.ln4_mx(mx)           
    
        return mol_annot, prot_annot, mol_adj, prot_adj
    
class TransformerDecoder(nn.Module):
    def __init__(self, dim,  depth, heads, mlp_ratio=4, drop_rate=0.):
        super().__init__()
        
        self.Decoder_Blocks = nn.ModuleList([     
            Decoder_Block(dim, heads, mlp_ratio, drop_rate)
            for i in range(depth)])
        
    def forward(self, mol_annot, prot_annot, mol_adj, prot_adj):
        
        for Decoder_Block in self.Decoder_Blocks:
            mol_annot, prot_annot, mol_adj, prot_adj  = Decoder_Block(mol_annot, prot_annot, mol_adj, prot_adj)
            
        return mol_annot, prot_annot,mol_adj, prot_adj