tfep.nn.transformers.moebius.moebius_transformer

tfep.nn.transformers.moebius.moebius_transformer(x: Tensor, w: Tensor, max_radius: float = 0.99, unit_sphere: bool = False, return_log_det_J: bool = True) → tuple[Tensor][source]

Moebius transformer.

This implements a generalization of the Moebius transformation proposed in [1, 2] to non-unit spheres. The transformer will expand/contract the distribution on the sphere of radius \(r\), where \(r\) is the norm of the input vector.

The transformation has the form

\(y = \frac{||x||^2 - ||w||^2}{||x - w||^2} (x - w) - w\)

where \(y, x, w\) are all dimension-dimensional vectors and \(||w|| < ||x||\). The function automatically rescales the w argument following the same strategy as in [2] to satisfy the condition on the norm. Consequently, ``w``s of any norm can be passed.

The implementation of the transformation on the unit sphere is slightly more efficient and can be toggled with the unit_sphere argument.

Parameters:

  • x (torch.Tensor) – Shape (batch_size, n_vectors, dimension). Input coordinates.

  • w (torch.Tensor) – Shape (batch_size, n_vectors, dimension). The transformation parameters. These parameter vectors are automatically rescaled so that |w| < |x|.

  • max_radius (float) – Must be strictly less than 1. Rescaling of the w vectors will be performed so that its maximum norm will be max_radius * |x|.

  • unit_sphere (bool) – If True, the input vectors x are assumed to be on the unit sphere, which makes the implementation slightly faster.

  • return_log_det_J (bool, optional) – Whether to return the log_det_J value. Default is True.

Returns:

  • y (torch.Tensor) – Shape (batch_size, n_vectors, dimension). The transformed vectors.

  • log_det_J (torch.Tensor, optional) – Shape (batch_size,). The logarithm of the absolute value of the Jacobian determinant dy / dx. This is returned only if return_log_det_J is True.

References

[1] Kato S, McCullagh P. Moebius transformation and a Cauchy family

on the sphere. arXiv preprint arXiv:1510.07679. 2015 Oct 26.

[2] Rezende DJ, Papamakarios G, Racanière S, Albergo MS, Kanwar G,

Shanahan PE, Cranmer K. Normalizing Flows on Tori and Spheres. arXiv preprint arXiv:2002.02428. 2020 Feb 6.