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 square_root_velocity(self, curve):
"""Compute the square root velocity representation of a curve.
The velocity is computed using the log map. The case of several curves
is handled through vectorization. In that case, an index selection
procedure allows to get rid of the log between the end point of
curve[k, :, :] and the starting point of curve[k + 1, :, :].
Parameters
----------
curve :
Returns
-------
srv :
"""
curve = gs.to_ndarray(curve, to_ndim=3)
n_curves, n_sampling_points, n_coords = curve.shape
srv_shape = (n_curves, n_sampling_points - 1, n_coords)
curve = gs.reshape(curve, (n_curves * n_sampling_points, n_coords))
coef = gs.cast(gs.array(n_sampling_points - 1), gs.float32)
velocity = coef * self.ambient_metric.log(point=curve[1:, :],
base_point=curve[:-1, :])
velocity_norm = self.ambient_metric.norm(velocity, curve[:-1, :])
srv = velocity / gs.sqrt(velocity_norm)
index = gs.arange(n_curves * n_sampling_points - 1)
mask = ~gs.equal((index + 1) % n_sampling_points, 0)
index_select = gs.gather(index, gs.squeeze(gs.where(mask)))
srv = gs.reshape(gs.gather(srv, index_select), srv_shape)
return srv
def assertAllEqual(self, a, b):
if tf_backend():
return tf.test.TestCase().assertAllEqual(a, b)
elif np_backend() or autograd_backend():
np.testing.assert_array_equal(a, b)
else:
self.assertTrue(gs.equal(a, b))
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 _exponential_matrix(self, rot_vec):
"""Compute exponential of rotation matrix represented by rot_vec.
Parameters
----------
rot_vec : array-like, shape=[..., 3]
Returns
-------
exponential_mat : Matrix exponential of rot_vec
"""
# TODO (nguigs): find usecase for this method
rot_vec = self.rotations.regularize(rot_vec)
n_rot_vecs = 1 if rot_vec.ndim == 1 else len(rot_vec)
angle = gs.linalg.norm(rot_vec, axis=-1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
skew_rot_vec = self.rotations.skew_matrix_from_vector(rot_vec)
coef_1 = gs.empty_like(angle)
coef_2 = gs.empty_like(coef_1)
mask_0 = gs.equal(angle, 0)
mask_0 = gs.squeeze(mask_0, axis=1)
mask_close_to_0 = gs.isclose(angle, 0)
mask_close_to_0 = gs.squeeze(mask_close_to_0, axis=1)
mask_else = ~mask_0 & ~mask_close_to_0
coef_1[mask_close_to_0] = (1. / 2.
- angle[mask_close_to_0] ** 2 / 24.)
coef_2[mask_close_to_0] = (1. / 6.
- angle[mask_close_to_0] ** 3 / 120.)
# TODO (nina): Check if the discontinuity at 0 is expected.
coef_1[mask_0] = 0
coef_2[mask_0] = 0
coef_1[mask_else] = (angle[mask_else] ** (-2)
* (1. - gs.cos(angle[mask_else])))
coef_2[mask_else] = (angle[mask_else] ** (-2)
* (1. - (gs.sin(angle[mask_else])
/ angle[mask_else])))
term_1 = gs.zeros((n_rot_vecs, self.n, self.n))
term_2 = gs.zeros_like(term_1)
for i in range(n_rot_vecs):
term_1[i] = gs.eye(self.n) + skew_rot_vec[i] * coef_1[i]
term_2[i] = gs.matmul(skew_rot_vec[i], skew_rot_vec[i]) * coef_2[i]
exponential_mat = term_1 + term_2
return exponential_mat
def belongs(self, mat):
"""
Check if mat belongs to GL(n).
"""
mat = gs.to_ndarray(mat, to_ndim=3)
n_mats, _, _ = mat.shape
mat_rank = gs.zeros((n_mats, 1))
mat_rank = gs.linalg.matrix_rank(mat)
mat_rank = gs.to_ndarray(mat_rank, to_ndim=1)
return gs.equal(mat_rank, self.n)
def exponential_matrix(self, rot_vec):
"""
Compute the exponential of the rotation matrix
represented by rot_vec.
:param rot_vec: 3D rotation vector
:returns exponential_mat: 3x3 matrix
"""
rot_vec = self.rotations.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)
skew_rot_vec = so_group.skew_matrix_from_vector(rot_vec)
coef_1 = gs.empty_like(angle)
coef_2 = gs.empty_like(coef_1)
mask_0 = gs.equal(angle, 0)
mask_0 = gs.squeeze(mask_0, axis=1)
mask_close_to_0 = gs.isclose(angle, 0)
mask_close_to_0 = gs.squeeze(mask_close_to_0, axis=1)
mask_else = ~mask_0 & ~mask_close_to_0
coef_1[mask_close_to_0] = (1. / 2.
- angle[mask_close_to_0] ** 2 / 24.)
coef_2[mask_close_to_0] = (1. / 6.
- angle[mask_close_to_0] ** 3 / 120.)
# TODO(nina): check if the discountinuity as 0 is expected.
coef_1[mask_0] = 0
coef_2[mask_0] = 0
coef_1[mask_else] = (angle[mask_else] ** (-2)
* (1. - gs.cos(angle[mask_else])))
coef_2[mask_else] = (angle[mask_else] ** (-2)
* (1. - (gs.sin(angle[mask_else])
/ angle[mask_else])))
term_1 = gs.zeros((n_rot_vecs, self.n, self.n))
term_2 = gs.zeros_like(term_1)
for i in range(n_rot_vecs):
term_1[i] = gs.eye(self.n) + skew_rot_vec[i] * coef_1[i]
term_2[i] = gs.matmul(skew_rot_vec[i], skew_rot_vec[i]) * coef_2[i]
exponential_mat = term_1 + term_2
assert exponential_mat.ndim == 3
return exponential_mat
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 dist(self, point_a, point_b):
"""
Geodesic distance between two points.
"""
if gs.all(gs.equal(point_a, point_b)):
return 0.
sq_norm_a = self.embedding_metric.squared_norm(point_a)
sq_norm_b = self.embedding_metric.squared_norm(point_b)
inner_prod = self.embedding_metric.inner_product(point_a, point_b)
cosh_angle = -inner_prod / gs.sqrt(sq_norm_a * sq_norm_b)
cosh_angle = gs.clip(cosh_angle, 1, None)
dist = gs.arccosh(cosh_angle)
return dist
def dist(self, point_a, point_b):
"""
Compute the distance induced on the hyperbolic
space, from its embedding in the Minkowski space.
"""
if gs.all(gs.equal(point_a, point_b)):
return 0.
sq_norm_a = self.embedding_metric.squared_norm(point_a)
sq_norm_b = self.embedding_metric.squared_norm(point_b)
inner_prod = self.embedding_metric.inner_product(point_a, point_b)
cosh_angle = -inner_prod / gs.sqrt(sq_norm_a * sq_norm_b)
cosh_angle = gs.clip(cosh_angle, 1, None)
dist = gs.arccosh(cosh_angle)
return dist
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.0, 1.0)
angle = gs.arccos(cos_angle)
mask_0 = gs.isclose(angle, 0.0)
mask_else = gs.equal(mask_0, gs.cast(gs.array(False), gs.int8))
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 log(self, point, base_point):
"""Compute the Riemannian logarithm of a point.
Parameters
----------
point : array-like, shape=[..., dim + 1]
Point on the hypersphere.
base_point : array-like, shape=[..., dim + 1]
Point on the hypersphere.
Returns
-------
log : array-like, shape=[..., dim + 1]
Tangent vector at the base point equal to the Riemannian logarithm
of point at the base point.
"""
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('...i,...j->...j', coef_1, point) -
gs.einsum('...i,...j->...j', 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 group_exp_from_identity(self, tangent_vec):
"""
Compute the group exponential of the tangent vector at the identity.
"""
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
rotations = self.rotations
dim_rotations = rotations.dimension
rot_vec = tangent_vec[:, :dim_rotations]
rot_vec = self.rotations.regularize(rot_vec)
translation = tangent_vec[:, dim_rotations:]
angle = gs.linalg.norm(rot_vec, axis=1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
mask_close_pi = gs.isclose(angle, gs.pi)
mask_close_pi = gs.squeeze(mask_close_pi, axis=1)
rot_vec[mask_close_pi] = rotations.regularize(
rot_vec[mask_close_pi])
skew_mat = so_group.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_0 = gs.squeeze(mask_0, axis=1)
mask_close_0 = gs.squeeze(mask_close_0, axis=1)
mask_else = ~mask_0 & ~mask_close_0
coef_1 = gs.zeros_like(angle)
coef_2 = gs.zeros_like(angle)
coef_1[mask_0] = 1. / 2.
coef_2[mask_0] = 1. / 6.
coef_1[mask_close_0] = (1. / 2. - angle[mask_close_0] ** 2 / 24.
+ angle[mask_close_0] ** 4 / 720.
- angle[mask_close_0] ** 6 / 40320.)
coef_2[mask_close_0] = (1. / 6. - angle[mask_close_0] ** 2 / 120.
+ angle[mask_close_0] ** 4 / 5040.
- angle[mask_close_0] ** 6 / 362880.)
coef_1[mask_else] = ((1. - gs.cos(angle[mask_else]))
/ angle[mask_else] ** 2)
coef_2[mask_else] = ((angle[mask_else] - gs.sin(angle[mask_else]))
/ angle[mask_else] ** 3)
n_tangent_vecs, _ = tangent_vec.shape
group_exp_translation = gs.zeros((n_tangent_vecs, self.n))
for i in range(n_tangent_vecs):
translation_i = translation[i]
term_1_i = coef_1[i] * gs.dot(translation_i,
gs.transpose(skew_mat[i]))
term_2_i = coef_2[i] * gs.dot(translation_i,
gs.transpose(sq_skew_mat[i]))
group_exp_translation[i] = translation_i + term_1_i + term_2_i
group_exp = gs.zeros_like(tangent_vec)
group_exp[:, :dim_rotations] = rot_vec
group_exp[:, dim_rotations:] = group_exp_translation
group_exp = self.regularize(group_exp)
return group_exp
def exp_from_identity(self, tangent_vec, point_type=None):
"""Compute group exponential of the tangent vector at the identity.
Parameters
----------
tangent_vec: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
point_type: str, {'vector', 'matrix'}, optional
default: self.default_point_type
Returns
-------
group_exp: array-like, shape=[n_samples, {dim, [n + 1, n + 1]}]
the group exponential of the tangent vectors calculated
at the identity
"""
if point_type == 'vector':
rotations = self.rotations
dim_rotations = rotations.dim
rot_vec = tangent_vec[:, :dim_rotations]
rot_vec = self.rotations.regularize(rot_vec, point_type=point_type)
translation = tangent_vec[:, dim_rotations:]
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)
angle += mask_0_float * gs.ones_like(angle)
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)
coef_1 += mask_else_float * ((1. - gs.cos(angle)) / angle**2)
coef_2 += mask_else_float * ((angle - gs.sin(angle)) / angle**3)
n_tangent_vecs, _ = tangent_vec.shape
exp_translation = gs.zeros((n_tangent_vecs, self.n))
for i in range(n_tangent_vecs):
translation_i = translation[i]
term_1_i = coef_1[i] * gs.dot(translation_i,
gs.transpose(skew_mat[i]))
term_2_i = coef_2[i] * gs.dot(translation_i,
gs.transpose(sq_skew_mat[i]))
mask_i_float = gs.get_mask_i_float(i, n_tangent_vecs)
exp_translation += gs.outer(
mask_i_float, translation_i + term_1_i + term_2_i)
group_exp = gs.concatenate([rot_vec, exp_translation], axis=1)
group_exp = self.regularize(group_exp, point_type=point_type)
return group_exp
if point_type == 'matrix':
return GeneralLinear.exp(tangent_vec)
raise ValueError('Invalid point_type, expected \'vector\' or '
'\'matrix\'.')
def group_exp_from_identity(self, tangent_vec, point_type=None):
"""
Compute the group exponential of the tangent vector at the identity.
"""
if point_type is None:
point_type = self.default_point_type
if point_type == 'vector':
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
rotations = self.rotations
dim_rotations = rotations.dimension
rot_vec = tangent_vec[:, :dim_rotations]
rot_vec = self.rotations.regularize(rot_vec, point_type=point_type)
translation = tangent_vec[:, dim_rotations:]
angle = gs.linalg.norm(rot_vec, axis=1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
mask_close_pi = gs.isclose(angle, gs.pi)
mask_close_pi = gs.squeeze(mask_close_pi, axis=1)
rot_vec[mask_close_pi] = rotations.regularize(
rot_vec[mask_close_pi],
point_type=point_type)
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_0 = gs.squeeze(mask_0, axis=1)
mask_close_0 = gs.squeeze(mask_close_0, axis=1)
mask_else = ~mask_0 & ~mask_close_0
coef_1 = gs.zeros_like(angle)
coef_2 = gs.zeros_like(angle)
coef_1[mask_0] = 1. / 2.
coef_2[mask_0] = 1. / 6.
coef_1[mask_close_0] = (
TAYLOR_COEFFS_1_AT_0[0]
+ TAYLOR_COEFFS_1_AT_0[2] * angle[mask_close_0] ** 2
+ TAYLOR_COEFFS_1_AT_0[4] * angle[mask_close_0] ** 4
+ TAYLOR_COEFFS_1_AT_0[6] * angle[mask_close_0] ** 6)
coef_2[mask_close_0] = (
TAYLOR_COEFFS_2_AT_0[0]
+ TAYLOR_COEFFS_2_AT_0[2] * angle[mask_close_0] ** 2
+ TAYLOR_COEFFS_2_AT_0[4] * angle[mask_close_0] ** 4
+ TAYLOR_COEFFS_2_AT_0[6] * angle[mask_close_0] ** 6)
coef_1[mask_else] = ((1. - gs.cos(angle[mask_else]))
/ angle[mask_else] ** 2)
coef_2[mask_else] = ((angle[mask_else] - gs.sin(angle[mask_else]))
/ angle[mask_else] ** 3)
n_tangent_vecs, _ = tangent_vec.shape
group_exp_translation = gs.zeros((n_tangent_vecs, self.n))
for i in range(n_tangent_vecs):
translation_i = translation[i]
term_1_i = coef_1[i] * gs.dot(translation_i,
gs.transpose(skew_mat[i]))
term_2_i = coef_2[i] * gs.dot(translation_i,
gs.transpose(sq_skew_mat[i]))
group_exp_translation[i] = translation_i + term_1_i + term_2_i
group_exp = gs.zeros_like(tangent_vec)
group_exp[:, :dim_rotations] = rot_vec
group_exp[:, dim_rotations:] = group_exp_translation
group_exp = self.regularize(group_exp, point_type=point_type)
return group_exp
elif point_type == 'matrix':
raise NotImplementedError()
def group_exp_from_identity(self, tangent_vec, point_type=None):
"""
Compute the group exponential of the tangent vector at the identity.
"""
if point_type is None:
point_type = self.default_point_type
if point_type == 'vector':
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
rotations = self.rotations
dim_rotations = rotations.dimension
rot_vec = tangent_vec[:, :dim_rotations]
rot_vec = self.rotations.regularize(rot_vec, point_type=point_type)
translation = tangent_vec[:, dim_rotations:]
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)
angle += mask_0_float * gs.ones_like(angle)
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)
coef_1 += mask_else_float * ((1. - gs.cos(angle)) / angle**2)
coef_2 += mask_else_float * ((angle - gs.sin(angle)) / angle**3)
n_tangent_vecs, _ = tangent_vec.shape
group_exp_translation = gs.zeros((n_tangent_vecs, self.n))
for i in range(n_tangent_vecs):
translation_i = translation[i]
term_1_i = coef_1[i] * gs.dot(translation_i,
gs.transpose(skew_mat[i]))
term_2_i = coef_2[i] * gs.dot(translation_i,
gs.transpose(sq_skew_mat[i]))
mask_i_float = gs.get_mask_i_float(i, n_tangent_vecs)
group_exp_translation += mask_i_float * (translation_i +
term_1_i + term_2_i)
group_exp = gs.concatenate([rot_vec, group_exp_translation],
axis=1)
group_exp = self.regularize(group_exp, point_type=point_type)
return group_exp
elif point_type == 'matrix':
raise NotImplementedError()
def get_mask_i_float(i, n):
range_n = gs.arange(n)
i_float = gs.cast(gs.array([i]), gs.int32)[0]
mask_i = gs.equal(range_n, i_float)
mask_i_float = gs.cast(mask_i, gs.float32)
return mask_i_float
def exp_from_identity(self, tangent_vec):
"""Compute group exponential of the tangent vector at the identity.
Parameters
----------
tangent_vec: array-like, shape=[..., 3]
Returns
-------
group_exp: array-like, shape=[..., 3]
The group exponential of the tangent vectors calculated
at the identity.
"""
rotations = self.rotations
dim_rotations = rotations.dim
rot_vec = tangent_vec[..., :dim_rotations]
rot_vec = self.rotations.regularize(rot_vec)
translation = tangent_vec[..., dim_rotations:]
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)
angle += mask_0_float * gs.ones_like(angle)
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)
coef_1 += mask_else_float * ((1. - gs.cos(angle)) / angle**2)
coef_2 += mask_else_float * ((angle - gs.sin(angle)) / angle**3)
n_tangent_vecs, _ = tangent_vec.shape
exp_translation = gs.zeros((n_tangent_vecs, self.n))
for i in range(n_tangent_vecs):
translation_i = translation[i]
term_1_i = coef_1[i] * gs.dot(translation_i,
gs.transpose(skew_mat[i]))
term_2_i = coef_2[i] * gs.dot(translation_i,
gs.transpose(sq_skew_mat[i]))
mask_i_float = gs.get_mask_i_float(i, n_tangent_vecs)
exp_translation += gs.outer(mask_i_float,
translation_i + term_1_i + term_2_i)
group_exp = gs.concatenate([rot_vec, exp_translation], axis=1)
group_exp = self.regularize(group_exp)
return group_exp
