DFT of Time Shifting Square Wave

SquareWave.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "otherwise-screw",
   "metadata": {},
   "source": [
    "## Shifting Square Wave\n",
    "### <a href = \"http://www.ctralie.com\">Chris Tralie</a>\n",
    "\n",
    "The purpose of this notebook is to create an animation of a time shifting square wave.  A square wave can be written as\n",
    "\n",
    "#### $y(t) = sgn(\\sin(2 \\pi f t))$\n",
    "\n",
    "where\n",
    "\n",
    "#### $ sgn(x) = \\left\\{  \\begin{array}{cc} 1 & x > 0 \\\\ 0 & x = 0 \\\\ -1 & x < 0   \\end{array} \\right\\} $\n",
    "\n",
    "\n",
    "Its fourier series is\n",
    "\n",
    "### $y(t) = \\frac{4}{\\pi} \\sum_{k = 0}^{\\infty} \\frac{\\sin(2 \\pi (2k+1)f t)}{2k+1} $\n",
    "\n",
    "If we shift the signal by an amount $s$ and evalate $y(t - s)$, then this is\n",
    "\n",
    "### $y(t-s) = \\frac{4}{\\pi} \\sum_{k = 0}^{\\infty} \\frac{\\sin(2 \\pi (2k+1)f (t-s))}{2k+1} = \\frac{4}{\\pi} \\sum_{k = 0}^{\\infty} \\frac{\\sin(2 \\pi (2k+1)f t - 2 \\pi (2k+1)f s))}{2k+1} $\n",
    "\n",
    "We notice that the phase here is \n",
    "\n",
    "$ - 2 \\pi (2k+1)f s $\n",
    "\n",
    "So for the same shift in time, the higher frequency indices $k$ get a larger phase shift.  \n",
    "\n",
    "If we shift this phase linearly and do a similar experiment with the discrete fourier transform, where we're careful to keep the frequency of the square wave an integer number of cycles over the discrete inveral, then this leads to beautiful patterns with the phase where the higher frequency phases change much more quickly than the lower frequency phases.  The code below accomplishes this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "inappropriate-operation",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "potential-fishing",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dft(x):\n",
    "    \"\"\"\n",
    "    Return the non-redundant real/imaginary components\n",
    "    of the DFT, expressed as amplitudes of sines/cosines\n",
    "    involved\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: ndarray(N)\n",
    "        A signal\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    cos_sum: ndarray(ceil(N/2)), sin_sums(ndarray(ceil(N/2)))\n",
    "        DFT cosine/sine amplitudes\n",
    "    \"\"\"\n",
    "    N = len(x)\n",
    "    t = np.linspace(0, 1, N+1)[0:N]\n",
    "    n_freqs = int(np.ceil(N/2))\n",
    "    f = np.fft.fft(x)\n",
    "    cos_sums = np.real(f)[0:n_freqs]/(N/2)\n",
    "    sin_sums = -np.imag(f)[0:n_freqs]/(N/2)\n",
    "    return cos_sums, sin_sums"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "personal-caution",
   "metadata": {},
   "outputs": [],
   "source": [
    "def time_shift_square_anim(N, max_freq_index, n_shifts, f):\n",
    "    \"\"\"\n",
    "    Create an animation of a time shifting by one period\n",
    "    of a square wave, while plotting its cosine/sine components\n",
    "    of the discrete Fourier Transform, as well as the associated \n",
    "    amplitudes/phases\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    N: int\n",
    "        Number of samples to take in the square wave\n",
    "    max_freq_index: int\n",
    "        Maximum frequency index to show in the plots\n",
    "    n_shifts: int\n",
    "        Number of time shifts to show in the animation\n",
    "    f: int\n",
    "        Frequency of the sinusoid, in units of cycles per\n",
    "        interval N\n",
    "    \"\"\"\n",
    "    t = np.linspace(0, 1, N+1)[0:N]\n",
    "    plt.figure(figsize=(16, 8))\n",
    "    for i, shift in enumerate(np.linspace(0, 1/f, n_shifts)):\n",
    "        x = np.sign(np.sin(2*np.pi*f*(t-shift)))\n",
    "        c, s = get_dft(x)\n",
    "        amps = np.sqrt(c**2+s**2)\n",
    "        phases = np.arctan2(s, c)\n",
    "        plt.clf()\n",
    "        plt.subplot2grid((2, 4), (0, 0), rowspan=1, colspan=4)\n",
    "        plt.plot(t, x)\n",
    "        plt.ylim([-1.1, 1.1])\n",
    "        plt.title(\"square$(2 \\pi {} (t - {:.3f}))$\".format(f, shift))\n",
    "        plt.subplot(245)\n",
    "        plt.stem(c[0:max_freq_index+2])\n",
    "        plt.xlabel(\"Frequency Index\")\n",
    "        plt.ylabel(\"Dot Product\")\n",
    "        plt.title(\"Cosines\")\n",
    "        plt.ylim([-4/np.pi-0.1, 4/np.pi+0.1])\n",
    "        plt.yticks([-4/np.pi, 0, 4/np.pi], [\"$-4/\\\\pi$\", \"0\", \"$4/\\\\pi$\"])\n",
    "        plt.xlim([-0.5, max_freq_index+0.5])\n",
    "        plt.subplot(246)\n",
    "        plt.stem(s[0:max_freq_index+2])\n",
    "        plt.xlabel(\"Frequency Index\")\n",
    "        plt.ylabel(\"Dot Product\")\n",
    "        plt.title(\"Sines\")\n",
    "        plt.ylim([-4/np.pi-0.1, 4/np.pi+0.1])\n",
    "        plt.yticks([-4/np.pi, 0, 4/np.pi], [\"$-4/\\\\pi$\", \"0\", \"$4/\\\\pi$\"])\n",
    "        plt.xlim([-0.5, max_freq_index+0.5])\n",
    "        plt.subplot(247)\n",
    "        plt.stem(amps[0:max_freq_index+2])\n",
    "        plt.xlabel(\"Frequency Index\")\n",
    "        plt.ylabel(\"Dot Product\")\n",
    "        plt.title(\"Amplitudes\")\n",
    "        plt.ylim([0, 4/np.pi+0.1])\n",
    "        plt.yticks([0, 4/np.pi/3, 4/np.pi/5, 4/np.pi/7, 4/np.pi], [\"0\", \"$(4/\\\\pi)/3$\", \"$(4/\\\\pi)/5$\", \"$(4/\\\\pi)/7$\", \"$4/\\\\pi$\"])\n",
    "        plt.xlim([-0.5, max_freq_index+0.5])\n",
    "        plt.grid(linestyle='--', linewidth=2, axis='y')\n",
    "        plt.subplot(248)\n",
    "        plt.stem((amps[0:max_freq_index+2]>1e-3)*phases[0:max_freq_index+2])\n",
    "        plt.xlim([-0.5, max_freq_index+0.5])\n",
    "        plt.ylim([-np.pi-0.1, np.pi+0.1])\n",
    "        plt.yticks([-np.pi, -np.pi/2, 0, np.pi/2, np.pi], [\"$-\\\\pi$\", \"$\\\\pi/2$\", \"0\", \"$\\\\pi/2$\", \"$\\\\pi$\"])\n",
    "        plt.xlabel(\"Frequency Index\")\n",
    "        plt.ylabel(\"Dot Product\")\n",
    "        plt.title(\"Phase\")\n",
    "        plt.tight_layout()\n",
    "        plt.savefig(\"{}.png\".format(i), bbox_inches='tight', facecolor='white', transparent=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "arbitrary-falls",
   "metadata": {},
   "outputs": [],
   "source": [
    "time_shift_square_anim(100000, 80, 500*6, 2)"
   ]
  }
 ],
 "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": 5
}