{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fourier Transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import numpy as np\n",
    "import scipy.fft as ft\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "import matplotlib.ticker as tck\n",
    "rng = np.random.default_rng()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Fourier transform (FT) of a quantity $f(t)$ is defined as\n",
    "\n",
    "$$\n",
    "\\hat{f}(\\omega) = F(f(t)) = \\frac{1}{\\sqrt{2\\pi}} \\int_{-\\infty}^{\\infty} \\mathrm{d}t\\, f(t)\\, \\exp{ \\left(-2 \\pi i f t \\right)}\n",
    ".\n",
    "$$\n",
    "\n",
    "In signal processing, taking the integral from $t = -\\infty$ to $t = \\infty$ is not practical. Together with a finite sample space $t$, the integral can be turned into a (discrete) sum, giving a Discrete Fourier Transform (DFT)\n",
    "\n",
    "$$\n",
    "F(t) = \\frac{1}{\\sqrt{2\\pi}} \\sum_{k=0}^{N-1} f(t_k)\\, \\exp{ \\left(-i \\omega t_k \\right)}\n",
    "\\label{eq:DFT} \\tag{1}\n",
    ",\n",
    "$$\n",
    "where the sum iterates over each sample $t_k$ up to the end ($k=N-1$) of the signal."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interdependence $\\Delta t$, $T$, $\\Delta f$, $f_{min}$, $f_{max}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The highest frequency $f_\\mathrm{max}$ (also known as the Nyquist frequency) that can be (reliably) measured within a signal is dependent only on the sampling rate ($1/\\Delta t$) of the signal with\n",
    "$$\n",
    "    f_\\mathrm{max} = \\frac{1}{2} f_\\mathrm{sampling} = \\frac{1}{2} \\frac{1}{\\Delta t}\n",
    "    .\n",
    "$$\n",
    "Higher frequencies will fall in between two adjacent samples. This means that these frequencies are then 'measured\\` with differing phases in such a way that they are 'folded' around the Nyquist frequency and appear as aliased frequencies below it.\n",
    "\n",
    "Likewise, the lowest frequency depends on the length of the signal $T$, such that over the whole signal one oscillation can be found\n",
    "$$\n",
    "    f_\\mathrm{min} \\sim \\frac{1}{T}\n",
    "    .\n",
    "$$\n",
    "Even lower frequencies will no longer be detected as a frequency. Instead, they will introduce a bias for the $0 Hz$ component in the spectrum.\n",
    "\n",
    "\n",
    "$$\n",
    "\\delta f = \\frac{1}{T}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FTs of ideal signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ft_spectrum( signal, sample_rate, fft=None, freq=None, mask_bias=False):\n",
    "    \"\"\"Return a FT of $signal$, with corresponding frequencies\"\"\"\n",
    "    n_samples = len(signal)\n",
    "    real_signal = np.isrealobj(signal)\n",
    "    \n",
    "    if fft is None:\n",
    "        if real_signal:\n",
    "            fft = ft.rfft\n",
    "            freq = ft.rfftfreq\n",
    "        else:\n",
    "            fft = ft.fft\n",
    "            freq = ft.fftfreq\n",
    "\n",
    "    if freq is None:\n",
    "        freq = ft.fftfreq\n",
    "    \n",
    "    spectrum = fft(signal) / sample_rate\n",
    "    freqs = freq(n_samples, 1/sample_rate)\n",
    "    \n",
    "    if not mask_bias:\n",
    "        return spectrum, freqs\n",
    "    else:\n",
    "        return spectrum[1:], freqs[1:]\n",
    "\n",
    "    \n",
    "def plot_spectrum( ax, spectrum, freqs, plot_complex=False, plot_power=False, plot_amplitude=None):\n",
    "    \"\"\" Plot a signal's spectrum on an Axis object\"\"\"\n",
    "    plot_amplitude = plot_amplitude or (not plot_power and not plot_complex)\n",
    "    alpha = 1\n",
    "    \n",
    "    ax.set_title(\"Spectrum\")\n",
    "    ax.set_xlabel(\"f (Hz)\")\n",
    "    ylabel = \"\"\n",
    "    if plot_amplitude or plot_complex:\n",
    "        ylabel = \"Amplitude\"\n",
    "    if plot_power:\n",
    "        if ylabel:\n",
    "            ylabel += \"|\"\n",
    "        ylabel += \"Power\"\n",
    "    ax.set_ylabel(ylabel)\n",
    "\n",
    "    if plot_complex:\n",
    "        alpha = 0.5\n",
    "        ax.plot(freqs, np.real(spectrum), '.-', label='Real', alpha=alpha)\n",
    "        ax.plot(freqs, np.imag(spectrum), '.-', label='Imag', alpha=alpha)\n",
    "\n",
    "    if plot_power:\n",
    "        ax.plot(freqs, np.abs(spectrum)**2, '.-', label='Power', alpha=alpha)\n",
    "    \n",
    "    if plot_amplitude:\n",
    "        ax.plot(freqs, np.abs(spectrum), '.-', label='Abs', alpha=alpha)\n",
    "\n",
    "    ax.legend()\n",
    "\n",
    "    return ax\n",
    "\n",
    "\n",
    "def plot_phase( ax, spectrum, freqs, ylim_epsilon=0.5):\n",
    "    ax.set_ylabel(\"Phase\")\n",
    "    ax.set_xlabel(\"f (Hz)\")\n",
    "\n",
    "    ax.plot(freqs, np.angle(spectrum), '.-')\n",
    "    ax.set_ylim(-1*np.pi - ylim_epsilon, np.pi + ylim_epsilon)\n",
    "    \n",
    "    return ax\n",
    "\n",
    "\n",
    "def plot_combined_spectrum(spectrum, freqs, \n",
    "                           spectrum_kwargs={}, fig=None, gs=None):\n",
    "    \"\"\"Plot both the frequencies and phase in one figure.\"\"\"\n",
    "    \n",
    "    # configure plotting layout\n",
    "    if fig is None:\n",
    "        fig = plt.figure(figsize=(8, 16))\n",
    "\n",
    "    if gs is None:\n",
    "        gs = gridspec.GridSpec(2, 1, figure=fig, height_ratios=[3,1], hspace=0)\n",
    "\n",
    "    ax1 = fig.add_subplot(gs[:-1, -1])\n",
    "    ax2 = fig.add_subplot(gs[-1, -1], sharex=ax1)\n",
    "\n",
    "    axes = np.array([ax1, ax2])\n",
    "    \n",
    "    # plot the spectrum \n",
    "    plot_spectrum(ax1, spectrum, freqs, **spectrum_kwargs)\n",
    "\n",
    "    # plot the phase\n",
    "    plot_phase(ax2, spectrum, freqs)\n",
    "\n",
    "    ax1.xaxis.tick_top()\n",
    "    [label.set_visible(False) for label in ax1.get_xticklabels()]\n",
    "    \n",
    "    return fig, axes\n",
    "    \n",
    "\n",
    "def plot_signal( ax, signal, sample_rate = 1):\n",
    "    ax.set_title(\"Signal\")\n",
    "    ax.set_xlabel(\"t (s)\")\n",
    "    ax.set_ylabel(\"A(t)\")\n",
    "\n",
    "    ax.plot(np.arange(len(signal)) / sample_rate, signal)\n",
    "    \n",
    "    return ax\n",
    "\n",
    "\n",
    "def plot_signal_and_spectrum( signal, sample_rate, title=None,\n",
    "                             ft_kwargs = {}, spectrum_kwargs = {}):\n",
    "    \"\"\"Plot a signal and its spectrum in one go.\"\"\"\n",
    "    fig = plt.figure(figsize=(16, 4))\n",
    "    \n",
    "    if title:\n",
    "        fig.suptitle(title)\n",
    "\n",
    "    # setup plot layout\n",
    "    gs0 = gridspec.GridSpec(1, 2, figure=fig)\n",
    "    gs00 = gs0[0].subgridspec(1, 1)\n",
    "    gs01 = gs0[1].subgridspec(2, 1, height_ratios=[3,1], hspace=0)\n",
    "      \n",
    "    # plot the signal\n",
    "    ax1 = fig.add_subplot(gs00[0, 0])\n",
    "    plot_signal(ax1, signal)\n",
    "    \n",
    "    # plot spectrum\n",
    "    signal_fft, freqs = ft_spectrum(signal, sample_rate, **ft_kwargs)\n",
    "    _, (ax2, ax3) = plot_combined_spectrum(signal_fft, freqs, spectrum_kwargs = spectrum_kwargs, fig=fig, gs=gs01)\n",
    "\n",
    "    # return the axes\n",
    "    axes = np.array([ax1, ax2, ax3])\n",
    "    \n",
    "    return (fig, axes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_rate = 1/1e-3 # s\n",
    "n_samples = int(1e3)\n",
    "t = np.arange(n_samples) / sample_rate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FTs: constant signal\n",
    "\n",
    "A constant signal (also termed offset or bias) results in only a single component in the Fourier transform, at the zero frequency.\n",
    "\n",
    "Since often such a component is included in a signal (for example when a signal is not on average zero), it can help to mask out the zero frequency when plotting the spectrum.\n",
    "Doing this for a constant signal, we get a totally flat spectrum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 20728 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 19548 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "signal_constant = np.ones(int(n_samples))\n",
    "\n",
    "fig, axes = plot_signal_and_spectrum(signal_constant, sample_rate, \"Constant signal\")\n",
    "axes.flat[1].set_xlim(0,10);\n",
    "\n",
    "fig, axes = plot_signal_and_spectrum(signal_constant, sample_rate, \"Constant signal (masked zero)\", ft_kwargs={'mask_bias':True})\n",
    "axes.flat[1].set_xlim(0,10);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FTs: noise (normal distributed)\n",
    "\n",
    "A FT of normal distributed noise will not have easily identifiable peaks in the spectrum.\n",
    "Instead, the spectrum is on average a bit flat, with only low power in any individual peak."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 113848 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "signal_normal_noise = rng.normal(size=n_samples)\n",
    "\n",
    "plot_signal_and_spectrum(signal_normal_noise, sample_rate, \"Noise (normal distributed)\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FTs: sine wave"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(95.0, 105.0)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 90056 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "f = 1e2 # Hz\n",
    "if f > sample_rate/2:\n",
    "    print(\"Sampling a frequency above the sample_rate/2 gives problems\")\n",
    "\n",
    "signal_sine = np.sin( 2*np.pi*f*t )\n",
    "\n",
    "if False:\n",
    "    # aliased peak\n",
    "    f_alias = (1.1* sample_rate/2)\n",
    "    signal_sine += 1/2*np.sin( 2*np.pi* f_alias *t)\n",
    "\n",
    "fig, axes = plot_signal_and_spectrum(signal_sine, sample_rate, \"Single frequency at $f = {:g}$MHz\".format(f/10**6))\n",
    "axes.flat[1].axvline(sample_rate/2, color='r', label='Nyquist Frequency');\n",
    "axes.flat[1].legend();\n",
    "axes.flat[1].set_xlim(f-5, f+5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A sine wave will give very clear peaks at the frequency of the sine wave."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 124756 characters]\n",
      "text/plain": [
       "<Figure size 1152x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "with_spectrum = True\n",
    "phase_offsets = np.linspace(0, 2*np.pi, 8, endpoint=False) # rad\n",
    "f = 99 # Hz\n",
    "\n",
    "if f > sample_rate/2:\n",
    "    print(\"Sampling a frequency above the sample_rate/2 gives problems\")\n",
    "\n",
    "if with_spectrum:\n",
    "    fig, axes = plt.subplots(2,1, sharex=True, figsize=(16,8))\n",
    "else:\n",
    "    fig, ax = plt.subplots(figsize=(16,4))\n",
    "    axes = np.array([None, ax])\n",
    "    \n",
    "if with_spectrum:\n",
    "    axes[0].set_ylabel(\"Amplitude\")\n",
    "    axes[0].axvline(f, alpha=0.3, label='f = {:g}Hz'.format(f))\n",
    "\n",
    "axes[1].axvline(f, alpha=0.3, label='f = {:g}Hz'.format(f))\n",
    "axes[1].set_ylabel(\"Phase\")\n",
    "axes[1].set_xlabel(\"f\")\n",
    "axes[1].set_xlim( f-2, f+2)\n",
    "axes[1].set_ylim( -1 - 0.2, 1+ 0.2)\n",
    "axes[1].grid()\n",
    "axes[1].yaxis.set_major_formatter(tck.FormatStrFormatter('%g $\\pi$'))\n",
    "axes[1].yaxis.set_major_locator(tck.MultipleLocator(base=0.5))\n",
    "\n",
    "for phase in phase_offsets:\n",
    "    signal = np.sin( 2*np.pi*f*t + phase )\n",
    "    \n",
    "    ft_signal, freqs = ft_spectrum(signal, sample_rate)\n",
    "    \n",
    "    if with_spectrum:\n",
    "        axes[0].plot(freqs, np.real(ft_signal),  '.-', label='Re $\\\\varphi = {}\\\\pi$'.format(phase/np.pi))\n",
    "    \n",
    "    axes[1].plot(freqs, np.angle(ft_signal)/np.pi, '.-', label='$\\\\varphi = {}\\\\pi$'.format(phase/np.pi))\n",
    "\n",
    "# trigger legends\n",
    "if with_spectrum:\n",
    "    axes[0].legend(loc='center right');\n",
    "axes[1].legend(loc='center right');\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a sine wave, the phase of the spectrum at it's frequency is $\\theta( \\varphi = 0 ) = - \\dfrac{\\pi}{2}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FTs: step function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 37068 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 41328 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "t_0 = t[n_samples//4]\n",
    "signal_step = np.heaviside(t - t_0, 1)\n",
    "fig, axes = plot_signal_and_spectrum(signal_step, sample_rate, \"Step function $t_0 = {:g}ns$\".format(t_0 * 10**9), ft_kwargs={'mask_bias':True});\n",
    "axes.flat[1].set_xlim(0, 20);\n",
    "\n",
    "t_0 = t[n_samples//2]\n",
    "signal_step = np.heaviside(t - t_0, 1)\n",
    "fig, axes = plot_signal_and_spectrum(signal_step, sample_rate, \"Step function $t_0 = {:g}ns$\".format(t_0 * 10**9), ft_kwargs={'mask_bias':True});\n",
    "axes.flat[1].set_xlim(0, 20);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FTs: delta peak\n",
    "\n",
    "All frequencies are accounted for with the same amplitude, but the phase will change depending on the position of the delta peak within the signal.\n",
    "\n",
    "Still, the individual phases are linearly related to their corresponding frequencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 27180 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 23604 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 64420 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 56692 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 64400 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "for n in [1, -1, 4, 2, -4]:\n",
    "    index_t0 = min(n_samples//n, n_samples-1)\n",
    "\n",
    "    t_0 = t[index_t0]\n",
    "    signal_step = np.zeros(n_samples)\n",
    "    signal_step[index_t0] = 1\n",
    "\n",
    "    fig, axes = plot_signal_and_spectrum(signal_step, sample_rate, ft_kwargs={'mask_bias':True}, spectrum_kwargs={'plot_complex':True, 'plot_amplitude':True});\n",
    "    fig.suptitle(\"Delta function at ${}/{}$th of the sample\".format(np.sign(n), np.abs(n)))\n",
    "    axes.flat[1].set_xlim(0, 20);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## FTs: gaussian with sine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 65112 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "f = 1e1 # Hz\n",
    "if f > sample_rate/2:\n",
    "    print(\"Sampling a frequency above the sample_rate/2 gives problems\")\n",
    "\n",
    "def gaussian(x, mu=0, sigma=1 ):\n",
    "    return 1/(sigma*np.sqrt(2*np.pi)) * np.exp(- (x-mu)**2 / sigma**2 )\n",
    "\n",
    "def dgaussian(x, mu=0, sigma=1, dmu=1, dsigma=0):\n",
    "    return gaussian(x, mu, sigma)   + gaussian(x, mu+dmu, sigma+dsigma)\n",
    "\n",
    "if not True:\n",
    "    signal = gaussian(t, t[2*len(t)//5], 100/len(t)) - gaussian(t, t[3*len(t)//5], 100/len(t))\n",
    "else:\n",
    "    signal = gaussian(t, t[2*len(t)//5], 100/len(t)) * np.sin(2*np.pi*f*t)\n",
    "\n",
    "fig, axes = plot_signal_and_spectrum(signal, sample_rate, \"Gaussian pulse\", spectrum_kwargs={'plot_complex':True, 'plot_amplitude':True})\n",
    "axes.flat[1].axvline(sample_rate/2, color='r', label='Nyquist Frequency')\n",
    "axes.flat[2].axvline(sample_rate/2, color='r')\n",
    "axes.flat[1].legend();\n",
    "\n",
    "if not True:\n",
    "    axes.flat[1].set_xlim(0, 20);\n",
    "else:\n",
    "    axes.flat[1].set_xlim(max(0, f - max(f/5, 10)), f + max(f/5, 10));\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### FTs: delta peak: group velocity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above plots, the phases follow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Peculiarities $F^{-1}( F(A(t)) ) \\sim A(t)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38336 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "t_0 = t[n_samples//4]\n",
    "signal_step = np.heaviside(t - t_0, 1) + 1\n",
    "\n",
    "signal_fft = ft.rfft(signal_step) / sample_rate\n",
    "signal_ifft = ft.irfft(signal_fft, len(signal_fft))\n",
    "ft_freqs = ft.rfftfreq(n_samples, 1/sample_rate)\n",
    "\n",
    "\n",
    "alpha = 0.8\n",
    "fig, (ax1, ax2) = plt.subplots(1,2, figsize=(16,4))\n",
    "ax1.set_title(\"Spectrum\")\n",
    "ax1.set_xlabel(\"f (Hz)\")\n",
    "ax1.set_ylabel(\"Amplitude\")\n",
    "ax1.set_xlim(-1,20)\n",
    "ax1.plot(ft_freqs, np.real(signal_fft), '.-', label='Real', alpha=alpha)\n",
    "ax1.plot(ft_freqs, np.imag(signal_fft), '.--', label='Imag', alpha=alpha)\n",
    "ax1.plot(ft_freqs, np.abs(signal_fft),  '.-', label='Abs', alpha=alpha)\n",
    "ax1.legend()\n",
    "\n",
    "# Signal and Inversed FFT\n",
    "ax2.set_title(\"Signal and Inversed FFT\")\n",
    "ax2.set_xlabel(\"t (s)\")\n",
    "ax2.set_ylabel(\"A(t)\")\n",
    "ax2.plot(np.arange(len(signal_ifft)) / sample_rate * 2, np.abs(signal_ifft)*sample_rate/2, label=\"Inversed FFT\")\n",
    "ax2.plot(np.arange(len(signal_step)) / sample_rate, signal_step, label=\"Original\")\n",
    "ax2.legend()\n",
    "\n",
    "if False:\n",
    "    ax2.set_xlim(t_0*(1-0.2), t_0*(1+0.2))\n",
    "\n",
    "plt.show();\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DFT\n",
    "\n",
    "A naive implementation of a DFT is simple to setup by implementing Eq. (1)\n",
    "\n",
    "$$\n",
    "\\hat{f}(\\omega) = F(t) = \\frac{1}{\\sqrt{2\\pi}} \\sum_{k=0}^{N-1} f(t_k)\\, \\exp{ \\left(-i \\omega t_k \\right)}\n",
    ".\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def naive_DFT( signal ):\n",
    "    \"\"\"\n",
    "    Calculate DFT in a naive way, returns a complex DFT.\n",
    "    \"\"\"\n",
    "\n",
    "    N = np.size(signal)\n",
    "\n",
    "    x_k = 2*np.pi/N*np.arange(N)\n",
    "    ft = np.zeros(N, dtype=np.complex128)\n",
    "    for j in range(N):\n",
    "        ft[j] = np.dot(signal, np.exp(- 1j * j * x_k ))\n",
    "\n",
    "    return ft"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DFT: Cooley and Tukey (or Divide and Conquer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A smarter DFT (Cooley and Tukey) can be implemented by splitting the sum in Eq. (1) into two parts: $k$ even and $k$ odd.\n",
    "\n",
    "With a uniform sample space $t$, the label $t_k$ can be rewritten to $t_k = k \\Delta t = $.\n",
    "\n",
    "This gives\n",
    "$$\n",
    "F(t) = \n",
    "    \\frac{1}{\\sqrt{2\\pi}} \\sum_{k=0}^{N/2-1} f(t_{2k})\\, \\exp{ \\left(-i \\omega \\, t_{2k} \\right)}\n",
    "    +\n",
    "    \\frac{1}{\\sqrt{2\\pi}} \\sum_{k=0}^{N/2-1} f(t_{2k+1})\\, \\exp{ \\left(-i \\omega \\, t_{2k+1} \\right)}\n",
    "    \\\\\n",
    "    =\n",
    "    \\frac{1}{\\sqrt{2\\pi}} \\sum_{k=0}^{N/2-1} f(t_{2k})\\, \\exp{ \\left(-i \\omega \\, 2k \\Delta t \\right)}\n",
    "    +\n",
    "    \\frac{1}{\\sqrt{2\\pi}} \\sum_{k=0}^{N/2-1} f(t_{2k+1})\\, \\exp{ \\left(-i \\omega \\, 2k \\Delta t \\right)}\\exp{ \\left(-i \\omega \\, \\Delta t \\right)}\n",
    "    \\\\\n",
    "    =\n",
    "    \\frac{1}{\\sqrt{2\\pi}} \\sum_{k=0}^{N/2-1} f(t_{2k})\\, \\exp{ \\left(-i \\omega \\, 2k \\Delta t \\right)}\n",
    "    +\n",
    "    \\frac{\\exp{ \\left(-i \\omega \\, \\Delta t \\right)} }{\\sqrt{2\\pi}} \\sum_{k=0}^{N/2-1} f(t_{2k+1})\\, \\exp{ \\left(-i \\omega \\, 2k \\Delta t \\right)}\n",
    "    \\\\\n",
    "    =\n",
    "    F(t_{even}) + F(t_{odd}) \\exp{ \\left(-i \\omega \\, \\Delta t \\right)}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cooley_tukey_DFT( signal ):\n",
    "    \"\"\"\n",
    "    Calculate a DFT according to the Cooley and Tukey algorithm (radix-2)\n",
    "    also known as Divide and Conquer.\n",
    "    \n",
    "    It divides the terms into 2 streams with odd and even index.\n",
    "    \"\"\"\n",
    "    radix = 2\n",
    "\n",
    "    N = np.size(signal)\n",
    "    \n",
    "    if N <= radix:\n",
    "        return naive_DFT(signal)\n",
    "    \n",
    "    if N % radix > 0:\n",
    "        raise ValueError(\"Provide a signal with  N%{radix} == 0 values.\".format(radix=radix))\n",
    "    \n",
    "    # divide and conquer\n",
    "    DFT_even = cooley_tukey_DFT(signal[::radix])\n",
    "    DFT_odd  = cooley_tukey_DFT(signal[1::radix])\n",
    "    \n",
    "    phase_terms = np.exp( -2j * np.pi * np.arange(0, N) / N )\n",
    "    \n",
    "    return np.tile(DFT_even, 2) + np.tile(DFT_odd,2) * phase_terms\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 81516 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "dft_sample_rate = 1/1e-3 # s\n",
    "dft_n_samples = 2**10\n",
    "dft_t = np.arange(dft_n_samples) / dft_sample_rate\n",
    "\n",
    "f = 1e2 # Hz\n",
    "if f > dft_sample_rate/2:\n",
    "    print(\"Sampling a frequency above the sample_rate/2 gives problems\")\n",
    "\n",
    "signal_sine = np.sin( 2*np.pi*f*dft_t )\n",
    "\n",
    "fig, axes = plot_signal_and_spectrum(signal_sine, dft_sample_rate, \"Single frequency at $f = {:g}$MHz\".format(f/10**6), ft_kwargs={'fft': cooley_tukey_DFT})\n",
    "axes.flat[1].axvline(dft_sample_rate/2, color='r', label='Nyquist Frequency');"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
