# Source code for pymanopt.manifolds.group

import numpy as np
import scipy.special

from pymanopt.manifolds.manifold import RiemannianSubmanifold
from pymanopt.tools import extend_docstring
from pymanopt.tools.multi import (
multiexpm,
multihconj,
multiherm,
multilogm,
multiqr,
multiskew,
multiskewh,
multitransp,
)

class _UnitaryBase(RiemannianSubmanifold):
_n: int
_k: int

def __init__(self, name, dimension, retraction):
super().__init__(name, dimension)

try:
self._retraction = getattr(self, f"_retraction_{retraction}")
except AttributeError:
raise ValueError(f"Invalid retraction type '{retraction}'")

def inner_product(self, point, tangent_vector_a, tangent_vector_b):
return np.tensordot(
tangent_vector_a.conj(),
tangent_vector_b,
axes=tangent_vector_a.ndim,
)

def norm(self, point, tangent_vector):
return np.linalg.norm(tangent_vector)

@property
def typical_dist(self):
return np.pi * np.sqrt(self._n * self._k)

def dist(self, point_a, point_b):
return self.norm(point_a, self.log(point_a, point_b))

def projection(self, point, vector):
return multiskew(multihconj(point) @ vector)

def to_tangent_space(self, point, vector):
return multiskewh(vector)

def embedding(self, point, tangent_vector):
return point @ tangent_vector

def euclidean_to_riemannian_hessian(
):
Xt = multihconj(point)
Xtehess = Xt @ euclidean_hessian
return multiskewh(Xtehess - tangent_vector @ symXtegrad)

def retraction(self, point, tangent_vector):
return self._retraction(point, tangent_vector)

def _retraction_qr(self, point, tangent_vector):
Y = point + point @ tangent_vector
q, _ = multiqr(Y)
return q

def _retraction_polar(self, point, tangent_vector):
Y = point + point @ tangent_vector
u, _, vt = np.linalg.svd(Y)
return u @ vt

def exp(self, point, tangent_vector):
return point @ multiexpm(tangent_vector)

def log(self, point_a, point_b):
return multiskewh(multilogm(multihconj(point_a) @ point_b))

def zero_vector(self, point):
zero = np.zeros((self._k, self._n, self._n))
if self._k == 1:
return zero[0]
return zero

def transport(self, point_a, point_b, tangent_vector_a):
return tangent_vector_a

def pair_mean(self, point_a, point_b):
return self.exp(point_a, self.log(point_a, point_b) / 2)

DOCSTRING_NOTE = """
Args:
n: The dimension of the space that elements of the group act on.
k: The number of elements in the product of groups.
retraction: The type of retraction to use.
Possible choices are qr and polar.

Note:
The default QR-based retraction is only a first-order approximation of
the exponential map.
Use of an SVD-based second-order retraction can be enabled by setting
the retraction argument to "polar".

The procedure to generate random points on the manifold sampled
uniformly from the Haar measure is detailed in [Mez2006]_.
"""

[docs]@extend_docstring(DOCSTRING_NOTE)
class SpecialOrthogonalGroup(_UnitaryBase):
r"""The (product) manifold of rotation matrices.

The special orthgonal group :math:\SO(n).
Points on the manifold are matrices :math:\vmQ \in \R^{n
\times n} such that each matrix is orthogonal with determinant 1, i.e.,
:math:\transp{\vmQ}\vmQ = \vmQ\transp{\vmQ} = \Id_n and :math:\det(\vmQ)
= 1.
For k > 1, the class can be used to optimize over the product manifold
of rotation matrices :math:\SO(n)^k.
In that case points on the manifold are represented as arrays of shape
(k, n, n).

The metric is the usual Euclidean one inherited from the embedding space
:math:(\R^{n \times n})^k.
As such :math:\SO(n)^k forms a Riemannian submanifold.

The tangent space :math:\tangent{\vmQ}\SO(n) at a point :math:\vmQ is
given by :math:\tangent{\vmQ}\SO(n) = \set{\vmQ \vmOmega \in \R^{n \times
n} \mid \vmOmega = -\transp{\vmOmega}} = \vmQ \Skew(n), where
:math:\Skew(n) denotes the set of skew-symmetric matrices.
This corresponds to the Lie algebra of :math:\SO(n), a fact which is used
here to conveniently represent tangent vectors numerically by their
skew-symmetric factor.
The method :meth:embedding can be used to transform a tangent vector from
its Lie algebra representation to the embedding space representation.
"""

def __init__(self, n: int, *, k: int = 1, retraction: str = "qr"):
self._n = n
self._k = k

if k == 1:
name = f"Special orthogonal group SO({n})"
elif k > 1:
name = f"Sphecial orthogonal group SO({n})^{k}"
else:
raise ValueError("k must be an integer no less than 1.")
dimension = int(k * scipy.special.comb(n, 2))
super().__init__(name, dimension, retraction)

[docs]    def random_point(self):
n, k = self._n, self._k
if n == 1:
point = np.ones((k, 1, 1))
else:
point, _ = multiqr(np.random.normal(size=(k, n, n)))
# Swap the first two columns of matrices where det(point) < 0 to
# flip the sign of their determinants.
negative_det, *_ = np.where(np.linalg.det(point) < 0)
negative_det = np.expand_dims(negative_det, (-2, -1))
point[negative_det, :, [0, 1]] = point[negative_det, :, [1, 0]]
if k == 1:
return point[0]
return point

[docs]    def random_tangent_vector(self, point):
vector = _random_skew_symmetric_matrix(self._n, self._k)
if self._k == 1:
vector = vector[0]
return vector / self.norm(point, vector)

[docs]@extend_docstring(DOCSTRING_NOTE)
class UnitaryGroup(_UnitaryBase):
r"""The (product) manifold of unitary matrices (i.e., the unitary group).

The unitary group :math:\U(n).
Points on the manifold are matrices :math:\vmX \in \C^{n
\times n} such that each matrix is unitary, i.e.,
:math:\transp{\conj{\vmX}}\vmX = \adj{\vmX}\vmX = \Id_n.
For k > 1, the class represents the product manifold
of unitary matrices :math:\U(n)^k.
In that case points on the manifold are represented as arrays of shape
(k, n, n).

The metric is the usual Euclidean one inherited from the embedding space
:math:(\C^{n \times n})^k, i.e., :math:\inner{\vmA}{\vmB} =
\Re\tr(\adj{\vmA}\vmB).
As such :math:\U(n)^k forms a Riemannian submanifold.

The tangent space :math:\tangent{\vmX}\U(n) at a point :math:\vmX is
given by :math:\tangent{\vmX}\U(n) = \set{\vmX \vmOmega \in \C^{n \times
n} \mid \vmOmega = -\adj{\vmOmega}} = \vmX \adj{\Skew}(n), where
:math:\adj{\Skew}(n) denotes the set of skew-Hermitian matrices.
This corresponds to the Lie algebra of :math:\U(n), a fact which is used
here to conveniently represent tangent vectors numerically by their
skew-Hermitian factor.
The method :meth:embedding can be used to convert a tangent vector from
its Lie algebra representation to the embedding space representation.
"""

def __init__(self, n: int, *, k: int = 1, retraction: str = "qr"):
self._n = n
self._k = k

if k == 1:
name = f"Unitary group U({n})"
elif k > 1:
name = f"Unitary group U({n})^{k}"
else:
raise ValueError("k must be an integer no less than 1.")
dimension = int(k * n**2)
super().__init__(name, dimension, retraction)

[docs]    def random_point(self):
n, k = self._n, self._k
if n == 1:
point = np.ones((k, 1, 1)) + 1j * np.ones((k, 1, 1))
point /= np.abs(point)
else:
point, _ = multiqr(
np.random.normal(size=(k, n, n))
+ 1j * np.random.normal(size=(k, n, n))
)
if k == 1:
return point[0]
return point

[docs]    def random_tangent_vector(self, point):
n, k = self._n, self._k
vector = (
_random_skew_symmetric_matrix(n, k)
+ 1j * _random_symmetric_matrix(n, k)
) / np.sqrt(2)
if k == 1:
vector = vector[0]
return vector / self.norm(point, vector)

def _random_skew_symmetric_matrix(n, k):
if n == 1:
return np.zeros((k, 1, 1))
vector = _random_upper_triangular_matrix(n, k)
return vector - multitransp(vector)

def _random_symmetric_matrix(n, k):
if n == 1:
return np.random.normal(size=(k, 1, 1))
vector = _random_upper_triangular_matrix(n, k)
vector = vector + multitransp(vector)
# The diagonal elements get scaled by a factor of 2 by the previous
# operation so re-draw them so every entry of the returned matrix follows a
# standard normal distribution.
indices = np.arange(n)
vector[:, indices, indices] = np.random.normal(size=(k, n))
return vector

def _random_upper_triangular_matrix(n, k):
if n < 2:
raise ValueError("Matrix dimension cannot be less than 2")
indices = np.triu_indices(n, 1)
vector = np.zeros((k, n, n))
vector[(slice(None), *indices)] = np.random.normal(
size=(k, n * (n - 1) // 2)
)
return vector