Bootcamp 4 - Fourier Analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "64dcdbbd-16e6-471c-a1a9-0f90be0bb92c",
   "metadata": {},
   "source": [
    "<h1>Scientific coding bootcamp notebook 4: Fourier analysis and the FFT.</h1>\n",
    "\n",
    "The big idea of Fourier analysis is that complex exponentials like $e^{2πinx}$ ($n$ fixed, $x$ a variable) for an orthogonal basis for many interesting function spaces.  The change of basis operation for representing a function in this basis is some sort of \"Fourier transform\", though the exact nature (and name) of this operation depends on whether you are working with $x$ continuous vs. discrete and bounded vs. unbounded.  Here's a table of the different cases: \n",
    "\n",
    "\n",
    "|  | Discrete | Continuous |\n",
    "| :-: | :-: | :-: |\n",
    "| **Bounded** | DFT | Fourier series |\n",
    "| **Unbounded** | z-transform | Fourier transform |\n",
    "\n",
    "We will start with Fourier series, as they are maybe the easiest to understand, but our main goal is the fully discrete case (DFT, which stands for \"Discrete Fourier Transform\"), as this is the most relevant for numerical work."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4196ec4-1834-485c-aad8-e01183d94d47",
   "metadata": {},
   "source": [
    "<h2>Fourier series</h2>\n",
    "\n",
    "In all versions of Fourier analysis, the key questions are \"What is the vector space?\" and \"What is the inner product?\".  For Fourier series, the vector space is functions on some bounded interval, which will always take to be the unit interval $[0,1]$.  The inner product is the obvious one,\n",
    "$$\n",
    "\\langle f |g \\rangle := \\int_0^1 f(x) g^*(x) dx.\n",
    "$$\n",
    "\n",
    "In this setting, complex exponentials like $e^{2πinx}$, with $n$ an integer, form an orthonormal basis for periodic 1D functions on the unit interval $[0,1]$.  That is,\n",
    "$$\n",
    "\\int_0^1 e^{2\\pi i n x} e^{-2\\pi i m x} dx = \\delta_{nm}.\n",
    "$$\n",
    "Let's expand some functions on $[0,1]$ with respect to this basis to see how it works. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e46130b5-73ed-44c9-b7f3-825c75772594",
   "metadata": {},
   "outputs": [],
   "source": [
    "using FFTW, Plots, LaTeXStrings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49661202-c9ad-45f9-b7ce-976d0eb52423",
   "metadata": {},
   "outputs": [],
   "source": [
    "const Lattice = AbstractRange\n",
    "\n",
    "function expOverlap(grid::Lattice,f::Vector{T},n::Real) where {T<:Number}\n",
    "    integrand = f .* exp.(-2π*im*grid*n)\n",
    "    return sum(integrand) * step(grid)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdbcc238-15ee-457e-982b-ac9afa350527",
   "metadata": {},
   "outputs": [],
   "source": [
    "#= Fill in the code below to decompose the function f as a sum of complex exponentials.  \n",
    "Check that it matches f reasonably well, and experiment with different grid spacings dx and numbers of modes. \n",
    "=#\n",
    "dx = 0.01\n",
    "grid = 0:dx:1-dx/2\n",
    "f = (x -> x < 0.5 ? x : 1.0-x).(grid)      # The  \"... ? ... : ...\" is called a ternary operator.  Look up how it works if it's unfamiliar.\n",
    "modes = -5:5\n",
    "coeffs = # Your code here.\n",
    "plot(grid,f,label=\"f\")\n",
    "plot!(grid,real.(sum(exp.(2pi*im*grid*modes[n])*coeffs[n] for n=1:length(modes))),label=\"Real{Mode Sum}\")\n",
    "plot!(grid,imag.(sum(exp.(2pi*im*grid*modes[n])*coeffs[n] for n=1:length(modes))),label=\"Imag{Mode Sum}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53250492-9cc3-4788-92ca-c1663b83e9bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#= Repeat the above for the following functions:\n",
    "f = (x -> sin(2pi*x^2)).(grid)\n",
    "f = (x -> x < 0.5 ? 1.0 : 0.0).(grid)    # This is called a square wave. \n",
    "=#"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21386297-4c66-421b-ad0e-7dd246ee6322",
   "metadata": {},
   "outputs": [],
   "source": [
    "#= For the square wave function f, see what happens as you increase the number of modes from -40:40 to -50:50 to -80:80\n",
    "at a step size of dx=0.01. Then make the grid spacing smaller and repeat.  \n",
    "(The wiggles at the jumps are called Gibbs' phenomenon, and are not the main interest here.) \n",
    "Can you understand the behavior you observe in terms of the dimension of the relevant vector space? \n",
    "=#"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca8fea52-730d-4b86-8d5c-abb489da7188",
   "metadata": {},
   "source": [
    "<h2>Discrete Fourier Transform</h2>\n",
    "\n",
    "When we expanded the square wave, we saw some numerical artifacts due to the fact that we are discretizing continuous integrals onto finite grids.  A convenient way to deal with this is to treat functions and integrals as discrete from the start.  That is, instead of functions $f(x)$ of a continuous variable $x\\in [0,1]$, we consider *sequences* $f_j = f(x_j)$ of the functions' values sampled at a grid of points $x_j = j/N$, $j=0,1,\\dots,N-1$.  Then instead of integrals $\\int_0^1 f(x) dx$, we consider sums $\\sum_{j=0}^N f_j$. \n",
    "\n",
    "Sampled complex exponentials are still orthonormal in the following sense: \n",
    "$$\n",
    "\\frac{1}{N}\\sum_{j=0}^N e^{2\\pi i n j/N} e^{-2\\pi i m j/N} = \\delta_{nm}.\n",
    "$$\n",
    "This means we can use sampled complex exponentials of this form as a basis for sequences of length $N$. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f78d0476-7141-429d-aa26-a5f694185d37",
   "metadata": {},
   "source": [
    "**Question:** How many such complex exponentials are needed to form a basis?  What can you say about the linear dependence of a collection of complex exponentials with more than this number? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6d416b6-f948-4c2f-b5dd-46482d166e8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#= Plot exp.(2pi * im * n * j/N) for some reasonable value of N and several values of n.  \n",
    "Find some minimal collection of n's which form a linearly dependent set.  \n",
    "What is the linear dependence relationship?\n",
    "=#"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a045aabb-9aef-46eb-a118-9cced24fc7e9",
   "metadata": {},
   "source": [
    "Two complex exponentials $e^{2\\pi i n j/N}$ with different frequencies (i.e. different values of $n$) which have the same values at all grid points are called *aliases* of each other.  For which values of $n$ does $e^{2\\pi i n j/N}$ have an alias $e^{2\\pi i m j/N}$ with $|m| < |n|$? \n",
    "\n",
    "This is essentially the content of the \"Whittaker–Nyquist–Shannon sampling theorem\".  It says that if you want to accurately measure a function with frequency $f$, you need to sample it at a rate of at least $2f$.  Otherwise, the function will be *aliased*, which means it will look like a lower frequency function. \n",
    "\n",
    "(The full theorem goes a bit further and talks about perfectly reconstructing functions with bounded frequency content using infinitely many samples.  I don't find this particularly useful in practice.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f925001-e57b-4b1b-b3a3-83868a899913",
   "metadata": {},
   "source": [
    "<h3>DFT as a unitary matrix.</h3>\n",
    "\n",
    "Complex exponentials $\\frac{1}{\\sqrt{N}} e^{2\\pi i n j/N}$, $n=0,\\dots,N-1$, form a basis (the \"Fourier basis\") for length $N$ sequences.  Another basis is the \"Dirac basis\" $\\delta_{nj}$, $n=0,\\dots,N-1$.  We are thinking of $\\delta_{nj}$ here as a sequence of length $N$, with a $1$ in the $n$-th slot and zeros everywhere else.  E.g. \n",
    "$$\n",
    "\\delta_{2,j} = (0,0,1,0,0,\\dots,0,0,0)\n",
    "$$\n",
    "When you write a sequence as $(f_0, f_1, \\dots, f_{N-1})$, you are effectively giving its coefficients with respect to the Dirac basis.  To expand the same function in the complex exponential basis as \n",
    "$$\n",
    "f_j = \\sum_n c_n \\frac{1}{\\sqrt{N}} e^{2\\pi i n j / N}\n",
    "$$\n",
    "you can compute $c_n$ via the overlap sum \n",
    "$$\n",
    "c_n = \\sum_j f_j \\frac{1}{\\sqrt{N}} e^{-2\\pi i n j/N}.\n",
    "$$\n",
    "Alternatively, we can compute a change of basis matrix $U$ with entries $$U_{nm} = \\sum_j \\delta_{nj} \\frac{1}{\\sqrt{N}} e^{-2\\pi i m j/N} = \\frac{1}{\\sqrt{N}} e^{-2\\pi i m n / N}.$$  Note that this is the matrix of inner products between the Dirac basis and the Fourier basis.  In terms of this matrix, the coefficients of $(f_0, f_1, \\dots, f_{N-1})$ are just \n",
    "$$\n",
    "\\begin{bmatrix}\n",
    "           c_0 \\\\ c_1 \\\\ \\vdots \\\\ c_{N-1}\n",
    "         \\end{bmatrix} = U \\begin{bmatrix}\n",
    "           f_0 \\\\ f_1 \\\\ \\vdots \\\\ f_{N-1}\n",
    "         \\end{bmatrix}.\n",
    "$$\n",
    "\n",
    "**Problem:** Justify this equation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c29e206a-f669-4f5f-95b4-d43124bd868b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#= Compute the U matrix for N=9 and use it to get the Fourier coefficients of sin(4pi * x), \n",
    "where x is sampled on 0:1/9:8/9.\n",
    "Check that the result makes sense analytically by writing sin(4pi * x) as a \n",
    "sum of complex exponentials directly (Euler's formula). \n",
    "=#"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8020d63a-5b6f-4c21-a828-1c1e42b1e109",
   "metadata": {},
   "source": [
    "The process of changing from the Dirac basis to the Fourier basis is called the \"Discrete Fourier Transform\" or DFT.  As a point of nomenclature, \"DFT\" usually refers to the coefficient formula \n",
    "$$\n",
    "c_n = \\sum_j f_j \\frac{1}{\\sqrt{N}} e^{-2\\pi i n j/N},\n",
    "$$\n",
    "but it could just as well refer to the matrix $U$.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c841b78-5b04-4017-8a72-ba117f9175ea",
   "metadata": {},
   "source": [
    "<h2>Fast Fourier Transform (FFT)</h2>\n",
    "\n",
    "Computing the DFT via overlap integrals or via the matrix $U$ can be slow for large sampling grids.  It turns out there is a much faster way to compute it, called the Fast Fourier Transform (FFT).  We won't get into how it works (look it up!), but we will see that it does work. \n",
    "\n",
    "An important note about FFT implementations is that there are several conventions for how they are normalized.  The standard Julia FFT implementation differs from the convention we used above by a factor of $1/\\sqrt{N}$.  That is, `U * fs = fft(fs)/sqrt(N)`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d34742d6-b586-4140-b5c0-4e48683843cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "cs2 = fft(fs)/sqrt(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a38248c-1e0a-4da6-af69-48933374fa1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's see how fast the FFT is compared to the matrix multiplication DFT method. \n",
    "using BenchmarkTools\n",
    "N = 1000\n",
    "xs = 0:1/N:1-1/N\n",
    "fs = sin.(4pi * xs)\n",
    "U = [exp(-2π * im * n * m / N)/sqrt(N) for n=0:N-1,m=0:N-1];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b23e40f3-e35a-4b97-96b0-8fff53fd2e24",
   "metadata": {},
   "outputs": [],
   "source": [
    "@benchmark U*fs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "835331c1-34b7-43b8-9e86-dbf2b103e62a",
   "metadata": {},
   "outputs": [],
   "source": [
    "@benchmark fft(fs) ./ sqrt(N)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76812e6b-863d-4727-9ef6-7a31d930d371",
   "metadata": {},
   "source": [
    "<h2>Appications of FFT.</h2>\n",
    "\n",
    "<h3>(De)Convolution.</h3>\n",
    "\n",
    "A common application of FFTs is fast (de)convolution. Suppose we have an image $I_{jk}$, $j,k=0,1,\\dots,N-1$, that we want to blur.  A way to accomplish this is to replace each pixel by the average of the neighboring pixels, out to a distance $n$ in each direction (wrapping around at the boundary).  This is called convolution, and it can be written mathematically as \n",
    "$$\n",
    "(I*K)_{jk} = \\sum_{pq} I_{j-p,k-q} K_{pq},\n",
    "$$\n",
    "where the *kernel* $K_{pq} = \\frac{1}{(2n+1)^2}$ if $j,k$ lie in a box of side length $2n+1$ centered at $(0,0)$ and wrapping around boundaries.  In the sum $I_{j-p,k-q}$ is also extended periodically when the indices go outside of the range $[0,N-1]$.  \n",
    "\n",
    "According to the *convolution theorem*, the convolution $I*K$ can also be computed using FFTs: \n",
    "$$\n",
    "(I*K)_{jk} = \\text{ifft}(\\text{fft}(I) \\text{ .* } \\text{fft}(K))_{jk},\n",
    "$$\n",
    "where $\\text{ifft}$ is the inverse FFT, and the product is taken pointwise.  Note that an FFT of a 2D grid is just an FFT along rows followed by another FFT along columns."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1c57a90-3fcb-4fbd-bc19-ac1c5a0bd51e",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 100\n",
    "dx = 1/N\n",
    "grid = (0:dx:1-dx,0:dx:1-dx)\n",
    "I = [(0.1 <= sqrt((x-1/2)^2 + (y-1/2)^2) <= 0.2 ? 1.0 : 0.0) for x in grid[1], y in grid[2]]\n",
    "heatmap(I)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8620a329-8533-4664-9c6b-a651bbe5e8e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 15\n",
    "K = [(j <= n || j >= N-n) && (k <= n || k >= N-n) ? 1.0 : 0.0 for j=0:N-1,k=0:N-1] / (2n+1)^2\n",
    "heatmap(K)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6bba8c3-951d-40b0-8210-ff0781cf8041",
   "metadata": {},
   "outputs": [],
   "source": [
    "zblurredImg = real(ifft(fft(I) .* fft(K)))\n",
    "heatmap(blurredImg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35525fdb-0c3c-44ef-af4b-302efee7787d",
   "metadata": {},
   "source": [
    "Do you think information has been lost in this process?  Perhaps surprisingly, it has not.  You can get the original image back by *doconvolution*, which amounts to simply dividing by the FFT of the kernal, rather than multiplying. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a73137c3-9048-4100-ae53-f425c1282e2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "recoveredImg = # Your code here.\n",
    "heatmap(recoveredImg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b8ed7eb-09a2-434f-b460-530d0129e28a",
   "metadata": {},
   "source": [
    "<h3>FFT as approximation to the Fourier transform.</h3>\n",
    "\n",
    "The (continuous) Fourier transform is a sort of limit of Fourier series over a very large interval, rather than just $[0,1]$.  Given a function $f:\\mathbb{R}\\rightarrow \\mathbb{R}$, provided it is sufficiently integrable, it's Fourier transform is defined as \n",
    "$$\n",
    "\\tilde f(y) = \\mathcal{F}\\left[f\\right](y) := \\int_{-\\infty}^{\\infty} f(x) e^{-2\\pi i x y} dx.\n",
    "$$\n",
    "The function $f$ can be recovered from its Fourier transform $\\tilde f$ via \n",
    "$$\n",
    "f(x) = \\int_{-\\infty}^{\\infty} \\tilde f(y) e^{2\\pi i x y} dy.\n",
    "$$\n",
    "Fourier transforms are very useful for e.g. solving differential equations, because of properties like \n",
    "$$\n",
    "\\mathcal{F}\\left[\\frac{df}{dx}\\right](y) = \\int_{-\\infty}^{\\infty} \\frac{df}{dx} e^{-2\\pi i xy} dx = (-2\\pi i y) \\int_{-\\infty}^{\\infty} f(x) e^{2\\pi i x y} dx,\n",
    "$$\n",
    "which comes from integrating by parts in the middle expression.  This says that Fourier transformation converts derivatives into multiplication.  This reduces differential equations to algebraic equations.  We won't dwell on this point, but if it's unfamiliar you should spend some time playing with it. \n",
    "\n",
    "For us the big question is: Does the FFT approximate the Fourier transform?  The answer is yes, but with caveats.  The FFT of $f(x_0), f(x_0 + dx), \\dots, f(x_0 + (N-1) dx)$ approximates the Fourier transform of $f(x)$ in the interval $\\left[\\frac{-1}{2\\; dx}, \\frac{1}{2\\; dx}\\right]$.  Outside of this range, the FFT is periodic, while the continuous Fourier transform is not.  \n",
    "\n",
    "If you want to use FFTs to approximate Fourier transforms, you need to be careful with units.  My recommendation is this: For given sequence length $N$, use the grid `(-floor(N / 2):floor((N - 1) / 2)) ./ sqrt(N)` for sampling your functions.  The nice thing about this grid is that the FFT points which match the Fourier transform live on the same grid.  So by using this grid, you don't need to keep track of two separate grids.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e23b77d1-594e-4af2-8fd0-13c3e8dca53c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function whose Fourier transform you can compute analytically and compare it to the FFT output. \n",
    "N = 128 \n",
    "grid = (-floor(N / 2):floor((N - 1) / 2)) ./ sqrt(N)\n",
    "fs = # Your code here.\n",
    "Fs = real.(ifftshift(fft(fftshift(fs))) ./ sqrt(N))   # The shifts are to map negative frequencies to equivalent positive frequencies and back. \n",
    "plot(#= Your code here =#)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd09c6b4-e69b-4de6-9c28-fa2215a497d7",
   "metadata": {},
   "source": [
    "<h3>Derivatives and integrals via FFT.</h3>\n",
    "\n",
    "Fourier methods make differentiation and integration very easy.  A derivative is just multiplication by $(-2\\pi i y)$ in the Fourier domain.  Integration is a bit trickier because the result will always be periodic; so you need to apply some linear shifts to make everything work out.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9232b8e6-f709-49d4-ac56-c8fa4f0212d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "sft(x) = ifftshift(fft(fftshift(x)))\n",
    "isft(x) = ifftshift(ifft(fftshift(x)))\n",
    "dfs = real.(isft( #= Your code here =# .* sft(fs)))\n",
    "plot(grid,[fs,dfs])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6a0b2c7-5796-4b57-bfb8-e50c8a1a4cf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "ifs = real.(isft( 1 ./(2pi*im * grid .+ im*eps()) .* sft(fs .- sum(fs)))) .+ (0:1/N:1-1/N) .* sum(fs) .* step(grid)\n",
    "plot(grid,[fs,ifs])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40e0ee9f-2fe7-4fd0-8eba-ba37c05b222c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.10.3",
   "language": "julia",
   "name": "julia-1.10"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}