DeepAb.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "nfLJKttO3fM4"
            },
            "source": [
                "# DeepAb antibody structure prediction notebook\n",
                "\n",
                "Official notebook for [DeepAb](https://www.sciencedirect.com/science/article/pii/S2666389921002804): Antibody structure prediction using interpretable deep learning. This notebook is made freely available for non-commercial use. Commercial use is available under the terms of the [Rosetta-DL license](https://github.com/RosettaCommons/Rosetta-DL).\n",
                "\n",
                "For more details about DeepAb, please read:\n",
                "[Ruffolo, J. A., Sulam, J., & Gray, J. J. (2022). Antibody structure prediction using interpretable deep learning. _Patterns_, 3(2), 100406.](https://www.sciencedirect.com/science/article/pii/S2666389921002804)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "cellView": "form",
                "id": "Xm2ACqt_3fM-"
            },
            "outputs": [],
            "source": [
                "#@title Input antibody Fv sequences then press `Runtime` -> `Run all`\n",
                "\n",
                "#@markdown To predict single-chain antibodies (nanobodies) or individual light chains, just omit the other chain sequence.\n",
                "\n",
                "import os\n",
                "\n",
                "name = \"my_antibody\" #@param {type:\"string\"}\n",
                "pred_dir = name\n",
                "os.makedirs(pred_dir, exist_ok=True)\n",
                "\n",
                "heavy_sequence = 'EIQLQQSGPELVKPGASVKISCKASGYSFTDYIMLWVKQSHGKSLEWIGNINPYYGSTSYNLKFKGKATLTVDKSSSTAYMQLNSLTSEDSAVYYCARKNYYGSSLDYWGQGTTLTVS' #@param {type:\"string\"}\n",
                "light_sequence = 'DVVMTQTPFSLPVSLGDQASISCRSSQSLVHSNGNTYLHWYLQKPGQSPKLLIYKVSNRFSGVPDRFSGSGSGTDFTLKISRVEAEDLGVYFCSQSTHVPYTFGGGTKLEIK' #@param {type:\"string\"}\n",
                "\n",
                "decoys = 1 #@param {type:\"integer\"}\n",
                "single_chain = len(heavy_sequence) == 0 or len(light_sequence) == 0\n",
                "renumber = True #@param {type:\"boolean\"}"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "cellView": "form",
                "id": "8dMQ7USW3fNB"
            },
            "outputs": [],
            "source": [
                "#@title Install DeepAb code and dependencies\n",
                "\n",
                "#@markdown This step will clone the DeepAb repository and install the required dependencies.\n",
                "#@markdown Be sure to obtain the appropriate [PyRosetta license](https://www.pyrosetta.org/home/licensing-pyrosetta). This step may request access to your Google Drive and require restarting your runtime. If this is your first time running the notebook, this step may take several minutes.\n",
                "\n",
                "! git clone http://github.com/rosettacommons/deepab\n",
                "%pip install -r deepab/requirements.txt\n",
                "\n",
                "import sys\n",
                "from glob import glob\n",
                "from tqdm.contrib.concurrent import process_map\n",
                "import torch\n",
                "import numpy as np\n",
                "\n",
                "sys.path.insert(0, \"deepab\")\n",
                "\n",
                "OUT_DIR = \"predict_structure\"\n",
                "os.makedirs(OUT_DIR, exist_ok=True)\n",
                "\n",
                "# Notebook setup\n",
                "if 'google.colab' in sys.modules:\n",
                "  !pip install pyrosettacolabsetup\n",
                "  import pyrosettacolabsetup\n",
                "  if not os.path.exists(\"/content/google_drive/MyDrive/PyRosetta/colab.bin\"):\n",
                "    pyrosettacolabsetup.install_pyrosetta()\n",
                "  else:\n",
                "    pyrosettacolabsetup.mount_pyrosetta_install()\n",
                "\n",
                "import pyrosetta\n",
                "pyrosetta.init(\n",
                "    \"-mute all -detect_disulf true -detect_disulf_tolerance 1.5 -check_cdr_chainbreaks false\"\n",
                ")\n",
                "\n",
                "from predict import load_model, ModelEnsemble, get_cst_defs, build_structure, renumber_pdb\n",
                "from annotate_attention import cdr_indices, annotate_structure"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "cellView": "form",
                "id": "6C0kI-6Y3fND"
            },
            "outputs": [],
            "source": [
                "#@title Download trained models\n",
                "\n",
                "#@markdown This step may take several minutes if models are not found locally.\n",
                "\n",
                "model_files = [\n",
                "    \"ensemble_abresnet/rs0.pt\",\n",
                "    \"ensemble_abresnet/rs1.pt\",\n",
                "    \"ensemble_abresnet/rs2.pt\",\n",
                "    \"ensemble_abresnet/rs3.pt\",\n",
                "    \"ensemble_abresnet/rs4.pt\",\n",
                "]\n",
                "\n",
                "do_download = all([not os.path.exists(f) for f in model_files])\n",
                "if do_download:\n",
                "    os.system(\"wget https://data.graylab.jhu.edu/ensemble_abresnet_v1.tar.gz\")\n",
                "    os.system(\"tar -xf ensemble_abresnet_v1.tar.gz\")\n",
                "\n",
                "device_type = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
                "device = torch.device(device_type)\n",
                "model = ModelEnsemble(model_files=model_files,\n",
                "                          load_model=load_model,\n",
                "                          eval_mode=True,\n",
                "                          device=device)\n",
                "single_model = load_model(model_files[0],\n",
                "                          eval_mode=True,\n",
                "                          device=device)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "cellView": "form",
                "id": "FdgOBLv93fND"
            },
            "outputs": [],
            "source": [
                "#@title Predict antibody structure with DeepAb\n",
                "\n",
                "#@markdown Antibody structure prediction with DeepAb involves three main steps: constraint generation, initial structure prediction by multi-dimensional scaling, and structural refinement with PyRosetta.\n",
                "#@markdown Expected runtime is approximately ten minutes for five decoys.\n",
                "\n",
                "fasta_file = os.path.join(pred_dir, f\"{name}.fasta\")\n",
                "with open(fasta_file, \"w\") as f:\n",
                "    f.write(\">:H\\n{}\\n>:L\\n{}\\n\".format(heavy_sequence, light_sequence))\n",
                "cst_defs = get_cst_defs(model, fasta_file, device=device)\n",
                "\n",
                "pred_pdb = build_structure(model,\n",
                "                                fasta_file,\n",
                "                                cst_defs,\n",
                "                                pred_dir,\n",
                "                                target=name,\n",
                "                                num_decoys=decoys,\n",
                "                                num_procs=decoys,\n",
                "                                single_chain=single_chain,\n",
                "                                device=device)\n",
                "\n",
                "if renumber:\n",
                "    renumber_pdb(pred_pdb, pred_pdb)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "cellView": "form",
                "id": "wlTrQp0Z3fNE"
            },
            "outputs": [],
            "source": [
                "#@title Annotate predicted structure with output attention.\n",
                "\n",
                "#@markdown Visualize the inter-residue attention from an output geometry branch from DeepAb. The attention scores will replace the B-factors in the predicted structure.\n",
                "\n",
                "cdr_loop = \"h3\" #@param [\"h1\", \"h2\", \"h3\", \"l1\", \"l2\", \"l3\"]\n",
                "cdr_i = cdr_indices(pred_pdb, cdr_loop)\n",
                "\n",
                "attn_branch = \"ca\" #@param [\"ca\", \"cb\", \"no\", \"omega\", \"theta\", \"phi\"]\n",
                "\n",
                "attn_pdb_file = pred_pdb.replace(\".pdb\", f\".{cdr_loop}_attn.pdb\")\n",
                "os.system(f\"cp {pred_pdb} {attn_pdb_file}\")\n",
                "\n",
                "annotate_structure(single_model, fasta_file, attn_pdb_file, cdr_i, attn_branch)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "cellView": "form",
                "id": "1WydyJMv3fNF"
            },
            "outputs": [],
            "source": [
                "from google.colab import files\n",
                "\n",
                "#@title Download results\n",
                "\n",
                "#@markdown Download zip file containing structure prediction and annotation results. If download fails, results are also accessible from file explorer on the left panel of the notebook.\n",
                "\n",
                "!zip -FSr $name\".result.zip\" $pred_dir/\n",
                "files.download(f\"{name}.result.zip\")"
            ]
        }
    ],
    "metadata": {
        "interpreter": {
            "hash": "782c5e8831d9fdd1f686474d0afe590ac2fe73cb65eedbd6b51c5a6ccffb32a5"
        },
        "kernelspec": {
            "display_name": "Python 3.7.13 ('venv': venv)",
            "language": "python",
            "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.7.13"
        },
        "orig_nbformat": 4,
        "colab": {
            "name": "Copy of DeepAb.ipynb",
            "provenance": [
                {
                    "file_id": "https://github.com/RosettaCommons/DeepAb/blob/structure-prediction-notebook/DeepAb.ipynb",
                    "timestamp": 1650037043700
                }
            ],
            "collapsed_sections": []
        }
    },
    "nbformat": 4,
    "nbformat_minor": 0
}