retrieval_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Efficient Retrieval of Music Recordings Using Graph-Based Index Structures\n",
    "\n",
    "This notebook accompanies the following paper:\n",
    "\n",
    "<ul>    \n",
    "<li><span style=\"color:black\">\n",
    "Frank Zalkow, Julian Brandner, and Meinard Müller. <strong>Efficient retrieval of music recordings using graph-based index structures.</strong> <em>Signals</em>, 2(2):336–352, 2021. doi: 10.3390/signals2020021.\n",
    "    <a class=\"btn btn-default btn-xs\" href=\"https://www.mdpi.com/2624-6120/2/2/21\">Paper</a>\n",
    "    <a class=\"btn btn-default btn-xs\" href=\"https://www.audiolabs-erlangen.de/resources/MIR/2020_signals-indexing/reference_paper.bib\">Bibtex</a>\n",
    "    <a class=\"btn btn-default btn-xs\" href=\"https://www.audiolabs-erlangen.de/resources/MIR/2020_signals-indexing\">Website</a>\n",
    "</ul>\n",
    "    \n",
    "The paper compares indexing strategies based on $K$-d trees and HNSW graphs in a cross-modal music retrieval application. See more details in the paper. The notebook shows how to load the features for a provided example dataset, generate shingles, reduce the shingle dimensionality using PCA, construct index structures ($K$-d trees and HNSW graphs), and search for the nearest shingles to a given query in the database.\n",
    "    \n",
    "We start by importing some Python packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import timeit\n",
    "\n",
    "import numpy as np\n",
    "import scipy\n",
    "import pandas as pd\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.neighbors import NearestNeighbors, KDTree\n",
    "import IPython.display as ipd\n",
    "import hnswlib\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "from common import compute_features, generate_shingles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters\n",
    "\n",
    "We now set various parameters used throughout this notebook.\n",
    "First, we decide whether we use dimensionality reduction (binary variable `dimensionality_reduction`).\n",
    "Second, we fix the used dimensionality $K$ (`K`, changed to 240 later if no dimensionality reduction is applied).\n",
    "Third, we set the number $\\nu$ of neighbors to search for (`nu`).\n",
    "\n",
    "Then, we have three more parameters for the HNSW graph, i.e., number $\\nu'$ of neighbor candidates during the search (`nu_prime`), the minimum number $M$ of edges for each node in the graph (`M`), and the number $M'$ of edge candidates during the construction of the graph (`M_prime`).\n",
    "For details about the meaning of these parameters, we refer to our paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dimensionality_reduction = True\n",
    "K = 30\n",
    "nu = 10\n",
    "\n",
    "nu_prime = 100\n",
    "M = 5\n",
    "M_prime = 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Dataset\n",
    "\n",
    "We next load the metadata and the CENS features for our dataset.\n",
    "Then, we generate shingles (as vectors of 240 dimensions) from the feature sequences and an index that specifies the original document for each shingle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata = pd.read_csv(os.path.join('data', 'metadata.csv'), sep=';')\n",
    "\n",
    "shingles = []\n",
    "index = []\n",
    " \n",
    "for _, (cur_idx, cur_composer, cur_work, cur_version, cur_fn) in metadata.iterrows():\n",
    "    cur_fn_npz = os.path.join('data', 'features_cens', cur_fn)\n",
    "    cur_X = np.load(cur_fn_npz)['f_chroma']\n",
    "    cur_shingles = generate_shingles(cur_X)\n",
    "    \n",
    "    shingles.append(cur_shingles)\n",
    "    index.append(np.repeat(cur_idx, cur_shingles.shape[0]))\n",
    "    \n",
    "shingles = np.concatenate(shingles, axis=0)\n",
    "index = np.concatenate(index)\n",
    "\n",
    "db_size = shingles.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To give an overview of the dataset, we create a table with the dataset's musical works, showing the number of versions for each work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>ComposerID</th>\n",
       "      <th>WorkID</th>\n",
       "      <th># Versions</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>Beethoven</td>\n",
       "      <td>Op055-01</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>Beethoven</td>\n",
       "      <td>Op055-02</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>Beethoven</td>\n",
       "      <td>Op055-03</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>Beethoven</td>\n",
       "      <td>Op055-04</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>53</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op024No1</td>\n",
       "      <td>61</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op033No2</td>\n",
       "      <td>67</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op033No3</td>\n",
       "      <td>50</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op068No4</td>\n",
       "      <td>62</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>Vivaldi</td>\n",
       "      <td>RV269-01</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>Vivaldi</td>\n",
       "      <td>RV269-02</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>Vivaldi</td>\n",
       "      <td>RV269-03</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "works_with_count = metadata[['ComposerID', 'WorkID']].value_counts(sort=False).to_frame()\n",
    "works_with_count = pd.DataFrame(works_with_count.to_records())\n",
    "works_with_count.columns = ['ComposerID', 'WorkID', '# Versions']\n",
    "\n",
    "ipd.display(ipd.HTML(works_with_count.to_html(index=False)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To give a concrete idea about the dataset content, we now select a random shingle from the dataset, show its corresponding metadata, and visualize the shingle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Index</th>\n",
       "      <th>ComposerID</th>\n",
       "      <th>WorkID</th>\n",
       "      <th>VersionID</th>\n",
       "      <th>Filename</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>309</th>\n",
       "      <td>309</td>\n",
       "      <td>Vivaldi</td>\n",
       "      <td>RV269-01</td>\n",
       "      <td>Abbado</td>\n",
       "      <td>Vivaldi_RV269-01_Abbado.npz</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAU8AAACsCAYAAAD7ai5qAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8vihELAAAACXBIWXMAAAsTAAALEwEAmpwYAAAZQ0lEQVR4nO3debRdZZnn8e8vc0gCNwOEEIIMBcQIEkOkEBFwKItALZAquwUspS00K7WIrbWWC6m2WrHtKsS5VDSdVtrgAFa1Q1EaBhtRqIWUEAghAcEwCGFICGGohJBwk6f/2PuGk5Mz7Hffu2/Ouff3WeusnLP3u/f77nsOD+8e3vdRRGBmZmlG7O0GmJl1IwdPM7MSHDzNzEpw8DQzK8HB08ysBAdPM7MSHDzNbEiTdKWkDZJWN1kvSV+VtFbSKknziuzXwdPMhrrvAKe3WL8AODJ/LQS+WWSnDp5mNqRFxC3AphZFzgauisztQI+kGe32O2qgGjhQenp6YsaMtu3eZdy4ccl1jBw5Mnkbs4G0c+fOpPLbtm2rqCWvGjt2bPI2I0ak9b9WrFixMSL2T64IkNRsOOQa4OWaz0sjYmnCrmcCj9d8Xpcve6rVRh0XPGfMmMGyZcsKlz/66KOT65g0aVJS+dQfSBmpw2TLDKuVVGl5K27Lli1J5R9++OHkOlK/vyOOOCK5jvHjxyeVl/SH5Ep2336PZRHxckTM789uGyxr+x9YpcFT0g7gXrLG7QAWR8RtVdZpZkNXo47Mjh07+rvbdcCsms8HA0+2bUt/a21ja0TMjYjjgL8FLqu4PjMboiQxYsSIPV4D4Frg/fld9xOBFyKi5Sk7DO5p+77Ac4NYn5kNMWWCpaSrgdOAaZLWAZ8CRgNExBJgOXAGsBZ4CfhAkf1WHTzHS1oJjANmAG+ruD4zG8LKBM+IOK/N+gAuSt1v1cFza0TMBZD0JuAqScdE3d0OSQvJnq/iwAMPrLhJZtaN+k7bO8WgtSQifgNMA/Z4TCEilkbE/IiY39PTM1hNMrMuU9E1z1IG7ZqnpNnASODZwarTzIaOTut5DtY1T8geV7ogIvr9XIGZDU/DJnhGRPJQngkTJnDCCSdU0ZyO5gfYu1vqs4aPPPJIUvkNGzYklQc46qijksqnPvC+Nwyb4GlmNlCG22m7mdmA6aR5KRw8zawruOdpZlaSg6eZWSL3PM3MSnLwNDMrwcHTzCyRT9vNzEpy8DQzS+SeZxvr16/ny1/+cuHyU6dOTa4jNe/R008/nVzHK6+8klT+xBNPTCpfJvHdAKQraOu559Lmu079j+GAAw5IKg8wGDN1Pf/880nlt2/fnlT+kEMOSSoPcNBBByVv0+k6KXh2TkvMzNooOyWdpNMlPSBpraRLGqzfT9K/SrpH0hpJbWeTrzx4SjpHUuRT0pmZlVI2h5GkkcAVwAJgDnCepDl1xS4C7svzrZ0GfFHSmFb7HYye53nAvwHnDkJdZjaElex5ngCsjYiHI2I7cA1wdl2ZACYpm65sIrAJ6G3ZlvTmFydpIvBm4EIcPM2sHyQxcuTIPV5kid3urHktrNt0JvB4zed1+bJaXwdeS5Zy+F7gIxGxs1V7qr5h9C7g+oh4UNImSfMi4q76QrU5jCZPnlxxk8ysWzXpaW6MiPktNms0+W3Uff5TYCVZksojgF9IujUiXmzaltZN7bfzyLrI5P82zGJXm8NowoQJFTfJzLpRP/K2rwNm1Xw+mKyHWesDwI8jsxZ4BGh5n6aynqekqWRR/BhJQZa/KCRdXJ8908ysiJKPKt0BHCnpMOAJskuI59eVeQx4O3CrpOnA0cDDLdtSpiUFvRu4KiJeExGHRsQssmh+coV1mtkQVqbnGRG9wGLgBuB+4J8iYo2kRZIW5cU+A5wk6V7gJuDjEbGx1X6rvOZ5HvDZumU/Iov4t1ZYr5kNQf0ZYRQRy4HldcuW1Lx/Enhnyj4rC54RcVqDZV9tt11vby/PPPNM4XpefvnltIYBN910U1L51NEjkCWyS/Hkk/WXYFobM6blI2gNbd68ufI6Uq/IpI5IWr9+fVJ5SE+WN3HixOQ6Uh177LFJ5c8888zkOlJTVqSOitsbOmmEUccNzzQza8Rj283MSnLwNDMrwcHTzCyRT9vNzEpy8DQzS+Sep5lZSamPX1XJwdPMuoJ7nmZmJTl4ttDT08PZZ9fPU9rcfvvtl1zH6tWrk8qnjv6B9JwzJ510UlL5adOmJZWHzvrh9dm6dWtS+TKjvVJNmTIleZuxY8dW0JLBNXr06ORtBiMvVh/3PM3MSuqk4Fn1TPLTJf1A0sOSVkj6jaRzqqzTzIausgngKmlLVTvOc4H8FLglIg6PiOPJ5tE7uKo6zWzo6sdkyJWo8rT9bcD2ummf/gB8rcI6zWwI67rTdkknSrpD0mZJ2yXtkNQ0t0fudcAe+YrMzMrotJ5n0Zq/Tja58e+B8cAHSexBSroiTyh/R4N1C/sy3w3G3VQz605lg6ek0yU9IGmtpEualDlN0kpJayT9um1bijY6T4o0MiJ2RMT/Ad7aZpM1wLya7S8iyxGyf4N970oA19PTU7RJZjbMlAmekkYCVwALgDnAeZLm1JXpAb4BnBURrwP+U9u2FGzzS5LGACslfU7S3wDtpkr/JTBO0l/XLNunYH1mZrvpx2n7CcDaiHg4IraTZfKtf5j8fLLsmY8BRMSGdjstGjzfR5b9cjGwhSyN51+02iDPkPku4FRJj0j6LbAM+HjBOs3MdtMkeE7ru+yXvxbWbTYTeLzm87p8Wa2jgMmSfpU/Vvn+dm0pdLc9v0sOsBX4dJFt8u2eIns8ycysX1qMMNoYEfNbbdpgWX2yrVHA8WSXFscDv5F0e0Q82GynhYKnpD8jS835mnwbkXUu9y2yfYpNmzbxgx/8oHD52bNb5qVv6OST07IfH3jggcl1pCaAS01SVmbIaGpytjJSZ73ZuXNnpeUhfbhlJw61LHPcqVJ/gwDbtm2roCXNlZxVaR3Z2XKfg4H6/4DWkQXhLcAWSbcAxwFNg2fR0/avABcAUyNi34iYVEXgNDNrph/XPO8AjpR0WH7v5lzg2roy/wK8RdIoSfsAf0yW472pog/JPw6sjsHoupiZNVHmuc6I6JW0GLiB7N7NlRGxRtKifP2SiLhf0vXAKmAn8K2IaDmDUNHgeTGwPH/2aVc/PSK+lHwkZmYllX0oPiKWA8vrli2p+/x54PNF91k0eP49sBkYB4wpunMzs4HSrVPSTYmId1baEjOzNjopeBZtyf+T5OBpZntNp41tL9rzvAi4WNJ24JV8WSWPKpmZNdNJPc+iD8lPqrohZmbtdF3wBJB0FnBK/vFXEfGzappkZranrrxhJOmzwBuB7+eLPiLp5IhoOLVTf8ycOZPLLruscPknnngiuY599kmbn+S1r31tch2jRjk9lA2cTgoatcaNGzeo9XXS36Hof+FnAHMjYieApGXA3cCAB08zs0a6sueZ6wE25e8L5/uVtAO4t2bRuyLi0YR6zcyA7ux5/gNwt6SbySYFOQX424Lbbo2IuSXaZma2i6SyE4NUom3wlDSCbKzniWTXPQV8PCKerrhtZma76aqeZ0TslLQ4Iv6JPWciKWK8pJX5+0ciYo+87fnkpQsBZs2aVb/azAzosuCZ+4WkjwE/JJtJHoCI2NR8k13anrZHxFJgKcC8efM8c5OZ7aFbbxj9Vf7vRTXLAjh8YJtjZtZc1wXPiDis6oaYmbXSrT1PJJ0EHFq7TURcVUGbzMwaKhs8JZ0O/CPZZMjfiojPNin3RuB24D0R8X9b7bPoCKPvAkcAK4Ed+eIA2gbPiJhYpI4+L730EnfeeWfh8ill+8yf3ypX1J7233+PVPMDLjUXzJYtW9oXqrNjx472hWr09vYm15Haru3btyfXkSo1h9HUqVOT60j9/jZu3JhUfvz48UnlIf13WyYwPfbYY8nb9EeZNtbkbf8TslxFd0i6NiLua1DucrIZ59sq2vOcD8xxGg4z21v6cdq+K297vp++vO331ZX7MPAjskcy2yraktVAegpJshFGklbWvDyk08xKqSpvu6SZwDnAbqk5WmnZ85T0r2Sn55OA+yT9lt1zGJ1VoA6PMDKzfqs4b/tXyAb/7Ciagrndafu1wHTg1rrlpwLp0xmZmfVDydP2Innb5wPX5IFzGnCGpN6I+GmznbYLnmcD/y0iVtUulLQF+BTw7QINrx1hBHBZRPywwHZmZrspGTx35W0n6/SdC5xfW6D2cUxJ3wF+1ipwQvvgeWh94MwrulPSoYWaXeC0vXZ45vTp0wvu1syGk7I3jIrkbS/TnnbBs9VMp+nPTjRROzxz9uzZvqNvZnvoz6xKRfK21yz/L0X22S6M3yHpQ/ULJV0IrChSgZnZQOmm7JkfBX4i6b28GiznA2PIbusXUX/N8/oq0neY2dDXNcMzI2I9cJKktwLH5It/HhG/LFpBRHTO7KVm1rW6cmx7RNwM3FxxW4Ds/yyTJhXPdDx79uzkOl588cWk8rfddltyHZMnT04qP3PmzPaFapRJvDV69Oik8mV+qKlDG19++eWk8s8//3xSeUgflppaHmDixKRRyMnlO1VPT8+g1td1wdPMbG/ryp6nmVkncPA0MyvBwdPMLJFP283MSnLwNDNL5J6nmVlJDp5mZiU4eJqZJerPxCBV6LjgOW7cuKRRQ2X+mKmpmFJH/wC88MILSeUff/zx9oVqzJgxI6k8pI9qKZOyKvX7SK2jzMic1ORpY8eOTa5juPIIIzOzDjdsbhhJ2gHcC4wGeoFlwFciYmdVdZrZ0NZJwbPKlmyNiLkR8TqyfMlnkKXuMDNL1tfzLDOfp6TTJT0gaW2jDL6S3itpVf66TdJx7fY5KGE8IjaQpdlYrKKp6czM6pQJnpJGAlcAC4A5wHmS5tQVewQ4NSJeD3yGPLNFy7Ykt76kPOH8COCA+nWSFvblXN64ceNgNcnMukzJnucJwNqIeDgitgPXkCW33CUibouI5/KPt5Nl2GzdlsS291fDXmdELI2I+RExf9q0aYPcJDPrBi1O26f1db7y18K6TWcCtY+zrMuXNXMhcF279gza3XZJhwM7gA2DVaeZDS1NepobI2J+i80addoaPiOXZ824EDi5XVsGJXhK2h9YAnw9yjw8aGbDXj8eVVoHzKr5fDDwZIP9vx74FrAgIp5tt9Mqg2df4re+R5W+C3ypwvrMbIgrGTzvAI6UdBjwBHAucH5tAUmHAD8G3hcRDxbZaWXBs2zit97eXp59tm3Q3+Whhx5KruPBBwv9bXY56qijkuu4++67k8pfffXVSeVTc/8AjBqV9nXv3Jn+SG7qj3vfffdNKj9lypSk8gDbtm1LKv/cc8+1L9RPW7ZsqbQ8wNatW5PKl/lNDbYywTMieiUtBm4ARgJXRsQaSYvy9UuATwJTgW/kDwT1trkU4BFGZtYd+jPCKCKWA8vrli2pef9B4IMp+3TwNLOu4YlBzMwSDZux7WZmA83B08wskXueZmYlOXiamZXg4Glmlsin7WZmJTl4mpklcs+zjTFjxnDwwW2n0tulTCK01LlJygyNmzt3blL5RYsWJZWfMGFCUnlIT2xWJhFa6lzXmzdvTir/zDPPJJWH9OGZ06dPT65j8uTJydukKDPP7bp165LK77fffsl1HHLIIUnlU4cI1+uk4Fl5SyQdKOkaSQ9Juk/Scknpg8XNbNgrm4ajCpX2PPOUGz8BlkXEufmyucB0IG12DjMb1obbaftbgVfqBuCvrLhOMxuihlPwPAZYUXEdZjYMDLeeZyF5zpGFkH4B2syGj06aVanqML4GOL5dodoEcPvvv3/FTTKzbtSfvO1VqLrmXwJjJX2ob4GkN0o6teJ6zWwIKhs8JZ0u6QFJayVd0mC9JH01X79K0ry2bSnR/sLyZG/nAH+SP6q0BriUBsmXzMxaKdvzlDQSuAJYAMwBzpM0p67YAuDI/LUQ+Ga7/VZ+zTMingT+c9X1mNnQV/I0/QRgbUQ8DCDpGuBs4L6aMmcDV+Udvtsl9UiaERFPNdtpR9wwqrVixYqNkv7QYNU0IH2YxcBw3a7bdQ+M15TdcMWKFTeMGDFiWoNV4yTdWfN5aUQsrfk8E3i85vM64I/r9tGozEyge4JnRDS8YyTpznbZ7Kriul236977IuL0kps2GjNcP0a7SJnddM5DU2Zm1VgHzKr5fDB73ncpUmY3Dp5mNtTdARwp6TBJY4BzgWvrylwLvD+/634i8EKr653QgaftLSxtX8R1u27X3aV1VyYieiUtBm4ARgJXRsQaSYvy9UvIcrqfAawFXgI+0G6/Sp2ezczMfNpuZlaKg6eZWQkdFzyrGEZVsN5Zkm6WdL+kNZI+0qDMaZJekLQyf31yIOrO9/2opHvz/d7ZYH1Vx310zfGslPSipI/WlRmw45Z0paQNklbXLJsi6ReSfp//23Ba9na/jZJ1f17S7/K/6U8k9TTZtuX3U7LuSyU9UfN3PaPJtlUc9w9r6n1U0som2/bruIe0iOiYF9nF3IeAw4ExwD3AnLoyZwDXkT2XdSLw7wNU9wxgXv5+EtlkzfV1nwb8rKJjfxSY1mJ9Jcfd4O//NPCaqo4bOAWYB6yuWfY54JL8/SXA5WV+GyXrficwKn9/eaO6i3w/Jeu+FPhYge9kwI+7bv0XgU9WcdxD+dVpPc9dw6giYjvQN4yq1q5hVBFxO9AjKT2RUZ2IeCoi7srf/wdwP9kIg05RyXHXeTvwUEQ0GuE1ICLiFmBT3eKzgWX5+2XAuxpsWuS3kVx3RNwYEb35x9vJnu8bcE2Ou4hKjruPJJENn766RNuGtU4Lns2GSKWW6RdJhwJvAP69weo3SbpH0nWSXjeA1QZwo6QVyuY3rVf5cZM9/9bsP6KqjhtgeuTP1OX/HtCgzGAc/1+R9e4baff9lLU4v2RwZZPLFVUf91uA9RHx+ybrqzrurtdpwbOSYVRJDZAmAj8CPhoRL9atvovslPY44GvATweqXuDNETGPbHaXiySdUt+0BtsM5HGPAc4C/rnB6iqPu6iqj/8TQC/w/SZF2n0/ZXwTOAKYSzaG+ouNmtZg2UA+X3gerXudVRz3kNBpwbOSYVRFSRpNFji/HxE/rl8fES9GxOb8/XJgtKRGExUki2z2KSJiA1nSvBPqilR23LkFwF0Rsb5B2yo77tz6vksQ+b8bGpSp8nu/APgz4L2RX+irV+D7SRYR6yNiR0TsBP53k31WedyjgD8HftiijQN+3ENFpwXPSoZRFZFf+/k2cH9EfKlJmQPzckg6gezv9+wA1D1B0qS+92Q3MVbXFavkuGs07YFUddw1rgUuyN9fAPxLgzJFfhvJJJ0OfBw4KyJealKmyPdTpu7aa9bnNNlnJcedewfwu4homOC9quMeMvb2Hav6F9ld5QfJ7jB+Il+2CFiUvxfZxKYPAfcC8weo3pPJTodWASvz1xl1dS8mSy1yD9nNhZMGqO7D833ek+9/0I473/c+ZMFwv5pllRw3WYB+CniFrFd1ITAVuAn4ff7vlLzsQcDyVr+NAah7Ldk1xb7vfEl93c2+nwGo+7v5d7mKLCDOGKzjzpd/p+87rik7oMc9lF8enmlmVkKnnbabmXUFB08zsxIcPM3MSnDwNDMrwcHTzKwEB89hQtLUmll0nq6ZzWezpG9UVOdHJb2/in2Xkc8Q1PThfknXSDpyMNtk3cuPKg1Dki4FNkfEFyqsYxTZsM558erEG3uVpEfJno9tmF5X0qnAX0bEhwa1YdaV3PMc5pTN1fmz/P2lkpZJujHvpf25pM/l8zlenw9fRdLxkn6dTxZxQ5PZnd5GNtyzN9/mv0q6L58E45p82YR8Qow7JN0t6ex8+UhJX8jrXSXpw/nyt+fl7s23G5svf1TSpyXdla+bnS+fmh/L3ZL+F/k48bzenyub6GS1pPfkbb4VeEce+M1acvC0ekcAZ5JNe/Y94OaIOBbYCpyZB9CvAe+OiOOBK4G/b7CfNwMraj5fArwhIl5PNnoJ4BPALyPijcBbgc/nwwAXAofVlP++pHFkI2Lek7dnFPDXNfvfGNkEFt8EPpYv+xTwbxHxBrIRPIfky08HnoyI4yLiGOB6gMjGmK8Fjkv5g9nw5OBp9a6LiFfIhg2OJA8s+edDgaOBY4BfKJt9/O9oPAfmDOCZms+ryILgX5LNXgTZWOlL8v38ChhHFuDeQTZMshcgIjbl9T4SEQ/m2y4jm+S3T99ELivydpKv/16+j58Dz9UcyzskXS7pLRHxQs1+NpANUTRryacnVm8bZL0wSa/EqxfFd5L9XgSsiYg3tdnPVrJg2OdMsmB2FvDflc0JKuAvIuKB2g3zSUiKTEW4R7uBHez+u97jon5EPCjpeLIx45dJujEi/ke+elzedrOW3PO0VA8A+0t6E2TT+Knx5Mj3A3+UlxkBzIqIm4GLgR5gIlke7Q/XzNj0hnzbG4FFfdceJU0BfgccKumP8jLvA37dpq23AO/N97EAmJy/Pwh4KSK+B3yBLEVFn6PIJsEwa8nB05JElgri3cDlku4hm4nopAZFr+PV0+qRwPck3QvcDXw5Ip4HPgOMBlYpS072mbz8t4DH8uX3AOdHxMvAB4B/zvezE1jSprmfBk6RdBfZJYLH8uXHAr/NLxd8AvifAJKmA1tjYKf6syHKjypZZST9BLg4mqd46CiS/gZ4MSK+vbfbYp3PPU+r0iVkN466xfO8mojOrCX3PM3MSnDP08ysBAdPM7MSHDzNzEpw8DQzK8HB08yshP8PcZ3ognVszdgAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 360x180 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "example_idx = np.random.randint(0, shingles.shape[0])\n",
    "example_item = shingles[example_idx, :]\n",
    "\n",
    "ipd.display(ipd.HTML(metadata[metadata['Index'] == index[example_idx]].to_html()))\n",
    "\n",
    "plt.figure(figsize=(5, 2.5))\n",
    "plt.imshow(example_item.reshape(12, 20), aspect='auto', origin='lower', cmap='gray_r', vmin=0, vmax=1)\n",
    "plt.yticks([0, 2, 4, 5, 7, 9, 11], 'CDEFGAB')\n",
    "plt.ylabel('Chroma')\n",
    "plt.xlabel('Time (seconds)')\n",
    "plt.colorbar()\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dimensionality Reduction\n",
    "\n",
    "If required, we perform dimensionality reduction to the shingle using PCA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "if dimensionality_reduction:\n",
    "    pca = PCA(n_components=K)\n",
    "    pca.fit(shingles)\n",
    "    items = pca.transform(shingles)\n",
    "else:\n",
    "    items = shingles\n",
    "    K = shingles.shape[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Query\n",
    "\n",
    "In this notebook, we use a query from the database to demonstrate the usage of the index structures.\n",
    "You may also use another query by computing the features for an audio file.\n",
    "To do this, we provide the functions `common.compute_features` and `common.generate_shingles`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Index</th>\n",
       "      <th>ComposerID</th>\n",
       "      <th>WorkID</th>\n",
       "      <th>VersionID</th>\n",
       "      <th>Filename</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>46</th>\n",
       "      <td>46</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>Chopin_Op007No1_Kushner-1990_pid9118-05.npz</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "query_idx = np.random.randint(0, items.shape[0])\n",
    "query_item = items[query_idx, :]\n",
    "\n",
    "print('Query')\n",
    "ipd.display(ipd.HTML(metadata[metadata['Index'] == index[query_idx]].to_html()))\n",
    "\n",
    "# to use another query, you can compute a query item from an audio file\n",
    "# import librosa\n",
    "# fn_audio = '...'  # please provide the path to an audio file\n",
    "# X_cens = compute_features(fn_audio)\n",
    "# query_shingle = generate_shingles(X_cens)[0, :]  # using the first shingle from the file\n",
    "# query_item = pca.transform(query_shingle.reshape(1, -1))[0] if dimensionality_reduction else query_shingle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exhaustive Nearest Neighbor Search\n",
    "\n",
    "Now, we search for the $\\nu$ nearest shingles in our dataset using an exhaustive search.\n",
    "Here, we compute a matrix of pairwise distances between the query and all database items and select the items corresponding to the $\\nu$ smallest distances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def exhaustive_search(database, queries, n_neighbors):\n",
    "    C = scipy.spatial.distance.cdist(database, queries, metric='euclidean')\n",
    "    neighbor_idx = np.argsort(C, axis=0)[:n_neighbors, :]\n",
    "    neighbor_dist = C[np.unravel_index(neighbor_idx, C.shape)]\n",
    "    return neighbor_idx, neighbor_dist\n",
    "    \n",
    "neighbor_idx, neighbor_dist = exhaustive_search(items, query_item.reshape(1, -1), nu)\n",
    "neighbor_dist = np.square(neighbor_dist)  # show squared Euclidean distance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We next show the results of the search using the function `show_result_table`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Results of exhaustive search\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>Match</th>\n",
       "      <th>ComposerID</th>\n",
       "      <th>WorkID</th>\n",
       "      <th>VersionID</th>\n",
       "      <th>ShingleIdx</th>\n",
       "      <th>Distance</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>01</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15251</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>02</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15226</td>\n",
       "      <td>0.342614</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>03</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15252</td>\n",
       "      <td>0.393929</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>04</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15250</td>\n",
       "      <td>0.408689</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>05</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15227</td>\n",
       "      <td>0.488261</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>06</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Ohlsson-1999_pid9153-05</td>\n",
       "      <td>16364</td>\n",
       "      <td>0.587715</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>07</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Czerny-Stefanska-1949_pid9086-05</td>\n",
       "      <td>13113</td>\n",
       "      <td>0.635510</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>08</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Czerny-Stefanska-1949_pid9086-05</td>\n",
       "      <td>13141</td>\n",
       "      <td>0.635921</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>09</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Uninsky-1959_pid9061-05</td>\n",
       "      <td>17532</td>\n",
       "      <td>0.656634</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Uninsky-1959_pid9061-05</td>\n",
       "      <td>17557</td>\n",
       "      <td>0.665450</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def show_result_table(neighbor_idx, neighbor_dist, index=index, metadata=metadata):\n",
    "    neighbor_idx = neighbor_idx.ravel()\n",
    "    neighbor_dist = neighbor_dist.ravel()\n",
    "    \n",
    "    table = []\n",
    "    for i, (cur_item_idx, cur_dist) in enumerate(zip(neighbor_idx, neighbor_dist)):\n",
    "        cur_version_idx = index[cur_item_idx]\n",
    "        cur_meta_row = metadata[metadata['Index'] == cur_version_idx].iloc[0]\n",
    "        \n",
    "        table.append([f'{i + 1:02d}', cur_meta_row['ComposerID'], cur_meta_row['WorkID'],\n",
    "                      cur_meta_row['VersionID'], cur_item_idx, cur_dist])\n",
    "    \n",
    "    cols = ['Match', 'ComposerID', 'WorkID', 'VersionID', 'ShingleIdx', 'Distance']\n",
    "    ipd.display(ipd.HTML(pd.DataFrame(table, columns=cols).to_html(index=False)))\n",
    "    \n",
    "print('Results of exhaustive search')\n",
    "show_result_table(neighbor_idx, neighbor_dist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Nearest Neighbor Search with $K$-d tree\n",
    "\n",
    "Next, we construct a $K$-d tree for our dataset and search for the $\\nu$ nearest shingles using this tree.\n",
    "We then show the results as a table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Results of search with K-d tree\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>Match</th>\n",
       "      <th>ComposerID</th>\n",
       "      <th>WorkID</th>\n",
       "      <th>VersionID</th>\n",
       "      <th>ShingleIdx</th>\n",
       "      <th>Distance</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>01</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15251</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>02</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15226</td>\n",
       "      <td>0.342614</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>03</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15252</td>\n",
       "      <td>0.393929</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>04</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15250</td>\n",
       "      <td>0.408689</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>05</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15227</td>\n",
       "      <td>0.488261</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>06</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Ohlsson-1999_pid9153-05</td>\n",
       "      <td>16364</td>\n",
       "      <td>0.587715</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>07</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Czerny-Stefanska-1949_pid9086-05</td>\n",
       "      <td>13113</td>\n",
       "      <td>0.635510</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>08</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Czerny-Stefanska-1949_pid9086-05</td>\n",
       "      <td>13141</td>\n",
       "      <td>0.635921</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>09</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Uninsky-1959_pid9061-05</td>\n",
       "      <td>17532</td>\n",
       "      <td>0.656634</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Uninsky-1959_pid9061-05</td>\n",
       "      <td>17557</td>\n",
       "      <td>0.665450</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "neighbors = NearestNeighbors(n_neighbors=nu, metric='euclidean')\n",
    "neighbors.fit(items)\n",
    "\n",
    "assert isinstance(neighbors._tree, KDTree)\n",
    "\n",
    "neighbor_dist, neighbor_idx = neighbors.kneighbors(query_item.reshape(1, -1), nu)\n",
    "neighbor_dist = np.square(neighbor_dist)  # show squared Euclidean distance\n",
    "\n",
    "print('Results of search with K-d tree')\n",
    "show_result_table(neighbor_idx, neighbor_dist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Nearest Neighbor Search with HNSW Graph\n",
    "\n",
    "Finally, we construct an HNSW graph for our dataset and search for the $\\nu$ nearest shingles using this graph.\n",
    "We then show the results as a table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Results of search with HNSW graph\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>Match</th>\n",
       "      <th>ComposerID</th>\n",
       "      <th>WorkID</th>\n",
       "      <th>VersionID</th>\n",
       "      <th>ShingleIdx</th>\n",
       "      <th>Distance</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>01</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15251</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>02</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15226</td>\n",
       "      <td>0.342614</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>03</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15252</td>\n",
       "      <td>0.393929</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>04</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15250</td>\n",
       "      <td>0.408689</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>05</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Kushner-1990_pid9118-05</td>\n",
       "      <td>15227</td>\n",
       "      <td>0.488261</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>06</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Ohlsson-1999_pid9153-05</td>\n",
       "      <td>16364</td>\n",
       "      <td>0.587715</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>07</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Czerny-Stefanska-1949_pid9086-05</td>\n",
       "      <td>13113</td>\n",
       "      <td>0.635511</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>08</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Czerny-Stefanska-1949_pid9086-05</td>\n",
       "      <td>13141</td>\n",
       "      <td>0.635921</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>09</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Uninsky-1959_pid9061-05</td>\n",
       "      <td>17532</td>\n",
       "      <td>0.656634</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>Chopin</td>\n",
       "      <td>Op007No1</td>\n",
       "      <td>Uninsky-1959_pid9061-05</td>\n",
       "      <td>17557</td>\n",
       "      <td>0.665450</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "graph = hnswlib.Index(space='l2', dim=K)\n",
    "graph.init_index(max_elements=items.shape[0], ef_construction=M_prime, M=M)\n",
    "graph.set_ef(nu_prime)\n",
    "graph.add_items(items, np.arange(items.shape[0]))\n",
    "\n",
    "neighbor_idx, neighbor_dist = graph.knn_query(query_item.reshape(1, -1), k=nu)\n",
    "\n",
    "print('Results of search with HNSW graph')\n",
    "show_result_table(neighbor_idx, neighbor_dist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrieval Runtime\n",
    "\n",
    "The aim of using index structures in our music application is to increase the efficiency of the retrieval runtime.\n",
    "In the next code cell, we compare the runtimes of the different search strategies.\n",
    "First, we define functions with a consistent interface for the exhaustive search, the $K$-d tree search, and the HNSW-based search.\n",
    "Then we measure the runtimes of these functions using the Python module `timeit`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Full Search: 4.942 ms +/- 1.0 ms\n",
      "KD:          0.797 ms +/- 0.5 ms\n",
      "HNSW:        0.023 ms +/- 0.0 ms\n"
     ]
    }
   ],
   "source": [
    "def kd(neighbors, queries, n_neighbors):\n",
    "    neighbor_dist, neighbor_idx = neighbors.kneighbors(queries, nu)\n",
    "    return neighbor_idx, neighbor_dist\n",
    "\n",
    "def hnsw(graph, queries, n_neighbors):\n",
    "    neighbor_idx, neighbor_dist = graph.knn_query(queries, k=nu)\n",
    "    return neighbor_idx, neighbor_dist\n",
    "    \n",
    "# measure runtimes\n",
    "times_fullsearch = timeit.repeat('exhaustive_search(items, query_item.reshape(1, -1), nu)',\n",
    "                                 globals=globals(), number=1, repeat=100)\n",
    "times_kd = timeit.repeat('kd(neighbors, query_item.reshape(1, -1), nu)',\n",
    "                         globals=globals(), number=1, repeat=100)\n",
    "times_hnsw = timeit.repeat('hnsw(graph, query_item.reshape(1, -1), nu)',\n",
    "                           globals=globals(), number=1, repeat=100)\n",
    "\n",
    "# convert from second to milliseconds\n",
    "times_fullsearch = np.array(times_fullsearch) * 1000\n",
    "times_kd = np.array(times_kd) * 1000\n",
    "times_hnsw = np.array(times_hnsw) * 1000\n",
    "\n",
    "# report results\n",
    "print(f'Full Search: {np.mean(times_fullsearch):.3f} ms +/- {np.std(times_fullsearch):.1f} ms')\n",
    "print(f'KD:          {np.mean(times_kd):.3f} ms +/- {np.std(times_kd):.1f} ms')\n",
    "print(f'HNSW:        {np.mean(times_hnsw):.3f} ms +/- {np.std(times_hnsw):.1f} ms')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert\" style=\"background-color:#F5F5F5; border-color:#C8C8C8\">\n",
    "    <strong>Acknowledgment:</strong> This notebook was created by <a href=\"https://www.audiolabs-erlangen.de/fau/assistant/zalkow/\">Frank Zalkow</a>.\n",
    "</div>"
   ]
  }
 ],
 "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}