day3_1_STA.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spike Triggered Average (STA)\n",
    "\n",
    "So far, we have familiarised ourselves with the alignment of spike trains to stimuli, binning them and calculating the firing rate. Today, you will try to predict the response of cells to a new kind of visual stimulus, and in the process build your first model - the linear-nonlinear (LN) model. But before you can start working on the model itself, you will need to compute the spike triggered average (STA) for each cell.\n",
    "\n",
    "The STA represents the sequence of stimuli that make a cell happy (and cause it to fire). The more similar the stimulus is to the cell's STA, the more the cell will fire. This is why you need it to predict the cell's responses to stimuli."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**To compute the STA, you will need to:**\n",
    "\n",
    "* Load and align the stimulus and pulses\n",
    "* Collect snippets of the stimulus preceding each spike\n",
    "* Find the average stimulus that evokes spiking"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's import the packages we will need:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%config InlineBackend.rc={'figure.figsize': (12, 6), 'font.size': 14 }\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "### HELPER FUNCTIONS ###\n",
    "# Please run this cell #\n",
    "# before starting with #\n",
    "#     the notebook     #\n",
    "########################\n",
    "\n",
    "\n",
    "def plot_rasters(\n",
    "    response_matrix: np.ndarray,\n",
    "    ax: plt.Axes,\n",
    "    sampling_rate: float,\n",
    "    **kwargs\n",
    "):\n",
    "    for t, trial in enumerate(response_matrix):\n",
    "        spike_list = []\n",
    "        for c, spike_count in enumerate(trial):\n",
    "            if spike_count > 0:\n",
    "                spike_list.extend((np.random.random(size=spike_count) + c).tolist())\n",
    "\n",
    "        if spike_list:\n",
    "            ax.plot(\n",
    "                np.array(spike_list) / sampling_rate,\n",
    "                [1 + t] * len(spike_list),\n",
    "                **kwargs\n",
    "            )\n",
    "    ax.invert_yaxis()\n",
    "\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1 - Full-field flicker stimulus\n",
    "\n",
    "The data you are going to work with comes from multielectrode-array recordings of retinal ganglion cells that were stimulated with a full-field flicker stimulus. The full-field flicker consists of a screen-wide presentation of contrast levels that change with a given frequency (here 75Hz). The contrast value of each presentation is calculated from a number that's been drawn from a Gaussian distribution of mean zero and standard deviation 0.3, with positive values representing brighter, negative values darker presentations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is what it looks like:\n",
    "\n",
    "<center><img src=\"images/fff.gif\" width=\"200\"></center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The (reasonable) assumption here is that the stimulus seen by the cell just before it spiked is actually what caused it to spike. Hence, we are going to collect the sequences of contrast values that were presented in a fixed time window preceding **each spike** and average them. The result will be the spike-triggered average, i.e. the average stimulus that elicits a spike."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center><img src=\"images/sta.png\" width=\"600\"></center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2 - Load the stimulus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us first load the pulses (times of stimulus presentation), the stimulus values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pulses_filepath = 'data_LN_model/frametimes_fullfieldnoise.txt'\n",
    "pulses = np.loadtxt(pulses_filepath)\n",
    "\n",
    "## Each pulse marks a change in stimulus presentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stimulus_filepath = 'data_LN_model/stimulus_fullfieldnoise.txt'\n",
    "stimulus = np.loadtxt(stimulus_filepath)\n",
    "\n",
    "## Each value is a contrast value of the screen."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us take a look at the dimensions of `pulses` and `stimulus`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(pulses.shape)\n",
    "print(stimulus.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The stimulus shown to the retina consists of a **training set** and a **test set**. The **training set** will be used to compute the STA and tune the model's parameters so it learns to predict a cell's responses, while the **test set** will be used to evaluate the model's ability to predict that cell's responses to unseen stimuli.\n",
    "\n",
    "The training set consists of randomly drawn contrast values from a Gaussian distribution (as mentioned before) and it **does not repeat in trials** whereas the test set consists of one sequence of randomly drawn contrast values that **remains the same across trials**. The test set is repeated in every trial to account for any variability, for eg. from experimental noise, that may arise in the cell's responses. The model will be evaluated on the average response of the cell to the test stimulus across trials."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p><center><img src=\"images/FFFstimstructure.png\" width=\"600\" height=\"400\"></center></p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each trial in our dataset contains 1800 training frames and 600 test frames. The `pulses` and `stimulus`, however, are simply long lists of numbers. You will need to reshape those lists to match the `(num_trials x num_frames)` sturcture of the presented stimulus."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Reshape the values in `pulses` and `stimulus` into trials.\n",
    "\n",
    "*Hint 1:* `pulses` and `stimulus` must have the same length, and the generated `stimulus` is longer than needed.\n",
    "\n",
    "*Hint 2:* You will need to compute the number of trials from the information you have so far."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nframes_training = 1800\n",
    "nframes_test = 600\n",
    "\n",
    "## START CODE HERE ##\n",
    "\n",
    "# Make sure stimulus and pulses have the same length\n",
    "stimulus = stimulus[:len(pulses)]\n",
    "\n",
    "# Calculate the number of trials and reshape pulses accordingly\n",
    "trial_frames = nframes_training + nframes_test\n",
    "num_trials = int(np.floor(len(pulses) / trial_frames))\n",
    "\n",
    "reshaped_pulses = pulses[:num_trials * trial_frames]\n",
    "reshaped_pulses = np.reshape(reshaped_pulses, (num_trials, trial_frames))\n",
    "\n",
    "reshaped_stimulus = stimulus[:num_trials * trial_frames]\n",
    "reshaped_stimulus = np.reshape(reshaped_stimulus, (num_trials, trial_frames))\n",
    "\n",
    "## END CODE HERE ##"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let us check if your code is correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(reshaped_pulses.shape)\n",
    "print(reshaped_stimulus.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:**\n",
    "\n",
    "```\n",
    "(41, 2400)\n",
    "(41, 2400)\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(reshaped_pulses[0, :10])  # first ten pulse timings\n",
    "print(reshaped_pulses[-1, -10:]) # last ten pulse timings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:** \n",
    "\n",
    "```\n",
    "[1.7499  1.76318 1.77658 1.78986 1.80322 1.81654 1.82986 1.84318 1.85658\n",
    " 1.8699 ]\n",
    "[1313.70406 1313.71738 1313.7307  1313.74398 1313.75738 1313.7707\n",
    " 1313.78402 1313.79734 1313.81054 1313.82386]\n",
    " ```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an additional step, let us separate the training set and test set in both `pulses` and `stimulus`. \n",
    "This will help us later while calculating the STA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pulses_training = reshaped_pulses[:, :nframes_training]\n",
    "pulses_test = reshaped_pulses[:, nframes_training:]\n",
    "\n",
    "stimulus_training = reshaped_stimulus[:, :nframes_training,]\n",
    "stimulus_test = reshaped_stimulus[:, nframes_training:]\n",
    "\n",
    "print(pulses_test.shape)\n",
    "print(pulses_training.shape)\n",
    "\n",
    "print(stimulus_test.shape)\n",
    "print(stimulus_training.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3 - Load the spike times"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `pulses` and `stimulus` are now organized into training and test sets. Let us now load the spiking activity of a single cell in response to the stimulus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st_filepath = 'data_LN_model/fullfieldnoise_C1.txt'\n",
    "spike_times = np.loadtxt(st_filepath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(spike_times.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are the time stamps (in seconds) of when the cell fired during the whole recording. Let us take a look at a few example spikes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(spike_times[:10])  # first ten spikes\n",
    "print(spike_times[-10:])  # last ten spikes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you compare the first and last spike times with the first and last pulse timings from above, you'll notice that there are spikes were recorded before the first pulse and after the last pulse. We are interested in the spikes which occur only during the stimulus presentation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Complete the code to remove the spikes that occur outside the stimulus presentation\n",
    "\n",
    "* *Hint:* Use `array.ravel()` to work with a flattened version of the array in-line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stimulus_spiketimes = spike_times[  # select...\n",
    "    ## COMPLETE CODE HERE\n",
    "    \n",
    "    # ... all spikes after the first pulse\n",
    "    # ... and all spikes before the last pulse\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## SOLUTION\n",
    "\n",
    "stimulus_spiketimes = spike_times[  # select...\n",
    "    (spike_times > reshaped_pulses.ravel()[0]) &  # ... all spikes after the first pulse\n",
    "    (spike_times < reshaped_pulses.ravel()[-1])  # ... and all spikes before the last pulse\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(stimulus_spiketimes[:10])  # first ten spike times\n",
    "print(stimulus_spiketimes[-10:])  # last ten spike times"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:**\n",
    "\n",
    "```\n",
    "[2.0142  2.01996 2.0236  2.02652 2.03316 2.03592 2.06432 2.07016 2.07772\n",
    " 2.33756] \n",
    "[1313.44004 1313.47076 1313.47576 1313.72976 1313.74104 1313.74892\n",
    " 1313.75636 1313.75916 1313.763   1313.81872]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4 - Bin the spikes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the stimulus updates about every 13.3ms (i.e. at 75Hz), it will make your life a little bit easier now (and later, in the next session) if you bin the spiking responses, such that instead of working with individual spike times, you deal with a count how many spikes each new stimulus frame elicited. If the cell sees a stimulus it likes, it'll fire more. *While it is possible to compute the STA without binning the spikes, we will stick with binning during this course.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The bin edges are defined by the **times where the stimulus was updated**, and each bin contains the **number of spikes that occured in that time window**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Complete the code to bin the `stimulus_spiketimes` according to the bin edges defined in `reshaped_pulses`.\n",
    "\n",
    "- *Alternative solution*: You could use the numpy functions `digitize` and `bincount` to compute the binned spikes. Read the corresponding numpy documentation and see if you can figure it out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "binned_spikes = np.zeros_like(reshaped_pulses.ravel(), dtype=int)\n",
    "for p in range(1, reshaped_pulses.size):\n",
    "    selected_spikes = stimulus_spiketimes[  # select...\n",
    "        ## COMPLETE THIS CODE  ##\n",
    "        \n",
    "        # ... all spiketimes before pulse p\n",
    "        # ... all spiketimes after pulse p-1\n",
    "    ]\n",
    "    binned_spikes[p-1] = selected_spikes.size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## SOLUTION\n",
    "\n",
    "binned_spikes = np.zeros_like(reshaped_pulses.ravel(), dtype=int)\n",
    "for p in range(1, reshaped_pulses.size):\n",
    "    selected_spikes = stimulus_spiketimes[  # select...\n",
    "        (stimulus_spiketimes <= reshaped_pulses.ravel()[p]) &  # ... all spiketimes before pulse p\n",
    "        (stimulus_spiketimes >  reshaped_pulses.ravel()[p-1])  # ... all spiketimes after pulse p-1\n",
    "    ]\n",
    "    binned_spikes[p-1] = selected_spikes.size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Vectorized alternative solution\n",
    "\n",
    "digi_spikes = np.digitize(\n",
    "    x=stimulus_spiketimes, bins=reshaped_pulses.ravel()\n",
    ") - 1  # important to make sure the bin ids of the spikes align with the stimulus\n",
    "binned_spikes = np.bincount(\n",
    "    digi_spikes, minlength=reshaped_pulses.size\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(binned_spikes.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:**\n",
    "\n",
    "```\n",
    "(98400,)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(binned_spikes[1000:1010])  # spikes in ten bins in the middle\n",
    "print(binned_spikes[-10:])  # spikes in last ten bins"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:**\n",
    "\n",
    "```\n",
    "[1 2 1 0 0 1 0 0 0 0]\n",
    "[0 1 1 2 2 0 0 0 1 0]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Questions:** can you explain what happens inside the `for` loop? Why does the loop run from 1 to `reshaped_pulses.size`? Why is `selected_spikes` assigned to the `p-1`th position of `binned_spikes`?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Arrange the `binned_spikes` into trials, just like with the pulses and stimuli. Then split them into `training_spikes` and `test_spikes`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## START CODE HERE ##\n",
    "\n",
    "spikes_in_trials = binned_spikes.reshape(reshaped_pulses.shape)\n",
    "\n",
    "training_spikes = spikes_in_trials[:, :nframes_training]\n",
    "test_spikes = spikes_in_trials[:, nframes_training:]\n",
    "\n",
    "## END CODE HERE ##"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(training_spikes.shape)\n",
    "print(test_spikes.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:**\n",
    "\n",
    "```\n",
    "(41, 1800)\n",
    "(41, 600)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, lets visualize the cell's responses to get an idea of what you just computed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(\n",
    "    ncols=2, \n",
    "    figsize=(12, 4),\n",
    "    sharex=False,\n",
    "    sharey=True,\n",
    "    layout=\"constrained\",\n",
    ")\n",
    "\n",
    "axs[0] = plot_rasters(\n",
    "    response_matrix=training_spikes,\n",
    "    ax=axs[0],\n",
    "    sampling_rate=75.,\n",
    "    c='k',\n",
    "    marker=\".\",\n",
    "    markersize=1.,\n",
    "    lw=0,\n",
    ")\n",
    "axs[0].set_title(\"training responses\")\n",
    "axs[0].set_xlabel(\"time (s)\")\n",
    "axs[0].set_ylabel(\"trials\")\n",
    "\n",
    "axs[1] = plot_rasters(\n",
    "    response_matrix=test_spikes,\n",
    "    ax=axs[1],\n",
    "    sampling_rate=75.,\n",
    "    c='k',\n",
    "    marker=\".\",\n",
    "    markersize=1.,\n",
    "    lw=0,\n",
    ")\n",
    "axs[1].set_title(\"test responses\")\n",
    "axs[1].set_xlabel(\"time (s)\")\n",
    "\n",
    "axs[1].invert_yaxis()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Questions:** What is each dot in these plots? What differences do you notice between the two plots? Why do those differences arise?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5 - Wrap things into functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You are currently working with the responses of a single cell, so it's alright to work with the code we have written above. However, once you begin to work with more cells at the same time, it is much easier if you have your code organized into functions. Functions allow you to repeat your analysis for different inputs without needing to rewrite / modify the core code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**CAUTION**: You want to avoid giving variables the same names inside and outside the function!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is an example in which the reshaping of the stimulus and pulses is wrapped up into one function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_reshaped_stimulus_and_pulses(stimulus_array, pulses_array, training_frames, test_frames):\n",
    "    \"\"\"\n",
    "    Return the stimulus and pulses arrays reshaped to (trials x frames).\n",
    "    \"\"\"    \n",
    "    # compute the total number of trials in the experiment\n",
    "    frames_per_trial = training_frames + test_frames\n",
    "    total_trials = int(np.floor(pulses_array.size / frames_per_trial))\n",
    "    \n",
    "    # truncate stimulus and pulses to match the total number of complete trials\n",
    "    pulses_array = pulses_array[:total_trials * frames_per_trial]\n",
    "    stimulus_array = stimulus_array[:total_trials * frames_per_trial]\n",
    "    \n",
    "    # reshape stimulus and pulses\n",
    "    pulses_reshaped = np.reshape(pulses_array, (total_trials, frames_per_trial))\n",
    "    stimulus_reshaped = np.reshape(stimulus_array, (total_trials, frames_per_trial))\n",
    "    \n",
    "    return stimulus_reshaped, pulses_reshaped"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is another example of a more abstract function! It splits a two dimensional array along it's second dimension at a specified location, and returns the smaller arrays. You did this three times in the code so far: once each for the pulses, the stimuli and the spike times. Writing a good general function saves you the effort of re-writing code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_into_two(two_dimensional_array, split_id):\n",
    "    \"\"\"\n",
    "    Split the input 2d-array along its second dimension into two sets at split_id.\n",
    "    \"\"\"\n",
    "    first_array = two_dimensional_array[:, :split_id]\n",
    "    second_array = two_dimensional_array[:, split_id:]\n",
    "    \n",
    "    return first_array, second_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a, b = split_into_two(reshaped_stimulus, nframes_training)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a.shape  # same shape as training_stimulus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b.shape  # same shape as test stimulus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now try this out yourself and write functions that do some of the analysis we already did earlier."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Write a function that computes the binned spike times, taking as input a spike time array and the reshaped pulses array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_binned_spikes(spike_times_array, reshaped_pulses_array):\n",
    "    \n",
    "    ## COMPLETE THIS FUNCTION ##\n",
    "    \n",
    "    return binned_spike_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SOLUTION (same as in code above, the vectorized solution should also work)\n",
    "\n",
    "def compute_binned_spikes(spike_times_array, reshaped_pulses_array):\n",
    "    \"\"\"\n",
    "    Bin spikes according to the bin edges defined by the pulses.\n",
    "    \"\"\"\n",
    "    binned_spike_array = np.zeros_like(reshaped_pulses_array.ravel(), dtype=int)\n",
    "    for p in range(1, reshaped_pulses_array.size):\n",
    "        selected_spikes = spike_times_array[  # select...\n",
    "            (spike_times_array <= reshaped_pulses_array.ravel()[p]) &  # ... all spiketimes before pulse p\n",
    "            (spike_times_array >  reshaped_pulses_array.ravel()[p-1])  # ... all spiketimes after pulse p-1\n",
    "        ]\n",
    "        binned_spike_array[p-1] = selected_spikes.size\n",
    "    \n",
    "    return binned_spike_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_output = compute_binned_spikes(stimulus_spiketimes, reshaped_pulses)\n",
    "assert(np.array_equal(test_output, binned_spikes))  # test if function output matches what you got earlier"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 6 - Collect stimuli before each spike"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The spike-triggered average is computed by collecting stimulus values before each spike and averaging them. But how many stimulus values before each spike should be collected? This is determined by a `time_window` parameter.\n",
    "\n",
    "The `time window` is typically chosen to be between 500ms and 1000ms, depending on the animal model. For our data, 600ms is a good window to choose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "time_window = 0.6  # defined in seconds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What does this mean in terms of number of pulses? The stimulus was presented at a frequency of 75Hz and each presentation was marked by a pulse. Knowing this, you can calculate the `duration` of each pulse and the number of pulses in a `time_window`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise**: Calculate `num_past_pulses` based on `time_window`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## START CODE HERE##\n",
    "\n",
    "mean_duration = np.mean(np.diff(pulses))\n",
    "num_past_pulses = round(time_window / mean_duration)\n",
    "\n",
    "# OR\n",
    "\n",
    "num_past_pulses = int(np.floor(time_window * 75))\n",
    "\n",
    "## END CODE HERE ##\n",
    "\n",
    "print(num_past_pulses)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:** `45`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This means that we will have to go `45` frames/values back in time for each spike. However, for the very first spike, probably occuring in the first pulse duration itself, there are no stimulus values back in time. Similarly for all spikes occuring within the first few frames, there are not enough stimulus values back in time.\n",
    "\n",
    "Thus, we will have to **ignore the first few spikes** for the calculation of STA. More specificially, we will have to **ignore all spikes occurring before the first `num_past_pulses` pulses** in ***each trial***."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Complete the function that takes as input the **binned training spikes, the reshaped stimulus and the number of past frames**, and returns an array with shape (N x `num_past_pulses`), where N is the total number of spikes.\n",
    "\n",
    "*Hint 1:*  Start with an empty list, loop through each trial and each bin, and append past stimulus values to the list for every spike.\n",
    "\n",
    "*Hint 2:*  For bins with more than 1 spike, append the stimulus snippet as many times as the number of spikes in that bin.\n",
    "\n",
    "*Hint 3:*  Use `enumerate` to write more concise for-loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## FILL MISSING PIECES OF THE FUNCTION ##\n",
    "\n",
    "def get_stimulus_stack(training_spikes_array, training_stimulus, past_frames):\n",
    "    \"\"\"\n",
    "    Extract stimulus values before each spike in the training set, ignoring\n",
    "    the first past_frames bins.\n",
    "    \"\"\"\n",
    "    snippets = []\n",
    "    \n",
    "    for t, trial in enumerate(   ):  # loop over trials\n",
    "        for b, bin_count in enumerate(   ):  # loop over bins in trial\n",
    "            # only for bins after first past_frames frames\n",
    "            if (    ) \\\n",
    "               and (    ):  # and only for bins which have at least 1 spike\n",
    "                snippet = training_stimulus[\n",
    "                    \n",
    "                ]  # cut out past_frames frames before the spike from the stimulus\n",
    "                for i in range(bin_count):  # append once to list for every spike in bin\n",
    "                    snippets.append(snippet)\n",
    "    \n",
    "    return np.array(snippets)  # return a numpy array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## SOLUTION\n",
    "\n",
    "def get_stimulus_stack(training_spikes_array, training_stimulus, past_frames):\n",
    "    \"\"\"\n",
    "    Extract stimulus values before each spike in the training set, ignoring\n",
    "    the first past_frames bins.\n",
    "    \"\"\"\n",
    "    snippets = []  # create an empty list\n",
    "    \n",
    "    for t, trial in enumerate(training_spikes_array):  # loop over trials\n",
    "        for b, bin_count in enumerate(trial):  # loop over bins in trial\n",
    "            # only for bins after first past_frames frames\n",
    "            if (b + 1 >= past_frames) \\\n",
    "               and (bin_count > 0):  # and only for bins which have at least 1 spike\n",
    "                snippet = training_stimulus[\n",
    "                    t, b + 1 - past_frames:b + 1\n",
    "                ]  # cut out past_frames frames before the spike from the stimulus\n",
    "                for i in range(bin_count):  # append once to list for every spike in bin\n",
    "                    snippets.append(snippet)\n",
    "    \n",
    "    return np.array(snippets)  # return a numpy array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stimulus_stack = get_stimulus_stack(training_spikes, stimulus_training, num_past_pulses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(stimulus_stack.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:**\n",
    "\n",
    "```\n",
    "(16083, 45)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How do you know if the first number is the number of spikes? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_spikes[:, num_past_pulses-1:].sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 7 - Calculate the STA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, you are ready to calculate the spike-triggered average!! Yoohoo!!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last step of the process is to compute the average over all the stimulus snippets you extracted in the previous section. So, go ahead and do that!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Take the mean over the first axis of `stimulus_stack` and plot it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "## COMPLETE THE FOLLOWING LINE OF CODE ##\n",
    "\n",
    "sta = stimulus_stack.mean(axis=0)\n",
    "\n",
    "ax.plot(sta[::-1], lw=2)\n",
    "ax.set_xlabel(\"time before spike (bins)\")\n",
    "ax.set_ylabel(\"contrast values\")\n",
    "ax.set_title(\"first sta!\")\n",
    "ax.invert_xaxis()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations! That looks great!\n",
    "\n",
    "But, what exactly do you see? Do you understand the plot?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The STA is showing us what stimulus the cell \"prefers\" in the temporal domain. Depending on the data set you have, it might be that the cells prefers decrease and then increase in contrast before spiking; increase and then decrease in contrast; or just increase in contrast. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You're more or less done with this session!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last thing you need to do is write a that function wraps up everything we did in this notebook so far."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Complete the function `calculate_sta` that takes as input the original stimuli, pulses, the cell's spiketimes, along with the training and test frame lengths and the num_past_pulses, and returns the sta for that cell. \n",
    "\n",
    "You're free to copy code you've written above, and you shouldn't have to write any new code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_sta(\n",
    "    spike_times_array,\n",
    "    stimulus_array,\n",
    "    pulses_array, \n",
    "    training_frames,\n",
    "    test_frames,\n",
    "    past_frames,\n",
    "):\n",
    "    \"\"\"\n",
    "    Compute a cell's spike-triggered average (STA) given it's responses to a visual stimulus.\n",
    "    \"\"\"\n",
    "    # reshape the stimulus and pulses using the training_frames and test_frames\n",
    "    \n",
    "    \n",
    "    # select only those spikes which occur during the stimulus presentation\n",
    "    \n",
    "    \n",
    "    # bin the spike train using the reshaped pulses\n",
    "    \n",
    "    \n",
    "    # reshape the binned spikes to the same shape as reshaped_pulses\n",
    "    \n",
    "    \n",
    "    # split spikes into training and test\n",
    "    \n",
    "    \n",
    "    # split stimulus into training and test\n",
    "    \n",
    "    \n",
    "    # get stimulus snippets past_frames before each spike\n",
    "    \n",
    "    \n",
    "    # average over the stimulus snippets\n",
    "    \n",
    "    \n",
    "    \n",
    "    return sta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_sta(\n",
    "    spike_times_array,\n",
    "    stimulus_array,\n",
    "    pulses_array, \n",
    "    training_frames,\n",
    "    test_frames,\n",
    "    past_frames,\n",
    "):\n",
    "    \"\"\"\n",
    "    Compute a cell's spike-triggered average (STA) given it's responses to a visual stimulus.\n",
    "    \"\"\"\n",
    "    # reshape the stimulus and pulses using the training_frames and test_frames\n",
    "    stimulus_reshaped, pulses_reshaped = get_reshaped_stimulus_and_pulses(\n",
    "        stimulus_array,\n",
    "        pulses_array,\n",
    "        training_frames,\n",
    "        test_frames\n",
    "    )\n",
    "    \n",
    "    # select only those spikes which occur during the stimulus presentation\n",
    "    spikes_in_stimulus = spike_times_array[  # select...\n",
    "        (spike_times_array > pulses_reshaped.ravel()[0]) &  # ... all spikes after the first pulse\n",
    "        (spike_times_array < pulses_reshaped.ravel()[-1])  # ... and all spikes before the last pulse\n",
    "    ]\n",
    "    \n",
    "    # bin the spike train using the reshaped pulses\n",
    "    spikes_binned = compute_binned_spikes(\n",
    "        spikes_in_stimulus,\n",
    "        pulses_reshaped\n",
    "    )\n",
    "    \n",
    "    # reshape the binned spikes to the same shape as reshaped_pulses\n",
    "    spikes_in_trials = np.reshape(spikes_binned, reshaped_pulses.shape)\n",
    "    \n",
    "    # split spikes into training and test\n",
    "    training_spikes, _ = split_into_two(\n",
    "        spikes_in_trials,\n",
    "        training_frames\n",
    "    )\n",
    "    \n",
    "    # split stimulus into training and test\n",
    "    training_stimulus, _ = split_into_two(\n",
    "        stimulus_reshaped,\n",
    "        training_frames\n",
    "    )\n",
    "    \n",
    "    # get stimulus snippets past_frames before each spike\n",
    "    stimulus_snippets = get_stimulus_stack(\n",
    "        training_spikes,\n",
    "        training_stimulus,\n",
    "        past_frames\n",
    "    )\n",
    "    \n",
    "    # average over the stimulus snippets\n",
    "    sta = stimulus_snippets.mean(axis=0)\n",
    "    \n",
    "    return sta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_sta = calculate_sta(\n",
    "    spike_times,\n",
    "    stimulus,\n",
    "    pulses,\n",
    "    nframes_training,\n",
    "    nframes_test,\n",
    "    num_past_pulses\n",
    ")\n",
    "\n",
    "print(np.array_equal(sta, test_sta))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Expected output:** True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 8. Compare multiple cells"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far, we calculated the STA of one example cell. However, different cells can have very different STAs. The shape of the STA usually depends on the cell-type (ON-cell, OFF-cell etc.).\n",
    "\n",
    "Let us look at a few more cells and examine the differences between their STAs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stas = []\n",
    "\n",
    "filepaths = sorted(Path('data_LN_model').glob('fullfieldnoise*.txt'))\n",
    "\n",
    "for filepath in filepaths:\n",
    "    spike_times = np.loadtxt(filepath)\n",
    "    sta = calculate_sta(\n",
    "        spike_times,\n",
    "        stimulus,\n",
    "        pulses,\n",
    "        nframes_training,\n",
    "        nframes_test,\n",
    "        num_past_pulses\n",
    "    )\n",
    "    stas.append(sta)\n",
    "\n",
    "stas = np.array(stas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stas.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The array `stas` now contains the stas of 8 cells. Let us plot them and see what they look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(nrows=2, ncols=4, sharex=True, sharey=True, layout=\"constrained\", figsize=(12, 6))\n",
    "\n",
    "for s, sta in enumerate(stas):\n",
    "    ax = axs.flat[s]\n",
    "    ax.plot(sta[::-1], lw=3)\n",
    "    ax.set_title(f\"cell {s+1}\")\n",
    "\n",
    "axs.flat[-1].invert_xaxis()\n",
    "fig.text(0.5, -0.03, \"time before spike (bins)\", ha=\"center\", va=\"center\")\n",
    "fig.text(-0.01, 0.5, \"contrast values\", ha=\"center\", va=\"center\", rotation=\"vertical\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}