def mstep_gpcm(
code: str,
W: torch.Tensor,
nk: torch.Tensor,
d: int,
state: Optional[GPCMResult] = None,
reg_covar: float = 0.0,
gpcm_tol: float = 1e-6,
gpcm_max_iter: int = 20,
) -> GPCMResult:
r"""
One M-step update for GPCM covariance code `code`.
Parameters
----------
code : str
One of the 14 mclust codes (case-insensitive), see `GPCM_CODES`.
W : torch.Tensor
Per-component scatter, shape `(K, d, d)` -- `W_k = sum_i r_ik
(x_i - ybar_k)(x_i - ybar_k)^T` (Phase 0's `stats.s2` for `full`
covariance).
nk : torch.Tensor
Effective per-component sample counts, shape `(K,)`.
d : int
Number of features.
state : GPCMResult, optional
Previous M-step's result, used to warm-start the 5 iterative
codes (`VEI`, `VEE`, `VEV`, `EVE`, `VVE`). Ignored by the other 9.
reg_covar : float, optional
Added to the assembled covariance's diagonal (not the
decomposition). (default: 0.0)
gpcm_tol : float, optional
Relative-change stopping tolerance for the iterative codes'
complete-data objective. (default: 1e-6)
gpcm_max_iter : int, optional
Maximum inner iterations for the iterative codes. (default: 20)
Returns
-------
GPCMResult
"""
code = code.upper()
if code not in GPCM_CODES:
raise ValueError(f"code must be one of {GPCM_CODES}, got {code!r}")
K = W.shape[0]
n = nk.sum()
dtype, device = W.dtype, W.device
I_d = torch.eye(d, dtype=dtype, device=device)
diagW = torch.diagonal(W, dim1=-2, dim2=-1) # (K, d) -- diag(W_k)
# ===================================================================
# Diagonal (axis-aligned, orientation = I) models: EII, VII, EEI, VVI
# (the last 4 already have dedicated, tested code in tgmm/gmm.py --
# implemented here too so tests can exercise all 14 codes uniformly).
# ===================================================================
if code == 'EII':
lam = (diagW.sum()) / (n * d)
volume, shape, orientation = lam, torch.ones(d, dtype=dtype, device=device), I_d
elif code == 'VII':
lam_k = diagW.sum(dim=-1) / (nk * d)
volume, shape, orientation = lam_k, torch.ones(d, dtype=dtype, device=device), I_d
elif code == 'EEI':
b = diagW.sum(dim=0) / n # (d,) -- this IS lambda * A already (free diagonal)
volume, shape, orientation = torch.tensor(1.0, dtype=dtype, device=device), b, I_d
elif code == 'VVI':
b_k = diagW / nk.unsqueeze(-1) # (K, d)
volume, shape, orientation = torch.ones(K, dtype=dtype, device=device), b_k, I_d
# ===================================================================
# VEI: volume varies, shape shared, axis-aligned. Iterative.
# ===================================================================
elif code == 'VEI':
A = state.shape.clone() if (state is not None and state.shape.dim() == 1) else torch.ones(d, dtype=dtype, device=device)
lam_k = None
prev_J = None
for _ in range(gpcm_max_iter):
A_inv = 1.0 / A
lam_k = (diagW * A_inv).sum(dim=-1) / (nk * d)
B = (diagW / lam_k.unsqueeze(-1)).sum(dim=0)
A = B / B.prod().clamp_min(1e-300).pow(1.0 / d)
J = (nk * d * torch.log(lam_k)).sum() + (diagW * (1.0 / A) / lam_k.unsqueeze(-1)).sum()
if prev_J is not None and abs(prev_J - J) < gpcm_tol * abs(prev_J).clamp_min(1e-300):
break
prev_J = J
volume, shape, orientation = lam_k, A, I_d
# ===================================================================
# EVI: volume shared, shape varies, axis-aligned. Closed form.
# ===================================================================
elif code == 'EVI':
B_k = diagW # (K, d)
det_B_k = B_k.prod(dim=-1).clamp_min(1e-300) # (K,)
A_k = B_k / det_B_k.unsqueeze(-1).pow(1.0 / d)
lam = det_B_k.pow(1.0 / d).sum() / n
volume, shape, orientation = lam, A_k, I_d
# ===================================================================
# EEV: volume & shape shared, orientation varies. Closed form.
# ===================================================================
elif code == 'EEV':
eigvals, eigvecs = _eigh_descending(W) # (K, d), (K, d, d)
Omega = eigvals.sum(dim=0) # (d,)
det_Omega = Omega.prod().clamp_min(1e-300)
A = Omega / det_Omega.pow(1.0 / d)
lam = det_Omega.pow(1.0 / d) / n
volume, shape, orientation = lam, A, eigvecs
# ===================================================================
# VEV: shape shared, volume & orientation vary. Iterative (D fixed
# per-k from the data; lambda_k/A alternate).
# ===================================================================
elif code == 'VEV':
eigvals, eigvecs = _eigh_descending(W) # (K, d), (K, d, d) -- Omega_k, D_k
A = state.shape.clone() if (state is not None and state.shape.dim() == 1) else torch.ones(d, dtype=dtype, device=device)
lam_k = None
prev_J = None
for _ in range(gpcm_max_iter):
A_inv = 1.0 / A
lam_k = (eigvals * A_inv).sum(dim=-1) / (nk * d)
Omega_sum = (eigvals / lam_k.unsqueeze(-1)).sum(dim=0)
A = Omega_sum / Omega_sum.prod().clamp_min(1e-300).pow(1.0 / d)
J = (nk * d * torch.log(lam_k)).sum() + (eigvals * (1.0 / A) / lam_k.unsqueeze(-1)).sum()
if prev_J is not None and abs(prev_J - J) < gpcm_tol * abs(prev_J).clamp_min(1e-300):
break
prev_J = J
volume, shape, orientation = lam_k, A, eigvecs
# ===================================================================
# EVV: volume shared, shape & orientation vary. Closed form.
# ===================================================================
elif code == 'EVV':
det_W_k = torch.linalg.det(W).clamp_min(1e-300) # (K,)
C_k = W / det_W_k.view(K, 1, 1).pow(1.0 / d)
lam = det_W_k.pow(1.0 / d).sum() / n
eigvals, eigvecs = _eigh_descending(C_k) # exposed shape_/orientation_ only
volume, shape, orientation = lam, eigvals, eigvecs
cov_full = lam * C_k + reg_covar * I_d.unsqueeze(0)
return GPCMResult(volume=volume, shape=shape, orientation=orientation, cov_full=cov_full)
# ===================================================================
# VEE: shape & orientation shared (C = D A D^T, det 1), volume
# varies. Iterative.
# ===================================================================
elif code == 'VEE':
C = state.orientation.clone() @ torch.diag_embed(state.shape.clone()) @ state.orientation.clone().transpose(-1, -2) \
if (state is not None and state.orientation.dim() == 2) else I_d.clone()
lam_k = None
prev_J = None
for _ in range(gpcm_max_iter):
C_inv = torch.linalg.inv(C)
lam_k = torch.einsum('kij,ji->k', W, C_inv) / (nk * d)
B = (W / lam_k.view(K, 1, 1)).sum(dim=0)
det_B = torch.linalg.det(B).clamp_min(1e-300)
C = B / det_B.pow(1.0 / d)
J = (nk * d * torch.log(lam_k)).sum() + torch.einsum('kij,ji->k', W, C_inv).div(lam_k).sum()
if prev_J is not None and abs(prev_J - J) < gpcm_tol * abs(prev_J).clamp_min(1e-300):
break
prev_J = J
eigvals, eigvecs = _eigh_descending(C)
volume, shape, orientation = lam_k, eigvals, eigvecs
cov_full = lam_k.view(K, 1, 1) * C.unsqueeze(0) + reg_covar * I_d.unsqueeze(0)
return GPCMResult(volume=volume, shape=shape, orientation=orientation, cov_full=cov_full)
# ===================================================================
# EVE / VVE: orientation shared, shape varies, volume shared (EVE) or
# varies (VVE). Iterative; the orientation sub-step is a Stiefel-
# manifold optimization (no closed form).
# ===================================================================
elif code in ('EVE', 'VVE'):
D = state.orientation.clone() if (state is not None and state.orientation.dim() == 2) else I_d.clone()
A_k = None
lam = None
lam_k = None
prev_J = None
for _ in range(gpcm_max_iter):
rotated = D.T @ W @ D # (K, d, d)
B_k = torch.diagonal(rotated, dim1=-2, dim2=-1) # (K, d)
det_B_k = B_k.prod(dim=-1).clamp_min(1e-300)
A_k = B_k / det_B_k.unsqueeze(-1).pow(1.0 / d)
A_k_inv = 1.0 / A_k
if code == 'EVE':
lam = (A_k_inv * B_k).sum() / (n * d)
lam_vec = lam.expand(K)
else: # VVE
lam_k = (A_k_inv * B_k).sum(dim=-1) / (nk * d)
lam_vec = lam_k
D = _optimize_orientation(W, A_k_inv, lam_vec, D)
rotated = D.T @ W @ D
B_k_check = torch.diagonal(rotated, dim1=-2, dim2=-1)
J = (nk * d * torch.log(lam_vec)).sum() + (A_k_inv * B_k_check / lam_vec.unsqueeze(-1)).sum()
if prev_J is not None and abs(prev_J - J) < gpcm_tol * abs(prev_J).clamp_min(1e-300):
break
prev_J = J
volume = lam if code == 'EVE' else lam_k
shape, orientation = A_k, D
# ===================================================================
# EEE, VVV: existing tgmm types ('tied_full', 'full'), implemented
# here too only so every code is available through one interface for
# testing -- tgmm/gmm.py keeps using its own dedicated code.
# ===================================================================
elif code == 'EEE':
C = W.sum(dim=0) / n
eigvals, eigvecs = _eigh_descending(C)
volume, shape, orientation = torch.tensor(1.0, dtype=dtype, device=device), eigvals, eigvecs
cov_full = C.unsqueeze(0).expand(K, d, d).clone() + reg_covar * I_d.unsqueeze(0)
return GPCMResult(volume=volume, shape=shape, orientation=orientation, cov_full=cov_full)
elif code == 'VVV':
C_k = W / nk.view(K, 1, 1)
eigvals, eigvecs = _eigh_descending(C_k)
volume, shape, orientation = torch.ones(K, dtype=dtype, device=device), eigvals, eigvecs
cov_full = C_k + reg_covar * I_d.unsqueeze(0)
return GPCMResult(volume=volume, shape=shape, orientation=orientation, cov_full=cov_full)
else: # pragma: no cover -- exhaustive above
raise AssertionError(code)
cov_full = _assemble_full(volume, shape, orientation, K, d, reg_covar)
return GPCMResult(volume=volume, shape=shape, orientation=orientation, cov_full=cov_full)