molearn_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MOLEARN ANALYSIS TUTORIAL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "[molearn](https://github.com/Degiacomi-Lab/molearn) is a Python package streamlining the implementation of machine learning models trainable on protein conformational spaces. This tutorial is dedicated to demonstrating the first principles of analysing the performance of a trained neural network. Implementations of what described below is available in the `MolearnAnalysis` class. You can find out how to use this class in the [molearn_GUI.ipynb](molearn_GUI.ipynb) notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial is divided in the following sections:\n",
    "\n",
    "[1. Load neural network and training set](#ref1)\n",
    "   - [1.1. The data](#ref11)\n",
    "   - [1.2. Data loading](#ref12)\n",
    "   - [1.3. Loading the neural network](#ref13)\n",
    "   \n",
    "   \n",
    "[2. Projection into latent space and back](#ref2)\n",
    " \n",
    " \n",
    "[3. Dataset analysis](#ref3)\n",
    "   -  [3.1. RMSD (or MSE) of input vs autoencoded](#ref31)\n",
    "   -  [3.2. DOPE score of input vs autoencoded](#ref32)\n",
    "   \n",
    "   \n",
    "[4. Full latent space characterization](#ref4)\n",
    "   - [4.1. Evaluate re-encoding error](#ref41)\n",
    "   - [4.2. Evaluate DOPE score](#ref42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get started, let's load some packages!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os, glob\n",
    "from copy import deepcopy\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim\n",
    "\n",
    "import nglview as nv\n",
    "import MDAnalysis as mda\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import modeller\n",
    "from modeller import *\n",
    "from modeller.scripts import complete_pdb\n",
    "\n",
    "import biobox as bb\n",
    "\n",
    "from tqdm import tqdm\n",
    "#edit path as required for your computer (or remove, if you installed molearn via conda-forge)\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.models.foldingnet import AutoEncoder\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Load neural network and training set <a class=\"anchor\" id=\"ref1\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1. the data <a class=\"anchor\" id=\"ref11\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MurD is a 47-kDa ATP-driven ligase responsible for the biosynthesis of a bacterial peptidoglycan precursor (UDP-N-acetylmuramoyl-L-alanyl-D-glutamate). When bound to its ligand, UDP-N-acetylmuramoyl-L-alanyl-D-alanine, MurD is stabilized in a closed conformation (PDB: [3UAG](https://www.rcsb.org/structure/3UAG)). In the absence of UDP-N-acetylmuramoyl-L-alanyl-D-alanine, MurD takes instead an open conformation (PDB: [1E0D](https://www.rcsb.org/structure/1E0D)). We have carried out molecular dynamics simulations of both states. For this tutorial, we will load a multi-PDB file containing a subset of conformations of both trajectories."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by defining the files containing the training set (variable `open_file` and `closed_file`), and test set (variable `test_set_file`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "open_file = f'data{os.sep}MurD_open_selection.pdb'\n",
    "closed_file = f'data{os.sep}MurD_closed_selection.pdb'\n",
    "test_set_file = f'data{os.sep}MurD_closed_apo_selection.pdb'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The open_file and closed_file contain training set MD simulations of the closed and open states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "universe_training = mda.Universe(open_file)\n",
    "w = nv.show_mdanalysis(universe_training)\n",
    "w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "universe_training = mda.Universe(closed_file)\n",
    "w = nv.show_mdanalysis(universe_training)\n",
    "w"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By taking the closed conformation and manually removing the ligand, the protein switches to an open conformation. This is the simulation we will used as test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "universe_test = mda.Universe(test_set_file)\n",
    "w = nv.show_mdanalysis(universe_test)\n",
    "w"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2. data loading <a class=\"anchor\" id=\"ref12\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now prepare the training set, transforming into a `torch.Tensor` making sure to extract from it only the atoms actually used during neural network training (CA, C, N, O, CB). The following cell yields a PDBData class with attributes:\n",
    " - `PDBData.dataset`, a pyTorch `Tensor` containing a desired subset of *normalized* atomic coordinates ready for submission to the neural network.\n",
    " - `PDBData.mean` and `PDBData.std`, the mean and standard deviation of the original input data, useful to rescale generated structures into atomic coordinates.\n",
    " - `PDBData.get_atominfo()` method that returns a list containing the names, resid, and resname of atoms in the training set (in the same order)\n",
    " - `PDBData.mol`, a  `biobox` instance corresponding to the coordinates in `PDBData.dataset`, useful as a PDB writer and offering useful helper functions, e.g. RMSD calculation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = molearn.PDBData()\n",
    "data.import_pdb(open_file)\n",
    "data.import_pdb(closed_file) # this will append closed_file structures onto open_files\n",
    "data.atomselect(atoms=['CA', 'C', 'CB', 'N', 'O'])\n",
    "data.prepare_dataset()\n",
    "\n",
    "training_set = data.dataset\n",
    "stdval = data.std\n",
    "meanval = data.mean\n",
    "mol = data.mol\n",
    "atom_names = data.get_atominfo()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`training_set` takes the form of a Nx3xM normalised `Tensor` object, where N is the number of examples, M the number of atoms, and 3 the x, y, z coordinates of each atom. This is the order of dimensions required by the neural network, a typical multiPDB has columns in the NxMx3 order."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(training_set.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now load the test set in the same way as the training set. This time, we will only retain the coordinates of atoms, normalised and ready to be submitted to the neural network (discarding all other outputs). The test set must contain the same atoms as the training set, in the same order. The two datasets should also be aligned, since *molearn* is not rototranslation invariant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = molearn.PDBData()\n",
    "test_data.import_pdb(test_set_file)\n",
    "test_data.atomselect(atoms=data.atoms) # use same atoms as in data\n",
    "test_data.std, test_data.mean = data.std, data.mean # use same std and mean as data\n",
    "test_data.prepare_dataset()\n",
    "\n",
    "test_set = data.dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3. loading the neural network <a class=\"anchor\" id=\"ref13\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now load the parameters of a trained neural network saved in the following file (variable `networkfile`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "networkfile = f'data{os.sep}checkpoint_no_optimizer_state_dict_epoch167_loss0.003259085263643.ckpt'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the parameters passed to the `AutoEncoder` costructor need to be the same as those used to build the trained neural network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "checkpoint = torch.load(networkfile, map_location=device)\n",
    "\n",
    "network = AutoEncoder(**checkpoint['network_kwargs'])\n",
    "network.load_state_dict(checkpoint['model_state_dict'])\n",
    "\n",
    "network.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(next(network.parameters()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Projection into latent space and back <a class=\"anchor\" id=\"ref2\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by projecting the whole training set into the latent space. Due to memory constraints we often can't simply encode or decode the entire dataset in one go like:\n",
    "```\n",
    "with torch.no_grad():\n",
    "    z = network.encode(training_set.float())\n",
    "```\n",
    "We will need to do this in batches instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 8\n",
    "\n",
    "def batched_encode(dataset):\n",
    "    # Preallocate space for the results to go\n",
    "    z = torch.empty(dataset.shape[0], 2, 1, dtype=dataset.dtype, device = dataset.device)\n",
    "    with torch.no_grad():\n",
    "        for i in tqdm(range(0, z.shape[0], BATCH_SIZE)):\n",
    "            z[i:i+BATCH_SIZE] = network.encode(dataset[i:i+BATCH_SIZE].float())\n",
    "    return z\n",
    "\n",
    "def batched_decode(latent_vector, n_atoms):\n",
    "    decoded = torch.empty(latent_vector.shape[0], 3, n_atoms, dtype=latent_vector.dtype, device = latent_vector.device)\n",
    "    with torch.no_grad():\n",
    "        for i in tqdm(range(0, decoded.shape[0], BATCH_SIZE)):\n",
    "            decoded[i:i+BATCH_SIZE] = network.decode(latent_vector[i:i+BATCH_SIZE].float())[:,:,:n_atoms]\n",
    "    return decoded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = batched_encode(dataset = training_set.float())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`z` is a (Nx2x1) `Tensor` object, featuring the projection in the latent space of each of the N examples in the training set. If this notebook is running in a computer with a GPU, it will be stored in its memory. To get the data so we can plot it, we should move it back to the CPU, and remove one dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_training = z.data.cpu().numpy()[:, :, 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`z` can be decoded back into structures that, hopefully, will closely resemble the training set. Note that we will have to clip the output to the size of dataset's degrees of freedom (an extra padding is present). As for the encoding part, while in principle it is possible to decode all the encoded training set in one shot as follows\n",
    "\n",
    "```\n",
    "with torch.no_grad():  \n",
    "    decoded_training = network.decode(z)[:, :, :training_set.shape[2]]\n",
    "```\n",
    "it is often impossible to do so owing to memory constraints. Let's decode our dateset in batches!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoded_training = batched_decode(latent_vector = z, n_atoms = training_set.shape[2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`decoded_training` is a now `Tensor` object stored in the GPU with the same dimensionality as the input `training_set`, and features normalised atomic coordinates. While it can already be directly compared with `training_set`, to obtain an actual protein structure we can observe, we should reorder the columns, move the data onto the CPU and \"un-normalise\" them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoded_protein = decoded_training[:, :, :training_set.shape[2]].data.cpu().numpy().transpose(0, 2, 1)\n",
    "print(decoded_protein.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ok, now let's repeat encoding and decoding with the test set!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#the following commented lines encode and decode all the test set in one shot\n",
    "#with torch.no_grad():\n",
    "#    z = network.encode(test_set.float())\n",
    "#    z_test = z.data.cpu().numpy()[:, :, 0]\n",
    "#    decoded_test = network.decode(z)[:, :, :test_set.shape[2]]\n",
    "\n",
    "z = batched_encode(dataset = test_set.float())\n",
    "z_test = z.data.cpu().numpy()[:, :, 0]\n",
    "decoded_test = batched_decode(latent_vector = z, n_atoms = test_set.shape[2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. dataset analysis <a class=\"anchor\" id=\"ref3\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1. RMSD (or MSE) of input vs autoencoded <a class=\"anchor\" id=\"ref31\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A way to evaluate the performance of the network, is to compare protein conformations before they are submitted to the neural network, and after they get encoded and decoded. Note that reporting such results for the training set only is poor practice, an independent test set should be submitted to evaluate the ability of the network to generalize beyond the examples provided for training. Let's calculate the MSE of both training and test sets!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_error(network, dataset, meanval, stdval, align=False, mol=\"\"):\n",
    "    '''\n",
    "    Calculate the reconstruction error of a dataset encoded and decoded by a trained neural network\n",
    "    '''\n",
    "\n",
    "    #z = network.encode(dataset.float())\n",
    "    z = batched_encode(dataset.float())\n",
    "    #decoded = network.decode(z)[:,:,:dataset.shape[2]]\n",
    "    decoded = batched_decode(z, dataset.shape[2])\n",
    "    err = []\n",
    "    for i in tqdm(range(dataset.shape[0])):\n",
    "\n",
    "        crd_ref = dataset[i].permute(1,0).unsqueeze(0).data.cpu().numpy()*stdval + meanval\n",
    "        crd_mdl = decoded[i].permute(1,0).unsqueeze(0).data.cpu().numpy()[:, :dataset.shape[2]]*stdval + meanval #clip the padding of models  \n",
    "        if align: # use Molecule Biobox class to calculate RMSD\n",
    "            mol.coordinates = deepcopy(crd_ref)\n",
    "            mol.set_current(0)\n",
    "            mol.add_xyz(crd_mdl[0])\n",
    "            rmsd = mol.rmsd(0, 1)\n",
    "        else:\n",
    "            rmsd = np.sqrt(np.sum((crd_ref.flatten()-crd_mdl.flatten())**2)/crd_mdl.shape[1]) # Cartesian L2 norm\n",
    "\n",
    "        err.append(rmsd)\n",
    "\n",
    "    return np.array(err)\n",
    "    \n",
    "rmsd_training = get_error(network, training_set, meanval, stdval)\n",
    "rmsd_test = get_error(network, test_set, meanval, stdval)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Time to produce two nice violin plots, reporting on the error of training and test set. Typically the latter will be slightly higher than the former, though ideally this difference should be as small as possible, indicating that the neural network can generalize well to structures it has never seen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(1, 1, 1)\n",
    "ax.violinplot([rmsd_training, rmsd_test], showmeans=True)\n",
    "ax.set_ylabel(\"RMSD / $\\AA$\")\n",
    "ax.set_xticks([1, 2])\n",
    "ax.set_xticklabels([\"training set\", \"test set\"]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2. DOPE score of input vs autoencoded <a class=\"anchor\" id=\"ref32\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The DOPE score is a heuristic function used to estimate the energy of a protein conformation. It is used by software like Modeller to define whether a protein model is accurate. Here, we start by defining a function calculating the DOPE score of a PDB file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dope_score(fname):\n",
    "    env = Environ()\n",
    "    env.libs.topology.read(file='$(LIB)/top_heav.lib')\n",
    "    env.libs.parameters.read(file='$(LIB)/par.lib')\n",
    "    mdl = complete_pdb(env, fname)\n",
    "    atmsel = Selection(mdl.chains[0])\n",
    "    score = atmsel.assess_dope()\n",
    "    return score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A positive value for a DOPE score indicates that the distriution of atoms is not what one should expect from a correct protein conformation (i.e. the protein structure is bad). However, a negative DOPE score does not immediately tell us how good the structure is. To tell whether the models generated by the neural network are good, we will compare their DOPE score with the DOPE score of the input dataset, that we can assume being good. We now calculate the DOPE score of the training and test sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dope(network, dataset, meanval, stdval, mol):\n",
    "    \n",
    "    # set residues names with protonated histidines back to generic HIS name (needed by DOPE score function)\n",
    "    testH = mol.data[\"resname\"].values\n",
    "    testH[testH == \"HIE\"] = \"HIS\"\n",
    "    testH[testH == \"HID\"] = \"HIS\"\n",
    "    mol.data[\"resname\"] = testH\n",
    "\n",
    "    #z = network.encode(dataset.float())\n",
    "    z = batched_encode(dataset.float())\n",
    "    #decoded = network.decode(z)[:,:,:dataset.shape[2]]\n",
    "    decoded = batched_decode(z, dataset.shape[2])\n",
    "    \n",
    "    dope_dataset = []\n",
    "    dope_decoded = []\n",
    "    for i in tqdm(range(dataset.shape[0])):\n",
    "\n",
    "        # calculate DOPE score of input dataset\n",
    "        crd_ref = dataset[i].permute(1,0).unsqueeze(0).data.cpu().numpy()*stdval + meanval\n",
    "        mol.coordinates = deepcopy(crd_ref)\n",
    "        mol.write_pdb(\"tmp.pdb\")\n",
    "        s = dope_score(\"tmp.pdb\")\n",
    "        dope_dataset.append(s)\n",
    "\n",
    "        # calculate DOPE score of decoded counterpart\n",
    "        crd_mdl = decoded[i].permute(1,0).unsqueeze(0).data.cpu().numpy()[:, :dataset.shape[2]]*stdval + meanval  \n",
    "        mol.coordinates = deepcopy(crd_mdl)\n",
    "        mol.write_pdb(\"tmp.pdb\")\n",
    "        s = dope_score(\"tmp.pdb\")\n",
    "        dope_decoded.append(s)\n",
    "    \n",
    "    return dope_dataset, dope_decoded\n",
    " \n",
    "dope_training, dope_training_decoded = get_dope(network, training_set, meanval, stdval, mol)\n",
    "dope_test, dope_test_decoded = get_dope(network, test_set, meanval, stdval, mol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(1, 1, 1)\n",
    "violins = ax.violinplot([dope_training, dope_training_decoded, dope_test, dope_test_decoded],\n",
    "                        showmeans = True, showextrema = True)\n",
    "\n",
    "ax.set_ylabel(\"DOPE / a.u.\")\n",
    "ax.set_xticks([1, 2, 3, 4])\n",
    "ax.set_xticklabels([\"training set\", \"training decoded\", \"test set\", \"test decoded\"]);\n",
    "\n",
    "colours=[\"red\", \"orange\", \"blue\", \"cyan\"]\n",
    "for i, violin in enumerate(violins['bodies']):\n",
    "    violin.set_color(colours[i])\n",
    "    violin.set_alpha(1)\n",
    "\n",
    "for partname in ('cbars','cmins','cmaxes','cmeans'):\n",
    "    vp = violins[partname]\n",
    "    vp.set_edgecolor(\"k\")\n",
    "    vp.set_linewidth(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. full latent space characterization <a class=\"anchor\" id=\"ref4\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we will grid sample the latent space. To this end, we start by identifying the region of interest defined by the coordinates of the training set projected into the latent space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = 50\n",
    "bx = (np.max(z_training[:, 0]) - np.min(z_training[:, 0]))*0.1 # 10% margins on x-axis\n",
    "by = (np.max(z_training[:, 1]) - np.min(z_training[:, 1]))*0.1 # 10% margins on y-axis\n",
    "xvals = np.linspace(np.min(z_training[:, 0])-bx, np.max(z_training[:, 0])+bx, samples)\n",
    "yvals = np.linspace(np.min(z_training[:, 1])-by, np.max(z_training[:, 1])+by, samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "<b>Warning:</b> depending on the sampling granularity, this operation can take a long time! Calculating the DOPE score of a 50x50 grid can take ~20 minutes. \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1. Evaluate re-encoding error <a class=\"anchor\" id=\"ref41\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each position in the latent space is decoded, then encoded, the decoded yet again. We assess the drift in latent space upon this iteration and the difference in two consecutive decoded structures with L2 norms. This acts as a heuristic for the precision of the neural network: the smaller these numbers, the more precise the network. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "surf_z = np.zeros((len(xvals), len(yvals))) # L2 norms in latent space (\"drift\")\n",
    "surf_c = np.zeros((len(xvals), len(yvals))) # L2 norms in Cartesian space\n",
    "\n",
    "with torch.no_grad():\n",
    "    \n",
    "    for x, i in tqdm(enumerate(xvals)):\n",
    "        for y, j in enumerate(yvals):\n",
    "    \n",
    "            # take latent space coordinate (1) and decode it (2)\n",
    "            z1 = torch.tensor([[[i,j]]]).float()\n",
    "            s1 = network.decode(z1)[:,:,:training_set.shape[2]]\n",
    "            \n",
    "            # take the decoded structure, re-encode it (3) and then decode it (4)\n",
    "            z2 = network.encode(s1)\n",
    "            s2 = network.decode(z2)[:,:,:training_set.shape[2]]\n",
    "            \n",
    "            surf_z[x,y] = np.sum((z2.numpy().flatten()-z1.numpy().flatten())**2) # Latent space L2, i.e. (1) vs (3)\n",
    "            surf_c[x,y] = np.sum((s2.numpy().flatten()-s1.numpy().flatten())**2) # Cartesian L2, i.e. (2) vs (4)\n",
    "            \n",
    "surf_c = np.sqrt(surf_c)\n",
    "surf_z = np.sqrt(surf_z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot error surfaces for L2 norms in the latent and Cartesian space, and compare the two. For visualisation purposes, we take the log of both quantities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X, Y = np.meshgrid(xvals, yvals)\n",
    "\n",
    "fig = plt.figure(figsize=(15, 5))\n",
    "ax = fig.add_subplot(1, 3, 1)\n",
    "ax.set_title(\"Latent space L2\")\n",
    "ax.pcolormesh(X, Y, np.log(surf_z))\n",
    "\n",
    "ax = fig.add_subplot(1, 3, 2)\n",
    "ax.set_title(\"3D space L2\")\n",
    "ax.pcolormesh(X, Y, np.log(surf_c))\n",
    "\n",
    "ax = fig.add_subplot(1, 3, 3)\n",
    "ax.set_xlabel(\"Latent space L2\")\n",
    "ax.set_ylabel(\"3D space L2\")\n",
    "ax.plot(surf_z.flatten(), surf_c.flatten(), \".\", alpha=0.2);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will notice that, although in general large drift corresponds to large RMSD, there is not a 1:1 relationship between the two quantities. This is because the latent space is not linear: the same magnitude of displacement in two different regions of the latent space may be associated to displacements of different magnitude in the 3D space. You may see several trends, overlayed, if the neural network has been trained with a collection of different molecular dynamics simulations projecting in different regions of the latent space."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2. Evaluate DOPE score <a class=\"anchor\" id=\"ref42\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now calculate the DOPE score of every conformation decoded from the latent space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "surf_dope = np.zeros((len(xvals), len(yvals)))\n",
    "\n",
    "with torch.no_grad():\n",
    "    \n",
    "    for x, i in tqdm(enumerate(xvals)):\n",
    "        for y, j in enumerate(yvals):\n",
    "    \n",
    "            # take latent space coordinate (1) and decode it (2)\n",
    "            z1 = torch.tensor([[[i,j]]]).float()\n",
    "            s1 = network.decode(z1)[:,:,:training_set.shape[2]]\n",
    "            \n",
    "            crd_mdl = s1[0].permute(1,0).unsqueeze(0).data.cpu().numpy()[:, :training_set.shape[2]]*stdval + meanval  \n",
    "            mol.coordinates = deepcopy(crd_mdl)\n",
    "            mol.write_pdb(\"tmp.pdb\")\n",
    "            s = dope_score(\"tmp.pdb\")\n",
    "            surf_dope[x,y] = s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = np.meshgrid(xvals, yvals)\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(1, 1, 1)\n",
    "ax.set_xlabel(\"Latent vector 1 / a.u.\")\n",
    "ax.set_ylabel(\"Latent vector 2 / a.u.\")\n",
    "c = ax.pcolormesh(x, y, surf_dope.T)\n",
    "#ax.plot(z_training[:, 0], z_training[:, 1], color=\"white\", alpha=0.5)\n",
    "cbar = fig.colorbar(c, ax=ax)\n",
    "cbar.set_label(\"DOPE score / a.u.\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the end of this overview. The methods described here are implemented in `MolearnAnalysis.py`. Have a look at the Jupyter notebook [molearn_GUI.ipynb](molearn_GUI.ipynb) to see it in action!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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": 4
}