Approximate Nearest Neighbor Search

import numpy as np

class RandomProjectionLSH:
    def __init__(self, dim, num_tables=10, hash_len=8):
        self.num_tables = num_tables
        self.hash_len = hash_len
        # Random projection vectors
        self.tables = []
        for _ in range(num_tables):
            r = np.random.randn(dim, hash_len)
            self.tables.append(r)
        self.buckets = [{} for _ in range(num_tables)]
    
    def _hash(self, table_idx, vector):
        r = self.tables[table_idx]
        projection = vector @ r  # Dot product
        return tuple((projection > 0).astype(int))  # Binary hash
    
    def insert(self, item_id, vector):
        for t in range(self.num_tables):
            h = self._hash(t, vector)
            if h not in self.buckets[t]:
                self.buckets[t][h] = []
            self.buckets[t][h].append(item_id)
    
    def query(self, vector, k=10):
        candidates = set()
        for t in range(self.num_tables):
            h = self._hash(t, vector)
            if h in self.buckets[t]:
                candidates.update(self.buckets[t][h])
        # Now compute exact distances only among candidates
        return candidates

class ProductQuantizer:
    def __init__(self, dim, num_subvectors=8, num_centroids=256):
        self.M = num_subvectors  # Number of sub-vectors (m)
        self.K = num_centroids  # Number of centroids per sub-vector (k*)
        self.Ds = dim // self.M  # Dimension per sub-vector (d*)
        self.codebooks = None
    
    def fit(self, vectors):
        """Learn centroids via k-means per sub-space."""
        n = vectors.shape[0]
        subspaces = vectors.reshape(n, self.M, self.Ds)
        
        self.codebooks = []
        for m in range(self.M):
            sub_vectors = subspaces[:, m, :]
            # K-means with K centroids
            centroids = self._kmeans(sub_vectors, self.K)
            self.codebooks.append(centroids)
    
    def encode(self, vector):
        """Encode as M centroid indices."""
        sub_vectors = vector.reshape(self.M, self.Ds)
        codes = []
        for m in range(self.M):
            # Nearest centroid
            dists = np.linalg.norm(self.codebooks[m] - sub_vectors[m], axis=1)
            codes.append(np.argmin(dists))
        return codes  # M bytes

def search(hnsw_graph, query, ef=100):
    """HNSW search with ef = number of candidates to explore."""
    current = hnsw_graph.entry_point
    
    # Phase 1: Coarse search through layers (top → bottom)
    for layer in reversed(range(1, len(hnsw_graph.layers))):
        current = greedy_search(hnsw_graph, query, current, layer)
    
    # Phase 2: Fine search at bottom layer
    candidates = set()
    candidates.add(current)
    visited = set()
    results = MinHeap()
    results.push(current, distance(query, current.vector))
    
    while candidates and len(visited) < ef:
        node = candidates.pop()
        visited.add(node)
        for neighbor in hnsw_graph.bottom_layer[node]:
            if neighbor not in visited:
                d = distance(query, neighbor.vector)
                results.push(neighbor, d)
                candidates.add(neighbor)
    
    return results.k_smallest(k)

class IVF:
    def __init__(self, n_clusters=100):
        self.n_clusters = n_clusters
        self.centroids = None
        self.lists = None
    
    def fit(self, vectors, ids):
        from sklearn.cluster import KMeans
        kmeans = KMeans(n_clusters=self.n_clusters)
        labels = kmeans.fit_predict(vectors)
        self.centroids = kmeans.cluster_centers_
        
        self.lists = [[] for _ in range(self.n_clusters)]
        for idx, label in enumerate(labels):
            self.lists[label].append((vectors[idx], ids[idx]))
    
    def query(self, vector, k=10, nprobe=10):
        # Find nprobe nearest clusters
        dists = np.linalg.norm(self.centroids - vector, axis=1)
        nearest = np.argsort(dists)[:nprobe]
        
        # Search those clusters exhaustively
        candidates = []
        for cluster_id in nearest:
            candidates.extend(self.lists[cluster_id])
        
        # Exact search among candidates
        candidates.sort(key=lambda x: np.linalg.norm(x[0] - vector))
        return candidates[:k]

import faiss
import numpy as np

# Generate data
d = 128  # Dimension
n = 1000000  # 1M vectors
xb = np.random.random((n, d)).astype('float32')

# Build index
index = faiss.IndexFlatL2(d)  # Exact search (baseline)
# index = faiss.IndexIVFFlat(index, d, 100)  # IVF (faster)
# index = faiss.IndexHNSWFlat(d, 32)  # HNSW (fast, accurate)
# index = faiss.IndexPQ(d, 8, 8)  # Product Quantization (compressed)

index.train(xb)
index.add(xb)

# Search
xq = np.random.random((1, d)).astype('float32')
D, I = index.search(xq, k=10)  # 10 nearest neighbors