{
 "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"
   ]
  },
  {
   "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`. To define molecules, a list of lists should be passed under the `specie_indices` keyword. The outer list has one entry per molecule and each inter list has the indices of that molecule. Only identical molecules are supported. The masses of the atoms in the molecules can be specified with `masses`. This must be a list with the same length as a molecule (the length of one of the inner lists in `specie_indices`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61cc6b61-7e5b-4b87-ab8c-58ba55def62e",
   "metadata": {},
   "outputs": [],
   "source": [
    "molecules = [[289, 290, 291,   292,   293,   294],\n",
    "             [285, 296, 297,   298,   299,   300]]\n",
    "masses    =  [12,  12,  1.008, 1.008, 1.008, 1.008]\n",
    "\n",
    "p_parms = {'specie': None,\n",
    "           'time_step': 1.2e-03,\n",
    "           'step_skip': 100,\n",
    "           'specie_indices': molecules,\n",
    "           'masses': masses,\n",
    "           'progress': False\n",
    "           }\n",
    "\n",
    "u_params = {'progress': False}"
   ]
  },
  {
   "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, parser_params=p_parms, uncertainty_params=u_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "043ab741-2814-4289-8ed3-5e7a6a17590b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.errorbar(diff.dt, diff.msd, diff.msd_std)\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": [
    "diff.diffusion(50, {'progress': False})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4787d956-9134-4116-904d-3b6bb8840f78",
   "metadata": {},
   "outputs": [],
   "source": [
    "diff.D.n, diff.D.ci()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8d462cb-97c9-45c9-bc42-cecfe2e98cb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "diff.intercept.n, diff.intercept.ci()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efe816d8-9994-4e1d-89f9-2c85f1e5695f",
   "metadata": {},
   "outputs": [],
   "source": [
    "diff.D.samples"
   ]
  },
  {
   "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",
    "plt.plot(diff.dt, diff.msd, 'k-')\n",
    "for i, ci in enumerate(credible_intervals):\n",
    "    plt.fill_between(diff.dt,\n",
    "                     *np.percentile(diff.distribution, ci, 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,
   "id": "b94b1262-a061-45d3-b576-c550df6028d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(diff.D.samples, density=True)\n",
    "plt.axvline(diff.D.n, c='k')\n",
    "plt.xlabel('$D$/cm$^2$s$^{-1}$')\n",
    "plt.ylabel('$p(D$/cm$^2$s$^{-1})$')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}