05_analysis/03_alignment_principal_component_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Alignment and principal component analysis (PCA)\n",
    "\n",
    "This notebook starts with alignment of a protein trajectory to a reference geometry, by minimizing the RMSD. This is needed to remove (random) global rotation and translation from the protein, before analyzing the principal components that describe the global motions of the protein.\n",
    "\n",
    "This notebook shows how to perform PCA with bare Numpy operations, to let you better understand each step. The MDTraj documentation also explains how to perform the [PCA analysis with scikit-learn](http://mdtraj.org/latest/examples/pca.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import mdtraj\n",
    "import nglview\n",
    "import numpy as np\n",
    "import pandas\n",
    "from openmm import *\n",
    "from openmm.app import *\n",
    "from openmm.unit import *\n",
    "from pdbfixer import PDBFixer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next cell loads the trajectory and strips water molecules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traj = mdtraj.load(\"traj.dcd\", top=\"init.pdb\")\n",
    "traj.restrict_atoms(traj.topology.select(\"protein\"))\n",
    "df = pandas.read_csv(\"scalars.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. RMSD Alignment\n",
    "\n",
    "We will use the [superpose](http://mdtraj.org/latest/api/generated/mdtraj.Trajectory.superpose.html) method from MDTraj to remove the global rotation and translation of the protein. All frames will be aligned with the first geometry of the trajectory, by setting the second argument to zero:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traj.superpose(traj, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This was relatively easy. We can also (optionally) compute the actual RMSD value at each time step, with the [rmsd](http://mdtraj.org/latest/api/generated/mdtraj.rmsd.html) function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmsds = mdtraj.rmsd(traj, traj, 0)\n",
    "plt.close(\"rmsd\")  # This is needed to rerun the code cell correctly\n",
    "fig, ax = plt.subplots(num=\"rmsd\")\n",
    "ax.plot(df[\"Time (ps)\"], rmsds)\n",
    "ax.set_xlabel(\"Time [ps]\")\n",
    "ax.set_ylabel(\"RMSD [nm]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The rapid initial increase of the RMSD represents a sudden change in geometry caused by the energy minimization before the actual MD run. Over time, the protein geometry slowly drifts away from its initial structure. In the following visualization, the effect of the alignment is clear: the geometry does not slowly rotate as in the previous notebook. Instead it just seems to wiggle in place."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "view = nglview.show_mdtraj(traj)\n",
    "view.add_representation(\"ball+stick\", selection=\"protein\")\n",
    "view"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. The covariance matrix\n",
    "\n",
    "For the remaining part of the analysis, the initialization phase will first be removed from the trajectory.\n",
    "\n",
    "The first step in a principal component analysis is to compute the covariance matrix of the atomic positions, for which we will use NumPy's [cov](https://docs.scipy.org/doc/numpy/reference/generated/numpy.cov.html) function. In this notebook, all atoms are included, but one may also work on a subset, e.g. the alpha carbons, instead.\n",
    "\n",
    "Because this is very short MD simulation of a rather rigid protein, a direct visualization of the covariance does not show much."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We first need to change the shape of the array with atomic positions.\n",
    "print(traj.xyz.shape)\n",
    "xyz = traj.xyz.reshape((-1, traj.n_atoms * 3))\n",
    "print(xyz.shape)\n",
    "\n",
    "# rowvar=True indicates that every row is one time step (or frame).\n",
    "# The first 250 steps are skipped because they cover the initialization phase.\n",
    "covar = np.cov(xyz[250:], rowvar=False)\n",
    "print(covar.shape)\n",
    "\n",
    "# Select the matrix elements corresponding to carbon atoms\n",
    "selection = traj.top.select(\"name CA\")\n",
    "\n",
    "# Use the selection to slice the covariance matrix\n",
    "selection_xyz = np.array([3 * selection, 3 * selection + 1, 3 * selection + 2]).T.ravel()\n",
    "sub_covar = covar[selection_xyz][:, selection_xyz]\n",
    "\n",
    "# Show the covariance matrix\n",
    "plt.close(\"covar\")\n",
    "fig, ax = plt.subplots(num=\"covar\")\n",
    "absvmax = abs(sub_covar).max() / 2\n",
    "m = ax.matshow(sub_covar, vmin=-absvmax, vmax=absvmax, cmap=\"PiYG\")\n",
    "ax.set_xlabel(\"Cartesian coordinate index i\")\n",
    "ax.set_ylabel(\"Cartesian coordinate index j\")\n",
    "plt.colorbar(m, label=\"Covariance\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The diagonal of the covariance matrix contains the variances of the atomic positions. When taking the square root of the diagonal, one obtains the so-called root-mean-square fluctuation (RMSF) of the atomic positions. A high RMSF corresponds to more flexible atoms. To make the RMSF plot legible, the following example only visualizes the RMSF for the alpha carbons atoms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The diagonal elements corresponding to X, Y and Z coordinates\n",
    "# of each atom are summed, such that one value part alpha\n",
    "# carbon is retained.\n",
    "rmsf = np.sqrt(np.diag(covar)[selection_xyz].reshape(-1, 3).sum(axis=1))\n",
    "plt.close(\"rmsf\")  # This is needed to rerun the code cell correctly\n",
    "fig, ax = plt.subplots(num=\"rmsf\")\n",
    "ax.plot(rmsf)\n",
    "ax.set_xlabel(\"Alpha carbon index\")\n",
    "ax.set_ylabel(\"RMSF\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A common result for this plot is that the termini are more flexible than the central part of the protein chain. Also flexible loops can be identified with this analysis. (Keep in mind that this example is based on a relatively short MD simulation, such that the results may not yet be converged.)\n",
    "\n",
    "One can also compute the RMSF directly, without first computing the whole covariance matrix and then extracting only the diagonal elements, as shown the following example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmsf_alt = np.sqrt(np.var(xyz[250:], axis=0)[selection_xyz].reshape(-1, 3).sum(axis=1))\n",
    "print(rmsf[:5])\n",
    "print(rmsf_alt[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Diagonalization of the covariance matrix\n",
    "\n",
    "The next step is to obtain the eigenvalues and eigenvectors, for which we use NumPy's [eigh](https://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.eigh.html) function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evals, evecs = np.linalg.eigh(covar)\n",
    "# Print the10 most significant eigenvalues.\n",
    "print(evals[-10:])\n",
    "plt.close(\"eigvals\")  # This is needed to rerun the code cell correctly\n",
    "fig, ax = plt.subplots(num=\"eigvals\")\n",
    "ax.semilogy(evals)\n",
    "ax.set_xlabel(\"Eigenvalue index\")\n",
    "ax.set_ylabel(\"Eigenvalue [nm^2]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first six eigenvalues are practically zero and correspond to the frozen global rotation and translation. The remaining 1746 modes represent internal motions of the protein with increasing magnitude. In most modes, there is also no motion, which is partially caused by the short simulation time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Projection on a PCA mode\n",
    "\n",
    "Once we have obtained the PCA modes (the eigenvectors), we can project the trajectory on one mode and follow the displacement along each mode. The following cell plots the projection on the first 5 modes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.close(\"projpca\")  # This is needed to rerun the code cell correctly\n",
    "fig, ax = plt.subplots(num=\"projpca\")\n",
    "\n",
    "# Compute the average geometry\n",
    "av = xyz.mean(axis=0)\n",
    "# Plot projections on the three-most significant principal modes.\n",
    "for i in range(3):\n",
    "    proj = np.dot(xyz - av.reshape(1, -1), evecs[:, -i - 1])\n",
    "    ax.plot(df[\"Time (ps)\"], proj, label=f\"Proj. {i}\")\n",
    "ax.set_xlabel(\"Time [ps]\")\n",
    "ax.set_ylabel(\"Displacement [nm]\")\n",
    "ax.legend(loc=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The [cosine content](https://doi.org/10.1103/PhysRevE.65.031910) is clearly visible in these projections, showing that the MD simulation is far too short to completely explore the conformational space. So far, the MD simulation essentially performed an undirected random walk in the conformational degrees of freedom of the protein."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Filtering the MD trajectory\n",
    "\n",
    "One may visualize the motion in one or more modes by filtering the atomic positions, only retaining displacements along modes. To do so, one collects the projected modes and reconstructs the trajectory using only a few selected modes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The filtering algorithm\n",
    "xyz_filter = 0\n",
    "# Motions along the three most significant modes are retained.\n",
    "for i in 0, 1, 2:\n",
    "    evec = evecs[:, -i - 1]\n",
    "    proj = np.dot(xyz - av.reshape(1, -1), evec)\n",
    "    xyz_filter += np.outer(proj, evec)\n",
    "xyz_filter += av\n",
    "\n",
    "# Finally, we cast the result back into an MDTraj trajectory and visualize it.\n",
    "traj_filter = mdtraj.Trajectory(xyz_filter.reshape(-1, traj.n_atoms, 3), traj.topology)\n",
    "view_filter = nglview.show_mdtraj(traj_filter)\n",
    "view_filter.add_representation(\"ball+stick\", selection=\"protein\")\n",
    "view_filter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A careful inspect of the filtered trajectory should reveal the following:\n",
    "\n",
    "- The average structure (to which the filtered modes are added) does not always make chemical sense. For example, methyl hydrogens overlap when a methyl group is freely rotating. Conversely, this also means PCA becomes meaningless with extreme conformational changes, because then the average structure is no longer meaningful.\n",
    "\n",
    "- Most of the motion is close to the termini and some side chains, while the majority of the protein remains locked in place."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Extrapolation of a PCA mode\n",
    "\n",
    "One may also visualize individual PCA modes by simply performing a linear extrapolation of the mode relative to the average structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The amplitude is taken from the maximum eigenvalue\n",
    "# of the covariance matrix. This will always have the\n",
    "# right order of magnitude.\n",
    "amplitude = np.sqrt(evals.max())\n",
    "# With evecs[:, -1], the last eigenmode,\n",
    "# i.e. the first principal mode, is selected.\n",
    "xyz_extra = np.outer(np.linspace(-amplitude, amplitude, 100), evecs[:, -1]) + av\n",
    "\n",
    "# Again, we cast the result back into an MDTraj trajectory and visualize it.\n",
    "traj_extra = mdtraj.Trajectory(xyz_extra.reshape(-1, traj.n_atoms, 3), traj.topology)\n",
    "view_extra = nglview.show_mdtraj(traj_extra)\n",
    "view_extra.add_representation(\"ball+stick\", selection=\"protein\")\n",
    "view_extra"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also in this analysis, the individual frames are not necessarily realistic geometries."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Hinging mode in Adenylate Kinase\n",
    "\n",
    "In the above examples, the principal modes were not informative because the MD simulation was too short. The villin headpiece is also not the best example to demonstrate global modes.\n",
    "\n",
    "**<span style=\"color:#A03;font-size:14pt\">\n",
    "&#x270B; HANDS-ON! &#x1F528;\n",
    "</span>**\n",
    "\n",
    "> The protein structures [2RGX](https://www.rcsb.org/structure/2rgx) and [2RH5](https://www.rcsb.org/structure/2rh5) represent two different conformations of one\n",
    "and the same protein: Adenylate Kinase from Aquifex aeolicus. This is a phosphotransferase enzyme that\n",
    "catalyzes the interconversion of adenine nucleotides. Basically, it is able to convert two molecules of\n",
    "adenosinediphosphate (ADP) to one adenosinemonophosphate (AMP) and one adenosinetriphosphate\n",
    "(ATP), and vice versa. As such, it plays an important role in cellular energy homeostasis.\n",
    "2RH5 describes the conformation of this enzyme in the absence of a substrate. 2RGX contains a substrate\n",
    "analogue (bis(adenosine)-5'-pentaphosphate) and hence can be considered to describe the conformation of\n",
    "the enzyme in the presence of a substrate. The main difference between both conformations resembles the\n",
    "action of pac-man: the enzyme opens up two lid-regions in the absence of a ligand, and closes them once a\n",
    "ligand is present.\n",
    ">\n",
    "> ![pacman](pacman.png)\n",
    ">\n",
    "> 1. Load both structures with MDTraj, align one to the other and visualize both on top of each other with nglview.\n",
    ">\n",
    "> 2. Perform 2 nanoseconds of MD simulation using only one of the two states as initial structure. Compute the RMSD as function of time, using the two crystal structures as references. Also perform a PCA on this trajectory to identify and visualize the hinging mode."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When using PDB files from the RCSB Protein Data Bank, it is common that some minor parts (atoms, terminial residues, terminations, ...) are missing. These issues can be fixed with a tool called PDBFixer. General usage instructions can be found on the [PDBFixer](https://htmlpreview.github.io/?https://github.com/openmm/pdbfixer/blob/master/Manual.html) homepage.\n",
    "It is also capable of a few other typical tasks, such as replacing non-standard residues or applying point mutations.\n",
    "\n",
    "The following cell shows how to use PDBFixer in thise case, with a for loop to fix both PDB files with the same code. PDBFixer downloads the PDB files, applies the requested fixed and writes out the corrected PDB. All this will take a few seconds.\n",
    "\n",
    "The fixed PDB files will be suitable for superposition (next code cell) and for starting points for an MD simulations, which you have to implement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for pdbid in \"2rgx\", \"2rh5\":\n",
    "    # Fix the PDB file with the PDB fixer\n",
    "    print(\"Loading and fixing\", pdbid)\n",
    "    # Remove all chains except the first\n",
    "    fixer = PDBFixer(pdbid=pdbid)\n",
    "    nchain = len(list(fixer.topology.chains()))\n",
    "    fixer.removeChains(range(1, nchain))\n",
    "    # Remove everything except water and protein\n",
    "    fixer.removeHeterogens(True)\n",
    "    # Fill in missing atoms\n",
    "    fixer.findMissingResidues()\n",
    "    fixer.findMissingAtoms()\n",
    "    fixer.addMissingAtoms()\n",
    "    fixer.addMissingHydrogens(7.0)\n",
    "    with open(f\"{pdbid}_fixed.pdb\", \"w\") as f:\n",
    "        PDBFile.writeFile(fixer.topology, fixer.positions, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code shows how to load the two PDB files as single-frame trajectories in MDTraj. They are aligned by minimizing the RMSD and then visualized with nglview."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load trajectories\n",
    "traj1 = mdtraj.load(\"2rgx_fixed.pdb\")\n",
    "traj2 = mdtraj.load(\"2rh5_fixed.pdb\")\n",
    "# Superpose\n",
    "traj2_aligned = traj2.superpose(traj1)\n",
    "# View\n",
    "view = nglview.show_mdtraj(traj1)\n",
    "view.add_component(traj2_aligned)\n",
    "view[0].clear_representations()\n",
    "view[1].clear_representations()\n",
    "view[0].add_cartoon(color=\"blue\")\n",
    "view[1].add_cartoon(color=\"red\")\n",
    "view"
   ]
  }
 ],
 "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"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": true,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}