|
|
import math |
|
|
import torch |
|
|
|
|
|
|
|
|
def tor_expmap(x: torch.Tensor, u: torch.Tensor) -> torch.Tensor: |
|
|
return (x + u) % (2 * math.pi) |
|
|
|
|
|
def tor_logmap(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
|
|
return torch.atan2(torch.sin(y - x), torch.cos(y - x)) |
|
|
|
|
|
def tor_projx(x: torch.Tensor) -> torch.Tensor: |
|
|
return x % (2 * math.pi) |
|
|
|
|
|
def tor_random_uniform(*size, dtype=None, device=None) -> torch.Tensor: |
|
|
z = torch.rand(*size, dtype=dtype, device=device) |
|
|
return z * 2 * math.pi |
|
|
|
|
|
def tor_uniform_logprob(x): |
|
|
dim = x.shape[-1] |
|
|
return torch.full_like(x[..., 0], -dim * math.log(2 * math.pi)) |
|
|
|
|
|
def tor_geodesic_t(t, angles_1, angles_0): |
|
|
|
|
|
tangent_vec = t * tor_logmap(angles_0, angles_1) |
|
|
points_at_time_t = tor_expmap(angles_0, tangent_vec) |
|
|
return points_at_time_t |
|
|
|
|
|
if __name__ =='__main__': |
|
|
a = tor_random_uniform((2,3,5)) |
|
|
b = tor_random_uniform((2,3,5)) |
|
|
t = torch.ones((2,1)) * 0.2 |
|
|
c = tor_geodesic_t(t[...,None],a,b) |
|
|
print(c) |
|
|
print(c.shape) |
