OpenFold.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "OpenFold.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pc5-mbsX9PZC"
      },
      "source": [
        "# OpenFold Colab\n",
        "\n",
        "Runs a simplified version of [OpenFold](https://github.com/aqlaboratory/openfold) on a target sequence. Adapted from DeepMind's [official AlphaFold Colab](https://colab.research.google.com/github/deepmind/alphafold/blob/main/notebooks/AlphaFold.ipynb).\n",
        "\n",
        "**Differences to AlphaFold v2.0**\n",
        "\n",
        "OpenFold is a trainable PyTorch reimplementation of AlphaFold 2. For the purposes of inference, it is practically identical to the original (\"practically\" because ensembling is excluded from OpenFold (recycling is enabled, however)).\n",
        "\n",
        "In this notebook, OpenFold is run with your choice of our original OpenFold parameters or DeepMind's publicly released parameters for AlphaFold 2.\n",
        "\n",
        "**Note**\n",
        "\n",
        "Like DeepMind's official Colab, this notebook uses **no templates (homologous structures)** and a selected portion of the full [BFD database](https://bfd.mmseqs.com/).\n",
        "\n",
        "**Citing this work**\n",
        "\n",
        "Any publication that discloses findings arising from using this notebook should [cite](https://github.com/deepmind/alphafold/#citing-this-work) DeepMind's [AlphaFold paper](https://doi.org/10.1038/s41586-021-03819-2).\n",
        "\n",
        "**Licenses**\n",
        "\n",
        "This Colab supports inference with the [AlphaFold model parameters](https://github.com/deepmind/alphafold/#model-parameters-license), made available under the Creative Commons Attribution 4.0 International ([CC BY 4.0](https://creativecommons.org/licenses/by/4.0/legalcode)) license. The Colab itself is provided under the [Apache 2.0 license](https://www.apache.org/licenses/LICENSE-2.0). See the full license statement below.\n",
        "\n",
        "**More information**\n",
        "\n",
        "You can find more information about how AlphaFold/OpenFold works in DeepMind's two Nature papers:\n",
        "\n",
        "*   [AlphaFold methods paper](https://www.nature.com/articles/s41586-021-03819-2)\n",
        "*   [AlphaFold predictions of the human proteome paper](https://www.nature.com/articles/s41586-021-03828-1)\n",
        "\n",
        "FAQ on how to interpret AlphaFold/OpenFold predictions are [here](https://alphafold.ebi.ac.uk/faq)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "woIxeCPygt7K",
        "cellView": "form"
      },
      "source": [
        "#@title Install third-party software\n",
        "\n",
        "#@markdown Please execute this cell by pressing the _Play_ button \n",
        "#@markdown on the left to download and import third-party software \n",
        "#@markdown in this Colab notebook. (See the [acknowledgements](https://github.com/deepmind/alphafold/#acknowledgements) in DeepMind's README.)\n",
        "\n",
        "#@markdown **Note**: This installs the software on the Colab \n",
        "#@markdown notebook in the cloud and not on your computer.\n",
        "\n",
        "from IPython.utils import io\n",
        "import os\n",
        "import subprocess\n",
        "import tqdm.notebook\n",
        "\n",
        "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
        "\n",
        "try:\n",
        "  with tqdm.notebook.tqdm(total=100, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "    with io.capture_output() as captured:\n",
        "      # Uninstall default Colab version of PyTorch.\n",
        "      # %shell pip uninstall -y torch\n",
        "\n",
        "      %shell sudo apt install --quiet --yes hmmer\n",
        "      pbar.update(6)\n",
        "\n",
        "      # Install py3dmol.\n",
        "      %shell pip install py3dmol\n",
        "      pbar.update(2)\n",
        "\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(9)\n",
        "\n",
        "      PATH=%env PATH\n",
        "      %env PATH=/opt/conda/bin:{PATH}\n",
        "      pbar.update(80)\n",
        "\n",
        "      # Create a ramdisk to store a database chunk to make Jackhmmer run fast.\n",
        "      %shell sudo mkdir -m 777 --parents /tmp/ramdisk\n",
        "      %shell sudo mount -t tmpfs -o size=9G ramdisk /tmp/ramdisk\n",
        "      pbar.update(2)\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",
        "\n",
        "      # Install git-lfs\n",
        "      %shell sudo apt-get install git-lfs\n",
        "      %shell git lfs install\n",
        "\n",
        "except subprocess.CalledProcessError as captured:\n",
        "  print(captured)\n",
        "  raise"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VzJ5iMjTtoZw",
        "cellView": "form"
      },
      "source": [
        "#@title Download OpenFold\n",
        "\n",
        "#@markdown Please execute this cell by pressing the *Play* button on \n",
        "#@markdown the left.\n",
        "\n",
        "GIT_REPO = 'https://github.com/aqlaboratory/openfold'\n",
        "\n",
        "OPENFOLD_PARAM_SOURCE_URL = \"https://huggingface.co/nz/OpenFold\"\n",
        "ALPHAFOLD_PARAM_SOURCE_URL = 'https://storage.googleapis.com/alphafold/alphafold_params_2022-01-19.tar'\n",
        "OPENFOLD_PARAMS_DIR = './openfold/openfold/resources/openfold_params'\n",
        "ALPHAFOLD_PARAMS_DIR = './openfold/openfold/resources/params'\n",
        "ALPHAFOLD_PARAMS_PATH = os.path.join(\n",
        "  ALPHAFOLD_PARAMS_DIR, os.path.basename(ALPHAFOLD_PARAM_SOURCE_URL)\n",
        ")\n",
        "\n",
        "try:\n",
        "  with tqdm.notebook.tqdm(total=100, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "    with io.capture_output() as captured:\n",
        "      %shell rm -rf openfold\n",
        "      %shell git clone {GIT_REPO} openfold\n",
        "      pbar.update(8)\n",
        "      # Install the required versions of all dependencies.\n",
        "      %shell conda env update -n base --file openfold/environment.yml\n",
        "      \n",
        "      %shell mkdir -p /content/openfold/openfold/resources\n",
        "      %shell cp -f /content/stereo_chemical_props.txt /content/openfold/openfold/resources\n",
        "      \n",
        "      # Run setup.py to install only Openfold.\n",
        "      %shell pip3 install --no-dependencies ./openfold\n",
        "      pbar.update(10)\n",
        "\n",
        "      # Apply OpenMM patch.\n",
        "      %shell pushd /opt/conda/lib/python3.7/site-packages/ && \\\n",
        "          patch -p0 < /content/openfold/lib/openmm.patch && \\\n",
        "          popd\n",
        "\n",
        "      %shell mkdir --parents \"{ALPHAFOLD_PARAMS_DIR}\"\n",
        "      %shell wget -O {ALPHAFOLD_PARAMS_PATH} {ALPHAFOLD_PARAM_SOURCE_URL}\n",
        "      pbar.update(27)\n",
        "\n",
        "      %shell tar --extract --verbose --file=\"{ALPHAFOLD_PARAMS_PATH}\" \\\n",
        "        --directory=\"{ALPHAFOLD_PARAMS_DIR}\" --preserve-permissions\n",
        "      %shell rm \"{ALPHAFOLD_PARAMS_PATH}\"\n",
        "\n",
        "      %shell mkdir --parents \"{OPENFOLD_PARAMS_DIR}\"\n",
        "      %shell git clone \"{OPENFOLD_PARAM_SOURCE_URL}\" \"{OPENFOLD_PARAMS_DIR}\"\n",
        "      pbar.update(55)\n",
        "except subprocess.CalledProcessError:\n",
        "  print(captured)\n",
        "  raise"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Import Python packages\n",
        "#@markdown Please execute this cell by pressing the *Play* button on \n",
        "#@markdown the left.\n",
        "\n",
        "import sys\n",
        "sys.path.append('/opt/conda/lib/python3.7/site-packages')\n",
        "import os\n",
        "\n",
        "from urllib import request\n",
        "from concurrent import futures\n",
        "from google.colab import files\n",
        "import json\n",
        "from matplotlib import gridspec\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import py3Dmol\n",
        "import torch\n",
        "\n",
        "# A filthy hack to avoid slow Linear layer initialization\n",
        "import openfold.model.primitives\n",
        "\n",
        "def __default_linear_init__(self, *args, **kwargs):\n",
        "    return torch.nn.Linear.__init__(\n",
        "      self, \n",
        "      *args[:2], \n",
        "      **{k:v for k,v in kwargs.items() if k == \"bias\"}\n",
        "    )\n",
        "\n",
        "openfold.model.primitives.Linear.__init__ = __default_linear_init__\n",
        "\n",
        "from openfold import config\n",
        "from openfold.data import feature_pipeline\n",
        "from openfold.data import parsers\n",
        "from openfold.data import data_pipeline\n",
        "from openfold.data.tools import jackhmmer\n",
        "from openfold.model import model\n",
        "from openfold.np import protein\n",
        "from openfold.np.relax import relax\n",
        "from openfold.np.relax import utils\n",
        "from openfold.utils.import_weights import import_jax_weights_\n",
        "from openfold.utils.tensor_utils import tensor_tree_map\n",
        "\n",
        "from IPython import display\n",
        "from ipywidgets import GridspecLayout\n",
        "from ipywidgets import Output"
      ],
      "metadata": {
        "cellView": "form",
        "id": "_FpxxMo-mvcP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W4JpOs6oA-QS"
      },
      "source": [
        "## Making a prediction\n",
        "\n",
        "Please paste the sequence of your protein in the text box below, then run the remaining cells via _Runtime_ > _Run after_. You can also run the cells individually by pressing the _Play_ button on the left.\n",
        "\n",
        "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 (see FAQ below)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rowN0bVYLe9n",
        "cellView": "form"
      },
      "source": [
        "#@markdown ### Enter the amino acid sequence to fold ⬇️\n",
        "sequence = 'MAAHKGAEHHHKAAEHHEQAAKHHHAAAEHHEKGEHEQAAHHADTAYAHHKHAEEHAAQAAKHDAEHHAPKPH'  #@param {type:\"string\"}\n",
        "\n",
        "MIN_SEQUENCE_LENGTH = 16\n",
        "MAX_SEQUENCE_LENGTH = 2500\n",
        "\n",
        "#@markdown ### Choose between OpenFold and AlphaFold model parameters ⬇️\n",
        "\n",
        "weight_set = 'OpenFold' #@param [\"OpenFold\", \"AlphaFold\"]\n",
        "\n",
        "# Remove all whitespaces, tabs and end lines; upper-case\n",
        "sequence = sequence.translate(str.maketrans('', '', ' \\n\\t')).upper()\n",
        "aatypes = set('ACDEFGHIKLMNPQRSTVWY')  # 20 standard aatypes\n",
        "if not set(sequence).issubset(aatypes):\n",
        "  raise Exception(f'Input sequence contains non-amino acid letters: {set(sequence) - aatypes}. AlphaFold only supports 20 standard amino acids as inputs.')\n",
        "if len(sequence) < MIN_SEQUENCE_LENGTH:\n",
        "  raise Exception(f'Input sequence is too short: {len(sequence)} amino acids, while the minimum is {MIN_SEQUENCE_LENGTH}')\n",
        "if len(sequence) > MAX_SEQUENCE_LENGTH:\n",
        "  raise Exception(f'Input sequence is too long: {len(sequence)} amino acids, while the maximum is {MAX_SEQUENCE_LENGTH}. Please use the full AlphaFold system for long sequences.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2tTeTTsLKPjB",
        "cellView": "form"
      },
      "source": [
        "#@title Search against genetic databases\n",
        "\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 OpenFold. In particular, \n",
        "#@markdown you’ll see how well each residue is covered by similar \n",
        "#@markdown sequences in the MSA.\n",
        "\n",
        "# Color bands for visualizing plddt\n",
        "PLDDT_BANDS = [(0, 50, '#FF7D45'),\n",
        "               (50, 70, '#FFDB13'),\n",
        "               (70, 90, '#65CBF3'),\n",
        "               (90, 100, '#0053D6')]\n",
        "\n",
        "# --- Find the closest source ---\n",
        "test_url_pattern = 'https://storage.googleapis.com/alphafold-colab{:s}/latest/uniref90_2021_03.fasta.1'\n",
        "ex = futures.ThreadPoolExecutor(3)\n",
        "def fetch(source):\n",
        "  request.urlretrieve(test_url_pattern.format(source))\n",
        "  return source\n",
        "fs = [ex.submit(fetch, source) for source in ['', '-europe', '-asia']]\n",
        "source = None\n",
        "for f in futures.as_completed(fs):\n",
        "  source = f.result()\n",
        "  ex.shutdown()\n",
        "  break\n",
        "\n",
        "# --- Search against genetic databases ---\n",
        "with open('target.fasta', 'wt') as f:\n",
        "  f.write(f'>query\\n{sequence}')\n",
        "\n",
        "# Run the search against chunks of genetic databases (since the genetic\n",
        "# databases don't fit in Colab ramdisk).\n",
        "\n",
        "jackhmmer_binary_path = '/usr/bin/jackhmmer'\n",
        "dbs = []\n",
        "\n",
        "num_jackhmmer_chunks = {'uniref90': 59, 'smallbfd': 17, 'mgnify': 71}\n",
        "total_jackhmmer_chunks = sum(num_jackhmmer_chunks.values())\n",
        "with tqdm.notebook.tqdm(total=total_jackhmmer_chunks, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "  def jackhmmer_chunk_callback(i):\n",
        "    pbar.update(n=1)\n",
        "\n",
        "  pbar.set_description('Searching uniref90')\n",
        "  jackhmmer_uniref90_runner = jackhmmer.Jackhmmer(\n",
        "      binary_path=jackhmmer_binary_path,\n",
        "      database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/uniref90_2021_03.fasta',\n",
        "      get_tblout=True,\n",
        "      num_streamed_chunks=num_jackhmmer_chunks['uniref90'],\n",
        "      streaming_callback=jackhmmer_chunk_callback,\n",
        "      z_value=135301051)\n",
        "  dbs.append(('uniref90', jackhmmer_uniref90_runner.query('target.fasta')))\n",
        "\n",
        "  pbar.set_description('Searching smallbfd')\n",
        "  jackhmmer_smallbfd_runner = jackhmmer.Jackhmmer(\n",
        "      binary_path=jackhmmer_binary_path,\n",
        "      database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/bfd-first_non_consensus_sequences.fasta',\n",
        "      get_tblout=True,\n",
        "      num_streamed_chunks=num_jackhmmer_chunks['smallbfd'],\n",
        "      streaming_callback=jackhmmer_chunk_callback,\n",
        "      z_value=65984053)\n",
        "  dbs.append(('smallbfd', jackhmmer_smallbfd_runner.query('target.fasta')))\n",
        "\n",
        "  pbar.set_description('Searching mgnify')\n",
        "  jackhmmer_mgnify_runner = jackhmmer.Jackhmmer(\n",
        "      binary_path=jackhmmer_binary_path,\n",
        "      database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/mgy_clusters_2019_05.fasta',\n",
        "      get_tblout=True,\n",
        "      num_streamed_chunks=num_jackhmmer_chunks['mgnify'],\n",
        "      streaming_callback=jackhmmer_chunk_callback,\n",
        "      z_value=304820129)\n",
        "  dbs.append(('mgnify', jackhmmer_mgnify_runner.query('target.fasta')))\n",
        "\n",
        "\n",
        "# --- Extract the MSAs and visualize ---\n",
        "# Extract the MSAs from the Stockholm files.\n",
        "# NB: deduplication happens later in data_pipeline.make_msa_features.\n",
        "\n",
        "mgnify_max_hits = 501\n",
        "\n",
        "msas = []\n",
        "deletion_matrices = []\n",
        "full_msa = []\n",
        "for db_name, db_results in dbs:\n",
        "  unsorted_results = []\n",
        "  for i, result in enumerate(db_results):\n",
        "    msa, deletion_matrix, target_names = parsers.parse_stockholm(result['sto'])\n",
        "    e_values_dict = parsers.parse_e_values_from_tblout(result['tbl'])\n",
        "    e_values = [e_values_dict[t.split('/')[0]] for t in target_names]\n",
        "    zipped_results = zip(msa, deletion_matrix, target_names, e_values)\n",
        "    if i != 0:\n",
        "      # Only take query from the first chunk\n",
        "      zipped_results = [x for x in zipped_results if x[2] != 'query']\n",
        "    unsorted_results.extend(zipped_results)\n",
        "  sorted_by_evalue = sorted(unsorted_results, key=lambda x: x[3])\n",
        "  db_msas, db_deletion_matrices, _, _ = zip(*sorted_by_evalue)\n",
        "  if db_msas:\n",
        "    if db_name == 'mgnify':\n",
        "      db_msas = db_msas[:mgnify_max_hits]\n",
        "      db_deletion_matrices = db_deletion_matrices[:mgnify_max_hits]\n",
        "    full_msa.extend(db_msas)\n",
        "    msas.append(db_msas)\n",
        "    deletion_matrices.append(db_deletion_matrices)\n",
        "    msa_size = len(set(db_msas))\n",
        "    print(f'{msa_size} Sequences Found in {db_name}')\n",
        "\n",
        "deduped_full_msa = list(dict.fromkeys(full_msa))\n",
        "total_msa_size = len(deduped_full_msa)\n",
        "print(f'\\n{total_msa_size} Sequences Found in Total\\n')\n",
        "\n",
        "aa_map = {restype: i for i, restype in enumerate('ABCDEFGHIJKLMNOPQRSTUVWXYZ-')}\n",
        "msa_arr = np.array([[aa_map[aa] for aa in seq] for seq in deduped_full_msa])\n",
        "num_alignments, num_res = msa_arr.shape\n",
        "\n",
        "fig = plt.figure(figsize=(12, 3))\n",
        "plt.title('Per-Residue Count of Non-Gap Amino Acids in the MSA')\n",
        "plt.plot(np.sum(msa_arr != aa_map['-'], axis=0), color='black')\n",
        "plt.ylabel('Non-Gap Count')\n",
        "plt.yticks(range(0, num_alignments + 1, max(1, int(num_alignments / 3))))\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XUo6foMQxwS2",
        "cellView": "form"
      },
      "source": [
        "#@title Run OpenFold and 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",
        "# --- Run the model ---\n",
        "model_names = ['model_1', 'model_2', 'model_3', 'model_4', 'model_5', 'model_1_ptm']\n",
        "\n",
        "def _placeholder_template_feats(num_templates_, num_res_):\n",
        "  return {\n",
        "      'template_aatype': torch.zeros(num_templates_, num_res_, 22).long(),\n",
        "      'template_all_atom_positions': torch.zeros(num_templates_, num_res_, 37, 3),\n",
        "      'template_all_atom_mask': torch.zeros(num_templates_, num_res_, 37),\n",
        "      'template_domain_names': torch.zeros(num_templates_),\n",
        "      'template_sum_probs': torch.zeros(num_templates_, 1),\n",
        "  }\n",
        "\n",
        "output_dir = 'prediction'\n",
        "os.makedirs(output_dir, exist_ok=True)\n",
        "\n",
        "plddts = {}\n",
        "pae_outputs = {}\n",
        "unrelaxed_proteins = {}\n",
        "\n",
        "with tqdm.notebook.tqdm(total=len(model_names) + 1, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "  for model_name in model_names:\n",
        "    pbar.set_description(f'Running {model_name}')\n",
        "    num_templates = 1 # dummy number --- is ignored\n",
        "    num_res = len(sequence)\n",
        "\n",
        "    feature_dict = {}\n",
        "    feature_dict.update(data_pipeline.make_sequence_features(sequence, 'test', num_res))\n",
        "    feature_dict.update(data_pipeline.make_msa_features(msas, deletion_matrices=deletion_matrices))\n",
        "    feature_dict.update(_placeholder_template_feats(num_templates, num_res))\n",
        "\n",
        "    cfg = config.model_config(model_name)\n",
        "    openfold_model = model.AlphaFold(cfg)\n",
        "    openfold_model = openfold_model.eval()\n",
        "    if(weight_set == \"AlphaFold\"):\n",
        "      params_name = os.path.join(ALPHAFOLD_PARAMS_DIR, f\"params_{model_name}.npz\")\n",
        "      import_jax_weights_(openfold_model, params_name, version=model_name)\n",
        "    elif(weight_set == \"OpenFold\"):\n",
        "      model_name_spl = model_name.split(\"_\")\n",
        "      if(model_name_spl[-1] == \"ptm\"):\n",
        "        of_model_name = \"finetuning_ptm_2.pt\"\n",
        "      else:\n",
        "        of_model_name = f\"finetuning_{model_name_spl[-1]}.pt\"\n",
        "      params_name = os.path.join(\n",
        "        OPENFOLD_PARAMS_DIR,\n",
        "        of_model_name\n",
        "      )\n",
        "      d = torch.load(params_name)\n",
        "      openfold_model.load_state_dict(d)\n",
        "    else:\n",
        "      raise ValueError(f\"Invalid weight set: {weight_set}\")\n",
        "\n",
        "    openfold_model = openfold_model.cuda()\n",
        "\n",
        "    pipeline = feature_pipeline.FeaturePipeline(cfg.data)\n",
        "    processed_feature_dict = pipeline.process_features(\n",
        "      feature_dict, mode='predict'\n",
        "    )\n",
        "\n",
        "    processed_feature_dict = tensor_tree_map(\n",
        "        lambda t: t.cuda(), processed_feature_dict\n",
        "    )\n",
        "\n",
        "    with torch.no_grad():\n",
        "      prediction_result = openfold_model(processed_feature_dict)\n",
        "\n",
        "    # Move the batch and output to np for further processing\n",
        "    processed_feature_dict = tensor_tree_map(\n",
        "      lambda t: np.array(t[..., -1].cpu()), processed_feature_dict\n",
        "    )\n",
        "    prediction_result = tensor_tree_map(\n",
        "      lambda t: np.array(t.cpu()), prediction_result\n",
        "    )\n",
        "\n",
        "    mean_plddt = prediction_result['plddt'].mean()\n",
        "\n",
        "    if 'predicted_aligned_error' in prediction_result:\n",
        "      pae_outputs[model_name] = (\n",
        "          prediction_result['predicted_aligned_error'],\n",
        "          prediction_result['max_predicted_aligned_error']\n",
        "      )\n",
        "    else:\n",
        "      # Get the pLDDT confidence metrics. Do not put pTM models here as they\n",
        "      # should never get selected.\n",
        "      plddts[model_name] = prediction_result['plddt']\n",
        "\n",
        "    # Set the b-factors to the per-residue plddt.\n",
        "    final_atom_mask = prediction_result['final_atom_mask']\n",
        "    b_factors = prediction_result['plddt'][:, None] * final_atom_mask\n",
        "    unrelaxed_protein = protein.from_prediction(\n",
        "      processed_feature_dict, prediction_result, b_factors=b_factors\n",
        "    )\n",
        "    unrelaxed_proteins[model_name] = unrelaxed_protein\n",
        "\n",
        "    # Delete unused outputs to save memory.\n",
        "    del openfold_model\n",
        "    del processed_feature_dict\n",
        "    del prediction_result\n",
        "    pbar.update(n=1)\n",
        "\n",
        "  # --- AMBER relax the best model ---\n",
        "  pbar.set_description(f'AMBER relaxation')\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",
        "      use_gpu=True,\n",
        "  )\n",
        "  # Find the best model according to the mean pLDDT.\n",
        "  best_model_name = max(plddts.keys(), key=lambda x: plddts[x].mean())\n",
        "  relaxed_pdb, _, _ = amber_relaxer.process(\n",
        "      prot=unrelaxed_proteins[best_model_name])\n",
        "  pbar.update(n=1)  # Finished AMBER relax.\n",
        "\n",
        "# Construct multiclass b-factors to indicate confidence bands\n",
        "# 0=very low, 1=low, 2=confident, 3=very high\n",
        "banded_b_factors = []\n",
        "for plddt in plddts[best_model_name]:\n",
        "  for idx, (min_val, max_val, _) in enumerate(PLDDT_BANDS):\n",
        "    if plddt >= min_val and plddt <= max_val:\n",
        "      banded_b_factors.append(idx)\n",
        "      break\n",
        "banded_b_factors = np.array(banded_b_factors)[:, None] * final_atom_mask\n",
        "to_visualize_pdb = utils.overwrite_b_factors(relaxed_pdb, banded_b_factors)\n",
        "\n",
        "\n",
        "# Write out the prediction\n",
        "pred_output_path = os.path.join(output_dir, 'selected_prediction.pdb')\n",
        "with open(pred_output_path, 'w') as f:\n",
        "  f.write(relaxed_pdb)\n",
        "\n",
        "\n",
        "# --- Visualise the prediction & confidence ---\n",
        "show_sidechains = True\n",
        "def plot_plddt_legend():\n",
        "  \"\"\"Plots the legend for pLDDT.\"\"\"\n",
        "  thresh = [\n",
        "            'Very low (pLDDT < 50)',\n",
        "            'Low (70 > pLDDT > 50)',\n",
        "            'Confident (90 > pLDDT > 70)',\n",
        "            'Very high (pLDDT > 90)']\n",
        "\n",
        "  colors = [x[2] for x in PLDDT_BANDS]\n",
        "\n",
        "  plt.figure(figsize=(2, 2))\n",
        "  for c in colors:\n",
        "    plt.bar(0, 0, color=c)\n",
        "  plt.legend(thresh, frameon=False, loc='center', fontsize=20)\n",
        "  plt.xticks([])\n",
        "  plt.yticks([])\n",
        "  ax = plt.gca()\n",
        "  ax.spines['right'].set_visible(False)\n",
        "  ax.spines['top'].set_visible(False)\n",
        "  ax.spines['left'].set_visible(False)\n",
        "  ax.spines['bottom'].set_visible(False)\n",
        "  plt.title('Model Confidence', fontsize=20, pad=20)\n",
        "  return plt\n",
        "\n",
        "# Color the structure by per-residue pLDDT\n",
        "color_map = {i: bands[2] for i, bands in enumerate(PLDDT_BANDS)}\n",
        "view = py3Dmol.view(width=800, height=600)\n",
        "view.addModelsAsFrames(to_visualize_pdb)\n",
        "style = {'cartoon': {\n",
        "    'colorscheme': {\n",
        "        'prop': 'b',\n",
        "        'map': color_map}\n",
        "        }}\n",
        "if show_sidechains:\n",
        "  style['stick'] = {}\n",
        "view.setStyle({'model': -1}, style)\n",
        "view.zoomTo()\n",
        "\n",
        "grid = GridspecLayout(1, 2)\n",
        "out = Output()\n",
        "with out:\n",
        "  view.show()\n",
        "grid[0, 0] = out\n",
        "\n",
        "out = Output()\n",
        "with out:\n",
        "  plot_plddt_legend().show()\n",
        "grid[0, 1] = out\n",
        "\n",
        "display.display(grid)\n",
        "\n",
        "# Display pLDDT and predicted aligned error (if output by the model).\n",
        "if pae_outputs:\n",
        "  num_plots = 2\n",
        "else:\n",
        "  num_plots = 1\n",
        "\n",
        "plt.figure(figsize=[8 * num_plots, 6])\n",
        "plt.subplot(1, num_plots, 1)\n",
        "plt.plot(plddts[best_model_name])\n",
        "plt.title('Predicted LDDT')\n",
        "plt.xlabel('Residue')\n",
        "plt.ylabel('pLDDT')\n",
        "\n",
        "if num_plots == 2:\n",
        "  plt.subplot(1, 2, 2)\n",
        "  pae, max_pae = list(pae_outputs.values())[0]\n",
        "  plt.imshow(pae, vmin=0., vmax=max_pae, cmap='Greens_r')\n",
        "  plt.colorbar(fraction=0.046, pad=0.04)\n",
        "  plt.title('Predicted Aligned Error')\n",
        "  plt.xlabel('Scored residue')\n",
        "  plt.ylabel('Aligned residue')\n",
        "\n",
        "# Save pLDDT and predicted aligned error (if it exists)\n",
        "pae_output_path = os.path.join(output_dir, 'predicted_aligned_error.json')\n",
        "if pae_outputs:\n",
        "  # Save predicted aligned error in the same format as the AF EMBL DB\n",
        "  rounded_errors = np.round(pae.astype(np.float64), 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)\n",
        "\n",
        "\n",
        "# --- Download the predictions ---\n",
        "!zip -q -r {output_dir}.zip {output_dir}\n",
        "files.download(f'{output_dir}.zip')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lUQAn5LYC5n4"
      },
      "source": [
        "### Interpreting the prediction\n",
        "\n",
        "Please see the [AlphaFold methods paper](https://www.nature.com/articles/s41586-021-03819-2) and the [AlphaFold predictions of the human proteome paper](https://www.nature.com/articles/s41586-021-03828-1), as well as [DeepMind's FAQ](https://alphafold.ebi.ac.uk/faq) on how to interpret AlphaFold/OpenFold predictions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jeb2z8DIA4om"
      },
      "source": [
        "## FAQ & Troubleshooting\n",
        "\n",
        "\n",
        "*   How do I get a predicted protein structure for my protein?\n",
        "    *   Click on the _Connect_ button on the top right to get started.\n",
        "    *   Paste the amino acid sequence of your protein (without any headers) into the “Enter the amino acid sequence to fold”.\n",
        "    *   Run all cells in the Colab, either by running them individually (with the play button on the left side) or via _Runtime_ > _Run all._\n",
        "    *   The predicted protein structure will be downloaded once all cells have been executed. Note: This can take minutes to hours - see below.\n",
        "*   How long will this take?\n",
        "    *   Downloading the OpenFold source code can take up to a few minutes.\n",
        "    *   Downloading and installing the third-party software can take up to a few minutes.\n",
        "    *   The search against genetic databases can take minutes to hours.\n",
        "    *   Running OpenFold and generating the prediction can take minutes to hours, depending on the length of your protein and on which GPU-type Colab has assigned you.\n",
        "*   My Colab no longer seems to be doing anything, what should I do?\n",
        "    *   Some steps may take minutes to hours to complete.\n",
        "    *   Sometimes, running the \"installation\" cells more than once can corrupt the OpenFold installation.\n",
        "    *   If nothing happens or if you receive an error message, try restarting your Colab runtime via _Runtime_ > _Restart runtime_.\n",
        "    *   If this doesn’t help, reset your Colab runtime via _Runtime_ > _Factory reset runtime_.\n",
        "*   How does what's run in this notebook compare to the full versions of Alphafold/Openfold?\n",
        "    *   This Colab version of OpenFold searches a selected portion of the BFD dataset and currently doesn’t use templates, so its accuracy is reduced in comparison to the full version, which is analogous to what's described in the [AlphaFold paper](https://doi.org/10.1038/s41586-021-03819-2) and [Github repo](https://github.com/deepmind/alphafold/). The full version of OpenFold can be run from our own [GitHub repo](https://github.com/aqlaboratory/openfold).\n",
        "*   What is a Colab?\n",
        "    *   See the [Colab FAQ](https://research.google.com/colaboratory/faq.html).\n",
        "*   I received a warning “Notebook requires high RAM”, what do I do?\n",
        "    *   The resources allocated to your Colab vary. See the [Colab FAQ](https://research.google.com/colaboratory/faq.html) for more details.\n",
        "    *   You can execute the Colab nonetheless.\n",
        "*   I received an error “Colab CPU runtime not supported” or “No GPU/TPU found”, what do I do?\n",
        "    *   Colab CPU runtime is not supported. Try changing your runtime via _Runtime_ > _Change runtime type_ > _Hardware accelerator_ > _GPU_.\n",
        "    *   The type of GPU allocated to your Colab varies. See the [Colab FAQ](https://research.google.com/colaboratory/faq.html) for more details.\n",
        "    *   If you receive “Cannot connect to GPU backend”, you can try again later to see if Colab allocates you a GPU.\n",
        "    *   [Colab Pro](https://colab.research.google.com/signup) offers priority access to GPUs. \n",
        "*   Does this tool install anything on my computer?\n",
        "    *   No, everything happens in the cloud on Google Colab.\n",
        "    *   At the end of the Colab execution a zip-archive with the obtained prediction will be automatically downloaded to your computer.\n",
        "*   How should I share feedback and bug reports?\n",
        "    *   Please share any feedback and bug reports as an [issue](https://github.com/aqlaboratory/openfold/issues) on Github.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YfPhvYgKC81B"
      },
      "source": [
        "# License and Disclaimer\n",
        "\n",
        "This Colab notebook and other information provided is for theoretical modelling only, caution should be exercised in its use. It is provided ‘as-is’ without any warranty of any kind, whether expressed or implied. Information is not intended to be a substitute for professional medical advice, diagnosis, or treatment, and does not constitute medical or other professional advice.\n",
        "\n",
        "## AlphaFold/OpenFold Code License\n",
        "\n",
        "Copyright 2021 AlQuraishi Laboratory\n",
        "\n",
        "Copyright 2021 DeepMind Technologies Limited.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at https://www.apache.org/licenses/LICENSE-2.0.\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.\n",
        "\n",
        "## Model Parameters License\n",
        "\n",
        "DeepMind's AlphaFold parameters are made available under the terms of the Creative Commons Attribution 4.0 International (CC BY 4.0) license. You can find details at: https://creativecommons.org/licenses/by/4.0/legalcode\n",
        "\n",
        "\n",
        "## Third-party software\n",
        "\n",
        "Use of the third-party software, libraries or code referred to in this notebook may be governed by separate terms and conditions or license provisions. Your use of the third-party software, libraries or code is subject to any such terms and you should check that you can comply with any applicable restrictions or terms and conditions before use.\n",
        "\n",
        "\n",
        "## Mirrored Databases\n",
        "\n",
        "The following databases have been mirrored by DeepMind, and are available with reference to the following:\n",
        "* UniRef90: v2021\\_03 (unmodified), by The UniProt Consortium, available under a [Creative Commons Attribution-NoDerivatives 4.0 International License](http://creativecommons.org/licenses/by-nd/4.0/).\n",
        "* MGnify: v2019\\_05 (unmodified), by Mitchell AL et al., available free of all copyright restrictions and made fully and freely available for both non-commercial and commercial use under [CC0 1.0 Universal (CC0 1.0) Public Domain Dedication](https://creativecommons.org/publicdomain/zero/1.0/).\n",
        "* BFD: (modified), by Steinegger M. and Söding J., modified by DeepMind, available under a [Creative Commons Attribution-ShareAlike 4.0 International License](https://creativecommons.org/licenses/by/4.0/). See the Methods section of the [AlphaFold proteome paper](https://www.nature.com/articles/s41586-021-03828-1) for details."
      ]
    }
  ]
}