"""Create, estimate, and sample from a Probabilistic PCA (PPCA) latent-factor model.
Defines the ProbabilisticPCADistribution, ProbabilisticPCASampler, ProbabilisticPCAAccumulatorFactory,
ProbabilisticPCAAccumulator, ProbabilisticPCAEstimator, and the ProbabilisticPCADataEncoder classes for
use with mixle.
Data type: np.ndarray[float] (a length-d real vector).
Probabilistic PCA is the latent linear-Gaussian model
z ~ N(0, I_q), x | z ~ N(W z + mu, sigma2 * I_d),
so marginally ``x ~ N(mu, C)`` with the structured covariance ``C = W W^T + sigma2 * I_d`` (a rank-q
factor structure plus isotropic noise). It is the probabilistic foundation of PCA / factor analysis and
gives a generative model, a likelihood, and a posterior over the latent factors
``E[z | x] = M^{-1} W^T (x - mu)`` (the low-dimensional embedding, exposed by ``transform``), with
``M = W^T W + sigma2 * I_q``.
Scoring uses the Woodbury identity, so the d-by-d inverse and log-determinant are obtained from a small
q-by-q solve (``C^{-1} = (I_d - W M^{-1} W^T) / sigma2`` and ``log|C| = (d-q) log sigma2 + log|M|``); the
reduction is engine-neutral, so the model scores on NumPy and Torch. Estimation is the **closed-form**
maximum-likelihood solution of Tipping & Bishop (1999): ``sigma2`` is the mean of the discarded
eigenvalues of the sample covariance and ``W`` is built from its top-q eigenpairs -- no EM iteration.
"""
from collections.abc import Sequence
from typing import Any
import numpy as np
from numpy.random import RandomState
from mixle.stats.compute.pdist import (
DataSequenceEncoder,
DistributionSampler,
ParameterEstimator,
SequenceEncodableProbabilityDistribution,
SequenceEncodableStatisticAccumulator,
StatisticAccumulatorFactory,
)
_MIN_SIGMA2 = 1.0e-12
_LOG_2PI = float(np.log(2.0 * np.pi))
[docs]
class ProbabilisticPCADistribution(SequenceEncodableProbabilityDistribution):
"""Probabilistic PCA: x ~ N(mu, W W^T + sigma2 I) with q latent factors."""
[docs]
@classmethod
def compute_capabilities(cls):
from mixle.stats.compute.capabilities import DistributionCapabilities
return DistributionCapabilities(engine_ready=("numpy", "torch"), kernel_status="generic")
def __init__(
self,
w: Sequence[Sequence[float]] | np.ndarray,
mu: Sequence[float] | np.ndarray,
sigma2: float,
name: str | None = None,
keys: str | None = None,
) -> None:
"""ProbabilisticPCADistribution object.
Args:
w (Union[Sequence[Sequence[float]], np.ndarray]): d-by-q factor-loading matrix W.
mu (Union[Sequence[float], np.ndarray]): Length-d mean vector.
sigma2 (float): Positive isotropic noise variance.
name (Optional[str]): Optional name for object instance.
keys (Optional[str]): Optional key for merging sufficient statistics.
Attributes:
w (np.ndarray): Factor-loading matrix (d, q).
mu (np.ndarray): Mean vector (d,).
sigma2 (float): Noise variance.
dim (int): Observation dimension d.
latent_dim (int): Number of factors q.
inv_covar (np.ndarray): Cached C^{-1} (d, d) via Woodbury.
log_det (float): Cached log|C|.
"""
w = np.asarray(w, dtype=float)
mu = np.asarray(mu, dtype=float).copy()
if w.ndim != 2 or w.shape[0] != len(mu):
raise ValueError("ProbabilisticPCADistribution requires W of shape (d, q) matching mu of length d.")
if sigma2 <= 0.0 or not np.isfinite(sigma2):
raise ValueError("ProbabilisticPCADistribution requires sigma2 > 0.")
self.w = w
self.mu = mu
self.sigma2 = float(sigma2)
self.dim = w.shape[0]
self.latent_dim = w.shape[1]
q = self.latent_dim
m = w.T @ w + self.sigma2 * np.eye(q) # (q, q)
self._m_inv = np.linalg.inv(m)
# Woodbury: C^{-1} = (I_d - W M^{-1} W^T) / sigma2
self.inv_covar = (np.eye(self.dim) - w @ self._m_inv @ w.T) / self.sigma2
sign, log_det_m = np.linalg.slogdet(m)
self.log_det = float((self.dim - q) * np.log(self.sigma2) + log_det_m)
self.name = name
self.keys = keys
def __str__(self) -> str:
"""Return string representation of ProbabilisticPCADistribution object."""
return "ProbabilisticPCADistribution(%s, %s, %s, name=%s, keys=%s)" % (
repr([[float(v) for v in row] for row in self.w]),
repr([float(v) for v in self.mu]),
repr(self.sigma2),
repr(self.name),
repr(self.keys),
)
[docs]
def density(self, x: Sequence[float] | np.ndarray) -> float:
"""Return the probability density at a single observation."""
return float(np.exp(self.log_density(x)))
[docs]
def log_density(self, x: Sequence[float] | np.ndarray) -> float:
"""Return the log-density at a single observation."""
diff = np.asarray(x, dtype=float) - self.mu
mahal = float(diff @ self.inv_covar @ diff)
return -0.5 * (self.dim * _LOG_2PI + self.log_det + mahal)
[docs]
def seq_log_density(self, x: np.ndarray) -> np.ndarray:
"""Return vectorized log-density values for sequence-encoded observations."""
diff = x - self.mu
mahal = np.einsum("ij,jk,ik->i", diff, self.inv_covar, diff)
return -0.5 * (self.dim * _LOG_2PI + self.log_det + mahal)
[docs]
def backend_seq_log_density(self, x: np.ndarray, engine: Any) -> Any:
"""Engine-neutral vectorized log-density for encoded data."""
diff = engine.asarray(x) - engine.asarray(self.mu)
mahal = engine.sum(engine.matmul(diff, engine.asarray(self.inv_covar)) * diff, axis=-1)
const = engine.asarray(self.dim * _LOG_2PI + self.log_det)
return engine.asarray(-0.5) * (const + mahal)
[docs]
def sampler(self, seed: int | None = None) -> "ProbabilisticPCASampler":
"""Return a sampler for drawing observations from this distribution."""
return ProbabilisticPCASampler(self, seed)
[docs]
def estimator(self, pseudo_count: float | None = None) -> "ProbabilisticPCAEstimator":
"""Return a closed-form ML estimator with the latent dimension fixed at this model's q."""
return ProbabilisticPCAEstimator(latent_dim=self.latent_dim, dim=self.dim, name=self.name, keys=self.keys)
[docs]
def dist_to_encoder(self) -> "ProbabilisticPCADataEncoder":
"""Return the data encoder used by this distribution for vectorized methods."""
return ProbabilisticPCADataEncoder()
[docs]
class ProbabilisticPCASampler(DistributionSampler):
"""Draw iid observations x = mu + W z + sigma * eps from a PPCA model."""
def __init__(self, dist: ProbabilisticPCADistribution, seed: int | None = None) -> None:
self.rng = RandomState(seed)
self.dist = dist
[docs]
def sample(self, size: int | None = None) -> np.ndarray:
"""Draw ``size`` iid vectors (shape (d,) when size is None, else (size, d))."""
sz = 1 if size is None else size
d, q = self.dist.dim, self.dist.latent_dim
z = self.rng.standard_normal(size=(sz, q))
noise = np.sqrt(self.dist.sigma2) * self.rng.standard_normal(size=(sz, d))
rv = self.dist.mu[None, :] + z @ self.dist.w.T + noise
return rv[0] if size is None else rv
[docs]
class ProbabilisticPCAAccumulator(SequenceEncodableStatisticAccumulator):
"""Accumulate the weighted count, mean, and second-moment matrix (the PPCA sufficient statistics)."""
def __init__(self, dim: int | None = None, keys: str | None = None) -> None:
self.dim = dim
self.count = 0.0
self.sum = np.zeros(dim) if dim is not None else None
self.sum2 = np.zeros((dim, dim)) if dim is not None else None
self.keys = keys
def _ensure_dim(self, d: int) -> None:
if self.dim is None:
self.dim = d
if self.sum is None:
self.sum = np.zeros(self.dim)
self.sum2 = np.zeros((self.dim, self.dim))
[docs]
def update(self, x: np.ndarray, weight: float, estimate: ProbabilisticPCADistribution | None) -> None:
xx = np.asarray(x, dtype=float)
self._ensure_dim(len(xx))
self.count += weight
self.sum += weight * xx
self.sum2 += weight * np.outer(xx, xx)
[docs]
def initialize(self, x: np.ndarray, weight: float, rng: RandomState | None) -> None:
self.update(x, weight, None)
[docs]
def seq_update(self, x: np.ndarray, weights: np.ndarray, estimate: ProbabilisticPCADistribution | None) -> None:
self._ensure_dim(x.shape[1])
self.count += float(np.sum(weights, dtype=np.float64))
self.sum += x.T @ weights
self.sum2 += (x * weights[:, None]).T @ x
[docs]
def seq_initialize(self, x: np.ndarray, weights: np.ndarray, rng: RandomState | None) -> None:
self.seq_update(x, weights, None)
[docs]
def combine(self, suff_stat: tuple[float, np.ndarray | None, np.ndarray | None]) -> "ProbabilisticPCAAccumulator":
count, s, s2 = suff_stat
if s is not None:
self._ensure_dim(len(s))
self.sum += s
self.sum2 += s2
self.count += count
return self
[docs]
def value(self) -> tuple[float, np.ndarray | None, np.ndarray | None]:
return self.count, self.sum, self.sum2
[docs]
def from_value(self, x: tuple[float, np.ndarray | None, np.ndarray | None]) -> "ProbabilisticPCAAccumulator":
self.count, self.sum, self.sum2 = x
self.dim = None if x[1] is None else len(x[1])
return self
[docs]
def key_merge(self, stats_dict: dict[str, Any]) -> None:
if self.keys is not None:
if self.keys in stats_dict:
stats_dict[self.keys].combine(self.value())
else:
stats_dict[self.keys] = self
[docs]
def key_replace(self, stats_dict: dict[str, Any]) -> None:
if self.keys is not None and self.keys in stats_dict:
self.from_value(stats_dict[self.keys].value())
[docs]
def acc_to_encoder(self) -> "ProbabilisticPCADataEncoder":
return ProbabilisticPCADataEncoder()
[docs]
class ProbabilisticPCAAccumulatorFactory(StatisticAccumulatorFactory):
"""Factory for ProbabilisticPCAAccumulator."""
def __init__(self, dim: int | None = None, keys: str | None = None) -> None:
self.dim = dim
self.keys = keys
[docs]
def make(self) -> ProbabilisticPCAAccumulator:
return ProbabilisticPCAAccumulator(dim=self.dim, keys=self.keys)
[docs]
class ProbabilisticPCAEstimator(ParameterEstimator):
"""Closed-form maximum-likelihood estimator for PPCA (Tipping & Bishop eigen-solution)."""
def __init__(
self,
latent_dim: int,
dim: int | None = None,
min_sigma2: float = _MIN_SIGMA2,
name: str | None = None,
keys: str | None = None,
) -> None:
if latent_dim is None or latent_dim < 1:
raise ValueError("ProbabilisticPCAEstimator requires latent_dim >= 1.")
self.latent_dim = int(latent_dim)
self.dim = dim
self.min_sigma2 = min_sigma2
self.name = name
self.keys = keys
[docs]
def accumulator_factory(self) -> ProbabilisticPCAAccumulatorFactory:
return ProbabilisticPCAAccumulatorFactory(dim=self.dim, keys=self.keys)
[docs]
def estimate(
self, nobs: float | None, suff_stat: tuple[float, np.ndarray | None, np.ndarray | None]
) -> ProbabilisticPCADistribution:
count, s, s2 = suff_stat
if s is None or count <= 0.0:
d = self.dim if self.dim is not None else self.latent_dim
return ProbabilisticPCADistribution(
np.zeros((d, self.latent_dim)), np.zeros(d), 1.0, name=self.name, keys=self.keys
)
d = len(s)
q = min(self.latent_dim, d)
mu = s / count
cov = s2 / count - np.outer(mu, mu)
cov = 0.5 * (cov + cov.T)
eigvals, eigvecs = np.linalg.eigh(cov)
order = np.argsort(eigvals)[::-1]
eigvals = np.clip(eigvals[order], 0.0, None)
eigvecs = eigvecs[:, order]
# sigma2 = mean of the discarded eigenvalues (the isotropic residual variance).
sigma2 = float(np.mean(eigvals[q:])) if q < d else 0.0
sigma2 = max(sigma2, self.min_sigma2)
# W = U_q (Lambda_q - sigma2 I)^{1/2}; padded with zero columns if q < latent_dim.
scale = np.sqrt(np.clip(eigvals[:q] - sigma2, 0.0, None))
w = eigvecs[:, :q] * scale[None, :]
if q < self.latent_dim:
w = np.hstack([w, np.zeros((d, self.latent_dim - q))])
return ProbabilisticPCADistribution(w, mu, sigma2, name=self.name, keys=self.keys)
[docs]
class ProbabilisticPCADataEncoder(DataSequenceEncoder):
"""Encode a sequence of length-d real vectors into an (n, d) float array."""
def __str__(self) -> str:
return "ProbabilisticPCADataEncoder"
def __eq__(self, other: object) -> bool:
return isinstance(other, ProbabilisticPCADataEncoder)
[docs]
def seq_encode(self, x: Sequence[Sequence[float]] | np.ndarray) -> np.ndarray:
rv = np.asarray(x, dtype=np.float64)
if rv.ndim != 2:
rv = rv.reshape((len(x), -1))
if rv.size and not np.all(np.isfinite(rv)):
raise ValueError("ProbabilisticPCADistribution requires finite real-vector observations.")
return rv