Step_1_Predict_target_protein_structure.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "3bd5922bb247ea84",
      "metadata": {
        "id": "3bd5922bb247ea84"
      },
      "source": [
        "# Step 1. Predict target protein structure with ESMFold\n",
        "\n",
        "\n",
        "The example uses malic enzyme from Trypanosoma cruzi, which is the cause of Chagas disease. Ligand (PDB ID SZD) is one of the inhibitors in Scheme 1 in [this paper](https://pubs.acs.org/doi/epdf/10.1021/acsinfecdis.1c00231).\n",
        "\n",
        "Even though we do have the crystal structure of the protein-ligand complex (PDB ID 6W59), we will still use ESMFold to predict the protein structure for demo purpose.\n",
        "\n",
        "The sequence of the protein is as follows:\n",
        "\n",
        "```\n",
        ">6W59_SZD_protein_A\n",
        "AILTDRYINRGTAFTMEERQKLHILGRLPPVVETLEEQVARVYGQVKKYEKPINRYQHLVSVHSTNTTLY\n",
        "YATILAHLEEMLPIIYTPTVGEACMEYSHLFFRERGVYFNRLYKGQFRNIMRDAGYQKVEVVVITDGSRI\n",
        "LGLGDLGSNGIGISIGKCSLYVAGAGIDPRLIVPVILDVGTNNERYLQDKDYLGMREKRLGDEEFYELLD\n",
        "EFMEAASAEWPNAVIQFEDFSNNHCFDIMERYQKKYRCFNDDIQGTGAVIAAGFLNAIKLSGVSPLQQRI\n",
        "VVFGAGSAAVGVANNIAALAARMYKFPVQDLVKTFYLVDTKGLVTTTRGDQLAAHKKLLARTDVSAEDSA\n",
        "KLRTLEEIVRFVKPTTLLGLGGVGPAFTEEIVKMVMQNTERPIIFPLSNPTSKAEVTPENAYKWTNGAAI\n",
        "VASGSPFPPTTIGGKTFKPSQGNNLYVFPGVGLGCALAQPTHIPEELLLTASESLNLLTTEGDLREGRLY\n",
        "PPLEDIHNISANVATDVILEAQRMKIDNNKKLPRTRDELLAFVKKAMWKPVYSG\n",
        "```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "JHBgANrD6l0z",
      "metadata": {
        "id": "JHBgANrD6l0z"
      },
      "source": [
        "## 1.1 Set up environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "RE_Fiwpr-eLJ",
      "metadata": {
        "id": "RE_Fiwpr-eLJ"
      },
      "outputs": [],
      "source": [
        "!pip install pandas numpy matplotlib loguru py3dmol"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "TNoybLa8bAOE",
      "metadata": {
        "id": "TNoybLa8bAOE"
      },
      "outputs": [],
      "source": [
        "import os, shutil\n",
        "from google.colab import userdata\n",
        "import requests\n",
        "import py3Dmol\n",
        "import os"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9d84f6b8c77d89db",
      "metadata": {
        "id": "9d84f6b8c77d89db"
      },
      "source": [
        "## 1.2 Set up output directory"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "sUBUv5skrE_X",
      "metadata": {
        "id": "sUBUv5skrE_X"
      },
      "outputs": [],
      "source": [
        "def prepare_output_directory(output):\n",
        "    \"\"\"\n",
        "    Prepare the output directory\n",
        "    output: str, the output directory\n",
        "    return: None\n",
        "    \"\"\"\n",
        "    # overwrite the output directory\n",
        "    if os.path.exists(output):\n",
        "        shutil.rmtree(output)\n",
        "    os.makedirs(output)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "OtY-VQxEZVHO",
      "metadata": {
        "id": "OtY-VQxEZVHO"
      },
      "outputs": [],
      "source": [
        "output_dir = \"/content/output/esmfold_result\"\n",
        "prepare_output_directory(output_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "N-UF2GAZ69WT",
      "metadata": {
        "id": "N-UF2GAZ69WT"
      },
      "source": [
        "## 1.3 Set up URL and API key"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9a60df8cee8c3a42",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-04-25T14:09:39.816181Z",
          "start_time": "2024-04-25T14:09:39.812917Z"
        },
        "id": "9a60df8cee8c3a42"
      },
      "outputs": [],
      "source": [
        "invoke_url = \"https://health.api.nvidia.com/v1/biology/nvidia/esmfold\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ulQ6vdjlY1k6",
      "metadata": {
        "id": "ulQ6vdjlY1k6"
      },
      "outputs": [],
      "source": [
        "API_KEY = userdata.get('API_KEY')\n",
        "headers = {\n",
        "    \"Authorization\": f\"Bearer {API_KEY}\",\n",
        "    \"Accept\": \"application/json\",\n",
        "}\n",
        "\n",
        "print(headers)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9ddb918c7c3b3fc7",
      "metadata": {
        "id": "9ddb918c7c3b3fc7"
      },
      "source": [
        "## 1.4 Predict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "49492d5b3354227",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-04-25T14:09:43.098471Z",
          "start_time": "2024-04-25T14:09:43.094950Z"
        },
        "id": "49492d5b3354227"
      },
      "outputs": [],
      "source": [
        "# sequence of the protein\n",
        "sequence = \"AILTDRYINRGTAFTMEERQKLHILGRLPPVVETLEEQVARVYGQVKKYEKPINRYQHLVSVHSTNTTLYYATILAHLEEMLPIIYTPTVGEACMEYSHLFFRERGVYFNRLYKGQFRNIMRDAGYQKVEVVVITDGSRILGLGDLGSNGIGISIGKCSLYVAGAGIDPRLIVPVILDVGTNNERYLQDKDYLGMREKRLGDEEFYELLDEFMEAASAEWPNAVIQFEDFSNNHCFDIMERYQKKYRCFNDDIQGTGAVIAAGFLNAIKLSGVSPLQQRIVVFGAGSAAVGVANNIAALAARMYKFPVQDLVKTFYLVDTKGLVTTTRGDQLAAHKKLLARTDVSAEDSAKLRTLEEIVRFVKPTTLLGLGGVGPAFTEEIVKMVMQNTERPIIFPLSNPTSKAEVTPENAYKWTNGAAIVASGSPFPPTTIGGKTFKPSQGNNLYVFPGVGLGCALAQPTHIPEELLLTASESLNLLTTEGDLREGRLYPPLEDIHNISANVATDVILEAQRMKIDNNKKLPRTRDELLAFVKKAMWKPVYSG\"\n",
        "\n",
        "print(f\"length of protein: {len(sequence)} aa\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7RqxPtfkh1Hf",
      "metadata": {
        "id": "7RqxPtfkh1Hf"
      },
      "outputs": [],
      "source": [
        "payload = {\n",
        "  \"sequence\": sequence\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "600d272ab0fe45fa",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-04-25T14:10:12.886175Z",
          "start_time": "2024-04-25T14:09:43.855932Z"
        },
        "id": "600d272ab0fe45fa"
      },
      "outputs": [],
      "source": [
        "# initiate a session to re-use connections\n",
        "session = requests.Session()\n",
        "response = session.post(invoke_url, headers=headers, json=payload)\n",
        "response.raise_for_status()\n",
        "result = response.json()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9f68a5eea7ecdc22",
      "metadata": {
        "id": "9f68a5eea7ecdc22"
      },
      "source": [
        "## 1.5 Save the result"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cdccac646bec20c1",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-04-25T14:10:15.867607Z",
          "start_time": "2024-04-25T14:10:15.863096Z"
        },
        "id": "cdccac646bec20c1"
      },
      "outputs": [],
      "source": [
        "result.keys()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8704adf6b50de5e5",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-04-25T14:10:16.487982Z",
          "start_time": "2024-04-25T14:10:16.485611Z"
        },
        "id": "8704adf6b50de5e5"
      },
      "outputs": [],
      "source": [
        "# only 1 sequence is allowed each time\n",
        "assert len(result[\"pdbs\"]) == 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "efac8aba836c6ebc",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-04-25T14:10:17.007357Z",
          "start_time": "2024-04-25T14:10:17.004376Z"
        },
        "id": "efac8aba836c6ebc"
      },
      "outputs": [],
      "source": [
        "# Write PDB file\n",
        "fp = os.path.join(output_dir, f\"predicted_protein.pdb\")\n",
        "with open(fp, \"w\") as f:\n",
        "    f.write(result[\"pdbs\"][0]) # the str representation of the PDB file"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7thHIH6h9VId",
      "metadata": {
        "id": "7thHIH6h9VId"
      },
      "outputs": [],
      "source": [
        "# Download the predicted protein PDB file, which needs to be uploaded in Step 3 DiffDock inference\n",
        "from google.colab import files\n",
        "pred_pdb_file = '/content/output/esmfold_result/predicted_protein.pdb'\n",
        "files.download(pred_pdb_file)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "rNwT_CTX8aEj",
      "metadata": {
        "id": "rNwT_CTX8aEj"
      },
      "source": [
        "## 1.6 Visualize the result"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7f870cb68f0154a7",
      "metadata": {
        "id": "7f870cb68f0154a7"
      },
      "source": [
        "We can take a look at the predicted protein structure and its pLDDT score. The score is stored as B-factor in the returned results.\n",
        "\n",
        "To align with the [EBI/AlphaFold database](https://alphafold.ebi.ac.uk/), we will use the following colors for the pLDDT score\n",
        "\n",
        "![plddt](https://res.cloudinary.com/dpfqlyh21/image/upload/v1705026011/obsidian/izrfmiepbzpnzm2aoqwh.png)\n",
        "\n",
        "We display the predicted structure with these colors using py3Dmol.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "z9nW3hzrmpR5",
      "metadata": {
        "id": "z9nW3hzrmpR5"
      },
      "outputs": [],
      "source": [
        "def load_protein(pdb_file_path, width=800, height=600):\n",
        "\n",
        "    \"\"\"\n",
        "    Load a protein structure from a PDB file and display it using py3Dmol.\n",
        "    pdb_file_path: str, path to the PDB file\n",
        "    width: int, width of the viewer in pixels\n",
        "    height: int, height of the viewer in pixels\n",
        "    return: py3Dmol.view object\n",
        "    \"\"\"\n",
        "\n",
        "    with open(pdb_file_path) as ifile:\n",
        "        pdb_data = \"\".join([x for x in ifile])\n",
        "\n",
        "    view = py3Dmol.view(width=width, height=height)\n",
        "    view.addModelsAsFrames(pdb_data)\n",
        "\n",
        "    for line in pdb_data.split(\"\\n\"):\n",
        "        split = line.split()\n",
        "        if len(split) == 0 or split[0] != \"ATOM\":\n",
        "            continue\n",
        "        # Assuming the B-factor is at position 10 (you may need to adjust this based on your PDB format)\n",
        "        b_factor = float(split[10])\n",
        "        if b_factor > 90:\n",
        "            color = \"blue\"\n",
        "        elif 70 <= b_factor <= 90:\n",
        "            color = \"cyan\"\n",
        "        elif 50 <= b_factor < 70:\n",
        "            color = \"yellow\"\n",
        "        else:\n",
        "            color = \"orange\"\n",
        "\n",
        "        # Atom serial numbers typically start from 1, hence idx should be used directly\n",
        "        idx = int(split[1])\n",
        "\n",
        "        # Style should be set per atom id\n",
        "        view.setStyle({'model': -1, 'serial': idx}, {\"cartoon\": {'color': color}})\n",
        "    view.zoomTo()\n",
        "    return view"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e9abf0ca5d3e4e7f",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-04-25T14:10:32.426898Z",
          "start_time": "2024-04-25T14:10:32.227978Z"
        },
        "id": "e9abf0ca5d3e4e7f"
      },
      "outputs": [],
      "source": [
        "view = load_protein(\n",
        "    pdb_file_path = '/content/output/esmfold_result/predicted_protein.pdb',\n",
        "    width=800,\n",
        "    height=500\n",
        ")\n",
        "view.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "42583ed12170800e",
      "metadata": {
        "id": "42583ed12170800e"
      },
      "source": [
        "## 1.7 Visualize the superimposed structure\n",
        "To show how the predicted structure aligns with the crystal structure, we will use the crystal structure from (PDB ID 6W59)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bTC8EOheceEL",
      "metadata": {
        "id": "bTC8EOheceEL"
      },
      "source": [
        "In the next two cells below, install PyMOL in Colab and then test if installed successfully by following tutorial here # https://github.com/MooersLab/colabOpenSourcePyMOLpySnips?tab=readme-ov-file. PyMOL will be used to align the predicted and the ground truth pdb files."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "exnh2FgjpHvX",
      "metadata": {
        "id": "exnh2FgjpHvX"
      },
      "outputs": [],
      "source": [
        "from IPython.utils import io\n",
        "import tqdm.notebook\n",
        "import os\n",
        "\"\"\"The PyMOL installation is done inside two nested context managers. This approach\n",
        "was inspired by Dr. Christopher Schlick's (of the Phenix group at\n",
        "Lawrence Berkeley National Laboratory) method for installing cctbx\n",
        "in a Colab Notebook. He presented his work on September 1, 2021 at the IUCr\n",
        "Crystallographic Computing School. I adapted Chris's approach here. It replaces my first approach\n",
        "that requires seven steps. My approach was presentated at the SciPy2021 conference\n",
        "in July 2021 and published in the\n",
        "[proceedings](http://conference.scipy.org/proceedings/scipy2021/blaine_mooers.html).\n",
        "The new approach is easier for beginners to use. The old approach is easier to debug\n",
        "and could be used as a back-up approach.\n",
        "\n",
        "Thank you to Professor David Oppenheimer of the University of Florida for suggesting the use mamba and of Open Source PyMOL.\n",
        "\"\"\"\n",
        "total = 100\n",
        "with tqdm.notebook.tqdm(total=total) as pbar:\n",
        "    with io.capture_output() as captured:\n",
        "\n",
        "        !pip install -q condacolab\n",
        "        import condacolab\n",
        "        condacolab.install()\n",
        "        pbar.update(10)\n",
        "\n",
        "        import sys\n",
        "        sys.path.append('/usr/local/lib/python3.10/site-packages/')\n",
        "        pbar.update(20)\n",
        "\n",
        "        # Install PyMOL\n",
        "        %shell mamba install pymol-open-source --yes\n",
        "\n",
        "        pbar.update(100)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ZjoEunEtpKpH",
      "metadata": {
        "id": "ZjoEunEtpKpH"
      },
      "outputs": [],
      "source": [
        "# test if installation is successful\n",
        "# If the test failed, try to run the installation cell above again\n",
        "from pymol import cmd\n",
        "from IPython.display import Image\n",
        "cmd.fetch(\"1lw9\")\n",
        "cmd.orient(\"1lw9\")\n",
        "cmd.rotate(\"z\",\"270\")\n",
        "cmd.png(\"test.png\")\n",
        "Image(filename = \"test.png\", unconfined=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "zWk1UI-Mpw6X",
      "metadata": {
        "id": "zWk1UI-Mpw6X"
      },
      "outputs": [],
      "source": [
        "import py3Dmol\n",
        "from pymol import cmd\n",
        "import os\n",
        "\n",
        "def align_protein(pred_pdb_file, true_pdb_file, output_dir, pred_color = 'cyan', true_color = 'green', width=800, height=600):\n",
        "    \"\"\"\n",
        "    Align two protein structures and display them using py3Dmol.\n",
        "    pred_pdb_file: str, path to the PDB file of the predicted protein\n",
        "    true_pdb_file: str, path to the PDB file of the true protein\n",
        "    output_dir: str, path to the output directory to save the aligned PDB files\n",
        "    pred_color: str, color of the predicted protein in the viewer\n",
        "    true_color: str, color of the true protein in the viewer\n",
        "    width: int, width of the viewer in pixels\n",
        "    height: int, height of the viewer in pixels\n",
        "    return: py3Dmol.view object\n",
        "    \"\"\"\n",
        "\n",
        "    # always reinitialize!\n",
        "    cmd.reinitialize()\n",
        "\n",
        "    cmd.load(pred_pdb_file, 'predicted_protein')\n",
        "    cmd.load(true_pdb_file, 'true_protein')\n",
        "\n",
        "    # Align the structures on the backbone atoms\n",
        "    cmd.align('predicted_protein & backbone', 'true_protein & backbone', quiet=0)\n",
        "\n",
        "    # Save the aligned structures\n",
        "    pred_aligned_fp = os.path.join(output_dir, 'aligned_predicted_protein.pdb')\n",
        "    true_aligned_fp = os.path.join(output_dir, 'aligned_true_protein.pdb')\n",
        "    cmd.save(pred_aligned_fp, 'predicted_protein')\n",
        "    cmd.save(true_aligned_fp, 'true_protein')\n",
        "\n",
        "    # Read PDB files into strings (assuming they are in the current directory)\n",
        "    with open(pred_aligned_fp, 'r') as file:\n",
        "        pred_pdb_data = file.read()\n",
        "\n",
        "    with open(true_aligned_fp, 'r') as file:\n",
        "        true_pdb_data = file.read()\n",
        "\n",
        "    # Create a py3Dmol view\n",
        "    view = py3Dmol.view(width=width, height=height)\n",
        "\n",
        "    # Add the predicted protein model\n",
        "    view.addModel(pred_pdb_data, 'pdb')\n",
        "    # Set a specific color for the predicted protein, e.g., green\n",
        "    view.setStyle({'model': 0}, {'cartoon': {'color': pred_color}})\n",
        "\n",
        "    # Add the true protein model\n",
        "    view.addModel(true_pdb_data, 'pdb')\n",
        "    # Set a different specific color for the true protein, e.g., blue\n",
        "    view.setStyle({'model': 1}, {'cartoon': {'color': true_color}})\n",
        "\n",
        "    view.zoomTo()\n",
        "    return view"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "XV1OYku6qF1v",
      "metadata": {
        "id": "XV1OYku6qF1v"
      },
      "outputs": [],
      "source": [
        "# download the ground truth pdb \"6W59_SZD_protein.pd\"\n",
        "!wget https://raw.githubusercontent.com/hw-ju/bionemo_nim/refs/heads/main/6W59_SZD_protein.pdb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c8dffc70d4f6cf31",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2024-04-25T14:10:45.721882Z",
          "start_time": "2024-04-25T14:10:45.597791Z"
        },
        "id": "c8dffc70d4f6cf31"
      },
      "outputs": [],
      "source": [
        "view = align_protein(\n",
        "    pred_pdb_file = '/content/output/esmfold_result/predicted_protein.pdb', # prediction\n",
        "    true_pdb_file = '/content/6W59_SZD_protein.pdb', # true crystal structure\n",
        "    output_dir = '/content/output/esmfold_result', # output directory to store the aligned poses so that py3dmol can load back\n",
        "    pred_color = 'cyan', # color for the predicted structure\n",
        "    true_color = 'green', # color for the true structure\n",
        "    width=800,\n",
        "    height=500\n",
        ")\n",
        "view.show()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.11.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}