notebooks/prediction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "611a67cb-71c1-4a16-bfef-757ce22f5bd9",
   "metadata": {},
   "source": [
    "<img src=\"https://github.com/Multiomics-Analytics-Group/course_protein_language_modeling/blob/main/img/nb_logo.png?raw=1\" width=\"600\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ee7457c-15d3-4c69-812a-2bdaa68392c4",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Multiomics-Analytics-Group/course_protein_language_modeling/blob/main/notebooks/prediction.ipynb)\n",
    "\n",
    "\n",
    "This is a version of the notebook from [Meta Research](https://research.facebook.com/) --- [here](https://github.com/facebookresearch/esm/blob/main/examples/sup_variant_prediction.ipynb) using the output from the [Embeddings notebook](https://colab.research.google.com/github/Multiomics-Analytics-Group/course_protein_language_modeling/blob/main/notebooks/embeddings.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3a0883e-7dc9-498b-8e3d-74a7252eb0f9",
   "metadata": {
    "id": "QMoeBQnUCK_E"
   },
   "outputs": [],
   "source": [
    "# Install requirements\n",
    "!pip install h5py > /dev/null\n",
    "!wget https://github.com/Multiomics-Analytics-Group/course_protein_language_modeling/raw/main/data/per_protein_embeddings.h5.zip -O per_protein_embeddings.h5.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57f2e944-def3-40e1-93e2-03cd105d3239",
   "metadata": {},
   "source": [
    "# Predicting Varient Effect from Sequence Embeddings\n",
    "\n",
    "In this notebook we will use the embeddings we generated in embeddings.ipynb to train and optimize various machine learning models in sklearn.\n",
    "\n",
    "Each observation in our dataset $--$ which we created in embeddings.ipynb (or can be found in the git repo at _data/per_protein_embeddings.h5.zip_)$--$contains: \n",
    "- value: an embedded representation of the mutated ß-lactamase sequence\n",
    "- key: `{index}|beta-lactamase_{mutation}|{scaled_varient_effect}` where the target value is the scaled_varient_effect, which describes the scaled effect of the mutation\n",
    "\n",
    "**Goal:**\n",
    "Train a regression model in to predict the \"effect\" score of a $\\beta$-$lactamase$ variant given the embedding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8771ca2f-4791-4c52-8c26-79455a1eaf7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "import h5py\n",
    "import zipfile\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy\n",
    "\n",
    "# for fine-tuning\n",
    "from sklearn.model_selection import GridSearchCV, train_test_split\n",
    "from sklearn.pipeline import Pipeline\n",
    "\n",
    "# dimensionality reduction\n",
    "from sklearn.decomposition import PCA\n",
    "\n",
    "# models\n",
    "from sklearn.neighbors import KNeighborsRegressor\n",
    "from sklearn.svm import SVR\n",
    "\n",
    "# path to files\n",
    "zip_path = 'per_protein_embeddings.h5.zip' # local path\n",
    "#zip_path = 'per_protein_embeddings.h5.zip' # on collab\n",
    "filename = 'per_protein_embeddings.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cb13342-7467-4dae-a3b2-adce66844685",
   "metadata": {},
   "outputs": [],
   "source": [
    "# unzips the file and puts it in the current working directory\n",
    "with zipfile.ZipFile(zip_path, 'r') as zip_ref:\n",
    "    # extract the filename of interest\n",
    "    zip_ref.extract(filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fda0ea3f-0343-434f-a62a-b80d9ef2e1fc",
   "metadata": {
    "id": "vuqI2vnujG88"
   },
   "outputs": [],
   "source": [
    "# functions to help us with reading in our dataset\n",
    "\n",
    "def read_hdf5(path:str) -> dict:\n",
    "    '''\n",
    "    read in the h5 file to a dictionary\n",
    "    '''\n",
    "    # empty dict and list\n",
    "    weights = {}\n",
    "    \n",
    "    # open file\n",
    "    with h5py.File(path, 'r') as f:\n",
    "        # append to dict\n",
    "        for key in f.keys():\n",
    "            weights[key] = list(f[key])\n",
    "            \n",
    "    return weights\n",
    "\n",
    "def emb_to_df(emb:dict) -> pd.DataFrame:\n",
    "    '''\n",
    "    Takes the dictionary from read_hdf5() to create a dataframe. \n",
    "    This function is super specific to per_protein_embeddings.h5\n",
    "    '''\n",
    "    # to dataframe\n",
    "    df_seq = pd.DataFrame.from_dict(emb, orient='index').reset_index()\n",
    "    \n",
    "    # additional formatting for our purposes\n",
    "    # making each part of key its own column\n",
    "    header = df_seq['index'].str.split('|', expand=True).rename(\n",
    "        columns={\n",
    "            0:'index_value',\n",
    "            1:'mutation',\n",
    "            2:'scaled_varient_effect'\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    # combining with sequence embeddings\n",
    "    df = pd.concat([header, df_seq.drop('index', axis=1)], axis=1)\n",
    "\n",
    "    # target column to float dtype\n",
    "    df['scaled_varient_effect'] = df['scaled_varient_effect'].astype(float)\n",
    "            \n",
    "    return df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b5f592b-281f-494b-b965-1c30f278109b",
   "metadata": {},
   "source": [
    "### The Dataset\n",
    "\n",
    "Here, we read in the embeddings as a dataframe and take a look. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9f5b506-8c45-4d35-9573-68e583c4b737",
   "metadata": {
    "id": "NUkgnuY0qdz0"
   },
   "outputs": [],
   "source": [
    "# load in embeddings\n",
    "embeddings = read_hdf5(filename)\n",
    "\n",
    "# to dataframe\n",
    "df = emb_to_df(embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06df3f26-5f81-4c2e-972d-707b12d9e9fb",
   "metadata": {
    "id": "_v6l02vzq11W"
   },
   "outputs": [],
   "source": [
    "# What does our data look like? \n",
    "print(df.shape)\n",
    "print(f'Missing data? {df.isna().any().any()}')\n",
    "display(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e75eca58-8de7-4f79-b711-6b219f80ee6a",
   "metadata": {},
   "source": [
    "Further getting the dataset ready by separating the features (embedding only) and target (scaled_varient_effect). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfd179ed-7c18-433b-a880-67ff4fff0130",
   "metadata": {},
   "outputs": [],
   "source": [
    "# target\n",
    "y = df['scaled_varient_effect']\n",
    "\n",
    "# isolating features\n",
    "X = df.drop(['index_value', 'mutation', 'scaled_varient_effect'], axis=1)\n",
    "\n",
    "# check\n",
    "print(X.shape, y.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd2d5554-38ee-450c-943c-e1776b8cd569",
   "metadata": {
    "id": "86riejVNsD5_"
   },
   "source": [
    "### Train / Test Split\n",
    "\n",
    "Here we choose to follow the Envision paper, using 80% of the data for training, but we actually found that pre-trained ESM embeddings require fewer downstream training examples to reach the same level of performance.\n",
    "\n",
    "The test set will be used at the end of our notebook to assess how the model performs on never before seen data. \n",
    "\n",
    "The training set will be used to train and fine tune hyperparameters via a validation set (i.e., an intermediary 'test' set), which we will incorporate later in the notebook via cross-validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "074e3a7d-d1cd-41d3-9b2e-efc3e557ac1e",
   "metadata": {
    "id": "B0IBi8QDsEib"
   },
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    X, \n",
    "    y, \n",
    "    test_size=0.2, \n",
    "    random_state=42)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31df3580-1bd1-4516-b633-dcc9f3e8c678",
   "metadata": {
    "id": "MjCKwf3UtrV_"
   },
   "source": [
    "### PCA\n",
    "\n",
    "Principal Component Analysis (PCA) is a popular technique for dimensionality reduction. Given `n_features` (1024 in our case), PCA computes a new set of `components` that \"best explain the data\" by capturing the variance in the data. \n",
    "\n",
    "\n",
    "Using a subset of `components` reduces the number of dimensons (e.g., columns), which enables downstream models to be trained faster with minimal loss in performance.  \n",
    "\n",
    "For the below example, we arbitrarily set `components` to 50, but feel free to change it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4bbc21d-ce18-41f7-9956-43502ee01687",
   "metadata": {
    "id": "q_DfS0BFsG2C"
   },
   "outputs": [],
   "source": [
    "num_pca_components = 50\n",
    "\n",
    "# instantiate\n",
    "pca = PCA(num_pca_components)\n",
    "\n",
    "# fit to the data and keep only num_pca_componenrs\n",
    "Xs_train_pca = pca.fit_transform(X_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6dc23a0-4617-46cd-b84f-3574ff74e196",
   "metadata": {},
   "outputs": [],
   "source": [
    "# how much variance in the data is captured in the selected components?\n",
    "\n",
    "# calculating variance captured\n",
    "exp_var = sum(pca.explained_variance_ratio_)\n",
    "\n",
    "print(f'{exp_var} of variance is captured in {num_pca_components} PCs.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0e1fc8f-5e15-42ab-be52-bcbd4927b2ef",
   "metadata": {
    "id": "uUMAc-jbtxQq"
   },
   "source": [
    "## Visualize Embeddings\n",
    "\n",
    "Here, we plot the first two principal components on the x- and y- axes. Each point is then colored by its scaled effect (what we want to predict).\n",
    "\n",
    "Visually, we can see a separation based on color/effect, suggesting that our representations are useful for this task, without any task-specific training!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baec2b47-6487-46f7-b1ee-73f15b1c7aa0",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 559
    },
    "id": "YiZjjqMOtuRY",
    "outputId": "6df68e14-6923-484e-fc20-9a2db2a1a374"
   },
   "outputs": [],
   "source": [
    "fig_dims = (7, 6)\n",
    "fig, ax = plt.subplots(figsize=fig_dims)\n",
    "sc = ax.scatter(Xs_train_pca[:,0], Xs_train_pca[:,1], c=y_train, marker='.')\n",
    "ax.set_xlabel('PCA first principal component')\n",
    "ax.set_ylabel('PCA second principal component')\n",
    "plt.colorbar(sc, label='Variant Effect')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0facaa12-7c13-4413-acab-a3e507cc15db",
   "metadata": {},
   "source": [
    "### Setting up our pipeline process\n",
    "\n",
    "Sklearn has a [`Pipeline` class](https://scikit-learn.org/stable/modules/generated/sklearn.pipeline.Pipeline.html) that allows us to chain together pre-processing and model training steps. The purpose of the pipeline is to assemble several steps that can be cross-validated together while setting different parameters. We will later create a list of `params` consisting of the different hyperparameters we wish to fine-tune using cross-validation.\n",
    "\n",
    "The sequence in the `Pipeline()` will be:\n",
    "\n",
    "- A Dimensionality Reduction technique to reduce the number of dimensions ( `PCA()` ), and\n",
    "- Training a regressor model on the training dataset\n",
    "\n",
    "where the data will be passed from one transformer to the next in that order. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e2b9c40-39ca-4cf1-90a6-77461608ba73",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = Pipeline(\n",
    "    [\n",
    "        ('reduce_dim', PCA(n_components=50)),\n",
    "        ('model', KNeighborsRegressor()) \n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82a8b271-1d57-4f95-bde9-986fa2a1b119",
   "metadata": {},
   "source": [
    "With no additional fine-tuning of the we can see how the model performs by fitting our pipeline to the train and test data. Here, a [score](https://scikit-learn.org/stable/modules/generated/sklearn.neighbors.KNeighborsRegressor.html#sklearn.neighbors.KNeighborsRegressor.score) closer to 1.0 is better. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5b85739-6ad3-4cf2-a068-47fc2dd6ab50",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe.fit(X_train, y_train)\n",
    "print(pipe.score(X_train, y_train))\n",
    "print(pipe.score(X_test, y_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae0056ee-0962-4dcc-bde7-d0b31928f276",
   "metadata": {},
   "source": [
    "This is an okay score and all.. but can we train a model to achieve a better score? Specifically:\n",
    "- Did we include a good amount of components in our data?\n",
    "- What about the KNeighborsRegressor() hyperparameters of n_neighbors, leaf_size, etc.? Can these be fine-tuned to get a model that has a better score on the test set?\n",
    "- And, is KNighborsRegressor() the best fit? What about another algorithm like a support vector regression (SVR) and its parameters?\n",
    "\n",
    "### Initialize parameter-grid for fine-tuning\n",
    "\n",
    "We will let [`GridSearchCV`](https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.GridSearchCV.html) investigate which algorithm and algorithm-specific hyperparameters result in the best model (i.e., highest score) by completing a cross-validated grid-search over a parameter grid we provide.\n",
    "\n",
    "We will need to provide the options for `GridSearchCV` to explore. When selecting the optimal set of parameters to include we should consider the algorithm, our problem and what we know about the data. \n",
    "\n",
    "From [the documentation](https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.GridSearchCV.html) we know that GridSearchCV() expects the `param_grid` to be a list of dictionaries of lists. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0923ed07-83ca-494c-a83e-90bbecb1ed85",
   "metadata": {},
   "outputs": [],
   "source": [
    "# creating a dictionary for each model \n",
    "\n",
    "kn_grid = {\n",
    "    'model': [KNeighborsRegressor()],\n",
    "    'model__n_neighbors': [5, 10],\n",
    "    'model__weights': ['uniform', 'distance'],\n",
    "    'model__leaf_size' : [15, 30],\n",
    "    'model__p' : [1, 2],\n",
    "    'reduce_dim__n_components': [50, 100]\n",
    "}\n",
    "\n",
    "svr_grid = {\n",
    "    'model': [SVR(gamma='scale', degree=3)],\n",
    "    'model__C': [0.01, 0.1, 1.0, 10.0],\n",
    "    'model__kernel': ['linear', 'poly', 'rbf', 'sigmoid'],\n",
    "    'reduce_dim__n_components': [50, 100]\n",
    "}\n",
    "\n",
    "# putting in a list\n",
    "params = [kn_grid, svr_grid]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83dd2526-ee94-4a07-b578-002913f61ca6",
   "metadata": {
    "id": "nx0gTUW9uRdT"
   },
   "source": [
    "### Run Grid Search\n",
    "\n",
    "(This will take a few minutes on a single core)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c55d5357-10a5-4294-b6d0-17e2ef0ac8a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# instantiating the grid search\n",
    "best_model = GridSearchCV(\n",
    "    estimator=pipe, \n",
    "    param_grid=params, \n",
    "    scoring='r2',\n",
    "    cv=5, # 5 folds (validation)\n",
    "    verbose=3,\n",
    "    n_jobs=-1 #if wanting verbosity during run than n_jobs must be 1 (default)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05b0f788-4c14-45fd-96a2-cf03b2c4a92c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# running\n",
    "best_model.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec3d67bd-d316-48d3-a6b5-9839b03c7c26",
   "metadata": {},
   "source": [
    "So what was the best model and hyperparameters as determined by the grid search?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63b520dc-ccab-40fb-9d08-c78f9e0bbbab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# look at parameters\n",
    "best_model.best_estimator_.get_params()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9bb457e-e587-47ad-a885-f990343fbdbf",
   "metadata": {},
   "source": [
    "SVR with a rbf kernel and a $C$ value of 1.0 trained on a dataset of 100 principal components (that's only 10% of the number of features we started with!) was selected as the best estimator. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25f8e2dc-ef79-4dfc-b67c-d53a858c7d97",
   "metadata": {},
   "source": [
    "Using the GridSearchCV `.predict()` method will [call predict on the best found parameters.](https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.GridSearchCV.html#sklearn.model_selection.GridSearchCV.predict)\n",
    "\n",
    "Let's take a look at how the model scores on the train and (most importantly) the test set, sequences it has never seen before.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f84b974b-b225-417a-bbd5-5eb6c5f011a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'training set score {best_model.score(X_train, y_train)}')\n",
    "print(f'test set score {best_model.score(X_test, y_test)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1fea9e0-8974-473d-94fe-c4f321237587",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "preds = best_model.predict(X_test)\n",
    "print(f'{scipy.stats.spearmanr(y_test, preds)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac917d67-7b7c-4309-9b6c-c3d63e70038b",
   "metadata": {},
   "source": [
    "We achieved a spearman rho of 0.815 on the test set!\n",
    "\n",
    "In conclusion, our downstream model was able to use fixed pre-trained ESM embedding representations and obtain a decent result."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}