From 2a8de873db5241949929282cd9798fafaf236cf9 Mon Sep 17 00:00:00 2001 From: Tim Head Date: Tue, 31 Jan 2023 10:19:33 +0000 Subject: [PATCH] Add docstring --- python/pylibraft/pylibraft/cluster/kmeans.pyx | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/python/pylibraft/pylibraft/cluster/kmeans.pyx b/python/pylibraft/pylibraft/cluster/kmeans.pyx index 04a76a96b6..6b72df6d35 100644 --- a/python/pylibraft/pylibraft/cluster/kmeans.pyx +++ b/python/pylibraft/pylibraft/cluster/kmeans.pyx @@ -203,6 +203,34 @@ def compute_new_centroids(X, @auto_sync_handle @auto_convert_output def init_plus_plus(X, n_clusters=None, seed=None, handle=None, centroids=None): + """ + Compute initial centroids using the "kmeans++" algorithm. + + Parameters + ---------- + + X : Input CUDA array interface compliant matrix shape (m, k) + n_clusters : Number of clusters to select + seed : Controls the random sampling of centroids + centroids : Optional writable CUDA array interface compliant matrix shape + (n_clusters, k). Use instead of passing `n_clusters`. + {handle_docstring} + + Examples + -------- + + >>> import cupy as cp + >>> from pylibraft.cluster.kmeans import init_plus_plus + + >>> n_samples = 5000 + >>> n_features = 50 + >>> n_clusters = 3 + + >>> X = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + + >>> centroids = init_plus_plus(X, n_clusters) + """ if n_clusters is not None and centroids is not None: msg = ("Parameters 'n_clusters' and 'centroids' are exclusive. Only " + "pass one at a time.")