ESMFold.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/ESMFold.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#**ESMFold**\n",
        "for more details see: [Github](https://github.com/facebookresearch/esm/tree/main/esm), [Preprint](https://www.biorxiv.org/content/10.1101/2022.07.20.500902v1)\n",
        "\n",
        "#### **Tips and Instructions**\n",
        "- click the little ▶ play icon to the left of each cell below.\n",
        "- use \"/\" to specify chainbreaks, (eg. sequence=\"AAA/AAA\")\n",
        "- for homo-oligomeric predictions, set copies > 1\n",
        "- See [experimental notebook](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/ESMFold_advanced.ipynb) for more advanced options (like sampling).\n",
        "\n",
        "#### **Colab Limitations**\n",
        "- For short monomeric proteins under the length 400, consider using [ESMFold API](https://esmatlas.com/resources?action=fold) (no need for GPU, super fast!)\n",
        "- On Tesla T4 (typical free colab GPU), max total length ~ 900"
      ],
      "metadata": {
        "id": "POQBeXf2Xoxo"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "boFQEwsNQ4Qt"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "#@title install\n",
        "#@markdown install ESMFold, OpenFold and download Params (~2min 30s)\n",
        "\n",
        "import os, time\n",
        "if not os.path.isfile(\"esmfold.model\"):\n",
        "  # download esmfold params\n",
        "  os.system(\"apt-get install aria2 -qq\")\n",
        "  os.system(\"aria2c -q -x 16 https://colabfold.steineggerlab.workers.dev/esm/esmfold.model &\")\n",
        "\n",
        "  # install libs\n",
        "  os.system(\"pip install -q omegaconf pytorch_lightning biopython ml_collections einops py3Dmol\")\n",
        "  os.system(\"pip install -q git+https://github.com/NVIDIA/dllogger.git\")\n",
        "\n",
        "  # install openfold\n",
        "  commit = \"6908936b68ae89f67755240e2f588c09ec31d4c8\"\n",
        "  os.system(f\"pip install -q git+https://github.com/aqlaboratory/openfold.git@{commit}\")\n",
        "\n",
        "  # install esmfold\n",
        "  os.system(f\"pip install -q git+https://github.com/sokrypton/esm.git\")\n",
        "\n",
        "  # wait for Params to finish downloading...\n",
        "  if not os.path.isfile(\"esmfold.model\"):\n",
        "    # backup source!\n",
        "    os.system(\"aria2c -q -x 16 https://files.ipd.uw.edu/pub/esmfold/esmfold.model\")\n",
        "  else:\n",
        "    while os.path.isfile(\"esmfold.model.aria2\"):\n",
        "      time.sleep(5)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##run **ESMFold**\n",
        "%%time\n",
        "from string import ascii_uppercase, ascii_lowercase\n",
        "import hashlib, re, os\n",
        "import numpy as np\n",
        "from jax.tree_util import tree_map\n",
        "import matplotlib.pyplot as plt\n",
        "from scipy.special import softmax\n",
        "\n",
        "def parse_output(output):\n",
        "  pae = (output[\"aligned_confidence_probs\"][0] * np.arange(64)).mean(-1) * 31\n",
        "  plddt = output[\"plddt\"][0,:,1]\n",
        "  \n",
        "  bins = np.append(0,np.linspace(2.3125,21.6875,63))\n",
        "  sm_contacts = softmax(output[\"distogram_logits\"],-1)[0]\n",
        "  sm_contacts = sm_contacts[...,bins<8].sum(-1)\n",
        "  xyz = output[\"positions\"][-1,0,:,1]\n",
        "  mask = output[\"atom37_atom_exists\"][0,:,1] == 1\n",
        "  o = {\"pae\":pae[mask,:][:,mask],\n",
        "       \"plddt\":plddt[mask],\n",
        "       \"sm_contacts\":sm_contacts[mask,:][:,mask],\n",
        "       \"xyz\":xyz[mask]}\n",
        "  return o\n",
        "\n",
        "def get_hash(x): return hashlib.sha1(x.encode()).hexdigest()\n",
        "alphabet_list = list(ascii_uppercase+ascii_lowercase)\n",
        "\n",
        "jobname = \"test\" #@param {type:\"string\"}\n",
        "jobname = re.sub(r'\\W+', '', jobname)[:50]\n",
        "\n",
        "sequence = \"GWSTELEKHREELKEFLKKEGITNVEIRIDNGRLEVRVEGGTERLKRFLEELRQKLEKKGYTVDIKIE\" #@param {type:\"string\"}\n",
        "sequence = re.sub(\"[^A-Z:]\", \"\", sequence.replace(\"/\",\":\").upper())\n",
        "sequence = re.sub(\":+\",\":\",sequence)\n",
        "sequence = re.sub(\"^[:]+\",\"\",sequence)\n",
        "sequence = re.sub(\"[:]+$\",\"\",sequence)\n",
        "copies = 1 #@param {type:\"integer\"}\n",
        "if copies == \"\" or copies <= 0: copies = 1\n",
        "sequence = \":\".join([sequence] * copies)\n",
        "num_recycles = 3 #@param [\"0\", \"1\", \"2\", \"3\", \"6\", \"12\", \"24\"] {type:\"raw\"}\n",
        "chain_linker = 25 \n",
        "\n",
        "ID = jobname+\"_\"+get_hash(sequence)[:5]\n",
        "seqs = sequence.split(\":\")\n",
        "lengths = [len(s) for s in seqs]\n",
        "length = sum(lengths)\n",
        "print(\"length\",length)\n",
        "\n",
        "u_seqs = list(set(seqs))\n",
        "if len(seqs) == 1: mode = \"mono\"\n",
        "elif len(u_seqs) == 1: mode = \"homo\"\n",
        "else: mode = \"hetero\"\n",
        "\n",
        "if \"model\" not in dir():\n",
        "  import torch\n",
        "  model = torch.load(\"esmfold.model\")\n",
        "  model.eval().cuda().requires_grad_(False)\n",
        "\n",
        "# optimized for Tesla T4\n",
        "if length > 700:\n",
        "  model.set_chunk_size(64)\n",
        "else:\n",
        "  model.set_chunk_size(128)\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "output = model.infer(sequence,\n",
        "                     num_recycles=num_recycles,\n",
        "                     chain_linker=\"X\"*chain_linker,\n",
        "                     residue_index_offset=512)\n",
        "\n",
        "pdb_str = model.output_to_pdb(output)[0]\n",
        "output = tree_map(lambda x: x.cpu().numpy(), output)\n",
        "ptm = output[\"ptm\"][0]\n",
        "plddt = output[\"plddt\"][0,...,1].mean()\n",
        "O = parse_output(output)\n",
        "print(f'ptm: {ptm:.3f} plddt: {plddt:.3f}')\n",
        "os.system(f\"mkdir -p {ID}\")\n",
        "prefix = f\"{ID}/ptm{ptm:.3f}_r{num_recycles}_default\"\n",
        "np.savetxt(f\"{prefix}.pae.txt\",O[\"pae\"],\"%.3f\")\n",
        "with open(f\"{prefix}.pdb\",\"w\") as out:\n",
        "  out.write(pdb_str)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "CcyNpAvhTX6q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title display (optional) {run: \"auto\"}\n",
        "import py3Dmol\n",
        "pymol_color_list = [\"#33ff33\",\"#00ffff\",\"#ff33cc\",\"#ffff00\",\"#ff9999\",\"#e5e5e5\",\"#7f7fff\",\"#ff7f00\",\n",
        "                    \"#7fff7f\",\"#199999\",\"#ff007f\",\"#ffdd5e\",\"#8c3f99\",\"#b2b2b2\",\"#007fff\",\"#c4b200\",\n",
        "                    \"#8cb266\",\"#00bfbf\",\"#b27f7f\",\"#fcd1a5\",\"#ff7f7f\",\"#ffbfdd\",\"#7fffff\",\"#ffff7f\",\n",
        "                    \"#00ff7f\",\"#337fcc\",\"#d8337f\",\"#bfff3f\",\"#ff7fff\",\"#d8d8ff\",\"#3fffbf\",\"#b78c4c\",\n",
        "                    \"#339933\",\"#66b2b2\",\"#ba8c84\",\"#84bf00\",\"#b24c66\",\"#7f7f7f\",\"#3f3fa5\",\"#a5512b\"]\n",
        "\n",
        "def show_pdb(pdb_str, show_sidechains=False, show_mainchains=False,\n",
        "             color=\"pLDDT\", chains=None, vmin=50, vmax=90,\n",
        "             size=(800,480), hbondCutoff=4.0,\n",
        "             Ls=None,\n",
        "             animate=False):\n",
        "  \n",
        "  if chains is None:\n",
        "    chains = 1 if Ls is None else len(Ls)\n",
        "  view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js', width=size[0], height=size[1])\n",
        "  if animate:\n",
        "    view.addModelsAsFrames(pdb_str,'pdb',{'hbondCutoff':hbondCutoff})\n",
        "  else:\n",
        "    view.addModel(pdb_str,'pdb',{'hbondCutoff':hbondCutoff})\n",
        "  if color == \"pLDDT\":\n",
        "    view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':vmin,'max':vmax}}})\n",
        "  elif color == \"rainbow\":\n",
        "    view.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "  elif color == \"chain\":\n",
        "    for n,chain,color in zip(range(chains),alphabet_list,pymol_color_list):\n",
        "       view.setStyle({'chain':chain},{'cartoon': {'color':color}})\n",
        "  if show_sidechains:\n",
        "    BB = ['C','O','N']\n",
        "    view.addStyle({'and':[{'resn':[\"GLY\",\"PRO\"],'invert':True},{'atom':BB,'invert':True}]},\n",
        "                  {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"GLY\"},{'atom':'CA'}]},\n",
        "                  {'sphere':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
        "                  {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})  \n",
        "  if show_mainchains:\n",
        "    BB = ['C','O','N','CA']\n",
        "    view.addStyle({'atom':BB},{'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "  view.zoomTo()\n",
        "  if animate: view.animate()\n",
        "  return view\n",
        "\n",
        "color = \"confidence\" #@param [\"confidence\", \"rainbow\", \"chain\"]\n",
        "if color == \"confidence\": color = \"pLDDT\"\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "show_pdb(pdb_str, color=color,\n",
        "         show_sidechains=show_sidechains,\n",
        "         show_mainchains=show_mainchains,\n",
        "         Ls=lengths).show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "JM5ciSmeTZKO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title plot confidence (optional)\n",
        "\n",
        "dpi = 100 #@param {type:\"integer\"}\n",
        "\n",
        "def plot_ticks(Ls):\n",
        "  Ln = sum(Ls)\n",
        "  L_prev = 0\n",
        "  for L_i in Ls[:-1]:\n",
        "    L = L_prev + L_i\n",
        "    L_prev += L_i\n",
        "    plt.plot([0,Ln],[L,L],color=\"black\")\n",
        "    plt.plot([L,L],[0,Ln],color=\"black\")\n",
        "  ticks = np.cumsum([0]+Ls)\n",
        "  ticks = (ticks[1:] + ticks[:-1])/2\n",
        "  plt.yticks(ticks,alphabet_list[:len(ticks)])\n",
        "\n",
        "def plot_confidence(O, Ls=None, dpi=100):\n",
        "  if \"lm_contacts\" in O:\n",
        "    plt.figure(figsize=(20,4), dpi=dpi)\n",
        "    plt.subplot(1,4,1)\n",
        "  else:\n",
        "    plt.figure(figsize=(15,4), dpi=dpi)\n",
        "    plt.subplot(1,3,1)\n",
        "\n",
        "  plt.title('Predicted lDDT')\n",
        "  plt.plot(O[\"plddt\"])\n",
        "  if Ls is not None:\n",
        "    L_prev = 0\n",
        "    for L_i in Ls[:-1]:\n",
        "      L = L_prev + L_i\n",
        "      L_prev += L_i\n",
        "      plt.plot([L,L],[0,100],color=\"black\")\n",
        "  plt.xlim(0,O[\"plddt\"].shape[0])\n",
        "  plt.ylim(0,100)\n",
        "  plt.ylabel('plDDT')\n",
        "  plt.xlabel('position')\n",
        "  plt.subplot(1,4 if \"lm_contacts\" in O else 3,2)\n",
        "\n",
        "  plt.title('Predicted Aligned Error')\n",
        "  Ln = O[\"pae\"].shape[0]\n",
        "  plt.imshow(O[\"pae\"],cmap=\"bwr\",vmin=0,vmax=30,extent=(0, Ln, Ln, 0))\n",
        "  if Ls is not None and len(Ls) > 1: plot_ticks(Ls)\n",
        "  plt.colorbar()\n",
        "  plt.xlabel('Scored residue')\n",
        "  plt.ylabel('Aligned residue')\n",
        "\n",
        "  if \"lm_contacts\" in O:\n",
        "    plt.subplot(1,4,3)\n",
        "    plt.title(\"contacts from LM\")\n",
        "    plt.imshow(O[\"lm_contacts\"],cmap=\"Greys\",vmin=0,vmax=1,extent=(0, Ln, Ln, 0))\n",
        "    if Ls is not None and len(Ls) > 1: plot_ticks(Ls)\n",
        "    plt.subplot(1,4,4)\n",
        "  else:\n",
        "    plt.subplot(1,3,3)\n",
        "  plt.title(\"contacts from Structure Module\")\n",
        "  plt.imshow(O[\"sm_contacts\"],cmap=\"Greys\",vmin=0,vmax=1,extent=(0, Ln, Ln, 0))\n",
        "  if Ls is not None and len(Ls) > 1: plot_ticks(Ls)\n",
        "  return plt\n",
        "\n",
        "plot_confidence(O, Ls=lengths, dpi=dpi)\n",
        "plt.savefig(f'{prefix}.png',bbox_inches='tight')\n",
        "plt.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "HGFBl0QYYQpR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title download predictions\n",
        "from google.colab import files\n",
        "os.system(f\"zip {ID}.zip {ID}/*\")\n",
        "files.download(f'{ID}.zip')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "hkMp_ZwRYfAQ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}