protein_language_solutions.ipynb

{
  "metadata": {
    "kernelspec": {
      "language": "python",
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.7.12",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "colab": {
      "name": "protein_embeddings_solutions.ipynb",
      "provenance": []
    },
    "accelerator": "GPU"
  },
  "nbformat_minor": 0,
  "nbformat": 4,
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Protein embeddings in practice"
      ],
      "metadata": {
        "id": "8QW5I6LYPCkb"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this notebook, we will practically explore protein language models and the embeddings that they produce. We will use two freely available software packages: the comprehensive *'bio_embeddings'* package and the more recent *'protein_sequence_models'* package."
      ],
      "metadata": {
        "id": "60zW_iWwgnru"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 0. Installation and libraries\n",
        "---\n",
        "\n",
        "First, let's install the necessary software and load the desired packages. Note: the pip's dependency resolver ERROR can be safely ignored."
      ],
      "metadata": {
        "id": "-aOQeo6fPiag"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Install software for Google Colab\n",
        "!pip3 install -U pip > /dev/null\n",
        "!pip3 install -U bio_embeddings[all] > /dev/null\n",
        "!pip install scikit_learn==1.0.2\n",
        "!pip install pyyaml==5.4.1"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Yq4KntowoogC",
        "outputId": "a2dc56d5-b46f-472b-a9d6-7538a89a0375"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
            "\u001b[0m\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "yellowbrick 1.4 requires scikit-learn>=1.0.0, but you have scikit-learn 0.24.2 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
            "\u001b[0m"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Install software for Kaggle\n",
        "!pip install bio_embeddings==0.2.2\n",
        "!pip install torchvision==0.10.1"
      ],
      "metadata": {
        "id": "mtiN3UkvO_m-",
        "outputId": "6456465c-0e01-40cc-d077-4b553b318d75",
        "execution": {
          "iopub.status.busy": "2022-05-22T20:00:35.829905Z",
          "iopub.execute_input": "2022-05-22T20:00:35.830225Z",
          "iopub.status.idle": "2022-05-22T20:03:08.083147Z",
          "shell.execute_reply.started": "2022-05-22T20:00:35.830130Z",
          "shell.execute_reply": "2022-05-22T20:03:08.082293Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: bio_embeddings==0.2.2 in /usr/local/lib/python3.7/dist-packages (0.2.2)\n",
            "Requirement already satisfied: plotly<6.0.0,>=5.1.0 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (5.5.0)\n",
            "Requirement already satisfied: tqdm<5.0.0,>=4.45.0 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (4.64.0)\n",
            "Requirement already satisfied: ruamel.yaml<0.18.0,>=0.17.10 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (0.17.21)\n",
            "Requirement already satisfied: appdirs<2.0.0,>=1.4.4 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (1.4.4)\n",
            "Requirement already satisfied: h5py<4.0.0,>=3.2.1 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (3.6.0)\n",
            "Requirement already satisfied: scipy<2.0.0,>=1.4.1 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (1.4.1)\n",
            "Requirement already satisfied: pandas<2.0.0,>=1.2.0 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (1.3.5)\n",
            "Requirement already satisfied: gensim<4.0.0,>=3.8.2 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (3.8.3)\n",
            "Requirement already satisfied: torch<=1.10.0,>=1.8.0 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (1.9.1)\n",
            "Requirement already satisfied: python-slugify<6.0.0,>=5.0.2 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (5.0.2)\n",
            "Requirement already satisfied: atomicwrites<2.0.0,>=1.4.0 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (1.4.0)\n",
            "Requirement already satisfied: umap-learn<0.6.0,>=0.5.1 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (0.5.3)\n",
            "Collecting scikit-learn<0.25.0,>=0.24.0\n",
            "  Using cached scikit_learn-0.24.2-cp37-cp37m-manylinux2010_x86_64.whl (22.3 MB)\n",
            "Requirement already satisfied: lock<2019.0.0,>=2018.3.25 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (2018.3.25.2110)\n",
            "Requirement already satisfied: importlib_metadata<5.0.0,>=4.6.1 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (4.11.3)\n",
            "Requirement already satisfied: matplotlib<4.0.0,>=3.2.1 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (3.2.2)\n",
            "Requirement already satisfied: biopython<2.0,>=1.79 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (1.79)\n",
            "Requirement already satisfied: humanize<4.0.0,>=3.2.0 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (3.14.0)\n",
            "Requirement already satisfied: numpy<2.0.0,>=1.18.3 in /usr/local/lib/python3.7/dist-packages (from bio_embeddings==0.2.2) (1.21.6)\n",
            "Requirement already satisfied: smart-open>=1.8.1 in /usr/local/lib/python3.7/dist-packages (from gensim<4.0.0,>=3.8.2->bio_embeddings==0.2.2) (6.0.0)\n",
            "Requirement already satisfied: six>=1.5.0 in /usr/local/lib/python3.7/dist-packages (from gensim<4.0.0,>=3.8.2->bio_embeddings==0.2.2) (1.15.0)\n",
            "Requirement already satisfied: cached-property in /usr/local/lib/python3.7/dist-packages (from h5py<4.0.0,>=3.2.1->bio_embeddings==0.2.2) (1.5.2)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib_metadata<5.0.0,>=4.6.1->bio_embeddings==0.2.2) (3.8.0)\n",
            "Requirement already satisfied: typing-extensions>=3.6.4 in /usr/local/lib/python3.7/dist-packages (from importlib_metadata<5.0.0,>=4.6.1->bio_embeddings==0.2.2) (4.2.0)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib<4.0.0,>=3.2.1->bio_embeddings==0.2.2) (0.11.0)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib<4.0.0,>=3.2.1->bio_embeddings==0.2.2) (2.8.2)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib<4.0.0,>=3.2.1->bio_embeddings==0.2.2) (1.4.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<4.0.0,>=3.2.1->bio_embeddings==0.2.2) (3.0.9)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas<2.0.0,>=1.2.0->bio_embeddings==0.2.2) (2022.1)\n",
            "Requirement already satisfied: tenacity>=6.2.0 in /usr/local/lib/python3.7/dist-packages (from plotly<6.0.0,>=5.1.0->bio_embeddings==0.2.2) (8.0.1)\n",
            "Requirement already satisfied: text-unidecode>=1.3 in /usr/local/lib/python3.7/dist-packages (from python-slugify<6.0.0,>=5.0.2->bio_embeddings==0.2.2) (1.3)\n",
            "Requirement already satisfied: ruamel.yaml.clib>=0.2.6 in /usr/local/lib/python3.7/dist-packages (from ruamel.yaml<0.18.0,>=0.17.10->bio_embeddings==0.2.2) (0.2.6)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn<0.25.0,>=0.24.0->bio_embeddings==0.2.2) (1.1.0)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-learn<0.25.0,>=0.24.0->bio_embeddings==0.2.2) (3.1.0)\n",
            "Requirement already satisfied: pynndescent>=0.5 in /usr/local/lib/python3.7/dist-packages (from umap-learn<0.6.0,>=0.5.1->bio_embeddings==0.2.2) (0.5.7)\n",
            "Requirement already satisfied: numba>=0.49 in /usr/local/lib/python3.7/dist-packages (from umap-learn<0.6.0,>=0.5.1->bio_embeddings==0.2.2) (0.51.2)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from numba>=0.49->umap-learn<0.6.0,>=0.5.1->bio_embeddings==0.2.2) (57.4.0)\n",
            "Requirement already satisfied: llvmlite<0.35,>=0.34.0.dev0 in /usr/local/lib/python3.7/dist-packages (from numba>=0.49->umap-learn<0.6.0,>=0.5.1->bio_embeddings==0.2.2) (0.34.0)\n",
            "Installing collected packages: scikit-learn\n",
            "  Attempting uninstall: scikit-learn\n",
            "    Found existing installation: scikit-learn 1.0.2\n",
            "    Uninstalling scikit-learn-1.0.2:\n",
            "      Successfully uninstalled scikit-learn-1.0.2\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "yellowbrick 1.4 requires scikit-learn>=1.0.0, but you have scikit-learn 0.24.2 which is incompatible.\u001b[0m\n",
            "Successfully installed scikit-learn-0.24.2\n",
            "Requirement already satisfied: torchvision==0.10.1 in /usr/local/lib/python3.7/dist-packages (0.10.1)\n",
            "Requirement already satisfied: pillow>=5.3.0 in /usr/local/lib/python3.7/dist-packages (from torchvision==0.10.1) (7.1.2)\n",
            "Requirement already satisfied: torch==1.9.1 in /usr/local/lib/python3.7/dist-packages (from torchvision==0.10.1) (1.9.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torchvision==0.10.1) (1.21.6)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch==1.9.1->torchvision==0.10.1) (4.2.0)\n",
            "\u001b[31mERROR: Invalid requirement: 'scikit_learn=1.0.2'\n",
            "Hint: = is not a valid operator. Did you mean == ?\u001b[0m\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# import libraries\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from tqdm import tqdm\n",
        "from Bio import SeqIO\n",
        "from bio_embeddings.embed import ProtTransBertBFDEmbedder"
      ],
      "metadata": {
        "id": "8jrtaDJdPEiH",
        "outputId": "b2575f2a-2f0b-47cb-bf6b-95d26c836a10",
        "execution": {
          "iopub.status.busy": "2022-05-22T20:03:12.243603Z",
          "iopub.execute_input": "2022-05-22T20:03:12.243974Z",
          "iopub.status.idle": "2022-05-22T20:03:38.300318Z",
          "shell.execute_reply.started": "2022-05-22T20:03:12.243932Z",
          "shell.execute_reply": "2022-05-22T20:03:38.299565Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/distributed/config.py:20: YAMLLoadWarning: calling yaml.load() without Loader=... is deprecated, as the default Loader is unsafe. Please read https://msg.pyyaml.org/load for full details.\n",
            "  defaults = yaml.load(f)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 1. Very simple introduction\n",
        "---\n",
        "\n",
        "#### The basics\n",
        "\n",
        "Let's use the ProtTransBertBFDEMbedder from the bio_embeddings package and compute embeddings for a simple test sequence. The basic workflow to compute an embedding is the following: (1) define an embedder (a pretrained model to compute embeddings from), (2) define your sequence as a string and (3) compute the embedding and reduce if desired. Note: The first time an embedder is loaded into memory takes a minute or two."
      ],
      "metadata": {
        "id": "a7V7GVzdjoHf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "embedder = ProtTransBertBFDEmbedder()\n",
        "sequence = \"THISISASEQUENCE\"\n",
        "embedding = embedder.embed(sequence)\n",
        "reduced_embedding = embedder.reduce_per_protein(embedding)"
      ],
      "metadata": {
        "id": "zmYDCoD_etFx",
        "outputId": "c33e17be-8a38-432e-a7fe-76a9d4a8dddb",
        "execution": {
          "iopub.status.busy": "2022-05-22T14:29:40.242354Z",
          "iopub.execute_input": "2022-05-22T14:29:40.242822Z",
          "iopub.status.idle": "2022-05-22T14:30:38.946136Z",
          "shell.execute_reply.started": "2022-05-22T14:29:40.242785Z",
          "shell.execute_reply": "2022-05-22T14:30:38.945436Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Some weights of the model checkpoint at /root/.cache/bio_embeddings/prottrans_bert_bfd/model_directory were not used when initializing BertModel: ['cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.decoder.weight', 'cls.predictions.decoder.bias', 'cls.seq_relationship.weight', 'cls.predictions.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.dense.bias']\n",
            "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "> **To-do and questions**\n",
        "> - Look at the embedding and reduced_embedding\n",
        "> - How do these two differ? And what do they signify?\n",
        "> - Recalling the theory slides, how does the reduction work? Can you think of other ways of doing this?\n",
        "> - Compare the embedding of the first 'E' with the embedding of the last 'E', are they the same?"
      ],
      "metadata": {
        "id": "3sE-yV66kDeb"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Computing many embeddings in a loop\n",
        "\n",
        "If you have many sequences available for which you want to compute embeddings, it's probably easier to put them in a FASTA file and loop over them.\n",
        "\n",
        "> **To-do**\n",
        ">\n",
        "> 1. Write a small function that parses a FASTA file and computes embeddings for each of the sequences.\n",
        "> \n",
        "> 2. Upload the two FASTA files in the GitHub repo onto the Kaggle or Colab platform. On Kaggle, see 'Add data' button in the upper right corner. This file will be located at '../input/a_folder_you_designate'. In Colab, files can be uploaded just as easy. Click on the *folder\b* icon in the left bar, then select the location where you want to store your files and click on the *upload* icon.\n",
        ">\n",
        "> 3. Run your function on one of the uploaded files and inspect the results.\n",
        ">\n",
        "> Tip 1: you can use 'SeqIO.parse' and a simple list comprehension to compute your embeddings in a loop.\n",
        ">\n",
        "> Tip 2: to track progress, you can use the 'tqdm' function in the list comprehension."
      ],
      "metadata": {
        "id": "HYYy8m0AZwMv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def compute_protein_embeddings(fasta_file):\n",
        "    embedder = ProtTransBertBFDEmbedder()\n",
        "    names = [record.id for record in SeqIO.parse(fasta_file, 'fasta')]\n",
        "    sequences = [str(record.seq) for record in SeqIO.parse(fasta_file, 'fasta')]\n",
        "    embeddings = [embedder.reduce_per_protein(embedder.embed(sequence)) for sequence in tqdm(sequences)]\n",
        "    #embeddings_df = pd.concat([pd.DataFrame({'ID':names}), pd.DataFrame(embeddings)], axis=1)\n",
        "    return names, embeddings"
      ],
      "metadata": {
        "id": "qAmGL0pUZD9F",
        "execution": {
          "iopub.status.busy": "2022-05-22T20:03:44.462240Z",
          "iopub.execute_input": "2022-05-22T20:03:44.462709Z",
          "iopub.status.idle": "2022-05-22T20:03:44.467779Z",
          "shell.execute_reply.started": "2022-05-22T20:03:44.462671Z",
          "shell.execute_reply": "2022-05-22T20:03:44.467131Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "embeddings = compute_protein_embeddings('your_fasta_file')"
      ],
      "metadata": {
        "id": "dIBij5Lnk-nl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Which is the best protein language model?\n",
        "\n",
        "\n",
        "What's clear from the [bio_embeddings GitHub repo](https://github.com/sacdallago/bio_embeddings) is that a variety of different protein language models exist and can be used to compute embeddings with. However, which of these embeddings is most informative? As mentioned in the 'bio_embeddings' ReadMe: \"The models prottrans_t5_xl_u50, esm1b, esm, prottrans_bert_bfd, prottrans_albert_bfd, seqvec and prottrans_xlnet_uniref100 were all trained with the goal of systematic predictions. From this pool, we believe the optimal model to be prottrans_t5_xl_u50, followed by esm1b.\"\n",
        "\n",
        "Below, try to load an instance of the t5_xl_u50 model and see what you get. \n",
        "\n",
        "**Warning**: this takes a while... \n",
        "\n",
        "**Spoiler**: it crashes the CPU..."
      ],
      "metadata": {
        "id": "EF9m3Bymk-nm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from bio_embeddings.embed import ProtTransT5XLU50Embedder\n",
        "embedder = ProtTransT5XLU50Embedder()"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-22T14:43:26.327811Z",
          "iopub.execute_input": "2022-05-22T14:43:26.328085Z"
        },
        "trusted": true,
        "id": "W3iEdv-1k-nn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Unfortunately, notebook environments like this are not suited to run the largest, most performant models (yet), due to their size and necessary allocation on the CPU and GPU. For more info, see [the init file](https://github.com/sacdallago/bio_embeddings/blob/develop/bio_embeddings/embed/__init__.py) of the embedders."
      ],
      "metadata": {
        "id": "CzPufaP0k-nn"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 2. Visualizing embeddings\n",
        "---\n",
        "\n",
        "#### Individual amino acid embeddings\n",
        "\n",
        "At the most basic level we would assume that, in some way at least, similar amino acids have similar embeddings. A simple way to check this is to make visualizations of the embeddings in two or three dimensions. We can simply use our favorite dimensionality reduction technique (PCA, t-SNE, UMAP) for this.\n",
        "\n",
        "> **To-do**\n",
        ">\n",
        "> - Compute an embedding for every naturally occuring amino acid and project them into lower dimensions.\n",
        "> - Visualize these lower dimension projections. You can use the 'render_scatter_plotly' function in bio_embeddings.visualize for this. Tip: help(render_scatter_plotly) to see what is needed as input.\n",
        "> - Color the amino acids by polarity and charge to identify any clusters. Tip: see the 'label' input for the render_scatter_plotly function.\n"
      ],
      "metadata": {
        "id": "UZhWNljPk-no"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from bio_embeddings.visualize import render_scatter_plotly\n",
        "from sklearn.manifold import TSNE\n",
        "import umap"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-22T20:08:14.974489Z",
          "iopub.execute_input": "2022-05-22T20:08:14.974989Z",
          "iopub.status.idle": "2022-05-22T20:08:14.979110Z",
          "shell.execute_reply.started": "2022-05-22T20:08:14.974952Z",
          "shell.execute_reply": "2022-05-22T20:08:14.978327Z"
        },
        "trusted": true,
        "id": "ArJzWsAzk-no"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "amino_acids = ['A','C','D','E','F','G','H','I','K','L','M','N','P','Q','R','S','T','V','W','Y']\n",
        "groups = {'polar': ['S', 'T', 'Y', 'N', 'Q'], 'non-polar':['G', 'A', 'V', 'C', 'P', 'L', 'I', 'M', 'W', 'F'], \n",
        "         'pos+': ['L', 'R', 'H'], 'neg-':['D', 'E']}\n",
        "help(render_scatter_plotly)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-22T15:02:58.543915Z",
          "iopub.execute_input": "2022-05-22T15:02:58.544385Z",
          "iopub.status.idle": "2022-05-22T15:02:58.551770Z",
          "shell.execute_reply.started": "2022-05-22T15:02:58.544348Z",
          "shell.execute_reply": "2022-05-22T15:02:58.550871Z"
        },
        "trusted": true,
        "id": "c-wWlzDGk-np"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "embedder = ProtTransBertBFDEmbedder()\n",
        "embeddings = [embedder.embed(AA)[0] for AA in amino_acids]\n",
        "labels = [key for AA in amino_acids for (key, value) in groups.items() if AA in value]\n",
        "embed_array = np.asarray(embeddings)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-22T15:14:10.315361Z",
          "iopub.execute_input": "2022-05-22T15:14:10.316093Z",
          "iopub.status.idle": "2022-05-22T15:14:29.353690Z",
          "shell.execute_reply.started": "2022-05-22T15:14:10.316054Z",
          "shell.execute_reply": "2022-05-22T15:14:29.352930Z"
        },
        "trusted": true,
        "id": "waJNwdRik-np"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "reducer = umap.UMAP()\n",
        "embed_umap = reducer.fit_transform(embed_array)\n",
        "embed_tsne = TSNE(n_components=2, init='random').fit_transform(embed_array)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-19T12:26:20.367573Z",
          "iopub.execute_input": "2022-05-19T12:26:20.368325Z",
          "iopub.status.idle": "2022-05-19T12:26:22.576029Z",
          "shell.execute_reply.started": "2022-05-19T12:26:20.368284Z",
          "shell.execute_reply": "2022-05-19T12:26:22.575238Z"
        },
        "trusted": true,
        "id": "8XcUuE5Ik-nq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "embeddings_umap = pd.concat([pd.DataFrame({'label':labels}), pd.DataFrame({'component_0': embed_umap[:,0]}), \n",
        "                          pd.DataFrame({'component_1': embed_umap[:,1]})], axis=1)\n",
        "embeddings_tsne = pd.concat([pd.DataFrame({'label':labels}), pd.DataFrame({'component_0': embed_tsne[:,0]}), \n",
        "                          pd.DataFrame({'component_1': embed_tsne[:,1]})], axis=1)\n",
        "fig = render_scatter_plotly(embeddings_umap)\n",
        "fig"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-19T12:26:33.30623Z",
          "iopub.execute_input": "2022-05-19T12:26:33.306506Z",
          "iopub.status.idle": "2022-05-19T12:26:33.391095Z",
          "shell.execute_reply.started": "2022-05-19T12:26:33.306477Z",
          "shell.execute_reply": "2022-05-19T12:26:33.390369Z"
        },
        "trusted": true,
        "id": "if5JAePck-nq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "> **To-do and Questions**\n",
        ">\n",
        "> - Do we see what we expect to see? What is in line with expectations, what is not?\n",
        "> - Explore other embedders to see if they produce different results: BeplerEmbedder, SeqVecEmbedder, ..."
      ],
      "metadata": {
        "id": "K8I3hx3hk-nq"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Full protein sequence embeddings related to *Saccharomyces cerevisiae* and *Homo sapiens*\n",
        "\n",
        "Now, we'll do the same for full proteins. More specifically, we will visualize the protein sequences in the FASTA files, collected from UniProt and related to *Saccharomyces cerevisiae* and *Homo sapiens*.\n",
        "\n",
        "> **To-do**\n",
        "> \n",
        "> - Let's get back to the two FASTA files we've seen before, upload both of them from the GitHub repo onto Kaggle or Colab if you haven't already.\n",
        "> - Repeat the process above: (1) use your previously implemented function to compute embeddings for all of the protein sequences, (2) project those into lower dimensions and (3) visualize them.\n"
      ],
      "metadata": {
        "id": "09ksK3C4k-nr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "names_sapiens, embeddings_sapiens = compute_protein_embeddings('../input/proteindata/uniprot_sapiens.fasta')\n",
        "names_yeast, embeddings_yeast = compute_protein_embeddings('../input/proteindata/uniprot_yeast.fasta')"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-22T20:03:59.443049Z",
          "iopub.execute_input": "2022-05-22T20:03:59.443322Z",
          "iopub.status.idle": "2022-05-22T20:06:55.390279Z",
          "shell.execute_reply.started": "2022-05-22T20:03:59.443293Z",
          "shell.execute_reply": "2022-05-22T20:06:55.389430Z"
        },
        "trusted": true,
        "id": "0HrTSvuGk-nr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "labels = ['sapiens']*len(embeddings_sapiens) + ['yeast']*len(embeddings_yeast)\n",
        "embeddings = embeddings_sapiens + embeddings_yeast\n",
        "embed_array = np.asarray(embeddings)\n",
        "reducer = umap.UMAP()\n",
        "embed_umap = reducer.fit_transform(embed_array)\n",
        "embeddings_umap = pd.concat([pd.DataFrame({'label':labels}), pd.DataFrame({'component_0': embed_umap[:,0]}), \n",
        "                          pd.DataFrame({'component_1': embed_umap[:,1]})], axis=1)\n",
        "fig = render_scatter_plotly(embeddings_umap)\n",
        "fig"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-22T20:08:21.044298Z",
          "iopub.execute_input": "2022-05-22T20:08:21.045054Z",
          "iopub.status.idle": "2022-05-22T20:08:31.528825Z",
          "shell.execute_reply.started": "2022-05-22T20:08:21.045013Z",
          "shell.execute_reply": "2022-05-22T20:08:31.528191Z"
        },
        "trusted": true,
        "id": "dBLwUoE-k-nr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 3. Stacking a supervised model\n",
        "---\n",
        "\n",
        "#### A simple binary classifier\n",
        "\n",
        "Recalling the theory slides, we can stack a supervised machine learning model on these computed embeddings to perform a specific task. This means we use our computed embeddings as input for a supervised model where labels are available.\n",
        "\n",
        "> **To-do**\n",
        "> \n",
        "> Build a binary classifier on top of the computed embeddings that discriminate between proteins from *Homo sapiens* and *Saccharomyces cerevisiae*.\n",
        ">\n",
        "> Tip: you can probably build a simple classifier in Scikit-learn with [just a few lines of code](https://scikit-learn.org/stable/tutorial/basic/tutorial.html). If you're stuck, you can have a look at a [related example](https://github.com/sacdallago/bio_embeddings/blob/develop/notebooks/deeploc_machine_learning.ipynb) from the *bio_embeddings* package.\n"
      ],
      "metadata": {
        "id": "IiyrI4-4k-ns"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from sklearn.ensemble import RandomForestClassifier\n",
        "from sklearn.model_selection import KFold, train_test_split\n",
        "from sklearn.metrics import accuracy_score, pairwise_distances"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-22T20:46:47.444248Z",
          "iopub.execute_input": "2022-05-22T20:46:47.444934Z",
          "iopub.status.idle": "2022-05-22T20:46:47.449060Z",
          "shell.execute_reply.started": "2022-05-22T20:46:47.444896Z",
          "shell.execute_reply": "2022-05-22T20:46:47.448212Z"
        },
        "trusted": true,
        "id": "igopFpAbk-ns"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "X_train, X_test, y_train, y_test = train_test_split(embed_array, labels, test_size=0.3, random_state=42)\n",
        "clf = RandomForestClassifier(max_depth=2, random_state=0)\n",
        "clf.fit(X_train, y_train)\n",
        "preds = clf.predict(X_test)\n",
        "accuracy = accuracy_score(y_test, preds)\n",
        "print('A simple Random Forest yields an accuracy of: ', round(accuracy, 3))"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-22T16:38:46.893532Z",
          "iopub.execute_input": "2022-05-22T16:38:46.894033Z",
          "iopub.status.idle": "2022-05-22T16:38:47.351878Z",
          "shell.execute_reply.started": "2022-05-22T16:38:46.893994Z",
          "shell.execute_reply": "2022-05-22T16:38:47.351124Z"
        },
        "trusted": true,
        "id": "BjDo7B1Xk-nt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "> **Question**\n",
        ">\n",
        "> How would this approach change if we were to use individual amino acid embeddings to make some predictions (e.g. predict secondary structures at the AA level)?\n",
        ">"
      ],
      "metadata": {
        "id": "YUbLWfr_k-nt"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### K-nearest neighbors\n",
        "\n",
        "Let's try a different approach. You can also implement a simple K-nearest neigbor classifier by computing distances in the embeddings space. \n",
        "\n",
        "> **To-do**\n",
        "> \n",
        "> 1. Split the computed embeddings into a training and test part.\n",
        "> 2. Compute the pairwise distances between the training embeddings vectors and each of the test vectors.\n",
        "> 3. For each of the test vectors, collect the labels of the k closest training vectors.\n",
        "> 4. Apply a majority vote to predict the label.\n",
        "> 5. Compare the accuracy with that of your previous classifier."
      ],
      "metadata": {
        "id": "yOLDX-L7k-nt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# 1. split\n",
        "X_train, X_test, y_train, y_test = train_test_split(embed_array, labels, test_size=0.3, random_state=42)\n",
        "\n",
        "# 2. compute distances\n",
        "dist_matrix = pairwise_distances(X_train, X_test)\n",
        "\n",
        "# 3. Collect nearest neighbor labels\n",
        "k = 10\n",
        "predictions = []\n",
        "for j in range(dist_matrix.shape[1]):\n",
        "    distances = list(dist_matrix[:,j])\n",
        "    nearest_labels = [(x, y) for y, x in sorted(zip(distances, y_train))][:k]\n",
        "    predictions.append(nearest_labels)\n",
        "    \n",
        "# 4. majority vote\n",
        "preds = [max(set(item), key=item.count)[0] for item in predictions]\n",
        "\n",
        "# 5. accuracy\n",
        "accuracy = accuracy_score(y_test, preds)\n",
        "print('A simple Nearest Neighbor yields an accuracy of: ', round(accuracy, 3))"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-22T20:47:00.232395Z",
          "iopub.execute_input": "2022-05-22T20:47:00.232651Z",
          "iopub.status.idle": "2022-05-22T20:47:00.530877Z",
          "shell.execute_reply.started": "2022-05-22T20:47:00.232625Z",
          "shell.execute_reply": "2022-05-22T20:47:00.530179Z"
        },
        "trusted": true,
        "id": "gRXu3r5Kk-nu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 4. Beyond the *bio_embeddings* package\n",
        "---\n",
        "\n",
        "For those who hadn't had enough already, a few more interesting publicly available packages exist. Two that stand out are Microsoft's [*protein-sequence-models*](https://github.com/microsoft/protein-sequence-models) package and Brandes' [*protein_bert*](https://github.com/nadavbra/protein_bert) package.\n",
        "\n",
        "\n",
        "> **To-do**\n",
        ">\n",
        "> Explore the pretrained CARP models in the *protein-sequence-models* package and how you can compute embeddings with those. Redo some of the visualizations and/or model stacking and compare.\n",
        ">"
      ],
      "metadata": {
        "id": "pOCFsDXAk-nu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install sequence-models\n",
        "!pip install wget"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-22T19:32:57.360195Z",
          "iopub.execute_input": "2022-05-22T19:32:57.360470Z",
          "iopub.status.idle": "2022-05-22T19:33:16.783787Z",
          "shell.execute_reply.started": "2022-05-22T19:32:57.360440Z",
          "shell.execute_reply": "2022-05-22T19:33:16.782936Z"
        },
        "trusted": true,
        "id": "4LCZ6sNuk-nv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from sequence_models.pretrained import load_model_and_alphabet\n",
        "model, collater = load_model_and_alphabet('carp_76M')"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-22T19:35:18.725877Z",
          "iopub.execute_input": "2022-05-22T19:35:18.726181Z",
          "iopub.status.idle": "2022-05-22T19:35:38.475017Z",
          "shell.execute_reply.started": "2022-05-22T19:35:18.726134Z",
          "shell.execute_reply": "2022-05-22T19:35:38.474072Z"
        },
        "trusted": true,
        "id": "_IVw1QSnk-nv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "seqs = [['MDREQ'], ['MGTRRLLP']]\n",
        "x = collater(seqs)  # (n, max_len)\n",
        "rep = model(x[0])  # (n, max_len, d_model)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2022-05-22T19:37:33.986689Z",
          "iopub.execute_input": "2022-05-22T19:37:33.987439Z",
          "iopub.status.idle": "2022-05-22T19:37:34.303988Z",
          "shell.execute_reply.started": "2022-05-22T19:37:33.987399Z",
          "shell.execute_reply": "2022-05-22T19:37:34.303111Z"
        },
        "trusted": true,
        "id": "LuBcLKZfk-nv"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}