hipvs/cluster/kmeans/
params.rs

1/*
2 * Copyright (c) 2025, NVIDIA CORPORATION.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use crate::distance_type::DistanceType;
18use crate::error::{check_cuvs, Result};
19use std::fmt;
20use std::io::{stderr, Write};
21
22pub struct Params(pub ffi::cuvsKMeansParams_t);
23
24impl Params {
25    /// Returns a new Params
26    pub fn new() -> Result<Params> {
27        unsafe {
28            let mut params = std::mem::MaybeUninit::<ffi::cuvsKMeansParams_t>::uninit();
29            check_cuvs(ffi::cuvsKMeansParamsCreate(params.as_mut_ptr()))?;
30            Ok(Params(params.assume_init()))
31        }
32    }
33
34    /// DistanceType to use for fitting kmeans
35    pub fn set_metric(self, metric: DistanceType) -> Params {
36        unsafe {
37            (*self.0).metric = metric;
38        }
39        self
40    }
41
42    /// The number of clusters to form as well as the number of centroids to generate (default:8).
43    pub fn set_n_clusters(self, n_clusters: i32) -> Params {
44        unsafe {
45            (*self.0).n_clusters = n_clusters;
46        }
47        self
48    }
49
50    /// Maximum number of iterations of the k-means algorithm for a single run.
51    pub fn set_max_iter(self, max_iter: i32) -> Params {
52        unsafe {
53            (*self.0).max_iter = max_iter;
54        }
55        self
56    }
57
58    /// Relative tolerance with regards to inertia to declare convergence.
59    pub fn set_tol(self, tol: f64) -> Params {
60        unsafe {
61            (*self.0).tol = tol;
62        }
63        self
64    }
65
66    /// Number of instance k-means algorithm will be run with different seeds.
67    pub fn set_n_init(self, n_init: i32) -> Params {
68        unsafe {
69            (*self.0).n_init = n_init;
70        }
71        self
72    }
73
74    /// Oversampling factor for use in the k-means|| algorithm
75    pub fn set_oversampling_factor(self, oversampling_factor: f64) -> Params {
76        unsafe {
77            (*self.0).oversampling_factor = oversampling_factor;
78        }
79        self
80    }
81
82    /**
83     * batch_samples and batch_centroids are used to tile 1NN computation which is
84     * useful to optimize/control the memory footprint
85     * Default tile is [batch_samples x n_clusters] i.e. when batch_centroids is 0
86     * then don't tile the centroids.
87     */
88    pub fn set_batch_samples(self, batch_samples: i32) -> Params {
89        unsafe {
90            (*self.0).batch_samples = batch_samples;
91        }
92        self
93    }
94    /// if 0 then batch_centroids = n_clusters
95    pub fn set_batch_centroids(self, batch_centroids: i32) -> Params {
96        unsafe {
97            (*self.0).batch_centroids = batch_centroids;
98        }
99        self
100    }
101
102    /// Whether to use hierarchical (balanced) kmeans or not
103    pub fn set_hierarchical(self, hierarchical: bool) -> Params {
104        unsafe {
105            (*self.0).hierarchical = hierarchical;
106        }
107        self
108    }
109
110    /// For hierarchical k-means , defines the number of training iterations
111    pub fn set_hierarchical_n_iters(self, hierarchical_n_iters: i32) -> Params {
112        unsafe {
113            (*self.0).hierarchical_n_iters = hierarchical_n_iters;
114        }
115        self
116    }
117}
118
119impl fmt::Debug for Params {
120    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
121        // custom debug trait here, default value will show the pointer address
122        // for the inner params object which isn't that useful.
123        write!(f, "Params({:?})", unsafe { *self.0 })
124    }
125}
126
127impl Drop for Params {
128    fn drop(&mut self) {
129        if let Err(e) = check_cuvs(unsafe { ffi::cuvsKMeansParamsDestroy(self.0) }) {
130            write!(stderr(), "failed to call cuvsKMeansParamsDestroy {:?}", e)
131                .expect("failed to write to stderr");
132        }
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[test]
141    fn test_params() {
142        let params = Params::new()
143            .unwrap()
144            .set_n_clusters(128)
145            .set_hierarchical(true);
146
147        unsafe {
148            assert_eq!((*params.0).n_clusters, 128);
149            assert_eq!((*params.0).hierarchical, true);
150        }
151    }
152}
hipvs/cluster/kmeans/params.rs

hipvs/cluster/kmeans/
params.rs
