tests/test_lsh.py

"""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)