def _exp_translation_transform(self, rot_vec):
n_samples = rot_vec.shape[0]
base_1 = gs.array([gs.eye(2)] * n_samples)
base_2 = -gs.array(
[self.rotations.skew_matrix_from_vector(gs.ones(1))] * n_samples)
mask_close_0 = gs.isclose(rot_vec, 0.)
mask_else = ~mask_close_0
mask_close_0_float = gs.cast(mask_close_0, gs.float32)
mask_else_float = gs.cast(mask_else, gs.float32)
cos_coef = gs.zeros_like(rot_vec)
sin_coef = gs.zeros_like(rot_vec)
cos_coef += mask_close_0_float * (
TAYLOR_COEFFS_1_AT_0[0] * rot_vec
+ TAYLOR_COEFFS_1_AT_0[1] * rot_vec ** 3
+ TAYLOR_COEFFS_1_AT_0[2] * rot_vec ** 5)
sin_coef += mask_close_0_float * (
1
- TAYLOR_COEFFS_2_AT_0[0] * rot_vec ** 2
- TAYLOR_COEFFS_2_AT_0[2] * rot_vec ** 4)
rot_vec = rot_vec + mask_close_0_float * 1e-6
cos_coef += mask_else_float * ((1. - gs.cos(rot_vec)) / rot_vec)
sin_coef += mask_else_float * (gs.sin(rot_vec) / rot_vec)
sin_term = gs.einsum('...i,...jk->...jk', sin_coef, base_1)
cos_term = gs.einsum('...i,...jk->...jk', cos_coef, base_2)
transform = sin_term + cos_term
return transform
def test_integrability_tensor_derivative_is_alternate(self):
r"""Integrability tensor derivatives is alternate in pre-shape.
For two horizontal vector fields :math:`X,Y` the integrability
tensor (hence its derivatives) is alternate:
:math:`\nabla_X ( A_Y Z + A_Z Y ) = 0`.
"""
nabla_x_a_y_z, a_y_z = self.space.integrability_tensor_derivative(
self.hor_x,
self.hor_y,
self.nabla_x_y,
self.hor_z,
self.nabla_x_z,
self.base_point,
)
nabla_x_a_z_y, a_z_y = self.space.integrability_tensor_derivative(
self.hor_x,
self.hor_z,
self.nabla_x_z,
self.hor_y,
self.nabla_x_y,
self.base_point,
)
result = nabla_x_a_y_z + nabla_x_a_z_y
self.assertAllClose(a_y_z + a_z_y, gs.zeros_like(result))
self.assertAllClose(result, gs.zeros_like(result))
def log(self, point, base_point):
"""
Riemannian logarithm of a point wrt a base point.
"""
point = gs.to_ndarray(point, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
norm_base_point = self.embedding_metric.norm(base_point)
norm_point = self.embedding_metric.norm(point)
inner_prod = self.embedding_metric.inner_product(base_point, point)
cos_angle = inner_prod / (norm_base_point * norm_point)
cos_angle = gs.clip(cos_angle, -1., 1.)
angle = gs.arccos(cos_angle)
angle = gs.to_ndarray(angle, to_ndim=1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
mask_0 = gs.isclose(angle, 0.)
mask_else = gs.equal(mask_0, gs.array(False))
mask_0_float = gs.cast(mask_0, gs.float32)
mask_else_float = gs.cast(mask_else, gs.float32)
coef_1 = gs.zeros_like(angle)
coef_2 = gs.zeros_like(angle)
coef_1 += mask_0_float * (1. + INV_SIN_TAYLOR_COEFFS[1] * angle**2 +
INV_SIN_TAYLOR_COEFFS[3] * angle**4 +
INV_SIN_TAYLOR_COEFFS[5] * angle**6 +
INV_SIN_TAYLOR_COEFFS[7] * angle**8)
coef_2 += mask_0_float * (1. + INV_TAN_TAYLOR_COEFFS[1] * angle**2 +
INV_TAN_TAYLOR_COEFFS[3] * angle**4 +
INV_TAN_TAYLOR_COEFFS[5] * angle**6 +
INV_TAN_TAYLOR_COEFFS[7] * angle**8)
# This avoids division by 0.
angle += mask_0_float * 1.
coef_1 += mask_else_float * angle / gs.sin(angle)
coef_2 += mask_else_float * angle / gs.tan(angle)
log = (gs.einsum('ni,nj->nj', coef_1, point) -
gs.einsum('ni,nj->nj', coef_2, base_point))
mask_same_values = gs.isclose(point, base_point)
mask_else = gs.equal(mask_same_values, gs.array(False))
mask_else_float = gs.cast(mask_else, gs.float32)
mask_else_float = gs.to_ndarray(mask_else_float, to_ndim=1)
mask_else_float = gs.to_ndarray(mask_else_float, to_ndim=2)
mask_not_same_points = gs.sum(mask_else_float, axis=1)
mask_same_points = gs.isclose(mask_not_same_points, 0.)
mask_same_points = gs.cast(mask_same_points, gs.float32)
mask_same_points = gs.to_ndarray(mask_same_points, to_ndim=2, axis=1)
mask_same_points_float = gs.cast(mask_same_points, gs.float32)
log -= mask_same_points_float * log
return log
def test_assignment_with_booleans_single_index(self):
np_array = _np.array([[2., 5.]])
gs_array = gs.array([[2., 5.]])
np_mask = _np.array([True])
gs_mask = gs.array([True])
np_array[np_mask] = _np.zeros_like(np_array[np_mask])
np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask])
np_result = np_array
values_mask = gs.zeros_like(gs_array[gs_mask])
gs_result = gs.assignment(gs_array, values_mask, gs_mask)
gs_result = gs.assignment(gs_result,
4 * gs.ones_like(gs_array[~gs_mask]),
~gs_mask)
self.assertAllCloseToNp(gs_result, np_result)
np_array = _np.array([[2., 5.]])
gs_array = gs.array([[2., 5.]])
np_mask = _np.array([True])
gs_mask = gs.array([True])
np_array[np_mask] = _np.zeros_like(np_array[np_mask])
np_array[~np_mask] = 4 * np_array[~np_mask]
np_result = np_array
values_mask = gs.zeros_like(gs_array[gs_mask])
gs_result = gs.assignment(gs_array, values_mask, gs_mask)
gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask)
self.assertAllCloseToNp(gs_result, np_result)
np_array = _np.array([[2., 5.]])
gs_array = gs.array([[2., 5.]])
np_mask = _np.array([False])
gs_mask = gs.array([False])
np_array[np_mask] = _np.zeros_like(np_array[np_mask])
np_array[~np_mask] = 4 * _np.ones_like(np_array[~np_mask])
np_result = np_array
values_mask = gs.zeros_like(gs_array[gs_mask])
gs_result = gs.assignment(gs_array, values_mask, gs_mask)
gs_result = gs.assignment(gs_result,
4 * gs.ones_like(gs_array[~gs_mask]),
~gs_mask)
self.assertAllCloseToNp(gs_result, np_result)
np_array = _np.array([[2., 5.]])
gs_array = gs.array([[2., 5.]])
np_mask = _np.array([False])
gs_mask = gs.array([False])
np_array[np_mask] = _np.zeros_like(np_array[np_mask])
np_array[~np_mask] = 4 * np_array[~np_mask]
np_result = np_array
values_mask = gs.zeros_like(gs_array[gs_mask])
gs_result = gs.assignment(gs_array, values_mask, gs_mask)
gs_result = gs.assignment(gs_result, 4 * gs_array[~gs_mask], ~gs_mask)
self.assertAllCloseToNp(gs_result, np_result)
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 log(self, point, base_point):
"""
Riemannian logarithm of a point wrt a base point.
"""
point = gs.to_ndarray(point, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
angle = self.dist(base_point, point)
angle = gs.to_ndarray(angle, to_ndim=1)
angle = gs.to_ndarray(angle, to_ndim=2)
mask_0 = gs.isclose(angle, 0)
mask_else = ~mask_0
coef_1 = gs.zeros_like(angle)
coef_2 = gs.zeros_like(angle)
coef_1[mask_0] = (1. + INV_SINH_TAYLOR_COEFFS[1] * angle[mask_0]**2 +
INV_SINH_TAYLOR_COEFFS[3] * angle[mask_0]**4 +
INV_SINH_TAYLOR_COEFFS[5] * angle[mask_0]**6 +
INV_SINH_TAYLOR_COEFFS[7] * angle[mask_0]**8)
coef_2[mask_0] = (1. + INV_TANH_TAYLOR_COEFFS[1] * angle[mask_0]**2 +
INV_TANH_TAYLOR_COEFFS[3] * angle[mask_0]**4 +
INV_TANH_TAYLOR_COEFFS[5] * angle[mask_0]**6 +
INV_TANH_TAYLOR_COEFFS[7] * angle[mask_0]**8)
coef_1[mask_else] = angle[mask_else] / gs.sinh(angle[mask_else])
coef_2[mask_else] = angle[mask_else] / gs.tanh(angle[mask_else])
log = (gs.einsum('ni,nj->nj', coef_1, point) -
gs.einsum('ni,nj->nj', coef_2, base_point))
return log
def test_integrability_tensor_derivative_reverses_hor_ver(self):
r"""Integrability tensor derivatives exchanges hor & ver in pre-shape.
For :math:`X,Y,Z` horizontal and :math:`V,W` vertical, the
integrability tensor (and thus its derivative) reverses horizontal
and vertical subspaces: :math:`\nabla_X < A_Y Z, H > = 0` and
:math:`nabla_X < A_Y V, W > = 0`.
"""
scal = self.space.ambient_metric.inner_product
nabla_x_a_y_z, a_y_z = self.space.integrability_tensor_derivative(
self.hor_x,
self.hor_y,
self.nabla_x_y,
self.hor_z,
self.nabla_x_z,
self.base_point,
)
result = scal(nabla_x_a_y_z, self.hor_h) + scal(a_y_z, self.nabla_x_h)
self.assertAllClose(result, gs.zeros_like(result))
nabla_x_a_y_v, a_y_v = self.space.integrability_tensor_derivative(
self.hor_x,
self.hor_y,
self.nabla_x_y,
self.ver_v,
self.nabla_x_v,
self.base_point,
)
result = scal(nabla_x_a_y_v, self.ver_w) + scal(a_y_v, self.nabla_x_w)
self.assertAllClose(result, gs.zeros_like(result))
def regularize_tangent_vec_at_identity(
self, tangent_vec, metric=None):
"""Regularize a tangent vector at the identify.
In 3D, regularize a tangent_vector by getting its norm at the identity,
determined by the metric, to be less than pi.
Parameters
----------
tangent_vec : array-like, shape=[..., 3]]
metric : RiemannianMetric, optional
default: self.left_canonical_metric
Returns
-------
regularized_vec : array-like, shape=[..., 3]]
"""
if metric is None:
metric = self.left_canonical_metric
tangent_vec_metric_norm = metric.norm(tangent_vec)
tangent_vec_canonical_norm = gs.linalg.norm(tangent_vec, axis=-1)
mask_norm_0 = gs.isclose(tangent_vec_metric_norm, 0.)
mask_canonical_norm_0 = gs.isclose(tangent_vec_canonical_norm, 0.)
mask_0 = mask_norm_0 | mask_canonical_norm_0
mask_else = ~mask_0
# This avoids division by 0.
mask_0_float = gs.cast(mask_0, gs.float32) + self.epsilon
mask_else_float = gs.cast(mask_else, gs.float32) + self.epsilon
regularized_vec = gs.zeros_like(tangent_vec)
regularized_vec += gs.einsum(
'...,...i->...i', mask_0_float, tangent_vec)
tangent_vec_canonical_norm += mask_0_float
coef = gs.zeros_like(tangent_vec_metric_norm)
coef += mask_else_float * (
tangent_vec_metric_norm
/ tangent_vec_canonical_norm)
coef_tangent_vec = gs.einsum(
'...,...i->...i', coef, tangent_vec)
regularized_vec += gs.einsum(
'...,...i->...i',
mask_else_float,
self.regularize(coef_tangent_vec))
coef += mask_0_float
regularized_vec = gs.einsum(
'...,...i->...i', 1. / coef, regularized_vec)
regularized_vec = gs.einsum(
'...,...i->...i', mask_else_float, regularized_vec)
return regularized_vec
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.
"""
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_0 = gs.equal(angle, 0.)
mask_close_0 = gs.isclose(angle, 0.) & ~mask_0
mask_else = ~mask_0 & ~mask_close_0
mask_0_float = gs.cast(mask_0, gs.float32)
mask_close_0_float = gs.cast(mask_close_0, gs.float32)
mask_else_float = gs.cast(mask_else, gs.float32)
coef_1 = gs.zeros_like(angle)
coef_2 = gs.zeros_like(angle)
coef_1 += mask_0_float * 1. / 2. * gs.ones_like(angle)
coef_2 += mask_0_float * 1. / 6. * gs.ones_like(angle)
coef_1 += mask_close_0_float * (
TAYLOR_COEFFS_1_AT_0[0]
+ TAYLOR_COEFFS_1_AT_0[2] * angle ** 2
+ TAYLOR_COEFFS_1_AT_0[4] * angle ** 4
+ TAYLOR_COEFFS_1_AT_0[6] * angle ** 6)
coef_2 += mask_close_0_float * (
TAYLOR_COEFFS_2_AT_0[0]
+ TAYLOR_COEFFS_2_AT_0[2] * angle ** 2
+ TAYLOR_COEFFS_2_AT_0[4] * angle ** 4
+ TAYLOR_COEFFS_2_AT_0[6] * angle ** 6)
angle += mask_0_float * 1.
coef_1 += mask_else_float * ((1. - gs.cos(angle)) / angle ** 2)
coef_2 += mask_else_float * ((angle - gs.sin(angle)) / angle ** 3)
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 exp(self, tangent_vec, base_point):
"""
Riemannian exponential of a tangent vector wrt to a base point.
Parameters
----------
tangent_vec : array-like, shape=[n_samples, dimension + 1]
or shape=[1, dimension + 1]
base_point : array-like, shape=[n_samples, dimension + 1]
or shape=[1, dimension + 1]
Returns
-------
exp : array-like, shape=[n_samples, dimension + 1]
or shape=[1, dimension + 1]
"""
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
sq_norm_tangent_vec = self.embedding_metric.squared_norm(
tangent_vec)
norm_tangent_vec = gs.sqrt(sq_norm_tangent_vec)
mask_0 = gs.isclose(sq_norm_tangent_vec, 0.)
mask_0 = gs.to_ndarray(mask_0, to_ndim=1)
mask_else = ~mask_0
mask_else = gs.to_ndarray(mask_else, to_ndim=1)
mask_0_float = gs.cast(mask_0, gs.float32)
mask_else_float = gs.cast(mask_else, gs.float32)
coef_1 = gs.zeros_like(norm_tangent_vec)
coef_2 = gs.zeros_like(norm_tangent_vec)
coef_1 += mask_0_float * (
1. + COSH_TAYLOR_COEFFS[2] * norm_tangent_vec ** 2
+ COSH_TAYLOR_COEFFS[4] * norm_tangent_vec ** 4
+ COSH_TAYLOR_COEFFS[6] * norm_tangent_vec ** 6
+ COSH_TAYLOR_COEFFS[8] * norm_tangent_vec ** 8)
coef_2 += mask_0_float * (
1. + SINH_TAYLOR_COEFFS[3] * norm_tangent_vec ** 2
+ SINH_TAYLOR_COEFFS[5] * norm_tangent_vec ** 4
+ SINH_TAYLOR_COEFFS[7] * norm_tangent_vec ** 6
+ SINH_TAYLOR_COEFFS[9] * norm_tangent_vec ** 8)
# This avoids dividing by 0.
norm_tangent_vec += mask_0_float * 1.0
coef_1 += mask_else_float * (gs.cosh(norm_tangent_vec))
coef_2 += mask_else_float * (
(gs.sinh(norm_tangent_vec) / (norm_tangent_vec)))
exp = (gs.einsum('ni,nj->nj', coef_1, base_point)
+ gs.einsum('ni,nj->nj', coef_2, tangent_vec))
hyperbolic_space = HyperbolicSpace(dimension=self.dimension)
exp = hyperbolic_space.regularize(exp)
return exp
def exp(self, tangent_vec, base_point):
"""Compute the Riemannian exponential of a tangent vector.
Parameters
----------
tangent_vec : array-like, shape=[..., dim + 1]
Tangent vector at a base point.
base_point : array-like, shape=[..., dim + 1]
Point in hyperbolic space.
Returns
-------
exp : array-like, shape=[..., dim + 1]
Point in hyperbolic space equal to the Riemannian exponential
of tangent_vec at the base point.
"""
sq_norm_tangent_vec = self.embedding_metric.squared_norm(
tangent_vec)
sq_norm_tangent_vec = gs.clip(sq_norm_tangent_vec, 0, math.inf)
norm_tangent_vec = gs.sqrt(sq_norm_tangent_vec)
mask_0 = gs.isclose(sq_norm_tangent_vec, 0.)
mask_0 = gs.to_ndarray(mask_0, to_ndim=1)
mask_else = ~mask_0
mask_else = gs.to_ndarray(mask_else, to_ndim=1)
mask_0_float = gs.cast(mask_0, gs.float32)
mask_else_float = gs.cast(mask_else, gs.float32)
coef_1 = gs.zeros_like(norm_tangent_vec)
coef_2 = gs.zeros_like(norm_tangent_vec)
coef_1 += mask_0_float * (
1. + COSH_TAYLOR_COEFFS[2] * norm_tangent_vec ** 2
+ COSH_TAYLOR_COEFFS[4] * norm_tangent_vec ** 4
+ COSH_TAYLOR_COEFFS[6] * norm_tangent_vec ** 6
+ COSH_TAYLOR_COEFFS[8] * norm_tangent_vec ** 8)
coef_2 += mask_0_float * (
1. + SINH_TAYLOR_COEFFS[3] * norm_tangent_vec ** 2
+ SINH_TAYLOR_COEFFS[5] * norm_tangent_vec ** 4
+ SINH_TAYLOR_COEFFS[7] * norm_tangent_vec ** 6
+ SINH_TAYLOR_COEFFS[9] * norm_tangent_vec ** 8)
# This avoids dividing by 0.
norm_tangent_vec += mask_0_float * 1.0
coef_1 += mask_else_float * (gs.cosh(norm_tangent_vec))
coef_2 += mask_else_float * (
(gs.sinh(norm_tangent_vec) / (norm_tangent_vec)))
exp = (
gs.einsum('...,...j->...j', coef_1, base_point)
+ gs.einsum('...,...j->...j', coef_2, tangent_vec))
hyperbolic_space = Hyperboloid(dim=self.dim)
exp = hyperbolic_space.regularize(exp)
return exp
def log(self, point, base_point):
"""Compute Riemannian logarithm of a point wrt a base point.
If point_type = 'poincare' then base_point belongs
to the Poincare ball and point is a vector in the Euclidean
space of the same dimension as the ball.
Parameters
----------
point : array-like, shape=[..., dim + 1]
Point in hyperbolic space.
base_point : array-like, shape=[..., dim + 1]
Point in hyperbolic space.
Returns
-------
log : array-like, shape=[..., dim + 1]
Tangent vector at the base point equal to the Riemannian logarithm
of point at the base point.
"""
angle = self.dist(base_point, point) / self.scale
angle = gs.to_ndarray(angle, to_ndim=1)
mask_0 = gs.isclose(angle, 0.)
mask_else = ~mask_0
mask_0_float = gs.cast(mask_0, gs.float32)
mask_else_float = gs.cast(mask_else, gs.float32)
coef_1 = gs.zeros_like(angle)
coef_2 = gs.zeros_like(angle)
coef_1 += mask_0_float * (
1. + INV_SINH_TAYLOR_COEFFS[1] * angle ** 2
+ INV_SINH_TAYLOR_COEFFS[3] * angle ** 4
+ INV_SINH_TAYLOR_COEFFS[5] * angle ** 6
+ INV_SINH_TAYLOR_COEFFS[7] * angle ** 8)
coef_2 += mask_0_float * (
1. + INV_TANH_TAYLOR_COEFFS[1] * angle ** 2
+ INV_TANH_TAYLOR_COEFFS[3] * angle ** 4
+ INV_TANH_TAYLOR_COEFFS[5] * angle ** 6
+ INV_TANH_TAYLOR_COEFFS[7] * angle ** 8)
# This avoids dividing by 0.
angle += mask_0_float * 1.
coef_1 += mask_else_float * (angle / gs.sinh(angle))
coef_2 += mask_else_float * (angle / gs.tanh(angle))
log_term_1 = gs.einsum('...,...j->...j', coef_1, point)
log_term_2 = - gs.einsum('...,...j->...j', coef_2, base_point)
log = log_term_1 + log_term_2
return log
def log(self, point, base_point):
"""
Riemannian logarithm of a point wrt a base point.
Parameters
----------
point : array-like, shape=[n_samples, dimension + 1]
or shape=[1, dimension + 1]
base_point : array-like, shape=[n_samples, dimension + 1]
or shape=[1, dimension + 1]
Returns
-------
log : array-like, shape=[n_samples, dimension + 1]
or shape=[1, dimension + 1]
"""
point = gs.to_ndarray(point, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
angle = self.dist(base_point, point)
angle = gs.to_ndarray(angle, to_ndim=1)
angle = gs.to_ndarray(angle, to_ndim=2)
mask_0 = gs.isclose(angle, 0.)
mask_else = ~mask_0
mask_0_float = gs.cast(mask_0, gs.float32)
mask_else_float = gs.cast(mask_else, gs.float32)
coef_1 = gs.zeros_like(angle)
coef_2 = gs.zeros_like(angle)
coef_1 += mask_0_float * (
1. + INV_SINH_TAYLOR_COEFFS[1] * angle ** 2
+ INV_SINH_TAYLOR_COEFFS[3] * angle ** 4
+ INV_SINH_TAYLOR_COEFFS[5] * angle ** 6
+ INV_SINH_TAYLOR_COEFFS[7] * angle ** 8)
coef_2 += mask_0_float * (
1. + INV_TANH_TAYLOR_COEFFS[1] * angle ** 2
+ INV_TANH_TAYLOR_COEFFS[3] * angle ** 4
+ INV_TANH_TAYLOR_COEFFS[5] * angle ** 6
+ INV_TANH_TAYLOR_COEFFS[7] * angle ** 8)
# This avoids dividing by 0.
angle += mask_0_float * 1.
coef_1 += mask_else_float * (angle / gs.sinh(angle))
coef_2 += mask_else_float * (angle / gs.tanh(angle))
log = (gs.einsum('ni,nj->nj', coef_1, point)
- gs.einsum('ni,nj->nj', coef_2, base_point))
return log
def test_assignment(self):
gs_array_1 = gs.ones(3)
self.assertRaises(ValueError, gs.assignment, gs_array_1, [.1, 2., 1.],
[0, 1])
np_array_1 = _np.ones(3)
gs_array_1 = gs.ones_like(gs.array(np_array_1))
np_array_1[2] = 1.5
gs_result = gs.assignment(gs_array_1, 1.5, 2)
self.assertAllCloseToNp(gs_result, np_array_1)
np_array_1_list = _np.ones(3)
gs_array_1_list = gs.ones_like(gs.array(np_array_1_list))
indices = [1, 2]
np_array_1_list[indices] = 1.5
gs_result = gs.assignment(gs_array_1_list, 1.5, indices)
self.assertAllCloseToNp(gs_result, np_array_1_list)
np_array_2 = _np.zeros((3, 2))
gs_array_2 = gs.zeros_like(gs.array(np_array_2))
np_array_2[0, :] = 1
gs_result = gs.assignment(gs_array_2, 1, 0, axis=1)
self.assertAllCloseToNp(gs_result, np_array_2)
np_array_3 = _np.zeros((3, 3))
gs_array_3 = gs.zeros_like(gs.array(np_array_3))
np_array_3[0, 1] = 1
gs_result = gs.assignment(gs_array_3, 1, (0, 1))
self.assertAllCloseToNp(gs_result, np_array_3)
np_array_4 = _np.zeros((3, 3, 2))
gs_array_4 = gs.zeros_like(gs.array(np_array_4))
np_array_4[0, :, 1] = 1
gs_result = gs.assignment(gs_array_4, 1, (0, 1), axis=1)
self.assertAllCloseToNp(gs_result, np_array_4)
gs_array_4_arr = gs.zeros_like(gs.array(np_array_4))
gs_result = gs.assignment(gs_array_4_arr, 1, gs.array((0, 1)), axis=1)
self.assertAllCloseToNp(gs_result, np_array_4)
np_array_4_list = _np.zeros((3, 3, 2))
gs_array_4_list = gs.zeros_like(gs.array(np_array_4_list))
np_array_4_list[(0, 1), :, (1, 1)] = 1
gs_result = gs.assignment(gs_array_4_list, 1, [(0, 1), (1, 1)], axis=1)
self.assertAllCloseToNp(gs_result, np_array_4_list)
def matrix_from_rotation_vector(self, rot_vec):
"""Convert rotation vector to rotation matrix.
Parameters
----------
rot_vec: array-like, shape=[..., 3]
Returns
-------
rot_mat: array-like, shape=[..., 3]
"""
rot_vec = self.regularize(rot_vec)
angle = gs.linalg.norm(rot_vec, axis=1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
skew_rot_vec = self.skew_matrix_from_vector(rot_vec)
coef_1 = gs.zeros_like(angle)
coef_2 = gs.zeros_like(angle)
# This avoids dividing by 0.
mask_0 = gs.isclose(angle, 0.)
mask_0_float = gs.cast(mask_0, gs.float32) + self.epsilon
coef_1 += mask_0_float * (1. - (angle ** 2) / 6.)
coef_2 += mask_0_float * (1. / 2. - angle ** 2)
# This avoids dividing by 0.
mask_else = ~mask_0
mask_else_float = gs.cast(mask_else, gs.float32) + self.epsilon
angle += mask_0_float
coef_1 += mask_else_float * (gs.sin(angle) / angle)
coef_2 += mask_else_float * (
(1. - gs.cos(angle)) / (angle ** 2))
coef_1 = gs.squeeze(coef_1, axis=1)
coef_2 = gs.squeeze(coef_2, axis=1)
term_1 = (gs.eye(self.dim)
+ gs.einsum('n,njk->njk', coef_1, skew_rot_vec))
squared_skew_rot_vec = gs.einsum(
'nij,njk->nik', skew_rot_vec, skew_rot_vec)
term_2 = gs.einsum('n,njk->njk', coef_2, squared_skew_rot_vec)
return term_1 + term_2
def log(self, point, base_point):
"""
Compute the Riemannian logarithm at point base_point,
of point wrt the metric obtained by
embedding of the n-dimensional sphere
in the (n+1)-dimensional euclidean space.
This gives a tangent vector at point base_point.
:param base_point: point on the n-dimensional sphere
:param point: point on the n-dimensional sphere
:return log: tangent vector at base_point
"""
point = gs.to_ndarray(point, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
norm_base_point = self.embedding_metric.norm(base_point)
norm_point = self.embedding_metric.norm(point)
inner_prod = self.embedding_metric.inner_product(base_point, point)
cos_angle = inner_prod / (norm_base_point * norm_point)
cos_angle = gs.clip(cos_angle, -1.0, 1.0)
angle = gs.arccos(cos_angle)
mask_0 = gs.isclose(angle, 0.0)
mask_else = gs.equal(mask_0, False)
coef_1 = gs.zeros_like(angle)
coef_2 = gs.zeros_like(angle)
coef_1[mask_0] = (
1. + INV_SIN_TAYLOR_COEFFS[1] * angle[mask_0] ** 2
+ INV_SIN_TAYLOR_COEFFS[3] * angle[mask_0] ** 4
+ INV_SIN_TAYLOR_COEFFS[5] * angle[mask_0] ** 6
+ INV_SIN_TAYLOR_COEFFS[7] * angle[mask_0] ** 8)
coef_2[mask_0] = (
1. + INV_TAN_TAYLOR_COEFFS[1] * angle[mask_0] ** 2
+ INV_TAN_TAYLOR_COEFFS[3] * angle[mask_0] ** 4
+ INV_TAN_TAYLOR_COEFFS[5] * angle[mask_0] ** 6
+ INV_TAN_TAYLOR_COEFFS[7] * angle[mask_0] ** 8)
coef_1[mask_else] = angle[mask_else] / gs.sin(angle[mask_else])
coef_2[mask_else] = angle[mask_else] / gs.tan(angle[mask_else])
log = (gs.einsum('ni,nj->nj', coef_1, point)
- gs.einsum('ni,nj->nj', coef_2, base_point))
return log
def test_fit_transform_hyperbolic(self):
point = gs.array([2., 1., 1., 1.])
points = gs.array([point, point])
transformer = ToTangentSpace(geometry=self.hyperbolic.metric)
result = transformer.fit_transform(X=points)
expected = gs.zeros_like(points)
self.assertAllClose(expected, result)
def curvature_test_data(self):
group = self.matrix_so3
metric = InvariantMetric(group)
x, y, z = metric.normal_basis(group.lie_algebra.basis)
smoke_data = [
dict(
group=group,
tangent_vec_a=x,
tangent_vec_b=y,
tangent_vec_c=x,
expected=1.0 / 8 * y,
),
dict(
group=group,
tangent_vec_a=gs.stack([x, x]),
tangent_vec_b=gs.stack([y] * 2),
tangent_vec_c=gs.stack([x, x]),
expected=gs.array([1.0 / 8 * y] * 2),
),
dict(
group=group,
tangent_vec_a=y,
tangent_vec_b=y,
tangent_vec_c=z,
expected=gs.zeros_like(z),
),
]
return self.generate_tests(smoke_data)
def rotation_vector_from_quaternion(self, quaternion):
"""
Convert a unit quaternion into a rotation vector.
"""
assert self.n == 3, ('The quaternion representation does not exist'
' for rotations in %d dimensions.' % self.n)
quaternion = gs.to_ndarray(quaternion, to_ndim=2)
n_quaternions, _ = quaternion.shape
cos_half_angle = quaternion[:, 0]
cos_half_angle = gs.clip(cos_half_angle, -1, 1)
half_angle = gs.arccos(cos_half_angle)
half_angle = gs.to_ndarray(half_angle, to_ndim=2, axis=1)
assert half_angle.shape == (n_quaternions, 1)
rot_vec = gs.zeros_like(quaternion[:, 1:])
mask_0 = gs.isclose(half_angle, 0)
mask_0 = gs.squeeze(mask_0, axis=1)
mask_not_0 = ~mask_0
rotation_axis = (quaternion[mask_not_0, 1:] /
gs.sin(half_angle[mask_not_0]))
rot_vec[mask_not_0] = (2 * half_angle[mask_not_0] * rotation_axis)
rot_vec = self.regularize(rot_vec)
return rot_vec
def quaternion_from_rotation_vector(self, rot_vec):
"""
Convert a rotation vector into a unit quaternion.
"""
assert self.n == 3, ('The quaternion representation does not exist'
' for rotations in %d dimensions.' % self.n)
rot_vec = self.regularize(rot_vec)
n_rot_vecs, _ = rot_vec.shape
angle = gs.linalg.norm(rot_vec, axis=1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
rotation_axis = gs.zeros_like(rot_vec)
mask_0 = gs.isclose(angle, 0)
mask_0 = gs.squeeze(mask_0, axis=1)
mask_not_0 = ~mask_0
rotation_axis[mask_not_0] = rot_vec[mask_not_0] / angle[mask_not_0]
n_quaternions, _ = rot_vec.shape
quaternion = gs.zeros((n_quaternions, 4))
quaternion[:, :1] = gs.cos(angle / 2)
quaternion[:, 1:] = gs.sin(angle / 2) * rotation_axis[:]
return quaternion
def closest_rotation_matrix(mat):
"""
Compute the closest rotation matrix of a given matrix mat,
in terms of the Frobenius norm.
"""
mat = gs.to_ndarray(mat, to_ndim=3)
n_mats, mat_dim_1, mat_dim_2 = mat.shape
assert mat_dim_1 == mat_dim_2
if mat_dim_1 == 3:
mat_unitary_u, diag_s, mat_unitary_v = gs.linalg.svd(mat)
rot_mat = gs.matmul(mat_unitary_u, mat_unitary_v)
mask = gs.where(gs.linalg.det(rot_mat) < 0)
new_mat_diag_s = gs.tile(gs.diag([1, 1, -1]), len(mask))
rot_mat[mask] = gs.matmul(
gs.matmul(mat_unitary_u[mask], new_mat_diag_s),
mat_unitary_v[mask])
else:
aux_mat = gs.matmul(gs.transpose(mat, axes=(0, 2, 1)), mat)
inv_sqrt_mat = gs.zeros_like(mat)
for i in range(n_mats):
sym_mat = aux_mat[i]
assert spd_matrices_space.is_symmetric(sym_mat)
inv_sqrt_mat[i] = gs.linalg.inv(spd_matrices_space.sqrtm(sym_mat))
rot_mat = gs.matmul(mat, inv_sqrt_mat)
assert rot_mat.ndim == 3
return rot_mat
def regularize(self, point):
"""
In 3D, regularize the norm of the rotation vector,
to be between 0 and pi, following the axis-angle
representation's convention.
If the angle angle is between pi and 2pi,
the function computes its complementary in 2pi and
inverts the direction of the rotation axis.
"""
point = gs.to_ndarray(point, to_ndim=2)
assert self.belongs(point)
n_points, vec_dim = point.shape
if vec_dim == 3:
angle = gs.linalg.norm(point, axis=1)
regularized_point = point.astype('float64')
mask_not_0 = angle != 0
k = gs.floor(angle / (2 * gs.pi) + .5)
norms_ratio = gs.zeros_like(angle).astype('float64')
norms_ratio[mask_not_0] = (
1. - 2. * gs.pi * k[mask_not_0] / angle[mask_not_0])
norms_ratio[angle == 0] = 1
for i in range(n_points):
regularized_point[i, :] = norms_ratio[i] * point[i]
else:
# TODO(nina): regularization needed in nD?
regularized_point = point
assert regularized_point.ndim == 2
return regularized_point
def test_curvature(self):
group = self.matrix_so3
metric = InvariantMetric(group=group)
x, y, z = metric.normal_basis(group.lie_algebra.basis)
result = metric.curvature_at_identity(x, y, x)
expected = 1.0 / 8 * y
self.assertAllClose(result, expected)
tan_a = gs.stack([x, x])
tan_b = gs.stack([y] * 2)
result = metric.curvature(tan_a, tan_b, tan_a)
self.assertAllClose(result, gs.array([expected] * 2))
point = group.random_uniform()
translation_map = group.tangent_translation_map(point)
tan_a = translation_map(x)
tan_b = translation_map(y)
result = metric.curvature(tan_a, tan_b, tan_a, point)
expected = translation_map(expected)
self.assertAllClose(result, expected)
result = metric.curvature(y, y, z)
expected = gs.zeros_like(z)
self.assertAllClose(result, expected)
def test_curvature_derivative(self):
group = self.matrix_so3
metric = InvariantMetric(group=group)
x, y, z = metric.normal_basis(group.lie_algebra.basis)
result = metric.curvature_derivative(x, y, z, x)
expected = gs.zeros_like(x)
self.assertAllClose(result, expected)
point = group.random_uniform()
translation_map = group.tangent_translation_map(point)
tan_a = translation_map(x)
tan_b = translation_map(y)
tan_c = translation_map(z)
result = metric.curvature_derivative(tan_a, tan_b, tan_c, tan_a, point)
expected = gs.zeros_like(x)
self.assertAllClose(result, expected)
def test_fit_riemannian_se2(self):
init = (self.y_se2[0], gs.zeros_like(self.y_se2[0]))
gr = GeodesicRegression(
self.se2,
metric=self.metric_se2,
center_X=False,
method="riemannian",
max_iter=50,
init_step_size=0.1,
verbose=True,
initialization=init,
)
gr.fit(self.X_se2, self.y_se2, compute_training_score=True)
intercept_hat, coef_hat = gr.intercept_, gr.coef_
training_score = gr.training_score_
self.assertAllClose(intercept_hat.shape, self.shape_se2)
self.assertAllClose(coef_hat.shape, self.shape_se2)
self.assertAllClose(training_score, 1.0, atol=1e-4)
self.assertAllClose(intercept_hat, self.intercept_se2_true, atol=1e-4)
tangent_vec_of_transport = self.se2.metric.log(
self.intercept_se2_true, base_point=intercept_hat)
transported_coef_hat = self.se2.metric.parallel_transport(
tangent_vec=coef_hat,
base_point=intercept_hat,
direction=tangent_vec_of_transport,
)
self.assertAllClose(transported_coef_hat, self.coef_se2_true, atol=0.6)
def inverse(self, point):
"""
Compute the group inverse in SE(n).
Formula:
(R, t)^{-1} = (R^{-1}, R^{-1}.(-t))
"""
rotations = self.rotations
dim_rotations = rotations.dimension
point = self.regularize(point)
n_points, _ = point.shape
rot_vec = point[:, :dim_rotations]
translation = point[:, dim_rotations:]
inverse_point = gs.zeros_like(point)
inverse_rotation = -rot_vec
inv_rot_mat = rotations.matrix_from_rotation_vector(inverse_rotation)
inverse_translation = gs.zeros((n_points, self.n))
for i in range(n_points):
inverse_translation[i] = gs.dot(-translation[i],
gs.transpose(inv_rot_mat[i]))
inverse_point[:, :dim_rotations] = inverse_rotation
inverse_point[:, dim_rotations:] = inverse_translation
inverse_point = self.regularize(inverse_point)
return inverse_point
def regularize(self, point, point_type=None):
"""
Regularize a point to the canonical representation
chosen for SE(n).
"""
if point_type is None:
point_type = self.default_point_type
if point_type == 'vector':
point = gs.to_ndarray(point, to_ndim=2)
assert self.belongs(point, point_type=point_type)
rotations = self.rotations
dim_rotations = rotations.dimension
regularized_point = gs.zeros_like(point)
rot_vec = point[:, :dim_rotations]
regularized_point[:, :dim_rotations] = rotations.regularize(
rot_vec, point_type=point_type)
regularized_point[:, dim_rotations:] = point[:, dim_rotations:]
elif point_type == 'matrix':
point = gs.to_ndarray(point, to_ndim=3)
# TODO(nina): regularization for matrices?
regularized_point = gs.copy(point)
return regularized_point
def compose(self, point_a, point_b):
"""Compute the group product of elements `point_a` and `point_b`.
Parameters
----------
point_a : array-like, shape=[..., 3]
Left factor in the product.
point_b : array-like, shape=[..., 3]
Right factor in the product.
Returns
-------
point_ab : array-like, shape=[..., 3]
Product of point_a and point_b along the first dimension.
"""
point_ab = point_a + point_b
point_ab_additional_term = gs.array(
1 / 2 * (point_a[..., 0] * point_b[..., 1] -
point_a[..., 1] * point_b[..., 0]))
point_ab = point_ab + gs.concatenate(
[
gs.zeros_like(point_ab[..., :2]),
point_ab_additional_term[..., None]
],
axis=-1,
)
return point_ab
def regularize_tangent_vec(self, tangent_vec, base_point, metric=None):
if metric is None:
metric = self.left_canonical_metric
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
rotations = self.rotations
dim_rotations = rotations.dimension
rot_tangent_vec = tangent_vec[:, :dim_rotations]
rot_base_point = base_point[:, :dim_rotations]
metric_mat = metric.inner_product_mat_at_identity
rot_metric_mat = metric_mat[:, :dim_rotations, :dim_rotations]
rot_metric = InvariantMetric(
group=rotations,
inner_product_mat_at_identity=rot_metric_mat,
left_or_right=metric.left_or_right)
regularized_vec = gs.zeros_like(tangent_vec)
regularized_vec[:, :dim_rotations] = rotations.regularize_tangent_vec(
tangent_vec=rot_tangent_vec,
base_point=rot_base_point,
metric=rot_metric)
regularized_vec[:, dim_rotations:] = tangent_vec[:, dim_rotations:]
return regularized_vec
def _log_translation_transform(self, rot_vec):
rot_vec = gs.to_ndarray(rot_vec, to_ndim=2, axis=1)
exp_transform = self._exp_translation_transform(rot_vec)
if rot_vec.dtype == gs.float32:
mask_close_0 = gs.isclose(rot_vec, 0., atol=1e-6)
else:
mask_close_0 = gs.isclose(rot_vec, 0.)
mask_else = ~mask_close_0
mask_close_0_float = gs.cast(mask_close_0, gs.float32)
mask_else_float = gs.cast(mask_else, gs.float32)
inv_determinant = gs.zeros_like(rot_vec)
inv_determinant += 0.5 * mask_close_0_float / (
TAYLOR_COEFFS_1_AT_0[0]
+ TAYLOR_COEFFS_1_AT_0[1] * rot_vec ** 2
+ TAYLOR_COEFFS_1_AT_0[2] * rot_vec ** 6)
rot_vec = rot_vec + 1e-3 * mask_close_0_float
inv_determinant += mask_else_float * (
rot_vec ** 2 / (2 * (1 - gs.cos(rot_vec))))
transform = gs.einsum(
'il, ijk -> ijk', inv_determinant,
gs.transpose(exp_transform, axes=[0, 2, 1]))
return transform
