{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ASSET: Analysis of Sequences of Synchronous Events in Massively Parallel Spike Trains\n",
    "\n",
    "The tutorial is based on the paper \\[[1](https://journals.plos.org/ploscompbiol/article?id=10.1371/journal.pcbi.1004939)\\] and demonstrates a method of finding patterns of synchronous spike times (synfire chains) which cannot be revealed by measuring neuronal firing rates only.\n",
    "\n",
    "In this tutorial, we use 50 neurons per group in 10 successive groups of a synfire chain embedded in a balanced network simulation. For more information about the data and ASSET algorithm, refer to \\[[1](https://journals.plos.org/ploscompbiol/article?id=10.1371/journal.pcbi.1004939)\\].\n",
    "\n",
    "#### References\n",
    "\n",
    "    [1] Torre E, Canova C, Denker M, Gerstein G, Helias M, Grün S (2016) ASSET: Analysis of Sequences of     Synchronous Events in Massively Parallel Spike Trains. PLoS Comput Biol 12(7): e1004939. https://doi.org/10.1371/journal.pcbi.1004939"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Explore the data and postulate the problem\n",
    "\n",
    "We start by importing the required packages, setting up matplotlib and loading the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import quantities as pq\n",
    "from elephant import asset\n",
    "from elephant.datasets import load_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "\n",
    "plt.style.use('dark_background')\n",
    "\n",
    "plt.rcParams['figure.autolayout'] = False\n",
    "plt.rcParams['figure.figsize'] = 20, 12\n",
    "plt.rcParams['axes.labelsize'] = 18\n",
    "plt.rcParams['axes.titlesize'] = 20\n",
    "plt.rcParams['font.size'] = 14\n",
    "plt.rcParams['lines.linewidth'] = 1.0\n",
    "plt.rcParams['lines.markersize'] = 8\n",
    "plt.rcParams['legend.fontsize'] = 14\n",
    "\n",
    "plt.rcParams['text.latex.preamble'] = r\"\\usepackage{subdepth}, \\usepackage{type1cm}\"\n",
    "plt.rcParams['mathtext.fontset'] = 'cm'\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we load the data with the spiking activity from the 500 neurons used in this tutorial. \n",
    "\n",
    "The spike trains are stored in a dataset saved with the NIX file format. The data in the NIX file is represented as a `neo.Block` with one `neo.Segment` inside, which contains raw `neo.SpikeTrain`s, each corresponding to the activity of one neuron.\n",
    "\n",
    "The NIX file will be automatically downloaded from https://datasets.python-elephant.org and the `neo.Segment` object loaded into the variable `segment`. Spike trains are accessible with the `.spiketrains` attribute (`segment.spiketrains`).\n",
    "\n",
    "For more information on `neo.Block`, `neo.Segment`, and `neo.SpikeTrain` refer to https://neo.readthedocs.io/en/stable/api_reference.html."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the `neo.Segment` containing all the data for the ASSET analysis\n",
    "segment = load_data('asset')\n",
    "\n",
    "# Access the spike trains with the activity of the 500 neurons\n",
    "spiketrains = segment.spiketrains"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The raster plot below shows the activity of all 500 neurons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.eventplot([st.magnitude for st in spiketrains], linewidths=5, linelengths=5)\n",
    "plt.xlabel('time [ms]')\n",
    "plt.ylabel('neuron id')\n",
    "plt.title('Raw spikes')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even though we see an increase of the firing rate, we cannot find a propagating activity just by looking at the raster plot above.\n",
    "\n",
    "We want to find a permuation of rows (neurons) in `spiketrains` such that a pattern (synfire chain) appears.\n",
    "The true unknown permutation is stored in the `segment.annotations['spiketrain_ordering']`. **The goal is to recreate this permutation from raw data with the statistical method ASSET.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Applying ASSET\n",
    "\n",
    "### 2.1. Intersection matrix\n",
    "\n",
    "The first step is to compute the intersection matrix `imat`, the `(i,j)` entry of which represents the number of neurons spiking both at bins `i` and `j` after the binning is applied. The resultant symmetric matrix `imat` shows one off-diagonal synfire chain pattern (see the picture below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2.1.1) create ASSET analysis object\n",
    "# hint: try different bin sizes, e.g. bin_size=2.5, 3.5, 4.0 ms\n",
    "asset_obj = asset.ASSET(spiketrains, bin_size=3*pq.ms)\n",
    "\n",
    "# 2.1.2) compute the intersection matrix\n",
    "imat = asset_obj.intersection_matrix()\n",
    "plt.matshow(imat)\n",
    "plt.colorbar();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2. Analytical probability matrix\n",
    "\n",
    "The second step is to estimate the probability $P_{null}$ that non-zero entries in `imat` occurred by chance. The resultant `pmat` matrix is defined as the probability of having strictly fewer coincident spikes at bins `i` and `j` strictly than the observed overlap (`imat`) under the null hypothesis of independence of the input spike trains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pmat = asset_obj.probability_matrix_analytical(imat, kernel_width=50*pq.ms)\n",
    "plt.matshow(pmat)\n",
    "plt.colorbar();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3. Joint probability matrix\n",
    "\n",
    "The third step is postprocessing of the analytical probability matrix `pmat`, obtained from the previous step. Centered at each (i,j) entry of `pmat` matrix, we apply a diagonal kernel with shape `filter_shape` and select the top `nr_largest` probabilities of (i,j) neighborhood (defined by `filter_shape`), and compute the significance of these `nr_largest` joint neighbor probabilities. The resultant `jmat` matrix is a \"dilated\" version of `imat`.\n",
    "\n",
    "This step is most time consuming. If you have PyCUDA or PyOpenCL installed, set `ELEPHANT_USE_CUDA` or `ELEPHANT_USE_OPENCL` environment flag to `1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['ELEPHANT_USE_OPENCL'] = '0'\n",
    "# try different filter_shapes, e.g. filter_shape=(7,3)\n",
    "jmat = asset_obj.joint_probability_matrix(pmat, filter_shape=(11, 3), n_largest=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.matshow(jmat)\n",
    "plt.colorbar();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4. Mask matrix\n",
    "\n",
    "After setting significance thresholds $\\alpha_P$ and $\\alpha_J$ for the corresponding matrices $P$ (probability matrix `pmat`) and $J$ (joint probability matrix `jmat`), we check the entries for significance. The resultant boolean mask matrix `mmat` is then defined as\n",
    "\n",
    "$$\n",
    "M_{ij} = 1_{P_{ij} > \\alpha_P} \\cdot 1_{J_{ij} > \\alpha_J}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hint: try different alphas for pmat and jmat\n",
    "# hint: try alphas in range [0.99, 1-1e-6]\n",
    "# hint: you can call 'asset.ASSET.mask_matrices(...)' without creating the asset_obj\n",
    "alpha = .99\n",
    "mmat = asset_obj.mask_matrices([pmat, jmat], [alpha, alpha])\n",
    "plt.matshow(mmat);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.5. Find clusters in the mask matrix\n",
    "\n",
    "Each entry (i,j) of the mask matrix $M$ from the previous step is assigned to a cluster id. A cluster is constrained to have at least `min_neighbors` number of associated elements with at most `eps` intra-distance (within the group). The cluster index, or the (i,j) entry of the resultant `cmat` matrix, is:\n",
    "\n",
    "1. a positive int (cluster id), if (i,j) is part of a cluster;\n",
    "2. `0`, if $M_{ij}$ is non-positive;\n",
    "3. `-1`, if the element (i,j) does not belong to any cluster.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hint: you can call asset.ASSET.cluster_matrix_entries(...) without creating the asset_obj\n",
    "cmat = asset_obj.cluster_matrix_entries(mmat, max_distance=11, min_neighbors=10, stretch=5)\n",
    "plt.matshow(cmat)\n",
    "plt.colorbar();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.6. Sequences of synchronous events\n",
    "\n",
    "Given the input spike trains, two arrays of bin edges from step 2.2 and the clustered intersection matrix `cmat` from step 2.5, extract the sequences of synchronous events (synfire chains).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sses = asset_obj.extract_synchronous_events(cmat)\n",
    "sses.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_id = 1\n",
    "cluster_chain = []\n",
    "for chain in sses[cluster_id].values():\n",
    "    cluster_chain.extend(chain)\n",
    "\n",
    "_, st_indices = np.unique(cluster_chain, return_index=True)\n",
    "st_indices = np.take(cluster_chain, np.sort(st_indices))\n",
    "\n",
    "reordered_sts = [spiketrains[idx] for idx in st_indices]\n",
    "spiketrains_not_a_pattern = [spiketrains[idx] for idx in range(len(spiketrains))\n",
    "                      if idx not in st_indices]\n",
    "reordered_sts.extend(spiketrains_not_a_pattern)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.eventplot([st.magnitude for st in reordered_sts], linewidths=5, linelengths=5)\n",
    "plt.xlabel('time [ms]')\n",
    "plt.ylabel('reordered neuron id')\n",
    "plt.title('Reconstructed ordering of the neurons (y-axis) with synfire chains');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### With the sequences of synchronous events `sses` we found a permutation of the spiketrains that reveals the synfire chain as in the ground truth ordering in the last Figure, shown below.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ordering_true = segment.annotations['spiketrain_ordering']\n",
    "spiketrains_ordered = [spiketrains[idx] for idx in ordering_true]\n",
    "\n",
    "plt.figure()\n",
    "plt.eventplot([st.magnitude for st in spiketrains_ordered], linewidths=5, linelengths=5)\n",
    "plt.xlabel('time [ms]')\n",
    "plt.ylabel('neuron id')\n",
    "plt.title('True (unknown) ordering of the neurons (y-axis)')\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.8.10"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}