{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github", "colab_type": "text" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "metadata": { "id": "OA2k3sAYuiXe" }, "source": [ "# AfDesign - hallucination\n", "For a given length, generate/hallucinate a protein sequence that AlphaFold thinks folds into a well structured protein (high plddt, low pae, many contacts).\n", "\n", "**WARNING**\n", "1. This notebook is in active development and was designed for demonstration purposes only.\n", "2. Using AfDesign as the only \"loss\" function for design might be a bad idea, you may find adversarial sequences (aka. sequences that trick AlphaFold)." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "cellView": "form", "id": "-AXy0s_4cKaK" }, "outputs": [], "source": [ "#@title setup\n", "%%time\n", "import os\n", "if not os.path.isdir(\"params\"):\n", " # get code\n", " os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n", " # for debugging\n", " os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n", " # download params\n", " os.system(\"mkdir params\")\n", " os.system(\"apt-get install aria2 -qq\")\n", " os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n", " os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n", "\n", "import warnings\n", "warnings.simplefilter(action='ignore', category=FutureWarning)\n", "\n", "import os\n", "from colabdesign import mk_afdesign_model, clear_mem\n", "from IPython.display import HTML\n", "from google.colab import files\n", "import numpy as np\n", "\n", "def get_pdb(pdb_code=\"\"):\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\"" ] }, { "cell_type": "code", "source": [ "#@title Hallucination Options\n", "length = 100#@param {type:\"integer\"}\n", "copies = 1#@param {type:\"integer\"}\n", "#@markdown ####Weights\n", "#@markdown - Minimizing `pae` or maximizing `plddt` often results in a single helix.\n", "#@markdown To avoid this, we start with a random sequence and instead try to optimize \n", "#@markdown defined `num`ber of `con`tacts per position. \n", "pae = 0.1 #@param [\"0.01\", \"0.1\", \"0.5\", \"1.0\"] {type:\"raw\"}\n", "plddt = 0.1 #@param [\"0.01\", \"0.1\", \"0.5\", \"1.0\"] {type:\"raw\"}\n", "con = 1.0 #@param [\"0.01\", \"0.1\", \"0.5\", \"1.0\"] {type:\"raw\"}\n", "#@markdown ####Contact Definition\n", "#@markdown - The contact definition is based on Cb-Cb diststance `cutoff`. To avoid \n", "#@markdown biasing towards helical contact, only contacts with sequence seperation > \n", "#@markdown `seqsep` are considered.\n", "\n", "\n", "seqsep = 9 #@param [\"0\",\"5\",\"9\"] {type:\"raw\"}\n", "cutoff = \"14\" #@param [\"8\", \"14\", \"max\"]\n", "num = \"2\" #@param [\"1\", \"2\", \"4\", \"8\", \"max\"]\n", "binary = False #@param {type:\"boolean\"}\n", "if cutoff == \"max\": cutoff = 21.6875\n", "if num == \"max\": num = length\n", "\n", "opt = {\"con\":{\"seqsep\":int(seqsep),\"cutoff\":float(cutoff),\"num\":int(num),\n", " \"binary\":binary}}\n", "weights = {\"con\":float(con), \"pae\":float(pae),\"plddt\":float(plddt)}\n", "\n", "if \"model\" not in dir() or model._len != length or model._args[\"copies\"] != copies:\n", " clear_mem()\n", " model = mk_afdesign_model(protocol=\"hallucination\")\n", " model.prep_inputs(length=length, copies=copies)\n", "\n", "#@markdown ####Optimizer settings\n", "pre_iters = 100 #@param {type:\"integer\"}\n", "soft_iters = 50 #@param {type:\"integer\"}\n", "temp_iters = 50 #@param {type:\"integer\"}\n", "hard_iters = 10 #@param {type:\"integer\"}\n", "\n", "# pre-design with gumbel initialization and softmax activation\n", "model.restart(mode=\"gumbel\",opt=opt,weights=weights)\n", "model.design_soft(pre_iters)\n", "\n", "# three stage design \n", "model.set_seq(model.aux[\"seq\"][\"pseudo\"])\n", "model.design_3stage(soft_iters,temp_iters,hard_iters)" ], "metadata": { "cellView": "form", "id": "eCGc3J663NGz" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "A1GxeLZdTTya", "cellView": "form" }, "outputs": [], "source": [ "#@markdown ## display hallucinated protein {run: \"auto\"}\n", "color = \"pLDDT\" #@param [\"chain\", \"pLDDT\", \"rainbow\"]\n", "show_sidechains = False #@param {type:\"boolean\"}\n", "show_mainchains = False #@param {type:\"boolean\"}\n", "model.save_pdb(f\"{model.protocol}.pdb\")\n", "model.plot_pdb(show_sidechains=show_sidechains,\n", " show_mainchains=show_mainchains,\n", " color=color)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "L2E9Tn2Acchj" }, "outputs": [], "source": [ "HTML(model.animate())" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "YSKWYu0_GlUH" }, "outputs": [], "source": [ "model.get_seqs()" ] } ], "metadata": { "accelerator": "GPU", "colab": { "collapsed_sections": [ "q4qiU9I0QHSz" ], "name": "hallucination.ipynb", "provenance": [], "include_colab_link": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }