use std::fmt;
use super::{common::VectorRelationBinding, definition::CuvsOperation};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum VectorKMeansInit {
KMeansPlusPlus,
Random,
Relation,
}
impl VectorKMeansInit {
pub(crate) const fn as_str(self) -> &'static str {
match self {
Self::KMeansPlusPlus => "kmeans_plus_plus",
Self::Random => "random",
Self::Relation => "relation",
}
}
}
impl fmt::Display for VectorKMeansInit {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct VectorKMeansParams {
pub input: VectorRelationBinding,
pub n_clusters: usize,
pub max_iterations: u32,
pub tolerance: f64,
pub n_init: u32,
pub init: VectorKMeansInit,
pub initial_centroids: Option<VectorRelationBinding>,
pub metric: VectorKMeansMetric,
}
impl VectorKMeansParams {
pub fn new(input: VectorRelationBinding, n_clusters: usize) -> Self {
let descriptor = CuvsOperation::KMeans.descriptor();
Self {
input,
n_clusters,
max_iterations: descriptor
.option("max_iter")
.and_then(|option| option.kind.default_u64())
.expect("canonical KMeans max_iter default") as u32,
tolerance: descriptor
.option("tol")
.and_then(|option| option.kind.default_f64())
.expect("canonical KMeans tol default"),
n_init: descriptor
.option("n_init")
.and_then(|option| option.kind.default_u64())
.expect("canonical KMeans n_init default") as u32,
init: match descriptor.option("init").and_then(|option| option.kind.default_str()) {
Some("random") => VectorKMeansInit::Random,
Some("kmeans++") => VectorKMeansInit::KMeansPlusPlus,
_ => panic!("canonical KMeans init default"),
},
initial_centroids: None,
metric: match descriptor.option("metric").and_then(|option| option.kind.default_str()) {
Some("l2_sqrt_expanded") => VectorKMeansMetric::L2SqrtExpanded,
Some("l2_expanded") => VectorKMeansMetric::L2Expanded,
_ => panic!("canonical KMeans metric default"),
},
}
}
pub fn with_initial_centroids(mut self, centroids: VectorRelationBinding) -> Self {
self.init = VectorKMeansInit::Relation;
self.initial_centroids = Some(centroids);
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum VectorKMeansMetric {
L2Expanded,
L2SqrtExpanded,
}
impl VectorKMeansMetric {
pub(crate) const fn as_str(self) -> &'static str {
match self {
Self::L2Expanded => "l2_expanded",
Self::L2SqrtExpanded => "l2_sqrt_expanded",
}
}
}
impl fmt::Display for VectorKMeansMetric {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.as_str())
}
}