DiffDock.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZAH5NYgIWIlI"
      },
      "source": [
        "# DiffDock\n",
        "Dock a PDB files and a SMILES with [DiffDock](https://github.com/gcorso/DiffDock).\n",
        "\n",
        "Select Runtime / Run all to run an example PDB file and SMILES.\n",
        "\n",
        "v2 improvements:\n",
        "- works on proteins >1000 aas (with flag added to extract.py)\n",
        "- works on standard GPU (by reducing batch size from 10 to 6)\n",
        "- works with AlphaFold ids (AF-...) as well as PDB ids\n",
        "- works with comma-delimited PDB_ids and/or SMILES\n",
        "- runs smina to generate affinities (as DiffDock posed, or with smina minimization)\n",
        "- shows results in a py3DMol view\n",
        "\n",
        "v3:\n",
        "- fix py3dmol incompatibility\n",
        "- downgrade pytorch since the colab version takes forever to install pytorch_geometric (no binary so it has to compile)\n",
        "\n",
        "v4:\n",
        "- replace smina with the more accurate gnina\n",
        "- remove the annoying run twice thing now there is a binary\n",
        "\n",
        "v5:\n",
        "- add ProLIF fingerprint\n",
        "\n",
        "v6:\n",
        "- bugfixes due to ProLIF version changes. prolif is now locked at 2.0.1\n",
        "\n",
        "colab by [@btnaughton](https://twitter.com/btnaughton)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "M4xQfLgqWQsZ"
      },
      "outputs": [],
      "source": [
        "#@title PDB + SMILES input\n",
        "\n",
        "PDB_id = '' #@param {type:\"string\"}\n",
        "SMILES_or_pubchem_id = '' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Download a tar file containing all results?\n",
        "download_results = True #@param {type:\"boolean\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2Ssuid52VT-i"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import requests\n",
        "import time\n",
        "from random import random\n",
        "\n",
        "def download_pdb_file(pdb_id: str) -> str:\n",
        "    \"\"\"Download pdb file as a string from rcsb.org\"\"\"\n",
        "    PDB_DIR =\"/tmp/pdb/\"\n",
        "    os.makedirs(PDB_DIR, exist_ok=True)\n",
        "\n",
        "    # url or pdb_id\n",
        "    if pdb_id.startswith('http'):\n",
        "        url = pdb_id\n",
        "        filename = url.split('/')[-1]\n",
        "    elif pdb_id.endswith(\".pdb\"):\n",
        "        return pdb_id\n",
        "    else:\n",
        "        if pdb_id.startswith(\"AF\"):\n",
        "            url = f\"https://alphafold.ebi.ac.uk/files/{pdb_id}-model_v3.pdb\"\n",
        "        else:\n",
        "            url = f\"http://files.rcsb.org/view/{pdb_id}.pdb\"\n",
        "        filename = f'{pdb_id}.pdb'\n",
        "\n",
        "    cache_path = os.path.join(PDB_DIR, filename)\n",
        "    if os.path.exists(cache_path):\n",
        "        return cache_path\n",
        "\n",
        "    pdb_req = requests.get(url)\n",
        "    pdb_req.raise_for_status()\n",
        "    open(cache_path, 'w').write(pdb_req.text)\n",
        "    return cache_path\n",
        "\n",
        "def download_smiles_str(pubchem_id: str, retries:int = 2) -> str:\n",
        "    \"\"\"Given a pubchem id, get a smiles string\"\"\"\n",
        "    while True:\n",
        "        req = requests.get(f\"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/CID/{pubchem_id}/property/CanonicalSMILES/CSV\")\n",
        "        smiles_url_csv = req.text if req.status_code == 200 else None\n",
        "        if smiles_url_csv is not None:\n",
        "            break\n",
        "        if retries == 0:\n",
        "            return None\n",
        "        time.sleep(1+random())\n",
        "        retries -= 1\n",
        "\n",
        "    return smiles_url_csv.splitlines()[1].split(',')[1].strip('\"').strip(\"'\") if smiles_url_csv is not None else None"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BYTBennzXxb1"
      },
      "outputs": [],
      "source": [
        "if not PDB_id or not SMILES_or_pubchem_id:\n",
        "    PDB_id = \"6agt\"\n",
        "    SMILES_or_pubchem_id = \"COc(cc1)ccc1C#N\"\n",
        "    print(f\"No input supplied. Using example data: {PDB_id} and {SMILES_or_pubchem_id}\")\n",
        "\n",
        "# to run many PDB+smiles at once, fill in a list of PDB_files and smiles here...\n",
        "pdb_files = [download_pdb_file(_PDB_id) for _PDB_id in PDB_id.split(\",\")]\n",
        "smiless = [download_smiles_str(_SMILES_or_pubchem_id) if str(_SMILES_or_pubchem_id).isnumeric() else _SMILES_or_pubchem_id\n",
        "           for _SMILES_or_pubchem_id in SMILES_or_pubchem_id.split(',') ]\n",
        "\n",
        "with open(\"/tmp/input_protein_ligand.csv\", 'w') as out:\n",
        "    out.write(\"protein_path,ligand\\n\")\n",
        "    for pdb_file in pdb_files:\n",
        "        for smiles in smiless:\n",
        "            out.write(f\"{pdb_file},{smiles}\\n\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ptul2hJM2HmA"
      },
      "outputs": [],
      "source": [
        "# clear out old results if running multiple times -- hopefully they have been downloaded already\n",
        "!rm -rf /content/DiffDock/results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dq9PwzHUWDlj"
      },
      "source": [
        "## Install prerequisites"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FSbikKk6fXAu"
      },
      "outputs": [],
      "source": [
        "!pip install ipython-autotime --quiet\n",
        "%load_ext autotime"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hTQPGM8ptbRy"
      },
      "outputs": [],
      "source": [
        "if not os.path.exists(\"/content/DiffDock\"):\n",
        "    %cd /content\n",
        "    !git clone https://github.com/gcorso/DiffDock.git\n",
        "    %cd /content/DiffDock\n",
        "    !git checkout a6c5275 # remove/update for more up to date code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0vSCAwwrtdr8"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import biopandas\n",
        "except:\n",
        "    !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.5.3 --quiet\n",
        "    !pip install biopandas==0.4.1 --quiet\n",
        "    # not diffdock-specific:\n",
        "    !pip install prolif==2.0.1 --quiet # no real version??\n",
        "    !pip install py3dmol==2.0.3 --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_RZ9qNmHcLMX"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "print(torch.__version__)\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 # @ 15573f4674b2a37b1b9adc967df69ef6eee573ea"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wpIBFqW7eNC7"
      },
      "source": [
        "## Install ESM and prepare PDB file for ESM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2lduAr-WWPyX"
      },
      "outputs": [],
      "source": [
        "if not os.path.exists(\"/content/DiffDock/esm\"):\n",
        "    %cd /content/DiffDock\n",
        "    !git clone https://github.com/facebookresearch/esm\n",
        "    %cd /content/DiffDock/esm\n",
        "    !git checkout ca8a710 # remove/update for more up to date code\n",
        "    !sudo pip install -e .\n",
        "    %cd /content/DiffDock"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qbxUtRgHwoIx"
      },
      "outputs": [],
      "source": [
        "%cd /content/DiffDock\n",
        "!python datasets/esm_embedding_preparation.py --protein_ligand_csv /tmp/input_protein_ligand.csv --out_file data/prepared_for_esm.fasta"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cKUU2PKDxTkX"
      },
      "outputs": [],
      "source": [
        "%cd /content/DiffDock\n",
        "%env HOME=esm/model_weights\n",
        "%env PYTHONPATH=$PYTHONPATH:/content/DiffDock/esm\n",
        "!python /content/DiffDock/esm/scripts/extract.py esm2_t33_650M_UR50D data/prepared_for_esm.fasta data/esm2_output --repr_layers 33 --include per_tok --truncation_seq_length 30000"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7bJTXj5SeYzG"
      },
      "source": [
        "## Run DiffDock"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_D5gqUcabTAg"
      },
      "outputs": [],
      "source": [
        "%cd /content/DiffDock\n",
        "!python -m inference --protein_ligand_csv /tmp/input_protein_ligand.csv --out_dir results/user_predictions_small --inference_steps 20 --samples_per_complex 40 --batch_size 6"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zlPOKLIBsiPU"
      },
      "source": [
        "# Post-process and download results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VqWHbQdeXbe0"
      },
      "outputs": [],
      "source": [
        "%cd /content/DiffDock\n",
        "!wget https://sourceforge.net/projects/smina/files/smina.static/download -O smina && chmod +x smina\n",
        "!wget https://github.com/gnina/gnina/releases/download/v1.0.3/gnina -O gnina && chmod +x gnina"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "81xcfUI-tU0M"
      },
      "source": [
        "## Run gnina to estimate affinity"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XpjCe9IWjz5M"
      },
      "outputs": [],
      "source": [
        "import re\n",
        "import pandas as pd\n",
        "from glob import glob\n",
        "from shlex import quote\n",
        "from datetime import datetime\n",
        "from tqdm.auto import tqdm\n",
        "from google.colab import files\n",
        "\n",
        "%cd /content/DiffDock/results/user_predictions_small\n",
        "results_dirs = glob(\"./index*\")\n",
        "\n",
        "rows = []\n",
        "for results_dir in tqdm(results_dirs, desc=\"runs\"):\n",
        "    results_pdb_file = \"/tmp/pdb/\" + re.findall(\"tmp-pdb-(.+\\.pdb)\", results_dir)[0]\n",
        "    results_smiles = re.findall(\"pdb_+(.+)\", results_dir)[0]\n",
        "    results_sdfs = [os.path.join(results_dir, f) for f in os.listdir(results_dir) if \"confidence\" in f and f.endswith(\".sdf\")]\n",
        "\n",
        "    results_pdb_file_no_hetatms = f\"{results_pdb_file}_nohet.pdb\"\n",
        "    !grep -v \"^HETATM\" {results_pdb_file} > {results_pdb_file_no_hetatms}\n",
        "    !cp {results_pdb_file} .\n",
        "\n",
        "    for results_sdf in tqdm(results_sdfs, leave=False, desc=\"files\"):\n",
        "        confidence = re.findall(\"confidence([\\-\\.\\d]+)\\.sdf\", results_sdf)[0]\n",
        "\n",
        "        scored_stdout = !/content/DiffDock/gnina --score_only -r \"{results_pdb_file_no_hetatms}\" -l \"{results_sdf}\"\n",
        "        scored_affinity = re.findall(\"Affinity:\\s*([\\-\\.\\d+]+)\", '\\n'.join(scored_stdout))[0]\n",
        "        minimized_stdout = !/content/DiffDock/gnina --local_only --minimize -r \"{results_pdb_file_no_hetatms}\" -l \"{results_sdf}\" --autobox_ligand \"{results_sdf}\" --autobox_add 2\n",
        "        minimized_affinity = re.findall(\"Affinity:\\s*([\\-\\.\\d+]+)\", '\\n'.join(minimized_stdout))[0]\n",
        "\n",
        "        rows.append((results_pdb_file.split('/')[-1], results_smiles, float(confidence), float(scored_affinity), float(minimized_affinity), results_sdf))\n",
        "\n",
        "df_results = pd.DataFrame(rows, columns=[\"pdb_file\", \"smiles\", \"diffdock_confidence\", \"gnina_scored_affinity\", \"gnina_minimized_affinity\", \"sdf_file\"])\n",
        "df_results_tsv = \"df_diffdock_results.tsv\"\n",
        "df_results.to_csv(df_results_tsv, sep='\\t', index=None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F4ej2vQVtrE9"
      },
      "source": [
        "### Top hit by DiffDock confidence for any PDB/smiles combination"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gw6REL-8s6Ii"
      },
      "outputs": [],
      "source": [
        "top_hit = df_results.sort_values(\"diffdock_confidence\", ascending=False).iloc[0]\n",
        "df_results.sort_values(\"diffdock_confidence\", ascending=False).head(3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xnzZfgTXtbDO"
      },
      "source": [
        "## Show prolif fingerprint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pCFhotSSsw5o"
      },
      "outputs": [],
      "source": [
        "import prolif as plf\n",
        "from prolif.plotting.network import LigNetwork\n",
        "from rdkit import Chem\n",
        "\n",
        "# load protein\n",
        "prot = Chem.MolFromPDBFile(top_hit.pdb_file, removeHs=False)\n",
        "prot = plf.Molecule(prot)\n",
        "\n",
        "# load ligands\n",
        "lig_suppl = list(plf.sdf_supplier(top_hit.sdf_file))\n",
        "\n",
        "# generate fingerprint\n",
        "fp = plf.Fingerprint()\n",
        "fp.run_from_iterable(lig_suppl, prot)\n",
        "\n",
        "df_prolif = fp.to_dataframe()\n",
        "display(df_prolif)\n",
        "\n",
        "prolif_fingerprint_html = \"prolif_fingerprint.html\"\n",
        "net = LigNetwork.from_fingerprint(fp, lig_suppl[0], kind=\"frame\", frame=0, rotation=0)\n",
        "net.save(prolif_fingerprint_html)\n",
        "net.display()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EAuG_ORIsssn"
      },
      "source": [
        "## Create dataframe, tar file and download"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Wg4auY_VsqYh"
      },
      "outputs": [],
      "source": [
        "out_pdbs = ' '.join(set(df_results.pdb_file.apply(quote)))\n",
        "out_sdfs = ' '.join(df_results.sdf_file.apply(quote))\n",
        "\n",
        "if download_results:\n",
        "    tarname = f\"diffdock_{datetime.now().isoformat()[2:10].replace('-','')}\"\n",
        "    _ = !tar cvf {tarname}.tar --transform 's,^,{tarname}/,' --transform 's,\\./,,' {out_pdbs} {out_sdfs} {df_results_tsv} {prolif_fingerprint_html}\n",
        "\n",
        "    files.download(f\"{tarname}.tar\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PIgH3MSt_dWx"
      },
      "source": [
        "## Compare gnina affinities with DiffDock confidences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fqNY2KK-Xn3t"
      },
      "outputs": [],
      "source": [
        "import seaborn as sns\n",
        "import matplotlib.pyplot as plt\n",
        "from scipy.stats import linregress\n",
        "%config InlineBackend.figure_format='retina'\n",
        "\n",
        "for (pdb_file, smiles), df_group in df_results.groupby([\"pdb_file\", \"smiles\"]):\n",
        "    f, ax = plt.subplots(1, 2, figsize=(20,8))\n",
        "    sns.regplot(data=df_group, x=\"diffdock_confidence\", y=\"gnina_scored_affinity\", ax=ax[0]);\n",
        "    sns.regplot(data=df_group, x=\"diffdock_confidence\", y=\"gnina_minimized_affinity\", ax=ax[1]);\n",
        "\n",
        "    slope, intercept, r_value_scored, p_value, std_err = linregress(df_group[\"diffdock_confidence\"], df_group[\"gnina_scored_affinity\"])\n",
        "    slope, intercept, r_value_minimized, p_value, std_err = linregress(df_group[\"diffdock_confidence\"], df_group[\"gnina_minimized_affinity\"])\n",
        "    ax[0].set_title(f\"{pdb_file} {smiles[:30]} gnina scored r={r_value_scored:.3f}\");\n",
        "    ax[1].set_title(f\"{pdb_file} {smiles[:30]} gnina minimized r={r_value_minimized:.3f}\");"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GbeaGfBEAn3D"
      },
      "source": [
        "# Visualize top hit (highest confidence) in 3D"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qCwsUBpmYof5"
      },
      "outputs": [],
      "source": [
        "from IPython.display import HTML\n",
        "import py3Dmol\n",
        "\n",
        "resid_hover = \"\"\"\n",
        "function(atom,viewer) {\n",
        "    if(!atom.label) {\n",
        "        atom.label = viewer.addLabel(atom.chain+\" \"+atom.resn+\" \"+atom.resi,\n",
        "            {position: atom, backgroundColor: 'mintcream', fontColor:'black', fontSize:12});\n",
        "    }\n",
        "}\"\"\"\n",
        "unhover_func = \"\"\"\n",
        "function(atom,viewer) {\n",
        "    if(atom.label) {\n",
        "        viewer.removeLabel(atom.label);\n",
        "        delete atom.label;\n",
        "    }\n",
        "}\"\"\"\n",
        "\n",
        "view = py3Dmol.view(width=800, height=800)\n",
        "view.setCameraParameters({'fov': 35, 'z': 100});\n",
        "\n",
        "\n",
        "# add sdf\n",
        "view.addModel(open(top_hit.sdf_file).read(), \"sdf\")\n",
        "view.setStyle({\"model\": 0}, {'stick':{\"color\":\"#ff0000\"}})\n",
        "view.setViewStyle({\"model\": 0}, {'style':'outline','color':'black','width':0.1})\n",
        "view.zoomTo();\n",
        "\n",
        "# add pdb\n",
        "view.addModel(open(top_hit.pdb_file).read(), \"pdb\");\n",
        "view.setStyle({\"model\": 1}, {\"cartoon\":{\"color\":\"spectrum\"}})\n",
        "view.setStyle({\"model\": 1, \"hetflag\":True}, {'stick':{\"color\":\"spectrum\"}})\n",
        "\n",
        "model = view.getModel()\n",
        "model.setHoverable({}, True, resid_hover, unhover_func)\n",
        "\n",
        "view"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuClass": "premium",
      "gpuType": "A100",
      "machine_shape": "hm",
      "provenance": [],
      "toc_visible": true
    },
    "gpuClass": "premium",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}