notebooks/rastermap_widefield.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_widefield.ipynb)\n",
    "\n",
    "# Rastermap sorting of widefield neural activity\n",
    "\n",
    "We will use a widefield imaging recording from mouse cortex from [Musall*, Kaufman* et al 2019](https://doi.org/10.1038/s41593-019-0502-4). The full dataset is available [here](https://labshare.cshl.edu/shares/library/repository/38599/). The imaging was collected while mice performed a decision-making task, the imaging rate was 30 Hz. Instead of single neurons here, each sample is a voxel. We filtered out voxels on the edge of the imaging that had low variance. We will run Rastermap using the principal components of the recording because the full recording is too large to fit in memory."
   ]
  },
  {
   "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/5d8q7)."
   ]
  },
  {
   "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=\"widefield\")\n",
    "\n",
    "dat = np.load(filename)\n",
    "\n",
    "# U0 is voxels by number of components (left singular vectors)\n",
    "# Vsv is time by number of components (right singular vectors weighted by the singular values)\n",
    "# (each timepoint is 33 ms)\n",
    "# sv are the singular values\n",
    "U0, sv, Vsv = dat[\"U0\"], dat[\"sv\"], dat[\"Vsv\"]\n",
    "n_voxels, n_time = U0.shape[0], Vsv.shape[0]\n",
    "print(f\"{n_voxels} voxels by {n_time} timepoints\")\n",
    "\n",
    "# XY position of each neuron in the recording\n",
    "xpos, ypos = dat[\"xpos\"], dat[\"ypos\"]\n",
    "\n",
    "# load the stimulus times\n",
    "stims = dat[\"stims\"]\n",
    "reward_times = dat[\"reward_times\"]\n",
    "reward_color = [0,0.5,0]\n",
    "fcolor = np.array([[0,0.5,1], [0,1,1], [1,0,0], \n",
    "                [1,0.5,0], [0.8,0.5,0.7]])\n",
    "\n",
    "\n",
    "# load the behavior and task variables\n",
    "regressors = dat[\"regressors\"]\n",
    "behav_idx = dat[\"behav_idx\"] # which regressors are behavioral (not task variables)"
   ]
  },
  {
   "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=U0.shape[1], # number of PCs (precomputed in U0 and Vsv)\n",
    "                  locality=0.5, # locality in sorting is low here to get more global sorting (this is a value from 0-1)\n",
    "                  time_lag_window=10, # use future timepoints to compute correlation\n",
    "                  grid_upsample=10, # default value, 10 is good for large recordings\n",
    "                  ).fit(Usv = U0 * sv, # left singular vectors weighted by the singular values\n",
    "                        Vsv = Vsv)     # right singular vectors weighted by the singular values\n",
    "                                                          \n",
    "y = model.embedding # neurons x 1\n",
    "isort = model.isort\n",
    "Vsv_sub = model.Vsv # these are the PCs across time with the mean across voxels subtracted"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create superneurons from Rastermap -- we sort the data and then sum over neighboring voxels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nbin = 200 # number of voxels to bin over \n",
    "U_sn = utils.bin1d(U0[isort], bin_size=nbin, axis=0) # bin over voxel axis\n",
    "sn = U_sn @ Vsv_sub.T\n",
    "sn = zscore(sn, axis=1)"
   ]
  },
  {
   "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": [
    "from matplotlib import patches\n",
    "\n",
    "def plot_events(ax, stims, reward_times, xmin, xmax):\n",
    "    \"\"\" shade stim times and plot reward times \"\"\"\n",
    "    nn = sn.shape[0]\n",
    "    for k in range(4):\n",
    "        starts = stims[stims[:,1]==k,0].copy()\n",
    "        starts = starts[np.logical_and(starts>=xmin, starts<xmax)]\n",
    "        starts -= xmin\n",
    "        for n in range(len(starts)):\n",
    "            start = starts[n]+1\n",
    "            width = 1.6*30\n",
    "            # add stimulus patch\n",
    "            ax.add_patch(\n",
    "                    patches.Rectangle(xy=(start, 0), width=width,\n",
    "                            height=nn, facecolor=fcolor[k], \n",
    "                            edgecolor=None, alpha=0.2))\n",
    "            \n",
    "    for reward_time in reward_times:\n",
    "        if reward_time >= xmin and reward_time < xmax:\n",
    "            ax.plot((reward_time-xmin)*np.ones(2), [0, nn],\n",
    "                        color=reward_color, lw=2)\n",
    "    ax.set_ylim([0,nn])\n",
    "    ax.invert_yaxis()\n",
    "\n",
    "# timepoints to visualize\n",
    "xmin=7810\n",
    "xmax=8720\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 superneuron activity\n",
    "ax = plt.subplot(grid[1:, :-1])\n",
    "ax.imshow(sn[:, xmin:xmax], cmap=\"gray_r\", vmin=0, vmax=1.1, aspect=\"auto\")\n",
    "ax.set_xlabel(\"time\")\n",
    "ax.set_ylabel(\"superneurons\")\n",
    "plot_events(ax, stims, reward_times, xmin, xmax)\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\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Color the voxels by their position in the rastermap:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(5, 5))\n",
    "plt.scatter(xpos, ypos, 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": [
    "### Neural activity prediction from behavior\n",
    "\n",
    "We can use ridge regression to predict superneuron activity from the task variables and the behavioral variables. We will first install the code package with the ridge regression and helper functions implemented:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install neuropop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Predict from both task and behavioral variables, and predict from behavioral variables alone:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from neuropop import linear_prediction\n",
    "import torch\n",
    "\n",
    "ve, _, sn_pred, itest = linear_prediction.prediction_wrapper(regressors, sn.T, lam=1e4, \n",
    "                                                             device=torch.device(\"cpu\"))\n",
    "sn_pred = sn_pred.T\n",
    "itest = itest.flatten()\n",
    "print(f\"prediction from task and behavioral variables: \\t{ve:.3f}\")\n",
    "\n",
    "ve, _, sn_pred_beh, itest = linear_prediction.prediction_wrapper(regressors[:,behav_idx], \n",
    "                                                                 sn.T, lam=1e4,\n",
    "                                                                 device=torch.device(\"cpu\"))\n",
    "itest = itest.flatten()\n",
    "sn_pred_beh = sn_pred_beh.T\n",
    "print(f\"prediction from only behavioral variables: \\t{ve:.3f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the predictions on test data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# timepoints to visualize\n",
    "xmin=820\n",
    "xmax=1730\n",
    "\n",
    "# make figure with grid for easy plotting\n",
    "fig = plt.figure(figsize=(10,14), dpi=200)\n",
    "grid = plt.GridSpec(4, 1, figure=fig, wspace = 0.05, hspace = 0.3)\n",
    "\n",
    "titles = [\"(i) widefield imaging during a decision-making task\",\n",
    "                \"(ii) prediction of activity from task and behavior variables\",\n",
    "                \"(iii) prediction of activity from behavior variables only\",\n",
    "                \"(iv) difference between (ii) and (iii)\"]\n",
    "\n",
    "for j in range(4):\n",
    "    # plot superneuron activity\n",
    "    ax = plt.subplot(grid[j])\n",
    "    if j==0:\n",
    "        sp = sn[:, itest[xmin:xmax]]\n",
    "    elif j==1:\n",
    "        sp = sn_pred[:, xmin:xmax]\n",
    "    elif j==2:\n",
    "        sp = sn_pred_beh[:, xmin:xmax]\n",
    "    else:\n",
    "        sp = sn_pred[:, xmin:xmax] - sn_pred_beh[:, xmin:xmax]\n",
    "\n",
    "    ax.imshow(sp, cmap=\"gray_r\" if j<3 else \"RdBu_r\", \n",
    "                vmin=-1.1*(j==3), vmax=1.1, aspect=\"auto\")\n",
    "    if j==3:\n",
    "        ax.set_xlabel(\"time\")\n",
    "    ax.set_ylabel(\"superneurons\")\n",
    "    ax.set_title(titles[j])\n",
    "    if j<3:\n",
    "        plot_events(ax, stims, reward_times, itest[xmin], itest[xmax])\n"
   ]
  },
  {
   "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
}