{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as pat\n",
    "import scipy.linalg as sl\n",
    "import numpy as np\n",
    "import librosa\n",
    "import librosa.display\n",
    "import IPython.display as ipd\n",
    "\n",
    "from copy import deepcopy\n",
    "import tqdm\n",
    "\n",
    "# additionally requires libcpab (https://github.com/SkafteNicki/libcpab)\n",
    "import nwarp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mfcc(filename, delta=False, segments=False, start=None, stop=None):\n",
    "    y, sr = librosa.core.load(filename, sr=None)\n",
    "    if segments:\n",
    "        y = y[start:stop]\n",
    "\n",
    "    # stretch samples so that K << T\n",
    "    y = librosa.effects.time_stretch(y, 1/250)\n",
    "\n",
    "    # Mel-frequency cepstral coefficients 1-12\n",
    "    mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=12, hop_length=512, n_mels=128, fmax=sr/2)\n",
    "    if delta:\n",
    "        mfccs = np.vstack([mfccs, librosa.feature.delta(mfccs), librosa.feature.delta(mfccs, order=2)])\n",
    "\n",
    "    # normalize\n",
    "    mfccs = (mfccs.transpose()-mfccs.mean(axis=1))/mfccs.std(axis=1)\n",
    "    mfccs = mfccs.transpose()\n",
    "\n",
    "    return torch.from_numpy(np.array(mfccs)[np.newaxis].swapaxes(1,2))\n",
    "\n",
    "def play_sequence(filepath, start, end, stretch=1):\n",
    "    '''\n",
    "    makes WAV playable in jupyter, returns array with amplitude and sample rate from file.\n",
    "    Option to stretch sound. stretch < 1 -> longer sample.\n",
    "    '''\n",
    "    y, sr = librosa.load(filepath, sr=None)\n",
    "    y = librosa.effects.time_stretch(y, stretch)\n",
    "    ipd.display(ipd.Audio(y[start:end], rate=sr))\n",
    "    return y, sr\n",
    "\n",
    "def almat_transition_points(almat):\n",
    "    '''\n",
    "    find the transition points in the alignment matrices, \n",
    "    padding to add start and stop \"transition\", and scale to [0,1]\n",
    "    '''\n",
    "    segments = (almat[:, :-1, :]\n",
    "           - almat[:, 1:, :]).abs().sum(axis=2).sign().nonzero()[:, 1:].view(almat.shape[0], -1) / float(almat.shape[1])\n",
    "    return np.pad(segments[0], pad_width=(1,1), constant_values=(0,1))\n",
    "\n",
    "def alignment(transitionPointsEstimated, transitionPointsTrue):\n",
    "    '''\n",
    "    distance from each true phoneme transition to the closest predicted one\n",
    "    '''\n",
    "    distances = np.abs(transitionPointsEstimated - transitionPointsTrue[:, None])\n",
    "    return distances.min(axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load and preprocess data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "        0      1    2\n",
      "0       0   5800   h#\n",
      "1    5800   9800  pau\n",
      "2    9800  10120   dh\n",
      "3   10120  11613   eh\n",
      "4   11613  12440    n\n",
      "5   12440  12600   dh\n",
      "6   12600  13560   ix\n",
      "7   13560  14600  kcl\n",
      "8   14600  15640    k\n",
      "9   15640  16710   ao\n",
      "10  16710  17503    r\n",
      "11  17503  18673   iy\n",
      "12  18673  21080   aa\n",
      "13  21080  22200  gcl\n",
      "14  22200  22600    g\n",
      "15  22600  23320  axr\n",
      "16  23320  25400    f\n",
      "17  25400  26920   er\n",
      "18  26920  28532    m\n",
      "19  28532  29535   ah\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 32912 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "N = 1\n",
      "T = 6016\n",
      "Dout = 12\n"
     ]
    }
   ],
   "source": [
    "# Please obtain the TIMIT corpus from https://catalog.ldc.upenn.edu/LDC93S1.\n",
    "\n",
    "filename  = \"TIMIT/SI1086\" # the full path in the corpus is \"timit/TRAIN/DR3/FALK0/SI1086\"\n",
    "phn = pd.read_csv(filename + \".PHN\", delimiter=' ', header=None)\n",
    "print(phn.head(20))\n",
    "\n",
    "# we only use part of the speech signal\n",
    "sequence       = [8, 18]\n",
    "segments       = sequence[1] - sequence[0]\n",
    "sequenceStart  = phn.iloc[sequence[0], 0]\n",
    "sequenceEnd    = phn.iloc[sequence[1], 0]\n",
    "\n",
    "y, sr = play_sequence(filename + '.WAV', sequenceStart, sequenceEnd)\n",
    "toleranceLevel = sr * 0.02 / (sequenceEnd - sequenceStart)\n",
    "\n",
    "# true transition points, scaled to [0,1] and MFCCs\n",
    "transitionPointsTrue = np.array([phn.iloc[i, 0] for i in range(sequence[0], sequence[1]+1)])\n",
    "transitionPointsTrue = (transitionPointsTrue - sequenceStart) / (sequenceEnd - sequenceStart)\n",
    "transitionPointsTrue = transitionPointsTrue[1:-1]\n",
    "y_train = mfcc(filename + \".WAV\", \n",
    "                segments=True, \n",
    "                delta=False,\n",
    "                start=sequenceStart, \n",
    "                stop=sequenceEnd)\n",
    "\n",
    "N, T, Dout = y_train.shape\n",
    "print(f\"N = {N}\\n\"\n",
    "      f\"T = {T}\\n\"\n",
    "      f\"Dout = {Dout}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ParameterWarp(\n",
      "  (warp): TSPStepWarp(\n",
      "    (loc_net): Constant()\n",
      "  )\n",
      "  (resample): Resample()\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# PIECEWISE CONSTANT: simply output the warped parameter vector\n",
    "# TSP-based parameter warping with constant mode vector\n",
    "K = len(transitionPointsTrue)+1 # number of segments\n",
    "Dparamg = 0 # number of global parameters (same in every segment)\n",
    "Dparaml = Dout # number of local parameters (different in every segment)\n",
    "paramwarp = nwarp.ParameterWarp(K, Dparamg, Dparaml,\n",
    "               nwarp.TSPStepWarp(nwarp.Constant((K-1,)),\n",
    "                                 width=0.125, power=16.,\n",
    "                                 min_step=0.0001, max_step=0.9999))\n",
    "\n",
    "print(paramwarp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_restarts    =  10 # number of randomized restarts\n",
    "n_epochs      = 300 # total number of epochs\n",
    "n_epochs_hard = 100 # use hard segmentation for the last X epochs\n",
    "show_plots    = False\n",
    "loss_fn       = nn.MSELoss(reduction='mean')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "restart  1/10: 100%|██████████| 300/300 [00:16<00:00, 18.39it/s, max=1.0296, cur=0.1904]\n",
      "restart  2/10: 100%|██████████| 300/300 [00:16<00:00, 18.05it/s, max=1.0175, cur=0.1954]\n",
      "restart  3/10: 100%|██████████| 300/300 [00:15<00:00, 19.94it/s, max=1.0048, cur=0.1928]\n",
      "restart  4/10: 100%|██████████| 300/300 [00:17<00:00, 16.86it/s, max=1.0578, cur=0.2156]\n",
      "restart  5/10: 100%|██████████| 300/300 [00:17<00:00, 17.30it/s, max=1.0334, cur=0.1956]\n",
      "restart  6/10: 100%|██████████| 300/300 [00:15<00:00, 19.82it/s, max=1.0706, cur=0.1985]\n",
      "restart  7/10: 100%|██████████| 300/300 [00:15<00:00, 18.88it/s, max=1.0405, cur=0.1977]\n",
      "restart  8/10: 100%|██████████| 300/300 [00:15<00:00, 18.85it/s, max=1.0075, cur=0.1983]\n",
      "restart  9/10: 100%|██████████| 300/300 [00:17<00:00, 17.27it/s, max=1.0336, cur=0.1940]\n",
      "restart 10/10: 100%|██████████| 300/300 [00:16<00:00, 17.70it/s, max=1.0220, cur=0.1851]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best loss = 0.1851\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "best_loss = np.inf\n",
    "for r in range(n_restarts):\n",
    "    # reset everything\n",
    "    optimizer = torch.optim.Adam(paramwarp.parameters(), lr=1e-1)\n",
    "    param_norm = []\n",
    "    grad_norm = []\n",
    "    train_losses = []\n",
    "    resample_kernel = \"linear\"\n",
    "    epoch_counter = tqdm.tqdm(range(n_epochs), desc=f\"restart {(r+1):2d}/{n_restarts:2d}\")\n",
    "\n",
    "    # initialize parameters\n",
    "    _ = paramwarp.apply(nwarp.reset_parameters)\n",
    "    nn.init.uniform_(paramwarp.warp.loc_net.const, -1., 0.) # segmentation\n",
    "\n",
    "    # perform training\n",
    "    paramwarp.train()\n",
    "    for epoch in epoch_counter:\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        if epoch == n_epochs - n_epochs_hard:\n",
    "            resample_kernel = \"integer\"\n",
    "\n",
    "        param_hat_train = paramwarp(y_train, resample_kernel=resample_kernel)[0]\n",
    "        y_hat_train = param_hat_train\n",
    "\n",
    "        loss = loss_fn(y_hat_train, y_train)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        train_losses.append(loss.item())\n",
    "        param_norm.append([sl.norm(p.detach()) for p in paramwarp.parameters() if len(p)>0])\n",
    "        grad_norm.append([sl.norm(p.grad.detach()) for p in paramwarp.parameters() if len(p)>0])\n",
    "        epoch_counter.set_postfix({\"max\": f\"{max(train_losses):.4f}\", \"cur\": f\"{loss.item():.4f}\"})\n",
    "\n",
    "    if train_losses[-1] < best_loss:\n",
    "        best_paramwarp_state = deepcopy(paramwarp.state_dict())\n",
    "        best_loss = train_losses[-1]\n",
    "\n",
    "    if show_plots:\n",
    "        plt.figure(figsize=(15,2))\n",
    "        plt.subplot(131)\n",
    "        plt.title(\"loss\")\n",
    "        plt.ylim(np.min(train_losses), np.percentile(train_losses, 95))\n",
    "        plt.plot(train_losses)\n",
    "\n",
    "        plt.subplot(132)\n",
    "        plt.title(\"parameter norm\")\n",
    "        lines = plt.plot(np.array(param_norm)/np.array(param_norm).max(axis=0))\n",
    "        plt.legend(lines, [\" x \".join([str(d) for d in p.size()]) for p in paramwarp.parameters() if len(p)>0])\n",
    "\n",
    "        plt.subplot(133)\n",
    "        plt.title(\"gradient norm\")\n",
    "        normalized_grad_norm = np.array(grad_norm)/np.array(grad_norm).max(axis=0)\n",
    "        lines = plt.plot(normalized_grad_norm)\n",
    "        plt.legend(lines, [\" x \".join([str(d) for d in p.size()]) for p in paramwarp.parameters() if len(p)>0])\n",
    "        plt.ylim(np.min(normalized_grad_norm), np.percentile(normalized_grad_norm, 95))\n",
    "\n",
    "        plt.show()\n",
    "\n",
    "paramwarp.eval()\n",
    "paramwarp.load_state_dict(best_paramwarp_state)\n",
    "print(f\"best loss = {best_loss:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_hat_train, almat_hat_train, gamma_hat_train = paramwarp(y_train, resample_kernel=resample_kernel)\n",
    "y_hat_train = param_hat_train\n",
    "\n",
    "transitionPointsEstimated = almat_transition_points(almat_hat_train)[1:-1]\n",
    "yEstimated   = np.zeros_like(transitionPointsEstimated)\n",
    "yTrue        = np.ones_like(transitionPointsTrue)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 137008 characters]\n",
      "text/plain": [
       "<Figure size 864x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "inp = np.array(y_train[0, :,:12].t().detach())\n",
    "out = np.array(y_hat_train[0, :, :12].t().detach())\n",
    "plt.figure(figsize=(12,4))\n",
    "plt.subplot(211, title=\"true\")\n",
    "plt.imshow(inp, aspect='auto', cmap=\"magma\")\n",
    "for xc in transitionPointsTrue*6016:\n",
    "    plt.axvline(x=xc, color='w', linestyle='--')\n",
    "    \n",
    "plt.subplot(212, title=\"predicted\")\n",
    "plt.imshow(out, aspect='auto', cmap=\"magma\")\n",
    "for xc in transitionPointsEstimated*6016:\n",
    "    plt.axvline(x=xc, color='w', linestyle='--')\n",
    "    \n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7 out of 9 segment boundaries detected with a tolerance of 20 ms\n",
      "random segmentation is worse in 9932 out of 10000 trials (99.32%)\n"
     ]
    }
   ],
   "source": [
    "criterion = np.sum(alignment(transitionPointsEstimated, transitionPointsTrue) < toleranceLevel)\n",
    "print(f\"{criterion} out of {len(transitionPointsTrue)} segment boundaries detected with a tolerance of 20 ms\")\n",
    "\n",
    "random_trials = 10000\n",
    "random_detectors = np.array([np.sum(alignment(np.sort(np.random.uniform(size=9)), transitionPointsTrue) < toleranceLevel) for i in range(random_trials)])\n",
    "random_worse = np.sum(random_detectors < criterion)\n",
    "print(f\"random segmentation is worse in {random_worse} out of {random_trials} trials ({random_worse/random_trials*100}%)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
