{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Equilibrating simulations\n",
    "\n",
    "In this tutorial we will explain what equilibration of a simulation is and how MDMC\n",
    "can ensure enough equilibration has been done before running the simulation."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is equilibration?\n",
    "Equilibration of a simulation involves running the simulation for a short period\n",
    "of time to make it ready for a 'production run' (i.e. one in which we record trajectories).\n",
    "\n",
    "When the simulation is first set up, it has usually been configured in a certain way\n",
    "such as having atoms spaced in a grid at certain densities. We don't want this to affect\n",
    "our trajectories, and we also want properties like kinetic and potential energies to be\n",
    "distributed around the simulation so that it better reflects the 'real-life' behaviour\n",
    "of molecules.\n",
    "\n",
    "First, let us create and minimize an MDMC `Simulation`. Minimization, like equilibration, \n",
    "is a short run of the MD simulation which resolves issues like atoms overlapping\n",
    "(which would affect energy readings for the simulation)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from MDMC.MD import Universe, Simulation, Atom, Dispersion, LennardJones\n",
    "\n",
    "# create the topology\n",
    "universe = Universe(dimensions=23.0668)\n",
    "Ar = Atom('Ar', charge=0., mass=36.0)\n",
    "universe.fill(Ar, num_density=0.0176)\n",
    "\n",
    "# define intermolecular forces\n",
    "Ar_dispersion = Dispersion(universe,\n",
    "                           (Ar.atom_type, Ar.atom_type),\n",
    "                           cutoff=8.,\n",
    "                           function=LennardJones(epsilon=1.0243, sigma=3.36))\n",
    "\n",
    "# MD Engine setup\n",
    "simulation = Simulation(universe,\n",
    "                        engine=\"lammps\",\n",
    "                        time_step=10.18893,\n",
    "                        temperature=120.,\n",
    "                        traj_step=15)\n",
    "\n",
    "simulation.minimize(n_steps=50)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then equilibrate the simulation. This is done by calling `simulation.run` with `equilibration=True`.\n",
    "Normally, to run for, say, 5000 steps, call `simulation.run(n_steps=5000, equilibration=True)`, but here\n",
    "we will put it in a `for` loop that will record how the temperature and potential energy of the simulation\n",
    "changes over time and run the equilibration in smaller chunks.\n",
    "\n",
    "Note that `engine.eval()` evaluates variables in the simulation's MD engine; `temp` and `pe` are\n",
    "the LAMMPS names for temperature and potential energy. It may vary in other MD engines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "temperatures = []\n",
    "potentials = []\n",
    "\n",
    "for _ in range(500):\n",
    "    simulation.run(n_steps=10, equilibration=True)\n",
    "    temperatures.append(simulation.engine.eval('temp'))\n",
    "    potentials.append(simulation.engine.eval('pe'))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This can then be plotted as graphs of how these properties change over time\n",
    "with matplotlib. We create a plot function so we can use it later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "def plot_vars(temperatures, potentials, line=0):\n",
    "    plt.figure()\n",
    "    plt.subplot(211)\n",
    "    plt.plot(np.arange(len(temperatures))*10, temperatures)\n",
    "    plt.ylabel('Temperature (K)')\n",
    "    if line:\n",
    "        plt.axvline(x=line, linestyle='--', color='lime')\n",
    "\n",
    "    plt.subplot(212)\n",
    "    plt.plot(np.arange(len(potentials))*10, potentials, color=\"orange\")\n",
    "    plt.xlabel('time (steps)')\n",
    "    plt.ylabel('Potential energy (kcal/mol)')\n",
    "    if line:\n",
    "        plt.axvline(x=line, linestyle='--', color='lime')\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "plot_vars(temperatures, potentials)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that at first there is a lot of change, and around 2000 steps both quantities cease to \n",
    "increase or decrease in any meaningful way; there is, of course, slight oscillation and uncertainty\n",
    "in the readings (which is unavoidable in molecular dynamics simulation), but ultimately it is\n",
    "stationary. Once this stationary behaviour has been reached, we would consider the simulation\n",
    "to have equilibrated successfully."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Avoiding under- or over- equilibrating"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Indeed, the next question would be 'how much equilibration is necessary?'. If we under-equilibrate,\n",
    "then it will affect our observed trajectory when we run the simulation. If we over-equilibrate, then\n",
    "we can waste time. How can we detect, on the fly, when our system is equilibrated?\n",
    "\n",
    "MDMC features 'auto-equilibration', which analyses how variables change as equilibration progresses,\n",
    "and automatically halts the equilibration when it has determined these variables to be stationary.\n",
    "User-defined parameters can determine the sensitivity of this, as well as what is analysed.\n",
    "\n",
    "To auto-equilibrate, create a simulation and then run `simulation.auto_equilibrate()`. It returns\n",
    "the number of steps it used to equilibrate as well as data for the tracked variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulation = Simulation(universe,\n",
    "                        engine=\"lammps\",\n",
    "                        time_step=10.18893,\n",
    "                        temperature=120.,\n",
    "                        traj_step=15)\n",
    "\n",
    "simulation.minimize(n_steps=50)\n",
    "total_steps, vals_dict = simulation.auto_equilibrate()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before, we can graph these variables. Note that the graph seems a lot less 'noisy';\n",
    "this is purely due to the different x-scale as we finish much sooner!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_temps = vals_dict['temp']\n",
    "auto_potentials = vals_dict['pe']\n",
    "\n",
    "plot_vars(auto_temps, auto_potentials)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally, we can see on the original equilibration graphs where the auto-equilibration\n",
    "considered the graph to be stationary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_vars(temperatures, potentials, line=total_steps)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}