sklears_python/
clustering.rs

1//! Python bindings for clustering algorithms
2//!
3//! This module provides Python bindings for sklears clustering algorithms,
4//! offering scikit-learn compatible interfaces with performance improvements.
5
6use crate::linear::common::{core_array2_to_py, pyarray_to_core_array2, PyValueError};
7use numpy::{PyArray1, PyArray2, PyReadonlyArray2};
8use pyo3::prelude::*;
9use scirs2_core::ndarray::{Array1, Array2};
10use sklears_clustering::kmeans::KMeansFitted;
11use sklears_clustering::{KMeans, KMeansConfig, KMeansInit, DBSCAN};
12use sklears_core::traits::{Fit, Predict, Trained};
13
14/// K-Means clustering.
15///
16/// Partitions data into `n_clusters` clusters by iteratively assigning
17/// points to the nearest centroid and recomputing centroids, using
18/// K-means++ initialization for the starting centroids.
19#[pyclass(name = "KMeans")]
20pub struct PyKMeans {
21    n_clusters: usize,
22    max_iter: usize,
23    tol: f64,
24    random_state: Option<u64>,
25    fitted: Option<KMeansFitted>,
26}
27
28impl PyKMeans {
29    fn config(&self) -> KMeansConfig {
30        KMeansConfig {
31            n_clusters: self.n_clusters,
32            init: KMeansInit::KMeansPlusPlus,
33            max_iter: self.max_iter,
34            tolerance: self.tol,
35            random_seed: self.random_state,
36        }
37    }
38
39    /// Core fit logic operating on plain `ndarray` types (no PyO3
40    /// dependency), so it is directly unit-testable without a live Python
41    /// interpreter. This crate builds with pyo3's `extension-module`
42    /// feature (required so the compiled `cdylib` can be imported from
43    /// Python), which means `Python::with_gil` cannot be used from a
44    /// standalone `cargo test` binary -- so all `#[cfg(test)]` coverage in
45    /// this file goes through `*_core` helpers like this one instead of
46    /// exercising the `#[pymethods]` directly.
47    ///
48    /// `pub` (rather than crate-private) so that
49    /// `benches/core_helpers_benchmarks.rs` -- which compiles as a
50    /// separate crate -- can call it directly for the same reason; this is
51    /// a minimal exposed surface for benchmarking, not part of the stable
52    /// Python-facing API.
53    pub fn fit_core(&mut self, x: &Array2<f64>) -> PyResult<()> {
54        let dummy_y = Array1::<f64>::zeros(x.nrows());
55        let fitted = KMeans::new(self.config())
56            .fit(x, &dummy_y)
57            .map_err(|e| PyValueError::new_err(format!("Failed to fit KMeans: {e}")))?;
58        self.fitted = Some(fitted);
59        Ok(())
60    }
61
62    /// Core predict logic; see `fit_core` for why this is split out and
63    /// `pub` (exposed for benchmarking from `benches/`, not part of the
64    /// stable Python-facing API).
65    pub fn predict_core(&self, x: &Array2<f64>) -> PyResult<Vec<i32>> {
66        let fitted = self
67            .fitted
68            .as_ref()
69            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
70        fitted
71            .predict(x)
72            .map_err(|e| PyValueError::new_err(format!("Prediction failed: {e}")))
73    }
74}
75
76#[pymethods]
77impl PyKMeans {
78    /// `pub` in addition to being reachable from Python via `#[new]`, so
79    /// `benches/core_helpers_benchmarks.rs` (a separate crate) can
80    /// construct instances to call `fit_core`/`predict_core` on; not part
81    /// of the stable Python-facing API surface.
82    #[new]
83    #[pyo3(signature = (n_clusters=8, max_iter=300, tol=1e-4, random_state=None))]
84    pub fn new(n_clusters: usize, max_iter: usize, tol: f64, random_state: Option<u64>) -> Self {
85        Self {
86            n_clusters,
87            max_iter,
88            tol,
89            random_state,
90            fitted: None,
91        }
92    }
93
94    /// Compute K-means clustering from training data.
95    fn fit(&mut self, x: PyReadonlyArray2<f64>) -> PyResult<()> {
96        let x_arr = pyarray_to_core_array2(x)?;
97        self.fit_core(&x_arr)
98    }
99
100    /// Predict the closest cluster each sample in `x` belongs to.
101    fn predict(&self, py: Python<'_>, x: PyReadonlyArray2<f64>) -> PyResult<Py<PyArray1<i32>>> {
102        let x_arr = pyarray_to_core_array2(x)?;
103        let labels = self.predict_core(&x_arr)?;
104        Ok(PyArray1::from_vec(py, labels).unbind())
105    }
106
107    /// Fit the model to `x`, then return the cluster label assigned to
108    /// each training sample.
109    fn fit_predict(
110        &mut self,
111        py: Python<'_>,
112        x: PyReadonlyArray2<f64>,
113    ) -> PyResult<Py<PyArray1<i32>>> {
114        let x_arr = pyarray_to_core_array2(x)?;
115        self.fit_core(&x_arr)?;
116        let fitted = self
117            .fitted
118            .as_ref()
119            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
120        Ok(PyArray1::from_vec(py, fitted.labels.clone()).unbind())
121    }
122
123    /// Cluster labels for the training data (available after `fit`).
124    #[getter]
125    fn labels_(&self, py: Python<'_>) -> PyResult<Py<PyArray1<i32>>> {
126        let fitted = self
127            .fitted
128            .as_ref()
129            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
130        Ok(PyArray1::from_vec(py, fitted.labels.clone()).unbind())
131    }
132
133    /// Coordinates of cluster centers (available after `fit`).
134    #[getter]
135    fn cluster_centers_(&self, py: Python<'_>) -> PyResult<Py<PyArray2<f64>>> {
136        let fitted = self
137            .fitted
138            .as_ref()
139            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
140        core_array2_to_py(py, &fitted.centroids)
141    }
142
143    /// Sum of squared distances of samples to their closest cluster center
144    /// (available after `fit`).
145    #[getter]
146    fn inertia_(&self) -> PyResult<f64> {
147        let fitted = self
148            .fitted
149            .as_ref()
150            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
151        Ok(fitted.inertia)
152    }
153
154    /// Number of iterations run before convergence (available after
155    /// `fit`).
156    #[getter]
157    fn n_iter_(&self) -> PyResult<usize> {
158        let fitted = self
159            .fitted
160            .as_ref()
161            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit() first."))?;
162        Ok(fitted.n_iterations)
163    }
164
165    fn __repr__(&self) -> String {
166        format!(
167            "KMeans(n_clusters={}, fitted={})",
168            self.n_clusters,
169            self.fitted.is_some()
170        )
171    }
172}
173
174/// DBSCAN (Density-Based Spatial Clustering of Applications with Noise).
175///
176/// Like scikit-learn's implementation, DBSCAN is transductive: it has no
177/// `predict()` for unseen data. Use `fit_predict()` to cluster the data it
178/// is fit on, and the `labels_` attribute to retrieve the result again
179/// afterwards (`-1` marks noise points).
180#[pyclass(name = "DBSCAN")]
181pub struct PyDBSCAN {
182    eps: f64,
183    min_samples: usize,
184    fitted: Option<DBSCAN<Trained>>,
185}
186
187impl PyDBSCAN {
188    /// Core fit_predict logic; see `PyKMeans::fit_core` for why this is
189    /// split out from the `#[pymethods]` wrapper and `pub` (exposed for
190    /// benchmarking from `benches/`, not part of the stable Python-facing
191    /// API).
192    pub fn fit_predict_core(&mut self, x: &Array2<f64>) -> PyResult<Vec<i32>> {
193        let model = DBSCAN::new().eps(self.eps).min_samples(self.min_samples);
194        let fitted = model
195            .fit(x, &())
196            .map_err(|e| PyValueError::new_err(format!("Failed to fit DBSCAN: {e}")))?;
197        let labels = fitted.labels().to_vec();
198        self.fitted = Some(fitted);
199        Ok(labels)
200    }
201
202    /// Core labels_ logic; see `PyKMeans::fit_core` for why this is split
203    /// out from the `#[pymethods]` wrapper.
204    fn labels_core(&self) -> PyResult<Vec<i32>> {
205        let fitted = self
206            .fitted
207            .as_ref()
208            .ok_or_else(|| PyValueError::new_err("Model not fitted. Call fit_predict() first."))?;
209        Ok(fitted.labels().to_vec())
210    }
211}
212
213#[pymethods]
214impl PyDBSCAN {
215    /// `pub` in addition to being reachable from Python via `#[new]`, so
216    /// `benches/core_helpers_benchmarks.rs` (a separate crate) can
217    /// construct instances to call `fit_predict_core` on; not part of the
218    /// stable Python-facing API surface.
219    #[new]
220    #[pyo3(signature = (eps=0.5, min_samples=5))]
221    pub fn new(eps: f64, min_samples: usize) -> Self {
222        Self {
223            eps,
224            min_samples,
225            fitted: None,
226        }
227    }
228
229    /// Fit DBSCAN to `x` and return the cluster label assigned to each
230    /// sample (`-1` for noise points).
231    fn fit_predict(
232        &mut self,
233        py: Python<'_>,
234        x: PyReadonlyArray2<f64>,
235    ) -> PyResult<Py<PyArray1<i32>>> {
236        let x_arr = pyarray_to_core_array2(x)?;
237        let labels = self.fit_predict_core(&x_arr)?;
238        Ok(PyArray1::from_vec(py, labels).unbind())
239    }
240
241    /// Cluster labels for the training data (available after
242    /// `fit_predict`).
243    #[getter]
244    fn labels_(&self, py: Python<'_>) -> PyResult<Py<PyArray1<i32>>> {
245        let labels = self.labels_core()?;
246        Ok(PyArray1::from_vec(py, labels).unbind())
247    }
248
249    fn __repr__(&self) -> String {
250        format!(
251            "DBSCAN(eps={}, min_samples={}, fitted={})",
252            self.eps,
253            self.min_samples,
254            self.fitted.is_some()
255        )
256    }
257}
258
259#[cfg(test)]
260mod tests {
261    use super::*;
262
263    /// Two well-separated blobs: 3 points near (0, 0) and 3 points near
264    /// (10, 10). Any working clustering algorithm run with 2 clusters
265    /// should separate these into 2 groups.
266    fn two_blob_data() -> Array2<f64> {
267        Array2::from_shape_vec(
268            (6, 2),
269            vec![
270                0.0, 0.0, 0.1, 0.1, 0.2, -0.1, 10.0, 10.0, 10.1, 10.1, 9.9, 10.2,
271            ],
272        )
273        .expect("shape matches data length")
274    }
275
276    #[test]
277    fn kmeans_fit_predict_round_trip_finds_two_clusters() {
278        let mut model = PyKMeans::new(2, 300, 1e-4, Some(42));
279        let x = two_blob_data();
280
281        model
282            .fit_core(&x)
283            .expect("fit should succeed on well-separated data");
284        let labels = model
285            .predict_core(&x)
286            .expect("predict should succeed after fit");
287
288        assert_eq!(labels.len(), 6);
289
290        let distinct: std::collections::HashSet<i32> = labels.iter().copied().collect();
291        assert_eq!(
292            distinct.len(),
293            2,
294            "expected exactly 2 distinct cluster labels, got {distinct:?}"
295        );
296
297        // The first 3 points must share one label and the last 3 must
298        // share a different label -- the old stub (always predicting [0])
299        // would fail this.
300        assert_eq!(labels[0], labels[1]);
301        assert_eq!(labels[1], labels[2]);
302        assert_eq!(labels[3], labels[4]);
303        assert_eq!(labels[4], labels[5]);
304        assert_ne!(labels[0], labels[3]);
305    }
306
307    #[test]
308    fn kmeans_fitted_attributes_work_post_fit() {
309        let mut model = PyKMeans::new(2, 300, 1e-4, Some(42));
310        let x = two_blob_data();
311        model.fit_core(&x).expect("fit should succeed");
312
313        let fitted = model.fitted.as_ref().expect("model was just fitted");
314        assert_eq!(fitted.labels.len(), 6);
315        assert_eq!(fitted.centroids.nrows(), 2);
316        assert_eq!(fitted.centroids.ncols(), 2);
317        assert!(fitted.inertia >= 0.0);
318        assert!(fitted.n_iterations >= 1);
319    }
320
321    #[test]
322    fn kmeans_predict_before_fit_errors_instead_of_panicking() {
323        let model = PyKMeans::new(2, 300, 1e-4, Some(42));
324        let x = two_blob_data();
325        assert!(model.predict_core(&x).is_err());
326    }
327
328    #[test]
329    fn dbscan_fit_predict_finds_more_than_one_group() {
330        let mut model = PyDBSCAN::new(1.0, 2);
331        let x = two_blob_data();
332
333        let labels = model
334            .fit_predict_core(&x)
335            .expect("fit_predict should succeed on well-separated data");
336
337        assert_eq!(labels.len(), 6);
338        assert!(
339            !labels.iter().all(|&l| l == 0),
340            "labels must not all be 0 (this is what the old stub always returned): {labels:?}"
341        );
342
343        let distinct: std::collections::HashSet<i32> = labels.iter().copied().collect();
344        assert!(
345            distinct.len() >= 2,
346            "expected at least 2 distinct labels/groups, got {distinct:?}"
347        );
348    }
349
350    #[test]
351    fn dbscan_labels_getter_errors_before_fit_and_works_after() {
352        let mut model = PyDBSCAN::new(1.0, 2);
353        assert!(
354            model.labels_core().is_err(),
355            "labels_ must error before fit_predict, not panic"
356        );
357
358        let x = two_blob_data();
359        model
360            .fit_predict_core(&x)
361            .expect("fit_predict should succeed");
362
363        let labels = model.labels_core().expect("labels_ should work after fit");
364        assert_eq!(labels.len(), 6);
365    }
366}
sklears_python/clustering.rs

sklears_python/
clustering.rs
