def exp(self, tangent_vec, base_point=None):
"""Compute Riemannian exponential of tan. vector wrt to base point.
Parameters
----------
tangent_vec : array-like, shape=[n_samples, dimension]
Tangent vector at a base point.
base_point : array-like, shape=[n_samples, dimension]
Point in the group.
Returns
-------
exp : array-like, shape=[n_samples, dimension]
Point in the group equal to the Riemannian exponential
of tangent_vec at the base point.
"""
if base_point is None:
base_point = self.group.identity
base_point = self.group.regularize(base_point)
if base_point is self.group.identity:
return self.exp_from_identity(tangent_vec)
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
assert gs.ndim(tangent_vec) == 2
n_tangent_vecs, _ = tangent_vec.shape
n_base_points, _ = base_point.shape
if n_tangent_vecs == 1:
tangent_vec = gs.tile(tangent_vec, (n_base_points, 1))
if n_base_points == 1:
base_point = gs.tile(base_point, (n_tangent_vecs, 1))
jacobian = self.group.jacobian_translation(
point=base_point,
left_or_right=self.left_or_right)
assert gs.ndim(jacobian) == 3
inv_jacobian = gs.linalg.inv(jacobian)
inv_jacobian_transposed = gs.transpose(inv_jacobian, axes=(0, 2, 1))
tangent_vec_at_id = gs.einsum(
'ni,nij->nj', tangent_vec, inv_jacobian_transposed)
exp_from_id = self.exp_from_identity(tangent_vec_at_id)
if self.left_or_right == 'left':
exp = self.group.compose(base_point, exp_from_id)
else:
exp = self.group.compose(exp_from_id, base_point)
exp = self.group.regularize(exp)
return exp
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 : The matrix exponential of rot_vec
"""
# TODO (nguigs): find usecase for this method
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 = 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 test_srv_metric_dist_and_geod(self):
"""
Test that the length of the geodesic gives the distance.
N.B: Here curve_a and curve_b are seen as curves in R3 and not S2.
"""
geod = self.srv_metric_r3.geodesic(initial_curve=self.curve_a,
end_curve=self.curve_b)
geod = geod(self.times)
srv = self.srv_metric_r3.square_root_velocity(geod)
srv_derivative = self.n_discretized_curves * (srv[1:, :] - srv[:-1, :])
norms = self.l2_metric_r3.norm(srv_derivative, geod[:-1, :-1, :])
result = gs.sum(norms, 0) / self.n_discretized_curves
result = gs.to_ndarray(result, to_ndim=1)
result = gs.to_ndarray(result, to_ndim=2, axis=1)
expected = self.srv_metric_r3.dist(self.curve_a, self.curve_b)
expected = gs.to_ndarray(expected, to_ndim=1)
expected = gs.to_ndarray(expected, to_ndim=2, axis=1)
self.assertAllClose(result, expected)
def inner_product_matrix(self, base_point=None):
"""
Inner product matrix, independent of the base point.
Parameters
----------
base_point: array-like, shape=[n_samples, dimension]
Returns
-------
inner_prod_mat: array-like, shape=[n_samples, dimension, dimension]
"""
inner_prod_mat = gs.eye(self.dimension - 1, self.dimension - 1)
first_row = gs.array([0.] * (self.dimension - 1))
first_row = gs.to_ndarray(first_row, to_ndim=2, axis=1)
inner_prod_mat = gs.vstack([gs.transpose(first_row), inner_prod_mat])
first_column = gs.array([-1.] + [0.] * (self.dimension - 1))
first_column = gs.to_ndarray(first_column, to_ndim=2, axis=1)
inner_prod_mat = gs.hstack([first_column, inner_prod_mat])
return inner_prod_mat
def group_exp(self, tangent_vec, base_point=None, point_type=None):
"""
Compute the group exponential at point base_point
of tangent vector tangent_vec.
"""
if point_type is None:
point_type = self.default_point_type
identity = self.get_identity(point_type=point_type)
identity = self.regularize(identity, point_type=point_type)
if base_point is None:
base_point = identity
base_point = self.regularize(base_point, point_type=point_type)
if point_type == 'vector':
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
if point_type == 'matrix':
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=3)
base_point = gs.to_ndarray(base_point, to_ndim=3)
n_tangent_vecs = tangent_vec.shape[0]
n_base_points = base_point.shape[0]
assert (tangent_vec.shape == base_point.shape or n_tangent_vecs == 1
or n_base_points == 1)
if n_tangent_vecs == 1:
tangent_vec = gs.array([tangent_vec[0]] * n_base_points)
if n_base_points == 1:
base_point = gs.array([base_point[0]] * n_tangent_vecs)
result = gs.cond(pred=gs.allclose(base_point, identity),
true_fn=lambda: self.group_exp_from_identity(
tangent_vec, point_type=point_type),
false_fn=lambda: self.group_exp_not_from_identity(
tangent_vec, base_point, point_type))
return result
def test_point_to_pdf(self, dim, point, n_samples):
point = gs.to_ndarray(point, 2)
n_points = point.shape[0]
pdf = self.space(dim).point_to_pdf(point)
alpha = gs.ones(dim)
samples = self.space(dim).sample(alpha, n_samples)
result = pdf(samples)
pdf = []
for i in range(n_points):
pdf.append(
gs.array([dirichlet.pdf(x, point[i, :]) for x in samples]))
expected = gs.squeeze(gs.stack(pdf, axis=0))
self.assertAllClose(result, expected)
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 exp(self, tangent_vec, base_point):
"""
Riemannian exponential of a tangent vector wrt to a base point.
"""
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)
coef_1 = gs.zeros_like(norm_tangent_vec)
coef_2 = gs.zeros_like(norm_tangent_vec)
coef_1[mask_0] = (1. +
COSH_TAYLOR_COEFFS[2] * norm_tangent_vec[mask_0]**2 +
COSH_TAYLOR_COEFFS[4] * norm_tangent_vec[mask_0]**4 +
COSH_TAYLOR_COEFFS[6] * norm_tangent_vec[mask_0]**6 +
COSH_TAYLOR_COEFFS[8] * norm_tangent_vec[mask_0]**8)
coef_2[mask_0] = (1. +
SINH_TAYLOR_COEFFS[3] * norm_tangent_vec[mask_0]**2 +
SINH_TAYLOR_COEFFS[5] * norm_tangent_vec[mask_0]**4 +
SINH_TAYLOR_COEFFS[7] * norm_tangent_vec[mask_0]**6 +
SINH_TAYLOR_COEFFS[9] * norm_tangent_vec[mask_0]**8)
coef_1[mask_else] = gs.cosh(norm_tangent_vec[mask_else])
coef_2[mask_else] = (gs.sinh(norm_tangent_vec[mask_else]) /
norm_tangent_vec[mask_else])
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 dist(self, curve_a, curve_b):
"""
Geodesic distance between two discretized curves.
"""
assert curve_a.shape == curve_b.shape
curve_a = gs.to_ndarray(curve_a, to_ndim=3)
curve_b = gs.to_ndarray(curve_b, to_ndim=3)
n_curves, n_sampling_points, n_coords = curve_a.shape
curve_a = gs.reshape(curve_a, (n_curves * n_sampling_points, n_coords))
curve_b = gs.reshape(curve_b, (n_curves * n_sampling_points, n_coords))
dist = self.embedding_metric.dist(curve_a, curve_b)
dist = gs.reshape(dist, (n_curves, n_sampling_points))
n_sampling_points_float = gs.array(n_sampling_points)
n_sampling_points_float = gs.cast(n_sampling_points_float, gs.float32)
dist = gs.sqrt(gs.sum(dist**2, -1) / n_sampling_points_float)
dist = gs.to_ndarray(dist, to_ndim=1)
dist = gs.to_ndarray(dist, to_ndim=2, axis=1)
return dist
def exp(self, tangent_vec, base_curve):
"""
Riemannian exponential of a tangent vector wrt to a base curve.
"""
if not isinstance(self.embedding_metric, EuclideanMetric):
raise AssertionError('The exponential map is only implemented '
'for dicretized curves embedded in a '
'Euclidean space.')
base_curve = gs.to_ndarray(base_curve, to_ndim=3)
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=3)
n_sampling_points = base_curve.shape[1]
base_curve_srv = self.square_root_velocity(base_curve)
tangent_vec_derivative = (n_sampling_points - 1) * (
tangent_vec[:, 1:, :] - tangent_vec[:, :-1, :])
base_curve_velocity = (n_sampling_points - 1) * (base_curve[:, 1:, :] -
base_curve[:, :-1, :])
base_curve_velocity_norm = self.pointwise_norm(base_curve_velocity,
base_curve[:, :-1, :])
inner_prod = self.pointwise_inner_product(tangent_vec_derivative,
base_curve_velocity,
base_curve[:, :-1, :])
coef_1 = 1 / gs.sqrt(base_curve_velocity_norm)
coef_2 = -1 / (2 * base_curve_velocity_norm**(5 / 2)) * inner_prod
term_1 = gs.einsum('ij,ijk->ijk', coef_1, tangent_vec_derivative)
term_2 = gs.einsum('ij,ijk->ijk', coef_2, base_curve_velocity)
srv_initial_derivative = term_1 + term_2
end_curve_srv = self.l2_metric.exp(tangent_vec=srv_initial_derivative,
base_curve=base_curve_srv)
end_curve_starting_point = self.embedding_metric.exp(
tangent_vec=tangent_vec[:, 0, :], base_point=base_curve[:, 0, :])
end_curve = self.square_root_velocity_inverse(
end_curve_srv, end_curve_starting_point)
return end_curve
def curve_on_geodesic(t):
t = gs.cast(t, gs.float32)
t = gs.to_ndarray(t, to_ndim=1)
t = gs.to_ndarray(t, to_ndim=2, axis=1)
new_initial_curve = gs.to_ndarray(initial_curve,
to_ndim=curve_ndim + 1)
new_initial_tangent_vec = gs.to_ndarray(initial_tangent_vec,
to_ndim=curve_ndim + 1)
tangent_vecs = gs.einsum('il,nkm->ikm', t, new_initial_tangent_vec)
def point_on_curve(tangent_vec):
assert gs.ndim(tangent_vec) >= 2
exp = self.exp(tangent_vec=tangent_vec,
base_curve=new_initial_curve)
return exp
curve_at_time_t = gs.vectorize(tangent_vecs,
point_on_curve,
signature='(i,j)->(i,j)')
return curve_at_time_t
def tangent_sphere_to_simplex(tangent_vec, base_point):
"""Send tangent vector of the sphere to tangent space of simplex.
This is the differential of the sphere_to_simplex map.
Parameters
----------
tangent_vec : array-like, shape=[..., dim + 1]
Tangent vec to the sphere at base point.
base_point : array-like, shape=[..., dim + 1]
Point of the sphere.
Returns
-------
tangent_vec_simplex : array-like, shape=[..., dim + 1]
Tangent vec to the simplex at the image of
base point by sphere_to_simplex.
"""
base_point = gs.to_ndarray(base_point, to_ndim=2)
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
tangent_vec_simplex = gs.einsum("...i,...i->...i", tangent_vec, 2 * base_point)
return gs.squeeze(tangent_vec_simplex)
def vectorize_kwargs(input_types, kwargs):
"""Vectorize input kwargs.
Transform input array-like kwargs into their fully-vectorized form,
where "fully-vectorized" means that:
- one scalar has shape [1, 1],
- n scalars have shape [n, 1],
- one d-D vector has shape [1, d],
- n d-D vectors have shape [n, d],
etc.
Parameters
----------
input_types :list
Point types corresponding to the args.
kwargs : dict
Kwargs of a function.
Returns
-------
vect_kwargs : dict
Kwargs of the function in their fully-vectorized form.
"""
vect_kwargs = {}
for key_arg, input_type in zip(kwargs.keys(), input_types):
arg = kwargs[key_arg]
if input_type == "scalar":
vect_arg = gs.to_ndarray(arg, to_ndim=1)
vect_arg = gs.to_ndarray(vect_arg, to_ndim=2, axis=1)
elif input_type in POINT_TYPES and arg is not None:
vect_arg = gs.to_ndarray(arg,
to_ndim=POINT_TYPES_TO_NDIMS[input_type])
elif input_type in OTHER_TYPES or arg is None:
vect_arg = arg
else:
raise ValueError(ERROR_MSG % input_type)
vect_kwargs[key_arg] = vect_arg
return vect_kwargs
def group_exp(self, tangent_vec, base_point=None, point_type=None):
"""
Compute the group exponential at point base_point
of tangent vector tangent_vec.
"""
if point_type is None:
point_type = self.default_point_type
identity = self.get_identity(point_type=point_type)
identity = self.regularize(identity, point_type=point_type)
if base_point is None:
base_point = identity
base_point = self.regularize(base_point, point_type=point_type)
if gs.allclose(base_point, identity):
return self.group_exp_from_identity(tangent_vec,
point_type=point_type)
jacobian = self.jacobian_translation(point=base_point,
left_or_right='left',
point_type=point_type)
if point_type == 'vector':
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
inv_jacobian = gs.linalg.inv(jacobian)
tangent_vec_at_id = gs.einsum(
'ij,ijk->ik', tangent_vec,
gs.transpose(inv_jacobian, axes=(0, 2, 1)))
group_exp_from_identity = self.group_exp_from_identity(
tangent_vec=tangent_vec_at_id, point_type=point_type)
group_exp = self.compose(base_point,
group_exp_from_identity,
point_type=point_type)
group_exp = self.regularize(group_exp, point_type=point_type)
return group_exp
elif point_type == 'matrix':
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=3)
raise NotImplementedError()
def exp(self, tangent_vec, base_point):
"""Compute exp map of a base point in tangent vector direction.
The Riemannian exponential is vector addition in the Euclidean space.
Parameters
----------
tangent_vec: array-like, shape=[n_samples, dimension]
or shape=[1, dimension]
base_point: array-like, shape=[n_samples, dimension]
or shape=[1, dimension]
Returns
-------
exp: array-like, shape=[n_samples, dimension]
or shape-[n_samples, dimension]
"""
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
exp = base_point + tangent_vec
return exp
def log(self, point, base_point):
"""Compute log map using a base point and other point.
The Riemannian logarithm is the subtraction in the Euclidean space.
Parameters
----------
point: array-like, shape=[n_samples, dimension]
or shape=[1, dimension]
base_point: array-like, shape=[n_samples, dimension]
or shape=[1, dimension]
Returns
-------
log: array-like, shape=[n_samples, dimension]
or shape-[n_samples, dimension]
"""
point = gs.to_ndarray(point, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
log = point - base_point
return log
def exp(self, tangent_vec, base_point, n_steps=N_STEPS):
"""Exponential map associated to the Fisher information metric.
Exponential map at base_point of tangent_vec computed by integration
of the geodesic equation (initial value problem), using the
christoffel symbols.
Parameters
----------
tangent_vec : array-like, shape=[..., dim]
Tangent vector at base point.
base_point : array-like, shape=[..., dim]
Base point.
n_steps : int
Number of steps for integration.
Optional, default: 100.
Returns
-------
exp : array-like, shape=[..., dim]
Riemannian exponential.
"""
base_point = gs.to_ndarray(base_point, to_ndim=2)
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
def ivp(state, _):
"""Reformat the initial value problem geodesic ODE."""
position, velocity = state[:2], state[2:]
eq = self.geodesic_equation(velocity=velocity, position=position)
return gs.hstack([velocity, eq])
times = gs.linspace(0, 1, n_steps + 1)
exp = []
for point, vec in zip(base_point, tangent_vec):
initial_state = gs.hstack([point, vec])
geodesic = odeint(
ivp, initial_state, times, (), rtol=1e-6)
exp.append(geodesic[-1, :2])
return exp[0] if len(base_point) == 1 else gs.stack(exp)
def plot_and_save_video(geodesics,
loss,
size=20,
fps=10,
dpi=100,
out='out.mp4',
color='red'):
"""Render a set of geodesics and save it to an mpeg 4 file."""
FFMpegWriter = animation.writers['ffmpeg']
writer = FFMpegWriter(fps=fps)
fig = plt.figure(figsize=(size, size))
ax = fig.add_subplot(111, projection='3d', aspect='equal')
sphere = visualization.Sphere()
sphere.plot_heatmap(ax, loss)
points = gs.to_ndarray(geodesics[0], to_ndim=2)
sphere.add_points(points)
sphere.draw(ax, color=color, marker='.')
with writer.saving(fig, out, dpi=dpi):
for points in geodesics[1:]:
points = gs.to_ndarray(points, to_ndim=2)
sphere.draw_points(ax, points=points, color=color, marker='.')
writer.grab_frame()
def exp_not_from_identity(self, tangent_vec, base_point, point_type):
"""Calculate the group exponential at base_point.
Parameters
----------
tangent_vec : array-like, shape=[n_samples, {dimension,[n,n]}]
base_point : array-like, shape=[n_samples, {dimension,[n,n]}]
point_type : str, {'vector', 'matrix'}
default: the default point type
Returns
-------
exp : array-like, shape=[n_samples, {dimension,[n,n]}]
the computed exponential
"""
jacobian = self.jacobian_translation(point=base_point,
left_or_right="left",
point_type=point_type)
if point_type == "vector":
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=2)
inv_jacobian = gs.linalg.inv(jacobian)
tangent_vec_at_id = gs.einsum(
"ni,nij->nj",
tangent_vec,
gs.transpose(inv_jacobian, axes=(0, 2, 1)),
)
exp_from_identity = self.exp_from_identity(
tangent_vec=tangent_vec_at_id, point_type=point_type)
exp = self.compose(base_point,
exp_from_identity,
point_type=point_type)
exp = self.regularize(exp, point_type=point_type)
return exp
elif point_type == "matrix":
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=3)
raise NotImplementedError()
def intrinsic_to_extrinsic_coords(self, point_intrinsic):
"""
Convert the parameterization of a point on the Hyperbolic space
from its intrinsic coordinates, to its extrinsic coordinates
in Minkowski space.
Parameters
----------
point_intrinsic : array-like, shape=[n_samples, dimension]
Returns
-------
point_extrinsic : array-like, shape=[n_samples, dimension + 1]
"""
point_intrinsic = gs.to_ndarray(point_intrinsic, to_ndim=2)
coord_0 = gs.sqrt(1. + gs.linalg.norm(point_intrinsic, axis=-1) ** 2)
coord_0 = gs.to_ndarray(coord_0, to_ndim=2, axis=1)
point_extrinsic = gs.concatenate([coord_0, point_intrinsic], axis=-1)
return point_extrinsic
def quaternion_from_matrix(self, rot_mat):
"""
Convert a rotation matrix into a unit quaternion.
"""
assert self.n == 3, ('The quaternion representation does not exist'
' for rotations in %d dimensions.' % self.n)
rot_mat = gs.to_ndarray(rot_mat, to_ndim=3)
rot_vec = self.rotation_vector_from_matrix(rot_mat)
quaternion = self.quaternion_from_rotation_vector(rot_vec)
assert quaternion.ndim == 2
return quaternion
def main():
y_pred = gs.array([1., 1.5, -0.3])
y_true = gs.array([0.1, 1.8, -0.1])
loss_rot_vec = loss(y_pred, y_true)
grad_rot_vec = grad(y_pred, y_true)
logging.info('The loss between the rotation vectors is: {}'.format(
loss_rot_vec))
logging.info('The Riemannian gradient is: {}'.format(
grad_rot_vec))
angle = gs.array(gs.pi / 6)
cos = gs.cos(angle / 2)
sin = gs.sin(angle / 2)
u = gs.array([1., 2., 3.])
u = u / gs.linalg.norm(u)
scalar = gs.to_ndarray(cos, to_ndim=1)
vec = sin * u
y_pred_quaternion = gs.concatenate([scalar, vec], axis=0)
angle = gs.array(gs.pi / 7)
cos = gs.cos(angle / 2)
sin = gs.sin(angle / 2)
u = gs.array([1., 2., 3.])
u = u / gs.linalg.norm(u)
scalar = gs.to_ndarray(cos, to_ndim=1)
vec = sin * u
y_true_quaternion = gs.concatenate([scalar, vec], axis=0)
loss_quaternion = loss(y_pred_quaternion, y_true_quaternion,
representation='quaternion')
grad_quaternion = grad(y_pred_quaternion, y_true_quaternion,
representation='quaternion')
logging.info('The loss between the quaternions is: {}'.format(
loss_quaternion))
logging.info('The Riemannian gradient is: {}'.format(
grad_quaternion))
def inner_product(self, tangent_vec_a, tangent_vec_b, base_landmarks):
"""Compute inner product between tangent vectors at base landmark set.
Parameters
----------
tangent_vec_a
tangent_vec_b
base_landmarks
Returns
-------
inner_prod
"""
if not (tangent_vec_a.shape == tangent_vec_b.shape
and tangent_vec_a.shape == base_landmarks.shape):
raise NotImplementedError
n_landmark_sets, n_landmarks_per_set, n_coords = tangent_vec_a.shape
new_dim = n_landmark_sets * n_landmarks_per_set
tangent_vec_a = gs.reshape(tangent_vec_a, (new_dim, n_coords))
tangent_vec_b = gs.reshape(tangent_vec_b, (new_dim, n_coords))
base_landmarks = gs.reshape(base_landmarks, (new_dim, n_coords))
inner_prod = self.ambient_metric.inner_product(tangent_vec_a,
tangent_vec_b,
base_landmarks)
inner_prod = gs.reshape(inner_prod,
(n_landmark_sets, n_landmarks_per_set))
inner_prod = gs.sum(inner_prod, -1)
n_landmarks_per_set_float = gs.array(n_landmarks_per_set)
n_landmarks_per_set_float = gs.cast(n_landmarks_per_set_float,
gs.float32)
inner_prod = inner_prod / n_landmarks_per_set_float
inner_prod = gs.to_ndarray(inner_prod, to_ndim=1)
inner_prod = gs.to_ndarray(inner_prod, to_ndim=2, axis=1)
return inner_prod
def log(self, point, base_point=None, point_type=None):
"""Compute the Riemannian logarithm of a point.
Parameters
----------
point : array-like, shape=[n_samples, dimension]
Point on the manifold.
base_point : array-like, shape=[n_samples, dimension]
Point on the manifold.
point_type : str, {'vector', 'matrix'}
Type of representation used for points.
Returns
-------
log : array-like, shape=[n_samples, dimension]
Tangent vector at the base point equal to the Riemannian logarithm
of point at the base point.
"""
if base_point is None:
base_point = [None, ] * self.n_metrics
if point_type is None:
point_type = self.default_point_type
if point_type == 'vector':
point = gs.to_ndarray(point, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
intrinsic = self._is_intrinsic(base_point)
args = {'point': point, 'base_point': base_point}
log = self._iterate_over_metrics('log', args, intrinsic)
return gs.hstack(log)
elif point_type == 'matrix':
point = gs.to_ndarray(point, to_ndim=3)
base_point = gs.to_ndarray(base_point, to_ndim=3)
return gs.stack(
[self.metrics[i].log(point[:, i], base_point[:, i])
for i in range(self.n_metrics)], axis=1)
else:
raise ValueError('invalid point_type argument: {}, expected '
'either matrix of vector'.format(point_type))
def test_exp_vectorization_single_samples(self):
dim = self.dimension + 1
one_vec = self.space.random_uniform()
one_base_point = self.space.random_uniform()
one_tangent_vec = self.space.to_tangent(one_vec,
base_point=one_base_point)
result = self.metric.exp(one_tangent_vec, one_base_point)
self.assertAllClose(gs.shape(result), (dim, ))
one_base_point = gs.to_ndarray(one_base_point, to_ndim=2)
result = self.metric.exp(one_tangent_vec, one_base_point)
self.assertAllClose(gs.shape(result), (1, dim))
one_tangent_vec = gs.to_ndarray(one_tangent_vec, to_ndim=2)
result = self.metric.exp(one_tangent_vec, one_base_point)
self.assertAllClose(gs.shape(result), (1, dim))
one_base_point = self.space.random_uniform()
result = self.metric.exp(one_tangent_vec, one_base_point)
self.assertAllClose(gs.shape(result), (1, dim))
def exp(self, tangent_vec, base_point=None, point_type=None):
"""Compute the Riemannian exponential of a tangent vector.
Parameters
----------
tangent_vec : array-like, shape=[n_samples, dimension]
Tangent vector at a base point.
base_point : array-like, shape=[n_samples, dimension]
Point on the manifold.
point_type : str, {'vector', 'matrix'}
Type of representation used for points.
Returns
-------
exp : array-like, shape=[n_samples, dimension]
Point on the manifold equal to the Riemannian exponential
of tangent_vec at the base point.
"""
if base_point is None:
base_point = [None, ] * self.n_metrics
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)
base_point = gs.to_ndarray(base_point, to_ndim=2)
intrinsic = self._is_intrinsic(base_point)
args = {'tangent_vec': tangent_vec, 'base_point': base_point}
exp = self._iterate_over_metrics('exp', args, intrinsic)
return gs.hstack(exp)
elif point_type == 'matrix':
tangent_vec = gs.to_ndarray(tangent_vec, to_ndim=3)
base_point = gs.to_ndarray(base_point, to_ndim=3)
return gs.stack([
self.metrics[i].exp(tangent_vec[:, i], base_point[:, i])
for i in range(self.n_metrics)], axis=1)
else:
raise ValueError('invalid point_type argument: {}, expected '
'either matrix of vector'.format(point_type))
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 landmarks_on_geodesic(t):
t = gs.cast(t, gs.float32)
t = gs.to_ndarray(t, to_ndim=1)
t = gs.to_ndarray(t, to_ndim=2, axis=1)
new_initial_landmarks = gs.to_ndarray(initial_landmarks,
to_ndim=landmarks_ndim + 1)
new_initial_tangent_vec = gs.to_ndarray(initial_tangent_vec,
to_ndim=landmarks_ndim + 1)
tangent_vecs = gs.einsum('il,nkm->ikm', t, new_initial_tangent_vec)
def point_on_landmarks(tangent_vec):
assert gs.ndim(tangent_vec) >= 2
exp = self.exp(tangent_vec=tangent_vec,
base_landmarks=new_initial_landmarks)
return exp
landmarks_at_time_t = gs.vectorize(tangent_vecs,
point_on_landmarks,
signature='(i,j)->(i,j)')
return landmarks_at_time_t
def inverse(self, mat):
"""Compute matrix inverse.
Parameters
----------
mat
Returns
-------
inverse
"""
mat = gs.to_ndarray(mat, to_ndim=3)
return gs.linalg.inv(mat)
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 on the hypersphere.
Returns
-------
exp : array-like, shape=[..., dim + 1]
Point on the hypersphere equal to the Riemannian exponential
of tangent_vec at the base point.
"""
_, extrinsic_dim = base_point.shape
n_tangent_vecs, _ = tangent_vec.shape
hypersphere = Hypersphere(dim=extrinsic_dim - 1)
proj_tangent_vec = hypersphere.to_tangent(tangent_vec, base_point)
norm_tangent_vec = self.embedding_metric.norm(proj_tangent_vec)
norm_tangent_vec = gs.to_ndarray(norm_tangent_vec, to_ndim=1)
mask_0 = gs.isclose(norm_tangent_vec, 0.)
mask_non0 = ~mask_0
coef_1 = gs.zeros((n_tangent_vecs, ))
coef_2 = gs.zeros((n_tangent_vecs, ))
norm2 = norm_tangent_vec[mask_0]**2
norm4 = norm2**2
norm6 = norm2**3
coef_1 = gs.assignment(coef_1,
1. - norm2 / 2. + norm4 / 24. - norm6 / 720.,
mask_0)
coef_2 = gs.assignment(coef_2,
1. - norm2 / 6. + norm4 / 120. - norm6 / 5040.,
mask_0)
coef_1 = gs.assignment(coef_1, gs.cos(norm_tangent_vec[mask_non0]),
mask_non0)
coef_2 = gs.assignment(
coef_2,
gs.sin(norm_tangent_vec[mask_non0]) / norm_tangent_vec[mask_non0],
mask_non0)
exp = (gs.einsum('...,...j->...j', coef_1, base_point) +
gs.einsum('...,...j->...j', coef_2, proj_tangent_vec))
return exp
