{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Week 12 Day 2: Signal filtering\n",
    "\n",
    "## Objectives\n",
    "\n",
    "* Review Fourier transforms\n",
    "* Filter signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional\n",
    "# %matplotlib notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convolutions and Fourier transforms\n",
    "\n",
    "Let's say you have a sample with N points over time T."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 1_000\n",
    "T = 2*np.pi\n",
    "\n",
    "t = np.linspace(0, T, N, endpoint=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define a plotting function that shows the signal and it's Fourier transform. Unlike before, let's get the $x$-axis units right."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_with_fft(t, signal):\n",
    "    f = np.fft.rfft(signal)\n",
    "    ft = np.fft.rfftfreq(len(t), t[1] - t[0])\n",
    "    \n",
    "    fig, axs = plt.subplots(1,2, figsize=(7,3))\n",
    "    axs[0].plot(t,signal)\n",
    "    axs[0].set_xlabel('t [s]')\n",
    "    axs[1].plot(ft,np.abs(f), color='black')\n",
    "    axs[1].set_xlabel('freq [Hz]')\n",
    "    axs[1].set_yscale('log')\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = 1 / (1 - .9 * np.sin(t))\n",
    "#y += np.random.rand(N)*2-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_with_fft(t, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convolution\n",
    "\n",
    "$$\n",
    "g(t) = \\int^{+\\infty}_{-\\infty} f(\\tau) \\, h(t-\\tau) \\, d \\tau\n",
    "\\equiv f(t) * h(t)\n",
    "\\tag{1}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Autocorrelation is just the above equation with both $h(t)=f(t)$. (Assuming a real signal - otherwise correlation needs the complex conjugate of one of the functions.)\n",
    "\n",
    "Now we compute the autocorrelation of the signal, and plot again. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = np.convolve(y, np.tile(y,2), 'valid')[:-1]\n",
    "tz = np.linspace(-T/2, T/2, len(z))\n",
    "print(len(z), z[0], z[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_with_fft(tz, z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The convolution theorem\n",
    "\n",
    "A convolution in real space is a multiplication in Fourier space!\n",
    "\n",
    "$$\n",
    "g(t) = f(t) * h(t)\n",
    "$$\n",
    "\n",
    "$$\n",
    "G(\\omega) = \\sqrt{2\\pi} F(\\omega) H(\\omega)\n",
    "$$\n",
    "\n",
    "Note that the $\\sqrt{2\\pi}$ is from the definition we have of the FT. I believe numpy and most software put $2\\pi$ in the iFT instead of splitting it up."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(y)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = np.fft.rfft(y)\n",
    "f**=2\n",
    "yy = np.fft.irfft(f)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(yy)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filters\n",
    "\n",
    "Now, let's define a filter $H(\\omega)$. We can try physically inspired ones from RC circuits:\n",
    "\n",
    "$$\n",
    "H_\\textrm{low}(\\omega) =\n",
    "  \\frac{1}{1 + i \\omega \\tau_\\mathrm{RC}}\n",
    "$$\n",
    "\n",
    "$$\n",
    "H_\\textrm{high}(\\omega) =\n",
    "  \\frac{i \\omega \\tau_\\mathrm{RC}}{1 + i \\omega \\tau_\\mathrm{RC}}\n",
    "$$\n",
    "\n",
    "Or, we can build them:\n",
    "\n",
    "$$\n",
    "H(\\omega) =\n",
    "  \\sum_{n=0}^{N} c_n e^{-i n \\omega \\tau}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make a noisy signal and try a filter. Let's start with a rectangular lowpass filter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = 1 / (1 - .9 * np.sin(t))\n",
    "y += np.random.rand(N)*2-1\n",
    "ft = np.fft.rfftfreq(len(t), t[1] - t[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = np.fft.rfft(y)\n",
    "f[ft > 6.5] = 0\n",
    "yy = np.fft.irfft(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1,2,figsize=(6,3))\n",
    "axs[0].plot(t, y)\n",
    "axs[1].plot(t, yy)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can do this in the time domain with convolutions as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "B = 1\n",
    "ift = np.linspace(-N/T/2, N/T/2, len(y)*2)\n",
    "\n",
    "H = 2*B * np.sinc(2*B*ift)\n",
    "z = np.convolve(y, H, 'valid')[:-1]\n",
    "i_H = np.abs(np.fft.ifft(H))\n",
    "\n",
    "\n",
    "fig, axs = plt.subplots(2,2,figsize=(6,6))\n",
    "axs[0,0].plot(t, y)\n",
    "axs[0,1].plot(t, z)\n",
    "axs[1,0].plot(ift, H)\n",
    "axs[1,1].plot(ift, np.fft.fftshift(i_H))\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See https://docs.scipy.org/doc/scipy/reference/signal.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.signal as signal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try a low-pass butterworth filter. https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.lfilter.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b, a = signal.butter(3, 0.05)\n",
    "zi = signal.lfilter_zi(b, a)\n",
    "z, _ = signal.lfilter(b, a, y, zi=zi*y[0])\n",
    "z2, _ = signal.lfilter(b, a, z, zi=zi*z[0])\n",
    "Y = signal.filtfilt(b, a, y)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(t,y)\n",
    "plt.plot(t,Y)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Aside: Plotting the wave packet in the book"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(0, T, N, endpoint=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k_0 = 10\n",
    "σ_0 = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ψ = np.exp(-1/2 * ((x - 5)/(σ_0))**2 ) * np.exp(1j * k_0 * x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(x, ψ.real, color='blue')\n",
    "ax.plot(x, ψ.imag, color='red')\n",
    "ax.plot(x, np.abs(ψ), color='black')\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('ψ')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.plot(t, ψ.real, ψ.imag)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "compclass",
   "language": "python",
   "name": "compclass"
  },
  "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}