notebooks/30_design_dashboard.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Summary\n",
    "\n",
    "In this notebook we load a network trained to reconstruct protein sequences and use this network to design sequences that fit the geometry of CATH domain [1n5uA03](http://www.cathdb.info/version/latest/domain/1n5uA03).\n",
    "\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "import logging\n",
    "import subprocess\n",
    "from pathlib import Path\n",
    "\n",
    "import ipywidgets as widgets\n",
    "import nglview\n",
    "from IPython.display import HTML, display, set_matplotlib_formats\n",
    "from ipywidgets import Layout\n",
    "\n",
    "import proteinsolver\n",
    "import proteinsolver.dashboard\n",
    "import torch\n",
    "from kmbio import PDB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib agg\n",
    "\n",
    "set_matplotlib_formats(\"svg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEBUG = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if DEBUG:\n",
    "    %load_ext autoreload\n",
    "    %autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NOTEBOOK_NAME = \"30_design_dashboard\"\n",
    "\n",
    "NOTEBOOK_PATH = Path(NOTEBOOK_NAME).resolve()\n",
    "NOTEBOOK_PATH.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "NOTEBOOK_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "JUPYTER_DATA_PATH = Path(subprocess.check_output([\"jupyter\", \"--data-dir\"], universal_newlines=True).strip()).resolve(strict=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "JUPYTER_DATA_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "STATIC_DATA_PATH = JUPYTER_DATA_PATH.joinpath(\"voila\", \"templates\", \"mytemplate\", \"static\").resolve(strict=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "STATIC_DATA_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "UNIQUE_ID = \"191f05de\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BEST_STATE_FILES = {\n",
    "    #\n",
    "    \"191f05de\": \"protein_train/191f05de/e53-s1952148-d93703104.state\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LOGGER = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "global_state = proteinsolver.dashboard.global_state\n",
    "global_state.output_folder = STATIC_DATA_PATH"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Custom CSS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HTML(\"\"\"\\\n",
    "<style>\n",
    "h1 {\n",
    "    margin-top: -20px !important;\n",
    "}\n",
    "\n",
    ".myheading {\n",
    "    font-size: large;\n",
    "}\n",
    "\n",
    ".mytext > .widget-label {\n",
    "    font-family: monospace;\n",
    "    font-size: small;\n",
    "    width: 60px;\n",
    "}\n",
    "\n",
    ".mysequence > .widget-label {\n",
    "    line-height: 1rem;\n",
    "    min-height: calc( var(--jp-widgets-inline-height) + 12);\n",
    "}\n",
    "\n",
    "/*\n",
    ".mynglout {\n",
    "    z-index: -100\n",
    "}\n",
    "*/\n",
    "\n",
    ".biojs_msa_labelblock {\n",
    "    overflow: hidden !important;\n",
    "}\n",
    "\n",
    ".ui-resizable {\n",
    "    height: 250px !important;\n",
    "}\n",
    "\n",
    "img {\n",
    "    max-width: 100% !important;\n",
    "}\n",
    "</style>\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Net = importlib.import_module(f\"protein_train.{UNIQUE_ID}.model\").Net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "global_state.net_class = Net\n",
    "global_state.state_file = BEST_STATE_FILES[UNIQUE_ID]\n",
    "global_state.net_kwargs = {\"x_input_size\": 21, \"adj_input_size\": 2, \"hidden_size\": 128, \"output_size\": 20}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Design pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "### Load protein sequence and geometry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "structure = PDB.load(Path(proteinsolver.__path__[0]).resolve(strict=True).joinpath(\"data\", \"inputs\", \"1n5uA03.pdb\"))\n",
    "proteinsolver.dashboard.load_structure(structure)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "### Test ProteinSolverProcess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "data = global_state.data\n",
    "data.x = torch.tensor(\n",
    "    [proteinsolver.utils.AMINO_ACID_TO_IDX[aa] for aa in global_state.target_sequence],\n",
    "    dtype=torch.long,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "global_state.proteinsolver_process = proteinsolver.dashboard.ProteinSolverProcess(\n",
    "    global_state.net_class,\n",
    "    global_state.state_file,\n",
    "    global_state.data,\n",
    "    num_designs=2,\n",
    "    net_kwargs=global_state.net_kwargs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "if DEBUG:\n",
    "    global_state.proteinsolver_process.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "if DEBUG:\n",
    "    global_state.generated_sequences = [global_state.proteinsolver_process.output_queue.get() for _ in range(2)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "if DEBUG:\n",
    "    print(global_state.generated_sequences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "outputs": [],
   "source": [
    "if DEBUG:\n",
    "    global_state.proteinsolver_process.join()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "## Dashboard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "### Structure view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ngl_stage = nglview.NGLWidget(\n",
    "#     height=\"300px\",\n",
    "#     width=\"300px\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ngl_stage_out = widgets.Output(\n",
    "    layout=widgets.Layout(width=\"20%\", margin=\"0 0 0 0\"))\n",
    "ngl_stage_out.add_class(\"mynglout\")\n",
    "\n",
    "with ngl_stage_out:\n",
    "    display(ngl_stage)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "### Distance matrix view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "distance_matrix_out = widgets.Output()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "hide"
    ]
   },
   "source": [
    "### Info panel view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_selection_out = widgets.Output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sequence_generation_out = widgets.Output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_structure_button = proteinsolver.dashboard.create_load_structure_button(\n",
    "    ngl_stage, distance_matrix_out, target_selection_out, sequence_generation_out\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_distance_matrix_button = proteinsolver.dashboard.create_load_distance_matrix_button(\n",
    "    ngl_stage, distance_matrix_out, target_selection_out, sequence_generation_out\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_example_buttons = proteinsolver.dashboard.create_load_example_buttons(\n",
    "    ngl_stage, distance_matrix_out, target_selection_out, sequence_generation_out\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_structure_out = widgets.Output(\n",
    "    layout=widgets.Layout(width=\"60%\", margin=\"0px 0px 0px 0px\"))\n",
    "\n",
    "with load_structure_out:\n",
    "    display(HTML('<p class=\"myheading\">1. Load a reference protein structure</p>'))\n",
    "    display(HTML('<p>Upload the distance matrix or the structure of a reference protein. An example distance matrix, corresponding to the protein <code>1n5uA03</code>, can be found <a href=\"voila/files/1n5uA03-distance-matrix.txt\">here</a>. </p>'))\n",
    "    display(widgets.HBox([load_structure_button, load_distance_matrix_button]))\n",
    "    display(HTML('<p style=\"margin-top: 16px\">Alternatively, try one of the examples provided below:</p>'))\n",
    "    display(load_example_buttons)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate protein sequences matching a predetermined geometry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = widgets.VBox(\n",
    "    [\n",
    "        widgets.HBox(\n",
    "            [\n",
    "                load_structure_out,\n",
    "                ngl_stage_out,\n",
    "                widgets.HBox(\n",
    "                    [distance_matrix_out],\n",
    "                    layout=widgets.Layout(\n",
    "                        width=\"19%\",\n",
    "                        margin=\"0 0 0 1%\",\n",
    "                        flex_flow=\"col nowrap\",\n",
    "                        justify_content=\"space-around\",\n",
    "                        align_items=\"center\"\n",
    "                    ),\n",
    "                ),\n",
    "            ],\n",
    "            layout=widgets.Layout(\n",
    "                width=\"auto\",\n",
    "                margin=\"0 0 0 0\",\n",
    "#                 flex_flow=\"row nowrap\",\n",
    "#                 justify_content=\"space-between\",\n",
    "            ),\n",
    "        ),\n",
    "        target_selection_out,\n",
    "        sequence_generation_out,\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# proteinsolver.dashboard.on_run_ps_button_clicked(run_ps_button, num_designs_field)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(HTML(\"\"\"\\\n",
    "<hr>\n",
    "<p>Running into problems? Please contact us at <a href=\"help@proteinsolver.org\">help@proteinsolver.org</a>.\n",
    "<br>\n",
    "<em>This website works best using the latest versions of Firefox or Chrome web browsers.</em>\n",
    "</p>\n",
    "\"\"\"))"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": false,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}