#!/usr/bin/env python3 """Hands-on P02 — an ANN index, assembled from seven lego blocks.""" import heapq, math, time import numpy as np from _harness import block, run_all rng = np.random.default_rng(0) @block(1, "Vectors and the metric decision", "normalise once, and three metrics collapse into one") def b1(s, show): def norm(x): return x / np.linalg.norm(x, axis=-1, keepdims=True) n, d = 8000, 48 data = norm(rng.standard_normal((n, d), dtype=np.float32)) q = norm(rng.standard_normal((64, d), dtype=np.float32)) if show: dot = data @ q[0] l2 = np.linalg.norm(data - q[0], axis=1) ident = np.abs(l2**2 - (2 - 2*dot)).max() print(f" {n} vectors, d={d}, unit norm") print(f" ||a-b||^2 == 2-2: max deviation {ident:.2e}") print(f" argsort by -dot == argsort by L2: " f"{np.array_equal(np.argsort(-dot)[:20], np.argsort(l2)[:20])}") print(" so ONE distance function serves cosine, dot and L2 -- but only") print(" because we normalised. Skip it and rankings silently diverge.") return {"data": data, "q": q, "norm": norm, "n": n, "d": d} @block(2, "Relative contrast", "measure how hard your dataset is BEFORE benchmarking an index") def b2(s, show): data, q = s["data"], s["q"] def rc(data, q, k=10): dist = np.sqrt(np.maximum(0, 2 - 2 * (data @ q.T))) srt = np.sort(dist, axis=0) return float(np.mean(dist.mean(axis=0) / srt[0])) r = rc(data, q) if show: print(f" RC = mean distance / nearest distance = {r:.3f}") for dd in (8, 48, 256): sub = s["norm"](rng.standard_normal((3000, dd), dtype=np.float32)) sq = s["norm"](rng.standard_normal((32, dd), dtype=np.float32)) print(f" d={dd:>4}: RC={rc(sub, sq):.3f}") print(" RC falls toward 1 with dimension: every point equidistant, no") print(" gradient for greedy search. RC predicts difficulty; d does not.") return {"rc": r} @block(3, "Brute force = the oracle", "you cannot measure recall without exact ground truth") def b3(s, show): data, q = s["data"], s["q"] def brute(qv, k=10): sims = data @ qv idx = np.argpartition(-sims, k)[:k] return idx[np.argsort(-sims[idx])].tolist() t0 = time.perf_counter() truth = [brute(v) for v in q] ms = (time.perf_counter() - t0) / len(q) * 1e3 if show: naive = sorted(range(s["n"]), key=lambda i: -float(data[i] @ q[0]))[:10] print(f" agrees with a naive sort: {naive == truth[0]}") print(f" {ms:.3f} ms/query, {1000/ms:>6.0f} qps, recall 1.0 BY CONSTRUCTION") print(f" cost model: {s['n']}x{s['d']} = {s['n']*s['d']:,} MACs, " f"{s['n']*s['d']*4/1e6:.1f} MB streamed -- one BLAS call, prefetchable") return {"brute": brute, "truth": truth, "brute_ms": ms} @block(4, "recall@k", "the contract: what exactly did approximation cost?") def b4(s, show): def recall_at_k(got, want, k): return len(set(got[:k]) & set(want[:k])) / k if show: print(f" perfect: {recall_at_k([1,2,3],[1,2,3],3):.2f} " f"two of three: {recall_at_k([1,2,9],[1,2,3],3):.2f} " f"none: {recall_at_k([7,8,9],[1,2,3],3):.2f}") print(" always against EXACT ground truth, same metric, same data, stated k") return {"recall_at_k": recall_at_k} @block(5, "A random graph, and why it fails", "the naive design, measured, so the fix is motivated") def b5(s, show): data, n = s["data"], s["n"] deg = 16 g = [rng.choice(n, deg, replace=False).tolist() for _ in range(n)] def greedy(qv, graph, entry=0, ef=32): seen = {entry}; d0 = 1 - float(data[entry] @ qv) cand = [(d0, entry)]; res = [(-d0, entry)]; nd = 1 while cand: d, node = heapq.heappop(cand) if -res[0][0] < d and len(res) >= ef: break for nb in graph[node]: if nb in seen: continue seen.add(nb); dn = 1 - float(data[nb] @ qv); nd += 1 if len(res) < ef or dn < -res[0][0]: heapq.heappush(cand, (dn, nb)); heapq.heappush(res, (-dn, nb)) if len(res) > ef: heapq.heappop(res) return [i for _, i in heapq.nsmallest(10, [(-a, b) for a, b in res])], nd rec = np.mean([s["recall_at_k"](greedy(v, g)[0], t, 10) for v, t in zip(s["q"], s["truth"])]) if show: print(f" {deg} RANDOM edges per node, greedy beam ef=32") print(f" recall@10 = {rec:.4f} <- near useless, and that is the point") print(" a random graph has no locality, so greedy descent has nothing to") print(" descend. The fix is not a bigger beam; it is better EDGES.") return {"greedy": greedy, "random_recall": float(rec)} @block(6, "NSW: edges that mean something", "insert by searching what you have built so far") def b6(s, show): data, n = s["data"], s["n"] M, efC = 12, 60 graph = [[] for _ in range(n)] order = rng.permutation(n) entry = int(order[0]) for count, node in enumerate(order[1:], 1): node = int(node) found, _ = s["greedy"](data[node], graph, entry, efC) cands = sorted(found, key=lambda i: 1 - float(data[i] @ data[node]))[:M] graph[node] = cands for c in cands: graph[c].append(node) if len(graph[c]) > 2 * M: # degree cap dd = data[graph[c]] @ data[c] graph[c] = [graph[c][i] for i in np.argsort(-dd)[:2 * M]] if show: edges = sum(len(x) for x in graph) print(f" M={M} efConstruction={efC}, {edges/n:.1f} edges/node") print(" early insertions land in a near-empty graph, so their edges are") print(" LONG. Those accidental long links are the small-world property.") return {"graph": graph, "entry": entry, "M": M} @block(7, "The efSearch knob", "recall is concave in ef; pick the recall FIRST") def b7(s, show): if show: print(f" {'ef':>5}{'recall@10':>11}{'ms':>9}{'dists/q':>10}{'vs brute':>10}") for ef in (10, 32, 64, 128): t0 = time.perf_counter(); recs = []; nds = 0 for v, t in zip(s["q"], s["truth"]): got, nd = s["greedy"](v, s["graph"], s["entry"], ef) recs.append(s["recall_at_k"](got, t, 10)); nds += nd ms = (time.perf_counter() - t0) / len(s["q"]) * 1e3 print(f" {ef:>5}{np.mean(recs):>11.4f}{ms:>9.3f}" f"{nds/len(s['q']):>10.0f}{s['brute_ms']/ms:>9.2f}x") return {} def assembly(s): print("\nThe seven blocks are an ANN index. Now the measurement that matters.\n") ef = 64 t0 = time.perf_counter(); recs = []; nds = 0 for v, t in zip(s["q"], s["truth"]): got, nd = s["greedy"](v, s["graph"], s["entry"], ef) recs.append(s["recall_at_k"](got, t, 10)); nds += nd ms = (time.perf_counter() - t0) / len(s["q"]) * 1e3 dpq = nds / len(s["q"]) algo = s["n"] / dpq const = (ms * 1e6 / dpq) / (s["brute_ms"] * 1e6 / s["n"]) print(f" at efSearch={ef}: recall {np.mean(recs):.4f}, {ms:.3f} ms/query") print(f" random-graph recall was {s['random_recall']:.4f} -> NSW is " f"{np.mean(recs)/max(s['random_recall'],1e-9):.0f}x better on the SAME search code.\n") print(" THE TWO-FACTOR MODEL -- decompose before you tune anything:") print(f" algorithmic win : {s['n']:,} / {dpq:.0f} distances = {algo:>7.1f}x fewer") print(f" constant factor : {ms*1e6/dpq:>7.1f} ns/dist (python) vs " f"{s['brute_ms']*1e6/s['n']:.1f} ns (BLAS) = {const:>6.1f}x slower each") print(f" predicted speedup: {algo:.1f} / {const:.1f} = {algo/const:.2f}x") print(f" measured speedup : {s['brute_ms']/ms:.2f}x") print("\n The algorithm is RIGHT and the implementation is WRONG. Those are") print(" different bugs. Only the distance counter can tell you which you have.") print(f"\n Dataset difficulty: RC = {s['rc']:.3f}. Report it with every recall") print(" number, or the result does not transfer to anyone else's corpus.") print("\n Built: metric choice -> RC -> oracle -> recall -> random graph ->") print(" NSW -> ef sweep -> the decomposition.") print(" Missing, and on the project page: HNSW layers (m6), Algorithm 4 (m7),") print(" a compiled inner loop (m8), persistence (m9), hnswlib comparison (E12).") if __name__ == "__main__": run_all(assembly, "HANDS-ON P02 — ANN index, block by block")