def is_tangent(self, vector, base_point, tangent_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.
tangent_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(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.0, tangent_atol), axis=(-2, -1))
return result
def left_log_from_identity(self, point):
"""Compute Riemannian log of a point wrt. id of left-invar. metric.
Compute Riemannian logarithm of a point wrt 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
----------
point : array-like, shape=[n_samples, dimension]
Point in the group.
Returns
-------
log : array-like, shape=[n_samples, dimension]
Tangent vector at the identity equal to the Riemannian logarithm
of point at the identity.
"""
point = self.group.regularize(point)
inner_prod_mat = self.inner_product_mat_at_identity
inv_inner_prod_mat = gs.linalg.inv(inner_prod_mat)
sqrt_inv_inner_prod_mat = gs.linalg.sqrtm(inv_inner_prod_mat)
assert sqrt_inv_inner_prod_mat.shape == ((1,)
+ (self.group.dimension,) * 2)
aux = gs.squeeze(sqrt_inv_inner_prod_mat, axis=0)
log = gs.matmul(point, aux)
log = self.group.regularize_tangent_vec_at_identity(
tangent_vec=log, metric=self)
assert gs.ndim(log) == 2
return log
def test_alignment_is_symmetric(self, k_landmarks, m_ambient, point,
base_point):
space = self.space(k_landmarks, m_ambient)
aligned = space.align(point, base_point)
alignment = gs.matmul(Matrices.transpose(aligned), base_point)
result = gs.all(Matrices.is_symmetric(alignment))
self.assertTrue(result)
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 log(self, point, base_point):
"""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.
"""
product = gs.matmul(base_point, point)
sqrt_product = gs.linalg.sqrtm(product)
transp_sqrt_product = Matrices.transpose(sqrt_product)
result = sqrt_product + transp_sqrt_product - 2 * base_point
return result
def test_inner_product_matrix_and_its_inverse(self):
inner_prod_mat = self.left_diag_metric.inner_product_mat_at_identity
inv_inner_prod_mat = gs.linalg.inv(inner_prod_mat)
result = gs.matmul(inv_inner_prod_mat, inner_prod_mat)
expected = gs.eye(self.group.dimension)
expected = gs.to_ndarray(expected, to_ndim=3, axis=0)
self.assertTrue(gs.allclose(result, expected))
def squared_dist(self, point_a, point_b):
"""Compute the Bures-Wasserstein squared distance.
Compute the Riemannian squared distance between point_a and point_b.
Parameters
----------
point_a : array-like, shape=[..., n, n]
Point.
point_b : array-like, shape=[..., n, n]
Point.
Returns
-------
squared_dist : array-like, shape=[...]
Riemannian squared distance.
"""
product = gs.matmul(point_a, point_b)
sqrt_product = gs.linalg.sqrtm(product)
trace_a = gs.trace(point_a)
trace_b = gs.trace(point_b)
trace_prod = gs.trace(sqrt_product)
result = trace_a + trace_b - 2 * trace_prod
return result
def test_skew_matrix_from_vector(self):
rot_vec = gs.array([0.9])
skew_matrix = self.group.skew_matrix_from_vector(rot_vec)
result = gs.matmul(skew_matrix, skew_matrix)
diag = gs.array([-0.81, -0.81])
expected = algebra_utils.from_vector_to_diagonal_matrix(diag)
self.assertAllClose(result, expected)
def transform(self, X, y=None):
"""Project X on the principal components.
Parameters
----------
X : array-like, shape=[..., n_features]
Data, where n_samples is the number of samples
and n_features is the number of features.
y : Ignored (Compliance with scikit-learn interface)
Returns
-------
X_new : array-like, shape=[..., n_components]
Projected data.
"""
tangent_vecs = self.metric.log(X, base_point=self.base_point_fit)
if self.point_type == "matrix":
if Matrices.is_symmetric(tangent_vecs).all():
X = SymmetricMatrices.to_vector(tangent_vecs)
else:
X = gs.reshape(tangent_vecs, (len(X), -1))
else:
X = tangent_vecs
X = X - self.mean_
X_transformed = gs.matmul(X, gs.transpose(self.components_))
return X_transformed
def inverse_transform(self, X):
"""Low-dimensional reconstruction of X.
The reconstruction will match X_original whose transform would be X
if `n_components=min(n_samples, n_features)`.
Parameters
----------
X : array-like, shape=[..., n_components]
New data, where n_samples is the number of samples
and n_components is the number of components.
Returns
-------
X_original : array-like, shape=[..., n_features]
Original data.
"""
scores = self.mean_ + gs.matmul(X, self.components_)
if self.point_type == "matrix":
if Matrices.is_symmetric(self.base_point_fit).all():
scores = SymmetricMatrices(self.base_point_fit.shape[-1]).from_vector(
scores
)
else:
dim = self.base_point_fit.shape[-1]
scores = gs.reshape(scores, (len(scores), dim, dim))
return self.metric.exp(scores, self.base_point_fit)
def inner_product_at_identity(self, tangent_vec_a, tangent_vec_b):
"""
Inner product matrix at the tangent space at the identity.
"""
assert self.group.point_representation in ('vector', 'matrix')
if self.group.point_representation == 'vector':
tangent_vec_a = gs.to_ndarray(tangent_vec_a, to_ndim=2)
tangent_vec_b = gs.to_ndarray(tangent_vec_b, to_ndim=2)
inner_prod = gs.einsum('ij,ijk,ik->i', tangent_vec_a,
self.inner_product_mat_at_identity,
tangent_vec_b)
inner_prod = gs.to_ndarray(inner_prod, to_ndim=2, axis=1)
elif self.group.point_representation == 'matrix':
logging.warning(
'Only the canonical inner product -Frobenius inner product-'
' is implemented for Lie groups whose elements are represented'
' by matrices.')
tangent_vec_a = gs.to_ndarray(tangent_vec_a, to_ndim=3)
tangent_vec_b = gs.to_ndarray(tangent_vec_b, to_ndim=3)
aux_prod = gs.matmul(gs.transpose(tangent_vec_a, axes=(0, 2, 1)),
tangent_vec_b)
inner_prod = gs.trace(aux_prod)
return inner_prod
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:])
j_matrix = gs.eye(p)
for i in range(1, p):
matrix_rd = Matrices.mul(
matrix_r, j_matrix, Matrices.transpose(matrix_d))
sub_matrix_v = gs.matmul(matrix_v[..., :, p:], matrix_rd)
matrix_v = gs.concatenate([
gs.concatenate([matrix_m, matrix_n], axis=-2),
sub_matrix_v], axis=-1)
det = gs.linalg.det(matrix_v)
if gs.all(det > 0):
break
ones = gs.ones(p)
reflection_vec = gs.concatenate(
[ones[:-i], gs.array([-1.] * i)], axis=0)
mask = gs.cast(det < 0, gs.float32)
sign = (mask[..., None] * reflection_vec
+ (1. - mask)[..., None] * ones)
j_matrix = algebra_utils.from_vector_to_diagonal_matrix(sign)
return matrix_v
def exp_domain(self, tangent_vec, base_point):
base_point = gs.to_ndarray(base_point, to_ndim=3)
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=3)
invsqrt_base_point = gs.linalg.powerm(base_point, -.5)
reduced_vec = gs.matmul(invsqrt_base_point, tangent_vec)
reduced_vec = gs.matmul(reduced_vec, invsqrt_base_point)
eigvals = gs.linalg.eigvalsh(reduced_vec)
min_eig = gs.amin(eigvals, axis=1)
max_eig = gs.amax(eigvals, axis=1)
inf_value = gs.where(max_eig <= 0, -math.inf, -1 / max_eig)
inf_value = gs.to_ndarray(inf_value, to_ndim=2)
sup_value = gs.where(min_eig >= 0, math.inf, -1 / min_eig)
sup_value = gs.to_ndarray(sup_value, to_ndim=2)
domain = gs.concatenate((inf_value, sup_value), axis=1)
return domain
def left_exp_from_identity(self, tangent_vec):
"""
Compute the *left* Riemannian exponential from the identity of the
Lie group of tangent vector tangent_vec.
The left Riemannian exponential has a special role since the
left Riemannian exponential of the canonical metric parameterizes
the points.
Note: In the case where the method is called by a right-invariant
metric, it used the left-invariant metric associated to the same
inner-product at the identity.
"""
import geomstats.spd_matrices_space as spd_matrices_space
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
tangent_vec = self.group.regularize_tangent_vec_at_identity(
tangent_vec=tangent_vec, metric=self)
sqrt_inner_product_mat = spd_matrices_space.sqrtm(
self.inner_product_mat_at_identity)
mat = gs.transpose(sqrt_inner_product_mat, axes=(0, 2, 1))
exp = gs.matmul(tangent_vec, mat)
exp = gs.squeeze(exp, axis=0)
exp = self.group.regularize(exp)
return exp
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 _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 _exp_translation_transform(self, rot_vec):
"""Compute matrix associated to rot_vec for the translation part in exp.
Parameters
----------
rot_vec : array-like, shape=[..., 3]
Returns
-------
transform : array-like, shape=[..., 3, 3]
Matrix to be applied to the translation part in exp.
"""
sq_angle = gs.sum(rot_vec**2, axis=-1)
skew_mat = self.rotations.skew_matrix_from_vector(rot_vec)
sq_skew_mat = gs.matmul(skew_mat, skew_mat)
coef_1_ = utils.taylor_exp_even_func(sq_angle,
utils.cosc_close_0,
order=4)
coef_2_ = utils.taylor_exp_even_func(sq_angle,
utils.var_sinc_close_0,
order=4)
term_1 = gs.einsum('...,...ij->...ij', coef_1_, skew_mat)
term_2 = gs.einsum('...,...ij->...ij', coef_2_, sq_skew_mat)
term_id = gs.eye(3)
transform = term_id + term_1 + term_2
return transform
def lie_bracket(self, matrix_a, matrix_b):
"""Compute the Lie_bracket (commutator) of two matrices.
Notice that inputs have to be given in matrix form, no conversion
between basis and matrix representation is attempted.
Parameters
----------
matrix_a: array-like, shape=[n_sample, n, n]
matrix_b: array-like, shape=[n_sample, n, n]
Returns
-------
bracket: shape=[n_sample, n, n]
"""
return gs.matmul(matrix_a, matrix_b) - gs.matmul(matrix_b, matrix_a)
def test_inner_product(self):
base_point = gs.array([
[1., 2., 3.],
[0., 0., 0.],
[3., 1., 1.]])
tangent_vector_1 = gs.array([
[1., 2., 3.],
[0., -10., 0.],
[30., 1., 1.]])
tangent_vector_2 = gs.array([
[1., 4., 3.],
[5., 0., 0.],
[3., 1., 1.]])
result = self.metric.inner_product(
tangent_vector_1,
tangent_vector_2,
base_point=base_point)
expected = gs.trace(
gs.matmul(
gs.transpose(tangent_vector_1),
tangent_vector_2))
self.assertAllClose(result, expected)
def _make_b(i, matrix, list_matrices_r):
b = gs.ones(i + 1)
for j in range(i):
b[j] = -gs.matmul(matrix[i, :j + 1], list_matrices_r[j])
return b
def setUp(self):
"""
Tangent vectors constructed following:
http://noodle.med.yale.edu/hdtag/notes/steifel_notes.pdf
"""
warnings.filterwarnings("ignore")
gs.random.seed(1234)
self.p = 3
self.n = 4
self.space = Stiefel(self.n, self.p)
self.n_samples = 10
self.dimension = int(self.p * self.n - (self.p * (self.p + 1) / 2))
self.point_a = gs.array(
[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 0.0]]
)
self.point_b = gs.array(
[
[1.0 / gs.sqrt(2.0), 0.0, 0.0],
[0.0, 1.0, 0.0],
[0.0, 0.0, 1.0],
[1.0 / gs.sqrt(2.0), 0.0, 0.0],
]
)
point_perp = gs.array([[0.0], [0.0], [0.0], [1.0]])
matrix_a_1 = gs.array([[0.0, 2.0, -5.0], [-2.0, 0.0, -1.0], [5.0, 1.0, 0.0]])
matrix_b_1 = gs.array([[-2.0, 1.0, 4.0]])
matrix_a_2 = gs.array([[0.0, 2.0, -5.0], [-2.0, 0.0, -1.0], [5.0, 1.0, 0.0]])
matrix_b_2 = gs.array([[-2.0, 1.0, 4.0]])
self.tangent_vector_1 = gs.matmul(self.point_a, matrix_a_1) + gs.matmul(
point_perp, matrix_b_1
)
self.tangent_vector_2 = gs.matmul(self.point_a, matrix_a_2) + gs.matmul(
point_perp, matrix_b_2
)
self.metric = self.space.canonical_metric
def grad(y_pred, y_true,
metric=SO3.bi_invariant_metric,
representation='vector'):
"""Closed-form for the gradient of pose_loss.
Parameters
----------
y_pred : array-like
Prediction on SO(3).
y_true : array-like
Ground-truth on SO(3).
metric : RiemannianMetric
Metric used to compute the loss and gradient.
representation : str, {'vector', 'matrix'}
Representation chosen for points in SE(3).
Returns
-------
lie_grad : array-like
Tangent vector at point y_pred.
"""
y_pred = gs.expand_dims(y_pred, axis=0)
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'vector':
lie_grad = lie_group.grad(y_pred, y_true, SO3, metric)
if representation == 'quaternion':
quat_scalar = y_pred[:, :1]
quat_vec = y_pred[:, 1:]
quat_vec_norm = gs.linalg.norm(quat_vec, axis=1)
quat_sq_norm = quat_vec_norm ** 2 + quat_scalar ** 2
quat_arctan2 = gs.arctan2(quat_vec_norm, quat_scalar)
differential_scalar = - 2 * quat_vec / (quat_sq_norm)
differential_scalar = gs.to_ndarray(differential_scalar, to_ndim=2)
differential_scalar = gs.transpose(differential_scalar)
differential_vec = (2 * (quat_scalar / quat_sq_norm
- 2 * quat_arctan2 / quat_vec_norm)
* (gs.einsum('ni,nj->nij', quat_vec, quat_vec)
/ quat_vec_norm ** 2)
+ 2 * quat_arctan2 / quat_vec_norm * gs.eye(3))
differential_vec = gs.squeeze(differential_vec)
differential = gs.concatenate(
[differential_scalar, differential_vec],
axis=1)
y_pred = SO3.rotation_vector_from_quaternion(y_pred)
y_true = SO3.rotation_vector_from_quaternion(y_true)
lie_grad = lie_group.grad(y_pred, y_true, SO3, metric)
lie_grad = gs.matmul(lie_grad, differential)
lie_grad = gs.squeeze(lie_grad, axis=0)
return lie_grad
def compose(self, point_1, point_2):
"""
Compose two elements of group SE(3).
Formula:
point_1 . point_2 = [R1 * R2, (R1 * t2) + t1]
where:
R1, R2 are rotation matrices,
t1, t2 are translation vectors.
:param point_1, point_2: 6d vectors elements of SE(3)
:return composition: composition of point_1 and point_2
"""
rotations = self.rotations
dim_rotations = rotations.dimension
point_1 = self.regularize(point_1)
point_2 = self.regularize(point_2)
n_points_1, _ = point_1.shape
n_points_2, _ = point_2.shape
assert (point_1.shape == point_2.shape
or n_points_1 == 1
or n_points_2 == 1)
rot_vec_1 = point_1[:, :dim_rotations]
rot_mat_1 = rotations.matrix_from_rotation_vector(rot_vec_1)
rot_mat_1 = so_group.closest_rotation_matrix(rot_mat_1)
rot_vec_2 = point_2[:, :dim_rotations]
rot_mat_2 = rotations.matrix_from_rotation_vector(rot_vec_2)
rot_mat_2 = so_group.closest_rotation_matrix(rot_mat_2)
translation_1 = point_1[:, dim_rotations:]
translation_2 = point_2[:, dim_rotations:]
n_compositions = gs.maximum(n_points_1, n_points_2)
composition_rot_mat = gs.matmul(rot_mat_1, rot_mat_2)
composition_rot_vec = rotations.rotation_vector_from_matrix(
composition_rot_mat)
composition_translation = gs.zeros((n_compositions, self.n))
for i in range(n_compositions):
translation_1_i = (translation_1[0] if n_points_1 == 1
else translation_1[i])
rot_mat_1_i = (rot_mat_1[0] if n_points_1 == 1
else rot_mat_1[i])
translation_2_i = (translation_2[0] if n_points_2 == 1
else translation_2[i])
composition_translation[i] = (gs.dot(translation_2_i,
gs.transpose(rot_mat_1_i))
+ translation_1_i)
composition = gs.zeros((n_compositions, self.dimension))
composition[:, :dim_rotations] = composition_rot_vec
composition[:, dim_rotations:] = composition_translation
composition = self.regularize(composition)
return composition
def inner_product(self, tangent_vec_a, tangent_vec_b, base_point):
"""Compute the Log-Euclidean inner product.
Compute the inner product of tangent_vec_a and tangent_vec_b
at point base_point using the log-Euclidean metric.
Parameters
----------
tangent_vec_a : array-like, shape=[n_samples, n, n]
tangent_vec_b : array-like, shape=[n_samples, n, n]
base_point : array-like, shape=[n_samples, n, n]
Returns
-------
inner_product : float
"""
tangent_vec_a = gs.to_ndarray(tangent_vec_a, to_ndim=3)
n_tangent_vecs_a, _, _ = tangent_vec_a.shape
tangent_vec_b = gs.to_ndarray(tangent_vec_b, to_ndim=3)
n_tangent_vecs_b, _, _ = tangent_vec_b.shape
base_point = gs.to_ndarray(base_point, to_ndim=3)
n_base_points, _, _ = base_point.shape
spd_space = self.space
assert (n_tangent_vecs_a == n_tangent_vecs_b == n_base_points
or n_tangent_vecs_a == n_tangent_vecs_b and n_base_points == 1
or n_base_points == n_tangent_vecs_a and n_tangent_vecs_b == 1
or n_base_points == n_tangent_vecs_b and n_tangent_vecs_a == 1
or n_tangent_vecs_a == 1 and n_tangent_vecs_b == 1
or n_base_points == 1 and n_tangent_vecs_a == 1
or n_base_points == 1 and n_tangent_vecs_b == 1)
if n_tangent_vecs_a == 1:
tangent_vec_a = gs.tile(
tangent_vec_a,
(gs.maximum(n_base_points, n_tangent_vecs_b), 1, 1))
if n_tangent_vecs_b == 1:
tangent_vec_b = gs.tile(
tangent_vec_b,
(gs.maximum(n_base_points, n_tangent_vecs_a), 1, 1))
if n_base_points == 1:
base_point = gs.tile(
base_point,
(gs.maximum(n_tangent_vecs_a, n_tangent_vecs_b), 1, 1))
modified_tangent_vec_a = spd_space.differential_log(
tangent_vec_a, base_point)
modified_tangent_vec_b = spd_space.differential_log(
tangent_vec_b, base_point)
product = gs.matmul(modified_tangent_vec_a, modified_tangent_vec_b)
inner_product = gs.trace(product, axis1=1, axis2=2)
inner_product = gs.to_ndarray(inner_product, to_ndim=2, axis=1)
return inner_product
def test_cong(self):
base_point = gs.array([
[1., 2., 3.],
[0., 0., 0.],
[3., 1., 1.]])
tangent_vector = gs.array([
[1., 2., 3.],
[0., -10., 0.],
[30., 1., 1.]])
result = self.space.congruent(tangent_vector, base_point)
expected = gs.matmul(
tangent_vector, gs.transpose(base_point))
expected = gs.matmul(base_point, expected)
self.assertAllClose(result, expected)
def _aux_inner_product(self, tangent_vec_a, tangent_vec_b, inv_base_point):
"""Compute the inner product (auxiliary).
Parameters
----------
tangent_vec_a : array-like, shape=[n_samples, n, n]
tangent_vec_b : array-like, shape=[n_samples, n, n]
inv_base_point : array-like, shape=[n_samples, n, n]
Returns
-------
inner_product : array-like, shape=[n_samples, n, n]
"""
aux_a = gs.matmul(inv_base_point, tangent_vec_a)
aux_b = gs.matmul(inv_base_point, tangent_vec_b)
inner_product = gs.trace(gs.matmul(aux_a, aux_b), axis1=1, axis2=2)
return inner_product
def inner_product(self, tangent_vec_a, tangent_vec_b, base_point):
r"""Compute the inner-product of two tangent vectors at a base point.
Canonical inner-product on the tangent space at `base_point`,
which is different from the inner-product induced by the embedding
(see [RLSMRZ2017]_).
.. math::
\langle\Delta, \tilde{\Delta}\rangle_{U}=\operatorname{tr}
\left(\Delta^{T}\left(I-\frac{1}{2} U U^{T}\right)
\tilde{\Delta}\right)
References
----------
.. [RLSMRZ2017] R Zimmermann. A matrix-algebraic algorithm for the
Riemannian logarithm on the Stiefel manifold under the canonical
metric. SIAM Journal on Matrix Analysis and Applications 38 (2),
322-342, 2017. https://epubs.siam.org/doi/pdf/10.1137/16M1074485
Parameters
----------
tangent_vec_a : array-like, shape=[n_samples, n, p]
First tangent vector at base point.
tangent_vec_b : array-like, shape=[n_samples, n, p]
Second tangent vector at base point.
base_point : array-like, shape=[n_samples, n, p]
Point in the Stiefel manifold.
Returns
-------
inner_prod : array-like, shape=[n_samples, 1]
Inner-product of the two tangent vectors.
"""
tangent_vec_a = gs.to_ndarray(tangent_vec_a, to_ndim=3)
tangent_vec_b = gs.to_ndarray(tangent_vec_b, to_ndim=3)
base_point = gs.to_ndarray(base_point, to_ndim=3)
base_point_transpose = gs.transpose(base_point, axes=(0, 2, 1))
aux = gs.matmul(
gs.transpose(tangent_vec_a, axes=(0, 2, 1)),
gs.eye(self.n) - 0.5 * gs.matmul(base_point, base_point_transpose))
inner_prod = gs.trace(gs.matmul(aux, tangent_vec_b), axis1=1, axis2=2)
inner_prod = gs.to_ndarray(inner_prod, to_ndim=2, axis=1)
return inner_prod
def _log_translation_transform(self, rot_vec):
"""Compute matrix associated to rot_vec for the translation part in log.
Parameters
----------
rot_vec : array-like, shape=[..., 3]
Returns
-------
transform : array-like, shape=[..., 3, 3]
Matrix to be applied to the translation part in log
"""
n_samples = rot_vec.shape[0]
angle = gs.linalg.norm(rot_vec, axis=1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
skew_mat = self.rotations.skew_matrix_from_vector(rot_vec)
sq_skew_mat = gs.matmul(skew_mat, skew_mat)
mask_close_0 = gs.isclose(angle, 0.0)
mask_close_pi = gs.isclose(angle, gs.pi)
mask_else = ~mask_close_0 & ~mask_close_pi
mask_close_0_float = gs.cast(mask_close_0, gs.float32)
mask_close_pi_float = gs.cast(mask_close_pi, gs.float32)
mask_else_float = gs.cast(mask_else, gs.float32)
mask_0 = gs.isclose(angle, 0.0, atol=1e-7)
mask_0_float = gs.cast(mask_0, gs.float32)
angle += mask_0_float * gs.ones_like(angle)
coef_1 = -0.5 * gs.ones_like(angle)
coef_2 = gs.zeros_like(angle)
coef_2 += mask_close_0_float * (1.0 / 12.0 + angle**2 / 720.0 +
angle**4 / 30240.0 +
angle**6 / 1209600.0)
delta_angle = angle - gs.pi
coef_2 += mask_close_pi_float * (
1.0 / PI2 + (PI2 - 8.0) * delta_angle / (4.0 * PI3) -
((PI2 - 12.0) * delta_angle**2 / (4.0 * PI4)) +
((-192.0 + 12.0 * PI2 + PI4) * delta_angle**3 / (48.0 * PI5)) -
((-240.0 + 12.0 * PI2 + PI4) * delta_angle**4 / (48.0 * PI6)) +
((-2880.0 + 120.0 * PI2 + 10.0 * PI4 + PI6) * delta_angle**5 /
(480.0 * PI7)) -
((-3360 + 120.0 * PI2 + 10.0 * PI4 + PI6) * delta_angle**6 /
(480.0 * PI8)))
psi = 0.5 * angle * gs.sin(angle) / (1 - gs.cos(angle))
coef_2 += mask_else_float * (1 - psi) / (angle**2)
term_1 = gs.einsum("...i,...ij->...ij", coef_1, skew_mat)
term_2 = gs.einsum("...i,...ij->...ij", coef_2, sq_skew_mat)
term_id = gs.array([gs.eye(3)] * n_samples)
transform = term_id + term_1 + term_2
return transform
def compose(self, point_1, point_2, point_type=None):
"""
Compose two elements of SE(n).
Formula:
point_1 . point_2 = [R1 * R2, (R1 * t2) + t1]
where:
R1, R2 are rotation matrices,
t1, t2 are translation vectors.
"""
if point_type is None:
point_type = self.default_point_type
rotations = self.rotations
dim_rotations = rotations.dimension
point_1 = self.regularize(point_1, point_type=point_type)
point_2 = self.regularize(point_2, point_type=point_type)
if point_type == 'vector':
n_points_1, _ = point_1.shape
n_points_2, _ = point_2.shape
assert (point_1.shape == point_2.shape or n_points_1 == 1
or n_points_2 == 1)
if n_points_1 == 1:
point_1 = gs.stack([point_1[0]] * n_points_2)
if n_points_2 == 1:
point_2 = gs.stack([point_2[0]] * n_points_1)
rot_vec_1 = point_1[:, :dim_rotations]
rot_mat_1 = rotations.matrix_from_rotation_vector(rot_vec_1)
rot_mat_1 = rotations.projection(rot_mat_1)
rot_vec_2 = point_2[:, :dim_rotations]
rot_mat_2 = rotations.matrix_from_rotation_vector(rot_vec_2)
rot_mat_2 = rotations.projection(rot_mat_2)
translation_1 = point_1[:, dim_rotations:]
translation_2 = point_2[:, dim_rotations:]
composition_rot_mat = gs.matmul(rot_mat_1, rot_mat_2)
composition_rot_vec = rotations.rotation_vector_from_matrix(
composition_rot_mat)
composition_translation = gs.einsum('ij,ikj->ik', translation_2,
rot_mat_1) + translation_1
composition = gs.concatenate(
(composition_rot_vec, composition_translation), axis=1)
elif point_type == 'matrix':
raise NotImplementedError()
composition = self.regularize(composition, point_type=point_type)
return composition
def compose(self, point_a, point_b, point_type=None):
r"""Compose two elements of SE(n).
Parameters
----------
point_1 : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
point_2 : array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
point_type: str, {'vector', 'matrix'}, optional
default: self.default_point_type
Equation
---------
(:math: `(R_1, t_1) \\cdot (R_2, t_2) = (R_1 R_2, R_1 t_2 + t_1)`)
Returns
-------
composition : the composition of point_1 and point_2
"""
rotations = self.rotations
dim_rotations = rotations.dim
point_a = self.regularize(point_a, point_type=point_type)
point_b = self.regularize(point_b, point_type=point_type)
if point_type == 'vector':
n_points_a, _ = point_a.shape
n_points_b, _ = point_b.shape
if not (point_a.shape == point_b.shape or n_points_a == 1
or n_points_b == 1):
raise ValueError()
rot_vec_a = point_a[:, :dim_rotations]
rot_mat_a = rotations.matrix_from_rotation_vector(rot_vec_a)
rot_vec_b = point_b[:, :dim_rotations]
rot_mat_b = rotations.matrix_from_rotation_vector(rot_vec_b)
translation_a = point_a[:, dim_rotations:]
translation_b = point_b[:, dim_rotations:]
composition_rot_mat = gs.matmul(rot_mat_a, rot_mat_b)
composition_rot_vec = rotations.rotation_vector_from_matrix(
composition_rot_mat)
composition_translation = gs.einsum(
'...j,...kj->...k', translation_b, rot_mat_a) + translation_a
composition = gs.concatenate(
(composition_rot_vec, composition_translation), axis=-1)
return self.regularize(composition, point_type=point_type)
if point_type == 'matrix':
return GeneralLinear.compose(point_a, point_b)
raise ValueError('Invalid point_type, expected \'vector\' or '
'\'matrix\'.')
