Clustering Kmeans

kmeans

CLASS	DESCRIPTION
KMeans	K-Means clustering with k-means++ initialization.

FUNCTION	DESCRIPTION
kmeans	Perform k-means clustering (functional interface).

Classes

KMeans(n_states: int, init: Literal['k-means++', 'random'] = 'k-means++', n_init: int = 1, n_iter: int = 100, tol: float = 0.0001, device: torch.device | str | None = None, dtype: torch.dtype = torch.float32)

Bases: Estimator

K-Means clustering with k-means++ initialization.

Implements the standard k-means clustering algorithm with support for multiple initializations and both random and k-means++ initialization strategies.

PARAMETER	DESCRIPTION
n_states	Number of states (clusters) to form. Must be at least 1. TYPE: int
init	Initialization method for cluster centers. TYPE: Literal["k-means++", "random"], by default "k-means++" DEFAULT: 'k-means++'
n_init	Number of random initializations to try. TYPE: int, by default 1 DEFAULT: 1
n_iter	Maximum number of iterations for a single run. TYPE: int, by default 100 DEFAULT: 100
tol	Relative tolerance for convergence. TYPE: float, by default 1e-4 DEFAULT: 0.0001
device	Device to run computations on. If None, automatically selects CUDA if available, otherwise CPU. TYPE: torch.device | str | None, by default None DEFAULT: None
dtype	Data type to use for computations. TYPE: torch.dtype, by default torch.float32 DEFAULT: float32

Properties

means : torch.Tensor Cluster centers of shape (n_states, in_features).

labels : torch.Tensor Labels of each training sample indicating cluster assignment of shape (n_samples,).

inertia : float Sum of squared distances from samples to their closest cluster center.

Examples:

Basic K-Means clustering:

>>> kmeans = KMeans(n_states=3)
>>> X = torch.randn(100, 2)
>>> labels = kmeans.fit_predict(X)
>>> labels.shape
torch.Size([100])

Multiple initializations:

>>> kmeans = KMeans(n_states=3, n_init=10)
>>> kmeans.fit(X)
>>> print(f"Best inertia: {kmeans.inertia:.2f}")

METHOD	DESCRIPTION
fit	Fit k-means clustering to input data.
predict	Predict cluster labels for input data.
fit_predict	Fit k-means and predict cluster labels.
score	Compute the negative inertia (higher is better).
get_results	Get all fitted parameters as a dictionary.

ATTRIBUTE	DESCRIPTION
means	Cluster centers of shape (n_states, in_features). TYPE: Tensor
labels	Training sample labels of shape (n_samples,). TYPE: Tensor
inertia	Sum of squared distances to nearest cluster center. TYPE: Tensor

Source code in spectre/clustering/kmeans.py

def __init__(
    self,
    n_states: int,
    init: Literal["k-means++", "random"] = "k-means++",
    n_init: int = 1,
    n_iter: int = 100,
    tol: float = 1e-4,
    device: torch.device | str | None = None,
    dtype: torch.dtype = torch.float32,
):
    super().__init__()

    if not isinstance(n_states, int):
        raise TypeError(
            f"Expected `n_states` to be an integer, got {type(n_states).__name__}."
        )
    check_in_interval(n_states, "[1, inf)")
    self.n_states = n_states
    self.out_features = n_states

    allowed_init = ["k-means++", "random"]
    if init not in allowed_init:
        raise ValueError(
            f"Expected `init` to be one of {allowed_init}, got {init}."
        )
    self.init = init

    if not isinstance(n_init, int):
        raise TypeError(
            f"Expected `n_init` to be an integer, got {type(n_init).__name__}."
        )
    check_in_interval(n_init, "[1, inf)")
    self.n_init = n_init

    if not isinstance(n_iter, int):
        raise TypeError(
            f"Expected `n_iter` to be an integer, got {type(n_iter).__name__}."
        )
    check_in_interval(n_iter, "[1, inf)")
    self.n_iter = n_iter

    if not isinstance(tol, (int, float)):
        raise TypeError(f"Expected `tol` to be a number, got {type(tol).__name__}.")
    check_in_interval(tol, "(0, inf)")
    self.tol = float(tol)

    if device is None:
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        self.device = torch.device(device)

    self.dtype = dtype

    # Internal fitted attributes
    self._means = None
    self._labels = None
    self._inertia = None

Attributes

means: torch.Tensor property

Cluster centers of shape (n_states, in_features).

labels: torch.Tensor property

Training sample labels of shape (n_samples,).

inertia: torch.Tensor property

Sum of squared distances to nearest cluster center.

Functions

fit(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> KMeans

Fit k-means clustering to input data.

PARAMETER	DESCRIPTION
X	Training data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights of shape (n_samples,). Weights are used to compute weighted cluster centers and weighted inertia. Weights are automatically normalized to sum to n_samples to preserve effective sample size. If None, all samples are weighted equally. TYPE: torch.Tensor | None, by default None DEFAULT: None
target	Target values. Not supported. TYPE: torch.Tensor | None, by default None DEFAULT: None

RETURNS	DESCRIPTION
KMeans	Returns self for method chaining.

Source code in spectre/clustering/kmeans.py

def fit(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> "KMeans":
    """
    Fit k-means clustering to input data.

    Parameters
    ----------
    X : torch.Tensor
        Training data of shape (n_samples, in_features).

    weights : torch.Tensor | None, by default None
        Sample weights of shape (n_samples,). Weights are used to compute
        weighted cluster centers and weighted inertia. Weights are automatically
        normalized to sum to n_samples to preserve effective sample size.
        If None, all samples are weighted equally.

    target : torch.Tensor | None, by default None
        Target values. Not supported.

    Returns
    -------
    KMeans
        Returns self for method chaining.
    """
    check_is_tensor(X)
    check_2d(X)

    if X.size(0) < self.n_states:
        raise ValueError(
            f"Number of samples ({X.size(0)}) must be at least "
            f"number of states ({self.n_states})."
        )

    X = X.to(device=self.device, dtype=self.dtype)

    if weights is not None:
        check_is_tensor(weights)
        check_1d(weights)
        check_same_shape(X, weights, axis=0)
        check_sample_weights(weights)
        weights = weights.to(device=self.device, dtype=self.dtype)
        # Normalize weights to sum to n_samples (maintains effective sample size)
        weights = weights * (X.shape[0] / weights.sum().clamp(min=1e-10))
    else:
        weights = torch.ones(X.shape[0], device=self.device, dtype=self.dtype)

    self.in_features = X.shape[1]

    b_inertia = torch.tensor(float("inf"), device=self.device)
    b_means = None
    b_labels = None

    for init_idx in range(self.n_init):
        mu = self._initialize(X)

        for iter_idx in range(self.n_iter):
            labels, inertia = self._predict_with_inertia(X, mu, weights)
            n_mu = self._update_means(X, labels, weights)

            # Compute weighted variance for convergence check
            weighted_var = (
                (X - X.mean(dim=0)) ** 2 * weights.unsqueeze(1)
            ).sum() / weights.sum()

            if torch.norm(n_mu - mu, dim=1).sum() <= weighted_var.mean() * self.tol:
                mu = n_mu
                break

            mu = n_mu

        if inertia < b_inertia:
            b_inertia = inertia
            b_means = mu
            b_labels = labels

    self._means = b_means
    self._labels = b_labels
    self._inertia = b_inertia

    self.is_fitted = True

    return self

predict(X: torch.Tensor) -> torch.Tensor

Predict cluster labels for input data.

PARAMETER	DESCRIPTION
X	Input data of shape (n_samples, in_features). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Predicted cluster labels of shape (n_samples,).

Source code in spectre/clustering/kmeans.py

def predict(self, X: torch.Tensor) -> torch.Tensor:
    """
    Predict cluster labels for input data.

    Parameters
    ----------
    X : torch.Tensor
        Input data of shape (n_samples, in_features).

    Returns
    -------
    torch.Tensor
        Predicted cluster labels of shape (n_samples,).
    """
    self.check_fitted()
    check_is_tensor(X)
    check_2d(X)

    if X.size(1) != self.in_features:
        raise ValueError(f"Expected {self.in_features} features, got {X.size(1)}.")

    X = X.to(device=self.device, dtype=self.dtype)

    labels, _ = self._predict_with_inertia(X, self._means)

    return labels

fit_predict(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Fit k-means and predict cluster labels.

PARAMETER	DESCRIPTION
X	Training data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights of shape (n_samples,). Weights are used to compute weighted cluster centers and weighted inertia. If None, all samples are weighted equally. TYPE: torch.Tensor | None, by default None DEFAULT: None
target	Target values. Not supported. TYPE: torch.Tensor | None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Predicted cluster labels of shape (n_samples,).

Source code in spectre/clustering/kmeans.py

def fit_predict(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Fit k-means and predict cluster labels.

    Parameters
    ----------
    X : torch.Tensor
        Training data of shape (n_samples, in_features).

    weights : torch.Tensor | None, by default None
        Sample weights of shape (n_samples,). Weights are used to compute
        weighted cluster centers and weighted inertia. If None, all samples
        are weighted equally.

    target : torch.Tensor | None, by default None
        Target values. Not supported.

    Returns
    -------
    torch.Tensor
        Predicted cluster labels of shape (n_samples,).
    """
    return self.fit(X, weights=weights, target=target).predict(X)

score(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> float

Compute the negative inertia (higher is better).

PARAMETER	DESCRIPTION
X	Input data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights of shape (n_samples,). If provided, computes weighted inertia. TYPE: torch.Tensor | None, by default None DEFAULT: None
target	Target values. Not supported. TYPE: torch.Tensor | None, by default None DEFAULT: None

RETURNS	DESCRIPTION
float	Negative inertia score (higher values indicate better clustering).

Source code in spectre/clustering/kmeans.py

def score(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> float:
    """
    Compute the negative inertia (higher is better).

    Parameters
    ----------
    X : torch.Tensor
        Input data of shape (n_samples, in_features).

    weights : torch.Tensor | None, by default None
        Sample weights of shape (n_samples,). If provided, computes weighted
        inertia.

    target : torch.Tensor | None, by default None
        Target values. Not supported.

    Returns
    -------
    float
        Negative inertia score (higher values indicate better clustering).
    """
    self.check_fitted()

    X = X.to(device=self.device, dtype=self.dtype)

    if weights is not None:
        check_is_tensor(weights)
        check_1d(weights)
        check_same_shape(X, weights, axis=0)
        check_sample_weights(weights)
        weights = weights.to(device=self.device, dtype=self.dtype)

    _, inertia = self._predict_with_inertia(X, self._means, weights)

    return -float(inertia)

get_results() -> dict[str, torch.Tensor]

Get all fitted parameters as a dictionary.

RETURNS	DESCRIPTION
dict[str, Tensor]	Dictionary containing computed results.

Source code in spectre/clustering/kmeans.py

def get_results(self) -> dict[str, torch.Tensor]:
    """
    Get all fitted parameters as a dictionary.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary containing computed results.
    """
    self.check_fitted()

    return {
        "means": self._means,
        "labels": self._labels,
        "inertia": self._inertia,
    }

Functions

kmeans(X: torch.Tensor, *, n_states: int, init: Literal['k-means++', 'random'] = 'k-means++', n_init: int = 1, n_iter: int = 100, tol: float = 0.0001, weights: torch.Tensor | None = None, device: torch.device | str | None = None, dtype: torch.dtype = torch.float32) -> dict[str, torch.Tensor | float | list]

Perform k-means clustering (functional interface).

PARAMETER	DESCRIPTION
X	Input data of shape (n_samples, in_features). TYPE: Tensor
n_states	Number of states (clusters) to form. TYPE: int
init	Initialization method for cluster centers. TYPE: Literal["k-means++", "random"], by default "k-means++" DEFAULT: 'k-means++'
n_init	Number of random initializations to try. TYPE: int, by default 1 DEFAULT: 1
n_iter	Maximum number of iterations for a single run. TYPE: int, by default 100 DEFAULT: 100
tol	Relative tolerance for convergence. TYPE: float, by default 1e-4 DEFAULT: 0.0001
weights	Sample weights of shape (n_samples,). Weights are used to compute weighted cluster centers and weighted inertia. If None, all samples are weighted equally. TYPE: torch.Tensor | None, by default None DEFAULT: None
device	Device on which to perform computations. TYPE: torch.device | str | None, by default None DEFAULT: None
dtype	Data type for computations. TYPE: torch.dtype, by default torch.float32 DEFAULT: float32

RETURNS	DESCRIPTION
dict[str, Tensor | float | list]	Dictionary with clustering results.

Source code in spectre/clustering/kmeans.py

def kmeans(
    X: torch.Tensor,
    *,
    n_states: int,
    init: Literal["k-means++", "random"] = "k-means++",
    n_init: int = 1,
    n_iter: int = 100,
    tol: float = 1e-4,
    weights: torch.Tensor | None = None,
    device: torch.device | str | None = None,
    dtype: torch.dtype = torch.float32,
) -> dict[str, torch.Tensor | float | list]:
    """
    Perform k-means clustering (functional interface).

    Parameters
    ----------
    X : torch.Tensor
        Input data of shape (n_samples, in_features).

    n_states : int
        Number of states (clusters) to form.

    init : Literal["k-means++", "random"], by default "k-means++"
        Initialization method for cluster centers.

    n_init : int, by default 1
        Number of random initializations to try.

    n_iter : int, by default 100
        Maximum number of iterations for a single run.

    tol : float, by default 1e-4
        Relative tolerance for convergence.

    weights : torch.Tensor | None, by default None
        Sample weights of shape (n_samples,). Weights are used to compute
        weighted cluster centers and weighted inertia. If None, all samples
        are weighted equally.

    device : torch.device | str | None, by default None
        Device on which to perform computations.

    dtype : torch.dtype, by default torch.float32
        Data type for computations.

    Returns
    -------
    dict[str, torch.Tensor | float | list]
        Dictionary with clustering results.
    """
    model = KMeans(
        n_states=n_states,
        init=init,
        n_init=n_init,
        n_iter=n_iter,
        tol=tol,
        device=device,
        dtype=dtype,
    )

    model.fit(X, weights=weights, target=None)

    return model.get_results()