ruvector_core/advanced_features/

opq.rs

//! Optimized Product Quantization (OPQ) with learned rotation matrix.
//!
//! OPQ improves upon standard PQ by learning an orthogonal rotation matrix R
//! that decorrelates vector dimensions before quantization. This reduces
//! quantization error by 10-30% and yields significant recall improvements,
//! especially when vector dimensions have unequal variance.
//!
//! Training alternates between PQ codebook learning and rotation update via
//! the Procrustes solution (SVD). ADC precomputes per-subspace distance tables
//! so each database lookup costs O(num_subspaces) instead of O(d).

use crate::error::{Result, RuvectorError};
use crate::types::DistanceMetric;
use serde::{Deserialize, Serialize};

/// Configuration for Optimized Product Quantization.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OPQConfig {
    /// Number of subspaces to split the (rotated) vector into.
    pub num_subspaces: usize,
    /// Codebook size per subspace (max 256 for u8 codes).
    pub codebook_size: usize,
    /// Number of k-means iterations for codebook training.
    pub num_iterations: usize,
    /// Number of outer OPQ iterations (rotation + PQ alternation).
    pub num_opq_iterations: usize,
    /// Distance metric used for codebook training and search.
    pub metric: DistanceMetric,
}

impl Default for OPQConfig {
    fn default() -> Self {
        Self {
            num_subspaces: 8, codebook_size: 256, num_iterations: 20,
            num_opq_iterations: 10, metric: DistanceMetric::Euclidean,
        }
    }
}

impl OPQConfig {
    /// Validate configuration parameters.
    pub fn validate(&self) -> Result<()> {
        if self.codebook_size > 256 {
            return Err(RuvectorError::InvalidParameter(format!(
                "Codebook size {} exceeds u8 max 256", self.codebook_size)));
        }
        if self.num_subspaces == 0 {
            return Err(RuvectorError::InvalidParameter("num_subspaces must be > 0".into()));
        }
        if self.num_opq_iterations == 0 {
            return Err(RuvectorError::InvalidParameter("num_opq_iterations must be > 0".into()));
        }
        Ok(())
    }
}

// -- Dense matrix (row-major, internal only) ----------------------------------

#[derive(Debug, Clone)]
struct Mat { rows: usize, cols: usize, data: Vec<f32> }

impl Mat {
    fn zeros(r: usize, c: usize) -> Self { Self { rows: r, cols: c, data: vec![0.0; r * c] } }
    fn identity(n: usize) -> Self {
        let mut m = Self::zeros(n, n);
        for i in 0..n { m.data[i * n + i] = 1.0; }
        m
    }
    #[inline] fn get(&self, r: usize, c: usize) -> f32 { self.data[r * self.cols + c] }
    #[inline] fn set(&mut self, r: usize, c: usize, v: f32) { self.data[r * self.cols + c] = v; }

    fn transpose(&self) -> Self {
        let mut t = Self::zeros(self.cols, self.rows);
        for r in 0..self.rows { for c in 0..self.cols { t.set(c, r, self.get(r, c)); } }
        t
    }
    fn mul(&self, b: &Mat) -> Mat {
        assert_eq!(self.cols, b.rows);
        let mut out = Mat::zeros(self.rows, b.cols);
        for i in 0..self.rows {
            for k in 0..self.cols {
                let a = self.get(i, k);
                for j in 0..b.cols { let c = out.get(i, j); out.set(i, j, c + a * b.get(k, j)); }
            }
        }
        out
    }
    fn from_rows(vecs: &[Vec<f32>]) -> Self {
        let (rows, cols) = (vecs.len(), vecs[0].len());
        let mut data = Vec::with_capacity(rows * cols);
        for v in vecs { data.extend_from_slice(v); }
        Self { rows, cols, data }
    }
    fn row(&self, i: usize) -> Vec<f32> { self.data[i * self.cols..(i + 1) * self.cols].to_vec() }
}

// -- SVD via power iteration + deflation --------------------------------------

/// Rank-1 SVD: returns (u, sigma, v) for the largest singular triplet.
fn svd_rank1(a: &Mat, max_iters: usize) -> (Vec<f32>, f32, Vec<f32>) {
    let ata = a.transpose().mul(a);
    let n = ata.cols;
    let mut v = vec![1.0 / (n as f32).sqrt(); n];
    for _ in 0..max_iters {
        let mut nv = vec![0.0; n];
        for i in 0..n { for j in 0..n { nv[i] += ata.get(i, j) * v[j]; } }
        let norm: f32 = nv.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm < 1e-12 { break; }
        for x in nv.iter_mut() { *x /= norm; }
        v = nv;
    }
    let mut av = vec![0.0; a.rows];
    for i in 0..a.rows { for j in 0..a.cols { av[i] += a.get(i, j) * v[j]; } }
    let sigma: f32 = av.iter().map(|x| x * x).sum::<f32>().sqrt();
    let u = if sigma > 1e-12 { av.iter().map(|x| x / sigma).collect() } else { vec![0.0; a.rows] };
    (u, sigma, v)
}

/// Full SVD by repeated rank-1 extraction + deflation.
fn svd_full(a: &Mat, iters: usize) -> (Mat, Vec<f32>, Mat) {
    let n = a.rows;
    let mut res = a.clone();
    let (mut uc, mut sv, mut vc) = (Vec::new(), Vec::new(), Vec::new());
    for _ in 0..n {
        let (u, s, v) = svd_rank1(&res, iters);
        if s > 1e-10 {
            for i in 0..res.rows { for j in 0..res.cols {
                let c = res.get(i, j); res.set(i, j, c - s * u[i] * v[j]);
            }}
        }
        uc.push(u); sv.push(s); vc.push(v);
    }
    let (mut um, mut vm) = (Mat::zeros(n, n), Mat::zeros(n, n));
    for j in 0..n { for i in 0..n { um.set(i, j, uc[j][i]); vm.set(i, j, vc[j][i]); } }
    (um, sv, vm)
}

/// Procrustes: find orthogonal R minimising ||Y - X @ R||_F.
fn procrustes(x: &Mat, y: &Mat) -> Mat {
    let m = x.transpose().mul(y);
    let (u, _s, v) = svd_full(&m, 100);
    v.mul(&u.transpose())
}

// -- Rotation matrix ----------------------------------------------------------

/// Orthogonal rotation matrix R (d x d) that decorrelates dimensions before PQ.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RotationMatrix { pub dim: usize, pub data: Vec<f32> }

impl RotationMatrix {
    /// Identity rotation (no-op).
    pub fn identity(dim: usize) -> Self {
        let mut data = vec![0.0; dim * dim];
        for i in 0..dim { data[i * dim + i] = 1.0; }
        Self { dim, data }
    }
    /// Rotate vector: y = x @ R.
    pub fn rotate(&self, v: &[f32]) -> Vec<f32> {
        let d = self.dim;
        (0..d).map(|j| (0..d).map(|i| v[i] * self.data[i * d + j]).sum()).collect()
    }
    /// Inverse rotate: x = y @ R^T.
    pub fn inverse_rotate(&self, v: &[f32]) -> Vec<f32> {
        let d = self.dim;
        (0..d).map(|j| (0..d).map(|i| v[i] * self.data[j * d + i]).sum()).collect()
    }
    fn from_mat(m: &Mat) -> Self { Self { dim: m.rows, data: m.data.clone() } }
}

// -- OPQ Index ----------------------------------------------------------------

/// OPQ index: learns rotation R + PQ codebooks, supports ADC search.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OPQIndex {
    pub config: OPQConfig,
    pub rotation: RotationMatrix,
    /// Codebooks: `[subspace][centroid][subspace_dim]`.
    pub codebooks: Vec<Vec<Vec<f32>>>,
    pub dimensions: usize,
}

impl OPQIndex {
    /// Train OPQ via alternating rotation update and PQ codebook learning.
    pub fn train(vectors: &[Vec<f32>], config: OPQConfig) -> Result<Self> {
        config.validate()?;
        if vectors.is_empty() {
            return Err(RuvectorError::InvalidParameter("Training set cannot be empty".into()));
        }
        let d = vectors[0].len();
        if d % config.num_subspaces != 0 {
            return Err(RuvectorError::InvalidParameter(format!(
                "Dimensions {} not divisible by num_subspaces {}", d, config.num_subspaces)));
        }
        for v in vectors { if v.len() != d {
            return Err(RuvectorError::DimensionMismatch { expected: d, actual: v.len() });
        }}
        let x_mat = Mat::from_rows(vectors);
        let mut r = Mat::identity(d);
        let mut codebooks: Vec<Vec<Vec<f32>>> = Vec::new();
        let sub_dim = d / config.num_subspaces;
        for _ in 0..config.num_opq_iterations {
            let x_rot = x_mat.mul(&r);
            let rotated: Vec<Vec<f32>> = (0..vectors.len()).map(|i| x_rot.row(i)).collect();
            codebooks = train_pq_codebooks(&rotated, config.num_subspaces,
                config.codebook_size, config.num_iterations, config.metric)?;
            let mut x_hat = Mat::zeros(vectors.len(), d);
            for (i, rv) in rotated.iter().enumerate() {
                let codes = encode_vec(rv, &codebooks, sub_dim, config.metric)?;
                let recon = decode_vec(&codes, &codebooks);
                for (j, &val) in recon.iter().enumerate() { x_hat.set(i, j, val); }
            }
            r = procrustes(&x_mat, &x_hat);
        }
        Ok(Self { config, rotation: RotationMatrix::from_mat(&r), codebooks, dimensions: d })
    }

    /// Encode a vector: rotate then PQ-quantize.
    pub fn encode(&self, vector: &[f32]) -> Result<Vec<u8>> {
        self.check_dim(vector.len())?;
        let rotated = self.rotation.rotate(vector);
        encode_vec(&rotated, &self.codebooks,
            self.dimensions / self.config.num_subspaces, self.config.metric)
    }

    /// Decode PQ codes back to approximate vector (with inverse rotation).
    pub fn decode(&self, codes: &[u8]) -> Result<Vec<f32>> {
        if codes.len() != self.config.num_subspaces {
            return Err(RuvectorError::InvalidParameter(format!(
                "Expected {} codes, got {}", self.config.num_subspaces, codes.len())));
        }
        Ok(self.rotation.inverse_rotate(&decode_vec(codes, &self.codebooks)))
    }

    /// ADC search: precompute distance tables then sum lookups per database vector.
    pub fn search_adc(&self, query: &[f32], codes_db: &[Vec<u8>], top_k: usize,
    ) -> Result<Vec<(usize, f32)>> {
        self.check_dim(query.len())?;
        let rq = self.rotation.rotate(query);
        let sub_dim = self.dimensions / self.config.num_subspaces;
        let tables: Vec<Vec<f32>> = (0..self.config.num_subspaces).map(|s| {
            let q_sub = &rq[s * sub_dim..(s + 1) * sub_dim];
            self.codebooks[s].iter().map(|c| dist(q_sub, c, self.config.metric)).collect()
        }).collect();
        let mut dists: Vec<(usize, f32)> = codes_db.iter().enumerate().map(|(idx, codes)| {
            let d: f32 = codes.iter().enumerate().map(|(s, &c)| tables[s][c as usize]).sum();
            (idx, d)
        }).collect();
        dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
        dists.truncate(top_k);
        Ok(dists)
    }

    /// Mean squared quantization error over a set of vectors.
    pub fn quantization_error(&self, vectors: &[Vec<f32>]) -> Result<f32> {
        if vectors.is_empty() { return Ok(0.0); }
        let mut total = 0.0f64;
        for v in vectors {
            let recon = self.decode(&self.encode(v)?)?;
            total += v.iter().zip(&recon).map(|(a, b)| ((a - b) as f64).powi(2)).sum::<f64>();
        }
        Ok((total / vectors.len() as f64) as f32)
    }

    fn check_dim(&self, len: usize) -> Result<()> {
        if len != self.dimensions {
            Err(RuvectorError::DimensionMismatch { expected: self.dimensions, actual: len })
        } else { Ok(()) }
    }
}

// -- PQ helpers ---------------------------------------------------------------

fn dist(a: &[f32], b: &[f32], m: DistanceMetric) -> f32 {
    match m {
        DistanceMetric::Euclidean =>
            a.iter().zip(b).map(|(x, y)| { let d = x - y; d * d }).sum::<f32>().sqrt(),
        DistanceMetric::Cosine => {
            let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
            let na = a.iter().map(|x| x * x).sum::<f32>().sqrt();
            let nb = b.iter().map(|x| x * x).sum::<f32>().sqrt();
            if na == 0.0 || nb == 0.0 { 1.0 } else { 1.0 - dot / (na * nb) }
        }
        DistanceMetric::DotProduct => -a.iter().zip(b).map(|(x, y)| x * y).sum::<f32>(),
        DistanceMetric::Manhattan => a.iter().zip(b).map(|(x, y)| (x - y).abs()).sum(),
    }
}

fn train_pq_codebooks(vecs: &[Vec<f32>], nsub: usize, k: usize, iters: usize,
    metric: DistanceMetric) -> Result<Vec<Vec<Vec<f32>>>> {
    let sub_dim = vecs[0].len() / nsub;
    (0..nsub).map(|s| {
        let sv: Vec<Vec<f32>> = vecs.iter().map(|v| v[s*sub_dim..(s+1)*sub_dim].to_vec()).collect();
        kmeans(&sv, k.min(sv.len()), iters, metric)
    }).collect()
}

fn encode_vec(v: &[f32], cbs: &[Vec<Vec<f32>>], sub_dim: usize, m: DistanceMetric,
) -> Result<Vec<u8>> {
    cbs.iter().enumerate().map(|(s, cb)| {
        let sub = &v[s * sub_dim..(s + 1) * sub_dim];
        cb.iter().enumerate()
            .min_by(|(_, a), (_, b)| dist(sub, a, m).partial_cmp(&dist(sub, b, m)).unwrap())
            .map(|(i, _)| i as u8)
            .ok_or_else(|| RuvectorError::Internal("Empty codebook".into()))
    }).collect()
}

fn decode_vec(codes: &[u8], cbs: &[Vec<Vec<f32>>]) -> Vec<f32> {
    codes.iter().enumerate().flat_map(|(s, &c)| cbs[s][c as usize].iter().copied()).collect()
}

fn kmeans(vecs: &[Vec<f32>], k: usize, iters: usize, metric: DistanceMetric,
) -> Result<Vec<Vec<f32>>> {
    use rand::seq::SliceRandom;
    if vecs.is_empty() || k == 0 {
        return Err(RuvectorError::InvalidParameter("Cannot cluster empty set or k=0".into()));
    }
    let dim = vecs[0].len();
    let mut rng = rand::thread_rng();
    let mut cents: Vec<Vec<f32>> = vecs.choose_multiple(&mut rng, k).cloned().collect();
    for _ in 0..iters {
        let (mut sums, mut counts) = (vec![vec![0.0f32; dim]; k], vec![0usize; k]);
        for v in vecs {
            let b = cents.iter().enumerate()
                .min_by(|(_, a), (_, b)| dist(v, a, metric).partial_cmp(&dist(v, b, metric)).unwrap())
                .map(|(i, _)| i).unwrap_or(0);
            counts[b] += 1;
            for (j, &val) in v.iter().enumerate() { sums[b][j] += val; }
        }
        for (i, c) in cents.iter_mut().enumerate() {
            if counts[i] > 0 { for j in 0..dim { c[j] = sums[i][j] / counts[i] as f32; } }
        }
    }
    Ok(cents)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_data(n: usize, d: usize) -> Vec<Vec<f32>> {
        let mut seed: u64 = 42;
        (0..n).map(|_| (0..d).map(|_| {
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
            ((seed >> 33) as f32) / (u32::MAX as f32) * 2.0 - 1.0
        }).collect()).collect()
    }
    fn cfg() -> OPQConfig {
        OPQConfig { num_subspaces: 2, codebook_size: 4, num_iterations: 5,
            num_opq_iterations: 3, metric: DistanceMetric::Euclidean }
    }

    #[test]
    fn test_rotation_orthogonality() {
        let r = RotationMatrix::identity(4);
        let v = vec![1.0, 2.0, 3.0, 4.0];
        let back = r.inverse_rotate(&r.rotate(&v));
        for i in 0..4 { assert!((v[i] - back[i]).abs() < 1e-6); }
    }
    #[test]
    fn test_rotation_preserves_norm() {
        let data = make_data(30, 4);
        let idx = OPQIndex::train(&data, cfg()).unwrap();
        let v = vec![1.0, 2.0, 3.0, 4.0];
        let n1: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        let n2: f32 = idx.rotation.rotate(&v).iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!((n1 - n2).abs() < 0.1, "norms: {} vs {}", n1, n2);
    }
    #[test]
    fn test_pq_encoding_roundtrip() {
        let data = make_data(30, 4);
        let idx = OPQIndex::train(&data, cfg()).unwrap();
        let codes = idx.encode(&data[0]).unwrap();
        assert_eq!(codes.len(), 2);
        assert_eq!(idx.decode(&codes).unwrap().len(), 4);
    }
    #[test]
    fn test_opq_training_convergence() {
        // Verify that OPQ training produces finite, non-negative quantization
        // error and that trained index can encode/decode without degradation.
        // Note: with small data and few centroids, more OPQ iterations do not
        // guarantee monotone error decrease due to stochastic k-means.
        let data = make_data(100, 4);
        let idx = OPQIndex::train(&data, cfg()).unwrap();
        let err = idx.quantization_error(&data).unwrap();
        assert!(err.is_finite() && err >= 0.0, "error must be finite non-negative: {}", err);
        // Verify round-trip through encode/decode does not explode.
        for v in &data {
            let codes = idx.encode(v).unwrap();
            let decoded = idx.decode(&codes).unwrap();
            assert_eq!(decoded.len(), v.len());
            for x in &decoded { assert!(x.is_finite()); }
        }
    }
    #[test]
    fn test_adc_correctness() {
        let data = make_data(30, 4);
        let idx = OPQIndex::train(&data, cfg()).unwrap();
        let db: Vec<Vec<u8>> = data.iter().map(|v| idx.encode(v).unwrap()).collect();
        let res = idx.search_adc(&[0.5, -0.5, 0.5, -0.5], &db, 3).unwrap();
        assert_eq!(res.len(), 3);
        for w in res.windows(2) { assert!(w[0].1 <= w[1].1 + 1e-6); }
    }
    #[test]
    fn test_quantization_error_reduction() {
        let data = make_data(50, 4);
        let err = OPQIndex::train(&data, cfg()).unwrap().quantization_error(&data).unwrap();
        assert!(err >= 0.0 && err.is_finite() && err < 10.0, "err={}", err);
    }
    #[test]
    fn test_svd_correctness() {
        let a = Mat { rows: 2, cols: 2, data: vec![3.0, 0.0, 0.0, 2.0] };
        let (u, s, v) = svd_full(&a, 200);
        for i in 0..2 { for j in 0..2 {
            let r: f32 = (0..2).map(|k| u.get(i, k) * s[k] * v.get(j, k)).sum();
            assert!((a.get(i, j) - r).abs() < 0.1, "SVD fail ({},{}): {} vs {}", i, j, a.get(i, j), r);
        }}
    }
    #[test]
    fn test_identity_rotation_baseline() {
        let data = make_data(30, 4);
        let idx = OPQIndex::train(&data, OPQConfig { num_opq_iterations: 1, ..cfg() }).unwrap();
        let recon = idx.decode(&idx.encode(&data[0]).unwrap()).unwrap();
        assert_eq!(recon.len(), data[0].len());
    }
    #[test]
    fn test_search_accuracy() {
        let data = make_data(40, 4);
        let idx = OPQIndex::train(&data, cfg()).unwrap();
        let db: Vec<Vec<u8>> = data.iter().map(|v| idx.encode(v).unwrap()).collect();
        let ids: Vec<usize> = idx.search_adc(&data[0], &db, 5).unwrap().iter().map(|r| r.0).collect();
        assert!(ids.contains(&0), "vector 0 should be in its own top-5");
    }
    #[test]
    fn test_config_validation() {
        assert!(OPQConfig { codebook_size: 300, ..cfg() }.validate().is_err());
        assert!(OPQConfig { num_subspaces: 0, ..cfg() }.validate().is_err());
        assert!(OPQConfig { num_opq_iterations: 0, ..cfg() }.validate().is_err());
    }
    #[test]
    fn test_dimension_mismatch_errors() {
        let idx = OPQIndex::train(&make_data(30, 4), cfg()).unwrap();
        assert!(idx.encode(&[1.0, 2.0]).is_err());
        assert!(idx.search_adc(&[1.0], &[], 1).is_err());
    }
}