Pole-Zero-Interaction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interactive IIR filter design plot\n",
    "\n",
    "*Brian McFee* <brian.mcfee@nyu.edu>, 2020-12-10\n",
    "\n",
    "\n",
    "This notebook gives a basic tool for interactively designing linear IIR systems by placing poles and zeros in the complex plane.\n",
    "\n",
    "Instructions:\n",
    "\n",
    "- To place a pole, use the left mouse button.\n",
    "- To place a zero, use the right mouse button.\n",
    "- To clear the plot and start over, use the middle mouse button.\n",
    "\n",
    "Sampling is assumed to be at 44.1 KHz, and poles/zeros are constrained to come in conjugate pairs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.signal\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "import matplotlib.colors\n",
    "import matplotlib.cm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c0545ca7e794c1f98d36055d2151ecb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# A mesh grid for the complex plane\n",
    "x = np.linspace(-1.5, 1.5, num=1000)\n",
    "y = 1j * np.linspace(-1.5, 1.5, num=1000)\n",
    "z = np.add.outer(x, y).T\n",
    "\n",
    "fs = 44100\n",
    "\n",
    "# Initially start with no poles or zeros\n",
    "zeros = []\n",
    "poles = []\n",
    "\n",
    "\n",
    "# Evaluate transfer function on the grid\n",
    "def compute_H(z, zeros, poles):\n",
    "\n",
    "    H = np.ones_like(z)\n",
    "    for _z in np.asarray(zeros):\n",
    "        H *= (z - _z)\n",
    "    \n",
    "    for _p in np.asarray(poles):\n",
    "        H /= (z - _p)\n",
    "        \n",
    "    return H\n",
    "\n",
    "H = compute_H(z, zeros, poles)\n",
    "\n",
    "# Make a mosaic grid\n",
    "fig = plt.figure(constrained_layout=True, figsize=(8, 8))\n",
    "axd = fig.subplot_mosaic(\n",
    "\"\"\"\n",
    "ZT\n",
    "FF\n",
    "\"\"\"\n",
    ")\n",
    "axpz = axd['Z']\n",
    "axtf = axd['T']\n",
    "axf = axd['F']\n",
    "axpz.sharex(axtf)\n",
    "axpz.sharey(axtf)\n",
    "\n",
    "# Scaffold the pole-zero panel\n",
    "axpz.set_aspect('equal')\n",
    "axpz.grid(True, zorder=-10)\n",
    "axpz.set(xlabel='Real', ylabel='Imaginary', title='Pole-Zero plot')\n",
    "circ = patches.Ellipse((0, 0), 2, 2, edgecolor='k', linewidth=3, fill=False, zorder=1)\n",
    "\n",
    "axpz.add_patch(circ)\n",
    "\n",
    "zero_scat = axpz.scatter(np.asarray(zeros).real, np.asarray(zeros).imag, marker='o', facecolor='none', \n",
    "                         edgecolor='g', linewidth=3, zorder=2, s=50)\n",
    "pole_scat = axpz.scatter(np.asarray(poles).real, np.asarray(poles).imag, marker='x', color='r', zorder=2, s=50)\n",
    "axpz.set(xlim=[x.min(), x.max()], ylim=[x.min(), x.max()])\n",
    "\n",
    "\n",
    "# Scaffold the frequency response panel\n",
    "freq, h = scipy.signal.freqz_zpk(zeros, poles, 1, fs=fs)\n",
    "#frcurve,  = axf.plot(freq, 20 * np.log10(np.abs(h) + 1e-12), linewidth=3)\n",
    "h_db = 20 * np.log10(np.abs(h) + 1e-12)\n",
    "frcurve = axf.scatter(freq, h_db, c=h_db, cmap='twilight_shifted', vmin=-80, vmax=80, marker='o', edgecolors='face', linewidth=10, s=20)\n",
    "axf.set(title='Frequency response', xlabel='Frequency [Hz]', ylabel='$|H(z)|$ [dB]')\n",
    "axf.axis('tight')\n",
    "axf.set(ylim=[-80, 80], xlim=[0, fs/2])\n",
    "axf.grid(True, zorder=-10)\n",
    "\n",
    "# And the transfer function  plot\n",
    "tfmesh = axtf.pcolormesh(x, x, 20 * np.log10(np.abs(H) + 1e-12), shading='nearest', cmap='twilight_shifted')\n",
    "tfmesh.set_clim([-80, 80])\n",
    "axtf.set(xlabel='Real', title='Transfer function magnitude')\n",
    "circ = patches.Ellipse((0, 0), 2, 2, edgecolor='k', linewidth=2, fill=False, zorder=1, alpha=0.5)\n",
    "\n",
    "axtf.add_patch(circ)\n",
    "axtf.set_aspect('equal')\n",
    "fig.colorbar(tfmesh, label='$|H(z)|$ [dB]', ax=axtf)\n",
    "\n",
    "# Set up the interaction\n",
    "def add_or_remove_point(event):\n",
    "    global zeros\n",
    "    global poles\n",
    "    \n",
    "    \n",
    "    \n",
    "    # Get the coordinates of the new point   \n",
    "    new_point = event.xdata + 1j * event.ydata\n",
    "    \n",
    "    # Button 1 = add a zero\n",
    "    if event.inaxes in (axtf, axpz):\n",
    "        \n",
    "        if event.button == 1:\n",
    "            zeros.append(new_point)\n",
    "            if event.ydata != 0:\n",
    "                zeros.append(np.conj(new_point))\n",
    "\n",
    "        # Button 3 = add a pole\n",
    "        elif event.button == 3:\n",
    "            poles.append(new_point)\n",
    "            if event.ydata != 0:\n",
    "                poles.append(np.conj(new_point))\n",
    "        \n",
    "    # Button 2 is reset\n",
    "    if event.button == 2:\n",
    "        zeros = []\n",
    "        poles = []\n",
    "        \n",
    "    Z = np.asarray(zeros)\n",
    "    zero_scat.set_offsets(np.vstack([Z.real, Z.imag]).T)\n",
    "    \n",
    "    P = np.asarray(poles)\n",
    "    pole_scat.set_offsets(np.vstack([P.real, P.imag]).T)\n",
    "        \n",
    "    # Update the transfer function\n",
    "    H = compute_H(z, zeros, poles)\n",
    "    H_db = 20 * np.log10(np.abs(H) + 1e-12).ravel()\n",
    "    tfmesh.set_array(H_db)\n",
    "    vmax = max(40, max(np.abs(H_db)))\n",
    "    tfmesh.set_clim([-vmax, vmax])\n",
    "    \n",
    "    # Update the frequency response curve\n",
    "    freq, h = scipy.signal.freqz_zpk(zeros, poles, 1, fs=44100)\n",
    "    h_db = 20 * np.log10(np.abs(h) + 1e-12)\n",
    "    \n",
    "    frcurve.set_offsets(np.vstack([freq, h_db]).T)\n",
    "    \n",
    "    n = matplotlib.colors.Normalize(vmin = -vmax, vmax = vmax)\n",
    "    m = matplotlib.cm.ScalarMappable(norm=n, cmap=matplotlib.cm.twilight_shifted)\n",
    "    frcurve.set_color(m.to_rgba(h_db))    \n",
    "    frcurve.set_clim(vmin=-vmax, vmax=vmax)\n",
    "        \n",
    "    axf.set(ylim=[-vmax, vmax])\n",
    "\n",
    "    plt.draw()\n",
    "    \n",
    "fig.canvas.mpl_connect('button_press_event', add_or_remove_point);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bloody",
   "language": "python",
   "name": "bloody"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}