{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5AEsv1z5KXXA"
      },
      "source": [
        "# **IgFold**: Fast, accurate antibody structure prediction\n",
        "\n",
        "Official notebook for [IgFold](https://www.biorxiv.org/content/10.1101/2022.04.20.488972): Fast, accurate antibody structure prediction from deep learning on massive set of natural antibodies.  The code, data, and weights for this work are made available for non-commercial use. For commercial inquiries, please contact `jruffolo[at]jhu.edu`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "0PsLNGK57LDq"
      },
      "outputs": [],
      "source": [
        "#@title Input antibody Fv sequences then press `Runtime` -> `Run all`\n",
        "\n",
        "import os\n",
        "import sys\n",
        "\n",
        "python_version = f\"{sys.version_info.major}.{sys.version_info.minor}\"\n",
        "\n",
        "name = \"my_antibody\" #@param {type:\"string\"}\n",
        "pred_dir = name\n",
        "os.makedirs(pred_dir, exist_ok=True)\n",
        "\n",
        "#@markdown Enter antibody sequences for structure prediction. To predict a nanobody structure (or an individual heavy or light chain), simply provide one sequence.\n",
        "heavy_sequence = \"EVQLVQSGPEVKKPGTSVKVSCKASGFTFMSSAVQWVRQARGQRLEWIGWIVIGSGNTNYAQKFQERVTITRDMSTSTAYMELSSLRSEDTAVYYCAAPYCSSISCNDGFDIWGQGTMVTVS\" #@param {type:\"string\"}\n",
        "light_sequence = \"DVVMTQTPFSLPVSLGDQASISCRSSQSLVHSNGNTYLHWYLQKPGQSPKLLIYKVSNRFSGVPDRFSGSGSGTDFTLKISRVEAEDLGVYFCSQSTHVPYTFGGGTKLEIK\" #@param {type:\"string\"}\n",
        "\n",
        "sequences = {}\n",
        "if len(heavy_sequence) > 0:\n",
        "    sequences[\"H\"] = heavy_sequence\n",
        "if len(light_sequence) > 0:\n",
        "    sequences[\"L\"] = light_sequence\n",
        "\n",
        "#@markdown Perform structural refinement with OpenMM\n",
        "do_refine = True #@param {type:\"boolean\"}\n",
        "#@markdown Renumber predicted antibody structure (Chothia) with AbNumber\n",
        "do_renum = False #@param {type:\"boolean\"}\n",
        "#@markdown Use only a single model for predictions (instead of model ensemble)\n",
        "single_model = False #@param {type:\"boolean\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "LsJNdVE87Go2"
      },
      "outputs": [],
      "source": [
        "#@title Install dependencies\n",
        "\n",
        "PYTHON_VERSION = python_version\n",
        "\n",
        "if not os.path.isfile(\"CONDA_READY\"):\n",
        "  print(\"installing conda...\")\n",
        "  os.system(\"wget -qnc https://github.com/jaimergp/miniforge/releases/latest/download/Mambaforge-colab-Linux-x86_64.sh\")\n",
        "  os.system(\"bash Mambaforge-colab-Linux-x86_64.sh -bfp /usr/local\")\n",
        "  os.system(\"mamba config --set auto_update_conda false\")\n",
        "  os.system(\"touch CONDA_READY\")\n",
        "\n",
        "if not os.path.isfile(\"CODE_READY\"):\n",
        "  print(\"installing igfold...\")\n",
        "  torch_string = \"torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html\"\n",
        "  os.system(f\"pip3 install {torch_string}\")\n",
        "  os.system(f\"pip install 'igfold>=0.3.0' {torch_string}\")\n",
        "  os.system(\"pip install -q --no-warn-conflicts 'py3Dmol>=2.0.1' matplotlib seaborn\")\n",
        "  os.system(\"touch CODE_READY\")\n",
        "\n",
        "if do_refine and not os.path.isfile(\"AMBER_READY\"):\n",
        "  print(\"installing amber...\")\n",
        "  os.system(f\"mamba install -y -q -c conda-forge openmm=7.7.0 python='{PYTHON_VERSION}' pdbfixer 2>&1 1>/dev/null\")\n",
        "  os.system(\"touch AMBER_READY\")\n",
        "\n",
        "if do_renum and not os.path.isfile(\"ABNUMBER_READY\"):\n",
        "  print(\"installing abnumber...\")\n",
        "  os.system(f\"mamba install -y -q -c bioconda abnumber python='{PYTHON_VERSION}' 2>&1 1>/dev/null\")\n",
        "  os.system(\"pip install pandas --force-reinstall\")\n",
        "  os.system(\"touch ABNUMBER_READY\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "a2a3BsiE9AXI"
      },
      "outputs": [],
      "source": [
        "#@title Predict antibody structure with IgFold\n",
        "\n",
        "if f\"/usr/local/lib/python{python_version}/site-packages/\" not in sys.path:\n",
        "    sys.path.insert(0, f\"/usr/local/lib/python{python_version}/site-packages/\")\n",
        "\n",
        "from igfold.utils.visualize import *\n",
        "from igfold import IgFoldRunner\n",
        "\n",
        "num_models = 1 if single_model else 4\n",
        "igfold = IgFoldRunner(num_models=num_models)\n",
        "\n",
        "pred_pdb = os.path.join(pred_dir, f\"{name}.pdb\")\n",
        "pred = igfold.fold(\n",
        "    pred_pdb,\n",
        "    sequences=sequences,\n",
        "    do_refine=do_refine,\n",
        "    use_openmm=True,\n",
        "    do_renum=do_renum,\n",
        ")\n",
        "show_pdb(pred_pdb, len(sequences), bb_sticks=False, sc_sticks=True, color=\"rainbow\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xFOTYxsP9Cz1"
      },
      "outputs": [],
      "source": [
        "#@title Plot per-residue predicted RMSD\n",
        "\n",
        "prmsd_fig_file = os.path.join(pred_dir, f\"{name}_prmsd.png\")\n",
        "plot_prmsd(sequences, pred.prmsd.cpu(), prmsd_fig_file, shade_cdr=do_renum, pdb_file=pred_pdb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ajyElWbZ9EFF"
      },
      "outputs": [],
      "source": [
        "#@title Show predicted structure with predicted RMSD\n",
        "\n",
        "#@markdown Structure is colored from low (blue) to high (red) pRMSD.\n",
        "\n",
        "show_pdb(pred_pdb, len(sequences), bb_sticks=False, sc_sticks=True, color=\"b\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "gZBzjpMdJ77q"
      },
      "outputs": [],
      "source": [
        "#@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",
        "from google.colab import files\n",
        "import locale\n",
        "locale.getpreferredencoding = lambda: \"UTF-8\"\n",
        "\n",
        "!zip -FSr $name\".result.zip\" $pred_dir/ &> /dev/null\n",
        "files.download(f\"{name}.result.zip\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "IgFold.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3.9.12 ('igfold_public')",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.9.12"
    },
    "vscode": {
      "interpreter": {
        "hash": "84181e5f1827f203c248bfcd3a60e7e3a4ffc08f0a7dd8a443bd855d4ab14b5d"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}