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",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os,re,math\n",
    "\n",
    "from degpred_model import DegpredEmbedder, get_degpred_model\n",
    "\n",
    "device = \"cuda:0\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### load TAPE bert-based pre-trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "embed = DegpredEmbedder(device=device)\n",
    "# define the architecture\n",
    "degpred_models = [get_degpred_model(fold, device=device) for fold in range(1, 6)]\n",
    "\n",
    "# average the outputs of five models\n",
    "def degpred(seq_bert):\n",
    "    preds = [model(seq_bert) for model in degpred_models]\n",
    "    return torch.stack(preds, dim=0).mean(dim=0)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### predict a sequence (example: P53)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "seq = 'MEEPQSDPSVEPPLSQETFSDLWKLLPENNVLSPLPSQAMDDLMLSPDDIEQWFTEDPGPDEAPRMPEAAPPVAPAPAAPTPAAPAPAPSWPLSSSVPSQKTYQGSYGFRLGFLHSGTAKSVTCTYSPALNKMFCQLAKTCPVQLWVDSTPPPGTRVRAMAIYKQSQHMTEVVRRCPHHERCSDSDGLAPPQHLIRVEGNLRVEYLDDRNTFRHSVVVPYEPPEVGSDCTTIHYNYMCNSSCMGGMNRRPILTIITLEDSSGNLLGRNSFEVRVCACPGRDRRTEEENLRKKGEPHHELPPGSTKRALPNNTSSSPQPKKKPLDGEYFTLQIRGRERFEMFRELNEALELKDAQAGKEPGGSRAHSSHLKSKKGQSTSRHKKLMFKTEGPDSD'\n",
    "seq_bert = embed(seq)\n",
    "pred = degpred(seq_bert.unsqueeze(0)).squeeze().cpu().detach().numpy()[1:-1]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### find degrons on the sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_connected_ranges_and_mask(pred, prob_thres=0.3, connect_thres=3):\n",
    "    num_list = torch.where(pred > prob_thres)[0].tolist()\n",
    "    \n",
    "    ranges=[]\n",
    "    s=1\n",
    "    while s <= len(num_list)-1:\n",
    "        if num_list[s] - num_list[s-1] <= connect_thres:\n",
    "            flag=s-1\n",
    "            while (s<=len(num_list)-1) and (num_list[s]-num_list[s-1] <= connect_thres):\n",
    "                s+=1\n",
    "            ranges.append(range(num_list[flag], num_list[s-1]))\n",
    "        else:\n",
    "            if abs(num_list[s-1]-num_list[s-2]) > connect_thres:\n",
    "                ranges.append(range(num_list[s-1], num_list[s-1]+1))\n",
    "            s+=1\n",
    "            \n",
    "    mask = torch.zeros(len(pred), dtype=torch.bool, device=pred.device)\n",
    "    for rng in ranges:\n",
    "        mask[rng] = True\n",
    "    return ranges, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'numpy.ndarray' object has no attribute 'device'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[13], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m deg_ranges, deg_mask \u001b[39m=\u001b[39m get_connected_ranges_and_mask(pred, prob_thres\u001b[39m=\u001b[39;49m\u001b[39m0.3\u001b[39;49m, connect_thres\u001b[39m=\u001b[39;49m\u001b[39m3\u001b[39;49m)\n\u001b[1;32m      3\u001b[0m \u001b[39mfor\u001b[39;00m rng \u001b[39min\u001b[39;00m deg_ranges:\n\u001b[1;32m      4\u001b[0m     \u001b[39mprint\u001b[39m(\u001b[39m'\u001b[39m\u001b[39mstart:\u001b[39m\u001b[39m'\u001b[39m, rng\u001b[39m.\u001b[39mstart, \u001b[39m'\u001b[39m\u001b[39mend: \u001b[39m\u001b[39m'\u001b[39m, rng\u001b[39m.\u001b[39mstop, \u001b[39m'\u001b[39m\u001b[39mdegron_seq:\u001b[39m\u001b[39m'\u001b[39m, seq[rng])\n",
      "Cell \u001b[0;32mIn[12], line 17\u001b[0m, in \u001b[0;36mget_connected_ranges_and_mask\u001b[0;34m(pred, prob_thres, connect_thres)\u001b[0m\n\u001b[1;32m     14\u001b[0m             ranges\u001b[39m.\u001b[39mappend(\u001b[39mrange\u001b[39m(num_list[s\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m], num_list[s\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m]\u001b[39m+\u001b[39m\u001b[39m1\u001b[39m))\n\u001b[1;32m     15\u001b[0m         s\u001b[39m+\u001b[39m\u001b[39m=\u001b[39m\u001b[39m1\u001b[39m\n\u001b[0;32m---> 17\u001b[0m mask \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mzeros(\u001b[39mlen\u001b[39m(pred), dtype\u001b[39m=\u001b[39mtorch\u001b[39m.\u001b[39mbool, device\u001b[39m=\u001b[39mpred\u001b[39m.\u001b[39;49mdevice)\n\u001b[1;32m     18\u001b[0m \u001b[39mfor\u001b[39;00m rng \u001b[39min\u001b[39;00m ranges:\n\u001b[1;32m     19\u001b[0m     mask[rng] \u001b[39m=\u001b[39m \u001b[39mTrue\u001b[39;00m\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'numpy.ndarray' object has no attribute 'device'"
     ]
    }
   ],
   "source": [
    "\n",
    "deg_ranges, deg_mask = get_connected_ranges_and_mask(pred, prob_thres=0.3, connect_thres=3)\n",
    "\n",
    "for rng in deg_ranges:\n",
    "    print('start:', rng.start, 'end: ', rng.stop, 'degron_seq:', seq[rng])"
   ]
  },
  {
   "attachments": {},
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}