{
 "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",
    "import copy\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from MDMC.control import Control\n",
    "from MDMC.MD import Atom, NonBonded, Simulation, Universe\n",
    "from MDMC.MD.interactions import NonBondedForce\n",
    "from MDMC.readers.observables.xml_SQw import XML_SQw\n",
    "from MDMC.refinement.FoM.FoM_abs import ObservablePair\n",
    "from MDMC.trajectory_analysis.observables.mdanse_observable import MDANSEObservable, get_default_mdanse_settings\n"
   ]
  },
  {
   "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": [
    "NonBondedForce(\n",
    "        universe,\n",
    "        Ar.atom_type,\n",
    "        cutoff=10.0,\n",
    "        ewald=1e-6,\n",
    "        function=NonBonded(charge=0.0, epsilon=1.0, sigma=3.0)\n",
    "    )"
   ]
  },
  {
   "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"
   ]
  },
  {
   "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}]\n",
    "\n",
    "start_params = get_default_mdanse_settings(\"SQw\")\n",
    "\n",
    "data_parser = XML_SQw('data/Well_s_q_omega_Ar_data.xml')\n",
    "\n",
    "exp_observable = MDANSEObservable(mdanse_job_type=\"SQw\")\n",
    "exp_observable.read_from_file(data_parser)\n",
    "md_observable = MDANSEObservable(mdanse_job_type=\"SQw\")\n",
    "md_observable.origin = 'MD'\n",
    "md_observable.independent_variables = copy.deepcopy(\n",
    "    exp_observable.independent_variables)\n",
    "\n",
    "observable_pair = ObservablePair(exp_obs=exp_observable,\n",
    "                                 MD_obs=md_observable,\n",
    "                                 weight=1.0,\n",
    "                                 rescale_factor=1.0,\n",
    "                                 auto_scale=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this stage we check and adjust the analysis parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_settings = md_observable.initial_parameters()\n",
    "print(\"Original parameters:\")\n",
    "for name, value in job_settings.items():\n",
    "    print(f\"{name}: {value}\")\n",
    "\n",
    "new_settings = {\n",
    "    \"running_mode\": (\"multicore\", -8)\n",
    "}\n",
    "md_observable.set_parameters(new_settings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the observables to check if the current parameters of the simulation and the observable calculation produce a result that covers the range of the experimental data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, axis in exp_observable.independent_variables.items():\n",
    "    print(f\"{name}: {axis}\")\n",
    "\n",
    "TIME_STEP_FS = 5.0\n",
    "FRAME_STEP = 15\n",
    "CORR_FRAMES = 200\n",
    "TOTAL_FRAMES = 10000\n",
    "\n",
    "for output_axis in md_observable.predict_output(time_step = TIME_STEP_FS,\n",
    "                                                frame_step = FRAME_STEP,\n",
    "                                                total_frames= TOTAL_FRAMES,\n",
    "                                                correlation_frames = CORR_FRAMES,\n",
    "                                                universe=universe):\n",
    "    print(f\"{output_axis[0]}: {output_axis[1]} {output_axis[2]}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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 and trajectory step which we just determined; \n",
    "- the temperature of the Universe (used to calculate atom velocity);\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MD Engine setup\n",
    "simulation = Simulation(universe,\n",
    "                        engine=\"openmm\",\n",
    "                        time_step=TIME_STEP_FS,\n",
    "                        temperature=120.,\n",
    "                        traj_step=FRAME_STEP,\n",
    "                        )\n",
    "\n",
    "md_observable.set_parameters({\"correlation_frames\": CORR_FRAMES})"
   ]
  },
  {
   "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.run(n_steps=10000, equilibration=True)"
   ]
  },
  {
   "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 is CMA-ES (Covariance Matrix Adaptation Evolution Strategy). During the optimisation procedure, both the MD trajectory and the simulated observable can be dumped to a file for further examination by loading it into an external package like MDANSE. File dumping can be set with the `file_dump` parameters seen in the code below. Each parameter controls how often the files are dumped, what their file names or locations will be:\n",
    "- `file_dump_frequency`: Defines how often the trajectory should be dumped to a H5MD file. Options available are \"best\" (only save the trajectory with the lowest FoM), \"none\" (don't save any trajectories), and \"every\" (save every trajectory).\n",
    "- `file_dump_extent`: Which files should be written out in the dump. Options available are \"traj\" (H5MD trajectory file), \"obs\" (MDA MDANSE observable file), and \"both\" (both the H5MD trajectory and MDA observable files).\n",
    "- `file_dump_loc`: Location that the H5MD and/or MDA files should be stored at.\n",
    "- `file_dump_timestamped`: Whether a time stamp should be added to the output file names.\n",
    "- `file_dump_prefix`: The name the dumped H5MD and/or MDA file should be."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "control = Control(simulation=simulation,\n",
    "                  exp_datasets=exp_datasets,\n",
    "                  fit_parameters=fit_parameters,\n",
    "                  observable_pairs= [observable_pair],\n",
    "                  equilibration_steps=9000,\n",
    "                  MD_steps=TOTAL_FRAMES,\n",
    "                  cont_slicing=True,\n",
    "                  data_printer='ipython',\n",
    "                  file_dump_extent=\"all\",\n",
    "                  file_dump_frequency=\"every\",\n",
    "                  FoM_options={'error': 'none'},\n",
    "                  file_dump_prefix=\"mdmc_argon_mdanse\",\n",
    "                  conv_tol = 1e-6,)\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=5000)\n",
    "control.equilibrate(n_steps=15000)"
   ]
  },
  {
   "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 `2` 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": [
    "N_STEPS=2\n",
    "control.refine(n_steps=N_STEPS)  # Increase this number for a real run\n",
    "if N_STEPS > 10:\n",
    "    control.plot_results()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "thirteen (3.13.13)",
   "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.13.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}