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