{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fa4f0e4d",
   "metadata": {},
   "source": [
    "## Fourier transform\n",
    "\n",
    "Numpy documentation at https://numpy.org/doc/stable/reference/routines.fft.html . Scipy documentation and examples at https://docs.scipy.org/doc/scipy/tutorial/fft.html\n",
    "\n",
    "The data files used in this notebook can be downlodaded from the \"Classrom exercises\" section on moodle"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b51c626",
   "metadata": {},
   "source": [
    "### Example: signal + noise\n",
    "\n",
    "In this example we download from the file pitch.txt a signal made of a basic wave, in which we recognize sone high freqency noise component and we perform basic denoising by Fourier transforming, removing the high-k coefficients and transforming back."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a935b8a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy as sp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9be08a29",
   "metadata": {},
   "source": [
    "##### 1) Downloading and plotting signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdccbddf",
   "metadata": {},
   "outputs": [],
   "source": [
    "pitch = np.genfromtxt('pitch.txt')\n",
    "\n",
    "plt.plot(pitch)\n",
    "plt.xlabel('$t/\\Delta t$', size = 16)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd76a9c4",
   "metadata": {},
   "source": [
    "##### 2) FT and power spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75e1af1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ck = np.fft.rfft(pitch)\n",
    "power = (abs(ck))**2\n",
    "\n",
    "plt.plot(power)\n",
    "plt.xlabel('k', size = 14)\n",
    "plt.ylabel('$P(k) = |c(k)|^2$', size = 14)\n",
    "plt.yscale('log')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a502bc6",
   "metadata": {},
   "source": [
    "##### 3) Denoising"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "797519c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "kmax = np.argmax(power)\n",
    "\n",
    "# Removing high frequency coefficients\n",
    "Delta = 100\n",
    "dck = ck.copy()\n",
    "dck[kmax + Delta:] = 0.0\n",
    "denoised_pitch = np.fft.irfft(dck)\n",
    "plt.xlabel('$t/ \\Delta t$', size = 14)\n",
    "plt.title('$\\Delta = $' + str(Delta), size = 16)\n",
    "plt.plot(denoised_pitch)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ae633a1",
   "metadata": {},
   "source": [
    "### Example: unblurring\n",
    "\n",
    "In this example we download a blurred image from the file blur.txt. Knowing that the point spread function is a Gaussian with $\\sigma$ = 25, we unblur it via PSF deconvolution, which is just a division by the PSF in Fourier space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eb1503d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy as sp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "777ee9a9",
   "metadata": {},
   "source": [
    "#### Reading file and displaying image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cae1ebf",
   "metadata": {},
   "outputs": [],
   "source": [
    "image = np.genfromtxt('blur.txt')\n",
    "plt.imshow(image)\n",
    "plt.show()\n",
    "K, L = image.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d072cd2",
   "metadata": {},
   "source": [
    "#### Generating and plotting the PSF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d70a34a",
   "metadata": {},
   "outputs": [],
   "source": [
    "size = len(image)\n",
    "sigma = 23\n",
    "\n",
    "def PSF(sigma,size):\n",
    "    # Point spread function sampled in image grid\n",
    "    def gauss(x,y,sigma):\n",
    "        norm = 1./(2.*sigma**2)\n",
    "        f = norm*np.exp(-(x**2 + y**2)/(2.*sigma**2))\n",
    "        return f\n",
    "    gauss = np.vectorize(gauss)\n",
    "    x = np.linspace(0,size,size)\n",
    "    y = np.linspace(0,size,size)\n",
    "    X,Y = np.meshgrid(x,y)\n",
    "    W = gauss(X,Y,sigma) + gauss(X-size,Y,sigma) + gauss(X,Y-size,sigma) + gauss(X-size,Y-size,sigma)\n",
    "    return W\n",
    "\n",
    "W = PSF(sigma,size)\n",
    "\n",
    "plt.imshow(W,cmap='Greys')\n",
    "plt.title('PSF (with periodic boundaries)')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "#######################################################\n",
    "# Calculating PSF with loops, instead of vectorizing\n",
    "\n",
    "#def PSF(x,y,sigma):\n",
    "#    norm = 1./(2.*sigma**2)\n",
    "#    f = norm*np.exp(-(x**2 + y**2)/(2.*sigma**2))\n",
    "#    return f\n",
    "\n",
    "#W = np.empty_like(image)\n",
    "#size = len(W)\n",
    "#sigma = 23\n",
    "\n",
    "# Calculating PSF with loops\n",
    "#for x in range (len(W)):\n",
    "#    for y in range(len(W)):\n",
    "#        W[x,y] = PSF(x,y,sigma) + PSF(x-size,y,sigma) + PSF(x,y-size,sigma) + PSF(x-size,y-size,sigma)\n",
    "\n",
    "#plt.imshow(W,cmap='Greys')\n",
    "#plt.title('PSF (with periodic boundaries)')\n",
    "#plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33a1fd4c",
   "metadata": {},
   "source": [
    "#### Deconvolving the PSF and comparing images before and after unblurring"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e5e596f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ck_img = np.fft.rfft2(image)\n",
    "ck_W = np.fft.rfft2(W)\n",
    "\n",
    "K, L = ck_W.shape\n",
    "\n",
    "# Deconvolving PSF by division in Fourier space \n",
    "ck_unblur = np.copy(ck_img)\n",
    "for i in range(K):\n",
    "    for j in range(L):        \n",
    "        if(abs(ck_W[i,j])>1e-3):\n",
    "            ck_unblur[i,j] = ck_img[i,j]/(ck_W[i,j])\n",
    "\n",
    "# Fourier transforming back to get the unblurred image\n",
    "img_unblur = np.fft.irfft2(ck_unblur)          \n",
    "            \n",
    "# Plotting    \n",
    "compare = display(image,img_unblur)\n",
    "\n",
    "def display(original,filtered):\n",
    "    fig, axs = plt.subplots(1,2,figsize=(10,6))\n",
    "    axs[0].set_title('Before')\n",
    "    axs[0].imshow(original)\n",
    "    axs[1].imshow(filtered)\n",
    "    axs[1].set_title('After')\n",
    "    # remove the x and y ticks\n",
    "    for ax in axs:\n",
    "        ax.set_xticks([])\n",
    "        ax.set_yticks([])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61c901f3",
   "metadata": {},
   "source": [
    "### Example. Adding noise\n",
    "\n",
    "We now add noise to the image, besides the blurring. If we deconvolve the PSF, in this case we incorretly deconvolve also the noise. This can lead to catastrophic effects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0c421d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy as sp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ca73662",
   "metadata": {},
   "outputs": [],
   "source": [
    "def display(original,filtered):\n",
    "    fig, axs = plt.subplots(1,2,figsize=(10,6))\n",
    "    axs[0].set_title('Before')\n",
    "    axs[0].imshow(original)\n",
    "    axs[1].imshow(filtered)\n",
    "    axs[1].set_title('After')\n",
    "    # remove the x and y ticks\n",
    "    for ax in axs:\n",
    "        ax.set_xticks([])\n",
    "        ax.set_yticks([])\n",
    "    plt.show()\n",
    "\n",
    "image = np.genfromtxt('blur.txt')\n",
    "K, L = image.shape\n",
    "\n",
    "print(K,L)\n",
    "\n",
    "# Adding noise\n",
    "sigma = 100\n",
    "frac = 0.\n",
    "def add_noise(image,std,highpass):\n",
    "    # Adding noise to the image\n",
    "    # Generating Gaussian noise in pixel space\n",
    "    K, L = image.shape\n",
    "    mean = 0\n",
    "    seed = 2938645\n",
    "    np.random.seed(seed)\n",
    "    noise = np.random.normal(mean,std,size=(K,L))\n",
    "    # Removing low frequency noise components in Fourier space\n",
    "    nk = np.fft.rfft2(noise)\n",
    "    keep_fraction = highpass\n",
    "    nk[0:int(K*highpass)] = 0\n",
    "    nk[int(K*(1-highpass)):] = 0\n",
    "    nk[:,0:int(K*highpass)] = 0\n",
    "    # Fourier transforming back to get filtered noise field\n",
    "    # in pixel space\n",
    "    noise = np.fft.irfft2(nk)\n",
    "    plt.imshow(abs(nk))\n",
    "    return image + noise\n",
    "    \n",
    "image2 = add_noise(image,sigma,frac)\n",
    "\n",
    "compare = display(image,image2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "035f95f0",
   "metadata": {},
   "source": [
    "##### Point spread function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ee69f71",
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma = 23\n",
    "size=1024\n",
    "\n",
    "def PSF(sigma,size):\n",
    "    # Point spread function sampled in image grid\n",
    "    def gauss(x,y,sigma):\n",
    "        norm = 1./(2.*sigma**2)\n",
    "        f = norm*np.exp(-(x**2 + y**2)/(2.*sigma**2))\n",
    "        return f\n",
    "    gauss = np.vectorize(gauss)\n",
    "    x = np.linspace(0,size,size)\n",
    "    y = np.linspace(0,size,size)\n",
    "    X,Y = np.meshgrid(x,y)\n",
    "    W = gauss(X,Y,sigma) + gauss(X-size,Y,sigma) + gauss(X,Y-size,sigma) + gauss(X-size,Y-size,sigma)\n",
    "    return W\n",
    "\n",
    "W = PSF(sigma,size)\n",
    "\n",
    "#plt.imshow(W,cmap='Greys')\n",
    "#plt.title('PSF (with periodic boundaries)')\n",
    "#plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08095ba4",
   "metadata": {},
   "source": [
    "#### Deconvolution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0930fb54",
   "metadata": {},
   "outputs": [],
   "source": [
    "ck_W = np.fft.rfft2(W)\n",
    "ck_img = np.fft.rfft2(image2)\n",
    "\n",
    "K, L = ck_W.shape\n",
    "# Deconvolving PSF by division in Fourier space \n",
    "ck_unblur = np.copy(ck_img)\n",
    "for i in range(K):\n",
    "    for j in range(L):        \n",
    "        if(abs(ck_W[i,j])>1e-3):\n",
    "            ck_unblur[i,j] = ck_img[i,j]/(ck_W[i,j])\n",
    "\n",
    "# Fourier transforming back to get the unblurred image\n",
    "img_unblur = np.fft.irfft2(ck_unblur).real          \n",
    "            \n",
    "# Plotting    \n",
    "compare = display(image2,img_unblur) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5bf3cdc",
   "metadata": {},
   "source": [
    "### Exercise: musical instruments \n",
    "\n",
    "Files piano.txt and trumpet.txt contain a sampling at 44100 Hz of a single note, played by a piano and a trumpet.\n",
    "\n",
    "1) Load the waveforms and plot them\n",
    "\n",
    "2) Apply fft and plot the first 10000 coefficients of the power spectrum, for both instruments\n",
    "\n",
    "3) Which note where the two instruments playing? [hint: the note middle C ('do') has a frequency of 216 Hz]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5477ec99",
   "metadata": {},
   "source": [
    "###### 1) Loading and plotting waveforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b8f8552",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a60fe7cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "trumpet = np.genfromtxt('trumpet.txt')\n",
    "plt.plot(trumpet)\n",
    "plt.title('Trumpet waveform', size = 14)\n",
    "plt.xlabel('Samples', size = 12)\n",
    "plt.show()\n",
    "\n",
    "piano = np.genfromtxt('piano.txt')\n",
    "plt.plot(piano)\n",
    "plt.title('Piano waveform', size = 14)\n",
    "plt.xlabel('Samples', size = 12)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf7d3726",
   "metadata": {},
   "source": [
    "##### 2) Doing FFT and plotting spectra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea45118f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Trumpet\n",
    "trumpetk = np.fft.rfft(trumpet)\n",
    "plt.xlim(0,10000)\n",
    "plt.plot(abs(trumpetk)**2)\n",
    "plt.yscale('log')\n",
    "plt.title('Trumpet', size = 14)\n",
    "plt.xlabel('k', size = 14)\n",
    "plt.show()\n",
    "\n",
    "# Piano\n",
    "pianok = np.fft.rfft(piano)\n",
    "plt.xlim(0,10000)\n",
    "plt.yscale('log')\n",
    "plt.title('Piano', size = 14)\n",
    "plt.xlabel('k', size = 14)\n",
    "plt.plot(abs(pianok)**2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78efda26",
   "metadata": {},
   "source": [
    "##### 3) Finding which note the instruments where playing\n",
    "(to do that: find the maximum of the power spectrum and convert to frequency, knowing that the sampling frequency was 44100 Hz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c13de2c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Trumpet\n",
    "sampling_freq = 44100\n",
    "kmax = np.argmax(abs(trumpetk)**2)\n",
    "f = (kmax/(len(trumpetk)))*44100\n",
    "print('Trumpet largest frequency component (in units of 261 Hz), f = ', f/261)\n",
    "\n",
    "# Piano\n",
    "kmax = np.argmax(abs(pianok)**2)\n",
    "f = (kmax/(len(pianok)))*44100\n",
    "print('Piano largest frequency component (in units of 261 Hz), f = ', f/261)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b35310f",
   "metadata": {},
   "source": [
    "### Exercise: sunsposts\n",
    "\n",
    "Download the sunspot data from the file sunspots.txt. The file contains two columns. The first column shows time, measured in months, starting from January 1749; the second column contains the observed number of sunspots in each month.\n",
    "\n",
    "1) Plot the number of sunspots as a function of time. Get a by eye estimate of the length of the cycle.\n",
    "\n",
    "2) Fourier transform the data and plot the power spectrum ($|c_k|^2$ as a function of k)\n",
    "\n",
    "3) Find the power spectrum peak for k>0. What does the $|c(0)|^2$ value represent?\n",
    "\n",
    "4) Find the period of the wave-term corresponding to the peak measured in 3) and compare to your estimate in 1)\n",
    "\n",
    "(if you want, you can solve it using your own DFT script, and compare with FFT numpy results, verifying also the difference in performance betwen the two algorithms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eba81bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc324811",
   "metadata": {},
   "source": [
    "#### 1) Plotting data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54fbc7e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "spots = np.genfromtxt('sunspots.txt')\n",
    "\n",
    "# Numbe of data\n",
    "N = len(spots[:,0])\n",
    "print(N)\n",
    "\n",
    "plt.xlabel('Month', size = 14)\n",
    "plt.ylabel('Number of sunspots',size = 14)\n",
    "#plt.xlim(2000,2500)\n",
    "plt.plot(spots[:,0],spots[:,1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "785081a4",
   "metadata": {},
   "source": [
    "#### Zooming to get an estimate of the period\n",
    "\n",
    "We get an estimate of about 100-140 months "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2925805",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel('Month', size = 14)\n",
    "plt.ylabel('Number of sunspots',size = 14)\n",
    "plt.xlim(2000,2500)\n",
    "plt.plot(spots[:,0],spots[:,1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd2b832c",
   "metadata": {},
   "source": [
    "#### 2) Fourier transform and power spectrum\n",
    "\n",
    "By plotting |c(k)|^2 we would see a large peak for k=0. This is just a constant offset (monopole), giving the average value \n",
    "of sunspots over the entire period of observation. Since we are interested in cycles, we remove it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8eb3a47",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fourier transform\n",
    "ck = np.fft.rfft(spots[:,1])\n",
    "# Removing monopole\n",
    "ck[0] = 0\n",
    "\n",
    "\n",
    "plt.plot(abs(ck)**2)\n",
    "#plt.yscale('log')\n",
    "plt.xlim(-1,200)\n",
    "plt.xlabel('k',size = 14)\n",
    "plt.ylabel('$|c(k)|^2$',size=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb4c729d",
   "metadata": {},
   "source": [
    "#### 3) Finding maximum of power spectrum and plotting corresponding mode, in comparison with the observed data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b4d0412",
   "metadata": {},
   "outputs": [],
   "source": [
    "kmax = np.argmax(abs(ck))\n",
    "\n",
    "filter = np.zeros_like(ck)\n",
    "filter[kmax] = ck[kmax]\n",
    "trend = np.fft.irfft(filter)\n",
    "\n",
    "plt.plot(trend)\n",
    "plt.plot(spots[:,1])\n",
    "plt.xlabel('Months')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22609cb7",
   "metadata": {},
   "source": [
    "#### 4) Getting period corresponding to frequency of maximum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f9acaa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "freq = kmax/N \n",
    "period = 1./freq\n",
    "print('Period (yrs) = ', period/12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bac5a61a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}