Covid_Data_Creation_github.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K5Ianj720LwK"
      },
      "source": [
        "#Intitialization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hUg9gqAoobKi"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "from datetime import datetime\n",
        "from dateutil.parser import parse as dateparse\n",
        "import time\n",
        "from collections import Counter\n",
        "from tqdm.notebook import tqdm\n",
        "import pandas as pd\n",
        "from pandas import read_csv\n",
        "from pandas import DataFrame\n",
        "from datetime import datetime\n",
        "from dateutil.parser import parse as dateparse\n",
        "import pickle\n",
        "import csv\n",
        "\n",
        "%tensorflow_version 2.x\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CYy877h90eA1"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive, files\n",
        "# drive.mount('/content/drive')\n",
        "\n",
        "FILELOC = \"/content/drive/My Drive/COVID_Python/\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SLknMJwnPCfK",
        "outputId": "f2e0e133-a1c6-4ec1-d631-30ef262251bb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Not connected to a TPU runtime.\n"
          ]
        }
      ],
      "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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0mBDnd0TDn42"
      },
      "source": [
        "#Dataset creation Pipelines"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mPJhCtJs0e_4"
      },
      "source": [
        "#Read files and preprocess patient data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dT5x-05IobKk"
      },
      "outputs": [],
      "source": [
        "metadf = pd.read_csv(FILELOC + \"patient_spike_ism_20210811.tsv\", \"\\t\")\n",
        "metadf['seqs'] = metadf['MaskedSeq'].apply(lambda x:np.array(list(x)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ejAqJKGpNXyN"
      },
      "outputs": [],
      "source": [
        "metadf = pd.read_csv(FILELOC + 'align0912_spike.csv')\n",
        "print(len(metadf))\n",
        "patdf = pd.read_csv(FILELOC + 'GISAID_patient_metadata_20210812_0912.tsv', '\\t')\n",
        "print(len(patdf))\n",
        "metadf = pd.merge(left=metadf[['SequenceID','MaskedSeq']], right = patdf,\n",
        "                  left_on = 'SequenceID', right_on='Accession ID')\n",
        "metadf['seqs'] = metadf['MaskedSeq'].apply(lambda x:np.array(list(x))) \n",
        "print(len(metadf))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v3B-VX0SkUVW"
      },
      "outputs": [],
      "source": [
        "!tar -xf \"/content/drive/My Drive/COVID_Python/metadata_tsv_2021_09_15.tar.xz\"\n",
        "mdf = pd.read_csv(\"metadata.tsv\", \"\\t\")\n",
        "sdf = pd.read_csv(FILELOC + \"ism0801/ism0801_Spike.csv\")\n",
        "mdf = pd.merge(left=mdf, right=sdf[['SequenceID','MaskedSeq']], left_on=\"Accession ID\", right_on=\"SequenceID\", how=\"inner\")\n",
        "mdf['Country'] = mdf['Location'].apply(lambda x: x.split(' / ')[1])\n",
        "from dateutil.parser import parse as dateparse\n",
        "REFDATE = '2019-12-01'\n",
        "mdf['Relative date'] = mdf['Collection date'].apply(lambda x: (dateparse(x) - dateparse(REFDATE)).days)\n",
        "with open(FILELOC + 'covid_spike_0801_alldata.pkl', 'wb') as f:\n",
        "    pickle.dump(mdf, f)\n",
        "print(\"Dropping invalid sequence values\", len(mdf))\n",
        "mdf.drop(mdf[mdf.MaskedSeq.str.contains('\\*')].index, inplace=True)\n",
        "print(len(mdf))\n",
        "mdf.drop(mdf[mdf.MaskedSeq.str.contains('X')].index, inplace=True)\n",
        "print(len(mdf))\n",
        "mdf.reset_index(drop=True, inplace=True)\n",
        "print(\"Dropping invalid date values\")\n",
        "print(len(mdf))\n",
        "mdf.drop(mdf[mdf[\"Relative date\"] < 0].index, inplace=True)\n",
        "print(len(mdf))\n",
        "mdf.reset_index(drop=True, inplace=True)\n",
        "sdf = mdf.groupby([\"MaskedSeq\"])[\"Relative date\"].apply(list).to_frame().reset_index(drop=False)\n",
        "sdf.to_csv(FILELOC + \"covid_spike_0801_reldate.csv\", index=False)\n",
        "\n",
        "data = pd.read_csv(FILELOC + \"covid_spike_0801_reldate.csv\")\n",
        "data.rename(columns={'MaskedSeq':'ISM'}, inplace=True)\n",
        "data['Relative date'] = data['Relative date'].apply(ast.literal_eval)\n",
        "data[\"Mindate\"] = data[\"Relative date\"].apply(min)\n",
        "data[\"Datebin\"] = np.digitize(data['Mindate'], np.histogram(data['Mindate'], bins=200)[1])\n",
        "data[\"Label\"] = (data.Datebin - min(data.Datebin)) / (max(data.Datebin) - min(data.Datebin))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iubI5pp8ypNN"
      },
      "outputs": [],
      "source": [
        "metadf['seqs'] = metadf['MaskedSeq'].apply(lambda x:np.array(list(x)))\n",
        "\n",
        "metadf['ISM'] = metadf['MaskedSeq']\n",
        "\n",
        "data = metadf[['seqs', 'ISM', 'Pango lineage', 'Country', 'Relative date', 'Gender',\n",
        "               'Patient age', 'Clade', 'AA Substitutions']].copy()\n",
        "data.rename(columns={'seqs':'FullSeq', 'Pango lineage':'Pango',\n",
        "                     'Patient age':'Age', 'AA Substitutions':'SNPs',\n",
        "                     'Lineage':'Pango', 'Relative date':'Date',\n",
        "                     }, inplace=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qglM2e4mh60_"
      },
      "outputs": [],
      "source": [
        "with open(FILELOC+\"covid_spike_data_0801.pkl\", \"wb\") as f:\n",
        "    pickle.dump(data, f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "emWQyoekobKn"
      },
      "outputs": [],
      "source": [
        "metadf['ISM'] = metadf['MaskedSeq']\n",
        "\n",
        "entropylist = pd.read_csv(FILELOC + \"entropy0801/entropy0801_Spike.csv\")\n",
        "entropylist = np.array(entropylist['Spike'])\n",
        "\n",
        "ismlist = np.where(entropylist < -0.01)[0]\n",
        "print(len(ismlist))\n",
        "metadf['ISM'] = metadf['seqs'].apply(lambda x: ''.join(x[ismlist]))\n",
        "metadf['ISM'] = metadf['MaskedSeq'].apply(lambda x:x[400:800])\n",
        "\n",
        "metadf['ISM'] = metadf['seqs'].apply(lambda x: ''.join(x[np.where(attmean_all1 - 1/ismlen > 0)[0]]))\n",
        "\n",
        "rng = np.random.default_rng(seed = 6127)\n",
        "rng = np.random.default_rng(seed = 172)\n",
        "rng = np.random.default_rng(seed = 12996)\n",
        "num_indices = 373\n",
        "rand_indices = np.sort(rng.choice(range(len(metadf.MaskedSeq[0])), size = num_indices, replace=False))\n",
        "metadf['ISM'] = metadf['seqs'].apply(lambda x: ''.join(x[rand_indices]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D8gSE-GFNaIZ"
      },
      "outputs": [],
      "source": [
        "data = metadf[['seqs', 'ISM', 'Pango lineage', 'Country', 'reldate', 'Gender_x', 'Patient age_x',\n",
        "               'Clade_x', 'AA Substitutions_x', 'Corrected status']].copy()\n",
        "data.rename(columns={'seqs':'FullSeq', 'Pango lineage':'Pango',\n",
        "                     'Gender_x':'Gender_Raw', 'Patient age_x':'Age_Raw',\n",
        "                     'Clade_x': 'Clade', 'AA Substitutions_x':'SNPs',\n",
        "                     'reldate':'Date', 'Corrected status':'Status',\n",
        "                    }, inplace=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gji-1X8qobKo"
      },
      "outputs": [],
      "source": [
        "data = metadf.copy()\n",
        "data.rename(columns={'Corrected status':'Status', 'Gender':'Gender_Raw',\n",
        "                     'Patient age':'Age_Raw'}, inplace=True)\n",
        "Location_Parsed = data['Location'].apply(lambda x: x.split('/'))\n",
        "data['Country'] = [x[1].strip() for x in Location_Parsed]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z_094gJsNdWn"
      },
      "outputs": [],
      "source": [
        "statusmapdf = pd.read_csv(FILELOC + \"patient_status_map_20210811.tsv\", \"\\t\")\n",
        "catmapdf = pd.read_csv(FILELOC + \"patient_status_reduced_map_20210811.csv\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "80m-rWS5obKp"
      },
      "outputs": [],
      "source": [
        "# statusmapdf = pd.read_csv(FILELOC + \"patient_status_map_20210812_0912.tsv\", \"\\t\")\n",
        "# catmapdf = pd.read_csv(FILELOC + \"patient_status_reduced_map_20210811.csv\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dy2u10uFobKr"
      },
      "outputs": [],
      "source": [
        "statusmap = dict(zip(statusmapdf['Original'].apply(lambda x: x.lower()),\n",
        "                     statusmapdf['Map1'].apply(lambda x: x.lower())))\n",
        "catmap = dict(zip(catmapdf['Status'].apply(lambda x: x.lower()),\n",
        "                  catmapdf['Category'].apply(lambda x: x.lower())))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OHWdN2EBobKr"
      },
      "outputs": [],
      "source": [
        "data['Category'] = data['Status'].apply(lambda x: x.lower()).map(statusmap).map(catmap)\n",
        "data['Category'] = data['Category'].fillna('unknown')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C30aL5wYobKs"
      },
      "outputs": [],
      "source": [
        "labelmap = {'alive' : -1,\n",
        "            'asymptomatic' : 0,\n",
        "            'dead' : 1,\n",
        "            'hospitalized' : 1,\n",
        "            'mild' : 0,\n",
        "            'moderate' : 1,\n",
        "            'released' : 1,\n",
        "            'screening' : 0,\n",
        "            'severe' : 1,\n",
        "            'symptomatic' : -1,\n",
        "            'unknown' : -1,\n",
        "           }\n",
        "\n",
        "ORIGINAL\n",
        "labelmap = {'alive' : 0,\n",
        "            'asymptomatic' : 0,\n",
        "            'dead' : 1,\n",
        "            'hospitalized' : 1,\n",
        "            'mild' : 0,\n",
        "            'moderate' : 1,\n",
        "            'released' : 1,\n",
        "            'screening' : 0,\n",
        "            'severe' : 1,\n",
        "            'symptomatic' : 1,\n",
        "            'unknown' : -1,\n",
        "           }"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QPm2UMpkobKs"
      },
      "outputs": [],
      "source": [
        "data['Label'] = data['Category'].map(labelmap).astype(int)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0eTLLbbiLjLN"
      },
      "outputs": [],
      "source": [
        "for clade in data.Clade.value_counts().index:\n",
        "    print(clade, ':', 'Label 0 =', len(data[(data.Clade==clade) & (data.Label==0)]),\n",
        "          'Label 1 =', len(data[(data.Clade==clade) & (data.Label==1)]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ifx79AgHobKu"
      },
      "outputs": [],
      "source": [
        "def f(x):\n",
        "    if x[0].lower() == 'f':\n",
        "        return 'Female'\n",
        "    elif x[0].lower() == 'm':\n",
        "        return 'Male'\n",
        "    else:\n",
        "        return 'Unknown'\n",
        "\n",
        "print('Invalid gender values =', len(data[data.Gender_Raw.isna()]))\n",
        "data['Gender_Raw'] = data['Gender_Raw'].fillna('Unknown')\n",
        "data['Gender'] = data['Gender_Raw'].apply(f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dWqDGE-0obKt"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    agemapdf\n",
        "except:\n",
        "    # the map file should define any \"unknown\" as -1\n",
        "    agemapdf = pd.read_csv(FILELOC + 'age_map_20210811.tsv', '\\t', encoding='ISO-8859-1', dtype={'Original':'str', 'Processed':'Int64'})\n",
        "    agemap = dict(zip(agemapdf.Original, agemapdf.Processed))\n",
        "    \n",
        "data['Age'] = data['Age_Raw'].map(agemap).fillna(-1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XVaij-n_cNDb"
      },
      "outputs": [],
      "source": [
        "# try:\n",
        "#     agemapdf\n",
        "# except:\n",
        "#     # the map file should define any \"unknown\" as -1\n",
        "#     agemapdf = pd.read_csv(FILELOC + 'age_map_20210812_0912.tsv', '\\t', encoding='ISO-8859-1', dtype={'Original':'str', 'Processed':'Int64'})\n",
        "#     agemap = dict(zip(agemapdf.Original, agemapdf.Processed))\n",
        "    \n",
        "# data['Age'] = data['Age_Raw'].map(agemap).fillna(-1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AFZXp1Nyxm46"
      },
      "outputs": [],
      "source": [
        "some countries may have an unexpected lower or upper case, so standardize\n",
        "\n",
        "data['Country'] = data['Country'].apply(lambda x: x.title())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KHjE3Lqyit_a"
      },
      "outputs": [],
      "source": [
        "data_backup = data.copy()\n",
        "data = data_backup.copy()\n",
        "\n",
        "train_backup = train.copy(); test_backup = test.copy()\n",
        "train = train_backup.copy(); test = test_backup.copy()\n",
        "print(len(train)); train.drop(train[train.Age==-1].index, inplace=True); print(len(train))\n",
        "print(len(test)); test.drop(test[test.Age==-1].index, inplace=True); print(len(test))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GtrXMSqIobKz"
      },
      "outputs": [],
      "source": [
        "data_backup = data.copy()\n",
        "\n",
        "print(\"Dropping invalid ISM values\", len(data))\n",
        "data.drop(data[data.ISM.str.contains('\\*')].index, inplace=True)\n",
        "print(len(data))\n",
        "data.drop(data[data.ISM.str.contains('X')].index, inplace=True)\n",
        "print(len(data))\n",
        "data.reset_index(drop=True, inplace=True)\n",
        "\n",
        "# print(\"Dropping invalid age values:\", len(data))\n",
        "# data.drop(data.loc[data.Age < 0].index, inplace=True)\n",
        "# data.reset_index(drop=True, inplace=True)\n",
        "# print(len(data))\n",
        "\n",
        "# print(\"Dropping invalid gender values:\", len(data))\n",
        "# data.drop(data.loc[data.Gender == \"Unknown\"].index, inplace=True)\n",
        "# data.reset_index(drop=True, inplace=True)\n",
        "# print(len(data))\n",
        "\n",
        "print(\"Dropping invalid label values\", len(data))\n",
        "data.drop(data[data.Label==-1].index, inplace=True)\n",
        "print(len(data))\n",
        "data.reset_index(drop=True, inplace=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sCkTeNhnxUR9"
      },
      "outputs": [],
      "source": [
        "countdata = data.groupby([\"ISM\",\"Label\"]).size().to_frame().rename(columns={0:\"Count\"})\n",
        "newdata = pd.DataFrame(data.ISM.value_counts()).rename(columns={\"ISM\":\"Count\"})\n",
        "newdata['Count_Dict'] = [countdata.loc[ind].to_dict()[\"Count\"] for ind in newdata.index]\n",
        "newdata = newdata.reset_index(drop=False).rename(columns={\"index\":\"ISM\"})\n",
        "label_list = np.zeros(len(newdata))\n",
        "for ind in newdata.index:\n",
        "    elem = newdata.Count_Dict[ind]\n",
        "    if len(elem) == 2:\n",
        "        label_list[ind] = newdata.Count_Dict[ind][1] / newdata.Count[ind]\n",
        "    elif 0 in elem:\n",
        "        label_list[ind] = 0\n",
        "    elif 1 in elem:\n",
        "        label_list[ind] = 1\n",
        "newdata['Label'] = label_list\n",
        "totalcount = np.sum(newdata.Count)\n",
        "newdata['sample_weight'] = newdata['Count'].apply(lambda x: x/totalcount*10)\n",
        "\n",
        "data = newdata.copy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EjLl8Ua4cR8a"
      },
      "outputs": [],
      "source": [
        "data.to_csv(FILELOC + 'covid_crossval_20210919_AllData.csv', index=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QRv3aqcj1I9r"
      },
      "outputs": [],
      "source": [
        "with open(FILELOC+\"covid_spike_data_0801.pkl\", \"rb\") as f:\n",
        "    data = pickle.load(f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "svjZ0NDz_QfS"
      },
      "outputs": [],
      "source": [
        "entropylist = pd.read_csv(FILELOC + \"entropy0801/entropy0801_Spike.csv\")\n",
        "entropylist = np.array(entropylist['Spike'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yWyihr6M_HJQ"
      },
      "outputs": [],
      "source": [
        "ismlist = np.where(entropylist < -0.01)[0]\n",
        "print(len(ismlist))\n",
        "data['ISM'] = data['FullSeq'].apply(lambda x: ''.join(x[ismlist]))\n",
        "\n",
        "print(\"Dropping invalid ISM values\", len(data))\n",
        "data.drop(data[data.ISM.str.contains('\\*')].index, inplace=True)\n",
        "print(len(data))\n",
        "data.drop(data[data.ISM.str.contains('X')].index, inplace=True)\n",
        "print(len(data))\n",
        "data.reset_index(drop=True, inplace=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LxGKmRYWGCTZ"
      },
      "source": [
        "###Combine and pre-process patient metadata files"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "t5nZ-GLmqUiT",
        "outputId": "fc8f5649-22c1-4478-8516-73f88b5f49cb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "113680 23073 18792\n"
          ]
        }
      ],
      "source": [
        "origdf = pd.read_csv(FILELOC + \"GISAID_patient_metadata_20210811.tsv\", '\\t')\n",
        "patdf = pd.read_csv(FILELOC + 'GISAID_patient_metadata_20210812_0912.tsv', '\\t')\n",
        "newdf = pd.read_csv(FILELOC + 'GISAID_patient_metadata_20210913_1004_raw.tsv', '\\t')\n",
        "print(len(origdf), len(patdf), len(newdf))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZFdj6oeHbsYj"
      },
      "outputs": [],
      "source": [
        "def proc_df(mdf2, patientstatusmapfile, agemapfile):\n",
        "\n",
        "    labelmap = {'alive' : -1,\n",
        "                'asymptomatic' : 0,\n",
        "                'dead' : 1,\n",
        "                'hospitalized' : 1,\n",
        "                'mild' : 0,\n",
        "                'moderate' : 0,\n",
        "                'released' : 1,\n",
        "                'screening' : 0,\n",
        "                'severe' : 1,\n",
        "                'symptomatic' : -1,\n",
        "                'unknown' : -1,\n",
        "            }\n",
        "\n",
        "    mdf2.rename(columns={'Corrected status':'Status', 'Gender':'Gender_Raw',\n",
        "                        'Patient age':'Age_Raw'}, inplace=True)\n",
        "    Location_Parsed = mdf2['Location'].apply(lambda x: x.split('/'))\n",
        "    mdf2['Country'] = [x[1].strip() for x in Location_Parsed]\n",
        "    agemapdf = pd.read_csv(agemapfile, '\\t', encoding='ISO-8859-1', dtype={'Original':'str', 'Processed':'Int64'})\n",
        "    agemap = dict(zip(agemapdf.Original, agemapdf.Processed))\n",
        "    mdf2['Age'] = mdf2['Age_Raw'].map(agemap).fillna(-1)\n",
        "\n",
        "    statusmapdf = pd.read_csv(patientstatusmapfile, \"\\t\")\n",
        "    catmapdf = pd.read_csv(FILELOC + \"patient_status_reduced_map_20210811.csv\")\n",
        "    statusmap = dict(zip(statusmapdf['Original'].apply(lambda x: x.lower()),\n",
        "                        statusmapdf['Map1'].apply(lambda x: x.lower())))\n",
        "    catmap = dict(zip(catmapdf['Status'].apply(lambda x: x.lower()),\n",
        "                    catmapdf['Category'].apply(lambda x: x.lower())))\n",
        "    mdf2['Category'] = mdf2['Status'].apply(lambda x: x.lower()).map(statusmap).map(catmap)\n",
        "    mdf2['Category'] = mdf2['Category'].fillna('unknown')\n",
        "    mdf2['Label'] = mdf2['Category'].map(labelmap).astype(int)\n",
        "\n",
        "    return mdf2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1Vb6zkxRbsNH"
      },
      "outputs": [],
      "source": [
        "odf = proc_df(origdf, FILELOC+'patient_status_map_20210811.tsv',\n",
        "              FILELOC+'age_map_20210811.tsv')\n",
        "pdf = proc_df(patdf, FILELOC+'patient_status_map_20210812_0912.tsv',\n",
        "              FILELOC+'age_map_20210812_0912.tsv')\n",
        "ndf = proc_df(newdf, FILELOC+'patient_status_map_20210913_1004.tsv',\n",
        "              FILELOC+'age_map_20210913_1004.tsv')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4IieNq9_qoJM"
      },
      "outputs": [],
      "source": [
        "# !tar -xvf '/content/drive/My Drive/COVID_Python/results1001_Spike.tar.gz' -C '/content/drive/My Drive/COVID_Python/'\n",
        "# seqs = pd.read_csv(FILELOC + 'results1001_Spike.csv')\n",
        "# seqlist = pd.concat([origdf['Accession ID'],patdf['Accession ID'],newdf['Accession ID']],\n",
        "#                     axis=0, ignore_index=True).to_frame()\n",
        "# seqlist.loc[len(seqlist)] = 'EPI_ISL_402124' # add reference sequence\n",
        "# print(len(seqs), len(seqlist))\n",
        "# selseqs = pd.merge(left=seqs,right=seqlist,left_on='SequenceID',right_on='Accession ID')\n",
        "# print(len(selseqs))\n",
        "# selseqs.to_csv(FILELOC + 'results_1001_Spike_selected.csv', index=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5u-xALxl1Vpv"
      },
      "outputs": [],
      "source": [
        "alldf = pd.concat([odf, pdf, ndf], ignore_index=True, axis=0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_owsGSxxOMkA"
      },
      "source": [
        "##Create cross-validation datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JuxuP2NDSUSI"
      },
      "outputs": [],
      "source": [
        "data = pd.read_csv(FILELOC + 'covid_rawseqs_crossval/covid_patient_0912_raw_sequences.csv')\n",
        "data.reset_index(drop=True, inplace=True)\n",
        "np.savetxt(FILELOC + 'covid_rawseqs_crossval/TrainLabels.csv', data.Label)\n",
        "with open(FILELOC + 'covid_rawseqs_crossval/TrainData.csv', 'w') as f:\n",
        "    csv.writer(f).writerows(data.ISM.values)\n",
        "\n",
        "random_state_split = 32687     # for splitting\n",
        "NSplits = 10  ## Number of folds\n",
        "skf = KFold(n_splits = NSplits, shuffle = True, random_state = random_state_split)\n",
        "fold=0\n",
        "for train_index, test_index in skf.split(data):\n",
        "    np.savetxt(FILELOC + 'covid_rawseqs_crossval/TrainIndex_'+str(fold)+'.csv', train_index, '%i', delimiter=',')\n",
        "    np.savetxt(FILELOC + 'covid_rawseqs_crossval/TestIndex_'+str(fold)+'.csv', test_index, '%i', delimiter=',')\n",
        "    fold += 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2XyMEUVDTEad"
      },
      "outputs": [],
      "source": [
        "data = pd.read_csv(FILELOC + 'covid_aligned_crossval/covid_patient_0912_aligned_sequences.csv')\n",
        "data.reset_index(drop=True, inplace=True)\n",
        "np.savetxt(FILELOC + 'covid_aligned_crossval/TrainLabels.csv', data.Label)\n",
        "with open(FILELOC + 'covid_aligned_crossval/TrainData.csv', 'w') as f:\n",
        "    csv.writer(f).writerows(data.ISM.values)\n",
        "\n",
        "random_state_split = 32687     # for splitting\n",
        "NSplits = 10  ## Number of folds\n",
        "skf = KFold(n_splits = NSplits, shuffle = True, random_state = random_state_split)\n",
        "fold=0\n",
        "for train_index, test_index in skf.split(data):\n",
        "    np.savetxt(FILELOC + 'covid_aligned_crossval/TrainIndex_'+str(fold)+'.csv', train_index, '%i', delimiter=',')\n",
        "    np.savetxt(FILELOC + 'covid_aligned_crossval/TestIndex_'+str(fold)+'.csv', test_index, '%i', delimiter=',')\n",
        "    fold += 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vrak9sZpSSlg"
      },
      "outputs": [],
      "source": [
        "data = pd.read_csv(FILELOC + 'covid_filtered_crossval/covid_patient_0912_filtered_sequences.csv')\n",
        "data.reset_index(drop=True, inplace=True)\n",
        "np.savetxt(FILELOC + 'covid_filtered_crossval/TrainLabels.csv', data.Label)\n",
        "with open(FILELOC + 'covid_filtered_crossval/TrainData.csv', 'w') as f:\n",
        "    csv.writer(f).writerows(data.ISM.values)\n",
        "\n",
        "random_state_split = 32687     # for splitting\n",
        "NSplits = 10  ## Number of folds\n",
        "skf = KFold(n_splits = NSplits, shuffle = True, random_state = random_state_split)\n",
        "fold=0\n",
        "for train_index, test_index in skf.split(data):\n",
        "    np.savetxt(FILELOC + 'covid_filtered_crossval/TrainIndex_'+str(fold)+'.csv', train_index, '%i', delimiter=',')\n",
        "    np.savetxt(FILELOC + 'covid_filtered_crossval/TestIndex_'+str(fold)+'.csv', test_index, '%i', delimiter=',')\n",
        "    fold += 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dy8bpP78Vw5k"
      },
      "source": [
        "##Generate patient datasets (aligned, raw, filtered)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oEgdVyMfAFlM"
      },
      "outputs": [],
      "source": [
        " alldf = pd.read_csv(FILELOC + 'covid_alldata_merged_align1001.csv')\n",
        "# alldf = pd.read_csv(FILELOC + 'covid_alldata_merged_rawseq1001.csv')\n",
        "\n",
        "alldf.rename(columns={'MaskedSeq':'ISM'}, inplace=True)\n",
        "# alldf.rename(columns={'Spike':'ISM'}, inplace=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3cvCaCR1Pwte",
        "outputId": "d2088c8c-9059-45ef-fb18-4e7a396ba98a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "147168\n",
            "105240\n"
          ]
        }
      ],
      "source": [
        "print(len(alldf))\n",
        "alldf.drop(alldf[alldf['ISM'].str.contains('\\*')].index, inplace=True)\n",
        "alldf.drop(alldf[alldf['ISM'].str.contains('X')].index, inplace=True)\n",
        "print(len(alldf))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LlrjQqD5NbB_",
        "outputId": "470ed174-a729-47b1-bd79-fe12a0ff0af8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Dropping invalid label values 105240\n",
            "54081\n"
          ]
        }
      ],
      "source": [
        "print(\"Dropping invalid label values\", len(alldf))\n",
        "alldf.drop(alldf[alldf.Label==-1].index, inplace=True)\n",
        "print(len(alldf))\n",
        "alldf.reset_index(drop=True, inplace=True)\n",
        "\n",
        "# print(\"Merging with sequences\")\n",
        "# print(len(alldf))\n",
        "# alldf_rawseqs = pd.merge(left=alldf, right=seqs, left_on='Accession ID', right_on='SequenceID')\n",
        "# # alldf_rawseqs = pd.merge(left=alldf, right=df, left_on='Accession ID', right_on='SequenceID')\n",
        "# print(len(alldf_rawseqs))\n",
        "# alldf_rawseqs.rename(columns={'MaskedSeq':'ISM'}, inplace=True)\n",
        "# alldf_rawseqs.rename(columns={'Spike':'ISM'}, inplace=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SxTaXG1Fa5Jl",
        "outputId": "81707ef7-5c2f-455c-df49-e47b25b5e08f"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py:2718: DtypeWarning: Columns (8,39) have mixed types.Specify dtype option on import or set low_memory=False.\n",
            "  interactivity=interactivity, compiler=compiler, result=result)\n"
          ]
        }
      ],
      "source": [
        "olddata = pd.read_csv(FILELOC + \"patient_spike_ism_20210811.tsv\", '\\t')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wP4Nb3StH2s-",
        "outputId": "366beab7-a0a3-46b6-bef9-e42fe390b8ce"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "22167 23073 22166\n"
          ]
        }
      ],
      "source": [
        "mdf2 = pd.read_csv(FILELOC + 'align0912_spike.csv')\n",
        "patdf = pd.read_csv(FILELOC + 'GISAID_patient_metadata_20210812_0912.tsv', '\\t')\n",
        "olddata2 = pd.merge(left=mdf2[['SequenceID','MaskedSeq']], right = patdf,\n",
        "                    left_on = 'SequenceID', right_on='Accession ID')\n",
        "print(len(mdf2), len(patdf), len(olddata2))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nASTmm6Hb_v2",
        "outputId": "9b4eb68e-7bd0-4bfb-e547-265fbc284807"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "54081 104776 22166\n",
            "46328\n"
          ]
        }
      ],
      "source": [
        "print(len(alldf), len(olddata), len(olddata2))\n",
        "earlydata = alldf[alldf['Accession ID'].isin(olddata['Accession ID']) | alldf['Accession ID'].isin(olddata2['Accession ID']) ]\n",
        "# earlydata = alldf_rawseqs[alldf_rawseqs['Accession ID'].isin(olddata['Accession ID'])]\n",
        "print(len(earlydata))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "14hHEMa2QqSz"
      },
      "outputs": [],
      "source": [
        "# seldata = alldf.copy()\n",
        "seldata = earlydata.copy()\n",
        "\n",
        "countdata = seldata.groupby([\"ISM\",\"Label\"]).size().to_frame().rename(columns={0:\"Count\"})\n",
        "newdata = pd.DataFrame(seldata.ISM.value_counts()).rename(columns={\"ISM\":\"Count\"})\n",
        "newdata['Count_Dict'] = [countdata.loc[ind].to_dict()[\"Count\"] for ind in newdata.index]\n",
        "newdata = newdata.reset_index(drop=False).rename(columns={\"index\":\"ISM\"})\n",
        "label_list = np.zeros(len(newdata))\n",
        "for ind in newdata.index:\n",
        "    elem = newdata.Count_Dict[ind]\n",
        "    if len(elem) == 2:\n",
        "        label_list[ind] = newdata.Count_Dict[ind][1] / newdata.Count[ind]\n",
        "    elif 0 in elem:\n",
        "        label_list[ind] = 0\n",
        "    elif 1 in elem:\n",
        "        label_list[ind] = 1\n",
        "newdata['Label'] = label_list\n",
        "totalcount = np.sum(newdata.Count)\n",
        "newdata['sample_weight'] = newdata['Count'].apply(lambda x: x/totalcount*10)\n",
        "\n",
        "seldata = newdata.copy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gn4JMgRfBU3N"
      },
      "outputs": [],
      "source": [
        "basedata = seldata.copy()\n",
        "# alldata = seldata.copy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hTE1gJecShyu",
        "outputId": "3c96557f-0122-4cdc-ccc0-f49ae62235d9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "7557 716 6841 96\n"
          ]
        }
      ],
      "source": [
        "bd = basedata[basedata.Count >= 5]\n",
        "ad = alldata[~alldata.ISM.isin(bd.ISM)]\n",
        "print(len(alldata), len(bd), len(ad), len(ad[ad.Count >= 5]))\n",
        "\n",
        "# alldata.to_csv(FILELOC+\"covid_patient_all_raw_sequences_1001.csv\", index=False)\n",
        "# bd.to_csv(FILELOC + 'covid_patient_0912_raw_sequences.csv', index=False)\n",
        "# ad.to_csv(FILELOC + 'covid_patient_0912_valid1001_raw_sequences.csv', index=False)\n",
        "# alldata.to_csv(FILELOC+\"covid_patient_all_aligned_sequences_1001.csv\", index=False)\n",
        "# bd.to_csv(FILELOC + 'covid_patient_0912_aligned_sequences.csv', index=False)\n",
        "# ad.to_csv(FILELOC + 'covid_patient_0912_valid1001_aligned_sequences.csv', index=False)\n",
        "alldata.to_csv(FILELOC+\"covid_patient_all_filtered_sequences_1001.csv\", index=False)\n",
        "bd.to_csv(FILELOC + 'covid_patient_0912_filtered_sequences.csv', index=False)\n",
        "ad.to_csv(FILELOC + 'covid_patient_0912_valid1001_filtered_sequences.csv', index=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lhdtjxl1crse"
      },
      "source": [
        "## B.1.640 Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HRekxhpscu_3",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ef288df8-f4c4-4b6c-da94-ba1b01df7ca8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: scikit-bio in /usr/local/lib/python3.7/dist-packages (0.5.6)\n",
            "Requirement already satisfied: IPython>=3.2.0 in /usr/local/lib/python3.7/dist-packages (from scikit-bio) (5.5.0)\n",
            "Requirement already satisfied: natsort>=4.0.3 in /usr/local/lib/python3.7/dist-packages (from scikit-bio) (5.5.0)\n",
            "Requirement already satisfied: lockfile>=0.10.2 in /usr/local/lib/python3.7/dist-packages (from scikit-bio) (0.12.2)\n",
            "Requirement already satisfied: scipy>=1.3.0 in /usr/local/lib/python3.7/dist-packages (from scikit-bio) (1.4.1)\n",
            "Requirement already satisfied: matplotlib>=1.4.3 in /usr/local/lib/python3.7/dist-packages (from scikit-bio) (3.2.2)\n",
            "Requirement already satisfied: pandas>=1.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-bio) (1.1.5)\n",
            "Requirement already satisfied: scikit-learn>=0.19.1 in /usr/local/lib/python3.7/dist-packages (from scikit-bio) (1.0.2)\n",
            "Requirement already satisfied: hdmedians>=0.13 in /usr/local/lib/python3.7/dist-packages (from scikit-bio) (0.14.2)\n",
            "Requirement already satisfied: decorator>=3.4.2 in /usr/local/lib/python3.7/dist-packages (from scikit-bio) (4.4.2)\n",
            "Requirement already satisfied: numpy>=1.9.2 in /usr/local/lib/python3.7/dist-packages (from scikit-bio) (1.19.5)\n",
            "Requirement already satisfied: CacheControl>=0.11.5 in /usr/local/lib/python3.7/dist-packages (from scikit-bio) (0.12.10)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from CacheControl>=0.11.5->scikit-bio) (2.23.0)\n",
            "Requirement already satisfied: msgpack>=0.5.2 in /usr/local/lib/python3.7/dist-packages (from CacheControl>=0.11.5->scikit-bio) (1.0.3)\n",
            "Requirement already satisfied: Cython>=0.23 in /usr/local/lib/python3.7/dist-packages (from hdmedians>=0.13->scikit-bio) (0.29.26)\n",
            "Requirement already satisfied: pygments in /usr/local/lib/python3.7/dist-packages (from IPython>=3.2.0->scikit-bio) (2.6.1)\n",
            "Requirement already satisfied: pexpect in /usr/local/lib/python3.7/dist-packages (from IPython>=3.2.0->scikit-bio) (4.8.0)\n",
            "Requirement already satisfied: pickleshare in /usr/local/lib/python3.7/dist-packages (from IPython>=3.2.0->scikit-bio) (0.7.5)\n",
            "Requirement already satisfied: simplegeneric>0.8 in /usr/local/lib/python3.7/dist-packages (from IPython>=3.2.0->scikit-bio) (0.8.1)\n",
            "Requirement already satisfied: prompt-toolkit<2.0.0,>=1.0.4 in /usr/local/lib/python3.7/dist-packages (from IPython>=3.2.0->scikit-bio) (1.0.18)\n",
            "Requirement already satisfied: setuptools>=18.5 in /usr/local/lib/python3.7/dist-packages (from IPython>=3.2.0->scikit-bio) (57.4.0)\n",
            "Requirement already satisfied: traitlets>=4.2 in /usr/local/lib/python3.7/dist-packages (from IPython>=3.2.0->scikit-bio) (5.1.1)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib>=1.4.3->scikit-bio) (1.3.2)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib>=1.4.3->scikit-bio) (3.0.7)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib>=1.4.3->scikit-bio) (2.8.2)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib>=1.4.3->scikit-bio) (0.11.0)\n",
            "Requirement already satisfied: pytz>=2017.2 in /usr/local/lib/python3.7/dist-packages (from pandas>=1.0.0->scikit-bio) (2018.9)\n",
            "Requirement already satisfied: six>=1.9.0 in /usr/local/lib/python3.7/dist-packages (from prompt-toolkit<2.0.0,>=1.0.4->IPython>=3.2.0->scikit-bio) (1.15.0)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.7/dist-packages (from prompt-toolkit<2.0.0,>=1.0.4->IPython>=3.2.0->scikit-bio) (0.2.5)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn>=0.19.1->scikit-bio) (1.1.0)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-learn>=0.19.1->scikit-bio) (3.0.0)\n",
            "Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.7/dist-packages (from pexpect->IPython>=3.2.0->scikit-bio) (0.7.0)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->CacheControl>=0.11.5->scikit-bio) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->CacheControl>=0.11.5->scikit-bio) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->CacheControl>=0.11.5->scikit-bio) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->CacheControl>=0.11.5->scikit-bio) (2021.10.8)\n",
            "Collecting biopython\n",
            "  Downloading biopython-1.79-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.whl (2.3 MB)\n",
            "\u001b[K     |████████████████████████████████| 2.3 MB 5.2 MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from biopython) (1.19.5)\n",
            "Installing collected packages: biopython\n",
            "Successfully installed biopython-1.79\n"
          ]
        }
      ],
      "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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xQnIiPcLdnSl"
      },
      "source": [
        "In the code below, OMICRONFILE can be any specific dataset (this was taken from the code used to obtain the Omicron spike sequences)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OP6Y24fqd4ad"
      },
      "outputs": [],
      "source": [
        "DIR = FILELOC + 'coronavirus_spike/'\n",
        "SPIKEFILE = DIR + 'EPI_ISL_402124-S.fasta'\n",
        "# OMICRONFILE = DIR + 'gisaid_hcov-19_2022_01_07_21_B.1.640.fasta'\n",
        "OMICRONFILE = DIR+'gisaid_hcov-19_2022_01_28_00.fasta' # BA.2 sequences before Dec. 15, 2022\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",
        "# 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]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1UlvD-Docu9t",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d91e875e-b266-4d8a-e11a-29a7f29ada05"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "20\n"
          ]
        }
      ],
      "source": [
        "omicrondf = pd.DataFrame.from_dict({'Metadata':omicron_metadata})\n",
        "def f(x):\n",
        "    try:\n",
        "        return x.split('|')[1]\n",
        "    except:\n",
        "        return x\n",
        "omicrondf['SeqID'] = omicrondf.Metadata.apply(f)\n",
        "omicrondf['Raw'] = omicron_seqs\n",
        "omicrondf['seqlen'] = omicrondf['Raw'].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 + 'B1640_spike_20220107.csv',index=False)\n",
        "omicrondf.to_csv(FILELOC + 'omicron_spike_BA2_20220127.csv',index=False)\n",
        "print(len(omicrondf))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dp0LwnmTcopA"
      },
      "outputs": [],
      "source": [
        "data = pd.read_csv(FILELOC + 'covid_newpatient_crossval/covid_newpatient_alldata_20210930.csv')\n",
        "data.reset_index(drop=True, inplace=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "8d10b70aaad64928ae32907df2fedc0c",
            "7f01b2ec0a484b9f9f9f6fa0c1e4deaa",
            "4d4a74df9e0047d5ac6b0dad69938426",
            "f561674caab5462f8722619d6e491100",
            "476604c6b11a4ed89b5f67fc7a2ebddc",
            "4e16e44c57f54c7caec7377d5b89b101",
            "afac531fc8c74fe992ce558131dc80a7",
            "db53d8a9c8ea4b0699647e1917d35aa4",
            "a1e4317363614913a234437c7ff4078c",
            "003ae6c4d1234125a1e9dd5f61070e96",
            "bc3a1a78fd6c4cceaef22ee516c5e978"
          ]
        },
        "id": "0icDx4_DcopG",
        "outputId": "e3e31482-1f6d-46ba-e7f5-4ce489fca3a6"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "8d10b70aaad64928ae32907df2fedc0c",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/10 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "train_index_list = []; test_index_list = []\n",
        "train_labels_list = []; test_labels_list = []\n",
        "train_predict_list = []; test_predict_list = []\n",
        "\n",
        "for fold in tqdm(range(10)):\n",
        "    train_index = np.loadtxt(FILELOC + 'covid_newpatient_crossval/TrainIndex_'+str(fold)+'.csv', dtype='i')\n",
        "    test_index = np.loadtxt(FILELOC + 'covid_newpatient_crossval/TestIndex_'+str(fold)+'.csv', dtype='i')\n",
        "    train_index_list.append(train_index)\n",
        "    test_index_list.append(test_index)\n",
        "\n",
        "    train_labels_list.append(data.loc[train_index]['Label'].values)\n",
        "    test_labels_list.append(data.loc[test_index]['Label'].values)\n",
        "    pred = np.loadtxt(FILELOC + 'covid_newpatient_output/covid_patient_crossval_20211001_predict_'+str(fold)+'.csv')\n",
        "    train_predict_list.append(pred[train_index])\n",
        "    test_predict_list.append(pred[test_index])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mOuypiqmcopH",
        "outputId": "32507a74-1140-4d98-e46e-3db8e95ac461"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "0 0.014638042230325044 0.3139640916609181 0.33004959332636913\n",
            "1 0.014895349535692779 0.34497391324343424 0.2971476132486225\n",
            "2 0.012205553933830781 0.3441414695162451 0.3306383144403651\n",
            "3 0.009035614463175728 0.3467450902640182 0.3377550571062478\n",
            "4 0.01512524239340003 0.30797196999688636 0.3456063164946015\n",
            "5 0.015713130225866217 0.33105579123557455 0.3430189319652821\n",
            "6 0.011228193564955096 0.33255374081796013 0.363040634449386\n",
            "7 0.013227983584382522 0.3356880046993152 0.359428105632646\n",
            "8 0.013675567825118714 0.3097147075583797 0.3265607025365889\n",
            "9 0.014458072489585418 0.31494888333001947 0.3291354305407129\n"
          ]
        }
      ],
      "source": [
        "for fold in range(10):\n",
        "    print(fold,\n",
        "          mean_squared_error(train_labels_list[fold], train_predict_list[fold]),\n",
        "          np.sqrt(mean_squared_error(test_labels_list[fold], test_predict_list[fold])),\n",
        "          np.sqrt(mean_squared_error(test_labels_list[fold], np.mean(train_labels_list[fold])*np.ones(len(test_labels_list[fold]))))\n",
        "          )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mJ1i4F0l3X0P",
        "outputId": "a27efb3d-a69a-4270-e40f-973416d376ac"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Mean RMSE using Attention NN model = 0.31494888333001947\n",
            "Mean RMSE guessing average of train data = 0.3291354305407129\n"
          ]
        }
      ],
      "source": [
        "print('Mean RMSE using Attention NN model =',\n",
        "      np.mean(np.sqrt(mean_squared_error(test_labels_list[fold], test_predict_list[fold]))))\n",
        "print('Mean RMSE guessing average of train data =',\n",
        "      np.mean(np.sqrt(mean_squared_error(test_labels_list[fold], np.mean(train_labels_list[fold])*np.ones(len(test_labels_list[fold]))))))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "Lhdtjxl1crse"
      ],
      "name": "Covid_Data_Creation_github.ipynb",
      "toc_visible": true,
      "provenance": [],
      "mount_file_id": "1KNeUn28gsB_tz4yVsdQkfi9rd4koMljr",
      "authorship_tag": "ABX9TyO1ECUu71MB+tzu15v1bWcM"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "003ae6c4d1234125a1e9dd5f61070e96": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "476604c6b11a4ed89b5f67fc7a2ebddc": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4d4a74df9e0047d5ac6b0dad69938426": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_db53d8a9c8ea4b0699647e1917d35aa4",
            "max": 10,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_a1e4317363614913a234437c7ff4078c",
            "value": 10
          }
        },
        "4e16e44c57f54c7caec7377d5b89b101": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7f01b2ec0a484b9f9f9f6fa0c1e4deaa": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_4e16e44c57f54c7caec7377d5b89b101",
            "placeholder": "​",
            "style": "IPY_MODEL_afac531fc8c74fe992ce558131dc80a7",
            "value": "100%"
          }
        },
        "8d10b70aaad64928ae32907df2fedc0c": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_7f01b2ec0a484b9f9f9f6fa0c1e4deaa",
              "IPY_MODEL_4d4a74df9e0047d5ac6b0dad69938426",
              "IPY_MODEL_f561674caab5462f8722619d6e491100"
            ],
            "layout": "IPY_MODEL_476604c6b11a4ed89b5f67fc7a2ebddc"
          }
        },
        "a1e4317363614913a234437c7ff4078c": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "afac531fc8c74fe992ce558131dc80a7": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "bc3a1a78fd6c4cceaef22ee516c5e978": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "db53d8a9c8ea4b0699647e1917d35aa4": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f561674caab5462f8722619d6e491100": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_003ae6c4d1234125a1e9dd5f61070e96",
            "placeholder": "​",
            "style": "IPY_MODEL_bc3a1a78fd6c4cceaef22ee516c5e978",
            "value": " 10/10 [00:03&lt;00:00,  3.22it/s]"
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}