add cool plots

docs/plots/extended/poincare/distance.py

@@ -0,0 +1,27 @@
+import geoopt.manifolds.poincare.math as pmath
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+sns.set_style("white")
+radius = 1
+coords = np.linspace(-radius, radius, 100)
+x = torch.tensor([-0.75, 0])
+xx, yy = np.meshgrid(coords, coords)
+dist2 = xx ** 2 + yy ** 2
+mask = dist2 <= radius ** 2
+grid = np.stack([xx, yy], axis=-1)
+dists = pmath.dist(torch.from_numpy(grid).float(), x)
+dists[(~mask).nonzero()] = np.nan
+circle = plt.Circle((0, 0), 1, fill=False, color="b")
+plt.gca().add_artist(circle)
+plt.xlim(-1.1, 1.1)
+plt.ylim(-1.1, 1.1)
+plt.gca().set_aspect("equal")
+plt.contourf(
+    grid[..., 0], grid[..., 1], dists.log().numpy(), levels=100, cmap="inferno"
+)
+plt.colorbar()
+plt.title("log distance to ($-$0.75, 0)")
+plt.show()

docs/plots/extended/poincare/distance2plane.py

@@ -0,0 +1,31 @@
+import geoopt.manifolds.poincare.math as pmath
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+sns.set_style("white")
+radius = 1
+coords = np.linspace(-radius, radius, 100)
+x = torch.tensor([-0.75, 0])
+v = torch.tensor([0.1 / 3, -1 / 3])
+xx, yy = np.meshgrid(coords, coords)
+dist2 = xx ** 2 + yy ** 2
+mask = dist2 <= radius ** 2
+grid = np.stack([xx, yy], axis=-1)
+dists = pmath.dist2plane(torch.from_numpy(grid).float(), v, x)
+dists[(~mask).nonzero()] = np.nan
+circle = plt.Circle((0, 0), 1, fill=False, color="b")
+plt.gca().add_artist(circle)
+plt.xlim(-1.1, 1.1)
+plt.ylim(-1.1, 1.1)
+
+plt.gca().set_aspect("equal")
+plt.contourf(
+    grid[..., 0], grid[..., 1], dists.log().numpy(), levels=100, cmap="inferno"
+)
+plt.colorbar()
+plt.scatter(*x, color="g")
+plt.arrow(*x, *v, color="g", width=0.01)
+plt.title(r"log distance to $\tilde{H}_{a, p}$")
+plt.show()

docs/plots/extended/poincare/gyrovector_parallel_transport.py

@@ -0,0 +1,48 @@
+import geoopt.manifolds.poincare.math as pmath
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+sns.set_style("white")
+
+x = torch.tensor((-0.25, -0.75))
+xv1 = torch.tensor((np.sin(np.pi / 3), np.cos(np.pi / 3))) / 5
+xv2 = torch.tensor((np.sin(-np.pi / 3), np.cos(np.pi / 3))) / 5
+t = torch.linspace(0, 1, 10)[:, None]
+
+y = torch.tensor((0.65, -0.55))
+xy = pmath.logmap(x, y)
+path = pmath.geodesic(t, x, y)
+yv1 = pmath.parallel_transport(x, y, xv1)
+yv2 = pmath.parallel_transport(x, y, xv2)
+
+xgv1 = pmath.geodesic_unit(t, x, xv1)
+xgv2 = pmath.geodesic_unit(t, x, xv2)
+
+ygv1 = pmath.geodesic_unit(t, y, yv1)
+ygv2 = pmath.geodesic_unit(t, y, yv2)
+
+
+def plot_gv(gv, **kwargs):
+    plt.plot(*gv.t().numpy(), **kwargs)
+    plt.arrow(*gv[-2], *(gv[-1] - gv[-2]), width=0.01, **kwargs)
+
+
+circle = plt.Circle((0, 0), 1, fill=False, color="b")
+plt.gca().add_artist(circle)
+plt.xlim(-1.1, 1.1)
+plt.ylim(-1.1, 1.1)
+plt.gca().set_aspect("equal")
+plt.annotate("x", x - 0.09, fontsize=15)
+plt.annotate("y", y - 0.09, fontsize=15)
+plt.annotate(r"$\vec{v}$", x + torch.tensor([0.3, 0.5]), fontsize=15)
+plot_gv(xgv1, color="r")
+plot_gv(xgv2, color="b")
+plt.arrow(*x, *xy, width=0.01, color="g")
+plot_gv(ygv1, color="r")
+plot_gv(ygv2, color="b")
+
+plt.plot(*path.t().numpy(), color="g")
+plt.title(r"gyrovector parallel transport $P_{x\to y}$")
+plt.show()

docs/plots/extended/poincare/parallel_transport.py

@@ -0,0 +1,35 @@
+import geoopt.manifolds.poincare.math as pmath
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+sns.set_style("white")
+
+x = torch.tensor((-0.25, -0.75))
+v1 = torch.tensor((np.sin(np.pi / 3), np.cos(np.pi / 3))) / 5
+v2 = torch.tensor((np.sin(-np.pi / 3), np.cos(np.pi / 3))) / 5
+y = torch.tensor((0.65, -0.55))
+t = torch.linspace(0, 1)
+xy = pmath.logmap(x, y)
+path = pmath.geodesic(t[:, None], x, y)
+yv1 = pmath.parallel_transport(x, y, v1)
+yv2 = pmath.parallel_transport(x, y, v2)
+
+
+circle = plt.Circle((0, 0), 1, fill=False, color="b")
+plt.gca().add_artist(circle)
+plt.xlim(-1.1, 1.1)
+plt.ylim(-1.1, 1.1)
+plt.gca().set_aspect("equal")
+plt.annotate("x", x - 0.07, fontsize=15)
+plt.annotate("y", y - 0.07, fontsize=15)
+plt.annotate(r"$\vec{v}$", x + torch.tensor([0.3, 0.5]), fontsize=15)
+plt.arrow(*x, *v1, width=0.01, color="r")
+plt.arrow(*x, *xy, width=0.01, color="g")
+plt.arrow(*x, *v2, width=0.01, color="b")
+plt.arrow(*y, *yv1, width=0.01, color="r")
+plt.arrow(*y, *yv2, width=0.01, color="b")
+plt.plot(*path.t().numpy(), color="g")
+plt.title(r"parallel transport $P^c_{x\to y}$")
+plt.show()

geoopt/manifolds/poincare/math.py

@@ -336,6 +336,8 @@ def dist(x, y, *, c=1.0, keepdim=False):
 
         d_c(x, y) = \frac{2}{\sqrt{c}}\tanh^{-1}(\sqrt{c}\|(-x)\oplus_c y\|_2)
 
+    .. plot:: plots/extended/poincare/distance.py
+
     Parameters
     ----------
     x : tensor
@@ -840,6 +842,8 @@ def dist2plane(x, a, p, *, c=1.0, keepdim=False):
     Distance from :math:`x` to a hyperbolic hyperplane in Poincare ball
     that is orthogonal to :math:`a` and contains :math:`p`.
 
+    .. plot:: plots/extended/poincare/distance2plane.py
+
     To form an intuition what is a hyperbolic hyperplane, let's first consider Euclidean hyperplane
 
     .. math::
@@ -1004,6 +1008,8 @@ def parallel_transport(x, y, v, *, c=1.0):
     Parallel transport is essential for adaptive algorithms in Riemannian manifolds.
     For Hyperbolic spaces parallel transport is expressed via gyration.
 
+    .. plot:: plots/extended/poincare/gyrovector_parallel_transport.py
+
     To recover parallel transport we first need to study isomorphism between gyrovectors and vectors.
     The reason is that originally, parallel transport is well defined for gyrovectors as
 
@@ -1029,6 +1035,8 @@ def parallel_transport(x, y, v, *, c=1.0):
         P^c_{x\to y}(v) = (U^c_y)^{-1}\left(\operatorname{gyr}[y, -x] U^c_x(v)\right)\\
         = \operatorname{gyr}[y, -x] v \lambda^c_x / \lambda^c_y
 
+    .. plot:: plots/extended/poincare/parallel_transport.py
+
 
     Parameters
     ----------

requirements-dev.txt

@@ -7,3 +7,4 @@ pymanopt
 twine
 wheel
 sphinx
+seaborn