{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ThinkDSP\n",
    "\n",
    "This notebook contains solutions to exercises in Chapter 7: Discrete Fourier Transform\n",
    "\n",
    "Copyright 2015 Allen Downey\n",
    "\n",
    "License: [Creative Commons Attribution 4.0 International](http://creativecommons.org/licenses/by/4.0/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "PI2 = 2 * np.pi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1\n",
    "\n",
    "In this chapter, I showed how we can express the DFT and inverse DFT\n",
    "as matrix multiplications.  These operations take time proportional to\n",
    "$N^2$, where $N$ is the length of the wave array.  That is fast enough\n",
    "for many applications, but there is a faster\n",
    "algorithm, the Fast Fourier Transform (FFT), which takes time\n",
    "proportional to $N \\log N$.\n",
    "\n",
    "The key to the FFT is the Danielson-Lanczos lemma:\n",
    "\n",
    "$DFT(y)[n] = DFT(e)[n] + \\exp(-2 \\pi i n / N) DFT(o)[n]$\n",
    "\n",
    "Where $ DFT(y)[n]$ is the $n$th element of the DFT of $y$; $e$ is the even elements of $y$, and $o$ is the odd elements of $y$.\n",
    "\n",
    "This lemma suggests a recursive algorithm for the DFT:\n",
    "\n",
    "1. Given a wave array, $y$, split it into its even elements, $e$, and its odd elements, $o$.\n",
    "\n",
    "2. Compute the DFT of $e$ and $o$ by making recursive calls.\n",
    "\n",
    "3. Compute $DFT(y)$ for each value of $n$ using the Danielson-Lanczos lemma.\n",
    "\n",
    "For the base case of this recursion, you could wait until the length\n",
    "of $y$ is 1.  In that case, $DFT(y) = y$.  Or if the length of $y$\n",
    "is sufficiently small, you could compute its DFT by matrix multiplication,\n",
    "possibly using a precomputed matrix.\n",
    "\n",
    "Hint: I suggest you implement this algorithm incrementally by starting\n",
    "with a version that is not truly recursive.  In Step 2, instead of\n",
    "making a recursive call, use `dft` or `np.fft.fft`.  Get Step 3 working,\n",
    "and confirm that the results are consistent with the other\n",
    "implementations.  Then add a base case and confirm that it works.\n",
    "Finally, replace Step 2 with recursive calls.\n",
    "\n",
    "One more hint: Remember that the DFT is periodic; you might find `np.tile` useful.\n",
    "\n",
    "You can read more about the FFT at https://en.wikipedia.org/wiki/Fast_Fourier_transform."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As the test case, I'll start with a small real signal and compute its FFT:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.2+0.j  -1.2-0.2j  0.2+0.j  -1.2+0.2j]\n"
     ]
    }
   ],
   "source": [
    "ys = [-0.5, 0.1, 0.7, -0.1]\n",
    "hs = np.fft.fft(ys)\n",
    "print(hs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's my implementation of DFT from the book:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def dft(ys):\n",
    "    N = len(ys)\n",
    "    ts = np.arange(N) / N\n",
    "    freqs = np.arange(N)\n",
    "    args = np.outer(ts, freqs)\n",
    "    M = np.exp(1j * PI2 * args)\n",
    "    amps = M.conj().transpose().dot(ys)\n",
    "    return amps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can confirm that this implementation gets the same result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5.864775846765962e-16"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hs2 = dft(ys)\n",
    "np.sum(np.abs(hs - hs2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a step toward making a recursive FFT, I'll start with a version that splits the input array and uses np.fft.fft to compute the FFT of the halves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def fft_norec(ys):\n",
    "    N = len(ys)\n",
    "    He = np.fft.fft(ys[::2])\n",
    "    Ho = np.fft.fft(ys[1::2])\n",
    "    \n",
    "    ns = np.arange(N)\n",
    "    W = np.exp(-1j * PI2 * ns / N)\n",
    "    \n",
    "    return np.tile(He, 2) + W * np.tile(Ho, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we get the same results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hs3 = fft_norec(ys)\n",
    "np.sum(np.abs(hs - hs3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can replace `np.fft.fft` with recursive calls, and add a base case:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def fft(ys):\n",
    "    N = len(ys)\n",
    "    if N == 1:\n",
    "        return ys\n",
    "    \n",
    "    He = fft(ys[::2])\n",
    "    Ho = fft(ys[1::2])\n",
    "    \n",
    "    ns = np.arange(N)\n",
    "    W = np.exp(-1j * PI2 * ns / N)\n",
    "    \n",
    "    return np.tile(He, 2) + W * np.tile(Ho, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we get the same results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.6653345369377348e-16"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hs4 = fft(ys)\n",
    "np.sum(np.abs(hs - hs4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This implementation of FFT takes time proportional to $n \\log n$.  It also takes space proportional to $n \\log n$, and it wastes some time making and copying arrays.  It can be improved to run \"in place\"; in that case, it requires no additional space, and spends less time on overhead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}