notebooks/rastermap_zebrafish.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MouseLand/rastermap/blob/main/notebooks/rastermap_zebrafish.ipynb)\n",
    "\n",
    "# Rastermap sorting of zebrafish neural activity\n",
    "\n",
    "We will use a zebrafish wholebrain neural activity recording from [Chen*, Mu*, Hu*, Kuan* et al 2018](https://doi.org/10.1016/j.neuron.2018.09.042). The full dataset is available [here](https://doi.org/10.25378/janelia.7272617). The recordings were performed at a rate of 2.1 Hz. We took the neurons with the highest variance signals and deconvolved them to reduce long timescales in the data from the calcium sensor."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we will install the required packages, if not already installed. If on google colab, it will require you to click the \"RESTART RUNTIME\" button because we are updating numpy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install numpy>=1.24 # (required for google colab)\n",
    "!pip install rastermap \n",
    "!pip install matplotlib"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load data and import libraries\n",
    "\n",
    "If not already downloaded, the following cell will automatically download the processed data stored [here](https://osf.io/2w8pa)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "# importing rastermap\n",
    "# (this will be slow the first time since it is compiling the numba functions)\n",
    "from rastermap import Rastermap, utils\n",
    "from scipy.stats import zscore\n",
    "\n",
    "# download spontaneous activity\n",
    "filename = utils.download_data(data_type=\"fish\")\n",
    "\n",
    "dat = np.load(filename)\n",
    "\n",
    "# spks is neurons by time\n",
    "# (each timepoint is 476 ms)\n",
    "spks = dat[\"spks\"]\n",
    "n_neurons, n_time = spks.shape\n",
    "print(f\"{n_neurons} neurons by {n_time} timepoints\")\n",
    "\n",
    "# zscore activity (each neuron activity trace is then mean 0 and standard-deviation 1)\n",
    "spks = zscore(spks, axis=1)\n",
    "\n",
    "# XYZ position of each neuron in the recording\n",
    "xyz = dat[\"xyz\"]\n",
    "\n",
    "# load the stimulus times\n",
    "stims = dat[\"stims\"]\n",
    "# stim colors\n",
    "fcolor = np.zeros((stims.max()+1, 4))\n",
    "fcolor[0:3] = np.array([[0., 0.5, 1.0, 1.0], [1.0, 0.0, 1.0, 1.0], \n",
    "                        [1., 1., 1., 1.]])\n",
    "fcolor[8:12] = np.array([[1.0,0.0,0,1],\n",
    "                        [0.8,1.0,0,1], [0,0,1,1], [0,1,1,1]])\n",
    "\n",
    "# load the fictive swimming\n",
    "swimming = dat[\"swimming\"]\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run Rastermap\n",
    "\n",
    "Let's sort the single neurons with Rastermap, with clustering and upsampling:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Rastermap(n_clusters=100, # number of clusters to compute\n",
    "                  n_PCs=200, # number of PCs to use\n",
    "                  locality=0.1, # locality in sorting is low here to get more global sorting (this is a value from 0-1)\n",
    "                  time_lag_window=5, # use future timepoints to compute correlation\n",
    "                  grid_upsample=10, # default value, 10 is good for large recordings\n",
    "                ).fit(spks)\n",
    "y = model.embedding # neurons x 1\n",
    "isort = model.isort"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create superneurons from Rastermap -- we sort the data and then sum over neighboring neurons:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nbin = 50 # number of neurons to bin over \n",
    "sn = utils.bin1d(spks[isort], bin_size=nbin, axis=0) # bin over neuron axis"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualization\n",
    "\n",
    "Use the Rastermap sorting to visualize the neural activity (see Figure 4 from the paper for the stimulus legend):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# timepoints to visualize\n",
    "xmin = 5700\n",
    "xmax = 7860\n",
    "\n",
    "# make figure with grid for easy plotting\n",
    "fig = plt.figure(figsize=(12,6), dpi=200)\n",
    "grid = plt.GridSpec(9, 20, figure=fig, wspace = 0.05, hspace = 0.3)\n",
    "\n",
    "# plot swimming speed\n",
    "ax = plt.subplot(grid[0, :-1])\n",
    "ax.plot(swimming[xmin:xmax, 0], color=fcolor[11])\n",
    "ax.plot(swimming[xmin:xmax, 1], color=fcolor[10])\n",
    "ax.set_xlim([0, xmax-xmin])\n",
    "ax.axis(\"off\")\n",
    "ax.set_title(\"swimming speed\")\n",
    "\n",
    "# plot superneuron activity\n",
    "ax = plt.subplot(grid[1:, :-1])\n",
    "ax.imshow(sn[:, xmin:xmax], cmap=\"gray_r\", vmin=0, vmax=0.8, aspect=\"auto\")\n",
    "ax.set_xlabel(\"time\")\n",
    "ax.set_ylabel(\"superneurons\")\n",
    "\n",
    "# color time periods by stimuli\n",
    "from matplotlib import patches\n",
    "nn = sn.shape[0]\n",
    "stims_t = stims[xmin:xmax]\n",
    "starts = np.nonzero(np.diff(stims_t))\n",
    "starts = np.append(np.array([0]), starts)\n",
    "starts = np.append(starts, np.array([len(stims_t)-1]))\n",
    "for n in range(len(starts)-1):\n",
    "    start = starts[n]+1\n",
    "    stype = stims_t[start]\n",
    "    if stype!=3:\n",
    "        width = starts[n+1] - start + min(0, start)\n",
    "        start = max(0, start)\n",
    "        ax.add_patch(\n",
    "                patches.Rectangle(xy=(start, 0), width=width,\n",
    "                            height=nn, facecolor=fcolor[stype], \n",
    "                            edgecolor=None, alpha=0.15*(stype!=2)))\n",
    "\n",
    "\n",
    "ax = plt.subplot(grid[1:, -1])\n",
    "ax.imshow(np.arange(0, len(sn))[:,np.newaxis], cmap=\"gist_ncar\", aspect=\"auto\")\n",
    "ax.axis(\"off\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Color the neurons by their position in the rastermap:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(5, 5))\n",
    "plt.scatter(xyz[:,1], xyz[:,0], s=1, c=y, cmap=\"gist_ncar\", alpha=0.25)\n",
    "plt.xlabel('X position')\n",
    "plt.ylabel('Y position')\n",
    "plt.axis(\"square\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also divide the rastermap into sections to more easily visualize spatial relations (as in Figure 4):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ny, nx = 3, 6\n",
    "nxy = nx * ny\n",
    "\n",
    "# divide into nxy sections\n",
    "nb = len(isort) // nxy\n",
    "colors = plt.get_cmap(\"gist_ncar\")(np.linspace(0, 0.9, nxy))\n",
    "\n",
    "# make figure with grid for easy plotting\n",
    "fig = plt.figure(figsize=(12,6), dpi=200)\n",
    "\n",
    "grid = plt.GridSpec(ny, nx, figure=fig, wspace = 0.25, hspace = 0.1)\n",
    "for j in range(nx):\n",
    "    for k in range(ny):\n",
    "        ax = plt.subplot(grid[k,j])\n",
    "        # plot all neurons\n",
    "        subsample = 25\n",
    "        ax.scatter(xyz[:,1][::subsample], xyz[:,0][::subsample], s=2, alpha=1, \n",
    "                    color=0.9*np.ones(3), rasterized=True)\n",
    "        ip = j + k*nx\n",
    "        ix = isort[ip*nb : (ip+1)*nb]\n",
    "        subsample = 1\n",
    "        ax.scatter(xyz[ix,1][::subsample], xyz[ix,0][::subsample],\n",
    "                    s=0.5, alpha=0.3, color=colors[ip])\n",
    "        ax.axis(\"off\")\n",
    "        ax.axis(\"square\")\n",
    "        ax.text(0.1,0,str(ip+1), transform=ax.transAxes, ha=\"right\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Settings\n",
    "\n",
    "You can see all the rastermap settings with `Rastermap?`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Rastermap?"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Outputs\n",
    "\n",
    "All the attributes assigned to the Rastermap `model` are listed with `Rastermap.fit?`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Rastermap.fit?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.16 ('rastermap')",
   "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.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "998540cc2fc2836a46e99cd3ca3c37c375205941b23fd1eb4b203c48f2be758f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}