{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "provenance": [], "include_colab_link": true }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" }, "accelerator": "GPU", "gpuClass": "standard" }, "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github", "colab_type": "text" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "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", "**WARNING** This notebook is in development, we are still working on adding all the options from the manuscript above.\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'` `copies=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\n", "\n" ], "metadata": { "id": "tSgCPxIZ1T_A" } }, { "cell_type": "code", "source": [ "#@title setup **RFdiffusion** (~2m30S)\n", "%%time\n", "import os, time\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 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 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\")\n", " os.system(\"pip install dgl==1.0.2+cu116 -f https://data.dgl.ai/wheels/cu116/repo.html\")\n", " os.system(\"cd RFdiffusion/env/SE3Transformer; pip -q install --no-cache-dir -r requirements.txt; pip -q install .\")\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\"]\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", "\n", "import sys, random, string, re\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", "from colabdesign.rf.utils import fix_contigs, fix_partial_contigs, fix_pdb\n", "from inference.utils import parse_pdb\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", " os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n", " return f\"{pdb_code}.pdb\"\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_diffusion(contigs, path, pdb=None, iterations=50,\n", " symmetry=\"cyclic\", copies=1, hotspot=None):\n", " # determine mode\n", " contigs = contigs.replace(\",\",\" \").replace(\":\",\" \").split()\n", " is_fixed, is_free = False, False\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.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_filename = get_pdb(pdb)\n", " parsed_pdb = parse_pdb(pdb_filename)\n", " opts = f\" inference.input_pdb={pdb_filename}\"\n", " if mode in [\"partial\"]:\n", " partial_T = int(80 * (iterations / 200))\n", " opts += f\" diffuser.partial_T={partial_T}\"\n", " contigs = fix_partial_contigs(contigs, parsed_pdb)\n", " else:\n", " opts += f\" diffuser.T={iterations}\"\n", " contigs = fix_contigs(contigs, parsed_pdb)\n", " else:\n", " opts = 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", " opts += f\" ppi.hotspot_res=[{hotspot}]\"\n", "\n", " # setup symmetry\n", " if copies > 1:\n", " sym = {\"cyclic\":\"c\",\"dihedral\":\"d\"}[symmetry] + str(copies)\n", " sym_opts = f\"--config-name symmetry inference.symmetry={sym} \\\n", " '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 = f\"{sym_opts} {opts}\"\n", " if symmetry == \"dihedral\": copies *= 2\n", " contigs = sum([contigs] * copies,[])\n", "\n", " opts = f\"{opts} 'contigmap.contigs=[{' '.join(contigs)}]'\"\n", "\n", " print(\"mode:\", mode)\n", " print(\"output:\", f\"outputs/{path}\")\n", " print(\"contigs:\", contigs)\n", "\n", " cmd = f\"./RFdiffusion/run_inference.py {opts} inference.output_prefix=outputs/{path} inference.num_designs=1\"\n", " print(cmd)\n", " !{cmd}\n", "\n", " # fix pdbs\n", " pdbs = [f\"outputs/traj/{path}_0_pX0_traj.pdb\",\n", " f\"outputs/traj/{path}_0_Xt-1_traj.pdb\",\n", " f\"outputs/{path}_0.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", " return contigs, copies" ], "metadata": { "cellView": "form", "id": "pZQnHLuDCsZm" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@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", "copies = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\", \"10\", \"11\", \"12\"] {type:\"raw\"}\n", "#@markdown ---\n", "#@markdown **advanced** settings\n", "iterations = 50 #@param [\"50\", \"100\", \"150\", \"200\"] {type:\"raw\"}\n", "symmetry = \"cyclic\" #@param [\"cyclic\", \"dihedral\"]\n", "hotspot = \"\" #@param {type:\"string\"}\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", " \"copies\":copies,\n", " \"iterations\":iterations,\n", " \"symmetry\":symmetry,\n", " \"hotspot\":hotspot,\n", " \"path\":path}\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)" ], "metadata": { "id": "TuRUfQJZ4vkM", "cellView": "form" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@title Display 3D structure {run: \"auto\"}\n", "import py3Dmol\n", "from colabdesign.shared.plot import pymol_color_list\n", "\n", "from string import ascii_uppercase,ascii_lowercase\n", "alphabet_list = list(ascii_uppercase+ascii_lowercase)\n", "\n", "show_mainchains = False \n", "animate = False #@param {type:\"boolean\"}\n", "color = \"chain\" #@param [\"rainbow\", \"chain\"]\n", "hbondCutoff = 4.0\n", "view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js')\n", "\n", "if animate:\n", " pdb = f\"/content/outputs/traj/{path}_0_pX0_traj.pdb\"\n", " pdb_str = open(pdb,'r').read()\n", " view.addModelsAsFrames(pdb_str,'pdb',{'hbondCutoff':hbondCutoff})\n", "else:\n", " pdb = f\"/content/outputs/{path}_0.pdb\"\n", " pdb_str = open(pdb,'r').read()\n", " view.addModel(pdb_str,'pdb',{'hbondCutoff':hbondCutoff})\n", "\n", "if color == \"rainbow\":\n", " view.setStyle({'cartoon': {'color':'spectrum'}})\n", "elif color == \"chain\":\n", " for n,chain,color in zip(range(len(contigs)),\n", " alphabet_list,\n", " pymol_color_list):\n", " view.setStyle({'chain':chain},{'cartoon': {'color':color}})\n", "\n", "view.zoomTo()\n", "if animate:\n", " view.animate({'loop': 'backAndForth'})\n", "view.show()" ], "metadata": { "id": "wqEi03_qi_g2", "cellView": "form" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "%%time\n", "#@title run **ProteinMPNN** to generate a sequence and **AlphaFold** to validate\n", "num_seqs = 8 #@param [\"8\", \"16\", \"32\", \"64\"] {type:\"raw\"}\n", "initial_guess = False #@param {type:\"boolean\"}\n", "num_recycles = 1 #@param [\"0\", \"1\", \"2\", \"3\", \"6\", \"12\"] {type:\"raw\"}\n", "use_multimer = False #@param {type:\"boolean\"}\n", "rm_aa = \"C\" #@param {type:\"string\"}\n", "#@markdown - for **binder** design, we recommend `initial_guess=True num_recycles=3`\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", "if initial_guess: opts.append(\"--initial_guess\")\n", "if use_multimer: opts.append(\"--use_multimer\")\n", "opts = ' '.join(opts)\n", "!python colabdesign/rf/designability_test.py {opts}" ], "metadata": { "id": "rES3p-q6j4tc", "cellView": "form" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@title Display best result\n", "import py3Dmol\n", "hbondCutoff = 4.0\n", "view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js')\n", "\n", "pdb_str = open(f\"outputs/{path}_0.pdb\",'r').read()\n", "view.addModel(pdb_str,'pdb',{'hbondCutoff':hbondCutoff})\n", "pdb_str = open(f\"outputs/{path}/best.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()" ], "metadata": { "cellView": "form", "id": "DUNKRBNSvk6_" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "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": { "cellView": "form", "id": "tVAE0BrnZoRR" }, "execution_count": null, "outputs": [] } ] }