step1_gcnn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb162838",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import package\n",
    "import numpy as np\n",
    "import random\n",
    "from numpy.core.numeric import ones_like\n",
    "from torch_geometric.data import Data, DataLoader\n",
    "import torch, os\n",
    "from tqdm import tqdm\n",
    "from sklearn.metrics import precision_recall_curve\n",
    "from sklearn.metrics import auc\n",
    "from sklearn.metrics import roc_auc_score\n",
    "from torch_geometric.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ab3f8128",
   "metadata": {},
   "outputs": [],
   "source": [
    "TF = 'ctcf'\n",
    "DATAPATH = '/DATA/yogesh/encodeDream/%s/'%(TF)\n",
    "POSITIVE = DATAPATH+'processed_out/positive/'\n",
    "NEGATIVE = DATAPATH+'processed_out/negative/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "389728ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import Linear\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCNConv,GENConv\n",
    "from torch_geometric.nn import global_mean_pool\n",
    "\n",
    "class GCN(torch.nn.Module):\n",
    "    def __init__(self, hidden_channels):\n",
    "        super(GCN, self).__init__()\n",
    "        torch.manual_seed(5)\n",
    "        # self.conv1 = GCNConv(train_loader.dataset[0].x.shape[1], hidden_channels)\n",
    "        # self.conv2 = GCNConv(hidden_channels, hidden_channels)\n",
    "        # self.conv3 = GCNConv(hidden_channels, hidden_channels)\n",
    "        self.conv1 = GENConv(8, hidden_channels)\n",
    "        self.conv2 = GENConv(hidden_channels, hidden_channels)\n",
    "        self.conv3 = GENConv(hidden_channels, hidden_channels)\n",
    "        self.lin = Linear(hidden_channels, 2)\n",
    "        \n",
    "    def forward(self, x, edge_index, batch):\n",
    "        # Get node embeeding\n",
    "        x = self.conv1(x, edge_index)\n",
    "        x = x.relu()\n",
    "        x = self.conv2(x, edge_index)\n",
    "        x = x.relu()\n",
    "        x = self.conv3(x, edge_index)\n",
    "        \n",
    "        # Readout alayer\n",
    "        x = global_mean_pool(x, batch)\n",
    "        \n",
    "        # Out layer\n",
    "        # x = F.dropout(x, p=DROPOUT, training=self.training)\n",
    "        x = self.lin(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8938bb48",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda:1\n",
      "Current cuda device ID: 1\n",
      "Current cuda device name: Tesla V100-PCIE-32GB\n",
      "GCN(\n",
      "  (conv1): GENConv(8, 64, aggr=softmax)\n",
      "  (conv2): GENConv(64, 64, aggr=softmax)\n",
      "  (conv3): GENConv(64, 64, aggr=softmax)\n",
      "  (lin): Linear(in_features=64, out_features=2, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda:1\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)\n",
    "torch.cuda.set_device(1)\n",
    "print('Current cuda device ID:',torch.cuda.current_device())\n",
    "print('Current cuda device name:', torch.cuda.get_device_name())\n",
    "\n",
    "# Train/test\n",
    "model = GCN(hidden_channels=64)\n",
    "model.to(device)\n",
    "print(model)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)\n",
    "loss_func = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "def to_device(data, device):\n",
    "    return data.to(device, non_blocking=True)\n",
    "\n",
    "def train():\n",
    "    model.train()\n",
    "    for data in train_loader:\n",
    "        a = to_device(data.x, device)\n",
    "        b = to_device(data.edge_index, device)\n",
    "        c = to_device(data.batch, device)\n",
    "        d = to_device(data.y, device)\n",
    "        \n",
    "        out = model(a, b, c)\n",
    "        #print(out)\n",
    "        \n",
    "        loss = loss_func(out, d)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "    \n",
    "def test(loader):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    aupr, auroc, loss=[], [], []\n",
    "    for data in loader:\n",
    "        a = to_device(data.x, device)\n",
    "        b = to_device(data.edge_index, device)\n",
    "        c = to_device(data.batch, device)\n",
    "        d = to_device(data.y, device)\n",
    "        \n",
    "        out = model(a, b, c)\n",
    "        prob = F.softmax(out, dim=1)\n",
    "        \n",
    "        # AUPR\n",
    "        precision, recall, thresholds = precision_recall_curve(d.detach().cpu().clone().numpy(), prob[:,1].detach().cpu().clone().numpy())\n",
    "        aupr.append(auc(recall, precision))\n",
    "        \n",
    "        # AUROC\n",
    "        try:\n",
    "            auroc.append(roc_auc_score(d.detach().cpu().clone().numpy(), prob[:,1].detach().cpu().clone().numpy()))\n",
    "        except:\n",
    "            pass\n",
    "        \n",
    "        # LOSS\n",
    "        loss_ = loss_func(out, d)\n",
    "        loss.append(loss_.detach().cpu().clone().numpy())\n",
    "        \n",
    "        # print(prob)\n",
    "        \n",
    "        pred = out.argmax(dim=1)\n",
    "        correct += int((pred==d).sum())\n",
    "        \n",
    "    return correct/len(loader.dataset), np.nanmean(aupr), np.nanmean(auroc), np.nanmean(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94c34cb8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch:001\n",
      "Train::ACC:0.6686409723281447 AUPR:0.7343817641450936 AUROC:0.7343336758309548 LOSS:0.6028475761413574\n",
      "Test ::ACC:0.6677456023736585 AUPR:0.7351739265459599 AUROC:0.7332531018809447 LOSS:0.6036773920059204\n",
      "\n",
      "Epoch:002\n",
      "Train::ACC:0.6733227901450625 AUPR:0.7386024878334434 AUROC:0.7408843569228478 LOSS:0.5972746014595032\n",
      "Test ::ACC:0.6718044557475804 AUPR:0.7358877390732771 AUROC:0.7382350877822781 LOSS:0.5995912551879883\n",
      "\n",
      "Epoch:003\n",
      "Train::ACC:0.6744563166809829 AUPR:0.7389475306947504 AUROC:0.741745697827416 LOSS:0.5970451235771179\n",
      "Test ::ACC:0.6721384120378399 AUPR:0.7351387072121461 AUROC:0.7390835576409297 LOSS:0.5989888906478882\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Import positeve dataset\n",
    "dum2 = torch.load(POSITIVE+os.listdir(POSITIVE)[0])\n",
    "np.seterr(divide='ignore', invalid='ignore')\n",
    "\n",
    "# Training\n",
    "train_acc_, test_acc_, train_aupr_, test_aupr_, train_auroc_, test_auroc_ = [], [], [], [], [], []\n",
    "for epoch in range(1, len(os.listdir(NEGATIVE))):\n",
    "    # Get data\n",
    "    dum1 = torch.load(NEGATIVE+os.listdir(NEGATIVE)[epoch])\n",
    "    DATALIST = dum1+dum2\n",
    "\n",
    "    # Split data\n",
    "    torch.manual_seed(11)\n",
    "    random.shuffle(DATALIST)\n",
    "\n",
    "    split_n = int(0.8*(len(DATALIST)))\n",
    "\n",
    "    train_dataset = DATALIST[:split_n]\n",
    "    test_dataset = DATALIST[split_n:]\n",
    "    \n",
    "    # Prepare dataloader\n",
    "    train_loader = DataLoader(train_dataset, batch_size=512, shuffle=True)\n",
    "    test_loader = DataLoader(test_dataset, batch_size=512, shuffle=True)\n",
    "    \n",
    "    # Start training\n",
    "    train()\n",
    "    train_acc, train_aupr, train_auroc, train_loss = test(train_loader)\n",
    "    test_acc, test_aupr, test_auroc, test_loss = test(test_loader)\n",
    "    \n",
    "    train_acc_.append(train_acc)\n",
    "    test_acc_.append(test_acc)\n",
    "    train_aupr_.append(train_aupr)\n",
    "    test_aupr_.append(test_aupr)\n",
    "    train_auroc_.append(train_auroc)\n",
    "    test_auroc_.append(test_auroc)\n",
    "    \n",
    "    print(f'Epoch:{epoch:03d}')\n",
    "    print(f'Train::ACC:{train_acc} AUPR:{train_aupr} AUROC:{train_auroc} LOSS:{train_loss}')\n",
    "    print(f'Test ::ACC:{test_acc} AUPR:{test_aupr} AUROC:{test_auroc} LOSS:{test_loss}')\n",
    "    print()\n",
    "    \n",
    "    # print(f'Epoch:{epoch:03d},Train acc:{train_acc:.4f},Train aupr:{train_aupr:.4f},Train auroc:{train_auroc:.4f},Test acc:{test_acc:.4f},Test aupr:{test_aupr:.4f},Test auroc:{test_auroc:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b3c464f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "mpl.rcParams['axes.linewidth'] = 1.2\n",
    "\n",
    "fig, ax = plt.subplots(1,3, figsize=(9,3))\n",
    "ax[0].plot(train_acc_, color='red', marker='o', linewidth=3, markersize=3, label='train')\n",
    "ax[0].plot(test_acc_, color='black', marker='o', linewidth=3, markersize=3, label='test')\n",
    "ax[0].legend()\n",
    "ax[0].set_ylim([0.5,1])\n",
    "ax[0].set_ylabel('Accuracy', fontsize='14')\n",
    "ax[0].set_xlabel('Epoch', fontsize='14')\n",
    "#ax[0].grid()\n",
    "\n",
    "ax[1].plot(train_aupr_, color='red', marker='o', linewidth=3, markersize=3, label='train')\n",
    "ax[1].plot(test_aupr_, color='black', marker='o', linewidth=3, markersize=3, label='test')\n",
    "ax[1].set_ylim([0,1])\n",
    "ax[1].legend()\n",
    "ax[1].set_ylabel('AUPR', fontsize='14')\n",
    "ax[1].set_xlabel('Epoch', fontsize='14')\n",
    "#ax[1].grid()\n",
    "\n",
    "ax[2].plot(train_auroc_, color='red', marker='o', linewidth=3, markersize=3, label='train')\n",
    "ax[2].plot(test_auroc_, color='black', marker='o', linewidth=3, markersize=3, label='test')\n",
    "ax[2].set_ylim([0,1])\n",
    "ax[2].legend()\n",
    "ax[2].set_ylabel('AUROC', fontsize='14')\n",
    "ax[2].set_xlabel('Epoch', fontsize='14')\n",
    "\n",
    "#ax[2].grid()\n",
    "\n",
    "fig.suptitle(TF.upper()+': Performance (Internal validation)', fontsize=14)\n",
    "plt.tight_layout()\n",
    "\n",
    "# plt.savefig(PATH_+'processed_out/'+'10lac_'+TF+'_'+str(kmer)+'_'+str(DROPOUT)+'.png', bbbox_inches =\"tight\", dpi=600)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}