lstm_partial_grid_rri.v.9.py

#CUDA_VISIBLE_DEVICES=1
import torch
import torch.autograd as autograd
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.nn import init
import numpy as np
import json
import os.path
import subprocess
import random
from operator import itemgetter
import sklearn.metrics as metrics 
import sys

from benchmark_tools import *

np.set_printoptions(linewidth=1000000000)
torch.cuda.manual_seed(1)

verbose = False#True
file_name = sys.argv[1]
results_path = sys.argv[2]

training = dict()
testing = dict()

#training['features'], training['sequences'] = read_seq_pssm("pssm_dimers_466_list.tsv","dimers_466")
#training['contacts'] = read_contacts("rri_dimers_466_list.tsv","dimers_466_rri",training['features'])
#training['scop'] = read_scop("466_dimers_list.tsv")

#training['features'], training['sequences'] = read_seq_pssm("pssm_dimers_1550_list.tsv","dimers_1550")
#training['contacts'] = read_contacts("rri_dimers_1550_list.tsv","dimers_1550_rri",training['features'])
#training['scop'] = read_scop("1550_polymers_list.tsv")

training['features'], training['sequences'] = read_seq_pssm("pssm_dimers_450_list.tsv","dimers_450")
training['contacts'] = read_contacts("rri_dimers_450_list.tsv","dimers_450_rri",training['features'])
training['scop'] = read_scop("450_dimers_list.tsv")

#training['features'], training['sequences'] = read_seq_pssm("pssm_dimers_230_list.tsv","dimers_230")
#training['contacts'] = read_contacts("rri_dimers_230_list.tsv","dimers_230_rri",training['features'])
#training['scop'] = read_scop("230_dimers_list.tsv")

testing['features'], testing['sequences'] = read_seq_pssm("pssm_dimers_bipspi_list.tsv","dimers_bipspi")
testing['contacts'] = read_contacts("rri_dimers_bipspi_list.tsv","dimers_bipspi_rri",testing['features'])


class BiLSTM(nn.Module):

    def __init__( self, input_dim=22, lstm_hidden_dim=256, hidden_1_dim=1024, hidden_2_dim=512, hidden_3_dim=1024, rri_size=2 ):
        super(BiLSTM, self).__init__()

        self.input_dim = input_dim
        self.lstm_hidden_dim = lstm_hidden_dim
        self.hidden_1_dim = hidden_1_dim
        self.hidden_2_dim = hidden_2_dim
        self.hidden_3_dim = hidden_3_dim
        self.rri_size = rri_size

        self.lstm_h0 = None
        self.lstm_c0 = None
        self.update_lstm_hidden()

        self.LSTM = nn.LSTM(input_dim, lstm_hidden_dim, num_layers=2, bidirectional=True, dropout=0.5)

        self.drop_hidden_1 = nn.Dropout(p=0.5)
        self.lstm2hidden_1 = nn.Linear(2*lstm_hidden_dim, hidden_1_dim)
        
        self.drop_hidden_2 = nn.Dropout(p=0.5)
        self.hidden2hidden_2 = nn.Linear(hidden_1_dim, hidden_2_dim)

        self.drop_hidden_3 = nn.Dropout(p=0.5)
        self.hidden2hidden_3 = nn.Linear(2*hidden_2_dim, hidden_3_dim)

        self.hidden2out = nn.Linear(hidden_3_dim, rri_size)

    def update_lstm_hidden(self):
        self.lstm_h0 = autograd.Variable(torch.zeros(4, 1, self.lstm_hidden_dim)).cuda()
        self.lstm_c0 = autograd.Variable(torch.zeros(4, 1, self.lstm_hidden_dim)).cuda()
        

    def prepare_data(self, pdb, sequence, features):
        list_pssm = []
        for aa in sequence:
          v = list(features[pdb][aa]["pssm"])
          list_pssm.append(v)
        return autograd.Variable( torch.unsqueeze(torch.FloatTensor(list_pssm),dim=1) ).cuda() 

    def forward( self, pdb, sequence_r, sequence_l, features, contacts, ch_r=None, ch_l=None, Flag=True ):

        v_r = self.prepare_data( pdb, sequence_r, features )
        v_l = self.prepare_data( pdb, sequence_l, features )

        self.update_lstm_hidden()
        out_LSTM_r, (hidden_LSTM_r, content_LSTM_r) = self.LSTM( v_r, (self.lstm_h0, self.lstm_c0))

        self.update_lstm_hidden()
        out_LSTM_l, (hidden_LSTM_l, content_LSTM_l) = self.LSTM( v_l, (self.lstm_h0, self.lstm_c0))

        hidden_r_1 = self.lstm2hidden_1( out_LSTM_r.view(len(sequence_r), -1) )
        hidden_r_1 = self.drop_hidden_1(hidden_r_1)
        out_hidden_r_1 = F.relu(hidden_r_1)

        hidden_r_2 = self.hidden2hidden_2( out_hidden_r_1 )
        hidden_r_2 = self.drop_hidden_2(hidden_r_2)
        out_hidden_r_2 = F.relu(hidden_r_2)

        hidden_l_1 = self.lstm2hidden_1( out_LSTM_l.view(len(sequence_l), -1) )
        hidden_l_1 = self.drop_hidden_1(hidden_l_1)
        out_hidden_l_1 = F.relu(hidden_l_1)

        hidden_l_2 = self.hidden2hidden_2( out_hidden_l_1 )
        hidden_l_2 = self.drop_hidden_2(hidden_l_2)
        out_hidden_l_2 = F.relu(hidden_l_2)

        rl = 0
        N_r = len(sequence_r)
        N_l = len(sequence_l)
        native_rri = []

        if Flag:
          res_rri = list(contacts['rri_ch_ch'][pdb][ch_r+":"+ch_l].keys())
          N_rri = len(res_rri)
          v_in = autograd.Variable( torch.FloatTensor(N_rri+len(res_rri),2*self.hidden_2_dim) ).cuda()
          v_in_t = autograd.Variable( torch.FloatTensor(N_rri+len(res_rri),2*self.hidden_2_dim) ).cuda()

          for rr in res_rri:
            R = rr.split(":")
            try:
              i_r = sequence_r.index(R[0])
            except ValueError as err:
              raise err
            try:
              j_l = sequence_l.index(R[1])
            except ValueError as err:
              raise err

            w_rc = out_hidden_r_2[i_r,:]
            w_lc = out_hidden_l_2[j_l,:]

            v_in[rl,:] = torch.cat( [w_rc, w_lc], dim=0 )
            v_in_t[rl,:] = torch.cat( [w_lc, w_rc], dim=0 )
            rl += 1
            native_rri.append(1)

          while N_rri>0:
            i_r = random.randint(0,N_r-1)
            j_l = random.randint(0,N_l-1)
            if( not sequence_r[i_r]+":"+sequence_l[j_l] in contacts['rri'][pdb]):
              w_rc = out_hidden_r_2[i_r,:]
              w_lc = out_hidden_l_2[j_l,:]

              v_in[rl,:] = torch.cat( [w_rc,w_lc], dim=0 )
              v_in_t[rl,:] = torch.cat( [w_lc,w_rc], dim=0 )

              rl += 1
              N_rri -= 1
              native_rri.append(0)

          hidden_3 = self.hidden2hidden_3(v_in)
          hidden_3 = self.drop_hidden_3(hidden_3)

          hidden_3_t = self.hidden2hidden_3(v_in_t)
          hidden_3_t = self.drop_hidden_3(hidden_3_t)

          out_hidden_3 = F.relu(0.5*(hidden_3+hidden_3_t))

          rri_out = self.hidden2out( out_hidden_3 )
          rri_out = F.log_softmax( rri_out, dim=1 )

          native_rri = autograd.Variable(torch.LongTensor(native_rri)).cuda()

          return rri_out, native_rri, out_hidden_3

        else:
          v_in = autograd.Variable( torch.FloatTensor(N_r*N_l,2*self.hidden_2_dim) ).cuda()
          v_in_t = autograd.Variable( torch.FloatTensor(N_r*N_l,2*self.hidden_2_dim) ).cuda()
          for i_r in range( N_r ):
            for j_l in range( N_l ): 
              if sequence_r[i_r]+":"+sequence_l[j_l] in contacts['rri'][pdb]:
                native_rri.append(1)
              else:
                native_rri.append(0)

              w_rc = out_hidden_r_2[i_r,:]
              w_lc = out_hidden_l_2[j_l,:]

              v_in[rl,:] = torch.cat( [w_rc,w_lc], dim=0 )
              v_in_t[rl,:] = torch.cat( [w_lc,w_rc], dim=0 )
              rl += 1

          hidden_3 = self.hidden2hidden_3(v_in)
          hidden_3 = self.drop_hidden_3(hidden_3)

          hidden_3_t = self.hidden2hidden_3(v_in_t)
          hidden_3_t = self.drop_hidden_3(hidden_3_t)

          out_hidden_3 = F.relu(0.5*(hidden_3+hidden_3_t))

          rri_out = self.hidden2out( out_hidden_3 )
          rri_out = F.log_softmax( rri_out, dim=1 )

          native_rri = autograd.Variable(torch.LongTensor(native_rri)).cuda()

          return rri_out, native_rri, out_hidden_3
          
input_dim=22
lstm_hidden_dim=256
hidden_1_dim=512
hidden_2_dim=256
hidden_3_dim=1024

#input_dim=22
#lstm_hidden_dim=16
#hidden_1_dim=32
#hidden_2_dim=16
#hidden_3_dim=64

model = BiLSTM( input_dim=input_dim, lstm_hidden_dim=lstm_hidden_dim, hidden_1_dim=hidden_1_dim, hidden_2_dim=hidden_2_dim, hidden_3_dim=hidden_3_dim, rri_size=2 )
model.cuda()
if(verbose): print(model)
loss_function = nn.NLLLoss()

if(verbose): print("Neural networking ...")

TRAGETS = list(map(str.strip, open(file_name,"r").readlines()))

LOSS_ = []
for target_ in TRAGETS:
  
  r = target_.split("\t")
  target = r[0]
  scop_r = r[3]
  scop_l = r[4]

  if os.path.isfile(results_path+"/"+target+".tsv"):
    continue 

  lr  = 0.1
  model = BiLSTM( input_dim=input_dim, lstm_hidden_dim=lstm_hidden_dim, hidden_1_dim=hidden_1_dim, hidden_2_dim=hidden_2_dim, hidden_3_dim=hidden_3_dim, rri_size=2 )
  model.cuda()

  for epoch in range(50):
    optimizer = optim.SGD(model.parameters(), lr=lr)
    lr *= 0.99
    N_current = training['contacts']['N_cci']-1
    cci_ = list(training['contacts']['cci'].keys())
    random.shuffle(cci_)

    for pdb in cci_:
      for ch_ch in training['contacts']['cci'][pdb]:
        R = ch_ch.split(":")
        ch_r = R[0]
        ch_l = R[1]
        if training['scop'][pdb][ch_ch] == scop_r+":::"+scop_l or training['scop'][pdb][ch_ch] == scop_l+":::"+scop_r:
          N_current -= 1
          continue

        local_sequence_r = training['sequences'][pdb][ch_r]
        local_sequence_l = training['sequences'][pdb][ch_l]

        model.zero_grad()
        optimizer.zero_grad()

        predicted_rri, native_rri, out_hidden_3 = model( pdb, local_sequence_r, local_sequence_l, training['features'], training['contacts'], ch_r=ch_r, ch_l=ch_l, Flag=True )
        l1 = 1e-6*torch.abs(out_hidden_3).sum()
        loss = loss_function( predicted_rri, native_rri ) + l1
        loss.backward()
        optimizer.step()
        LOSS_.append( loss.data[0] )
        if(verbose): print("%d - %s %s:%s - LOSS: %0.4f - L1: %0.4f         \r"  %(N_current, pdb,ch_r,ch_l,loss.data[0],l1.data[0]),end="")
        N_current -= 1

    model.train(mode=False)
    ##TRAINING SET AUC SCORE FOR EACH EPOCH
    AUC = []
    for pdb in cci_:
      for ch_ch in training['contacts']['cci'][pdb]:
        R = ch_ch.split(":")
        ch_r = R[0]
        ch_l = R[1]

        local_sequence_r = training['sequences'][pdb][ch_r]
        local_sequence_l = training['sequences'][pdb][ch_l]

        predicted_rri, native_rri, out_hidden_3 = model( pdb, local_sequence_r, local_sequence_l, training['features'], training['contacts'], ch_r=ch_r, ch_l=ch_l, Flag=True )
        np_class = native_rri.data.cpu().numpy() 
        np_prediction = predicted_rri.data.cpu()[:,1].numpy()
  
        fpr, tpr, thresholds = metrics.roc_curve(np_class, np_prediction, pos_label=1)
        new_auc = metrics.auc(fpr, tpr)
        AUC.append(new_auc)
    training_auc =  np.mean(AUC)

    ##TESTING SET AUC FOR EACH EPOCH
    for ch_ch in testing['contacts']['cci'][target]:
      R = ch_ch.split(":")
      ch_r = R[0]
      ch_l = R[1]

      local_sequence_r = testing['sequences'][target][ch_r]
      local_sequence_l = testing['sequences'][target][ch_l]

      N_r = len(local_sequence_r)
      N_l = len(local_sequence_l)

      model.zero_grad()
      optimizer.zero_grad()

      predicted_rri, native_rri, out_hidden_3 = model( target, local_sequence_r, local_sequence_l, testing['features'], testing['contacts'], ch_r=ch_r, ch_l=ch_l, Flag=False )

      np_class = native_rri.data.cpu().numpy() 
      np_prediction = predicted_rri.data.cpu()[:,1].numpy()

      np_all = np.stack((np_prediction,np_class),  axis=-1)
      np_all = np.insert(np_all, 0, np.array((N_r,N_l)), 0) 
      np.savetxt(results_path+"/predictions/"+target+"."+ch_r+":"+ch_l+"."+str(epoch)+".tsv",np_all)

      fpr, tpr, thresholds = metrics.roc_curve(np_class, np_prediction, pos_label=1)
      testing_auc = metrics.auc(fpr, tpr)

      loss_ = np.array(LOSS_).mean()
      if(verbose): print( "%d - %s %s:%s - AUC=%0.4f - TRAINING_AUC=%0.4f - LOSS=%0.4f" % (epoch, target,ch_r,ch_l,testing_auc,training_auc, loss_) )
      open(results_path+"/"+target+".tsv", "a").write("%d - %s %s:%s - AUC=%0.4f - TRAINING_AUC=%0.4f\n" % (epoch, target,ch_r,ch_l,testing_auc,training_auc) )
    model.train(mode=True)