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