def compute_gain(self, observation):
"""Compute the Kalman gain given the observation model.
Given the observation Jacobian H and covariance N (not necessarily
equal to that of the sensor), and the current covariance P, the Kalman
gain is K = P H^T(H P H^T + N)^{-1}.
Parameters
----------
observation : array-like, shape=[dim_obs]
Obtained measurement.
Returns
-------
gain : array-like, shape=[model.dim, model.dim_obs]
Kalman gain.
"""
obs_cov = self.model.get_measurement_noise_cov(self.state,
self.measurement_noise)
obs_jac = self.model.observation_jacobian(self.state, observation)
expected_cov = Matrices.mul(obs_jac, self.covariance,
Matrices.transpose(obs_jac))
innovation_cov = expected_cov + obs_cov
return Matrices.mul(self.covariance, Matrices.transpose(obs_jac),
gs.linalg.inv(innovation_cov))
def exp(self, tangent_vec, base_point, **kwargs):
"""Compute the Riemannian exponential of a tangent vector.
Parameters
----------
tangent_vec : array-like, shape=[..., n, p]
Tangent vector at a base point.
base_point : array-like, shape=[..., n, p]
Point in the Stiefel manifold.
Returns
-------
exp : array-like, shape=[..., n, p]
Point in the Stiefel manifold equal to the Riemannian exponential
of tangent_vec at the base point.
"""
p = self.p
matrix_a = Matrices.mul(Matrices.transpose(base_point), tangent_vec)
matrix_k = tangent_vec - Matrices.mul(base_point, matrix_a)
matrix_q, matrix_r = gs.linalg.qr(matrix_k)
matrix_ar = gs.concatenate([matrix_a, -Matrices.transpose(matrix_r)],
axis=-1)
zeros = gs.zeros_like(tangent_vec)[..., :p, :p]
matrix_rz = gs.concatenate([matrix_r, zeros], axis=-1)
block = gs.concatenate([matrix_ar, matrix_rz], axis=-2)
matrix_mn_e = gs.linalg.expm(block)
exp = Matrices.mul(base_point,
matrix_mn_e[..., :p, :p]) + Matrices.mul(
matrix_q, matrix_mn_e[..., p:, :p])
return exp
def test_horizontal_projection(self, k_landmarks, m_ambient, tangent_vec, point):
space = self.space(k_landmarks, m_ambient)
horizontal = space.horizontal_projection(tangent_vec, point)
transposed_point = Matrices.transpose(point)
result = gs.matmul(transposed_point, horizontal)
expected = Matrices.transpose(result)
self.assertAllClose(result, expected)
def _procrustes_preprocessing(p, matrix_v, matrix_m, matrix_n):
"""Procrustes preprocessing.
Parameters
----------
matrix_v : array-like
matrix_m : array-like
matrix_n : array-like
Returns
-------
matrix_v : array-like
"""
[matrix_d, _, matrix_r] = gs.linalg.svd(matrix_v[..., p:, p:])
matrix_v_final = gs.copy(matrix_v)
for i in range(1, p + 1):
matrix_rd = Matrices.mul(matrix_r, Matrices.transpose(matrix_d))
sub_matrix_v = gs.matmul(matrix_v[..., :, p:], matrix_rd)
matrix_v_final = gs.concatenate(
[gs.concatenate([matrix_m, matrix_n], axis=-2), sub_matrix_v],
axis=-1)
det = gs.linalg.det(matrix_v_final)
if gs.all(det > 0):
break
ones = gs.ones(p)
reflection_vec = gs.concatenate(
[ones[:-i], gs.array([-1.0] * i)], axis=0)
mask = gs.cast(det < 0, matrix_v.dtype)
sign = mask[..., None] * reflection_vec + (1.0 - mask)[...,
None] * ones
matrix_d = gs.einsum("...ij,...i->...ij",
Matrices.transpose(matrix_d), sign)
return matrix_v_final
def random_uniform(self, n_samples=1):
"""Sample random points from a uniform distribution.
Following [Chikuse03]_, :math: `n_samples * n * k` scalars are sampled
from a standard normal distribution and reshaped to matrices,
the projectors on their first k columns follow a uniform distribution.
Parameters
----------
n_samples : int
The number of points to sample
Optional. default: 1.
Returns
-------
projectors : array-like, shape=[..., n, n]
Points following a uniform distribution.
References
----------
.. [Chikuse03] Yasuko Chikuse, Statistics on special manifolds,
New York: Springer-Verlag. 2003, 10.1007/978-0-387-21540-2
"""
points = gs.random.normal(size=(n_samples, self.n, self.k))
full_rank = Matrices.mul(Matrices.transpose(points), points)
projector = Matrices.mul(points, GeneralLinear.inverse(full_rank),
Matrices.transpose(points))
return projector[0] if n_samples == 1 else projector
def submersion(point, k):
r"""Submersion that defines the Grassmann manifold.
The Grassmann manifold is defined here as embedded in the set of
symmetric matrices, as the pre-image of the function defined around the
projector on the space spanned by the first k columns of the identity
matrix by (see Exercise E.25 in [Pau07]_).
.. math:
\begin{pmatrix} I_k + A & B^T \\ B & D \end{pmatrix} \mapsto
(D - B(I_k + A)^{-1}B^T, A + A^2 + B^TB
This map is a submersion and its zero space is the set of orthogonal
rank-k projectors.
References
----------
.. [Pau07] Paulin, Frédéric. “Géométrie différentielle élémentaire,” 2007.
https://www.imo.universite-paris-saclay.fr/~paulin
/notescours/cours_geodiff.pdf.
"""
_, eigvecs = gs.linalg.eigh(point)
eigvecs = gs.flip(eigvecs, -1)
flipped_point = Matrices.mul(Matrices.transpose(eigvecs), point, eigvecs)
b = flipped_point[..., k:, :k]
d = flipped_point[..., k:, k:]
a = flipped_point[..., :k, :k] - gs.eye(k)
first = d - Matrices.mul(b, GeneralLinear.inverse(a + gs.eye(k)),
Matrices.transpose(b))
second = a + Matrices.mul(a, a) + Matrices.mul(Matrices.transpose(b), b)
row_1 = gs.concatenate([first, gs.zeros_like(b)], axis=-1)
row_2 = gs.concatenate([Matrices.transpose(gs.zeros_like(b)), second],
axis=-1)
return gs.concatenate([row_1, row_2], axis=-2)
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=[..., k_landmarks, m_ambient]
Tangent vector to the pre-shape space at `base_point`.
base_point : array-like, shape=[..., k_landmarks, m_ambient]
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=[..., k_landmarks, m_ambient]
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
def integrability_tensor(self, tangent_vec_x, tangent_vec_e, base_point):
r"""Compute the fundamental tensor A of the submersion.
The fundamental tensor A is defined for tangent vectors of the total
space by [O'Neill]_ :math:`A_X Y = ver\nabla^M_{hor X} (hor Y)
+ hor \nabla^M_{hor X}( ver Y)`
where :math:`hor, ver` are the horizontal and vertical projections.
For the Kendall shape space, we have the closed-form expression at
base-point P [Pennec]_:
:math:`A_X E = P Sylv_P(E^\top hor(X)) + F + <F,P> P` where
:math:`F = hor(X) Sylv_P(P^\top E)` and :math:`Sylv_P(B)` is the
unique skew-symmetric matrix :math:`\Omega` solution of
:math:`P^\top P \Omega + \Omega P^\top P = B - B^\top`.
Parameters
----------
tangent_vec_x : array-like, shape=[..., k_landmarks, m_ambient]
Tangent vector at `base_point`.
tangent_vec_e : array-like, shape=[..., k_landmarks, m_ambient]
Tangent vector at `base_point`.
base_point : array-like, shape=[..., k_landmarks, m_ambient]
Point of the total space.
Returns
-------
vector : array-like, shape=[..., k_landmarks, m_ambient]
Tangent vector at `base_point`, result of the A tensor applied to
`tangent_vec_x` and `tangent_vec_e`.
References
----------
.. [O'Neill] O’Neill, Barrett. The Fundamental Equations of a
Submersion, Michigan Mathematical Journal 13, no. 4 (December 1966):
459–69. https://doi.org/10.1307/mmj/1028999604.
.. [Pennec] Pennec, Xavier. Computing the curvature and its gradient
in Kendall shape spaces. Unpublished.
"""
hor_x = self.horizontal_projection(tangent_vec_x, base_point)
p_top = Matrices.transpose(base_point)
p_top_p = gs.matmul(p_top, base_point)
def sylv_p(mat_b):
"""Solves Sylvester equation for vertical component."""
return gs.linalg.solve_sylvester(
p_top_p, p_top_p, mat_b - Matrices.transpose(mat_b)
)
e_top_hor_x = gs.matmul(Matrices.transpose(tangent_vec_e), hor_x)
sylv_e_top_hor_x = sylv_p(e_top_hor_x)
p_top_e = gs.matmul(p_top, tangent_vec_e)
sylv_p_top_e = sylv_p(p_top_e)
result = gs.matmul(base_point, sylv_e_top_hor_x) + gs.matmul(
hor_x, sylv_p_top_e
)
return result
def align(self, point, base_point, **kwargs):
"""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=[..., k_landmarks, m_ambient]
Point on the manifold.
base_point : array-like, shape=[..., k_landmarks, m_ambient]
Point on the manifold.
Returns
-------
aligned : array-like, shape=[..., k_landmarks, m_ambient]
R.point.
"""
mat = gs.matmul(Matrices.transpose(point), base_point)
left, singular_values, right = gs.linalg.svd(mat)
det = gs.linalg.det(mat)
conditioning = (
singular_values[..., -2] + gs.sign(det) * singular_values[..., -1]
) / singular_values[..., 0]
if gs.any(conditioning < gs.atol):
logging.warning(
f"Singularity close, ill-conditioned matrix "
f"encountered: "
f"{conditioning[conditioning < 1e-10]}"
)
if gs.any(gs.isclose(conditioning, 0.0)):
logging.warning("Alignment matrix is not unique.")
flipped = flip_determinant(Matrices.transpose(right), det)
return Matrices.mul(point, left, Matrices.transpose(flipped))
def vertical_projection(tangent_vec, base_point):
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`
Then Ax is returned.
Parameters
----------
tangent_vec : array-like, shape=[..., k_landmarks, m_ambient]
Tangent vector to the pre-shape space at `base_point`.
base_point : array-like, shape=[..., k_landmarks, m_ambient]
Point on the pre-shape space.
Returns
-------
vertical : array-like, shape=[..., k_landmarks, 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)
return -gs.matmul(base_point, skew)
def test_horizontal_projection(self):
vector = gs.random.rand(self.k_landmarks, self.m_ambient)
point = self.space.random_point()
tan = self.space.to_tangent(vector, point)
horizontal = self.space.horizontal_projection(tan, point)
transposed_point = Matrices.transpose(point)
result = gs.matmul(transposed_point, horizontal)
expected = Matrices.transpose(result)
self.assertAllClose(result, expected)
def integrability_tensor_old(self, tangent_vec_a, tangent_vec_b,
base_point):
r"""Compute the fundamental tensor A of the submersion (old).
The fundamental tensor A is defined for tangent vectors of the total
space by [O'Neill]_ :math:`A_X Y = ver\nabla^M_{hor X} (hor Y)
+ hor \nabla^M_{hor X}( ver Y)` where :math:`hor,ver` are the
horizontal and vertical projections.
For the pre-shape space, we have closed-form expressions and the result
does not depend on the vertical part of :math:`X`.
Parameters
----------
tangent_vec_a : array-like, shape=[..., k_landmarks, m_ambient]
Tangent vector at `base_point`.
tangent_vec_b : array-like, shape=[..., k_landmarks, m_ambient]
Tangent vector at `base_point`.
base_point : array-like, shape=[..., k_landmarks, m_ambient]
Point of the total space.
Returns
-------
vector : array-like, shape=[..., k_landmarks, m_ambient]
Tangent vector at `base_point`, result of the A tensor applied to
`tangent_vec_a` and `tangent_vec_b`.
References
----------
.. [O'Neill] O’Neill, Barrett. The Fundamental Equations of a
Submersion, Michigan Mathematical Journal 13, no. 4
(December 1966): 459–69. https://doi.org/10.1307/mmj/1028999604.
"""
# Only the horizontal part of a counts
horizontal_a = self.horizontal_projection(tangent_vec_a, base_point)
vertical_b, skew = self.vertical_projection(tangent_vec_b,
base_point,
return_skew=True)
horizontal_b = tangent_vec_b - vertical_b
# For the horizontal part of b
transposed_point = Matrices.transpose(base_point)
sigma = gs.matmul(transposed_point, base_point)
alignment = gs.matmul(Matrices.transpose(horizontal_a), horizontal_b)
right_term = alignment - Matrices.transpose(alignment)
skew_hor = gs.linalg.solve_sylvester(sigma, sigma, right_term)
vertical = -gs.matmul(base_point, skew_hor)
# For the vertical part of b
vert_part = -gs.matmul(horizontal_a, skew)
tangent_vert = self.to_tangent(vert_part, base_point)
horizontal_ = self.horizontal_projection(tangent_vert, base_point)
return vertical + horizontal_
def test_vertical_projection(self, k_landmarks, m_ambient, tangent_vec, point):
space = self.space(k_landmarks, m_ambient)
vertical = space.vertical_projection(tangent_vec, point)
transposed_point = Matrices.transpose(point)
tmp_expected = gs.matmul(transposed_point, tangent_vec)
expected = Matrices.transpose(tmp_expected) - tmp_expected
tmp_result = gs.matmul(transposed_point, vertical)
result = Matrices.transpose(tmp_result) - tmp_result
self.assertAllClose(result, expected)
def force(state, time):
gamma_t = self.ambient_metric.exp(time * horizontal_b, base_point)
speed = self.ambient_metric.parallel_transport(
horizontal_b, base_point, time * horizontal_b)
coef = self.inner_product(speed, state, gamma_t)
normal = gs.einsum("...,...ij->...ij", coef, gamma_t)
align = gs.matmul(Matrices.transpose(speed), state)
right = align - Matrices.transpose(align)
left = gs.matmul(Matrices.transpose(gamma_t), gamma_t)
skew_ = gs.linalg.solve_sylvester(left, left, right)
vertical_ = -gs.matmul(gamma_t, skew_)
return vertical_ - normal
def test_vertical_projection(self):
vector = gs.random.rand(self.k_landmarks, self.m_ambient)
point = self.space.random_point()
tan = self.space.to_tangent(vector, point)
vertical = self.space.vertical_projection(tan, point)
transposed_point = Matrices.transpose(point)
tmp_expected = gs.matmul(transposed_point, tan)
expected = Matrices.transpose(tmp_expected) - tmp_expected
tmp_result = gs.matmul(transposed_point, vertical)
result = Matrices.transpose(tmp_result) - tmp_result
self.assertAllClose(result, expected)
def test_to_tangent_vec_vectorization(self):
n = self.group.n
tangent_vecs = gs.arange(self.n_samples * (n + 1) ** 2)
tangent_vecs = gs.cast(tangent_vecs, gs.float32)
tangent_vecs = gs.reshape(tangent_vecs, (self.n_samples,) + (n + 1,) * 2)
point = self.group.random_point(self.n_samples)
tangent_vecs = Matrices.mul(point, tangent_vecs)
regularized = self.group.to_tangent(tangent_vecs, point)
result = Matrices.mul(Matrices.transpose(point), regularized) + Matrices.mul(
Matrices.transpose(regularized), point
)
result = result[:, :n, :n]
expected = gs.zeros_like(result)
self.assertAllClose(result, expected)
def force(state, time):
horizontal_geodesic_t = (
1 - time) * square_root_bp + time * horizontal_velocity
geodesic_t = ((1 - time)**2 * base_point + time *
(1 - time) * partial_horizontal_velocity +
time**2 * end_point)
align = Matrices.mul(
horizontal_geodesic_t,
Matrices.transpose(horizontal_velocity - square_root_bp),
state,
)
right = align + Matrices.transpose(align)
return gs.linalg.solve_sylvester(geodesic_t, geodesic_t, -right)
def test_align_vectorization(self):
base_point = self.space.random_point()
point = self.space.random_point(2)
aligned = self.space.align(point, base_point)
alignment = gs.matmul(Matrices.transpose(aligned), base_point)
result = Matrices.is_symmetric(alignment)
self.assertTrue(gs.all(result))
base_point = self.space.random_point(2)
point = self.space.random_point()
aligned = self.space.align(point, base_point)
alignment = gs.matmul(Matrices.transpose(aligned), base_point)
result = Matrices.is_symmetric(alignment)
self.assertTrue(gs.all(result))
def log_not_from_identity(self, point, base_point):
"""Compute the group logarithm of `point` from `base_point`.
Parameters
----------
point : array-like, shape=[..., {dim, [n, n]}]
base_point : array-like, shape=[..., {dim, [n, n]}]
Returns
-------
tangent_vec : array-like, shape=[..., {dim, [n, n]}]
"""
if self.default_point_type == 'vector':
jacobian = self.jacobian_translation(point=base_point,
left_or_right='left')
point_near_id = self.compose(self.inverse(base_point), point)
log_from_id = self.log_from_identity(point=point_near_id)
log = gs.einsum('...i,...ij->...j', log_from_id,
Matrices.transpose(jacobian))
return log
lie_point = self.compose(self.inverse(base_point), point)
return self.compose(base_point, self.log_from_identity(lie_point))
def log(self, point, base_point, **kwargs):
"""Compute the Bures-Wasserstein logarithm map.
Compute the Riemannian logarithm at point base_point,
of point wrt the Bures-Wasserstein metric.
This gives a tangent vector at point base_point.
Parameters
----------
point : array-like, shape=[..., n, n]
Point.
base_point : array-like, shape=[..., n, n]
Base point.
Returns
-------
log : array-like, shape=[..., n, n]
Riemannian logarithm.
"""
# compute B^1/2(B^-1/2 A B^-1/2)B^-1/2 instead of sqrtm(AB^-1)
sqrt_bp, inv_sqrt_bp = SymmetricMatrices.powerm(
base_point, [0.5, -0.5])
pdt = SymmetricMatrices.powerm(Matrices.mul(sqrt_bp, point, sqrt_bp),
0.5)
sqrt_product = Matrices.mul(sqrt_bp, pdt, inv_sqrt_bp)
transp_sqrt_product = Matrices.transpose(sqrt_product)
return sqrt_product + transp_sqrt_product - 2 * base_point
def random_tangent_vec(self, n_samples=1, base_point=None):
"""Sample on the tangent space of SPD(n) from the uniform distribution.
Parameters
----------
n_samples : int
Number of samples.
Optional, default: 1.
base_point : array-like, shape=[..., n, n]
Base point of the tangent space.
Optional, default: None.
Returns
-------
samples : array-like, shape=[..., n, n]
Points sampled in the tangent space at base_point.
"""
n = self.n
size = (n_samples, n, n) if n_samples != 1 else (n, n)
if base_point is None:
base_point = gs.eye(n)
sqrt_base_point = gs.linalg.sqrtm(base_point)
tangent_vec_at_id = 2 * gs.random.rand(*size) - 1
tangent_vec_at_id += Matrices.transpose(tangent_vec_at_id)
tangent_vec = Matrices.mul(
sqrt_base_point, tangent_vec_at_id, sqrt_base_point)
return tangent_vec
def metric_matrix(self, base_point=None):
"""Compute inner product matrix at the tangent space at a base point.
Parameters
----------
base_point : array-like, shape=[..., dim], optional
Point in the group (the default is identity).
Returns
-------
metric_mat : array-like, shape=[..., dim, dim]
Metric matrix at base_point.
"""
if base_point is None:
return self.metric_mat_at_identity
base_point = self.group.regularize(base_point)
jacobian = self.group.jacobian_translation(
point=base_point, left_or_right=self.left_or_right)
inv_jacobian = GeneralLinear.inverse(jacobian)
inv_jacobian_transposed = Matrices.transpose(inv_jacobian)
metric_mat = Matrices.mul(
inv_jacobian_transposed, self.metric_mat_at_identity, inv_jacobian)
return metric_mat
def left_exp_from_identity(self, tangent_vec):
"""Compute the exponential from identity with the left-invariant metric.
Compute Riemannian exponential of a tangent vector at the identity
associated to the left-invariant metric.
If the method is called by a right-invariant metric, it uses the
left-invariant metric associated to the same inner-product matrix
at the identity.
Parameters
----------
tangent_vec : array-like, shape=[..., dim]
Tangent vector at identity.
Returns
-------
exp : array-like, shape=[..., dim]
Point in the group.
"""
tangent_vec = self.group.regularize_tangent_vec_at_identity(
tangent_vec=tangent_vec, metric=self)
sqrt_inner_product_mat = gs.linalg.sqrtm(self.metric_mat_at_identity)
mat = Matrices.transpose(sqrt_inner_product_mat)
exp = gs.einsum('...i,...ij->...j', tangent_vec, mat)
exp = self.group.regularize(exp)
return exp
def apply_func_to_eigvals(x, function, check_positive=False):
"""
Apply function to eigenvalues and reconstruct the matrix.
Parameters
----------
x : array_like, shape=[..., n, n]
Symmetric matrix.
function : callable
Function to apply to eigenvalues.
Returns
-------
x : array_like, shape=[..., n, n]
Symmetric matrix.
"""
eigvals, eigvecs = gs.linalg.eigh(x)
if check_positive:
if gs.any(gs.cast(eigvals, gs.float32) < 0.):
logging.warning('Negative eigenvalue encountered in'
' {}'.format(function.__name__))
eigvals = function(eigvals)
eigvals = algebra_utils.from_vector_to_diagonal_matrix(eigvals)
transp_eigvecs = Matrices.transpose(eigvecs)
reconstuction = gs.matmul(eigvecs, eigvals)
reconstuction = gs.matmul(reconstuction, transp_eigvecs)
return reconstuction
def random_uniform(self, n_samples=1):
r"""Sample on St(n,p) from the uniform distribution.
If :math:`Z(p,n) \sim N(0,1)`, then :math:`St(n,p) \sim U`,
according to Haar measure:
:math:`St(n,p) := Z(Z^TZ)^{-1/2}`.
Parameters
----------
n_samples : int, optional
Number of samples.
Returns
-------
samples : array-like, shape=[..., n, p]
Samples on the Stiefel manifold.
"""
n, p = self.n, self.p
size = (n_samples, n, p) if n_samples != 1 else (n, p)
std_normal = gs.random.normal(size=size)
std_normal_transpose = Matrices.transpose(std_normal)
aux = gs.einsum('...ij,...jk->...ik', std_normal_transpose, std_normal)
sqrt_aux = gs.linalg.sqrtm(aux)
inv_sqrt_aux = gs.linalg.inv(sqrt_aux)
samples = gs.einsum('...ij,...jk->...ik', std_normal, inv_sqrt_aux)
return samples
def belongs(self, point, tolerance=TOLERANCE):
"""Test if a point belongs to St(n,p).
Test whether the point is a p-frame in n-dimensional space,
and it is orthonormal.
Parameters
----------
point : array-like, shape=[..., n, p]
Point.
tolerance : float, optional
Tolerance at which to evaluate.
Returns
-------
belongs : array-like, shape=[..., 1]
Array of booleans evaluating if the corresponding points
belong to the Stiefel manifold.
"""
n_points, n, p = point.shape
if (n, p) != (self.n, self.p):
return gs.array([False] * n_points)
point_transpose = Matrices.transpose(point)
identity = gs.eye(p)
diff = gs.einsum('...ij,...jk->...ik', point_transpose,
point) - identity
diff_norm = gs.linalg.norm(diff, axis=(-2, -1))
belongs = gs.less_equal(diff_norm, tolerance)
belongs = gs.to_ndarray(belongs, to_ndim=1)
return belongs
def parallel_transport(self, tangent_vec_a, tangent_vec_b, base_point):
r"""Compute the parallel transport of a tangent vector.
Closed-form solution for the parallel transport of a tangent vector a
along the geodesic defined by :math: `t \mapsto exp_(base_point)(t*
tangent_vec_b)`. As the special Euclidean group endowed with its
canonical left-invariant metric is a symmetric space, parallel
transport is achieved by a geodesic symmetry, or equivalently, one step
of the pole ladder scheme.
Parameters
----------
tangent_vec_a : array-like, shape=[..., n + 1, n + 1]
Tangent vector at base point to be transported.
tangent_vec_b : array-like, shape=[..., n + 1, n + 1]
Tangent vector at base point, along which the parallel transport
is computed.
base_point : array-like, shape=[..., n + 1, n + 1]
Point on the hypersphere.
Returns
-------
transported_tangent_vec: array-like, shape=[..., n + 1, n + 1]
Transported tangent vector at `exp_(base_point)(tangent_vec_b)`.
"""
midpoint = self.exp(1. / 2. * tangent_vec_b, base_point)
transposed = Matrices.transpose(tangent_vec_a)
transported_vec = Matrices.mul(midpoint, transposed, midpoint)
return (-1.) * transported_vec
def is_tangent(self, vector, base_point):
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.
Returns
-------
belongs : array-like, shape=[...,]
Boolean denoting if vector belongs to tangent space
at base_point.
"""
vector_sym = Matrices(self.n, self.n).to_symmetric(vector)
_, r = gs.linalg.eigh(base_point)
r_ort = r[..., :, self.n - self.rank:self.n]
r_ort_t = Matrices.transpose(r_ort)
rr = gs.matmul(r_ort, r_ort_t)
candidates = Matrices.mul(rr, vector_sym, rr)
result = gs.all(gs.isclose(candidates, 0., gs.atol), axis=(-2, -1))
return result
def exp(self, tangent_vec, base_point):
"""Compute the Bures-Wasserstein exponential map.
Parameters
----------
tangent_vec : array-like, shape=[..., n, n]
Tangent vector at base point.
base_point : array-like, shape=[..., n, n]
Base point.
Returns
-------
exp : array-like, shape=[...,]
Riemannian exponential.
"""
eigvals, eigvecs = gs.linalg.eigh(base_point)
transp_eigvecs = Matrices.transpose(eigvecs)
rotated_tangent_vec = Matrices.mul(transp_eigvecs, tangent_vec,
eigvecs)
coefficients = 1 / (eigvals[..., :, None] + eigvals[..., None, :])
rotated_sylvester = rotated_tangent_vec * coefficients
rotated_hessian = gs.einsum(
'...ij,...j->...ij', rotated_sylvester, eigvals)
rotated_hessian = Matrices.mul(rotated_hessian, rotated_sylvester)
hessian = Matrices.mul(eigvecs, rotated_hessian, transp_eigvecs)
return base_point + tangent_vec + hessian
def metric_matrix(self, base_point=None):
"""Compute inner product matrix at the tangent space at a base point.
Parameters
----------
base_point : array-like, shape=[..., dim], optional
Point in the group (the default is identity).
Returns
-------
metric_mat : array-like, shape=[..., dim, dim]
Metric matrix at base_point.
"""
if self.group.default_point_type == 'matrix':
raise NotImplementedError(
'inner_product_matrix not implemented for Lie groups'
' whose elements are represented as matrices.')
if base_point is None:
base_point = self.group.identity
else:
base_point = self.group.regularize(base_point)
jacobian = self.group.jacobian_translation(
point=base_point, left_or_right=self.left_or_right)
inv_jacobian = GeneralLinear.inverse(jacobian)
inv_jacobian_transposed = Matrices.transpose(inv_jacobian)
metric_mat = gs.einsum('...ij,...jk->...ik', inv_jacobian_transposed,
self.metric_mat_at_identity)
metric_mat = gs.einsum('...ij,...jk->...ik', metric_mat, inv_jacobian)
return metric_mat
