notebooks/tutorial.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/tutorial.ipynb)\n",
    "\n",
    "\n",
    "# What to do with thousands of neurons in cortex?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use a spontaneous activity recording from [Syeda et al, 2023](https://www.biorxiv.org/content/10.1101/2022.11.03.515121v1.abstract). We recorded 34,086 neurons from mouse sensorimotor cortex for 2+ hours using two-photon calcium imaging at a rate of 3.2Hz. FYI to make the download of the dataset faster, we are analyzing only the first half of the recording. During the recording, the mouse was free to run on an air floating ball, and we recorded the mouse face with a camera at a rate of 50Hz and tracked keypoints on the mouse face."
   ]
  },
  {
   "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. Also if you are on google colab, select the GPU runtime if you want to fit the neural network to predict neural activity from behaviors:\n",
    "**Runtime > Change runtime type > Hardware accelerator = GPU**"
   ]
  },
  {
   "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\n",
    "!pip install neuropop"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Loading and preparing. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import the Python libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os # os stands for \"operating system\" and includes read/write routines etc. \n",
    "import numpy as np # by far the most used library for everyday computation\n",
    "from scipy import stats # here we import a whole sub-library of stats functions\n",
    "from matplotlib import pyplot as plt # all of our plotting is done with plt\n",
    "%matplotlib inline \n",
    "# %matplotlib notebook # if you need to zoom into a figure, this is the \"interactive\" mode of IPython\n",
    "\n",
    "from rastermap import utils, Rastermap\n",
    "\n",
    "# download the spontaneous activity recording\n",
    "filename = utils.download_data(data_type=\"spont2\")\n",
    "\n",
    "# load the neural data and the processed behavioral data\n",
    "dat = np.load(filename)\n",
    "spks, U, sv, V = dat[\"spks\"], dat[\"U\"], dat[\"sv\"], dat[\"V\"]\n",
    "xpos, ypos = dat[\"xpos\"], dat[\"ypos\"]\n",
    "tcam, tneural = dat[\"tcam\"], dat[\"tneural\"]\n",
    "run, beh, beh_names = dat[\"run\"], dat[\"beh\"], dat[\"beh_names\"]\n",
    "\n",
    "n_neurons, n_time = spks.shape\n",
    "\n",
    "print(f\"{n_neurons} neurons by {n_time} timepoints\")\n",
    "\n",
    "# we will z-score each neuron so that the activity is standard deviation 1 and mean 0 for each neuron\n",
    "spks = stats.zscore(spks, axis=1)\n",
    "\n",
    "# colors for the behaviors\n",
    "kp_colors = np.array([[0.55,0.55,0.55], [0.,0.,1],\n",
    "                      [0.8,0,0], [1.,0.4,0.2],\n",
    "                      [0,0.6,0.4], [0.2,1,0.5],\n",
    "                      ])\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Visualizing neural activity with Rastermap"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first look at the positions of the neurons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# POSITIONS OF ALL NEURONS\n",
    "plt.figure(figsize=(4, 4))\n",
    "plt.scatter(ypos, xpos, s = 1)\n",
    "plt.xlabel('X position (um)')\n",
    "plt.ylabel('Y position (um)')\n",
    "plt.axis(\"square\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's run Rastermap. Rastermap re-arranges neurons in the raster plot based on similarity of activity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rastermap import Rastermap\n",
    "\n",
    "### run rastermap\n",
    "model = Rastermap(n_clusters=100, n_PCs=128, locality=0.6,\n",
    "                  time_lag_window=5).fit(spks)\n",
    "cc_nodes = model.cc.copy()\n",
    "isort = model.isort"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create superneurons from Rastermap: sort the data and then sum over neighboring neurons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nbin = 200 # number of neurons to bin over \n",
    "ndiv = (n_neurons//nbin) * nbin\n",
    "# group sorted matrix into rows of length nbin\n",
    "sn = spks[isort][:ndiv].reshape(ndiv//nbin, nbin, -1)\n",
    "# take mean over neurons in a bin\n",
    "sn = sn.mean(axis=1)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize the activity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# timepoints to visualize\n",
    "xmin = 0 \n",
    "xmax = xmin + 1000\n",
    "\n",
    "fig = plt.figure(figsize=(12,6), dpi=200)\n",
    "grid = plt.GridSpec(9, 10, figure=fig, wspace = 0.05, hspace = 0.3)\n",
    "\n",
    "# plot running speed\n",
    "ax = plt.subplot(grid[0, :-1])\n",
    "ax.plot(run[xmin:xmax], color=kp_colors[0])\n",
    "ax.set_xlim([0, xmax-xmin])\n",
    "ax.axis(\"off\")\n",
    "ax.set_title(\"running speed\", color=kp_colors[0])\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",
    "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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "color the neurons by their position in the rastermap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# POSITIONS OF ALL NEURONS\n",
    "plt.figure(figsize=(4, 4))\n",
    "plt.scatter(ypos, xpos, s = 1, c=model.embedding[:,0], cmap=\"gist_ncar\")\n",
    "plt.xlabel('X position (um)')\n",
    "plt.ylabel('Y position (um)')\n",
    "plt.axis(\"square\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Dimensionality reduction: what are the dominant patterns of activity?\n",
    "\n",
    "Let's compute the top PCs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.decomposition import TruncatedSVD\n",
    "\n",
    "# this function returns the left singular vectors scaled by the singular values\n",
    "Vsv = TruncatedSVD(n_components = 128).fit_transform(spks.T)\n",
    "\n",
    "# compute the other singular vectors\n",
    "U = spks @ (Vsv / (Vsv**2).sum(axis=0)**0.5)\n",
    "U /= (U**2).sum(axis=0)**0.5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "plot the PCs in time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12,6), dpi=200)\n",
    "grid = plt.GridSpec(9, 1, figure=fig, hspace = 0.4)\n",
    "\n",
    "# plot running speed\n",
    "ax = plt.subplot(grid[0, 0])\n",
    "ax.plot(run[xmin:xmax], color=kp_colors[0])\n",
    "ax.set_xlim([0, xmax-xmin])\n",
    "ax.axis(\"off\")\n",
    "ax.set_title(\"running speed\", color=kp_colors[0])\n",
    "\n",
    "pc_colors = plt.get_cmap(\"viridis\")(np.linspace(0,0.9,8))\n",
    "for j in range(8):\n",
    "    ax = plt.subplot(grid[j+1])\n",
    "    ax.plot(Vsv[xmin:xmax, j], color=pc_colors[j])\n",
    "    ax.set_xlim([0, xmax-xmin])\n",
    "    ax.axis(\"off\")\n",
    "    ax.set_title(f\"PC {j+1}\", color=pc_colors[j])\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "we will predict the PCs from the behavior (many fewer of them)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Behavioral analysis\n",
    "\n",
    "We will figure out the sorts of behaviors the superneurons care about!\n",
    "\n",
    "The behavioral video is at 50Hz while the neural data is at 3.2 Hz.\n",
    "\n",
    "Here are the behaviors we tracked:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12,6), dpi=200)\n",
    "grid = plt.GridSpec(len(beh_names), 1, figure=fig, hspace = 0.4)\n",
    "\n",
    "n_beh = beh.shape[1]\n",
    "for j in range(n_beh):\n",
    "    ax = plt.subplot(grid[j])\n",
    "    ax.plot(beh[17*xmin:17*xmax, j], color=kp_colors[j])\n",
    "    ax.set_xlim([0, 17*(xmax-xmin)])\n",
    "    ax.axis(\"off\")\n",
    "    ax.set_title(beh_names[j], color=kp_colors[j])\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the sampling rate and timing of neural activity and the behavior are different, we will first downsample the behavior data to the timestamps of the neural activity. \n",
    "\n",
    "The easiest way to do this is with interpolation: we know when in time each behavior frame happened (`tcam`), and then we sample it at each time the neural activity happened (`tneural`). But there are a lot of fast things going on in the behavior, so to get the average over timepoints at the neural activity time, we smooth the behavioral data first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.ndimage import gaussian_filter1d # here we import a smoothing function\n",
    "from scipy.interpolate import interp1d # importing an interpolation function\n",
    "\n",
    "# initialize empty matrix\n",
    "beh_ds = np.zeros((len(tneural), n_beh), \"float32\")\n",
    "\n",
    "for j in range(n_beh):\n",
    "    # filter the data\n",
    "    # (smoothing scale proportional to difference in sampling rate)\n",
    "    bsmooth = gaussian_filter1d(beh[:,j], 50/3.2) \n",
    "    # interpolate\n",
    "    f = interp1d(tcam, bsmooth)\n",
    "    beh_ds[:,j] = f(tneural)\n",
    "\n",
    "print(beh_ds.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the traces again\n",
    "fig = plt.figure(figsize=(12,6), dpi=200)\n",
    "grid = plt.GridSpec(len(beh_names), 1, figure=fig, hspace = 0.4)\n",
    "for j in range(n_beh):\n",
    "    ax = plt.subplot(grid[j])\n",
    "    ax.plot(beh_ds[xmin:xmax, j], color=kp_colors[j])\n",
    "    ax.set_xlim([0, (xmax-xmin)])\n",
    "    ax.axis(\"off\")\n",
    "    ax.set_title(beh_names[j], color=kp_colors[j]) \n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, to do prediction, we have to do a train-test split. You always want to train your model on a subset of data and test its performance on another set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# split into train-test\n",
    "# * use interleaved segments *\n",
    "n_segs = 20\n",
    "n_len  = n_time / n_segs\n",
    "sinds = np.linspace(0, n_time - n_len, n_segs).astype(int)\n",
    "itest = (sinds[:,np.newaxis] + np.arange(0, n_len*0.25, 1, int)).flatten()\n",
    "itrain = np.ones(n_time, \"bool\")\n",
    "itrain[itest] = 0\n",
    "itest = ~itrain\n",
    "\n",
    "plt.plot(itrain)\n",
    "plt.title(\"train times\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Can you determine how I did the split above? Why might I have split into segments rather than randomly interleaving time-points?\n",
    "\n",
    "## Linear regression prediction\n",
    "\n",
    "Use linear regression to perform the prediction, predict PCs $Y$ using behaviors $X$:\n",
    "\n",
    "$$ A = (X_\\text{train}^\\top X_\\text{train})^{-1} (X_\\text{train}^\\top Y_\\text{train})$$\n",
    "\n",
    "$X$ is behavioral components by time, $Y$ is neural components by time. If you want to regularize the linear regression:\n",
    "\n",
    "$$ A = (X_\\text{train}^\\top X_\\text{train} + \\lambda I)^{-1} (X_\\text{train}^\\top Y_\\text{train})$$\n",
    "\n",
    "Then the prediction on time points is:\n",
    "\n",
    "$$ \\hat Y_\\text{test} = X_\\text{test} A $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## predict using behavior traces\n",
    "# regularized linear regression from behavior to neural PCs\n",
    "\n",
    "XtX = beh_ds[itrain].T @ beh_ds[itrain]\n",
    "XtY = beh_ds[itrain].T @ Vsv[itrain]\n",
    "lam = 1e1 # regularizer\n",
    "XtX += lam * np.eye(n_beh)\n",
    "\n",
    "# regression matrix\n",
    "A = np.linalg.solve(XtX, XtY)\n",
    "\n",
    "# prediction\n",
    "Vpred = beh_ds[itest] @ A\n",
    "\n",
    "# variance explained per PC\n",
    "residual = ((Vpred - Vsv[itest])**2).sum(axis=0)\n",
    "varexp_PC = 1 - residual / (Vsv[itest]**2).sum(axis=0)\n",
    "\n",
    "# overall varexp\n",
    "varexp = 1 - residual.sum() / (Vsv[itest]**2).sum()\n",
    "\n",
    "print(f\"overall variance explained = {varexp}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "plot PCs and prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12,6), dpi=200)\n",
    "grid = plt.GridSpec(8, 1, figure=fig, hspace = 0.4)\n",
    "\n",
    "for j in range(8):\n",
    "    ax = plt.subplot(grid[j])\n",
    "    ax.plot(Vsv[itest][xmin:xmax, j], color=pc_colors[j])\n",
    "    ax.plot(Vpred[xmin:xmax, j], color=\"k\", linestyle=\"--\")\n",
    "    ax.set_xlim([0, xmax-xmin])\n",
    "    ax.axis(\"off\")\n",
    "    ax.set_title(f\"PC {j+1}, varexp = {varexp_PC[j]:.2f}\", color=pc_colors[j])\n",
    "    if j==0:\n",
    "        ax.legend([\"PC\", \"prediction\"], loc=\"upper right\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predicting with a 1D convolution layer\n",
    "\n",
    "There are finer temportal features in the behavioral features that we aren't capturing by smoothing and using the smoothed traces.\n",
    "\n",
    "Instead we can learn the temporal features by using a 1D convolution layer with various filters -- called kernels.\n",
    "\n",
    "See below a nice illustration of a convolution from this [webpage](https://e2eml.school/convolution_one_d.html). This kernel is a gaussian, you can see how it smooths the data. But a neural network can learn whatever kernels help with prediction.\n",
    "\n",
    "![conv_gif](https://e2eml.school/images/conv1d/stride_1.gif)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use **pytorch** for this. It allows us to write a network and then it automatically finds the gradients to optimize it to our data.\n",
    "\n",
    "We will create a model with a linear input layer, a one-dimensional convolutional layer, and a linear output layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from neuropop import nn_prediction\n",
    "import torch\n",
    "\n",
    "# ideally we have a GPU we can use (\"cuda\" option)\n",
    "device = torch.device(\"cuda\")\n",
    "\n",
    "# declare the model\n",
    "pred_model = nn_prediction.PredictionNetwork(n_in=beh.shape[-1], n_kp=22, identity=False, \n",
    "                                             n_filt=10, n_latents=0,\n",
    "                                             n_out=Vsv.shape[-1], n_core_layers=1,\n",
    "                                             relu_wavelets=False, relu_latents=False)\n",
    "# put model on the GPU\n",
    "pred_model.to(device);\n",
    "\n",
    "print(pred_model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should see the layers, can you figure out what some of these variables mean?\n",
    "\n",
    "Now let's train the model with a function we have that includes the data splitting and the gradient descent. We will see the variance explained on each iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred_all, ve_all, itest = pred_model.train_model(beh, Vsv, tcam, tneural, delay=-1,\n",
    "                                                    learning_rate=1e-3, n_iter=400,\n",
    "                                                    device=device, verbose=True)\n",
    "Vpred_conv = y_pred_all\n",
    "itest = itest.flatten() # we run using batches of data so we now flatten it\n",
    "\n",
    "# variance explained per PC\n",
    "residual = ((Vpred_conv - Vsv[itest])**2).sum(axis=0)\n",
    "varexp_PC_conv = 1 - residual / (Vsv[itest]**2).sum(axis=0)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This fit better! Let's see what the prediction looks like, it seems like the higher PCs are better captured especially."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12,6), dpi=200)\n",
    "grid = plt.GridSpec(8, 1, figure=fig, hspace = 0.4)\n",
    "\n",
    "for j in range(8):\n",
    "    ax = plt.subplot(grid[j])\n",
    "    ax.plot(Vsv[itest][xmin:xmax, j], color=pc_colors[j])\n",
    "    ax.plot(Vpred_conv[xmin:xmax, j], color=\"k\", linestyle=\"--\")\n",
    "    ax.set_xlim([0, xmax-xmin])\n",
    "    ax.axis(\"off\")\n",
    "    ax.set_title(f\"PC {j+1}, varexp = {varexp_PC_conv[j]:.2f}\", color=pc_colors[j])\n",
    "    if j==0:\n",
    "        ax.legend([\"PC\", \"prediction\"], loc=\"upper right\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Receptive fields of superneurons\n",
    "\n",
    "We can use this linear model to estimate the receptive fields of the superneurons.\n",
    "\n",
    "First we need to get the superneuron PC weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sort and bin PCs for maxstim estimation\n",
    "U_sn = U[isort][:ndiv].reshape(ndiv//nbin, nbin, -1).mean(axis=1)\n",
    "\n",
    "print(U_sn.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ms_model = nn_prediction.MaxStimModel(pred_model)\n",
    "ms_model.requires_grad = False\n",
    "u = torch.from_numpy(U_sn).to(device)\n",
    "u.requires_grad = False\n",
    "xr = ms_model.train_batch(u, n_iter=200, learning_rate=1e-2)\n",
    "rfs = xr.detach().cpu().numpy()\n",
    "\n",
    "print(rfs.shape)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize a subset of receptive fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(12,6), dpi=200)\n",
    "grid = plt.GridSpec(12, 21, figure=fig, wspace = 0.05, hspace = 0.0)\n",
    "\n",
    "vmax = 8\n",
    "ks = np.linspace(5, len(U_sn)-5, 12*3).astype(\"int\")\n",
    "for i, k in enumerate(ks):\n",
    "    for j in range(n_beh):\n",
    "        ax = plt.subplot(grid[i%12, j + 6*(i//12) + (i//12)])\n",
    "        ax.plot(rfs[k, 100:-100, j], color=kp_colors[j])\n",
    "        ax.set_ylim([-vmax, vmax])\n",
    "        ax.axis(\"off\")\n",
    "        if i==0:\n",
    "            ax.set_title(beh_names[j], color=kp_colors[j], rotation=45)\n",
    "            \n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the receptive fields with the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# timepoints to visualize\n",
    "xmin = 0 \n",
    "xmax = xmin + 1000\n",
    "\n",
    "fig = plt.figure(figsize=(12,6), dpi=200)\n",
    "grid = plt.GridSpec(9, 18, figure=fig, wspace = 0.35, hspace = 0.3)\n",
    "\n",
    "# plot running speed\n",
    "ax = plt.subplot(grid[0, :12])\n",
    "ax.plot(run[xmin:xmax], color=kp_colors[0])\n",
    "ax.set_xlim([0, xmax-xmin])\n",
    "ax.axis(\"off\")\n",
    "ax.set_title(\"running speed\", color=kp_colors[0])\n",
    "\n",
    "# plot superneuron activity\n",
    "ax = plt.subplot(grid[1:, :12])\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",
    "for j in range(n_beh):\n",
    "    ax = plt.subplot(grid[1:, j+12])\n",
    "    ax.imshow(rfs[:,100:-100,j], aspect=\"auto\", vmin=-vmax, vmax=vmax, cmap=\"RdBu_r\")\n",
    "    ax.axis(\"off\")\n",
    "    ax.set_title(beh_names[j], color=kp_colors[j], rotation=45)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Nonlinear prediction with 1D convolutions\n",
    "\n",
    "We can put non-linearities in our neural network models to better model non-linear aspects of the data. We also add another layer to make it more complex. The network below is the default network from the Facemap paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_model = nn_prediction.PredictionNetwork(n_in=beh.shape[-1], n_kp=22, \n",
    "                                             n_out=Vsv.shape[-1])\n",
    "# put model on the GPU\n",
    "pred_model.to(device);\n",
    "\n",
    "print(pred_model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "y_pred_all, ve_all, itest = pred_model.train_model(beh, Vsv, tcam, tneural, delay=-1,\n",
    "                                                        learning_rate=1e-3, n_iter=400,\n",
    "                                                    device=device, verbose=True)\n",
    "Vpred_nl = y_pred_all\n",
    "itest = itest.flatten() # we run using batches of data so we now flatten it\n",
    "\n",
    "# variance explained per PC\n",
    "residual = ((Vpred_nl - Vsv[itest])**2).sum(axis=0)\n",
    "varexp_PC_nl = 1 - residual / (Vsv[itest]**2).sum(axis=0)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the prediction. Remember we need to project using the PCs into the neuron space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn_pred = U_sn @ Vpred_nl.T\n",
    "\n",
    "\n",
    "fig = plt.figure(figsize=(12,12), dpi=200)\n",
    "grid = plt.GridSpec(13, 1, figure=fig, wspace = 0.35, hspace = 0.6)\n",
    "\n",
    "# plot running speed\n",
    "ax = plt.subplot(grid[0, 0])\n",
    "ax.plot(run[itest][xmin:xmax], color=kp_colors[0])\n",
    "ax.set_xlim([0, xmax-xmin])\n",
    "ax.axis(\"off\")\n",
    "ax.set_title(\"running speed\", color=kp_colors[0])\n",
    "\n",
    "# plot superneuron activity\n",
    "ax = plt.subplot(grid[1:7, 0])\n",
    "ax.imshow(sn[:, itest][:, xmin:xmax], cmap=\"gray_r\", vmin=0, vmax=0.85, aspect=\"auto\")\n",
    "ax.set_ylabel(\"superneurons\")\n",
    "ax.set_xticks([])\n",
    "ax.set_title(\"neural activity\")\n",
    "\n",
    "# plot superneuron prediction\n",
    "ax = plt.subplot(grid[7:, 0])\n",
    "ax.imshow(sn_pred[:, xmin:xmax], cmap=\"gray_r\", vmin=0, vmax=0.85, aspect=\"auto\")\n",
    "ax.set_xlabel(\"time\")\n",
    "ax.set_ylabel(\"superneurons\")\n",
    "ax.set_title(\"behavior prediction\")"
   ]
  }
 ],
 "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"
  },
  "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
  },
  "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()) "
    }
   },
   "position": {
    "height": "908px",
    "left": "1679px",
    "right": "20px",
    "top": "112px",
    "width": "479px"
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": true
  },
  "vscode": {
   "interpreter": {
    "hash": "998540cc2fc2836a46e99cd3ca3c37c375205941b23fd1eb4b203c48f2be758f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}