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 __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
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} 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 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 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')