Navigable small world graph in Python
A proximity graph in which greedy routing from any node reaches near neighbors of a query quickly.
- File
navigable_small_world.py, 77 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/navigable-small-world/. The authors are listed in prov/XDDYQ2.jsonl.
import heapq
import math
import random
from typing import Callable, Hashable, Sequence
class NSW:
"""Navigable small world graph (Malkov et al.): incremental insertion with links to the f nearest found."""
def __init__(self, d: Callable[[object, object], float] = math.dist, f: int = 3,
ef_construction: int = 10, entries: int = 1, seed: int = 0) -> None:
if f < 1 or entries < 1:
raise ValueError("f and entries must be at least 1")
self.d, self.f, self.ef_construction, self.entries = d, f, ef_construction, entries
self.rng = random.Random(seed)
self.points: list = []
self.adj: list[set[int]] = []
self.evaluations = 0
def _dist(self, i: int, q) -> float:
self.evaluations += 1
return self.d(self.points[i], q)
def _entry_points(self) -> list[int]:
n = len(self.points)
extra = [self.rng.randrange(n) for _ in range(self.entries - 1)]
return list(dict.fromkeys([0] + extra))
def search(self, q, k: int, ef: int, entries: Sequence[int] | None = None) -> list[tuple[float, int]]:
"""The k nearest (distance, index) pairs among the nodes the beam search evaluates, nearest first."""
if not self.points or k < 1:
return []
ef = max(ef, k)
starts = self._entry_points() if entries is None else list(dict.fromkeys(entries))
visited: set[int] = set()
cand: list[tuple[float, int]] = []
best: list[tuple[float, int]] = [] # @line L1
for e in starts: # @line L2
de = self._dist(e, q)
visited.add(e)
heapq.heappush(cand, (de, e)) # @line L3
heapq.heappush(best, (-de, -e))
if len(best) > ef:
heapq.heappop(best)
while cand: # @line L4
dc, c = heapq.heappop(cand) # @line L5
if len(best) == ef and (dc, c) > (-best[0][0], -best[0][1]): # @line L6
break
for e in sorted(self.adj[c]): # @line L7
if e in visited:
continue
visited.add(e)
de = self._dist(e, q) # @line L8
if len(best) < ef or (de, e) < (-best[0][0], -best[0][1]): # @line L9
heapq.heappush(cand, (de, e))
heapq.heappush(best, (-de, -e))
if len(best) > ef:
heapq.heappop(best)
self.last_visited = visited
return sorted((-nd, -ni) for nd, ni in best)[:k] # @line L10
def insert(self, p) -> int:
i = len(self.points)
if i == 0: # @line L11
self.points.append(p)
self.adj.append(set())
return 0
found = self.search(p, self.f, max(self.ef_construction, self.f)) # @line L12
self.points.append(p)
self.adj.append(set())
for _, j in found: # @line L13
self.adj[i].add(j)
self.adj[j].add(i)
return i
test_navigable_small_world.py 80 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/navigable-small-world/. The authors are listed in prov/XDDYQ2.jsonl.
import math
import random
from navigable_small_world import NSW
P = [(9, 5), (9, 1), (0, 4), (8, 9), (5, 4), (2, 5), (2, 2), (1, 6), (5, 2), (6, 9), (6, 7), (7, 3)]
NAMES = "abcdefghijkl"
def build(points, f=2, ef_c=2, entries=1, seed=0):
g = NSW(f=f, ef_construction=ef_c, entries=entries, seed=seed)
for p in points:
g.insert(p)
return g
def test_example() -> None:
g = build(P)
adj = {NAMES[i]: "".join(NAMES[j] for j in sorted(a)) for i, a in enumerate(g.adj)}
assert adj == {"a": "bcdej", "b": "acdei", "c": "abfgh", "d": "abjk", "e": "abfil", "f": "cegh",
"g": "cf", "h": "cf", "i": "bel", "j": "adk", "k": "dj", "l": "ei"}
q = (2, 9)
g.evaluations = 0
assert [i for _, i in g.search(q, 1, 1, [0])] == [9]
assert g.evaluations == 7
g.evaluations = 0
assert [i for _, i in g.search(q, 1, 3, [0])] == [7]
assert g.evaluations == 10
def test_edges() -> None:
g = NSW()
assert g.search((0, 0), 3, 5) == []
g.insert((1.0, 1.0))
assert g.search((0, 0), 3, 5) == [(math.dist((1, 1), (0, 0)), 0)]
assert g.search((0, 0), 0, 5) == []
try:
NSW(f=0)
raise AssertionError("f = 0 accepted")
except ValueError:
pass
def brute(points, q, k):
return sorted((math.dist(p, q), i) for i, p in enumerate(points))[:k]
def test_random_against_brute_force() -> None:
rng = random.Random(2014)
for _ in range(500):
n = rng.randint(1, 40)
dim = rng.randint(1, 5)
pts = [tuple(rng.uniform(0, 10) for _ in range(dim)) for _ in range(n)]
f = rng.randint(1, 4)
g = build(pts, f=f, ef_c=rng.randint(1, 6), entries=rng.randint(1, 3), seed=rng.randrange(10**6))
# invariants: undirected, links at insertion, connected
for i, a in enumerate(g.adj):
assert i not in a
for j in a:
assert i in g.adj[j]
assert sum(1 for j in a if j < i) >= min(f, i)
seen, stack = {0}, [0]
while stack:
for j in g.adj[stack.pop()]:
if j not in seen:
seen.add(j)
stack.append(j)
assert len(seen) == n
q = tuple(rng.uniform(0, 10) for _ in range(dim))
k = rng.randint(1, 5)
# a full beam is exact
assert g.search(q, k, n) == brute(pts, q, k)
# any beam returns the best of what it evaluated
ef = rng.randint(1, 8)
got = g.search(q, k, ef)
seen_pts = sorted((math.dist(pts[i], q), i) for i in g.last_visited)
assert got == seen_pts[:k]
What the next level needs
- the tests did not run in this build