Session3-Atomistic_MD/SPC_E_Water.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Atomistic Molecular Dynamics: Bulk Liquids - SPC/E Water\n",
    "\n",
    "Today's workshop session will involve performing atomistic molecular dynamics of bulk liquid systems. The first system we will consider is a bulk water system where we will use the extended simple point charge model (SPC/E)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build the system\n",
    "\n",
    "As in our prior workshops, we will use mBuild to construct our system. Here, we will import mBuild and define a class for a single water molecule, which we load from a PDB structure file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mbuild as mb\n",
    "\n",
    "class H2O(mb.Compound):\n",
    "    def __init__(self):\n",
    "        super(H2O, self).__init__()\n",
    "        mb.load('utils/h2o.pdb', compound=self)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before putting together the bulk system, let's take a look at an instance of our water molecule class to make sure things look reasonable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "water = H2O()\n",
    "water.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### SPC/E water model\n",
    "\n",
    "Because we are going to be using the SPC/E water model, we need only define three particles (for more complex water models additional sites are often used). Below is a quick overview of this model that shows a diagram of a water molecule along with the parameters for the model.\n",
    "\n",
    "<img src=\"utils/spce-diagram.png\" alt=\"SPCE Diagram\" width=\"300px\"/>\n",
    "<img src=\"utils/spce-table.png\" alt=\"SPCE Diagram\" width=\"200px\"/>\n",
    "(images from http://www.sklogwiki.org/SklogWiki/index.php/SPC/E_model_of_water)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we've verified our `H2O` class is good, we'll define a class for a box of water. We'll set this up so that we can toggle the number of water molecules via a `n_molecules` argument.\n",
    "\n",
    "The recipe for this class is relatively straightfoward. We create an instance of the `H2O` class and use the `fill_box` function to fill a box with copies of this molecule at a density of 1000 kg/m^3. We also rename each molecule copy to `Water` as this speeds up the atom-typing process (whereby only a single copy will need to be atom-typed and those types can then be mapped onto all subsequent copies)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WaterBox(mb.Compound):\n",
    "    \"\"\"An box of water.\"\"\"\n",
    "    def __init__(self, n_molecules):\n",
    "        \"\"\"Initialize an WaterBox Compound.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_molecules : int\n",
    "            Number of molecules to place in the box\n",
    "        \"\"\"\n",
    "        super(WaterBox, self).__init__()\n",
    "        \n",
    "        # Create a prototype for water using the class above\n",
    "        water = H2O()\n",
    "\n",
    "        # Fill a box with water at 1000 kg/m^3\n",
    "        box_of_waters = mb.fill_box(compound=water, n_compounds=n_molecules,\n",
    "                                    density=1000)\n",
    "        \n",
    "        # Rename all molecules to `water`, this speeds up the atom-typing process\n",
    "        for child in box_of_waters.children:\n",
    "            child.name = 'Water'\n",
    "        self.add(box_of_waters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a water box with 500 molecules and visualize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "water_box = WaterBox(500)\n",
    "water_box.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to run the simulations today in GROMACS. Like HOOMD, GROMACS is an open-source molecular simulation code, and, like HOOMD, it is optimized for GPUs. GROMACS does not have the flexibility of HOOMD (only a few potential styles are supported for force field terms), but GROMACS has a dedicated team of full-time developers that have optimized the code to yield perhaps the best performance of any freely available molecular dynamics code.\n",
    "\n",
    "GROMACS expects two data formats when setting up a simulation:\n",
    "  - **GRO** format specifies particle coordinates and box information\n",
    "  - **TOP** format specifies the topology and force field information\n",
    "\n",
    "mBuild (via the ParmEd package) can write to both of these formats. First, we'll output our water box to GRO format. Because we are only writing positions and box information, we do not have to atom-type the system just yet. However, we will use the `residues` argument. Although originating from the protein community, we can think of residues as molecules. This does not really provide much advantage in writing the GRO file, but it will when we save our topology."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "water_box.save('spce.gro', overwrite=True, residues='Water')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a quick look at some parts of the GRO file we've just written."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!head spce.gro"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Following a header line, we see that the GRO format specifies the number of particles in the system, and then contains a line for each particle specifying the residue name, the atom name (in our case the element), the atom index (starting from 1...), and the XYZ coordinates (note that GROMACS uses nanometers as its distance unit)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tail spce.gro"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At the end of the file we see three values: `2.46388    2.46388    2.46388`. These represent our three box lengths in nanometers.\n",
    "\n",
    "Okay, now that we've checked out the GRO file, let's write our TOP file. This is where we need to define a force field that can be applied to our system. Since we're using the SPC/E water model, we will use a force field that corresponds to these model parameters. Let's check this out real quick:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat utils/spce.xml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, now let's write our TOP file, again providing the `residues` argument, and now also providing a `forcefield_files` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "water_box.save('spce.top', forcefield_files='utils/spce.xml', overwrite=True, residues='Water')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we're ready to run a simulation! We'll split our simulation up into two stages: an equilibration stage and a production stage. Both of which will be performed under the NPT ensemble."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### GROMACS MDP files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unlike HOOMD, whose simulation script is in Python, GROMACS has its own format for defining simulation parameters. Specifically, these are housed within another file, called an MDP file (Molecular Dynamics Parameters). There are numerous commands available for you to fine-tune how you would like to run your simulation (see the complete list [here](http://manual.gromacs.org/online/mdp_opt.html)). Let's take a look at the MDP file we're going to be using for equilibration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat utils/npt-water.mdp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll walk through section-by-section to see what's going on. You should find many similarities between the sections in the GROMACS MDP file and the HOOMD run script. As you become more familiar with various simulation codes, you'll find that they all share many of the same qualities, each with its own unique spin on how simulations should be defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sed -n '1,14 p' utils/npt-water.mdp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the section printed above, we define various control parameters for our run. We tell GROMACS which integrator to use (here we choose leap-frog), as well as the timestep and the number of timesteps to perform.\n",
    "\n",
    "We additionally tell GROMACS to perform removal of center-of-mass translation at a defined frequency. This is important as rounding errors can lead to build-up of kinetic energy of the center-of-mass, which can lead to [undesired results](https://en.wikipedia.org/wiki/Flying_ice_cube)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sed -n '17,24 p' utils/npt-water.mdp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above section, we can tell GROMACS how often certain quantities should be written to a file. Here, we tell GROMACS to generate three files, a log file (which will contain run information as well as various quantities like temperature and energies), an energy file (a binary file with various energies and other measured quantities), and a coordinate file. Note that the coordinate file we are producing is designed to be lightweight via lossy compression. Depending on what quantities you are measuring, this may or may not be desirable for analysis. Large, more precise, files can also be generated if desired."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sed -n '27,43 p' utils/npt-water.mdp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above code, we give GROMACS information on the neighborlist that should be created. We can specify the neighborlist type, the update frequency, and the skin distance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sed -n '46,59 p' utils/npt-water.mdp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we've told GROMACS to use the particle-mesh Ewald algorithm for computing long-range electrostatics. We also provide the cutoffs for VDW and short-range Coulombic interactions, and tell Gromacs to shift the VDW potential to zero at the cutoff."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sed -n '62,90 p' utils/npt-water.mdp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we provide information as to any thermostats or barostats that should be employed. Because we want to run under the NPT ensemble, we need to specify both a thermostat and a barostat, as well as coupling constants (which dictate how aggressively the temperature/pressure is controlled) and target values.\n",
    "\n",
    "Here, we will be using the Berendsen barostat and v-rescale thermostat. It should be noted that the Berendsen algorithm does not reproduce the desired ensemble, and thus should normally not be used for production simulations. It is a useful algorithm for equilibrating systems, however. For our simulation here, we will use the Berendsen barostat for the entire simulation (as we are running only for a short period of time for demonstration purposes), but do note that if we were to gather data we wanted to publish, we would want to switch to a better barostat (such as Parinello-Rahman) after equilibrating."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sed -n '93,96 p' utils/npt-water.mdp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we tell GROMACS to generate velocities for our system upon startup, based on a desired temperature and random seed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### GROMPP-ing\n",
    "\n",
    "[Documentation](http://manual.gromacs.org/documentation/2018/onlinehelp/gmx-grompp.html)\n",
    "\n",
    "GROMACS simulations are executed in two stages, each using a separate GROMACS utility. The first utility is called `grompp`, which stands for GROMACS preprocessor. From the GROMACS documentation, `grompp` \"reads a molecular topology file, checks the validity of the file, and expands the topology from a molecular description to an atomic description\". The resulting file, a TPR file, a portable binary input file that can be used to perform a simulation.\n",
    "\n",
    "The `grompp` utility is useful both as a result of the portable file it produces (that can be stored to easily reproduce a given simulation) and because of the variety of validity checks it performs. Common issues with simulation input that `grompp` will find include:\n",
    "  - Attempting to use a timestep that is too large\n",
    "  - Attempting to use certain barostats/thermostats that are not designed for equilibration\n",
    "  \n",
    "Because the result from `grompp` is a single file (the TPR file) that contains all of the information necessary to perform a simulation, this utility requires several files as inputs. These are the MDP parameter file, the GRO coordinate file, and the TOP topology file. We have already generated each of these files, so we are ready to use `grompp` to prepare our system for simulation.\n",
    "\n",
    "Execute the cell below to run `grompp`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gmx grompp -v -f utils/npt-water.mdp -c spce.gro -p spce.top -o npt-water"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### MDrun\n",
    "\n",
    "[Documentation](http://manual.gromacs.org/documentation/5.1/onlinehelp/gmx-mdrun.html)\n",
    "\n",
    "The second utility required to perform a GROMACS simulation is `mdrun`. From the documentation, `mdrun` \"is the main computational chemistry engine within GROMACS.\" It takes the TPR file from `grompp` and actually runs the simulation. Arguments can be provided to control specifications such as the number of MPI process or OpenMP threads utilized. By default GROMACS will attempt to guess the optimal parameters based on your hardware.\n",
    "\n",
    "When executing `mdrun` GROMACS begins each simulation with a short optimization period, whereby various aspects of the domain decomposition and long-range electrostatics solver are optimized for maximum throughput.\n",
    "\n",
    "Execute the cell below to run `mdrun` to perform our simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!gmx mdrun -v -s npt-water.tpr -o -x -deffnm npt-water"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's view the trajectory!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nglview as nv\n",
    "import mdtraj as md\n",
    "\n",
    "water_box.save('spce.mol2', overwrite=True)\n",
    "t = md.load(\"npt-water.xtc\", top=\"spce.mol2\")\n",
    "nv.show_mdtraj(t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to `grompp` and `mdrun`, GROMACS contains many other built-in utilities, most of these for analysis. Here, we will use the `gmx energy` utility to extract output from our simulation to a file in column format that we can read with Numpy and plot.\n",
    "\n",
    "`gmx energy` is interactive and queries the user as to which quantities should be extracted. Each quantity is assigned a number. Here, we will combine all of this into a single command and tell GROMACS to extract quantities 6, 7, 8, 12, and 13 which correspond to the total energy, temperature, pressure, volume, and density respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!echo 6 8 9 13 14\\\\n0 | gmx energy -f npt-water.edr -o npt-water.xvg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll take the extracted quantities from `gmx energy` and plot them with `matplotlib`. Execute the cell below and examine the plots that are created. What trends do you notice? Does our measured density match the expected value of 1000 kg/m^3?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "data = np.loadtxt('npt-water.xvg', skiprows=27)\n",
    "\n",
    "# Strip out the very beginning of the run\n",
    "data = data[10:]\n",
    "\n",
    "fig, ax = plt.subplots(5, 1)\n",
    "\n",
    "properties = ['Total energy, kJ/mol', 'Temperature, K', 'Pressure, atm', 'Volume, nm^3', 'Density, kg/m^3']\n",
    "\n",
    "for i, sub_ax in enumerate(ax):\n",
    "    sub_ax.plot(data[:, 0], data[:, i + 1])\n",
    "    sub_ax.set_title(properties[i])\n",
    "    sub_ax.set_ylabel(properties[i])\n",
    "    \n",
    "fig.subplots_adjust(hspace=1.0)\n",
    "fig.set_size_inches(4, 12)\n",
    "plt.show()"
   ]
  }
 ],
 "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.5.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}