{
 "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 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": [
    "import numpy as np\n",
    "\n",
    "from MDMC.MD import Universe, Simulation, Atom, NonBonded, NonBondedForce\n",
    "from openmm import unit\n",
    "\n",
    "# create the topology\n",
    "universe = Universe(dimensions=38.4441)\n",
    "Ar = Atom('Ar', charge=0., mass=36.0)\n",
    "n_ar_atoms = int(0.0176 * np.prod(universe.dimensions))\n",
    "print(n_ar_atoms)\n",
    "universe.fill(Ar, num_struc_units=(n_ar_atoms))\n",
    "\n",
    "# define intermolecular forces\n",
    "Ar_dispersion = NonBondedForce(universe,\n",
    "                          Ar.atom_type,\n",
    "                          cutoff=8.,\n",
    "                          ewald=1e-6,\n",
    "                          function=NonBonded(charge=0.0, epsilon=1.0243, sigma=3.36))\n",
    "\n",
    "# MD Engine setup\n",
    "simulation = Simulation(universe,\n",
    "                        engine=\"openmm\",\n",
    "                        time_step=10.18893,\n",
    "                        temperature=120.,\n",
    "                        traj_step=15,        \n",
    "                        openmm_ensembles=[\n",
    "                            # equilibration stage - equilibrate the cell volume and temperature\n",
    "                            # with high friction and frequent Monte Carlo pressure changes\n",
    "                            {\n",
    "                                \"integrator\": \"LangevinMiddle\",\n",
    "                                \"frictionCoeff\": 1.0 / unit.picoseconds,\n",
    "                                \"barostat\": {\n",
    "                                    \"barostat\": \"MonteCarlo\",\n",
    "                                    # https://journals.aps.org/pra/pdf/10.1103/PhysRevA.31.3391\n",
    "                                    # table 2 measurement (a) 2.01 MPa\n",
    "                                    \"defaultPressure\": 20.1 * unit.bar,\n",
    "                                    \"frequency\": 25\n",
    "                                },\n",
    "                                # runs auto-equilibration using the KPSS test on certain properties\n",
    "                                # this tuple can be replaced with an int if you prefer to run\n",
    "                                # a specific number of steps instead.\n",
    "                                # auto-equilibration parameters are as follows:\n",
    "                                # (ensemble [NPT runs KPSS test on volume and temperature],\n",
    "                                # max number of steps, steps per iteration, kpss window,\n",
    "                                # kpss tolerance)\n",
    "                                \"n_steps\": (\"NPT\", 100000, 100, 1000, 0.01)\n",
    "                            },\n",
    "                            # production stage - not used here\n",
    "                            {\n",
    "                                \"integrator\": \"LangevinMiddle\",\n",
    "                                \"frictionCoeff\": 1.0 / unit.picoseconds,\n",
    "                            }\n",
    "                        ])\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then equilibrate the simulation. This is done by setting `\"n_steps\": (\"NPT\", 100000, 100, 1000, 0.01)` above and using with `equilibration=True` in `simulation.run(n_steps=1, equilibration=True)` note that the `n_steps` in `simulation.run` is ignored as the maximum steps is specified in the settings tuple."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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(volumes, temperatures, energies, line=0):\n",
    "    plt.figure()\n",
    "    plt.subplot(311)\n",
    "    plt.plot(np.arange(len(volumes))*10, volumes)\n",
    "    plt.ylabel('Volumes')\n",
    "    if line:\n",
    "        plt.axvline(x=line, linestyle='--', color='lime')\n",
    "\n",
    "    plt.subplot(312)\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(313)\n",
    "    plt.plot(np.arange(len(energies))*10, energies, color=\"orange\")\n",
    "    plt.xlabel('time (steps)')\n",
    "    plt.ylabel('Total energy (kJ/mol)')\n",
    "    if line:\n",
    "        plt.axvline(x=line, linestyle='--', color='lime')\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Avoiding under- or over- equilibrating"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we under-equilibrate, 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.engine.autoequilibrate`. 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.engine.add_barostat(simulation.engine.openmm_ensembles[0][\"barostat\"])\n",
    "total_steps, vals_dict = simulation.engine.autoequilibrate(*simulation.engine.openmm_ensembles[0][\"n_steps\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_vols = vals_dict['volumes']\n",
    "auto_temperatures = vals_dict['temperatures']\n",
    "auto_energies = vals_dict['total_energies']\n",
    "\n",
    "plot_vars(auto_vols, auto_temperatures, auto_energies)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try changing the values in `\"n_steps\": (\"NPT\", 100000, 100, 1000, 0.01)` or increasing you simulation size `universe = Universe(dimensions=38.4441)` and rerunning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.11.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}