community/reproduce/AlphaFold2-Chinese/Fold_CN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "AlphaFold2中文版蛋白质预测模型使用指南(基于Deepmind与Mindspore开源框架).ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.10"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G4yBrceuFbf3"
      },
      "source": [
        "<img src=\"https://raw.githubusercontent.com/sokrypton/ColabFold/main/.github/ColabFold_Marv_Logo_Small.png\" height=\"200\" align=\"right\" style=\"height:240px\">\n",
        "\n",
        "##AlphaFold2_CN: AlphaFold2 with MMseqs2\n",
        "\n",
        "简单的中文版蛋白质结构预测操作指南（中文版），基于[AlphaFold2](https://www.nature.com/articles/s41586-021-03819-2)和[Alphafold2-multimer](https://www.biorxiv.org/content/10.1101/2021.10.04.463034v1). 序列比对方式基于[MMseqs2](mmseqs.com)和[HHsearch](https://github.com/soedinglab/hh-suite)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kOblAo-xetgx",
        "cellView": "form"
      },
      "source": [
        "#@title 输入蛋白质序列（默认为视频中的测试序列）\n",
        "from google.colab import files\n",
        "import os.path\n",
        "import re\n",
        "import hashlib\n",
        "import random\n",
        "\n",
        "def add_hash(x,y):\n",
        "  return x+\"_\"+hashlib.sha1(y.encode()).hexdigest()[:5]\n",
        "\n",
        "query_sequence = 'MAAHKGAEHHHKAAEHHEQAAKHHHAAAEHHEKGEHEQAAHHADTAYAHHKHAEEHAAQAAKHDAEHHAPKPH' #@param {type:\"string\"}\n",
        "#@markdown  - Use `:` to specify inter-protein chainbreaks for **modeling complexes** (supports homo- and hetro-oligomers). For example **PI...SK:PI...SK** for a homodimer\n",
        "\n",
        "# remove whitespaces\n",
        "query_sequence = \"\".join(query_sequence.split())\n",
        "\n",
        "jobname = 'test' #@param {type:\"string\"}\n",
        "# remove whitespaces\n",
        "basejobname = \"\".join(jobname.split())\n",
        "basejobname = re.sub(r'\\W+', '', basejobname)\n",
        "jobname = add_hash(basejobname, query_sequence)\n",
        "while os.path.isfile(f\"{jobname}.csv\"):\n",
        "  jobname = add_hash(basejobname, ''.join(random.sample(query_sequence,len(query_sequence))))\n",
        "\n",
        "with open(f\"{jobname}.csv\", \"w\") as text_file:\n",
        "    text_file.write(f\"id,sequence\\n{jobname},{query_sequence}\")\n",
        "\n",
        "queries_path=f\"{jobname}.csv\"\n",
        "\n",
        "# number of models to use\n",
        "use_amber = False #@param {type:\"boolean\"}\n",
        "template_mode = \"none\" #@param [\"none\", \"pdb70\",\"custom\"]\n",
        "#@markdown - \"none\" = no template information is used, \"pdb70\" = detect templates in pdb70, \"custom\" - upload and search own templates (PDB or mmCIF format, see [notes below](#custom_templates))\n",
        "\n",
        "if template_mode == \"pdb70\":\n",
        "  use_templates = True\n",
        "  custom_template_path = None\n",
        "elif template_mode == \"custom\":\n",
        "  custom_template_path = f\"{jobname}_template\"\n",
        "  os.mkdir(custom_template_path)\n",
        "  uploaded = files.upload()\n",
        "  use_templates = True\n",
        "  for fn in uploaded.keys():\n",
        "    os.rename(fn, f\"{jobname}_template/{fn}\")\n",
        "else:\n",
        "  custom_template_path = None\n",
        "  use_templates = False\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ### MSA选项(custom MSA upload, single sequence, pairing mode)\n",
        "msa_mode = \"MMseqs2 (UniRef+Environmental)\" #@param [\"MMseqs2 (UniRef+Environmental)\", \"MMseqs2 (UniRef only)\",\"single_sequence\",\"custom\"]\n",
        "pair_mode = \"unpaired+paired\" #@param [\"unpaired+paired\",\"paired\",\"unpaired\"] {type:\"string\"}\n",
        "#@markdown - \"unpaired+paired\" = pair sequences from same species + unpaired MSA, \"unpaired\" = separate MSA for each chain, \"paired\" - only use paired sequences.\n",
        "\n",
        "# decide which a3m to use\n",
        "if msa_mode.startswith(\"MMseqs2\"):\n",
        "  a3m_file = f\"{jobname}.a3m\"\n",
        "elif msa_mode == \"custom\":\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 not line.rstrip():\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",
        "    queries_path=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)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "C2_sh2uAonJH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ### 参数设置\n",
        "model_type = \"auto\" #@param [\"auto\", \"AlphaFold2-ptm\", \"AlphaFold2-multimer-v1\", \"AlphaFold2-multimer-v2\"]\n",
        "#@markdown - \"auto\" = protein structure prediction using \"AlphaFold2-ptm\" and complex prediction \"AlphaFold-multimer-v2\". For complexes \"AlphaFold-multimer-v[1,2]\" and \"AlphaFold-ptm\" can be used.\n",
        "num_recycles = 3 #@param [1,3,6,12,24,48] {type:\"raw\"}\n",
        "save_to_google_drive = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown -  if the save_to_google_drive option was selected, the result zip will be uploaded to your Google Drive\n",
        "dpi = 200 #@param {type:\"integer\"}\n",
        "#@markdown - set dpi for image resolution\n",
        "\n",
        "#@markdown Don't forget to hit `Runtime` -> `Run all` after updating the form.\n",
        "\n",
        "\n",
        "if save_to_google_drive:\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!\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "ADDuaolKmjGW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iccGdbe_Pmt9",
        "pycharm": {
          "name": "#%%\n"
        },
        "cellView": "form"
      },
      "source": [
        "#@title 环境安装\n",
        "%%bash -s $use_amber $use_templates\n",
        "\n",
        "set -e\n",
        "\n",
        "USE_AMBER=$1\n",
        "USE_TEMPLATES=$2\n",
        "\n",
        "if [ ! -f COLABFOLD_READY ]; then\n",
        "  # install dependencies\n",
        "  # We have to use \"--no-warn-conflicts\" because colab already has a lot preinstalled with requirements different to ours\n",
        "  pip install -q --no-warn-conflicts \"colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold\"\n",
        "  # high risk high gain\n",
        "  pip install -q \"jax[cuda11_cudnn805]>=0.3.8,<0.4\" -f https://storage.googleapis.com/jax-releases/jax_releases.html\n",
        "  touch COLABFOLD_READY\n",
        "fi\n",
        "\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 kalign2=2.04 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",
        "  touch AMBER_READY\n",
        "fi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_sztQyz29DIC",
        "cellView": "form"
      },
      "source": [
        "#@title 模型预测\n",
        "\n",
        "import sys\n",
        "\n",
        "from colabfold.download import download_alphafold_params, default_data_dir\n",
        "from colabfold.utils import setup_logging\n",
        "from colabfold.batch import get_queries, run, set_model_type\n",
        "K80_chk = !nvidia-smi | grep \"Tesla K80\" | wc -l\n",
        "if \"1\" in K80_chk:\n",
        "  print(\"WARNING: found GPU Tesla K80: limited to total length < 1000\")\n",
        "  if \"TF_FORCE_UNIFIED_MEMORY\" in os.environ:\n",
        "    del os.environ[\"TF_FORCE_UNIFIED_MEMORY\"]\n",
        "  if \"XLA_PYTHON_CLIENT_MEM_FRACTION\" in os.environ:\n",
        "    del os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"]\n",
        "\n",
        "from colabfold.colabfold import plot_protein\n",
        "from pathlib import Path\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "# For some reason we need that to get pdbfixer to import\n",
        "if use_amber and '/usr/local/lib/python3.7/site-packages/' not in sys.path:\n",
        "    sys.path.insert(0, '/usr/local/lib/python3.7/site-packages/')\n",
        "\n",
        "def prediction_callback(unrelaxed_protein, length, prediction_result, input_features, type):\n",
        "  fig = plot_protein(unrelaxed_protein, Ls=length, dpi=150)\n",
        "  plt.show()\n",
        "  plt.close()\n",
        "\n",
        "result_dir=\".\"\n",
        "setup_logging(Path(\".\").joinpath(\"log.txt\"))\n",
        "queries, is_complex = get_queries(queries_path)\n",
        "model_type = set_model_type(is_complex, model_type)\n",
        "download_alphafold_params(model_type, Path(\".\"))\n",
        "run(\n",
        "    queries=queries,\n",
        "    result_dir=result_dir,\n",
        "    use_templates=use_templates,\n",
        "    custom_template_path=custom_template_path,\n",
        "    use_amber=use_amber,\n",
        "    msa_mode=msa_mode,    \n",
        "    model_type=model_type,\n",
        "    num_models=5,\n",
        "    num_recycles=num_recycles,\n",
        "    model_order=[1, 2, 3, 4, 5],\n",
        "    is_complex=is_complex,\n",
        "    data_dir=Path(\".\"),\n",
        "    keep_existing_results=False,\n",
        "    recompile_padding=1.0,\n",
        "    rank_by=\"auto\",\n",
        "    pair_mode=pair_mode,\n",
        "    stop_at_score=float(100),\n",
        "    prediction_callback=prediction_callback,\n",
        "    dpi=dpi\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KK7X9T44pWb7",
        "cellView": "form"
      },
      "source": [
        "#@title 展示3维结构 {run: \"auto\"}\n",
        "import py3Dmol\n",
        "import glob\n",
        "import matplotlib.pyplot as plt\n",
        "from colabfold.colabfold import plot_plddt_legend\n",
        "rank_num = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\"] {type:\"raw\"}\n",
        "color = \"lDDT\" #@param [\"chain\", \"lDDT\", \"rainbow\"]\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "\n",
        "jobname_prefix = \".custom\" if msa_mode == \"custom\" else \"\"\n",
        "if use_amber:\n",
        "  pdb_filename = f\"{jobname}{jobname_prefix}_relaxed_rank_{rank_num}_model_*.pdb\"\n",
        "else:\n",
        "  pdb_filename = f\"{jobname}{jobname_prefix}_unrelaxed_rank_{rank_num}_model_*.pdb\"\n",
        "\n",
        "pdb_file = glob.glob(pdb_filename)\n",
        "\n",
        "def show_pdb(rank_num=1, show_sidechains=False, show_mainchains=False, color=\"lDDT\"):\n",
        "  model_name = f\"rank_{rank_num}\"\n",
        "  view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js',)\n",
        "  view.addModel(open(pdb_file[0],'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",
        "    chains = len(queries[0][1]) + 1 if is_complex else 1\n",
        "    for n,chain,color in zip(range(chains),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",
        "\n",
        "show_pdb(rank_num,show_sidechains, show_mainchains, color).show()\n",
        "if color == \"lDDT\":\n",
        "  plot_plddt_legend().show() "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "11l8k--10q0C",
        "cellView": "form"
      },
      "source": [
        "#@title 图表展示 {run: \"auto\"}\n",
        "from IPython.display import display, HTML\n",
        "import base64\n",
        "from html import escape\n",
        "\n",
        "# see: https://stackoverflow.com/a/53688522\n",
        "def image_to_data_url(filename):\n",
        "  ext = filename.split('.')[-1]\n",
        "  prefix = f'data:image/{ext};base64,'\n",
        "  with open(filename, 'rb') as f:\n",
        "    img = f.read()\n",
        "  return prefix + base64.b64encode(img).decode('utf-8')\n",
        "\n",
        "pae = image_to_data_url(f\"{jobname}{jobname_prefix}_PAE.png\")\n",
        "cov = image_to_data_url(f\"{jobname}{jobname_prefix}_coverage.png\")\n",
        "plddt = image_to_data_url(f\"{jobname}{jobname_prefix}_plddt.png\")\n",
        "display(HTML(f\"\"\"\n",
        "<style>\n",
        "  img {{\n",
        "    float:left;\n",
        "  }}\n",
        "  .full {{\n",
        "    max-width:100%;\n",
        "  }}\n",
        "  .half {{\n",
        "    max-width:50%;\n",
        "  }}\n",
        "  @media (max-width:640px) {{\n",
        "    .half {{\n",
        "      max-width:100%;\n",
        "    }}\n",
        "  }}\n",
        "</style>\n",
        "<div style=\"max-width:90%; padding:2em;\">\n",
        "  <h1>Plots for {escape(jobname)}</h1>\n",
        "  <img src=\"{pae}\" class=\"full\" />\n",
        "  <img src=\"{cov}\" class=\"half\" />\n",
        "  <img src=\"{plddt}\" class=\"half\" />\n",
        "</div>\n",
        "\"\"\"))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "33g5IIegij5R",
        "cellView": "form"
      },
      "source": [
        "#@title结果下载\n",
        "#@markdown If you are 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",
        "if msa_mode == \"custom\":\n",
        "  print(\"Don't forget to cite your custom MSA generation method.\")\n",
        "\n",
        "!zip -FSr $jobname\".result.zip\" config.json $jobname*\".json\" $jobname*\".a3m\" $jobname*\"relaxed_rank_\"*\".pdb\" \"cite.bibtex\" $jobname*\".png\"\n",
        "files.download(f\"{jobname}.result.zip\")\n",
        "\n",
        "if save_to_google_drive == True and drive:\n",
        "  uploaded = drive.CreateFile({'title': f\"{jobname}.result.zip\"})\n",
        "  uploaded.SetContentFile(f\"{jobname}.result.zip\")\n",
        "  uploaded.Upload()\n",
        "  print(f\"Uploaded {jobname}.result.zip to Google Drive with ID {uploaded.get('id')}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UGUBLzB3C6WN",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "# 操作指南 <a name=\"Instructions\"></a>\n",
        "**Quick start**\n",
        "1. 把你要预测的氨基酸序列复制进输入框.\n",
        "2. 点击\"Runtime\" -> \"Run all\".\n",
        "3. 目前的模型预测包括5个模块，最终生成一个3维结构图.\n",
        "\n",
        "**生成文件**\n",
        "\n",
        "1. PDB格式的模型结构文件.\n",
        "2. 模型质量图\n",
        "3. 模型MSA覆盖率.\n",
        "4. 其他.\n",
        "\n",
        "**Acknowledgments**\n",
        "- We thank the AlphaFold team for developing an excellent model and open sourcing the software. \n",
        "\n",
        "- [Söding Lab](https://www.mpibpc.mpg.de/soeding) for providing the computational resources for the MMseqs2 server\n",
        "\n",
        "- Richard Evans for helping to benchmark the ColabFold's Alphafold-multimer support\n",
        "\n",
        "- [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",
        "- Do-Yoon Kim for creating the ColabFold logo.\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"
      ]
    }
  ]
}