{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [PyBroMo](http://tritemio.github.io/PyBroMo/) - 1. Disk-single-core - Simulate 3D trajectories"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<small>\n",
    "*This notebook is part of [PyBroMo](http://tritemio.github.io/PyBroMo/) a \n",
    "python-based single-molecule Brownian motion diffusion simulator \n",
    "that simulates confocal [smFRET](http://en.wikipedia.org/wiki/Single-molecule_FRET)\n",
    "experiments. You can find the full list of notebooks in \n",
    "[Usage Examples](http://tritemio.github.io/PyBroMo/#usage-examples).*\n",
    "</small>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## *Overview*\n",
    "\n",
    "*In this notebook we show how to perform a 3-D trajectories simulation of a set of freely diffusing molecules. The simulation computes (and saves!) 3-D trajectories and emission rates due to a confocal excitation PSF for each single molecule. Depending on the simulation length, the required disk space can be significant (~ 750MB per minute of simulated diffusion).*\n",
    "\n",
    "*For more info see [PyBroMo Homepage](http://tritemio.github.io/PyBroMo/)*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Together with a few standard python libraries we import **PyBroMo** using the short name `pbm`. \n",
    "All **PyBroMo** functions will be available as `pbm.`*something*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import tables\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import pybromo as pbm\n",
    "print('Numpy version:', np.__version__)\n",
    "print('PyTables version:', tables.__version__)\n",
    "print('PyBroMo version:', pbm.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we define the simulation parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Initialize the random state\n",
    "rs = np.random.RandomState(seed=1)\n",
    "print('Initial random state:', pbm.hash_(rs.get_state()))\n",
    "\n",
    "# Diffusion coefficient\n",
    "Du = 12.0            # um^2 / s\n",
    "D1 = Du*(1e-6)**2    # m^2 / s\n",
    "D2 = D1/2\n",
    "\n",
    "# Simulation box definition\n",
    "box = pbm.Box(x1=-4.e-6, x2=4.e-6, y1=-4.e-6, y2=4.e-6, z1=-6e-6, z2=6e-6)\n",
    "\n",
    "# PSF definition\n",
    "psf = pbm.NumericPSF()\n",
    "\n",
    "# Particles definition\n",
    "P = pbm.Particles(num_particles=20, D=D1, box=box, rs=rs)\n",
    "P.add(num_particles=15, D=D2)\n",
    "\n",
    "# Simulation time step (seconds)\n",
    "t_step = 0.5e-6\n",
    "\n",
    "# Time duration of the simulation (seconds)\n",
    "t_max = 1\n",
    "\n",
    "# Particle simulation definition\n",
    "S = pbm.ParticlesSimulation(t_step=t_step, t_max=t_max, \n",
    "                            particles=P, box=box, psf=psf)\n",
    "\n",
    "print('Current random state:', pbm.hash_(rs.get_state()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most important line is the last line which creates an object `S` \n",
    "that contains all the simulation parameters (it also contains methods to run \n",
    "the simulation). You can print `S` and check the current parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "S"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "or check the required RAM for the current parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "S.print_sizes()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **NOTE:** This is the maximum in-memory array size when using a single chunk. \n",
    "> In the following, we simulate the diffusion in smaller time windows (chunks), \n",
    "> thus requiring only a few tens MB of RAM, regardless of the simulated duration."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Brownian motion simulation\n",
    "\n",
    "In the brownian motion simulation we keep using the same random state object `rs`. \n",
    "Initial and final state are saved so the same simulation can be reproduced. \n",
    "See [PyBroMo - A1. Reference - Data format and internals.ipynb](PyBroMo - A1. Reference - Data format and internals.ipynb) \n",
    "for more info on the random state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('Current random state:', pbm.hash_(rs.get_state()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#%%timeit -n1 -r1\n",
    "S.simulate_diffusion(total_emission=False, save_pos=True, verbose=True, rs=rs, chunksize=2**19, chunkslice='times')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('Current random state:', pbm.hash_(rs.get_state()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The normalized emission rate (peaks at 1) for each particle is stored \n",
    "in a 2D pytable array and accessible through the `emission` attribute:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('Simulation file size: %d MB' % (S.store.h5file.get_filesize()/1024./1024.))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "S.compact_name()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "S.particles.diffusion_coeff_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "S.store.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load trajectories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "S = pbm.ParticlesSimulation.from_datafile('0168', mode='w')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Plotting the emission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#from IPython.display import display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def plot_emission(S, s=0, size=2e6, slice_=None, em_th=0.01, save=False, figsize=(9, 4.5)):\n",
    "    if slice_ is None:\n",
    "        slice_ = (s*size, (s+1)*size)\n",
    "    slice_ = slice(*slice_)\n",
    "    em = S.emission[:, slice_]\n",
    "    dec = 1 if slice_.step is None else slice_.step\n",
    "    t_step = S.t_step*dec\n",
    "    t = np.arange(em.shape[1])*(t_step*1e3)\n",
    "    fig, ax = plt.subplots(figsize=figsize)\n",
    "    for ip, em_ip in enumerate(em):\n",
    "        if em_ip.max() < em_th: continue\n",
    "        plt.plot(t, em_ip, label='P%d' % ip)\n",
    "    ax.set_xlabel('Time (ms)')\n",
    "    \n",
    "    rs_hash = pbm.hash_(S.traj_group._v_attrs['init_random_state'])[:3]\n",
    "    ax.set_title('%ds ID-EID: %d-%d, sim rs = %s, part rs = %s' %\\\n",
    "              (s, S.ID, S.EID, rs_hash, S.particles.rs_hash[:3]))\n",
    "    ax.legend(bbox_to_anchor=(1.03, 1), loc=2, borderaxespad=0.)\n",
    "    if save:\n",
    "        plt.savefig('em %ds ID-EID %d-%d, rs=%s' %\\\n",
    "                (s, S.ID, S.EID, rs_hash), \n",
    "                dpi=200, bbox_inches='tight')\n",
    "    #plt.close(fig)\n",
    "    #display(fig)\n",
    "    #fig.clear()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plot_emission(S, slice_=(0, 2e6, 10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def plot_tracks(S, slice_=None, particles=None):\n",
    "    if slice_ is None:\n",
    "        slice_ = (0, 100e3, 100)\n",
    "    duration = (slice_[1] - slice_[0])*S.t_step\n",
    "    slice_ = slice(*slice_)\n",
    "    \n",
    "    if particles is None:\n",
    "        particles = range(S.num_particles)\n",
    "    \n",
    "    fig, AX = plt.subplots(1, 2, figsize=(11, 5), sharey=True)\n",
    "    plt.subplots_adjust(left=0.05, right=0.93, top=0.95, bottom=0.09,\n",
    "                        wspace=0.05)\n",
    "    plt.suptitle(\"Total: %.1f s, Visualized: %.2f ms\" % (\n",
    "                 S.t_step*S.n_samples, duration*1e3))\n",
    "\n",
    "    for ip in particles:\n",
    "        x, y, z = S.position[ip, :, slice_]\n",
    "        x0, y0, z0 = S.particles[ip].r0\n",
    "        plot_kwargs = dict(ls='', marker='o', mew=0, ms=2, alpha=0.5, \n",
    "                           label='P%d' % ip)\n",
    "        l, = AX[0].plot(x*1e6, y*1e6, **plot_kwargs)\n",
    "        AX[1].plot(z*1e6, y*1e6, color=l.get_color(), **plot_kwargs)\n",
    "        #AX[1].plot([x0*1e6], [y0*1e6], 'o', color=l.get_color())\n",
    "        #AX[0].plot([x0*1e6], [z0*1e6], 'o', color=l.get_color())\n",
    "\n",
    "    AX[0].set_xlabel(\"x (um)\")\n",
    "    AX[0].set_ylabel(\"y (um)\")\n",
    "    AX[1].set_xlabel(\"z (um)\")\n",
    "\n",
    "    sig = np.array([0.2, 0.2, 0.6])*1e-6\n",
    "    ## Draw an outline of the PSF\n",
    "    a = np.arange(360)/360.*2*np.pi\n",
    "    rx, ry, rz = (sig)  # draw radius at 3 sigma\n",
    "    AX[0].plot((rx*np.cos(a))*1e6, (ry*np.sin(a))*1e6, lw=2, color='k')\n",
    "    AX[1].plot((rz*np.cos(a))*1e6, (ry*np.sin(a))*1e6, lw=2, color='k')\n",
    "    \n",
    "    AX[0].set_xlim(-4, 4)\n",
    "    AX[0].set_ylim(-4, 4)\n",
    "    AX[1].set_xlim(-4, 4)\n",
    "    \n",
    "    if len(particles) <= 20:\n",
    "        AX[1].legend(bbox_to_anchor=(1.01, 1), loc=2, borderaxespad=0.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plot_tracks(S) #particles=[2, 5, 7, 22, 30])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulate timestamps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "S.simulate_timestamps(max_rate=200e3, bg_rate=1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "S.simulate_timestamps(max_rate=300e3, bg_rate=2e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "S.simulate_timestamps(max_rate=150e3, bg_rate=5e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "S.timestamp_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def plot_timetrace(S, n=-1, max_rate=200e3, bg_rate=1e3, binwidth=1e-3, scroll_gui=False):\n",
    "    fig, ax = plt.subplots()\n",
    "    plt.xlabel(\"Time (s)\")\n",
    "    plt.ylabel(\"Counts\")\n",
    "    times, part = S.get_timestamps(max_rate, bg_rate)\n",
    "    clk_p = times._v_attrs['clk_p']\n",
    "    binwidth_clk = binwidth/clk_p\n",
    "    \n",
    "    bins = np.arange(0, times[n] + 1, binwidth_clk)\n",
    "    counts, _ = np.histogram(times[:n], bins=bins)\n",
    "    \n",
    "    bins *= (clk_p*1e3)\n",
    "    plt.plot(bins[:-1], counts)\n",
    "    plt.xlabel('Time (ms)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plot_timetrace(S, max_rate=200e3, bg_rate=1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plot_timetrace(S, max_rate=300e3, bg_rate=2e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plot_timetrace(S, max_rate=150e3, bg_rate=5e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plot_emission(S, slice_=(0, 2e6, 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "S.get_timestamps(200e3, 1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}