Method train

nodedb-codec/src/vector_quant/opq.rs:74–160

Train an OPQ codec using the Non-Para OPQ algorithm. Alternates between a codebook step (Lloyd's k-means on the rotated training set) and a Procrustes step (SVD-based rotation update to minimize reconstruction error) for `opq_iters` iterations. - `opq_iters`: number of alternating Procrustes+codebook iterations. - `kmeans_iters`: Lloyd's k-means iterations per subspace per OPQ iter.

(
        vectors: &[&[f32]],
        dim: usize,
        m: usize,
        k: usize,
        opq_iters: usize,
        kmeans_iters: usize,
    )

    /// - `opq_iters`: number of alternating Procrustes+codebook iterations.
    /// - `kmeans_iters`: Lloyd's k-means iterations per subspace per OPQ iter.
    pub fn train(
        vectors: &[&[f32]],
        dim: usize,
        m: usize,
        k: usize,
        opq_iters: usize,
        kmeans_iters: usize,
    ) -> Self {
        assert!(!vectors.is_empty(), "training set must be non-empty");
        assert!(dim > 0 && m > 0 && k > 0, "dim/m/k must be positive");
        assert!(
            dim.is_multiple_of(m),
            "dim ({dim}) must be divisible by m ({m})"
        );
        let sub_dim = dim / m;
        let seed = dim as u64 ^ ((m as u64) << 16) ^ ((k as u64) << 32);

        let mut rotation = identity(dim);
        let mut codebooks: Vec<Vec<Vec<f32>>> = Vec::new();

        let iters = opq_iters.max(1);

        for iter in 0..iters {
            // Codebook step: train PQ on the current rotated training set.
            let rotated: Vec<Vec<f32>> =
                vectors.iter().map(|v| matvec(&rotation, v, dim)).collect();
            codebooks = train_codebooks(&rotated, m, k, sub_dim, kmeans_iters, seed ^ iter as u64);

            // Procrustes step: find R minimising ‖R·X - Y‖_F where Y is
            // the dequantized reconstruction of R·X.
            //
            // Closed-form solution (Ge et al. CVPR 2013, §3.2):
            //   M = X · Yᵀ   (dim × dim)
            //   SVD(M) = U Σ Vᵀ
            //   R_new = V · Uᵀ
            //
            // Skip rotation update on the last iteration — codebooks were
            // already retrained with the current R.
            if iter + 1 < iters {
                let n = vectors.len();
                // Build dim×N matrices X (original) and Y (reconstructed).
                // DMatrix is column-major; we store column j = vector j.
                let x_mat = DMatrix::from_fn(dim, n, |row, col| vectors[col][row]);
                let y_mat = {
                    let recon: Vec<Vec<f32>> = rotated
                        .iter()
                        .map(|rv| {
                            let codes = pq_encode(rv, &codebooks, m, sub_dim);
                            dequantize_codes(&codes, &codebooks)
                        })
                        .collect();
                    DMatrix::from_fn(dim, n, |row, col| recon[col][row])
                };

                // M = X · Yᵀ  (dim × dim)
                let m_mat = &x_mat * y_mat.transpose();

                // Guard: skip rotation update if M contains NaN (degenerate