def tait_bryan_angles_from_quaternion_intrinsic_zyx(self, quaternion):
assert self.n == 3, ('The quaternion representation'
' and the Tait-Bryan angles representation'
' do not exist'
' for rotations in %d dimensions.' % self.n)
quaternion = gs.to_ndarray(quaternion, to_ndim=2)
w, x, y, z = gs.hsplit(quaternion, 4)
angle_1 = gs.arctan2(y * z + w * x, 1. / 2. - (x**2 + y**2))
angle_2 = gs.arcsin(-2. * (x * z - w * y))
angle_3 = gs.arctan2(x * y + w * z, 1. / 2. - (y**2 + z**2))
tait_bryan_angles = gs.concatenate([angle_3, angle_2, angle_1], axis=1)
return tait_bryan_angles
def tait_bryan_angles_from_quaternion_intrinsic_xyz(self, quaternion):
assert self.n == 3, ('The quaternion representation'
' and the Tait-Bryan angles representation'
' do not exist'
' for rotations in %d dimensions.' % self.n)
quaternion = gs.to_ndarray(quaternion, to_ndim=2)
w, x, y, z = gs.hsplit(quaternion, 4)
angle_3 = gs.arctan2(2. * (x * y + w * z),
w * w + x * x - y * y - z * z)
angle_2 = gs.arcsin(-2. * (x * z - w * y))
angle_1 = gs.arctan2(2. * (y * z + w * x),
w * w - x * x - y * y + z * z)
tait_bryan_angles = gs.concatenate([angle_1, angle_2, angle_3], axis=1)
return tait_bryan_angles
def extrinsic_to_spherical(self, point_extrinsic):
"""Convert point from extrinsic to spherical coordinates.
Convert from the extrinsic coordinates, i.e. embedded in Euclidean
space of dim 3 to spherical coordinates in the hypersphere.
Spherical coordinates are defined from the north pole, i.e.
angles [0., 0.] correspond to point [0., 0., 1.].
Only implemented in dimension 2.
Parameters
----------
point_extrinsic : array-like, shape=[..., dim]
Point on the sphere, in extrinsic coordinates.
Returns
-------
point_spherical : array_like, shape=[..., dim + 1]
Point on the sphere, in spherical coordinates relative to the
north pole.
"""
if self.dim != 2:
raise NotImplementedError(
"The conversion from to extrinsic coordinates "
"spherical coordinates is implemented"
" only in dimension 2.")
theta = gs.arccos(point_extrinsic[..., -1])
x = point_extrinsic[..., 0]
y = point_extrinsic[..., 1]
phi = gs.arctan2(y, x)
phi = gs.where(phi < 0, phi + 2 * gs.pi, phi)
return gs.stack([theta, phi], axis=-1)
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 tait_bryan_angles_from_quaternion_intrinsic_zyx(quaternion):
"""Convert quaternion to tait bryan representation of order zyx.
Parameters
----------
quaternion : array-like, shape=[..., 4]
Returns
-------
tait_bryan_angles : array-like, shape=[..., 3]
"""
w, x, y, z = gs.hsplit(quaternion, 4)
angle_1 = gs.arctan2(y * z + w * x, 1. / 2. - (x**2 + y**2))
angle_2 = gs.arcsin(-2. * (x * z - w * y))
angle_3 = gs.arctan2(x * y + w * z, 1. / 2. - (y**2 + z**2))
tait_bryan_angles = gs.concatenate([angle_1, angle_2, angle_3], axis=1)
return tait_bryan_angles
def grad(y_pred,
y_true,
metric=SE3.left_canonical_metric,
representation='vector'):
"""
Closed-form for the gradient of pose_loss.
:return: tangent vector at point y_pred.
"""
if gs.ndim(y_pred) == 1:
y_pred = gs.expand_dims(y_pred, axis=0)
if gs.ndim(y_true) == 1:
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'vector':
grad = lie_group.grad(y_pred, y_true, SE3, metric)
if representation == 'quaternion':
y_pred_rot_vec = SO3.rotation_vector_from_quaternion(y_pred[:, :4])
y_pred_pose = gs.hstack([y_pred_rot_vec, y_pred[:, 4:]])
y_true_rot_vec = SO3.rotation_vector_from_quaternion(y_true[:, :4])
y_true_pose = gs.hstack([y_true_rot_vec, y_true[:, 4:]])
grad = lie_group.grad(y_pred_pose, y_true_pose, SE3, metric)
quat_scalar = y_pred[:, :1]
quat_vec = y_pred[:, 1:4]
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_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 *
quat_vec_norm) + 2 * quat_arctan2 / quat_vec_norm * gs.eye(3))
differential_scalar_t = gs.transpose(differential_scalar, axes=(1, 0))
upper_left_block = gs.hstack(
(differential_scalar_t, differential_vec[0]))
upper_right_block = gs.zeros((3, 3))
lower_right_block = gs.eye(3)
lower_left_block = gs.zeros((3, 4))
top = gs.hstack((upper_left_block, upper_right_block))
bottom = gs.hstack((lower_left_block, lower_right_block))
differential = gs.vstack((top, bottom))
differential = gs.expand_dims(differential, axis=0)
grad = gs.einsum('ni,nij->ni', grad, differential)
grad = gs.squeeze(grad, axis=0)
return grad
def grad(y_pred,
y_true,
metric=SE3.left_canonical_metric,
representation='vector'):
"""
Closed-form for the gradient of pose_loss.
:return: tangent vector at point y_pred.
"""
if y_pred.ndim == 1:
y_pred = gs.expand_dims(y_pred, axis=0)
if y_true.ndim == 1:
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'vector':
grad = lie_group.grad(y_pred, y_true, SE3, metric)
if representation == 'quaternion':
y_pred_rot_vec = SO3.rotation_vector_from_quaternion(y_pred[:, :4])
y_pred_pose = gs.hstack([y_pred_rot_vec, y_pred[:, 4:]])
y_true_rot_vec = SO3.rotation_vector_from_quaternion(y_true[:, :4])
y_true_pose = gs.hstack([y_true_rot_vec, y_true[:, 4:]])
grad = lie_group.grad(y_pred_pose, y_true_pose, SE3, metric)
differential = gs.zeros((1, 6, 7))
upper_left_block = gs.zeros((1, 3, 4))
lower_right_block = gs.zeros((1, 3, 3))
quat_scalar = y_pred[:, :1]
quat_vec = y_pred[:, 1:4]
quat_vec_norm = gs.linalg.norm(quat_vec, axis=1)
quat_sq_norm = quat_vec_norm**2 + quat_scalar**2
# TODO(nina): check that this sq norm is 1?
quat_arctan2 = gs.arctan2(quat_vec_norm, quat_scalar)
differential_scalar = -2 * quat_vec / (quat_sq_norm)
differential_vec = (
2 *
(quat_scalar / quat_sq_norm - 2 * quat_arctan2 / quat_vec_norm) *
gs.outer(quat_vec, quat_vec) / quat_vec_norm**2 +
2 * quat_arctan2 / quat_vec_norm * gs.eye(3))
upper_left_block[0, :, :1] = differential_scalar.transpose()
upper_left_block[0, :, 1:] = differential_vec
lower_right_block[0, :, :] = gs.eye(3)
differential[0, :3, :4] = upper_left_block
differential[0, 3:, 4:] = lower_right_block
grad = gs.matmul(grad, differential)
grad = gs.squeeze(grad, axis=0)
return grad
def convert_to_polar_coordinates(self, points):
"""Assign polar coordinates to given pre-shapes."""
aligned_points = S32.align(points, self.pole)
speeds = S32.ambient_metric.log(aligned_points, self.pole)
coords_theta = gs.arctan2(
S32.ambient_metric.inner_product(speeds, self.na),
S32.ambient_metric.inner_product(speeds, self.ua))
coords_phi = 2. * S32.ambient_metric.dist(self.pole, aligned_points)
return coords_theta, coords_phi
def convert_to_polar_coordinates(self, points):
"""Assign polar coordinates to given pre-shapes."""
aligned_points = S33.align(points, self.centre)
aligned_points2d = aligned_points[..., :, :2]
speeds = S32.ambient_metric.log(aligned_points2d, self.pole)
coords_r = S32.ambient_metric.dist(self.pole, aligned_points2d)
coords_theta = gs.arctan2(
S32.ambient_metric.inner_product(speeds, self.na),
S32.ambient_metric.inner_product(speeds, self.ua))
return coords_r, coords_theta
def tait_bryan_angles_from_quaternion_intrinsic_xyz(quaternion):
"""Convert quaternion to tait bryan representation of order xyz.
Parameters
----------
quaternion : array-like, shape=[..., 4]
Returns
-------
tait_bryan_angles : array-like, shape=[..., 3]
"""
w, x, y, z = gs.hsplit(quaternion, 4)
angle_1 = gs.arctan2(2. * (-x * y + w * z),
w * w + x * x - y * y - z * z)
angle_2 = gs.arcsin(2 * (x * z + w * y))
angle_3 = gs.arctan2(2. * (-y * z + w * x),
w * w + z * z - x * x - y * y)
tait_bryan_angles = gs.concatenate([angle_1, angle_2, angle_3], axis=1)
return tait_bryan_angles
def convert_to_klein_coordinates(points):
poincare_coords = points[:, 1:] / (1 + points[:, :1])
poincare_radius = gs.linalg.norm(poincare_coords, axis=1)
poincare_angle = gs.arctan2(poincare_coords[:, 1], poincare_coords[:, 0])
klein_radius = 2 * poincare_radius / (1 + poincare_radius ** 2)
klein_angle = poincare_angle
coords_0 = gs.expand_dims(klein_radius * gs.cos(klein_angle), axis=1)
coords_1 = gs.expand_dims(klein_radius * gs.sin(klein_angle), axis=1)
klein_coords = gs.concatenate([coords_0, coords_1], axis=1)
return klein_coords
def convert_to_klein_coordinates(self, points):
poincare_coords = points[:, 1:] / (1 + points[:, :1])
poincare_radius = gs.linalg.norm(poincare_coords, axis=1)
poincare_angle = gs.arctan2(poincare_coords[:, 1], poincare_coords[:,
0])
klein_radius = 2 * poincare_radius / (1 + poincare_radius**2)
klein_angle = poincare_angle
klein_coords = gs.zeros_like(poincare_coords)
klein_coords[:, 0] = klein_radius * gs.cos(klein_angle)
klein_coords[:, 1] = klein_radius * gs.sin(klein_angle)
return klein_coords
def extrinsic_to_angle(point_extrinsic):
"""Compute the angle of a point in the plane.
Convert from the extrinsic coordinates in the 2d plane to angle in
radians. Angle 0 corresponds to point [1., 0.]. This method is only
implemented in dimension 1.
Parameters
----------
point_extrinsic : array-like, shape=[...]
Point on the circle, in extrinsic coordinates in Euclidean space.
Returns
-------
point_angle : array_like, shape=[..., 2]
Point on the circle, i.e. an angle in radians in [-Pi, Pi].
"""
return gs.arctan2(point_extrinsic[..., 1], point_extrinsic[..., 0])
def grad(y_pred, y_true,
metric=SO3.bi_invariant_metric,
representation='vector'):
y_pred = gs.expand_dims(y_pred, axis=0)
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'vector':
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)
grad = lie_group.grad(y_pred, y_true, SO3, metric)
grad = gs.matmul(grad, differential)
grad = gs.squeeze(grad, axis=0)
return grad
def grad(y_pred, y_true,
metric=SO3.bi_invariant_metric,
representation='vector'):
y_pred = gs.expand_dims(y_pred, axis=0)
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'vector':
grad = lie_group.grad(y_pred, y_true, SO3, metric)
if representation == 'quaternion':
differential = gs.zeros((1, 6, 7))
differential = gs.zeros((1, 3, 4))
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_vec = (2 * (quat_scalar / quat_sq_norm
- 2 * quat_arctan2 / quat_vec_norm)
* gs.outer(quat_vec, quat_vec) / quat_vec_norm ** 2
+ 2 * quat_arctan2 / quat_vec_norm * gs.eye(3))
differential[0, :, :1] = differential_scalar.transpose()
differential[0, :, 1:] = differential_vec
y_pred = SO3.rotation_vector_from_quaternion(y_pred)
y_true = SO3.rotation_vector_from_quaternion(y_true)
grad = lie_group.grad(y_pred, y_true, SO3, metric)
grad = gs.matmul(grad, differential)
grad = gs.squeeze(grad, axis=0)
return grad
def grad(y_pred,
y_true,
metric=SE3.left_canonical_metric,
representation='vector'):
"""Closed-form for the gradient of pose_loss.
Parameters
----------
y_pred : array-like
Prediction on SE(3).
y_true : array-like
Ground-truth on SE(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.
"""
if gs.ndim(y_pred) == 1:
y_pred = gs.expand_dims(y_pred, axis=0)
if gs.ndim(y_true) == 1:
y_true = gs.expand_dims(y_true, axis=0)
if representation == 'vector':
lie_grad = lie_group.grad(y_pred, y_true, SE3, metric)
if representation == 'quaternion':
y_pred_rot_vec = SO3.rotation_vector_from_quaternion(y_pred[:, :4])
y_pred_pose = gs.hstack([y_pred_rot_vec, y_pred[:, 4:]])
y_true_rot_vec = SO3.rotation_vector_from_quaternion(y_true[:, :4])
y_true_pose = gs.hstack([y_true_rot_vec, y_true[:, 4:]])
lie_grad = lie_group.grad(y_pred_pose, y_true_pose, SE3, metric)
quat_scalar = y_pred[:, :1]
quat_vec = y_pred[:, 1:4]
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_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 *
quat_vec_norm) + 2 * quat_arctan2 / quat_vec_norm * gs.eye(3))
differential_scalar_t = gs.transpose(differential_scalar, axes=(1, 0))
upper_left_block = gs.hstack(
(differential_scalar_t, differential_vec[0]))
upper_right_block = gs.zeros((3, 3))
lower_right_block = gs.eye(3)
lower_left_block = gs.zeros((3, 4))
top = gs.hstack((upper_left_block, upper_right_block))
bottom = gs.hstack((lower_left_block, lower_right_block))
differential = gs.vstack((top, bottom))
differential = gs.expand_dims(differential, axis=0)
lie_grad = gs.einsum('ni,nij->ni', lie_grad, differential)
lie_grad = gs.squeeze(lie_grad, axis=0)
return lie_grad