{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"
"
]
},
{
"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",
"**NOTE:** 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
}