{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "235dee9e-d50a-4cb4-8c25-fa0f3fc33c50",
   "metadata": {},
   "source": [
    "# Diffusion coefficient of molecules using center of mass\n",
    "\n",
    "`Kinisi` includes the ability to calculate the mean-squared displacement and diffusion coefficient of the center of mass (or geometry) of molecules. This can be done for a single molecule or a collection of molecules. It is important to note that inclusion of rotational motion in the calculation of diffusion coeffiencents can lead to erronious results. This rotation can be eliminated from the calculation by taking the center of mass for each molecule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc9626ec-5a7d-4be9-a23c-6721f6f4addc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This initialises the duecredit framework\n",
    "%env DUECREDIT_ENABLE=yes\n",
    "%env DUECREDIT_FILE=ase_COG.duecredit.p\n",
    "import duecredit\n",
    "duecredit.due.activate()\n",
    "\n",
    "# Other imports for the notebook\n",
    "import numpy as np\n",
    "from ase.io import read\n",
    "import matplotlib.pyplot as plt\n",
    "from kinisi.analyze import DiffusionAnalyzer\n",
    "import scipp as sc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "144679f3-d48e-49af-a7bb-4f9a4d5cd02e",
   "metadata": {},
   "source": [
    "We will use a simulation of ethene in ZSM-5 zeolite. This was run in DL_POLY, so we will use `ASE` to load in the trajectory (HISTORY) file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7c6a6ab-8511-4f61-9b1c-c062dae23d23",
   "metadata": {},
   "outputs": [],
   "source": [
    "traj = read('ethene_zeo_HISTORY.gz', format='dlp-history', index=':')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9afa4b93-7e49-4a83-9426-fbef4aaf8ecc",
   "metadata": {},
   "source": [
    "We want to calculate the diffusion of the center of mass of the ethene molecule. This can be done by setting `specie` to None and specifying the indices of the molecules of interest in `specie_indices` as a scipp array. To define molecules, a scipp array should be passed under the `specie_indices` keyword with two dimensions `atom` and `group_of_atoms`. \n",
    "\n",
    "The `particle` dimension should be the same size as the number of molecules, `particle` is the generic term in kinisi for a particle or group of particles. \n",
    "\n",
    "The `atoms in particle` dimension should be the same size as the number of atoms in each each molecule. \n",
    "\n",
    "In the example below we are calculating the msd and diffusion for 2 ethene molecules. \n",
    "\n",
    " Only identical molecules are supported. The masses of the atoms in the molecules can be specified with `masses`. This must be a scipp array with the same length and dimension `atoms in particle` as a molecule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61cc6b61-7e5b-4b87-ab8c-58ba55def62e",
   "metadata": {},
   "outputs": [],
   "source": [
    "molecules = [[288, 289, 290, 291, 292, 293],\n",
    "             [284, 295, 296, 297, 298, 299]]\n",
    "mass    =  [12, 12, 1.008, 1.008, 1.008, 1.008]\n",
    "\n",
    "params = {'specie': None,\n",
    "           'time_step': 1.2e-03 * sc.Unit('ps'),\n",
    "           'step_skip': 100 * sc.Unit('dimensionless'),\n",
    "           'specie_indices': sc.array(dims=['particle', 'atoms in particle'], \n",
    "                                      values=molecules, \n",
    "                                      unit=sc.Unit('dimensionless')),\n",
    "           'masses': sc.array(dims = ['atoms in particle'], \n",
    "                              values = mass),\n",
    "           'progress': False\n",
    "           }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9b5ea9a-2ba1-4d2b-8c27-e6a479721289",
   "metadata": {},
   "source": [
    "With the parameters set, we now calculate the mean squared-displacement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "018a6b82-7eb3-4229-ae25-d143ca4cd382",
   "metadata": {},
   "outputs": [],
   "source": [
    "diff = DiffusionAnalyzer.from_ase(traj, **params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "043ab741-2814-4289-8ed3-5e7a6a17590b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.errorbar(diff.dt.values[::20], \n",
    "             diff.msd.values[::20], \n",
    "             diff.msd.variances[::20]**0.5)\n",
    "plt.ylabel('MSD/Å$^2$')\n",
    "plt.xlabel(r'$\\Delta t$/ps')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc9db243-e5f4-4401-b337-f37eb9e9e4da",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_of_diffusion = 60 * sc.Unit('ps')\n",
    "diff.diffusion(start_of_diffusion, progress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4787d956-9134-4116-904d-3b6bb8840f78",
   "metadata": {},
   "outputs": [],
   "source": [
    "diff.D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ae2cfb1",
   "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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24f70169-6ce7-48c0-ba7b-8c4d04070a35",
   "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",
   "id": "3891ecc9",
   "metadata": {},
   "source": [
    "The `duecredit` framework can then be used to output the relevant references that should be cited. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45fd779d",
   "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": 5
}