Basic analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CCP-BioSim Training Course:\n",
    "### Basic MD data analysis with MDTraj\n",
    "\n",
    "IPython (Jupyter) notebooks are a great environment in which to do exploratory analysis of MD data. This notebook illustrates how two key Python packages - *numpy* and *matplotlib* - can be used with the MD analysis package [*mdtraj*](http://mdtraj.org/1.9.0/) for some simple standard analysis tasks.\n",
    "\n",
    "I am assuming you have a basic knowledge of Python and Jupyter notebooks; if anything is unfamiliar a quick web search should enable you to fill in the gaps.\n",
    "\n",
    "Along with this notebook you should have downloaded two data files:\n",
    "\n",
    "1. **bpti.gro** : Structure file for the protein BPTI in Gromacs format.\n",
    "\n",
    "2. **bpti.xtc** : Gromacs trajectory file for BPTI.\n",
    "\n",
    "-----\n",
    "To begin with, run the code in the following cell to check that you have all the data files and Python packages required:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os.path as op\n",
    "problems = False\n",
    "top_file = 'data/bpti.gro'\n",
    "traj_file = 'data/bpti.xtc'\n",
    "try:\n",
    "    import mdtraj as mdt\n",
    "except ImportError:\n",
    "    print('Error: you don\\'t seem to have mdtraj installed - use pip or similar to get it then try again.')\n",
    "    problems = True\n",
    "try:\n",
    "    import numpy as np\n",
    "except ImportError:\n",
    "    print('Error: you don\\'t seem to have numpy installed - use pip or similar to get it then try again.')\n",
    "    problems = True\n",
    "try:\n",
    "    import matplotlib.pyplot as plt\n",
    "    # This next line makes matplotlib show its graphs within the notebook.\n",
    "    %matplotlib inline \n",
    "except ImportError:\n",
    "    print('Error: you don\\'t seem to have matplotlib installed - use pip or similar to get it then try again.')\n",
    "    problems = True\n",
    "try:\n",
    "    import nglview as nv\n",
    "except ImportError:\n",
    "    print('Error: you don\\'t seem to have nglview installed - use pip or similar to get it then try again.')\n",
    "    problems = True\n",
    "\n",
    "if not op.exists(top_file):\n",
    "    print('Error: you don\\'t seem to have the data file {} in this directory.'.format(top_file))\n",
    "    problems = True\n",
    "if not op.exists(traj_file):\n",
    "    print('Error: you don\\'t seem to have the data file {} in this directory.'.format(traj_file))\n",
    "    problems = True\n",
    "\n",
    "if problems:\n",
    "    print('Fix the errors above then re-run this cell')\n",
    "else:\n",
    "    print('Success! You seem to have everything ready to run the notebook.')\n",
    "plt.rcParams.update({'font.size': 15}) #This sets a better default label size for plots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparation: loading the trajectory data.\n",
    "\n",
    "Success? If so let's start by loading the MD data.\n",
    "\n",
    "In the cell below, we create an MDTraj *Trajectory* object from the trajectory and topology files, then use the nglviewer app to have a quick look at the trajectory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traj = mdt.load(traj_file, top=top_file)\n",
    "import nglview as nv\n",
    "view = nv.show_mdtraj(traj)\n",
    "view"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having got a feel for the molecule we are dealing with, and its dynamics, let's look at the Python representation of it.\n",
    "\n",
    "Let's have a little dig around the trajectory object *traj*: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(traj)\n",
    "print(type(traj.xyz))\n",
    "print(traj.xyz.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So you can see that the trajectory coordinate data is stored in the numpy array *traj.xyz*. This is a three-dimensional array of shape [n_frames, n_atoms, 3]. \n",
    "\n",
    "**NB:** MDTraj stores coordinates in nanometer units.\n",
    "\n",
    "---\n",
    "\n",
    "## Example 1: Calculate and plot the end-to-end distance as a function of time.\n",
    "Let's start with something simple. The cell below defines a function that calculates, for each frame, the distance between the first atom and the last. Then we apply the function to the trajectory object, and plot the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def end_to_end(xyz):\n",
    "    '''\n",
    "    Calculate end to end distances\n",
    "    \n",
    "    Arguments:\n",
    "        xyz : numpy array of shape (n_frames, n_atoms, 3)\n",
    "    \n",
    "    Returns:\n",
    "        numpy array of shape (n_frames)\n",
    "    '''\n",
    "    d = np.zeros(len(xyz)) # create an array of zeros of length n_frames\n",
    "    for i in range(len(xyz)):\n",
    "        dxyz = xyz[i, 0] - xyz[i, -1] # note use of '-1' to get the last element of the array (the last atom)\n",
    "        # The next two lines are just Pythagoras:\n",
    "        dxyz2 = dxyz * dxyz\n",
    "        d[i] = np.sqrt(dxyz2.sum()) # If you are not familiar with numpy sum() you might want to look it up\n",
    "    return d\n",
    "\n",
    "e2e = end_to_end(traj.xyz)\n",
    "plt.plot(e2e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Follow-up challenges:\n",
    "\n",
    "Tweak the code in the cell above to:\n",
    "\n",
    "1. Add suitable x- and y-axis labels to the plot, and maybe a title (hint: the matplotlib documentation is [here](https://matplotlib.org)).\n",
    "2. Plot the end-to-end distance in angstroms instead of nanometers.\n",
    "\n",
    "---\n",
    "\n",
    "## Example 2: Plotting RMSD data.\n",
    "### 2a: The easy way.\n",
    "\n",
    "*MDTraj* includes many useful trajectory analysis tools. Here we use the *rmsd* function to calculate the RMSD of each trajectory frame from the first, and then we plot it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmsd = mdt.rmsd(traj, traj[0])\n",
    "plt.plot(rmsd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2b: The harder way.\n",
    "Let's try using *numpy* to calculate the same thing. \n",
    "\n",
    "If you are quite new to *numpy* you may want to research what the *axis* argument to the *sum()* function does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def myrmsd(xyz):\n",
    "    '''\n",
    "    Calculate rmsd for a trajectory\n",
    "    \n",
    "    Arguments:\n",
    "        xyz: numpy array of shape (n_frames, n_atoms, 3)\n",
    "        \n",
    "    Returns:\n",
    "        numpy array of shape (n_frames)\n",
    "    '''\n",
    "    # calculate all displacements relative to the first frame:\n",
    "    dxyz = xyz - xyz[0]\n",
    "    # square them\n",
    "    dxyz2 = dxyz * dxyz\n",
    "    # Pythagoras: r2 = x2 + y2 + z2\n",
    "    dr2 = dxyz2.sum(axis=2)\n",
    "    # square root of mean value of r2 for each snapshot:\n",
    "    rmsd = np.sqrt(dr2.mean(axis=1))\n",
    "    return rmsd\n",
    "\n",
    "rmsd2 = myrmsd(traj.xyz)\n",
    "plt.plot(rmsd2) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Whoops! That doesn't look right, does it? This is because it is standard practice when calculating RMSDs to least-squares fit the two snapshots together first, to remove translation and rotation. *MDTraj* did this, but the simple *numpy* code did not, so the RMSD values are much bigger. \n",
    "\n",
    "There is no simple way to fix this with *numpy*, but with *MDTraj* it is easy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "traj.superpose(traj[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the trajectory data is all least-squares fitted to the first snapshot (\"traj[0]\"), we can try the simple *numpy* function again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rmsd3 = myrmsd(traj.xyz)\n",
    "plt.plot(rmsd3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Follow-up challenge:\n",
    "Well it looks like the *MDTraj* result, doesn't it? But to check, write some code in the cell below to calculate and plot the difference between the RMSDs calculated using the two approaches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write your code here:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Example 3: Calculating RMSFs.\n",
    "### 3a: The simple approach\n",
    "While *MDTraj* has a simple function to calculate RMSDs, there is no equivalent for RMSFs (root-mean-square fluctuations) (Yes, there is an approach if you dig around a bit, but it's a bit clunky, IMHO). But with a bit of *numpy* its easy to create our own:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def myrmsf(xyz):\n",
    "    '''\n",
    "    Calculate RMSF from trajectory data\n",
    "    \n",
    "    Arguments:\n",
    "        xyz : numpy array of shape (n_frames, n_atoms, 3)\n",
    "    \n",
    "    Returns:\n",
    "        numpy array of shape (n_atoms)\n",
    "    '''\n",
    "    # calculate all displacements relative to the mean structure:\n",
    "    dxyz = xyz - xyz.mean(axis=0)\n",
    "    # square them\n",
    "    dxyz2 = dxyz * dxyz\n",
    "    # Pythagoras: r2 = x2 + y2 + z2\n",
    "    dr2 = dxyz2.sum(axis=2)\n",
    "    # Square root of mean value of r2 for each atom, over all frames:\n",
    "    rmsf = np.sqrt(dr2.mean(axis=0))\n",
    "    return rmsf\n",
    "\n",
    "rmsf = myrmsf(traj.xyz)\n",
    "plt.plot(rmsf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Follow-up challenge:\n",
    "\n",
    "Add suitable axis labels to the plot.\n",
    "\n",
    "### 3b: A more advanced approach.\n",
    "Instead of plotting the RMSF for each atom, how about the average RMSF for each residue? To do this we need to:\n",
    "\n",
    "1. Find the index of the first atom in each residue\n",
    "2. Average the atomic RMSFs from this atom to one less than the first atom of the next residue, or the last atom if it is the last residue.\n",
    "\n",
    "To do the first you will see a bit of slightly more advanced wrangling with *MDTraj*, then the second is pretty straightforward:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_atoms = [r.atom(0).index for r in traj.topology.residues] # You may want to do a bit of research on this line.\n",
    "def residue_rmsf(rmsf, first_atoms):\n",
    "    '''\n",
    "    Convert atom-resolution RMSFs into residue-resolution values\n",
    "    \n",
    "    Arguments:\n",
    "        rmsf : rmsf data as numpy array of shape (n_atoms)\n",
    "        first_atoms : indices of first atoms in each residue as numpy array of shape (n_residues)\n",
    "        \n",
    "    Returns:\n",
    "        numpy array of rmsf values of shape (n_residues)\n",
    "    '''\n",
    "    n_res = len(first_atoms)\n",
    "    result = np.zeros(n_res)\n",
    "    for i in range(n_res - 1):\n",
    "        result[i] = rmsf[first_atoms[i]:first_atoms[i+1]].mean()\n",
    "    result[n_res - 1] = rmsf[first_atoms[n_res -1]:].mean()\n",
    "    return result\n",
    "\n",
    "res_rmsf = residue_rmsf(rmsf, first_atoms)\n",
    "plt.plot(res_rmsf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Follow-up challenge:\n",
    "\n",
    "You can get a numpy array with the indices of all the backbone (N, CA, C and O) atoms using the command:\n",
    "\n",
    "```\n",
    "backbone_list = traj.topology.select('backbone')\n",
    "```\n",
    "\n",
    "In the cell below write some code to calculate and plot the residue-based RMSF values for backbone atoms only:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write your code here:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Example 4: Secondary structure analysis.\n",
    "*MDTraj* includes the DSSP algorithm to assign secondary structure. In its simplest form, this returns an [n_frames, n_residues] array where in each frame, each residue is labelled as 'H' (helical), 'E' (beta-strand), or 'C' (random coil):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ss = mdt.compute_dssp(traj)\n",
    "print (ss.shape)\n",
    "print (ss[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot the number of helical residues in each snapshot. If you are fairly new to *numpy* then note that when the sum() function is applied to an array of logical (Boolean) values, it counts the number of True elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_helical = (ss == 'H').sum(axis=1)\n",
    "plt.plot(n_helical)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the cell below, we get a little cleverer and calculate the percentage native structure for each snapshot (assuming 'native' means 'same as the first snapshot'):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot((ss == ss[0]).sum(axis=1) * 100 / 58)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Follow-up challenge:\n",
    "This graph might look nicer if the data was smoothed over, for example, a running window of 5 points. In the cell below create a function to do this and plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write your code here:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Hopefully this brief introduction has convinced you that iPython/Jupyter notebooks are a very useful approach to MD data analysis. We have only scratched the surface of what *MDTraj* can do, and not even begun to look at how, with a bit of extra work, you can get publication-quality graphs out of *matplotlib*."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}