examples/lm-design/fixed_backbone.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4f84941e",
   "metadata": {},
   "source": [
    "# Fixed Backbone design from LM\n",
    "\n",
    "This notebook demonstrates the Fixed Backbone design task from the paper [Language models generalize beyond natural proteins\n",
    "](https://www.biorxiv.org/content/10.1101/2022.12.21.521521v1).\n",
    "\n",
    "Given an input structure as .pdb file, the LM is used iteratively in an MCMC optimization to find a sequence that folds to that structure\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d378b7f4-0792-446b-9e95-f7025bee5bec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First install additional dependencies\n",
    "!pip install -r additional_requirements.txt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfd13d6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "import os\n",
    "import time\n",
    "import hydra\n",
    "import py3Dmol\n",
    "from lm_design import Designer\n",
    "\n",
    "# Params\n",
    "pdb_fn = os.getcwd() + '/2N2U.pdb'\n",
    "seed = 0  # Use different seeds to get different sequence designs for the same structure\n",
    "TASK = \"fixedbb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "989996bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load hydra config from config.yaml\n",
    "with hydra.initialize_config_module(config_module=\"conf\"):\n",
    "    cfg = hydra.compose(\n",
    "        config_name=\"config\", \n",
    "        overrides=[\n",
    "            f\"task={TASK}\", \n",
    "            f\"seed={seed}\", \n",
    "            f\"pdb_fn={pdb_fn}\", \n",
    "            # 'tasks.fixedbb.num_iter=100'  # DEBUG - use a smaller number of iterations\n",
    "        ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63178538",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a designer from configuration\n",
    "des = Designer(cfg, pdb_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86d25575",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Run the designer\n",
    "start_time = time.time()\n",
    "des.run_from_cfg()\n",
    "print(\"finished after %s hours\", (time.time() - start_time) / 3600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6d9f742",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Output seq:\", des.output_seq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba6c8c66",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Fold output with ESMFold API\n",
    "output_seq = des.output_seq\n",
    "# Fold with api:\n",
    "#  curl -X POST --data \"GENGEIPLEIRATTGAEVDTRAVTAVEMTEGTLGIFRLPEEDYTALENFRYNRVAGENWKPASTVIYVGGTYARLCAYAPYNSVEFKNSSLKTEAGLTMQTYAAEKDMRFAVSGGDEVWKKTPTANFELKRAYARLVLSVVRDATYPNTCKITKAKIEAFTGNIITANTVDISTGTEGSGTQTPQYIHTVTTGLKDGFAIGLPQQTFSGGVVLTLTVDGMEYSVTIPANKLSTFVRGTKYIVSLAVKGGKLTLMSDKILIDKDWAEVQTGTGGSGDDYDTSFN\" https://api.esmatlas.com/foldSequence/v1/pdb/\n",
    "import requests\n",
    "import json\n",
    "url = 'https://api.esmatlas.com/foldSequence/v1/pdb/'\n",
    "r = requests.post(url, data=output_seq)\n",
    "output_struct = r.text\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5c06ab3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize output structure\n",
    "view = py3Dmol.view(width=800, height=800)\n",
    "view.addModel(output_struct, 'pdb')\n",
    "view.setStyle({'cartoon': {'color': 'spectrum'}})\n",
    "view.zoomTo()\n",
    "view.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7247225",
   "metadata": {},
   "outputs": [],
   "source": [
    "des.x_logits.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8e5c184",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize wild type structure\n",
    "wt_struct_file = pdb_fn\n",
    "view = py3Dmol.view(width=800, height=800)\n",
    "view.addModel(open(wt_struct_file).read(), 'pdb')\n",
    "view.setStyle({'cartoon': {'color': 'spectrum'}})\n",
    "view.zoomTo()\n",
    "view.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "222ec344",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "5502aca739f2549ad2771378ffc455b2bbb8b06f1a91617971f7097758a3cf84"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}