Gaussian Mixture Model
A
"""
README, Author - Md Ruman Islam (mailto:ruman23.github.io)
Requirements:
- numpy
- matplotlib
Python:
- 3.8+
Inputs:
- data : a 2D numpy array of features.
- n_components : number of Gaussian distributions (clusters) to fit.
- max_iter : maximum number of EM iterations.
- tol : convergence tolerance.
Usage:
1. define 'n_components' value and 'data' features array
2. initialize model:
gmm = GaussianMixture(n_components=3, max_iter=100)
3. fit model to data:
gmm.fit(data)
4. get cluster predictions:
labels = gmm.predict(data)
5. visualize results:
gmm.plot_results(data)
"""
import warnings
import matplotlib.pyplot as plt
import numpy as np
from numpy.typing import NDArray
from scipy.stats import multivariate_normal
warnings.filterwarnings("ignore")
TAG = "GAUSSIAN-MIXTURE/ "
class GaussianMixture:
"""
Gaussian Mixture Model implemented using the Expectation-Maximization algorithm.
"""
def __init__(
self,
n_components: int = 2,
max_iter: int = 100,
tol: float = 1e-4,
seed: int | None = None,
) -> None:
self.n_components: int = n_components
self.max_iter: int = max_iter
self.tol: float = tol
self.seed: int | None = seed
# parameters
self.weights_: NDArray[np.float64] | None = None
self.means_: NDArray[np.float64] | None = None
self.covariances_: NDArray[np.float64] | None = None
self.log_likelihoods_: list[float] = []
def _initialize_parameters(self, data: NDArray[np.float64]) -> None:
"""Randomly initialize means, covariances, and mixture weights.
Examples
--------
>>> sample = np.array(
... [[0.0, 0.5], [1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]
... )
>>> model = GaussianMixture(n_components=2, seed=0)
>>> model._initialize_parameters(sample)
>>> model.means_.shape
(2, 2)
>>> bool(np.isclose(model.weights_.sum(), 1.0))
True
"""
rng = np.random.default_rng(self.seed)
n_samples, _ = data.shape
indices = rng.choice(n_samples, self.n_components, replace=False)
self.means_ = data[indices]
identity = np.eye(data.shape[1]) * 1e-6
self.covariances_ = np.array(
[np.cov(data, rowvar=False) + identity for _ in range(self.n_components)]
)
self.weights_ = np.ones(self.n_components) / self.n_components
def _e_step(self, data: NDArray[np.float64]) -> NDArray[np.float64]:
"""Compute responsibilities (posterior probabilities).
Examples
--------
>>> sample = np.array(
... [[0.0, 0.5], [1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]
... )
>>> model = GaussianMixture(n_components=2, seed=0)
>>> model._initialize_parameters(sample)
>>> resp = model._e_step(sample)
>>> resp.shape
(4, 2)
>>> bool(np.allclose(resp.sum(axis=1), 1.0))
True
"""
if self.weights_ is None or self.means_ is None or self.covariances_ is None:
raise ValueError(
"Model parameters must be initialized before running the E-step."
)
n_samples = data.shape[0]
responsibilities = np.zeros((n_samples, self.n_components))
weights = self.weights_
means = self.means_
covariances = self.covariances_
for k in range(self.n_components):
rv = multivariate_normal(
mean=means[k], cov=covariances[k], allow_singular=True
)
responsibilities[:, k] = weights[k] * rv.pdf(data)
# Normalize to get probabilities
responsibilities /= responsibilities.sum(axis=1, keepdims=True)
return responsibilities
def _m_step(
self,
data: NDArray[np.float64],
responsibilities: NDArray[np.float64],
) -> None:
"""Update weights, means, and covariances.
Note: assumes the model parameters are already initialized.
Examples
--------
>>> sample = np.array(
... [[0.0, 0.5], [1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]
... )
>>> model = GaussianMixture(n_components=2, seed=0)
>>> model._initialize_parameters(sample)
>>> resp = model._e_step(sample)
>>> model._m_step(sample, resp)
>>> bool(np.isclose(model.weights_.sum(), 1.0))
True
"""
n_samples, n_features = data.shape
component_counts = responsibilities.sum(axis=0)
self.weights_ = component_counts / n_samples
self.means_ = (responsibilities.T @ data) / component_counts[:, np.newaxis]
if self.covariances_ is None or self.means_ is None:
raise ValueError(
"Model parameters must be initialized before running the M-step."
)
covariances = self.covariances_
means = self.means_
for k in range(self.n_components):
diff = data - means[k]
covariances[k] = (responsibilities[:, k][:, np.newaxis] * diff).T @ diff
covariances[k] /= component_counts[k]
# Add small regularization term for numerical stability
covariances[k] += np.eye(n_features) * 1e-6
def _compute_log_likelihood(self, data: NDArray[np.float64]) -> float:
"""Compute total log-likelihood of the model.
Note: assumes the model parameters are already initialized.
Examples
--------
>>> sample = np.array(
... [[0.0, 0.5], [1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]
... )
>>> model = GaussianMixture(n_components=2, seed=0)
>>> model._initialize_parameters(sample)
>>> bool(np.isfinite(model._compute_log_likelihood(sample)))
True
"""
if self.weights_ is None or self.means_ is None or self.covariances_ is None:
raise ValueError(
"Model parameters must be initialized before computing likelihood."
)
n_samples = data.shape[0]
total_pdf = np.zeros((n_samples, self.n_components))
weights = self.weights_
means = self.means_
covariances = self.covariances_
for k in range(self.n_components):
rv = multivariate_normal(
mean=means[k], cov=covariances[k], allow_singular=True
)
total_pdf[:, k] = weights[k] * rv.pdf(data)
log_likelihood = np.sum(np.log(np.sum(total_pdf, axis=1) + 1e-12))
return log_likelihood
def fit(self, data: NDArray[np.float64]) -> None:
"""Fit the Gaussian Mixture Model to data using the EM algorithm.
Examples
--------
>>> sample = np.array(
... [[0.0, 0.5], [1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]
... )
>>> model = GaussianMixture(n_components=2, max_iter=5, tol=1e-3, seed=0)
>>> model.fit(sample) # doctest: +ELLIPSIS
GAUSSIAN-MIXTURE/ ...
>>> len(model.log_likelihoods_) > 0
True
"""
self._initialize_parameters(data)
prev_log_likelihood = None
for i in range(self.max_iter):
# E-step
responsibilities = self._e_step(data)
# M-step
self._m_step(data, responsibilities)
# Log-likelihood
log_likelihood = self._compute_log_likelihood(data)
self.log_likelihoods_.append(log_likelihood)
if (
prev_log_likelihood is not None
and abs(log_likelihood - prev_log_likelihood) < self.tol
):
print(f"{TAG}Converged at iteration {i}.")
break
prev_log_likelihood = log_likelihood
print(f"{TAG}Training complete. Final log-likelihood: {log_likelihood:.4f}")
def predict(self, data: NDArray[np.float64]) -> NDArray[np.int_]:
"""Predict cluster assignment for each data point.
Note: assumes the model parameters are already initialized.
Examples
--------
>>> sample = np.array(
... [[0.0, 0.5], [1.0, 1.5], [2.0, 2.5], [3.0, 3.5]]
... )
>>> model = GaussianMixture(n_components=2, max_iter=5, tol=1e-3, seed=0)
>>> model.fit(sample) # doctest: +ELLIPSIS
GAUSSIAN-MIXTURE/ ...
>>> labels = model.predict(sample)
>>> labels.shape
(4,)
"""
responsibilities = self._e_step(data)
return np.argmax(responsibilities, axis=1)
def plot_results(self, data: NDArray[np.float64]) -> None:
"""Visualize GMM clustering results (2D only).
Note: This method assumes self.means_ is initialized.
Examples
--------
>>> sample = np.ones((3, 3))
>>> model = GaussianMixture()
>>> model.plot_results(sample)
GAUSSIAN-MIXTURE/ Plotting only supported for 2D data.
"""
if data.shape[1] != 2:
print(f"{TAG}Plotting only supported for 2D data.")
return
labels = self.predict(data)
if self.means_ is None:
raise ValueError("Model means must be initialized before plotting.")
plt.scatter(data[:, 0], data[:, 1], c=labels, cmap="viridis", s=30)
plt.scatter(self.means_[:, 0], self.means_[:, 1], c="red", s=100, marker="x")
plt.title("Gaussian Mixture Model Clustering")
plt.xlabel("Feature 1")
plt.ylabel("Feature 2")
plt.show()
# Mock test
if __name__ == "__main__":
from sklearn.datasets import make_blobs
sample_data, _ = make_blobs(
n_samples=300, centers=3, cluster_std=1.2, random_state=42
)
gmm = GaussianMixture(n_components=3, max_iter=100, seed=42)
gmm.fit(sample_data)
labels = gmm.predict(sample_data)
gmm.plot_results(sample_data)
