🔎 向量检索 ANN

# LSH 简易实现 - 基于随机超平面
import numpy as np
from collections import defaultdict

class LSHTable:
    """LSH 哈希表 - 随机超平面哈希"""
    
    def __init__(self, num_hyperplanes, dim):
        # 随机超平面法向量
        self.hyperplanes = np.random.randn(num_hyperplanes, dim)
        self.hash_table = defaultdict(list)
    
    def _hash(self, vector):
        """计算向量的哈希值（二进制）"""
        projections = np.dot(self.hyperplanes, vector)
        # 根据投影正负生成二进制哈希
        return tuple((projections > 0).astype(int))
    
    def insert(self, idx, vector):
        """插入向量"""
        hash_key = self._hash(vector)
        self.hash_table[hash_key].append(idx)
    
    def query(self, vector):
        """查询候选集"""
        hash_key = self._hash(vector)
        return self.hash_table[hash_key]


class LSHIndex:
    """LSH 索引 - 多哈希表增强"""
    
    def __init__(self, dim, num_tables=10, num_hyperplanes=16):
        self.tables = [
            LSHTable(num_hyperplanes, dim) 
            for _ in range(num_tables)
        ]
    
    def build(self, vectors):
        """构建索引"""
        for idx, vec in enumerate(vectors):
            for table in self.tables:
                table.insert(idx, vec)
    
    def search(self, query_vector, k=5):
        """搜索最近邻"""
        # 收集所有候选
        candidates = set()
        for table in self.tables:
            candidates.update(table.query(query_vector))
        
        # 在候选集中精确计算距离
        scores = []
        for idx in candidates:
            dist = np.linalg.norm(query_vector - self.vectors[idx])
            scores.append((idx, dist))
        
        # 返回 top-k
        scores.sort(key=lambda x: x[1])
        return scores[:k]

# 使用示例
np.random.seed(42)
n_samples, dim = 10000, 128
vectors = np.random.randn(n_samples, dim)

# 构建索引
lsh = LSHIndex(dim=dim, num_tables=10)
lsh.vectors = vectors  # 存储原始向量
lsh.build(vectors)

# 查询
query = np.random.randn(dim)
results = lsh.search(query, k=5)
print(f"Top-5 最近邻：{results}")

# FAISS 使用示例
import faiss
import numpy as np

# 准备数据
d = 128  # 向量维度
nb = 100000  # 数据库大小
np.random.seed(42)
xb = np.random.random((nb, d)).astype('float32')

# ========== 1. 暴力搜索（精确） ==========
index_flat = faiss.IndexFlatL2(d)
index_flat.add(xb)

# 查询
nq = 10
xq = np.random.random((nq, d)).astype('float32')
D, I = index_flat.search(xq, k=5)
print(f"暴力搜索 - 距离：{D[0]}, 索引：{I[0]}")

# ========== 2. IVF 倒排索引 ==========
nlist = 100  # 聚类中心数量
quantizer = faiss.IndexFlatL2(d)
index_ivf = faiss.IndexIVFFlat(quantizer, d, nlist)

# 训练（需要）
index_ivf.train(xb)
index_ivf.add(xb)

# 设置查询参数
index_ivf.nprobe = 10  # 搜索的聚类数量
D, I = index_ivf.search(xq, k=5)
print(f"IVF 搜索 - 距离：{D[0]}, 索引：{I[0]}")

# ========== 3. IVF + PQ（压缩） ==========
m = 8  # 子向量数量
nbits = 8  # 每子向量比特数
index_pq = faiss.IndexIVFPQ(quantizer, d, nlist, m, nbits)

index_pq.train(xb)
index_pq.add(xb)
index_pq.nprobe = 10

D, I = index_pq.search(xq, k=5)
print(f"IVF+PQ 搜索 - 距离：{D[0]}, 索引：{I[0]}")

# ========== 4. HNSW ==========
M = 16  # 连接数
index_hnsw = faiss.IndexHNSWFlat(d, M)
index_hnsw.hnsw.efSearch = 64  # 查询参数
index_hnsw.add(xb)

D, I = index_hnsw.search(xq, k=5)
print(f"HNSW 搜索 - 距离：{D[0]}, 索引：{I[0]}")

# ========== 5. 索引持久化 ==========
faiss.write_index(index_hnsw, "hnsw.index")
loaded_index = faiss.read_index("hnsw.index")