RFdiffusion_peptide_binder_design_example.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tSgCPxIZ1T_A"
      },
      "source": [
        "#**Peptide binder design using RFdiffusion**\n",
        "\n",
        "아래 RFdiffusion 공식 Github Repo에 게시된 RFdiffusion example notebook을 기반으로 작성되었습니다 (2024.01.31 기준).\n",
        "https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/rf/examples/diffusion.ipynb\n",
        "\n",
        "초안은 서울대학교 약학대학 이주용 교수님 연구실에서 작성하였으며, 실습목적에 맞게 편집되었습니다.\n",
        "\n",
        "\n",
        "RFdiffusion은 단백질 구조 설계를 위해 개발된 인공지능으로, 아래와 같은 폭넓은 응용이 가능합니다.\n",
        " - Unconditional generation: 다양한 단백질 구조의 생성\n",
        " - Motif scaffolding: 기능 모티프를 가지는 단백질 구조의 생성\n",
        " - Symmetric protein generation: 대칭성을 가지는 단백질 (e.g. homo-oligomer) 구조의 생성\n",
        " - Protein/peptide binder design: 표적 단백질에 결합하는 단백질/펩타이드 디자인\n",
        " - Fold conditioning: 원하는 폴드를 가지는 단백질 구조 생성\n",
        " - Partial Diffusion: 주어진 단백질 구조의 다양화 (diversification)\n",
        "\n",
        "방법에 대한 자세한 내용은 링크된 [논문](https://www.nature.com/articles/s41586-023-06415-8)을 참고하시기 바랍니다.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "pZQnHLuDCsZm"
      },
      "outputs": [],
      "source": [
        "#@title setup **RFdiffusion** (2~10m)\n",
        "%%time\n",
        "import os, time, signal\n",
        "import sys, random, string, re\n",
        "if not os.path.isdir(\"params\"):\n",
        "  os.system(\"apt-get install aria2\")\n",
        "  os.system(\"mkdir params\")\n",
        "  # send param download into background\n",
        "  os.system(\"(\\\n",
        "  aria2c -q -x 16 https://files.ipd.uw.edu/krypton/schedules.zip; \\\n",
        "  aria2c -q -x 16 http://files.ipd.uw.edu/pub/RFdiffusion/6f5902ac237024bdd0c176cb93063dc4/Base_ckpt.pt; \\\n",
        "  aria2c -q -x 16 http://files.ipd.uw.edu/pub/RFdiffusion/e29311f6f1bf1af907f9ef9f44b8328b/Complex_base_ckpt.pt; \\\n",
        "  aria2c -q -x 16 http://files.ipd.uw.edu/pub/RFdiffusion/f572d396fae9206628714fb2ce00f72e/Complex_beta_ckpt.pt; \\\n",
        "  aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar; \\\n",
        "  tar -xf alphafold_params_2022-12-06.tar -C params; \\\n",
        "  touch params/done.txt) &\")\n",
        "\n",
        "if not os.path.isdir(\"RFdiffusion\"):\n",
        "  print(\"installing RFdiffusion...\")\n",
        "  os.system(\"git clone https://github.com/sokrypton/RFdiffusion.git\")\n",
        "  os.system(\"pip -q install jedi omegaconf hydra-core icecream pyrsistent\")\n",
        "  os.system(\"pip install dgl -f https://data.dgl.ai/wheels/cu121/repo.html\")\n",
        "  os.system(\"cd RFdiffusion/env/SE3Transformer; pip -q install --no-cache-dir -r requirements.txt; pip -q install .\")\n",
        "  os.system(\"wget -qnc https://files.ipd.uw.edu/krypton/ananas\")\n",
        "  os.system(\"chmod +x ananas\")\n",
        "\n",
        "if not os.path.isdir(\"colabdesign\"):\n",
        "  print(\"installing ColabDesign...\")\n",
        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git\")\n",
        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
        "\n",
        "if not os.path.isdir(\"RFdiffusion/models\"):\n",
        "  print(\"downloading RFdiffusion params...\")\n",
        "  os.system(\"mkdir RFdiffusion/models\")\n",
        "  models = [\"Base_ckpt.pt\",\"Complex_base_ckpt.pt\",\"Complex_beta_ckpt.pt\"]\n",
        "  for m in models:\n",
        "    while os.path.isfile(f\"{m}.aria2\"):\n",
        "      time.sleep(5)\n",
        "  os.system(f\"mv {' '.join(models)} RFdiffusion/models\")\n",
        "  os.system(\"unzip schedules.zip; rm schedules.zip\")\n",
        "\n",
        "if 'RFdiffusion' not in sys.path:\n",
        "  os.environ[\"DGLBACKEND\"] = \"pytorch\"\n",
        "  sys.path.append('RFdiffusion')\n",
        "\n",
        "from google.colab import files\n",
        "import json\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython.display import display, HTML\n",
        "import ipywidgets as widgets\n",
        "import py3Dmol\n",
        "\n",
        "from inference.utils import parse_pdb\n",
        "from colabdesign.rf.utils import get_ca\n",
        "from colabdesign.rf.utils import fix_contigs, fix_partial_contigs, fix_pdb, sym_it\n",
        "from colabdesign.shared.protein import pdb_to_string\n",
        "from colabdesign.shared.plot import plot_pseudo_3D\n",
        "\n",
        "def get_pdb(pdb_code=None):\n",
        "  if pdb_code is None or pdb_code == \"\":\n",
        "    upload_dict = files.upload()\n",
        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
        "    return \"tmp.pdb\"\n",
        "  elif os.path.isfile(pdb_code):\n",
        "    return pdb_code\n",
        "  elif len(pdb_code) == 4:\n",
        "    if not os.path.isfile(f\"{pdb_code}.pdb1\"):\n",
        "      os.system(f\"wget -qnc https://files.rcsb.org/download/{pdb_code}.pdb1.gz\")\n",
        "      os.system(f\"gunzip {pdb_code}.pdb1.gz\")\n",
        "    return f\"{pdb_code}.pdb1\"\n",
        "  else:\n",
        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\"\n",
        "\n",
        "def run_ananas(pdb_str, path, sym=None):\n",
        "  pdb_filename = f\"outputs/{path}/ananas_input.pdb\"\n",
        "  out_filename = f\"outputs/{path}/ananas.json\"\n",
        "  with open(pdb_filename,\"w\") as handle:\n",
        "    handle.write(pdb_str)\n",
        "\n",
        "  cmd = f\"./ananas {pdb_filename} -u -j {out_filename}\"\n",
        "  if sym is None: os.system(cmd)\n",
        "  else: os.system(f\"{cmd} {sym}\")\n",
        "\n",
        "  # parse results\n",
        "  try:\n",
        "    out = json.loads(open(out_filename,\"r\").read())\n",
        "    results,AU = out[0], out[-1][\"AU\"]\n",
        "    group = AU[\"group\"]\n",
        "    chains = AU[\"chain names\"]\n",
        "    rmsd = results[\"Average_RMSD\"]\n",
        "    print(f\"AnAnaS detected {group} symmetry at RMSD:{rmsd:.3}\")\n",
        "\n",
        "    C = np.array(results['transforms'][0]['CENTER'])\n",
        "    A = [np.array(t[\"AXIS\"]) for t in results['transforms']]\n",
        "\n",
        "    # apply symmetry and filter to the asymmetric unit\n",
        "    new_lines = []\n",
        "    for line in pdb_str.split(\"\\n\"):\n",
        "      if line.startswith(\"ATOM\"):\n",
        "        chain = line[21:22]\n",
        "        if chain in chains:\n",
        "          x = np.array([float(line[i:(i+8)]) for i in [30,38,46]])\n",
        "          if group[0] == \"c\":\n",
        "            x = sym_it(x,C,A[0])\n",
        "          if group[0] == \"d\":\n",
        "            x = sym_it(x,C,A[1],A[0])\n",
        "          coord_str = \"\".join([\"{:8.3f}\".format(a) for a in x])\n",
        "          new_lines.append(line[:30]+coord_str+line[54:])\n",
        "      else:\n",
        "        new_lines.append(line)\n",
        "    return results, \"\\n\".join(new_lines)\n",
        "\n",
        "  except:\n",
        "    return None, pdb_str\n",
        "\n",
        "def run(command, steps, num_designs=1, visual=\"none\"):\n",
        "\n",
        "  def run_command_and_get_pid(command):\n",
        "    pid_file = '/dev/shm/pid'\n",
        "    os.system(f'nohup {command} & echo $! > {pid_file}')\n",
        "    with open(pid_file, 'r') as f:\n",
        "      pid = int(f.read().strip())\n",
        "    os.remove(pid_file)\n",
        "    return pid\n",
        "  def is_process_running(pid):\n",
        "    try:\n",
        "      os.kill(pid, 0)\n",
        "    except OSError:\n",
        "      return False\n",
        "    else:\n",
        "      return True\n",
        "\n",
        "  run_output = widgets.Output()\n",
        "  progress = widgets.FloatProgress(min=0, max=1, description='running', bar_style='info')\n",
        "  display(widgets.VBox([progress, run_output]))\n",
        "\n",
        "  # clear previous run\n",
        "  for n in range(steps):\n",
        "    if os.path.isfile(f\"/dev/shm/{n}.pdb\"):\n",
        "      os.remove(f\"/dev/shm/{n}.pdb\")\n",
        "\n",
        "  pid = run_command_and_get_pid(command)\n",
        "  try:\n",
        "    fail = False\n",
        "    for _ in range(num_designs):\n",
        "\n",
        "      # for each step check if output generated\n",
        "      for n in range(steps):\n",
        "        wait = True\n",
        "        while wait and not fail:\n",
        "          time.sleep(0.1)\n",
        "          if os.path.isfile(f\"/dev/shm/{n}.pdb\"):\n",
        "            pdb_str = open(f\"/dev/shm/{n}.pdb\").read()\n",
        "            if pdb_str[-3:] == \"TER\":\n",
        "              wait = False\n",
        "            elif not is_process_running(pid):\n",
        "              fail = True\n",
        "          elif not is_process_running(pid):\n",
        "            fail = True\n",
        "\n",
        "        if fail:\n",
        "          progress.bar_style = 'danger'\n",
        "          progress.description = \"failed\"\n",
        "          break\n",
        "\n",
        "        else:\n",
        "          progress.value = (n+1) / steps\n",
        "          if visual != \"none\":\n",
        "            with run_output:\n",
        "              run_output.clear_output(wait=True)\n",
        "              if visual == \"image\":\n",
        "                xyz, bfact = get_ca(f\"/dev/shm/{n}.pdb\", get_bfact=True)\n",
        "                fig = plt.figure()\n",
        "                fig.set_dpi(100);fig.set_figwidth(6);fig.set_figheight(6)\n",
        "                ax1 = fig.add_subplot(111);ax1.set_xticks([]);ax1.set_yticks([])\n",
        "                plot_pseudo_3D(xyz, c=bfact, cmin=0.5, cmax=0.9, ax=ax1)\n",
        "                plt.show()\n",
        "              if visual == \"interactive\":\n",
        "                view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js')\n",
        "                view.addModel(pdb_str,'pdb')\n",
        "                view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':0.5,'max':0.9}}})\n",
        "                view.zoomTo()\n",
        "                view.show()\n",
        "        if os.path.exists(f\"/dev/shm/{n}.pdb\"):\n",
        "          os.remove(f\"/dev/shm/{n}.pdb\")\n",
        "      if fail:\n",
        "        progress.bar_style = 'danger'\n",
        "        progress.description = \"failed\"\n",
        "        break\n",
        "\n",
        "    while is_process_running(pid):\n",
        "      time.sleep(0.1)\n",
        "\n",
        "  except KeyboardInterrupt:\n",
        "    os.kill(pid, signal.SIGTERM)\n",
        "    progress.bar_style = 'danger'\n",
        "    progress.description = \"stopped\"\n",
        "\n",
        "def run_diffusion(contigs, path, pdb=None, iterations=50,\n",
        "                  symmetry=\"none\", order=1, hotspot=None,\n",
        "                  chains=None, add_potential=False, partial_T=\"auto\",\n",
        "                  num_designs=1, use_beta_model=False, visual=\"none\"):\n",
        "\n",
        "  full_path = f\"outputs/{path}\"\n",
        "  os.makedirs(full_path, exist_ok=True)\n",
        "  opts = [f\"inference.output_prefix={full_path}\",\n",
        "          f\"inference.num_designs={num_designs}\"]\n",
        "\n",
        "  if chains == \"\": chains = None\n",
        "\n",
        "  # determine symmetry type\n",
        "  if symmetry in [\"auto\",\"cyclic\",\"dihedral\"]:\n",
        "    if symmetry == \"auto\":\n",
        "      sym, copies = None, 1\n",
        "    else:\n",
        "      sym, copies = {\"cyclic\":(f\"c{order}\",order),\n",
        "                     \"dihedral\":(f\"d{order}\",order*2)}[symmetry]\n",
        "  else:\n",
        "    symmetry = None\n",
        "    sym, copies = None, 1\n",
        "\n",
        "  # determine mode\n",
        "  contigs = contigs.replace(\",\",\" \").replace(\":\",\" \").split()\n",
        "  is_fixed, is_free = False, False\n",
        "  fixed_chains = []\n",
        "  for contig in contigs:\n",
        "    for x in contig.split(\"/\"):\n",
        "      a = x.split(\"-\")[0]\n",
        "      if a[0].isalpha():\n",
        "        is_fixed = True\n",
        "        if a[0] not in fixed_chains:\n",
        "          fixed_chains.append(a[0])\n",
        "      if a.isnumeric():\n",
        "        is_free = True\n",
        "  if len(contigs) == 0 or not is_free:\n",
        "    mode = \"partial\"\n",
        "  elif is_fixed:\n",
        "    mode = \"fixed\"\n",
        "  else:\n",
        "    mode = \"free\"\n",
        "\n",
        "  # fix input contigs\n",
        "  if mode in [\"partial\",\"fixed\"]:\n",
        "    pdb_str = pdb_to_string(get_pdb(pdb), chains=chains)\n",
        "    if symmetry == \"auto\":\n",
        "      a, pdb_str = run_ananas(pdb_str, path)\n",
        "      if a is None:\n",
        "        print(f'ERROR: no symmetry detected')\n",
        "        symmetry = None\n",
        "        sym, copies = None, 1\n",
        "      else:\n",
        "        if a[\"group\"][0] == \"c\":\n",
        "          symmetry = \"cyclic\"\n",
        "          sym, copies = a[\"group\"], int(a[\"group\"][1:])\n",
        "        elif a[\"group\"][0] == \"d\":\n",
        "          symmetry = \"dihedral\"\n",
        "          sym, copies = a[\"group\"], 2 * int(a[\"group\"][1:])\n",
        "        else:\n",
        "          print(f'ERROR: the detected symmetry ({a[\"group\"]}) not currently supported')\n",
        "          symmetry = None\n",
        "          sym, copies = None, 1\n",
        "\n",
        "    elif mode == \"fixed\":\n",
        "      pdb_str = pdb_to_string(pdb_str, chains=fixed_chains)\n",
        "\n",
        "    pdb_filename = f\"{full_path}/input.pdb\"\n",
        "    with open(pdb_filename, \"w\") as handle:\n",
        "      handle.write(pdb_str)\n",
        "\n",
        "    parsed_pdb = parse_pdb(pdb_filename)\n",
        "    opts.append(f\"inference.input_pdb={pdb_filename}\")\n",
        "    if mode in [\"partial\"]:\n",
        "      if partial_T == \"auto\":\n",
        "        iterations = int(80 * (iterations / 200))\n",
        "      else:\n",
        "        iterations = int(partial_T)\n",
        "      opts.append(f\"diffuser.partial_T={iterations}\")\n",
        "      contigs = fix_partial_contigs(contigs, parsed_pdb)\n",
        "    else:\n",
        "      opts.append(f\"diffuser.T={iterations}\")\n",
        "      contigs = fix_contigs(contigs, parsed_pdb)\n",
        "  else:\n",
        "    opts.append(f\"diffuser.T={iterations}\")\n",
        "    parsed_pdb = None\n",
        "    contigs = fix_contigs(contigs, parsed_pdb)\n",
        "\n",
        "  if hotspot is not None and hotspot != \"\":\n",
        "    hotspot = \",\".join(hotspot.replace(\",\",\" \").split())\n",
        "    opts.append(f\"ppi.hotspot_res='[{hotspot}]'\")\n",
        "\n",
        "  # setup symmetry\n",
        "  if sym is not None:\n",
        "    sym_opts = [\"--config-name symmetry\", f\"inference.symmetry={sym}\"]\n",
        "    if add_potential:\n",
        "      sym_opts += [\"'potentials.guiding_potentials=[\\\"type:olig_contacts,weight_intra:1,weight_inter:0.1\\\"]'\",\n",
        "                   \"potentials.olig_intra_all=True\",\"potentials.olig_inter_all=True\",\n",
        "                   \"potentials.guide_scale=2\",\"potentials.guide_decay=quadratic\"]\n",
        "    opts = sym_opts + opts\n",
        "    contigs = sum([contigs] * copies,[])\n",
        "\n",
        "  opts.append(f\"'contigmap.contigs=[{' '.join(contigs)}]'\")\n",
        "  opts += [\"inference.dump_pdb=True\",\"inference.dump_pdb_path='/dev/shm'\"]\n",
        "  if use_beta_model:\n",
        "    opts += [\"inference.ckpt_override_path=./RFdiffusion/models/Complex_beta_ckpt.pt\"]\n",
        "\n",
        "  print(\"mode:\", mode)\n",
        "  print(\"output:\", full_path)\n",
        "  print(\"contigs:\", contigs)\n",
        "\n",
        "  opts_str = \" \".join(opts)\n",
        "  cmd = f\"./RFdiffusion/run_inference.py {opts_str}\"\n",
        "  print(cmd)\n",
        "\n",
        "  # RUN\n",
        "  run(cmd, iterations, num_designs, visual=visual)\n",
        "\n",
        "  # fix pdbs\n",
        "  for n in range(num_designs):\n",
        "    pdbs = [f\"outputs/traj/{path}_{n}_pX0_traj.pdb\",\n",
        "            f\"outputs/traj/{path}_{n}_Xt-1_traj.pdb\",\n",
        "            f\"{full_path}_{n}.pdb\"]\n",
        "    for pdb in pdbs:\n",
        "      with open(pdb,\"r\") as handle: pdb_str = handle.read()\n",
        "      with open(pdb,\"w\") as handle: handle.write(fix_pdb(pdb_str, contigs))\n",
        "\n",
        "  return contigs, copies"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Parameters for RFdiffusion\n",
        "\n",
        "아래 섹션에서는 RFdiffusion을 돌리기 위해 필요한 input parameter들이 무엇이 있는지, 각각 어떤 의미를 가지는지에 대해 설명하고 있습니다. 예제를 따라 돌려보면서 인풋 파라미터가 바뀌면 결과가 어떻게 달라지는지 확인해보세요.\n",
        "\n",
        "----\n",
        "### **공통 인풋 파라미터** ###\n",
        "\n",
        "1. **`name`**: 프로젝트 이름으로 디자인 결과 파일이 해당 이름으로 시작됩니다.\n",
        "\n",
        "2. **`contig`**: contig option은 contiguous의 줄임말로 연속적인 단백질 체인을 정의하기 위해 사용됩니다. \":\" 구분자는 여러 contigs(즉 여러개의 체인)를 정의하기 위해 사용되며, \"/\" 구분자는 하나의 체인 안에 있는 여러 부분들(예: motif로 주어진 부분과 디자인할 부분)을 구분하여 정의하기 위해 사용됩니다. 본 예제에서는 GLP-1R 구조에서 R chain에 결합하는 binder를 디자인 하고자 하는 것이 목표입니다. 따라서 contig option을 **`R:30`**과 같이 주게됩니다. 이는 R chain에 결합하는 30개짜리 peptide를 디자인하라는 뜻입니다. 이 외의 사용법에 대해서는 아래의 예시를 참고하시기 바랍니다.\n",
        "\n",
        "  >**unconditional**\n",
        "  - `contigs='100'` - diffuse **monomer** of length 100\n",
        "  - `contigs='50-100'` - diffuse **monomer** having length between 50 and 100\n",
        "  - `contigs='50:100'` - diffuse **hetero-oligomer** of lengths 50 and 100\n",
        "  - `contigs='50'` `symmetry='cyclic'` `order=2` - make two copies of the defined contig(s) and add a symmetry constraint, for **homo-oligomeric** diffusion.\n",
        "\n",
        "  >**binder design**\n",
        "  - `contigs='A:50'` `pdb='4N5T'` - diffuse a **binder** of length 50 to chain A of defined PDB.\n",
        "  - `contigs='E6-155:70-100'` `pdb='5KQV'` `hotspot='E64,E88,E96'` - diffuse a **binder** of length 70 to 100 (sampled randomly) to chain E and defined hotspot(s).\n",
        "\n",
        "  >**motif scaffolding**\n",
        "  - `contigs='40/A163-181/40'` `pdb='5TPN'`\n",
        "  - `contigs='A3-30/36/A33-68'` `pdb='6MRR'` - diffuse a loop of length 36 between two segments of defined PDB ranges.\n",
        "\n",
        "  >**partial diffusion**\n",
        "  - `contigs=''` `pdb='6MRR'` - noise all coordinates\n",
        "  - `contigs='A1-10'` `pdb='6MRR'` - keep first 10 positions fixed, noise the rest\n",
        "  - `contigs='A'` `pdb='1SSC'` - fix chain A, noise the rest\n",
        "\n",
        "3. **`pdb`**: 인풋으로 활용한 PDB ID를 적거나 업로드할 수 있습니다. PDB 파일 업로드를 원하는 경우, 빈칸으로 두고 실행시키면 업로드를 위한 프롬프트 창이 열립니다.\n",
        "\n",
        "\n",
        "4. **`iterations`**: fully noisy conformation에서 시작하여 몇 번의 denoising을 통해 구조를 생성할 것인지를 설정합니다. 많은 denoising을 할 수록 더 실제 단백질 구조에 가까운 구조들이 생성되지만 계산량이 많아지게 됩니다. 따라서 자신이 관심있는 시스템의 크기와 난이도에 맞게 이 값을 조정하는 것이 좋습니다. 본 예제에서는 25로 설정해두었습니다.\n",
        "\n",
        "5. **`hotspot`**: 특정 residue에 반드시 결합하도록 제한 조건을 주는 option입니다. 디자인하는 binder가 결합할 표적단백질의 부위를 제한하는 역할을 합니다. 예를 들어, `A30,A33,A35`와 같이 지정해주면, 생성되는 binder가 해당 레지듀에 결합하도록 디자인을 수행하게 됩니다.\n",
        "\n",
        "6. **`num_designs`**: 몇개의 구조를 생성할지를 정합니다. 본 예제에서는 1개만 생성해보겠지만, 실제 디자인을 수행할때는 수천~수만개 이상의 디자인 구조 샘플링이 필요합니다.\n",
        "\n",
        "7. **`visual`**: 생성한 단백질 구조를 interactive model viewer로 확인할지, 그림 형태로 확인할지, 혹은 보지 않을지를 정할 수 있습니다.\n",
        "\n",
        "----\n",
        "### **Symmetry 관련 인풋 파라미터** ###\n",
        "RFdiffusion은 symmetric oligomer 구조를 생성할 수 있습니다.\n",
        "\n",
        "**`symmetry`**의 option은 아래와 같습니다.\n",
        "1. `none`: symmetry가 없음.\n",
        "2. `auto`: 인풋 모티프 혹은 인풋 PDB가 가진 symmetry를 그대로 유지. 인풋을 기반으로 현재 가진 symmetry를 유추하여 사용하게 됩니다.\n",
        "3. `Cyclic`: [rotational symmetry](https://en.wikipedia.org/wiki/Cyclic_symmetry_in_three_dimensions)\n",
        "4. `Dihedral`: [rotational & reflection symmetry](https://en.wikipedia.org/wiki/Dihedral_group)\n",
        "\n",
        "<img src=\"https://s3-eu-west-1.amazonaws.com/ppreviews-plos-725668748/955260/preview.jpg\" width=400>\n",
        "\n",
        "**`order`**: 몇 개의 subunit으로 symmetry를 형성할 것인지를 지정해줍니다.\n",
        "\n",
        "**`chain`**: 지정된 chain들을 이용해서 symmetric 단백질 구조를 생성합니다.\n",
        "\n",
        "----\n",
        "### **Advanced settings** ###\n",
        "**`partial_T`**: **partial diffusion**을 하는 경우, 얼마나 많은 noise step을 사용할 것인지 지정합니다. 이 값이 클 수록 초기 구조에서 많이 벗어난 다양한 구조를 생성하게 됩니다.\n",
        "**`use_beta`**: 현재 RFdiffusion 모델은 대체적으로 helical 구조를 선호합니다.\n",
        "만일 beta-sheet 구조를 가지는 단백질을 디자인 하고 싶다면 이 옵션을 사용해보셔도 좋습니다."
      ],
      "metadata": {
        "id": "BupJuX3YY7Ej"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "TuRUfQJZ4vkM"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "#@title run **RFdiffusion** to generate a backbone\n",
        "name = \"6X18\" #@param {type:\"string\"}\n",
        "contigs = \"R:30\" #@param {type:\"string\"}\n",
        "pdb = \"6X18\" #@param {type:\"string\"}\n",
        "iterations = 25 #@param [\"25\", \"50\", \"100\", \"150\", \"200\"] {type:\"raw\"}\n",
        "hotspot = \"\" #@param {type:\"string\"}\n",
        "num_designs = 1 #@param [\"1\", \"2\", \"4\", \"8\", \"16\", \"32\"] {type:\"raw\"}\n",
        "visual = \"interactive\" #@param [\"none\", \"image\", \"interactive\"]\n",
        "#@markdown ---\n",
        "#@markdown **symmetry** settings\n",
        "#@markdown ---\n",
        "symmetry = \"none\" #@param [\"none\", \"auto\", \"cyclic\", \"dihedral\"]\n",
        "order = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\", \"10\", \"11\", \"12\"] {type:\"raw\"}\n",
        "chains = \"\" #@param {type:\"string\"}\n",
        "add_potential = True #@param {type:\"boolean\"}\n",
        "#@markdown - `symmetry='auto'` enables automatic symmetry dectection with [AnAnaS](https://team.inria.fr/nano-d/software/ananas/).\n",
        "#@markdown - `chains=\"A,B\"` filter PDB input to these chains (may help auto-symm detector)\n",
        "#@markdown - `add_potential` to discourage clashes between chains\n",
        "#@markdown ---\n",
        "#@markdown **advanced** settings\n",
        "#@markdown ---\n",
        "partial_T = \"auto\" # @param [\"auto\", \"10\", \"20\", \"40\", \"60\", \"80\"]\n",
        "#@markdown - specify number of noising steps (only used for the partial diffusion protocol)\n",
        "use_beta_model = False #@param {type:\"boolean\"}\n",
        "#@markdown - if you are seeing lots of helices, switch to the \"beta\" params for a better SSE balance.\n",
        "\n",
        "# determine where to save\n",
        "path = name\n",
        "while os.path.exists(f\"outputs/{path}_0.pdb\"):\n",
        "  path = name + \"_\" + ''.join(random.choices(string.ascii_lowercase + string.digits, k=5))\n",
        "\n",
        "flags = {\"contigs\":contigs,\n",
        "         \"pdb\":pdb,\n",
        "         \"order\":order,\n",
        "         \"iterations\":iterations,\n",
        "         \"symmetry\":symmetry,\n",
        "         \"hotspot\":hotspot,\n",
        "         \"path\":path,\n",
        "         \"chains\":chains,\n",
        "         \"add_potential\":add_potential,\n",
        "         \"num_designs\":num_designs,\n",
        "         \"use_beta_model\":use_beta_model,\n",
        "         \"visual\":visual,\n",
        "         \"partial_T\":partial_T}\n",
        "\n",
        "for k,v in flags.items():\n",
        "  if isinstance(v,str):\n",
        "    flags[k] = v.replace(\"'\",\"\").replace('\"','')\n",
        "\n",
        "contigs, copies = run_diffusion(**flags)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wqEi03_qi_g2"
      },
      "outputs": [],
      "source": [
        "#@title Display 3D structure {run: \"auto\"}\n",
        "animate = \"interactive\" #@param [\"none\", \"movie\", \"interactive\"]\n",
        "color = \"chain\" #@param [\"rainbow\", \"chain\", \"plddt\"]\n",
        "denoise = True\n",
        "dpi = 100 #@param [\"100\", \"200\", \"400\"] {type:\"raw\"}\n",
        "from colabdesign.shared.plot import pymol_color_list\n",
        "from colabdesign.rf.utils import get_ca, get_Ls, make_animation\n",
        "from string import ascii_uppercase,ascii_lowercase\n",
        "alphabet_list = list(ascii_uppercase+ascii_lowercase)\n",
        "\n",
        "def plot_pdb(num=0):\n",
        "  if denoise:\n",
        "    pdb_traj = f\"outputs/traj/{path}_{num}_pX0_traj.pdb\"\n",
        "  else:\n",
        "    pdb_traj = f\"outputs/traj/{path}_{num}_Xt-1_traj.pdb\"\n",
        "  if animate in [\"none\",\"interactive\"]:\n",
        "    hbondCutoff = 4.0\n",
        "    view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js')\n",
        "    if animate == \"interactive\":\n",
        "      pdb_str = open(pdb_traj,'r').read()\n",
        "      view.addModelsAsFrames(pdb_str,'pdb',{'hbondCutoff':hbondCutoff})\n",
        "    else:\n",
        "      pdb = f\"outputs/{path}_{num}.pdb\"\n",
        "      pdb_str = open(pdb,'r').read()\n",
        "      view.addModel(pdb_str,'pdb',{'hbondCutoff':hbondCutoff})\n",
        "    if color == \"rainbow\":\n",
        "      view.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "    elif color == \"chain\":\n",
        "      for n,chain,c in zip(range(len(contigs)),\n",
        "                              alphabet_list,\n",
        "                              pymol_color_list):\n",
        "          view.setStyle({'chain':chain},{'cartoon': {'color':c}})\n",
        "    else:\n",
        "      view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':0.5,'max':0.9}}})\n",
        "    view.zoomTo()\n",
        "    if animate == \"interactive\":\n",
        "      view.animate({'loop': 'backAndForth'})\n",
        "    view.show()\n",
        "  else:\n",
        "    Ls = get_Ls(contigs)\n",
        "    xyz, bfact = get_ca(pdb_traj, get_bfact=True)\n",
        "    xyz = xyz.reshape((-1,sum(Ls),3))[::-1]\n",
        "    bfact = bfact.reshape((-1,sum(Ls)))[::-1]\n",
        "    if color == \"chain\":\n",
        "      display(HTML(make_animation(xyz, Ls=Ls, dpi=dpi, ref=-1)))\n",
        "    elif color == \"rainbow\":\n",
        "      display(HTML(make_animation(xyz, dpi=dpi, ref=-1)))\n",
        "    else:\n",
        "      display(HTML(make_animation(xyz, plddt=bfact*100, dpi=dpi, ref=-1)))\n",
        "\n",
        "\n",
        "if num_designs > 1:\n",
        "  output = widgets.Output()\n",
        "  def on_change(change):\n",
        "    if change['name'] == 'value':\n",
        "      with output:\n",
        "        output.clear_output(wait=True)\n",
        "        plot_pdb(change['new'])\n",
        "  dropdown = widgets.Dropdown(\n",
        "      options=[(f'{k}',k) for k in range(num_designs)],\n",
        "      value=0, description='design:',\n",
        "  )\n",
        "  dropdown.observe(on_change)\n",
        "  display(widgets.VBox([dropdown, output]))\n",
        "  with output:\n",
        "    plot_pdb(dropdown.value)\n",
        "else:\n",
        "  plot_pdb()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "여기서 성공적으로 디자인된 binder의 backbone 구조를 확인할 수 있습니다.\n",
        "\n",
        "**주의**: 현재 디자인 된 구조는 **backbone**구조만 가지고 있습니다.\n",
        "따라서 전체 아미노산이 모두 glycine으로 구성되어 있다고 볼 수 있습니다.\n",
        "\n",
        "디자인된 구조를 기반으로 서열 디자인은 아래 단계에서 수행하겠습니다."
      ],
      "metadata": {
        "id": "w3IAWRNAiMb9"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "rES3p-q6j4tc"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "#@title run **ProteinMPNN** to generate a sequence and **AlphaFold** to validate\n",
        "#@markdown ProteinMPNN Settings\n",
        "num_seqs = 8 #@param [\"1\", \"2\", \"4\", \"8\", \"16\", \"32\", \"64\"] {type:\"raw\"}\n",
        "mpnn_sampling_temp = 0.1 #@param [\"0.0001\", \"0.1\", \"0.15\", \"0.2\", \"0.25\", \"0.3\", \"0.5\", \"1.0\"] {type:\"raw\"}\n",
        "rm_aa = \"C\" #@param {type:\"string\"}\n",
        "use_solubleMPNN = False #@param {type:\"boolean\"}\n",
        "#@markdown - `mpnn_sampling_temp` - control diversity of sampled sequences. (higher = more diverse).\n",
        "#@markdown - `rm_aa='C'` - do not use [C]ysteines.\n",
        "#@markdown - `use_solubleMPNN` - use weights trained only on soluble proteins. See [preprint](https://www.biorxiv.org/content/10.1101/2023.05.09.540044v2).\n",
        "#@markdown\n",
        "#@markdown AlphaFold Settings\n",
        "initial_guess = True #@param {type:\"boolean\"}\n",
        "#@markdown - soft initialization with desired coordinates, see [paper](https://www.nature.com/articles/s41467-023-38328-5).\n",
        "num_recycles = 1 #@param [\"0\", \"1\", \"2\", \"3\", \"6\", \"12\"] {type:\"raw\"}\n",
        "#@markdown - for **binder** design, we recommend `initial_guess=True num_recycles=3`\n",
        "use_multimer = True #@param {type:\"boolean\"}\n",
        "#@markdown - `use_multimer` - use AlphaFold Multimer v3 params for prediction.\n",
        "\n",
        "if not os.path.isfile(\"params/done.txt\"):\n",
        "  print(\"downloading AlphaFold params...\")\n",
        "  while not os.path.isfile(\"params/done.txt\"):\n",
        "    time.sleep(5)\n",
        "\n",
        "contigs_str = \":\".join(contigs)\n",
        "opts = [f\"--pdb=outputs/{path}_0.pdb\",\n",
        "        f\"--loc=outputs/{path}\",\n",
        "        f\"--contig={contigs_str}\",\n",
        "        f\"--copies={copies}\",\n",
        "        f\"--num_seqs={num_seqs}\",\n",
        "        f\"--num_recycles={num_recycles}\",\n",
        "        f\"--rm_aa={rm_aa}\",\n",
        "        f\"--mpnn_sampling_temp={mpnn_sampling_temp}\",\n",
        "        f\"--num_designs={num_designs}\"]\n",
        "if initial_guess: opts.append(\"--initial_guess\")\n",
        "if use_multimer: opts.append(\"--use_multimer\")\n",
        "if use_solubleMPNN: opts.append(\"--use_soluble\")\n",
        "opts = ' '.join(opts)\n",
        "!python colabdesign/rf/designability_test.py {opts}"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Analyzing results\n",
        "------\n",
        "\n",
        "RFdiffusion으로 만들어진 구조를 얼마나 믿을 수 있을까요?\n",
        "\n",
        "RFdiffusion으로 생성된 구조가 실제로 결합할지 여부를 판단하기 위해서 많이 사용되는 방법은 **RFdiffusion으로 생성된 구조가 AlphaFold2를 이용해서 예측된 구조와 일치**하는지를 살펴보는 것입니다.\n",
        "\n",
        "두 구조가 유사하다면 두 가지 서로 다른 방법이 비슷한 결과를 준다는 의미이므로 그 결과를 좀 더 믿을만 할 것입니다.\n",
        "\n",
        "\n",
        "위 코드를 수행하면 RFdiffusion으로 얻어진 backbone 구조를 바탕으로 ProteinMPNN을 수행하고 그 결과를 바탕으로 AF2를 실행한 결과를 보여줍니다.\n",
        "\n",
        "위 결과에서 `mpnn`은 ProteinMPNN score를 의미하며, 낮을 수록 좋습니다.\n",
        "\n",
        "`plddt`는 AF2의 결과가 얼마나 믿을만 한가?를 보여줍니다. 이 값이 높을수록 좋습니다.\n",
        "\n",
        "`i_ptm`: interaction predicted TMscore, 예측된 interface 구조의 정확도이며 높을수록 좋습니다.\n",
        "\n",
        "`i_pae`: interaction predicted aligned error, 예측된 domain 사이의 정렬 오차. 이 값이 낮을수록 좋으며, RFdiffusion의 저자들은 이 값이 [10 이하](https://github.com/RosettaCommons/RFdiffusion)일 때 실험하는 것이 좋다고 추천하고 있습니다.\n",
        "\n",
        "`rmsd`: RFdiffusion결과와 AF2 결과 사이의 root mean square error. 낮을수록 좋습니다.\n",
        "\n",
        "위의 수치들을 종합하여 어떤 서열을 실험할지를 정하는 것이 좋습니다.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "k8TdLdkOhBBy"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DUNKRBNSvk6_",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Display best result\n",
        "import py3Dmol\n",
        "def plot_pdb(num = \"best\"):\n",
        "  if num == \"best\":\n",
        "    with open(f\"outputs/{path}/best.pdb\",\"r\") as f:\n",
        "      # REMARK 001 design {m} N {n} RMSD {rmsd}\n",
        "      info = f.readline().strip('\\n').split()\n",
        "    num = info[3]\n",
        "  hbondCutoff = 4.0\n",
        "  view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js')\n",
        "  pdb_str = open(f\"outputs/{path}_{num}.pdb\",'r').read()\n",
        "  view.addModel(pdb_str,'pdb',{'hbondCutoff':hbondCutoff})\n",
        "  pdb_str = open(f\"outputs/{path}/best_design{num}.pdb\",'r').read()\n",
        "  view.addModel(pdb_str,'pdb',{'hbondCutoff':hbondCutoff})\n",
        "\n",
        "  view.setStyle({\"model\":0},{'cartoon':{}}) #: {'colorscheme': {'prop':'b','gradient': 'roygb','min':0,'max':100}}})\n",
        "  view.setStyle({\"model\":1},{'cartoon':{'colorscheme': {'prop':'b','gradient': 'roygb','min':0,'max':100}}})\n",
        "  view.zoomTo()\n",
        "  view.show()\n",
        "\n",
        "if num_designs > 1:\n",
        "  def on_change(change):\n",
        "    if change['name'] == 'value':\n",
        "      with output:\n",
        "        output.clear_output(wait=True)\n",
        "        plot_pdb(change['new'])\n",
        "  dropdown = widgets.Dropdown(\n",
        "    options=[\"best\"] + [str(k) for k in range(num_designs)],\n",
        "    value=\"best\",\n",
        "    description='design:',\n",
        "  )\n",
        "  dropdown.observe(on_change)\n",
        "  output = widgets.Output()\n",
        "  display(widgets.VBox([dropdown, output]))\n",
        "  with output:\n",
        "    plot_pdb(dropdown.value)\n",
        "else:\n",
        "  plot_pdb()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tVAE0BrnZoRR"
      },
      "outputs": [],
      "source": [
        "#@title Package and download results\n",
        "#@markdown If you are having issues downloading the result archive,\n",
        "#@markdown try disabling your adblocker and run this cell again.\n",
        "#@markdown  If that fails click on the little folder icon to the\n",
        "#@markdown  left, navigate to file: `name.result.zip`,\n",
        "#@markdown  right-click and select \\\"Download\\\"\n",
        "#@markdown (see [screenshot](https://pbs.twimg.com/media/E6wRW2lWUAEOuoe?format=jpg&name=small)).\n",
        "!zip -r {path}.result.zip outputs/{path}* outputs/traj/{path}*\n",
        "files.download(f\"{path}.result.zip\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}