Upload tests/test_lsh.py with huggingface_hub

tests/test_lsh.py

@@ -0,0 +1,137 @@
+"""Tests for the Hamming LSH index (``palimseste.lsh``).
+
+Key properties to verify:
+  - exact neighbors are always retrievable within radius (recall on self)
+  - near-duplicates collide far more often than random pairs (selectivity)
+  - query_candidates returns a small fraction of |M| (sub-linear scaling)
+  - rebuild is consistent with incremental insert
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from palimseste import hv, lsh
+
+
+def _flip_bits(x: hv.HV, n: int, rng) -> hv.HV:
+    """Return a copy of x with n random bit positions flipped."""
+    signs = hv.bits_to_signs(x)
+    pos = rng.choice(x.D, size=n, replace=False)
+    signs[pos] = -signs[pos]
+    return hv.signs_to_bits(signs)
+
+
+def test_insert_and_size():
+    cfg = lsh.LSHConfig(D=500, K=8, L=4)
+    idx = lsh.LSHIndex(config=cfg)
+    rng = np.random.default_rng(0)
+    for _ in range(10):
+        idx.insert(idx.size, hv.random_hv(D=500, rng=rng))
+    assert idx.size == 10
+
+
+def test_self_query_returns_self():
+    cfg = lsh.LSHConfig(D=1000, K=6, L=6)
+    idx = lsh.LSHIndex(config=cfg)
+    rng = np.random.default_rng(1)
+    hvs = [hv.random_hv(D=1000, rng=rng) for _ in range(50)]
+    for i, h in enumerate(hvs):
+        idx.insert(i, h)
+    # query with one of the stored HVs at radius 0 -> should find itself
+    for i in (0, 17, 49):
+        nbs = idx.query_neighbors(hvs[i], radius=0, items=hvs)
+        assert i in nbs
+
+
+def test_near_neighbor_found():
+    cfg = lsh.LSHConfig(D=2000, K=10, L=12)
+    idx = lsh.LSHIndex(config=cfg)
+    rng = np.random.default_rng(2)
+    base = hv.random_hv(D=2000, rng=rng)
+    # near neighbor: 5 bits different
+    near = _flip_bits(base, 5, rng)
+    # fill with unrelated random HVs
+    hvs = [base] + [hv.random_hv(D=2000, rng=rng) for _ in range(200)]
+    hvs.append(near)
+    near_id = len(hvs) - 1
+    for i, h in enumerate(hvs):
+        idx.insert(i, h)
+    found = idx.query_neighbors(near, radius=10, items=hvs)
+    # both base and near should be within radius 10 of `near`
+    assert 0 in found
+    assert near_id in found
+
+
+def test_candidate_set_is_small():
+    # With reasonable (K,L), candidates << |M|
+    cfg = lsh.LSHConfig(D=2000, K=14, L=8)
+    idx = lsh.LSHIndex(config=cfg)
+    rng = np.random.default_rng(3)
+    N = 1000
+    hvs = [hv.random_hv(D=2000, rng=rng) for _ in range(N)]
+    for i, h in enumerate(hvs):
+        idx.insert(i, h)
+    q = hv.random_hv(D=2000, rng=rng)
+    cand = idx.query_candidates(q)
+    # should be well under N
+    assert len(cand) < N * 0.5
+
+
+def test_rebuild_matches_insert():
+    # Two indices built with identical rng seeds must have identical
+    # projections, so rebuild() and incremental insert() produce the
+    # same bucket membership.
+    cfg = lsh.LSHConfig(D=500, K=6, L=4)
+    rng = np.random.default_rng(4)
+    hvs = [hv.random_hv(D=500, rng=rng) for _ in range(30)]
+
+    idx = lsh.LSHIndex(config=cfg, _rng=np.random.default_rng(123))
+    for i, h in enumerate(hvs):
+        idx.insert(i, h)
+    cand_before = idx.query_candidates(hvs[5])
+
+    idx2 = lsh.LSHIndex(config=cfg, _rng=np.random.default_rng(123))
+    idx2.rebuild(hvs)
+    cand_after = idx2.query_candidates(hvs[5])
+    assert cand_before == cand_after
+
+
+def test_rebuild_clears_old_entries():
+    cfg = lsh.LSHConfig(D=500, K=6, L=4)
+    idx = lsh.LSHIndex(config=cfg, _rng=np.random.default_rng(7))
+    rng = np.random.default_rng(0)
+    # insert some, then rebuild with a smaller set -> size must reflect new set
+    for _ in range(50):
+        idx.insert(idx.size, hv.random_hv(D=500, rng=rng))
+    new_hvs = [hv.random_hv(D=500, rng=rng) for _ in range(5)]
+    idx.rebuild(new_hvs)
+    assert idx.size == 5
+
+
+def test_tune_config_valid():
+    cfg = lsh.LSHConfig.tune(D=10000, target_radius=0.1, recall=0.9)
+    assert 1 <= cfg.K <= cfg.D
+    assert cfg.L >= 1
+    # larger recall -> at least as many tables
+    cfg_hi = lsh.LSHConfig.tune(D=10000, target_radius=0.1, recall=0.99)
+    assert cfg_hi.L >= cfg.L
+
+
+def test_dimension_mismatch_raises():
+    cfg = lsh.LSHConfig(D=100)
+    idx = lsh.LSHIndex(config=cfg)
+    with pytest.raises(ValueError):
+        idx.insert(0, hv.random_hv(D=200))
+
+
+def test_invalid_config():
+    with pytest.raises(ValueError):
+        lsh.LSHConfig(D=0)
+    with pytest.raises(ValueError):
+        lsh.LSHConfig(D=100, K=0)
+    with pytest.raises(ValueError):
+        lsh.LSHConfig(D=100, L=0)
+    with pytest.raises(ValueError):
+        lsh.LSHConfig(D=100, K=101)