incremental_rs/
minibatch_kmeans.rs

1use crate::error::IncrementalError;
2use crate::IncrementalUnsupervisedEstimator;
3use ndarray::{Array1, Array2};
4use rand::Rng;
5
6#[derive(Debug)]
7pub struct MiniBatchKMeans {
8    n_clusters: usize,
9    centroids: Option<Array2<f64>>,
10    counts: Array1<f64>, // Cumulative point counts assigned to each cluster
11    n_features: Option<usize>,
12}
13
14impl MiniBatchKMeans {
15    pub fn new(n_clusters: usize) -> Self {
16        Self {
17            n_clusters,
18            centroids: None,
19            counts: Array1::zeros(n_clusters),
20            n_features: None,
21        }
22    }
23
24    /// Calculates Euclidean distance squared between a single row and a centroid row.
25    fn sq_euclidean(a: &ndarray::ArrayView1<f64>, b: &ndarray::ArrayView1<f64>) -> f64 {
26        a.iter()
27            .zip(b.iter())
28            .map(|(x, y)| (x - y).powi(2))
29            .sum()
30    }
31
32    /// Initialize centroids randomly from the first mini-batch.
33    fn init_centroids(&mut self, batch_x: &Array2<f64>) -> Result<(), IncrementalError> {
34        let n_samples = batch_x.nrows();
35        let n_features = batch_x.ncols();
36
37        if n_samples < self.n_clusters {
38            return Err(IncrementalError::EmptyBatch);
39        }
40
41        let mut rng = rand::thread_rng();
42        let mut centroids = Array2::zeros((self.n_clusters, n_features));
43        let mut chosen_indices = Vec::with_capacity(self.n_clusters);
44
45        while chosen_indices.len() < self.n_clusters {
46            let idx = rng.gen_range(0..n_samples);
47            if !chosen_indices.contains(&idx) {
48                chosen_indices.push(idx);
49            }
50        }
51
52        for (c_idx, &s_idx) in chosen_indices.iter().enumerate() {
53            centroids.row_mut(c_idx).assign(&batch_x.row(s_idx));
54        }
55
56        self.centroids = Some(centroids);
57        self.n_features = Some(n_features);
58        Ok(())
59    }
60
61    fn validate_batch(&self, batch_x: &Array2<f64>) -> Result<(), IncrementalError> {
62        if batch_x.nrows() == 0 {
63            return Err(IncrementalError::EmptyBatch);
64        }
65        if batch_x.iter().any(|v| !v.is_finite()) {
66            return Err(IncrementalError::NonFiniteInput);
67        }
68        if let Some(n_f) = self.n_features {
69            if batch_x.ncols() != n_f {
70                return Err(IncrementalError::DimensionMismatch {
71                    expected: n_f,
72                    actual: batch_x.ncols(),
73                });
74            }
75        }
76        Ok(())
77    }
78
79    /// Access calculated cluster centroids.
80    pub fn centroids(&self) -> Option<&Array2<f64>> {
81        self.centroids.as_ref()
82    }
83}
84
85impl IncrementalUnsupervisedEstimator for MiniBatchKMeans {
86    fn partial_fit(&mut self, batch_x: &Array2<f64>) -> Result<(), IncrementalError> {
87        self.validate_batch(batch_x)?;
88
89        if self.centroids.is_none() {
90            self.init_centroids(batch_x)?;
91        }
92
93        let centroids = self.centroids.as_mut().unwrap();
94
95        // 1. Assign samples to closest centroids
96        let mut assignments = Vec::with_capacity(batch_x.nrows());
97        for row in batch_x.rows() {
98            let mut best_cluster = 0;
99            let mut min_dist = f64::INFINITY;
100
101            for (c_idx, c_row) in centroids.rows().into_iter().enumerate() {
102                let dist = Self::sq_euclidean(&row, &c_row);
103                if dist < min_dist {
104                    min_dist = dist;
105                    best_cluster = c_idx;
106                }
107            }
108            assignments.push(best_cluster);
109        }
110
111        // 2. Update centroids using Sculley's per-cluster learning rate
112        for (row_idx, &c_idx) in assignments.iter().enumerate() {
113            self.counts[c_idx] += 1.0;
114            let eta = 1.0 / self.counts[c_idx]; // Per-cluster step size[cite: 1]
115
116            let x_i = batch_x.row(row_idx);
117            let mut c_i = centroids.row_mut(c_idx);
118
119            // c_i = (1 - eta) * c_i + eta * x_i
120            c_i.zip_mut_with(&x_i, |c_val, &x_val| {
121                *c_val = (1.0 - eta) * (*c_val) + eta * x_val;
122            });
123        }
124
125        Ok(())
126    }
127
128    fn predict_labels(&self, x: &Array2<f64>) -> Result<Array1<usize>, IncrementalError> {
129        self.validate_batch(x)?;
130
131        let centroids = match &self.centroids {
132            Some(c) => c,
133            None => return Err(IncrementalError::EmptyBatch),
134        };
135
136        let mut predictions = Vec::with_capacity(x.nrows());
137        for row in x.rows() {
138            let mut best_cluster = 0;
139            let mut min_dist = f64::INFINITY;
140
141            for (c_idx, c_row) in centroids.rows().into_iter().enumerate() {
142                let dist = Self::sq_euclidean(&row, &c_row);
143                if dist < min_dist {
144                    min_dist = dist;
145                    best_cluster = c_idx;
146                }
147            }
148            predictions.push(best_cluster);
149        }
150
151        Ok(Array1::from(predictions))
152    }
153}
incremental_rs/minibatch_kmeans.rs

incremental_rs/
minibatch_kmeans.rs
