class HyperbolicMetric(RiemannianMetric):
"""Class that defines operations using a hyperbolic metric.
Parameters
----------
dimension : int
Dimension of the hyperbolic space.
point_type : str, {'extrinsic', 'intrinsic', etc}, optional
Default coordinates to represent points in hyperbolic space.
scale : int, optional
Scale of the hyperbolic space, defined as the set of points
in Minkowski space whose squared norm is equal to -scale.
"""
def __init__(self, dimension, point_type='extrinsic', scale=1):
super(HyperbolicMetric, self).__init__(dimension=dimension,
signature=(dimension, 0, 0))
self.embedding_metric = MinkowskiMetric(dimension + 1)
self.point_type = point_type
assert scale > 0, 'The scale should be strictly positive'
self.scale = scale
def inner_product(self, tangent_vec_a, tangent_vec_b, base_point=None):
"""Compute the inner-product of two tangent vectors at a base point.
Parameters
----------
tangent_vec_a : array-like, shape=[n_samples, dimension + 1]
First tangent vector at base point.
tangent_vec_b : array-like, shape=[n_samples, dimension + 1]
Second tangent vector at base point.
base_point : array-like, shape=[n_samples, dimension + 1], optional
Point in hyperbolic space.
Returns
-------
inner_prod : array-like, shape=[n_samples, 1]
Inner-product of the two tangent vectors.
"""
inner_prod = self.embedding_metric.inner_product(
tangent_vec_a, tangent_vec_b, base_point)
inner_prod *= self.scale**2
return inner_prod
def squared_norm(self, vector, base_point=None):
"""Compute the squared norm of a vector.
Squared norm of a vector associated with the inner-product
at the tangent space at a base point.
Parameters
----------
vector : array-like, shape=[n_samples, dimension + 1]
Vector on the tangent space of the hyperbolic space at base point.
base_point : array-like, shape=[n_samples, dimension + 1], optional
Point in hyperbolic space in extrinsic coordinates.
Returns
-------
sq_norm : array-like, shape=[n_samples, 1]
Squared norm of the vector.
"""
sq_norm = self.embedding_metric.squared_norm(vector)
sq_norm *= self.scale**2
return sq_norm
def exp(self, tangent_vec, base_point):
"""Compute the Riemannian exponential of a tangent vector.
Parameters
----------
tangent_vec : array-like, shape=[n_samples, dimension + 1]
Tangent vector at a base point.
base_point : array-like, shape=[n_samples, dimension + 1]
Point in hyperbolic space.
Returns
-------
exp : array-like, shape=[n_samples, dimension + 1]
Point in hyperbolic space equal to the Riemannian exponential
of tangent_vec at the base point.
"""
if self.point_type == 'extrinsic':
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)
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('ni,nj->nj', coef_1, base_point) +
gs.einsum('ni,nj->nj', coef_2, tangent_vec))
hyperbolic_space = Hyperbolic(dimension=self.dimension)
exp = hyperbolic_space.regularize(exp)
return exp
elif self.point_type == 'ball':
norm_base_point = gs.to_ndarray(gs.linalg.norm(base_point, -1), 2,
-1)
norm_base_point = gs.repeat(norm_base_point, base_point.shape[-1],
-1)
den = 1 - norm_base_point**2
norm_tan = gs.to_ndarray(gs.linalg.norm(tangent_vec, axis=-1), 2,
-1)
norm_tan = gs.repeat(norm_tan, base_point.shape[-1], -1)
lambda_base_point = 1 / den
direction = tangent_vec / norm_tan
factor = gs.tanh(lambda_base_point * norm_tan)
exp = self.mobius_add(base_point, direction * factor)
return exp
else:
raise NotImplementedError(
'exp is only implemented for ball and extrinsic')
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=[n_samples, dimension + 1]
Point in hyperbolic space.
base_point : array-like, shape=[n_samples, dimension + 1]
Point in hyperbolic space.
Returns
-------
log : array-like, shape=[n_samples, dimension + 1]
Tangent vector at the base point equal to the Riemannian logarithm
of point at the base point.
"""
if self.point_type == 'extrinsic':
point = gs.to_ndarray(point, to_ndim=2)
base_point = gs.to_ndarray(base_point, to_ndim=2)
angle = self.dist(base_point, point) / self.scale
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
elif self.point_type == 'ball':
add_base_point = self.mobius_add(-base_point, point)
norm_add = gs.to_ndarray(gs.linalg.norm(add_base_point, axis=-1),
2, -1)
norm_add = gs.repeat(norm_add, base_point.shape[-1], -1)
norm_base_point = gs.to_ndarray(
gs.linalg.norm(base_point, axis=-1), 2, -1)
norm_base_point = gs.repeat(norm_base_point, base_point.shape[-1],
-1)
log = (1 - norm_base_point**2) * gs.arctanh(norm_add)\
* (add_base_point / norm_add)
mask_0 = gs.all(gs.isclose(norm_add, 0.))
log[mask_0] = 0
return log
else:
raise NotImplementedError(
'log is only implemented for ball and extrinsic')
def mobius_add(self, point_a, point_b):
r"""Compute the Mobius addition of two points.
Mobius addition operation that is a necessary operation
to compute the log and exp using the 'ball' representation.
.. math::
a\oplus b=\frac{(1+2\langle a,b\rangle + ||b||^2)a+
(1-||a||^2)b}{1+2\langle a,b\rangle + ||a||^2||b||^2}
Parameters
----------
point_a : array-like, shape=[n_samples, dimension + 1]
Point in hyperbolic space.
point_b : array-like, shape=[n_samples, dimension + 1]
Point in hyperbolic space.
Returns
-------
mobius_add : array-like, shape=[n_samples, 1]
Result of the Mobius addition.
"""
norm_point_a = gs.sum(point_a**2, axis=-1, keepdims=True)
# to redefine to use autograd
norm_point_a = gs.repeat(norm_point_a, point_a.shape[-1], -1)
norm_point_b = gs.sum(point_b**2, axis=-1, keepdims=True)
norm_point_b = gs.repeat(norm_point_b, point_a.shape[-1], -1)
sum_prod_a_b = gs.sum(point_a * point_b, axis=-1, keepdims=True)
sum_prod_a_b = gs.repeat(sum_prod_a_b, point_a.shape[-1], -1)
add_nominator = ((1 + 2 * sum_prod_a_b + norm_point_b) * point_a +
(1 - norm_point_a) * point_b)
add_denominator = (1 + 2 * sum_prod_a_b + norm_point_a * norm_point_b)
mobius_add = add_nominator / add_denominator
return mobius_add
def dist(self, point_a, point_b):
"""Compute the geodesic distance between two points.
Parameters
----------
point_a : array-like, shape=[n_samples, dimension + 1]
First point in hyperbolic space.
point_b : array-like, shape=[n_samples, dimension + 1]
Second point in hyperbolic space.
Returns
-------
dist : array-like, shape=[n_samples, 1]
Geodesic distance between the two points.
"""
if self.point_type == 'extrinsic':
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.0, 1e24)
dist = gs.arccosh(cosh_angle)
dist *= self.scale
return dist
elif self.point_type == 'ball':
point_a_norm = gs.clip(gs.sum(point_a**2, -1), 0., 1 - EPSILON)
point_b_norm = gs.clip(gs.sum(point_b**2, -1), 0., 1 - EPSILON)
diff_norm = gs.sum((point_a - point_b)**2, -1)
norm_function = 1 + 2 * \
diff_norm / ((1 - point_a_norm) * (1 - point_b_norm))
dist = gs.log(norm_function + gs.sqrt(norm_function**2 - 1))
dist = gs.to_ndarray(dist, to_ndim=1)
dist = gs.to_ndarray(dist, to_ndim=2, axis=1)
dist *= self.scale
return dist
else:
raise NotImplementedError(
'dist is only implemented for ball and extrinsic')
class HyperboloidMetric(HyperbolicMetric):
"""Class that defines operations using a hyperbolic metric.
Parameters
----------
dim : int
Dimension of the hyperbolic space.
point_type : str, {'extrinsic', 'intrinsic', etc}
Default coordinates to represent points in hyperbolic space.
Optional, default: 'extrinsic'.
scale : int
Scale of the hyperbolic space, defined as the set of points
in Minkowski space whose squared norm is equal to -scale.
Optional, default: 1.
"""
default_point_type = "vector"
default_coords_type = "extrinsic"
def __init__(self, dim, coords_type="extrinsic", scale=1):
super(HyperboloidMetric, self).__init__(dim=dim, scale=scale)
self.embedding_metric = MinkowskiMetric(dim + 1)
self.coords_type = coords_type
self.point_type = HyperbolicMetric.default_point_type
self.scale = scale
def metric_matrix(self, base_point=None):
"""Compute the inner product matrix.
Parameters
----------
base_point: array-like, shape=[..., dim + 1]
Base point.
Optional, default: None.
Returns
-------
inner_prod_mat: array-like, shape=[..., dim+1, dim + 1]
Inner-product matrix.
"""
self.embedding_metric.metric_matrix(base_point)
def _inner_product(self, tangent_vec_a, tangent_vec_b, base_point=None):
"""Compute the inner-product of two tangent vectors at a base point.
Parameters
----------
tangent_vec_a : array-like, shape=[..., dim + 1]
First tangent vector at base point.
tangent_vec_b : array-like, shape=[..., dim + 1]
Second tangent vector at base point.
base_point : array-like, shape=[..., dim + 1], optional
Point in hyperbolic space.
Returns
-------
inner_prod : array-like, shape=[...,]
Inner-product of the two tangent vectors.
"""
inner_prod = self.embedding_metric.inner_product(
tangent_vec_a, tangent_vec_b, base_point)
return inner_prod
def _squared_norm(self, vector, base_point=None):
"""Compute the squared norm of a vector.
Squared norm of a vector associated with the inner-product
at the tangent space at a base point.
Parameters
----------
vector : array-like, shape=[..., dim + 1]
Vector on the tangent space of the hyperbolic space at base point.
base_point : array-like, shape=[..., dim + 1], optional
Point in hyperbolic space in extrinsic coordinates.
Returns
-------
sq_norm : array-like, shape=[...,]
Squared norm of the vector.
"""
sq_norm = self.embedding_metric.squared_norm(vector)
return sq_norm
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)
coef_1 = utils.taylor_exp_even_func(sq_norm_tangent_vec,
utils.cosh_close_0,
order=5)
coef_2 = utils.taylor_exp_even_func(sq_norm_tangent_vec,
utils.sinch_close_0,
order=5)
exp = gs.einsum("...,...j->...j", coef_1, base_point) + gs.einsum(
"...,...j->...j", coef_2, tangent_vec)
exp = Hyperboloid(dim=self.dim).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
coef_1_ = utils.taylor_exp_even_func(angle**2,
utils.inv_sinch_close_0,
order=4)
coef_2_ = utils.taylor_exp_even_func(angle**2,
utils.inv_tanh_close_0,
order=4)
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 dist(self, point_a, point_b):
"""Compute the geodesic distance between two points.
Parameters
----------
point_a : array-like, shape=[..., dim + 1]
First point in hyperbolic space.
point_b : array-like, shape=[..., dim + 1]
Second point in hyperbolic space.
Returns
-------
dist : array-like, shape=[...,]
Geodesic distance between the two points.
"""
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.0, 1e24)
dist = gs.arccosh(cosh_angle)
dist *= self.scale
return dist
def parallel_transport(self,
tangent_vec,
base_point,
direction=None,
end_point=None):
r"""Compute the parallel transport of a tangent vector.
Closed-form solution for the parallel transport of a tangent vector
along the geodesic between two points `base_point` and `end_point`
or alternatively defined by :math:`t\mapsto exp_(base_point)(
t*direction)`.
Parameters
----------
tangent_vec : array-like, shape=[..., dim + 1]
Tangent vector at base point to be transported.
base_point : array-like, shape=[..., dim + 1]
Point on the hyperboloid.
direction : array-like, shape=[..., dim + 1]
Tangent vector at base point, along which the parallel transport
is computed.
Optional, default : None.
end_point : array-like, shape=[..., dim + 1]
Point on the hyperboloid. Point to transport to. Unused if `tangent_vec_b`
is given.
Optional, default : None.
Returns
-------
transported_tangent_vec: array-like, shape=[..., dim + 1]
Transported tangent vector at exp_(base_point)(tangent_vec_b).
"""
if direction is None:
if end_point is not None:
direction = self.log(end_point, base_point)
else:
raise ValueError(
"Either an end_point or a tangent_vec_b must be given to define the"
" geodesic along which to transport.")
theta = self.embedding_metric.norm(direction)
eps = gs.where(theta == 0.0, 1.0, theta)
normalized_b = gs.einsum("...,...i->...i", 1 / eps, direction)
pb = self.embedding_metric.inner_product(tangent_vec, normalized_b)
p_orth = tangent_vec - gs.einsum("...,...i->...i", pb, normalized_b)
transported = (gs.einsum("...,...i->...i",
gs.sinh(theta) * pb, base_point) +
gs.einsum("...,...i->...i",
gs.cosh(theta) * pb, normalized_b) + p_orth)
return transported
class HyperboloidMetric(HyperbolicMetric):
"""Class that defines operations using a hyperbolic metric.
Parameters
----------
dim : int
Dimension of the hyperbolic space.
point_type : str, {'extrinsic', 'intrinsic', etc}, optional
Default coordinates to represent points in hyperbolic space.
scale : int, optional
Scale of the hyperbolic space, defined as the set of points
in Minkowski space whose squared norm is equal to -scale.
"""
default_point_type = 'vector'
default_coords_type = 'extrinsic'
def __init__(self, dim, coords_type='extrinsic', scale=1):
super(HyperboloidMetric, self).__init__(
dim=dim,
scale=scale)
self.embedding_metric = MinkowskiMetric(dim + 1)
self.coords_type = coords_type
self.point_type = HyperbolicMetric.default_point_type
self.scale = scale
def inner_product_matrix(self, base_point=None):
"""Compute the inner product matrix.
Parameters
----------
base_point: array-like, shape=[..., dim+1]
Returns
-------
inner_prod_mat: array-like, shape=[..., dim+1, dim+1]
"""
self.embedding_metric.inner_product_matrix(base_point)
def _inner_product(self, tangent_vec_a, tangent_vec_b, base_point=None):
"""Compute the inner-product of two tangent vectors at a base point.
Parameters
----------
tangent_vec_a : array-like, shape=[..., dim + 1]
First tangent vector at base point.
tangent_vec_b : array-like, shape=[..., dim + 1]
Second tangent vector at base point.
base_point : array-like, shape=[..., dim + 1], optional
Point in hyperbolic space.
Returns
-------
inner_prod : array-like, shape=[..., 1]
Inner-product of the two tangent vectors.
"""
inner_prod = self.embedding_metric.inner_product(
tangent_vec_a, tangent_vec_b, base_point)
return inner_prod
def _squared_norm(self, vector, base_point=None):
"""Compute the squared norm of a vector.
Squared norm of a vector associated with the inner-product
at the tangent space at a base point.
Parameters
----------
vector : array-like, shape=[..., dim + 1]
Vector on the tangent space of the hyperbolic space at base point.
base_point : array-like, shape=[..., dim + 1], optional
Point in hyperbolic space in extrinsic coordinates.
Returns
-------
sq_norm : array-like, shape=[..., 1]
Squared norm of the vector.
"""
sq_norm = self.embedding_metric.squared_norm(vector)
return sq_norm
@geomstats.vectorization.decorator(['else', 'vector', 'vector'])
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
@geomstats.vectorization.decorator(['else', 'vector', 'vector'])
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 dist(self, point_a, point_b):
"""Compute the geodesic distance between two points.
Parameters
----------
point_a : array-like, shape=[..., dim + 1]
First point in hyperbolic space.
point_b : array-like, shape=[..., dim + 1]
Second point in hyperbolic space.
Returns
-------
dist : array-like, shape=[..., 1]
Geodesic distance between the two points.
"""
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.0, 1e24)
dist = gs.arccosh(cosh_angle)
dist *= self.scale
return dist
class HyperboloidMetric(HyperbolicMetric):
"""Class that defines operations using a hyperbolic metric.
Parameters
----------
dim : int
Dimension of the hyperbolic space.
point_type : str, {'extrinsic', 'intrinsic', etc}
Default coordinates to represent points in hyperbolic space.
Optional, default: 'extrinsic'.
scale : int
Scale of the hyperbolic space, defined as the set of points
in Minkowski space whose squared norm is equal to -scale.
Optional, default: 1.
"""
default_point_type = 'vector'
default_coords_type = 'extrinsic'
def __init__(self, dim, coords_type='extrinsic', scale=1):
super(HyperboloidMetric, self).__init__(dim=dim, scale=scale)
self.embedding_metric = MinkowskiMetric(dim + 1)
self.coords_type = coords_type
self.point_type = HyperbolicMetric.default_point_type
self.scale = scale
def inner_product_matrix(self, base_point=None):
"""Compute the inner product matrix.
Parameters
----------
base_point: array-like, shape=[..., dim + 1]
Base point.
Optional, default: None.
Returns
-------
inner_prod_mat: array-like, shape=[..., dim+1, dim + 1]
Inner-product matrix.
"""
self.embedding_metric.inner_product_matrix(base_point)
def _inner_product(self, tangent_vec_a, tangent_vec_b, base_point=None):
"""Compute the inner-product of two tangent vectors at a base point.
Parameters
----------
tangent_vec_a : array-like, shape=[..., dim + 1]
First tangent vector at base point.
tangent_vec_b : array-like, shape=[..., dim + 1]
Second tangent vector at base point.
base_point : array-like, shape=[..., dim + 1], optional
Point in hyperbolic space.
Returns
-------
inner_prod : array-like, shape=[...,]
Inner-product of the two tangent vectors.
"""
inner_prod = self.embedding_metric.inner_product(
tangent_vec_a, tangent_vec_b, base_point)
return inner_prod
def _squared_norm(self, vector, base_point=None):
"""Compute the squared norm of a vector.
Squared norm of a vector associated with the inner-product
at the tangent space at a base point.
Parameters
----------
vector : array-like, shape=[..., dim + 1]
Vector on the tangent space of the hyperbolic space at base point.
base_point : array-like, shape=[..., dim + 1], optional
Point in hyperbolic space in extrinsic coordinates.
Returns
-------
sq_norm : array-like, shape=[...,]
Squared norm of the vector.
"""
sq_norm = self.embedding_metric.squared_norm(vector)
return sq_norm
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)
coef_1 = utils.taylor_exp_even_func(sq_norm_tangent_vec,
utils.cosh_close_0,
order=5)
coef_2 = utils.taylor_exp_even_func(sq_norm_tangent_vec,
utils.sinch_close_0,
order=5)
exp = (gs.einsum('...,...j->...j', coef_1, base_point) +
gs.einsum('...,...j->...j', coef_2, tangent_vec))
exp = Hyperboloid(dim=self.dim).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
coef_1_ = utils.taylor_exp_even_func(angle**2,
utils.inv_sinch_close_0,
order=4)
coef_2_ = utils.taylor_exp_even_func(angle**2,
utils.inv_tanh_close_0,
order=4)
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 dist(self, point_a, point_b):
"""Compute the geodesic distance between two points.
Parameters
----------
point_a : array-like, shape=[..., dim + 1]
First point in hyperbolic space.
point_b : array-like, shape=[..., dim + 1]
Second point in hyperbolic space.
Returns
-------
dist : array-like, shape=[...,]
Geodesic distance between the two points.
"""
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.0, 1e24)
dist = gs.arccosh(cosh_angle)
dist *= self.scale
return dist
class HyperbolicMetric(RiemannianMetric):
"""Class for the Hyperbolic metric."""
def __init__(self, dimension, point_type='extrinsic', scale=1):
super(HyperbolicMetric, self).__init__(dimension=dimension,
signature=(dimension, 0, 0))
self.embedding_metric = MinkowskiMetric(dimension + 1)
self.point_type = point_type
assert scale > 0, 'The scale should be strictly positive'
self.scale = scale
def inner_product(self, tangent_vec_a, tangent_vec_b, base_point=None):
"""Compute the inner product of two tangent vectors at a base point.
Parameters
----------
tangent_vec_a : array-like, shape=[n_samples, dimension + 1]
or shape=[1, dimension + 1]
tangent_vec_b : 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
-------
inner_prod : array-like, shape=[n_samples, 1]
or shape=[1, 1]
"""
inner_prod = self.scale**2 * self.embedding_metric.inner_product(
tangent_vec_a, tangent_vec_b, base_point)
return inner_prod
def squared_norm(self, vector, base_point=None):
"""Compute the squared norm of a vector at a given base point.
Squared norm of a vector associated with the inner product
at the tangent space at a base point. Extrinsic base point only
Parameters
----------
vector : 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
-------
sq_norm : array-like, shape=[n_samples, 1]
or shape=[1, 1]
"""
sq_norm = self.scale**2 * self.embedding_metric.squared_norm(vector)
return sq_norm
def exp(self, tangent_vec, base_point):
"""Compute Riemannian exponential of tangent vector wrt to 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]
"""
if self.point_type == 'extrinsic':
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 = Hyperbolic(dimension=self.dimension)
exp = hyperbolic_space.regularize(exp)
return exp
elif self.point_type == 'ball':
norm_base_point = gs.to_ndarray(gs.norm(base_point, -1), 2, -1)
norm_base_point = gs.repeat(norm_base_point, base_point.shape[-1],
-1)
den = 1 - norm_base_point**2
norm_tan = gs.to_ndarray(gs.norm(tangent_vec, -1), 2, -1)
norm_tan = gs.repeat(norm_tan, base_point.shape[-1], -1)
lambda_base_point = 1 / den
direction = tangent_vec / norm_tan
factor = gs.tanh(lambda_base_point * norm_tan)
exp = self.mobius_add(base_point, direction * factor)
return exp
else:
raise NotImplementedError(
'exp is only implemented for ball and extrinsic')
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=[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]
"""
if self.point_type == 'extrinsic':
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
elif self.point_type == 'ball':
add_base_point = self.mobius_add(-base_point, point)
norm_add = gs.to_ndarray(gs.norm(add_base_point, -1), 2, -1)
norm_add = gs.repeat(norm_add, base_point.shape[-1], -1)
norm2_base_point = gs.to_ndarray(gs.sum(base_point**2, -1), 2, -1)
norm2_base_point = gs.repeat(norm2_base_point,
base_point.shape[-1], -1)
log = (1 - norm2_base_point) * gs.arctanh(norm_add)\
* (add_base_point / norm_add)
mask_0 = gs.all(gs.isclose(norm_add, 0))
log[mask_0] = 0
return log
else:
raise NotImplementedError(
'log is only implemented for ball and extrinsic')
def mobius_add(self, point_a, point_b):
"""Compute the mobius addition of two points.
Mobius addition is necessary for computation of the log and exp
using the 'poincare' representation set as point_type.
Parameters
----------
point_a : array-like, shape=[n_samples, dimension + 1]
or shape=[1, dimension + 1]
point_b : array-like, shape=[n_samples, dimension + 1]
or shape=[1, dimension + 1]
Returns
-------
mobius_add : array-like, shape=[n_samples, 1]
or shape=[1, 1]
"""
norm_point_a = gs.sum(point_a**2, axis=-1, keepdims=True)
# to redefine to use autograd
norm_point_a = gs.repeat(norm_point_a, point_a.shape[-1], -1)
norm_point_b = gs.sum(point_b**2, axis=-1, keepdims=True)
norm_point_b = gs.repeat(norm_point_b, point_a.shape[-1], -1)
sum_prod_a_b = (point_a * point_b).sum(-1, keepdims=True)
sum_prod_a_b = gs.repeat(sum_prod_a_b, point_a.shape[-1], -1)
add_nominator = ((1 + 2 * sum_prod_a_b + norm_point_b) * point_a +
(1 - norm_point_a) * point_b)
add_denominator = (1 + 2 * sum_prod_a_b + norm_point_a * norm_point_b)
mobius_add = add_nominator / add_denominator
return mobius_add
def dist(self, point_a, point_b):
"""Compute the geodesic distance between two points.
Parameters
----------
point_a : array-like, shape=[n_samples, dimension + 1]
or shape=[1, dimension + 1]
point_b : array-like, shape=[n_samples, dimension + 1]
or shape=[1, dimension + 1]
Returns
-------
dist : array-like, shape=[n_samples, 1]
or shape=[1, 1]
"""
if self.point_type == 'extrinsic':
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.0, 1e24)
dist = gs.arccosh(cosh_angle)
return self.scale * dist
elif self.point_type == 'ball':
point_a_norm = gs.clip(gs.sum(point_a**2, -1), 0., 1 - EPSILON)
point_b_norm = gs.clip(gs.sum(point_b**2, -1), 0., 1 - EPSILON)
diff_norm = gs.sum((point_a - point_b)**2, -1)
norm_function = 1 + 2 * \
diff_norm / ((1 - point_a_norm) * (1 - point_b_norm))
dist = gs.log(norm_function + gs.sqrt(norm_function**2 - 1))
dist = gs.to_ndarray(dist, to_ndim=1)
dist = gs.to_ndarray(dist, to_ndim=2, axis=1)
return self.scale * dist
else:
raise NotImplementedError(
'dist is only implemented for ball and extrinsic')