AlphaFold2_complexes.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "AlphaFold2_complexes.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_complexes.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_complexes\n",
        "\n",
        "---------\n",
        "**UPDATE** (Aug. 13, 2021)\n",
        "\n",
        "This notebook is being retired and no longer updated. The functionality for complex prediction (including going beyond dimers) has been integrated in our [new advanced notebook](https://github.com/sokrypton/ColabFold/blob/main/beta/AlphaFold2_advanced.ipynb).\n",
        "\n",
        "---------\n",
        "\n",
        "Credit to Minkyung Baek @minkbaek and Yoshitaka Moriwaki @Ag_smith for initially showing protein-complex prediction works in alphafold2.\n",
        "- https://twitter.com/minkbaek/status/1417538291709071362\n",
        "- https://twitter.com/Ag_smith/status/1417063635000598528\n",
        "\n",
        "- [script](https://github.com/RosettaCommons/RoseTTAFold/blob/main/example/complex_modeling/make_joint_MSA_bacterial.py) from rosettafold for paired alignment generation\n",
        "\n",
        "**Instructions**\n",
        "- For *monomers* and *homo-oligomers*, see this [notebook](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb).\n",
        "- For prokaryotic protein complexes (found in operons), we recommend using the `pair_msa` option.\n",
        "\n",
        "\n",
        "**Limitations**\n",
        "- This notebook does NOT use templates or amber relax at the end for refinement.\n",
        "- For a typical Google-Colab-GPU (16G) session, the max total length is **1400 residues**.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kOblAo-xetgx",
        "cellView": "form"
      },
      "source": [
        "#@title Input protein sequences\n",
        "import os\n",
        "os.environ['TF_FORCE_UNIFIED_MEMORY'] = '1'\n",
        "os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '2.0'\n",
        "\n",
        "from google.colab import files\n",
        "import os.path\n",
        "import re\n",
        "\n",
        "import hashlib\n",
        "def add_hash(x,y):\n",
        "  return x+\"_\"+hashlib.sha1(y.encode()).hexdigest()[:5]\n",
        "\n",
        "query_sequence_a = 'AVLKIIQGALDTRELLKAYQEEACAKNFGAFCVFVGIVRKEDNIQGLSFDIYEALLKTWFEKWHHKAKDLGVVLKMAHSLGDVLIGQSSFLCVSMGKNRKNALELYENFIEDFKHNAPIWKYDLIHNKRIYAKERSHPLKGSGLLA' #@param {type:\"string\"}\n",
        "query_sequence_a = \"\".join(query_sequence_a.split())\n",
        "query_sequence_a = re.sub(r'[^A-Z]','', query_sequence_a.upper())\n",
        "\n",
        "query_sequence_b = 'MMVEVRFFGPIKEENFFIKANDLKELRAILQEKEGLKEWLGVCAIALNDHLIDNLNTPLKDGDVISLLPPVCGG' #@param {type:\"string\"}\n",
        "query_sequence_b = \"\".join(query_sequence_b.split())\n",
        "query_sequence_b = re.sub(r'[^A-Z]','', query_sequence_b.upper())\n",
        "\n",
        "# Using trick from @onoda_hiroki\n",
        "# https://twitter.com/onoda_hiroki/status/1420068104239910915\n",
        "# \"U\" indicates an \"UNKNOWN\" residue and it will not be modeled\n",
        "# But we need linker of at least length 32\n",
        "query_sequence_a = re.sub(r'U+',\"U\"*32,query_sequence_a)\n",
        "query_sequence_b = re.sub(r'U+',\"U\"*32,query_sequence_b)\n",
        "\n",
        "query_sequence = query_sequence_a + query_sequence_b\n",
        "\n",
        "if len(query_sequence) > 1400:\n",
        "  print(f\"WARNING: For a typical Google-Colab-GPU (16G) session, the max total length is 1400 residues. You are at {len(query_sequence)}!\")\n",
        "\n",
        "jobname = 'test' #@param {type:\"string\"}\n",
        "jobname = \"\".join(jobname.split())\n",
        "jobname = re.sub(r'\\W+', '', jobname)\n",
        "jobname = add_hash(jobname, 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\"]\n",
        "use_msa = True if msa_mode == \"MMseqs2\" else False\n",
        "pair_msa = False #@param {type:\"boolean\"}\n",
        "disable_mmseqs2_filter = pair_msa\n",
        "\n",
        "#@markdown ---\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_msa=%s\\n\" % use_msa)\n",
        "    text_file.write(\"msa_mode=%s\\n\" % msa_mode)\n",
        "    text_file.write(\"pair_msa=%s\\n\" % pair_msa)\n",
        "    text_file.write(\"disable_mmseqs2_filter=%s\\n\" % disable_mmseqs2_filter)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iccGdbe_Pmt9",
        "cellView": "form"
      },
      "source": [
        "#@title Install dependencies\n",
        "%%bash -s $use_msa\n",
        "USE_MSA=$1\n",
        "if [ ! -f AF2_READY ]; then\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",
        "  wget -qnc https://raw.githubusercontent.com/sokrypton/ColabFold/main/beta/colabfold.py\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",
        "  # 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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JPWfhGssZdTb",
        "cellView": "form"
      },
      "source": [
        "#@title Import libraries\n",
        "# setup the model\n",
        "if \"IMPORTED\" not in dir():\n",
        "\n",
        "  import time\n",
        "  import requests\n",
        "  import tarfile\n",
        "  import sys\n",
        "  import numpy as np\n",
        "  import pickle\n",
        "\n",
        "  from string import ascii_uppercase\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",
        "  import colabfold as cf\n",
        "\n",
        "  # plotting libraries\n",
        "  import py3Dmol\n",
        "  import matplotlib.pyplot as plt\n",
        "  IMPORTED = True\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, random_seed=0, num_models=5):  \n",
        "  \"\"\"Predicts structure using AlphaFold for the given sequence.\"\"\"\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 = []\n",
        "  paes = []\n",
        "  unrelaxed_pdb_lines = []\n",
        "  relaxed_pdb_lines = []\n",
        "\n",
        "  model_names = [\"model_4\",\"model_1\",\"model_2\",\"model_3\",\"model_5\"][:num_models]\n",
        "  for n,model_name in enumerate(model_names):\n",
        "    model_config = config.model_config(model_name+\"_ptm\")\n",
        "    model_config.data.eval.num_ensemble = 1\n",
        "    model_params = data.get_model_haiku_params(model_name+\"_ptm\", data_dir=\".\")\n",
        "\n",
        "    if model_name == \"model_4\":\n",
        "      model_runner = model.RunModel(model_config, model_params)\n",
        "      processed_feature_dict = model_runner.process_features(feature_dict,random_seed=0)\n",
        "    else:\n",
        "      # swap params\n",
        "      for k in model_runner.params.keys():\n",
        "        model_runner.params[k] = model_params[k]\n",
        "\n",
        "    print(f\"running model_{n+1}\")\n",
        "    prediction_result = model_runner.predict(processed_feature_dict)\n",
        "\n",
        "    # cleanup to save memory\n",
        "    if model_name == \"model_5\": del model_runner\n",
        "    del model_params\n",
        "    \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",
        "    # Delete unused outputs to save memory.\n",
        "    del prediction_result\n",
        "\n",
        "  # rerank models based on predicted lddt\n",
        "  lddt_rank = np.mean(plddts,-1).argsort()[::-1]\n",
        "  plddts_ranked = {}\n",
        "  paes_ranked = {}\n",
        "  print(\"model\\tplldt\\tpae_ab\")\n",
        "  L = Ls[0]\n",
        "  for n,r in enumerate(lddt_rank):\n",
        "    plddt = plddts[r].mean()\n",
        "    pae_ab = (paes[r][L:,:L].mean() + paes[r][:L,L:].mean()) / 2\n",
        "    print(f\"model_{n+1}\\t{plddt:.2f}\\t{pae_ab:.2f}\")\n",
        "    \n",
        "    unrelaxed_pdb_path = f'{prefix}_unrelaxed_model_{n+1}.pdb'    \n",
        "    with open(unrelaxed_pdb_path, 'w') as f:\n",
        "      f.write(unrelaxed_pdb_lines[r])\n",
        "    set_bfactor(unrelaxed_pdb_path, plddts[r], idx_res, chains)\n",
        "\n",
        "    plddts_ranked[f\"model_{n+1}\"] = plddts[r]\n",
        "    paes_ranked[f\"model_{n+1}\"] = paes[r]\n",
        "\n",
        "  return plddts_ranked, paes_ranked\n",
        "\n",
        "# CODE FROM MINKYUNG/ROSETTAFOLD\n",
        "def read_a3m(a3m_lines):\n",
        "  '''parse an a3m files as a dictionary {label->sequence}'''\n",
        "  seq = []\n",
        "  lab = []\n",
        "  is_first = True\n",
        "  for line in a3m_lines.splitlines():\n",
        "    if line[0] == '>':\n",
        "      label = line.rstrip().split()[0][1:]\n",
        "      is_incl = True\n",
        "      if is_first: # include first sequence (query)\n",
        "        is_first = False\n",
        "        lab.append(label)\n",
        "        continue\n",
        "      if \"UniRef\" in label:\n",
        "        code = label.split()[0].split('_')[-1]\n",
        "        if code.startswith(\"UPI\"): # UniParc identifier -- exclude\n",
        "          is_incl = False\n",
        "          continue\n",
        "      elif label.startswith(\"tr|\"):\n",
        "        code = label.split('|')[1]\n",
        "      else:\n",
        "        is_incl = False\n",
        "        continue\n",
        "      lab.append(code)\n",
        "    else:\n",
        "      if is_incl:\n",
        "        seq.append(line.rstrip())\n",
        "      else:\n",
        "        continue\n",
        "  return seq, lab\n",
        "\n",
        "# https://www.uniprot.org/help/accession_numbers\n",
        "def uni2idx(ids):\n",
        "  '''convert uniprot ids into integers according to the structure\n",
        "  of uniprot accession numbers'''\n",
        "  ids2 = [i.split(\"-\")[0] for i in ids]\n",
        "  ids2 = [i+'AAA0' if len(i)==6 else i for i in ids2]\n",
        "  arr = np.array([list(s) for s in ids2], dtype='|S1').view(np.uint8)\n",
        "  for i in [1,5,9]:\n",
        "    arr[:,i] -= ord('0')\n",
        "  arr[arr>=ord('A')] -= ord('A')\n",
        "  arr[arr>=ord('0')] -= ord('0')-26\n",
        "  arr[:,0][arr[:,0]>ord('Q')-ord('A')] -= 3\n",
        "  arr = arr.astype(np.int64)\n",
        "  coef = np.array([23,10,26,36,36,10,26,36,36,1], dtype=np.int64)\n",
        "  coef = np.tile(coef[None,:],[len(ids),1])\n",
        "  c1 = [i for i,id_ in enumerate(ids) if id_[0] in 'OPQ' and len(id_)==6]\n",
        "  c2 = [i for i,id_ in enumerate(ids) if id_[0] not in 'OPQ' and len(id_)==6]\n",
        "  coef[c1] = np.array([3, 10,36,36,36,1,1,1,1,1])\n",
        "  coef[c2] = np.array([23,10,26,36,36,1,1,1,1,1])\n",
        "  for i in range(1,10):\n",
        "    coef[:,-i-1] *= coef[:,-i]\n",
        "  return np.sum(arr*coef,axis=-1)\n",
        "\n",
        "def run_mmseqs2(query_sequence, prefix, use_env=True, filter=False):\n",
        "    def submit(query_sequence, mode):\n",
        "      res = requests.post('https://a3m.mmseqs.com/ticket/msa', data={'q':f\">1\\n{query_sequence}\", 'mode': mode})\n",
        "      return res.json()\n",
        "    def status(ID):\n",
        "      res = requests.get(f'https://a3m.mmseqs.com/ticket/{ID}')\n",
        "      return res.json()\n",
        "    def download(ID, path):\n",
        "      res = requests.get(f'https://a3m.mmseqs.com/result/download/{ID}')\n",
        "      with open(path,\"wb\") as out: out.write(res.content)\n",
        "      \n",
        "    if filter:\n",
        "      mode = \"env\" if use_env else \"all\"\n",
        "    else:\n",
        "      mode = \"env-nofilter\" if use_env else \"nofilter\"\n",
        "    \n",
        "    path = f\"{prefix}_{mode}\"\n",
        "    if not os.path.isdir(path): os.mkdir(path)\n",
        "\n",
        "    # call mmseqs2 api\n",
        "    tar_gz_file = f'{path}/out.tar.gz'\n",
        "    if not os.path.isfile(tar_gz_file):\n",
        "      out = submit(query_sequence, mode)\n",
        "      while out[\"status\"] in [\"RUNNING\",\"PENDING\"]:\n",
        "        time.sleep(1)\n",
        "        out = status(out[\"id\"])    \n",
        "      download(out[\"id\"], tar_gz_file)\n",
        "    \n",
        "    # parse a3m files\n",
        "    a3m_lines = []\n",
        "    a3m = f\"{prefix}_{mode}.a3m\"\n",
        "    if not os.path.isfile(a3m):\n",
        "      with tarfile.open(tar_gz_file) as tar_gz: tar_gz.extractall(path)\n",
        "      a3m_files = [f\"{path}/uniref.a3m\"]\n",
        "      if use_env: a3m_files.append(f\"{path}/bfd.mgnify30.metaeuk30.smag30.a3m\")\n",
        "      a3m_out = open(a3m,\"w\")\n",
        "      for a3m_file in a3m_files:\n",
        "        for line in open(a3m_file,\"r\"):\n",
        "          line = line.replace(\"\\x00\",\"\")\n",
        "          if len(line) > 0:\n",
        "            a3m_lines.append(line)\n",
        "            a3m_out.write(line)\n",
        "    else:\n",
        "      a3m_lines = open(a3m).readlines()\n",
        "    return \"\".join(a3m_lines), len(a3m_lines)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A9tUpDaikPC8",
        "cellView": "form"
      },
      "source": [
        "#@title Call MMseqs2 to get MSA for each gene \n",
        "\n",
        "Ls = [len(query_sequence_a),len(query_sequence_b)]\n",
        "msas = []\n",
        "deletion_matrices = []\n",
        "if use_msa:\n",
        "  os.makedirs('tmp', exist_ok=True)\n",
        "  prefix = hashlib.sha1(query_sequence.encode()).hexdigest()\n",
        "  prefix = os.path.join('tmp',prefix)\n",
        "  print(f\"running mmseqs2 (use_env={True} filter={True})\")\n",
        "  a3m_lines = cf.run_mmseqs2([query_sequence_a, query_sequence_b], prefix, use_env=True, filter=True)\n",
        "  if pair_msa:\n",
        "    a3m_lines.append([])\n",
        "\n",
        "    print(f\"running mmseqs2 for pair_msa (use_env={False} filter={False})\")\n",
        "    a3m_lines_pair = cf.run_mmseqs2([query_sequence_a, query_sequence_b], prefix, use_env=False, filter=False)\n",
        "\n",
        "    # CODE FROM MINKYUNG/ROSETTAFOLD\n",
        "    msa1, lab1 = read_a3m(a3m_lines_pair[0])\n",
        "    msa2, lab2 = read_a3m(a3m_lines_pair[1])\n",
        "    if len(lab1) > 1 and len(lab2) > 1:\n",
        "      # convert uniprot ids into integers\n",
        "      hash1 = uni2idx(lab1[1:])\n",
        "      hash2 = uni2idx(lab2[1:])\n",
        "\n",
        "      # find pairs of uniprot ids which are separated by at most 10\n",
        "      idx1, idx2 = np.where(np.abs(hash1[:,None]-hash2[None,:]) < 10)\n",
        "      if idx1.shape[0] > 0:\n",
        "        a3m_lines[2] = ['>query\\n%s%s\\n'%(msa1[0],msa2[0])]\n",
        "        for i,j in zip(idx1,idx2):\n",
        "           a3m_lines[2].append(\">%s_%s\\n%s%s\\n\"%(lab1[i+1],lab2[j+1],msa1[i+1],msa2[j+1]))\n",
        "        \n",
        "        msa, deletion_matrix = pipeline.parsers.parse_a3m(\"\".join(a3m_lines[2]))\n",
        "        msas.append(msa)\n",
        "        deletion_matrices.append(deletion_matrix)\n",
        "        print(\"pairs found:\",len(msa))\n",
        "  \n",
        "  msa, deletion_matrix = pipeline.parsers.parse_a3m(a3m_lines[0])\n",
        "  msas.append([seq+\"-\"*Ls[1] for seq in msa])\n",
        "  deletion_matrices.append([mtx+[0]*Ls[1] for mtx in deletion_matrix])\n",
        "\n",
        "  msa, deletion_matrix = pipeline.parsers.parse_a3m(a3m_lines[1])\n",
        "  msas.append([\"-\"*Ls[0]+seq for seq in msa])\n",
        "  deletion_matrices.append([[0]*Ls[0]+mtx for mtx in deletion_matrix])\n",
        "\n",
        "else:\n",
        "  msas.append([query_sequence])\n",
        "  deletion_matrices.append([[0]*len(query_sequence)])\n",
        "\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=msas, deletion_matrices=deletion_matrices),\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "6sgAOSxCLIfC"
      },
      "source": [
        "#@title Plot Number of Sequences per Position\n",
        "dpi =  100#@param {type:\"integer\"}\n",
        "# confidence per position\n",
        "plt.figure(dpi=dpi)\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": "hUYApPElB30u",
        "cellView": "form"
      },
      "source": [
        "#@title Predict structure\n",
        "plddts, paes = predict_structure(jobname, feature_dict, Ls=Ls, num_models=num_models)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oCZxmXKKTAxt",
        "cellView": "form"
      },
      "source": [
        "#@title Plot Predicted Alignment Error\n",
        "dpi =  100#@param {type:\"integer\"}\n",
        "\n",
        "# confidence per position\n",
        "plt.figure(figsize=(3*num_models,2), dpi=dpi)\n",
        "for n,(model_name,value) in enumerate(paes.items()):\n",
        "  plt.subplot(1,num_models,n+1)\n",
        "  plt.title(model_name)\n",
        "  plt.imshow(value,label=model_name,cmap=\"bwr\",vmin=0,vmax=30)\n",
        "  plt.colorbar()\n",
        "plt.savefig(jobname+\"_PAE.png\")\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "exKwNxDxF7IO",
        "cellView": "form"
      },
      "source": [
        "#@title Plot lDDT per residue\n",
        "# confidence per position\n",
        "dpi =  100#@param {type:\"integer\"}\n",
        "plt.figure(dpi=dpi)\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": "g-rPnOXdjf18",
        "cellView": "form"
      },
      "source": [
        "#@title Display 3D structure {run: \"auto\"}\n",
        "model_num = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\"] {type:\"raw\"}\n",
        "color = \"chain\" #@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(plddts[model_name])\n",
        "  for x in [len(query_sequence_a)]:\n",
        "    plt.plot([x,x],[0,100],color=\"black\")\n",
        "\n",
        "  plt.ylabel('plDDT')\n",
        "  plt.xlabel('position')\n",
        "  #########################################\n",
        "  plt.subplot(1,2,2);plt.title('Predicted Aligned Error')\n",
        "  plt.imshow(paes[model_name], 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",
        "  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(2),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": "33g5IIegij5R",
        "cellView": "form"
      },
      "source": [
        "#@title Package and download results\n",
        "!zip -FSr $jobname\".result.zip\" $jobname\".log\" $jobname\"_msa_coverage.png\" $jobname\"_\"*\"relaxed_model_\"*\".pdb\" $jobname\"_lDDT.png\"  $jobname\"_PAE.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\"."
      ]
    }
  ]
}