{
 "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 calcuation of diffusion coeffiencents can lead to erronious results. This rotation can be elminated 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": [
    "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 calcuate 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()"
   ]
  }
 ],
 "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": 5
}