mojolearn.KMeans#

class mojolearn.KMeans(n_clusters=8, init='k-means++', n_init=1, max_iter=300, tol=0.0001, random_state=0, init_centroids=None, metric='euclidean', oversampling_factor=2.0)[source]#

Bases: NumericModeMixin

K-means clustering.

Partitions the rows of X into n_clusters groups by minimizing the within-cluster sum of squared distances (Lloyd iterations after k-means++ or random seeding). X is read as float32, C order; a NaN or an infinity in X, sample_weight or init_centroids raises ValueError.

The defaults follow cuVS, not scikit-learn, and one of them changes results: n_init is 1 here and 10 in scikit-learn. To compare with scikit-learn, set n_init to the same value on both.

Parameters:
  • n_clusters (int, default 8) – Number of clusters. Must not exceed the number of rows.

  • init ({‘k-means++’, ‘random’, ‘array’}, array or callable, default ‘k-means++’) – How the starting centers are chosen. ‘array’ takes them from init_centroids. An array of shape (n_clusters, n_features) is used directly as the starting centers (init_centroids is then ignored). A callable is called as init(X, n_clusters, random_state) and must return such an array.

  • n_init (int, default 1) – Number of restarts. Each restart draws different starting centers; the one with the lowest final inertia is kept.

  • max_iter (int, default 300) – Maximum number of iterations per restart.

  • tol (float, default 1e-4) – Stop when the centers move less than this between iterations.

  • random_state (int, default 0) – Seed for the starting centers.

  • init_centroids (array of shape (n_clusters, n_features), default None) – Starting centers for init='array'. Required in that case.

  • metric ({‘euclidean’, ‘l2_expanded’, ‘l2_sqrt_expanded’}, default ‘euclidean’) – ‘euclidean’ and ‘l2_expanded’ are the same: squared Euclidean distance, so inertia_ and transform use squared distances. ‘l2_sqrt_expanded’ uses the Euclidean distance itself, which changes inertia_ and transform but not the labels. Any other value raises ValueError.

  • oversampling_factor (float, default 2.0) – Selects the k-means++ seeding algorithm: 0.0 runs classic sequential k-means++, any positive value runs scalable k-means|| with this oversampling. Negative values raise ValueError. Only used with init='k-means++'.

Variables:
  • cluster_centers_ (Array of shape (n_clusters, n_features), float32) – The fitted centers.

  • labels_ (Array of shape (n_samples,), int32) – The cluster of each training row, assigned against the final centers.

  • inertia_ (float) – Weighted sum of distances (under metric) from each row to its final center. With n_init > 1, the kept restart’s value.

  • n_iter_ (int) – Number of iterations run.

  • weight_scale_ (sum_scale_,) – The fixed-point accumulator scales chosen for your data. Reproducing a result exactly needs them.

  • n_features_in_ (int) – Number of features seen during fit.

Notes

Attributes are mojolearn Array`s; `np.asarray(model.labels_) gives a zero-copy NumPy view.

Examples

>>> import numpy as np, mojolearn
>>> X = np.random.default_rng(0).random((1000, 8), dtype=np.float32)
>>> model = mojolearn.KMeans(n_clusters=4, n_init=3).fit(X)
>>> labels = model.predict(X)
>>> distances = model.transform(X)
__init__(n_clusters=8, init='k-means++', n_init=1, max_iter=300, tol=0.0001, random_state=0, init_centroids=None, metric='euclidean', oversampling_factor=2.0)[source]#

Methods

fit(X[, y, sample_weight])

Compute the k-means clustering of X.

fit_predict(X[, y, sample_weight])

Fit the model and return labels_, the cluster of each row of X.

fit_transform(X[, y, sample_weight])

Fit the model, then return transform(X).

get_params([deep])

load(path)

Load a model written by save.

numeric_mode_used()

The numeric mode this estimator runs in.

predict(X)

Return the index of the nearest fitted center for each row of X.

save(path)

Save the fitted model to path as an .npz file.

set_params(**params)

Set constructor parameters.

transform(X)

Return the distance from each row of X to each fitted center.

vendor_used()

The GPU backend this estimator runs on.

fit(X, y=None, sample_weight=None)[source]#

Compute the k-means clustering of X.

Parameters:
  • X (array-like of shape (n_samples, n_features)) – Training data, read as float32.

  • y (ignored) – Present for scikit-learn API compatibility.

  • sample_weight (array-like of shape (n_samples,), default None) – Weight of each row. None weighs every row equally.

Returns:

self (KMeans) – The fitted estimator.

Raises:
  • ValueError – For an unknown init or metric, n_clusters greater than the number of rows, init='array' without init_centroids, starting centers of the wrong shape, a sample_weight of the wrong length, or a NaN or infinity in any input.

  • TypeError – If oversampling_factor is not a number.

fit_predict(X, y=None, sample_weight=None)[source]#

Fit the model and return labels_, the cluster of each row of X.

Returns:

labels (Array of shape (n_samples,), int32)

predict(X)[source]#

Return the index of the nearest fitted center for each row of X.

Uses the same assignment as fit, so predict on the training rows returns labels_ exactly. A tie goes to the lowest center index. On a CPU-only install, predict runs on the CPU.

Parameters:

X (array-like of shape (n_samples, n_features)) – Read as float32. Must have the fitted number of features and no NaN or infinity.

Returns:

labels (Array of shape (n_samples,), int32)

Raises:
  • RuntimeError – If the model is not fitted.

  • ValueError – For a feature count other than the fit’s, or a non-finite value.

transform(X)[source]#

Return the distance from each row of X to each fitted center.

Distances use the model’s metric: with the default ‘euclidean’ they are squared distances, with ‘l2_sqrt_expanded’ they are Euclidean distances. scikit-learn’s transform always returns Euclidean distances; fit with metric='l2_sqrt_expanded' to get the same meaning. transform(X)[i, predict(X)[i]] is exactly the minimum of row i. On a CPU-only install, transform runs on the CPU.

Parameters:

X (array-like of shape (n_samples, n_features)) – Read as float32.

Returns:

distances (Array of shape (n_samples, n_clusters), float32)

Raises:
  • RuntimeError – If the model is not fitted.

  • ValueError – For a feature count other than the fit’s, or a non-finite value.

fit_transform(X, y=None, sample_weight=None)[source]#

Fit the model, then return transform(X).

Returns:

distances (Array of shape (n_samples, n_clusters), float32)

save(path)[source]#

Save the fitted model to path as an .npz file.

The file holds the centers, the training labels, the metric, the init method and the fitted scalars. Saving the same model twice gives the same file. A model fitted on a GPU can be loaded with load, or with mojolearn.host_model(path) to predict and transform on a machine with no GPU, giving the same answers as the GPU.

init_centroids is not saved: it is an input to a fit, not part of the fitted model. A loaded init='array' model that is fitted again raises ValueError for the missing init_centroids.

Parameters:

path (str or path-like) – Destination file.

Raises:
  • RuntimeError – If the model is not fitted.

  • ValueError – If metric or init is not one the model supports.

classmethod load(path)[source]#

Load a model written by save.

Arrays keep their saved dtypes. The result supports predict and transform, and carries labels_, inertia_, n_iter_, sum_scale_ and weight_scale_ as the fit set them.

Parameters:

path (str or path-like) – File written by KMeans.save.

Returns:

model (KMeans)

Raises:

ValueError – If the file is of another format or estimator, is truncated or inconsistent (field counts, center shape, label count), or names a metric or init the class does not support.