{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Argon A-to-Z"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial aims to take the user from no familiarity with MDMC to enough competence that they can create a simple simulation and run a refinement, by walking through a simulation and refinement for the liquid argon data from [van Well et al. (1985)](https://doi.org/10.1103/PhysRevA.31.3391).  For more details on specific parts of this, please see the how-to guides!\n",
    "\n",
    "If you'd like to learn more about a specific object, use the Python `help()` command - for example `help(Atom)` provides information on the MDMC `Atom`.\n",
    "\n",
    "We first import all of the objects we require, and set some environmental variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from MDMC.control import Control\n",
    "from MDMC.MD import Atom, Dispersion, LennardJones, Simulation, Universe\n",
    "\n",
    "# This setting tells MDMC how many simultaneous processes on your computer it should use\n",
    "# for simulation and refinement calculations.\n",
    "os.environ[\"OMP_NUM_THREADS\"] = \"4\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting up a configuration and simulation\n",
    "\n",
    "We now build our `Universe`. For this tutorial we're creating a cube of argon-36 atoms; our side length is 23.0668Å (angstrom), and the atoms fill it with a density of 0.0176 atoms per cubic angstrom. This simulation matches the experimental data that we'd like to refine against - you will see more about this below.\n",
    "\n",
    "Atoms are created using MDMC `Atom` objects, and then the universe is filled with them via the `universe.fill` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "universe = Universe(dimensions=23.0668)\n",
    "\n",
    "Ar = Atom('Ar[36]', charge=0.)\n",
    "universe.fill(Ar, num_density=0.0176)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that at this point, there are no interaction forces between the argon atoms! The simulation doesn't know how these atoms should interact with each other. In the cell below an appropriate (for argon) force-field interaction potential is defined; we use a dispersive interaction with potential energy calculated by the [Lennard-Jones potential](https://en.wikipedia.org/wiki/Lennard-Jones_potential). This interaction has two parameters, epsilon and sigma, which determine the strength of the potential energy between atoms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Ar_dispersion = Dispersion(universe,  # the universe our interaction applies to\n",
    "                           (Ar.atom_type, Ar.atom_type),  # the types of atoms to which it applies (only one type here!)\n",
    "                           cutoff=8.,  # the cutoff distance\n",
    "                           function=LennardJones(epsilon=1.0243, sigma=3.36))  # the function which calculates the potential energy between two atoms"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A `cutoff` distance, past which atoms do not interact, is chosen arbitrarily (see `help(Dispersion)` for more info). A [rule of thumb for Lennard-Jones](https://en.wikipedia.org/wiki/Lennard-Jones_potential#Lennard-Jones_truncated_&_shifted_(LJTS)_potential) is to pick `cutoff=2.5*sigma`. The value for argon is recommended to be between 8 and 12 ang. Ideally, and for any system you want to pick at value of the `cutoff` which is small while not compromising accuracy. For this system, picking a value between 8 and 12 ang is found to give near identical results to the experimental data.\n",
    "\n",
    "Next (and before starting the refinement), we set up the Simulation, which contains:\n",
    "- the `Universe` that the simulation is based on;\n",
    "- the MD engine used to run the simulations;\n",
    "- the time-step in femtoseconds between each simulation frame; \n",
    "- the temperature of the Universe (used to calculate atom velocity);\n",
    "- the trajectory step, which is how many time steps should occur for each 'step' of the refinement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MD Engine setup\n",
    "simulation = Simulation(universe,\n",
    "                        engine=\"lammps\",\n",
    "                        time_step=10.18893,\n",
    "                        temperature=120.,\n",
    "                        traj_step=15)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then minimize and equilibrate the simulation; minimising the simulation avoids it getting 'stuck' in local minima, and equilibration runs the simulation until it reaches a state with a physically feasible temperature/energy distribution. This ensures our simulation isn't affected by the initial arrangement of the atoms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Energy Minimization and equilibration\n",
    "simulation.minimize(n_steps=50)\n",
    "simulation.run(n_steps=10000, equilibration=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Refining our data\n",
    "\n",
    "Our simulation is now fully set up. Now we need some data to which we fit our simulation; here we use an experimental [dynamic structure factor](https://en.wikipedia.org/wiki/Dynamic_structure_factor) $S(Q, \\omega)$ for liquid argon. We call these dynamical properties an 'observable'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# exp_datasets is a list of dictionaries with one dictionary per experimental\n",
    "# dataset\n",
    "# Dataset from: van Well et al. (1985). Physical Review A, 31(5), 3391-3414\n",
    "# resolution is None as the original author already accounted for instrument resolution\n",
    "exp_datasets = [{'file_name':'data/Well_s_q_omega_Ar_data.xml',\n",
    "                 'type':'SQw',\n",
    "                 'reader':'xml_SQw',\n",
    "                 'weight':1.,\n",
    "                 'auto_scale':True,\n",
    "                 'resolution':800}]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then need to create our parameters to fit against. In this case, we take all the universe parameters (which here are just the sigma and epsilon values that the Lennard-Jones potential depends on). Note that above when we set our initial `LennardJones` function in the `Dispersion` object, the epsilon and sigma were our \"initial guesses\" that the refinement will start from.\n",
    "\n",
    "We also set constraints for our fitting parameters; we bound our sigma values to be between 2.8 and 3.8, and our epsilon to be between 0.6 and 1.4."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit_parameters = universe.parameters\n",
    "fit_parameters['sigma'].constraints = [2.7,3.8]\n",
    "fit_parameters['epsilon'].constraints = [0.5, 1.5]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create our `Control` object. This object oversees the refinement; it brings the simulation, dataset, and the fitting parameters together, and then does the following:\n",
    "1. Run the simulation with the current parameters.\n",
    "2. Calculate the simulated observable from the simulation trajectory.\n",
    "3. Compare it to the experimental observable.\n",
    "4. Use a minimizer (optimisation process) to refine the parameters, bringing them closer to the experimental observable.\n",
    "5. Repeat with new parameters.\n",
    "\n",
    "The minimizer we're using in this tutorial (see the parameter `minimizer_type`) uses a [Gaussian process optimiser](https://en.wikipedia.org/wiki/Gaussian_process). If you'd like to compare them, MDMC also provides the [Metropolis-Hastings](https://en.wikipedia.org/wiki/Metropolis%E2%80%93Hastings_algorithm) Monte Carlo algorithm, which you can select by setting `minimizer_type=\"MMC\"`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "control = Control(simulation=simulation,\n",
    "                  exp_datasets=exp_datasets,\n",
    "                  fit_parameters=fit_parameters,\n",
    "                  minimizer_type=\"GPO\",\n",
    "                  reset_config=True,\n",
    "                  MD_steps=4000,\n",
    "                  equilibration_steps=4000,\n",
    "                  data_printer='ipython')\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the dataset has been specified, and used to configure various processes and parameters, the system can be equilibrated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Energy Minimization and equilibration\n",
    "control.minimize(n_steps=50)\n",
    "control.equilibrate(n_steps=10000)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The number of `MD_steps` specified must be large enough to allow for statistically reasonable calculation of all observables. This depends the `type` of the dataset provided and the value of the `traj_step` (specified when creating the `Simulation`). If a value for `MD_steps` is not provided, then the minimum number needed will be used automatically.\n",
    "\n",
    "Additionally, some observables will have an upper limit on the number of MD_steps that can be used in calculating their dependent variable(s). In these cases, the number of `MD_steps` is rounded down to a multiple of this upper limit so that we only run steps that will be useful. For example, if we use 1000 `MD_steps` in calculation, but a value of 2500 is provided, then we will run 2000 steps and use this to calculate the variable twice, without wasting time performing an additional 500 steps."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, start the refinement! `n_steps` has been set to `25` just so you can see what a refinement looks like; it will take many more steps to fully refine a dataset. Bump it up to a higher number when you're ready. Results can also be plotted via the `control.plot_results` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "control.refine(n_steps=25)\n",
    "control.plot_results()"
   ]
  }
 ],
 "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.2"
  },
  "vscode": {
   "interpreter": {
    "hash": "949777d72b0d2535278d3dc13498b2535136f6dfe0678499012e853ee9abcab1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}