AlphaFold2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "AlphaFold2PredictStructure.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G4yBrceuFbf3"
      },
      "source": [
        "#Protein structure prediction with AlphaFold2 and MMseqs2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UGUBLzB3C6WN"
      },
      "source": [
        "Easy to use version of AlphaFold 2 (Jumper et al. 2021, Nature) using an API hosted at the Södinglab based on the MMseqs2 server (Mirdita et al. 2019, Bioinformatics) for the multiple sequence alignment creation.\n",
        "\n",
        "**Quickstart**\n",
        "1. Change the runtime type to GPU at \"Runtime\" -> \"Change runtime type\" (improves speed).\n",
        "2. Paste your protein sequence in the input field below.\n",
        "3. Press \"Runtime\" -> \"Run all\".\n",
        "4. The pipeline consists of 10 steps. The currently running steps is indicated by a circle with a stop sign next to it.\n",
        "\n",
        "**Result**\n",
        "\n",
        "1. PDB formated structures sorted by avg. pIDDT. (relaxed, unrelaxed).\n",
        "2. Plots of the model quality.\n",
        "3. Plots of the msa coverage.\n",
        "4. Parameter log file.\n",
        "5. A3M formatted input MSA.\n",
        "\n",
        "At the end of the job a download modal box will popup with a `jobname.result.zip` file.\n",
        "\n",
        "**Using a custom MSA as input**\n",
        "\n",
        "To predict the structure with a custom MSA (A3M formatted) (1) Past the input sequence into the \"query_sequence\" field. (2) Change the msa_mode: to \"custom\", (3) \"Runtime\" -> \"Run all\". (4) A upload box will appear at the end of the \"Input Protein ...\" box. Upload your A3M. \n",
        "\n",
        "To generate good input MSAs HHblits can be used here: https://toolkit.tuebingen.mpg.de/tools/hhblits\n",
        "Hit submit you query -> click \"Query Template MSA\" -> \"Download Full A3M\".\n",
        "Download the a3m file and upload it to the notebook!\n",
        "\n",
        "**Troubleshooting**\n",
        "* Try to restart the session \"Runntime\" -> \"Factory reset runtime\".\n",
        "* Check your input sequence.\n",
        "\n",
        "**Known issues**\n",
        "* Colab assigns different type of GPUs, some might have not enough memory to predict the structure.\n",
        "* Browser can block the download of the result.\n",
        "\n",
        "**Limitations**\n",
        "* MSAs: MMseqs2 is very precise and sensitive but might find less hits compared to HHblits/HMMer searched against BFD or Mgnify.\n",
        "* Computing resources: MMseqs2 can probably handle >20k requests per day since we run it only on 16 cores.\n",
        "\n",
        "For best results, we recommend using the full pipeline: https://github.com/deepmind/alphafold\n",
        "\n",
        "**Acknowledgements**\n",
        "\n",
        "Most of the python code was written by Sergey Ovchinnikov (@sokrypton). The API is hosted at the Södinglab (@SoedingL) and maintained by Milot Mirdita (@milot_mirdita). Martin Steinegger (@thesteinegger) integrated everything. Also credit to [David Koes](https://github.com/dkoes) for his awesome [py3Dmol](https://3dmol.csb.pitt.edu/) plugin, without whom these notebooks would be quite boring!\n",
        "\n",
        "For related notebooks see: [ColabFold](https://github.com/sokrypton/ColabFold)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kOblAo-xetgx",
        "cellView": "form"
      },
      "source": [
        "#@title Input protein sequence here before you \"Run all\"\n",
        "from google.colab import files\n",
        "import os\n",
        "import os.path\n",
        "import re\n",
        "\n",
        "query_sequence = 'MAKTIKITQTRSAIGRLPKHKATLLGLGLRRIGHTVEREDTPAIRGMINAVSFMVKVEE' #@param {type:\"string\"}\n",
        "# remove whitespaces\n",
        "query_sequence = \"\".join(query_sequence.split())\n",
        "query_sequence = re.sub(r'[^a-zA-Z]','', query_sequence).upper()\n",
        "\n",
        "jobname = 'RL30_ECOLI' #@param {type:\"string\"}\n",
        "# remove whitespaces\n",
        "jobname = \"\".join(jobname.split())\n",
        "jobname = re.sub(r'\\W+', '', jobname)\n",
        "\n",
        "with open(f\"{jobname}.fasta\", \"w\") as text_file:\n",
        "    text_file.write(\">1\\n%s\" % query_sequence)\n",
        "\n",
        "# number of models to use\n",
        "#@markdown ---\n",
        "#@markdown ### Advanced settings\n",
        "num_models = 5 #@param [1,2,3,4,5] {type:\"raw\"}\n",
        "msa_mode = \"MMseqs2\" #@param [\"MMseqs2\",\"single_sequence\",\"custom\"]\n",
        "use_msa = True if msa_mode == \"MMseqs2\" else False\n",
        "use_custom_msa = True if msa_mode == \"custom\" else False\n",
        "use_amber = False #@param {type:\"boolean\"}\n",
        "use_templates = False #@param {type:\"boolean\"}\n",
        "#@markdown ---\n",
        "\n",
        "with open(f\"{jobname}.log\", \"w\") as text_file:\n",
        "    text_file.write(\"num_models=%s\\n\" % num_models)\n",
        "    text_file.write(\"use_amber=%s\\n\" % use_amber)\n",
        "    text_file.write(\"use_msa=%s\\n\" % use_msa)\n",
        "    text_file.write(\"msa_mode=%s\\n\" % msa_mode)\n",
        "    text_file.write(\"use_templates=%s\\n\" % use_templates)\n",
        "\n",
        "# decide which a3m to use\n",
        "if use_msa:\n",
        "  a3m_file = f\"{jobname}.a3m\"\n",
        "elif use_custom_msa:\n",
        "  a3m_file = f\"{jobname}.custom.a3m\"\n",
        "  if not os.path.isfile(a3m_file):\n",
        "    custom_msa_dict = files.upload()\n",
        "    custom_msa = list(custom_msa_dict.keys())[0]\n",
        "    os.rename(custom_msa, a3m_file)\n",
        "    print(f\"moving {custom_msa} to {a3m_file}\")\n",
        "else:\n",
        "  a3m_file = f\"{jobname}.single_sequence.a3m\"\n",
        "  with open(a3m_file, \"w\") as text_file:\n",
        "    text_file.write(\">1\\n%s\" % query_sequence)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iccGdbe_Pmt9",
        "cellView": "form"
      },
      "source": [
        "#@title Install dependencies\n",
        "%%bash -s $use_amber $use_msa $use_templates\n",
        "\n",
        "USE_AMBER=$1\n",
        "USE_MSA=$2\n",
        "USE_TEMPLATES=$3\n",
        "\n",
        "if [ ! -f AF2_READY ]; then\n",
        "  # install dependencies\n",
        "  pip -q install biopython\n",
        "  pip -q install dm-haiku\n",
        "  pip -q install ml-collections\n",
        "  pip -q install py3Dmol\n",
        "\n",
        "  # download model\n",
        "  if [ ! -d \"alphafold/\" ]; then\n",
        "    git clone https://github.com/deepmind/alphafold.git --quiet\n",
        "    mv alphafold alphafold_\n",
        "    mv alphafold_/alphafold .\n",
        "    # remove \"END\" from PDBs, otherwise biopython complains\n",
        "    sed -i \"s/pdb_lines.append('END')//\" /content/alphafold/common/protein.py\n",
        "    sed -i \"s/pdb_lines.append('ENDMDL')//\" /content/alphafold/common/protein.py\n",
        "  fi\n",
        "\n",
        "  # download model params (~1 min)\n",
        "  if [ ! -d \"params/\" ]; then\n",
        "    wget -qnc https://storage.googleapis.com/alphafold/alphafold_params_2021-07-14.tar\n",
        "    mkdir params\n",
        "    tar -xf alphafold_params_2021-07-14.tar -C params/\n",
        "    rm alphafold_params_2021-07-14.tar\n",
        "  fi\n",
        "  touch AF2_READY\n",
        "fi\n",
        "# download libraries for interfacing with MMseqs2 API\n",
        "if [ ${USE_MSA} == \"True\" ] || [ ${USE_TEMPLATES} == \"True\" ]; then\n",
        "  if [ ! -f MMSEQ2_READY ]; then\n",
        "    apt-get -qq -y update 2>&1 1>/dev/null\n",
        "    apt-get -qq -y install jq curl zlib1g gawk 2>&1 1>/dev/null\n",
        "    touch MMSEQ2_READY\n",
        "  fi\n",
        "fi\n",
        "# setup conda\n",
        "if [ ${USE_AMBER} == \"True\" ] || [ ${USE_TEMPLATES} == \"True\" ]; then\n",
        "  if [ ! -f CONDA_READY ]; then\n",
        "    wget -qnc https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh\n",
        "    bash Miniconda3-latest-Linux-x86_64.sh -bfp /usr/local 2>&1 1>/dev/null\n",
        "    rm Miniconda3-latest-Linux-x86_64.sh\n",
        "    touch CONDA_READY\n",
        "  fi\n",
        "fi\n",
        "# setup template search\n",
        "if [ ${USE_TEMPLATES} == \"True\" ] && [ ! -f HH_READY ]; then\n",
        "  conda install -y -q -c conda-forge -c bioconda kalign3=3.2.2 hhsuite=3.3.0 python=3.7 2>&1 1>/dev/null\n",
        "  touch HH_READY\n",
        "fi\n",
        "# setup openmm for amber refinement\n",
        "if [ ${USE_AMBER} == \"True\" ] && [ ! -f AMBER_READY ]; then\n",
        "  conda install -y -q -c conda-forge openmm=7.5.1 python=3.7 pdbfixer 2>&1 1>/dev/null\n",
        "  (cd /usr/local/lib/python3.7/site-packages; patch -s -p0 < /content/alphafold_/docker/openmm.patch)\n",
        "  wget -qnc https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt\n",
        "  mv stereo_chemical_props.txt alphafold/common/\n",
        "  touch AMBER_READY\n",
        "fi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A9tUpDaikPC8",
        "cellView": "form"
      },
      "source": [
        "#@title Call MMseqs2 to get MSA/templates\n",
        "%%bash -s $use_amber $use_msa $use_templates $jobname\n",
        "USE_AMBER=$1\n",
        "USE_MSA=$2\n",
        "USE_TEMPLATES=$3\n",
        "NAME=$4\n",
        "if [ ${USE_MSA} == \"True\" ] || [ ${USE_TEMPLATES} == \"True\" ]; then\n",
        "  if [ ! -f ${NAME}.mmseqs2.tar.gz ]; then\n",
        "    # query MMseqs2 webserver\n",
        "    echo \"submitting job\"\n",
        "    ID=$(curl -s -F q=@${NAME}.fasta -F mode=all https://a3m.mmseqs.com/ticket/msa | jq -r '.id')\n",
        "    STATUS=$(curl -s https://a3m.mmseqs.com/ticket/${ID} | jq -r '.status')\n",
        "    while [ \"${STATUS}\" == \"RUNNING\" ]; do\n",
        "      STATUS=$(curl -s https://a3m.mmseqs.com/ticket/${ID} | jq -r '.status')\n",
        "      sleep 1\n",
        "    done\n",
        "    if [ \"${STATUS}\" == \"COMPLETE\" ]; then\n",
        "      curl -s https://a3m.mmseqs.com/result/download/${ID} > ${NAME}.mmseqs2.tar.gz\n",
        "      tar xzf ${NAME}.mmseqs2.tar.gz\n",
        "      tr -d '\\000' < uniref.a3m > ${NAME}.a3m\n",
        "      rm uniref.a3m\n",
        "      mv pdb70.m8 ${NAME}.m8\n",
        "    else\n",
        "      echo \"MMseqs2 server did not return a valid result.\"\n",
        "      exit 1\n",
        "    fi\n",
        "  fi\n",
        "  if [ ${USE_MSA} == \"True\" ]; then\n",
        "    echo \"Found $(grep -c \">\" ${NAME}.a3m) sequences (after redundacy filtering)\"\n",
        "  fi\n",
        "  if [ ${USE_TEMPLATES} == \"True\" ] && [ ! -f ${NAME}_hhm.ffindex ]; then\n",
        "    echo \"getting templates\"\n",
        "    if [ ! -d templates ]; then\n",
        "      mkdir templates/\n",
        "    fi\n",
        "    printf \"pdb\\tevalue\\n\"\n",
        "    head -n 20 ${NAME}.m8 | awk '{print $2\"\\t\"$11}'\n",
        "    TMPL=$(head -n 20 ${NAME}.m8 | awk '{printf $2\",\"}')\n",
        "    curl -s https://a3m-templates.mmseqs.com/template/${TMPL} | tar xzf - -C templates/\n",
        "    mv templates/pdb70_a3m.ffdata ${NAME}_a3m.ffdata\n",
        "    mv templates/pdb70_a3m.ffindex ${NAME}_a3m.ffindex\n",
        "    mv templates/pdb70_hhm.ffdata ${NAME}_hhm.ffdata\n",
        "    mv templates/pdb70_hhm.ffindex ${NAME}_hhm.ffindex\n",
        "    cp ${NAME}_a3m.ffindex ${NAME}_cs219.ffindex\n",
        "    touch ${NAME}_cs219.ffdata\n",
        "  fi\n",
        "fi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JPWfhGssZdTb",
        "cellView": "form"
      },
      "source": [
        "#@title Import libraries and setup model\n",
        "# the following code is written by Sergey Ovchinnikov\n",
        "# setup the model\n",
        "if \"model\" not in dir():\n",
        "\n",
        "  # hiding warning messages\n",
        "  import warnings\n",
        "  from absl import logging\n",
        "  import os\n",
        "  import tensorflow as tf\n",
        "  warnings.filterwarnings('ignore')\n",
        "  logging.set_verbosity(\"error\")\n",
        "  os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'\n",
        "  tf.get_logger().setLevel('ERROR')\n",
        "\n",
        "  import sys\n",
        "  import numpy as np\n",
        "  import pickle\n",
        "  from alphafold.common import protein\n",
        "  from alphafold.data import pipeline\n",
        "  from alphafold.data import templates\n",
        "  from alphafold.model import data\n",
        "  from alphafold.model import config\n",
        "  from alphafold.model import model\n",
        "  from alphafold.data.tools import hhsearch\n",
        "\n",
        "  # plotting libraries\n",
        "  import py3Dmol\n",
        "  import matplotlib.pyplot as plt\n",
        "  import ipywidgets\n",
        "  from ipywidgets import interact, fixed\n",
        "\n",
        "if use_amber and \"relax\" not in dir():\n",
        "  sys.path.insert(0, '/usr/local/lib/python3.7/site-packages/')\n",
        "  from alphafold.relax import relax\n",
        "\n",
        "if \"model_params\" not in dir(): model_params = {}\n",
        "for model_name in [\"model_1\",\"model_2\",\"model_3\",\"model_4\",\"model_5\"][:num_models]:\n",
        "  if model_name not in model_params:\n",
        "    model_config = config.model_config(model_name)\n",
        "    model_config.data.eval.num_ensemble = 1\n",
        "    model_params[model_name] = data.get_model_haiku_params(model_name=model_name, data_dir=\".\")\n",
        "    if model_name == \"model_1\":\n",
        "      model_runner_1 = model.RunModel(model_config, model_params[model_name])\n",
        "    if model_name == \"model_3\":\n",
        "      model_runner_3 = model.RunModel(model_config, model_params[model_name])\n",
        "\n",
        "def mk_mock_template(query_sequence):\n",
        "  # since alphafold's model requires a template input\n",
        "  # we create a blank example w/ zero input, confidence -1\n",
        "  ln = len(query_sequence)\n",
        "  output_templates_sequence = \"-\"*ln\n",
        "  output_confidence_scores = np.full(ln,-1)\n",
        "  templates_all_atom_positions = np.zeros((ln, templates.residue_constants.atom_type_num, 3))\n",
        "  templates_all_atom_masks = np.zeros((ln, templates.residue_constants.atom_type_num))\n",
        "  templates_aatype = templates.residue_constants.sequence_to_onehot(output_templates_sequence,\n",
        "                                                                    templates.residue_constants.HHBLITS_AA_TO_ID)\n",
        "  template_features = {'template_all_atom_positions': templates_all_atom_positions[None],\n",
        "                       'template_all_atom_masks': templates_all_atom_masks[None],\n",
        "                       'template_sequence': [f'none'.encode()],\n",
        "                       'template_aatype': np.array(templates_aatype)[None],\n",
        "                       'template_confidence_scores': output_confidence_scores[None],\n",
        "                       'template_domain_names': [f'none'.encode()],\n",
        "                       'template_release_date': [f'none'.encode()]}\n",
        "  return template_features\n",
        "\n",
        "def mk_template(jobname):\n",
        "  template_featurizer = templates.TemplateHitFeaturizer(\n",
        "      mmcif_dir=\"templates/\",\n",
        "      max_template_date=\"2100-01-01\",\n",
        "      max_hits=20,\n",
        "      kalign_binary_path=\"kalign\",\n",
        "      release_dates_path=None,\n",
        "      obsolete_pdbs_path=None)\n",
        "\n",
        "  hhsearch_pdb70_runner = hhsearch.HHSearch(binary_path=\"hhsearch\",databases=[jobname])\n",
        "\n",
        "  a3m_lines = \"\\n\".join(open(f\"{jobname}.a3m\",\"r\").readlines())\n",
        "  hhsearch_result = hhsearch_pdb70_runner.query(a3m_lines)\n",
        "  hhsearch_hits = pipeline.parsers.parse_hhr(hhsearch_result)\n",
        "  templates_result = template_featurizer.get_templates(query_sequence=query_sequence,\n",
        "                                                       query_pdb_code=None,\n",
        "                                                       query_release_date=None,\n",
        "                                                       hhr_hits=hhsearch_hits)\n",
        "  return templates_result.features\n",
        "\n",
        "def set_bfactor(pdb_filename, bfac):\n",
        "  I = open(pdb_filename,\"r\").readlines()\n",
        "  O = open(pdb_filename,\"w\")\n",
        "  for line in I:\n",
        "    if line[0:6] == \"ATOM  \":\n",
        "      seq_id = int(line[23:26].strip()) - 1\n",
        "      O.write(\"{prefix}{bfac:6.2f}{suffix}\".format(prefix=line[:60], bfac=bfac[seq_id], suffix=line[66:]))\n",
        "  O.close()\n",
        "\n",
        "def predict_structure(prefix, feature_dict, do_relax=True, random_seed=0):  \n",
        "  \"\"\"Predicts structure using AlphaFold for the given sequence.\"\"\"\n",
        "\n",
        "  # Run the models.\n",
        "  plddts = []\n",
        "  unrelaxed_pdb_lines = []\n",
        "  relaxed_pdb_lines = []\n",
        "\n",
        "  for model_name, params in model_params.items():\n",
        "    print(f\"running {model_name}\")\n",
        "    # swap params to avoid recompiling\n",
        "    # note: models 1,2 have diff number of params compared to models 3,4,5\n",
        "    if any(str(m) in model_name for m in [1,2]): model_runner = model_runner_1\n",
        "    if any(str(m) in model_name for m in [3,4,5]): model_runner = model_runner_3\n",
        "    model_runner.params = params\n",
        "    \n",
        "    processed_feature_dict = model_runner.process_features(feature_dict, random_seed=random_seed)\n",
        "    prediction_result = model_runner.predict(processed_feature_dict)\n",
        "    unrelaxed_protein = protein.from_prediction(processed_feature_dict,prediction_result)\n",
        "    unrelaxed_pdb_lines.append(protein.to_pdb(unrelaxed_protein))\n",
        "    plddts.append(prediction_result['plddt'])\n",
        "\n",
        "    if do_relax:\n",
        "      # Relax the prediction.\n",
        "      amber_relaxer = relax.AmberRelaxation(max_iterations=0,tolerance=2.39,\n",
        "                                            stiffness=10.0,exclude_residues=[],\n",
        "                                            max_outer_iterations=20)      \n",
        "      relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)\n",
        "      relaxed_pdb_lines.append(relaxed_pdb_str)\n",
        "\n",
        "  # rerank models based on predicted lddt\n",
        "  lddt_rank = np.mean(plddts,-1).argsort()[::-1]\n",
        "  plddts_ranked = {}\n",
        "  for n,r in enumerate(lddt_rank):\n",
        "    print(f\"model_{n+1} {np.mean(plddts[r])}\")\n",
        "\n",
        "    unrelaxed_pdb_path = f'{prefix}_unrelaxed_model_{n+1}.pdb'    \n",
        "    with open(unrelaxed_pdb_path, 'w') as f: f.write(unrelaxed_pdb_lines[r])\n",
        "    set_bfactor(unrelaxed_pdb_path,plddts[r]/100)\n",
        "\n",
        "    if do_relax:\n",
        "      relaxed_pdb_path = f'{prefix}_relaxed_model_{n+1}.pdb'\n",
        "      with open(relaxed_pdb_path, 'w') as f: f.write(relaxed_pdb_lines[r])\n",
        "      set_bfactor(relaxed_pdb_path,plddts[r]/100)\n",
        "\n",
        "    plddts_ranked[f\"model_{n+1}\"] = plddts[r]\n",
        "\n",
        "  return plddts_ranked"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hUYApPElB30u",
        "cellView": "form"
      },
      "source": [
        "#@title Gather input features, predict structure\n",
        "# parse TEMPLATES\n",
        "if use_templates: template_features = mk_template(jobname)\n",
        "else: template_features = mk_mock_template(query_sequence)\n",
        "\n",
        "# parse MSA\n",
        "a3m_lines = \"\".join(open(a3m_file,\"r\").readlines())\n",
        "msa, deletion_matrix = pipeline.parsers.parse_a3m(a3m_lines)\n",
        "\n",
        "if len(query_sequence) != len(msa[0]):\n",
        "  print(f\"ERROR: the length of query ({len(query_sequence)}) does not match length of MSA sequences ({len(msa[0])})\")\n",
        "\n",
        "# gather features\n",
        "feature_dict = {\n",
        "    **pipeline.make_sequence_features(sequence=query_sequence,\n",
        "                                      description=\"none\",\n",
        "                                      num_res=len(query_sequence)),\n",
        "    **pipeline.make_msa_features(msas=[msa],deletion_matrices=[deletion_matrix]),\n",
        "    **template_features\n",
        "}\n",
        "plddts = predict_structure(jobname, feature_dict, do_relax=use_amber)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "exKwNxDxF7IO",
        "cellView": "form"
      },
      "source": [
        "#@title Plot lDDT per residue\n",
        "# confidence per position\n",
        "plt.figure(dpi=100)\n",
        "for model_name,value in plddts.items():\n",
        "  plt.plot(value,label=model_name)\n",
        "plt.legend()\n",
        "plt.ylim(0,100)\n",
        "plt.ylabel(\"predicted lDDT\")\n",
        "plt.xlabel(\"positions\")\n",
        "plt.savefig(jobname+\"_lDDT.png\")\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6xbvRNrwnJqj",
        "cellView": "form"
      },
      "source": [
        "#@title Plot Number of Sequences per Position\n",
        "# confidence per position\n",
        "plt.figure(dpi=100)\n",
        "plt.plot((feature_dict[\"msa\"] != 21).sum(0))\n",
        "plt.xlabel(\"positions\")\n",
        "plt.ylabel(\"number of sequences\")\n",
        "plt.savefig(jobname+\"_msa_coverage.png\")\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g-rPnOXdjf18",
        "cellView": "form"
      },
      "source": [
        "#@title Show 3D structure\n",
        "def show_pdb(model_name,\n",
        "             show_sidechains=False,\n",
        "             show_mainchain=False,\n",
        "             color=\"None\"):\n",
        "\n",
        "  def mainchain(p, color=\"white\", model=0):\n",
        "    BB = ['C','O','N','CA']\n",
        "    p.addStyle({\"model\":model,'atom':BB},\n",
        "                       {'stick':{'colorscheme':f\"{color}Carbon\",'radius':0.4}})\n",
        "\n",
        "  def sidechain(p, model=0):\n",
        "    HP = [\"ALA\",\"GLY\",\"VAL\",\"ILE\",\"LEU\",\"PHE\",\"MET\",\"PRO\",\"TRP\",\"CYS\",\"TYR\"]\n",
        "    BB = ['C','O','N']\n",
        "    p.addStyle({\"model\":model,'and':[{'resn':HP},{'atom':BB,'invert':True}]},\n",
        "              {'stick':{'colorscheme':\"yellowCarbon\",'radius':0.4}})\n",
        "    p.addStyle({\"model\":model,'and':[{'resn':\"GLY\"},{'atom':'CA'}]},\n",
        "              {'sphere':{'colorscheme':\"yellowCarbon\",'radius':0.4}})\n",
        "    p.addStyle({\"model\":model,'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
        "              {'stick':{'colorscheme':\"yellowCarbon\",'radius':0.4}})  \n",
        "    p.addStyle({\"model\":model,'and':[{'resn':HP,'invert':True},{'atom':BB,'invert':True}]},\n",
        "              {'stick':{'colorscheme':\"whiteCarbon\",'radius':0.4}})\n",
        "\n",
        "  if use_amber:\n",
        "    pdb_filename = f\"{jobname}_relaxed_{model_name}.pdb\"\n",
        "  else:\n",
        "    pdb_filename = f\"{jobname}_unrelaxed_{model_name}.pdb\"\n",
        "\n",
        "  p = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js')\n",
        "  p.addModel(open(pdb_filename,'r').read(),'pdb')\n",
        "  if color == \"lDDT\":\n",
        "    p.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':0,'max':1}}})\n",
        "  elif color == \"rainbow\":\n",
        "    p.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "  else:\n",
        "    p.setStyle({'cartoon':{}})\n",
        "\n",
        "  if show_sidechains: sidechain(p)\n",
        "  if show_mainchain: mainchain(p)\n",
        "  p.zoomTo()\n",
        "  return p.show()\n",
        "\n",
        "interact(show_pdb,\n",
        "         model_name=ipywidgets.Dropdown(options=model_params.keys(), value='model_1'),\n",
        "         show_sidechains=ipywidgets.Checkbox(value=False),\n",
        "         show_mainchain=ipywidgets.Checkbox(value=False),\n",
        "         color=ipywidgets.Dropdown(options=['None', 'rainbow', 'lDDT'], value='lDDT'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "33g5IIegij5R",
        "cellView": "form"
      },
      "source": [
        "#@title Package and download results\n",
        "!zip -FSr $jobname\".result.zip\" $jobname\".log\" $a3m_file $jobname\"_msa_coverage.png\" $jobname\"_\"*\"relaxed_model_\"*\".pdb\" $jobname\"_lDDT.png\"\n",
        "files.download(f\"{jobname}.result.zip\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vP1m6jI7BgMT"
      },
      "source": [
        "# Instructions\n",
        "- If you having issues downloading results, try disable adblocker and run the last cell again. If that fails click on the little folder icon to the left, navigate to file:`jobname.result.zip`, right-click and select \"download\"."
      ]
    }
  ]
}