{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github", "colab_type": "text" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "metadata": { "id": "OA2k3sAYuiXe" }, "source": [ "# AfDesign - peptide binder design\n", "For a given protein target and protein binder length, generate/hallucinate a protein binder sequence AlphaFold thinks will bind to the target structure. To do this, we maximize number of contacts at the interface and maximize pLDDT of the binder.\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", "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 colabdesign.shared.utils import copy_dict\n", "from colabdesign.af.alphafold.common import residue_constants\n", "\n", "from IPython.display import HTML\n", "from google.colab import files\n", "import numpy as np\n", "\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 **prep inputs**\n", "import re\n", "#@markdown ---\n", "#@markdown **target info**\n", "pdb = \"4N5T\" #@param {type:\"string\"}\n", "#@markdown - enter PDB code or UniProt code (to fetch AlphaFoldDB model) or leave blink to upload your own\n", "target_chain = \"A\" #@param {type:\"string\"}\n", "target_hotspot = \"\" #@param {type:\"string\"}\n", "if target_hotspot == \"\": target_hotspot = None\n", "#@markdown - restrict loss to predefined positions on target (eg. \"1-10,12,15\")\n", "target_flexible = False #@param {type:\"boolean\"}\n", "#@markdown - allow backbone of target structure to be flexible\n", "\n", "#@markdown ---\n", "#@markdown **binder info**\n", "binder_len = 14 #@param {type:\"integer\"}\n", "#@markdown - length of binder to hallucination\n", "binder_seq = \"\" #@param {type:\"string\"}\n", "binder_seq = re.sub(\"[^A-Z]\", \"\", binder_seq.upper())\n", "if len(binder_seq) > 0:\n", " binder_len = len(binder_seq)\n", "else:\n", " binder_seq = None\n", "#@markdown - if defined, will initialize design with this sequence\n", "\n", "binder_chain = \"\" #@param {type:\"string\"}\n", "if binder_chain == \"\": binder_chain = None\n", "#@markdown - if defined, supervised loss is used (binder_len is ignored)\n", "\n", "#@markdown ---\n", "#@markdown **model config**\n", "use_multimer = False #@param {type:\"boolean\"}\n", "#@markdown - use alphafold-multimer for design\n", "num_recycles = 0 #@param [\"0\", \"1\", \"3\", \"6\"] {type:\"raw\"}\n", "num_models = \"2\" #@param [\"1\", \"2\", \"3\", \"4\", \"5\", \"all\"]\n", "num_models = 5 if num_models == \"all\" else int(num_models)\n", "#@markdown - number of trained models to use during optimization\n", "\n", "\n", "x = {\"pdb_filename\":pdb,\n", " \"chain\":target_chain,\n", " \"binder_len\":binder_len,\n", " \"binder_chain\":binder_chain,\n", " \"hotspot\":target_hotspot,\n", " \"use_multimer\":use_multimer,\n", " \"rm_target_seq\":target_flexible}\n", " \n", "x[\"pdb_filename\"] = get_pdb(x[\"pdb_filename\"]) \n", "\n", "if \"x_prev\" not in dir() or x != x_prev:\n", " clear_mem()\n", " model = mk_afdesign_model(protocol=\"binder\",\n", " use_multimer=x[\"use_multimer\"],\n", " num_recycles=num_recycles,\n", " recycle_mode=\"sample\")\n", " model.prep_inputs(**x,\n", " ignore_missing=False)\n", " x_prev = copy_dict(x)\n", " print(\"target length:\", model._target_len)\n", " print(\"binder length:\", model._binder_len)\n", " binder_len = model._binder_len" ], "metadata": { "id": "HSgE99WALOE-", "cellView": "form" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "#@title **run AfDesign**\n", "from scipy.special import softmax\n", "\n", "optimizer = \"pssm_semigreedy\" #@param [\"pssm_semigreedy\", \"3stage\", \"semigreedy\", \"pssm\", \"logits\", \"soft\", \"hard\"]\n", "#@markdown - `pssm_semigreedy` - uses the designed PSSM to bias semigreedy opt. (Recommended)\n", "#@markdown - `3stage` - gradient based optimization (GD) (logits → soft → hard)\n", "#@markdown - `pssm` - GD optimize (logits → soft) to get a sequence profile (PSSM).\n", "#@markdown - `semigreedy` - tries X random mutations, accepts those that decrease loss\n", "#@markdown - `logits` - GD optimize logits inputs (continious)\n", "#@markdown - `soft` - GD optimize softmax(logits) inputs (probabilities)\n", "#@markdown - `hard` - GD optimize one_hot(logits) inputs (discrete)\n", "\n", "#@markdown WARNING: The output sequence from `pssm`,`logits`,`soft` is not one_hot. To get a valid sequence use the other optimizers, or redesign the output backbone with another protocol like ProteinMPNN.\n", "\n", "#@markdown ----\n", "#@markdown #### advanced GD settings\n", "GD_method = \"sgd\" #@param [\"adabelief\", \"adafactor\", \"adagrad\", \"adam\", \"adamw\", \"fromage\", \"lamb\", \"lars\", \"noisy_sgd\", \"dpsgd\", \"radam\", \"rmsprop\", \"sgd\", \"sm3\", \"yogi\"]\n", "learning_rate = 0.1 #@param {type:\"raw\"}\n", "norm_seq_grad = True #@param {type:\"boolean\"}\n", "dropout = True #@param {type:\"boolean\"}\n", "\n", "model.restart(seq=binder_seq)\n", "model.set_optimizer(optimizer=GD_method,\n", " learning_rate=learning_rate,\n", " norm_seq_grad=norm_seq_grad)\n", "models = model._model_names[:num_models]\n", "\n", "flags = {\"num_recycles\":num_recycles,\n", " \"models\":models,\n", " \"dropout\":dropout}\n", "\n", "if optimizer == \"3stage\":\n", " model.design_3stage(120, 60, 10, **flags)\n", " pssm = softmax(model._tmp[\"seq_logits\"],-1)\n", "\n", "if optimizer == \"pssm_semigreedy\":\n", " model.design_pssm_semigreedy(120, 32, **flags)\n", " pssm = softmax(model._tmp[\"seq_logits\"],1)\n", "\n", "if optimizer == \"semigreedy\":\n", " model.design_pssm_semigreedy(0, 32, **flags)\n", " pssm = None\n", "\n", "if optimizer == \"pssm\":\n", " model.design_logits(120, e_soft=1.0, num_models=1, ramp_recycles=True, **flags)\n", " model.design_soft(32, num_models=1, **flags)\n", " flags.update({\"dropout\":False,\"save_best\":True})\n", " model.design_soft(10, num_models=num_models, **flags)\n", " pssm = softmax(model.aux[\"seq\"][\"logits\"],-1)\n", "\n", "O = {\"logits\":model.design_logits,\n", " \"soft\":model.design_soft,\n", " \"hard\":model.design_hard}\n", "\n", "if optimizer in O:\n", " O[optimizer](120, num_models=1, ramp_recycles=True, **flags)\n", " flags.update({\"dropout\":False,\"save_best\":True})\n", " O[optimizer](10, num_models=num_models, **flags)\n", " pssm = softmax(model.aux[\"seq\"][\"logits\"],-1)\n", "\n", "model.save_pdb(f\"{model.protocol}.pdb\")" ], "metadata": { "cellView": "form", "id": "60qmxpzno0yV" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "A1GxeLZdTTya", "cellView": "form" }, "outputs": [], "source": [ "#@title 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", "color_HP = False #@param {type:\"boolean\"}\n", "animate = True #@param {type:\"boolean\"}\n", "model.plot_pdb(show_sidechains=show_sidechains,\n", " show_mainchains=show_mainchains,\n", " color=color, color_HP=color_HP, animate=animate)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "L2E9Tn2Acchj" }, "outputs": [], "source": [ "HTML(model.animate(dpi=100))" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "YSKWYu0_GlUH" }, "outputs": [], "source": [ "model.save_pdb(f\"{model.protocol}.pdb\")\n", "model.get_seqs()" ] }, { "cell_type": "code", "source": [ "#@markdown ### Amino acid probabilties\n", "import plotly.express as px\n", "alphabet = \"ACDEFGHIKLMNPQRSTVWY\"\n", "if \"pssm\" in dir() and pssm is not None:\n", " fig = px.imshow(pssm.mean(0).T,\n", " labels=dict(x=\"positions\", y=\"amino acids\", color=\"probability\"),\n", " y=residue_constants.restypes,\n", " zmin=0,\n", " zmax=1,\n", " template=\"simple_white\",\n", " )\n", " fig.update_xaxes(side=\"top\")\n", " fig.show()" ], "metadata": { "cellView": "form", "id": "3pvptBM55GHU" }, "execution_count": null, "outputs": [] }, { "cell_type": "code", "source": [ "# log\n", "model._tmp[\"best\"][\"aux\"][\"log\"]" ], "metadata": { "id": "1SGmdJKLNKvb" }, "execution_count": null, "outputs": [] } ], "metadata": { "accelerator": "GPU", "colab": { "collapsed_sections": [], "provenance": [], "include_colab_link": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }