iris/ml/
mod.rs

1use crate::error::{IrisError, Result};
2use burn::tensor::{Int, Tensor, TensorData, backend::Backend};
3
4/// Classical Machine Learning algorithms powered by Burn.
5pub struct KMeans<B: Backend> {
6    pub k: usize,
7    pub max_iter: usize,
8    pub centroids: Option<Tensor<B, 2>>,
9}
10
11impl<B: Backend> KMeans<B> {
12    /// Creates a new `KMeans` model.
13    #[must_use]
14    pub fn new(k: usize, max_iter: usize) -> Self {
15        Self {
16            k,
17            max_iter,
18            centroids: None,
19        }
20    }
21
22    /// Fits K-Means on the input data of shape [N, D] (N samples, D features).
23    pub fn fit(&mut self, data: &Tensor<B, 2>) -> Result<()> {
24        let dims = data.dims();
25        let n = dims[0];
26        let d = dims[1];
27
28        if n < self.k {
29            return Err(IrisError::InvalidParameter(
30                "Number of data points must be >= K".into(),
31            ));
32        }
33
34        let device = &data.device();
35
36        // 1. Initialize centroids by picking the first K points
37        let initial_centroids = data.clone().slice([0..self.k, 0..d]);
38        let mut centroids = initial_centroids;
39
40        // 2. Iterate
41        for _ in 0..self.max_iter {
42            // Compute Euclidean distances between all N points and K centroids
43            // points shape: [N, 1, D], centroids shape: [1, K, D]
44            let p_unsqueezed = data.clone().unsqueeze_dim::<3>(1); // [N, 1, D]
45            let c_unsqueezed = centroids.clone().unsqueeze_dim::<3>(0); // [1, K, D]
46
47            let diff = p_unsqueezed.sub(c_unsqueezed); // [N, K, D]
48            let squared_diff = diff.powf_scalar(2.0);
49            let dists = squared_diff.sum_dim(2).squeeze::<2>(); // [N, K]
50
51            // Find closest centroid for each point
52            let assignments = dists.argmin(1).squeeze::<1>(); // [N]
53
54            // Update centroids
55            // For a pure tensor implementation, we can group and mean, or do it on CPU for exactness
56            let assignments_data = assignments.into_data();
57            let assignments_vec: Vec<i32> = assignments_data.iter::<i32>().collect();
58            let data_data = data.clone().into_data();
59            let flat_data: Vec<f32> = data_data.iter::<f32>().collect();
60
61            let mut new_centroids_data = vec![0.0f32; self.k * d];
62            let mut counts = vec![0.0f32; self.k];
63
64            for i in 0..n {
65                let cluster = assignments_vec[i] as usize;
66                counts[cluster] += 1.0;
67                for j in 0..d {
68                    new_centroids_data[cluster * d + j] += flat_data[i * d + j];
69                }
70            }
71
72            for k in 0..self.k {
73                let count = counts[k].max(1.0);
74                for j in 0..d {
75                    new_centroids_data[k * d + j] /= count;
76                }
77            }
78
79            centroids =
80                Tensor::<B, 2>::from_data(TensorData::new(new_centroids_data, [self.k, d]), device);
81        }
82
83        self.centroids = Some(centroids);
84        Ok(())
85    }
86
87    /// Predicts closest cluster assignments for inputs of shape [N, D].
88    pub fn predict(&self, data: &Tensor<B, 2>) -> Result<Tensor<B, 1, Int>> {
89        let centroids = self.centroids.as_ref().ok_or_else(|| {
90            IrisError::Generic("K-Means centroids are not initialized. Fit the model first.".into())
91        })?;
92
93        // points shape: [N, 1, D], centroids shape: [1, K, D]
94        let p_unsqueezed = data.clone().unsqueeze_dim::<3>(1); // [N, 1, D]
95        let c_unsqueezed = centroids.clone().unsqueeze_dim::<3>(0); // [1, K, D]
96
97        let diff = p_unsqueezed.sub(c_unsqueezed); // [N, K, D]
98        let squared_diff = diff.powf_scalar(2.0);
99        let dists = squared_diff.sum_dim(2).squeeze::<2>(); // [N, K]
100
101        let assignments = dists.argmin(1).squeeze::<1>(); // [N]
102        Ok(assignments)
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use crate::test_helpers::{TestBackend, test_device};
110
111    #[test]
112    fn test_kmeans_clustering() {
113        let device = test_device();
114        let data = Tensor::<TestBackend, 2>::from_data(
115            TensorData::new(vec![1.0f32, 1.0, 1.1, 1.1, 10.0, 10.0, 10.2, 10.2], [4, 2]),
116            &device,
117        );
118
119        let mut km = KMeans::new(2, 5);
120        km.fit(&data).unwrap();
121
122        assert!(km.centroids.is_some());
123        let assignments = km.predict(&data).unwrap();
124        assert_eq!(assignments.dims(), [4]);
125    }
126}
iris/ml/mod.rs

iris/ml/
mod.rs
