tutorial1_DeepLDA.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "dc758926-6c33-4469-b02a-bf655e38d119",
    "deepnote_cell_height": 231.1875,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "# 1. DeepLDA: design CVs from equilibrium fluctuations\n",
    "\n",
    "Reference paper: _Bonati, Rizzi and Parrinello, [JCPL](https://pubs.acs.org/doi/10.1021/acs.jpclett.0c00535) (2020)_ [[arXiv]](https://arxiv.org/abs/2002.06562). \n",
    "\n",
    "The aim of this tutorial is to illustrate how we can design collective variables in a data-driven way, starting from local fluctuations of a set of physical descriptors in the metastable states. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00001-b0f589bf-04c6-46e5-ab89-9db34dcd34bf",
    "deepnote_cell_height": 62,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "### Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00002-e9a6185a-2896-4878-a5e1-1a5495cccd03",
    "deepnote_cell_height": 756.765625,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "**DeepLDA summary**\n",
    "\n",
    "To this extent, we resort to a statistical method called Linear Discriminant Analysis ([LDA](https://en.wikipedia.org/wiki/Linear_discriminant_analysis)). LDA searches for the linear projection $s=\\bar{w}^T x$ of the input data $x$ such that the classes are maximally separated. This is measured by the so called Fisher's ratio:\n",
    "\n",
    "$$\\bar{w} = \\text{argmax}_w {\\frac{wS_b w^T}{wS_w w^T}}$$\n",
    "\n",
    "where $S_b$ is the scatter matrix between the classes and $S_w$ the one between them. In the simple case of two classes (states) A and B these can be easily computed from the mean and covariance matrices of the data in the states:\n",
    "\n",
    "$$S_w = \\Sigma_A + \\Sigma_B $$\n",
    "\n",
    "$$S_b = (\\mu_A-\\mu_b)(\\mu_A-\\mu_b)^T $$\n",
    "\n",
    "![LDA](images/lda.png)\n",
    "\n",
    "From a practical perspective, the vector $\\bar{w}$ is found by solving the generalized eigenvalue problem: $S_b\\bar{w} = v S_w \\bar{w} $, where the eigenvalue $v$ measures the amount of separation between the states.\n",
    "\n",
    "Here we employ a non-linear generalization of LDA in which the mapping function is a neural network (NN). This greatly increases the discriminative power of the model, by learning a set of latent variables in which the metastable states are linearly separable. This is achieved by performing a nonlinear featurization of the inputs via a NN, and then LDA is applied to the outputs of the network. During the training, the parameters of the NN are optimized as to maximize the LDA eigenvalue $v$. In other words, we are transforming the input space in such a way that the discrimination between the states is maximal. \n",
    "\n",
    "Furthermore, in the context of enhanced sampling calculation, NNs lend themselves well because they provide a continuous and differentiable representation, and have no trouble in handling several descriptors. \n",
    "\n",
    "![DeepLDA scheme](images/DeepLDA_scheme.png)\n",
    "\n",
    "To achieve this we use the following loss function:\n",
    "\n",
    "$$ \\mathcal{L} = -v - \\alpha \\frac{1}{1+(\\bar{s}^2-1)^2} $$\n",
    "\n",
    "where the first term correspond to the maximization of the eigenvalue which describe the amount of separation of the two states and the second is a lorentzian regularization of the average value of the output CV, which keeps it close to 1.\n",
    "\n",
    "Note that to stabilize the learning we also regularize the calculation of $S_w$ by adding the identity matrix multiplied by a parameter $\\lambda$, as: $S_w '= S_w+\\lambda\\mathbb{1}$ . The two regularization parameters $\\alpha$ and $\\lambda$ are not independent, as the former affects the numerator of the Fisher's ratio above and the latter the denominator. Hence, in the following, we will choose only $\\lambda$ and set $\\alpha=\\frac{2}{\\lambda}$.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00003-ff7e2089-ed67-4594-b59f-c793d8aaebd0",
    "deepnote_cell_height": 178.390625,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "To apply this method, the only knowledge that is required as initial inputs are the snapshots of the system in different metastable states. These could be for instance the reactant and product of a chemical reaction, a crystalline and a liquid configurations of a material, the bound and unbound states of a ligand into a protein, etc...\n",
    "\n",
    "Of course, a variable built only by compressing the equilibrium fluctuations in the local minima of the FES into a CV will not be perfect. Rather, it should be seen as a first step which allows to observe some transitions and learn something new on the system that we can later use to refine it and find better CVs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00004-e8926b61-df1c-487c-a0b0-6e19755b45ad",
    "deepnote_cell_height": 400.78125,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "**Outline**\n",
    "\n",
    "In this tutorial, we take as example **alanine dipeptide** in vacuum, which is often used as a toy model for enhanced sampling methods. This molecule has two metastable states, called $C7_{eq}$ and $C7_{ax}$, which in the following we will refer to as A and B, respectively. \n",
    "\n",
    "To mimic a realistic scenario, we assume that we do not know anything about this system but for two realizations of states A and B. \n",
    "\n",
    "![Ala 2 metastable states](images/ala2-belfast-2-transition.png)\n",
    "\n",
    "Using these two realizations of the molecule we proceed as follows:\n",
    "\n",
    "1. We perform short unbiased MD simulations in the metastable states and evaluate a set of physical descriptors (e.g. interatomic distances between heavy atoms)\n",
    "2. We train the DeepLDA CV and inspect it\n",
    "3. Finally, we apply a bias potential to enhance the fluctuations of the DeepLDA CV and (hopefully) drive the system back and forth between A to B."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00005-56057813-a052-4c4b-b989-4c94e028a809",
    "deepnote_cell_height": 70,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00006-5eecd276-bd9b-4f6a-8d2e-20a3192eaeea",
    "deepnote_cell_height": 264,
    "deepnote_cell_type": "code"
   },
   "outputs": [],
   "source": [
    "import mlcvs\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import subprocess\n",
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "\n",
    "# delete outputs of simulations from `folder``\n",
    "def clean(folder='./'):\n",
    "    subprocess.run(\"rm bck.* COLVAR KERNELS alanine.*\", cwd=folder, shell=True)\n",
    "\n",
    "# execute bash command in the given folder\n",
    "def execute(command, folder, background=False):\n",
    "    cmd = subprocess.run(command, cwd=folder, shell=True, capture_output = True, text=True, close_fds=background)\n",
    "    if cmd.returncode == 0:\n",
    "        print(f'Completed: {command}')\n",
    "    else:\n",
    "        print(cmd.stderr)\n",
    "\n",
    "#GMX_CMD = '. /work/sourceme.sh && gmx_mpi'\n",
    "GMX_CMD = 'gmx_mpi'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00007-134f3833-5ee1-4854-8b7b-5bd9fd8a32d9",
    "deepnote_cell_height": 570,
    "deepnote_cell_type": "code"
   },
   "outputs": [],
   "source": [
    "# Define a few plotting functions\n",
    "\n",
    "def plot_ramachandran(x,y,z,scatter=None, ax=None):\n",
    "    # Setup plot\n",
    "    if ax is None:\n",
    "        _, ax = plt.subplots(figsize=(5,4.), dpi=100)\n",
    "        ax.set_title(f'Ramachandran plot')\n",
    "\n",
    "    # Plot countour plot\n",
    "    h = ax.hexbin(x,y,C=z,cmap='fessa')\n",
    "    cbar = plt.colorbar(h,ax=ax)\n",
    "    cbar.set_label(f'Deep-LDA CV')\n",
    "\n",
    "    ax.set_xlabel(r'$\\phi$ [rad]')\n",
    "    ax.set_ylabel(r'$\\psi$ [rad]')\n",
    "\n",
    "def plot_cv_histogram(s,label=None,ax=None,**kwargs):\n",
    "    # Setup plot\n",
    "    if ax is None:\n",
    "        _, ax = plt.subplots(figsize=(5,4.), dpi=100)\n",
    "        ax.set_title('Histogram')\n",
    "\n",
    "    if (type(s)==torch.Tensor):\n",
    "        s = s.squeeze(1).detach().numpy()\n",
    "\n",
    "    # Plot histogram\n",
    "    ax.hist(s,**kwargs)\n",
    "    if label is not None:\n",
    "        ax.set_xlabel(label)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00008-58dda8c1-e5e5-4488-89c8-8a79ab7a58d0",
    "deepnote_cell_height": 70,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "## 1.0 Unbiased simulations in the metastable states"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00009-e880b39b-6b4c-4192-a184-7727610eaadf",
    "deepnote_cell_height": 97.1875,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "First, we perform short MD simulations starting from the two snapshots of the molecule, and characterizing each state with a set of descriptors. To proceed in a blind way, we choose to use as input features all the distances between heavy atoms (the list of such descriptors are in the file `plumed-distances.dat`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00010-6a64a81c-4090-4ce3-b33b-5b912e6987dd",
    "deepnote_cell_height": 52.390625,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "**State A**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00011-db784728-6b80-4b91-a20f-e5b28c3827ea",
    "deepnote_cell_height": 1275.1875,
    "deepnote_cell_type": "code",
    "deepnote_output_heights": [
     null,
     21.1875
    ]
   },
   "outputs": [],
   "source": [
    "# CREATE FOLDER AND COPY INPUTS\n",
    "folder = '1_DeepLDA/0_unbiased-sA/'\n",
    "Path(folder).mkdir(parents=True, exist_ok=True)\n",
    "execute(f\"cp ../md_inputs/input.ala2.pdb ../md_inputs/input.tpr .\", folder=folder)\n",
    "\n",
    "# WRITE PLUMED INPUT FILE\n",
    "with open(folder+\"plumed.dat\",\"w\") as f:\n",
    "    print(\"\"\"\n",
    "# vim:ft=plumed\n",
    "\n",
    "# Compute torsion angles, as well as energy\n",
    "MOLINFO STRUCTURE=input.ala2.pdb\n",
    "phi: TORSION ATOMS=@phi-2\n",
    "psi: TORSION ATOMS=@psi-2\n",
    "theta: TORSION ATOMS=6,5,7,9\n",
    "xi: TORSION ATOMS=16,15,17,19\n",
    "ene: ENERGY\n",
    "\n",
    "# Compute descriptors\n",
    "INCLUDE FILE=../plumed-distances.dat\n",
    "\n",
    "# Print \n",
    "PRINT FMT=%g STRIDE=100 FILE=COLVAR ARG=*\n",
    "\n",
    "ENDPLUMED\n",
    "\"\"\",file=f)\n",
    "\n",
    "## RUN GROMACS\n",
    "num_steps=500000\n",
    "\n",
    "clean(folder) # note: this deletes all previous results in folder!\n",
    "execute(f\"{GMX_CMD} mdrun -s input.tpr -deffnm alanine -plumed plumed.dat -ntomp 1 -nsteps {num_steps} > alanine.out\", folder=folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00012-2cb94712-0541-427b-a7dd-e029a6aeea69",
    "deepnote_cell_height": 52.390625,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "**State B**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00013-9bbbd5ff-d532-4438-a25f-d78b187ece14",
    "deepnote_cell_height": 1257.1875,
    "deepnote_cell_type": "code",
    "deepnote_output_heights": [
     null,
     21.1875
    ]
   },
   "outputs": [],
   "source": [
    "folder = '1_DeepLDA/0_unbiased-sB/'\n",
    "Path(folder).mkdir(parents=True, exist_ok=True)\n",
    "execute(f\"cp ../md_inputs/input.ala2.pdb .\", folder=folder)\n",
    "execute(f\"cp ../md_inputs/input.sB.tpr input.tpr\", folder=folder)\n",
    "\n",
    "with open(folder+\"plumed.dat\",\"w\") as f:\n",
    "    print(\"\"\"\n",
    "# vim:ft=plumed\n",
    "\n",
    "# Compute torsion angles, as well as energy\n",
    "MOLINFO STRUCTURE=input.ala2.pdb\n",
    "phi: TORSION ATOMS=@phi-2\n",
    "psi: TORSION ATOMS=@psi-2\n",
    "theta: TORSION ATOMS=6,5,7,9\n",
    "xi: TORSION ATOMS=16,15,17,19\n",
    "ene: ENERGY\n",
    "\n",
    "INCLUDE FILE=../plumed-distances.dat\n",
    "\n",
    "PRINT FMT=%g STRIDE=100 FILE=COLVAR ARG=*\n",
    "\n",
    "ENDPLUMED\n",
    "\"\"\",file=f)\n",
    "\n",
    "## RUN GROMACS\n",
    "num_steps=500000\n",
    "\n",
    "clean(folder)\n",
    "execute(f\"{GMX_CMD} mdrun -s input.tpr -deffnm alanine -plumed plumed.dat -ntomp 1 -nsteps {num_steps} > alanine.out\", folder=folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00014-9a111a5f-73f6-4da7-b168-fc9d46d9194b",
    "deepnote_cell_height": 70,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "## 1.1 DeepLDA CV on pairwise distances (heavy atoms) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00015-9545aff6-bd5b-42be-847c-a3b98aed474a",
    "deepnote_cell_height": 62,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "### (a) Train CV"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00016-9c68db43-0319-40f3-b005-7ead48c4071d",
    "deepnote_cell_height": 97.1875,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "To load the PLUMED output of the two unbiased MD runs we can use the [load_dataframe](https://mlcvs.readthedocs.io/en/latest/autosummary/mlcvs.utils.io.load_dataframe.html) function.\n",
    "From this data, we build our training dataset. Since this is a supervised learning task, the dataset will be of the form (`X,y`), in which `X` are the input samples and `y` the corresponding labels (the states to which they belong to)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00017-c809af4f-236a-4968-84f7-28a0b7936b52",
    "deepnote_cell_height": 462,
    "deepnote_cell_type": "code"
   },
   "outputs": [],
   "source": [
    "from mlcvs.utils.io import load_dataframe\n",
    "# load state A and assign label 0\n",
    "folder = '1_DeepLDA/0_unbiased-sA/'\n",
    "colvarA = load_dataframe(folder+\"COLVAR\")\n",
    "colvarA['state']=np.full(len(colvarA),'A')\n",
    "colvarA['label']=np.full(len(colvarA),0)\n",
    "\n",
    "# load stateB and assign label 1\n",
    "folder = '1_DeepLDA/0_unbiased-sB/'\n",
    "colvarB = load_dataframe(folder+\"COLVAR\")\n",
    "colvarB['state']=np.full(len(colvarB),'B')\n",
    "colvarB['label']=np.full(len(colvarB),1)\n",
    "\n",
    "# concatenate data into a single dataframe\n",
    "colvar = pd.concat([colvarA,colvarB.reset_index(drop=True)])\n",
    "\n",
    "# create training dataset \n",
    "X = colvar.filter(regex='d_').values\n",
    "y = colvar['label'].values\n",
    "\n",
    "# transform them into torch.tensors \n",
    "X = torch.Tensor(X)\n",
    "y = torch.Tensor(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00018-b7ed746e-1845-433d-8d18-a6078abcbf64",
    "deepnote_cell_height": 88.796875,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "We can take a look at the descriptors, by computing their histogram in the two states. \n",
    "\n",
    "--> **Question:** Is there any descriptor that is able to discriminate by its own between the states?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00019-1c3ae9ef-87cb-4e8e-8300-7fa2e012233a",
    "deepnote_cell_height": 662.140625,
    "deepnote_cell_type": "code",
    "deepnote_output_heights": [
     421.140625
    ]
   },
   "outputs": [],
   "source": [
    "descriptors_names = colvar.filter(regex='d_').columns.values\n",
    "\n",
    "fig,axs = plt.subplots(5,9,figsize=(20,10),sharey=True)\n",
    "\n",
    "for ax,desc in zip(axs.flatten(),descriptors_names):\n",
    "    colvar.pivot(columns='state')[desc].plot.hist(bins=50,alpha=0.5,ax=ax,legend=False)\n",
    "    ax.set_title(desc)\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00020-d6c1cafc-f2e4-40e2-a823-0a810d4a6a2a",
    "deepnote_cell_height": 97.1875,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "Here we use the [mlcvs](https://mlcvs.readthedocs.io) package to train a DeepLDA CV out of this data. This can be as simple as follows: define the network architecture, specify when to stop to training (e.g. by using early stopping on validation score) and call the fit method. This will output the training and validation score along the training. Let's give it a try!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00021-0dc8f8b6-ec9e-41be-8f83-2f5da66bfffb",
    "deepnote_cell_height": 414.953125,
    "deepnote_cell_type": "code"
   },
   "outputs": [],
   "source": [
    "from mlcvs.lda import DeepLDA_CV\n",
    "nodes = [X.size(1),30,30,5]\n",
    "\n",
    "model = DeepLDA_CV(nodes)\n",
    "\n",
    "# TRAIN\n",
    "model.set_earlystopping(patience=20,min_delta=0.1)\n",
    "model.fit(X=X,y=y,log_every=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00022-c0ed2ef9-7fe6-4d92-b386-08661eab45a0",
    "deepnote_cell_height": 119.59375,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "However, to better understand what we are doing, we shall consider a more detailed example, in which we analyze the different steps and options. First we create a `TensorDataset` which we then divide into training and validation set. From them, we construct `Dataloader`-like objects. The definition of such auxiliary objects is a standard PyTorch practice, which allows us to easily train the models on different devices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00023-d92930e9-6e50-4f71-b694-9186b3e18976",
    "deepnote_cell_height": 228,
    "deepnote_cell_type": "code"
   },
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset,random_split\n",
    "from mlcvs.utils.data import FastTensorDataLoader\n",
    "\n",
    "dataset = TensorDataset(X,y)\n",
    "train_size = int(0.8 * len(dataset))\n",
    "valid_size = len(dataset) - train_size\n",
    "\n",
    "train_data, valid_data = random_split(dataset,[train_size,valid_size])\n",
    "train_loader = FastTensorDataLoader(train_data,batch_size=0,shuffle=True) # here 0 means to use a single batch\n",
    "valid_loader = FastTensorDataLoader(valid_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00024-c20b6349-b72b-4204-9089-9f63e9b50b3b",
    "deepnote_cell_height": 454.765625,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "Then, we need to inizialize the neural network and the optimizer and define when to stop the training (EarlyStopping or after a given number of epochs). The following is a list with all the parameters and their explanation. Note that we also standardize the inputs such that their range is betwen the -1 and 1 in the training set. \n",
    "\n",
    "| Parameter | Type | Description |\n",
    "| :- | :- | :- |\n",
    "| **Neural network** |\n",
    "| nodes | list | NN architecture (last value equal to the number of hidden layers which are input of LDA) |\n",
    "| activation | string | Activation function (relu,tanh,elu,linear) |\n",
    "| **Optimization** |\n",
    "| lrate | float | Learning rate |\n",
    "| sw_reg | float | S_w matrix regularization ($\\lambda$)| \n",
    "| l2_reg | float | L2 regularization |\n",
    "| num_epochs | int | Number of epochs |\n",
    "| **Early Stopping** |\n",
    "| es_patience | int | Number of epochs before stopping |\n",
    "| es_consecutive | bool | Whether es_patience should count consecutive (True) or cumulative patience |\n",
    "| es_min_delta | float | Minimum decrease of validation loss |\n",
    "| **Log** |\n",
    "| log_every | int | How often print the train/valid loss during training |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00025-ecbb56c9-6014-47cf-95ca-7fa481d9809a",
    "deepnote_cell_height": 874.46875,
    "deepnote_cell_type": "code"
   },
   "outputs": [],
   "source": [
    "from mlcvs.lda import DeepLDA_CV\n",
    "\n",
    "#------------- PARAMETERS -------------\n",
    "nodes             = [X.size(1),30,30,5]\n",
    "activation        = 'tanh'\n",
    "\n",
    "lrate             = 0.001\n",
    "sw_reg            = 0.05\n",
    "l2_reg            = 1e-5\n",
    "\n",
    "num_epochs        = 1000\n",
    "earlystop         = True\n",
    "es_patience       = 20\n",
    "es_consecutive    = True\n",
    "es_min_delta      = 0.1\n",
    "\n",
    "log_every         = 100\n",
    "#--------------------------------------\n",
    "\n",
    "# DEVICE: check if there is a GPU\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# MODEL: initialize the DeepLDA CV object\n",
    "model = DeepLDA_CV(nodes,activation)\n",
    "model.to(device)\n",
    "\n",
    "# OPTIMIZER: here we use adam\n",
    "opt = torch.optim.Adam(model.parameters(), lr=lrate, weight_decay=l2_reg)\n",
    "model.set_optimizer(opt)\n",
    "# set criterion for stopping the learning --> avoid overfitting\n",
    "model.set_earlystopping(patience=es_patience,consecutive=es_consecutive,min_delta=es_min_delta)\n",
    "\n",
    "# REGULARIZATION: add regularization to the calculation of the S_w matrix\n",
    "model.set_regularization(sw_reg=sw_reg)\n",
    "\n",
    "# TRAIN: fit the model to maximize Fisher's discriminant ratio\n",
    "model.fit(train_loader,valid_loader, \n",
    "            standardize_inputs = True, \n",
    "            standardize_outputs = False,\n",
    "            log_every=log_every)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00026-05c5f35d-69aa-410e-ae1e-6dbb327f702b",
    "deepnote_cell_height": 133.59375,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "After the training we can plot the learning curve to see the training and validation score.\n",
    "\n",
    "--> **Exercise**: You can try to repeat the training in the cell above and make sure that you find similar values for the loss function. You can examine the behaviour of the NN when changing the NN architecture, and in particular the number of outputs which regulates the space in which LDA is applied."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00027-66660d74-b8c6-47f5-a4a9-d42b8aa4dba4",
    "deepnote_cell_height": 777,
    "deepnote_cell_type": "code",
    "deepnote_output_heights": [
     392
    ]
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(5,4),dpi=100)\n",
    "\n",
    "loss_train = [x.cpu().numpy() for x in model.loss_train]\n",
    "loss_valid = [x.cpu().numpy() for x in model.loss_valid]\n",
    "\n",
    "# Loss function\n",
    "ax.plot(loss_train,'-',label='Train')\n",
    "ax.plot(loss_valid,'--',label='Valid')\n",
    "ax.set_ylabel('Loss Function')\n",
    "\n",
    "#if model.earlystopping_.early_stop:\n",
    "#    ax.axvline(model.earlystopping_.best_epoch,ls='dotted',color='grey',alpha=0.5,label='Early Stopping')\n",
    "ax.set_xlabel('#Epochs')\n",
    "ax.legend(ncol=2)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00028-0436f832-648e-4e85-befd-3dfc488836c1",
    "deepnote_cell_height": 88.796875,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "To better understand what the network is doing, we can inspect the output of the NN, before the application of LDA. \n",
    "\n",
    "--> **Question**: What are the differences with the input descriptors?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00029-24ef25a5-2c0e-4eb5-93e5-acd0e8898d72",
    "deepnote_cell_height": 593.328125,
    "deepnote_cell_type": "code",
    "deepnote_output_heights": [
     208.328125
    ]
   },
   "outputs": [],
   "source": [
    "model.output_hidden=True\n",
    "with torch.no_grad():\n",
    "    hidden = model(X).numpy()\n",
    "df = pd.DataFrame( hidden )\n",
    "df ['label'] = y\n",
    "model.output_hidden=False\n",
    "\n",
    "# Plot histogram in the two state\n",
    "fig,axs = plt.subplots(1,nodes[-1],figsize=(16,4),sharey=True)\n",
    "\n",
    "hidden_names = [i for i in range(nodes[-1])]\n",
    "\n",
    "for ax,desc in zip(axs.flatten(),hidden_names):\n",
    "    df.pivot(columns='label')[desc].plot.hist(bins=50,alpha=0.5,ax=ax,legend=False)\n",
    "    ax.set_title(desc)\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00030-c4d53d74-c9de-4c7e-a749-33fa233fe276",
    "deepnote_cell_height": 97.1875,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "Finally, we can look at the histogram of the CV (right), which shows that the two states are mapped around -1 and +1. To appreciate the discriminative power, we can also inspect the Ramachandran plot (left) of the two torsion angles phi and psi where we colored the points according to the value of the Deep-LDA CV."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00031-9443c60c-9112-4711-9fe8-d6226c102e0c",
    "deepnote_cell_height": 647.421875,
    "deepnote_cell_type": "code",
    "deepnote_output_heights": [
     334.421875
    ]
   },
   "outputs": [],
   "source": [
    "_, axs = plt.subplots(1,2, figsize=(10,4.), dpi=100)\n",
    "\n",
    "axs[0].set_title('Ramachandran plot')\n",
    "with torch.no_grad():\n",
    "    s = model(X)\n",
    "    \n",
    "plot_ramachandran(colvar['phi'],colvar['psi'],s,ax=axs[0])\n",
    "\n",
    "# Calculate CV values over training set\n",
    "axs[1].set_title(f'Deep-LDA Histogram')\n",
    "plot_cv_histogram(s,label=model.name_,bins=50, ax=axs[1] )\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Number of NN outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nn_out = [10,5,3,1]\n",
    "\n",
    "fig, axs = plt.subplots(2,len(nn_out),figsize=(16,8),dpi=100)\n",
    "\n",
    "for i,out in enumerate(nn_out):\n",
    "    print(f'NN outputs = {out}')\n",
    "    # ARCHITECTURE\n",
    "    nodes = [X.size(1),30,30,out]\n",
    "    model = DeepLDA_CV(nodes)\n",
    "\n",
    "    # TRAIN\n",
    "    model.set_earlystopping(patience=20,min_delta=0.05,consecutive=True)\n",
    "    model.fit(X=X,y=y,log_every=100)\n",
    "\n",
    "    # PREDICT\n",
    "    with torch.no_grad():\n",
    "        s = model(X)\n",
    "    \n",
    "    # PLOT\n",
    "    axs[0][i].set_title(f'NN outputs = {out}')\n",
    "    plot_ramachandran(colvar['phi'],colvar['psi'],s,ax=axs[0][i])\n",
    "    axs[1][i].set_title(f'Deep-LDA Histogram')\n",
    "    plot_cv_histogram(s,label=model.name_,bins=50, ax=axs[1][i] )\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Extra: features relevance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compute the relevance of the input features from the derivatives of the cv $s$ with respect to the descriptors $d_i$:\n",
    "$$r_i = \\sum_{j=1} ^{n} |\\frac{\\partial s^{(j)}}{\\partial d_i ^{(j)}}|\\ \\sigma(d_i)$$\n",
    "where the sum is over the configurations in the training set and $\\sigma$ is the standard deviation of the descriptor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters\n",
    "multiply_by_stddev = True #whether to multiply derivatives by std dev of inputs\n",
    "order_by_importance = True #plot results ordered by importance\n",
    "\n",
    "#input names \n",
    "input_names = colvar.filter(regex='d_').columns.values\n",
    "n_input = len(input_names)\n",
    "\n",
    "#init arrays\n",
    "in_num=np.arange(n_input)\n",
    "rank=torch.zeros(n_input)\n",
    "\n",
    "#compute input std dev\n",
    "if multiply_by_stddev:\n",
    "    in_std=torch.std(X,axis=0).numpy()\n",
    "\n",
    "#compute the derivatives of the outputs w.r.t. inputs over all training set\n",
    "for (x,Y) in train_loader:    \n",
    "    x.requires_grad=True\n",
    "    for x_i,y_i in zip(x,Y):\n",
    "        # calculate cv \n",
    "        s_i = model(x_i) \n",
    "        # calculate derivatives\n",
    "        grad_i = torch.autograd.grad(s_i,x_i)\n",
    "        # accumulate them\n",
    "        rank += grad_i[0].abs()\n",
    "\n",
    "rank = rank.numpy()\n",
    "\n",
    "#multiply by std dev\n",
    "if multiply_by_stddev:\n",
    "    rank = rank * in_std\n",
    "\n",
    "#normalize to 1\n",
    "rank/= np.sum(rank)\n",
    "\n",
    "#sort\n",
    "if order_by_importance:\n",
    "    index= rank.argsort()\n",
    "    input_names = input_names[index]\n",
    "    rank = rank[index]\n",
    "\n",
    "#plot\n",
    "fig=plt.figure(figsize=(5,0.25*n_input), dpi=100)\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "if order_by_importance:\n",
    "    ax.barh(in_num, rank,color='fessa1',edgecolor = 'fessa0',linewidth=0.3)\n",
    "    ax.set_yticklabels(input_names,fontsize=9)\n",
    "else:\n",
    "    ax.barh(in_num[::-1], rank[::-1],color='fessa1',edgecolor = 'fessa0',linewidth=0.3)\n",
    "    ax.set_yticklabels(input_names[::-1],fontsize=9)\n",
    "\n",
    "ax.set_xlabel('Relevance')\n",
    "ax.set_ylabel('Inputs')\n",
    "ax.set_yticks(in_num)\n",
    "ax.yaxis.tick_right()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 4\n",
    "\n",
    "fig, axs = plt.subplots(1,n,figsize=(n*5,4),dpi=100)\n",
    "\n",
    "for i in range(n):\n",
    "    ax = axs[i]\n",
    "    name = input_names[::-1][i]\n",
    "    colvar.pivot(columns='state')[name].plot.hist(bins=50,alpha=0.5,ax=ax,legend=False)\n",
    "    ax.set_title(name)\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00032-46d9677b-4a3e-4661-82ff-a583ec9fd308",
    "deepnote_cell_height": 62,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "### (b) Bias DeepLDA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00033-6f86c101-2068-447e-a3ef-44f09473d1ea",
    "deepnote_cell_height": 363.171875,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "Once we have designed our CV from the data, we can use to to enhance the sampling. To do so, we compile the model using `torch.jit` (this is done inside the `export` function) which creates a python-independent file which can be loaded in PLUMED via the `PYTORCH_MODEL` function. If you look in the export folder below, you will find two files: `model_checkpoint.pt` which can be used to load the model back to Python, and a compiled one, `model.ptc` which we will load in the PLUMED input file. \n",
    "\n",
    "If you are curious about how to export pythorch-based functions with jit have a look at the source code of the [export](https://mlcvs.readthedocs.io/en/latest/autosummary/mlcvs.models.NeuralNetworkCV.html#mlcvs.models.NeuralNetworkCV.export) method.\n",
    "\n",
    "PLUMED will use the pytorch C++ APIs to load the model and evaluate it together with its derivatives with automatic differentiation. The outputs of the model will be stored in components called `deep.node-0,deep.node-1,...`, where `deep` is the label assigned to the PLUMED function (see input below).\n",
    "\n",
    "To apply the bias potential, our method of choice is [OPES](https://www.plumed.org/doc-master/user-doc/html/_o_p_e_s.html), as it as several advantages to metadynamics (fewer parameters, quickly converges to a quasi-static regime, performs a kernel merging which allows to bias more efficiently multiple CVs, well-behaved adaptive sigma, it can limit the amount of bias deposited...). However, this variable can be used with any other CV-based methods.\n",
    "\n",
    "--> **Exercise**: look at the PLUMED input file below and fill the missing parameters, then run gromacs.\n",
    "\n",
    "**Note**: on Deepnote this could require quite some time. If you prefer you can open the terminal and run gromacs on the background by executing the command in the argument of `execute` while you keep on playing around with the CVs training. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00034-2c6ee548-50d5-43ab-b05c-f447f11a78cf",
    "deepnote_cell_height": 768,
    "deepnote_cell_type": "code"
   },
   "outputs": [],
   "source": [
    "folder = '1_DeepLDA/1_opes-deeplda/'\n",
    "Path(folder).mkdir(parents=True, exist_ok=True)\n",
    "execute(f\"cp ../0_unbiased-sA/input* .\", folder=folder)\n",
    "\n",
    "# export model\n",
    "model.export(folder)\n",
    "\n",
    "# write plumed input\n",
    "with open(folder+\"plumed.dat\",\"w\") as f:\n",
    "    print(\"\"\"\n",
    "# vim:ft=plumed\n",
    "\n",
    "# Compute torsion angles, as well as energy\n",
    "MOLINFO STRUCTURE=input.ala2.pdb\n",
    "phi: TORSION ATOMS=@phi-2\n",
    "psi: TORSION ATOMS=@psi-2\n",
    "theta: TORSION ATOMS=6,5,7,9\n",
    "xi: TORSION ATOMS=16,15,17,19\n",
    "ene: ENERGY\n",
    "\n",
    "# Compute descriptors\n",
    "INCLUDE FILE=../plumed-distances.dat\n",
    "\n",
    "# Compute DeepLDA CV\n",
    "deep: PYTORCH_MODEL FILE=____FILL____ ARG=d_2_5,d_2_6,d_2_7,d_2_9,d_2_11,d_2_15,d_2_16,d_2_17,d_2_19,d_5_6,d_5_7,d_5_9,d_5_11,d_5_15,d_5_16,d_5_17,d_5_19,d_6_7,d_6_9,d_6_11,d_6_15,d_6_16,d_6_17,d_6_19,d_7_9,d_7_11,d_7_15,d_7_16,d_7_17,d_7_19,d_9_11,d_9_15,d_9_16,d_9_17,d_9_19,d_11_15,d_11_16,d_11_17,d_11_19,d_15_16,d_15_17,d_15_19,d_16_17,d_16_19,d_17_19\n",
    "# Apply OPES bias \n",
    "opes: OPES_METAD ARG=deep.node-0 PACE=500 BARRIER=30\n",
    "\n",
    "# Print \n",
    "PRINT FMT=%g STRIDE=500 FILE=COLVAR ARG=*\n",
    "\n",
    "ENDPLUMED\n",
    "\"\"\",file=f)\n",
    "\n",
    "## RUN GROMACS\n",
    "num_steps=2500000\n",
    "\n",
    "clean(folder)\n",
    "execute(f\"{GMX_CMD} mdrun -s input.tpr -deffnm alanine -plumed plumed.dat -ntomp 1 -nsteps {num_steps} > alanine.out\", folder=folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00036-ddeb21f0-dfad-4b74-9fd2-8f0785a49525",
    "deepnote_cell_height": 162.796875,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "Once the simulation is over, we can plot the time evolution of the Deep-LDA CV, in which several transitions between the states A and B (-1 and 1) can be observed. Furthermore, we can look at the Ramachandran plot to see the explored region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00037-04448e12-e51f-4031-b21a-0600965dd4be",
    "deepnote_cell_height": 705.15625,
    "deepnote_cell_type": "code",
    "deepnote_output_heights": [
     338.15625
    ]
   },
   "outputs": [],
   "source": [
    "folder = '1_DeepLDA/1_opes-deeplda/'\n",
    "colvar = load_dataframe(folder+'COLVAR')\n",
    "\n",
    "fig,axs = plt.subplots(1,2,figsize=(10,4),dpi=100)\n",
    "# Time evolution (DeepLDA)\n",
    "colvar.plot.scatter('time','deep.node-0',s=1,ax=axs[0])\n",
    "axs[1].set_xlabel('Time [ps]')\n",
    "axs[1].set_xlabel('DeepLDA')\n",
    "# 2D scatter plot colored with DeepLDA\n",
    "colvar.plot.scatter('phi','psi',c='deep.node-0',s=1,cmap='fessa',ax=axs[1])\n",
    "axs[1].set_xlabel(r'$\\phi$ [rad]')\n",
    "axs[1].set_ylabel(r'$\\psi$ [rad]')\n",
    "axs[1].set_aspect('equal')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00038-a2d6b69c-7f56-402c-9e5a-cbb2a9eef055",
    "deepnote_cell_height": 111.1875,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "From this simulation we can compute the free energy surface as a function of the DeepLDA CV. To do so, we use the `compute_fes` function which performs a (weighted) block average. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00039-8984f94f-4edf-4989-95f6-e7978178eb87",
    "deepnote_cell_height": 773.1875,
    "deepnote_cell_type": "code",
    "deepnote_output_heights": [
     21.1875,
     372
    ]
   },
   "outputs": [],
   "source": [
    "from mlcvs.utils.fes import compute_fes\n",
    "\n",
    "s = colvar['deep.node-0'].values\n",
    "\n",
    "# compute weights\n",
    "kbT = 2.5\n",
    "w = np.exp(colvar['opes.bias'].values/kbT)\n",
    "\n",
    "fig,ax = plt.subplots(figsize=(6,4),dpi=100)\n",
    "fes,grid,bounds,error = compute_fes(s, weights=w, kbt=kbT, \n",
    "                                    blocks=5, bandwidth=0.01, \n",
    "                                    plot=True, ax = ax)\n",
    "ax.set_xlabel('DeepLDA')\n",
    "ax.set_ylabel('FES [kJ/mol]')\n",
    "ax.set_ylim(0,50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "--> **Exercise**: calculate the FES as a function of other variables, such as the Ramachandran angles `phi` and `psi`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: the function `compute_fes` uses the function `KernelDensity` from scikit-learn package to perform Kernel Density Estimation (KDE). While this allows us to easily perform KDE on many dimensions, we need to take into account the periodicity of the CV. One simple solution is to perform data augmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def augment_periodic(cv,weights=None,bandwidth=0.1):\n",
    "    \"\"\"Add points across the periodic boundaries (-np.pi and np.pi)\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    cv : np.array\n",
    "    weigths : np.array, optional\n",
    "    bandwidth : float, optional\n",
    "    \"\"\"\n",
    "    mask = (cv < -np.pi + 3*bandwidth)\n",
    "    index = np.argwhere(mask)[:,0]\n",
    "    cv = np.insert( cv, index, cv[mask] + 2*np.pi )\n",
    "    if weights is not None:\n",
    "        weights = np.insert( weights, index, weights[mask] )\n",
    "\n",
    "    mask = (cv > np.pi  - 5*bandwidth) & ( cv < np.pi )\n",
    "    index = np.argwhere(mask)[:,0]\n",
    "    cv = np.insert( cv, np.argwhere(mask)[:,0], cv[mask] - 2*np.pi )\n",
    "    if weights is not None:\n",
    "        weights = np.insert( weights, index, weights[mask] )\n",
    "\n",
    "    return cv, weights\n",
    "\n",
    "phi = colvar['phi'].values\n",
    "psi = colvar['psi'].values\n",
    "w = np.exp(colvar['opes.bias'].values/kbT)\n",
    "\n",
    "bandwidth = 0.05\n",
    "\n",
    "phi, w_phi = augment_periodic(phi,w,bandwidth)\n",
    "psi, w_psi = augment_periodic(psi,w,bandwidth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlcvs.utils.fes import compute_fes\n",
    "\n",
    "# compute weights\n",
    "kbT = 2.5\n",
    "\n",
    "#w = np.exp(colvar['opes.bias'].values/kbT)\n",
    "\n",
    "#phi = colvar['phi'].values\n",
    "\n",
    "fig,axs = plt.subplots(1,2,figsize=(12,4),dpi=100)\n",
    "ax = axs[0]\n",
    "fes,grid,bounds,error = compute_fes(phi, weights=w_phi, kbt=kbT, \n",
    "                                    blocks=5, bandwidth=bandwidth, \n",
    "                                    plot=True, ax = ax)\n",
    "ax.set_xlabel(r'$\\phi$ [rad]')\n",
    "ax.set_ylabel('FES [kJ/mol]')\n",
    "ax.set_xlim(-np.pi,np.pi)\n",
    "ax.set_ylim(0,50)\n",
    "\n",
    "ax = axs[1]\n",
    "fes,grid,bounds,error = compute_fes(psi, weights=w_psi, kbt=kbT, \n",
    "                                    blocks=5, bandwidth=bandwidth, \n",
    "                                    plot=True, ax = ax)\n",
    "ax.set_xlabel(r'$\\psi$ [rad]')\n",
    "ax.set_ylabel('FES [kJ/mol]')\n",
    "ax.set_xlim(-np.pi,np.pi)\n",
    "ax.set_ylim(0,50)\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00040-87de779f-38dc-4f0f-aa92-f305b05841b9",
    "deepnote_cell_height": 155.984375,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "Once we have obtained a reasonable sampling, we can try to understand what the NN has learnt. A naive guess would be to compute the correlation between the DeepLDA CV and all the input distances. \n",
    "\n",
    "--> **Question**: Can we identify some features that are more correlated than others? You can also try to use `method='spearman'` in the correlation function, which rather than looking for a linear correlation only assesses how well the relationship between two variables can be described using a monotonic function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00041-cdb7e08a-7cde-49c6-9c0c-dfc3f55d29f9",
    "deepnote_cell_height": 480.0625,
    "deepnote_cell_type": "code",
    "deepnote_output_heights": [
     239.0625
    ]
   },
   "outputs": [],
   "source": [
    "cols = ['deep.node-0']\n",
    "cols.extend(colvar.filter(regex='d_').columns)\n",
    "corr = colvar[cols].corr(method='pearson') \n",
    "\n",
    "fig,ax = plt.subplots(figsize=(16,4),dpi=100)\n",
    "\n",
    "corr['deep.node-0'].drop('deep.node-0').plot(kind='bar', ax=ax, rot=35)\n",
    "ax.set_ylabel('Correlation with DeepLDA')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00042-2bbba021-dfcc-4e52-8bdc-6af177befd3f",
    "deepnote_cell_height": 133.59375,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "Since the relationship between the input distances and the DeepLDA CV is not linear, understanding what the NN has learnt might not be an easy job. Hovewer, we could try to see whether we detect any correlation with respect to other physical descriptors which might play a role in the transition between the two states. \n",
    "\n",
    "--> **Exercise**: compute the correlation between the DeepLDA CV and the torsion angles computed in the PLUMED input file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00043-e175f154-4fde-490c-b5f1-43d7cc95bd5b",
    "deepnote_cell_height": 246,
    "deepnote_cell_type": "code"
   },
   "outputs": [],
   "source": [
    "# select deelda and input distances, as well as dihedral angles\n",
    "cols = ['deep.node-0', ______FILL______ ]\n",
    "\n",
    "# compute correlation\n",
    "corr = colvar[cols].corr(method='spearman')\n",
    "\n",
    "# plot\n",
    "fig,ax = plt.subplots(figsize=(4,4),dpi=100)\n",
    "corr['deep.node-0'].drop('deep.node-0').plot(kind='bar', ax=ax, rot=35)\n",
    "ax.set_ylabel('Correlation with DeepLDA')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,axs = plt.subplots(1,len(cols)-1,figsize=(16,4),dpi=100)\n",
    "\n",
    "state = np.zeros(len(colvar))\n",
    "\n",
    "state[colvar['deep.node-0']<-1.1] = -1\n",
    "state[colvar['deep.node-0']>1.1] = 1\n",
    "\n",
    "colvar['state'] = state\n",
    "\n",
    "for i,desc in enumerate(cols[1:]):\n",
    "    colvar.plot.scatter('deep.node-0',desc,c='state',s=1,ax=axs[i],cmap='fessa',colorbar=False)\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00045-abdfa292-d7ef-42b2-b1da-543aff4f215a",
    "deepnote_cell_height": 62,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "### Bonus exercise: stretching the CV"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "cell_id": "00046-2f2e4581-7c3f-4d97-895a-71d072ee9f08",
    "deepnote_cell_height": 259.578125,
    "deepnote_cell_type": "markdown"
   },
   "source": [
    "As you might have noticed, the DeepLDA CV acts as a powerful classifier, which maps the equilibrium fluctuations into very narrow distributions. Although we were using OPES with adaptive bandwitdh estimation, this might still lead to artifacts in the enhanced sampling dynamics. To address this behaviour we can work on the regularization of the NN (e.g. penalty functions to the loss function) or, in this case, we can more simply stretch it. Since the states are mapped (due to the lorentzian regularization) on a sphere of radius ~ 1.1, if we use a function of the kind `s' = s + s^N` we obtain a transformation of the CV which is approximately linear around zero and amplifies the fluctuations around 1. In the picture below you can see the case for N=3. \n",
    "\n",
    "You can try to repeat the enhanced sampling simulation biasing a function of the CV, using the `CUSTOM` PLUMED keyword, e.g.:\n",
    "\n",
    "`deep_mod: CUSTOM ARG=deep.node-0 FUNC=x+x^3 PERIODIC=NO`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cell_id": "00047-881e703b-c95e-47d0-a77b-5f90dd00d4d6",
    "deepnote_cell_height": 401,
    "deepnote_cell_type": "code",
    "deepnote_output_heights": [
     250
    ]
   },
   "outputs": [],
   "source": [
    "x = np.linspace(-1.2,1.2,100)\n",
    "plt.plot(x,x)\n",
    "plt.plot(x,x+x**3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "folder = '1_DeepLDA/2_opes-deeplda-stretch/'\n",
    "Path(folder).mkdir(parents=True, exist_ok=True)\n",
    "execute(f\"cp ../0_unbiased-sA/input* .\", folder=folder)\n",
    "\n",
    "# export model\n",
    "#model.export(folder)\n",
    "execute(f\"cp ../1_opes-deeplda/model* .\", folder=folder)\n",
    "\n",
    "# write plumed input\n",
    "with open(folder+\"plumed.dat\",\"w\") as f:\n",
    "    print(\"\"\"\n",
    "# vim:ft=plumed\n",
    "\n",
    "# Compute torsion angles, as well as energy\n",
    "MOLINFO STRUCTURE=input.ala2.pdb\n",
    "phi: TORSION ATOMS=@phi-2\n",
    "psi: TORSION ATOMS=@psi-2\n",
    "theta: TORSION ATOMS=6,5,7,9\n",
    "xi: TORSION ATOMS=16,15,17,19\n",
    "ene: ENERGY\n",
    "\n",
    "# Compute descriptors\n",
    "INCLUDE FILE=../plumed-distances.dat\n",
    "\n",
    "# Compute DeepLDA CV\n",
    "deep: PYTORCH_MODEL FILE=model.ptc    ARG=d_2_5,d_2_6,d_2_7,d_2_9,d_2_11,d_2_15,d_2_16,d_2_17,d_2_19,d_5_6,d_5_7,d_5_9,d_5_11,d_5_15,d_5_16,d_5_17,d_5_19,d_6_7,d_6_9,d_6_11,d_6_15,d_6_16,d_6_17,d_6_19,d_7_9,d_7_11,d_7_15,d_7_16,d_7_17,d_7_19,d_9_11,d_9_15,d_9_16,d_9_17,d_9_19,d_11_15,d_11_16,d_11_17,d_11_19,d_15_16,d_15_17,d_15_19,d_16_17,d_16_19,d_17_19\n",
    "\n",
    "# Stretch the CV\n",
    "deep_mod: CUSTOM ARG=deep.node-0 FUNC=x+x^3 PERIODIC=NO\n",
    "\n",
    "# Apply OPES bias \n",
    "opes: OPES_METAD ARG=deep_mod PACE=500 BARRIER=30\n",
    "\n",
    "# Print \n",
    "PRINT FMT=%g STRIDE=500 FILE=COLVAR ARG=*\n",
    "\n",
    "ENDPLUMED\n",
    "\"\"\",file=f)\n",
    "\n",
    "## RUN GROMACS\n",
    "num_steps=5000000\n",
    "\n",
    "clean(folder)\n",
    "#execute(f\"{GMX_CMD} mdrun -s input.tpr -deffnm alanine -plumed plumed.dat -ntomp 1 -nsteps {num_steps} > alanine.out\", folder=folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "folder = '1_DeepLDA/2_opes-deeplda-stretch/'\n",
    "colvar = load_dataframe(folder+'COLVAR')\n",
    "\n",
    "fig,axs = plt.subplots(1,2,figsize=(10,4),dpi=100)\n",
    "# Time evolution (DeepLDA)\n",
    "colvar.plot.scatter('time','deep_mod',s=1,ax=axs[0])\n",
    "axs[1].set_xlabel('Time [ps]')\n",
    "axs[1].set_xlabel('DeepLDA')\n",
    "# 2D scatter plot colored with DeepLDA\n",
    "colvar.plot.scatter('phi','psi',c='deep_mod',s=1,cmap='fessa',ax=axs[1])\n",
    "axs[1].set_xlabel(r'$\\phi$ [rad]')\n",
    "axs[1].set_ylabel(r'$\\psi$ [rad]')\n",
    "axs[1].set_aspect('equal')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlcvs.utils.fes import compute_fes\n",
    "\n",
    "s = colvar['deep.node-0'].values\n",
    "s2 = colvar['deep_mod'].values\n",
    "\n",
    "# compute weights\n",
    "kbT = 2.5\n",
    "w = np.exp(colvar['opes.bias'].values/kbT)\n",
    "\n",
    "fig,axs = plt.subplots(1,2,figsize=(12,4),dpi=100)\n",
    "ax = axs[0]\n",
    "fes,grid,bounds,error = compute_fes(s, weights=w, kbt=kbT, \n",
    "                                    blocks=5, bandwidth=0.01, \n",
    "                                    plot=True, ax = ax)\n",
    "ax.set_xlabel('DeepLDA')\n",
    "ax.set_ylabel('FES [kJ/mol]')\n",
    "ax.set_ylim(0,50)\n",
    "\n",
    "ax = axs[1]\n",
    "fes,grid,bounds,error = compute_fes(s2, weights=w, kbt=kbT, \n",
    "                                    blocks=5, bandwidth=0.02, \n",
    "                                    plot=True, ax = ax)\n",
    "ax.set_xlabel('DeepLDA mod (x+x^3)')\n",
    "ax.set_ylabel('FES [kJ/mol]')\n",
    "ax.set_ylim(0,50)"
   ]
  }
 ],
 "metadata": {
  "deepnote": {},
  "deepnote_execution_queue": [],
  "deepnote_notebook_id": "43e8f693-9b9d-48d5-9d02-272c9ea44de0",
  "interpreter": {
   "hash": "1cbeac1d7079eaeba64f3210ccac5ee24400128e300a45ae35eee837885b08b3"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 ('pytorch')",
   "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.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}