tm_estimate_gui.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "eee83991",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\richa\\anaconda3\\envs\\KaggleComptetion_2022\\lib\\site-packages\\scipy\\__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.23.4\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    }
   ],
   "source": [
    "import PySimpleGUI as sg\n",
    "import xgboost as xgb\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import gc\n",
    "import torch, esm\n",
    "from sklearn.decomposition import PCA\n",
    "import pickle as pk\n",
    "import re"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e83afa8",
   "metadata": {},
   "source": [
    "# Defining functions to use with GUI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cbe316e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Preloads Facebook ESM model\n",
    "\n",
    "def torch_batch_converter ():\n",
    "    token_map = {'L': 0, 'A': 1, 'G': 2, 'V': 3, 'S': 4, 'E': 5, 'R': 6, 'T': 7, 'I': 8, 'D': 9, 'P': 10, \n",
    "             'K': 11, 'Q': 12, 'N': 13, 'F': 14, 'Y': 15, 'M': 16, 'H': 17, 'W': 18, 'C': 19}\n",
    "    t_model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "    batch_converter = alphabet.get_batch_converter()\n",
    "    t_model.eval()  # disables dropout for deterministic results\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    t_model.to(device)\n",
    "    \n",
    "    return t_model, batch_converter, device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "eed6a3fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Passes protein sequence into Facebook ESM model and extracts the last layer\n",
    "\n",
    "def seq_analyse (seq, t_model, batch_converter, device):\n",
    "    all_pdb_embed_pool_test = []\n",
    "\n",
    "    from scipy.special import softmax \n",
    "\n",
    "    data = [(\"protein1\", seq)]\n",
    "    batch_labels, batch_strs, batch_tokens = batch_converter(data)\n",
    "    batch_tokens = batch_tokens.to(device)\n",
    "    with torch.no_grad():\n",
    "        results = t_model(batch_tokens, repr_layers=[33])\n",
    "\n",
    "    results = results[\"representations\"][33].detach().cpu().numpy()    \n",
    "    torch_model_results = np.mean( results[0,:,:],axis=0 )       \n",
    "\n",
    "    del batch_tokens, results, batch_labels, batch_strs\n",
    "    gc.collect(); torch.cuda.empty_cache()\n",
    "    \n",
    "    return pd.DataFrame (torch_model_results.reshape(1,-1), columns = ['val_{}'.format(index) for index in range(1280)], index = ['0'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ceff8ba8",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Runs a pre-loaded PCA model on ECM dataset\n",
    "\n",
    "def pca_seq_df (df, pca_model):\n",
    "    pca_list = [col for col in df.columns.tolist() if col.__contains__('val')]\n",
    "    pca_pool = pk.load(open(pca_model, 'rb'))\n",
    "    pca_esm = pca_pool.transform(df[pca_list])\n",
    "    \n",
    "    return pd.DataFrame (pca_esm, columns = ['pca_val_{}'.format(i) for i in range (pca_esm.shape[1])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4f330004",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Runs a preloaded XGBoost model to predict Tm. Returns melting temperature, rounded to 1 digit\n",
    "\n",
    "def xgb_predict (df, xgb_model):\n",
    "    \n",
    "    dpredict = xgb.DMatrix (data=df)\n",
    "    \n",
    "    bst = xgb.Booster()\n",
    "    bst.load_model (xgb_model)\n",
    "    tm = bst.predict (dpredict)\n",
    "    \n",
    "    return tm[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b83997b",
   "metadata": {},
   "source": [
    "# Main GUI code\n",
    "\n",
    "Code will create a new window in which the user can enter in an amino acid sequence (using the one-letter code for the amino aicd) for proteins up to 2000 amino acids of size. \n",
    "\n",
    "PCA and XGBoost models are already pre-trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b7355b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading the ESM Model. This may take a few minutes.\n",
      "Extraction of features from ESM model is occuring. This may take a few minutes if the protein is very large.\n",
      "Features from ESM are undergoing PCA!\n",
      "Features are running through the XGBoost model. Almost there!\n"
     ]
    }
   ],
   "source": [
    "def main ():\n",
    "    \n",
    "    \n",
    "    #Identifies loading locations for XGBoost and PCA models\n",
    "    load_model_loc = 'XGBoost Model'\n",
    "    xgb_model = load_model_loc + '/XGB_model.xgb'\n",
    "    pca_model = load_model_loc+'/pca.pkl'\n",
    "    \n",
    "    predict_tm_bool = False\n",
    "    \n",
    "    #Creates UI\n",
    "    layout = [\n",
    "             [sg.Text('Please input protein sequence - use the one-letter code for each amino acid.'), sg.InputText('', key = 'input_text')],\n",
    "             [sg.Text('', key = 'instruction')],\n",
    "             [sg.Button ('Preload ESM model'),sg.Button('Predict Tm', disabled = True), sg.Button ('Exit')]\n",
    "             ]\n",
    "    \n",
    "    window = sg.Window ('Protein Stability Prediction', layout, modal = True)\n",
    "    \n",
    "    while True:\n",
    "        event, values = window.read()\n",
    "        \n",
    "        if event == 'Preload ESM model':\n",
    "            try:\n",
    "                #Preloads ESM model\n",
    "                print ('Loading the ESM Model. This may take a few minutes.')\n",
    "                t_model, batch_converter, device = torch_batch_converter ()\n",
    "                window['instruction'].update ('ESM Model loaded successfully. You can now predict the melting temperature.')\n",
    "                predict_tm_bool = True\n",
    "            except ExceptionError as e:\n",
    "                window['instruction'].update ('Error. Could not load ESM model.')\n",
    "        elif event == 'Predict Tm':\n",
    "            \n",
    "            #Error checks entered protein sequence to ensure that no incompatible amino acids, or unexpected characters are entered.\n",
    "            protein_seq = values ['input_text'].upper()\n",
    "            \n",
    "            protein_seq_regex = r'[^A-Z]+'\n",
    "            \n",
    "            if protein_seq != '':\n",
    "                if len (protein_seq) < 2000:\n",
    "                    if re.search (protein_seq_regex, protein_seq) == None:\n",
    "                        invalid_aa_regex = r'[BJOUXZ]'\n",
    "                        if re.search (invalid_aa_regex, protein_seq) == None:\n",
    "                            print ('Extraction of features from ESM model is occuring. This may take a few minutes if the protein is very large.')\n",
    "                            df = seq_analyse (protein_seq, t_model, batch_converter, device)\n",
    "                            print ('Features from ESM are undergoing PCA!')\n",
    "                            pca_df = pca_seq_df (df, pca_model)\n",
    "                            print ('Features are running through the XGBoost model. Almost there!')\n",
    "                            tm_estimate = xgb_predict (pca_df, xgb_model)\n",
    "                            \n",
    "                            tm_estimate = str (round(tm_estimate,1))                            \n",
    "                            \n",
    "                            window['instruction'].update ('Analysis complete. Estimated melting temperature is {}.'.format (tm_estimate))\n",
    "                            protein_seq = ''\n",
    "                        else:\n",
    "                            window['instruction'].update ('Error. Invalid Amino Acid code was entered. B, J, O, U, X, and Z are invalid amino acids.')\n",
    "                            window['input_text'].update ('')\n",
    "                    else:\n",
    "                        window['instruction'].update ('Error. Improper sequence entered. Please only use letters. B, J, O, U, X, and Z are not valid amino acids to enter for analysis.')\n",
    "                        window['input_text'].update ('')\n",
    "                else:\n",
    "                    window['instruction'].update ('Error. Protein sequence is too large for analysis. Please enter sequences less than 2000 amino acids in length.')\n",
    "                    window['input_text'].update ('')\n",
    "            else:\n",
    "                window['instruction'].update ('Error. Please enter a protein sequence.')\n",
    "            \n",
    "            \n",
    "            \n",
    "        elif event == sg.WIN_CLOSED or event == 'Exit':\n",
    "            break\n",
    "            \n",
    "        if predict_tm_bool:\n",
    "            window['Predict Tm'].update (disabled = False)\n",
    "    \n",
    "    window.close()\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}