cuvs/
kmeans.rs

1use std::ptr;
2
3use crate::{
4    brute_force::{matrix_layout, validate_matrix_binding},
5    device::{DeviceBuffer, MatrixLayout, OwnedMatrix, ValidatedDeviceTensor},
6    dlpack::DlpackTensor,
7    elements::KMeansElement,
8    error::{Result, check_cuvs, invalid_input, invalid_state, require_nonzero},
9    params_handle::ParamsHandle,
10    resources::CuvsResources,
11};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum KMeansInit {
15    KMeansPlusPlus,
16    Random,
17    Array,
18}
19
20impl KMeansInit {
21    fn raw(self) -> cuvs_sys::cuvsKMeansInitMethod {
22        match self {
23            KMeansInit::KMeansPlusPlus => cuvs_sys::cuvsKMeansInitMethod::KMeansPlusPlus,
24            KMeansInit::Random => cuvs_sys::cuvsKMeansInitMethod::Random,
25            KMeansInit::Array => cuvs_sys::cuvsKMeansInitMethod::Array,
26        }
27    }
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum KMeansMetric {
32    L2Expanded,
33    L2SqrtExpanded,
34}
35
36impl KMeansMetric {
37    fn raw(self) -> cuvs_sys::cuvsDistanceType {
38        match self {
39            Self::L2Expanded => cuvs_sys::cuvsDistanceType::L2Expanded,
40            Self::L2SqrtExpanded => cuvs_sys::cuvsDistanceType::L2SqrtExpanded,
41        }
42    }
43}
44
45#[derive(Debug, Clone)]
46pub struct KMeansParams {
47    pub n_clusters: usize,
48    pub max_iter: i32,
49    pub tol: f64,
50    pub n_init: i32,
51    pub init: KMeansInit,
52    pub metric: KMeansMetric,
53}
54
55impl KMeansParams {
56    pub fn new(n_clusters: usize) -> Self {
57        Self {
58            n_clusters,
59            max_iter: 100,
60            tol: 1e-4,
61            n_init: 1,
62            init: KMeansInit::KMeansPlusPlus,
63            metric: KMeansMetric::L2Expanded,
64        }
65    }
66}
67
68#[derive(Debug, Clone, Copy)]
69pub struct KMeansFit {
70    pub inertia: f64,
71    pub n_iter: i32,
72}
73
74#[derive(Debug, Clone, Copy)]
75pub struct KMeansPrediction {
76    pub rows: usize,
77    pub inertia: f64,
78}
79
80#[derive(Debug)]
81pub struct KMeansModel<T: KMeansElement> {
82    resources: CuvsResources,
83    params: KMeansParams,
84    centroids: OwnedMatrix<T>,
85    n_features: usize,
86    fit: KMeansFit,
87}
88
89fn create_params(params: &KMeansParams, operation: &'static str) -> Result<ParamsHandle<cuvs_sys::cuvsKMeansParams>> {
90    validate_params(params, operation)?;
91    let handle = ParamsHandle::new(
92        cuvs_sys::cuvsKMeansParamsCreate,
93        cuvs_sys::cuvsKMeansParamsDestroy,
94        operation,
95        "KMeans params",
96    )?;
97    unsafe {
98        let native = handle.as_ptr();
99        (*native).metric = params.metric.raw();
100        (*native).n_clusters = params.n_clusters as i32;
101        (*native).init = params.init.raw();
102        (*native).max_iter = params.max_iter;
103        (*native).tol = params.tol;
104        (*native).n_init = params.n_init;
105        (*native).hierarchical = false;
106        (*native).device_buffer_samples = 0;
107        (*native).init_size = 0;
108    }
109    Ok(handle)
110}
111
112impl<T: KMeansElement> KMeansModel<T> {
113    pub fn fit(
114        resources: &CuvsResources,
115        data: &ValidatedDeviceTensor<'_, T>,
116        params: KMeansParams,
117        initial_centroids: Option<DeviceBuffer<T>>,
118    ) -> Result<Self> {
119        let operation = "KMeansModel::fit";
120        validate_matrix_binding(resources, data, operation)?;
121        validate_params(&params, operation)?;
122        let n_features = data.columns()?;
123        let initial_centroids = initial_centroids
124            .map(|centroids| {
125                OwnedMatrix::from_buffer(
126                    resources,
127                    centroids,
128                    params.n_clusters,
129                    n_features,
130                    MatrixLayout::RowMajor,
131                    operation,
132                )
133            })
134            .transpose()?;
135        Self::fit_with_initial_centroids(resources, data, params, initial_centroids)
136    }
137
138    /// Fit KMeans with an optional owned row-major centroid relation.
139    ///
140    /// Consumed centroids satisfy the exclusive mutable ownership required
141    /// by `KMeansInit::Array`.
142    pub fn fit_with_initial_centroids(
143        resources: &CuvsResources,
144        data: &ValidatedDeviceTensor<'_, T>,
145        params: KMeansParams,
146        initial_centroids: Option<OwnedMatrix<T>>,
147    ) -> Result<Self> {
148        let operation = "KMeansModel::fit_with_initial_centroids";
149        validate_matrix_binding(resources, data, operation)?;
150        if matrix_layout(data, operation)? != MatrixLayout::RowMajor {
151            return Err(invalid_input(operation, "KMeans input must be contiguous row-major"));
152        }
153        validate_params(&params, operation)?;
154        let rows = data.rows()?;
155        validate_row_count(rows, operation)?;
156        let n_features = data.columns()?;
157        validate_feature_count(n_features, operation)?;
158        if params.n_clusters > rows {
159            return Err(invalid_input(
160                operation,
161                format!("n_clusters ({}) must be <= rows ({rows})", params.n_clusters),
162            ));
163        }
164        if matches!(params.init, KMeansInit::Array) != initial_centroids.is_some() {
165            return Err(invalid_input(
166                operation,
167                "initial centroids must be supplied exactly when init is KMeansInit::Array",
168            ));
169        }
170        let mut centroids = match initial_centroids {
171            Some(centroids) => {
172                if !centroids.is_input_validated() {
173                    return Err(invalid_input(operation, "initial centroids must carry finite-input validation"));
174                }
175                if centroids.rows() != params.n_clusters || centroids.columns() != n_features {
176                    return Err(invalid_input(
177                        operation,
178                        format!(
179                            "initial centroid shape [{}, {}] must equal [{}, {n_features}]",
180                            centroids.rows(),
181                            centroids.columns(),
182                            params.n_clusters
183                        ),
184                    ));
185                }
186                if centroids.layout() != MatrixLayout::RowMajor {
187                    return Err(invalid_input(operation, "initial centroids must be a contiguous row-major relation"));
188                }
189                if centroids.binding_id() != resources.binding_id() {
190                    return Err(invalid_input(
191                        operation,
192                        "initial centroids and cuVS resources must use the same exact runtime binding",
193                    )
194                    .with_context("matrix_binding_id", centroids.binding_id())
195                    .with_context("resource_binding_id", resources.binding_id()));
196                }
197                centroids
198            },
199            None => OwnedMatrix::new(resources, params.n_clusters, n_features, MatrixLayout::RowMajor)?,
200        };
201        let mut centroids_view = centroids.as_tensor_mut(resources)?;
202        let mut inertia = 0.0;
203        let mut n_iter = 0;
204        resources.with_native(operation, |native| {
205            let mut params_handle = create_params(&params, operation)?;
206            let mut data_tensor = DlpackTensor::from_const(data.as_tensor());
207            let mut centroids_tensor = DlpackTensor::from_mut(&mut centroids_view);
208            let call_result = unsafe {
209                check_cuvs(
210                    cuvs_sys::cuvsKMeansFit(
211                        native,
212                        params_handle.as_ptr(),
213                        data_tensor.as_mut_ptr(),
214                        ptr::null_mut(),
215                        centroids_tensor.as_mut_ptr(),
216                        &mut inertia,
217                        &mut n_iter,
218                    ),
219                    operation,
220                )
221            };
222            finish_params_call(call_result, params_handle.close(operation))
223        })?;
224        if !inertia.is_finite() || inertia < 0.0 {
225            return Err(invalid_state(operation, "native KMeans fit returned non-finite or negative inertia"));
226        }
227        if n_iter < 0 || n_iter > params.max_iter {
228            return Err(invalid_state(
229                operation,
230                format!("native KMeans fit returned invalid iteration count {n_iter}"),
231            ));
232        }
233        centroids.validate_native_output(resources, operation)?;
234
235        Ok(Self { resources: resources.retain(), params, centroids, n_features, fit: KMeansFit { inertia, n_iter } })
236    }
237
238    pub fn predict_into(
239        &self,
240        data: &ValidatedDeviceTensor<'_, T>,
241        labels: &mut DeviceBuffer<i32>,
242    ) -> Result<KMeansPrediction> {
243        let operation = "KMeansModel::predict_into";
244        validate_matrix_binding(&self.resources, data, operation)?;
245        if matrix_layout(data, operation)? != MatrixLayout::RowMajor {
246            return Err(invalid_input(operation, "KMeans input must be contiguous row-major"));
247        }
248        if data.columns()? != self.n_features {
249            return Err(invalid_input(
250                operation,
251                format!("input columns {} must match model features {}", data.columns()?, self.n_features),
252            ));
253        }
254        validate_feature_count(self.n_features, operation)?;
255        let rows = data.rows()?;
256        validate_row_count(rows, operation)?;
257        if labels.len() != rows {
258            return Err(invalid_input(
259                operation,
260                format!("labels contain {} elements, expected one per input row ({rows})", labels.len()),
261            ));
262        }
263
264        let mut labels_view = labels.as_tensor_mut(&self.resources, &[rows], None, 0)?;
265        let centroids_view = self.centroids.as_validated_tensor(&self.resources)?;
266        let mut inertia = 0.0;
267        self.resources.with_native(operation, |native| {
268            let mut params_handle = create_params(&self.params, operation)?;
269            let mut data_tensor = DlpackTensor::from_const(data.as_tensor());
270            let mut centroids_tensor = DlpackTensor::from_const(centroids_view.as_tensor());
271            let mut labels_tensor = DlpackTensor::from_mut(&mut labels_view);
272            let call_result = unsafe {
273                check_cuvs(
274                    cuvs_sys::cuvsKMeansPredict(
275                        native,
276                        params_handle.as_ptr(),
277                        data_tensor.as_mut_ptr(),
278                        ptr::null_mut(),
279                        centroids_tensor.as_mut_ptr(),
280                        labels_tensor.as_mut_ptr(),
281                        false,
282                        &mut inertia,
283                    ),
284                    operation,
285                )
286            };
287            finish_params_call(call_result, params_handle.close(operation))
288        })?;
289        if !inertia.is_finite() || inertia < 0.0 {
290            return Err(invalid_state(operation, "native KMeans prediction returned non-finite or negative inertia"));
291        }
292        Ok(KMeansPrediction { rows, inertia })
293    }
294
295    pub fn params(&self) -> &KMeansParams {
296        &self.params
297    }
298
299    pub fn fit_result(&self) -> KMeansFit {
300        self.fit
301    }
302
303    pub fn centroids(&self) -> &OwnedMatrix<T> {
304        &self.centroids
305    }
306
307    pub fn n_features(&self) -> usize {
308        self.n_features
309    }
310}
311
312fn validate_params(params: &KMeansParams, operation: &'static str) -> Result<()> {
313    require_nonzero(params.n_clusters, operation, "n_clusters")?;
314    if params.n_clusters > i32::MAX as usize {
315        return Err(invalid_input(operation, "n_clusters exceeds the cuVS i32 parameter range"));
316    }
317    if params.max_iter <= 0 {
318        return Err(invalid_input(operation, "max_iter must be greater than zero"));
319    }
320    if params.n_init <= 0 {
321        return Err(invalid_input(operation, "n_init must be greater than zero"));
322    }
323    if !params.tol.is_finite() || params.tol < 0.0 {
324        return Err(invalid_input(operation, "tol must be finite and non-negative"));
325    }
326    Ok(())
327}
328
329fn validate_row_count(rows: usize, operation: &'static str) -> Result<()> {
330    if rows > i32::MAX as usize {
331        return Err(invalid_input(operation, "KMeans row count exceeds the cuVS i32 range"));
332    }
333    Ok(())
334}
335
336fn validate_feature_count(columns: usize, operation: &'static str) -> Result<()> {
337    if columns > i32::MAX as usize {
338        return Err(invalid_input(operation, "KMeans feature count exceeds the cuVS prediction i32 range"));
339    }
340    Ok(())
341}
342
343fn finish_params_call(call: Result<()>, close: Result<()>) -> Result<()> {
344    match (call, close) {
345        (Err(error), _) => Err(error),
346        (Ok(()), close) => close,
347    }
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353
354    #[test]
355    fn row_count_accepts_cuvs_i32_maximum() {
356        validate_row_count(i32::MAX as usize, "test").expect("i32::MAX rows are representable");
357    }
358
359    #[test]
360    fn row_count_rejects_one_past_cuvs_i32_maximum() {
361        let error = validate_row_count(i32::MAX as usize + 1, "test").expect_err("row count must fit i32");
362        assert!(error.message().contains("row count exceeds"));
363    }
364
365    #[test]
366    fn feature_count_rejects_one_past_the_prediction_i32_maximum() {
367        let error =
368            validate_feature_count(i32::MAX as usize + 1, "test").expect_err("feature count must fit prediction i32");
369        assert!(error.message().contains("feature count exceeds"));
370    }
371}
cuvs/kmeans.rs

cuvs/
kmeans.rs
