-
Notifications
You must be signed in to change notification settings - Fork 239
/
hyperboloid.py
451 lines (371 loc) · 15.5 KB
/
hyperboloid.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
"""The n-dimensional hyperbolic space.
The n-dimensional hyperbolic space embedded with
the hyperboloid representation (embedded in minkowsky space).
"""
import math
import geomstats.algebra_utils as utils
import geomstats.backend as gs
import geomstats.vectorization
from geomstats.geometry._hyperbolic import HyperbolicMetric, _Hyperbolic
from geomstats.geometry.base import LevelSet
from geomstats.geometry.minkowski import Minkowski, MinkowskiMetric
class Hyperboloid(_Hyperbolic, LevelSet):
"""Class for the n-dimensional hyperboloid space.
Class for the n-dimensional hyperboloid space as embedded in (
n+1)-dimensional Minkowski space. For other representations of
hyperbolic spaces see the `Hyperbolic` class.
The coords_type parameter allows to choose the
representation of the points as input.
Parameters
----------
dim : int
Dimension of the hyperbolic space.
coords_type : str, {'extrinsic', 'intrinsic'}
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_coords_type = "extrinsic"
default_point_type = "vector"
def __init__(self, dim, coords_type="extrinsic", scale=1):
minkowski = Minkowski(dim + 1)
super(Hyperboloid, self).__init__(
dim=dim,
embedding_space=minkowski,
submersion=minkowski.metric.squared_norm,
value=-1.0,
tangent_submersion=minkowski.metric.inner_product,
scale=scale,
)
self.coords_type = coords_type
self.point_type = Hyperboloid.default_point_type
self.metric = HyperboloidMetric(self.dim, self.coords_type, self.scale)
def belongs(self, point, atol=gs.atol):
"""Test if a point belongs to the hyperbolic space.
Test if a point belongs to the hyperbolic space in
its hyperboloid representation.
Parameters
----------
point : array-like, shape=[..., dim]
Point to be tested.
atol : float, optional
Tolerance at which to evaluate how close the squared norm
is to the reference value.
Optional, default: backend atol.
Returns
-------
belongs : array-like, shape=[...,]
Array of booleans indicating whether the corresponding points
belong to the hyperbolic space.
"""
point_dim = point.shape[-1]
if point_dim is not self.dim + 1:
belongs = False
if point_dim is self.dim and self.coords_type == "intrinsic":
belongs = True
if gs.ndim(point) == 2:
belongs = gs.tile([belongs], (point.shape[0],))
return belongs
return super(Hyperboloid, self).belongs(point, atol)
def projection(self, point):
"""Project a point in space on the hyperboloid.
Parameters
----------
point : array-like, shape=[..., dim + 1]
Point in embedding Euclidean space.
Returns
-------
projected_point : array-like, shape=[..., dim + 1]
Point projected on the hyperboloid.
"""
belongs = self.belongs(point)
# avoid dividing by 0
factor = gs.where(point[..., 0] == 0.0, 1.0, point[..., 0] + gs.atol)
first_coord = gs.where(belongs, 1.0, 1.0 / factor)
intrinsic = gs.einsum("...,...i->...i", first_coord, point)[..., 1:]
return self.intrinsic_to_extrinsic_coords(intrinsic)
def regularize(self, point):
"""Regularize a point to the canonical representation.
Regularize a point to the canonical representation chosen
for the hyperbolic space, to avoid numerical issues.
Parameters
----------
point : array-like, shape=[..., dim + 1]
Point.
Returns
-------
projected_point : array-like, shape=[..., dim + 1]
Point in hyperbolic space in canonical representation
in extrinsic coordinates.
"""
if self.coords_type == "intrinsic":
point = self.intrinsic_to_extrinsic_coords(point)
sq_norm = self.embedding_metric.squared_norm(point)
if not gs.all(sq_norm):
raise ValueError(
"Cannot project a vector of norm 0. in the "
"Minkowski space to the hyperboloid"
)
real_norm = gs.sqrt(gs.abs(sq_norm))
projected_point = gs.einsum("...i,...->...i", point, 1.0 / real_norm)
return projected_point
@geomstats.vectorization.decorator(["else", "vector", "vector"])
def to_tangent(self, vector, base_point):
"""Project a vector to a tangent space of the hyperbolic space.
Project a vector in Minkowski space on the tangent space
of the hyperbolic space at a base point.
Parameters
----------
vector : array-like, shape=[..., dim + 1]
Vector in Minkowski space to be projected.
base_point : array-like, shape=[..., dim + 1]
Point in hyperbolic space.
Returns
-------
tangent_vec : array-like, shape=[..., dim + 1]
Tangent vector at the base point, equal to the projection of
the vector in Minkowski space.
"""
if self.coords_type == "intrinsic":
base_point = self.intrinsic_to_extrinsic_coords(base_point)
sq_norm = self.embedding_metric.squared_norm(base_point)
inner_prod = self.embedding_metric.inner_product(base_point, vector)
coef = inner_prod / sq_norm
tangent_vec = vector - gs.einsum("...,...j->...j", coef, base_point)
return tangent_vec
def is_tangent(self, vector, base_point, atol=gs.atol):
"""Check whether the vector is tangent at base_point.
Parameters
----------
vector : array-like, shape=[..., dim + 1]
Vector.
base_point : array-like, shape=[..., dim + 1]
Point on the manifold.
atol : float
Absolute tolerance.
Optional, default: backend atol.
Returns
-------
is_tangent : bool
Boolean denoting if vector is a tangent vector at the base point.
"""
product = self.embedding_metric.inner_product(vector, base_point)
return gs.isclose(product, 0.0)
def intrinsic_to_extrinsic_coords(self, point_intrinsic):
"""Convert from intrinsic to extrinsic coordinates.
Parameters
----------
point_intrinsic : array-like, shape=[..., dim]
Point in the embedded manifold in intrinsic coordinates.
Returns
-------
point_extrinsic : array-like, shape=[..., dim + 1]
Point in the embedded manifold in extrinsic coordinates.
"""
if self.dim != point_intrinsic.shape[-1]:
raise NameError(
"Wrong intrinsic dimension: "
+ str(point_intrinsic.shape[-1])
+ " instead of "
+ str(self.dim)
)
return _Hyperbolic.change_coordinates_system(
point_intrinsic, "intrinsic", "extrinsic"
)
def extrinsic_to_intrinsic_coords(self, point_extrinsic):
"""Convert from extrinsic to intrinsic coordinates.
Parameters
----------
point_extrinsic : array-like, shape=[..., dim + 1]
Point in the embedded manifold in extrinsic coordinates,
i. e. in the coordinates of the embedding manifold.
Returns
-------
point_intrinsic : array-like, shape=[..., dim]
Point in intrinsic coordinates.
"""
belong_point = self.belongs(point_extrinsic)
if not gs.all(belong_point):
raise NameError("Point that does not belong to the hyperboloid " "found")
return _Hyperbolic.change_coordinates_system(
point_extrinsic, "extrinsic", "intrinsic"
)
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_a, tangent_vec_b, base_point):
"""Compute the parallel transport of a tangent vector.
Closed-form solution for the parallel transport of a tangent vector a
along the geodesic defined by exp_(base_point)(tangent_vec_b).
Parameters
----------
tangent_vec_a : array-like, shape=[..., dim + 1]
Tangent vector at base point to be transported.
tangent_vec_b : array-like, shape=[..., dim + 1]
Tangent vector at base point, along which the parallel transport
is computed.
base_point : array-like, shape=[..., dim + 1]
Point on the hyperboloid.
Returns
-------
transported_tangent_vec: array-like, shape=[..., dim + 1]
Transported tangent vector at exp_(base_point)(tangent_vec_b).
"""
theta = self.embedding_metric.norm(tangent_vec_b)
eps = gs.where(theta == 0.0, 1.0, theta)
normalized_b = gs.einsum("...,...i->...i", 1 / eps, tangent_vec_b)
pb = self.embedding_metric.inner_product(tangent_vec_a, normalized_b)
p_orth = tangent_vec_a - 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