molearn_GUI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7c6c338e",
   "metadata": {},
   "source": [
    "# VISUALIZING A TRAINED NEURAL NETWORK"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84de871e",
   "metadata": {},
   "source": [
    "In this notebook, we will visualize the latent space of a trained neural network, and use it to generate protein interpolations. \n",
    "This notebook focusses on describing how to use the `MolearnAnalysis` and `MolearnGUI` classes. To find out more on their implementation, please see the [molearn_analysis.ipynb](molearn_analysis.ipynb) notebook. Let's start by loading a few useful classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ea02ce3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import sys, os\n",
    "\n",
    "#uncomment/edit as appropriate for your system\n",
    "#sys.path.insert(0, f\"C:{os.sep}Users{os.sep}xdzl45{os.sep}workspace{os.sep}molearn{os.sep}src\") \n",
    "\n",
    "import molearn\n",
    "from molearn.analysis import MolearnAnalysis\n",
    "from molearn.data import PDBData\n",
    "from molearn.models.foldingnet import AutoEncoder\n",
    "\n",
    "import torch\n",
    "from copy import deepcopy\n",
    "import biobox as bb\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cfdd9a6",
   "metadata": {},
   "source": [
    "## 1. Neural network analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee675e66",
   "metadata": {},
   "source": [
    "### 1.1. Data loading"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2836312",
   "metadata": {},
   "source": [
    "The `MolearnAnalysis` class features a series of methods to characterize the latent space of a neural network, assessing its performance, and generating new protein conformations. Its constructor requires two parameters: neural network parameters and a multiPDB containing the coordinates of the examples the neural network was trained with. Optional parameters enable the user to define the neural network architecture (if other than default) and atoms selected from training (backbone and beta carbon by default). We start by loading a trained neural network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "416c09e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "networkfile = f'data{os.sep}checkpoint_no_optimizer_state_dict_epoch167_loss0.003259085263643.ckpt'\n",
    "\n",
    "checkpoint = torch.load(networkfile, map_location=torch.device('cpu'))\n",
    "net = AutoEncoder(**checkpoint['network_kwargs'])\n",
    "net.load_state_dict(checkpoint['model_state_dict'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0306d0ad",
   "metadata": {},
   "source": [
    "Now we can create a `MolearnAnalysis` object, and load the network in it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6af6c1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "MA = MolearnAnalysis()\n",
    "MA.set_network(net)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34b4f251",
   "metadata": {},
   "source": [
    "Now we prepare the data used to train and test the network. Calling `PDBData.import_pdb` multiple times enables appending multiple datasets (these must contain the same number of atoms). After loading, the names of terminal oxygens renamed to \"O\" with `PDBData.atomselect`, and then atoms of interests can be selected via the `PDBData.import_pdb` method. One datasets are loaded, we can prepare them as inputs for a neural network, i.e. normalizing it and transforming it into a tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2e1d667",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = PDBData()\n",
    "data.import_pdb(f'data{os.sep}MurD_closed_selection.pdb')\n",
    "data.import_pdb(f'data{os.sep}MurD_open_selection.pdb')\n",
    "data.fix_terminal()\n",
    "data.atomselect(atoms = ['CA', 'C', 'N', 'CB', 'O'])\n",
    "data.prepare_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6716b0ac",
   "metadata": {},
   "source": [
    "When training the neural network, it is likely the dataset we just prepared was split into a training and a test set. We can replicate the way it was split with `PDBData.split`. By default, 10% examples are randomly selected as training set. To ensure that the same random split is produced, the parameter `manual_seed` can be defined (same splits are obtained when setting the same number). Once a split is produced, let's add our two resulting datasets to `MolearnAnalysis`, assigning them to an appropriate label name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9840b8b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_train, data_valid = data.split(manual_seed=25)\n",
    "MA.set_dataset(\"training\", data_train)\n",
    "MA.set_dataset(\"validation\", data_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94a6337b",
   "metadata": {},
   "source": [
    "We now also load separate simulation data, featuring MurD spontaneously switching from the closed to the open state (a phenomenon produced when when the closed state is simulated without ligand bound). This can act as our test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30190f46",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_test = PDBData()\n",
    "data_test.import_pdb(f'data{os.sep}MurD_closed_apo_selection.pdb')\n",
    "data_test.fix_terminal()\n",
    "data_test.atomselect(atoms = ['CA', 'C', 'N', 'CB', 'O'])\n",
    "data_test.std, data_test.mean = data.std, data.mean # use same std and mean as data used for training\n",
    "data_test.prepare_dataset()\n",
    "MA.set_dataset(\"test\", data_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e88d65b1",
   "metadata": {},
   "source": [
    "### 1.2. Dataset analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2d6cb02",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "<b>Warning:</b> Calculations can be made faster, but more memory intensive, via two attributes <code>MolearnAnalysis.batch_size</code> defines the number of examples that are simultaneously encoded/decoded, <code>MolearnAnalysis.processes</code> defined how many parallel threads to spawn when calculating DOPE and Ramachandran scores. If only a low amount of RAM is available, the user may want to consider keeping these numbers low (default is 1). \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "419d1d28",
   "metadata": {},
   "outputs": [],
   "source": [
    "MA.batch_size = 8\n",
    "MA.processes = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a63bd7fe",
   "metadata": {},
   "source": [
    "The following methods yield information on RMSD, DOPE, and Ramachandran score of training and test sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8ec172d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# RMSDs\n",
    "err_train = MA.get_error('training')\n",
    "err_test = MA.get_error('validation')\n",
    "\n",
    "#DOPE score\n",
    "dope_train_data = MA.get_dope('training')\n",
    "dope_test_data = MA.get_dope('validation')\n",
    "\n",
    "# Ramachandran scores\n",
    "rama_train_data = MA.get_ramachandran('training') \n",
    "rama_test_data = MA.get_ramachandran('validation')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd7f6792",
   "metadata": {},
   "source": [
    "### 1.3. Latent space scan"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "370c723a",
   "metadata": {},
   "source": [
    "the following methods perform a grid search of the latent space. The output of these commands is also stored internally in the `MolearnAnalysis` object. Here, we will generate a 50x50 grid surrounding the latent space projection of each dataset loaded, plus/minus 10%. Each of these grid points will be converted in its protein conformation counterpart, i.e. a new `grid` dataset is added to the `MolearnAnalysis` object."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec2762c7",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "<b>Warning:</b> depending on the sampling granularity and number of processors used, the operations in this section can take a long time! For instance, calculating the DOPE score of a 50x50 grid with a single processor can take >10 minutes.\n",
    "If you are in a rush, you can skip these steps, or comment out one of the two lines below. The rest of this notebook will still run, though you will not be able to visualize coloured landscapes in the GUI presented in the next section.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79a59a9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "MA.setup_grid(30)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4c46b51",
   "metadata": {},
   "source": [
    "We will assess the L2 norm in latent and 3D space, as well as the local Ramachandran scores (see Jupyter notebook `molearn_analysis.ipynb`). These act as heuristic for the percision of the neural network. Each landscape calculated with the methods described hereafter is both returned to the caller, and stored in the `MolearnAnalysis.surfaces` dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8aef916",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "landscape_err_latent, landscape_err_3d, xaxis, yaxis = MA.scan_error()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6be6f47e",
   "metadata": {},
   "outputs": [],
   "source": [
    "landscape_ramachandran, xaxis, yaxis = MA.scan_ramachandran()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2d03ed8",
   "metadata": {},
   "source": [
    "To produce a DOPE score landscape, the `MA.scan_dope` method is available. Compared to the two previous methods, this takes one more extra parameter, `refine`. If True (default), structures generated will be energy minimised before DOPE scoring. Here, we will just calculate the raw DOPE score landscape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4d9b9d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "landscape_dope_unrefined, xaxis, yaxis = MA.scan_dope(refine=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72be411f",
   "metadata": {},
   "source": [
    "An interesting plot, is one where the latent space is coloured as a function of how closely structures generated resemble a target structure. In this example, we pick the random conformation (we pick the conformation at position 42 in the test set tensor) from the test set and compare a grid of the latent space to it in terms of RMSD. As for other landscape scanning methods above, the result of this scan will be both returned to the user and stored internally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3436b52e",
   "metadata": {},
   "outputs": [],
   "source": [
    "landscape_target_rmsd, xaxis, yaxis = MA.scan_error_from_target(\"validation\", index=42)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "387ea945",
   "metadata": {},
   "source": [
    "The user can also create custom functions to quantify a desired property through the latent space. Any custom function must take as first parameter a (1xNx3) array describing the coordinates of N atoms. Additional parameters can the also be provided. The custom function is used by calling `MolearnAnalysis.scan_custom`. This takes as three parameters the custom function, a list of extra parameters to pass the function, and a string used as label. This label will be used to retrieve the data associated with this function, stored in the dictionary `MolearnAnalysis.custom_data`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2e8e1ab",
   "metadata": {},
   "source": [
    " In the example below, we measure the radius of gyration by implementing a custom function called `f_rgyr`. The function takes one extra parameter: a `Biobox` object describing the protein of interest (available in the `MolearnAnalysis` upon data loading)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e6ad817",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_rgyr(crd, M):\n",
    "    M.coordinates = deepcopy(crd)\n",
    "    M.set_current(0)\n",
    "    return bb.rgyr(M)\n",
    "\n",
    "rgyr, xaxis, yaxis = MA.scan_custom(f_rgyr, [MA.mol], \"rgyr\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a89cf50d",
   "metadata": {},
   "source": [
    "In this second example, we measure the distance between two atoms of interest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47d9ef2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_distance(crd, M, idx1, idx2):\n",
    "    M.coordinates = deepcopy(crd)\n",
    "    M.set_current(0)\n",
    "    return np.sqrt(np.sum(M.points[idx2]-M.points[idx1])**2)\n",
    "\n",
    "_, idx1 = MA.mol.atomselect(\"*\", 1, \"CA\", get_index=True) # find index of first atom of interest (alpha carbon of resid 1)\n",
    "_, idx2 = MA.mol.atomselect(\"*\", 300, \"CA\", get_index=True)  # find index of second atom of interest (alpha carbon of resid 300)\n",
    "dist, xaxis, yaxis = MA.scan_custom(f_distance, [MA.mol, idx1, idx2], \"dist\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efb61996",
   "metadata": {},
   "source": [
    "### 2. Neural Network visualisation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48fdbceb",
   "metadata": {},
   "source": [
    "Now that we have loaded some information in our `MolearnAnalysis` object, it's time to explore its contents! To this end, we will use a `MolearnGUI` object, creating an interactive interface!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0123215",
   "metadata": {},
   "outputs": [],
   "source": [
    "from molearn.analysis import MolearnGUI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c93c7851",
   "metadata": {},
   "outputs": [],
   "source": [
    "MolearnGUI(MA);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d3ea2bb",
   "metadata": {},
   "source": [
    "The interface is divided in three areas: a control panel (left), a 2D latent space representation (centre) and a 3D protein view (right, initially empty). Here are instructions on how to use it:\n",
    "\n",
    "* the **datasets drop down menu** enable displaying the projection in the latent space of the loaded datasets.\n",
    "\n",
    "* the **drop down menu and scroller** enable colouring the latent space surface in different ways. Colouring styles available will depend on whether grid sampling methods have been called in the MolearnAnalysis object, before the GUI was started.\n",
    "\n",
    "* the **2D surface is clickable**. Clicking multiple times enables defining a path. \n",
    "The coordinates of clicked points appear in the **editable text box** to the bottom left as (x1, y1, x2, y2, ...).\n",
    "\n",
    "* There are two different methods to define paths: as straight lines or as low-energy path. Method selection is controlled by the **Path** drop down menu. If the Euclidean method is selected, paths will be straight lines regularly sampled with a number of points defined by the **sampling** box. If the A* method is selected, the A* shortest path algorithm is chosen. The number of sampling point will be equal to the number of grid points traversed (note, a size 3 running average is applied to the path).\n",
    "\n",
    "* When clicking the latent space, a **3D protein structure** will appear on the right. The representation will feature a number of conformations equal to the number of path sampling points represented in the latent space.\n",
    "\n",
    "* The interpolation visible on the right can be saved into a multiPDB file via the button labelled \"**Save PDB**\" to the bottom of the menu on the left.\n",
    "\n",
    "* The buttons \"**Save state**\" and \"**Load state**\" enable you to save/load the current state of the GUI. In this way, you do not have to re-calculate everything again, the next time you execute this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1bfd4c8",
   "metadata": {},
   "source": [
    "MolearnGUI can be launched without any parameter. This is useful if you want to start by loading a previous GUI state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f64dbf4",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "MG = MolearnGUI();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "986330e8",
   "metadata": {},
   "source": [
    "Note that state files created via the button **\"Save PDB\"** use the Python package `pickle`. You can open them yourself with the following command (where `MA` is a `MolearnAnalysis` object and `waypoints` are the waypoints shown in the GUI):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6c73d4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = f'data{os.sep}state.p' #name of your state file\n",
    "\n",
    "import pickle\n",
    "MA, waypoints = pickle.load( open( fname, \"rb\" ) )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c3f61d3",
   "metadata": {},
   "source": [
    "You got to the end of this notebook! If you are interested in knowing what is happening inside the MolearnAnalysis object, see the Juptyer notebook `molearn_analysis.ipynb`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89e9d18b",
   "metadata": {},
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}