File size: 7,567 Bytes
d766458 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | {
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/hallucination.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"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
} |