{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating an Observable from a Simulation \n",
    "To run a refinement, MDMC makes quantitative comparisons between a dynamical property derived from an experiment, and the equivalent dynamical property calculated from our MD simulation. Within MDMC each of these properties is called an `Observable`. Examples of these observables are the [dynamic structure factor](https://en.wikipedia.org/wiki/Dynamic_structure_factor), $S(Q,\\omega)$, and the [pair distribution function](https://en.wikipedia.org/wiki/Pair_distribution_function), $G(r)$.\n",
    "\n",
    "Within a refinement, each `Observable` is calculated from MD automatically. However it is also useful to be able to calculate an `Observable` from an MD simulation and plot it, which is demonstrated in this tutorial."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial requires matplotlib, the Python plotting library, to be installed; if it is not installed, install it in the terminal with the command `pip3 install matplotlib`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    %matplotlib notebook\n",
    "    import matplotlib.pyplot as plt\n",
    "    from mpl_toolkits.mplot3d import Axes3D\n",
    "except ImportError:\n",
    "    print('Please install matplotlib and restart the kernel!')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running a simulation\n",
    "\n",
    "Below we setup and run a simulation of liquid argon so that we have a `CompactTrajectory` from which to calculate an `Observable`.\n",
    "\n",
    "If the building of the universe and simulation is confusing to you, please see the guides for how to [create atomic configurations](creating-atomic-configurations.ipynb), [define molecule interactions](defining-molecule-interactions.ipynb) and [run a simulation](running-a-simulation.ipynb) - or the [Argon a-to-z](../../../tutorials/Argon-a-to-z.ipynb) tutorial.\n",
    "\n",
    "**As this is minimizing, equilibrating, and running a simulation, this should take ~3 minutes to execute**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from MDMC.MD import *\n",
    "import numpy as np\n",
    "\n",
    "# Build universe\n",
    "universe = Universe(dimensions=38.4441)\n",
    "Ar = Atom('Ar', charge=0.)\n",
    "# Fill box with 1000 atoms, and hence atoms per AA^-3 density = 0.0176\n",
    "universe.fill(Ar, num_struc_units=1000)\n",
    "print('Universe contains {} Ar atoms'.format(len(universe.atoms)))\n",
    "Ar_dispersion = Dispersion(universe,\n",
    "                           (Ar.atom_type, Ar.atom_type),\n",
    "                           cutoff=8.,\n",
    "                           vdw_tail_correction=True,\n",
    "                           function=LennardJones(1.0243, 3.36))\n",
    "\n",
    "# MD Engine setup\n",
    "simulation = Simulation(universe,\n",
    "                        engine=\"lammps\",\n",
    "                        time_step=10.0,\n",
    "                        temperature=120.,\n",
    "                        traj_step=25)\n",
    "\n",
    "simulation.minimize(n_steps=50)  # energy minimisation\n",
    "simulation.run(n_steps=20000, equilibration=True)  # equilibration\n",
    "\n",
    "simulation.run(12025)  # the actual run!\n",
    "trajectory = simulation.trajectory"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before our excursion into calculating the dynamic structure factor from this trajectory, we define a helper function `plot_surface` which will let us easily visualise our observable as a 3D plot. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_surface(x, y, z):\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(projection='3d')\n",
    "    X, Y = np.meshgrid(x, y)\n",
    "    surf = ax.plot_surface(X, Y, z)\n",
    "    plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculating $S(Q,\\omega)$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To see which observables can be calculated, use the `help` function to look at the observables module documentation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from MDMC.trajectory_analysis import observables\n",
    "help(observables)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each observable may have one or more aliases - both `DynamicStructureFactor` or `SQw` can be used for calculating $S(Q,\\omega)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some of the details given below are specific to calculating the dynamic structure factor, however the general method can be followed for any observable."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Argon has comparable scattering lengths for both coherent and incoherent scattering, so the total scattering has clear contributions from each.  For the purposes of this tutorial we will simply calculate $S(Q,\\omega)_{coh}$ and compare it with experimental data; calculating the incoherent and total scattering can be done similarly with the objects `observables.SQwIncoh` or `observables.SQw`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sqw_coh = observables.SQwCoh()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dynamic structure factor has two independent variables; the scattering wavevector $Q$, and the energy transfer $\\omega$ (or $E/\\hbar$ where $\\hbar$ denotes the reduced Planck constant).\n",
    "\n",
    "The size of the simulation box defines the spacing (and therefore the lower bound) of $Q$ values; however any Q values can be specified for calculating $S(Q,\\omega)$ and MDMC will simply calculate it for valid values.\n",
    "\n",
    "Here, we use the `np.arange` function to create an array of Q-values, evenly spaced in steps of 0.3 between 0.42 and 3.8."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q = np.arange(start=0.42, stop=3.8, step=0.3)\n",
    "print('Q will be calculated for values of: {}'.format(Q))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The total duration of the trajectory defines the smallest spacing of the $E$ values.  If $E$ values are not specified, MDMC will simply calculate the $E$ values with this spacing.  However for this setting there will only be two points in the trajectory contributing to each energy bin and the statistics will be poor.\n",
    "\n",
    "**should take ~30s to execute**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the Q values for the Q independent variable\n",
    "sqw_coh.independent_variables = {'Q':Q}\n",
    "\n",
    "# Calculate the coherent dynamic structure factor\n",
    "sqw_coh.calculate_from_MD(trajectory)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and indeed, the plot betrays our poor statistics:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_surface(sqw_coh.E, sqw_coh.Q, sqw_coh.SQw[0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Therefore it is strongly recommended that you specify the $E$ values for calculating the dynamic structure factor.\n",
    "\n",
    "The allowed values of $E$ (in meV) are determined by trajectory times.  To aid in calculating these, the `calculate_E` method \n",
    "can be used, which requires the time step of the trajectory (in fs) and the total number of E values to be calculated.  The number of E values should be changed to improve the statistics, but somewhere in the range of 1/4 to 1/8 of the total number of trajectory configurations (i.e. the number of times for which a configuration was recorded) will often be reasonable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This subtracts the time of the zeroeth step from the time of the first step\n",
    "# It should be equal to the time_step * traj_step passed to Simulation\n",
    "trajectory_timestep = trajectory.times[1] - trajectory.times[0]\n",
    "print('CompactTrajectory time step: {} fs'.format(trajectory_timestep))\n",
    "\n",
    "n_configurations = len(trajectory.times)\n",
    "print('Number of trajectory configurations: {}'.format(n_configurations))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "E = sqw_coh.calculate_E((n_configurations - 1) / 6, trajectory_timestep)\n",
    "print('Calculated E values are: \\n{} meV'.format(E))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Of course if the $E$ resolution required is lower (i.e. step between the `E` values can be larger), the number of E values specified can be reduced:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('A reduced E resolution: \\n{} meV'.format(sqw_coh.calculate_E((n_configurations - 1) / 24, trajectory_timestep)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use `Q` and `E` to set the indepedendent variables of `sqw_coh`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sqw_coh.independent_variables = {'E':E,\n",
    "                                 'Q':Q}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using this binning, we recalculate the dynamic structure factor, which is smoother due to better statistics.\n",
    "\n",
    " **This calculation should take ~30s.**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sqw_coh.calculate_from_MD(trajectory)\n",
    "plot_surface(sqw_coh.E, sqw_coh.Q, sqw_coh.SQw[0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compare this with experimental data measured using quasi-elastic neutron scattering (QENS) data (specifically here the [van Well et al. (1985)](https://doi.org/10.1103/PhysRevA.31.3391) data, seen in detail in the [Argon a-to-z](../../../tutorials/Argon-a-to-z.ipynb) tutorial) by reading the data into an SQwCoh object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sqw_exp = observables.SQwCoh().read_from_file(reader='xml_SQw', file_name='./data/Well_s_q_omega_Ar_data.xml')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The independent and dependent variables are read into the same attributes (`sqw_exp.E`, `sqw_exp.Q`, and `sqw_exp.SQw`) as if they had been calculated from MD (as in `sqw_coh` above)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For many QENS datasets (if not most) the instrument resolution has not been removed, and to compare with MD simulation this requires the simulation output to be smoothed to the instrument resolution. One way this can be done is using the optional `energy_resolution`:   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# energy_resolution is the FWHM of the instrumentation's resolution function expressed in ueV (micro eV).\n",
    "# It is used to smooth the calculated S(Q,w) with a Gaussian window in order to match the values obtained experimentally.\n",
    "sqw_coh.calculate_from_MD(trajectory, energy_resolution={'gaussian': 1.0e3})\n",
    "\n",
    "plot_surface(sqw_coh.E, sqw_coh.Q, sqw_coh.SQw[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other observables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same applies for the observables, such as the pair distribution function $G(r)$.  To determine how to calculate another observables from an MD trajectory, please see the associated help documentation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(observables.PDF.calculate_from_MD)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.11.2 64-bit",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}