htFuncLib.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "P54aQ2yupRNk"
   },
   "source": [
    "# htFuncLib\n",
    "\n",
    "In order to use this notebook, first use the [FuncLib](https://funclib.weizmann.ac.il/step/fl_terms/) server to calculate the scores of the different multipoint mutants.   \n",
    "\n",
    "---\n",
    "\n",
    "\n",
    "**IMPORTANT** even if the file was changed after FuncLib, only use a file with the original column names and not any other additional ones.\n",
    "\n",
    "\n",
    "---\n",
    "\n",
    "\n",
    "\n",
    "Wherever the user's input is required, or a parameter may be changed, it is labeled **EDIT HERE** with an explanation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VrtQ2tlcpRNm"
   },
   "source": [
    "## 01 Imports and initial setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-12T20:07:56.229245Z",
     "iopub.status.busy": "2022-11-12T20:07:56.228669Z",
     "iopub.status.idle": "2022-11-12T20:07:56.353482Z",
     "shell.execute_reply": "2022-11-12T20:07:56.352601Z",
     "shell.execute_reply.started": "2022-11-12T20:07:56.229179Z"
    },
    "id": "v2nu4dekpRNn",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "from IPython.display import display, Markdown\n",
    "\n",
    "import matplotlib.patches as mpatches\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import rc\n",
    "import seaborn as sns\n",
    "\n",
    "from sklearn import preprocessing\n",
    "from sklearn.metrics import classification_report, confusion_matrix\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.neural_network import MLPClassifier\n",
    "import joblib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:11:44.470915Z",
     "iopub.status.busy": "2022-11-10T12:11:44.470650Z",
     "iopub.status.idle": "2022-11-10T12:11:44.547673Z",
     "shell.execute_reply": "2022-11-10T12:11:44.546845Z",
     "shell.execute_reply.started": "2022-11-10T12:11:44.470884Z"
    },
    "id": "5buPCSNApRNo",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def display_all(df, n=1000):\n",
    "    with pd.option_context(\"display.max_rows\", n, \"display.max_columns\", 1000):\n",
    "        display(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ATz_0_Z5pRNo"
   },
   "source": [
    "## 02 Loading data from FuncLib server\n",
    "This is where you use the files you got from FuncLib.\n",
    "\n",
    "Make sure all runs used the same refined structure and PSSM files (generated in the first FuncLib you ran).\n",
    "\n",
    "\n",
    "\n",
    "*   ***Minimal number of mutations per design*** was suppose to be set to 1\n",
    "*   ***Maximal number of mutations per design*** was suppose to be set to 100000 (or the number of positions in your library)\n",
    "*   ***Difference between clustered variants*** was suppose to be set to 1 (to get all designs calculated by Rosetta)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0rVR3Xoh5FBq"
   },
   "source": [
    "**EDIT HERE**   \n",
    "In the following block code, the scores of the different mutants calculated by FuncLib are being loaded. Either upload the files if using Google Colab or edit the file paths if running locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wiltDW2qfre4"
   },
   "outputs": [],
   "source": [
    "###################\n",
    "# for google colab\n",
    "###################\n",
    "# from google.colab import files\n",
    "# import io\n",
    "\n",
    "# uploaded = files.upload() # load the csv files generated by FuncLib\n",
    "# scores = [pd.read_csv(io.BytesIO(v)) for k, v in uploaded.items()]\n",
    "\n",
    "###################\n",
    "# for running locally\n",
    "###################\n",
    "paths = [] # add here the paths to the csv files generated by FuncLib\n",
    "scores = [pd.read_csv(p) for p in paths]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Th6ibmPMhBIw"
   },
   "outputs": [],
   "source": [
    "def count_mutations(serial_number):\n",
    "    pairs = [serial_number[i:i+2] for i in range(0, len(serial_number), 2)]\n",
    "    return sum([1 for i in pairs if i != '01'])\n",
    "\n",
    "scores = pd.concat(scores, ignore_index=True)\n",
    "\n",
    "scores.sort_values('total_score', inplace=True)\n",
    "scores.drop_duplicates(subset=['serial_number'], keep='last', inplace=True)\n",
    "\n",
    "scores.serial_number = scores.serial_number.apply(lambda x: str(x[1:]))\n",
    "scores['num_muts'] = scores.serial_number.apply(lambda x: count_mutations(x))\n",
    "\n",
    "library_positions = set(scores.columns) - set(['serial_number', 'num_muts', 'total_score'])\n",
    "library_positions = sorted(list(library_positions))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ebmaP6cHpRNp"
   },
   "source": [
    "## 03 Prepare data for training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XC25KFMRpRNq"
   },
   "source": [
    "### Sort by total score and find the WT score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-12T20:05:34.694950Z",
     "iopub.status.busy": "2022-11-12T20:05:34.694424Z",
     "iopub.status.idle": "2022-11-12T20:05:34.830517Z",
     "shell.execute_reply": "2022-11-12T20:05:34.829624Z",
     "shell.execute_reply.started": "2022-11-12T20:05:34.694889Z"
    },
    "id": "k8x7AqU5pRNq",
    "tags": []
   },
   "outputs": [],
   "source": [
    "scores.sort_values('total_score', inplace=True)\n",
    "WT = scores.loc[scores[\"num_muts\"] == 0, :].iloc[0]\n",
    "assert WT.num_muts == 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mOYFa1T2pRNq"
   },
   "source": [
    "### Set labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:13:02.721914Z",
     "iopub.status.busy": "2022-11-10T12:13:02.721346Z",
     "iopub.status.idle": "2022-11-10T12:13:02.854894Z",
     "shell.execute_reply": "2022-11-10T12:13:02.853751Z",
     "shell.execute_reply.started": "2022-11-10T12:13:02.721852Z"
    },
    "id": "WCE0GFUVpRNq",
    "tags": []
   },
   "outputs": [],
   "source": [
    "LABEL = \"y\"\n",
    "NLABEL = \"negative\"\n",
    "N = len(scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:13:03.991093Z",
     "iopub.status.busy": "2022-11-10T12:13:03.990515Z",
     "iopub.status.idle": "2022-11-10T12:13:04.110626Z",
     "shell.execute_reply": "2022-11-10T12:13:04.109542Z",
     "shell.execute_reply.started": "2022-11-10T12:13:03.991028Z"
    },
    "id": "-4sDwFbIpRNq",
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(WT.total_score)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LiQH1wfl7QVy"
   },
   "source": [
    "**EDIT HERE**.   \n",
    "The parameter ***delta_wt*** can be adjusted to control the selection stringency.\n",
    "\n",
    "delta_wt > 0 will allow mutants with a total score worse than the WT to be considered \"good.\" We do not recommend using a value higher than 5\n",
    "\n",
    "delta_wt < 0 will enforce a stringent selection of  \"good\" multipoint mutants. Use this option if a large part of the enumerated sequence space has a total score better than the WT sequence.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:13:19.103243Z",
     "iopub.status.busy": "2022-11-10T12:13:19.102663Z",
     "iopub.status.idle": "2022-11-10T12:13:19.226039Z",
     "shell.execute_reply": "2022-11-10T12:13:19.224618Z",
     "shell.execute_reply.started": "2022-11-10T12:13:19.103177Z"
    },
    "id": "lFZuiQoUpRNq"
   },
   "outputs": [],
   "source": [
    "delta_wt = 0 # edit here (see above)\n",
    "\n",
    "scores[LABEL] = 0\n",
    "scores.loc[scores['total_score'] <= WT.total_score + delta_wt, LABEL] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:13:25.282735Z",
     "iopub.status.busy": "2022-11-10T12:13:25.282144Z",
     "iopub.status.idle": "2022-11-10T12:13:25.406165Z",
     "shell.execute_reply": "2022-11-10T12:13:25.404690Z",
     "shell.execute_reply.started": "2022-11-10T12:13:25.282670Z"
    },
    "id": "sxGNbGD-pRNq"
   },
   "outputs": [],
   "source": [
    "label1 = sum(scores[LABEL])\n",
    "percent = 100 * (label1 / N)\n",
    "print(f'{label1} out of {N} multipoint mutants are labeled as 1 ({percent:.2f})')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l3sQH9ZM9WNP"
   },
   "source": [
    "We considere the bottom 50% of multipoint mutants to be \"bad\" designs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:13:54.053537Z",
     "iopub.status.busy": "2022-11-10T12:13:54.052933Z",
     "iopub.status.idle": "2022-11-10T12:13:54.173929Z",
     "shell.execute_reply": "2022-11-10T12:13:54.172951Z",
     "shell.execute_reply.started": "2022-11-10T12:13:54.053472Z"
    },
    "id": "A7nBtRXkpRNq",
    "tags": []
   },
   "outputs": [],
   "source": [
    "nlabel = np.zeros(N)\n",
    "nlabel[int(0.5 * N) :] = 1\n",
    "scores[NLABEL] = nlabel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "both_labels = len(scores[(scores.y == 1) & (scores.negative == 1)])\n",
    "if both_labels:\n",
    "    msg = f'<span style=\"color: #ff0000\">There are {both_labels} designs labeled as both good and bad </span>'\n",
    "    msg += '<br /> Please lower the **delta_wt** parameter.'\n",
    "    display(Markdown(msg))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate some statistics on the labeling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:14:13.363513Z",
     "iopub.status.busy": "2022-11-10T12:14:13.363283Z",
     "iopub.status.idle": "2022-11-10T12:14:13.455923Z",
     "shell.execute_reply": "2022-11-10T12:14:13.455291Z",
     "shell.execute_reply.started": "2022-11-10T12:14:13.363490Z"
    },
    "id": "66i-F-XipRNq"
   },
   "outputs": [],
   "source": [
    "label_rates = pd.DataFrame(columns=[\"stat\", \"value\"])\n",
    "label_rates.loc[len(label_rates)] = [\"Total number of mutants\", N]\n",
    "\n",
    "label_rates.loc[len(label_rates)] = ['WT score', WT.total_score]\n",
    "label_rates.loc[len(label_rates)] = ['limit for label 1', WT.total_score + delta_wt]\n",
    "\n",
    "l1 = len(scores[scores.y == 1])\n",
    "label_rates.loc[len(label_rates)] = ['Labeled 1',\n",
    "                                     f'{l1} ({(l1/N)*100:.2f})']\n",
    "\n",
    "l_neg = len(scores[scores[NLABEL] == 1])\n",
    "label_rates.loc[len(label_rates)] = ['Labeled negative',\n",
    "                                     f'{l_neg} ({(l_neg/N)*100:.2f})']\n",
    "\n",
    "non = len(scores[(scores.y == 0) & (scores.negative == 0)])\n",
    "label_rates.loc[len(label_rates)] = ['Non classified',\n",
    "                                     f'{non} ({(non/N)*100:.2f})']\n",
    "print('Stats on the labeling')\n",
    "label_rates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:16:24.915327Z",
     "iopub.status.busy": "2022-11-10T12:16:24.914742Z",
     "iopub.status.idle": "2022-11-10T12:16:25.734764Z",
     "shell.execute_reply": "2022-11-10T12:16:25.733873Z",
     "shell.execute_reply.started": "2022-11-10T12:16:24.915260Z"
    },
    "id": "parG6VlfpRNr",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# plot the histogram\n",
    "plt.figure(facecolor=\"w\")\n",
    "dist = sns.histplot(\n",
    "    scores[\"total_score\"],\n",
    "    color=\"#99ccff\",\n",
    "    kde=True,\n",
    "    bins=101,\n",
    ")\n",
    "\n",
    "top_negative = min(scores[scores[NLABEL] == 1][\"total_score\"])\n",
    "# color less than WT\n",
    "for p in dist.patches:\n",
    "    if p.get_x() < WT[\"total_score\"] + delta_wt:\n",
    "        p.set_facecolor(\"#66ff66\")\n",
    "    elif p.get_x() >= top_negative:\n",
    "        p.set_facecolor(\"#ff6666\")\n",
    "\n",
    "\n",
    "plt.gca().set_yticks([])\n",
    "\n",
    "patches = []\n",
    "patches.append(mpatches.Patch(color=\"#66ff66\", label=\"Labeled good\"))\n",
    "patches.append(mpatches.Patch(color=\"#99ccff\", label=\"non classified\"))\n",
    "patches.append(mpatches.Patch(color=\"#ff6666\", label=\"50% bottom\"))\n",
    "plt.legend(handles=patches)\n",
    "plt.title('Scores distribution')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9iJdcG-CpRNr"
   },
   "source": [
    "### Features preparation\n",
    "Creating the features for the NN to train on. We use one-hot encoding of the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:16:43.575117Z",
     "iopub.status.busy": "2022-11-10T12:16:43.574525Z",
     "iopub.status.idle": "2022-11-10T12:16:43.810112Z",
     "shell.execute_reply": "2022-11-10T12:16:43.809566Z",
     "shell.execute_reply.started": "2022-11-10T12:16:43.575048Z"
    },
    "id": "jg8-L3wvpRNr"
   },
   "outputs": [],
   "source": [
    "encoded_table = pd.get_dummies(scores[library_positions])\n",
    "encoded_table = encoded_table.loc[:, ~encoded_table.columns.duplicated()]\n",
    "FEATURES = encoded_table.columns\n",
    "scores[FEATURES] = encoded_table\n",
    "scores.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7TJuAeWlpRNr"
   },
   "source": [
    "### Split to train and test set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZQytGGFmDds4"
   },
   "source": [
    "**EDIT HERE**\n",
    "\n",
    "\n",
    "1.   **non labeled** - As the \"non labeled\" mutants are not sure enough, they are discarded from training.\n",
    "For using all multipoint mutants in training and testing uncomment the line below\n",
    "2.   **size of test set** - the size of the test portion of the data can be set using ***test_size***\n",
    "3.   **random state** - the random state for splitting the data can be set using ***random_state***\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:17:29.974572Z",
     "iopub.status.busy": "2022-11-10T12:17:29.973979Z",
     "iopub.status.idle": "2022-11-10T12:17:30.133567Z",
     "shell.execute_reply": "2022-11-10T12:17:30.132084Z",
     "shell.execute_reply.started": "2022-11-10T12:17:29.974507Z"
    },
    "id": "c3EBmW1npRNr"
   },
   "outputs": [],
   "source": [
    "positive = scores[scores[LABEL] == 1]\n",
    "negative = scores[scores[NLABEL] == 1]\n",
    "training = pd.concat([positive, negative])\n",
    "# training = scores # to include the non-labeled mutants in the training data\n",
    "\n",
    "test_size = 0.2\n",
    "random_state = 42\n",
    "train, test = train_test_split(training, test_size=test_size, random_state=random_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_gS4JzmupRNr"
   },
   "source": [
    "Leaving only the columns needed for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:17:49.837334Z",
     "iopub.status.busy": "2022-11-10T12:17:49.836743Z",
     "iopub.status.idle": "2022-11-10T12:17:49.962173Z",
     "shell.execute_reply": "2022-11-10T12:17:49.961134Z",
     "shell.execute_reply.started": "2022-11-10T12:17:49.837268Z"
    },
    "id": "uudZy_HIpRNr"
   },
   "outputs": [],
   "source": [
    "X_train = train[FEATURES]\n",
    "y_train = pd.DataFrame(train[LABEL], columns=[LABEL])\n",
    "\n",
    "X_test = test[FEATURES]\n",
    "y_test = pd.DataFrame(test[LABEL], columns=[LABEL])\n",
    "\n",
    "print(FEATURES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YbgKG60spRNr"
   },
   "source": [
    "### Save Data\n",
    "Saving all datasets as csv files\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:18:40.910851Z",
     "iopub.status.busy": "2022-11-10T12:18:40.910252Z",
     "iopub.status.idle": "2022-11-10T12:18:41.318173Z",
     "shell.execute_reply": "2022-11-10T12:18:41.317318Z",
     "shell.execute_reply.started": "2022-11-10T12:18:40.910785Z"
    },
    "id": "YLi4U7VupRNr"
   },
   "outputs": [],
   "source": [
    "scores.to_csv(\"all_data.csv\")\n",
    "\n",
    "X_train.to_pickle( \"X_train.pkl\")\n",
    "X_test.to_pickle(\"X_test.pkl\")\n",
    "\n",
    "y_train.to_pickle(\"y_train.pkl\")\n",
    "y_test.to_pickle(\"y_test.pkl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MqsE2mAJpRNr"
   },
   "source": [
    "### *Restore* all data and reinitialize whatever variables needed for later\n",
    "\n",
    "Can be skipped"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:18:52.391490Z",
     "iopub.status.busy": "2022-11-10T12:18:52.390898Z",
     "iopub.status.idle": "2022-11-10T12:18:52.675699Z",
     "shell.execute_reply": "2022-11-10T12:18:52.674809Z",
     "shell.execute_reply.started": "2022-11-10T12:18:52.391423Z"
    },
    "id": "hzJv7ounpRNr",
    "tags": []
   },
   "outputs": [],
   "source": [
    "scores = pd.read_csv(\"all_data.csv\")\n",
    "\n",
    "X_train = pd.read_pickle(\"X_train.pkl\")\n",
    "X_test = pd.read_pickle(\"X_test.pkl\")\n",
    "\n",
    "y_train = pd.read_pickle(\"y_train.pkl\")\n",
    "y_test = pd.read_pickle(\"y_test.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:19:19.845669Z",
     "iopub.status.busy": "2022-11-10T12:19:19.845065Z",
     "iopub.status.idle": "2022-11-10T12:19:19.967440Z",
     "shell.execute_reply": "2022-11-10T12:19:19.966492Z",
     "shell.execute_reply.started": "2022-11-10T12:19:19.845602Z"
    },
    "id": "maOksaCWpRNr",
    "tags": []
   },
   "outputs": [],
   "source": [
    "FEATURES = X_train.columns\n",
    "library_positions = sorted(list(set([i.split('_')[0] for i in FEATURES]))) ==\n",
    "\n",
    "LABEL = \"y\"\n",
    "NLABEL = \"negative\"\n",
    "N = len(scores)\n",
    "\n",
    "WT = scores.loc[scores[\"num_muts\"] == 0, :].iloc[1]\n",
    "assert WT.num_muts == 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZNZ1u2FXpRNr"
   },
   "source": [
    "## 04 Training the neural net model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:19:24.164538Z",
     "iopub.status.busy": "2022-11-10T12:19:24.163948Z",
     "iopub.status.idle": "2022-11-10T12:19:34.098368Z",
     "shell.execute_reply": "2022-11-10T12:19:34.097291Z",
     "shell.execute_reply.started": "2022-11-10T12:19:24.164471Z"
    },
    "id": "kw6IxYfipRNr",
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "mlp = MLPClassifier(\n",
    "    hidden_layer_sizes=(len(library_positions)),\n",
    "    activation=\"logistic\",\n",
    "    learning_rate=\"invscaling\",\n",
    "    random_state=42,\n",
    "    verbose=True,\n",
    "    max_iter=20000,\n",
    ")\n",
    "mlp.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c0Ynaq2kGt1s"
   },
   "source": [
    "Plotting the loss curve, Should monotonously decline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:19:43.458524Z",
     "iopub.status.busy": "2022-11-10T12:19:43.457905Z",
     "iopub.status.idle": "2022-11-10T12:19:43.704682Z",
     "shell.execute_reply": "2022-11-10T12:19:43.703947Z",
     "shell.execute_reply.started": "2022-11-10T12:19:43.458454Z"
    },
    "id": "B3mFrGpCpRNs"
   },
   "outputs": [],
   "source": [
    "plt.figure(facecolor=\"w\")\n",
    "plt.ion()\n",
    "ax = plt.gca()\n",
    "ax.set_title(\"loss curve\")\n",
    "plt.plot(mlp.loss_curve_)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UMALUkCYpRNs"
   },
   "source": [
    "### Save/ Load the model using joblib\n",
    "Uncomment the loading model line if restoring a previous run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:20:10.180443Z",
     "iopub.status.busy": "2022-11-10T12:20:10.179852Z",
     "iopub.status.idle": "2022-11-10T12:20:10.312093Z",
     "shell.execute_reply": "2022-11-10T12:20:10.311266Z",
     "shell.execute_reply.started": "2022-11-10T12:20:10.180361Z"
    },
    "id": "VLqUi9nqpRNs",
    "tags": []
   },
   "outputs": [],
   "source": [
    "filename = \"model.sav\"\n",
    "joblib.dump(mlp, filename)\n",
    "\n",
    "# mlp = joblib.load(filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "082i7wympRNw"
   },
   "source": [
    "### Analyzing predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:21:28.699616Z",
     "iopub.status.busy": "2022-11-10T12:21:28.699025Z",
     "iopub.status.idle": "2022-11-10T12:21:28.890450Z",
     "shell.execute_reply": "2022-11-10T12:21:28.888102Z",
     "shell.execute_reply.started": "2022-11-10T12:21:28.699550Z"
    },
    "id": "9j5qYpuGpRNw"
   },
   "outputs": [],
   "source": [
    "predictions = mlp.predict(scores[FEATURES])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:21:39.284192Z",
     "iopub.status.busy": "2022-11-10T12:21:39.283620Z",
     "iopub.status.idle": "2022-11-10T12:21:39.416738Z",
     "shell.execute_reply": "2022-11-10T12:21:39.415872Z",
     "shell.execute_reply.started": "2022-11-10T12:21:39.284127Z"
    },
    "id": "OAZnudbtpRNw",
    "tags": []
   },
   "outputs": [],
   "source": [
    "a = confusion_matrix(scores[LABEL], predictions)\n",
    "rates = pd.DataFrame(\n",
    "    [[a[0, 0], a[0, 1], a[1, 0], a[1, 1]]], columns=[\"True Negative\", \"False Positive\", \"False Negative\", \"True Positive\"]\n",
    ")\n",
    "rates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:21:52.026395Z",
     "iopub.status.busy": "2022-11-10T12:21:52.025862Z",
     "iopub.status.idle": "2022-11-10T12:21:52.165115Z",
     "shell.execute_reply": "2022-11-10T12:21:52.164229Z",
     "shell.execute_reply.started": "2022-11-10T12:21:52.026323Z"
    },
    "id": "DZPUrr4zpRNw"
   },
   "outputs": [],
   "source": [
    "print(classification_report(scores[LABEL], predictions))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JuN_iSWQpRNw"
   },
   "source": [
    "## 05 Rank mutaitons by EpiNNet\n",
    "the following blocks use the EpiNNet model to rank all tested mutations,\n",
    "refer to the weights_activations_1 column as rank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:22:02.664921Z",
     "iopub.status.busy": "2022-11-10T12:22:02.664328Z",
     "iopub.status.idle": "2022-11-10T12:22:02.926641Z",
     "shell.execute_reply": "2022-11-10T12:22:02.925860Z",
     "shell.execute_reply.started": "2022-11-10T12:22:02.664857Z"
    },
    "id": "juS27mO7pRNw"
   },
   "outputs": [],
   "source": [
    "pred_diag = pd.DataFrame([{col: 1} for col in FEATURES], columns=FEATURES) # just a diagonal matrix\n",
    "pred_diag.fillna(0, inplace=True)\n",
    "\n",
    "weights_activations = mlp.predict_proba(pred_diag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T12:22:26.214884Z",
     "iopub.status.busy": "2022-11-10T12:22:26.214293Z",
     "iopub.status.idle": "2022-11-10T12:22:26.348946Z",
     "shell.execute_reply": "2022-11-10T12:22:26.347984Z",
     "shell.execute_reply.started": "2022-11-10T12:22:26.214818Z"
    },
    "id": "WdS-O7c3pRNw",
    "tags": []
   },
   "outputs": [],
   "source": [
    "results = pd.DataFrame([i.split(\"_\") for i in FEATURES], columns=[\"position\", \"AA\"])\n",
    "results[\"weights_activations_0\"] = weights_activations[:, 0]\n",
    "results[\"weights_activations_1\"] = weights_activations[:, 1]\n",
    "\n",
    "results.sort_values(\"weights_activations_1\", ascending=False, inplace=True)\n",
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FowRYMWjpRNx"
   },
   "source": [
    "## 06 Find the best sequence space\n",
    "These blocks will help you construct the final sequence space.\n",
    "it iteratively adds additional mutations (from the top ranked mutaitons, as ranked by the EpiNNet model).\n",
    "for each potential sequence space, the size of the encoded library, and it's estimated good Vs. bad number of deisgns are evaluated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-11T13:22:29.711541Z",
     "iopub.status.busy": "2022-11-11T13:22:29.711122Z",
     "iopub.status.idle": "2022-11-11T13:22:59.737291Z",
     "shell.execute_reply": "2022-11-11T13:22:59.735803Z",
     "shell.execute_reply.started": "2022-11-11T13:22:29.711494Z"
    },
    "id": "Iaiw_EUipRNx",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_seq_space_upto_row(sub_results):\n",
    "    \"\"\"\n",
    "    Get a sequence space of the N top ranked mutations\n",
    "    :param sub_results: a pandas dataframe of the top ranked mutations\n",
    "    \"\"\"\n",
    "    seq_space = {}\n",
    "    for p, wt in WT[library_positions].items():\n",
    "        aas = sub_results.loc[sub_results[\"position\"] == p, \"AA\"].values\n",
    "        seq_space[p] = [wt] + [aa for aa in aas if aa != wt]\n",
    "    return seq_space\n",
    "\n",
    "def get_sub_score_df(sc_df, seq_space):\n",
    "    \"\"\"Get the scores of the mutation spanned by a sequece space\"\"\"\n",
    "    sc_df = sc_df.copy()\n",
    "    for p, aas in seq_space.items():\n",
    "        sc_df = sc_df.loc[sc_df[p].isin(aas)]\n",
    "    return sc_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-10T13:24:35.125597Z",
     "iopub.status.busy": "2022-11-10T13:24:35.125014Z",
     "iopub.status.idle": "2022-11-10T13:24:41.308637Z",
     "shell.execute_reply": "2022-11-10T13:24:41.307887Z",
     "shell.execute_reply.started": "2022-11-10T13:24:35.125532Z"
    },
    "id": "HY0zUp1ypRNx",
    "tags": []
   },
   "outputs": [],
   "source": [
    "results.sort_values(\"weights_activations_1\", ascending=False, inplace=True)\n",
    "\n",
    "rows = list()\n",
    "for i in tqdm(range(len(results))):\n",
    "    seq_space = get_seq_space_upto_row(results.iloc[:i])\n",
    "    sub_scores = get_sub_score_df(scores, seq_space)\n",
    "\n",
    "    row = {p: ''.join(AAs) for p, AAs in seq_space.items()}\n",
    "    row[\"i\"] = i\n",
    "    row[\"good\"] = len(sub_scores.loc[sub_scores[\"y\"] == 1])\n",
    "    row[\"bad\"] = len(sub_scores.loc[sub_scores[\"negative\"] == 1])\n",
    "    row[\"total\"] = len(sub_scores)\n",
    "    rows.append(row)\n",
    "\n",
    "summ_df = pd.DataFrame(rows)\n",
    "\n",
    "summ_df[\"lib_size\"] = summ_df.apply(lambda r: np.product([len(r[p]) for p in library_positions]), axis=1)\n",
    "summ_df[\"rest\"] = summ_df.total - summ_df.good - summ_df.bad\n",
    "summ_df[\"good_rel\"] = summ_df[\"good\"] / summ_df[\"total\"]\n",
    "summ_df[\"bad_rel\"] = summ_df[\"bad\"] / summ_df[\"total\"]\n",
    "summ_df[\"rest_rel\"] = summ_df[\"rest\"] / summ_df[\"total\"]\n",
    "display_all(summ_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RtqJyO8_pRNx"
   },
   "source": [
    "### Sub-sequence space bar plot\n",
    "The following bar plot shows the good/bad/unknown rates for all sub sequence spaces, as ranked by EpiNNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-11T12:15:13.040170Z",
     "iopub.status.busy": "2022-11-11T12:15:13.039593Z",
     "iopub.status.idle": "2022-11-11T12:15:13.675818Z",
     "shell.execute_reply": "2022-11-11T12:15:13.675272Z",
     "shell.execute_reply.started": "2022-11-11T12:15:13.040105Z"
    },
    "id": "Y3q4SJP5pRNx",
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n",
    "bottoms = [0] * len(summ_df)\n",
    "\n",
    "for j, status in enumerate([\"good\", \"bad\", \"rest\"]):\n",
    "    ax.bar(\n",
    "        x=summ_df.i,\n",
    "        height=summ_df[f\"{status}_rel\"],\n",
    "        bottom=bottoms,\n",
    "        width=0.8,\n",
    "        label=status,\n",
    "        facecolor=sns.color_palette(\"pastel\", 3)[j],\n",
    "        edgecolor=\"k\",\n",
    "    )\n",
    "    bottoms = [b + s for b, s in zip(bottoms, summ_df[f\"{status}_rel\"])]\n",
    "ax.set_xlabel(\"mutation index\")\n",
    "ax.set_ylabel(\"frequency in modeled data\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QWpD-DleYDxw"
   },
   "source": [
    "The above figure shows that the more positions are included, the fraction of bad mutants included in the library is increasing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HufpjM2qpRNx"
   },
   "source": [
    "### Twick the lower and upper limits to find your favorite sequence space\n",
    "The numbers in the next block will help you examine sequence spaces encoding library sizes in the range you specify.\n",
    "\n",
    "**EDIT HERE**.    \n",
    "change the min and max library size and see also the fraction of good and bad designs included"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-11T12:17:35.631357Z",
     "iopub.status.busy": "2022-11-11T12:17:35.630777Z",
     "iopub.status.idle": "2022-11-11T12:17:35.788979Z",
     "shell.execute_reply": "2022-11-11T12:17:35.788176Z",
     "shell.execute_reply.started": "2022-11-11T12:17:35.631293Z"
    },
    "id": "Gtb5ECUjpRNx",
    "tags": []
   },
   "outputs": [],
   "source": [
    "min_library_size = 1000\n",
    "max_library_size = 100000\n",
    "summ_df[summ_df.lib_size.between(min_library_size, max_library_size)][['good_rel', 'bad_rel', 'lib_size']]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ne8Fhan8aSCL"
   },
   "source": [
    "**EDIT HERE**.     \n",
    "Change the variable ***mutation_index*** to see the library generated by the top N residues selected by EpiNNet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-11T12:36:06.346689Z",
     "iopub.status.busy": "2022-11-11T12:36:06.346098Z",
     "iopub.status.idle": "2022-11-11T12:36:06.484468Z",
     "shell.execute_reply": "2022-11-11T12:36:06.483584Z",
     "shell.execute_reply.started": "2022-11-11T12:36:06.346624Z"
    },
    "id": "edLad8xJpRNy",
    "tags": []
   },
   "outputs": [],
   "source": [
    "mutation_index = 20 # change here to different values and see the resulting library\n",
    "\n",
    "chosen_seq_space = get_seq_space_upto_row(results.iloc[:mutation_index])\n",
    "sub_scores = get_sub_score_df(scores, chosen_seq_space)\n",
    "\n",
    "lib_size = np.product(list(map(len, chosen_seq_space.values())))\n",
    "goods = len(sub_scores.loc[sub_scores[\"y\"] == 1])\n",
    "bads = len(sub_scores.loc[sub_scores[\"negative\"] == 1])\n",
    "N = len(sub_scores)\n",
    "\n",
    "print(\"position\\tAAs\")\n",
    "print('-'*30)\n",
    "for p, aas in chosen_seq_space.items():\n",
    "    print(f'{p}\\t\\t{\"\".join(aas)}')\n",
    "\n",
    "print(f\"the chosen sequence space encodes a library of size {lib_size:,}\")\n",
    "print(f'the library includes an estimate of {goods/N * 100:.2f}% good designs')\n",
    "print(f'the library includes an estimate of {bads/N * 100:.2f}% bad designs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lBkTiS9KpRNy"
   },
   "source": [
    "## 07 Test random designs from the library\n",
    "The library was constructed based on multipoint mutants from the general sequence space, and maybe only seen designs from bubbles and a small percentage of the possible sequence space.\n",
    "\n",
    "To estimate better the rate of good and bad designs in the chosen library, run a final FuncLib calculation with the selected sequence space.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2zWQYSUfpRNy"
   },
   "source": [
    "### Load chosen sequence space mutants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-12T14:19:16.152308Z",
     "iopub.status.busy": "2022-11-12T14:19:16.152094Z",
     "iopub.status.idle": "2022-11-12T14:19:16.351997Z",
     "shell.execute_reply": "2022-11-12T14:19:16.351442Z",
     "shell.execute_reply.started": "2022-11-12T14:19:16.152287Z"
    },
    "id": "MokY3fQxpRNy",
    "tags": []
   },
   "outputs": [],
   "source": [
    "###################\n",
    "# for google colab\n",
    "###################\n",
    "uploaded_chosen = files.upload() # load the csv files generated by FuncLib\n",
    "chosen = pd.read_csv(io.BytesIO(list(uploaded_chosen.values())[0]))\n",
    "\n",
    "###################\n",
    "# for running locally\n",
    "###################\n",
    "# path = '' # add here the path to the csv files generated by FuncLib\n",
    "# chosen = pd.read_csv(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PQOYERr4gk6J"
   },
   "source": [
    "Select a random set of designs from the general sequence space to compare"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pt-mD5BZgf6Y"
   },
   "outputs": [],
   "source": [
    "N = chosen.shape[0]\n",
    "general = scores.sample(N)\n",
    "\n",
    "chosen['seq_space'] = 'chosen'\n",
    "general['seq_space'] = 'general'\n",
    "test_10k_df = pd.concat([chosen, general])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NMwRDAQPh_-z"
   },
   "source": [
    "Plotting the score distributions. The chosen library will have a distribution shifted to the left (i.e. lower Rosetta scores)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-12T19:52:50.937480Z",
     "iopub.status.busy": "2022-11-12T19:52:50.936898Z",
     "iopub.status.idle": "2022-11-12T19:52:51.309784Z",
     "shell.execute_reply": "2022-11-12T19:52:51.309216Z",
     "shell.execute_reply.started": "2022-11-12T19:52:50.937416Z"
    },
    "id": "2EbwyYpTpRNy",
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(10, 5), facecolor=\"w\")\n",
    "sns.kdeplot(\n",
    "    data=test_10k_df,\n",
    "    x=\"total_score\",\n",
    "    hue=\"seq_space\",\n",
    "    palette=\"pastel\",\n",
    "    fill=True,\n",
    "    ax=ax,\n",
    "    linewidth=2,\n",
    ")\n",
    "plt.axvline(WT.total_score, linestyle=\"--\", c=\"k\")\n",
    "\n",
    "for i, (name, height) in enumerate(zip([\"chosen\", \"general\"], [0.8, 0.7])):\n",
    "    for over_under, left_right in zip([\"over\", \"under\"], [0.65, 0.05]):\n",
    "        if over_under == \"under\":\n",
    "            top = len(\n",
    "                test_10k_df.loc[\n",
    "                    (test_10k_df[\"seq_space\"] == name)\n",
    "                    & (test_10k_df[\"total_score\"] < WT.total_score)\n",
    "                ]\n",
    "            )\n",
    "        else:\n",
    "            top = len(\n",
    "                test_10k_df.loc[\n",
    "                    (test_10k_df[\"seq_space\"] == name)\n",
    "                    & (test_10k_df[\"total_score\"] > WT.total_score)\n",
    "                ]\n",
    "            )\n",
    "        ax.text(\n",
    "            left_right,\n",
    "            height,\n",
    "            f'{100*top/len(test_10k_df.loc[test_10k_df[\"seq_space\"] == name]):.1f}%',\n",
    "            color=sns.color_palette(\"pastel\")[i],\n",
    "            size=12,\n",
    "            transform=ax.transAxes,\n",
    "        )\n",
    "\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}