sample_notebook.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ZAH5NYgIWIlI"
      },
      "source": [
        "# DiffDock - Sample Notebook\n",
        "Dock a PDB files and a SMILES with [DiffDock](https://github.com/gcorso/DiffDock).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "M4xQfLgqWQsZ"
      },
      "outputs": [],
      "source": [
        "## PDB + SMILES input\n",
        "\n",
        "PDB_id = ''\n",
        "SMILES_or_pubchem_id = ''"
      ]
    },
    {
      "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",
        "    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)]\n",
        "smiless = [download_smiles_str(SMILES_or_pubchem_id) if str(SMILES_or_pubchem_id).isnumeric() else SMILES_or_pubchem_id]\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\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Prepare PDB file with ESM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "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": {},
      "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"
      ]
    },
    {
      "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"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Post-process (Get smina binding affinities)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "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",
        "\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/smina --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/smina --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))"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generate Results CSV"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "## Ceate dataframe, tar file and download\n",
        "df_results = pd.DataFrame(rows, columns=[\"pdb_file\", \"smiles\", \"diffdock_confidence\", \"smina_scored_affinity\", \"smina_minimized_affinity\", \"sdf_file\"])\n",
        "df_results_tsv = \"df_diffdock_results.tsv\"\n",
        "df_results.to_csv(df_results_tsv, sep='\\t', index=None)\n",
        "\n",
        "df_results"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compare smina affinities with DiffDock confidences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "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=\"smina_scored_affinity\", ax=ax[0]);\n",
        "    sns.regplot(data=df_group, x=\"diffdock_confidence\", y=\"smina_minimized_affinity\", ax=ax[1]);\n",
        "\n",
        "    slope, intercept, r_value_scored, p_value, std_err = linregress(df_group[\"diffdock_confidence\"], df_group[\"smina_scored_affinity\"])\n",
        "    slope, intercept, r_value_minimized, p_value, std_err = linregress(df_group[\"diffdock_confidence\"], df_group[\"smina_minimized_affinity\"])\n",
        "    ax[0].set_title(f\"{pdb_file} {smiles[:30]} smina scored r={r_value_scored:.3f}\");\n",
        "    ax[1].set_title(f\"{pdb_file} {smiles[:30]} smina minimized r={r_value_minimized:.3f}\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "df_results.sort_values(\"diffdock_confidence\", ascending=False).head(3)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualize top hit (highest confidence) in 3D"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "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",
        "# top hit for any pdb file and any smiles\n",
        "top_hit = df_results.sort_values(\"diffdock_confidence\", ascending=False).iloc[0]\n",
        "print(\"top hit:\")\n",
        "display(top_hit)\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": {
      "provenance": []
    },
    "gpuClass": "premium",
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.9.13 (tags/v3.9.13:6de2ca5, May 17 2022, 16:36:42) [MSC v.1929 64 bit (AMD64)]"
    },
    "vscode": {
      "interpreter": {
        "hash": "3105bb9dbd1db35e33253469a2fbb23a7e0a062113aef81588e5cd6cc0939227"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}