{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [PyBroMo](http://tritemio.github.io/PyBroMo/) - GUI Trajectory explorer"
   ]
  },
  {
   "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 implements an interactive 3-D trajectories visualizer. To visualize trajectories you need simulatte the trajectories first.*\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
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import tables\n",
    "import matplotlib.pyplot as plt\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": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib qt\n",
    "import seaborn as sns\n",
    "sns.set_style('whitegrid')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load trajectories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "S = pbm.manage.load_trajectories('pybromo_4bce*')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Plotting the emission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from PyQt4 import QtGui, QtCore\n",
    "\n",
    "class TrackPlotter:\n",
    "    def __init__(self, S, particles=None, duration=0.01, decimate=100, scale=1, page_step=100):\n",
    "        if particles is None:\n",
    "            particles = range(S.num_particles)\n",
    "        self.particles = particles\n",
    "        self.S = S\n",
    "        self.position = S.position\n",
    "        self.duration = duration\n",
    "        self.decimate = decimate\n",
    "        self.scale = scale\n",
    "        self.page_step = page_step\n",
    "        self.duration_steps = duration//S.t_step\n",
    "        self.num_points = self.duration_steps//decimate\n",
    "        self.fig, self.AX = plt.subplots(1, 2, figsize=(11, 5), sharey=True)\n",
    "        \n",
    "        # Retrive the QMainWindow used by current figure and add a toolbar\n",
    "        # to host the new widgets\n",
    "        QMainWin = self.fig.canvas.parent()\n",
    "        toolbar = QtGui.QToolBar(QMainWin)\n",
    "        #QMainWin.addToolBar(QtCore.Qt.BottomToolBarArea, toolbar)\n",
    "        QMainWin.addToolBar(QtCore.Qt.TopToolBarArea, toolbar)\n",
    "        # Create the slider and spinbox for x-axis scrolling in toolbar\n",
    "        self.set_slider(toolbar)\n",
    "        \n",
    "        self.setup_plot()\n",
    "        self.update()\n",
    "\n",
    "    def set_limits(self):\n",
    "        self.smin = 0\n",
    "        self.smax = self.position.shape[-1]\n",
    "        self.width = self.num_points*self.decimate\n",
    "\n",
    "    def set_slider(self, parent):\n",
    "        self.set_limits()\n",
    "        self.slider = QtGui.QSlider(QtCore.Qt.Horizontal, parent=parent)\n",
    "        self.slider.setTickPosition(QtGui.QSlider.TicksAbove)\n",
    "        \n",
    "        self.slider.setTickInterval((self.smax - self.smin) / 10. * self.scale)\n",
    "        self.slider.setMinimum(self.smin * self.scale)\n",
    "        if (self.smax - self.width) > 0:\n",
    "            self.slider.setMaximum((self.smax - self.width) * self.scale)\n",
    "        else:\n",
    "            self.slider.setMaximum(self.smax * self.scale)\n",
    "        self.slider.setSingleStep(self.width * self.scale/4.)\n",
    "        self.slider.setPageStep(self.page_step * self.width * self.scale)\n",
    "        self.slider.setValue(self.smin * self.scale) # set the initial position\n",
    "        self.slider.valueChanged.connect(self.xpos_changed)\n",
    "        parent.addWidget(self.slider)\n",
    "    \n",
    "    def xpos_changed(self, pos):\n",
    "        pos /= self.scale\n",
    "        slice_ = (pos, pos + self.duration_steps, self.decimate)\n",
    "        self.update(slice_)\n",
    "\n",
    "    def setup_plot(self):\n",
    "        fig, AX = self.fig, self.AX\n",
    "        plt.subplots_adjust(left=0.05, right=0.93, top=0.95, bottom=0.09,\n",
    "                            wspace=0.05)\n",
    "        self.title_suffix = \"Total: %.1f s, Visualized: %.2f ms\" % (\n",
    "                            self.S.t_step*self.S.n_samples, self.duration*1e3)\n",
    "\n",
    "        AX[1].set_xlabel(\"z (um)\")\n",
    "        AX[0].set_xlabel(\"x (um)\")\n",
    "        AX[0].set_ylabel(\"y (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",
    "        lines_xy, lines_zy = [], []\n",
    "        for ip in self.particles:\n",
    "            plot_kwargs = dict(ls='', marker='o', mew=0, ms=2, \n",
    "                               alpha=0.5, label='P%d' % ip)\n",
    "            l_xy, = AX[0].plot([], [], **plot_kwargs)\n",
    "            l_zy, = AX[1].plot([], [], color=l_xy.get_color(), **plot_kwargs)\n",
    "            lines_xy.append(l_xy)\n",
    "            lines_zy.append(l_zy)\n",
    "        self.lines_xy = lines_xy\n",
    "        self.lines_zy = lines_zy\n",
    "    \n",
    "        if len(self.particles) <= 20:\n",
    "            AX[1].legend(bbox_to_anchor=(1.01, 1), loc=2, borderaxespad=0.)\n",
    "\n",
    "        AX[0].set_xlim(-4, 4)\n",
    "        AX[0].set_ylim(-4, 4)\n",
    "        AX[1].set_xlim(-4, 4)\n",
    "        for ax in AX:\n",
    "            ax.autoscale(False)\n",
    "        fig.canvas.draw()\n",
    "        self.background = fig.canvas.copy_from_bbox(fig.bbox)\n",
    "        self.title = plt.suptitle('')\n",
    "        \n",
    "    def update(self, slice_=None):\n",
    "        if slice_ is None:\n",
    "            slice_ = (0, self.duration_steps, self.decimate)\n",
    "        slice_ = slice(*slice_)\n",
    "        pos = self.position[:, :, slice_]\n",
    "        \n",
    "        self.fig.canvas.restore_region(self.background)\n",
    "        for ip, l_xy, l_zy in zip(self.particles, self.lines_xy, self.lines_zy):\n",
    "            x, y, z = pos[ip, 0], pos[ip, 1], pos[ip, 2]\n",
    "            l_xy.set_data(x*1e6, y*1e6)\n",
    "            l_zy.set_data(z*1e6, y*1e6)\n",
    "            self.AX[0].draw_artist(l_xy)\n",
    "            self.AX[1].draw_artist(l_zy)\n",
    "        time = slice_.start * self.S.t_step\n",
    "        self.title.set_text(\"t = %5.1fs %s\" % (time, self.title_suffix))\n",
    "        self.fig.draw_artist(self.title)\n",
    "        self.fig.canvas.blit(self.fig.bbox)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "p = TrackPlotter(S, duration=0.02, decimate=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import matplotlib.gridspec as gridspec\n",
    "from matplotlib import cm\n",
    "from PyQt4 import QtGui, QtCore\n",
    "\n",
    "class TrackEmPlotter:\n",
    "    def __init__(self, S, particles=None, duration=0.01, decimate=100, scale=1, page_step=100):\n",
    "        if particles is None:\n",
    "            particles = range(S.num_particles)\n",
    "        self.particles = particles\n",
    "        self.S = S\n",
    "        self.position = S.position\n",
    "        self.duration = duration\n",
    "        self.decimate = decimate\n",
    "        self.scale = scale\n",
    "        self.page_step = page_step\n",
    "        self.duration_steps = duration//S.t_step\n",
    "        self.num_points = self.duration_steps//decimate\n",
    "        self.create_figure()\n",
    "    \n",
    "        # Retrive the QMainWindow used by current figure and add a toolbar\n",
    "        # to host the new widgets\n",
    "        QMainWin = self.fig.canvas.parent()\n",
    "        toolbar = QtGui.QToolBar(QMainWin)\n",
    "        #QMainWin.addToolBar(QtCore.Qt.BottomToolBarArea, toolbar)\n",
    "        QMainWin.addToolBar(QtCore.Qt.TopToolBarArea, toolbar)\n",
    "        # Create the slider and spinbox for x-axis scrolling in toolbar\n",
    "        self.set_slider(toolbar)\n",
    "        self.init_plot()\n",
    "        self.update()\n",
    "\n",
    "    def create_figure(self):\n",
    "        self.fig = plt.figure(figsize=(11,8))\n",
    "        gs = gridspec.GridSpec(2, 2, height_ratios=[2, 1])\n",
    "        ax1 = self.fig.add_subplot(gs[0, 0])\n",
    "        ax2 = self.fig.add_subplot(gs[0, 1], sharey=ax1)\n",
    "        plt.setp(ax2.get_yticklabels(), visible=False)\n",
    "        ax3 = self.fig.add_subplot(gs[1, :])\n",
    "        self.AX = [ax1, ax2, ax3]\n",
    "\n",
    "    def set_limits(self):\n",
    "        self.smin = 0\n",
    "        self.smax = self.position.shape[-1]\n",
    "        self.width = self.num_points*self.decimate\n",
    "\n",
    "    def set_slider(self, parent):\n",
    "        self.set_limits()\n",
    "        self.slider = QtGui.QSlider(QtCore.Qt.Horizontal, parent=parent)\n",
    "        self.slider.setTickPosition(QtGui.QSlider.TicksAbove)\n",
    "        \n",
    "        self.slider.setTickInterval((self.smax - self.smin) / 10. * self.scale)\n",
    "        self.slider.setMinimum(self.smin * self.scale)\n",
    "        if (self.smax - self.width) > 0:\n",
    "            self.slider.setMaximum((self.smax - self.width) * self.scale)\n",
    "        else:\n",
    "            self.slider.setMaximum(self.smax * self.scale)\n",
    "        self.slider.setSingleStep(self.width * self.scale/4.)\n",
    "        self.slider.setPageStep(self.page_step * self.width * self.scale)\n",
    "        self.slider.setValue(self.smin * self.scale) # set the initial position\n",
    "        self.slider.valueChanged.connect(self.xpos_changed)\n",
    "        parent.addWidget(self.slider)\n",
    "    \n",
    "    def xpos_changed(self, pos):\n",
    "        pos /= self.scale\n",
    "        slice_ = (pos, pos + self.duration_steps, self.decimate)\n",
    "        self.update(slice_)\n",
    "    def plot_psf(self):\n",
    "        psf = np.concatenate((S.psf.hdata[:, :0:-1], S.psf.hdata), axis=1)\n",
    "        cmap = cm.YlGnBu\n",
    "        cmap.set_under(alpha=0)\n",
    "        kwargs = dict(interpolation='nearest', cmap=cmap, vmin=1e-1, zorder=1)\n",
    "        self.AX[1].imshow(psf.T, extent=(-6, 6, -4, 4), **kwargs)\n",
    "        \n",
    "        x = np.concatenate((-self.S.psf.xi[:0:-1], self.S.psf.xi))*1e-6\n",
    "        X, Y = np.meshgrid(x, x)\n",
    "        R = np.sqrt(X**2 + Y**2)\n",
    "        psf_xy = S.psf.eval_xz(R, 0)\n",
    "        self.AX[0].imshow(psf_xy, extent=(-4, 4, -4, 4), **kwargs)\n",
    "        \n",
    "    def init_plot(self):\n",
    "        fig, AX = self.fig, self.AX\n",
    "        plt.subplots_adjust(left=0.05, right=0.93, top=0.95, bottom=0.09,\n",
    "                            wspace=0.05)\n",
    "        self.title_suffix = \"Total: %.1f s, Visualized: %.2f ms\" % (\n",
    "                            self.S.t_step*self.S.n_samples, self.duration*1e3)\n",
    "\n",
    "        AX[1].set_xlabel(\"z (um)\")\n",
    "        AX[0].set_xlabel(\"x (um)\")\n",
    "        AX[0].set_ylabel(\"y (um)\")\n",
    "        AX[2].set_xlabel(\"Time (ms)\")\n",
    "        self.plot_psf()\n",
    "        \n",
    "        pal = sns.color_palette()\n",
    "        colors = pal[1:]\n",
    "        par_counts = [c[1] for c in self.S.particles.diffusion_coeff_counts]\n",
    "        \n",
    "        em_y = np.zeros(self.num_points)\n",
    "        em_x = np.arange(self.num_points)*self.S.t_step*self.decimate*1e3\n",
    "        lines_xy, lines_zy, lines_em = [], [], []\n",
    "        for ip in self.particles:\n",
    "            color = colors[0] if ip < par_counts[0] else colors[1]\n",
    "            plot_kwargs = dict(ls='', marker='o', mew=0, ms=2, color=color,\n",
    "                               alpha=0.5, label='P%d' % ip)\n",
    "            l_xy, = AX[0].plot([], [], **plot_kwargs)\n",
    "            l_zy, = AX[1].plot([], [], **plot_kwargs)\n",
    "            em_kw = dict(alpha=0.8, ls='-', color=color)\n",
    "            l_em, = AX[2].plot(em_x, em_y, **em_kw)\n",
    "            lines_xy.append(l_xy)\n",
    "            lines_zy.append(l_zy)\n",
    "            lines_em.append(l_em)\n",
    "        self.lines_xy = lines_xy\n",
    "        self.lines_zy = lines_zy\n",
    "        self.lines_em = lines_em\n",
    "    \n",
    "        if len(self.particles) <= 20:\n",
    "            AX[1].legend(bbox_to_anchor=(1.01, 1), loc=2, borderaxespad=0.)\n",
    "\n",
    "        AX[0].set_xlim(-4, 4)\n",
    "        AX[0].set_ylim(-4, 4)\n",
    "        AX[1].set_xlim(-4, 4)\n",
    "        AX[2].set_ylim(0, 1)\n",
    "        for ax in AX:\n",
    "            ax.autoscale(False)\n",
    "            ax.set_axisbelow(True)\n",
    "        fig.canvas.draw()\n",
    "        self.background = fig.canvas.copy_from_bbox(fig.bbox)\n",
    "        self.title = plt.suptitle('')\n",
    "        \n",
    "    def update(self, slice_=None):\n",
    "        if slice_ is None:\n",
    "            slice_ = (0, self.duration_steps, self.decimate)\n",
    "        slice_ = slice(*slice_)\n",
    "        assert (slice_.stop - slice_.start)//self.decimate == self.num_points\n",
    "        pos = self.position[:, :, slice_]\n",
    "        emission = self.S.emission[:, slice_]\n",
    "        \n",
    "        self.fig.canvas.restore_region(self.background)\n",
    "        for ip, l_xy, l_zy, l_em in zip(self.particles, \n",
    "                                        self.lines_xy, self.lines_zy, \n",
    "                                        self.lines_em):\n",
    "            x, y, z = pos[ip, 0], pos[ip, 1], pos[ip, 2]\n",
    "            l_xy.set_data(x*1e6, y*1e6)\n",
    "            l_zy.set_data(z*1e6, y*1e6)\n",
    "            l_em.set_ydata(emission[ip])\n",
    "            self.AX[0].draw_artist(l_xy)\n",
    "            self.AX[1].draw_artist(l_zy)\n",
    "            self.AX[2].draw_artist(l_em)\n",
    "        time = slice_.start * self.S.t_step\n",
    "        self.title.set_text(\"t = %5.1fs %s\" % (time, self.title_suffix))\n",
    "        self.fig.draw_artist(self.title)\n",
    "        self.fig.canvas.blit(self.fig.bbox)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "p = TrackEmPlotter(S)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": 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
}