visualize_ridge_plot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2efef589",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib qt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "final-marine",
   "metadata": {},
   "source": [
    "# Ridge plot\n",
    "Ridge plots show each trace from your calcium session offset a little bit on the y-axis. They can be a useful way to get an initial big-picture view of your data. In this notebook, we will recreate the following plot: \n",
    "\n",
    "![Example](data/ridge_example.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a19a3f3",
   "metadata": {},
   "source": [
    "## Load data and build cnmf object\n",
    "Start cluster and build object from `hdf5` file saved from data already analyzed in the CNMFE demo in caiman:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a5e0bf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import caiman as cm\n",
    "from caiman.source_extraction.cnmf.cnmf import load_CNMF\n",
    "from caiman.source_extraction import cnmf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1afd9042",
   "metadata": {},
   "outputs": [],
   "source": [
    "cnmfe_results_path = r'./data/cnmfe_demo.hdf5'\n",
    "num_cpus = 2 # use however many you need: just leave a few so your RAM stays happy\n",
    "# note if a cluster already exists it will be closed so a new session will be opened\n",
    "if 'dview' in locals():  # locals contains list of current local variables\n",
    "    cm.stop_server(dview=dview)\n",
    "c, dview, n_processes = cm.cluster.setup_cluster(backend='local', \n",
    "                                                 n_processes=num_cpus, \n",
    "                                                 single_thread=False,\n",
    "                                                 ignore_preexisting=True)\n",
    "#Number of nodes in cluster \n",
    "print(f\"Cluster has {n_processes} processes in the pool\");\n",
    "cnm = load_CNMF(cnmfe_results_path, \n",
    "                n_processes, \n",
    "                dview=dview)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2133b468",
   "metadata": {},
   "source": [
    "## Extract basic info from estimates object\n",
    "Get denoised traces from cnmf object. Also, reconstruct array of frame times. \n",
    "\n",
    "Note we will focus on the \"good\" components in `cnm.estimates.idx_components` (for more on interpreting the properties in `cnm.estimates`, see https://caiman.readthedocs.io/en/master/Getting_Started.html#result-interpretation)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c69613a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "good_inds = cnm.estimates.idx_components\n",
    "denoised_traces = cnm.estimates.C[good_inds,:]  #num_components x num_frames\n",
    "\n",
    "num_components = denoised_traces.shape[0]\n",
    "num_samples = denoised_traces.shape[1]\n",
    "frame_rate = cnm.params.data['fr']\n",
    "sampling_pd = 1/frame_rate\n",
    "num_frames = num_samples\n",
    "frame_times = np.arange(0, sampling_pd*num_samples, sampling_pd)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02bc661b",
   "metadata": {},
   "source": [
    "## Define ridge plot\n",
    "This is a pretty straightforward function: just plot each trace, slightly offset from its neighbor. As discussed below, the main wrinkles have to do with how to handle the spacing of the traces and the ytick labels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f88b050",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ridge_plot(array, \n",
    "               times,\n",
    "               trace_spacing=5,\n",
    "               ytick_spacing=10,\n",
    "               title=None,\n",
    "               color='black',\n",
    "               alpha=0.5,\n",
    "               line_width=1, \n",
    "               ax = None):\n",
    "    \"\"\"\n",
    "    Plot ridge plot of all rows in array, given x values.\n",
    "    Inputs:\n",
    "        array: num_components x num_times array of traces\n",
    "        times: 1-d array \n",
    "        trace_spacing (scalar): distance between each plot on y axis\n",
    "        ytick_spacing (int): period between ytick labels (every ytick_spacing traces)\n",
    "        color: line color (r,g,b) or standard matplotlib color string\n",
    "        alpha (scalar): alpha for each trace, if you want them more see-through when density is high.\n",
    "        line_width (scalar): how wide?\n",
    "        ax: axes object if you want to draw in pre-existing axes (None if you want new axes)\n",
    "    \n",
    "    Outputs:\n",
    "        ax: axes object with lines drawn\n",
    "    \"\"\"\n",
    "    num_traces = array.shape[0]\n",
    "    num_yticks = int(1+num_traces//ytick_spacing)\n",
    "    \n",
    "    # set y position of each trace\n",
    "    y_position_traces = np.linspace(0, num_traces*trace_spacing, num=num_traces) \n",
    "    \n",
    "    # set y tick properties \n",
    "    y_ticks = np.linspace(0, (num_traces-1)*trace_spacing, num=num_yticks)\n",
    "    y_tick_labels = np.arange(0, num_traces+2*ytick_spacing, ytick_spacing, dtype=np.uint8) # +2*y_tick_spacing just for insurance\n",
    "    y_tick_labels = y_tick_labels[:num_yticks]\n",
    "    \n",
    "    if ax is None:\n",
    "        f, ax = plt.subplots()\n",
    "    for ind, line in enumerate(array):\n",
    "        ax.plot(times, \n",
    "                line+y_position_traces[ind], \n",
    "                color=color,\n",
    "                alpha=alpha,\n",
    "                linewidth=line_width)\n",
    "    # only show left/bottom axis lines\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ylims = ax.get_ylim()\n",
    "    # set ylimits to make it pretty (this could use some tweaking probably)\n",
    "    ax.set_ylim(0.1*ylims[0], ylims[1]-0.05*ylims[1])\n",
    "    ax.set_yticks(y_ticks)\n",
    "    ax.set_yticklabels(y_tick_labels)\n",
    "    plt.autoscale(enable=True, axis='x', tight=True)\n",
    "    return ax\n",
    "print('done')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ebcc2b7",
   "metadata": {},
   "source": [
    "## Using ridge_plot()\n",
    "As mentioned above, there are two spacing parameters that are the only real wrinkles here: `trace_spacing` and `ytick_spacing`.  \n",
    "\n",
    "The `ytick_spacing` parameter sets how frequently to display ytick labels (labeling every 10th trace is the default). The `trace_spacing` parameter determines how far apart traces are from each other on the y axis. The best value will depend on the scale of your traces.  You will likely want to get a sense of the maximum values in your array and scale it based on that, as I do in the following example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d776661",
   "metadata": {},
   "outputs": [],
   "source": [
    "component_spacing = int(np.max(denoised_traces)*0.1)\n",
    "print(f\"Traces will be spread {component_spacing} units apart\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfaf1127",
   "metadata": {},
   "outputs": [],
   "source": [
    "f, traces_ax = plt.subplots(figsize=(19,10))\n",
    "traces_ax = ridge_plot(denoised_traces, \n",
    "                       frame_times, \n",
    "                       color='black',\n",
    "                       trace_spacing=component_spacing, \n",
    "                       ytick_spacing=5,\n",
    "                       alpha=0.6,\n",
    "                       line_width=0.75,\n",
    "                       ax=traces_ax)\n",
    "traces_ax.set_xlabel(\"Time (s)\", fontsize=12)\n",
    "traces_ax.set_ylabel(\"Component #\", fontsize=12)\n",
    "traces_ax.set_title(\"Ridge Plot Example\", fontsize=16)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e734413",
   "metadata": {},
   "source": [
    "If this was a plotting library we would likely just wrap that trace spacing calculation into `ridge_plot()`, but for now it's just a parameter you can explore."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63c0c2d7",
   "metadata": {},
   "source": [
    "## Doing things with your ridge plot\n",
    "Once you have a ridge plot you like, there are a few simple things you can do. \n",
    "\n",
    "For instance, you can add a rectangular highlight (e.g., when a stimulus was presented):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fbb65ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "rect_overlay = traces_ax.axvspan(25, 45, color='yellow', alpha=0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f6afb4b",
   "metadata": {},
   "source": [
    "You can remove the rectangle easily if you want to build animations with scrolling rectangles that sweep across the ridge plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9361c9e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "rect_overlay.remove()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b88476b0",
   "metadata": {},
   "source": [
    "If you ever want to build an oscilloscope-like animation, you will need an efficient way to remove lines between sweeps. (Note Matplotlib will not be fast enough for real-time display of data, but if you just want to create a high-quality animation for a presentation, it is fine).\n",
    "\n",
    "The following function will remove all the lines from your ridge plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bed19d93",
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_lines(ax):\n",
    "    for line in ax.get_lines(): # ax.lines:\n",
    "        line.remove()\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68b449da",
   "metadata": {},
   "outputs": [],
   "source": [
    "traces_ax = remove_lines(traces_ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29f03e4f",
   "metadata": {},
   "source": [
    "It is left as an exercise to write a function to add lines back to the axes object. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56bca7ea",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}