consnet.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/rohanjuneja/ARM-Simulator/blob/master/consnet.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iyND_I7A3oo7"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.datasets as datasets\n",
        "from torchvision.utils import save_image\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "import numpy as np\n",
        "import os\n",
        "import glob\n",
        "import PIL\n",
        "from PIL import Image\n",
        "from torch.utils import data as D\n",
        "from torch.utils.data.sampler import SubsetRandomSampler\n",
        "import random\n",
        "import pandas as pd\n",
        "import time\n",
        "import torch.autograd as autograd\n",
        "from torch.autograd import Variable"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jkQIZpTx3opA"
      },
      "outputs": [],
      "source": [
        "from matplotlib import pyplot as plt\n",
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cbf78QmW3opB"
      },
      "outputs": [],
      "source": [
        "# Define the NN architecture\n",
        "\n",
        "class ConsNet(nn.Module):\n",
        "    def __init__(self, hidden_dim, fc1_dim, fc2_dim):\n",
        "        super(ConsNet, self).__init__()\n",
        "        self.hidden_dim = hidden_dim\n",
        "\n",
        "        self.lstm = nn.LSTM(4, hidden_dim) \n",
        "        self.fc1 = nn.Linear(hidden_dim, fc1_dim)\n",
        "        self.dropout1 = nn.Dropout(p=0.2)\n",
        "        self.fc2 = nn.Linear(fc1_dim, fc2_dim)\n",
        "        self.dropout2 = nn.Dropout(p=0.1)\n",
        "        self.fc3 = nn.Linear(fc2_dim, 4)\n",
        "        self.dropout3 = nn.Dropout(p=0.1)\n",
        "        self.hidden_init_values = None\n",
        "        self.hidden = self.init_hidden()\n",
        "        nn.init.xavier_uniform_(self.fc1.weight)\n",
        "        nn.init.xavier_uniform_(self.fc2.weight)\n",
        "        nn.init.xavier_uniform_(self.fc3.weight)\n",
        "        \n",
        "    def init_hidden(self):\n",
        "        if self.hidden_init_values == None:\n",
        "            self.hidden_init_values = (autograd.Variable(torch.randn(1, 1, self.hidden_dim)),\n",
        "                                       autograd.Variable(torch.randn(1, 1, self.hidden_dim)))\n",
        "        return self.hidden_init_values\n",
        "\n",
        "    def forward(self, input_seq):\n",
        "        lstm_out, self.hidden = self.lstm(\n",
        "            input_seq.view(len(input_seq), 1, -1), self.hidden)\n",
        "        tmp1 = F.relu(self.dropout1(self.fc1(lstm_out.view(len(input_seq), -1))))\n",
        "        tmp2 = F.relu(self.dropout2(self.fc2(tmp1)))\n",
        "        _out = self.dropout3(self.fc3(tmp2))\n",
        "        x = _out\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UcOKQ5hQ3opD",
        "outputId": "24f0915f-be6f-4dcb-969b-f8650971d672"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "one-hot encoding for DNA bases\n",
            "A: [[1, 0, 0, 0]]\n",
            "C: [[0, 1, 0, 0]]\n",
            "G: [[0, 0, 1, 0]]\n",
            "T: [[0, 0, 0, 1]]\n"
          ]
        }
      ],
      "source": [
        "import random\n",
        "torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
        "\n",
        "bmap = {\"A\":0, \"C\":1, \"G\":2, \"T\":3}\n",
        "def one_hot(b):\n",
        "    t = [[0,0,0,0]]\n",
        "    i = bmap[b]\n",
        "    t[0][i] = 1\n",
        "    return t\n",
        "\n",
        "print(\"one-hot encoding for DNA bases\")\n",
        "print(\"A:\", one_hot(\"A\"))\n",
        "print(\"C:\", one_hot(\"C\"))\n",
        "print(\"G:\", one_hot(\"G\"))\n",
        "print(\"T:\", one_hot(\"T\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NEYaJhg-3opE"
      },
      "outputs": [],
      "source": [
        "def sim_error(seq, pi=0.05, pd=0.05, ps=0.01):\n",
        "    \"\"\"\n",
        "    Given an input sequence `seq`, generating another\n",
        "    sequence with errors. \n",
        "    pi: insertion error rate\n",
        "    pd: deletion error rate\n",
        "    ps: substitution error rate\n",
        "    \"\"\"\n",
        "    out_seq = []\n",
        "    for c in seq:\n",
        "        while 1:\n",
        "            r = random.uniform(0,1)\n",
        "            if r < pi:\n",
        "                out_seq.append(random.choice([\"A\",\"C\",\"G\",\"T\"]))\n",
        "            else:\n",
        "                break\n",
        "        r -= pi\n",
        "        if r < pd:\n",
        "            continue\n",
        "        r -= pd\n",
        "        if r < ps:\n",
        "            out_seq.append(random.choice([\"A\",\"C\",\"G\",\"T\"]))\n",
        "            continue\n",
        "        out_seq.append(c)\n",
        "    return \"\".join(out_seq)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G29Wk0Xc3opF"
      },
      "source": [
        "### Generate training data and train the model "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UjqWDQJC3opH",
        "outputId": "0685fcb8-35d0-4fd9-f486-4a1232ab5ee0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "CTGACACGAGGACCCGTTTCGACCTGCAGCGTGCGACTTATTTTTTCAAACGACGCCAAAACCTCGAACCTAAAGTGCGACTTTTCGGATTGCCTCGATATCATGCGGGTGGAAACACAT\n",
            "10000\n",
            "10000\n"
          ]
        }
      ],
      "source": [
        "num_clusters = 10000\n",
        "train_consensus_strands = []\n",
        "train_target = []\n",
        "for i in range(num_clusters):\n",
        "    strand = [random.choice([\"A\",\"C\",\"G\",\"T\"]) for _ in range(120)]\n",
        "    train_consensus_strands.append(strand)\n",
        "    strand_t = [one_hot(c) for c in strand]\n",
        "#     strand_t = Variable(torch.FloatTensor([one_hot(c) for c in strand]))\n",
        "    train_target.append(strand_t)\n",
        "print(\"\".join(strand))\n",
        "\n",
        "print(len(train_consensus_strands))\n",
        "print(len(train_target))\n",
        "\n",
        "train_target = torch.Tensor(train_target).cuda()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "scrolled": true,
        "id": "4AnG9K3M3opI",
        "outputId": "a44da320-14ee-444f-83cd-bee3963d31bd"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "torch.Size([10000, 120, 1, 4])\n"
          ]
        }
      ],
      "source": [
        "print(train_target.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Dv-4cZU3opJ",
        "outputId": "717bae4e-5b83-4f4a-cdfe-5a8dc1bda625"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "10000\n"
          ]
        }
      ],
      "source": [
        "train_data = []\n",
        "\n",
        "for i in range(num_clusters):\n",
        "    train_clusters = []\n",
        "    seq = train_consensus_strands[i]\n",
        "    for j in range(10):\n",
        "        noisy_seq = sim_error(seq, pi=random.uniform(0.04, 0.06), pd=random.uniform(0.04, 0.06), \n",
        "        ps=random.uniform(0.01, 0.03))\n",
        "        \n",
        "        noisy_seq_t = [one_hot(c) for c in strand]\n",
        "#         noisy_seq_t = Variable(torch.FloatTensor([one_hot(c) for c in noisy_seq])).cuda()\n",
        "        train_clusters.append(noisy_seq_t)\n",
        "        \n",
        "    train_data.append(train_clusters)\n",
        "    \n",
        "print(len(train_data))\n",
        "train_data = torch.Tensor(train_data).cuda()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JazdF7mn3opJ",
        "outputId": "ade1ee55-a06d-4a19-8ae2-7be1f05205cf"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "torch.Size([10000, 10, 120, 1, 4])\n"
          ]
        }
      ],
      "source": [
        "print(train_data.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r8vtyufQ3opJ",
        "outputId": "a3c2a76f-62cd-42e1-8869-a190697ac0a8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "ConsNet(\n",
            "  (lstm): LSTM(4, 32)\n",
            "  (fc1): Linear(in_features=32, out_features=12, bias=True)\n",
            "  (dropout1): Dropout(p=0.2, inplace=False)\n",
            "  (fc2): Linear(in_features=12, out_features=12, bias=True)\n",
            "  (dropout2): Dropout(p=0.1, inplace=False)\n",
            "  (fc3): Linear(in_features=12, out_features=4, bias=True)\n",
            "  (dropout3): Dropout(p=0.1, inplace=False)\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "# initialize the model\n",
        "model = ConsNet(32, 12, 12)\n",
        "model.cuda()\n",
        "print(model)\n",
        "model.zero_grad()\n",
        "model.hidden = model.init_hidden()\n",
        "\n",
        "# initial the paramerters in the DCNet\n",
        "for name, param in model.named_parameters():\n",
        "    if 'bias' in name:\n",
        "        nn.init.constant_(param, 0.0)\n",
        "    elif 'weight' in name:\n",
        "        nn.init.xavier_normal_(param)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dMHNgM6t3opK"
      },
      "outputs": [],
      "source": [
        "loss = nn.MSELoss()\n",
        "initial_lr = 0.1\n",
        "lr=initial_lr\n",
        "optimizer = optim.SGD(model.parameters(), lr=initial_lr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r9yGJAqk3opK",
        "outputId": "880bf128-882e-4bb7-8350-a662ebcd2907"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch: 0 Training loss: 3.6179283261299135e-05 learning rate: 0.1\n",
            "Epoch: 1 Training loss: 3.968231379985809e-05 learning rate: 0.095\n",
            "Epoch: 2 Training loss: 3.3966299891471864e-05 learning rate: 0.095\n",
            "Epoch: 3 Training loss: 4.36478853225708e-05 learning rate: 0.09025\n",
            "Epoch: 4 Training loss: 4.020006060600281e-05 learning rate: 0.09025\n",
            "Epoch: 5 Training loss: 3.661943376064301e-05 learning rate: 0.0857375\n",
            "Epoch: 6 Training loss: 3.6713653802871704e-05 learning rate: 0.0857375\n",
            "Epoch: 7 Training loss: 3.752442896366119e-05 learning rate: 0.08145062499999998\n",
            "Epoch: 8 Training loss: 4.580929279327393e-05 learning rate: 0.08145062499999998\n",
            "Epoch: 9 Training loss: 3.692469894886017e-05 learning rate: 0.07737809374999999\n",
            "Epoch: 10 Training loss: 3.812678754329681e-05 learning rate: 0.07737809374999999\n",
            "Epoch: 11 Training loss: 3.5812777280807496e-05 learning rate: 0.07350918906249998\n",
            "Epoch: 12 Training loss: 3.5603234171867374e-05 learning rate: 0.07350918906249998\n",
            "Epoch: 13 Training loss: 4.046078622341156e-05 learning rate: 0.06983372960937498\n",
            "Epoch: 14 Training loss: 3.47099244594574e-05 learning rate: 0.06983372960937498\n",
            "Epoch: 15 Training loss: 3.4385251998901364e-05 learning rate: 0.06634204312890622\n",
            "Epoch: 16 Training loss: 3.757069110870361e-05 learning rate: 0.06634204312890622\n",
            "Epoch: 17 Training loss: 3.584400117397309e-05 learning rate: 0.0630249409724609\n",
            "Epoch: 18 Training loss: 3.782275915145874e-05 learning rate: 0.0630249409724609\n",
            "Epoch: 19 Training loss: 3.67660254240036e-05 learning rate: 0.05987369392383786\n"
          ]
        }
      ],
      "source": [
        "# Training the model\n",
        "num_epochs = 20\n",
        "\n",
        "range_ = (1, 121)\n",
        "\n",
        "for epoch in range(num_epochs):\n",
        "    model.train()\n",
        "    for i in range(int(len(train_data))):\n",
        "        train_loss = 0\n",
        "        s, e = range_\n",
        "        optimizer.zero_grad()\n",
        "        for seq in train_data[i]:\n",
        "            model.hidden = model.init_hidden()\n",
        "            model.zero_grad()\n",
        "            # Noisy clusters (training data)\n",
        "            seq = seq[s-1:e]\n",
        "            seq_ = seq.view(-1,4)\n",
        "            out = model(seq_)\n",
        "            # Original string (training target)\n",
        "            seq_target = train_target[i][s-1:e]\n",
        "            seq_target = seq_target.view(-1, 4)\n",
        "            # Loss computation\n",
        "            train_loss += loss(out, seq_target)\n",
        "            \n",
        "        # Backward propagation operation\n",
        "        train_loss.backward()\n",
        "        optimizer.step()\n",
        "        \n",
        "    print(\"Epoch:\", epoch, \"Training loss:\", train_loss.cpu().item()/len(train_data), \"learning rate:\", lr)\n",
        "        \n",
        "    # Learning rate decay\n",
        "    if epoch % 2 ==0:\n",
        "        lr *= 0.95\n",
        "        optimizer = optim.SGD(model.parameters(), lr=lr)\n",
        "    \n",
        "if (num_epochs > 0):\n",
        "    torch.save(model.state_dict(), \"consnet.pt\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GdoHPhRU3opL",
        "outputId": "f9403061-418f-43f0-f11e-1fefbee1d1e5"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model.load_state_dict(torch.load(\"consnet.pt\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0wla4nrf3opL"
      },
      "source": [
        "### Generate the test data "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YAbeqnPD3opL",
        "outputId": "9440b17c-5830-409f-eab4-c60d203d062d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "ACATATTCTTCTGAGGCTTCCCTGACTTCTCGAGCCTATTTGGTTGCCGCTCGGCTCGGAAGGAGACCATCGGGAGAGGCACTATAAGGTAAGCGGGACCGTCGGAGTTTACAGTCTCTA\n",
            "100\n",
            "100\n"
          ]
        }
      ],
      "source": [
        "# DNA clusters to test model\n",
        "num_test_clusters = 1000\n",
        "test_consensus_strands = []\n",
        "test_target = []\n",
        "\n",
        "for i in range(num_test_clusters):\n",
        "    strand = [random.choice([\"A\",\"C\",\"G\",\"T\"]) for _ in range(120)]\n",
        "    test_consensus_strands.append(strand)\n",
        "    strand_t = [one_hot(c) for c in strand]\n",
        "#     strand_t = Variable(torch.FloatTensor([one_hot(c) for c in seq])).cuda()\n",
        "    test_target.append(strand_t)\n",
        "    \n",
        "print(\"\".join(strand))\n",
        "\n",
        "print(len(test_consensus_strands))\n",
        "print(len(test_target))\n",
        "\n",
        "test_target = torch.Tensor(test_target)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "csW1LzzB3opM",
        "outputId": "2bdf9f9f-4413-4694-ac32-cd079be7b745"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "100\n"
          ]
        }
      ],
      "source": [
        "test_data = []\n",
        "\n",
        "for i in range(num_test_clusters):\n",
        "    test_clusters = []\n",
        "    seq = test_consensus_strands[i]\n",
        "    for j in range(10):\n",
        "        noisy_seq = sim_error(seq, pi=random.uniform(0.04, 0.06), pd=random.uniform(0.04, 0.06), \n",
        "        ps=random.uniform(0.01, 0.03))\n",
        "        \n",
        "        noisy_seq_t = [one_hot(c) for c in strand]\n",
        "#         noisy_seq_t = Variable(torch.FloatTensor([one_hot(c) for c in noisy_seq])).cuda()\n",
        "        test_clusters.append(noisy_seq_t)\n",
        "        \n",
        "    test_data.append(test_clusters)\n",
        "    \n",
        "print(len(test_data))\n",
        "test_data = torch.Tensor(test_data).cuda()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "scrolled": true,
        "id": "DoksrxyH3opM",
        "outputId": "32440ee1-4eb5-4821-f907-09cdc88fe2f5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "torch.Size([100, 10, 120, 1, 4])\n"
          ]
        }
      ],
      "source": [
        "print(test_data.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ex-fUu3f3opM",
        "outputId": "e7ff215b-913b-48d4-95be-e8f0315462ff"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "torch.Size([100, 120, 1, 4])\n"
          ]
        }
      ],
      "source": [
        "print(test_target.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Oyg2gcDP3opN",
        "outputId": "a0408535-24d9-4a2c-e5be-0cd257113f0b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "torch.Size([4])\n",
            "torch.float32\n",
            "Vectors for one hot of DNA bases\n",
            "A: 0\n",
            "C: 1\n",
            "G: 2\n",
            "T: 3\n"
          ]
        }
      ],
      "source": [
        "#{\"A\":0, \"C\":1, \"G\":2, \"T\":3}\n",
        "\n",
        "def one_hot_to_vec(one_hot):\n",
        "    one_hot = np.array(one_hot)\n",
        "    vec = np.argmax(one_hot, axis = 0)\n",
        "    return vec\n",
        "\n",
        "vec = torch.zeros(4).cpu()\n",
        "vec[0] = 1\n",
        "print(vec.shape)\n",
        "print(vec.dtype)\n",
        "\n",
        "print(\"Vectors for one hot of DNA bases\")\n",
        "print(\"A:\", one_hot_to_vec(vec))\n",
        "print(\"C:\", one_hot_to_vec([0, 1, 0, 0]))\n",
        "print(\"G:\", one_hot_to_vec([0, 0, 1, 0]))\n",
        "print(\"T:\", one_hot_to_vec([0, 0, 0, 1]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lFJPRVKv3opN",
        "outputId": "ae2eb3cb-c8d7-41d5-a5ef-abf5b26cf609"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Time for forward prop per DNA sequence = 0.0019271373748779297 sec\n",
            "Time for accuracy calculation per DNA sequence = 0.00529789924621582 sec\n",
            "Accuracy = 0.25125\n"
          ]
        }
      ],
      "source": [
        "range_ = (1, 121)\n",
        "\n",
        "model.hidden = model.init_hidden()\n",
        "accuracy = 0\n",
        "\n",
        "model.eval()\n",
        "# Run again\n",
        "for i in range(len(test_data)):\n",
        "    test_loss = 0\n",
        "    s, e = range_\n",
        "    for seq in test_data[i]:\n",
        "        start_time = time.time()\n",
        "        model.hidden = model.init_hidden()\n",
        "        model.zero_grad()\n",
        "        # Noisy clusters (test data)\n",
        "        seq = seq[s-1:e]\n",
        "        seq_ = seq.view(-1,4)\n",
        "        out = model(seq_)\n",
        "        check_time = time.time()\n",
        "        # Original string (test target)\n",
        "        seq_target = test_target[i][s-1:e]\n",
        "        seq_target = seq_target.view(-1, 4)\n",
        "        \n",
        "        for j in range(120):\n",
        "            if (one_hot_to_vec(out[j].detach().cpu().numpy()) == one_hot_to_vec(seq_target[j].detach().cpu().numpy())):\n",
        "                accuracy = accuracy + 1\n",
        "        end_time = time.time()\n",
        "        \n",
        "print(\"Time for forward prop per DNA sequence = \" + str(check_time - start_time) + \" sec\")\n",
        "print(\"Time for accuracy calculation per DNA sequence = \" + str(end_time - check_time) + \" sec\")\n",
        "                \n",
        "print(\"Accuracy = \" + str(accuracy/(test_data.shape[0] * test_data.shape[1] * 120)))"
      ]
    }
  ],
  "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.8.5"
    },
    "colab": {
      "name": "consnet.ipynb",
      "provenance": [],
      "include_colab_link": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}