Representations_AlphaFold2_v3.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Representations_AlphaFold2_v2.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G4yBrceuFbf3"
      },
      "source": [
        "#AlphaFold2 w/ MMseqs2\n",
        "Easy to use version of AlphaFold 2 [(Jumper et al. 2021, Nature)](https://www.nature.com/articles/s41586-021-03819-2) a protein structure prediction pipeline, with an API hosted at the Södinglab based on the MMseqs2 server [(Mirdita et al. 2019, Bioinformatics)](https://academic.oup.com/bioinformatics/article/35/16/2856/5280135) for the multiple sequence alignment creation. \n",
        "\n",
        "**WARNING**: this notebook does NOT use the AlphaFold2 pipeline for MSA/template generation. It may give better or worse results depending on number of sequences that can be found. Check out the [full AlphaFold2 pipeline](https://github.com/deepmind/alphafold) or Deepmind's official [google-colab notebook](https://colab.research.google.com/github/deepmind/alphafold/blob/main/notebooks/AlphaFold.ipynb).\n",
        "\n",
        "<strong>For detailed instructions, see <a href=\"#Instructions\">bottom</a> of notebook!</strong>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kOblAo-xetgx",
        "cellView": "form",
        "colab": {
          "resources": {
            "http://localhost:8080/nbextensions/google.colab/files.js": {
              "data": "Ly8gQ29weXJpZ2h0IDIwMTcgR29vZ2xlIExMQwovLwovLyBMaWNlbnNlZCB1bmRlciB0aGUgQXBhY2hlIExpY2Vuc2UsIFZlcnNpb24gMi4wICh0aGUgIkxpY2Vuc2UiKTsKLy8geW91IG1heSBub3QgdXNlIHRoaXMgZmlsZSBleGNlcHQgaW4gY29tcGxpYW5jZSB3aXRoIHRoZSBMaWNlbnNlLgovLyBZb3UgbWF5IG9idGFpbiBhIGNvcHkgb2YgdGhlIExpY2Vuc2UgYXQKLy8KLy8gICAgICBodHRwOi8vd3d3LmFwYWNoZS5vcmcvbGljZW5zZXMvTElDRU5TRS0yLjAKLy8KLy8gVW5sZXNzIHJlcXVpcmVkIGJ5IGFwcGxpY2FibGUgbGF3IG9yIGFncmVlZCB0byBpbiB3cml0aW5nLCBzb2Z0d2FyZQovLyBkaXN0cmlidXRlZCB1bmRlciB0aGUgTGljZW5zZSBpcyBkaXN0cmlidXRlZCBvbiBhbiAiQVMgSVMiIEJBU0lTLAovLyBXSVRIT1VUIFdBUlJBTlRJRVMgT1IgQ09ORElUSU9OUyBPRiBBTlkgS0lORCwgZWl0aGVyIGV4cHJlc3Mgb3IgaW1wbGllZC4KLy8gU2VlIHRoZSBMaWNlbnNlIGZvciB0aGUgc3BlY2lmaWMgbGFuZ3VhZ2UgZ292ZXJuaW5nIHBlcm1pc3Npb25zIGFuZAovLyBsaW1pdGF0aW9ucyB1bmRlciB0aGUgTGljZW5zZS4KCi8qKgogKiBAZmlsZW92ZXJ2aWV3IEhlbHBlcnMgZm9yIGdvb2dsZS5jb2xhYiBQeXRob24gbW9kdWxlLgogKi8KKGZ1bmN0aW9uKHNjb3BlKSB7CmZ1bmN0aW9uIHNwYW4odGV4dCwgc3R5bGVBdHRyaWJ1dGVzID0ge30pIHsKICBjb25zdCBlbGVtZW50ID0gZG9jdW1lbnQuY3JlYXRlRWxlbWVudCgnc3BhbicpOwogIGVsZW1lbnQudGV4dENvbnRlbnQgPSB0ZXh0OwogIGZvciAoY29uc3Qga2V5IG9mIE9iamVjdC5rZXlzKHN0eWxlQXR0cmlidXRlcykpIHsKICAgIGVsZW1lbnQuc3R5bGVba2V5XSA9IHN0eWxlQXR0cmlidXRlc1trZXldOwogIH0KICByZXR1cm4gZWxlbWVudDsKfQoKLy8gTWF4IG51bWJlciBvZiBieXRlcyB3aGljaCB3aWxsIGJlIHVwbG9hZGVkIGF0IGEgdGltZS4KY29uc3QgTUFYX1BBWUxPQURfU0laRSA9IDEwMCAqIDEwMjQ7CgpmdW5jdGlvbiBfdXBsb2FkRmlsZXMoaW5wdXRJZCwgb3V0cHV0SWQpIHsKICBjb25zdCBzdGVwcyA9IHVwbG9hZEZpbGVzU3RlcChpbnB1dElkLCBvdXRwdXRJZCk7CiAgY29uc3Qgb3V0cHV0RWxlbWVudCA9IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKG91dHB1dElkKTsKICAvLyBDYWNoZSBzdGVwcyBvbiB0aGUgb3V0cHV0RWxlbWVudCB0byBtYWtlIGl0IGF2YWlsYWJsZSBmb3IgdGhlIG5leHQgY2FsbAogIC8vIHRvIHVwbG9hZEZpbGVzQ29udGludWUgZnJvbSBQeXRob24uCiAgb3V0cHV0RWxlbWVudC5zdGVwcyA9IHN0ZXBzOwoKICByZXR1cm4gX3VwbG9hZEZpbGVzQ29udGludWUob3V0cHV0SWQpOwp9CgovLyBUaGlzIGlzIHJvdWdobHkgYW4gYXN5bmMgZ2VuZXJhdG9yIChub3Qgc3VwcG9ydGVkIGluIHRoZSBicm93c2VyIHlldCksCi8vIHdoZXJlIHRoZXJlIGFyZSBtdWx0aXBsZSBhc3luY2hyb25vdXMgc3RlcHMgYW5kIHRoZSBQeXRob24gc2lkZSBpcyBnb2luZwovLyB0byBwb2xsIGZvciBjb21wbGV0aW9uIG9mIGVhY2ggc3RlcC4KLy8gVGhpcyB1c2VzIGEgUHJvbWlzZSB0byBibG9jayB0aGUgcHl0aG9uIHNpZGUgb24gY29tcGxldGlvbiBvZiBlYWNoIHN0ZXAsCi8vIHRoZW4gcGFzc2VzIHRoZSByZXN1bHQgb2YgdGhlIHByZXZpb3VzIHN0ZXAgYXMgdGhlIGlucHV0IHRvIHRoZSBuZXh0IHN0ZXAuCmZ1bmN0aW9uIF91cGxvYWRGaWxlc0NvbnRpbnVlKG91dHB1dElkKSB7CiAgY29uc3Qgb3V0cHV0RWxlbWVudCA9IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKG91dHB1dElkKTsKICBjb25zdCBzdGVwcyA9IG91dHB1dEVsZW1lbnQuc3RlcHM7CgogIGNvbnN0IG5leHQgPSBzdGVwcy5uZXh0KG91dHB1dEVsZW1lbnQubGFzdFByb21pc2VWYWx1ZSk7CiAgcmV0dXJuIFByb21pc2UucmVzb2x2ZShuZXh0LnZhbHVlLnByb21pc2UpLnRoZW4oKHZhbHVlKSA9PiB7CiAgICAvLyBDYWNoZSB0aGUgbGFzdCBwcm9taXNlIHZhbHVlIHRvIG1ha2UgaXQgYXZhaWxhYmxlIHRvIHRoZSBuZXh0CiAgICAvLyBzdGVwIG9mIHRoZSBnZW5lcmF0b3IuCiAgICBvdXRwdXRFbGVtZW50Lmxhc3RQcm9taXNlVmFsdWUgPSB2YWx1ZTsKICAgIHJldHVybiBuZXh0LnZhbHVlLnJlc3BvbnNlOwogIH0pOwp9CgovKioKICogR2VuZXJhdG9yIGZ1bmN0aW9uIHdoaWNoIGlzIGNhbGxlZCBiZXR3ZWVuIGVhY2ggYXN5bmMgc3RlcCBvZiB0aGUgdXBsb2FkCiAqIHByb2Nlc3MuCiAqIEBwYXJhbSB7c3RyaW5nfSBpbnB1dElkIEVsZW1lbnQgSUQgb2YgdGhlIGlucHV0IGZpbGUgcGlja2VyIGVsZW1lbnQuCiAqIEBwYXJhbSB7c3RyaW5nfSBvdXRwdXRJZCBFbGVtZW50IElEIG9mIHRoZSBvdXRwdXQgZGlzcGxheS4KICogQHJldHVybiB7IUl0ZXJhYmxlPCFPYmplY3Q+fSBJdGVyYWJsZSBvZiBuZXh0IHN0ZXBzLgogKi8KZnVuY3Rpb24qIHVwbG9hZEZpbGVzU3RlcChpbnB1dElkLCBvdXRwdXRJZCkgewogIGNvbnN0IGlucHV0RWxlbWVudCA9IGRvY3VtZW50LmdldEVsZW1lbnRCeUlkKGlucHV0SWQpOwogIGlucHV0RWxlbWVudC5kaXNhYmxlZCA9IGZhbHNlOwoKICBjb25zdCBvdXRwdXRFbGVtZW50ID0gZG9jdW1lbnQuZ2V0RWxlbWVudEJ5SWQob3V0cHV0SWQpOwogIG91dHB1dEVsZW1lbnQuaW5uZXJIVE1MID0gJyc7CgogIGNvbnN0IHBpY2tlZFByb21pc2UgPSBuZXcgUHJvbWlzZSgocmVzb2x2ZSkgPT4gewogICAgaW5wdXRFbGVtZW50LmFkZEV2ZW50TGlzdGVuZXIoJ2NoYW5nZScsIChlKSA9PiB7CiAgICAgIHJlc29sdmUoZS50YXJnZXQuZmlsZXMpOwogICAgfSk7CiAgfSk7CgogIGNvbnN0IGNhbmNlbCA9IGRvY3VtZW50LmNyZWF0ZUVsZW1lbnQoJ2J1dHRvbicpOwogIGlucHV0RWxlbWVudC5wYXJlbnRFbGVtZW50LmFwcGVuZENoaWxkKGNhbmNlbCk7CiAgY2FuY2VsLnRleHRDb250ZW50ID0gJ0NhbmNlbCB1cGxvYWQnOwogIGNvbnN0IGNhbmNlbFByb21pc2UgPSBuZXcgUHJvbWlzZSgocmVzb2x2ZSkgPT4gewogICAgY2FuY2VsLm9uY2xpY2sgPSAoKSA9PiB7CiAgICAgIHJlc29sdmUobnVsbCk7CiAgICB9OwogIH0pOwoKICAvLyBXYWl0IGZvciB0aGUgdXNlciB0byBwaWNrIHRoZSBmaWxlcy4KICBjb25zdCBmaWxlcyA9IHlpZWxkIHsKICAgIHByb21pc2U6IFByb21pc2UucmFjZShbcGlja2VkUHJvbWlzZSwgY2FuY2VsUHJvbWlzZV0pLAogICAgcmVzcG9uc2U6IHsKICAgICAgYWN0aW9uOiAnc3RhcnRpbmcnLAogICAgfQogIH07CgogIGNhbmNlbC5yZW1vdmUoKTsKCiAgLy8gRGlzYWJsZSB0aGUgaW5wdXQgZWxlbWVudCBzaW5jZSBmdXJ0aGVyIHBpY2tzIGFyZSBub3QgYWxsb3dlZC4KICBpbnB1dEVsZW1lbnQuZGlzYWJsZWQgPSB0cnVlOwoKICBpZiAoIWZpbGVzKSB7CiAgICByZXR1cm4gewogICAgICByZXNwb25zZTogewogICAgICAgIGFjdGlvbjogJ2NvbXBsZXRlJywKICAgICAgfQogICAgfTsKICB9CgogIGZvciAoY29uc3QgZmlsZSBvZiBmaWxlcykgewogICAgY29uc3QgbGkgPSBkb2N1bWVudC5jcmVhdGVFbGVtZW50KCdsaScpOwogICAgbGkuYXBwZW5kKHNwYW4oZmlsZS5uYW1lLCB7Zm9udFdlaWdodDogJ2JvbGQnfSkpOwogICAgbGkuYXBwZW5kKHNwYW4oCiAgICAgICAgYCgke2ZpbGUudHlwZSB8fCAnbi9hJ30pIC0gJHtmaWxlLnNpemV9IGJ5dGVzLCBgICsKICAgICAgICBgbGFzdCBtb2RpZmllZDogJHsKICAgICAgICAgICAgZmlsZS5sYXN0TW9kaWZpZWREYXRlID8gZmlsZS5sYXN0TW9kaWZpZWREYXRlLnRvTG9jYWxlRGF0ZVN0cmluZygpIDoKICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgICAgJ24vYSd9IC0gYCkpOwogICAgY29uc3QgcGVyY2VudCA9IHNwYW4oJzAlIGRvbmUnKTsKICAgIGxpLmFwcGVuZENoaWxkKHBlcmNlbnQpOwoKICAgIG91dHB1dEVsZW1lbnQuYXBwZW5kQ2hpbGQobGkpOwoKICAgIGNvbnN0IGZpbGVEYXRhUHJvbWlzZSA9IG5ldyBQcm9taXNlKChyZXNvbHZlKSA9PiB7CiAgICAgIGNvbnN0IHJlYWRlciA9IG5ldyBGaWxlUmVhZGVyKCk7CiAgICAgIHJlYWRlci5vbmxvYWQgPSAoZSkgPT4gewogICAgICAgIHJlc29sdmUoZS50YXJnZXQucmVzdWx0KTsKICAgICAgfTsKICAgICAgcmVhZGVyLnJlYWRBc0FycmF5QnVmZmVyKGZpbGUpOwogICAgfSk7CiAgICAvLyBXYWl0IGZvciB0aGUgZGF0YSB0byBiZSByZWFkeS4KICAgIGxldCBmaWxlRGF0YSA9IHlpZWxkIHsKICAgICAgcHJvbWlzZTogZmlsZURhdGFQcm9taXNlLAogICAgICByZXNwb25zZTogewogICAgICAgIGFjdGlvbjogJ2NvbnRpbnVlJywKICAgICAgfQogICAgfTsKCiAgICAvLyBVc2UgYSBjaHVua2VkIHNlbmRpbmcgdG8gYXZvaWQgbWVzc2FnZSBzaXplIGxpbWl0cy4gU2VlIGIvNjIxMTU2NjAuCiAgICBsZXQgcG9zaXRpb24gPSAwOwogICAgZG8gewogICAgICBjb25zdCBsZW5ndGggPSBNYXRoLm1pbihmaWxlRGF0YS5ieXRlTGVuZ3RoIC0gcG9zaXRpb24sIE1BWF9QQVlMT0FEX1NJWkUpOwogICAgICBjb25zdCBjaHVuayA9IG5ldyBVaW50OEFycmF5KGZpbGVEYXRhLCBwb3NpdGlvbiwgbGVuZ3RoKTsKICAgICAgcG9zaXRpb24gKz0gbGVuZ3RoOwoKICAgICAgY29uc3QgYmFzZTY0ID0gYnRvYShTdHJpbmcuZnJvbUNoYXJDb2RlLmFwcGx5KG51bGwsIGNodW5rKSk7CiAgICAgIHlpZWxkIHsKICAgICAgICByZXNwb25zZTogewogICAgICAgICAgYWN0aW9uOiAnYXBwZW5kJywKICAgICAgICAgIGZpbGU6IGZpbGUubmFtZSwKICAgICAgICAgIGRhdGE6IGJhc2U2NCwKICAgICAgICB9LAogICAgICB9OwoKICAgICAgbGV0IHBlcmNlbnREb25lID0gZmlsZURhdGEuYnl0ZUxlbmd0aCA9PT0gMCA/CiAgICAgICAgICAxMDAgOgogICAgICAgICAgTWF0aC5yb3VuZCgocG9zaXRpb24gLyBmaWxlRGF0YS5ieXRlTGVuZ3RoKSAqIDEwMCk7CiAgICAgIHBlcmNlbnQudGV4dENvbnRlbnQgPSBgJHtwZXJjZW50RG9uZX0lIGRvbmVgOwoKICAgIH0gd2hpbGUgKHBvc2l0aW9uIDwgZmlsZURhdGEuYnl0ZUxlbmd0aCk7CiAgfQoKICAvLyBBbGwgZG9uZS4KICB5aWVsZCB7CiAgICByZXNwb25zZTogewogICAgICBhY3Rpb246ICdjb21wbGV0ZScsCiAgICB9CiAgfTsKfQoKc2NvcGUuZ29vZ2xlID0gc2NvcGUuZ29vZ2xlIHx8IHt9OwpzY29wZS5nb29nbGUuY29sYWIgPSBzY29wZS5nb29nbGUuY29sYWIgfHwge307CnNjb3BlLmdvb2dsZS5jb2xhYi5fZmlsZXMgPSB7CiAgX3VwbG9hZEZpbGVzLAogIF91cGxvYWRGaWxlc0NvbnRpbnVlLAp9Owp9KShzZWxmKTsK",
              "ok": true,
              "headers": [
                [
                  "content-type",
                  "application/javascript"
                ]
              ],
              "status": 200,
              "status_text": ""
            }
          },
          "base_uri": "https://localhost:8080/",
          "height": 89
        },
        "outputId": "f078b90b-f02f-41c5-f632-5e2831c991f9"
      },
      "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 = 'MTVKTEAAKGTLTYSRMRGMVAILIAFMKQRRMGLNDFIQKIANNSYACKHPEVQSILKISQPQEPELMNANPSPPPSPSQQINLGPSSNPHAKPSDFHFLKVIGKGSFGKVLLARHKAEEVFYAVKVLQKKAILKKKEEKHIMSERNVLLKNVKHPFLVGLHFSFQTADKLYFVLDYINGGELFYHLQRERCFLEPRARFYAAEIASALGYLHSLNIVYRDLKPENILLDSQGHIVLTDFGLCKENIEHNSTTSTFCGTPEYLAPEVLHKQPYDRTVDWWCLGAVLYEMLYGLPPFYSRNTAEMYDNILNKPLQLKPNITNSARHLLEGLLQKDRTKRLGAKDDFMEIKSHVFFSLINWDDLINKKITPPFNPNVSGPNDLRHFDPEFTEEPVPNSIGKSPDSVLVTASVKEAAEAFLGFSYAPPTDSFL' #@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 = 'O00141' #@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 = \"custom\" #@param [\"MMseqs2 (UniRef+Environmental)\", \"MMseqs2 (UniRef only)\",\"single_sequence\",\"custom\"]\n",
        "num_models = 1 #@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 = True #@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": 50,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "     <input type=\"file\" id=\"files-64047f67-45fa-416a-9fa7-fe85e3f7ccdc\" name=\"files[]\" multiple disabled\n",
              "        style=\"border:none\" />\n",
              "     <output id=\"result-64047f67-45fa-416a-9fa7-fe85e3f7ccdc\">\n",
              "      Upload widget is only available when the cell has been executed in the\n",
              "      current browser session. Please rerun this cell to enable.\n",
              "      </output>\n",
              "      <script src=\"/nbextensions/google.colab/files.js\"></script> "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Saving 002_O00141.a3m to 002_O00141.a3m\n",
            "moving 002_O00141.a3m to O00141_26d52.custom.a3m\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "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",
        "    git clone https://github.com/xinformatics/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": 51,
      "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]/100, 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]/100, idx_res, chains)\n",
        "\n",
        "    out[f\"model_{n+1}\"] = {\"plddt\":plddts[r], \"pae\":paes[r]}\n",
        "  return out, prediction_result"
      ],
      "execution_count": 52,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A9tUpDaikPC8",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "cellView": "form",
        "outputId": "bee71638-134f-4636-f33c-b1c26648e4f9"
      },
      "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://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\" ] || [ \"${STATUS}\" == \"PENDING\" ]; 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",
        "      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": 53,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "submitting job\n",
            "getting templates\n",
            "pdb\tevalue\n",
            "2R5T_A\t3.860E-116\n",
            "3HDM_A\t3.860E-116\n",
            "3TXO_A\t2.374E-106\n",
            "2UZV_A\t3.856E-102\n",
            "3PFQ_A\t6.221E-98\n",
            "5F9E_A\t9.212E-96\n",
            "1XJD_A\t1.720E-95\n",
            "4WB7_B\t1.720E-95\n",
            "6BYR_C\t7.292E-94\n",
            "3WF7_A\t6.483E-93\n",
            "4L44_A\t8.858E-93\n",
            "4RA4_A\t2.259E-92\n",
            "3A8X_A\t3.087E-92\n",
            "4WOT_B\t1.076E-91\n",
            "2JDO_A\t5.119E-91\n",
            "5LI9_A\t1.783E-90\n",
            "4AW2_A\t6.209E-90\n",
            "5N1F_A\t1.159E-89\n",
            "5N1G_A\t1.159E-89\n",
            "1FOT_A\t1.583E-89\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hUYApPElB30u",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "cellView": "form",
        "outputId": "be28a409-99d1-47f0-f9f8-f771f7099989"
      },
      "source": [
        "#@title Gather input features, predict structure { vertical-output: true }\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, prediction_result = 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": [
        {
          "output_type": "stream",
          "text": [
            "running model_1\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6xbvRNrwnJqj",
        "cellView": "form"
      },
      "source": [
        "# #@title Make plots\n",
        "\n",
        "# ##################################################################\n",
        "# plt.figure(figsize=(12,4),dpi=100)\n",
        "# plt.subplot(1,2,1); plt.title(\"Number of sequences per position\")\n",
        "# plt.plot((feature_dict[\"msa\"] != 21).sum(0))\n",
        "# plt.xlabel(\"positions\")\n",
        "# plt.ylabel(\"number of sequences\")\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",
        "# 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": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "svaADwocVdwl",
        "cellView": "form"
      },
      "source": [
        "#@title Display 3D structure\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_name):\n",
        "  plt.figure(figsize=(6, 10))\n",
        "  \"\"\"Plots the legend for plDDT.\"\"\"\n",
        "  #########################################\n",
        "  plt.subplot(2,1,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.xlabel('Residue')\n",
        "  plt.ylabel('plDDT')\n",
        "  #########################################\n",
        "  plt.subplot(2,1,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_name, show_sidechains, show_mainchain, color):\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(width=800, height=600)\n",
        "  view.addModel(open(pdb_filename,'r').read(),'pdb')\n",
        "  if color == \"lDDT\":\n",
        "    view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':0.5,'max':0.9}}})\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_mainchain:\n",
        "    BB = ['C','O','N','CA']\n",
        "    view.addStyle({'atom':BB},{'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "\n",
        "  view.zoomTo()\n",
        "  if color == \"lDDT\":\n",
        "    grid = GridspecLayout(6, 2)\n",
        "    out = Output()\n",
        "    with out: view.show()\n",
        "    grid[:-1,0] = out\n",
        "    out = Output()\n",
        "    with out: plot_confidence(model_name).show()\n",
        "    grid[:,1] = out\n",
        "    out = Output()\n",
        "    with out: plot_plddt_legend().show()\n",
        "    grid[-1,0] = out\n",
        "    return grid\n",
        "  else:\n",
        "    return view.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=['rainbow', 'lDDT', 'chain'], value='lDDT'))"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "33g5IIegij5R",
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 170
        },
        "outputId": "5bee7cfd-af0d-4d2c-8d5e-2797925af403"
      },
      "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",
        "}\n",
        "\n",
        "to_cite = [ \"Jumper2021\", \"Ovchinnikov2021\" ]\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\"\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": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Found 5 citations for tools or databases.\n",
            "Don't forget to cite your custom MSA generation method.\n",
            "  adding: B2RXH2_5b7a0.log (deflated 18%)\n",
            "  adding: B2RXH2_5b7a0.custom.a3m (deflated 77%)\n",
            "  adding: B2RXH2_5b7a0_relaxed_model_1.pdb (deflated 78%)\n",
            "  adding: B2RXH2_5b7a0_unrelaxed_model_1.pdb (deflated 78%)\n",
            "  adding: B2RXH2_5b7a0_coverage_lDDT.png (deflated 5%)\n",
            "  adding: B2RXH2_5b7a0.bibtex (deflated 50%)\n",
            "  adding: B2RXH2_5b7a0_PAE.png (deflated 1%)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/javascript": [
              "\n",
              "    async function download(id, filename, size) {\n",
              "      if (!google.colab.kernel.accessAllowed) {\n",
              "        return;\n",
              "      }\n",
              "      const div = document.createElement('div');\n",
              "      const label = document.createElement('label');\n",
              "      label.textContent = `Downloading \"${filename}\": `;\n",
              "      div.appendChild(label);\n",
              "      const progress = document.createElement('progress');\n",
              "      progress.max = size;\n",
              "      div.appendChild(progress);\n",
              "      document.body.appendChild(div);\n",
              "\n",
              "      const buffers = [];\n",
              "      let downloaded = 0;\n",
              "\n",
              "      const channel = await google.colab.kernel.comms.open(id);\n",
              "      // Send a message to notify the kernel that we're ready.\n",
              "      channel.send({})\n",
              "\n",
              "      for await (const message of channel.messages) {\n",
              "        // Send a message to notify the kernel that we're ready.\n",
              "        channel.send({})\n",
              "        if (message.buffers) {\n",
              "          for (const buffer of message.buffers) {\n",
              "            buffers.push(buffer);\n",
              "            downloaded += buffer.byteLength;\n",
              "            progress.value = downloaded;\n",
              "          }\n",
              "        }\n",
              "      }\n",
              "      const blob = new Blob(buffers, {type: 'application/binary'});\n",
              "      const a = document.createElement('a');\n",
              "      a.href = window.URL.createObjectURL(blob);\n",
              "      a.download = filename;\n",
              "      div.appendChild(a);\n",
              "      a.click();\n",
              "      div.remove();\n",
              "    }\n",
              "  "
            ],
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "application/javascript": [
              "download(\"download_8737b838-c9b3-46ef-ac43-8f019b39cc33\", \"B2RXH2_5b7a0.result.zip\", 400786)"
            ],
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UGUBLzB3C6WN"
      },
      "source": [
        "# Instructions <a name=\"Instructions\"></a>\n",
        "**Quick start**\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 zip file contents**\n",
        "\n",
        "1. PDB formatted 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",
        "6. BibTeX file with citations for all used tools and databases.\n",
        "\n",
        "At the end of the job a download modal box will pop up with a `jobname.result.zip` file. Additionally, if the `save_to_google_drive` option was selected, the `jobname.result.zip` will be uploaded to your Google Drive.\n",
        "\n",
        "**Using a custom MSA as input**\n",
        "\n",
        "To predict the structure with a custom MSA (A3M formatted): (1) Change the msa_mode: to \"custom\", (2) Wait for an upload box to appear at the end of the \"Input Protein ...\" box. Upload your A3M. The first fasta entry of the A3M must be the query sequence without gaps.\n",
        "\n",
        "To generate good input MSAs the HHblits server can be used here: https://toolkit.tuebingen.mpg.de/tools/hhblits\n",
        "\n",
        "After submitting your query, click \"Query Template MSA\" -> \"Download Full A3M\". Download the a3m file and upload it to the notebook.\n",
        "\n",
        "**Troubleshooting**\n",
        "* Try to restart the session \"Runtime\" -> \"Factory reset runtime\".\n",
        "* Check your input sequence.\n",
        "\n",
        "**Known issues**\n",
        "* Colab assigns different types of GPUs with varying amount of memory. Some might have not enough memory to predict the structure.\n",
        "* Your browser can block the pop-up for downloading the result file. You can choose the `save_to_google_drive` option to upload to Google Drive instead or manually download the result file: 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",
        "**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: Our MMseqs2 API can probably handle ~20k requests per day.\n",
        "* For best results, we recommend using the full pipeline: https://github.com/deepmind/alphafold\n",
        "\n",
        "**Description of the plots**\n",
        "*   **Number of sequences per position** - We want to see at least 30 sequences per position, for best performance, ideally 100 sequences.\n",
        "*   **Predicted lDDT per position** - model confidence (out of 100) at each position. Higher the better.\n",
        "*   **Predicted Alignment Error** - For homooligomers, this could be a useful metric to assess how confident the model is about the interface. Lower the better.\n",
        "\n",
        "\n",
        "\n",
        "**Bugs**\n",
        "- If you encounter any bugs, please report the issue to https://github.com/sokrypton/ColabFold/issues\n",
        "\n",
        "\n",
        "**Acknowledgments**\n",
        "- We would like to thank the AlphaFold team for developing an excellent model and open sourcing the software. \n",
        "\n",
        "- A colab by Sergey Ovchinnikov ([@sokrypton](https://twitter.com/sokrypton)), Milot Mirdita ([@milot_mirdita](https://twitter.com/milot_mirdita)) and Martin Steinegger ([@thesteinegger](https://twitter.com/thesteinegger)).\n",
        "\n",
        "- Minkyung Baek ([@minkbaek](https://twitter.com/minkbaek)) and Yoshitaka Moriwaki ([@Ag_smith](https://twitter.com/Ag_smith)) for protein-complex prediction proof-of-concept in AlphaFold2.\n",
        "\n",
        "- 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)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J4k3xA6ghTWY"
      },
      "source": [
        "####### Representations"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "el92tlerhTb6",
        "outputId": "0379d3bb-97ff-4790-a1c8-331cf06dd21f"
      },
      "source": [
        "prediction_result.keys()"
      ],
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dict_keys(['distogram', 'experimentally_resolved', 'masked_msa', 'predicted_aligned_error', 'predicted_lddt', 'representations', 'structure_module', 'plddt', 'aligned_confidence_probs', 'max_predicted_aligned_error', 'ptm'])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 30
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dOrqNZC7hTf2"
      },
      "source": [
        "embeddings = prediction_result['representations']"
      ],
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "33R9j9NfhTjy",
        "outputId": "2734a137-9835-4d97-a5f6-4967c82e8889"
      },
      "source": [
        "embeddings.keys()"
      ],
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "dict_keys(['msa', 'msa_first_row', 'pair', 'single', 'structure_module'])"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 32
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9FhrpDekhTne",
        "outputId": "0427336d-d7a2-428a-cb5a-12bce9d05495"
      },
      "source": [
        "embeddings['msa'].shape"
      ],
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(508, 505, 256)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 33
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5aWv1D0Dvm-T",
        "outputId": "81cc4ceb-8a99-46a3-b720-59847962c06c"
      },
      "source": [
        "prediction_result['representations']['msa_first_row'].shape"
      ],
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(505, 256)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 34
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yeEx1dC5wA1H",
        "outputId": "52b435bc-22ed-4ef3-f3f0-1f275374df8d"
      },
      "source": [
        "prediction_result['representations']['pair'].shape"
      ],
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(505, 505, 128)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Y51_Sw-YwGDk",
        "outputId": "79cfd5a0-00a3-4293-e84a-facd7feb6277"
      },
      "source": [
        "prediction_result['representations']['single'].shape"
      ],
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(505, 384)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-iofwRUy1bHa"
      },
      "source": [
        "embed = prediction_result['representations']['single']"
      ],
      "execution_count": 47,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QguYLr8y1jc1"
      },
      "source": [
        "name = jobname.split('_')[0]"
      ],
      "execution_count": 39,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fRyCeffI1sJ5"
      },
      "source": [
        "import pickle"
      ],
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MyPUq28A262X"
      },
      "source": [
        "with open(name+'.pickle', 'wb') as handle:\n",
        "    pickle.dump(embed, handle, protocol=pickle.HIGHEST_PROTOCOL)"
      ],
      "execution_count": 48,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ziPFNfOt3GDc",
        "outputId": "a388670c-c80b-44e2-8ac0-f4239a02c2db"
      },
      "source": [
        "!ls -sh"
      ],
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "total 26M\n",
            "   0 AF2_READY\t\t\t\t 760K B2RXH2.pickle\n",
            "4.0K alphafold\t\t\t\t    0 CONDA_READY\n",
            "4.0K alphafold_\t\t\t\t    0 HH_READY\n",
            "460K B2RXH2_5b7a0.a3m\t\t\t    0 MMSEQ2_READY\n",
            " 22M B2RXH2_5b7a0_a3m.ffdata\t\t 4.0K params\n",
            "4.0K B2RXH2_5b7a0_a3m.ffindex\t\t 4.0K sample_data\n",
            "   0 B2RXH2_5b7a0_cs219.ffdata\t\t 4.0K temp_1783c.a3m\n",
            "4.0K B2RXH2_5b7a0_cs219.ffindex\t\t  40K temp_1783c_coverage_lDDT.png\n",
            "160K B2RXH2_5b7a0.custom.a3m\t\t 4.0K temp_1783c.fasta\n",
            "4.0K B2RXH2_5b7a0.fasta\t\t\t 4.0K temp_1783c.log\n",
            "1.5M B2RXH2_5b7a0_hhm.ffdata\t\t    0 temp_1783c.m8\n",
            "4.0K B2RXH2_5b7a0_hhm.ffindex\t\t 4.0K temp_1783c.mmseqs2.tar.gz\n",
            "4.0K B2RXH2_5b7a0.log\t\t\t 8.0K temp_1783c_PAE.png\n",
            "4.0K B2RXH2_5b7a0.m8\t\t\t 8.0K temp_1783c_unrelaxed_model_1.pdb\n",
            "104K B2RXH2_5b7a0.mmseqs2.tar.gz\t 4.0K templates\n",
            "316K B2RXH2_5b7a0_unrelaxed_model_1.pdb\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Kkln9JSi3VJa"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}