Method train

nodedb-vector/src/multivec/plaid.rs:208–254

Train the pruner from a `MultiVectorStore`. `num_centroids` — number of K-means clusters. `kmeans_iters` — Lloyd iterations. `seed` — deterministic seed for centroid initialisation.

(
        store: &MultiVectorStore,
        num_centroids: u16,
        kmeans_iters: usize,
        seed: u64,
    )

    /// * `kmeans_iters` — Lloyd iterations.
    /// * `seed` — deterministic seed for centroid initialisation.
    pub fn train(
        store: &MultiVectorStore,
        num_centroids: u16,
        kmeans_iters: usize,
        seed: u64,
    ) -> Self {
        let dim = store.dim;
        let nc = num_centroids as usize;

        // Collect all document vectors for K-means training.
        let all_vectors: Vec<Vec<f32>> = store
            .iter()
            .flat_map(|doc| doc.vectors.iter().cloned())
            .collect();

        let centroids = kmeans(&all_vectors, nc, kmeans_iters, seed, dim);

        // Encode each document as a sorted, deduplicated bag of centroid IDs.
        let doc_centroids: HashMap<u32, Vec<u16>> = store
            .iter()
            .map(|doc| {
                let mut ids: Vec<u16> = doc
                    .vectors
                    .iter()
                    .map(|v| {
                        centroids
                            .iter()
                            .enumerate()
                            .map(|(i, c)| (i as u16, scalar_distance(v, c, DistanceMetric::L2)))
                            .min_by(|a, b| {
                                a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)
                            })
                            .map(|(i, _)| i)
                            .unwrap_or(0)
                    })
                    .collect();
                ids.sort_unstable();
                ids.dedup();
                (doc.doc_id, ids)
            })
            .collect();

        Self {
            centroids,
            doc_centroids,
        }
    }

    /// Return candidate doc IDs whose centroid bag overlaps the query's
    /// centroid bag.