predict_degrons_and_binding_E3s.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from tape import ProteinBertModel, TAPETokenizer\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os,re,math\n",
    "\n",
    "device = \"cuda:0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### load TAPE bert-based pre-trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load pretrained model https://github.com/songlab-cal/tape\n",
    "model = ProteinBertModel.from_pretrained('bert-base').to(device)\n",
    "\n",
    "tokenizer = TAPETokenizer(vocab='iupac')\n",
    "\n",
    "def seq_embed(seq):\n",
    "    seq_tensor = torch.tensor([tokenizer.encode(seq)]).to(device)\n",
    "    seq_bert = model(seq_tensor)[0][0]\n",
    "    return seq_bert"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Degpred architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the architecture\n",
    "class DEG_LSTM(nn.Module):\n",
    "    def __init__(self, input_size, deg_lstm_hidden_size, fc1_output_size, output_size):\n",
    "        super().__init__()\n",
    "        self.deg_lstm = nn.LSTM(input_size, deg_lstm_hidden_size, 2, bidirectional=True, batch_first = True)\n",
    "        self.deg_fc1 = nn.Linear(deg_lstm_hidden_size*2, fc1_output_size)\n",
    "        self.deg_fc2 = nn.Linear(fc1_output_size, output_size)\n",
    "\n",
    "    def forward(self, input):\n",
    "        output, (h, c) = self.deg_lstm(input)\n",
    "        output = self.deg_fc1(output)\n",
    "        output = torch.sigmoid(self.deg_fc2(output))\n",
    "        return output\n",
    "\n",
    "# load five models\n",
    "lstm1 = torch.load('model/five_model/degpred_model1.pkl').to(device)\n",
    "lstm2 = torch.load('model/five_model/degpred_model2.pkl').to(device)\n",
    "lstm3 = torch.load('model/five_model/degpred_model3.pkl').to(device)\n",
    "lstm4 = torch.load('model/five_model/degpred_model4.pkl').to(device)\n",
    "lstm5 = torch.load('model/five_model/degpred_model5.pkl').to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# average the outputs of five models\n",
    "def degpred(seq_bert):\n",
    "    pred1 = lstm1(seq_bert)\n",
    "    pred2 = lstm2(seq_bert)\n",
    "    pred3 = lstm3(seq_bert)\n",
    "    pred4 = lstm4(seq_bert)\n",
    "    pred5 = lstm5(seq_bert)\n",
    "    pred = (pred1 + pred2 + pred3 + pred4 + pred5)/5\n",
    "    return pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### predict a sequence (example: P53)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "seq = 'MEEPQSDPSVEPPLSQETFSDLWKLLPENNVLSPLPSQAMDDLMLSPDDIEQWFTEDPGPDEAPRMPEAAPPVAPAPAAPTPAAPAPAPSWPLSSSVPSQKTYQGSYGFRLGFLHSGTAKSVTCTYSPALNKMFCQLAKTCPVQLWVDSTPPPGTRVRAMAIYKQSQHMTEVVRRCPHHERCSDSDGLAPPQHLIRVEGNLRVEYLDDRNTFRHSVVVPYEPPEVGSDCTTIHYNYMCNSSCMGGMNRRPILTIITLEDSSGNLLGRNSFEVRVCACPGRDRRTEEENLRKKGEPHHELPPGSTKRALPNNTSSSPQPKKKPLDGEYFTLQIRGRERFEMFRELNEALELKDAQAGKEPGGSRAHSSHLKSKKGQSTSRHKKLMFKTEGPDSD'\n",
    "seq_bert = seq_embed(seq)\n",
    "pred = degpred(seq_bert.unsqueeze(0)).squeeze().cpu().detach().numpy()[1:-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### find degrons on the sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def continusFind(num_list):\n",
    "    s=1\n",
    "    find_list=[]\n",
    "    have_list=[]\n",
    "    while s <= len(num_list)-1:\n",
    "        if num_list[s] - num_list[s-1] < 4:\n",
    "            flag=s-1\n",
    "            while (s<=len(num_list)-1) and (num_list[s]-num_list[s-1] < 4):\n",
    "                s+=1\n",
    "            find_list.append([num_list[flag], num_list[s-1]])\n",
    "            have_list+=num_list[flag:s]\n",
    "        else:\n",
    "            if abs(num_list[s-1]-num_list[s-2]) > 3:\n",
    "                find_list.append(num_list[s-1:s])\n",
    "            s+=1\n",
    "    return find_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "start: 16 end:  27 degron_seq: ETFSDLWKLLPE\n",
      "start: 256 end:  262 degron_seq: LEDSSGN\n",
      "start: 282 end:  290 degron_seq: RTEEENLRK\n"
     ]
    }
   ],
   "source": [
    "thre = 0.3\n",
    "deg_position = list(np.where(pred > thre)[0])\n",
    "deg_interval = [i for i in continusFind(deg_position) if i[-1] - i[0] > 2]\n",
    "\n",
    "for j in deg_interval:\n",
    "    print('start:', j[0], 'end: ', j[1], 'degron_seq:', seq[j[0]: j[1]+1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### predict binding E3s of the degrons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "e3 = pd.read_csv('motifs/pssm_cutoffs.csv')\n",
    "pssms = []\n",
    "for i in e3.index:\n",
    "    a = pd.read_table('motifs/' + e3.loc[i, 'E3_entry'] + '_' + str(e3.loc[i, 'length']) + '_pssm.txt', index_col=0)\n",
    "    a.columns = a.columns.astype(int)\n",
    "    pssms.append(a)\n",
    "e3['pssm'] = pssms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def findE3(dseq):\n",
    "    e3s = []\n",
    "    for i in e3.index:\n",
    "        scores = [0,]\n",
    "        ps = e3.loc[i, 'pssm']\n",
    "        length = e3.loc[i, 'length']\n",
    "        try:\n",
    "            for k in range(len(dseq) - length + 1):\n",
    "                p = dseq[k : k+length]\n",
    "                s = 0\n",
    "                for j in range(length):\n",
    "                    s += ps.loc[p[j], j+1]\n",
    "                scores.append(s)\n",
    "            if max(scores) > e3.loc[i, 'thre1000']:\n",
    "                e3s.append(e3.loc[i, 'E3'])\n",
    "        except:\n",
    "            print('error in degron sequence', dseq)\n",
    "    return e3s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in deg_interval:\n",
    "    print(i, seq[i[0]: i[1]+1], findE3(seq[max(0, i[0]-3): min(i[1]+4, len(seq))]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}