{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import bayesmsd\n",
    "import noctiluca as nl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating trajectories\n",
    "It is quite straight-forward to sample trajectories from the Gaussian processes that ``bayesmsd`` uses for fitting. To define such a process, it suffices to specify its MSD. For this example, we will sample trajectories from an MSD that interpolates two powerlaws:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@bayesmsd.deco.MSDfun\n",
    "def msd(dt):\n",
    "    return dt**0.3 + 0.03*dt**1.7\n",
    "\n",
    "np.random.seed(1548423571)\n",
    "data = bayesmsd.gp.generate((msd, 1, 2), T=100, n=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the use of the ``bayesmsd.deco.MSDfun`` decorator; this ensures proper handling of the lag time argument, i.e. $\\text{MSD}(-\\Delta t) \\equiv \\text{MSD}(\\Delta t)$ and $\\text{MSD}(0) = 0$.\n",
    "\n",
    "Let's plot this and check that it actually exhibits the MSD we aimed for:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dt = np.logspace(0, 2, 100)\n",
    "\n",
    "# Asymptotes\n",
    "plt.plot(dt, dt**0.3, color='gray', label='asymptotes')\n",
    "plt.plot(dt, 0.03*dt**1.7, color='gray')\n",
    "\n",
    "# Theoretical MSD\n",
    "plt.plot(dt, msd(dt), color='k', linewidth=2, label='theoretical MSD')\n",
    "\n",
    "# Measured from the data\n",
    "emp_msd = nl.analysis.MSD(data)\n",
    "plt.plot(np.arange(1, len(emp_msd)), emp_msd[1:], color='tab:red', label='empirical MSD')\n",
    "\n",
    "plt.legend()\n",
    "plt.title('MSD with ground truth')\n",
    "plt.xlabel('lag time [frames]')\n",
    "plt.ylabel('MSD [a.u.]')\n",
    "plt.ylim([0.8, 100])\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Being able to generate trajectories from a given MSD comes in very handy when comparing a real data set to the null hypothesis that it is sampled from a Gaussian process; we can simply fit the MSD and see what a data set with that MSD would look like, if it was drawn from a homogeneous Gaussian process. This allows us (for example) to start looking for heterogeneity in a given data set.\n",
    "\n",
    "# Detecting trajectory level heterogeneity\n",
    "\n",
    "Let's generate a data set with trajectory to trajectory variation in the anomalous diffusion constant and see, whether we can detect that heterogeneity; for simplicity we stick to powerlaw MSDs of the form $\\Gamma(\\Delta t)^\\alpha$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trajectory_with_random_Gamma(T, p_miss_frame=0.2):\n",
    "    G = np.exp(np.random.normal())\n",
    "    traj = bayesmsd.gp.generate((bayesmsd.deco.MSDfun(lambda dt: G*dt**0.63), 1, 1), T)[0]\n",
    "    \n",
    "    # randomly missing a few observations\n",
    "    traj.data[:, np.random.rand(len(traj)) < p_miss_frame, :] = np.nan\n",
    "    \n",
    "    return traj\n",
    "\n",
    "# Assemble a whole dataset of these, with varying length\n",
    "data = nl.TaggedSet((trajectory_with_random_Gamma(T)\n",
    "                     for T in np.random.geometric(1/100, size=50) + 10),\n",
    "                    hasTags=False,\n",
    "                   )\n",
    "\n",
    "# Fit MSD\n",
    "fit = bayesmsd.lib.NPXFit(data, ss_order=1)\n",
    "res = fit.run(show_progress=True)\n",
    "\n",
    "# Generate control data sets with fitted MSD\n",
    "data_control1 = bayesmsd.gp.generate_dataset_like(data, (fit, res))\n",
    "data_control2 = bayesmsd.gp.generate_dataset_like(data, (fit, res))\n",
    "\n",
    "# Let's plot both and compare\n",
    "fig, axs = plt.subplots(1, 3,\n",
    "                        figsize=[15, 4],\n",
    "                        sharex=True, sharey=True,\n",
    "                       )\n",
    "\n",
    "for ax, dat, title in zip(axs,\n",
    "                          [data, data_control1, data_control2],\n",
    "                          ['original data', 'homogeneous control I', 'homogeneous control II'],\n",
    "                         ):\n",
    "    nl.plot.msd_overview(dat, ax=ax)\n",
    "    ax.set_title(title)\n",
    "\n",
    "ax = axs[0]\n",
    "ax.set_ylim([1e-1, None])\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use ``bayesmsd.gp.generate_dataset_like`` to generate the control datasets. This function creates a dataset that has the same number of trajectories as the original, with the same lengths and missing frames, but otherwise is just sampled from the Gaussian Process defined by the given MSD (here in form of a ``Fit`` object and associated fit result). We can clearly see that the broad scatter at early times in the original data is not consistent with a homogeneous Gaussian process; thus there has to be some real heterogeneity in the data, which here of course was put in by hand."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting fit results\n",
    "In the above example, we would like to add the powerlaw fit to the plot. For this purpose, ``Fit`` objects provide the ``MSD()`` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Just copied from above\n",
    "fig, axs = plt.subplots(1, 3,\n",
    "                        figsize=[15, 4],\n",
    "                        sharex=True, sharey=True,\n",
    "                       )\n",
    "\n",
    "for ax, dat, title in zip(axs,\n",
    "                          [data, data_control1, data_control2],\n",
    "                          ['original data', 'homogeneous control I', 'homogeneous control II'],\n",
    "                         ):\n",
    "    nl.plot.msd_overview(dat, ax=ax)\n",
    "    ax.set_title(title)\n",
    "\n",
    "ax = axs[0]\n",
    "ax.set_ylim([1e-1, None])\n",
    "\n",
    "# Now add the fitted MSD\n",
    "msd = fit.MSD(res['params']) # `msd` is now a callable function representing MSD(Δt)\n",
    "dt = np.logspace(0, 3)\n",
    "for ax in axs:\n",
    "    ax.plot(dt, msd(dt),\n",
    "            color='r', linewidth=3,\n",
    "            label='fitted MSD',\n",
    "           )\n",
    "    ax.legend()\n",
    "\n",
    "    \n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}