GeoDock.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Graylab/GeoDock/blob/main/GeoDock.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7hJOcHgk85Oe"
      },
      "source": [
        "# **GeoDock**: Fast, Flexible Protein-Protein Docking\n",
        "\n",
        "Official notebook for [GeoDock](www.biorxiv.org/content/10.1101/2023.06.29.547134v1): Flexible Protein-Protein Docking with a Multi-Track Iterative Transformer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QhJ4lxFrwZ3l",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Hit `Runtime` -> `Run all`, then input docking partner1 and partner2\n",
        "\n",
        "from google.colab import files\n",
        "\n",
        "out_name = 'test' #@param {type:\"string\"}\n",
        "\n",
        "##@markdown Perform structural refinement with OpenMM\n",
        "#do_refine = False #@param {type:\"boolean\"}\n",
        "# refinement is not yet supported\n",
        "do_refine = False\n",
        "\n",
        "upload1 = files.upload()\n",
        "upload2 = files.upload()\n",
        "partner1 = list(upload1.keys())[0]\n",
        "partner2 = list(upload2.keys())[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lscm2F0chz1D",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Install dependencies (~1 min)\n",
        "%%time\n",
        "import sys\n",
        "import torch\n",
        "\n",
        "pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
        "version_str=\"\".join([\n",
        "    f\"py3{sys.version_info.minor}_cu\",\n",
        "    torch.version.cuda.replace(\".\",\"\"),\n",
        "    f\"_pyt{pyt_version_str}\"\n",
        "])\n",
        "\n",
        "def format_pytorch_version(version):\n",
        "  return version.split('+')[0]\n",
        "\n",
        "TORCH_version = torch.__version__\n",
        "TORCH = format_pytorch_version(TORCH_version)\n",
        "\n",
        "def format_cuda_version(version):\n",
        "  return 'cu' + version.replace('.', '')\n",
        "\n",
        "CUDA_version = torch.version.cuda\n",
        "CUDA = format_cuda_version(CUDA_version)\n",
        "\n",
        "!pip install geodock\n",
        "!pip install torch-geometric\n",
        "!pip install torch-scatter     -f https://pytorch-geometric.com/whl/torch-{TORCH}+{CUDA}.html\n",
        "!pip install torch-sparse      -f https://pytorch-geometric.com/whl/torch-{TORCH}+{CUDA}.html\n",
        "!pip install torch-cluster     -f https://pytorch-geometric.com/whl/torch-{TORCH}+{CUDA}.html\n",
        "!pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-{TORCH}+{CUDA}.html\n",
        "!pip install py3Dmol\n",
        "\n",
        "# download model weights\n",
        "!git clone https://github.com/Graylab/GeoDock.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CAY_ZAn_u7iY"
      },
      "outputs": [],
      "source": [
        "#@title Run Prediction\n",
        "from geodock.GeoDockRunner import GeoDockRunner\n",
        "torch.cuda.empty_cache()\n",
        "ckpt_file = \"/content/GeoDock/geodock/weights/dips_0.3.ckpt\"\n",
        "geodock = GeoDockRunner(ckpt_file=ckpt_file)\n",
        "pred = geodock.dock(\n",
        "    partner1=partner1,\n",
        "    partner2=partner2,\n",
        "    out_name=out_name,\n",
        "    do_refine=do_refine,\n",
        "    use_openmm=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "m1dXzzjPOWU-"
      },
      "outputs": [],
      "source": [
        "#@title Display 3D structure {run: \"auto\"}\n",
        "import py3Dmol\n",
        "\n",
        "color = \"chain\" #@param [\"chain\", \"lDDT\", \"rainbow\"]\n",
        "\n",
        "view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
        "view.addModel(open(out_name+'.pdb','r').read(),'pdb')\n",
        "if color == 'chain':\n",
        "  view.setStyle({'chain': 'A'},{'cartoon': {'color': 'green'}})\n",
        "  view.setStyle({'chain': 'B'},{'cartoon': {'color': 'blue'}})\n",
        "elif color == 'lDDT':\n",
        "  view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':50,'max':90}}})\n",
        "elif color == 'rainbow':\n",
        "  view.setStyle({'model': -1}, {\"cartoon\": {'color': 'spectrum'}})\n",
        "view.zoomTo()\n",
        "view.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kkdmk8Xx7o3X"
      },
      "outputs": [],
      "source": [
        "#@title Download results\n",
        "\n",
        "#@markdown Download the predicted pdb file. If download fails, results are also accessible from file explorer on the left panel of the notebook.\n",
        "\n",
        "from google.colab import files\n",
        "files.download(f\"{out_name}.pdb\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Instructions <a name=\"Instructions\"></a>\n",
        "**Quick start**\n",
        "1. Enter the name of your output.\n",
        "2. Press \"Runtime\" -> \"Run all\".\n",
        "3. Upload docking partner1 and partner2.\n",
        "\n",
        "**Known issues**\n",
        "* Google Colab assigns different types of GPUs with varying amount of memory. Some might not have enough memory to predict the large protein complexes."
      ],
      "metadata": {
        "id": "d3FLjIxYx2Ud"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "mount_file_id": "1qNN74eL3oay8qF3Uh-NFYwFV7OoBpCU6",
      "authorship_tag": "ABX9TyM6+SslON4X9H5xKqsv7V0n",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}