lectures/05_deconv_spike_sorting.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spike Sorting by Deconvolution\n",
    "\n",
    "The last chapter framed spike sorting as a matrix factorization problem &mdash; specifically, a _semi-nonnegative matrix factorization (semi-NMF)_ problem.  However, that model only makes sense under the simplifying assumption that the raw voltage is downsampled to $\\sim$ 500Hz.\n",
    "Otherwise, spike waveforms would be spread over many time bins.\n",
    "\n",
    "In this chapter we'll relax that assumption and develop a more realistic model using _convolutional_ matrix factorization. This model is inspired by Kilosort, a state-of-the-art spike sorting algorithm {cite}`pachitariu2023solving`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convolutional model\n",
    "\n",
    "Like before, let $x_{n,t}$ denote the voltage on channel $c$ and time sample $t$, but now consider the data at its native resolution of around 30 kHz (i.e. voltage is sampled every $\\sim 0.03$ ms. At this sampling frequency, a spike waveform typically at least 60-90 time steps. \n",
    "\n",
    "Let $\\mathbf{W}_k \\in \\mathbb{R}^{N \\times D}$ denote a **waveform**. In this model, it is a **matrix** for each neuron, where $D$ denotes the number of time steps that a spike waveform persists in the voltage recording. \n",
    "\n",
    "Let $w_{k,n,d}$ denote the entries of waveform $\\mathbf{W}_k$, and let $\\mathbf{W} = \\{\\mathbf{W}_k\\}_{k=1}^K$ be shorthand for the set of waveforms for all $K$ neurons. Let $\\mathbf{A} \\in \\mathbb{R}_+^{T \\times K}$ denote the matrix of spike amplitudes, as before. Now, $a_{k,t} = 1$ denotes the **start** of a unit-amplitude spike with waveform $\\mathbf{W}_k$ at time $t$.\n",
    "\n",
    "The new model's likelihood is,\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "p(\\mathbf{X} \\mid \\mathbf{W}, \\mathbf{A})\n",
    "&= \\prod_{n=1}^N \\prod_{t=1}^T \\mathcal{N} \\left( x_{n,t} \\, \\bigg| \\, \\sum_{k=1}^K \\sum_{d=1}^D a_{k, t-d} \\, w_{k,n,d}, \\sigma^2 \\right) \\\\\n",
    "&= \\prod_{n=1}^N \\prod_{t=1}^T \\mathcal{N} \\left( x_{n,t} \\, \\bigg| \\, \\sum_{k=1}^K [\\mathbf{a}_{k} \\circledast \\mathbf{w}_{k,n}]_t, \\sigma^2 \\right) \\\\\n",
    "&= \\prod_{t=1}^T \\mathcal{N} \\left( \\mathbf{x}_{t} \\, \\bigg| \\, \\sum_{k=1}^K [\\mathbf{a}_{k} \\circledast \\mathbf{W}_{k}]_{:,t}, \\sigma^2 \\mathbf{I} \\right)\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "where $\\circledast$ denotes the discrete time **convolution**."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convolution and cross-correlation\n",
    "\n",
    "The discrete time **convolution** of a signal $\\mathbf{x}$ with a filter $\\mathbf{f}$ is defined as \n",
    "\n",
    "$$\n",
    "[\\mathbf{x} \\circledast \\mathbf{f}]_t = \\sum_{d = -\\infty}^{\\infty} x_{t - d} f_d.\n",
    "$$\n",
    "\n",
    "Of course, in practice we're dealing with finite-length vectors $\\mathbf{x} \\in \\mathbb{R}^T$ and filters $\\mathbf{f} \\in \\mathbb{R}^D$, so we need to decide how to deal with boundary effects. One possibility is to pad $\\mathbf{x}$ with $D-1$ zeros; another is to return only the \"valid\" section of the convolution. Yet another is to assume the signal is periodic, so that the convolution wraps around when the index $t-d$ is negative. That is called a **circular convolution**.\n",
    "\n",
    "The **cross-correlation** is \n",
    "\n",
    "$$\n",
    "[\\mathbf{x} \\star f]_t = \\sum_{d = -\\infty}^{\\infty} x_{t + d} f_d.\n",
    "$$ \n",
    "\n",
    "Thus, **convolution is equivalent to cross-correlation with a reversed filter**.  \n",
    "\n",
    "```{warning}\n",
    "Unfortunately, the definition of cross-correlation is not unique; our definition consistent with Numpy's `correlate` function, but it's what [Wikipedia](https://en.wikipedia.org/wiki/Cross-correlation) would call $[f \\star x]_t$ instead (note the order is swapped). \n",
    "\n",
    "To make matters more confusing, the \"convolution\" operation performed by most neural network libraries is actually a cross-correlation (with Wikipedia's semantics).\n",
    "```\n",
    "\n",
    "Since cross-correlations (convolutions in machine learning parlance) are such fundamental building blocks of modern neural networks, libraries like PyTorch have flexible APIs for performing a variety of types of convolutions. For example, with [`torch.nn.functional.conv1d`](https://pytorch.org/docs/stable/generated/torch.nn.functional.conv1d.html) you can cross-correlate a bank of 1D signal $\\mathbf{X} \\in \\mathbb{R}^{N \\times T}$ with a bank of filters $\\mathbf{F} \\in \\mathbb{R}^{N \\times D}$ by varying the number of `in_channels` and `out_channels`. \n",
    "\n",
    "```{admonition} Notation\n",
    "The notation for convolutions with multiple input/output channels is less standardized. We will let $\\mathbf{Y} = \\mathbf{x} \\star \\mathbf{F}$ denote the cross-correlation of a signal $\\mathbf{x} \\in \\mathbb{R}^T$ with a bank of filters $\\mathbf{F} \\in \\mathbb{R}^{N \\times D}$, which yields a bank of outputs $\\mathbf{Y} \\in \\mathbb{R}^{N \\times T'}$ (the length depends on the padding strategy).\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scale invariance through the waveform prior\n",
    "\n",
    "Just like before, there is a scale invariance between $\\mathbf{a}_k$ and $\\mathbf{W}_k$. Last chapter, we placed constrained the waveform vector $\\mathbf{w}_k \\in \\mathbb{R}^N$ to have unit _Euclidean_ ($\\ell_2$) norm. Now that the waveforms are matrices $\\mathbf{W}_k \\in \\mathbb{R}^{N \\times D}$, the natural generalization is to constrain the _Frobenius_ norm of the waveforms,\n",
    "\n",
    "$$\n",
    "\\|\\mathbf{W}_k\\|_{\\mathrm{F}} = 1.\n",
    "$$\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Frobenius norm and the SVD\n",
    "\n",
    "The Frobenius norm can be rewritten in many ways. \n",
    "1. It is equal to the Euclidean ($\\ell_2$) norm of the _vectorized_ matrix,\n",
    "\n",
    "    $$\n",
    "    \\|\\mathbf{W}_k\\|_{\\mathrm{F}} = \\|\\mathrm{vec}(\\mathbf{W}_k)\\|_2\n",
    "    $$\n",
    "\n",
    "2. It is the norm induced by the _Frobenius inner product_ of a matrix with itself,\n",
    "\n",
    "    $$\n",
    "    \\|\\mathbf{W}_k\\|_{\\mathrm{F}} = \\sqrt{\\langle \\mathbf{W}_k, \\mathbf{W}_k \\rangle_{\\mathrm{F}} }\n",
    "    $$\n",
    "\n",
    "    where\n",
    "\n",
    "    $$\n",
    "    \\langle \\mathbf{A}, \\mathbf{B} \\rangle_{\\mathrm{F}} = \\mathrm{Tr}(\\mathbf{A}^\\top \\mathbf{B})\n",
    "    $$\n",
    "\n",
    "3. It is the Euclidean norm of the vector of singular values of the matrix. Let $\\mathbf{W}_k = \\mathbf{U}_k \\mathbf{S}_k \\mathbf{V}_k^\\top$ where $\\mathbf{U}_k$ and $\\mathbf{V}_k$ are semi-orthogonal matrices, and where $\\mathbf{S}_k = \\mathrm{diag}(\\mathbf{s}_k)$ is the diagonal matrix of singular values, $\\mathbf{s}_k = (s_{k,1}, \\ldots, s_{k,R})$. Then,\n",
    "\n",
    "    $$\n",
    "    \\begin{align*}\n",
    "    \\|\\mathbf{W}_k\\|_{\\mathrm{F}} \n",
    "    &= \\sqrt{\\mathrm{Tr}(\\mathbf{W}_k^\\top \\mathbf{W}_k)} \\\\\n",
    "    &= \\sqrt{\\mathrm{Tr}(\\mathbf{V}_k \\mathbf{S}_k \\mathbf{U}_k^\\top \\mathbf{U}_k \\mathbf{S}_k \\mathbf{V}_k^\\top)} \\\\\n",
    "    &= \\sqrt{\\mathrm{Tr}(\\mathbf{V}_k^\\top \\mathbf{V}_k \\mathbf{S}_k \\mathbf{U}_k^\\top \\mathbf{U}_k \\mathbf{S}_k )} \\\\\n",
    "    &= \\sqrt{\\mathrm{Tr}(\\mathbf{S}_k^2)} \\\\\n",
    "    &= \\|\\mathbf{s}_k\\|_2.\n",
    "    \\end{align*}\n",
    "    $$\n",
    "\n",
    "```{admonition} Singular Value Decomposition (SVD)\n",
    "Recall that the (compact) [**singular value decomposition (SVD)**](https://en.wikipedia.org/wiki/Singular_value_decomposition) of a real valued matrix $\\mathbf{W} \\in \\mathbb{R}^{N \\times D}$ is a factorization of the form,\n",
    "\n",
    "$$\n",
    "\\mathbf{W} = \\mathbf{U} \\mathbf{S} \\mathbf{V}^\\top\n",
    "$$\n",
    "\n",
    "where $\\mathbf{U} \\in \\mathbb{R}^{N \\times R}$ and $\\mathbf{V} \\in \\mathbb{R}^{D \\times R}$ with $R \\leq \\min\\{N, D\\}$ are real semi-orthogonal matrices ($\\mathbf{U}^\\top \\mathbf{U} = \\mathbf{I}$ and $\\mathbf{V}^\\top \\mathbf{V} = \\mathbf{I}$). The diagonal matrix $\\mathbf{S} = \\mathrm{diag}(\\mathbf{s})$ contains the **singular values** $\\mathbf{s} = (s_1, \\ldots, s_R)$. The number of nonzero singular values $R$ is the **rank** of the matrix $\\mathbf{W}$.\n",
    "\n",
    "Equivalently, the SVD can be written as a sum of outer products,\n",
    "\n",
    "$$\n",
    "\\mathbf{W} = \\sum_{r=1}^R s_r \\mathbf{u}_r \\mathbf{v}_r^\\top.\n",
    "$$\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constraining the rank of of the waveform matrices\n",
    "\n",
    "Thinking of the Frobenius norm constraint in terms of a constraint on the singular values leads to a natural extension. Rather than just constraining the norm, _constrain the rank_ of the waveform matrices as well.\n",
    "\n",
    "There are at least two reasons why this is sensible:\n",
    "\n",
    "1. Manually identified spike waveforms are well approximated as outer product of a **spatial footprint** $\\mathbf{u}_k \\in \\mathbb{S}_{N-1}$ and a **temporal profile** $\\mathbf{v}_k \\in \\mathbb{S}_{D-1}$,\n",
    "\n",
    "    $$\n",
    "    \\mathbf{W}_k \\approx \\mathbf{u}_k \\mathbf{v}_k^\\top.\n",
    "    $$\n",
    "\n",
    "    _Note that this is a rank $R=1$ matrix._\n",
    "\n",
    "2. Constraining the waveform rank can dramatically reduce the number of free parameters, which is good from a statistical estimation standpoint. For example, if we constrain the wave forms to be rank 1 then the waveforms have only $\\mathcal{O}(N + D)$ free parameters in contrast to $\\mathcal{O}(ND)$ free parameters in the full-rank model.\n",
    "\n",
    "    ```{admonition} Exercise\n",
    "    :class: tip\n",
    "    I used big-O notation because the norm constraints remove additional degrees of freedom. How many degrees of freedom do the rank-1 and full rank models truly have?\n",
    "    ```\n",
    "\n",
    "We will constrain the waveforms to be rank $R$ via a uniform prior \n",
    "\n",
    "$$\\mathbf{W}_k \\sim \\mathrm{Unif}(\\mathbb{S}_R^{N,D})$$ \n",
    "\n",
    "where \n",
    "\n",
    "$$\n",
    "\\mathbb{S}_R^{N,D} = \\left\\{\\mathbf{W}: \\mathbf{W} \\in \\mathbb{R}^{N \\times D}, \\mathrm{rank}(\\mathbf{W}) = R, \\|\\mathbf{W}\\|_{\\mathrm{F}} = 1 \\right\\}\n",
    "$$\n",
    "\n",
    "is the set of unit-norm, rank-$R$ matrices in $\\mathbb{R}^{N \\times D}$.\n",
    "\n",
    "When $R=1$, these matrices can be expressed as $\\mathbf{W}_k = \\mathbf{u}_k \\mathbf{v}_k^\\top$, where  $\\mathbf{u}_k \\in \\mathbb{S}_{N-1}$ and $\\mathbf{v}_k \\in \\mathbb{S}_{D-1}$.\n",
    "\n",
    "We will use the same exponential prior on the amplitudes as in the previous chapter."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Maximum _a posteriori_ (MAP) estimation\n",
    "\n",
    "Like before, we will fit the model by using coordinate ascent to maximize the posterior probability, which is proportional to the joint probability. Again, that will entail updating the amplitudes given the waveforms, and then the waveforms given the amplitudes. When updating the parameters for neuron $k$, the solutions will depend on the residual,\n",
    "\n",
    "$$\n",
    "\\mathbf{R} = \\mathbf{X} - \\sum_{j \\neq k} \\mathbf{a}_j \\circledast \\mathbf{W}_j.\n",
    "$$\n",
    "\n",
    "where $\\mathbf{R} \\in \\mathbb{R}^{N \\times T}$ has columns $\\mathbf{r}_{t}$ and entries $r_{n,t}$.\n",
    "\n",
    "```{warning}\n",
    "Technically, we should refer to the residual matrix as $\\mathbf{R}_k$ since it is the residual when updating that neuron, but the notation gets a bit cumbersome, and it will be clear from context.\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimizing the amplitudes\n",
    "\n",
    "As a function of the waveform $\\mathbf{a}_k$ for neuron $k$, the log joint probability is,\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\log p(\\mathbf{X}, \\mathbf{A}, \\mathbf{W})\n",
    "&= \\sum_{t=1}^T \\sum_{n=1}^N \\log \\mathcal{N}\\left(\\mathbf{r}_{n,t} \\,\\bigg|\\, [\\mathbf{a}_k \\circledast \\mathbf{w}_{k,n}]_{t}, \\sigma^2 \\right) + \\sum_{t=1}^T \\mathrm{Exp}(a_{k,t}; \\lambda) \\\\\n",
    "&= -\\frac{1}{2\\sigma^2} \\| \\mathbf{R} - \\mathbf{a}_k \\circledast \\mathbf{W}_k \\|_{\\mathrm{F}}^2 - \\sum_{t=1}^T \\lambda a_{k,t} \\\\\n",
    "&= \\underbrace{-\\frac{1}{2\\sigma^2} \\| \\mathbf{a}_k \\circledast \\mathbf{W}_k \\|_{\\mathrm{F}}^2}_{\\mathcal{L}_2(\\mathbf{a}_k)} + \\underbrace{\\frac{1}{\\sigma^2} \\langle \\mathbf{R}, \\mathbf{a}_k \\circledast \\mathbf{W}_k \\rangle_{\\mathrm{F}}}_{\\mathcal{L}_1(\\mathbf{a}_k)} - \\sum_{t=1}^T \\lambda a_{k,t}.\n",
    "\\end{align*}\n",
    "$$\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### The linear term\n",
    "Lets start by unpacking the linear term,\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\mathcal{L}_1(\\mathbf{a}_k) \n",
    "&= \\frac{1}{\\sigma^2} \\langle \\mathbf{R}, \\mathbf{a}_k \\circledast \\mathbf{W}_k \\rangle \\\\\n",
    "&= \\frac{1}{\\sigma^2} \\sum_{t=1}^T \\sum_{n=1}^N r_{n,t} [\\mathbf{a}_k \\circledast \\mathbf{w}_{k,n}]_t \\\\\n",
    "&= \\frac{1}{\\sigma^2} \\sum_{t=1}^T \\sum_{n=1}^N \\sum_{d=1}^D a_{k,t-d} r_{n,t} w_{k,n,d} \\\\\n",
    "&= \\frac{1}{\\sigma^2} \\sum_{t=1}^T a_{k,t} \\sum_{n=1}^N \\sum_{d=1}^D r_{n,t+d} w_{k,n,d} \\\\\n",
    "&= \\frac{1}{\\sigma^2} \\sum_{t=1}^T a_{k,t} [\\mathbf{R} \\star \\mathbf{W}_k]_t\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "where $\\mathbf{R} \\star \\mathbf{W}_k$ denotes a **2D cross-correlation**, which maps $\\mathbb{R}^{N \\times T} \\times \\mathbb{R}^{N \\times D} \\mapsto \\mathbb{R}^{T}$ (with appropriate padding).\n",
    "\n",
    "```{note}\n",
    "In this particular case the \"signal\" $\\mathbf{R} \\in \\mathbb{R}^{N \\times T}$ and the \"filter\" $\\mathbf{W}_k \\in \\mathbb{R}^{N \\times D}$ have the same number of rows. We can implement this 2D cross-correlation using PyTorch's 1-D convolutions by taking advantage of the in- and out-channels, as we'll see in lab.\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The quadratic term\n",
    "\n",
    "Now unpack the quadratic term,\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\mathcal{L}_2(\\mathbf{a}_k)\n",
    "&= -\\frac{1}{2\\sigma^2} \\| \\mathbf{a}_k \\circledast \\mathbf{W}_k \\|_{\\mathrm{F}}^2 \\\\\n",
    "&= -\\frac{1}{2\\sigma^2} \\sum_{n=1}^N \\sum_{t=1}^T [\\mathbf{a}_k \\circledast \\mathbf{w}_{k,n}]_{t}^2 \\\\\n",
    "&= -\\frac{1}{2\\sigma^2} \\sum_{n=1}^N \\sum_{t=1}^T \\left[\\sum_{d=1}^D a_{k,t-d} w_{k,n,d} \\right]^2 \\\\\n",
    "&= -\\frac{1}{2\\sigma^2} \\sum_{n=1}^N \\sum_{t=1}^T \\left[\\sum_{d=1}^D a_{k,t-d}^2 w_{k,n,d}^2  + 2 \\sum_{d=1}^D \\sum_{d'=1}^{d-1} a_{k,t-d} a_{k,t-d'} w_{k,n,d} w_{k,n,d'} \\right] \\\\\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "The second term has **interactions** between $a_{k,t}$ and $a_{k,t'}$ whenever $|t-t'|<D$, which makes the coordinate ascent update for the vector $\\mathbf{a}_k$ hard!\n",
    "\n",
    "However, remember that the waveform width $D$ is roughly the duration of one spike. Thus, it is highly unlikely for two spikes to occur within $D$ timesteps of each other. We will **assume that the nonzero entries in $\\mathbf{a}_k$ are separated by at least $D$ timesteps.**\n",
    "\n",
    "Under this assumption, the quadratic term reduces to,\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\mathcal{L}_2(\\mathbf{a}_k)\n",
    "&= -\\frac{1}{2\\sigma^2} \\sum_{n=1}^N \\sum_{t=1}^T \\sum_{d=1}^D a_{k,t-d}^2 w_{k,n,d}^2\\\\\n",
    "&= -\\frac{1}{2\\sigma^2} \\sum_{n=1}^N \\sum_{t=1}^T \\sum_{d=1}^D a_{k,t}^2 w_{k,n,d}^2\\\\\n",
    "&= -\\frac{1}{2\\sigma^2} \\sum_{t=1}^T a_{k,t}^2 \\|\\mathbf{W}_k\\|_{\\mathrm{F}}^2 \\\\\n",
    "&= -\\frac{1}{2\\sigma^2} \\sum_{t=1}^T a_{k,t}^2,\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "just like in the previous chapter.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Finishing the optimization\n",
    "\n",
    "We have once again reduced the coordinate update for the amplitudes to solving a bunch of independent, scalar, quadratic optimization problems subject to non-negativity constraints. For $a_{k,t}$, the problem reduces to,\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "a_{k,t}^\\star &= \\text{arg} \\, \\max_{a_{k,t} \\in \\mathbb{R}_+} \\; f(a_{k,t}) = -\\frac{\\alpha}{2} a_{k,t}^2 + \\beta a_{k,t} \n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\alpha &= \\frac{1}{\\sigma^2} \\\\\n",
    "\\beta &= \\frac{[\\mathbf{R} \\star \\mathbf{W}]_t}{\\sigma^2} - \\lambda.\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "The solution is,\n",
    "\n",
    "$$\n",
    "a_{k,t}^\\star = \\max \\left\\{ 0, \\, [\\mathbf{R} \\star \\mathbf{W}]_t - \\lambda \\sigma^2 \\right\\}.\n",
    "$$\n",
    "\n",
    "```{warning}\n",
    "Note that this solution does not guarantee that the resulting nonzero amplitudes will be separated by at least $D$ time steps! In practice, we can enforce this constraint via the following heuristic: after solving for the optimal amplitudes, use the [`scipy.signal.find_peaks`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.find_peaks.html) function to keep only a subset of nonzero amplitudes that are separated by a distance of $D$.\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimizing the waveforms\n",
    "\n",
    "As a function of the waveform $\\mathbf{W}_k$ for neuron $k$, the log joint probability is,\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\log p(\\mathbf{X}, \\mathbf{A}, \\mathbf{W})\n",
    "&= \\frac{1}{\\sigma^2} \\langle \\mathbf{R}, \\mathbf{a}_k \\circledast \\mathbf{W}_k \\rangle_{\\mathrm{F}} + c'\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "We can simplify this expression a bit by introducing notation for _windows_ of the residual matrix. Let,\n",
    "\n",
    "$$\n",
    "\\mathbf{R}_{t} = \n",
    "\\begin{bmatrix}\n",
    "r_{1,t} & \\ldots & r_{1,t+D} \\\\\n",
    "\\vdots    &        & \\vdots \\\\\n",
    "r_{n,t} & \\ldots & r_{n,t+D}\n",
    "\\end{bmatrix}.\n",
    "$$\n",
    "\n",
    "(In code, this is a slice of the residual matrix `R[:,t:t+D]`.)\n",
    "\n",
    "Once again assuming that the nonzero amplitudes are separated by at least $D$ time steps, the log probability above simplifies to,\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\log p(\\mathbf{X}, \\mathbf{A}, \\mathbf{W})\n",
    "&= \\frac{1}{2\\sigma^2} \\sum_{t=1}^T  \\langle a_{k,t} \\mathbf{R}_t, \\mathbf{W}_k \\rangle + c' \\\\\n",
    "&= \\frac{1}{2\\sigma^2} \\left \\langle \\sum_{t=1}^T  a_{k,t} \\mathbf{R}_t, \\mathbf{W}_k \\right \\rangle + c' \n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "This is analogous to the norm-constrained optimization problem for vector waveforms from the previous chapter!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solving the optimization\n",
    "\n",
    "We want to maximize this log joint probability over the space of low-rank, unit-norm matrices $\\mathbb{S}_R^{N,D}$,\n",
    "\n",
    "$$\n",
    "\\mathbf{W}_k^\\star = \\text{arg} \\, \\max_{\\mathbf{W}_k \\in \\mathbb{S}_R^{N,D}} \\left \\langle \\sum_{t=1}^T a_{k,t} \\mathbf{R}_t, \\mathbf{W}_k \\right \\rangle\n",
    "$$\n",
    "\n",
    "Such optimization problems come up frequently with dealing with low-rank approximations. \n",
    "\n",
    "Recall that when we had vector waveforms in the previous chapter, the solution was to set the waveform proportional to the weighted sum of residuals (the other vector in the inner product). Here, the solution is to set the waveform matrix \"proportional to\" the weighted sum of residual matrices by taking its SVD and renormaling the singular values. \n",
    "\n",
    "Let $\\mathbf{U} \\mathbf{S} \\mathbf{V}^\\top$ where $\\mathbf{S} = \\mathrm{diag}(\\mathbf{s})$ be the SVD of the matrix $\\sum_{t=1}^T a_{k,t} \\mathbf{R}_t$. Furthermore, assume the singular values $\\mathbf{s}= (s_1, \\ldots, s_{\\min \\{N,D\\}})$ are sorted in _descending_ order. The optimal waveform update is,\n",
    "\n",
    "$$\n",
    "\\mathbf{W}_k^\\star = \\sum_{r=1}^R \\bar{s}_r \\mathbf{u}_r \\mathbf{v}_r^\\top\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\n",
    "\\bar{s}_r = \\frac{s_r}{\\sqrt{\\sum_{r'=1}^R s_{r'}^2}}.\n",
    "$$\n",
    "\n",
    "\n",
    "```{admonition} Deriving the solution\n",
    ":class: dropdown\n",
    "\n",
    "To check this solution, consider a general low-rank approximation problem,\n",
    "\n",
    "$$\n",
    "\\text{arg} \\, \\max_{\\mathbf{W} \\in \\mathbb{S}_R^{N,D}} \\; f(\\mathbf{W}) = \\left \\langle \\mathbf{C}, \\mathbf{W} \\right \\rangle.\n",
    "$$\n",
    "\n",
    "Let $\\mathbf{U} \\mathbf{S} \\mathbf{V}^\\top$ be the SVD of $\\mathbf{C}$, as above, and let $\\tilde{\\mathbf{U}} \\tilde{\\mathbf{S}} \\tilde{\\mathbf{V}}^\\top$ be the SVD of $\\mathbf{W}$. By constraint, $\\tilde{\\mathbf{S}}$ can have only $R$ nonzero singular values.\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "f(\\mathbf{W}) \n",
    "&= \\mathrm{Tr}(\\mathbf{W}^\\top \\mathbf{C}) \\\\\n",
    "&= \\mathrm{Tr}(\\tilde{\\mathbf{V}} \\tilde{\\mathbf{S}} \\tilde{\\mathbf{U}}^\\top \\mathbf{U} \\mathbf{S} \\mathbf{V}^\\top) \\\\\n",
    "&= \\mathrm{Tr}(\\mathbf{V}^\\top \\tilde{\\mathbf{V}} \\tilde{\\mathbf{S}} \\tilde{\\mathbf{U}}^\\top \\mathbf{U} \\mathbf{S}) \\\\\n",
    "&= \\sum_{r=1}^R \\sum_{m=1}^{\\min\\{N,D\\}} \\tilde{s}_r s_m \\tilde{\\mathbf{u}}_r^\\top \\mathbf{u}_m \\tilde{\\mathbf{v}}_r^\\top \\mathbf{v}_m.\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "Since $\\mathbf{U}$, $\\mathbf{V}$, $\\tilde{\\mathbf{U}}$, and $\\tilde{\\mathbf{V}}$ are all semi-orthogonal matrices, the inner products $\\tilde{\\mathbf{u}}_r^\\top \\mathbf{u}_m$ and $\\tilde{\\mathbf{v}}_r^\\top \\mathbf{v}_m$ can be at most one, and that is achieved when $\\tilde{\\mathbf{u}}_r = \\mathbf{u}_m$. \n",
    "\n",
    "The objective is maximized when $\\tilde{\\mathbf{u}}_r = \\mathbf{u}_r$ and $\\tilde{\\mathbf{v}}_r = \\mathbf{v}_r$ for $r=1,\\ldots,R$. With these left and right singular vectors, the objective is, \n",
    "\n",
    "$$\n",
    "f(\\mathbf{W}) = \\sum_{r=1}^R s_r \\tilde{s}_r = \\langle \\mathbf{s}_{:R}, \\tilde{\\mathbf{s}} \\rangle.\n",
    "$$\n",
    "\n",
    "Since $\\tilde{\\mathbf{s}}$ is constrained to be unit-norm, this is maximized when \n",
    "\n",
    "$$\n",
    "\\tilde{\\mathbf{s}} = \\frac{\\mathbf{s}_{:R}}{\\|\\mathbf{s}_{:R}\\|_2}.\n",
    "$$\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## More efficient computation\n",
    "\n",
    "Recall that a key term in the amplitude updates was the cross-correlation of the residual and the waveforms. We can compute that more efficiently by leveraging the fact that the waveforms are low rank,\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "[\\mathbf{R} \\star \\mathbf{W}_k]_t \n",
    "&= \\sum_{n=1}^N \\sum_{d=1}^D r_{n,t+d} w_{k,n,d} \\\\\n",
    "&= \\sum_{d=1}^D \\mathbf{r}_{t+d}^\\top \\mathbf{w}_{k,:,d} \\\\\n",
    "&= \\sum_{d=1}^D \\mathbf{r}_{t+d}^\\top \\mathbf{U}_k \\mathbf{S}_k \\mathbf{v}_{k,:,d} \\\\\n",
    "&= \\sum_{d=1}^D (\\mathbf{U}_k^\\top \\mathbf{r}_{t+d})^\\top [\\mathbf{S}_k \\mathbf{V}_k^\\top]_{:,d} \\\\\n",
    "&= [(\\mathbf{U}_k^\\top \\mathbf{R}) \\star (\\mathbf{S}_k \\mathbf{V}_k^\\top)]_t\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "Note that $\\mathbf{U}_k^\\top \\mathbf{r}_t$ is a projection of the residual onto the $R$-dimensional subspace spanned by the columns of $\\mathbf{U}_k$. This equality shows that we can perform the cross-correlation between residual and waveform in this lower dimensional space instead. \n",
    "In particular, when $R=1$, it reduces to a 1-dimensional cross-correlation of the projected residual $\\mathbf{u}_k^\\top \\mathbf{R}$ and the waveform's temporal profile $\\mathbf{v}_k$. This can yield a huge performance boost!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing\n",
    "\n",
    "In practice, the raw voltage recordings are lightly preprocessed to create the matrix $\\mathbf{X}$. \n",
    "\n",
    "1. Sometimes electrical recordings have artifacts from the environment in which the recording is performed or from nearby electronics. One step toward reducing these artifacts is **common average referencing**, where we first subtract the mean across time, then subtract the median across channels.\n",
    "\n",
    "2. Spikes and the resulting EAPs are only a few milliseconds long. Real voltage recordings also have slower signals like local field potentials (LFPs), which have time scales of 3ms to 500ms. Since we are interested in spikes, we typically **bandpass filter** each channel $\\mathbf{x}_{n} = (x_{n,1}, \\ldots, x_{n,t})$ to focus on content in the [300 Hz, 2000 Hz] frequency range; i.e. signals that vary over 0.5 to 3 ms.\n",
    "\n",
    "3. In many recordings, especially those from freely moving animals, the electrode may move slightly over time. This movement results in **drift** of the spike waveforms. State-of-the-art spike sorting software like Kilosort {cite}`pachitariu2023solving` tries to correct for drift in preprocessing.\n",
    "\n",
    "4. Since the channels are so closely spaced, noise tends to be correlated across channels. Since the noise is assumed to be conditionally independent in the convolutional semi-NMF model described above, it is common to **whiten** the data before analysis. After bandpass filtering, the data should be mean zero. Thus, the empirical covariance is $\\hat{\\mathbf{C}} = \\frac{1}{T} \\sum_{t=1}^T \\mathbf{x}_t \\mathbf{x}_t^\\top$. Whiten the data by left-multiplying by the **inverse square root** of the covariance matrix, $\\mathbf{X} \\leftarrow \\hat{\\mathbf{C}}^{-\\frac{1}{2}} \\mathbf{X}$. \n",
    "\n",
    "5. The MAP estimation problem is nonconvex, and the solution found by coordinate ascent will depend on the initialization procedure. There is no right answer for how to initialize the templates. An approach used by Kilosort (which this chapter is based on) is to initialize with a library of \"universal\" templates."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Post processing\n",
    "\n",
    "The waveforms extracted by convolutional semi-NMF usually still need a bit of post-processing. For example, sometimes the model finds two waveforms for the same neuron. Alternatively, it may assign two neurons to the same waveform if their spike timing is highly correlated. \n",
    "\n",
    "These types of errors can be addressed with a post-processing step to split or merge clusters. This step is not unique to spike sorting &mdash; it's a common postprocessing step in many unsupervised clustering analyses. For spike sorting, we can bring extra domain knowledge to bear on the problem. For example, we expect the spatial footprints to be localized, and we expect the temporal profiles to have a single downward deflection. Modern libraries incorporate checks like these into the postprocessing steps."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "This chapter extended the previous one by allowing waveforms that extend in time. The generalized model is a form of **convolutional semi-NMF**. Along the way, we picked up some new skills:\n",
    "\n",
    "- **Convolution and Cross-Correlation:** the building blocks of many machine learning models, and we'll return to them multiple times in this course.\n",
    "\n",
    "- **Frobenius norm and inner product**: basic tools for dealing with matrix-valued variables\n",
    "\n",
    "- **Singular value decomposition**: a crucial matrix factorization with lots of applications in low-rank approximation.\n",
    "\n",
    "- **More MAP estimation!** by now, you're quite familiar with framing estimation problems as maximizing the log probability and then deriving coordinate ascent algorithms. Here, the coordinate updates were particularl interesting, as they involved making some simplifying assumptions and optimizing over manifolds of low-rank matrices.\n",
    "\n",
    "Next time, we'll consider an analogous problem for working with calcium imaging data. Many of the models and tools we've developed will transfer."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Further reading\n",
    "\n",
    "The algorithm presented in this chapter is similar to Kilosort {cite}`pachitariu2016kilosort,steinmetz2021neuropixels`. The exact algorithms employed by Kilosort change from version to version, but the convolutional generative model is central. Complete details of Kilosort and the differences from one version to the next can be found in {cite:t}`pachitariu2023solving`. \n",
    "\n",
    "Of course, there are other spike sorting algorithms and implementations as well, like YASS {cite}`lee2020yass` and MountainSort {cite}`chung2017fully`. Each has its own unique aspects, and it is interesting to compare and contrast different methods. For Neuropixels users, Kilosort appears to be the go-to method.\n",
    "\n",
    "Is spike sorting really necessary though? For some questions of interest, like population decoding or state space analysis, it may not be. For example, {cite:t}`deng2015clusterless` showed improved performance using a \"clusterless\" decoding approach that uses extracted spike waveforms but does not try to assign them neuron labels. Similarly, {cite:t}`trautmann2019accurate` showed that you can identify low dimensional states and dynamics from electrophysiological recordings without spike sorting. However, if your scientific objectives involve understanding individual neurons' coding properties, then spike sorting is a necessary step.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.11 (default, Jul 27 2021, 07:03:16) \n[Clang 10.0.0 ]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "5b28c5bd4ee93d765ebe901023d5522822fb8ad083dac3187c5545022f913719"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}