Assignments/A3-TimeFrequencyFiltering.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Name\n",
    "- PID\n",
    "- COGS118C - Assignment 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This notebook has [50 + 5 Bonus] points in total \n",
    "##### The number of points for each question is denoted by []. Make sure you've answered all the questions and that the point total add up."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lab 3 - Time-Frequency Analysis and Filtering\n",
    "\n",
    "In this lab, we will cover time-frequency (TF) analysis, and the two main ways of doing this (short-time Fourier transform (STFT) and filtering). In particular, we will motivate why TF analysis is performed, rather than just applying Fourier Transform to your whole recording once. We will code for ourselves the operations required to perform STFT and filtering, then explore parameter choice considerations when applying these methods. From the time-frequency representation, we will compute the spectrogram, power spectral density, and phase coherence. We will also explore the effect of windowing on your spectral estimates. Finally, there is a practical tutorial on filtering and visualizing the frequency response of filters.\n",
    "\n",
    "Key concepts:\n",
    "\n",
    "- Short Time Fourier Transform (STFT)\n",
    "- (Symmetry and orthogonality of DFT complex exponentials)\n",
    "- Time-frequency uncertainty principle\n",
    "- Power spectral density (PSD), event-related spectral response (ERSP) and phase coherence\n",
    "- Windowing\n",
    "- Convolution and filtering\n",
    "- Filter parameter choices\n",
    "\n",
    "**Answers for questions requiring written responses can be entered in the cell immediately below the question, so that when you write your response, it doesn't screw up the formatting of the question.** "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# make the imports\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy import io, signal # we will also import the signal module, from scipy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def plot_spectrogram(spg, t, f, freq_lims=[0,100], plot_db=False):\n",
    "    \"\"\"\n",
    "    Utility function for plotting the spectrogram for you.\n",
    "    \n",
    "    spg: spectrogram, 2D real-numbered array, dimensions are [frequency x time]\n",
    "    t: time axis of spectrogram\n",
    "    f: frequency axis of spectrogram\n",
    "    freq_lims (optional): limits the frequency axis, defaults to 0-100Hz\n",
    "    \"\"\"\n",
    "    plt.figure(figsize=(15,4))\n",
    "    if plot_db:\n",
    "        plt.imshow(10*np.log10(spg), aspect='auto', extent=[t[0], t[-1], f[-1], f[0]])\n",
    "    else:\n",
    "        plt.imshow(spg, aspect='auto', extent=[t[0], t[-1], f[-1], f[0]])\n",
    "    plt.xlabel('Time'); plt.ylabel('Frequency(Hz)');\n",
    "    plt.ylim(freq_lims)\n",
    "    plt.colorbar()\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Time and Frequency Resolution\n",
    "At the heart of time-frequency analysis is the ability to resolve how frequencies change over time, whereas a single Fourier transform (or power spectrum) is unable to. Below, we will simulate a 20-second signal to analyze. I plot the data in time for you already."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "T, fs = 20, 1000\n",
    "t = np.arange(0,T,1/fs)\n",
    "# simulate a signal\n",
    "# refer to the function documentation for f0, t1, f1\n",
    "sig = signal.chirp(t, f0=10, t1=20,f1=30)\n",
    "\n",
    "# plot it\n",
    "plt.figure(figsize=(15,3))\n",
    "plt.plot(t,sig)\n",
    "plt.xlim([0,5])\n",
    "plt.xlabel('Time (s)'); plt.ylabel('Voltage (V)');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [5] Q1: Computing Power Spectrum\n",
    "[3] 1.1: Compute the power spectrum of the signal (store in variable `ps`), starting from its Fourier transform (`np.fft.fft`). Plot the power spectrum and label your axes. **Remember, plot x-axis as frequency in Hz, this is now an implicit requirement moving forward for all frequency domain plots unless otherwise stated.** You can use `np.fft.fftfreq`). Zoom into 0-50Hz.\n",
    "\n",
    "[1] 1.2: Comment on the frequency content of the signal based on the power spectrum (i.e., what frequencies are present and/or domninant).\n",
    "\n",
    "[1] 1.3: Does what you see in the time-seres plot (above) match what you see in the power spectrum, in terms of frequency content? Feel free to manipulate the x-limits of the time series plot to explore."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Response for 1.2:** ANSWER\n",
    "\n",
    "**Response for 1.3:** ANSWER\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#_YOUR_CODE_HERE\n",
    "ps = \n",
    "freqs = \n",
    "plt.figure(figsize=(5,5))\n",
    "plt.plot(freqs, ps, '.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [5] Q1 - Continued\n",
    "[1] 1.4: You'll be doing this a lot in the following section, so create a function that computes the power spectrum (`pwsp`) given a signal and it's sampling frequency, as well as returning the frequency axis (`freqs`).\n",
    "\n",
    "[1] 1.5: Using your function, compute the power spectrum of the **first second** of the signal, i.e., where t=[0,1) (this notation means inclusive of t=0, and exclusive of t=1, i.e. t>=0 and t\\<1.). Hint: this is the first 1000 points of the signal, since `fs=1000`. Plot the power spectrum and zoom into 0-50Hz (you'll be doing this a lot too, so consider making an addition functional for plotting, though this is not required).\n",
    "\n",
    "[1] 1.6: Repeat 1.5, but for the last second of the data. You can plot them on the same plot. **Remember to label your traces as well using `plt.legend()`.**\n",
    "\n",
    "[2] 1.7 Comment on the similarities and differences between the two power spectrum plotted for 1.5 and 1.6, and your interpretation of this in terms of the frequency content. Compare this plot with the plot for Question 1.1, what are their similarities and differences?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Response for 1.7:** ANSWER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def compute_pwsp(sig, fs):\n",
    "    #_YOUR_CODE_HERE    \n",
    "    return pwsp, freqs\n",
    "    \n",
    "pwsp_first, freqs_win = compute_pwsp() #_YOUR_CODE_HERE\n",
    "pwsp_last, _ = compute_pwsp()\n",
    "\n",
    "plt.figure(figsize=(6,4))\n",
    "#_YOUR_CODE_HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Symmetry and Nyquist Frequency in DFT\n",
    "A quick note to explain it in code, see L9 slides for derivation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ft = np.fft.fft(sig[0:1000])\n",
    "plt.figure(figsize=(12,8))\n",
    "plt.subplot(3,1,1)\n",
    "plt.plot(freqs_win, ft_win.real, '-o')\n",
    "plt.xlim([-50,50]); plt.ylabel('Real Component')\n",
    "\n",
    "plt.subplot(3,1,2)\n",
    "plt.plot(freqs_win, ft_win.imag, '-o')\n",
    "plt.xlim([-50,50]); plt.ylabel('Imaginary Component')\n",
    "\n",
    "plt.subplot(3,1,3)\n",
    "plt.plot(freqs_win, abs(ft_win)**2, '-o')\n",
    "plt.xlim([-50,50]); \n",
    "plt.xlabel('Frequency(Hz)'); plt.ylabel('Power');\n",
    "\n",
    "# print a few pairs just to see that it's true\n",
    "print(ft[10], ft[-10])\n",
    "print(ft[20], ft[-20])\n",
    "print(ft[100], ft[-100])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "If you recall, in A2, you were asked to compute the frequency in Hz from the wave numbers, k. There was a bonus question to make your frequencies match `np.fft.fftfreq()`, as that function returns negative frequencies after `fs/2`. `fs/2` is also known as Nyquist frequency, the fastest frequency you can resolve given a sampling rate of `fs` (not to be confused with Nyquist rate, which is the minimum sampling rate, 2B, required to accurately capture a signal with a maximum frequency of B).  \n",
    "\n",
    "In addition, you were asked to plot the real and imaginary component of the Fourier Transform. I didn't make a comment on it then, but you may have noticed something a little funny, and again in the plots above as you plot the power spectra before zooming into 0 to 50Hz. In the cell above, I plot the real and imaginary components of the FT, and power spectrum you computed in Q1.1, zooming into -50 to 50Hz instead.\n",
    "\n",
    "Notice that in the first plot (real component), the dots are symmetric about 0Hz, as if they are mirrored. The imaginary components (second plot), on the other hand, are anti-symmetric - mirrored but flipped about 0Hz. This means that for any frequency and its negative, their Fourier coefficients form a conjugate pair (symmetric about the real number line). As a byproduct, since the power is just the power squared magnitude of the vectors, it's also symmetric, while the phase (not plotted) are again anti-symmetric. This is a property of the DFT when applied to **real-valued signals**. We will go through the details in lecture, but for a formal proof, see [here](https://dsp.stackexchange.com/questions/8715/symmetry-of-real-and-imaginary-parts-in-fft). \n",
    "\n",
    "**Practically**, what this means is that half of the Fourier transform for a real-valued signal (the only kind you will work with in this class) is redundant. We can simply reconstruct the negative half from the positive half by inverting the imaginary component (or phase). **This will become relevant in Q3, where I return for you only the positive halves of the Fourier coefficents, up to `fs/2`, as that's the last (fastest) positive frequency represented.**\n",
    "\n",
    "![](./img/complex_conjugates.png)\n",
    "In the diagram above, the complex vectors **x** and **x\\*** represent the complex coefficients at one frequency and its negative, while **y** and **y\\*** are another pair of frequencies. Diagram from [here](https://betterexplained.com/articles/intuitive-arithmetic-with-complex-numbers/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Time-Varying Frequencies\n",
    "Okay, back on track. Q1 should have demonstrated to you that simply taking a power spectrum is not sufficient for some signals. In particular, when a signal violates stationarity, it may have frequency content that changes over time. This is **important**, as most signals in nature are NOT stationary. It would be like taking an average of all the frames in a movie to see what the movie is about.\n",
    "\n",
    "How do we overcome this? Using **Short Time Fourier Transforms (STFT)**. The principle behind this method is already demonstrated in the last part of Q1. Instead of taking one big Fourier transform across time, STFT takes snippets of the signal to measure its frequency content over time, giving us both frequency and time resolution (hence the name, time-frequency analysis). \n",
    "\n",
    "Instead of picking two arbitrary time windows, however, STFT does this more systematically. Specifically, it's parameterized by two additional parameters: window length (`len_win`), which corresponds to the length of the small window of data you want to compute over, and step length (`len_step`), which corresponds to how much you step forward to compute the FT of the next window. In the diagram below, `len_win`=M, and `len_step`=R. \n",
    "\n",
    "Instead of step length, some implementations use the parameter overlap length (`len_overlap`), which is how much the next window overlaps with the current one (L in diagram). This is simply the complement of R, i.e., L = M-R. This is what `scipy.signal.stft()` uses, so we will adopt that as well to minimize confusion.\n",
    "\n",
    "![](./img/stft.png) (Image from [MATLAB](https://www.mathworks.com/help/signal/ref/stft.html))\n",
    "\n",
    "---\n",
    "Below, we will walk through the steps of constructing the STFT yourself. I will first ask you some math questions to consolidate the concept, then guide you through implementing the operation in code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [8] Q2: Step By Step\n",
    "There's really nothing fancy about the STFT: forget about the Fourier transform part for a second. The fundamental component (and difference to the regular FT) is the stepping operation, taking a window (or chunk) of your data at regularly sized intervals - that's all! This operation is exactly analogous to what you did in A1 to find trial indices for ERPs, except at a regular interval. However, instead of acquiring the time indices from the experiment itself, you will compute them based on the parameters of the STFT. Note: for the rest of this question, start indexing both the window and datapoint at 0, i.e., the 0th datapoint of the 0th window has index 0.\n",
    "\n",
    "[1] 2.1: In Q1.5, we took a 1-second window of data, which is 1000-points (since `fs=1000`Hz). If we were to compute the STFT, how many windows would the entire signal be divided into with no overlap? (Remember, `T=20` seconds)\n",
    "\n",
    "[1] 2.2: What would be the timestamp (in seconds) of the zeroth data point in the 10th, 11th. and 12th window? \n",
    "\n",
    "[1] 2.3: What would be the corresponding array indices (in terms of the full data array) of those points from Q2.2? What is the formula you used to compute this?\n",
    "\n",
    "[1] 2.4: If now we want to overlap each window by 0.5 seconds (500 data points), how many **full** windows would we end up with? What would be the timestamp of the zeroth data point in the 10th, 11th, and 12th window then?\n",
    "\n",
    "[1] 2.5: What would be the array index (in terms of the full data array) of those points in Q2.4?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Response for 2.1:** ANSWER\n",
    "\n",
    "**Response for 2.2:** ANSWER\n",
    "\n",
    "**Response for 2.3:** ANSWER\n",
    "\n",
    "**Response for 2.4:** ANSWER\n",
    "\n",
    "**Response for 2.5:** ANSWER\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "[1] 2.6: Complete the function `slide_window_time()` below, which returns an array that contains the zeroth timestamp of every window. E.g., if `len_win=1s` and `len_overlap=0.5s`, the first 3 elements of the array should be [0., 0.5, 1.0 ...]. Note that the last timestamp should be the start of a valid (full) data window, i.e., it has to be at least `len_win` less than the total signal time, `T`. Hint: one easy way to do this is to call `np.arange()`, and throw out all the times values that cross into the last window.\n",
    "\n",
    "[1] 2.7: Complete the function `slide_window_index()` below, which returns an array (of integers) that contains index of the zeroth data point of each window. E.g., if `len_win=1s` and `len_overlap=0.5s`, the first 3 elements of the array should be [0, 500, 1000 ...]. Note that it requires an additional argument, `fs`. Again, you can do this by stepping through as you do in 2.6, or just use the output of `slide_window_time()`.\n",
    "\n",
    "[1] 2.8: Use the above two functions to confirm your response for Q2.2-2.5. `t_stft` should store the array of timestamps, while `ind_stft`, indices.\n",
    "\n",
    "---\n",
    "Check to make sure that your two functions return the correct answer for a range of parameters, especially for `len_overlap`. Some people have experienced strange rounding issues from `np.arange()`, which caused data windows with unequal lengths to be returned as a result. You may need to use `np.round()` to round to integer values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def slide_window_time(T,len_win,len_overlap):\n",
    "    # T is total signal time, len_win is window length in seconds, len_overlap is overlap length in seconds\n",
    "    #_YOUR_CODE_HERE\n",
    "    t_steps\n",
    "    return t_steps\n",
    "    \n",
    "def slide_window_index(T,fs,len_win,len_overlap):\n",
    "    #_YOUR_CODE_HERE\n",
    "    return \n",
    "    \n",
    "    \n",
    "# Q2.2-3\n",
    "len_win = #_YOUR_CODE_HERE # seconds\n",
    "len_overlap = #_YOUR_CODE_HERE # seconds\n",
    "t_stft = slide_window_time(T, len_win, len_overlap)\n",
    "ind_stft = slide_window_index(T, fs, len_win, len_overlap)\n",
    "print() # print the number of windows\n",
    "print() #  print the zeroth timestamps and corresponding indices of windows 10-12\n",
    "\n",
    "# Q2.4-5\n",
    "len_win = #_YOUR_CODE_HERE # seconds\n",
    "len_overlap = #_YOUR_CODE_HERE # seconds\n",
    "t_stft = slide_window_time(T, len_win, len_overlap)\n",
    "ind_stft = slide_window_index(T, fs, len_win, len_overlap)\n",
    "print(len(t_stft)) # print the number of windows\n",
    "print() #  print the zeroth timestamps and corresponding indices of windows 10-12"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [2] Q3: Short-Time Fourier Transform\n",
    "At this point, you have an array of data indices in `ind_stft`, which represents the index of the first data point of every window. The next step is very similar to what you did for computing an ERP: grab a window of data around those indices. Only it's even simpler here: you don't grab any data before that index, since it's the start of the window, and you grab `len_win` seconds worth of data following. \n",
    "\n",
    "[1] 3.1: Using array indexing and the values stored in `ind_stft` from Q2.5 (`len_win=1` & `len_overlap=0.5`), grab the last data window, which should be `len_win` seconds long. Store that in the variable `data_win`. Do the same for the time indices (from `t`), and store it in `t_win`. Plot that window of data against its time indices. (Hint: it should start at t=19.)\n",
    "\n",
    "[1] 3.2: Compute it's FT using `numpy.fft.fft()`, and plot the power spectrum. It should look the same as your (orange) trace from Q1.6."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# grab a window of the signal and the corresponding time vector\n",
    "sig_win = #_YOUR_CODE_HERE\n",
    "t_win = #_YOUR_CODE_HERE\n",
    "\n",
    "# plot time series window\n",
    "plt.figure(figsize=(8,4))\n",
    "#_YOUR_CODE_HERE\n",
    "\n",
    "# compute power spectrum\n",
    "#_YOUR_CODE_HERE\n",
    "\n",
    "# plot power spectrum\n",
    "plt.figure(figsize=(6,4))\n",
    "#_YOUR_CODE_HERE\n",
    "plt.xlim([0,50])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [6] Q3 - Continued\n",
    "[3] 3.3: The last two steps guides you on how to grab a single time window given it's starting index and compute the Fourier transform (and power spectrum). Complete the function `my_stft()` below, which should iterate over all window-start indices and grab the corresponding window of data and compute each Fourier Transform. Use the functions you've built previously. \n",
    "\n",
    "You need to return 3 output arguments: `f_stft`, the corresponding frequency axis for the STFT; `t_stft`, the time stamps for the zeroth data point of each window; and `stft`, the short time Fourier transform (2D matrix). `stft` should have shape (`len(f_stft), len(t_stft)`). Note here `f_stft` and `stft` should contain only the positive frequencies (I actually do that for you in the function, assuming you computed the variables correctly).\n",
    "\n",
    "You can either loop through the time series and compute the Fourier transform of each window and store that in a matrix, or collect all the time series data in one go, and call `np.fft.fft()` once - it will perform a 1-dimensional FFT along the axis specified.\n",
    "\n",
    "[1] 3.4: Compute the spectrogram from the STFT, which is the squared magnitude of every value in STFT. Store that in `spg`.\n",
    "\n",
    "[1] 3.5: Plot the spectrogram you have computed using the helper function `plot_spectrogram()`. Set `freq_lims` to [0,50] Hz. I plot the spectrogram generated using `signal.spectrogram` for you, they should look identical. \n",
    "\n",
    "[1] 3.6: Comment on the spectrogram and what you can interpret from it about the signal.\n",
    "\n",
    "[2] 3.7: Compute the power spectral density (PSD) from your spectrogram `spg`, as well as scipy's spectrogram `spg_sp` and plot them. They should look identical. Remember to label the traces and the axes appropriately. Note: don't worry about the normalization for spectral density for now, just take the average over time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Response for 3.6:** ANSWER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def my_stft(data, fs, len_win, len_overlap):\n",
    "    T = #_YOUR_CODE_HERE\n",
    "    inds_stft = #_YOUR_CODE_HERE\n",
    "    t_stft = #_YOUR_CODE_HERE\n",
    "    f_stft = #_YOUR_CODE_HERE\n",
    "    stft = #_YOUR_CODE_HERE\n",
    "    \n",
    "    # clip the frequency axis to return just the non-negative frequencies\n",
    "    # np.fft.fftfreq returns the nyquist frequency as negative, which we also need to keep\n",
    "    positive_fs = np.logical_or(f_stft>=0, f_stft==-fs/2)    \n",
    "    \n",
    "    # I return for you just the positive frequencies \n",
    "    return abs(f_stft[positive_fs]), t_stft, stft[positive_fs, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "len_win = 1.\n",
    "len_overlap = 0.5\n",
    "\n",
    "# computing and plotting your spectrogram\n",
    "f_stft, t_stft, stft = my_stft() #_YOUR_CODE_HERE \n",
    "print(len(f_stft), len(t_stft), stft.shape) # print the shape of the STFT to confirm it is the right dimensions\n",
    "spg = #_YOUR_CODE_HERE\n",
    "plot_spectrogram(spg, t_stft, f_stft, freq_lims=[0,50], plot_db=False)\n",
    "plt.title('My Spectrogram')\n",
    "\n",
    "# I compute the spectrogram using scipy and plot it for you\n",
    "f_sp, t_sp, spg_sp = signal.spectrogram(sig, fs, window='rect', nperseg=int(fs*len_win), noverlap=int(fs*len_overlap), detrend=False)\n",
    "# I undo scipy's normalization for you so they will look the same\n",
    "spg_sp = spg_sp*fs**2*len_win/2\n",
    "plot_spectrogram(spg_sp, t_sp, f_sp, freq_lims=[0,50], plot_db=False)\n",
    "plt.title('scipy Spectrogram');\n",
    "\n",
    "# computing and plotting your PSDs\n",
    "plt.figure(figsize=(5,5))\n",
    "#_YOUR_CODE_HERE\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Now let's try it on some real brain signals\n",
    "We'll be working with the same LFP data as last time, recorded in the rat hippocampus. This dataset comes from an openly accessible neuroscience database. For more information on this particular dataset, see [here](https://crcns.org/data-sets/hc/hc-2/about-hc-2). \n",
    "\n",
    "We will use a slightly longer snippet of data to start (`lfp_short`, 2min), for the purpose of demonstrating the spectrogram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = io.loadmat('./data/LFP.mat', squeeze_me=True)\n",
    "print(data.keys())\n",
    "\n",
    "# unpack the variables\n",
    "fs = data['fs'] # sampling rate \n",
    "print('Sampling rate = %iHz'%fs)\n",
    "\n",
    "lfp = data['lfp'][0,:]/1000 # this file contains two channels, we'll only work with the first one\n",
    "lfp_short = lfp[:int(120*fs)] # make a variable that has only the first two seconds of the LFP\n",
    "t_short = np.arange(0, len(lfp_short)/fs, 1/fs) # create the corresponding time vector\n",
    "\n",
    "plt.figure(figsize=(15,3))\n",
    "plt.plot(t_short,lfp_short, 'k', alpha=0.8)\n",
    "plt.xlim([0,10])\n",
    "plt.xlabel('Time (s)');plt.ylabel('Voltage');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [6 + 2 Bonus] Q4: Applying STFT & PSD\n",
    "[3] 4.1: **Using a window length of 1s, and overlap of 0.8s**, compute the STFT (using `my_stft`), spectrogram (**normalized/spectral density**), and PSD. Plot the spectrogram and PSD. I compute the PSD using sp.signal for you in `f_welch` and `p_welch`. Plot that to confirm you have the right answer (the PSDs should overlap). Zoom into 0 to 20Hz for both the spectrogram and PSD.\n",
    "\n",
    "[1] 4.2: Repeat the above analyses and plots in a new cell, but **using a window length of 5s, and overlap of 4s**. You can literally copy and paste the entire block of code, changing just the variable values for `len_win` and `len_overlap`.\n",
    "\n",
    "[2] 4.3: Discuss similarities and differences for two spectrograms, especially with respect to temporal and frequency resolution. Also compare the two sets of PSDs, noting the frequency content of the signal represented in both cases.\n",
    "\n",
    "[BONUS: 2] 4.4: There is a known cause for the change in frequency for this neural oscillation in rats. What is the name of this oscillation, and what is the cause for the frequency change? Reference your source here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Response for 4.3:** ANSWER\n",
    "\n",
    "**Response for 4.4:** ANSWER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "len_win = #_YOUR_CODE_HERE\n",
    "len_overlap = #_YOUR_CODE_HERE\n",
    "f_stft, t_stft, stft = #_YOUR_CODE_HERE\n",
    "spg = #_YOUR_CODE_HERE\n",
    "psd = #_YOUR_CODE_HERE\n",
    "\n",
    "# use plot_spectrogram\n",
    "plot_spectrogram(spg, t_stft, f_stft, plot_db=False, freq_lims=[0,20])\n",
    "\n",
    "# computing welch's spectrogram for you\n",
    "f_welch, p_welch = signal.welch(lfp_short,fs,window='boxcar',nperseg=int(fs*len_win),noverlap=int(fs*len_overlap))\n",
    "\n",
    "# again, undoing scipy's normalization, one of the following should give you the exact same plot\n",
    "# p_welch *= 0.5*len_win \n",
    "# p_welch *= (fs**2*len_win/2)\n",
    "\n",
    "# plot the PSDs\n",
    "plt.figure(figsize=(5,5))\n",
    "#_YOUR_CODE_HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Intermission & scipy.signal\n",
    "Q4 is an example of what practical (neural) signal processing is about: making informed parameter decisions based on your data and the question you want to answer. We've covered most of the math details already, so from this point on, the labs will shift focus from implementing the analysis toolkits to applying them and making sound judgements for your analysis parameters.\n",
    "\n",
    "---\n",
    "In the exercises above, I call functions from the `scipy.signal` module to check your answers for you, which has most of the functions relevant for digital signal processing, especially for this class. As we progress through more complex material, you will be asked to use them directly to speed things up, and certainly for your final project, so it's important to be familiar with the documentation for them. To view documentation, put your cursor inside the brackets of the function call in a code cell, and press shift+tab+tab.\n",
    "\n",
    "In particular, the following functions are relevant for you for this lab:\n",
    "- `signal.stft()`: computing STFT\n",
    "- `signal.welch()`: computing PSD using STFT with optional overlapping windows\n",
    "- `signal.windows`: a sub-module that contains windowing functions\n",
    "- `signal.firwin()`: computing FIR filter coefficients\n",
    "- `signal.hilbert()`: computing the Hilbert transform\n",
    "\n",
    "Note: most of the functions above take in the **number of points** as it's length arguments, such as `nperseg` (number of data points per window) and `noverlap`. We still want to define our length parameters in units of seconds because that's more intuitive, so take note to convert them to number of points (and they have to be integers) by multiplying by the sampling rate."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# [6] Q5: Windowing & PSD\n",
    "This question illustrates how windowing with different window functions can give rise to slightly different spectral estimates. As a reminder, windowing is multiplying your short-time data segments with a windowing function, and is done regardless whether you perform it explicitly (in that case, a \"boxcar\" window is equivalently used). This is equivalent to convolving the \"true\" signal power spectrum with the spectral profile of the window.\n",
    "\n",
    "[1] 5.1: Since all the `scipy.signal` functions specify window length in terms of points, complete the function `secs_to_points` that converts the length variables in time (`len_win`, `len_overlap`) to number of points. Remember to return `ints`. We will still specify parameters in units of second, but convert them to points before using for `scipy` calls from this point on.\n",
    "\n",
    "[1] 5.2: With window length of 2s and overlap length of 1s, compute the PSD using `signal.spectrogram()` by averaging the spectrogram, as well as using `signal.welch()`, specifying `fs`, `nperseg`, and `noverlap` from above. Plot the two PSDs and zoom into 0-30Hz. They should **not look identical**, and I print the percentage difference at the maximum point for you, it should be around 3.49%.\n",
    "\n",
    "[1] 5.3: Find the default windows used `signal.welch()` and `signal.spectrogram()`. Create them both using `signal.windows`, where `n=nperseg`, and plot them.\n",
    "\n",
    "[1] 5.4: Estimate their spectral response by simply taking the squared magnitude of their Fourier Transform. Plot them in frequency domain and zoom into [-5, 5] Hz. Remember to label the traces and the axes for this and the above plots!\n",
    "\n",
    "[2] 5.5: Describe the differences between the two spectral responses. In particular, note the height of the center point, and the main lobe. Is this consistent with the differences in the PSD estimates from Q5.2? Why or why not? Hint: convolution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Response for 5.5:** ANSWER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def secs_to_points(fs, len_win, len_overlap):\n",
    "    nperseg = #_YOUR_CODE_HERE\n",
    "    noverlap = #_YOUR_CODE_HERE\n",
    "    return nperseg, noverlap\n",
    "\n",
    "nperseg, noverlap = secs_to_points(fs, len_win=2, len_overlap=1)\n",
    "\n",
    "f_spg, t_spg, spg = signal.spectrogram() #_YOUR_CODE_HERE\n",
    "p_spg = #_YOUR_CODE_HERE for computing your PSD\n",
    "f_welch, p_welch = signal.welch()\n",
    "print(100*(p_spg[16]-p_welch[16])/p_welch[16])\n",
    "\n",
    "plt.figure(figsize=(15,5))\n",
    "plt.subplot(1,3,1)\n",
    "#_YOUR_CODE_HERE for plotting PSDs\n",
    "\n",
    "plt.legend()\n",
    "plt.xlim([0,30])\n",
    "\n",
    "win_tukey = #_YOUR_CODE_HERE\n",
    "win_hann = #_YOUR_CODE_HERE\n",
    "\n",
    "plt.subplot(1,3,2)\n",
    "#_YOUR_CODE_HERE for plotting windows\n",
    "plt.legend(); plt.title('Windows in time')\n",
    "\n",
    "plt.subplot(1,3,3)\n",
    "#_YOUR_CODE_HERE for plotting windows in frequency\n",
    "plt.xlim([-5,5])\n",
    "plt.legend(); plt.title('Windows spectral profile');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [8 + 2 Bonus] Q6: Coherence\n",
    "Now we will look at inter-trial coherence: in this experiment, the rat is freely running around. At some points, there will be events that prompts it to do something, or be given a reward, similar to the stimulus in the ERP experiment. I have given you the time points of those events in `t_trials`. We will look at how stimulus in those trials affect the strong rhythm we saw in the PSD, at 8.5Hz.\n",
    "\n",
    "[1] 6.1: Compute the STFT using `signal.stft`, using a 'hann' window.\n",
    "\n",
    "[1] 6.2: We will be looking at the coherence of the dominant oscillation, which is about 8.5Hz. Find the index of the frequency vector where it equals to 8.5Hz. Equivalently, compute the wave number k for the 8.5Hz oscillation.\n",
    "8.5Hz.\n",
    "\n",
    "[1] 6.3: To start, we will look at the coherence of 7 random trials. This is often done to establish a null for what to expect by chance. You will notice that the individual complex vectors (dashed) have different lengths. To compute coherence, which we will define to be the magnitude (length) of the averege complex vector, we want to remove the effect of individual trial powers. In the for loop, add a line to normalize the complex coefficient of each trial, `X_trial`, by its length. Do the same for the average, `X_mean`.\n",
    "\n",
    "[1] 6.4: Compute and print the coherence value. Again, we define that here as the magnitude of the average **length-normalized** vector. This should not be very high (in the 0.2 neighborhood).\n",
    "\n",
    "[1] 6.5: Repeat the analyses for the real trials, given the trial times `t_trials`. This is similar to what you did for the ERP analysis. To make your life easier, these trial times can be matched exactly to the STFT window times. Your first task is to find the STFT indices (in the time dimension) that correspond to these trial times. Store that in `trial_inds`.\n",
    "\n",
    "[2] 6.6: Repeat the plotting and coherence computation in Q7.3-4, but using the trial indices you just found. Plot these in the same figure, but in red. (Just change the `k` in `'k.--'` to `r`). Don't forget the normalization. The single trial vectors should be cluster around pi/2, and the coherence value should be very high (0.988).\n",
    "\n",
    "[BONUS: 2] 6.7: Repeat 6.5-6, but for the STFT window immediately before the onset of the trial. This should allow you to compute the *change* in coherence.\n",
    "\n",
    "[1] 6.8: What is your interpretation of this result, in the context of the experiment? i.e., how does the stimulus affect the oscillation, compared to the randomly chosen time windows?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Response for 6.8:** ANSWER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Trial times\n",
    "t_trials = np.array([6., 22., 45., 63., 78., 111., 118.])\n",
    "\n",
    "# Use these length settings for STFT\n",
    "nperseg, noverlap = secs_to_points(fs, len_win=2, len_overlap=1)\n",
    "\n",
    "# compute stft\n",
    "f_stft, t_stft, stft = #_YOUR_CODE_HERE\n",
    "\n",
    "# find the index (k) of 8.5Hz.\n",
    "f_ind = #_YOUR_CODE_HERE\n",
    "\n",
    "\n",
    "# grab random slices of the STFT to compute coherence \n",
    "np.random.seed(0)\n",
    "trial_inds = np.random.randint(0,len(t_stft), size=len(t_trials))\n",
    "\n",
    "plt.figure(figsize=(5,5))\n",
    "plt.xlabel('Real'); plt.ylabel('Imag')\n",
    "# plot all trial vectors\n",
    "for trial_ind in trial_inds:\n",
    "    X_trial = stft[f_ind, trial_ind].squeeze()\n",
    "    \n",
    "    # normalize by vector length\n",
    "    #X_trial = #_YOUR_CODE_HERE\n",
    "    \n",
    "    plt.plot([0,X_trial.real], [0, X_trial.imag], 'k.--', alpha=0.5)\n",
    "    \n",
    "# plot mean vector\n",
    "X_mean = #_YOUR_CODE_HERE\n",
    "plt.plot([0,X_mean.real], [0, X_mean.imag], 'ko-', alpha=0.5)\n",
    "\n",
    "coherence = #_YOUR_CODE_HERE\n",
    "print(coherence)\n",
    "\n",
    "# ---------- now do it for real -------------------\n",
    "# get trial indices from trial times\n",
    "trial_inds = #_YOUR_CODE_HERE\n",
    "\n",
    "for trial_ind in trial_inds:\n",
    "    #_YOUR_CODE_HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FIR Filtering & Filter Response\n",
    "There are 4 types of filters: lowpass, highpass, bandpass, and bandstop. They refer to the frequency response of the filter, e.g., lowpass means to allow low frequencies through (from 0Hz to the cutoff) and filter out high frequencies, the opposite for highpass. Bandpass allows through a narrow band frequency and filters out the rest, while bandstop does the opposite, which is commonly used for filtering out a specific frequency of noise (such as 60Hz line noise). Filter design is an art that will take many such labs to cover extensively, so we will just introduce the idea here and give you the tools to explore that at a later stage (such as for your project).\n",
    "\n",
    "### Finite Impulse Response \n",
    "The filter response can be examined in both time and frequency domain. If we plot the coefficients of an FIR filter in time, that's quite literally its impulse response function, i.e., if you tried to filter a delta with this function, it will output itself. FIR stands for finite impulse response, which means the impulse response function has finite time. Infinite impulse response (IIR) filters, on the other hand, have feedback, and thus will continue out to infinity even for a single delta input. We will only be using FIR filters for this course.\n",
    "\n",
    "---\n",
    "### `signal.firwin()`\n",
    "A quick tutorial on `signal.firwin()`: this function designs the FIR filter based on your frequency requirements, and return the filter coefficients. The 4 critical parameters are `numtaps`, `fs`, `cutoff`, and `pass_zero`. \n",
    "\n",
    "- `numtaps` is the filter \"order\", basically, how many points are in the filter. The longer the filter is, the better frequency resolution you will have, but worse temporal resolution.\n",
    "- `fs` is the sampling rate of your signal\n",
    "- `cutoff` defines the frequency to pass/block\n",
    "- `pass_zero` defines whether 0Hz is passed or blocked\n",
    "\n",
    "`cutoff` and `pass_zero`, in conjunction, defines the filter type. If your cutoff is at 20Hz and define `pass_zero=True`, then `firwin()` interprets that to be a lowpass filter. If `pass_zero=False`, then it's a highpass filter. Same idea applies to bandpass and bandstop, except `cutoff` is now required to be a tuple.\n",
    "\n",
    "---\n",
    "### `np.convolve()`\n",
    "Finally, to apply the filter, we simply convolve the signal with the filter, using `np.convolve()`. Remember, convolution in time domain is multiplication in frequency domain, and actually, I believe `np.convolve()` is implemented via FFT and multiplication in the frequency domain.\n",
    "\n",
    "Below, I will demonstrate an example of a lowpass filter, plotting its IRF in time and frequency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cutoff = (20, 40) #Hz\n",
    "\n",
    "# we typically want a filter order to be at least as long as 3 periods (cycles) of \n",
    "# the slowest frequency in the cutoff. \n",
    "# so if cutoff is 20Hz, we want 3*0.05s = 0.15s long, which is 0.15*fs points\n",
    "# the larger this number is, the better frequency resolution you will have\n",
    "filt_order = int(3*fs/cutoff[0])\n",
    "filt_coefs = signal.firwin(filt_order, cutoff, fs=fs, pass_zero=True)\n",
    "\n",
    "# compute the magnitude and phase response of the filter\n",
    "freq_resp = np.fft.fft(filt_coefs, n=int(fs))\n",
    "mag_resp = abs(freq_resp)**2\n",
    "ph_resp = np.angle(freq_resp)\n",
    "freqs = np.fft.fftfreq(int(fs),1/fs)\n",
    "\n",
    "plt.figure(figsize=(15,5))\n",
    "plt.subplot(1,3,1)\n",
    "# plot impulse response\n",
    "t_filt = np.arange(0,len(filt_coefs))/fs\n",
    "plt.plot(t_filt,filt_coefs, '.')\n",
    "plt.xlabel('Time(s)'); plt.ylabel('Amplitude'); \n",
    "\n",
    "plt.subplot(1,3,2)\n",
    "plt.plot(freqs, mag_resp, '.')\n",
    "plt.xlabel('Frequency(Hz)'); plt.ylabel('Power'); \n",
    "plt.xlim([0,100])\n",
    "\n",
    "plt.subplot(1,3,3)\n",
    "plt.plot(freqs, ph_resp, '.')\n",
    "plt.xlabel('Frequency(Hz)'); plt.ylabel('Phase'); \n",
    "plt.xlim([0,100])\n",
    "\n",
    "plt.figure(figsize=(15,4))\n",
    "lfp_filt = np.convolve(lfp_short, filt_coefs, mode='same')\n",
    "plt.plot(t_short, lfp_short, alpha=0.8, label='Original')\n",
    "plt.plot(t_short, lfp_filt, alpha=0.8, label='Filtered')\n",
    "plt.xlim([0,1])\n",
    "plt.legend()\n",
    "plt.xlabel('Time(s)');plt.ylabel('Voltage (V)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [4 + 1 Bonus] Q7: Filter Decisions\n",
    "Following the template from above, do the following:\n",
    "- construct filter coefficients using `signal.firwin()`\n",
    "- plot the IRF in time and frequency domain (both magnitude and phase)\n",
    "- filter the lfp signal using `np.convolve`, and plot both signals in time (zoom into first 5 seconds)\n",
    "- plot the power spectrum of both the original and filtered signal. (I didn't do this for you)\n",
    "\n",
    "[1] 7.1: a band-pass filter, with cut-off between 4-12Hz.\n",
    "\n",
    "[1] 7.2: a band-stop filter, with cut-off between 4-12Hz.\n",
    "\n",
    "[1] 7.3: a high-pass filter, with a cut-off at 0.1Hz.\n",
    "\n",
    "[1] 7.4: Which of the above is most suitable for isolating (keeping) the dominant frequency in the LFP.\n",
    "\n",
    "[BONUS 1] 7.5: since you will be doing the same sequence of operations above multiple times, it's useful to think about how to construct them into functions. You will be rewarded 2 bonus points if, after your function definition, each of the points above is completed in one line of code (to your function call). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Response for 7.4:** ANSWER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Code for 7.1-3. You can use more than one cell."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Congratulations!\n",
    "You have just performed your first time-frequency analysis on neural data! While Fourier analysis from the last lab is the basis of signal analysis (and of the methods today), the set of tools presented here are the ones practicing neuroscientists actually use on a daily basis, including computing PSD and coherence. In addition, filtering has broad applications across domains, especially in sound engineering. Almost every lab in Cognitive Science and Neuroscience that record electrophysiological data here at UCSD will employ these tools for their analysis, so understanding these well will certainly position you to become an effective research assistant, or even graduate student researcher."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# End Survery\n",
    "Please take a few minutes to fill out the following as it will help us to improve the following assignments & lectures."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Content:\n",
    "What was one thing you learned from this lab & associated lectures?\n",
    "\n",
    "ANSWER: \n",
    "\n",
    "---\n",
    "What was one thing that you still found confusing after the lab, and need clarification?\n",
    "\n",
    "ANSWER:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Style:\n",
    "What was one thing you enjoyed about the formatting of this assignment (e.g., clarity, structure, guidance, etc.)?\n",
    "\n",
    "ANSWER:\n",
    "\n",
    "---\n",
    "What was one thing that you thought could use improvements on?\n",
    "\n",
    "ANSWER:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Thank you!"
   ]
  }
 ],
 "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}