PROSTATA_tool.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/AIRI-Institute/PROSTATA/blob/main/PROSTATA_tool.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f59Ujuujn___"
   },
   "source": [
    "Code is provided according with Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License\n",
    "# Install dependecies and download weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "apiUcTpNTnlU",
    "outputId": "d7491c10-dae0-4701-844d-1c76d0353a04",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install fair-esm\n",
    "!pip install biopython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "Bsyfz4BrSxMN"
   },
   "outputs": [],
   "source": [
    "from io import StringIO, BytesIO\n",
    "from urllib.request import urlretrieve\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "import esm\n",
    "from esm.pretrained import load_model_and_alphabet_hub\n",
    "from Bio import SeqIO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "N3Pjcm9HvRMC"
   },
   "outputs": [],
   "source": [
    "model_names = [\n",
    "    \"ESMForSingleMutationPosOuter\",\n",
    "    \"ESMForSingleMutationPosConcat\",\n",
    "    \"ESMForSingleMutation_pos_cat_cls\",\n",
    "    \"ESMForSingleMutation_pos\",\n",
    "    \"ESMForSingleMutation_cls\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "for model_name in model_names:\n",
    "    urlretrieve(\n",
    "        f\"https://a025generative-modeling-for-design.obs.ru-moscow-1.hc.sbercloud.ru/prostata/mix_ds_s669_weights/{model_name}_mix_ds_s669\",\n",
    "        model_name,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "hYLdBMG7UUx3"
   },
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "HIDDEN_UNITS_POS_CONTACT = 5\n",
    "\n",
    "\n",
    "class ESMForSingleMutationPosConcat(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.esm2, _ = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "        self.fc1 = nn.Linear(1280 * 2, HIDDEN_UNITS_POS_CONTACT)\n",
    "        self.fc2 = nn.Linear(HIDDEN_UNITS_POS_CONTACT, 1)\n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):\n",
    "        outputs1 = self.esm2.forward(token_ids1, repr_layers=[33])[\"representations\"][33]\n",
    "        outputs2 = self.esm2.forward(token_ids2, repr_layers=[33])[\"representations\"][33]\n",
    "        outputs1_pos = outputs1[:, pos + 1]\n",
    "        outputs2_pos = outputs2[:, pos + 1]\n",
    "        outputs_pos_concat = torch.cat((outputs1_pos, outputs2_pos), 2)\n",
    "        fc1_outputs = F.relu(self.fc1(outputs_pos_concat))\n",
    "        logits = self.fc2(fc1_outputs)\n",
    "        return logits\n",
    "\n",
    "\n",
    "HIDDEN_UNITS_POS_OUTER = 5\n",
    "\n",
    "\n",
    "class ESMForSingleMutationPosOuter(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.esm2, _ = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "        self._freeze_esm2_layers()\n",
    "        self.fc1 = nn.Linear(1280 * 1280, HIDDEN_UNITS_POS_OUTER)\n",
    "        self.fc2 = nn.Linear(HIDDEN_UNITS_POS_OUTER, 1)\n",
    "\n",
    "    def _freeze_esm2_layers(self):\n",
    "        total_blocks = 33\n",
    "        initial_layers = 2\n",
    "        layers_per_block = 16\n",
    "        num_freeze_blocks = total_blocks - 3\n",
    "        for _, param in list(self.esm2.named_parameters())[: initial_layers + layers_per_block * num_freeze_blocks]:\n",
    "            param.requires_grad = False\n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):\n",
    "        outputs1 = self.esm2.forward(token_ids1, repr_layers=[33])[\"representations\"][33]\n",
    "        outputs2 = self.esm2.forward(token_ids2, repr_layers=[33])[\"representations\"][33]\n",
    "        outputs1_pos = outputs1[:, pos + 1]\n",
    "        outputs2_pos = outputs2[:, pos + 1]\n",
    "        outer_prod = outputs1_pos.unsqueeze(3) @ outputs2_pos.unsqueeze(2)\n",
    "        outer_prod_view = outer_prod.view(outer_prod.shape[0], outer_prod.shape[1], -1)\n",
    "        fc1_outputs = F.relu(self.fc1(outer_prod_view))\n",
    "        logits = self.fc2(fc1_outputs)\n",
    "        return logits\n",
    "\n",
    "\n",
    "class ESMForSingleMutation_pos(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.esm1v, self.esm1v_alphabet = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "        self.classifier = nn.Linear(1280, 1)\n",
    "        self.const1 = torch.nn.Parameter(torch.ones((1, 1280)))\n",
    "        self.const2 = torch.nn.Parameter(-1 * torch.ones((1, 1280)))\n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):\n",
    "        outputs1 = self.esm1v.forward(token_ids1, repr_layers=[33])[\"representations\"][33]\n",
    "        outputs2 = self.esm1v.forward(token_ids2, repr_layers=[33])[\"representations\"][33]\n",
    "        outputs = self.const1 * outputs1[:, pos + 1, :] + self.const2 * outputs2[:, pos + 1, :]\n",
    "        logits = self.classifier(outputs)\n",
    "        return logits\n",
    "\n",
    "\n",
    "class ESMForSingleMutation_cls(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.esm1v, self.esm1v_alphabet = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "        self.classifier = nn.Linear(1280, 1)\n",
    "        self.const1 = torch.nn.Parameter(torch.ones((1, 1280)))\n",
    "        self.const2 = torch.nn.Parameter(-1 * torch.ones((1, 1280)))\n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):\n",
    "        outputs1 = self.esm1v.forward(token_ids1, repr_layers=[33])[\"representations\"][33]\n",
    "        outputs2 = self.esm1v.forward(token_ids2, repr_layers=[33])[\"representations\"][33]\n",
    "        outputs = self.const1 * outputs1[:, 0, :] + self.const2 * outputs2[:, 0, :]\n",
    "        logits = self.classifier(outputs.unsqueeze(0))\n",
    "        return logits\n",
    "\n",
    "\n",
    "class ESMForSingleMutation_pos_cat_cls(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.esm1v, self.esm1v_alphabet = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "        self.classifier = nn.Linear(1280 * 2, 1)\n",
    "        self.const1 = torch.nn.Parameter(torch.ones((1, 1280)))\n",
    "        self.const2 = torch.nn.Parameter(-1 * torch.ones((1, 1280)))\n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):\n",
    "        outputs1 = self.esm1v.forward(token_ids1, repr_layers=[33])[\"representations\"][33]\n",
    "        outputs2 = self.esm1v.forward(token_ids2, repr_layers=[33])[\"representations\"][33]\n",
    "        cls_out = self.const1 * outputs1[:, 0, :] + self.const2 * outputs2[:, 0, :]\n",
    "        pos_out = self.const1 * outputs1[:, pos + 1, :] + self.const2 * outputs2[:, pos + 1, :]\n",
    "        outputs = torch.cat([cls_out.unsqueeze(0), pos_out], axis=-1)\n",
    "        logits = self.classifier(outputs)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TAp3NQyaupxE"
   },
   "source": [
    "# Compute DeltaDDG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hardcoded seq from p53 use for testing\n",
    "seq = \"SVPSQKTYQGSYGFRLGFLHSGTAKSVTCTYSPALNKMFCQLAKTCPVQLWVDSTPPPGTRVRAMAIYKQSQHMTEVVRRCPHHERCSDSDGLAPPQHLIRVEGNLRVEYLDDRNTFRHSVVVPYEPPEVGSDCTTIHYNYMCNSSCMGGMNRRPILTIITLEDSSGNLLGRNSFEVRVCACPGRDRRTEEENL\"  # @param {type:\"string\"}\n",
    "mutation_code = \"Q9H\"  # @param {type:\"string\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get wildtype sequence, mutation position and mutated sequence\n",
    "wt_aa = mutation_code[0]\n",
    "mut_aa = mutation_code[-1]\n",
    "mut_pos = int(mutation_code[1:-1]) - 1\n",
    "\n",
    "wt = seq\n",
    "tt = list(seq)\n",
    "tt[mut_pos] = mut_aa\n",
    "mut = \"\".join(tt)\n",
    "\n",
    "model = torch.load(\"ESMForSingleMutation_cls\", map_location=torch.device(\"cpu\"))\n",
    "esm2_alphabet = model.esm1v_alphabet\n",
    "esm2batch_converter = esm2_alphabet.get_batch_converter()\n",
    "_, _, esm2_batch_tokens1 = esm2batch_converter([(\"\", wt[:1022])])\n",
    "_, _, esm2_batch_tokens2 = esm2batch_converter([(\"\", mut[:1022])])\n",
    "esm2_batch_tokens1 = esm2_batch_tokens1.cuda()\n",
    "esm2_batch_tokens2 = esm2_batch_tokens2.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "_1qjq3i-VCEz"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model ESMForSingleMutationPosOuter DDG prediction is -0.24246525764465332\n",
      "Model ESMForSingleMutationPosConcat DDG prediction is -0.17090177536010742\n",
      "Model ESMForSingleMutation_pos_cat_cls DDG prediction is -0.04388764873147011\n",
      "Model ESMForSingleMutation_pos DDG prediction is 0.22296708822250366\n",
      "Model ESMForSingleMutation_cls DDG prediction is -0.06753194332122803\n"
     ]
    }
   ],
   "source": [
    "res = []\n",
    "for model_name in model_names:\n",
    "    model = torch.load(model_name, map_location=torch.device(\"cpu\"))\n",
    "    model.eval()\n",
    "    model.cuda()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        res.append(\n",
    "            model(token_ids1=esm2_batch_tokens1, token_ids2=esm2_batch_tokens2, pos=torch.LongTensor([mut_pos]))\n",
    "            .cpu()\n",
    "            .numpy()\n",
    "        )\n",
    "    print(f\"Model {model_name} DDG prediction is {res[-1][0,0,0]}\")\n",
    "res = np.mean(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "4by6H7tRdcQO"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted DDG for the mutation Q9H is -0.06036390736699104\n"
     ]
    }
   ],
   "source": [
    "print(f\"Predicted DDG for the mutation {mutation_code} is {res}\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "name": "Protein Stability Assessment using Transformers",
   "provenance": [],
   "toc_visible": true
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}