{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({\n",
    "    \"text.usetex\": True,\n",
    "    \"font.family\": \"Helvetica\",\n",
    "    \"font.size\": 20\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_lpcnet_features(feature_file, version=2):\n",
    "    if version == 2 or version == 3:\n",
    "        layout = {\n",
    "            'cepstrum': [0,18],\n",
    "            'periods': [18, 19],\n",
    "            'pitch_corr': [19, 20],\n",
    "            'lpc': [20, 36]\n",
    "            }\n",
    "        frame_length = 36\n",
    "\n",
    "    elif version == 1:\n",
    "        layout = {\n",
    "            'cepstrum': [0,18],\n",
    "            'periods': [36, 37],\n",
    "            'pitch_corr': [37, 38],\n",
    "            'lpc': [39, 55],\n",
    "            }\n",
    "        frame_length = 55\n",
    "    else:\n",
    "        raise ValueError(f'unknown feature version: {version}')\n",
    "\n",
    "\n",
    "    raw_features = torch.from_numpy(np.fromfile(feature_file, dtype='float32'))\n",
    "    raw_features = raw_features.reshape((-1, frame_length))\n",
    "\n",
    "    features = torch.cat(\n",
    "        [\n",
    "            raw_features[:, layout['cepstrum'][0]   : layout['cepstrum'][1]],\n",
    "            raw_features[:, layout['pitch_corr'][0] : layout['pitch_corr'][1]]\n",
    "        ],\n",
    "        dim=1\n",
    "    )\n",
    "\n",
    "    lpcs = raw_features[:, layout['lpc'][0]   : layout['lpc'][1]]\n",
    "    if version < 3:\n",
    "        periods = (0.1 + 50 * raw_features[:, layout['periods'][0] : layout['periods'][1]] + 100).long()\n",
    "    else:\n",
    "        periods = torch.round(torch.clip(256./2**(raw_features[:, layout['periods'][0] : layout['periods'][1]] + 1.5), 32, 256)).long()\n",
    "\n",
    "    return {'features' : features, 'periods' : periods, 'lpcs' : lpcs}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "lpcnet_features = load_lpcnet_features('lp/features.f32')\n",
    "\n",
    "features = lpcnet_features['features'].numpy()\n",
    "periods = lpcnet_features['periods'].squeeze(-1).numpy()\n",
    "lpcs = lpcnet_features['lpcs'].numpy()\n",
    "\n",
    "x = np.fromfile('data/critical_3.pcm', dtype=np.int16).astype(np.float32) / 2**15\n",
    "x = np.concatenate((np.zeros(80), x, np.zeros(320)))\n",
    "x_preemph = x.copy()\n",
    "x_preemph[1:] -= 0.85 * x_preemph[:-1]\n",
    "\n",
    "num_frames = features.shape[0]\n",
    "x = x[:160 * num_frames]\n",
    "x_preemph = x_preemph[:160 * num_frames]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def stp(A, alpha=0.8, beta=0.5):\n",
    "    A_num = A * (beta ** np.arange(len(A)))\n",
    "    A_den = A * (alpha ** np.arange(len(A)))\n",
    "    \n",
    "    return A_num, A_den"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame_idx = 31"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7ff3342887f0>]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 18376 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "A = np.concatenate((np.ones(1), lpcs[frame_idx]))\n",
    "A_num, A_den = stp(A)\n",
    "w, h = scipy.signal.freqz([1], A, fs=16000)\n",
    "w, h_stp = scipy.signal.freqz(A_num, A_den, fs=16000)\n",
    "plt.plot(w/1000, 10*np.log10(np.abs(h)))\n",
    "plt.plot(w/1000, 10*np.log10(np.abs(h_stp)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_frame = x_preemph[frame_idx * 160 - 80: (frame_idx + 1) * 160 + 80]\n",
    "\n",
    "window = scipy.signal.get_window('hamming', 320)\n",
    "spec = np.fft.fft(x_frame, n=1024)\n",
    "\n",
    "log_mag = 10*np.log10(np.abs(spec[:512]) + 1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 75208 characters]",
      "text/plain": [
       "<Figure size 648x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(9, 6))\n",
    "\n",
    "plt.plot(w/1000, log_mag - log_mag.mean(), alpha=0.7, label='spectrum')\n",
    "plt.plot(w/1000, 10*np.log10(np.abs(h)), \"r--\", label='LPC spectrum', linewidth=3, alpha=0.9)\n",
    "plt.plot(w/1000, 10*np.log10(np.abs(h_stp)), \"k--\", label='short-term postfilter', linewidth=3, alpha=0.9)\n",
    "plt.legend()\n",
    "plt.xlabel('kHz')\n",
    "plt.ylabel('Amplitude (dB)')\n",
    "plt.show()\n",
    "fig.savefig('plots/short_term_postfilter.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "periods[frame_idx]\n",
    "\n",
    "p = int(periods[frame_idx])\n",
    "ltp_num = np.zeros(p+1)\n",
    "ltp_den = np.zeros(p+1)\n",
    "\n",
    "ltp_num[0] = ltp_den[0] = 1\n",
    "ltp_num[p] = 0.25\n",
    "ltp_den[p] = -0.25\n",
    "\n",
    "w, h_ltp = scipy.signal.freqz(ltp_num, ltp_den, fs=16000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 78740 characters]",
      "text/plain": [
       "<Figure size 648x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(9, 6))\n",
    "plt.plot(w/1000, log_mag - log_mag.mean(), alpha=0.7, label='spectrum')\n",
    "plt.plot(w/1000, 10*np.log10(np.abs(h_ltp)), \"k\", alpha=1, label='long-term postfilter', linewidth=1.5)\n",
    "plt.legend()\n",
    "plt.xlabel('kHz')\n",
    "plt.ylabel('Amplitude (dB)')\n",
    "plt.show()\n",
    "fig.savefig('plots/long_term_postfilter.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
}
