DiscobaAlphaFold2HMMER.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "DiscobaAlphaFold2HMMER.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/zephyris/discoba_alphafold/blob/main/DiscobaAlphaFold2HMMER.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": [
        "#AlphaFold2 w/ MMseqs2 & HMMER Discoba\n",
        "=== This version is non-functional due to jaxlib incompatibilities, but kept unchanged as a permanent record ===\n",
        "=== Please try [this new experimental notebook](https://colab.research.google.com/github/zephyris/discoba_alphafold/blob/main/DiscobaAlphaFold2HMMERv2.ipynb) ===\n",
        "\n",
        "Easy-to-use version of AlphaFold2 incorporating searching of a database of Discoba species, ideal for analysing proteins from kinetoplastid species like _Leishmania_ and _Trypanosoma_.\n",
        "\n",
        "This is an adapted version of [ColabFold](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb), with everything except generating an a3m from the Discoba database taken _directly_ from that notebook. It appends the Discoba alignment (made using HMMER) to the result of the remote MMseqs2 search (specified in msa_mode).\n",
        "\n",
        "**Please read their [documentation](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb#scrollTo=UGUBLzB3C6WN)** for full instructions, full credit/attribution/citation information and caveats about using an unofficial MSA/template generation pipeline.\n",
        "\n",
        "A full description of this resource is available in [this paper](https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0259871).\n",
        "If you make use of this resource please cite: Wheeler RJ. \"A resource for improved predictions of Trypanosoma and Leishmania protein three-dimensional structure\" _PLoS One_, doi:&nbsp;[10.1371/journal.pone.0259871](https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0259871) (2021). You should also cite MMSeqs2 and the databases used, please see the information on their [Github page](https://github.com/soedinglab/MMseqs2).\n",
        "\n",
        "To use HMMER to generate only an a3m alignment from the Discoba database, for use in a different AlphaFold2 implementation, please use [this Colab notebook](https://colab.research.google.com/github/zephyris/discoba_alphafold/blob/main/DiscobaHMMER.ipynb)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kOblAo-xetgx",
        "cellView": "form"
      },
      "source": [
        "#@title Input protein sequence, then hit `Runtime` -> `Run all`\n",
        "from google.colab import files\n",
        "import os\n",
        "import os.path\n",
        "import re\n",
        "import hashlib\n",
        "\n",
        "def add_hash(x,y):\n",
        "  return x+\"_\"+hashlib.sha1(y.encode()).hexdigest()[:5]\n",
        "\n",
        "query_sequence = 'MESKPNHRRVWVTSAVVLTALVTLIVRRRQQKRQLARRTAADAMHEELAEALKQTLLPVLVRPEDLPAGWGVAPPVFYPPYCAAIPLVLPREMQESERDSSGVRSASTSRFQGGGADGLAVAPTAFLFCSYTGYTAGPFATRQDEEVFLTEVLAHHPVLRERLAENPRQWSPLPESAPHASTDAKTAAAVTEGRRQNSFPRTFSHVLVSDDCVLLAGMNGPYTVVAVLTDFAGAWPLSRTARASADGEESGRRRTGSAAPPTIAEETELSDAIEGIACATINVPLATLGSDSAALAFRVPGTCFAAHQGYYRVVCAREGRELELCVPPEWTMRSECVNTSCTTPAGPPVPARKPVSTVGEDAESKGIILTLSFTPSSFMSEGRVDVCISAELFSALYEEPQAAAATLWAASGATDVCNPVAKATLGALTARPLPASSHARTNVVTMVYVQPKFGVLFSVHPRSAVVYEPWMTEQPTILYYPLGDAADDEGSPRMTIEYVVELPKTWEVFARDDEEFVHNVLFHFTSGEAAAISTTLTEISGIRCAMFHETRESRRCRTYVLPRGATLLVIRWETLAESWDKDLPVFQQTLDTLHIDAAAVIQ' #@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 = 'A4IBK2' #@param {type:\"string\"}\n",
        "# remove whitespaces\n",
        "jobname = \"\".join(jobname.split())\n",
        "jobname = re.sub(r'\\W+', '', jobname)\n",
        "jobname = add_hash(jobname, query_sequence)\n",
        "\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",
        "msa_mode = \"MMseqs2 (UniRef+Environmental)\" #@param [\"MMseqs2 (UniRef+Environmental)\", \"MMseqs2 (UniRef only)\",\"single_sequence\",\"custom\"]\n",
        "num_models = 5 #@param [1,2,3,4,5] {type:\"raw\"}\n",
        "use_msa = True if msa_mode.startswith(\"MMseqs2\") else False\n",
        "use_env = True if msa_mode == \"MMseqs2 (UniRef+Environmental)\" 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",
        "#@markdown ### Experimental options\n",
        "homooligomer = 1 #@param [1,2,3,4,5,6,7,8] {type:\"raw\"}\n",
        "save_to_google_drive = False #@param {type:\"boolean\"}\n",
        "#@markdown ---\n",
        "#@markdown Don't forget to hit `Runtime` -> `Run all` after updating form\n",
        "\n",
        "\n",
        "if homooligomer > 1:\n",
        "  if use_amber:\n",
        "    print(\"amber disabled: amber is not currently supported for homooligomers\")\n",
        "    use_amber = False\n",
        "  if use_templates:\n",
        "    print(\"templates disabled: templates are not currently supported for homooligomers\")\n",
        "    use_templates = False\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",
        "    text_file.write(\"homooligomer=%s\\n\" % homooligomer)\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",
        "    header = 0\n",
        "    import fileinput\n",
        "    for line in fileinput.FileInput(custom_msa,inplace=1):\n",
        "      if line.startswith(\">\"):\n",
        "         header = header + 1 \n",
        "      if line.startswith(\"#\"):\n",
        "        continue\n",
        "      if line.rstrip() == False:\n",
        "        continue\n",
        "      if line.startswith(\">\") == False and header == 1:\n",
        "         query_sequence = line.rstrip() \n",
        "      print(line, end='')\n",
        "\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)\n",
        "\n",
        "if save_to_google_drive == True:\n",
        "  from pydrive.drive import GoogleDrive\n",
        "  from pydrive.auth import GoogleAuth\n",
        "  from google.colab import auth\n",
        "  from oauth2client.client import GoogleCredentials\n",
        "  auth.authenticate_user()\n",
        "  gauth = GoogleAuth()\n",
        "  gauth.credentials = GoogleCredentials.get_application_default()\n",
        "  drive = GoogleDrive(gauth)\n",
        "  print(\"You are logged into Google Drive and are good to go!\")"
      ],
      "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",
        "  #Ensure compatable jax library versions are installed\n",
        "  pip uninstall -y jax --quiet\n",
        "  pip uninstall -y jaxlib --quiet\n",
        "  pip install jax==0.2.27 --quiet\n",
        "  #pip install --upgrade --quiet jaxlib==0.1.76+cuda11.cudnn805 -f https://storage.googleapis.com/jax-releases/jax_releases.html\n",
        "  # No longer listed, so manually fetch and install\n",
        "  wget --quiet https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.1.76+cuda11.cudnn805-cp37-none-manylinux2010_x86_64.whl\n",
        "  pip install --quiet jaxlib-0.1.76+cuda11.cudnn805-cp37-none-manylinux2010_x86_64.whl\n",
        "\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",
        "    (cd alphafold; git checkout 0bab1bf84d9d887aba5cfb6d09af1e8c3ecbc408 --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": "bOBqlac_wYEh",
        "cellView": "form"
      },
      "source": [
        "#@title Install Hmmer, fetch Discoba database\n",
        "\n",
        "#@markdown _New section for Discoba alignments_\n",
        "\n",
        "#@markdown Downloads a database database which comprehensively gathers Discoba protein sequence databases.\n",
        "\n",
        "#@markdown Many of these sequences are not in the UniRef100 database, which is commonly used by online database search/alignment tools.\n",
        "\n",
        "%%bash -s\n",
        "#Install Hmmer, for search and alignment\n",
        "if [ ! -f HMMER_READY ]; then\n",
        "  apt-get install hmmer > /dev/null 2>&1\n",
        "  touch HMMER_READY\n",
        "fi\n",
        "\n",
        "#Download the custom Discoba database\n",
        "if [ ! -f DISCOBA_READY ]; then\n",
        "  if [ -d discoba ]; then\n",
        "    rm -r discoba\n",
        "  fi\n",
        "  mkdir discoba\n",
        "  cd discoba\n",
        "    curl http://wheelerlab.net/discobaStats.txt\n",
        "    curl http://wheelerlab.net/discoba.fasta.gz -s -L -o discoba.fasta.gz\n",
        "    gzip -d discoba.fasta.gz\n",
        "  cd ..\n",
        "  touch DISCOBA_READY\n",
        "fi\n",
        "\n",
        "#Install hh-suite\n",
        "if [ ! -f HHSUITE_READY ]; then\n",
        "  if [ -d hh-suite ]; then\n",
        "    rm -r hh-suite\n",
        "  fi\n",
        "  git clone https://github.com/soedinglab/hh-suite\n",
        "  touch HHSUITE_READY\n",
        "fi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JPWfhGssZdTb",
        "cellView": "form"
      },
      "source": [
        "#@title Import libraries\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, GridspecLayout, Output\n",
        "\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",
        "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",
        "                                                       hits=hhsearch_hits)\n",
        "  return templates_result.features\n",
        "\n",
        "def set_bfactor(pdb_filename, bfac, idx_res, chains):\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[22:26].strip()) - 1\n",
        "      seq_id = np.where(idx_res == seq_id)[0][0]\n",
        "      O.write(f\"{line[:21]}{chains[seq_id]}{line[22:60]}{bfac[seq_id]:6.2f}{line[66:]}\")\n",
        "  O.close()\n",
        "\n",
        "def predict_structure(prefix, feature_dict, Ls, model_params, use_model, do_relax=False, random_seed=0):  \n",
        "  \"\"\"Predicts structure using AlphaFold for the given sequence.\"\"\"\n",
        "\n",
        "  # Minkyung's code\n",
        "  # add big enough number to residue index to indicate chain breaks\n",
        "  idx_res = feature_dict['residue_index']\n",
        "  L_prev = 0\n",
        "  # Ls: number of residues in each chain\n",
        "  for L_i in Ls[:-1]:\n",
        "      idx_res[L_prev+L_i:] += 200\n",
        "      L_prev += L_i  \n",
        "  chains = list(\"\".join([ascii_uppercase[n]*L for n,L in enumerate(Ls)]))\n",
        "  feature_dict['residue_index'] = idx_res\n",
        "\n",
        "  # Run the models.\n",
        "  plddts,paes = [],[]\n",
        "  unrelaxed_pdb_lines = []\n",
        "  relaxed_pdb_lines = []\n",
        "\n",
        "  for model_name, params in model_params.items():\n",
        "    if model_name in use_model:\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",
        "      paes.append(prediction_result['predicted_aligned_error'])\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",
        "  out = {}\n",
        "  print(\"reranking models based on avg. predicted lDDT\")\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], idx_res, chains)\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], idx_res, chains)\n",
        "\n",
        "    out[f\"model_{n+1}\"] = {\"plddt\":plddts[r], \"pae\":paes[r]}\n",
        "  return out"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "139gnocjv2HZ",
        "cellView": "form"
      },
      "source": [
        "#@title Run Hmmer to get Discoba MSA\n",
        "\n",
        "#@markdown _New section for Discoba alignments_\n",
        "\n",
        "#@markdown It is normal for this to take several minutes.\n",
        "%%bash -s $jobname\n",
        "\n",
        "if [ ! -f $1.hmm.a3m ]; then\n",
        "  jackhmmer -A $1.hmm.sto -o $1.hmm.out $1.fasta discoba/discoba.fasta\n",
        "  perl hh-suite/scripts/reformat.pl sto a3m $1.hmm.sto $1.hmm.a3m\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 $use_env\n",
        "USE_AMBER=$1\n",
        "USE_MSA=$2\n",
        "USE_TEMPLATES=$3\n",
        "NAME=$4\n",
        "USE_ENV=$5\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",
        "    MODE=all\n",
        "    if [ ${USE_ENV} == \"True\" ]; then\n",
        "      MODE=env\n",
        "    fi\n",
        "    ID=$(curl -s -F q=@${NAME}.fasta -F mode=${MODE} https://api.colabfold.com/ticket/msa | jq -r '.id')\n",
        "    STATUS=$(curl -s https://api.colabfold.com/ticket/${ID} | jq -r '.status')\n",
        "    while [ \"${STATUS}\" == \"RUNNING\" ] || [ \"${STATUS}\" == \"PENDING\" ]; do\n",
        "      STATUS=$(curl -s https://api.colabfold.com/ticket/${ID} | jq -r '.status')\n",
        "      sleep 1\n",
        "    done\n",
        "    if [ \"${STATUS}\" == \"COMPLETE\" ]; then\n",
        "      curl -s https://api.colabfold.com/result/download/${ID} > ${NAME}.mmseqs2.tar.gz\n",
        "      tar xzf ${NAME}.mmseqs2.tar.gz\n",
        "      if [ ${USE_ENV} == \"True\" ]; then\n",
        "        cat uniref.a3m bfd.mgnify30.metaeuk30.smag30.a3m > tmp.a3m\n",
        "        tr -d '\\000' < tmp.a3m > ${NAME}.a3m\n",
        "        rm uniref.a3m bfd.mgnify30.metaeuk30.smag30.a3m tmp.a3m\n",
        "      else\n",
        "        tr -d '\\000' < uniref.a3m > ${NAME}.a3m\n",
        "        rm uniref.a3m\n",
        "      fi\n",
        "      mv pdb70.m8 ${NAME}.m8\n",
        "    else\n",
        "      echo \"MMseqs2 server did not return a valid result.\"\n",
        "      cp ${NAME}.fasta ${NAME}.a3m\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 [ -s ${NAME}.m8 ]; then\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",
        "    else\n",
        "      echo \"no templates found\"\n",
        "    fi\n",
        "  fi\n",
        "fi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "IxpdE8-vV7XK"
      },
      "source": [
        "#@title Concatenate HMMER a3m with remote\n",
        "\n",
        "#@markdown _New section for Discoba alignments_\n",
        "\n",
        "#@markdown Not running this section will leave the MSA unchanged\n",
        "#@markdown (ie. use only the remote MMSeqs UniRef/UniRef+Environmental alignment).\n",
        "%%bash -s $jobname\n",
        "\n",
        "if [ -f \"$1.original.a3m\" ]; then\n",
        "  cp \"$1.original.a3m\" \"$1.a3m\"\n",
        "else\n",
        "  perl hh-suite/scripts/reformat.pl a3m a3m $1.a3m $1.a3m > /dev/null 2>&1\n",
        "  cp \"$1.a3m\" \"$1.original.a3m\"\n",
        "fi\n",
        "\n",
        "COUNT=$(grep \">\" \"$1.hmm.a3m\" | wc -l)\n",
        "cat \"$1.hmm.a3m\" >> \"$1.a3m\"\n",
        "echo \"Added $COUNT sequences (no redundancy filtering)\"\n",
        "\n",
        "COUNT=$(grep \">\" \"$1.a3m\" | wc -l)\n",
        "echo \"$COUNT total sequences\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hUYApPElB30u",
        "cellView": "form"
      },
      "source": [
        "#@title Gather input features, predict structure\n",
        "from string import ascii_uppercase\n",
        "\n",
        "# collect model weights\n",
        "use_model = {}\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",
        "  use_model[model_name] = True\n",
        "  if model_name not in model_params:\n",
        "    model_params[model_name] = data.get_model_haiku_params(model_name=model_name+\"_ptm\", data_dir=\".\")\n",
        "    if model_name == \"model_1\":\n",
        "      model_config = config.model_config(model_name+\"_ptm\")\n",
        "      model_config.data.eval.num_ensemble = 1\n",
        "      model_runner_1 = model.RunModel(model_config, model_params[model_name])\n",
        "    if model_name == \"model_3\":\n",
        "      model_config = config.model_config(model_name+\"_ptm\")\n",
        "      model_config.data.eval.num_ensemble = 1\n",
        "      model_runner_3 = model.RunModel(model_config, model_params[model_name])\n",
        "\n",
        "# parse TEMPLATES\n",
        "if use_templates and os.path.isfile(f\"{jobname}_hhm.ffindex\"):\n",
        "  template_features = mk_template(jobname)\n",
        "else:\n",
        "  template_features = mk_mock_template(query_sequence * homooligomer)\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 homooligomer == 1:\n",
        "  msas = [msa]\n",
        "  deletion_matrices = [deletion_matrix]\n",
        "else:\n",
        "  # make multiple copies of msa for each copy\n",
        "  # AAA------\n",
        "  # ---AAA---\n",
        "  # ------AAA\n",
        "  #\n",
        "  # note: if you concat the sequences (as below), it does NOT work\n",
        "  # AAAAAAAAA\n",
        "  msas = []\n",
        "  deletion_matrices = []\n",
        "  Ln = len(query_sequence)\n",
        "  for o in range(homooligomer):\n",
        "    L = Ln * o\n",
        "    R = Ln * (homooligomer-(o+1))\n",
        "    msas.append([\"-\"*L+seq+\"-\"*R for seq in msa])\n",
        "    deletion_matrices.append([[0]*L+mtx+[0]*R for mtx in deletion_matrix])\n",
        "\n",
        "# gather features\n",
        "feature_dict = {\n",
        "    **pipeline.make_sequence_features(sequence=query_sequence*homooligomer,\n",
        "                                      description=\"none\",\n",
        "                                      num_res=len(query_sequence)*homooligomer),\n",
        "    **pipeline.make_msa_features(msas=msas,deletion_matrices=deletion_matrices),\n",
        "    **template_features\n",
        "}\n",
        "outs = predict_structure(jobname, feature_dict,\n",
        "                         Ls=[len(query_sequence)]*homooligomer,\n",
        "                         model_params=model_params, use_model=use_model,\n",
        "                         do_relax=use_amber)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6xbvRNrwnJqj",
        "cellView": "form"
      },
      "source": [
        "#@title Make plots\n",
        "\n",
        "# gather MSA info\n",
        "deduped_full_msa = list(dict.fromkeys(msa))\n",
        "msa_arr = np.array([list(seq) for seq in deduped_full_msa])\n",
        "seqid = (np.array(list(query_sequence)) == msa_arr).mean(-1)\n",
        "seqid_sort = seqid.argsort() #[::-1]\n",
        "non_gaps = (msa_arr != \"-\").astype(float)\n",
        "non_gaps[non_gaps == 0] = np.nan\n",
        "\n",
        "##################################################################\n",
        "plt.figure(figsize=(14,4),dpi=100)\n",
        "##################################################################\n",
        "plt.subplot(1,2,1); plt.title(\"Sequence coverage\")\n",
        "plt.imshow(non_gaps[seqid_sort]*seqid[seqid_sort,None],\n",
        "           interpolation='nearest', aspect='auto',\n",
        "           cmap=\"rainbow_r\", vmin=0, vmax=1, origin='lower')\n",
        "plt.plot((msa_arr != \"-\").sum(0), color='black')\n",
        "plt.xlim(-0.5,msa_arr.shape[1]-0.5)\n",
        "plt.ylim(-0.5,msa_arr.shape[0]-0.5)\n",
        "plt.colorbar(label=\"Sequence identity to query\",)\n",
        "plt.xlabel(\"Positions\")\n",
        "plt.ylabel(\"Sequences\")\n",
        "\n",
        "##################################################################\n",
        "plt.subplot(1,2,2); plt.title(\"Predicted lDDT per position\")\n",
        "for model_name,value in outs.items():\n",
        "  plt.plot(value[\"plddt\"],label=model_name)\n",
        "if homooligomer > 0:\n",
        "  for n in range(homooligomer+1):\n",
        "    x = n*(len(query_sequence)-1)\n",
        "    plt.plot([x,x],[0,100],color=\"black\")\n",
        "plt.legend()\n",
        "plt.ylim(0,100)\n",
        "plt.ylabel(\"Predicted lDDT\")\n",
        "plt.xlabel(\"Positions\")\n",
        "plt.savefig(jobname+\"_coverage_lDDT.png\")\n",
        "##################################################################\n",
        "plt.show()\n",
        "\n",
        "print(\"Predicted Alignment Error\")\n",
        "##################################################################\n",
        "plt.figure(figsize=(3*num_models,2), dpi=100)\n",
        "for n,(model_name,value) in enumerate(outs.items()):\n",
        "  plt.subplot(1,num_models,n+1)\n",
        "  plt.title(model_name)\n",
        "  plt.imshow(value[\"pae\"],label=model_name,cmap=\"bwr\",vmin=0,vmax=30)\n",
        "  plt.colorbar()\n",
        "plt.savefig(jobname+\"_PAE.png\")\n",
        "plt.show()\n",
        "##################################################################"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KK7X9T44pWb7",
        "cellView": "form"
      },
      "source": [
        "#@title Display 3D structure {run: \"auto\"}\n",
        "model_num = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\"] {type:\"raw\"}\n",
        "color = \"lDDT\" #@param [\"chain\", \"lDDT\", \"rainbow\"]\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "\n",
        "def plot_plddt_legend():\n",
        "  thresh = ['plDDT:','Very low (<50)','Low (60)','OK (70)','Confident (80)','Very high (>90)']\n",
        "  plt.figure(figsize=(1,0.1),dpi=100)\n",
        "  ########################################\n",
        "  for c in [\"#FFFFFF\",\"#FF0000\",\"#FFFF00\",\"#00FF00\",\"#00FFFF\",\"#0000FF\"]:\n",
        "    plt.bar(0, 0, color=c)\n",
        "  plt.legend(thresh, frameon=False,\n",
        "             loc='center', ncol=6,\n",
        "             handletextpad=1,\n",
        "             columnspacing=1,\n",
        "             markerscale=0.5,)\n",
        "  plt.axis(False)\n",
        "  return plt\n",
        "\n",
        "def plot_confidence(model_num=1):\n",
        "  model_name = f\"model_{model_num}\"\n",
        "  plt.figure(figsize=(10,3),dpi=100)\n",
        "  \"\"\"Plots the legend for plDDT.\"\"\"\n",
        "  #########################################\n",
        "  plt.subplot(1,2,1); plt.title('Predicted lDDT')\n",
        "  plt.plot(outs[model_name][\"plddt\"])\n",
        "  for n in range(homooligomer+1):\n",
        "    x = n*(len(query_sequence))\n",
        "    plt.plot([x,x],[0,100],color=\"black\")\n",
        "  plt.ylabel('plDDT')\n",
        "  plt.xlabel('position')\n",
        "  #########################################\n",
        "  plt.subplot(1,2,2);plt.title('Predicted Aligned Error')\n",
        "  plt.imshow(outs[model_name][\"pae\"], cmap=\"bwr\",vmin=0,vmax=30)\n",
        "  plt.colorbar()\n",
        "  plt.xlabel('Scored residue')\n",
        "  plt.ylabel('Aligned residue')\n",
        "  #########################################\n",
        "  return plt\n",
        "\n",
        "def show_pdb(model_num=1, show_sidechains=False, show_mainchains=False, color=\"lDDT\"):\n",
        "  model_name = f\"model_{model_num}\"\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",
        "  view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
        "  view.addModel(open(pdb_filename,'r').read(),'pdb')\n",
        "\n",
        "  if color == \"lDDT\":\n",
        "    view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':50,'max':90}}})\n",
        "  elif color == \"rainbow\":\n",
        "    view.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "  elif color == \"chain\":\n",
        "    for n,chain,color in zip(range(homooligomer),list(\"ABCDEFGH\"),\n",
        "                     [\"lime\",\"cyan\",\"magenta\",\"yellow\",\"salmon\",\"white\",\"blue\",\"orange\"]):\n",
        "       view.setStyle({'chain':chain},{'cartoon': {'color':color}})\n",
        "  if show_sidechains:\n",
        "    BB = ['C','O','N']\n",
        "    view.addStyle({'and':[{'resn':[\"GLY\",\"PRO\"],'invert':True},{'atom':BB,'invert':True}]},\n",
        "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"GLY\"},{'atom':'CA'}]},\n",
        "                        {'sphere':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
        "                        {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})  \n",
        "  if show_mainchains:\n",
        "    BB = ['C','O','N','CA']\n",
        "    view.addStyle({'atom':BB},{'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "\n",
        "  view.zoomTo()\n",
        "  return view\n",
        "\n",
        "show_pdb(model_num,show_sidechains, show_mainchains, color).show()\n",
        "if color == \"lDDT\": plot_plddt_legend().show()  \n",
        "plot_confidence(model_num).show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BOKHYdSeatkj",
        "cellView": "form"
      },
      "source": [
        "#@title Prepare JSON summary of plDDT and pAE\n",
        "\n",
        "#@markdown _New section_\n",
        "\n",
        "import json\n",
        "\n",
        "result={}\n",
        "for model_num in range(1,num_models+1):\n",
        "  model_name = f\"model_{model_num}\"\n",
        "  result[model_num]={}\n",
        "  result[model_num][\"plddt\"]=outs[model_name][\"plddt\"].tolist()\n",
        "  result[model_num][\"pae\"]=outs[model_name][\"pae\"].tolist()\n",
        "\n",
        "with open(f\"{jobname}.json\", 'w') as outfile:\n",
        "    json.dump(result, outfile)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "33g5IIegij5R",
        "cellView": "form"
      },
      "source": [
        "#@title Package and download results\n",
        "#@markdown If you having issues downloading the result archive, try disabling your adblocker and run this 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\\\" (see [screenshot](https://pbs.twimg.com/media/E6wRW2lWUAEOuoe?format=jpg&name=small)).\n",
        "\n",
        "citations = {\n",
        "  \"Ovchinnikov2021\":  \"\"\"@software{Ovchinnikov2021,\n",
        "author = {Ovchinnikov, Sergey and Steinegger, Martin and Mirdita, Milot},\n",
        "title = {{ColabFold - Making Protein folding accessible to all via Google Colab}},\n",
        "year = {2021},\n",
        "publisher = {Zenodo},\n",
        "version = {v1.0-alpha},\n",
        "doi = {10.5281/zenodo.5123297},\n",
        "url = {https://doi.org/10.5281/zenodo.5123297},\n",
        "comment = {The AlphaFold notebook}\n",
        "}\"\"\",\n",
        "  \"LevyKarin2020\": \"\"\"@article{LevyKarin2020,\n",
        "author = {{Levy Karin}, Eli and Mirdita, Milot and S{\\\"{o}}ding, Johannes},\n",
        "doi = {10.1186/s40168-020-00808-x},\n",
        "journal = {Microbiome},\n",
        "number = {1},\n",
        "title = {{MetaEuk—sensitive, high-throughput gene discovery, and annotation for large-scale eukaryotic metagenomics}},\n",
        "volume = {8},\n",
        "year = {2020},\n",
        "comment = {MetaEuk database}\n",
        "}\"\"\",\n",
        "  \"Delmont2020\": \"\"\"@article{Delmont2020,\n",
        "author = {Delmont, Tom O. and Gaia, Morgan and Hinsinger, Damien D. and Fremont, Paul and Guerra, Antonio Fernandez and Eren, A. Murat and Vanni, Chiara and Kourlaiev, Artem and D'Agata, Leo and Clayssen, Quentin and Villar, Emilie and Labadie, Karine and Cruaud, Corinne and Poulain, Julie and da Silva, Corinne and Wessner, Marc and Noel, Benjamin and Aury, Jean Marc and de Vargas, Colomban and Bowler, Chris and Karsenti, Eric and Pelletier, Eric and Wincker, Patrick and Jaillon, Olivier and Sunagawa, Shinichi and Acinas, Silvia G. and Bork, Peer and Karsenti, Eric and Bowler, Chris and Sardet, Christian and Stemmann, Lars and de Vargas, Colomban and Wincker, Patrick and Lescot, Magali and Babin, Marcel and Gorsky, Gabriel and Grimsley, Nigel and Guidi, Lionel and Hingamp, Pascal and Jaillon, Olivier and Kandels, Stefanie and Iudicone, Daniele and Ogata, Hiroyuki and Pesant, St{\\'{e}}phane and Sullivan, Matthew B. and Not, Fabrice and Karp-Boss, Lee and Boss, Emmanuel and Cochrane, Guy and Follows, Michael and Poulton, Nicole and Raes, Jeroen and Sieracki, Mike and Speich, Sabrina},\n",
        "journal = {bioRxiv},\n",
        "title = {{Functional repertoire convergence of distantly related eukaryotic plankton lineages revealed by genome-resolved metagenomics}},\n",
        "year = {2020},\n",
        "comment = {SMAG database}\n",
        "}\"\"\",\n",
        "  \"Mitchell2019\": \"\"\"@article{Mitchell2019,\n",
        "author = {Mitchell, Alex L and Almeida, Alexandre and Beracochea, Martin and Boland, Miguel and Burgin, Josephine and Cochrane, Guy and Crusoe, Michael R and Kale, Varsha and Potter, Simon C and Richardson, Lorna J and Sakharova, Ekaterina and Scheremetjew, Maxim and Korobeynikov, Anton and Shlemov, Alex and Kunyavskaya, Olga and Lapidus, Alla and Finn, Robert D},\n",
        "doi = {10.1093/nar/gkz1035},\n",
        "journal = {Nucleic Acids Res.},\n",
        "title = {{MGnify: the microbiome analysis resource in 2020}},\n",
        "year = {2019},\n",
        "comment = {MGnify database}\n",
        "}\"\"\",\n",
        "  \"Eastman2017\": \"\"\"@article{Eastman2017,\n",
        "author = {Eastman, Peter and Swails, Jason and Chodera, John D. and McGibbon, Robert T. and Zhao, Yutong and Beauchamp, Kyle A. and Wang, Lee-Ping and Simmonett, Andrew C. and Harrigan, Matthew P. and Stern, Chaya D. and Wiewiora, Rafal P. and Brooks, Bernard R. and Pande, Vijay S.},\n",
        "doi = {10.1371/journal.pcbi.1005659},\n",
        "journal = {PLOS Comput. Biol.},\n",
        "number = {7},\n",
        "title = {{OpenMM 7: Rapid development of high performance algorithms for molecular dynamics}},\n",
        "volume = {13},\n",
        "year = {2017},\n",
        "comment = {Amber relaxation}\n",
        "}\"\"\",\n",
        "  \"Jumper2021\": \"\"\"@article{Jumper2021,\n",
        "author = {Jumper, John and Evans, Richard and Pritzel, Alexander and Green, Tim and Figurnov, Michael and Ronneberger, Olaf and Tunyasuvunakool, Kathryn and Bates, Russ and {\\v{Z}}{\\'{i}}dek, Augustin and Potapenko, Anna and Bridgland, Alex and Meyer, Clemens and Kohl, Simon A. A. and Ballard, Andrew J. and Cowie, Andrew and Romera-Paredes, Bernardino and Nikolov, Stanislav and Jain, Rishub and Adler, Jonas and Back, Trevor and Petersen, Stig and Reiman, David and Clancy, Ellen and Zielinski, Michal and Steinegger, Martin and Pacholska, Michalina and Berghammer, Tamas and Bodenstein, Sebastian and Silver, David and Vinyals, Oriol and Senior, Andrew W. and Kavukcuoglu, Koray and Kohli, Pushmeet and Hassabis, Demis},\n",
        "doi = {10.1038/s41586-021-03819-2},\n",
        "journal = {Nature},\n",
        "pmid = {34265844},\n",
        "title = {{Highly accurate protein structure prediction with AlphaFold.}},\n",
        "year = {2021},\n",
        "comment = {AlphaFold2 + BFD Database}\n",
        "}\"\"\",\n",
        "  \"Mirdita2019\": \"\"\"@article{Mirdita2019,\n",
        "author = {Mirdita, Milot and Steinegger, Martin and S{\\\"{o}}ding, Johannes},\n",
        "doi = {10.1093/bioinformatics/bty1057},\n",
        "journal = {Bioinformatics},\n",
        "number = {16},\n",
        "pages = {2856--2858},\n",
        "pmid = {30615063},\n",
        "title = {{MMseqs2 desktop and local web server app for fast, interactive sequence searches}},\n",
        "volume = {35},\n",
        "year = {2019},\n",
        "comment = {MMseqs2 search server}\n",
        "}\"\"\",\n",
        "  \"Steinegger2019\": \"\"\"@article{Steinegger2019,\n",
        "author = {Steinegger, Martin and Meier, Markus and Mirdita, Milot and V{\\\"{o}}hringer, Harald and Haunsberger, Stephan J. and S{\\\"{o}}ding, Johannes},\n",
        "doi = {10.1186/s12859-019-3019-7},\n",
        "journal = {BMC Bioinform.},\n",
        "number = {1},\n",
        "pages = {473},\n",
        "pmid = {31521110},\n",
        "title = {{HH-suite3 for fast remote homology detection and deep protein annotation}},\n",
        "volume = {20},\n",
        "year = {2019},\n",
        "comment = {PDB70 database}\n",
        "}\"\"\",\n",
        "  \"Mirdita2017\": \"\"\"@article{Mirdita2017,\n",
        "author = {Mirdita, Milot and von den Driesch, Lars and Galiez, Clovis and Martin, Maria J. and S{\\\"{o}}ding, Johannes and Steinegger, Martin},\n",
        "doi = {10.1093/nar/gkw1081},\n",
        "journal = {Nucleic Acids Res.},\n",
        "number = {D1},\n",
        "pages = {D170--D176},\n",
        "pmid = {27899574},\n",
        "title = {{Uniclust databases of clustered and deeply annotated protein sequences and alignments}},\n",
        "volume = {45},\n",
        "year = {2017},\n",
        "comment = {Uniclust30/UniRef30 database},\n",
        "}\"\"\",\n",
        "  \"Berman2003\": \"\"\"@misc{Berman2003,\n",
        "author = {Berman, Helen and Henrick, Kim and Nakamura, Haruki},\n",
        "booktitle = {Nat. Struct. Biol.},\n",
        "doi = {10.1038/nsb1203-980},\n",
        "number = {12},\n",
        "pages = {980},\n",
        "pmid = {14634627},\n",
        "title = {{Announcing the worldwide Protein Data Bank}},\n",
        "volume = {10},\n",
        "year = {2003},\n",
        "comment = {templates downloaded from wwPDB server}\n",
        "}\"\"\",\n",
        "  \"Wheeler2021\":  \"\"\"@article{Wheeler2021,\n",
        "author = {Wheeler, Richard John},\n",
        "title = {{A resource for improved predictions of Trypanosoma and Leishmania protein three-dimensional structure}},\n",
        "year = {2021},\n",
        "publisher = {PloS One},\n",
        "doi = {10.1371/journal.pone.0259871},\n",
        "url = {https://doi.org/10.1371/journal.pone.0259871},\n",
        "comment = {The Discoba database}\n",
        "}\"\"\",\n",
        "}\n",
        "\n",
        "to_cite = [ \"Jumper2021\", \"Ovchinnikov2021\", \"Wheeler2021\" ]\n",
        "if use_msa:       to_cite += [\"Mirdita2019\"]\n",
        "if use_msa:       to_cite += [\"Mirdita2017\"]\n",
        "if use_env:       to_cite += [\"Mitchell2019\"]\n",
        "if use_env:       to_cite += [\"Delmont2020\"]\n",
        "if use_env:       to_cite += [\"LevyKarin2020\"]\n",
        "if use_templates: to_cite += [\"Steinegger2019\"]\n",
        "if use_templates: to_cite += [\"Berman2003\"]\n",
        "if use_amber:     to_cite += [\"Eastman2017\"]\n",
        "\n",
        "with open(f\"{jobname}.bibtex\", 'w') as writer:\n",
        "  for i in to_cite:\n",
        "    writer.write(citations[i])\n",
        "    writer.write(\"\\n\")\n",
        "\n",
        "print(f\"Found {len(to_cite)} citation{'s' if len(to_cite) > 1 else ''} for tools or databases.\")\n",
        "if use_custom_msa:\n",
        "  print(\"Don't forget to cite your custom MSA generation method.\")\n",
        "\n",
        "!zip -FSr $jobname\".result.zip\" $jobname\".log\" $a3m_file $jobname\"_\"*\"relaxed_model_\"*\".pdb\" $jobname\"_coverage_lDDT.png\" $jobname\".bibtex\" $jobname\"_PAE.png\" $jobname\".json\"\n",
        "files.download(f\"{jobname}.result.zip\")\n",
        "\n",
        "if save_to_google_drive == True and drive != None:\n",
        "  uploaded = drive.CreateFile({'title': f\"{jobname}.result.zip\"})\n",
        "  uploaded.SetContentFile(f\"{jobname}.result.zip\")\n",
        "  uploaded.Upload()\n",
        "  print(f\"Uploaded {jobname}.result.zip to Google Drive with ID {uploaded.get('id')}\")"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}