cookbook/tutorials/3_gfp_design.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zWXOAcBB8h3z"
   },
   "source": [
    "# Design a GFP Candidate with ESM3\n",
    "\n",
    "This notebook walks through the computational methods used to design esmGFP in [Hayes et al., 2024](https://doi.org/10.1101/2024.07.01.600583). esmGFP has similar brightness and spectral properties as GFPs found in nature despite being a far distance (58% identity) from known fluorescent proteins, but we also found many other bright new GFPs with similar or higher sequence identity. One can likely design a lot more new GFPs with the approach sketched in this notebook!\n",
    "\n",
    "This notebook implements the core prompt used to begin the chain of thought used to create esmGFP. The overall process we used differs in two keys ways:\n",
    "\n",
    "1. We continued the generation process beyond what is shown in this notebook\n",
    "   to do a joint optimization of the generated sequence and structure.\n",
    "2. We used significantly more compute than is easy to do with a notebook to\n",
    "   generate many designs and filter them with a set of computational filters and\n",
    "   ranking mechanisms.\n",
    "\n",
    "And we validated a small number of the generated designs in a wet lab, which of course you can also do... but this notebook isn't very helpful with that!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kZRkkEKv-8YW"
   },
   "source": [
    "## Set up the notebook and model (via the Forge API).\n",
    "\n",
    "We begin by installing the [esm package](https://github.com/evolutionaryscale/esm) and py3Dmol, which will allow us to visualize our generations, and then importing necessary packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VgTZdaIMQ44H"
   },
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "\n",
    "!pip install git+https://github.com/evolutionaryscale/esm.git\n",
    "!pip install py3Dmol\n",
    "\n",
    "clear_output()  # Suppress pip install log lines after installation is complete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "poK5NTaXRGcX"
   },
   "outputs": [],
   "source": [
    "from getpass import getpass\n",
    "\n",
    "import biotite.sequence as seq\n",
    "import biotite.sequence.align as align\n",
    "import biotite.sequence.graphics as graphics\n",
    "import matplotlib.pyplot as pl\n",
    "import py3Dmol\n",
    "import torch\n",
    "from esm.sdk import client\n",
    "from esm.sdk.api import ESMProtein, GenerationConfig\n",
    "from esm.utils.structure.protein_chain import ProteinChain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vmVYm2uQ7m-5"
   },
   "source": [
    "\n",
    "ESM3 is a frontier generative model for biology. It is  scalable due to its ability to tokenize sequence, structure, and function and use a (nearly standard) transformer architecture while still being able to reason across all modalities simulateously.\n",
    "\n",
    "The largest ESM3 (98 billion parameters) was trained with 1.07e24 FLOPs on 2.78 billion proteins and 771 billion unique tokens. To create esmGFP we used the 7 billion parameter variant of ESM3. We'll use this model via the [EvolutionaryScale Forge](https://forge.evolutionaryscale.ai) API.\n",
    "\n",
    "Grab a token from [the Forge console](https://forge.evolutionaryscale.ai/console) and add it below. Note that your token is like a password for your account and you should take care to protect it. For this reason it is recommended to frequently create a new token and delete old, unused ones. It is also recommended to paste the token directly into an environment variable or use a utility like `getpass` as shown below so tokens are not accidentally shared or checked into code repositories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zNrU9Q2SYonX"
   },
   "outputs": [],
   "source": [
    "token = getpass(\"Token from Forge console: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9jIc4OZyh2oE"
   },
   "source": [
    "We then create a model stub that behaves somewhat like a PyTorch model, but under the hood it sends the inputs to the Forge server, runs the inputs through the neural network weights on that remote server, and then returns the output tensors here in this notebook. This stub can also be used in the EvolutionaryScale SDK to simplify a lot of the operations around generation, folding, and generally using the sampling. This is important because iterative sampling is key to getting the best performance from ESM3, and the SDK manages a lot of the complexity around implementing these standard routines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Tna_mjGOjdXA"
   },
   "outputs": [],
   "source": [
    "model = client(\n",
    "    model=\"esm3-medium-2024-03\", url=\"https://forge.evolutionaryscale.ai\", token=token\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZaLBVJlzTxdT"
   },
   "source": [
    "## Construct the GFP Prompt\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KTn1Z4hCORVR"
   },
   "source": [
    "ESM3 is a generative model. To access the generative capabilities we need to get comfortable with constructing prompts. ESM3 jointly reasons across sequence, structure, and function of proteins, so we can construct new types of prompts that can guide the model to generate proteins with higher levels of control than many other biological language models.\n",
    "\n",
    "Sequence, structure, and function modalities are represented as tracks of discrete tokens that are present at both the input and output of the model and fused into a single latent space within the model. ESM3 is trained\n",
    "with a generative masked language modeling objective with variable mask rates, so we can prompt with a fully or partially masked context and different points of conditioning across the various tracks. This gives us an opportunity to be highly creative with how we specify our prompts!\n",
    "\n",
    "Prompt engineering is a bit of an art and a bit of a science, so one typically needs to experiment to get a prompt that produces a desired result. Also because we use sampling to generate from the model the results of different generations from the same prompt will vary. Some prompts tend to have higher success rates requiring only a few generations to get a candidate protein design. Other more difficult prompts may require thousands of generations! The models are more controllable with alignment.\n",
    "\n",
    "The model we will be using is the raw pretrained (unaligned) model, but we've worked a lot on this prompt so one can typically get an interesting design with only a few generations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qtwnyA1BngWy"
   },
   "source": [
    "We'll construct our prompt from fragments of the [1qy3](https://www.rcsb.org/structure/1qy3) sequence and structure from the PDB. The following code fetches data from the PDB and then uses ESM3's tokenizers to convert the sequence and structure to tokens that can be passed into the model. Once can see that both the amino acid type and the coordinates are converted into one discrete token per sequence position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cDWcXKmlbC1z"
   },
   "outputs": [],
   "source": [
    "template_gfp = ESMProtein.from_protein_chain(\n",
    "    ProteinChain.from_rcsb(\"1qy3\", chain_id=\"A\")\n",
    ")\n",
    "\n",
    "template_gfp_tokens = model.encode(template_gfp)\n",
    "\n",
    "print(\"Sequence tokens:\")\n",
    "print(\n",
    "    \"    \", \", \".join([str(token) for token in template_gfp_tokens.sequence.tolist()])\n",
    ")\n",
    "\n",
    "print(\"Structure tokens:\")\n",
    "print(\n",
    "    \"    \", \", \".join([str(token) for token in template_gfp_tokens.structure.tolist()])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xUxwNTuWqxks"
   },
   "source": [
    "We'll now build a prompt. Specifically we'll specify 4 amino acid identities at positions near where we want the chromophore to form, and 2 amino acid identities on the beta barrel that are known to support chromophore formation.\n",
    "\n",
    "Furthermore we'll specify the structure should be similar to the 1qy3 structure at all these positions by adding tokens from the encoded 1qy3 structure to the structure track of our prompt. We'll also specify a few more positions (along the alpha helix kink)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YBfYwRFGDKjU"
   },
   "outputs": [],
   "source": [
    "prompt_sequence = [\"_\"] * len(template_gfp.sequence)\n",
    "prompt_sequence[59] = \"T\"\n",
    "prompt_sequence[62] = \"T\"\n",
    "prompt_sequence[63] = \"Y\"\n",
    "prompt_sequence[64] = \"G\"\n",
    "prompt_sequence[93] = \"R\"\n",
    "prompt_sequence[219] = \"E\"\n",
    "prompt_sequence = \"\".join(prompt_sequence)\n",
    "\n",
    "print(template_gfp.sequence)\n",
    "print(prompt_sequence)\n",
    "\n",
    "prompt = model.encode(ESMProtein(sequence=prompt_sequence))\n",
    "\n",
    "# We construct an empty structure track like |<bos> <mask> ... <mask> <eos>|...\n",
    "prompt.structure = torch.full_like(prompt.sequence, 4096)\n",
    "prompt.structure[0] = 4098\n",
    "prompt.structure[-1] = 4097\n",
    "# ... and then we fill in structure tokens at key residues near the alpha helix\n",
    "# kink and at the stabilizing R and E positions on the beta barrel.\n",
    "prompt.structure[55:70] = template_gfp_tokens.structure[56:71]\n",
    "prompt.structure[93] = template_gfp_tokens.structure[93]\n",
    "prompt.structure[219] = template_gfp_tokens.structure[219]\n",
    "\n",
    "print(\"\".join([\"✔\" if st < 4096 else \"_\" for st in prompt.structure]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "107URq1bpA4_"
   },
   "source": [
    "The output shows the original 1qy3 sequence and the our prompt sequence track amino acid identities and the positions that have a token on the structure track. ESM3 will then be tasked with filling in the structure and sequence at the remaining masked (underscore) positions.\n",
    "\n",
    "One small note, we introduced the mutation A93R in our prompt. This isn't a mistake. Using Alanine at this position causes the chromophore to mature extremely slowly (which is how we are able to measure the precyclized structure of GFP!). However we don't want to wait around for our GFPs to glow so we go with Arginine at this position."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DE0BvswIATN8"
   },
   "source": [
    "## Generate a Structure\n",
    "\n",
    "We then prompt the model and decode the structure tokens track. This is similar to creating a backbone scaffold for an active site prompt, but there are some subtle differences. For example, since we've already specified some of the structure tokens (e.g., around the active site and key corresponding residues) the model literally generates around this structure.\n",
    "\n",
    "Tokens are iteratively sampled from ESM3. They\n",
    "can be sampled one at a time, or in parallel, in any order, until all positions are fully unmasked. The generate() function in the EvolutionaryScale SDK implements one recipe we think is effective for sampling from the model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yatAF6kYHZdm"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "num_tokens_to_decode = (prompt.structure == 4096).sum().item()\n",
    "\n",
    "structure_generation = model.generate(\n",
    "    prompt,\n",
    "    GenerationConfig(\n",
    "        # Generate a structure.\n",
    "        track=\"structure\",\n",
    "        # Sample one token per forward pass of the model.\n",
    "        num_steps=num_tokens_to_decode,\n",
    "        # Sampling temperature trades perplexity with diversity.\n",
    "        temperature=1.0,\n",
    "    ),\n",
    ")\n",
    "\n",
    "print(\"These are the structure tokens corresponding to our new design:\")\n",
    "print(\n",
    "    \"    \", \", \".join([str(token) for token in structure_generation.structure.tolist()])\n",
    ")\n",
    "\n",
    "# Decodes structure tokens to backbone coordinates.\n",
    "structure_generation_protein = model.decode(structure_generation)\n",
    "\n",
    "print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0HARel94tJfI"
   },
   "source": [
    "Now let's visualize our generated structure. This will probably look like the familiar GFP beta barrel around an alpha helix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "D30KQC6ffrrH"
   },
   "outputs": [],
   "source": [
    "view = py3Dmol.view(width=1000, height=500)\n",
    "view.addModel(\n",
    "    structure_generation_protein.to_protein_chain().infer_oxygen().to_pdb_string(),\n",
    "    \"pdb\",\n",
    ")\n",
    "view.setStyle({\"cartoon\": {\"color\": \"lightgreen\"}})\n",
    "view.zoomTo()\n",
    "view.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "K55y-5BeDRdr"
   },
   "source": [
    "At this point we only want to continue the generation if this design is a close match to a wildtype GFP at the active site, has some structural difference across the full protein (otherwise it would end up being very sequence-similar to wildtype GFP), and overall still looks like the classic GFP alpha helix in a beta barrel structure.\n",
    "\n",
    "Of course when generating many designs we cannot look at each one manually, so we adopt some automated rejection sampling criteria based on the overall structure RMSD and the constrained site RMSD for our generated structure being faithful to the prompt. If these checks pass then we'll try to design a sequence for this structure. If not, one should go back up a few cells and design another structure until it passes these computational screens. (Or not... this is your GFP design!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aalzUw39t2O1"
   },
   "outputs": [],
   "source": [
    "constrained_site_positions = [59, 62, 63, 64, 93, 219]\n",
    "\n",
    "template_chain = template_gfp.to_protein_chain()\n",
    "generation_chain = structure_generation_protein.to_protein_chain()\n",
    "\n",
    "constrained_site_rmsd = template_chain[constrained_site_positions].rmsd(\n",
    "    generation_chain[constrained_site_positions]\n",
    ")\n",
    "backbone_rmsd = template_chain.rmsd(generation_chain)\n",
    "\n",
    "c_pass = \"✅\" if constrained_site_rmsd < 1.5 else \"❌\"\n",
    "b_pass = \"✅\" if backbone_rmsd > 1.5 else \"❌\"\n",
    "\n",
    "print(f\"Constrained site RMSD: {constrained_site_rmsd:.2f} Ang {c_pass}\")\n",
    "print(f\"Backbone RMSD: {backbone_rmsd:.2f} Ang {b_pass}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6Bfbm8UiEqya"
   },
   "source": [
    "# Sequence Design\n",
    "\n",
    "Now we have a backbone with some structural variation but that also matches the GFP constrained site, and we want to design a sequence that folds to this structure. We can use the prior generation, which is exactly our original prompt plus the new structure tokens representing the backbone, to prompt ESM3 again.\n",
    "\n",
    "One we have designed a sequence we'll want to confirm that sequence is a match for our structure, so we'll remove all other conditioning from the prompt and fold the sequence. Conveniently with ESM3, folding a sequence is simply generating a set of structure tokens conditioned on the amino acid sequence. In this case we want the model's highest confidence generation (with no diversity) so we sample with a temperature of zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GOrWSEVTnOq0"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "num_tokens_to_decode = (prompt.sequence == 32).sum().item()\n",
    "\n",
    "sequence_generation = model.generate(\n",
    "    # Generate a sequence.\n",
    "    structure_generation,\n",
    "    GenerationConfig(track=\"sequence\", num_steps=num_tokens_to_decode, temperature=1.0),\n",
    ")\n",
    "\n",
    "# Refold\n",
    "sequence_generation.structure = None\n",
    "length_of_sequence = sequence_generation.sequence.numel() - 2\n",
    "sequence_generation = model.generate(\n",
    "    sequence_generation,\n",
    "    GenerationConfig(track=\"structure\", num_steps=length_of_sequence, temperature=0.0),\n",
    ")\n",
    "\n",
    "# Decode to AA string and coordinates.\n",
    "sequence_generation_protein = model.decode(sequence_generation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v_zK7TDCzEX3"
   },
   "source": [
    "We now have a candidate GFP sequence!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ao_n0-R5r2fT"
   },
   "outputs": [],
   "source": [
    "sequence_generation_protein.sequence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LBvQYpR_zQAK"
   },
   "source": [
    "We can align this sequence against the original template to see how similar it is to avGFP. One might also want to search against all known fluorescent proteins to assess the novelty of this potential GFP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JuNegp37JRyD"
   },
   "outputs": [],
   "source": [
    "seq1 = seq.ProteinSequence(template_gfp.sequence)\n",
    "seq2 = seq.ProteinSequence(sequence_generation_protein.sequence)\n",
    "\n",
    "alignments = align.align_optimal(\n",
    "    seq1, seq2, align.SubstitutionMatrix.std_protein_matrix(), gap_penalty=(-10, -1)\n",
    ")\n",
    "\n",
    "alignment = alignments[0]\n",
    "\n",
    "identity = align.get_sequence_identity(alignment)\n",
    "print(f\"Sequence identity: {100*identity:.2f}%\")\n",
    "\n",
    "print(\"\\nSequence alignment:\")\n",
    "fig = pl.figure(figsize=(8.0, 4.0))\n",
    "ax = fig.add_subplot(111)\n",
    "graphics.plot_alignment_similarity_based(\n",
    "    ax, alignment, symbols_per_line=45, spacing=2, show_numbers=True\n",
    ")\n",
    "fig.tight_layout()\n",
    "pl.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QQHMYfJzz13w"
   },
   "source": [
    "We now recheck our computational metrics for the constrained site. If we see the constrained site is not a match then we'd want to try designing the sequence again. If many attempts to design a sequence that matches the structure fail, then it's likely the structure is not easily designable and we may want to reject this structure generation as well!\n",
    "\n",
    "At this point the backbone RMSD doesn't matter very much to us, so long as the sequence is adequately distant to satisfy our scientific curiosity!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HUGQ7L4_z1BV"
   },
   "outputs": [],
   "source": [
    "template_chain = template_gfp.to_protein_chain()\n",
    "generation_chain = sequence_generation_protein.to_protein_chain()\n",
    "\n",
    "constrained_site_rmsd = template_chain[constrained_site_positions].rmsd(\n",
    "    generation_chain[constrained_site_positions]\n",
    ")\n",
    "backbone_rmsd = template_chain.rmsd(generation_chain)\n",
    "\n",
    "c_pass = \"✅\" if constrained_site_rmsd < 1.5 else \"❌\"\n",
    "b_pass = \"🤷‍♂️\"\n",
    "\n",
    "print(f\"Constrained site RMSD: {constrained_site_rmsd:.2f} Ang {c_pass}\")\n",
    "print(f\"Backbone RMSD: {backbone_rmsd:.2f} Ang {b_pass}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0cIeC4Lg1Bz9"
   },
   "source": [
    "An now we can visualize the final structure prediction of our candidate GFP design."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WTGRyt-es2sJ"
   },
   "outputs": [],
   "source": [
    "view = py3Dmol.view(width=600, height=600)\n",
    "view.addModel(sequence_generation_protein.to_pdb_string(), \"pdb\")\n",
    "view.setStyle({\"cartoon\": {\"color\": \"lightgreen\"}})\n",
    "view.zoomTo()\n",
    "view.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VrNuZHeHRWuP"
   },
   "source": [
    "Before considering this sequence for wet lab validation, we run a joint optimization of the sequence and structure. The outputs of that process are then passed through stringent computational filters and then many designs from many starting points are ranked by a number of computational scores to select the final designs sent for testing. We'll walk through that process in a different notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c3jSQrJa1Tfi"
   },
   "source": [
    "If you've made it this far it's worth noting that this isn't the only method to prompt ESM3 to design a GFP, it's just the one we used to report the successful generation of esmGFP in our paper. We hope you'll try different techniques to generate from ESM3. We're interested to hear what works for you!"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}