{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation\n",
    "from scipy.io import wavfile\n",
    "import scipy.signal\n",
    "import torch\n",
    "\n",
    "from playback import make_playback_animation\n",
    "from endoscopy import read_data\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({\n",
    "    \"text.usetex\": True,\n",
    "    \"font.family\": \"Helvetica\",\n",
    "    \"font.size\": 24\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 174976 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 8232 characters]",
      "text/plain": [
       "<Figure size 513.2x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 191308 characters]",
      "text/plain": [
       "<Figure size 511.2x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# comparison opus@6kb/s vs. LACE@6kb/s\n",
    "for name in ['opus_6kbps', 'lace_6kbps']:\n",
    "    fs, y = wavfile.read(f'lace_demo/{name}.wav')\n",
    "    spec, freqs, t, im = plt.specgram(y, NFFT=512, cmap='inferno', noverlap=256 + 128, pad_to=4096)\n",
    "    spec = 10*np.log10(spec)\n",
    "    \n",
    "    make_playback_animation(f'lace_demo/{name}_spec.mp4', spec, len(y)/16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = read_data('lace_demo/endoscopy/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_adaconv_impz(data, prefix, frame_size=80, overlap_size=40):\n",
    "    win1 = .5 + .5 * np.cos((np.arange(overlap_size) + 0.5) * torch.pi / overlap_size)\n",
    "    win1 = win1[:, np.newaxis]\n",
    "    win2 = win1[::-1]\n",
    "    kernels = data[prefix + '_kernels']['data'][0] * data[prefix + '_gains']['data'][0, :, np.newaxis]\n",
    "    kernels = kernels[:, ::-1]\n",
    "    num_frames = len(kernels)\n",
    "    impz = np.repeat(kernels, frame_size, 0)\n",
    "    for i in range(num_frames - 2, 0, -1):\n",
    "        idx = i * frame_size\n",
    "        impz[idx : idx + overlap_size] = win2 * impz[idx : idx + overlap_size] + win1 * impz[idx - overlap_size : idx]\n",
    "        \n",
    "    return impz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_adacomb_impz(data, prefix, frame_size=80, overlap_size=40):\n",
    "    win1 = .5 + .5 * np.cos((np.arange(overlap_size) + 0.5) * torch.pi / overlap_size)\n",
    "    win1 = win1[:, np.newaxis]\n",
    "    win2 = win1[::-1]\n",
    "    kernels = data[prefix + '_kernels']['data'][0]\n",
    "    gg = data[prefix + '_global_conv_gains']['data'][0]\n",
    "    g  = data[prefix + '_gains']['data'][0]\n",
    "    lags = data[prefix + '_lags']['data'][0]\n",
    "    \n",
    "    num_frames = len(kernels)\n",
    "    max_lag = int(lags.max())\n",
    "    kernel_size = kernels.shape[1]\n",
    "    padding = kernel_size // 2\n",
    "    \n",
    "    impz = np.zeros((num_frames, max_lag + padding + 1))\n",
    "    for i in range(num_frames):\n",
    "        p = int(lags[i])\n",
    "        impz[i, 0] = gg[i]\n",
    "        impz[i, p - padding : p - padding + kernel_size] = gg[i] * g[i] * kernels[i, ::-1]\n",
    "    \n",
    "    impz = np.repeat(impz, frame_size, 0)\n",
    "    for i in range(num_frames - 2, 0, -1):\n",
    "        idx = i * frame_size\n",
    "        impz[idx : idx + overlap_size] = win2 * impz[idx : idx + overlap_size] + win1 * impz[idx - overlap_size : idx]\n",
    "    \n",
    "    return impz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "impz_comb1 = get_adacomb_impz(data, 'limited_adaptive_comb1d_1')\n",
    "impz_comb2 = get_adacomb_impz(data, 'limited_adaptive_comb1d_2')\n",
    "impz_conv1 = get_adaconv_impz(data, 'limited_adaptive_conv1d_1')\n",
    "\n",
    "phi = data['hidden_features']['data'][0, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 119776 characters]",
      "text/plain": [
       "<Figure size 864x1296 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# LACE animation\n",
    "\n",
    "fig, axs = plt.subplots(4, 2)\n",
    "frame_duration=40\n",
    "fig.set_size_inches(12, 18)\n",
    "axs[0, 0].set_xlabel('Tap')\n",
    "axs[0, 0].set_ylabel('Amplitude')\n",
    "axs[0, 1].set_xlabel('Frequency (kHz)')\n",
    "axs[0, 1].set_ylabel('Amplitude (dB)')\n",
    "axs[1, 0].set_xlabel('Tap')\n",
    "axs[1, 0].set_ylabel('Amplitude')\n",
    "axs[1, 1].set_xlabel('Frequency (kHz)')\n",
    "axs[1, 1].set_ylabel('Amplitude (dB)')\n",
    "axs[2, 0].set_xlabel('Tap')\n",
    "axs[2, 0].set_ylabel('Amplitude')\n",
    "axs[2, 1].set_xlabel('Frequency (kHz)')\n",
    "axs[2, 1].set_ylabel('Amplitude (dB)')\n",
    "axs[3, 0].set_xlabel('Dim')\n",
    "axs[3, 1].set_visible(False)\n",
    "fig.tight_layout()\n",
    "\n",
    "frames = []\n",
    "for i in range(12800, 14000, 4):\n",
    "    f = impz_comb1[i]\n",
    "    w, h = scipy.signal.freqz(f, fs=16000)\n",
    "    frame = axs[0, 0].plot(f, 'b') + axs[0, 1].plot(w/1000, 10 * np.log10(np.abs(h)), 'b')\n",
    "    \n",
    "    f = impz_comb2[i]\n",
    "    w, h = scipy.signal.freqz(f, fs=16000)\n",
    "    frame += axs[1, 0].plot(f, 'b') + axs[1, 1].plot(w/1000, 10 * np.log10(np.abs(h)), 'b')\n",
    "    \n",
    "    f = impz_conv1[i]\n",
    "    w, h = scipy.signal.freqz(f, fs=16000)\n",
    "    frame += axs[2, 0].plot(f, 'b') + axs[2, 1].plot(w/1000, 10 * np.log10(np.abs(h)), 'b')\n",
    "    \n",
    "    frame += axs[3, 0].plot(phi[i//80], 'b')\n",
    "    \n",
    "    frames.append(frame)\n",
    "    \n",
    "ani = matplotlib.animation.ArtistAnimation(fig, frames, blit=True, interval=frame_duration)\n",
    "ani.save('lace_demo/responses2.mp4', dpi=720)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 124156 characters]",
      "text/plain": [
       "<Figure size 864x1296 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# static plot for Jean-Marc's on-site presentation\n",
    "\n",
    "fig, axs = plt.subplots(4, 2)\n",
    "frame_duration=40\n",
    "fig.set_size_inches(12, 18)\n",
    "axs[0, 0].set_xlabel('Tap')\n",
    "axs[0, 0].set_ylabel('Amplitude')\n",
    "axs[0, 1].set_xlabel('Frequency (kHz)')\n",
    "axs[0, 1].set_ylabel('Amplitude (dB)')\n",
    "axs[1, 0].set_xlabel('Tap')\n",
    "axs[1, 0].set_ylabel('Amplitude')\n",
    "axs[1, 1].set_xlabel('Frequency (kHz)')\n",
    "axs[1, 1].set_ylabel('Amplitude (dB)')\n",
    "axs[2, 0].set_xlabel('Tap')\n",
    "axs[2, 0].set_ylabel('Amplitude')\n",
    "axs[2, 1].set_xlabel('Frequency (kHz)')\n",
    "axs[2, 1].set_ylabel('Amplitude (dB)')\n",
    "axs[3, 0].set_xlabel('Tap')\n",
    "axs[3, 0].set_ylabel('Amplitude')\n",
    "axs[3, 1].set_xlabel('Frequency (kHz)')\n",
    "axs[3, 1].set_ylabel('Amplitude (dB)')\n",
    "fig.tight_layout()\n",
    "i=10*80\n",
    "\n",
    "f = impz_comb1[i]\n",
    "w, h = scipy.signal.freqz(f, fs=16000)\n",
    "axs[0, 0].plot(f, 'b')\n",
    "axs[0, 1].plot(w/1000, 10 * np.log10(np.abs(h)), 'b')\n",
    "\n",
    "f = impz_comb2[i]\n",
    "w, h = scipy.signal.freqz(f, fs=16000)\n",
    "axs[1, 0].plot(f, 'b')\n",
    "axs[1, 1].plot(w/1000, 10 * np.log10(np.abs(h)), 'b')\n",
    "\n",
    "f = impz_conv1[i]\n",
    "w, h = scipy.signal.freqz(f, fs=16000)\n",
    "axs[2, 0].plot(f, 'b')\n",
    "axs[2, 1].plot(w/1000, 10 * np.log10(np.abs(h)), 'b')\n",
    "\n",
    "f = np.convolve(np.convolve(impz_comb1[i], impz_comb2[i], mode='full'), impz_conv1[i])\n",
    "w, h = scipy.signal.freqz(f, fs=16000)\n",
    "axs[3, 0].plot(f, 'b')\n",
    "axs[3, 1].plot(w/1000, 10 * np.log10(np.abs(h)), 'b')\n",
    "\n",
    "fig.savefig('plots/lace_snapshot_unvoiced.png')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "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": 2
}
