Covid_Predict_Omicron_Resistance.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "name": "Covid_Predict_Omicron_Resistance.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "mount_file_id": "1-5bfBr79Dr-tCjSTFeApEdyEcWmccjp7",
      "authorship_tag": "ABX9TyMFG68Y1WoOlI33rdr7WHBP",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/bahrad/Covid/blob/main/Covid_Predict_Omicron_Resistance.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UzkSry5x4VaL"
      },
      "source": [
        "#Initialization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zLa737hB7e4_"
      },
      "source": [
        "%tensorflow_version 2.x\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "\n",
        "import numpy as np\n",
        "import os\n",
        "import csv\n",
        "\n",
        "import pandas as pd\n",
        "import pickle"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IiNZfPTu0_em"
      },
      "source": [
        "from google.colab import drive, files\n",
        "# drive.mount('/content/drive')\n",
        "\n",
        "FILELOC = \"/content/drive/My Drive/COVID_Python/\""
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "INmsA86dCdMV",
        "outputId": "0254b7c0-b715-458a-928e-a653541b0764"
      },
      "source": [
        "try:\n",
        "    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection\n",
        "    print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])\n",
        "    tf.config.experimental_connect_to_cluster(tpu)\n",
        "    tf.tpu.experimental.initialize_tpu_system(tpu)\n",
        "    tpu_strategy = tf.distribute.TPUStrategy(tpu)\n",
        "    tpu_env=True\n",
        "except ValueError:\n",
        "    print('Not connected to a TPU runtime.')\n",
        "    tpu_env=False"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Running on TPU  ['10.99.211.42:8470']\n",
            "INFO:tensorflow:Deallocate tpu buffers before initializing tpu system.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:Deallocate tpu buffers before initializing tpu system.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:Initializing the TPU system: grpc://10.99.211.42:8470\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:Initializing the TPU system: grpc://10.99.211.42:8470\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:Finished initializing TPU system.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:Finished initializing TPU system.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:Found TPU system:\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:Found TPU system:\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores: 8\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores: 8\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Num TPU Workers: 1\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Num TPU Workers: 1\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Num TPU Cores Per Worker: 8\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZXhwnqGgqqjV"
      },
      "source": [
        "##Define models"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3EfE13Hp8AJ9"
      },
      "source": [
        "def reset_model(regress, singleclass, multiclass, output_multiheadatt, use_att, nclasses=4):\n",
        "\n",
        "    if output_multiheadatt:\n",
        "        model_fn = AttMod_2\n",
        "    else:\n",
        "        model_fn = AttModel\n",
        "    model = model_fn(L=ismlen,\n",
        "                     vocab_size=len(aa_list)+1,\n",
        "                     embdim = ENCDIM,\n",
        "                     numheads = NHEADS,\n",
        "                     ffdim = FFDIM,\n",
        "                     num_dense = NDENSE,\n",
        "                     mask_zero=True,\n",
        "                     dropout_rate = DROPRATE,\n",
        "                     trans_drop = TRANSDROPRATE,\n",
        "                     Nt = NT,\n",
        "                     W = 1, Nc = NC, Nl = NL,\n",
        "                     regress=regress, singleclass=singleclass,\n",
        "                     multiclass=multiclass, use_att=use_att,\n",
        "                     nclasses=nclasses,\n",
        "                     )\n",
        "    \n",
        "    optimizer = keras.optimizers.Adam(learning_rate=LEARN_RATE)\n",
        "    if regress:\n",
        "        loss = keras.losses.MeanSquaredError()\n",
        "        metrics = [keras.metrics.MeanSquaredError(name='mse'),\n",
        "            keras.metrics.MeanSquaredLogarithmicError(name='msle'),\n",
        "            keras.losses.MeanAbsoluteError(name='mae')\n",
        "            ]\n",
        "    if singleclass:\n",
        "        loss = keras.losses.BinaryCrossentropy()\n",
        "        metrics = [keras.metrics.BinaryAccuracy(name='acc')]\n",
        "    if multiclass:    \n",
        "        loss = keras.losses.SparseCategoricalCrossentropy()\n",
        "        metrics = [keras.metrics.SparseCategoricalAccuracy(name='acc')]\n",
        "    model.compile(loss=loss, optimizer=optimizer, metrics=metrics,)\n",
        "                #   steps_per_execution = STEPS_PER_EXECUTION,)\n",
        "    \n",
        "    return model"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SGYlfhd95NKk"
      },
      "source": [
        "class TransformerBlock(keras.layers.Layer):\n",
        "    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):\n",
        "        super(TransformerBlock, self).__init__()\n",
        "        self.att = keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)\n",
        "        self.ffn = keras.Sequential(\n",
        "            [keras.layers.Dense(ff_dim, activation=\"relu\"), keras.layers.Dense(embed_dim),]\n",
        "        )\n",
        "        self.layernorm1 = keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "        self.layernorm2 = keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "        self.dropout1 = keras.layers.Dropout(rate)\n",
        "        self.dropout2 = keras.layers.Dropout(rate)\n",
        "\n",
        "    def call(self, inputs, training):\n",
        "        attn_output = self.att(inputs, inputs)\n",
        "        attn_output = self.dropout1(attn_output, training=training)\n",
        "        out1 = self.layernorm1(inputs + attn_output)\n",
        "        ffn_output = self.ffn(out1)\n",
        "        ffn_output = self.dropout2(ffn_output, training=training)\n",
        "        return self.layernorm2(out1 + ffn_output)\n",
        "\n",
        "class TokenAndPositionEmbedding(keras.layers.Layer):\n",
        "    def __init__(self, maxlen, vocab_size, embed_dim, mask_zero=False):\n",
        "        super(TokenAndPositionEmbedding, self).__init__()\n",
        "        self.token_emb = keras.layers.Embedding(input_dim=vocab_size,\n",
        "                                                output_dim=embed_dim,\n",
        "                                                mask_zero=mask_zero)\n",
        "        self.pos_emb = keras.layers.Embedding(input_dim=maxlen, output_dim=embed_dim,\n",
        "                                              mask_zero=mask_zero)\n",
        "\n",
        "    def call(self, x):\n",
        "        maxlen = tf.shape(x)[-1]\n",
        "        positions = tf.range(start=0, limit=maxlen, delta=1)\n",
        "        positions = self.pos_emb(positions)\n",
        "        x = self.token_emb(x)\n",
        "        return x + positions\n",
        "\n",
        "def linear01(x):\n",
        "    return tf.clip_by_value(x, clip_value_min=0, clip_value_max=1)"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gPCCClpo7Rim"
      },
      "source": [
        "def AttMod_2(L, vocab_size, embdim, numheads, ffdim, num_dense=False,\n",
        "             mask_zero=False, dropout_rate=False, trans_drop=0.1,\n",
        "             Nt=1, W=False, Nc=False, Nl=False,\n",
        "             regress=True, singleclass=False, multiclass=False, use_att=True,\n",
        "             nclasses=4):\n",
        "\n",
        "    inpTensor = keras.Input(shape=(L,))\n",
        "    x = inpTensor\n",
        "\n",
        "    if mask_zero:\n",
        "        x = keras.layers.Masking(mask_value=0)(x)   \n",
        "\n",
        "    x = TokenAndPositionEmbedding(L, vocab_size, embdim, mask_zero)(x)\n",
        "\n",
        "    if W and Nc and Nl:\n",
        "        for n in range(Nl):\n",
        "            x = keras.layers.Conv1D(filters = Nc,\n",
        "                                kernel_size = W,\n",
        "                                activation = 'relu',\n",
        "                                padding = 'same',\n",
        "                                )(x)\n",
        "            if n > 1 and n < Nl-1:\n",
        "                x = keras.layers.BatchNormalization()(x)\n",
        "\n",
        "    y, attout = keras.layers.MultiHeadAttention(num_heads=numheads, key_dim=embdim,\n",
        "                                                )(x, x, return_attention_scores=True)\n",
        "    y = keras.layers.Dropout(trans_drop)(y)\n",
        "    z = keras.layers.LayerNormalization(epsilon=1e-6)(x + y)\n",
        "    z1 = keras.Sequential( [keras.layers.Dense(ffdim, activation=\"relu\"), keras.layers.Dense(embdim),])\n",
        "    z1 = keras.layers.Dropout(trans_drop)(z)\n",
        "    x = keras.layers.LayerNormalization(epsilon=1e-6)(z + z1)\n",
        "\n",
        "    if use_att:\n",
        "        # Attention layer\n",
        "        h = keras.layers.TimeDistributed(keras.layers.Dense(Nc, activation='tanh'))(x)\n",
        "        attention = keras.layers.TimeDistributed(keras.layers.Dense(1, activation='tanh'))(h)\n",
        "        attention = keras.layers.Flatten()(attention)  \n",
        "        attention = keras.layers.Softmax(axis=1, name='attention')(attention) # normalize attention values\n",
        "        attention = keras.layers.RepeatVector(Nc)(attention)\n",
        "        attention = keras.layers.Permute([2, 1])(attention)\n",
        "        representation = keras.layers.multiply([h, attention])\n",
        "        representation = tf.math.reduce_sum(representation, axis = 1)\n",
        "        x = representation\n",
        "    else:\n",
        "        x = keras.layers.GlobalAveragePooling1D()(x)\n",
        "\n",
        "    if num_dense:\n",
        "        x = keras.layers.Dense(num_dense, activation = 'relu')(x)\n",
        "    if dropout_rate:\n",
        "        x = keras.layers.Dropout(Params[dropout_rate])(x)\n",
        "\n",
        "    if regress:\n",
        "        finalOut = keras.layers.Dense(1, activation=linear01)(x)\n",
        "    if singleclass:\n",
        "        finalOut = keras.layers.Dense(1, activation='sigmoid')(x)\n",
        "    if multiclass:\n",
        "        finalOut = keras.layers.Dense(nclasses, activation='softmax')(x)\n",
        "\n",
        "    # define the model's start and end points    \n",
        "    model = keras.Model(inpTensor,finalOut)\n",
        "\n",
        "    return model"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KK05TjTo79ql"
      },
      "source": [
        "def AttModel(L, vocab_size, embdim, numheads, ffdim, num_dense=False,\n",
        "             mask_zero=False, dropout_rate=False, trans_drop=0.1,\n",
        "             Nt=1, W=False, Nc=False, Nl=False,\n",
        "             regress=True, singleclass=False, multiclass=False, use_att=True,\n",
        "             nclasses=4):\n",
        "\n",
        "    inpTensor = keras.Input(shape=(L,))\n",
        "    x = inpTensor\n",
        "\n",
        "    if mask_zero:\n",
        "        x = keras.layers.Masking(mask_value=0)(x)   \n",
        "\n",
        "    x = TokenAndPositionEmbedding(L, vocab_size, embdim, mask_zero)(x)\n",
        "\n",
        "    if W and Nc and Nl:\n",
        "        for n in range(Nl):\n",
        "            x = keras.layers.Conv1D(filters = Nc,\n",
        "                                kernel_size = W,\n",
        "                                activation = 'relu',\n",
        "                                padding = 'same',\n",
        "                                )(x)\n",
        "            if n > 1 and n < Nl-1:\n",
        "                x = keras.layers.BatchNormalization()(x)\n",
        "\n",
        "    for n in range(Nt):\n",
        "        x = TransformerBlock(Nc, numheads, ffdim, rate=trans_drop)(x)\n",
        "\n",
        "    if use_att:\n",
        "        # Attention layer\n",
        "        h = keras.layers.TimeDistributed(keras.layers.Dense(Nc, activation='tanh'))(x)\n",
        "        attention = keras.layers.TimeDistributed(keras.layers.Dense(1, activation='tanh'))(h)\n",
        "        attention = keras.layers.Flatten()(attention)  \n",
        "        attention = keras.layers.Softmax(axis=1, name='attention')(attention) # normalize attention values\n",
        "        attention = keras.layers.RepeatVector(Nc)(attention)\n",
        "        attention = keras.layers.Permute([2, 1])(attention)\n",
        "        representation = keras.layers.multiply([h, attention])\n",
        "        representation = tf.math.reduce_sum(representation, axis = 1)\n",
        "        x = representation\n",
        "    else:\n",
        "        x = keras.layers.GlobalAveragePooling1D()(x)\n",
        "\n",
        "    if num_dense:\n",
        "        x = keras.layers.Dense(num_dense, activation = 'relu')(x)\n",
        "    if dropout_rate:\n",
        "        x = keras.layers.Dropout(Params[dropout_rate])(x)\n",
        "\n",
        "    if regress:\n",
        "        # finalOut = keras.layers.Dense(1, activation=linear01)(x)\n",
        "        finalOut = keras.layers.Dense(1, activation='sigmoid')(x)       \n",
        "    if singleclass:\n",
        "        finalOut = keras.layers.Dense(1, activation='sigmoid')(x)\n",
        "    if multiclass:\n",
        "        finalOut = keras.layers.Dense(nclasses, activation='softmax')(x)\n",
        "\n",
        "    # define the model's start and end points    \n",
        "    model = keras.Model(inpTensor,finalOut)\n",
        "\n",
        "    return model"
      ],
      "execution_count": 50,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NN3ETSv0Covu"
      },
      "source": [
        "ismlen = 1273\n",
        "\n",
        "def tokenize_sequences(data_dataframe, SeqCol='ISM', seqlen=1273):\n",
        "    def f(x):\n",
        "        if len(x) < seqlen:\n",
        "            return x + '*'*(seqlen-len(x))\n",
        "        elif len(x) > seqlen:\n",
        "            return x[:seqlen]\n",
        "        else:\n",
        "            return x\n",
        "    data = np.vstack(data_dataframe[SeqCol].apply(f).apply(lambda x: np.array(list(x))))\n",
        "    aa_list = ['A', 'R', 'N', 'D', 'C', 'Q', 'E',\n",
        "            'G', 'H', 'I', 'L', 'K', 'M', 'F',\n",
        "            'P', 'S', 'T', 'W', 'Y', 'V', '-',\n",
        "            ]\n",
        "    aa_tokenizer = {aa_list[k]:k+1 for k in range(len(aa_list))}\n",
        "    aa_tokenizer['*'] = 0\n",
        "    aa_tokenizer['X'] = 0\n",
        "    # optionally handle B, J, Z ambiguities\n",
        "    # Asx\tB\tAspartic acid or Asparagine (D or N)\n",
        "    # Glx\tZ\tGlutamic acid or Glutamine (E or Q)\n",
        "    # Xaa\tX\tAny amino acid\n",
        "    # Xle\tJ\tLeucine or Isoleucine (L or I)\n",
        "    aa_tokenizer['B'] = 0\n",
        "    aa_tokenizer['Z'] = 0\n",
        "    aa_tokenizer['J'] = 0\n",
        "\n",
        "    return np.vectorize(aa_tokenizer.get)(data)\n",
        "\n",
        "aa_list = ['A', 'R', 'N', 'D', 'C', 'Q', 'E',\n",
        "        'G', 'H', 'I', 'L', 'K', 'M', 'F',\n",
        "        'P', 'S', 'T', 'W', 'Y', 'V', '-',\n",
        "        ]"
      ],
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lsAfrYxP75be"
      },
      "source": [
        "# These parameters are currently hard-coded\n",
        "ENCDIM = 1500\n",
        "NC = 300\n",
        "NL = 1\n",
        "NT = 1\n",
        "NHEADS = 8\n",
        "FFDIM = 64\n",
        "NDENSE = 64\n",
        "TRANSDROPRATE = 0.4\n",
        "DROPRATE = 0\n",
        "\n",
        "LEARN_RATE = 0.0001\n",
        "\n",
        "BATCH_SIZE = 48\n",
        "\n",
        "STEPS_PER_EXECUTION = 50"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tK6lGYVzqkV-"
      },
      "source": [
        "##Prepare sequences for training / validation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gNaPWU5tqymM"
      },
      "source": [
        "Based on protein sequence data downloaded from GISAID (http://www.gisaid.org) that was current as of October 1, 2021."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fJfOLZK8rDGb"
      },
      "source": [
        "Aligned sequences (aligned using BLOSUM62 with local pairwise SSW method implemented in scikit-bio)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_dqY5aUWon_k"
      },
      "source": [
        "# mdf = pd.read_csv(FILELOC+'metadata.tsv', '\\t')\n",
        "# seqs = pd.read_csv(FILELOC+'align1001_all.csv')\n",
        "# df = pd.merge(left=mdf[['Accession ID', 'Collection date', 'Clade', 'Pango lineage', 'Location']],\n",
        "#               right=seqs, left_on='Accession ID', right_on=\"SequenceID\", how='inner')\n",
        "# df.rename(columns={'Accession ID':'SequenceID', 'Pango lineage':'Lineage'}, inplace=True)\n",
        "# df = df.loc[:,~df.columns.duplicated()]\n",
        "# df['Country'] = df['Location'].apply(lambda x: x.split('/')[1].strip())\n",
        "# from dateutil.parser import parse as dateparse\n",
        "# df = df[['SequenceID', 'Date', 'Clade', 'Lineage', 'Country', 'MaskedSeq', 'Collection date', 'Location']]\n",
        "# # with open(FILELOC + 'Spike1001_all.pkl', 'wb') as f:\n",
        "# #     pickle.dump(df, f)\n",
        "# with open(FILELOC + 'Spike1001_all.pkl', 'rb') as f:\n",
        "#     df = pickle.load(f)\n",
        "# df = df[df.Date > 0]\n",
        "# df_reldate = df.groupby(\"MaskedSeq\")[\"Date\"].apply(list).to_frame()\n",
        "# df_reldate['First_Date'] = df_reldate.Date.apply(np.amin)\n",
        "# df_reldate['Last_Date'] = df_reldate.Date.apply(np.amax)\n",
        "# df_reldate['Peak_Date'] = df_reldate.Date.apply(lambda x: np.median(np.argwhere(np.bincount(x)==np.amax(np.bincount(x)))))\n",
        "# df_reldate['Count'] = df_reldate.Date.apply(len)\n",
        "# max_first_date = max(df_reldate.First_Date)\n",
        "# df_reldate['relfirstdate'] = df_reldate['First_Date'].apply(lambda x: x/max_first_date)\n",
        "# with open(FILELOC + 'Spike1001_allseqs_reldate.pkl', 'wb') as f:\n",
        "#     pickle.dump(df_reldate, f)\n",
        "# ddf = df.join(df_reldate, on='MaskedSeq', lsuffix='', rsuffix='_List')\n",
        "# with open(FILELOC + 'Spike1001_allseqs_reldate_metadata.pkl', 'wb') as f:\n",
        "#     pickle.dump(ddf, f)\n",
        "# with open(FILELOC + 'Spike1001_allseqs_reldate.pkl', 'rb') as f:\n",
        "#     alignseqs = pickle.load(f)\n",
        "# with open(FILELOC + 'Spike1001_allseqs_reldate_metadata.pkl', 'rb') as f:\n",
        "#     aligndf = pickle.load(f)\n",
        "# print(len(aligndf))\n",
        "# aligndf.drop(aligndf[aligndf.Lineage.isna()].index, inplace=True)\n",
        "# print(len(aligndf))\n",
        "# adf = aligndf.groupby('MaskedSeq')['Lineage'].apply(lambda x: x.value_counts().index[0]).to_frame()\n",
        "# adfm = adf.join(alignseqs, how='left')\n",
        "# with open(FILELOC + 'aligned_sequence_lineages_1001.pkl', 'wb') as f:\n",
        "#     pickle.dump([adf, adfm], f)"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UnbcbihbrdZi"
      },
      "source": [
        "Raw sequence data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d_D4KILlCRri"
      },
      "source": [
        "# df = pd.read_csv(FILELOC + 'results1001_Spike.csv')\n",
        "# df.drop(columns='index', inplace=True)\n",
        "# mdf = pd.read_csv(FILELOC+'metadata.tsv', '\\t')\n",
        "# mdf = pd.merge(left=mdf[['Accession ID', 'Collection date', 'Clade', 'Pango lineage', 'Location']],\n",
        "#               right=df, left_on='Accession ID', right_on=\"SequenceID\", how='inner')\n",
        "# mdf.rename(columns={'Pango lineage':'Lineage'}, inplace=True)\n",
        "# print(len(mdf))\n",
        "# mdf['Country'] = mdf['Location'].apply(lambda x: x.split('/')[1].strip())\n",
        "# mdf['Country'].fillna('Unknown', inplace=True)\n",
        "# REFDATE = '2019-12-01'; refdt = dateparse(REFDATE)\n",
        "# def f(x):\n",
        "#     return (dateparse(x) - refdt).days\n",
        "# mdf['reldate'] = mdf['Collection date'].apply(f)\n",
        "# # with open('RawSequence_Metadata_1001.pkl', 'wb') as f:\n",
        "# #     pickle.dump(mdf, f)\n",
        "# # mdf.to_csv('RawSequence_Metadata_1001.csv', index=False)\n",
        "# # !tar -cvzf RawSequence_Metadata_1001.csv.tar.gz RawSequence_Metadata_1001.csv\n",
        "# df_reldate = mdf.groupby(\"Spike\")[\"reldate\"].apply(list).to_frame()\n",
        "# df_reldate.rename(columns={\"reldate\":\"Date\"}, inplace=True)\n",
        "# df_reldate['First_Date'] = df_reldate.Date.apply(np.amin)\n",
        "# df_reldate['Last_Date'] = df_reldate.Date.apply(np.amax)\n",
        "# df_reldate['Count'] = df_reldate.Date.apply(len)\n",
        "# max_first_date = max(df_reldate.First_Date)\n",
        "# df_reldate['relfirstdate'] = df_reldate['First_Date'].apply(lambda x: x/max_first_date)\n",
        "# df_reldate.reset_index().to_csv(FILELOC + 'RawSequence1001_RelDate.csv', index=False)\n",
        "# rawseqs = pd.read_csv(FILELOC + 'RawSequence1001_RelDate.csv')\n",
        "# with open(FILELOC + 'RawSequence_Metadata_1001.pkl', 'rb') as f:\n",
        "#     rawdf = pickle.load(f)\n",
        "# print(len(rawdf))\n",
        "# rawdf.drop(rawdf[rawdf.Lineage.isna()].index, inplace=True)\n",
        "# print(len(rawdf))\n",
        "# rdf = rawdf.groupby('Spike')['Lineage'].apply(lambda x: x.value_counts().index[0]).to_frame()\n",
        "# rsdf = rawseqs.set_index('Spike', drop=True)\n",
        "# rdfm = rdf.join(rsdf, how='left')\n",
        "# with open(FILELOC + 'raw_sequence_lineages_1001.pkl', 'wb') as f:\n",
        "#     pickle.dump([rdf, rdfm], f)"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TZXONmmr3tgo"
      },
      "source": [
        "##Get Omicron spike sequences"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xzS4Sh18VBFt"
      },
      "source": [
        "# !pip install scikit-bio\n",
        "# import skbio\n",
        "# from skbio import TabularMSA, DNA, Protein\n",
        "# from skbio.alignment import local_pairwise_align_ssw\n",
        "# !pip install biopython\n",
        "# from Bio.Align import substitution_matrices\n",
        "# import itertools"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ROurf6JoxeVo"
      },
      "source": [
        "DIR = FILELOC + 'coronavirus_spike/'\n",
        "SPIKEFILE = DIR + 'EPI_ISL_402124-S.fasta'\n",
        "# FASTAFILE = DIR + 'gisaid_hcov-19_2021_11_28_08.fasta'\n",
        "OMICRONFILE = DIR + 'gisaid_hcov-19_2021_12_01_22.fasta'\n",
        "\n",
        "seqs = [seq for seq in skbio.io.read(OMICRONFILE, format='fasta')]\n",
        "spike = DNA.read(SPIKEFILE)  # reference sequence from NCBI\n",
        "align = [local_pairwise_align_ssw(DNA(seq),spike) for seq in seqs]\n",
        "# align is a triplet of (alignment, score, start_end_positions)\n",
        "spike_prot = ''.join([x.decode() for x in spike.translate().values])\n",
        "\n",
        "# # Two ways of identifying aligned sequences (generate the same results)\n",
        "\n",
        "# find aligned spike sequences without degenerates to translate\n",
        "no_degenerates = [k for k in range(len(seqs)) if not align[k][0][0].has_degenerates()]\n",
        "omicron_protein_seqs = [align[k][0][0].degap().translate() for k in no_degenerates]\n",
        "omicron_seqs = [''.join([x.decode() for x in seq.values]) for seq in omicron_protein_seqs]\n",
        "omicron_metadata = [seq.metadata['id'] for seq in omicron_protein_seqs]\n",
        "\n",
        "# # find aligned spike sequences without degenerates to translate\n",
        "# no_degenerates = [k for k in range(len(seqs)) if not align[k][0][0].has_degenerates()]\n",
        "# spikelocs = [align[k][2][0] for k in range(len(seqs))]\n",
        "# omicron_spike_seqs = [seqs[k][spikelocs[k][0]:spikelocs[k][1]+1] for k in no_degenerates]\n",
        "# omicron_protein_seqs = [DNA(seq).translate() for seq in omicron_spike_seqs]\n",
        "# omicron_seqs = [''.join([x.decode() for x in seq.values]) for seq in omicron_protein_seqs]\n",
        "# omicron_metadata = [seq.metadata['id'] for seq in omicron_protein_seqs]\n",
        "\n",
        "omicrondf = pd.DataFrame.from_dict({'Metadata':omicron_metadata})\n",
        "omicrondf['SeqID'] = omicrondf.Metadata.apply(lambda x:x.split('|')[1])\n",
        "omicrondf['Raw'] = omicron_seqs\n",
        "# omicrondf['seqlen'] = omicrondf['Spike'].apply(len)\n",
        "\n",
        "# # generate modified aligned sequences where there are insertions relative to\n",
        "# # reference sequence, i.e., where the inserted residues are deleted\n",
        "\n",
        "blosum62 = substitution_matrices.load(\"BLOSUM62\")\n",
        "def f(x):\n",
        "    alignment,score,startend = local_pairwise_align_ssw(Protein(x),\n",
        "                                                        Protein(spike_prot),\n",
        "                                                        substitution_matrix=blosum62)\n",
        "    seqalign = ''.join([x.decode() for x in alignment[0].values])\n",
        "    refalign = ''.join([x.decode() for x in alignment[1].values])\n",
        "    return (seqalign, refalign)\n",
        "\n",
        "alignments = omicrondf['Raw'].apply(f).tolist()\n",
        "omicrondf['seqalign'] = [a[0] for a in alignments]\n",
        "omicrondf['refalign'] = [a[1] for a in alignments]\n",
        "\n",
        "def f(x):\n",
        "    seqalign, refalign = x\n",
        "    try:\n",
        "        if '-' in seqalign: # gap in aligned reference sequence indicates insertion\n",
        "            gaps = [0 if c == '-' else 1 for c in refalign]\n",
        "            return ('').join(itertools.compress(seqalign, gaps))\n",
        "        else:\n",
        "            return seqalign\n",
        "    except:\n",
        "        print('Error')\n",
        "        return np.nan\n",
        "\n",
        "# also remove stopcodons\n",
        "MAXLEN = 1273\n",
        "def remstop(x):\n",
        "    if x[-1] == '*' and len(x) > MAXLEN:\n",
        "        return x[:-1]\n",
        "    else:\n",
        "        return x\n",
        "omicrondf['Align'] = omicrondf[['seqalign','refalign']].apply(f, axis=1).apply(remstop)\n",
        "\n",
        "omicrondf.to_csv(FILELOC + 'omicron_spike_20211201.csv',index=False)"
      ],
      "execution_count": 132,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sYEwQVdzH7Hb"
      },
      "source": [
        "#Predict vaccine resistance"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iaun2ALSBUZw"
      },
      "source": [
        "Cromer D, Steain M, Reynaldi A, Schlub TE, Wheatley AK, Juno JA, Kent SJ, Triccas JA, Khoury DS, Davenport MP. Neutralising antibody titres as predictors of protection against SARS-CoV-2 variants and the impact of boosting: a meta-analysis. Lancet Microbe. 2021 Nov 15. doi: 10.1016/S2666-5247(21)00267-6. Epub ahead of print. PMID: 34806056; PMCID: PMC8592563.\n",
        "\n",
        "https://pubmed.ncbi.nlm.nih.gov/34806056/\n",
        "\n",
        "Kappa (B.1.617.1) has 6.8x less susceptible, while Delta (B.1.617.2) is 2.9x. Edara VV, Lai L, Sahoo MK, Floyd K, Sibai M, Solis D, Flowers MW, et al. Infection and vaccine-induced neutralizing antibody responses to the SARS-CoV-2 B.1.617.1 variant. bioRxiv [Preprint]. 2021 May 10:2021.05.09.443299. doi: 10.1101/2021.05.09.443299. Update in: N Engl J Med. 2021 Jul 7;: PMID: 34013272; PMCID: PMC8132229. https://www.nejm.org/doi/full/10.1056/NEJMc2107799\n",
        "\n",
        "Potentially relevant to AY.4:\n",
        "Arora, P., Kempf, A., Nehlmeier, I. et al. Delta variant (B.1.617.2) sublineages do not show increased neutralization resistance. Cell Mol Immunol 18, 2557–2559 (2021). https://doi.org/10.1038/s41423-021-00772-y\n",
        "\n",
        "Uriu K, Kimura I, Shirakawa K, Takaori-Kondo A, Nakada TA, Kaneda A, Nakagawa S, Sato K; Genotype to Phenotype Japan (G2P-Japan) Consortium. Neutralization of the SARS-CoV-2 Mu Variant by Convalescent and Vaccine Serum. N Engl J Med. 2021 Nov 3:NEJMc2114706. doi: 10.1056/NEJMc2114706. Epub ahead of print. PMID: 34731554; PMCID: PMC8609602. https://www.nejm.org/doi/full/10.1056/NEJMc2114706\n",
        "\n",
        "\"Virus neutralization assays, performed with the use of serum samples obtained from 13 persons who had recovered from Covid-19 who were infected early in the pandemic (April through September 2020), showed that the mu variant was 10.6 times as resistant to neutralization as the B.1 lineage virus (parental virus), which bears the D614G mutation (Figure 1B). Assays performed with serum samples obtained from 14 persons who had received the BNT162b2 vaccine showed that the mu variant was 9.1 as resistant as the parental virus (Figure 1C). Although the beta variant (a variant of concern) was thought to be the most resistant variant to date,3,4 the mu variant was 2.0 as resistant to neutralization by convalescent serum (Figure 1B) and 1.5 times as resistant to neutralization by vaccine serum as the beta variant (Figure 1C).\"\n",
        "For below: Parental, Alpha, Beta, Gamma, Delta, Epsilon (B.1.429, B.1.427), Lambda, Mu\n",
        "Convalescent: -, 2.3, 7.0, 2.9, 3.8, 2.1, 3.4, 10.6x reduction\n",
        "Vaccine: -, 2.3, 7.6, 2.9, 2.9, 3.0, 1.9, 9.1\n",
        "\n",
        "For Lambda:\n",
        "Infectivity and immune escape of the new SARS-CoV-2 variant of interest Lambda\n",
        "Mónica L. Acevedo, Luis Alonso-Palomares, Andrés Bustamante, Aldo Gaggero, Fabio Paredes, Claudia P. Cortés, Fernando Valiente-Echeverría, Ricardo Soto-Rifo\n",
        "medRxiv 2021.06.28.21259673; doi: https://doi.org/10.1101/2021.06.28.21259673\n",
        "\n",
        "\"We observed an increased infectivity mediated by the Lambda spike protein that was even higher than that of the D614G (lineage B) or the Alpha and Gamma variants. Compared to the Wild type (lineage A), neutralization was decreased by 3.05-fold for the Lambda variant while it was 2.33-fold for the Gamma variant and 2.03-fold for the Alpha variant.\"\n",
        "\n",
        "Delta variants\n",
        "\n",
        "Baoling Ying, Bradley Whitener, Laura A. VanBlargan, Ahmed O. Hassan, Swathi Shrihari, Chieh-Yu Liang, Courtney E. Karl, Samantha Mackin, Rita E. Chen, Natasha M. Kafai, Samuel H. Wilks, Derek J. Smith, Juan Manuel Carreño, Gagandeep Singh, Florian Krammer, Andrea Carfi, Sayda M. Elbashir, Darin K. Edwards, Larissa B. Thackray, Michael S. Diamond,  Protective activity of mRNA vaccines against ancestral and variant SARS-CoV-2 strains, Science Translational Medicine, 0, 0, (undefined). /doi/10.1126/scitranslmed.abm3302\n",
        "\n",
        "\"Geometric mean neutralizing Ab titers (GMTs) for the Pfizer/BioNtech BNT162b2-elicited plasma were respectively reduced 3.0-, 2.4-, and 4.1-fold for B.1.617.1, B.1.617.2, and B.1.617.2+ S (GMTs 86, 110, and 63), compared to G614 S (GMT: 260) (Fig. 1A and figs. S1 and S2). Moderna mRNA-1273-elicited plasma GMTs were reduced 4.1-, 2.6-, and 9.5-fold for B.1.617.1, B.1.617.2, and B.1.617.2+ S (GMTs 170, 270, and 74), respectively, compared to G614 S (GMT: 700) (Fig. 1B and figs. S1 and S2). The average neutralization potency of the Janssen Ad26.COV2.S-elicited plasma was reduced 3.0-, 2.4-, and 3.5-fold for B.1.617.1, B.1.617.2, and B.1.617.2+ S (GMTs 10, 13, and 8.8), respectively, compared to G614 S (GMT: 31)\""
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y_D-uFkQ3CWe"
      },
      "source": [
        "### Lineage -> Impact on neutralization mapping\n",
        "\n",
        "# Cromer D, Steain M, et al. + Edara VV, Lai L, et al. for Kappa (based on Cromer) +\n",
        "# Uriu et al. for B lineages (D614G) and Mu (B.1.621), Epsilon (B.1.427/429)\n",
        "lins = ['A', 'A.1', 'A.2.5',\n",
        "        'B', 'B.1', 'B.1.1.7',\n",
        "        'B.1.351', 'B.1.351.2', 'B.1.351.5',\n",
        "        'P.1', 'B.1.617.2', 'B.1.617.1',\n",
        "        'B.1.621', 'B.1.427', 'B.1.429',\n",
        "        'C.37']\n",
        "neutr = [1, 1, 1,\n",
        "         1/1.17, 1/1.17, 1/1.6,\n",
        "         1/8.8, 1/8.8, 1/8.8,\n",
        "         1/3.5, 1/3.9, 1/6.5,\n",
        "         1/10.5, 1/3.0, 1/3.0,\n",
        "         1/4.6]\n",
        "\n",
        "# Uriu et al. for all lineages (convalescent response) \n",
        "# Parental, Alpha, Beta, Gamma, Delta, Epsilon (B.1.429, B.1.427), Lambda, Mu Convalescent\n",
        "\n",
        "lins = ['B.1', 'B.1.1.7', 'B.1.351', 'P.1', 'B.1.617.2', 'B.1.427', 'B.1.429',\n",
        "        'C.37', 'C.37.1', 'C39', 'C40', 'B.1.621']\n",
        "neutr = 1 / np.array([1, 2.3, 7.0, 2.9, 3.8, 2.1, 2.1,\n",
        "                      3.4, 3.4, 3.4, 3.4, 10.6])\n",
        "\n",
        "# Uriu et al. and Ying et al.\n",
        "lins = ['B.1', 'B.1.1.7', 'B.1.351', 'P.1', 'B.1.617.2', 'B.1.427', 'B.1.429',\n",
        "        'C.37', 'B.1.621', 'B.1.617.1', 'B.1.617.2+']\n",
        "neutr = 1 / np.array([1, 2.3, 8.8, 2.9, 3.8, 2.1, 2.1,\n",
        "                      3.4, 13.3, 5.35, 9])\n",
        "\n",
        "\n",
        "assert len(lins) == len(neutr)\n",
        "numlins = len(lins)\n",
        "\n",
        "labelmap = {lins[k]:neutr[k] for k in range(numlins)}"
      ],
      "execution_count": 113,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KR0PQgrjIMMT"
      },
      "source": [
        "with open(FILELOC + 'aligned_sequence_lineages_1001.pkl', 'rb') as f:\n",
        "    _, tdf = pickle.load(f)\n",
        "tdf.reset_index(drop=False, inplace=True)\n",
        "\n",
        "# with open(FILELOC + 'raw_sequence_lineages_1001.pkl', 'rb') as f:\n",
        "#     _, tdf = pickle.load(f)\n",
        "# datadf.reset_index(drop=False, inplace=True)"
      ],
      "execution_count": 114,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vbd3twTRqdOr"
      },
      "source": [
        "tdf['Lineage_Orig'] = tdf['Lineage']\n",
        "def f(x):\n",
        "    if x[483] != \"E\":\n",
        "        return 'B.1.617.2+'\n",
        "    else:\n",
        "        return x\n",
        "tdf.loc[tdf[tdf.Lineage_Orig=='B.1.617.2'].index, 'Lineage'] = tdf.loc[tdf.Lineage_Orig=='B.1.617.2']['MaskedSeq'].apply(f)"
      ],
      "execution_count": 115,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2QkzNilJLu7U",
        "outputId": "6cadca14-62e0-45b4-91a2-5312ab92c164"
      },
      "source": [
        "tdf = tdf[tdf.Lineage.isin(lins)]\n",
        "print([len(tdf[tdf.Count >= c]) for c in range(1,6)])\n",
        "# filter sequences by removing any with 'X' or '*' characters\n",
        "tdf = tdf[~tdf.MaskedSeq.str.contains('X|\\*')]\n",
        "print([len(tdf[tdf.Count >= c]) for c in range(1,6)])"
      ],
      "execution_count": 116,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[220603, 30331, 17778, 12938, 10210]\n",
            "[28659, 12006, 8103, 6277, 5176]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oRxe35iUhYF3",
        "outputId": "929a21cb-fa9c-4c34-828c-041e16bcca9f"
      },
      "source": [
        "MINCOUNT = 1\n",
        "tdf[tdf.Count >= MINCOUNT]['Lineage'].map(labelmap).value_counts()"
      ],
      "execution_count": 117,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.434783    18816\n",
              "0.344828     3052\n",
              "1.000000     2568\n",
              "0.476190     1672\n",
              "0.113636      865\n",
              "0.294118      730\n",
              "0.186916      442\n",
              "0.075188      413\n",
              "0.111111      101\n",
              "Name: Lineage, dtype: int64"
            ]
          },
          "metadata": {},
          "execution_count": 117
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8Jx2hdJamtpC"
      },
      "source": [
        "traindf = tdf[tdf.Count >= MINCOUNT].copy()\n",
        "traindf['label'] = traindf.Lineage.map(labelmap)"
      ],
      "execution_count": 118,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yiIZ1m-lid1q",
        "outputId": "1a942312-fb13-42cd-e650-5a357493cf8a"
      },
      "source": [
        "MAX_SAMPLE = 200\n",
        "\n",
        "# from imblearn.under_sampling import RandomUnderSampler, EditedNearestNeighbours\n",
        "# SRS = 2077\n",
        "# sm = RandomUnderSampler(sampling_strategy='majoriy', random_state=SRS, replacement=False)\n",
        "# xt, yt = sm.fit_resample(xtrain, traindf.Lineage.map({lins[ind]:ind for ind in range(len(lins))}))\n",
        "\n",
        "### Manual undersampling\n",
        "# dropindices = {}\n",
        "# for label in np.unique(traindf.label):\n",
        "#     n = len(traindf[traindf.label==label])\n",
        "#     if n > MAX_SAMPLE:\n",
        "#         print(1/label, len(traindf))\n",
        "#         dropindices[label] = np.random.choice(traindf[traindf.label==label].index, n-MAX_SAMPLE, replace=False)\n",
        "#         traindf.drop(dropindices[label], inplace=True)\n",
        "#         print(len(traindf))\n",
        "\n",
        "## Take the top samples\n",
        "dropindices = {}\n",
        "for label in np.unique(traindf.label):\n",
        "    n = len(traindf[traindf.label==label])\n",
        "    if n > MAX_SAMPLE:\n",
        "        print(1/label, len(traindf))\n",
        "        dropindices[label] = traindf[traindf.label==label].sort_values(by='Count',\n",
        "                                                                       ascending=False).index[MAX_SAMPLE:]   \n",
        "        traindf.drop(dropindices[label], inplace=True)\n",
        "        print(len(traindf))"
      ],
      "execution_count": 120,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "13.3 28659\n",
            "28446\n",
            "8.8 28446\n",
            "27781\n",
            "5.35 27781\n",
            "27539\n",
            "3.4 27539\n",
            "27009\n",
            "2.9 27009\n",
            "24157\n",
            "2.3 24157\n",
            "5541\n",
            "2.1 5541\n",
            "4069\n",
            "1.0 4069\n",
            "1701\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dFEKPhX8myk8"
      },
      "source": [
        "xtrain = tokenize_sequences(traindf, 'MaskedSeq', ismlen)\n",
        "# xtrain = tokenize_sequences(traindf, 'Spike', ismlen)\n",
        "\n",
        "ytrain = traindf.Lineage.map(labelmap).values"
      ],
      "execution_count": 122,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C51zCr_ychel"
      },
      "source": [
        "NUM_EPOCHS = 200\n",
        "PATIENCE = 50   # 10\n",
        "\n",
        "early_stopping = tf.keras.callbacks.EarlyStopping(\n",
        "    monitor = 'loss',\n",
        "    verbose = 1,\n",
        "    patience = PATIENCE,\n",
        "    mode = 'auto',\n",
        "    min_delta = 0,\n",
        "    baseline = 0.2,\n",
        "    restore_best_weights = True\n",
        "    )"
      ],
      "execution_count": 123,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eHaM98wKMKUf"
      },
      "source": [
        "n_epochs = 0\n",
        "while (n_epochs < (PATIENCE + 5)):\n",
        "    tf.keras.backend.clear_session()\n",
        "    with tpu_strategy.scope():\n",
        "        model = reset_model(regress=True, singleclass=False, multiclass=False,\n",
        "                            output_multiheadatt=False, use_att=True)\n",
        "    history = model.fit(xtrain, ytrain,\n",
        "                        batch_size = BATCH_SIZE,\n",
        "                        epochs = NUM_EPOCHS,\n",
        "                        verbose = 2,\n",
        "                        callbacks = [early_stopping],\n",
        "                        )\n",
        "    n_epochs = len(history.history['loss'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X7RdMVyLWP0r"
      },
      "source": [
        "WEIGHTFILE = \"coronavirus_vaxresist_20211201.h5\"\n",
        "\n",
        "# model.save_weights(FILELOC + WEIGHTFILE, save_format='h5', overwrite=True)\n",
        "# with open(FILELOC + WEIGHTFILE + '_setup.pkl','wb') as f:\n",
        "#     pickle.dump([tdf, traindf, dropindices, xtrain, ytrain, MINCOUNT, MAX_SAMPLE, labelmap], f)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q388COdM9onj"
      },
      "source": [
        "with tpu_strategy.scope():\n",
        "    model = reset_model(regress=True, singleclass=False, multiclass=False,\n",
        "                        output_multiheadatt=True, use_att=True)\n",
        "    model.load_weights(FILELOC + WEIGHTFILE)\n",
        "    model.compile()\n",
        "with open(FILELOC + WEIGHTFILE + '_setup.pkl','rb') as f:\n",
        "    tdf, traindf, dropindices, xtrain, ytrain, MINCOUNT, MAX_SAMPLE, labelmap = pickle.load(f)"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CsVRlQZaWjxE"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KI4BPoADWfyF"
      },
      "source": [
        "from sklearn.metrics import mean_squared_error\n",
        "from collections import Counter"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vzf2r9ojH9QH"
      },
      "source": [
        "# omicron = pd.read_csv(FILELOC + \"omicron_spike_20211128.csv\")\n",
        "omicron = pd.read_csv(FILELOC + \"omicron_spike_20211201.csv\")\n",
        "omicron['seqlen'] = omicron.Align.apply(len)\n",
        "omicron.drop(omicron[omicron.seqlen<1273].index, inplace=True)\n",
        "omicron.reset_index(inplace=True, drop=True)\n",
        "otok = tokenize_sequences(omicron,'Align',ismlen)"
      ],
      "execution_count": 141,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O4OuG898uDZ7"
      },
      "source": [
        "omicron['Predict'] = model.predict(otok)"
      ],
      "execution_count": 142,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3zAk09t7Y1cp",
        "outputId": "2f5f0aa9-8ea9-4cd9-fb6a-c2bfab24a660"
      },
      "source": [
        "np.mean(omicron.Predict)"
      ],
      "execution_count": 143,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.10882622003555298"
            ]
          },
          "metadata": {},
          "execution_count": 143
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bO2DoQROY4n3",
        "outputId": "45a94356-48a9-45e1-f40d-aa2374272140"
      },
      "source": [
        "np.max(omicron.Predict)"
      ],
      "execution_count": 144,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.1921105980873108"
            ]
          },
          "metadata": {},
          "execution_count": 144
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GBnN5EdJY24t",
        "outputId": "4f97d949-8de3-4141-e4cf-b1cdf3dcb8b8"
      },
      "source": [
        "np.min(omicron.Predict)"
      ],
      "execution_count": 145,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.05763649940490723"
            ]
          },
          "metadata": {},
          "execution_count": 145
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gZhrPVaCY6on",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "103a34aa-ca2e-4b4e-8855-a2d8206de3d7"
      },
      "source": [
        "Counter(omicron.Predict)"
      ],
      "execution_count": 147,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Counter({0.05763649940490723: 12,\n",
              "         0.11000391840934753: 1,\n",
              "         0.11113205552101135: 1,\n",
              "         0.1113305389881134: 2,\n",
              "         0.11178094148635864: 87,\n",
              "         0.17104101181030273: 1,\n",
              "         0.17186003923416138: 2,\n",
              "         0.1859753429889679: 1,\n",
              "         0.1921105980873108: 1})"
            ]
          },
          "metadata": {},
          "execution_count": 147
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rANlj87iH9Hc"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}