stab_ESM_IF.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "L4",
      "toc_visible": true,
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/deepsatflow/adaptyv-bio-pdc/blob/main/stab_ESM_IF.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# <b><font color='#009e74'>Absolute folding stabiity prediction via generative models</font></b>\n",
        "Colaboratory implementation of : **Cagiada M., Ovchinnikov S. & Lindorff-Larsen K.** - [Predicting absolute protein folding stability using generative models\n",
        "](https://doi.org/10.1038/s41467-023-39909-0). Source code is available on the [Github](https://github.com/KULL-Centre/_2024_cagiada_stability) page of the project.\n",
        "\n",
        "Prediction of the absolute stability for a target folding given a protein sequence.\n",
        "\n",
        "The Colab uses ESM-IF and the following measure:\n",
        "\n",
        "$\\Delta G_{f-u} = \\sum_{i}^{N} \\mathscr{L}_i^{WT}$\n",
        "\n",
        "where $\\mathscr{L}_i^{WT}$ is the amino acid likelihood extracted from ESM-IF for the wild-type amino acid at position i, to evaluate the absolute stability ($\\Delta G_{f-u}$) for a specific protein folding.\n",
        "\n"
      ],
      "metadata": {
        "id": "AcjxZ9mXikrK"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Additional notes:\n",
        "\n",
        "- Run this notebook  <b><font color='#d55c00'>preferabilly</font></b> in a Colab GPU session (go to page menu: `Runtime`->  `Change runtime type` -> select `GPU` and confirm.\n",
        "\n",
        "- Cells labelled <b><font color='#f0e422'>PRELIMINARY OPERATIONS </font></b>  have to be run <b><font color='#d55c00'>ONCE</font></b>  at the start and skipped for new predictions.\n",
        "\n",
        "- Multiple predictions can be run in a single session, but only <b><font color='#d55c00'>ONE</font></b> pdb at a time will be processed by the notebook.\n",
        "\n",
        "- A <b><font color='#d55c00'>new run</font></b> can be perform input direcly the new structure in the pdb upload cell and run the prediction cell again (you don't need to run the <b><font color='#f0e422'>PRELIMINARY OPERATIONS </font></b> again.\n",
        "\n",
        "- If you wish to download the predictions run the <b><font color='#56b4e9'>DOWNLOAD RESULTS </font></b> and <b>ALL</b>\n",
        " the predictions made in the session will be dowloaded.\n",
        "\n",
        "- To evaluate absolute stability the model uses an input structure in PDB format, usually in the form of an [AlphaFold2](https://www.nature.com/articles/s41586-021-03819-2) predicted structure, if you have an experimentally resolved structure please consider generating an AF2 prediction using your structure as a template.\n",
        "\n",
        "- The absolute stability can only be predicted for <b><font color='#d55c00'>one</font></b> chain at a time, but the input structure protein can be either a single chain or a multi-chain complex (the remaining chains would not be taken into account during the prediction).\n",
        "\n",
        "- The output of a prediction is displayed below the prediction cell and saved in an output file which includes the absolute stability (as a sum of likelihoods and in kcal/mol) and the contribution to the total stability for each residue (as a likelihood).\n",
        "\n",
        "- An **alternative sequence** can be used in the input instead of the sequence extracted from the PDB file, **HOWEVER** it must be the same length as the original sequence.\n",
        "****"
      ],
      "metadata": {
        "id": "a4dANSUAinaS"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "40_BsFRU0VEO",
        "outputId": "8b02cfea-e20a-4dc1-f225-478615379a70",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "installing libs...\n"
          ]
        }
      ],
      "source": [
        "\n",
        "#@title <b><font color='#f0e422'>PRELIMINARY OPERATIONS:</font> Install dependencies\n",
        "\n",
        "#@markdown Run the cell to install all the extra necessaries packages <b>(~1min 30s)</b>, including:\n",
        "#@markdown - ESM-IF (library and parameters)\n",
        "#@markdown - Torch libraries: torch-scatter,-sparse,-cluster,spline-conv,-geometric\n",
        "#@markdown - Python libraries: biopython, biotite\n",
        "\n",
        "%%time\n",
        "import os,time,subprocess,re,sys,shutil\n",
        "from google.colab import files\n",
        "import torch\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "def format_pytorch_version(version):\n",
        "  return version.split('+')[0]\n",
        "\n",
        "def format_cuda_version(version):\n",
        "  return 'cu' + version.replace('.', '')\n",
        "\n",
        "TORCH_version = torch.__version__\n",
        "TORCH = format_pytorch_version(TORCH_version)\n",
        "CUDA_version = torch.version.cuda\n",
        "CUDA = format_cuda_version(CUDA_version)\n",
        "\n",
        "IF_model_name = \"esm_if1_gvp4_t16_142M_UR50.pt\"\n",
        "\n",
        "if not os.path.isfile(IF_model_name):\n",
        "  # download esmfold params\n",
        "  os.system(\"apt-get install aria2 -qq\")\n",
        "  os.system(f\"aria2c -x 16 https://sid.erda.dk/share_redirect/eIZVVNEd8B --out={IF_model_name} &\")\n",
        "\n",
        "  if not os.path.isfile(\"finished_install\"):\n",
        "    # install libs\n",
        "    print(\"installing libs...\")\n",
        "\n",
        "    os.system(f\"pip install torch-scatter -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html\")\n",
        "    os.system(f\"pip install torch-sparse -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html\")\n",
        "    os.system(f\"pip install torch-cluster -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html\")\n",
        "    os.system(f\"pip install torch-spline-conv -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html\")\n",
        "    os.system(f\"pip install torch-geometric\")\n",
        "    os.system(f\"pip install biopython\")\n",
        "    os.system(f\"pip install biotite\")\n",
        "\n",
        "    print(\"installing esmfold...\")\n",
        "    # install esmfold\n",
        "    os.system(f\"pip install git+https://github.com/matteo-cagiada/esm.git\")\n",
        "    os.system(\"touch finished_install\")\n",
        "\n",
        "    #wait for Params to finish downloading...\n",
        "    while not os.path.isfile(IF_model_name):\n",
        "      time.sleep(5)\n",
        "    if os.path.isfile(f\"{IF_model_name}.aria2\"):\n",
        "      print(\"downloading params...\")\n",
        "    while os.path.isfile(f\"{IF_model_name}.aria2\"):\n",
        "      time.sleep(5)\n",
        "\n",
        "## Verify that pytorch-geometric is correctly installed\n",
        "\n",
        "import esm\n",
        "\n",
        "from esm.inverse_folding.util import load_structure, extract_coords_from_structure,CoordBatchConverter\n",
        "from esm.inverse_folding.multichain_util import extract_coords_from_complex,_concatenate_coords,load_complex_coords\n",
        "\n",
        "\n",
        "print(\"importing the model\")\n",
        "\n",
        "model, alphabet = esm.pretrained.load_model_and_alphabet(IF_model_name)\n",
        "model.eval().cuda().requires_grad_(False)\n",
        "\n",
        "print(\"--> Installations succeeded\")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title <b><font color='#f0e422'>PRELIMINARY OPERATIONS:</font> Load EXTRA functions\n",
        "#@markdown Run the cell to load the required functions\n",
        "\n",
        "def run_model(coords,sequence,model,cmplx=False,chain_target='A'):\n",
        "\n",
        "    device = next(model.parameters()).device\n",
        "\n",
        "    batch_converter = CoordBatchConverter(alphabet)\n",
        "    batch = [(coords, None, sequence)]\n",
        "    coords, confidence, strs, tokens, padding_mask = batch_converter(\n",
        "        batch, device=device)\n",
        "\n",
        "    prev_output_tokens = tokens[:, :-1].to(device)\n",
        "    target = tokens[:, 1:]\n",
        "    target_padding_mask = (target == alphabet.padding_idx)\n",
        "\n",
        "    logits, _ = model.forward(coords, padding_mask, confidence, prev_output_tokens)\n",
        "\n",
        "    logits_swapped=torch.swapaxes(logits,1,2)\n",
        "    token_probs = torch.softmax(logits_swapped, dim=-1)\n",
        "\n",
        "    return token_probs\n",
        "\n",
        "def score_variants(sequence,token_probs,alphabet):\n",
        "\n",
        "    aa_list=[]\n",
        "    wt_scores=[]\n",
        "    skip_pos=0\n",
        "\n",
        "    alphabetAA_L_D={'-':0,'_' :0,'A':1,'C':2,'D':3,'E':4,'F':5,'G':6,'H':7,'I':8,'K':9,'L':10,'M':11,'N':12,'P':13,'Q':14,'R':15,'S':16,'T':17,'V':18,'W':19,'Y':20}\n",
        "    alphabetAA_D_L={v: k for k, v in alphabetAA_L_D.items()}\n",
        "\n",
        "    for i,n in enumerate(sequence):\n",
        "      aa_list.append(n+str(i+1))\n",
        "      score_pos=[]\n",
        "      for j in range(1,21):\n",
        "          score_pos.append(masked_absolute(alphabetAA_D_L[j],i, token_probs, alphabet))\n",
        "          if n == alphabetAA_D_L[j]:\n",
        "            WT_score_pos=score_pos[-1]\n",
        "\n",
        "      wt_scores.append(WT_score_pos)\n",
        "\n",
        "    return aa_list, wt_scores\n",
        "\n",
        "def masked_absolute(mut, idx, token_probs, alphabet):\n",
        "\n",
        "    mt_encoded = alphabet.get_idx(mut)\n",
        "\n",
        "    score = token_probs[0,idx, mt_encoded]\n",
        "    return score.item()"
      ],
      "metadata": {
        "id": "lWXN38psLY3u",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title <b><font color='#56b4e9'> DATA UPLOADING</font>\n",
        "#@markdown Fill in the fields and run the cell to set up the job name, import the structure, select the chain and upload an alternative sequence (not mandatory).\n",
        "jobname='lys_ecoly'#@param {type:\"string\"}\n",
        "\n",
        "#@markdown Choose between <b><font color='#d55c00'> ONE</font></b> of the possible input sources for the target pdb and <b><font color='#d55c00'>leave the other cells empty or unmarked</font></b>\n",
        "#@markdown - AlphaFold2 PDB (v4) via Uniprot ID:\n",
        "AF_ID =''#@param {type:\"string\"}\n",
        "#@markdown - Upload custom PDB\n",
        "AF_custom =False#@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "#@markdown Select target chain (default A)\n",
        "chain_id='' #@param {type:'string'}\n",
        "\n",
        "#@markdown Upload an alternative sequence for the structure (leave empty if not used)\n",
        "alternative_sequence='SESELEAKRQELKELMEKTEEKARELLAKGDVAGARAAYGAYLLKAAE' #@param {type:'string'}\n",
        "\n",
        "input_path = f\"/content/inputs\"\n",
        "if not os.path.exists(input_path):\n",
        "  os.mkdir(input_path)\n",
        "\n",
        "output_path = f\"/content/outputs\"\n",
        "if not os.path.exists(output_path):\n",
        "  os.mkdir(output_path)\n",
        "\n",
        "if AF_custom:\n",
        "  print('Upload PDB file:')\n",
        "  uploaded_AF = files.upload()\n",
        "  for fn in uploaded_AF.keys():\n",
        "    os.rename(fn, f\"/content/inputs/query_protein.pdb\")\n",
        "    output_name_pdb=fn\n",
        "    print('... PDB file correctly loaded')\n",
        "elif (AF_ID !='') and (len(AF_ID)>=6) :\n",
        "    subprocess.call(['curl','-s','-f',f'https://alphafold.ebi.ac.uk/files/AF-{AF_ID}-F1-model_v4.pdb','-o','/content/inputs/query_protein.pdb'])\n",
        "    output_name_pdb=f'AF-{AF_ID}-F1-model_v4.pdb'\n",
        "else:\n",
        "  sys.exit(f'ERROR: any PDB uploaded, please select one of the above inputs')\n",
        "\n",
        "structure = load_structure(f\"/content/inputs/query_protein.pdb\", chain_id)\n",
        "coords_structure, sequence_structure = extract_coords_from_structure(structure)\n",
        "\n",
        "if alternative_sequence != '':\n",
        "  if ' ' in alternative_sequence:\n",
        "    sys.exit ('!!!! Run interrupted: please check input sequence before proceeding space characters detected!!!!')\n",
        "\n",
        "  assert len(alternative_sequence) == len(sequence_structure), \"Alternative sequence length doesn't match pdb sequence length, run interrupted!\"\n",
        "\n",
        "  sequence_structure = alternative_sequence\n",
        "  print('... Alternative sequence loaded correctly')\n",
        "\n",
        "print('... Target sequence:', sequence_structure)\n",
        "#@markdown ****"
      ],
      "metadata": {
        "id": "fbmpNhrvbKil",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title <b><font color='#56b4e9'> MODEL RUN</font>\n",
        "#@markdown Run this cell to evaluate the ΔG for the selected structure and sequence\n",
        "\n",
        "#@markdown **N.B:** the ΔG value will be output in the scale of the chosen metric and also in kcal/mol (see the manuscript for how we converted the scale)\n",
        "a=0.10413378327743603 ## fitting param from the manuscript to convert IF score scale to kcal/mol\n",
        "b=0.6162549378400894 ## fitting param from the manuscript to convert IF score scale to kcal/mol\n",
        "\n",
        "prob_tokens = run_model(coords_structure,sequence_structure,model,chain_target=chain_id)\n",
        "aa_list, wt_scores = score_variants(sequence_structure,prob_tokens,alphabet)\n",
        "\n",
        "dg_IF= np.nansum(wt_scores)\n",
        "print('ΔG predicted (likelihoods sum): ',dg_IF)\n",
        "\n",
        "dg_kcalmol= a * dg_IF + b\n",
        "\n",
        "print('ΔG predicted (kcal/mol): ', dg_kcalmol)\n",
        "\n",
        "aa_list_export=aa_list+['dG_IF','dG_kcalmol']\n",
        "wt_scores_export=wt_scores+[dg_IF,dg_kcalmol]\n",
        "\n",
        "df_export=pd.DataFrame({'Residue':aa_list_export,'score':wt_scores_export})\n",
        "\n",
        "df_export.to_csv(f\"outputs/\"+f\"{jobname}_dG_pos_scores_and_total.csv\",sep=',')\n",
        "## move pdb to output folder\n",
        "try:\n",
        "  os.rename(f\"inputs/query_protein.pdb\",f\"outputs/{output_name_pdb}\")\n",
        "except:\n",
        "  print('!!!! Data not saved, please re-upload the structure by running the uploading cell')\n"
      ],
      "metadata": {
        "id": "ocOnty1TP9Ec",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title <b><font color='#56b4e9'>DOWNLOAD RESULTS </font></b>\n",
        "#@markdown **N.B:** This will download **ALL** the predictions produced during the current session as zip file\n",
        "os.system( \"zip -r {} {}\".format( f\"dG_runs.zip\" , f\"outputs\" ) )\n",
        "files.download(f\"dG_runs.zip\")\n"
      ],
      "metadata": {
        "id": "HdIPxqSE6zWi",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "<b><font color='#56b4e9'>EXTRA </font></b>\n",
        "\n",
        "\\\\\n",
        "\n",
        "**Output of the colab**\n",
        "\n",
        "When a prediction is complete, the output files generated with the run are stored in the `/content/outputs` folder. When the <b><font color='#56b4e9'>DOWNLOAD RESULTS </font></b> cell is executed, all files are downloaded at once.\n",
        "\n",
        "The output files stored for each run are\n",
        "- the pdb used as input\n",
        "- the fasta file of the query sequence\n",
        "- the prediction output file\n",
        "\n",
        "**Prediction file Format**\n",
        "\n",
        "The output csv file consists of two columns, where the wild-type amino acid probability for each position is reported, and then at the bottom both the $\\Delta G$ as the sum of the wild-type probabilities and in kcal/mol are reported.\n",
        "\\\\\n",
        "\n",
        ">OUTPUT FILE EXAMPLE:\n",
        "\n",
        ">For a target protein with 45 residues, the scores file should be formatted like this:\n",
        "\n",
        ">Residue  Likelihood\n",
        "\n",
        ">M1               0.4  \n",
        "A2                0.2\n",
        "D3                0.3  \n",
        "C5                0.9   \n",
        "..  \n",
        "..  \n",
        "Y45               0.3\n",
        "dG_IF       201\n",
        "dG_kcalmol  13\n",
        "\n",
        "\\\\\n",
        "\n",
        "**Known problems:**\n",
        "\n",
        "- Predictions on multi-domain proteins or proteins with complex folding kinetics show an absolute stability overestimated compared to the real one.\n",
        "\n",
        "- Predictions are limited to proteins with 1023 residues (Max protein size for the ESM-IF language model)\n",
        "\n",
        "\\\\\n",
        "\n",
        "**License:**\n",
        "\n",
        "The $ΔG$ predictor, and ESM-IF source code and parameters are licensed under the permissive Apache Licence, Version 2.0.\n",
        "\n",
        "\\\\\n",
        "\n",
        "**Bugs:**\n",
        "\n",
        "For any bugs please report the issue on the project [Github](https://github.com/KULL-Centre/_2024_cagiada_stability) or contact one of the listed authors in the connected [manuscript](https://doi.org/).\n",
        "\n",
        "\\\\\n",
        "\n",
        "**Citing this work:**\n",
        "\n",
        "If you use our model please cite:\n",
        "\n",
        "citation here.\n",
        "\n",
        "```\n",
        "\n",
        "citation here\n",
        "\n",
        "```\n"
      ],
      "metadata": {
        "id": "0NULd4ycW4dw"
      }
    }
  ]
}