{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Comparison with MDAnalysis\n",
    "\n",
    "`MDAnalysis` also contains a tool for calculating the mean-squared displacements. So why use `kinisi` over `MDAnalysis`?\n",
    "\n",
    "Well, the approach taken by `kinisi`, which is outlined in the [methodology](./methodology.html), uses a high precision approach to estimate the diffusion coefficent and offers an accurate estimate of the variance in the mean-squared displacements and diffusion coefficient from a single simulation.\n",
    "\n",
    "In this notebook, we will compare the results from `MDAnalysis` and `kinisi`. \n",
    "First we will import the `kinisi.analyze.DiffusionAnalyzer` and `MDAnalysis.analysis.msd` classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "warnings.filterwarnings('ignore', category=DeprecationWarning)\n",
    "warnings.filterwarnings('ignore', category=RuntimeWarning)\n",
    "warnings.filterwarnings('ignore', category=UserWarning)\n",
    "\n",
    "import numpy as np\n",
    "import scipp as sc\n",
    "\n",
    "import kinisi\n",
    "from kinisi.analyze import DiffusionAnalyzer\n",
    "import MDAnalysis.analysis.msd as msd\n",
    "from MDAnalysis.transformations.nojump import NoJump"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we are going to need to do a little 'magic' to get `MDAnalysis` to read an `extended xyz` file.\n",
    "Luckly `ASE` can help us out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ase.io import read\n",
    "from MDAnalysis import Universe\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since this trajectory has a triclinic cell; we need to extract the cell dimensions and pass them to `MDAnalysis`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The file we want to load is in the test cases for kinisi, \n",
    "# so will need to get the path to it.\n",
    "path_to_file = os.path.join(os.path.dirname(kinisi.__file__), \n",
    "                            'tests/inputs/LiPS.exyz')\n",
    "atoms = read(path_to_file, format='extxyz', index=':')\n",
    "\n",
    "cell_dimensions = []\n",
    "for frame in atoms:\n",
    "    lengths = frame.cell.lengths()\n",
    "    angles = frame.cell.angles()\n",
    "    cell = [*lengths, *angles]\n",
    "    cell_dimensions.append(cell)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have the cell dimensions ($a, b, c, \\alpha, \\beta, \\gamma$) in an array. We now need to create an `MDAnalysis` `universe` and apply these dimension to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = Universe(path_to_file, \n",
    "             path_to_file, \n",
    "             format='XYZ',\n",
    "             topology_format='XYZ', \n",
    "             dt=20.0/1000)\n",
    "for ts, dims in zip(u.trajectory, cell_dimensions):\n",
    "    ts.dimensions = dims"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we add an unwrapping transformation to the `universe`, create the `MSD` object, run the analysis, and extract the results as a timeseries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u.trajectory.add_transformations(NoJump())\n",
    "MSD = msd.EinsteinMSD(u, \n",
    "                      select='type LI', \n",
    "                      msd_type='xyz', \n",
    "                      fft=True, \n",
    "                      verbose=False)\n",
    "MSD.run(verbose=False)\n",
    "\n",
    "mda_MSD = MSD.results.timeseries\n",
    "mda_dt = np.linspace(u.trajectory.dt, \n",
    "                     u.trajectory.dt * len(mda_MSD), \n",
    "                     len(mda_MSD))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the result from `MDAnalysis` in one hand we can now do the same thing in `kinisi` with the other.\n",
    "Here, we use the custom `dt` functionality that is available in `kinisi`, with this we pass a `scipp` array of time intervals to compute the MSD over.\n",
    "Note, that the custom `dt` must be represented within possible simulation time intervals, or an error will be raised.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = {'specie': 'LI',\n",
    "          'time_step': 0.001 * sc.Unit('ps'),\n",
    "          'step_skip': 20 * sc.units.dimensionless,\n",
    "          'progress': False,\n",
    "          'dt': sc.arange(dim='time interval', start=0.1 * sc.Unit('ps'),\n",
    "                         stop=4.1 * sc.Unit('ps'),\n",
    "                         step=0.1 * sc.Unit('ps'))\n",
    "          }\n",
    "\n",
    "kinisi_from_universe = DiffusionAnalyzer.from_universe(u, **params)\n",
    "kinisi_MSD = kinisi_from_universe.msd\n",
    "time = kinisi_from_universe.dt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Putting the two analyses' together, we can plot for comparison."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.plot(time, kinisi_MSD, label='Kinisi MSD', c='#ff7f0e')\n",
    "plt.plot(mda_dt, mda_MSD, label='MDanalysis', ls='--')\n",
    "plt.ylabel('MSD (Å$^2$)')\n",
    "plt.xlabel('Diffusion time (ps)')\n",
    "plt.legend()\n",
    "\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The results overlap almost entirely.\n",
    "\n",
    "We can now extract an accurate estimation of the variance in the observed MSD from `kinisi`. For clarity, we will use array indexing to remove every other data point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.errorbar(kinisi_from_universe.dt.values,\n",
    "             kinisi_from_universe.msd.values,\n",
    "             np.sqrt(kinisi_from_universe.msd.variances),\n",
    "             c='#ff7f0e')\n",
    "plt.ylabel(r'MSD/Å$^2$')\n",
    "plt.xlabel(r'$\\Delta t$/ps')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also calculate estimated diffusion coefficient, and the associated uncertainty, with `kinisi`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kinisi_from_universe.diffusion(1.5 * sc.Unit('ps'))\n",
    "kinisi_from_universe.D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(kinisi_from_universe.D.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "credible_intervals = [[16, 84], [2.5, 97.5], [0.15, 99.85]]\n",
    "alpha = [0.6, 0.4, 0.2]\n",
    "\n",
    "plt.plot(kinisi_from_universe.dt, kinisi_from_universe.msd, 'k-')\n",
    "for i, ci in enumerate(credible_intervals):\n",
    "    plt.fill_between(kinisi_from_universe.dt.values,\n",
    "                     *np.percentile(kinisi_from_universe.distributions, \n",
    "                                    ci, \n",
    "                                    axis=1),\n",
    "                     alpha=alpha[i],\n",
    "                     color='#0173B2',\n",
    "                     lw=0)\n",
    "plt.ylabel('MSD/Å$^2$')\n",
    "plt.xlabel(r'$\\Delta t$/ps')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kinisi_from_universe.da"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "kinisi-dev",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}