Clustering Gaussian Mixture

gaussian_mixture

CLASS	DESCRIPTION
GaussianMixture	Gaussian Mixture Model clustering using Expectation-Maximization algorithm.

FUNCTION	DESCRIPTION
gaussian_mixture	Perform Gaussian Mixture Model clustering (functional interface).

Classes

GaussianMixture(n_states: int, covariance_type: Literal['diag', 'spherical', 'full', 'tied'] = 'diag', init: Literal['random', 'k-means++'] = 'random', n_init: int = 1, n_iter: int = 100, eps: float = torch.finfo(torch.float32).eps, tol: float = 0.001, device: torch.device | str | None = None, dtype: torch.dtype = torch.float32)

Bases: Estimator

Gaussian Mixture Model clustering using Expectation-Maximization algorithm.

A Gaussian Mixture Model represents data as a combination of multiple Gaussian distributions. This implementation uses the Expectation-Maximization (EM) algorithm to iteratively estimate the parameters of each Gaussian component.

The fit method accepts sample weights to handle class imbalance, importance sampling, or uncertainty weighting. Weights are incorporated into both the E-step (responsibilities) and M-step (parameter updates):

• Responsibilities weighted by sample weights
• Component means computed as weighted averages
• Covariances computed with weighted samples
• Mixture weights normalized based on effective sample counts

Weights are automatically normalized to sum to n_samples. Zero or negative weights raise an error.

PARAMETER	DESCRIPTION
n_states	Number of Gaussian components (states) to form. Must be at least 1. TYPE: int
covariance_type	Type of covariance. "diag": Diagonal covariance matrices (features assumed independent) "spherical": Single variance per component (isotropic covariances) "full": Full covariance matrices (arbitrary correlations) "tied": Single covariance matrix shared across all components TYPE: Literal["diag", "spherical", "full", "tied"], by default "diag" DEFAULT: 'diag'
init	Initialization strategy for component parameters. "random": Random assignment of samples to components "k-means++": Use k-means++ to initialize component centers TYPE: Literal["random", "k-means++"], by default "random" DEFAULT: 'random'
n_init	Number of EM initializations to perform. Best result is kept. TYPE: int, by default 1 DEFAULT: 1
n_iter	Maximum number of EM iterations per initialization. TYPE: int, by default 100 DEFAULT: 100
eps	Regularization term added to diagonal of covariance matrices to ensure numerical stability. TYPE: float, by default torch.finfo(torch.float32).eps DEFAULT: eps
tol	Convergence tolerance. EM stops when log-likelihood improvement is below this threshold. TYPE: float, by default 1e-3 DEFAULT: 0.001
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 Mean parameters for each mixture component of shape (n_states, in_features).

covariances : torch.Tensor Covariance parameters for each component. Shape depends on covariance_type:

- "diag": (n_states, in_features)
- "spherical": (n_states, 1)
- "full": (n_states, in_features, in_features)
- "tied": (in_features, in_features)

weights : torch.Tensor Mixing weights for each component of shape (n_states,).

labels : torch.Tensor Component labels for each training sample.

lower_bound : float Log-likelihood lower bound achieved by the model.

converged : bool Whether the EM algorithm converged for the best fit.

Examples:

Basic Gaussian mixture modeling:

>>> from spectre.utils.compute import deterministic
>>> deterministic(42)
>>> gmm = GaussianMixture(n_states=3, covariance_type="diag")
>>> X = torch.randn(100, 2)
>>> labels = gmm.fit_predict(X)
>>> labels.shape
torch.Size([100])

Accessing fitted parameters:

>>> gmm.fit(X)
>>> print(f"Component means: {gmm.means_.shape}")
>>> print(f"Component weights: {gmm.weights_}")

Using spherical covariances:

>>> gmm = GaussianMixture(n_states=2, covariance_type="spherical", n_init=5)
>>> labels = gmm.fit_predict(X)

Evaluating clustering quality:

>>> score = gmm.score(X)  # Silhouette score
>>> print(f"Clustering quality: {score:.3f}")

Soft clustering probabilities:

>>> gmm = GaussianMixture(n_states=2, covariance_type="diag")
>>> gmm.fit(X)
>>> probabilities = gmm.predict_proba(X)  # Returns: (n_samples, n_states)
>>> # Get samples with high uncertainty
>>> max_probs = probabilities.max(dim=1)[0]
>>> uncertain = X[max_probs < 0.7]

Weighted samples for class imbalance or importance weighting:

>>> # Emphasize certain samples during fitting
>>> weights = torch.ones(100)
>>> weights[:50] = 2.0  # Double weight for first half
>>> gmm = GaussianMixture(n_states=2, covariance_type="diag")
>>> gmm.fit(X, weights=weights)
>>> # Component weights will reflect the sample weighting
>>> print(gmm.weights_)  # Biased toward states containing weighted samples

METHOD	DESCRIPTION
fit	Fit Gaussian Mixture Model to input data.
predict	Predict component labels for input data.
fit_predict	Fit Gaussian Mixture Model and predict component labels.
predict_proba	Predict component membership probabilities for input data.
score	Return the silhouette score of the clustering.
get_results	Get all fitted parameters as a dictionary.

ATTRIBUTE	DESCRIPTION
means	Component means. TYPE: Tensor
covariances	Component covariances. TYPE: Tensor
mixing_weights	Component mixing weights. TYPE: Tensor
labels	Training sample labels. TYPE: Tensor
lower_bound	Log-likelihood lower bound. TYPE: float
converged	Whether EM algorithm converged. TYPE: bool

Source code in spectre/clustering/gaussian_mixture.py

def __init__(
    self,
    n_states: int,
    covariance_type: Literal["diag", "spherical", "full", "tied"] = "diag",
    init: Literal["random", "k-means++"] = "random",
    n_init: int = 1,
    n_iter: int = 100,
    eps: float = torch.finfo(torch.float32).eps,
    tol: float = 1e-3,
    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
    # Assume number of output features is the number of states.
    self.out_features = n_states

    allowed_covariance_types = ["diag", "spherical", "full", "tied"]
    if covariance_type not in allowed_covariance_types:
        raise ValueError(
            f"Expected `covariance_type` to be one of {allowed_covariance_types}, "
            f"got {covariance_type}."
        )
    self.covariance_type = covariance_type

    allowed_init = ["random", "k-means++"]
    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(eps, (int, float)):
        raise TypeError(f"Expected `eps` to be a number, got {type(eps).__name__}.")
    check_in_interval(eps, "(0, inf)")
    self.eps = float(eps)

    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

    self._log_2pi = torch.tensor(2.0 * torch.pi, device=self.device).log()

    # Internal attributes.

    # Mixture
    self._means = None
    self._covariances = None
    self._mixing_weights = None

    # Clustering results
    self._labels = None

    self._lower_bound = None
    self._converged = None

Attributes

means: torch.Tensor property

Component means.

covariances: torch.Tensor property

Component covariances.

mixing_weights: torch.Tensor property

Component mixing weights.

labels: torch.Tensor property

Training sample labels.

lower_bound: float property

Log-likelihood lower bound.

converged: bool property

Whether EM algorithm converged.

Functions

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

Fit Gaussian Mixture Model to input data.

PARAMETER	DESCRIPTION
X	Training data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights of shape (n_samples,). If provided, samples are weighted during EM parameter estimation. Weights are normalized to sum to n_samples. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
target	Target values for supervised learning. Not supported. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
GaussianMixture	Return self.

Source code in spectre/clustering/gaussian_mixture.py

def fit(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> "GaussianMixture":
    """
    Fit Gaussian Mixture Model to input data.

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

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape (n_samples,). If provided, samples are weighted
        during EM parameter estimation. Weights are normalized to sum to n_samples.

    target : torch.Tensor | None, optional, by default None
        Target values for supervised learning. Not supported.

    Returns
    -------
    GaussianMixture
        Return self.
    """
    check_is_tensor(X)
    check_2d(X)

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

    if weights is not None:
        check_is_tensor(weights)
        check_1d(weights)
        check_same_len(X, weights)
        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=self.eps))
    else:
        # Uniform weights
        weights = torch.ones(X.shape[0], device=self.device, dtype=self.dtype)

    self.in_features = X.shape[1]

    b_lower_bound = float("-inf")
    b_mixture = None
    b_labels = None
    b_converged = False

    for _ in range(self.n_init):
        means, covariances, mixing_weights = self._initialize_mixture(
            X, weights=weights
        )

        lower_bound_prev = float("-inf")
        converged = False

        for iter_idx in range(self.n_iter):
            # E-step: compute responsibilities.
            log_resp = self._compute_log_responsibilities(
                X, means, covariances, mixing_weights
            )
            resp = torch.softmax(log_resp, dim=1)
            lower_bound = self._compute_lower_bound(log_resp, weights)

            if abs(lower_bound - lower_bound_prev) < self.tol:
                converged = True
                break

            lower_bound_prev = lower_bound

            # M-step: update parameters with sample weights.
            means, covariances, mixing_weights = self._update_mixture(
                X, resp, weights
            )

            # Check for ill-conditioned covariances based on type
            if self.covariance_type in ["diag", "spherical"]:
                if torch.any(covariances <= 0):
                    raise ValueError(
                        "Fitting failed due to ill-conditioned covariance matrices. "
                        "Try increasing eps or reducing `n_states`."
                    )

        if lower_bound > b_lower_bound:
            b_lower_bound = lower_bound
            b_mixture = (means, covariances, mixing_weights)
            b_labels = torch.argmax(resp, dim=1)
            b_converged = converged

    self._means, self._covariances, self._mixing_weights = b_mixture
    self._labels = b_labels
    self._lower_bound = b_lower_bound
    self._converged = b_converged

    self.is_fitted = True

    return self

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

Predict component labels for input data.

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

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

Source code in spectre/clustering/gaussian_mixture.py

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

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

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

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

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

    log_resp = self._compute_log_responsibilities(
        X, self._means, self._covariances, self._mixing_weights
    )

    return torch.argmax(log_resp, dim=1)

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

Fit Gaussian Mixture Model and predict component labels.

PARAMETER	DESCRIPTION
X	Training data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights of shape (n_samples, ). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
target	Target values for supervised learning. Currently not supported. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

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

Source code in spectre/clustering/gaussian_mixture.py

def fit_predict(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Fit Gaussian Mixture Model and predict component labels.

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

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape (n_samples, ).

    target : torch.Tensor | None, optional, by default None
        Target values for supervised learning. Currently not supported.

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

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

Predict component membership probabilities for input data.

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

RETURNS	DESCRIPTION
Tensor	Component membership probabilities of shape (n_samples, n_states).

Source code in spectre/clustering/gaussian_mixture.py

def predict_proba(self, X: torch.Tensor) -> torch.Tensor:
    """
    Predict component membership probabilities for input data.

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

    Returns
    -------
    torch.Tensor
        Component membership probabilities of shape (n_samples, n_states).
    """
    self.check_fitted()
    check_2d(X)
    check_is_tensor(X)

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

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

    log_resp = self._compute_log_responsibilities(
        X, self._means, self._covariances, self._mixing_weights
    )

    return torch.softmax(log_resp, dim=1)

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

Return the silhouette score of the clustering.

Score is between -1 and 1. Higher values indicate better clustering

PARAMETER	DESCRIPTION
X	Data to score. TYPE: Tensor
weights	Sample weights. Currently ignored for silhouette score computation. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
target	Target values. Ignored for clustering (compatibility). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Silhouette score.

Source code in spectre/clustering/gaussian_mixture.py

def score(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Return the silhouette score of the clustering.

    Score is between -1 and 1. Higher values indicate better clustering

    Parameters
    ----------
    X : torch.Tensor
        Data to score.

    weights : torch.Tensor | None, optional, by default None
        Sample weights. Currently ignored for silhouette score computation.

    target : torch.Tensor | None, optional, by default None
        Target values. Ignored for clustering (compatibility).

    Returns
    -------
    torch.Tensor
        Silhouette score.
    """
    self.check_fitted()

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

    return self._silhouette_score(X, self._labels)

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

Get all fitted parameters as a dictionary.

RETURNS	DESCRIPTION
dict	Dictionary containing computed results.

Source code in spectre/clustering/gaussian_mixture.py

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

    Returns
    -------
    dict
        Dictionary containing computed results.
    """
    self.check_fitted()

    return {
        "means": self._means,
        "covariances": self._covariances,
        "mixing_weights": self._mixing_weights,
        "labels": self._labels,
        "lower_bound": self._lower_bound,
        "converged": self._converged,
    }

Functions

gaussian_mixture(X: torch.Tensor, weights: torch.Tensor | None = None, *, n_states: int, covariance_type: Literal['diag', 'spherical', 'full', 'tied'] = 'diag', init: Literal['random', 'k-means++'] = 'random', n_init: int = 1, n_iter: int = 100, eps: float = torch.finfo(torch.float32).eps, tol: float = 0.001, device: torch.device | str | None = None, dtype: torch.dtype = torch.float32) -> dict[str, torch.Tensor | float | bool]

Perform Gaussian Mixture Model clustering (functional interface).

PARAMETER	DESCRIPTION
X	Input data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights. TYPE: torch.Tensor | None, by default None DEFAULT: None
n_states	Number of Gaussian components. TYPE: int
covariance_type	Type of covariance parameters. See GaussianMixture for details. TYPE: Literal["diag", "spherical", "full", "tied"], by default "diag" DEFAULT: 'diag'
init	Initialization strategy. TYPE: Literal["random", "k-means++"], by default "random" DEFAULT: 'random'
n_init	Number of random initializations. TYPE: int, by default 1 DEFAULT: 1
n_iter	Maximum EM iterations. TYPE: int, by default 100 DEFAULT: 100
eps	Covariance regularization. TYPE: float, by default torch.finfo(torch.float32).eps DEFAULT: eps
tol	Convergence tolerance. TYPE: float, by default 1e-3 DEFAULT: 0.001
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 | bool]	Clustering results.

Source code in spectre/clustering/gaussian_mixture.py

def gaussian_mixture(
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    *,
    n_states: int,
    covariance_type: Literal["diag", "spherical", "full", "tied"] = "diag",
    init: Literal["random", "k-means++"] = "random",
    n_init: int = 1,
    n_iter: int = 100,
    eps: float = torch.finfo(torch.float32).eps,
    tol: float = 1e-3,
    device: torch.device | str | None = None,
    dtype: torch.dtype = torch.float32,
) -> dict[str, torch.Tensor | float | bool]:
    """
    Perform Gaussian Mixture Model clustering (functional interface).

    Parameters
    ----------
    X : torch.Tensor
        Input data of shape (n_samples, in_features).
    weights : torch.Tensor | None, by default None
        Sample weights.
    n_states : int
        Number of Gaussian components.
    covariance_type : Literal["diag", "spherical", "full", "tied"], by default "diag"
        Type of covariance parameters. See `GaussianMixture` for details.
    init : Literal["random", "k-means++"], by default "random"
        Initialization strategy.
    n_init : int, by default 1
        Number of random initializations.
    n_iter : int, by default 100
        Maximum EM iterations.
    eps : float, by default torch.finfo(torch.float32).eps
        Covariance regularization.
    tol : float, by default 1e-3
        Convergence tolerance.
    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 | bool]
        Clustering results.
    """
    gmm = GaussianMixture(
        n_states=n_states,
        covariance_type=covariance_type,
        init=init,
        n_init=n_init,
        n_iter=n_iter,
        eps=eps,
        tol=tol,
        device=device,
        dtype=dtype,
    )
    gmm.fit(X, weights=weights, target=None)

    return gmm.get_results()