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,
)