LEP-AD.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/dagaa/anaconda3/envs/graphdta/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import sys, os\n",
    "from scipy import stats\n",
    "import wandb\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "from random import shuffle\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import OrderedDict\n",
    "import json,pickle\n",
    "from rdkit import RDLogger\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import MolFromSmiles\n",
    "from lifelines.utils import concordance_index\n",
    "\n",
    "# Model\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch_geometric.data import InMemoryDataset, DataLoader, Batch\n",
    "import torch.nn.functional as F\n",
    "import torch_geometric\n",
    "from torch_geometric.nn import TransformerConv,GATConv, GCNConv,global_max_pool as gmp, global_add_pool as gap,global_mean_pool as gep,global_sort_pool\n",
    "from torch_geometric.utils import dropout_adj\n",
    "from torch_geometric import data as DATA\n",
    "import torchvision.transforms as T\n",
    "import warnings\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "warnings.filterwarnings(\"ignore\")\n",
    "wandb.init(project=\"LEP-AD\", entity=\"daga06\")\n",
    "\n",
    "dataset = 'davis'\n",
    "# load dataset\n",
    "dataset_path = 'data/' + dataset + '/'\n",
    "\n",
    "ligands = json.load(open(dataset_path + 'drug_dict.txt'), object_pairs_hook=OrderedDict)\n",
    "proteins = json.load(open(dataset_path + 'protein_dict.txt'), object_pairs_hook=OrderedDict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "drugs = []\n",
    "drug_smiles = []\n",
    "\n",
    "RDLogger.DisableLog('rdApp.*')\n",
    "# smiles\n",
    "for d in ligands.keys():\n",
    "    lg = Chem.MolToSmiles(Chem.MolFromSmiles(ligands[d]), isomericSmiles=True)\n",
    "    drugs.append(lg)\n",
    "    drug_smiles.append(ligands[d])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# one ont encoding\n",
    "def one_of_k_encoding(x, allowable_set):\n",
    "    if x not in allowable_set:\n",
    "        # print(x)\n",
    "        raise Exception('input {0} not in allowable set{1}:'.format(x, allowable_set))\n",
    "    return list(map(lambda s: x == s, allowable_set))\n",
    "\n",
    "\n",
    "def one_of_k_encoding_unk(x, allowable_set):\n",
    "    '''Maps inputs not in the allowable set to the last element.'''\n",
    "    if x not in allowable_set:\n",
    "        x = allowable_set[-1]\n",
    "    return list(map(lambda s: x == s, allowable_set))\n",
    "    \n",
    "# mol atom feature for mol graph\n",
    "def atom_features(atom):\n",
    "    # 44 +11 +11 +11 +1\n",
    "    return np.array(one_of_k_encoding_unk(atom.GetSymbol(),\n",
    "                                          ['C', 'N', 'O', 'S', 'F', 'Si', 'P', 'Cl', 'Br', 'Mg', 'Na', 'Ca', 'Fe', 'As',\n",
    "                                           'Al', 'I', 'B', 'V', 'K', 'Tl', 'Yb', 'Sb', 'Sn', 'Ag', 'Pd', 'Co', 'Se',\n",
    "                                           'Ti', 'Zn', 'H', 'Li', 'Ge', 'Cu', 'Au', 'Ni', 'Cd', 'In', 'Mn', 'Zr', 'Cr',\n",
    "                                           'Pt', 'Hg', 'Pb', 'X']) +\n",
    "                    one_of_k_encoding(atom.GetDegree(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +\n",
    "                    one_of_k_encoding_unk(atom.GetTotalNumHs(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +\n",
    "                    one_of_k_encoding_unk(atom.GetImplicitValence(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) +\n",
    "                    [atom.GetIsAromatic()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def smile_to_graph(smile):\n",
    "    mol = Chem.MolFromSmiles(smile)\n",
    "\n",
    "    c_size = mol.GetNumAtoms()\n",
    "\n",
    "    features = []\n",
    "    for atom in mol.GetAtoms():\n",
    "        feature = atom_features(atom)\n",
    "        features.append(feature / sum(feature))\n",
    "\n",
    "    edges = []\n",
    "    for bond in mol.GetBonds():\n",
    "        edges.append([bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()])\n",
    "    g = nx.Graph(edges).to_directed()\n",
    "    edge_index = []\n",
    "    mol_edge_weight=[]\n",
    "    mol_adj = np.zeros((c_size, c_size))\n",
    "    for e1, e2 in g.edges:\n",
    "        mol_adj[e1, e2] = 1\n",
    "        # edge_index.append([e1, e2])\n",
    "    mol_adj += np.matrix(np.eye(mol_adj.shape[0]))\n",
    "    index_row, index_col = np.where(mol_adj >= 0.5)\n",
    "    for i, j in zip(index_row, index_col):\n",
    "        edge_index.append([i, j])\n",
    "        mol_edge_weight.append([1])\n",
    "    # print('smile_to_graph')\n",
    "    # print(np.array(features).shape)\n",
    "    return c_size, features, edge_index,mol_edge_weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 14min\n",
    "compound_iso_smiles = drugs\n",
    "if not os.path.exists('data/'+dataset+'/temp/smile_graph.pickle'):\n",
    "    os.mkdir('data/'+dataset+'/temp')\n",
    "    # create smile graph\n",
    "    smile_graph = {}\n",
    "    i = 0\n",
    "    for smile in compound_iso_smiles:\n",
    "        g = smile_to_graph(smile)\n",
    "        smile_graph[smile] = g\n",
    "        if i%1000==0:\n",
    "            print(i)\n",
    "        i+=1\n",
    "    with open('data/'+ dataset +'/temp/smile_graph.pickle','wb') as handle:\n",
    "        pickle.dump(smile_graph, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "        \n",
    "with open('data/'+ dataset +'/temp/smile_graph.pickle','rb') as f:\n",
    "    smile_graph = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('data/'+dataset+'/temp/protein_rep.pickle'):\n",
    "    esm_emb_path = 'data/' + dataset + '/proteins_emb_esm2/UniRef50_'\n",
    "    protein_dict = json.load(open('data/davis/protein_dict.txt'))\n",
    "    target_reps_dict = {}\n",
    "    i=0\n",
    "    for key in proteins.keys():\n",
    "        target_reps_dict[protein_dict[key]] = torch.load(esm_emb_path+ key + '.pt')['mean_representations'][36]\n",
    "        if i%1000==0:\n",
    "            print(i)\n",
    "        i+=1\n",
    "\n",
    "    with open('data/'+ dataset +'/temp/protein_rep.pickle','wb') as handle:\n",
    "        pickle.dump(target_reps_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "\n",
    "with open('data/'+ dataset +'/temp/protein_rep.pickle', 'rb') as handle:\n",
    "    target_reps_dict = pickle.load(handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# initialize the dataset\n",
    "class DTADataset(InMemoryDataset):\n",
    "    def __init__(self, root='/tmp', dataset='davis',\n",
    "                 xd=None, y=None, transform= None,\n",
    "                 pre_transform=None, smile_graph=None, target_key=None, target_rep=None):\n",
    "        super(DTADataset, self).__init__(root, transform, pre_transform)\n",
    "        self.dataset = dataset\n",
    "        self.process(xd, target_key, y, smile_graph, target_rep)\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        pass\n",
    "        # return ['some_file_1', 'some_file_2', ...]\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self):\n",
    "        return [self.dataset + '_data_mol.pt', self.dataset + '_data_pro.pt']\n",
    "\n",
    "    def _process(self):\n",
    "        if not os.path.exists(self.processed_dir):\n",
    "            os.makedirs(self.processed_dir)\n",
    "\n",
    "    def process(self, xd, target_key, y, smile_graph, target_rep):\n",
    "        assert (len(xd) == len(target_key) and len(xd) == len(y)), 'The three lists must be the same length!'\n",
    "        data_list_mol = []\n",
    "        data_list_pro = []\n",
    "        data_len = len(xd)\n",
    "        for i in range(data_len):\n",
    "            entity1 = xd[i]\n",
    "\n",
    "            labels = y[i]\n",
    "            if entity1 in smile_graph.keys():\n",
    "                c_size, features, edge_index,edge_weight = smile_graph[entity1]\n",
    "            else:\n",
    "                c_size, features, edge_index,edge_weight = smile_to_graph(entity1)\n",
    "            GCNData_mol = DATA.Data(x=torch.Tensor(np.array(features)),\n",
    "                                    edge_index=torch.LongTensor(edge_index).transpose(1, 0),\n",
    "                                    y=torch.FloatTensor([labels])\n",
    "                                    )\n",
    "            GCNData_mol.__setitem__('c_size', torch.LongTensor([c_size]))\n",
    "            data_list_mol.append(GCNData_mol)\n",
    "\n",
    "            data_list_pro.append(torch.Tensor(target_rep[target_key[i]]))\n",
    "            if i%10000==0:\n",
    "                print(i)\n",
    "            \n",
    "        if self.pre_filter is not None:\n",
    "            data_list_mol = [data for data in data_list_mol if self.pre_filter(data)]\n",
    "        if self.pre_transform is not None:\n",
    "            data_list_mol = [self.pre_transform(data) for data in data_list_mol]\n",
    "        self.data_mol = data_list_mol\n",
    "        self.data_pro = data_list_pro\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data_mol)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.data_mol[idx], self.data_pro[idx]\n",
    "        \n",
    "def collate(batch):\n",
    "    graphs = Batch.from_data_list([item[0] for item in batch])\n",
    "    tensors = [item[1] for item in batch]\n",
    "    tensors = torch.stack(tensors)\n",
    "\n",
    "    return graphs,tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train_fold = pd.read_csv('data/' + dataset + '/'+ dataset+'_' + 'train' + '.csv')\n",
    "train_drugs, train_prot_keys, train_Y = list(df_train_fold['compound_iso_smiles']), list(df_train_fold['target_sequence']), list(df_train_fold['affinity'])\n",
    "train_drugs, train_prot_keys, train_Y = np.asarray(train_drugs), np.asarray(train_prot_keys), np.asarray(train_Y)\n",
    "\n",
    "train_dataset = DTADataset(root='data', dataset=dataset + '_' + 'train', xd=train_drugs, target_key=train_prot_keys,\n",
    "                            y=train_Y, smile_graph=smile_graph, target_rep=target_reps_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "cuda_name = 'cuda:0'\n",
    "print('cuda_name:', cuda_name)\n",
    "fold = [0, 1, 2, 3, 4][0]\n",
    "cross_validation_flag = True\n",
    "\n",
    "TRAIN_BATCH_SIZE = 512\n",
    "TEST_BATCH_SIZE = 512\n",
    "LR = 0.0005\n",
    "NUM_EPOCHS = 3000\n",
    "from sklearn.model_selection import train_test_split\n",
    "train_data,valid_data=train_test_split(train_dataset,shuffle=True,test_size=0.2)\n",
    "train_loader = torch.utils.data.DataLoader(train_data, batch_size=TRAIN_BATCH_SIZE, shuffle=True,num_workers=8,\n",
    "                                            collate_fn=collate)\n",
    "valid_loader = torch.utils.data.DataLoader(valid_data, batch_size=TEST_BATCH_SIZE, shuffle=False,num_workers=4,\n",
    "                                            collate_fn=collate)\n",
    "                                            \n",
    "# next(iter(train_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GCN based model\n",
    "class GNNNet(torch.nn.Module):\n",
    "    def __init__(self, n_output=1, num_features_mol=78, output_dim=128, dropout=0.2):\n",
    "        super(GNNNet, self).__init__()\n",
    "\n",
    "        print('GNNNet Loaded')\n",
    "        self.n_output = n_output\n",
    "        self.mol_conv1 = TransformerConv(num_features_mol, num_features_mol,heads=4)\n",
    "        self.mol_conv2 = TransformerConv(4*num_features_mol, num_features_mol * 2,heads=2)\n",
    "        self.mol_conv3 = TransformerConv(num_features_mol * 4, num_features_mol * 4,heads=1)\n",
    "        self.mol_fc_g1 = torch.nn.Linear(num_features_mol * 4, 1024)\n",
    "        self.mol_fc_g2 = torch.nn.Linear(1024, output_dim)\n",
    "\n",
    "        self.relu = nn.ReLU()\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "        # combined layers\n",
    "        self.fc1 = nn.Linear(2688, 1024)\n",
    "        self.fc2 = nn.Linear(1024, 512)\n",
    "        self.out = nn.Linear(512, self.n_output)\n",
    "\n",
    "    def forward(self, data_mol, data_pro):\n",
    "        mol_x, mol_edge_index, mol_batch = data_mol.x, data_mol.edge_index, data_mol.batch\n",
    "\n",
    "        x = self.mol_conv1(mol_x, mol_edge_index)\n",
    "        x = self.relu(x)\n",
    "\n",
    "        x = self.mol_conv2(x, mol_edge_index)\n",
    "        x = self.relu(x)\n",
    "\n",
    "        x = self.mol_conv3(x, mol_edge_index)\n",
    "        x = self.relu(x)\n",
    "\n",
    "        x = gep(x, mol_batch)  # global pooling\n",
    "\n",
    "        # flatten\n",
    "        x = self.relu(self.mol_fc_g1(x))\n",
    "        x = self.dropout(x)\n",
    "        x = self.mol_fc_g2(x)\n",
    "        x = self.dropout(x)\n",
    "\n",
    "        # concat\n",
    "        xc = torch.cat((x, data_pro), 1)\n",
    "        # add some dense layers\n",
    "        xc = self.fc1(xc)\n",
    "        xc = self.relu(xc)\n",
    "        xc = self.dropout(xc)\n",
    "        xc = self.fc2(xc)\n",
    "        xc = self.relu(xc)\n",
    "        xc = self.dropout(xc)\n",
    "        out = self.out(xc)\n",
    "        return out\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Learning rate: ', LR)\n",
    "print('Epochs: ', NUM_EPOCHS)\n",
    "\n",
    "models_dir = 'models'\n",
    "results_dir = 'results'\n",
    "\n",
    "if not os.path.exists(models_dir):\n",
    "    os.makedirs(models_dir)\n",
    "\n",
    "if not os.path.exists(results_dir):\n",
    "    os.makedirs(results_dir)\n",
    "\n",
    "# Main program: iterate over different datasets\n",
    "result_str = ''\n",
    "USE_CUDA = torch.cuda.is_available()\n",
    "device = torch.device(cuda_name if USE_CUDA else 'cpu')\n",
    "model = GNNNet()\n",
    "model.to(device)\n",
    "\n",
    "model_st = GNNNet.__name__\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=LR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# training function at each epoch\n",
    "scaler = torch.cuda.amp.GradScaler()\n",
    "def train(model, device, train_loader, optimizer, epoch):\n",
    "    print('Training on {} samples...'.format(len(train_loader.dataset)))\n",
    "    model.train()\n",
    "    LOG_INTERVAL = 100\n",
    "    TRAIN_BATCH_SIZE = 512\n",
    "    loss_fn = torch.nn.MSELoss()\n",
    "    \n",
    "    for batch_idx, data in enumerate(train_loader):\n",
    "        data_mol = data[0].to(device)\n",
    "        data_pro = data[1].to(device)\n",
    "        optimizer.zero_grad()\n",
    "        with torch.cuda.amp.autocast():\n",
    "            output = model(data_mol, data_pro)\n",
    "            labels = data_mol.y.view(-1, 1)\n",
    "            loss = loss_fn(output, labels)\n",
    "        scaler.scale(loss).backward()\n",
    "        wandb.log({\"loss per batch\": loss})\n",
    "        #optimizer.step()\n",
    "        scaler.step(optimizer)\n",
    "        scaler.update()\n",
    "        if batch_idx % LOG_INTERVAL == 0:\n",
    "            print('Train epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(epoch,\n",
    "                                                                           batch_idx * TRAIN_BATCH_SIZE,\n",
    "                                                                           len(train_loader.dataset),\n",
    "                                                                           100. * batch_idx / len(train_loader),\n",
    "                                                                           loss.item()))\n",
    "\n",
    "# predict\n",
    "def predicting(model, device, loader):\n",
    "    model.eval()\n",
    "    total_preds = torch.Tensor()\n",
    "    total_labels = torch.Tensor()\n",
    "    print('Make prediction for {} samples...'.format(len(loader.dataset)))\n",
    "    with torch.no_grad():\n",
    "        for data in loader:\n",
    "            data_mol = data[0].to(device)\n",
    "            data_pro = data[1].to(device)\n",
    "            output = model(data_mol, data_pro)\n",
    "            labels = data_mol.y.view(-1, 1)\n",
    "            total_preds = torch.cat((total_preds, output.cpu()), 0)\n",
    "\n",
    "            total_labels = torch.cat((total_labels, labels.cpu()), 0)\n",
    "    return total_labels.numpy().flatten(), total_preds.numpy().flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def load_model(model_path):\n",
    "    model = torch.load(model_path)\n",
    "    return model\n",
    "\n",
    "\n",
    "import time\n",
    "def calculate_metrics(Y, P, dataset='davis'):\n",
    "    # # aupr = get_aupr(Y, P)\n",
    "    # t = time.time()\n",
    "    \n",
    "    \n",
    "    # cindex = get_cindex(Y, P)\n",
    "    # print(cindex)\n",
    "    # print(concordance_index(Y, P))  # DeepDTAget_cindex(Y, P)\n",
    "    \n",
    "    \n",
    "    cindex2 = concordance_index(Y, P)  # GraphDTA\n",
    "    rm2 = get_rm2(Y, P)  # DeepDTA\n",
    "    mse = get_mse(Y, P)\n",
    "    # pearson = get_pearson(Y, P)\n",
    "    # spearman = get_spearman(Y, P)\n",
    "    # rmse = get_rmse(Y, P)\n",
    "\n",
    "    print('metrics for ', dataset)\n",
    "    # print('aupr:', aupr)\n",
    "    # print('cindex:', cindex)\n",
    "    print('cindex2', cindex2)\n",
    "    print('rm2:', rm2)\n",
    "    print('mse:', mse)\n",
    "    result_str = ''\n",
    "    result_str += dataset + '\\r\\n'\n",
    "    result_str += ' ' + ' mse:' + str(mse) + ' ' + ' '+ ' ' + 'ci:' + str(cindex2)\n",
    "    print(result_str)\n",
    "    return mse,cindex2,rm2\n",
    "\n",
    "def plot_density(Y, P, fold=0, dataset='davis'):\n",
    "    plt.figure(figsize=(10, 5))\n",
    "    plt.grid(linestyle='--')\n",
    "    ax = plt.gca()\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "\n",
    "    plt.scatter(P, Y, color='blue', s=40)\n",
    "    plt.title('density of ' + dataset, fontsize=30, fontweight='bold')\n",
    "    plt.xlabel('predicted', fontsize=30, fontweight='bold')\n",
    "    plt.ylabel('measured', fontsize=30, fontweight='bold')\n",
    "    # plt.xlim(0, 21)\n",
    "    # plt.ylim(0, 21)\n",
    "    if dataset == 'davis':\n",
    "        plt.plot([5, 11], [5, 11], color='black')\n",
    "    else:\n",
    "        plt.plot([6, 16], [6, 16], color='black')\n",
    "    # plt.legend()\n",
    "    plt.legend(loc=0, numpoints=1)\n",
    "    leg = plt.gca().get_legend()\n",
    "    ltext = leg.get_texts()\n",
    "    plt.setp(ltext, fontsize=12, fontweight='bold')\n",
    "    plt.savefig(os.path.join('results', dataset + '_' + str(fold) + '.png'), dpi=500, bbox_inches='tight')\n",
    "\n",
    "\n",
    "    # plot_density(Y, P, fold, dataset)\n",
    "import numpy as np\n",
    "import subprocess\n",
    "from math import sqrt\n",
    "from sklearn.metrics import average_precision_score\n",
    "from scipy import stats\n",
    "\n",
    "\n",
    "def get_aupr(Y, P, threshold=7.0):\n",
    "    # print(Y.shape,P.shape)\n",
    "    Y = np.where(Y >= 7.0, 1, 0)\n",
    "    P = np.where(P >= 7.0, 1, 0)\n",
    "    aupr = average_precision_score(Y, P)\n",
    "    return aupr\n",
    "\n",
    "\n",
    "def get_cindex(Y, P):\n",
    "    summ = 0\n",
    "    pair = 0\n",
    "\n",
    "    for i in range(1, len(Y)):\n",
    "        for j in range(0, i):\n",
    "            if i is not j:\n",
    "                if (Y[i] > Y[j]):\n",
    "                    pair += 1\n",
    "                    summ += 1 * (P[i] > P[j]) + 0.5 * (P[i] == P[j])\n",
    "\n",
    "    if pair != 0:\n",
    "        return summ / pair\n",
    "    else:\n",
    "        return 0\n",
    "\n",
    "\n",
    "def r_squared_error(y_obs, y_pred):\n",
    "    y_obs = np.array(y_obs)\n",
    "    y_pred = np.array(y_pred)\n",
    "    y_obs_mean = [np.mean(y_obs) for y in y_obs]\n",
    "    y_pred_mean = [np.mean(y_pred) for y in y_pred]\n",
    "\n",
    "    mult = sum((y_pred - y_pred_mean) * (y_obs - y_obs_mean))\n",
    "    mult = mult * mult\n",
    "\n",
    "    y_obs_sq = sum((y_obs - y_obs_mean) * (y_obs - y_obs_mean))\n",
    "    y_pred_sq = sum((y_pred - y_pred_mean) * (y_pred - y_pred_mean))\n",
    "\n",
    "    return mult / float(y_obs_sq * y_pred_sq)\n",
    "\n",
    "\n",
    "def get_k(y_obs, y_pred):\n",
    "    y_obs = np.array(y_obs)\n",
    "    y_pred = np.array(y_pred)\n",
    "\n",
    "    return sum(y_obs * y_pred) / float(sum(y_pred * y_pred))\n",
    "\n",
    "\n",
    "def squared_error_zero(y_obs, y_pred):\n",
    "    k = get_k(y_obs, y_pred)\n",
    "\n",
    "    y_obs = np.array(y_obs)\n",
    "    y_pred = np.array(y_pred)\n",
    "    y_obs_mean = [np.mean(y_obs) for y in y_obs]\n",
    "    upp = sum((y_obs - (k * y_pred)) * (y_obs - (k * y_pred)))\n",
    "    down = sum((y_obs - y_obs_mean) * (y_obs - y_obs_mean))\n",
    "\n",
    "    return 1 - (upp / float(down))\n",
    "\n",
    "\n",
    "def get_rm2(ys_orig, ys_line):\n",
    "    r2 = r_squared_error(ys_orig, ys_line)\n",
    "    r02 = squared_error_zero(ys_orig, ys_line)\n",
    "\n",
    "    return r2 * (1 - np.sqrt(np.absolute((r2 * r2) - (r02 * r02))))\n",
    "\n",
    "\n",
    "def get_rmse(y, f):\n",
    "    rmse = sqrt(((y - f) ** 2).mean(axis=0))\n",
    "    return rmse\n",
    "\n",
    "\n",
    "def get_mse(y, f):\n",
    "    mse = ((y - f) ** 2).mean(axis=0)\n",
    "    return mse\n",
    "\n",
    "\n",
    "def get_pearson(y, f):\n",
    "    rp = np.corrcoef(y, f)[0, 1]\n",
    "    return rp\n",
    "\n",
    "\n",
    "def get_spearman(y, f):\n",
    "    rs = stats.spearmanr(y, f)[0]\n",
    "    return rs\n",
    "\n",
    "\n",
    "def get_ci(y, f):\n",
    "    ind = np.argsort(y)\n",
    "    y = y[ind]\n",
    "    f = f[ind]\n",
    "    i = len(y) - 1\n",
    "    j = i - 1\n",
    "    z = 0.0\n",
    "    S = 0.0\n",
    "    while i > 0:\n",
    "        while j >= 0:\n",
    "            if y[i] > y[j]:\n",
    "                z = z + 1\n",
    "                u = f[i] - f[j]\n",
    "                if u > 0:\n",
    "                    S = S + 1\n",
    "                elif u == 0:\n",
    "                    S = S + 0.5\n",
    "            j = j - 1\n",
    "        i = i - 1\n",
    "        j = i - 1\n",
    "    ci = S / z\n",
    "    return ci"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#188 minutes\n",
    "best_mse = 1000\n",
    "best_test_mse = 1000\n",
    "best_epoch = -1\n",
    "# model_file_name = 'models/model_' + model_st + '_' + dataset + '_' + str(fold) + '.model'\n",
    "\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    train(model, device, train_loader, optimizer, epoch + 1)\n",
    "    print('predicting for valid data')\n",
    "    G, P = predicting(model, device, valid_loader)\n",
    "    val_mse,val_ci,val_rm2 = calculate_metrics(G, P, dataset)\n",
    "    wandb.log({\"val_ci\": val_ci})\n",
    "    wandb.log({\"val_mse\": val_mse})\n",
    "    wandb.log({\"val_rm2\": val_rm2})\n",
    "\n",
    "    print('predicting for test data')\n",
    "    \n",
    "    if val_mse < best_mse:\n",
    "        best_mse = val_mse\n",
    "        best_epoch = epoch + 1\n",
    "        # torch.save(model.state_dict(), model_file_name)\n",
    "        print('rmse improved at epoch ', best_epoch, '; best_test_mse', best_mse, model_st, dataset, fold)\n",
    "    else:\n",
    "        print('No improvement since epoch ', best_epoch, '; best_test_mse', best_mse, model_st, dataset, fold)\n",
    "\n",
    "\n",
    "#reaching optimzation\n",
    "torch.save({\n",
    "    'epoch':epoch,\n",
    "    'model_state_dict':model.state_dict(),\n",
    "    'optimizer_state_dict':optimizer.state_dict()\n",
    "},'saved_models/esm-2-stitch.pt')\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test_fold = pd.read_csv('data/' + dataset + '/'+ dataset+'_' + 'test' + '.csv')\n",
    "test_drugs, test_prot_keys, test_Y = list(df_test_fold['compound_iso_smiles']), list(df_test_fold['target_sequence']), list(df_test_fold['affinity'])\n",
    "test_drugs, test_prot_keys, test_Y = np.asarray(test_drugs), np.asarray(test_prot_keys), np.asarray(test_Y)\n",
    "\n",
    "test_data = DTADataset(root='data', dataset=dataset + '_' + 'test', xd=test_drugs, target_key=test_prot_keys,\n",
    "                            y=test_Y, smile_graph=smile_graph, target_rep=target_reps_dict)\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=TEST_BATCH_SIZE, shuffle=False,num_workers=4,\n",
    "                                            collate_fn=collate)\n",
    "G, P = predicting(model, device, test_loader)\n",
    "test_mse,test_ci,test_rm2 = calculate_metrics(G, P, dataset)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "LEP-AD",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "8374b1f9cd6cf8b7bee55b14bca9022ee9d6a251b66f39376ae0f3b8e6f8508e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}