<reponame>t-zhong/CPI_prediction
import pickle
import sys
from tqdm import tqdm

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from sklearn.metrics import precision_score, recall_score, roc_auc_score
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.dataset import random_split

torch.manual_seed(1234)


class CPIPrediction(nn.Module):
    def __init__(self):
        super(CPIPrediction, self).__init__()
        self.embed_fingerprint = nn.Embedding(n_fingerprint, dim) # molecule embedding
        self.embed_word = nn.Embedding(n_word, dim) # protein embedding
        self.W_gnn = nn.ModuleList([nn.Linear(dim, dim) #
                                    for _ in range(layer_gnn)])
        self.W_cnn = nn.ModuleList([nn.Conv2d(
                     in_channels=1, out_channels=1, kernel_size=2*window+1,
                     stride=1, padding=window) for _ in range(layer_cnn)])
        self.W_attention = nn.Linear(dim, dim)
        self.W_out = nn.ModuleList([nn.Linear(2*dim, 2*dim)
                                    for _ in range(layer_output)])
        self.W_interaction = nn.Linear(2*dim, 1)

    def gnn(self, xs, A, layer):
        for i in range(layer):
            hs = torch.relu(self.W_gnn[i](xs))
            xs = xs + torch.matmul(A, hs)
        # return torch.unsqueeze(torch.sum(xs, 0), 0)
        return torch.mean(xs, 1)

    def attention_cnn(self, x, xs, layer):
        """The attention mechanism is applied to the last layer of CNN."""

        xs = torch.unsqueeze(xs, 1)
        for i in range(layer):
            xs = torch.relu(self.W_cnn[i](xs))
        xs = torch.squeeze(xs, 1)

        h = torch.relu(self.W_attention(x))
        hs = torch.relu(self.W_attention(xs))
        weights = torch.tanh(F.linear(h, hs))
        ys = torch.t(weights) * hs

        # return torch.unsqueeze(torch.sum(ys, 0), 0)
        return torch.mean(ys, 1)

    def forward(self, fingerprints, adjacency, words):

        # fingerprints, adjacency, words = inputs

        """Compound vector with GNN."""
        fingerprint_vectors = self.embed_fingerprint(fingerprints)
        compound_vector = self.gnn(fingerprint_vectors, adjacency, layer_gnn)

        """Protein vector with attention-CNN."""
        word_vectors = self.embed_word(words)
        protein_vector = self.attention_cnn(compound_vector,
                                            word_vectors, layer_cnn)

        """Concatenate the above two vectors and output the interaction."""
        cat_vector = torch.cat((compound_vector, protein_vector), 1)
        for j in range(layer_output):
            cat_vector = torch.relu(self.W_out[j](cat_vector))
        interaction = self.W_interaction(cat_vector)

        return F.softmax(interaction, dim=1)

class CPIDataset(Dataset):
    
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        return self.data[index]     


def load_tensor(file_name, dtype):
    return [dtype(d).to(device) for d in np.load(file_name + '.npy', allow_pickle=True)]


def load_pickle(file_name):
    with open(file_name, 'rb') as f:
        return pickle.load(f)


def train_step(model, batch):
    model.train()
    compounds, adjacencies, proteins, true_y = batch
    pred_y = model(compounds, adjacencies, proteins)
    loss = F.binary_cross_entropy(pred_y, true_y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

def val_step(model, batch):
    model.eval()
    compounds, adjacencies, proteins, true_y = batch
    pred_y = model(compounds, adjacencies, proteins)
    loss = F.binary_cross_entropy(pred_y, true_y)
    return loss.item()


if __name__ == "__main__":

    """Hyperparameters."""
    (DATASET, radius, ngram, dim, layer_gnn, window, layer_cnn, layer_output,
     lr, lr_decay, decay_interval, weight_decay, iteration,
     setting) = sys.argv[1:]
    (dim, layer_gnn, window, layer_cnn, layer_output, decay_interval,
     iteration) = map(int, [dim, layer_gnn, window, layer_cnn, layer_output,
                            decay_interval, iteration])
    lr, lr_decay, weight_decay = map(float, [lr, lr_decay, weight_decay])

    """CPU or GPU."""
    if torch.cuda.is_available():
        device = torch.device('cuda:1')
        print('The code uses GPU...')
    else:
        device = torch.device('cpu')
        print('The code uses CPU!!!')

    """Load preprocessed data."""
    dir_input = ('../dataset/' + DATASET + '/input/'
                 'radius' + radius + '_ngram' + ngram + '/')
    compounds = load_tensor(dir_input + 'compounds', torch.LongTensor)
    adjacencies = load_tensor(dir_input + 'adjacencies', torch.FloatTensor)
    proteins = load_tensor(dir_input + 'proteins', torch.LongTensor)
    interactions = load_tensor(dir_input + 'interactions', torch.LongTensor)
    fingerprint_dict = load_pickle(dir_input + 'fingerprint_dict.pickle')
    word_dict = load_pickle(dir_input + 'word_dict.pickle')
    n_fingerprint = len(fingerprint_dict)
    n_word = len(word_dict)

    """Create a dataset and split it into train/dev/test."""
    dataset = list(zip(compounds, adjacencies, proteins, interactions))
    
    count = len(dataset)
    train_len = int(count * 0.8)
    val_len = count - train_len
    train_set, val_set = random_split(dataset, [train_len, val_len])
    trainloader = DataLoader(train_set, batch_size=1, shuffle=True)
    valloader = DataLoader(val_set, batch_size=1, shuffle=False)


    """Set a model."""
    model = CPIPrediction().to(device)

    """Start training."""
    print('Training...')
    optimizer = optim.Adam(model.parameters(), lr=0.1)
    lr_scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.99)
    for epoch in range(iteration):
        for batch in tqdm(trainloader):
            loss = train_step(model, batch)

    

import pickle
import sys
from tqdm import tqdm

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from sklearn.metrics import precision_score, recall_score, roc_auc_score
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.dataset import random_split

torch.manual_seed(1234)


class CPIPrediction(nn.Module):
    def __init__(self):
        super(CPIPrediction, self).__init__()
        self.embed_fingerprint = nn.Embedding(n_fingerprint, dim) # molecule embedding
        self.embed_word = nn.Embedding(n_word, dim) # protein embedding
        self.W_gnn = nn.ModuleList([nn.Linear(dim, dim) #
                                    for _ in range(layer_gnn)])
        self.W_cnn = nn.ModuleList([nn.Conv2d(
                     in_channels=1, out_channels=1, kernel_size=2*window+1,
                     stride=1, padding=window) for _ in range(layer_cnn)])
        self.W_attention = nn.Linear(dim, dim)
        self.W_out = nn.ModuleList([nn.Linear(2*dim, 2*dim)
                                    for _ in range(layer_output)])
        self.W_interaction = nn.Linear(2*dim, 1)

    def gnn(self, xs, A, layer):
        for i in range(layer):
            hs = torch.relu(self.W_gnn[i](xs))
            xs = xs + torch.matmul(A, hs)
        # return torch.unsqueeze(torch.sum(xs, 0), 0)
        return torch.mean(xs, 1)

    def attention_cnn(self, x, xs, layer):
        """The attention mechanism is applied to the last layer of CNN."""

        xs = torch.unsqueeze(xs, 1)
        for i in range(layer):
            xs = torch.relu(self.W_cnn[i](xs))
        xs = torch.squeeze(xs, 1)

        h = torch.relu(self.W_attention(x))
        hs = torch.relu(self.W_attention(xs))
        weights = torch.tanh(F.linear(h, hs))
        ys = torch.t(weights) * hs

        # return torch.unsqueeze(torch.sum(ys, 0), 0)
        return torch.mean(ys, 1)

    def forward(self, fingerprints, adjacency, words):

        # fingerprints, adjacency, words = inputs

        """Compound vector with GNN."""
        fingerprint_vectors = self.embed_fingerprint(fingerprints)
        compound_vector = self.gnn(fingerprint_vectors, adjacency, layer_gnn)

        """Protein vector with attention-CNN."""
        word_vectors = self.embed_word(words)
        protein_vector = self.attention_cnn(compound_vector,
                                            word_vectors, layer_cnn)

        """Concatenate the above two vectors and output the interaction."""
        cat_vector = torch.cat((compound_vector, protein_vector), 1)
        for j in range(layer_output):
            cat_vector = torch.relu(self.W_out[j](cat_vector))
        interaction = self.W_interaction(cat_vector)

        return F.softmax(interaction, dim=1)

class CPIDataset(Dataset):
    
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        return self.data[index]     


def load_tensor(file_name, dtype):
    return [dtype(d).to(device) for d in np.load(file_name + '.npy', allow_pickle=True)]


def load_pickle(file_name):
    with open(file_name, 'rb') as f:
        return pickle.load(f)


def train_step(model, batch):
    model.train()
    compounds, adjacencies, proteins, true_y = batch
    pred_y = model(compounds, adjacencies, proteins)
    loss = F.binary_cross_entropy(pred_y, true_y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

def val_step(model, batch):
    model.eval()
    compounds, adjacencies, proteins, true_y = batch
    pred_y = model(compounds, adjacencies, proteins)
    loss = F.binary_cross_entropy(pred_y, true_y)
    return loss.item()


if __name__ == "__main__":

    """Hyperparameters."""
    (DATASET, radius, ngram, dim, layer_gnn, window, layer_cnn, layer_output,
     lr, lr_decay, decay_interval, weight_decay, iteration,
     setting) = sys.argv[1:]
    (dim, layer_gnn, window, layer_cnn, layer_output, decay_interval,
     iteration) = map(int, [dim, layer_gnn, window, layer_cnn, layer_output,
                            decay_interval, iteration])
    lr, lr_decay, weight_decay = map(float, [lr, lr_decay, weight_decay])

    """CPU or GPU."""
    if torch.cuda.is_available():
        device = torch.device('cuda:1')
        print('The code uses GPU...')
    else:
        device = torch.device('cpu')
        print('The code uses CPU!!!')

    """Load preprocessed data."""
    dir_input = ('../dataset/' + DATASET + '/input/'
                 'radius' + radius + '_ngram' + ngram + '/')
    compounds = load_tensor(dir_input + 'compounds', torch.LongTensor)
    adjacencies = load_tensor(dir_input + 'adjacencies', torch.FloatTensor)
    proteins = load_tensor(dir_input + 'proteins', torch.LongTensor)
    interactions = load_tensor(dir_input + 'interactions', torch.LongTensor)
    fingerprint_dict = load_pickle(dir_input + 'fingerprint_dict.pickle')
    word_dict = load_pickle(dir_input + 'word_dict.pickle')
    n_fingerprint = len(fingerprint_dict)
    n_word = len(word_dict)

    """Create a dataset and split it into train/dev/test."""
    dataset = list(zip(compounds, adjacencies, proteins, interactions))
    
    count = len(dataset)
    train_len = int(count * 0.8)
    val_len = count - train_len
    train_set, val_set = random_split(dataset, [train_len, val_len])
    trainloader = DataLoader(train_set, batch_size=1, shuffle=True)
    valloader = DataLoader(val_set, batch_size=1, shuffle=False)


    """Set a model."""
    model = CPIPrediction().to(device)

    """Start training."""
    print('Training...')
    optimizer = optim.Adam(model.parameters(), lr=0.1)
    lr_scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.99)
    for epoch in range(iteration):
        for batch in tqdm(trainloader):
            loss = train_step(model, batch)
# molecule embedding # protein embedding # # return torch.unsqueeze(torch.sum(xs, 0), 0) The attention mechanism is applied to the last layer of CNN. # return torch.unsqueeze(torch.sum(ys, 0), 0) # fingerprints, adjacency, words = inputs Compound vector with GNN. Protein vector with attention-CNN. Concatenate the above two vectors and output the interaction. Hyperparameters. CPU or GPU. Load preprocessed data. Create a dataset and split it into train/dev/test. Set a model. Start training.