Fourier Analysis II (Problem Version).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e36feea8",
   "metadata": {},
   "source": [
    "## Problem Set for Lecture II: Filtering and Aliasing (Bryan's Version)\n",
    "\n",
    "Bryan Scott, CIERA/Northwestern\n",
    "\n",
    "Presented at DSFP Session 20: Time Series Methods in Astronomy\n",
    "\n",
    "Exercises adapted and expanded from Statistics, Data Mining, and Machine Learning in Astronomy, by Ivezić, Connolly, VanderPlas, and Gray\n",
    "\n",
    "version 0.1, February 16, 2024\n",
    "\n",
    "The goal of this notebook is to walk you through filtering and aliasing so that you can learn how to perform optimal filtering of an input signal and gain intuition for the impact that filtering has on a signal. The final part will help you build intuition for why aliasing occurs and what impact it has on your ability to characterize signals of interest. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c94158aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "import matplotlib.pyplot as plt \n",
    "import scipy.fft as fft\n",
    "import scipy.stats as stats"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9df06b1",
   "metadata": {},
   "source": [
    "### Problem 1: Quick Simulated Data and Low Pass Filtering\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "736ede41",
   "metadata": {},
   "source": [
    "### Part A) Simulated data and familiarizing yourself with the FFT\n",
    "\n",
    "Let's begin by generating a mock time series consisting of a Gaussian signal that we'd like to measure and a large noise term. After you plot the signal, plot the FFT of the signal using the np.fft.fft function. When you plot, make sure you understand the structure and symmetries of the object FFT returns. We will discuss why this object behaves this way in Session 21 - it is related to deep facts about the Fourier transform and its discretization for running on digital processors. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a76dd71",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the noisy data\n",
    "\n",
    "np.random.seed(5)\n",
    "N = 2000\n",
    "dt = 0.05\n",
    "noise = 0.01\n",
    "\n",
    "t =  #N evenly spaced samples dt apart \n",
    "\n",
    "sr =  # samples per interval \n",
    "\n",
    "freq =  # N * 1/dt, by hand\n",
    "\n",
    "freq = # using convenience function in np.fft \n",
    "\n",
    "h1 = np.exp(-0.5 * ((t - 20.) / 1.0) ** 2) #gaussian signal\n",
    "h2 = 0.3*np.sin(2*np.pi*0.5*t) # periodic signal \n",
    "hN1 = h1 +  # add noise to the first signal\n",
    "hN2 = h2 +  #add noise to the second signal \n",
    "\n",
    "\n",
    "fig = plt.figure(figsize=(10, 6.75))\n",
    "\n",
    "\n",
    "# First plot: noisy signal\n",
    "ax = fig.add_subplot(211)\n",
    "ax.plot(t, hN1, '-', c='gray', label = 'signal + noise')\n",
    "ax.plot(t, h1, c = 'dodgerblue', label = 'signal')\n",
    "\n",
    "ax.set_xlabel(r'$\\lambda$')\n",
    "ax.set_ylabel('flux')\n",
    "\n",
    "ax = fig.add_subplot(212)\n",
    "ax.plot(t, hN2, '-', c='gray', label = 'signal + noise')\n",
    "ax.plot(t, h2, c = 'dodgerblue', label = 'signal')\n",
    "\n",
    "# plot the FFTs\n",
    "\n",
    "#plotting code goes here \n",
    "\n",
    "plt.xlabel('Freq (Hz)')\n",
    "plt.ylabel('FFT Amplitude |X(freq)|')\n",
    "plt.xlim(0, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "929d7ab9",
   "metadata": {},
   "source": [
    "## Part b) Low Pass Filtering \n",
    "\n",
    "Write a function that returns 1 for frquencies below some critical value and 0 for values above that critical value. The numpy piecewise and lambda functions may be helpful, as before, an example of the syntax is in the documentation for numpy.piecewise. After filtering your signal from part a plot the original signal and the inverse transform of the filtered signal. Repeat this exercise for the different functions you plotted in part a. How is the recovered signal different? Repeat again for varying noise amplitudes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a3b81ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def low_pass(nu, nu_c):\n",
    "    return \n",
    "\n",
    "freq = \n",
    "\n",
    "fig = plt.figure(figsize=(10, 6.75))\n",
    "\n",
    "ax = fig.add_subplot(211)\n",
    "\n",
    "ax.set_xlabel(r'$\\lambda$')\n",
    "ax.set_ylabel('flux')\n",
    "\n",
    "# ax.plot(t, hN2, '-', c='gray', label = 'signal + noise')\n",
    "\n",
    "\n",
    "\n",
    "ax.plot()\n",
    "ax.plot(t, h2, label = 'signal', c = 'orange', ls = '-.')\n",
    "\n",
    "ax = fig.add_subplot(212)\n",
    "\n",
    "ax.plot()\n",
    "ax.plot(t, h1, label = 'signal', c = 'orange', ls = '-.')\n",
    "\n",
    "ax.set_xlim(0, 90)\n",
    "\n",
    "ax.set_xlabel(r'$\\lambda$')\n",
    "ax.set_ylabel('flux')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4e00f14",
   "metadata": {},
   "source": [
    "### Problem 2: Optimal Filtering \n",
    "\n",
    "In the lecture, we discussed high pass and low pass filters. Another common filter you may encounter is the \"Wiener Filter\" which is an 'optimal filter' in the following sense. \n",
    "\n",
    "Suppose you have a signal you're interested in which we can call u(t). Rather than being a pure measurement of some underlying process, suppose u(t) is the product of a measurement made by some imperfect instrument with a response function u(t). The total observed signal is then the convolution of the signal we're interested in and the instrument response,\n",
    "\n",
    "$$ \n",
    "s(t) = \\int r(\\tau) u(t-\\tau) d\\tau\n",
    "$$\n",
    "\n",
    "Let us also suppose, for complexity's sake, that the signal s(t) is not measured independent of noise contributions.That is, we want to infer u(t) from c(t), where \n",
    "\n",
    "$$\n",
    "c(t) = s(t) + n(t) \n",
    "$$\n",
    "\n",
    "where n(t) is some noise signal. How might we proceed? We want to find a filter that retains as much information from the original signal as possible while discarding information from the noise and deconvolving with respect to the instrument response. We can quantify our error in the least squares sense via\n",
    "\n",
    "$$\n",
    "\\int | \\tilde{u}(t) - u(t) |^2 dt = \\int |\\tilde{U}(f) - U(f)|^2 df \n",
    "$$\n",
    "\n",
    "where the reconstructed signal is given by \n",
    "\n",
    "$$\n",
    "\\tilde{U}(f) = \\frac{C(f) \\Phi(f)}{R(f)}\n",
    "$$\n",
    "\n",
    "Minimizing the least squares error and plugging in all of the above expressions yields an 'optimal' filter:\n",
    "\n",
    "$$\n",
    "\\Phi(f) = \\frac{ |S(f)|^2 }{|S(f)|^2 + |N(f)|^2} \n",
    "$$\n",
    "\n",
    "#### Part A) Why is this an optimal filter? Hint: In what limits does each term in the filter 'matter more'? "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c74295cc",
   "metadata": {},
   "source": [
    "#### Part B) Estimates of the Noise Power"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a106708",
   "metadata": {},
   "source": [
    "Using the Wiener filter in practice requires you to have an estimate of the signal and noise spectra (and their sum) in advance. While this is theoretically problematic, in practice it's rarely an issue. We can proceed in two ways. A more sophisticated method is to define a signal and noise function with a functional form close to those expected in your simulated signal from part a. Then compute the filter and filter the total signal to recover the component of interest. Alternatively, it is often easier to just approximate the signal PSD as the total PSD of the observed signal and assume the noise is flat in frequency space. Try both. Some code for achieving this is outlined below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c86ef975",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute the PSD of the input\n",
    "N = len(t)\n",
    "Df =  \n",
    "f = fftpack.ifftshift(Df * (np.arange(N) - N / 2)) # shifts the central frequency to zero \n",
    "\n",
    "H = np.fft.fft()\n",
    "PSD = \n",
    "\n",
    "def signal():  \n",
    "    return \n",
    "\n",
    "def noise():\n",
    "    return\n",
    "\n",
    "# make some intelligent choices for 'guess' for each of the above parameters of your model functions:\n",
    "# hint, what do all of the terms in the PSD tell you about the signal?\n",
    "\n",
    "amp_guess = \n",
    "width_guess = \n",
    "signal_params = (amp_guess, width_guess)\n",
    "\n",
    "noise_params = \n",
    "\n",
    "# now fit for the signal - noise:\n",
    "\n",
    "def min_func(v): \n",
    "    # what should the minimum function return? What are you trying to minimize when you fit?\n",
    "    return \n",
    "\n",
    "v0 = tuple(signal_params) + tuple(noise_params)\n",
    "v = optimize.minimize(min_func, v0, method='Nelder-Mead')['x']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ebf78d7",
   "metadata": {},
   "source": [
    "#### Part C) Implementing a Wiener Filter on a Noisy Signal "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7581707e",
   "metadata": {},
   "source": [
    "Now implement both versions of the Wiener filter - based on the fit or directly from the PSD of the total observed signal and compare the results. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e8db856",
   "metadata": {},
   "outputs": [],
   "source": [
    "# P_S = signal(f, v[0], v[1])\n",
    "\n",
    "P_S = \n",
    "\n",
    "P_N = \n",
    "Phi = \n",
    "Phi[0] = 1  # correct for DC offset\n",
    "\n",
    "# Use Phi to filter and smooth the values\n",
    "h_smooth = fftpack.ifft(Phi * H)\n",
    "\n",
    "plt.plot(f, P_S)\n",
    "plt.show()\n",
    "plt.plot(f, P_N)\n",
    "plt.show()\n",
    "\n",
    "plt.plot(t, h_smooth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56f907e2",
   "metadata": {},
   "source": [
    "### Problem 3: Sampling and Aliasing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71df3d07",
   "metadata": {},
   "source": [
    "We saw in the lecture that signals can be reconstructed if they are sampled at rates greater than the Nyquist frequency. What happens if we sample at rates that are 'too low'? The answer is that power from frequencies too low to be captured by a given sampling rate is misidentified as coming from frequencies that are measured. Let's consider two similar functions,\n",
    "\n",
    "$$\n",
    "f(t) = \\cos\\left(\\frac{3 \\pi}{2} t\\right) \n",
    "$$\n",
    "\n",
    "and \n",
    "\n",
    "$$\n",
    "g(t) = \\cos\\left(\\frac{\\pi}{2} t\\right)\n",
    "$$\n",
    "\n",
    "For different numbers of samples. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78ad2c77",
   "metadata": {},
   "source": [
    "### Part a) Starting in the Time Domain, plot the functions f(t) and g(t) over the interval (-4, 4) for n = 9 samples and n = 18 samples. From your plot, make sure you understand why undersampling leads to aliasing. Explain the origin of the term 'aliasing'. How did I pick this interval and this number of samples?\n",
    "\n",
    "Hint: The word alias is defined as \"an assumed or additional name that a person sometimes uses\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "204c5518",
   "metadata": {},
   "outputs": [],
   "source": [
    "t1 = \n",
    "t2 = \n",
    "\n",
    "plt.plot(t1, )\n",
    "plt.plot(t1, )\n",
    "\n",
    "plt.plot(t2, )\n",
    "plt.plot(t2, )\n",
    "\n",
    "# plt.grid()\n",
    "plt.xlabel('t')\n",
    "plt.ylabel('f(t), g(t)')\n",
    "plt.legend(loc = 'upper right')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a570723b",
   "metadata": {},
   "source": [
    "### Part B) Let's now plot the Fourier Transforms of each signal. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9403421",
   "metadata": {},
   "outputs": [],
   "source": [
    "N1 = \n",
    "N2 = \n",
    "\n",
    "t1,dt1 = np.linspace(retstep =)\n",
    "t2,dt2 = np.linspace(retstep =)\n",
    "\n",
    "freq1 = \n",
    "freq2 = \n",
    "\n",
    "H1 = np.fft\n",
    "H2 = np.fft\n",
    "\n",
    "H3 = np.fft\n",
    "H4 = np.fft\n",
    "\n",
    "plt.stem()\n",
    "\n",
    "plt.axvline()\n",
    "plt.axvline()\n",
    "plt.xlim(-1.5, 1.5)\n",
    "plt.show()\n",
    "\n",
    "plt.axvline()\n",
    "plt.axvline()\n",
    "plt.stem()\n",
    "plt.xlim()\n",
    "\n",
    "plt.show()\n",
    "\n",
    "\n",
    "plt.stem()\n",
    "plt.axvline()\n",
    "plt.axvline()\n",
    "plt.xlim()\n",
    "\n",
    "plt.show()\n",
    "\n",
    "plt.stem()\n",
    "\n",
    "plt.axvline()\n",
    "plt.axvline()\n",
    "plt.xlim()\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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
}