Source code for geomstats.geometry.rank_k_psd_matrices

r"""The manifold of Positive Semi Definite matrices of rank k PSD(n,k).

Lead author: Anna Calissano.
"""

import geomstats.backend as gs
from geomstats.geometry.fiber_bundle import FiberBundle
from geomstats.geometry.full_rank_matrices import FullRankMatrices
from geomstats.geometry.general_linear import GeneralLinear
from geomstats.geometry.group_action import SpecialOrthogonalComposeAction
from geomstats.geometry.manifold import Manifold, register_quotient
from geomstats.geometry.matrices import Matrices, MatricesMetric
from geomstats.geometry.quotient_metric import QuotientMetric
from geomstats.geometry.spd_matrices import SPDEuclideanMetric, SPDMatrices
from geomstats.geometry.symmetric_matrices import SymmetricMatrices


[docs] class RankKPSDMatrices(Manifold): r"""Class for PSD(n,k). The manifold of Positive Semi Definite matrices of rank k: PSD(n,k). Parameters ---------- n : int Integer representing the shape of the matrices: n x n. k : int Integer representing the rank of the matrix (k<n). """ def __init__(self, n, k, equip=True): if not k < n: raise ValueError("k must be lower than n") super().__init__( dim=int(k * n - k * (k + 1) / 2), shape=(n, n), equip=equip, intrinsic=False, ) self.n = n self.rank = k self.sym = SymmetricMatrices(self.n, equip=False)
[docs] @staticmethod def default_metric(): """Metric to equip the space with if equip is True.""" return PSDEuclideanMetric
[docs] def belongs(self, point, atol=gs.atol): r"""Check if the matrix belongs to the space. Parameters ---------- mat : array-like, shape=[..., n, n] Matrix to be checked. atol : float Tolerance. Optional, default: backend atol. Returns ------- belongs : array-like, shape=[...,] Boolean denoting if mat is an SPD matrix. """ shape_belongs = self.shape == point.shape[-self.point_ndim :] if not shape_belongs: shape = point.shape[: -self.point_ndim] return gs.zeros(shape, dtype=bool) is_symmetric = self.sym.belongs(point, atol) eigvalues = gs.linalg.eigvalsh(point) is_semipositive = gs.all(eigvalues > -atol, axis=-1) sum_eig = gs.sum(gs.where(eigvalues < atol, 0, 1), axis=-1) is_rankk = gs.isclose(sum_eig, self.rank, atol=atol) return gs.logical_and(gs.logical_and(is_symmetric, is_semipositive), is_rankk)
[docs] def projection(self, point): r"""Project a matrix to the space of PSD matrices of rank k. The nearest symmetric positive semidefinite matrix in the Frobenius norm to an arbitrary real matrix A is shown to be (B + H)/2, where H is the symmetric polar factor of B=(A + A')/2. As [Higham1988]_ is turning the matrix into a PSD, the rank is then forced to be k. Parameters ---------- point : array-like, shape=[..., n, n] Matrix to project. Returns ------- projected : array-like, shape=[..., n, n] PSD matrix rank k. References ---------- .. [Higham1988] Highamm, N. J. “Computing a nearest symmetric positive semidefinite matrix.” Linear Algebra and Its Applications 103 (May 1, 1988): 103-118. https://doi.org/10.1016/0024-3795(88)90223-6 """ sym = Matrices.to_symmetric(point) _, s, v = gs.linalg.svd(sym) h = gs.matmul(Matrices.transpose(v), s[..., None] * v) sym_proj = (sym + h) / 2 eigvals, eigvecs = gs.linalg.eigh(sym_proj) k = self.n - self.rank i = gs.array([0.0] * k + [2 * gs.atol] * self.rank) regularized = ( gs.hstack([gs.zeros(point.shape[:-2] + (k,)), eigvals[..., k:]]) + i ) reconstruction = gs.einsum("...ij,...j->...ij", eigvecs, regularized) return Matrices.mul(reconstruction, Matrices.transpose(eigvecs))
[docs] def random_point(self, n_samples=1, bound=1.0): r"""Sample in PSD(n,k) from the log-uniform distribution. Parameters ---------- n_samples : int Number of samples. Optional, default: 1. bound : float Bound of the interval in which to sample in the tangent space. Optional, default: 1. Returns ------- samples : array-like, shape=[..., n, n] Points sampled in PSD(n,k). """ n = self.n size = (n_samples, n, n) if n_samples != 1 else (n, n) mat = bound * (2 * gs.random.rand(*size) - 1) spd_mat = GeneralLinear.exp(Matrices.to_symmetric(mat)) return self.projection(spd_mat)
[docs] def is_tangent(self, vector, base_point, atol=gs.atol): r"""Check if the vector belongs to the tangent space. Parameters ---------- vector : array-like, shape=[..., n, n] Matrix to check if it belongs to the tangent space. base_point : array-like, shape=[..., n, n] Base point of the tangent space. Optional, default: None. atol : float Absolute tolerance. Optional, default: backend atol. Returns ------- belongs : array-like, shape=[...,] Boolean denoting if vector belongs to tangent space at base_point. """ vector_sym = Matrices.to_symmetric(vector) r, _, _ = gs.linalg.svd(base_point) r_ort = r[..., :, -(self.n - self.rank) :] r_ort_t = Matrices.transpose(r_ort) rr = gs.matmul(r_ort, r_ort_t) candidates = Matrices.mul(rr, vector_sym, rr) return gs.all(gs.isclose(candidates, 0.0, atol=atol), axis=(-2, -1))
[docs] def to_tangent(self, vector, base_point): r"""Project to the tangent space of PSD(n,k) at base_point. Parameters ---------- vector : array-like, shape=[..., n, n] Matrix to check if it belongs to the tangent space. base_point : array-like, shape=[..., n, n] Base point of the tangent space. Optional, default: None. Returns ------- tangent : array-like, shape=[..., n, n] Projection of the tangent vector at base_point. """ vector_sym = Matrices.to_symmetric(vector) r, _, _ = gs.linalg.svd(base_point) r_ort = r[..., :, : (self.rank - 1)] r_ort_t = Matrices.transpose(r_ort) rr = gs.matmul(r_ort, r_ort_t) return vector_sym - Matrices.mul(rr, vector_sym, rr)
PSDEuclideanMetric = SPDEuclideanMetric
[docs] class PSDMatrices: r"""Class for the psd matrices. The class is redirecting to the correct embedding manifold. The stratum PSD rank k if the matrix is not full rank. The top stratum SPD if the matrix is full rank. The whole stratified space of PSD if no rank is specified. Parameters ---------- n : int Integer representing the shapes of the matrices : n x n. k : int Integer representing the rank of the matrices. """ def __new__( cls, n, k, equip=True, ): """Instantiate class from one of the parent classes.""" if n > k: return RankKPSDMatrices(n, k, equip=equip) if n == k: return SPDMatrices(n, equip=equip) raise NotImplementedError("The PSD matrices is not implemented yet.")
[docs] class BuresWassersteinBundle(FiberBundle): """Class for the quotient structure on PSD matrices.""" def __init__(self, total_space): super().__init__(total_space=total_space, aligner=True)
[docs] @staticmethod def riemannian_submersion(point): """Project.""" return Matrices.mul(point, Matrices.transpose(point))
[docs] def tangent_riemannian_submersion(self, tangent_vec, base_point): """Differential.""" product = Matrices.mul(base_point, Matrices.transpose(tangent_vec)) return product + Matrices.transpose(product)
[docs] def lift(self, point): """Find a representer in top space.""" k = self._total_space.k eigvals, eigvecs = gs.linalg.eigh(point) return gs.einsum( "...ij,...j->...ij", eigvecs[..., -k:], eigvals[..., -k:] ** 0.5 )
[docs] def horizontal_lift(self, tangent_vec, base_point=None, fiber_point=None): """Horizontal lift of a tangent vector.""" n = self._total_space.n if fiber_point is None: fiber_point = self.lift(base_point) transposed_point = Matrices.transpose(fiber_point) alignment = Matrices.mul(transposed_point, fiber_point) projector = Matrices.mul(fiber_point, GeneralLinear.inverse(alignment)) right_term = Matrices.mul(transposed_point, tangent_vec, fiber_point) sylvester = gs.linalg.solve_sylvester(alignment, alignment, right_term) skew_term = Matrices.mul(projector, sylvester) orth_proj = gs.eye(n) - Matrices.mul(projector, transposed_point) orth_part = Matrices.mul(orth_proj, tangent_vec, projector) return skew_term + orth_part
[docs] def vertical_projection(self, tangent_vec, base_point, return_skew=False): r"""Project to vertical subspace. Compute the vertical component of a tangent vector :math:`w` at a base point :math:`x` by solving the sylvester equation: .. math:: Axx^T + xx^TA = wx^T - xw^T where A is skew-symmetric. Then Ax is the vertical projection of w. Parameters ---------- tangent_vec : array-like, shape=[..., n, k] Tangent vector to the pre-shape space at `base_point`. base_point : array-like, shape=[..., n, k] Point on the pre-shape space. return_skew : bool Whether to return the skew-symmetric matrix A. Optional, default: False Returns ------- vertical : array-like, shape=[..., n, k] Vertical component of `tangent_vec`. skew : array-like, shape=[..., m_ambient, m_ambient] Vertical component of `tangent_vec`. """ transposed_point = Matrices.transpose(base_point) left_term = gs.matmul(transposed_point, base_point) alignment = gs.matmul(Matrices.transpose(tangent_vec), base_point) right_term = alignment - Matrices.transpose(alignment) skew = gs.linalg.solve_sylvester(left_term, left_term, right_term) vertical = -gs.matmul(base_point, skew) return (vertical, skew) if return_skew else vertical
[docs] def align(self, point, base_point): """Align point to base_point. Find the optimal rotation R in SO(m) such that the base point and R.point are well positioned. Parameters ---------- point : array-like, shape=[..., n, k] Point on the manifold. base_point : array-like, shape=[..., n, k] Point on the manifold. Returns ------- aligned : array-like, shape=[..., n, k] R.point. """ return Matrices.align_matrices(point, base_point)
[docs] class PSDBuresWassersteinMetric(QuotientMetric): """Bures-Wasserstein metric for fixed rank PSD matrices.""" def __init__(self, space, total_space=None): if total_space is None: k = space.rank if hasattr(space, "rank") else space.n total_space = FullRankMatrices(space.n, k, equip=False) total_space.equip_with_metric(MatricesMetric) if not hasattr(total_space, "group_action"): total_space.equip_with_group_action( SpecialOrthogonalComposeAction(total_space.k) ) if not hasattr(total_space, "quotient"): total_space.equip_with_quotient() super().__init__(space=space, total_space=total_space)
register_quotient( Space=FullRankMatrices, Metric=MatricesMetric, GroupAction=SpecialOrthogonalComposeAction, FiberBundle=BuresWassersteinBundle, )