4-Receptive-Fields.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Week 6: Receptive Fields\n",
    "\n",
    "In the previous lesson, we discussed linear time-invariant systems, which can serve as a simple model of how the rate of a neuron can depend on a time-varying stimulus.\n",
    "\n",
    "This week, we'll look at some empirical findings about how neurons in the visual system respond to simple stimuli, and then consider how to represent these responses using an LTI model.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Group Members**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# load matplotlib inline mode\n",
    "%matplotlib inline\n",
    "\n",
    "# import some useful libraries\n",
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np                # numerical analysis linear algebra\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt   # plotting\n",
    "from comp_neurosci_uva import signal, data\n",
    "\n",
    "# set some style options\n",
    "mpl.rcParams['image.origin'] = 'lower'\n",
    "mpl.rcParams['image.aspect'] = 'auto'\n",
    "mpl.rcParams['image.cmap'] = 'jet'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Example\n",
    "\n",
    "Let's look at an RF with one spatial and one temporal dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filts = np.load(data.data_path / \"filters.npz\")\n",
    "k1 = filts['wb2'] * 18\n",
    "plt.imshow(k1, aspect='equal')\n",
    "plt.xlabel(\"tau (ms)\")\n",
    "plt.ylabel(\"x (px)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try to describe what kinds of visual stimuli would best excite a neuron with this RF."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "We're going to simulate the response to 50 seconds of Gaussian white noise. This is equivalent to presenting a field of vertical or horizontal bars that are various shades of gray. I'm only going to show the first 1000 samples (1 second) so you have a better idea of what this looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nx, nt = k1.shape\n",
    "stim_dt = 1         # sampling rate of the stimulus in samples/ms\n",
    "nsamples = 50000    # duration of the stimulus in samples \n",
    "np.random.seed(1)\n",
    "stim_raw = np.random.randn(nx, nsamples)\n",
    "stim = stim_raw * 0.1\n",
    "stim[:,:500] = 0\n",
    "fig, axes = plt.subplots(nrows=2, ncols=1, sharex=True, figsize=(12, 4))\n",
    "axes[0].imshow(stim, cmap=\"gray\")\n",
    "axes[0].set_ylabel(\"x (px)\")\n",
    "axes[0].set_title(\"stimulus (s)\")\n",
    "axes[1].plot(stim[0,:])\n",
    "axes[1].set_ylabel(\"pixel 0\")\n",
    "axes[1].set_xlim(0, 1000)\n",
    "axes[1].set_xlabel(\"time (ms)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Now let's do the convolution. Remember, to implement in Python we need to convolve each spatial position and then sum them up:\n",
    "\n",
    "\\begin{align}\n",
    "r_{i}(t) & = \\sum_k h(x_i, \\tau_k) s(x_i, t - \\tau_k) \\\\\n",
    "r(t) & = \\sum_i r_{i}(t)\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convolve each pixel (row) separately\n",
    "convrows = [np.convolve(k1[i], stim[i]) for i in range(nx)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sum across rows\n",
    "conv = np.row_stack(convrows).sum(0)\n",
    "\n",
    "fig, axes = plt.subplots(nrows=2, ncols=1, sharex=True, figsize=(12, 4))\n",
    "axes[0].imshow(stim, cmap=\"gray\")\n",
    "axes[0].set_ylabel(\"channel\")\n",
    "axes[0].set_title(\"stimulus (s)\")\n",
    "axes[1].plot(conv[:nsamples])\n",
    "axes[1].set_title(\"convolution\")\n",
    "axes[1].set_ylabel(\"s * k1\")\n",
    "axes[1].set_xlim(0, 1000)\n",
    "axes[1].set_xlabel(\"time (ms)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the output of the convolution is much smoother than the input. This is because the linear filter of the model neuron is averaging over space and time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Next comes the \"N\" phase of the LNP model: we have to convert the convolution to an estimated rate by passing it through a nonlinear function. We'll use `exp` with a constant offset ($r_0$ = 1 Hz)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is the baseline firing rate\n",
    "r_0 = 1\n",
    "r_est = np.exp(conv[:nsamples] + np.log(r_0))\n",
    "fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(12, 4))\n",
    "axes.plot(r_est)\n",
    "axes.set_ylabel(\"rate (Hz)\")\n",
    "axes.set_xlim(0, 1000)\n",
    "axes.set_xlabel(\"time (ms)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Then we'll generate spikes using the Bernoulli approximation we used last time: the probability of spiking in each bin is compared to a sample from a uniform distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# multiply rate by bin size to get intensity (probability) of the Poisson process\n",
    "lmb = r_est * 0.001\n",
    "n_trials = 10\n",
    "spikes = []\n",
    "fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(12, 4))\n",
    "for i in range(n_trials):\n",
    "    # draw one sample for each time bin\n",
    "    runif = np.random.uniform(size=nsamples)\n",
    "    # if lambda is greater than the random value, then a spike is emitted\n",
    "    # we can represent the spikes as a time series of zeros and ones\n",
    "    spk_v = (lmb > runif)\n",
    "    # if we want the spike times, then we have to find the bins where there are ones \n",
    "    spk_t = spk_v.nonzero()[0]\n",
    "    axes.vlines(spk_t, i - 0.4, i + 0.4)\n",
    "    spikes.append(spk_t)\n",
    "axes.set_xlim(0, 1000)\n",
    "axes.set_ylabel(\"trial\")\n",
    "axes.set_xlabel(\"time (ms)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise 1**\n",
    "\n",
    "A. Calculate and plot a PSTH of the response, averaging over trials. Label your axes for full credit. How does this compare to the estimated rate function?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "B. Use the PSTH to calculate the average firing rate of the neuron across the whole stimulus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "C. Generate 10 more trials, but with a constant offset of -0.1 added to the output of the convolution. Plot the PSTH and compare it to the result in A."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Estimating RFs\n",
    "\n",
    "As Chapter 2 in Dayan and Abbott discusses in more detail, receptive fields can predict many of the tuning properties of visual neurons. This suggests two key conclusions:\n",
    "\n",
    "1. Visual neurons (up to a point) are pretty linear.\n",
    "2. Linear models have a lot of explanatory power.\n",
    "\n",
    "As a consequence, we often want to try to estimate the linear RF of sensory neurons. Even if the neuron isn't totally linear, it's a good place to start. But how do we infer the kernel in a system where we can't stimulate with a delta function?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Example\n",
    "\n",
    "Let's see how we would do spike-triggered averaging using the response we just simulated. First, let's look at the stimulus that precedes one of the spikes in the response:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = spikes[0][5]\n",
    "sts = stim[:,t-nt:t]\n",
    "plt.imshow(sts, aspect=\"equal\", cmap=\"gray\", extent=(50, 0, 0, 50))\n",
    "plt.xlabel(\"tau (ms)\")\n",
    "plt.ylabel(\"x (px)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Doesn't look like much, right? "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "But what happens if we average over many, many spikes?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "# using an accumulator pattern to count spikes \n",
    "n_spikes = 0\n",
    "# each stimulus will be added into this empty array\n",
    "sta = np.zeros_like(k1)\n",
    "for trial in spikes:\n",
    "    for t in trial:\n",
    "        # exclude spikes that are too close to the start\n",
    "        if t - nt < 0: continue\n",
    "        n_spikes += 1\n",
    "        sta += stim[:,t-nt:t]\n",
    "# average by dividing the sum by N\n",
    "sta /= n_spikes\n",
    "fig, axes = plt.subplots(nrows=1, ncols=2, sharey=True, figsize=(9, 4))\n",
    "axes[0].imshow(sta, aspect=\"equal\", extent=(50, 0, 0, 50))\n",
    "axes[0].set_title(\"STA (n=%d)\" % n_spikes)\n",
    "axes[1].imshow(k1, aspect=\"equal\")\n",
    "axes[1].set_title(\"RF (true)\")\n",
    "axes[0].set_xlabel(\"tau (ms)\")\n",
    "axes[0].set_ylabel(\"x (px)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Voila! The average stimulus preceding the spikes produced by this neuron looks like a time-flipped version of the receptive field.\n",
    "\n",
    "If the stimulus is random and uncorrelated, given enough spikes the reverse correlation will always recover the receptive field.\n",
    "\n",
    "What happens if we reduce the average firing rate (for example by setting the baseline $r_0$ to less than 1 Hz)? What happens if we \"record\" for more or less time?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Stimulus correlations\n",
    "\n",
    "Real-world stimuli rarely look like white noise. In the visual domain, stimuli are dominated by low frequencies, because objects tend to be fairly similar intensity surrounded by sharp edges.\n",
    "\n",
    "In other words, pixels tend to be correlated with neighboring pixels and tend to change slowly over time. What happens if we try to use a stimulus that has some correlations?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import ndimage\n",
    "stim_filt = ndimage.gaussian_filter(stim_raw, sigma=(3, 50))\n",
    "stim_filt[:,:500] = 0\n",
    "fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(12, 4))\n",
    "axes.imshow(stim_filt, cmap=\"gray\")\n",
    "axes.set_xlabel(\"time (ms)\")\n",
    "axes.set_ylabel(\"x (px)\")\n",
    "axes.set_xlim(0, 2000);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "**Exercise 2**\n",
    "\n",
    "A. Simulate 10 trials of the neuron's response to the filtered stimulus, adding a constant offset of 1.0 to the convolution (otherwise the model won't fire). Use the code above as a model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(nrows=4, ncols=1, sharex=True, figsize=(12, 6))\n",
    "# Plot the stimulus in axes[0], the convolution in axes[1], the estimated rate in axes[2], and the spikes in axes[3]\n",
    "# Label your axes for full credit.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "B. Calculate and plot the spike-triggered average from this response. Compare it with the original kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "What happened?? The STA looks nothing like the kernel, even though we averaged data from roughly the same number of spikes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "**Very important conclusion**: The spike-triggered average only recovers the kernel of an LTI system if there are no correlations in the stimulus. \n",
    "\n",
    "To see why this is, think about how the average would be affected if the ensemble of black points is not spherical:\n",
    "\n",
    "<img src=\"https://gracula.psyc.virginia.edu/public/courseware/comp-neurosci/images/l8_sta_ensemble.png\" alt=\"spike-triggered average\" style=\"width: 450px;\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Kernel estimation as linear regression\n",
    "\n",
    "The solution to the foregoing problem is to correct for the correlations in the stimulus. How do we do this? It took the field a suprisingly long time to figure this out.\n",
    "\n",
    "It can be helpful to recast the problem as one of linear regression. Remember the expansion of convolution:\n",
    "\n",
    "$$r(t) = h_1 s(t) + h_2 s(t-\\Delta) + h_3 s(t-2\\Delta) + \\cdots + h_N(t-N\\Delta)$$ \n",
    "\n",
    "Here we've represented the kernel with subscripts starting with 1 but kept the functional notation for the stimulus.\n",
    "\n",
    "We can rewrite this sum as a dot product,\n",
    "\n",
    "$$r(t) = \\mathbf{s}(t) \\cdot \\mathbf{h}$$\n",
    "\n",
    "where $\\mathbf{s}(t)$ refers to the **time-lagged** stimulus. That is $\\mathbf{s}(t) = \\{s(t), s(t - \\Delta), \\ldots, s(t - N\\Delta)\\}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Example\n",
    "\n",
    "Let's look at a concrete example in numpy using a really short stimulus and kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the kernel\n",
    "h = np.asarray([0, 1, -0.5])\n",
    "# the full stimulus\n",
    "stim = np.random.randn(20)\n",
    "plt.plot(stim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's calculate the response at $t = 2$. $\\mathbf{h}$ is time-invariant, so we leave it as is. $\\mathbf{s}(2)$ is `[stim[0], stim[1], stim[2]]` but **in reverse**. In slice notation, this is `stim[2::-1]` (i.e., start at index 2 and go to the beginning in steps of -1)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"r_2 =\", np.dot(h, stim[2::-1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Vectorizing the model\n",
    "\n",
    "Instead of writing the response and stimulus as functions of time, let's think of them as a series of observations made at discrete intervals. Thus,\n",
    "\n",
    "$$r(t) = \\mathbf{s}(t) \\cdot \\mathbf{h}$$\n",
    "\n",
    "becomes\n",
    "\n",
    "$$r_i = \\mathbf{s}_i \\cdot \\mathbf{h}$$\n",
    "\n",
    "where $r_i$ is the rate at time $t_i$ and $\\mathbf{s}_i$ is the stimulus from $t_{i-N}$ to $t_i$.\n",
    "\n",
    "We can further simplify our notation by stacking all the observations in a vector $\\mathbf{r}$:\n",
    "\n",
    "$$\\mathbf{r} = \\mathbf{S} \\mathbf{h}$$\n",
    "\n",
    "The matrix $\\mathbf{S}$ has as many rows as there are time points in the observation vector $\\mathbf{r}$ and as many columns as there are time points in the kernel $\\mathbf{h}$. In each row, it contains the stimulus at the current time and at a set of previous lags. \n",
    "\n",
    "Let's see what this would look like for our toy problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "from IPython.display import display\n",
    "\n",
    "print(\"stimulus:\")\n",
    "display(stim)\n",
    "ntau = len(h)\n",
    "nt = len(stim)\n",
    "S = np.zeros((nt - ntau, ntau))\n",
    "for i in range(ntau, nt):\n",
    "    S[i-ntau] = stim[i-ntau:i][::-1]\n",
    "print(\"stimulus matrix:\")\n",
    "display(S)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take a close look at the stimulus matrix. Notice how the values shift over by one in each row?\n",
    "\n",
    "This kind of matrix is called a [Toeplitz matrix](https://en.wikipedia.org/wiki/Toeplitz_matrix), and it's the linear algebra equivalent of convolution.\n",
    "\n",
    "There's a function in `scipy` that can generate this matrix and even do some padding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.linalg import toeplitz\n",
    "S = toeplitz(stim, np.asarray([stim[0], 0, 0]))\n",
    "print(S)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "With our Toeplitz matrix, convolution is as simple as matrix multiplication:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the slow way\n",
    "r1 = np.convolve(stim, h)[:nt]\n",
    "# and the fast way\n",
    "r2 = np.dot(S, h)\n",
    "# should give the same result....\n",
    "assert np.all(r1 == r2)\n",
    "display(r1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Is $r_2$ the same as what we manually calculated above?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Linear regression with time-varying stimuli\n",
    "\n",
    "What does this have to do with linear regression? The advantage of this notation is that it's easy to see how our data consist of a series of observations made for different values of $\\mathbf{s}$. Some of the variance in these observations is explained by variation in the stimulus; the rest is random noise. We represent this noise as an additional variable that's added at each time point. \n",
    "\n",
    "$$r_i = \\mathbf{s}_i \\cdot \\mathbf{h} + \\varepsilon_i$$\n",
    "\n",
    "Or equivalently,\n",
    "\n",
    "$$\\mathbf{r} = \\mathbf{S} \\mathbf{h} + \\mathbf{\\varepsilon}$$\n",
    "\n",
    "Given this model, our goal is to estimate $\\mathbf{h}$. That is, to find values that maximize the amount of variance explained by the stimulus and minimize the amount of random error ($\\mathbf{\\varepsilon}$)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "If we assume that the noise is normally distributed and independent, then what we need to do is minimize the sum of the squares of the error. This is called **ordinary least squares** (OLS).\n",
    "\n",
    "Hopefully, you've seen that our model is a pretty vanilla linear model. The OLS solution is well-known:\n",
    "\n",
    "$$\\hat{h} = (\\mathbf{S}^T \\mathbf{S})^{-1} \\mathbf{S}^T \\mathbf{r}$$\n",
    "\n",
    "- The first term is the autocovariance matrix for the independent variables. \n",
    "- The second term is the covariance between the dependent and independent variables. \n",
    "- Here, these correspond to the **autocorrelation** and the **cross-correlation**\n",
    "\n",
    "If the stimulus is white noise, the correlation between any two instants in time is zero, so $\\mathbf{S}^T\\mathbf{S} = \\sigma^2\\mathbf{I}$ ($\\mathbf{I}$ is the identity matrix and $\\sigma^2$ is the variance of the stimulus).\n",
    "\n",
    "The $^{-1}$ operator stands for **matrix inversion**. This is what \"undoes\" the effects of the correlations in the stimulus. Matrix inversion is computationally expensive and numerically hairy. There are some tricks for dealing with the latter, but for now let's look at a quick example:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Example\n",
    "\n",
    "Let's use a nice alpha function as our kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tau = 5\n",
    "dt = 1.0\n",
    "kt = np.arange(0, 75, dt)\n",
    "k = kt / tau**2 * np.exp(-kt / tau)\n",
    "nkt = len(k)\n",
    "plt.plot(kt, k);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "And let's consider two stimuli: one that's uncorrelated and one that is correlated (i.e. smoothed)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "nt = 10000\n",
    "t = np.arange(0, nt, 1.0)\n",
    "stim_u = np.random.randn(nt)\n",
    "stim_c = ndimage.gaussian_filter(stim_u, 5)\n",
    "plt.plot(t, stim_u, t, stim_c)\n",
    "plt.xlim(0, 200);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "We'll generate our Toeplitz matrices for both stimuli and use them to simulate a response. Note that we're adding a bit of noise to the response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "S_u = toeplitz(stim_u[nkt:], stim_u[nkt-1::-1])\n",
    "S_c = toeplitz(stim_c[nkt:], stim_c[nkt-1::-1])\n",
    "\n",
    "nrt = nt - nkt\n",
    "eps = np.random.randn(nrt) * 0.05\n",
    "r_u = np.dot(S_u, k) + eps\n",
    "r_c = np.dot(S_c, k) + eps\n",
    "plt.plot(r_u)\n",
    "plt.plot(r_c)\n",
    "plt.xlim(0, 200);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "The STA is just the inverse of the convolution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "STA_u = np.dot(S_u.T, r_u)\n",
    "STA_c = np.dot(S_c.T, r_c)\n",
    "plt.plot(kt, STA_u / nt, kt, STA_c / nt, kt, k, 'k')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the STA for the uncorrelated stimulus is pretty good, but the STA for the correlated one is badly distorted. To see why this is happening, let's look at the autocorrelation of the stimulis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cov_u = np.dot(S_u.T, S_u)\n",
    "cov_c = np.dot(S_c.T, S_c)\n",
    "fig, axes = plt.subplots(nrows=1, ncols=2, sharey=True, figsize=(9, 4))\n",
    "axes[0].imshow(cov_u, aspect=\"equal\")\n",
    "axes[0].set_title(\"uncorrelated\")\n",
    "axes[1].imshow(cov_c, aspect=\"equal\")\n",
    "axes[1].set_title(\"correlated\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Notice that for the uncorrelated stimulus, the autocorrelation matrix is pretty much zero everywhere except the diagonal. In contrast, for the correlated stimulus, there are values off the diagonal that are greater than zero - this reflects the fact that for the smoothed stimulus, values near each other tend to be more similar. \n",
    "Now let's find the OLS estimate by dividing out the stimulus autocorrelations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import linalg\n",
    "\n",
    "# this works because the covariance matrix is relatively easy to invert\n",
    "ols_u = np.dot(linalg.inv(cov_u), STA_u)\n",
    "\n",
    "# we have to do a little regularization to avoid numerical issues\n",
    "ols_c = np.dot(linalg.inv(cov_c + 10 * np.eye(nkt)), STA_c)\n",
    "\n",
    "plt.plot(kt, ols_u, kt, ols_c, kt, k, 'k')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the estimate from the correlated stimulus is much closer to the true kernel!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Summary\n",
    "\n",
    "- The receptive field of a neuron can be approximated as the linear kernel of a linear-nonlinear-Poisson model\n",
    "- If the stimulus distribution is spherical and Gaussian, the RF/kernel can be estimated using spike-triggered averaging\n",
    "- If the stimulus distribution is not spherical, the RF/kernel can be estimated using linear correlation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise 3**\n",
    "\n",
    "Use linear regression to recover the visual receptive field from the response to `stim_filt` in Exercise 2.\n",
    "\n",
    "First, here's a little help in building your design matrix $\\mathbf{S}$. If your receptive field has $m$ channels and $k$ time points, then $\\mathbf{h}$ will be dimension $mk$. This is called \"unwrapping\" or \"flattening\" the receptive field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "M, K = k1.shape\n",
    "display((M, K))\n",
    "h = k1.flatten()\n",
    "display(h)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get your receptive field back, you have to *reshape* $\\mathbf{h}$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(h.reshape(M, K), aspect='equal')\n",
    "plt.xlabel(\"tau (ms)\")\n",
    "plt.ylabel(\"x (px)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your stimulus will have to have $m$ channels to match the receptive field. If it has $n$ time points, then $\\mathbf{S}$ will be dimension $n \\times mk$. This is called a *block Toeplitz* matrix, and it's formed by concatenating the Toeplitz matrices for each channel along the columns dimension.\n",
    "\n",
    "$$\n",
    "\\mathbf{S} = \\begin{pmatrix}\n",
    "s_1(0) & s_1(0 - \\Delta) & \\ldots & s_1(0 - k\\Delta) & s_2(0) & \\ldots & s_m(0 - k\\Delta) \\\\\n",
    "s_1(1) & s_1(1 - \\Delta) & \\ldots & s_1(1 - k\\Delta) & s_2(1) & \\ldots & s_m(1 - k\\Delta) \\\\\n",
    "\\vdots & & \\vdots & & & \\vdots \\\\\n",
    "s_1(n) & s_1(n - \\Delta) & \\ldots & s_1(n - k\\Delta) & s_2(n) & \\ldots & s_m(n - k\\Delta)\n",
    "\\end{pmatrix}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert stim_filt.shape[0] == M, \"Number of channels doesn't match between stim and RF\"\n",
    "M, N = stim_filt.shape\n",
    "toeps = []\n",
    "for i in range(M):\n",
    "    first = np.zeros(K)\n",
    "    first[0] = stim_filt[i, 0]\n",
    "    Xi = toeplitz(stim_filt[i], first)\n",
    "    toeps.append(Xi)\n",
    "X = np.column_stack(toeps)\n",
    "display(X.shape)\n",
    "assert X.shape[1] == h.size, \"design matrix does not have the right shape\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**A.** Calculate the spike-triggered average of the response to `stim_filt` (from Exercise 2) using `np.dot`. Plot the result and confirm that you get the same result as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**B.** Use the code in the previous section to calculate the OLS estimate of the RF. You will need to use regularization. Plot your estimate and the original RF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "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.11.4"
  },
  "rise": {
   "scroll": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}