def matrix_from_vector(self, vec):
"""Convert point in vector point-type to matrix.
Parameters
----------
vec : array-like, shape=[..., dimension]
Vector.
Returns
-------
mat : array-like, shape=[..., n+1, n+1]
Matrix.
"""
vec = self.regularize(vec)
n_vecs, _ = vec.shape
rot_vec = vec[:, :self.rotations.dim]
trans_vec = vec[:, self.rotations.dim:]
rot_mat = self.rotations.matrix_from_rotation_vector(rot_vec)
trans_vec = gs.reshape(trans_vec, (n_vecs, self.n, 1))
mat = gs.concatenate((rot_mat, trans_vec), axis=2)
last_lines = gs.array(gs.get_mask_i_float(self.n, self.n + 1))
last_lines = gs.to_ndarray(last_lines, to_ndim=2)
last_lines = gs.to_ndarray(last_lines, to_ndim=3)
mat = gs.concatenate((mat, last_lines), axis=1)
return mat
def matrix_from_quaternion(self, quaternion):
"""Convert a unit quaternion into a rotation vector.
Parameters
----------
quaternion : array-like, shape=[..., 4]
Returns
-------
rot_mat : array-like, shape=[..., 3]
"""
n_quaternions, _ = quaternion.shape
w, x, y, z = gs.hsplit(quaternion, 4)
rot_mat = gs.zeros((n_quaternions, ) + (self.n, ) * 2)
for i in range(n_quaternions):
# TODO(nina): Vectorize by applying the composition of
# quaternions to the identity matrix
column_1 = [
w[i]**2 + x[i]**2 - y[i]**2 - z[i]**2,
2 * x[i] * y[i] - 2 * w[i] * z[i],
2 * x[i] * z[i] + 2 * w[i] * y[i]
]
column_2 = [
2 * x[i] * y[i] + 2 * w[i] * z[i],
w[i]**2 - x[i]**2 + y[i]**2 - z[i]**2,
2 * y[i] * z[i] - 2 * w[i] * x[i]
]
column_3 = [
2 * x[i] * z[i] - 2 * w[i] * y[i],
2 * y[i] * z[i] + 2 * w[i] * x[i],
w[i]**2 - x[i]**2 - y[i]**2 + z[i]**2
]
mask_i = gs.get_mask_i_float(i, n_quaternions)
rot_mat_i = gs.transpose(gs.hstack([column_1, column_2, column_3]))
rot_mat_i = gs.to_ndarray(rot_mat_i, to_ndim=3)
rot_mat += gs.einsum('...,...ij->...ij', mask_i, rot_mat_i)
return rot_mat
def group_log_from_identity(self, point, point_type=None):
"""
Compute the group logarithm of the point at the identity.
"""
if point_type is None:
point_type = self.default_point_type
point = self.regularize(point, point_type=point_type)
rotations = self.rotations
dim_rotations = rotations.dimension
if point_type == 'vector':
rot_vec = point[:, :dim_rotations]
angle = gs.linalg.norm(rot_vec, axis=1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
translation = point[:, dim_rotations:]
skew_rot_vec = rotations.skew_matrix_from_vector(rot_vec)
sq_skew_rot_vec = gs.matmul(skew_rot_vec, skew_rot_vec)
mask_close_0 = gs.isclose(angle, 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., atol=1e-6)
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. / 12. + angle**2 / 720. +
angle**4 / 30240. +
angle**6 / 1209600.)
delta_angle = angle - gs.pi
coef_2 += mask_close_pi_float * (
1. / PI2 + (PI2 - 8.) * delta_angle / (4. * PI3) -
((PI2 - 12.) * delta_angle**2 / (4. * PI4)) +
((-192. + 12. * PI2 + PI4) * delta_angle**3 / (48. * PI5)) -
((-240. + 12. * PI2 + PI4) * delta_angle**4 / (48. * PI6)) +
((-2880. + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**5 /
(480. * PI7)) -
((-3360 + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**6 /
(480. * PI8)))
psi = 0.5 * angle * gs.sin(angle) / (1 - gs.cos(angle))
coef_2 += mask_else_float * (1 - psi) / (angle**2)
n_points, _ = point.shape
group_log_translation = gs.zeros((n_points, self.n))
for i in range(n_points):
translation_i = translation[i]
term_1_i = coef_1[i] * gs.dot(translation_i,
gs.transpose(skew_rot_vec[i]))
term_2_i = coef_2[i] * gs.dot(translation_i,
gs.transpose(sq_skew_rot_vec[i]))
mask_i_float = gs.get_mask_i_float(i, n_points)
group_log_translation += mask_i_float * (translation_i +
term_1_i + term_2_i)
group_log = gs.concatenate([rot_vec, group_log_translation],
axis=1)
assert gs.ndim(group_log) == 2
elif point_type == 'matrix':
raise NotImplementedError()
return group_log
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 log_from_identity(self, point):
"""Compute the group logarithm of the point at the identity.
Parameters
----------
point: array-like, shape=[..., 3]
Returns
-------
group_log: array-like, shape=[..., 3]
the group logarithm in the Lie algbra
"""
point = self.regularize(point)
rotations = self.rotations
dim_rotations = rotations.dim
rot_vec = point[:, :dim_rotations]
angle = gs.linalg.norm(rot_vec, axis=1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
translation = point[:, dim_rotations:]
skew_rot_vec = rotations.skew_matrix_from_vector(rot_vec)
sq_skew_rot_vec = gs.matmul(skew_rot_vec, skew_rot_vec)
mask_close_0 = gs.isclose(angle, 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., atol=1e-6)
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. / 12. + angle**2 / 720. + angle**4 /
30240. + angle**6 / 1209600.)
delta_angle = angle - gs.pi
coef_2 += mask_close_pi_float * (
1. / PI2 + (PI2 - 8.) * delta_angle / (4. * PI3) -
((PI2 - 12.) * delta_angle**2 / (4. * PI4)) +
((-192. + 12. * PI2 + PI4) * delta_angle**3 / (48. * PI5)) -
((-240. + 12. * PI2 + PI4) * delta_angle**4 / (48. * PI6)) +
((-2880. + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**5 /
(480. * PI7)) -
((-3360 + 120. * PI2 + 10. * PI4 + PI6) * delta_angle**6 /
(480. * PI8)))
psi = 0.5 * angle * gs.sin(angle) / (1 - gs.cos(angle))
coef_2 += mask_else_float * (1 - psi) / (angle**2)
n_points, _ = point.shape
log_translation = gs.zeros((n_points, self.n))
for i in range(n_points):
translation_i = translation[i]
term_1_i = coef_1[i] * gs.dot(translation_i,
gs.transpose(skew_rot_vec[i]))
term_2_i = coef_2[i] * gs.dot(translation_i,
gs.transpose(sq_skew_rot_vec[i]))
mask_i_float = gs.get_mask_i_float(i, n_points)
log_translation += gs.outer(mask_i_float,
translation_i + term_1_i + term_2_i)
return gs.concatenate([rot_vec, log_translation], axis=1)
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
def jacobian_translation(self, point, left_or_right='left'):
"""Compute the jacobian matrix corresponding to translation.
Compute the jacobian matrix of the differential
of the left/right translations from the identity to point in SO(3).
Parameters
----------
point : array-like, shape=[..., 3]
left_or_right : str, {'left', 'right'}, optional
default: 'left'
point_type : str, {'vector', 'matrix'}, optional
default: self.default_point_type
Returns
-------
jacobian : array-like, shape=[..., 3, 3]
"""
geomstats.error.check_parameter_accepted_values(
left_or_right, 'left_or_right', ['left', 'right'])
point = self.regularize(point)
n_points, _ = point.shape
angle = gs.linalg.norm(point, axis=-1)
angle = gs.expand_dims(angle, axis=-1)
coef_1 = gs.zeros([n_points, 1])
coef_2 = gs.zeros([n_points, 1])
# 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 * (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_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)
# This avoids dividing by 0.
mask_pi = gs.isclose(angle, gs.pi)
mask_pi_float = gs.cast(mask_pi, gs.float32) + self.epsilon
delta_angle = angle - gs.pi
coef_1 += mask_pi_float * (TAYLOR_COEFFS_1_AT_PI[1] * delta_angle +
TAYLOR_COEFFS_1_AT_PI[2] * delta_angle**2 +
TAYLOR_COEFFS_1_AT_PI[3] * delta_angle**3 +
TAYLOR_COEFFS_1_AT_PI[4] * delta_angle**4 +
TAYLOR_COEFFS_1_AT_PI[5] * delta_angle**5 +
TAYLOR_COEFFS_1_AT_PI[6] * delta_angle**6)
angle += mask_0_float
coef_2 += mask_pi_float * ((1 - coef_1) / angle**2)
# This avoids dividing by 0.
mask_else = ~mask_0 & ~mask_pi
mask_else_float = gs.cast(mask_else, gs.float32) + self.epsilon
# This avoids division by 0.
angle += mask_pi_float
coef_1 += mask_else_float * ((angle / 2) / gs.tan(angle / 2))
coef_2 += mask_else_float * ((1 - coef_1) / angle**2)
jacobian = gs.zeros((n_points, self.dim, self.dim))
n_points_tensor = gs.array(n_points)
for i in range(n_points):
# This avoids dividing by 0.
mask_i_float = (gs.get_mask_i_float(i, n_points_tensor) +
self.epsilon)
sign = -1.
if left_or_right == 'left':
sign = +1.
jacobian_i = (coef_1[i] * gs.eye(self.dim) +
coef_2[i] * gs.outer(point[i], point[i]) +
sign * self.skew_matrix_from_vector(point[i]) / 2.)
jacobian += gs.einsum('n,ij->nij', mask_i_float, jacobian_i)
return jacobian
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\'.')
