{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Diffusion coefficient from a VASP file\n",
    "\n",
    "[Previously](./vasp_msd.html), we looked at obtaining accurate estimates for the mean-squared displacement with `kinisi`. \n",
    "Here, we show that the same `DiffusionAnalyzer` can be used to evaluate the diffusion coefficient, using the `kinisi` [methodology](./methodology.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This initialises the duecredit framework\n",
    "%env DUECREDIT_ENABLE=yes\n",
    "%env DUECREDIT_FILE=vasp_d.duecredit.p\n",
    "import duecredit\n",
    "duecredit.due.activate()\n",
    "\n",
    "# Other imports for the notebook\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipp as sc\n",
    "from pymatgen.io.vasp import Xdatcar\n",
    "from kinisi.analyze import DiffusionAnalyzer\n",
    "rng = np.random.RandomState(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As with the [previous example](./vasp_msd.html), `params` dictionary will describe the details about the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = {'specie': 'Li',\n",
    "          'time_step': 2.0 * sc.Unit('fs'),\n",
    "          'step_skip': 50 * sc.Unit('dimensionless'),\n",
    "          'progress': False\n",
    "          }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, we will add an additional key-value pair to the dictionary. \n",
    "This argument means that the diffusion coefficient will only be calculated in the *xy* plane of the simulation box. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params['dimension']= 'xy'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As with the previous example, we now use the `from_xdatcar` class method to construct the `DiffusionAnalyzer` object. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xd = Xdatcar('./example_XDATCAR.gz')\n",
    "diff = DiffusionAnalyzer.from_xdatcar(xd, **params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above cells, we parse and determine the uncertainty on the mean-squared displacement as a function of the timestep. \n",
    "We should visualise this, to check that we are observing diffusion in our material and to determine the timescale at which this diffusion begins. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.errorbar(diff.dt.values, diff.msd.values, np.sqrt(diff.msd.variances))\n",
    "ax.set_xlabel(f'Time / {diff.dt.unit}')\n",
    "ax.set_ylabel(f'MSD / {diff.msd.unit}')\n",
    "ax.set_xlim(0, None)\n",
    "ax.set_ylim(0, None)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualise this on a log-log scale, which helps to reveal the diffusive regime (the region where the gradient stops changing)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.errorbar(diff.dt.values, diff.msd.values, np.sqrt(diff.msd.variances))\n",
    "ax.axvline(3000, color='g')\n",
    "ax.set_xlabel(f'Time / {diff.dt.unit}')\n",
    "ax.set_ylabel(f'MSD / {diff.msd.unit}')\n",
    "ax.set_xscale('log')\n",
    "ax.set_yscale('log')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The green line at 3000 fs appears to be a reasonable estimate of the start of the diffusive regime. \n",
    "Therefore, we want to pass `3000 * sc.Units('fs')` as the argument to the diffusion analysis below. \n",
    "At this stage, we pass the `random_state` argument to ensure reproducibility. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_of_diffusion = 3000 * sc.Unit('fs')\n",
    "diff.diffusion(start_of_diffusion, progress=False, random_state=rng)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This method estimates the correlation matrix between the timesteps and uses posterior sampling to find the self-diffusion coefficient, $D*$ and intercept.\n",
    "We can find the mean of the marginal posterior samples of $D*$: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diff.D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same for the intercept can be found from `diff.intercept`. \n",
    "A histogram of the marginal posterior probability distribution for $D*$ can be plotted as shown below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.hist(diff.D.values, density=True)\n",
    "ax.axvline(sc.mean(diff.D).value, c='k')\n",
    "ax.set_xlabel(f'D* / [{diff.D.unit}]')\n",
    "ax.set_ylabel(f'p(D*) / [{(1 / diff.D.unit).unit}]')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is also possible to plot the posterior distribution of the models on the data. \n",
    "We represent this distribution with 1σ-, 2σ-, and 3σ-credible intervals. \n",
    "Here we remove the measured error bars, in the interest of clarity."
   ]
  },
  {
   "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",
    "fig, ax = plt.subplots()\n",
    "ax.plot(diff.dt.values, diff.msd.values, 'k-')\n",
    "for i, ci in enumerate(credible_intervals):\n",
    "    ax.fill_between(diff.dt.values,\n",
    "                      *np.percentile(diff.distributions, ci, axis=1),\n",
    "                      alpha=alpha[i],\n",
    "                      color='#0173B2',\n",
    "                      lw=0)\n",
    "ax.set_xlabel(f'Time / {diff.dt.unit}')\n",
    "ax.set_ylabel(f'MSD / {diff.msd.unit}')\n",
    "ax.set_xlim(0, None)\n",
    "ax.set_ylim(0, None)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, the joint posterior probability distribution for the diffusion coefficient and intercept can be visualised with the `corner` library. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from corner import corner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corner(np.array([i.values for i in diff.flatchain.values()]).T, \n",
    "       labels=['/'.join([k, str(v.unit)]) for k, v in diff.flatchain.items()])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `duecredit` framework can then be used to output the relevant references that should be cited. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from duecredit import due\n",
    "\n",
    "due.dump()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "kinisi-docs",
   "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.12.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}