DiffDock_SingleComplex.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/suneelbvs/DiffDock/blob/main/DiffDock_SingleComplex.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# DiffDock\n",
        "Dock a small molecules on to protein structures using DiffDock approach\n",
        "\n",
        "1.   This notebook allows you to run diffdock on single protein/ligands and also multiple proteins/ligands.\n",
        "\n",
        "2.   Colab basic version works fine with single simulations. \"Premium GPU\" (colab pro), and even then it may fail on large complexes.\n",
        "\n",
        "## References:\n",
        "\n",
        "[Research Article](https://arxiv.org/abs/2210.01776)\n",
        "\n",
        "[Github](https://github.com/gcorso/DiffDock)\n",
        "\n",
        "[Interactive Online tool by Simon Duerr](https://huggingface.co/spaces/simonduerr/diffdock)\n",
        "\n",
        "[Colab Notebook by Brian Naughton](https://colab.research.google.com/drive/1nvCyQkbO-TwXZKJ0RCShVEym1aFWxlkX). The current notebook revised from Brain's work/code.\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "s4ov_2_SdX8J"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Start with mapping Google Drive to Colab\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CE1D_307TPZc",
        "outputId": "13d3f3e9-1ed4-4c15-f40e-7714ceef6564"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Step 1**: Setup working directory named \"DiffDock_V2\" in your Google Drive and update directory path.\n",
        "\n",
        "Copy or move this colab notebook to the current directory."
      ],
      "metadata": {
        "id": "8wsKRnvhMqSE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## Enable this code inorder to create DiffDock_V2 directory\n",
        "## Pls ignore this step if you have already created one \n",
        "#%cd /content/drive/MyDrive\n",
        "#%mkdir DiffDock_V2\n",
        "#%cd DiffDock_V2\n",
        "#%ls"
      ],
      "metadata": {
        "id": "tMtmKxmNbC3N"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "If you have already created or would like to work on different directory; please update the path accordingly"
      ],
      "metadata": {
        "id": "Av8ynSKRb3zc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/drive/MyDrive/DiffDock_V2\n",
        "%ls"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qOUZ0HWiiGPg",
        "outputId": "d02afcd1-8adb-4801-ff48-62ef9b6a2787"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/drive/MyDrive/DiffDock_V2\n",
            "\u001b[0m\u001b[01;34mDiffDock\u001b[0m/  DiffDock_V2.ipynb  DiffDock_V3.ipynb  \u001b[01;34mprolif\u001b[0m/  Untitled0.ipynb\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZAH5NYgIWIlI"
      },
      "source": [
        "## Step 2: \n",
        "Install the dependencies for DiffDock "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dq9PwzHUWDlj"
      },
      "source": [
        "## Install prerequisites"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "FSbikKk6fXAu",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "80d44276-3f0f-4b4f-809f-a41dba0d0531"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting ipython-autotime\n",
            "  Downloading ipython_autotime-0.3.1-py2.py3-none-any.whl (6.8 kB)\n",
            "Requirement already satisfied: ipython in /usr/local/lib/python3.7/dist-packages (from ipython-autotime) (7.9.0)\n",
            "Requirement already satisfied: setuptools>=18.5 in /usr/local/lib/python3.7/dist-packages (from ipython->ipython-autotime) (57.4.0)\n",
            "Collecting jedi>=0.10\n",
            "  Downloading jedi-0.18.1-py2.py3-none-any.whl (1.6 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.6 MB 7.0 MB/s \n",
            "\u001b[?25hRequirement already satisfied: backcall in /usr/local/lib/python3.7/dist-packages (from ipython->ipython-autotime) (0.2.0)\n",
            "Requirement already satisfied: pexpect in /usr/local/lib/python3.7/dist-packages (from ipython->ipython-autotime) (4.8.0)\n",
            "Requirement already satisfied: pickleshare in /usr/local/lib/python3.7/dist-packages (from ipython->ipython-autotime) (0.7.5)\n",
            "Requirement already satisfied: pygments in /usr/local/lib/python3.7/dist-packages (from ipython->ipython-autotime) (2.6.1)\n",
            "Requirement already satisfied: traitlets>=4.2 in /usr/local/lib/python3.7/dist-packages (from ipython->ipython-autotime) (5.1.1)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.7/dist-packages (from ipython->ipython-autotime) (4.4.2)\n",
            "Requirement already satisfied: prompt-toolkit<2.1.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from ipython->ipython-autotime) (2.0.10)\n",
            "Requirement already satisfied: parso<0.9.0,>=0.8.0 in /usr/local/lib/python3.7/dist-packages (from jedi>=0.10->ipython->ipython-autotime) (0.8.3)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.7/dist-packages (from prompt-toolkit<2.1.0,>=2.0.0->ipython->ipython-autotime) (0.2.5)\n",
            "Requirement already satisfied: six>=1.9.0 in /usr/local/lib/python3.7/dist-packages (from prompt-toolkit<2.1.0,>=2.0.0->ipython->ipython-autotime) (1.15.0)\n",
            "Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.7/dist-packages (from pexpect->ipython->ipython-autotime) (0.7.0)\n",
            "Installing collected packages: jedi, ipython-autotime\n",
            "Successfully installed ipython-autotime-0.3.1 jedi-0.18.1\n",
            "time: 824 µs (started: 2022-10-24 01:35:20 +00:00)\n"
          ]
        }
      ],
      "source": [
        "!pip install ipython-autotime\n",
        "%load_ext autotime"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "hTQPGM8ptbRy",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2c409b05-705b-49b0-f041-950533be8324"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/drive/MyDrive/DiffDock_V2\n",
            "fatal: destination path 'DiffDock' already exists and is not an empty directory.\n",
            "/content/drive/MyDrive/DiffDock_V2/DiffDock\n",
            "HEAD is now at 0f9c419 improve README\n",
            "time: 24 s (started: 2022-10-24 01:35:20 +00:00)\n"
          ]
        }
      ],
      "source": [
        "%cd /content/drive/MyDrive/DiffDock_V2\n",
        "!git clone https://github.com/gcorso/DiffDock.git\n",
        "%cd /content/drive/MyDrive/DiffDock_V2/DiffDock\n",
        "!git checkout 0f9c419 # remove/update for more up to date code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "0vSCAwwrtdr8",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e9c697eb-463c-4e85-f925-f56c3a1f4b09"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[?25l\r\u001b[K     |█████                           | 10 kB 15.5 MB/s eta 0:00:01\r\u001b[K     |██████████                      | 20 kB 16.9 MB/s eta 0:00:01\r\u001b[K     |███████████████▏                | 30 kB 21.4 MB/s eta 0:00:01\r\u001b[K     |████████████████████▏           | 40 kB 6.7 MB/s eta 0:00:01\r\u001b[K     |█████████████████████████▏      | 51 kB 6.7 MB/s eta 0:00:01\r\u001b[K     |██████████████████████████████▎ | 61 kB 7.7 MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 65 kB 2.1 MB/s \n",
            "\u001b[?25h  Building wheel for pyg (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for pkgtools (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[K     |████████████████████████████████| 2.6 MB 7.3 MB/s \n",
            "\u001b[K     |████████████████████████████████| 36.8 MB 24 kB/s \n",
            "\u001b[K     |████████████████████████████████| 117 kB 7.1 MB/s \n",
            "\u001b[K     |████████████████████████████████| 878 kB 7.1 MB/s \n",
            "\u001b[K     |████████████████████████████████| 5.7 MB 5.6 MB/s \n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for nglview (PEP 517) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[K     |████████████████████████████████| 16.1 MB 7.7 MB/s \n",
            "\u001b[?25htime: 1min 19s (started: 2022-10-24 01:35:44 +00:00)\n"
          ]
        }
      ],
      "source": [
        "!pip install pyg==0.7.1 --quiet\n",
        "!pip install pyyaml==6.0 --quiet\n",
        "!pip install scipy==1.7.3 --quiet\n",
        "!pip install networkx==2.6.3 --quiet\n",
        "!pip install biopython==1.79 --quiet\n",
        "!pip install rdkit-pypi==2022.03.5 --quiet\n",
        "!pip install e3nn==0.5.0 --quiet\n",
        "!pip install spyrmsd==0.5.2 --quiet\n",
        "!pip install pandas==1.3.5 --quiet\n",
        "!pip install biopandas==0.4.1 --quiet\n",
        "!pip install torch==1.12.1+cu113 --quiet\n",
        "!pip install nglview --quiet\n",
        "!pip install -q nglview pytraj --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "_RZ9qNmHcLMX",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "1b9e9205-ad8b-48cd-e2e1-0377e428301e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[33mWARNING: Skipping torch-scatter as it is not installed.\u001b[0m\n",
            "\u001b[33mWARNING: Skipping torch-sparse as it is not installed.\u001b[0m\n",
            "\u001b[33mWARNING: Skipping torch-geometric as it is not installed.\u001b[0m\n",
            "\u001b[33mWARNING: Skipping torch-cluster as it is not installed.\u001b[0m\n",
            "\u001b[K     |████████████████████████████████| 7.9 MB 8.8 MB/s \n",
            "\u001b[K     |████████████████████████████████| 3.5 MB 8.4 MB/s \n",
            "\u001b[K     |████████████████████████████████| 2.4 MB 7.8 MB/s \n",
            "\u001b[?25h  Building wheel for torch-geometric (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "time: 29.3 s (started: 2022-10-24 01:37:04 +00:00)\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "\n",
        "try:\n",
        "    import torch_geometric\n",
        "except ModuleNotFoundError:\n",
        "    !pip uninstall torch-scatter torch-sparse torch-geometric torch-cluster  --y\n",
        "    !pip install torch-scatter -f https://data.pyg.org/whl/torch-{torch.__version__}.html --quiet\n",
        "    !pip install torch-sparse -f https://data.pyg.org/whl/torch-{torch.__version__}.html --quiet\n",
        "    !pip install torch-cluster -f https://data.pyg.org/whl/torch-{torch.__version__}.html --quiet\n",
        "    !pip install git+https://github.com/pyg-team/pytorch_geometric.git  --quiet # no version for some reason??"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VgNOaZ3YdKsg"
      },
      "source": [
        "### Download 2GB PDBBind dataset\n",
        "unnecessary for inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "e72N-MyCuN8K",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "970dc2c5-cc2f-42f5-f793-ee8f77a3fb64"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "time: 1.54 ms (started: 2022-10-24 01:37:33 +00:00)\n"
          ]
        }
      ],
      "source": [
        "#!test -d /content/DiffDock/data/PDBBind_processed || (wget https://zenodo.org/record/6034088/files/PDBBind.zip && unzip -q PDBBind.zip && mv PDBBind_processed /content/DiffDock/data/)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Upload Input files\n",
        "\n"
      ],
      "metadata": {
        "id": "OJMaJ969h8Ol"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Step 3:** \n",
        "\n",
        "1.   Upload protein and ligand file in data directory.\n",
        "2.   DiffDock supports .pdb file format for protein\n",
        "3.   and it supports, .sdf or .mol2, and SMILES format for ligand\n",
        "4.   For example, i have saved protein as 'protein.pdb' and ligand as 'ligand.sdf'. \n",
        "5.   Update the respective file names in esm embedding preparation and inference steps.\n",
        "6.   Alternatively, you can also provide SMILES as an input. For example, **--ligand \"COc(cc1)ccc1C#N\"** instead of *--ligand ligand.sdf*\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "5dUfLI9pXSNy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /data\n",
        "from google.colab import files\n",
        "uploaded = files.upload()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 125
        },
        "id": "8bNlEvK8PJmV",
        "outputId": "3341050d-8f66-4cc4-d4ed-c3ab1593eb56"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[Errno 2] No such file or directory: '/data'\n",
            "/content/drive/MyDrive/DiffDock_V2/DiffDock\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "     <input type=\"file\" id=\"files-fd489e7a-2cb5-4425-bd51-36927190f905\" name=\"files[]\" multiple disabled\n",
              "        style=\"border:none\" />\n",
              "     <output id=\"result-fd489e7a-2cb5-4425-bd51-36927190f905\">\n",
              "      Upload widget is only available when the cell has been executed in the\n",
              "      current browser session. Please rerun this cell to enable.\n",
              "      </output>\n",
              "      <script>// Copyright 2017 Google LLC\n",
              "//\n",
              "// Licensed under the Apache License, Version 2.0 (the \"License\");\n",
              "// you may not use this file except in compliance with the License.\n",
              "// You may obtain a copy of the License at\n",
              "//\n",
              "//      http://www.apache.org/licenses/LICENSE-2.0\n",
              "//\n",
              "// Unless required by applicable law or agreed to in writing, software\n",
              "// distributed under the License is distributed on an \"AS IS\" BASIS,\n",
              "// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
              "// See the License for the specific language governing permissions and\n",
              "// limitations under the License.\n",
              "\n",
              "/**\n",
              " * @fileoverview Helpers for google.colab Python module.\n",
              " */\n",
              "(function(scope) {\n",
              "function span(text, styleAttributes = {}) {\n",
              "  const element = document.createElement('span');\n",
              "  element.textContent = text;\n",
              "  for (const key of Object.keys(styleAttributes)) {\n",
              "    element.style[key] = styleAttributes[key];\n",
              "  }\n",
              "  return element;\n",
              "}\n",
              "\n",
              "// Max number of bytes which will be uploaded at a time.\n",
              "const MAX_PAYLOAD_SIZE = 100 * 1024;\n",
              "\n",
              "function _uploadFiles(inputId, outputId) {\n",
              "  const steps = uploadFilesStep(inputId, outputId);\n",
              "  const outputElement = document.getElementById(outputId);\n",
              "  // Cache steps on the outputElement to make it available for the next call\n",
              "  // to uploadFilesContinue from Python.\n",
              "  outputElement.steps = steps;\n",
              "\n",
              "  return _uploadFilesContinue(outputId);\n",
              "}\n",
              "\n",
              "// This is roughly an async generator (not supported in the browser yet),\n",
              "// where there are multiple asynchronous steps and the Python side is going\n",
              "// to poll for completion of each step.\n",
              "// This uses a Promise to block the python side on completion of each step,\n",
              "// then passes the result of the previous step as the input to the next step.\n",
              "function _uploadFilesContinue(outputId) {\n",
              "  const outputElement = document.getElementById(outputId);\n",
              "  const steps = outputElement.steps;\n",
              "\n",
              "  const next = steps.next(outputElement.lastPromiseValue);\n",
              "  return Promise.resolve(next.value.promise).then((value) => {\n",
              "    // Cache the last promise value to make it available to the next\n",
              "    // step of the generator.\n",
              "    outputElement.lastPromiseValue = value;\n",
              "    return next.value.response;\n",
              "  });\n",
              "}\n",
              "\n",
              "/**\n",
              " * Generator function which is called between each async step of the upload\n",
              " * process.\n",
              " * @param {string} inputId Element ID of the input file picker element.\n",
              " * @param {string} outputId Element ID of the output display.\n",
              " * @return {!Iterable<!Object>} Iterable of next steps.\n",
              " */\n",
              "function* uploadFilesStep(inputId, outputId) {\n",
              "  const inputElement = document.getElementById(inputId);\n",
              "  inputElement.disabled = false;\n",
              "\n",
              "  const outputElement = document.getElementById(outputId);\n",
              "  outputElement.innerHTML = '';\n",
              "\n",
              "  const pickedPromise = new Promise((resolve) => {\n",
              "    inputElement.addEventListener('change', (e) => {\n",
              "      resolve(e.target.files);\n",
              "    });\n",
              "  });\n",
              "\n",
              "  const cancel = document.createElement('button');\n",
              "  inputElement.parentElement.appendChild(cancel);\n",
              "  cancel.textContent = 'Cancel upload';\n",
              "  const cancelPromise = new Promise((resolve) => {\n",
              "    cancel.onclick = () => {\n",
              "      resolve(null);\n",
              "    };\n",
              "  });\n",
              "\n",
              "  // Wait for the user to pick the files.\n",
              "  const files = yield {\n",
              "    promise: Promise.race([pickedPromise, cancelPromise]),\n",
              "    response: {\n",
              "      action: 'starting',\n",
              "    }\n",
              "  };\n",
              "\n",
              "  cancel.remove();\n",
              "\n",
              "  // Disable the input element since further picks are not allowed.\n",
              "  inputElement.disabled = true;\n",
              "\n",
              "  if (!files) {\n",
              "    return {\n",
              "      response: {\n",
              "        action: 'complete',\n",
              "      }\n",
              "    };\n",
              "  }\n",
              "\n",
              "  for (const file of files) {\n",
              "    const li = document.createElement('li');\n",
              "    li.append(span(file.name, {fontWeight: 'bold'}));\n",
              "    li.append(span(\n",
              "        `(${file.type || 'n/a'}) - ${file.size} bytes, ` +\n",
              "        `last modified: ${\n",
              "            file.lastModifiedDate ? file.lastModifiedDate.toLocaleDateString() :\n",
              "                                    'n/a'} - `));\n",
              "    const percent = span('0% done');\n",
              "    li.appendChild(percent);\n",
              "\n",
              "    outputElement.appendChild(li);\n",
              "\n",
              "    const fileDataPromise = new Promise((resolve) => {\n",
              "      const reader = new FileReader();\n",
              "      reader.onload = (e) => {\n",
              "        resolve(e.target.result);\n",
              "      };\n",
              "      reader.readAsArrayBuffer(file);\n",
              "    });\n",
              "    // Wait for the data to be ready.\n",
              "    let fileData = yield {\n",
              "      promise: fileDataPromise,\n",
              "      response: {\n",
              "        action: 'continue',\n",
              "      }\n",
              "    };\n",
              "\n",
              "    // Use a chunked sending to avoid message size limits. See b/62115660.\n",
              "    let position = 0;\n",
              "    do {\n",
              "      const length = Math.min(fileData.byteLength - position, MAX_PAYLOAD_SIZE);\n",
              "      const chunk = new Uint8Array(fileData, position, length);\n",
              "      position += length;\n",
              "\n",
              "      const base64 = btoa(String.fromCharCode.apply(null, chunk));\n",
              "      yield {\n",
              "        response: {\n",
              "          action: 'append',\n",
              "          file: file.name,\n",
              "          data: base64,\n",
              "        },\n",
              "      };\n",
              "\n",
              "      let percentDone = fileData.byteLength === 0 ?\n",
              "          100 :\n",
              "          Math.round((position / fileData.byteLength) * 100);\n",
              "      percent.textContent = `${percentDone}% done`;\n",
              "\n",
              "    } while (position < fileData.byteLength);\n",
              "  }\n",
              "\n",
              "  // All done.\n",
              "  yield {\n",
              "    response: {\n",
              "      action: 'complete',\n",
              "    }\n",
              "  };\n",
              "}\n",
              "\n",
              "scope.google = scope.google || {};\n",
              "scope.google.colab = scope.google.colab || {};\n",
              "scope.google.colab._files = {\n",
              "  _uploadFiles,\n",
              "  _uploadFilesContinue,\n",
              "};\n",
              "})(self);\n",
              "</script> "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Saving ligand.sdf to ligand.sdf\n",
            "time: 29.5 s (started: 2022-10-24 01:37:33 +00:00)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "For demo files refer my [github profile](https://github.com/suneelbvs/DiffDock)"
      ],
      "metadata": {
        "id": "LA0e0STv3nNC"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wpIBFqW7eNC7"
      },
      "source": [
        "## Install ESM and prepare PDB file for ESM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "2lduAr-WWPyX",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0a391038-6f50-44f6-db19-ef3e61d21315"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/drive/MyDrive/DiffDock_V2/DiffDock\n",
            "fatal: destination path 'esm' already exists and is not an empty directory.\n",
            "/content/drive/MyDrive/DiffDock_V2/DiffDock/esm\n",
            "HEAD is now at f07aed6 fix fairscale inference example (#298)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Obtaining file:///content/drive/MyDrive/DiffDock_V2/DiffDock/esm\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n",
            "Installing collected packages: fair-esm\n",
            "  Running setup.py develop for fair-esm\n",
            "Successfully installed fair-esm\n",
            "/content/drive/MyDrive/DiffDock_V2/DiffDock\n",
            "time: 57.8 s (started: 2022-10-24 01:38:03 +00:00)\n"
          ]
        }
      ],
      "source": [
        "%cd /content/drive/MyDrive/DiffDock_V2/DiffDock\n",
        "!git clone https://github.com/facebookresearch/esm\n",
        "%cd /content/drive/MyDrive/DiffDock_V2/DiffDock/esm\n",
        "!git checkout f07aed6 # remove/update for more up to date code\n",
        "!sudo pip install -e .\n",
        "%cd /content/drive/MyDrive/DiffDock_V2/DiffDock"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "qbxUtRgHwoIx",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3a99f839-82ba-42af-b6ac-4b30eca65e2a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/drive/MyDrive/DiffDock_V2/DiffDock\n",
            "6w70.pdb              \u001b[0m\u001b[01;34mesm\u001b[0m/                                \u001b[01;34m__pycache__\u001b[0m/\n",
            "\u001b[01;34mbaselines\u001b[0m/            evaluate_confidence_calibration.py  README.md\n",
            "\u001b[01;34mconfidence\u001b[0m/           evaluate.py                         \u001b[01;34mresults\u001b[0m/\n",
            "\u001b[01;34mdata\u001b[0m/                 inference.py                        train.py\n",
            "\u001b[01;34mdatasets\u001b[0m/             LICENSE                             \u001b[01;34mutils\u001b[0m/\n",
            "diffdock_results.tar  ligand.sdf                          \u001b[01;34mvisualizations\u001b[0m/\n",
            "environment.yml       \u001b[01;34mmodels\u001b[0m/                             \u001b[01;34mworkdir\u001b[0m/\n",
            "100% 1/1 [00:00<00:00, 22.83it/s]\n",
            "time: 1.37 s (started: 2022-10-24 01:43:46 +00:00)\n"
          ]
        }
      ],
      "source": [
        "%cd /content/drive/MyDrive/DiffDock_V2/DiffDock\n",
        "!python datasets/esm_embedding_preparation.py --protein_path data/protein.pdb --out_file data/prepared_for_esm.fasta "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cKUU2PKDxTkX",
        "outputId": "8df17d7c-2444-4701-f48b-2c1e841a9bf9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/drive/MyDrive/DiffDock_V2/DiffDock\n",
            "env: HOME=esm/model_weights\n",
            "env: PYTHONPATH=$PYTHONPATH:/content/drive/MyDrive/DiffDock_V2/DiffDock/esm\n",
            "Read data/prepared_for_esm.fasta with 1 sequences\n",
            "Processing 1 of 1 batches (1 sequences)\n",
            "time: 25.1 s (started: 2022-10-24 01:43:51 +00:00)\n"
          ]
        }
      ],
      "source": [
        "%cd /content/drive/MyDrive/DiffDock_V2/DiffDock\n",
        "%env HOME=esm/model_weights\n",
        "%env PYTHONPATH=$PYTHONPATH:/content/drive/MyDrive/DiffDock_V2/DiffDock/esm\n",
        "!python /content/drive/MyDrive/DiffDock_V2/DiffDock/esm/scripts/extract.py esm2_t33_650M_UR50D data/prepared_for_esm.fasta data/esm2_output --repr_layers 33 --include per_tok"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7bJTXj5SeYzG"
      },
      "source": [
        "## Run DiffDock"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/drive/MyDrive/DiffDock_V2/DiffDock\n",
        "!python -m inference --protein_path data/protein.pdb --ligand data/ligand.sdf --out_dir results/singlecomplx --inference_steps 20 --samples_per_complex 40 --batch_size 10 --actual_steps 18 --no_final_step_noise\n",
        "#!mv 'index0_data-testing-6w70.pdb____data-testing-6w70_ligand.sdf' out #update the folder name, if you provide custom names for inputs\n",
        "#%cd ./out\n",
        "#%ls"
      ],
      "metadata": {
        "id": "LLMDimzDWWqy",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "16ae2773-1040-4d2b-bf59-abda564d6e49"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/drive/MyDrive/DiffDock_V2/DiffDock\n",
            "Reading molecules and generating local structures with RDKit\n",
            "100% 1/1 [00:00<00:00, 17.53it/s]\n",
            "Reading language model embeddings.\n",
            "Generating graphs for ligands and proteins\n",
            "loading complexes: 100% 1/1 [00:00<00:00,  3.80it/s]\n",
            "loading data from memory:  data/cache_torsion/limit0_INDEX_maxLigSizeNone_H0_recRad15.0_recMax24_esmEmbeddings1577801764/heterographs.pkl\n",
            "Number of complexes:  1\n",
            "radius protein: mean 23.78305435180664, std 0.0, max 23.78305435180664\n",
            "radius molecule: mean 4.922807216644287, std 0.0, max 4.922807216644287\n",
            "distance protein-mol: mean 75.22794342041016, std 0.0, max 75.22794342041016\n",
            "rmsd matching: mean 0.0, std 0.0, max 0\n",
            "HAPPENING | confidence model uses different type of graphs than the score model. Loading (or creating if not existing) the data for the confidence model now.\n",
            "Reading molecules and generating local structures with RDKit\n",
            "100% 1/1 [00:00<00:00, 14.73it/s]\n",
            "Reading language model embeddings.\n",
            "Generating graphs for ligands and proteins\n",
            "loading complexes: 100% 1/1 [00:00<00:00,  1.90it/s]\n",
            "loading data from memory:  data/cache_torsion_allatoms/limit0_INDEX_maxLigSizeNone_H0_recRad15.0_recMax24_atomRad5_atomMax8_esmEmbeddings1577801764/heterographs.pkl\n",
            "Number of complexes:  1\n",
            "radius protein: mean 23.78305435180664, std 0.0, max 23.78305435180664\n",
            "radius molecule: mean 4.969401836395264, std 0.0, max 4.969401836395264\n",
            "distance protein-mol: mean 75.30010223388672, std 0.0, max 75.30010223388672\n",
            "rmsd matching: mean 0.0, std 0.0, max 0\n",
            "common t schedule [1.   0.95 0.9  0.85 0.8  0.75 0.7  0.65 0.6  0.55 0.5  0.45 0.4  0.35\n",
            " 0.3  0.25 0.2  0.15 0.1  0.05]\n",
            "Size of test dataset:  1\n",
            "1it [22:43, 1363.02s/it]\n",
            "Failed for 0 complexes\n",
            "Skipped 0 complexes\n",
            "Results are in results/singlecomplx\n",
            "time: 23min 23s (started: 2022-10-24 01:44:17 +00:00)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rVn2g1AcdPaM"
      },
      "source": [
        "## Download results"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd ./results/singlecomplx\n",
        "!mv 'index0_data-protein.pdb____data-ligand.sdf' out\n",
        "#%cp ./data/*.*pdb\n",
        "%cd ./out\n",
        "%ls"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qLpm-a_trP1Z",
        "outputId": "d0894444-2a0b-4899-c849-8b86a0c7153b"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/drive/MyDrive/DiffDock_V2/DiffDock/results/singlecomplx\n",
            "/content/drive/MyDrive/DiffDock_V2/DiffDock/results/singlecomplx/out\n",
            "rank10_confidence0.01.sdf   rank29_confidence-1.25.sdf\n",
            "rank11_confidence0.00.sdf   rank2_confidence0.38.sdf\n",
            "rank12_confidence-0.08.sdf  rank30_confidence-1.33.sdf\n",
            "rank13_confidence-0.14.sdf  rank31_confidence-1.33.sdf\n",
            "rank14_confidence-0.21.sdf  rank32_confidence-1.47.sdf\n",
            "rank15_confidence-0.24.sdf  rank33_confidence-1.53.sdf\n",
            "rank16_confidence-0.26.sdf  rank34_confidence-1.64.sdf\n",
            "rank17_confidence-0.27.sdf  rank35_confidence-1.93.sdf\n",
            "rank18_confidence-0.31.sdf  rank36_confidence-2.18.sdf\n",
            "rank19_confidence-0.35.sdf  rank37_confidence-2.87.sdf\n",
            "rank1_confidence0.44.sdf    rank38_confidence-2.96.sdf\n",
            "rank1.sdf                   rank39_confidence-3.29.sdf\n",
            "rank20_confidence-0.36.sdf  rank3_confidence0.38.sdf\n",
            "rank21_confidence-0.44.sdf  rank40_confidence-3.41.sdf\n",
            "rank22_confidence-0.45.sdf  rank4_confidence0.37.sdf\n",
            "rank23_confidence-0.50.sdf  rank5_confidence0.36.sdf\n",
            "rank24_confidence-0.52.sdf  rank6_confidence0.34.sdf\n",
            "rank25_confidence-0.60.sdf  rank7_confidence0.30.sdf\n",
            "rank26_confidence-0.63.sdf  rank8_confidence0.06.sdf\n",
            "rank27_confidence-0.77.sdf  rank9_confidence0.06.sdf\n",
            "rank28_confidence-0.84.sdf\n",
            "time: 283 ms (started: 2022-10-24 02:12:31 +00:00)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import output\n",
        "output.enable_custom_widget_manager()"
      ],
      "metadata": {
        "id": "13Rg9T7Sa-1U",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "55106ede-efac-4fc5-f5d2-4beb26be68f1"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "time: 4 ms (started: 2022-10-24 01:40:34 +00:00)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%ls"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vdPxZZGFxRXi",
        "outputId": "57f5eef3-e899-4e60-9868-e66e94189ed0"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "rank10_confidence0.01.sdf   rank29_confidence-1.25.sdf\n",
            "rank11_confidence0.00.sdf   rank2_confidence0.38.sdf\n",
            "rank12_confidence-0.08.sdf  rank30_confidence-1.33.sdf\n",
            "rank13_confidence-0.14.sdf  rank31_confidence-1.33.sdf\n",
            "rank14_confidence-0.21.sdf  rank32_confidence-1.47.sdf\n",
            "rank15_confidence-0.24.sdf  rank33_confidence-1.53.sdf\n",
            "rank16_confidence-0.26.sdf  rank34_confidence-1.64.sdf\n",
            "rank17_confidence-0.27.sdf  rank35_confidence-1.93.sdf\n",
            "rank18_confidence-0.31.sdf  rank36_confidence-2.18.sdf\n",
            "rank19_confidence-0.35.sdf  rank37_confidence-2.87.sdf\n",
            "rank1_confidence0.44.sdf    rank38_confidence-2.96.sdf\n",
            "rank1.sdf                   rank39_confidence-3.29.sdf\n",
            "rank20_confidence-0.36.sdf  rank3_confidence0.38.sdf\n",
            "rank21_confidence-0.44.sdf  rank40_confidence-3.41.sdf\n",
            "rank22_confidence-0.45.sdf  rank4_confidence0.37.sdf\n",
            "rank23_confidence-0.50.sdf  rank5_confidence0.36.sdf\n",
            "rank24_confidence-0.52.sdf  rank6_confidence0.34.sdf\n",
            "rank25_confidence-0.60.sdf  rank7_confidence0.30.sdf\n",
            "rank26_confidence-0.63.sdf  rank8_confidence0.06.sdf\n",
            "rank27_confidence-0.77.sdf  rank9_confidence0.06.sdf\n",
            "rank28_confidence-0.84.sdf\n",
            "time: 136 ms (started: 2022-10-24 02:22:18 +00:00)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "# Work In Progress: Analysis Part\n",
        "**bold text**\n",
        "\n"
      ],
      "metadata": {
        "id": "X1_mWbA322rx"
      }
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}