{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to SciPy\n",
    "Tutorial at EuroSciPy 2019, Bilbao\n",
    "## 2. Signal analysis – Rocking motion of a TGV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy import fftpack, signal\n",
    "%matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "from mpl_toolkits.mplot3d import Axes3D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Importing the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.genfromtxt('data/TGV_data.csv.bz2', delimiter=',', names=True)\n",
    "time = data['Time_s']\n",
    "omega_y = data['Gyroscope_y_rads']\n",
    "_ = plt.plot(time, omega_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Time spacing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before analyzing the data any further, let us take a look at the time elapsed between subsequent measurements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "time_intervals = time[1:]-time[:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.hist(time_intervals, bins=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will ignore these small differences and assume the data to be equally spaced in time with the mean time difference between subsequent data points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta_t = np.mean(time_intervals)\n",
    "delta_t"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spectral analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fourier transformation of a discrete real signal with `rfft`:\n",
    "$$y_j = \\sum_{k=0}^{n-1} x_k \\exp\\left(-\\text{i}\\frac{2\\pi jk}{n}\\right)$$\n",
    "With exception of $y_0$ and, for even $n$, $y_{n/2}$ all $y_j$ are complex.\n",
    "`rfft` returns an array $y_0, \\text{Re}(y_1), \\text{Im}(y_1),\\ldots$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fft_y = fftpack.rfft(omega_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.plot(fft_y[::2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.plot(fft_y[1::2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "freqs = fftpack.rfftfreq(omega_y.shape[0], delta_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.plot(freqs[::2], fft_y[::2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.plot(freqs[1::2], fft_y[1::2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"color: DarkBlue;background-color: LightYellow\">\n",
    "<i>Exercise:</i> Given the Fourier transform, how does one get back to the original data?\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Power spectral density"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fft_y = np.insert(fft_y, 0, 0)\n",
    "fft_square = fft_y*fft_y\n",
    "power = fft_square[::2]+fft_square[1::2]\n",
    "power = 2*delta_t/fft_square.shape[0]*power\n",
    "_ = plt.plot(freqs[::2], power)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.plot(*signal.periodogram(omega_y, 1/delta_t))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filter out the rocking signal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are interested in the signal around 1.4 Hz. Filter out frequencies beyond 3 Hz."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filter_coeffs = signal.firwin(301, 3, pass_zero=True, fs=1/delta_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.plot(filter_coeffs, '.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "freqs, response = signal.freqz(filter_coeffs, fs=1/delta_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.plot(freqs, 20*np.log10(np.abs(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.polar(np.angle(response), np.abs(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "omega_y_filtered = signal.convolve(omega_y, filter_coeffs, mode='valid')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "omega_y.shape, omega_y_filtered.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.plot(time[150:-150], omega_y[150:-150])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.plot(time[150:-150], omega_y_filtered)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Filter out anything but the range from 53-70 Hz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filter_coeffs = signal.firwin(301, [53, 70], pass_zero='bandpass', fs=1/delta_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "freqs, response = signal.freqz(filter_coeffs, fs=1/delta_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.plot(freqs, 20*np.log10(abs(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "omega_y_filtered = signal.convolve(omega_y, filter_coeffs, mode='valid')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.plot(time[150:-150], omega_y[150:-150])\n",
    "_ = plt.plot(time[150:-150], omega_y_filtered)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = plt.plot(fftpack.rfftfreq(omega_y_filtered.shape[0], delta_t)[::2],\n",
    "             fftpack.rfft(omega_y_filtered)[::2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spectrogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "freq, sp_time, Sxx = signal.spectrogram(omega_y, fs=1/delta_t)\n",
    "_ = plt.pcolormesh(sp_time, freq, Sxx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"color: DarkBlue;background-color: LightYellow\">\n",
    "<i>Exercise:</i> Try to obtain more structure by plotting the logarithm of the spectrogram.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = Axes3D(fig)\n",
    "_ = ax.plot_surface(sp_time, freq[:, np.newaxis], Sxx, cmap=cm.viridis)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
