def test_h2_to_sspd2(seed, n, d):
spd = SPD(2)
lorentz = Lorentz(3)
x = lorentz.rand(n, ir=1.0).mul_(d)
y = h2_to_sspd2(x)
assert_allclose(sspd2_to_h2(y), x / d, atol=1e-4)
hyp_dists = sspd2_hyp_radius_ * lorentz.pdist(x / d)
assert_allclose(spd.pdist(y), hyp_dists, atol=1e-4)
def test_hyp_sph_mapping(seed, n):
hyp = Lorentz(3)
sph = Sphere(3)
x = hyp.rand(n, ir=1e-2, out=torch.empty(n, 3, dtype=torch.float64))
hyp_dists = hyp.dist(hyp.zero(n, out=x.new()), x)
y = hyperboloid_to_sphere(x)
sph_dists = sph.dist(sph.zero(n, out=x.new()), y)
assert_allclose(sph_dists, hyp_dists, atol=1e-4)
def test_sspd2_to_h2(seed, n):
spd = SPD(2)
lorentz = Lorentz(3)
x = spd.rand(n, ir=1.0)
x.div_(x.det().sqrt_().reshape(-1, 1, 1)) # unit determinant
assert_allclose(x.det(), torch.ones(n), atol=1e-4)
assert_allclose(x, spd.projx(x), atol=1e-4)
y = sspd2_to_h2(x)
hyp_dists = sspd2_hyp_radius_ * lorentz.pdist(y)
assert_allclose(spd.pdist(x), hyp_dists, atol=1e-4)
def _hyp(self, x, writer, epoch):
if self.manifold.dim == 2:
y = Lorentz.to_poincare_ball(x.detach())
self._add_2d_embedding(y.cpu().detach().numpy(),
writer,
epoch,
circle=True)
elif self.manifold.dim == 3:
y = Lorentz.to_poincare_ball(x.detach())
self._add_3d_embedding(y.cpu().detach().numpy(),
writer,
epoch,
sphere=True)
def build_manifold(*names):
from graphembed.manifolds import (Euclidean, Grassmann, Lorentz,
SymmetricPositiveDefinite,
SpecialOrthogonalGroup, Sphere)
factors = []
for name in names:
parts = name.split('_')
identifier = parts[0]
if identifier in ['euc', 'sph', 'hyp', 'so', 'spd', 'spdstein']:
n = int(parts[1])
elif identifier in ['grass']:
n1 = int(parts[1])
n2 = int(parts[2])
else:
raise ValueError(f'Unkown manifold identifier {identifier}')
if identifier == 'euc':
man = Euclidean(n)
elif identifier == 'sph':
man = Sphere(n)
elif identifier == 'hyp':
man = Lorentz(n)
elif identifier == 'so':
man = SpecialOrthogonalGroup(n)
elif identifier == 'spd':
man = SymmetricPositiveDefinite(n)
elif identifier == 'spdstein':
man = SymmetricPositiveDefinite(n, use_stein_div=True)
elif identifier == 'grass':
man = Grassmann(n1, n2)
factors.append(man)
return factors
def test_sspd2_to_h2_nonconst_factor(seed, n, d):
spd = SPD(2)
lorentz = Lorentz(3)
x = spd.rand(n, ir=1.0)
x.div_(x.det().sqrt_().reshape(-1, 1, 1)) # unit determinant
x.mul_(d) # d**2 determinant
dets = torch.empty(n).fill_(d**2)
assert_allclose(x.det(), dets, atol=1e-4)
assert_allclose(x, spd.projx(x), atol=1e-4)
y = sspd2_to_h2(x)
hyp_dists = sspd2_hyp_radius_ * lorentz.pdist(y)
# The determinant essentially does not affect the curvatures, they are all
# isometric to the 2-dimensional hyperbolic space of -1/2 constant sectional
# curvature.
assert_allclose(spd.pdist(x), hyp_dists, atol=1e-4)
def main():
args = parse_args()
# Fix the random seeds.
set_seeds(args.random_seed)
# Default torch settings.
torch.set_default_dtype(torch.float64)
if torch.cuda.is_available():
torch.set_default_tensor_type(torch.cuda.DoubleTensor)
# load data
gpdists, g = load_graph_pdists(args.input_graph,
cache_dir='.cached_pdists')
n_nodes = g.number_of_nodes()
ds = GraphDataset(gpdists)
fp = FastPrecision(g)
# run hyp2
hyp = Lorentz(3)
emb = ManifoldEmbedding(n_nodes, [hyp] * args.n_factors)
for i in range(args.n_factors):
emb.scales[i] = torch.nn.Parameter(torch.tensor(2.0))
man_name = '_'.join('hyp2' for _ in range(args.n_factors))
save_dir = os.path.join(args.save_dir, man_name)
if args.hyp_snapshot or args.hyp_pretrained:
logging.info('Loading embedding for %s', man_name)
load_embedding(emb, save_dir)
if not args.hyp_pretrained:
train(ds, fp, emb, args.n_epochs, save_dir)
# map it to SPD
spd = SPD(2 * args.n_factors)
spd_emb = ManifoldEmbedding(n_nodes, [spd])
save_dir = os.path.join(args.save_dir, 'spd{}'.format(spd.dim))
if args.spd_snapshot:
logging.info('Loading embedding for SPD%d', spd.dim)
load_embedding(spd_emb, save_dir)
else:
with torch.no_grad():
spd_emb.xs[0] = ManifoldParameter(block_diag([
h2_to_sspd2(emb.xs[i].mul(math.sqrt(2)))
for i in range(args.n_factors)
]),
manifold=spd)
hyp_dists = emb.to('cpu').compute_dists(None)
spd_dists = spd_emb.compute_dists(None).to('cpu')
assert torch.allclose(hyp_dists, spd_dists, atol=1e-4)
# run spd2
train(ds, fp, spd_emb, args.n_epochs, save_dir, args.n_epochs)
def main():
args = parse_args()
# Fix the random seeds.
set_seeds(args.random_seed)
# Default torch settings.
torch.set_default_dtype(torch.float64)
if torch.cuda.is_available():
torch.set_default_tensor_type(torch.cuda.DoubleTensor)
# load data
gpdists, g = load_graph_pdists(args.input_graph,
cache_dir='.cached_pdists')
n_nodes = g.number_of_nodes()
ds = GraphDataset(gpdists)
fp = FastPrecision(g)
# run hyp2
emb = ManifoldEmbedding(n_nodes, [Lorentz(3)])
path = os.path.join(args.save_dir, 'hyp2')
train(ds, fp, emb, args.n_epochs, path)
curvature_sq = 1 / emb.scales[0]
# map it to SSPD
sspd_emb = ManifoldEmbedding(n_nodes, [SPD(2)])
sspd_emb.xs[0] = ManifoldParameter(h2_to_sspd2(emb.xs[0] /
curvature_sq.sqrt()),
manifold=sspd_emb.manifolds[0])
sspd_emb.scales[0] = torch.nn.Parameter(1 / curvature_sq / 2)
assert torch.allclose(emb.compute_dists(None),
sspd_emb.compute_dists(None),
atol=1e-4)
# run spd2
path = os.path.join(args.save_dir, 'spd2')
train(ds, fp, sspd_emb, args.n_epochs, path, args.n_epochs)