python/tests/ECIR_Tuturial_Image_Search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "34k2vOrCPlDD"
   },
   "source": [
    "# SETUP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 8256,
     "status": "ok",
     "timestamp": 1743711898665,
     "user": {
      "displayName": "Nico Hezel",
      "userId": "06522085364649969825"
     },
     "user_tz": -120
    },
    "id": "GspBXJAxX2Ne",
    "outputId": "208a63e4-ab80-4d51-e899-163e029db29c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting faiss-cpu\n",
      "  Downloading faiss_cpu-1.10.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (4.4 kB)\n",
      "Collecting deglib\n",
      "  Downloading deglib-0.1.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (7.3 kB)\n",
      "Requirement already satisfied: numpy<3.0,>=1.25.0 in /usr/local/lib/python3.11/dist-packages (from faiss-cpu) (2.0.2)\n",
      "Requirement already satisfied: packaging in /usr/local/lib/python3.11/dist-packages (from faiss-cpu) (24.2)\n",
      "Requirement already satisfied: psutil in /usr/local/lib/python3.11/dist-packages (from deglib) (5.9.5)\n",
      "Downloading faiss_cpu-1.10.0-cp311-cp311-manylinux_2_28_x86_64.whl (30.7 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m30.7/30.7 MB\u001b[0m \u001b[31m18.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading deglib-0.1.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (551 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m551.9/551.9 kB\u001b[0m \u001b[31m21.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: faiss-cpu, deglib\n",
      "Successfully installed deglib-0.1.2 faiss-cpu-1.10.0\n"
     ]
    }
   ],
   "source": [
    "!pip install faiss-cpu deglib\n",
    "# !pip install --no-binary=deglib --no-deps --force-reinstall deglib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4972,
     "status": "ok",
     "timestamp": 1743711903648,
     "user": {
      "displayName": "Nico Hezel",
      "userId": "06522085364649969825"
     },
     "user_tz": -120
    },
    "id": "PUPkbiN7nrUn",
    "outputId": "15b9dbfd-ad07-4cd1-a9ab-d3e7e8761abc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "100 18.0M  100 18.0M    0     0  6187k      0  0:00:02  0:00:02 --:--:-- 6186k\n",
      "Archive:  sift_subset_100k.zip\n",
      "   creating: sift_subset_100k/\n",
      "  inflating: sift_subset_100k/sift_query.fvecs  \n",
      "  inflating: sift_subset_100k/sift_groundtruth.ivecs  \n",
      "  inflating: sift_subset_100k/sift_base.fvecs  \n",
      "sift_base.fvecs  sift_groundtruth.ivecs  sift_query.fvecs\n"
     ]
    }
   ],
   "source": [
    "!curl -L -o sift_subset_100k.zip https://static.visual-computing.com/paper/DEG/sift_subset_100k.zip\n",
    "!unzip sift_subset_100k.zip\n",
    "!ls sift_subset_100k"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_7OqMcJJ493N"
   },
   "source": [
    "# Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "executionInfo": {
     "elapsed": 162,
     "status": "ok",
     "timestamp": 1743711903811,
     "user": {
      "displayName": "Nico Hezel",
      "userId": "06522085364649969825"
     },
     "user_tz": -120
    },
    "id": "C77_SybePhkq"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import time\n",
    "import faiss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 96,
     "status": "ok",
     "timestamp": 1743711903911,
     "user": {
      "displayName": "Nico Hezel",
      "userId": "06522085364649969825"
     },
     "user_tz": -120
    },
    "id": "vebB5MJFXV0h",
    "outputId": "8437e722-dd6e-40a5-f49b-b2e886d31b0f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Base shape: (100000, 128), Query shape: (10000, 128)\n"
     ]
    }
   ],
   "source": [
    "# Read .fvecs and .ivecs files\n",
    "def read_fvecs(filename):\n",
    "    with open(filename, 'rb') as f:\n",
    "        dim = np.fromfile(f, dtype=np.int32, count=1)[0]\n",
    "    return np.fromfile(filename, dtype=np.float32).reshape(-1, dim + 1)[:, 1:]\n",
    "\n",
    "def read_ivecs(filename):\n",
    "    with open(filename, 'rb') as f:\n",
    "        dim = np.fromfile(f, dtype=np.int32, count=1)[0]\n",
    "    return np.fromfile(filename, dtype=np.int32).reshape(-1, dim + 1)[:, 1:]\n",
    "\n",
    "# Load data\n",
    "sift_base = read_fvecs(\"./sift_subset_100k/sift_base.fvecs\")              # base vectors\n",
    "sift_query = read_fvecs(\"./sift_subset_100k/sift_query.fvecs\")            # queries\n",
    "sift_query_gt = read_ivecs(\"./sift_subset_100k/sift_groundtruth.ivecs\")   # ground truth indices\n",
    "\n",
    "print(f\"Base shape: {sift_base.shape}, Query shape: {sift_query.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9jwEErfOPhki"
   },
   "source": [
    "## Nearest Neighbor Search with FAISS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 17496,
     "status": "ok",
     "timestamp": 1743711921406,
     "user": {
      "displayName": "Nico Hezel",
      "userId": "06522085364649969825"
     },
     "user_tz": -120
    },
    "id": "fK_vnK5mXeSM",
    "outputId": "2ce98390-1de2-4b66-8b9a-0f23b647bf39"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FAISS index creation took 0.00 seconds.\n",
      "FAISS adding data took 0.11 seconds.\n",
      "FAISS kNNS (1 thread) took 17.32 seconds.\n",
      "Nearest neighbor indices shape: (10000, 100)\n"
     ]
    }
   ],
   "source": [
    "# Restrict to 1 thread\n",
    "faiss.omp_set_num_threads(1)\n",
    "\n",
    "# Get dimensionality\n",
    "d = sift_base.shape[1]\n",
    "\n",
    "# Create a FAISS index for L2 distance\n",
    "start = time.time()\n",
    "index = faiss.IndexFlatL2(d)\n",
    "end = time.time()\n",
    "print(f\"FAISS index creation took {end - start:.2f} seconds.\")\n",
    "\n",
    "# Add the base vectors to the index\n",
    "start = time.time()\n",
    "index.add(sift_base)\n",
    "end = time.time()\n",
    "print(f\"FAISS adding data took {end - start:.2f} seconds.\")\n",
    "\n",
    "# Search\n",
    "k = 100  # number of nearest neighbors\n",
    "start = time.time()\n",
    "distances, indices = index.search(sift_query, k)\n",
    "end = time.time()\n",
    "print(f\"FAISS kNNS (1 thread) took {end - start:.2f} seconds.\")\n",
    "print(f\"Nearest neighbor indices shape: {indices.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Gryhvw3VYiMi"
   },
   "source": [
    "# Approximated Nearest Neighbor Search (4 approaches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "executionInfo": {
     "elapsed": 18,
     "status": "ok",
     "timestamp": 1743711921407,
     "user": {
      "displayName": "Nico Hezel",
      "userId": "06522085364649969825"
     },
     "user_tz": -120
    },
    "id": "pm48na4-aeoF"
   },
   "outputs": [],
   "source": [
    "def recall_at_k(retrieved_indices, ground_truth, k=100):\n",
    "    \"\"\"\n",
    "    Computes Recall@k for nearest neighbor search.\n",
    "\n",
    "    Parameters:\n",
    "        retrieved_indices: np.ndarray of shape (num_queries, k)\n",
    "        ground_truth: np.ndarray of shape (num_queries, 1) – true nearest neighbor indices\n",
    "        k: int – number of top predictions to check\n",
    "\n",
    "    Returns:\n",
    "        recall: float – fraction of queries where ground truth was in top-k\n",
    "    \"\"\"\n",
    "    correct = 0\n",
    "    for i in range(ground_truth.shape[0]):\n",
    "        if ground_truth[i, 0] in retrieved_indices[i, :k]:\n",
    "            correct += 1\n",
    "    return correct / ground_truth.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BPnqYwWzagkY"
   },
   "source": [
    "## ANNS with FAISS Inverted File (IVF) Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 8180,
     "status": "ok",
     "timestamp": 1743711929571,
     "user": {
      "displayName": "Nico Hezel",
      "userId": "06522085364649969825"
     },
     "user_tz": -120
    },
    "id": "orj-IQKlZMLZ",
    "outputId": "3017f720-25a3-43cf-d8a5-c1fe5794e03a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[FAISS-IVF] nlist= 100, nprobe=  1, k=100 | Recall@100:  56.84% | QPS: 7779.44 | Create:  0.00s | Train:  0.40s | Add:  0.23s | Search:   1.29s\n",
      "[FAISS-IVF] nlist= 100, nprobe=  2, k=100 | Recall@100:  76.11% | QPS: 6079.76 | Create:  0.00s | Train:  0.31s | Add:  0.16s | Search:   1.64s\n",
      "[FAISS-IVF] nlist= 100, nprobe=  5, k=100 | Recall@100:  93.43% | QPS: 2817.35 | Create:  0.00s | Train:  0.25s | Add:  0.15s | Search:   3.55s\n"
     ]
    }
   ],
   "source": [
    "def run_faiss_ivf_ann(sift_base, sift_query, sift_query_gt, nlist=100, nprobe=10, k=100):\n",
    "    \"\"\"Run FAISS IVFFlat ANN search with timing and recall reporting.\"\"\"\n",
    "\n",
    "    # Restrict to 1 thread\n",
    "    faiss.omp_set_num_threads(1)\n",
    "\n",
    "    # Dimensionality\n",
    "    d = sift_base.shape[1]\n",
    "\n",
    "    # Index creation\n",
    "    t0 = time.time()\n",
    "    quantizer = faiss.IndexFlatL2(d)\n",
    "    index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)\n",
    "    t1 = time.time()\n",
    "    creation_time = t1 - t0\n",
    "\n",
    "    # Training\n",
    "    t0 = time.time()\n",
    "    index.train(sift_base)\n",
    "    t1 = time.time()\n",
    "    training_time = t1 - t0\n",
    "\n",
    "    # Add vectors\n",
    "    t0 = time.time()\n",
    "    index.add(sift_base)\n",
    "    t1 = time.time()\n",
    "    add_time = t1 - t0\n",
    "\n",
    "    # Search\n",
    "    index.nprobe = nprobe\n",
    "    t0 = time.time()\n",
    "    distances, indices = index.search(sift_query, k)\n",
    "    t1 = time.time()\n",
    "    search_time = t1 - t0\n",
    "    qps = len(sift_query) / search_time\n",
    "\n",
    "    # Recall\n",
    "    recall = recall_at_k(indices, sift_query_gt, k)\n",
    "\n",
    "    # Pretty print\n",
    "    print(\n",
    "        f\"[FAISS-IVF] nlist={nlist:4d}, nprobe={nprobe:3d}, k={k:3d} | \"\n",
    "        f\"Recall@{k:<3d}: {recall * 100:6.2f}% | \"\n",
    "        f\"QPS: {qps:7.2f} | \"\n",
    "        f\"Create: {creation_time:5.2f}s | \"\n",
    "        f\"Train: {training_time:5.2f}s | \"\n",
    "        f\"Add: {add_time:5.2f}s | \"\n",
    "        f\"Search: {search_time:6.2f}s\"\n",
    "    )\n",
    "\n",
    "    return recall, qps\n",
    "\n",
    "for nprobe in [1, 2, 5]:\n",
    "    run_faiss_ivf_ann(sift_base, sift_query, sift_query_gt, nlist=100, nprobe=nprobe, k=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sKoy42mKa37J"
   },
   "source": [
    "## ANNS with FAISS Inverted File with Product Quantization (IVFPQ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 24317,
     "status": "ok",
     "timestamp": 1743711953875,
     "user": {
      "displayName": "Nico Hezel",
      "userId": "06522085364649969825"
     },
     "user_tz": -120
    },
    "id": "HPfIYk2lbE6c",
    "outputId": "8115d4ac-29b9-40f9-c55f-110a103842e4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[FAISS-IVFPQ] nlist= 100, nprobe=  1, k=100, m=16, nbits= 8 | Recall@100:  56.84% | QPS: 11101.75 | Create:  0.00s | Train:  5.44s | Add:  1.17s | Search:   0.90s\n",
      "[FAISS-IVFPQ] nlist= 100, nprobe=  2, k=100, m=16, nbits= 8 | Recall@100:  76.11% | QPS: 7969.58 | Create:  0.00s | Train:  5.26s | Add:  0.81s | Search:   1.25s\n",
      "[FAISS-IVFPQ] nlist= 100, nprobe=  5, k=100, m=16, nbits= 8 | Recall@100:  93.42% | QPS: 4043.35 | Create:  0.00s | Train:  6.05s | Add:  0.79s | Search:   2.47s\n"
     ]
    }
   ],
   "source": [
    "def run_faiss_ivfpq_ann(sift_base, sift_query, sift_query_gt, nlist=100, nprobe=10, k=100, m=16, nbits=8):\n",
    "    \"\"\"\n",
    "    Run FAISS IVFPQ ANN search.\n",
    "    Parameters:\n",
    "        - sift_base: database vectors (N, d)\n",
    "        - sift_query: query vectors (Q, d)\n",
    "        - sift_query_gt: ground truth indices (Q, 1)\n",
    "        - nlist: number of coarse clusters\n",
    "        - nprobe: number of clusters to search\n",
    "        - k: top-k neighbors to retrieve\n",
    "        - m: number of subquantizers (must divide d)\n",
    "        - nbits: number of bits per subquantizer\n",
    "    Returns:\n",
    "        - recall@k\n",
    "        - queries per second\n",
    "    \"\"\"\n",
    "    faiss.omp_set_num_threads(1)\n",
    "    d = sift_base.shape[1]\n",
    "\n",
    "    # Index creation\n",
    "    t0 = time.time()\n",
    "    quantizer = faiss.IndexFlatL2(d)\n",
    "    index = faiss.IndexIVFPQ(quantizer, d, nlist, m, nbits)\n",
    "    t1 = time.time()\n",
    "    creation_time = t1 - t0\n",
    "\n",
    "    # Training\n",
    "    t0 = time.time()\n",
    "    index.train(sift_base)\n",
    "    t1 = time.time()\n",
    "    training_time = t1 - t0\n",
    "\n",
    "    # Add data\n",
    "    t0 = time.time()\n",
    "    index.add(sift_base)\n",
    "    t1 = time.time()\n",
    "    add_time = t1 - t0\n",
    "\n",
    "    # Search\n",
    "    index.nprobe = nprobe\n",
    "    t0 = time.time()\n",
    "    distances, indices = index.search(sift_query, k)\n",
    "    t1 = time.time()\n",
    "    search_time = t1 - t0\n",
    "    qps = len(sift_query) / search_time\n",
    "\n",
    "    # Recall\n",
    "    recall = recall_at_k(indices, sift_query_gt, k)\n",
    "\n",
    "    # Pretty print\n",
    "    print(\n",
    "        f\"[FAISS-IVFPQ] nlist={nlist:4d}, nprobe={nprobe:3d}, k={k:3d}, m={m:2d}, nbits={nbits:2d} | \"\n",
    "        f\"Recall@{k:<3d}: {recall * 100:6.2f}% | \"\n",
    "        f\"QPS: {qps:7.2f} | \"\n",
    "        f\"Create: {creation_time:5.2f}s | \"\n",
    "        f\"Train: {training_time:5.2f}s | \"\n",
    "        f\"Add: {add_time:5.2f}s | \"\n",
    "        f\"Search: {search_time:6.2f}s\"\n",
    "    )\n",
    "\n",
    "    return recall, qps\n",
    "\n",
    "for nprobe in [1, 2, 5]:\n",
    "    run_faiss_ivfpq_ann(sift_base, sift_query, sift_query_gt, nlist=100, nprobe=nprobe, k=100, m=16, nbits=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "if1wDpTqgKcT"
   },
   "source": [
    "## ANNS with FAISS Hierarchical Navigable Small World (HNSW)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 28687,
     "status": "ok",
     "timestamp": 1743711982564,
     "user": {
      "displayName": "Nico Hezel",
      "userId": "06522085364649969825"
     },
     "user_tz": -120
    },
    "id": "L0xggFFNgQRA",
    "outputId": "6dbe7e0c-3574-48ec-a7d5-d012350f36ba"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[FAISS-HNSW] m= 32, efC=  16, efS=   8, k=100 | Recall@100:  80.38% | QPS: 11128.39 | Create:  0.00s | Add:  8.81s | Search:   0.90s\n",
      "[FAISS-HNSW] m= 32, efC=  16, efS=  16, k=100 | Recall@100:  89.56% | QPS: 6398.19 | Create:  0.00s | Add:  8.81s | Search:   1.56s\n",
      "[FAISS-HNSW] m= 32, efC=  16, efS=  32, k=100 | Recall@100:  95.59% | QPS: 4882.09 | Create:  0.00s | Add:  8.81s | Search:   2.05s\n",
      "[FAISS-HNSW] m= 32, efC=  16, efS=  64, k=100 | Recall@100:  98.38% | QPS: 3748.09 | Create:  0.00s | Add:  8.81s | Search:   2.67s\n",
      "[FAISS-HNSW] m= 32, efC=  16, efS= 128, k=100 | Recall@100:  99.47% | QPS: 2024.98 | Create:  0.00s | Add:  8.81s | Search:   4.94s\n",
      "[FAISS-HNSW] m= 32, efC=  16, efS= 256, k=100 | Recall@100:  99.79% | QPS: 1346.19 | Create:  0.00s | Add:  8.81s | Search:   7.43s\n"
     ]
    }
   ],
   "source": [
    "def run_faiss_hnsw_ann(\n",
    "    sift_base,\n",
    "    sift_query,\n",
    "    sift_query_gt,\n",
    "    k=100,\n",
    "    hnsw_m=32,\n",
    "    ef_construction=200,\n",
    "    ef_search_list=[16, 32, 64, 128]\n",
    "):\n",
    "    \"\"\"\n",
    "    Run FAISS HNSW ANN search across multiple ef_search values.\n",
    "    Parameters:\n",
    "        - sift_base: (N, d) base vectors\n",
    "        - sift_query: (Q, d) query vectors\n",
    "        - sift_query_gt: (Q, 1) ground truth\n",
    "        - k: top-k neighbors to retrieve\n",
    "        - hnsw_m: number of edges per node\n",
    "        - ef_construction: controls indexing quality/speed\n",
    "        - ef_search_list: list of efSearch values to benchmark\n",
    "    Returns:\n",
    "        - recalls: list of recall@k values\n",
    "        - qps_list: list of queries/sec values\n",
    "    \"\"\"\n",
    "    faiss.omp_set_num_threads(1)\n",
    "    d = sift_base.shape[1]\n",
    "\n",
    "    # Create HNSW index\n",
    "    t0 = time.time()\n",
    "    index = faiss.IndexHNSWFlat(d, hnsw_m, faiss.METRIC_L2)\n",
    "    index.hnsw.efConstruction = ef_construction\n",
    "    t1 = time.time()\n",
    "    creation_time = t1 - t0\n",
    "\n",
    "    # Add base vectors\n",
    "    t0 = time.time()\n",
    "    index.add(sift_base)\n",
    "    t1 = time.time()\n",
    "    add_time = t1 - t0\n",
    "\n",
    "    recalls = []\n",
    "    qps_list = []\n",
    "\n",
    "    for ef in ef_search_list:\n",
    "        index.hnsw.efSearch = ef\n",
    "\n",
    "        # Search\n",
    "        t0 = time.time()\n",
    "        distances, indices = index.search(sift_query, k)\n",
    "        t1 = time.time()\n",
    "        search_time = t1 - t0\n",
    "        qps = len(sift_query) / search_time\n",
    "\n",
    "        # Recall\n",
    "        recall = recall_at_k(indices, sift_query_gt, k)\n",
    "        recalls.append(recall)\n",
    "        qps_list.append(qps)\n",
    "\n",
    "        # Pretty print\n",
    "        print(\n",
    "            f\"[FAISS-HNSW] m={hnsw_m:3d}, efC={ef_construction:4d}, efS={ef:4d}, k={k:3d} | \"\n",
    "            f\"Recall@{k:<3d}: {recall * 100:6.2f}% | \"\n",
    "            f\"QPS: {qps:7.2f} | \"\n",
    "            f\"Create: {creation_time:5.2f}s | \"\n",
    "            f\"Add: {add_time:5.2f}s | \"\n",
    "            f\"Search: {search_time:6.2f}s\"\n",
    "        )\n",
    "\n",
    "    return index, recalls, qps_list\n",
    "\n",
    "index, recalls, qps = run_faiss_hnsw_ann(\n",
    "    sift_base, sift_query, sift_query_gt,\n",
    "    k=100,\n",
    "    hnsw_m=32,\n",
    "    ef_construction=16,\n",
    "    ef_search_list=[8, 16, 32, 64, 128, 256]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 61480,
     "status": "ok",
     "timestamp": 1743712044052,
     "user": {
      "displayName": "Nico Hezel",
      "userId": "06522085364649969825"
     },
     "user_tz": -120
    },
    "id": "4ORpy9XboNFy",
    "outputId": "aa7a6a55-4033-4827-95a6-f88762b296fb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[FAISS-HNSW] m= 32, efC= 160, efS=   8, k=100 | Recall@100:  86.40% | QPS: 11496.43 | Create:  0.00s | Add: 49.93s | Search:   0.87s\n",
      "[FAISS-HNSW] m= 32, efC= 160, efS=  16, k=100 | Recall@100:  94.46% | QPS: 8577.57 | Create:  0.00s | Add: 49.93s | Search:   1.17s\n",
      "[FAISS-HNSW] m= 32, efC= 160, efS=  32, k=100 | Recall@100:  98.33% | QPS: 5852.80 | Create:  0.00s | Add: 49.93s | Search:   1.71s\n",
      "[FAISS-HNSW] m= 32, efC= 160, efS=  64, k=100 | Recall@100:  99.61% | QPS: 3078.60 | Create:  0.00s | Add: 49.93s | Search:   3.25s\n",
      "[FAISS-HNSW] m= 32, efC= 160, efS= 128, k=100 | Recall@100:  99.94% | QPS: 2317.30 | Create:  0.00s | Add: 49.93s | Search:   4.32s\n"
     ]
    }
   ],
   "source": [
    "index, recalls, qps = run_faiss_hnsw_ann(\n",
    "    sift_base, sift_query, sift_query_gt,\n",
    "    k=100,\n",
    "    hnsw_m=32,\n",
    "    ef_construction=160,\n",
    "    ef_search_list=[8, 16, 32, 64, 128]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1u4bt2B40iY8"
   },
   "source": [
    "## ANNS with the Dynamic Exploration Graph (DEG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 88377,
     "status": "ok",
     "timestamp": 1743712132421,
     "user": {
      "displayName": "Nico Hezel",
      "userId": "06522085364649969825"
     },
     "user_tz": -120
    },
    "id": "G2N5Qj1Fp8p0",
    "outputId": "3aa60418-74c5-461d-c4c7-8d6d7c2f4022"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100.00% [############################################################] (100000 / 100000)\n",
      "[DEG] d=16, k_ext=32, eps_ext=0.100, eps=0.000, k=100 | Recall@100:  98.94% | QPS:  430.70 | Build: 17.33s | Search:  23.22s\n",
      "[DEG] d=16, k_ext=32, eps_ext=0.100, eps=0.001, k=100 | Recall@100:  98.97% | QPS:  423.76 | Build: 17.33s | Search:  23.60s\n",
      "[DEG] d=16, k_ext=32, eps_ext=0.100, eps=0.010, k=100 | Recall@100:  99.16% | QPS:  417.41 | Build: 17.33s | Search:  23.96s\n"
     ]
    }
   ],
   "source": [
    "import deglib\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "def run_deg_ann(\n",
    "    base_vectors,\n",
    "    query_vectors,\n",
    "    ground_truth,\n",
    "    k=100,\n",
    "    edges_per_vertex=30,\n",
    "    extend_k=60,\n",
    "    extend_eps=0.1,\n",
    "    eps_list=[0.001, 0.01, 0.05, 0.1],\n",
    "):\n",
    "    \"\"\"\n",
    "    Run DEG ANN search across multiple eps values.\n",
    "\n",
    "    Parameters:\n",
    "        - base_vectors: (N, d) base vectors\n",
    "        - query_vectors: (Q, d) query vectors\n",
    "        - ground_truth: (Q, k) ground truth indices\n",
    "        - k: top-k neighbors to retrieve\n",
    "        - edges_per_vertex: number of edges per vertex in the graph\n",
    "        - extend_k: candidate pool size during graph build\n",
    "        - extend_eps: effort during graph build\n",
    "        - eps_list: list of eps values to benchmark\n",
    "    Returns:\n",
    "        - graph: the DEG search graph\n",
    "        - recalls: list of recall@k values\n",
    "        - qps_list: list of queries/sec values\n",
    "    \"\"\"\n",
    "\n",
    "    # Ensure C-contiguous arrays\n",
    "    base_vectors = np.ascontiguousarray(base_vectors, dtype=np.float32)\n",
    "    query_vectors = np.ascontiguousarray(query_vectors, dtype=np.float32)\n",
    "    ground_truth = np.ascontiguousarray(ground_truth, dtype=np.int32)\n",
    "\n",
    "    # Build graph\n",
    "    t0 = time.time()\n",
    "\n",
    "    from deglib.distances import Metric\n",
    "    from deglib.graph import SizeBoundedGraph\n",
    "    from deglib.builder import EvenRegularGraphBuilder, LID\n",
    "    from deglib.std import Mt19937\n",
    "\n",
    "    capacity = base_vectors.shape[0]\n",
    "    graph = SizeBoundedGraph.create_empty(capacity, base_vectors.shape[1], edges_per_vertex, Metric.L2)\n",
    "    builder = EvenRegularGraphBuilder(\n",
    "        graph, None, lid=LID.Low, extend_k=extend_k, extend_eps=extend_eps, improve_k=0,\n",
    "        improve_eps=0, max_path_length=0, swap_tries=0,\n",
    "        additional_swap_tries=0\n",
    "    )\n",
    "    labels = np.arange(base_vectors.shape[0], dtype=np.uint32)\n",
    "    builder.set_thread_count(1)\n",
    "    builder.set_batch_size(1)\n",
    "    builder.add_entry(labels, base_vectors)\n",
    "    builder.build(callback=\"progress\")\n",
    "\n",
    "    # graph = deglib.builder.build_from_data(\n",
    "    #     data=base_vectors,\n",
    "    #     edges_per_vertex=edges_per_vertex,\n",
    "    #     extend_k=extend_k,\n",
    "    #     extend_eps=extend_eps,\n",
    "    #     improve_k=0,\n",
    "    #     lid=deglib.builder.LID.Low,\n",
    "    #     callback=\"progress\"\n",
    "    # )\n",
    "    t1 = time.time()\n",
    "    build_time = t1 - t0\n",
    "\n",
    "    recalls = []\n",
    "    qps_list = []\n",
    "\n",
    "    for eps in eps_list:\n",
    "        # Vectorized search\n",
    "        t0 = time.time()\n",
    "        results, _ = graph.search(query=query_vectors, eps=eps, k=k, threads=1)\n",
    "        t1 = time.time()\n",
    "        search_time = t1 - t0\n",
    "        qps = len(query_vectors) / search_time\n",
    "\n",
    "        recall = recall_at_k(results, ground_truth, k)\n",
    "        recalls.append(recall)\n",
    "        qps_list.append(qps)\n",
    "\n",
    "        print(\n",
    "            f\"[DEG] d={edges_per_vertex}, k_ext={extend_k}, eps_ext={extend_eps:.3f}, \"\n",
    "            f\"eps={eps:.3f}, k={k:3d} | \"\n",
    "            f\"Recall@{k:<3d}: {recall * 100:6.2f}% | \"\n",
    "            f\"QPS: {qps:7.2f} | \"\n",
    "            f\"Build: {build_time:5.2f}s | \"\n",
    "            f\"Search: {search_time:6.2f}s\"\n",
    "        )\n",
    "\n",
    "    return graph, recalls, qps_list\n",
    "\n",
    "\n",
    "graph, recalls, qps = run_deg_ann(\n",
    "    sift_base, sift_query, sift_query_gt,\n",
    "    k=100,\n",
    "    edges_per_vertex=16,            # Graph degree (d)\n",
    "    extend_k=32,                    # Build: candidate pool size\n",
    "    extend_eps=0.1,                 # Build: search radius parameter\n",
    "    eps_list=[0.0001, 0.001, 0.01], # Try multiple search radius settings\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 25,
     "status": "ok",
     "timestamp": 1743712132423,
     "user": {
      "displayName": "Nico Hezel",
      "userId": "06522085364649969825"
     },
     "user_tz": -120
    },
    "id": "PDx2vMKp-mob",
    "outputId": "f016682e-d79e-4bd5-a32d-7a3244536743"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "import deglib_cpp\n",
    "print(deglib_cpp.avx_usable())\n",
    "print(deglib_cpp.sse_usable())"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.3"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}