{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating an Observable from a Simulation \n",
    "When running a refinement, MDMC makes quantitative comparisons between properties measured experimentally and calculated from MD simulations: within MDMC each of these properties is an `Observable`. Examples of these observables are the dynamic structure factor, $S(Q,\\omega)$, and the 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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial requires matplotlib to be installed:"
   ]
  },
  {
   "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",
    "    %pip install matplotlib\n",
    "    try:\n",
    "        %matplotlib notebook\n",
    "        import matplotlib.pyplot as plt\n",
    "        from mpl_toolkits.mplot3d import Axes3D\n",
    "    except ImportError:\n",
    "        print('Please restart the kernel so that matplotlib can be imported.')"
   ]
  },
  {
   "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",
    "**As this is minimizing, equilibrating, and running a production run, this should take ~3 minutes to execute**. Minimization lowers the potential energy of the simulated system by adjusting atomic positions. Running a simulation with `equilibration` set to `True` will not record a `CompactTrajectory`, and since we specified neither a thermostat or barastat for the simulation a Berendsen thermostat will be applied for the duration of the equilibration. It will not be present for the second run, when the `CompactTrajectory` is recorded."
   ]
  },
  {
   "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",
    "# Energy Minimization and equilibration\n",
    "simulation.minimize(n_steps=50)\n",
    "simulation.run(n_steps=20000, equilibration=True)\n",
    "simulation.run(12025)\n",
    "trajectory = simulation.trajectory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3D surface plotting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we define a simple function for plotting 3D surfaces, which we will use for plotting the observables."
   ]
  },
  {
   "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()"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As is mentioned in the Examples section, each observable may have one or more aliases.  So 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."
   ]
  },
  {
   "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; the exercise of calculating the incoherent and total scattering is left as an exercise for the reader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sqw_coh = observables.SQwCoh()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dynamic structure factor has two independent variables, the scattering wavevector, $Q$, and $E$ (or $\\hbar\\omega$) is the energy transfer.\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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using the arange function from NumPy - this creates an uniform array of floats with a\n",
    "# start (0.42), stop (3.8) and step (0.3) specified below\n",
    "Q = np.arange(0.42, 3.8, 0.3)\n",
    "print('Q will be calculated for values of: {}'.format(Q))"
   ]
  },
  {
   "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 (**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)\n",
    "\n",
    "# Use the plotting function defined above to plot the result\n",
    "plot_surface(sqw_coh.E, sqw_coh.Q, sqw_coh.SQw[0])"
   ]
  },
  {
   "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 `calulate_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 be probably 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}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using this binning, recalculate the dynamic structure factor, which is smoother due to better statistics **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])"
   ]
  },
  {
   "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. Phys. Rev. A **31** 3391 (1985) data, which the authors corrected for instrument resolution) by reading the data into an SQwCoh object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a second SQwCoh object, this time called sqw_exp\n",
    "sqw_exp = observables.SQwCoh()\n",
    "\n",
    "# Read the data from the file ./data/Well_s_q_omega_Ar_data.xml using the `xml_SQw` reader\n",
    "sqw_exp.read_from_file('xml_SQw', './data/Well_s_q_omega_Ar_data.xml')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The independent and dependent variables are read into the same attributes (`sqw_coh.E`, `sqw_exp.Q`, and `sqw_exp.SQw`) as if they had been calculated from MD (as in `sqw_coh` above)."
   ]
  },
  {
   "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 convolution with 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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}