AlphaFold_+_B2B_Tools.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Bio2Byte/public_notebooks/blob/main/AlphaFold_%2B_B2B_Tools.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pc5-mbsX9PZC"
      },
      "source": [
        "# AlphaFold Colab + Bio2Byte biophysical predictions\n",
        "\n",
        "This Colab notebook allows you to easily predict the structure of a protein using a slightly simplified version of [AlphaFold v2.1.0](https://doi.org/10.1038/s41586-021-03819-2). In addition, on top of the best predicted model \n",
        "this notebook allows you to visualize the biophysical predictions obteined after\n",
        "using the Bio2Byte tools (https://pypi.org/project/b2bTools/).\n",
        "\n",
        "This Colab notebook is based on this Deepmind's one: https://colab.research.google.com/github/deepmind/alphafold/blob/main/notebooks/AlphaFold.ipynb#scrollTo=XUo6foMQxwS2\n",
        "\n",
        "\n",
        "## From the original Alphafold notebook\n",
        "\n",
        "**Differences to AlphaFold v2.1.0**\n",
        "\n",
        "In comparison to AlphaFold v2.1.0, this Colab notebook uses **no templates (homologous structures)** and a selected portion of the [BFD database](https://bfd.mmseqs.com/). We have validated these changes on several thousand recent PDB structures. While accuracy will be near-identical to the full AlphaFold system on many targets, a small fraction have a large drop in accuracy due to the smaller MSA and lack of templates. For best reliability, we recommend instead using the [full open source AlphaFold](https://github.com/deepmind/alphafold/), or the [AlphaFold Protein Structure Database](https://alphafold.ebi.ac.uk/).\n",
        "\n",
        "**This Colab has a small drop in average accuracy for multimers compared to local AlphaFold installation, for full multimer accuracy it is highly recommended to run [AlphaFold locally](https://github.com/deepmind/alphafold#running-alphafold).** Moreover, the AlphaFold-Multimer requires searching for MSA for every unique sequence in the complex, hence it is substantially slower. If your notebook times-out due to slow multimer MSA search, we recommend either using Colab Pro or running AlphaFold locally.\n",
        "\n",
        "Please note that this Colab notebook is provided as an early-access prototype and is not a finished product. It is provided for theoretical modelling only and caution should be exercised in its use. \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) the [AlphaFold paper](https://doi.org/10.1038/s41586-021-03819-2).\n",
        "\n",
        "**Licenses**\n",
        "\n",
        "This Colab uses the [AlphaFold model parameters](https://github.com/deepmind/alphafold/#model-parameters-license) which are subject to 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 works in the following 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",
        "*   [AlphaFold-Multimer paper](https://www.biorxiv.org/content/10.1101/2021.10.04.463034v1)\n",
        "\n",
        "FAQ on how to interpret AlphaFold predictions are [here](https://alphafold.ebi.ac.uk/faq)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uC1dKAwk2eyl"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Start by running the 2 cells below to set up AlphaFold and all required software."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "woIxeCPygt7K"
      },
      "outputs": [],
      "source": [
        "#@title 1. 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 our 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 TF.\n",
        "      %shell pip uninstall -y tensorflow\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",
        "      %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(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",
        "except subprocess.CalledProcessError:\n",
        "  print(captured)\n",
        "  raise"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "VzJ5iMjTtoZw"
      },
      "outputs": [],
      "source": [
        "#@title 2. Download AlphaFold\n",
        "\n",
        "#@markdown Please execute this cell by pressing the *Play* button on \n",
        "#@markdown the left.\n",
        "\n",
        "GIT_REPO = 'https://github.com/deepmind/alphafold'\n",
        "\n",
        "SOURCE_URL = 'https://storage.googleapis.com/alphafold/alphafold_params_colab_2022-01-19.tar'\n",
        "PARAMS_DIR = './alphafold/data/params'\n",
        "PARAMS_PATH = os.path.join(PARAMS_DIR, os.path.basename(SOURCE_URL))\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 alphafold\n",
        "      %shell git clone --branch main {GIT_REPO} alphafold\n",
        "      pbar.update(8)\n",
        "      # Install the required versions of all dependencies.\n",
        "      %shell pip3 install -r ./alphafold/requirements.txt\n",
        "      # Run setup.py to install only AlphaFold.\n",
        "      %shell pip3 install --no-dependencies ./alphafold\n",
        "      pbar.update(10)\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",
        "      # Make sure stereo_chemical_props.txt is in all locations where it could be searched for.\n",
        "      %shell mkdir -p /content/alphafold/alphafold/common\n",
        "      %shell cp -f /content/stereo_chemical_props.txt /content/alphafold/alphafold/common\n",
        "      %shell mkdir -p /opt/conda/lib/python3.7/site-packages/alphafold/common/\n",
        "      %shell cp -f /content/stereo_chemical_props.txt /opt/conda/lib/python3.7/site-packages/alphafold/common/\n",
        "\n",
        "      %shell mkdir --parents \"{PARAMS_DIR}\"\n",
        "      %shell wget -O \"{PARAMS_PATH}\" \"{SOURCE_URL}\"\n",
        "      pbar.update(27)\n",
        "\n",
        "      %shell tar --extract --verbose --file=\"{PARAMS_PATH}\" \\\n",
        "        --directory=\"{PARAMS_DIR}\" --preserve-permissions\n",
        "      %shell rm \"{PARAMS_PATH}\"\n",
        "      pbar.update(55)\n",
        "except subprocess.CalledProcessError:\n",
        "  print(captured)\n",
        "  raise\n",
        "\n",
        "import jax\n",
        "if jax.local_devices()[0].platform == 'tpu':\n",
        "  raise RuntimeError('Colab TPU runtime not supported. Change it to GPU via Runtime -> Change Runtime Type -> Hardware accelerator -> GPU.')\n",
        "elif jax.local_devices()[0].platform == 'cpu':\n",
        "  raise RuntimeError('Colab CPU runtime not supported. Change it to GPU via Runtime -> Change Runtime Type -> Hardware accelerator -> GPU.')\n",
        "else:\n",
        "  print(f'Running with {jax.local_devices()[0].device_kind} GPU')\n",
        "\n",
        "# Make sure everything we need is on the path.\n",
        "import sys\n",
        "sys.path.append('/opt/conda/lib/python3.7/site-packages')\n",
        "sys.path.append('/content/alphafold')\n",
        "\n",
        "# Make sure all necessary environment variables are set.\n",
        "import os\n",
        "os.environ['TF_FORCE_UNIFIED_MEMORY'] = '1'\n",
        "os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '2.0'"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 2.bis Installing Bio2Byte tools\n",
        "\n",
        "#@markdown Please execute this cell by pressing the *Play* button on \n",
        "#@markdown the left.\n",
        "!pip install b2bTools==3.0.5"
      ],
      "metadata": {
        "cellView": "form",
        "id": "AbDViTxRie8e"
      },
      "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",
      "execution_count": null,
      "metadata": {
        "id": "rowN0bVYLe9n",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title 3. Enter the amino acid sequence(s) to fold ⬇️\n",
        "#@markdown Enter the amino acid sequence(s) to fold:\n",
        "#@markdown * If you enter only a single sequence, the monomer model will be used.\n",
        "#@markdown * If you enter multiple sequences, the multimer model will be used.\n",
        "\n",
        "from alphafold.notebooks import notebook_utils\n",
        "\n",
        "sequence_1 = 'NLYIQWLKDGGPSSGRPPPS'  #@param {type:\"string\"}\n",
        "sequence_2 = ''  #@param {type:\"string\"}\n",
        "sequence_3 = ''  #@param {type:\"string\"}\n",
        "sequence_4 = ''  #@param {type:\"string\"}\n",
        "sequence_5 = ''  #@param {type:\"string\"}\n",
        "sequence_6 = ''  #@param {type:\"string\"}\n",
        "sequence_7 = ''  #@param {type:\"string\"}\n",
        "sequence_8 = ''  #@param {type:\"string\"}\n",
        "\n",
        "input_sequences = (sequence_1, sequence_2, sequence_3, sequence_4,\n",
        "                    sequence_5, sequence_6, sequence_7, sequence_8)\n",
        "\n",
        "#@markdown If folding a complex target and all the input sequences are\n",
        "#@markdown prokaryotic then set `is_prokaryotic` to `True`. Set to `False`\n",
        "#@markdown otherwise or if the origin is unknown.\n",
        "\n",
        "is_prokaryote = False  #@param {type:\"boolean\"}\n",
        "\n",
        "MIN_SINGLE_SEQUENCE_LENGTH = 16\n",
        "MAX_SINGLE_SEQUENCE_LENGTH = 2500\n",
        "MAX_MULTIMER_LENGTH = 2500\n",
        "\n",
        "# Validate the input.\n",
        "sequences, model_type_to_use = notebook_utils.validate_input(\n",
        "    input_sequences=input_sequences,\n",
        "    min_length=MIN_SINGLE_SEQUENCE_LENGTH,\n",
        "    max_length=MAX_SINGLE_SEQUENCE_LENGTH,\n",
        "    max_multimer_length=MAX_MULTIMER_LENGTH)\n",
        "\n",
        "print(sequences)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2tTeTTsLKPjB"
      },
      "outputs": [],
      "source": [
        "#@title 4. 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 AlphaFold. In particular, \n",
        "#@markdown you’ll see how well each residue is covered by similar \n",
        "#@markdown sequences in the MSA.\n",
        "\n",
        "# --- Python imports ---\n",
        "import collections\n",
        "import copy\n",
        "from concurrent import futures\n",
        "import json\n",
        "import random\n",
        "\n",
        "from urllib import request\n",
        "from google.colab import files\n",
        "from matplotlib import gridspec\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import py3Dmol\n",
        "\n",
        "from alphafold.model import model\n",
        "from alphafold.model import config\n",
        "from alphafold.model import data\n",
        "\n",
        "from alphafold.data import feature_processing\n",
        "from alphafold.data import msa_pairing\n",
        "from alphafold.data import pipeline\n",
        "from alphafold.data import pipeline_multimer\n",
        "from alphafold.data.tools import jackhmmer\n",
        "\n",
        "from alphafold.common import protein\n",
        "\n",
        "from alphafold.relax import relax\n",
        "from alphafold.relax import utils\n",
        "\n",
        "from IPython import display\n",
        "from ipywidgets import GridspecLayout\n",
        "from ipywidgets import Output\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",
        "JACKHMMER_BINARY_PATH = '/usr/bin/jackhmmer'\n",
        "DB_ROOT_PATH = f'https://storage.googleapis.com/alphafold-colab{source}/latest/'\n",
        "# The z_value is the number of sequences in a database.\n",
        "MSA_DATABASES = [\n",
        "    {'db_name': 'uniref90',\n",
        "     'db_path': f'{DB_ROOT_PATH}uniref90_2021_03.fasta',\n",
        "     'num_streamed_chunks': 59,\n",
        "     'z_value': 135_301_051},\n",
        "    {'db_name': 'smallbfd',\n",
        "     'db_path': f'{DB_ROOT_PATH}bfd-first_non_consensus_sequences.fasta',\n",
        "     'num_streamed_chunks': 17,\n",
        "     'z_value': 65_984_053},\n",
        "    {'db_name': 'mgnify',\n",
        "     'db_path': f'{DB_ROOT_PATH}mgy_clusters_2019_05.fasta',\n",
        "     'num_streamed_chunks': 71,\n",
        "     'z_value': 304_820_129},\n",
        "]\n",
        "\n",
        "# Search UniProt and construct the all_seq features only for heteromers, not homomers.\n",
        "if model_type_to_use == notebook_utils.ModelType.MULTIMER and len(set(sequences)) > 1:\n",
        "  MSA_DATABASES.extend([\n",
        "      # Swiss-Prot and TrEMBL are concatenated together as UniProt.\n",
        "      {'db_name': 'uniprot',\n",
        "       'db_path': f'{DB_ROOT_PATH}uniprot_2021_03.fasta',\n",
        "       'num_streamed_chunks': 98,\n",
        "       'z_value': 219_174_961 + 565_254},\n",
        "  ])\n",
        "\n",
        "TOTAL_JACKHMMER_CHUNKS = sum([cfg['num_streamed_chunks'] for cfg in MSA_DATABASES])\n",
        "\n",
        "MAX_HITS = {\n",
        "    'uniref90': 10_000,\n",
        "    'smallbfd': 5_000,\n",
        "    'mgnify': 501,\n",
        "    'uniprot': 50_000,\n",
        "}\n",
        "\n",
        "\n",
        "def get_msa(fasta_path):\n",
        "  \"\"\"Searches for MSA for the given sequence using chunked Jackhmmer search.\"\"\"\n",
        "\n",
        "  # Run the search against chunks of genetic databases (since the genetic\n",
        "  # databases don't fit in Colab disk).\n",
        "  raw_msa_results = collections.defaultdict(list)\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",
        "    for db_config in MSA_DATABASES:\n",
        "      db_name = db_config['db_name']\n",
        "      pbar.set_description(f'Searching {db_name}')\n",
        "      jackhmmer_runner = jackhmmer.Jackhmmer(\n",
        "          binary_path=JACKHMMER_BINARY_PATH,\n",
        "          database_path=db_config['db_path'],\n",
        "          get_tblout=True,\n",
        "          num_streamed_chunks=db_config['num_streamed_chunks'],\n",
        "          streaming_callback=jackhmmer_chunk_callback,\n",
        "          z_value=db_config['z_value'])\n",
        "      # Group the results by database name.\n",
        "      raw_msa_results[db_name].extend(jackhmmer_runner.query(fasta_path))\n",
        "\n",
        "  return raw_msa_results\n",
        "\n",
        "\n",
        "features_for_chain = {}\n",
        "raw_msa_results_for_sequence = {}\n",
        "for sequence_index, sequence in enumerate(sequences, start=1):\n",
        "  print(f'\\nGetting MSA for sequence {sequence_index}')\n",
        "\n",
        "  fasta_path = f'target_{sequence_index}.fasta'\n",
        "  with open(fasta_path, 'wt') as f:\n",
        "    f.write(f'>query\\n{sequence}')\n",
        "\n",
        "  # Don't do redundant work for multiple copies of the same chain in the multimer.\n",
        "  if sequence not in raw_msa_results_for_sequence:\n",
        "    raw_msa_results = get_msa(fasta_path=fasta_path)\n",
        "    raw_msa_results_for_sequence[sequence] = raw_msa_results\n",
        "  else:\n",
        "    raw_msa_results = copy.deepcopy(raw_msa_results_for_sequence[sequence])\n",
        "\n",
        "  # Extract the MSAs from the Stockholm files.\n",
        "  # NB: deduplication happens later in pipeline.make_msa_features.\n",
        "  single_chain_msas = []\n",
        "  uniprot_msa = None\n",
        "  for db_name, db_results in raw_msa_results.items():\n",
        "    merged_msa = notebook_utils.merge_chunked_msa(\n",
        "        results=db_results, max_hits=MAX_HITS.get(db_name))\n",
        "    if merged_msa.sequences and db_name != 'uniprot':\n",
        "      single_chain_msas.append(merged_msa)\n",
        "      msa_size = len(set(merged_msa.sequences))\n",
        "      print(f'{msa_size} unique sequences found in {db_name} for sequence {sequence_index}')\n",
        "    elif merged_msa.sequences and db_name == 'uniprot':\n",
        "      uniprot_msa = merged_msa\n",
        "\n",
        "  notebook_utils.show_msa_info(single_chain_msas=single_chain_msas, sequence_index=sequence_index)\n",
        "\n",
        "  # Turn the raw data into model features.\n",
        "  feature_dict = {}\n",
        "  feature_dict.update(pipeline.make_sequence_features(\n",
        "      sequence=sequence, description='query', num_res=len(sequence)))\n",
        "  feature_dict.update(pipeline.make_msa_features(msas=single_chain_msas))\n",
        "  # We don't use templates in AlphaFold Colab notebook, add only empty placeholder features.\n",
        "  feature_dict.update(notebook_utils.empty_placeholder_template_features(\n",
        "      num_templates=0, num_res=len(sequence)))\n",
        "\n",
        "  # Construct the all_seq features only for heteromers, not homomers.\n",
        "  if model_type_to_use == notebook_utils.ModelType.MULTIMER and len(set(sequences)) > 1:\n",
        "    valid_feats = msa_pairing.MSA_FEATURES + (\n",
        "        'msa_uniprot_accession_identifiers',\n",
        "        'msa_species_identifiers',\n",
        "    )\n",
        "    all_seq_features = {\n",
        "        f'{k}_all_seq': v for k, v in pipeline.make_msa_features([uniprot_msa]).items()\n",
        "        if k in valid_feats}\n",
        "    feature_dict.update(all_seq_features)\n",
        "\n",
        "  features_for_chain[protein.PDB_CHAIN_IDS[sequence_index - 1]] = feature_dict\n",
        "\n",
        "\n",
        "# Do further feature post-processing depending on the model type.\n",
        "if model_type_to_use == notebook_utils.ModelType.MONOMER:\n",
        "  np_example = features_for_chain[protein.PDB_CHAIN_IDS[0]]\n",
        "\n",
        "elif model_type_to_use == notebook_utils.ModelType.MULTIMER:\n",
        "  all_chain_features = {}\n",
        "  for chain_id, chain_features in features_for_chain.items():\n",
        "    all_chain_features[chain_id] = pipeline_multimer.convert_monomer_features(\n",
        "        chain_features, chain_id)\n",
        "\n",
        "  all_chain_features = pipeline_multimer.add_assembly_features(all_chain_features)\n",
        "\n",
        "  np_example = feature_processing.pair_and_merge(\n",
        "      all_chain_features=all_chain_features, is_prokaryote=is_prokaryote)\n",
        "\n",
        "  # Pad MSA to avoid zero-sized extra_msa.\n",
        "  np_example = pipeline_multimer.pad_msa(np_example, min_num_seq=512)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "XUo6foMQxwS2"
      },
      "outputs": [],
      "source": [
        "#@title 5. Run AlphaFold 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",
        "#@markdown In case you are having issues with the relaxation stage, you can disable it below.\n",
        "#@markdown Warning: This means that the prediction might have distracting\n",
        "#@markdown small stereochemical violations.\n",
        "\n",
        "run_relax = True  #@param {type:\"boolean\"}\n",
        "\n",
        "# --- Run the model ---\n",
        "if model_type_to_use == notebook_utils.ModelType.MONOMER:\n",
        "  model_names = config.MODEL_PRESETS['monomer'] + ('model_2_ptm',)\n",
        "elif model_type_to_use == notebook_utils.ModelType.MULTIMER:\n",
        "  model_names = config.MODEL_PRESETS['multimer']\n",
        "\n",
        "output_dir = 'prediction'\n",
        "os.makedirs(output_dir, exist_ok=True)\n",
        "\n",
        "plddts = {}\n",
        "ranking_confidences = {}\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",
        "\n",
        "    cfg = config.model_config(model_name)\n",
        "    if model_type_to_use == notebook_utils.ModelType.MONOMER:\n",
        "      cfg.data.eval.num_ensemble = 1\n",
        "    elif model_type_to_use == notebook_utils.ModelType.MULTIMER:\n",
        "      cfg.model.num_ensemble_eval = 1\n",
        "    params = data.get_model_haiku_params(model_name, './alphafold/data')\n",
        "    model_runner = model.RunModel(cfg, params)\n",
        "    processed_feature_dict = model_runner.process_features(np_example, random_seed=0)\n",
        "    prediction = model_runner.predict(processed_feature_dict, random_seed=random.randrange(sys.maxsize))\n",
        "\n",
        "    mean_plddt = prediction['plddt'].mean()\n",
        "\n",
        "    if model_type_to_use == notebook_utils.ModelType.MONOMER:\n",
        "      if 'predicted_aligned_error' in prediction:\n",
        "        pae_outputs[model_name] = (prediction['predicted_aligned_error'],\n",
        "                                   prediction['max_predicted_aligned_error'])\n",
        "      else:\n",
        "        # Monomer models are sorted by mean pLDDT. Do not put monomer pTM models here as they\n",
        "        # should never get selected.\n",
        "        ranking_confidences[model_name] = prediction['ranking_confidence']\n",
        "        plddts[model_name] = prediction['plddt']\n",
        "    elif model_type_to_use == notebook_utils.ModelType.MULTIMER:\n",
        "      # Multimer models are sorted by pTM+ipTM.\n",
        "      ranking_confidences[model_name] = prediction['ranking_confidence']\n",
        "      plddts[model_name] = prediction['plddt']\n",
        "      pae_outputs[model_name] = (prediction['predicted_aligned_error'],\n",
        "                                 prediction['max_predicted_aligned_error'])\n",
        "\n",
        "    # Set the b-factors to the per-residue plddt.\n",
        "    final_atom_mask = prediction['structure_module']['final_atom_mask']\n",
        "    b_factors = prediction['plddt'][:, None] * final_atom_mask\n",
        "    unrelaxed_protein = protein.from_prediction(\n",
        "        processed_feature_dict,\n",
        "        prediction,\n",
        "        b_factors=b_factors,\n",
        "        remove_leading_feature_dimension=(\n",
        "            model_type_to_use == notebook_utils.ModelType.MONOMER))\n",
        "    unrelaxed_proteins[model_name] = unrelaxed_protein\n",
        "\n",
        "    # Delete unused outputs to save memory.\n",
        "    del model_runner\n",
        "    del params\n",
        "    del prediction\n",
        "    pbar.update(n=1)\n",
        "\n",
        "  # --- AMBER relax the best model ---\n",
        "\n",
        "  # Find the best model according to the mean pLDDT.\n",
        "  best_model_name = max(ranking_confidences.keys(), key=lambda x: ranking_confidences[x])\n",
        "\n",
        "  if run_relax:\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=3,\n",
        "        use_gpu=True)\n",
        "    relaxed_pdb, _, _ = amber_relaxer.process(prot=unrelaxed_proteins[best_model_name])\n",
        "  else:\n",
        "    print('Warning: Running without the relaxation stage.')\n",
        "    relaxed_pdb = protein.to_pdb(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 = ['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",
        "# Show the structure coloured by chain if the multimer model has been used.\n",
        "if model_type_to_use == notebook_utils.ModelType.MULTIMER:\n",
        "  multichain_view = py3Dmol.view(width=800, height=600)\n",
        "  multichain_view.addModelsAsFrames(to_visualize_pdb)\n",
        "  multichain_style = {'cartoon': {'colorscheme': 'chain'}}\n",
        "  multichain_view.setStyle({'model': -1}, multichain_style)\n",
        "  multichain_view.zoomTo()\n",
        "  multichain_view.show()\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': {'colorscheme': {'prop': 'b', 'map': color_map}}}\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",
        "\n",
        "  # Display lines at chain boundaries.\n",
        "  best_unrelaxed_prot = unrelaxed_proteins[best_model_name]\n",
        "  total_num_res = best_unrelaxed_prot.residue_index.shape[-1]\n",
        "  chain_ids = best_unrelaxed_prot.chain_index\n",
        "  for chain_boundary in np.nonzero(chain_ids[:-1] - chain_ids[1:]):\n",
        "    if chain_boundary.size:\n",
        "      plt.plot([0, total_num_res], [chain_boundary, chain_boundary], color='red')\n",
        "      plt.plot([chain_boundary, chain_boundary], [0, total_num_res], color='red')\n",
        "\n",
        "  plt.title('Predicted Aligned Error')\n",
        "  plt.xlabel('Scored residue')\n",
        "  plt.ylabel('Aligned residue')\n",
        "\n",
        "# Save the 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",
        "  pae_data = notebook_utils.get_pae_json(pae=pae, max_pae=max_pae.item())\n",
        "  with open(pae_output_path, 'w') as f:\n",
        "    f.write(pae_data)\n",
        "\n",
        "# --- Download the predictions ---\n",
        "!zip -q -r {output_dir}.zip {output_dir}\n",
        "files.download(f'{output_dir}.zip')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 6. Get predicted biophysical features using the Bio2Byte tools \n",
        "\n",
        "#@markdown Once this cell has been executed, a prediction dictionary will be\n",
        "#@markdown generated containing all the predicted features.\n",
        "\n",
        "#@markdown ## About bio2Byte\n",
        "\n",
        "#@markdown <img src=\"https://pbs.twimg.com/profile_images/1247824923546079232/B9b_Yg7n.jpg\" alt=\"b2b\" width=\"48\"/>\n",
        "#@markdown <img src=\"https://upload.wikimedia.org/wikipedia/commons/thumb/3/31/Vrije_Universiteit_Brussel_logo.svg/1200px-Vrije_Universiteit_Brussel_logo.svg.png\" alt=\"vub\" width=\"120\"/>\n",
        "\n",
        "#@markdown We investigate how the dynamics, conformational states and available experimental data of proteins relates to their amino acid sequence. Underlying physical and chemical principles are computationally unravelled through data integration, analysis and machine learning, so connecting them to biological events and improving our understanding of the way proteins work.\n",
        "#@markdown Online predictors are available on https://bio2byte.be/b2btools/\n",
        "\n",
        "#@markdown ### Tools and predictors\n",
        "\n",
        "#@markdown * **DynaMine** is a fast predictor of protein backbone dynamics using only sequence information as input. The version here also predicts side-chain dynamics and secondary structure predictors using the same principle. This tool returns the following predictions: \n",
        "#@markdown   * DynaMine backbone dynamics\n",
        "#@markdown   * DynaMine side chain dynamics\n",
        "#@markdown   * DynaMine conformational propensities (polyproline II)\n",
        "#@markdown   * DynaMine conformational propensities (coil)\n",
        "#@markdown   * DynaMine conformational propensities (sheet)\n",
        "#@markdown   * DynaMine conformational propensities (helix)\n",
        "\n",
        "#@markdown * **EFoldMine** is a method that predicts from the primary amino acid sequence of a protein, which amino acids are likely involved in early folding events. This tool returns the following predictions:\n",
        "#@markdown   * EFoldMine earlyFolding propensity\n",
        "\n",
        "#@markdown * **DisoMine** method predicts protein disorder with recurrent neural networks not directly from the amino acid sequence, but instead from more generic predictions of key biophysical properties, here protein dynamics, secondary structure and early folding. This tool returns the following predictions:\n",
        "#@markdown   * Disomine disorder\n",
        "\n",
        "#@markdown * **AgMata** is a single-sequence based predictor of protein regions that are likely to cause beta-aggregation. This tool returns the following predictions:\n",
        "#@markdown   * Agmata aggregation propensity\n",
        "\n",
        "from b2bTools import SingleSeq\n",
        "with open(\"/content/input.fasta\", 'w') as fasta_file:\n",
        "  for sequence, number in zip(sequences, range(1, 9)):\n",
        "    if sequence != '':\n",
        "      print(\">sequence_{}\\n\".format(number))\n",
        "      fasta_file.write(f\">sequence_{number}\\n{sequence}\\n\")\n",
        "\n",
        "single_seq = SingleSeq(\"/content/input.fasta\")\n",
        "single_seq.predict(tools=['dynamine', 'agmata', 'efoldmine', 'disomine'])\n",
        "\n",
        "predictions = single_seq.get_all_predictions()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "yjUM-f_lykZ9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 7. Plot predicted biophysical features using the Bio2Byte tools \n",
        "\n",
        "#@markdown Once this cell has been executed, a plot will be rendered\n",
        "#@markdown using the library called MatPlotLib\n",
        "\n",
        "from matplotlib import pyplot as plt\n",
        "print(predictions.keys())\n",
        "for sequence_key in predictions.keys():\n",
        "  sequence = predictions[sequence_key]\n",
        "\n",
        "  predictions_values_keys = (x for x in sequence.keys() if x not in ['seq'])\n",
        "  residues = sequence['seq']\n",
        "  residues_count = len(residues)\n",
        "  fig, (ax1, ax2) = plt.subplots(1, 2, sharex='col', figsize=(30, 5))\n",
        "\n",
        "  for prediction_value_key in predictions_values_keys:\n",
        "    prediction_value = sequence[prediction_value_key]\n",
        "    if (prediction_value_key != 'agmata'):\n",
        "      ax1.plot(prediction_value, label=prediction_value_key)\n",
        "    else:\n",
        "      ax2.plot(prediction_value, label=prediction_value_key)\n",
        "\n",
        "  ax1.legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
        "  ax1.set_xlabel('Residues')  \n",
        "  ax1.set_ylabel('Values')\n",
        "  ax1.set_title(\"Bio2Byte biophysical predictions ({})\".format(sequence_key))\n",
        "\n",
        "  ax2.set_xlabel('Residues')  \n",
        "  ax2.set_ylabel('Values')\n",
        "  ax2.set_title(\"Bio2Byte biophysical AgMata ({})\".format(sequence_key))\n",
        "\n",
        "  \n",
        "  ax1.set_xticks(range(residues_count))\n",
        "  ax1.set_xticklabels(residues)\n",
        "  ax1.xaxis.grid(True)\n",
        "\n",
        "  ax2.set_xticks(range(residues_count))\n",
        "  ax2.set_xticklabels(residues)\n",
        "  ax2.xaxis.grid(True)\n",
        "\n",
        "\n",
        "  ax1.set_ylim([-0.1, 1.1])\n",
        "  plt.show()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "OBLe0pTb9_BL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Plot 3D structures colored by biophysical predictions\n",
        "#@markdown Please execute this cell by pressing the _Play_ button \n",
        "\n",
        "import numpy as np\n",
        "\n",
        "VALUES_CONSTANTS = {\n",
        "    \"backbone\": {\n",
        "        \"colors\": [\n",
        "            (0, 0.69, \"#FF7D45\"),\n",
        "            (0.69, 0.80, \"#FFDB13\"),\n",
        "            (0.80, 1.00, \"#65CBF3\"),\n",
        "            (1.00, 2.00, \"#0053D6\"),\n",
        "        ],\n",
        "        \"legends\": [\n",
        "            \"Flexible region (value < 0.69)\",\n",
        "            \"Context dependent region (0.69 < value < 0.80)\",\n",
        "            \"Rigid region (0.80 < value < 1.00)\",\n",
        "            \"Membrane spanning region (value > 1.00)\",\n",
        "        ],\n",
        "        \"title\": \"DynaMine backbone dynamics\",\n",
        "    },\n",
        "    \"sidechain\": {\n",
        "        \"colors\": [\n",
        "            (0, 0.69, \"#FF7D45\"),\n",
        "            (0.69, 0.80, \"#FFDB13\"),\n",
        "            (0.80, 1.00, \"#65CBF3\"),\n",
        "            (1.00, 2.00, \"#0053D6\"),\n",
        "        ],\n",
        "        \"legends\": [\n",
        "            \"Flexible region (value < 0.69)\",\n",
        "            \"Context dependent region (0.69 < value < 0.80)\",\n",
        "            \"Rigid region (0.80 < value < 1.00)\",\n",
        "            \"Membrane spanning region (value > 1.00)\",\n",
        "        ],\n",
        "        \"title\": \"DynaMine sidechain dynamics\",\n",
        "    },\n",
        "    \"ppII\": {\n",
        "        \"colors\": [   \n",
        "            (0.00, 0.25, \"#ff0100\"),\n",
        "            (0.25, 0.50, \"#ffff00\"),\n",
        "            (0.50, 0.75, \"#7fff00\"),\n",
        "            (0.75, 1.00, \"#01ff00\"),\n",
        "        ],\n",
        "        \"legends\": [\n",
        "            \"Lowest propensity (value = 0.00)\",\n",
        "            \"Low propensity (0.00 < value < 0.25)\",\n",
        "            \"Mid propensity (0.25 < value < 0.50)\",\n",
        "            \"High propensity (0.50 < value < 0.75)\",\n",
        "            \"Very high propensity(0.75 < value < 1.00)\",\n",
        "        ],\n",
        "        \"title\": \"DynaMine conformational propensities (polyproline II)\",\n",
        "    },\n",
        "    \"coil\": {\n",
        "        \"colors\": [   \n",
        "            (0.00, 0.25, \"#ff0100\"),\n",
        "            (0.25, 0.50, \"#ffff00\"),\n",
        "            (0.50, 0.75, \"#7fff00\"),\n",
        "            (0.75, 1.00, \"#01ff00\"),\n",
        "        ],\n",
        "        \"legends\": [\n",
        "            \"Lowest propensity (value = 0.00)\",\n",
        "            \"Low propensity (0.00 < value < 0.25)\",\n",
        "            \"Mid propensity (0.25 < value < 0.50)\",\n",
        "            \"High propensity (0.50 < value < 0.75)\",\n",
        "            \"Very high propensity(0.75 < value < 1.00)\",\n",
        "        ],\n",
        "        \"title\": \"DynaMine conformational propensities (coil)\",\n",
        "    },\n",
        "    \"sheet\": {\n",
        "        \"colors\": [   \n",
        "            (0.00, 0.25, \"#ff0100\"),\n",
        "            (0.25, 0.50, \"#ffff00\"),\n",
        "            (0.50, 0.75, \"#7fff00\"),\n",
        "            (0.75, 1.00, \"#01ff00\"),\n",
        "        ],\n",
        "        \"legends\": [\n",
        "            \"Lowest propensity (value = 0.00)\",\n",
        "            \"Low propensity (0.00 < value < 0.25)\",\n",
        "            \"Mid propensity (0.25 < value < 0.50)\",\n",
        "            \"High propensity (0.50 < value < 0.75)\",\n",
        "            \"Very high propensity(0.75 < value < 1.00)\",\n",
        "        ],\n",
        "        \"title\": \"DynaMine conformational propensities (sheet)\",\n",
        "    },\n",
        "    \"helix\": {\n",
        "        \"colors\": [   \n",
        "            (0.00, 0.25, \"#ff0100\"),\n",
        "            (0.25, 0.50, \"#ffff00\"),\n",
        "            (0.50, 0.75, \"#7fff00\"),\n",
        "            (0.75, 1.00, \"#01ff00\"),\n",
        "        ],\n",
        "        \"legends\": [\n",
        "            \"Lowest propensity (value = 0.00)\",\n",
        "            \"Low propensity (0.00 < value < 0.25)\",\n",
        "            \"Mid propensity (0.25 < value < 0.50)\",\n",
        "            \"High propensity (0.50 < value < 0.75)\",\n",
        "            \"Very high propensity(0.75 < value < 1.00)\",\n",
        "        ],\n",
        "        \"title\": \"DynaMine conformational propensities (helix)\",\n",
        "        \"subtitle\": \"DynaMine is a fast predictor of protein backbone dynamics using only sequence information as input. The version here also predicts side-chain dynamics and secondary structure predictors using the same principle.\",\n",
        "    },\n",
        "    \"earlyFolding\": {\n",
        "        \"colors\": [\n",
        "            (0, 0.169, \"#FF7D45\"),\n",
        "            (0.169 ,2, \"#0053D6\"),\n",
        "        ],\n",
        "        \"legends\": [\n",
        "            \"Not likely to start the protein folding process (value < 0.169)\",\n",
        "            \"Likely to start the protein folding process (value > 0.169)\",\n",
        "        ],\n",
        "        \"title\": \"EFoldMine earlyFolding propensity\",\n",
        "    },\n",
        "    \"disoMine\": {\n",
        "        \"colors\": [\n",
        "            (0, 0.50, \"#FF7D45\"),\n",
        "            (0.50, 1.00, \"#0053D6\"),\n",
        "        ],\n",
        "        \"legends\": [\n",
        "            \"Not likely a disordered residue (value < 0.50)\",\n",
        "            \"Likely a disordered residue (value > 0.50)\",\n",
        "        ],\n",
        "        \"title\": \"Disomine disorder\",\n",
        "    },\n",
        "    \"agmata\": {\n",
        "        \"colors\": [\n",
        "            (0.00, 0.25, \"#ff0100\"),\n",
        "            (0.50, 0.75, \"#ffff00\"),\n",
        "            (0.75, 1.00, \"#01ff00\"),\n",
        "        ],\n",
        "        \"legends\": [\n",
        "            \"Not likely to be involved in beta-sheet aggregation (Trough)\",\n",
        "            \"Likely to be involved in beta-sheet aggregation (Mid Peak)\",\n",
        "            \"Most likely to be involved in beta-sheet aggregation (Higher Peak)\",\n",
        "        ],\n",
        "        \"title\": \"Agmata aggregation propensity\",\n",
        "    },\n",
        "}\n",
        "\n",
        "B2B_TOOLS = [\"dynamine\", \"disomine\", \"efoldmine\", \"agmata\"]\n",
        "\n",
        "B2B_TOOLS_RESULTS = np.array(\n",
        "    [\n",
        "        [\"backbone\", \"sidechain\", \"ppII\"],\n",
        "        [\"coil\", \"sheet\", \"helix\"],\n",
        "        [\"earlyFolding\", \"disoMine\", \"agmata\"],\n",
        "    ]\n",
        ")\n",
        "\n",
        "# PLDDT_COLOR_MAP = {i: bands[2] for i, bands in enumerate(VALUES_CONSTANTS[\"pLDDT\"][\"colors\"])}\n",
        "# PLDDT_DEFAULT_ATOM_STYLE = {\"cartoon\": {\"colorscheme\": {\"prop\": \"b\", \"map\": PLDDT_COLOR_MAP}}}\n",
        "\n",
        "class Atom(dict):\n",
        "    def __init__(self, line):\n",
        "        self[\"type\"] = line[0:6].strip()\n",
        "        self[\"idx\"] = line[6:11].strip()\n",
        "        self[\"name\"] = line[12:16].strip()\n",
        "        self[\"resname\"] = line[17:20].strip()\n",
        "        self[\"resid\"] = int(int(line[22:26]))\n",
        "        self[\"x\"] = float(line[30:38])\n",
        "        self[\"y\"] = float(line[38:46])\n",
        "        self[\"z\"] = float(line[46:54])\n",
        "        self[\"sym\"] = line[76:78].strip()\n",
        "\n",
        "    def __str__(self):\n",
        "        line = list(\" \" * 80)\n",
        "\n",
        "        line[0:6] = self[\"type\"].ljust(6)\n",
        "        line[6:11] = self[\"idx\"].ljust(5)\n",
        "        line[12:16] = self[\"name\"].ljust(4)\n",
        "        line[17:20] = self[\"resname\"].ljust(3)\n",
        "        line[22:26] = str(self[\"resid\"]).ljust(4)\n",
        "        line[30:38] = str(self[\"x\"]).rjust(8)\n",
        "        line[38:46] = str(self[\"y\"]).rjust(8)\n",
        "        line[46:54] = str(self[\"z\"]).rjust(8)\n",
        "        line[76:78] = self[\"sym\"].rjust(2)\n",
        "        return \"\".join(line) + \"\\n\"\n",
        "\n",
        "\n",
        "class Molecule(list):\n",
        "    def __init__(self, file):\n",
        "        for line in file:\n",
        "            if \"ATOM\" in line or \"HETATM\" in line:\n",
        "                self.append(Atom(line))\n",
        "\n",
        "    def __str__(self):\n",
        "        outstr = \"\"\n",
        "        for at in self:\n",
        "            outstr += str(at)\n",
        "\n",
        "        return outstr\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "import tqdm.notebook\n",
        "\n",
        "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
        "\n",
        "def get_value_color(result_key, value, min_value=0, max_value=1):\n",
        "    if (result_key == \"ppII\" \n",
        "        or result_key == \"coil\" \n",
        "        or result_key == \"sheet\" \n",
        "        or result_key == \"helix\" \n",
        "        or result_key == \"agmata\"):\n",
        "        return get_spectrum_color(value, min_value, max_value)\n",
        "    else:\n",
        "        for band in VALUES_CONSTANTS[result_key][\"colors\"]:\n",
        "            band_floor, band_ceiling, band_color = band\n",
        "            if value > band_floor and value <= band_ceiling:\n",
        "                return band_color\n",
        "\n",
        "def get_spectrum_color(real, min_value, max_value):\n",
        "    color_r = 0\n",
        "    color_g = 0\n",
        "    color_b = 0\n",
        "\n",
        "    value = real\n",
        "    base = max_value - min_value\n",
        "    if (value < min_value):\n",
        "        value = min_value \n",
        "    elif value > max_value:\n",
        "        value = max_value\n",
        "\n",
        "    if base == 0:\n",
        "        value = 100\n",
        "    else:\n",
        "        value = (value - min_value) / base * 100\n",
        "    \n",
        "    if value < 50:\n",
        "        color_r = 255\n",
        "        off_g = 5.1 * value\n",
        "        color_g = max(1, off_g)\n",
        "    else:\n",
        "        color_g = 255\n",
        "        off_r = 510 - 5.1 * value\n",
        "        color_r = max(1, off_r)\n",
        "  \n",
        "    return '#%02x%02x%02x' % (int(color_r), int(color_g), int(color_b))\n",
        "\n",
        "def create_py3dmol_view(molecule, color_key):\n",
        "    view = py3Dmol.view(width=800, height=600)\n",
        "    view.setBackgroundColor(\"white\")\n",
        "    view.addModelsAsFrames(str(molecule))\n",
        "\n",
        "    default_color_style = {\"cartoon\": {\"color\": \"black\"}}\n",
        "    for atom_index, atom in enumerate(molecule):\n",
        "        view.setStyle(\n",
        "            {\"model\": -1, \"serial\": atom_index + 1}, \n",
        "            atom.get(color_key, default_color_style)\n",
        "        )\n",
        "\n",
        "    view.zoomTo()\n",
        "    return view\n",
        "\n",
        "\n",
        "def plot_result_legend(result_key):\n",
        "    \"\"\"Plots the legend for a predicted result.\"\"\"\n",
        "\n",
        "    value_constants = VALUES_CONSTANTS[result_key]\n",
        "    colors = [x[2] for x in value_constants[\"colors\"]]\n",
        "\n",
        "    plt.figure(figsize=(2, 2))\n",
        "    for c in colors:\n",
        "        plt.bar(0, 0, color=c)\n",
        "\n",
        "    plt.legend(value_constants[\"legends\"], frameon=False, loc=\"center\", fontsize=14)\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.suptitle(value_constants[\"title\"], fontsize=18)\n",
        "    plt.title(value_constants[\"title\"], fontsize=16, pad=16)\n",
        "\n",
        "    return plt\n",
        "\n",
        "def plot_structure(relaxed_pdb, b2b_predictions_results):\n",
        "    with tqdm.notebook.tqdm(total=100, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "        b2b_agmata_result = b2b_predictions_results['agmata']\n",
        "        agmata_min_value = min(b2b_agmata_result)\n",
        "        agmata_max_value = max(b2b_agmata_result)\n",
        "\n",
        "        pbar.update(2.5)\n",
        "\n",
        "        molecule = Molecule(relaxed_pdb)\n",
        "    \n",
        "        for atom in molecule:\n",
        "            res_id = int(atom[\"resid\"]) - 1\n",
        "            for result_key in VALUES_CONSTANTS.keys():\n",
        "                atom[result_key] = b2b_predictions_results[result_key][res_id]\n",
        "\n",
        "                if result_key == \"agmata\":\n",
        "                    atom_color = get_value_color(\n",
        "                        result_key, \n",
        "                        atom[result_key], \n",
        "                        min_value=agmata_min_value, \n",
        "                        max_value=agmata_max_value\n",
        "                    )\n",
        "                else:\n",
        "                    atom_color = get_value_color(result_key, atom[result_key])\n",
        "\n",
        "                atom[\"{0}_color\".format(result_key)] = {\n",
        "                    \"cartoon\": {\n",
        "                        \"color\": atom_color\n",
        "                      }\n",
        "                }\n",
        "        \n",
        "        pbar.update(2.5)\n",
        "        grid = GridspecLayout(len(np.matrix.flatten(B2B_TOOLS_RESULTS)) + 1, 2)\n",
        "\n",
        "        pbar.update(5)\n",
        "\n",
        "        predictors_count = B2B_TOOLS_RESULTS.shape[0] * B2B_TOOLS_RESULTS.shape[1]\n",
        "        for index, result_key in enumerate(np.matrix.flatten(B2B_TOOLS_RESULTS)):\n",
        "            output_plot = Output()\n",
        "            with output_plot:\n",
        "                create_py3dmol_view(\n",
        "                    molecule, \n",
        "                    \"{0}_color\".format(result_key)\n",
        "                ).show()\n",
        "\n",
        "            grid[index + 1, 0] = output_plot\n",
        "\n",
        "            output_legend = Output()\n",
        "            with output_legend:\n",
        "                plot_result_legend(result_key).show()\n",
        "\n",
        "            grid[index + 1, 1] = output_legend\n",
        "            pbar.update(90 / predictors_count)\n",
        "\n",
        "        display.display(grid)\n",
        "\n",
        "for sequence_key in predictions.keys():\n",
        "  sequence_predictions = predictions[sequence_key]\n",
        "  print(f\"3D STRUCTURES FOR {sequence_key}\")\n",
        "  plot_structure(relaxed_pdb.split(\"\\n\"), sequence_predictions)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "MiT3nTOoFOHi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lUQAn5LYC5n4"
      },
      "source": [
        "### Interpreting the prediction\n",
        "\n",
        "In general predicted LDDT (pLDDT) is best used for intra-domain confidence, whereas Predicted Aligned Error (PAE) is best used for determining between domain or between chain confidence.\n",
        "\n",
        "Please see the [AlphaFold methods paper](https://www.nature.com/articles/s41586-021-03819-2), the [AlphaFold predictions of the human proteome paper](https://www.nature.com/articles/s41586-021-03828-1), and the [AlphaFold-Multimer paper](https://www.biorxiv.org/content/10.1101/2021.10.04.463034v1) as well as [our FAQ](https://alphafold.ebi.ac.uk/faq) on how to interpret AlphaFold 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._ Make sure you run all 5 cells in order.\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 AlphaFold 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 AlphaFold 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",
        "    *   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, try resetting your Colab runtime via _Runtime_ > _Factory reset runtime_.\n",
        "*   How does this compare to the open-source version of AlphaFold?\n",
        "    *   This Colab version of AlphaFold 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 of AlphaFold that is 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 is available via the inference script).\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",
        "*   I received an error “ModuleNotFoundError: No module named ...”, even though I ran the cell that imports it, what do I do?\n",
        "    *   Colab notebooks on the free tier time out after a certain amount of time. See the [Colab FAQ](https://research.google.com/colaboratory/faq.html#idle-timeouts). Try rerunning the whole notebook from the beginning.\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/deepmind/alphafold/issues) on Github.\n",
        "\n",
        "\n",
        "## Related work\n",
        "\n",
        "Take a look at these Colab notebooks provided by the community (please note that these notebooks may vary from our validated AlphaFold system and we cannot guarantee their accuracy):\n",
        "\n",
        "*   The [ColabFold AlphaFold2 notebook](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb) by Sergey Ovchinnikov, Milot Mirdita and Martin Steinegger, which uses an API hosted at the Södinglab based on the MMseqs2 server ([Mirdita et al. 2019, Bioinformatics](https://academic.oup.com/bioinformatics/article/35/16/2856/5280135)) for the multiple sequence alignment creation.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YfPhvYgKC81B"
      },
      "source": [
        "# Alphafold License and Disclaimer\n",
        "\n",
        "This is not an officially-supported Google product.\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",
        "Copyright 2021 DeepMind Technologies Limited.\n",
        "\n",
        "\n",
        "## AlphaFold Code License\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",
        "The 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 the [Acknowledgements section](https://github.com/deepmind/alphafold/#acknowledgements) in the AlphaFold README 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",
        "* UniProt: 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",
        "* 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.\n",
        "\n",
        "# Bio2Byte\n",
        "\n",
        "**Important note:** If you need to run these analysis again, please restart the Runtime from the menu \"Runtime\"\n",
        "\n",
        "Download the Bio2Byte's tools package from our PyPi repository: https://pypi.org/project/b2bTools/.\n",
        "\n",
        "## Citations\n",
        "\n",
        "DynaMine:\n",
        "> Elisa Cilia, Rita Pancsa, Peter Tompa, Tom Lenaerts, and Wim Vranken From protein sequence to dynamics and disorder with DynaMine Nature Communications 4:2741(2013)\n",
        "doi: 10.1038/ncomms3741\n",
        "\n",
        "EFoldMine:\n",
        "> Raimondi, D., Orlando, G., Pancsa, R. et al. Exploring the Sequence-based Prediction of Folding Initiation Sites in Proteins. Sci Rep 7, 8826 (2017).\n",
        "doi: https://doi.org/10.1038/s41598-017-08366-3\n",
        "\n",
        "Disomine:\n",
        "> Gabriele Orlando, Daniele Raimondi, Francesco Codice, Francesco Tabaro, Wim Vranken bioRxiv 2020.05.25.115253;\n",
        "doi: https://doi.org/10.1101/2020.05.25.115253\n",
        "\n",
        "AgMata:\n",
        "> O.G, S.A, MR.S, R.D, V.W.; Accurate prediction of protein beta-aggregation with generalized statistical potentials. doi: 10.1093/bioinformatics/btz912."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "uC1dKAwk2eyl",
        "W4JpOs6oA-QS"
      ],
      "name": "AlphaFold + B2B Tools.ipynb",
      "private_outputs": true,
      "provenance": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}