peptide_binder_design.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/deepsatflow/adaptyv-bio-pdc/blob/main/peptide_binder_design.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 - 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 = \"6ARU\" #@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 = \"305-504\" #@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 = 48 #@param {type:\"integer\"}\n",
        "#@markdown - length of binder to hallucination\n",
        "binder_seq = \"SESELEAKRQELKELMEKTEEKARELLAKGDVAGARAAYGAYLLKAAE\" #@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 = True #@param {type:\"boolean\"}\n",
        "#@markdown - use alphafold-multimer for design\n",
        "num_recycles = 3 #@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": {
      "provenance": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}