model/PXDesignBench/ColabDesign/rf/examples/diffusion.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/rf/examples/diffusion.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tSgCPxIZ1T_A"
      },
      "source": [
        "#**RFdiffusion**\n",
        "RFdiffusion is a method for structure generation, with or without conditional information (a motif, target etc). It can perform a whole range of protein design challenges as we have outlined in the RFdiffusion [manuscript](https://www.biorxiv.org/content/10.1101/2022.12.09.519842v2).\n",
        "\n",
        "**<font color=\"red\">NOTE:</font>** This notebook is in development, we are still working on adding all the options from the manuscript above.\n",
        "\n",
        "For **instructions**, see end of Notebook.\n",
        "\n",
        "See [diffusion_foldcond](https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/rf/examples/diffusion_foldcond.ipynb) for fold conditioning functionality.\n",
        "\n",
        "See [original version](https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/rf/examples/diffusion_ori.ipynb) of this notebook (from 31Mar2023).\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "pZQnHLuDCsZm"
      },
      "outputs": [],
      "source": [
        "#@title setup **RFdiffusion** (~3min)\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 install jedi omegaconf hydra-core icecream pyrsistent pynvml decorator\")\n",
        "  os.system(\"pip install git+https://github.com/NVIDIA/dllogger#egg=dllogger\")\n",
        "  # 17Mar2024: adding --no-dependencies to avoid installing nvidia-cuda-* dependencies\n",
        "  # 25Aug2025: updating dgi install to work with latest pytorch\n",
        "  os.system(\"pip install --no-dependencies dgl -f https://data.dgl.ai/wheels/torch-2.4/cu124/repo.html\")\n",
        "  os.system(\"pip install --no-dependencies e3nn==0.5.5 opt_einsum_fx\")\n",
        "  os.system(\"cd RFdiffusion/env/SE3Transformer; pip 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": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "TuRUfQJZ4vkM"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "#@title run **RFdiffusion** to generate a backbone\n",
        "name = \"test\" #@param {type:\"string\"}\n",
        "contigs = \"100\" #@param {type:\"string\"}\n",
        "pdb = \"\" #@param {type:\"string\"}\n",
        "iterations = 50 #@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 = \"image\" #@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 = \"none\" #@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": "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 = False #@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 = False #@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": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "DUNKRBNSvk6_"
      },
      "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\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Instructions**\n",
        "---\n",
        "---\n",
        "\n",
        "Use `contigs` to define continious chains. Use a `:` to define multiple contigs and a `/` to define mutliple segments within a contig.\n",
        "For example:\n",
        "\n",
        "**unconditional**\n",
        "- `contigs='100'` - diffuse **monomer** of length 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",
        "*hints and tips*\n",
        "- `pdb=''` leave blank to get an upload prompt\n",
        "- `contigs='50-100'` use dash to specify a range of lengths to sample from"
      ],
      "metadata": {
        "id": "DKQXlWEjIOsf"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}