{
  "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/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 (v1.1.1)\n",
        "Backprop through AlphaFold for protein design.\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": "markdown",
      "metadata": {
        "id": "UUfKrOzT0gOS"
      },
      "source": [
        "# fixed backbone design (fixbb)\n",
        "For a given protein backbone, generate/design a new sequence that AlphaFold thinks folds into that conformation. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qLd1DsnKzxBJ"
      },
      "outputs": [],
      "source": [
        "clear_mem()\n",
        "af_model = mk_afdesign_model(protocol=\"fixbb\")\n",
        "af_model.prep_inputs(pdb_filename=get_pdb(\"1TEN\"), chain=\"A\")\n",
        "\n",
        "print(\"length\",  af_model._len)\n",
        "print(\"weights\", af_model.opt[\"weights\"])"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "af_model.restart()\n",
        "af_model.design_3stage()"
      ],
      "metadata": {
        "id": "u0AwskJ84NGx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "af_model.plot_traj()  "
      ],
      "metadata": {
        "id": "8FB1v7dn1LL6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YEApO8YzBoS0"
      },
      "outputs": [],
      "source": [
        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
        "af_model.plot_pdb()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cW1KQiHKJpfp"
      },
      "outputs": [],
      "source": [
        "HTML(af_model.animate())"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "af_model.get_seqs()"
      ],
      "metadata": {
        "id": "YDrChASGVUUx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 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)."
      ],
      "metadata": {
        "id": "qLwS2s_xcjRI"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sZnYfCbfEvol"
      },
      "outputs": [],
      "source": [
        "clear_mem()\n",
        "af_model = mk_afdesign_model(protocol=\"hallucination\")\n",
        "af_model.prep_inputs(length=100)\n",
        "\n",
        "print(\"length\",af_model._len)\n",
        "print(\"weights\",af_model.opt[\"weights\"])"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# pre-design with gumbel initialization and softmax activation\n",
        "af_model.restart(mode=\"gumbel\")\n",
        "af_model.design_soft(50)\n",
        "\n",
        "# three stage design  \n",
        "af_model.set_seq(af_model.aux[\"seq\"][\"pseudo\"])\n",
        "af_model.design_3stage(50,50,10)"
      ],
      "metadata": {
        "id": "f76xqCkw0vj9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A1GxeLZdTTya"
      },
      "outputs": [],
      "source": [
        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
        "af_model.plot_pdb()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L2E9Tn2Acchj"
      },
      "outputs": [],
      "source": [
        "HTML(af_model.animate())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YSKWYu0_GlUH"
      },
      "outputs": [],
      "source": [
        "af_model.get_seqs()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dXfm4B8ISLuL"
      },
      "source": [
        "# binder hallucination\n",
        "For a given protein target and protein binder length, generate/hallucinate a protein binder sequence AlphaFold thinks will bind to the target structure.\n",
        "To do this, we minimize PAE and maximize number of contacts at the interface and within the binder, and we maximize pLDDT of the binder.\n",
        "By default, AlphaFold-ptm with residue index offset hack is used. To enable AlphaFold-multimer set: mk_afdesign_model(use_multimer=True).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3XLJPiRKx5Mw"
      },
      "outputs": [],
      "source": [
        "clear_mem()\n",
        "af_model = mk_afdesign_model(protocol=\"binder\")\n",
        "af_model.prep_inputs(pdb_filename=get_pdb(\"4MZK\"), chain=\"A\", binder_len=19)\n",
        "\n",
        "print(\"target_length\",af_model._target_len)\n",
        "print(\"binder_length\",af_model._binder_len)\n",
        "print(\"weights\",af_model.opt[\"weights\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u6VxjuinyCZa"
      },
      "outputs": [],
      "source": [
        "af_model.restart()\n",
        "af_model.design_3stage(100,100,10)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
        "af_model.plot_pdb()"
      ],
      "metadata": {
        "id": "sTlS7_L8Zfwf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9cARoviGyIKb"
      },
      "outputs": [],
      "source": [
        "HTML(af_model.animate())"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "af_model.get_seqs()"
      ],
      "metadata": {
        "id": "RzE137NDZdZc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#**ADVANCED**"
      ],
      "metadata": {
        "id": "SK0IJLoen_EC"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## partial hallucination + custom Radius of Gyration (rg) loss\n",
        "mix supervised (fixbb) and unsupervised (hallucination) losses to constrain the halluciation process."
      ],
      "metadata": {
        "id": "zl6JGTUzXRnk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "from colabdesign.af.alphafold.common import residue_constants\n",
        "\n",
        "# first off, let's implement a custom Radius of Gyration loss function\n",
        "def rg_loss(inputs, outputs):\n",
        "  positions = outputs[\"structure_module\"][\"final_atom_positions\"]\n",
        "  ca = positions[:,residue_constants.atom_order[\"CA\"]]\n",
        "  center = ca.mean(0)\n",
        "  rg = jnp.sqrt(jnp.square(ca - center).sum(-1).mean() + 1e-8)\n",
        "  rg_th = 2.38 * ca.shape[0] ** 0.365\n",
        "  rg = jax.nn.elu(rg - rg_th)\n",
        "  return {\"rg\":rg}"
      ],
      "metadata": {
        "id": "spec3m8BlGer"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "clear_mem()\n",
        "af_model = mk_afdesign_model(protocol=\"partial\",\n",
        "                             loss_callback=rg_loss, # add rg_loss\n",
        "                             use_templates=False)   # set True to constrain positions using template input\n",
        "\n",
        "af_model.opt[\"weights\"][\"rg\"] = 0.1  # optional: specify weight for rg_loss\n",
        "\n",
        "af_model.prep_inputs(pdb_filename=get_pdb(\"6MRR\"),\n",
        "                     chain=\"A\",\n",
        "                     pos=\"3-30,33-68\",  # define positions to contrain\n",
        "                     length=100)          # total length if different from input pdb\n",
        "\n",
        "af_model.rewire(loops=[36]) # set loop length between segments                     "
      ],
      "metadata": {
        "id": "h_BvzwbAKo6V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# initialize with wildtype seq, fill in the rest with soft_gumbel distribution\n",
        "af_model.restart(mode=[\"soft\",\"gumbel\",\"wildtype\"])\n",
        "af_model.design_3stage(100, 100, 10)"
      ],
      "metadata": {
        "id": "5Unr9u2GYKRD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
        "af_model.plot_pdb()"
      ],
      "metadata": {
        "id": "BFweaqNWYuF0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "HTML(af_model.animate())"
      ],
      "metadata": {
        "id": "GSu2lB9HYw-t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "af_model.get_seqs()"
      ],
      "metadata": {
        "id": "2EG2t2_KY4Td"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "rTGKbhsI0t8k"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "q4qiU9I0QHSz"
      ],
      "name": "design.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}