{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "089a0f58",
   "metadata": {},
   "source": [
    "### Introduction: Running OpenMM simulations\n",
    "OpenMM is a bit different from other MD simulation programs because there is no executable OpenMM program; instead it has an easy to use Python API for setting up and running simulations.  One could think of the OpenMM input file as being a short Python script or Jupyter notebook.\n",
    "\n",
    "#### Resources\n",
    "The [OpenMM Script Builder](http://builder.openmm.org/) is a web app that assists in creating Python scripts for OpenMM simulations.  Once you are done with this notebook and ready to use features like protein force fields, periodic boundary conditions and temperature / pressure controls, you may use the Builder to generate a script with the correct OpenMM API calls.\n",
    "\n",
    "The [OpenMM documentation page](http://openmm.org/documentation.html) is very high-quality. Before reading the docs you're encouraged to watch the two short videos on the landing page given by Peter Eastman, the OpenMM developer.  There are also various OpenMM-related videos that you can find on YouTube from past workshops given at Stanford, including some from Lee-Ping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e4f0931",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import nglview\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "# OpenMM imports\n",
    "from openmm.app import *\n",
    "from openmm import *\n",
    "from openmm.unit import *\n",
    "# Import the sys module because sys.stdout is the file object\n",
    "# that represents printing to the console.\n",
    "import sys"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cb9ee4b",
   "metadata": {},
   "source": [
    "#### Create a `PDBFile` object by passing a `.pdb` file.\n",
    "This object will contain the atomic positions in `pdb.positions` and the topology object in `pdb.topology`. \n",
    "\n",
    "This `.pdb` file contains 27 xenon atoms arranged in a cube; it's more of a demo-sized simulation than a research-sized one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bc408bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1: Create PDB object that contains positions and topology.\n",
    "pdb = PDBFile('xe-27.pdb')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e353b0f",
   "metadata": {},
   "source": [
    "#### Create a `ForceField` object by passing a force field `.xml` file.\n",
    "This object represents the force field definition, including parameters and residue templates. It will contain the method `forcefield.createSystem(topology)` for creating the `System` object given `pdb.topology`.\n",
    "\n",
    "It is possible to read and understand very simple force field `.xml` files so you are encouraged to open the files in the text editor to see what information it contains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24aaa953",
   "metadata": {},
   "outputs": [],
   "source": [
    "forcefield = ForceField('xenon.xml')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57bb5b4a",
   "metadata": {},
   "source": [
    "#### Create a `System` object.\n",
    "This object represents the parameterized system containing all particles and their individual interactions. In our case, the only interactions are the Lennard-Jones interactions between each pair of xenon atoms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4bceb0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "system = forcefield.createSystem(pdb.topology)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78b42192",
   "metadata": {},
   "source": [
    "#### Create an `Integrator` object.\n",
    "\n",
    "This object represents the algorithm for evolving the system in time. The `VerletIntegrator` uses a leapfrog Verlet algorithm that is similar (but not identical) to the velocity Verlet algorithm covered in lecture.\n",
    "\n",
    "Note the use of `picosecond`: this is an OpenMM unit object that we got from importing `simtk.unit` into the global namespace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aa7f82c",
   "metadata": {},
   "outputs": [],
   "source": [
    "integrator = VerletIntegrator(0.001*picosecond)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9496529b",
   "metadata": {},
   "source": [
    "#### Create a `Platform` object.\n",
    "This object represents which set of low-level codes are being used to actually evaluate the interactions. OpenMM contains a simple `Reference` platform, an optimized `CPU` platform, an optimized `CUDA` platform for NVidia GPUs, and an optimized `OpenCL` platform for AMD GPUs. If a `Platform` is not explicitly created and provided to the creation of the `Simulation` object, a default will be used.\n",
    "\n",
    "Because this system is so small, the `Reference` platform in fact gives the best performance, so that's what we'll use here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07aa7ff4",
   "metadata": {},
   "outputs": [],
   "source": [
    "platform = Platform.getPlatformByName('Reference')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db0af576",
   "metadata": {},
   "source": [
    "#### Create a `Simulation` object.\n",
    "This object is the central one that you interact with when running simulations. It is created by passing in the various objects that you created above, and it contains methods for carrying out basic simulation tasks (see diagram below for how various classes are related)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdb6178b",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulation = Simulation(pdb.topology, system, integrator, platform)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "786e0579",
   "metadata": {},
   "source": [
    "![OpenMM diagram](https://i.imgur.com/BI6M5pp.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08944801",
   "metadata": {},
   "source": [
    "#### Set the initial positions.\n",
    "Up to this point we haven't done anything with the positions.  The positions from the `PDBFile` must be explicitly passed to `Simulation` before any actual simulating can take place.\n",
    "\n",
    "Note that the method is actually contained in `simulation.context`.  The `Context` object contains most of the functions of `Simulation`; technically speaking  `Context` is a C++ API layer object, and `Simulation` is defined in the Python layer and adds a few convenience functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b6a08ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulation.context.setPositions(pdb.positions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98adeb11",
   "metadata": {},
   "source": [
    "#### Compute the energy.\n",
    "To compute the energy in OpenMM, we use `simulation.context.getState()`. For performance reasons, this function only retrieves properties that the user asks for, so you should look at the API for how to retrieve other properties such as the positions and velocities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8ae89e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "state = simulation.context.getState(getEnergy=True)\n",
    "print(state.getPotentialEnergy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c076bd90",
   "metadata": {},
   "source": [
    "#### Create Reporter objects and add them to the Simulation.\n",
    "Reporter objects periodically return data as the simulation is running. `StateDataReporter` prints current values of variables to a file or the console, and `PDBReporter` writes the current positions to a trajectory `.pdb` file.\n",
    "\n",
    "Reporter objects must be added to the list `simulation.reporters` after being created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1a3aa35",
   "metadata": {},
   "outputs": [],
   "source": [
    "reporter1 = StateDataReporter(sys.stdout, 100, step=True, potentialEnergy=True, kineticEnergy=True, totalEnergy=True)\n",
    "reporter2 = PDBReporter('xenon-output.pdb', 100)\n",
    "simulation.reporters.append(reporter1)\n",
    "simulation.reporters.append(reporter2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "754a0041",
   "metadata": {},
   "source": [
    "#### Run the MD simulation.\n",
    "Running the simulation itself is as easy as calling `Simulation.step(number_of_steps)`. Here, one million steps are carried out using a time step of 0.001 ps (defined above) to give a total of 1 ns of simulation time.  \n",
    "\n",
    "Once the simulation is started by calling this method, it will print information to the console regarding current values of the simulation time, potential / kinetic / total energy, and a .pdb file will be written to the current folder that you will be able to view in VMD.\n",
    "\n",
    "It is also possible to use the `nglview` package to view trajectories but the `pytraj` package may be required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e894b09",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulation.step(1000000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d9f7cd4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}