import numpy as np
import sklearn
import torch
from torch import Tensor
from torch.func import vmap
[docs]
class GaussianMixture:
"""
Gaussian Mixture Class.
Parameters (6):
- n_components (int)
- dim (int)
- weights (Tensor) shape (n_components,)
- means (Tensor) shape (n_components, dim)
- covariances (Tensor) shape (n_components, dim, dim)
- device (torch.device or str)
"""
def __init__( # noqa: PLR0913
self,
n_components: int = 1,
dim: int = 1,
weights: Tensor | None = None,
means: Tensor | None = None,
covariances: Tensor | None = None,
device: torch.device | str = "cpu",
) -> None:
self.device = torch.device(device)
self.n_components = n_components
self.dim = dim
if weights is None:
weights = torch.full((n_components,), 1.0 / n_components, device=self.device)
if means is None:
means = torch.zeros((n_components, dim), device=self.device)
if covariances is None:
covariances = torch.stack(
[torch.eye(dim, device=self.device) for _ in range(n_components)]
)
self.n_components = self._validate_n_components(n_components)
self.dim = self._validate_dim(dim)
self.weights = self._validate_weights(weights)
self.means = self._validate_means(means)
self.covariances = self._validate_covariances(covariances)
[docs]
def prob_hyperrec(self, event: Tensor, num_samples: int = 100_00) -> Tensor:
"""
Compute the integral of the Gaussian Mixture over a hyperrectangle known as the "event".
Parameters
----------
- E (Tensor): hyperrectangle bounds of shape (dim, 2)
Returns
-------
- Tensor: integral value over the hyperrectangle
"""
prob_hyperrec = torch.zeros(1, device=self.device)
event = event.to(self.device)
samples = self.sample(num_samples)
counter = 0
def in_hyperrectangle(sample: Tensor, event: Tensor) -> Tensor:
return torch.all((sample >= event[:, 0]) & (sample <= event[:, 1]))
in_hyperrectangle_vmap = vmap(lambda sample: in_hyperrectangle(sample, event))
mask = in_hyperrectangle_vmap(samples)
counter = mask.sum().item()
prob_hyperrec[0] = float(counter) / float(num_samples)
return prob_hyperrec
[docs]
def expectation(self) -> Tensor:
"""
Compute the expectation value of the Gaussian Mixture.
Returns
-------
- Tensor: expectation value of shape (dim)
"""
expectation = torch.zeros((self.dim), device=self.device)
for i in range(self.n_components):
expectation += self.weights[i] * self.means[i, :]
return expectation
[docs]
def covariance(self) -> Tensor:
"""
Compute the covariance matrix of the Gaussian Mixture.
Returns
-------
- Tensor: covariance matrix of shape (dim, dim)
"""
covariance = torch.zeros((self.dim, self.dim), device=self.device)
for i in range(self.n_components):
covariance += self.weights[i] * (
self.covariances[i, :, :]
+ torch.tensordot(self.means[i, :], self.means[i, :], dims=0)
) - 1 / self.n_components * torch.tensordot(
self.expectation(), self.expectation(), dims=0
)
return covariance
[docs]
def sample(self, n_samples: int, seed: int | None = None) -> Tensor:
"""
Generate samples from the Gaussian Mixture.
Parameters
----------
- num_samples (int): number of samples to generate
- seed (int, optional): random seed for reproducibility
Returns
-------
- Tensor: samples of shape (num_samples, dim)
"""
if seed is not None:
rng = sklearn.utils.check_random_state(seed)
else:
rng = sklearn.utils.check_random_state(None)
n_samples_comp = rng.multinomial(n_samples, self.weights.to("cpu").numpy())
samples_numpy = np.vstack(
[
rng.multivariate_normal(mean, covariance, int(sample))
for (mean, covariance, sample) in zip(
self.means.to("cpu").numpy(),
self.covariances.to("cpu").numpy(),
n_samples_comp,
strict=True,
)
]
)
samples = torch.tensor(samples_numpy, device=self.device).detach().clone()
return samples
def _validate_n_components(self, k: int) -> int:
if not isinstance(k, int) or k <= 0:
msg = "n_components must be a positive integer"
raise ValueError(msg)
return k
def _validate_dim(self, d: int) -> int:
if not isinstance(d, int) or d <= 0:
msg = "dim must be a positive integer"
raise ValueError(msg)
return d
def _validate_weights(self, w: Tensor) -> Tensor:
if w.ndim != 1 or w.shape[0] != self.n_components:
msg = "weights must have shape (n_components)"
raise ValueError(msg)
w = w.to(self.device)
w_sum = w.sum()
if not torch.isclose(w_sum, torch.tensor(1.0, device=self.device)):
msg = "weights must sum to 1.0"
raise ValueError(msg)
if torch.any(w < 0):
msg = "weights must be non-negative"
raise ValueError(msg)
if w_sum <= 0:
msg = "weights must sum to a positive value"
raise ValueError(msg)
return (w / w_sum).clone()
def _validate_means(self, m: Tensor) -> Tensor:
mean_tensor_order = 2
if m.ndim != mean_tensor_order or m.shape != (self.n_components, self.dim):
msg = "means must have shape (n_components, dim)"
raise ValueError(msg)
return m.to(self.device).clone()
def _validate_covariances(self, c: Tensor) -> Tensor:
covariance_tensor_order = 3
if c.ndim != covariance_tensor_order or c.shape != (self.n_components, self.dim, self.dim):
msg = "covariances must have shape (n_components, dim, dim)"
raise ValueError(msg)
c = c.to(self.device).clone()
# Basic positive-definite check (diagonal > 0 if diagonal matrix)
for i in range(self.n_components):
matrix = c[i, :, :]
eigenvalues = torch.linalg.eigvalsh(matrix)
if torch.any(eigenvalues <= 0):
msg = "Covariance matrices must be positive definite"
raise ValueError(msg)
return c
def __repr__(self) -> str:
return (
f"GaussianMixture(n_components={self.n_components}, "
f"dim={self.dim} ,"
f"weights_shape={tuple(self.weights.shape)}, "
f"means_shape={tuple(self.means.shape)}, "
f"covariances_shape={tuple(self.covariances.shape)}, device={self.device})"
)