def inverse(self, z_hyper, edge_index=None):
z = inverse_exp_map_mu0(z_hyper, self.radius)
z_mu0 = z[..., 1:]
log_det_J, x = z_mu0.new_zeros(z_mu0.shape[0]), z_mu0
log_det_J = logmap_logdet(z, self.radius)
preclamp_norm_list = []
for i in range(0, self.n_blocks):
x_ = x * self.mask[i]
if self.layer_type != 'Linear':
s = self.s[i](x_, edge_index)
t_out = self.t[i](x_, edge_index)
else:
s = self.s[i](x_)
t_out = self.t[i](x_)
t_proj = proj_vec(t_out, self.radius)
t1, t_rest = t_proj[:, 0].unsqueeze(1), t_proj[:, 1:]
t = self.create_masked_t((1 - self.mask[i]), t1, t_rest)
# (1-b) \odot \tilde{x} \odot exp(s(b \odot \tilde{x}))
x_pt_arg = expand_proj_dims((1 - self.mask[i]) * x * torch.exp(s))
# (1-b) \odot \textnormal{PT}_{\textbf{o}\to t(b \odot \tilde{x})
pt = parallel_transport_mu0(x_pt_arg, dst=t, radius=self.radius)
preclamp_norm = pt.max()
pt = clamp(pt, min=-max_clamp_norm, max=max_clamp_norm)
if pt.max() == max_clamp_norm:
preclamp_norm_list.append(preclamp_norm)
x_t = exp_map(x=pt, at_point=t, radius=self.radius)
log_det_J += _logdet(pt, self.radius, subdim=(self.mask[i]).sum())
preclamp_norm = x_t.max()
x_t = clamp(x_t, min=-max_clamp_norm, max=max_clamp_norm)
if x_t.max() == max_clamp_norm:
preclamp_norm_list.append(preclamp_norm)
#\log_{\textbf{o}}(\textnormal{exp}_{t()}(\textnormal{PT}_{\textbf{o}\to t()))
x_0_full = inverse_exp_map_mu0(x_t, self.radius)
x_0 = x_0_full[..., 1:]
log_det_J += logmap_logdet(x_0_full,
self.radius,
subdim=(self.mask[i]).sum())
x = x_ + (1 - self.mask[i]) * x_0
log_det_J += ((1 - self.mask[i]) * s).sum(dim=1) # log det dx/du
preclamp_norm = x.max()
x = clamp(x, min=-max_clamp_norm, max=max_clamp_norm)
if x.max() == max_clamp_norm:
preclamp_norm_list.append(preclamp_norm)
x_mu0 = expand_proj_dims(x)
# Project back to Manifold
x = exp_map_mu0(x_mu0, self.radius)
log_det_J += _logdet(x_mu0, self.radius)
self.preclamp_norm = torch.Tensor([
sum(preclamp_norm_list) / len(preclamp_norm_list)
]) if preclamp_norm_list else self.preclamp_norm
return x, log_det_J
def forward(self, x_hyper):
x = inverse_exp_map_mu0(x_hyper, self.radius)
x_mu0 = x[..., 1:]
log_det_J, z = x.new_zeros(x_mu0.shape[0]), x_mu0
log_det_J = -1 * logmap_logdet(x, self.radius)
for i in reversed(range(0, self.n_blocks)):
if i > 0:
# Project between Flow Layers
z_proj_mu0 = inverse_exp_map_mu0(z, self.radius)
z = z_proj_mu0[..., 1:]
log_det_J -= logmap_logdet(z_proj_mu0, self.radius)
z_ = self.mask[i] * z
if self.layer_type != 'Linear':
s = self.s[i](z_, edge_index)
t = self.t[i](z_, edge_index)
else:
s = self.s[i](z_)
t = self.t[i](z_)
z = (1 - self.mask[i]) * (z - t) * torch.exp(-s) + z_
log_det_J -= ((1 - self.mask[i]) * s).sum(dim=1)
z_mu0 = expand_proj_dims(z)
# Project back to Manifold
z = exp_map_mu0(z_mu0, self.radius)
log_det_J -= _logdet(z_mu0, self.radius)
return z, log_det_J
def inverse(self, z_hyper):
z = inverse_exp_map_mu0(z_hyper, self.radius)
z_mu0 = z[..., 1:]
log_det_J, x = z_mu0.new_zeros(z_mu0.shape[0]), z_mu0
log_det_J = logmap_logdet(z, self.radius)
for i in range(0, self.n_blocks):
if i > 0:
# Project between Flow Layers
x_proj_mu0 = inverse_exp_map_mu0(x, self.radius)
x = x_proj_mu0[..., 1:]
log_det_J += logmap_logdet(x_proj_mu0, self.radius)
x_ = x * self.mask[i]
if self.layer_type != 'Linear':
s = self.s[i](x_, edge_index)
t = self.t[i](x_, edge_index)
else:
s = self.s[i](x_)
t = self.t[i](x_)
x = x_ + (1 - self.mask[i]) * (x * torch.exp(s) + t)
self.preclamp_norm = x.max()
x = clamp(x, min=-max_clamp_norm, max=max_clamp_norm)
log_det_J += ((1 - self.mask[i]) * s).sum(dim=1) # log det dx/du
x_mu0 = expand_proj_dims(x)
# Project back to Manifold
x = exp_map_mu0(x_mu0, self.radius)
log_det_J += _logdet(x_mu0, self.radius)
return x, log_det_J
def some_density(args):
radius = torch.Tensor([args.radius]).cuda()
n_pts = 100
f1 = lambda z: torch.sin(6 * math.pi * z[:, 0] / 4)
f2 = lambda z: 3 * torch.exp(-0.5 * ((z[:, 0] - 1) / 0.6)**2)
f3 = lambda z: 3 * torch.sigmoid((z[:, 0] - 1) / 0.3)
xx, yy, zz = setup_grid(5, n_pts)
base_prob_dist = -f1(zz)
# Map x, y coordinates on tangent space at origin to manifold (Lorentz model).
twodim = zz
threedim = expand_proj_dims(twodim).cuda()
clamped_threedim = clamp(threedim, min=-max_clamp_norm,
max=max_clamp_norm).cuda()
on_mani = exp_map_mu0(clamped_threedim, radius)
# Calculate densities of x, y coords on Lorentz model.
log_det = _logdet(clamped_threedim, radius)
log_probs = base_prob_dist - log_det
probs = torch.exp(log_probs)
# Calculate the poincare coordinates
xy_poincare = lorentz_to_poincare(on_mani.squeeze(), radius)
plot_density(xy_poincare, probs, radius, args.namestr)
if args.flow != 'none':
plot_flow(args, radius, args.flow, f1, args.namestr)
def forward(self, x_hyper, edge_index=None):
x = inverse_exp_map_mu0(x_hyper, self.radius)
x_mu0 = x[..., 1:]
log_det_J, z = x.new_zeros(x_mu0.shape[0]), x_mu0
log_det_J = -1 * logmap_logdet(x, self.radius)
for i in reversed(range(0, self.n_blocks)):
z_ = self.mask[i] * z
if self.layer_type != 'Linear':
s = self.s[i](z_, edge_index)
t_out = self.t[i](z_, edge_index)
else:
s = self.s[i](z_)
t_out = self.t[i](z_)
t_proj = proj_vec(t_out, self.radius)
t1, t_rest = t_proj[:, 0].unsqueeze(1), t_proj[:, 1:]
t = self.create_masked_t((1 - self.mask[i]), t1, t_rest)
z_2 = expand_proj_dims((1 - self.mask[i]) * z)
z_2 = clamp(z_2, min=-max_clamp_norm, max=max_clamp_norm)
z_exp_2 = exp_map_mu0(z_2, self.radius)
log_det_J -= _logdet(z_2, self.radius, subdim=(self.mask[i]).sum())
z_exp_2 = clamp(z_exp_2, min=-max_clamp_norm, max=max_clamp_norm)
z_inv_pt_arg = inverse_exp_map(x=z_exp_2,
at_point=t,
radius=self.radius)
log_det_J -= logmap_logdet(z_inv_pt_arg,
self.radius,
subdim=(self.mask[i]).sum())
z_inv_pt_arg = clamp(z_inv_pt_arg,
min=-max_clamp_norm,
max=max_clamp_norm)
pt = inverse_parallel_transport_mu0(z_inv_pt_arg,
src=t,
radius=self.radius)
pt = pt[..., 1:]
z = (1 - self.mask[i]) * pt * torch.exp(-s) + z_
log_det_J -= ((1 - self.mask[i]) * s).sum(dim=1)
z_mu0 = expand_proj_dims(z)
z = exp_map_mu0(z_mu0, self.radius)
log_det_J -= _logdet(z_mu0, self.radius)
return z, log_det_J