use crate::core::{ConfigError, FormicaXError};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KMeansVariant {
Lloyd,
Elkan,
Hamerly,
MiniBatch,
}
impl Default for KMeansVariant {
fn default() -> Self {
Self::Lloyd
}
}
#[derive(Debug, Clone)]
pub struct KMeansConfig {
pub k: usize,
pub variant: KMeansVariant,
pub parallel: bool,
pub max_iterations: usize,
pub tolerance: f64,
pub random_seed: Option<u64>,
pub simd: bool,
}
impl KMeansConfig {
pub fn builder() -> KMeansConfigBuilder {
KMeansConfigBuilder::default()
}
}
#[derive(Debug, Clone, Default)]
pub struct KMeansConfigBuilder {
k: Option<usize>,
variant: Option<KMeansVariant>,
parallel: Option<bool>,
max_iterations: Option<usize>,
tolerance: Option<f64>,
random_seed: Option<u64>,
}
impl KMeansConfigBuilder {
pub fn k(mut self, k: usize) -> Self {
self.k = Some(k);
self
}
pub fn variant(mut self, variant: KMeansVariant) -> Self {
self.variant = Some(variant);
self
}
pub fn parallel(mut self, parallel: bool) -> Self {
self.parallel = Some(parallel);
self
}
pub fn max_iterations(mut self, max_iterations: usize) -> Self {
self.max_iterations = Some(max_iterations);
self
}
pub fn tolerance(mut self, tolerance: f64) -> Self {
self.tolerance = Some(tolerance);
self
}
pub fn random_seed(mut self, random_seed: u64) -> Self {
self.random_seed = Some(random_seed);
self
}
pub fn build(self) -> Result<KMeansConfig, FormicaXError> {
let k = self
.k
.ok_or(FormicaXError::Config(ConfigError::MissingField {
field: "k".to_string(),
}))?;
if k == 0 {
return Err(FormicaXError::Config(ConfigError::InvalidValue {
field: "k".to_string(),
message: "must be greater than 0".to_string(),
}));
}
Ok(KMeansConfig {
k,
variant: self.variant.unwrap_or(KMeansVariant::Lloyd),
parallel: self.parallel.unwrap_or(false),
max_iterations: self.max_iterations.unwrap_or(100),
tolerance: self.tolerance.unwrap_or(1e-8),
random_seed: self.random_seed,
simd: false,
})
}
}