esm/source/examples/protein-programming-language/tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# High-level protein programming language tutorial\n",
    "\n",
    "In this notebook, we will walk through creating a simple program encoding a symmetric protein and running an optimization run.\n",
    "At the end of the notebook, we will also link to examples of more complex design programs.\n",
    "\n",
    "We design proteins guided by [ESMFold](https://github.com/facebookresearch/esm#esmfold), so let's get set up by loading the ESMFold model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from language import EsmFoldv1\n",
    "\n",
    "folding_callback = EsmFoldv1()\n",
    "folding_callback.load(device=\"cuda:0\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's write a program.\n",
    "Our high-level programming language is organized by a *syntax tree*, where *leaf nodes* correspond to protein sequence segments. To begin, let's create a leaf node defined on a fixed-length sequence segment with 25 amino acid residues:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from language import FixedLengthSequenceSegment\n",
    "from language import ProgramNode\n",
    "\n",
    "protomer = FixedLengthSequenceSegment(25)\n",
    "def make_protomer_node():\n",
    "    return ProgramNode(sequence_segment=protomer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this tutorial, we will make a symmetric protein by repeating the above protomer sequence $N$ times, where $N$ is the desired rotational symmetry.\n",
    "\n",
    "Our programming language enables *internal nodes* in the syntax tree that let us hierarchically control multiple child nodes. To program a protein with 3-fold symmetry, let's create an internal node with $N = 3$ children."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 3\n",
    "_node = ProgramNode(\n",
    "    children=[make_protomer_node() for _ in range(N)],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To complete our programming language, we require a set of *constraints*.\n",
    "An important set of constraints to apply to all proteins is to maximize ESMFold's confidence in the model structure, which we quantify by the pTM and mean pLDDT scores.\n",
    "We also want to constrain the protein to have rotational symmetry.\n",
    "\n",
    "Adding those three constraints (on pTM, pLDDT, and symmetry) completes the simple program for a symmetric protein.\n",
    "We can also weight different constraints differently; for this tutorial, we will double the weight on the symmetry constraint.\n",
    "The constraints are compiled into a single *energy function* (right now, via a linear combination of constraint functions), which we will use to guide the optimization algorithm\n",
    "\n",
    "The code below will add the constraints and weights, compile the program, and print out a summary of the energy function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root:MaximizePTM = 1.0 * 0.84\n",
      "root:MaximizePLDDT = 1.0 * 0.65\n",
      "root:SymmetryRing = 1.0 * 0.57\n"
     ]
    }
   ],
   "source": [
    "from language import MaximizePTM, MaximizePLDDT, SymmetryRing\n",
    "\n",
    "# Define the program.\n",
    "program = ProgramNode(\n",
    "    energy_function_terms=[MaximizePTM(), MaximizePLDDT(), SymmetryRing()],\n",
    "    energy_function_weights=[1.0, 1.0, 1.0],\n",
    "    children=[make_protomer_node() for _ in range(N)],\n",
    ")\n",
    "\n",
    "# Set up the program.\n",
    "sequence, residue_indices = program.get_sequence_and_set_residue_index_ranges()\n",
    "\n",
    "# Compute and print the energy function.\n",
    "energy_terms = program.get_energy_term_functions()\n",
    "folding_output = folding_callback.fold(sequence, residue_indices)\n",
    "for name, weight, energy_fn in energy_terms:\n",
    "    print(f\"{name} = {weight:.1f} * {energy_fn(folding_output):.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to use this program to guide an optimization loop to design a protein sequence and structure.\n",
    "We currently perform simulated annealing of an MCMC optimization algorithm to do design.\n",
    "The code below will set up the simulated annealing algorithm, including the relevant parameters, and run the optimization loop. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2883ee40821e4f5ca4f699b96fe837e2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">/private/home/brianhie/protein-bert-model/examples/protein-programming-language/language/optimize.py:56: \n",
       "RuntimeWarning: overflow encountered in exp\n",
       "  np.exp(energy_differential / temperature),\n",
       "</pre>\n"
      ],
      "text/plain": [
       "/private/home/brianhie/protein-bert-model/examples/protein-programming-language/language/optimize.py:56: \n",
       "RuntimeWarning: overflow encountered in exp\n",
       "  np.exp(energy_differential / temperature),\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from language import run_simulated_annealing\n",
    "\n",
    "optimized_program = run_simulated_annealing(\n",
    "    program=program,\n",
    "    initial_temperature=1.0,\n",
    "    annealing_rate=0.97,\n",
    "    total_num_steps=10_000,\n",
    "    folding_callback=folding_callback,\n",
    "    display_progress=True,\n",
    ")\n",
    "print(\"Final sequence = {}\".format(optimized_program.get_sequence_and_set_residue_index_ranges()[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This program provides a minimal example of symmetric protein design, but we can also design much more complex proteins, as described in [our paper](https://www.biorxiv.org/content/10.1101/2022.12.21.521526v1).\n",
    "\n",
    "We provide example programs corresponding to our paper figures, as described in the table below:\n",
    "\n",
    "| Design task                                 | Figure in paper | Program file                                                                |\n",
    "|:--------------------------------------------|:---------------|:----------------------------------------------------------------------------|\n",
    "| Free hallucination                          | Figure 2A       | [free_hallucination.py](programs/free_hallucination.py)                     |\n",
    "| Fixed backbone design                       | Figure 2D       | [fixed_backbone.py](programs/fixed_backbone.py)                             |\n",
    "| Secondary structure design                  | Figure 2G       | [secondary_structure.py](programs/secondary_structure.py)                   |\n",
    "| Functional site scaffolding                 | Figure 2H       | [functional_site_scaffolding.py](programs/functional_site_scaffolding.py)   |\n",
    "| Symmetric monomer design                    | Figure 3A       | [symmetric_monomer.py](programs/symmetric_monomer.py)                       |\n",
    "| Two-level symmetric<br>homo-oligomer design | Figure 4A       | [symmetric_two_level_multimer.py](programs/symmetric_two_level_multimer.py) |\n",
    "| Symmetric binding site<br>scaffolding       | Figure 5A       | [symmetric_binding.py](programs/symmetric_binding.py)                       |\n",
    "\n",
    "These can easily be imported and run.\n",
    "For example, the following code imports a symmetric binding site scaffolding program and computes the energy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from programs.symmetric_binding import symmetric_binding_il10\n",
    "\n",
    "program = symmetric_binding_il10()\n",
    "\n",
    "sequence, residue_indices = program.get_sequence_and_set_residue_index_ranges()\n",
    "\n",
    "folding_output = folding_callback.fold(sequence, residue_indices)\n",
    "\n",
    "energy_terms = program.get_energy_term_functions()\n",
    "for name, weight, energy_fn in energy_terms:\n",
    "    print(f\"{name} = {weight:.1f} * {energy_fn(folding_output):.2f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}